diff --git a/.asf.yaml b/.asf.yaml new file mode 100644 index 0000000000000..c4bd5e6f2992d --- /dev/null +++ b/.asf.yaml @@ -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. + +notifications: + commits: commits@kafka.apache.org + issues: jira@kafka.apache.org + pullrequests: jira@kafka.apache.org + jira_options: link label + +jenkins: + github_whitelist: + - ConcurrencyPractitioner + - ableegoldman + - cadonna diff --git a/.gitignore b/.gitignore index 60883492f3127..b5ca1a624ed43 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,6 @@ dist *classes +*.class target/ build/ build_eclipse/ @@ -33,8 +34,7 @@ Vagrantfile.local config/server-* config/zookeeper-* core/data/* -gradle/wrapper/* -gradlew +gradle/wrapper/*.jar gradlew.bat results @@ -47,3 +47,14 @@ tests/venv docs/generated/ .release-settings.json + +kafkatest.egg-info/ +systest/ +*.swp +clients/src/generated +clients/src/generated-test +jmh-benchmarks/generated +jmh-benchmarks/src/main/generated +streams/src/generated +raft/src/generated +core/src/generated diff --git a/.travis.yml b/.travis.yml index 8a22c9bc4cd4a..a8cd4876b9a23 100644 --- a/.travis.yml +++ b/.travis.yml @@ -10,10 +10,17 @@ # See the License for the specific language governing permissions and # limitations under the License. -sudo: required -dist: trusty +arch: amd64 +os: linux +dist: bionic language: java +services: + - docker + +jdk: + - openjdk8 + env: - _DUCKTAPE_OPTIONS="--subset 0 --subsets 15" - _DUCKTAPE_OPTIONS="--subset 1 --subsets 15" @@ -31,22 +38,14 @@ env: - _DUCKTAPE_OPTIONS="--subset 13 --subsets 15" - _DUCKTAPE_OPTIONS="--subset 14 --subsets 15" -jdk: - - oraclejdk8 - -before_install: - - gradle wrapper - script: - ./gradlew rat - ./gradlew systemTestLibs && /bin/bash ./tests/docker/run_tests.sh -services: - - docker - before_cache: - rm -f $HOME/.gradle/caches/modules-2/modules-2.lock - rm -fr $HOME/.gradle/caches/*/plugin-resolution/ + cache: directories: - "$HOME/.m2/repository" 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/Jenkinsfile b/Jenkinsfile new file mode 100644 index 0000000000000..87840d54b79a9 --- /dev/null +++ b/Jenkinsfile @@ -0,0 +1,164 @@ +/* + * + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + */ + +def setupGradle() { + // Delete gradle cache to workaround cache corruption bugs, see KAFKA-3167 + dir('.gradle') { + deleteDir() + } + sh './gradlew -version' +} + +def doValidation() { + sh ''' + ./gradlew -PscalaVersion=$SCALA_VERSION clean compileJava compileScala compileTestJava compileTestScala \ + spotlessScalaCheck checkstyleMain checkstyleTest spotbugsMain rat \ + --profile --no-daemon --continue -PxmlSpotBugsReport=true + ''' +} + +def doTest() { + sh ''' + ./gradlew -PscalaVersion=$SCALA_VERSION unitTest integrationTest \ + --profile --no-daemon --continue -PtestLoggingEvents=started,passed,skipped,failed \ + -PignoreFailures=true -PmaxParallelForks=2 -PmaxTestRetries=1 -PmaxTestRetryFailures=5 + ''' + junit '**/build/test-results/**/TEST-*.xml' +} + +def doStreamsArchetype() { + echo 'Verify that Kafka Streams archetype compiles' + + sh ''' + ./gradlew streams:install clients:install connect:json:install connect:api:install \ + || { echo 'Could not install kafka-streams.jar (and dependencies) locally`'; exit 1; } + ''' + + VERSION = sh(script: 'grep "^version=" gradle.properties | cut -d= -f 2', returnStdout: true).trim() + + dir('streams/quickstart') { + sh ''' + mvn clean install -Dgpg.skip \ + || { echo 'Could not `mvn install` streams quickstart archetype'; exit 1; } + ''' + + dir('test-streams-archetype') { + // Note the double quotes for variable interpolation + sh """ + echo "Y" | mvn archetype:generate \ + -DarchetypeCatalog=local \ + -DarchetypeGroupId=org.apache.kafka \ + -DarchetypeArtifactId=streams-quickstart-java \ + -DarchetypeVersion=${VERSION} \ + -DgroupId=streams.examples \ + -DartifactId=streams.examples \ + -Dversion=0.1 \ + -Dpackage=myapps \ + || { echo 'Could not create new project using streams quickstart archetype'; exit 1; } + """ + + dir('streams.examples') { + sh ''' + mvn compile \ + || { echo 'Could not compile streams quickstart archetype project'; exit 1; } + ''' + } + } + } +} + +def tryStreamsArchetype() { + try { + doStreamsArchetype() + } catch(err) { + echo 'Failed to build Kafka Streams archetype, marking this build UNSTABLE' + currentBuild.result = 'UNSTABLE' + } +} + + +pipeline { + agent none + stages { + stage('Build') { + parallel { + stage('JDK 8') { + agent { label 'ubuntu' } + tools { + jdk 'jdk_1.8_latest' + maven 'maven_3_latest' + } + options { + timeout(time: 8, unit: 'HOURS') + timestamps() + } + environment { + SCALA_VERSION=2.12 + } + steps { + setupGradle() + doValidation() + doTest() + tryStreamsArchetype() + } + } + + stage('JDK 11') { + agent { label 'ubuntu' } + tools { + jdk 'jdk_11_latest' + } + options { + timeout(time: 8, unit: 'HOURS') + timestamps() + } + environment { + SCALA_VERSION=2.13 + } + steps { + setupGradle() + doValidation() + doTest() + echo 'Skipping Kafka Streams archetype test for Java 11' + } + } + + stage('JDK 15') { + agent { label 'ubuntu' } + tools { + jdk 'jdk_15_latest' + } + options { + timeout(time: 8, unit: 'HOURS') + timestamps() + } + environment { + SCALA_VERSION=2.13 + } + steps { + setupGradle() + doValidation() + doTest() + echo 'Skipping Kafka Streams archetype test for Java 15' + } + } + } + } + } +} diff --git a/LICENSE b/LICENSE index bf7fe1c487b38..b0115dabb4d0f 100644 --- a/LICENSE +++ b/LICENSE @@ -201,130 +201,349 @@ See the License for the specific language governing permissions and limitations under the License. -This distribution has a binary dependency on jersey, which is available under the CDDL +------------------------------------------------------------------------------------ +This distribution has a binary dependency on jersey, which is available under the EPLv2 License as described below. -COMMON DEVELOPMENT AND DISTRIBUTION LICENSE (CDDL - Version 1.1) -1. Definitions. -1.1. “Contributor” means each individual or entity that creates or contributes to the creation of Modifications. - -1.2. “Contributor Version” means the combination of the Original Software, prior Modifications used by a Contributor (if any), and the Modifications made by that particular Contributor. - -1.3. “Covered Software” means (a) the Original Software, or (b) Modifications, or (c) the combination of files containing Original Software with files containing Modifications, in each case including portions thereof. - -1.4. “Executable” means the Covered Software in any form other than Source Code. - -1.5. “Initial Developer” means the individual or entity that first makes Original Software available under this License. - -1.6. “Larger Work” means a work which combines Covered Software or portions thereof with code not governed by the terms of this License. - -1.7. “License” means this document. - -1.8. “Licensable” means having the right to grant, to the maximum extent possible, whether at the time of the initial grant or subsequently acquired, any and all of the rights conveyed herein. - -1.9. “Modifications” means the Source Code and Executable form of any of the following: - -A. Any file that results from an addition to, deletion from or modification of the contents of a file containing Original Software or previous Modifications; - -B. Any new file that contains any part of the Original Software or previous Modification; or - -C. Any new file that is contributed or otherwise made available under the terms of this License. - -1.10. “Original Software” means the Source Code and Executable form of computer software code that is originally released under this License. - -1.11. “Patent Claims” means any patent claim(s), now owned or hereafter acquired, including without limitation, method, process, and apparatus claims, in any patent Licensable by grantor. - -1.12. “Source Code” means (a) the common form of computer software code in which modifications are made and (b) associated documentation included in or with such code. - -1.13. “You” (or “Your”) means an individual or a legal entity exercising rights under, and complying with all of the terms of, this License. For legal entities, “You” includes any entity which controls, is controlled by, or is under common control with You. For purposes of this definition, “control” means (a) the power, direct or indirect, to cause the direction or management of such entity, whether by contract or otherwise, or (b) ownership of more than fifty percent (50%) of the outstanding shares or beneficial ownership of such entity. - -2. License Grants. -2.1. The Initial Developer Grant. - -Conditioned upon Your compliance with Section 3.1 below and subject to third party intellectual property claims, the Initial Developer hereby grants You a world-wide, royalty-free, non-exclusive license: - -(a) under intellectual property rights (other than patent or trademark) Licensable by Initial Developer, to use, reproduce, modify, display, perform, sublicense and distribute the Original Software (or portions thereof), with or without Modifications, and/or as part of a Larger Work; and - -(b) under Patent Claims infringed by the making, using or selling of Original Software, to make, have made, use, practice, sell, and offer for sale, and/or otherwise dispose of the Original Software (or portions thereof). - -(c) The licenses granted in Sections 2.1(a) and (b) are effective on the date Initial Developer first distributes or otherwise makes the Original Software available to a third party under the terms of this License. - -(d) Notwithstanding Section 2.1(b) above, no patent license is granted: (1) for code that You delete from the Original Software, or (2) for infringements caused by: (i) the modification of the Original Software, or (ii) the combination of the Original Software with other software or devices. - -2.2. Contributor Grant. - -Conditioned upon Your compliance with Section 3.1 below and subject to third party intellectual property claims, each Contributor hereby grants You a world-wide, royalty-free, non-exclusive license: - -(a) under intellectual property rights (other than patent or trademark) Licensable by Contributor to use, reproduce, modify, display, perform, sublicense and distribute the Modifications created by such Contributor (or portions thereof), either on an unmodified basis, with other Modifications, as Covered Software and/or as part of a Larger Work; and - -(b) under Patent Claims infringed by the making, using, or selling of Modifications made by that Contributor either alone and/or in combination with its Contributor Version (or portions of such combination), to make, use, sell, offer for sale, have made, and/or otherwise dispose of: (1) Modifications made by that Contributor (or portions thereof); and (2) the combination of Modifications made by that Contributor with its Contributor Version (or portions of such combination). - -(c) The licenses granted in Sections 2.2(a) and 2.2(b) are effective on the date Contributor first distributes or otherwise makes the Modifications available to a third party. - -(d) Notwithstanding Section 2.2(b) above, no patent license is granted: (1) for any code that Contributor has deleted from the Contributor Version; (2) for infringements caused by: (i) third party modifications of Contributor Version, or (ii) the combination of Modifications made by that Contributor with other software (except as part of the Contributor Version) or other devices; or (3) under Patent Claims infringed by Covered Software in the absence of Modifications made by that Contributor. - -3. Distribution Obligations. -3.1. Availability of Source Code. - -Any Covered Software that You distribute or otherwise make available in Executable form must also be made available in Source Code form and that Source Code form must be distributed only under the terms of this License. You must include a copy of this License with every copy of the Source Code form of the Covered Software You distribute or otherwise make available. You must inform recipients of any such Covered Software in Executable form as to how they can obtain such Covered Software in Source Code form in a reasonable manner on or through a medium customarily used for software exchange. - -3.2. Modifications. - -The Modifications that You create or to which You contribute are governed by the terms of this License. You represent that You believe Your Modifications are Your original creation(s) and/or You have sufficient rights to grant the rights conveyed by this License. - -3.3. Required Notices. - -You must include a notice in each of Your Modifications that identifies You as the Contributor of the Modification. You may not remove or alter any copyright, patent or trademark notices contained within the Covered Software, or any notices of licensing or any descriptive text giving attribution to any Contributor or the Initial Developer. - -3.4. Application of Additional Terms. - -You may not offer or impose any terms on any Covered Software in Source Code form that alters or restricts the applicable version of this License or the recipients’ rights hereunder. You may choose to offer, and to charge a fee for, warranty, support, indemnity or liability obligations to one or more recipients of Covered Software. However, you may do so only on Your own behalf, and not on behalf of the Initial Developer or any Contributor. You must make it absolutely clear that any such warranty, support, indemnity or liability obligation is offered by You alone, and You hereby agree to indemnify the Initial Developer and every Contributor for any liability incurred by the Initial Developer or such Contributor as a result of warranty, support, indemnity or liability terms You offer. - -3.5. Distribution of Executable Versions. - -You may distribute the Executable form of the Covered Software under the terms of this License or under the terms of a license of Your choice, which may contain terms different from this License, provided that You are in compliance with the terms of this License and that the license for the Executable form does not attempt to limit or alter the recipient’s rights in the Source Code form from the rights set forth in this License. If You distribute the Covered Software in Executable form under a different license, You must make it absolutely clear that any terms which differ from this License are offered by You alone, not by the Initial Developer or Contributor. You hereby agree to indemnify the Initial Developer and every Contributor for any liability incurred by the Initial Developer or such Contributor as a result of any such terms You offer. - -3.6. Larger Works. - -You may create a Larger Work by combining Covered Software with other code not governed by the terms of this License and distribute the Larger Work as a single product. In such a case, You must make sure the requirements of this License are fulfilled for the Covered Software. - -4. Versions of the License. -4.1. New Versions. - -Oracle is the initial license steward and may publish revised and/or new versions of this License from time to time. Each version will be given a distinguishing version number. Except as provided in Section 4.3, no one other than the license steward has the right to modify this License. - -4.2. Effect of New Versions. - -You may always continue to use, distribute or otherwise make the Covered Software available under the terms of the version of the License under which You originally received the Covered Software. If the Initial Developer includes a notice in the Original Software prohibiting it from being distributed or otherwise made available under any subsequent version of the License, You must distribute and make the Covered Software available under the terms of the version of the License under which You originally received the Covered Software. Otherwise, You may also choose to use, distribute or otherwise make the Covered Software available under the terms of any subsequent version of the License published by the license steward. - -4.3. Modified Versions. - -When You are an Initial Developer and You want to create a new license for Your Original Software, You may create and use a modified version of this License if You: (a) rename the license and remove any references to the name of the license steward (except to note that the license differs from this License); and (b) otherwise make it clear that the license contains terms which differ from this License. - -5. DISCLAIMER OF WARRANTY. -COVERED SOFTWARE IS PROVIDED UNDER THIS LICENSE ON AN “AS IS” BASIS, WITHOUT WARRANTY OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, WITHOUT LIMITATION, WARRANTIES THAT THE COVERED SOFTWARE IS FREE OF DEFECTS, MERCHANTABLE, FIT FOR A PARTICULAR PURPOSE OR NON-INFRINGING. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE COVERED SOFTWARE IS WITH YOU. SHOULD ANY COVERED SOFTWARE PROVE DEFECTIVE IN ANY RESPECT, YOU (NOT THE INITIAL DEVELOPER OR ANY OTHER CONTRIBUTOR) ASSUME THE COST OF ANY NECESSARY SERVICING, REPAIR OR CORRECTION. THIS DISCLAIMER OF WARRANTY CONSTITUTES AN ESSENTIAL PART OF THIS LICENSE. NO USE OF ANY COVERED SOFTWARE IS AUTHORIZED HEREUNDER EXCEPT UNDER THIS DISCLAIMER. - -6. TERMINATION. -6.1. This License and the rights granted hereunder will terminate automatically if You fail to comply with terms herein and fail to cure such breach within 30 days of becoming aware of the breach. Provisions which, by their nature, must remain in effect beyond the termination of this License shall survive. - -6.2. If You assert a patent infringement claim (excluding declaratory judgment actions) against Initial Developer or a Contributor (the Initial Developer or Contributor against whom You assert such claim is referred to as “Participant”) alleging that the Participant Software (meaning the Contributor Version where the Participant is a Contributor or the Original Software where the Participant is the Initial Developer) directly or indirectly infringes any patent, then any and all rights granted directly or indirectly to You by such Participant, the Initial Developer (if the Initial Developer is not the Participant) and all Contributors under Sections 2.1 and/or 2.2 of this License shall, upon 60 days notice from Participant terminate prospectively and automatically at the expiration of such 60 day notice period, unless if within such 60 day period You withdraw Your claim with respect to the Participant Software against such Participant either unilaterally or pursuant to a written agreement with Participant. - -6.3. If You assert a patent infringement claim against Participant alleging that the Participant Software directly or indirectly infringes any patent where such claim is resolved (such as by license or settlement) prior to the initiation of patent infringement litigation, then the reasonable value of the licenses granted by such Participant under Sections 2.1 or 2.2 shall be taken into account in determining the amount or value of any payment or license. - -6.4. In the event of termination under Sections 6.1 or 6.2 above, all end user licenses that have been validly granted by You or any distributor hereunder prior to termination (excluding licenses granted to You by any distributor) shall survive termination. - -7. LIMITATION OF LIABILITY. -UNDER NO CIRCUMSTANCES AND UNDER NO LEGAL THEORY, WHETHER TORT (INCLUDING NEGLIGENCE), CONTRACT, OR OTHERWISE, SHALL YOU, THE INITIAL DEVELOPER, ANY OTHER CONTRIBUTOR, OR ANY DISTRIBUTOR OF COVERED SOFTWARE, OR ANY SUPPLIER OF ANY OF SUCH PARTIES, BE LIABLE TO ANY PERSON FOR ANY INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES OF ANY CHARACTER INCLUDING, WITHOUT LIMITATION, DAMAGES FOR LOSS OF GOODWILL, WORK STOPPAGE, COMPUTER FAILURE OR MALFUNCTION, OR ANY AND ALL OTHER COMMERCIAL DAMAGES OR LOSSES, EVEN IF SUCH PARTY SHALL HAVE BEEN INFORMED OF THE POSSIBILITY OF SUCH DAMAGES. THIS LIMITATION OF LIABILITY SHALL NOT APPLY TO LIABILITY FOR DEATH OR PERSONAL INJURY RESULTING FROM SUCH PARTY’S NEGLIGENCE TO THE EXTENT APPLICABLE LAW PROHIBITS SUCH LIMITATION. SOME JURISDICTIONS DO NOT ALLOW THE EXCLUSION OR LIMITATION OF INCIDENTAL OR CONSEQUENTIAL DAMAGES, SO THIS EXCLUSION AND LIMITATION MAY NOT APPLY TO YOU. - -8. U.S. GOVERNMENT END USERS. -The Covered Software is a “commercial item,” as that term is defined in 48 C.F.R. 2.101 (Oct. 1995), consisting of “commercial computer software” (as that term is defined at 48 C.F.R. § 252.227-7014(a)(1)) and “commercial computer software documentation” as such terms are used in 48 C.F.R. 12.212 (Sept. 1995). Consistent with 48 C.F.R. 12.212 and 48 C.F.R. 227.7202-1 through 227.7202-4 (June 1995), all U.S. Government End Users acquire Covered Software with only those rights set forth herein. This U.S. Government Rights clause is in lieu of, and supersedes, any other FAR, DFAR, or other clause or provision that addresses Government rights in computer software under this License. - -9. MISCELLANEOUS. -This License represents the complete agreement concerning subject matter hereof. If any provision of this License is held to be unenforceable, such provision shall be reformed only to the extent necessary to make it enforceable. This License shall be governed by the law of the jurisdiction specified in a notice contained within the Original Software (except to the extent applicable law, if any, provides otherwise), excluding such jurisdiction’s conflict-of-law provisions. Any litigation relating to this License shall be subject to the jurisdiction of the courts located in the jurisdiction and venue specified in a notice contained within the Original Software, with the losing party responsible for costs, including, without limitation, court costs and reasonable attorneys’ fees and expenses. The application of the United Nations Convention on Contracts for the International Sale of Goods is expressly excluded. Any law or regulation which provides that the language of a contract shall be construed against the drafter shall not apply to this License. You agree that You alone are responsible for compliance with the United States export administration regulations (and the export control laws and regulation of any other countries) when You use, distribute or otherwise make available any Covered Software. - -10. RESPONSIBILITY FOR CLAIMS. -As between Initial Developer and the Contributors, each party is responsible for claims and damages arising, directly or indirectly, out of its utilization of rights under this License and You agree to work with Initial Developer and Contributors to distribute such responsibility on an equitable basis. Nothing herein is intended or shall be deemed to constitute any admission of liability. - -NOTICE PURSUANT TO SECTION 9 OF THE COMMON DEVELOPMENT AND DISTRIBUTION LICENSE (CDDL) - -The code released under the CDDL shall be governed by the laws of the State of California (excluding conflict-of-law provisions). Any litigation relating to this License shall be subject to the jurisdiction of the Federal Courts of the Northern District of California and the state courts of the State of California, with venue lying in Santa Clara County, California. +Eclipse Public License - v 2.0 + + THE ACCOMPANYING PROGRAM IS PROVIDED UNDER THE TERMS OF THIS ECLIPSE + PUBLIC LICENSE ("AGREEMENT"). ANY USE, REPRODUCTION OR DISTRIBUTION + OF THE PROGRAM CONSTITUTES RECIPIENT'S ACCEPTANCE OF THIS AGREEMENT. + + 1. DEFINITIONS + + "Contribution" means: + + a) in the case of the initial Contributor, the initial content + Distributed under this Agreement, and + + b) in the case of each subsequent Contributor: + i) changes to the Program, and + ii) additions to the Program; + where such changes and/or additions to the Program originate from + and are Distributed by that particular Contributor. A Contribution + "originates" from a Contributor if it was added to the Program by + such Contributor itself or anyone acting on such Contributor's behalf. + Contributions do not include changes or additions to the Program that + are not Modified Works. + + "Contributor" means any person or entity that Distributes the Program. + + "Licensed Patents" mean patent claims licensable by a Contributor which + are necessarily infringed by the use or sale of its Contribution alone + or when combined with the Program. + + "Program" means the Contributions Distributed in accordance with this + Agreement. + + "Recipient" means anyone who receives the Program under this Agreement + or any Secondary License (as applicable), including Contributors. + + "Derivative Works" shall mean any work, whether in Source Code or other + form, that is based on (or derived from) the Program and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. + + "Modified Works" shall mean any work in Source Code or other form that + results from an addition to, deletion from, or modification of the + contents of the Program, including, for purposes of clarity any new file + in Source Code form that contains any contents of the Program. Modified + Works shall not include works that contain only declarations, + interfaces, types, classes, structures, or files of the Program solely + in each case in order to link to, bind by name, or subclass the Program + or Modified Works thereof. + + "Distribute" means the acts of a) distributing or b) making available + in any manner that enables the transfer of a copy. + + "Source Code" means the form of a Program preferred for making + modifications, including but not limited to software source code, + documentation source, and configuration files. + + "Secondary License" means either the GNU General Public License, + Version 2.0, or any later versions of that license, including any + exceptions or additional permissions as identified by the initial + Contributor. + + 2. GRANT OF RIGHTS + + a) Subject to the terms of this Agreement, each Contributor hereby + grants Recipient a non-exclusive, worldwide, royalty-free copyright + license to reproduce, prepare Derivative Works of, publicly display, + publicly perform, Distribute and sublicense the Contribution of such + Contributor, if any, and such Derivative Works. + + b) Subject to the terms of this Agreement, each Contributor hereby + grants Recipient a non-exclusive, worldwide, royalty-free patent + license under Licensed Patents to make, use, sell, offer to sell, + import and otherwise transfer the Contribution of such Contributor, + if any, in Source Code or other form. This patent license shall + apply to the combination of the Contribution and the Program if, at + the time the Contribution is added by the Contributor, such addition + of the Contribution causes such combination to be covered by the + Licensed Patents. The patent license shall not apply to any other + combinations which include the Contribution. No hardware per se is + licensed hereunder. + + c) Recipient understands that although each Contributor grants the + licenses to its Contributions set forth herein, no assurances are + provided by any Contributor that the Program does not infringe the + patent or other intellectual property rights of any other entity. + Each Contributor disclaims any liability to Recipient for claims + brought by any other entity based on infringement of intellectual + property rights or otherwise. As a condition to exercising the + rights and licenses granted hereunder, each Recipient hereby + assumes sole responsibility to secure any other intellectual + property rights needed, if any. For example, if a third party + patent license is required to allow Recipient to Distribute the + Program, it is Recipient's responsibility to acquire that license + before distributing the Program. + + d) Each Contributor represents that to its knowledge it has + sufficient copyright rights in its Contribution, if any, to grant + the copyright license set forth in this Agreement. + + e) Notwithstanding the terms of any Secondary License, no + Contributor makes additional grants to any Recipient (other than + those set forth in this Agreement) as a result of such Recipient's + receipt of the Program under the terms of a Secondary License + (if permitted under the terms of Section 3). + + 3. REQUIREMENTS + + 3.1 If a Contributor Distributes the Program in any form, then: + + a) the Program must also be made available as Source Code, in + accordance with section 3.2, and the Contributor must accompany + the Program with a statement that the Source Code for the Program + is available under this Agreement, and informs Recipients how to + obtain it in a reasonable manner on or through a medium customarily + used for software exchange; and + + b) the Contributor may Distribute the Program under a license + different than this Agreement, provided that such license: + i) effectively disclaims on behalf of all other Contributors all + warranties and conditions, express and implied, including + warranties or conditions of title and non-infringement, and + implied warranties or conditions of merchantability and fitness + for a particular purpose; + + ii) effectively excludes on behalf of all other Contributors all + liability for damages, including direct, indirect, special, + incidental and consequential damages, such as lost profits; + + iii) does not attempt to limit or alter the recipients' rights + in the Source Code under section 3.2; and + + iv) requires any subsequent distribution of the Program by any + party to be under a license that satisfies the requirements + of this section 3. + + 3.2 When the Program is Distributed as Source Code: + + a) it must be made available under this Agreement, or if the + Program (i) is combined with other material in a separate file or + files made available under a Secondary License, and (ii) the initial + Contributor attached to the Source Code the notice described in + Exhibit A of this Agreement, then the Program may be made available + under the terms of such Secondary Licenses, and + + b) a copy of this Agreement must be included with each copy of + the Program. + + 3.3 Contributors may not remove or alter any copyright, patent, + trademark, attribution notices, disclaimers of warranty, or limitations + of liability ("notices") contained within the Program from any copy of + the Program which they Distribute, provided that Contributors may add + their own appropriate notices. + + 4. COMMERCIAL DISTRIBUTION + + Commercial distributors of software may accept certain responsibilities + with respect to end users, business partners and the like. While this + license is intended to facilitate the commercial use of the Program, + the Contributor who includes the Program in a commercial product + offering should do so in a manner which does not create potential + liability for other Contributors. Therefore, if a Contributor includes + the Program in a commercial product offering, such Contributor + ("Commercial Contributor") hereby agrees to defend and indemnify every + other Contributor ("Indemnified Contributor") against any losses, + damages and costs (collectively "Losses") arising from claims, lawsuits + and other legal actions brought by a third party against the Indemnified + Contributor to the extent caused by the acts or omissions of such + Commercial Contributor in connection with its distribution of the Program + in a commercial product offering. The obligations in this section do not + apply to any claims or Losses relating to any actual or alleged + intellectual property infringement. In order to qualify, an Indemnified + Contributor must: a) promptly notify the Commercial Contributor in + writing of such claim, and b) allow the Commercial Contributor to control, + and cooperate with the Commercial Contributor in, the defense and any + related settlement negotiations. The Indemnified Contributor may + participate in any such claim at its own expense. + + For example, a Contributor might include the Program in a commercial + product offering, Product X. That Contributor is then a Commercial + Contributor. If that Commercial Contributor then makes performance + claims, or offers warranties related to Product X, those performance + claims and warranties are such Commercial Contributor's responsibility + alone. Under this section, the Commercial Contributor would have to + defend claims against the other Contributors related to those performance + claims and warranties, and if a court requires any other Contributor to + pay any damages as a result, the Commercial Contributor must pay + those damages. + + 5. NO WARRANTY + + EXCEPT AS EXPRESSLY SET FORTH IN THIS AGREEMENT, AND TO THE EXTENT + PERMITTED BY APPLICABLE LAW, THE PROGRAM IS PROVIDED ON AN "AS IS" + BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, EITHER EXPRESS OR + IMPLIED INCLUDING, WITHOUT LIMITATION, ANY WARRANTIES OR CONDITIONS OF + TITLE, NON-INFRINGEMENT, MERCHANTABILITY OR FITNESS FOR A PARTICULAR + PURPOSE. Each Recipient is solely responsible for determining the + appropriateness of using and distributing the Program and assumes all + risks associated with its exercise of rights under this Agreement, + including but not limited to the risks and costs of program errors, + compliance with applicable laws, damage to or loss of data, programs + or equipment, and unavailability or interruption of operations. + + 6. DISCLAIMER OF LIABILITY + + EXCEPT AS EXPRESSLY SET FORTH IN THIS AGREEMENT, AND TO THE EXTENT + PERMITTED BY APPLICABLE LAW, NEITHER RECIPIENT NOR ANY CONTRIBUTORS + SHALL HAVE ANY LIABILITY FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, + EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING WITHOUT LIMITATION LOST + PROFITS), HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN + CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) + ARISING IN ANY WAY OUT OF THE USE OR DISTRIBUTION OF THE PROGRAM OR THE + EXERCISE OF ANY RIGHTS GRANTED HEREUNDER, EVEN IF ADVISED OF THE + POSSIBILITY OF SUCH DAMAGES. + + 7. GENERAL + + If any provision of this Agreement is invalid or unenforceable under + applicable law, it shall not affect the validity or enforceability of + the remainder of the terms of this Agreement, and without further + action by the parties hereto, such provision shall be reformed to the + minimum extent necessary to make such provision valid and enforceable. + + If Recipient institutes patent litigation against any entity + (including a cross-claim or counterclaim in a lawsuit) alleging that the + Program itself (excluding combinations of the Program with other software + or hardware) infringes such Recipient's patent(s), then such Recipient's + rights granted under Section 2(b) shall terminate as of the date such + litigation is filed. + + All Recipient's rights under this Agreement shall terminate if it + fails to comply with any of the material terms or conditions of this + Agreement and does not cure such failure in a reasonable period of + time after becoming aware of such noncompliance. If all Recipient's + rights under this Agreement terminate, Recipient agrees to cease use + and distribution of the Program as soon as reasonably practicable. + However, Recipient's obligations under this Agreement and any licenses + granted by Recipient relating to the Program shall continue and survive. + + Everyone is permitted to copy and distribute copies of this Agreement, + but in order to avoid inconsistency the Agreement is copyrighted and + may only be modified in the following manner. The Agreement Steward + reserves the right to publish new versions (including revisions) of + this Agreement from time to time. No one other than the Agreement + Steward has the right to modify this Agreement. The Eclipse Foundation + is the initial Agreement Steward. The Eclipse Foundation may assign the + responsibility to serve as the Agreement Steward to a suitable separate + entity. Each new version of the Agreement will be given a distinguishing + version number. The Program (including Contributions) may always be + Distributed subject to the version of the Agreement under which it was + received. In addition, after a new version of the Agreement is published, + Contributor may elect to Distribute the Program (including its + Contributions) under the new version. + + Except as expressly stated in Sections 2(a) and 2(b) above, Recipient + receives no rights or licenses to the intellectual property of any + Contributor under this Agreement, whether expressly, by implication, + estoppel or otherwise. All rights in the Program not expressly granted + under this Agreement are reserved. Nothing in this Agreement is intended + to be enforceable by any entity that is not a Contributor or Recipient. + No third-party beneficiary rights are created under this Agreement. + + Exhibit A - Form of Secondary Licenses Notice + + "This Source Code may also be made available under the following + Secondary Licenses when the conditions for such availability set forth + in the Eclipse Public License, v. 2.0 are satisfied: {name license(s), + version(s), and exceptions or additional permissions here}." + + Simply including a copy of this Agreement, including this Exhibit A + is not sufficient to license the Source Code under Secondary Licenses. + + If it is not possible or desirable to put the notice in a particular + file, then You may include the notice in a location (such as a LICENSE + file in a relevant directory) where a recipient would be likely to + look for such a notice. + + You may add additional accurate notices of copyright ownership. + +------------------------------------------------------------------------------------ +This distribution has a binary dependency on zstd, which is available under the BSD 3-Clause License as described below. + +BSD License + +For Zstandard software + +Copyright (c) 2016-present, Facebook, Inc. All rights reserved. + +Redistribution and use in source and binary forms, with or without modification, +are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + + * Neither the name Facebook nor the names of its contributors may be used to + endorse or promote products derived from this software without specific + prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR +ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON +ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +------------------------------------------------------------------------------------ +This distribution has a binary dependency on zstd-jni, which is available under the BSD 2-Clause License +as described below. + +Zstd-jni: JNI bindings to Zstd Library + +Copyright (c) 2015-2016, Luben Karavelov/ All rights reserved. + +BSD License + +Redistribution and use in source and binary forms, with or without modification, +are permitted provided that the following conditions are met: + +* Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +* Redistributions in binary form must reproduce the above copyright notice, this + list of conditions and the following disclaimer in the documentation and/or + other materials provided with the distribution. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR +ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON +ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/NOTICE b/NOTICE index 7a62b7fed60cd..f5505f39ad215 100644 --- a/NOTICE +++ b/NOTICE @@ -1,8 +1,8 @@ Apache Kafka -Copyright 2017 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 3a1d82a0e002c..34e13b6ea7da7 100644 --- a/README.md +++ b/README.md @@ -1,23 +1,19 @@ Apache Kafka ================= -See our [web site](http://kafka.apache.org) for details on the project. +See our [web site](https://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. +You need to have [Java](http://www.oracle.com/technetwork/java/javase/downloads/index.html) installed. -Kafka requires Gradle 2.0 or higher. +We build and test Apache Kafka with Java 8, 11 and 15. We set the `release` parameter in javac and scalac +to `8` to ensure the generated binaries are compatible with Java 8 or higher (independently of the Java version +used for compilation). -Java 7 should be used for building in order to support both Java 7 and Java 8 at runtime. - -### First bootstrap and download the wrapper ### - cd kafka_source_dir - gradle - -Now everything else will work. +Scala 2.13 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/quickstart ### Build source jar ### ./gradlew srcJar @@ -43,7 +39,7 @@ Follow instructions in http://kafka.apache.org/documentation.html#quickstart ./gradlew cleanTest integrationTest ### Running a particular unit/integration test ### - ./gradlew -Dtest.single=RequestResponseSerializationTest core:test + ./gradlew clients:test --tests RequestResponseTest ### Running a particular test method within a unit/integration test ### ./gradlew core:test --tests kafka.api.ProducerFailureHandlingTest.testCannotSendToInternalTopic @@ -52,40 +48,58 @@ Follow instructions in http://kafka.apache.org/documentation.html#quickstart ### Running a particular unit/integration test with log4j output ### Change the log4j setting in either `clients/src/test/resources/log4j.properties` or `core/src/test/resources/log4j.properties` - ./gradlew -i -Dtest.single=RequestResponseSerializationTest core:test + ./gradlew clients:test --tests RequestResponseTest + +### Specifying test retries ### +By default, each failed test is retried once up to a maximum of five retries per test run. Tests are retried at the end of the test task. Adjust these parameters in the following way: + + ./gradlew test -PmaxTestRetries=1 -PmaxTestRetryFailures=5 + +See [Test Retry Gradle Plugin](https://github.com/gradle/test-retry-gradle-plugin) for more details. ### Generating test coverage reports ### Generate coverage reports for the whole project: - ./gradlew reportCoverage + ./gradlew reportCoverage -PenableTestCoverage=true Generate coverage for a single module, i.e.: - ./gradlew clients:reportCoverage + ./gradlew clients:reportCoverage -PenableTestCoverage=true ### Building a binary release gzipped tar ball ### - ./gradlew clean - ./gradlew releaseTarGz + ./gradlew clean releaseTarGz The above command will fail if you haven't set up the signing key. To bypass signing the artifact, you can run: - ./gradlew releaseTarGz -x signArchives + ./gradlew clean releaseTarGz -x signArchives The release file can be found inside `./core/build/distributions/`. +### Building auto generated messages ### +Sometimes it is only necessary to rebuild the RPC auto-generated message data when switching between branches, as they could +fail due to code changes. You can just run: + + ./gradlew processMessages processTestMessages + ### Cleaning the build ### ./gradlew clean -### Running a task on a particular version of Scala (either 2.11.x or 2.12.x) ### -*Note that if building the jars with a version other than 2.11.12, you need to set the `SCALA_VERSION` variable or change it in `bin/kafka-run-class.sh` to run the quick start.* +### Running a task with one of the Scala versions available (2.12.x or 2.13.x) ### +*Note that if building the jars with a version other than 2.13.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): + + ./gradlew -PscalaVersion=2.12 jar + ./gradlew -PscalaVersion=2.12 test + ./gradlew -PscalaVersion=2.12 releaseTarGz -You can pass either the major version (eg 2.11) or the full version (eg 2.11.12): +### Running a task with all the scala versions enabled by default ### - ./gradlew -PscalaVersion=2.11 jar - ./gradlew -PscalaVersion=2.11 test - ./gradlew -PscalaVersion=2.11 releaseTarGz +Invoke the `gradlewAll` script followed by the task(s): -Scala 2.12.x requires Java 8. + ./gradlewAll test + ./gradlewAll jar + ./gradlewAll releaseTarGz ### Running a task for a specific project ### This is for `core`, `examples` and `clients` @@ -93,6 +107,10 @@ This is for `core`, `examples` and `clients` ./gradlew core:jar ./gradlew core:test +Streams has multiple sub-projects, but you can run all the tests: + + ./gradlew :streams:testAll + ### Listing all gradle tasks ### ./gradlew tasks @@ -106,17 +124,8 @@ The `eclipse` task has been configured to use `${project_dir}/build_eclipse` as build directory (`${project_dir}/bin`) clashes with Kafka's scripts directory and we don't use Gradle's build directory to avoid known issues with this configuration. -### Building the jar for all scala versions and for all projects ### - ./gradlew jarAll - -### Running unit/integration tests for all scala versions and for all projects ### - ./gradlew testAll - -### Building a binary release gzipped tar ball for all scala versions ### - ./gradlew releaseTarGzAll - ### Publishing the jar for all version of Scala and for all projects to maven ### - ./gradlew uploadArchivesAll + ./gradlewAll uploadArchives Please note for this to work you should create/update `${GRADLE_USER_HOME}/gradle.properties` (typically, `~/.gradle/gradle.properties`) and assign the following variables @@ -158,7 +167,7 @@ Please note for this to work you should create/update user maven settings (typic ### Installing the jars to the local Maven repository ### - ./gradlew installAll + ./gradlewAll install ### Building the test jar ### ./gradlew testJar @@ -170,7 +179,7 @@ Please note for this to work you should create/update user maven settings (typic ./gradlew dependencyUpdates ### Running code quality checks ### -There are two code quality analysis tools that we regularly run, findbugs and checkstyle. +There are two code quality analysis tools that we regularly run, spotbugs and checkstyle. #### Checkstyle #### Checkstyle enforces a consistent coding style in Kafka. @@ -179,16 +188,21 @@ You can run checkstyle using: ./gradlew checkstyleMain checkstyleTest The checkstyle warnings will be found in `reports/checkstyle/reports/main.html` and `reports/checkstyle/reports/test.html` files in the -subproject build directories. They are also are printed to the console. The build will fail if Checkstyle fails. +subproject build directories. They are also printed to the console. The build will fail if Checkstyle fails. + +#### Spotbugs #### +Spotbugs uses static analysis to look for bugs in the code. +You can run spotbugs using: -#### Findbugs #### -Findbugs uses static analysis to look for bugs in the code. -You can run findbugs using: + ./gradlew spotbugsMain spotbugsTest -x test - ./gradlew findbugsMain findbugsTest -x test +The spotbugs warnings will be found in `reports/spotbugs/main.html` and `reports/spotbugs/test.html` files in the subproject build +directories. Use -PxmlSpotBugsReport=true to generate an XML report instead of an HTML one. -The findbugs warnings will be found in `reports/findbugs/main.html` and `reports/findbugs/test.html` files in the subproject build -directories. Use -PxmlFindBugsReport=true to generate an XML report instead of an HTML one. +### JMH microbenchmarks ### +We use [JMH](https://openjdk.java.net/projects/code-tools/jmh/) to write microbenchmarks that produce reliable results in the JVM. + +See [jmh-benchmarks/README.md](https://github.com/apache/kafka/blob/trunk/jmh-benchmarks/README.md) for details on how to run the microbenchmarks. ### Common build options ### @@ -197,10 +211,32 @@ The following options should be set with a `-P` switch, for example `./gradlew - * `commitId`: sets the build commit ID as .git/HEAD might not be correct if there are local commits added for build purposes. * `mavenUrl`: sets the URL of the maven deployment repository (`file://path/to/repo` can be used to point to a local repository). * `maxParallelForks`: limits the maximum number of processes for each task. +* `ignoreFailures`: ignore test failures from junit * `showStandardStreams`: shows standard out and standard error of the test JVM(s) on the console. * `skipSigning`: skips signing of artifacts. * `testLoggingEvents`: unit test events to be logged, separated by comma. For example `./gradlew -PtestLoggingEvents=started,passed,skipped,failed test`. -* `xmlFindBugsReport`: enable XML reports for findBugs. This also disables HTML reports as only one can be enabled at a time. +* `xmlSpotBugsReport`: enable XML reports for spotBugs. This also disables HTML reports as only one can be enabled at a time. +* `maxTestRetries`: the maximum number of retries for a failing test case. +* `maxTestRetryFailures`: maximum number of test failures before retrying is disabled for subsequent tests. +* `enableTestCoverage`: enables test coverage plugins and tasks, including bytecode enhancement of classes required to track said +coverage. Note that this introduces some overhead when running tests and hence why it's disabled by default (the overhead +varies, but 15-20% is a reasonable estimate). + +### 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 ### @@ -211,4 +247,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 new file mode 100644 index 0000000000000..3891857562f93 --- /dev/null +++ b/TROGDOR.md @@ -0,0 +1,189 @@ +Trogdor +======================================== +Trogdor is a test framework for Apache Kafka. + +Trogdor can run benchmarks and other workloads. Trogdor can also inject faults in order to stress test the system. + +Quickstart +========================================================= +First, we want to start a single-node Kafka cluster with a ZooKeeper and a broker. + +Running ZooKeeper: + + > ./bin/zookeeper-server-start.sh ./config/zookeeper.properties &> /tmp/zookeeper.log & + +Running Kafka: + + > ./bin/kafka-server-start.sh ./config/server.properties &> /tmp/kafka.log & + +Then, we want to run a Trogdor Agent, plus a Trogdor Coordinator. + +To run the Trogdor Agent: + + > ./bin/trogdor.sh agent -c ./config/trogdor.conf -n node0 &> /tmp/trogdor-agent.log & + +To run the Trogdor Coordinator: + + > ./bin/trogdor.sh coordinator -c ./config/trogdor.conf -n node0 &> /tmp/trogdor-coordinator.log & + +Let's confirm that all of the daemons are running: + + > jps + 116212 Coordinator + 115188 QuorumPeerMain + 116571 Jps + 115420 Kafka + 115694 Agent + +Now, we can submit a test job to Trogdor. + + > ./bin/trogdor.sh client createTask -t localhost:8889 -i produce0 --spec ./tests/spec/simple_produce_bench.json + Sent CreateTaskRequest for task produce0. + +We can run showTask to see what the task's status is: + + > ./bin/trogdor.sh client showTask -t localhost:8889 -i produce0 + 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: { + "totalSent" : 50000, + "averageLatencyMs" : 17.83388, + "p50LatencyMs" : 12, + "p95LatencyMs" : 75, + "p99LatencyMs" : 96, + "transactionsCommitted" : 0 + } + +Trogdor Architecture +======================================== +Trogdor has a single coordinator process which manages multiple agent processes. Each agent process is responsible for a single cluster node. + +The Trogdor coordinator manages tasks. A task is anything we might want to do on a cluster, such as running a benchmark, injecting a fault, or running a workload. In order to implement each task, the coordinator creates workers on one or more agent nodes. + +The Trogdor agent process implements the tasks. For example, when running a workload, the agent process is the process which produces and consumes messages. + +Both the coordinator and the agent expose a REST interface that accepts objects serialized via JSON. There is also a command-line program which makes it easy to send messages to either one without manually crafting the JSON message body. + +All Trogdor RPCs are idempotent except the shutdown requests. Sending an idempotent RPC twice in a row has the same effect as sending the RPC once. + +Tasks +======================================== +Tasks are described by specifications containing: + +* A "class" field describing the task type. This contains a full Java class name. +* A "startMs" field describing when the task should start. This is given in terms of milliseconds since the UNIX epoch. +* A "durationMs" field describing how long the task should last. This is given in terms of milliseconds. +* Other fields which are task-specific. + +The task specification is usually written as JSON. For example, this task specification describes a network partition between nodes 1 and 2, and 3: + + { + "class": "org.apache.kafka.trogdor.fault.NetworkPartitionFaultSpec", + "startMs": 1000, + "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. + +Task specifications are immutable; they do not change after the task has been created. + +Tasks can be in several states: +* PENDING, when task is waiting to execute, +* RUNNING, when the task is running, +* STOPPING, when the task is in the process of stopping, +* DONE, when the task is done. + +Tasks that are DONE also have an error field which will be set if the task failed. + +Workloads +======================================== +Trogdor can run several workloads. Workloads perform operations on the cluster and measure their performance. Workloads fail when the operations cannot be performed. + +### 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. + +### ConsumeBench +ConsumeBench starts one or more Kafka consumers on a single agent node. Depending on the passed in configuration (see ConsumeBenchSpec), the consumers either subscribe to a set of topics (leveraging consumer group functionality and dynamic partition assignment) or manually assign partitions to themselves. +The workload measures the average produce latency, as well as the median, 95th percentile, and 99th percentile latency. + +Faults +======================================== +Trogdor can run several faults which deliberately break something in the cluster. + +### ProcessStopFault +ProcessStopFault stops a process by sending it a SIGSTOP signal. When the fault ends, the process is resumed with SIGCONT. + +### 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. + +When using exec mode, you must pass in a Task specification to use. The Agent will try to start this task. + +For example: + + > ./bin/trogdor.sh agent -n node0 -c ./config/trogdor.conf --exec ./tests/spec/simple_produce_bench.json + +When using exec mode, the Agent will exit once the task is complete. diff --git a/Vagrantfile b/Vagrantfile index 2179a3c37ecb1..ee08487be66cd 100644 --- a/Vagrantfile +++ b/Vagrantfile @@ -40,8 +40,10 @@ ec2_keypair_file = nil ec2_region = "us-east-1" ec2_az = nil # Uses set by AWS -ec2_ami = "ami-9eaa1cf6" +ec2_ami = "ami-29ebb519" ec2_instance_type = "m3.medium" +ec2_spot_instance = ENV['SPOT_INSTANCE'] ? ENV['SPOT_INSTANCE'] == 'true' : true +ec2_spot_max_price = "0.113" # On-demand price for instance type ec2_user = "ubuntu" ec2_instance_name_prefix = "kafka-vagrant" ec2_security_groups = nil @@ -50,6 +52,9 @@ ec2_subnet_id = nil # are running Vagrant from within that VPC as well. ec2_associate_public_ip = nil +jdk_major = '8' +jdk_full = '8u202-linux-x64' + local_config_file = File.join(File.dirname(__FILE__), "Vagrantfile.local") if File.exists?(local_config_file) then eval(File.read(local_config_file), binding, "Vagrantfile.local") @@ -73,15 +78,6 @@ Vagrant.configure(VAGRANTFILE_API_VERSION) do |config| if Vagrant.has_plugin?("vagrant-cachier") override.cache.scope = :box - # Besides the defaults, we use a custom cache to handle the Oracle JDK - # download, which downloads via wget during an apt install. Because of the - # way the installer ends up using its cache directory, we need to jump - # through some hoops instead of just specifying a cache directly -- we - # share to a temporary location and the provisioning scripts symlink data - # to the right location. - override.cache.enable :generic, { - "oracle-jdk7" => { cache_dir: "/tmp/oracle-jdk7-installer-cache" }, - } end end @@ -133,6 +129,10 @@ Vagrant.configure(VAGRANTFILE_API_VERSION) do |config| else aws.associate_public_ip = ec2_associate_public_ip end + aws.region_config ec2_region do |region| + region.spot_instance = ec2_spot_instance + region.spot_max_price = ec2_spot_max_price + end # Exclude some directories that can grow very large from syncing override.vm.synced_folder ".", "/vagrant", type: "rsync", rsync__exclude: ['.git', 'core/data/', 'logs/', 'tests/results/', 'results/'] @@ -141,7 +141,10 @@ Vagrant.configure(VAGRANTFILE_API_VERSION) do |config| def name_node(node, name, ec2_instance_name_prefix) node.vm.hostname = name node.vm.provider :aws do |aws| - aws.tags = { 'Name' => ec2_instance_name_prefix + "-" + Socket.gethostname + "-" + name } + aws.tags = { + 'Name' => ec2_instance_name_prefix + "-" + Socket.gethostname + "-" + name, + 'JenkinsBuildUrl' => ENV['BUILD_URL'] + } end end @@ -160,7 +163,7 @@ Vagrant.configure(VAGRANTFILE_API_VERSION) do |config| name_node(zookeeper, name, ec2_instance_name_prefix) ip_address = "192.168.50." + (10 + i).to_s assign_local_ip(zookeeper, ip_address) - zookeeper.vm.provision "shell", path: "vagrant/base.sh" + zookeeper.vm.provision "shell", path: "vagrant/base.sh", env: {"JDK_MAJOR" => jdk_major, "JDK_FULL" => jdk_full} zk_jmx_port = enable_jmx ? (8000 + i).to_s : "" zookeeper.vm.provision "shell", path: "vagrant/zk.sh", :args => [i.to_s, num_zookeepers, zk_jmx_port] end @@ -177,7 +180,7 @@ Vagrant.configure(VAGRANTFILE_API_VERSION) do |config| # host DNS isn't setup, we shouldn't use hostnames -- IP addresses must be # used to support clients running on the host. zookeeper_connect = zookeepers.map{ |zk_addr| zk_addr + ":2181"}.join(",") - broker.vm.provision "shell", path: "vagrant/base.sh" + broker.vm.provision "shell", path: "vagrant/base.sh", env: {"JDK_MAJOR" => jdk_major, "JDK_FULL" => jdk_full} kafka_jmx_port = enable_jmx ? (9000 + i).to_s : "" broker.vm.provision "shell", path: "vagrant/broker.sh", :args => [i.to_s, enable_dns ? name : ip_address, zookeeper_connect, kafka_jmx_port] end @@ -189,7 +192,7 @@ Vagrant.configure(VAGRANTFILE_API_VERSION) do |config| name_node(worker, name, ec2_instance_name_prefix) ip_address = "192.168.50." + (100 + i).to_s assign_local_ip(worker, ip_address) - worker.vm.provision "shell", path: "vagrant/base.sh" + worker.vm.provision "shell", path: "vagrant/base.sh", env: {"JDK_MAJOR" => jdk_major, "JDK_FULL" => jdk_full} end } diff --git a/bin/connect-distributed.sh b/bin/connect-distributed.sh index 08fc57b9c51ac..b8088ad923451 100755 --- a/bin/connect-distributed.sh +++ b/bin/connect-distributed.sh @@ -1,4 +1,4 @@ -#!/bin/sh +#!/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. diff --git a/bin/connect-mirror-maker.sh b/bin/connect-mirror-maker.sh new file mode 100755 index 0000000000000..8e2b2e162daac --- /dev/null +++ b/bin/connect-mirror-maker.sh @@ -0,0 +1,45 @@ +#!/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. + +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/connect-standalone.sh b/bin/connect-standalone.sh index 931cc37c9a43b..441069fed3139 100755 --- a/bin/connect-standalone.sh +++ b/bin/connect-standalone.sh @@ -1,4 +1,4 @@ -#!/bin/sh +#!/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. diff --git a/bin/kafka-delegation-tokens.sh b/bin/kafka-delegation-tokens.sh new file mode 100755 index 0000000000000..49cb276ab3182 --- /dev/null +++ b/bin/kafka-delegation-tokens.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.DelegationTokenCommand "$@" diff --git a/bin/kafka-dump-log.sh b/bin/kafka-dump-log.sh new file mode 100755 index 0000000000000..a97ea7d3d9f8c --- /dev/null +++ b/bin/kafka-dump-log.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.tools.DumpLogSegments "$@" diff --git a/bin/kafka-features.sh b/bin/kafka-features.sh new file mode 100755 index 0000000000000..9dd9f16fd1b05 --- /dev/null +++ b/bin/kafka-features.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.FeatureCommand "$@" 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-replay-log-producer.sh b/bin/kafka-replay-log-producer.sh deleted file mode 100755 index bba3241d75f0d..0000000000000 --- a/bin/kafka-replay-log-producer.sh +++ /dev/null @@ -1,17 +0,0 @@ -#!/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.tools.ReplayLogProducer "$@" diff --git a/bin/kafka-run-class.sh b/bin/kafka-run-class.sh index fc89f25d9d255..f485f5b04459b 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 @@ -32,7 +32,7 @@ if [ -z "$INCLUDE_TEST_JARS" ]; then fi # Exclude jars not necessary for running commands. -regex="(-(test|src|scaladoc|javadoc)\.jar|jar.asc)$" +regex="(-(test|test-sources|src|scaladoc|javadoc)\.jar|jar.asc)$" should_include_file() { if [ "$INCLUDE_TEST_JARS" = true ]; then return 0 @@ -48,7 +48,10 @@ should_include_file() { base_dir=$(dirname $0)/.. if [ -z "$SCALA_VERSION" ]; then - SCALA_VERSION=2.11.12 + SCALA_VERSION=2.13.4 + if [[ -f "$base_dir/gradle.properties" ]]; then + SCALA_VERSION=`grep "^scalaVersion=" "$base_dir/gradle.properties" | cut -d= -f 2` + fi fi if [ -z "$SCALA_BINARY_VERSION" ]; then @@ -57,14 +60,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 - if [ -z "$CLASSPATH" ] ; then - CLASSPATH="$dir/*" - else +if [ -z "$UPGRADE_KAFKA_STREAMS_TEST_VERSION" ]; then + for dir in "$base_dir"/core/build/dependant-libs-${SCALA_VERSION}*; + do CLASSPATH="$CLASSPATH:$dir/*" - fi -done + done +fi for file in "$base_dir"/examples/build/libs/kafka-examples*.jar; do @@ -73,28 +74,63 @@ do fi done -for file in "$base_dir"/clients/build/libs/kafka-clients*.jar; +if [ -z "$UPGRADE_KAFKA_STREAMS_TEST_VERSION" ]; then + clients_lib_dir=$(dirname $0)/../clients/build/libs + streams_lib_dir=$(dirname $0)/../streams/build/libs + streams_dependant_clients_lib_dir=$(dirname $0)/../streams/build/dependant-libs-${SCALA_VERSION} +else + clients_lib_dir=/opt/kafka-$UPGRADE_KAFKA_STREAMS_TEST_VERSION/libs + streams_lib_dir=$clients_lib_dir + streams_dependant_clients_lib_dir=$streams_lib_dir +fi + + +for file in "$clients_lib_dir"/kafka-clients*.jar; do if should_include_file "$file"; then CLASSPATH="$CLASSPATH":"$file" fi done -for file in "$base_dir"/streams/build/libs/kafka-streams*.jar; +for file in "$streams_lib_dir"/kafka-streams*.jar; do if should_include_file "$file"; then CLASSPATH="$CLASSPATH":"$file" fi done -for file in "$base_dir"/streams/examples/build/libs/kafka-streams-examples*.jar; -do - if should_include_file "$file"; then - CLASSPATH="$CLASSPATH":"$file" +if [ -z "$UPGRADE_KAFKA_STREAMS_TEST_VERSION" ]; then + for file in "$base_dir"/streams/examples/build/libs/kafka-streams-examples*.jar; + do + if should_include_file "$file"; then + CLASSPATH="$CLASSPATH":"$file" + fi + done +else + VERSION_NO_DOTS=`echo $UPGRADE_KAFKA_STREAMS_TEST_VERSION | sed 's/\.//g'` + SHORT_VERSION_NO_DOTS=${VERSION_NO_DOTS:0:((${#VERSION_NO_DOTS} - 1))} # remove last char, ie, bug-fix number + for file in "$base_dir"/streams/upgrade-system-tests-$SHORT_VERSION_NO_DOTS/build/libs/kafka-streams-upgrade-system-tests*.jar; + do + if should_include_file "$file"; then + 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 "$streams_dependant_clients_lib_dir"/rocksdb*.jar; +do + CLASSPATH="$CLASSPATH":"$file" done -for file in "$base_dir"/streams/build/dependant-libs-${SCALA_VERSION}/rocksdb*.jar; +for file in "$streams_dependant_clients_lib_dir"/*hamcrest*.jar; do CLASSPATH="$CLASSPATH":"$file" done @@ -111,7 +147,7 @@ do CLASSPATH="$CLASSPATH:$dir/*" done -for cc_pkg in "api" "transforms" "runtime" "file" "json" "tools" +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 @@ -141,7 +177,7 @@ done shopt -u nullglob if [ -z "$CLASSPATH" ] ; then - echo "Classpath is empty. Please build the project first e.g. by running './gradlew jar -Pscala_version=$SCALA_VERSION'" + echo "Classpath is empty. Please build the project first e.g. by running './gradlew jar -PscalaVersion=$SCALA_VERSION'" exit 1 fi @@ -216,11 +252,11 @@ if [ -z "$KAFKA_HEAP_OPTS" ]; then fi # JVM performance options +# MaxInlineLevel=15 is the default since JDK 14 and can be removed once older JDKs are no longer supported 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" + KAFKA_JVM_PERFORMANCE_OPTS="-server -XX:+UseG1GC -XX:MaxGCPauseMillis=20 -XX:InitiatingHeapOccupancyPercent=35 -XX:+ExplicitGCInvokesConcurrent -XX:MaxInlineLevel=15 -Djava.awt.headless=true" fi - while [ $# -gt 0 ]; do COMMAND=$1 case $COMMAND in @@ -250,22 +286,34 @@ GC_FILE_SUFFIX='-gc.log' GC_LOG_FILE_NAME='' if [ "x$GC_LOG_ENABLED" = "xtrue" ]; then GC_LOG_FILE_NAME=$DAEMON_NAME$GC_FILE_SUFFIX - # the first segment of the version number, which is '1' for releases before Java 9 + + # The first segment of the version number, which is '1' for releases before Java 9 # it then becomes '9', '10', ... - JAVA_MAJOR_VERSION=$($JAVA -version 2>&1 | sed -E -n 's/.* version "([^.-]*).*"/\1/p') + # Some examples of the first line of `java --version`: + # 8 -> java version "1.8.0_152" + # 9.0.4 -> java version "9.0.4" + # 10 -> java version "10" 2018-03-20 + # 10.0.1 -> java version "10.0.1" 2018-04-17 + # We need to match to the end of the line to prevent sed from printing the characters that do not match + JAVA_MAJOR_VERSION=$("$JAVA" -version 2>&1 | sed -E -n 's/.* version "([0-9]*).*$/\1/p') if [[ "$JAVA_MAJOR_VERSION" -ge "9" ]] ; then - KAFKA_GC_LOG_OPTS="-Xlog:gc*:file=$LOG_DIR/$GC_LOG_FILE_NAME:time,tags:filecount=10,filesize=102400" + KAFKA_GC_LOG_OPTS="-Xlog:gc*:file=$LOG_DIR/$GC_LOG_FILE_NAME:time,tags:filecount=10,filesize=100M" else KAFKA_GC_LOG_OPTS="-Xloggc:$LOG_DIR/$GC_LOG_FILE_NAME -verbose:gc -XX:+PrintGCDetails -XX:+PrintGCDateStamps -XX:+PrintGCTimeStamps -XX:+UseGCLogFileRotation -XX:NumberOfGCLogFiles=10 -XX:GCLogFileSize=100M" fi fi +# Remove a possible colon prefix from the classpath (happens at lines like `CLASSPATH="$CLASSPATH:$file"` when CLASSPATH is blank) +# Syntax used on the right side is native Bash string manipulation; for more details see +# http://tldp.org/LDP/abs/html/string-manipulation.html, specifically the section titled "Substring Removal" +CLASSPATH=${CLASSPATH#:} + # If Cygwin is detected, classpath is converted to Windows format. (( CYGWIN )) && CLASSPATH=$(cygpath --path --mixed "${CLASSPATH}") # Launch mode if [ "x$DAEMON_MODE" = "xtrue" ]; then - nohup $JAVA $KAFKA_HEAP_OPTS $KAFKA_JVM_PERFORMANCE_OPTS $KAFKA_GC_LOG_OPTS $KAFKA_JMX_OPTS $KAFKA_LOG4J_OPTS -cp $CLASSPATH $KAFKA_OPTS "$@" > "$CONSOLE_OUTPUT_FILE" 2>&1 < /dev/null & + nohup "$JAVA" $KAFKA_HEAP_OPTS $KAFKA_JVM_PERFORMANCE_OPTS $KAFKA_GC_LOG_OPTS $KAFKA_JMX_OPTS $KAFKA_LOG4J_OPTS -cp "$CLASSPATH" $KAFKA_OPTS "$@" > "$CONSOLE_OUTPUT_FILE" 2>&1 < /dev/null & else - exec $JAVA $KAFKA_HEAP_OPTS $KAFKA_JVM_PERFORMANCE_OPTS $KAFKA_GC_LOG_OPTS $KAFKA_JMX_OPTS $KAFKA_LOG4J_OPTS -cp $CLASSPATH $KAFKA_OPTS "$@" + exec "$JAVA" $KAFKA_HEAP_OPTS $KAFKA_JVM_PERFORMANCE_OPTS $KAFKA_GC_LOG_OPTS $KAFKA_JMX_OPTS $KAFKA_LOG4J_OPTS -cp "$CLASSPATH" $KAFKA_OPTS "$@" fi diff --git a/bin/kafka-server-stop.sh b/bin/kafka-server-stop.sh index d3c660cab5a77..4c86fddf43eb8 100755 --- a/bin/kafka-server-stop.sh +++ b/bin/kafka-server-stop.sh @@ -1,24 +1,35 @@ -#!/bin/sh +#!/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. -PIDS=$(ps ax | grep -i 'kafka\.Kafka' | grep java | grep -v grep | awk '{print $1}') +SIGNAL=${SIGNAL:-TERM} + +OSNAME=$(uname -s) +if [[ "$OSNAME" == "OS/390" ]]; then + if [ -z $JOBNAME ]; then + JOBNAME="KAFKSTRT" + fi + PIDS=$(ps -A -o pid,jobname,comm | grep -i $JOBNAME | grep java | grep -v grep | awk '{print $1}') +elif [[ "$OSNAME" == "OS400" ]]; then + PIDS=$(ps -af | grep -i 'kafka\.Kafka' | grep java | grep -v grep | awk '{print $2}') +else + PIDS=$(ps ax | grep ' kafka\.Kafka ' | grep java | grep -v grep | awk '{print $1}') +fi if [ -z "$PIDS" ]; then echo "No kafka server to stop" exit 1 -else - kill -s TERM $PIDS +else + kill -s $SIGNAL $PIDS fi - diff --git a/bin/kafka-simple-consumer-shell.sh b/bin/kafka-simple-consumer-shell.sh deleted file mode 100755 index 27e386ad7ee87..0000000000000 --- a/bin/kafka-simple-consumer-shell.sh +++ /dev/null @@ -1,17 +0,0 @@ -#!/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.tools.SimpleConsumerShell "$@" diff --git a/bin/trogdor.sh b/bin/trogdor.sh index b21120932dc20..3324c4ea8ec3e 100755 --- a/bin/trogdor.sh +++ b/bin/trogdor.sh @@ -1,4 +1,4 @@ -#!/usr/bin/env bash +#!/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. diff --git a/bin/windows/kafka-consumer-offset-checker.bat b/bin/windows/kafka-consumer-offset-checker.bat deleted file mode 100644 index 2baa1b84061a0..0000000000000 --- a/bin/windows/kafka-consumer-offset-checker.bat +++ /dev/null @@ -1,17 +0,0 @@ -@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.tools.ConsumerOffsetChecker %* diff --git a/bin/windows/kafka-delegation-tokens.bat b/bin/windows/kafka-delegation-tokens.bat new file mode 100644 index 0000000000000..996537f8c020c --- /dev/null +++ b/bin/windows/kafka-delegation-tokens.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.DelegationTokenCommand %* 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-dump-log.bat b/bin/windows/kafka-dump-log.bat new file mode 100644 index 0000000000000..3a1473dc61bc7 --- /dev/null +++ b/bin/windows/kafka-dump-log.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.tools.DumpLogSegments %* 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-replay-log-producer.bat b/bin/windows/kafka-replay-log-producer.bat deleted file mode 100644 index 7b51302a005d3..0000000000000 --- a/bin/windows/kafka-replay-log-producer.bat +++ /dev/null @@ -1,17 +0,0 @@ -@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.tools.ReplayLogProducer %* diff --git a/bin/windows/kafka-run-class.bat b/bin/windows/kafka-run-class.bat index 1dfd9a5abfede..3490588e37ded 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.11.12 + set SCALA_VERSION=2.13.4 ) IF ["%SCALA_BINARY_VERSION%"] EQU [""] ( @@ -111,7 +111,7 @@ IF ["%JMX_PORT%"] NEQ [""] ( rem Log directory to use IF ["%LOG_DIR%"] EQU [""] ( - set LOG_DIR="%BASE_DIR~%/logs" + set LOG_DIR=%BASE_DIR%/logs ) rem Log4j settings @@ -176,7 +176,7 @@ IF not defined CLASSPATH ( EXIT /B 2 ) -set COMMAND=%JAVA% %KAFKA_HEAP_OPTS% %KAFKA_JVM_PERFORMANCE_OPTS% %KAFKA_JMX_OPTS% %KAFKA_LOG4J_OPTS% -cp %CLASSPATH% %KAFKA_OPTS% %* +set COMMAND=%JAVA% %KAFKA_HEAP_OPTS% %KAFKA_JVM_PERFORMANCE_OPTS% %KAFKA_JMX_OPTS% %KAFKA_LOG4J_OPTS% -cp "%CLASSPATH%" %KAFKA_OPTS% %* rem echo. rem echo %COMMAND% rem echo. diff --git a/bin/windows/kafka-simple-consumer-shell.bat b/bin/windows/kafka-simple-consumer-shell.bat deleted file mode 100644 index 8836128b8b05a..0000000000000 --- a/bin/windows/kafka-simple-consumer-shell.bat +++ /dev/null @@ -1,17 +0,0 @@ -@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.tools.SimpleConsumerShell %* 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/bin/windows/zookeeper-shell.bat b/bin/windows/zookeeper-shell.bat index ea3c3987ccd94..f1c86c430c170 100644 --- a/bin/windows/zookeeper-shell.bat +++ b/bin/windows/zookeeper-shell.bat @@ -15,8 +15,8 @@ rem See the License for the specific language governing permissions and rem limitations under the License. IF [%1] EQU [] ( - echo USAGE: %0 zookeeper_host:port[/path] [args...] + echo USAGE: %0 zookeeper_host:port[/path] [-zk-tls-config-file file] [args...] EXIT /B 1 ) -"%~dp0kafka-run-class.bat" org.apache.zookeeper.ZooKeeperMain -server %* +"%~dp0kafka-run-class.bat" org.apache.zookeeper.ZooKeeperMainWithTlsSupportForKafka -server %* diff --git a/bin/zookeeper-server-stop.sh b/bin/zookeeper-server-stop.sh index f771064cb5508..725c02d00d914 100755 --- a/bin/zookeeper-server-stop.sh +++ b/bin/zookeeper-server-stop.sh @@ -1,24 +1,35 @@ -#!/bin/sh +#!/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. -PIDS=$(ps ax | grep java | grep -i QuorumPeerMain | grep -v grep | awk '{print $1}') +SIGNAL=${SIGNAL:-TERM} + +OSNAME=$(uname -s) +if [[ "$OSNAME" == "OS/390" ]]; then + if [ -z $JOBNAME ]; then + JOBNAME="ZKEESTRT" + fi + PIDS=$(ps -A -o pid,jobname,comm | grep -i $JOBNAME | grep java | grep -v grep | awk '{print $1}') +elif [[ "$OSNAME" == "OS400" ]]; then + PIDS=$(ps -af | grep java | grep -i QuorumPeerMain | grep -v grep | awk '{print $2}') +else + PIDS=$(ps ax | grep java | grep -i QuorumPeerMain | grep -v grep | awk '{print $1}') +fi if [ -z "$PIDS" ]; then echo "No zookeeper server to stop" exit 1 else - kill -s TERM $PIDS + kill -s $SIGNAL $PIDS fi - diff --git a/bin/zookeeper-shell.sh b/bin/zookeeper-shell.sh index 95007faeeabdb..2f1d0f2c61670 100755 --- a/bin/zookeeper-shell.sh +++ b/bin/zookeeper-shell.sh @@ -1,4 +1,4 @@ -#!/bin/sh +#!/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. @@ -16,8 +16,8 @@ if [ $# -lt 1 ]; then - echo "USAGE: $0 zookeeper_host:port[/path] [args...]" + echo "USAGE: $0 zookeeper_host:port[/path] [-zk-tls-config-file file] [args...]" exit 1 fi -exec $(dirname $0)/kafka-run-class.sh org.apache.zookeeper.ZooKeeperMain -server "$@" +exec $(dirname $0)/kafka-run-class.sh org.apache.zookeeper.ZooKeeperMainWithTlsSupportForKafka -server "$@" diff --git a/build.gradle b/build.gradle index 1223150a7b24d..28ad728079b4c 100644 --- a/build.gradle +++ b/build.gradle @@ -15,39 +15,54 @@ import org.ajoberstar.grgit.Grgit +import java.nio.charset.StandardCharsets + buildscript { repositories { mavenCentral() jcenter() + maven { + url "https://plugins.gradle.org/m2/" + } } 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.15.0' - classpath 'org.scoverage:gradle-scoverage:2.1.0' - classpath 'com.github.jengelman.gradle.plugins:shadow:2.0.1' - classpath 'org.owasp:dependency-check-gradle:3.0.1' + 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.snom:spotbugs-gradle-plugin:$versions.spotbugsPlugin" + classpath "org.gradle:test-retry-gradle-plugin:$versions.testRetryPlugin" + } +} + +apply plugin: "com.diffplug.spotless" +spotless { + scala { + target 'streams/**/*.scala' + scalafmt("$versions.scalafmt").configFile('checkstyle/.scalafmt.conf') } } -apply from: "$rootDir/gradle/dependencies.gradle" allprojects { repositories { mavenCentral() } - + apply plugin: 'idea' - apply plugin: "jacoco" apply plugin: 'org.owasp.dependencycheck' apply plugin: 'com.github.ben-manes.versions' dependencyUpdates { revision="release" - resolutionStrategy = { + resolutionStrategy { componentSelection { rules -> rules.all { ComponentSelection selection -> boolean rejected = ['snap', 'alpha', 'beta', 'rc', 'cr', 'm'].any { qualifier -> @@ -59,34 +74,52 @@ allprojects { } } } - configurations { - runtime { + configurations.all { + // zinc is the Scala incremental compiler, it has a configuration for its own dependencies + // that are unrelated to the project dependencies, we should not change them + if (name != "zinc") { resolutionStrategy { - force "com.fasterxml.jackson.core:jackson-annotations:$versions.jackson" + force( + // ensure we have a single version of scala jars in the classpath, we enable inlining + // in the scala compiler for the `core` module so binary compatibility is only + // guaranteed if the exact same version of the scala jars is used for compilation + // and at runtime + libs.scalaLibrary, + libs.scalaReflect, + // ensures we have a single version of jackson-annotations in the classpath even if + // some modules only have a transitive reference to an older version + libs.jacksonAnnotations, + // be explicit about the Netty dependency version instead of relying on the version + // set by ZooKeeper (potentially older and containing CVEs) + libs.nettyHandler, + libs.nettyTransportNativeEpoll + ) } } } } - if (JavaVersion.current().isJava8Compatible()) { - tasks.withType(Javadoc) { - // disable the crazy super-strict doclint tool in Java 8 - // noinspection SpellCheckingInspection - options.addStringOption('Xdoclint:none', '-quiet') - } + tasks.withType(Javadoc) { + // disable the crazy super-strict doclint tool in Java 8 + // noinspection SpellCheckingInspection + options.addStringOption('Xdoclint:none', '-quiet') } } ext { - gradleVersion = "4.2.1" + gradleVersion = versions.gradle + minJavaVersion = "8" buildVersionFileName = "kafka-version.properties" - maxPermSizeArgs = [] - if (!JavaVersion.current().isJava8Compatible()) - maxPermSizeArgs += '-XX:MaxPermSize=512m' + defaultMaxHeapSize = "2g" + defaultJvmArgs = ["-Xss4m", "-XX:+UseParallelGC"] userMaxForks = project.hasProperty('maxParallelForks') ? maxParallelForks.toInteger() : null + userIgnoreFailures = project.hasProperty('ignoreFailures') ? ignoreFailures : false + + userMaxTestRetries = project.hasProperty('maxTestRetries') ? maxTestRetries.toInteger() : 0 + userMaxTestRetryFailures = project.hasProperty('maxTestRetryFailures') ? maxTestRetryFailures.toInteger() : 0 skipSigning = project.hasProperty('skipSigning') && skipSigning.toBoolean() shouldSign = !skipSigning && !version.endsWith("SNAPSHOT") && project.gradle.startParameter.taskNames.any { it.contains("upload") } @@ -99,6 +132,8 @@ ext { userTestLoggingEvents = project.hasProperty("testLoggingEvents") ? Arrays.asList(testLoggingEvents.split(",")) : null + userEnableTestCoverage = project.hasProperty("enableTestCoverage") ? enableTestCoverage : false + generatedDocsDir = new File("${project.rootDir}/docs/generated") commitId = project.hasProperty('commitId') ? commitId : null @@ -106,12 +141,12 @@ ext { apply from: file('wrapper.gradle') -if (new File('.git').exists()) { +if (file('.git').exists()) { apply from: file('gradle/rat.gradle') 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.file('.')) + 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([ @@ -121,16 +156,29 @@ if (new File('.git').exists()) { 'PULL_REQUEST_TEMPLATE.md', 'gradlew', 'gradlew.bat', + 'gradle/wrapper/gradle-wrapper.properties', + 'TROGDOR.md', '**/README.md', '**/id_rsa', '**/id_rsa.pub', 'checkstyle/suppressions.xml', - 'streams/quickstart/java/src/test/resources/projects/basic/goal.txt' + 'streams/quickstart/java/src/test/resources/projects/basic/goal.txt', + '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')) { @@ -140,24 +188,25 @@ subprojects { apply plugin: 'maven' apply plugin: 'signing' apply plugin: 'checkstyle' + apply plugin: "com.github.spotbugs" + apply plugin: 'org.gradle.test-retry' - if (!JavaVersion.current().isJava9Compatible()) - apply plugin: 'findbugs' - - sourceCompatibility = 1.7 - targetCompatibility = 1.7 + sourceCompatibility = minJavaVersion + targetCompatibility = minJavaVersion - compileJava { + tasks.withType(JavaCompile) { options.encoding = 'UTF-8' - options.compilerArgs << "-Xlint:deprecation" - // -Xlint:unchecked is too buggy in Java 7, so we only enable for Java 8 or higher - if (JavaVersion.current().isJava8Compatible()) - options.compilerArgs << "-Xlint:unchecked" + options.compilerArgs << "-Xlint:all" + // temporary exclusions until all the warnings are fixed + options.compilerArgs << "-Xlint:-rawtypes" + options.compilerArgs << "-Xlint:-serial" + options.compilerArgs << "-Xlint:-try" + options.compilerArgs << "-Werror" // --release is the recommended way to select the target release, but it's only supported in Java 9 so we also // set --source and --target via `sourceCompatibility` and `targetCompatibility`. If/when Gradle supports `--release` // natively (https://github.com/gradle/gradle/issues/2510), we should switch to that. if (JavaVersion.current().isJava9Compatible()) - options.compilerArgs << "--release" << "7" + options.compilerArgs << "--release" << minJavaVersion } uploadArchives { @@ -177,11 +226,11 @@ subprojects { pom.project { name 'Apache Kafka' packaging 'jar' - url 'http://kafka.apache.org' + url 'https://kafka.apache.org' 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' } } @@ -192,61 +241,174 @@ subprojects { } } + def shouldUseJUnit5 = ["tools", "raft"].contains(it.project.name) def testLoggingEvents = ["passed", "skipped", "failed"] def testShowStandardStreams = false def testExceptionFormat = 'full' + // Gradle built-in logging only supports sending test output to stdout, which generates a lot + // of noise, especially for passing tests. We really only want output for failed tests. This + // hooks into the output and logs it (so we don't have to buffer it all in memory) and only + // saves the output for failing tests. Directory and filenames are such that you can, e.g., + // create a Jenkins rule to collect failed test output. + def logTestStdout = { + def testId = { TestDescriptor descriptor -> + "${descriptor.className}.${descriptor.name}".toString() + } + + def logFiles = new HashMap() + def logStreams = new HashMap() + beforeTest { TestDescriptor td -> + def tid = testId(td) + // truncate the file name if it's too long + def logFile = new File( + "${projectDir}/build/reports/testOutput/${tid.substring(0, Math.min(tid.size(),240))}.test.stdout" + ) + logFile.parentFile.mkdirs() + logFiles.put(tid, logFile) + logStreams.put(tid, new FileOutputStream(logFile)) + } + onOutput { TestDescriptor td, TestOutputEvent toe -> + def tid = testId(td) + // Some output can happen outside the context of a specific test (e.g. at the class level) + // and beforeTest/afterTest seems to not be invoked for these cases (and similarly, there's + // a TestDescriptor hierarchy that includes the thread executing the test, Gradle tasks, + // etc). We see some of these in practice and it seems like something buggy in the Gradle + // test runner since we see it *before* any tests and it is frequently not related to any + // 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 || 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 { + logStreams.get(tid).write(toe.message.getBytes(StandardCharsets.UTF_8)) + } catch (Exception e) { + println "ERROR: Failed to write output for test ${tid}" + e.printStackTrace() + } + } + afterTest { TestDescriptor td, TestResult tr -> + def tid = testId(td) + try { + logStreams.get(tid).close() + if (tr.resultType != TestResult.ResultType.FAILURE) { + logFiles.get(tid).delete() + } else { + def file = logFiles.get(tid) + println "${tid} failed, log available in ${file}" + } + } catch (Exception e) { + println "ERROR: Failed to close stdout file for ${tid}" + e.printStackTrace() + } finally { + logFiles.remove(tid) + logStreams.remove(tid) + } + } + } test { maxParallelForks = userMaxForks ?: Runtime.runtime.availableProcessors() + ignoreFailures = userIgnoreFailures - minHeapSize = "256m" - maxHeapSize = "2048m" - jvmArgs = maxPermSizeArgs + maxHeapSize = defaultMaxHeapSize + jvmArgs = defaultJvmArgs testLogging { events = userTestLoggingEvents ?: testLoggingEvents showStandardStreams = userShowStandardStreams ?: testShowStandardStreams exceptionFormat = testExceptionFormat } + logTestStdout.rehydrate(delegate, owner, this)() + // 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' + + if (shouldUseJUnit5) + useJUnitPlatform() + + retry { + maxRetries = userMaxTestRetries + maxFailures = userMaxTestRetryFailures + } } task integrationTest(type: Test, dependsOn: compileJava) { maxParallelForks = userMaxForks ?: Runtime.runtime.availableProcessors() + ignoreFailures = userIgnoreFailures + + maxHeapSize = defaultMaxHeapSize + jvmArgs = defaultJvmArgs - minHeapSize = "256m" - maxHeapSize = "2048m" - jvmArgs = maxPermSizeArgs testLogging { events = userTestLoggingEvents ?: testLoggingEvents showStandardStreams = userShowStandardStreams ?: testShowStandardStreams exceptionFormat = testExceptionFormat } + logTestStdout.rehydrate(delegate, owner, this)() - useJUnit { - includeCategories 'org.apache.kafka.test.IntegrationTest' + // 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' + + if (shouldUseJUnit5) { + useJUnitPlatform { + includeTags "integration" + } + } else { + useJUnit { + includeCategories 'org.apache.kafka.test.IntegrationTest' + } } + retry { + maxRetries = userMaxTestRetries + maxFailures = userMaxTestRetryFailures + } } task unitTest(type: Test, dependsOn: compileJava) { maxParallelForks = userMaxForks ?: Runtime.runtime.availableProcessors() + ignoreFailures = userIgnoreFailures - minHeapSize = "256m" - maxHeapSize = "2048m" - jvmArgs = maxPermSizeArgs + maxHeapSize = defaultMaxHeapSize + jvmArgs = defaultJvmArgs testLogging { events = userTestLoggingEvents ?: testLoggingEvents showStandardStreams = userShowStandardStreams ?: testShowStandardStreams exceptionFormat = testExceptionFormat } + logTestStdout.rehydrate(delegate, owner, this)() + + // 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' - useJUnit { - excludeCategories 'org.apache.kafka.test.IntegrationTest' + if (shouldUseJUnit5) { + useJUnitPlatform { + excludeTags "integration" + } + } else { + if (it.project.name != 'generator') { + useJUnit { + excludeCategories 'org.apache.kafka.test.IntegrationTest' + } + } + } + + retry { + maxRetries = userMaxTestRetries + maxFailures = userMaxTestRetryFailures } - } jar { @@ -305,7 +467,12 @@ subprojects { } plugins.withType(ScalaPlugin) { - task scaladocJar(type:Jar) { + + scala { + zincVersion = versions.zinc + } + + task scaladocJar(type:Jar, dependsOn: scaladoc) { classifier = 'scaladoc' from "$rootDir/LICENSE" from "$rootDir/NOTICE" @@ -315,9 +482,6 @@ subprojects { //documentation task should also trigger building scala doc jar docsJar.dependsOn scaladocJar - artifacts { - archives scaladocJar - } } tasks.withType(ScalaCompile) { @@ -330,11 +494,10 @@ subprojects { "-language:postfixOps", "-language:implicitConversions", "-language:existentials", - "-Xlint:by-name-right-associative", + "-Xlint:constant", "-Xlint:delayedinit-select", "-Xlint:doc-detached", "-Xlint:missing-interpolator", - "-Xlint:nullary-override", "-Xlint:nullary-unit", "-Xlint:option-implicit", "-Xlint:package-object-classes", @@ -342,69 +505,139 @@ subprojects { "-Xlint:private-shadow", "-Xlint:stars-align", "-Xlint:type-parameter-shadow", - "-Xlint:unsound-match", + "-Xlint:unused" ] - if (versions.baseScala != '2.11') { + // Inline more aggressively when compiling the `core` jar since it's not meant to be used as a library. + // More specifically, inline classes from the Scala library so that we can inline methods like `Option.exists` + // and avoid lambda allocations. This is only safe if the Scala library version is the same at compile time + // and runtime. We cannot guarantee this for libraries like kafka streams, so only inline classes from the + // Kafka project in that case. + List inlineFrom + if (project.name.equals('core')) + inlineFrom = ["-opt-inline-from:scala.**", "-opt-inline-from:kafka.**", "-opt-inline-from:org.apache.kafka.**"] + else + inlineFrom = ["-opt-inline-from:org.apache.kafka.**"] + + // Somewhat confusingly, `-opt:l:inline` enables all optimizations. `inlineFrom` configures what can be inlined. + // See https://www.lightbend.com/blog/scala-inliner-optimizer for more information about the optimizer. + scalaCompileOptions.additionalParameters += ["-opt:l:inline"] + scalaCompileOptions.additionalParameters += inlineFrom + + if (versions.baseScala != '2.12') { + scalaCompileOptions.additionalParameters += ["-opt-warnings", "-Xlint:strict-unsealed-patmat"] + // Scala 2.13.2 introduces compiler warnings suppression, which is a pre-requisite for -Xfatal-warnings + scalaCompileOptions.additionalParameters += ["-Xfatal-warnings"] + } + + // 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 == '2.12') { scalaCompileOptions.additionalParameters += [ - "-Xlint:constant", - "-Xlint:unused" + "-Xlint:by-name-right-associative", + "-Xlint:nullary-override", + "-Xlint:unsound-match" ] } + // Scalac's `-release` requires Java 9 or higher + if (JavaVersion.current().isJava9Compatible()) + scalaCompileOptions.additionalParameters += ["-release", minJavaVersion] + configure(scalaCompileOptions.forkOptions) { - memoryMaximumSize = '1g' - jvmArgs = ['-Xss2m'] + maxPermSizeArgs + memoryMaximumSize = defaultMaxHeapSize + jvmArgs = defaultJvmArgs } } checkstyle { configFile = new File(rootDir, "checkstyle/checkstyle.xml") - configProperties = [importControlFile: "$rootDir/checkstyle/import-control.xml"] - // version 7.x requires Java 8 - toolVersion = '6.19' + configProperties = checkstyleConfigProperties("import-control.xml") + toolVersion = versions.checkstyle + } + + configure(checkstyleMain) { + group = 'Verification' + description = 'Run checkstyle on all main Java sources' } + + configure(checkstyleTest) { + group = 'Verification' + description = 'Run checkstyle on all test Java sources' + } + test.dependsOn('checkstyleMain', 'checkstyleTest') - if (!JavaVersion.current().isJava9Compatible()) { - findbugs { - toolVersion = "3.0.1" - excludeFilter = file("$rootDir/gradle/findbugs-exclude.xml") - ignoreFailures = false - } - test.dependsOn('findbugsMain') + spotbugs { + toolVersion = versions.spotbugs + excludeFilter = file("$rootDir/gradle/spotbugs-exclude.xml") + ignoreFailures = false + } + test.dependsOn('spotbugsMain') - tasks.withType(FindBugs) { - reports { - xml.enabled(project.hasProperty('xmlFindBugsReport')) - html.enabled(!project.hasProperty('xmlFindBugsReport')) - } + tasks.withType(com.github.spotbugs.snom.SpotBugsTask) { + reports { + // Continue supporting `xmlFindBugsReport` for compatibility + xml.enabled(project.hasProperty('xmlSpotBugsReport') || project.hasProperty('xmlFindBugsReport')) + html.enabled(!project.hasProperty('xmlSpotBugsReport') && !project.hasProperty('xmlFindBugsReport')) } + maxHeapSize = defaultMaxHeapSize + jvmArgs = defaultJvmArgs } // Ignore core since its a scala project if (it.path != ':core') { - // NOTE: Gradles Jacoco plugin does not support "offline instrumentation" this means that classes mocked by PowerMock - // may report 0 coverage, since the source was modified after initial instrumentation. - // See https://github.com/jacoco/jacoco/issues/51 - jacocoTestReport { - dependsOn tasks.test - sourceSets sourceSets.main - reports { - html.enabled = true - xml.enabled = true - csv.enabled = false + if (userEnableTestCoverage) { + apply plugin: "jacoco" + + jacoco { + toolVersion = versions.jacoco + } + + // NOTE: Jacoco Gradle plugin does not support "offline instrumentation" this means that classes mocked by PowerMock + // may report 0 coverage, since the source was modified after initial instrumentation. + // See https://github.com/jacoco/jacoco/issues/51 + jacocoTestReport { + dependsOn tasks.test + sourceSets sourceSets.main + reports { + html.enabled = true + xml.enabled = true + csv.enabled = false + } } + } } - def coverageGen = it.path == ':core' ? 'reportScoverage' : 'jacocoTestReport' - task reportCoverage(dependsOn: [coverageGen]) + if (userEnableTestCoverage) { + def coverageGen = it.path == ':core' ? 'reportScoverage' : 'jacocoTestReport' + task reportCoverage(dependsOn: [coverageGen]) + } + + 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" + } + } } gradle.taskGraph.whenReady { taskGraph -> - taskGraph.getAllTasks().findAll { it.name.contains('findbugsScoverage') || it.name.contains('findbugsTest') }.each { task -> + taskGraph.getAllTasks().findAll { it.name.contains('spotbugsScoverage') || it.name.contains('spotbugsTest') }.each { task -> task.enabled = false } } @@ -425,8 +658,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 */ @@ -440,125 +673,107 @@ def fineTuneEclipseClasspathFile(eclipse, project) { } } -// Aggregates all jacoco results into the root project directory -task jacocoRootReport(type: org.gradle.testing.jacoco.tasks.JacocoReport) { - def javaProjects = subprojects.findAll { it.path != ':core' } - - description = 'Generates an aggregate report from all subprojects' - dependsOn(javaProjects.test) - - additionalSourceDirs = files(javaProjects.sourceSets.main.allSource.srcDirs) - sourceDirectories = files(javaProjects.sourceSets.main.allSource.srcDirs) - classDirectories = files(javaProjects.sourceSets.main.output) - executionData = files(javaProjects.jacocoTestReport.executionData) - - reports { - html.enabled = true - xml.enabled = true - } - - // workaround to ignore projects that don't have any tests at all - onlyIf = { true } - doFirst { - executionData = files(executionData.findAll { it.exists() }) - } +def checkstyleConfigProperties(configFileName) { + [importControlFile: "$rootDir/checkstyle/$configFileName", + suppressionsFile: "$rootDir/checkstyle/suppressions.xml", + headerFile: "$rootDir/checkstyle/java.header"] } -task reportCoverage(dependsOn: ['jacocoRootReport', 'core:reportCoverage']) - -for ( sv in availableScalaVersions ) { - String taskSuffix = sv.replaceAll("\\.", "_") - - tasks.create(name: "jar_core_${taskSuffix}", type: GradleBuild) { - startParameter = project.getGradle().getStartParameter().newInstance() - startParameter.projectProperties += [scalaVersion: "${sv}"] - tasks = ['core:jar'] - } - - tasks.create(name: "test_core_${taskSuffix}", type: GradleBuild) { - startParameter = project.getGradle().getStartParameter().newInstance() - startParameter.projectProperties += [scalaVersion: "${sv}"] - tasks = ['core:test'] - } - - tasks.create(name: "srcJar_${taskSuffix}", type: GradleBuild) { - startParameter = project.getGradle().getStartParameter().newInstance() - startParameter.projectProperties += [scalaVersion: "${sv}"] - tasks = ['core:srcJar'] - } +// Aggregates all jacoco results into the root project directory +if (userEnableTestCoverage) { + task jacocoRootReport(type: org.gradle.testing.jacoco.tasks.JacocoReport) { + def javaProjects = subprojects.findAll { it.path != ':core' } - tasks.create(name: "docsJar_${taskSuffix}", type: GradleBuild) { - startParameter = project.getGradle().getStartParameter().newInstance() - startParameter.projectProperties += [scalaVersion: "${sv}"] - tasks = ['core:docsJar'] - } + description = 'Generates an aggregate report from all subprojects' + dependsOn(javaProjects.test) - tasks.create(name: "install_${taskSuffix}", type: GradleBuild) { - startParameter = project.getGradle().getStartParameter().newInstance() - startParameter.projectProperties += [scalaVersion: "${sv}"] - tasks = ['install'] - } + additionalSourceDirs.from = javaProjects.sourceSets.main.allSource.srcDirs + sourceDirectories.from = javaProjects.sourceSets.main.allSource.srcDirs + classDirectories.from = javaProjects.sourceSets.main.output + executionData.from = javaProjects.jacocoTestReport.executionData - tasks.create(name: "releaseTarGz_${taskSuffix}", type: GradleBuild) { - startParameter = project.getGradle().getStartParameter().newInstance() - startParameter.projectProperties += [scalaVersion: "${sv}"] - tasks = ['releaseTarGz'] - } + reports { + html.enabled = true + xml.enabled = true + } - tasks.create(name: "uploadCoreArchives_${taskSuffix}", type: GradleBuild) { - startParameter = project.getGradle().getStartParameter().newInstance() - startParameter.projectProperties += [scalaVersion: "${sv}"] - tasks = ['core:uploadArchives'] + // workaround to ignore projects that don't have any tests at all + onlyIf = { true } + doFirst { + executionData = files(executionData.findAll { it.exists() }) + } } } -def connectPkgs = ['connect:api', 'connect:runtime', 'connect:transforms', 'connect:json', 'connect:file'] -def pkgs = ['clients', 'examples', 'log4j-appender', 'tools', 'streams', 'streams:examples'] + connectPkgs - -/** Create one task per default Scala version */ -def withDefScalaVersions(taskName) { - defaultScalaVersions.collect { taskName + '_' + it.replaceAll('\\.', '_') } +if (userEnableTestCoverage) { + task reportCoverage(dependsOn: ['jacocoRootReport', 'core:reportCoverage']) } -tasks.create(name: "jarConnect", dependsOn: connectPkgs.collect { it + ":jar" }) {} -tasks.create(name: "jarAll", dependsOn: withDefScalaVersions('jar_core') + pkgs.collect { it + ":jar" }) { } - -tasks.create(name: "srcJarAll", dependsOn: withDefScalaVersions('srcJar') + pkgs.collect { it + ":srcJar" }) { } +def connectPkgs = [ + 'connect:api', + 'connect:basic-auth-extension', + 'connect:file', + 'connect:json', + 'connect:runtime', + 'connect:transforms', + 'connect:mirror', + 'connect:mirror-client' +] -tasks.create(name: "docsJarAll", dependsOn: withDefScalaVersions('docsJar') + pkgs.collect { it + ":docsJar" }) { } +tasks.create(name: "jarConnect", dependsOn: connectPkgs.collect { it + ":jar" }) {} tasks.create(name: "testConnect", dependsOn: connectPkgs.collect { it + ":test" }) {} -tasks.create(name: "testAll", dependsOn: withDefScalaVersions('test_core') + pkgs.collect { it + ":test" }) { } - -tasks.create(name: "installAll", dependsOn: withDefScalaVersions('install') + pkgs.collect { it + ":install" }) { } - -tasks.create(name: "releaseTarGzAll", dependsOn: withDefScalaVersions('releaseTarGz')) { } - -tasks.create(name: "uploadArchivesAll", dependsOn: withDefScalaVersions('uploadCoreArchives') + pkgs.collect { it + ":uploadArchives" }) { } project(':core') { println "Building project 'core' with Scala version ${versions.scala}" apply plugin: 'scala' - apply plugin: "org.scoverage" + + // scaladoc generation is configured at the sub-module level with an artifacts + // block (cf. see streams-scala). If scaladoc generation is invoked explicitly + // for the `core` module, this ensures the generated jar doesn't include scaladoc + // files since the `core` module doesn't include public APIs. + scaladoc { + enabled = false + } + if (userEnableTestCoverage) + apply plugin: "org.scoverage" archivesBaseName = "kafka_${versions.baseScala}" dependencies { compile project(':clients') + compile project(':raft') compile libs.jacksonDatabind + compile libs.jacksonModuleScala + compile libs.jacksonDataformatCsv + compile libs.jacksonJDK8Datatypes compile libs.joptSimple compile libs.metrics - compile libs.scala - compile libs.slf4jlog4j + 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.zkclient - compile libs.zookeeper + compile libs.slf4jApi + compile(libs.zookeeper) { + exclude module: 'slf4j-log4j12' + exclude module: 'log4j' + } + // ZooKeeperMain depends on commons-cli but declares the dependency as `provided` + compile libs.commonsCli + + compileOnly libs.log4j testCompile project(':clients').sourceSets.test.output testCompile libs.bcpkix + testCompile libs.mockitoCore testCompile libs.easymock testCompile(libs.apacheda) { exclude group: 'xml-apis', module: 'xml-apis' + // `mina-core` is a transitive dependency for `apacheds` and `apacheda`. + // It is safer to use from `apacheds` since that is the implementation. + exclude module: 'mina-core' } testCompile libs.apachedsCoreApi testCompile libs.apachedsInterceptorKerberos @@ -568,23 +783,21 @@ project(':core') { testCompile libs.apachedsLdifPartition testCompile libs.apachedsMavibotPartition testCompile libs.apachedsJdbmPartition - testCompile libs.junit + testCompile libs.junitJupiterApi + testCompile libs.junitVintageEngine testCompile libs.scalatest + testCompile libs.slf4jlog4j testCompile libs.jfreechart - - scoverage libs.scoveragePlugin - scoverage libs.scoverageRuntime } - jacocoTestReport.enabled = false - scoverage { - reportDir = file("${rootProject.buildDir}/scoverage") - highlighting = false - } - checkScoverage { - minimumRate = 0.0 + if (userEnableTestCoverage) { + scoverage { + scoverageVersion = versions.scoverage + reportDir = file("${rootProject.buildDir}/scoverage") + highlighting = false + minimumRate = 0.0 + } } - checkScoverage.shouldRunAfter('test') configurations { // manually excludes some unnecessary dependencies @@ -594,7 +807,6 @@ project(':core') { compile.exclude module: 'jmxri' compile.exclude module: 'jmxtools' compile.exclude module: 'mail' - compile.exclude module: 'netty' // To prevent a UniqueResourceException due the same resource existing in both // org.apache.directory.api/api-all and org.apache.directory.api/api-ldap-schema-data testCompile.exclude module: 'api-ldap-schema-data' @@ -603,6 +815,7 @@ project(':core') { tasks.create(name: "copyDependantLibs", type: Copy) { from (configurations.testRuntime) { include('slf4j-log4j12*') + include('log4j*jar') } from (configurations.runtime) { exclude('kafka-clients*') @@ -611,6 +824,20 @@ project(':core') { duplicatesStrategy 'exclude' } + task processMessages(type:JavaExec) { + main = "org.apache.kafka.message.MessageGenerator" + classpath = project(':generator').sourceSets.main.runtimeClasspath + args = [ "-p", "kafka.internals.generated", + "-o", "src/generated/java/kafka/internals/generated", + "-i", "src/main/resources/common/message", + "-m", "MessageDataGenerator" + ] + inputs.dir("src/main/resources/common/message") + outputs.dir("src/generated/java/kafka/internals/generated") + } + + compileJava.dependsOn 'processMessages' + task genProtocolErrorDocs(type: JavaExec) { classpath = sourceSets.main.runtimeClasspath main = 'org.apache.kafka.common.protocol.Errors' @@ -618,6 +845,13 @@ project(':core') { standardOutput = new File(generatedDocsDir, "protocol_errors.html").newOutputStream() } + task genProtocolTypesDocs(type: JavaExec) { + classpath = sourceSets.main.runtimeClasspath + main = 'org.apache.kafka.common.protocol.types.Type' + if( !generatedDocsDir.exists() ) { generatedDocsDir.mkdirs() } + standardOutput = new File(generatedDocsDir, "protocol_types.html").newOutputStream() + } + task genProtocolApiKeyDocs(type: JavaExec) { classpath = sourceSets.main.runtimeClasspath main = 'org.apache.kafka.common.protocol.ApiKeys' @@ -681,10 +915,12 @@ project(':core') { standardOutput = new File(generatedDocsDir, "producer_metrics.html").newOutputStream() } - task siteDocsTar(dependsOn: ['genProtocolErrorDocs', 'genProtocolApiKeyDocs', 'genProtocolMessageDocs', + task siteDocsTar(dependsOn: ['genProtocolErrorDocs', 'genProtocolTypesDocs', 'genProtocolApiKeyDocs', 'genProtocolMessageDocs', 'genAdminClientConfigDocs', 'genProducerConfigDocs', 'genConsumerConfigDocs', 'genKafkaConfigDocs', 'genTopicConfigDocs', ':connect:runtime:genConnectConfigDocs', ':connect:runtime:genConnectTransformationDocs', + ':connect:runtime:genConnectPredicateDocs', + ':connect:runtime:genSinkConnectorConfigDocs', ':connect:runtime:genSourceConnectorConfigDocs', ':streams:genStreamsConfigDocs', 'genConsumerMetricsDocs', 'genProducerMetricsDocs', ':connect:runtime:genConnectMetricsDocs'], type: Tar) { classifier = 'site-docs' @@ -716,8 +952,18 @@ project(':core') { from(project(':connect:json').configurations.runtime) { into("libs/") } from(project(':connect:file').jar) { into("libs/") } 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/") } + from(project(':streams:streams-scala').configurations.runtime) { into("libs/") } + from(project(':streams:test-utils').jar) { into("libs/") } + from(project(':streams:test-utils').configurations.runtime) { into("libs/") } from(project(':streams:examples').jar) { into("libs/") } from(project(':streams:examples').configurations.runtime) { into("libs/") } duplicatesStrategy 'exclude' @@ -738,13 +984,29 @@ project(':core') { include('*.jar') } into "$buildDir/dependant-testlibs" + //By default gradle does not handle test dependencies between the sub-projects + //This line is to include clients project test jar to dependant-testlibs + from (project(':clients').testJar ) { "$buildDir/dependant-testlibs" } duplicatesStrategy 'exclude' } systemTestLibs.dependsOn('jar', 'testJar', 'copyDependantTestLibs') checkstyle { - configProperties = [importControlFile: "$rootDir/checkstyle/import-control-core.xml"] + configProperties = checkstyleConfigProperties("import-control-core.xml") + } + + sourceSets { + main { + java { + srcDirs = ["src/generated/java", "src/main/java"] + } + } + test { + java { + srcDirs = ["src/generated/java", "src/test/java"] + } + } } } @@ -760,44 +1022,59 @@ project(':examples') { } checkstyle { - configProperties = [importControlFile: "$rootDir/checkstyle/import-control-core.xml"] + configProperties = checkstyleConfigProperties("import-control-core.xml") + } +} + +project(':generator') { + dependencies { + compile libs.argparse4j + compile libs.jacksonDatabind + compile libs.jacksonJDK8Datatypes + compile libs.jacksonJaxrsJsonProvider + testCompile libs.junitJupiterApi + testCompile libs.junitVintageEngine + } + + integrationTest { + enabled = false + } + + javadoc { + enabled = false } } project(':clients') { archivesBaseName = "kafka-clients" + configurations { + jacksonDatabindConfig + } + + // add jacksonDatabindConfig as provided scope config with high priority (1000) + conf2ScopeMappings.addMapping(1000, configurations.jacksonDatabindConfig, "provided") + dependencies { + compile libs.zstd compile libs.lz4 compile libs.snappy compile libs.slf4jApi + compileOnly libs.jacksonDatabind // for SASL/OAUTHBEARER bearer token parsing + compileOnly libs.jacksonJDK8Datatypes + + jacksonDatabindConfig libs.jacksonDatabind // to publish as provided scope dependency. + testCompile libs.bcpkix - testCompile libs.junit - testCompile libs.easymock - testCompile libs.powermockJunit4 - testCompile libs.powermockEasymock + testCompile libs.junitJupiterApi + testCompile libs.junitVintageEngine + testCompile libs.mockitoCore testRuntime libs.slf4jlog4j - } - - 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" - } + testRuntime libs.jacksonDatabind + testRuntime libs.jacksonJDK8Datatypes + testCompile libs.jacksonJaxrsJsonProvider } task createVersionFile(dependsOn: determineCommitId) { @@ -827,6 +1104,48 @@ project(':clients') { delete "$buildDir/kafka/" } + task processMessages(type:JavaExec) { + main = "org.apache.kafka.message.MessageGenerator" + classpath = project(':generator').sourceSets.main.runtimeClasspath + args = [ "-p", "org.apache.kafka.common.message", + "-o", "src/generated/java/org/apache/kafka/common/message", + "-i", "src/main/resources/common/message", + "-t", "ApiMessageTypeGenerator", + "-m", "MessageDataGenerator", "JsonConverterGenerator" + ] + 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 = [ "-p", "org.apache.kafka.common.message", + "-o", "src/generated-test/java/org/apache/kafka/common/message", + "-i", "src/test/resources/common/message", + "-m", "MessageDataGenerator", "JsonConverterGenerator" + ] + inputs.dir("src/test/resources/common/message") + outputs.dir("src/generated-test/java/org/apache/kafka/common/message") + } + + sourceSets { + main { + java { + srcDirs = ["src/generated/java", "src/main/java"] + } + } + 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/*" @@ -836,11 +1155,94 @@ project(':clients') { include "**/org/apache/kafka/common/annotation/*" include "**/org/apache/kafka/common/errors/*" include "**/org/apache/kafka/common/header/*" + include "**/org/apache/kafka/common/metrics/*" + include "**/org/apache/kafka/common/metrics/stats/*" include "**/org/apache/kafka/common/resource/*" include "**/org/apache/kafka/common/serialization/*" include "**/org/apache/kafka/common/config/*" + include "**/org/apache/kafka/common/config/provider/*" include "**/org/apache/kafka/common/security/auth/*" + include "**/org/apache/kafka/common/security/plain/*" + 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/*" + } +} + +project(':raft') { + archivesBaseName = "kafka-raft" + + dependencies { + compile project(':clients') + compile libs.slf4jApi + compile libs.jacksonDatabind + + testCompile project(':clients') + testCompile project(':clients').sourceSets.test.output + testCompile libs.junitJupiter + testCompile libs.mockitoCore + + testRuntime libs.slf4jlog4j + } + + task createVersionFile(dependsOn: determineCommitId) { + ext.receiptFile = file("$buildDir/kafka/$buildVersionFileName") + 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") + } + } + + task processMessages(type:JavaExec) { + main = "org.apache.kafka.message.MessageGenerator" + classpath = project(':generator').sourceSets.main.runtimeClasspath + args = [ "-p", "org.apache.kafka.raft.generated", + "-o", "src/generated/java/org/apache/kafka/raft/generated", + "-i", "src/main/resources/common/message", + "-m", "MessageDataGenerator", "JsonConverterGenerator"] + inputs.dir("src/main/resources/common/message") + outputs.dir("src/generated/java/org/apache/kafka/raft/generated") + } + + sourceSets { + main { + java { + srcDirs = ["src/generated/java", "src/main/java"] + } + } + test { + java { + srcDirs = ["src/generated/java", "src/test/java"] + } + } + } + + compileJava.dependsOn 'processMessages' + + jar { + dependsOn createVersionFile + from("$buildDir") { + include "kafka/$buildVersionFileName" + } + } + + clean.doFirst { + delete "$buildDir/kafka/" + } + + javadoc { + enabled = false } } @@ -852,20 +1254,24 @@ project(':tools') { compile project(':log4j-appender') compile libs.argparse4j compile libs.jacksonDatabind - compile libs.slf4jlog4j + compile libs.jacksonJDK8Datatypes + compile libs.slf4jApi compile libs.jacksonJaxrsJsonProvider compile libs.jerseyContainerServlet + compile libs.jerseyHk2 + compile libs.jaxbApi // Jersey dependency that was available in the JDK before Java 9 + compile libs.activation // Jersey dependency that was available in the JDK before Java 9 compile libs.jettyServer compile libs.jettyServlet compile libs.jettyServlets testCompile project(':clients') - testCompile libs.junit + testCompile libs.junitJupiter testCompile project(':clients').sourceSets.test.output - testCompile libs.easymock - testCompile libs.powermockJunit4 - testCompile libs.powermockEasymock + testCompile libs.mockitoInline // supports mocking static methods, final classes, etc. + + testRuntime libs.slf4jlog4j } javadoc { @@ -875,6 +1281,7 @@ project(':tools') { tasks.create(name: "copyDependantLibs", type: Copy) { from (configurations.testRuntime) { include('slf4j-log4j12*') + include('log4j*jar') } from (configurations.runtime) { exclude('kafka-clients*') @@ -890,31 +1297,75 @@ project(':tools') { project(':streams') { archivesBaseName = "kafka-streams" + ext.buildStreamsVersionFileName = "kafka-streams-version.properties" dependencies { compile project(':clients') - compile project(':connect:json') // this dependency should be removed after we unify data API + + // this dependency should be removed after we unify data API + compile(project(':connect:json')) { + // this transitive dependency is not used in Streams, and it breaks SBT builds + exclude module: 'javax.ws.rs-api' + } + compile libs.slf4jApi compile libs.rocksDBJni + // testCompileOnly prevents streams from exporting a dependency on test-utils, which would cause a dependency cycle + testCompileOnly project(':streams:test-utils') testCompile project(':clients').sourceSets.test.output testCompile project(':core') testCompile project(':core').sourceSets.test.output - testCompile libs.junit + testCompile libs.log4j + testCompile libs.junitJupiterApi + testCompile libs.junitVintageEngine testCompile libs.easymock + testCompile libs.powermockJunit4 + testCompile libs.powermockEasymock testCompile libs.bcpkix + testCompile libs.hamcrest + testRuntimeOnly project(':streams:test-utils') testRuntime libs.slf4jlog4j } - javadoc { - include "**/org/apache/kafka/streams/**" - exclude "**/internals/**" + task processMessages(type:JavaExec) { + main = "org.apache.kafka.message.MessageGenerator" + classpath = project(':generator').sourceSets.main.runtimeClasspath + args = [ "-p", "org.apache.kafka.streams.internals.generated", + "-o", "src/generated/java/org/apache/kafka/streams/internals/generated", + "-i", "src/main/resources/common/message", + "-m", "MessageDataGenerator" + ] + inputs.dir("src/main/resources/common/message") + outputs.dir("src/generated/java/org/apache/kafka/streams/internals/generated") } - tasks.create(name: "copyDependantLibs", type: Copy) { + sourceSets { + main { + java { + srcDirs = ["src/generated/java", "src/main/java"] + } + } + test { + java { + srcDirs = ["src/generated/java", "src/test/java"] + } + } + } + + compileJava.dependsOn 'processMessages' + + javadoc { + include "**/org/apache/kafka/streams/**" + exclude "**/internals/**" + } + + tasks.create(name: "copyDependantLibs", type: Copy) { from (configurations.testRuntime) { include('slf4j-log4j12*') + include('log4j*jar') + include('*hamcrest*') } from (configurations.runtime) { exclude('kafka-clients*') @@ -923,7 +1374,27 @@ project(':streams') { duplicatesStrategy 'exclude' } + 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' } @@ -937,6 +1408,111 @@ project(':streams') { if( !generatedDocsDir.exists() ) { generatedDocsDir.mkdirs() } standardOutput = new File(generatedDocsDir, "streams_config.html").newOutputStream() } + + task testAll( + dependsOn: [ + ':streams:test', + ':streams:test-utils:test', + ':streams:streams-scala:test', + ':streams:upgrade-system-tests-0100:test', + ':streams:upgrade-system-tests-0101:test', + ':streams:upgrade-system-tests-0102:test', + ':streams:upgrade-system-tests-0110:test', + ':streams:upgrade-system-tests-10:test', + ':streams:upgrade-system-tests-11:test', + ':streams:upgrade-system-tests-20:test', + ':streams:upgrade-system-tests-21:test', + ':streams:upgrade-system-tests-22:test', + ':streams:upgrade-system-tests-23:test', + ':streams:upgrade-system-tests-24:test', + ':streams:upgrade-system-tests-25:test', + ':streams:upgrade-system-tests-26:test', + ':streams:examples:test' + ] + ) +} + +project(':streams:streams-scala') { + println "Building project 'streams-scala' with Scala version ${versions.scala}" + apply plugin: 'scala' + archivesBaseName = "kafka-streams-scala_${versions.baseScala}" + dependencies { + compile project(':streams') + + compile libs.scalaLibrary + compile libs.scalaCollectionCompat + + testCompile project(':core') + testCompile project(':core').sourceSets.test.output + testCompile project(':streams').sourceSets.test.output + testCompile project(':clients').sourceSets.test.output + testCompile project(':streams:test-utils') + + testCompile libs.junitJupiterApi + testCompile libs.junitVintageEngine + testCompile libs.scalatest + testCompile libs.easymock + testCompile libs.hamcrest + + testRuntime libs.slf4jlog4j + } + + javadoc { + include "**/org/apache/kafka/streams/scala/**" + } + + tasks.create(name: "copyDependantLibs", type: Copy) { + from (configurations.runtime) { + exclude('kafka-streams*') + } + into "$buildDir/dependant-libs-${versions.scala}" + duplicatesStrategy 'exclude' + } + + jar { + dependsOn 'copyDependantLibs' + } + + artifacts { + archives scaladocJar + } + + test.dependsOn(':spotlessScalaCheck') +} + +project(':streams:test-utils') { + archivesBaseName = "kafka-streams-test-utils" + + dependencies { + compile project(':streams') + compile project(':clients') + + testCompile project(':clients').sourceSets.test.output + testCompile libs.junitJupiterApi + testCompile libs.junitVintageEngine + testCompile libs.easymock + testCompile libs.hamcrest + + testRuntime libs.slf4jlog4j + } + + javadoc { + include "**/org/apache/kafka/streams/test/**" + exclude "**/internals/**" + } + + tasks.create(name: "copyDependantLibs", type: Copy) { + from (configurations.runtime) { + exclude('kafka-streams*') + } + into "$buildDir/dependant-libs-${versions.scala}" + duplicatesStrategy 'exclude' + } + + jar { + dependsOn 'copyDependantLibs' + } + } project(':streams:examples') { @@ -945,7 +1521,12 @@ project(':streams:examples') { dependencies { compile project(':streams') compile project(':connect:json') // this dependency should be removed after we unify data API - compile libs.slf4jlog4j // this dependency should be removed after KIP-4 + compile libs.slf4jlog4j + + testCompile project(':streams:test-utils') + testCompile project(':clients').sourceSets.test.output // for org.apache.kafka.test.IntegrationTest + testCompile libs.junitJupiterApi + testCompile libs.junitVintageEngine } javadoc { @@ -965,6 +1546,162 @@ project(':streams:examples') { } } +project(':streams:upgrade-system-tests-0100') { + archivesBaseName = "kafka-streams-upgrade-system-tests-0100" + + dependencies { + testCompile libs.kafkaStreams_0100 + } + + systemTestLibs { + dependsOn testJar + } +} + +project(':streams:upgrade-system-tests-0101') { + archivesBaseName = "kafka-streams-upgrade-system-tests-0101" + + dependencies { + testCompile libs.kafkaStreams_0101 + } + + systemTestLibs { + dependsOn testJar + } +} + +project(':streams:upgrade-system-tests-0102') { + archivesBaseName = "kafka-streams-upgrade-system-tests-0102" + + dependencies { + testCompile libs.kafkaStreams_0102 + } + + systemTestLibs { + dependsOn testJar + } +} + +project(':streams:upgrade-system-tests-0110') { + archivesBaseName = "kafka-streams-upgrade-system-tests-0110" + + dependencies { + testCompile libs.kafkaStreams_0110 + } + + systemTestLibs { + dependsOn testJar + } +} + +project(':streams:upgrade-system-tests-10') { + archivesBaseName = "kafka-streams-upgrade-system-tests-10" + + dependencies { + testCompile libs.kafkaStreams_10 + } + + systemTestLibs { + dependsOn testJar + } +} + +project(':streams:upgrade-system-tests-11') { + archivesBaseName = "kafka-streams-upgrade-system-tests-11" + + dependencies { + testCompile libs.kafkaStreams_11 + } + + systemTestLibs { + dependsOn testJar + } +} + +project(':streams:upgrade-system-tests-20') { + archivesBaseName = "kafka-streams-upgrade-system-tests-20" + + dependencies { + testCompile libs.kafkaStreams_20 + } + + systemTestLibs { + dependsOn testJar + } +} + +project(':streams:upgrade-system-tests-21') { + archivesBaseName = "kafka-streams-upgrade-system-tests-21" + + dependencies { + testCompile libs.kafkaStreams_21 + } + + systemTestLibs { + dependsOn testJar + } +} + +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(':streams:upgrade-system-tests-24') { + archivesBaseName = "kafka-streams-upgrade-system-tests-24" + + dependencies { + testCompile libs.kafkaStreams_24 + } + + systemTestLibs { + dependsOn testJar + } +} + +project(':streams:upgrade-system-tests-25') { + archivesBaseName = "kafka-streams-upgrade-system-tests-25" + + dependencies { + testCompile libs.kafkaStreams_25 + } + + systemTestLibs { + dependsOn testJar + } +} + +project(':streams:upgrade-system-tests-26') { + archivesBaseName = "kafka-streams-upgrade-system-tests-26" + + dependencies { + testCompile libs.kafkaStreams_26 + } + + systemTestLibs { + dependsOn testJar + } +} + project(':jmh-benchmarks') { apply plugin: 'com.github.johnrengelman.shadow' @@ -976,11 +1713,25 @@ project(':jmh-benchmarks') { } dependencies { + compile(project(':core')) { + // jmh requires jopt 4.x while `core` depends on 5.0, they are not binary compatible + exclude group: 'net.sf.jopt-simple', module: 'jopt-simple' + } compile project(':clients') compile project(':streams') + compile project(':core') + compile project(':clients').sourceSets.test.output + compile project(':core').sourceSets.test.output compile libs.jmhCore - compile libs.jmhGeneratorAnnProcess + annotationProcessor libs.jmhGeneratorAnnProcess compile libs.jmhCoreBenchmarks + compile libs.mockitoCore + compile libs.slf4jlog4j + } + + tasks.withType(JavaCompile) { + // Suppress warning caused by code generated by jmh: `warning: [cast] redundant cast to long` + options.compilerArgs << "-Xlint:-cast" } jar { @@ -989,6 +1740,9 @@ project(':jmh-benchmarks') { } } + checkstyle { + configProperties = checkstyleConfigProperties("import-control-jmh-benchmarks.xml") + } task jmh(type: JavaExec, dependsOn: [':jmh-benchmarks:clean', ':jmh-benchmarks:shadowJar']) { @@ -996,7 +1750,7 @@ project(':jmh-benchmarks') { doFirst { if (System.getProperty("jmhArgs")) { - args System.getProperty("jmhArgs").split(',') + args System.getProperty("jmhArgs").split(' ') } args = [shadowJar.archivePath, *args] } @@ -1015,7 +1769,9 @@ project(':log4j-appender') { compile libs.slf4jlog4j testCompile project(':clients').sourceSets.test.output - testCompile libs.junit + testCompile libs.junitJupiterApi + testCompile libs.junitVintageEngine + testCompile libs.easymock } javadoc { @@ -1030,21 +1786,28 @@ project(':connect:api') { dependencies { compile project(':clients') compile libs.slf4jApi + compile libs.jaxrsApi - testCompile libs.junit + testCompile libs.junitJupiterApi + testCompile libs.junitVintageEngine testRuntime libs.slf4jlog4j testCompile project(':clients').sourceSets.test.output } javadoc { - include "**/org/apache/kafka/connect/**" // needed for the `javadocAll` task - options.links "http://docs.oracle.com/javase/7/docs/api/" + include "**/org/apache/kafka/connect/**" // needed for the `aggregatedJavadoc` task + // 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) { from (configurations.testRuntime) { include('slf4j-log4j12*') + include('log4j*jar') } from (configurations.runtime) { exclude('kafka-clients*') @@ -1067,7 +1830,8 @@ project(':connect:transforms') { compile libs.slf4jApi testCompile libs.easymock - testCompile libs.junit + testCompile libs.junitJupiterApi + testCompile libs.junitVintageEngine testCompile libs.powermockJunit4 testCompile libs.powermockEasymock @@ -1082,6 +1846,7 @@ project(':connect:transforms') { tasks.create(name: "copyDependantLibs", type: Copy) { from (configurations.testRuntime) { include('slf4j-log4j12*') + include('log4j*jar') } from (configurations.runtime) { exclude('kafka-clients*') @@ -1102,10 +1867,12 @@ project(':connect:json') { dependencies { compile project(':connect:api') compile libs.jacksonDatabind + compile libs.jacksonJDK8Datatypes compile libs.slf4jApi testCompile libs.easymock - testCompile libs.junit + testCompile libs.junitJupiterApi + testCompile libs.junitVintageEngine testCompile libs.powermockJunit4 testCompile libs.powermockEasymock @@ -1120,6 +1887,7 @@ project(':connect:json') { tasks.create(name: "copyDependantLibs", type: Copy) { from (configurations.testRuntime) { include('slf4j-log4j12*') + include('log4j*jar') } from (configurations.runtime) { exclude('kafka-clients*') @@ -1138,28 +1906,38 @@ project(':connect:runtime') { archivesBaseName = "connect-runtime" dependencies { + compile project(':connect:api') - compile project(":connect:transforms") compile project(':clients') compile project(':tools') - compile libs.slf4jApi + compile project(':connect:json') + compile project(':connect:transforms') + compile libs.slf4jApi compile libs.jacksonJaxrsJsonProvider compile libs.jerseyContainerServlet + compile libs.jerseyHk2 + compile libs.jaxbApi // Jersey dependency that was available in the JDK before Java 9 + compile libs.activation // Jersey dependency that was available in the JDK before Java 9 compile libs.jettyServer compile libs.jettyServlet compile libs.jettyServlets + compile libs.jettyClient compile(libs.reflections) compile(libs.mavenArtifact) testCompile project(':clients').sourceSets.test.output testCompile libs.easymock - testCompile libs.junit + testCompile libs.junitJupiterApi + testCompile libs.junitVintageEngine testCompile libs.powermockJunit4 testCompile libs.powermockEasymock + testCompile libs.mockitoCore + testCompile libs.httpclient - testCompile project(":connect:json") testCompile project(':clients').sourceSets.test.output + testCompile project(':core') + testCompile project(':core').sourceSets.test.output testRuntime libs.slf4jlog4j } @@ -1171,6 +1949,7 @@ project(':connect:runtime') { tasks.create(name: "copyDependantLibs", type: Copy) { from (configurations.testRuntime) { include('slf4j-log4j12*') + include('log4j*jar') } from (configurations.runtime) { exclude('kafka-clients*') @@ -1191,6 +1970,20 @@ project(':connect:runtime') { standardOutput = new File(generatedDocsDir, "connect_config.html").newOutputStream() } + task genSinkConnectorConfigDocs(type: JavaExec) { + classpath = sourceSets.main.runtimeClasspath + main = 'org.apache.kafka.connect.runtime.SinkConnectorConfig' + if( !generatedDocsDir.exists() ) { generatedDocsDir.mkdirs() } + standardOutput = new File(generatedDocsDir, "sink_connector_config.html").newOutputStream() + } + + task genSourceConnectorConfigDocs(type: JavaExec) { + classpath = sourceSets.main.runtimeClasspath + main = 'org.apache.kafka.connect.runtime.SourceConnectorConfig' + if( !generatedDocsDir.exists() ) { generatedDocsDir.mkdirs() } + standardOutput = new File(generatedDocsDir, "source_connector_config.html").newOutputStream() + } + task genConnectTransformationDocs(type: JavaExec) { classpath = sourceSets.main.runtimeClasspath main = 'org.apache.kafka.connect.tools.TransformationDoc' @@ -1198,6 +1991,13 @@ project(':connect:runtime') { standardOutput = new File(generatedDocsDir, "connect_transforms.html").newOutputStream() } + task genConnectPredicateDocs(type: JavaExec) { + classpath = sourceSets.main.runtimeClasspath + main = 'org.apache.kafka.connect.tools.PredicateDoc' + if( !generatedDocsDir.exists() ) { generatedDocsDir.mkdirs() } + standardOutput = new File(generatedDocsDir, "connect_predicates.html").newOutputStream() + } + task genConnectMetricsDocs(type: JavaExec) { classpath = sourceSets.test.runtimeClasspath main = 'org.apache.kafka.connect.runtime.ConnectMetrics' @@ -1215,12 +2015,99 @@ project(':connect:file') { compile libs.slf4jApi testCompile libs.easymock - testCompile libs.junit + testCompile libs.junitJupiterApi + testCompile libs.junitVintageEngine + testCompile libs.powermockJunit4 + testCompile libs.powermockEasymock + + testRuntime libs.slf4jlog4j + testCompile project(':clients').sourceSets.test.output + } + + 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:basic-auth-extension') { + archivesBaseName = "connect-basic-auth-extension" + + dependencies { + compile project(':connect:api') + compile libs.slf4jApi + + testCompile libs.bcpkix + testCompile libs.easymock + testCompile libs.junitJupiterApi + testCompile libs.junitVintageEngine testCompile libs.powermockJunit4 testCompile libs.powermockEasymock + testCompile project(':clients').sourceSets.test.output testRuntime libs.slf4jlog4j + testRuntime libs.jerseyContainerServlet + } + + 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') { + 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.junitJupiterApi + testCompile libs.junitVintageEngine + testCompile libs.mockitoCore 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 { @@ -1230,6 +2117,43 @@ project(':connect:file') { 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.junitJupiterApi + testCompile libs.junitVintageEngine + 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*') @@ -1250,5 +2174,9 @@ task aggregatedJavadoc(type: Javadoc) { 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/.scalafmt.conf b/checkstyle/.scalafmt.conf new file mode 100644 index 0000000000000..057e3b930962e --- /dev/null +++ b/checkstyle/.scalafmt.conf @@ -0,0 +1,20 @@ +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +docstrings = JavaDoc +maxColumn = 120 +continuationIndent.defnSite = 2 +assumeStandardLibraryStripMargin = true +danglingParentheses = true +rewrite.rules = [SortImports, RedundantBraces, RedundantParens, SortModifiers] \ No newline at end of file diff --git a/checkstyle/checkstyle.xml b/checkstyle/checkstyle.xml index ccab85ce4fecb..91045adc60856 100644 --- a/checkstyle/checkstyle.xml +++ b/checkstyle/checkstyle.xml @@ -25,7 +25,7 @@ - + @@ -71,6 +71,12 @@ + + + + + + @@ -97,6 +103,14 @@ + + + + + + + + @@ -105,7 +119,7 @@ - + @@ -114,7 +128,7 @@ - + @@ -131,6 +145,6 @@ - + diff --git a/checkstyle/import-control-core.xml b/checkstyle/import-control-core.xml index bf06a19de34b7..6e5042fd35d8a 100644 --- a/checkstyle/import-control-core.xml +++ b/checkstyle/import-control-core.xml @@ -38,26 +38,14 @@ - - - - - - - - - - - - - + + + - - - @@ -68,9 +56,6 @@ - - - diff --git a/checkstyle/import-control-jmh-benchmarks.xml b/checkstyle/import-control-jmh-benchmarks.xml new file mode 100644 index 0000000000000..f1538df3bcb4a --- /dev/null +++ b/checkstyle/import-control-jmh-benchmarks.xml @@ -0,0 +1,55 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/checkstyle/import-control.xml b/checkstyle/import-control.xml index 403cae2b7c23e..5368ef217bfbf 100644 --- a/checkstyle/import-control.xml +++ b/checkstyle/import-control.xml @@ -28,6 +28,7 @@ + @@ -47,9 +48,11 @@ + + @@ -65,6 +68,14 @@ + + + + + + + + @@ -92,27 +103,42 @@ + + + + + + + + + + + + + + + @@ -125,13 +151,19 @@ - + + + + + + + @@ -142,7 +174,11 @@ - + + + + + @@ -163,6 +199,8 @@ + + @@ -187,6 +225,7 @@ + @@ -198,6 +237,13 @@ + + + + + + + @@ -208,55 +254,55 @@ - + + + + + - + - - - + + + + + + - + + + + + - - - - - - - - - - - @@ -267,12 +313,35 @@ + + + + + + + + + + + + + + + + + + + + + + + @@ -290,11 +359,45 @@ - + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + @@ -302,11 +405,20 @@ + + + + + + + + + @@ -315,11 +427,13 @@ + + @@ -327,6 +441,22 @@ + + + + + + + + + + + + + + + + diff --git a/checkstyle/suppressions.xml b/checkstyle/suppressions.xml index 6218794f62b19..7fb20eca2a3c7 100644 --- a/checkstyle/suppressions.xml +++ b/checkstyle/suppressions.xml @@ -1,4 +1,4 @@ - + + + + + + + + + + + + + files="(Fetcher|Sender|SenderTest|ConsumerCoordinator|KafkaConsumer|KafkaProducer|Utils|TransactionManager|TransactionManagerTest|KafkaAdminClient|NetworkClient|Admin|KafkaRaftClient|KafkaRaftClientTest).java"/> + files="KerberosLogin.java|RequestResponseTest.java|ConnectMetricsRegistry.java|KafkaConsumer.java"/> + files="(NetworkClient|FieldSpec|KafkaRaftClient).java"/> - + files="(KafkaConsumer|ConsumerCoordinator|Fetcher|KafkaProducer|AbstractRequest|AbstractResponse|TransactionManager|Admin|KafkaAdminClient|MockAdminClient|KafkaRaftClient|KafkaRaftClientTest).java"/> + files="(Errors|SaslAuthenticatorTest|AgentTest|CoordinatorTest).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|DefaultSslEngineFactory).java"/> + files="(AbstractRequest|AbstractResponse|KerberosLogin|WorkerSinkTaskTest|TransactionManagerTest|SenderTest|KafkaAdminClient|ConsumerCoordinatorTest|KafkaAdminClientTest).java"/> - + + + + + + + files="MessageTest.java"/> - + + files="(Sender|Fetcher|KafkaConsumer|Metrics|RequestResponse|TransactionManager|KafkaAdminClient|Message|KafkaProducer)Test.java"/> + + + files="MockAdminClient.java"/> + files="RequestResponseTest.java|FetcherTest.java|KafkaAdminClientTest.java"/> + + + + - + files="(DistributedHerder|Worker).java"/> + - + files="(KafkaConfigBackingStore|Values|IncrementalCooperativeAssignor).java"/> - + files="Worker(SinkTask|SourceTask|Coordinator).java"/> @@ -93,111 +124,128 @@ files="JsonConverter.java"/> + files="(FileStreamSourceTask|DistributedHerder|KafkaConfigBackingStore).java"/> - - - + files="(ConnectRecord|JsonConverter|Values|ConnectHeader|ConnectHeaders).java"/> + files="(KafkaConfigBackingStore|Values).java"/> - - - - + files="(DistributedHerder|RestClient|RestServer|JsonConverter|KafkaConfigBackingStore|FileStreamSourceTask|TopicAdmin).java"/> + files="(DistributedHerder|KafkaBasedLog|WorkerSourceTaskWithTopicCreation)Test.java"/> + + + + + + files="(KafkaStreams|KStreamImpl|KTableImpl|StreamsPartitionAssignor).java"/> + files="KTableImpl.java"/> - + files="StreamThread.java"/> + files="(KStreamImpl|KTableImpl).java"/> - - + files="(StreamsPartitionAssignor|StreamThread|TaskManager).java"/> - + - - + files="(KafkaStreams|StreamsPartitionAssignor|StreamThread|TaskManager|GlobalStateManagerImpl).java"/> - + + + + + + + + + + + + + + + + files="(StreamsPartitionAssignorTest|StreamThreadTest|StreamTaskTest|TopologyTestDriverTest).java"/> - - + files="(EosBetaUpgradeIntegrationTest|KStreamKStreamJoinTest|RocksDBWindowStoreTest).java"/> - - - + files="(EosBetaUpgradeIntegrationTest|KStreamKStreamJoinTest|KTableKTableForeignKeyJoinIntegrationTest|RocksDBGenericOptionsToDbOptionsColumnFamilyOptionsAdapterTest|RelationalSmokeTest|MockProcessorContextStateStoreTest).java"/> - + files="(EosBetaUpgradeIntegrationTest|KStreamKStreamJoinTest|TaskManagerTest).java"/> - + files="(EosBetaUpgradeIntegrationTest|EosTestDriver|KStreamKStreamJoinTest|KTableKTableForeignKeyJoinIntegrationTest|RelationalSmokeTest|MockProcessorContextStateStoreTest).java"/> + + + + + + + + + + + + + - - + + files="(ProducerPerformance|StreamsResetter|Agent|TransactionalMessageCopier).java"/> + + + + requestBuilder() { @@ -101,4 +112,8 @@ public long createdTimeMs() { public int correlationId() { return correlationId; } + + public int requestTimeoutMs() { + return requestTimeoutMs; + } } diff --git a/clients/src/main/java/org/apache/kafka/clients/ClientResponse.java b/clients/src/main/java/org/apache/kafka/clients/ClientResponse.java index 0ff30e93515eb..446bf44010fc5 100644 --- a/clients/src/main/java/org/apache/kafka/clients/ClientResponse.java +++ b/clients/src/main/java/org/apache/kafka/clients/ClientResponse.java @@ -16,6 +16,7 @@ */ package org.apache.kafka.clients; +import org.apache.kafka.common.errors.AuthenticationException; import org.apache.kafka.common.errors.UnsupportedVersionException; import org.apache.kafka.common.requests.AbstractResponse; import org.apache.kafka.common.requests.RequestHeader; @@ -33,6 +34,7 @@ public class ClientResponse { private final long latencyMs; private final boolean disconnected; private final UnsupportedVersionException versionMismatch; + private final AuthenticationException authenticationException; private final AbstractResponse responseBody; /** @@ -53,6 +55,7 @@ public ClientResponse(RequestHeader requestHeader, long receivedTimeMs, boolean disconnected, UnsupportedVersionException versionMismatch, + AuthenticationException authenticationException, AbstractResponse responseBody) { this.requestHeader = requestHeader; this.callback = callback; @@ -61,6 +64,7 @@ public ClientResponse(RequestHeader requestHeader, this.latencyMs = receivedTimeMs - createdTimeMs; this.disconnected = disconnected; this.versionMismatch = versionMismatch; + this.authenticationException = authenticationException; this.responseBody = responseBody; } @@ -76,6 +80,10 @@ public UnsupportedVersionException versionMismatch() { return versionMismatch; } + public AuthenticationException authenticationException() { + return authenticationException; + } + public RequestHeader requestHeader() { return requestHeader; } 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 ea4c4dba56f3a..5e5286e329efc 100644 --- a/clients/src/main/java/org/apache/kafka/clients/ClientUtils.java +++ b/clients/src/main/java/org/apache/kafka/clients/ClientUtils.java @@ -16,29 +16,39 @@ */ package org.apache.kafka.clients; -import java.io.Closeable; -import java.net.InetSocketAddress; -import java.util.ArrayList; -import java.util.List; -import java.util.concurrent.atomic.AtomicReference; - import org.apache.kafka.common.config.AbstractConfig; +import org.apache.kafka.common.config.ConfigException; +import org.apache.kafka.common.config.SaslConfigs; +import org.apache.kafka.common.network.ChannelBuilder; import org.apache.kafka.common.network.ChannelBuilders; import org.apache.kafka.common.security.JaasContext; import org.apache.kafka.common.security.auth.SecurityProtocol; -import org.apache.kafka.common.network.ChannelBuilder; -import org.apache.kafka.common.config.ConfigException; -import org.apache.kafka.common.config.SaslConfigs; +import org.apache.kafka.common.utils.LogContext; +import org.apache.kafka.common.utils.Time; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +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 static org.apache.kafka.common.utils.Utils.getHost; import static org.apache.kafka.common.utils.Utils.getPort; -public class ClientUtils { +public final class ClientUtils { private static final Logger log = LoggerFactory.getLogger(ClientUtils.class); - public static List parseAndValidateAddresses(List urls) { + private ClientUtils() { + } + + public static List parseAndValidateAddresses(List urls, String clientDnsLookupConfig) { + return parseAndValidateAddresses(urls, ClientDnsLookup.forConfig(clientDnsLookupConfig)); + } + + public static List parseAndValidateAddresses(List urls, ClientDnsLookup clientDnsLookup) { List addresses = new ArrayList<>(); for (String url : urls) { if (url != null && !url.isEmpty()) { @@ -48,15 +58,30 @@ public static List parseAndValidateAddresses(List url if (host == null || port == null) throw new ConfigException("Invalid url in " + CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG + ": " + url); - InetSocketAddress address = new InetSocketAddress(host, port); - - if (address.isUnresolved()) { - log.warn("Removing server {} from {} as DNS resolution failed for {}", url, CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, host); + if (clientDnsLookup == ClientDnsLookup.RESOLVE_CANONICAL_BOOTSTRAP_SERVERS_ONLY) { + InetAddress[] inetAddresses = InetAddress.getAllByName(host); + for (InetAddress inetAddress : inetAddresses) { + 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); + } else { + addresses.add(address); + } + } } else { - addresses.add(address); + InetSocketAddress address = new InetSocketAddress(host, port); + if (address.isUnresolved()) { + log.warn("Couldn't resolve server {} from {} as DNS resolution failed for {}", url, CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, host); + } else { + addresses.add(address); + } } + } catch (IllegalArgumentException e) { throw new ConfigException("Invalid port in " + CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG + ": " + url); + } catch (UnknownHostException e) { + throw new ConfigException("Unknown host in " + CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG + ": " + url); } } } @@ -65,25 +90,54 @@ public static List parseAndValidateAddresses(List url 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); - } - } - } - /** + * Create a new channel builder from the provided configuration. + * * @param config client configs + * @param time the time implementation + * @param logContext the logging context + * * @return configured ChannelBuilder based on the configs. */ - public static ChannelBuilder createChannelBuilder(AbstractConfig config) { + public static ChannelBuilder createChannelBuilder(AbstractConfig config, Time time, LogContext logContext) { SecurityProtocol securityProtocol = SecurityProtocol.forName(config.getString(CommonClientConfigs.SECURITY_PROTOCOL_CONFIG)); String clientSaslMechanism = config.getString(SaslConfigs.SASL_MECHANISM); return ChannelBuilders.clientChannelBuilder(securityProtocol, JaasContext.Type.CLIENT, config, null, - clientSaslMechanism, true); + clientSaslMechanism, time, true, logContext); + } + + static List resolve(String host, ClientDnsLookup clientDnsLookup) throws UnknownHostException { + InetAddress[] addresses = InetAddress.getAllByName(host); + + switch (clientDnsLookup) { + case DEFAULT: + return Collections.singletonList(addresses[0]); + case USE_ALL_DNS_IPS: + case RESOLVE_CANONICAL_BOOTSTRAP_SERVERS_ONLY: + return filterPreferredAddresses(addresses); + } + + throw new IllegalStateException("Unhandled ClientDnsLookup instance: " + clientDnsLookup); + } + + /** + * Return a list containing the first address in `allAddresses` and subsequent addresses + * that are a subtype of the first address. + * + * The outcome is that all returned addresses are either IPv4 or IPv6 (InetAddress has two + * subclasses: Inet4Address and Inet6Address). + */ + static List filterPreferredAddresses(InetAddress[] allAddresses) { + List preferredAddresses = new ArrayList<>(); + Class clazz = null; + for (InetAddress address : allAddresses) { + if (clazz == null) { + clazz = address.getClass(); + } + if (clazz.isInstance(address)) { + preferredAddresses.add(address); + } + } + return preferredAddresses; } } 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 89e958d7c469e..1b20dcc70d716 100644 --- a/clients/src/main/java/org/apache/kafka/clients/ClusterConnectionStates.java +++ b/clients/src/main/java/org/apache/kafka/clients/ClusterConnectionStates.java @@ -16,11 +16,20 @@ */ package org.apache.kafka.clients; -import java.util.concurrent.ThreadLocalRandom; +import java.util.HashSet; +import java.util.Set; +import java.util.stream.Collectors; import org.apache.kafka.common.errors.AuthenticationException; +import org.apache.kafka.common.utils.ExponentialBackoff; +import org.apache.kafka.common.utils.LogContext; +import org.slf4j.Logger; +import java.net.InetAddress; +import java.net.UnknownHostException; +import java.util.Collections; import java.util.HashMap; +import java.util.List; import java.util.Map; /** @@ -28,17 +37,32 @@ * */ final class ClusterConnectionStates { - private final long reconnectBackoffInitMs; - private final long reconnectBackoffMaxMs; - private final static int RECONNECT_BACKOFF_EXP_BASE = 2; - private final double reconnectBackoffMaxExp; + final static int RECONNECT_BACKOFF_EXP_BASE = 2; + final static double RECONNECT_BACKOFF_JITTER = 0.2; + final static int CONNECTION_SETUP_TIMEOUT_EXP_BASE = 2; + final static double CONNECTION_SETUP_TIMEOUT_JITTER = 0.2; private final Map nodeState; + private final Logger log; + private Set connectingNodes; + private ExponentialBackoff reconnectBackoff; + private ExponentialBackoff connectionSetupTimeout; - public ClusterConnectionStates(long reconnectBackoffMs, long reconnectBackoffMaxMs) { - this.reconnectBackoffInitMs = reconnectBackoffMs; - this.reconnectBackoffMaxMs = reconnectBackoffMaxMs; - this.reconnectBackoffMaxExp = Math.log(this.reconnectBackoffMaxMs / (double) Math.max(reconnectBackoffMs, 1)) / Math.log(RECONNECT_BACKOFF_EXP_BASE); + public ClusterConnectionStates(long reconnectBackoffMs, long reconnectBackoffMaxMs, + long connectionSetupTimeoutMs, long connectionSetupTimeoutMaxMs, + LogContext logContext) { + this.log = logContext.logger(ClusterConnectionStates.class); + this.reconnectBackoff = new ExponentialBackoff( + reconnectBackoffMs, + RECONNECT_BACKOFF_EXP_BASE, + reconnectBackoffMaxMs, + RECONNECT_BACKOFF_JITTER); + this.connectionSetupTimeout = new ExponentialBackoff( + connectionSetupTimeoutMs, + CONNECTION_SETUP_TIMEOUT_EXP_BASE, + connectionSetupTimeoutMaxMs, + CONNECTION_SETUP_TIMEOUT_JITTER); this.nodeState = new HashMap<>(); + this.connectingNodes = new HashSet<>(); } /** @@ -64,11 +88,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; } /** @@ -81,8 +103,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 @@ -101,31 +123,111 @@ public boolean isConnecting(String id) { } /** - * Enter the connecting state for the given connection. + * 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 - * @param now the current time + * @param now the current time in ms + * @param host the host of the connection, to be resolved internally if needed + * @param clientDnsLookup the mode of DNS lookup to use when resolving the {@code host} */ - public void connecting(String id, long now) { - if (nodeState.containsKey(id)) { - NodeConnectionState node = nodeState.get(id); - node.lastConnectAttemptMs = now; - node.state = ConnectionState.CONNECTING; - } else { - nodeState.put(id, new NodeConnectionState(ConnectionState.CONNECTING, now, - this.reconnectBackoffInitMs)); + public void connecting(String id, long now, String host, ClientDnsLookup clientDnsLookup) { + NodeConnectionState connectionState = nodeState.get(id); + if (connectionState != null && connectionState.host().equals(host)) { + connectionState.lastConnectAttemptMs = now; + connectionState.state = ConnectionState.CONNECTING; + // Move to next resolved address, or if addresses are exhausted, mark node to be re-resolved + connectionState.moveToNextAddress(); + connectingNodes.add(id); + return; + } else if (connectionState != null) { + log.info("Hostname for node {} changed from {} to {}.", id, connectionState.host(), host); } + + // Create a new NodeConnectionState if nodeState does not already contain one + // for the specified id or if the hostname associated with the node id changed. + nodeState.put(id, new NodeConnectionState(ConnectionState.CONNECTING, now, + reconnectBackoff.backoff(0), connectionSetupTimeout.backoff(0), host, clientDnsLookup)); + connectingNodes.add(id); + } + + /** + * Returns a resolved address for the given connection, resolving it if necessary. + * @param id the id of the connection + * @throws UnknownHostException if the address was not resolvable + */ + public InetAddress currentAddress(String id) throws UnknownHostException { + return nodeState(id).currentAddress(); } /** * Enter the disconnected state for the given node. * @param id the connection we have disconnected - * @param now the current time + * @param now the current time in ms */ public void disconnected(String id, long now) { NodeConnectionState nodeState = nodeState(id); - nodeState.state = ConnectionState.DISCONNECTED; nodeState.lastConnectAttemptMs = now; updateReconnectBackoff(nodeState); + if (nodeState.state == ConnectionState.CONNECTING) { + updateConnectionSetupTimeout(nodeState); + connectingNodes.remove(id); + } else { + resetConnectionSetupTimeout(nodeState); + } + nodeState.state = ConnectionState.DISCONNECTED; + } + + /** + * Indicate that the connection is throttled until the specified deadline. + * @param id the connection to be throttled + * @param throttleUntilTimeMs the throttle deadline in milliseconds + */ + public void throttle(String id, long throttleUntilTimeMs) { + NodeConnectionState state = nodeState.get(id); + // The throttle deadline should never regress. + if (state != null && state.throttleUntilTimeMs < throttleUntilTimeMs) { + state.throttleUntilTimeMs = throttleUntilTimeMs; + } + } + + /** + * Return the remaining throttling delay in milliseconds if throttling is in progress. Return 0, otherwise. + * @param id the connection to check + * @param now the current time in ms + */ + public long throttleDelayMs(String id, long now) { + NodeConnectionState state = nodeState.get(id); + if (state != null && state.throttleUntilTimeMs > now) { + return state.throttleUntilTimeMs - now; + } else { + return 0; + } + } + + /** + * Return the number of milliseconds to wait, based on the connection state and the throttle time, before + * attempting to send data. If the connection has been established but being throttled, return throttle delay. + * Otherwise, return connection delay. + * @param id the connection to check + * @param now the current time in ms + */ + public long pollDelayMs(String id, long now) { + long throttleDelayMs = throttleDelayMs(id, now); + if (isConnected(id) && throttleDelayMs > 0) { + return throttleDelayMs; + } else { + return connectionDelay(id, now); + } } /** @@ -135,6 +237,9 @@ public void disconnected(String id, long now) { public void checkingApiVersions(String id) { NodeConnectionState nodeState = nodeState(id); nodeState.state = ConnectionState.CHECKING_API_VERSIONS; + resetReconnectBackoff(nodeState); + resetConnectionSetupTimeout(nodeState); + connectingNodes.remove(id); } /** @@ -144,13 +249,16 @@ public void checkingApiVersions(String id) { public void ready(String id) { NodeConnectionState nodeState = nodeState(id); nodeState.state = ConnectionState.READY; + nodeState.authenticationException = null; resetReconnectBackoff(nodeState); + resetConnectionSetupTimeout(nodeState); + connectingNodes.remove(id); } /** * Enter the authentication failed state for the given node. * @param id the connection identifier - * @param now the current time + * @param now the current time in ms * @param exception the authentication exception */ public void authenticationFailed(String id, long now, AuthenticationException exception) { @@ -162,26 +270,43 @@ public void authenticationFailed(String id, long now, AuthenticationException ex } /** - * Return true if the connection is ready. + * Return true if the connection is in the READY state and currently not throttled. + * * @param id the connection identifier + * @param now the current time in ms */ - public boolean isReady(String id) { - NodeConnectionState state = nodeState.get(id); - return state != null && state.state == ConnectionState.READY; + public boolean isReady(String id, long now) { + return isReady(nodeState.get(id), now); + } + + private boolean isReady(NodeConnectionState state, long now) { + return state != null && state.state == ConnectionState.READY && state.throttleUntilTimeMs <= now; } /** - * Return true if there is at least one node with connection in ready state and false otherwise. + * Return true if there is at least one node with connection in the READY state and not throttled. Returns false + * otherwise. + * + * @param now the current time in ms */ - public boolean hasReadyNodes() { + public boolean hasReadyNodes(long now) { for (Map.Entry entry : nodeState.entrySet()) { - NodeConnectionState state = entry.getValue(); - if (state != null && state.state == ConnectionState.READY) + if (isReady(entry.getValue(), now)) { return true; + } } return false; } + /** + * Return true if the connection has been established + * @param id The id of the node to check + */ + public boolean isConnected(String id) { + NodeConnectionState state = nodeState.get(id); + return state != null && state.state.isConnected(); + } + /** * Return true if the connection has been disconnected * @param id The id of the node to check @@ -208,26 +333,43 @@ public AuthenticationException authenticationException(String id) { */ private void resetReconnectBackoff(NodeConnectionState nodeState) { nodeState.failedAttempts = 0; - nodeState.reconnectBackoffMs = this.reconnectBackoffInitMs; + nodeState.reconnectBackoffMs = reconnectBackoff.backoff(0); + } + + /** + * Resets the failure count for a node and sets the connection setup timeout to the base + * value configured via socket.connection.setup.timeout.ms + * + * @param nodeState The node state object to update + */ + private void resetConnectionSetupTimeout(NodeConnectionState nodeState) { + nodeState.failedConnectAttempts = 0; + nodeState.connectionSetupTimeoutMs = connectionSetupTimeout.backoff(0); } /** - * Update the node reconnect backoff exponentially. + * Increment the failure counter, update the node reconnect backoff exponentially, + * and record the current timestamp. * The delay is reconnect.backoff.ms * 2**(failures - 1) * (+/- 20% random jitter) * Up to a (pre-jitter) maximum of reconnect.backoff.max.ms * * @param nodeState The node state object to update */ private void updateReconnectBackoff(NodeConnectionState nodeState) { - if (this.reconnectBackoffMaxMs > this.reconnectBackoffInitMs) { - nodeState.failedAttempts += 1; - double backoffExp = Math.min(nodeState.failedAttempts - 1, this.reconnectBackoffMaxExp); - double backoffFactor = Math.pow(RECONNECT_BACKOFF_EXP_BASE, backoffExp); - long reconnectBackoffMs = (long) (this.reconnectBackoffInitMs * backoffFactor); - // Actual backoff is randomized to avoid connection storms. - double randomFactor = ThreadLocalRandom.current().nextDouble(0.8, 1.2); - nodeState.reconnectBackoffMs = (long) (randomFactor * reconnectBackoffMs); - } + nodeState.reconnectBackoffMs = reconnectBackoff.backoff(nodeState.failedAttempts); + nodeState.failedAttempts++; + } + + /** + * Increment the failure counter and update the node connection setup timeout exponentially. + * The delay is socket.connection.setup.timeout.ms * 2**(failures) * (+/- 20% random jitter) + * Up to a (pre-jitter) maximum of reconnect.backoff.max.ms + * + * @param nodeState The node state object to update + */ + private void updateConnectionSetupTimeout(NodeConnectionState nodeState) { + nodeState.failedConnectAttempts++; + nodeState.connectionSetupTimeoutMs = connectionSetupTimeout.backoff(nodeState.failedConnectAttempts); } /** @@ -261,6 +403,55 @@ private NodeConnectionState nodeState(String id) { return state; } + /** + * Get the id set of nodes which are in CONNECTING state + */ + // package private for testing only + Set connectingNodes() { + return this.connectingNodes; + } + + /** + * Get the timestamp of the latest connection attempt of a given node + * @param id the connection to fetch the state for + */ + public long lastConnectAttemptMs(String id) { + NodeConnectionState nodeState = this.nodeState.get(id); + return nodeState == null ? 0 : nodeState.lastConnectAttemptMs; + } + + /** + * Get the current socket connection setup timeout of the given node. + * The base value is defined via socket.connection.setup.timeout. + * @param id the connection to fetch the state for + */ + public long connectionSetupTimeoutMs(String id) { + NodeConnectionState nodeState = this.nodeState(id); + return nodeState.connectionSetupTimeoutMs; + } + + /** + * Test if the connection to the given node has reached its timeout + * @param id the connection to fetch the state for + * @param now the current time in ms + */ + public boolean isConnectionSetupTimeout(String id, long now) { + NodeConnectionState nodeState = this.nodeState(id); + if (nodeState.state != ConnectionState.CONNECTING) + throw new IllegalStateException("Node " + id + " is not in connecting state"); + return now - lastConnectAttemptMs(id) > connectionSetupTimeoutMs(id); + } + + /** + * Return the List of nodes whose connection setup has timed out. + * @param now the current time in ms + */ + public List nodesWithConnectionSetupTimeout(long now) { + return connectingNodes.stream() + .filter(id -> isConnectionSetupTimeout(id, now)) + .collect(Collectors.toList()); + } + /** * The state of our connection to a node. */ @@ -270,18 +461,65 @@ private static class NodeConnectionState { AuthenticationException authenticationException; long lastConnectAttemptMs; long failedAttempts; + long failedConnectAttempts; long reconnectBackoffMs; + long connectionSetupTimeoutMs; + // Connection is being throttled if current time < throttleUntilTimeMs. + long throttleUntilTimeMs; + private List addresses; + private int addressIndex; + private final String host; + private final ClientDnsLookup clientDnsLookup; - public NodeConnectionState(ConnectionState state, long lastConnectAttempt, long reconnectBackoffMs) { + private NodeConnectionState(ConnectionState state, long lastConnectAttempt, long reconnectBackoffMs, + long connectionSetupTimeoutMs, String host, ClientDnsLookup clientDnsLookup) { this.state = state; + this.addresses = Collections.emptyList(); + this.addressIndex = -1; this.authenticationException = null; this.lastConnectAttemptMs = lastConnectAttempt; this.failedAttempts = 0; this.reconnectBackoffMs = reconnectBackoffMs; + this.connectionSetupTimeoutMs = connectionSetupTimeoutMs; + this.throttleUntilTimeMs = 0; + this.host = host; + this.clientDnsLookup = clientDnsLookup; + } + + public String host() { + return host; + } + + /** + * Fetches the current selected IP address for this node, resolving {@link #host()} if necessary. + * @return the selected address + * @throws UnknownHostException if resolving {@link #host()} fails + */ + private InetAddress currentAddress() throws UnknownHostException { + if (addresses.isEmpty()) { + // (Re-)initialize list + addresses = ClientUtils.resolve(host, clientDnsLookup); + addressIndex = 0; + } + + return addresses.get(addressIndex); + } + + /** + * Jumps to the next available resolved address for this node. If no other addresses are available, marks the + * list to be refreshed on the next {@link #currentAddress()} call. + */ + private void moveToNextAddress() { + if (addresses.isEmpty()) + return; // Avoid div0. List will initialize on next currentAddress() call + + addressIndex = (addressIndex + 1) % addresses.size(); + if (addressIndex == 0) + addresses = Collections.emptyList(); // Exhausted list. Re-resolve on next currentAddress() call } public String toString() { - return "NodeState(" + state + ", " + lastConnectAttemptMs + ", " + failedAttempts + ")"; + return "NodeState(" + state + ", " + lastConnectAttemptMs + ", " + failedAttempts + ", " + throttleUntilTimeMs + ")"; } } } 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 7da7a60b4cb0b..1930cf31f4877 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); @@ -41,18 +41,37 @@ public class CommonClientConfigs { + "discover the full cluster membership (which may change dynamically), this list need not contain the full set of " + "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, connect to each returned IP " + + "address in sequence until a successful connection is established. " + + "After a disconnection, the next IP is used. Once all IPs have been " + + "used once, the client resolves the IP(s) from the hostname again " + + "(both the JVM and the OS cache DNS name lookups, however). " + + "If set to resolve_canonical_bootstrap_servers_only, " + + "resolve each bootstrap address into a list of canonical names. After " + + "the bootstrap phase, this behaves the same as use_all_dns_ips. " + + "If set to default (deprecated), attempt to connect to the " + + "first IP address returned by the lookup, even if the lookup returns multiple " + + "IP addresses."; + 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 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; public static final String RECEIVE_BUFFER_CONFIG = "receive.buffer.bytes"; public static final String RECEIVE_BUFFER_DOC = "The size of the TCP receive buffer (SO_RCVBUF) to use when reading data. If the value is -1, the OS default will be used."; + public static final int RECEIVE_BUFFER_LOWER_BOUND = -1; 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."; @@ -60,7 +79,8 @@ public class CommonClientConfigs { public static final String RECONNECT_BACKOFF_MAX_MS_DOC = "The maximum amount of time in milliseconds to wait when reconnecting to a broker that has repeatedly failed to connect. If provided, the backoff per host will increase exponentially for each consecutive connection failure, up to this maximum. After calculating the backoff increase, 20% random jitter is added to avoid connection storms."; public static final String RETRIES_CONFIG = "retries"; - public static final String RETRIES_DOC = "Setting a value greater than zero will cause the client to resend any request that fails with a potentially transient error."; + public static final String RETRIES_DOC = "Setting a value greater than zero will cause the client to resend any request that fails with a potentially transient error." + + " It is recommended to set the value to either zero or `MAX_VALUE` and use corresponding timeout parameters to control how long a client should retry a request."; public static final String RETRY_BACKOFF_MS_CONFIG = "retry.backoff.ms"; public static final String RETRY_BACKOFF_MS_DOC = "The amount of time to wait before attempting to retry a failed request to a given topic partition. This avoids repeatedly sending requests in a tight loop under some failure scenarios."; @@ -77,11 +97,21 @@ public class CommonClientConfigs { 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."; + public static final String METRICS_CONTEXT_PREFIX = "metrics.context."; + public static final String SECURITY_PROTOCOL_CONFIG = "security.protocol"; public static final String SECURITY_PROTOCOL_DOC = "Protocol used to communicate with brokers. Valid values are: " + Utils.join(SecurityProtocol.names(), ", ") + "."; public static final String DEFAULT_SECURITY_PROTOCOL = "PLAINTEXT"; + public static final String SOCKET_CONNECTION_SETUP_TIMEOUT_MS_CONFIG = "socket.connection.setup.timeout.ms"; + public static final String SOCKET_CONNECTION_SETUP_TIMEOUT_MS_DOC = "The amount of time the client will wait for the socket connection to be established. If the connection is not built before the timeout elapses, clients will close the socket channel."; + public static final Long DEFAULT_SOCKET_CONNECTION_SETUP_TIMEOUT_MS = 10 * 1000L; + + public static final String SOCKET_CONNECTION_SETUP_TIMEOUT_MAX_MS_CONFIG = "socket.connection.setup.timeout.max.ms"; + public static final String SOCKET_CONNECTION_SETUP_TIMEOUT_MAX_MS_DOC = "The maximum amount of time the client will wait for the socket connection to be established. The connection setup timeout will increase exponentially for each consecutive connection failure up to this maximum. To avoid connection storms, a randomization factor of 0.2 will be applied to the timeout resulting in a random range between 20% below and 20% above the computed value."; + public static final Long DEFAULT_SOCKET_CONNECTION_SETUP_TIMEOUT_MAX_MS = 127 * 1000L; + public static final String CONNECTIONS_MAX_IDLE_MS_CONFIG = "connections.max.idle.ms"; public static final String CONNECTIONS_MAX_IDLE_MS_DOC = "Close idle connections after the number of milliseconds specified by this config."; @@ -91,6 +121,50 @@ 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 DEFAULT_API_TIMEOUT_MS_CONFIG = "default.api.timeout.ms"; + public static final String DEFAULT_API_TIMEOUT_MS_DOC = "Specifies the timeout (in milliseconds) for client APIs. " + + "This configuration is used as the default timeout for all client operations that do not specify a timeout parameter."; + /** * Postprocess the configuration so that exponential backoff is disabled when reconnect backoff * is explicitly configured but the maximum reconnect backoff is not explicitly configured. @@ -105,10 +179,19 @@ public static Map postProcessReconnectBackoffConfigs(AbstractCon HashMap rval = new HashMap<>(); if ((!config.originals().containsKey(RECONNECT_BACKOFF_MAX_MS_CONFIG)) && config.originals().containsKey(RECONNECT_BACKOFF_MS_CONFIG)) { - log.debug("Disabling exponential reconnect backoff because " + RECONNECT_BACKOFF_MS_CONFIG + - " is set, but " + RECONNECT_BACKOFF_MAX_MS_CONFIG + " is not."); + log.debug("Disabling exponential reconnect backoff because {} is set, but {} is not.", + RECONNECT_BACKOFF_MS_CONFIG, RECONNECT_BACKOFF_MAX_MS_CONFIG); rval.put(RECONNECT_BACKOFF_MAX_MS_CONFIG, parsedValues.get(RECONNECT_BACKOFF_MS_CONFIG)); } return rval; } + + public static void warnIfDeprecatedDnsLookupValue(AbstractConfig config) { + String clientDnsLookupValue = config.getString(CLIENT_DNS_LOOKUP_CONFIG); + if (clientDnsLookupValue.equals(ClientDnsLookup.DEFAULT.toString())) + log.warn("Configuration '{}' with value '{}' is deprecated and will be removed in " + + "future version. Please use '{}' or another non-deprecated value.", + CLIENT_DNS_LOOKUP_CONFIG, ClientDnsLookup.DEFAULT, + ClientDnsLookup.USE_ALL_DNS_IPS); + } } diff --git a/clients/src/main/java/org/apache/kafka/clients/ConnectionState.java b/clients/src/main/java/org/apache/kafka/clients/ConnectionState.java index 28b43d67fb1d7..f92c7fa037b0e 100644 --- a/clients/src/main/java/org/apache/kafka/clients/ConnectionState.java +++ b/clients/src/main/java/org/apache/kafka/clients/ConnectionState.java @@ -31,4 +31,8 @@ public enum ConnectionState { public boolean isDisconnected() { return this == AUTHENTICATION_FAILED || this == DISCONNECTED; } + + public boolean isConnected() { + return this == CHECKING_API_VERSIONS || this == READY; + } } diff --git a/clients/src/main/java/org/apache/kafka/clients/FetchSessionHandler.java b/clients/src/main/java/org/apache/kafka/clients/FetchSessionHandler.java new file mode 100644 index 0000000000000..50252698c456a --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/FetchSessionHandler.java @@ -0,0 +1,484 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF 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.TopicPartition; +import org.apache.kafka.common.protocol.Errors; +import org.apache.kafka.common.requests.FetchMetadata; +import org.apache.kafka.common.requests.FetchRequest.PartitionData; +import org.apache.kafka.common.requests.FetchResponse; +import org.apache.kafka.common.utils.LogContext; +import org.apache.kafka.common.utils.Utils; +import org.slf4j.Logger; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.Collections; +import java.util.Iterator; +import java.util.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Map.Entry; +import java.util.Set; + +import static org.apache.kafka.common.requests.FetchMetadata.INVALID_SESSION_ID; + +/** + * FetchSessionHandler maintains the fetch session state for connecting to a broker. + * + * Using the protocol outlined by KIP-227, clients can create incremental fetch sessions. + * These sessions allow the client to fetch information about a set of partition over + * and over, without explicitly enumerating all the partitions in the request and the + * response. + * + * FetchSessionHandler tracks the partitions which are in the session. It also + * determines which partitions need to be included in each fetch request, and what + * the attached fetch session metadata should be for each request. The corresponding + * class on the receiving broker side is FetchManager. + */ +public class FetchSessionHandler { + private final Logger log; + + private final int node; + + /** + * The metadata for the next fetch request. + */ + private FetchMetadata nextMetadata = FetchMetadata.INITIAL; + + public FetchSessionHandler(LogContext logContext, int node) { + this.log = logContext.logger(FetchSessionHandler.class); + this.node = node; + } + + /** + * All of the partitions which exist in the fetch request session. + */ + private LinkedHashMap sessionPartitions = + new LinkedHashMap<>(0); + + public static class FetchRequestData { + /** + * The partitions to send in the fetch request. + */ + private final Map toSend; + + /** + * The partitions to send in the request's "forget" list. + */ + private final List toForget; + + /** + * All of the partitions which exist in the fetch request session. + */ + private final Map sessionPartitions; + + /** + * The metadata to use in this fetch request. + */ + private final FetchMetadata metadata; + + FetchRequestData(Map toSend, + List toForget, + Map sessionPartitions, + FetchMetadata metadata) { + this.toSend = toSend; + this.toForget = toForget; + this.sessionPartitions = sessionPartitions; + this.metadata = metadata; + } + + /** + * Get the set of partitions to send in this fetch request. + */ + public Map toSend() { + return toSend; + } + + /** + * Get a list of partitions to forget in this fetch request. + */ + public List toForget() { + return toForget; + } + + /** + * Get the full set of partitions involved in this fetch request. + */ + public Map sessionPartitions() { + return sessionPartitions; + } + + public FetchMetadata metadata() { + return metadata; + } + + @Override + public String toString() { + if (metadata.isFull()) { + StringBuilder bld = new StringBuilder("FullFetchRequest("); + String prefix = ""; + for (TopicPartition partition : toSend.keySet()) { + bld.append(prefix); + bld.append(partition); + prefix = ", "; + } + bld.append(")"); + return bld.toString(); + } else { + StringBuilder bld = new StringBuilder("IncrementalFetchRequest(toSend=("); + String prefix = ""; + for (TopicPartition partition : toSend.keySet()) { + bld.append(prefix); + bld.append(partition); + prefix = ", "; + } + bld.append("), toForget=("); + prefix = ""; + for (TopicPartition partition : toForget) { + bld.append(prefix); + bld.append(partition); + prefix = ", "; + } + bld.append("), implied=("); + prefix = ""; + for (TopicPartition partition : sessionPartitions.keySet()) { + if (!toSend.containsKey(partition)) { + bld.append(prefix); + bld.append(partition); + prefix = ", "; + } + } + bld.append("))"); + return bld.toString(); + } + } + } + + public class Builder { + /** + * The next partitions which we want to fetch. + * + * It is important to maintain the insertion order of this list by using a LinkedHashMap rather + * than a regular Map. + * + * One reason is that when dealing with FULL fetch requests, if there is not enough response + * space to return data from all partitions, the server will only return data from partitions + * early in this list. + * + * Another reason is because we make use of the list ordering to optimize the preparation of + * incremental fetch requests (see below). + */ + private LinkedHashMap next; + private final boolean copySessionPartitions; + + Builder() { + this.next = new LinkedHashMap<>(); + this.copySessionPartitions = true; + } + + Builder(int initialSize, boolean copySessionPartitions) { + this.next = new LinkedHashMap<>(initialSize); + this.copySessionPartitions = copySessionPartitions; + } + + /** + * Mark that we want data from this partition in the upcoming fetch. + */ + public void add(TopicPartition topicPartition, PartitionData data) { + next.put(topicPartition, data); + } + + public FetchRequestData build() { + if (nextMetadata.isFull()) { + if (log.isDebugEnabled()) { + log.debug("Built full fetch {} for node {} with {}.", + nextMetadata, node, partitionsToLogString(next.keySet())); + } + sessionPartitions = next; + next = null; + Map toSend = + Collections.unmodifiableMap(new LinkedHashMap<>(sessionPartitions)); + return new FetchRequestData(toSend, Collections.emptyList(), toSend, nextMetadata); + } + + List added = new ArrayList<>(); + List removed = new ArrayList<>(); + List altered = new ArrayList<>(); + for (Iterator> iter = + sessionPartitions.entrySet().iterator(); iter.hasNext(); ) { + Entry entry = iter.next(); + TopicPartition topicPartition = entry.getKey(); + PartitionData prevData = entry.getValue(); + PartitionData nextData = next.remove(topicPartition); + if (nextData != null) { + if (!prevData.equals(nextData)) { + // Re-add the altered partition to the end of 'next' + next.put(topicPartition, nextData); + entry.setValue(nextData); + altered.add(topicPartition); + } + } else { + // Remove this partition from the session. + iter.remove(); + // Indicate that we no longer want to listen to this partition. + removed.add(topicPartition); + } + } + // Add any new partitions to the session. + for (Entry entry : next.entrySet()) { + TopicPartition topicPartition = entry.getKey(); + PartitionData nextData = entry.getValue(); + if (sessionPartitions.containsKey(topicPartition)) { + // In the previous loop, all the partitions which existed in both sessionPartitions + // and next were moved to the end of next, or removed from next. Therefore, + // once we hit one of them, we know there are no more unseen entries to look + // at in next. + break; + } + sessionPartitions.put(topicPartition, nextData); + added.add(topicPartition); + } + 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(next); + Map curSessionPartitions = copySessionPartitions + ? Collections.unmodifiableMap(new LinkedHashMap<>(sessionPartitions)) + : Collections.unmodifiableMap(sessionPartitions); + next = null; + return new FetchRequestData(toSend, Collections.unmodifiableList(removed), + curSessionPartitions, nextMetadata); + } + } + + public Builder newBuilder() { + return new Builder(); + } + + + /** A builder that allows for presizing the PartitionData hashmap, and avoiding making a + * secondary copy of the sessionPartitions, in cases where this is not necessarily. + * This builder is primarily for use by the Replica Fetcher + * @param size the initial size of the PartitionData hashmap + * @param copySessionPartitions boolean denoting whether the builder should make a deep copy of + * session partitions + */ + public Builder newBuilder(int size, boolean copySessionPartitions) { + return new Builder(size, copySessionPartitions); + } + + private String partitionsToLogString(Collection partitions) { + if (!log.isTraceEnabled()) { + return String.format("%d partition(s)", partitions.size()); + } + return "(" + Utils.join(partitions, ", ") + ")"; + } + + /** + * Return some partitions which are expected to be in a particular set, but which are not. + * + * @param toFind The partitions to look for. + * @param toSearch The set of partitions to search. + * @return null if all partitions were found; some of the missing ones + * in string form, if not. + */ + static Set findMissing(Set toFind, Set toSearch) { + Set ret = new LinkedHashSet<>(); + for (TopicPartition partition : toFind) { + if (!toSearch.contains(partition)) { + ret.add(partition); + } + } + return ret; + } + + /** + * Verify that a full fetch response contains all the partitions in the fetch session. + * + * @param response The response. + * @return True if the full fetch response partitions are valid. + */ + String verifyFullFetchResponsePartitions(FetchResponse response) { + StringBuilder bld = new StringBuilder(); + Set extra = + findMissing(response.responseData().keySet(), sessionPartitions.keySet()); + Set omitted = + findMissing(sessionPartitions.keySet(), response.responseData().keySet()); + if (!omitted.isEmpty()) { + bld.append("omitted=(").append(Utils.join(omitted, ", ")).append(", "); + } + if (!extra.isEmpty()) { + bld.append("extra=(").append(Utils.join(extra, ", ")).append(", "); + } + if ((!omitted.isEmpty()) || (!extra.isEmpty())) { + bld.append("response=(").append(Utils.join(response.responseData().keySet(), ", ")).append(")"); + return bld.toString(); + } + return null; + } + + /** + * Verify that the partitions in an incremental fetch response are contained in the session. + * + * @param response The response. + * @return True if the incremental fetch response partitions are valid. + */ + String verifyIncrementalFetchResponsePartitions(FetchResponse response) { + Set extra = + findMissing(response.responseData().keySet(), sessionPartitions.keySet()); + if (!extra.isEmpty()) { + StringBuilder bld = new StringBuilder(); + bld.append("extra=(").append(Utils.join(extra, ", ")).append("), "); + bld.append("response=(").append( + Utils.join(response.responseData().keySet(), ", ")).append("), "); + return bld.toString(); + } + return null; + } + + /** + * Create a string describing the partitions in a FetchResponse. + * + * @param response The FetchResponse. + * @return The string to log. + */ + private String responseDataToLogString(FetchResponse response) { + if (!log.isTraceEnabled()) { + int implied = sessionPartitions.size() - response.responseData().size(); + if (implied > 0) { + return String.format(" with %d response partition(s), %d implied partition(s)", + response.responseData().size(), implied); + } else { + return String.format(" with %d response partition(s)", + response.responseData().size()); + } + } + StringBuilder bld = new StringBuilder(); + bld.append(" with response=("). + append(Utils.join(response.responseData().keySet(), ", ")). + append(")"); + String prefix = ", implied=("; + String suffix = ""; + for (TopicPartition partition : sessionPartitions.keySet()) { + if (!response.responseData().containsKey(partition)) { + bld.append(prefix); + bld.append(partition); + prefix = ", "; + suffix = ")"; + } + } + bld.append(suffix); + return bld.toString(); + } + + /** + * Handle the fetch response. + * + * @param response The response. + * @return True if the response is well-formed; false if it can't be processed + * because of missing or unexpected partitions. + */ + public boolean handleResponse(FetchResponse response) { + if (response.error() != Errors.NONE) { + log.info("Node {} was unable to process the fetch request with {}: {}.", + node, nextMetadata, response.error()); + if (response.error() == Errors.FETCH_SESSION_ID_NOT_FOUND) { + nextMetadata = FetchMetadata.INITIAL; + } else { + nextMetadata = nextMetadata.nextCloseExisting(); + } + return false; + } + if (nextMetadata.isFull()) { + if (response.responseData().isEmpty() && response.throttleTimeMs() > 0) { + // Normally, an empty full fetch response would be invalid. However, KIP-219 + // specifies that if the broker wants to throttle the client, it will respond + // to a full fetch request with an empty response and a throttleTimeMs + // value set. We don't want to log this with a warning, since it's not an error. + // However, the empty full fetch response can't be processed, so it's still appropriate + // to return false here. + if (log.isDebugEnabled()) { + log.debug("Node {} sent a empty full fetch response to indicate that this " + + "client should be throttled for {} ms.", node, response.throttleTimeMs()); + } + nextMetadata = FetchMetadata.INITIAL; + return false; + } + String problem = verifyFullFetchResponsePartitions(response); + if (problem != null) { + log.info("Node {} sent an invalid full fetch response with {}", node, problem); + nextMetadata = FetchMetadata.INITIAL; + return false; + } else if (response.sessionId() == INVALID_SESSION_ID) { + 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. + 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; + } + } else { + String problem = verifyIncrementalFetchResponsePartitions(response); + if (problem != null) { + log.info("Node {} sent an invalid incremental fetch response with {}", node, problem); + nextMetadata = nextMetadata.nextCloseExisting(); + return false; + } else if (response.sessionId() == INVALID_SESSION_ID) { + // The incremental fetch session was closed by the server. + 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. + // We don't have to do anything special here to support KIP-219, since an empty incremental + // fetch request is perfectly valid. + if (log.isDebugEnabled()) + log.debug("Node {} sent an incremental fetch response with throttleTimeMs = {} " + + "for session {}{}", node, response.throttleTimeMs(), response.sessionId(), + responseDataToLogString(response)); + nextMetadata = nextMetadata.nextIncremental(); + return true; + } + } + } + + /** + * Handle an error sending the prepared request. + * + * When a network error occurs, we close any existing fetch session on our next request, + * and try to create a new session. + * + * @param t The exception. + */ + public void handleError(Throwable t) { + 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/InFlightRequests.java b/clients/src/main/java/org/apache/kafka/clients/InFlightRequests.java index 3689a09a11712..efca453b6c781 100644 --- a/clients/src/main/java/org/apache/kafka/clients/InFlightRequests.java +++ b/clients/src/main/java/org/apache/kafka/clients/InFlightRequests.java @@ -17,12 +17,13 @@ package org.apache.kafka.clients; import java.util.ArrayDeque; +import java.util.ArrayList; import java.util.Collections; import java.util.Deque; import java.util.HashMap; -import java.util.LinkedList; import java.util.List; import java.util.Map; +import java.util.concurrent.atomic.AtomicInteger; /** * The set of requests which have been sent or are being sent but haven't yet received a response @@ -31,6 +32,8 @@ final class InFlightRequests { private final int maxInFlightRequestsPerConnection; private final Map> requests = new HashMap<>(); + /** Thread safe total number of in flight requests. */ + private final AtomicInteger inFlightRequestCount = new AtomicInteger(0); public InFlightRequests(int maxInFlightRequestsPerConnection) { this.maxInFlightRequestsPerConnection = maxInFlightRequestsPerConnection; @@ -47,6 +50,7 @@ public void add(NetworkClient.InFlightRequest request) { this.requests.put(destination, reqs); } reqs.addFirst(request); + inFlightRequestCount.incrementAndGet(); } /** @@ -63,7 +67,9 @@ private Deque requestQueue(String node) { * Get the oldest request (the one that will be completed next) for the given node */ public NetworkClient.InFlightRequest completeNext(String node) { - return requestQueue(node).pollLast(); + NetworkClient.InFlightRequest inFlightRequest = requestQueue(node).pollLast(); + inFlightRequestCount.decrementAndGet(); + return inFlightRequest; } /** @@ -80,7 +86,9 @@ public NetworkClient.InFlightRequest lastSent(String node) { * @return The request */ public NetworkClient.InFlightRequest completeLastSent(String node) { - return requestQueue(node).pollFirst(); + NetworkClient.InFlightRequest inFlightRequest = requestQueue(node).pollFirst(); + inFlightRequestCount.decrementAndGet(); + return inFlightRequest; } /** @@ -114,13 +122,10 @@ public boolean isEmpty(String node) { } /** - * Count all in-flight requests for all nodes + * Count all in-flight requests for all nodes. This method is thread safe, but may lag the actual count. */ public int count() { - int total = 0; - for (Deque deque : this.requests.values()) - total += deque.size(); - return total; + return inFlightRequestCount.get(); } /** @@ -141,29 +146,38 @@ public boolean isEmpty() { * @return All the in-flight requests for that node that have been removed */ public Iterable clearAll(String node) { - Deque reqs = requests.remove(node); - return (reqs == null) ? Collections.emptyList() : reqs; + Deque reqs = requests.get(node); + if (reqs == null) { + return Collections.emptyList(); + } else { + final Deque clearedRequests = requests.remove(node); + inFlightRequestCount.getAndAdd(-clearedRequests.size()); + return () -> clearedRequests.descendingIterator(); + } + } + + private Boolean hasExpiredRequest(long now, Deque deque) { + for (NetworkClient.InFlightRequest request : deque) { + long timeSinceSend = Math.max(0, now - request.sendTimeMs); + if (timeSinceSend > request.requestTimeoutMs) + return true; + } + return false; } /** * Returns a list of nodes with pending in-flight request, that need to be timed out * * @param now current time in milliseconds - * @param requestTimeoutMs max time to wait for the request to be completed * @return list of nodes */ - public List getNodesWithTimedOutRequests(long now, int requestTimeoutMs) { - List nodeIds = new LinkedList<>(); + public List nodesWithTimedOutRequests(long now) { + List nodeIds = new ArrayList<>(); for (Map.Entry> requestEntry : requests.entrySet()) { String nodeId = requestEntry.getKey(); Deque deque = requestEntry.getValue(); - - if (!deque.isEmpty()) { - NetworkClient.InFlightRequest request = deque.peekLast(); - long timeSinceSend = now - request.sendTimeMs; - if (timeSinceSend > requestTimeoutMs) - nodeIds.add(nodeId); - } + if (hasExpiredRequest(now, deque)) + nodeIds.add(nodeId); } return nodeIds; } diff --git a/clients/src/main/java/org/apache/kafka/clients/KafkaClient.java b/clients/src/main/java/org/apache/kafka/clients/KafkaClient.java index 0a9b51981e7ff..18a7eefe2022b 100644 --- a/clients/src/main/java/org/apache/kafka/clients/KafkaClient.java +++ b/clients/src/main/java/org/apache/kafka/clients/KafkaClient.java @@ -48,7 +48,7 @@ public interface KafkaClient extends Closeable { boolean ready(Node node, long now); /** - * Returns the number of milliseconds to wait, based on the connection state, before attempting to send data. When + * Return the number of milliseconds to wait, based on the connection state, before attempting to send data. When * disconnected, this respects the reconnect backoff time. When connecting or connected, this handles slow/stalled * connections. * @@ -58,6 +58,16 @@ public interface KafkaClient extends Closeable { */ long connectionDelay(Node node, long now); + /** + * Return the number of milliseconds to wait, based on the connection state and the throttle time, before + * attempting to send data. If the connection has been established but being throttled, return throttle delay. + * Otherwise, return connection delay. + * + * @param node the connection to check + * @param now the current time in ms + */ + long pollDelayMs(Node node, long now); + /** * Check if the connection of the node has failed, based on the connection state. Such connection failure are * usually transient and can be resumed in the next {@link #ready(org.apache.kafka.common.Node, long)} } @@ -145,9 +155,12 @@ public interface KafkaClient extends Closeable { boolean hasInFlightRequests(String nodeId); /** - * Return true if there is at least one node with connection in ready state and false otherwise. + * Return true if there is at least one node with connection in the READY state and not throttled. Returns false + * otherwise. + * + * @param now the current time */ - boolean hasReadyNodes(); + boolean hasReadyNodes(long now); /** * Wake up the client if it is currently blocked waiting for I/O @@ -172,9 +185,32 @@ ClientRequest newClientRequest(String nodeId, AbstractRequest.Builder request * @param requestBuilder the request builder to use * @param createdTimeMs the time in milliseconds to use as the creation time of the request * @param expectResponse true iff we expect a response + * @param requestTimeoutMs Upper bound time in milliseconds to await a response before disconnecting the socket and + * cancelling the request. The request may get cancelled sooner if the socket disconnects + * for any reason including if another pending request to the same node timed out first. * @param callback the callback to invoke when we get a response */ - ClientRequest newClientRequest(String nodeId, AbstractRequest.Builder requestBuilder, long createdTimeMs, - boolean expectResponse, RequestCompletionHandler callback); + ClientRequest newClientRequest(String nodeId, + AbstractRequest.Builder requestBuilder, + long createdTimeMs, + boolean expectResponse, + int requestTimeoutMs, + RequestCompletionHandler callback); + + + + /** + * Initiates shutdown of this client. This method may be invoked from another thread while this + * client is being polled. No further requests may be sent using the client. The current poll() + * will be terminated using wakeup(). The client should be explicitly shutdown using {@link #close()} + * after poll returns. Note that {@link #close()} should not be invoked concurrently while polling. + */ + void initiateClose(); + + /** + * Returns true if the client is still active. Returns false if {@link #initiateClose()} or {@link #close()} + * was invoked for this client. + */ + boolean active(); } 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 8252cf3a9cd12..3d5154994fcda 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,13 +36,10 @@ * 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() { - this(new ArrayList(0)); + this(new ArrayList<>(0)); } public ManualMetadataUpdater(List nodes) { @@ -69,24 +66,22 @@ public long maybeUpdate(long now) { } @Override - public void handleDisconnection(String destination) { - // Do nothing + 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 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 handleFailedRequest(long now, Optional maybeFatalException) { + // Do nothing } @Override - public void handleCompletedMetadataResponse(RequestHeader requestHeader, long now, MetadataResponse response) { + public void handleSuccessfulResponse(RequestHeader requestHeader, long now, MetadataResponse response) { // Do nothing } @Override - public void requestUpdate() { - // Do nothing + public void close() { } } 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 3b8c18a6d8dd6..1dc1f39734bac 100644 --- a/clients/src/main/java/org/apache/kafka/clients/Metadata.java +++ b/clients/src/main/java/org/apache/kafka/clients/Metadata.java @@ -17,24 +17,33 @@ package org.apache.kafka.clients; import org.apache.kafka.common.Cluster; -import org.apache.kafka.common.internals.ClusterResourceListeners; +import org.apache.kafka.common.KafkaException; import org.apache.kafka.common.Node; -import org.apache.kafka.common.PartitionInfo; -import org.apache.kafka.common.errors.AuthenticationException; -import org.apache.kafka.common.errors.TimeoutException; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.errors.InvalidMetadataException; +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.protocol.Errors; +import org.apache.kafka.common.requests.MetadataRequest; +import org.apache.kafka.common.requests.MetadataResponse; +import org.apache.kafka.common.utils.LogContext; import org.slf4j.Logger; -import org.slf4j.LoggerFactory; +import java.io.Closeable; +import java.net.InetSocketAddress; 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.List; import java.util.Map; import java.util.Objects; +import java.util.Optional; import java.util.Set; +import java.util.function.Supplier; + +import static org.apache.kafka.common.record.RecordBatch.NO_PARTITION_LEADER_EPOCH; /** * A class encapsulating some of the logic around metadata. @@ -48,244 +57,424 @@ * is removed from the metadata refresh set after an update. Consumers disable topic expiry since they explicitly * manage topics while producers rely on topic expiry to limit the refresh set. */ -public final class Metadata { - - private static final Logger log = LoggerFactory.getLogger(Metadata.class); - - public static final long TOPIC_EXPIRY_MS = 5 * 60 * 1000; - private static final long TOPIC_EXPIRY_NEEDS_UPDATE = -1L; - +public class Metadata implements Closeable { + private final Logger log; private final long refreshBackoffMs; private final long metadataExpireMs; - private int version; + private int updateVersion; // bumped on every metadata response + private int requestVersion; // bumped on every new topic addition private long lastRefreshMs; private long lastSuccessfulRefreshMs; - private AuthenticationException authenticationException; - private Cluster cluster; - private boolean needUpdate; - /* Topics with expiry time */ - private final Map topics; - private final List listeners; + private KafkaException fatalException; + private Set invalidTopics; + private Set unauthorizedTopics; + private MetadataCache cache = MetadataCache.empty(); + private boolean needFullUpdate; + private boolean needPartialUpdate; private final ClusterResourceListeners clusterResourceListeners; - private boolean needMetadataForAllTopics; - private final boolean allowAutoTopicCreation; - private final boolean topicExpiryEnabled; - - public Metadata(long refreshBackoffMs, long metadataExpireMs, boolean allowAutoTopicCreation) { - this(refreshBackoffMs, metadataExpireMs, allowAutoTopicCreation, false, new ClusterResourceListeners()); - } + private boolean isClosed; + private final Map lastSeenLeaderEpochs; /** * Create a new Metadata instance - * @param refreshBackoffMs The minimum amount of time that must expire between metadata refreshes to avoid busy - * polling - * @param metadataExpireMs The maximum amount of time that metadata can be retained without refresh - * @param allowAutoTopicCreation If this and the broker config 'auto.create.topics.enable' are true, topics that - * don't exist will be created by the broker when a metadata request is sent - * @param topicExpiryEnabled If true, enable expiry of unused topics + * + * @param refreshBackoffMs The minimum amount of time that must expire between metadata refreshes to avoid busy + * polling + * @param metadataExpireMs The maximum amount of time that metadata can be retained without refresh + * @param logContext Log context corresponding to the containing client * @param clusterResourceListeners List of ClusterResourceListeners which will receive metadata updates. */ - public Metadata(long refreshBackoffMs, long metadataExpireMs, boolean allowAutoTopicCreation, - boolean topicExpiryEnabled, ClusterResourceListeners clusterResourceListeners) { + public Metadata(long refreshBackoffMs, + long metadataExpireMs, + LogContext logContext, + ClusterResourceListeners clusterResourceListeners) { + this.log = logContext.logger(Metadata.class); this.refreshBackoffMs = refreshBackoffMs; this.metadataExpireMs = metadataExpireMs; - this.allowAutoTopicCreation = allowAutoTopicCreation; - this.topicExpiryEnabled = topicExpiryEnabled; this.lastRefreshMs = 0L; this.lastSuccessfulRefreshMs = 0L; - this.version = 0; - this.cluster = Cluster.empty(); - this.needUpdate = false; - this.topics = new HashMap<>(); - this.listeners = new ArrayList<>(); + this.requestVersion = 0; + this.updateVersion = 0; + this.needFullUpdate = false; + this.needPartialUpdate = false; this.clusterResourceListeners = clusterResourceListeners; - this.needMetadataForAllTopics = false; + this.isClosed = false; + this.lastSeenLeaderEpochs = new HashMap<>(); + this.invalidTopics = Collections.emptySet(); + this.unauthorizedTopics = Collections.emptySet(); } /** * Get the current cluster info without blocking */ public synchronized Cluster fetch() { - return this.cluster; + return cache.cluster(); } /** - * Add the topic to maintain in the metadata. If topic expiry is enabled, expiry time - * will be reset on the next update. + * Return the next time when the current cluster info can be updated (i.e., backoff time has elapsed). + * + * @param nowMs current time in ms + * @return remaining time in ms till the cluster info can be updated again */ - public synchronized void add(String topic) { - Objects.requireNonNull(topic, "topic cannot be null"); - if (topics.put(topic, TOPIC_EXPIRY_NEEDS_UPDATE) == null) { - requestUpdateForNewTopics(); - } + public synchronized long timeToAllowUpdate(long nowMs) { + return Math.max(this.lastRefreshMs + this.refreshBackoffMs - nowMs, 0); } /** * The next time to update the cluster info is the maximum of the time the current info will expire and the time the * current info can be updated (i.e. backoff time has elapsed); If an update has been request then the expiry time * is now + * + * @param nowMs current time in ms + * @return remaining time in ms till updating the cluster info */ public synchronized long timeToNextUpdate(long nowMs) { - long timeToExpire = needUpdate ? 0 : Math.max(this.lastSuccessfulRefreshMs + this.metadataExpireMs - nowMs, 0); - long timeToAllowUpdate = this.lastRefreshMs + this.refreshBackoffMs - nowMs; - return Math.max(timeToExpire, timeToAllowUpdate); + long timeToExpire = updateRequested() ? 0 : Math.max(this.lastSuccessfulRefreshMs + this.metadataExpireMs - nowMs, 0); + return Math.max(timeToExpire, timeToAllowUpdate(nowMs)); + } + + public long metadataExpireMs() { + return this.metadataExpireMs; } /** - * Request an update of the current cluster metadata info, return the current version before the update + * Request an update of the current cluster metadata info, return the current updateVersion before the update */ public synchronized int requestUpdate() { - this.needUpdate = true; - return this.version; + this.needFullUpdate = true; + return this.updateVersion; + } + + public synchronized int requestUpdateForNewTopics() { + // Override the timestamp of last refresh to let immediate update. + this.lastRefreshMs = 0; + this.needPartialUpdate = true; + this.requestVersion++; + return this.updateVersion; + } + + /** + * Request an update for the partition metadata iff we have seen a newer leader epoch. This is called by the client + * any time it handles a response from the broker that includes leader epoch, except for UpdateMetadata which + * follows a different code path ({@link #update}). + * + * @param topicPartition + * @param leaderEpoch + * @return true if we updated the last seen epoch, false otherwise + */ + public synchronized boolean updateLastSeenEpochIfNewer(TopicPartition topicPartition, int leaderEpoch) { + Objects.requireNonNull(topicPartition, "TopicPartition cannot be null"); + if (leaderEpoch < 0) + throw new IllegalArgumentException("Invalid leader epoch " + leaderEpoch + " (must be non-negative)"); + + Integer oldEpoch = lastSeenLeaderEpochs.get(topicPartition); + log.trace("Determining if we should replace existing epoch {} with new epoch {} for partition {}", oldEpoch, leaderEpoch, topicPartition); + + final boolean updated; + if (oldEpoch == null) { + log.debug("Not replacing null epoch with new epoch {} for partition {}", leaderEpoch, topicPartition); + updated = false; + } else if (leaderEpoch > oldEpoch) { + log.debug("Updating last seen epoch from {} to {} for partition {}", oldEpoch, leaderEpoch, topicPartition); + lastSeenLeaderEpochs.put(topicPartition, leaderEpoch); + updated = true; + } else { + log.debug("Not replacing existing epoch {} with new epoch {} for partition {}", oldEpoch, leaderEpoch, topicPartition); + updated = false; + } + + this.needFullUpdate = this.needFullUpdate || updated; + return updated; + } + + public Optional lastSeenLeaderEpoch(TopicPartition topicPartition) { + return Optional.ofNullable(lastSeenLeaderEpochs.get(topicPartition)); } /** * Check whether an update has been explicitly requested. + * * @return true if an update was requested, false otherwise */ public synchronized boolean updateRequested() { - return this.needUpdate; + return this.needFullUpdate || this.needPartialUpdate; } /** - * If any non-retriable authentication exceptions were encountered during - * metadata update, clear and return the exception. + * Return the cached partition info if it exists and a newer leader epoch isn't known about. */ - public synchronized AuthenticationException getAndClearAuthenticationException() { - if (authenticationException != null) { - AuthenticationException exception = authenticationException; - authenticationException = null; - return exception; - } else - return null; + synchronized Optional partitionMetadataIfCurrent(TopicPartition topicPartition) { + Integer epoch = lastSeenLeaderEpochs.get(topicPartition); + Optional partitionMetadata = cache.partitionMetadata(topicPartition); + if (epoch == null) { + // old cluster format (no epochs) + return partitionMetadata; + } else { + return partitionMetadata.filter(metadata -> + metadata.leaderEpoch.orElse(NO_PARTITION_LEADER_EPOCH).equals(epoch)); + } + } + + public synchronized LeaderAndEpoch currentLeader(TopicPartition topicPartition) { + Optional maybeMetadata = partitionMetadataIfCurrent(topicPartition); + if (!maybeMetadata.isPresent()) + return new LeaderAndEpoch(Optional.empty(), Optional.ofNullable(lastSeenLeaderEpochs.get(topicPartition))); + + MetadataResponse.PartitionMetadata partitionMetadata = maybeMetadata.get(); + Optional leaderEpochOpt = partitionMetadata.leaderEpoch; + Optional leaderNodeOpt = partitionMetadata.leaderId.flatMap(cache::nodeById); + return new LeaderAndEpoch(leaderNodeOpt, leaderEpochOpt); + } + + public synchronized void bootstrap(List addresses) { + this.needFullUpdate = true; + this.updateVersion += 1; + this.cache = MetadataCache.bootstrap(addresses); + } + + /** + * Update metadata assuming the current request version. + * + * For testing only. + */ + public synchronized void updateWithCurrentRequestVersion(MetadataResponse response, boolean isPartialUpdate, long nowMs) { + this.update(this.requestVersion, response, isPartialUpdate, nowMs); } /** - * Wait for metadata update until the current version is larger than the last version we know of + * Updates the cluster metadata. If topic expiry is enabled, expiry time + * is set for topics if required and expired topics are removed from the metadata. + * + * @param requestVersion The request version corresponding to the update response, as provided by + * {@link #newMetadataRequestAndVersion(long)}. + * @param response metadata response received from the broker + * @param isPartialUpdate whether the metadata request was for a subset of the active topics + * @param nowMs current time in milliseconds */ - public synchronized void awaitUpdate(final int lastVersion, final long maxWaitMs) throws InterruptedException { - if (maxWaitMs < 0) { - throw new IllegalArgumentException("Max time to wait for metadata updates should not be < 0 milliseconds"); + public synchronized void update(int requestVersion, MetadataResponse response, boolean isPartialUpdate, long nowMs) { + Objects.requireNonNull(response, "Metadata response cannot be null"); + if (isClosed()) + throw new IllegalStateException("Update requested after metadata close"); + + this.needPartialUpdate = requestVersion < this.requestVersion; + this.lastRefreshMs = nowMs; + this.updateVersion += 1; + if (!isPartialUpdate) { + this.needFullUpdate = false; + this.lastSuccessfulRefreshMs = nowMs; } - long begin = System.currentTimeMillis(); - long remainingWaitMs = maxWaitMs; - while (this.version <= lastVersion) { - AuthenticationException ex = getAndClearAuthenticationException(); - if (ex != null) - throw ex; - if (remainingWaitMs != 0) - wait(remainingWaitMs); - long elapsed = System.currentTimeMillis() - begin; - if (elapsed >= maxWaitMs) - throw new TimeoutException("Failed to update metadata after " + maxWaitMs + " ms."); - remainingWaitMs = maxWaitMs - elapsed; + + String previousClusterId = cache.clusterResource().clusterId(); + + this.cache = handleMetadataResponse(response, isPartialUpdate, nowMs); + + Cluster cluster = cache.cluster(); + maybeSetMetadataError(cluster); + + this.lastSeenLeaderEpochs.keySet().removeIf(tp -> !retainTopic(tp.topic(), false, nowMs)); + + String newClusterId = cache.clusterResource().clusterId(); + if (!Objects.equals(previousClusterId, newClusterId)) { + log.info("Cluster ID: {}", newClusterId); + } + clusterResourceListeners.onUpdate(cache.clusterResource()); + + log.debug("Updated cluster metadata updateVersion {} to {}", this.updateVersion, this.cache); + } + + private void maybeSetMetadataError(Cluster cluster) { + clearRecoverableErrors(); + checkInvalidTopics(cluster); + checkUnauthorizedTopics(cluster); + } + + private void checkInvalidTopics(Cluster cluster) { + if (!cluster.invalidTopics().isEmpty()) { + log.error("Metadata response reported invalid topics {}", 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()); + unauthorizedTopics = new HashSet<>(cluster.unauthorizedTopics()); } } /** - * Replace the current set of topics maintained to the one provided. - * If topic expiry is enabled, expiry time of the topics will be - * reset on the next update. - * @param topics + * Transform a MetadataResponse into a new MetadataCache instance. */ - public synchronized void setTopics(Collection topics) { - if (!this.topics.keySet().containsAll(topics)) { - requestUpdateForNewTopics(); + private MetadataCache handleMetadataResponse(MetadataResponse metadataResponse, boolean isPartialUpdate, long nowMs) { + // All encountered topics. + Set topics = new HashSet<>(); + + // Retained topics to be passed to the metadata cache. + Set internalTopics = new HashSet<>(); + Set unauthorizedTopics = new HashSet<>(); + Set invalidTopics = new HashSet<>(); + + List partitions = new ArrayList<>(); + for (MetadataResponse.TopicMetadata metadata : metadataResponse.topicMetadata()) { + topics.add(metadata.topic()); + + if (!retainTopic(metadata.topic(), metadata.isInternal(), nowMs)) + continue; + + if (metadata.isInternal()) + internalTopics.add(metadata.topic()); + + if (metadata.error() == Errors.NONE) { + for (MetadataResponse.PartitionMetadata partitionMetadata : metadata.partitionMetadata()) { + // Even if the partition's metadata includes an error, we need to handle + // the update to catch new epochs + updateLatestMetadata(partitionMetadata, metadataResponse.hasReliableLeaderEpochs()) + .ifPresent(partitions::add); + + if (partitionMetadata.error.exception() instanceof InvalidMetadataException) { + log.debug("Requesting metadata update for partition {} due to error {}", + partitionMetadata.topicPartition, partitionMetadata.error); + requestUpdate(); + } + } + } else { + if (metadata.error().exception() instanceof InvalidMetadataException) { + log.debug("Requesting metadata update for topic {} due to error {}", metadata.topic(), metadata.error()); + requestUpdate(); + } + + if (metadata.error() == Errors.INVALID_TOPIC_EXCEPTION) + invalidTopics.add(metadata.topic()); + else if (metadata.error() == Errors.TOPIC_AUTHORIZATION_FAILED) + unauthorizedTopics.add(metadata.topic()); + } } - this.topics.clear(); - for (String topic : topics) - this.topics.put(topic, TOPIC_EXPIRY_NEEDS_UPDATE); + + Map nodes = metadataResponse.brokersById(); + if (isPartialUpdate) + return this.cache.mergeWith(metadataResponse.clusterId(), nodes, partitions, + unauthorizedTopics, invalidTopics, internalTopics, metadataResponse.controller(), + (topic, isInternal) -> !topics.contains(topic) && retainTopic(topic, isInternal, nowMs)); + else + return new MetadataCache(metadataResponse.clusterId(), nodes, partitions, + unauthorizedTopics, invalidTopics, internalTopics, metadataResponse.controller()); } /** - * Get the list of topics we are currently maintaining metadata for + * Compute the latest partition metadata to cache given ordering by leader epochs (if both + * available and reliable). */ - public synchronized Set topics() { - return new HashSet<>(this.topics.keySet()); + private Optional updateLatestMetadata( + MetadataResponse.PartitionMetadata partitionMetadata, + boolean hasReliableLeaderEpoch) { + TopicPartition tp = partitionMetadata.topicPartition; + 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 + Integer currentEpoch = lastSeenLeaderEpochs.get(tp); + if (currentEpoch == null || newEpoch >= currentEpoch) { + log.debug("Updating last seen epoch for partition {} from {} to epoch {} from new metadata", tp, currentEpoch, newEpoch); + lastSeenLeaderEpochs.put(tp, newEpoch); + return Optional.of(partitionMetadata); + } else { + // Otherwise ignore the new metadata and use the previously cached info + log.debug("Got metadata for an older epoch {} (current is {}) for partition {}, not updating", newEpoch, currentEpoch, tp); + return cache.partitionMetadata(tp); + } + } else { + // Handle old cluster formats as well as error responses where leader and epoch are missing + lastSeenLeaderEpochs.remove(tp); + return Optional.of(partitionMetadata.withoutLeaderEpoch()); + } } /** - * Check if a topic is already in the topic set. - * @param topic topic to check - * @return true if the topic exists, false otherwise + * 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 boolean containsTopic(String topic) { - return this.topics.containsKey(topic); + public synchronized void maybeThrowAnyException() { + clearErrorsAndMaybeThrowException(this::recoverableException); } /** - * Updates the cluster metadata. If topic expiry is enabled, expiry time - * is set for topics if required and expired topics are removed from the metadata. - * - * @param cluster the cluster containing metadata for topics with valid metadata - * @param unavailableTopics topics which are non-existent or have one or more partitions whose - * leader is not known - * @param now current time in milliseconds + * 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 update(Cluster cluster, Set unavailableTopics, long now) { - Objects.requireNonNull(cluster, "cluster should not be null"); - - this.needUpdate = false; - this.lastRefreshMs = now; - this.lastSuccessfulRefreshMs = now; - this.version += 1; - - if (topicExpiryEnabled) { - // Handle expiry of topics from the metadata refresh set. - for (Iterator> it = topics.entrySet().iterator(); it.hasNext(); ) { - Map.Entry entry = it.next(); - long expireMs = entry.getValue(); - if (expireMs == TOPIC_EXPIRY_NEEDS_UPDATE) - entry.setValue(now + TOPIC_EXPIRY_MS); - else if (expireMs <= now) { - it.remove(); - log.debug("Removing unused topic {} from the metadata list, expiryMs {} now {}", entry.getKey(), expireMs, now); - } - } + protected synchronized void maybeThrowFatalException() { + KafkaException metadataException = this.fatalException; + if (metadataException != null) { + fatalException = null; + throw metadataException; } + } - for (Listener listener: listeners) - listener.onMetadataUpdate(cluster, unavailableTopics); + /** + * 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)); + } - String previousClusterId = cluster.clusterResource().clusterId(); + private void clearErrorsAndMaybeThrowException(Supplier recoverableExceptionSupplier) { + KafkaException metadataException = Optional.ofNullable(fatalException).orElseGet(recoverableExceptionSupplier); + fatalException = null; + clearRecoverableErrors(); + if (metadataException != null) + throw metadataException; + } - if (this.needMetadataForAllTopics) { - // the listener may change the interested topics, which could cause another metadata refresh. - // If we have already fetched all topics, however, another fetch should be unnecessary. - this.needUpdate = false; - this.cluster = getClusterForCurrentTopics(cluster); - } else { - this.cluster = cluster; - } + // 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; + } - // The bootstrap cluster is guaranteed not to have any useful information - if (!cluster.isBootstrapConfigured()) { - String clusterId = cluster.clusterResource().clusterId(); - if (clusterId == null ? previousClusterId != null : !clusterId.equals(previousClusterId)) - log.info("Cluster ID: {}", cluster.clusterResource().clusterId()); - clusterResourceListeners.onUpdate(cluster.clusterResource()); - } + 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; + } - notifyAll(); - log.debug("Updated cluster metadata version {} to {}", this.version, this.cluster); + 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; - if (authenticationException != null) - this.notifyAll(); } /** - * @return The current metadata version + * 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; + } + + /** + * @return The current metadata updateVersion */ - public synchronized int version() { - return this.version; + public synchronized int updateVersion() { + return this.updateVersion; } /** @@ -295,84 +484,122 @@ public synchronized long lastSuccessfulUpdate() { return this.lastSuccessfulRefreshMs; } - public boolean allowAutoTopicCreation() { - return allowAutoTopicCreation; + /** + * Close this metadata instance to indicate that metadata updates are no longer possible. + */ + @Override + public synchronized void close() { + this.isClosed = true; } /** - * Set state to indicate if metadata for all topics in Kafka cluster is required or not. - * @param needMetadataForAllTopics boolean indicating need for metadata of all topics in cluster. + * Check if this metadata instance has been closed. See {@link #close()} for more information. + * + * @return True if this instance has been closed; false otherwise */ - public synchronized void needMetadataForAllTopics(boolean needMetadataForAllTopics) { - if (needMetadataForAllTopics && !this.needMetadataForAllTopics) { - requestUpdateForNewTopics(); + public synchronized boolean isClosed() { + return this.isClosed; + } + + public synchronized MetadataRequestAndVersion newMetadataRequestAndVersion(long nowMs) { + MetadataRequest.Builder request = null; + boolean isPartialUpdate = false; + + // Perform a partial update only if a full update hasn't been requested, and the last successful + // hasn't exceeded the metadata refresh time. + if (!this.needFullUpdate && this.lastSuccessfulRefreshMs + this.metadataExpireMs > nowMs) { + request = newMetadataRequestBuilderForNewTopics(); + isPartialUpdate = true; } - this.needMetadataForAllTopics = needMetadataForAllTopics; + if (request == null) { + request = newMetadataRequestBuilder(); + isPartialUpdate = false; + } + return new MetadataRequestAndVersion(request, requestVersion, isPartialUpdate); } /** - * Get whether metadata for all topics is needed or not + * Constructs and returns a metadata request builder for fetching cluster data and all active topics. + * + * @return the constructed non-null metadata builder */ - public synchronized boolean needMetadataForAllTopics() { - return this.needMetadataForAllTopics; + protected MetadataRequest.Builder newMetadataRequestBuilder() { + return MetadataRequest.Builder.allTopics(); } /** - * Add a Metadata listener that gets notified of metadata updates + * Constructs and returns a metadata request builder for fetching cluster data and any uncached topics, + * otherwise null if the functionality is not supported. + * + * @return the constructed metadata builder, or null if not supported */ - public synchronized void addListener(Listener listener) { - this.listeners.add(listener); + protected MetadataRequest.Builder newMetadataRequestBuilderForNewTopics() { + return null; } - /** - * Stop notifying the listener of metadata updates - */ - public synchronized void removeListener(Listener listener) { - this.listeners.remove(listener); + protected boolean retainTopic(String topic, boolean isInternal, long nowMs) { + return true; + } + + public static class MetadataRequestAndVersion { + public final MetadataRequest.Builder requestBuilder; + public final int requestVersion; + public final boolean isPartialUpdate; + + private MetadataRequestAndVersion(MetadataRequest.Builder requestBuilder, + int requestVersion, + boolean isPartialUpdate) { + this.requestBuilder = requestBuilder; + this.requestVersion = requestVersion; + this.isPartialUpdate = isPartialUpdate; + } } /** - * MetadataUpdate Listener + * Represents current leader state known in metadata. It is possible that we know the leader, but not the + * epoch if the metadata is received from a broker which does not support a sufficient Metadata API version. + * It is also possible that we know of the leader epoch, but not the leader when it is derived + * from an external source (e.g. a committed offset). */ - public interface Listener { - /** - * Callback invoked on metadata update. - * - * @param cluster the cluster containing metadata for topics with valid metadata - * @param unavailableTopics topics which are non-existent or have one or more partitions whose - * leader is not known - */ - void onMetadataUpdate(Cluster cluster, Set unavailableTopics); - } - - private synchronized void requestUpdateForNewTopics() { - // Override the timestamp of last refresh to let immediate update. - this.lastRefreshMs = 0; - requestUpdate(); - } + public static class LeaderAndEpoch { + private static final LeaderAndEpoch NO_LEADER_OR_EPOCH = new LeaderAndEpoch(Optional.empty(), Optional.empty()); - private Cluster getClusterForCurrentTopics(Cluster cluster) { - Set unauthorizedTopics = new HashSet<>(); - Collection partitionInfos = new ArrayList<>(); - List nodes = Collections.emptyList(); - Set internalTopics = Collections.emptySet(); - Node controller = null; - String clusterId = null; - if (cluster != null) { - clusterId = cluster.clusterResource().clusterId(); - internalTopics = cluster.internalTopics(); - unauthorizedTopics.addAll(cluster.unauthorizedTopics()); - unauthorizedTopics.retainAll(this.topics.keySet()); - - for (String topic : this.topics.keySet()) { - List partitionInfoList = cluster.partitionsForTopic(topic); - if (!partitionInfoList.isEmpty()) { - partitionInfos.addAll(partitionInfoList); - } - } - nodes = cluster.nodes(); - controller = cluster.controller(); + public final Optional leader; + public final Optional epoch; + + public LeaderAndEpoch(Optional 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") + + '}'; } - return new Cluster(clusterId, nodes, partitionInfos, unauthorizedTopics, internalTopics, controller); } } diff --git a/clients/src/main/java/org/apache/kafka/clients/MetadataCache.java b/clients/src/main/java/org/apache/kafka/clients/MetadataCache.java new file mode 100644 index 0000000000000..2da61ab2c6287 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/MetadataCache.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.clients; + +import org.apache.kafka.common.Cluster; +import org.apache.kafka.common.ClusterResource; +import org.apache.kafka.common.Node; +import org.apache.kafka.common.PartitionInfo; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.requests.MetadataResponse; +import org.apache.kafka.common.requests.MetadataResponse.PartitionMetadata; + +import java.net.InetSocketAddress; +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.Optional; +import java.util.Set; +import java.util.function.BiPredicate; +import java.util.function.Predicate; +import java.util.stream.Collectors; + +/** + * An internal mutable cache of nodes, topics, and partitions in the Kafka cluster. This keeps an up-to-date Cluster + * instance which is optimized for read access. + */ +public class MetadataCache { + private final String clusterId; + private final Map nodes; + private final Set unauthorizedTopics; + private final Set invalidTopics; + private final Set internalTopics; + private final Node controller; + private final Map metadataByPartition; + + private Cluster clusterInstance; + + MetadataCache(String clusterId, + Map nodes, + Collection partitions, + Set unauthorizedTopics, + Set invalidTopics, + Set internalTopics, + Node controller) { + this(clusterId, nodes, partitions, unauthorizedTopics, invalidTopics, internalTopics, controller, null); + } + + private MetadataCache(String clusterId, + Map nodes, + Collection partitions, + Set unauthorizedTopics, + Set invalidTopics, + Set internalTopics, + Node controller, + Cluster clusterInstance) { + this.clusterId = clusterId; + this.nodes = nodes; + this.unauthorizedTopics = unauthorizedTopics; + this.invalidTopics = invalidTopics; + this.internalTopics = internalTopics; + this.controller = controller; + + this.metadataByPartition = new HashMap<>(partitions.size()); + for (PartitionMetadata p : partitions) { + this.metadataByPartition.put(p.topicPartition, p); + } + + if (clusterInstance == null) { + computeClusterView(); + } else { + this.clusterInstance = clusterInstance; + } + } + + Optional partitionMetadata(TopicPartition topicPartition) { + return Optional.ofNullable(metadataByPartition.get(topicPartition)); + } + + Optional nodeById(int id) { + return Optional.ofNullable(nodes.get(id)); + } + + Cluster cluster() { + if (clusterInstance == null) { + throw new IllegalStateException("Cached Cluster instance should not be null, but was."); + } else { + return clusterInstance; + } + } + + ClusterResource clusterResource() { + return new ClusterResource(clusterId); + } + + /** + * Merges the metadata cache's contents with the provided metadata, returning a new metadata cache. The provided + * metadata is presumed to be more recent than the cache's metadata, and therefore all overlapping metadata will + * be overridden. + * + * @param newClusterId the new cluster Id + * @param newNodes the new set of nodes + * @param addPartitions partitions to add + * @param addUnauthorizedTopics unauthorized topics to add + * @param addInternalTopics internal topics to add + * @param newController the new controller node + * @param retainTopic returns whether a topic's metadata should be retained + * @return the merged metadata cache + */ + MetadataCache mergeWith(String newClusterId, + Map newNodes, + Collection addPartitions, + Set addUnauthorizedTopics, + Set addInvalidTopics, + Set addInternalTopics, + Node newController, + BiPredicate retainTopic) { + + Predicate shouldRetainTopic = topic -> retainTopic.test(topic, internalTopics.contains(topic)); + + Map newMetadataByPartition = new HashMap<>(addPartitions.size()); + for (PartitionMetadata partition : addPartitions) { + newMetadataByPartition.put(partition.topicPartition, partition); + } + for (Map.Entry entry : metadataByPartition.entrySet()) { + if (shouldRetainTopic.test(entry.getKey().topic())) { + newMetadataByPartition.putIfAbsent(entry.getKey(), entry.getValue()); + } + } + + Set newUnauthorizedTopics = fillSet(addUnauthorizedTopics, unauthorizedTopics, shouldRetainTopic); + Set newInvalidTopics = fillSet(addInvalidTopics, invalidTopics, shouldRetainTopic); + Set newInternalTopics = fillSet(addInternalTopics, internalTopics, shouldRetainTopic); + + return new MetadataCache(newClusterId, newNodes, newMetadataByPartition.values(), newUnauthorizedTopics, + newInvalidTopics, newInternalTopics, newController); + } + + /** + * Copies {@code baseSet} and adds all non-existent elements in {@code fillSet} such that {@code predicate} is true. + * In other words, all elements of {@code baseSet} will be contained in the result, with additional non-overlapping + * elements in {@code fillSet} where the predicate is true. + * + * @param baseSet the base elements for the resulting set + * @param fillSet elements to be filled into the resulting set + * @param predicate tested against the fill set to determine whether elements should be added to the base set + */ + private static Set fillSet(Set baseSet, Set fillSet, Predicate predicate) { + Set result = new HashSet<>(baseSet); + for (T element : fillSet) { + if (predicate.test(element)) { + result.add(element); + } + } + return result; + } + + private void computeClusterView() { + List partitionInfos = metadataByPartition.values() + .stream() + .map(metadata -> MetadataResponse.toPartitionInfo(metadata, nodes)) + .collect(Collectors.toList()); + this.clusterInstance = new Cluster(clusterId, nodes.values(), partitionInfos, unauthorizedTopics, + invalidTopics, internalTopics, controller); + } + + static MetadataCache bootstrap(List addresses) { + Map nodes = new HashMap<>(); + int nodeId = -1; + for (InetSocketAddress address : addresses) { + nodes.put(nodeId, new Node(nodeId, address.getHostString(), address.getPort())); + nodeId--; + } + return new MetadataCache(null, nodes, Collections.emptyList(), + Collections.emptySet(), Collections.emptySet(), Collections.emptySet(), + null, Cluster.bootstrap(addresses)); + } + + static MetadataCache empty() { + return new MetadataCache(null, Collections.emptyMap(), Collections.emptyList(), + Collections.emptySet(), Collections.emptySet(), Collections.emptySet(), null, Cluster.empty()); + } + + @Override + public String toString() { + return "MetadataCache{" + + "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 cb821d6a86064..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,12 +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 @@ -29,7 +33,7 @@ *

* This class is not thread-safe! */ -interface MetadataUpdater { +public interface MetadataUpdater extends Closeable { /** * Gets the current cluster info without blocking. @@ -45,7 +49,7 @@ interface MetadataUpdater { * 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 @@ -54,32 +58,36 @@ interface MetadataUpdater { long maybeUpdate(long now); /** - * If `request` is a metadata request, handles it and return `true`. Otherwise, returns `false`. + * 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); /** - * If `request` is a metadata request, handles it and returns `true`. Otherwise, returns `false`. + * Handle responses for metadata requests. * * 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); + void handleSuccessfulResponse(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). + * Close this updater. */ - void requestUpdate(); + @Override + void close(); } 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 0654a91c8b223..840185da464b6 100644 --- a/clients/src/main/java/org/apache/kafka/clients/NetworkClient.java +++ b/clients/src/main/java/org/apache/kafka/clients/NetworkClient.java @@ -17,41 +17,52 @@ 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.NetworkSend; 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.types.Struct; +import org.apache.kafka.common.protocol.types.SchemaException; 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.CorrelationIdMismatchException; import org.apache.kafka.common.requests.MetadataRequest; 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; import org.slf4j.Logger; import java.io.IOException; +import java.net.InetAddress; import java.net.InetSocketAddress; +import java.nio.BufferUnderflowException; 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; /** * A network client for asynchronous request/response network i/o. This is an internal class used to implement the @@ -61,6 +72,12 @@ */ public class NetworkClient implements KafkaClient { + private enum State { + ACTIVE, + CLOSING, + CLOSED + } + private final Logger log; /* the selector used to perform network i/o */ @@ -88,12 +105,14 @@ public class NetworkClient implements KafkaClient { /* the current correlation id to use when sending requests to servers */ private int correlation; - /* max time in ms for the producer to wait for acknowledgement from server*/ - private final int requestTimeoutMs; + /* default timeout for individual requests to await acknowledgement from servers */ + private final int defaultRequestTimeoutMs; /* time in ms to wait before retrying to create connection to a server */ private final long reconnectBackoffMs; + private final ClientDnsLookup clientDnsLookup; + private final Time time; /** @@ -109,6 +128,8 @@ public class NetworkClient implements KafkaClient { private final Sensor throttleTimeSensor; + private final AtomicReference state; + public NetworkClient(Selectable selector, Metadata metadata, String clientId, @@ -117,15 +138,32 @@ public NetworkClient(Selectable selector, long reconnectBackoffMax, int socketSendBuffer, int socketReceiveBuffer, - int requestTimeoutMs, + int defaultRequestTimeoutMs, + long connectionSetupTimeoutMs, + long connectionSetupTimeoutMaxMs, + ClientDnsLookup clientDnsLookup, Time time, boolean discoverBrokerVersions, ApiVersions apiVersions, LogContext logContext) { - this(null, metadata, selector, clientId, maxInFlightRequestsPerConnection, - reconnectBackoffMs, reconnectBackoffMax, - socketSendBuffer, socketReceiveBuffer, requestTimeoutMs, time, - discoverBrokerVersions, apiVersions, null, logContext); + this(null, + metadata, + selector, + clientId, + maxInFlightRequestsPerConnection, + reconnectBackoffMs, + reconnectBackoffMax, + socketSendBuffer, + socketReceiveBuffer, + defaultRequestTimeoutMs, + connectionSetupTimeoutMs, + connectionSetupTimeoutMaxMs, + clientDnsLookup, + time, + discoverBrokerVersions, + apiVersions, + null, + logContext); } public NetworkClient(Selectable selector, @@ -136,16 +174,33 @@ public NetworkClient(Selectable selector, long reconnectBackoffMax, int socketSendBuffer, int socketReceiveBuffer, - int requestTimeoutMs, + int defaultRequestTimeoutMs, + long connectionSetupTimeoutMs, + long connectionSetupTimeoutMaxMs, + ClientDnsLookup clientDnsLookup, Time time, boolean discoverBrokerVersions, ApiVersions apiVersions, Sensor throttleTimeSensor, LogContext logContext) { - this(null, metadata, selector, clientId, maxInFlightRequestsPerConnection, - reconnectBackoffMs, reconnectBackoffMax, - socketSendBuffer, socketReceiveBuffer, requestTimeoutMs, time, - discoverBrokerVersions, apiVersions, throttleTimeSensor, logContext); + this(null, + metadata, + selector, + clientId, + maxInFlightRequestsPerConnection, + reconnectBackoffMs, + reconnectBackoffMax, + socketSendBuffer, + socketReceiveBuffer, + defaultRequestTimeoutMs, + connectionSetupTimeoutMs, + connectionSetupTimeoutMaxMs, + clientDnsLookup, + time, + discoverBrokerVersions, + apiVersions, + throttleTimeSensor, + logContext); } public NetworkClient(Selectable selector, @@ -156,15 +211,32 @@ public NetworkClient(Selectable selector, long reconnectBackoffMax, int socketSendBuffer, int socketReceiveBuffer, - int requestTimeoutMs, + int defaultRequestTimeoutMs, + long connectionSetupTimeoutMs, + long connectionSetupTimeoutMaxMs, + ClientDnsLookup clientDnsLookup, Time time, boolean discoverBrokerVersions, ApiVersions apiVersions, LogContext logContext) { - this(metadataUpdater, null, selector, clientId, maxInFlightRequestsPerConnection, - reconnectBackoffMs, reconnectBackoffMax, - socketSendBuffer, socketReceiveBuffer, requestTimeoutMs, time, - discoverBrokerVersions, apiVersions, null, logContext); + this(metadataUpdater, + null, + selector, + clientId, + maxInFlightRequestsPerConnection, + reconnectBackoffMs, + reconnectBackoffMax, + socketSendBuffer, + socketReceiveBuffer, + defaultRequestTimeoutMs, + connectionSetupTimeoutMs, + connectionSetupTimeoutMaxMs, + clientDnsLookup, + time, + discoverBrokerVersions, + apiVersions, + null, + logContext); } private NetworkClient(MetadataUpdater metadataUpdater, @@ -176,7 +248,10 @@ private NetworkClient(MetadataUpdater metadataUpdater, long reconnectBackoffMax, int socketSendBuffer, int socketReceiveBuffer, - int requestTimeoutMs, + int defaultRequestTimeoutMs, + long connectionSetupTimeoutMs, + long connectionSetupTimeoutMaxMs, + ClientDnsLookup clientDnsLookup, Time time, boolean discoverBrokerVersions, ApiVersions apiVersions, @@ -196,18 +271,22 @@ private NetworkClient(MetadataUpdater metadataUpdater, this.selector = selector; this.clientId = clientId; this.inFlightRequests = new InFlightRequests(maxInFlightRequestsPerConnection); - this.connectionStates = new ClusterConnectionStates(reconnectBackoffMs, reconnectBackoffMax); + this.connectionStates = new ClusterConnectionStates( + reconnectBackoffMs, reconnectBackoffMax, + connectionSetupTimeoutMs, connectionSetupTimeoutMaxMs, logContext); this.socketSendBuffer = socketSendBuffer; this.socketReceiveBuffer = socketReceiveBuffer; this.correlation = 0; this.randOffset = new Random(); - this.requestTimeoutMs = requestTimeoutMs; + this.defaultRequestTimeoutMs = defaultRequestTimeoutMs; this.reconnectBackoffMs = reconnectBackoffMs; this.time = time; this.discoverBrokerVersions = discoverBrokerVersions; this.apiVersions = apiVersions; this.throttleTimeSensor = throttleTimeSensor; this.log = logContext.logger(NetworkClient.class); + this.clientDnsLookup = clientDnsLookup; + this.state = new AtomicReference<>(State.ACTIVE); } /** @@ -249,24 +328,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)); - } - } + + 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()); + } } } @@ -280,9 +364,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); } @@ -300,6 +383,22 @@ public long connectionDelay(Node node, long now) { return connectionStates.connectionDelay(node.idString(), now); } + // Return the remaining throttling delay in milliseconds if throttling is in progress. Return 0, otherwise. + // This is for testing. + public long throttleDelayMs(Node node, long now) { + return connectionStates.throttleDelayMs(node.idString(), now); + } + + /** + * Return the poll delay in milliseconds based on both connection and throttle delay. + * @param node the connection to check + * @param now the current time in ms + */ + @Override + public long pollDelayMs(Node node, long now) { + return connectionStates.pollDelayMs(node.idString(), now); + } + /** * Check if the connection of the node has failed, based on the connection state. Such connection failure are * usually transient and can be resumed in the next {@link #ready(org.apache.kafka.common.Node, long)} } @@ -336,16 +435,18 @@ public AuthenticationException authenticationException(Node node) { public boolean isReady(Node node, long now) { // if we need to update our metadata now declare all requests unready to make metadata requests first // priority - return !metadataUpdater.isUpdateDue(now) && canSendRequest(node.idString()); + return !metadataUpdater.isUpdateDue(now) && canSendRequest(node.idString(), now); } /** * Are we connected and ready and able to send more requests to the given connection? * * @param node The node + * @param now the current timestamp */ - private boolean canSendRequest(String node) { - return connectionStates.isReady(node) && selector.isChannelReady(node) && inFlightRequests.canSendMore(node); + private boolean canSendRequest(String node, long now) { + return connectionStates.isReady(node, now) && selector.isChannelReady(node) && + inFlightRequests.canSendMore(node); } /** @@ -358,13 +459,14 @@ 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); } private void doSend(ClientRequest clientRequest, boolean isInternalRequest, long now) { + ensureActive(); String nodeId = clientRequest.destination(); if (!isInternalRequest) { // If this request came from outside the NetworkClient, validate @@ -373,7 +475,7 @@ private void doSend(ClientRequest clientRequest, boolean isInternalRequest, long // will be slightly different for some internal requests (for // example, ApiVersionsRequests can be sent prior to being in // READY state.) - if (!canSendRequest(nodeId)) + if (!canSendRequest(nodeId, now)) throw new IllegalStateException("Attempt to send a request to node " + nodeId + " which is not ready."); } AbstractRequest.Builder builder = clientRequest.requestBuilder(); @@ -395,44 +497,39 @@ private void doSend(ClientRequest clientRequest, boolean isInternalRequest, long // The call to build may also throw UnsupportedVersionException, if there are essential // fields that cannot be represented in the chosen version. doSend(clientRequest, isInternalRequest, now, builder.build(version)); - } catch (UnsupportedVersionException e) { + } catch (UnsupportedVersionException unsupportedVersionException) { // If the version is not supported, skip sending the request over the wire. // Instead, simply add it to the local queue of aborted requests. log.debug("Version mismatch when attempting to send {} with correlation id {} to {}", builder, - clientRequest.correlationId(), clientRequest.destination(), e); + clientRequest.correlationId(), clientRequest.destination(), unsupportedVersionException); ClientResponse clientResponse = new ClientResponse(clientRequest.makeHeader(builder.latestAllowedVersion()), clientRequest.callback(), clientRequest.destination(), now, now, - false, e, null); - abortedSends.add(clientResponse); + false, unsupportedVersionException, null, null); + + if (!isInternalRequest) + abortedSends.add(clientResponse); + else if (clientRequest.apiKey() == ApiKeys.METADATA) + metadataUpdater.handleFailedRequest(now, Optional.of(unsupportedVersionException)); } } private void doSend(ClientRequest clientRequest, boolean isInternalRequest, long now, AbstractRequest request) { - String nodeId = clientRequest.destination(); + String destination = clientRequest.destination(); RequestHeader header = clientRequest.makeHeader(request.version()); if (log.isDebugEnabled()) { - int latestClientVersion = clientRequest.apiKey().latestVersion(); - if (header.apiVersion() == latestClientVersion) { - log.trace("Sending {} {} with correlation id {} to node {}", clientRequest.apiKey(), request, - clientRequest.correlationId(), nodeId); - } else { - log.debug("Using older server API v{} to send {} {} with correlation id {} to node {}", - header.apiVersion(), clientRequest.apiKey(), request, clientRequest.correlationId(), nodeId); - } + log.debug("Sending {} request with header {} and timeout {} to node {}: {}", + clientRequest.apiKey(), header, clientRequest.requestTimeoutMs(), destination, request); } - Send send = request.toSend(nodeId, header); + Send send = request.toSend(header); InFlightRequest inFlightRequest = new InFlightRequest( + clientRequest, header, - clientRequest.createdTimeMs(), - clientRequest.destination(), - clientRequest.callback(), - clientRequest.expectResponse(), isInternalRequest, request, send, now); this.inFlightRequests.add(inFlightRequest); - selector.send(inFlightRequest.send); + selector.send(new NetworkSend(clientRequest.destination(), send)); } /** @@ -446,6 +543,8 @@ private void doSend(ClientRequest clientRequest, boolean isInternalRequest, long */ @Override public List poll(long timeout, long now) { + ensureActive(); + if (!abortedSends.isEmpty()) { // If there are aborted sends because of unsupported version exceptions or disconnects, // handle them immediately without waiting for Selector#poll. @@ -457,7 +556,7 @@ public List poll(long timeout, long now) { long metadataTimeout = metadataUpdater.maybeUpdate(now); try { - this.selector.poll(Utils.min(timeout, metadataTimeout, requestTimeoutMs)); + this.selector.poll(Utils.min(timeout, metadataTimeout, defaultRequestTimeoutMs)); } catch (IOException e) { log.error("Unexpected error during I/O", e); } @@ -470,6 +569,7 @@ public List poll(long timeout, long now) { handleDisconnections(responses, updatedNow); handleConnections(); handleInitiateApiVersionRequests(updatedNow); + handleTimedOutConnections(responses, updatedNow); handleTimedOutRequests(responses, updatedNow); completeResponses(responses); @@ -513,8 +613,8 @@ public boolean hasInFlightRequests(String node) { } @Override - public boolean hasReadyNodes() { - return connectionStates.hasReadyNodes(); + public boolean hasReadyNodes(long now) { + return connectionStates.hasReadyNodes(now); } /** @@ -525,69 +625,120 @@ public void wakeup() { this.selector.wakeup(); } + @Override + public void initiateClose() { + if (state.compareAndSet(State.ACTIVE, State.CLOSING)) { + wakeup(); + } + } + + @Override + public boolean active() { + return state.get() == State.ACTIVE; + } + + private void ensureActive() { + if (!active()) + throw new DisconnectException("NetworkClient is no longer active, state is " + state); + } + /** * Close the network client */ @Override public void close() { - this.selector.close(); + state.compareAndSet(State.ACTIVE, State.CLOSING); + if (state.compareAndSet(State.CLOSING, State.CLOSED)) { + this.selector.close(); + this.metadataUpdater.close(); + } else { + log.warn("Attempting to close NetworkClient that has already been closed."); + } } /** * 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. + * connection if all existing connections are in use. If no connection exists, this method will prefer a node + * with least recent connection attempts. 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, or an active + * connection which is being throttled. * * @return The node with the fewest in-flight requests. */ @Override public Node leastLoadedNode(long now) { List nodes = this.metadataUpdater.fetchNodes(); + 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)) { + if (foundCanConnect == null || + this.connectionStates.lastConnectAttemptMs(foundCanConnect.idString()) > + this.connectionStates.lastConnectAttemptMs(node.idString())) { + 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) { - Struct responseStruct = parseStructMaybeUpdateThrottleTimeMetrics(responseBuffer, requestHeader, null, 0); - return AbstractResponse.parseResponse(requestHeader.apiKey(), responseStruct); - } - - private static Struct parseStructMaybeUpdateThrottleTimeMetrics(ByteBuffer responseBuffer, RequestHeader requestHeader, - Sensor throttleTimeSensor, long now) { - ResponseHeader responseHeader = ResponseHeader.parse(responseBuffer); - // 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); - if (throttleTimeSensor != null && responseBody.hasField(CommonFields.THROTTLE_TIME_MS)) - throttleTimeSensor.record(responseBody.get(CommonFields.THROTTLE_TIME_MS), now); - return responseBody; + try { + return AbstractResponse.parseResponse(responseBuffer, requestHeader); + } catch (BufferUnderflowException e) { + throw new SchemaException("Buffer underflow while parsing response for request with header " + requestHeader, e); + } catch (CorrelationIdMismatchException e) { + if (SaslClientAuthenticator.isReserved(requestHeader.correlationId()) + && !SaslClientAuthenticator.isReserved(e.responseCorrelationId())) + throw new SchemaException("The response is unrelated to Sasl request since its correlation id is " + + e.responseCorrelationId() + " and the reserved range for Sasl request is [ " + + SaslClientAuthenticator.MIN_RESERVED_CORRELATION_ID + "," + + SaslClientAuthenticator.MAX_RESERVED_CORRELATION_ID + "]"); + else { + throw e; + } + } } /** @@ -596,38 +747,38 @@ private static Struct parseStructMaybeUpdateThrottleTimeMetrics(ByteBuffer respo * @param responses The list of responses to update * @param nodeId Id of the node to be disconnected * @param now The current time + * @param disconnectState The state of the disconnected channel */ - private void processDisconnection(List responses, String nodeId, long now, ChannelState disconnectState) { + private void processDisconnection(List responses, + String nodeId, + long now, + ChannelState disconnectState) { connectionStates.disconnected(nodeId, now); apiVersions.remove(nodeId); nodesNeedingApiVersionsFetch.remove(nodeId); switch (disconnectState.state()) { case AUTHENTICATION_FAILED: - connectionStates.authenticationFailed(nodeId, now, disconnectState.exception()); - log.error("Connection to node {} failed authentication due to: {}", nodeId, disconnectState.exception().getMessage()); + AuthenticationException exception = disconnectState.exception(); + connectionStates.authenticationFailed(nodeId, now, exception); + log.error("Connection to node {} ({}) failed authentication due to: {}", nodeId, + disconnectState.remoteAddress(), exception.getMessage()); break; case AUTHENTICATE: - // This warning applies to older brokers which dont provide feedback on authentication failures - log.warn("Connection to node {} terminated during authentication. This may indicate " + - "that authentication failed due to invalid credentials.", nodeId); + log.warn("Connection to node {} ({}) terminated during authentication. This may happen " + + "due to any of the following reasons: (1) Authentication failed due to invalid " + + "credentials with brokers older than 1.0.0, (2) Firewall blocking Kafka TLS " + + "traffic (eg it may only allow HTTPS traffic), (3) Transient network issue.", + nodeId, disconnectState.remoteAddress()); break; case NOT_CONNECTED: - log.warn("Connection to node {} could not be established. Broker may not be available.", nodeId); + log.warn("Connection to node {} ({}) could not be established. Broker may not be available.", nodeId, disconnectState.remoteAddress()); break; 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.request, - request.header.correlationId(), nodeId); - if (request.isInternalRequest && request.header.apiKey() == ApiKeys.METADATA) - metadataUpdater.handleDisconnection(request.destination); - else - responses.add(request.disconnected(now)); - } - AuthenticationException authenticationException = connectionStates.authenticationException(nodeId); - if (authenticationException != null) - metadataUpdater.handleAuthenticationFailure(authenticationException); + + cancelInFlightRequests(nodeId, now, responses); + metadataUpdater.handleServerDisconnect(now, nodeId, Optional.ofNullable(disconnectState.exception())); } /** @@ -638,17 +789,13 @@ private void processDisconnection(List responses, String nodeId, * @param now The current time */ private void handleTimedOutRequests(List responses, long now) { - List nodeIds = this.inFlightRequests.getNodesWithTimedOutRequests(now, this.requestTimeoutMs); + List nodeIds = this.inFlightRequests.nodesWithTimedOutRequests(now); for (String nodeId : nodeIds) { // close connection to the node this.selector.close(nodeId); 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) { @@ -656,6 +803,27 @@ private void handleAbortedSends(List responses) { abortedSends.clear(); } + /** + * Handle socket channel connection timeout. The timeout will hit iff a connection + * stays at the ConnectionState.CONNECTING state longer than the timeout value, + * as indicated by ClusterConnectionStates.NodeConnectionState. + * + * @param responses The list of responses to update + * @param now The current time + */ + private void handleTimedOutConnections(List responses, long now) { + List nodes = connectionStates.nodesWithConnectionSetupTimeout(now); + for (String nodeId : nodes) { + this.selector.close(nodeId); + log.debug( + "Disconnecting from node {} due to socket connection setup timeout. " + + "The timeout value is {} ms.", + nodeId, + connectionStates.connectionSetupTimeoutMs(nodeId)); + processDisconnection(responses, nodeId, now, ChannelState.LOCAL_CLOSE); + } + } + /** * Handle any completed request send. In particular if no response is expected consider the request complete. * @@ -664,15 +832,33 @@ private void handleAbortedSends(List responses) { */ private void handleCompletedSends(List responses, long now) { // if no response is expected then when the send is completed, return it - for (Send send : this.selector.completedSends()) { - InFlightRequest request = this.inFlightRequests.lastSent(send.destination()); + for (NetworkSend send : this.selector.completedSends()) { + InFlightRequest request = this.inFlightRequests.lastSent(send.destinationId()); if (!request.expectResponse) { - this.inFlightRequests.completeLastSent(send.destination()); + this.inFlightRequests.completeLastSent(send.destinationId()); responses.add(request.completed(null, now)); } } } + /** + * If a response from a node includes a non-zero throttle delay and client-side throttling has been enabled for + * the connection to the node, throttle the connection for the specified delay. + * + * @param response the response + * @param apiVersion the API version of the response + * @param nodeId the id of the node + * @param now The current time + */ + private void maybeThrottle(AbstractResponse response, short apiVersion, String nodeId, long now) { + int throttleTimeMs = response.throttleTimeMs(); + if (throttleTimeMs > 0 && response.shouldClientThrottle(apiVersion)) { + connectionStates.throttle(nodeId, now + throttleTimeMs); + log.trace("Connection to node {} is throttled for {} ms until timestamp {}", nodeId, throttleTimeMs, + now + throttleTimeMs); + } + } + /** * Handle any completed receives and update the response list with the responses received. * @@ -683,42 +869,55 @@ private void handleCompletedReceives(List responses, long now) { for (NetworkReceive receive : this.selector.completedReceives()) { String source = receive.source(); InFlightRequest req = inFlightRequests.completeNext(source); - Struct responseStruct = parseStructMaybeUpdateThrottleTimeMetrics(receive.payload(), req.header, - throttleTimeSensor, now); - if (log.isTraceEnabled()) { - log.trace("Completed receive from node {} for {} with correlation id {}, received {}", req.destination, - req.header.apiKey(), req.header.correlationId(), responseStruct); + + AbstractResponse response = parseResponse(receive.payload(), req.header); + if (throttleTimeSensor != null) + throttleTimeSensor.record(response.throttleTimeMs(), now); + + if (log.isDebugEnabled()) { + log.debug("Received {} response from node {} for request with header {}: {}", + req.header.apiKey(), req.destination, req.header, response); } - AbstractResponse body = AbstractResponse.parseResponse(req.header.apiKey(), responseStruct); - if (req.isInternalRequest && body instanceof MetadataResponse) - metadataUpdater.handleCompletedMetadataResponse(req.header, now, (MetadataResponse) body); - else if (req.isInternalRequest && body instanceof ApiVersionsResponse) - handleApiVersionsResponse(responses, req, now, (ApiVersionsResponse) body); + + // If the received response includes a throttle delay, throttle the connection. + maybeThrottle(response, req.header.apiVersion(), req.destination, now); + if (req.isInternalRequest && response instanceof MetadataResponse) + metadataUpdater.handleSuccessfulResponse(req.header, now, (MetadataResponse) response); + else if (req.isInternalRequest && response instanceof ApiVersionsResponse) + handleApiVersionsResponse(responses, req, now, (ApiVersionsResponse) response); else - responses.add(req.completed(body, now)); + responses.add(req.completed(response, now)); } } 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); - if (log.isDebugEnabled()) { - log.debug("Recorded API versions for node {}: {}", node, nodeVersionInfo); - } + log.debug("Recorded API versions for node {}: {}", node, nodeVersionInfo); } /** @@ -733,9 +932,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(); } /** @@ -743,7 +939,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. @@ -773,33 +969,27 @@ 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()) - throw new IllegalStateException("Correlation id for response (" + responseHeader.correlationId() - + ") does not match request (" + requestHeader.correlationId() + "), request header: " + requestHeader); - } - /** * Initiate a connection to the given node + * @param node the node to connect to + * @param now current time in epoch milliseconds */ private void initiateConnect(Node node, long now) { String nodeConnectionId = node.idString(); try { - log.debug("Initiating connection to node {}", node); - this.connectionStates.connecting(nodeConnectionId, now); + connectionStates.connecting(nodeConnectionId, now, node.host(), clientDnsLookup); + InetAddress address = connectionStates.currentAddress(nodeConnectionId); + log.debug("Initiating connection to node {} using address {}", node, address); selector.connect(nodeConnectionId, - new InetSocketAddress(node.host(), node.port()), - this.socketSendBuffer, - this.socketReceiveBuffer); + new InetSocketAddress(address, node.port()), + this.socketSendBuffer, + this.socketReceiveBuffer); } catch (IOException e) { - /* attempt failed, we'll try again after the backoff */ + log.warn("Error connecting to node {}", node, e); + // Attempt failed, we'll try again after the backoff connectionStates.disconnected(nodeConnectionId, now); - /* maybe the problem is our metadata, update it */ - metadataUpdater.requestUpdate(); - log.debug("Error connecting to node {}", node, e); + // Notify metadata updater of the connection failure + metadataUpdater.handleServerDisconnect(now, nodeConnectionId, Optional.empty()); } } @@ -808,12 +998,12 @@ class DefaultMetadataUpdater implements MetadataUpdater { /* the current cluster metadata */ private final Metadata metadata; - /* true iff there is a metadata request that has been sent and for which we have not yet received a response */ - private boolean metadataFetchInProgress; + // Defined if there is a request in progress, null otherwise + private InProgressData inProgress; DefaultMetadataUpdater(Metadata metadata) { this.metadata = metadata; - this.metadataFetchInProgress = false; + this.inProgress = null; } @Override @@ -823,14 +1013,18 @@ public List fetchNodes() { @Override public boolean isUpdateDue(long now) { - return !this.metadataFetchInProgress && this.metadata.timeToNextUpdate(now) == 0; + return !hasFetchInProgress() && this.metadata.timeToNextUpdate(now) == 0; + } + + private boolean hasFetchInProgress() { + return inProgress != null; } @Override public long maybeUpdate(long now) { // should we update our metadata? long timeToNextMetadataUpdate = metadata.timeToNextUpdate(now); - long waitForMetadataFetch = this.metadataFetchInProgress ? requestTimeoutMs : 0; + long waitForMetadataFetch = hasFetchInProgress() ? defaultRequestTimeoutMs : 0; long metadataTimeout = Math.max(timeToNextMetadataUpdate, waitForMetadataFetch); if (metadataTimeout > 0) { @@ -849,51 +1043,73 @@ public long maybeUpdate(long now) { } @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); } - metadataFetchInProgress = false; + // 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) { - metadataFetchInProgress = false; - if (metadata.updateRequested()) - metadata.failedUpdate(time.milliseconds(), exception); + public void handleFailedRequest(long now, Optional maybeFatalException) { + maybeFatalException.ifPresent(metadata::fatalError); + metadata.failedUpdate(now); + inProgress = null; } @Override - public void handleCompletedMetadataResponse(RequestHeader requestHeader, long now, MetadataResponse response) { - this.metadataFetchInProgress = false; - Cluster cluster = response.cluster(); - // check if any topics metadata failed to get updated + 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) + .map(partitionMetadata -> new TopicPartition(topicMetadata.topic(), partitionMetadata.partition()))) + .collect(Collectors.toList()); + if (!missingListenerPartitions.isEmpty()) { + int count = missingListenerPartitions.size(); + log.warn("{} partitions have leader brokers without a matching listener, including {}", + count, missingListenerPartitions.subList(0, Math.min(10, count))); + } + + // 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 (cluster.nodes().size() > 0) { - this.metadata.update(cluster, response.unavailableTopics(), now); - } else { + 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(inProgress.requestVersion, response, inProgress.isPartialUpdate, now); } + + inProgress = null; } @Override - public void requestUpdate() { - this.metadata.requestUpdate(); + public void close() { + this.metadata.close(); } /** @@ -914,19 +1130,13 @@ private boolean isAnyNodeConnecting() { private long maybeUpdate(long now, Node node) { String nodeConnectionId = node.idString(); - if (canSendRequest(nodeConnectionId)) { - this.metadataFetchInProgress = true; - MetadataRequest.Builder metadataRequest; - if (metadata.needMetadataForAllTopics()) - metadataRequest = MetadataRequest.Builder.allTopics(); - else - metadataRequest = new MetadataRequest.Builder(new ArrayList<>(metadata.topics()), - metadata.allowAutoTopicCreation()); - - + if (canSendRequest(nodeConnectionId, now)) { + Metadata.MetadataRequestAndVersion requestAndVersion = metadata.newMetadataRequestAndVersion(now); + MetadataRequest.Builder metadataRequest = requestAndVersion.requestBuilder; log.debug("Sending metadata request {} to node {}", metadataRequest, node); sendInternalMetadataRequest(metadataRequest, nodeConnectionId, now); - return requestTimeoutMs; + inProgress = new InProgressData(requestAndVersion.requestVersion, requestAndVersion.isPartialUpdate); + return defaultRequestTimeoutMs; } // If there's any connection establishment underway, wait until it completes. This prevents @@ -939,7 +1149,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; @@ -951,19 +1161,48 @@ private long maybeUpdate(long now, Node node) { return Long.MAX_VALUE; } + public class InProgressData { + public final int requestVersion; + public final boolean isPartialUpdate; + + private InProgressData(int requestVersion, boolean isPartialUpdate) { + this.requestVersion = requestVersion; + this.isPartialUpdate = isPartialUpdate; + } + }; + } @Override - public ClientRequest newClientRequest(String nodeId, AbstractRequest.Builder requestBuilder, long createdTimeMs, + public ClientRequest newClientRequest(String nodeId, + AbstractRequest.Builder requestBuilder, + long createdTimeMs, boolean expectResponse) { - return newClientRequest(nodeId, requestBuilder, createdTimeMs, expectResponse, null); + 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, long createdTimeMs, - boolean expectResponse, RequestCompletionHandler callback) { - return new ClientRequest(nodeId, requestBuilder, correlation++, clientId, createdTimeMs, expectResponse, - callback); + public ClientRequest newClientRequest(String nodeId, + AbstractRequest.Builder requestBuilder, + long createdTimeMs, + boolean expectResponse, + int requestTimeoutMs, + RequestCompletionHandler callback) { + return new ClientRequest(nodeId, requestBuilder, nextCorrelationId(), clientId, createdTimeMs, expectResponse, + requestTimeoutMs, callback); + } + + public boolean discoverBrokerVersions() { + return discoverBrokerVersions; } static class InFlightRequest { @@ -976,8 +1215,28 @@ static class InFlightRequest { final Send send; final long sendTimeMs; final long createdTimeMs; + final long requestTimeoutMs; + + public InFlightRequest(ClientRequest clientRequest, + RequestHeader header, + boolean isInternalRequest, + AbstractRequest request, + Send send, + long sendTimeMs) { + this(header, + clientRequest.requestTimeoutMs(), + clientRequest.createdTimeMs(), + clientRequest.destination(), + clientRequest.callback(), + clientRequest.expectResponse(), + isInternalRequest, + request, + send, + sendTimeMs); + } public InFlightRequest(RequestHeader header, + int requestTimeoutMs, long createdTimeMs, String destination, RequestCompletionHandler callback, @@ -987,6 +1246,8 @@ public InFlightRequest(RequestHeader header, Send send, long sendTimeMs) { this.header = header; + this.requestTimeoutMs = requestTimeoutMs; + this.createdTimeMs = createdTimeMs; this.destination = destination; this.callback = callback; this.expectResponse = expectResponse; @@ -994,15 +1255,16 @@ public InFlightRequest(RequestHeader header, this.request = request; this.send = send; this.sendTimeMs = sendTimeMs; - this.createdTimeMs = createdTimeMs; } public ClientResponse completed(AbstractResponse response, long timeMs) { - return new ClientResponse(header, callback, destination, createdTimeMs, timeMs, false, null, response); + return new ClientResponse(header, callback, destination, createdTimeMs, timeMs, + false, null, null, response); } - public ClientResponse disconnected(long timeMs) { - return new ClientResponse(header, callback, destination, createdTimeMs, timeMs, true, null, null); + public ClientResponse disconnected(long timeMs, AuthenticationException authenticationException) { + return new ClientResponse(header, callback, destination, createdTimeMs, timeMs, + true, null, authenticationException, null); } @Override @@ -1019,7 +1281,4 @@ public String toString() { } } - public boolean discoverBrokerVersions() { - return discoverBrokerVersions; - } } diff --git a/clients/src/main/java/org/apache/kafka/clients/NetworkClientUtils.java b/clients/src/main/java/org/apache/kafka/clients/NetworkClientUtils.java index c4559a48ca84d..c952b82462def 100644 --- a/clients/src/main/java/org/apache/kafka/clients/NetworkClientUtils.java +++ b/clients/src/main/java/org/apache/kafka/clients/NetworkClientUtils.java @@ -18,6 +18,7 @@ package org.apache.kafka.clients; import org.apache.kafka.common.Node; +import org.apache.kafka.common.errors.DisconnectException; import org.apache.kafka.common.utils.Time; import java.io.IOException; @@ -26,7 +27,9 @@ /** * Provides additional utilities for {@link NetworkClient} (e.g. to implement blocking behaviour). */ -public class NetworkClientUtils { +public final class NetworkClientUtils { + + private NetworkClientUtils() {} /** * Checks whether the node is currently connected, first calling `client.poll` to ensure that any pending @@ -81,25 +84,35 @@ public static boolean awaitReady(KafkaClient client, Node node, Time time, long * disconnection happens (which can happen for a number of reasons including a request timeout). * * In case of a disconnection, an `IOException` is thrown. + * If shutdown is initiated on the client during this method, an IOException is thrown. * * This method is useful for implementing blocking behaviour on top of the non-blocking `NetworkClient`, use it with * care. */ public static ClientResponse sendAndReceive(KafkaClient client, ClientRequest request, Time time) throws IOException { - client.send(request, time.milliseconds()); - while (true) { - List responses = client.poll(Long.MAX_VALUE, time.milliseconds()); - for (ClientResponse response : responses) { - if (response.requestHeader().correlationId() == request.correlationId()) { - if (response.wasDisconnected()) { - throw new IOException("Connection to " + response.destination() + " was disconnected before the response was read"); - } - if (response.versionMismatch() != null) { - throw response.versionMismatch(); + try { + client.send(request, time.milliseconds()); + while (client.active()) { + List responses = client.poll(Long.MAX_VALUE, time.milliseconds()); + for (ClientResponse response : responses) { + if (response.requestHeader().correlationId() == request.correlationId()) { + if (response.wasDisconnected()) { + throw new IOException("Connection to " + response.destination() + " was disconnected before the response was read"); + } + if (response.versionMismatch() != null) { + throw response.versionMismatch(); + } + return response; } - return response; } } + throw new IOException("Client was shutdown before response was read"); + } catch (DisconnectException e) { + if (client.active()) + throw e; + else + throw new IOException("Client was shutdown before response was read"); + } } } 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..e2727f5358d6d 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); @@ -58,7 +60,7 @@ public static NodeApiVersions create() { */ public static NodeApiVersions create(Collection overrides) { List apiVersions = new LinkedList<>(overrides); - for (ApiKeys apiKey : ApiKeys.values()) { + for (ApiKeys apiKey : ApiKeys.enabledApis()) { boolean exists = false; for (ApiVersion apiVersion : apiVersions) { if (apiVersion.apiKey == apiKey.id) { @@ -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)) { @@ -140,7 +167,7 @@ public String toString(boolean lineBreaks) { // Also handle the case where some apiKey types are not specified at all in the given ApiVersions, // which may happen when the remote is too old. - for (ApiKeys apiKey : ApiKeys.values()) { + for (ApiKeys apiKey : ApiKeys.enabledApis()) { if (!apiKeysText.containsKey(apiKey.id)) { StringBuilder bld = new StringBuilder(); bld.append(apiKey.name).append("("). diff --git a/clients/src/main/java/org/apache/kafka/clients/RequestCompletionHandler.java b/clients/src/main/java/org/apache/kafka/clients/RequestCompletionHandler.java index 4e08ddf328c3e..add623f72cbd1 100644 --- a/clients/src/main/java/org/apache/kafka/clients/RequestCompletionHandler.java +++ b/clients/src/main/java/org/apache/kafka/clients/RequestCompletionHandler.java @@ -22,6 +22,5 @@ */ public interface RequestCompletionHandler { - public void onComplete(ClientResponse response); - + void onComplete(ClientResponse response); } diff --git a/clients/src/main/java/org/apache/kafka/clients/StaleMetadataException.java b/clients/src/main/java/org/apache/kafka/clients/StaleMetadataException.java new file mode 100644 index 0000000000000..dafc2d5483c4c --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/StaleMetadataException.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 cannot be used. This is often used as a way to trigger a metadata + * update before retrying another operation. + * + * Note: this is not a public API. + */ +public class StaleMetadataException extends InvalidMetadataException { + private static final long serialVersionUID = 1L; + + public StaleMetadataException() {} + + public StaleMetadataException(String message) { + super(message); + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/AbstractOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/AbstractOptions.java index d3085c3f104b3..2312fe4b81dd2 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/AbstractOptions.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/AbstractOptions.java @@ -23,10 +23,10 @@ */ public abstract class AbstractOptions { - private Integer timeoutMs = null; + protected Integer timeoutMs = null; /** - * Set the request timeout in milliseconds for this operation or {@code null} if the default request timeout for the + * Set the timeout in milliseconds for this operation or {@code null} if the default api timeout for the * AdminClient should be used. */ @SuppressWarnings("unchecked") @@ -36,7 +36,7 @@ public T timeoutMs(Integer timeoutMs) { } /** - * The request timeout in milliseconds for this operation or {@code null} if the default request timeout for the + * The timeout in milliseconds for this operation or {@code null} if the default api timeout for the * AdminClient should be used. */ public Integer timeoutMs() { 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..90e9b0b64be68 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/Admin.java @@ -0,0 +1,1451 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF 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.List; +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.clients.consumer.OffsetAndMetadata; +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.errors.FeatureUpdateFailedException; +import org.apache.kafka.common.quota.ClientQuotaAlteration; +import org.apache.kafka.common.quota.ClientQuotaFilter; +import org.apache.kafka.common.requests.LeaveGroupResponse; + +/** + * The administrative client for Kafka, which supports managing and inspecting topics, brokers, configurations and ACLs. + *

+ * Instances returned from the {@code create} methods of this interface are guaranteed to be thread safe. + *

+ * The operations exposed by Admin follow a consistent pattern: + *

    + *
  • Admin instances should be created using {@link Admin#create(Properties)} or {@link Admin#create(Map)}
    • + *
  • Each operation typically has two overloaded methods, one which uses a default set of options and an + * overloaded method where the last parameter is an explicit options object. + *
  • The operation method's first parameter is a {@code Collection} of items to perform + * the operation on. Batching multiple requests into a single call is more efficient and should be + * preferred over multiple calls to the same method. + *
  • The operation methods execute asynchronously. + *
  • Each {@code xxx} operation method returns an {@code XxxResult} class with methods which expose + * {@link org.apache.kafka.common.KafkaFuture} for accessing the result(s) of the operation. + *
  • Typically an {@code all()} method is provided for getting the overall success/failure of the batch and a + * {@code values()} method provided access to each item in a request batch. + * Other methods may also be provided. + *
  • For synchronous behaviour use {@link org.apache.kafka.common.KafkaFuture#get()} + *
+ *

+ * Here is a simple example of using an Admin client instance to create a new topic: + * 

+ * {@code
+ * Properties props = new Properties();
+ * props.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092");
+ *
+ * try (Admin admin = Admin.create(props)) {
+ *   String topicName = "my-topic";
+ *   int partitions = 12;
+ *   short replicationFactor = 3;
+ *   // Create a compacted topic
+ *   CreateTopicsResult result = admin.createTopics(Collections.singleton(
+ *     new NewTopic(topicName, partitions, replicationFactor)
+ *       .configs(Collections.singletonMap(TopicConfig.CLEANUP_POLICY_CONFIG, TopicConfig.CLEANUP_POLICY_COMPACT)));
+ *
+ *   // Call values() to get the result for a specific topic
+ *   KafkaFuture future = result.values().get(topicName);
+ *
+ *   // Call get() to block until the topic creation is complete or has failed
+ *   // if creation failed the ExecutionException wraps the underlying cause.
+ *   future.get();
+ * }
+ * }
+ *
+ * + *

Bootstrap and balancing

+ *

+ * The {@code bootstrap.servers} config in the {@code Map} or {@code Properties} passed + * to {@link Admin#create(Properties)} is only used for discovering the brokers in the cluster, + * which the client will then connect to as needed. + * As such, it is sufficient to include only two or three broker addresses to cope with the possibility of brokers + * being unavailable. + *

+ * Different operations necessitate requests being sent to different nodes in the cluster. For example + * {@link #createTopics(Collection)} communicates with the controller, but {@link #describeTopics(Collection)} + * can talk to any broker. When the recipient does not matter the instance will try to use the broker with the + * fewest outstanding requests. + *

+ * The client will transparently retry certain errors which are usually transient. + * For example if the request for {@code createTopics()} get sent to a node which was not the controller + * the metadata would be refreshed and the request re-sent to the controller. + * + *

Broker Compatibility

+ *

+ * 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 #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 {@link org.apache.kafka.common.errors.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 {@link org.apache.kafka.common.errors.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 {@link CreateTopicsResult} returns + * success for all the brokers to become aware that the topics have been created. + * During this time, {@link #listTopics()} and {@link #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 #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 {@link DeleteTopicsResult} returns + * success for all the brokers to become aware that the topics are gone. + * During this time, {@link #listTopics()} and {@link #describeTopics(Collection)} + * 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 #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 #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 #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 #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 {@code createAcls} or {@code deleteAcls} to be reflected + * in the output of {@code 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 #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 #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 #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 #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 #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 AlterConfigsResult + */ + 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 {@link AlterConfigsResult}: + *

    + *
  • {@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.UnknownTopicOrPartitionException} + * if the Topic doesn't exist.
    • + *
  • {@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 AlterConfigsResult + */ + 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. For detailed result, inspect the returned {@link AlterReplicaLogDirsResult} instance. + *

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

+ * This is a convenience method for {@link #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. For detailed result, inspect the returned {@link AlterReplicaLogDirsResult} instance. + *

+ * 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 #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 #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 #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 {@link 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 partitions. + * @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 {@link 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 {@link 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 {@link 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 {@link 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 #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 #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 #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 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 #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 election 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 {@link 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. See {@link NewPartitionReassignment}. + * @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); + + /** + *

Alters offsets for the specified group. In order to succeed, the group must be empty. + * + *

This is a convenience method for {@link #alterConsumerGroupOffsets(String, Map, AlterConsumerGroupOffsetsOptions)} with default options. + * See the overload for more details. + * + * @param groupId The group for which to alter offsets. + * @param offsets A map of offsets by partition with associated metadata. + * @return The AlterOffsetsResult. + */ + default AlterConsumerGroupOffsetsResult alterConsumerGroupOffsets(String groupId, Map offsets) { + return alterConsumerGroupOffsets(groupId, offsets, new AlterConsumerGroupOffsetsOptions()); + } + + /** + *

Alters offsets for the specified group. In order to succeed, the group must be empty. + * + *

This operation is not transactional so it may succeed for some partitions while fail for others. + * + * @param groupId The group for which to alter offsets. + * @param offsets A map of offsets by partition with associated metadata. Partitions not specified in the map are ignored. + * @param options The options to use when altering the offsets. + * @return The AlterOffsetsResult. + */ + AlterConsumerGroupOffsetsResult alterConsumerGroupOffsets(String groupId, Map offsets, AlterConsumerGroupOffsetsOptions options); + + /** + *

List offset for the specified partitions and OffsetSpec. This operation enables to find + * the beginning offset, end offset as well as the offset matching a timestamp in partitions. + * + *

This is a convenience method for {@link #listOffsets(Map, ListOffsetsOptions)} + * + * @param topicPartitionOffsets The mapping from partition to the OffsetSpec to look up. + * @return The ListOffsetsResult. + */ + default ListOffsetsResult listOffsets(Map topicPartitionOffsets) { + return listOffsets(topicPartitionOffsets, new ListOffsetsOptions()); + } + + /** + *

List offset for the specified partitions. This operation enables to find + * the beginning offset, end offset as well as the offset matching a timestamp in partitions. + * + * @param topicPartitionOffsets The mapping from partition to the OffsetSpec to look up. + * @param options The options to use when retrieving the offsets + * @return The ListOffsetsResult. + */ + ListOffsetsResult listOffsets(Map topicPartitionOffsets, ListOffsetsOptions options); + + /** + * Describes all entities matching the provided filter that have at least one client quota configuration + * value defined. + *

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

+ * This operation is supported by brokers with version 2.6.0 or higher. + * + * @param filter the filter to apply to match entities + * @return the DescribeClientQuotasResult containing the result + */ + default DescribeClientQuotasResult describeClientQuotas(ClientQuotaFilter filter) { + return describeClientQuotas(filter, new DescribeClientQuotasOptions()); + } + + /** + * Describes all entities matching the provided filter that have at least one client quota configuration + * value defined. + *

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

    + *
  • {@link org.apache.kafka.common.errors.ClusterAuthorizationException} + * If the authenticated user didn't have describe access to the cluster.
    • + *
  • {@link org.apache.kafka.common.errors.InvalidRequestException} + * If the request details are invalid. e.g., an invalid entity type was specified.
    • + *
  • {@link org.apache.kafka.common.errors.TimeoutException} + * If the request timed out before the describe could finish.
    • + *
+ *

+ * This operation is supported by brokers with version 2.6.0 or higher. + * + * @param filter the filter to apply to match entities + * @param options the options to use + * @return the DescribeClientQuotasResult containing the result + */ + DescribeClientQuotasResult describeClientQuotas(ClientQuotaFilter filter, DescribeClientQuotasOptions options); + + /** + * Alters client quota configurations with the specified alterations. + *

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

+ * This operation is supported by brokers with version 2.6.0 or higher. + * + * @param entries the alterations to perform + * @return the AlterClientQuotasResult containing the result + */ + default AlterClientQuotasResult alterClientQuotas(Collection entries) { + return alterClientQuotas(entries, new AlterClientQuotasOptions()); + } + + /** + * Alters client quota configurations with the specified alterations. + *

+ * Alterations for a single entity are atomic, but across entities is not guaranteed. The resulting + * per-entity error code should be evaluated to resolve the success or failure of all updates. + *

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

    + *
  • {@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.InvalidRequestException} + * If the request details are invalid. e.g., a configuration key was specified more than once for an entity.
    • + *
  • {@link org.apache.kafka.common.errors.TimeoutException} + * If the request timed out before the alterations could finish. It cannot be guaranteed whether the update + * succeed or not.
    • + *
+ *

+ * This operation is supported by brokers with version 2.6.0 or higher. + * + * @param entries the alterations to perform + * @return the AlterClientQuotasResult containing the result + */ + AlterClientQuotasResult alterClientQuotas(Collection entries, AlterClientQuotasOptions options); + + /** + * Describe all SASL/SCRAM credentials. + * + *

This is a convenience method for {@link #describeUserScramCredentials(List, DescribeUserScramCredentialsOptions)} + * + * @return The DescribeUserScramCredentialsResult. + */ + default DescribeUserScramCredentialsResult describeUserScramCredentials() { + return describeUserScramCredentials(null, new DescribeUserScramCredentialsOptions()); + } + + /** + * Describe SASL/SCRAM credentials for the given users. + * + *

This is a convenience method for {@link #describeUserScramCredentials(List, DescribeUserScramCredentialsOptions)} + * + * @param users the users for which credentials are to be described; all users' credentials are described if null + * or empty. + * @return The DescribeUserScramCredentialsResult. + */ + default DescribeUserScramCredentialsResult describeUserScramCredentials(List users) { + return describeUserScramCredentials(users, new DescribeUserScramCredentialsOptions()); + } + + /** + * Describe SASL/SCRAM credentials. + *

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

    + *
  • {@link org.apache.kafka.common.errors.ClusterAuthorizationException} + * If the authenticated user didn't have describe access to the cluster.
    • + *
  • {@link org.apache.kafka.common.errors.ResourceNotFoundException} + * If the user did not exist/had no SCRAM credentials.
    • + *
  • {@link org.apache.kafka.common.errors.DuplicateResourceException} + * If the user was requested to be described more than once in the original request.
    • + *
  • {@link org.apache.kafka.common.errors.TimeoutException} + * If the request timed out before the describe operation could finish.
    • + *
+ *

+ * This operation is supported by brokers with version 2.7.0 or higher. + * + * @param users the users for which credentials are to be described; all users' credentials are described if null + * or empty. + * @param options The options to use when describing the credentials + * @return The DescribeUserScramCredentialsResult. + */ + DescribeUserScramCredentialsResult describeUserScramCredentials(List users, DescribeUserScramCredentialsOptions options); + + /** + * Alter SASL/SCRAM credentials for the given users. + * + *

This is a convenience method for {@link #alterUserScramCredentials(List, AlterUserScramCredentialsOptions)} + * + * @param alterations the alterations to be applied + * @return The AlterUserScramCredentialsResult. + */ + default AlterUserScramCredentialsResult alterUserScramCredentials(List alterations) { + return alterUserScramCredentials(alterations, new AlterUserScramCredentialsOptions()); + } + + /** + * Alter SASL/SCRAM credentials. + * + *

+ * The following exceptions can be anticipated when calling {@code get()} any of the futures from the + * returned {@link AlterUserScramCredentialsResult}: + *

    + *
  • {@link org.apache.kafka.common.errors.NotControllerException} + * If the request is not sent to the Controller broker.
    • + *
  • {@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.UnsupportedByAuthenticationException} + * If the user authenticated with a delegation token.
    • + *
  • {@link org.apache.kafka.common.errors.UnsupportedSaslMechanismException} + * If the requested SCRAM mechanism is unrecognized or otherwise unsupported.
    • + *
  • {@link org.apache.kafka.common.errors.UnacceptableCredentialException} + * If the username is empty or the requested number of iterations is too small or too large.
    • + *
  • {@link org.apache.kafka.common.errors.TimeoutException} + * If the request timed out before the describe could finish.
    • + *
+ *

+ * This operation is supported by brokers with version 2.7.0 or higher. + * + * @param alterations the alterations to be applied + * @param options The options to use when altering the credentials + * @return The AlterUserScramCredentialsResult. + */ + AlterUserScramCredentialsResult alterUserScramCredentials(List alterations, + AlterUserScramCredentialsOptions options); + /** + * Describes finalized as well as supported features. + *

+ * This is a convenience method for {@link #describeFeatures(DescribeFeaturesOptions)} with default options. + * See the overload for more details. + * + * @return the {@link DescribeFeaturesResult} containing the result + */ + default DescribeFeaturesResult describeFeatures() { + return describeFeatures(new DescribeFeaturesOptions()); + } + + /** + * Describes finalized as well as supported features. The request is issued to any random + * broker. + *

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

    + *
  • {@link org.apache.kafka.common.errors.TimeoutException} + * If the request timed out before the describe operation could finish.
    • + *
+ *

+ * + * @param options the options to use + * @return the {@link DescribeFeaturesResult} containing the result + */ + DescribeFeaturesResult describeFeatures(DescribeFeaturesOptions options); + + /** + * Applies specified updates to finalized features. This operation is not transactional so some + * updates may succeed while the rest may fail. + *

+ * The API takes in a map of finalized feature names to {@link FeatureUpdate} that needs to be + * applied. Each entry in the map specifies the finalized feature to be added or updated or + * deleted, along with the new max feature version level value. This request is issued only to + * the controller since the API is only served by the controller. The return value contains an + * error code for each supplied {@link FeatureUpdate}, and the code indicates if the update + * succeeded or failed in the controller. + *

    + *
  • Downgrade of feature version level is not a regular operation/intent. It is only allowed + * in the controller if the {@link FeatureUpdate} has the allowDowngrade flag set. Setting this + * flag conveys user intent to attempt downgrade of a feature max version level. Note that + * despite the allowDowngrade flag being set, certain downgrades may be rejected by the + * controller if it is deemed impossible.
    • + *
  • Deletion of a finalized feature version is not a regular operation/intent. It could be + * done by setting the allowDowngrade flag to true in the {@link FeatureUpdate}, and, setting + * the max version level to a value less than 1.
    • + *
+ *

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

    + *
  • {@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.InvalidRequestException} + * If the request details are invalid. e.g., a non-existing finalized feature is attempted + * to be deleted or downgraded.
    • + *
  • {@link org.apache.kafka.common.errors.TimeoutException} + * If the request timed out before the updates could finish. It cannot be guaranteed whether + * the updates succeeded or not.
    • + *
  • {@link FeatureUpdateFailedException} + * This means there was an unexpected error encountered when the update was applied on + * the controller. There is no guarantee on whether the update succeeded or failed. The best + * way to find out is to issue a {@link Admin#describeFeatures(DescribeFeaturesOptions)} + * request.
    • + *
+ *

+ * This operation is supported by brokers with version 2.7.0 or higher. + + * @param featureUpdates the map of finalized feature name to {@link FeatureUpdate} + * @param options the options to use + * @return the {@link UpdateFeaturesResult} containing the result + */ + UpdateFeaturesResult updateFeatures(Map featureUpdates, UpdateFeaturesOptions 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 abe666a43870a..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,522 +17,36 @@ package org.apache.kafka.clients.admin; -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.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), 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), 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. - */ - public abstract void close(long duration, TimeUnit unit); - - /** - * 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()); + return (AdminClient) Admin.create(conf); } - - /** - * 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 this method 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 AdminClient#deleteTopics 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. This API is currently only useful if it is used - * before the replica has been created on the broker. It will support moving replicas that have already been created after - * KIP-113 is fully implemented. - * - * 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.0.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. This API is currently only useful if it is used - * before the replica has been created on the broker. It will support moving replicas that have already been created after - * KIP-113 is fully implemented. - * - * 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.0.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()); - } - - /** - *

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); } + 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 b5ca15a7a6fdd..814f6e211ac93 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 @@ -17,17 +17,20 @@ package org.apache.kafka.clients.admin; +import org.apache.kafka.clients.ClientDnsLookup; import org.apache.kafka.clients.CommonClientConfigs; 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 java.util.Map; import java.util.Set; import static org.apache.kafka.common.config.ConfigDef.Range.atLeast; +import static org.apache.kafka.common.config.ConfigDef.Range.between; import static org.apache.kafka.common.config.ConfigDef.ValidString.in; /** @@ -42,6 +45,12 @@ public class AdminClientConfig extends AbstractConfig { public static final String BOOTSTRAP_SERVERS_CONFIG = CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG; private static final String BOOTSTRAP_SERVERS_DOC = CommonClientConfigs.BOOTSTRAP_SERVERS_DOC; + /** + * client.dns.lookup + */ + public static final String CLIENT_DNS_LOOKUP_CONFIG = CommonClientConfigs.CLIENT_DNS_LOOKUP_CONFIG; + private static final String CLIENT_DNS_LOOKUP_DOC = CommonClientConfigs.CLIENT_DNS_LOOKUP_DOC; + /** * reconnect.backoff.ms */ @@ -62,6 +71,12 @@ public class AdminClientConfig extends AbstractConfig { "retry a failed request. This avoids repeatedly sending requests in a tight loop under " + "some failure scenarios."; + /** socket.connection.setup.timeout.ms */ + public static final String SOCKET_CONNECTION_SETUP_TIMEOUT_MS_CONFIG = CommonClientConfigs.SOCKET_CONNECTION_SETUP_TIMEOUT_MS_CONFIG; + + /** socket.connection.setup.timeout.max.ms */ + public static final String SOCKET_CONNECTION_SETUP_TIMEOUT_MAX_MS_CONFIG = CommonClientConfigs.SOCKET_CONNECTION_SETUP_TIMEOUT_MAX_MS_CONFIG; + /** connections.max.idle.ms */ public static final String CONNECTIONS_MAX_IDLE_MS_CONFIG = CommonClientConfigs.CONNECTIONS_MAX_IDLE_MS_CONFIG; private static final String CONNECTIONS_MAX_IDLE_MS_DOC = CommonClientConfigs.CONNECTIONS_MAX_IDLE_MS_DOC; @@ -99,6 +114,13 @@ public class AdminClientConfig extends AbstractConfig { private static final String METRICS_RECORDING_LEVEL_DOC = CommonClientConfigs.METRICS_RECORDING_LEVEL_DOC; public static final String RETRIES_CONFIG = CommonClientConfigs.RETRIES_CONFIG; + public static final String DEFAULT_API_TIMEOUT_MS_CONFIG = CommonClientConfigs.DEFAULT_API_TIMEOUT_MS_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, @@ -129,10 +151,20 @@ public class AdminClientConfig extends AbstractConfig { RETRY_BACKOFF_MS_DOC) .define(REQUEST_TIMEOUT_MS_CONFIG, Type.INT, - 120000, + 30000, atLeast(0), Importance.MEDIUM, REQUEST_TIMEOUT_MS_DOC) + .define(SOCKET_CONNECTION_SETUP_TIMEOUT_MS_CONFIG, + Type.LONG, + CommonClientConfigs.DEFAULT_SOCKET_CONNECTION_SETUP_TIMEOUT_MS, + Importance.MEDIUM, + CommonClientConfigs.SOCKET_CONNECTION_SETUP_TIMEOUT_MS_DOC) + .define(SOCKET_CONNECTION_SETUP_TIMEOUT_MAX_MS_CONFIG, + Type.LONG, + CommonClientConfigs.DEFAULT_SOCKET_CONNECTION_SETUP_TIMEOUT_MAX_MS, + Importance.MEDIUM, + CommonClientConfigs.SOCKET_CONNECTION_SETUP_TIMEOUT_MAX_MS_DOC) .define(CONNECTIONS_MAX_IDLE_MS_CONFIG, Type.LONG, 5 * 60 * 1000, @@ -140,10 +172,16 @@ public class AdminClientConfig extends AbstractConfig { CONNECTIONS_MAX_IDLE_MS_DOC) .define(RETRIES_CONFIG, Type.INT, - 5, - atLeast(0), + Integer.MAX_VALUE, + between(0, Integer.MAX_VALUE), Importance.LOW, CommonClientConfigs.RETRIES_DOC) + .define(DEFAULT_API_TIMEOUT_MS_CONFIG, + Type.INT, + 60000, + atLeast(0), + Importance.MEDIUM, + CommonClientConfigs.DEFAULT_API_TIMEOUT_MS_DOC) .define(METRICS_SAMPLE_WINDOW_MS_CONFIG, Type.LONG, 30000, @@ -155,10 +193,23 @@ public class AdminClientConfig extends AbstractConfig { .define(METRICS_RECORDING_LEVEL_CONFIG, Type.STRING, Sensor.RecordingLevel.INFO.toString(), - in(Sensor.RecordingLevel.INFO.toString(), Sensor.RecordingLevel.DEBUG.toString()), + in(Sensor.RecordingLevel.INFO.toString(), Sensor.RecordingLevel.DEBUG.toString(), Sensor.RecordingLevel.TRACE.toString()), Importance.LOW, METRICS_RECORDING_LEVEL_DOC) + .define(CLIENT_DNS_LOOKUP_CONFIG, + Type.STRING, + ClientDnsLookup.USE_ALL_DNS_IPS.toString(), + in(ClientDnsLookup.DEFAULT.toString(), + ClientDnsLookup.USE_ALL_DNS_IPS.toString(), + ClientDnsLookup.RESOLVE_CANONICAL_BOOTSTRAP_SERVERS_ONLY.toString()), + Importance.MEDIUM, + CLIENT_DNS_LOOKUP_DOC) // security support + .define(SECURITY_PROVIDERS_CONFIG, + Type.STRING, + null, + Importance.LOW, + SECURITY_PROVIDERS_DOC) .define(SECURITY_PROTOCOL_CONFIG, Type.STRING, DEFAULT_SECURITY_PROTOCOL, @@ -170,19 +221,28 @@ public class AdminClientConfig extends AbstractConfig { @Override protected Map postProcessParsedConfig(final Map parsedValues) { + CommonClientConfigs.warnIfDeprecatedDnsLookupValue(this); return CommonClientConfigs.postProcessReconnectBackoffConfigs(this, parsedValues); } public AdminClientConfig(Map props) { - super(CONFIG, props); + this(props, false); + } + + protected AdminClientConfig(Map props, boolean doLog) { + super(CONFIG, props, doLog); } 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(4, config -> "adminclientconfigs_" + config)); } } diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/AlterClientQuotasOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/AlterClientQuotasOptions.java new file mode 100644 index 0000000000000..3cdaa97762a4d --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/AlterClientQuotasOptions.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.clients.admin; + +import org.apache.kafka.common.annotation.InterfaceStability; + +/** + * Options for {@link Admin#alterClientQuotas(Collection, AlterClientQuotasOptions)}. + * + * The API of this class is evolving, see {@link Admin} for details. + */ +@InterfaceStability.Evolving +public class AlterClientQuotasOptions extends AbstractOptions { + + private boolean validateOnly = false; + + /** + * Returns whether the request should be validated without altering the configs. + */ + public boolean validateOnly() { + return this.validateOnly; + } + + /** + * Sets whether the request should be validated without altering the configs. + */ + public AlterClientQuotasOptions validateOnly(boolean validateOnly) { + this.validateOnly = validateOnly; + return this; + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/AlterClientQuotasResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/AlterClientQuotasResult.java new file mode 100644 index 0000000000000..63c6b3e714043 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/AlterClientQuotasResult.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.KafkaFuture; +import org.apache.kafka.common.annotation.InterfaceStability; +import org.apache.kafka.common.quota.ClientQuotaEntity; + +import java.util.Map; + +/** + * The result of the {@link Admin#alterClientQuotas(Collection, AlterClientQuotasOptions)} call. + * + * The API of this class is evolving, see {@link Admin} for details. + */ +@InterfaceStability.Evolving +public class AlterClientQuotasResult { + + private final Map> futures; + + /** + * Maps an entity to its alteration result. + * + * @param futures maps entity to its alteration result + */ + public AlterClientQuotasResult(Map> futures) { + this.futures = futures; + } + + /** + * Returns a map from quota entity to a future which can be used to check the status of the operation. + */ + public Map> values() { + return futures; + } + + /** + * Returns a future which succeeds only if all quota alterations succeed. + */ + public KafkaFuture all() { + return KafkaFuture.allOf(futures.values().toArray(new KafkaFuture[0])); + } +} 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..1131e12dd5c9c --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/AlterConfigOp.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.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 the value of the configuration entry. + */ + SET((byte) 0), + /** + * Revert the configuration entry to the default value (possibly null). + */ + DELETE((byte) 1), + /** + * (For list-type configuration entries only.) Add the specified values to the + * current value of the configuration entry. If the configuration value has not been set, + * adds to the default value. + */ + APPEND((byte) 2), + /** + * (For list-type configuration entries only.) Removes the specified values from the current + * value of the configuration entry. It is legal to remove values that are not currently in the + * configuration entry. Removing all entries from the current configuration value leaves an empty + * list and does NOT revert to the default value of the entry. + */ + 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 7c84e05d479e7..198a4eab62a70 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,15 +22,26 @@ import java.util.Map; /** - * Options for {@link AdminClient#alterConfigs(Map)}. + * Options for {@link Admin#incrementalAlterConfigs(Map)} and {@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 { private boolean validateOnly = false; + /** + * Set the timeout in milliseconds for this operation or {@code null} if the default api timeout for the + * AdminClient should be used. + * + */ + // This method is retained to keep binary compatibility with 0.11 + public AlterConfigsOptions timeoutMs(Integer timeoutMs) { + this.timeoutMs = timeoutMs; + return this; + } + /** * Return true if the request should be validated without altering the configs. */ 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/AlterConsumerGroupOffsetsOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/AlterConsumerGroupOffsetsOptions.java new file mode 100644 index 0000000000000..f630b4847c46c --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/AlterConsumerGroupOffsetsOptions.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 org.apache.kafka.common.annotation.InterfaceStability; + +import java.util.Map; + +/** + * Options for the {@link AdminClient#alterConsumerGroupOffsets(String, Map, AlterConsumerGroupOffsetsOptions)} call. + * + * The API of this class is evolving, see {@link AdminClient} for details. + */ +@InterfaceStability.Evolving +public class AlterConsumerGroupOffsetsOptions extends AbstractOptions { +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/AlterConsumerGroupOffsetsResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/AlterConsumerGroupOffsetsResult.java new file mode 100644 index 0000000000000..38ee14a15e60a --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/AlterConsumerGroupOffsetsResult.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 java.util.List; +import java.util.Map; +import java.util.stream.Collectors; + +import org.apache.kafka.common.KafkaFuture; +import org.apache.kafka.common.KafkaFuture.BaseFunction; +import org.apache.kafka.common.KafkaFuture.BiConsumer; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.annotation.InterfaceStability; +import org.apache.kafka.common.internals.KafkaFutureImpl; +import org.apache.kafka.common.protocol.Errors; + +/** + * The result of the {@link AdminClient#alterConsumerGroupOffsets(String, Map)} call. + * + * The API of this class is evolving, see {@link AdminClient} for details. + */ +@InterfaceStability.Evolving +public class AlterConsumerGroupOffsetsResult { + + private final KafkaFuture> future; + + AlterConsumerGroupOffsetsResult(KafkaFuture> future) { + this.future = future; + } + + /** + * Return a future which can be used to check the result for a given partition. + */ + public KafkaFuture partitionResult(final TopicPartition partition) { + final KafkaFutureImpl result = new KafkaFutureImpl<>(); + + this.future.whenComplete(new BiConsumer, Throwable>() { + @Override + public void accept(final Map topicPartitions, final Throwable throwable) { + if (throwable != null) { + result.completeExceptionally(throwable); + } else if (!topicPartitions.containsKey(partition)) { + result.completeExceptionally(new IllegalArgumentException( + "Alter offset for partition \"" + partition + "\" was not attempted")); + } else { + final Errors error = topicPartitions.get(partition); + if (error == Errors.NONE) { + result.complete(null); + } else { + result.completeExceptionally(error.exception()); + } + } + + } + }); + + return result; + } + + /** + * Return a future which succeeds if all the alter offsets succeed. + */ + public KafkaFuture all() { + return this.future.thenApply(new BaseFunction, Void>() { + @Override + public Void apply(final Map topicPartitionErrorsMap) { + List partitionsFailed = topicPartitionErrorsMap.entrySet() + .stream() + .filter(e -> e.getValue() != Errors.NONE) + .map(Map.Entry::getKey) + .collect(Collectors.toList()); + for (Errors error : topicPartitionErrorsMap.values()) { + if (error != Errors.NONE) { + throw error.exception( + "Failed altering consumer group offsets for the following partitions: " + partitionsFailed); + } + } + return null; + } + }); + } +} 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..81eb0ea73bc4f 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 @@ -16,14 +16,24 @@ */ package org.apache.kafka.clients.admin; +import java.util.Map; +import java.util.concurrent.CancellationException; +import java.util.concurrent.ExecutionException; import org.apache.kafka.common.KafkaFuture; import org.apache.kafka.common.TopicPartitionReplica; import org.apache.kafka.common.annotation.InterfaceStability; -import java.util.Map; - +import org.apache.kafka.common.errors.ClusterAuthorizationException; +import org.apache.kafka.common.errors.InvalidTopicException; +import org.apache.kafka.common.errors.KafkaStorageException; +import org.apache.kafka.common.errors.LogDirNotFoundException; +import org.apache.kafka.common.errors.ReplicaNotAvailableException; +import org.apache.kafka.common.errors.UnknownServerException; /** - * The result of {@link AdminClient#alterReplicaLogDirs(Map, AlterReplicaLogDirsOptions)}. + * The result of {@link Admin#alterReplicaLogDirs(Map, AlterReplicaLogDirsOptions)}. + * + * To retrieve the detailed result per specified {@link TopicPartitionReplica}, use {@link #values()}. To retrieve the + * overall result only, use {@link #all()}. */ @InterfaceStability.Evolving public class AlterReplicaLogDirsResult { @@ -34,22 +44,34 @@ public class AlterReplicaLogDirsResult { } /** + * Return a map from {@link TopicPartitionReplica} to {@link KafkaFuture} which holds the status of individual + * replica movement. * - * Return a map from replica to future which can be used to check the status of individual replica movement. - * - * Possible error code: + * To check the result of individual replica movement, call {@link KafkaFuture#get()} from the value contained + * in the returned map. If there is no error, it will return silently; if not, an {@link Exception} will be thrown + * like the following: * - * LOG_DIR_NOT_FOUND (57) - * KAFKA_STORAGE_ERROR (56) - * REPLICA_NOT_AVAILABLE (9) - * UNKNOWN (-1) + *
    + *
  • {@link CancellationException}: The task was canceled.
    • + *
  • {@link InterruptedException}: Interrupted while joining I/O thread.
    • + *
  • {@link ExecutionException}: Execution failed with the following causes:
    • + *
      + *
    • {@link ClusterAuthorizationException}: Authorization failed. (CLUSTER_AUTHORIZATION_FAILED, 31)
      • + *
    • {@link InvalidTopicException}: The specified topic name is too long. (INVALID_TOPIC_EXCEPTION, 17)
      • + *
    • {@link LogDirNotFoundException}: The specified log directory is not found in the broker. (LOG_DIR_NOT_FOUND, 57)
      • + *
    • {@link ReplicaNotAvailableException}: The replica does not exist on the broker. (REPLICA_NOT_AVAILABLE, 9)
      • + *
    • {@link KafkaStorageException}: Disk error occurred. (KAFKA_STORAGE_ERROR, 56)
      • + *
    • {@link UnknownServerException}: Unknown. (UNKNOWN_SERVER_ERROR, -1)
      • + *
    + *
*/ public Map> values() { return futures; } /** - * Return a future which succeeds if all the replica movement have succeeded + * Return a {@link KafkaFuture} which succeeds on {@link KafkaFuture#get()} if all the replica movement have succeeded. + * if not, it throws an {@link Exception} described in {@link #values()} method. */ public KafkaFuture all() { return KafkaFuture.allOf(futures.values().toArray(new KafkaFuture[0])); diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/AlterUserScramCredentialsOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/AlterUserScramCredentialsOptions.java new file mode 100644 index 0000000000000..23a0b0a4338fb --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/AlterUserScramCredentialsOptions.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.List; + +/** + * Options for {@link AdminClient#alterUserScramCredentials(List, AlterUserScramCredentialsOptions)} + * + * The API of this class is evolving. See {@link AdminClient} for details. + */ +@InterfaceStability.Evolving +public class AlterUserScramCredentialsOptions extends AbstractOptions { +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/AlterUserScramCredentialsResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/AlterUserScramCredentialsResult.java new file mode 100644 index 0000000000000..a0ce013e4277b --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/AlterUserScramCredentialsResult.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.clients.admin; + +import org.apache.kafka.common.KafkaFuture; +import org.apache.kafka.common.annotation.InterfaceStability; + +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.Objects; + +/** + * The result of the {@link Admin#alterUserScramCredentials(List)} call. + * + * The API of this class is evolving, see {@link Admin} for details. + */ +@InterfaceStability.Evolving +public class AlterUserScramCredentialsResult { + private final Map> futures; + + /** + * + * @param futures the required map from user names to futures representing the results of the alteration(s) + * for each user + */ + public AlterUserScramCredentialsResult(Map> futures) { + this.futures = Collections.unmodifiableMap(Objects.requireNonNull(futures)); + } + + /** + * Return a map from user names to futures, which can be used to check the status of the alteration(s) + * for each user. + */ + public Map> values() { + return this.futures; + } + + /** + * Return a future which succeeds only if all the user SCRAM credential alterations succeed. + */ + public KafkaFuture all() { + return KafkaFuture.allOf(futures.values().toArray(new KafkaFuture[0])); + } +} 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 56f9c27079295..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 @@ -21,39 +21,40 @@ import java.util.Collection; import java.util.Collections; +import java.util.HashMap; +import java.util.Map; /** * 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 { - private final Collection entries; + private final Map entries = new HashMap<>(); /** * Create a configuration instance with the provided entries. */ public Config(Collection entries) { - this.entries = Collections.unmodifiableCollection(entries); + for (ConfigEntry entry : entries) { + this.entries.put(entry.name(), entry); + } } /** * Configuration entries for a resource. */ public Collection entries() { - return entries; + return Collections.unmodifiableCollection(entries.values()); } /** * Get the configuration entry with the provided name or null if there isn't one. */ public ConfigEntry get(String name) { - for (ConfigEntry entry : entries) - if (entry.name().equals(name)) - return entry; - return null; + return entries.get(name); } @Override @@ -75,6 +76,6 @@ public int hashCode() { @Override public String toString() { - return "Config(entries=" + entries + ")"; + return "Config(entries=" + entries.values() + ")"; } } 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 90ba5d4ed7291..b6d947f1b577a 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 @@ -19,21 +19,26 @@ import org.apache.kafka.common.annotation.InterfaceStability; +import java.util.Collections; +import java.util.List; import java.util.Objects; /** * 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 { private final String name; private final String value; - private final boolean isDefault; + private final ConfigSource source; private final boolean isSensitive; private final boolean isReadOnly; + private final List synonyms; + private final ConfigType type; + private final String documentation; /** * Create a configuration entry with the provided values. @@ -53,14 +58,41 @@ public ConfigEntry(String name, String value) { * @param isDefault whether the config value is the default or if it's been explicitly set * @param isSensitive whether the config value is sensitive, the broker never returns the value if it is sensitive * @param isReadOnly whether the config is read-only and cannot be updated + * @deprecated since 1.1.0. This constructor will be removed in a future release. */ + @Deprecated public ConfigEntry(String name, String value, boolean isDefault, boolean isSensitive, boolean isReadOnly) { + this(name, + value, + isDefault ? ConfigSource.DEFAULT_CONFIG : ConfigSource.UNKNOWN, + isSensitive, + isReadOnly, + Collections.emptyList(), + ConfigType.UNKNOWN, + null); + } + + /** + * Create a configuration with the provided values. + * + * @param name the non-null config name + * @param value the config value or null + * @param source the source of this config entry + * @param isSensitive whether the config value is sensitive, the broker never returns the value if it is sensitive + * @param isReadOnly whether the config is read-only and cannot be updated + * @param synonyms Synonym configs in order of precedence + */ + ConfigEntry(String name, String value, ConfigSource source, boolean isSensitive, boolean isReadOnly, + List synonyms, ConfigType type, String documentation) { Objects.requireNonNull(name, "name should not be null"); this.name = name; this.value = value; - this.isDefault = isDefault; + this.source = source; this.isSensitive = isSensitive; this.isReadOnly = isReadOnly; + this.synonyms = synonyms; + this.type = type; + this.documentation = documentation; } /** @@ -77,11 +109,18 @@ public String value() { return value; } + /** + * Return the source of this configuration entry. + */ + public ConfigSource source() { + return source; + } + /** * Return whether the config value is the default or if it's been explicitly set. */ public boolean isDefault() { - return isDefault; + return source == ConfigSource.DEFAULT_CONFIG; } /** @@ -99,14 +138,164 @@ public boolean isReadOnly() { return isReadOnly; } + /** + * Returns all config values that may be used as the value of this config along with their source, + * in the order of precedence. The list starts with the value returned in this ConfigEntry. + * The list is empty if synonyms were not requested using {@link DescribeConfigsOptions#includeSynonyms(boolean)} + */ + public List synonyms() { + return synonyms; + } + + /** + * Return the config data type. + */ + public ConfigType type() { + return type; + } + + /** + * Return the config documentation. + */ + public String documentation() { + return documentation; + } + + @Override + public boolean equals(Object o) { + if (this == o) + return true; + if (o == null || getClass() != o.getClass()) + return false; + + ConfigEntry that = (ConfigEntry) o; + + return this.name.equals(that.name) && + this.value != null ? this.value.equals(that.value) : that.value == null && + this.isSensitive == that.isSensitive && + this.isReadOnly == that.isReadOnly && + this.source == that.source && + Objects.equals(this.synonyms, that.synonyms); + } + + @Override + public int hashCode() { + final int prime = 31; + int result = 1; + result = prime * result + name.hashCode(); + result = prime * result + ((value == null) ? 0 : value.hashCode()); + result = prime * result + (isSensitive ? 1 : 0); + result = prime * result + (isReadOnly ? 1 : 0); + result = prime * result + source.hashCode(); + result = prime * result + synonyms.hashCode(); + return result; + } + @Override public String toString() { return "ConfigEntry(" + "name=" + name + ", value=" + value + - ", isDefault=" + isDefault + + ", source=" + source + ", isSensitive=" + isSensitive + ", isReadOnly=" + isReadOnly + + ", synonyms=" + synonyms + ")"; } + + /** + * Data type of configuration entry. + */ + public enum ConfigType { + UNKNOWN, + BOOLEAN, + STRING, + INT, + SHORT, + LONG, + DOUBLE, + LIST, + CLASS, + PASSWORD + } + + /** + * Source of configuration entries. + */ + 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) + DEFAULT_CONFIG, // built-in default configuration for configs that have a default value + UNKNOWN // source unknown e.g. in the ConfigEntry used for alter requests where source is not set + } + + /** + * Class representing a configuration synonym of a {@link ConfigEntry}. + */ + public static class ConfigSynonym { + + private final String name; + private final String value; + private final ConfigSource source; + + /** + * Create a configuration synonym with the provided values. + * + * @param name Configuration name (this may be different from the name of the associated {@link ConfigEntry} + * @param value Configuration value + * @param source {@link ConfigSource} of this configuraton + */ + ConfigSynonym(String name, String value, ConfigSource source) { + this.name = name; + this.value = value; + this.source = source; + } + + /** + * Returns the name of this configuration. + */ + public String name() { + return name; + } + + /** + * Returns the value of this configuration, which may be null if the configuration is sensitive. + */ + public String value() { + return value; + } + + /** + * Returns the source of this configuration. + */ + public ConfigSource source() { + return source; + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + + ConfigSynonym that = (ConfigSynonym) o; + return Objects.equals(name, that.name) && Objects.equals(value, that.value) && source == that.source; + } + + @Override + public int hashCode() { + return Objects.hash(name, value, source); + } + + @Override + public String toString() { + return "ConfigSynonym(" + + "name=" + name + + ", value=" + value + + ", source=" + source + + ")"; + } + } } 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 new file mode 100644 index 0000000000000..32bd165817457 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/ConsumerGroupDescription.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.clients.admin; + +import org.apache.kafka.common.ConsumerGroupState; +import org.apache.kafka.common.Node; +import org.apache.kafka.common.acl.AclOperation; +import org.apache.kafka.common.utils.Utils; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.Collections; +import java.util.Objects; +import java.util.Set; + +/** + * A detailed description of a single consumer group in the cluster. + */ +public class ConsumerGroupDescription { + private final String groupId; + private final boolean isSimpleConsumerGroup; + private final Collection members; + private final String partitionAssignor; + private final ConsumerGroupState state; + private final Node coordinator; + private final Set authorizedOperations; + + public ConsumerGroupDescription(String groupId, + boolean isSimpleConsumerGroup, + Collection members, + String partitionAssignor, + ConsumerGroupState state, + 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() : + Collections.unmodifiableList(new ArrayList<>(members)); + this.partitionAssignor = partitionAssignor == null ? "" : partitionAssignor; + this.state = state; + this.coordinator = coordinator; + this.authorizedOperations = authorizedOperations; + } + + @Override + public boolean equals(final Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + final ConsumerGroupDescription that = (ConsumerGroupDescription) o; + return isSimpleConsumerGroup == that.isSimpleConsumerGroup && + Objects.equals(groupId, that.groupId) && + Objects.equals(members, that.members) && + Objects.equals(partitionAssignor, that.partitionAssignor) && + state == that.state && + Objects.equals(coordinator, that.coordinator) && + Objects.equals(authorizedOperations, that.authorizedOperations); + } + + @Override + public int hashCode() { + return Objects.hash(groupId, isSimpleConsumerGroup, members, partitionAssignor, state, coordinator, authorizedOperations); + } + + /** + * The id of the consumer group. + */ + public String groupId() { + return groupId; + } + + /** + * If consumer group is simple or not. + */ + public boolean isSimpleConsumerGroup() { + return isSimpleConsumerGroup; + } + + /** + * A list of the members of the consumer group. + */ + public Collection members() { + return members; + } + + /** + * The consumer group partition assignor. + */ + public String partitionAssignor() { + return partitionAssignor; + } + + /** + * The consumer group state, or UNKNOWN if the state is too new for us to parse. + */ + public ConsumerGroupState state() { + return state; + } + + /** + * The consumer group coordinator, or null if the coordinator is not known. + */ + public Node coordinator() { + return coordinator; + } + + /** + * authorizedOperations for this group, or null if that information is not known. + */ + public Set authorizedOperations() { + return authorizedOperations; + } + + @Override + public String toString() { + return "(groupId=" + groupId + + ", isSimpleConsumerGroup=" + isSimpleConsumerGroup + + ", members=" + Utils.join(members, ",") + + ", partitionAssignor=" + partitionAssignor + + ", state=" + state + + ", coordinator=" + coordinator + + ", authorizedOperations=" + authorizedOperations + + ")"; + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/ConsumerGroupListing.java b/clients/src/main/java/org/apache/kafka/clients/admin/ConsumerGroupListing.java new file mode 100644 index 0000000000000..0abc3e01ca9de --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/ConsumerGroupListing.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.clients.admin; + +import java.util.Objects; +import java.util.Optional; + +import org.apache.kafka.common.ConsumerGroupState; + +/** + * A listing of a consumer group in the cluster. + */ +public class ConsumerGroupListing { + private final String groupId; + private final boolean isSimpleConsumerGroup; + private final Optional state; + + /** + * Create an instance with the specified parameters. + * + * @param groupId Group Id + * @param isSimpleConsumerGroup If consumer group is simple or not. + */ + public ConsumerGroupListing(String groupId, boolean isSimpleConsumerGroup) { + this(groupId, isSimpleConsumerGroup, Optional.empty()); + } + + /** + * Create an instance with the specified parameters. + * + * @param groupId Group Id + * @param isSimpleConsumerGroup If consumer group is simple or not. + * @param state The state of the consumer group + */ + public ConsumerGroupListing(String groupId, boolean isSimpleConsumerGroup, Optional state) { + this.groupId = groupId; + this.isSimpleConsumerGroup = isSimpleConsumerGroup; + this.state = Objects.requireNonNull(state); + } + + /** + * Consumer Group Id + */ + public String groupId() { + return groupId; + } + + /** + * If Consumer Group is simple or not. + */ + public boolean isSimpleConsumerGroup() { + return isSimpleConsumerGroup; + } + + /** + * Consumer Group state + */ + public Optional state() { + return state; + } + + @Override + public String toString() { + return "(" + + "groupId='" + groupId + '\'' + + ", isSimpleConsumerGroup=" + isSimpleConsumerGroup + + ", state=" + state + + ')'; + } + + @Override + public int hashCode() { + return Objects.hash(groupId, isSimpleConsumerGroup, state); + } + + @Override + public boolean equals(Object obj) { + if (this == obj) + return true; + if (obj == null) + return false; + if (getClass() != obj.getClass()) + return false; + ConsumerGroupListing other = (ConsumerGroupListing) obj; + if (groupId == null) { + if (other.groupId != null) + return false; + } else if (!groupId.equals(other.groupId)) + return false; + if (isSimpleConsumerGroup != other.isSimpleConsumerGroup) + return false; + if (state == null) { + if (other.state != null) + return false; + } else if (!state.equals(other.state)) + return false; + return true; + } + +} 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 410f079f839a6..ad4ae74ff26ec 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,11 +22,22 @@ 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 { + /** + * Set the timeout in milliseconds for this operation or {@code null} if the default api timeout for the + * AdminClient should be used. + * + */ + // This method is retained to keep binary compatibility with 0.11 + public CreateAclsOptions timeoutMs(Integer timeoutMs) { + this.timeoutMs = timeoutMs; + return this; + } + } 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 new file mode 100644 index 0000000000000..6a082d499bbb4 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/CreateDelegationTokenOptions.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 org.apache.kafka.clients.admin; + +import java.util.LinkedList; +import java.util.List; + +import org.apache.kafka.common.annotation.InterfaceStability; +import org.apache.kafka.common.security.auth.KafkaPrincipal; + +/** + * Options for {@link Admin#createDelegationToken(CreateDelegationTokenOptions)}. + * + * The API of this class is evolving, see {@link Admin} for details. + */ +@InterfaceStability.Evolving +public class CreateDelegationTokenOptions extends AbstractOptions { + private long maxLifeTimeMs = -1; + private List renewers = new LinkedList<>(); + + public CreateDelegationTokenOptions renewers(List renewers) { + this.renewers = renewers; + return this; + } + + public List renewers() { + return renewers; + } + + public CreateDelegationTokenOptions maxlifeTimeMs(long maxLifeTimeMs) { + this.maxLifeTimeMs = maxLifeTimeMs; + return this; + } + + public long maxlifeTimeMs() { + return maxLifeTimeMs; + } +} 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 new file mode 100644 index 0000000000000..7aa48049226e0 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/CreateDelegationTokenResult.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.annotation.InterfaceStability; +import org.apache.kafka.common.security.token.delegation.DelegationToken; + +/** + * The result of the {@link KafkaAdminClient#createDelegationToken(CreateDelegationTokenOptions)} call. + * + * The API of this class is evolving, see {@link Admin} for details. + */ +@InterfaceStability.Evolving +public class CreateDelegationTokenResult { + private final KafkaFuture delegationToken; + + CreateDelegationTokenResult(KafkaFuture delegationToken) { + this.delegationToken = delegationToken; + } + + /** + * Returns a future which yields a delegation token + */ + public KafkaFuture delegationToken() { + return delegationToken; + } +} 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..5a183bd3dfdf4 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,14 +22,15 @@ 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 { private boolean validateOnly = false; + private boolean retryOnQuotaViolation = true; public CreatePartitionsOptions() { } @@ -48,4 +49,19 @@ public CreatePartitionsOptions validateOnly(boolean validateOnly) { this.validateOnly = validateOnly; return this; } -} \ No newline at end of file + + /** + * Set to true if quota violation should be automatically retried. + */ + public CreatePartitionsOptions retryOnQuotaViolation(boolean retryOnQuotaViolation) { + this.retryOnQuotaViolation = retryOnQuotaViolation; + return this; + } + + /** + * Returns true if quota violation should be automatically retried. + */ + public boolean shouldRetryOnQuotaViolation() { + return retryOnQuotaViolation; + } +} 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 7d4bd9e0bdda1..c897f03ca0806 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,14 +22,26 @@ 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 { private boolean validateOnly = false; + private boolean retryOnQuotaViolation = true; + + /** + * Set the timeout in milliseconds for this operation or {@code null} if the default api timeout for the + * AdminClient should be used. + * + */ + // This method is retained to keep binary compatibility with 0.11 + public CreateTopicsOptions timeoutMs(Integer timeoutMs) { + this.timeoutMs = timeoutMs; + return this; + } /** * Set to true if the request should be validated without creating the topic. @@ -46,4 +58,19 @@ public boolean shouldValidateOnly() { return validateOnly; } + + /** + * Set to true if quota violation should be automatically retried. + */ + public CreateTopicsOptions retryOnQuotaViolation(boolean retryOnQuotaViolation) { + this.retryOnQuotaViolation = retryOnQuotaViolation; + return this; + } + + /** + * Returns true if quota violation should be automatically retried. + */ + public boolean shouldRetryOnQuotaViolation() { + return retryOnQuotaViolation; + } } 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..eac086a34c7bc 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; + + protected 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); + } + + public 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 ca57978964232..7c250e10b3e27 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,11 +22,22 @@ 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 { + /** + * Set the timeout in milliseconds for this operation or {@code null} if the default api timeout for the + * AdminClient should be used. + * + */ + // This method is retained to keep binary compatibility with 0.11 + public DeleteAclsOptions timeoutMs(Integer timeoutMs) { + this.timeoutMs = timeoutMs; + return this; + } + } 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 90bc2970e1307..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 { @@ -101,29 +101,26 @@ public Map> values() { * Note that it if the filters don't match any ACLs, this is not considered an error. */ public KafkaFuture> all() { - return KafkaFuture.allOf(futures.values().toArray(new KafkaFuture[0])).thenApply( - new KafkaFuture.Function>() { - @Override - public Collection apply(Void v) { - List acls = new ArrayList<>(); - for (Map.Entry> entry : futures.entrySet()) { - FilterResults results; - try { - results = entry.getValue().get(); - } catch (Throwable e) { - // This should be unreachable, since the future returned by KafkaFuture#allOf should - // have failed if any Future failed. - throw new KafkaException("DeleteAclsResult#all: internal error", e); - } - for (FilterResult result : results.values()) { - if (result.exception() != null) { - throw result.exception(); - } - acls.add(result.binding()); - } - } - return acls; - } - }); + return KafkaFuture.allOf(futures.values().toArray(new KafkaFuture[0])).thenApply(v -> getAclBindings(futures)); + } + + private List getAclBindings(Map> futures) { + List acls = new ArrayList<>(); + for (KafkaFuture value: futures.values()) { + FilterResults results; + try { + results = value.get(); + } catch (Throwable e) { + // This should be unreachable, since the future returned by KafkaFuture#allOf should + // have failed if any Future failed. + throw new KafkaException("DeleteAclsResult#all: internal error", e); + } + for (FilterResult result : results.values()) { + if (result.exception() != null) + throw result.exception(); + acls.add(result.binding()); + } + } + return acls; } } 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 new file mode 100644 index 0000000000000..081aeab6e750a --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/DeleteConsumerGroupsOptions.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.Collection; + +/** + * Options for the {@link Admin#deleteConsumerGroups(Collection)} call. + * + * 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 new file mode 100644 index 0000000000000..c7d7a5ab672a5 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/DeleteConsumerGroupsResult.java @@ -0,0 +1,52 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.kafka.clients.admin; + +import org.apache.kafka.common.KafkaFuture; +import org.apache.kafka.common.annotation.InterfaceStability; + +import java.util.Collection; +import java.util.Map; + +/** + * The result of the {@link Admin#deleteConsumerGroups(Collection)} call. + * + * The API of this class is evolving, see {@link Admin} for details. + */ +@InterfaceStability.Evolving +public class DeleteConsumerGroupsResult { + private final Map> futures; + + DeleteConsumerGroupsResult(final Map> futures) { + this.futures = futures; + } + + /** + * Return a map from group id to futures which can be used to check the status of + * individual deletions. + */ + public Map> deletedGroups() { + return futures; + } + + /** + * Return a future which succeeds only if all the consumer group deletions succeed. + */ + public KafkaFuture all() { + return KafkaFuture.allOf(futures.values().toArray(new KafkaFuture[0])); + } +} 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 44c5252633ce0..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,16 +24,16 @@ 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 { private final Map> futures; - DeleteRecordsResult(Map> futures) { + public DeleteRecordsResult(Map> futures) { this.futures = futures; } 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 d7c5af3afa42e..2711aff5c9c12 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,11 +22,38 @@ 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 { + private boolean retryOnQuotaViolation = true; + + /** + * Set the timeout in milliseconds for this operation or {@code null} if the default api timeout for the + * AdminClient should be used. + * + */ + // This method is retained to keep binary compatibility with 0.11 + public DeleteTopicsOptions timeoutMs(Integer timeoutMs) { + this.timeoutMs = timeoutMs; + return this; + } + + /** + * Set to true if quota violation should be automatically retried. + */ + public DeleteTopicsOptions retryOnQuotaViolation(boolean retryOnQuotaViolation) { + this.retryOnQuotaViolation = retryOnQuotaViolation; + return this; + } + + /** + * Returns true if quota violation should be automatically retried. + */ + public boolean shouldRetryOnQuotaViolation() { + return retryOnQuotaViolation; + } } 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 097cd19c66349..e44d58473df92 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,11 +21,22 @@ 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 { + /** + * Set the timeout in milliseconds for this operation or {@code null} if the default api timeout for the + * AdminClient should be used. + * + */ + // This method is retained to keep binary compatibility with 0.11 + public DescribeAclsOptions timeoutMs(Integer timeoutMs) { + this.timeoutMs = timeoutMs; + return this; + } + } 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/DescribeClientQuotasOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeClientQuotasOptions.java new file mode 100644 index 0000000000000..14e7e451219e6 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeClientQuotasOptions.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#describeClientQuotas(ClientQuotaFilter, DescribeClientQuotasOptions)}. + * + * The API of this class is evolving, see {@link Admin} for details. + */ +@InterfaceStability.Evolving +public class DescribeClientQuotasOptions extends AbstractOptions { +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/DescribeClientQuotasResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeClientQuotasResult.java new file mode 100644 index 0000000000000..b4855901a03de --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeClientQuotasResult.java @@ -0,0 +1,52 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.clients.admin; + +import org.apache.kafka.common.KafkaFuture; +import org.apache.kafka.common.annotation.InterfaceStability; +import org.apache.kafka.common.quota.ClientQuotaEntity; + +import java.util.Map; + +/** + * The result of the {@link Admin#describeClientQuotas(ClientQuotaFilter, DescribeClientQuotasOptions)} call. + * + * The API of this class is evolving, see {@link Admin} for details. + */ +@InterfaceStability.Evolving +public class DescribeClientQuotasResult { + + private final KafkaFuture>> entities; + + /** + * Maps an entity to its configured quota value(s). Note if no value is defined for a quota + * type for that entity's config, then it is not included in the resulting value map. + * + * @param entities future for the collection of entities that matched the filter + */ + public DescribeClientQuotasResult(KafkaFuture>> entities) { + this.entities = entities; + } + + /** + * Returns a map from quota entity to a future which can be used to check the status of the operation. + */ + public KafkaFuture>> entities() { + return entities; + } +} 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 cb5652bc6f53b..2eac1f055f6e0 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,11 +20,36 @@ 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 timeout in milliseconds for this operation or {@code null} if the default api timeout for the + * AdminClient should be used. + * + */ + // This method is retained to keep binary compatibility with 0.11 + public DescribeClusterOptions timeoutMs(Integer timeoutMs) { + this.timeoutMs = 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 bb37e6bb22306..bfb9c18b02234 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,11 +22,54 @@ 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 { + private boolean includeSynonyms = false; + private boolean includeDocumentation = false; + + /** + * Set the timeout in milliseconds for this operation or {@code null} if the default api timeout for the + * AdminClient should be used. + * + */ + // This method is retained to keep binary compatibility with 0.11 + public DescribeConfigsOptions timeoutMs(Integer timeoutMs) { + this.timeoutMs = timeoutMs; + return this; + } + + /** + * Return true if synonym configs should be returned in the response. + */ + public boolean includeSynonyms() { + return includeSynonyms; + } + + /** + * Return true if config documentation should be returned in the response. + */ + public boolean includeDocumentation() { + return includeDocumentation; + } + + /** + * Set to true if synonym configs should be returned in the response. + */ + public DescribeConfigsOptions includeSynonyms(boolean includeSynonyms) { + this.includeSynonyms = includeSynonyms; + return this; + } + + /** + * Set to true if config documentation should be returned in the response. + */ + public DescribeConfigsOptions includeDocumentation(boolean includeDocumentation) { + this.includeDocumentation = includeDocumentation; + return this; + } } 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 478bf055cfd24..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 { @@ -53,7 +53,7 @@ public Map> values() { */ public KafkaFuture> all() { return KafkaFuture.allOf(futures.values().toArray(new KafkaFuture[0])). - thenApply(new KafkaFuture.Function>() { + thenApply(new KafkaFuture.BaseFunction>() { @Override public Map apply(Void v) { Map configs = new HashMap<>(futures.size()); 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 new file mode 100644 index 0000000000000..70238a8b1d5f3 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeConsumerGroupsOptions.java @@ -0,0 +1,41 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.clients.admin; + +import org.apache.kafka.common.annotation.InterfaceStability; + +import java.util.Collection; + +/** + * Options for {@link Admin#describeConsumerGroups(Collection, DescribeConsumerGroupsOptions)}. + *

+ * The API of this class is evolving, see {@link Admin} for details. + */ +@InterfaceStability.Evolving +public class DescribeConsumerGroupsOptions extends AbstractOptions { + private boolean includeAuthorizedOperations; + + public DescribeConsumerGroupsOptions includeAuthorizedOperations(boolean includeAuthorizedOperations) { + this.includeAuthorizedOperations = includeAuthorizedOperations; + return this; + } + + public boolean includeAuthorizedOperations() { + return includeAuthorizedOperations; + } +} \ No newline at end of file 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 new file mode 100644 index 0000000000000..2eddbba305db7 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeConsumerGroupsResult.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.clients.admin; + +import org.apache.kafka.common.KafkaFuture; +import org.apache.kafka.common.annotation.InterfaceStability; + +import java.util.Collection; +import java.util.HashMap; +import java.util.Map; +import java.util.concurrent.ExecutionException; + + +/** + * The result of the {@link KafkaAdminClient#describeConsumerGroups(Collection, DescribeConsumerGroupsOptions)}} call. + * + * The API of this class is evolving, see {@link Admin} for details. + */ +@InterfaceStability.Evolving +public class DescribeConsumerGroupsResult { + + private final Map> futures; + + public DescribeConsumerGroupsResult(final Map> futures) { + this.futures = futures; + } + + /** + * Return a map from group id to futures which yield group descriptions. + */ + public Map> describedGroups() { + return futures; + } + + /** + * Return a future which yields all ConsumerGroupDescription objects, if all the describes succeed. + */ + public KafkaFuture> all() { + return KafkaFuture.allOf(futures.values().toArray(new KafkaFuture[0])).thenApply( + new KafkaFuture.BaseFunction>() { + @Override + public Map apply(Void v) { + try { + Map descriptions = new HashMap<>(futures.size()); + for (Map.Entry> entry : futures.entrySet()) { + descriptions.put(entry.getKey(), entry.getValue().get()); + } + return descriptions; + } catch (InterruptedException | ExecutionException e) { + // This should be unreachable, since the KafkaFuture#allOf already ensured + // that all of the futures completed successfully. + throw new RuntimeException(e); + } + } + }); + } +} 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 new file mode 100644 index 0000000000000..ef9f105850a5f --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeDelegationTokenOptions.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.clients.admin; + +import java.util.List; + +import org.apache.kafka.common.annotation.InterfaceStability; +import org.apache.kafka.common.security.auth.KafkaPrincipal; + +/** + * Options for {@link Admin#describeDelegationToken(DescribeDelegationTokenOptions)}. + * + * The API of this class is evolving, see {@link Admin} for details. + */ +@InterfaceStability.Evolving +public class DescribeDelegationTokenOptions extends AbstractOptions { + private List owners; + + /** + * if owners is null, all the user owned tokens and tokens where user have Describe permission + * will be returned. + * @param owners + * @return this instance + */ + public DescribeDelegationTokenOptions owners(List owners) { + this.owners = owners; + return this; + } + + public List owners() { + return owners; + } +} 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 new file mode 100644 index 0000000000000..47b2530328199 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeDelegationTokenResult.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.clients.admin; + +import java.util.List; + +import org.apache.kafka.common.KafkaFuture; +import org.apache.kafka.common.annotation.InterfaceStability; +import org.apache.kafka.common.security.token.delegation.DelegationToken; + +/** + * The result of the {@link KafkaAdminClient#describeDelegationToken(DescribeDelegationTokenOptions)} call. + * + * The API of this class is evolving, see {@link Admin} for details. + */ +@InterfaceStability.Evolving +public class DescribeDelegationTokenResult { + private final KafkaFuture> delegationTokens; + + DescribeDelegationTokenResult(KafkaFuture> delegationTokens) { + this.delegationTokens = delegationTokens; + } + + /** + * Returns a future which yields list of delegation tokens + */ + public KafkaFuture> delegationTokens() { + return delegationTokens; + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/DescribeFeaturesOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeFeaturesOptions.java new file mode 100644 index 0000000000000..a51ca74eb6422 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeFeaturesOptions.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.clients.admin; + +import org.apache.kafka.common.annotation.InterfaceStability; + +/** + * Options for {@link AdminClient#describeFeatures(DescribeFeaturesOptions)}. + * + * The API of this class is evolving. See {@link Admin} for details. + */ +@InterfaceStability.Evolving +public class DescribeFeaturesOptions extends AbstractOptions { +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/DescribeFeaturesResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeFeaturesResult.java new file mode 100644 index 0000000000000..c48dc19143077 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeFeaturesResult.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.clients.admin; + +import org.apache.kafka.common.KafkaFuture; + +/** + * The result of the {@link Admin#describeFeatures(DescribeFeaturesOptions)} call. + * + * The API of this class is evolving, see {@link Admin} for details. + */ +public class DescribeFeaturesResult { + + private final KafkaFuture future; + + DescribeFeaturesResult(KafkaFuture future) { + this.future = future; + } + + public KafkaFuture featureMetadata() { + return future; + } +} 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 de186fd751df9..96a81f08b02f2 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 @@ -23,48 +23,96 @@ import java.util.Collection; import java.util.Map; import java.util.concurrent.ExecutionException; -import org.apache.kafka.common.requests.DescribeLogDirsResponse.LogDirInfo; +import java.util.stream.Collectors; +import java.util.stream.Stream; + +import org.apache.kafka.common.protocol.Errors; +import org.apache.kafka.common.requests.DescribeLogDirsResponse; /** - * 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 { - private final Map>> futures; + private final Map>> futures; - DescribeLogDirsResult(Map>> futures) { + DescribeLogDirsResult(Map>> futures) { this.futures = futures; } /** - * Return a map from brokerId to future which can be used to check the information of partitions on each individual broker + * Return a map from brokerId to future which can be used to check the information of partitions on each individual broker. + * @deprecated Deprecated Since Kafka 2.7. Use {@link #descriptions()}. */ - public Map>> values() { + @Deprecated + @SuppressWarnings("deprecation") + public Map>> values() { + return descriptions().entrySet().stream() + .collect(Collectors.toMap( + Map.Entry::getKey, + entry -> entry.getValue().thenApply(map -> convertMapValues(map)))); + } + + @SuppressWarnings("deprecation") + private Map convertMapValues(Map map) { + Stream> stream = map.entrySet().stream(); + return stream.collect(Collectors.toMap( + Map.Entry::getKey, + infoEntry -> { + LogDirDescription logDir = infoEntry.getValue(); + return new DescribeLogDirsResponse.LogDirInfo(logDir.error() == null ? Errors.NONE : Errors.forException(logDir.error()), + logDir.replicaInfos().entrySet().stream().collect(Collectors.toMap( + Map.Entry::getKey, + replicaEntry -> new DescribeLogDirsResponse.ReplicaInfo( + replicaEntry.getValue().size(), + replicaEntry.getValue().offsetLag(), + replicaEntry.getValue().isFuture()) + ))); + })); + } + + /** + * Return a map from brokerId to future which can be used to check the information of partitions on each individual broker. + * The result of the future is a map from broker log directory path to a description of that log directory. + */ + public Map>> descriptions() { return futures; } /** * Return a future which succeeds only if all the brokers have responded without error + * @deprecated Deprecated Since Kafka 2.7. Use {@link #allDescriptions()}. + */ + @Deprecated + @SuppressWarnings("deprecation") + public KafkaFuture>> all() { + return allDescriptions().thenApply(map -> map.entrySet().stream().collect(Collectors.toMap( + entry -> entry.getKey(), + entry -> convertMapValues(entry.getValue()) + ))); + } + + /** + * Return a future which succeeds only if all the brokers have responded without error. + * The result of the future is a map from brokerId to a map from broker log directory path + * to a description of that log directory. */ - public KafkaFuture>> all() { + public KafkaFuture>> allDescriptions() { return KafkaFuture.allOf(futures.values().toArray(new KafkaFuture[0])). - thenApply(new KafkaFuture.Function>>() { - @Override - public Map> apply(Void v) { - Map> descriptions = new HashMap<>(futures.size()); - for (Map.Entry>> entry : futures.entrySet()) { - try { - descriptions.put(entry.getKey(), entry.getValue().get()); - } catch (InterruptedException | ExecutionException e) { - // This should be unreachable, because allOf ensured that all the futures completed successfully. - throw new RuntimeException(e); - } + thenApply(v -> { + Map> descriptions = new HashMap<>(futures.size()); + for (Map.Entry>> entry : futures.entrySet()) { + try { + descriptions.put(entry.getKey(), entry.getValue().get()); + } catch (InterruptedException | ExecutionException e) { + // This should be unreachable, because allOf ensured that all the futures completed successfully. + throw new RuntimeException(e); } - return descriptions; } + return descriptions; }); } } 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 401b4aa7b9d43..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 { @@ -52,7 +52,7 @@ public Map> values() { */ public KafkaFuture> all() { return KafkaFuture.allOf(futures.values().toArray(new KafkaFuture[0])). - thenApply(new KafkaFuture.Function>() { + thenApply(new KafkaFuture.BaseFunction>() { @Override public Map apply(Void v) { Map replicaLogDirInfos = new HashMap<>(); 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 64ead4866907f..299aaea3626dc 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,11 +22,33 @@ 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 timeout in milliseconds for this operation or {@code null} if the default api timeout for the + * AdminClient should be used. + * + */ + // This method is retained to keep binary compatibility with 0.11 + public DescribeTopicsOptions timeoutMs(Integer timeoutMs) { + this.timeoutMs = 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 18f5f9d20cdad..7753984a7bda7 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,13 +28,13 @@ /** * 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 { private final Map> futures; - DescribeTopicsResult(Map> futures) { + protected DescribeTopicsResult(Map> futures) { this.futures = futures; } @@ -51,21 +51,18 @@ public Map> values() { */ public KafkaFuture> all() { return KafkaFuture.allOf(futures.values().toArray(new KafkaFuture[0])). - thenApply(new KafkaFuture.Function>() { - @Override - public Map apply(Void v) { - Map descriptions = new HashMap<>(futures.size()); - for (Map.Entry> entry : futures.entrySet()) { - try { - descriptions.put(entry.getKey(), entry.getValue().get()); - } catch (InterruptedException | ExecutionException e) { - // This should be unreachable, because allOf ensured that all the futures - // completed successfully. - throw new RuntimeException(e); - } + thenApply(v -> { + Map descriptions = new HashMap<>(futures.size()); + for (Map.Entry> entry : futures.entrySet()) { + try { + descriptions.put(entry.getKey(), entry.getValue().get()); + } catch (InterruptedException | ExecutionException e) { + // This should be unreachable, because allOf ensured that all the futures + // completed successfully. + throw new RuntimeException(e); } - return descriptions; } + return descriptions; }); } } diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/DescribeUserScramCredentialsOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeUserScramCredentialsOptions.java new file mode 100644 index 0000000000000..1d1af47288cc3 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeUserScramCredentialsOptions.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.List; + +/** + * Options for {@link AdminClient#describeUserScramCredentials(List, DescribeUserScramCredentialsOptions)} + * + * The API of this class is evolving. See {@link AdminClient} for details. + */ +@InterfaceStability.Evolving +public class DescribeUserScramCredentialsOptions extends AbstractOptions { +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/DescribeUserScramCredentialsResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeUserScramCredentialsResult.java new file mode 100644 index 0000000000000..2eddd7ee28c0c --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/DescribeUserScramCredentialsResult.java @@ -0,0 +1,150 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF 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.annotation.InterfaceStability; +import org.apache.kafka.common.errors.ResourceNotFoundException; +import org.apache.kafka.common.internals.KafkaFutureImpl; +import org.apache.kafka.common.message.DescribeUserScramCredentialsResponseData; +import org.apache.kafka.common.protocol.Errors; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import java.util.stream.Collectors; + +/** + * The result of the {@link Admin#describeUserScramCredentials()} call. + * + * The API of this class is evolving, see {@link Admin} for details. + */ +@InterfaceStability.Evolving +public class DescribeUserScramCredentialsResult { + private final KafkaFuture dataFuture; + + /** + * Package-private constructor + * + * @param dataFuture the future indicating response data from the call + */ + DescribeUserScramCredentialsResult(KafkaFuture dataFuture) { + this.dataFuture = Objects.requireNonNull(dataFuture); + } + + /** + * + * @return a future for the results of all described users with map keys (one per user) being consistent with the + * contents of the list returned by {@link #users()}. The future will complete successfully only if all such user + * descriptions complete successfully. + */ + public KafkaFuture> all() { + final KafkaFutureImpl> retval = new KafkaFutureImpl<>(); + dataFuture.whenComplete((data, throwable) -> { + if (throwable != null) { + retval.completeExceptionally(throwable); + } else { + /* Check to make sure every individual described user succeeded. Note that a successfully described user + * is one that appears with *either* a NONE error code or a RESOURCE_NOT_FOUND error code. The + * RESOURCE_NOT_FOUND means the client explicitly requested a describe of that particular user but it could + * not be described because it does not exist; such a user will not appear as a key in the returned map. + */ + Optional optionalFirstFailedDescribe = + data.results().stream().filter(result -> + result.errorCode() != Errors.NONE.code() && result.errorCode() != Errors.RESOURCE_NOT_FOUND.code()).findFirst(); + if (optionalFirstFailedDescribe.isPresent()) { + retval.completeExceptionally(Errors.forCode(optionalFirstFailedDescribe.get().errorCode()).exception(optionalFirstFailedDescribe.get().errorMessage())); + } else { + Map retvalMap = new HashMap<>(); + data.results().stream().forEach(userResult -> + retvalMap.put(userResult.user(), new UserScramCredentialsDescription(userResult.user(), + getScramCredentialInfosFor(userResult)))); + retval.complete(retvalMap); + } + } + }); + return retval; + } + + /** + * + * @return a future indicating the distinct users that meet the request criteria and that have at least one + * credential. The future will not complete successfully if the user is not authorized to perform the describe + * operation; otherwise, it will complete successfully as long as the list of users with credentials can be + * successfully determined within some hard-coded timeout period. Note that the returned list will not include users + * that do not exist/have no credentials: a request to describe an explicit list of users, none of which existed/had + * a credential, will result in a future that returns an empty list being returned here. A returned list will + * include users that have a credential but that could not be described. + */ + public KafkaFuture> users() { + final KafkaFutureImpl> retval = new KafkaFutureImpl<>(); + dataFuture.whenComplete((data, throwable) -> { + if (throwable != null) { + retval.completeExceptionally(throwable); + } else { + retval.complete(data.results().stream() + .filter(result -> result.errorCode() != Errors.RESOURCE_NOT_FOUND.code()) + .map(result -> result.user()).collect(Collectors.toList())); + } + }); + return retval; + } + + /** + * + * @param userName the name of the user description being requested + * @return a future indicating the description results for the given user. The future will complete exceptionally if + * the future returned by {@link #users()} completes exceptionally. Note that if the given user does not exist in + * the list of described users then the returned future will complete exceptionally with + * {@link org.apache.kafka.common.errors.ResourceNotFoundException}. + */ + public KafkaFuture description(String userName) { + final KafkaFutureImpl retval = new KafkaFutureImpl<>(); + dataFuture.whenComplete((data, throwable) -> { + if (throwable != null) { + retval.completeExceptionally(throwable); + } else { + // it is possible that there is no future for this user (for example, the original describe request was + // for users 1, 2, and 3 but this is looking for user 4), so explicitly take care of that case + Optional optionalUserResult = + data.results().stream().filter(result -> result.user().equals(userName)).findFirst(); + if (!optionalUserResult.isPresent()) { + retval.completeExceptionally(new ResourceNotFoundException("No such user: " + userName)); + } else { + DescribeUserScramCredentialsResponseData.DescribeUserScramCredentialsResult userResult = optionalUserResult.get(); + if (userResult.errorCode() != Errors.NONE.code()) { + // RESOURCE_NOT_FOUND is included here + retval.completeExceptionally(Errors.forCode(userResult.errorCode()).exception(userResult.errorMessage())); + } else { + retval.complete(new UserScramCredentialsDescription(userResult.user(), getScramCredentialInfosFor(userResult))); + } + } + } + }); + return retval; + } + + private static List getScramCredentialInfosFor( + DescribeUserScramCredentialsResponseData.DescribeUserScramCredentialsResult userResult) { + return userResult.credentialInfos().stream().map(c -> + new ScramCredentialInfo(ScramMechanism.fromType(c.mechanism()), c.iterations())) + .collect(Collectors.toList()); + } +} 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..ae03ebe2b8453 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/ElectLeadersOptions.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.clients.admin; + +import org.apache.kafka.common.ElectionType; +import org.apache.kafka.common.annotation.InterfaceStability; + +import java.util.Set; + +/** + * 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..186c5848f4a6f --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/ElectLeadersResult.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.clients.admin; + + +import java.util.Map; +import java.util.Optional; +import java.util.Set; + +import org.apache.kafka.common.ElectionType; +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 new file mode 100644 index 0000000000000..e5e385932ac39 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/ElectPreferredLeadersOptions.java @@ -0,0 +1,36 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.clients.admin; + +import org.apache.kafka.common.ElectionType; +import org.apache.kafka.common.annotation.InterfaceStability; + +import java.util.Collection; +import java.util.Set; + +/** + * Options for {@link Admin#electPreferredLeaders(Collection, ElectPreferredLeadersOptions)}. + *

+ * The API of this class is evolving, see {@link Admin} 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 new file mode 100644 index 0000000000000..bf630f3f4db07 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/ElectPreferredLeadersResult.java @@ -0,0 +1,114 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.clients.admin; + + +import java.util.Collection; +import java.util.Map; +import java.util.Optional; +import java.util.Set; + +import org.apache.kafka.common.ElectionType; +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.UnknownTopicOrPartitionException; +import org.apache.kafka.common.internals.KafkaFutureImpl; + +/** + * The result of {@link Admin#electPreferredLeaders(Collection, ElectPreferredLeadersOptions)} + * + * The API of this class is evolving, see {@link Admin} for details. + * + * @deprecated Since 2.4.0. Use {@link Admin#electLeaders(ElectionType, Set, ElectLeadersOptions)}. + */ +@InterfaceStability.Evolving +@Deprecated +public class ElectPreferredLeadersResult { + private final ElectLeadersResult electionResult; + + ElectPreferredLeadersResult(ElectLeadersResult electionResult) { + this.electionResult = electionResult; + } + + /** + * Get the result of the election for the given {@code partition}. + * If there was not an election triggered for the given {@code partition}, the + * returned future will complete with an error. + */ + public KafkaFuture partitionResult(final TopicPartition partition) { + 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 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; + } + + /** + *

Get a future for the topic partitions for which a leader election + * was attempted. A partition will be present in this result if + * an election was attempted even if the election was not successful.

+ * + *

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

+ */ + public KafkaFuture> partitions() { + 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()); + } + } + }); + + return result; + } + + /** + * Return a future which succeeds if all the topic elections succeed. + */ + public KafkaFuture all() { + 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 new file mode 100644 index 0000000000000..3bf9489a05e32 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/ExpireDelegationTokenOptions.java @@ -0,0 +1,39 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.clients.admin; + +import org.apache.kafka.common.annotation.InterfaceStability; + +/** + * Options for {@link Admin#expireDelegationToken(byte[], ExpireDelegationTokenOptions)}. + * + * The API of this class is evolving, see {@link Admin} for details. + */ +@InterfaceStability.Evolving +public class ExpireDelegationTokenOptions extends AbstractOptions { + private long expiryTimePeriodMs = -1L; + + public ExpireDelegationTokenOptions expiryTimePeriodMs(long expiryTimePeriodMs) { + this.expiryTimePeriodMs = expiryTimePeriodMs; + return this; + } + + public long expiryTimePeriodMs() { + return expiryTimePeriodMs; + } +} 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 new file mode 100644 index 0000000000000..59b17140ae791 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/ExpireDelegationTokenResult.java @@ -0,0 +1,42 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.clients.admin; + +import org.apache.kafka.common.KafkaFuture; +import org.apache.kafka.common.annotation.InterfaceStability; + +/** + * The result of the {@link KafkaAdminClient#expireDelegationToken(byte[], ExpireDelegationTokenOptions)} call. + * + * The API of this class is evolving, see {@link Admin} for details. + */ +@InterfaceStability.Evolving +public class ExpireDelegationTokenResult { + private final KafkaFuture expiryTimestamp; + + ExpireDelegationTokenResult(KafkaFuture expiryTimestamp) { + this.expiryTimestamp = expiryTimestamp; + } + + /** + * Returns a future which yields expiry timestamp + */ + public KafkaFuture expiryTimestamp() { + return expiryTimestamp; + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/FeatureMetadata.java b/clients/src/main/java/org/apache/kafka/clients/admin/FeatureMetadata.java new file mode 100644 index 0000000000000..815f9e3b97ca6 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/FeatureMetadata.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.admin; + +import static java.util.stream.Collectors.joining; + +import java.util.HashMap; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; + +/** + * Encapsulates details about finalized as well as supported features. This is particularly useful + * to hold the result returned by the {@link Admin#describeFeatures(DescribeFeaturesOptions)} API. + */ +public class FeatureMetadata { + + private final Map finalizedFeatures; + + private final Optional finalizedFeaturesEpoch; + + private final Map supportedFeatures; + + FeatureMetadata(final Map finalizedFeatures, + final Optional finalizedFeaturesEpoch, + final Map supportedFeatures) { + this.finalizedFeatures = new HashMap<>(finalizedFeatures); + this.finalizedFeaturesEpoch = finalizedFeaturesEpoch; + this.supportedFeatures = new HashMap<>(supportedFeatures); + } + + /** + * Returns a map of finalized feature versions. Each entry in the map contains a key being a + * feature name and the value being a range of version levels supported by every broker in the + * cluster. + */ + public Map finalizedFeatures() { + return new HashMap<>(finalizedFeatures); + } + + /** + * The epoch for the finalized features. + * If the returned value is empty, it means the finalized features are absent/unavailable. + */ + public Optional finalizedFeaturesEpoch() { + return finalizedFeaturesEpoch; + } + + /** + * Returns a map of supported feature versions. Each entry in the map contains a key being a + * feature name and the value being a range of versions supported by a particular broker in the + * cluster. + */ + public Map supportedFeatures() { + return new HashMap<>(supportedFeatures); + } + + @Override + public boolean equals(Object other) { + if (this == other) { + return true; + } + if (!(other instanceof FeatureMetadata)) { + return false; + } + + final FeatureMetadata that = (FeatureMetadata) other; + return Objects.equals(this.finalizedFeatures, that.finalizedFeatures) && + Objects.equals(this.finalizedFeaturesEpoch, that.finalizedFeaturesEpoch) && + Objects.equals(this.supportedFeatures, that.supportedFeatures); + } + + @Override + public int hashCode() { + return Objects.hash(finalizedFeatures, finalizedFeaturesEpoch, supportedFeatures); + } + + private static String mapToString(final Map featureVersionsMap) { + return String.format( + "{%s}", + featureVersionsMap + .entrySet() + .stream() + .map(entry -> String.format("(%s -> %s)", entry.getKey(), entry.getValue())) + .collect(joining(", ")) + ); + } + + @Override + public String toString() { + return String.format( + "FeatureMetadata{finalizedFeatures:%s, finalizedFeaturesEpoch:%s, supportedFeatures:%s}", + mapToString(finalizedFeatures), + finalizedFeaturesEpoch.map(Object::toString).orElse(""), + mapToString(supportedFeatures)); + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/FeatureUpdate.java b/clients/src/main/java/org/apache/kafka/clients/admin/FeatureUpdate.java new file mode 100644 index 0000000000000..38753af3fe7fc --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/FeatureUpdate.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.clients.admin; + +import java.util.Objects; + +/** + * Encapsulates details about an update to a finalized feature. + */ +public class FeatureUpdate { + private final short maxVersionLevel; + private final boolean allowDowngrade; + + /** + * @param maxVersionLevel the new maximum version level for the finalized feature. + * a value < 1 is special and indicates that the update is intended to + * delete the finalized feature, and should be accompanied by setting + * the allowDowngrade flag to true. + * @param allowDowngrade - true, if this feature update was meant to downgrade the existing + * maximum version level of the finalized feature. + * - false, otherwise. + */ + public FeatureUpdate(final short maxVersionLevel, final boolean allowDowngrade) { + if (maxVersionLevel < 1 && !allowDowngrade) { + throw new IllegalArgumentException(String.format( + "The allowDowngrade flag should be set when the provided maxVersionLevel:%d is < 1.", + maxVersionLevel)); + } + this.maxVersionLevel = maxVersionLevel; + this.allowDowngrade = allowDowngrade; + } + + public short maxVersionLevel() { + return maxVersionLevel; + } + + public boolean allowDowngrade() { + return allowDowngrade; + } + + @Override + public boolean equals(Object other) { + if (this == other) { + return true; + } + + if (!(other instanceof FeatureUpdate)) { + return false; + } + + final FeatureUpdate that = (FeatureUpdate) other; + return this.maxVersionLevel == that.maxVersionLevel && this.allowDowngrade == that.allowDowngrade; + } + + @Override + public int hashCode() { + return Objects.hash(maxVersionLevel, allowDowngrade); + } + + @Override + public String toString() { + return String.format("FeatureUpdate{maxVersionLevel:%d, allowDowngrade:%s}", maxVersionLevel, allowDowngrade); + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/FinalizedVersionRange.java b/clients/src/main/java/org/apache/kafka/clients/admin/FinalizedVersionRange.java new file mode 100644 index 0000000000000..aa0401a8a86eb --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/FinalizedVersionRange.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.clients.admin; + +import java.util.Objects; + +/** + * Represents a range of version levels supported by every broker in a cluster for some feature. + */ +public class FinalizedVersionRange { + private final short minVersionLevel; + + private final short maxVersionLevel; + + /** + * Raises an exception unless the following condition is met: + * minVersionLevel >= 1 and maxVersionLevel >= 1 and maxVersionLevel >= minVersionLevel. + * + * @param minVersionLevel The minimum version level value. + * @param maxVersionLevel The maximum version level value. + * + * @throws IllegalArgumentException Raised when the condition described above is not met. + */ + FinalizedVersionRange(final short minVersionLevel, final short maxVersionLevel) { + if (minVersionLevel < 1 || maxVersionLevel < 1 || maxVersionLevel < minVersionLevel) { + throw new IllegalArgumentException( + String.format( + "Expected minVersionLevel >= 1, maxVersionLevel >= 1 and" + + " maxVersionLevel >= minVersionLevel, but received" + + " minVersionLevel: %d, maxVersionLevel: %d", minVersionLevel, maxVersionLevel)); + } + this.minVersionLevel = minVersionLevel; + this.maxVersionLevel = maxVersionLevel; + } + + public short minVersionLevel() { + return minVersionLevel; + } + + public short maxVersionLevel() { + return maxVersionLevel; + } + + @Override + public boolean equals(Object other) { + if (this == other) { + return true; + } + if (!(other instanceof FinalizedVersionRange)) { + return false; + } + + final FinalizedVersionRange that = (FinalizedVersionRange) other; + return this.minVersionLevel == that.minVersionLevel && + this.maxVersionLevel == that.maxVersionLevel; + } + + @Override + public int hashCode() { + return Objects.hash(minVersionLevel, maxVersionLevel); + } + + @Override + public String toString() { + return String.format( + "FinalizedVersionRange[min_version_level:%d, max_version_level:%d]", + minVersionLevel, + maxVersionLevel); + } +} 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 cf48846d91a4d..4558e86d8b07f 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 @@ -18,18 +18,32 @@ package org.apache.kafka.clients.admin; import org.apache.kafka.clients.ApiVersions; +import org.apache.kafka.clients.ClientDnsLookup; import org.apache.kafka.clients.ClientRequest; import org.apache.kafka.clients.ClientResponse; import org.apache.kafka.clients.ClientUtils; +import org.apache.kafka.clients.CommonClientConfigs; import org.apache.kafka.clients.KafkaClient; -import org.apache.kafka.clients.Metadata; 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.ListOffsetsResult.ListOffsetsResultInfo; +import org.apache.kafka.clients.admin.OffsetSpec.TimestampSpec; +import org.apache.kafka.clients.admin.internals.AdminMetadataManager; +import org.apache.kafka.clients.admin.internals.ConsumerGroupOperationContext; +import org.apache.kafka.clients.admin.internals.MetadataOperationContext; +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.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; +import org.apache.kafka.common.MetricName; import org.apache.kafka.common.Node; import org.apache.kafka.common.PartitionInfo; import org.apache.kafka.common.TopicPartition; @@ -37,67 +51,202 @@ 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.annotation.InterfaceStability; +import org.apache.kafka.common.config.ConfigException; import org.apache.kafka.common.config.ConfigResource; import org.apache.kafka.common.errors.ApiException; import org.apache.kafka.common.errors.AuthenticationException; -import org.apache.kafka.common.errors.BrokerNotAvailableException; import org.apache.kafka.common.errors.DisconnectException; +import org.apache.kafka.common.errors.InvalidGroupIdException; import org.apache.kafka.common.errors.InvalidRequestException; import org.apache.kafka.common.errors.InvalidTopicException; +import org.apache.kafka.common.errors.KafkaStorageException; import org.apache.kafka.common.errors.RetriableException; +import org.apache.kafka.common.errors.ThrottlingQuotaExceededException; import org.apache.kafka.common.errors.TimeoutException; +import org.apache.kafka.common.errors.UnacceptableCredentialException; import org.apache.kafka.common.errors.UnknownServerException; +import org.apache.kafka.common.errors.UnknownTopicOrPartitionException; +import org.apache.kafka.common.errors.UnsupportedSaslMechanismException; 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.AlterReplicaLogDirsRequestData; +import org.apache.kafka.common.message.AlterReplicaLogDirsRequestData.AlterReplicaLogDir; +import org.apache.kafka.common.message.AlterReplicaLogDirsRequestData.AlterReplicaLogDirTopic; +import org.apache.kafka.common.message.AlterReplicaLogDirsResponseData.AlterReplicaLogDirPartitionResult; +import org.apache.kafka.common.message.AlterReplicaLogDirsResponseData.AlterReplicaLogDirTopicResult; +import org.apache.kafka.common.message.AlterUserScramCredentialsRequestData; +import org.apache.kafka.common.message.ApiVersionsResponseData.FinalizedFeatureKey; +import org.apache.kafka.common.message.ApiVersionsResponseData.SupportedFeatureKey; +import org.apache.kafka.common.message.CreateAclsRequestData; +import org.apache.kafka.common.message.CreateAclsRequestData.AclCreation; +import org.apache.kafka.common.message.CreateAclsResponseData.AclCreationResult; +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.CreatePartitionsRequestData; +import org.apache.kafka.common.message.CreatePartitionsRequestData.CreatePartitionsAssignment; +import org.apache.kafka.common.message.CreatePartitionsRequestData.CreatePartitionsTopic; +import org.apache.kafka.common.message.CreatePartitionsRequestData.CreatePartitionsTopicCollection; +import org.apache.kafka.common.message.CreatePartitionsResponseData.CreatePartitionsTopicResult; +import org.apache.kafka.common.message.CreateTopicsRequestData; +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.DeleteAclsRequestData; +import org.apache.kafka.common.message.DeleteAclsRequestData.DeleteAclsFilter; +import org.apache.kafka.common.message.DeleteAclsResponseData; +import org.apache.kafka.common.message.DeleteAclsResponseData.DeleteAclsFilterResult; +import org.apache.kafka.common.message.DeleteAclsResponseData.DeleteAclsMatchingAcl; +import org.apache.kafka.common.message.DeleteGroupsRequestData; +import org.apache.kafka.common.message.DeleteRecordsRequestData; +import org.apache.kafka.common.message.DeleteRecordsRequestData.DeleteRecordsPartition; +import org.apache.kafka.common.message.DeleteRecordsRequestData.DeleteRecordsTopic; +import org.apache.kafka.common.message.DeleteRecordsResponseData; +import org.apache.kafka.common.message.DeleteRecordsResponseData.DeleteRecordsTopicResult; +import org.apache.kafka.common.message.DeleteTopicsRequestData; +import org.apache.kafka.common.message.DeleteTopicsResponseData.DeletableTopicResult; +import org.apache.kafka.common.message.DescribeConfigsRequestData; +import org.apache.kafka.common.message.DescribeConfigsResponseData; +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.DescribeLogDirsRequestData; +import org.apache.kafka.common.message.DescribeLogDirsRequestData.DescribableLogDirTopic; +import org.apache.kafka.common.message.DescribeLogDirsResponseData; +import org.apache.kafka.common.message.DescribeUserScramCredentialsRequestData; +import org.apache.kafka.common.message.DescribeUserScramCredentialsRequestData.UserName; +import org.apache.kafka.common.message.DescribeUserScramCredentialsResponseData; +import org.apache.kafka.common.message.ExpireDelegationTokenRequestData; +import org.apache.kafka.common.message.FindCoordinatorRequestData; +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.ListOffsetRequestData.ListOffsetPartition; +import org.apache.kafka.common.message.ListOffsetRequestData.ListOffsetTopic; +import org.apache.kafka.common.message.ListOffsetResponseData.ListOffsetPartitionResponse; +import org.apache.kafka.common.message.ListOffsetResponseData.ListOffsetTopicResponse; +import org.apache.kafka.common.message.ListPartitionReassignmentsRequestData; +import org.apache.kafka.common.message.MetadataRequestData; +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.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.message.UpdateFeaturesRequestData; +import org.apache.kafka.common.message.UpdateFeaturesResponseData.UpdatableFeatureResult; import org.apache.kafka.common.metrics.JmxReporter; +import org.apache.kafka.common.metrics.KafkaMetricsContext; import org.apache.kafka.common.metrics.MetricConfig; import org.apache.kafka.common.metrics.Metrics; +import org.apache.kafka.common.metrics.MetricsContext; import org.apache.kafka.common.metrics.MetricsReporter; import org.apache.kafka.common.metrics.Sensor; import org.apache.kafka.common.network.ChannelBuilder; import org.apache.kafka.common.network.Selector; import org.apache.kafka.common.protocol.Errors; +import org.apache.kafka.common.quota.ClientQuotaAlteration; +import org.apache.kafka.common.quota.ClientQuotaEntity; +import org.apache.kafka.common.quota.ClientQuotaFilter; import org.apache.kafka.common.requests.AbstractRequest; import org.apache.kafka.common.requests.AbstractResponse; +import org.apache.kafka.common.requests.AlterClientQuotasRequest; +import org.apache.kafka.common.requests.AlterClientQuotasResponse; 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.AlterUserScramCredentialsRequest; +import org.apache.kafka.common.requests.AlterUserScramCredentialsResponse; import org.apache.kafka.common.requests.ApiError; +import org.apache.kafka.common.requests.ApiVersionsRequest; +import org.apache.kafka.common.requests.ApiVersionsResponse; import org.apache.kafka.common.requests.CreateAclsRequest; -import org.apache.kafka.common.requests.CreateAclsRequest.AclCreation; import org.apache.kafka.common.requests.CreateAclsResponse; -import org.apache.kafka.common.requests.CreateAclsResponse.AclCreationResponse; +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.CreateTopicsRequest; import org.apache.kafka.common.requests.CreateTopicsResponse; import org.apache.kafka.common.requests.DeleteAclsRequest; import org.apache.kafka.common.requests.DeleteAclsResponse; -import org.apache.kafka.common.requests.DeleteAclsResponse.AclDeletionResult; -import org.apache.kafka.common.requests.DeleteAclsResponse.AclFilterResponse; +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.DescribeClientQuotasRequest; +import org.apache.kafka.common.requests.DescribeClientQuotasResponse; 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.DescribeGroupsRequest; +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.DescribeUserScramCredentialsRequest; +import org.apache.kafka.common.requests.DescribeUserScramCredentialsResponse; +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.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.ListOffsetRequest; +import org.apache.kafka.common.requests.ListOffsetResponse; +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.Resource; -import org.apache.kafka.common.requests.ResourceType; +import org.apache.kafka.common.requests.OffsetCommitRequest; +import org.apache.kafka.common.requests.OffsetCommitResponse; +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.requests.UpdateFeaturesRequest; +import org.apache.kafka.common.requests.UpdateFeaturesResponse; +import org.apache.kafka.common.security.auth.KafkaPrincipal; +import org.apache.kafka.common.security.scram.internals.ScramFormatter; +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; 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; +import java.nio.ByteBuffer; +import java.security.InvalidKeyException; +import java.security.NoSuchAlgorithmException; +import java.time.Duration; import java.util.ArrayList; import java.util.Arrays; import java.util.Collection; @@ -109,18 +258,36 @@ 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.Function; +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. + *

+ * This class is thread-safe. + *

+ * The API of this class is evolving, see {@link Admin} for details. */ @InterfaceStability.Evolving public class KafkaAdminClient extends AdminClient { @@ -140,12 +307,22 @@ public class KafkaAdminClient extends AdminClient { */ private static final long INVALID_SHUTDOWN_TIME = -1; + /** + * Thread name prefix for admin client network thread + */ + static final String NETWORK_THREAD_PREFIX = "kafka-admin-client-thread"; + private final Logger log; /** * The default timeout to use for an operation. */ - private final int defaultTimeoutMs; + private final int defaultApiTimeoutMs; + + /** + * The timeout to use for a single request. + */ + private final int requestTimeoutMs; /** * The name of this AdminClient instance. @@ -158,9 +335,9 @@ public class KafkaAdminClient extends AdminClient { private final Time time; /** - * The cluster metadata used by the KafkaClient. + * The cluster metadata manager used by the KafkaClient. */ - private final Metadata metadata; + private final AdminMetadataManager metadataManager; /** * The metrics for this KafkaAdminClient. @@ -195,6 +372,8 @@ public class KafkaAdminClient extends AdminClient { private final int maxRetries; + private final long retryBackoffMs; + /** * Get or create a list value from a map. * @@ -205,12 +384,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<>()); } /** @@ -221,9 +395,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)); } /** @@ -267,7 +450,7 @@ static String generateClientId(AdminClientConfig config) { private long calcDeadlineMs(long now, Integer optionTimeoutMs) { if (optionTimeoutMs != null) return now + Math.max(0, optionTimeoutMs); - return now + defaultTimeoutMs; + return now + defaultApiTimeoutMs; } /** @@ -299,24 +482,34 @@ static KafkaAdminClient createInternal(AdminClientConfig config, TimeoutProcesso try { // Since we only request node information, it's safe to pass true for allowAutoTopicCreation (and it // simplifies communication with older brokers) - Metadata metadata = new Metadata(config.getLong(AdminClientConfig.RETRY_BACKOFF_MS_CONFIG), - config.getLong(AdminClientConfig.METADATA_MAX_AGE_CONFIG), true); + AdminMetadataManager metadataManager = new AdminMetadataManager(logContext, + config.getLong(AdminClientConfig.RETRY_BACKOFF_MS_CONFIG), + config.getLong(AdminClientConfig.METADATA_MAX_AGE_CONFIG)); + List addresses = ClientUtils.parseAndValidateAddresses( + config.getList(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG), + config.getString(AdminClientConfig.CLIENT_DNS_LOOKUP_CONFIG)); + metadataManager.update(Cluster.bootstrap(addresses), time.milliseconds()); List reporters = config.getConfiguredInstances(AdminClientConfig.METRIC_REPORTER_CLASSES_CONFIG, - MetricsReporter.class); + MetricsReporter.class, + Collections.singletonMap(AdminClientConfig.CLIENT_ID_CONFIG, clientId)); Map metricTags = Collections.singletonMap("client-id", clientId); MetricConfig metricConfig = new MetricConfig().samples(config.getInt(AdminClientConfig.METRICS_NUM_SAMPLES_CONFIG)) .timeWindow(config.getLong(AdminClientConfig.METRICS_SAMPLE_WINDOW_MS_CONFIG), TimeUnit.MILLISECONDS) .recordLevel(Sensor.RecordingLevel.forName(config.getString(AdminClientConfig.METRICS_RECORDING_LEVEL_CONFIG))) .tags(metricTags); - reporters.add(new JmxReporter(JMX_PREFIX)); - metrics = new Metrics(metricConfig, reporters, time); + JmxReporter jmxReporter = new JmxReporter(); + jmxReporter.configure(config.originals()); + reporters.add(jmxReporter); + MetricsContext metricsContext = new KafkaMetricsContext(JMX_PREFIX, + config.originalsWithPrefix(CommonClientConfigs.METRICS_CONTEXT_PREFIX)); + metrics = new Metrics(metricConfig, reporters, time, metricsContext); String metricGrpPrefix = "admin-client"; - channelBuilder = ClientUtils.createChannelBuilder(config); + channelBuilder = ClientUtils.createChannelBuilder(config, time, logContext); selector = new Selector(config.getLong(AdminClientConfig.CONNECTIONS_MAX_IDLE_MS_CONFIG), metrics, time, metricGrpPrefix, channelBuilder, logContext); networkClient = new NetworkClient( selector, - metadata, + metadataManager.updater(), clientId, 1, config.getLong(AdminClientConfig.RECONNECT_BACKOFF_MS_CONFIG), @@ -324,71 +517,105 @@ static KafkaAdminClient createInternal(AdminClientConfig config, TimeoutProcesso config.getInt(AdminClientConfig.SEND_BUFFER_CONFIG), config.getInt(AdminClientConfig.RECEIVE_BUFFER_CONFIG), (int) TimeUnit.HOURS.toMillis(1), + config.getLong(AdminClientConfig.SOCKET_CONNECTION_SETUP_TIMEOUT_MS_CONFIG), + config.getLong(AdminClientConfig.SOCKET_CONNECTION_SETUP_TIMEOUT_MAX_MS_CONFIG), + ClientDnsLookup.forConfig(config.getString(AdminClientConfig.CLIENT_DNS_LOOKUP_CONFIG)), time, true, apiVersions, logContext); - return new KafkaAdminClient(config, clientId, time, metadata, metrics, networkClient, + return new KafkaAdminClient(config, clientId, time, metadataManager, metrics, networkClient, timeoutProcessorFactory, logContext); } catch (Throwable exc) { closeQuietly(metrics, "Metrics"); closeQuietly(networkClient, "NetworkClient"); closeQuietly(selector, "Selector"); closeQuietly(channelBuilder, "ChannelBuilder"); - throw new KafkaException("Failed create new KafkaAdminClient", exc); + throw new KafkaException("Failed to create new KafkaAdminClient", exc); } } - static KafkaAdminClient createInternal(AdminClientConfig config, KafkaClient client, Metadata metadata, Time time) { + static KafkaAdminClient createInternal(AdminClientConfig config, + AdminMetadataManager metadataManager, + KafkaClient client, + Time time) { Metrics metrics = null; String clientId = generateClientId(config); try { - metrics = new Metrics(new MetricConfig(), new LinkedList(), time); - return new KafkaAdminClient(config, clientId, time, metadata, metrics, client, null, - createLogContext(clientId)); + metrics = new Metrics(new MetricConfig(), new LinkedList<>(), time); + LogContext logContext = createLogContext(clientId); + return new KafkaAdminClient(config, clientId, time, metadataManager, metrics, + client, null, logContext); } catch (Throwable exc) { closeQuietly(metrics, "Metrics"); - throw new KafkaException("Failed create new KafkaAdminClient", exc); + throw new KafkaException("Failed to create new KafkaAdminClient", exc); } } - private static LogContext createLogContext(String clientId) { + static LogContext createLogContext(String clientId) { return new LogContext("[AdminClient clientId=" + clientId + "] "); } - private KafkaAdminClient(AdminClientConfig config, String clientId, Time time, Metadata metadata, - Metrics metrics, KafkaClient client, TimeoutProcessorFactory timeoutProcessorFactory, - LogContext logContext) { - this.defaultTimeoutMs = config.getInt(AdminClientConfig.REQUEST_TIMEOUT_MS_CONFIG); + private KafkaAdminClient(AdminClientConfig config, + String clientId, + Time time, + AdminMetadataManager metadataManager, + Metrics metrics, + KafkaClient client, + TimeoutProcessorFactory timeoutProcessorFactory, + LogContext logContext) { this.clientId = clientId; this.log = logContext.logger(KafkaAdminClient.class); + this.requestTimeoutMs = config.getInt(AdminClientConfig.REQUEST_TIMEOUT_MS_CONFIG); + this.defaultApiTimeoutMs = configureDefaultApiTimeoutMs(config); this.time = time; - this.metadata = metadata; - List addresses = ClientUtils.parseAndValidateAddresses( - config.getList(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG)); - this.metadata.update(Cluster.bootstrap(addresses), Collections.emptySet(), time.milliseconds()); + this.metadataManager = metadataManager; this.metrics = metrics; this.client = client; this.runnable = new AdminClientRunnable(); - String threadName = "kafka-admin-client-thread | " + clientId; + String threadName = NETWORK_THREAD_PREFIX + " | " + clientId; this.thread = new KafkaThread(threadName, runnable, true); this.timeoutProcessorFactory = (timeoutProcessorFactory == null) ? new TimeoutProcessorFactory() : timeoutProcessorFactory; 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(); } - Time time() { - return time; + /** + * If a default.api.timeout.ms has been explicitly specified, raise an error if it conflicts with request.timeout.ms. + * If no default.api.timeout.ms has been configured, then set its value as the max of the default and request.timeout.ms. Also we should probably log a warning. + * Otherwise, use the provided values for both configurations. + * + * @param config The configuration + */ + private int configureDefaultApiTimeoutMs(AdminClientConfig config) { + int requestTimeoutMs = config.getInt(AdminClientConfig.REQUEST_TIMEOUT_MS_CONFIG); + int defaultApiTimeoutMs = config.getInt(AdminClientConfig.DEFAULT_API_TIMEOUT_MS_CONFIG); + + if (defaultApiTimeoutMs < requestTimeoutMs) { + if (config.originals().containsKey(AdminClientConfig.DEFAULT_API_TIMEOUT_MS_CONFIG)) { + throw new ConfigException("The specified value of " + AdminClientConfig.DEFAULT_API_TIMEOUT_MS_CONFIG + + " must be no smaller than the value of " + AdminClientConfig.REQUEST_TIMEOUT_MS_CONFIG + "."); + } else { + log.warn("Overriding the default value for {} ({}) with the explicitly configured request timeout {}", + AdminClientConfig.DEFAULT_API_TIMEOUT_MS_CONFIG, this.defaultApiTimeoutMs, + requestTimeoutMs); + return requestTimeoutMs; + } + } + return defaultApiTimeoutMs; } @Override - public void close(long duration, TimeUnit unit) { - long waitTimeMs = unit.toMillis(duration); + public void close(Duration timeout) { + long waitTimeMs = timeout.toMillis(); + if (waitTimeMs < 0) + throw new IllegalArgumentException("The timeout cannot be negative."); waitTimeMs = Math.min(TimeUnit.DAYS.toMillis(365), waitTimeMs); // Limit the timeout to a year. long now = time.milliseconds(); long newHardShutdownTimeMs = now + waitTimeMs; @@ -415,11 +642,12 @@ public void close(long duration, TimeUnit unit) { log.debug("Waiting for the I/O thread to exit. Hard shutdown in {} ms.", deltaMs); } try { - // Wait for the thread to be joined. - thread.join(); - - AppInfoParser.unregisterAppInfo(JMX_PREFIX, clientId, metrics); - + // close() can be called by AdminClient thread when it invokes callback. That will + // cause deadlock, so check for that condition. + if (Thread.currentThread() != thread) { + // Wait for the thread to be joined. + thread.join(waitTimeMs); + } log.debug("Kafka admin client closed."); } catch (InterruptedException e) { log.debug("Interrupted while joining I/O thread", e); @@ -434,6 +662,13 @@ private interface NodeProvider { Node provide(); } + private class MetadataUpdateNodeIdProvider implements NodeProvider { + @Override + public Node provide() { + return client.leastLoadedNode(time.milliseconds()); + } + } + private class ConstantNodeIdProvider implements NodeProvider { private final int nodeId; @@ -443,7 +678,16 @@ private class ConstantNodeIdProvider implements NodeProvider { @Override public Node provide() { - return metadata.fetch().nodeById(nodeId); + if (metadataManager.isReady() && + (metadataManager.nodeById(nodeId) != null)) { + return metadataManager.nodeById(nodeId); + } + // If we can't find the node with the given constant ID, we schedule a + // metadata update and hope it appears. This behavior is useful for avoiding + // flaky behavior in tests when the cluster is starting up and not all nodes + // have appeared. + metadataManager.requestUpdate(); + return null; } } @@ -453,7 +697,12 @@ public Node provide() { private class ControllerNodeProvider implements NodeProvider { @Override public Node provide() { - return metadata.fetch().controller(); + if (metadataManager.isReady() && + (metadataManager.controller() != null)) { + return metadataManager.controller(); + } + metadataManager.requestUpdate(); + return null; } } @@ -463,23 +712,41 @@ public Node provide() { private class LeastLoadedNodeProvider implements NodeProvider { @Override public Node provide() { - return client.leastLoadedNode(time.milliseconds()); + if (metadataManager.isReady()) { + // This may return null if all nodes are busy. + // In that case, we will postpone node assignment. + return client.leastLoadedNode(time.milliseconds()); + } + metadataManager.requestUpdate(); + return null; } } abstract class Call { + private final boolean internal; private final String callName; private final long deadlineMs; private final NodeProvider nodeProvider; private int tries = 0; private boolean aborted = false; + private Node curNode = null; + private long nextAllowedTryMs = 0; - Call(String callName, long deadlineMs, NodeProvider nodeProvider) { + Call(boolean internal, String callName, long deadlineMs, NodeProvider nodeProvider) { + this.internal = internal; this.callName = callName; this.deadlineMs = deadlineMs; this.nodeProvider = nodeProvider; } + Call(String callName, long deadlineMs, NodeProvider nodeProvider) { + this(false, callName, deadlineMs, nodeProvider); + } + + protected Node curNode() { + return curNode; + } + /** * Handle a failure. * @@ -496,11 +763,7 @@ final void fail(long now, Throwable throwable) { // TimeoutException. In this case, we do not get any more retries - the call has // failed. We increment tries anyway in order to display an accurate log message. tries++; - if (log.isDebugEnabled()) { - log.debug("{} aborted at {} after {} attempt(s)", this, now, tries, - new Exception(prettyPrintException(throwable))); - } - handleFailure(new TimeoutException("Aborted due to timeout.")); + failWithTimeout(now, throwable); return; } // If this is an UnsupportedVersionException that we can retry, do so. Note that a @@ -508,21 +771,19 @@ final void fail(long now, Throwable throwable) { // this RPC. That is why 'tries' is not incremented. if ((throwable instanceof UnsupportedVersionException) && handleUnsupportedVersionException((UnsupportedVersionException) throwable)) { - log.trace("{} attempting protocol downgrade.", this); + log.debug("{} attempting protocol downgrade and then retry.", this); runnable.enqueue(this, now); return; } tries++; + nextAllowedTryMs = now + retryBackoffMs; + // If the call has timed out, fail. if (calcTimeoutMsRemainingAsInt(now, deadlineMs) < 0) { - if (log.isDebugEnabled()) { - log.debug("{} timed out at {} after {} attempt(s)", this, now, tries, - new Exception(prettyPrintException(throwable))); - } - handleFailure(throwable); + failWithTimeout(now, throwable); return; } - // If the exception is not retryable, fail. + // If the exception is not retriable, fail. if (!(throwable instanceof RetriableException)) { if (log.isDebugEnabled()) { log.debug("{} failed with non-retriable exception after {} attempt(s)", this, tries, @@ -533,11 +794,7 @@ final void fail(long now, Throwable throwable) { } // If we are out of retries, fail. if (tries > maxRetries) { - if (log.isDebugEnabled()) { - log.debug("{} failed after {} attempt(s)", this, tries, - new Exception(prettyPrintException(throwable))); - } - handleFailure(throwable); + failWithTimeout(now, throwable); return; } if (log.isDebugEnabled()) { @@ -547,6 +804,15 @@ final void fail(long now, Throwable throwable) { runnable.enqueue(this, now); } + private void failWithTimeout(long now, Throwable cause) { + if (log.isDebugEnabled()) { + log.debug("{} timed out at {} after {} attempt(s)", this, now, tries, + new Exception(prettyPrintException(cause))); + } + handleFailure(new TimeoutException(this + " timed out at " + now + + " after " + tries + " attempt(s)", cause)); + } + /** * Create an AbstractRequest.Builder for this Call. * @@ -554,6 +820,7 @@ final void fail(long now, Throwable throwable) { * * @return The AbstractRequest builder. */ + @SuppressWarnings("rawtypes") abstract AbstractRequest.Builder createRequest(int timeoutMs); /** @@ -566,7 +833,7 @@ final void fail(long now, Throwable throwable) { /** * Handle a failure. This will only be called if the failure exception was not - * retryable, or if we hit a timeout. + * retriable, or if we hit a timeout. * * @param throwable The exception. */ @@ -585,7 +852,12 @@ boolean handleUnsupportedVersionException(UnsupportedVersionException exception) @Override public String toString() { - return "Call(callName=" + callName + ", deadlineMs=" + deadlineMs + ")"; + return "Call(callName=" + callName + ", deadlineMs=" + deadlineMs + + ", tries=" + tries + ", nextAllowedTryMs=" + nextAllowedTryMs + ")"; + } + + public boolean isInternal() { + return internal; } } @@ -631,7 +903,7 @@ int handleTimeouts(Collection calls, String msg) { Call call = iter.next(); int remainingMs = calcTimeoutMsRemainingAsInt(now, call.deadlineMs); if (remainingMs < 0) { - call.fail(now, new TimeoutException(msg)); + call.fail(now, new TimeoutException(msg + " Call: " + call.callName)); iter.remove(); numTimedOut++; } else { @@ -664,58 +936,52 @@ int nextTimeoutMs() { private final class AdminClientRunnable implements Runnable { /** - * Pending calls. Protected by the object monitor. - * This will be null only if the thread has shut down. + * Calls which have not yet been assigned to a node. + * Only accessed from this thread. */ - private List newCalls = new LinkedList<>(); + private final ArrayList pendingCalls = new ArrayList<>(); /** - * Check if the AdminClient metadata is ready. - * We need to know who the controller is, and have a non-empty view of the cluster. - * - * @param prevMetadataVersion The previous metadata version which wasn't usable. - * @return null if the metadata is usable; the current metadata - * version otherwise + * Maps nodes to calls that we want to send. + * Only accessed from this thread. */ - private Integer checkMetadataReady(Integer prevMetadataVersion) { - if (prevMetadataVersion != null) { - if (prevMetadataVersion == metadata.version()) - return prevMetadataVersion; - } - Cluster cluster = metadata.fetch(); - if (cluster.nodes().isEmpty()) { - log.trace("Metadata is not ready yet. No cluster nodes found."); - return metadata.requestUpdate(); - } - if (cluster.controller() == null) { - log.trace("Metadata is not ready yet. No controller found."); - return metadata.requestUpdate(); - } - if (prevMetadataVersion != null) { - log.trace("Metadata is now ready."); - } - return null; - } + private final Map> callsToSend = new HashMap<>(); + + /** + * Maps node ID strings to calls that have been sent. + * Only accessed from this thread. + */ + private final Map> callsInFlight = new HashMap<>(); + + /** + * Maps correlation IDs to calls that have been sent. + * Only accessed from this thread. + */ + private final Map correlationIdToCalls = new HashMap<>(); + + /** + * Pending calls. Protected by the object monitor. + * This will be null only if the thread has shut down. + */ + private List newCalls = new LinkedList<>(); /** - * Time out the elements in the newCalls list which are expired. + * Time out the elements in the pendingCalls list which are expired. * * @param processor The timeout processor. */ - private synchronized void timeoutNewCalls(TimeoutProcessor processor) { - int numTimedOut = processor.handleTimeouts(newCalls, - "Timed out waiting for a node assignment."); + private void timeoutPendingCalls(TimeoutProcessor processor) { + int numTimedOut = processor.handleTimeouts(pendingCalls, "Timed out waiting for a node assignment."); if (numTimedOut > 0) - log.debug("Timed out {} new calls.", numTimedOut); + log.debug("Timed out {} pending calls.", numTimedOut); } /** * Time out calls which have been assigned to nodes. * * @param processor The timeout processor. - * @param callsToSend A map of nodes to the calls they need to handle. */ - private void timeoutCallsToSend(TimeoutProcessor processor, Map> callsToSend) { + private int timeoutCallsToSend(TimeoutProcessor processor) { int numTimedOut = 0; for (List callList : callsToSend.values()) { numTimedOut += processor.handleTimeouts(callList, @@ -723,65 +989,81 @@ private void timeoutCallsToSend(TimeoutProcessor processor, Map } if (numTimedOut > 0) log.debug("Timed out {} call(s) with assigned nodes.", numTimedOut); + return numTimedOut; } /** - * Choose nodes for the calls in the callsToSend list. + * Drain all the calls from newCalls into pendingCalls. * * This function holds the lock for the minimum amount of time, to avoid blocking * users of AdminClient who will also take the lock to add new calls. + */ + private synchronized void drainNewCalls() { + if (!newCalls.isEmpty()) { + pendingCalls.addAll(newCalls); + newCalls.clear(); + } + } + + /** + * Choose nodes for the calls in the pendingCalls list. * * @param now The current time in milliseconds. - * @param callsToSend A map of nodes to the calls they need to handle. - * + * @return The minimum time until a call is ready to be retried if any of the pending + * calls are backing off after a failure */ - private void chooseNodesForNewCalls(long now, Map> callsToSend) { - List newCallsToAdd = null; - synchronized (this) { - if (newCalls.isEmpty()) { - return; + private long maybeDrainPendingCalls(long now) { + long pollTimeout = Long.MAX_VALUE; + log.trace("Trying to choose nodes for {} at {}", pendingCalls, now); + + Iterator pendingIter = pendingCalls.iterator(); + while (pendingIter.hasNext()) { + Call call = pendingIter.next(); + + // If the call is being retried, await the proper backoff before finding the node + if (now < call.nextAllowedTryMs) { + pollTimeout = Math.min(pollTimeout, call.nextAllowedTryMs - now); + } else if (maybeDrainPendingCall(call, now)) { + pendingIter.remove(); } - newCallsToAdd = newCalls; - newCalls = new LinkedList<>(); - } - for (Call call : newCallsToAdd) { - chooseNodeForNewCall(now, callsToSend, call); } + return pollTimeout; } /** - * Choose a node for a new call. - * - * @param now The current time in milliseconds. - * @param callsToSend A map of nodes to the calls they need to handle. - * @param call The call. + * Check whether a pending call can be assigned a node. Return true if the pending call was either + * transferred to the callsToSend collection or if the call was failed. Return false if it + * should remain pending. */ - private void chooseNodeForNewCall(long now, Map> callsToSend, Call call) { - Node node = call.nodeProvider.provide(); - if (node == null) { - call.fail(now, new BrokerNotAvailableException( - String.format("Error choosing node for %s: no node found.", call.callName))); - return; + private boolean maybeDrainPendingCall(Call call, long now) { + try { + Node node = call.nodeProvider.provide(); + if (node != null) { + log.trace("Assigned {} to node {}", call, node); + call.curNode = node; + getOrCreateListValue(callsToSend, node).add(call); + return true; + } else { + log.trace("Unable to assign {} to a node.", call); + return false; + } + } catch (Throwable t) { + // Handle authentication errors while choosing nodes. + log.debug("Unable to choose node for {}", call, t); + call.fail(now, t); + return true; } - log.trace("Assigned {} to {}", call, node); - getOrCreateListValue(callsToSend, node).add(call); } /** * Send the calls which are ready. * * @param now The current time in milliseconds. - * @param callsToSend The calls to send, by node. - * @param correlationIdToCalls A map of correlation IDs to calls. - * @param callsInFlight A map of nodes to the calls they have in flight. - * * @return The minimum timeout we need for poll(). */ - private long sendEligibleCalls(long now, Map> callsToSend, - Map correlationIdToCalls, Map> callsInFlight) { + private long sendEligibleCalls(long now) { long pollTimeout = Long.MAX_VALUE; - for (Iterator>> iter = callsToSend.entrySet().iterator(); - iter.hasNext(); ) { + for (Iterator>> iter = callsToSend.entrySet().iterator(); iter.hasNext(); ) { Map.Entry> entry = iter.next(); List calls = entry.getValue(); if (calls.isEmpty()) { @@ -790,23 +1072,25 @@ private long sendEligibleCalls(long now, Map> callsToSend, } Node node = entry.getKey(); if (!client.ready(node, now)) { - long nodeTimeout = client.connectionDelay(node, now); + long nodeTimeout = client.pollDelayMs(node, now); pollTimeout = Math.min(pollTimeout, nodeTimeout); log.trace("Client is not ready to send to {}. Must delay {} ms", node, nodeTimeout); continue; } Call call = calls.remove(0); - int timeoutMs = calcTimeoutMsRemainingAsInt(now, call.deadlineMs); - AbstractRequest.Builder requestBuilder = null; + int requestTimeoutMs = Math.min(KafkaAdminClient.this.requestTimeoutMs, + calcTimeoutMsRemainingAsInt(now, call.deadlineMs)); + AbstractRequest.Builder requestBuilder; try { - requestBuilder = call.createRequest(timeoutMs); + requestBuilder = call.createRequest(requestTimeoutMs); } catch (Throwable throwable) { call.fail(now, new KafkaException(String.format( "Internal error sending %s to %s.", call.callName, node))); continue; } - ClientRequest clientRequest = client.newClientRequest(node.idString(), requestBuilder, now, true); - log.trace("Sending {} to {}. correlationId={}", requestBuilder, node, clientRequest.correlationId()); + ClientRequest clientRequest = client.newClientRequest(node.idString(), requestBuilder, now, + true, requestTimeoutMs, null); + log.debug("Sending {} to {}. correlationId={}", requestBuilder, node, clientRequest.correlationId()); client.send(clientRequest, now); getOrCreateListValue(callsInFlight, node.idString()).add(call); correlationIdToCalls.put(clientRequest.correlationId(), call); @@ -822,9 +1106,8 @@ private long sendEligibleCalls(long now, Map> callsToSend, * to time them out is to close the entire connection. * * @param processor The timeout processor. - * @param callsInFlight A map of nodes to the calls they have in flight. */ - private void timeoutCallsInFlight(TimeoutProcessor processor, Map> callsInFlight) { + private void timeoutCallsInFlight(TimeoutProcessor processor) { int numTimedOut = 0; for (Map.Entry> entry : callsInFlight.entrySet()) { List contexts = entry.getValue(); @@ -852,51 +1135,17 @@ private void timeoutCallsInFlight(TimeoutProcessor processor, Map> callsToSend) { - AuthenticationException authenticationException = metadata.getAndClearAuthenticationException(); - if (authenticationException == null) { - for (Node node : callsToSend.keySet()) { - authenticationException = client.authenticationException(node); - if (authenticationException != null) - break; - } - } - if (authenticationException != null) { - synchronized (this) { - failCalls(now, newCalls, authenticationException); - } - for (List calls : callsToSend.values()) { - failCalls(now, calls, authenticationException); - } - callsToSend.clear(); - } - } - - private void failCalls(long now, List calls, AuthenticationException authenticationException) { - for (Call call : calls) { - call.fail(now, authenticationException); - } - calls.clear(); - } - /** * Handle responses from the server. * * @param now The current time in milliseconds. * @param responses The latest responses from KafkaClient. - * @param correlationIdToCall A map of correlation IDs to calls. - * @param callsInFlight A map of nodes to the calls they have in flight. - **/ - private void handleResponses(long now, List responses, Map> callsInFlight, - Map correlationIdToCall) { + **/ + private void handleResponses(long now, List responses) { for (ClientResponse response : responses) { int correlationId = response.requestHeader().correlationId(); - Call call = correlationIdToCall.get(correlationId); + Call call = correlationIdToCalls.get(correlationId); if (call == null) { // If the server returns information about a correlation ID we didn't use yet, // an internal server error has occurred. Close the connection and log an error message. @@ -908,7 +1157,7 @@ private void handleResponses(long now, List responses, Map calls = callsInFlight.get(response.destination()); if ((calls == null) || (!calls.remove(call))) { log.error("Internal server error on {}: ignoring call {} in correlationIdToCall " + @@ -917,19 +1166,23 @@ private void handleResponses(long now, List responses, Map responses, Map> callsToSend, Map correlationIdToCalls) { - if (newCalls.isEmpty() && callsToSend.isEmpty() && correlationIdToCalls.isEmpty()) { + /** + * Unassign calls that have not yet been sent based on some predicate. For example, this + * is used to reassign the calls that have been assigned to a disconnected node. + * + * @param shouldUnassign Condition for reassignment. If the predicate is true, then the calls will + * be put back in the pendingCalls collection and they will be reassigned + */ + private void unassignUnsentCalls(Predicate shouldUnassign) { + for (Iterator>> iter = callsToSend.entrySet().iterator(); iter.hasNext(); ) { + Map.Entry> entry = iter.next(); + Node node = entry.getKey(); + List awaitingCalls = entry.getValue(); + + if (awaitingCalls.isEmpty()) { + iter.remove(); + } else if (shouldUnassign.test(node)) { + pendingCalls.addAll(awaitingCalls); + iter.remove(); + } + } + } + + private boolean hasActiveExternalCalls(Collection calls) { + for (Call call : calls) { + if (!call.isInternal()) { + return true; + } + } + return false; + } + + /** + * Return true if there are currently active external calls. + */ + private boolean hasActiveExternalCalls() { + if (hasActiveExternalCalls(pendingCalls)) { + return true; + } + for (List callList : callsToSend.values()) { + if (hasActiveExternalCalls(callList)) { + return true; + } + } + return hasActiveExternalCalls(correlationIdToCalls.values()); + } + + private boolean threadShouldExit(long now, long curHardShutdownTimeMs) { + if (!hasActiveExternalCalls()) { log.trace("All work has been completed, and the I/O thread is now exiting."); return true; } - if (now > curHardShutdownTimeMs) { + if (now >= curHardShutdownTimeMs) { log.info("Forcing a hard I/O thread shutdown. Requests in progress will be aborted."); return true; } @@ -955,79 +1253,87 @@ private synchronized boolean threadShouldExit(long now, long curHardShutdownTime @Override public void run() { - /* - * Maps nodes to calls that we want to send. - */ - Map> callsToSend = new HashMap<>(); - - /* - * Maps node ID strings to calls that have been sent. - */ - Map> callsInFlight = new HashMap<>(); - - /* - * Maps correlation IDs to calls that have been sent. - */ - Map correlationIdToCalls = new HashMap<>(); - - /* - * The previous metadata version which wasn't usable, or null if there is none. - */ - Integer prevMetadataVersion = null; + log.trace("Thread starting"); + try { + processRequests(); + } finally { + AppInfoParser.unregisterAppInfo(JMX_PREFIX, clientId, metrics); + + int numTimedOut = 0; + TimeoutProcessor timeoutProcessor = new TimeoutProcessor(Long.MAX_VALUE); + synchronized (this) { + numTimedOut += timeoutProcessor.handleTimeouts(newCalls, "The AdminClient thread has exited."); + newCalls = null; + } + numTimedOut += timeoutProcessor.handleTimeouts(pendingCalls, "The AdminClient thread has exited."); + numTimedOut += timeoutCallsToSend(timeoutProcessor); + numTimedOut += timeoutProcessor.handleTimeouts(correlationIdToCalls.values(), + "The AdminClient thread has exited."); + if (numTimedOut > 0) { + log.debug("Timed out {} remaining operation(s).", numTimedOut); + } + closeQuietly(client, "KafkaClient"); + closeQuietly(metrics, "Metrics"); + log.debug("Exiting AdminClientRunnable thread."); + } + } + private void processRequests() { long now = time.milliseconds(); - log.trace("Thread starting"); while (true) { + // Copy newCalls into pendingCalls. + drainNewCalls(); + // Check if the AdminClient thread should shut down. long curHardShutdownTimeMs = hardShutdownTimeMs.get(); - if ((curHardShutdownTimeMs != INVALID_SHUTDOWN_TIME) && - threadShouldExit(now, curHardShutdownTimeMs, callsToSend, correlationIdToCalls)) + if ((curHardShutdownTimeMs != INVALID_SHUTDOWN_TIME) && threadShouldExit(now, curHardShutdownTimeMs)) break; // Handle timeouts. TimeoutProcessor timeoutProcessor = timeoutProcessorFactory.create(now); - timeoutNewCalls(timeoutProcessor); - timeoutCallsToSend(timeoutProcessor, callsToSend); - timeoutCallsInFlight(timeoutProcessor, callsInFlight); + timeoutPendingCalls(timeoutProcessor); + timeoutCallsToSend(timeoutProcessor); + timeoutCallsInFlight(timeoutProcessor); long pollTimeout = Math.min(1200000, timeoutProcessor.nextTimeoutMs()); if (curHardShutdownTimeMs != INVALID_SHUTDOWN_TIME) { pollTimeout = Math.min(pollTimeout, curHardShutdownTimeMs - now); } - // Handle new calls and metadata update requests. - prevMetadataVersion = checkMetadataReady(prevMetadataVersion); - if (prevMetadataVersion == null) { - chooseNodesForNewCalls(now, callsToSend); - pollTimeout = Math.min(pollTimeout, - sendEligibleCalls(now, callsToSend, correlationIdToCalls, callsInFlight)); + // Choose nodes for our pending calls. + pollTimeout = Math.min(pollTimeout, maybeDrainPendingCalls(now)); + long metadataFetchDelayMs = metadataManager.metadataFetchDelayMs(now); + if (metadataFetchDelayMs == 0) { + metadataManager.transitionToUpdatePending(now); + Call metadataCall = makeMetadataCall(now); + // Create a new metadata fetch call and add it to the end of pendingCalls. + // Assign a node for just the new call (we handled the other pending nodes above). + + if (!maybeDrainPendingCall(metadataCall, now)) + pendingCalls.add(metadataCall); + } + pollTimeout = Math.min(pollTimeout, sendEligibleCalls(now)); + + if (metadataFetchDelayMs > 0) { + pollTimeout = Math.min(pollTimeout, metadataFetchDelayMs); } + // Ensure that we use a small poll timeout if there are pending calls which need to be sent + if (!pendingCalls.isEmpty()) + pollTimeout = Math.min(pollTimeout, retryBackoffMs); + // Wait for network responses. log.trace("Entering KafkaClient#poll(timeout={})", pollTimeout); List responses = client.poll(pollTimeout, now); log.trace("KafkaClient#poll retrieved {} response(s)", responses.size()); + // unassign calls to disconnected nodes + unassignUnsentCalls(client::connectionFailed); + // Update the current time and handle the latest responses. now = time.milliseconds(); - handleAuthenticationException(now, callsToSend); - handleResponses(now, responses, callsInFlight, correlationIdToCalls); - } - int numTimedOut = 0; - TimeoutProcessor timeoutProcessor = new TimeoutProcessor(Long.MAX_VALUE); - synchronized (this) { - numTimedOut += timeoutProcessor.handleTimeouts(newCalls, - "The AdminClient thread has exited."); - newCalls = null; + handleResponses(now, responses); } - numTimedOut += timeoutProcessor.handleTimeouts(correlationIdToCalls.values(), - "The AdminClient thread has exited."); - if (numTimedOut > 0) { - log.debug("Timed out {} remaining operations.", numTimedOut); - } - closeQuietly(client, "KafkaClient"); - closeQuietly(metrics, "Metrics"); - log.debug("Exiting AdminClientRunnable thread."); } /** @@ -1041,6 +1347,11 @@ public void run() { * @param now The current time in milliseconds. */ void enqueue(Call call, long now) { + if (call.tries > maxRetries) { + log.debug("Max retries {} for {} reached", maxRetries, call); + call.fail(time.milliseconds(), new TimeoutException()); + return; + } if (log.isDebugEnabled()) { log.debug("Queueing {} with a timeout {} ms from now.", call, call.deadlineMs - now); } @@ -1075,114 +1386,312 @@ void call(Call call, long now) { enqueue(call, now); } } + + /** + * Create a new metadata call. + */ + private Call makeMetadataCall(long now) { + return new Call(true, "fetchMetadata", calcDeadlineMs(now, requestTimeoutMs), + new MetadataUpdateNodeIdProvider()) { + @Override + public MetadataRequest.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(new MetadataRequestData() + .setTopics(Collections.emptyList()) + .setAllowAutoTopicCreation(true)); + } + + @Override + public void handleResponse(AbstractResponse abstractResponse) { + MetadataResponse response = (MetadataResponse) abstractResponse; + long now = time.milliseconds(); + metadataManager.update(response.cluster(), now); + + // Unassign all unsent requests after a metadata refresh to allow for a new + // destination to be selected from the new metadata + unassignUnsentCalls(node -> true); + } + + @Override + public void handleFailure(Throwable e) { + metadataManager.updateFailed(e); + } + }; + } + } + + /** + * Returns true if a topic name cannot be represented in an RPC. This function does NOT check + * whether the name is too long, contains invalid characters, etc. It is better to enforce + * those policies on the server, so that they can be changed in the future if needed. + */ + private static boolean topicNameIsUnrepresentable(String topicName) { + return topicName == null || topicName.isEmpty(); + } + + private static boolean groupIdIsUnrepresentable(String groupId) { + return groupId == null; + } + + // for testing + int numPendingCalls() { + return runnable.pendingCalls.size(); + } + + /** + * Fail futures in the given stream which are not done. + * Used when a response handler expected a result for some entity but no result was present. + */ + private static void completeUnrealizedFutures( + Stream>> futures, + Function messageFormatter) { + futures.filter(entry -> !entry.getValue().isDone()).forEach(entry -> + entry.getValue().completeExceptionally(new ApiException(messageFormatter.apply(entry.getKey())))); + } + + /** + * Fail futures in the given Map which were retried due to exceeding quota. We propagate + * the initial error back to the caller if the request timed out. + */ + private static void maybeCompleteQuotaExceededException( + boolean shouldRetryOnQuotaViolation, + Throwable throwable, + Map> futures, + Map quotaExceededExceptions, + int throttleTimeDelta) { + if (shouldRetryOnQuotaViolation && throwable instanceof TimeoutException) { + quotaExceededExceptions.forEach((key, value) -> futures.get(key).completeExceptionally( + new ThrottlingQuotaExceededException( + Math.max(0, value.throttleTimeMs() - throttleTimeDelta), + value.getMessage()))); + } } @Override public CreateTopicsResult createTopics(final Collection newTopics, final CreateTopicsOptions options) { - final Map> topicFutures = new HashMap<>(newTopics.size()); - final Map topicsMap = new HashMap<>(newTopics.size()); + final Map> topicFutures = new HashMap<>(newTopics.size()); + final CreatableTopicCollection topics = new CreatableTopicCollection(); for (NewTopic newTopic : newTopics) { - if (topicFutures.get(newTopic.name()) == null) { - topicFutures.put(newTopic.name(), new KafkaFutureImpl()); - topicsMap.put(newTopic.name(), newTopic.convertToTopicDetails()); + if (topicNameIsUnrepresentable(newTopic.name())) { + 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); + } else if (!topicFutures.containsKey(newTopic.name())) { + topicFutures.put(newTopic.name(), new KafkaFutureImpl<>()); + topics.add(newTopic.convertToCreatableTopic()); } } - final long now = time.milliseconds(); - runnable.call(new Call("createTopics", calcDeadlineMs(now, options.timeoutMs()), - new ControllerNodeProvider()) { + if (!topics.isEmpty()) { + final long now = time.milliseconds(); + final long deadline = calcDeadlineMs(now, options.timeoutMs()); + final Call call = getCreateTopicsCall(options, topicFutures, topics, + Collections.emptyMap(), now, deadline); + runnable.call(call, now); + } + return new CreateTopicsResult(new HashMap<>(topicFutures)); + } + private Call getCreateTopicsCall(final CreateTopicsOptions options, + final Map> futures, + final CreatableTopicCollection topics, + final Map quotaExceededExceptions, + final long now, + final long deadline) { + return new Call("createTopics", deadline, new ControllerNodeProvider()) { @Override - public AbstractRequest.Builder createRequest(int timeoutMs) { - return new CreateTopicsRequest.Builder(topicsMap, timeoutMs, options.shouldValidateOnly()); + public CreateTopicsRequest.Builder createRequest(int timeoutMs) { + return new CreateTopicsRequest.Builder( + new CreateTopicsRequestData() + .setTopics(topics) + .setTimeoutMs(timeoutMs) + .setValidateOnly(options.shouldValidateOnly())); } @Override public void handleResponse(AbstractResponse abstractResponse) { - CreateTopicsResponse response = (CreateTopicsResponse) abstractResponse; + // Check for controller change + handleNotControllerError(abstractResponse); // Handle server responses for particular topics. - for (Map.Entry entry : response.errors().entrySet()) { - KafkaFutureImpl future = topicFutures.get(entry.getKey()); + final CreateTopicsResponse response = (CreateTopicsResponse) abstractResponse; + final CreatableTopicCollection retryTopics = new CreatableTopicCollection(); + final Map retryTopicQuotaExceededExceptions = new HashMap<>(); + for (CreatableTopicResult result : response.data().topics()) { + KafkaFutureImpl future = futures.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(); - if (exception != null) { - future.completeExceptionally(exception); + ApiError error = new ApiError(result.errorCode(), result.errorMessage()); + if (error.isFailure()) { + if (error.is(Errors.THROTTLING_QUOTA_EXCEEDED)) { + ThrottlingQuotaExceededException quotaExceededException = new ThrottlingQuotaExceededException( + response.throttleTimeMs(), error.messageWithFallback()); + if (options.shouldRetryOnQuotaViolation()) { + retryTopics.add(topics.find(result.name()).duplicate()); + retryTopicQuotaExceededExceptions.put(result.name(), quotaExceededException); + } else { + future.completeExceptionally(quotaExceededException); + } + } else { + future.completeExceptionally(error.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(this::configEntry) + .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(); - if (!future.isDone()) { - future.completeExceptionally(new ApiException("The server response did not " + - "contain a reference to node " + entry.getKey())); - } + // If there are topics to retry, retry them; complete unrealized futures otherwise. + if (retryTopics.isEmpty()) { + // The server should send back a response for every topic. But do a sanity check anyway. + completeUnrealizedFutures(futures.entrySet().stream(), + topic -> "The controller response did not contain a result for topic " + topic); + } else { + final long now = time.milliseconds(); + final Call call = getCreateTopicsCall(options, futures, retryTopics, + retryTopicQuotaExceededExceptions, now, deadline); + runnable.call(call, now); } } + private ConfigEntry configEntry(CreatableTopicConfigs config) { + return new ConfigEntry( + config.name(), + config.value(), + configSource(DescribeConfigsResponse.ConfigSource.forId(config.configSource())), + config.isSensitive(), + config.readOnly(), + Collections.emptyList(), + null, + null); + } + @Override void handleFailure(Throwable throwable) { - completeAllExceptionally(topicFutures.values(), throwable); + // If there were any topics retries due to a quota exceeded exception, we propagate + // the initial error back to the caller if the request timed out. + maybeCompleteQuotaExceededException(options.shouldRetryOnQuotaViolation(), + throwable, futures, quotaExceededExceptions, (int) (time.milliseconds() - now)); + // Fail all the other remaining futures + completeAllExceptionally(futures.values(), throwable); } - }, now); - return new CreateTopicsResult(new HashMap>(topicFutures)); + }; } @Override public DeleteTopicsResult deleteTopics(final Collection topicNames, - DeleteTopicsOptions options) { + final DeleteTopicsOptions options) { final Map> topicFutures = new HashMap<>(topicNames.size()); + final List validTopicNames = new ArrayList<>(topicNames.size()); for (String topicName : topicNames) { - if (topicFutures.get(topicName) == null) { - topicFutures.put(topicName, new KafkaFutureImpl()); + if (topicNameIsUnrepresentable(topicName)) { + KafkaFutureImpl future = new KafkaFutureImpl<>(); + future.completeExceptionally(new InvalidTopicException("The given topic name '" + + topicName + "' cannot be represented in a request.")); + topicFutures.put(topicName, future); + } else if (!topicFutures.containsKey(topicName)) { + topicFutures.put(topicName, new KafkaFutureImpl<>()); + validTopicNames.add(topicName); } } - final long now = time.milliseconds(); - runnable.call(new Call("deleteTopics", calcDeadlineMs(now, options.timeoutMs()), - new ControllerNodeProvider()) { + if (!validTopicNames.isEmpty()) { + final long now = time.milliseconds(); + final long deadline = calcDeadlineMs(now, options.timeoutMs()); + final Call call = getDeleteTopicsCall(options, topicFutures, validTopicNames, + Collections.emptyMap(), now, deadline); + runnable.call(call, now); + } + return new DeleteTopicsResult(new HashMap<>(topicFutures)); + } + private Call getDeleteTopicsCall(final DeleteTopicsOptions options, + final Map> futures, + final List topics, + final Map quotaExceededExceptions, + final long now, + final long deadline) { + return new Call("deleteTopics", deadline, new ControllerNodeProvider()) { @Override - AbstractRequest.Builder createRequest(int timeoutMs) { - return new DeleteTopicsRequest.Builder(new HashSet<>(topicNames), timeoutMs); + DeleteTopicsRequest.Builder createRequest(int timeoutMs) { + return new DeleteTopicsRequest.Builder( + new DeleteTopicsRequestData() + .setTopicNames(topics) + .setTimeoutMs(timeoutMs)); } @Override void handleResponse(AbstractResponse abstractResponse) { - DeleteTopicsResponse response = (DeleteTopicsResponse) abstractResponse; + // Check for controller change + handleNotControllerError(abstractResponse); // Handle server responses for particular topics. - for (Map.Entry entry : response.errors().entrySet()) { - KafkaFutureImpl future = topicFutures.get(entry.getKey()); + final DeleteTopicsResponse response = (DeleteTopicsResponse) abstractResponse; + final List retryTopics = new ArrayList<>(); + final Map retryTopicQuotaExceededExceptions = new HashMap<>(); + for (DeletableTopicResult result : response.data().responses()) { + KafkaFutureImpl future = futures.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(); - if (exception != null) { - future.completeExceptionally(exception); + ApiError error = new ApiError(result.errorCode(), result.errorMessage()); + if (error.isFailure()) { + if (error.is(Errors.THROTTLING_QUOTA_EXCEEDED)) { + ThrottlingQuotaExceededException quotaExceededException = new ThrottlingQuotaExceededException( + response.throttleTimeMs(), error.messageWithFallback()); + if (options.shouldRetryOnQuotaViolation()) { + retryTopics.add(result.name()); + retryTopicQuotaExceededExceptions.put(result.name(), quotaExceededException); + } else { + future.completeExceptionally(quotaExceededException); + } + } else { + future.completeExceptionally(error.exception()); + } } else { future.complete(null); } } } - // 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(); - if (!future.isDone()) { - future.completeExceptionally(new ApiException("The server response did not " + - "contain a reference to node " + entry.getKey())); - } + // If there are topics to retry, retry them; complete unrealized futures otherwise. + if (retryTopics.isEmpty()) { + // The server should send back a response for every topic. But do a sanity check anyway. + completeUnrealizedFutures(futures.entrySet().stream(), + topic -> "The controller response did not contain a result for topic " + topic); + } else { + final long now = time.milliseconds(); + final Call call = getDeleteTopicsCall(options, futures, retryTopics, + retryTopicQuotaExceededExceptions, now, deadline); + runnable.call(call, now); } } @Override void handleFailure(Throwable throwable) { - completeAllExceptionally(topicFutures.values(), throwable); + // If there were any topics retries due to a quota exceeded exception, we propagate + // the initial error back to the caller if the request timed out. + maybeCompleteQuotaExceededException(options.shouldRetryOnQuotaViolation(), + throwable, futures, quotaExceededExceptions, (int) (time.milliseconds() - now)); + // Fail all the other remaining futures + completeAllExceptionally(futures.values(), throwable); } - }, now); - return new DeleteTopicsResult(new HashMap>(topicFutures)); + }; } @Override @@ -1193,19 +1702,19 @@ public ListTopicsResult listTopics(final ListTopicsOptions options) { new LeastLoadedNodeProvider()) { @Override - AbstractRequest.Builder createRequest(int timeoutMs) { + MetadataRequest.Builder createRequest(int timeoutMs) { return MetadataRequest.Builder.allTopics(); } @Override void handleResponse(AbstractResponse abstractResponse) { MetadataResponse response = (MetadataResponse) abstractResponse; - Cluster cluster = response.cluster(); Map topicListing = new HashMap<>(); - for (String topicName : cluster.topics()) { - boolean internal = cluster.internalTopics().contains(topicName); - if (!internal || options.shouldListInternal()) - topicListing.put(topicName, new TopicListing(topicName, internal)); + for (MetadataResponse.TopicMetadata topicMetadata : response.topicMetadata()) { + String topicName = topicMetadata.topic(); + boolean isInternal = topicMetadata.isInternal(); + if (!topicMetadata.isInternal() || options.shouldListInternal()) + topicListing.put(topicName, new TopicListing(topicName, isInternal)); } topicListingFuture.complete(topicListing); } @@ -1223,21 +1732,29 @@ public DescribeTopicsResult describeTopics(final Collection topicNames, final Map> topicFutures = new HashMap<>(topicNames.size()); final ArrayList topicNamesList = new ArrayList<>(); for (String topicName : topicNames) { - if (!topicFutures.containsKey(topicName)) { - topicFutures.put(topicName, new KafkaFutureImpl()); + if (topicNameIsUnrepresentable(topicName)) { + KafkaFutureImpl future = new KafkaFutureImpl<>(); + future.completeExceptionally(new InvalidTopicException("The given topic name '" + + topicName + "' cannot be represented in a request.")); + topicFutures.put(topicName, future); + } else if (!topicFutures.containsKey(topicName)) { + topicFutures.put(topicName, new KafkaFutureImpl<>()); topicNamesList.add(topicName); } } final long now = time.milliseconds(); - runnable.call(new Call("describeTopics", calcDeadlineMs(now, options.timeoutMs()), - new ControllerNodeProvider()) { + Call call = new Call("describeTopics", calcDeadlineMs(now, options.timeoutMs()), + new LeastLoadedNodeProvider()) { private boolean supportsDisablingTopicCreation = true; @Override - AbstractRequest.Builder createRequest(int timeoutMs) { + MetadataRequest.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(); } @@ -1246,17 +1763,18 @@ AbstractRequest.Builder createRequest(int timeoutMs) { void handleResponse(AbstractResponse abstractResponse) { MetadataResponse response = (MetadataResponse) abstractResponse; // Handle server responses for particular topics. + Cluster cluster = response.cluster(); + Map errors = response.errors(); for (Map.Entry> entry : topicFutures.entrySet()) { String topicName = entry.getKey(); KafkaFutureImpl future = entry.getValue(); - Errors topicError = response.errors().get(topicName); + Errors topicError = errors.get(topicName); if (topicError != null) { future.completeExceptionally(topicError.exception()); continue; } - Cluster cluster = response.cluster(); if (!cluster.topics().contains(topicName)) { - future.completeExceptionally(new InvalidTopicException("Topic " + topicName + " not found.")); + future.completeExceptionally(new UnknownTopicOrPartitionException("Topic " + topicName + " not found.")); continue; } boolean isInternal = cluster.internalTopics().contains(topicName); @@ -1268,13 +1786,9 @@ void handleResponse(AbstractResponse abstractResponse) { Arrays.asList(partitionInfo.inSyncReplicas())); partitions.add(topicPartitionInfo); } - Collections.sort(partitions, new Comparator() { - @Override - public int compare(TopicPartitionInfo tp1, TopicPartitionInfo tp2) { - return Integer.compare(tp1.partition(), tp2.partition()); - } - }); - TopicDescription topicDescription = new TopicDescription(topicName, isInternal, partitions); + partitions.sort(Comparator.comparingInt(TopicPartitionInfo::partition)); + TopicDescription topicDescription = new TopicDescription(topicName, isInternal, partitions, + validAclOperations(response.topicAuthorizedOperations(topicName).get())); future.complete(topicDescription); } } @@ -1298,8 +1812,11 @@ boolean handleUnsupportedVersionException(UnsupportedVersionException exception) void handleFailure(Throwable throwable) { completeAllExceptionally(topicFutures.values(), throwable); } - }, now); - return new DescribeTopicsResult(new HashMap>(topicFutures)); + }; + if (!topicNamesList.isEmpty()) { + runnable.call(call, now); + } + return new DescribeTopicsResult(new HashMap<>(topicFutures)); } @Override @@ -1307,15 +1824,20 @@ 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()) { @Override - AbstractRequest.Builder createRequest(int timeoutMs) { + MetadataRequest.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 @@ -1324,6 +1846,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) { @@ -1337,21 +1861,29 @@ 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 public DescribeAclsResult describeAcls(final AclBindingFilter filter, DescribeAclsOptions options) { + if (filter.isUnknown()) { + KafkaFutureImpl> future = new KafkaFutureImpl<>(); + future.completeExceptionally(new InvalidRequestException("The AclBindingFilter " + + "must not contain UNKNOWN elements.")); + return new DescribeAclsResult(future); + } final long now = time.milliseconds(); final KafkaFutureImpl> future = new KafkaFutureImpl<>(); runnable.call(new Call("describeAcls", calcDeadlineMs(now, options.timeoutMs()), new LeastLoadedNodeProvider()) { @Override - AbstractRequest.Builder createRequest(int timeoutMs) { + DescribeAclsRequest.Builder createRequest(int timeoutMs) { return new DescribeAclsRequest.Builder(filter); } @@ -1361,7 +1893,7 @@ void handleResponse(AbstractResponse abstractResponse) { if (response.error().isFailure()) { future.completeExceptionally(response.error().exception()); } else { - future.complete(response.acls()); + future.complete(DescribeAclsResponse.aclBindings(response.acls())); } } @@ -1378,44 +1910,48 @@ public CreateAclsResult createAcls(Collection acls, CreateAclsOption final long now = time.milliseconds(); final Map> futures = new HashMap<>(); final List aclCreations = new ArrayList<>(); + final List aclBindingsSent = new ArrayList<>(); for (AclBinding acl : acls) { if (futures.get(acl) == null) { KafkaFutureImpl future = new KafkaFutureImpl<>(); futures.put(acl, future); String indefinite = acl.toFilter().findIndefiniteField(); if (indefinite == null) { - aclCreations.add(new AclCreation(acl)); + aclCreations.add(CreateAclsRequest.aclCreation(acl)); + aclBindingsSent.add(acl); } else { future.completeExceptionally(new InvalidRequestException("Invalid ACL creation: " + indefinite)); } } } + final CreateAclsRequestData data = new CreateAclsRequestData().setCreations(aclCreations); runnable.call(new Call("createAcls", calcDeadlineMs(now, options.timeoutMs()), new LeastLoadedNodeProvider()) { @Override - AbstractRequest.Builder createRequest(int timeoutMs) { - return new CreateAclsRequest.Builder(aclCreations); + CreateAclsRequest.Builder createRequest(int timeoutMs) { + return new CreateAclsRequest.Builder(data); } @Override void handleResponse(AbstractResponse abstractResponse) { CreateAclsResponse response = (CreateAclsResponse) abstractResponse; - List responses = response.aclCreationResponses(); - Iterator iter = responses.iterator(); - for (AclCreation aclCreation : aclCreations) { - KafkaFutureImpl future = futures.get(aclCreation.acl()); + List responses = response.results(); + Iterator iter = responses.iterator(); + for (AclBinding aclBinding : aclBindingsSent) { + KafkaFutureImpl future = futures.get(aclBinding); if (!iter.hasNext()) { future.completeExceptionally(new UnknownServerException( - "The broker reported no creation result for the given ACL.")); + "The broker reported no creation result for the given ACL: " + aclBinding)); } else { - AclCreationResponse creation = iter.next(); - if (creation.error().isFailure()) { - future.completeExceptionally(creation.error().exception()); - } else { + AclCreationResult creation = iter.next(); + Errors error = Errors.forCode(creation.errorCode()); + ApiError apiError = new ApiError(error, creation.errorMessage()); + if (apiError.isFailure()) + future.completeExceptionally(apiError.exception()); + else future.complete(null); - } } } } @@ -1425,46 +1961,53 @@ void handleFailure(Throwable throwable) { completeAllExceptionally(futures.values(), throwable); } }, now); - return new CreateAclsResult(new HashMap>(futures)); + return new CreateAclsResult(new HashMap<>(futures)); } @Override public DeleteAclsResult deleteAcls(Collection filters, DeleteAclsOptions options) { final long now = time.milliseconds(); final Map> futures = new HashMap<>(); - final List filterList = new ArrayList<>(); + final List aclBindingFiltersSent = new ArrayList<>(); + final List deleteAclsFilters = new ArrayList<>(); for (AclBindingFilter filter : filters) { if (futures.get(filter) == null) { - filterList.add(filter); - futures.put(filter, new KafkaFutureImpl()); + aclBindingFiltersSent.add(filter); + deleteAclsFilters.add(DeleteAclsRequest.deleteAclsFilter(filter)); + futures.put(filter, new KafkaFutureImpl<>()); } } + final DeleteAclsRequestData data = new DeleteAclsRequestData().setFilters(deleteAclsFilters); runnable.call(new Call("deleteAcls", calcDeadlineMs(now, options.timeoutMs()), new LeastLoadedNodeProvider()) { @Override - AbstractRequest.Builder createRequest(int timeoutMs) { - return new DeleteAclsRequest.Builder(filterList); + DeleteAclsRequest.Builder createRequest(int timeoutMs) { + return new DeleteAclsRequest.Builder(data); } @Override void handleResponse(AbstractResponse abstractResponse) { DeleteAclsResponse response = (DeleteAclsResponse) abstractResponse; - List responses = response.responses(); - Iterator iter = responses.iterator(); - for (AclBindingFilter filter : filterList) { - KafkaFutureImpl future = futures.get(filter); + List results = response.filterResults(); + Iterator iter = results.iterator(); + for (AclBindingFilter bindingFilter : aclBindingFiltersSent) { + KafkaFutureImpl future = futures.get(bindingFilter); if (!iter.hasNext()) { future.completeExceptionally(new UnknownServerException( "The broker reported no deletion result for the given filter.")); } else { - AclFilterResponse deletion = iter.next(); - if (deletion.error().isFailure()) { - future.completeExceptionally(deletion.error().exception()); + DeleteAclsFilterResult filterResult = iter.next(); + ApiError error = new ApiError(Errors.forCode(filterResult.errorCode()), filterResult.errorMessage()); + if (error.isFailure()) { + future.completeExceptionally(error.exception()); } else { List filterResults = new ArrayList<>(); - for (AclDeletionResult deletionResult : deletion.deletions()) { - filterResults.add(new FilterResult(deletionResult.acl(), deletionResult.error().exception())); + for (DeleteAclsMatchingAcl matchingAcl : filterResult.matchingAcls()) { + ApiError aclError = new ApiError(Errors.forCode(matchingAcl.errorCode()), + matchingAcl.errorMessage()); + AclBinding aclBinding = DeleteAclsResponse.aclBinding(matchingAcl); + filterResults.add(new FilterResult(aclBinding, aclError.exception())); } future.complete(new FilterResults(filterResults)); } @@ -1477,156 +2020,171 @@ void handleFailure(Throwable throwable) { completeAllExceptionally(futures.values(), throwable); } }, now); - return new DeleteAclsResult(new HashMap>(futures)); + return new DeleteAclsResult(new HashMap<>(futures)); } @Override public DescribeConfigsResult describeConfigs(Collection configResources, final DescribeConfigsOptions options) { - final Map> unifiedRequestFutures = new HashMap<>(); - final Map> brokerFutures = new HashMap<>(configResources.size()); - - // The BROKER resources which we want to describe. We must make a separate DescribeConfigs - // request for every BROKER resource we want to describe. - final Collection brokerResources = new ArrayList<>(); - - // The non-BROKER resources which we want to describe. These resources can be described by a - // single, unified DescribeConfigs request. - final Collection unifiedRequestResources = new ArrayList<>(configResources.size()); + // Partition the requested config resources based on which broker they must be sent to with the + // null broker being used for config resources which can be obtained from any broker + final Map>> brokerFutures = new HashMap<>(configResources.size()); for (ConfigResource resource : configResources) { - if (resource.type() == ConfigResource.Type.BROKER) { - brokerFutures.put(resource, new KafkaFutureImpl()); - brokerResources.add(configResourceToResource(resource)); - } else { - unifiedRequestFutures.put(resource, new KafkaFutureImpl()); - unifiedRequestResources.add(configResourceToResource(resource)); - } + Integer broker = nodeFor(resource); + brokerFutures.compute(broker, (key, value) -> { + if (value == null) { + value = new HashMap<>(); + } + value.put(resource, new KafkaFutureImpl<>()); + return value; + }); } final long now = time.milliseconds(); - if (!unifiedRequestResources.isEmpty()) { + for (Map.Entry>> entry : brokerFutures.entrySet()) { + Integer broker = entry.getKey(); + Map> unified = entry.getValue(); + runnable.call(new Call("describeConfigs", calcDeadlineMs(now, options.timeoutMs()), - new LeastLoadedNodeProvider()) { + broker != null ? new ConstantNodeIdProvider(broker) : new LeastLoadedNodeProvider()) { @Override - AbstractRequest.Builder createRequest(int timeoutMs) { - return new DescribeConfigsRequest.Builder(unifiedRequestResources); + DescribeConfigsRequest.Builder createRequest(int timeoutMs) { + return new DescribeConfigsRequest.Builder(new DescribeConfigsRequestData() + .setResources(unified.keySet().stream() + .map(config -> + new DescribeConfigsRequestData.DescribeConfigsResource() + .setResourceName(config.name()) + .setResourceType(config.type().id()) + .setConfigurationKeys(null)) + .collect(Collectors.toList())) + .setIncludeSynonyms(options.includeSynonyms()) + .setIncludeDocumentation(options.includeDocumentation())); } @Override void handleResponse(AbstractResponse abstractResponse) { DescribeConfigsResponse response = (DescribeConfigsResponse) abstractResponse; - for (Map.Entry> entry : unifiedRequestFutures.entrySet()) { + for (Map.Entry entry : response.resultMap().entrySet()) { ConfigResource configResource = entry.getKey(); - KafkaFutureImpl future = entry.getValue(); - DescribeConfigsResponse.Config config = response.config(configResourceToResource(configResource)); - if (config == null) { - future.completeExceptionally(new UnknownServerException( - "Malformed broker response: missing config for " + configResource)); - continue; - } - if (config.error().isFailure()) { - future.completeExceptionally(config.error().exception()); - continue; - } - List configEntries = new ArrayList<>(); - for (DescribeConfigsResponse.ConfigEntry configEntry : config.entries()) { - configEntries.add(new ConfigEntry(configEntry.name(), configEntry.value(), - configEntry.isDefault(), configEntry.isSensitive(), configEntry.isReadOnly())); + DescribeConfigsResponseData.DescribeConfigsResult describeConfigsResult = entry.getValue(); + KafkaFutureImpl future = unified.get(configResource); + if (future == null) { + if (broker != null) { + log.warn("The config {} in the response from broker {} is not in the request", + configResource, broker); + } else { + log.warn("The config {} in the response from the least loaded broker is not in the request", + configResource); + } + } else { + if (describeConfigsResult.errorCode() != Errors.NONE.code()) { + future.completeExceptionally(Errors.forCode(describeConfigsResult.errorCode()) + .exception(describeConfigsResult.errorMessage())); + } else { + future.complete(describeConfigResult(describeConfigsResult)); + } } - future.complete(new Config(configEntries)); } + completeUnrealizedFutures( + unified.entrySet().stream(), + configResource -> "The broker response did not contain a result for config resource " + configResource); } @Override void handleFailure(Throwable throwable) { - completeAllExceptionally(unifiedRequestFutures.values(), throwable); + completeAllExceptionally(unified.values(), throwable); } }, now); } - for (Map.Entry> entry : brokerFutures.entrySet()) { - final KafkaFutureImpl brokerFuture = entry.getValue(); - final Resource resource = configResourceToResource(entry.getKey()); - final int nodeId = Integer.parseInt(resource.name()); - runnable.call(new Call("describeBrokerConfigs", calcDeadlineMs(now, options.timeoutMs()), - new ConstantNodeIdProvider(nodeId)) { - - @Override - AbstractRequest.Builder createRequest(int timeoutMs) { - return new DescribeConfigsRequest.Builder(Collections.singleton(resource)); - } - - @Override - void handleResponse(AbstractResponse abstractResponse) { - DescribeConfigsResponse response = (DescribeConfigsResponse) abstractResponse; - DescribeConfigsResponse.Config config = response.configs().get(resource); - - if (config == null) { - brokerFuture.completeExceptionally(new UnknownServerException( - "Malformed broker response: missing config for " + resource)); - return; - } - if (config.error().isFailure()) - brokerFuture.completeExceptionally(config.error().exception()); - else { - List configEntries = new ArrayList<>(); - for (DescribeConfigsResponse.ConfigEntry configEntry : config.entries()) { - configEntries.add(new ConfigEntry(configEntry.name(), configEntry.value(), - configEntry.isDefault(), configEntry.isSensitive(), configEntry.isReadOnly())); - } - brokerFuture.complete(new Config(configEntries)); - } - } + return new DescribeConfigsResult(new HashMap<>(brokerFutures.entrySet().stream() + .flatMap(x -> x.getValue().entrySet().stream()) + .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)))); + } - @Override - void handleFailure(Throwable throwable) { - brokerFuture.completeExceptionally(throwable); - } - }, now); - } - final Map> allFutures = new HashMap<>(); - allFutures.putAll(brokerFutures); - allFutures.putAll(unifiedRequestFutures); - return new DescribeConfigsResult(allFutures); + private Config describeConfigResult(DescribeConfigsResponseData.DescribeConfigsResult describeConfigsResult) { + return new Config(describeConfigsResult.configs().stream().map(config -> new ConfigEntry( + config.name(), + config.value(), + DescribeConfigsResponse.ConfigSource.forId(config.configSource()).source(), + config.isSensitive(), + config.readOnly(), + (config.synonyms().stream().map(synonym -> new ConfigEntry.ConfigSynonym(synonym.name(), synonym.value(), + DescribeConfigsResponse.ConfigSource.forId(synonym.source()).source()))).collect(Collectors.toList()), + DescribeConfigsResponse.ConfigType.forId(config.configType()).type(), + config.documentation() + )).collect(Collectors.toList())); } - private Resource configResourceToResource(ConfigResource configResource) { - ResourceType resourceType; - switch (configResource.type()) { - case TOPIC: - resourceType = ResourceType.TOPIC; + private ConfigEntry.ConfigSource configSource(DescribeConfigsResponse.ConfigSource source) { + ConfigEntry.ConfigSource configSource; + switch (source) { + case TOPIC_CONFIG: + configSource = ConfigEntry.ConfigSource.DYNAMIC_TOPIC_CONFIG; + break; + case DYNAMIC_BROKER_CONFIG: + configSource = ConfigEntry.ConfigSource.DYNAMIC_BROKER_CONFIG; + break; + case DYNAMIC_DEFAULT_BROKER_CONFIG: + configSource = ConfigEntry.ConfigSource.DYNAMIC_DEFAULT_BROKER_CONFIG; + break; + case STATIC_BROKER_CONFIG: + configSource = ConfigEntry.ConfigSource.STATIC_BROKER_CONFIG; break; - case BROKER: - resourceType = ResourceType.BROKER; + case DYNAMIC_BROKER_LOGGER_CONFIG: + configSource = ConfigEntry.ConfigSource.DYNAMIC_BROKER_LOGGER_CONFIG; + break; + case DEFAULT_CONFIG: + configSource = ConfigEntry.ConfigSource.DEFAULT_CONFIG; break; default: - throw new IllegalArgumentException("Unexpected resource type " + configResource.type()); + throw new IllegalArgumentException("Unexpected config source " + source); } - return new Resource(resourceType, configResource.name()); + return configSource; } @Override + @Deprecated public AlterConfigsResult alterConfigs(Map configs, final AlterConfigsOptions options) { - final Map> futures = new HashMap<>(configs.size()); - for (ConfigResource configResource : configs.keySet()) { - futures.put(configResource, new KafkaFutureImpl()); + 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()) { + Integer node = nodeFor(resource); + if (node != null) { + NodeProvider nodeProvider = new ConstantNodeIdProvider(node); + allFutures.putAll(alterConfigs(configs, options, Collections.singleton(resource), nodeProvider)); + } else + unifiedRequestResources.add(resource); } - final Map requestMap = new HashMap<>(configs.size()); - for (Map.Entry entry : configs.entrySet()) { + if (!unifiedRequestResources.isEmpty()) + allFutures.putAll(alterConfigs(configs, options, unifiedRequestResources, new LeastLoadedNodeProvider())); + return new AlterConfigsResult(new HashMap<>(allFutures)); + } + + private Map> alterConfigs(Map configs, + final AlterConfigsOptions options, + Collection resources, + NodeProvider nodeProvider) { + final Map> futures = new HashMap<>(); + final Map requestMap = new HashMap<>(resources.size()); + for (ConfigResource resource : resources) { List configEntries = new ArrayList<>(); - for (ConfigEntry configEntry: entry.getValue().entries()) + for (ConfigEntry configEntry: configs.get(resource).entries()) configEntries.add(new AlterConfigsRequest.ConfigEntry(configEntry.name(), configEntry.value())); - ConfigResource resource = entry.getKey(); - requestMap.put(configResourceToResource(resource), new AlterConfigsRequest.Config(configEntries)); + requestMap.put(resource, new AlterConfigsRequest.Config(configEntries)); + futures.put(resource, new KafkaFutureImpl<>()); } final long now = time.milliseconds(); - runnable.call(new Call("alterConfigs", calcDeadlineMs(now, options.timeoutMs()), - new LeastLoadedNodeProvider()) { + runnable.call(new Call("alterConfigs", calcDeadlineMs(now, options.timeoutMs()), nodeProvider) { @Override - public AbstractRequest.Builder createRequest(int timeoutMs) { + public AlterConfigsRequest.Builder createRequest(int timeoutMs) { return new AlterConfigsRequest.Builder(requestMap, options.shouldValidateOnly()); } @@ -1635,7 +2193,69 @@ public void handleResponse(AbstractResponse abstractResponse) { AlterConfigsResponse response = (AlterConfigsResponse) abstractResponse; for (Map.Entry> entry : futures.entrySet()) { KafkaFutureImpl future = entry.getValue(); - ApiException exception = response.errors().get(configResourceToResource(entry.getKey())).exception(); + ApiException exception = response.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; + } + + @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()) { + Integer node = nodeFor(resource); + if (node != null) { + NodeProvider nodeProvider = new ConstantNodeIdProvider(node); + 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 IncrementalAlterConfigsRequest.Builder createRequest(int timeoutMs) { + return new IncrementalAlterConfigsRequest.Builder(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 { @@ -1649,7 +2269,7 @@ void handleFailure(Throwable throwable) { completeAllExceptionally(futures.values(), throwable); } }, now); - return new AlterConfigsResult(new HashMap>(futures)); + return futures; } @Override @@ -1657,98 +2277,136 @@ public AlterReplicaLogDirsResult alterReplicaLogDirs(Map> futures = new HashMap<>(replicaAssignment.size()); for (TopicPartitionReplica replica : replicaAssignment.keySet()) - futures.put(replica, new KafkaFutureImpl()); + futures.put(replica, new KafkaFutureImpl<>()); - Map> replicaAssignmentByBroker = new HashMap<>(); + Map replicaAssignmentByBroker = new HashMap<>(); for (Map.Entry entry: replicaAssignment.entrySet()) { TopicPartitionReplica replica = entry.getKey(); String logDir = entry.getValue(); int brokerId = replica.brokerId(); - TopicPartition topicPartition = new TopicPartition(replica.topic(), replica.partition()); - if (!replicaAssignmentByBroker.containsKey(brokerId)) - replicaAssignmentByBroker.put(brokerId, new HashMap()); - replicaAssignmentByBroker.get(brokerId).put(topicPartition, logDir); + AlterReplicaLogDirsRequestData value = replicaAssignmentByBroker.computeIfAbsent(brokerId, + key -> new AlterReplicaLogDirsRequestData()); + AlterReplicaLogDir alterReplicaLogDir = value.dirs().find(logDir); + if (alterReplicaLogDir == null) { + alterReplicaLogDir = new AlterReplicaLogDir(); + alterReplicaLogDir.setPath(logDir); + value.dirs().add(alterReplicaLogDir); + } + AlterReplicaLogDirTopic alterReplicaLogDirTopic = alterReplicaLogDir.topics().find(replica.topic()); + if (alterReplicaLogDirTopic == null) { + alterReplicaLogDirTopic = new AlterReplicaLogDirTopic().setName(replica.topic()); + alterReplicaLogDir.topics().add(alterReplicaLogDirTopic); + } + alterReplicaLogDirTopic.partitions().add(replica.partition()); } final long now = time.milliseconds(); - for (Map.Entry> entry: replicaAssignmentByBroker.entrySet()) { + for (Map.Entry entry: replicaAssignmentByBroker.entrySet()) { final int brokerId = entry.getKey(); - final Map assignment = entry.getValue(); + final AlterReplicaLogDirsRequestData assignment = entry.getValue(); runnable.call(new Call("alterReplicaLogDirs", calcDeadlineMs(now, options.timeoutMs()), new ConstantNodeIdProvider(brokerId)) { @Override - public AbstractRequest.Builder createRequest(int timeoutMs) { + public AlterReplicaLogDirsRequest.Builder createRequest(int timeoutMs) { return new AlterReplicaLogDirsRequest.Builder(assignment); } @Override public void handleResponse(AbstractResponse abstractResponse) { AlterReplicaLogDirsResponse response = (AlterReplicaLogDirsResponse) abstractResponse; - for (Map.Entry responseEntry: response.responses().entrySet()) { - TopicPartition tp = responseEntry.getKey(); - Errors error = responseEntry.getValue(); - TopicPartitionReplica replica = new TopicPartitionReplica(tp.topic(), tp.partition(), brokerId); - KafkaFutureImpl future = futures.get(replica); - if (future == null) { - handleFailure(new IllegalStateException( - "The partition " + tp + " in the response from broker " + brokerId + " is not in the request")); - } else if (error == Errors.NONE) { - future.complete(null); - } else { - future.completeExceptionally(error.exception()); + for (AlterReplicaLogDirTopicResult topicResult: response.data().results()) { + for (AlterReplicaLogDirPartitionResult partitionResult: topicResult.partitions()) { + TopicPartitionReplica replica = new TopicPartitionReplica( + topicResult.topicName(), partitionResult.partitionIndex(), brokerId); + KafkaFutureImpl future = futures.get(replica); + if (future == null) { + log.warn("The partition {} in the response from broker {} is not in the request", + new TopicPartition(topicResult.topicName(), partitionResult.partitionIndex()), + brokerId); + } else if (partitionResult.errorCode() == Errors.NONE.code()) { + future.complete(null); + } else { + future.completeExceptionally(Errors.forCode(partitionResult.errorCode()).exception()); + } } } + // The server should send back a response for every replica. But do a sanity check anyway. + completeUnrealizedFutures( + futures.entrySet().stream().filter(entry -> entry.getKey().brokerId() == brokerId), + replica -> "The response from broker " + brokerId + + " did not contain a result for replica " + replica); } @Override void handleFailure(Throwable throwable) { - completeAllExceptionally(futures.values(), throwable); + // Only completes the futures of brokerId + completeAllExceptionally( + futures.entrySet().stream() + .filter(entry -> entry.getKey().brokerId() == brokerId) + .map(Map.Entry::getValue), + throwable); } }, now); } - return new AlterReplicaLogDirsResult(new HashMap>(futures)); + return new AlterReplicaLogDirsResult(new HashMap<>(futures)); } @Override public DescribeLogDirsResult describeLogDirs(Collection brokers, DescribeLogDirsOptions options) { - final Map>> futures = new HashMap<>(brokers.size()); - - for (Integer brokerId: brokers) { - futures.put(brokerId, new KafkaFutureImpl>()); - } + final Map>> futures = new HashMap<>(brokers.size()); final long now = time.milliseconds(); - for (final Integer brokerId: brokers) { + for (final Integer brokerId : brokers) { + KafkaFutureImpl> future = new KafkaFutureImpl<>(); + futures.put(brokerId, future); + runnable.call(new Call("describeLogDirs", calcDeadlineMs(now, options.timeoutMs()), new ConstantNodeIdProvider(brokerId)) { @Override - public AbstractRequest.Builder createRequest(int timeoutMs) { + public DescribeLogDirsRequest.Builder createRequest(int timeoutMs) { // Query selected partitions in all log directories - return new DescribeLogDirsRequest.Builder(null); + return new DescribeLogDirsRequest.Builder(new DescribeLogDirsRequestData().setTopics(null)); } + @SuppressWarnings("deprecation") @Override public void handleResponse(AbstractResponse abstractResponse) { DescribeLogDirsResponse response = (DescribeLogDirsResponse) abstractResponse; - KafkaFutureImpl> future = futures.get(brokerId); - if (response.logDirInfos().size() > 0) { - future.complete(response.logDirInfos()); + Map descriptions = logDirDescriptions(response); + if (descriptions.size() > 0) { + future.complete(descriptions); } else { - // response.logDirInfos() will be empty if and only if the user is not authorized to describe clsuter resource. + // descriptions will be empty if and only if the user is not authorized to describe cluster resource. future.completeExceptionally(Errors.CLUSTER_AUTHORIZATION_FAILED.exception()); } } @Override void handleFailure(Throwable throwable) { - completeAllExceptionally(futures.values(), throwable); + future.completeExceptionally(throwable); } }, now); } - return new DescribeLogDirsResult(new HashMap>>(futures)); + return new DescribeLogDirsResult(new HashMap<>(futures)); + } + + private static Map logDirDescriptions(DescribeLogDirsResponse response) { + Map result = new HashMap<>(response.data().results().size()); + for (DescribeLogDirsResponseData.DescribeLogDirsResult logDirResult : response.data().results()) { + Map replicaInfoMap = new HashMap<>(); + for (DescribeLogDirsResponseData.DescribeLogDirsTopic t : logDirResult.topics()) { + for (DescribeLogDirsResponseData.DescribeLogDirsPartition p : t.partitions()) { + replicaInfoMap.put( + new TopicPartition(t.name(), p.partitionIndex()), + new ReplicaInfo(p.partitionSize(), p.offsetLag(), p.isFutureKey())); + } + } + result.put(logDirResult.logDir(), new LogDirDescription(Errors.forCode(logDirResult.errorCode()).exception(), replicaInfoMap)); + } + return result; } @Override @@ -1756,30 +2414,42 @@ public DescribeReplicaLogDirsResult describeReplicaLogDirs(Collection> futures = new HashMap<>(replicas.size()); for (TopicPartitionReplica replica : replicas) { - futures.put(replica, new KafkaFutureImpl()); + futures.put(replica, new KafkaFutureImpl<>()); } - Map> partitionsByBroker = new HashMap<>(); + Map partitionsByBroker = new HashMap<>(); for (TopicPartitionReplica replica: replicas) { - if (!partitionsByBroker.containsKey(replica.brokerId())) - partitionsByBroker.put(replica.brokerId(), new HashSet()); - partitionsByBroker.get(replica.brokerId()).add(new TopicPartition(replica.topic(), replica.partition())); + DescribeLogDirsRequestData requestData = partitionsByBroker.computeIfAbsent(replica.brokerId(), + brokerId -> new DescribeLogDirsRequestData()); + DescribableLogDirTopic describableLogDirTopic = requestData.topics().find(replica.topic()); + if (describableLogDirTopic == null) { + List partitionIndex = new ArrayList<>(); + partitionIndex.add(replica.partition()); + describableLogDirTopic = new DescribableLogDirTopic().setTopic(replica.topic()) + .setPartitionIndex(partitionIndex); + requestData.topics().add(describableLogDirTopic); + } else { + describableLogDirTopic.partitionIndex().add(replica.partition()); + } } final long now = time.milliseconds(); - for (Map.Entry> entry: partitionsByBroker.entrySet()) { + for (Map.Entry entry: partitionsByBroker.entrySet()) { final int brokerId = entry.getKey(); - final Set topicPartitions = entry.getValue(); + final DescribeLogDirsRequestData topicPartitions = entry.getValue(); final Map replicaDirInfoByPartition = new HashMap<>(); - for (TopicPartition topicPartition: topicPartitions) - replicaDirInfoByPartition.put(topicPartition, new ReplicaLogDirInfo()); + for (DescribableLogDirTopic topicPartition: topicPartitions.topics()) { + for (Integer partitionId : topicPartition.partitionIndex()) { + replicaDirInfoByPartition.put(new TopicPartition(topicPartition.topic(), partitionId), new ReplicaLogDirInfo()); + } + } runnable.call(new Call("describeReplicaLogDirs", calcDeadlineMs(now, options.timeoutMs()), new ConstantNodeIdProvider(brokerId)) { @Override - public AbstractRequest.Builder createRequest(int timeoutMs) { + public DescribeLogDirsRequest.Builder createRequest(int timeoutMs) { // Query selected partitions in all log directories return new DescribeLogDirsRequest.Builder(topicPartitions); } @@ -1787,32 +2457,31 @@ public AbstractRequest.Builder createRequest(int timeoutMs) { @Override public void handleResponse(AbstractResponse abstractResponse) { DescribeLogDirsResponse response = (DescribeLogDirsResponse) abstractResponse; - for (Map.Entry responseEntry: response.logDirInfos().entrySet()) { + for (Map.Entry responseEntry: logDirDescriptions(response).entrySet()) { String logDir = responseEntry.getKey(); - DescribeLogDirsResponse.LogDirInfo logDirInfo = responseEntry.getValue(); + LogDirDescription logDirInfo = responseEntry.getValue(); // No replica info will be provided if the log directory is offline - if (logDirInfo.error == Errors.KAFKA_STORAGE_ERROR) + if (logDirInfo.error() instanceof KafkaStorageException) continue; - if (logDirInfo.error != Errors.NONE) + if (logDirInfo.error() != null) handleFailure(new IllegalStateException( - "The error " + logDirInfo.error + " for log directory " + logDir + " in the response from broker " + brokerId + " is illegal")); + "The error " + logDirInfo.error().getClass().getName() + " for log directory " + logDir + " in the response from broker " + brokerId + " is illegal")); - for (Map.Entry replicaInfoEntry: logDirInfo.replicaInfos.entrySet()) { + for (Map.Entry replicaInfoEntry: logDirInfo.replicaInfos().entrySet()) { TopicPartition tp = replicaInfoEntry.getKey(); - DescribeLogDirsResponse.ReplicaInfo replicaInfo = replicaInfoEntry.getValue(); + ReplicaInfo replicaInfo = replicaInfoEntry.getValue(); ReplicaLogDirInfo replicaLogDirInfo = replicaDirInfoByPartition.get(tp); if (replicaLogDirInfo == null) { - handleFailure(new IllegalStateException( - "The partition " + tp + " in the response from broker " + brokerId + " is not in the request")); - } else if (replicaInfo.isFuture) { + log.warn("Server response from broker {} mentioned unknown partition {}", brokerId, tp); + } else if (replicaInfo.isFuture()) { replicaDirInfoByPartition.put(tp, new ReplicaLogDirInfo(replicaLogDirInfo.getCurrentReplicaLogDir(), replicaLogDirInfo.getCurrentReplicaOffsetLag(), logDir, - replicaInfo.offsetLag)); + replicaInfo.offsetLag())); } else { replicaDirInfoByPartition.put(tp, new ReplicaLogDirInfo(logDir, - replicaInfo.offsetLag, + replicaInfo.offsetLag(), replicaLogDirInfo.getFutureReplicaLogDir(), replicaLogDirInfo.getFutureReplicaOffsetLag())); } @@ -1832,47 +2501,112 @@ void handleFailure(Throwable throwable) { }, now); } - return new DescribeReplicaLogDirsResult(new HashMap>(futures)); + return new DescribeReplicaLogDirsResult(new HashMap<>(futures)); } - public CreatePartitionsResult createPartitions(Map newPartitions, + @Override + public CreatePartitionsResult createPartitions(final Map newPartitions, final CreatePartitionsOptions options) { final Map> futures = new HashMap<>(newPartitions.size()); - for (String topic : newPartitions.keySet()) { - futures.put(topic, new KafkaFutureImpl()); + final CreatePartitionsTopicCollection topics = new CreatePartitionsTopicCollection(newPartitions.size()); + for (Map.Entry entry : newPartitions.entrySet()) { + final String topic = entry.getKey(); + final NewPartitions newPartition = entry.getValue(); + List> newAssignments = newPartition.assignments(); + List assignments = newAssignments == null ? null : + newAssignments.stream() + .map(brokerIds -> new CreatePartitionsAssignment().setBrokerIds(brokerIds)) + .collect(Collectors.toList()); + topics.add(new CreatePartitionsTopic() + .setName(topic) + .setCount(newPartition.totalCount()) + .setAssignments(assignments)); + futures.put(topic, new KafkaFutureImpl<>()); } - final Map requestMap = new HashMap<>(newPartitions); - - final long now = time.milliseconds(); - runnable.call(new Call("createPartitions", calcDeadlineMs(now, options.timeoutMs()), - new ControllerNodeProvider()) { + if (!topics.isEmpty()) { + final long now = time.milliseconds(); + final long deadline = calcDeadlineMs(now, options.timeoutMs()); + final Call call = getCreatePartitionsCall(options, futures, topics, + Collections.emptyMap(), now, deadline); + runnable.call(call, now); + } + return new CreatePartitionsResult(new HashMap<>(futures)); + } + private Call getCreatePartitionsCall(final CreatePartitionsOptions options, + final Map> futures, + final CreatePartitionsTopicCollection topics, + final Map quotaExceededExceptions, + final long now, + final long deadline) { + return new Call("createPartitions", deadline, new ControllerNodeProvider()) { @Override - public AbstractRequest.Builder createRequest(int timeoutMs) { - return new CreatePartitionsRequest.Builder(requestMap, timeoutMs, options.validateOnly()); + public CreatePartitionsRequest.Builder createRequest(int timeoutMs) { + return new CreatePartitionsRequest.Builder( + new CreatePartitionsRequestData() + .setTopics(topics) + .setValidateOnly(options.validateOnly()) + .setTimeoutMs(timeoutMs)); } @Override public void handleResponse(AbstractResponse abstractResponse) { - CreatePartitionsResponse response = (CreatePartitionsResponse) abstractResponse; - for (Map.Entry result : response.errors().entrySet()) { - KafkaFutureImpl future = futures.get(result.getKey()); - if (result.getValue().isSuccess()) { - future.complete(null); + // Check for controller change + handleNotControllerError(abstractResponse); + // Handle server responses for particular topics. + final CreatePartitionsResponse response = (CreatePartitionsResponse) abstractResponse; + final CreatePartitionsTopicCollection retryTopics = new CreatePartitionsTopicCollection(); + final Map retryTopicQuotaExceededExceptions = new HashMap<>(); + for (CreatePartitionsTopicResult result : response.data().results()) { + KafkaFutureImpl future = futures.get(result.name()); + if (future == null) { + log.warn("Server response mentioned unknown topic {}", result.name()); } else { - future.completeExceptionally(result.getValue().exception()); + ApiError error = new ApiError(result.errorCode(), result.errorMessage()); + if (error.isFailure()) { + if (error.is(Errors.THROTTLING_QUOTA_EXCEEDED)) { + ThrottlingQuotaExceededException quotaExceededException = new ThrottlingQuotaExceededException( + response.throttleTimeMs(), error.messageWithFallback()); + if (options.shouldRetryOnQuotaViolation()) { + retryTopics.add(topics.find(result.name()).duplicate()); + retryTopicQuotaExceededExceptions.put(result.name(), quotaExceededException); + } else { + future.completeExceptionally(quotaExceededException); + } + } else { + future.completeExceptionally(error.exception()); + } + } else { + future.complete(null); + } } } + // If there are topics to retry, retry them; complete unrealized futures otherwise. + if (retryTopics.isEmpty()) { + // The server should send back a response for every topic. But do a sanity check anyway. + completeUnrealizedFutures(futures.entrySet().stream(), + topic -> "The controller response did not contain a result for topic " + topic); + } else { + final long now = time.milliseconds(); + final Call call = getCreatePartitionsCall(options, futures, retryTopics, + retryTopicQuotaExceededExceptions, now, deadline); + runnable.call(call, now); + } } @Override void handleFailure(Throwable throwable) { + // If there were any topics retries due to a quota exceeded exception, we propagate + // the initial error back to the caller if the request timed out. + maybeCompleteQuotaExceededException(options.shouldRetryOnQuotaViolation(), + throwable, futures, quotaExceededExceptions, (int) (time.milliseconds() - now)); + // Fail all the other remaining futures completeAllExceptionally(futures.values(), throwable); } - }, now); - return new CreatePartitionsResult(new HashMap>(futures)); + }; } + @Override public DeleteRecordsResult deleteRecords(final Map recordsToDelete, final DeleteRecordsOptions options) { @@ -1881,7 +2615,7 @@ public DeleteRecordsResult deleteRecords(final Map> futures = new HashMap<>(recordsToDelete.size()); for (TopicPartition topicPartition: recordsToDelete.keySet()) { - futures.put(topicPartition, new KafkaFutureImpl()); + futures.put(topicPartition, new KafkaFutureImpl<>()); } // preparing topics list for asking metadata about them @@ -1897,8 +2631,10 @@ public DeleteRecordsResult deleteRecords(final Map(topics), false); + MetadataRequest.Builder createRequest(int timeoutMs) { + return new MetadataRequest.Builder(new MetadataRequestData() + .setTopics(convertToMetadataRequestTopic(topics)) + .setAllowAutoTopicCreation(false)); } @Override @@ -1908,68 +2644,78 @@ 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 - Map> leaders = new HashMap<>(); + // Group topic partitions by leader + Map> leaders = new HashMap<>(); for (Map.Entry entry: recordsToDelete.entrySet()) { + TopicPartition topicPartition = entry.getKey(); + KafkaFutureImpl future = futures.get(topicPartition); - // avoiding to send deletion request for topics with errors - if (!errors.containsKey(entry.getKey().topic())) { - - Node node = cluster.leaderFor(entry.getKey()); + // Fail partitions with topic errors + Errors topicError = errors.get(topicPartition.topic()); + if (errors.containsKey(topicPartition.topic())) { + future.completeExceptionally(topicError.exception()); + } else { + Node node = cluster.leaderFor(topicPartition); if (node != null) { - if (!leaders.containsKey(node)) - leaders.put(node, new HashMap()); - leaders.get(node).put(entry.getKey(), entry.getValue().beforeOffset()); + Map deletionsForLeader = leaders.computeIfAbsent( + node, key -> new HashMap<>()); + DeleteRecordsTopic deleteRecords = deletionsForLeader.get(topicPartition.topic()); + if (deleteRecords == null) { + deleteRecords = new DeleteRecordsTopic() + .setName(topicPartition.topic()); + deletionsForLeader.put(topicPartition.topic(), deleteRecords); + } + deleteRecords.partitions().add(new DeleteRecordsPartition() + .setPartitionIndex(topicPartition.partition()) + .setOffset(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, new ConstantNodeIdProvider(brokerId)) { @Override - AbstractRequest.Builder createRequest(int timeoutMs) { - return new DeleteRecordsRequest.Builder(timeoutMs, entry.getValue()); + DeleteRecordsRequest.Builder createRequest(int timeoutMs) { + return new DeleteRecordsRequest.Builder(new DeleteRecordsRequestData() + .setTimeoutMs(timeoutMs) + .setTopics(new ArrayList<>(partitionDeleteOffsets.values()))); } @Override void handleResponse(AbstractResponse abstractResponse) { DeleteRecordsResponse response = (DeleteRecordsResponse) abstractResponse; - for (Map.Entry result: response.responses().entrySet()) { - - KafkaFutureImpl future = futures.get(result.getKey()); - if (result.getValue().error == Errors.NONE) { - future.complete(new DeletedRecords(result.getValue().lowWatermark)); - } else { - future.completeExceptionally(result.getValue().error.exception()); + for (DeleteRecordsTopicResult topicResult: response.data().topics()) { + for (DeleteRecordsResponseData.DeleteRecordsPartitionResult partitionResult : topicResult.partitions()) { + KafkaFutureImpl future = futures.get(new TopicPartition(topicResult.name(), partitionResult.partitionIndex())); + if (partitionResult.errorCode() == Errors.NONE.code()) { + future.complete(new DeletedRecords(partitionResult.lowWatermark())); + } else { + future.completeExceptionally(Errors.forCode(partitionResult.errorCode()).exception()); + } } } } @Override void handleFailure(Throwable throwable) { - completeAllExceptionally(futures.values(), throwable); + Stream> callFutures = + partitionDeleteOffsets.values().stream().flatMap( + recordsToDelete -> + recordsToDelete.partitions().stream().map(partitionsToDelete -> + new TopicPartition(recordsToDelete.name(), partitionsToDelete.partitionIndex())) + ).map(futures::get); + completeAllExceptionally(callFutures, throwable); } - }, nowDelete); + }, deleteRecordsCallTimeMs); } } @@ -1979,6 +2725,1753 @@ void handleFailure(Throwable throwable) { } }, nowMetadata); - return new DeleteRecordsResult(new HashMap>(futures)); + return new DeleteRecordsResult(new HashMap<>(futures)); + } + + @Override + 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 + CreateDelegationTokenRequest.Builder createRequest(int timeoutMs) { + return new CreateDelegationTokenRequest.Builder( + new CreateDelegationTokenRequestData() + .setRenewers(renewers) + .setMaxLifetimeMs(options.maxlifeTimeMs())); + } + + @Override + void handleResponse(AbstractResponse abstractResponse) { + CreateDelegationTokenResponse response = (CreateDelegationTokenResponse) abstractResponse; + if (response.hasError()) { + delegationTokenFuture.completeExceptionally(response.error().exception()); + } else { + 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); + } + } + + @Override + void handleFailure(Throwable throwable) { + delegationTokenFuture.completeExceptionally(throwable); + } + }, now); + + return new CreateDelegationTokenResult(delegationTokenFuture); + } + + @Override + public RenewDelegationTokenResult renewDelegationToken(final byte[] hmac, final RenewDelegationTokenOptions options) { + final KafkaFutureImpl expiryTimeFuture = new KafkaFutureImpl<>(); + final long now = time.milliseconds(); + runnable.call(new Call("renewDelegationToken", calcDeadlineMs(now, options.timeoutMs()), + new LeastLoadedNodeProvider()) { + + @Override + RenewDelegationTokenRequest.Builder createRequest(int timeoutMs) { + return new RenewDelegationTokenRequest.Builder( + new RenewDelegationTokenRequestData() + .setHmac(hmac) + .setRenewPeriodMs(options.renewTimePeriodMs())); + } + + @Override + void handleResponse(AbstractResponse abstractResponse) { + RenewDelegationTokenResponse response = (RenewDelegationTokenResponse) abstractResponse; + if (response.hasError()) { + expiryTimeFuture.completeExceptionally(response.error().exception()); + } else { + expiryTimeFuture.complete(response.expiryTimestamp()); + } + } + + @Override + void handleFailure(Throwable throwable) { + expiryTimeFuture.completeExceptionally(throwable); + } + }, now); + + return new RenewDelegationTokenResult(expiryTimeFuture); + } + + @Override + public ExpireDelegationTokenResult expireDelegationToken(final byte[] hmac, final ExpireDelegationTokenOptions options) { + final KafkaFutureImpl expiryTimeFuture = new KafkaFutureImpl<>(); + final long now = time.milliseconds(); + runnable.call(new Call("expireDelegationToken", calcDeadlineMs(now, options.timeoutMs()), + new LeastLoadedNodeProvider()) { + + @Override + ExpireDelegationTokenRequest.Builder createRequest(int timeoutMs) { + return new ExpireDelegationTokenRequest.Builder( + new ExpireDelegationTokenRequestData() + .setHmac(hmac) + .setExpiryTimePeriodMs(options.expiryTimePeriodMs())); + } + + @Override + void handleResponse(AbstractResponse abstractResponse) { + ExpireDelegationTokenResponse response = (ExpireDelegationTokenResponse) abstractResponse; + if (response.hasError()) { + expiryTimeFuture.completeExceptionally(response.error().exception()); + } else { + expiryTimeFuture.complete(response.expiryTimestamp()); + } + } + + @Override + void handleFailure(Throwable throwable) { + expiryTimeFuture.completeExceptionally(throwable); + } + }, now); + + return new ExpireDelegationTokenResult(expiryTimeFuture); + } + + @Override + public DescribeDelegationTokenResult describeDelegationToken(final DescribeDelegationTokenOptions options) { + final KafkaFutureImpl> tokensFuture = new KafkaFutureImpl<>(); + final long now = time.milliseconds(); + runnable.call(new Call("describeDelegationToken", calcDeadlineMs(now, options.timeoutMs()), + new LeastLoadedNodeProvider()) { + + @Override + DescribeDelegationTokenRequest.Builder createRequest(int timeoutMs) { + return new DescribeDelegationTokenRequest.Builder(options.owners()); + } + + @Override + void handleResponse(AbstractResponse abstractResponse) { + DescribeDelegationTokenResponse response = (DescribeDelegationTokenResponse) abstractResponse; + if (response.hasError()) { + tokensFuture.completeExceptionally(response.error().exception()); + } else { + tokensFuture.complete(response.tokens()); + } + } + + @Override + void handleFailure(Throwable throwable) { + tokensFuture.completeExceptionally(throwable); + } + }, now); + + return new DescribeDelegationTokenResult(tokensFuture); + } + + private void rescheduleFindCoordinatorTask(ConsumerGroupOperationContext context, Supplier nextCall, Call failedCall) { + log.info("Node {} is no longer the Coordinator. Retrying with new coordinator.", + context.node().orElse(null)); + // Requeue the task so that we can try with new coordinator + context.setNode(null); + + Call call = nextCall.get(); + call.tries = failedCall.tries + 1; + call.nextAllowedTryMs = calculateNextAllowedRetryMs(); + + Call findCoordinatorCall = getFindCoordinatorCall(context, nextCall); + runnable.call(findCoordinatorCall, time.milliseconds()); + } + + private void rescheduleMetadataTask(MetadataOperationContext context, Supplier> nextCalls) { + log.info("Retrying to fetch metadata."); + // Requeue the task so that we can re-attempt fetching metadata + context.setResponse(Optional.empty()); + Call metadataCall = getMetadataCall(context, nextCalls); + runnable.call(metadataCall, 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 = 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 Map.Entry> entry : futures.entrySet()) { + // skip sending request for those futures that already failed. + if (entry.getValue().isCompletedExceptionally()) + continue; + + final String groupId = entry.getKey(); + + final long startFindCoordinatorMs = time.milliseconds(); + final long deadline = calcDeadlineMs(startFindCoordinatorMs, options.timeoutMs()); + ConsumerGroupOperationContext context = + new ConsumerGroupOperationContext<>(groupId, options, deadline, futures.get(groupId)); + Call findCoordinatorCall = getFindCoordinatorCall(context, + () -> getDescribeConsumerGroupsCall(context)); + runnable.call(findCoordinatorCall, startFindCoordinatorMs); + } + + return new DescribeConsumerGroupsResult(new HashMap<>(futures)); + } + + /** + * 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.deadline(), new LeastLoadedNodeProvider()) { + @Override + FindCoordinatorRequest.Builder createRequest(int timeoutMs) { + return new FindCoordinatorRequest.Builder( + new FindCoordinatorRequestData() + .setKeyType(CoordinatorType.GROUP.id()) + .setKey(context.groupId())); + } + + @Override + void handleResponse(AbstractResponse abstractResponse) { + final FindCoordinatorResponse response = (FindCoordinatorResponse) abstractResponse; + + if (handleGroupRequestError(response.error(), context.future())) + return; + + context.setNode(response.node()); + + runnable.call(nextCall.get(), time.milliseconds()); + } + + @Override + void handleFailure(Throwable throwable) { + context.future().completeExceptionally(throwable); + } + }; + } + + private Call getDescribeConsumerGroupsCall( + ConsumerGroupOperationContext context) { + return new Call("describeConsumerGroups", + context.deadline(), + new ConstantNodeIdProvider(context.node().get().id())) { + @Override + DescribeGroupsRequest.Builder createRequest(int timeoutMs) { + return new DescribeGroupsRequest.Builder( + new DescribeGroupsRequestData() + .setGroups(Collections.singletonList(context.groupId())) + .setIncludeAuthorizedOperations(context.options().includeAuthorizedOperations())); + } + + @Override + void handleResponse(AbstractResponse abstractResponse) { + final DescribeGroupsResponse response = (DescribeGroupsResponse) abstractResponse; + + List describedGroups = response.data().groups(); + if (describedGroups.isEmpty()) { + context.future().completeExceptionally( + new InvalidGroupIdException("No consumer group found for GroupId: " + context.groupId())); + return; + } + + if (describedGroups.size() > 1 || + !describedGroups.get(0).groupId().equals(context.groupId())) { + String ids = Arrays.toString(describedGroups.stream().map(DescribedGroup::groupId).toArray()); + context.future().completeExceptionally(new InvalidGroupIdException( + "DescribeConsumerGroup request for GroupId: " + context.groupId() + " returned " + ids)); + return; + } + + final DescribedGroup describedGroup = describedGroups.get(0); + + // If coordinator changed since we fetched it, retry + if (ConsumerGroupOperationContext.hasCoordinatorMoved(response)) { + Call call = getDescribeConsumerGroupsCall(context); + rescheduleFindCoordinatorTask(context, () -> call, this); + return; + } + + final Errors groupError = Errors.forCode(describedGroup.errorCode()); + if (handleGroupRequestError(groupError, context.future())) + 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.groupId(), protocolType.isEmpty(), + memberDescriptions, + describedGroup.protocolData(), + ConsumerGroupState.parse(describedGroup.groupState()), + context.node().get(), + authorizedOperations); + context.future().complete(consumerGroupDescription); + } else { + context.future().completeExceptionally(new IllegalArgumentException( + String.format("GroupId %s is not a consumer group (%s).", + context.groupId(), protocolType))); + } + } + + @Override + void handleFailure(Throwable throwable) { + context.future().completeExceptionally(throwable); + } + }; + } + + /** + * Returns a {@code Call} object to fetch the cluster metadata. Takes a List of Calls + * parameter to schedule actions that need to be taken using the metadata. The param is a Supplier + * so that it can be lazily created, so that it can use the results of the metadata call in its + * construction. + * + * @param The type of return value of the KafkaFuture, like ListOffsetsResultInfo, etc. + * @param The type of configuration option, like ListOffsetsOptions, etc + */ + private > Call getMetadataCall(MetadataOperationContext context, + Supplier> nextCalls) { + return new Call("metadata", context.deadline(), new LeastLoadedNodeProvider()) { + @Override + MetadataRequest.Builder createRequest(int timeoutMs) { + return new MetadataRequest.Builder(new MetadataRequestData() + .setTopics(convertToMetadataRequestTopic(context.topics())) + .setAllowAutoTopicCreation(false)); + } + + @Override + void handleResponse(AbstractResponse abstractResponse) { + MetadataResponse response = (MetadataResponse) abstractResponse; + MetadataOperationContext.handleMetadataErrors(response); + + context.setResponse(Optional.of(response)); + + for (Call call : nextCalls.get()) { + runnable.call(call, time.milliseconds()); + } + } + + @Override + void handleFailure(Throwable throwable) { + for (KafkaFutureImpl future : context.futures().values()) { + future.completeExceptionally(throwable); + } + } + }; + } + + private Set validAclOperations(final int authorizedOperations) { + if (authorizedOperations == MetadataResponse.AUTHORIZED_OPERATIONS_OMITTED) { + return null; + } + return Utils.from32BitField(authorizedOperations) + .stream() + .map(AclOperation::fromCode) + .filter(operation -> operation != AclOperation.UNKNOWN + && operation != AclOperation.ALL + && operation != AclOperation.ANY) + .collect(Collectors.toSet()); + } + + private boolean handleGroupRequestError(Errors error, KafkaFutureImpl future) { + if (error == Errors.COORDINATOR_LOAD_IN_PROGRESS || error == Errors.COORDINATOR_NOT_AVAILABLE) { + throw error.exception(); + } else if (error != Errors.NONE) { + future.completeExceptionally(error.exception()); + return true; + } + return false; + } + + private final static class ListConsumerGroupsResults { + private final List errors; + private final HashMap listings; + private final HashSet remaining; + private final KafkaFutureImpl> future; + + ListConsumerGroupsResults(Collection leaders, + KafkaFutureImpl> future) { + this.errors = new ArrayList<>(); + this.listings = new HashMap<>(); + this.remaining = new HashSet<>(leaders); + this.future = future; + tryComplete(); + } + + synchronized void addError(Throwable throwable, Node node) { + ApiError error = ApiError.fromThrowable(throwable); + if (error.message() == null || error.message().isEmpty()) { + errors.add(error.error().exception("Error listing groups on " + node)); + } else { + errors.add(error.error().exception("Error listing groups on " + node + ": " + error.message())); + } + } + + synchronized void addListing(ConsumerGroupListing listing) { + listings.put(listing.groupId(), listing); + } + + synchronized void tryComplete(Node leader) { + remaining.remove(leader); + tryComplete(); + } + + private synchronized void tryComplete() { + if (remaining.isEmpty()) { + ArrayList results = new ArrayList<>(listings.values()); + results.addAll(errors); + future.complete(results); + } + } + } + + @Override + public ListConsumerGroupsResult listConsumerGroups(ListConsumerGroupsOptions options) { + final KafkaFutureImpl> all = new KafkaFutureImpl<>(); + final long nowMetadata = time.milliseconds(); + final long deadline = calcDeadlineMs(nowMetadata, options.timeoutMs()); + runnable.call(new Call("findAllBrokers", deadline, new LeastLoadedNodeProvider()) { + @Override + MetadataRequest.Builder createRequest(int timeoutMs) { + return new MetadataRequest.Builder(new MetadataRequestData() + .setTopics(Collections.emptyList()) + .setAllowAutoTopicCreation(true)); + } + + @Override + void handleResponse(AbstractResponse abstractResponse) { + MetadataResponse metadataResponse = (MetadataResponse) abstractResponse; + Collection nodes = metadataResponse.brokers(); + if (nodes.isEmpty()) + throw new StaleMetadataException("Metadata fetch failed due to missing broker list"); + + HashSet allNodes = new HashSet<>(nodes); + final ListConsumerGroupsResults results = new ListConsumerGroupsResults(allNodes, all); + + for (final Node node : allNodes) { + final long nowList = time.milliseconds(); + runnable.call(new Call("listConsumerGroups", deadline, new ConstantNodeIdProvider(node.id())) { + @Override + ListGroupsRequest.Builder createRequest(int timeoutMs) { + List states = options.states() + .stream() + .map(s -> s.toString()) + .collect(Collectors.toList()); + return new ListGroupsRequest.Builder(new ListGroupsRequestData().setStatesFilter(states)); + } + + private void maybeAddConsumerGroup(ListGroupsResponseData.ListedGroup group) { + String protocolType = group.protocolType(); + if (protocolType.equals(ConsumerProtocol.PROTOCOL_TYPE) || protocolType.isEmpty()) { + final String groupId = group.groupId(); + final Optional state = group.groupState().equals("") + ? Optional.empty() + : Optional.of(ConsumerGroupState.parse(group.groupState())); + final ConsumerGroupListing groupListing = new ConsumerGroupListing(groupId, protocolType.isEmpty(), state); + results.addListing(groupListing); + } + } + + @Override + void handleResponse(AbstractResponse abstractResponse) { + final ListGroupsResponse response = (ListGroupsResponse) abstractResponse; + synchronized (results) { + 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 (ListGroupsResponseData.ListedGroup group : response.data().groups()) { + maybeAddConsumerGroup(group); + } + } + results.tryComplete(node); + } + } + + @Override + void handleFailure(Throwable throwable) { + synchronized (results) { + results.addError(throwable, node); + results.tryComplete(node); + } + } + }, nowList); + } + } + + @Override + void handleFailure(Throwable throwable) { + KafkaException exception = new KafkaException("Failed to find brokers to send ListGroups", throwable); + all.complete(Collections.singletonList(exception)); + } + }, nowMetadata); + + return new ListConsumerGroupsResult(all); + } + + @Override + 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()); + + 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.deadline(), + new ConstantNodeIdProvider(context.node().get().id())) { + @Override + OffsetFetchRequest.Builder createRequest(int timeoutMs) { + // Set the flag to false as for admin client request, + // we don't need to wait for any pending offset state to clear. + return new OffsetFetchRequest.Builder(context.groupId(), false, context.options().topicPartitions(), false); + } + + @Override + void handleResponse(AbstractResponse abstractResponse) { + final OffsetFetchResponse response = (OffsetFetchResponse) abstractResponse; + final Map groupOffsetsListing = new HashMap<>(); + + // If coordinator changed since we fetched it, retry + if (ConsumerGroupOperationContext.hasCoordinatorMoved(response)) { + Call call = getListConsumerGroupOffsetsCall(context); + rescheduleFindCoordinatorTask(context, () -> call, this); + return; + } + + if (handleGroupRequestError(response.error(), context.future())) + return; + + 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 { + groupOffsetsListing.put(topicPartition, new OffsetAndMetadata(offset, leaderEpoch, metadata)); + } + } else { + log.warn("Skipping return offset for {} due to error {}.", topicPartition, error); + } + } + context.future().complete(groupOffsetsListing); + } + + @Override + void handleFailure(Throwable throwable) { + context.future().completeExceptionally(throwable); + } + }; + } + + @Override + public DeleteConsumerGroupsResult deleteConsumerGroups(Collection groupIds, DeleteConsumerGroupsOptions options) { + + 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. + 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); + } + + return new DeleteConsumerGroupsResult(new HashMap<>(futures)); + } + + private Call getDeleteConsumerGroupsCall(ConsumerGroupOperationContext context) { + return new Call("deleteConsumerGroups", context.deadline(), new ConstantNodeIdProvider(context.node().get().id())) { + + @Override + DeleteGroupsRequest.Builder createRequest(int timeoutMs) { + return new DeleteGroupsRequest.Builder( + new DeleteGroupsRequestData() + .setGroupsNames(Collections.singletonList(context.groupId())) + ); + } + + @Override + void handleResponse(AbstractResponse abstractResponse) { + final DeleteGroupsResponse response = (DeleteGroupsResponse) abstractResponse; + + // If coordinator changed since we fetched it, retry + if (ConsumerGroupOperationContext.hasCoordinatorMoved(response)) { + Call call = getDeleteConsumerGroupsCall(context); + rescheduleFindCoordinatorTask(context, () -> call, this); + return; + } + + final Errors groupError = response.get(context.groupId()); + if (handleGroupRequestError(groupError, context.future())) + return; + + context.future().complete(null); + } + + @Override + void handleFailure(Throwable throwable) { + context.future().completeExceptionally(throwable); + } + }; + } + + @Override + public DeleteConsumerGroupOffsetsResult deleteConsumerGroupOffsets( + String groupId, + Set partitions, + DeleteConsumerGroupOffsetsOptions options) { + final KafkaFutureImpl> future = new KafkaFutureImpl<>(); + + if (groupIdIsUnrepresentable(groupId)) { + future.completeExceptionally(new InvalidGroupIdException("The given group id '" + + groupId + "' cannot be represented in a request.")); + return new DeleteConsumerGroupOffsetsResult(future, partitions); + } + + final long startFindCoordinatorMs = time.milliseconds(); + final long deadline = calcDeadlineMs(startFindCoordinatorMs, options.timeoutMs()); + ConsumerGroupOperationContext, DeleteConsumerGroupOffsetsOptions> context = + new ConsumerGroupOperationContext<>(groupId, options, deadline, future); + + Call findCoordinatorCall = getFindCoordinatorCall(context, + () -> getDeleteConsumerGroupOffsetsCall(context, partitions)); + runnable.call(findCoordinatorCall, startFindCoordinatorMs); + + return new DeleteConsumerGroupOffsetsResult(future, partitions); + } + + private Call getDeleteConsumerGroupOffsetsCall( + ConsumerGroupOperationContext, DeleteConsumerGroupOffsetsOptions> context, + Set partitions) { + return new Call("deleteConsumerGroupOffsets", context.deadline(), new ConstantNodeIdProvider(context.node().get().id())) { + + @Override + OffsetDeleteRequest.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 handleResponse(AbstractResponse abstractResponse) { + final OffsetDeleteResponse response = (OffsetDeleteResponse) abstractResponse; + + // If coordinator changed since we fetched it, retry + if (ConsumerGroupOperationContext.hasCoordinatorMoved(response)) { + Call call = getDeleteConsumerGroupOffsetsCall(context, partitions); + rescheduleFindCoordinatorTask(context, () -> call, this); + return; + } + + // 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.future())) + 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.future().complete(partitions); + } + + @Override + void handleFailure(Throwable throwable) { + context.future().completeExceptionally(throwable); + } + }; + } + + @Override + public Map metrics() { + return Collections.unmodifiableMap(this.metrics.metrics()); + } + + @Override + 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("electLeaders", calcDeadlineMs(now, options.timeoutMs()), + new ControllerNodeProvider()) { + + @Override + public ElectLeadersRequest.Builder createRequest(int timeoutMs) { + return new ElectLeadersRequest.Builder(electionType, topicPartitions, timeoutMs); + } + + @Override + public void handleResponse(AbstractResponse abstractResponse) { + 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 + void handleFailure(Throwable throwable) { + electionFuture.completeExceptionally(throwable); + } + }, now); + + 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 AlterPartitionReassignmentsRequest.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, response.data().errorMessage()).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 + ListPartitionReassignmentsRequest.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 long calculateNextAllowedRetryMs() { + return time.milliseconds() + retryBackoffMs; + } + + private void handleNotControllerError(AbstractResponse response) throws ApiException { + if (response.errorCounts().containsKey(Errors.NOT_CONTROLLER)) { + handleNotControllerError(Errors.NOT_CONTROLLER); + } + } + + private void handleNotControllerError(Errors error) throws ApiException { + metadataManager.clearController(); + metadataManager.requestUpdate(); + throw error.exception(); + } + + /** + * Returns the broker id pertaining to the given resource, or null if the resource is not associated + * with a particular broker. + */ + private Integer nodeFor(ConfigResource resource) { + if ((resource.type() == ConfigResource.Type.BROKER && !resource.isDefault()) + || resource.type() == ConfigResource.Type.BROKER_LOGGER) { + return Integer.valueOf(resource.name()); + } else { + return null; + } + } + + private List getMembersFromGroup(String groupId) { + Collection members; + try { + members = describeConsumerGroups(Collections.singleton(groupId)).describedGroups().get(groupId).get().members(); + } catch (Exception ex) { + throw new KafkaException("Encounter exception when trying to get members from group: " + groupId, ex); + } + + List membersToRemove = new ArrayList<>(); + for (final MemberDescription member : members) { + if (member.groupInstanceId().isPresent()) { + membersToRemove.add(new MemberIdentity().setGroupInstanceId(member.groupInstanceId().get())); + } else { + membersToRemove.add(new MemberIdentity().setMemberId(member.consumerId())); + } + } + return membersToRemove; + } + + @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); + + List members; + if (options.removeAll()) { + members = getMembersFromGroup(groupId); + } else { + members = options.members().stream().map(MemberToRemove::toMemberIdentity).collect(Collectors.toList()); + } + Call findCoordinatorCall = getFindCoordinatorCall(context, () -> getRemoveMembersFromGroupCall(context, members)); + runnable.call(findCoordinatorCall, startFindCoordinatorMs); + + return new RemoveMembersFromConsumerGroupResult(future, options.members()); + } + + private Call getRemoveMembersFromGroupCall(ConsumerGroupOperationContext, RemoveMembersFromConsumerGroupOptions> context, + List members) { + return new Call("leaveGroup", context.deadline(), new ConstantNodeIdProvider(context.node().get().id())) { + @Override + LeaveGroupRequest.Builder createRequest(int timeoutMs) { + return new LeaveGroupRequest.Builder(context.groupId(), members); + } + + @Override + void handleResponse(AbstractResponse abstractResponse) { + final LeaveGroupResponse response = (LeaveGroupResponse) abstractResponse; + + // If coordinator changed since we fetched it, retry + if (ConsumerGroupOperationContext.hasCoordinatorMoved(response)) { + Call call = getRemoveMembersFromGroupCall(context, members); + rescheduleFindCoordinatorTask(context, () -> call, this); + return; + } + + if (handleGroupRequestError(response.topLevelError(), context.future())) + return; + + final Map memberErrors = new HashMap<>(); + for (MemberResponse memberResponse : response.memberResponses()) { + memberErrors.put(new MemberIdentity() + .setMemberId(memberResponse.memberId()) + .setGroupInstanceId(memberResponse.groupInstanceId()), + Errors.forCode(memberResponse.errorCode())); + } + context.future().complete(memberErrors); + } + + @Override + void handleFailure(Throwable throwable) { + context.future().completeExceptionally(throwable); + } + }; + } + + @Override + public AlterConsumerGroupOffsetsResult alterConsumerGroupOffsets(String groupId, + Map offsets, + AlterConsumerGroupOffsetsOptions options) { + final KafkaFutureImpl> future = new KafkaFutureImpl<>(); + + final long startFindCoordinatorMs = time.milliseconds(); + final long deadline = calcDeadlineMs(startFindCoordinatorMs, options.timeoutMs()); + + ConsumerGroupOperationContext, AlterConsumerGroupOffsetsOptions> context = + new ConsumerGroupOperationContext<>(groupId, options, deadline, future); + + Call findCoordinatorCall = getFindCoordinatorCall(context, + () -> KafkaAdminClient.this.getAlterConsumerGroupOffsetsCall(context, offsets)); + runnable.call(findCoordinatorCall, startFindCoordinatorMs); + + return new AlterConsumerGroupOffsetsResult(future); + } + + private Call getAlterConsumerGroupOffsetsCall(ConsumerGroupOperationContext, + AlterConsumerGroupOffsetsOptions> context, + Map offsets) { + + return new Call("commitOffsets", context.deadline(), new ConstantNodeIdProvider(context.node().get().id())) { + + @Override + OffsetCommitRequest.Builder createRequest(int timeoutMs) { + List topics = new ArrayList<>(); + Map> offsetData = new HashMap<>(); + for (Map.Entry entry : offsets.entrySet()) { + String topic = entry.getKey().topic(); + OffsetAndMetadata oam = entry.getValue(); + offsetData.compute(topic, (key, value) -> { + if (value == null) { + value = new ArrayList<>(); + } + OffsetCommitRequestPartition partition = new OffsetCommitRequestPartition() + .setCommittedOffset(oam.offset()) + .setCommittedLeaderEpoch(oam.leaderEpoch().orElse(-1)) + .setCommittedMetadata(oam.metadata()) + .setPartitionIndex(entry.getKey().partition()); + value.add(partition); + return value; + }); + } + for (Map.Entry> entry : offsetData.entrySet()) { + OffsetCommitRequestTopic topic = new OffsetCommitRequestTopic() + .setName(entry.getKey()) + .setPartitions(entry.getValue()); + topics.add(topic); + } + OffsetCommitRequestData data = new OffsetCommitRequestData() + .setGroupId(context.groupId()) + .setTopics(topics); + return new OffsetCommitRequest.Builder(data); + } + + @Override + void handleResponse(AbstractResponse abstractResponse) { + final OffsetCommitResponse response = (OffsetCommitResponse) abstractResponse; + + Map errorCounts = response.errorCounts(); + // 1) If coordinator changed since we fetched it, retry + // 2) If there is a coordinator error, retry + if (ConsumerGroupOperationContext.hasCoordinatorMoved(errorCounts) || + ConsumerGroupOperationContext.shouldRefreshCoordinator(errorCounts)) { + Call call = getAlterConsumerGroupOffsetsCall(context, offsets); + rescheduleFindCoordinatorTask(context, () -> call, this); + return; + } + + final Map partitions = new HashMap<>(); + for (OffsetCommitResponseTopic topic : response.data().topics()) { + for (OffsetCommitResponsePartition partition : topic.partitions()) { + TopicPartition tp = new TopicPartition(topic.name(), partition.partitionIndex()); + Errors error = Errors.forCode(partition.errorCode()); + partitions.put(tp, error); + } + } + context.future().complete(partitions); + } + + @Override + void handleFailure(Throwable throwable) { + context.future().completeExceptionally(throwable); + } + }; + } + + @Override + public ListOffsetsResult listOffsets(Map topicPartitionOffsets, + ListOffsetsOptions options) { + + // preparing topics list for asking metadata about them + final Map> futures = new HashMap<>(topicPartitionOffsets.size()); + final Set topics = new HashSet<>(); + for (TopicPartition topicPartition : topicPartitionOffsets.keySet()) { + topics.add(topicPartition.topic()); + futures.put(topicPartition, new KafkaFutureImpl<>()); + } + + final long nowMetadata = time.milliseconds(); + final long deadline = calcDeadlineMs(nowMetadata, options.timeoutMs()); + + MetadataOperationContext context = + new MetadataOperationContext<>(topics, options, deadline, futures); + + Call metadataCall = getMetadataCall(context, + () -> KafkaAdminClient.this.getListOffsetsCalls(context, topicPartitionOffsets, futures)); + runnable.call(metadataCall, nowMetadata); + + return new ListOffsetsResult(new HashMap<>(futures)); + } + + private List getListOffsetsCalls(MetadataOperationContext context, + Map topicPartitionOffsets, + Map> futures) { + + MetadataResponse mr = context.response().orElseThrow(() -> new IllegalStateException("No Metadata response")); + List calls = new ArrayList<>(); + // grouping topic partitions per leader + Map> leaders = new HashMap<>(); + + for (Map.Entry entry: topicPartitionOffsets.entrySet()) { + + OffsetSpec offsetSpec = entry.getValue(); + TopicPartition tp = entry.getKey(); + KafkaFutureImpl future = futures.get(tp); + long offsetQuery = (offsetSpec instanceof TimestampSpec) + ? ((TimestampSpec) offsetSpec).timestamp() + : (offsetSpec instanceof OffsetSpec.EarliestSpec) + ? ListOffsetRequest.EARLIEST_TIMESTAMP + : ListOffsetRequest.LATEST_TIMESTAMP; + // avoid sending listOffsets request for topics with errors + if (!mr.errors().containsKey(tp.topic())) { + Node node = mr.cluster().leaderFor(tp); + if (node != null) { + Map leadersOnNode = leaders.computeIfAbsent(node, k -> new HashMap()); + ListOffsetTopic topic = leadersOnNode.computeIfAbsent(tp.topic(), k -> new ListOffsetTopic().setName(tp.topic())); + topic.partitions().add(new ListOffsetPartition().setPartitionIndex(tp.partition()).setTimestamp(offsetQuery)); + } else { + future.completeExceptionally(Errors.LEADER_NOT_AVAILABLE.exception()); + } + } else { + future.completeExceptionally(mr.errors().get(tp.topic()).exception()); + } + } + + for (final Map.Entry> entry : leaders.entrySet()) { + final int brokerId = entry.getKey().id(); + + calls.add(new Call("listOffsets on broker " + brokerId, context.deadline(), new ConstantNodeIdProvider(brokerId)) { + + final List partitionsToQuery = new ArrayList<>(entry.getValue().values()); + + @Override + ListOffsetRequest.Builder createRequest(int timeoutMs) { + return ListOffsetRequest.Builder + .forConsumer(true, context.options().isolationLevel()) + .setTargetTimes(partitionsToQuery); + } + + @Override + void handleResponse(AbstractResponse abstractResponse) { + ListOffsetResponse response = (ListOffsetResponse) abstractResponse; + Map retryTopicPartitionOffsets = new HashMap<>(); + + for (ListOffsetTopicResponse topic : response.topics()) { + for (ListOffsetPartitionResponse partition : topic.partitions()) { + TopicPartition tp = new TopicPartition(topic.name(), partition.partitionIndex()); + KafkaFutureImpl future = futures.get(tp); + Errors error = Errors.forCode(partition.errorCode()); + OffsetSpec offsetRequestSpec = topicPartitionOffsets.get(tp); + if (offsetRequestSpec == null) { + log.warn("Server response mentioned unknown topic partition {}", tp); + } else if (MetadataOperationContext.shouldRefreshMetadata(error)) { + retryTopicPartitionOffsets.put(tp, offsetRequestSpec); + } else if (error == Errors.NONE) { + Optional leaderEpoch = (partition.leaderEpoch() == ListOffsetResponse.UNKNOWN_EPOCH) + ? Optional.empty() + : Optional.of(partition.leaderEpoch()); + future.complete(new ListOffsetsResultInfo(partition.offset(), partition.timestamp(), leaderEpoch)); + } else { + future.completeExceptionally(error.exception()); + } + } + } + + if (retryTopicPartitionOffsets.isEmpty()) { + // The server should send back a response for every topic partition. But do a sanity check anyway. + for (ListOffsetTopic topic : partitionsToQuery) { + for (ListOffsetPartition partition : topic.partitions()) { + TopicPartition tp = new TopicPartition(topic.name(), partition.partitionIndex()); + ApiException error = new ApiException("The response from broker " + brokerId + + " did not contain a result for topic partition " + tp); + futures.get(tp).completeExceptionally(error); + } + } + } else { + Set retryTopics = retryTopicPartitionOffsets.keySet().stream().map( + TopicPartition::topic).collect(Collectors.toSet()); + MetadataOperationContext retryContext = + new MetadataOperationContext<>(retryTopics, context.options(), context.deadline(), futures); + rescheduleMetadataTask(retryContext, () -> getListOffsetsCalls(retryContext, retryTopicPartitionOffsets, futures)); + } + } + + @Override + void handleFailure(Throwable throwable) { + for (ListOffsetTopic topic : entry.getValue().values()) { + for (ListOffsetPartition partition : topic.partitions()) { + TopicPartition tp = new TopicPartition(topic.name(), partition.partitionIndex()); + KafkaFutureImpl future = futures.get(tp); + future.completeExceptionally(throwable); + } + } + } + }); + } + return calls; + } + + @Override + public DescribeClientQuotasResult describeClientQuotas(ClientQuotaFilter filter, DescribeClientQuotasOptions options) { + KafkaFutureImpl>> future = new KafkaFutureImpl<>(); + + final long now = time.milliseconds(); + runnable.call(new Call("describeClientQuotas", calcDeadlineMs(now, options.timeoutMs()), + new LeastLoadedNodeProvider()) { + + @Override + DescribeClientQuotasRequest.Builder createRequest(int timeoutMs) { + return new DescribeClientQuotasRequest.Builder(filter); + } + + @Override + void handleResponse(AbstractResponse abstractResponse) { + DescribeClientQuotasResponse response = (DescribeClientQuotasResponse) abstractResponse; + response.complete(future); + } + + @Override + void handleFailure(Throwable throwable) { + future.completeExceptionally(throwable); + } + }, now); + + return new DescribeClientQuotasResult(future); + } + + @Override + public AlterClientQuotasResult alterClientQuotas(Collection entries, AlterClientQuotasOptions options) { + Map> futures = new HashMap<>(entries.size()); + for (ClientQuotaAlteration entry : entries) { + futures.put(entry.entity(), new KafkaFutureImpl<>()); + } + + final long now = time.milliseconds(); + runnable.call(new Call("alterClientQuotas", calcDeadlineMs(now, options.timeoutMs()), + new LeastLoadedNodeProvider()) { + + @Override + AlterClientQuotasRequest.Builder createRequest(int timeoutMs) { + return new AlterClientQuotasRequest.Builder(entries, options.validateOnly()); + } + + @Override + void handleResponse(AbstractResponse abstractResponse) { + AlterClientQuotasResponse response = (AlterClientQuotasResponse) abstractResponse; + response.complete(futures); + } + + @Override + void handleFailure(Throwable throwable) { + completeAllExceptionally(futures.values(), throwable); + } + }, now); + + return new AlterClientQuotasResult(Collections.unmodifiableMap(futures)); + } + + @Override + public DescribeUserScramCredentialsResult describeUserScramCredentials(List users, DescribeUserScramCredentialsOptions options) { + final KafkaFutureImpl dataFuture = new KafkaFutureImpl<>(); + final long now = time.milliseconds(); + Call call = new Call("describeUserScramCredentials", calcDeadlineMs(now, options.timeoutMs()), + new LeastLoadedNodeProvider()) { + @Override + public DescribeUserScramCredentialsRequest.Builder createRequest(final int timeoutMs) { + final DescribeUserScramCredentialsRequestData requestData = new DescribeUserScramCredentialsRequestData(); + + if (users != null && !users.isEmpty()) { + final List userNames = new ArrayList<>(users.size()); + + for (final String user : users) { + if (user != null) { + userNames.add(new UserName().setName(user)); + } + } + + requestData.setUsers(userNames); + } + + return new DescribeUserScramCredentialsRequest.Builder(requestData); + } + + @Override + public void handleResponse(AbstractResponse abstractResponse) { + DescribeUserScramCredentialsResponse response = (DescribeUserScramCredentialsResponse) abstractResponse; + DescribeUserScramCredentialsResponseData data = response.data(); + short messageLevelErrorCode = data.errorCode(); + if (messageLevelErrorCode != Errors.NONE.code()) { + dataFuture.completeExceptionally(Errors.forCode(messageLevelErrorCode).exception(data.errorMessage())); + } else { + dataFuture.complete(data); + } + } + + @Override + void handleFailure(Throwable throwable) { + dataFuture.completeExceptionally(throwable); + } + }; + runnable.call(call, now); + return new DescribeUserScramCredentialsResult(dataFuture); + } + + @Override + public AlterUserScramCredentialsResult alterUserScramCredentials(List alterations, + AlterUserScramCredentialsOptions options) { + final long now = time.milliseconds(); + final Map> futures = new HashMap<>(); + for (UserScramCredentialAlteration alteration: alterations) { + futures.put(alteration.user(), new KafkaFutureImpl<>()); + } + final Map userIllegalAlterationExceptions = new HashMap<>(); + // We need to keep track of users with deletions of an unknown SCRAM mechanism + final String usernameMustNotBeEmptyMsg = "Username must not be empty"; + String passwordMustNotBeEmptyMsg = "Password must not be empty"; + final String unknownScramMechanismMsg = "Unknown SCRAM mechanism"; + alterations.stream().filter(a -> a instanceof UserScramCredentialDeletion).forEach(alteration -> { + final String user = alteration.user(); + if (user == null || user.isEmpty()) { + userIllegalAlterationExceptions.put(alteration.user(), new UnacceptableCredentialException(usernameMustNotBeEmptyMsg)); + } else { + UserScramCredentialDeletion deletion = (UserScramCredentialDeletion) alteration; + ScramMechanism mechanism = deletion.mechanism(); + if (mechanism == null || mechanism == ScramMechanism.UNKNOWN) { + userIllegalAlterationExceptions.put(user, new UnsupportedSaslMechanismException(unknownScramMechanismMsg)); + } + } + }); + // Creating an upsertion may throw InvalidKeyException or NoSuchAlgorithmException, + // so keep track of which users are affected by such a failure so we can fail all their alterations later + final Map> userInsertions = new HashMap<>(); + alterations.stream().filter(a -> a instanceof UserScramCredentialUpsertion) + .filter(alteration -> !userIllegalAlterationExceptions.containsKey(alteration.user())) + .forEach(alteration -> { + final String user = alteration.user(); + if (user == null || user.isEmpty()) { + userIllegalAlterationExceptions.put(alteration.user(), new UnacceptableCredentialException(usernameMustNotBeEmptyMsg)); + } else { + UserScramCredentialUpsertion upsertion = (UserScramCredentialUpsertion) alteration; + try { + byte[] password = upsertion.password(); + if (password == null || password.length == 0) { + userIllegalAlterationExceptions.put(user, new UnacceptableCredentialException(passwordMustNotBeEmptyMsg)); + } else { + ScramMechanism mechanism = upsertion.credentialInfo().mechanism(); + if (mechanism == null || mechanism == ScramMechanism.UNKNOWN) { + userIllegalAlterationExceptions.put(user, new UnsupportedSaslMechanismException(unknownScramMechanismMsg)); + } else { + userInsertions.putIfAbsent(user, new HashMap<>()); + userInsertions.get(user).put(mechanism, getScramCredentialUpsertion(upsertion)); + } + } + } catch (NoSuchAlgorithmException e) { + // we might overwrite an exception from a previous alteration, but we don't really care + // since we just need to mark this user as having at least one illegal alteration + // and make an exception instance available for completing the corresponding future exceptionally + userIllegalAlterationExceptions.put(user, new UnsupportedSaslMechanismException(unknownScramMechanismMsg)); + } catch (InvalidKeyException e) { + // generally shouldn't happen since we deal with the empty password case above, + // but we still need to catch/handle it + userIllegalAlterationExceptions.put(user, new UnacceptableCredentialException(e.getMessage(), e)); + } + } + }); + + // submit alterations only for users that do not have an illegal alteration as identified above + Call call = new Call("alterUserScramCredentials", calcDeadlineMs(now, options.timeoutMs()), + new ControllerNodeProvider()) { + @Override + public AlterUserScramCredentialsRequest.Builder createRequest(int timeoutMs) { + return new AlterUserScramCredentialsRequest.Builder( + new AlterUserScramCredentialsRequestData().setUpsertions(alterations.stream() + .filter(a -> a instanceof UserScramCredentialUpsertion) + .filter(a -> !userIllegalAlterationExceptions.containsKey(a.user())) + .map(a -> userInsertions.get(a.user()).get(((UserScramCredentialUpsertion) a).credentialInfo().mechanism())) + .collect(Collectors.toList())) + .setDeletions(alterations.stream() + .filter(a -> a instanceof UserScramCredentialDeletion) + .filter(a -> !userIllegalAlterationExceptions.containsKey(a.user())) + .map(d -> getScramCredentialDeletion((UserScramCredentialDeletion) d)) + .collect(Collectors.toList()))); + } + + @Override + public void handleResponse(AbstractResponse abstractResponse) { + AlterUserScramCredentialsResponse response = (AlterUserScramCredentialsResponse) abstractResponse; + // Check for controller change + for (Errors error : response.errorCounts().keySet()) { + if (error == Errors.NOT_CONTROLLER) { + handleNotControllerError(error); + } + } + /* Now that we have the results for the ones we sent, + * fail any users that have an illegal alteration as identified above. + * Be sure to do this after the NOT_CONTROLLER error check above + * so that all errors are consistent in that case. + */ + userIllegalAlterationExceptions.entrySet().stream().forEach(entry -> { + futures.get(entry.getKey()).completeExceptionally(entry.getValue()); + }); + response.data().results().forEach(result -> { + KafkaFutureImpl future = futures.get(result.user()); + if (future == null) { + log.warn("Server response mentioned unknown user {}", result.user()); + } else { + Errors error = Errors.forCode(result.errorCode()); + if (error != Errors.NONE) { + future.completeExceptionally(error.exception(result.errorMessage())); + } else { + future.complete(null); + } + } + }); + completeUnrealizedFutures( + futures.entrySet().stream(), + user -> "The broker response did not contain a result for user " + user); + } + + @Override + void handleFailure(Throwable throwable) { + completeAllExceptionally(futures.values(), throwable); + } + }; + runnable.call(call, now); + return new AlterUserScramCredentialsResult(new HashMap<>(futures)); + } + + private static AlterUserScramCredentialsRequestData.ScramCredentialUpsertion getScramCredentialUpsertion(UserScramCredentialUpsertion u) throws InvalidKeyException, NoSuchAlgorithmException { + AlterUserScramCredentialsRequestData.ScramCredentialUpsertion retval = new AlterUserScramCredentialsRequestData.ScramCredentialUpsertion(); + return retval.setName(u.user()) + .setMechanism(u.credentialInfo().mechanism().type()) + .setIterations(u.credentialInfo().iterations()) + .setSalt(u.salt()) + .setSaltedPassword(getSaltedPasword(u.credentialInfo().mechanism(), u.password(), u.salt(), u.credentialInfo().iterations())); + } + + private static AlterUserScramCredentialsRequestData.ScramCredentialDeletion getScramCredentialDeletion(UserScramCredentialDeletion d) { + return new AlterUserScramCredentialsRequestData.ScramCredentialDeletion().setName(d.user()).setMechanism(d.mechanism().type()); + } + + private static byte[] getSaltedPasword(ScramMechanism publicScramMechanism, byte[] password, byte[] salt, int iterations) throws NoSuchAlgorithmException, InvalidKeyException { + return new ScramFormatter(org.apache.kafka.common.security.scram.internals.ScramMechanism.forMechanismName(publicScramMechanism.mechanismName())) + .hi(password, salt, iterations); + } + + @Override + public DescribeFeaturesResult describeFeatures(final DescribeFeaturesOptions options) { + final KafkaFutureImpl future = new KafkaFutureImpl<>(); + final long now = time.milliseconds(); + final Call call = new Call( + "describeFeatures", calcDeadlineMs(now, options.timeoutMs()), new LeastLoadedNodeProvider()) { + + private FeatureMetadata createFeatureMetadata(final ApiVersionsResponse response) { + final Map finalizedFeatures = new HashMap<>(); + for (final FinalizedFeatureKey key : response.data().finalizedFeatures().valuesSet()) { + finalizedFeatures.put(key.name(), new FinalizedVersionRange(key.minVersionLevel(), key.maxVersionLevel())); + } + + Optional finalizedFeaturesEpoch; + if (response.data().finalizedFeaturesEpoch() >= 0L) { + finalizedFeaturesEpoch = Optional.of(response.data().finalizedFeaturesEpoch()); + } else { + finalizedFeaturesEpoch = Optional.empty(); + } + + final Map supportedFeatures = new HashMap<>(); + for (final SupportedFeatureKey key : response.data().supportedFeatures().valuesSet()) { + supportedFeatures.put(key.name(), new SupportedVersionRange(key.minVersion(), key.maxVersion())); + } + + return new FeatureMetadata(finalizedFeatures, finalizedFeaturesEpoch, supportedFeatures); + } + + @Override + ApiVersionsRequest.Builder createRequest(int timeoutMs) { + return new ApiVersionsRequest.Builder(); + } + + @Override + void handleResponse(AbstractResponse response) { + final ApiVersionsResponse apiVersionsResponse = (ApiVersionsResponse) response; + if (apiVersionsResponse.data().errorCode() == Errors.NONE.code()) { + future.complete(createFeatureMetadata(apiVersionsResponse)); + } else { + future.completeExceptionally(Errors.forCode(apiVersionsResponse.data().errorCode()).exception()); + } + } + + @Override + void handleFailure(Throwable throwable) { + completeAllExceptionally(Collections.singletonList(future), throwable); + } + }; + + runnable.call(call, now); + return new DescribeFeaturesResult(future); + } + + @Override + public UpdateFeaturesResult updateFeatures(final Map featureUpdates, + final UpdateFeaturesOptions options) { + if (featureUpdates.isEmpty()) { + throw new IllegalArgumentException("Feature updates can not be null or empty."); + } + + final Map> updateFutures = new HashMap<>(); + for (final Map.Entry entry : featureUpdates.entrySet()) { + final String feature = entry.getKey(); + if (feature.trim().isEmpty()) { + throw new IllegalArgumentException("Provided feature can not be empty."); + } + updateFutures.put(entry.getKey(), new KafkaFutureImpl<>()); + } + + final long now = time.milliseconds(); + final Call call = new Call("updateFeatures", calcDeadlineMs(now, options.timeoutMs()), + new ControllerNodeProvider()) { + + @Override + UpdateFeaturesRequest.Builder createRequest(int timeoutMs) { + final UpdateFeaturesRequestData.FeatureUpdateKeyCollection featureUpdatesRequestData + = new UpdateFeaturesRequestData.FeatureUpdateKeyCollection(); + for (Map.Entry entry : featureUpdates.entrySet()) { + final String feature = entry.getKey(); + final FeatureUpdate update = entry.getValue(); + final UpdateFeaturesRequestData.FeatureUpdateKey requestItem = + new UpdateFeaturesRequestData.FeatureUpdateKey(); + requestItem.setFeature(feature); + requestItem.setMaxVersionLevel(update.maxVersionLevel()); + requestItem.setAllowDowngrade(update.allowDowngrade()); + featureUpdatesRequestData.add(requestItem); + } + return new UpdateFeaturesRequest.Builder( + new UpdateFeaturesRequestData() + .setTimeoutMs(timeoutMs) + .setFeatureUpdates(featureUpdatesRequestData)); + } + + @Override + void handleResponse(AbstractResponse abstractResponse) { + final UpdateFeaturesResponse response = + (UpdateFeaturesResponse) abstractResponse; + + Errors topLevelError = Errors.forCode(response.data().errorCode()); + switch (topLevelError) { + case NONE: + for (final UpdatableFeatureResult result : response.data().results()) { + final KafkaFutureImpl future = updateFutures.get(result.feature()); + if (future == null) { + log.warn("Server response mentioned unknown feature {}", result.feature()); + } else { + final Errors error = Errors.forCode(result.errorCode()); + if (error == Errors.NONE) { + future.complete(null); + } else { + future.completeExceptionally(error.exception(result.errorMessage())); + } + } + } + // The server should send back a response for every feature, but we do a sanity check anyway. + completeUnrealizedFutures(updateFutures.entrySet().stream(), + feature -> "The controller response did not contain a result for feature " + feature); + break; + case NOT_CONTROLLER: + handleNotControllerError(topLevelError); + break; + default: + for (final Map.Entry> entry : updateFutures.entrySet()) { + final String errorMsg = response.data().errorMessage(); + entry.getValue().completeExceptionally(topLevelError.exception(errorMsg)); + } + break; + } + } + + @Override + void handleFailure(Throwable throwable) { + completeAllExceptionally(updateFutures.values(), throwable); + } + }; + + runnable.call(call, now); + return new UpdateFeaturesResult(new HashMap<>(updateFutures)); + } + + /** + * 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 new file mode 100644 index 0000000000000..af738ca209fb9 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/ListConsumerGroupOffsetsOptions.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 org.apache.kafka.clients.admin; + +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.annotation.InterfaceStability; + +import java.util.List; + +/** + * Options for {@link Admin#listConsumerGroupOffsets(String)}. + *

+ * The API of this class is evolving, see {@link Admin} for details. + */ +@InterfaceStability.Evolving +public class ListConsumerGroupOffsetsOptions extends AbstractOptions { + + private List topicPartitions = null; + + /** + * Set the topic partitions to list as part of the result. + * {@code null} includes all topic partitions. + * + * @param topicPartitions List of topic partitions to include + * @return This ListGroupOffsetsOptions + */ + public ListConsumerGroupOffsetsOptions topicPartitions(List topicPartitions) { + this.topicPartitions = topicPartitions; + return this; + } + + /** + * Returns a list of topic partitions to add as part of the result. + */ + public List topicPartitions() { + return topicPartitions; + } +} 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 new file mode 100644 index 0000000000000..48f4531418110 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/ListConsumerGroupOffsetsResult.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.clients.admin; + +import org.apache.kafka.clients.consumer.OffsetAndMetadata; +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 the {@link Admin#listConsumerGroupOffsets(String)} call. + *

+ * The API of this class is evolving, see {@link Admin} for details. + */ +@InterfaceStability.Evolving +public class ListConsumerGroupOffsetsResult { + + final KafkaFuture> future; + + ListConsumerGroupOffsetsResult(KafkaFuture> future) { + this.future = future; + } + + /** + * 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 new file mode 100644 index 0000000000000..9f1f38dd4a8e6 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/ListConsumerGroupsOptions.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 org.apache.kafka.clients.admin; + +import java.util.Collections; +import java.util.HashSet; +import java.util.Set; + +import org.apache.kafka.common.ConsumerGroupState; +import org.apache.kafka.common.annotation.InterfaceStability; + +/** + * Options for {@link Admin#listConsumerGroups()}. + * + * The API of this class is evolving, see {@link Admin} for details. + */ +@InterfaceStability.Evolving +public class ListConsumerGroupsOptions extends AbstractOptions { + + private Set states = Collections.emptySet(); + + /** + * If states is set, only groups in these states will be returned by listConsumerGroups() + * Otherwise, all groups are returned. + * This operation is supported by brokers with version 2.6.0 or later. + */ + public ListConsumerGroupsOptions inStates(Set states) { + this.states = (states == null) ? Collections.emptySet() : new HashSet<>(states); + return this; + } + + /** + * Returns the list of States that are requested or empty if no states have been specified + */ + public Set states() { + return states; + } +} 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 new file mode 100644 index 0000000000000..7732ec98c9638 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/ListConsumerGroupsResult.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.KafkaFuture; +import org.apache.kafka.common.annotation.InterfaceStability; +import org.apache.kafka.common.internals.KafkaFutureImpl; + +import java.util.ArrayList; +import java.util.Collection; + +/** + * The result of the {@link Admin#listConsumerGroups()} call. + *

+ * The API of this class is evolving, see {@link Admin} for details. + */ +@InterfaceStability.Evolving +public class ListConsumerGroupsResult { + private final KafkaFutureImpl> all; + private final KafkaFutureImpl> valid; + private final KafkaFutureImpl> errors; + + ListConsumerGroupsResult(KafkaFutureImpl> future) { + this.all = new KafkaFutureImpl<>(); + this.valid = new KafkaFutureImpl<>(); + this.errors = new KafkaFutureImpl<>(); + future.thenApply(new KafkaFuture.BaseFunction, Void>() { + @Override + public Void apply(Collection results) { + ArrayList curErrors = new ArrayList<>(); + ArrayList curValid = new ArrayList<>(); + for (Object resultObject : results) { + if (resultObject instanceof Throwable) { + curErrors.add((Throwable) resultObject); + } else { + curValid.add((ConsumerGroupListing) resultObject); + } + } + if (!curErrors.isEmpty()) { + all.completeExceptionally(curErrors.get(0)); + } else { + all.complete(curValid); + } + valid.complete(curValid); + errors.complete(curErrors); + return null; + } + }); + } + + /** + * Returns a future that yields either an exception, or the full set of consumer group + * listings. + * + * In the event of a failure, the future yields nothing but the first exception which + * occurred. + */ + public KafkaFuture> all() { + return all; + } + + /** + * Returns a future which yields just the valid listings. + * + * This future never fails with an error, no matter what happens. Errors are completely + * ignored. If nothing can be fetched, an empty collection is yielded. + * If there is an error, but some results can be returned, this future will yield + * those partial results. When using this future, it is a good idea to also check + * the errors future so that errors can be displayed and handled. + */ + public KafkaFuture> valid() { + return valid; + } + + /** + * Returns a future which yields just the errors which occurred. + * + * If this future yields a non-empty collection, it is very likely that elements are + * missing from the valid() set. + * + * This future itself never fails with an error. In the event of an error, this future + * will successfully yield a collection containing at least one exception. + */ + public KafkaFuture> errors() { + return errors; + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/ListOffsetsOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/ListOffsetsOptions.java new file mode 100644 index 0000000000000..684ad265de77d --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/ListOffsetsOptions.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.clients.admin; + +import org.apache.kafka.common.IsolationLevel; +import org.apache.kafka.common.annotation.InterfaceStability; + +import java.util.Map; + +/** + * Options for {@link AdminClient#listOffsets(Map)}. + * + * The API of this class is evolving, see {@link AdminClient} for details. + */ +@InterfaceStability.Evolving +public class ListOffsetsOptions extends AbstractOptions { + + private final IsolationLevel isolationLevel; + + public ListOffsetsOptions() { + this(IsolationLevel.READ_UNCOMMITTED); + } + + public ListOffsetsOptions(IsolationLevel isolationLevel) { + this.isolationLevel = isolationLevel; + } + + public IsolationLevel isolationLevel() { + return isolationLevel; + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/ListOffsetsResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/ListOffsetsResult.java new file mode 100644 index 0000000000000..5eb00deb0697a --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/ListOffsetsResult.java @@ -0,0 +1,107 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF 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.HashMap; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.ExecutionException; + +import org.apache.kafka.common.KafkaFuture; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.annotation.InterfaceStability; + +/** + * The result of the {@link AdminClient#listOffsets(Map)} call. + * + * The API of this class is evolving, see {@link AdminClient} for details. + */ +@InterfaceStability.Evolving +public class ListOffsetsResult { + + private final Map> futures; + + public ListOffsetsResult(Map> futures) { + this.futures = futures; + } + + /** + * Return a future which can be used to check the result for a given partition. + */ + public KafkaFuture partitionResult(final TopicPartition partition) { + KafkaFuture future = futures.get(partition); + if (future == null) { + throw new IllegalArgumentException( + "List Offsets for partition \"" + partition + "\" was not attempted"); + } + return future; + } + + /** + * Return a future which succeeds only if offsets for all specified partitions have been successfully + * retrieved. + */ + public KafkaFuture> all() { + return KafkaFuture.allOf(futures.values().toArray(new KafkaFuture[0])) + .thenApply(new KafkaFuture.BaseFunction>() { + @Override + public Map apply(Void v) { + Map offsets = new HashMap<>(futures.size()); + for (Map.Entry> entry : futures.entrySet()) { + try { + offsets.put(entry.getKey(), entry.getValue().get()); + } catch (InterruptedException | ExecutionException e) { + // This should be unreachable, because allOf ensured that all the futures completed successfully. + throw new RuntimeException(e); + } + } + return offsets; + } + }); + } + + public static class ListOffsetsResultInfo { + + private final long offset; + private final long timestamp; + private final Optional leaderEpoch; + + public ListOffsetsResultInfo(long offset, long timestamp, Optional leaderEpoch) { + this.offset = offset; + this.timestamp = timestamp; + this.leaderEpoch = leaderEpoch; + } + + public long offset() { + return offset; + } + + public long timestamp() { + return timestamp; + } + + public Optional leaderEpoch() { + return leaderEpoch; + } + + @Override + public String toString() { + return "ListOffsetsResultInfo(offset=" + offset + ", timestamp=" + timestamp + ", leaderEpoch=" + + leaderEpoch + ")"; + } + } +} 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 f656ff4443e50..5a8f2b4ad1514 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,15 +20,26 @@ 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 { private boolean listInternal = false; + /** + * Set the timeout in milliseconds for this operation or {@code null} if the default api timeout for the + * AdminClient should be used. + * + */ + // This method is retained to keep binary compatibility with 0.11 + public ListTopicsOptions timeoutMs(Integer timeoutMs) { + this.timeoutMs = timeoutMs; + return this; + } + /** * Set whether we should list internal topics. * 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 e54b3defe36ff..21540732d30d2 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 { @@ -48,23 +48,13 @@ public KafkaFuture> namesToListings() { * Return a future which yields a collection of TopicListing objects. */ public KafkaFuture> listings() { - return future.thenApply(new KafkaFuture.Function, Collection>() { - @Override - public Collection apply(Map namesToDescriptions) { - return namesToDescriptions.values(); - } - }); + return future.thenApply(namesToDescriptions -> namesToDescriptions.values()); } /** * Return a future which yields a collection of topic names. */ public KafkaFuture> names() { - return future.thenApply(new KafkaFuture.Function, Set>() { - @Override - public Set apply(Map namesToListings) { - return namesToListings.keySet(); - } - }); + return future.thenApply(namesToListings -> namesToListings.keySet()); } } diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/LogDirDescription.java b/clients/src/main/java/org/apache/kafka/clients/admin/LogDirDescription.java new file mode 100644 index 0000000000000..1c326ec43b926 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/LogDirDescription.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.clients.admin; + +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.errors.ApiException; + +import java.util.Map; + +import static java.util.Collections.unmodifiableMap; + +/** + * A description of a log directory on a particular broker. + */ +public class LogDirDescription { + private final Map replicaInfos; + private final ApiException error; + + public LogDirDescription(ApiException error, Map replicaInfos) { + this.error = error; + this.replicaInfos = replicaInfos; + } + + /** + * Returns `ApiException` if the log directory is offline or an error occurred, otherwise returns null. + *
    + *
  • KafkaStorageException - The log directory is offline. + *
  • UnknownServerException - The server experienced an unexpected error when processing the request. + *
+ */ + public ApiException error() { + return error; + } + + /** + * A map from topic partition to replica information for that partition + * in this log directory. + */ + public Map replicaInfos() { + return unmodifiableMap(replicaInfos); + } + + @Override + public String toString() { + return "LogDirDescription(" + + "replicaInfos=" + replicaInfos + + ", error=" + error + + ')'; + } +} 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 new file mode 100644 index 0000000000000..3305de02c757b --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/MemberAssignment.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 org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.utils.Utils; + +import java.util.Collections; +import java.util.HashSet; +import java.util.Objects; +import java.util.Set; + +/** + * A description of the assignments of a specific group member. + */ +public class MemberAssignment { + private final Set topicPartitions; + + /** + * Creates an instance with the specified parameters. + * + * @param topicPartitions List of topic partitions + */ + public MemberAssignment(Set topicPartitions) { + this.topicPartitions = topicPartitions == null ? Collections.emptySet() : + Collections.unmodifiableSet(new HashSet<>(topicPartitions)); + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + + MemberAssignment that = (MemberAssignment) o; + + return Objects.equals(topicPartitions, that.topicPartitions); + } + + @Override + public int hashCode() { + return topicPartitions != null ? topicPartitions.hashCode() : 0; + } + + /** + * The topic partitions assigned to a group member. + */ + public Set topicPartitions() { + return topicPartitions; + } + + @Override + public String toString() { + return "(topicPartitions=" + Utils.join(topicPartitions, ",") + ")"; + } +} 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 new file mode 100644 index 0000000000000..7bc6b14906fe2 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/MemberDescription.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.clients.admin; + +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, + 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; + } + + public MemberDescription(String memberId, + String clientId, + String host, + MemberAssignment assignment) { + this(memberId, Optional.empty(), clientId, host, assignment); + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + 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); + } + + @Override + public int hashCode() { + return Objects.hash(memberId, groupInstanceId, clientId, host, assignment); + } + + /** + * The consumer id of the group member. + */ + public String consumerId() { + return memberId; + } + + /** + * The instance id of the group member. + */ + public Optional groupInstanceId() { + return groupInstanceId; + } + + /** + * The client id of the group member. + */ + public String clientId() { + return clientId; + } + + /** + * The host where the group member is running. + */ + public String host() { + return host; + } + + /** + * The assignment of the group member. + */ + public MemberAssignment assignment() { + return 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 ff09579372878..2f335d02f2f2b 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,20 +17,27 @@ package org.apache.kafka.clients.admin; -import org.apache.kafka.common.requests.CreateTopicsRequest.TopicDetails; +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; +import org.apache.kafka.common.requests.CreateTopicsRequest; import java.util.Collection; import java.util.Collections; import java.util.List; import java.util.Map; +import java.util.Objects; +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 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; @@ -38,6 +45,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; @@ -53,8 +69,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); } @@ -69,14 +85,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(CreateTopicsRequest.NO_NUM_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(CreateTopicsRequest.NO_REPLICATION_FACTOR); } /** @@ -98,31 +114,63 @@ public NewTopic configs(Map configs) { return this; } - TopicDetails convertToTopicDetails() { + /** + * The configuration for the new topic or null if no configs ever specified. + */ + public Map configs() { + return configs; + } + + CreatableTopic convertToCreatableTopic() { + CreatableTopic creatableTopic = new CreatableTopic(). + setName(name). + setNumPartitions(numPartitions.orElse(CreateTopicsRequest.NO_NUM_PARTITIONS)). + setReplicationFactor(replicationFactor.orElse(CreateTopicsRequest.NO_REPLICATION_FACTOR)); if (replicasAssignments != null) { - if (configs != null) { - return new TopicDetails(replicasAssignments, configs); - } else { - return new TopicDetails(replicasAssignments); + for (Entry> entry : replicasAssignments.entrySet()) { + creatableTopic.assignments().add( + new CreatableReplicaAssignment(). + setPartitionIndex(entry.getKey()). + setBrokerIds(entry.getValue())); } - } else { - if (configs != null) { - return new TopicDetails(numPartitions, replicationFactor, configs); - } else { - return new TopicDetails(numPartitions, replicationFactor); + } + if (configs != null) { + for (Entry entry : configs.entrySet()) { + creatableTopic.configs().add( + new CreateableTopicConfig(). + setName(entry.getKey()). + setValue(entry.getValue())); } } + return creatableTopic; } @Override 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(")"); return bld.toString(); } + + @Override + public boolean equals(final Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + final NewTopic that = (NewTopic) o; + return Objects.equals(name, that.name) && + Objects.equals(numPartitions, that.numPartitions) && + Objects.equals(replicationFactor, that.replicationFactor) && + Objects.equals(replicasAssignments, that.replicasAssignments) && + Objects.equals(configs, that.configs); + } + + @Override + public int hashCode() { + return Objects.hash(name, numPartitions, replicationFactor, replicasAssignments, configs); + } } diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/OffsetSpec.java b/clients/src/main/java/org/apache/kafka/clients/admin/OffsetSpec.java new file mode 100644 index 0000000000000..339e9cf815e87 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/OffsetSpec.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.clients.admin; + +import java.util.Map; + +/** + * This class allows to specify the desired offsets when using {@link KafkaAdminClient#listOffsets(Map, ListOffsetsOptions)} + */ +public class OffsetSpec { + + public static class EarliestSpec extends OffsetSpec { } + public static class LatestSpec extends OffsetSpec { } + public static class TimestampSpec extends OffsetSpec { + private final long timestamp; + + TimestampSpec(long timestamp) { + this.timestamp = timestamp; + } + + long timestamp() { + return timestamp; + } + } + + /** + * Used to retrieve the latest offset of a partition + */ + public static OffsetSpec latest() { + return new LatestSpec(); + } + + /** + * Used to retrieve the earliest offset of a partition + */ + public static OffsetSpec earliest() { + return new EarliestSpec(); + } + + /** + * Used to retrieve the earliest offset whose timestamp is greater than + * or equal to the given timestamp in the corresponding partition + * @param timestamp in milliseconds + */ + public static OffsetSpec forTimestamp(long timestamp) { + return new TimestampSpec(timestamp); + } + +} 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 1981929fc8757..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 { @@ -51,6 +51,21 @@ public long beforeOffset() { return offset; } + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + + RecordsToDelete that = (RecordsToDelete) o; + + return this.offset == that.offset; + } + + @Override + public int hashCode() { + return (int) offset; + } + @Override public String toString() { return "(beforeOffset = " + offset + ")"; 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..322beec39e419 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/RemoveMembersFromConsumerGroupOptions.java @@ -0,0 +1,55 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF 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.Collections; +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) { + if (members.isEmpty()) { + throw new IllegalArgumentException("Invalid empty members has been provided"); + } + this.members = new HashSet<>(members); + } + + public RemoveMembersFromConsumerGroupOptions() { + this.members = Collections.emptySet(); + } + + public Set members() { + return members; + } + + public boolean removeAll() { + return members.isEmpty(); + } +} 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..3845e2f6aac62 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/RemoveMembersFromConsumerGroupResult.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.clients.admin; + +import org.apache.kafka.common.KafkaException; +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 { + if (removeAll()) { + for (Map.Entry entry: memberErrors.entrySet()) { + Exception exception = entry.getValue().exception(); + if (exception != null) { + Throwable ex = new KafkaException("Encounter exception when trying to remove: " + + entry.getKey(), exception); + result.completeExceptionally(ex); + return; + } + } + } 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 (removeAll()) { + throw new IllegalArgumentException("The method: memberResult is not applicable in 'removeAll' mode"); + } + 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; + } + } + + private boolean removeAll() { + return memberInfos.isEmpty(); + } +} 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 new file mode 100644 index 0000000000000..5c2b0d1afcea1 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/RenewDelegationTokenOptions.java @@ -0,0 +1,39 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.clients.admin; + +import org.apache.kafka.common.annotation.InterfaceStability; + +/** + * Options for {@link Admin#renewDelegationToken(byte[], RenewDelegationTokenOptions)}. + * + * The API of this class is evolving, see {@link Admin} for details. + */ +@InterfaceStability.Evolving +public class RenewDelegationTokenOptions extends AbstractOptions { + private long renewTimePeriodMs = -1; + + public RenewDelegationTokenOptions renewTimePeriodMs(long renewTimePeriodMs) { + this.renewTimePeriodMs = renewTimePeriodMs; + return this; + } + + public long renewTimePeriodMs() { + return renewTimePeriodMs; + } +} 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 new file mode 100644 index 0000000000000..74725d4d633b5 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/RenewDelegationTokenResult.java @@ -0,0 +1,42 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.clients.admin; + +import org.apache.kafka.common.KafkaFuture; +import org.apache.kafka.common.annotation.InterfaceStability; + +/** + * The result of the {@link KafkaAdminClient#expireDelegationToken(byte[], ExpireDelegationTokenOptions)} call. + * + * The API of this class is evolving, see {@link Admin} for details. + */ +@InterfaceStability.Evolving +public class RenewDelegationTokenResult { + private final KafkaFuture expiryTimestamp; + + RenewDelegationTokenResult(KafkaFuture expiryTimestamp) { + this.expiryTimestamp = expiryTimestamp; + } + + /** + * Returns a future which yields expiry timestamp + */ + public KafkaFuture expiryTimestamp() { + return expiryTimestamp; + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/ReplicaInfo.java b/clients/src/main/java/org/apache/kafka/clients/admin/ReplicaInfo.java new file mode 100644 index 0000000000000..b77375d59605d --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/ReplicaInfo.java @@ -0,0 +1,70 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.kafka.clients.admin; + +/** + * A description of a replica on a particular broker. + */ +public class ReplicaInfo { + + private final long size; + private final long offsetLag; + private final boolean isFuture; + + public ReplicaInfo(long size, long offsetLag, boolean isFuture) { + this.size = size; + this.offsetLag = offsetLag; + this.isFuture = isFuture; + } + + /** + * The total size of the log segments in this replica in bytes. + */ + public long size() { + return size; + } + + /** + * The lag of the log's LEO with respect to the partition's + * high watermark (if it is the current log for the partition) + * or the current replica's LEO (if it is the {@linkplain #isFuture() future log} + * for the partition). + */ + public long offsetLag() { + return offsetLag; + } + + /** + * Whether this replica has been created by a AlterReplicaLogDirsRequest + * but not yet replaced the current replica on the broker. + * + * @return true if this log is created by AlterReplicaLogDirsRequest and will replace the current log + * of the replica at some time in the future. + */ + public boolean isFuture() { + return isFuture; + } + + @Override + public String toString() { + return "ReplicaInfo(" + + "size=" + size + + ", offsetLag=" + offsetLag + + ", isFuture=" + isFuture + + ')'; + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/ScramCredentialInfo.java b/clients/src/main/java/org/apache/kafka/clients/admin/ScramCredentialInfo.java new file mode 100644 index 0000000000000..e8403b6e12823 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/ScramCredentialInfo.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.clients.admin; + +import java.util.Objects; + +/** + * Mechanism and iterations for a SASL/SCRAM credential associated with a user. + * + * @see KIP-554: Add Broker-side SCRAM Config API + */ +public class ScramCredentialInfo { + private final ScramMechanism mechanism; + private final int iterations; + + /** + * + * @param mechanism the required mechanism + * @param iterations the number of iterations used when creating the credential + */ + public ScramCredentialInfo(ScramMechanism mechanism, int iterations) { + this.mechanism = Objects.requireNonNull(mechanism); + this.iterations = iterations; + } + + /** + * + * @return the mechanism + */ + public ScramMechanism mechanism() { + return mechanism; + } + + /** + * + * @return the number of iterations used when creating the credential + */ + public int iterations() { + return iterations; + } + + @Override + public String toString() { + return "ScramCredentialInfo{" + + "mechanism=" + mechanism + + ", iterations=" + iterations + + '}'; + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + ScramCredentialInfo that = (ScramCredentialInfo) o; + return iterations == that.iterations && + mechanism == that.mechanism; + } + + @Override + public int hashCode() { + return Objects.hash(mechanism, iterations); + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/ScramMechanism.java b/clients/src/main/java/org/apache/kafka/clients/admin/ScramMechanism.java new file mode 100644 index 0000000000000..9869bc7c06cbb --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/ScramMechanism.java @@ -0,0 +1,87 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF 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.Arrays; + +/** + * Representation of a SASL/SCRAM Mechanism. + * + * @see KIP-554: Add Broker-side SCRAM Config API + */ +public enum ScramMechanism { + UNKNOWN((byte) 0), + SCRAM_SHA_256((byte) 1), + SCRAM_SHA_512((byte) 2); + + private static final ScramMechanism[] VALUES = values(); + + /** + * + * @param type the type indicator + * @return the instance corresponding to the given type indicator, otherwise {@link #UNKNOWN} + */ + public static ScramMechanism fromType(byte type) { + for (ScramMechanism scramMechanism : VALUES) { + if (scramMechanism.type == type) { + return scramMechanism; + } + } + return UNKNOWN; + } + + /** + * + * @param mechanismName the SASL SCRAM mechanism name + * @return the corresponding SASL SCRAM mechanism enum, otherwise {@link #UNKNOWN} + * @see + * Salted Challenge Response Authentication Mechanism (SCRAM) SASL and GSS-API Mechanisms, Section 4 + */ + public String mechanismName() { + return this.mechanismName; + } + + /** + * + * @return the type indicator for this SASL SCRAM mechanism + */ + public byte type() { + return this.type; + } + + private final byte type; + private final String mechanismName; + + private ScramMechanism(byte type) { + this.type = type; + this.mechanismName = toString().replace('_', '-'); + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/SupportedVersionRange.java b/clients/src/main/java/org/apache/kafka/clients/admin/SupportedVersionRange.java new file mode 100644 index 0000000000000..d71da31fb8200 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/SupportedVersionRange.java @@ -0,0 +1,82 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.kafka.clients.admin; + +import java.util.Objects; + +/** + * Represents a range of versions that a particular broker supports for some feature. + */ +public class SupportedVersionRange { + private final short minVersion; + + private final short maxVersion; + + /** + * Raises an exception unless the following conditions are met: + * 1 <= minVersion <= maxVersion. + * + * @param minVersion The minimum version value. + * @param maxVersion The maximum version value. + * + * @throws IllegalArgumentException Raised when the condition described above is not met. + */ + SupportedVersionRange(final short minVersion, final short maxVersion) { + if (minVersion < 1 || maxVersion < 1 || maxVersion < minVersion) { + throw new IllegalArgumentException( + String.format( + "Expected 1 <= minVersion <= maxVersion but received minVersion:%d, maxVersion:%d.", + minVersion, + maxVersion)); + } + this.minVersion = minVersion; + this.maxVersion = maxVersion; + } + + public short minVersion() { + return minVersion; + } + + public short maxVersion() { + return maxVersion; + } + + @Override + public boolean equals(Object other) { + if (this == other) { + return true; + } + + if (other == null || getClass() != other.getClass()) { + return false; + } + + final SupportedVersionRange that = (SupportedVersionRange) other; + return this.minVersion == that.minVersion && this.maxVersion == that.maxVersion; + } + + @Override + public int hashCode() { + return Objects.hash(minVersion, maxVersion); + } + + @Override + public String toString() { + return String.format("SupportedVersionRange[min_version:%d, max_version:%d]", minVersion, maxVersion); + } +} + 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 c2208923264e4..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,6 +33,23 @@ public class TopicDescription { private final String name; private final boolean internal; private final List partitions; + private final Set authorizedOperations; + + @Override + public boolean equals(final Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + 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() { + return Objects.hash(name, internal, partitions, authorizedOperations); + } /** * Create an instance with the specified parameters. @@ -39,9 +60,24 @@ public class TopicDescription { * 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; } /** @@ -67,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/UpdateFeaturesOptions.java b/clients/src/main/java/org/apache/kafka/clients/admin/UpdateFeaturesOptions.java new file mode 100644 index 0000000000000..7a9f2141b2ab1 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/UpdateFeaturesOptions.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 java.util.Map; +import org.apache.kafka.common.annotation.InterfaceStability; + +/** + * Options for {@link AdminClient#updateFeatures(Map, UpdateFeaturesOptions)}. + * + * The API of this class is evolving. See {@link Admin} for details. + */ +@InterfaceStability.Evolving +public class UpdateFeaturesOptions extends AbstractOptions { +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/UpdateFeaturesResult.java b/clients/src/main/java/org/apache/kafka/clients/admin/UpdateFeaturesResult.java new file mode 100644 index 0000000000000..6c484dc24d95b --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/UpdateFeaturesResult.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.clients.admin; + +import java.util.Map; +import org.apache.kafka.common.KafkaFuture; + +/** + * The result of the {@link Admin#updateFeatures(Map, UpdateFeaturesOptions)} call. + * + * The API of this class is evolving, see {@link Admin} for details. + */ +public class UpdateFeaturesResult { + private final Map> futures; + + /** + * @param futures a map from feature name to future, which can be used to check the status of + * individual feature updates. + */ + UpdateFeaturesResult(final Map> futures) { + this.futures = futures; + } + + public Map> values() { + return futures; + } + + /** + * Return a future which succeeds if all the feature updates succeed. + */ + public KafkaFuture all() { + return KafkaFuture.allOf(futures.values().toArray(new KafkaFuture[0])); + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/UserScramCredentialAlteration.java b/clients/src/main/java/org/apache/kafka/clients/admin/UserScramCredentialAlteration.java new file mode 100644 index 0000000000000..8293fe514df19 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/UserScramCredentialAlteration.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.clients.admin; + +import java.util.Objects; + +/** + * A request to alter a user's SASL/SCRAM credentials. + * + * @see KIP-554: Add Broker-side SCRAM Config API + */ +public abstract class UserScramCredentialAlteration { + protected final String user; + + /** + * + * @param user the mandatory user + */ + protected UserScramCredentialAlteration(String user) { + this.user = Objects.requireNonNull(user); + } + + /** + * + * @return the always non-null user + */ + public String user() { + return this.user; + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/UserScramCredentialDeletion.java b/clients/src/main/java/org/apache/kafka/clients/admin/UserScramCredentialDeletion.java new file mode 100644 index 0000000000000..633075aaf2f2b --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/UserScramCredentialDeletion.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.clients.admin; + +import java.util.Objects; + +/** + * A request to delete a SASL/SCRAM credential for a user. + * + * @see KIP-554: Add Broker-side SCRAM Config API + */ +public class UserScramCredentialDeletion extends UserScramCredentialAlteration { + private final ScramMechanism mechanism; + + /** + * @param user the mandatory user + * @param mechanism the mandatory mechanism + */ + public UserScramCredentialDeletion(String user, ScramMechanism mechanism) { + super(user); + this.mechanism = Objects.requireNonNull(mechanism); + } + + /** + * + * @return the always non-null mechanism + */ + public ScramMechanism mechanism() { + return mechanism; + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/UserScramCredentialUpsertion.java b/clients/src/main/java/org/apache/kafka/clients/admin/UserScramCredentialUpsertion.java new file mode 100644 index 0000000000000..5d5cf9cbad347 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/UserScramCredentialUpsertion.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.admin; + +import org.apache.kafka.common.security.scram.internals.ScramFormatter; + +import java.nio.charset.StandardCharsets; +import java.security.SecureRandom; +import java.util.Objects; + +/** + * A request to update/insert a SASL/SCRAM credential for a user. + * + * @see KIP-554: Add Broker-side SCRAM Config API + */ +public class UserScramCredentialUpsertion extends UserScramCredentialAlteration { + private final ScramCredentialInfo info; + private final byte[] salt; + private final byte[] password; + + /** + * Constructor that generates a random salt + * + * @param user the user for which the credential is to be updated/inserted + * @param credentialInfo the mechanism and iterations to be used + * @param password the password + */ + public UserScramCredentialUpsertion(String user, ScramCredentialInfo credentialInfo, String password) { + this(user, credentialInfo, password.getBytes(StandardCharsets.UTF_8)); + } + + /** + * Constructor that generates a random salt + * + * @param user the user for which the credential is to be updated/inserted + * @param credentialInfo the mechanism and iterations to be used + * @param password the password + */ + public UserScramCredentialUpsertion(String user, ScramCredentialInfo credentialInfo, byte[] password) { + this(user, credentialInfo, password, generateRandomSalt()); + } + + /** + * Constructor that accepts an explicit salt + * + * @param user the user for which the credential is to be updated/inserted + * @param credentialInfo the mechanism and iterations to be used + * @param password the password + * @param salt the salt to be used + */ + public UserScramCredentialUpsertion(String user, ScramCredentialInfo credentialInfo, byte[] password, byte[] salt) { + super(Objects.requireNonNull(user)); + this.info = Objects.requireNonNull(credentialInfo); + this.password = Objects.requireNonNull(password); + this.salt = Objects.requireNonNull(salt); + } + + /** + * + * @return the mechanism and iterations + */ + public ScramCredentialInfo credentialInfo() { + return info; + } + + /** + * + * @return the salt + */ + public byte[] salt() { + return salt; + } + + /** + * + * @return the password + */ + public byte[] password() { + return password; + } + + private static byte[] generateRandomSalt() { + return ScramFormatter.secureRandomBytes(new SecureRandom()); + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/UserScramCredentialsDescription.java b/clients/src/main/java/org/apache/kafka/clients/admin/UserScramCredentialsDescription.java new file mode 100644 index 0000000000000..97bc3588af6aa --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/UserScramCredentialsDescription.java @@ -0,0 +1,82 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.clients.admin; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Objects; + +/** + * Representation of all SASL/SCRAM credentials associated with a user that can be retrieved, or an exception indicating + * why credentials could not be retrieved. + * + * @see KIP-554: Add Broker-side SCRAM Config API + */ +public class UserScramCredentialsDescription { + private final String name; + private final List credentialInfos; + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + UserScramCredentialsDescription that = (UserScramCredentialsDescription) o; + return name.equals(that.name) && + credentialInfos.equals(that.credentialInfos); + } + + @Override + public int hashCode() { + return Objects.hash(name, credentialInfos); + } + + @Override + public String toString() { + return "UserScramCredentialsDescription{" + + "name='" + name + '\'' + + ", credentialInfos=" + credentialInfos + + '}'; + } + + /** + * + * @param name the required user name + * @param credentialInfos the required SASL/SCRAM credential representations for the user + */ + public UserScramCredentialsDescription(String name, List credentialInfos) { + this.name = Objects.requireNonNull(name); + this.credentialInfos = Collections.unmodifiableList(new ArrayList<>(credentialInfos)); + } + + /** + * + * @return the user name + */ + public String name() { + return name; + } + + /** + * + * @return the always non-null/unmodifiable list of SASL/SCRAM credential representations for the user + */ + public List credentialInfos() { + return credentialInfos; + } +} 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 new file mode 100644 index 0000000000000..6e834520c46f7 --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/internals/AdminMetadataManager.java @@ -0,0 +1,258 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF 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.internals; + +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; +import org.apache.kafka.common.requests.RequestHeader; +import org.apache.kafka.common.utils.LogContext; +import org.slf4j.Logger; + +import java.util.Collections; +import java.util.List; +import java.util.Optional; + +/** + * Manages the metadata for KafkaAdminClient. + * + * This class is not thread-safe. It is only accessed from the AdminClient + * service thread (which also uses the NetworkClient). + */ +public class AdminMetadataManager { + private Logger log; + + /** + * The minimum amount of time that we should wait between subsequent + * retries, when fetching metadata. + */ + private final long refreshBackoffMs; + + /** + * The minimum amount of time that we should wait before triggering an + * automatic metadata refresh. + */ + private final long metadataExpireMs; + + /** + * Used to update the NetworkClient metadata. + */ + private final AdminMetadataUpdater updater; + + /** + * The current metadata state. + */ + private State state = State.QUIESCENT; + + /** + * The time in wall-clock milliseconds when we last updated the metadata. + */ + private long lastMetadataUpdateMs = 0; + + /** + * The time in wall-clock milliseconds when we last attempted to fetch new + * metadata. + */ + private long lastMetadataFetchAttemptMs = 0; + + /** + * The current cluster information. + */ + private Cluster cluster = Cluster.empty(); + + /** + * If we got an authorization exception when we last attempted to fetch + * metadata, this is it; null, otherwise. + */ + private AuthenticationException authException = null; + + public class AdminMetadataUpdater implements MetadataUpdater { + @Override + public List fetchNodes() { + return cluster.nodes(); + } + + @Override + public boolean isUpdateDue(long now) { + return false; + } + + @Override + public long maybeUpdate(long now) { + return Long.MAX_VALUE; + } + + @Override + public void handleServerDisconnect(long now, String destinationId, Optional maybeFatalException) { + maybeFatalException.ifPresent(AdminMetadataManager.this::updateFailed); + AdminMetadataManager.this.requestUpdate(); + } + + @Override + public void handleFailedRequest(long now, Optional maybeFatalException) { + // Do nothing + } + + @Override + public void handleSuccessfulResponse(RequestHeader requestHeader, long now, MetadataResponse metadataResponse) { + // Do nothing + } + + @Override + public void close() { + } + } + + /** + * The current AdminMetadataManager state. + */ + enum State { + QUIESCENT, + UPDATE_REQUESTED, + UPDATE_PENDING + } + + public AdminMetadataManager(LogContext logContext, long refreshBackoffMs, long metadataExpireMs) { + this.log = logContext.logger(AdminMetadataManager.class); + this.refreshBackoffMs = refreshBackoffMs; + this.metadataExpireMs = metadataExpireMs; + this.updater = new AdminMetadataUpdater(); + } + + public AdminMetadataUpdater updater() { + return updater; + } + + public boolean isReady() { + if (authException != null) { + log.debug("Metadata is not usable: failed to get metadata.", authException); + throw authException; + } + if (cluster.nodes().isEmpty()) { + log.trace("Metadata is not ready: bootstrap nodes have not been " + + "initialized yet."); + return false; + } + if (cluster.isBootstrapConfigured()) { + log.trace("Metadata is not ready: we have not fetched metadata from " + + "the bootstrap nodes yet."); + return false; + } + log.trace("Metadata is ready to use."); + return true; + } + + public Node controller() { + return cluster.controller(); + } + + public Node nodeById(int nodeId) { + return cluster.nodeById(nodeId); + } + + public void requestUpdate() { + if (state == State.QUIESCENT) { + state = State.UPDATE_REQUESTED; + log.debug("Requesting metadata update."); + } + } + + public void clearController() { + if (cluster.controller() != null) { + log.trace("Clearing cached controller node {}.", cluster.controller()); + this.cluster = new Cluster(cluster.clusterResource().clusterId(), + cluster.nodes(), + Collections.emptySet(), + Collections.emptySet(), + Collections.emptySet(), + null); + } + } + + /** + * Determine if the AdminClient should fetch new metadata. + */ + public long metadataFetchDelayMs(long now) { + switch (state) { + case QUIESCENT: + // Calculate the time remaining until the next periodic update. + // We want to avoid making many metadata requests in a short amount of time, + // so there is a metadata refresh backoff period. + return Math.max(delayBeforeNextAttemptMs(now), delayBeforeNextExpireMs(now)); + case UPDATE_REQUESTED: + // Respect the backoff, even if an update has been requested + return delayBeforeNextAttemptMs(now); + default: + // An update is already pending, so we don't need to initiate another one. + return Long.MAX_VALUE; + } + } + + private long delayBeforeNextExpireMs(long now) { + long timeSinceUpdate = now - lastMetadataUpdateMs; + return Math.max(0, metadataExpireMs - timeSinceUpdate); + } + + private long delayBeforeNextAttemptMs(long now) { + long timeSinceAttempt = now - lastMetadataFetchAttemptMs; + return Math.max(0, refreshBackoffMs - timeSinceAttempt); + } + + /** + * Transition into the UPDATE_PENDING state. Updates lastMetadataFetchAttemptMs. + */ + public void transitionToUpdatePending(long now) { + this.state = State.UPDATE_PENDING; + this.lastMetadataFetchAttemptMs = now; + } + + public void updateFailed(Throwable exception) { + // We depend on pending calls to request another metadata update + this.state = State.QUIESCENT; + + if (exception instanceof AuthenticationException) { + log.warn("Metadata update failed due to authentication error", exception); + this.authException = (AuthenticationException) exception; + } else { + log.info("Metadata update failed", exception); + } + } + + /** + * Receive new metadata, and transition into the QUIESCENT state. + * Updates lastMetadataUpdateMs, cluster, and authException. + */ + public void update(Cluster cluster, long now) { + if (cluster.isBootstrapConfigured()) { + log.debug("Setting bootstrap cluster metadata {}.", cluster); + } else { + log.debug("Updating cluster metadata to {}", cluster); + this.lastMetadataUpdateMs = now; + } + + this.state = State.QUIESCENT; + this.authException = null; + + if (!cluster.nodes().isEmpty()) { + this.cluster = cluster; + } + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/internals/ConsumerGroupOperationContext.java b/clients/src/main/java/org/apache/kafka/clients/admin/internals/ConsumerGroupOperationContext.java new file mode 100644 index 0000000000000..175c0855c97ec --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/internals/ConsumerGroupOperationContext.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.clients.admin.internals; + +import java.util.Map; +import java.util.Optional; + +import org.apache.kafka.clients.admin.AbstractOptions; +import org.apache.kafka.common.Node; +import org.apache.kafka.common.internals.KafkaFutureImpl; +import org.apache.kafka.common.protocol.Errors; +import org.apache.kafka.common.requests.AbstractResponse; + +/** + * 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. + */ +public final 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 groupId() { + return groupId; + } + + public O options() { + return options; + } + + public long deadline() { + return deadline; + } + + public KafkaFutureImpl future() { + return future; + } + + public Optional node() { + return node; + } + + public void setNode(Node node) { + this.node = Optional.ofNullable(node); + } + + public static boolean hasCoordinatorMoved(AbstractResponse response) { + return hasCoordinatorMoved(response.errorCounts()); + } + + public static boolean hasCoordinatorMoved(Map errorCounts) { + return errorCounts.containsKey(Errors.NOT_COORDINATOR); + } + + public static boolean shouldRefreshCoordinator(Map errorCounts) { + return errorCounts.containsKey(Errors.COORDINATOR_LOAD_IN_PROGRESS) || + errorCounts.containsKey(Errors.COORDINATOR_NOT_AVAILABLE); + } +} diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/internals/MetadataOperationContext.java b/clients/src/main/java/org/apache/kafka/clients/admin/internals/MetadataOperationContext.java new file mode 100644 index 0000000000000..c05e5cfac0f3c --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/admin/internals/MetadataOperationContext.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.internals; + +import java.util.Collection; +import java.util.Map; +import java.util.Optional; + +import org.apache.kafka.clients.admin.AbstractOptions; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.errors.InvalidMetadataException; +import org.apache.kafka.common.internals.KafkaFutureImpl; +import org.apache.kafka.common.protocol.Errors; +import org.apache.kafka.common.requests.MetadataResponse; +import org.apache.kafka.common.requests.MetadataResponse.PartitionMetadata; +import org.apache.kafka.common.requests.MetadataResponse.TopicMetadata; + +/** + * Context class to encapsulate parameters of a call to fetch and use cluster metadata. + * Some of the parameters are provided at construction and are immutable whereas others are provided + * as "Call" are completed and values are available. + * + * @param The type of return value of the KafkaFuture + * @param The type of configuration option. + */ +public final class MetadataOperationContext> { + final private Collection topics; + final private O options; + final private long deadline; + final private Map> futures; + private Optional response; + + public MetadataOperationContext(Collection topics, + O options, + long deadline, + Map> futures) { + this.topics = topics; + this.options = options; + this.deadline = deadline; + this.futures = futures; + this.response = Optional.empty(); + } + + public void setResponse(Optional response) { + this.response = response; + } + + public Optional response() { + return response; + } + + public O options() { + return options; + } + + public long deadline() { + return deadline; + } + + public Map> futures() { + return futures; + } + + public Collection topics() { + return topics; + } + + public static void handleMetadataErrors(MetadataResponse response) { + for (TopicMetadata tm : response.topicMetadata()) { + for (PartitionMetadata pm : tm.partitionMetadata()) { + if (shouldRefreshMetadata(pm.error)) { + throw pm.error.exception(); + } + } + } + } + + public static boolean shouldRefreshMetadata(Errors error) { + return error.exception() instanceof InvalidMetadataException; + } +} \ No newline at end of file diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/CommitFailedException.java b/clients/src/main/java/org/apache/kafka/clients/consumer/CommitFailedException.java index c6006b7ea80d7..2040216cd02c9 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/CommitFailedException.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/CommitFailedException.java @@ -28,12 +28,16 @@ public class CommitFailedException extends KafkaException { private static final long serialVersionUID = 1L; + public CommitFailedException(final String message) { + super(message); + } + public CommitFailedException() { super("Commit cannot be completed since the group has already " + "rebalanced and assigned the partitions to another member. This means that 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 message processing. " + - "You can address this either by increasing the session timeout or by reducing the maximum " + + "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."); } } 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 0e27e1f038362..ee4a70770af97 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 @@ -22,6 +22,7 @@ import org.apache.kafka.common.TopicPartition; import java.io.Closeable; +import java.time.Duration; import java.util.Collection; import java.util.List; import java.util.Map; @@ -38,156 +39,239 @@ public interface Consumer extends Closeable { /** * @see KafkaConsumer#assignment() */ - public Set assignment(); + Set assignment(); /** * @see KafkaConsumer#subscription() */ - public Set subscription(); + Set subscription(); /** * @see KafkaConsumer#subscribe(Collection) */ - public void subscribe(Collection topics); + void subscribe(Collection topics); /** * @see KafkaConsumer#subscribe(Collection, ConsumerRebalanceListener) */ - public void subscribe(Collection topics, ConsumerRebalanceListener callback); + void subscribe(Collection topics, ConsumerRebalanceListener callback); /** * @see KafkaConsumer#assign(Collection) */ - public void assign(Collection partitions); + void assign(Collection partitions); /** * @see KafkaConsumer#subscribe(Pattern, ConsumerRebalanceListener) */ - public void subscribe(Pattern pattern, ConsumerRebalanceListener callback); + void subscribe(Pattern pattern, ConsumerRebalanceListener callback); /** * @see KafkaConsumer#subscribe(Pattern) */ - public void subscribe(Pattern pattern); + void subscribe(Pattern pattern); /** * @see KafkaConsumer#unsubscribe() */ - public void unsubscribe(); + void unsubscribe(); /** * @see KafkaConsumer#poll(long) */ - public ConsumerRecords poll(long timeout); + @Deprecated + ConsumerRecords poll(long timeout); + + /** + * @see KafkaConsumer#poll(Duration) + */ + ConsumerRecords poll(Duration timeout); /** * @see KafkaConsumer#commitSync() */ - public void commitSync(); + void commitSync(); + + /** + * @see KafkaConsumer#commitSync(Duration) + */ + void commitSync(Duration timeout); /** * @see KafkaConsumer#commitSync(Map) */ - public void commitSync(Map offsets); + void commitSync(Map offsets); + /** + * @see KafkaConsumer#commitSync(Map, Duration) + */ + void commitSync(final Map offsets, final Duration timeout); /** * @see KafkaConsumer#commitAsync() */ - public void commitAsync(); + void commitAsync(); /** * @see KafkaConsumer#commitAsync(OffsetCommitCallback) */ - public void commitAsync(OffsetCommitCallback callback); + void commitAsync(OffsetCommitCallback callback); /** * @see KafkaConsumer#commitAsync(Map, OffsetCommitCallback) */ - public void commitAsync(Map offsets, OffsetCommitCallback callback); + void commitAsync(Map offsets, OffsetCommitCallback callback); /** * @see KafkaConsumer#seek(TopicPartition, long) */ - public void seek(TopicPartition partition, long offset); + void seek(TopicPartition partition, long offset); + + /** + * @see KafkaConsumer#seek(TopicPartition, OffsetAndMetadata) + */ + void seek(TopicPartition partition, OffsetAndMetadata offsetAndMetadata); /** * @see KafkaConsumer#seekToBeginning(Collection) */ - public void seekToBeginning(Collection partitions); + void seekToBeginning(Collection partitions); /** * @see KafkaConsumer#seekToEnd(Collection) */ - public void seekToEnd(Collection partitions); + void seekToEnd(Collection partitions); /** * @see KafkaConsumer#position(TopicPartition) */ - public long position(TopicPartition partition); + long position(TopicPartition partition); + + /** + * @see KafkaConsumer#position(TopicPartition, Duration) + */ + long position(TopicPartition partition, final Duration timeout); /** * @see KafkaConsumer#committed(TopicPartition) */ - public OffsetAndMetadata committed(TopicPartition partition); + @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() */ - public Map metrics(); + Map metrics(); /** * @see KafkaConsumer#partitionsFor(String) */ - public List partitionsFor(String topic); + List partitionsFor(String topic); + + /** + * @see KafkaConsumer#partitionsFor(String, Duration) + */ + List partitionsFor(String topic, Duration timeout); /** * @see KafkaConsumer#listTopics() */ - public Map> listTopics(); + Map> listTopics(); + + /** + * @see KafkaConsumer#listTopics(Duration) + */ + Map> listTopics(Duration timeout); /** * @see KafkaConsumer#paused() */ - public Set paused(); + Set paused(); /** * @see KafkaConsumer#pause(Collection) */ - public void pause(Collection partitions); + void pause(Collection partitions); /** * @see KafkaConsumer#resume(Collection) */ - public void resume(Collection partitions); + void resume(Collection partitions); + + /** + * @see KafkaConsumer#offsetsForTimes(Map) + */ + Map offsetsForTimes(Map timestampsToSearch); + + /** + * @see KafkaConsumer#offsetsForTimes(Map, Duration) + */ + Map offsetsForTimes(Map timestampsToSearch, Duration timeout); /** - * @see KafkaConsumer#offsetsForTimes(java.util.Map) + * @see KafkaConsumer#beginningOffsets(Collection) */ - public Map offsetsForTimes(Map timestampsToSearch); + Map beginningOffsets(Collection partitions); /** - * @see KafkaConsumer#beginningOffsets(java.util.Collection) + * @see KafkaConsumer#beginningOffsets(Collection, Duration) */ - public Map beginningOffsets(Collection partitions); + Map beginningOffsets(Collection partitions, Duration timeout); /** - * @see KafkaConsumer#endOffsets(java.util.Collection) + * @see KafkaConsumer#endOffsets(Collection) */ - public Map endOffsets(Collection partitions); + Map endOffsets(Collection partitions); + + /** + * @see KafkaConsumer#endOffsets(Collection, Duration) + */ + Map endOffsets(Collection partitions, Duration timeout); + + /** + * @see KafkaConsumer#groupMetadata() + */ + ConsumerGroupMetadata groupMetadata(); + + /** + * @see KafkaConsumer#enforceRebalance() + */ + void enforceRebalance(); /** * @see KafkaConsumer#close() */ - public void close(); + void close(); /** * @see KafkaConsumer#close(long, TimeUnit) */ - public void close(long timeout, TimeUnit unit); + @Deprecated + void close(long timeout, TimeUnit unit); + + /** + * @see KafkaConsumer#close(Duration) + */ + void close(Duration timeout); /** * @see KafkaConsumer#wakeup() */ - public void wakeup(); + void wakeup(); } 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 be3077ffe17d5..c90bac5466b42 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 @@ -16,21 +16,27 @@ */ package org.apache.kafka.clients.consumer; +import org.apache.kafka.clients.ClientDnsLookup; import org.apache.kafka.clients.CommonClientConfigs; +import org.apache.kafka.common.IsolationLevel; 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.errors.InvalidConfigurationException; import org.apache.kafka.common.metrics.Sensor; -import org.apache.kafka.common.requests.IsolationLevel; +import org.apache.kafka.common.requests.JoinGroupRequest; import org.apache.kafka.common.serialization.Deserializer; import java.util.Collections; import java.util.HashMap; import java.util.Locale; import java.util.Map; +import java.util.Optional; import java.util.Properties; import java.util.Set; +import java.util.concurrent.atomic.AtomicInteger; import static org.apache.kafka.common.config.ConfigDef.Range.atLeast; import static org.apache.kafka.common.config.ConfigDef.ValidString.in; @@ -49,46 +55,42 @@ 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 */ 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; + /** * enable.auto.commit */ @@ -105,7 +107,18 @@ 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 partition assignment " + + "strategies that the client will use to distribute partition " + + "ownership amongst consumer instances when group management is " + + "used.

In addition to the default class specified below, " + + "you can use the " + + "org.apache.kafka.clients.consumer.RoundRobinAssignor" + + "class for round robin assignments of partitions to consumers. " + + "

Implementing the " + + "org.apache.kafka.clients.consumer.ConsumerPartitionAssignor" + + " interface allows you to plug in a custom assignment" + + "strategy."; /** * auto.offset.reset @@ -162,6 +175,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 */ @@ -211,6 +229,12 @@ public class ConsumerConfig extends AbstractConfig { public static final String VALUE_DESERIALIZER_CLASS_CONFIG = "value.deserializer"; public static final String VALUE_DESERIALIZER_CLASS_DOC = "Deserializer class for value that implements the org.apache.kafka.common.serialization.Deserializer interface."; + /** socket.connection.setup.timeout.ms */ + public static final String SOCKET_CONNECTION_SETUP_TIMEOUT_MS_CONFIG = CommonClientConfigs.SOCKET_CONNECTION_SETUP_TIMEOUT_MS_CONFIG; + + /** socket.connection.setup.timeout.max.ms */ + public static final String SOCKET_CONNECTION_SETUP_TIMEOUT_MAX_MS_CONFIG = CommonClientConfigs.SOCKET_CONNECTION_SETUP_TIMEOUT_MAX_MS_CONFIG; + /** connections.max.idle.ms */ public static final String CONNECTIONS_MAX_IDLE_MS_CONFIG = CommonClientConfigs.CONNECTIONS_MAX_IDLE_MS_CONFIG; @@ -218,6 +242,9 @@ public class ConsumerConfig extends AbstractConfig { public static final String REQUEST_TIMEOUT_MS_CONFIG = CommonClientConfigs.REQUEST_TIMEOUT_MS_CONFIG; private static final String REQUEST_TIMEOUT_MS_DOC = CommonClientConfigs.REQUEST_TIMEOUT_MS_DOC; + /** default.api.timeout.ms */ + public static final String DEFAULT_API_TIMEOUT_MS_CONFIG = CommonClientConfigs.DEFAULT_API_TIMEOUT_MS_CONFIG; + /** interceptor.classes */ public static final String INTERCEPTOR_CLASSES_CONFIG = "interceptor.classes"; public static final String INTERCEPTOR_CLASSES_DOC = "A list of classes to use as interceptors. " @@ -227,8 +254,8 @@ public class ConsumerConfig extends AbstractConfig { /** exclude.internal.topics */ public static final String EXCLUDE_INTERNAL_TOPICS_CONFIG = "exclude.internal.topics"; - private static final String EXCLUDE_INTERNAL_TOPICS_DOC = "Whether records from internal topics (such as offsets) should be exposed to the consumer. " - + "If set to true the only way to receive records from an internal topic is subscribing to it."; + private static final String EXCLUDE_INTERNAL_TOPICS_DOC = "Whether internal topics matching a subscribed pattern should " + + "be excluded from the subscription. It is always possible to explicitly subscribe to an internal topic."; public static final boolean DEFAULT_EXCLUDE_INTERNAL_TOPICS = true; /** @@ -242,24 +269,69 @@ public class ConsumerConfig extends AbstractConfig { */ static final String LEAVE_GROUP_ON_CLOSE_CONFIG = "internal.leave.group.on.close"; + /** + * internal.throw.on.fetch.stable.offset.unsupported + * Whether or not the consumer should throw when the new stable offset feature is supported. + * If set to true then the client shall crash upon hitting it. + * The purpose of this flag is to prevent unexpected broker downgrade which makes + * the offset fetch protection against pending commit invalid. The safest approach + * is to fail fast to avoid introducing correctness issue. + * + *

+ * Note: this is an internal configuration and could be changed in the future in a backward incompatible way + * + */ + static final String THROW_ON_FETCH_STABLE_OFFSET_UNSUPPORTED = "internal.throw.on.fetch.stable.offset.unsupported"; + /** 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" + + " 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" + " return the LSO"; 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 = true; + + /** + * security.providers + */ + public static final String SECURITY_PROVIDERS_CONFIG = SecurityConfig.SECURITY_PROVIDERS_CONFIG; + private static final String SECURITY_PROVIDERS_DOC = SecurityConfig.SECURITY_PROVIDERS_DOC; + + private static final AtomicInteger CONSUMER_CLIENT_ID_SEQUENCE = new AtomicInteger(1); + static { CONFIG = new ConfigDef().define(BOOTSTRAP_SERVERS_CONFIG, Type.LIST, + Collections.emptyList(), + new ConfigDef.NonNullValidator(), Importance.HIGH, CommonClientConfigs.BOOTSTRAP_SERVERS_DOC) - .define(GROUP_ID_CONFIG, Type.STRING, "", Importance.HIGH, GROUP_ID_DOC) + .define(CLIENT_DNS_LOOKUP_CONFIG, + Type.STRING, + ClientDnsLookup.USE_ALL_DNS_IPS.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(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, @@ -273,6 +345,7 @@ public class ConsumerConfig extends AbstractConfig { .define(PARTITION_ASSIGNMENT_STRATEGY_CONFIG, Type.LIST, Collections.singletonList(RangeAssignor.class), + new ConfigDef.NonNullValidator(), Importance.MEDIUM, PARTITION_ASSIGNMENT_STRATEGY_DOC) .define(METADATA_MAX_AGE_CONFIG, @@ -297,6 +370,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, @@ -306,13 +384,13 @@ public class ConsumerConfig extends AbstractConfig { .define(SEND_BUFFER_CONFIG, Type.INT, 128 * 1024, - atLeast(-1), + atLeast(CommonClientConfigs.SEND_BUFFER_LOWER_BOUND), Importance.MEDIUM, CommonClientConfigs.SEND_BUFFER_DOC) .define(RECEIVE_BUFFER_CONFIG, Type.INT, 64 * 1024, - atLeast(-1), + atLeast(CommonClientConfigs.RECEIVE_BUFFER_LOWER_BOUND), Importance.MEDIUM, CommonClientConfigs.RECEIVE_BUFFER_DOC) .define(FETCH_MIN_BYTES_CONFIG, @@ -377,12 +455,13 @@ public class ConsumerConfig extends AbstractConfig { .define(METRICS_RECORDING_LEVEL_CONFIG, Type.STRING, Sensor.RecordingLevel.INFO.toString(), - in(Sensor.RecordingLevel.INFO.toString(), Sensor.RecordingLevel.DEBUG.toString()), + in(Sensor.RecordingLevel.INFO.toString(), Sensor.RecordingLevel.DEBUG.toString(), Sensor.RecordingLevel.TRACE.toString()), Importance.LOW, CommonClientConfigs.METRICS_RECORDING_LEVEL_DOC) .define(METRIC_REPORTER_CLASSES_CONFIG, Type.LIST, - "", + Collections.emptyList(), + new ConfigDef.NonNullValidator(), Importance.LOW, CommonClientConfigs.METRIC_REPORTER_CLASSES_DOC) .define(KEY_DESERIALIZER_CLASS_CONFIG, @@ -395,10 +474,26 @@ public class ConsumerConfig extends AbstractConfig { VALUE_DESERIALIZER_CLASS_DOC) .define(REQUEST_TIMEOUT_MS_CONFIG, Type.INT, - 305000, // chosen to be higher than the default of max.poll.interval.ms + 30000, atLeast(0), Importance.MEDIUM, REQUEST_TIMEOUT_MS_DOC) + .define(DEFAULT_API_TIMEOUT_MS_CONFIG, + Type.INT, + 60 * 1000, + atLeast(0), + Importance.MEDIUM, + CommonClientConfigs.DEFAULT_API_TIMEOUT_MS_DOC) + .define(SOCKET_CONNECTION_SETUP_TIMEOUT_MS_CONFIG, + Type.LONG, + CommonClientConfigs.DEFAULT_SOCKET_CONNECTION_SETUP_TIMEOUT_MS, + Importance.MEDIUM, + CommonClientConfigs.SOCKET_CONNECTION_SETUP_TIMEOUT_MS_DOC) + .define(SOCKET_CONNECTION_SETUP_TIMEOUT_MAX_MS_CONFIG, + Type.LONG, + CommonClientConfigs.DEFAULT_SOCKET_CONNECTION_SETUP_TIMEOUT_MAX_MS, + Importance.MEDIUM, + CommonClientConfigs.SOCKET_CONNECTION_SETUP_TIMEOUT_MAX_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(CONNECTIONS_MAX_IDLE_MS_CONFIG, Type.LONG, @@ -407,7 +502,8 @@ public class ConsumerConfig extends AbstractConfig { CommonClientConfigs.CONNECTIONS_MAX_IDLE_MS_DOC) .define(INTERCEPTOR_CLASSES_CONFIG, Type.LIST, - null, + Collections.emptyList(), + new ConfigDef.NonNullValidator(), Importance.LOW, INTERCEPTOR_CLASSES_DOC) .define(MAX_POLL_RECORDS_CONFIG, @@ -428,16 +524,30 @@ 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) + .defineInternal(THROW_ON_FETCH_STABLE_OFFSET_UNSUPPORTED, + Type.BOOLEAN, + false, + 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, @@ -445,19 +555,44 @@ public class ConsumerConfig extends AbstractConfig { CommonClientConfigs.SECURITY_PROTOCOL_DOC) .withClientSslSupport() .withClientSaslSupport(); - } @Override protected Map postProcessParsedConfig(final Map parsedValues) { - return CommonClientConfigs.postProcessReconnectBackoffConfigs(this, parsedValues); + CommonClientConfigs.warnIfDeprecatedDnsLookupValue(this); + Map refinedConfigs = CommonClientConfigs.postProcessReconnectBackoffConfigs(this, parsedValues); + maybeOverrideClientId(refinedConfigs); + return refinedConfigs; + } + + private void maybeOverrideClientId(Map configs) { + final String clientId = this.getString(CLIENT_ID_CONFIG); + if (clientId == null || clientId.isEmpty()) { + final String groupId = this.getString(GROUP_ID_CONFIG); + String groupInstanceId = this.getString(GROUP_INSTANCE_ID_CONFIG); + if (groupInstanceId != null) + JoinGroupRequest.validateGroupInstanceId(groupInstanceId); + + String groupInstanceIdPart = groupInstanceId != null ? groupInstanceId : CONSUMER_CLIENT_ID_SEQUENCE.getAndIncrement() + ""; + String generatedClientId = String.format("consumer-%s-%s", groupId, groupInstanceIdPart); + configs.put(CLIENT_ID_CONFIG, generatedClientId); + } } + /** + * @deprecated Since 2.7.0. This will be removed in a future major release. + */ + @Deprecated public static Map addDeserializerToConfig(Map configs, Deserializer keyDeserializer, Deserializer valueDeserializer) { - Map newConfigs = new HashMap(); - newConfigs.putAll(configs); + return appendDeserializerToConfig(configs, keyDeserializer, valueDeserializer); + } + + static Map appendDeserializerToConfig(Map configs, + Deserializer keyDeserializer, + Deserializer valueDeserializer) { + Map newConfigs = new HashMap<>(configs); if (keyDeserializer != null) newConfigs.put(KEY_DESERIALIZER_CLASS_CONFIG, keyDeserializer.getClass()); if (valueDeserializer != null) @@ -465,6 +600,10 @@ public static Map addDeserializerToConfig(Map co return newConfigs; } + /** + * @deprecated Since 2.7.0. This will be removed in a future major release. + */ + @Deprecated public static Properties addDeserializerToConfig(Properties properties, Deserializer keyDeserializer, Deserializer valueDeserializer) { @@ -477,11 +616,28 @@ public static Properties addDeserializerToConfig(Properties properties, return newProperties; } - ConsumerConfig(Map props) { + boolean maybeOverrideEnableAutoCommit() { + Optional groupId = Optional.ofNullable(getString(CommonClientConfigs.GROUP_ID_CONFIG)); + boolean enableAutoCommit = getBoolean(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG); + if (!groupId.isPresent()) { // overwrite in case of default group id where the config is not explicitly provided + if (!originals().containsKey(ENABLE_AUTO_COMMIT_CONFIG)) { + enableAutoCommit = false; + } else if (enableAutoCommit) { + throw new InvalidConfigurationException(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG + " cannot be set to true when default group id (null) is used."); + } + } + return enableAutoCommit; + } + + public ConsumerConfig(Properties props) { + super(CONFIG, props); + } + + public ConsumerConfig(Map props) { super(CONFIG, props); } - ConsumerConfig(Map props, boolean doLog) { + protected ConsumerConfig(Map props, boolean doLog) { super(CONFIG, props, doLog); } @@ -489,8 +645,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(4, config -> "consumerconfigs_" + config)); } } 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..f9b7b284e1fae --- /dev/null +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/ConsumerGroupMetadata.java @@ -0,0 +1,92 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF 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.requests.JoinGroupRequest; + +import java.util.Objects; +import java.util.Optional; + +/** + * A metadata struct containing the consumer group information. + * Note: Any change to this class is considered public and requires a KIP. + */ +public class ConsumerGroupMetadata { + final private String groupId; + final private int generationId; + final private String memberId; + final private Optional groupInstanceId; + + public ConsumerGroupMetadata(String groupId, + int generationId, + String memberId, + Optional groupInstanceId) { + this.groupId = Objects.requireNonNull(groupId, "group.id can't be null"); + this.generationId = generationId; + this.memberId = Objects.requireNonNull(memberId, "member.id can't be null"); + this.groupInstanceId = Objects.requireNonNull(groupInstanceId, "group.instance.id can't be null"); + } + + public ConsumerGroupMetadata(String groupId) { + this(groupId, + JoinGroupRequest.UNKNOWN_GENERATION_ID, + JoinGroupRequest.UNKNOWN_MEMBER_ID, + Optional.empty()); + } + + public String groupId() { + return groupId; + } + + public int generationId() { + return generationId; + } + + public String memberId() { + return memberId; + } + + public Optional groupInstanceId() { + return groupInstanceId; + } + + @Override + public String toString() { + return String.format("GroupMetadata(groupId = %s, generationId = %d, memberId = %s, groupInstanceId = %s)", + groupId, + generationId, + memberId, + groupInstanceId.orElse("")); + } + + @Override + public boolean equals(final Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + final ConsumerGroupMetadata that = (ConsumerGroupMetadata) o; + return generationId == that.generationId && + Objects.equals(groupId, that.groupId) && + Objects.equals(memberId, that.memberId) && + Objects.equals(groupInstanceId, that.groupInstanceId); + } + + @Override + public int hashCode() { + return Objects.hash(groupId, generationId, memberId, groupInstanceId); + } + +} diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/ConsumerInterceptor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/ConsumerInterceptor.java index 2f4e310f7025e..6af47058e4e21 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/ConsumerInterceptor.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/ConsumerInterceptor.java @@ -35,14 +35,16 @@ * the user configures the interceptor with the wrong key and value type parameters, the consumer will not throw an exception, * just log the errors. *

- * ConsumerInterceptor callbacks are called from the same thread that invokes {@link org.apache.kafka.clients.consumer.KafkaConsumer#poll(long)}. + * ConsumerInterceptor callbacks are called from the same thread that invokes + * {@link org.apache.kafka.clients.consumer.KafkaConsumer#poll(java.time.Duration)}. *

* 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. */ -public interface ConsumerInterceptor extends Configurable { +public interface ConsumerInterceptor extends Configurable, AutoCloseable { /** - * This is called just before the records are returned by {@link org.apache.kafka.clients.consumer.KafkaConsumer#poll(long)} + * This is called just before the records are returned by + * {@link org.apache.kafka.clients.consumer.KafkaConsumer#poll(java.time.Duration)} *

* This method is allowed to modify consumer records, in which case the new records will be * returned. There is no limitation on number of records that could be returned from this 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 845bff3fc4931..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 @@ -44,13 +45,48 @@ * partition is reassigned it may want to automatically trigger a flush of this cache, before the new owner takes over * consumption. *

- * This callback will execute in the user thread as part of the {@link Consumer#poll(long) poll(long)} call whenever partition assignment changes. + * 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
@@ -67,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)
@@ -79,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
-     * invocation of {@link KafkaConsumer#poll(long)} in which this callback is being executed. This means it is not
+     * 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}
      */
@@ -102,23 +145,56 @@ public interface ConsumerRebalanceListener {
 
     /**
      * A callback method the user can implement to provide handling of customized offsets on completion of a successful
-     * partition re-assignment. This method will be called after an offset re-assignment completes and before the
-     * consumer starts fetching data.
+     * 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
-     * invocation of {@link KafkaConsumer#poll(long)} in which this callback is being executed. This means it is not
+     * 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 7f85246184427..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
@@ -22,6 +22,8 @@
 import org.apache.kafka.common.record.RecordBatch;
 import org.apache.kafka.common.record.TimestampType;
 
+import java.util.Optional;
+
 /**
  * A key/value pair to be received from Kafka. This also consists of a topic name and 
  * a partition number from which the record is being received, an offset that points 
@@ -42,6 +44,7 @@ public class ConsumerRecord {
     private final Headers headers;
     private final K key;
     private final V value;
+    private final Optional leaderEpoch;
 
     private volatile Long checksum;
 
@@ -120,8 +123,43 @@ public ConsumerRecord(String topic,
                           K key,
                           V value,
                           Headers headers) {
+        this(topic, partition, offset, timestamp, timestampType, checksum, serializedKeySize, serializedValueSize,
+                key, value, headers, Optional.empty());
+    }
+
+    /**
+     * Creates a record to be received from a specified topic and partition
+     *
+     * @param topic The topic this record is received from
+     * @param partition The partition of the topic this record is received from
+     * @param offset The offset of this record in the corresponding Kafka partition
+     * @param timestamp The timestamp of the record.
+     * @param timestampType The timestamp type
+     * @param checksum The checksum (CRC32) of the full record
+     * @param serializedKeySize The length of the serialized key
+     * @param serializedValueSize The length of the serialized value
+     * @param key The key of the record, if one exists (null is allowed)
+     * @param value The record contents
+     * @param headers The headers of the record
+     * @param leaderEpoch Optional leader epoch of the record (may be empty for legacy record formats)
+     */
+    public ConsumerRecord(String topic,
+                          int partition,
+                          long offset,
+                          long timestamp,
+                          TimestampType timestampType,
+                          Long checksum,
+                          int serializedKeySize,
+                          int serializedValueSize,
+                          K key,
+                          V value,
+                          Headers headers,
+                          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;
         this.offset = offset;
@@ -133,10 +171,11 @@ public ConsumerRecord(String topic,
         this.key = key;
         this.value = value;
         this.headers = headers;
+        this.leaderEpoch = leaderEpoch;
     }
 
     /**
-     * The topic this record is received from
+     * The topic this record is received from (never null)
      */
     public String topic() {
         return this.topic;
@@ -150,7 +189,7 @@ public int partition() {
     }
 
     /**
-     * The headers
+     * The headers (never null)
      */
     public Headers headers() {
         return headers;
@@ -225,13 +264,26 @@ public int serializedValueSize() {
         return this.serializedValueSize;
     }
 
+    /**
+     * Get the leader epoch for the record if available
+     *
+     * @return the leader epoch or empty for legacy record formats
+     */
+    public Optional leaderEpoch() {
+        return leaderEpoch;
+    }
+
     @Override
     public String toString() {
-        return "ConsumerRecord(topic = " + topic() + ", partition = " + partition() + ", offset = " + offset()
+        return "ConsumerRecord(topic = " + topic
+               + ", partition = " + partition
+               + ", leaderEpoch = " + leaderEpoch.orElse(null)
+               + ", offset = " + offset
                + ", " + timestampType + " = " + timestamp
                + ", serialized key size = "  + serializedKeySize
                + ", serialized value size = " + serializedValueSize
                + ", headers = " + headers
-               + ", key = " + key + ", value = " + value + ")";
+               + ", key = " + key
+               + ", value = " + value + ")";
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/ConsumerRecords.java b/clients/src/main/java/org/apache/kafka/clients/consumer/ConsumerRecords.java
index f2dc9bbc2ba31..4d0f62c3a1d15 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/ConsumerRecords.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/ConsumerRecords.java
@@ -29,7 +29,7 @@
 /**
  * A container that holds the list {@link ConsumerRecord} per partition for a
  * particular topic. There is one {@link ConsumerRecord} list for every topic
- * partition returned by a {@link Consumer#poll(long)} operation.
+ * partition returned by a {@link Consumer#poll(java.time.Duration)} operation.
  */
 public class ConsumerRecords implements Iterable> {
 
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 e9499cbd465bf..b6bebc1717ae0 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
@@ -17,44 +17,52 @@
 package org.apache.kafka.clients.consumer;
 
 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;
-import org.apache.kafka.clients.consumer.internals.ConsumerMetrics;
+import org.apache.kafka.clients.consumer.internals.ConsumerMetadata;
 import org.apache.kafka.clients.consumer.internals.ConsumerNetworkClient;
-import org.apache.kafka.clients.consumer.internals.ConsumerNetworkClient.PollCondition;
 import org.apache.kafka.clients.consumer.internals.Fetcher;
+import org.apache.kafka.clients.consumer.internals.FetcherMetricsRegistry;
+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.IsolationLevel;
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.Metric;
 import org.apache.kafka.common.MetricName;
 import org.apache.kafka.common.PartitionInfo;
 import org.apache.kafka.common.TopicPartition;
-import org.apache.kafka.common.config.ConfigException;
 import org.apache.kafka.common.errors.InterruptException;
+import org.apache.kafka.common.errors.InvalidGroupIdException;
+import org.apache.kafka.common.errors.TimeoutException;
 import org.apache.kafka.common.internals.ClusterResourceListeners;
 import org.apache.kafka.common.metrics.JmxReporter;
+import org.apache.kafka.common.metrics.KafkaMetricsContext;
 import org.apache.kafka.common.metrics.MetricConfig;
 import org.apache.kafka.common.metrics.Metrics;
+import org.apache.kafka.common.metrics.MetricsContext;
 import org.apache.kafka.common.metrics.MetricsReporter;
 import org.apache.kafka.common.metrics.Sensor;
 import org.apache.kafka.common.network.ChannelBuilder;
 import org.apache.kafka.common.network.Selector;
-import org.apache.kafka.common.requests.IsolationLevel;
 import org.apache.kafka.common.requests.MetadataRequest;
 import org.apache.kafka.common.serialization.Deserializer;
 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.Timer;
 import org.apache.kafka.common.utils.Utils;
 import org.slf4j.Logger;
 
 import java.net.InetSocketAddress;
+import java.time.Duration;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.ConcurrentModificationException;
@@ -63,6 +71,8 @@
 import java.util.List;
 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;
@@ -71,6 +81,8 @@
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.regex.Pattern;
 
+import static org.apache.kafka.clients.consumer.internals.PartitionAssignorAdapter.getAssignorInstances;
+
 /**
  * A client that consumes records from a Kafka cluster.
  * 

@@ -97,7 +109,7 @@
  * 

  * The {@link #position(TopicPartition) position} of the consumer gives the offset of the next record that will be given
  * out. It will be one larger than the highest offset the consumer has seen in that partition. It automatically advances
- * every time the consumer receives messages in a call to {@link #poll(long)}.
+ * every time the consumer receives messages in a call to {@link #poll(Duration)}.
  * 

  * The {@link #commitSync() committed position} is the last offset that has been stored securely. Should the
  * process fail and restart, this is the offset that the consumer will recover to. The consumer can either automatically commit
@@ -148,7 +160,7 @@
  *
  * 
Detecting Consumer Failures

  *
- * After subscribing to a set of topics, the consumer will automatically join the group when {@link #poll(long)} is
+ * After subscribing to a set of topics, the consumer will automatically join the group when {@link #poll(Duration)} is
  * invoked. The poll API is designed to ensure consumer liveness. As long as you continue to call poll, the consumer
  * will stay in the group and continue to receive messages from the partitions it was assigned. Underneath the covers,
  * the consumer sends periodic heartbeats to the server. If the consumer crashes or is unable to send heartbeats for
@@ -167,10 +179,10 @@
  * The consumer provides two configuration settings to control the behavior of the poll loop:
  * 
    
  *     
  1. max.poll.interval.ms: By increasing the interval between expected polls, you can give
- *     the consumer more time to handle a batch of records returned from {@link #poll(long)}. The drawback
+ *     the consumer more time to handle a batch of records returned from {@link #poll(Duration)}. The drawback
  *     is that increasing this value may delay a group rebalance since the consumer will only join the rebalance
  *     inside the call to poll. You can use this setting to bound the time to finish a rebalance, but
- *     you risk slower progress if the consumer cannot actually call {@link #poll(long) poll} often enough.
    2. 
+ *     you risk slower progress if the consumer cannot actually call {@link #poll(Duration) poll} often enough.
  *     
  2. max.poll.records: Use this setting to limit the total records returned from a single
  *     call to poll. This can make it easier to predict the maximum that must be handled within each poll
  *     interval. By tuning this value, you may be able to reduce the poll interval, which will reduce the
@@ -179,32 +191,32 @@
  * 
    
  * For use cases where message processing time varies unpredictably, neither of these options may be sufficient.
  * The recommended way to handle these cases is to move message processing to another thread, which allows
- * the consumer to continue calling {@link #poll(long) poll} while the processor is still working. Some care must be taken
- * to ensure that committed offsets do not get ahead of the actual position. Typically, you must disable automatic
- * commits and manually commit processed offsets for records only after the thread has finished handling them
- * (depending on the delivery semantics you need). Note also that you will need to {@link #pause(Collection) pause}
- * the partition so that no new records are received from poll until after thread has finished handling those
- * previously returned.
+ * the consumer to continue calling {@link #poll(Duration) poll} while the processor is still working.
+ * Some care must be taken to ensure that committed offsets do not get ahead of the actual position.
+ * Typically, you must disable automatic commits and manually commit processed offsets for records only after the
+ * thread has finished handling them (depending on the delivery semantics you need).
+ * Note also that you will need to {@link #pause(Collection) pause} the partition so that no new records are received
+ * from poll until after thread has finished handling those previously returned.
  *
  * 
    Usage Examples
    
  * The consumer APIs offer flexibility to cover a variety of consumption use cases. Here are some examples to
  * demonstrate how to use them.
  *
  * 
    Automatic Offset Committing
    
- * This example demonstrates a simple usage of Kafka's consumer api that relying on automatic offset committing.
+ * This example demonstrates a simple usage of Kafka's consumer api that relies on automatic offset committing.
  * 
    
  * 
      *     Properties props = new Properties();
- *     props.put("bootstrap.servers", "localhost:9092");
- *     props.put("group.id", "test");
- *     props.put("enable.auto.commit", "true");
- *     props.put("auto.commit.interval.ms", "1000");
- *     props.put("key.deserializer", "org.apache.kafka.common.serialization.StringDeserializer");
- *     props.put("value.deserializer", "org.apache.kafka.common.serialization.StringDeserializer");
+ *     props.setProperty("bootstrap.servers", "localhost:9092");
+ *     props.setProperty("group.id", "test");
+ *     props.setProperty("enable.auto.commit", "true");
+ *     props.setProperty("auto.commit.interval.ms", "1000");
+ *     props.setProperty("key.deserializer", "org.apache.kafka.common.serialization.StringDeserializer");
+ *     props.setProperty("value.deserializer", "org.apache.kafka.common.serialization.StringDeserializer");
  *     KafkaConsumer<String, String> consumer = new KafkaConsumer<>(props);
  *     consumer.subscribe(Arrays.asList("foo", "bar"));
  *     while (true) {
- *         ConsumerRecords<String, String> records = consumer.poll(100);
+ *         ConsumerRecords<String, String> records = consumer.poll(Duration.ofMillis(100));
  *         for (ConsumerRecord<String, String> record : records)
  *             System.out.printf("offset = %d, key = %s, value = %s%n", record.offset(), record.key(), record.value());
  *     }
@@ -233,17 +245,17 @@
  * 
    
  * 
      *     Properties props = new Properties();
- *     props.put("bootstrap.servers", "localhost:9092");
- *     props.put("group.id", "test");
- *     props.put("enable.auto.commit", "false");
- *     props.put("key.deserializer", "org.apache.kafka.common.serialization.StringDeserializer");
- *     props.put("value.deserializer", "org.apache.kafka.common.serialization.StringDeserializer");
+ *     props.setProperty("bootstrap.servers", "localhost:9092");
+ *     props.setProperty("group.id", "test");
+ *     props.setProperty("enable.auto.commit", "false");
+ *     props.setProperty("key.deserializer", "org.apache.kafka.common.serialization.StringDeserializer");
+ *     props.setProperty("value.deserializer", "org.apache.kafka.common.serialization.StringDeserializer");
  *     KafkaConsumer<String, String> consumer = new KafkaConsumer<>(props);
  *     consumer.subscribe(Arrays.asList("foo", "bar"));
  *     final int minBatchSize = 200;
  *     List<ConsumerRecord<String, String>> buffer = new ArrayList<>();
  *     while (true) {
- *         ConsumerRecords<String, String> records = consumer.poll(100);
+ *         ConsumerRecords<String, String> records = consumer.poll(Duration.ofMillis(100));
  *         for (ConsumerRecord<String, String> record : records) {
  *             buffer.add(record);
  *         }
@@ -257,7 +269,8 @@
  *
  * In this example we will consume a batch of records and batch them up in memory. When we have enough records
  * batched, we will insert them into a database. If we allowed offsets to auto commit as in the previous example, records
- * would be considered consumed after they were returned to the user in {@link #poll(long) poll}. It would then be possible
+ * would be considered consumed after they were returned to the user in {@link #poll(Duration) poll}. It would then be
+ * possible
  * for our process to fail after batching the records, but before they had been inserted into the database.
  * 
    
  * To avoid this, we will manually commit the offsets only after the corresponding records have been inserted into the
@@ -269,7 +282,7 @@
  * time but in failure cases could be duplicated.
  * 
    
  * Note: Using automatic offset commits can also give you "at-least-once" delivery, but the requirement is that
- * you must consume all data returned from each call to {@link #poll(long)} before any subsequent calls, or before
+ * you must consume all data returned from each call to {@link #poll(Duration)} before any subsequent calls, or before
  * {@link #close() closing} the consumer. If you fail to do either of these, it is possible for the committed offset
  * to get ahead of the consumed position, which results in missing records. The advantage of using manual offset
  * control is that you have direct control over when a record is considered "consumed."
@@ -281,7 +294,7 @@
  * 
      *     try {
  *         while(running) {
- *             ConsumerRecords<String, String> records = consumer.poll(Long.MAX_VALUE);
+ *             ConsumerRecords<String, String> records = consumer.poll(Duration.ofMillis(Long.MAX_VALUE));
  *             for (TopicPartition partition : records.partitions()) {
  *                 List<ConsumerRecord<String, String>> partitionRecords = records.records(partition);
  *                 for (ConsumerRecord<String, String> record : partitionRecords) {
@@ -324,7 +337,7 @@
  *     consumer.assign(Arrays.asList(partition0, partition1));
  * 
    
  *
- * Once assigned, you can call {@link #poll(long) poll} in a loop, just as in the preceding examples to consume
+ * Once assigned, you can call {@link #poll(Duration) poll} in a loop, just as in the preceding examples to consume
  * records. The group that the consumer specifies is still used for committing offsets, but now the set of partitions
  * will only change with another call to {@link #assign(Collection) assign}. Manual partition assignment does
  * not use group coordination, so consumer failures will not cause assigned partitions to be rebalanced. Each consumer
@@ -416,7 +429,7 @@
  * 
    
  * Kafka supports dynamic controlling of consumption flows by using {@link #pause(Collection)} and {@link #resume(Collection)}
  * to pause the consumption on the specified assigned partitions and resume the consumption
- * on the specified paused partitions respectively in the future {@link #poll(long)} calls.
+ * on the specified paused partitions respectively in the future {@link #poll(Duration)} calls.
  *
  * 
    Reading Transactional Messages
    
  *
@@ -463,11 +476,16 @@
  *     private final AtomicBoolean closed = new AtomicBoolean(false);
  *     private final KafkaConsumer consumer;
  *
+ *     public KafkaConsumerRunner(KafkaConsumer consumer) {
+ *       this.consumer = consumer;
+ *     }
+ *
+ *     {@literal}@Override
  *     public void run() {
  *         try {
  *             consumer.subscribe(Arrays.asList("topic"));
  *             while (!closed.get()) {
- *                 ConsumerRecords records = consumer.poll(10000);
+ *                 ConsumerRecords records = consumer.poll(Duration.ofMillis(10000));
  *                 // Handle new records
  *             }
  *         } catch (WakeupException e) {
@@ -542,16 +560,18 @@
  */
 public class KafkaConsumer implements Consumer {
 
+    private static final String CLIENT_ID_METRIC_TAG = "client-id";
     private static final long NO_CURRENT_THREAD = -1L;
-    private static final AtomicInteger CONSUMER_CLIENT_ID_SEQUENCE = new AtomicInteger(1);
     private static final String JMX_PREFIX = "kafka.consumer";
     static final long DEFAULT_CLOSE_TIMEOUT_MS = 30 * 1000;
 
     // Visible for testing
     final Metrics metrics;
+    final KafkaConsumerMetrics kafkaConsumerMetrics;
 
-    private final Logger log;
+    private Logger log;
     private final String clientId;
+    private final Optional groupId;
     private final ConsumerCoordinator coordinator;
     private final Deserializer keyDeserializer;
     private final Deserializer valueDeserializer;
@@ -561,11 +581,12 @@ public class KafkaConsumer implements Consumer {
     private final Time time;
     private final ConsumerNetworkClient client;
     private final SubscriptionState subscriptions;
-    private final Metadata metadata;
+    private final ConsumerMetadata metadata;
     private final long retryBackoffMs;
     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
@@ -573,13 +594,18 @@ public class KafkaConsumer implements Consumer {
     // refcount is used to allow reentrant access by the thread who has acquired currentThread
     private final AtomicInteger refcount = new AtomicInteger(0);
 
+    // to keep from repeatedly scanning subscriptions in poll(), cache the result during metadata updates
+    private boolean cachedSubscriptionHashAllFetchPositions;
+
     /**
      * A consumer is instantiated by providing a set of key-value pairs as configuration. Valid configuration strings
      * are documented here. Values can be
      * either strings or objects of the appropriate type (for example a numeric configuration would accept either the
      * string "42" or the integer 42).
      * 
    
-     * Valid configuration strings are documented at {@link ConsumerConfig}
+     * Valid configuration strings are documented at {@link ConsumerConfig}.
+     * 
    
+     * Note: after creating a {@code KafkaConsumer} you must always {@link #close()} it to avoid resource leaks.
      *
      * @param configs The consumer configs
      */
@@ -587,29 +613,12 @@ public KafkaConsumer(Map configs) {
         this(configs, null, null);
     }
 
-    /**
-     * A consumer is instantiated by providing a set of key-value pairs as configuration, and a key and a value {@link Deserializer}.
-     * 
    
-     * Valid configuration strings are documented at {@link ConsumerConfig}
-     *
-     * @param configs The consumer configs
-     * @param keyDeserializer The deserializer for key that implements {@link Deserializer}. The configure() method
-     *            won't be called in the consumer when the deserializer is passed in directly.
-     * @param valueDeserializer The deserializer for value that implements {@link Deserializer}. The configure() method
-     *            won't be called in the consumer when the deserializer is passed in directly.
-     */
-    public KafkaConsumer(Map configs,
-                         Deserializer keyDeserializer,
-                         Deserializer valueDeserializer) {
-        this(new ConsumerConfig(ConsumerConfig.addDeserializerToConfig(configs, keyDeserializer, valueDeserializer)),
-            keyDeserializer,
-            valueDeserializer);
-    }
-
     /**
      * A consumer is instantiated by providing a {@link java.util.Properties} object as configuration.
      * 
    
-     * Valid configuration strings are documented at {@link ConsumerConfig}
+     * Valid configuration strings are documented at {@link ConsumerConfig}.
+     * 
    
+     * Note: after creating a {@code KafkaConsumer} you must always {@link #close()} it to avoid resource leaks.
      *
      * @param properties The consumer configuration properties
      */
@@ -621,7 +630,9 @@ public KafkaConsumer(Properties properties) {
      * A consumer is instantiated by providing a {@link java.util.Properties} object as configuration, and a
      * key and a value {@link Deserializer}.
      * 
    
-     * Valid configuration strings are documented at {@link ConsumerConfig}
+     * Valid configuration strings are documented at {@link ConsumerConfig}.
+     * 
    
+     * Note: after creating a {@code KafkaConsumer} you must always {@link #close()} it to avoid resource leaks.
      *
      * @param properties The consumer configuration properties
      * @param keyDeserializer The deserializer for key that implements {@link Deserializer}. The configure() method
@@ -632,81 +643,104 @@ public KafkaConsumer(Properties properties) {
     public KafkaConsumer(Properties properties,
                          Deserializer keyDeserializer,
                          Deserializer valueDeserializer) {
-        this(new ConsumerConfig(ConsumerConfig.addDeserializerToConfig(properties, keyDeserializer, valueDeserializer)),
-             keyDeserializer,
-             valueDeserializer);
+        this(Utils.propsToMap(properties), keyDeserializer, valueDeserializer);
+    }
+
+    /**
+     * A consumer is instantiated by providing a set of key-value pairs as configuration, and a key and a value {@link Deserializer}.
+     * 
    
+     * Valid configuration strings are documented at {@link ConsumerConfig}.
+     * 
    
+     * Note: after creating a {@code KafkaConsumer} you must always {@link #close()} it to avoid resource leaks.
+     *
+     * @param configs The consumer configs
+     * @param keyDeserializer The deserializer for key that implements {@link Deserializer}. The configure() method
+     *            won't be called in the consumer when the deserializer is passed in directly.
+     * @param valueDeserializer The deserializer for value that implements {@link Deserializer}. The configure() method
+     *            won't be called in the consumer when the deserializer is passed in directly.
+     */
+    public KafkaConsumer(Map configs,
+                         Deserializer keyDeserializer,
+                         Deserializer valueDeserializer) {
+        this(new ConsumerConfig(ConsumerConfig.appendDeserializerToConfig(configs, keyDeserializer, valueDeserializer)),
+                keyDeserializer, valueDeserializer);
     }
 
     @SuppressWarnings("unchecked")
-    private KafkaConsumer(ConsumerConfig config,
-                          Deserializer keyDeserializer,
-                          Deserializer valueDeserializer) {
+    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;
-            String groupId = config.getString(ConsumerConfig.GROUP_ID_CONFIG);
+            GroupRebalanceConfig groupRebalanceConfig = new GroupRebalanceConfig(config,
+                    GroupRebalanceConfig.ProtocolType.CONSUMER);
+
+            this.groupId = Optional.ofNullable(groupRebalanceConfig.groupId);
+            this.clientId = config.getString(CommonClientConfigs.CLIENT_ID_CONFIG);
+
+            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.orElse("null") + "] ");
+            } else {
+                logContext = new LogContext("[Consumer clientId=" + clientId + ", groupId=" + groupId.orElse("null") + "] ");
+            }
 
-            LogContext logContext = new LogContext("[Consumer clientId=" + clientId + ", groupId=" + groupId + "] ");
             this.log = logContext.logger(getClass());
+            boolean enableAutoCommit = config.maybeOverrideEnableAutoCommit();
+            groupId.ifPresent(groupIdStr -> {
+                if (groupIdStr.isEmpty()) {
+                    log.warn("Support for using the empty group id by consumers is deprecated and will be removed in the next major release.");
+                }
+            });
 
             log.debug("Initializing the Kafka consumer");
             this.requestTimeoutMs = config.getInt(ConsumerConfig.REQUEST_TIMEOUT_MS_CONFIG);
-            int sessionTimeOutMs = config.getInt(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG);
-            int fetchMaxWaitMs = config.getInt(ConsumerConfig.FETCH_MAX_WAIT_MS_CONFIG);
-            if (this.requestTimeoutMs <= sessionTimeOutMs || this.requestTimeoutMs <= fetchMaxWaitMs)
-                throw new ConfigException(ConsumerConfig.REQUEST_TIMEOUT_MS_CONFIG + " should be greater than " + ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG + " and " + ConsumerConfig.FETCH_MAX_WAIT_MS_CONFIG);
+            this.defaultApiTimeoutMs = config.getInt(ConsumerConfig.DEFAULT_API_TIMEOUT_MS_CONFIG);
             this.time = Time.SYSTEM;
-
-            Map metricsTags = Collections.singletonMap("client-id", clientId);
-            MetricConfig metricConfig = new MetricConfig().samples(config.getInt(ConsumerConfig.METRICS_NUM_SAMPLES_CONFIG))
-                    .timeWindow(config.getLong(ConsumerConfig.METRICS_SAMPLE_WINDOW_MS_CONFIG), TimeUnit.MILLISECONDS)
-                    .recordLevel(Sensor.RecordingLevel.forName(config.getString(ConsumerConfig.METRICS_RECORDING_LEVEL_CONFIG)))
-                    .tags(metricsTags);
-            List reporters = config.getConfiguredInstances(ConsumerConfig.METRIC_REPORTER_CLASSES_CONFIG,
-                    MetricsReporter.class);
-            reporters.add(new JmxReporter(JMX_PREFIX));
-            this.metrics = new Metrics(metricConfig, reporters, time);
+            this.metrics = buildMetrics(config, time, clientId);
             this.retryBackoffMs = config.getLong(ConsumerConfig.RETRY_BACKOFF_MS_CONFIG);
 
-            // load interceptors and make sure they get clientId
-            Map userProvidedConfigs = config.originals();
-            userProvidedConfigs.put(ConsumerConfig.CLIENT_ID_CONFIG, clientId);
-            List> interceptorList = (List) (new ConsumerConfig(userProvidedConfigs, false)).getConfiguredInstances(ConsumerConfig.INTERCEPTOR_CLASSES_CONFIG,
-                    ConsumerInterceptor.class);
-            this.interceptors = interceptorList.isEmpty() ? null : new ConsumerInterceptors<>(interceptorList);
+            List> interceptorList = (List) config.getConfiguredInstances(
+                    ConsumerConfig.INTERCEPTOR_CLASSES_CONFIG,
+                    ConsumerInterceptor.class,
+                    Collections.singletonMap(ConsumerConfig.CLIENT_ID_CONFIG, clientId));
+            this.interceptors = new ConsumerInterceptors<>(interceptorList);
             if (keyDeserializer == null) {
-                this.keyDeserializer = config.getConfiguredInstance(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG,
-                        Deserializer.class);
-                this.keyDeserializer.configure(config.originals(), true);
+                this.keyDeserializer = config.getConfiguredInstance(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, Deserializer.class);
+                this.keyDeserializer.configure(config.originals(Collections.singletonMap(ConsumerConfig.CLIENT_ID_CONFIG, clientId)), true);
             } else {
                 config.ignore(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG);
                 this.keyDeserializer = keyDeserializer;
             }
             if (valueDeserializer == null) {
-                this.valueDeserializer = config.getConfiguredInstance(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG,
-                        Deserializer.class);
-                this.valueDeserializer.configure(config.originals(), false);
+                this.valueDeserializer = config.getConfiguredInstance(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, Deserializer.class);
+                this.valueDeserializer.configure(config.originals(Collections.singletonMap(ConsumerConfig.CLIENT_ID_CONFIG, clientId)), false);
             } else {
                 config.ignore(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG);
                 this.valueDeserializer = valueDeserializer;
             }
-            ClusterResourceListeners clusterResourceListeners = configureClusterResourceListeners(keyDeserializer, valueDeserializer, reporters, interceptorList);
-            this.metadata = new Metadata(retryBackoffMs, config.getLong(ConsumerConfig.METADATA_MAX_AGE_CONFIG),
-                    true, false, clusterResourceListeners);
-            List addresses = ClientUtils.parseAndValidateAddresses(config.getList(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG));
-            this.metadata.update(Cluster.bootstrap(addresses), Collections.emptySet(), 0);
+            OffsetResetStrategy offsetResetStrategy = OffsetResetStrategy.valueOf(config.getString(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG).toUpperCase(Locale.ROOT));
+            this.subscriptions = new SubscriptionState(logContext, offsetResetStrategy);
+            ClusterResourceListeners clusterResourceListeners = configureClusterResourceListeners(keyDeserializer,
+                    valueDeserializer, metrics.reporters(), interceptorList);
+            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);
             String metricGrpPrefix = "consumer";
-            ConsumerMetrics metricsRegistry = new ConsumerMetrics(metricsTags.keySet(), "consumer");
-            ChannelBuilder channelBuilder = ClientUtils.createChannelBuilder(config);
 
+            FetcherMetricsRegistry metricsRegistry = new FetcherMetricsRegistry(Collections.singleton(CLIENT_ID_METRIC_TAG), metricGrpPrefix);
+            ChannelBuilder channelBuilder = ClientUtils.createChannelBuilder(config, time, logContext);
             IsolationLevel isolationLevel = IsolationLevel.valueOf(
                     config.getString(ConsumerConfig.ISOLATION_LEVEL_CONFIG).toUpperCase(Locale.ROOT));
-            Sensor throttleTimeSensor = Fetcher.throttleTimeSensor(metrics, metricsRegistry.fetcherMetrics);
-
+            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,
@@ -717,9 +751,12 @@ private KafkaConsumer(ConsumerConfig config,
                     config.getInt(ConsumerConfig.SEND_BUFFER_CONFIG),
                     config.getInt(ConsumerConfig.RECEIVE_BUFFER_CONFIG),
                     config.getInt(ConsumerConfig.REQUEST_TIMEOUT_MS_CONFIG),
+                    config.getLong(ConsumerConfig.SOCKET_CONNECTION_SETUP_TIMEOUT_MS_CONFIG),
+                    config.getLong(ConsumerConfig.SOCKET_CONNECTION_SETUP_TIMEOUT_MAX_MS_CONFIG),
+                    ClientDnsLookup.forConfig(config.getString(ConsumerConfig.CLIENT_DNS_LOOKUP_CONFIG)),
                     time,
                     true,
-                    new ApiVersions(),
+                    apiVersions,
                     throttleTimeSensor,
                     logContext);
             this.client = new ConsumerNetworkClient(
@@ -730,29 +767,25 @@ private KafkaConsumer(ConsumerConfig config,
                     retryBackoffMs,
                     config.getInt(ConsumerConfig.REQUEST_TIMEOUT_MS_CONFIG),
                     heartbeatIntervalMs); //Will avoid blocking an extended period of time to prevent heartbeat thread starvation
-            OffsetResetStrategy offsetResetStrategy = OffsetResetStrategy.valueOf(config.getString(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG).toUpperCase(Locale.ROOT));
-            this.subscriptions = new SubscriptionState(offsetResetStrategy);
-            this.assignors = config.getConfiguredInstances(
-                    ConsumerConfig.PARTITION_ASSIGNMENT_STRATEGY_CONFIG,
-                    PartitionAssignor.class);
-            this.coordinator = new ConsumerCoordinator(logContext,
-                    this.client,
-                    groupId,
-                    config.getInt(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG),
-                    config.getInt(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG),
-                    heartbeatIntervalMs,
-                    assignors,
-                    this.metadata,
-                    this.subscriptions,
-                    metrics,
-                    metricGrpPrefix,
-                    this.time,
-                    retryBackoffMs,
-                    config.getBoolean(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG),
-                    config.getInt(ConsumerConfig.AUTO_COMMIT_INTERVAL_MS_CONFIG),
-                    this.interceptors,
-                    config.getBoolean(ConsumerConfig.EXCLUDE_INTERNAL_TOPICS_CONFIG),
-                    config.getBoolean(ConsumerConfig.LEAVE_GROUP_ON_CLOSE_CONFIG));
+
+            this.assignors = getAssignorInstances(config.getList(ConsumerConfig.PARTITION_ASSIGNMENT_STRATEGY_CONFIG),
+                    config.originals(Collections.singletonMap(ConsumerConfig.CLIENT_ID_CONFIG, clientId)));
+
+            // no coordinator will be constructed for the default (null) group id
+            this.coordinator = !groupId.isPresent() ? null :
+                new ConsumerCoordinator(groupRebalanceConfig,
+                        logContext,
+                        this.client,
+                        assignors,
+                        this.metadata,
+                        this.subscriptions,
+                        metrics,
+                        metricGrpPrefix,
+                        this.time,
+                        enableAutoCommit,
+                        config.getInt(ConsumerConfig.AUTO_COMMIT_INTERVAL_MS_CONFIG),
+                        this.interceptors,
+                        config.getBoolean(ConsumerConfig.THROW_ON_FETCH_STABLE_OFFSET_UNSUPPORTED));
             this.fetcher = new Fetcher<>(
                     logContext,
                     this.client,
@@ -762,24 +795,30 @@ private KafkaConsumer(ConsumerConfig config,
                     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),
                     this.keyDeserializer,
                     this.valueDeserializer,
                     this.metadata,
                     this.subscriptions,
                     metrics,
-                    metricsRegistry.fetcherMetrics,
+                    metricsRegistry,
                     this.time,
                     this.retryBackoffMs,
-                    isolationLevel);
+                    this.requestTimeoutMs,
+                    isolationLevel,
+                    apiVersions);
 
-            config.logUnused();
-            AppInfoParser.registerAppInfo(JMX_PREFIX, clientId, metrics);
+            this.kafkaConsumerMetrics = new KafkaConsumerMetrics(metrics, metricGrpPrefix);
 
+            config.logUnused();
+            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
-            close(0, true);
+            // call close methods if internal objects are already constructed; this is to prevent resource leak. see KAFKA-2121
+            // we do not need to call `close` at all when `log` is null, which means no internal objects were initialized.
+            if (this.log != null) {
+                close(0, true);
+            }
             // now propagate the exception
             throw new KafkaException("Failed to construct kafka consumer", t);
         }
@@ -797,17 +836,19 @@ private KafkaConsumer(ConsumerConfig config,
                   ConsumerNetworkClient client,
                   Metrics metrics,
                   SubscriptionState subscriptions,
-                  Metadata metadata,
+                  ConsumerMetadata metadata,
                   long retryBackoffMs,
                   long requestTimeoutMs,
-                  List assignors) {
+                  int defaultApiTimeoutMs,
+                  List assignors,
+                  String groupId) {
         this.log = logContext.logger(getClass());
         this.clientId = clientId;
         this.coordinator = coordinator;
         this.keyDeserializer = keyDeserializer;
         this.valueDeserializer = valueDeserializer;
         this.fetcher = fetcher;
-        this.interceptors = interceptors;
+        this.interceptors = Objects.requireNonNull(interceptors);
         this.time = time;
         this.client = client;
         this.metrics = metrics;
@@ -815,7 +856,26 @@ private KafkaConsumer(ConsumerConfig config,
         this.metadata = metadata;
         this.retryBackoffMs = retryBackoffMs;
         this.requestTimeoutMs = requestTimeoutMs;
+        this.defaultApiTimeoutMs = defaultApiTimeoutMs;
         this.assignors = assignors;
+        this.groupId = Optional.ofNullable(groupId);
+        this.kafkaConsumerMetrics = new KafkaConsumerMetrics(metrics, "consumer");
+    }
+
+    private static Metrics buildMetrics(ConsumerConfig config, Time time, String clientId) {
+        Map metricsTags = Collections.singletonMap(CLIENT_ID_METRIC_TAG, clientId);
+        MetricConfig metricConfig = new MetricConfig().samples(config.getInt(ConsumerConfig.METRICS_NUM_SAMPLES_CONFIG))
+                .timeWindow(config.getLong(ConsumerConfig.METRICS_SAMPLE_WINDOW_MS_CONFIG), TimeUnit.MILLISECONDS)
+                .recordLevel(Sensor.RecordingLevel.forName(config.getString(ConsumerConfig.METRICS_RECORDING_LEVEL_CONFIG)))
+                .tags(metricsTags);
+        List reporters = config.getConfiguredInstances(ConsumerConfig.METRIC_REPORTER_CLASSES_CONFIG,
+                MetricsReporter.class, Collections.singletonMap(ConsumerConfig.CLIENT_ID_CONFIG, clientId));
+        JmxReporter jmxReporter = new JmxReporter();
+        jmxReporter.configure(config.originals(Collections.singletonMap(ConsumerConfig.CLIENT_ID_CONFIG, clientId)));
+        reporters.add(jmxReporter);
+        MetricsContext metricsContext = new KafkaMetricsContext(JMX_PREFIX,
+                config.originalsWithPrefix(CommonClientConfigs.METRICS_CONTEXT_PREFIX));
+        return new Metrics(metricConfig, reporters, time, metricsContext);
     }
 
     /**
@@ -829,7 +889,7 @@ private KafkaConsumer(ConsumerConfig config,
     public Set assignment() {
         acquireAndEnsureOpen();
         try {
-            return Collections.unmodifiableSet(new HashSet<>(this.subscriptions.assignedPartitions()));
+            return Collections.unmodifiableSet(this.subscriptions.assignedPartitions());
         } finally {
             release();
         }
@@ -859,17 +919,20 @@ public Set subscription() {
      *
      * 
    
      * As part of group management, the consumer will keep track of the list of consumers that belong to a particular
-     * group and will trigger a rebalance operation if one of the following events trigger -
+     * group and will trigger a rebalance operation if any one of the following events are triggered:
      * 
      
-     * 
    • Number of partitions change for any of the subscribed list of topics
-     * 
    • Topic is created or deleted
-     * 
    • An existing member of the consumer group dies
-     * 
    • A new member is added to an existing consumer group via the join API
+     * 
    • Number of partitions change for any of the subscribed topics
+     * 
    • A subscribed topic is created or deleted
+     * 
    • An existing member of the consumer group is shutdown or fails
+     * 
    • A new member is added to the consumer group
      * 
    
      * 
    
      * When any of these events are triggered, the provided listener will be invoked first to indicate that
      * the consumer's assignment has been revoked, and then again when the new assignment has been received.
-     * Note that this listener will immediately override any listener set in a previous call to subscribe.
+     * Note that rebalances will only occur during an active call to {@link #poll(Duration)}, so callbacks will
+     * also only be invoked during that time.
+     *
+     * The provided listener will immediately override any listener set in a previous call to subscribe.
      * It is guaranteed, however, that the partitions revoked/assigned through this interface are from topics
      * subscribed in this call. See {@link ConsumerRebalanceListener} for more details.
      *
@@ -885,9 +948,10 @@ public Set subscription() {
     public void subscribe(Collection topics, ConsumerRebalanceListener listener) {
         acquireAndEnsureOpen();
         try {
-            if (topics == null) {
+            maybeThrowInvalidGroupIdException();
+            if (topics == null)
                 throw new IllegalArgumentException("Topic collection to subscribe to cannot be null");
-            } else if (topics.isEmpty()) {
+            if (topics.isEmpty()) {
                 // treat subscribing to empty topic list as the same as unsubscribing
                 this.unsubscribe();
             } else {
@@ -897,10 +961,10 @@ public void subscribe(Collection topics, ConsumerRebalanceListener liste
                 }
 
                 throwIfNoAssignorsConfigured();
-
-                log.debug("Subscribed to topic(s): {}", Utils.join(topics, ", "));
-                this.subscriptions.subscribe(new HashSet<>(topics), listener);
-                metadata.setTopics(subscriptions.groupSubscription());
+                fetcher.clearBufferedDataForUnassignedTopics(topics);
+                log.info("Subscribed to topic(s): {}", Utils.join(topics, ", "));
+                if (this.subscriptions.subscribe(new HashSet<>(topics), listener))
+                    metadata.requestUpdateForNewTopics();
             }
         } finally {
             release();
@@ -917,7 +981,7 @@ public void subscribe(Collection topics, ConsumerRebalanceListener liste
      *
      * 
    
      * This is a short-hand for {@link #subscribe(Collection, ConsumerRebalanceListener)}, which
-     * uses a noop listener. If you need the ability to seek to particular offsets, you should prefer
+     * uses a no-op listener. If you need the ability to seek to particular offsets, you should prefer
      * {@link #subscribe(Collection, ConsumerRebalanceListener)}, since group rebalances will cause partition offsets
      * to be reset. You should also provide your own listener if you are doing your own offset
      * management since the listener gives you an opportunity to commit offsets before a rebalance finishes.
@@ -935,17 +999,14 @@ public void subscribe(Collection topics) {
 
     /**
      * Subscribe to all topics matching specified pattern to get dynamically assigned partitions.
-     * The pattern matching will be done periodically against topic existing at the time of check.
+     * The pattern matching will be done periodically against all topics existing at the time of check.
+     * This can be controlled through the {@code metadata.max.age.ms} configuration: by lowering
+     * the max metadata age, the consumer will refresh metadata more often and check for matching topics.
      * 
    
-     * As part of group management, the consumer will keep track of the list of consumers that
-     * belong to a particular group and will trigger a rebalance operation if one of the
-     * following events trigger -
-     * 
      
-     * 
    • Number of partitions change for any of the subscribed list of topics
-     * 
    • Topic is created or deleted
-     * 
    • An existing member of the consumer group dies
-     * 
    • A new member is added to an existing consumer group via the join API
-     * 
    
+     * See {@link #subscribe(Collection, ConsumerRebalanceListener)} for details on the
+     * use of the {@link ConsumerRebalanceListener}. Generally rebalances are triggered when there
+     * is a change to the topics matching the provided pattern and when consumer group membership changes.
+     * Group rebalances only take place during an active call to {@link #poll(Duration)}.
      *
      * @param pattern Pattern to subscribe to
      * @param listener Non-null listener instance to get notifications on partition assignment/revocation for the
@@ -957,18 +1018,18 @@ public void subscribe(Collection topics) {
      */
     @Override
     public void subscribe(Pattern pattern, ConsumerRebalanceListener listener) {
+        maybeThrowInvalidGroupIdException();
+        if (pattern == null || pattern.toString().equals(""))
+            throw new IllegalArgumentException("Topic pattern to subscribe to cannot be " + (pattern == null ?
+                    "null" : "empty"));
+
         acquireAndEnsureOpen();
         try {
-            if (pattern == null)
-                throw new IllegalArgumentException("Topic pattern to subscribe to cannot be null");
-
             throwIfNoAssignorsConfigured();
-
-            log.debug("Subscribed to pattern: {}", pattern);
+            log.info("Subscribed to pattern: '{}'", pattern);
             this.subscriptions.subscribe(pattern, listener);
-            this.metadata.needMetadataForAllTopics(true);
             this.coordinator.updatePatternSubscription(metadata.fetch());
-            this.metadata.requestUpdate();
+            this.metadata.requestUpdateForNewTopics();
         } finally {
             release();
         }
@@ -979,7 +1040,7 @@ public void subscribe(Pattern pattern, ConsumerRebalanceListener listener) {
      * The pattern matching will be done periodically against topics existing at the time of check.
      * 
    
      * This is a short-hand for {@link #subscribe(Pattern, ConsumerRebalanceListener)}, which
-     * uses a noop listener. If you need the ability to seek to particular offsets, you should prefer
+     * uses a no-op listener. If you need the ability to seek to particular offsets, you should prefer
      * {@link #subscribe(Pattern, ConsumerRebalanceListener)}, since group rebalances will cause partition offsets
      * to be reset. You should also provide your own listener if you are doing your own offset
      * management since the listener gives you an opportunity to commit offsets before a rebalance finishes.
@@ -998,14 +1059,19 @@ 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 {
-            log.debug("Unsubscribed all topics or patterns and assigned partitions");
+            fetcher.clearBufferedDataForUnassignedPartitions(Collections.emptySet());
+            if (this.coordinator != null) {
+                this.coordinator.onLeavePrepare();
+                this.coordinator.maybeLeaveGroup("the consumer unsubscribed from all topics");
+            }
             this.subscriptions.unsubscribe();
-            this.coordinator.maybeLeaveGroup();
-            this.metadata.needMetadataForAllTopics(false);
+            log.info("Unsubscribed all topics or patterns and assigned partitions");
         } finally {
             release();
         }
@@ -1039,21 +1105,21 @@ public void assign(Collection partitions) {
             } else if (partitions.isEmpty()) {
                 this.unsubscribe();
             } else {
-                Set topics = new HashSet<>();
                 for (TopicPartition tp : partitions) {
                     String topic = (tp != null) ? tp.topic() : null;
                     if (topic == null || topic.trim().isEmpty())
                         throw new IllegalArgumentException("Topic partitions to assign to cannot have null or empty topic");
-                    topics.add(topic);
                 }
+                fetcher.clearBufferedDataForUnassignedPartitions(partitions);
 
                 // make sure the offsets of topic partitions the consumer is unsubscribing from
                 // are committed since there will be no following rebalance
-                this.coordinator.maybeAutoCommitOffsetsNow();
+                if (coordinator != null)
+                    this.coordinator.maybeAutoCommitOffsetsAsync(time.milliseconds());
 
-                log.debug("Subscribed to partition(s): {}", Utils.join(partitions, ", "));
-                this.subscriptions.assignFromUser(new HashSet<>(partitions));
-                metadata.setTopics(topics);
+                log.info("Subscribed to partition(s): {}", Utils.join(partitions, ", "));
+                if (this.subscriptions.assignFromUser(new HashSet<>(partitions)))
+                    metadata.requestUpdateForNewTopics();
             }
         } finally {
             release();
@@ -1069,7 +1135,7 @@ public void assign(Collection partitions) {
      * offset for the subscribed list of partitions
      *
      *
-     * @param timeout The time, in milliseconds, spent waiting in poll if data is not available in the buffer.
+     * @param timeoutMs The time, in milliseconds, spent waiting in poll if data is not available in the buffer.
      *            If 0, returns immediately with any records that are available currently in the buffer, else returns empty.
      *            Must not be negative.
      * @return map of topic to records since the last fetch for the subscribed list of topics and partitions
@@ -1088,22 +1154,87 @@ 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.
      */
+    @Deprecated
     @Override
-    public ConsumerRecords poll(long timeout) {
+    public ConsumerRecords poll(final long timeoutMs) {
+        return poll(time.timer(timeoutMs), false);
+    }
+
+    /**
+     * Fetch data for the topics or partitions specified using one of the subscribe/assign APIs. It is an error to not have
+     * subscribed to any topics or partitions before polling for data.
+     * 
    
+     * On each poll, consumer will try to use the last consumed offset as the starting offset and fetch sequentially. The last
+     * consumed offset can be manually set through {@link #seek(TopicPartition, long)} or automatically set as the last committed
+     * offset for the subscribed list of partitions
+     *
+     * 
    
+     * This method returns immediately if there are records available. Otherwise, it will await the passed timeout.
+     * If the timeout expires, an empty record set will be returned. Note that this method may block beyond the
+     * timeout in order to execute custom {@link ConsumerRebalanceListener} callbacks.
+     *
+     *
+     * @param timeout The maximum time to block (must not be greater than {@link Long#MAX_VALUE} milliseconds)
+     *
+     * @return map of topic to records since the last fetch for the subscribed list of topics and partitions
+     *
+     * @throws org.apache.kafka.clients.consumer.InvalidOffsetException if the offset for a partition or set of
+     *             partitions is undefined or out of range and no offset reset policy has been configured
+     * @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 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, 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.UnsupportedVersionException if the consumer attempts to fetch stable offsets
+     *             when the broker doesn't support this feature
+     * @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 {
-            if (timeout < 0)
-                throw new IllegalArgumentException("Timeout must not be negative");
+            this.kafkaConsumerMetrics.recordPollStart(timer.currentTimeMs());
 
-            if (this.subscriptions.hasNoSubscriptionOrUserAssignment())
+            if (this.subscriptions.hasNoSubscriptionOrUserAssignment()) {
                 throw new IllegalStateException("Consumer is not subscribed to any topics or assigned any partitions");
+            }
 
-            // poll for new data until the timeout expires
-            long start = time.milliseconds();
-            long remaining = timeout;
             do {
-                Map>> records = pollOnce(remaining);
+                client.maybeTriggerWakeup();
+
+                if (includeMetadataInTimeout) {
+                    // try to update assignment metadata BUT do not need to block on the timer for join group
+                    updateAssignmentMetadataIfNeeded(timer, false);
+                } else {
+                    while (!updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE), true)) {
+                        log.warn("Still waiting for metadata");
+                    }
+                }
+
+                final Map>> records = pollForFetches(timer);
                 if (!records.isEmpty()) {
                     // before returning the fetched records, we can send off the next round of fetches
                     // and avoid block waiting for their responses to enable pipelining while the user
@@ -1111,84 +1242,93 @@ public ConsumerRecords poll(long timeout) {
                     //
                     // 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();
+                    if (fetcher.sendFetches() > 0 || client.hasPendingRequests()) {
+                        client.transmitSends();
+                    }
 
-                    if (this.interceptors == null)
-                        return new ConsumerRecords<>(records);
-                    else
-                        return this.interceptors.onConsume(new ConsumerRecords<>(records));
+                    return this.interceptors.onConsume(new ConsumerRecords<>(records));
                 }
-
-                long elapsed = time.milliseconds() - start;
-                remaining = timeout - elapsed;
-            } while (remaining > 0);
+            } while (timer.notExpired());
 
             return ConsumerRecords.empty();
         } finally {
             release();
+            this.kafkaConsumerMetrics.recordPollEnd(timer.currentTimeMs());
         }
     }
 
+    boolean updateAssignmentMetadataIfNeeded(final Timer timer, final boolean waitForJoinGroup) {
+        if (coordinator != null && !coordinator.poll(timer, waitForJoinGroup)) {
+            return false;
+        }
+
+        return updateFetchPositions(timer);
+    }
+
     /**
-     * Do one round of polling. In addition to checking for new data, this does any needed offset commits
-     * (if auto-commit is enabled), and offset resets (if an offset reset policy is defined).
-     * @param timeout The maximum time to block in the underlying call to {@link ConsumerNetworkClient#poll(long)}.
-     * @return The fetched records (may be empty)
+     * @throws KafkaException if the rebalance callback throws exception
      */
-    private Map>> pollOnce(long timeout) {
-        client.maybeTriggerWakeup();
-        coordinator.poll(time.milliseconds(), timeout);
-
-        // fetch positions if we have partitions we're subscribed to that we
-        // don't know the offset for
-        if (!subscriptions.hasAllFetchPositions())
-            updateFetchPositions(this.subscriptions.missingFetchPositions());
+    private Map>> pollForFetches(Timer timer) {
+        long pollTimeout = coordinator == null ? timer.remainingMs() :
+                Math.min(coordinator.timeToNextPoll(timer.currentTimeMs()), timer.remainingMs());
 
         // if data is available already, return it immediately
-        Map>> records = fetcher.fetchedRecords();
-        if (!records.isEmpty())
+        final Map>> records = fetcher.fetchedRecords();
+        if (!records.isEmpty()) {
             return records;
+        }
 
         // send any new fetches (won't resend pending fetches)
         fetcher.sendFetches();
 
-        long now = time.milliseconds();
-        long pollTimeout = Math.min(coordinator.timeToNextPoll(now), timeout);
+        // We do not want to be stuck blocking in poll if we are missing some positions
+        // since the offset lookup may be backing off after a failure
 
-        client.poll(pollTimeout, now, new PollCondition() {
-            @Override
-            public boolean shouldBlock() {
-                // 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();
-            }
-        });
+        // NOTE: the use of cachedSubscriptionHashAllFetchPositions means we MUST call
+        // updateAssignmentMetadataIfNeeded before this method.
+        if (!cachedSubscriptionHashAllFetchPositions && pollTimeout > retryBackoffMs) {
+            pollTimeout = retryBackoffMs;
+        }
 
-        // after the long poll, we should check whether the group needs to rebalance
-        // prior to returning data so that the group can stabilize faster
-        if (coordinator.needRejoin())
-            return Collections.emptyMap();
+        log.trace("Polling for fetches with timeout {}", pollTimeout);
+
+        Timer pollTimer = time.timer(pollTimeout);
+        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.hasAvailableFetches();
+        });
+        timer.update(pollTimer.currentTimeMs());
 
         return fetcher.fetchedRecords();
     }
 
     /**
-     * Commit offsets returned on the last {@link #poll(long) poll()} for all the subscribed list of topics and partitions.
+     * Commit offsets returned on the last {@link #poll(Duration) poll()} for all the subscribed list of topics and
+     * partitions.
      * 
    
      * This commits offsets only to Kafka. The offsets committed using this API will be used on the first fetch after
      * every rebalance and also on startup. As such, if you need to store offsets in anything other than Kafka, this API
      * should not be used.
      * 
    
-     * This is a synchronous commits and will block until either the commit succeeds or an unrecoverable error is
-     * encountered (in which case it is thrown to the caller).
+     * This is a synchronous commit and will block until either the commit succeeds, 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).
      * 
    
      * Note that asynchronous offset commits sent previously with the {@link #commitAsync(OffsetCommitCallback)}
      * (or similar) are guaranteed to have their callbacks invoked prior to completion of this method.
      *
      * @throws org.apache.kafka.clients.consumer.CommitFailedException if the commit failed and cannot be retried.
-     *             This can only occur if you are using automatic group management with {@link #subscribe(Collection)},
-     *             or if there is an active group with the same groupId which is using group management.
+     *             This fatal error can only occur if you are using automatic group management with {@link #subscribe(Collection)},
+     *             or if there is an active group with the same group.id which is using group management. In such cases,
+     *             when you are trying to commit to partitions that are no longer assigned to this consumer because the
+     *             consumer is for example no longer part of the group this exception would be thrown.
+     * @throws org.apache.kafka.common.errors.RebalanceInProgressException if the consumer instance is in the middle of a rebalance
+     *            so it is not yet determined which partitions would be assigned to the consumer. In such cases you can first
+     *            complete the rebalance by calling {@link #poll(Duration)} and commit can be reconsidered afterwards.
+     *            NOTE when you reconsider committing after the rebalance, the assigned partitions may have changed,
+     *            and also for those partitions that are still assigned their fetch positions may have changed too
+     *            if more records are returned from the {@link #poll(Duration)} call.
      * @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
@@ -1198,15 +1338,56 @@ public boolean shouldBlock() {
      *             configured groupId. See the exception for more details
      * @throws org.apache.kafka.common.KafkaException for any other unrecoverable errors (e.g. if offset metadata
      *             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() {
-        acquireAndEnsureOpen();
-        try {
-            coordinator.commitOffsetsSync(subscriptions.allConsumed(), Long.MAX_VALUE);
-        } finally {
-            release();
-        }
+        commitSync(Duration.ofMillis(defaultApiTimeoutMs));
+    }
+
+    /**
+     * Commit offsets returned on the last {@link #poll(Duration) poll()} for all the subscribed list of topics and
+     * partitions.
+     * 
    
+     * This commits offsets only to Kafka. The offsets committed using this API will be used on the first fetch after
+     * every rebalance and also on startup. As such, if you need to store offsets in anything other than Kafka, this API
+     * should not be used.
+     * 
    
+     * This is a synchronous commit and will block until either the commit succeeds, an unrecoverable error is
+     * encountered (in which case it is thrown to the caller), or the passed timeout expires.
+     * 
    
+     * Note that asynchronous offset commits sent previously with the {@link #commitAsync(OffsetCommitCallback)}
+     * (or similar) are guaranteed to have their callbacks invoked prior to completion of this method.
+     *
+     * @throws org.apache.kafka.clients.consumer.CommitFailedException if the commit failed and cannot be retried.
+     *             This can only occur if you are using automatic group management with {@link #subscribe(Collection)},
+     *             or if there is an active group with the same group.id which is using group management. In such cases,
+     *             when you are trying to commit to partitions that are no longer assigned to this consumer because the
+     *             consumer is for example no longer part of the group this exception would be thrown.
+     * @throws org.apache.kafka.common.errors.RebalanceInProgressException if the consumer instance is in the middle of a rebalance
+     *            so it is not yet determined which partitions would be assigned to the consumer. In such cases you can first
+     *            complete the rebalance by calling {@link #poll(Duration)} and commit can be reconsidered afterwards.
+     *            NOTE when you reconsider committing after the rebalance, the assigned partitions may have changed,
+     *            and also for those partitions that are still assigned their fetch positions may have changed too
+     *            if more records are returned from the {@link #poll(Duration)} call.
+     * @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 (e.g. if offset metadata
+     *             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) {
+        commitSync(subscriptions.allConsumed(), timeout);
     }
 
     /**
@@ -1215,10 +1396,12 @@ public void commitSync() {
      * This commits offsets to Kafka. The offsets committed using this API will be used on the first fetch after every
      * rebalance and also on startup. As such, if you need to store offsets in anything other than Kafka, this API
      * should not be used. The committed offset should be the next message your application will consume,
-     * i.e. lastProcessedMessageOffset + 1.
+     * i.e. lastProcessedMessageOffset + 1. If automatic group management with {@link #subscribe(Collection)} is used,
+     * then the committed offsets must belong to the currently auto-assigned partitions.
      * 
    
-     * This is a synchronous commits and will block until either the commit succeeds or an unrecoverable error is
-     * encountered (in which case it is thrown to the caller).
+     * This is a synchronous commit and will block until either the commit succeeds or 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).
      * 
    
      * Note that asynchronous offset commits sent previously with the {@link #commitAsync(OffsetCommitCallback)}
      * (or similar) are guaranteed to have their callbacks invoked prior to completion of this method.
@@ -1226,7 +1409,16 @@ public void commitSync() {
      * @param offsets A map of offsets by partition with associated metadata
      * @throws org.apache.kafka.clients.consumer.CommitFailedException if the commit failed and cannot be retried.
      *             This can only occur if you are using automatic group management with {@link #subscribe(Collection)},
-     *             or if there is an active group with the same groupId which is using group management.
+     *             or if there is an active group with the same group.id which is using group management. In such cases,
+     *             when you are trying to commit to partitions that are no longer assigned to this consumer because the
+     *             consumer is for example no longer part of the group this exception would be thrown.
+     * @throws org.apache.kafka.common.errors.RebalanceInProgressException if the consumer instance is in the middle of a rebalance
+     *            so it is not yet determined which partitions would be assigned to the consumer. In such cases you can first
+     *            complete the rebalance by calling {@link #poll(Duration)} and commit can be reconsidered afterwards.
+     *            NOTE when you reconsider committing after the rebalance, the assigned partitions may have changed,
+     *            and also for those partitions that are still assigned their fetch positions may have changed too
+     *            if more records are returned from the {@link #poll(Duration)} call, so when you retry committing
+     *            you should consider updating the passed in {@code offset} parameter.
      * @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
@@ -1237,20 +1429,77 @@ public void commitSync() {
      * @throws java.lang.IllegalArgumentException if the committed offset is negative
      * @throws org.apache.kafka.common.KafkaException for any other unrecoverable errors (e.g. if offset metadata
      *             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) {
+        commitSync(offsets, Duration.ofMillis(defaultApiTimeoutMs));
+    }
+
+    /**
+     * Commit the specified offsets for the specified list of topics and partitions.
+     * 
    
+     * This commits offsets to Kafka. The offsets committed using this API will be used on the first fetch after every
+     * rebalance and also on startup. As such, if you need to store offsets in anything other than Kafka, this API
+     * should not be used. The committed offset should be the next message your application will consume,
+     * i.e. lastProcessedMessageOffset + 1. If automatic group management with {@link #subscribe(Collection)} is used,
+     * then the committed offsets must belong to the currently auto-assigned partitions.
+     * 
    
+     * This is a synchronous commit and will block until either the commit succeeds, an unrecoverable error is
+     * encountered (in which case it is thrown to the caller), or the timeout expires.
+     * 
    
+     * Note that asynchronous offset commits sent previously with the {@link #commitAsync(OffsetCommitCallback)}
+     * (or similar) are guaranteed to have their callbacks invoked prior to completion of this method.
+     *
+     * @param offsets A map of offsets by partition with associated metadata
+     * @param timeout The maximum amount of time to await completion of the offset commit
+     * @throws org.apache.kafka.clients.consumer.CommitFailedException if the commit failed and cannot be retried.
+     *             This can only occur if you are using automatic group management with {@link #subscribe(Collection)},
+     *             or if there is an active group with the same group.id which is using group management. In such cases,
+     *             when you are trying to commit to partitions that are no longer assigned to this consumer because the
+     *             consumer is for example no longer part of the group this exception would be thrown.
+     * @throws org.apache.kafka.common.errors.RebalanceInProgressException if the consumer instance is in the middle of a rebalance
+     *            so it is not yet determined which partitions would be assigned to the consumer. In such cases you can first
+     *            complete the rebalance by calling {@link #poll(Duration)} and commit can be reconsidered afterwards.
+     *            NOTE when you reconsider committing after the rebalance, the assigned partitions may have changed,
+     *            and also for those partitions that are still assigned their fetch positions may have changed too
+     *            if more records are returned from the {@link #poll(Duration)} call, so when you retry committing
+     *            you should consider updating the passed in {@code offset} parameter.
+     * @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 java.lang.IllegalArgumentException if the committed offset is negative
+     * @throws org.apache.kafka.common.KafkaException for any other unrecoverable errors (e.g. if offset metadata
+     *             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) {
         acquireAndEnsureOpen();
         try {
-            coordinator.commitOffsetsSync(new HashMap<>(offsets), Long.MAX_VALUE);
+            maybeThrowInvalidGroupIdException();
+            offsets.forEach(this::updateLastSeenEpochIfNewer);
+            if (!coordinator.commitOffsetsSync(new HashMap<>(offsets), time.timer(timeout))) {
+                throw new TimeoutException("Timeout of " + timeout.toMillis() + "ms expired before successfully " +
+                        "committing offsets " + offsets);
+            }
         } finally {
             release();
         }
     }
 
     /**
-     * Commit offsets returned on the last {@link #poll(long) poll()} for all the subscribed list of topics and partition.
+     * 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() {
@@ -1258,7 +1507,7 @@ public void commitAsync() {
     }
 
     /**
-     * Commit offsets returned on the last {@link #poll(long) poll()} for the subscribed list of topics and partitions.
+     * Commit offsets returned on the last {@link #poll(Duration) poll()} for the subscribed list of topics and partitions.
      * 
    
      * This commits offsets only to Kafka. The offsets committed using this API will be used on the first fetch after
      * every rebalance and also on startup. As such, if you need to store offsets in anything other than Kafka, this API
@@ -1273,15 +1522,11 @@ 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) {
-        acquireAndEnsureOpen();
-        try {
-            commitAsync(subscriptions.allConsumed(), callback);
-        } finally {
-            release();
-        }
+        commitAsync(subscriptions.allConsumed(), callback);
     }
 
     /**
@@ -1290,7 +1535,8 @@ public void commitAsync(OffsetCommitCallback callback) {
      * This commits offsets to Kafka. The offsets committed using this API will be used on the first fetch after every
      * rebalance and also on startup. As such, if you need to store offsets in anything other than Kafka, this API
      * should not be used. The committed offset should be the next message your application will consume,
-     * i.e. lastProcessedMessageOffset + 1.
+     * i.e. lastProcessedMessageOffset + 1. If automatic group management with {@link #subscribe(Collection)} is used,
+     * then the committed offsets must belong to the currently auto-assigned partitions.
      * 
    
      * This is an asynchronous call and will not block. Any errors encountered are either passed to the callback
      * (if provided) or discarded.
@@ -1303,12 +1549,15 @@ 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) {
         acquireAndEnsureOpen();
         try {
+            maybeThrowInvalidGroupIdException();
             log.debug("Committing offsets: {}", offsets);
+            offsets.forEach(this::updateLastSeenEpochIfNewer);
             coordinator.commitOffsetsAsync(new HashMap<>(offsets), callback);
         } finally {
             release();
@@ -1316,22 +1565,62 @@ public void commitAsync(final Map offsets, Of
     }
 
     /**
-     * Overrides the fetch offsets that the consumer will use on the next {@link #poll(long) poll(timeout)}. If this API
+     * Overrides the fetch offsets that the consumer will use on the next {@link #poll(Duration) poll(timeout)}. If this API
      * is invoked for the same partition more than once, the latest offset will be used on the next poll(). Note that
      * you may lose data if this API is arbitrarily used in the middle of consumption, to reset the fetch offsets
      *
-     * @throws IllegalArgumentException if the provided TopicPartition is not assigned to this consumer
-     *                                  or if provided offset is negative
+     * @throws IllegalArgumentException if the provided offset is negative
+     * @throws IllegalStateException if the provided TopicPartition is not assigned to this consumer
      */
     @Override
     public void seek(TopicPartition partition, long offset) {
+        if (offset < 0)
+            throw new IllegalArgumentException("seek offset must not be a negative number");
+
         acquireAndEnsureOpen();
         try {
-            if (offset < 0)
-                throw new IllegalArgumentException("seek offset must not be a negative number");
+            log.info("Seeking to offset {} for partition {}", offset, partition);
+            SubscriptionState.FetchPosition newPosition = new SubscriptionState.FetchPosition(
+                    offset,
+                    Optional.empty(), // This will ensure we skip validation
+                    this.metadata.currentLeader(partition));
+            this.subscriptions.seekUnvalidated(partition, newPosition);
+        } finally {
+            release();
+        }
+    }
 
-            log.debug("Seeking to offset {} for partition {}", offset, partition);
-            this.subscriptions.seek(partition, offset);
+    /**
+     * Overrides the fetch offsets that the consumer will use on the next {@link #poll(Duration) poll(timeout)}. If this API
+     * is invoked for the same partition more than once, the latest offset will be used on the next poll(). Note that
+     * you may lose data if this API is arbitrarily used in the middle of consumption, to reset the fetch offsets. This
+     * method allows for setting the leaderEpoch along with the desired offset.
+     *
+     * @throws IllegalArgumentException if the provided offset is negative
+     * @throws IllegalStateException if the provided TopicPartition is not assigned to this consumer
+     */
+    @Override
+    public void seek(TopicPartition partition, OffsetAndMetadata offsetAndMetadata) {
+        long offset = offsetAndMetadata.offset();
+        if (offset < 0) {
+            throw new IllegalArgumentException("seek offset must not be a negative number");
+        }
+
+        acquireAndEnsureOpen();
+        try {
+            if (offsetAndMetadata.leaderEpoch().isPresent()) {
+                log.info("Seeking to offset {} for partition {} with epoch {}",
+                        offset, partition, offsetAndMetadata.leaderEpoch().get());
+            } else {
+                log.info("Seeking to offset {} for partition {}", offset, partition);
+            }
+            Metadata.LeaderAndEpoch currentLeaderAndEpoch = this.metadata.currentLeader(partition);
+            SubscriptionState.FetchPosition newPosition = new SubscriptionState.FetchPosition(
+                    offsetAndMetadata.offset(),
+                    offsetAndMetadata.leaderEpoch(),
+                    currentLeaderAndEpoch);
+            this.updateLastSeenEpochIfNewer(partition, offsetAndMetadata);
+            this.subscriptions.seekUnvalidated(partition, newPosition);
         } finally {
             release();
         }
@@ -1339,22 +1628,21 @@ public void seek(TopicPartition partition, long offset) {
 
     /**
      * Seek to the first offset for each of the given partitions. This function evaluates lazily, seeking to the
-     * first offset in all partitions only when {@link #poll(long)} or {@link #position(TopicPartition)} are called.
+     * first offset in all partitions only when {@link #poll(Duration)} or {@link #position(TopicPartition)} are called.
      * If no partitions are provided, seek to the first offset for all of the currently assigned partitions.
      *
-     * @throws IllegalArgumentException if {@code partitions} is {@code null} or the provided TopicPartition is not assigned to this consumer
+     * @throws IllegalArgumentException if {@code partitions} is {@code null}
+     * @throws IllegalStateException if any of the provided partitions are not currently assigned to this consumer
      */
+    @Override
     public void seekToBeginning(Collection partitions) {
+        if (partitions == null)
+            throw new IllegalArgumentException("Partitions collection cannot be null");
+
         acquireAndEnsureOpen();
         try {
-            if (partitions == null) {
-                throw new IllegalArgumentException("Partitions collection cannot be null");
-            }
             Collection parts = partitions.size() == 0 ? this.subscriptions.assignedPartitions() : partitions;
-            for (TopicPartition tp : parts) {
-                log.debug("Seeking to beginning of partition {}", tp);
-                subscriptions.needOffsetReset(tp, OffsetResetStrategy.EARLIEST);
-            }
+            subscriptions.requestOffsetReset(parts, OffsetResetStrategy.EARLIEST);
         } finally {
             release();
         }
@@ -1362,25 +1650,24 @@ public void seekToBeginning(Collection partitions) {
 
     /**
      * Seek to the last offset for each of the given partitions. This function evaluates lazily, seeking to the
-     * final offset in all partitions only when {@link #poll(long)} or {@link #position(TopicPartition)} are called.
+     * final offset in all partitions only when {@link #poll(Duration)} or {@link #position(TopicPartition)} are called.
      * If no partitions are provided, seek to the final offset for all of the currently assigned partitions.
      * 
    
      * If {@code isolation.level=read_committed}, the end offset will be the Last Stable Offset, i.e., the offset
      * of the first message with an open transaction.
      *
-     * @throws IllegalArgumentException if {@code partitions} is {@code null} or the provided TopicPartition is not assigned to this consumer
+     * @throws IllegalArgumentException if {@code partitions} is {@code null}
+     * @throws IllegalStateException if any of the provided partitions are not currently assigned to this consumer
      */
+    @Override
     public void seekToEnd(Collection partitions) {
+        if (partitions == null)
+            throw new IllegalArgumentException("Partitions collection cannot be null");
+
         acquireAndEnsureOpen();
         try {
-            if (partitions == null) {
-                throw new IllegalArgumentException("Partitions collection cannot be null");
-            }
             Collection parts = partitions.size() == 0 ? this.subscriptions.assignedPartitions() : partitions;
-            for (TopicPartition tp : parts) {
-                log.debug("Seeking to end of partition {}", tp);
-                subscriptions.needOffsetReset(tp, OffsetResetStrategy.LATEST);
-            }
+            subscriptions.requestOffsetReset(parts, OffsetResetStrategy.LATEST);
         } finally {
             release();
         }
@@ -1391,11 +1678,12 @@ public void seekToEnd(Collection partitions) {
      * This method may issue a remote call to the server if there is no current position for the given partition.
      * 
    
      * This call will block until either the position could be determined or an unrecoverable error is
-     * encountered (in which case it is thrown to the caller).
+     * 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 partition The partition to get the position for
-     * @return The offset
-     * @throws IllegalArgumentException if the provided TopicPartition is not assigned to this consumer
+     * @return The current position of the consumer (that is, the offset of the next record to be fetched)
+     * @throws IllegalStateException if the provided TopicPartition is not assigned to this consumer
      * @throws org.apache.kafka.clients.consumer.InvalidOffsetException if no offset is currently defined for
      *             the partition
      * @throws org.apache.kafka.common.errors.WakeupException if {@link #wakeup()} is called before or while this
@@ -1405,20 +1693,61 @@ public void seekToEnd(Collection partitions) {
      * @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.errors.UnsupportedVersionException if the consumer attempts to fetch stable offsets
+     *             when the broker doesn't support this feature
      * @throws org.apache.kafka.common.KafkaException for any other unrecoverable errors
+     * @throws org.apache.kafka.common.errors.TimeoutException if the position cannot be determined before the
+     *             timeout specified by {@code default.api.timeout.ms} expires
      */
+    @Override
     public long position(TopicPartition partition) {
+        return position(partition, Duration.ofMillis(defaultApiTimeoutMs));
+    }
+
+    /**
+     * Get the offset of the next record that will be fetched (if a record with that offset exists).
+     * This method may issue a remote call to the server if there is no current position
+     * for the given partition.
+     * 
    
+     * 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 get the position for
+     * @param timeout The maximum amount of time to await determination of the current position
+     * @return The current position of the consumer (that is, the offset of the next record to be fetched)
+     * @throws IllegalStateException if the provided TopicPartition is not assigned to this consumer
+     * @throws org.apache.kafka.clients.consumer.InvalidOffsetException if no offset is currently defined for
+     *             the partition
+     * @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.TimeoutException if the position cannot be determined before the
+     *             passed timeout expires
+     * @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
+     */
+    @Override
+    public long position(TopicPartition partition, final Duration timeout) {
         acquireAndEnsureOpen();
         try {
             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) {
-                // batch update fetch positions for any partitions without a valid position
-                updateFetchPositions(subscriptions.assignedPartitions());
-                offset = this.subscriptions.position(partition);
-            }
-            return offset;
+                throw new IllegalStateException("You can only check the position for partitions assigned to this consumer.");
+
+            Timer timer = time.timer(timeout);
+            do {
+                SubscriptionState.FetchPosition position = this.subscriptions.validPosition(partition);
+                if (position != null)
+                    return position.offset;
+
+                updateFetchPositions(timer);
+                client.poll(timer);
+            } while (timer.notExpired());
+
+            throw new TimeoutException("Timeout of " + timeout.toMillis() + "ms expired before the position " +
+                    "for partition " + partition + " could be determined");
         } finally {
             release();
         }
@@ -1428,7 +1757,10 @@ public long position(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 do a remote call to get the latest committed offset from the server, and will block until the
+     * committed offset is 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 partition The partition to check
      * @return The last committed offset and metadata or null if there was no prior commit
@@ -1440,13 +1772,115 @@ public long position(TopicPartition partition) {
      * @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.
+     *
+     * @deprecated since 2.4 Use {@link #committed(Set)} instead
      */
+    @Deprecated
     @Override
     public OffsetAndMetadata committed(TopicPartition partition) {
+        return committed(partition, Duration.ofMillis(defaultApiTimeoutMs));
+    }
+
+    /**
+     * 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 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
+     * @return The last committed offset and metadata or null if there was no prior commit
+     * @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
+     *
+     * @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.errors.UnsupportedVersionException if the consumer attempts to fetch stable offsets
+     *             when the broker doesn't support this feature
+     * @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 {
-            Map offsets = coordinator.fetchCommittedOffsets(Collections.singleton(partition));
-            return offsets.get(partition);
+            maybeThrowInvalidGroupIdException();
+            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 partitions " + partitions + " could be determined. Try tuning default.api.timeout.ms " +
+                    "larger to relax the threshold.");
+            } else {
+                offsets.forEach(this::updateLastSeenEpochIfNewer);
+                return offsets;
+            }
         } finally {
             release();
         }
@@ -1465,6 +1899,7 @@ public OffsetAndMetadata committed(TopicPartition partition) {
      * does not already have any metadata about the given topic.
      *
      * @param topic The topic to get partition metadata for
+     *
      * @return The list of partitions
      * @throws org.apache.kafka.common.errors.WakeupException if {@link #wakeup()} is called before or while this
      *             function is called
@@ -1472,12 +1907,36 @@ public OffsetAndMetadata committed(TopicPartition partition) {
      *             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 specified topic. See the exception for more details
-     * @throws org.apache.kafka.common.errors.TimeoutException if the topic metadata could not be fetched before
-     *             expiration of the configured request timeout
      * @throws org.apache.kafka.common.KafkaException for any other unrecoverable errors
+     * @throws org.apache.kafka.common.errors.TimeoutException if the offset metadata could not be fetched before
+     *         the amount of time allocated by {@code default.api.timeout.ms} expires.
      */
     @Override
     public List partitionsFor(String topic) {
+        return partitionsFor(topic, Duration.ofMillis(defaultApiTimeoutMs));
+    }
+
+    /**
+     * Get metadata about the partitions for a given topic. This method will issue a remote call to the server if it
+     * does not already have any metadata about the given topic.
+     *
+     * @param topic The topic to get partition metadata for
+     * @param timeout The maximum of time to await topic metadata
+     *
+     * @return The list of partitions
+     * @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 specified topic. See
+     *             the exception for more details
+     * @throws org.apache.kafka.common.errors.TimeoutException if topic metadata cannot be fetched before expiration
+     *             of the passed timeout
+     * @throws org.apache.kafka.common.KafkaException for any other unrecoverable errors
+     */
+    @Override
+    public List partitionsFor(String topic, Duration timeout) {
         acquireAndEnsureOpen();
         try {
             Cluster cluster = this.metadata.fetch();
@@ -1485,8 +1944,9 @@ public List partitionsFor(String topic) {
             if (!parts.isEmpty())
                 return parts;
 
+            Timer timer = time.timer(timeout);
             Map> topicMetadata = fetcher.getTopicMetadata(
-                    new MetadataRequest.Builder(Collections.singletonList(topic), true), requestTimeoutMs);
+                    new MetadataRequest.Builder(Collections.singletonList(topic), metadata.allowAutoTopicCreation()), timer);
             return topicMetadata.get(topic);
         } finally {
             release();
@@ -1498,31 +1958,52 @@ public List partitionsFor(String topic) {
      * remote call to the server.
 
      * @return The map of topics and its partitions
+     *
      * @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.TimeoutException if the topic metadata could not be fetched before
-     *             expiration of the configured request timeout
      * @throws org.apache.kafka.common.KafkaException for any other unrecoverable errors
+     * @throws org.apache.kafka.common.errors.TimeoutException if the offset metadata could not be fetched before
+     *         the amount of time allocated by {@code default.api.timeout.ms} expires.
      */
     @Override
     public Map> listTopics() {
+        return listTopics(Duration.ofMillis(defaultApiTimeoutMs));
+    }
+
+    /**
+     * Get metadata about partitions for all topics that the user is authorized to view. This method will issue a
+     * remote call to the server.
+     *
+     * @param timeout The maximum time this operation will block to fetch topic metadata
+     *
+     * @return The map of topics and its partitions
+     * @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.TimeoutException if the topic metadata could not be fetched before
+     *             expiration of the passed timeout
+     * @throws org.apache.kafka.common.KafkaException for any other unrecoverable errors
+     */
+    @Override
+    public Map> listTopics(Duration timeout) {
         acquireAndEnsureOpen();
         try {
-            return fetcher.getAllTopicMetadata(requestTimeoutMs);
+            return fetcher.getAllTopicMetadata(time.timer(timeout));
         } finally {
             release();
         }
     }
 
     /**
-     * Suspend fetching from the requested partitions. Future calls to {@link #poll(long)} will not return
+     * Suspend fetching from the requested partitions. Future calls to {@link #poll(Duration)} will not return
      * any records from these partitions until they have been resumed using {@link #resume(Collection)}.
      * Note that this method does not affect partition subscription. In particular, it does not cause a group
      * rebalance when automatic assignment is used.
      * @param partitions The partitions which should be paused
-     * @throws IllegalStateException if one of the provided partitions is not assigned to this consumer
+     * @throws IllegalStateException if any of the provided partitions are not currently assigned to this consumer
      */
     @Override
     public void pause(Collection partitions) {
@@ -1539,10 +2020,10 @@ public void pause(Collection partitions) {
 
     /**
      * Resume specified partitions which have been paused with {@link #pause(Collection)}. New calls to
-     * {@link #poll(long)} will return records from these partitions if there are any to be fetched.
+     * {@link #poll(Duration)} will return records from these partitions if there are any to be fetched.
      * If the partitions were not previously paused, this method is a no-op.
      * @param partitions The partitions which should be resumed
-     * @throws IllegalStateException if one of the provided partitions is not assigned to this consumer
+     * @throws IllegalStateException if any of the provided partitions are not currently assigned to this consumer
      */
     @Override
     public void resume(Collection partitions) {
@@ -1580,22 +2061,48 @@ public Set paused() {
      * If the message format version in a partition is before 0.10.0, i.e. the messages do not have timestamps, null
      * will be returned for that partition.
      *
-     * Notice that this method may block indefinitely if the partition does not exist.
-     *
      * @param timestampsToSearch the mapping from partition to the timestamp to look up.
+     *
      * @return a mapping from partition to the timestamp and offset of the first message with timestamp greater
      *         than or equal to the target timestamp. {@code null} will be returned for the partition if there is no
      *         such message.
      * @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(s). See the exception for more details
-     * @throws IllegalArgumentException if the target timestamp is negative.
+     * @throws IllegalArgumentException if the target timestamp is negative
      * @throws org.apache.kafka.common.errors.TimeoutException if the offset metadata could not be fetched before
-     *         expiration of the configured request timeout
+     *         the amount of time allocated by {@code default.api.timeout.ms} expires.
      * @throws org.apache.kafka.common.errors.UnsupportedVersionException if the broker does not support looking up
-     *         the offsets by timestamp.
+     *         the offsets by timestamp
      */
     @Override
     public Map offsetsForTimes(Map timestampsToSearch) {
+        return offsetsForTimes(timestampsToSearch, Duration.ofMillis(defaultApiTimeoutMs));
+    }
+
+    /**
+     * Look up the offsets for the given partitions by timestamp. The returned offset for each partition is the
+     * earliest offset whose timestamp is greater than or equal to the given timestamp in the corresponding partition.
+     *
+     * This is a blocking call. The consumer does not have to be assigned the partitions.
+     * If the message format version in a partition is before 0.10.0, i.e. the messages do not have timestamps, null
+     * will be returned for that partition.
+     *
+     * @param timestampsToSearch the mapping from partition to the timestamp to look up.
+     * @param timeout The maximum amount of time to await retrieval of the offsets
+     *
+     * @return a mapping from partition to the timestamp and offset of the first message with timestamp greater
+     *         than or equal to the target timestamp. {@code null} will be returned for the partition if there is no
+     *         such message.
+     * @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(s). See the exception for more details
+     * @throws IllegalArgumentException if the target timestamp is negative
+     * @throws org.apache.kafka.common.errors.TimeoutException if the offset metadata could not be fetched before
+     *         expiration of the passed timeout
+     * @throws org.apache.kafka.common.errors.UnsupportedVersionException if the broker does not support looking up
+     *         the offsets by timestamp
+     */
+    @Override
+    public Map offsetsForTimes(Map timestampsToSearch, Duration timeout) {
         acquireAndEnsureOpen();
         try {
             for (Map.Entry entry : timestampsToSearch.entrySet()) {
@@ -1605,7 +2112,7 @@ public Map offsetsForTimes(Map offsetsForTimes(Map
-     * Notice that this method may block indefinitely if the partition does not exist.
      * This method does not change the current consumer position of the partitions.
      *
      * @see #seekToBeginning(Collection)
@@ -1624,28 +2130,48 @@ public Map offsetsForTimes(Map beginningOffsets(Collection partitions) {
+        return beginningOffsets(partitions, Duration.ofMillis(defaultApiTimeoutMs));
+    }
+
+    /**
+     * Get the first offset for the given partitions.
+     * 
    
+     * This method does not change the current consumer position of the partitions.
+     *
+     * @see #seekToBeginning(Collection)
+     *
+     * @param partitions the partitions to get the earliest offsets
+     * @param timeout The maximum amount of time to await retrieval of the beginning offsets
+     *
+     * @return The earliest available offsets for the given partitions
+     * @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(s). See the exception for more details
+     * @throws org.apache.kafka.common.errors.TimeoutException if the offset metadata could not be fetched before
+     *         expiration of the passed timeout
+     */
+    @Override
+    public Map beginningOffsets(Collection partitions, Duration timeout) {
         acquireAndEnsureOpen();
         try {
-            return fetcher.beginningOffsets(partitions, requestTimeoutMs);
+            return fetcher.beginningOffsets(partitions, time.timer(timeout));
         } finally {
             release();
         }
     }
 
     /**
-     * Get the last offset for the given partitions. The last offset of a partition is the offset of the upcoming
-     * message, i.e. the offset of the last available message + 1.
+     * Get the end offsets for the given partitions. In the default {@code read_uncommitted} isolation level, the end
+     * offset is the high watermark (that is, the offset of the last successfully replicated message plus one). For
+     * {@code read_committed} consumers, the end offset is the last stable offset (LSO), which is the minimum of
+     * the high watermark and the smallest offset of any open transaction. Finally, if the partition has never been
+     * written to, the end offset is 0.
+     *
      * 
    
-     * Notice that this method may block indefinitely if the partition does not exist.
      * This method does not change the current consumer position of the partitions.
-     * 
    
-     * When {@code isolation.level=read_committed} the last offset will be the Last Stable Offset (LSO).
-     * This is the offset of the first message with an open transaction. The LSO moves forward as transactions
-     * are completed.
      *
      * @see #seekToEnd(Collection)
      *
@@ -1654,23 +2180,98 @@ public Map beginningOffsets(Collection par
      * @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(s). See the exception for more details
      * @throws org.apache.kafka.common.errors.TimeoutException if the offset metadata could not be fetched before
-     *         expiration of the configured request timeout
+     *         the amount of time allocated by {@code request.timeout.ms} expires
      */
     @Override
     public Map endOffsets(Collection partitions) {
+        return endOffsets(partitions, Duration.ofMillis(requestTimeoutMs));
+    }
+
+    /**
+     * Get the end offsets for the given partitions. In the default {@code read_uncommitted} isolation level, the end
+     * offset is the high watermark (that is, the offset of the last successfully replicated message plus one). For
+     * {@code read_committed} consumers, the end offset is the last stable offset (LSO), which is the minimum of
+     * the high watermark and the smallest offset of any open transaction. Finally, if the partition has never been
+     * written to, the end offset is 0.
+     *
+     * 
    
+     * This method does not change the current consumer position of the partitions.
+     *
+     * @see #seekToEnd(Collection)
+     *
+     * @param partitions the partitions to get the end offsets.
+     * @param timeout The maximum amount of time to await retrieval of the end offsets
+     *
+     * @return The end offsets for the given partitions.
+     * @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(s). See the exception for more details
+     * @throws org.apache.kafka.common.errors.TimeoutException if the offsets could not be fetched before
+     *         expiration of the passed timeout
+     */
+    @Override
+    public Map endOffsets(Collection partitions, Duration timeout) {
+        acquireAndEnsureOpen();
+        try {
+            return fetcher.endOffsets(partitions, time.timer(timeout));
+        } finally {
+            release();
+        }
+    }
+
+    /**
+     * Return the current group metadata associated with this consumer.
+     *
+     * @return consumer group metadata
+     * @throws org.apache.kafka.common.errors.InvalidGroupIdException if consumer does not have a group
+     */
+    @Override
+    public ConsumerGroupMetadata groupMetadata() {
         acquireAndEnsureOpen();
         try {
-            return fetcher.endOffsets(partitions, requestTimeoutMs);
+            maybeThrowInvalidGroupIdException();
+            return coordinator.groupMetadata();
         } finally {
             release();
         }
+    }
 
+    /**
+     * Alert the consumer to trigger a new rebalance by rejoining the group. This is a nonblocking call that forces
+     * the consumer to trigger a new rebalance on the next {@link #poll(Duration)} call. Note that this API does not
+     * itself initiate the rebalance, so you must still call {@link #poll(Duration)}. If a rebalance is already in
+     * progress this call will be a no-op. If you wish to force an additional rebalance you must complete the current
+     * one by calling poll before retrying this API.
+     * 
    
+     * You do not need to call this during normal processing, as the consumer group will manage itself
+     * automatically and rebalance when necessary. However there may be situations where the application wishes to
+     * trigger a rebalance that would otherwise not occur. For example, if some condition external and invisible to
+     * the Consumer and its group changes in a way that would affect the userdata encoded in the
+     * {@link org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.Subscription Subscription}, the Consumer
+     * will not be notified and no rebalance will occur. This API can be used to force the group to rebalance so that
+     * the assignor can perform a partition reassignment based on the latest userdata. If your assignor does not use
+     * this userdata, or you do not use a custom
+     * {@link org.apache.kafka.clients.consumer.ConsumerPartitionAssignor ConsumerPartitionAssignor}, you should not
+     * use this API.
+     *
+     * @throws java.lang.IllegalStateException if the consumer does not use group subscription
+     */
+    @Override
+    public void enforceRebalance() {
+        acquireAndEnsureOpen();
+        try {
+            if (coordinator == null) {
+                throw new IllegalStateException("Tried to force a rebalance but consumer does not have a group.");
+            }
+            coordinator.requestRejoin();
+        } finally {
+            release();
+        }
     }
 
     /**
      * Close the consumer, waiting for up to the default timeout of 30 seconds for any needed cleanup.
      * If auto-commit is enabled, this will commit the current offsets if possible within the default
-     * timeout. See {@link #close(long, TimeUnit)} for details. Note that {@link #wakeup()}
+     * timeout. See {@link #close(Duration)} for details. Note that {@link #wakeup()}
      * cannot be used to interrupt close.
      *
      * @throws org.apache.kafka.common.errors.InterruptException if the calling thread is interrupted
@@ -1679,7 +2280,7 @@ public Map endOffsets(Collection partition
      */
     @Override
     public void close() {
-        close(DEFAULT_CLOSE_TIMEOUT_MS, TimeUnit.MILLISECONDS);
+        close(Duration.ofMillis(DEFAULT_CLOSE_TIMEOUT_MS));
     }
 
     /**
@@ -1696,17 +2297,43 @@ public void close() {
      * @throws IllegalArgumentException If the {@code timeout} is negative.
      * @throws InterruptException If the thread is interrupted before or while this function is called
      * @throws org.apache.kafka.common.KafkaException for any other error during close
+     *
+     * @deprecated Since 2.0. Use {@link #close(Duration)} or {@link #close()}.
      */
+    @Deprecated
+    @Override
     public void close(long timeout, TimeUnit timeUnit) {
-        if (timeout < 0)
+        close(Duration.ofMillis(timeUnit.toMillis(timeout)));
+    }
+
+    /**
+     * Tries to close the consumer cleanly within the specified timeout. This method waits up to
+     * {@code timeout} for the consumer to complete pending commits and leave the group.
+     * If auto-commit is enabled, this will commit the current offsets if possible within the
+     * timeout. If the consumer is unable to complete offset commits and gracefully leave the group
+     * before the timeout expires, the consumer is force closed. Note that {@link #wakeup()} cannot be
+     * used to interrupt close.
+     *
+     * @param timeout The maximum time to wait for consumer to close gracefully. The value must be
+     *                non-negative. Specifying a timeout of zero means do not wait for pending requests to complete.
+     *
+     * @throws IllegalArgumentException If the {@code timeout} is negative.
+     * @throws InterruptException If the thread is interrupted before or while this function is called
+     * @throws org.apache.kafka.common.KafkaException for any other error during close
+     */
+    @Override
+    public void close(Duration timeout) {
+        if (timeout.toMillis() < 0)
             throw new IllegalArgumentException("The timeout cannot be negative.");
         acquire();
         try {
             if (!closed) {
-                closed = true;
-                close(timeUnit.toMillis(timeout), false);
+                // 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();
         }
     }
@@ -1736,17 +2363,18 @@ private void close(long timeoutMs, boolean swallowException) {
         AtomicReference firstException = new AtomicReference<>();
         try {
             if (coordinator != null)
-                coordinator.close(Math.min(timeoutMs, requestTimeoutMs));
+                coordinator.close(time.timer(Math.min(timeoutMs, requestTimeoutMs)));
         } catch (Throwable t) {
             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();
@@ -1762,28 +2390,35 @@ private void close(long timeoutMs, boolean swallowException) {
      * Set the fetch position to the committed position (if there is one)
      * or reset it using the offset reset policy the user has configured.
      *
-     * @param partitions The partitions that needs updating fetch positions
      * @throws org.apache.kafka.common.errors.AuthenticationException if authentication fails. See the exception for more details
      * @throws NoOffsetForPartitionException If no offset is stored for a given partition and no offset reset policy is
      *             defined
+     * @return true iff the operation completed without timing out
      */
-    private void updateFetchPositions(Set partitions) {
-        // lookup any positions for partitions which are awaiting reset (which may be the
-        // case if the user called seekToBeginning or seekToEnd. We do this check first to
-        // avoid an unnecessary lookup of committed offsets (which typically occurs when
-        // the user is manually assigning partitions and managing their own offsets).
-        fetcher.resetOffsetsIfNeeded(partitions);
-
-        if (!subscriptions.hasAllFetchPositions(partitions)) {
-            // if we still don't have offsets for the given partitions, then we should either
-            // seek to the last committed position or reset using the auto reset policy
-
-            // first refresh commits for all assigned partitions
-            coordinator.refreshCommittedOffsetsIfNeeded();
-
-            // then do any offset lookups in case some positions are not known
-            fetcher.updateFetchPositions(partitions);
-        }
+    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;
+
+        // If there are any partitions which do not have a valid position and are not
+        // awaiting reset, then we need to fetch committed offsets. We will only do a
+        // coordinator lookup if there are partitions which have missing positions, so
+        // a consumer with manually assigned partitions can avoid a coordinator dependence
+        // by always ensuring that assigned partitions have an initial position.
+        if (coordinator != null && !coordinator.refreshCommittedOffsetsIfNeeded(timer)) return false;
+
+        // If there are partitions still needing a position and a reset policy is defined,
+        // request reset using the default policy. If no reset strategy is defined and there
+        // are partitions with a missing position, then we will raise an exception.
+        subscriptions.resetInitializingPositions();
+
+        // Finally send an asynchronous request to lookup and update the positions of any
+        // partitions which are awaiting reset.
+        fetcher.resetOffsetsIfNeeded();
+
+        return true;
     }
 
     /**
@@ -1824,4 +2459,24 @@ private void throwIfNoAssignorsConfigured() {
             throw new IllegalStateException("Must configure at least one partition assigner class name to " +
                 ConsumerConfig.PARTITION_ASSIGNMENT_STRATEGY_CONFIG + " configuration property");
     }
+
+    private void maybeThrowInvalidGroupIdException() {
+        if (!groupId.isPresent())
+            throw new InvalidGroupIdException("To use the group management or offset commit APIs, you must " +
+                    "provide a valid " + ConsumerConfig.GROUP_ID_CONFIG + " in the consumer configuration.");
+    }
+
+    private void updateLastSeenEpochIfNewer(TopicPartition topicPartition, OffsetAndMetadata offsetAndMetadata) {
+        if (offsetAndMetadata != null)
+            offsetAndMetadata.leaderEpoch().ifPresent(epoch -> metadata.updateLastSeenEpochIfNewer(topicPartition, epoch));
+    }
+
+    // Functions below are for testing only
+    String getClientId() {
+        return clientId;
+    }
+
+    boolean updateAssignmentMetadataIfNeeded(final Timer timer) {
+        return updateAssignmentMetadataIfNeeded(timer, true);
+    }
 }
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..336eed4a3b4bf
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/LogTruncationException.java
@@ -0,0 +1,55 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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 java.util.Collections;
+import java.util.Map;
+
+/**
+ * In the event of an 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 known to diverge from what the
+ * consumer previously read.
+ */
+public class LogTruncationException extends OffsetOutOfRangeException {
+
+    private final Map divergentOffsets;
+
+    public LogTruncationException(String message,
+                                  Map fetchOffsets,
+                                  Map divergentOffsets) {
+        super(message, fetchOffsets);
+        this.divergentOffsets = Collections.unmodifiableMap(divergentOffsets);
+    }
+
+    /**
+     * Get the divergent offsets for the partitions which were truncated. For each
+     * partition, this is the first offset which is known to diverge from what the
+     * consumer read.
+     *
+     * Note that there is no guarantee that this offset will be known. It is necessary
+     * to use {@link #partitions()} to see the set of partitions that were truncated
+     * and then check for the presence of a divergent offset in this map.
+     */
+    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 9b0c058d2689d..7bf4c3f16dc94 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,6 +16,7 @@
  */
 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;
@@ -24,7 +25,9 @@
 import org.apache.kafka.common.PartitionInfo;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.errors.WakeupException;
+import org.apache.kafka.common.utils.LogContext;
 
+import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
@@ -33,36 +36,41 @@
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Queue;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.regex.Pattern;
+import java.util.stream.Collectors;
 
 
 /**
  * A mock of the {@link Consumer} interface you can use for testing code that uses Kafka. This class is  not
  * threadsafe . However, you can use the {@link #schedulePollTask(Runnable)} method to write multithreaded tests
- * where a driver thread waits for {@link #poll(long)} to be called by a background thread and then can safely perform
+ * where a driver thread waits for {@link #poll(Duration)} to be called by a background thread and then can safely perform
  * operations during a callback.
  */
 public class MockConsumer implements Consumer {
 
     private final Map> partitions;
     private final SubscriptionState subscriptions;
-    private Map>> records;
-    private Set paused;
-    private boolean closed;
     private final Map beginningOffsets;
     private final Map endOffsets;
+    private final Map committed;
+    private final Queue pollTasks;
+    private final Set paused;
 
-    private Queue pollTasks;
-    private KafkaException exception;
-
+    private Map>> records;
+    private KafkaException pollException;
+    private KafkaException offsetsException;
     private AtomicBoolean wakeup;
+    private Duration lastPollTimeout;
+    private boolean closed;
+    private boolean shouldRebalance;
 
     public MockConsumer(OffsetResetStrategy offsetResetStrategy) {
-        this.subscriptions = new SubscriptionState(offsetResetStrategy);
+        this.subscriptions = new SubscriptionState(new LogContext(), offsetResetStrategy);
         this.partitions = new HashMap<>();
         this.records = new HashMap<>();
         this.paused = new HashSet<>();
@@ -70,8 +78,10 @@ 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<>();
+        this.shouldRebalance = false;
     }
 
     @Override
@@ -99,6 +109,7 @@ public synchronized void subscribe(Collection topics) {
     @Override
     public synchronized void subscribe(Pattern pattern, final ConsumerRebalanceListener listener) {
         ensureNotClosed();
+        committed.clear();
         this.subscriptions.subscribe(pattern, listener);
         Set topicsToSubscribe = new HashSet<>();
         for (String topic: partitions.keySet()) {
@@ -108,6 +119,14 @@ public synchronized void subscribe(Pattern pattern, final ConsumerRebalanceListe
         }
         ensureNotClosed();
         this.subscriptions.subscribeFromPattern(topicsToSubscribe);
+        final Set assignedPartitions = new HashSet<>();
+        for (final String topic : topicsToSubscribe) {
+            for (final PartitionInfo info : this.partitions.get(topic)) {
+                assignedPartitions.add(new TopicPartition(topic, info.partition()));
+            }
+
+        }
+        subscriptions.assignFromSubscribed(assignedPartitions);
     }
 
     @Override
@@ -118,25 +137,36 @@ public synchronized void subscribe(Pattern pattern) {
     @Override
     public synchronized void subscribe(Collection topics, final ConsumerRebalanceListener listener) {
         ensureNotClosed();
+        committed.clear();
         this.subscriptions.subscribe(new HashSet<>(topics), listener);
     }
 
     @Override
     public synchronized void assign(Collection partitions) {
         ensureNotClosed();
+        committed.clear();
         this.subscriptions.assignFromUser(new HashSet<>(partitions));
     }
 
     @Override
     public synchronized void unsubscribe() {
         ensureNotClosed();
+        committed.clear();
         subscriptions.unsubscribe();
     }
 
+    @Deprecated
     @Override
     public synchronized ConsumerRecords poll(long timeout) {
+        return poll(Duration.ofMillis(timeout));
+    }
+
+    @Override
+    public synchronized ConsumerRecords poll(final Duration timeout) {
         ensureNotClosed();
 
+        lastPollTimeout = timeout;
+
         // Synchronize around the entire execution so new tasks to be triggered on subsequent poll calls can be added in
         // the callback
         synchronized (pollTasks) {
@@ -150,60 +180,78 @@ public synchronized ConsumerRecords poll(long 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;
         }
 
         // Handle seeks that need to wait for a poll() call to be processed
-        for (TopicPartition tp : subscriptions.missingFetchPositions())
-            updateFetchPosition(tp);
+        for (TopicPartition tp : subscriptions.assignedPartitions())
+            if (!subscriptions.hasValidPosition(tp))
+                updateFetchPosition(tp);
 
         // update the consumed offset
         final Map>> results = new HashMap<>();
-        for (final TopicPartition topicPartition : records.keySet()) {
-            results.put(topicPartition, new ArrayList>());
-        }
+        final List toClear = new ArrayList<>();
 
         for (Map.Entry>> entry : this.records.entrySet()) {
             if (!subscriptions.isPaused(entry.getKey())) {
                 final List> recs = entry.getValue();
                 for (final ConsumerRecord rec : recs) {
-                    if (assignment().contains(entry.getKey()) && rec.offset() >= subscriptions.position(entry.getKey())) {
-                        results.get(entry.getKey()).add(rec);
-                        subscriptions.position(entry.getKey(), rec.offset() + 1);
+                    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() >= position) {
+                        results.computeIfAbsent(entry.getKey(), partition -> new ArrayList<>()).add(rec);
+                        Metadata.LeaderAndEpoch leaderAndEpoch = new Metadata.LeaderAndEpoch(Optional.empty(), rec.leaderEpoch());
+                        SubscriptionState.FetchPosition newPosition = new SubscriptionState.FetchPosition(
+                                rec.offset() + 1, rec.leaderEpoch(), leaderAndEpoch);
+                        subscriptions.position(entry.getKey(), newPosition);
                     }
                 }
+                toClear.add(entry.getKey());
             }
         }
-        this.records.clear();
+
+        toClear.forEach(p -> this.records.remove(p));
         return new ConsumerRecords<>(results);
     }
 
     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.get(tp);
-        if (recs == null) {
-            recs = new ArrayList>();
-            this.records.put(tp, recs);
-        }
+        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
     public synchronized void commitAsync(Map offsets, OffsetCommitCallback callback) {
         ensureNotClosed();
         for (Map.Entry entry : offsets.entrySet())
-            subscriptions.committed(entry.getKey(), entry.getValue());
+            committed.put(entry.getKey(), entry.getValue());
         if (callback != null) {
             callback.onComplete(offsets, null);
         }
@@ -230,6 +278,16 @@ public synchronized void commitSync() {
         commitSync(this.subscriptions.allConsumed());
     }
 
+    @Override
+    public synchronized void commitSync(Duration timeout) {
+        commitSync(this.subscriptions.allConsumed());
+    }
+
+    @Override
+    public void commitSync(Map offsets, final Duration timeout) {
+        commitSync(offsets);
+    }
+
     @Override
     public synchronized void seek(TopicPartition partition, long offset) {
         ensureNotClosed();
@@ -237,12 +295,36 @@ public synchronized void seek(TopicPartition partition, long offset) {
     }
 
     @Override
-    public synchronized OffsetAndMetadata committed(TopicPartition partition) {
+    public void seek(TopicPartition partition, OffsetAndMetadata offsetAndMetadata) {
         ensureNotClosed();
-        if (subscriptions.isAssigned(partition)) {
-            return subscriptions.committed(partition);
-        }
-        return new OffsetAndMetadata(0);
+        subscriptions.seek(partition, offsetAndMetadata.offset());
+    }
+
+    @Deprecated
+    @Override
+    public synchronized OffsetAndMetadata committed(final TopicPartition partition) {
+        return committed(Collections.singleton(partition)).get(partition);
+    }
+
+    @Deprecated
+    @Override
+    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
@@ -250,19 +332,23 @@ 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
+    public synchronized long position(TopicPartition partition, final Duration timeout) {
+        return position(partition);
     }
 
     @Override
     public synchronized void seekToBeginning(Collection partitions) {
         ensureNotClosed();
-        for (TopicPartition tp : partitions)
-            subscriptions.needOffsetReset(tp, OffsetResetStrategy.EARLIEST);
+        subscriptions.requestOffsetReset(partitions, OffsetResetStrategy.EARLIEST);
     }
 
     public synchronized void updateBeginningOffsets(Map newOffsets) {
@@ -272,11 +358,10 @@ public synchronized void updateBeginningOffsets(Map newOff
     @Override
     public synchronized void seekToEnd(Collection partitions) {
         ensureNotClosed();
-        for (TopicPartition tp : partitions)
-            subscriptions.needOffsetReset(tp, OffsetResetStrategy.LATEST);
+        subscriptions.requestOffsetReset(partitions, OffsetResetStrategy.LATEST);
     }
 
-    public synchronized void updateEndOffsets(Map newOffsets) {
+    public synchronized void updateEndOffsets(final Map newOffsets) {
         endOffsets.putAll(newOffsets);
     }
 
@@ -326,6 +411,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);
@@ -338,6 +428,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 = endOffsets.get(tp);
@@ -353,9 +448,9 @@ public synchronized void close() {
         close(KafkaConsumer.DEFAULT_CLOSE_TIMEOUT_MS, TimeUnit.MILLISECONDS);
     }
 
+    @Deprecated
     @Override
     public synchronized void close(long timeout, TimeUnit unit) {
-        ensureNotClosed();
         this.closed = true;
     }
 
@@ -369,7 +464,7 @@ public synchronized void wakeup() {
     }
 
     /**
-     * Schedule a task to be executed during a poll(). One enqueued task will be executed per {@link #poll(long)}
+     * Schedule a task to be executed during a poll(). One enqueued task will be executed per {@link #poll(Duration)}
      * invocation. You can use this repeatedly to mock out multiple responses to poll invocations.
      * @param task the task to be executed
      */
@@ -380,12 +475,7 @@ public synchronized void schedulePollTask(Runnable task) {
     }
 
     public synchronized void scheduleNopPollTask() {
-        schedulePollTask(new Runnable() {
-            @Override
-            public void run() {
-                // noop
-            }
-        });
+        schedulePollTask(() -> { });
     }
 
     public synchronized Set paused() {
@@ -400,11 +490,11 @@ private void ensureNotClosed() {
     private void updateFetchPosition(TopicPartition tp) {
         if (subscriptions.isOffsetResetNeeded(tp)) {
             resetOffsetPosition(tp);
-        } else if (subscriptions.committed(tp) == null) {
-            subscriptions.needOffsetReset(tp);
+        } else if (!committed.containsKey(tp)) {
+            subscriptions.requestOffsetReset(tp);
             resetOffsetPosition(tp);
         } else {
-            subscriptions.seek(tp, subscriptions.committed(tp).offset());
+            subscriptions.seek(tp, committed.get(tp).offset());
         }
     }
 
@@ -424,4 +514,57 @@ private void resetOffsetPosition(TopicPartition tp) {
         }
         seek(tp, offset);
     }
+
+    @Override
+    public List partitionsFor(String topic, Duration timeout) {
+        return partitionsFor(topic);
+    }
+
+    @Override
+    public Map> listTopics(Duration timeout) {
+        return listTopics();
+    }
+
+    @Override
+    public Map offsetsForTimes(Map timestampsToSearch,
+            Duration timeout) {
+        return offsetsForTimes(timestampsToSearch);
+    }
+
+    @Override
+    public Map beginningOffsets(Collection partitions, Duration timeout) {
+        return beginningOffsets(partitions);
+    }
+
+    @Override
+    public Map endOffsets(Collection partitions, Duration timeout) {
+        return endOffsets(partitions);
+    }
+
+    @Override
+    public ConsumerGroupMetadata groupMetadata() {
+        return new ConsumerGroupMetadata("dummy.group.id", 1, "1", Optional.empty());
+    }
+
+    @Override
+    public void enforceRebalance() {
+        shouldRebalance = true;
+    }
+
+    public boolean shouldRebalance() {
+        return shouldRebalance;
+    }
+
+    public void resetShouldRebalance() {
+        shouldRebalance = false;
+    }
+
+    public Duration lastPollTimeout() {
+        return lastPollTimeout;
+    }
+
+    @Override
+    public void close(Duration timeout) {
+        close();
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/OffsetAndMetadata.java b/clients/src/main/java/org/apache/kafka/clients/consumer/OffsetAndMetadata.java
index 262d8f8fc6c25..d6b3b947c209d 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/OffsetAndMetadata.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/OffsetAndMetadata.java
@@ -19,6 +19,8 @@
 import org.apache.kafka.common.requests.OffsetFetchResponse;
 
 import java.io.Serializable;
+import java.util.Objects;
+import java.util.Optional;
 
 /**
  * The Kafka offset commit API allows users to provide additional metadata (in the form of a string)
@@ -26,16 +28,30 @@
  * node made the commit, what time the commit was made, etc.
  */
 public class OffsetAndMetadata implements Serializable {
+    private static final long serialVersionUID = 2019555404968089681L;
+
     private final long offset;
     private final String metadata;
 
+    // We use null to represent the absence of a leader epoch to simplify serialization.
+    // I.e., older serializations of this class which do not have this field will automatically
+    // initialize its value to null.
+    private final Integer leaderEpoch;
+
     /**
      * Construct a new OffsetAndMetadata object for committing through {@link KafkaConsumer}.
+     *
      * @param offset The offset to be committed
+     * @param leaderEpoch Optional leader epoch of the last consumed record
      * @param metadata Non-null metadata
      */
-    public OffsetAndMetadata(long offset, String metadata) {
+    public OffsetAndMetadata(long offset, Optional leaderEpoch, String metadata) {
+        if (offset < 0)
+            throw new IllegalArgumentException("Invalid negative offset");
+
         this.offset = offset;
+        this.leaderEpoch = leaderEpoch.orElse(null);
+
         // The server converts null metadata to an empty string. So we store it as an empty string as well on the client
         // to be consistent.
         if (metadata == null)
@@ -44,6 +60,15 @@ public OffsetAndMetadata(long offset, String metadata) {
             this.metadata = metadata;
     }
 
+    /**
+     * Construct a new OffsetAndMetadata object for committing through {@link KafkaConsumer}.
+     * @param offset The offset to be committed
+     * @param metadata Non-null metadata
+     */
+    public OffsetAndMetadata(long offset, String metadata) {
+        this(offset, Optional.empty(), metadata);
+    }
+
     /**
      * Construct a new OffsetAndMetadata object for committing through {@link KafkaConsumer}. The metadata
      * associated with the commit will be empty.
@@ -61,29 +86,41 @@ public String metadata() {
         return metadata;
     }
 
+    /**
+     * Get the leader epoch of the previously consumed record (if one is known). Log truncation is detected
+     * if there exists a leader epoch which is larger than this epoch and begins at an offset earlier than
+     * the committed offset.
+     *
+     * @return the leader epoch or empty if not known
+     */
+    public Optional leaderEpoch() {
+        if (leaderEpoch == null || leaderEpoch < 0)
+            return Optional.empty();
+        return Optional.of(leaderEpoch);
+    }
+
     @Override
     public boolean equals(Object o) {
         if (this == o) return true;
         if (o == null || getClass() != o.getClass()) return false;
-
         OffsetAndMetadata that = (OffsetAndMetadata) o;
-
-        if (offset != that.offset) return false;
-        return metadata.equals(that.metadata);
+        return offset == that.offset &&
+                Objects.equals(metadata, that.metadata) &&
+                Objects.equals(leaderEpoch, that.leaderEpoch);
     }
 
     @Override
     public int hashCode() {
-        int result = (int) (offset ^ (offset >>> 32));
-        result = 31 * result + metadata.hashCode();
-        return result;
+        return Objects.hash(offset, metadata, leaderEpoch);
     }
 
     @Override
     public String toString() {
         return "OffsetAndMetadata{" +
                 "offset=" + offset +
+                ", leaderEpoch=" + leaderEpoch +
                 ", metadata='" + metadata + '\'' +
                 '}';
     }
+
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/OffsetAndTimestamp.java b/clients/src/main/java/org/apache/kafka/clients/consumer/OffsetAndTimestamp.java
index 3af057f9ce4ae..40d993074a25e 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/OffsetAndTimestamp.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/OffsetAndTimestamp.java
@@ -16,7 +16,8 @@
  */
 package org.apache.kafka.clients.consumer;
 
-import org.apache.kafka.common.utils.Utils;
+import java.util.Objects;
+import java.util.Optional;
 
 /**
  * A container class for offset and timestamp.
@@ -24,12 +25,22 @@
 public final class OffsetAndTimestamp {
     private final long timestamp;
     private final long offset;
+    private final Optional leaderEpoch;
 
     public OffsetAndTimestamp(long offset, long timestamp) {
+        this(offset, timestamp, Optional.empty());
+    }
+
+    public OffsetAndTimestamp(long offset, long timestamp, Optional leaderEpoch) {
+        if (offset < 0)
+            throw new IllegalArgumentException("Invalid negative offset");
+
+        if (timestamp < 0)
+            throw new IllegalArgumentException("Invalid negative timestamp");
+
         this.offset = offset;
-        assert this.offset >= 0;
         this.timestamp = timestamp;
-        assert this.timestamp >= 0;
+        this.leaderEpoch = leaderEpoch;
     }
 
     public long timestamp() {
@@ -40,21 +51,35 @@ public long offset() {
         return offset;
     }
 
+    /**
+     * Get the leader epoch corresponding to the offset that was found (if one exists).
+     * This can be provided to seek() to ensure that the log hasn't been truncated prior to fetching.
+     *
+     * @return The leader epoch or empty if it is not known
+     */
+    public Optional leaderEpoch() {
+        return leaderEpoch;
+    }
+
     @Override
     public String toString() {
-        return "(timestamp=" + timestamp + ", offset=" + offset + ")";
+        return "(timestamp=" + timestamp +
+                ", leaderEpoch=" + leaderEpoch.orElse(null) +
+                ", offset=" + offset + ")";
     }
 
     @Override
-    public int hashCode() {
-        return 31 * Utils.longHashcode(timestamp) + Utils.longHashcode(offset);
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        OffsetAndTimestamp that = (OffsetAndTimestamp) o;
+        return timestamp == that.timestamp &&
+                offset == that.offset &&
+                Objects.equals(leaderEpoch, that.leaderEpoch);
     }
 
     @Override
-    public boolean equals(Object o) {
-        if (o == null || !(o instanceof OffsetAndTimestamp))
-            return false;
-        OffsetAndTimestamp other = (OffsetAndTimestamp) o;
-        return this.timestamp == other.timestamp() && this.offset == other.offset();
+    public int hashCode() {
+        return Objects.hash(timestamp, offset, leaderEpoch);
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/OffsetCommitCallback.java b/clients/src/main/java/org/apache/kafka/clients/consumer/OffsetCommitCallback.java
index b217a63257485..53e8ae7b906a0 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/OffsetCommitCallback.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/OffsetCommitCallback.java
@@ -18,12 +18,13 @@
 
 import org.apache.kafka.common.TopicPartition;
 
+import java.time.Duration;
 import java.util.Collection;
 import java.util.Map;
 
 /**
  * A callback interface that the user can implement to trigger custom actions when a commit request completes. The callback
- * may be executed in any thread calling {@link Consumer#poll(long) poll()}.
+ * may be executed in any thread calling {@link Consumer#poll(java.time.Duration) poll()}.
  */
 public interface OffsetCommitCallback {
 
@@ -37,6 +38,9 @@ public interface OffsetCommitCallback {
      * @throws org.apache.kafka.clients.consumer.CommitFailedException if the commit failed and cannot be retried.
      *             This can only occur if you are using automatic group management with {@link KafkaConsumer#subscribe(Collection)},
      *             or if there is an active group with the same groupId which is using group management.
+     * @throws org.apache.kafka.common.errors.RebalanceInProgressException if the commit failed because
+     *            it is in the middle of a rebalance. In such cases
+     *            commit could be retried after the rebalance is completed with the {@link KafkaConsumer#poll(Duration)} call.
      * @throws org.apache.kafka.common.errors.WakeupException if {@link KafkaConsumer#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
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/OffsetOutOfRangeException.java b/clients/src/main/java/org/apache/kafka/clients/consumer/OffsetOutOfRangeException.java
index dae19b29f2d19..c98e22fc82587 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/OffsetOutOfRangeException.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/OffsetOutOfRangeException.java
@@ -31,10 +31,18 @@ public class OffsetOutOfRangeException extends InvalidOffsetException {
     private final Map offsetOutOfRangePartitions;
 
     public OffsetOutOfRangeException(Map offsetOutOfRangePartitions) {
-        super("Offsets out of range with no configured reset policy for partitions: " + offsetOutOfRangePartitions);
+        this("Offsets out of range with no configured reset policy for partitions: " +
+            offsetOutOfRangePartitions, offsetOutOfRangePartitions);
+    }
+
+    public OffsetOutOfRangeException(String message, Map offsetOutOfRangePartitions) {
+        super(message);
         this.offsetOutOfRangePartitions = offsetOutOfRangePartitions;
     }
 
+    /**
+     * Get a map of the topic partitions and the respective out-of-range fetch offsets.
+     */
     public Map offsetOutOfRangePartitions() {
         return offsetOutOfRangePartitions;
     }
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 d8d72ee601bb2..21eb47a19531c 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
@@ -26,17 +26,43 @@
 import java.util.Map;
 
 /**
- * The range assignor works on a per-topic basis. For each topic, we lay out the available partitions in numeric order
+ * 
    The range assignor works on a per-topic basis. For each topic, we lay out the available partitions in numeric order
  * and the consumers in lexicographic order. We then divide the number of partitions by the total number of
  * consumers to determine the number of partitions to assign to each consumer. If it does not evenly
  * divide, then the first few consumers will have one extra partition.
  *
- * 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:
- * C0: [t0p0, t0p1, t1p0, t1p1]
- * C1: [t0p2, t1p2]
+ * 
    The assignment will be:
+ * 
      
+ * 
    • 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 {
 
@@ -45,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());
+            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)
@@ -80,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/RetriableCommitFailedException.java b/clients/src/main/java/org/apache/kafka/clients/consumer/RetriableCommitFailedException.java
index 69f21a4d9d1f2..f44dce6187cda 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/RetriableCommitFailedException.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/RetriableCommitFailedException.java
@@ -22,13 +22,9 @@ public class RetriableCommitFailedException extends RetriableException {
 
     private static final long serialVersionUID = 1L;
 
-    public static RetriableCommitFailedException withUnderlyingMessage(String additionalMessage) {
-        return new RetriableCommitFailedException("Offset commit failed with a retriable exception. " +
-                "You should retry committing offsets. The underlying error was: " + additionalMessage);
-    }
-
     public RetriableCommitFailedException(Throwable t) {
-        super("Offset commit failed with a retriable exception. You should retry committing offsets.", t);
+        super("Offset commit failed with a retriable exception. You should retry committing " +
+                "the latest consumed offsets.", t);
     }
 
     public RetriableCommitFailedException(String message) {
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 7e8d6f2c8077e..edecce84baeaa 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
@@ -29,29 +29,73 @@
 import java.util.TreeSet;
 
 /**
- * The round robin assignor lays out all the available partitions and all the available consumers. It
+ * 
    The round robin assignor lays out all the available partitions and all the available consumers. It
  * then proceeds to do a round robin assignment from partition to consumer. If the subscriptions of all consumer
  * 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:
- * C0: [t0p0, t0p2, t1p1]
- * C1: [t0p1, t1p0, t1p2]
+ * 
    The assignment will be:
+ * 
      
+ * 
    • C0: [t0p0, t0p2, t1p1]
+ * 
    • C1: [t0p1, t1p0, t1p2]
+ * 
    
  *
- * When subscriptions differ across consumer instances, the assignment process still considers each
+ * 
    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.
  *
- * Tha assignment will be:
- * C0: [t0p0]
- * C1: [t1p0]
- * C2: [t1p1, t2p0, t2p1, t2p2]
+ * 
    That assignment will be:
+ * 
      
+ * 
    • C0: [t0p0]
+ * 
    • 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 {
 
@@ -59,22 +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 247b6194693fa..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,512 +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());
-            for (String topic: entry.getValue().topics()) {
-                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);
+    private int generation = DEFAULT_GENERATION; // consumer group generation
 
-        // 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());
-
-        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 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
-     */
-    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(topicPartition + " is assigned to more than one consumer.");
-                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(topicPartition + " can be moved from consumer " + otherConsumer + " to consumer " + consumer + " for a more balanced assignment.");
-                        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 " + consumer + " is assigned more partitions than the maximum possible.");
-
-        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);
-        }
-
-        fixedAssignments.clear();
-    }
-
-    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 '" + partition + "' to be assigned to a consumer");
-
-                // 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;
-            }
+    protected MemberData memberData(Subscription subscription) {
+        ByteBuffer userData = subscription.userData();
+        if (userData == null || !userData.hasRemaining()) {
+            return new MemberData(Collections.emptyList(), Optional.empty());
         }
-
-        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);
+        return deserializeTopicPartitionAssignment(userData);
     }
 
-    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.groupDataByTopic(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;
@@ -694,265 +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();
-        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 " + (path.size() - 1) + " was found: " + 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 " + topicMovements.getKey()
-                            + "\nPartition movements for this topic occurred among the following consumer pairs:"
-                            + "\n" + 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 8e56a5bead418..29842a90c3362 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,31 +17,45 @@
 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;
 import org.apache.kafka.common.errors.InterruptException;
+import org.apache.kafka.common.errors.MemberIdRequiredException;
 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.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;
@@ -51,13 +65,17 @@
 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.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;
+import java.util.Objects;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
 
@@ -81,7 +99,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
@@ -92,58 +110,58 @@
  */
 public abstract class AbstractCoordinator implements Closeable {
     public static final String HEARTBEAT_THREAD_PREFIX = "kafka-coordinator-heartbeat-thread";
+    public static final int JOIN_GROUP_TIMEOUT_LAPSE = 5000;
 
-    private enum MemberState {
-        UNJOINED,    // the client is not part of a group
-        REBALANCING, // the client has begun rebalancing
-        STABLE,      // the client has joined and is sending heartbeats
+    protected enum MemberState {
+        UNJOINED,             // the client is not part of a group
+        PREPARING_REBALANCE,  // the client has sent the join group request, but have not received response
+        COMPLETING_REBALANCE, // the client has received join group response, but have not received assignment
+        STABLE;               // the client has joined and is sending heartbeats
+
+        public boolean hasNotJoinedGroup() {
+            return equals(UNJOINED) || equals(PREPARING_REBALANCE);
+        }
     }
 
     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;
-    protected final ConsumerNetworkClient client;
+    private final GroupCoordinatorMetrics sensors;
+    private final GroupRebalanceConfig rebalanceConfig;
+
     protected final Time time;
-    protected final long retryBackoffMs;
+    protected final ConsumerNetworkClient client;
 
-    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 RequestFuture findCoordinatorFuture = null;
+    volatile private RuntimeException findCoordinatorException = null;
     private Generation generation = Generation.NO_GENERATION;
+    private long lastRebalanceStartMs = -1L;
+    private long lastRebalanceEndMs = -1L;
+
+    protected MemberState state = MemberState.UNJOINED;
+
 
-    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,
-                               int heartbeatIntervalMs,
                                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 = groupId;
-        this.rebalanceTimeoutMs = rebalanceTimeoutMs;
-        this.sessionTimeoutMs = sessionTimeoutMs;
-        this.leaveGroupOnClose = leaveGroupOnClose;
-        this.heartbeat = new Heartbeat(sessionTimeoutMs, heartbeatIntervalMs, rebalanceTimeoutMs, retryBackoffMs);
+        this.heartbeat = new Heartbeat(rebalanceConfig, time);
         this.sensors = new GroupCoordinatorMetrics(metrics, metricGrpPrefix);
-        this.retryBackoffMs = retryBackoffMs;
     }
 
     /**
@@ -161,7 +179,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
@@ -175,17 +193,17 @@ 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 is woken up (i.e.
-     * if the invocation raises {@link org.apache.kafka.common.errors.WakeupException}), then it
-     * will be retried on the next call to {@link #ensureActiveGroup()}.
+     * Invoked when a group member has successfully joined a group. If this call fails with an exception,
+     * then it will be retried using the same assignment state on the next call to {@link #ensureActiveGroup()}.
      *
      * @param generation The generation that was joined
      * @param memberId The identifier for the local member in the group
@@ -198,47 +216,51 @@ protected abstract void onJoinComplete(int generation,
                                            ByteBuffer memberAssignment);
 
     /**
-     * Block until the coordinator for this group is known and is ready to receive requests.
+     * 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)
      */
-    public synchronized void ensureCoordinatorReady() {
-        // Using zero as current time since timeout is effectively infinite
-        ensureCoordinatorReady(0, Long.MAX_VALUE);
-    }
+    protected void onLeavePrepare() {}
 
     /**
+     * Visible for testing.
+     *
      * Ensure that the coordinator is ready to receive requests.
-     * @param startTimeMs Current time in milliseconds
-     * @param timeoutMs Maximum time to wait to discover the coordinator
+     *
+     * @param timer Timer bounding how long this method can block
      * @return true If coordinator discovery and initial connection succeeded, false otherwise
      */
-    protected synchronized boolean ensureCoordinatorReady(long startTimeMs, long timeoutMs) {
-        long remainingMs = timeoutMs;
+    protected synchronized boolean ensureCoordinatorReady(final Timer timer) {
+        if (!coordinatorUnknown())
+            return true;
+
+        do {
+            if (findCoordinatorException != null && !(findCoordinatorException instanceof RetriableException)) {
+                final RuntimeException fatalException = findCoordinatorException;
+                findCoordinatorException = null;
+                throw fatalException;
+            }
+            final RequestFuture future = lookupCoordinator();
+            client.poll(future, timer);
 
-        while (coordinatorUnknown()) {
-            RequestFuture future = lookupCoordinator();
-            client.poll(future, remainingMs);
+            if (!future.isDone()) {
+                // ran out of time
+                break;
+            }
 
             if (future.failed()) {
                 if (future.isRetriable()) {
-                    remainingMs = timeoutMs - (time.milliseconds() - startTimeMs);
-                    if (remainingMs <= 0)
-                        break;
-
-                    log.debug("Coordinator discovery failed, refreshing metadata");
-                    client.awaitMetadataUpdate(remainingMs);
+                    log.debug("Coordinator discovery failed, refreshing metadata", future.exception());
+                    client.awaitMetadataUpdate(timer);
                 } else
                     throw future.exception();
-            } else if (coordinator != null && client.connectionFailed(coordinator)) {
+            } else if (coordinator != null && client.isUnavailable(coordinator)) {
                 // we found the coordinator, but the connection has failed, so mark
                 // it dead and backoff before retrying discovery
-                coordinatorDead();
-                time.sleep(retryBackoffMs);
+                markCoordinatorUnknown();
+                timer.sleep(rebalanceConfig.retryBackoffMs);
             }
-
-            remainingMs = timeoutMs - (time.milliseconds() - startTimeMs);
-            if (remainingMs <= 0)
-                break;
-        }
+        } while (coordinatorUnknown() && timer.notExpired());
 
         return !coordinatorUnknown();
     }
@@ -248,10 +270,22 @@ protected synchronized RequestFuture lookupCoordinator() {
             // find a node to ask about the coordinator
             Node node = this.client.leastLoadedNode();
             if (node == null) {
-                log.debug("No broker available to send GroupCoordinator request");
+                log.debug("No broker available to send FindCoordinator request");
                 return RequestFuture.noBrokersAvailable();
-            } else
-                findCoordinatorFuture = sendGroupCoordinatorRequest(node);
+            } else {
+                findCoordinatorFuture = sendFindCoordinatorRequest(node);
+                // remember the exception even after the future is cleared so that
+                // it can still be thrown by the ensureCoordinatorReady caller
+                findCoordinatorFuture.addListener(new RequestFutureListener() {
+                    @Override
+                    public void onSuccess(Void value) {} // do nothing
+
+                    @Override
+                    public void onFailure(RuntimeException e) {
+                        findCoordinatorException = e;
+                    }
+                });
+            }
         }
         return findCoordinatorFuture;
     }
@@ -261,15 +295,14 @@ private synchronized void clearFindCoordinatorFuture() {
     }
 
     /**
-     * Check whether the group should be rejoined (e.g. if metadata changes)
+     * Check whether the group should be rejoined (e.g. if metadata changes) or whether a
+     * rejoin request is already in flight and needs to be completed.
+     *
      * @return true if it should, false otherwise
      */
-    protected synchronized boolean needRejoin() {
-        return rejoinNeeded;
-    }
-
-    private synchronized boolean rejoinIncomplete() {
-        return joinFuture != null;
+    protected synchronized boolean rejoinNeededOrPending() {
+        // if there's a pending joinFuture, we should try to complete handling it.
+        return rejoinNeeded || joinFuture != null;
     }
 
     /**
@@ -278,6 +311,7 @@ private synchronized boolean rejoinIncomplete() {
      * to ensure that the member stays in the group. If an interval of time longer than the
      * provided rebalance timeout expires without calling this method, then the client will proactively
      * leave the group.
+     *
      * @param now current time in milliseconds
      * @throws RuntimeException for unexpected errors raised from the heartbeat thread
      */
@@ -299,8 +333,9 @@ protected synchronized void pollHeartbeat(long now) {
     }
 
     protected synchronized long timeToNextHeartbeat(long now) {
-        // if we have not joined the group, we don't need to send heartbeats
-        if (state == MemberState.UNJOINED)
+        // if we have not joined the group or we are preparing rebalance,
+        // we don't need to send heartbeats
+        if (state.hasNotJoinedGroup())
             return Long.MAX_VALUE;
         return heartbeat.timeToNextHeartbeat(now);
     }
@@ -309,11 +344,27 @@ protected synchronized long timeToNextHeartbeat(long now) {
      * Ensure that the group is active (i.e. joined and synced)
      */
     public void ensureActiveGroup() {
+        while (!ensureActiveGroup(time.timer(Long.MAX_VALUE))) {
+            log.warn("still waiting to ensure active group");
+        }
+    }
+
+    /**
+     * 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) {
         // always ensure that the coordinator is ready because we may have been disconnected
         // when sending heartbeats and does not necessarily require us to rejoin the group.
-        ensureCoordinatorReady();
+        if (!ensureCoordinatorReady(timer)) {
+            return false;
+        }
+
         startHeartbeatThreadIfNeeded();
-        joinGroupIfNeeded();
+        return joinGroupIfNeeded(timer);
     }
 
     private synchronized void startHeartbeatThreadIfNeeded() {
@@ -323,13 +374,8 @@ private synchronized void startHeartbeatThreadIfNeeded() {
         }
     }
 
-    private synchronized void disableHeartbeatThread() {
-        if (heartbeatThread != null)
-            heartbeatThread.disable();
-    }
-
     private void closeHeartbeatThread() {
-        HeartbeatThread thread = null;
+        HeartbeatThread thread;
         synchronized (this) {
             if (heartbeatThread == null)
                 return;
@@ -345,10 +391,26 @@ private void closeHeartbeatThread() {
         }
     }
 
-    // visible for testing. Joins the group without starting the heartbeat thread.
-    void joinGroupIfNeeded() {
-        while (needRejoin() || rejoinIncomplete()) {
-            ensureCoordinatorReady();
+    /**
+     * Joins the group without starting the heartbeat thread.
+     *
+     * If this function returns true, the state must always be in STABLE and heartbeat enabled.
+     * If this function returns false, the state can be in one of the following:
+     *  * UNJOINED: got error response but times out before being able to re-join, heartbeat disabled
+     *  * PREPARING_REBALANCE: not yet received join-group response before timeout, heartbeat disabled
+     *  * COMPLETING_REBALANCE: not yet received sync-group response before timeout, hearbeat enabled
+     *
+     * 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) {
+        while (rejoinNeededOrPending()) {
+            if (!ensureCoordinatorReady(timer)) {
+                return false;
+            }
 
             // call onJoinPrepare if needed. We set a flag to make sure that we do not call it a second
             // time if the client is woken up before a pending rebalance completes. This must be called
@@ -356,32 +418,74 @@ void joinGroupIfNeeded() {
             // 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);
             }
 
-            RequestFuture future = initiateJoinGroup();
-            client.poll(future);
+            final RequestFuture future = initiateJoinGroup();
+            client.poll(future, timer);
+            if (!future.isDone()) {
+                // we ran out of time
+                return false;
+            }
 
             if (future.succeeded()) {
-                onJoinComplete(generation.generationId, generation.memberId, generation.protocol, future.value());
+                Generation generationSnapshot;
+                MemberState stateSnapshot;
 
-                // 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;
+                    stateSnapshot = this.state;
+                }
+
+                if (!generationSnapshot.equals(Generation.NO_GENERATION) && stateSnapshot == MemberState.STABLE) {
+                    // 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.protocolName, 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 to {} and state is now {} before " +
+                         "the rebalance callback is triggered, marking this rebalance as failed and retry",
+                         generationSnapshot, stateSnapshot);
+                    resetStateAndRejoin();
+                    resetJoinGroupFuture();
+                }
             } else {
+                final RuntimeException exception = future.exception();
+
+                // we do not need to log error for memberId required,
+                // since it is not really an error and is transient
+                if (!(exception instanceof MemberIdRequiredException)) {
+                    log.info("Rebalance failed.", exception);
+                }
+
                 resetJoinGroupFuture();
-                RuntimeException exception = future.exception();
                 if (exception instanceof UnknownMemberIdException ||
-                        exception instanceof RebalanceInProgressException ||
-                        exception instanceof IllegalGenerationException)
+                    exception instanceof RebalanceInProgressException ||
+                    exception instanceof IllegalGenerationException ||
+                    exception instanceof MemberIdRequiredException)
                     continue;
                 else if (!future.isRetriable())
                     throw exception;
-                time.sleep(retryBackoffMs);
+
+                resetStateAndRejoin();
+                timer.sleep(rebalanceConfig.retryBackoffMs);
             }
         }
+        return true;
     }
 
     private synchronized void resetJoinGroupFuture() {
@@ -393,34 +497,25 @@ private synchronized RequestFuture initiateJoinGroup() {
         // rebalance in the call to poll below. This ensures that we do not mistakenly attempt
         // to rejoin before the pending rebalance has completed.
         if (joinFuture == null) {
-            // fence off the heartbeat thread explicitly so that it cannot interfere with the join group.
-            // Note that this must come after the call to onJoinPrepare since we must be able to continue
-            // sending heartbeats if that callback takes some time.
-            disableHeartbeatThread();
-
-            state = MemberState.REBALANCING;
+            state = MemberState.PREPARING_REBALANCE;
+            // 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
                 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();
-                    }
+                    // do nothing since all the handler logic are in SyncGroupResponseHandler already
                 }
 
                 @Override
                 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
+                    // after having been woken up, the exception is ignored and we will rejoin;
+                    // this can be triggered when either join or sync request failed
                     synchronized (AbstractCoordinator.this) {
-                        state = MemberState.UNJOINED;
+                        sensors.failedRebalanceSensor.record();
                     }
                 }
             });
@@ -430,76 +525,142 @@ 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
+     *
      * @return A request future which wraps the assignment returned from the group leader
      */
-    private RequestFuture sendJoinGroupRequest() {
+    RequestFuture sendJoinGroupRequest() {
         if (coordinatorUnknown())
             return RequestFuture.coordinatorNotAvailable();
 
         // 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);
-        return client.send(coordinator, requestBuilder)
-                .compose(new JoinGroupResponseHandler());
+
+        // 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(client.defaultRequestTimeoutMs(),
+            rebalanceConfig.rebalanceTimeoutMs + JOIN_GROUP_TIMEOUT_LAPSE);
+        return client.send(coordinator, requestBuilder, joinGroupTimeoutMs)
+                .compose(new JoinGroupResponseHandler(generation));
     }
 
     private class JoinGroupResponseHandler extends CoordinatorResponseHandler {
+        private JoinGroupResponseHandler(final Generation generation) {
+            super(generation);
+        }
+
         @Override
         public void handle(JoinGroupResponse joinResponse, RequestFuture future) {
             Errors error = joinResponse.error();
             if (error == Errors.NONE) {
-                log.debug("Received successful JoinGroup response: {}", joinResponse);
-                sensors.joinLatency.record(response.requestLatencyMs());
+                if (isProtocolTypeInconsistent(joinResponse.data().protocolType())) {
+                    log.error("JoinGroup failed: Inconsistent Protocol Type, received {} but expected {}",
+                        joinResponse.data().protocolType(), protocolType());
+                    future.raise(Errors.INCONSISTENT_GROUP_PROTOCOL);
+                } else {
+                    log.debug("Received successful JoinGroup response: {}", joinResponse);
+                    sensors.joinSensor.record(response.requestLatencyMs());
 
-                synchronized (AbstractCoordinator.this) {
-                    if (state != MemberState.REBALANCING) {
-                        // if the consumer was woken up before a rebalance completes, we may have already left
-                        // 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());
-                        if (joinResponse.isLeader()) {
-                            onJoinLeader(joinResponse).chain(future);
+                    synchronized (AbstractCoordinator.this) {
+                        if (state != MemberState.PREPARING_REBALANCE) {
+                            // if the consumer was woken up before a rebalance completes, we may have already left
+                            // the group. In this case, we do not want to continue with the sync group.
+                            future.raise(new UnjoinedGroupException());
                         } else {
-                            onJoinFollower().chain(future);
+                            state = MemberState.COMPLETING_REBALANCE;
+
+                            // we only need to enable heartbeat thread whenever we transit to
+                            // COMPLETING_REBALANCE state since we always transit from this state to STABLE
+                            if (heartbeatThread != null)
+                                heartbeatThread.enable();
+
+                            AbstractCoordinator.this.generation = new Generation(
+                                joinResponse.data().generationId(),
+                                joinResponse.data().memberId(), joinResponse.data().protocolName());
+
+                            log.info("Successfully joined group with generation {}", AbstractCoordinator.this.generation);
+
+                            if (joinResponse.isLeader()) {
+                                onJoinLeader(joinResponse).chain(future);
+                            } else {
+                                onJoinFollower().chain(future);
+                            }
                         }
                     }
                 }
             } else if (error == Errors.COORDINATOR_LOAD_IN_PROGRESS) {
-                log.debug("Attempt to join group rejected since coordinator {} is loading the group.", coordinator());
+                log.info("JoinGroup failed: Coordinator {} is loading the group.", coordinator());
                 // backoff and retry
                 future.raise(error);
             } else if (error == Errors.UNKNOWN_MEMBER_ID) {
-                // reset the member id and retry immediately
-                resetGeneration();
-                log.debug("Attempt to join group failed due to unknown member id.");
-                future.raise(Errors.UNKNOWN_MEMBER_ID);
+                log.info("JoinGroup failed: {} Need to re-join the group. Sent generation was {}",
+                         error.message(), sentGeneration);
+                // only need to reset the member id if generation has not been changed,
+                // then retry immediately
+                if (generationUnchanged())
+                    resetGenerationOnResponseError(ApiKeys.JOIN_GROUP, error);
+
+                future.raise(error);
             } else if (error == Errors.COORDINATOR_NOT_AVAILABLE
                     || error == Errors.NOT_COORDINATOR) {
                 // re-discover the coordinator and retry with backoff
-                coordinatorDead();
-                log.debug("Attempt to join group failed due to obsolete coordinator information: {}", error.message());
+                markCoordinatorUnknown();
+                log.info("JoinGroup failed: {} Marking coordinator unknown. Sent generation was {}",
+                          error.message(), sentGeneration);
+                future.raise(error);
+            } else if (error == Errors.FENCED_INSTANCE_ID) {
+                // for join-group request, even if the generation has changed we would not expect the instance id
+                // gets fenced, and hence we always treat this as a fatal error
+                log.error("JoinGroup failed: The group instance id {} has been fenced by another instance. " +
+                              "Sent generation was {}", rebalanceConfig.groupInstanceId, sentGeneration);
                 future.raise(error);
             } else if (error == Errors.INCONSISTENT_GROUP_PROTOCOL
                     || error == Errors.INVALID_SESSION_TIMEOUT
-                    || error == Errors.INVALID_GROUP_ID) {
+                    || 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());
+                log.error("JoinGroup failed due to fatal error: {}", error.message());
+                if (error == Errors.GROUP_MAX_SIZE_REACHED) {
+                    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(GroupAuthorizationException.forGroupId(rebalanceConfig.groupId));
+                } else {
+                    future.raise(error);
+                }
+            } else if (error == Errors.UNSUPPORTED_VERSION) {
+                log.error("JoinGroup 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.
+                String memberId = joinResponse.data().memberId();
+                log.debug("JoinGroup failed due to non-fatal error: {} Will set the member id as {} and then rejoin. " +
+                              "Sent generation was  {}", error, memberId, sentGeneration);
+                synchronized (AbstractCoordinator.this) {
+                    AbstractCoordinator.this.generation = new Generation(OffsetCommitRequest.DEFAULT_GENERATION_ID, memberId, null);
+                }
                 future.raise(error);
-            } else if (error == Errors.GROUP_AUTHORIZATION_FAILED) {
-                future.raise(new GroupAuthorizationException(groupId));
             } else {
                 // unexpected error, throw the exception
+                log.error("JoinGroup failed due to unexpected error: {}", error.message());
                 future.raise(new KafkaException("Unexpected error in join group response: " + error.message()));
             }
         }
@@ -508,21 +669,46 @@ 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)
+                                .setProtocolType(protocolType())
+                                .setProtocolName(generation.protocolName)
+                                .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)
+                                    .setProtocolType(protocolType())
+                                    .setProtocolName(generation.protocolName)
+                                    .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);
@@ -533,34 +719,87 @@ private RequestFuture sendSyncGroupRequest(SyncGroupRequest.Builder
         if (coordinatorUnknown())
             return RequestFuture.coordinatorNotAvailable();
         return client.send(coordinator, requestBuilder)
-                .compose(new SyncGroupResponseHandler());
+                .compose(new SyncGroupResponseHandler(generation));
     }
 
     private class SyncGroupResponseHandler extends CoordinatorResponseHandler {
+        private SyncGroupResponseHandler(final Generation generation) {
+            super(generation);
+        }
+
         @Override
         public void handle(SyncGroupResponse syncResponse,
                            RequestFuture future) {
             Errors error = syncResponse.error();
             if (error == Errors.NONE) {
-                sensors.syncLatency.record(response.requestLatencyMs());
-                future.complete(syncResponse.memberAssignment());
+                if (isProtocolTypeInconsistent(syncResponse.data().protocolType())) {
+                    log.error("SyncGroup failed due to inconsistent Protocol Type, received {} but expected {}",
+                        syncResponse.data().protocolType(), protocolType());
+                    future.raise(Errors.INCONSISTENT_GROUP_PROTOCOL);
+                } else {
+                    log.debug("Received successful SyncGroup response: {}", syncResponse);
+                    sensors.syncSensor.record(response.requestLatencyMs());
+
+                    synchronized (AbstractCoordinator.this) {
+                        if (!generation.equals(Generation.NO_GENERATION) && state == MemberState.COMPLETING_REBALANCE) {
+                            // check protocol name only if the generation is not reset
+                            final String protocolName = syncResponse.data().protocolName();
+                            final boolean protocolNameInconsistent = protocolName != null &&
+                                !protocolName.equals(generation.protocolName);
+
+                            if (protocolNameInconsistent) {
+                                log.error("SyncGroup failed due to inconsistent Protocol Name, received {} but expected {}",
+                                    protocolName, generation.protocolName);
+
+                                future.raise(Errors.INCONSISTENT_GROUP_PROTOCOL);
+                            } else {
+                                log.info("Successfully synced group in generation {}", generation);
+                                state = MemberState.STABLE;
+                                rejoinNeeded = false;
+                                // record rebalance latency
+                                lastRebalanceEndMs = time.milliseconds();
+                                sensors.successfulRebalanceSensor.record(lastRebalanceEndMs - lastRebalanceStartMs);
+                                lastRebalanceStartMs = -1L;
+
+                                future.complete(ByteBuffer.wrap(syncResponse.data().assignment()));
+                            }
+                        } else {
+                            log.info("Generation data was cleared by heartbeat thread to {} and state is now {} before " +
+                                "receiving SyncGroup response, marking this rebalance as failed and retry",
+                                generation, state);
+                            // use ILLEGAL_GENERATION error code to let it retry immediately
+                            future.raise(Errors.ILLEGAL_GENERATION);
+                        }
+                    }
+                }
             } 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 due to group rebalance");
+                    log.info("SyncGroup failed: The group began another rebalance. Need to re-join the group. " +
+                                 "Sent generation was {}", sentGeneration);
+                    future.raise(error);
+                } else if (error == Errors.FENCED_INSTANCE_ID) {
+                    // for sync-group request, even if the generation has changed we would not expect the instance id
+                    // gets fenced, and hence we always treat this as a fatal error
+                    log.error("SyncGroup failed: The group instance id {} has been fenced by another instance. " +
+                        "Sent generation was {}", rebalanceConfig.groupInstanceId, sentGeneration);
                     future.raise(error);
                 } else if (error == Errors.UNKNOWN_MEMBER_ID
                         || error == Errors.ILLEGAL_GENERATION) {
-                    log.debug("SyncGroup failed: {}", error.message());
-                    resetGeneration();
+                    log.info("SyncGroup failed: {} Need to re-join the group. Sent generation was {}",
+                            error.message(), sentGeneration);
+                    if (generationUnchanged())
+                        resetGenerationOnResponseError(ApiKeys.SYNC_GROUP, error);
+
                     future.raise(error);
                 } else if (error == Errors.COORDINATOR_NOT_AVAILABLE
                         || error == Errors.NOT_COORDINATOR) {
-                    log.debug("SyncGroup failed:", error.message());
-                    coordinatorDead();
+                    log.info("SyncGroup failed: {} Marking coordinator unknown. Sent generation was {}",
+                             error.message(), sentGeneration);
+                    markCoordinatorUnknown();
                     future.raise(error);
                 } else {
                     future.raise(new KafkaException("Unexpected error from SyncGroup: " + error.message()));
@@ -574,42 +813,46 @@ public void handle(SyncGroupResponse syncResponse,
      * one of the brokers. The returned future should be polled to get the result of the request.
      * @return A request future which indicates the completion of the metadata request
      */
-    private RequestFuture sendGroupCoordinatorRequest(Node node) {
+    private RequestFuture sendFindCoordinatorRequest(Node node) {
         // initiate the group metadata request
-        log.debug("Sending GroupCoordinator request to broker {}", node);
+        log.debug("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 GroupCoordinatorResponseHandler());
+                .compose(new FindCoordinatorResponseHandler());
     }
 
-    private class GroupCoordinatorResponseHandler extends RequestFutureAdapter {
+    private class FindCoordinatorResponseHandler extends RequestFutureAdapter {
 
         @Override
         public void onSuccess(ClientResponse resp, RequestFuture future) {
-            log.debug("Received GroupCoordinator response {}", resp);
+            log.debug("Received FindCoordinator response {}", resp);
+            clearFindCoordinatorFuture();
 
             FindCoordinatorResponse findCoordinatorResponse = (FindCoordinatorResponse) resp.responseBody();
-            // use MAX_VALUE - node.id as the coordinator id to mimic separate connections
-            // for the coordinator in the underlying network client layer
-            // TODO: this needs to be better handled in KAFKA-1935
             Errors error = findCoordinatorResponse.error();
-            clearFindCoordinatorFuture();
             if (error == Errors.NONE) {
                 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.data().nodeId();
+
                     AbstractCoordinator.this.coordinator = new Node(
-                            Integer.MAX_VALUE - findCoordinatorResponse.node().id(),
-                            findCoordinatorResponse.node().host(),
-                            findCoordinatorResponse.node().port());
-                    log.info("Discovered coordinator {}", coordinator);
+                            coordinatorConnectionId,
+                            findCoordinatorResponse.data().host(),
+                            findCoordinatorResponse.data().port());
+                    log.info("Discovered group coordinator {}", coordinator);
                     client.tryConnect(coordinator);
-                    heartbeat.resetTimeouts(time.milliseconds());
+                    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.debug("Group coordinator lookup failed: {}", findCoordinatorResponse.data().errorMessage());
                 future.raise(error);
             }
         }
@@ -626,84 +869,137 @@ public void onFailure(RuntimeException e, RequestFuture future) {
      * @return true if the coordinator is unknown
      */
     public boolean coordinatorUnknown() {
-        return coordinator() == null;
+        return checkAndGetCoordinator() == null;
     }
 
     /**
-     * Get the current coordinator
+     * Get the coordinator if its connection is still active. Otherwise mark it unknown and
+     * return null.
+     *
      * @return the current coordinator or null if it is unknown
      */
-    protected synchronized Node coordinator() {
-        if (coordinator != null && client.connectionFailed(coordinator)) {
-            coordinatorDead();
+    protected synchronized Node checkAndGetCoordinator() {
+        if (coordinator != null && client.isUnavailable(coordinator)) {
+            markCoordinatorUnknown(true);
             return null;
         }
         return this.coordinator;
     }
 
-    /**
-     * Mark the current coordinator as dead.
-     */
-    protected synchronized void coordinatorDead() {
+    private synchronized Node coordinator() {
+        return this.coordinator;
+    }
+
+    protected synchronized void markCoordinatorUnknown() {
+        markCoordinatorUnknown(false);
+    }
+
+    protected synchronized void markCoordinatorUnknown(boolean isDisconnected) {
         if (this.coordinator != null) {
-            log.info("Marking the coordinator {} dead", this.coordinator);
+            log.info("Group coordinator {} is unavailable or invalid, will attempt rediscovery", this.coordinator);
+            Node oldCoordinator = this.coordinator;
 
-            // Disconnect from the coordinator to ensure that there are no in-flight requests remaining.
-            // Pending callbacks will be invoked with a DisconnectException.
-            client.disconnect(this.coordinator);
+            // Mark the coordinator dead before disconnecting requests since the callbacks for any pending
+            // requests may attempt to do likewise. This also prevents new requests from being sent to the
+            // coordinator while the disconnect is in progress.
             this.coordinator = null;
+
+            // Disconnect from the coordinator to ensure that there are no in-flight requests remaining.
+            // Pending callbacks will be invoked with a DisconnectException on the next call to poll.
+            if (!isDisconnected)
+                client.disconnectAsync(oldCoordinator);
         }
     }
 
     /**
-     * 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;
     }
 
     /**
-     * Reset the generation and memberId because we have fallen out of the group.
+     * Get the current generation state if the group is stable, otherwise return null
+     *
+     * @return the current generation or null
      */
-    protected synchronized void resetGeneration() {
-        this.generation = Generation.NO_GENERATION;
-        this.rejoinNeeded = true;
-        this.state = MemberState.UNJOINED;
+    protected synchronized Generation generationIfStable() {
+        if (this.state != MemberState.STABLE)
+            return null;
+        return generation;
+    }
+
+    protected synchronized boolean rebalanceInProgress() {
+        return this.state == MemberState.PREPARING_REBALANCE || this.state == MemberState.COMPLETING_REBALANCE;
+    }
+
+    protected synchronized String memberId() {
+        return generation.memberId;
+    }
+
+    private synchronized void resetState() {
+        state = MemberState.UNJOINED;
+        generation = Generation.NO_GENERATION;
+    }
+
+    private synchronized void resetStateAndRejoin() {
+        resetState();
+        rejoinNeeded = true;
+    }
+
+    synchronized void resetGenerationOnResponseError(ApiKeys api, Errors error) {
+        log.debug("Resetting generation after encountering {} from {} response and requesting re-join", error, api);
+
+        resetStateAndRejoin();
     }
 
-    protected synchronized void requestRejoin() {
+    synchronized void resetGenerationOnLeaveGroup() {
+        log.debug("Resetting generation due to consumer pro-actively leaving the group");
+
+        resetStateAndRejoin();
+    }
+
+    public synchronized void requestRejoin() {
         this.rejoinNeeded = true;
     }
 
+    private boolean isProtocolTypeInconsistent(String protocolType) {
+        return protocolType != null && !protocolType.equals(protocolType());
+    }
+
     /**
      * Close the coordinator, waiting if needed to send LeaveGroup.
      */
     @Override
     public final void close() {
-        close(0);
+        close(time.timer(0));
     }
 
-    protected void close(long timeoutMs) {
+    /**
+     * @throws KafkaException if the rebalance callback throws exception
+     */
+    protected void close(Timer timer) {
         try {
             closeHeartbeatThread();
         } finally {
-
             // 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
                 // interrupted using wakeup) and the leave group request which have been queued, but not
                 // yet sent to the broker. Wait up to close timeout for these pending requests to be processed.
                 // If coordinator is not known, requests are aborted.
-                Node coordinator = coordinator();
-                if (coordinator != null && !client.awaitPendingRequests(coordinator, timeoutMs))
+                Node coordinator = checkAndGetCoordinator();
+                if (coordinator != null && !client.awaitPendingRequests(coordinator, timer))
                     log.warn("Close timed out with {} pending requests to coordinator, terminating client connections",
                             client.pendingRequestCount(coordinator));
             }
@@ -711,32 +1007,57 @@ protected void close(long timeoutMs) {
     }
 
     /**
-     * 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 != Generation.NO_GENERATION) {
+    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.debug("Sending LeaveGroup request to coordinator {}", coordinator);
-            LeaveGroupRequest.Builder request =
-                    new LeaveGroupRequest.Builder(groupId, 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(generation));
             client.pollNoWakeup();
         }
 
-        resetGeneration();
+        resetGenerationOnLeaveGroup();
+
+        return future;
+    }
+
+    protected boolean isDynamicMember() {
+        return !rebalanceConfig.groupInstanceId.isPresent();
     }
 
     private class LeaveGroupResponseHandler extends CoordinatorResponseHandler {
+        private LeaveGroupResponseHandler(final Generation generation) {
+            super(generation);
+        }
+
         @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");
+                log.debug("LeaveGroup response with {} returned successfully: {}", sentGeneration, response);
                 future.complete(null);
             } else {
-                log.debug("LeaveGroup request failed with error: {}", error.message());
+                log.error("LeaveGroup request with {} failed with error: {}", sentGeneration, error.message());
                 future.raise(error);
             }
         }
@@ -744,41 +1065,64 @@ public void handle(LeaveGroupResponse leaveResponse, RequestFuture future)
 
     // visible for testing
     synchronized RequestFuture sendHeartbeatRequest() {
-        log.debug("Sending Heartbeat request to coordinator {}", coordinator);
+        log.debug("Sending Heartbeat request with generation {} and member id {} to coordinator {}",
+            generation.generationId, generation.memberId, 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());
+                .compose(new HeartbeatResponseHandler(generation));
     }
 
     private class HeartbeatResponseHandler extends CoordinatorResponseHandler {
+        private HeartbeatResponseHandler(final Generation generation) {
+            super(generation);
+        }
+
         @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");
                 future.complete(null);
             } else if (error == Errors.COORDINATOR_NOT_AVAILABLE
                     || error == Errors.NOT_COORDINATOR) {
-                log.debug("Attempt to heartbeat since coordinator {} is either not started or not valid.",
+                log.info("Attempt to heartbeat failed since coordinator {} is either not started or not valid",
                         coordinator());
-                coordinatorDead();
+                markCoordinatorUnknown();
                 future.raise(error);
             } else if (error == Errors.REBALANCE_IN_PROGRESS) {
-                log.debug("Attempt to heartbeat failed since group is rebalancing");
-                requestRejoin();
-                future.raise(Errors.REBALANCE_IN_PROGRESS);
-            } else if (error == Errors.ILLEGAL_GENERATION) {
-                log.debug("Attempt to heartbeat failed since generation {} is not current", generation.generationId);
-                resetGeneration();
-                future.raise(Errors.ILLEGAL_GENERATION);
-            } else if (error == Errors.UNKNOWN_MEMBER_ID) {
-                log.debug("Attempt to heartbeat failed for since member id {} is not valid.", generation.memberId);
-                resetGeneration();
-                future.raise(Errors.UNKNOWN_MEMBER_ID);
+                // since we may be sending the request during rebalance, we should check
+                // this case and ignore the REBALANCE_IN_PROGRESS error
+                if (state == MemberState.STABLE) {
+                    log.info("Attempt to heartbeat failed since group is rebalancing");
+                    requestRejoin();
+                    future.raise(error);
+                } else {
+                    log.debug("Ignoring heartbeat response with error {} during {} state", error, state);
+                    future.complete(null);
+                }
+            } else if (error == Errors.ILLEGAL_GENERATION ||
+                       error == Errors.UNKNOWN_MEMBER_ID ||
+                       error == Errors.FENCED_INSTANCE_ID) {
+                if (generationUnchanged()) {
+                    log.info("Attempt to heartbeat with {} and group instance id {} failed due to {}, resetting generation",
+                        sentGeneration, rebalanceConfig.groupInstanceId, error);
+                    resetGenerationOnResponseError(ApiKeys.HEARTBEAT, error);
+                    future.raise(error);
+                } else {
+                    // if the generation has changed, then ignore this error
+                    log.info("Attempt to heartbeat with stale {} and group instance id {} failed due to {}, ignoring the error",
+                        sentGeneration, rebalanceConfig.groupInstanceId, error);
+                    future.complete(null);
+                }
             } 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()));
             }
@@ -786,7 +1130,12 @@ public void handle(HeartbeatResponse heartbeatResponse, RequestFuture futu
     }
 
     protected abstract class CoordinatorResponseHandler extends RequestFutureAdapter {
-        protected ClientResponse response;
+        CoordinatorResponseHandler(final Generation generation) {
+            this.sentGeneration = generation;
+        }
+
+        final Generation sentGeneration;
+        ClientResponse response;
 
         public abstract void handle(R response, RequestFuture future);
 
@@ -794,7 +1143,7 @@ protected abstract class CoordinatorResponseHandler extends RequestFutureA
         public void onFailure(RuntimeException e, RequestFuture future) {
             // mark the coordinator as dead
             if (e instanceof DisconnectException) {
-                coordinatorDead();
+                markCoordinatorUnknown(true);
             }
             future.raise(e);
         }
@@ -812,10 +1161,15 @@ public void onSuccess(ClientResponse clientResponse, RequestFuture future) {
             }
         }
 
+        boolean generationUnchanged() {
+            synchronized (AbstractCoordinator.this) {
+                return generation.equals(sentGeneration);
+            }
+        }
     }
 
     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,
@@ -825,65 +1179,123 @@ 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"));
-
-            Measurable lastHeartbeat =
-                new Measurable() {
-                    public double measure(MetricConfig config, long now) {
-                        return TimeUnit.SECONDS.convert(now - heartbeat.lastHeartbeatSend(), TimeUnit.MILLISECONDS);
-                    }
-                };
+                    "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 controller heartbeat was sent"),
+                "The number of seconds since the last coordinator heartbeat was sent"),
                 lastHeartbeat);
         }
     }
 
-    private class HeartbeatThread extends KafkaThread {
+    private class HeartbeatThread extends KafkaThread implements AutoCloseable {
         private boolean enabled = false;
         private boolean closed = false;
-        private AtomicReference failed = new AtomicReference<>(null);
+        private final 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() {
             synchronized (AbstractCoordinator.this) {
                 log.debug("Enabling heartbeat thread");
                 this.enabled = true;
-                heartbeat.resetTimeouts(time.milliseconds());
+                heartbeat.resetTimeouts();
                 AbstractCoordinator.this.notify();
             }
         }
@@ -924,9 +1336,10 @@ public void run() {
                             continue;
                         }
 
-                        if (state != MemberState.STABLE) {
-                            // the group is not stable (perhaps because we left the group or because the coordinator
-                            // kicked us out), so disable heartbeats and wait for the main thread to rejoin.
+                        // we do not need to heartbeat we are not part of a group yet;
+                        // also if we already have fatal error, the client will be
+                        // crashed soon, hence we do not need to continue heartbeating either
+                        if (state.hasNotJoinedGroup() || hasFailed()) {
                             disable();
                             continue;
                         }
@@ -938,27 +1351,32 @@ 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.
-                            coordinatorDead();
+                            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.
-                            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);
-
-                            sendHeartbeatRequest().addListener(new RequestFutureListener() {
+                            final RequestFuture heartbeatFuture = sendHeartbeatRequest();
+                            heartbeatFuture.addListener(new RequestFutureListener() {
                                 @Override
                                 public void onSuccess(Void value) {
                                     synchronized (AbstractCoordinator.this) {
-                                        heartbeat.receiveHeartbeat(time.milliseconds());
+                                        heartbeat.receiveHeartbeat();
                                     }
                                 }
 
@@ -970,10 +1388,12 @@ public void onFailure(RuntimeException e) {
                                             // ensures that the coordinator keeps the member in the group for as long
                                             // 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(time.milliseconds());
+                                            heartbeat.receiveHeartbeat();
+                                        } else if (e instanceof FencedInstanceIdException) {
+                                            log.error("Caught fenced group.instance.id {} error in heartbeat thread", rebalanceConfig.groupInstanceId);
+                                            heartbeatThread.failed.set(e);
                                         } else {
                                             heartbeat.failHeartbeat();
-
                                             // wake up the thread if it's sleeping to reschedule the heartbeat
                                             AbstractCoordinator.this.notify();
                                         }
@@ -987,15 +1407,18 @@ public void onFailure(RuntimeException e) {
                 log.error("An authentication error occurred in the heartbeat thread", e);
                 this.failed.set(e);
             } catch (GroupAuthorizationException e) {
-                log.error("A group authorization error occurred in the heartbeat thread for group {}", groupId, e);
+                log.error("A group authorization error occurred in the heartbeat thread", e);
                 this.failed.set(e);
             } catch (InterruptedException | InterruptException e) {
                 Thread.interrupted();
-                log.error("Unexpected interrupt received in heartbeat thread for group {}", groupId, e);
+                log.error("Unexpected interrupt received in heartbeat thread", e);
                 this.failed.set(new RuntimeException(e));
-            } catch (RuntimeException e) {
-                log.error("Heartbeat thread for group {} failed due to unexpected error", groupId, e);
-                this.failed.set(e);
+            } catch (Throwable e) {
+                log.error("Heartbeat thread failed due to unexpected error", e);
+                if (e instanceof RuntimeException)
+                    this.failed.set((RuntimeException) e);
+                else
+                    this.failed.set(new RuntimeException(e));
             } finally {
                 log.debug("Heartbeat thread has closed");
             }
@@ -1011,17 +1434,87 @@ protected static class Generation {
 
         public final int generationId;
         public final String memberId;
-        public final String protocol;
+        public final String protocolName;
 
-        public Generation(int generationId, String memberId, String protocol) {
+        public Generation(int generationId, String memberId, String protocolName) {
             this.generationId = generationId;
             this.memberId = memberId;
-            this.protocol = protocol;
+            this.protocolName = protocolName;
+        }
+
+        /**
+         * @return true if this generation has a valid member id, false otherwise. A member might have an id before
+         * it becomes part of a group generation.
+         */
+        public boolean hasMemberId() {
+            return !memberId.isEmpty();
+        }
+
+        @Override
+        public boolean equals(final Object o) {
+            if (this == o) return true;
+            if (o == null || getClass() != o.getClass()) return false;
+            final Generation that = (Generation) o;
+            return generationId == that.generationId &&
+                    Objects.equals(memberId, that.memberId) &&
+                    Objects.equals(protocolName, that.protocolName);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(generationId, memberId, protocolName);
+        }
+
+        @Override
+        public String toString() {
+            return "Generation{" +
+                    "generationId=" + generationId +
+                    ", memberId='" + memberId + '\'' +
+                    ", protocol='" + protocolName + '\'' +
+                    '}';
         }
     }
 
+    @SuppressWarnings("serial")
     private static class UnjoinedGroupException extends RetriableException {
 
     }
 
+    // For testing only below
+    final Heartbeat heartbeat() {
+        return heartbeat;
+    }
+
+    final synchronized 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.equals(Generation.NO_GENERATION) && generation.generationId == generationId;
+    }
+
+    final boolean hasUnknownGeneration() {
+        return generation.equals(Generation.NO_GENERATION);
+    }
+
+    /**
+     * @return true if the current generation's member ID is valid, false otherwise
+     */
+    final boolean hasValidMemberId() {
+        return !hasUnknownGeneration() && generation.hasMemberId();
+    }
+
+    final synchronized void setNewGeneration(final Generation generation) {
+        this.generation = generation;
+    }
+
+    final synchronized void setNewState(final MemberState state) {
+        this.state = state;
+    }
 }
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 8ec887ef3e340..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,20 +69,11 @@ 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) {
-        List list = map.get(key);
-        if (list == null) {
-            list = new ArrayList<>();
-            map.put(key, list);
-        }
+        List list = map.computeIfAbsent(key, k -> new ArrayList<>());
         list.add(value);
     }
 
@@ -94,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..7e42e44a7d4c8
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignor.java
@@ -0,0 +1,1006 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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) {
+        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;
+    }
+
+
+    /**
+     * This constrainedAssign optimizes the assignment algorithm when all consumers were subscribed to same set of topics.
+     * The method includes the following steps:
+     *
+     * 1. Reassign as many previously owned partitions as possible, up to the maxQuota
+     * 2. Fill remaining members up to minQuota
+     * 3. 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
+     * 4. 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
+     *
+     * @param partitionsPerTopic          The number of partitions for each subscribed topic
+     * @param consumerToOwnedPartitions   Each consumer's previously owned and still-subscribed partitions
+     *
+     * @return Map from each member to the list of partitions assigned to them.
+     */
+    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
+        // the consumers not yet at capacity
+        List unfilledMembers = new LinkedList<>();
+        // the members with exactly maxQuota partitions assigned
+        Queue maxCapacityMembers = new LinkedList<>();
+        // the members with exactly minQuota partitions assigned
+        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))));
+
+        // Reassign as many previously owned partitions as possible
+        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();
+
+        // Fill remaining members up to minQuota
+        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;
+    }
+
+    /**
+     * This generalAssign algorithm guarantees the assignment that is as balanced as possible.
+     * This method includes the following steps:
+     *
+     * 1. Preserving all the existing partition assignments
+     * 2. Removing all the partition assignments that have become invalid due to the change that triggers the reassignment
+     * 3. Assigning the unassigned partitions in a way that balances out the overall assignments of partitions to consumers
+     * 4. Further balancing out the resulting assignment by finding the partitions that can be reassigned
+     *    to another consumer towards an overall more balanced assignment.
+     *
+     * @param partitionsPerTopic         The number of partitions for each subscribed topic.
+     * @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.
+     */
+    private Map> generalAssign(Map partitionsPerTopic,
+                                                            Map subscriptions) {
+        Map> currentAssignment = new HashMap<>();
+        Map prevAssignment = new HashMap<>();
+        partitionMovements = new PartitionMovements();
+
+        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;
+        }
+    }
+}
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 1f5a41cb28068..80be7a9c858c3 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,37 +16,53 @@
  */
 package org.apache.kafka.clients.consumer.internals;
 
-import org.apache.kafka.clients.Metadata;
+import org.apache.kafka.clients.GroupRebalanceConfig;
 import org.apache.kafka.clients.consumer.CommitFailedException;
+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.RebalanceProtocol;
+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.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.UnstableOffsetCommitException;
+import org.apache.kafka.common.errors.RebalanceInProgressException;
 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.JoinGroupRequest;
 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.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.nio.ByteBuffer;
@@ -56,24 +72,29 @@
 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.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 Metadata metadata;
+    private final List assignors;
+    private final ConsumerMetadata metadata;
     private final ConsumerCoordinatorMetrics sensors;
     private final SubscriptionState subscriptions;
     private final OffsetCommitCallback defaultOffsetCommitCallback;
     private final boolean autoCommitEnabled;
     private final int autoCommitIntervalMs;
     private final ConsumerInterceptors interceptors;
-    private final boolean excludeInternalTopics;
     private final AtomicInteger pendingAsyncCommits;
 
     // this collection must be thread-safe because it is modified from the response handler
@@ -84,43 +105,60 @@ public final class ConsumerCoordinator extends AbstractCoordinator {
     private Set joinedSubscription;
     private MetadataSnapshot metadataSnapshot;
     private MetadataSnapshot assignmentSnapshot;
-    private long nextAutoCommitDeadline;
+    private Timer nextAutoCommitTimer;
+    private AtomicBoolean asyncCommitFenced;
+    private ConsumerGroupMetadata groupMetadata;
+    private final boolean throwOnFetchStableOffsetsUnsupported;
+
+    // hold onto request&future for committed offset requests to enable async calls.
+    private PendingCommittedOffsetRequest pendingCommittedOffsetRequest = null;
+
+    private static class PendingCommittedOffsetRequest {
+        private final Set requestedPartitions;
+        private final Generation requestedGeneration;
+        private final RequestFuture> response;
+
+        private PendingCommittedOffsetRequest(final Set requestedPartitions,
+                                              final Generation generationAtRequestTime,
+                                              final RequestFuture> response) {
+            this.requestedPartitions = Objects.requireNonNull(requestedPartitions);
+            this.response = Objects.requireNonNull(response);
+            this.requestedGeneration = generationAtRequestTime;
+        }
+
+        private boolean sameRequest(final Set currentRequest, final Generation currentGeneration) {
+            return Objects.equals(requestedGeneration, currentGeneration) && requestedPartitions.equals(currentRequest);
+        }
+    }
+
+    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,
-                               int heartbeatIntervalMs,
-                               List assignors,
-                               Metadata metadata,
+                               List assignors,
+                               ConsumerMetadata metadata,
                                SubscriptionState subscriptions,
                                Metrics metrics,
                                String metricGrpPrefix,
                                Time time,
-                               long retryBackoffMs,
                                boolean autoCommitEnabled,
                                int autoCommitIntervalMs,
                                ConsumerInterceptors interceptors,
-                               boolean excludeInternalTopics,
-                               final boolean leaveGroupOnClose) {
-        super(logContext,
+                               boolean throwOnFetchStableOffsetsUnsupported) {
+        super(rebalanceConfig,
+              logContext,
               client,
-              groupId,
-              rebalanceTimeoutMs,
-              sessionTimeoutMs,
-              heartbeatIntervalMs,
               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());
+        this.metadataSnapshot = new MetadataSnapshot(subscriptions, metadata.fetch(), metadata.updateVersion());
         this.subscriptions = subscriptions;
         this.defaultOffsetCommitCallback = new DefaultOffsetCommitCallback();
         this.autoCommitEnabled = autoCommitEnabled;
@@ -129,14 +167,42 @@ public ConsumerCoordinator(LogContext logContext,
         this.completedOffsetCommits = new ConcurrentLinkedQueue<>();
         this.sensors = new ConsumerCoordinatorMetrics(metrics, metricGrpPrefix);
         this.interceptors = interceptors;
-        this.excludeInternalTopics = excludeInternalTopics;
         this.pendingAsyncCommits = new AtomicInteger();
+        this.asyncCommitFenced = new AtomicBoolean(false);
+        this.groupMetadata = new ConsumerGroupMetadata(rebalanceConfig.groupId,
+            JoinGroupRequest.UNKNOWN_GENERATION_ID, JoinGroupRequest.UNKNOWN_MEMBER_ID, rebalanceConfig.groupInstanceId);
+        this.throwOnFetchStableOffsetsUnsupported = throwOnFetchStableOffsetsUnsupported;
 
         if (autoCommitEnabled)
-            this.nextAutoCommitDeadline = time.milliseconds() + autoCommitIntervalMs;
+            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();
-        addMetadataListener();
     }
 
     @Override
@@ -145,61 +211,133 @@ public String protocolType() {
     }
 
     @Override
-    public 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) {
-        final Set topicsToSubscribe = new HashSet<>();
+        final Set topicsToSubscribe = cluster.topics().stream()
+                .filter(subscriptions::matchesSubscribedPattern)
+                .collect(Collectors.toSet());
+        if (subscriptions.subscribeFromPattern(topicsToSubscribe))
+            metadata.requestUpdateForNewTopics();
+    }
 
-        for (String topic : cluster.topics())
-            if (subscriptions.subscribedPattern().matcher(topic).matches() &&
-                    !(excludeInternalTopics && cluster.internalTopics().contains(topic)))
-                topicsToSubscribe.add(topic);
+    private ConsumerPartitionAssignor lookupAssignor(String name) {
+        for (ConsumerPartitionAssignor assignor : this.assignors) {
+            if (assignor.name().equals(name))
+                return assignor;
+        }
+        return null;
+    }
+
+    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());
+            }
 
-        subscriptions.subscribeFromPattern(topicsToSubscribe);
+            if (!addedTopics.isEmpty()) {
+                Set newSubscription = new HashSet<>(subscriptions.subscription());
+                Set newJoinedSubscription = new HashSet<>(joinedSubscription);
+                newSubscription.addAll(addedTopics);
+                newJoinedSubscription.addAll(addedTopics);
 
-        // note we still need to update the topics contained in the metadata. Although we have
-        // specified that all topics should be fetched, only those set explicitly will be retained
-        metadata.setTopics(subscriptions.groupSubscription());
+                if (this.subscriptions.subscribeFromPattern(newSubscription))
+                    metadata.requestUpdateForNewTopics();
+                this.joinedSubscription = newJoinedSubscription;
+            }
+        }
     }
 
-    private void addMetadataListener() {
-        this.metadata.addListener(new Metadata.Listener() {
-            @Override
-            public void onMetadataUpdate(Cluster cluster, Set unavailableTopics) {
-                // if we encounter any unauthorized topics, raise an exception to the user
-                if (!cluster.unauthorizedTopics().isEmpty())
-                    throw new TopicAuthorizationException(new HashSet<>(cluster.unauthorizedTopics()));
-
-                if (subscriptions.hasPatternSubscription())
-                    updatePatternSubscription(cluster);
-
-                // check if there are any changes to the metadata which should trigger a rebalance
-                if (subscriptions.partitionsAutoAssigned()) {
-                    MetadataSnapshot snapshot = new MetadataSnapshot(subscriptions, cluster);
-                    if (!snapshot.equals(metadataSnapshot))
-                        metadataSnapshot = snapshot;
-                }
+    private Exception invokeOnAssignment(final ConsumerPartitionAssignor assignor, final Assignment assignment) {
+        log.info("Notifying assignor about the new {}", assignment);
 
-                if (!Collections.disjoint(metadata.topics(), unavailableTopics))
-                    metadata.requestUpdate();
-            }
-        });
+        try {
+            assignor.onAssignment(assignment, groupMetadata);
+        } catch (Exception e) {
+            return e;
+        }
+
+        return null;
     }
 
-    private PartitionAssignor lookupAssignor(String name) {
-        for (PartitionAssignor assignor : this.assignors) {
-            if (assignor.name().equals(name))
-                return assignor;
+    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;
+        }
+
+        return null;
+    }
+
+    private Exception invokePartitionsRevoked(final Set revokedPartitions) {
+        log.info("Revoke previously assigned partitions {}", Utils.join(revokedPartitions, ", "));
+
+        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;
+        }
+
+        return null;
+    }
+
+    private Exception invokePartitionsLost(final Set lostPartitions) {
+        log.info("Lost previously assigned partitions {}", Utils.join(lostPartitions, ", "));
+
+        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;
     }
 
@@ -208,64 +346,109 @@ protected void onJoinComplete(int generation,
                                   String memberId,
                                   String assignmentStrategy,
                                   ByteBuffer assignmentBuffer) {
-        // only the leader is responsible for monitoring for metadata changes (i.e. partition changes)
+        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);
 
+        // Give the assignor a chance to update internal state based on the received assignment
+        groupMetadata = new ConsumerGroupMetadata(rebalanceConfig.groupId, generation, memberId, rebalanceConfig.groupInstanceId);
+
+        Set ownedPartitions = new HashSet<>(subscriptions.assignedPartitions());
+
+        // should at least encode the short version
+        if (assignmentBuffer.remaining() < 2)
+            throw new IllegalStateException("There are insufficient bytes available to read assignment from the sync-group response (" +
+                "actual byte size " + assignmentBuffer.remaining() + ") , this is not expected; " +
+                "it is possible that the leader's assign function is buggy and did not return any assignment for this member, " +
+                "or because static member is configured and the protocol is buggy hence did not get the assignment for this member");
+
         Assignment assignment = ConsumerProtocol.deserializeAssignment(assignmentBuffer);
 
-        // set the flag to refresh last committed offsets
-        subscriptions.needRefreshCommits();
+        Set assignedPartitions = new HashSet<>(assignment.partitions());
 
-        // update partition assignment
-        subscriptions.assignFromSubscribed(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());
 
-        // 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
-        //
-        // TODO this part of the logic should be removed once we allow regex on leader assign
-        Set addedTopics = new HashSet<>();
-        for (TopicPartition tp : subscriptions.assignedPartitions()) {
-            if (!joinedSubscription.contains(tp.topic()))
-                addedTopics.add(tp.topic());
-        }
+            requestRejoin();
 
-        if (!addedTopics.isEmpty()) {
-            Set newSubscription = new HashSet<>(subscriptions.subscription());
-            Set newJoinedSubscription = new HashSet<>(joinedSubscription);
-            newSubscription.addAll(addedTopics);
-            newJoinedSubscription.addAll(addedTopics);
+            return;
+        }
 
-            this.subscriptions.subscribeFromPattern(newSubscription);
-            this.joinedSubscription = newJoinedSubscription;
+        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" +
+                    "\tAssigned partitions:                       {}\n" +
+                    "\tCurrent owned partitions:                  {}\n" +
+                    "\tAdded partitions (assigned - owned):       {}\n" +
+                    "\tRevoked partitions (owned - assigned):     {}\n",
+                assignedPartitions,
+                ownedPartitions,
+                addedPartitions,
+                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.info("Need to revoke partitions {} and re-join the group", revokedPartitions);
+                requestRejoin();
+            }
         }
 
-        // update the metadata and enforce a refresh to make sure the fetcher can start
-        // fetching data in the next iteration
-        this.metadata.setTopics(subscriptions.groupSubscription());
-        client.ensureFreshMetadata();
+        // 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);
+        // Catch any exception here to make sure we could complete the user callback.
+        firstException.compareAndSet(null, invokeOnAssignment(assignor, assignment));
 
-        // reschedule the auto commit starting from now
-        this.nextAutoCommitDeadline = time.milliseconds() + autoCommitIntervalMs;
+        // Reschedule the auto commit starting from now
+        if (autoCommitEnabled)
+            this.nextAutoCommitTimer.updateAndReset(autoCommitIntervalMs);
 
-        // execute the user's callback after rebalance
-        ConsumerRebalanceListener listener = subscriptions.listener();
-        log.info("Setting newly assigned partitions {}", subscriptions.assignedPartitions());
-        try {
-            Set assigned = new HashSet<>(subscriptions.assignedPartitions());
-            listener.onPartitionsAssigned(assigned);
-        } 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) {
+            if (firstException.get() instanceof KafkaException) {
+                throw (KafkaException) firstException.get();
+            } else {
+                throw new KafkaException("User rebalance callback throws an error", firstException.get());
+            }
+        }
+    }
+
+    void maybeUpdateSubscriptionMetadata() {
+        int version = metadata.updateVersion();
+        if (version > metadataSnapshot.version) {
+            Cluster cluster = metadata.fetch();
+
+            if (subscriptions.hasPatternSubscription())
+                updatePatternSubscription(cluster);
+
+            // Update the current snapshot, which will be used to check for subscription
+            // changes that would require a rebalance (e.g. new partitions).
+            metadataSnapshot = new MetadataSnapshot(subscriptions, cluster, version);
         }
     }
 
@@ -273,27 +456,62 @@ protected void onJoinComplete(int generation,
      * Poll for coordinator events. This ensures that the coordinator is known and that the consumer
      * has joined the group (if it is using group management). This also handles periodic offset commits
      * if they are enabled.
+     * 
    
+     * Returns early if the timeout expires or if waiting on rejoin is not required
      *
-     * @param now current time in milliseconds
+     * @param timer Timer bounding how long this method can block
+     * @param waitForJoinGroup Boolean flag indicating if we should wait until re-join group completes
+     * @throws KafkaException if the rebalance callback throws an exception
+     * @return true iff the operation succeeded
      */
-    public void poll(long now, long remainingMs) {
+    public boolean poll(Timer timer, boolean waitForJoinGroup) {
+        maybeUpdateSubscriptionMetadata();
+
         invokeCompletedOffsetCommitCallbacks();
 
-        if (subscriptions.partitionsAutoAssigned()) {
-            if (coordinatorUnknown()) {
-                ensureCoordinatorReady();
-                now = time.milliseconds();
+        if (subscriptions.hasAutoAssignedPartitions()) {
+            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());
+            if (coordinatorUnknown() && !ensureCoordinatorReady(timer)) {
+                return false;
             }
 
-            if (needRejoin()) {
+            if (rejoinNeededOrPending()) {
                 // due to a race condition between the initial metadata fetch and the initial rebalance,
                 // we need to ensure that the metadata is fresh before joining initially. This ensures
                 // that we have matched the pattern against the cluster's topics at least once before joining.
-                if (subscriptions.hasPatternSubscription())
-                    client.ensureFreshMetadata();
+                if (subscriptions.hasPatternSubscription()) {
+                    // For consumer group that uses pattern-based subscription, after a topic is created,
+                    // any consumer that discovers the topic after metadata refresh can trigger rebalance
+                    // across the entire consumer group. Multiple rebalances can be triggered after one topic
+                    // creation if consumers refresh metadata at vastly different times. We can significantly
+                    // reduce the number of rebalances caused by single topic creation by asking consumer to
+                    // refresh metadata before re-joining the group as long as the refresh backoff time has
+                    // passed.
+                    if (this.metadata.timeToAllowUpdate(timer.currentTimeMs()) == 0) {
+                        this.metadata.requestUpdate();
+                    }
+
+                    if (!client.ensureFreshMetadata(timer)) {
+                        return false;
+                    }
 
-                ensureActiveGroup();
-                now = time.milliseconds();
+                    maybeUpdateSubscriptionMetadata();
+                }
+
+                // if not wait for join group, we would just use a timer of 0
+                if (!ensureActiveGroup(waitForJoinGroup ? timer : time.timer(0L))) {
+                    // since we may use a different timer in the callee, we'd still need
+                    // to update the original timer's current time after the call
+                    timer.update(time.milliseconds());
+
+                    return false;
+                }
             }
         } else {
             // For manually assigned partitions, if there are no ready nodes, await metadata.
@@ -303,20 +521,17 @@ public void poll(long now, long remainingMs) {
             // awaitMetadataUpdate() initiates new connections with configured backoff and avoids the busy loop.
             // When group management is used, metadata wait is already performed for this scenario as
             // coordinator is unknown, hence this check is not required.
-            if (metadata.updateRequested() && !client.hasReadyNodes()) {
-                boolean metadataUpdated = client.awaitMetadataUpdate(remainingMs);
-                if (!metadataUpdated && !client.hasReadyNodes())
-                    return;
-                now = time.milliseconds();
+            if (metadata.updateRequested() && !client.hasReadyNodes(timer.currentTimeMs())) {
+                client.awaitMetadataUpdate(timer);
             }
         }
 
-        pollHeartbeat(now);
-        maybeAutoCommitOffsetsAsync(now);
+        maybeAutoCommitOffsetsAsync(timer.currentTimeMs());
+        return true;
     }
 
     /**
-     * Return the time to the next needed invocation of {@link #poll(long)}.
+     * Return the time to the next needed invocation of {@link ConsumerNetworkClient#poll(Timer)}.
      * @param now current time in milliseconds
      * @return the maximum time in milliseconds the caller should wait before the next invocation of poll()
      */
@@ -324,42 +539,58 @@ public long timeToNextPoll(long now) {
         if (!autoCommitEnabled)
             return timeToNextHeartbeat(now);
 
-        if (now > nextAutoCommitDeadline)
-            return 0;
+        return Math.min(nextAutoCommitTimer.remainingMs(), timeToNextHeartbeat(now));
+    }
 
-        return Math.min(nextAutoCommitDeadline - now, timeToNextHeartbeat(now));
+    private void updateGroupSubscription(Set topics) {
+        // the leader will begin watching for changes to any of the topics the group is interested in,
+        // which ensures that all metadata changes will eventually be seen
+        if (this.subscriptions.groupSubscribe(topics))
+            metadata.requestUpdateForNewTopics();
+
+        // update metadata (if needed) and keep track of the metadata used for assignment so that
+        // we can check after rebalance completion whether anything has changed
+        if (!client.ensureFreshMetadata(time.timer(Long.MAX_VALUE)))
+            throw new TimeoutException();
+
+        maybeUpdateSubscriptionMetadata();
     }
 
     @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,
         // which ensures that all metadata changes will eventually be seen
-        this.subscriptions.groupSubscribe(allSubscribedTopics);
-        metadata.setTopics(this.subscriptions.groupSubscription());
-
-        // update metadata (if needed) and keep track of the metadata used for assignment so that
-        // we can check after rebalance completion whether anything has changed
-        client.ensureFreshMetadata();
+        updateGroupSubscription(allSubscribedTopics);
 
         isLeader = true;
 
         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
@@ -367,9 +598,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());
         }
@@ -387,17 +618,15 @@ protected Map performAssignment(String leaderId,
                     "fetched from the brokers: {}", newlyAddedTopics);
 
             allSubscribedTopics.addAll(assignedTopics);
-            this.subscriptions.groupSubscribe(allSubscribedTopics);
-            metadata.setTopics(this.subscriptions.groupSubscription());
-            client.ensureFreshMetadata();
+            updateGroupSubscription(allSubscribedTopics);
         }
 
         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);
         }
@@ -405,106 +634,285 @@ 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(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.listener();
-        log.info("Revoking previously assigned partitions {}", subscriptions.assignedPartitions());
-        try {
-            Set revoked = new HashSet<>(subscriptions.assignedPartitions());
-            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.hasAutoAssignedPartitions() && !droppedPartitions.isEmpty()) {
+            final Exception e;
+            if (generation() == Generation.NO_GENERATION || rebalanceInProgress()) {
+                e = invokePartitionsLost(droppedPartitions);
+            } else {
+                e = invokePartitionsRevoked(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 needRejoin() {
-        if (!subscriptions.partitionsAutoAssigned())
+    public boolean rejoinNeededOrPending() {
+        if (!subscriptions.hasAutoAssignedPartitions())
             return false;
 
-        // we need to rejoin if we performed the assignment and metadata has changed
-        if (assignmentSnapshot != null && !assignmentSnapshot.equals(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)) {
+            log.info("Requesting to re-join the group and trigger rebalance since the assignment metadata has changed from {} to {}",
+                    assignmentSnapshot, 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())) {
+            log.info("Requesting to re-join the group and trigger rebalance since the subscription has changed from {} to {}",
+                joinedSubscription, subscriptions.subscription());
+
+            requestRejoin();
             return true;
+        }
 
-        return super.needRejoin();
+        return super.rejoinNeededOrPending();
     }
 
     /**
      * Refresh the committed offsets for provided partitions.
+     *
+     * @param timer Timer bounding how long this method can block
+     * @return true iff the operation completed within the timeout
      */
-    public void refreshCommittedOffsetsIfNeeded() {
-        if (subscriptions.refreshCommitsNeeded()) {
-            Map offsets = fetchCommittedOffsets(subscriptions.assignedPartitions());
-            for (Map.Entry entry : offsets.entrySet()) {
-                TopicPartition tp = entry.getKey();
-                // verify assignment is still active
-                if (subscriptions.isAssigned(tp))
-                    this.subscriptions.committed(tp, entry.getValue());
+    public boolean refreshCommittedOffsetsIfNeeded(Timer timer) {
+        final Set initializingPartitions = subscriptions.initializingPartitions();
+
+        final Map offsets = fetchCommittedOffsets(initializingPartitions, timer);
+        if (offsets == null) return false;
+
+        for (final Map.Entry entry : offsets.entrySet()) {
+            final TopicPartition tp = entry.getKey();
+            final OffsetAndMetadata offsetAndMetadata = entry.getValue();
+            if (offsetAndMetadata != null) {
+                // first update the epoch if necessary
+                entry.getValue().leaderEpoch().ifPresent(epoch -> this.metadata.updateLastSeenEpochIfNewer(entry.getKey(), epoch));
+
+                // it's possible that the partition is no longer assigned when the response is received,
+                // so we need to ignore seeking if that's the case
+                if (this.subscriptions.isAssigned(tp)) {
+                    final ConsumerMetadata.LeaderAndEpoch leaderAndEpoch = metadata.currentLeader(tp);
+                    final SubscriptionState.FetchPosition position = new SubscriptionState.FetchPosition(
+                            offsetAndMetadata.offset(), offsetAndMetadata.leaderEpoch(),
+                            leaderAndEpoch);
+
+                    this.subscriptions.seekUnvalidated(tp, position);
+
+                    log.info("Setting offset for partition {} to the committed offset {}", tp, position);
+                } else {
+                    log.info("Ignoring the returned {} since its partition {} is no longer assigned",
+                        offsetAndMetadata, tp);
+                }
             }
-            this.subscriptions.commitsRefreshed();
         }
+        return true;
     }
 
     /**
      * Fetch the current committed offsets from the coordinator for a set of partitions.
+     *
      * @param partitions The partitions to fetch offsets for
-     * @return A map from partition to the committed offset
+     * @return A map from partition to the committed offset or null if the operation timed out
      */
-    public Map fetchCommittedOffsets(Set partitions) {
-        while (true) {
-            ensureCoordinatorReady();
+    public Map fetchCommittedOffsets(final Set partitions,
+                                                                        final Timer timer) {
+        if (partitions.isEmpty()) return Collections.emptyMap();
+
+        final Generation generationForOffsetRequest = generationIfStable();
+        if (pendingCommittedOffsetRequest != null &&
+            !pendingCommittedOffsetRequest.sameRequest(partitions, generationForOffsetRequest)) {
+            // if we were waiting for a different request, then just clear it.
+            pendingCommittedOffsetRequest = null;
+        }
+
+        do {
+            if (!ensureCoordinatorReady(timer)) return null;
 
             // contact coordinator to fetch committed offsets
-            RequestFuture> future = sendOffsetFetchRequest(partitions);
-            client.poll(future);
+            final RequestFuture> future;
+            if (pendingCommittedOffsetRequest != null) {
+                future = pendingCommittedOffsetRequest.response;
+            } else {
+                future = sendOffsetFetchRequest(partitions);
+                pendingCommittedOffsetRequest = new PendingCommittedOffsetRequest(partitions, generationForOffsetRequest, future);
+            }
+            client.poll(future, timer);
 
-            if (future.succeeded())
-                return future.value();
+            if (future.isDone()) {
+                pendingCommittedOffsetRequest = null;
 
-            if (!future.isRetriable())
-                throw future.exception();
+                if (future.succeeded()) {
+                    return future.value();
+                } else if (!future.isRetriable()) {
+                    throw future.exception();
+                } else {
+                    timer.sleep(rebalanceConfig.retryBackoffMs);
+                }
+            } else {
+                return null;
+            }
+        } while (timer.notExpired());
+        return null;
+    }
 
-            time.sleep(retryBackoffMs);
-        }
+    /**
+     * Return the consumer group metadata.
+     *
+     * @return the current consumer group metadata
+     */
+    public ConsumerGroupMetadata groupMetadata() {
+        return groupMetadata;
     }
 
-    public void close(long timeoutMs) {
+    /**
+     * @throws KafkaException if the rebalance callback throws exception
+     */
+    public void close(final Timer timer) {
         // we do not need to re-enable wakeups since we are closing already
         client.disableWakeups();
-
-        long now = time.milliseconds();
-        long endTimeMs = now + timeoutMs;
         try {
-            maybeAutoCommitOffsetsSync(timeoutMs);
-            now = time.milliseconds();
-            if (pendingAsyncCommits.get() > 0 && endTimeMs > now) {
-                ensureCoordinatorReady(now, endTimeMs - now);
-                now = time.milliseconds();
+            maybeAutoCommitOffsetsSync(timer);
+            while (pendingAsyncCommits.get() > 0 && timer.notExpired()) {
+                ensureCoordinatorReady(timer);
+                client.poll(timer);
+                invokeCompletedOffsetCommitCallbacks();
             }
         } finally {
-            super.close(Math.max(0, endTimeMs - now));
+            super.close(timer);
         }
     }
 
     // visible for testing
     void invokeCompletedOffsetCommitCallbacks() {
+        if (asyncCommitFenced.get()) {
+            throw new FencedInstanceIdException("Get fenced exception for group.instance.id "
+                + rebalanceConfig.groupInstanceId.orElse("unset_instance_id")
+                + ", current member.id is " + memberId());
+        }
         while (true) {
             OffsetCommitCompletion completion = completedOffsetCommits.poll();
-            if (completion == null)
+            if (completion == null) {
                 break;
+            }
             completion.invoke();
         }
     }
@@ -527,13 +935,14 @@ public void commitOffsetsAsync(final Map offs
                 public void onSuccess(Void value) {
                     pendingAsyncCommits.decrementAndGet();
                     doCommitOffsetsAsync(offsets, callback);
+                    client.pollNoWakeup();
                 }
 
                 @Override
                 public void onFailure(RuntimeException e) {
                     pendingAsyncCommits.decrementAndGet();
                     completedOffsetCommits.add(new OffsetCommitCompletion(callback, offsets,
-                            RetriableCommitFailedException.withUnderlyingMessage(e.getMessage())));
+                            new RetriableCommitFailedException(e)));
                 }
             });
         }
@@ -545,7 +954,6 @@ public void onFailure(RuntimeException e) {
     }
 
     private void doCommitOffsetsAsync(final Map offsets, final OffsetCommitCallback callback) {
-        this.subscriptions.needRefreshCommits();
         RequestFuture future = sendOffsetCommitRequest(offsets);
         final OffsetCommitCallback cb = callback == null ? defaultOffsetCommitCallback : callback;
         future.addListener(new RequestFutureListener() {
@@ -553,7 +961,6 @@ private void doCommitOffsetsAsync(final Map o
             public void onSuccess(Void value) {
                 if (interceptors != null)
                     interceptors.onCommit(offsets);
-
                 completedOffsetCommits.add(new OffsetCommitCompletion(cb, offsets, null));
             }
 
@@ -561,10 +968,13 @@ public void onSuccess(Void value) {
             public void onFailure(RuntimeException e) {
                 Exception commitException = e;
 
-                if (e instanceof RetriableException)
-                    commitException = RetriableCommitFailedException.withUnderlyingMessage(e.getMessage());
-
+                if (e instanceof RetriableException) {
+                    commitException = new RetriableCommitFailedException(e);
+                }
                 completedOffsetCommits.add(new OffsetCommitCompletion(cb, offsets, commitException));
+                if (commitException instanceof FencedInstanceIdException) {
+                    asyncCommitFenced.set(true);
+                }
             }
         });
     }
@@ -576,28 +986,23 @@ public void onFailure(RuntimeException e) {
      * @throws org.apache.kafka.common.errors.AuthorizationException if the consumer is not authorized to the group
      *             or to any of the specified partitions. See the exception for more details
      * @throws CommitFailedException if an unrecoverable error occurs before the commit can be completed
+     * @throws FencedInstanceIdException if a static member gets fenced
      * @return If the offset commit was successfully sent and a successful response was received from
      *         the coordinator
      */
-    public boolean commitOffsetsSync(Map offsets, long timeoutMs) {
+    public boolean commitOffsetsSync(Map offsets, Timer timer) {
         invokeCompletedOffsetCommitCallbacks();
 
         if (offsets.isEmpty())
             return true;
 
-        long now = time.milliseconds();
-        long startMs = now;
-        long remainingMs = timeoutMs;
         do {
-            if (coordinatorUnknown()) {
-                if (!ensureCoordinatorReady(now, remainingMs))
-                    return false;
-
-                remainingMs = timeoutMs - (time.milliseconds() - startMs);
+            if (coordinatorUnknown() && !ensureCoordinatorReady(timer)) {
+                return false;
             }
 
             RequestFuture future = sendOffsetCommitRequest(offsets);
-            client.poll(future, remainingMs);
+            client.poll(future, timer);
 
             // We may have had in-flight offset commits when the synchronous commit began. If so, ensure that
             // the corresponding callbacks are invoked prior to returning in order to preserve the order that
@@ -613,55 +1018,47 @@ public boolean commitOffsetsSync(Map offsets,
             if (future.failed() && !future.isRetriable())
                 throw future.exception();
 
-            time.sleep(retryBackoffMs);
-
-            now = time.milliseconds();
-            remainingMs = timeoutMs - (now - startMs);
-        } while (remainingMs > 0);
+            timer.sleep(rebalanceConfig.retryBackoffMs);
+        } while (timer.notExpired());
 
         return false;
     }
 
-    private void maybeAutoCommitOffsetsAsync(long now) {
+    public void maybeAutoCommitOffsetsAsync(long now) {
         if (autoCommitEnabled) {
-            if (coordinatorUnknown()) {
-                this.nextAutoCommitDeadline = now + retryBackoffMs;
-            } else if (now >= nextAutoCommitDeadline) {
-                this.nextAutoCommitDeadline = now + autoCommitIntervalMs;
+            nextAutoCommitTimer.update(now);
+            if (nextAutoCommitTimer.isExpired()) {
+                nextAutoCommitTimer.reset(autoCommitIntervalMs);
                 doAutoCommitOffsetsAsync();
             }
         }
     }
 
-    public void maybeAutoCommitOffsetsNow() {
-        if (autoCommitEnabled && !coordinatorUnknown())
-            doAutoCommitOffsetsAsync();
-    }
-
     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) {
-                    log.warn("Asynchronous auto-commit of offsets {} failed: {}", offsets, exception.getMessage());
-                    if (exception instanceof RetriableException)
-                        nextAutoCommitDeadline = Math.min(time.milliseconds() + retryBackoffMs, nextAutoCommitDeadline);
+        commitOffsetsAsync(allConsumedOffsets, (offsets, exception) -> {
+            if (exception != null) {
+                if (exception instanceof RetriableCommitFailedException) {
+                    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);
             }
         });
     }
 
-    private void maybeAutoCommitOffsetsSync(long timeoutMs) {
+    private void maybeAutoCommitOffsetsSync(Timer timer) {
         if (autoCommitEnabled) {
             Map allConsumedOffsets = subscriptions.allConsumed();
             try {
                 log.debug("Sending synchronous auto-commit of offsets {}", allConsumedOffsets);
-                if (!commitOffsetsSync(allConsumedOffsets, timeoutMs))
+                if (!commitOffsetsSync(allConsumedOffsets, timer))
                     log.debug("Auto-commit of offsets {} timed out before completion", allConsumedOffsets);
             } catch (WakeupException | InterruptException e) {
                 log.debug("Auto-commit of offsets {} was interrupted before completion", allConsumedOffsets);
@@ -687,110 +1084,182 @@ public void onComplete(Map offsets, Exception
      * which returns a request future that can be polled in the case of a synchronous commit or ignored in the
      * asynchronous case.
      *
+     * NOTE: This is visible only for testing
+     *
      * @param offsets The list of offsets per partition that should be committed.
      * @return A request future whose value indicates whether the commit was successful or not
      */
-    private RequestFuture sendOffsetCommitRequest(final Map offsets) {
+    RequestFuture sendOffsetCommitRequest(final Map offsets) {
         if (offsets.isEmpty())
             return RequestFuture.voidSuccess();
 
-        Node coordinator = coordinator();
+        Node coordinator = checkAndGetCoordinator();
         if (coordinator == null)
             return RequestFuture.coordinatorNotAvailable();
 
         // create the offset commit request
-        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.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.hasAutoAssignedPartitions()) {
+            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");
+
+                if (rebalanceInProgress()) {
+                    // if the client knows it is already rebalancing, we can use RebalanceInProgressException instead of
+                    // CommitFailedException to indicate this is not a fatal error
+                    return RequestFuture.failure(new RebalanceInProgressException("Offset commit cannot be completed since the " +
+                        "consumer is undergoing a rebalance for auto partition assignment. You can try completing the rebalance " +
+                        "by calling poll() and then retry the operation."));
+                } else {
+                    return RequestFuture.failure(new CommitFailedException("Offset commit cannot be completed since the " +
+                        "consumer is not part of an active group for auto partition assignment; it is likely that the consumer " +
+                        "was kicked out of the group."));
+                }
+            }
+        } 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).
-                setRetentionTime(OffsetCommitRequest.DEFAULT_RETENTION_TIME);
+        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);
 
         return client.send(coordinator, builder)
-                .compose(new OffsetCommitResponseHandler(offsets));
+                .compose(new OffsetCommitResponseHandler(offsets, generation));
     }
 
     private class OffsetCommitResponseHandler extends CoordinatorResponseHandler {
-
         private final Map offsets;
 
-        private OffsetCommitResponseHandler(Map offsets) {
+        private OffsetCommitResponseHandler(Map offsets, Generation generation) {
+            super(generation);
             this.offsets = offsets;
         }
 
         @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();
-
-                Errors error = entry.getValue();
-                if (error == Errors.NONE) {
-                    log.debug("Committed offset {} for partition {}", offset, tp);
-                    if (subscriptions.isAssigned(tp))
-                        // update the local cache only if the partition is still assigned
-                        subscriptions.committed(tp, offsetAndMetadata);
-                } else {
-                    log.error("Offset commit failed on partition {} at offset {}: {}", tp, offset, error.message());
+            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);
 
-                    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) {
-                        // just retry
-                        future.raise(error);
-                        return;
-                    } else if (error == Errors.COORDINATOR_NOT_AVAILABLE
-                            || error == Errors.NOT_COORDINATOR
-                            || error == Errors.REQUEST_TIMED_OUT) {
-                        coordinatorDead();
-                        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;
-                    } else if (error == Errors.UNKNOWN_TOPIC_OR_PARTITION) {
-                        future.raise(new KafkaException("Topic or Partition " + tp + " does not exist"));
-                        return;
+                    long offset = offsetAndMetadata.offset();
+
+                    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.info("OffsetCommit failed with {} due to group instance id {} fenced", sentGeneration, rebalanceConfig.groupInstanceId);
+
+                            // if the generation has changed or we are not in rebalancing, do not raise the fatal error but rebalance-in-progress
+                            if (generationUnchanged()) {
+                                future.raise(error);
+                            } else {
+                                if (ConsumerCoordinator.this.state == MemberState.PREPARING_REBALANCE) {
+                                    future.raise(new RebalanceInProgressException("Offset commit cannot be completed since the " +
+                                        "consumer member's old generation is fenced by its group instance id, it is possible that " +
+                                        "this consumer has already participated another rebalance and got a new generation"));
+                                } else {
+                                    future.raise(new CommitFailedException());
+                                }
+                            }
+                            return;
+                        } else if (error == Errors.REBALANCE_IN_PROGRESS) {
+                            /* Consumer should not try 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 to indicate that we are still in the middle of a rebalance.
+                             * In this case we would throw a RebalanceInProgressException,
+                             * request re-join but do not reset generations. If the callers decide to retry they
+                             * can go ahead and call poll to finish up the rebalance first, and then try commit again.
+                             */
+                            requestRejoin();
+                            future.raise(new RebalanceInProgressException("Offset commit cannot be completed since the " +
+                                "consumer group is executing a rebalance at the moment. You can try completing the rebalance " +
+                                "by calling poll() and then retry commit again"));
+                            return;
+                        } else if (error == Errors.UNKNOWN_MEMBER_ID
+                                || error == Errors.ILLEGAL_GENERATION) {
+                            log.info("OffsetCommit failed with {}: {}", sentGeneration, error.message());
+
+                            // only need to reset generation and re-join group if generation has not changed or we are not in rebalancing;
+                            // otherwise only raise rebalance-in-progress error
+                            if (!generationUnchanged() && ConsumerCoordinator.this.state == MemberState.PREPARING_REBALANCE) {
+                                future.raise(new RebalanceInProgressException("Offset commit cannot be completed since the " +
+                                    "consumer member's generation is already stale, meaning it has already participated another rebalance and " +
+                                    "got a new generation. You can try completing the rebalance by calling poll() and then retry commit again"));
+                            } else {
+                                resetGenerationOnResponseError(ApiKeys.OFFSET_COMMIT, error);
+                                future.raise(new CommitFailedException());
+                            }
+                            return;
+                        } else {
+                            future.raise(new KafkaException("Unexpected error in commit: " + error.message()));
+                            return;
+                        }
                     }
                 }
             }
@@ -812,14 +1281,14 @@ public void handle(OffsetCommitResponse commitResponse, RequestFuture futu
      * @return A request future containing the committed offsets.
      */
     private RequestFuture> sendOffsetFetchRequest(Set partitions) {
-        Node coordinator = coordinator();
+        Node coordinator = checkAndGetCoordinator();
         if (coordinator == null)
             return RequestFuture.coordinatorNotAvailable();
 
         log.debug("Fetching committed offsets for partitions: {}", partitions);
         // construct the request
-        OffsetFetchRequest.Builder requestBuilder = new OffsetFetchRequest.Builder(this.groupId,
-                new ArrayList<>(partitions));
+        OffsetFetchRequest.Builder requestBuilder =
+            new OffsetFetchRequest.Builder(this.rebalanceConfig.groupId, true, new ArrayList<>(partitions), throwOnFetchStableOffsetsUnsupported);
 
         // send the request with a callback
         return client.send(coordinator, requestBuilder)
@@ -827,6 +1296,10 @@ private RequestFuture> sendOffsetFetchReq
     }
 
     private class OffsetFetchResponseHandler extends CoordinatorResponseHandler> {
+        private OffsetFetchResponseHandler() {
+            super(Generation.NO_GENERATION);
+        }
+
         @Override
         public void handle(OffsetFetchResponse response, RequestFuture> future) {
             if (response.hasError()) {
@@ -838,64 +1311,111 @@ public void handle(OffsetFetchResponse response, RequestFuture unauthorizedTopics = null;
             Map offsets = new HashMap<>(response.responseData().size());
+            Set unstableTxnOffsetTopicPartitions = new HashSet<>();
             for (Map.Entry entry : response.responseData().entrySet()) {
                 TopicPartition tp = entry.getKey();
-                OffsetFetchResponse.PartitionData data = entry.getValue();
-                if (data.hasError()) {
-                    Errors error = data.error;
+                OffsetFetchResponse.PartitionData partitionData = entry.getValue();
+                if (partitionData.hasError()) {
+                    Errors error = partitionData.error;
                     log.debug("Failed to fetch offset for partition {}: {}", tp, error.message());
 
                     if (error == Errors.UNKNOWN_TOPIC_OR_PARTITION) {
                         future.raise(new KafkaException("Topic or Partition " + tp + " does not exist"));
+                        return;
+                    } else if (error == Errors.TOPIC_AUTHORIZATION_FAILED) {
+                        if (unauthorizedTopics == null) {
+                            unauthorizedTopics = new HashSet<>();
+                        }
+                        unauthorizedTopics.add(tp.topic());
+                    } else if (error == Errors.UNSTABLE_OFFSET_COMMIT) {
+                        unstableTxnOffsetTopicPartitions.add(tp);
                     } else {
-                        future.raise(new KafkaException("Unexpected error in fetch offset response: " + error.message()));
+                        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)
-                    offsets.put(tp, new OffsetAndMetadata(data.offset, data.metadata));
+                } else if (partitionData.offset >= 0) {
+                    // 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(partitionData.offset, partitionData.leaderEpoch, partitionData.metadata));
                 } else {
-                    log.debug("Found no committed offset for partition {}", tp);
+                    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 if (!unstableTxnOffsetTopicPartitions.isEmpty()) {
+                // just retry
+                log.info("The following partitions still have unstable offsets " +
+                             "which are not cleared on the broker side: {}" +
+                             ", this could be either " +
+                             "transactional offsets waiting for completion, or " +
+                             "normal offsets waiting for replication after appending to local log", unstableTxnOffsetTopicPartitions);
+                future.raise(new UnstableOffsetCommitException("There are unstable offsets for the requested topic partitions"));
+            } 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 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"));
 
-            Measurable numParts =
-                new Measurable() {
-                    public double measure(MetricConfig config, long now) {
-                        return subscriptions.assignedPartitions().size();
-                    }
-                };
+            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());
+
+            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);
@@ -903,26 +1423,27 @@ public double measure(MetricConfig config, long now) {
     }
 
     private static class MetadataSnapshot {
+        private final int version;
         private final Map partitionsPerTopic;
 
-        private MetadataSnapshot(SubscriptionState subscription, Cluster cluster) {
+        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;
         }
 
-        @Override
-        public boolean equals(Object o) {
-            if (this == o) return true;
-            if (o == null || getClass() != o.getClass()) return false;
-            MetadataSnapshot that = (MetadataSnapshot) o;
-            return partitionsPerTopic != null ? partitionsPerTopic.equals(that.partitionsPerTopic) : that.partitionsPerTopic == null;
+        boolean matches(MetadataSnapshot other) {
+            return version == other.version || partitionsPerTopic.equals(other.partitionsPerTopic);
         }
 
         @Override
-        public int hashCode() {
-            return partitionsPerTopic != null ? partitionsPerTopic.hashCode() : 0;
+        public String toString() {
+            return "(version" + version + ": " + partitionsPerTopic + ")";
         }
     }
 
@@ -943,4 +1464,12 @@ public void invoke() {
         }
     }
 
+    /* test-only classes below */
+    RebalanceProtocol getProtocol() {
+        return protocol;
+    }
+
+    boolean poll(Timer timer) {
+        return poll(timer, true);
+    }
 }
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
new file mode 100644
index 0000000000000..ef7d92417471b
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerMetadata.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.clients.consumer.internals;
+
+import org.apache.kafka.clients.Metadata;
+import org.apache.kafka.common.internals.ClusterResourceListeners;
+import org.apache.kafka.common.requests.MetadataRequest;
+import org.apache.kafka.common.utils.LogContext;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+
+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.metadataTopics());
+        topics.addAll(transientTopics);
+        return new MetadataRequest.Builder(topics, allowAutoTopicCreation);
+    }
+
+    synchronized void addTransientTopics(Set topics) {
+        this.transientTopics.addAll(topics);
+        if (!fetch().topics().containsAll(topics))
+            requestUpdateForNewTopics();
+    }
+
+    synchronized void clearTransientTopics() {
+        this.transientTopics.clear();
+    }
+
+    @Override
+    protected synchronized boolean retainTopic(String topic, boolean isInternal, long nowMs) {
+        if (transientTopics.contains(topic) || subscription.needsMetadata(topic))
+            return true;
+
+        if (isInternal && !includeInternalTopics)
+            return false;
+
+        return subscription.matchesSubscribedPattern(topic);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerMetrics.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerMetrics.java
index 3492323170841..e58db82ee3cd6 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerMetrics.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerMetrics.java
@@ -37,8 +37,7 @@ public ConsumerMetrics(String metricGroupPrefix) {
     }
 
     private List getAllTemplates() {
-        List l = new ArrayList<>();
-        l.addAll(this.fetcherMetrics.getAllTemplates());
+        List l = new ArrayList<>(this.fetcherMetrics.getAllTemplates());
         return l;
     }
 
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 4b3ef68a70aa8..c591074ba9573 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
@@ -28,9 +28,9 @@
 import org.apache.kafka.common.errors.TimeoutException;
 import org.apache.kafka.common.errors.WakeupException;
 import org.apache.kafka.common.requests.AbstractRequest;
-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.Timer;
 import org.slf4j.Logger;
 
 import java.io.Closeable;
@@ -44,6 +44,7 @@
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.locks.ReentrantLock;
 
 /**
  * Higher level consumer access to the network layer with basic support for request futures. This class
@@ -62,13 +63,18 @@ public class ConsumerNetworkClient implements Closeable {
     private final Time time;
     private final long retryBackoffMs;
     private final int maxPollTimeoutMs;
-    private final long unsentExpiryMs;
+    private final int requestTimeoutMs;
     private final AtomicBoolean wakeupDisabled = new AtomicBoolean();
 
+    // We do not need high throughput, so use a fair lock to try to avoid starvation
+    private final ReentrantLock lock = new ReentrantLock(true);
+
     // when requests complete, they are transferred to this queue prior to invocation. The purpose
     // is to avoid invoking them while holding this object's monitor which can open the door for deadlocks.
     private final ConcurrentLinkedQueue pendingCompletion = new ConcurrentLinkedQueue<>();
 
+    private final ConcurrentLinkedQueue pendingDisconnects = new ConcurrentLinkedQueue<>();
+
     // this flag allows the client to be safely woken up without waiting on the lock above. It is
     // atomic to avoid the need to acquire the lock above in order to enable it concurrently.
     private final AtomicBoolean wakeup = new AtomicBoolean(false);
@@ -78,7 +84,7 @@ public ConsumerNetworkClient(LogContext logContext,
                                  Metadata metadata,
                                  Time time,
                                  long retryBackoffMs,
-                                 long requestTimeoutMs,
+                                 int requestTimeoutMs,
                                  int maxPollTimeoutMs) {
         this.log = logContext.logger(ConsumerNetworkClient.class);
         this.client = client;
@@ -86,12 +92,23 @@ public ConsumerNetworkClient(LogContext logContext,
         this.time = time;
         this.retryBackoffMs = retryBackoffMs;
         this.maxPollTimeoutMs = Math.min(maxPollTimeoutMs, MAX_POLL_TIMEOUT_MS);
-        this.unsentExpiryMs = requestTimeoutMs;
+        this.requestTimeoutMs = requestTimeoutMs;
+    }
+
+    public int defaultRequestTimeoutMs() {
+        return requestTimeoutMs;
+    }
+
+    /**
+     * Send a request with the default timeout. See {@link #send(Node, AbstractRequest.Builder, int)}.
+     */
+    public RequestFuture send(Node node, AbstractRequest.Builder requestBuilder) {
+        return send(node, requestBuilder, requestTimeoutMs);
     }
 
     /**
      * Send a new request. Note that the request is not actually transmitted on the
-     * network until one of the {@link #poll(long)} variants is invoked. At this
+     * network until one of the {@link #poll(Timer)} variants is invoked. At this
      * point the request will either be transmitted successfully or will fail.
      * Use the returned future to obtain the result of the send. Note that there is no
      * need to check for disconnects explicitly on the {@link ClientResponse} object;
@@ -99,13 +116,18 @@ public ConsumerNetworkClient(LogContext logContext,
      *
      * @param node The destination of the request
      * @param requestBuilder A builder for the request payload
+     * @param requestTimeoutMs Maximum time in milliseconds to await a response before disconnecting the socket and
+     *                         cancelling the request. The request may be cancelled sooner if the socket disconnects
+     *                         for any reason.
      * @return A future which indicates the result of the send.
      */
-    public RequestFuture send(Node node, AbstractRequest.Builder requestBuilder) {
+    public RequestFuture send(Node node,
+                                              AbstractRequest.Builder requestBuilder,
+                                              int requestTimeoutMs) {
         long now = time.milliseconds();
         RequestFutureCompletionHandler completionHandler = new RequestFutureCompletionHandler();
         ClientRequest clientRequest = client.newClientRequest(node.idString(), requestBuilder, now, true,
-                completionHandler);
+            requestTimeoutMs, completionHandler);
         unsent.put(node, clientRequest);
 
         // wakeup the client in case it is blocking in poll so that we can send the queued request
@@ -113,19 +135,22 @@ public RequestFuture send(Node node, AbstractRequest.Builder
         return completionHandler.future;
     }
 
-    public synchronized Node leastLoadedNode() {
-        return client.leastLoadedNode(time.milliseconds());
-    }
-
-    public synchronized boolean hasReadyNodes() {
-        return client.hasReadyNodes();
+    public Node leastLoadedNode() {
+        lock.lock();
+        try {
+            return client.leastLoadedNode(time.milliseconds());
+        } finally {
+            lock.unlock();
+        }
     }
 
-    /**
-     * Block until the metadata has been refreshed.
-     */
-    public void awaitMetadataUpdate() {
-        awaitMetadataUpdate(Long.MAX_VALUE);
+    public boolean hasReadyNodes(long now) {
+        lock.lock();
+        try {
+            return client.hasReadyNodes(now);
+        } finally {
+            lock.unlock();
+        }
     }
 
     /**
@@ -133,25 +158,25 @@ public void awaitMetadataUpdate() {
      *
      * @return true if update succeeded, false otherwise.
      */
-    public boolean awaitMetadataUpdate(long timeout) {
-        long startMs = time.milliseconds();
+    public boolean awaitMetadataUpdate(Timer timer) {
         int version = this.metadata.requestUpdate();
         do {
-            poll(timeout);
-            AuthenticationException ex = this.metadata.getAndClearAuthenticationException();
-            if (ex != null)
-                throw ex;
-        } while (this.metadata.version() == version && time.milliseconds() - startMs < timeout);
-        return this.metadata.version() > version;
+            poll(timer);
+        } while (this.metadata.updateVersion() == version && timer.notExpired());
+        return this.metadata.updateVersion() > version;
     }
 
     /**
      * Ensure our metadata is fresh (if an update is expected, this will block
      * until it has completed).
      */
-    public void ensureFreshMetadata() {
-        if (this.metadata.updateRequested() || this.metadata.timeToNextUpdate(time.milliseconds()) == 0)
-            awaitMetadataUpdate();
+    boolean ensureFreshMetadata(Timer timer) {
+        if (this.metadata.updateRequested() || this.metadata.timeToNextUpdate(timer.currentTimeMs()) == 0) {
+            return awaitMetadataUpdate(timer);
+        } else {
+            // the metadata is already fresh
+            return true;
+        }
     }
 
     /**
@@ -160,7 +185,7 @@ public void ensureFreshMetadata() {
      */
     public void wakeup() {
         // wakeup should be safe without holding the client lock since it simply delegates to
-        // Selector's wakeup, which is threadsafe
+        // Selector's wakeup, which is thread-safe
         log.debug("Received user wakeup");
         this.wakeup.set(true);
         this.client.wakeup();
@@ -174,80 +199,79 @@ public void wakeup() {
      */
     public void poll(RequestFuture future) {
         while (!future.isDone())
-            poll(Long.MAX_VALUE, time.milliseconds(), future);
+            poll(time.timer(Long.MAX_VALUE), future);
     }
 
     /**
      * Block until the provided request future request has finished or the timeout has expired.
      * @param future The request future to wait for
-     * @param timeout The maximum duration (in ms) to wait for the request
+     * @param timer Timer bounding how long this method can block
      * @return true if the future is done, false otherwise
      * @throws WakeupException if {@link #wakeup()} is called from another thread
      * @throws InterruptException if the calling thread is interrupted
      */
-    public boolean poll(RequestFuture future, long timeout) {
-        long begin = time.milliseconds();
-        long remaining = timeout;
-        long now = begin;
+    public boolean poll(RequestFuture future, Timer timer) {
         do {
-            poll(remaining, now, future);
-            now = time.milliseconds();
-            long elapsed = now - begin;
-            remaining = timeout - elapsed;
-        } while (!future.isDone() && remaining > 0);
+            poll(timer, future);
+        } while (!future.isDone() && timer.notExpired());
         return future.isDone();
     }
 
     /**
      * Poll for any network IO.
-     * @param timeout The maximum time to wait for an IO event.
+     * @param timer Timer bounding how long this method can block
      * @throws WakeupException if {@link #wakeup()} is called from another thread
      * @throws InterruptException if the calling thread is interrupted
      */
-    public void poll(long timeout) {
-        poll(timeout, time.milliseconds(), null);
+    public void poll(Timer timer) {
+        poll(timer, null);
     }
 
     /**
      * Poll for any network IO.
-     * @param timeout timeout in milliseconds
-     * @param now current time in milliseconds
+     * @param timer Timer bounding how long this method can block
+     * @param pollCondition Nullable blocking condition
      */
-    public void poll(long timeout, long now, PollCondition pollCondition) {
-        poll(timeout, now, pollCondition, false);
+    public void poll(Timer timer, PollCondition pollCondition) {
+        poll(timer, pollCondition, false);
     }
 
     /**
      * Poll for any network IO.
-     * @param timeout timeout in milliseconds
-     * @param now current time in milliseconds
+     * @param timer Timer bounding how long this method can block
+     * @param pollCondition Nullable blocking condition
      * @param disableWakeup If TRUE disable triggering wake-ups
      */
-    public void poll(long timeout, long now, PollCondition pollCondition, boolean disableWakeup) {
+    public void poll(Timer timer, PollCondition pollCondition, boolean disableWakeup) {
         // there may be handlers which need to be invoked if we woke up the previous call to poll
         firePendingCompletedRequests();
 
-        synchronized (this) {
+        lock.lock();
+        try {
+            // Handle async disconnects prior to attempting any sends
+            handlePendingDisconnects();
+
             // send all the requests we can send now
-            trySend(now);
+            long pollDelayMs = trySend(timer.currentTimeMs());
 
             // check whether the poll is still needed by the caller. Note that if the expected completion
             // condition becomes satisfied after the call to shouldBlock() (because of a fired completion
             // handler), the client will be woken up.
-            if (pollCondition == null || pollCondition.shouldBlock()) {
+            if (pendingCompletion.isEmpty() && (pollCondition == null || pollCondition.shouldBlock())) {
                 // if there are no requests in flight, do not block longer than the retry backoff
+                long pollTimeout = Math.min(timer.remainingMs(), pollDelayMs);
                 if (client.inFlightRequestCount() == 0)
-                    timeout = Math.min(timeout, retryBackoffMs);
-                client.poll(Math.min(maxPollTimeoutMs, timeout), now);
-                now = time.milliseconds();
+                    pollTimeout = Math.min(pollTimeout, retryBackoffMs);
+                client.poll(pollTimeout, timer.currentTimeMs());
             } else {
-                client.poll(0, now);
+                client.poll(0, timer.currentTimeMs());
             }
+            timer.update();
 
             // handle any disconnects by failing the active requests. note that disconnects must
             // be checked immediately following poll since any subsequent call to client.ready()
             // will reset the disconnect status
-            checkDisconnects(now);
+            checkDisconnects(timer.currentTimeMs());
             if (!disableWakeup) {
                 // trigger wakeups after checking for disconnects so that the callbacks will be ready
                 // to be fired on the next call to poll()
@@ -258,41 +282,61 @@ public void poll(long timeout, long now, PollCondition pollCondition, boolean di
 
             // try again to send requests since buffer space may have been
             // cleared or a connect finished in the poll
-            trySend(now);
+            trySend(timer.currentTimeMs());
 
             // fail requests that couldn't be sent if they have expired
-            failExpiredRequests(now);
+            failExpiredRequests(timer.currentTimeMs());
 
             // clean unsent requests collection to keep the map from growing indefinitely
             unsent.clean();
+        } finally {
+            lock.unlock();
         }
 
         // called without the lock to avoid deadlock potential if handlers need to acquire locks
         firePendingCompletedRequests();
+
+        metadata.maybeThrowAnyException();
     }
 
     /**
      * Poll for network IO and return immediately. This will not trigger wakeups.
      */
     public void pollNoWakeup() {
-        poll(0, time.milliseconds(), null, true);
+        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
-     * @param timeoutMs The maximum time in milliseconds to block
+     * @param timer Timer bounding how long this method can block
      * @return true If all requests finished, false if the timeout expired first
      */
-    public boolean awaitPendingRequests(Node node, long timeoutMs) {
-        long startMs = time.milliseconds();
-        long remainingMs = timeoutMs;
-
-        while (hasPendingRequests(node) && remainingMs > 0) {
-            poll(remainingMs);
-            remainingMs = timeoutMs - (time.milliseconds() - startMs);
+    public boolean awaitPendingRequests(Node node, Timer timer) {
+        while (hasPendingRequests(node) && timer.notExpired()) {
+            poll(timer);
         }
-
         return !hasPendingRequests(node);
     }
 
@@ -303,8 +347,11 @@ public boolean awaitPendingRequests(Node node, long timeoutMs) {
      * @return The number of pending requests
      */
     public int pendingRequestCount(Node node) {
-        synchronized (this) {
+        lock.lock();
+        try {
             return unsent.requestCount(node) + client.inFlightRequestCount(node.idString());
+        } finally {
+            lock.unlock();
         }
     }
 
@@ -317,8 +364,11 @@ public int pendingRequestCount(Node node) {
     public boolean hasPendingRequests(Node node) {
         if (unsent.hasRequests(node))
             return true;
-        synchronized (this) {
+        lock.lock();
+        try {
             return client.hasInFlightRequests(node.idString());
+        } finally {
+            lock.unlock();
         }
     }
 
@@ -328,8 +378,11 @@ public boolean hasPendingRequests(Node node) {
      * @return The total count of pending requests
      */
     public int pendingRequestCount() {
-        synchronized (this) {
+        lock.lock();
+        try {
             return unsent.requestCount() + client.inFlightRequestCount();
+        } finally {
+            lock.unlock();
         }
     }
 
@@ -341,8 +394,11 @@ public int pendingRequestCount() {
     public boolean hasPendingRequests() {
         if (unsent.hasRequests())
             return true;
-        synchronized (this) {
+        lock.lock();
+        try {
             return client.hasInFlightRequests();
+        } finally {
+            lock.unlock();
         }
     }
 
@@ -375,66 +431,80 @@ private void checkDisconnects(long now) {
                 for (ClientRequest request : requests) {
                     RequestFutureCompletionHandler handler = (RequestFutureCompletionHandler) request.callback();
                     AuthenticationException authenticationException = client.authenticationException(node);
-                    if (authenticationException != null)
-                        handler.onFailure(authenticationException);
-                    else
-                        handler.onComplete(new ClientResponse(request.makeHeader(request.requestBuilder().latestAllowedVersion()),
+                    handler.onComplete(new ClientResponse(request.makeHeader(request.requestBuilder().latestAllowedVersion()),
                             request.callback(), request.destination(), request.createdTimeMs(), now, true,
-                            null, null));
+                            null, authenticationException, null));
                 }
             }
         }
     }
 
-    public void disconnect(Node node) {
-        synchronized (this) {
-            failUnsentRequests(node, DisconnectException.INSTANCE);
-            client.disconnect(node.idString());
+    private void handlePendingDisconnects() {
+        lock.lock();
+        try {
+            while (true) {
+                Node node = pendingDisconnects.poll();
+                if (node == null)
+                    break;
+
+                failUnsentRequests(node, DisconnectException.INSTANCE);
+                client.disconnect(node.idString());
+            }
+        } finally {
+            lock.unlock();
         }
+    }
 
-        // We need to poll to ensure callbacks from in-flight requests on the disconnected socket are fired
-        pollNoWakeup();
+    public void disconnectAsync(Node node) {
+        pendingDisconnects.offer(node);
+        client.wakeup();
     }
 
     private void failExpiredRequests(long now) {
         // clear all expired unsent requests and fail their corresponding futures
-        Collection expiredRequests = unsent.removeExpiredRequests(now, unsentExpiryMs);
+        Collection expiredRequests = unsent.removeExpiredRequests(now);
         for (ClientRequest request : expiredRequests) {
             RequestFutureCompletionHandler handler = (RequestFutureCompletionHandler) request.callback();
-            handler.onFailure(new TimeoutException("Failed to send request after " + unsentExpiryMs + " ms."));
+            handler.onFailure(new TimeoutException("Failed to send request after " + request.requestTimeoutMs() + " ms."));
         }
     }
 
     private void failUnsentRequests(Node node, RuntimeException e) {
         // clear unsent requests to node and fail their corresponding futures
-        synchronized (this) {
+        lock.lock();
+        try {
             Collection unsentRequests = unsent.remove(node);
             for (ClientRequest unsentRequest : unsentRequests) {
                 RequestFutureCompletionHandler handler = (RequestFutureCompletionHandler) unsentRequest.callback();
                 handler.onFailure(e);
             }
+        } finally {
+            lock.unlock();
         }
-
-        // called without the lock to avoid deadlock potential
-        firePendingCompletedRequests();
     }
 
-    private boolean trySend(long now) {
-        // send any requests that can be sent now
-        boolean requestsSent = false;
+    // Visible for testing
+    long trySend(long now) {
+        long pollDelayMs = maxPollTimeoutMs;
 
+        // send any requests that can be sent now
         for (Node node : unsent.nodes()) {
             Iterator iterator = unsent.requestIterator(node);
+            if (iterator.hasNext())
+                pollDelayMs = Math.min(pollDelayMs, client.pollDelayMs(node, now));
+
             while (iterator.hasNext()) {
                 ClientRequest request = iterator.next();
                 if (client.ready(node, now)) {
                     client.send(request, now);
                     iterator.remove();
-                    requestsSent = true;
+                } else {
+                    // try next node when current node is not ready
+                    break;
                 }
             }
         }
-        return requestsSent;
+        return pollDelayMs;
     }
 
     public void maybeTriggerWakeup() {
@@ -457,19 +527,39 @@ public void disableWakeups() {
 
     @Override
     public void close() throws IOException {
-        synchronized (this) {
+        lock.lock();
+        try {
             client.close();
+        } finally {
+            lock.unlock();
+        }
+    }
+
+
+    /**
+     * Check if the code is disconnected and unavailable for immediate reconnection (i.e. if it is in
+     * reconnect backoff window following the disconnect).
+     */
+    public boolean isUnavailable(Node node) {
+        lock.lock();
+        try {
+            return client.connectionFailed(node) && client.connectionDelay(node, time.milliseconds()) > 0;
+        } finally {
+            lock.unlock();
         }
     }
 
     /**
-     * Find whether a previous connection has failed. Note that the failure state will persist until either
-     * {@link #tryConnect(Node)} or {@link #send(Node, AbstractRequest.Builder)} has been called.
-     * @param node Node to connect to if possible
+     * Check for an authentication error on a given node and raise the exception if there is one.
      */
-    public boolean connectionFailed(Node node) {
-        synchronized (this) {
-            return client.connectionFailed(node);
+    public void maybeThrowAuthFailure(Node node) {
+        lock.lock();
+        try {
+            AuthenticationException exception = client.authenticationException(node);
+            if (exception != null)
+                throw exception;
+        } finally {
+            lock.unlock();
         }
     }
 
@@ -480,8 +570,11 @@ public boolean connectionFailed(Node node) {
      * @param node The node to connect to
      */
     public void tryConnect(Node node) {
-        synchronized (this) {
+        lock.lock();
+        try {
             client.ready(node, time.milliseconds());
+        } finally {
+            lock.unlock();
         }
     }
 
@@ -497,11 +590,11 @@ private RequestFutureCompletionHandler() {
         public void fireCompletion() {
             if (e != null) {
                 future.raise(e);
+            } else if (response.authenticationException() != null) {
+                future.raise(response.authenticationException());
             } else if (response.wasDisconnected()) {
-                RequestHeader requestHeader = response.requestHeader();
-                int correlation = requestHeader.correlationId();
-                log.debug("Cancelled {} request {} with correlation id {} due to node {} being disconnected",
-                        requestHeader.apiKey(), requestHeader, correlation, response.destination());
+                log.debug("Cancelled request with header {} due to node {} being disconnected",
+                        response.requestHeader(), response.destination());
                 future.raise(DisconnectException.INSTANCE);
             } else if (response.versionMismatch() != null) {
                 future.raise(response.versionMismatch());
@@ -539,7 +632,7 @@ public interface PollCondition {
     }
 
     /*
-     * A threadsafe helper class to hold requests per node that have not been sent yet
+     * A thread-safe helper class to hold requests per node that have not been sent yet
      */
     private final static class UnsentRequests {
         private final ConcurrentMap> unsent;
@@ -584,13 +677,14 @@ public boolean hasRequests() {
             return false;
         }
 
-        public Collection removeExpiredRequests(long now, long unsentExpiryMs) {
+        private Collection removeExpiredRequests(long now) {
             List expiredRequests = new ArrayList<>();
             for (ConcurrentLinkedQueue requests : unsent.values()) {
                 Iterator requestIterator = requests.iterator();
                 while (requestIterator.hasNext()) {
                     ClientRequest request = requestIterator.next();
-                    if (request.createdTimeMs() < now - unsentExpiryMs) {
+                    long elapsedMs = Math.max(0, now - request.createdTimeMs());
+                    if (elapsedMs > request.requestTimeoutMs()) {
                         expiredRequests.add(request);
                         requestIterator.remove();
                     } else
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 920c2957c4dc9..a05e8715eae84 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,13 +16,15 @@
  */
 package org.apache.kafka.clients.consumer.internals;
 
+import java.nio.BufferUnderflowException;
+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;
-import org.apache.kafka.common.protocol.types.Schema;
+import org.apache.kafka.common.message.ConsumerProtocolAssignment;
+import org.apache.kafka.common.message.ConsumerProtocolSubscription;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.MessageUtil;
 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 org.apache.kafka.common.utils.CollectionUtils;
 
 import java.nio.ByteBuffer;
@@ -32,118 +34,143 @@
 
 /**
  * ConsumerProtocol contains the schemas for consumer subscriptions and assignments for use with
- * Kafka's generalized group management protocol. Below is the version 0 format:
- *
- * 
    - * Subscription => Version Topics
- *   Version    => Int16
- *   Topics     => [String]
- *   UserData   => Bytes
- *
- * Assignment => Version TopicPartitions
- *   Version         => int16
- *   TopicPartitions => [Topic Partitions]
- *     Topic         => String
- *     Partitions    => [int32]
- * 
    
+ * Kafka's generalized group management protocol.
  *
  * 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.
  */
 public class ConsumerProtocol {
-
     public static final String PROTOCOL_TYPE = "consumer";
 
-    public static final String VERSION_KEY_NAME = "version";
-    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 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 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);
-
-    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 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) {
-        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);
-        buffer.flip();
-        return buffer;
+    static {
+        // Safety check to ensure that both parts of the consumer protocol remain in sync.
+        if (ConsumerProtocolSubscription.LOWEST_SUPPORTED_VERSION
+                != ConsumerProtocolAssignment.LOWEST_SUPPORTED_VERSION)
+            throw new IllegalStateException("Subscription and Assignment schemas must have the " +
+                "same lowest version");
+
+        if (ConsumerProtocolSubscription.HIGHEST_SUPPORTED_VERSION
+                != ConsumerProtocolAssignment.HIGHEST_SUPPORTED_VERSION)
+            throw new IllegalStateException("Subscription and Assignment schemas must have the " +
+                "same highest version");
     }
 
-    public static PartitionAssignor.Subscription deserializeSubscription(ByteBuffer buffer) {
-        Struct header = CONSUMER_PROTOCOL_HEADER_SCHEMA.read(buffer);
-        Short version = header.getShort(VERSION_KEY_NAME);
-        checkVersionCompatibility(version);
-        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);
+    public static short deserializeVersion(final ByteBuffer buffer) {
+        try {
+            return buffer.getShort();
+        } catch (BufferUnderflowException e) {
+            throw new SchemaException("Buffer underflow while parsing consumer protocol's header", e);
+        }
     }
 
-    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);
-        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)) {
-                Integer partition = (Integer) partitionObj;
-                partitions.add(new TopicPartition(topic, partition));
+    public static ByteBuffer serializeSubscription(final Subscription subscription) {
+        return serializeSubscription(subscription, ConsumerProtocolSubscription.HIGHEST_SUPPORTED_VERSION);
+    }
+
+    public static ByteBuffer serializeSubscription(final Subscription subscription, short version) {
+        version = checkSubscriptionVersion(version);
+
+        ConsumerProtocolSubscription data = new ConsumerProtocolSubscription();
+        data.setTopics(subscription.topics());
+        data.setUserData(subscription.userData() != null ? subscription.userData().duplicate() : null);
+        Map> partitionsByTopic = CollectionUtils.groupPartitionsByTopic(subscription.ownedPartitions());
+        for (Map.Entry> topicEntry : partitionsByTopic.entrySet()) {
+            data.ownedPartitions().add(new ConsumerProtocolSubscription.TopicPartition()
+                .setTopic(topicEntry.getKey())
+                .setPartitions(topicEntry.getValue()));
+        }
+
+        return MessageUtil.toVersionPrefixedByteBuffer(version, data);
+    }
+
+    public static Subscription deserializeSubscription(final ByteBuffer buffer, short version) {
+        version = checkSubscriptionVersion(version);
+
+        try {
+            ConsumerProtocolSubscription data =
+                new ConsumerProtocolSubscription(new ByteBufferAccessor(buffer), version);
+
+            List ownedPartitions = new ArrayList<>();
+            for (ConsumerProtocolSubscription.TopicPartition tp : data.ownedPartitions()) {
+                for (Integer partition : tp.partitions()) {
+                    ownedPartitions.add(new TopicPartition(tp.topic(), partition));
+                }
             }
+
+            return new Subscription(
+                data.topics(),
+                data.userData() != null ? data.userData().duplicate() : null,
+                ownedPartitions);
+        } catch (BufferUnderflowException e) {
+            throw new SchemaException("Buffer underflow while parsing consumer protocol's subscription", e);
         }
-        return new PartitionAssignor.Assignment(partitions, userData);
     }
 
-    public static ByteBuffer serializeAssignment(PartitionAssignor.Assignment assignment) {
-        Struct struct = new Struct(ASSIGNMENT_V0);
-        struct.set(USER_DATA_KEY_NAME, assignment.userData());
-        List topicAssignments = new ArrayList<>();
-        Map> partitionsByTopic = CollectionUtils.groupDataByTopic(assignment.partitions());
+    public static Subscription deserializeSubscription(final ByteBuffer buffer) {
+        return deserializeSubscription(buffer, deserializeVersion(buffer));
+    }
+
+    public static ByteBuffer serializeAssignment(final Assignment assignment) {
+        return serializeAssignment(assignment, ConsumerProtocolAssignment.HIGHEST_SUPPORTED_VERSION);
+    }
+
+    public static ByteBuffer serializeAssignment(final Assignment assignment, short version) {
+        version = checkAssignmentVersion(version);
+
+        ConsumerProtocolAssignment data = new ConsumerProtocolAssignment();
+        data.setUserData(assignment.userData() != null ? assignment.userData().duplicate() : null);
+        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);
+            data.assignedPartitions().add(new ConsumerProtocolAssignment.TopicPartition()
+                .setTopic(topicEntry.getKey())
+                .setPartitions(topicEntry.getValue()));
         }
-        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);
-        buffer.flip();
-        return buffer;
+
+        return MessageUtil.toVersionPrefixedByteBuffer(version, data);
     }
 
-    private static void checkVersionCompatibility(short version) {
-        // check for invalid versions
-        if (version < CONSUMER_PROTOCOL_V0)
-            throw new SchemaException("Unsupported subscription version: " + version);
+    public static Assignment deserializeAssignment(final ByteBuffer buffer, short version) {
+        version = checkAssignmentVersion(version);
+
+        try {
+            ConsumerProtocolAssignment data =
+                new ConsumerProtocolAssignment(new ByteBufferAccessor(buffer), version);
+
+            List assignedPartitions = new ArrayList<>();
+            for (ConsumerProtocolAssignment.TopicPartition tp : data.assignedPartitions()) {
+                for (Integer partition : tp.partitions()) {
+                    assignedPartitions.add(new TopicPartition(tp.topic(), partition));
+                }
+            }
+
+            return new Assignment(
+                assignedPartitions,
+                data.userData() != null ? data.userData().duplicate() : null);
+        } catch (BufferUnderflowException e) {
+            throw new SchemaException("Buffer underflow while parsing consumer protocol's assignment", e);
+        }
+    }
 
-        // otherwise, assume versions can be parsed as V0
+    public static Assignment deserializeAssignment(final ByteBuffer buffer) {
+        return deserializeAssignment(buffer, deserializeVersion(buffer));
     }
 
+    private static short checkSubscriptionVersion(final short version) {
+        if (version < ConsumerProtocolSubscription.LOWEST_SUPPORTED_VERSION)
+            throw new SchemaException("Unsupported subscription version: " + version);
+        else if (version > ConsumerProtocolSubscription.HIGHEST_SUPPORTED_VERSION)
+            return ConsumerProtocolSubscription.HIGHEST_SUPPORTED_VERSION;
+        else
+            return version;
+    }
+
+    private static short checkAssignmentVersion(final short version) {
+        if (version < ConsumerProtocolAssignment.LOWEST_SUPPORTED_VERSION)
+            throw new SchemaException("Unsupported assignment version: " + version);
+        else if (version > ConsumerProtocolAssignment.HIGHEST_SUPPORTED_VERSION)
+            return ConsumerProtocolAssignment.HIGHEST_SUPPORTED_VERSION;
+        else
+            return version;
+    }
 }
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 4c68f1f84830f..7b1101a0a8a19 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,30 @@
  */
 package org.apache.kafka.clients.consumer.internals;
 
+import org.apache.kafka.clients.ApiVersion;
+import org.apache.kafka.clients.ApiVersions;
 import org.apache.kafka.clients.ClientResponse;
+import org.apache.kafka.clients.FetchSessionHandler;
 import org.apache.kafka.clients.Metadata;
+import org.apache.kafka.clients.NodeApiVersions;
+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.NoOffsetForPartitionException;
+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.clients.consumer.internals.OffsetsForLeaderEpochClient.OffsetForEpochResult;
+import org.apache.kafka.clients.consumer.internals.SubscriptionState.FetchPosition;
 import org.apache.kafka.common.Cluster;
+import org.apache.kafka.common.IsolationLevel;
 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.InvalidMetadataException;
+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,37 +48,45 @@
 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.message.ListOffsetRequestData.ListOffsetPartition;
+import org.apache.kafka.common.message.ListOffsetResponseData.ListOffsetPartitionResponse;
+import org.apache.kafka.common.message.ListOffsetResponseData.ListOffsetTopicResponse;
+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.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;
 import org.apache.kafka.common.record.TimestampType;
 import org.apache.kafka.common.requests.FetchRequest;
 import org.apache.kafka.common.requests.FetchResponse;
-import org.apache.kafka.common.requests.IsolationLevel;
 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.OffsetsForLeaderEpochRequest;
 import org.apache.kafka.common.serialization.Deserializer;
-import org.apache.kafka.common.serialization.ExtendedDeserializer;
 import org.apache.kafka.common.utils.CloseableIterator;
 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 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;
@@ -79,19 +97,42 @@
 import java.util.LinkedHashMap;
 import java.util.List;
 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;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Function;
+import java.util.stream.Collectors;
 
 import static java.util.Collections.emptyList;
-import static org.apache.kafka.common.serialization.ExtendedDeserializer.Wrapper.ensureExtended;
 
 /**
- * This class manage the fetching process with the brokers.
+ * This class manages the fetching process with the brokers.
+ * 
    
+ * Thread-safety:
+ * Requests and responses of Fetcher may be processed by different threads since heartbeat
+ * thread may process responses. Other operations are single-threaded and invoked only from
+ * the thread polling the consumer.
+ * 
      
+ *     
    • If a response handler accesses any shared state of the Fetcher (e.g. FetchSessionHandler),
+ *     all access to that state must be synchronized on the Fetcher instance.
      • 
+ *     
    • If a response handler accesses any shared state of the coordinator (e.g. SubscriptionState),
+ *     it is assumed that all access to that state is synchronized on the coordinator instance by
+ *     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;
     private final Time time;
     private final int minBytes;
@@ -99,19 +140,28 @@ public class Fetcher implements SubscriptionState.Listener, Closeable {
     private final int maxWaitMs;
     private final int fetchSize;
     private final long retryBackoffMs;
+    private final long requestTimeoutMs;
     private final int maxPollRecords;
     private final boolean checkCrcs;
-    private final Metadata metadata;
+    private final String clientRackId;
+    private final ConsumerMetadata metadata;
     private final FetchManagerMetrics sensors;
     private final SubscriptionState subscriptions;
     private final ConcurrentLinkedQueue completedFetches;
     private final BufferSupplier decompressionBufferSupplier = BufferSupplier.create();
-
-    private final ExtendedDeserializer keyDeserializer;
-    private final ExtendedDeserializer valueDeserializer;
+    private final Deserializer keyDeserializer;
+    private final Deserializer valueDeserializer;
     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 final AtomicInteger metadataUpdateVersion = new AtomicInteger(-1);
+
 
-    private PartitionRecords nextInLineRecords = null;
+    private CompletedFetch nextInLineFetch = null;
 
     public Fetcher(LogContext logContext,
                    ConsumerNetworkClient client,
@@ -121,16 +171,20 @@ public Fetcher(LogContext logContext,
                    int fetchSize,
                    int maxPollRecords,
                    boolean checkCrcs,
+                   String clientRackId,
                    Deserializer keyDeserializer,
                    Deserializer valueDeserializer,
-                   Metadata metadata,
+                   ConsumerMetadata metadata,
                    SubscriptionState subscriptions,
                    Metrics metrics,
                    FetcherMetricsRegistry metricsRegistry,
                    Time time,
                    long retryBackoffMs,
-                   IsolationLevel isolationLevel) {
+                   long requestTimeoutMs,
+                   IsolationLevel isolationLevel,
+                   ApiVersions apiVersions) {
         this.log = logContext.logger(Fetcher.class);
+        this.logContext = logContext;
         this.time = time;
         this.client = client;
         this.metadata = metadata;
@@ -141,50 +195,50 @@ public Fetcher(LogContext logContext,
         this.fetchSize = fetchSize;
         this.maxPollRecords = maxPollRecords;
         this.checkCrcs = checkCrcs;
-        this.keyDeserializer = ensureExtended(keyDeserializer);
-        this.valueDeserializer = ensureExtended(valueDeserializer);
+        this.clientRackId = clientRackId;
+        this.keyDeserializer = keyDeserializer;
+        this.valueDeserializer = valueDeserializer;
         this.completedFetches = new ConcurrentLinkedQueue<>();
         this.sensors = new FetchManagerMetrics(metrics, metricsRegistry);
         this.retryBackoffMs = retryBackoffMs;
+        this.requestTimeoutMs = requestTimeoutMs;
         this.isolationLevel = isolationLevel;
-
-        subscriptions.addListener(this);
+        this.apiVersions = apiVersions;
+        this.sessionHandlers = new HashMap<>();
+        this.offsetsForLeaderEpochClient = new OffsetsForLeaderEpochClient(client, logContext);
+        this.nodesWithPendingFetchRequests = new HashSet<>();
     }
 
     /**
      * Represents data about an offset returned by a broker.
      */
-    private static class OffsetData {
-        /**
-         * The offset
-         */
+    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
 
-        /**
-         * The timestamp.
-         *
-         * Will be null if the broker does not support returning timestamps.
-         */
-        final Long timestamp;
-
-        OffsetData(long offset, Long timestamp) {
+        ListOffsetData(long offset, Long timestamp, Optional leaderEpoch) {
             this.offset = offset;
             this.timestamp = timestamp;
+            this.leaderEpoch = leaderEpoch;
         }
     }
 
     /**
-     * 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();
     }
 
-    private boolean matchesRequestedPartitions(FetchRequest.Builder request, FetchResponse response) {
-        Set requestedPartitions = request.fetchData().keySet();
-        Set fetchedPartitions = response.responseData().keySet();
-        return fetchedPartitions.equals(requestedPartitions);
+    /**
+     * 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));
     }
 
     /**
@@ -192,126 +246,133 @@ private boolean matchesRequestedPartitions(FetchRequest.Builder request, FetchRe
      * an in-flight fetch or pending fetch data.
      * @return number of fetches sent
      */
-    public int sendFetches() {
-        Map fetchRequestMap = createFetchRequests();
-        for (Map.Entry fetchEntry : fetchRequestMap.entrySet()) {
-            final FetchRequest.Builder request = fetchEntry.getValue();
-            final Node fetchTarget = fetchEntry.getKey();
-
-            log.debug("Sending {} fetch for partitions {} to broker {}", isolationLevel, request.fetchData().keySet(),
-                    fetchTarget);
-            client.send(fetchTarget, request)
-                    .addListener(new RequestFutureListener() {
-                        @Override
-                        public void onSuccess(ClientResponse resp) {
-                            FetchResponse response = (FetchResponse) resp.responseBody();
-                            if (!matchesRequestedPartitions(request, response)) {
-                                // obviously we expect the broker to always send us valid responses, so this check
-                                // is mainly for test cases where mock fetch responses must be manually crafted.
-                                log.warn("Ignoring fetch response containing partitions {} since it does not match " +
-                                        "the requested partitions {}", response.responseData().keySet(),
-                                        request.fetchData().keySet());
+    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();
+            final FetchSessionHandler.FetchRequestData data = entry.getValue();
+            final FetchRequest.Builder request = FetchRequest.Builder
+                    .forConsumer(this.maxWaitMs, this.minBytes, data.toSend())
+                    .isolationLevel(isolationLevel)
+                    .setMaxBytes(this.maxBytes)
+                    .metadata(data.metadata())
+                    .toForget(data.toForget())
+                    .rackId(clientRackId);
+
+            if (log.isDebugEnabled()) {
+                log.debug("Sending {} {} to broker {}", isolationLevel, data.toString(), fetchTarget);
+            }
+            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;
                             }
 
                             Set partitions = new HashSet<>(response.responseData().keySet());
                             FetchResponseMetricAggregator metricAggregator = new FetchResponseMetricAggregator(sensors, partitions);
 
-                            for (Map.Entry entry : response.responseData().entrySet()) {
+                            for (Map.Entry> entry : response.responseData().entrySet()) {
                                 TopicPartition partition = entry.getKey();
-                                long fetchOffset = request.fetchData().get(partition).fetchOffset;
-                                FetchResponse.PartitionData fetchData = 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()));
+                                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, partitionData);
+
+                                    Iterator batches = partitionData.records().batches().iterator();
+                                    short responseVersion = resp.requestHeader().apiVersion();
+
+                                    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) {
-                            log.debug("Fetch request {} to {} failed", request.fetchData(), fetchTarget, 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();
-    }
-
-    /**
-     * Lookup and set offsets for any partitions which are awaiting an explicit reset.
-     * @param partitions the partitions to reset
-     */
-    public void resetOffsetsIfNeeded(Set partitions) {
-        final Set needsOffsetReset = new HashSet<>();
-        for (TopicPartition tp : partitions) {
-            if (subscriptions.isAssigned(tp) && subscriptions.isOffsetResetNeeded(tp))
-                needsOffsetReset.add(tp);
-        }
-        if (!needsOffsetReset.isEmpty()) {
-            resetOffsets(needsOffsetReset);
-        }
-    }
-
-    /**
-     * Update the fetch positions for the provided partitions.
-     * @param partitions the partitions to update positions for
-     * @throws NoOffsetForPartitionException If no offset is stored for a given partition and no reset policy is available
-     */
-    public void updateFetchPositions(Set partitions) {
-        final Set needsOffsetReset = new HashSet<>();
-        // reset the fetch position to the committed position
-        for (TopicPartition tp : partitions) {
-            if (!subscriptions.isAssigned(tp) || subscriptions.hasValidPosition(tp))
-                continue;
-
-            if (subscriptions.isOffsetResetNeeded(tp)) {
-                needsOffsetReset.add(tp);
-            } else if (subscriptions.committed(tp) == null) {
-                // there's no committed position, so we need to reset with the default strategy
-                subscriptions.needOffsetReset(tp);
-                needsOffsetReset.add(tp);
-            } else {
-                long committed = subscriptions.committed(tp).offset();
-                log.debug("Resetting offset for partition {} to the committed offset {}", tp, committed);
-                subscriptions.seek(tp, committed);
-            }
-        }
+                    }
+                }
+            });
 
-        if (!needsOffsetReset.isEmpty()) {
-            resetOffsets(needsOffsetReset);
         }
+        return fetchRequestMap.size();
     }
 
     /**
      * Get topic metadata for all topics in the cluster
-     * @param timeout time for which getting topic metadata is attempted
+     * @param timer Timer bounding how long this method can block
      * @return The map of topics with their partition information
      */
-    public Map> getAllTopicMetadata(long timeout) {
-        return getTopicMetadata(MetadataRequest.Builder.allTopics(), timeout);
+    public Map> getAllTopicMetadata(Timer timer) {
+        return getTopicMetadata(MetadataRequest.Builder.allTopics(), timer);
     }
 
     /**
      * Get metadata for all topics present in Kafka cluster
      *
      * @param request The MetadataRequest to send
-     * @param timeout time for which getting topic metadata is attempted
+     * @param timer Timer bounding how long this method can block
      * @return The map of topics with their partition information
      */
-    public Map> getTopicMetadata(MetadataRequest.Builder request, long timeout) {
+    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();
 
-        long start = time.milliseconds();
-        long remaining = timeout;
-
         do {
             RequestFuture future = sendMetadataRequest(request);
-            client.poll(future, remaining);
+            client.poll(future, timer);
 
             if (future.failed() && !future.isRetriable())
                 throw future.exception();
@@ -353,20 +414,13 @@ else if (error.exception() instanceof RetriableException)
                 if (!shouldRetry) {
                     HashMap> topicsPartitionInfos = new HashMap<>();
                     for (String topic : cluster.topics())
-                        topicsPartitionInfos.put(topic, cluster.availablePartitionsForTopic(topic));
+                        topicsPartitionInfos.put(topic, cluster.partitionsForTopic(topic));
                     return topicsPartitionInfos;
                 }
             }
 
-            long elapsed = time.milliseconds() - start;
-            remaining = timeout - elapsed;
-
-            if (remaining > 0) {
-                long backoff = Math.min(remaining, retryBackoffMs);
-                time.sleep(backoff);
-                remaining -= backoff;
-            }
-        } while (remaining > 0);
+            timer.sleep(retryBackoffMs);
+        } while (timer.notExpired());
 
         throw new TimeoutException("Timeout expired while fetching topic metadata");
     }
@@ -383,127 +437,155 @@ private RequestFuture sendMetadataRequest(MetadataRequest.Builde
             return client.send(node, request);
     }
 
-    private void offsetResetStrategyTimestamp(
-            final TopicPartition partition,
-            final Map output,
-            final Set partitionsWithNoOffsets) {
+    private Long offsetResetStrategyTimestamp(final TopicPartition partition) {
         OffsetResetStrategy strategy = subscriptions.resetStrategy(partition);
         if (strategy == OffsetResetStrategy.EARLIEST)
-            output.put(partition, ListOffsetRequest.EARLIEST_TIMESTAMP);
+            return ListOffsetRequest.EARLIEST_TIMESTAMP;
         else if (strategy == OffsetResetStrategy.LATEST)
-            output.put(partition, endTimestamp());
+            return ListOffsetRequest.LATEST_TIMESTAMP;
+        else
+            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
-            partitionsWithNoOffsets.add(partition);
+            return null;
     }
 
     /**
-     * Reset offsets for the given partition using the offset reset strategy.
+     * Reset offsets for all assigned partitions that require it.
      *
-     * @param partitions  The partitions that need offsets reset
      * @throws org.apache.kafka.clients.consumer.NoOffsetForPartitionException If no offset reset strategy is defined
+     *   and one or more partitions aren't awaiting a seekToBeginning() or seekToEnd().
      */
-    private void resetOffsets(final Set partitions) {
-        final Map offsetResets = new HashMap<>();
-        final Set partitionsWithNoOffsets = new HashSet<>();
+    public void resetOffsetsIfNeeded() {
+        // Raise exception from previous offset fetch if there is one
+        RuntimeException exception = cachedListOffsetsException.getAndSet(null);
+        if (exception != null)
+            throw exception;
+
+        Set partitions = subscriptions.partitionsNeedingReset(time.milliseconds());
+        if (partitions.isEmpty())
+            return;
+
+        final Map offsetResetTimestamps = new HashMap<>();
         for (final TopicPartition partition : partitions) {
-            offsetResetStrategyTimestamp(partition, offsetResets, partitionsWithNoOffsets);
-        }
-        final Map offsetsByTimes = retrieveOffsetsByTimes(offsetResets, Long.MAX_VALUE, false);
-        for (final TopicPartition partition : partitions) {
-            final OffsetData offsetData = offsetsByTimes.get(partition);
-            if (offsetData == null) {
-                partitionsWithNoOffsets.add(partition);
-                continue;
-            }
-            // we might lose the assignment while fetching the offset, so check it is still active
-            if (subscriptions.isAssigned(partition)) {
-                log.debug("Resetting offset for partition {} to offset {}.", partition, offsetData.offset);
-                this.subscriptions.seek(partition, offsetData.offset);
-            }
-        }
-        if (!partitionsWithNoOffsets.isEmpty()) {
-            throw new NoOffsetForPartitionException(partitionsWithNoOffsets);
+            Long timestamp = offsetResetStrategyTimestamp(partition);
+            if (timestamp != null)
+                offsetResetTimestamps.put(partition, timestamp);
         }
+
+        resetOffsetsAsync(offsetResetTimestamps);
     }
 
-    public Map getOffsetsByTimes(Map timestampsToSearch,
-                                                                     long timeout) {
-        Map offsetData = retrieveOffsetsByTimes(timestampsToSearch, timeout, true);
+    /**
+     * 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
+        // If we see a new metadata version, check all partitions
+        validatePositionsOnMetadataChange();
+
+        // Collect positions needing validation, with backoff
+        Map partitionsToValidate = subscriptions
+                .partitionsNeedingValidation(time.milliseconds())
+                .stream()
+                .filter(tp -> subscriptions.position(tp) != null)
+                .collect(Collectors.toMap(Function.identity(), subscriptions::position));
+
+        validateOffsetsAsync(partitionsToValidate);
+    }
 
-        HashMap offsetsByTimes = new HashMap<>(timestampsToSearch.size());
-        for (Map.Entry entry : timestampsToSearch.entrySet())
-            offsetsByTimes.put(entry.getKey(), null);
+    public Map offsetsForTimes(Map timestampsToSearch,
+                                                                   Timer timer) {
+        metadata.addTransientTopics(topicsForPartitions(timestampsToSearch.keySet()));
 
-        for (Map.Entry entry : offsetData.entrySet()) {
-            // 'entry.getValue().timestamp' will not be null since we are guaranteed
-            // to work with a v1 (or later) ListOffset request
-            offsetsByTimes.put(entry.getKey(), new OffsetAndTimestamp(entry.getValue().offset, entry.getValue().timestamp));
-        }
+        try {
+            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()) {
+                // 'entry.getValue().timestamp' will not be null since we are guaranteed
+                // to work with a v1 (or later) ListOffset request
+                ListOffsetData offsetData = entry.getValue();
+                offsetsByTimes.put(entry.getKey(), new OffsetAndTimestamp(offsetData.offset, offsetData.timestamp,
+                        offsetData.leaderEpoch));
+            }
 
-        return offsetsByTimes;
+            return offsetsByTimes;
+        } finally {
+            metadata.clearTransientTopics();
+        }
     }
 
-    private Map retrieveOffsetsByTimes(
-            Map timestampsToSearch, long timeout, boolean requireTimestamps) {
+    private ListOffsetResult fetchOffsetsByTimes(Map timestampsToSearch,
+                                                 Timer timer,
+                                                 boolean requireTimestamps) {
+        ListOffsetResult result = new ListOffsetResult();
         if (timestampsToSearch.isEmpty())
-            return Collections.emptyMap();
+            return result;
 
-        long startMs = time.milliseconds();
-        long remaining = timeout;
+        Map remainingToSearch = new HashMap<>(timestampsToSearch);
         do {
-            RequestFuture> future =
-                    sendListOffsetRequests(requireTimestamps, timestampsToSearch);
-            client.poll(future, remaining);
+            RequestFuture future = sendListOffsetsRequests(remainingToSearch, requireTimestamps);
+            client.poll(future, timer);
 
-            if (!future.isDone())
+            if (!future.isDone()) {
                 break;
-
-            if (future.succeeded())
-                return future.value();
-
-            if (!future.isRetriable())
+            } else if (future.succeeded()) {
+                ListOffsetResult value = future.value();
+                result.fetchedOffsets.putAll(value.fetchedOffsets);
+                remainingToSearch.keySet().retainAll(value.partitionsToRetry);
+            } else if (!future.isRetriable()) {
                 throw future.exception();
+            }
 
-            long elapsed = time.milliseconds() - startMs;
-            remaining = timeout - elapsed;
-            if (remaining <= 0)
-                break;
-
-            if (future.exception() instanceof InvalidMetadataException)
-                client.awaitMetadataUpdate(remaining);
-            else
-                time.sleep(Math.min(remaining, retryBackoffMs));
-
-            elapsed = time.milliseconds() - startMs;
-            remaining = timeout - elapsed;
-        } while (remaining > 0);
-        throw new TimeoutException("Failed to get offsets by times in " + timeout + " ms");
-    }
+            if (remainingToSearch.isEmpty()) {
+                return result;
+            } else {
+                client.awaitMetadataUpdate(timer);
+            }
+        } while (timer.notExpired());
 
-    public Map beginningOffsets(Collection partitions, long timeout) {
-        return beginningOrEndOffset(partitions, ListOffsetRequest.EARLIEST_TIMESTAMP, timeout);
+        throw new TimeoutException("Failed to get offsets by times in " + timer.elapsedMs() + "ms");
     }
 
-    public Map endOffsets(Collection partitions, long timeout) {
-        return beginningOrEndOffset(partitions, endTimestamp(), timeout);
+    public Map beginningOffsets(Collection partitions, Timer timer) {
+        return beginningOrEndOffset(partitions, ListOffsetRequest.EARLIEST_TIMESTAMP, timer);
     }
 
-    private long endTimestamp() {
-        return ListOffsetRequest.LATEST_TIMESTAMP;
+    public Map endOffsets(Collection partitions, Timer timer) {
+        return beginningOrEndOffset(partitions, ListOffsetRequest.LATEST_TIMESTAMP, timer);
     }
 
     private Map beginningOrEndOffset(Collection partitions,
                                                            long timestamp,
-                                                           long timeout) {
-        Map timestampsToSearch = new HashMap<>();
-        for (TopicPartition tp : partitions)
-            timestampsToSearch.put(tp, timestamp);
-        Map result = new HashMap<>();
-        for (Map.Entry entry :
-                retrieveOffsetsByTimes(timestampsToSearch, timeout, false).entrySet()) {
-            result.put(entry.getKey(), entry.getValue().offset);
+                                                           Timer timer) {
+        metadata.addTransientTopics(topicsForPartitions(partitions));
+        try {
+            Map timestampsToSearch = partitions.stream()
+                    .distinct()
+                    .collect(Collectors.toMap(Function.identity(), tp -> timestamp));
+
+            ListOffsetResult result = fetchOffsetsByTimes(timestampsToSearch, timer, false);
+
+            return result.fetchedOffsets.entrySet().stream()
+                    .collect(Collectors.toMap(Map.Entry::getKey, entry -> entry.getValue().offset));
+        } finally {
+            metadata.clearTransientTopics();
         }
-        return result;
     }
 
     /**
@@ -514,23 +596,49 @@ 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;
-
-                    nextInLineRecords = parseCompletedFetch(completedFetch);
+                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;
+                        }
+                    } 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);
@@ -550,107 +658,312 @@ 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);
+                    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",
+                    completedFetch.partition);
         } else {
-            // note that the consumed position should always be available as long as the partition is still assigned
-            long position = subscriptions.position(partitionRecords.partition);
-            if (!subscriptions.isFetchable(partitionRecords.partition)) {
-                // this can happen when a partition is paused before fetched records are returned to the consumer's poll call
-                log.debug("Not returning fetched records for assigned partition {} since it is no longer fetchable",
-                        partitionRecords.partition);
-            } else 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);
-
-                Long partitionLag = subscriptions.partitionLag(partitionRecords.partition, isolationLevel);
+            FetchPosition position = subscriptions.position(completedFetch.partition);
+            if (position == null) {
+                throw new IllegalStateException("Missing position for fetchable partition " + completedFetch.partition);
+            }
+
+            if (completedFetch.nextFetchOffset == position.offset) {
+                List> partRecords = completedFetch.fetchRecords(maxRecords);
+
+                log.trace("Returning {} fetched records at offset {} for assigned partition {}",
+                        partRecords.size(), position, completedFetch.partition);
+
+                if (completedFetch.nextFetchOffset > position.offset) {
+                    FetchPosition nextPosition = new FetchPosition(
+                            completedFetch.nextFetchOffset,
+                            completedFetch.lastEpoch,
+                            position.currentLeader);
+                    log.trace("Update fetching position to {} for partition {}", nextPosition, completedFetch.partition);
+                    subscriptions.position(completedFetch.partition, nextPosition);
+                }
+
+                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(completedFetch.partition);
+                if (lead != null) {
+                    this.sensors.recordPartitionLead(completedFetch.partition, lead);
+                }
 
                 return partRecords;
             } else {
                 // 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();
     }
 
+    // Visible for testing
+    void resetOffsetIfNeeded(TopicPartition partition, OffsetResetStrategy requestedResetStrategy, ListOffsetData offsetData) {
+        FetchPosition position = new FetchPosition(
+            offsetData.offset,
+            Optional.empty(), // This will ensure we skip validation
+            metadata.currentLeader(partition));
+        offsetData.leaderEpoch.ifPresent(epoch -> metadata.updateLastSeenEpochIfNewer(partition, epoch));
+        subscriptions.maybeSeekUnvalidated(partition, position, requestedResetStrategy);
+    }
+
+    private void resetOffsetsAsync(Map partitionResetTimestamps) {
+        Map> timestampsToSearchByNode =
+                groupListOffsetRequests(partitionResetTimestamps, new HashSet<>());
+        for (Map.Entry> entry : timestampsToSearchByNode.entrySet()) {
+            Node node = entry.getKey();
+            final Map resetTimestamps = entry.getValue();
+            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.requestFailed(result.partitionsToRetry, time.milliseconds() + retryBackoffMs);
+                        metadata.requestUpdate();
+                    }
+
+                    for (Map.Entry fetchedOffset : result.fetchedOffsets.entrySet()) {
+                        TopicPartition partition = fetchedOffset.getKey();
+                        ListOffsetData offsetData = fetchedOffset.getValue();
+                        ListOffsetPartition requestedReset = resetTimestamps.get(partition);
+                        resetOffsetIfNeeded(partition, timestampToOffsetResetStrategy(requestedReset.timestamp()), offsetData);
+                    }
+                }
+
+                @Override
+                public void onFailure(RuntimeException e) {
+                    subscriptions.requestFailed(resetTimestamps.keySet(), time.milliseconds() + retryBackoffMs);
+                    metadata.requestUpdate();
+
+                    if (!(e instanceof RetriableException) && !cachedListOffsetsException.compareAndSet(null, e))
+                        log.error("Discarding error in ListOffsetResponse because another error is pending", e);
+                }
+            });
+        }
+    }
+
+    static 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);
+
+        long nextResetTimeMs = time.milliseconds() + requestTimeoutMs;
+        regrouped.forEach((node, fetchPositions) -> {
+            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)",
+                    fetchPositions.keySet());
+                for (TopicPartition partition : fetchPositions.keySet()) {
+                    subscriptions.completeValidation(partition);
+                }
+                return;
+            }
+
+            subscriptions.setNextAllowedRetry(fetchPositions.keySet(), nextResetTimeMs);
+
+            RequestFuture future =
+                offsetsForLeaderEpochClient.sendAsyncRequest(node, fetchPositions);
+
+            future.addListener(new RequestFutureListener() {
+                @Override
+                public void onSuccess(OffsetForEpochResult offsetsResult) {
+                    List truncations = new ArrayList<>();
+                    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.
+                    //
+                    // In addition, check whether the returned offset and epoch are valid. If not, then we should reset
+                    // its offset if reset policy is configured, or throw out of range exception.
+                    offsetsResult.endOffsets().forEach((topicPartition, respEndOffset) -> {
+                        FetchPosition requestPosition = fetchPositions.get(topicPartition);
+                        Optional truncationOpt =
+                            subscriptions.maybeCompleteValidation(topicPartition, requestPosition, respEndOffset);
+                        truncationOpt.ifPresent(truncations::add);
+                    });
+
+                    if (!truncations.isEmpty()) {
+                        maybeSetOffsetForLeaderException(buildLogTruncationException(truncations));
+                    }
+                }
+
+                @Override
+                public void onFailure(RuntimeException e) {
+                    subscriptions.requestFailed(fetchPositions.keySet(), time.milliseconds() + retryBackoffMs);
+                    metadata.requestUpdate();
+
+                    if (!(e instanceof RetriableException)) {
+                        maybeSetOffsetForLeaderException(e);
+                    }
+                }
+            });
+        });
+    }
+
+    private LogTruncationException buildLogTruncationException(List truncations) {
+        Map divergentOffsets = new HashMap<>();
+        Map truncatedFetchOffsets = new HashMap<>();
+        for (SubscriptionState.LogTruncation truncation : truncations) {
+            truncation.divergentOffsetOpt.ifPresent(divergentOffset ->
+                divergentOffsets.put(truncation.topicPartition, divergentOffset));
+            truncatedFetchOffsets.put(truncation.topicPartition, truncation.fetchPosition.offset);
+        }
+        return new LogTruncationException("Detected truncated partitions: " + truncations,
+            truncatedFetchOffsets, divergentOffsets);
+    }
+
+    private void maybeSetOffsetForLeaderException(RuntimeException e) {
+        if (!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.
      *
+     * @param timestampsToSearch the mapping between partitions and target time
      * @param requireTimestamps true if we should fail with an UnsupportedVersionException if the broker does
      *                         not support fetching precise timestamps for offsets
-     * @param timestampsToSearch the mapping between partitions and target time
      * @return A response which can be polled to obtain the corresponding timestamps and offsets.
      */
-    private RequestFuture> sendListOffsetRequests(
-            final boolean requireTimestamps,
-            final Map timestampsToSearch) {
-        // Group the partitions by node.
-        final Map> timestampsToSearchByNode = new HashMap<>();
+    private RequestFuture sendListOffsetsRequests(final Map timestampsToSearch,
+                                                                    final boolean requireTimestamps) {
+        final Set partitionsToRetry = new HashSet<>();
+        Map> timestampsToSearchByNode =
+                groupListOffsetRequests(timestampsToSearch, partitionsToRetry);
+        if (timestampsToSearchByNode.isEmpty())
+            return RequestFuture.failure(new StaleMetadataException());
+
+        final RequestFuture listOffsetRequestsFuture = new RequestFuture<>();
+        final Map fetchedTimestampOffsets = new HashMap<>();
+        final AtomicInteger remainingResponses = new AtomicInteger(timestampsToSearchByNode.size());
+
+        for (Map.Entry> entry : timestampsToSearchByNode.entrySet()) {
+            RequestFuture future =
+                sendListOffsetRequest(entry.getKey(), entry.getValue(), requireTimestamps);
+            future.addListener(new RequestFutureListener() {
+                @Override
+                public void onSuccess(ListOffsetResult partialResult) {
+                    synchronized (listOffsetRequestsFuture) {
+                        fetchedTimestampOffsets.putAll(partialResult.fetchedOffsets);
+                        partitionsToRetry.addAll(partialResult.partitionsToRetry);
+
+                        if (remainingResponses.decrementAndGet() == 0 && !listOffsetRequestsFuture.isDone()) {
+                            ListOffsetResult result = new ListOffsetResult(fetchedTimestampOffsets, partitionsToRetry);
+                            listOffsetRequestsFuture.complete(result);
+                        }
+                    }
+                }
+
+                @Override
+                public void onFailure(RuntimeException e) {
+                    synchronized (listOffsetRequestsFuture) {
+                        if (!listOffsetRequestsFuture.isDone())
+                            listOffsetRequestsFuture.raise(e);
+                    }
+                }
+            });
+        }
+        return listOffsetRequestsFuture;
+    }
+
+    /**
+     * Groups timestamps to search by node for topic partitions in `timestampsToSearch` that have
+     * leaders available. Topic partitions from `timestampsToSearch` that do not have their leader
+     * available are added to `partitionsToRetry`
+     * @param timestampsToSearch The mapping from partitions ot the target timestamps
+     * @param partitionsToRetry A set of topic partitions that will be extended with partitions
+     *                          that need metadata update or re-connect to the leader.
+     */
+    private Map> groupListOffsetRequests(
+            Map timestampsToSearch,
+            Set partitionsToRetry) {
+        final Map partitionDataMap = new HashMap<>();
         for (Map.Entry entry: timestampsToSearch.entrySet()) {
             TopicPartition tp  = entry.getKey();
-            PartitionInfo info = metadata.fetch().partition(tp);
-            if (info == null) {
-                metadata.add(tp.topic());
-                log.debug("Partition {} is unknown for fetching offset, wait for metadata refresh", tp);
-                return RequestFuture.staleMetadata();
-            } else if (info.leader() == null) {
-                log.debug("Leader for partition {} unavailable for fetching offset, wait for metadata refresh", tp);
-                return RequestFuture.leaderNotAvailable();
+            Long offset = entry.getValue();
+            Metadata.LeaderAndEpoch leaderAndEpoch = metadata.currentLeader(tp);
+
+            if (!leaderAndEpoch.leader.isPresent()) {
+                log.debug("Leader for partition {} is unknown for fetching offset {}", tp, offset);
+                metadata.requestUpdate();
+                partitionsToRetry.add(tp);
             } else {
-                Node node = info.leader();
-                Map topicData = timestampsToSearchByNode.get(node);
-                if (topicData == null) {
-                    topicData = new HashMap<>();
-                    timestampsToSearchByNode.put(node, topicData);
+                Node leader = leaderAndEpoch.leader.get();
+                if (client.isUnavailable(leader)) {
+                    client.maybeThrowAuthFailure(leader);
+
+                    // The connection has failed and we need to await the backoff period before we can
+                    // try again. No need to request a metadata update since the disconnect will have
+                    // done so already.
+                    log.debug("Leader {} for partition {} is unavailable for fetching offset until reconnect backoff expires",
+                            leader, tp);
+                    partitionsToRetry.add(tp);
+                } else {
+                    int currentLeaderEpoch = leaderAndEpoch.epoch.orElse(ListOffsetResponse.UNKNOWN_EPOCH);
+                    partitionDataMap.put(tp, new ListOffsetPartition()
+                            .setPartitionIndex(tp.partition())
+                            .setTimestamp(offset)
+                            .setCurrentLeaderEpoch(currentLeaderEpoch));
                 }
-                topicData.put(entry.getKey(), entry.getValue());
             }
         }
-
-        final RequestFuture> listOffsetRequestsFuture = new RequestFuture<>();
-        final Map fetchedTimestampOffsets = new HashMap<>();
-        final AtomicInteger remainingResponses = new AtomicInteger(timestampsToSearchByNode.size());
-        for (Map.Entry> entry : timestampsToSearchByNode.entrySet()) {
-            sendListOffsetRequest(entry.getKey(), entry.getValue(), requireTimestamps)
-                    .addListener(new RequestFutureListener>() {
-                        @Override
-                        public void onSuccess(Map value) {
-                            synchronized (listOffsetRequestsFuture) {
-                                fetchedTimestampOffsets.putAll(value);
-                                if (remainingResponses.decrementAndGet() == 0 && !listOffsetRequestsFuture.isDone())
-                                    listOffsetRequestsFuture.complete(fetchedTimestampOffsets);
-                            }
-                        }
-
-                        @Override
-                        public void onFailure(RuntimeException e) {
-                            synchronized (listOffsetRequestsFuture) {
-                                // This may cause all the requests to be retried, but should be rare.
-                                if (!listOffsetRequestsFuture.isDone())
-                                    listOffsetRequestsFuture.raise(e);
-                            }
-                        }
-                    });
-        }
-        return listOffsetRequestsFuture;
+        return regroupPartitionMapByNode(partitionDataMap);
     }
 
     /**
@@ -661,28 +974,27 @@ public void onFailure(RuntimeException e) {
      * @param requireTimestamp  True if we require a timestamp in the response.
      * @return A response which can be polled to obtain the corresponding timestamps and offsets.
      */
-    private RequestFuture> sendListOffsetRequest(final Node node,
-                                                                                 final Map timestampsToSearch,
-                                                                                 boolean requireTimestamp) {
+    private RequestFuture sendListOffsetRequest(final Node node,
+                                                                  final Map timestampsToSearch,
+                                                                  boolean requireTimestamp) {
         ListOffsetRequest.Builder builder = ListOffsetRequest.Builder
                 .forConsumer(requireTimestamp, isolationLevel)
-                .setTargetTimes(timestampsToSearch);
+                .setTargetTimes(ListOffsetRequest.toListOffsetTopics(timestampsToSearch));
 
-        log.trace("Sending ListOffsetRequest {} to broker {}", builder, node);
+        log.debug("Sending ListOffsetRequest {} to broker {}", builder, node);
         return client.send(node, builder)
-                .compose(new RequestFutureAdapter>() {
+                .compose(new RequestFutureAdapter() {
                     @Override
-                    public void onSuccess(ClientResponse response, RequestFuture> future) {
+                    public void onSuccess(ClientResponse response, RequestFuture future) {
                         ListOffsetResponse lor = (ListOffsetResponse) response.responseBody();
                         log.trace("Received ListOffsetResponse {} from broker {}", lor, node);
-                        handleListOffsetResponse(timestampsToSearch, lor, future);
+                        handleListOffsetResponse(lor, future);
                     }
                 });
     }
 
     /**
      * Callback for the response of the list offset call above.
-     * @param timestampsToSearch The mapping from partitions to target timestamps
      * @param listOffsetResponse The response from the server.
      * @param future The future to be completed when the response returns. Note that any partition-level errors will
      *               generally fail the entire future result. The one exception is UNSUPPORTED_FOR_MESSAGE_FORMAT,
@@ -691,153 +1003,257 @@ public void onSuccess(ClientResponse response, RequestFuture timestampsToSearch,
-                                          ListOffsetResponse listOffsetResponse,
-                                          RequestFuture> future) {
-        Map timestampOffsetMap = new HashMap<>();
-        for (Map.Entry entry : timestampsToSearch.entrySet()) {
-            TopicPartition topicPartition = entry.getKey();
-            ListOffsetResponse.PartitionData partitionData = listOffsetResponse.responseData().get(topicPartition);
-            Errors error = partitionData.error;
-            if (error == Errors.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);
-                        timestampOffsetMap.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);
-                        timestampOffsetMap.put(topicPartition, offsetData);
-                    }
+    private void handleListOffsetResponse(ListOffsetResponse listOffsetResponse,
+                                          RequestFuture future) {
+        Map fetchedOffsets = new HashMap<>();
+        Set partitionsToRetry = new HashSet<>();
+        Set unauthorizedTopics = new HashSet<>();
+
+        for (ListOffsetTopicResponse topic : listOffsetResponse.topics()) {
+            for (ListOffsetPartitionResponse partition : topic.partitions()) {
+                TopicPartition topicPartition = new TopicPartition(topic.name(), partition.partitionIndex());
+                Errors error = Errors.forCode(partition.errorCode());
+                switch (error) {
+                    case NONE:
+                        if (!partition.oldStyleOffsets().isEmpty()) {
+                            // Handle v0 response with offsets
+                            long offset;
+                            if (partition.oldStyleOffsets().size() > 1) {
+                                future.raise(new IllegalStateException("Unexpected partitionData response of length " +
+                                        partition.oldStyleOffsets().size()));
+                                return;
+                            } else {
+                                offset = partition.oldStyleOffsets().get(0);
+                            }
+                            log.debug("Handling v0 ListOffsetResponse response for {}. Fetched offset {}",
+                                topicPartition, offset);
+                            if (offset != ListOffsetResponse.UNKNOWN_OFFSET) {
+                                ListOffsetData offsetData = new ListOffsetData(offset, null, Optional.empty());
+                                fetchedOffsets.put(topicPartition, offsetData);
+                            }
+                        } else {
+                            // Handle v1 and later response or v0 without offsets
+                            log.debug("Handling ListOffsetResponse response for {}. Fetched offset {}, timestamp {}",
+                                topicPartition, partition.offset(), partition.timestamp());
+                            if (partition.offset() != ListOffsetResponse.UNKNOWN_OFFSET) {
+                                Optional leaderEpoch = (partition.leaderEpoch() == ListOffsetResponse.UNKNOWN_EPOCH)
+                                        ? Optional.empty()
+                                        : Optional.of(partition.leaderEpoch());
+                                ListOffsetData offsetData = new ListOffsetData(partition.offset(), partition.timestamp(),
+                                    leaderEpoch);
+                                fetchedOffsets.put(topicPartition, offsetData);
+                            }
+                        }
+                        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_OR_FOLLOWER:
+                    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);
+                        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);
                 }
-            } else if (error == Errors.UNSUPPORTED_FOR_MESSAGE_FORMAT) {
-                // The message format on the broker side is before 0.10.0, we simply put null in the response.
-                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) {
-                log.debug("Attempt to fetch offsets for partition {} failed due to obsolete leadership information, retrying.",
-                        topicPartition);
-                future.raise(error);
-                return;
-            } else if (error == Errors.UNKNOWN_TOPIC_OR_PARTITION) {
-                log.warn("Received unknown topic or partition error in ListOffset request for partition {}. The topic/partition " +
-                        "may not exist or the user may not have Describe access to it.", topicPartition);
-                future.raise(error);
-                return;
-            } else {
-                log.warn("Attempt to fetch offsets for partition {} failed due to: {}", topicPartition, error.message());
-                future.raise(new StaleMetadataException());
-                return;
             }
         }
-        if (!future.isDone())
-            future.complete(timestampOffsetMap);
+
+        if (!unauthorizedTopics.isEmpty())
+            future.raise(new TopicAuthorizationException(unauthorizedTopics));
+        else
+            future.complete(new ListOffsetResult(fetchedOffsets, partitionsToRetry));
+    }
+
+    static class ListOffsetResult {
+        private final Map fetchedOffsets;
+        private final Set partitionsToRetry;
+
+        ListOffsetResult(Map fetchedOffsets, Set partitionsNeedingRetry) {
+            this.fetchedOffsets = fetchedOffsets;
+            this.partitionsToRetry = partitionsNeedingRetry;
+        }
+
+        ListOffsetResult() {
+            this.fetchedOffsets = new HashMap<>();
+            this.partitionsToRetry = new HashSet<>();
+        }
     }
 
     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;
+        }
+    }
+
+    /**
+     * If we have seen new metadata (as tracked by {@link org.apache.kafka.clients.Metadata#updateVersion()}), then
+     * we should check that all of the assignments have a valid position.
+     */
+    private void validatePositionsOnMetadataChange() {
+        int newMetadataUpdateVersion = metadata.updateVersion();
+        if (metadataUpdateVersion.getAndSet(newMetadataUpdateVersion) != newMetadataUpdateVersion) {
+            subscriptions.assignedPartitions().forEach(topicPartition -> {
+                ConsumerMetadata.LeaderAndEpoch leaderAndEpoch = metadata.currentLeader(topicPartition);
+                subscriptions.maybeValidatePositionForCurrentLeader(apiVersions, topicPartition, leaderAndEpoch);
+            });
+        }
     }
 
     /**
      * Create fetch requests for all nodes for which we have assigned partitions
      * that have no existing requests in flight.
      */
-    private Map createFetchRequests() {
-        // create the fetch info
-        Cluster cluster = metadata.fetch();
-        Map> fetchable = new LinkedHashMap<>();
+    private Map prepareFetchRequests() {
+        Map fetchable = new LinkedHashMap<>();
+
+        validatePositionsOnMetadataChange();
+
+        long currentTimeMs = time.milliseconds();
+
         for (TopicPartition partition : fetchablePartitions()) {
-            Node node = cluster.leaderFor(partition);
-            if (node == null) {
+            FetchPosition position = this.subscriptions.position(partition);
+            if (position == null) {
+                throw new IllegalStateException("Missing position for fetchable partition " + partition);
+            }
+
+            Optional leaderOpt = position.currentLeader.leader;
+            if (!leaderOpt.isPresent()) {
+                log.debug("Requesting metadata update for partition {} since the position {} is missing the current leader node", partition, position);
                 metadata.requestUpdate();
-            } else if (!this.client.hasPendingRequests(node)) {
+                continue;
+            }
+
+            // Use the preferred read replica if set, otherwise the position's leader
+            Node node = selectReadReplica(partition, leaderOpt.get(), currentTimeMs);
+            if (client.isUnavailable(node)) {
+                client.maybeThrowAuthFailure(node);
+
+                // If we try to send during the reconnect backoff 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 (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
-                LinkedHashMap fetch = fetchable.get(node);
-                if (fetch == null) {
-                    fetch = new LinkedHashMap<>();
-                    fetchable.put(node, fetch);
+                FetchSessionHandler.Builder builder = fetchable.get(node);
+                if (builder == null) {
+                    int id = node.id();
+                    FetchSessionHandler handler = sessionHandler(id);
+                    if (handler == null) {
+                        handler = new FetchSessionHandler(logContext, id);
+                        sessionHandlers.put(id, handler);
+                    }
+                    builder = handler.newBuilder();
+                    fetchable.put(node, builder);
                 }
 
-                long position = this.subscriptions.position(partition);
-                fetch.put(partition, new FetchRequest.PartitionData(position, FetchRequest.INVALID_LOG_START_OFFSET,
-                        this.fetchSize));
-                log.debug("Added {} fetch request for partition {} at offset {} to node {}", isolationLevel,
-                        partition, position, node);
-            } else {
-                log.trace("Skipping fetch for partition {} because there is an in-flight request to {}", partition, node);
+                builder.add(partition, new FetchRequest.PartitionData(position.offset,
+                    FetchRequest.INVALID_LOG_START_OFFSET, this.fetchSize,
+                    position.currentLeader.epoch, Optional.empty()));
+
+                log.debug("Added {} fetch request for partition {} at position {} to node {}", isolationLevel,
+                    partition, position, node);
             }
         }
 
-        // create the fetches
-        Map requests = new HashMap<>();
-        for (Map.Entry> entry : fetchable.entrySet()) {
-            Node node = entry.getKey();
-            FetchRequest.Builder fetch = FetchRequest.Builder.forConsumer(this.maxWaitMs, this.minBytes,
-                    entry.getValue(), isolationLevel)
-                    .setMaxBytes(this.maxBytes);
-            requests.put(node, fetch);
+        Map reqs = new LinkedHashMap<>();
+        for (Map.Entry entry : fetchable.entrySet()) {
+            reqs.put(entry.getKey(), entry.getValue().build());
         }
-        return requests;
+        return reqs;
+    }
+
+    private Map> regroupFetchPositionsByLeader(
+            Map partitionMap) {
+        return partitionMap.entrySet()
+                .stream()
+                .filter(entry -> entry.getValue().currentLeader.leader.isPresent())
+                .collect(Collectors.groupingBy(entry -> entry.getValue().currentLeader.leader.get(),
+                        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;
-        Errors error = partition.error;
+    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) {
+                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;
                 }
 
                 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);
+                        partition.records().sizeInBytes(), tp, position);
+                Iterator batches = partition.records().batches().iterator();
+                completedFetch = nextCompletedFetch;
 
-                if (!batches.hasNext() && partition.records.sizeInBytes() > 0) {
+                if (!batches.hasNext() && partition.records().sizeInBytes() > 0) {
                     if (completedFetch.responseVersion < 3) {
                         // Implement the pre KIP-74 behavior of throwing a RecordTooLargeException.
                         Map recordTooLargePartitions = Collections.singletonMap(tp, fetchOffset);
@@ -855,44 +1271,80 @@ private PartitionRecords parseCompletedFetch(CompletedFetch completedFetch) {
                     }
                 }
 
-                if (partition.highWatermark >= 0) {
-                    log.trace("Updating high watermark for partition {} to {}", tp, partition.highWatermark);
-                    subscriptions.updateHighWatermark(tp, partition.highWatermark);
+                if (partition.highWatermark() >= 0) {
+                    log.trace("Updating high watermark for partition {} to {}", tp, partition.highWatermark());
+                    subscriptions.updateHighWatermark(tp, partition.highWatermark());
+                }
+
+                if (partition.logStartOffset() >= 0) {
+                    log.trace("Updating log start offset for partition {} to {}", tp, partition.logStartOffset());
+                    subscriptions.updateLogStartOffset(tp, partition.logStartOffset());
+                }
+
+                if (partition.lastStableOffset() >= 0) {
+                    log.trace("Updating last stable offset for partition {} to {}", tp, partition.lastStableOffset());
+                    subscriptions.updateLastStableOffset(tp, partition.lastStableOffset());
                 }
 
-                if (partition.lastStableOffset >= 0) {
-                    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;
+                    });
                 }
-            } else if (error == Errors.NOT_LEADER_FOR_PARTITION ||
-                       error == Errors.KAFKA_STORAGE_ERROR) {
+
+                nextCompletedFetch.initialized = true;
+            } else if (error == Errors.NOT_LEADER_OR_FOLLOWER ||
+                       error == Errors.REPLICA_NOT_AVAILABLE ||
+                       error == Errors.KAFKA_STORAGE_ERROR ||
+                       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 {}. The topic/partition " +
-                        "may not exist or the user may not have Describe access to it", tp);
+                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.needOffsetReset(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
+                    FetchPosition position = subscriptions.position(tp);
+                    if (position == null || fetchOffset != position.offset) {
+                        log.debug("Discarding stale fetch response for partition {} since the fetched offset {} " +
+                                "does not match the current offset {}", tp, fetchOffset, position);
+                    } else {
+                        handleOffsetOutOfRange(position, tp);
+                    }
                 } 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) {
-                log.warn("Not authorized to read from topic {}.", tp.topic());
+                //we log the actual partition and not just the topic to help with ACL propagation issues in large clusters
+                log.warn("Not authorized to read from partition {}.", tp);
                 throw new TopicAuthorizationException(Collections.singleton(tp.topic()));
+            } else if (error == Errors.UNKNOWN_LEADER_EPOCH) {
+                log.debug("Received unknown leader epoch error in fetch for partition {}", tp);
             } else if (error == Errors.UNKNOWN_SERVER_ERROR) {
-                log.warn("Unknown error fetching data for topic-partition {}", tp);
+                log.warn("Unknown server error while fetching offset {} for topic-partition {}",
+                        fetchOffset, tp);
+            } else if (error == Errors.CORRUPT_MESSAGE) {
+                throw new KafkaException("Encountered corrupt message when fetching offset "
+                        + fetchOffset
+                        + " for topic-partition "
+                        + tp);
             } else {
-                throw new IllegalStateException("Unexpected error code " + error.code() + " while fetching data");
+                throw new IllegalStateException("Unexpected error code "
+                        + error.code()
+                        + " while fetching at offset "
+                        + fetchOffset
+                        + " from topic-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
@@ -900,9 +1352,20 @@ private PartitionRecords parseCompletedFetch(CompletedFetch completedFetch) {
                 subscriptions.movePartitionToEnd(tp);
         }
 
-        return partitionRecords;
+        return completedFetch;
     }
 
+    private void handleOffsetOutOfRange(FetchPosition fetchPosition, TopicPartition topicPartition) {
+        String errorMessage = "Fetch position " + fetchPosition + " is out of range for partition " + topicPartition;
+        if (subscriptions.hasDefaultOffsetResetPolicy()) {
+            log.info("{}, resetting offset", errorMessage);
+            subscriptions.requestOffsetReset(topicPartition);
+        } else {
+            log.info("{}, raising error to the application since no reset policy is configured", errorMessage);
+            throw new OffsetOutOfRangeException(errorMessage,
+                Collections.singletonMap(topicPartition, fetchPosition.offset));
+        }
+    }
 
     /**
      * Parse the record entry, deserializing the key / value fields if necessary
@@ -913,6 +1376,7 @@ private ConsumerRecord parseRecord(TopicPartition partition,
         try {
             long offset = record.offset();
             long timestamp = record.timestamp();
+            Optional leaderEpoch = maybeLeaderEpoch(batch.partitionLeaderEpoch());
             TimestampType timestampType = batch.timestampType();
             Headers headers = new RecordHeaders(record.headers());
             ByteBuffer keyBytes = record.key();
@@ -925,16 +1389,57 @@ private ConsumerRecord parseRecord(TopicPartition partition,
                                         timestamp, timestampType, record.checksumOrNull(),
                                         keyByteArray == null ? ConsumerRecord.NULL_SIZE : keyByteArray.length,
                                         valueByteArray == null ? ConsumerRecord.NULL_SIZE : valueByteArray.length,
-                                        key, value, headers);
+                                        key, value, headers, leaderEpoch);
         } catch (RuntimeException e) {
             throw new SerializationException("Error deserializing key/value for partition " + partition +
                     " at offset " + record.offset() + ". If needed, please seek past the record to continue consumption.", e);
         }
     }
 
-    @Override
-    public void onAssignment(Set assignment) {
-        sensors.updatePartitionLagSensors(assignment);
+    private Optional maybeLeaderEpoch(int leaderEpoch) {
+        return leaderEpoch == RecordBatch.NO_PARTITION_LEADER_EPOCH ? Optional.empty() : Optional.of(leaderEpoch);
+    }
+
+    /**
+     * 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 completedFetchesItr = completedFetches.iterator();
+        while (completedFetchesItr.hasNext()) {
+            CompletedFetch records = completedFetchesItr.next();
+            TopicPartition tp = records.partition;
+            if (!assignedPartitions.contains(tp)) {
+                records.drain();
+                completedFetchesItr.remove();
+            }
+        }
+
+        if (nextInLineFetch != null && !assignedPartitions.contains(nextInLineFetch.partition)) {
+            nextInLineFetch.drain();
+            nextInLineFetch = null;
+        }
+    }
+
+    /**
+     * Clear the buffered data which are not a part of newly assigned topics
+     *
+     * @param assignedTopics  newly assigned topics
+     */
+    public void clearBufferedDataForUnassignedTopics(Collection assignedTopics) {
+        Set currentTopicPartitions = new HashSet<>();
+        for (TopicPartition tp : subscriptions.assignedPartitions()) {
+            if (assignedTopics.contains(tp.topic())) {
+                currentTopicPartitions.add(tp);
+            }
+        }
+        clearBufferedDataForUnassignedPartitions(currentTopicPartitions);
+    }
+
+    // Visible for testing
+    protected FetchSessionHandler sessionHandler(int node) {
+        return sessionHandlers.get(node);
     }
 
     public static Sensor throttleTimeSensor(Metrics metrics, FetcherMetricsRegistry metricsRegistry) {
@@ -946,12 +1451,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;
@@ -959,27 +1466,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).
@@ -992,7 +1507,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());
                 }
@@ -1003,7 +1518,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());
                 }
@@ -1035,6 +1550,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()) {
@@ -1082,7 +1600,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<>();
@@ -1133,20 +1651,14 @@ private boolean isBatchAborted(RecordBatch batch) {
             return batch.isTransactional() && abortedProducerIds.contains(batch.producerId());
         }
 
-        private PriorityQueue abortedTransactions(FetchResponse.PartitionData partition) {
-            if (partition.abortedTransactions == null || partition.abortedTransactions.isEmpty())
+        private PriorityQueue abortedTransactions(FetchResponse.PartitionData partition) {
+            if (partition.abortedTransactions() == null || partition.abortedTransactions().isEmpty())
                 return null;
 
             PriorityQueue abortedTransactions = new PriorityQueue<>(
-                    partition.abortedTransactions.size(),
-                    new Comparator() {
-                        @Override
-                        public int compare(FetchResponse.AbortedTransaction o1, FetchResponse.AbortedTransaction o2) {
-                            return Long.compare(o1.firstOffset, o2.firstOffset);
-                        }
-                    }
+                    partition.abortedTransactions().size(), Comparator.comparingLong(o -> o.firstOffset)
             );
-            abortedTransactions.addAll(partition.abortedTransactions);
+            abortedTransactions.addAll(partition.abortedTransactions());
             return abortedTransactions;
         }
 
@@ -1156,31 +1668,14 @@ private boolean containsAbortMarker(RecordBatch batch) {
 
             Iterator batchIterator = batch.iterator();
             if (!batchIterator.hasNext())
-                throw new InvalidRecordException("Invalid batch for partition " + partition + " at offset " +
-                        batch.baseOffset() + " with control sequence set, but no records");
+                return false;
 
             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;
         }
     }
 
@@ -1221,8 +1716,8 @@ public void record(TopicPartition partition, int bytes, int records) {
 
             if (this.unrecordedPartitions.isEmpty()) {
                 // once all expected partitions from the fetch have reported in, record the metrics
-                this.sensors.bytesFetched.record(topicFetchMetric.fetchBytes);
-                this.sensors.recordsFetched.record(topicFetchMetric.fetchRecords);
+                this.sensors.bytesFetched.record(this.fetchMetrics.fetchBytes);
+                this.sensors.recordsFetched.record(this.fetchMetrics.fetchRecords);
 
                 // also record per-topic metrics
                 for (Map.Entry entry: this.topicFetchMetrics.entrySet()) {
@@ -1250,8 +1745,10 @@ private static class FetchManagerMetrics {
         private final Sensor recordsFetched;
         private final Sensor fetchLatency;
         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;
@@ -1271,11 +1768,14 @@ 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");
             this.recordsFetchLag.add(metrics.metricInstance(metricsRegistry.recordsLagMax), new Max());
+
+            this.recordsFetchLead = metrics.sensor("records-lead");
+            this.recordsFetchLead.add(metrics.metricInstance(metricsRegistry.recordsLeadMin), new Min());
         }
 
         private void recordTopicFetchMetrics(String topic, int bytes, int records) {
@@ -1311,14 +1811,50 @@ private void recordTopicFetchMetrics(String topic, int bytes, int records) {
             recordsFetched.record(records);
         }
 
-        private void updatePartitionLagSensors(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;
+            }
+        }
+
+        private void recordPartitionLead(TopicPartition tp, long lead) {
+            this.recordsFetchLead.record(lead);
+
+            String name = partitionLeadMetricName(tp);
+            Sensor recordsLead = this.metrics.getSensor(name);
+            if (recordsLead == null) {
+                Map metricTags = topicPartitionTags(tp);
+
+                recordsLead = this.metrics.sensor(name);
+
+                recordsLead.add(this.metrics.metricInstance(metricsRegistry.partitionRecordsLead, metricTags), new Value());
+                recordsLead.add(this.metrics.metricInstance(metricsRegistry.partitionRecordsLeadMin, metricTags), new Min());
+                recordsLead.add(this.metrics.metricInstance(metricsRegistry.partitionRecordsLeadAvg, metricTags), new Avg());
             }
-            this.assignedPartitions = assignedPartitions;
+            recordsLead.record(lead);
         }
 
         private void recordPartitionLag(TopicPartition tp, long lag) {
@@ -1327,16 +1863,12 @@ private void recordPartitionLag(TopicPartition tp, long lag) {
             String name = partitionLagMetricName(tp);
             Sensor recordsLag = this.metrics.getSensor(name);
             if (recordsLag == null) {
+                Map metricTags = topicPartitionTags(tp);
                 recordsLag = this.metrics.sensor(name);
-                recordsLag.add(this.metrics.metricName(name,
-                        metricsRegistry.partitionRecordsLag.group(),
-                        metricsRegistry.partitionRecordsLag.description()), new Value());
-                recordsLag.add(this.metrics.metricName(name + "-max",
-                        metricsRegistry.partitionRecordsLagMax.group(),
-                        metricsRegistry.partitionRecordsLagMax.description()), new Max());
-                recordsLag.add(this.metrics.metricName(name + "-avg",
-                        metricsRegistry.partitionRecordsLagAvg.group(),
-                        metricsRegistry.partitionRecordsLagAvg.description()), new Avg());
+
+                recordsLag.add(this.metrics.metricInstance(metricsRegistry.partitionRecordsLag, metricTags), new Value());
+                recordsLag.add(this.metrics.metricInstance(metricsRegistry.partitionRecordsLagMax, metricTags), new Max());
+                recordsLag.add(this.metrics.metricInstance(metricsRegistry.partitionRecordsLagAvg, metricTags), new Avg());
             }
             recordsLag.record(lag);
         }
@@ -1344,13 +1876,33 @@ private void recordPartitionLag(TopicPartition tp, long lag) {
         private static String partitionLagMetricName(TopicPartition tp) {
             return tp + ".records-lag";
         }
+
+        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();
     }
 
+    private Set topicsForPartitions(Collection partitions) {
+        return partitions.stream().map(TopicPartition::topic).collect(Collectors.toSet());
+    }
+
 }
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 acf42ec339f25..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
@@ -18,6 +18,7 @@
 
 import java.util.Arrays;
 import java.util.HashSet;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Set;
 
@@ -37,6 +38,7 @@ public class FetcherMetricsRegistry {
     public MetricNameTemplate fetchRequestRate;
     public MetricNameTemplate fetchRequestTotal;
     public MetricNameTemplate recordsLagMax;
+    public MetricNameTemplate recordsLeadMin;
     public MetricNameTemplate fetchThrottleTimeAvg;
     public MetricNameTemplate fetchThrottleTimeMax;
     public MetricNameTemplate topicFetchSizeAvg;
@@ -49,6 +51,10 @@ public class FetcherMetricsRegistry {
     public MetricNameTemplate partitionRecordsLag;
     public MetricNameTemplate partitionRecordsLagMax;
     public MetricNameTemplate partitionRecordsLagAvg;
+    public MetricNameTemplate partitionRecordsLead;
+    public MetricNameTemplate partitionRecordsLeadMin;
+    public MetricNameTemplate partitionRecordsLeadAvg;
+    public MetricNameTemplate partitionPreferredReadReplica;
 
     public FetcherMetricsRegistry() {
         this(new HashSet(), "");
@@ -59,75 +65,86 @@ public FetcherMetricsRegistry(String metricGrpPrefix) {
     }
 
     public FetcherMetricsRegistry(Set tags, String metricGrpPrefix) {
-        
+
         /***** Client level *****/
         String groupName = metricGrpPrefix + "-fetch-manager-metrics";
-                
-        this.fetchSizeAvg = new MetricNameTemplate("fetch-size-avg", groupName, 
+
+        this.fetchSizeAvg = new MetricNameTemplate("fetch-size-avg", groupName,
                 "The average number of bytes fetched per request", tags);
 
-        this.fetchSizeMax = new MetricNameTemplate("fetch-size-max", groupName, 
+        this.fetchSizeMax = new MetricNameTemplate("fetch-size-max", groupName,
                 "The maximum number of bytes fetched per request", tags);
-        this.bytesConsumedRate = new MetricNameTemplate("bytes-consumed-rate", groupName, 
+        this.bytesConsumedRate = new MetricNameTemplate("bytes-consumed-rate", groupName,
                 "The average number of bytes consumed per second", tags);
         this.bytesConsumedTotal = new MetricNameTemplate("bytes-consumed-total", groupName,
                 "The total number of bytes consumed", tags);
 
-        this.recordsPerRequestAvg = new MetricNameTemplate("records-per-request-avg", groupName, 
+        this.recordsPerRequestAvg = new MetricNameTemplate("records-per-request-avg", groupName,
                 "The average number of records in each request", tags);
-        this.recordsConsumedRate = new MetricNameTemplate("records-consumed-rate", groupName, 
+        this.recordsConsumedRate = new MetricNameTemplate("records-consumed-rate", groupName,
                 "The average number of records consumed per second", tags);
         this.recordsConsumedTotal = new MetricNameTemplate("records-consumed-total", groupName,
                 "The total number of records consumed", tags);
 
-        this.fetchLatencyAvg = new MetricNameTemplate("fetch-latency-avg", groupName, 
+        this.fetchLatencyAvg = new MetricNameTemplate("fetch-latency-avg", groupName,
                 "The average time taken for a fetch request.", tags);
-        this.fetchLatencyMax = new MetricNameTemplate("fetch-latency-max", groupName, 
+        this.fetchLatencyMax = new MetricNameTemplate("fetch-latency-max", groupName,
                 "The max time taken for any fetch request.", tags);
-        this.fetchRequestRate = new MetricNameTemplate("fetch-rate", groupName, 
+        this.fetchRequestRate = new MetricNameTemplate("fetch-rate", groupName,
                 "The number of fetch requests per second.", tags);
         this.fetchRequestTotal = new MetricNameTemplate("fetch-total", groupName,
                 "The total number of fetch requests.", tags);
 
-        this.recordsLagMax = new MetricNameTemplate("records-lag-max", groupName, 
+        this.recordsLagMax = new MetricNameTemplate("records-lag-max", groupName,
                 "The maximum lag in terms of number of records for any partition in this window", tags);
+        this.recordsLeadMin = new MetricNameTemplate("records-lead-min", groupName,
+                "The minimum lead in terms of number of records for any partition in this window", tags);
 
-        this.fetchThrottleTimeAvg = new MetricNameTemplate("fetch-throttle-time-avg", groupName, 
+        this.fetchThrottleTimeAvg = new MetricNameTemplate("fetch-throttle-time-avg", groupName,
                 "The average throttle time in ms", tags);
-        this.fetchThrottleTimeMax = new MetricNameTemplate("fetch-throttle-time-max", groupName, 
+        this.fetchThrottleTimeMax = new MetricNameTemplate("fetch-throttle-time-max", groupName,
                 "The maximum throttle time in ms", tags);
 
         /***** Topic level *****/
-        Set topicTags = new HashSet<>(tags);
+        Set topicTags = new LinkedHashSet<>(tags);
         topicTags.add("topic");
 
-        this.topicFetchSizeAvg = new MetricNameTemplate("fetch-size-avg", groupName, 
+        this.topicFetchSizeAvg = new MetricNameTemplate("fetch-size-avg", groupName,
                 "The average number of bytes fetched per request for a topic", topicTags);
-        this.topicFetchSizeMax = new MetricNameTemplate("fetch-size-max", groupName, 
+        this.topicFetchSizeMax = new MetricNameTemplate("fetch-size-max", groupName,
                 "The maximum number of bytes fetched per request for a topic", topicTags);
-        this.topicBytesConsumedRate = new MetricNameTemplate("bytes-consumed-rate", groupName, 
+        this.topicBytesConsumedRate = new MetricNameTemplate("bytes-consumed-rate", groupName,
                 "The average number of bytes consumed per second for a topic", topicTags);
         this.topicBytesConsumedTotal = new MetricNameTemplate("bytes-consumed-total", groupName,
                 "The total number of bytes consumed for a topic", topicTags);
 
-        this.topicRecordsPerRequestAvg = new MetricNameTemplate("records-per-request-avg", groupName, 
+        this.topicRecordsPerRequestAvg = new MetricNameTemplate("records-per-request-avg", groupName,
                 "The average number of records in each request for a topic", topicTags);
-        this.topicRecordsConsumedRate = new MetricNameTemplate("records-consumed-rate", groupName, 
+        this.topicRecordsConsumedRate = new MetricNameTemplate("records-consumed-rate", groupName,
                 "The average number of records consumed per second for a topic", topicTags);
         this.topicRecordsConsumedTotal = new MetricNameTemplate("records-consumed-total", groupName,
                 "The total number of records consumed for a topic", topicTags);
-        
+
         /***** Partition level *****/
-        this.partitionRecordsLag = new MetricNameTemplate("{topic}-{partition}.records-lag", groupName, 
-                "The latest lag of the partition", tags);
-        this.partitionRecordsLagMax = new MetricNameTemplate("{topic}-{partition}.records-lag-max", groupName, 
-                "The max lag of the partition", tags);
-        this.partitionRecordsLagAvg = new MetricNameTemplate("{topic}-{partition}.records-lag-avg", groupName, 
-                "The average lag of the partition", tags);
-        
-    
+        Set partitionTags = new HashSet<>(topicTags);
+        partitionTags.add("partition");
+        this.partitionRecordsLag = new MetricNameTemplate("records-lag", groupName,
+                "The latest lag of the partition", partitionTags);
+        this.partitionRecordsLagMax = new MetricNameTemplate("records-lag-max", groupName,
+                "The max lag of the partition", partitionTags);
+        this.partitionRecordsLagAvg = new MetricNameTemplate("records-lag-avg", groupName,
+                "The average lag of the partition", partitionTags);
+        this.partitionRecordsLead = new MetricNameTemplate("records-lead", groupName,
+                "The latest lead of the partition", partitionTags);
+        this.partitionRecordsLeadMin = new MetricNameTemplate("records-lead-min", groupName,
+                "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() {
         return Arrays.asList(
             fetchSizeAvg,
@@ -142,6 +159,7 @@ public List getAllTemplates() {
             fetchRequestRate,
             fetchRequestTotal,
             recordsLagMax,
+            recordsLeadMin,
             fetchThrottleTimeAvg,
             fetchThrottleTimeMax,
             topicFetchSizeAvg,
@@ -153,7 +171,11 @@ public List getAllTemplates() {
             topicRecordsConsumedTotal,
             partitionRecordsLag,
             partitionRecordsLagAvg,
-            partitionRecordsLagMax
+            partitionRecordsLagMax,
+            partitionRecordsLead,
+            partitionRecordsLeadMin,
+            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 38a7c78a599c4..dfb9f85144d2d 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,89 +16,120 @@
  */
 package org.apache.kafka.clients.consumer.internals;
 
+import org.apache.kafka.clients.GroupRebalanceConfig;
+import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.common.utils.Timer;
+
+import org.slf4j.Logger;
+
 /**
  * A helper class for managing the heartbeat to the coordinator
  */
 public final class Heartbeat {
-    private final long sessionTimeout;
-    private final long heartbeatInterval;
-    private final long maxPollInterval;
-    private final long retryBackoffMs;
-
-    private volatile long lastHeartbeatSend; // volatile since it is read by metrics
-    private long lastHeartbeatReceive;
-    private long lastSessionReset;
-    private long lastPoll;
-    private boolean heartbeatFailed;
-
-    public Heartbeat(long sessionTimeout,
-                     long heartbeatInterval,
-                     long maxPollInterval,
-                     long retryBackoffMs) {
-        if (heartbeatInterval >= sessionTimeout)
+    private final int maxPollIntervalMs;
+    private final GroupRebalanceConfig rebalanceConfig;
+    private final Time time;
+    private final Timer heartbeatTimer;
+    private final Timer sessionTimer;
+    private final Timer pollTimer;
+    private final Logger log;
+
+    private volatile long lastHeartbeatSend = 0L;
+    private volatile boolean heartbeatInFlight = false;
+
+    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.heartbeatTimer = time.timer(config.heartbeatIntervalMs);
+        this.sessionTimer = time.timer(config.sessionTimeoutMs);
+        this.maxPollIntervalMs = config.rebalanceTimeoutMs;
+        this.pollTimer = time.timer(maxPollIntervalMs);
+
+        final LogContext logContext = new LogContext("[Heartbeat groupID=" + config.groupId + "] ");
+        this.log = logContext.logger(getClass());
+    }
 
-        this.sessionTimeout = sessionTimeout;
-        this.heartbeatInterval = heartbeatInterval;
-        this.maxPollInterval = maxPollInterval;
-        this.retryBackoffMs = retryBackoffMs;
+    private void update(long now) {
+        heartbeatTimer.update(now);
+        sessionTimer.update(now);
+        pollTimer.update(now);
     }
 
     public void poll(long now) {
-        this.lastPoll = now;
+        update(now);
+        pollTimer.reset(maxPollIntervalMs);
+    }
+
+    boolean hasInflight() {
+        return heartbeatInFlight;
     }
 
-    public void sentHeartbeat(long now) {
-        this.lastHeartbeatSend = now;
-        this.heartbeatFailed = false;
+    void sentHeartbeat(long now) {
+        lastHeartbeatSend = now;
+        heartbeatInFlight = true;
+        update(now);
+        heartbeatTimer.reset(rebalanceConfig.heartbeatIntervalMs);
+
+        if (log.isTraceEnabled()) {
+            log.trace("Sending heartbeat request with {}ms remaining on timer", heartbeatTimer.remainingMs());
+        }
     }
 
-    public void failHeartbeat() {
-        this.heartbeatFailed = true;
+    void failHeartbeat() {
+        update(time.milliseconds());
+        heartbeatInFlight = false;
+        heartbeatTimer.reset(rebalanceConfig.retryBackoffMs);
+
+        log.trace("Heartbeat failed, reset the timer to {}ms remaining", heartbeatTimer.remainingMs());
     }
 
-    public void receiveHeartbeat(long now) {
-        this.lastHeartbeatReceive = now;
+    void receiveHeartbeat() {
+        update(time.milliseconds());
+        heartbeatInFlight = false;
+        sessionTimer.reset(rebalanceConfig.sessionTimeoutMs);
     }
 
-    public boolean shouldHeartbeat(long now) {
-        return timeToNextHeartbeat(now) == 0;
+    boolean shouldHeartbeat(long now) {
+        update(now);
+        return heartbeatTimer.isExpired();
     }
     
-    public long lastHeartbeatSend() {
+    long lastHeartbeatSend() {
         return this.lastHeartbeatSend;
     }
 
-    public long timeToNextHeartbeat(long now) {
-        long timeSinceLastHeartbeat = now - Math.max(lastHeartbeatSend, lastSessionReset);
-        final long delayToNextHeartbeat;
-        if (heartbeatFailed)
-            delayToNextHeartbeat = retryBackoffMs;
-        else
-            delayToNextHeartbeat = heartbeatInterval;
-
-        if (timeSinceLastHeartbeat > delayToNextHeartbeat)
-            return 0;
-        else
-            return delayToNextHeartbeat - timeSinceLastHeartbeat;
+    long timeToNextHeartbeat(long now) {
+        update(now);
+        return heartbeatTimer.remainingMs();
     }
 
-    public boolean sessionTimeoutExpired(long now) {
-        return now - Math.max(lastSessionReset, lastHeartbeatReceive) > sessionTimeout;
+    boolean sessionTimeoutExpired(long now) {
+        update(now);
+        return sessionTimer.isExpired();
     }
 
-    public long interval() {
-        return heartbeatInterval;
+    void resetTimeouts() {
+        update(time.milliseconds());
+        sessionTimer.reset(rebalanceConfig.sessionTimeoutMs);
+        pollTimer.reset(maxPollIntervalMs);
+        heartbeatTimer.reset(rebalanceConfig.heartbeatIntervalMs);
     }
 
-    public void resetTimeouts(long now) {
-        this.lastSessionReset = now;
-        this.lastPoll = now;
-        this.heartbeatFailed = false;
+    void resetSessionTimeout() {
+        update(time.milliseconds());
+        sessionTimer.reset(rebalanceConfig.sessionTimeoutMs);
     }
 
-    public boolean pollTimeoutExpired(long now) {
-        return now - lastPoll > maxPollInterval;
+    boolean pollTimeoutExpired(long now) {
+        update(now);
+        return pollTimer.isExpired();
     }
 
-}
\ No newline at end of file
+    long lastPollTime() {
+        return pollTimer.currentTimeMs();
+    }
+}
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..0e750abc99b68
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/KafkaConsumerMetrics.java
@@ -0,0 +1,87 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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 MetricName lastPollMetricName;
+    private final Sensor timeBetweenPollSensor;
+    private final Sensor pollIdleSensor;
+    private final Metrics metrics;
+    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..57650f52fe35a
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/OffsetsForLeaderEpochClient.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.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.message.OffsetForLeaderEpochRequestData.OffsetForLeaderPartition;
+import org.apache.kafka.common.message.OffsetForLeaderEpochRequestData.OffsetForLeaderTopic;
+import org.apache.kafka.common.message.OffsetForLeaderEpochRequestData.OffsetForLeaderTopicCollection;
+import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData.EpochEndOffset;
+import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData.OffsetForLeaderTopicResult;
+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.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) {
+        OffsetForLeaderTopicCollection topics = new OffsetForLeaderTopicCollection(requestData.size());
+        requestData.forEach((topicPartition, fetchPosition) ->
+            fetchPosition.offsetEpoch.ifPresent(fetchEpoch -> {
+                OffsetForLeaderTopic topic = topics.find(topicPartition.topic());
+                if (topic == null) {
+                    topic = new OffsetForLeaderTopic().setTopic(topicPartition.topic());
+                    topics.add(topic);
+                }
+                topic.partitions().add(new OffsetForLeaderPartition()
+                    .setPartition(topicPartition.partition())
+                    .setLeaderEpoch(fetchEpoch)
+                    .setCurrentLeaderEpoch(fetchPosition.currentLeader.epoch
+                        .orElse(RecordBatch.NO_PARTITION_LEADER_EPOCH))
+                );
+            })
+        );
+        return OffsetsForLeaderEpochRequest.Builder.forConsumer(topics);
+    }
+
+    @Override
+    protected OffsetForEpochResult handleResponse(
+            Node node,
+            Map requestData,
+            OffsetsForLeaderEpochResponse response) {
+
+        Set partitionsToRetry = new HashSet<>(requestData.keySet());
+        Set unauthorizedTopics = new HashSet<>();
+        Map endOffsets = new HashMap<>();
+
+        for (OffsetForLeaderTopicResult topic : response.data().topics()) {
+            for (EpochEndOffset partition : topic.partitions()) {
+                TopicPartition topicPartition = new TopicPartition(topic.topic(), partition.partition());
+
+                if (!requestData.containsKey(topicPartition)) {
+                    logger().warn("Received unrequested topic or partition {} from response, ignoring.", topicPartition);
+                    continue;
+                }
+
+                Errors error = Errors.forCode(partition.errorCode());
+                switch (error) {
+                    case NONE:
+                        logger().debug("Handling OffsetsForLeaderEpoch response for {}. Got offset {} for epoch {}.",
+                            topicPartition, partition.endOffset(), partition.leaderEpoch());
+                        endOffsets.put(topicPartition, partition);
+                        partitionsToRetry.remove(topicPartition);
+                        break;
+                    case NOT_LEADER_OR_FOLLOWER:
+                    case REPLICA_NOT_AVAILABLE:
+                    case KAFKA_STORAGE_ERROR:
+                    case OFFSET_NOT_AVAILABLE:
+                    case LEADER_NOT_AVAILABLE:
+                    case FENCED_LEADER_EPOCH:
+                    case UNKNOWN_LEADER_EPOCH:
+                        logger().debug("Attempt to fetch offsets for partition {} failed due to {}, retrying.",
+                            topicPartition, error);
+                        break;
+                    case UNKNOWN_TOPIC_OR_PARTITION:
+                        logger().warn("Received unknown topic or partition error in OffsetsForLeaderEpoch request for partition {}.",
+                            topicPartition);
+                        break;
+                    case TOPIC_AUTHORIZATION_FAILED:
+                        unauthorizedTopics.add(topicPartition.topic());
+                        partitionsToRetry.remove(topicPartition);
+                        break;
+                    default:
+                        logger().warn("Attempt to fetch offsets for partition {} failed due to: {}, retrying.",
+                            topicPartition, error.message());
+                }
+            }
+        }
+
+        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/RequestFuture.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/RequestFuture.java
index f7e8ca14af592..8a9d970182c80 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/RequestFuture.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/RequestFuture.java
@@ -18,13 +18,15 @@
 
 import org.apache.kafka.common.errors.RetriableException;
 import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.utils.Timer;
+
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.atomic.AtomicReference;
 
 /**
- * Result of an asynchronous request from {@link ConsumerNetworkClient}. Use {@link ConsumerNetworkClient#poll(long)}
+ * Result of an asynchronous request from {@link ConsumerNetworkClient}. Use {@link ConsumerNetworkClient#poll(Timer)}
  * (and variants) to finish a request future. Use {@link #isDone()} to check if the future is complete, and
  * {@link #succeeded()} to check if the request completed successfully. Typical usage might look like this:
  *
@@ -242,18 +244,10 @@ public static  RequestFuture coordinatorNotAvailable() {
         return failure(Errors.COORDINATOR_NOT_AVAILABLE.exception());
     }
 
-    public static  RequestFuture leaderNotAvailable() {
-        return failure(Errors.LEADER_NOT_AVAILABLE.exception());
-    }
-
     public static  RequestFuture noBrokersAvailable() {
         return failure(new NoAvailableBrokersException());
     }
 
-    public static  RequestFuture staleMetadata() {
-        return failure(new StaleMetadataException());
-    }
-
     @Override
     public boolean shouldBlock() {
         return !isDone();
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/StaleMetadataException.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/StaleMetadataException.java
deleted file mode 100644
index 53110d3838747..0000000000000
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/StaleMetadataException.java
+++ /dev/null
@@ -1,26 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements. See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License. You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.kafka.clients.consumer.internals;
-
-import org.apache.kafka.common.errors.InvalidMetadataException;
-
-/**
- * Thrown when metadata is old and needs to be refreshed.
- */
-public class StaleMetadataException extends InvalidMetadataException {
-    private static final long serialVersionUID = 1L;
-}
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 e852bd2497969..30491110a3d78 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,30 +16,47 @@
  */
 package org.apache.kafka.clients.consumer.internals;
 
+import org.apache.kafka.clients.ApiVersions;
+import org.apache.kafka.clients.Metadata;
+import org.apache.kafka.clients.NodeApiVersions;
 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.IsolationLevel;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.internals.PartitionStates;
-import org.apache.kafka.common.requests.IsolationLevel;
+import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData.EpochEndOffset;
+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.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
 import java.util.Set;
+import java.util.function.LongSupplier;
+import java.util.function.Predicate;
 import java.util.regex.Pattern;
 
+import static org.apache.kafka.clients.consumer.internals.Fetcher.hasUsableOffsetForLeaderEpochVersion;
+import static org.apache.kafka.common.requests.OffsetsForLeaderEpochResponse.UNDEFINED_EPOCH;
+import static org.apache.kafka.common.requests.OffsetsForLeaderEpochResponse.UNDEFINED_EPOCH_OFFSET;
+
 /**
  * A class for tracking the topics, partitions, and offsets for the consumer. A partition
  * is "assigned" either directly with {@link #assignFromUser(Set)} (manual assignment)
  * 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
@@ -49,14 +66,14 @@
  * 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.
  *
- * This class also maintains a cache of the latest commit position for each of the assigned
- * partitions. This is updated through {@link #committed(TopicPartition, OffsetAndMetadata)} and can be used
- * to set the initial fetch position (e.g. {@link Fetcher#resetOffset(TopicPartition)}.
+ * Thread Safety: this class is thread-safe.
  */
 public class SubscriptionState {
     private static final String SUBSCRIPTION_EXCEPTION_MESSAGE =
             "Subscription to topics, partitions and pattern are mutually exclusive";
 
+    private final Logger log;
+
     private enum SubscriptionType {
         NONE, AUTO_TOPICS, AUTO_PATTERN, USER_ASSIGNED
     }
@@ -70,34 +87,68 @@ private enum SubscriptionType {
     /* the list of topics the user has requested */
     private Set subscription;
 
-    /* the list of topics the group has subscribed to (set only for the leader on join group completion) */
-    private final Set groupSubscription;
+    /* 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 Set groupSubscription;
 
     /* the partitions that are currently assigned, note that the order of partition matters (see FetchBuilder for more details) */
     private final PartitionStates assignment;
 
-    /* do we need to request the latest committed offsets from the coordinator? */
-    private boolean needsFetchCommittedOffsets;
-
     /* Default offset reset strategy */
     private final OffsetResetStrategy defaultResetStrategy;
 
     /* User-provided listener to be invoked when assignment changes */
-    private ConsumerRebalanceListener listener;
+    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 + ")}";
+    }
 
-    /* Listeners provide a hook for internal state cleanup (e.g. metrics) on assignment changes */
-    private List listeners = new ArrayList<>();
+    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(OffsetResetStrategy defaultResetStrategy) {
+    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 = new HashSet<>();
-        this.needsFetchCommittedOffsets = true; // initialize to true for the consumers to fetch offset upon starting up
         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.
@@ -111,48 +162,53 @@ else if (this.subscriptionType != type)
             throw new IllegalStateException(SUBSCRIPTION_EXCEPTION_MESSAGE);
     }
 
-    public void 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);
+        return changeSubscription(topics);
+    }
 
-        this.listener = listener;
-
-        changeSubscription(topics);
+    public synchronized void subscribe(Pattern pattern, ConsumerRebalanceListener listener) {
+        registerRebalanceListener(listener);
+        setSubscriptionType(SubscriptionType.AUTO_PATTERN);
+        this.subscribedPattern = pattern;
     }
 
-    public void subscribeFromPattern(Set topics) {
+    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);
 
-        changeSubscription(topics);
+        return changeSubscription(topics);
     }
 
-    private void changeSubscription(Set topicsToSubscribe) {
-        if (!this.subscription.equals(topicsToSubscribe)) {
-            this.subscription = topicsToSubscribe;
-            this.groupSubscription.addAll(topicsToSubscribe);
-        }
+    private boolean changeSubscription(Set topicsToSubscribe) {
+        if (subscription.equals(topicsToSubscribe))
+            return false;
+
+        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 void groupSubscribe(Collection topics) {
-        if (this.subscriptionType == SubscriptionType.USER_ASSIGNED)
+    synchronized boolean groupSubscribe(Collection topics) {
+        if (!hasAutoAssignedPartitions())
             throw new IllegalStateException(SUBSCRIPTION_EXCEPTION_MESSAGE);
-        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();
     }
 
     /**
@@ -160,107 +216,151 @@ public void resetGroupSubscription() {
      * note this is different from {@link #assignFromSubscribed(Collection)}
      * whose input partitions are provided from the subscribed topics.
      */
-    public void assignFromUser(Set partitions) {
+    public synchronized boolean assignFromUser(Set partitions) {
         setSubscriptionType(SubscriptionType.USER_ASSIGNED);
 
-        if (!this.assignment.partitionSet().equals(partitions)) {
-            fireOnAssignment(partitions);
+        if (this.assignment.partitionSet().equals(partitions))
+            return false;
+
+        assignmentId++;
+
+        // update the subscribed topics
+        Set manualSubscribedTopics = new HashSet<>();
+        Map partitionToState = new HashMap<>();
+        for (TopicPartition partition : partitions) {
+            TopicPartitionState state = assignment.stateValue(partition);
+            if (state == null)
+                state = new TopicPartitionState();
+            partitionToState.put(partition, state);
+
+            manualSubscribedTopics.add(partition.topic());
+        }
+
+        this.assignment.set(partitionToState);
+        return changeSubscription(manualSubscribedTopics);
+    }
 
-            Map partitionToState = new HashMap<>();
-            for (TopicPartition partition : partitions) {
-                TopicPartitionState state = assignment.stateValue(partition);
-                if (state == null)
-                    state = new TopicPartitionState();
-                partitionToState.put(partition, state);
+    /**
+     * @return true if assignments matches subscription, otherwise false
+     */
+    public synchronized boolean checkAssignmentMatchedSubscription(Collection assignments) {
+        for (TopicPartition topicPartition : assignments) {
+            if (this.subscribedPattern != null) {
+                if (!this.subscribedPattern.matcher(topicPartition.topic()).matches()) {
+                    log.info("Assigned partition {} for non-subscribed topic regex pattern; subscription pattern is {}",
+                        topicPartition,
+                        this.subscribedPattern);
+
+                    return false;
+                }
+            } else {
+                if (!this.subscription.contains(topicPartition.topic())) {
+                    log.info("Assigned partition {} for non-subscribed topic; subscription is {}", topicPartition, this.subscription);
+
+                    return false;
+                }
             }
-            this.assignment.set(partitionToState);
-            this.needsFetchCommittedOffsets = true;
         }
+
+        return true;
     }
 
     /**
-     * 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
+     * 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 void assignFromSubscribed(Collection assignments) {
-        if (!this.partitionsAutoAssigned())
+    public synchronized void assignFromSubscribed(Collection assignments) {
+        if (!this.hasAutoAssignedPartitions())
             throw new IllegalArgumentException("Attempt to dynamically assign partitions while manual assignment in use");
 
-        Map assignedPartitionStates = partitionToStateMap(assignments);
-        fireOnAssignment(assignedPartitionStates.keySet());
-
-        if (this.subscribedPattern != null) {
-            for (TopicPartition tp : assignments) {
-                if (!this.subscribedPattern.matcher(tp.topic()).matches())
-                    throw new IllegalArgumentException("Assigned partition " + tp + " for non-subscribed topic regex pattern; subscription pattern is " + this.subscribedPattern);
-            }
-        } else {
-            for (TopicPartition tp : assignments)
-                if (!this.subscription.contains(tp.topic()))
-                    throw new IllegalArgumentException("Assigned partition " + tp + " for non-subscribed topic; subscription is " + this.subscription);
+        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);
         }
 
-        // after rebalancing, we always reinitialize the assignment value
+        assignmentId++;
         this.assignment.set(assignedPartitionStates);
-        this.needsFetchCommittedOffsets = true;
     }
 
-    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.listener = listener;
-        this.subscribedPattern = pattern;
+        this.rebalanceListener = listener;
     }
 
-    public boolean hasPatternSubscription() {
+    /**
+     * Check whether pattern subscription is in use.
+     *
+     */
+    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 = Collections.emptySet();
         this.assignment.clear();
         this.subscribedPattern = null;
         this.subscriptionType = SubscriptionType.NONE;
-        fireOnAssignment(Collections.emptySet());
+        this.assignmentId++;
     }
 
-    public Pattern subscribedPattern() {
-        return this.subscribedPattern;
+    /**
+     * Check whether a topic matches a subscribed pattern.
+     *
+     * @return true if pattern subscription is in use and the topic matches the subscribed pattern, false otherwise
+     */
+    synchronized boolean matchesSubscribedPattern(String topic) {
+        Pattern pattern = this.subscribedPattern;
+        if (hasPatternSubscription() && pattern != null)
+            return pattern.matcher(topic).matches();
+        return false;
     }
 
-    public Set subscription() {
-        return this.subscription;
+    public synchronized Set subscription() {
+        if (hasAutoAssignedPartitions())
+            return this.subscription;
+        return Collections.emptySet();
     }
 
-    public Set pausedPartitions() {
-        HashSet paused = new HashSet<>();
-        for (PartitionStates.PartitionState state : assignment.partitionStates()) {
-            if (state.value().paused) {
-                paused.add(state.topicPartition());
-            }
-        }
-        return paused;
+    public synchronized Set pausedPartitions() {
+        return collectPartitions(TopicPartitionState::isPaused);
     }
 
     /**
-     * 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).
+     *
      * @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;
+        }
+    }
+
+    synchronized boolean needsMetadata(String topic) {
+        return subscription.contains(topic) || groupSubscription.contains(topic);
     }
 
     private TopicPartitionState assignedState(TopicPartition tp) {
@@ -270,211 +370,563 @@ private TopicPartitionState assignedState(TopicPartition tp) {
         return state;
     }
 
-    public void committed(TopicPartition tp, OffsetAndMetadata offset) {
-        assignedState(tp).committed(offset);
+    private TopicPartitionState assignedStateOrNull(TopicPartition tp) {
+        return this.assignment.stateValue(tp);
+    }
+
+    public synchronized void seekValidated(TopicPartition tp, FetchPosition position) {
+        assignedState(tp).seekValidated(position);
     }
 
-    public OffsetAndMetadata committed(TopicPartition tp) {
-        return assignedState(tp).committed;
+    public void seek(TopicPartition tp, long offset) {
+        seekValidated(tp, new FetchPosition(offset));
     }
 
-    public void needRefreshCommits() {
-        this.needsFetchCommittedOffsets = true;
+    public void seekUnvalidated(TopicPartition tp, FetchPosition position) {
+        assignedState(tp).seekUnvalidated(position);
     }
 
-    public boolean refreshCommitsNeeded() {
-        return this.needsFetchCommittedOffsets;
+    synchronized void maybeSeekUnvalidated(TopicPartition tp, FetchPosition position, 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 position {}.", tp, position);
+            state.seekUnvalidated(position);
+        }
     }
 
-    public void commitsRefreshed() {
-        this.needsFetchCommittedOffsets = false;
+    /**
+     * @return a modifiable copy of the currently assigned partitions
+     */
+    public synchronized Set assignedPartitions() {
+        return new HashSet<>(this.assignment.partitionSet());
     }
 
-    public void seek(TopicPartition tp, long offset) {
-        assignedState(tp).seek(offset);
+    /**
+     * @return a modifiable copy of the currently assigned partitions as a list
+     */
+    public synchronized List assignedPartitionsList() {
+        return new ArrayList<>(this.assignment.partitionSet());
     }
 
-    public Set assignedPartitions() {
-        return this.assignment.partitionSet();
+    /**
+     * Provides the number of assigned partitions in a thread safe manner.
+     * @return the number of assigned partitions.
+     */
+    synchronized int numAssignedPartitions() {
+        return this.assignment.size();
     }
 
-    public List fetchablePartitions() {
-        List fetchable = new ArrayList<>(assignment.size());
-        for (PartitionStates.PartitionState state : assignment.partitionStates()) {
-            if (state.value().isFetchable())
-                fetchable.add(state.topicPartition());
-        }
-        return fetchable;
+    // Visible for testing
+    public synchronized List fetchablePartitions(Predicate isAvailable) {
+        // Since this is in the hot-path for fetching, we do this instead of using java.util.stream API
+        List result = new ArrayList<>();
+        assignment.forEach((topicPartition, topicPartitionState) -> {
+            // Cheap check is first to avoid evaluating the predicate if possible
+            if (topicPartitionState.isFetchable() && isAvailable.test(topicPartition)) {
+                result.add(topicPartition);
+            }
+        });
+        return result;
     }
 
-    public boolean partitionsAutoAssigned() {
+    public synchronized boolean hasAutoAssignedPartitions() {
         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);
+    }
+
+    /**
+     * Enter the offset validation state if the leader for this partition is known to support a usable version of the
+     * OffsetsForLeaderEpoch API. If the leader node does not support the API, simply complete the offset validation.
+     *
+     * @param apiVersions supported API versions
+     * @param tp topic partition to validate
+     * @param leaderAndEpoch leader epoch of the topic partition
+     * @return true if we enter the offset validation state
+     */
+    public synchronized boolean maybeValidatePositionForCurrentLeader(ApiVersions apiVersions,
+                                                                      TopicPartition tp,
+                                                                      Metadata.LeaderAndEpoch leaderAndEpoch) {
+        if (leaderAndEpoch.leader.isPresent()) {
+            NodeApiVersions nodeApiVersions = apiVersions.get(leaderAndEpoch.leader.get().idString());
+            if (nodeApiVersions == null || hasUsableOffsetForLeaderEpochVersion(nodeApiVersions)) {
+                return assignedState(tp).maybeValidatePosition(leaderAndEpoch);
+            } else {
+                // If the broker does not support a newer version of OffsetsForLeaderEpoch, we skip validation
+                assignedState(tp).updatePositionLeaderNoValidation(leaderAndEpoch);
+                return false;
+            }
+        } else {
+            return assignedState(tp).maybeValidatePosition(leaderAndEpoch);
+        }
+    }
+
+    /**
+     * Attempt to complete validation with the end offset returned from the OffsetForLeaderEpoch request.
+     * @return Log truncation details if 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() == UNDEFINED_EPOCH_OFFSET ||
+                        epochEndOffset.leaderEpoch() == UNDEFINED_EPOCH) {
+                if (hasDefaultOffsetResetPolicy()) {
+                    log.info("Truncation detected for partition {} at offset {}, resetting offset",
+                             tp, currentPosition);
+                    requestOffsetReset(tp);
+                } else {
+                    log.warn("Truncation detected for partition {} at offset {}, but no reset policy is set",
+                             tp, currentPosition);
+                    return Optional.of(new LogTruncation(tp, requestPosition, Optional.empty()));
+                }
+            } 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 {
+                    OffsetAndMetadata divergentOffset = new OffsetAndMetadata(epochEndOffset.endOffset(),
+                        Optional.of(epochEndOffset.leaderEpoch()), null);
+                    log.warn("Truncation detected for partition {} at offset {} (the end offset from the " +
+                             "broker is {}), but no reset policy is set", tp, currentPosition, divergentOffset);
+                    return Optional.of(new LogTruncation(tp, requestPosition, Optional.of(divergentOffset)));
+                }
+            } 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;
+    }
+
+    synchronized Long partitionLead(TopicPartition tp) {
+        TopicPartitionState topicPartitionState = assignedState(tp);
+        return topicPartitionState.logStartOffset == null ? null : topicPartitionState.position.offset - topicPartitionState.logStartOffset;
     }
 
-    public void updateHighWatermark(TopicPartition tp, long highWatermark) {
-        assignedState(tp).highWatermark = highWatermark;
+    synchronized void updateHighWatermark(TopicPartition tp, long highWatermark) {
+        assignedState(tp).highWatermark(highWatermark);
     }
 
-    public void updateLastStableOffset(TopicPartition tp, long lastStableOffset) {
-        assignedState(tp).lastStableOffset = lastStableOffset;
+    synchronized void updateLogStartOffset(TopicPartition tp, long logStartOffset) {
+        assignedState(tp).logStartOffset(logStartOffset);
     }
 
-    public Map allConsumed() {
-        Map allConsumed = new HashMap<>();
-        for (PartitionStates.PartitionState state : assignment.partitionStates()) {
-            if (state.value().hasValidPosition())
-                allConsumed.put(state.topicPartition(), new OffsetAndMetadata(state.value().position));
+    synchronized void updateLastStableOffset(TopicPartition tp, long lastStableOffset) {
+        assignedState(tp).lastStableOffset(lastStableOffset);
+    }
+
+    /**
+     * 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, LongSupplier timeMs) {
+        assignedState(tp).updatePreferredReadReplica(preferredReadReplicaId, timeMs);
+    }
+
+    /**
+     * 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);
         }
+    }
+
+    /**
+     * 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.forEach((topicPartition, partitionState) -> {
+            if (partitionState.hasValidPosition())
+                allConsumed.put(topicPartition, new OffsetAndMetadata(partitionState.position.offset,
+                        partitionState.position.offsetEpoch, ""));
+        });
         return allConsumed;
     }
 
-    public void needOffsetReset(TopicPartition partition, OffsetResetStrategy offsetResetStrategy) {
-        assignedState(partition).awaitReset(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 needOffsetReset(TopicPartition partition) {
-        needOffsetReset(partition, defaultResetStrategy);
+    synchronized void setNextAllowedRetry(Set partitions, long nextAllowResetTimeMs) {
+        for (TopicPartition partition : partitions) {
+            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(Collection partitions) {
-        for (TopicPartition partition : partitions)
-            if (!hasValidPosition(partition))
+    public synchronized boolean hasAllFetchPositions() {
+        // Since this is in the hot-path for fetching, we do this instead of using java.util.stream API
+        Iterator it = assignment.stateIterator();
+        while (it.hasNext()) {
+            if (!it.next().hasValidPosition()) {
                 return false;
+            }
+        }
         return true;
     }
 
-    public boolean hasAllFetchPositions() {
-        return hasAllFetchPositions(this.assignedPartitions());
+    public synchronized Set initializingPartitions() {
+        return collectPartitions(state -> state.fetchState.equals(FetchStates.INITIALIZING));
     }
 
-    public Set missingFetchPositions() {
-        Set missing = new HashSet<>();
-        for (PartitionStates.PartitionState state : assignment.partitionStates()) {
-            if (!state.value().hasValidPosition())
-                missing.add(state.topicPartition());
-        }
-        return missing;
+    private Set collectPartitions(Predicate filter) {
+        Set result = new HashSet<>();
+        assignment.forEach((topicPartition, topicPartitionState) -> {
+            if (filter.test(topicPartitionState)) {
+                result.add(topicPartition);
+            }
+        });
+        return result;
     }
 
-    public boolean isAssigned(TopicPartition tp) {
-        return assignment.contains(tp);
+
+    public synchronized void resetInitializingPositions() {
+        final Set partitionsWithNoOffsets = new HashSet<>();
+        assignment.forEach((tp, partitionState) -> {
+            if (partitionState.fetchState.equals(FetchStates.INITIALIZING)) {
+                if (defaultResetStrategy == OffsetResetStrategy.NONE)
+                    partitionsWithNoOffsets.add(tp);
+                else
+                    requestOffsetReset(tp);
+            }
+        });
+
+        if (!partitionsWithNoOffsets.isEmpty())
+            throw new NoOffsetForPartitionException(partitionsWithNoOffsets);
     }
 
-    public boolean isPaused(TopicPartition tp) {
-        return isAssigned(tp) && assignedState(tp).paused;
+    public synchronized Set partitionsNeedingReset(long nowMs) {
+        return collectPartitions(state -> state.awaitingReset() && !state.awaitingRetryBackoff(nowMs));
     }
 
-    public boolean isFetchable(TopicPartition tp) {
-        return isAssigned(tp) && assignedState(tp).isFetchable();
+    public synchronized Set partitionsNeedingValidation(long nowMs) {
+        return collectPartitions(state -> state.awaitingValidation() && !state.awaitingRetryBackoff(nowMs));
     }
 
-    public boolean hasValidPosition(TopicPartition tp) {
-        return isAssigned(tp) && assignedState(tp).hasValidPosition();
+    public synchronized boolean isAssigned(TopicPartition tp) {
+        return assignment.contains(tp);
     }
 
-    public void pause(TopicPartition tp) {
-        assignedState(tp).pause();
+    public synchronized boolean isPaused(TopicPartition tp) {
+        TopicPartitionState assignedOrNull = assignedStateOrNull(tp);
+        return assignedOrNull != null && assignedOrNull.isPaused();
     }
 
-    public void resume(TopicPartition tp) {
-        assignedState(tp).resume();
+    synchronized boolean isFetchable(TopicPartition tp) {
+        TopicPartitionState assignedOrNull = assignedStateOrNull(tp);
+        return assignedOrNull != null && assignedOrNull.isFetchable();
     }
 
-    public void movePartitionToEnd(TopicPartition tp) {
-        assignment.moveToEnd(tp);
+    public synchronized boolean hasValidPosition(TopicPartition tp) {
+        TopicPartitionState assignedOrNull = assignedStateOrNull(tp);
+        return assignedOrNull != null && assignedOrNull.hasValidPosition();
     }
 
-    public ConsumerRebalanceListener listener() {
-        return listener;
+    public synchronized void pause(TopicPartition tp) {
+        assignedState(tp).pause();
+    }
+
+    public synchronized void resume(TopicPartition tp) {
+        assignedState(tp).resume();
     }
 
-    public void addListener(Listener listener) {
-        listeners.add(listener);
+    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 fireOnAssignment(Set assignment) {
-        for (Listener listener : listeners)
-            listener.onAssignment(assignment);
+    synchronized void movePartitionToEnd(TopicPartition tp) {
+        assignment.moveToEnd(tp);
     }
 
-    private static Map partitionToStateMap(Collection assignments) {
-        Map map = new HashMap<>(assignments.size());
-        for (TopicPartition tp : assignments)
-            map.put(tp, new TopicPartitionState());
-        return map;
+    public synchronized ConsumerRebalanceListener rebalanceListener() {
+        return rebalanceListener;
     }
 
     private static class TopicPartitionState {
-        private Long position; // last consumed position
+
+        private FetchState fetchState;
+        private FetchPosition position; // last consumed position
+
         private Long highWatermark; // the high watermark from last fetch
+        private Long logStartOffset; // the log start offset
         private Long lastStableOffset;
-        private OffsetAndMetadata committed;  // last committed position
         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 nextRetryTimeMs;
+        private Integer preferredReadReplica;
+        private Long preferredReadReplicaExpireTimeMs;
 
-        public TopicPartitionState() {
+        TopicPartitionState() {
             this.paused = false;
+            this.fetchState = FetchStates.INITIALIZING;
             this.position = null;
             this.highWatermark = null;
+            this.logStartOffset = null;
             this.lastStableOffset = null;
-            this.committed = null;
             this.resetStrategy = null;
+            this.nextRetryTimeMs = null;
+            this.preferredReadReplica = null;
         }
 
-        private void awaitReset(OffsetResetStrategy strategy) {
-            this.resetStrategy = strategy;
-            this.position = null;
+        private void transitionState(FetchState newState, Runnable runIfTransitioned) {
+            FetchState nextState = this.fetchState.transitionTo(newState);
+            if (nextState.equals(newState)) {
+                this.fetchState = nextState;
+                runIfTransitioned.run();
+                if (this.position == null && nextState.requiresPosition()) {
+                    throw new IllegalStateException("Transitioned subscription state to " + nextState + ", but position is null");
+                } else if (!nextState.requiresPosition()) {
+                    this.position = null;
+                }
+            }
+        }
+
+        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, LongSupplier timeMs) {
+            if (this.preferredReadReplica == null || preferredReadReplica != this.preferredReadReplica) {
+                this.preferredReadReplica = preferredReadReplica;
+                this.preferredReadReplicaExpireTimeMs = timeMs.getAsLong();
+            }
+        }
+
+        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) {
+            transitionState(FetchStates.AWAIT_RESET, () -> {
+                this.resetStrategy = strategy;
+                this.nextRetryTimeMs = null;
+            });
+        }
+
+        /**
+         * Check if the position exists and needs to be validated. If so, enter the AWAIT_VALIDATION state. This method
+         * also will update the position with the current leader and epoch.
+         *
+         * @param currentLeaderAndEpoch leader and epoch to compare the offset with
+         * @return true if the position is now awaiting validation
+         */
+        private boolean maybeValidatePosition(Metadata.LeaderAndEpoch currentLeaderAndEpoch) {
+            if (this.fetchState.equals(FetchStates.AWAIT_RESET)) {
+                return false;
+            }
+
+            if (!currentLeaderAndEpoch.leader.isPresent()) {
+                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);
+        }
+
+        /**
+         * For older versions of the API, we cannot perform offset validation so we simply transition directly to FETCHING
+         */
+        private void updatePositionLeaderNoValidation(Metadata.LeaderAndEpoch currentLeaderAndEpoch) {
+            if (position != null) {
+                transitionState(FetchStates.FETCHING, () -> {
+                    this.position = new FetchPosition(position.offset, position.offsetEpoch, currentLeaderAndEpoch);
+                    this.nextRetryTimeMs = null;
+                });
+            }
+        }
+
+        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 awaitingRetryBackoff(long nowMs) {
+            return nextRetryTimeMs != null && nowMs < nextRetryTimeMs;
+        }
+
+        private boolean awaitingReset() {
+            return fetchState.equals(FetchStates.AWAIT_RESET);
+        }
+
+        private void setNextAllowedRetry(long nextAllowedRetryTimeMs) {
+            this.nextRetryTimeMs = nextAllowedRetryTimeMs;
+        }
+
+        private void requestFailed(long nextAllowedRetryTimeMs) {
+            this.nextRetryTimeMs = nextAllowedRetryTimeMs;
         }
 
-        public boolean awaitingReset() {
-            return resetStrategy != null;
+        private boolean hasValidPosition() {
+            return fetchState.hasValidPosition();
         }
 
-        public boolean hasValidPosition() {
+        private boolean hasPosition() {
             return position != null;
         }
 
-        private void seek(long offset) {
-            this.position = offset;
-            this.resetStrategy = null;
+        private boolean isPaused() {
+            return paused;
+        }
+
+        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 void committed(OffsetAndMetadata offset) {
-            this.committed = offset;
+        private FetchPosition validPosition() {
+            if (hasValidPosition()) {
+                return position;
+            } else {
+                return null;
+            }
         }
 
         private void pause() {
@@ -489,16 +941,208 @@ 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 {
+    /**
+     * 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;
+            }
+        }
+
         /**
-         * 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
+         * Return the valid states which this state can transition to
+         */
+        Collection validTransitions();
+
+        /**
+         * Test if this state requires a position to be set
+         */
+        boolean requiresPosition();
+
+        /**
+         * Test if this state is considered to have a valid position which can be used for fetching
          */
-        void onAssignment(Set assignment);
+        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 requiresPosition() {
+                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 requiresPosition() {
+                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 requiresPosition() {
+                return false;
+            }
+
+            @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 requiresPosition() {
+                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.
+     */
+    public static class FetchPosition {
+        public final long offset;
+        final Optional offsetEpoch;
+        final Metadata.LeaderAndEpoch currentLeader;
+
+        FetchPosition(long offset) {
+            this(offset, Optional.empty(), Metadata.LeaderAndEpoch.noLeaderOrEpoch());
+        }
+
+        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 +
+                    '}';
+        }
+    }
+
+    public static class LogTruncation {
+        public final TopicPartition topicPartition;
+        public final FetchPosition fetchPosition;
+        public final Optional divergentOffsetOpt;
+
+        public LogTruncation(TopicPartition topicPartition,
+                             FetchPosition fetchPosition,
+                             Optional divergentOffsetOpt) {
+            this.topicPartition = topicPartition;
+            this.fetchPosition = fetchPosition;
+            this.divergentOffsetOpt = divergentOffsetOpt;
+        }
+
+        @Override
+        public String toString() {
+            StringBuilder bldr = new StringBuilder()
+                .append("(partition=")
+                .append(topicPartition)
+                .append(", fetchOffset=")
+                .append(fetchPosition.offset)
+                .append(", fetchEpoch=")
+                .append(fetchPosition.offsetEpoch);
+
+            if (divergentOffsetOpt.isPresent()) {
+                OffsetAndMetadata divergentOffset = divergentOffsetOpt.get();
+                bldr.append(", divergentOffset=")
+                    .append(divergentOffset.offset())
+                    .append(", divergentEpoch=")
+                    .append(divergentOffset.leaderEpoch());
+            } else {
+                bldr.append(", divergentOffset=unknown")
+                    .append(", divergentEpoch=unknown");
+            }
+
+            return bldr.append(")").toString();
+
+        }
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/BufferExhaustedException.java b/clients/src/main/java/org/apache/kafka/clients/producer/BufferExhaustedException.java
index be840db282fac..292bb4eef1c19 100644
--- a/clients/src/main/java/org/apache/kafka/clients/producer/BufferExhaustedException.java
+++ b/clients/src/main/java/org/apache/kafka/clients/producer/BufferExhaustedException.java
@@ -16,13 +16,17 @@
  */
 package org.apache.kafka.clients.producer;
 
-import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.errors.TimeoutException;
 
 /**
- * This exception is thrown if the producer is in non-blocking mode and the rate of data production exceeds the rate at
- * which data can be sent for long enough for the allocated buffer to be exhausted.
+ * This exception is thrown if the producer cannot allocate memory for a record within max.block.ms due to the buffer
+ * being too full.
+ *
+ * In earlier versions a TimeoutException was thrown instead of this. To keep existing catch-clauses working
+ * this class extends TimeoutException.
+ *
  */
-public class BufferExhaustedException extends KafkaException {
+public class BufferExhaustedException extends TimeoutException {
 
     private static final long serialVersionUID = 1L;
 
diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/Callback.java b/clients/src/main/java/org/apache/kafka/clients/producer/Callback.java
index a70e4e9a68530..ee0610ebe4722 100644
--- a/clients/src/main/java/org/apache/kafka/clients/producer/Callback.java
+++ b/clients/src/main/java/org/apache/kafka/clients/producer/Callback.java
@@ -24,10 +24,11 @@ public interface Callback {
 
     /**
      * A callback method the user can implement to provide asynchronous handling of request completion. This method will
-     * be called when the record sent to the server has been acknowledged. Exactly one of the arguments will be
-     * non-null.
-     * @param metadata The metadata for the record that was sent (i.e. the partition and offset). Null if an error
-     *        occurred.
+     * be called when the record sent to the server has been acknowledged. When exception is not null in the callback,
+     * metadata will contain the special -1 value for all fields except for topicPartition, which will be valid.
+     *
+     * @param metadata The metadata for the record that was sent (i.e. the partition and offset). An empty metadata
+     *                 with -1 value for all fields except for topicPartition will be returned if an error occurred.
      * @param exception The exception thrown during processing of this record. Null if no error occurred.
      *                  Possible thrown exceptions include:
      *
@@ -38,6 +39,8 @@ public interface Callback {
      *                  RecordBatchTooLargeException
      *                  RecordTooLargeException
      *                  UnknownServerException
+     *                  UnknownProducerIdException
+     *                  InvalidProducerEpochException
      *
      *                  Retriable exceptions (transient, may be covered by increasing #.retries):
      *
@@ -49,5 +52,5 @@ public interface Callback {
      *                  TimeoutException
      *                  UnknownTopicOrPartitionException
      */
-    public void onCompletion(RecordMetadata metadata, Exception exception);
+    void onCompletion(RecordMetadata metadata, Exception exception);
 }
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 b3cff19114e4d..e0a961c3a731d 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
@@ -17,13 +17,18 @@
 package org.apache.kafka.clients.producer;
 
 import org.apache.kafka.clients.ApiVersions;
+import org.apache.kafka.clients.ClientDnsLookup;
 import org.apache.kafka.clients.ClientUtils;
-import org.apache.kafka.clients.Metadata;
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.clients.KafkaClient;
 import org.apache.kafka.clients.NetworkClient;
+import org.apache.kafka.clients.consumer.ConsumerGroupMetadata;
 import org.apache.kafka.clients.consumer.KafkaConsumer;
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.clients.consumer.OffsetCommitCallback;
+import org.apache.kafka.clients.producer.internals.BufferPool;
 import org.apache.kafka.clients.producer.internals.ProducerInterceptors;
+import org.apache.kafka.clients.producer.internals.ProducerMetadata;
 import org.apache.kafka.clients.producer.internals.ProducerMetrics;
 import org.apache.kafka.clients.producer.internals.RecordAccumulator;
 import org.apache.kafka.clients.producer.internals.Sender;
@@ -40,18 +45,20 @@
 import org.apache.kafka.common.errors.AuthenticationException;
 import org.apache.kafka.common.errors.AuthorizationException;
 import org.apache.kafka.common.errors.InterruptException;
+import org.apache.kafka.common.errors.InvalidTopicException;
 import org.apache.kafka.common.errors.ProducerFencedException;
 import org.apache.kafka.common.errors.RecordTooLargeException;
 import org.apache.kafka.common.errors.SerializationException;
 import org.apache.kafka.common.errors.TimeoutException;
-import org.apache.kafka.common.errors.TopicAuthorizationException;
 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.internals.ClusterResourceListeners;
 import org.apache.kafka.common.metrics.JmxReporter;
+import org.apache.kafka.common.metrics.KafkaMetricsContext;
 import org.apache.kafka.common.metrics.MetricConfig;
 import org.apache.kafka.common.metrics.Metrics;
+import org.apache.kafka.common.metrics.MetricsContext;
 import org.apache.kafka.common.metrics.MetricsReporter;
 import org.apache.kafka.common.metrics.Sensor;
 import org.apache.kafka.common.network.ChannelBuilder;
@@ -59,15 +66,17 @@
 import org.apache.kafka.common.record.AbstractRecords;
 import org.apache.kafka.common.record.CompressionType;
 import org.apache.kafka.common.record.RecordBatch;
-import org.apache.kafka.common.serialization.ExtendedSerializer;
+import org.apache.kafka.common.requests.JoinGroupRequest;
 import org.apache.kafka.common.serialization.Serializer;
 import org.apache.kafka.common.utils.AppInfoParser;
 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;
+import java.time.Duration;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
@@ -76,10 +85,8 @@
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 
-import static org.apache.kafka.common.serialization.ExtendedSerializer.Wrapper.ensureExtended;
 
 /**
  * A Kafka client that publishes records to the Kafka cluster.
@@ -95,9 +102,7 @@
  * props.put("bootstrap.servers", "localhost:9092");
  * props.put("acks", "all");
  * props.put("retries", 0);
- * props.put("batch.size", 16384);
  * props.put("linger.ms", 1);
- * props.put("buffer.memory", 33554432);
  * props.put("key.serializer", "org.apache.kafka.common.serialization.StringSerializer");
  * props.put("value.serializer", "org.apache.kafka.common.serialization.StringSerializer");
  *
@@ -230,9 +235,9 @@
 public class KafkaProducer implements Producer {
 
     private final Logger log;
-    private static final AtomicInteger PRODUCER_CLIENT_ID_SEQUENCE = new AtomicInteger(1);
     private static final String JMX_PREFIX = "kafka.producer";
     public static final String NETWORK_THREAD_PREFIX = "kafka-producer-network-thread";
+    public static final String PRODUCER_METRIC_GROUP_NAME = "producer-metrics";
 
     private final String clientId;
     // Visible for testing
@@ -240,18 +245,17 @@ public class KafkaProducer implements Producer {
     private final Partitioner partitioner;
     private final int maxRequestSize;
     private final long totalMemorySize;
-    private final Metadata metadata;
+    private final ProducerMetadata metadata;
     private final RecordAccumulator accumulator;
     private final Sender sender;
     private final Thread ioThread;
     private final CompressionType compressionType;
     private final Sensor errors;
     private final Time time;
-    private final ExtendedSerializer keySerializer;
-    private final ExtendedSerializer valueSerializer;
+    private final Serializer keySerializer;
+    private final Serializer valueSerializer;
     private final ProducerConfig producerConfig;
     private final long maxBlockTimeMs;
-    private final int requestTimeoutMs;
     private final ProducerInterceptors interceptors;
     private final ApiVersions apiVersions;
     private final TransactionManager transactionManager;
@@ -261,11 +265,13 @@ public class KafkaProducer implements Producer {
      * are documented here. Values can be
      * either strings or Objects of the appropriate type (for example a numeric configuration would accept either the
      * string "42" or the integer 42).
+     * 
    
+     * Note: after creating a {@code KafkaProducer} you must always {@link #close()} it to avoid resource leaks.
      * @param configs   The producer configs
      *
      */
-    public KafkaProducer(Map configs) {
-        this(new ProducerConfig(configs), null, null);
+    public KafkaProducer(final Map configs) {
+        this(configs, null, null);
     }
 
     /**
@@ -273,6 +279,8 @@ public KafkaProducer(Map configs) {
      * Valid configuration strings are documented here.
      * Values can be either strings or Objects of the appropriate type (for example a numeric configuration would accept
      * either the string "42" or the integer 42).
+     * 
    
+     * Note: after creating a {@code KafkaProducer} you must always {@link #close()} it to avoid resource leaks.
      * @param configs   The producer configs
      * @param keySerializer  The serializer for key that implements {@link Serializer}. The configure() method won't be
      *                       called in the producer when the serializer is passed in directly.
@@ -280,22 +288,26 @@ public KafkaProducer(Map configs) {
      *                         be called in the producer when the serializer is passed in directly.
      */
     public KafkaProducer(Map configs, Serializer keySerializer, Serializer valueSerializer) {
-        this(new ProducerConfig(ProducerConfig.addSerializerToConfig(configs, keySerializer, valueSerializer)),
-                keySerializer, valueSerializer);
+        this(new ProducerConfig(ProducerConfig.appendSerializerToConfig(configs, keySerializer, valueSerializer)),
+                keySerializer, valueSerializer, null, null, null, Time.SYSTEM);
     }
 
     /**
      * A producer is instantiated by providing a set of key-value pairs as configuration. Valid configuration strings
      * are documented here.
+     * 
    
+     * Note: after creating a {@code KafkaProducer} you must always {@link #close()} it to avoid resource leaks.
      * @param properties   The producer configs
      */
     public KafkaProducer(Properties properties) {
-        this(new ProducerConfig(properties), null, null);
+        this(properties, null, null);
     }
 
     /**
      * A producer is instantiated by providing a set of key-value pairs as configuration, a key and a value {@link Serializer}.
      * Valid configuration strings are documented here.
+     * 
    
+     * Note: after creating a {@code KafkaProducer} you must always {@link #close()} it to avoid resource leaks.
      * @param properties   The producer configs
      * @param keySerializer  The serializer for key that implements {@link Serializer}. The configure() method won't be
      *                       called in the producer when the serializer is passed in directly.
@@ -303,23 +315,26 @@ public KafkaProducer(Properties properties) {
      *                         be called in the producer when the serializer is passed in directly.
      */
     public KafkaProducer(Properties properties, Serializer keySerializer, Serializer valueSerializer) {
-        this(new ProducerConfig(ProducerConfig.addSerializerToConfig(properties, keySerializer, valueSerializer)),
-                keySerializer, valueSerializer);
+        this(Utils.propsToMap(properties), keySerializer, valueSerializer);
     }
 
+    // visible for testing
     @SuppressWarnings("unchecked")
-    private KafkaProducer(ProducerConfig config, Serializer keySerializer, Serializer valueSerializer) {
+    KafkaProducer(ProducerConfig config,
+                  Serializer keySerializer,
+                  Serializer valueSerializer,
+                  ProducerMetadata metadata,
+                  KafkaClient kafkaClient,
+                  ProducerInterceptors interceptors,
+                  Time time) {
         try {
-            Map userProvidedConfigs = config.originals();
             this.producerConfig = config;
-            this.time = Time.SYSTEM;
-            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.time = time;
+
+            String transactionalId = config.getString(ProducerConfig.TRANSACTIONAL_ID_CONFIG);
+
+            this.clientId = config.getString(ProducerConfig.CLIENT_ID_CONFIG);
+
             LogContext logContext;
             if (transactionalId == null)
                 logContext = new LogContext(String.format("[Producer clientId=%s] ", clientId));
@@ -334,195 +349,222 @@ private KafkaProducer(ProducerConfig config, Serializer keySerializer, Serial
                     .recordLevel(Sensor.RecordingLevel.forName(config.getString(ProducerConfig.METRICS_RECORDING_LEVEL_CONFIG)))
                     .tags(metricTags);
             List reporters = config.getConfiguredInstances(ProducerConfig.METRIC_REPORTER_CLASSES_CONFIG,
-                    MetricsReporter.class);
-            reporters.add(new JmxReporter(JMX_PREFIX));
-            this.metrics = new Metrics(metricConfig, reporters, time);
-            ProducerMetrics metricsRegistry = new ProducerMetrics(this.metrics);
-            this.partitioner = config.getConfiguredInstance(ProducerConfig.PARTITIONER_CLASS_CONFIG, Partitioner.class);
+                    MetricsReporter.class,
+                    Collections.singletonMap(ProducerConfig.CLIENT_ID_CONFIG, clientId));
+            JmxReporter jmxReporter = new JmxReporter();
+            jmxReporter.configure(config.originals(Collections.singletonMap(ProducerConfig.CLIENT_ID_CONFIG, clientId)));
+            reporters.add(jmxReporter);
+            MetricsContext metricsContext = new KafkaMetricsContext(JMX_PREFIX,
+                    config.originalsWithPrefix(CommonClientConfigs.METRICS_CONTEXT_PREFIX));
+            this.metrics = new Metrics(metricConfig, reporters, time, metricsContext);
+            this.partitioner = config.getConfiguredInstance(
+                    ProducerConfig.PARTITIONER_CLASS_CONFIG,
+                    Partitioner.class,
+                    Collections.singletonMap(ProducerConfig.CLIENT_ID_CONFIG, clientId));
             long retryBackoffMs = config.getLong(ProducerConfig.RETRY_BACKOFF_MS_CONFIG);
             if (keySerializer == null) {
-                this.keySerializer = ensureExtended(config.getConfiguredInstance(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG,
-                                                                                         Serializer.class));
-                this.keySerializer.configure(config.originals(), true);
+                this.keySerializer = config.getConfiguredInstance(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG,
+                                                                                         Serializer.class);
+                this.keySerializer.configure(config.originals(Collections.singletonMap(ProducerConfig.CLIENT_ID_CONFIG, clientId)), true);
             } else {
                 config.ignore(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG);
-                this.keySerializer = ensureExtended(keySerializer);
+                this.keySerializer = keySerializer;
             }
             if (valueSerializer == null) {
-                this.valueSerializer = ensureExtended(config.getConfiguredInstance(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG,
-                                                                                           Serializer.class));
-                this.valueSerializer.configure(config.originals(), false);
+                this.valueSerializer = config.getConfiguredInstance(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG,
+                                                                                           Serializer.class);
+                this.valueSerializer.configure(config.originals(Collections.singletonMap(ProducerConfig.CLIENT_ID_CONFIG, clientId)), false);
             } else {
                 config.ignore(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG);
-                this.valueSerializer = ensureExtended(valueSerializer);
+                this.valueSerializer = valueSerializer;
             }
 
-            // load interceptors and make sure they get clientId
-            userProvidedConfigs.put(ProducerConfig.CLIENT_ID_CONFIG, clientId);
-            List> interceptorList = (List) (new ProducerConfig(userProvidedConfigs, false)).getConfiguredInstances(ProducerConfig.INTERCEPTOR_CLASSES_CONFIG,
-                    ProducerInterceptor.class);
-            this.interceptors = interceptorList.isEmpty() ? null : new ProducerInterceptors<>(interceptorList);
-            ClusterResourceListeners clusterResourceListeners = configureClusterResourceListeners(keySerializer, valueSerializer, interceptorList, reporters);
-            this.metadata = new Metadata(retryBackoffMs, config.getLong(ProducerConfig.METADATA_MAX_AGE_CONFIG),
-                    true, true, clusterResourceListeners);
+            List> interceptorList = (List) config.getConfiguredInstances(
+                    ProducerConfig.INTERCEPTOR_CLASSES_CONFIG,
+                    ProducerInterceptor.class,
+                    Collections.singletonMap(ProducerConfig.CLIENT_ID_CONFIG, clientId));
+            if (interceptors != null)
+                this.interceptors = interceptors;
+            else
+                this.interceptors = new ProducerInterceptors<>(interceptorList);
+            ClusterResourceListeners clusterResourceListeners = configureClusterResourceListeners(keySerializer,
+                    valueSerializer, interceptorList, reporters);
             this.maxRequestSize = config.getInt(ProducerConfig.MAX_REQUEST_SIZE_CONFIG);
             this.totalMemorySize = config.getLong(ProducerConfig.BUFFER_MEMORY_CONFIG);
             this.compressionType = CompressionType.forName(config.getString(ProducerConfig.COMPRESSION_TYPE_CONFIG));
 
             this.maxBlockTimeMs = config.getLong(ProducerConfig.MAX_BLOCK_MS_CONFIG);
-            this.requestTimeoutMs = config.getInt(ProducerConfig.REQUEST_TIMEOUT_MS_CONFIG);
-            this.transactionManager = configureTransactionState(config, logContext, log);
-            int retries = configureRetries(config, transactionManager != null, log);
-            int maxInflightRequests = configureInflightRequests(config, transactionManager != null);
-            short acks = configureAcks(config, transactionManager != null, log);
+            int deliveryTimeoutMs = configureDeliveryTimeout(config, log);
 
             this.apiVersions = new ApiVersions();
+            this.transactionManager = configureTransactionState(config, logContext);
             this.accumulator = new RecordAccumulator(logContext,
                     config.getInt(ProducerConfig.BATCH_SIZE_CONFIG),
-                    this.totalMemorySize,
                     this.compressionType,
-                    config.getLong(ProducerConfig.LINGER_MS_CONFIG),
+                    lingerMs(config),
                     retryBackoffMs,
+                    deliveryTimeoutMs,
                     metrics,
+                    PRODUCER_METRIC_GROUP_NAME,
                     time,
                     apiVersions,
-                    transactionManager);
-            List addresses = ClientUtils.parseAndValidateAddresses(config.getList(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG));
-            this.metadata.update(Cluster.bootstrap(addresses), Collections.emptySet(), time.milliseconds());
-            ChannelBuilder channelBuilder = ClientUtils.createChannelBuilder(config);
-            Sensor throttleTimeSensor = Sender.throttleTimeSensor(metricsRegistry.senderMetrics);
-            NetworkClient client = new NetworkClient(
-                    new Selector(config.getLong(ProducerConfig.CONNECTIONS_MAX_IDLE_MS_CONFIG),
-                            this.metrics, time, "producer", channelBuilder, logContext),
-                    this.metadata,
-                    clientId,
-                    maxInflightRequests,
-                    config.getLong(ProducerConfig.RECONNECT_BACKOFF_MS_CONFIG),
-                    config.getLong(ProducerConfig.RECONNECT_BACKOFF_MAX_MS_CONFIG),
-                    config.getInt(ProducerConfig.SEND_BUFFER_CONFIG),
-                    config.getInt(ProducerConfig.RECEIVE_BUFFER_CONFIG),
-                    this.requestTimeoutMs,
-                    time,
-                    true,
-                    apiVersions,
-                    throttleTimeSensor,
-                    logContext);
-            this.sender = new Sender(logContext,
-                    client,
-                    this.metadata,
-                    this.accumulator,
-                    maxInflightRequests == 1,
-                    config.getInt(ProducerConfig.MAX_REQUEST_SIZE_CONFIG),
-                    acks,
-                    retries,
-                    metricsRegistry.senderMetrics,
-                    Time.SYSTEM,
-                    this.requestTimeoutMs,
-                    config.getLong(ProducerConfig.RETRY_BACKOFF_MS_CONFIG),
-                    this.transactionManager,
-                    apiVersions);
+                    transactionManager,
+                    new BufferPool(this.totalMemorySize, config.getInt(ProducerConfig.BATCH_SIZE_CONFIG), metrics, time, PRODUCER_METRIC_GROUP_NAME));
+
+            List addresses = ClientUtils.parseAndValidateAddresses(
+                    config.getList(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG),
+                    config.getString(ProducerConfig.CLIENT_DNS_LOOKUP_CONFIG));
+            if (metadata != null) {
+                this.metadata = metadata;
+            } else {
+                this.metadata = new ProducerMetadata(retryBackoffMs,
+                        config.getLong(ProducerConfig.METADATA_MAX_AGE_CONFIG),
+                        config.getLong(ProducerConfig.METADATA_MAX_IDLE_CONFIG),
+                        logContext,
+                        clusterResourceListeners,
+                        Time.SYSTEM);
+                this.metadata.bootstrap(addresses);
+            }
+            this.errors = this.metrics.sensor("errors");
+            this.sender = newSender(logContext, kafkaClient, this.metadata);
             String ioThreadName = NETWORK_THREAD_PREFIX + " | " + clientId;
             this.ioThread = new KafkaThread(ioThreadName, this.sender, true);
             this.ioThread.start();
-            this.errors = this.metrics.sensor("errors");
             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
-            close(0, TimeUnit.MILLISECONDS, true);
+            close(Duration.ofMillis(0), true);
             // now propagate the exception
             throw new KafkaException("Failed to construct kafka producer", t);
         }
     }
 
-    private static TransactionManager configureTransactionState(ProducerConfig config, LogContext logContext, Logger log) {
-
-        TransactionManager transactionManager = null;
+    // visible for testing
+    Sender newSender(LogContext logContext, KafkaClient kafkaClient, ProducerMetadata metadata) {
+        int maxInflightRequests = configureInflightRequests(producerConfig);
+        int requestTimeoutMs = producerConfig.getInt(ProducerConfig.REQUEST_TIMEOUT_MS_CONFIG);
+        ChannelBuilder channelBuilder = ClientUtils.createChannelBuilder(producerConfig, time, logContext);
+        ProducerMetrics metricsRegistry = new ProducerMetrics(this.metrics);
+        Sensor throttleTimeSensor = Sender.throttleTimeSensor(metricsRegistry.senderMetrics);
+        KafkaClient client = kafkaClient != null ? kafkaClient : new NetworkClient(
+                new Selector(producerConfig.getLong(ProducerConfig.CONNECTIONS_MAX_IDLE_MS_CONFIG),
+                        this.metrics, time, "producer", channelBuilder, logContext),
+                metadata,
+                clientId,
+                maxInflightRequests,
+                producerConfig.getLong(ProducerConfig.RECONNECT_BACKOFF_MS_CONFIG),
+                producerConfig.getLong(ProducerConfig.RECONNECT_BACKOFF_MAX_MS_CONFIG),
+                producerConfig.getInt(ProducerConfig.SEND_BUFFER_CONFIG),
+                producerConfig.getInt(ProducerConfig.RECEIVE_BUFFER_CONFIG),
+                requestTimeoutMs,
+                producerConfig.getLong(ProducerConfig.SOCKET_CONNECTION_SETUP_TIMEOUT_MS_CONFIG),
+                producerConfig.getLong(ProducerConfig.SOCKET_CONNECTION_SETUP_TIMEOUT_MAX_MS_CONFIG),
+                ClientDnsLookup.forConfig(producerConfig.getString(ProducerConfig.CLIENT_DNS_LOOKUP_CONFIG)),
+                time,
+                true,
+                apiVersions,
+                throttleTimeSensor,
+                logContext);
+        short acks = configureAcks(producerConfig, log);
+        return new Sender(logContext,
+                client,
+                metadata,
+                this.accumulator,
+                maxInflightRequests == 1,
+                producerConfig.getInt(ProducerConfig.MAX_REQUEST_SIZE_CONFIG),
+                acks,
+                producerConfig.getInt(ProducerConfig.RETRIES_CONFIG),
+                metricsRegistry.senderMetrics,
+                time,
+                requestTimeoutMs,
+                producerConfig.getLong(ProducerConfig.RETRY_BACKOFF_MS_CONFIG),
+                this.transactionManager,
+                apiVersions);
+    }
 
-        boolean userConfiguredIdempotence = false;
-        if (config.originals().containsKey(ProducerConfig.ENABLE_IDEMPOTENCE_CONFIG))
-            userConfiguredIdempotence = true;
+    private static int lingerMs(ProducerConfig config) {
+        return (int) Math.min(config.getLong(ProducerConfig.LINGER_MS_CONFIG), Integer.MAX_VALUE);
+    }
 
-        boolean userConfiguredTransactions = false;
-        if (config.originals().containsKey(ProducerConfig.TRANSACTIONAL_ID_CONFIG))
-            userConfiguredTransactions = true;
+    private static int configureDeliveryTimeout(ProducerConfig config, Logger log) {
+        int deliveryTimeoutMs = config.getInt(ProducerConfig.DELIVERY_TIMEOUT_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 < 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
+                    + " should be equal to or larger than " + ProducerConfig.LINGER_MS_CONFIG
+                    + " + " + ProducerConfig.REQUEST_TIMEOUT_MS_CONFIG);
+            } else {
+                // override deliveryTimeoutMs default value to lingerMs + requestTimeoutMs for backward compatibility
+                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);
+            }
+        }
+        return deliveryTimeoutMs;
+    }
 
-        boolean idempotenceEnabled = config.getBoolean(ProducerConfig.ENABLE_IDEMPOTENCE_CONFIG);
+    private TransactionManager configureTransactionState(ProducerConfig config,
+                                                         LogContext logContext) {
 
-        if (!idempotenceEnabled && userConfiguredIdempotence && userConfiguredTransactions)
-            throw new ConfigException("Cannot set a " + ProducerConfig.TRANSACTIONAL_ID_CONFIG + " without also enabling idempotence.");
+        TransactionManager transactionManager = null;
 
-        if (userConfiguredTransactions)
-            idempotenceEnabled = true;
+        final boolean userConfiguredIdempotence = config.originals().containsKey(ProducerConfig.ENABLE_IDEMPOTENCE_CONFIG);
+        final boolean userConfiguredTransactions = config.originals().containsKey(ProducerConfig.TRANSACTIONAL_ID_CONFIG);
+        if (userConfiguredTransactions && !userConfiguredIdempotence)
+            log.info("Overriding the default {} to true since {} is specified.", ProducerConfig.ENABLE_IDEMPOTENCE_CONFIG,
+                    ProducerConfig.TRANSACTIONAL_ID_CONFIG);
+
+        if (config.idempotenceEnabled()) {
+            final String transactionalId = config.getString(ProducerConfig.TRANSACTIONAL_ID_CONFIG);
+            final int transactionTimeoutMs = config.getInt(ProducerConfig.TRANSACTION_TIMEOUT_CONFIG);
+            final long retryBackoffMs = config.getLong(ProducerConfig.RETRY_BACKOFF_MS_CONFIG);
+            final boolean autoDowngradeTxnCommit = config.getBoolean(ProducerConfig.AUTO_DOWNGRADE_TXN_COMMIT);
+            transactionManager = new TransactionManager(
+                logContext,
+                transactionalId,
+                transactionTimeoutMs,
+                retryBackoffMs,
+                apiVersions,
+                autoDowngradeTxnCommit);
 
-        if (idempotenceEnabled) {
-            String transactionalId = config.getString(ProducerConfig.TRANSACTIONAL_ID_CONFIG);
-            int transactionTimeoutMs = config.getInt(ProducerConfig.TRANSACTION_TIMEOUT_CONFIG);
-            long retryBackoffMs = config.getLong(ProducerConfig.RETRY_BACKOFF_MS_CONFIG);
-            transactionManager = new TransactionManager(logContext, transactionalId, transactionTimeoutMs, retryBackoffMs);
             if (transactionManager.isTransactional())
                 log.info("Instantiated a transactional producer.");
             else
                 log.info("Instantiated an idempotent producer.");
         }
-
         return transactionManager;
     }
 
-    private static int configureRetries(ProducerConfig config, boolean idempotenceEnabled, Logger log) {
-        boolean userConfiguredRetries = false;
-        if (config.originals().containsKey(ProducerConfig.RETRIES_CONFIG)) {
-            userConfiguredRetries = true;
-        }
-        if (idempotenceEnabled && !userConfiguredRetries) {
-            // We recommend setting infinite retries when the idempotent producer is enabled, so it makes sense to make
-            // this the default.
-            log.info("Overriding the default retries config to the recommended value of {} since the idempotent " +
-                    "producer is enabled.", Integer.MAX_VALUE);
-            return Integer.MAX_VALUE;
-        }
-        if (idempotenceEnabled && config.getInt(ProducerConfig.RETRIES_CONFIG) == 0) {
-            throw new ConfigException("Must set " + ProducerConfig.RETRIES_CONFIG + " to non-zero when using the idempotent producer.");
-        }
-        return config.getInt(ProducerConfig.RETRIES_CONFIG);
-    }
-
-    private static int configureInflightRequests(ProducerConfig config, boolean idempotenceEnabled) {
-        if (idempotenceEnabled && 5 < config.getInt(ProducerConfig.MAX_IN_FLIGHT_REQUESTS_PER_CONNECTION)) {
+    private static int configureInflightRequests(ProducerConfig config) {
+        if (config.idempotenceEnabled() && 5 < config.getInt(ProducerConfig.MAX_IN_FLIGHT_REQUESTS_PER_CONNECTION)) {
             throw new ConfigException("Must set " + ProducerConfig.MAX_IN_FLIGHT_REQUESTS_PER_CONNECTION + " to at most 5" +
                     " to use the idempotent producer.");
         }
         return config.getInt(ProducerConfig.MAX_IN_FLIGHT_REQUESTS_PER_CONNECTION);
     }
 
-    private static short configureAcks(ProducerConfig config, boolean idempotenceEnabled, Logger log) {
-        boolean userConfiguredAcks = false;
-        short acks = (short) parseAcks(config.getString(ProducerConfig.ACKS_CONFIG));
-        if (config.originals().containsKey(ProducerConfig.ACKS_CONFIG)) {
-            userConfiguredAcks = true;
-        }
+    private static short configureAcks(ProducerConfig config, Logger log) {
+        boolean userConfiguredAcks = config.originals().containsKey(ProducerConfig.ACKS_CONFIG);
+        short acks = Short.parseShort(config.getString(ProducerConfig.ACKS_CONFIG));
 
-        if (idempotenceEnabled && !userConfiguredAcks) {
-            log.info("Overriding the default {} to all since idempotence is enabled.", ProducerConfig.ACKS_CONFIG);
-            return -1;
-        }
-
-        if (idempotenceEnabled && acks != -1) {
-            throw new ConfigException("Must set " + ProducerConfig.ACKS_CONFIG + " to all in order to use the idempotent " +
-                    "producer. Otherwise we cannot guarantee idempotence.");
+        if (config.idempotenceEnabled()) {
+            if (!userConfiguredAcks)
+                log.info("Overriding the default {} to all since idempotence is enabled.", ProducerConfig.ACKS_CONFIG);
+            else if (acks != -1)
+                throw new ConfigException("Must set " + ProducerConfig.ACKS_CONFIG + " to all in order to use the idempotent " +
+                        "producer. Otherwise we cannot guarantee idempotence.");
         }
         return acks;
     }
 
-    private static int parseAcks(String acksString) {
-        try {
-            return acksString.trim().equalsIgnoreCase("all") ? -1 : Integer.parseInt(acksString.trim());
-        } catch (NumberFormatException e) {
-            throw new ConfigException("Invalid configuration value for 'acks': " + acksString);
-        }
-    }
-
     /**
      * Needs to be called before any other methods when the transactional.id is set in the configuration.
      *
@@ -534,18 +576,26 @@ private static int parseAcks(String acksString) {
      *   2. Gets the internal producer id and epoch, used in all future transactional
      *      messages issued by the producer.
      *
+     * Note that this method will raise {@link TimeoutException} if the transactional state cannot
+     * be initialized before expiration of {@code max.block.ms}. Additionally, it will raise {@link InterruptException}
+     * if interrupted. It is safe to retry in either case, but once the transactional state has been successfully
+     * initialized, this method should no longer be used.
+     *
      * @throws IllegalStateException if no transactional.id has been configured
      * @throws org.apache.kafka.common.errors.UnsupportedVersionException fatal error indicating the broker
      *         does not support transactions (i.e. if its version is lower than 0.11.0.0)
      * @throws org.apache.kafka.common.errors.AuthorizationException fatal error indicating that the configured
      *         transactional.id is not authorized. See the exception for more details
      * @throws KafkaException if the producer has encountered a previous fatal error or for any other unexpected error
+     * @throws TimeoutException if the time taken for initialize the transaction has surpassed max.block.ms.
+     * @throws InterruptException if the thread is interrupted while blocked
      */
     public void initTransactions() {
         throwIfNoTransactionManager();
+        throwIfProducerClosed();
         TransactionalRequestResult result = transactionManager.initializeTransactions();
         sender.wakeup();
-        result.await();
+        result.await(maxBlockTimeMs, TimeUnit.MILLISECONDS);
     }
 
     /**
@@ -555,6 +605,8 @@ public void initTransactions() {
      * @throws IllegalStateException if no transactional.id has been configured or if {@link #initTransactions()}
      *         has not yet been invoked
      * @throws ProducerFencedException if another producer with the same transactional.id is active
+     * @throws org.apache.kafka.common.errors.InvalidProducerEpochException if the producer has attempted to produce with an old epoch
+     *         to the partition leader. See the exception for more details
      * @throws org.apache.kafka.common.errors.UnsupportedVersionException fatal error indicating the broker
      *         does not support transactions (i.e. if its version is lower than 0.11.0.0)
      * @throws org.apache.kafka.common.errors.AuthorizationException fatal error indicating that the configured
@@ -563,6 +615,7 @@ public void initTransactions() {
      */
     public void beginTransaction() throws ProducerFencedException {
         throwIfNoTransactionManager();
+        throwIfProducerClosed();
         transactionManager.beginTransaction();
     }
 
@@ -579,23 +632,72 @@ public void beginTransaction() throws ProducerFencedException {
      * and should also not commit offsets manually (via {@link KafkaConsumer#commitSync(Map) sync} or
      * {@link KafkaConsumer#commitAsync(Map, OffsetCommitCallback) async} commits).
      *
-     * @throws IllegalStateException if no transactional.id has been configured or no transaction has been started
+     * @throws IllegalStateException if no transactional.id has been configured, no transaction has been started
      * @throws ProducerFencedException fatal error indicating another producer with the same transactional.id is active
      * @throws org.apache.kafka.common.errors.UnsupportedVersionException fatal error indicating the broker
      *         does not support transactions (i.e. if its version is lower than 0.11.0.0)
-     * @throws org.apache.kafka.common.errors.UnsupportedForMessageFormatException  fatal error indicating the message
+     * @throws org.apache.kafka.common.errors.UnsupportedForMessageFormatException fatal error indicating the message
      *         format used for the offsets topic on the broker does not support transactions
      * @throws org.apache.kafka.common.errors.AuthorizationException fatal error indicating that the configured
-     *         transactional.id is not authorized. See the exception for more details
+     *         transactional.id is not authorized, or the consumer group id is not authorized.
+     * @throws org.apache.kafka.common.errors.InvalidProducerEpochException if the producer has attempted to produce with an old epoch
+     *         to the partition leader. See the exception for more details
      * @throws KafkaException if the producer has encountered a previous fatal or abortable error, or for any
      *         other unexpected error
      */
     public void sendOffsetsToTransaction(Map offsets,
                                          String consumerGroupId) throws ProducerFencedException {
+        sendOffsetsToTransaction(offsets, new ConsumerGroupMetadata(consumerGroupId));
+    }
+
+    /**
+     * Sends a list of specified offsets to the consumer group coordinator, and also marks
+     * those offsets as part of the current transaction. These offsets will be considered
+     * committed only if the transaction is committed successfully. The committed offset should
+     * be the next message your application will consume, i.e. lastProcessedMessageOffset + 1.
+     * 
    
+     * This method should be used when you need to batch consumed and produced messages
+     * together, typically in a consume-transform-produce pattern. Thus, the specified
+     * {@code groupMetadata} should be extracted from the used {@link KafkaConsumer consumer} via
+     * {@link KafkaConsumer#groupMetadata()} to leverage consumer group metadata for stronger fencing than
+     * {@link #sendOffsetsToTransaction(Map, String)} which only sends with consumer group id.
+     *
+     * 
    
+     * Note, that the consumer should have {@code enable.auto.commit=false} and should
+     * also not commit offsets manually (via {@link KafkaConsumer#commitSync(Map) sync} or
+     * {@link KafkaConsumer#commitAsync(Map, OffsetCommitCallback) async} commits).
+     * This method will raise {@link TimeoutException} if the producer cannot send offsets before expiration of {@code max.block.ms}.
+     * Additionally, it will raise {@link InterruptException} if interrupted.
+     *
+     * @throws IllegalStateException if no transactional.id has been configured or no transaction has been started.
+     * @throws ProducerFencedException fatal error indicating another producer with the same transactional.id is active
+     * @throws org.apache.kafka.common.errors.UnsupportedVersionException fatal error indicating the broker
+     *         does not support transactions (i.e. if its version is lower than 0.11.0.0) or
+     *         the broker doesn't support latest version of transactional API with consumer group metadata (i.e. if its version is
+     *         lower than 2.5.0).
+     * @throws org.apache.kafka.common.errors.UnsupportedForMessageFormatException fatal error indicating the message
+     *         format used for the offsets topic on the broker does not support transactions
+     * @throws org.apache.kafka.common.errors.AuthorizationException fatal error indicating that the configured
+     *         transactional.id is not authorized, or the consumer group id is not authorized.
+     * @throws org.apache.kafka.clients.consumer.CommitFailedException if the commit failed and cannot be retried
+     *         (e.g. if the consumer has been kicked out of the group). Users should handle this by aborting the transaction.
+     * @throws org.apache.kafka.common.errors.FencedInstanceIdException if this producer instance gets fenced by broker due to a
+     *                                                                  mis-configured consumer instance id within group metadata.
+     * @throws org.apache.kafka.common.errors.InvalidProducerEpochException if the producer has attempted to produce with an old epoch
+     *         to the partition leader. See the exception for more details
+     * @throws KafkaException if the producer has encountered a previous fatal or abortable error, or for any
+     *         other unexpected error
+     * @throws TimeoutException if the time taken for sending offsets has surpassed max.block.ms.
+     * @throws InterruptException if the thread is interrupted while blocked
+     */
+    public void sendOffsetsToTransaction(Map offsets,
+                                         ConsumerGroupMetadata groupMetadata) throws ProducerFencedException {
+        throwIfInvalidGroupMetadata(groupMetadata);
         throwIfNoTransactionManager();
-        TransactionalRequestResult result = transactionManager.sendOffsetsToTransaction(offsets, consumerGroupId);
+        throwIfProducerClosed();
+        TransactionalRequestResult result = transactionManager.sendOffsetsToTransaction(offsets, groupMetadata);
         sender.wakeup();
-        result.await();
+        result.await(maxBlockTimeMs, TimeUnit.MILLISECONDS);
     }
 
     /**
@@ -605,20 +707,30 @@ public void sendOffsetsToTransaction(Map offs
      * errors, this method will throw the last received exception immediately and the transaction will not be committed.
      * So all {@link #send(ProducerRecord)} calls in a transaction must succeed in order for this method to succeed.
      *
+     * Note that this method will raise {@link TimeoutException} if the transaction cannot be committed before expiration
+     * of {@code max.block.ms}. Additionally, it will raise {@link InterruptException} if interrupted.
+     * It is safe to retry in either case, but it is not possible to attempt a different operation (such as abortTransaction)
+     * since the commit may already be in the progress of completing. If not retrying, the only option is to close the producer.
+     *
      * @throws IllegalStateException if no transactional.id has been configured or no transaction has been started
      * @throws ProducerFencedException fatal error indicating another producer with the same transactional.id is active
      * @throws org.apache.kafka.common.errors.UnsupportedVersionException fatal error indicating the broker
      *         does not support transactions (i.e. if its version is lower than 0.11.0.0)
      * @throws org.apache.kafka.common.errors.AuthorizationException fatal error indicating that the configured
      *         transactional.id is not authorized. See the exception for more details
+     * @throws org.apache.kafka.common.errors.InvalidProducerEpochException if the producer has attempted to produce with an old epoch
+     *         to the partition leader. See the exception for more details
      * @throws KafkaException if the producer has encountered a previous fatal or abortable error, or for any
      *         other unexpected error
+     * @throws TimeoutException if the time taken for committing the transaction has surpassed max.block.ms.
+     * @throws InterruptException if the thread is interrupted while blocked
      */
     public void commitTransaction() throws ProducerFencedException {
         throwIfNoTransactionManager();
+        throwIfProducerClosed();
         TransactionalRequestResult result = transactionManager.beginCommit();
         sender.wakeup();
-        result.await();
+        result.await(maxBlockTimeMs, TimeUnit.MILLISECONDS);
     }
 
     /**
@@ -626,19 +738,30 @@ public void commitTransaction() throws ProducerFencedException {
      * This call will throw an exception immediately if any prior {@link #send(ProducerRecord)} calls failed with a
      * {@link ProducerFencedException} or an instance of {@link org.apache.kafka.common.errors.AuthorizationException}.
      *
+     * Note that this method will raise {@link TimeoutException} if the transaction cannot be aborted before expiration
+     * of {@code max.block.ms}. Additionally, it will raise {@link InterruptException} if interrupted.
+     * It is safe to retry in either case, but it is not possible to attempt a different operation (such as commitTransaction)
+     * since the abort may already be in the progress of completing. If not retrying, the only option is to close the producer.
+     *
      * @throws IllegalStateException if no transactional.id has been configured or no transaction has been started
      * @throws ProducerFencedException fatal error indicating another producer with the same transactional.id is active
+     * @throws org.apache.kafka.common.errors.InvalidProducerEpochException if the producer has attempted to produce with an old epoch
+     *         to the partition leader. See the exception for more details
      * @throws org.apache.kafka.common.errors.UnsupportedVersionException fatal error indicating the broker
      *         does not support transactions (i.e. if its version is lower than 0.11.0.0)
      * @throws org.apache.kafka.common.errors.AuthorizationException fatal error indicating that the configured
      *         transactional.id is not authorized. See the exception for more details
      * @throws KafkaException if the producer has encountered a previous fatal error or for any other unexpected error
+     * @throws TimeoutException if the time taken for aborting the transaction has surpassed max.block.ms.
+     * @throws InterruptException if the thread is interrupted while blocked
      */
     public void abortTransaction() throws ProducerFencedException {
         throwIfNoTransactionManager();
+        throwIfProducerClosed();
+        log.info("Aborting incomplete transaction");
         TransactionalRequestResult result = transactionManager.beginAbort();
         sender.wakeup();
-        result.await();
+        result.await(maxBlockTimeMs, TimeUnit.MILLISECONDS);
     }
 
     /**
@@ -749,28 +872,44 @@ public Future send(ProducerRecord record) {
      *
      * @throws AuthenticationException if authentication fails. See the exception for more details
      * @throws AuthorizationException fatal error indicating that the producer is not allowed to write
-     * @throws IllegalStateException if a transactional.id has been configured and no transaction has been started
+     * @throws IllegalStateException if a transactional.id has been configured and no transaction has been started, or
+     *                               when send is invoked after producer has been closed.
      * @throws InterruptException If the thread is interrupted while blocked
      * @throws SerializationException If the key or value are not valid objects given the configured serializers
-     * @throws TimeoutException If the time taken for fetching metadata or allocating memory for the record has surpassed max.block.ms.
      * @throws KafkaException If a Kafka related error occurs that does not belong to the public API exceptions.
-     *
      */
     @Override
     public Future send(ProducerRecord record, Callback callback) {
         // intercept the record, which can be potentially modified; this method does not throw exceptions
-        ProducerRecord interceptedRecord = this.interceptors == null ? record : this.interceptors.onSend(record);
+        ProducerRecord interceptedRecord = this.interceptors.onSend(record);
         return doSend(interceptedRecord, callback);
     }
 
+    // Verify that this producer instance has not been closed. This method throws IllegalStateException if the producer
+    // has already been closed.
+    private void throwIfProducerClosed() {
+        if (sender == null || !sender.isRunning())
+            throw new IllegalStateException("Cannot perform operation after producer has been closed");
+    }
+
     /**
      * Implementation of asynchronously send a record to a topic.
      */
     private Future doSend(ProducerRecord record, Callback callback) {
         TopicPartition tp = null;
         try {
+            throwIfProducerClosed();
             // first make sure the metadata for the topic is available
-            ClusterAndWaitTime clusterAndWaitTime = waitOnMetadata(record.topic(), record.partition(), maxBlockTimeMs);
+            long nowMs = time.milliseconds();
+            ClusterAndWaitTime clusterAndWaitTime;
+            try {
+                clusterAndWaitTime = waitOnMetadata(record.topic(), record.partition(), nowMs, maxBlockTimeMs);
+            } catch (KafkaException e) {
+                if (metadata.isClosed())
+                    throw new KafkaException("Producer closed while send in progress", e);
+                throw e;
+            }
+            nowMs += clusterAndWaitTime.waitedOnMetadataMs;
             long remainingWaitMs = Math.max(0, maxBlockTimeMs - clusterAndWaitTime.waitedOnMetadataMs);
             Cluster cluster = clusterAndWaitTime.cluster;
             byte[] serializedKey;
@@ -798,16 +937,37 @@ private Future doSend(ProducerRecord record, Callback call
             int serializedSize = AbstractRecords.estimateSizeInBytesUpperBound(apiVersions.maxUsableProduceMagic(),
                     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);
+            long timestamp = record.timestamp() == null ? nowMs : record.timestamp();
+            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 = this.interceptors == null ? callback : new InterceptorCallback<>(callback, this.interceptors, tp);
+            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, nowMs);
+
+            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, nowMs);
+            }
 
             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();
@@ -821,29 +981,19 @@ private Future doSend(ProducerRecord record, Callback call
             if (callback != null)
                 callback.onCompletion(null, e);
             this.errors.record();
-            if (this.interceptors != null)
-                this.interceptors.onSendError(record, tp, e);
+            this.interceptors.onSendError(record, tp, e);
             return new FutureFailure(e);
         } catch (InterruptedException e) {
             this.errors.record();
-            if (this.interceptors != null)
-                this.interceptors.onSendError(record, tp, e);
+            this.interceptors.onSendError(record, tp, e);
             throw new InterruptException(e);
-        } catch (BufferExhaustedException e) {
-            this.errors.record();
-            this.metrics.sensor("buffer-exhausted-records").record();
-            if (this.interceptors != null)
-                this.interceptors.onSendError(record, tp, e);
-            throw e;
         } catch (KafkaException e) {
             this.errors.record();
-            if (this.interceptors != null)
-                this.interceptors.onSendError(record, tp, e);
+            this.interceptors.onSendError(record, tp, e);
             throw e;
         } catch (Exception e) {
             // we notify interceptor about all exceptions, since onSend is called before anything else in this method
-            if (this.interceptors != null)
-                this.interceptors.onSendError(record, tp, e);
+            this.interceptors.onSendError(record, tp, e);
             throw e;
         }
     }
@@ -858,51 +1008,62 @@ private void setReadOnly(Headers headers) {
      * Wait for cluster metadata including partitions for the given topic to be available.
      * @param topic The topic we want metadata for
      * @param partition A specific partition expected to exist in metadata, or null if there's no preference
+     * @param nowMs The current time in ms
      * @param maxWaitMs The maximum time in ms for waiting on the metadata
      * @return The cluster containing topic metadata and the amount of time we waited in ms
+     * @throws TimeoutException if metadata could not be refreshed within {@code max.block.ms}
+     * @throws KafkaException for all Kafka-related exceptions, including the case where this method is called after producer close
      */
-    private ClusterAndWaitTime waitOnMetadata(String topic, Integer partition, long maxWaitMs) throws InterruptedException {
+    private ClusterAndWaitTime waitOnMetadata(String topic, Integer partition, long nowMs, long maxWaitMs) throws InterruptedException {
         // add topic to metadata topic list if it is not there already and reset expiry
-        metadata.add(topic);
         Cluster cluster = metadata.fetch();
+
+        if (cluster.invalidTopics().contains(topic))
+            throw new InvalidTopicException(topic);
+
+        metadata.add(topic, nowMs);
+
         Integer partitionsCount = cluster.partitionCountForTopic(topic);
         // Return cached metadata if we have it, and if the record's partition is either undefined
         // or within the known partition range
         if (partitionsCount != null && (partition == null || partition < partitionsCount))
             return new ClusterAndWaitTime(cluster, 0);
 
-        long begin = time.milliseconds();
         long remainingWaitMs = maxWaitMs;
-        long elapsed;
-        // Issue metadata requests until we have metadata for the topic or maxWaitTimeMs is exceeded.
-        // In case we already have cached metadata for the topic, but the requested partition is greater
-        // than expected, issue an update request only once. This is necessary in case the metadata
+        long elapsed = 0;
+        // Issue metadata requests until we have metadata for the topic and the requested partition,
+        // or until maxWaitTimeMs is exceeded. This is necessary in case the metadata
         // is stale and the number of partitions for this topic has increased in the meantime.
         do {
-            log.trace("Requesting metadata update for topic {}.", topic);
-            metadata.add(topic);
-            int version = metadata.requestUpdate();
+            if (partition != null) {
+                log.trace("Requesting metadata update for partition {} of topic {}.", partition, topic);
+            } else {
+                log.trace("Requesting metadata update for topic {}.", topic);
+            }
+            metadata.add(topic, nowMs + elapsed);
+            int version = metadata.requestUpdateForTopic(topic);
             sender.wakeup();
             try {
                 metadata.awaitUpdate(version, remainingWaitMs);
             } catch (TimeoutException ex) {
                 // Rethrow with original maxWaitMs to prevent logging exception with remainingWaitMs
-                throw new TimeoutException("Failed to update metadata after " + maxWaitMs + " ms.");
+                throw new TimeoutException(
+                        String.format("Topic %s not present in metadata after %d ms.",
+                                topic, maxWaitMs));
             }
             cluster = metadata.fetch();
-            elapsed = time.milliseconds() - begin;
-            if (elapsed >= maxWaitMs)
-                throw new TimeoutException("Failed to update metadata after " + maxWaitMs + " ms.");
-            if (cluster.unauthorizedTopics().contains(topic))
-                throw new TopicAuthorizationException(topic);
+            elapsed = time.milliseconds() - nowMs;
+            if (elapsed >= maxWaitMs) {
+                throw new TimeoutException(partitionsCount == null ?
+                        String.format("Topic %s not present in metadata after %d ms.",
+                                topic, maxWaitMs) :
+                        String.format("Partition %d of topic %s with partition count %d is not present in metadata after %d ms.",
+                                partition, topic, partitionsCount, maxWaitMs));
+            }
+            metadata.maybeThrowExceptionForTopic(topic);
             remainingWaitMs = maxWaitMs - elapsed;
             partitionsCount = cluster.partitionCountForTopic(topic);
-        } while (partitionsCount == null);
-
-        if (partition != null && partition >= partitionsCount) {
-            throw new KafkaException(
-                    String.format("Invalid partition given with record: %d is not in the range [0...%d).", partition, partitionsCount));
-        }
+        } while (partitionsCount == null || (partition != null && partition >= partitionsCount));
 
         return new ClusterAndWaitTime(cluster, elapsed);
     }
@@ -911,12 +1072,11 @@ private ClusterAndWaitTime waitOnMetadata(String topic, Integer partition, long
      * Validate that the record size isn't too large
      */
     private void ensureValidRecordSize(int size) {
-        if (size > this.maxRequestSize)
+        if (size > maxRequestSize)
             throw new RecordTooLargeException("The message is " + size +
-                    " bytes when serialized which is larger than the maximum request size you have configured with the " +
-                    ProducerConfig.MAX_REQUEST_SIZE_CONFIG +
-                    " configuration.");
-        if (size > this.totalMemorySize)
+                    " bytes when serialized which is larger than " + maxRequestSize + ", which is the value of the " +
+                    ProducerConfig.MAX_REQUEST_SIZE_CONFIG + " configuration.");
+        if (size > totalMemorySize)
             throw new RecordTooLargeException("The message is " + size +
                     " bytes when serialized which is larger than the total memory buffer you have configured with the " +
                     ProducerConfig.BUFFER_MEMORY_CONFIG +
@@ -942,7 +1102,7 @@ private void ensureValidRecordSize(int size) {
      * for(ConsumerRecord record: consumer.poll(100))
      *     producer.send(new ProducerRecord("my-topic", record.key(), record.value());
      * producer.flush();
-     * consumer.commit();
+     * consumer.commitSync();
      * }
      * 
    
      *
@@ -973,14 +1133,15 @@ public void flush() {
      * Get the partition metadata for the given topic. This can be used for custom partitioning.
      * @throws AuthenticationException if authentication fails. See the exception for more details
      * @throws AuthorizationException if not authorized to the specified topic. See the exception for more details
-     * @throws InterruptException If the thread is interrupted while blocked
+     * @throws InterruptException if the thread is interrupted while blocked
      * @throws TimeoutException if metadata could not be refreshed within {@code max.block.ms}
+     * @throws KafkaException for all Kafka-related exceptions, including the case where this method is called after producer close
      */
     @Override
     public List partitionsFor(String topic) {
         Objects.requireNonNull(topic, "topic cannot be null");
         try {
-            return waitOnMetadata(topic, null, maxBlockTimeMs).cluster.partitionsForTopic(topic);
+            return waitOnMetadata(topic, null, time.milliseconds(), maxBlockTimeMs).cluster.partitionsForTopic(topic);
         } catch (InterruptedException e) {
             throw new InterruptException(e);
         }
@@ -1003,55 +1164,62 @@ public List partitionsFor(String topic) {
      * block forever.
      * 
    
      *
-     * @throws InterruptException If the thread is interrupted while blocked
+     * @throws InterruptException If the thread is interrupted while blocked.
+     * @throws KafkaException If a unexpected error occurs while trying to close the client, this error should be treated
+     *                        as fatal and indicate the client is no longer functionable.
      */
     @Override
     public void close() {
-        close(Long.MAX_VALUE, TimeUnit.MILLISECONDS);
+        close(Duration.ofMillis(Long.MAX_VALUE));
     }
 
     /**
      * 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(0, TimeUnit.MILLISECONDS). This is done since no further sending will happen while
+     * close(Duration.ofMillis(0)). This is done since no further sending will happen while
      * blocking the I/O thread of the producer.
      *
      * @param timeout The maximum time to wait for producer to complete any pending requests. The value should be
      *                non-negative. Specifying a timeout of zero means do not wait for pending send requests to complete.
-     * @param timeUnit The time unit for the timeout
-     * @throws InterruptException If the thread is interrupted while blocked
+     * @throws InterruptException If the thread is interrupted while blocked.
+     * @throws KafkaException If a unexpected error occurs while trying to close the client, this error should be treated
+     *                        as fatal and indicate the client is no longer functionable.
      * @throws IllegalArgumentException If the timeout is negative.
+     *
      */
     @Override
-    public void close(long timeout, TimeUnit timeUnit) {
-        close(timeout, timeUnit, false);
+    public void close(Duration timeout) {
+        close(timeout, false);
     }
 
-    private void close(long timeout, TimeUnit timeUnit, boolean swallowException) {
-        if (timeout < 0)
+    private void close(Duration timeout, boolean swallowException) {
+        long timeoutMs = timeout.toMillis();
+        if (timeoutMs < 0)
             throw new IllegalArgumentException("The timeout cannot be negative.");
+        log.info("Closing the Kafka producer with timeoutMillis = {} ms.", timeoutMs);
 
-        log.info("Closing the Kafka producer with timeoutMillis = {} ms.", timeUnit.toMillis(timeout));
         // this will keep track of the first encountered exception
         AtomicReference firstException = new AtomicReference<>();
         boolean invokedFromCallback = Thread.currentThread() == this.ioThread;
-        if (timeout > 0) {
+        if (timeoutMs > 0) {
             if (invokedFromCallback) {
                 log.warn("Overriding close timeout {} ms to 0 ms in order to prevent useless blocking due to self-join. " +
-                        "This means you have incorrectly invoked close with a non-zero timeout from the producer call-back.", timeout);
+                        "This means you have incorrectly invoked close with a non-zero timeout from the producer call-back.",
+                        timeoutMs);
             } else {
                 // Try to close gracefully.
                 if (this.sender != null)
                     this.sender.initiateClose();
                 if (this.ioThread != null) {
                     try {
-                        this.ioThread.join(timeUnit.toMillis(timeout));
+                        this.ioThread.join(timeoutMs);
                     } catch (InterruptedException t) {
-                        firstException.compareAndSet(null, t);
+                        firstException.compareAndSet(null, new InterruptException(t));
                         log.error("Interrupted while joining ioThread", t);
                     }
                 }
@@ -1060,27 +1228,32 @@ private void close(long timeout, TimeUnit timeUnit, boolean swallowException) {
 
         if (this.sender != null && this.ioThread != null && this.ioThread.isAlive()) {
             log.info("Proceeding to force close the producer since pending requests could not be completed " +
-                    "within timeout {} ms.", timeout);
+                    "within timeout {} ms.", timeoutMs);
             this.sender.forceClose();
             // Only join the sender thread when not calling from callback.
             if (!invokedFromCallback) {
                 try {
                     this.ioThread.join();
                 } catch (InterruptedException e) {
-                    firstException.compareAndSet(null, e);
+                    firstException.compareAndSet(null, new InterruptException(e));
                 }
             }
         }
 
-        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);
+        Throwable exception = firstException.get();
+        if (exception != null && !swallowException) {
+            if (exception instanceof InterruptException) {
+                throw (InterruptException) exception;
+            }
+            throw new KafkaException("Failed to close kafka producer", exception);
+        }
         log.debug("Kafka producer has been closed");
-        if (firstException.get() != null && !swallowException)
-            throw new KafkaException("Failed to close kafka producer", firstException.get());
     }
 
     private ClusterResourceListeners configureClusterResourceListeners(Serializer keySerializer, Serializer valueSerializer, List... candidateLists) {
@@ -1106,12 +1279,26 @@ private int partition(ProducerRecord record, byte[] serializedKey, byte[]
                         record.topic(), record.key(), serializedKey, record.value(), serializedValue, cluster);
     }
 
+    private void throwIfInvalidGroupMetadata(ConsumerGroupMetadata groupMetadata) {
+        if (groupMetadata == null) {
+            throw new IllegalArgumentException("Consumer group metadata could not be null");
+        } else if (groupMetadata.generationId() > 0
+            && JoinGroupRequest.UNKNOWN_MEMBER_ID.equals(groupMetadata.memberId())) {
+            throw new IllegalArgumentException("Passed in group metadata " + groupMetadata + " has generationId > 0 but member.id ");
+        }
+    }
+
     private void throwIfNoTransactionManager() {
         if (transactionManager == null)
             throw new IllegalStateException("Cannot use transactional methods without enabling transactions " +
                     "by setting the " + ProducerConfig.TRANSACTIONAL_ID_CONFIG + " configuration property");
     }
 
+    // Visible for testing
+    String getClientId() {
+        return clientId;
+    }
+
     private static class ClusterAndWaitTime {
         final Cluster cluster;
         final long waitedOnMetadataMs;
@@ -1172,14 +1359,8 @@ private InterceptorCallback(Callback userCallback, ProducerInterceptors in
         }
 
         public void onCompletion(RecordMetadata metadata, Exception exception) {
-            if (this.interceptors != null) {
-                if (metadata == null) {
-                    this.interceptors.onAcknowledgement(new RecordMetadata(tp, -1, -1, RecordBatch.NO_TIMESTAMP,
-                                    Long.valueOf(-1L), -1, -1), exception);
-                } else {
-                    this.interceptors.onAcknowledgement(metadata, exception);
-                }
-            }
+            metadata = metadata != null ? metadata : new RecordMetadata(tp, -1, -1, RecordBatch.NO_TIMESTAMP, -1L, -1, -1);
+            this.interceptors.onAcknowledgement(metadata, exception);
             if (this.userCallback != null)
                 this.userCallback.onCompletion(metadata, exception);
         }
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 d2a84c66abeea..03c17d5bd4e5c 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,32 +16,31 @@
  */
 package org.apache.kafka.clients.producer;
 
+import org.apache.kafka.clients.consumer.ConsumerGroupMetadata;
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.clients.producer.internals.DefaultPartitioner;
 import org.apache.kafka.clients.producer.internals.FutureRecordMetadata;
 import org.apache.kafka.clients.producer.internals.ProduceRequestResult;
 import org.apache.kafka.common.Cluster;
+import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.Metric;
 import org.apache.kafka.common.MetricName;
 import org.apache.kafka.common.PartitionInfo;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.errors.ProducerFencedException;
 import org.apache.kafka.common.record.RecordBatch;
-import org.apache.kafka.common.serialization.ExtendedSerializer;
 import org.apache.kafka.common.serialization.Serializer;
+import org.apache.kafka.common.utils.Time;
 
+import java.time.Duration;
 import java.util.ArrayDeque;
 import java.util.ArrayList;
-import java.util.Collections;
 import java.util.Deque;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.Future;
-import java.util.concurrent.TimeUnit;
-
-import static org.apache.kafka.common.serialization.ExtendedSerializer.Wrapper.ensureExtended;
 
 /**
  * A mock of the producer interface you can use for testing code that uses Kafka.
@@ -59,8 +58,8 @@ public class MockProducer implements Producer {
     private final Map offsets;
     private final List>> consumerGroupOffsets;
     private Map> uncommittedConsumerGroupOffsets;
-    private final ExtendedSerializer keySerializer;
-    private final ExtendedSerializer valueSerializer;
+    private final Serializer keySerializer;
+    private final Serializer valueSerializer;
     private boolean autoComplete;
     private boolean closed;
     private boolean transactionInitialized;
@@ -70,6 +69,17 @@ public class MockProducer implements Producer {
     private boolean producerFenced;
     private boolean sentOffsets;
     private long commitCount = 0L;
+    private final Map mockMetrics;
+
+    public RuntimeException initTransactionException = null;
+    public RuntimeException beginTransactionException = null;
+    public RuntimeException sendOffsetsToTransactionException = null;
+    public RuntimeException commitTransactionException = null;
+    public RuntimeException abortTransactionException = null;
+    public RuntimeException sendException = null;
+    public RuntimeException flushException = null;
+    public RuntimeException partitionsForException = null;
+    public RuntimeException closeException = null;
 
     /**
      * Create a mock producer
@@ -77,7 +87,7 @@ public class MockProducer implements Producer {
      * @param cluster The cluster holding metadata for this producer
      * @param autoComplete If true automatically complete all requests successfully and execute the callback. Otherwise
      *        the user must call {@link #completeNext()} or {@link #errorNext(RuntimeException)} after
-     *        {@link #send(ProducerRecord) send()} to complete the call and unblock the @{link
+     *        {@link #send(ProducerRecord) send()} to complete the call and unblock the {@link
      *        java.util.concurrent.Future Future<RecordMetadata>} that is returned.
      * @param partitioner The partition strategy
      * @param keySerializer The serializer for key that implements {@link Serializer}.
@@ -91,14 +101,15 @@ public MockProducer(final Cluster cluster,
         this.cluster = cluster;
         this.autoComplete = autoComplete;
         this.partitioner = partitioner;
-        this.keySerializer = ensureExtended(keySerializer);
-        this.valueSerializer = ensureExtended(valueSerializer);
+        this.keySerializer = keySerializer;
+        this.valueSerializer = valueSerializer;
         this.offsets = new HashMap<>();
         this.sent = new ArrayList<>();
         this.uncommittedSends = new ArrayList<>();
         this.consumerGroupOffsets = new ArrayList<>();
         this.uncommittedConsumerGroupOffsets = new HashMap<>();
         this.completions = new ArrayDeque<>();
+        this.mockMetrics = new HashMap<>();
     }
 
     /**
@@ -139,13 +150,29 @@ public void initTransactions() {
         if (this.transactionInitialized) {
             throw new IllegalStateException("MockProducer has already been initialized for transactions.");
         }
+        if (this.initTransactionException != null) {
+            throw this.initTransactionException;
+        }
         this.transactionInitialized = true;
+        this.transactionInFlight = false;
+        this.transactionCommitted = false;
+        this.transactionAborted = false;
+        this.sentOffsets = false;
     }
 
     @Override
     public void beginTransaction() throws ProducerFencedException {
         verifyProducerState();
         verifyTransactionsInitialized();
+
+        if (this.beginTransactionException != null) {
+            throw this.beginTransactionException;
+        }
+
+        if (transactionInFlight) {
+            throw new IllegalStateException("Transaction already started");
+        }
+
         this.transactionInFlight = true;
         this.transactionCommitted = false;
         this.transactionAborted = false;
@@ -155,27 +182,40 @@ public void beginTransaction() throws ProducerFencedException {
     @Override
     public void sendOffsetsToTransaction(Map offsets,
                                          String consumerGroupId) throws ProducerFencedException {
+        Objects.requireNonNull(consumerGroupId);
         verifyProducerState();
         verifyTransactionsInitialized();
-        verifyNoTransactionInFlight();
-        Objects.requireNonNull(consumerGroupId);
+        verifyTransactionInFlight();
+
+        if (this.sendOffsetsToTransactionException != null) {
+            throw this.sendOffsetsToTransactionException;
+        }
+
         if (offsets.size() == 0) {
             return;
         }
-        Map uncommittedOffsets = this.uncommittedConsumerGroupOffsets.get(consumerGroupId);
-        if (uncommittedOffsets == null) {
-            uncommittedOffsets = new HashMap<>();
-            this.uncommittedConsumerGroupOffsets.put(consumerGroupId, uncommittedOffsets);
-        }
+        Map uncommittedOffsets =
+            this.uncommittedConsumerGroupOffsets.computeIfAbsent(consumerGroupId, k -> new HashMap<>());
         uncommittedOffsets.putAll(offsets);
         this.sentOffsets = true;
     }
 
+    @Override
+    public void sendOffsetsToTransaction(Map offsets,
+                                         ConsumerGroupMetadata groupMetadata) throws ProducerFencedException {
+        Objects.requireNonNull(groupMetadata);
+        sendOffsetsToTransaction(offsets, groupMetadata.groupId());
+    }
+
     @Override
     public void commitTransaction() throws ProducerFencedException {
         verifyProducerState();
         verifyTransactionsInitialized();
-        verifyNoTransactionInFlight();
+        verifyTransactionInFlight();
+
+        if (this.commitTransactionException != null) {
+            throw this.commitTransactionException;
+        }
 
         flush();
 
@@ -196,7 +236,12 @@ public void commitTransaction() throws ProducerFencedException {
     public void abortTransaction() throws ProducerFencedException {
         verifyProducerState();
         verifyTransactionsInitialized();
-        verifyNoTransactionInFlight();
+        verifyTransactionInFlight();
+
+        if (this.abortTransactionException != null) {
+            throw this.abortTransactionException;
+        }
+
         flush();
         this.uncommittedSends.clear();
         this.uncommittedConsumerGroupOffsets.clear();
@@ -205,7 +250,7 @@ public void abortTransaction() throws ProducerFencedException {
         this.transactionInFlight = false;
     }
 
-    private void verifyProducerState() {
+    private synchronized void verifyProducerState() {
         if (this.closed) {
             throw new IllegalStateException("MockProducer is already closed.");
         }
@@ -220,7 +265,7 @@ private void verifyTransactionsInitialized() {
         }
     }
 
-    private void verifyNoTransactionInFlight() {
+    private void verifyTransactionInFlight() {
         if (!this.transactionInFlight) {
             throw new IllegalStateException("There is no open transaction.");
         }
@@ -243,16 +288,33 @@ public synchronized Future send(ProducerRecord record) {
      */
     @Override
     public synchronized Future send(ProducerRecord record, Callback callback) {
-        verifyProducerState();
+        if (this.closed) {
+            throw new IllegalStateException("MockProducer is already closed.");
+        }
+
+        if (this.producerFenced) {
+            throw new KafkaException("MockProducer is fenced.", new ProducerFencedException("Fenced"));
+        }
+        if (this.sendException != null) {
+            throw this.sendException;
+        }
+
         int partition = 0;
         if (!this.cluster.partitionsForTopic(record.topic()).isEmpty())
             partition = partition(record, this.cluster);
+        else {
+            //just to throw ClassCastException if serializers are not the proper ones to serialize key/value
+            keySerializer.serialize(record.topic(), record.key());
+            valueSerializer.serialize(record.topic(), record.value());
+        }
+            
         TopicPartition topicPartition = new TopicPartition(record.topic(), partition);
         ProduceRequestResult result = new ProduceRequestResult(topicPartition);
-        FutureRecordMetadata future = new FutureRecordMetadata(result, 0, RecordBatch.NO_TIMESTAMP, 0L, 0, 0);
+        FutureRecordMetadata future = new FutureRecordMetadata(result, 0, RecordBatch.NO_TIMESTAMP,
+                0L, 0, 0, Time.SYSTEM);
         long offset = nextOffset(topicPartition);
         Completion completion = new Completion(offset, new RecordMetadata(topicPartition, 0, offset,
-                RecordBatch.NO_TIMESTAMP, Long.valueOf(0L), 0, 0), result, callback);
+                RecordBatch.NO_TIMESTAMP, 0L, 0, 0), result, callback);
 
         if (!this.transactionInFlight)
             this.sent.add(record);
@@ -284,28 +346,45 @@ private long nextOffset(TopicPartition tp) {
 
     public synchronized void flush() {
         verifyProducerState();
+
+        if (this.flushException != null) {
+            throw this.flushException;
+        }
+
         while (!this.completions.isEmpty())
             completeNext();
     }
 
     public List partitionsFor(String topic) {
+        if (this.partitionsForException != null) {
+            throw this.partitionsForException;
+        }
+
         return this.cluster.partitionsForTopic(topic);
     }
 
     public Map metrics() {
-        return Collections.emptyMap();
+        return mockMetrics;
+    }
+
+    /**
+     * Set a mock metric for testing purpose
+     */
+    public void setMockMetrics(MetricName name, Metric metric) {
+        mockMetrics.put(name, metric);
     }
 
     @Override
     public void close() {
-        close(0, null);
+        close(Duration.ofMillis(0));
     }
 
     @Override
-    public void close(long timeout, TimeUnit timeUnit) {
-        if (this.closed) {
-            throw new IllegalStateException("MockProducer is already closed.");
+    public void close(Duration timeout) {
+        if (this.closeException != null) {
+            throw this.closeException;
         }
+
         this.closed = true;
     }
 
@@ -313,7 +392,7 @@ public boolean closed() {
         return this.closed;
     }
 
-    public void fenceProducer() {
+    public synchronized void fenceProducer() {
         verifyProducerState();
         verifyTransactionsInitialized();
         this.producerFenced = true;
@@ -354,27 +433,32 @@ public synchronized List> history() {
         return new ArrayList<>(this.sent);
     }
 
+    public synchronized List> uncommittedRecords() {
+        return new ArrayList<>(this.uncommittedSends);
+    }
+
     /**
+     *
      * Get the list of committed consumer group offsets since the last call to {@link #clear()}
      */
     public synchronized List>> consumerGroupOffsetsHistory() {
         return new ArrayList<>(this.consumerGroupOffsets);
     }
+
+    public synchronized Map> uncommittedOffsets() {
+        return this.uncommittedConsumerGroupOffsets;
+    }
+
     /**
-     *
-     * Clear the stored history of sent records, consumer group offsets, and transactional state
+     * Clear the stored history of sent records, consumer group offsets
      */
     public synchronized void clear() {
         this.sent.clear();
         this.uncommittedSends.clear();
+        this.sentOffsets = false;
         this.completions.clear();
         this.consumerGroupOffsets.clear();
         this.uncommittedConsumerGroupOffsets.clear();
-        this.transactionInitialized = false;
-        this.transactionInFlight = false;
-        this.transactionCommitted = false;
-        this.transactionAborted = false;
-        this.producerFenced = false;
     }
 
     /**
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 49820333731be..15a62cd232cb5 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
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.clients.producer;
 
+import org.apache.kafka.clients.consumer.ConsumerGroupMetadata;
 import org.apache.kafka.common.Metric;
 import org.apache.kafka.common.MetricName;
 import org.apache.kafka.common.PartitionInfo;
@@ -24,6 +25,7 @@
 import org.apache.kafka.common.errors.ProducerFencedException;
 
 import java.io.Closeable;
+import java.time.Duration;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.Future;
@@ -52,6 +54,12 @@ public interface Producer extends Closeable {
     void sendOffsetsToTransaction(Map offsets,
                                   String consumerGroupId) throws ProducerFencedException;
 
+    /**
+     * See {@link KafkaProducer#sendOffsetsToTransaction(Map, ConsumerGroupMetadata)}
+     */
+    void sendOffsetsToTransaction(Map offsets,
+                                  ConsumerGroupMetadata groupMetadata) throws ProducerFencedException;
+
     /**
      * See {@link KafkaProducer#commitTransaction()}
      */
@@ -92,9 +100,13 @@ void sendOffsetsToTransaction(Map offsets,
      */
     void close();
 
+    @Deprecated
+    default void close(long timeout, TimeUnit unit) {
+        close(Duration.ofMillis(unit.toMillis(timeout)));
+    }
+
     /**
-     * See {@link KafkaProducer#close(long, TimeUnit)}
+     * See {@link KafkaProducer#close(Duration)}
      */
-    void close(long timeout, TimeUnit unit);
-
+    void close(Duration timeout);
 }
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 2dcf38caa3e47..2be8ced1f766f 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
@@ -16,19 +16,24 @@
  */
 package org.apache.kafka.clients.producer;
 
+import org.apache.kafka.clients.ClientDnsLookup;
 import org.apache.kafka.clients.CommonClientConfigs;
 import org.apache.kafka.clients.producer.internals.DefaultPartitioner;
 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.ConfigException;
+import org.apache.kafka.common.config.SecurityConfig;
 import org.apache.kafka.common.metrics.Sensor;
 import org.apache.kafka.common.serialization.Serializer;
 
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Properties;
 import java.util.Set;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import static org.apache.kafka.common.config.ConfigDef.Range.atLeast;
 import static org.apache.kafka.common.config.ConfigDef.Range.between;
@@ -50,10 +55,20 @@ 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;
+
     /** 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;
 
+    /** metadata.max.idle.ms */
+    public static final String METADATA_MAX_IDLE_CONFIG = "metadata.max.idle.ms";
+    private static final String METADATA_MAX_IDLE_DOC =
+            "Controls how long the producer will cache metadata for a topic that's idle. If the elapsed " +
+            "time since a topic was last produced to exceeds the metadata idle duration, then the topic's " +
+            "metadata is forgotten and the next access to it will force a metadata fetch request.";
+
     /** batch.size */
     public static final String BATCH_SIZE_CONFIG = "batch.size";
     private static final String BATCH_SIZE_DOC = "The producer will attempt to batch records together into fewer requests whenever multiple records are being sent"
@@ -77,13 +92,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."
                                            + " 
  3. 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."
                                            + " 
  4. 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";
@@ -98,6 +114,23 @@ public class ProducerConfig extends AbstractConfig {
                                                 + "specified time waiting for more records to show up. This setting defaults to 0 (i.e. no delay). Setting " + LINGER_MS_CONFIG + "=5, "
                                                 + "for example, would have the effect of reducing the number of requests sent but would add up to 5ms of latency to records sent in the absence of load.";
 
+    /** request.timeout.ms */
+    public static final String REQUEST_TIMEOUT_MS_CONFIG = CommonClientConfigs.REQUEST_TIMEOUT_MS_CONFIG;
+    private static final String REQUEST_TIMEOUT_MS_DOC = CommonClientConfigs.REQUEST_TIMEOUT_MS_DOC
+        + " This should be larger than replica.lag.time.max.ms (a broker configuration)"
+        + " to reduce the possibility of message duplication due to unnecessary producer retries.";
+
+    /** delivery.timeout.ms */
+    public static final String DELIVERY_TIMEOUT_MS_CONFIG = "delivery.timeout.ms";
+    private static final String DELIVERY_TIMEOUT_MS_DOC = "An upper bound on the time to report success or failure "
+            + "after a call to send() returns. This limits the total time that a record will be delayed "
+            + "prior to sending, the time to await acknowledgement from the broker (if expected), and the time allowed "
+            + "for retriable send failures. The producer may report failure to send a record earlier than this config if "
+            + "either an unrecoverable error is encountered, the retries have been exhausted, "
+            + "or the record is added to a batch which reached an earlier delivery expiration deadline. "
+            + "The value of this config should be greater than or equal to the sum of " + REQUEST_TIMEOUT_MS_CONFIG + " "
+            + "and " + LINGER_MS_CONFIG + ".";
+
     /** client.id */
     public static final String CLIENT_ID_CONFIG = CommonClientConfigs.CLIENT_ID_CONFIG;
 
@@ -109,10 +142,11 @@ public class ProducerConfig extends AbstractConfig {
 
     /** max.request.size */
     public static final String MAX_REQUEST_SIZE_CONFIG = "max.request.size";
-    private static final String MAX_REQUEST_SIZE_DOC = "The maximum size of a request in bytes. This setting will limit the number of record "
-                                                       + "batches the producer will send in a single request to avoid sending huge requests. "
-                                                       + "This is also effectively a cap on the maximum record batch size. Note that the server "
-                                                       + "has its own cap on record batch size which may be different from this.";
+    private static final String MAX_REQUEST_SIZE_DOC =
+        "The maximum size of a request in bytes. This setting will limit the number of record " +
+        "batches the producer will send in a single request to avoid sending huge requests. " +
+        "This is also effectively a cap on the maximum uncompressed record batch size. Note that the server " +
+        "has its own cap on the record batch size (after compression if compression is enabled) which may be different from this.";
 
     /** reconnect.backoff.ms */
     public static final String RECONNECT_BACKOFF_MS_CONFIG = CommonClientConfigs.RECONNECT_BACKOFF_MS_CONFIG;
@@ -122,9 +156,14 @@ public class ProducerConfig extends AbstractConfig {
 
     /** max.block.ms */
     public static final String MAX_BLOCK_MS_CONFIG = "max.block.ms";
-    private static final String MAX_BLOCK_MS_DOC = "The configuration controls how long KafkaProducer.send() and KafkaProducer.partitionsFor() will block."
-                                                    + "These methods can be blocked either because the buffer is full or metadata unavailable."
-                                                    + "Blocking in the user-supplied serializers or partitioner will not be counted against this timeout.";
+    private static final String MAX_BLOCK_MS_DOC = "The configuration controls how long the KafkaProducer's send(), partitionsFor(), "
+                                                    + "initTransactions(), sendOffsetsToTransaction(), commitTransaction() "
+                                                    + "and abortTransaction() methods will block. "
+                                                    + "For send() this timeout bounds the total time waiting for both metadata fetch and buffer allocation "
+                                                    + "(blocking in the user-supplied serializers or partitioner is not counted against this timeout). "
+                                                    + "For partitionsFor() this timeout bounds the time spent waiting for metadata if it is unavailable. "
+                                                    + "The transaction-related methods always block, but may timeout if "
+                                                    + "the transaction coordinator could not be discovered or did not respond within the timeout.";
 
     /** buffer.memory */
     public static final String BUFFER_MEMORY_CONFIG = "buffer.memory";
@@ -141,7 +180,7 @@ public class ProducerConfig extends AbstractConfig {
     /** compression.type */
     public static final String COMPRESSION_TYPE_CONFIG = "compression.type";
     private static final String COMPRESSION_TYPE_DOC = "The compression type for all data generated by the producer. The default is none (i.e. no compression). Valid "
-                                                       + " values are none, gzip, snappy, or lz4. "
+                                                       + " values are none, gzip, snappy, lz4, or zstd. "
                                                        + "Compression is of full batches of data, so the efficacy of batching will also impact the compression ratio (more batching means better compression).";
 
     /** metrics.sample.window.ms */
@@ -151,7 +190,7 @@ public class ProducerConfig extends AbstractConfig {
     public static final String METRICS_NUM_SAMPLES_CONFIG = CommonClientConfigs.METRICS_NUM_SAMPLES_CONFIG;
 
     /**
-     * metrics.log.level
+     * metrics.recording.level
      */
     public static final String METRICS_RECORDING_LEVEL_CONFIG = CommonClientConfigs.METRICS_RECORDING_LEVEL_CONFIG;
 
@@ -167,10 +206,14 @@ public class ProducerConfig extends AbstractConfig {
     /** retries */
     public static final String RETRIES_CONFIG = CommonClientConfigs.RETRIES_CONFIG;
     private static final String RETRIES_DOC = "Setting a value greater than zero will cause the client to resend any record whose send fails with a potentially transient error."
-                                              + " Note that this retry is no different than if the client resent the record upon receiving the error."
-                                              + " Allowing retries without setting " + MAX_IN_FLIGHT_REQUESTS_PER_CONNECTION + " to 1 will potentially change the"
-                                              + " ordering of records because if two batches are sent to a single partition, and the first fails and is retried but the second"
-                                              + " succeeds, then the records in the second batch may appear first.";
+            + " Note that this retry is no different than if the client resent the record upon receiving the error."
+            + " Allowing retries without setting " + MAX_IN_FLIGHT_REQUESTS_PER_CONNECTION + " to 1 will potentially change the"
+            + " ordering of records because if two batches are sent to a single partition, and the first fails and is retried but the second"
+            + " succeeds, then the records in the second batch may appear first. Note additionally that produce requests will be"
+            + " failed before the number of retries has been exhausted if the timeout configured by"
+            + " " + DELIVERY_TIMEOUT_MS_CONFIG + " expires first before successful acknowledgement. Users should generally"
+            + " prefer to leave this config unset and instead use " + DELIVERY_TIMEOUT_MS_CONFIG + " to control"
+            + " retry behavior.";
 
     /** key.serializer */
     public static final String KEY_SERIALIZER_CLASS_CONFIG = "key.serializer";
@@ -180,6 +223,12 @@ public class ProducerConfig extends AbstractConfig {
     public static final String VALUE_SERIALIZER_CLASS_CONFIG = "value.serializer";
     public static final String VALUE_SERIALIZER_CLASS_DOC = "Serializer class for value that implements the org.apache.kafka.common.serialization.Serializer interface.";
 
+    /** socket.connection.setup.timeout.ms */
+    public static final String SOCKET_CONNECTION_SETUP_TIMEOUT_MS_CONFIG = CommonClientConfigs.SOCKET_CONNECTION_SETUP_TIMEOUT_MS_CONFIG;
+
+    /** socket.connection.setup.timeout.max.ms */
+    public static final String SOCKET_CONNECTION_SETUP_TIMEOUT_MAX_MS_CONFIG = CommonClientConfigs.SOCKET_CONNECTION_SETUP_TIMEOUT_MAX_MS_CONFIG;
+
     /** connections.max.idle.ms */
     public static final String CONNECTIONS_MAX_IDLE_MS_CONFIG = CommonClientConfigs.CONNECTIONS_MAX_IDLE_MS_CONFIG;
 
@@ -187,12 +236,6 @@ public class ProducerConfig extends AbstractConfig {
     public static final String PARTITIONER_CLASS_CONFIG = "partitioner.class";
     private static final String PARTITIONER_CLASS_DOC = "Partitioner class that implements the org.apache.kafka.clients.producer.Partitioner interface.";
 
-    /** request.timeout.ms */
-    public static final String REQUEST_TIMEOUT_MS_CONFIG = CommonClientConfigs.REQUEST_TIMEOUT_MS_CONFIG;
-    private static final String REQUEST_TIMEOUT_MS_DOC = CommonClientConfigs.REQUEST_TIMEOUT_MS_DOC
-                                                        + " This should be larger than replica.lag.time.max.ms (a broker configuration)"
-                                                        + " to reduce the possibility of message duplication due to unnecessary producer retries.";
-
     /** interceptor.classes */
     public static final String INTERCEPTOR_CLASSES_CONFIG = "interceptor.classes";
     public static final String INTERCEPTOR_CLASSES_DOC = "A list of classes to use as interceptors. "
@@ -204,25 +247,57 @@ public class ProducerConfig extends AbstractConfig {
     public static final String ENABLE_IDEMPOTENCE_DOC = "When set to 'true', the producer will ensure that exactly one copy of each message is written in the stream. If 'false', producer "
                                                         + "retries due to broker failures, etc., may write duplicates of the retried message in the stream. "
                                                         + "Note that enabling idempotence requires " + MAX_IN_FLIGHT_REQUESTS_PER_CONNECTION + " to be less than or equal to 5, "
-                                                        + "" + RETRIES_CONFIG + " to be greater than 0 and " + ACKS_CONFIG + " must be 'all'. If these values "
+                                                        + "" + RETRIES_CONFIG + " to be greater than 0 and " + ACKS_CONFIG + " must be 'all'. If these values "
                                                         + "are not explicitly set by the user, suitable values will be chosen. If incompatible values are set, "
-                                                        + "a ConfigException will be thrown.";
+                                                        + "a ConfigException will be thrown.";
 
     /**  transaction.timeout.ms  */
     public static final String TRANSACTION_TIMEOUT_CONFIG = "transaction.timeout.ms";
     public static final String TRANSACTION_TIMEOUT_DOC = "The maximum amount of time in ms that the transaction coordinator will wait for a transaction status update from the producer before proactively aborting the ongoing transaction." +
-            "If this value is larger than the max.transaction.timeout.ms setting in the broker, the request will fail with a `InvalidTransactionTimeout` error.";
+            "If this value is larger than the transaction.max.timeout.ms setting in the broker, the request will fail with a InvalidTxnTimeoutException error.";
 
     /**  transactional.id  */
     public static final String TRANSACTIONAL_ID_CONFIG = "transactional.id";
     public static final String TRANSACTIONAL_ID_DOC = "The TransactionalId to use for transactional delivery. This enables reliability semantics which span multiple producer sessions since it allows the client to guarantee that transactions using the same TransactionalId have been completed prior to starting any new transactions. If no TransactionalId is provided, then the producer is limited to idempotent delivery. " +
-            "Note that enable.idempotence must be enabled if a TransactionalId is configured. " +
-            "The default is empty, which means transactions cannot be used.";
+            "If a TransactionalId is configured, enable.idempotence is implied. " +
+            "By default the TransactionId is not configured, 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;
+
+    /**
+     * internal.auto.downgrade.txn.commit
+     * Whether or not the producer should automatically downgrade the transactional commit request when the new group metadata
+     * feature is not supported by the broker.
+     * 
    
+     * The purpose of this flag is to make Kafka Streams being capable of working with old brokers when applying this new API.
+     * Non Kafka Streams users who are building their own EOS applications should be careful playing around
+     * with config as there is a risk of violating EOS semantics when turning on this flag.
+     *
+     * 
    
+     * Note: this is an internal configuration and could be changed in the future in a backward incompatible way
+     *
+     */
+    static final String AUTO_DOWNGRADE_TXN_COMMIT = "internal.auto.downgrade.txn.commit";
+
+    private static final AtomicInteger PRODUCER_CLIENT_ID_SEQUENCE = new AtomicInteger(1);
 
     static {
-        CONFIG = new ConfigDef().define(BOOTSTRAP_SERVERS_CONFIG, Type.LIST, Importance.HIGH, CommonClientConfigs.BOOTSTRAP_SERVERS_DOC)
+        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,
+                                        Type.STRING,
+                                        ClientDnsLookup.USE_ALL_DNS_IPS.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(BUFFER_MEMORY_CONFIG, Type.LONG, 32 * 1024 * 1024L, atLeast(0L), Importance.HIGH, BUFFER_MEMORY_DOC)
-                                .define(RETRIES_CONFIG, Type.INT, 0, between(0, Integer.MAX_VALUE), Importance.HIGH, RETRIES_DOC)
+                                .define(RETRIES_CONFIG, Type.INT, Integer.MAX_VALUE, between(0, Integer.MAX_VALUE), Importance.HIGH, RETRIES_DOC)
                                 .define(ACKS_CONFIG,
                                         Type.STRING,
                                         "1",
@@ -231,13 +306,14 @@ 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.LONG, 0, atLeast(0L), Importance.MEDIUM, LINGER_MS_DOC)
+                                .define(LINGER_MS_CONFIG, Type.LONG, 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(CLIENT_ID_CONFIG, Type.STRING, "", Importance.MEDIUM, CommonClientConfigs.CLIENT_ID_DOC)
-                                .define(SEND_BUFFER_CONFIG, Type.INT, 128 * 1024, atLeast(-1), Importance.MEDIUM, CommonClientConfigs.SEND_BUFFER_DOC)
-                                .define(RECEIVE_BUFFER_CONFIG, Type.INT, 32 * 1024, atLeast(-1), Importance.MEDIUM, CommonClientConfigs.RECEIVE_BUFFER_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)
                                 .define(MAX_REQUEST_SIZE_CONFIG,
                                         Type.INT,
-                                        1 * 1024 * 1024,
+                                        1024 * 1024,
                                         atLeast(0),
                                         Importance.MEDIUM,
                                         MAX_REQUEST_SIZE_DOC)
@@ -257,6 +333,12 @@ public class ProducerConfig extends AbstractConfig {
                                         Importance.MEDIUM,
                                         REQUEST_TIMEOUT_MS_DOC)
                                 .define(METADATA_MAX_AGE_CONFIG, Type.LONG, 5 * 60 * 1000, atLeast(0), Importance.LOW, METADATA_MAX_AGE_DOC)
+                                .define(METADATA_MAX_IDLE_CONFIG,
+                                        Type.LONG,
+                                        5 * 60 * 1000,
+                                        atLeast(5000),
+                                        Importance.LOW,
+                                        METADATA_MAX_IDLE_DOC)
                                 .define(METRICS_SAMPLE_WINDOW_MS_CONFIG,
                                         Type.LONG,
                                         30000,
@@ -267,12 +349,13 @@ public class ProducerConfig extends AbstractConfig {
                                 .define(METRICS_RECORDING_LEVEL_CONFIG,
                                         Type.STRING,
                                         Sensor.RecordingLevel.INFO.toString(),
-                                        in(Sensor.RecordingLevel.INFO.toString(), Sensor.RecordingLevel.DEBUG.toString()),
+                                        in(Sensor.RecordingLevel.INFO.toString(), Sensor.RecordingLevel.DEBUG.toString(), Sensor.RecordingLevel.TRACE.toString()),
                                         Importance.LOW,
                                         CommonClientConfigs.METRICS_RECORDING_LEVEL_DOC)
                                 .define(METRIC_REPORTER_CLASSES_CONFIG,
                                         Type.LIST,
-                                        "",
+                                        Collections.emptyList(),
+                                        new ConfigDef.NonNullValidator(),
                                         Importance.LOW,
                                         CommonClientConfigs.METRIC_REPORTER_CLASSES_DOC)
                                 .define(MAX_IN_FLIGHT_REQUESTS_PER_CONNECTION,
@@ -289,6 +372,16 @@ public class ProducerConfig extends AbstractConfig {
                                         Type.CLASS,
                                         Importance.HIGH,
                                         VALUE_SERIALIZER_CLASS_DOC)
+                                .define(SOCKET_CONNECTION_SETUP_TIMEOUT_MS_CONFIG,
+                                        Type.LONG,
+                                        CommonClientConfigs.DEFAULT_SOCKET_CONNECTION_SETUP_TIMEOUT_MS,
+                                        Importance.MEDIUM,
+                                        CommonClientConfigs.SOCKET_CONNECTION_SETUP_TIMEOUT_MS_DOC)
+                                .define(SOCKET_CONNECTION_SETUP_TIMEOUT_MAX_MS_CONFIG,
+                                        Type.LONG,
+                                        CommonClientConfigs.DEFAULT_SOCKET_CONNECTION_SETUP_TIMEOUT_MAX_MS,
+                                        Importance.MEDIUM,
+                                        CommonClientConfigs.SOCKET_CONNECTION_SETUP_TIMEOUT_MAX_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(CONNECTIONS_MAX_IDLE_MS_CONFIG,
                                         Type.LONG,
@@ -301,7 +394,8 @@ public class ProducerConfig extends AbstractConfig {
                                         Importance.MEDIUM, PARTITIONER_CLASS_DOC)
                                 .define(INTERCEPTOR_CLASSES_CONFIG,
                                         Type.LIST,
-                                        null,
+                                        Collections.emptyList(),
+                                        new ConfigDef.NonNullValidator(),
                                         Importance.LOW,
                                         INTERCEPTOR_CLASSES_DOC)
                                 .define(CommonClientConfigs.SECURITY_PROTOCOL_CONFIG,
@@ -309,6 +403,11 @@ 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)
                                 .withClientSslSupport()
                                 .withClientSaslSupport()
                                 .define(ENABLE_IDEMPOTENCE_CONFIG,
@@ -326,18 +425,86 @@ public class ProducerConfig extends AbstractConfig {
                                         null,
                                         new ConfigDef.NonEmptyString(),
                                         Importance.LOW,
-                                        TRANSACTIONAL_ID_DOC);
+                                        TRANSACTIONAL_ID_DOC)
+                                .defineInternal(AUTO_DOWNGRADE_TXN_COMMIT,
+                                        Type.BOOLEAN,
+                                        false,
+                                        Importance.LOW);
     }
 
     @Override
     protected Map postProcessParsedConfig(final Map parsedValues) {
-        return CommonClientConfigs.postProcessReconnectBackoffConfigs(this, parsedValues);
+        CommonClientConfigs.warnIfDeprecatedDnsLookupValue(this);
+        Map refinedConfigs = CommonClientConfigs.postProcessReconnectBackoffConfigs(this, parsedValues);
+        maybeOverrideEnableIdempotence(refinedConfigs);
+        maybeOverrideClientId(refinedConfigs);
+        maybeOverrideAcksAndRetries(refinedConfigs);
+        return refinedConfigs;
+    }
+
+    private void maybeOverrideClientId(final Map configs) {
+        String refinedClientId;
+        boolean userConfiguredClientId = this.originals().containsKey(CLIENT_ID_CONFIG);
+        if (userConfiguredClientId) {
+            refinedClientId = this.getString(CLIENT_ID_CONFIG);
+        } else {
+            String transactionalId = this.getString(TRANSACTIONAL_ID_CONFIG);
+            refinedClientId = "producer-" + (transactionalId != null ? transactionalId : PRODUCER_CLIENT_ID_SEQUENCE.getAndIncrement());
+        }
+        configs.put(CLIENT_ID_CONFIG, refinedClientId);
+    }
+
+    private void maybeOverrideEnableIdempotence(final Map configs) {
+        boolean userConfiguredIdempotence = this.originals().containsKey(ENABLE_IDEMPOTENCE_CONFIG);
+        boolean userConfiguredTransactions = this.originals().containsKey(TRANSACTIONAL_ID_CONFIG);
+
+        if (userConfiguredTransactions && !userConfiguredIdempotence) {
+            configs.put(ENABLE_IDEMPOTENCE_CONFIG, true);
+        }
+    }
+
+    private void maybeOverrideAcksAndRetries(final Map configs) {
+        final String acksStr = parseAcks(this.getString(ACKS_CONFIG));
+        configs.put(ACKS_CONFIG, acksStr);
+        // For idempotence producers, values for `RETRIES_CONFIG` and `ACKS_CONFIG` might need to be overridden.
+        if (idempotenceEnabled()) {
+            boolean userConfiguredRetries = this.originals().containsKey(RETRIES_CONFIG);
+            if (this.getInt(RETRIES_CONFIG) == 0) {
+                throw new ConfigException("Must set " + ProducerConfig.RETRIES_CONFIG + " to non-zero when using the idempotent producer.");
+            }
+            configs.put(RETRIES_CONFIG, userConfiguredRetries ? this.getInt(RETRIES_CONFIG) : Integer.MAX_VALUE);
+
+            boolean userConfiguredAcks = this.originals().containsKey(ACKS_CONFIG);
+            final short acks = Short.valueOf(acksStr);
+            if (userConfiguredAcks && acks != (short) -1) {
+                throw new ConfigException("Must set " + ACKS_CONFIG + " to all in order to use the idempotent " +
+                        "producer. Otherwise we cannot guarantee idempotence.");
+            }
+            configs.put(ACKS_CONFIG, "-1");
+        }
+    }
+
+    private static String parseAcks(String acksString) {
+        try {
+            return acksString.trim().equalsIgnoreCase("all") ? "-1" : Short.parseShort(acksString.trim()) + "";
+        } catch (NumberFormatException e) {
+            throw new ConfigException("Invalid configuration value for 'acks': " + acksString);
+        }
     }
 
+    /**
+     * @deprecated Since 2.7.0. This will be removed in a future major release.
+     */
+    @Deprecated
     public static Map addSerializerToConfig(Map configs,
                                                             Serializer keySerializer, Serializer valueSerializer) {
-        Map newConfigs = new HashMap<>();
-        newConfigs.putAll(configs);
+        return appendSerializerToConfig(configs, keySerializer, valueSerializer);
+    }
+
+    static Map appendSerializerToConfig(Map configs,
+            Serializer keySerializer,
+            Serializer valueSerializer) {
+        Map newConfigs = new HashMap<>(configs);
         if (keySerializer != null)
             newConfigs.put(KEY_SERIALIZER_CLASS_CONFIG, keySerializer.getClass());
         if (valueSerializer != null)
@@ -345,8 +512,13 @@ public static Map addSerializerToConfig(Map conf
         return newConfigs;
     }
 
+    /**
+     * @deprecated Since 2.7.0. This will be removed in a future major release.
+     */
+    @Deprecated
     public static Properties addSerializerToConfig(Properties properties,
-                                                   Serializer keySerializer, Serializer valueSerializer) {
+                                                   Serializer keySerializer,
+                                                   Serializer valueSerializer) {
         Properties newProperties = new Properties();
         newProperties.putAll(properties);
         if (keySerializer != null)
@@ -356,10 +528,24 @@ public static Properties addSerializerToConfig(Properties properties,
         return newProperties;
     }
 
-    ProducerConfig(Map props) {
+    public ProducerConfig(Properties props) {
+        super(CONFIG, props);
+    }
+
+    public ProducerConfig(Map props) {
         super(CONFIG, props);
     }
 
+    boolean idempotenceEnabled() {
+        boolean userConfiguredIdempotence = this.originals().containsKey(ENABLE_IDEMPOTENCE_CONFIG);
+        boolean userConfiguredTransactions = this.originals().containsKey(TRANSACTIONAL_ID_CONFIG);
+        boolean idempotenceEnabled = userConfiguredIdempotence && this.getBoolean(ENABLE_IDEMPOTENCE_CONFIG);
+
+        if (!idempotenceEnabled && userConfiguredIdempotence && userConfiguredTransactions)
+            throw new ConfigException("Cannot set a " + ProducerConfig.TRANSACTIONAL_ID_CONFIG + " without also enabling idempotence.");
+        return userConfiguredTransactions || idempotenceEnabled;
+    }
+
     ProducerConfig(Map props, boolean doLog) {
         super(CONFIG, props, doLog);
     }
@@ -368,8 +554,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(4, config -> "producerconfigs_" + config));
     }
 
 }
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 c8ff00bf4f135..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.
@@ -104,7 +106,7 @@ public ProducerRecord(String topic, Integer partition, Long timestamp, K key, V
      * @param value The record contents
      * @param headers The headers that will be included in the record
      */
-    public ProducerRecord(String topic, Integer partition, K key, V value,  Iterable
     headers) {
+    public ProducerRecord(String topic, Integer partition, K key, V value, Iterable
     headers) {
         this(topic, partition, null, key, value, headers);
     }
     
@@ -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/RecordMetadata.java b/clients/src/main/java/org/apache/kafka/clients/producer/RecordMetadata.java
index 0924244829b6f..1ae5cf654b8de 100644
--- a/clients/src/main/java/org/apache/kafka/clients/producer/RecordMetadata.java
+++ b/clients/src/main/java/org/apache/kafka/clients/producer/RecordMetadata.java
@@ -56,16 +56,6 @@ public RecordMetadata(TopicPartition topicPartition, long baseOffset, long relat
         this.topicPartition = topicPartition;
     }
 
-    /**
-     * @deprecated As of 0.11.0. Use @{@link RecordMetadata#RecordMetadata(TopicPartition, long, long, long, Long, int, int)}.
-     */
-    @Deprecated
-    public RecordMetadata(TopicPartition topicPartition, long baseOffset, long relativeOffset, long timestamp,
-                          long checksum, int serializedKeySize, int serializedValueSize) {
-        this(topicPartition, baseOffset, relativeOffset, timestamp, Long.valueOf(checksum), serializedKeySize,
-                serializedValueSize);
-    }
-
     /**
      * Indicates whether the record metadata includes the offset.
      * @return true if the offset is included in the metadata, false otherwise.
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 c5df2dacbca1e..ee84c7c168a4a 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
@@ -23,8 +23,9 @@
 import java.util.concurrent.locks.Condition;
 import java.util.concurrent.locks.ReentrantLock;
 
+import org.apache.kafka.clients.producer.BufferExhaustedException;
+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;
@@ -55,6 +56,7 @@ public class BufferPool {
     private final Metrics metrics;
     private final Time time;
     private final Sensor waitTime;
+    private boolean closed;
 
     /**
      * Create a new buffer pool
@@ -81,7 +83,14 @@ public BufferPool(long memory, int poolableSize, Metrics metrics, Time time, Str
         MetricName totalMetricName = metrics.metricName("bufferpool-wait-time-total",
                                                    metricGrpName,
                                                    "The total time an appender waits for space allocation.");
+
+        Sensor bufferExhaustedRecordSensor = metrics.sensor("buffer-exhausted-records");
+        MetricName bufferExhaustedRateMetricName = metrics.metricName("buffer-exhausted-rate", metricGrpName, "The average per-second number of record sends that are dropped due to buffer exhaustion");
+        MetricName bufferExhaustedTotalMetricName = metrics.metricName("buffer-exhausted-total", metricGrpName, "The total number of record sends that are dropped due to buffer exhaustion");
+        bufferExhaustedRecordSensor.add(new Meter(bufferExhaustedRateMetricName, bufferExhaustedTotalMetricName));
+
         this.waitTime.add(new Meter(TimeUnit.NANOSECONDS, rateMetricName, totalMetricName));
+        this.closed = false;
     }
 
     /**
@@ -104,6 +113,12 @@ 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())
@@ -135,11 +150,15 @@ public ByteBuffer allocate(int size, long maxTimeToBlockMs) throws InterruptedEx
                         } finally {
                             long endWaitNs = time.nanoseconds();
                             timeNs = Math.max(0L, endWaitNs - startWaitNs);
-                            this.waitTime.record(timeNs, time.milliseconds());
+                            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.");
+                            this.metrics.sensor("buffer-exhausted-records").record();
+                            throw new BufferExhaustedException("Failed to allocate memory within the configured max blocking time " + maxTimeToBlockMs + " ms.");
                         }
 
                         remainingTimeToBlockNs -= timeNs;
@@ -185,6 +204,11 @@ public ByteBuffer allocate(int size, long maxTimeToBlockMs) throws InterruptedEx
             return buffer;
     }
 
+    // Protected for testing
+    protected void recordWaitTime(long timeNs) {
+        this.waitTime.record(timeNs, time.milliseconds());
+    }
+
     /**
      * Allocate a buffer.  If buffer allocation fails (e.g. because of OOM) then return the size count back to
      * available memory and signal the next waiter if it exists.
@@ -311,4 +335,19 @@ 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();
+        }
+    }
 }
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..cf765d1eee6aa 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
@@ -16,28 +16,24 @@
  */
 package org.apache.kafka.clients.producer.internals;
 
-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;
-import org.apache.kafka.common.PartitionInfo;
 import org.apache.kafka.common.utils.Utils;
 
+import java.util.Map;
+
 /**
  * The default partitioning strategy:
  * 
        
  * 
      • 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 +48,36 @@ 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) {
-        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;
-        }
+        return partition(topic, key, keyBytes, value, valueBytes, cluster, cluster.partitionsForTopic(topic).size());
     }
 
-    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;
-            }
+    /**
+     * Compute the partition for the given record.
+     *
+     * @param topic The topic name
+     * @param numPartitions The number of partitions of the given {@code topic}
+     * @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,
+                         int numPartitions) {
+        if (keyBytes == null) {
+            return stickyPartitionCache.partition(topic, cluster);
         }
-        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/FutureRecordMetadata.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/FutureRecordMetadata.java
index 8fcc46ff3ff60..d1a643b319679 100644
--- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/FutureRecordMetadata.java
+++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/FutureRecordMetadata.java
@@ -16,13 +16,14 @@
  */
 package org.apache.kafka.clients.producer.internals;
 
+import org.apache.kafka.clients.producer.RecordMetadata;
+import org.apache.kafka.common.utils.Time;
+
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 
-import org.apache.kafka.clients.producer.RecordMetadata;
-
 /**
  * The future result of a record send
  */
@@ -34,16 +35,18 @@ public final class FutureRecordMetadata implements Future {
     private final Long checksum;
     private final int serializedKeySize;
     private final int serializedValueSize;
+    private final Time time;
     private volatile FutureRecordMetadata nextRecordMetadata = null;
 
     public FutureRecordMetadata(ProduceRequestResult result, long relativeOffset, long createTimestamp,
-                                Long checksum, int serializedKeySize, int serializedValueSize) {
+                                Long checksum, int serializedKeySize, int serializedValueSize, Time time) {
         this.result = result;
         this.relativeOffset = relativeOffset;
         this.createTimestamp = createTimestamp;
         this.checksum = checksum;
         this.serializedKeySize = serializedKeySize;
         this.serializedValueSize = serializedValueSize;
+        this.time = time;
     }
 
     @Override
@@ -67,13 +70,14 @@ public RecordMetadata get() throws InterruptedException, ExecutionException {
     @Override
     public RecordMetadata get(long timeout, TimeUnit unit) throws InterruptedException, ExecutionException, TimeoutException {
         // Handle overflow.
-        long now = System.currentTimeMillis();
-        long deadline = Long.MAX_VALUE - timeout < now ? Long.MAX_VALUE : now + timeout;
+        long now = time.milliseconds();
+        long timeoutMillis = unit.toMillis(timeout);
+        long deadline = Long.MAX_VALUE - timeoutMillis < now ? Long.MAX_VALUE : now + timeoutMillis;
         boolean occurred = this.result.await(timeout, unit);
-        if (nextRecordMetadata != null)
-            return nextRecordMetadata.get(deadline - System.currentTimeMillis(), TimeUnit.MILLISECONDS);
         if (!occurred)
-            throw new TimeoutException("Timeout after waiting for " + TimeUnit.MILLISECONDS.convert(timeout, unit) + " ms.");
+            throw new TimeoutException("Timeout after waiting for " + timeoutMillis + " ms.");
+        if (nextRecordMetadata != null)
+            return nextRecordMetadata.get(deadline - time.milliseconds(), TimeUnit.MILLISECONDS);
         return valueOrError();
     }
 
diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProduceRequestResult.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProduceRequestResult.java
index fbfef6171ed19..1e8c787963284 100644
--- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProduceRequestResult.java
+++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProduceRequestResult.java
@@ -29,7 +29,7 @@
  * partition in a produce request and it is shared by all the {@link RecordMetadata} instances that are batched together
  * for the same partition in the request.
  */
-public final class ProduceRequestResult {
+public class ProduceRequestResult {
 
     private final CountDownLatch latch = new CountDownLatch(1);
     private final TopicPartition topicPartition;
diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerBatch.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerBatch.java
index ea0f0f7dff953..cfd3a6794cc04 100644
--- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerBatch.java
+++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerBatch.java
@@ -17,10 +17,10 @@
 package org.apache.kafka.clients.producer.internals;
 
 import org.apache.kafka.clients.producer.Callback;
+import org.apache.kafka.common.utils.ProducerIdAndEpoch;
 import org.apache.kafka.clients.producer.RecordMetadata;
 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.header.Header;
 import org.apache.kafka.common.record.AbstractRecords;
 import org.apache.kafka.common.record.CompressionRatioEstimator;
@@ -32,6 +32,7 @@
 import org.apache.kafka.common.record.RecordBatch;
 import org.apache.kafka.common.record.TimestampType;
 import org.apache.kafka.common.requests.ProduceResponse;
+import org.apache.kafka.common.utils.Time;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -73,17 +74,16 @@ private enum FinalState { ABORTED, FAILED, SUCCEEDED }
     private long lastAttemptMs;
     private long lastAppendTime;
     private long drainedMs;
-    private String expiryErrorMessage;
     private boolean retry;
-    private boolean reopened = false;
+    private boolean reopened;
 
-    public ProducerBatch(TopicPartition tp, MemoryRecordsBuilder recordsBuilder, long now) {
-        this(tp, recordsBuilder, now, false);
+    public ProducerBatch(TopicPartition tp, MemoryRecordsBuilder recordsBuilder, long createdMs) {
+        this(tp, recordsBuilder, createdMs, false);
     }
 
-    public ProducerBatch(TopicPartition tp, MemoryRecordsBuilder recordsBuilder, long now, boolean isSplitBatch) {
-        this.createdMs = now;
-        this.lastAttemptMs = now;
+    public ProducerBatch(TopicPartition tp, MemoryRecordsBuilder recordsBuilder, long createdMs, boolean isSplitBatch) {
+        this.createdMs = createdMs;
+        this.lastAttemptMs = createdMs;
         this.recordsBuilder = recordsBuilder;
         this.topicPartition = tp;
         this.lastAppendTime = createdMs;
@@ -111,7 +111,8 @@ public FutureRecordMetadata tryAppend(long timestamp, byte[] key, byte[] value,
             FutureRecordMetadata future = new FutureRecordMetadata(this.produceFuture, this.recordCount,
                                                                    timestamp, checksum,
                                                                    key == null ? -1 : key.length,
-                                                                   value == null ? -1 : value.length);
+                                                                   value == null ? -1 : value.length,
+                                                                   Time.SYSTEM);
             // we have to keep every future returned to the users in case the batch needs to be
             // split to several new batches and resent.
             thunks.add(new Thunk(callback, future));
@@ -135,7 +136,8 @@ private boolean tryAppendForSplit(long timestamp, ByteBuffer key, ByteBuffer val
             FutureRecordMetadata future = new FutureRecordMetadata(this.produceFuture, this.recordCount,
                                                                    timestamp, thunk.future.checksumOrNull(),
                                                                    key == null ? -1 : key.remaining(),
-                                                                   value == null ? -1 : value.remaining());
+                                                                   value == null ? -1 : value.remaining(),
+                                                                   Time.SYSTEM);
             // Chain the future to the original thunk.
             thunk.future.chain(future);
             this.thunks.add(thunk);
@@ -158,7 +160,24 @@ public void abort(RuntimeException exception) {
     }
 
     /**
-     * Complete the request. If the batch was previously aborted, this is a no-op.
+     * Return `true` if {@link #done(long, long, RuntimeException)} has been invoked at least once, `false` otherwise.
+     */
+    public boolean isDone() {
+        return finalState() != null;
+    }
+
+    /**
+     * Finalize the state of a batch. Final state, once set, is immutable. This function may be called
+     * once or twice on a batch. It may be called twice if
+     * 1. An inflight batch expires before a response from the broker is received. The batch's final
+     * state is set to FAILED. But it could succeed on the broker and second time around batch.done() may
+     * try to set SUCCEEDED final state.
+     * 2. If a transaction abortion happens or if the producer is closed forcefully, the final state is
+     * ABORTED but again it could succeed if broker responds with a success.
+     *
+     * Attempted transitions from [FAILED | ABORTED] --> SUCCEEDED are logged.
+     * Attempted transitions from one failure state to the same or a different failed state are ignored.
+     * Attempted transitions from SUCCEEDED to the same or a failed state throw an exception.
      *
      * @param baseOffset The base offset of the messages assigned by the server
      * @param logAppendTime The log append time or -1 if CreateTime is being used
@@ -166,26 +185,34 @@ public void abort(RuntimeException exception) {
      * @return true if the batch was completed successfully and false if the batch was previously aborted
      */
     public boolean done(long baseOffset, long logAppendTime, RuntimeException exception) {
-        final FinalState finalState;
-        if (exception == null) {
+        final FinalState tryFinalState = (exception == null) ? FinalState.SUCCEEDED : FinalState.FAILED;
+
+        if (tryFinalState == FinalState.SUCCEEDED) {
             log.trace("Successfully produced messages to {} with base offset {}.", topicPartition, baseOffset);
-            finalState = FinalState.SUCCEEDED;
         } else {
-            log.trace("Failed to produce messages to {}.", topicPartition, exception);
-            finalState = FinalState.FAILED;
+            log.trace("Failed to produce messages to {} with base offset {}.", topicPartition, baseOffset, exception);
         }
 
-        if (!this.finalState.compareAndSet(null, finalState)) {
-            if (this.finalState.get() == FinalState.ABORTED) {
-                log.debug("ProduceResponse returned for {} after batch had already been aborted.", topicPartition);
-                return false;
+        if (this.finalState.compareAndSet(null, tryFinalState)) {
+            completeFutureAndFireCallbacks(baseOffset, logAppendTime, exception);
+            return true;
+        }
+
+        if (this.finalState.get() != FinalState.SUCCEEDED) {
+            if (tryFinalState == FinalState.SUCCEEDED) {
+                // Log if a previously unsuccessful batch succeeded later on.
+                log.debug("ProduceResponse returned {} for {} after batch with base offset {} had already been {}.",
+                    tryFinalState, topicPartition, baseOffset, this.finalState.get());
             } else {
-                throw new IllegalStateException("Batch has already been completed in final state " + this.finalState.get());
+                // FAILED --> FAILED and ABORTED --> FAILED transitions are ignored.
+                log.debug("Ignored state transition {} -> {} for {} batch with base offset {}",
+                    this.finalState.get(), tryFinalState, topicPartition, baseOffset);
             }
+        } else {
+            // A SUCCESSFUL batch must not attempt another state change.
+            throw new IllegalStateException("A " + this.finalState.get() + " batch must not attempt another state change to " + tryFinalState);
         }
-
-        completeFutureAndFireCallbacks(baseOffset, logAppendTime, exception);
-        return true;
+        return false;
     }
 
     private void completeFutureAndFireCallbacks(long baseOffset, long logAppendTime, RuntimeException exception) {
@@ -241,14 +268,17 @@ public Deque split(int splitBatchSize) {
             // A newly created batch can always host the first message.
             if (!batch.tryAppendForSplit(record.timestamp(), record.key(), record.value(), record.headers(), thunk)) {
                 batches.add(batch);
+                batch.closeForRecordAppends();
                 batch = createBatchOffAccumulatorForRecord(record, splitBatchSize);
                 batch.tryAppendForSplit(record.timestamp(), record.key(), record.value(), record.headers(), thunk);
             }
         }
 
         // Close the last batch and add it to the batch list after split.
-        if (batch != null)
+        if (batch != null) {
             batches.add(batch);
+            batch.closeForRecordAppends();
+        }
 
         produceFuture.set(ProduceResponse.INVALID_OFFSET, NO_TIMESTAMP, new RecordBatchTooLargeException());
         produceFuture.done();
@@ -299,37 +329,12 @@ public String toString() {
         return "ProducerBatch(topicPartition=" + topicPartition + ", recordCount=" + recordCount + ")";
     }
 
-    /**
-     * A batch whose metadata is not available should be expired if one of the following is true:
-     * 
          
-     *     
        1.  the batch is not in retry AND request timeout has elapsed after it is ready (full or linger.ms has reached).
-     *     
        2.  the batch is in retry AND request timeout has elapsed after the backoff period ended.
-     * 
        
-     * This methods closes this batch and sets {@code expiryErrorMessage} if the batch has timed out.
-     */
-    boolean maybeExpire(int requestTimeoutMs, long retryBackoffMs, long now, long lingerMs, boolean isFull) {
-        if (!this.inRetry() && isFull && requestTimeoutMs < (now - this.lastAppendTime))
-            expiryErrorMessage = (now - this.lastAppendTime) + " ms has passed since last append";
-        else if (!this.inRetry() && requestTimeoutMs < (createdTimeMs(now) - lingerMs))
-            expiryErrorMessage = (createdTimeMs(now) - lingerMs) + " ms has passed since batch creation plus linger time";
-        else if (this.inRetry() && requestTimeoutMs < (waitedTimeMs(now) - retryBackoffMs))
-            expiryErrorMessage = (waitedTimeMs(now) - retryBackoffMs) + " ms has passed since last attempt plus backoff time";
-
-        boolean expired = expiryErrorMessage != null;
-        if (expired)
-            abortRecordAppends();
-        return expired;
+    boolean hasReachedDeliveryTimeout(long deliveryTimeoutMs, long now) {
+        return deliveryTimeoutMs <= now - this.createdMs;
     }
 
-    /**
-     * If {@link #maybeExpire(int, long, long, long, boolean)} returned true, the sender will fail the batch with
-     * the exception returned by this method.
-     * @return An exception indicating the batch expired.
-     */
-    TimeoutException timeoutException() {
-        if (expiryErrorMessage == null)
-            throw new IllegalStateException("Batch has not expired");
-        return new TimeoutException("Expiring " + recordCount + " record(s) for " + topicPartition + ": " + expiryErrorMessage);
+    public FinalState finalState() {
+        return this.finalState.get();
     }
 
     int attempts() {
@@ -347,10 +352,6 @@ long queueTimeMs() {
         return drainedMs - createdMs;
     }
 
-    long createdTimeMs(long nowMs) {
-        return Math.max(0, nowMs - createdMs);
-    }
-
     long waitedTimeMs(long nowMs) {
         return Math.max(0, nowMs - lastAttemptMs);
     }
@@ -391,6 +392,8 @@ public void setProducerState(ProducerIdAndEpoch producerIdAndEpoch, int baseSequ
     }
 
     public void resetProducerState(ProducerIdAndEpoch producerIdAndEpoch, int baseSequence, boolean isTransactional) {
+        log.info("Resetting sequence number of batch with current sequence {} for partition {} to {}",
+                this.baseSequence(), this.topicPartition, baseSequence);
         reopened = true;
         recordsBuilder.reopenAndRewriteProducerState(producerIdAndEpoch.producerId, producerIdAndEpoch.epoch, baseSequence, isTransactional);
     }
@@ -456,6 +459,10 @@ public int baseSequence() {
         return recordsBuilder.baseSequence();
     }
 
+    public int lastSequence() {
+        return recordsBuilder.baseSequence() + recordsBuilder.numRecords() - 1;
+    }
+
     public boolean hasSequence() {
         return baseSequence() != RecordBatch.NO_SEQUENCE;
     }
@@ -467,5 +474,4 @@ public boolean isTransactional() {
     public boolean sequenceHasBeenReset() {
         return reopened;
     }
-
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerIdAndEpoch.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerIdAndEpoch.java
deleted file mode 100644
index ebd5cc3281c93..0000000000000
--- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerIdAndEpoch.java
+++ /dev/null
@@ -1,41 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements. See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License. You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.kafka.clients.producer.internals;
-
-import static org.apache.kafka.common.record.RecordBatch.NO_PRODUCER_EPOCH;
-import static org.apache.kafka.common.record.RecordBatch.NO_PRODUCER_ID;
-
-class ProducerIdAndEpoch {
-    static final ProducerIdAndEpoch NONE = new ProducerIdAndEpoch(NO_PRODUCER_ID, NO_PRODUCER_EPOCH);
-
-    public final long producerId;
-    public final short epoch;
-
-    ProducerIdAndEpoch(long producerId, short epoch) {
-        this.producerId = producerId;
-        this.epoch = epoch;
-    }
-
-    public boolean isValid() {
-        return NO_PRODUCER_ID < producerId;
-    }
-
-    @Override
-    public String toString() {
-        return "(producerId=" + producerId + ", epoch=" + epoch + ")";
-    }
-}
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
new file mode 100644
index 0000000000000..d7a88a10f8b29
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerMetadata.java
@@ -0,0 +1,159 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.clients.Metadata;
+import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.internals.ClusterResourceListeners;
+import org.apache.kafka.common.requests.MetadataRequest;
+import org.apache.kafka.common.requests.MetadataResponse;
+import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.common.utils.Time;
+import org.slf4j.Logger;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+
+public class ProducerMetadata extends Metadata {
+    // If a topic hasn't been accessed for this many milliseconds, it is removed from the cache.
+    private final long metadataIdleMs;
+
+    /* Topics with expiry time */
+    private final Map topics = new HashMap<>();
+    private final Set newTopics = new HashSet<>();
+    private final Logger log;
+    private final Time time;
+
+    public ProducerMetadata(long refreshBackoffMs,
+                            long metadataExpireMs,
+                            long metadataIdleMs,
+                            LogContext logContext,
+                            ClusterResourceListeners clusterResourceListeners,
+                            Time time) {
+        super(refreshBackoffMs, metadataExpireMs, logContext, clusterResourceListeners);
+        this.metadataIdleMs = metadataIdleMs;
+        this.log = logContext.logger(ProducerMetadata.class);
+        this.time = time;
+    }
+
+    @Override
+    public synchronized MetadataRequest.Builder newMetadataRequestBuilder() {
+        return new MetadataRequest.Builder(new ArrayList<>(topics.keySet()), true);
+    }
+
+    @Override
+    public synchronized MetadataRequest.Builder newMetadataRequestBuilderForNewTopics() {
+        return new MetadataRequest.Builder(new ArrayList<>(newTopics), true);
+    }
+
+    public synchronized void add(String topic, long nowMs) {
+        Objects.requireNonNull(topic, "topic cannot be null");
+        if (topics.put(topic, nowMs + metadataIdleMs) == null) {
+            newTopics.add(topic);
+            requestUpdateForNewTopics();
+        }
+    }
+
+    public synchronized int requestUpdateForTopic(String topic) {
+        if (newTopics.contains(topic)) {
+            return requestUpdateForNewTopics();
+        } else {
+            return requestUpdate();
+        }
+    }
+
+    // Visible for testing
+    synchronized Set topics() {
+        return topics.keySet();
+    }
+
+    // Visible for testing
+    synchronized Set newTopics() {
+        return newTopics;
+    }
+
+    public synchronized boolean containsTopic(String topic) {
+        return topics.containsKey(topic);
+    }
+
+    @Override
+    public synchronized boolean retainTopic(String topic, boolean isInternal, long nowMs) {
+        Long expireMs = topics.get(topic);
+        if (expireMs == null) {
+            return false;
+        } else if (newTopics.contains(topic)) {
+            return true;
+        } else if (expireMs <= nowMs) {
+            log.debug("Removing unused topic {} from the metadata list, expiryMs {} now {}", topic, expireMs, nowMs);
+            topics.remove(topic);
+            return false;
+        } else {
+            return true;
+        }
+    }
+
+    /**
+     * Wait for metadata update until the current version is larger than the last version we know of
+     */
+    public synchronized void awaitUpdate(final int lastVersion, final long timeoutMs) throws InterruptedException {
+        long currentTimeMs = time.milliseconds();
+        long deadlineMs = currentTimeMs + timeoutMs < 0 ? Long.MAX_VALUE : currentTimeMs + timeoutMs;
+        time.waitObject(this, () -> {
+            // Throw fatal exceptions, if there are any. Recoverable topic errors will be handled by the caller.
+            maybeThrowFatalException();
+            return updateVersion() > lastVersion || isClosed();
+        }, deadlineMs);
+
+        if (isClosed())
+            throw new KafkaException("Requested metadata update after close");
+    }
+
+    @Override
+    public synchronized void update(int requestVersion, MetadataResponse response, boolean isPartialUpdate, long nowMs) {
+        super.update(requestVersion, response, isPartialUpdate, nowMs);
+
+        // Remove all topics in the response that are in the new topic set. Note that if an error was encountered for a
+        // new topic's metadata, then any work to resolve the error will include the topic in a full metadata update.
+        if (!newTopics.isEmpty()) {
+            for (MetadataResponse.TopicMetadata metadata : response.topicMetadata()) {
+                newTopics.remove(metadata.topic());
+            }
+        }
+
+        notifyAll();
+    }
+
+    @Override
+    public synchronized void fatalError(KafkaException fatalException) {
+        super.fatalError(fatalException);
+        notifyAll();
+    }
+
+    /**
+     * Close this instance and notify any awaiting threads.
+     */
+    @Override
+    public synchronized void close() {
+        super.close();
+        notifyAll();
+    }
+
+}
diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerMetrics.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerMetrics.java
index f78f3d6647ca1..030a23299b93e 100644
--- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerMetrics.java
+++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/ProducerMetrics.java
@@ -36,8 +36,7 @@ public ProducerMetrics(Metrics metrics) {
     }
 
     private List getAllTemplates() {
-        List l = new ArrayList<>();
-        l.addAll(this.senderMetrics.allTemplates());
+        List l = new ArrayList<>(this.senderMetrics.allTemplates());
         return l;
     }
 
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 ba8c28ece27f6..378129727dbb9 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
@@ -16,9 +16,24 @@
  */
 package org.apache.kafka.clients.producer.internals;
 
+import java.nio.ByteBuffer;
+import java.util.ArrayDeque;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Deque;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.atomic.AtomicInteger;
 import org.apache.kafka.clients.ApiVersions;
 import org.apache.kafka.clients.producer.Callback;
+import org.apache.kafka.common.utils.ProducerIdAndEpoch;
 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;
@@ -28,36 +43,19 @@
 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.Meter;
 import org.apache.kafka.common.record.AbstractRecords;
 import org.apache.kafka.common.record.CompressionRatioEstimator;
 import org.apache.kafka.common.record.CompressionType;
-import org.apache.kafka.common.record.Record;
-import org.apache.kafka.common.record.RecordBatch;
 import org.apache.kafka.common.record.MemoryRecords;
 import org.apache.kafka.common.record.MemoryRecordsBuilder;
+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.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;
 
-import java.nio.ByteBuffer;
-import java.util.ArrayDeque;
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.Deque;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.Iterator;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
-import java.util.concurrent.ConcurrentMap;
-import java.util.concurrent.atomic.AtomicInteger;
-
 /**
  * This class acts as a queue that accumulates records into {@link MemoryRecords}
  * instances to be sent to the server.
@@ -73,8 +71,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 int deliveryTimeoutMs;
     private final BufferPool free;
     private final Time time;
     private final ApiVersions apiVersions;
@@ -84,13 +83,13 @@ public final class RecordAccumulator {
     private final Set muted;
     private int drainIndex;
     private final TransactionManager transactionManager;
+    private long nextBatchExpiryTimeMs = Long.MAX_VALUE; // the earliest time (absolute) a batch will expire.
 
     /**
      * Create a new record accumulator
      *
      * @param logContext The log context used for logging
      * @param batchSize The size to use when allocating {@link MemoryRecords} instances
-     * @param totalSize The maximum memory the record accumulator can use.
      * @param compression The compression codec for the records
      * @param lingerMs An artificial delay time to add before declaring a records instance that isn't full ready for
      *        sending. This allows time for more records to arrive. Setting a non-zero lingerMs will trade off some
@@ -105,14 +104,16 @@ public final class RecordAccumulator {
      */
     public RecordAccumulator(LogContext logContext,
                              int batchSize,
-                             long totalSize,
                              CompressionType compression,
-                             long lingerMs,
+                             int lingerMs,
                              long retryBackoffMs,
+                             int deliveryTimeoutMs,
                              Metrics metrics,
+                             String metricGrpName,
                              Time time,
                              ApiVersions apiVersions,
-                             TransactionManager transactionManager) {
+                             TransactionManager transactionManager,
+                             BufferPool bufferPool) {
         this.log = logContext.logger(RecordAccumulator.class);
         this.drainIndex = 0;
         this.closed = false;
@@ -122,9 +123,9 @@ public RecordAccumulator(LogContext logContext,
         this.compression = compression;
         this.lingerMs = lingerMs;
         this.retryBackoffMs = retryBackoffMs;
+        this.deliveryTimeoutMs = deliveryTimeoutMs;
         this.batches = new CopyOnWriteMap<>();
-        String metricGrpName = "producer-metrics";
-        this.free = new BufferPool(totalSize, batchSize, metrics, time, metricGrpName);
+        this.free = bufferPool;
         this.incomplete = new IncompleteBatches();
         this.muted = new HashSet<>();
         this.time = time;
@@ -157,11 +158,6 @@ public double measure(MetricConfig config, long now) {
             }
         };
         metrics.addMetric(metricName, availableBytes);
-
-        Sensor bufferExhaustedRecordSensor = metrics.sensor("buffer-exhausted-records");
-        MetricName rateMetricName = metrics.metricName("buffer-exhausted-rate", metricGrpName, "The average per-second number of record sends that are dropped due to buffer exhaustion");
-        MetricName totalMetricName = metrics.metricName("buffer-exhausted-total", metricGrpName, "The total number of record sends that are dropped due to buffer exhaustion");
-        bufferExhaustedRecordSensor.add(new Meter(rateMetricName, totalMetricName));
     }
 
     /**
@@ -177,6 +173,9 @@ 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 partitioner's onNewBatch method before trying to append again
+     * @param nowMs The current time, in milliseconds
      */
     public RecordAppendResult append(TopicPartition tp,
                                      long timestamp,
@@ -184,7 +183,9 @@ public RecordAppendResult append(TopicPartition tp,
                                      byte[] value,
                                      Header[] headers,
                                      Callback callback,
-                                     long maxTimeToBlock) throws InterruptedException {
+                                     long maxTimeToBlock,
+                                     boolean abortOnNewBatch,
+                                     long nowMs) throws InterruptedException {
         // We keep track of the number of appending thread to make sure we do not miss batches in
         // abortIncompleteBatches().
         appendsInProgress.incrementAndGet();
@@ -195,39 +196,47 @@ public RecordAppendResult append(TopicPartition tp,
             Deque dq = getOrCreateDeque(tp);
             synchronized (dq) {
                 if (closed)
-                    throw new IllegalStateException("Cannot send after the producer is closed.");
-                RecordAppendResult appendResult = tryAppend(timestamp, key, value, headers, callback, dq);
+                    throw new KafkaException("Producer closed while send in progress");
+                RecordAppendResult appendResult = tryAppend(timestamp, key, value, headers, callback, dq, nowMs);
                 if (appendResult != null)
                     return appendResult;
             }
 
             // 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());
+            log.trace("Allocating a new {} byte message buffer for topic {} partition {} with remaining timeout {}ms", size, tp.topic(), tp.partition(), maxTimeToBlock);
             buffer = free.allocate(size, maxTimeToBlock);
+
+            // Update the current time in case the buffer allocation blocked above.
+            nowMs = time.milliseconds();
             synchronized (dq) {
                 // Need to check if producer is closed again after grabbing the dequeue lock.
                 if (closed)
-                    throw new IllegalStateException("Cannot send after the producer is closed.");
+                    throw new KafkaException("Producer closed while send in progress");
 
-                RecordAppendResult appendResult = tryAppend(timestamp, key, value, headers, callback, dq);
+                RecordAppendResult appendResult = tryAppend(timestamp, key, value, headers, callback, dq, nowMs);
                 if (appendResult != null) {
                     // Somebody else found us a batch, return the one we waited for! Hopefully this doesn't happen often...
                     return appendResult;
                 }
 
                 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()));
+                ProducerBatch batch = new ProducerBatch(tp, recordsBuilder, nowMs);
+                FutureRecordMetadata future = Objects.requireNonNull(batch.tryAppend(timestamp, key, value, headers,
+                        callback, nowMs));
 
                 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)
@@ -239,7 +248,7 @@ public RecordAppendResult append(TopicPartition tp,
     private MemoryRecordsBuilder recordsBuilder(ByteBuffer buffer, byte maxUsableMagic) {
         if (transactionManager != null && maxUsableMagic < RecordBatch.MAGIC_VALUE_V2) {
             throw new UnsupportedVersionException("Attempting to use idempotence with a broker which does not " +
-                    "support the required message format (v2). The broker must be version 0.11 or later.");
+                "support the required message format (v2). The broker must be version 0.11 or later.");
         }
         return MemoryRecords.builder(buffer, maxUsableMagic, compression, TimestampType.CREATE_TIME, 0L);
     }
@@ -253,49 +262,55 @@ private MemoryRecordsBuilder recordsBuilder(ByteBuffer buffer, byte maxUsableMag
      *  if it is expired, or when the producer is closed.
      */
     private RecordAppendResult tryAppend(long timestamp, byte[] key, byte[] value, Header[] headers,
-                                         Callback callback, Deque deque) {
+                                         Callback callback, Deque deque, long nowMs) {
         ProducerBatch last = deque.peekLast();
         if (last != null) {
-            FutureRecordMetadata future = last.tryAppend(timestamp, key, value, headers, callback, time.milliseconds());
+            FutureRecordMetadata future = last.tryAppend(timestamp, key, value, headers, callback, nowMs);
             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) {
+        return muted.contains(tp);
+    }
+
+    public void resetNextBatchExpiryTime() {
+        nextBatchExpiryTimeMs = Long.MAX_VALUE;
+    }
+
+    public void maybeUpdateNextBatchExpiryTime(ProducerBatch batch) {
+        if (batch.createdMs + deliveryTimeoutMs  > 0) {
+            // the non-negative check is to guard us against potential overflow due to setting
+            // a large value for deliveryTimeoutMs
+            nextBatchExpiryTimeMs = Math.min(nextBatchExpiryTimeMs, batch.createdMs + deliveryTimeoutMs);
+        } else {
+            log.warn("Skipping next batch expiry time update due to addition overflow: "
+                + "batch.createMs={}, deliveryTimeoutMs={}", batch.createdMs, deliveryTimeoutMs);
+        }
+    }
+
     /**
      * Get a list of batches which have been sitting in the accumulator too long and need to be expired.
      */
-    public List expiredBatches(int requestTimeout, long now) {
+    public List expiredBatches(long now) {
         List expiredBatches = new ArrayList<>();
         for (Map.Entry> entry : this.batches.entrySet()) {
-            Deque dq = entry.getValue();
-            TopicPartition tp = entry.getKey();
-            // We only check if the batch should be expired if the partition does not have a batch in flight.
-            // This is to prevent later batches from being expired while an earlier batch is still in progress.
-            // Note that `muted` is only ever populated if `max.in.flight.request.per.connection=1` so this protection
-            // is only active in this case. Otherwise the expiration order is not guaranteed.
-            if (!muted.contains(tp)) {
-                synchronized (dq) {
-                    // iterate over the batches and expire them if they have been in the accumulator for more than requestTimeOut
-                    ProducerBatch lastBatch = dq.peekLast();
-                    Iterator batchIterator = dq.iterator();
-                    while (batchIterator.hasNext()) {
-                        ProducerBatch batch = batchIterator.next();
-                        boolean isFull = batch != lastBatch || batch.isFull();
-                        // Check if the batch has expired. Expired batches are closed by maybeExpire, but callbacks
-                        // are invoked after completing the iterations, since sends invoked from callbacks
-                        // may append more batches to the deque being iterated. The batch is deallocated after
-                        // callbacks are invoked.
-                        if (batch.maybeExpire(requestTimeout, retryBackoffMs, now, this.lingerMs, isFull)) {
-                            expiredBatches.add(batch);
-                            batchIterator.remove();
-                        } else {
-                            // Stop at the first batch that has not expired.
-                            break;
-                        }
+            // expire the batches in the order of sending
+            Deque deque = entry.getValue();
+            synchronized (deque) {
+                while (!deque.isEmpty()) {
+                    ProducerBatch batch = deque.getFirst();
+                    if (batch.hasReachedDeliveryTimeout(deliveryTimeoutMs, now)) {
+                        deque.poll();
+                        batch.abortRecordAppends();
+                        expiredBatches.add(batch);
+                    } else {
+                        maybeUpdateNextBatchExpiryTime(batch);
+                        break;
                     }
                 }
             }
@@ -303,8 +318,13 @@ public List expiredBatches(int requestTimeout, long now) {
         return expiredBatches;
     }
 
+    public long getDeliveryTimeoutMs() {
+        return deliveryTimeoutMs;
+    }
+
     /**
-     * Re-enqueue the given record batch in the accumulator to retry
+     * Re-enqueue the given record batch in the accumulator. In Sender.completeBatch method, we check
+     * whether the batch has reached deliveryTimeoutMs or not. Hence we do not do the delivery timeout check here.
      */
     public void reenqueue(ProducerBatch batch, long now) {
         batch.reenqueued(now);
@@ -348,8 +368,8 @@ public int splitAndReenqueue(ProducerBatch bigBatch) {
     }
 
     // We will have to do extra work to ensure the queue is in order when requests are being retried and there are
-    // multiple requests in flight to that partition. If the first inflight request fails to append, then all the subsequent
-    // in flight requests will also fail because the sequence numbers will not be accepted.
+    // multiple requests in flight to that partition. If the first in flight request fails to append, then all the
+    // subsequent in flight requests will also fail because the sequence numbers will not be accepted.
     //
     // Further, once batches are being retried, we are reduced to a single in flight request for that partition. So when
     // the subsequent batches come back in sequence order, they will have to be placed further back in the queue.
@@ -360,12 +380,12 @@ public int splitAndReenqueue(ProducerBatch bigBatch) {
     private void insertInSequenceOrder(Deque deque, ProducerBatch batch) {
         // When we are requeing and have enabled idempotence, the reenqueued batch must always have a sequence.
         if (batch.baseSequence() == RecordBatch.NO_SEQUENCE)
-            throw new IllegalStateException("Trying to reenqueue a batch which doesn't have a sequence even " +
-                    "though idempotence is enabled.");
+            throw new IllegalStateException("Trying to re-enqueue a batch which doesn't have a sequence even " +
+                "though idempotency is enabled.");
 
         if (transactionManager.nextBatchBySequence(batch.topicPartition) == null)
-            throw new IllegalStateException("We are reenqueueing a batch which is not tracked as part of the in flight " +
-                    "requests. batch.topicPartition: " + batch.topicPartition + "; batch.baseSequence: " + batch.baseSequence());
+            throw new IllegalStateException("We are re-enqueueing a batch which is not tracked as part of the in flight " +
+                "requests. batch.topicPartition: " + batch.topicPartition + "; batch.baseSequence: " + batch.baseSequence());
 
         ProducerBatch firstBatchInQueue = deque.peekFirst();
         if (firstBatchInQueue != null && firstBatchInQueue.hasSequence() && firstBatchInQueue.baseSequence() < batch.baseSequence()) {
@@ -382,7 +402,7 @@ private void insertInSequenceOrder(Deque deque, ProducerBatch bat
                 orderedBatches.add(deque.pollFirst());
 
             log.debug("Reordered incoming batch with sequence {} for partition {}. It was placed in the queue at " +
-                    "position {}", batch.baseSequence(), batch.topicPartition, orderedBatches.size());
+                "position {}", batch.baseSequence(), batch.topicPartition, orderedBatches.size());
             // Either we have reached a point where there are batches without a sequence (ie. never been drained
             // and are hence in order by default), or the batch at the front of the queue has a sequence greater
             // than the incoming batch. This is the right place to add the incoming batch.
@@ -427,18 +447,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) && !muted.contains(part)) {
-                    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)) {
                         long waitedTimeMs = batch.waitedTimeMs(nowMs);
                         boolean backingOff = batch.attempts() > 0 && waitedTimeMs < retryBackoffMs;
                         long timeToWaitMs = backingOff ? retryBackoffMs : lingerMs;
@@ -458,7 +479,6 @@ public ReadyCheckResult ready(Cluster cluster, long nowMs) {
                 }
             }
         }
-
         return new ReadyCheckResult(readyNodes, nextReadyCheckDelayMs, unknownLeaderTopics);
     }
 
@@ -476,6 +496,119 @@ public boolean hasUndrained() {
         return false;
     }
 
+    private boolean shouldStopDrainBatchesForPartition(ProducerBatch first, TopicPartition tp) {
+        ProducerIdAndEpoch producerIdAndEpoch = null;
+        if (transactionManager != null) {
+            if (!transactionManager.isSendToPartitionAllowed(tp))
+                return true;
+
+            producerIdAndEpoch = transactionManager.producerIdAndEpoch();
+            if (!producerIdAndEpoch.isValid())
+                // we cannot send the batch until we have refreshed the producer id
+                return true;
+
+            if (!first.hasSequence()) {
+                if (transactionManager.hasInflightBatches(tp)) {
+                    // Don't drain any new batches while the partition has in-flight batches with a different epoch
+                    // and/or producer ID. Otherwise, a batch with a new epoch and sequence number
+                    // 0 could be written before earlier batches complete, which would cause out of sequence errors
+                    ProducerBatch firstInFlightBatch = transactionManager.nextBatchBySequence(tp);
+
+                    if (firstInFlightBatch != null && transactionManager.producerIdOrEpochNotMatch(firstInFlightBatch)) {
+                        return true;
+                    }
+                }
+
+                if (transactionManager.hasUnresolvedSequence(first.topicPartition))
+                    // Don't drain any new batches while the state of previous sequence numbers
+                    // is unknown. The previous batches would be unknown if they were aborted
+                    // on the client after being sent to the broker at least once.
+                    return true;
+            }
+
+            int firstInFlightSequence = transactionManager.firstInFlightSequence(first.topicPartition);
+            if (firstInFlightSequence != RecordBatch.NO_SEQUENCE && first.hasSequence()
+                && first.baseSequence() != firstInFlightSequence)
+                // If the queued batch already has an assigned sequence, then it is being retried.
+                // In this case, we wait until the next immediate batch is ready and drain that.
+                // We only move on when the next in line batch is complete (either successfully or due to
+                // a fatal broker error). This effectively reduces our in flight request count to 1.
+                return true;
+        }
+        return false;
+    }
+
+    private List drainBatchesForOneNode(Cluster cluster, Node node, int maxSize, long now) {
+        int size = 0;
+        List parts = cluster.partitionsForNode(node.id());
+        List ready = new ArrayList<>();
+        /* to make starvation less likely this loop doesn't start at 0 */
+        int start = drainIndex = drainIndex % parts.size();
+        do {
+            PartitionInfo part = parts.get(drainIndex);
+            TopicPartition tp = new TopicPartition(part.topic(), part.partition());
+            this.drainIndex = (this.drainIndex + 1) % parts.size();
+
+            // Only proceed if the partition has no in-flight batches.
+            if (isMuted(tp))
+                continue;
+
+            Deque deque = getDeque(tp);
+            if (deque == null)
+                continue;
+
+            synchronized (deque) {
+                // invariant: !isMuted(tp,now) && deque != null
+                ProducerBatch first = deque.peekFirst();
+                if (first == null)
+                    continue;
+
+                // first != null
+                boolean backoff = first.attempts() > 0 && first.waitedTimeMs(now) < retryBackoffMs;
+                // Only drain the batch if it is not during backoff period.
+                if (backoff)
+                    continue;
+
+                if (size + first.estimatedSizeInBytes() > maxSize && !ready.isEmpty()) {
+                    // there is a rare case that a single batch size is larger than the request size due to
+                    // compression; in this case we will still eventually send this batch in a single request
+                    break;
+                } else {
+                    if (shouldStopDrainBatchesForPartition(first, tp))
+                        break;
+
+                    boolean isTransactional = transactionManager != null && transactionManager.isTransactional();
+                    ProducerIdAndEpoch producerIdAndEpoch =
+                        transactionManager != null ? transactionManager.producerIdAndEpoch() : null;
+                    ProducerBatch batch = deque.pollFirst();
+                    if (producerIdAndEpoch != null && !batch.hasSequence()) {
+                        // If the batch already has an assigned sequence, then we should not change the producer id and
+                        // sequence number, since this may introduce duplicates. In particular, the previous attempt
+                        // may actually have been accepted, and if we change the producer id and sequence here, this
+                        // attempt will also be accepted, causing a duplicate.
+                        //
+                        // Additionally, we update the next sequence number bound for the partition, and also have
+                        // the transaction manager track the batch so as to ensure that sequence ordering is maintained
+                        // even if we receive out of order responses.
+                        batch.setProducerState(producerIdAndEpoch, transactionManager.sequenceNumber(batch.topicPartition), isTransactional);
+                        transactionManager.incrementSequenceNumber(batch.topicPartition, batch.recordCount);
+                        log.debug("Assigned producerId {} and producerEpoch {} to batch with base sequence " +
+                                "{} being sent to partition {}", producerIdAndEpoch.producerId,
+                            producerIdAndEpoch.epoch, batch.baseSequence(), tp);
+
+                        transactionManager.addInFlightBatch(batch);
+                    }
+                    batch.close();
+                    size += batch.records().sizeInBytes();
+                    ready.add(batch);
+
+                    batch.drained(now);
+                }
+            }
+        } while (start != drainIndex);
+        return ready;
+    }
+
     /**
      * Drain all the data for the given nodes and collate them into a list of batches that will fit within the specified
      * size on a per-node basis. This method attempts to avoid choosing the same topic-node over and over.
@@ -486,106 +619,25 @@ public boolean hasUndrained() {
      * @param now The current unix time in milliseconds
      * @return A list of {@link ProducerBatch} for each node specified with total size less than the requested maxSize.
      */
-    public Map> drain(Cluster cluster,
-                                                   Set nodes,
-                                                   int maxSize,
-                                                   long now) {
+    public Map> drain(Cluster cluster, Set nodes, int maxSize, long now) {
         if (nodes.isEmpty())
             return Collections.emptyMap();
 
         Map> batches = new HashMap<>();
         for (Node node : nodes) {
-            int size = 0;
-            List parts = cluster.partitionsForNode(node.id());
-            List ready = new ArrayList<>();
-            /* to make starvation less likely this loop doesn't start at 0 */
-            int start = drainIndex = drainIndex % parts.size();
-            do {
-                PartitionInfo part = parts.get(drainIndex);
-                TopicPartition tp = new TopicPartition(part.topic(), part.partition());
-                // Only proceed if the partition has no in-flight batches.
-                if (!muted.contains(tp)) {
-                    Deque deque = getDeque(tp);
-                    if (deque != null) {
-                        synchronized (deque) {
-                            ProducerBatch first = deque.peekFirst();
-                            if (first != null) {
-                                boolean backoff = first.attempts() > 0 && first.waitedTimeMs(now) < retryBackoffMs;
-                                // Only drain the batch if it is not during backoff period.
-                                if (!backoff) {
-                                    if (size + first.estimatedSizeInBytes() > maxSize && !ready.isEmpty()) {
-                                        // there is a rare case that a single batch size is larger than the request size due
-                                        // to compression; in this case we will still eventually send this batch in a single
-                                        // request
-                                        break;
-                                    } else {
-                                        ProducerIdAndEpoch producerIdAndEpoch = null;
-                                        boolean isTransactional = false;
-                                        if (transactionManager != null) {
-                                            if (!transactionManager.isSendToPartitionAllowed(tp))
-                                                break;
-
-                                            producerIdAndEpoch = transactionManager.producerIdAndEpoch();
-                                            if (!producerIdAndEpoch.isValid())
-                                                // we cannot send the batch until we have refreshed the producer id
-                                                break;
-
-                                            isTransactional = transactionManager.isTransactional();
-
-                                            if (!first.hasSequence() && transactionManager.hasUnresolvedSequence(first.topicPartition))
-                                                // Don't drain any new batches while the state of previous sequence numbers
-                                                // is unknown. The previous batches would be unknown if they were aborted
-                                                // on the client after being sent to the broker at least once.
-                                                break;
-
-                                            int firstInFlightSequence = transactionManager.firstInFlightSequence(first.topicPartition);
-                                            if (firstInFlightSequence != RecordBatch.NO_SEQUENCE && first.hasSequence()
-                                                    && first.baseSequence() != firstInFlightSequence)
-                                                // If the queued batch already has an assigned sequence, then it is being
-                                                // retried. In this case, we wait until the next immediate batch is ready
-                                                // and drain that. We only move on when the next in line batch is complete (either successfully
-                                                // or due to a fatal broker error). This effectively reduces our
-                                                // in flight request count to 1.
-                                                break;
-                                        }
-
-                                        ProducerBatch batch = deque.pollFirst();
-                                        if (producerIdAndEpoch != null && !batch.hasSequence()) {
-                                            // If the batch already has an assigned sequence, then we should not change the producer id and
-                                            // sequence number, since this may introduce duplicates. In particular,
-                                            // the previous attempt may actually have been accepted, and if we change
-                                            // the producer id and sequence here, this attempt will also be accepted,
-                                            // causing a duplicate.
-                                            //
-                                            // Additionally, we update the next sequence number bound for the partition,
-                                            // and also have the transaction manager track the batch so as to ensure
-                                            // that sequence ordering is maintained even if we receive out of order
-                                            // responses.
-                                            batch.setProducerState(producerIdAndEpoch, transactionManager.sequenceNumber(batch.topicPartition), isTransactional);
-                                            transactionManager.incrementSequenceNumber(batch.topicPartition, batch.recordCount);
-                                            log.debug("Assigned producerId {} and producerEpoch {} to batch with base sequence " +
-                                                            "{} being sent to partition {}", producerIdAndEpoch.producerId,
-                                                    producerIdAndEpoch.epoch, batch.baseSequence(), tp);
-
-                                            transactionManager.addInFlightBatch(batch);
-                                        }
-                                        batch.close();
-                                        size += batch.records().sizeInBytes();
-                                        ready.add(batch);
-                                        batch.drained(now);
-                                    }
-                                }
-                            }
-                        }
-                    }
-                }
-                this.drainIndex = (this.drainIndex + 1) % parts.size();
-            } while (start != drainIndex);
+            List ready = drainBatchesForOneNode(cluster, node, maxSize, now);
             batches.put(node.id(), ready);
         }
         return batches;
     }
 
+    /**
+     * The earliest absolute time a batch will expire (in milliseconds)
+     */
+    public long nextExpiryTimeMs() {
+        return this.nextBatchExpiryTimeMs;
+    }
+
     private Deque getDeque(TopicPartition tp) {
         return batches.get(tp);
     }
@@ -693,7 +745,7 @@ public void abortIncompleteBatches() {
      * Go through incomplete batches and abort them.
      */
     private void abortBatches() {
-        abortBatches(new IllegalStateException("Producer is closed forcefully."));
+        abortBatches(new KafkaException("Producer is closed forcefully."));
     }
 
     /**
@@ -745,6 +797,7 @@ public void unmutePartition(TopicPartition tp) {
      */
     public void close() {
         this.closed = true;
+        this.free.close();
     }
 
     /*
@@ -754,11 +807,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;
         }
     }
 
@@ -776,5 +831,4 @@ public ReadyCheckResult(Set readyNodes, long nextReadyCheckDelayMs, Set> inFlightBatches;
+
     public Sender(LogContext logContext,
                   KafkaClient client,
-                  Metadata metadata,
+                  ProducerMetadata metadata,
                   RecordAccumulator accumulator,
                   boolean guaranteeMessageOrder,
                   int maxRequestSize,
@@ -130,7 +130,7 @@ public Sender(LogContext logContext,
                   int retries,
                   SenderMetricsRegistry metricsRegistry,
                   Time time,
-                  int requestTimeout,
+                  int requestTimeoutMs,
                   long retryBackoffMs,
                   TransactionManager transactionManager,
                   ApiVersions apiVersions) {
@@ -144,23 +144,102 @@ public Sender(LogContext logContext,
         this.acks = acks;
         this.retries = retries;
         this.time = time;
-        this.sensors = new SenderMetrics(metricsRegistry);
-        this.requestTimeout = requestTimeout;
+        this.sensors = new SenderMetrics(metricsRegistry, metadata, client, time);
+        this.requestTimeoutMs = requestTimeoutMs;
         this.retryBackoffMs = retryBackoffMs;
         this.apiVersions = apiVersions;
         this.transactionManager = transactionManager;
+        this.inFlightBatches = new HashMap<>();
+    }
+
+    public List inFlightBatches(TopicPartition tp) {
+        return inFlightBatches.containsKey(tp) ? inFlightBatches.get(tp) : new ArrayList<>();
+    }
+
+    private void maybeRemoveFromInflightBatches(ProducerBatch batch) {
+        List batches = inFlightBatches.get(batch.topicPartition);
+        if (batches != null) {
+            batches.remove(batch);
+            if (batches.isEmpty()) {
+                inFlightBatches.remove(batch.topicPartition);
+            }
+        }
+    }
+
+    private void maybeRemoveAndDeallocateBatch(ProducerBatch batch) {
+        maybeRemoveFromInflightBatches(batch);
+        this.accumulator.deallocate(batch);
+    }
+
+    /**
+     *  Get the in-flight batches that has reached delivery timeout.
+     */
+    private List getExpiredInflightBatches(long now) {
+        List expiredBatches = new ArrayList<>();
+
+        for (Iterator>> batchIt = inFlightBatches.entrySet().iterator(); batchIt.hasNext();) {
+            Map.Entry> entry = batchIt.next();
+            List partitionInFlightBatches = entry.getValue();
+            if (partitionInFlightBatches != null) {
+                Iterator iter = partitionInFlightBatches.iterator();
+                while (iter.hasNext()) {
+                    ProducerBatch batch = iter.next();
+                    if (batch.hasReachedDeliveryTimeout(accumulator.getDeliveryTimeoutMs(), now)) {
+                        iter.remove();
+                        // expireBatches is called in Sender.sendProducerData, before client.poll.
+                        // The !batch.isDone() invariant should always hold. An IllegalStateException
+                        // exception will be thrown if the invariant is violated.
+                        if (!batch.isDone()) {
+                            expiredBatches.add(batch);
+                        } else {
+                            throw new IllegalStateException(batch.topicPartition + " batch created at " +
+                                batch.createdMs + " gets unexpected final state " + batch.finalState());
+                        }
+                    } else {
+                        accumulator.maybeUpdateNextBatchExpiryTime(batch);
+                        break;
+                    }
+                }
+                if (partitionInFlightBatches.isEmpty()) {
+                    batchIt.remove();
+                }
+            }
+        }
+        return expiredBatches;
+    }
+
+    private void addToInflightBatches(List batches) {
+        for (ProducerBatch batch : batches) {
+            List inflightBatchList = inFlightBatches.get(batch.topicPartition);
+            if (inflightBatchList == null) {
+                inflightBatchList = new ArrayList<>();
+                inFlightBatches.put(batch.topicPartition, inflightBatchList);
+            }
+            inflightBatchList.add(batch);
+        }
+    }
+
+    public void addToInflightBatches(Map> batches) {
+        for (List batchList : batches.values()) {
+            addToInflightBatches(batchList);
+        }
+    }
+
+    private boolean hasPendingTransactionalRequests() {
+        return transactionManager != null && transactionManager.hasPendingRequests() && transactionManager.hasOngoingTransaction();
     }
 
     /**
      * The main run loop for the sender thread
      */
+    @Override
     public void run() {
         log.debug("Starting Kafka producer I/O thread.");
 
         // 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);
             }
@@ -169,18 +248,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();
         }
@@ -196,52 +293,42 @@ 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();
-
-                if (!transactionManager.isTransactional()) {
-                    // this is an idempotent producer, so make sure we have a producer id
-                    maybeWaitForProducerId();
-                } else if (transactionManager.hasUnresolvedSequences() && !transactionManager.hasFatalError()) {
-                    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);
-                    return;
-                }
+                transactionManager.maybeResolveSequences();
 
-                // do not continue sending if the transaction manager is in a failed state or if there
-                // is no producer id (for the idempotent case).
-                if (transactionManager.hasFatalError() || !transactionManager.hasProducerId()) {
+                // do not continue sending if the transaction manager is in a failed state
+                if (transactionManager.hasFatalError()) {
                     RuntimeException lastError = transactionManager.lastError();
                     if (lastError != null)
                         maybeAbortBatches(lastError);
-                    client.poll(retryBackoffMs, now);
+                    client.poll(retryBackoffMs, time.milliseconds());
+                    return;
+                }
+
+                // Check whether we need a new producerId. If so, we will enqueue an InitProducerId
+                // request which will be sent below
+                transactionManager.bumpIdempotentEpochAndResetIdIfNeeded();
+
+                if (maybeSendAndPollTransactionalRequest()) {
                     return;
-                } else if (transactionManager.hasAbortableError()) {
-                    accumulator.abortUndrainedBatches(transactionManager.lastError());
                 }
             } catch (AuthenticationException e) {
                 // This is already logged as error, but propagated here to perform any clean ups.
-                log.trace("Authentication exception while processing transactional request: {}", e);
+                log.trace("Authentication exception while processing transactional request", e);
                 transactionManager.authenticationFailed(e);
             }
         }
 
-        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) {
         Cluster cluster = metadata.fetch();
-
         // get the list of partitions with data ready to send
         RecordAccumulator.ReadyCheckResult result = this.accumulator.ready(cluster, now);
 
@@ -251,7 +338,10 @@ private long sendProducerData(long now) {
             // topics which may have expired. Add the topic again to metadata to ensure it is included
             // and request metadata update, since there are messages to send to the topic.
             for (String topic : result.unknownLeaderTopics)
-                this.metadata.add(topic);
+                this.metadata.add(topic, now);
+
+            log.debug("Requesting metadata update due to unknown leader topics from the batched records: {}",
+                result.unknownLeaderTopics);
             this.metadata.requestUpdate();
         }
 
@@ -262,13 +352,13 @@ private long sendProducerData(long now) {
             Node node = iter.next();
             if (!this.client.ready(node, now)) {
                 iter.remove();
-                notReadyTimeout = Math.min(notReadyTimeout, this.client.connectionDelay(node, now));
+                notReadyTimeout = Math.min(notReadyTimeout, this.client.pollDelayMs(node, now));
             }
         }
 
         // create produce requests
-        Map> batches = this.accumulator.drain(cluster, result.readyNodes,
-                this.maxRequestSize, now);
+        Map> batches = this.accumulator.drain(cluster, result.readyNodes, this.maxRequestSize, now);
+        addToInflightBatches(batches);
         if (guaranteeMessageOrder) {
             // Mute all the partitions drained
             for (List batchList : batches.values()) {
@@ -277,27 +367,35 @@ private long sendProducerData(long now) {
             }
         }
 
-        List expiredBatches = this.accumulator.expiredBatches(this.requestTimeout, now);
+        accumulator.resetNextBatchExpiryTime();
+        List expiredInflightBatches = getExpiredInflightBatches(now);
+        List expiredBatches = this.accumulator.expiredBatches(now);
+        expiredBatches.addAll(expiredInflightBatches);
+
         // Reset the producer id if an expired batch has previously been sent to the broker. Also update the metrics
-        // for expired batches. see the documentation of @TransactionState.resetProducerId to understand why
+        // for expired batches. see the documentation of @TransactionState.resetIdempotentProducerId to understand why
         // we need to reset the producer id here.
         if (!expiredBatches.isEmpty())
             log.trace("Expired {} batches in accumulator", expiredBatches.size());
         for (ProducerBatch expiredBatch : expiredBatches) {
-            failBatch(expiredBatch, -1, NO_TIMESTAMP, expiredBatch.timeoutException(), false);
+            String errorMessage = "Expiring " + expiredBatch.recordCount + " record(s) for " + expiredBatch.topicPartition
+                + ":" + (now - expiredBatch.createdMs) + " ms has passed since batch creation";
+            failBatch(expiredBatch, -1, NO_TIMESTAMP, new TimeoutException(errorMessage), false);
             if (transactionManager != null && expiredBatch.inRetry()) {
                 // This ensures that no new batches are drained until the current in flight batches are fully resolved.
-                transactionManager.markSequenceUnresolved(expiredBatch.topicPartition);
+                transactionManager.markSequenceUnresolved(expiredBatch);
             }
         }
-
         sensors.updateProduceRequestMetrics(batches);
 
         // If we have any nodes that are ready to send + have sendable data, poll with 0 timeout so this can immediately
-        // loop and try sending more data. Otherwise, the timeout is determined by nodes that have partitions with data
-        // that isn't yet sendable (e.g. lingering, backing off). Note that this specifically does not include nodes
-        // with sendable data that aren't ready to send since they would cause busy looping.
+        // loop and try sending more data. Otherwise, the timeout will be the smaller value between next batch expiry
+        // time, and the delay time for checking data availability. Note that the nodes may have data that isn't yet
+        // sendable due to lingering, backing off, etc. This specifically does not include nodes with sendable data
+        // that aren't ready to send since they would cause busy looping.
         long pollTimeout = Math.min(result.nextReadyCheckDelayMs, notReadyTimeout);
+        pollTimeout = Math.min(pollTimeout, this.accumulator.nextExpiryTimeMs() - now);
+        pollTimeout = Math.max(pollTimeout, 0);
         if (!result.readyNodes.isEmpty()) {
             log.trace("Nodes with data ready to send: {}", result.readyNodes);
             // if some partitions are already ready to be sent, the select time would be 0;
@@ -307,74 +405,98 @@ private long sendProducerData(long now) {
             pollTimeout = 0;
         }
         sendProduceRequests(batches, now);
-
         return pollTimeout;
     }
 
-    private boolean maybeSendTransactionalRequest(long now) {
-        if (transactionManager.isCompleting() && accumulator.hasIncomplete()) {
-            if (transactionManager.isAborting())
-                accumulator.abortUndrainedBatches(new KafkaException("Failing batch since transaction was aborted"));
+    /**
+     * Returns true if a transactional request is sent or polled, or if a FindCoordinator request is enqueued
+     */
+    private boolean maybeSendAndPollTransactionalRequest() {
+        if (transactionManager.hasInFlightRequest()) {
+            // as long as there are outstanding transactional requests, we simply wait for them to return
+            client.poll(retryBackoffMs, time.milliseconds());
+            return true;
+        }
+
+        if (transactionManager.hasAbortableError() || transactionManager.isAborting()) {
+            if (accumulator.hasIncomplete()) {
+                // Attempt to get the last error that caused this abort.
+                RuntimeException exception = transactionManager.lastError();
+                // If there was no error, but we are still aborting,
+                // then this is most likely a case where there was no fatal error.
+                if (exception == null) {
+                    exception = new TransactionAbortedException();
+                }
+                accumulator.abortUndrainedBatches(exception);
+            }
+        }
 
+        if (transactionManager.isCompleting() && !accumulator.flushInProgress()) {
             // There may still be requests left which are being retried. Since we do not know whether they had
             // been successfully appended to the broker log, we must resend them until their final status is clear.
             // If they had been appended and we did not receive the error, then our sequence number would no longer
             // be correct which would lead to an OutOfSequenceException.
-            if (!accumulator.flushInProgress())
-                accumulator.beginFlush();
+            accumulator.beginFlush();
         }
 
-        TransactionManager.TxnRequestHandler nextRequestHandler = transactionManager.nextRequestHandler(accumulator.hasIncomplete());
+        TransactionManager.TxnRequestHandler nextRequestHandler = transactionManager.nextRequest(accumulator.hasIncomplete());
         if (nextRequestHandler == null)
             return false;
 
         AbstractRequest.Builder requestBuilder = nextRequestHandler.requestBuilder();
-        while (true) {
-            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, requestTimeout)) {
-                        transactionManager.lookupCoordinator(nextRequestHandler);
-                        break;
-                    }
-                } else {
-                    targetNode = awaitLeastLoadedNodeReady(requestTimeout);
-                }
-
-                if (targetNode != null) {
-                    if (nextRequestHandler.isRetry())
-                        time.sleep(nextRequestHandler.retryBackoffMs());
-
-                    ClientRequest clientRequest = client.newClientRequest(targetNode.idString(),
-                            requestBuilder, now, true, nextRequestHandler);
-                    transactionManager.setInFlightTransactionalRequestCorrelationId(clientRequest.correlationId());
-                    log.debug("Sending transactional request {} to node {}", requestBuilder, targetNode);
-
-                    client.send(clientRequest, now);
+        Node targetNode = null;
+        try {
+            FindCoordinatorRequest.CoordinatorType coordinatorType = nextRequestHandler.coordinatorType();
+            targetNode = coordinatorType != null ?
+                    transactionManager.coordinator(coordinatorType) :
+                    client.leastLoadedNode(time.milliseconds());
+            if (targetNode != null) {
+                if (!awaitNodeReady(targetNode, coordinatorType)) {
+                    log.trace("Target node {} not ready within request timeout, will retry when node is ready.", targetNode);
+                    maybeFindCoordinatorAndRetry(nextRequestHandler);
                     return true;
                 }
-            } catch (IOException e) {
-                log.debug("Disconnect from {} while trying to send request {}. Going " +
-                        "to back off and retry", targetNode, requestBuilder);
-                if (nextRequestHandler.needsCoordinator()) {
-                    // We break here so that we pick up the FindCoordinator request immediately.
-                    transactionManager.lookupCoordinator(nextRequestHandler);
-                    break;
-                }
+            } else if (coordinatorType != null) {
+                log.trace("Coordinator not known for {}, will retry {} after finding coordinator.", coordinatorType, requestBuilder.apiKey());
+                maybeFindCoordinatorAndRetry(nextRequestHandler);
+                return true;
+            } else {
+                log.trace("No nodes available to send requests, will poll and retry when until a node is ready.");
+                transactionManager.retry(nextRequestHandler);
+                client.poll(retryBackoffMs, time.milliseconds());
+                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 {} with correlation ID {}", requestBuilder, targetNode, clientRequest.correlationId());
+            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.
+            maybeFindCoordinatorAndRetry(nextRequestHandler);
+            return true;
+        }
+    }
+
+    private void maybeFindCoordinatorAndRetry(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) {
@@ -403,54 +525,21 @@ public void forceClose() {
         initiateClose();
     }
 
-    private ClientResponse sendAndAwaitInitProducerIdRequest(Node node) throws IOException {
-        String nodeId = node.idString();
-        InitProducerIdRequest.Builder builder = new InitProducerIdRequest.Builder(null);
-        ClientRequest request = client.newClientRequest(nodeId, builder, time.milliseconds(), true, 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)) {
-            return node;
-        }
-        return null;
+    public boolean isRunning() {
+        return running;
     }
 
-    private void maybeWaitForProducerId() {
-        while (!transactionManager.hasProducerId() && !transactionManager.hasError()) {
-            try {
-                Node node = awaitLeastLoadedNodeReady(requestTimeout);
-                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());
-                        transactionManager.setProducerIdAndEpoch(producerIdAndEpoch);
-                        return;
-                    } else if (error.exception() instanceof RetriableException) {
-                        log.debug("Retriable error from InitProducerId response", error.message());
-                    } else {
-                        transactionManager.transitionToFatalError(error.exception());
-                        break;
-                    }
-                } else {
-                    log.debug("Could not find an available broker to send InitProducerIdRequest to. " +
-                            "We will back off and try again.");
-                }
-            } catch (UnsupportedVersionException e) {
-                transactionManager.transitionToFatalError(e);
-                break;
-            } catch (IOException e) {
-                log.debug("Broker {} disconnected while awaiting InitProducerId response", e);
+    private boolean awaitNodeReady(Node node, FindCoordinatorRequest.CoordinatorType coordinatorType) throws IOException {
+        if (NetworkClientUtils.awaitReady(client, node, time, requestTimeoutMs)) {
+            if (coordinatorType == FindCoordinatorRequest.CoordinatorType.TRANSACTION) {
+                // Indicate to the transaction manager that the coordinator is ready, allowing it to check ApiVersions
+                // This allows us to bump transactional epochs even if the coordinator is temporarily unavailable at
+                // the time when the abortable error is handled
+                transactionManager.handleCoordinatorReady();
             }
-            log.trace("Retry InitProducerIdRequest in {}ms.", retryBackoffMs);
-            time.sleep(retryBackoffMs);
-            metadata.requestUpdate();
+            return true;
         }
+        return false;
     }
 
     /**
@@ -461,9 +550,10 @@ private void handleProduceResponse(ClientResponse response, Map entry : produceResponse.responses().entrySet()) {
-                    TopicPartition tp = entry.getKey();
-                    ProduceResponse.PartitionResponse partResp = entry.getValue();
+                produceResponse.data().responses().forEach(r -> r.partitionResponses().forEach(p -> {
+                    TopicPartition tp = new TopicPartition(r.name(), p.index());
+                    ProduceResponse.PartitionResponse partResp = new ProduceResponse.PartitionResponse(
+                            Errors.forCode(p.errorCode()),
+                            p.baseOffset(),
+                            p.logAppendTimeMs(),
+                            p.logStartOffset(),
+                            p.recordErrors()
+                                .stream()
+                                .map(e -> new ProduceResponse.RecordError(e.batchIndex(), e.batchIndexErrorMessage()))
+                                .collect(Collectors.toList()),
+                            p.errorMessage());
                     ProducerBatch batch = batches.get(tp);
                     completeBatch(batch, partResp, correlationId, now);
-                }
+                }));
                 this.sensors.recordLatency(response.destination(), response.requestLatencyMs());
             } else {
                 // this is the acks = 0 case, just complete all requests
@@ -502,40 +603,30 @@ private void completeBatch(ProducerBatch batch, ProduceResponse.PartitionRespons
                                long now) {
         Errors error = response.error;
 
-        if (error == Errors.MESSAGE_TOO_LARGE && batch.recordCount > 1 &&
+        if (error == Errors.MESSAGE_TOO_LARGE && batch.recordCount > 1 && !batch.isDone() &&
                 (batch.magic() >= RecordBatch.MAGIC_VALUE_V2 || batch.isCompressed())) {
             // If the batch is too large, we split the batch and send the split batches again. We do not decrement
             // the retry attempts in this case.
-            log.warn("Got error produce response in correlation id {} on topic-partition {}, splitting and retrying ({} attempts left). Error: {}",
-                     correlationId,
-                     batch.topicPartition,
-                     this.retries - batch.attempts(),
-                     error);
+            log.warn(
+                "Got error produce response in correlation id {} on topic-partition {}, splitting and retrying ({} attempts left). Error: {}",
+                correlationId,
+                batch.topicPartition,
+                this.retries - batch.attempts(),
+                formatErrMsg(response));
             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)) {
-                log.warn("Got error produce response with correlation id {} on topic-partition {}, retrying ({} attempts left). Error: {}",
-                        correlationId,
-                        batch.topicPartition,
-                        this.retries - batch.attempts() - 1,
-                        error);
-                if (transactionManager == null) {
-                    reenqueueBatch(batch, now);
-                } else if (transactionManager.hasProducerIdAndEpoch(batch.producerId(), batch.producerEpoch())) {
-                    // If idempotence is enabled only retry the request if the current producer id is the same as
-                    // the producer id of the batch.
-                    log.debug("Retrying batch to topic-partition {}. ProducerId: {}; Sequence number : {}",
-                            batch.topicPartition, batch.producerId(), batch.baseSequence());
-                    reenqueueBatch(batch, now);
-                } else {
-                    failBatch(batch, response, new OutOfOrderSequenceException("Attempted to retry sending a " +
-                            "batch but the producer id changed from " + batch.producerId() + " to " +
-                            transactionManager.producerIdAndEpoch().producerId + " in the mean time. This batch will be dropped."), false);
-                }
+            if (canRetry(batch, response, now)) {
+                log.warn(
+                    "Got error produce response with correlation id {} on topic-partition {}, retrying ({} attempts left). Error: {}",
+                    correlationId,
+                    batch.topicPartition,
+                    this.retries - batch.attempts() - 1,
+                    formatErrMsg(response));
+                reenqueueBatch(batch, now);
             } else if (error == Errors.DUPLICATE_SEQUENCE_NUMBER) {
                 // If we have received a duplicate sequence error, it means that the sequence number has advanced beyond
                 // the sequence of the current batch, and we haven't retained batch metadata on the broker to return
@@ -546,23 +637,28 @@ 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
-                    exception = error.exception();
+                    exception = error.exception(response.errorMessage);
                 // tell the user the result of their request. We only adjust sequence numbers if the batch didn't exhaust
                 // its retries -- if it did, we don't know whether the sequence number was accepted or not, and
                 // thus it is not safe to reassign the sequence.
                 failBatch(batch, response, exception, batch.attempts() < this.retries);
             }
             if (error.exception() instanceof InvalidMetadataException) {
-                if (error.exception() instanceof UnknownTopicOrPartitionException)
+                if (error.exception() instanceof UnknownTopicOrPartitionException) {
                     log.warn("Received unknown topic or partition error in produce request on partition {}. The " +
-                            "topic/partition may not exist or the user may not have Describe access to it", batch.topicPartition);
+                            "topic-partition may not exist or the user may not have Describe access to it",
+                        batch.topicPartition);
+                } else {
+                    log.warn("Received invalid metadata error in produce request on partition {} due to {}. Going " +
+                            "to request metadata update now", batch.topicPartition,
+                            error.exception(response.errorMessage).toString());
+                }
                 metadata.requestUpdate();
             }
-
         } else {
             completeBatch(batch, response);
         }
@@ -572,71 +668,67 @@ else if (error == Errors.CLUSTER_AUTHORIZATION_FAILED)
             this.accumulator.unmutePartition(batch.topicPartition);
     }
 
+    /**
+     * Format the error from a {@link ProduceResponse.PartitionResponse} in a user-friendly string
+     * e.g "NETWORK_EXCEPTION. Error Message: Disconnected from node 0"
+     */
+    private String formatErrMsg(ProduceResponse.PartitionResponse response) {
+        String errorMessageSuffix = (response.errorMessage == null || response.errorMessage.isEmpty()) ?
+                "" : String.format(". Error Message: %s", response.errorMessage);
+        return String.format("%s%s", response.error, errorMessageSuffix);
+    }
+
     private void reenqueueBatch(ProducerBatch batch, long currentTimeMs) {
         this.accumulator.reenqueue(batch, currentTimeMs);
+        maybeRemoveFromInflightBatches(batch);
         this.sensors.recordRetries(batch.topicPartition.topic(), batch.recordCount);
     }
 
     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));
-            }
-            transactionManager.updateLastAckedOffset(response, batch);
-            transactionManager.removeInFlightBatch(batch);
+            transactionManager.handleCompletedBatch(batch, response);
         }
 
-        if (batch.done(response.baseOffset, response.logAppendTime, null))
-            this.accumulator.deallocate(batch);
+        if (batch.done(response.baseOffset, response.logAppendTime, null)) {
+            maybeRemoveAndDeallocateBatch(batch);
+        }
     }
 
-    private void failBatch(ProducerBatch batch, ProduceResponse.PartitionResponse response, RuntimeException exception, boolean adjustSequenceNumbers) {
+    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))
-            this.accumulator.deallocate(batch);
+        if (batch.done(baseOffset, logAppendTime, exception)) {
+            maybeRemoveAndDeallocateBatch(batch);
+        }
     }
 
     /**
      * We can retry a send if the error is transient and the number of attempts taken is fewer than the maximum allowed.
-     * We can also retry OutOfOrderSequence exceptions for future batches, since if the first batch has failed, the future
-     * batches are certain to fail with an OutOfOrderSequence exception.
+     * We can also retry OutOfOrderSequence exceptions for future batches, since if the first batch has failed, the
+     * future batches are certain to fail with an OutOfOrderSequence exception.
      */
-    private boolean canRetry(ProducerBatch batch, ProduceResponse.PartitionResponse response) {
-        return batch.attempts() < this.retries &&
-                ((response.error.exception() instanceof RetriableException) ||
-                        (transactionManager != null && transactionManager.canRetry(response, batch)));
+    private boolean canRetry(ProducerBatch batch, ProduceResponse.PartitionResponse response, long now) {
+        return !batch.hasReachedDeliveryTimeout(accumulator.getDeliveryTimeoutMs(), now) &&
+            batch.attempts() < this.retries &&
+            !batch.isDone() &&
+            (transactionManager == null ?
+                    response.error.exception() instanceof RetriableException :
+                    transactionManager.canRetry(response, batch));
     }
 
     /**
@@ -644,7 +736,7 @@ private boolean canRetry(ProducerBatch batch, ProduceResponse.PartitionResponse
      */
     private void sendProduceRequests(Map> collated, long now) {
         for (Map.Entry> entry : collated.entrySet())
-            sendProduceRequest(now, entry.getKey(), acks, requestTimeout, entry.getValue());
+            sendProduceRequest(now, entry.getKey(), acks, requestTimeoutMs, entry.getValue());
     }
 
     /**
@@ -654,7 +746,6 @@ private void sendProduceRequest(long now, int destination, short acks, int timeo
         if (batches.isEmpty())
             return;
 
-        Map produceRecordsByPartition = new HashMap<>(batches.size());
         final Map recordsByPartition = new HashMap<>(batches.size());
 
         // find the minimum magic version used when creating the record sets
@@ -663,7 +754,7 @@ private void sendProduceRequest(long now, int destination, short acks, int timeo
             if (batch.magic() < minUsedMagic)
                 minUsedMagic = batch.magic();
         }
-
+        ProduceRequestData.TopicProduceDataCollection tpd = new ProduceRequestData.TopicProduceDataCollection();
         for (ProducerBatch batch : batches) {
             TopicPartition tp = batch.topicPartition;
             MemoryRecords records = batch.records();
@@ -677,7 +768,14 @@ private void sendProduceRequest(long now, int destination, short acks, int timeo
             // which is supporting the new magic version to one which doesn't, then we will need to convert.
             if (!records.hasMatchingMagic(minUsedMagic))
                 records = batch.records().downConvert(minUsedMagic, 0, time).records();
-            produceRecordsByPartition.put(tp, records);
+            ProduceRequestData.TopicProduceData tpData = tpd.find(tp.topic());
+            if (tpData == null) {
+                tpData = new ProduceRequestData.TopicProduceData().setName(tp.topic());
+                tpd.add(tpData);
+            }
+            tpData.partitionData().add(new ProduceRequestData.PartitionProduceData()
+                    .setIndex(tp.partition())
+                    .setRecords(records));
             recordsByPartition.put(tp, batch);
         }
 
@@ -685,16 +783,18 @@ private void sendProduceRequest(long now, int destination, short acks, int timeo
         if (transactionManager != null && transactionManager.isTransactional()) {
             transactionalId = transactionManager.transactionalId();
         }
-        ProduceRequest.Builder requestBuilder = ProduceRequest.Builder.forMagic(minUsedMagic, acks, timeout,
-                produceRecordsByPartition, transactionalId);
-        RequestCompletionHandler callback = new RequestCompletionHandler() {
-            public void onComplete(ClientResponse response) {
-                handleProduceResponse(response, recordsByPartition, time.milliseconds());
-            }
-        };
+
+        ProduceRequest.Builder requestBuilder = ProduceRequest.forMagic(minUsedMagic,
+                new ProduceRequestData()
+                        .setAcks(acks)
+                        .setTimeoutMs(timeout)
+                        .setTransactionalId(transactionalId)
+                        .setTopicData(tpd));
+        RequestCompletionHandler callback = response -> handleProduceResponse(response, recordsByPartition, time.milliseconds());
 
         String nodeId = Integer.toString(destination);
-        ClientRequest clientRequest = client.newClientRequest(nodeId, requestBuilder, now, acks != 0, callback);
+        ClientRequest clientRequest = client.newClientRequest(nodeId, requestBuilder, now, acks != 0,
+                requestTimeoutMs, callback);
         client.send(clientRequest, now);
         log.trace("Sent produce request to {}: {}", nodeId, requestBuilder);
     }
@@ -716,7 +816,7 @@ public static Sensor throttleTimeSensor(SenderMetricsRegistry metrics) {
     /**
      * A collection of sensors for the sender
      */
-    private class SenderMetrics {
+    private static class SenderMetrics {
         public final Sensor retrySensor;
         public final Sensor errorSensor;
         public final Sensor queueTimeSensor;
@@ -727,9 +827,11 @@ private class SenderMetrics {
         public final Sensor maxRecordSizeSensor;
         public final Sensor batchSplitSensor;
         private final SenderMetricsRegistry metrics;
+        private final Time time;
 
-        public SenderMetrics(SenderMetricsRegistry metrics) {
+        public SenderMetrics(SenderMetricsRegistry metrics, Metadata metadata, KafkaClient client, Time time) {
             this.metrics = metrics;
+            this.time = time;
 
             this.batchSizeSensor = metrics.sensor("batch-size");
             this.batchSizeSensor.add(metrics.batchSizeAvg, new Avg());
@@ -760,16 +862,9 @@ public SenderMetrics(SenderMetricsRegistry metrics) {
             this.maxRecordSizeSensor.add(metrics.recordSizeMax, new Max());
             this.maxRecordSizeSensor.add(metrics.recordSizeAvg, new Avg());
 
-            this.metrics.addMetric(metrics.requestsInFlight, new Measurable() {
-                public double measure(MetricConfig config, long now) {
-                    return client.inFlightRequestCount();
-                }
-            });
-            metrics.addMetric(metrics.metadataAge, new Measurable() {
-                public double measure(MetricConfig config, long now) {
-                    return (now - metadata.lastSuccessfulUpdate()) / 1000.0;
-                }
-            });
+            this.metrics.addMetric(metrics.requestsInFlight, (config, now) -> client.inFlightRequestCount());
+            this.metrics.addMetric(metrics.metadataAge,
+                (config, now) -> (now - metadata.lastSuccessfulUpdate()) / 1000.0);
 
             this.batchSplitSensor = metrics.sensor("batch-split-rate");
             this.batchSplitSensor.add(new Meter(metrics.batchSplitRate, metrics.batchSplitTotal));
@@ -824,17 +919,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/SenderMetricsRegistry.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/SenderMetricsRegistry.java
index 21dbca618303d..643897375a726 100644
--- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/SenderMetricsRegistry.java
+++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/SenderMetricsRegistry.java
@@ -17,7 +17,7 @@
 package org.apache.kafka.clients.producer.internals;
 
 import java.util.ArrayList;
-import java.util.HashSet;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -70,12 +70,12 @@ public class SenderMetricsRegistry {
     
     private final Metrics metrics;
     private final Set tags;
-    private final HashSet topicTags;
+    private final LinkedHashSet topicTags;
 
     public SenderMetricsRegistry(Metrics metrics) {
         this.metrics = metrics;
         this.tags = this.metrics.config().tags().keySet();
-        this.allTemplates = new ArrayList();
+        this.allTemplates = new ArrayList<>();
         
         /***** Client level *****/
         
@@ -84,7 +84,8 @@ public SenderMetricsRegistry(Metrics metrics) {
         this.batchSizeMax = createMetricName("batch-size-max",
                 "The max number of bytes sent per partition per-request.");
         this.compressionRateAvg = createMetricName("compression-rate-avg",
-                "The average compression rate of record batches.");
+                "The average compression rate of record batches, defined as the average ratio of the " +
+                        "compressed batch size over the uncompressed size.");
         this.recordQueueTimeAvg = createMetricName("record-queue-time-avg",
                 "The average time in ms record batches spent in the send buffer.");
         this.recordQueueTimeMax = createMetricName("record-queue-time-max",
@@ -126,7 +127,7 @@ public SenderMetricsRegistry(Metrics metrics) {
                 "The maximum time in ms a request was throttled by a broker");
 
         /***** Topic level *****/
-        this.topicTags = new HashSet(tags);
+        this.topicTags = new LinkedHashSet<>(tags);
         this.topicTags.add("topic");
 
         // We can't create the MetricName up front for these, because we don't know the topic name yet.
@@ -139,7 +140,8 @@ public SenderMetricsRegistry(Metrics metrics) {
         this.topicByteTotal = createTopicTemplate("byte-total", 
                 "The total number of bytes sent for a topic.");
         this.topicCompressionRate = createTopicTemplate("compression-rate",
-                "The average compression rate of record batches for a topic.");
+                "The average compression rate of record batches for a topic, defined as the average ratio " +
+                        "of the compressed batch size over the uncompressed size.");
         this.topicRecordRetryRate = createTopicTemplate("record-retry-rate",
                 "The average per-second number of retried record sends for a topic");
         this.topicRecordRetryTotal = createTopicTemplate("record-retry-total",
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..b432009261b34
--- /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)) {
+                    int 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 006a12b1bfd4c..e58469b003a7e 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
@@ -16,16 +16,37 @@
  */
 package org.apache.kafka.clients.producer.internals;
 
+import org.apache.kafka.clients.ApiVersion;
+import org.apache.kafka.clients.ApiVersions;
 import org.apache.kafka.clients.ClientResponse;
+import org.apache.kafka.clients.NodeApiVersions;
 import org.apache.kafka.clients.RequestCompletionHandler;
+import org.apache.kafka.clients.consumer.CommitFailedException;
+import org.apache.kafka.clients.consumer.ConsumerGroupMetadata;
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
+import org.apache.kafka.common.errors.InvalidPidMappingException;
+import org.apache.kafka.common.errors.InvalidProducerEpochException;
+import org.apache.kafka.common.errors.RetriableException;
+import org.apache.kafka.common.errors.UnknownProducerIdException;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.utils.ProducerIdAndEpoch;
 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.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.AddOffsetsToTxnRequestData;
+import org.apache.kafka.common.message.EndTxnRequestData;
+import org.apache.kafka.common.message.FindCoordinatorRequestData;
+import org.apache.kafka.common.message.InitProducerIdRequestData;
 import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.record.DefaultRecordBatch;
 import org.apache.kafka.common.record.RecordBatch;
 import org.apache.kafka.common.requests.AbstractRequest;
 import org.apache.kafka.common.requests.AbstractResponse;
@@ -36,6 +57,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;
@@ -45,35 +67,129 @@
 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;
 import java.util.Comparator;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.Iterator;
+import java.util.Locale;
 import java.util.Map;
+import java.util.OptionalInt;
+import java.util.OptionalLong;
 import java.util.PriorityQueue;
 import java.util.Set;
-
-import static org.apache.kafka.common.record.RecordBatch.NO_PRODUCER_EPOCH;
-import static org.apache.kafka.common.record.RecordBatch.NO_PRODUCER_ID;
+import java.util.SortedSet;
+import java.util.TreeSet;
+import java.util.function.Consumer;
+import java.util.function.Supplier;
 
 /**
  * A class which maintains state for transactions. Also keeps the state necessary to ensure idempotent production.
  */
 public class TransactionManager {
     private static final int NO_INFLIGHT_REQUEST_CORRELATION_ID = -1;
+    private static final int NO_LAST_ACKED_SEQUENCE_NUMBER = -1;
 
     private final Logger log;
     private final String transactionalId;
     private final int transactionTimeoutMs;
+    private final ApiVersions apiVersions;
+    private final boolean autoDowngradeTxnCommit;
+
+    private static class TopicPartitionBookkeeper {
+
+        private final Map topicPartitions = new HashMap<>();
+
+        private TopicPartitionEntry getPartition(TopicPartition topicPartition) {
+            TopicPartitionEntry ent = topicPartitions.get(topicPartition);
+            if (ent == null)
+                throw new IllegalStateException("Trying to get the sequence number for " + topicPartition +
+                        ", but the sequence number was never set for this partition.");
+            return ent;
+        }
+
+        private void addPartition(TopicPartition topicPartition) {
+            this.topicPartitions.putIfAbsent(topicPartition, new TopicPartitionEntry());
+        }
+
+        private boolean contains(TopicPartition topicPartition) {
+            return topicPartitions.containsKey(topicPartition);
+        }
+
+        private void reset() {
+            topicPartitions.clear();
+        }
+
+        private OptionalLong lastAckedOffset(TopicPartition topicPartition) {
+            TopicPartitionEntry entry = topicPartitions.get(topicPartition);
+            if (entry != null && entry.lastAckedOffset != ProduceResponse.INVALID_OFFSET)
+                return OptionalLong.of(entry.lastAckedOffset);
+            else
+                return OptionalLong.empty();
+        }
+
+        private OptionalInt lastAckedSequence(TopicPartition topicPartition) {
+            TopicPartitionEntry entry = topicPartitions.get(topicPartition);
+            if (entry != null && entry.lastAckedSequence != NO_LAST_ACKED_SEQUENCE_NUMBER)
+                return OptionalInt.of(entry.lastAckedSequence);
+            else
+                return OptionalInt.empty();
+        }
+
+        private void startSequencesAtBeginning(TopicPartition topicPartition, ProducerIdAndEpoch newProducerIdAndEpoch) {
+            final PrimitiveRef.IntRef sequence = PrimitiveRef.ofInt(0);
+            TopicPartitionEntry topicPartitionEntry = getPartition(topicPartition);
+            topicPartitionEntry.resetSequenceNumbers(inFlightBatch -> {
+                inFlightBatch.resetProducerState(newProducerIdAndEpoch, sequence.value, inFlightBatch.isTransactional());
+                sequence.value += inFlightBatch.recordCount;
+            });
+            topicPartitionEntry.nextSequence = sequence.value;
+            topicPartitionEntry.lastAckedSequence = NO_LAST_ACKED_SEQUENCE_NUMBER;
+        }
+    }
+
+    private static class TopicPartitionEntry {
+
+        // The base sequence of the next batch bound for a given partition.
+        private int nextSequence;
+
+        // The sequence number of the last record of the last ack'd batch from the given partition. When there are no
+        // in flight requests for a partition, the lastAckedSequence(topicPartition) == nextSequence(topicPartition) - 1.
+        private int lastAckedSequence;
+
+        // Keep track of the in flight batches bound for a partition, ordered by sequence. This helps us to ensure that
+        // 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 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.
+        private long lastAckedOffset;
+
+        TopicPartitionEntry() {
+            this.nextSequence = 0;
+            this.lastAckedSequence = NO_LAST_ACKED_SEQUENCE_NUMBER;
+            this.lastAckedOffset = ProduceResponse.INVALID_OFFSET;
+            this.inflightBatchesBySequence = new TreeSet<>(Comparator.comparingInt(ProducerBatch::baseSequence));
+        }
+
+        void resetSequenceNumbers(Consumer resetSequence) {
+            TreeSet newInflights = new TreeSet<>(Comparator.comparingInt(ProducerBatch::baseSequence));
+            for (ProducerBatch inflightBatch : inflightBatchesBySequence) {
+                resetSequence.accept(inflightBatch);
+                newInflights.add(inflightBatch);
+            }
+            inflightBatchesBySequence = newInflights;
+        }
+    }
 
-    // The base sequence of the next batch bound for a given partition.
-    private final Map nextSequence;
+    private final TopicPartitionBookkeeper topicPartitionBookkeeper;
 
-    // The sequence of the last record of the last ack'd batch from the given partition. When there are no
-    // in flight requests for a partition, the lastAckedSequence(topicPartition) == nextSequence(topicPartition) - 1.
-    private final Map lastAckedSequence;
+    private final Map pendingTxnOffsetCommits;
 
     // If a batch bound for a partition expired locally after being sent at least once, the partition has is considered
     // to have an unresolved state. We keep track fo such partitions here, and cannot assign any more sequence numbers
@@ -81,23 +197,20 @@ public class TransactionManager {
     // successfully (indicating that the expired batch actually made it to the broker). If we don't get any successful
     // responses for the partition once the inflight request count falls to zero, we reset the producer id and
     // consequently clear this data structure as well.
-    private final Set partitionsWithUnresolvedSequences;
+    // The value of the map is the sequence number of the batch following the expired one, computed by adding its
+    // record count to its sequence number. This is used to tell if a subsequent batch is the one immediately following
+    // the expired one.
+    private final Map partitionsWithUnresolvedSequences;
 
-    // Keep track of the in flight batches bound for a partition, ordered by sequence. This helps us to ensure that
-    // 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 final Map> 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.
-    private final Map lastAckedOffset;
+    // The partitions that have received an error that triggers an epoch bump. When the epoch is bumped, these
+    // partitions will have the sequences of their in-flight batches rewritten
+    private final Set partitionsToRewriteSequences;
 
     private final PriorityQueue pendingRequests;
     private final Set newPartitionsInTransaction;
     private final Set pendingPartitionsInTransaction;
     private final Set partitionsInTransaction;
-    private final Map pendingTxnOffsetCommits;
+    private TransactionalRequestResult pendingResult;
 
     // This is used by the TxnRequestHandlers to control how long to back off before a given request is retried.
     // For instance, this value is lowered by the AddPartitionsToTxnHandler when it receives a CONCURRENT_TRANSACTIONS
@@ -111,11 +224,13 @@ public class TransactionManager {
     private int inFlightRequestCorrelationId = NO_INFLIGHT_REQUEST_CORRELATION_ID;
     private Node transactionCoordinator;
     private Node consumerGroupCoordinator;
+    private boolean coordinatorSupportsBumpingEpoch;
 
     private volatile State currentState = State.UNINITIALIZED;
     private volatile RuntimeException lastError = null;
     private volatile ProducerIdAndEpoch producerIdAndEpoch;
     private volatile boolean transactionStarted = false;
+    private volatile boolean epochBumpRequired = false;
 
     private enum State {
         UNINITIALIZED,
@@ -129,8 +244,10 @@ private enum State {
 
         private boolean isTransitionValid(State source, State target) {
             switch (target) {
+                case UNINITIALIZED:
+                    return source == READY;
                 case INITIALIZING:
-                    return source == UNINITIALIZED;
+                    return source == UNINITIALIZED || source == ABORTING_TRANSACTION;
                 case READY:
                     return source == INITIALIZING || source == COMMITTING_TRANSACTION || source == ABORTING_TRANSACTION;
                 case IN_TRANSACTION:
@@ -153,12 +270,14 @@ private boolean isTransitionValid(State source, State target) {
 
     // We use the priority to determine the order in which requests need to be sent out. For instance, if we have
     // a pending FindCoordinator request, that must always go first. Next, If we need a producer id, that must go second.
-    // The endTxn request must always go last.
+    // The endTxn request must always go last, unless we are bumping the epoch (a special case of InitProducerId) as
+    // part of ending the transaction.
     private enum Priority {
         FIND_COORDINATOR(0),
         INIT_PRODUCER_ID(1),
         ADD_PARTITIONS_OR_OFFSETS(2),
-        END_TXN(3);
+        END_TXN(3),
+        EPOCH_BUMP(4);
 
         final int priority;
 
@@ -167,10 +286,13 @@ private enum Priority {
         }
     }
 
-    public TransactionManager(LogContext logContext, String transactionalId, int transactionTimeoutMs, long retryBackoffMs) {
-        this.producerIdAndEpoch = new ProducerIdAndEpoch(NO_PRODUCER_ID, NO_PRODUCER_EPOCH);
-        this.nextSequence = new HashMap<>();
-        this.lastAckedSequence = new HashMap<>();
+    public TransactionManager(final LogContext logContext,
+                              final String transactionalId,
+                              final int transactionTimeoutMs,
+                              final long retryBackoffMs,
+                              final ApiVersions apiVersions,
+                              final boolean autoDowngradeTxnCommit) {
+        this.producerIdAndEpoch = ProducerIdAndEpoch.NONE;
         this.transactionalId = transactionalId;
         this.log = logContext.logger(TransactionManager.class);
         this.transactionTimeoutMs = transactionTimeoutMs;
@@ -179,34 +301,40 @@ public TransactionManager(LogContext logContext, String transactionalId, int tra
         this.newPartitionsInTransaction = new HashSet<>();
         this.pendingPartitionsInTransaction = new HashSet<>();
         this.partitionsInTransaction = new HashSet<>();
+        this.pendingRequests = new PriorityQueue<>(10, Comparator.comparingInt(o -> o.priority().priority));
         this.pendingTxnOffsetCommits = new HashMap<>();
-        this.pendingRequests = new PriorityQueue<>(10, new Comparator() {
-            @Override
-            public int compare(TxnRequestHandler o1, TxnRequestHandler o2) {
-                return Integer.compare(o1.priority().priority, o2.priority().priority);
-            }
-        });
-
-        this.partitionsWithUnresolvedSequences = new HashSet<>();
-        this.inflightBatchesBySequence = new HashMap<>();
-        this.lastAckedOffset = new HashMap<>();
-
+        this.partitionsWithUnresolvedSequences = new HashMap<>();
+        this.partitionsToRewriteSequences = new HashSet<>();
         this.retryBackoffMs = retryBackoffMs;
+        this.topicPartitionBookkeeper = new TopicPartitionBookkeeper();
+        this.apiVersions = apiVersions;
+        this.autoDowngradeTxnCommit = autoDowngradeTxnCommit;
     }
 
-    TransactionManager() {
-        this(new LogContext(), null, 0, 100);
+    public synchronized TransactionalRequestResult initializeTransactions() {
+        return initializeTransactions(ProducerIdAndEpoch.NONE);
     }
 
-    public synchronized TransactionalRequestResult initializeTransactions() {
-        ensureTransactional();
-        transitionTo(State.INITIALIZING);
-        setProducerIdAndEpoch(ProducerIdAndEpoch.NONE);
-        this.nextSequence.clear();
-        InitProducerIdRequest.Builder builder = new InitProducerIdRequest.Builder(transactionalId, transactionTimeoutMs);
-        InitProducerIdHandler handler = new InitProducerIdHandler(builder);
-        enqueueRequest(handler);
-        return handler.result;
+    synchronized TransactionalRequestResult initializeTransactions(ProducerIdAndEpoch producerIdAndEpoch) {
+        boolean isEpochBump = producerIdAndEpoch != ProducerIdAndEpoch.NONE;
+        return handleCachedTransactionRequestResult(() -> {
+            // If this is an epoch bump, we will transition the state as part of handling the EndTxnRequest
+            if (!isEpochBump) {
+                transitionTo(State.INITIALIZING);
+                log.info("Invoking InitProducerId for the first time in order to acquire a producer ID");
+            } else {
+                log.info("Invoking InitProducerId with current producer ID and epoch {} in order to bump the epoch", producerIdAndEpoch);
+            }
+            InitProducerIdRequestData requestData = new InitProducerIdRequestData()
+                    .setTransactionalId(transactionalId)
+                    .setTransactionTimeoutMs(transactionTimeoutMs)
+                    .setProducerId(producerIdAndEpoch.producerId)
+                    .setProducerEpoch(producerIdAndEpoch.epoch);
+            InitProducerIdHandler handler = new InitProducerIdHandler(new InitProducerIdRequest.Builder(requestData),
+                    isEpochBump);
+            enqueueRequest(handler);
+            return handler.result;
+        }, State.INITIALIZING);
     }
 
     public synchronized void beginTransaction() {
@@ -216,56 +344,78 @@ public synchronized void beginTransaction() {
     }
 
     public synchronized TransactionalRequestResult beginCommit() {
-        ensureTransactional();
-        maybeFailWithError();
-        transitionTo(State.COMMITTING_TRANSACTION);
-        return beginCompletingTransaction(TransactionResult.COMMIT);
+        return handleCachedTransactionRequestResult(() -> {
+            maybeFailWithError();
+            transitionTo(State.COMMITTING_TRANSACTION);
+            return beginCompletingTransaction(TransactionResult.COMMIT);
+        }, State.COMMITTING_TRANSACTION);
     }
 
     public synchronized TransactionalRequestResult beginAbort() {
-        ensureTransactional();
-        if (currentState != State.ABORTABLE_ERROR)
-            maybeFailWithError();
-        transitionTo(State.ABORTING_TRANSACTION);
+        return handleCachedTransactionRequestResult(() -> {
+            if (currentState != State.ABORTABLE_ERROR)
+                maybeFailWithError();
+            transitionTo(State.ABORTING_TRANSACTION);
 
-        // We're aborting the transaction, so there should be no need to add new partitions
-        newPartitionsInTransaction.clear();
-        return beginCompletingTransaction(TransactionResult.ABORT);
+            // We're aborting the transaction, so there should be no need to add new partitions
+            newPartitionsInTransaction.clear();
+            return beginCompletingTransaction(TransactionResult.ABORT);
+        }, State.ABORTING_TRANSACTION);
     }
 
     private TransactionalRequestResult beginCompletingTransaction(TransactionResult transactionResult) {
         if (!newPartitionsInTransaction.isEmpty())
             enqueueRequest(addPartitionsToTransactionHandler());
-        EndTxnRequest.Builder builder = new EndTxnRequest.Builder(transactionalId, producerIdAndEpoch.producerId,
-                producerIdAndEpoch.epoch, transactionResult);
-        EndTxnHandler handler = new EndTxnHandler(builder);
-        enqueueRequest(handler);
-        return handler.result;
+
+        // If the error is an INVALID_PRODUCER_ID_MAPPING error, the server will not accept an EndTxnRequest, so skip
+        // directly to InitProducerId. Otherwise, we must first abort the transaction, because the producer will be
+        // fenced if we directly call InitProducerId.
+        if (!(lastError instanceof InvalidPidMappingException)) {
+            EndTxnRequest.Builder builder = new EndTxnRequest.Builder(
+                    new EndTxnRequestData()
+                            .setTransactionalId(transactionalId)
+                            .setProducerId(producerIdAndEpoch.producerId)
+                            .setProducerEpoch(producerIdAndEpoch.epoch)
+                            .setCommitted(transactionResult.id));
+
+            EndTxnHandler handler = new EndTxnHandler(builder);
+            enqueueRequest(handler);
+            if (!epochBumpRequired) {
+                return handler.result;
+            }
+        }
+
+        return initializeTransactions(this.producerIdAndEpoch);
     }
 
-    public synchronized TransactionalRequestResult sendOffsetsToTransaction(Map offsets,
-                                                                            String consumerGroupId) {
+    public synchronized TransactionalRequestResult sendOffsetsToTransaction(final Map offsets,
+                                                                            final ConsumerGroupMetadata groupMetadata) {
         ensureTransactional();
         maybeFailWithError();
         if (currentState != State.IN_TRANSACTION)
             throw new KafkaException("Cannot send offsets to transaction either because the producer is not in an " +
                     "active transaction");
 
-        log.debug("Begin adding offsets {} for consumer group {} to transaction", offsets, consumerGroupId);
-        AddOffsetsToTxnRequest.Builder builder = new AddOffsetsToTxnRequest.Builder(transactionalId,
-                producerIdAndEpoch.producerId, producerIdAndEpoch.epoch, consumerGroupId);
-        AddOffsetsToTxnHandler handler = new AddOffsetsToTxnHandler(builder, offsets);
+        log.debug("Begin adding offsets {} for consumer group {} to transaction", offsets, groupMetadata);
+        AddOffsetsToTxnRequest.Builder builder = new AddOffsetsToTxnRequest.Builder(
+            new AddOffsetsToTxnRequestData()
+                .setTransactionalId(transactionalId)
+                .setProducerId(producerIdAndEpoch.producerId)
+                .setProducerEpoch(producerIdAndEpoch.epoch)
+                .setGroupId(groupMetadata.groupId())
+        );
+        AddOffsetsToTxnHandler handler = new AddOffsetsToTxnHandler(builder, offsets, groupMetadata);
+
         enqueueRequest(handler);
         return handler.result;
     }
 
     public synchronized void maybeAddPartitionToTransaction(TopicPartition topicPartition) {
-        failIfNotReadyForSend();
-
         if (isPartitionAdded(topicPartition) || isPartitionPendingAdd(topicPartition))
             return;
 
         log.debug("Begin adding new partition {} to transaction", topicPartition);
+        topicPartitionBookkeeper.addPartition(topicPartition);
         newPartitionsInTransaction.add(topicPartition);
     }
 
@@ -328,11 +478,18 @@ synchronized void transitionToAbortableError(RuntimeException exception) {
                     "aborted. Underlying exception: ", exception);
             return;
         }
+
+        log.info("Transiting to abortable error state due to {}", exception.toString());
         transitionTo(State.ABORTABLE_ERROR, exception);
     }
 
     synchronized void transitionToFatalError(RuntimeException exception) {
+        log.info("Transiting to fatal error state due to {}", exception.toString());
         transitionTo(State.FATAL_ERROR, exception);
+
+        if (pendingResult != null) {
+            pendingResult.fail(exception);
+        }
     }
 
     // visible for testing
@@ -355,85 +512,103 @@ ProducerIdAndEpoch producerIdAndEpoch() {
         return producerIdAndEpoch;
     }
 
-    boolean hasProducerId(long producerId) {
-        return producerIdAndEpoch.producerId == producerId;
-    }
-
-    boolean hasProducerIdAndEpoch(long producerId, short producerEpoch) {
+    boolean producerIdOrEpochNotMatch(ProducerBatch batch) {
         ProducerIdAndEpoch idAndEpoch = this.producerIdAndEpoch;
-        return idAndEpoch.producerId == producerId && idAndEpoch.epoch == producerEpoch;
+        return idAndEpoch.producerId != batch.producerId() || idAndEpoch.epoch != batch.producerEpoch();
     }
 
     /**
      * Set the producer id and epoch atomically.
      */
-    void setProducerIdAndEpoch(ProducerIdAndEpoch producerIdAndEpoch) {
+    private void setProducerIdAndEpoch(ProducerIdAndEpoch producerIdAndEpoch) {
         log.info("ProducerId set to {} with epoch {}", producerIdAndEpoch.producerId, producerIdAndEpoch.epoch);
         this.producerIdAndEpoch = producerIdAndEpoch;
     }
 
     /**
-     * This method is used when the producer needs to reset its internal state because of an irrecoverable exception
-     * from the broker.
-     *
-     * We need to reset the producer id and associated state when we have sent a batch to the broker, but we either get
-     * a non-retriable exception or we run out of retries, or the batch expired in the producer queue after it was already
-     * sent to the broker.
-     *
-     * In all of these cases, we don't know whether batch was actually committed on the broker, and hence whether the
-     * sequence number was actually updated. If we don't reset the producer state, we risk the chance that all future
-     * messages will return an OutOfOrderSequenceException.
-     *
-     * Note that we can't reset the producer state for the transactional producer as this would mean bumping the epoch
-     * for the same producer id. This might involve aborting the ongoing transaction during the initPidRequest, and the user
-     * would not have any way of knowing this happened. So for the transactional producer, it's best to return the
-     * produce error to the user and let them abort the transaction and close the producer explicitly.
+     * This method resets the producer ID and epoch and sets the state to UNINITIALIZED, which will trigger a new
+     * InitProducerId request. This method is only called when the producer epoch is exhausted; we will bump the epoch
+     * instead.
      */
-    synchronized void resetProducerId() {
+    private void resetIdempotentProducerId() {
         if (isTransactional())
             throw new IllegalStateException("Cannot reset producer state for a transactional producer. " +
                     "You must either abort the ongoing transaction or reinitialize the transactional producer instead");
+        log.debug("Resetting idempotent producer ID. ID and epoch before reset are {}", this.producerIdAndEpoch);
         setProducerIdAndEpoch(ProducerIdAndEpoch.NONE);
-        this.nextSequence.clear();
-        this.lastAckedSequence.clear();
-        this.inflightBatchesBySequence.clear();
+        transitionTo(State.UNINITIALIZED);
+    }
+
+    private void resetSequenceForPartition(TopicPartition topicPartition) {
+        topicPartitionBookkeeper.topicPartitions.remove(topicPartition);
+        this.partitionsWithUnresolvedSequences.remove(topicPartition);
+    }
+
+    private void resetSequenceNumbers() {
+        topicPartitionBookkeeper.reset();
         this.partitionsWithUnresolvedSequences.clear();
-        this.lastAckedOffset.clear();
+    }
+
+    synchronized void requestEpochBumpForPartition(TopicPartition tp) {
+        epochBumpRequired = true;
+        this.partitionsToRewriteSequences.add(tp);
+    }
+
+    private void bumpIdempotentProducerEpoch() {
+        if (this.producerIdAndEpoch.epoch == Short.MAX_VALUE) {
+            resetIdempotentProducerId();
+        } else {
+            setProducerIdAndEpoch(new ProducerIdAndEpoch(this.producerIdAndEpoch.producerId, (short) (this.producerIdAndEpoch.epoch + 1)));
+            log.debug("Incremented producer epoch, current producer ID and epoch are now {}", this.producerIdAndEpoch);
+        }
+
+        // When the epoch is bumped, rewrite all in-flight sequences for the partition(s) that triggered the epoch bump
+        for (TopicPartition topicPartition : this.partitionsToRewriteSequences) {
+            this.topicPartitionBookkeeper.startSequencesAtBeginning(topicPartition, this.producerIdAndEpoch);
+            this.partitionsWithUnresolvedSequences.remove(topicPartition);
+        }
+
+        this.partitionsToRewriteSequences.clear();
+        epochBumpRequired = false;
+    }
+
+    synchronized void bumpIdempotentEpochAndResetIdIfNeeded() {
+        if (!isTransactional()) {
+            if (epochBumpRequired) {
+                bumpIdempotentProducerEpoch();
+            }
+            if (currentState != State.INITIALIZING && !hasProducerId()) {
+                transitionTo(State.INITIALIZING);
+                InitProducerIdRequestData requestData = new InitProducerIdRequestData()
+                        .setTransactionalId(null)
+                        .setTransactionTimeoutMs(Integer.MAX_VALUE);
+                InitProducerIdHandler handler = new InitProducerIdHandler(new InitProducerIdRequest.Builder(requestData), false);
+                enqueueRequest(handler);
+            }
+        }
     }
 
     /**
      * Returns the next sequence number to be written to the given TopicPartition.
      */
     synchronized Integer sequenceNumber(TopicPartition topicPartition) {
-        Integer currentSequenceNumber = nextSequence.get(topicPartition);
-        if (currentSequenceNumber == null) {
-            currentSequenceNumber = 0;
-            nextSequence.put(topicPartition, currentSequenceNumber);
-        }
-        return currentSequenceNumber;
+        if (!isTransactional())
+            topicPartitionBookkeeper.addPartition(topicPartition);
+
+        return topicPartitionBookkeeper.getPartition(topicPartition).nextSequence;
     }
 
     synchronized void incrementSequenceNumber(TopicPartition topicPartition, int increment) {
-        Integer currentSequenceNumber = nextSequence.get(topicPartition);
-        if (currentSequenceNumber == null)
-            throw new IllegalStateException("Attempt to increment sequence number for a partition with no current sequence.");
+        Integer currentSequence = sequenceNumber(topicPartition);
 
-        currentSequenceNumber += increment;
-        nextSequence.put(topicPartition, currentSequenceNumber);
+        currentSequence = DefaultRecordBatch.incrementSequence(currentSequence, increment);
+        topicPartitionBookkeeper.getPartition(topicPartition).nextSequence = currentSequence;
     }
 
     synchronized void addInFlightBatch(ProducerBatch batch) {
         if (!batch.hasSequence())
             throw new IllegalStateException("Can't track batch for partition " + batch.topicPartition + " when sequence is not set.");
-        if (!inflightBatchesBySequence.containsKey(batch.topicPartition)) {
-            inflightBatchesBySequence.put(batch.topicPartition, new PriorityQueue<>(5, new Comparator() {
-                @Override
-                public int compare(ProducerBatch o1, ProducerBatch o2) {
-                    return o1.baseSequence() - o2.baseSequence();
-                }
-            }));
-        }
-        inflightBatchesBySequence.get(batch.topicPartition).offer(batch);
+        topicPartitionBookkeeper.getPartition(batch.topicPartition).inflightBatchesBySequence.add(batch);
     }
 
     /**
@@ -444,68 +619,144 @@ public int compare(ProducerBatch o1, ProducerBatch o2) {
      *         RecordBatch.NO_SEQUENCE.
      */
     synchronized int firstInFlightSequence(TopicPartition topicPartition) {
-        PriorityQueue inFlightBatches = inflightBatchesBySequence.get(topicPartition);
-        if (inFlightBatches == null)
+        if (!hasInflightBatches(topicPartition))
             return RecordBatch.NO_SEQUENCE;
 
-        ProducerBatch firstInFlightBatch = inFlightBatches.peek();
-        if (firstInFlightBatch == null)
+        SortedSet inflightBatches = topicPartitionBookkeeper.getPartition(topicPartition).inflightBatchesBySequence;
+        if (inflightBatches.isEmpty())
             return RecordBatch.NO_SEQUENCE;
-
-        return firstInFlightBatch.baseSequence();
+        else
+            return inflightBatches.first().baseSequence();
     }
 
     synchronized ProducerBatch nextBatchBySequence(TopicPartition topicPartition) {
-        PriorityQueue queue = inflightBatchesBySequence.get(topicPartition);
-        if (queue == null)
-            return null;
-        return queue.peek();
+        SortedSet queue = topicPartitionBookkeeper.getPartition(topicPartition).inflightBatchesBySequence;
+        return queue.isEmpty() ? null : queue.first();
     }
 
     synchronized void removeInFlightBatch(ProducerBatch batch) {
-        PriorityQueue queue = inflightBatchesBySequence.get(batch.topicPartition);
-        if (queue == null)
-            return;
-        queue.remove(batch);
+        if (hasInflightBatches(batch.topicPartition)) {
+            topicPartitionBookkeeper.getPartition(batch.topicPartition).inflightBatchesBySequence.remove(batch);
+        }
     }
 
-    synchronized void maybeUpdateLastAckedSequence(TopicPartition topicPartition, int sequence) {
-        if (sequence > lastAckedSequence(topicPartition))
-            lastAckedSequence.put(topicPartition, sequence);
+    private int maybeUpdateLastAckedSequence(TopicPartition topicPartition, int sequence) {
+        int lastAckedSequence = lastAckedSequence(topicPartition).orElse(NO_LAST_ACKED_SEQUENCE_NUMBER);
+        if (sequence > lastAckedSequence) {
+            topicPartitionBookkeeper.getPartition(topicPartition).lastAckedSequence = sequence;
+            return sequence;
+        }
+
+        return lastAckedSequence;
     }
 
-    synchronized int lastAckedSequence(TopicPartition topicPartition) {
-        Integer currentLastAckedSequence = lastAckedSequence.get(topicPartition);
-        if (currentLastAckedSequence == null)
-            return -1;
-        return currentLastAckedSequence;
+    synchronized OptionalInt lastAckedSequence(TopicPartition topicPartition) {
+        return topicPartitionBookkeeper.lastAckedSequence(topicPartition);
     }
 
-    synchronized long lastAckedOffset(TopicPartition topicPartition) {
-        Long offset = lastAckedOffset.get(topicPartition);
-        if (offset == null)
-            return ProduceResponse.INVALID_OFFSET;
-        return offset;
+    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;
-        if (lastOffset > lastAckedOffset(batch.topicPartition)) {
-            lastAckedOffset.put(batch.topicPartition, lastOffset);
+        OptionalLong lastAckedOffset = lastAckedOffset(batch.topicPartition);
+        // It might happen that the TransactionManager has been reset while a request was reenqueued and got a valid
+        // response for this. This can happen only if the producer is only idempotent (not transactional) and in
+        // this case there will be no tracked bookkeeper entry about it, so we have to insert one.
+        if (!lastAckedOffset.isPresent() && !isTransactional()) {
+            topicPartitionBookkeeper.addPartition(batch.topicPartition);
+        }
+        if (lastOffset > lastAckedOffset.orElse(ProduceResponse.INVALID_OFFSET)) {
+            topicPartitionBookkeeper.getPartition(batch.topicPartition).lastAckedOffset = lastOffset;
         } else {
             log.trace("Partition {} keeps lastOffset at {}", batch.topicPartition, lastOffset);
         }
     }
 
+    public synchronized void handleCompletedBatch(ProducerBatch batch, ProduceResponse.PartitionResponse response) {
+        int lastAckedSequence = maybeUpdateLastAckedSequence(batch.topicPartition, batch.lastSequence());
+        log.debug("ProducerId: {}; Set last ack'd sequence number for topic-partition {} to {}",
+                batch.producerId(),
+                batch.topicPartition,
+                lastAckedSequence);
+
+        updateLastAckedOffset(response, batch);
+        removeInFlightBatch(batch);
+
+        if (producerIdOrEpochNotMatch(batch) && !hasInflightBatches(batch.topicPartition)) {
+            // If the batch was on a different ID and/or epoch (due to an epoch bump) and all its in-flight batches
+            // have completed, reset the partition sequence so that the next batch (with the new epoch) starts from 0
+            topicPartitionBookkeeper.startSequencesAtBeginning(batch.topicPartition, this.producerIdAndEpoch);
+        }
+    }
+
+    private void maybeTransitionToErrorState(RuntimeException exception) {
+        if (exception instanceof ClusterAuthorizationException
+                || exception instanceof TransactionalIdAuthorizationException
+                || exception instanceof ProducerFencedException
+                || exception instanceof UnsupportedVersionException) {
+            transitionToFatalError(exception);
+        } else if (isTransactional()) {
+            if (canBumpEpoch() && !isCompleting()) {
+                epochBumpRequired = true;
+            }
+            transitionToAbortableError(exception);
+        }
+    }
+
+    synchronized void handleFailedBatch(ProducerBatch batch, RuntimeException exception, boolean adjustSequenceNumbers) {
+        maybeTransitionToErrorState(exception);
+        removeInFlightBatch(batch);
+
+        if (hasFatalError()) {
+            log.debug("Ignoring batch {} with producer id {}, epoch {}, and sequence number {} " +
+                            "since the producer is already in fatal error state", batch, batch.producerId(),
+                    batch.producerEpoch(), batch.baseSequence(), exception);
+            return;
+        }
+
+        if (producerIdOrEpochNotMatch(batch)) {
+            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());
+
+            // If we fail with an OutOfOrderSequenceException, we have a gap in the log. Bump the epoch for this
+            // partition, which will reset the sequence number to 0 and allow us to continue
+            requestEpochBumpForPartition(batch.topicPartition);
+        } else if (exception instanceof UnknownProducerIdException) {
+            // If we get an UnknownProducerId for a partition, then the broker has no state for that producer. It will
+            // therefore accept a write with sequence number 0. We reset the sequence number for the partition here so
+            // that the producer can continue after aborting the transaction. All inflight-requests to this partition
+            // will also fail with an UnknownProducerId error, so the sequence will remain at 0. Note that if the
+            // broker supports bumping the epoch, we will later reset all sequence numbers after calling InitProducerId
+            resetSequenceForPartition(batch.topicPartition);
+        } else {
+            if (adjustSequenceNumbers) {
+                if (!isTransactional()) {
+                    requestEpochBumpForPartition(batch.topicPartition);
+                } else {
+                    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) {
-        if (!this.nextSequence.containsKey(batch.topicPartition))
+    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.
             return;
@@ -514,39 +765,27 @@ 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 : inflightBatchesBySequence.get(batch.topicPartition)) {
+        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 : inflightBatchesBySequence.get(topicPartition)) {
-            log.info("Resetting sequence number of batch with current sequence {} for partition {} to {}",
-                    inFlightBatch.baseSequence(), inFlightBatch.topicPartition, sequence);
-            inFlightBatch.resetProducerState(new ProducerIdAndEpoch(inFlightBatch.producerId(),
-                    inFlightBatch.producerEpoch()), sequence, inFlightBatch.isTransactional());
-
-            sequence += inFlightBatch.recordCount;
-        }
-        setNextSequence(topicPartition, sequence);
-        lastAckedSequence.remove(topicPartition);
+        });
     }
 
     synchronized boolean hasInflightBatches(TopicPartition topicPartition) {
-        return inflightBatchesBySequence.containsKey(topicPartition) && !inflightBatchesBySequence.get(topicPartition).isEmpty();
+        return topicPartitionBookkeeper.contains(topicPartition)
+                && !topicPartitionBookkeeper.getPartition(topicPartition).inflightBatchesBySequence.isEmpty();
     }
 
     synchronized boolean hasUnresolvedSequences() {
@@ -554,51 +793,72 @@ synchronized boolean hasUnresolvedSequences() {
     }
 
     synchronized boolean hasUnresolvedSequence(TopicPartition topicPartition) {
-        return partitionsWithUnresolvedSequences.contains(topicPartition);
+        return partitionsWithUnresolvedSequences.containsKey(topicPartition);
     }
 
-    synchronized void markSequenceUnresolved(TopicPartition topicPartition) {
-        log.debug("Marking partition {} unresolved", topicPartition);
-        partitionsWithUnresolvedSequences.add(topicPartition);
+    synchronized void markSequenceUnresolved(ProducerBatch batch) {
+        int nextSequence = batch.lastSequence() + 1;
+        partitionsWithUnresolvedSequences.compute(batch.topicPartition,
+            (k, v) -> v == null ? nextSequence : Math.max(v, nextSequence));
+        log.debug("Marking partition {} unresolved with next sequence number {}", batch.topicPartition,
+                partitionsWithUnresolvedSequences.get(batch.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() {
-        if (isTransactional())
-            // We should not reset producer state if we are transactional. We will transition to a fatal error instead.
-            return false;
-        for (TopicPartition topicPartition : partitionsWithUnresolvedSequences) {
+    // Attempts to resolve unresolved sequences. If all in-flight requests are complete and some partitions are still
+    // unresolved, either bump the epoch if possible, or transition to a fatal error
+    synchronized void maybeResolveSequences() {
+        for (Iterator iter = partitionsWithUnresolvedSequences.keySet().iterator(); iter.hasNext(); ) {
+            TopicPartition topicPartition = iter.next();
             if (!hasInflightBatches(topicPartition)) {
-                // The partition has been fully drained. At this point, the last ack'd sequence should be once less than
+                // The partition has been fully drained. At this point, the last ack'd sequence should be one less than
                 // next sequence destined for the partition. If so, the partition is fully resolved. If not, we should
                 // reset the sequence number if necessary.
                 if (isNextSequence(topicPartition, sequenceNumber(topicPartition))) {
                     // This would happen when a batch was expired, but subsequent batches succeeded.
-                    partitionsWithUnresolvedSequences.remove(topicPartition);
+                    iter.remove();
                 } else {
                     // We would enter this branch if all in flight batches were ultimately expired in the producer.
-                    log.info("No inflight batches remaining for {}, last ack'd sequence for partition is {}, next sequence is {}. " +
-                            "Going to reset producer state.", topicPartition, lastAckedSequence(topicPartition), sequenceNumber(topicPartition));
-                    return true;
+                    if (isTransactional()) {
+                        // For the transactional producer, we bump the epoch if possible, otherwise we transition to a fatal error
+                        String unackedMessagesErr = "The client hasn't received acknowledgment for some previously " +
+                                "sent messages and can no longer retry them. ";
+                        if (canBumpEpoch()) {
+                            epochBumpRequired = true;
+                            KafkaException exception = new KafkaException(unackedMessagesErr + "It is safe to abort " +
+                                    "the transaction and continue.");
+                            transitionToAbortableError(exception);
+                        } else {
+                            KafkaException exception = new KafkaException(unackedMessagesErr + "It isn't safe to continue.");
+                            transitionToFatalError(exception);
+                        }
+                    } else {
+                        // For the idempotent producer, bump the epoch
+                        log.info("No inflight batches remaining for {}, last ack'd sequence for partition is {}, next sequence is {}. " +
+                                        "Going to bump epoch and reset sequence numbers.", topicPartition,
+                                lastAckedSequence(topicPartition).orElse(NO_LAST_ACKED_SEQUENCE_NUMBER), sequenceNumber(topicPartition));
+                        requestEpochBumpForPartition(topicPartition);
+                    }
+
+                    iter.remove();
                 }
             }
         }
-        return false;
     }
 
-    synchronized boolean isNextSequence(TopicPartition topicPartition, int sequence) {
-        return sequence - lastAckedSequence(topicPartition) == 1;
+    private boolean isNextSequence(TopicPartition topicPartition, int sequence) {
+        return sequence - lastAckedSequence(topicPartition).orElse(NO_LAST_ACKED_SEQUENCE_NUMBER) == 1;
+    }
+
+    private void setNextSequence(TopicPartition topicPartition, int sequence) {
+        topicPartitionBookkeeper.getPartition(topicPartition).nextSequence = sequence;
     }
 
-    private synchronized void setNextSequence(TopicPartition topicPartition, int sequence) {
-        if (!nextSequence.containsKey(topicPartition) && sequence != 0)
-            throw new IllegalStateException("Trying to set the sequence number for " + topicPartition + " to " + sequence +
-            ", but the sequence number was never set for this partition.");
-        nextSequence.put(topicPartition, sequence);
+    private boolean isNextSequenceForUnresolvedPartition(TopicPartition topicPartition, int sequence) {
+        return this.hasUnresolvedSequence(topicPartition) &&
+                sequence == this.partitionsWithUnresolvedSequences.get(topicPartition);
     }
 
-    synchronized TxnRequestHandler nextRequestHandler(boolean hasIncompleteBatches) {
+    synchronized TxnRequestHandler nextRequest(boolean hasIncompleteBatches) {
         if (!newPartitionsInTransaction.isEmpty())
             enqueueRequest(addPartitionsToTransactionHandler());
 
@@ -643,6 +903,15 @@ 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.fail(shutdownException);
+        }
+    }
+
     Node coordinator(FindCoordinatorRequest.CoordinatorType type) {
         switch (type) {
             case GROUP:
@@ -658,15 +927,15 @@ void lookupCoordinator(TxnRequestHandler request) {
         lookupCoordinator(request.coordinatorType(), request.coordinatorKey());
     }
 
-    void setInFlightTransactionalRequestCorrelationId(int correlationId) {
+    void setInFlightCorrelationId(int correlationId) {
         inFlightRequestCorrelationId = correlationId;
     }
 
-    void clearInFlightTransactionalRequestCorrelationId() {
+    private void clearInFlightCorrelationId() {
         inFlightRequestCorrelationId = NO_INFLIGHT_REQUEST_CORRELATION_ID;
     }
 
-    boolean hasInFlightTransactionalRequest() {
+    boolean hasInFlightRequest() {
         return inFlightRequestCorrelationId != NO_INFLIGHT_REQUEST_CORRELATION_ID;
     }
 
@@ -690,6 +959,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
@@ -697,23 +970,16 @@ synchronized boolean hasOngoingTransaction() {
     }
 
     synchronized boolean canRetry(ProduceResponse.PartitionResponse response, ProducerBatch batch) {
-        if (!hasProducerId(batch.producerId()))
-            return false;
-
         Errors error = response.error;
-        if (error == Errors.OUT_OF_ORDER_SEQUENCE_NUMBER && !hasUnresolvedSequence(batch.topicPartition) &&
-                (batch.sequenceHasBeenReset() || !isNextSequence(batch.topicPartition, batch.baseSequence())))
-            // We should retry the OutOfOrderSequenceException if the batch is _not_ the next batch, ie. its base
-            // sequence isn't the lastAckedSequence + 1. However, if the first in flight batch fails fatally, we will
-            // adjust the sequences of the other inflight batches to account for the 'loss' of the sequence range in
-            // the batch which failed. In this case, an inflight batch will have a base sequence which is
-            // the lastAckedSequence + 1 after adjustment. When this batch fails with an OutOfOrderSequence, we want to retry it.
-            // To account for the latter case, we check whether the sequence has been reset since the last drain.
-            // If it has, we will retry it anyway.
-            return true;
 
+        // An UNKNOWN_PRODUCER_ID means that we have lost the producer state on the broker. Depending on the log start
+        // offset, we may want to retry these, as described for each case below. If none of those apply, then for the
+        // idempotent producer, we will locally bump the epoch and reset the sequence numbers of in-flight batches from
+        // sequence 0, then retry the failed batch, which should now succeed. For the transactional producer, allow the
+        // batch to fail. When processing the failed batch, we will transition to an abortable error and set a flag
+        // indicating that we need to bump the epoch (if supported by the broker).
         if (error == Errors.UNKNOWN_PRODUCER_ID) {
-            if (response.logStartOffset == -1)
+            if (response.logStartOffset == -1) {
                 // We don't know the log start offset with this response. We should just retry the request until we get it.
                 // The UNKNOWN_PRODUCER_ID error code was added along with the new ProduceResponse which includes the
                 // logStartOffset. So the '-1' sentinel is not for backward compatibility. Instead, it is possible for
@@ -722,6 +988,7 @@ synchronized boolean canRetry(ProduceResponse.PartitionResponse response, Produc
                 // response was being constructed. In these cases, we should just retry the request: we are guaranteed
                 // to eventually get a logStartOffset once things settle down.
                 return true;
+            }
 
             if (batch.sequenceHasBeenReset()) {
                 // When the first inflight batch fails due to the truncation case, then the sequences of all the other
@@ -729,15 +996,46 @@ synchronized boolean canRetry(ProduceResponse.PartitionResponse response, Produc
                 // come back from the broker, they would also come with an UNKNOWN_PRODUCER_ID error. In this case, we should not
                 // reset the sequence numbers to the beginning.
                 return true;
-            } else if (lastAckedOffset(batch.topicPartition) < response.logStartOffset) {
+            } else if (lastAckedOffset(batch.topicPartition).orElse(NO_LAST_ACKED_SEQUENCE_NUMBER) < response.logStartOffset) {
                 // The head of the log has been removed, probably due to the retention time elapsing. In this case,
-                // we expect to lose the producer state. Reset the sequences of all inflight batches to be from the beginning
-                // and retry them.
-                startSequencesAtBeginning(batch.topicPartition);
+                // we expect to lose the producer state. For the transactional procducer, reset the sequences of all
+                // inflight batches to be from the beginning and retry them, so that the transaction does not need to
+                // be aborted. For the idempotent producer, bump the epoch to avoid reusing (sequence, epoch) pairs
+                if (isTransactional()) {
+                    topicPartitionBookkeeper.startSequencesAtBeginning(batch.topicPartition, this.producerIdAndEpoch);
+                } else {
+                    requestEpochBumpForPartition(batch.topicPartition);
+                }
+                return true;
+            }
+
+            if (!isTransactional()) {
+                // For the idempotent producer, always retry UNKNOWN_PRODUCER_ID errors. If the batch has the current
+                // producer ID and epoch, request a bump of the epoch. Otherwise just retry the produce.
+                requestEpochBumpForPartition(batch.topicPartition);
+                return true;
+            }
+        } else if (error == Errors.OUT_OF_ORDER_SEQUENCE_NUMBER) {
+            if (!hasUnresolvedSequence(batch.topicPartition) &&
+                    (batch.sequenceHasBeenReset() || !isNextSequence(batch.topicPartition, batch.baseSequence()))) {
+                // We should retry the OutOfOrderSequenceException if the batch is _not_ the next batch, ie. its base
+                // sequence isn't the lastAckedSequence + 1.
+                return true;
+            } else if (!isTransactional()) {
+                // For the idempotent producer, retry all OUT_OF_ORDER_SEQUENCE_NUMBER errors. If there are no
+                // unresolved sequences, or this batch is the one immediately following an unresolved sequence, we know
+                // there is actually a gap in the sequences, and we bump the epoch. Otherwise, retry without bumping
+                // and wait to see if the sequence resolves
+                if (!hasUnresolvedSequence(batch.topicPartition) ||
+                        isNextSequenceForUnresolvedPartition(batch.topicPartition, batch.baseSequence())) {
+                    requestEpochBumpForPartition(batch.topicPartition);
+                }
                 return true;
             }
         }
-        return false;
+
+        // If neither of the above cases are true, retry if the exception is retriable
+        return error.exception() instanceof RetriableException;
     }
 
     // visible for testing
@@ -745,11 +1043,22 @@ synchronized boolean isReady() {
         return isTransactional() && currentState == State.READY;
     }
 
+    void handleCoordinatorReady() {
+        NodeApiVersions nodeApiVersions = transactionCoordinator != null ?
+                apiVersions.get(transactionCoordinator.idString()) :
+                null;
+        ApiVersion initProducerIdVersion = nodeApiVersions != null ?
+                nodeApiVersions.apiVersion(ApiKeys.INIT_PRODUCER_ID) :
+                null;
+        this.coordinatorSupportsBumpingEpoch = initProducerIdVersion != null &&
+                initProducerIdVersion.maxVersion >= 3;
+    }
+
     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 "
@@ -758,7 +1067,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;
@@ -778,8 +1087,18 @@ 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 if (lastError instanceof InvalidProducerEpochException) {
+                throw new InvalidProducerEpochException("Producer attempted to produce with an old epoch " + producerIdAndEpoch);
+            } else {
+                throw new KafkaException("Cannot execute transactional method because we are in an error state", lastError);
+            }
+        }
     }
 
     private boolean maybeTerminateRequestWithError(TxnRequestHandler requestHandler) {
@@ -799,7 +1118,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;
@@ -811,40 +1130,87 @@ 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() {
-        transitionTo(State.READY);
-        lastError = null;
-        transactionStarted = false;
-        newPartitionsInTransaction.clear();
-        pendingPartitionsInTransaction.clear();
-        partitionsInTransaction.clear();
-    }
-
-    private synchronized TxnRequestHandler addPartitionsToTransactionHandler() {
+    private TxnRequestHandler addPartitionsToTransactionHandler() {
         pendingPartitionsInTransaction.addAll(newPartitionsInTransaction);
         newPartitionsInTransaction.clear();
-        AddPartitionsToTxnRequest.Builder builder = new AddPartitionsToTxnRequest.Builder(transactionalId,
-                producerIdAndEpoch.producerId, producerIdAndEpoch.epoch, new ArrayList<>(pendingPartitionsInTransaction));
+        AddPartitionsToTxnRequest.Builder builder =
+            new AddPartitionsToTxnRequest.Builder(transactionalId,
+                producerIdAndEpoch.producerId,
+                producerIdAndEpoch.epoch,
+                new ArrayList<>(pendingPartitionsInTransaction));
         return new AddPartitionsToTxnHandler(builder);
     }
 
     private TxnOffsetCommitHandler txnOffsetCommitHandler(TransactionalRequestResult result,
                                                           Map offsets,
-                                                          String consumerGroupId) {
+                                                          ConsumerGroupMetadata groupMetadata) {
         for (Map.Entry entry : offsets.entrySet()) {
             OffsetAndMetadata offsetAndMetadata = entry.getValue();
-            CommittedOffset committedOffset = new CommittedOffset(offsetAndMetadata.offset(), offsetAndMetadata.metadata());
+            CommittedOffset committedOffset = new CommittedOffset(offsetAndMetadata.offset(),
+                    offsetAndMetadata.metadata(), offsetAndMetadata.leaderEpoch());
             pendingTxnOffsetCommits.put(entry.getKey(), committedOffset);
         }
-        TxnOffsetCommitRequest.Builder builder = new TxnOffsetCommitRequest.Builder(transactionalId, consumerGroupId,
-                producerIdAndEpoch.producerId, producerIdAndEpoch.epoch, pendingTxnOffsetCommits);
+
+        final TxnOffsetCommitRequest.Builder builder =
+            new TxnOffsetCommitRequest.Builder(transactionalId,
+                groupMetadata.groupId(),
+                producerIdAndEpoch.producerId,
+                producerIdAndEpoch.epoch,
+                pendingTxnOffsetCommits,
+                groupMetadata.memberId(),
+                groupMetadata.generationId(),
+                groupMetadata.groupInstanceId(),
+                autoDowngradeTxnCommit
+            );
         return new TxnOffsetCommitHandler(result, builder);
     }
 
+    private TransactionalRequestResult handleCachedTransactionRequestResult(
+            Supplier transactionalRequestResultSupplier,
+            State targetState) {
+        ensureTransactional();
+
+        if (pendingResult != null && currentState == targetState) {
+            TransactionalRequestResult result = pendingResult;
+            if (result.isCompleted())
+                pendingResult = null;
+            return result;
+        }
+
+        pendingResult = transactionalRequestResultSupplier.get();
+        return pendingResult;
+    }
+
+    // package-private for testing
+    boolean canBumpEpoch() {
+        if (!isTransactional()) {
+            return true;
+        }
+
+        return coordinatorSupportsBumpingEpoch;
+    }
+
+    private void completeTransaction() {
+        if (epochBumpRequired) {
+            transitionTo(State.INITIALIZING);
+        } else {
+            transitionTo(State.READY);
+        }
+        lastError = null;
+        epochBumpRequired = false;
+        transactionStarted = false;
+        newPartitionsInTransaction.clear();
+        pendingPartitionsInTransaction.clear();
+        partitionsInTransaction.clear();
+    }
+
     abstract class TxnRequestHandler implements RequestCompletionHandler {
         protected final TransactionalRequestResult result;
         private boolean isRetry = false;
@@ -853,25 +1219,31 @@ abstract class TxnRequestHandler implements RequestCompletionHandler {
             this.result = result;
         }
 
-        TxnRequestHandler() {
-            this(new TransactionalRequestResult());
+        TxnRequestHandler(String operation) {
+            this(new TransactionalRequestResult(operation));
         }
 
         void fatalError(RuntimeException e) {
-            result.setError(e);
+            result.fail(e);
             transitionToFatalError(e);
-            result.done();
         }
 
         void abortableError(RuntimeException e) {
-            result.setError(e);
+            result.fail(e);
             transitionToAbortableError(e);
-            result.done();
+        }
+
+        void abortableErrorIfPossible(RuntimeException e) {
+            if (canBumpEpoch()) {
+                epochBumpRequired = true;
+                abortableError(e);
+            } else {
+                fatalError(e);
+            }
         }
 
         void fail(RuntimeException e) {
-            result.setError(e);
-            result.done();
+            result.fail(e);
         }
 
         void reenqueue() {
@@ -890,7 +1262,7 @@ public void onComplete(ClientResponse response) {
             if (response.requestHeader().correlationId() != inFlightRequestCorrelationId) {
                 fatalError(new RuntimeException("Detected more than one in-flight transactional request."));
             } else {
-                clearInFlightTransactionalRequestCorrelationId();
+                clearInFlightCorrelationId();
                 if (response.wasDisconnected()) {
                     log.debug("Disconnected from {}. Will retry.", response.destination());
                     if (this.needsCoordinator())
@@ -943,9 +1315,12 @@ boolean isEndTxn() {
 
     private class InitProducerIdHandler extends TxnRequestHandler {
         private final InitProducerIdRequest.Builder builder;
+        private final boolean isEpochBump;
 
-        private InitProducerIdHandler(InitProducerIdRequest.Builder builder) {
+        private InitProducerIdHandler(InitProducerIdRequest.Builder builder, boolean isEpochBump) {
+            super("InitProducerId");
             this.builder = builder;
+            this.isEpochBump = isEpochBump;
         }
 
         @Override
@@ -955,7 +1330,16 @@ InitProducerIdRequest.Builder requestBuilder() {
 
         @Override
         Priority priority() {
-            return Priority.INIT_PRODUCER_ID;
+            return this.isEpochBump ? Priority.EPOCH_BUMP : Priority.INIT_PRODUCER_ID;
+        }
+
+        @Override
+        FindCoordinatorRequest.CoordinatorType coordinatorType() {
+            if (isTransactional()) {
+                return FindCoordinatorRequest.CoordinatorType.TRANSACTION;
+            } else {
+                return null;
+            }
         }
 
         @Override
@@ -964,18 +1348,27 @@ 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;
+                if (this.isEpochBump) {
+                    resetSequenceNumbers();
+                }
                 result.done();
             } else if (error == Errors.NOT_COORDINATOR || error == Errors.COORDINATOR_NOT_AVAILABLE) {
                 lookupCoordinator(FindCoordinatorRequest.CoordinatorType.TRANSACTION, transactionalId);
                 reenqueue();
             } else if (error == Errors.COORDINATOR_LOAD_IN_PROGRESS || error == Errors.CONCURRENT_TRANSACTIONS) {
                 reenqueue();
-            } else if (error == Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED) {
+            } else if (error == Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED ||
+                    error == Errors.CLUSTER_AUTHORIZATION_FAILED) {
                 fatalError(error.exception());
+            } else if (error == Errors.INVALID_PRODUCER_EPOCH || error == Errors.PRODUCER_FENCED) {
+                // We could still receive INVALID_PRODUCER_EPOCH from old versioned transaction coordinator,
+                // just treat it the same as PRODUCE_FENCED.
+                fatalError(Errors.PRODUCER_FENCED.exception());
             } else {
                 fatalError(new KafkaException("Unexpected error in InitProducerIdResponse; " + error.message()));
             }
@@ -987,6 +1380,7 @@ private class AddPartitionsToTxnHandler extends TxnRequestHandler {
         private long retryBackoffMs;
 
         private AddPartitionsToTxnHandler(AddPartitionsToTxnRequest.Builder builder) {
+            super("AddPartitionsToTxn");
             this.builder = builder;
             this.retryBackoffMs = TransactionManager.this.retryBackoffMs;
         }
@@ -1026,14 +1420,15 @@ public void handleResponse(AbstractResponse response) {
                 } else if (error == Errors.COORDINATOR_LOAD_IN_PROGRESS || error == Errors.UNKNOWN_TOPIC_OR_PARTITION) {
                     reenqueue();
                     return;
-                } else if (error == Errors.INVALID_PRODUCER_EPOCH) {
-                    fatalError(error.exception());
+                } else if (error == Errors.INVALID_PRODUCER_EPOCH || error == Errors.PRODUCER_FENCED) {
+                    // We could still receive INVALID_PRODUCER_EPOCH from old versioned transaction coordinator,
+                    // just treat it the same as PRODUCE_FENCED.
+                    fatalError(Errors.PRODUCER_FENCED.exception());
                     return;
                 } else if (error == Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED) {
                     fatalError(error.exception());
                     return;
-                } else if (error == Errors.INVALID_PRODUCER_ID_MAPPING
-                        || error == Errors.INVALID_TXN_STATE) {
+                } else if (error == Errors.INVALID_TXN_STATE) {
                     fatalError(new KafkaException(error.exception()));
                     return;
                 } else if (error == Errors.TOPIC_AUTHORIZATION_FAILED) {
@@ -1042,6 +1437,9 @@ public void handleResponse(AbstractResponse response) {
                     log.debug("Did not attempt to add partition {} to transaction because other partitions in the " +
                             "batch had errors.", topicPartition);
                     hasPartitionErrors = true;
+                } else if (error == Errors.UNKNOWN_PRODUCER_ID || error == Errors.INVALID_PRODUCER_ID_MAPPING) {
+                    abortableErrorIfPossible(error.exception());
+                    return;
                 } else {
                     log.error("Could not add partition {} due to unexpected error {}", topicPartition, error);
                     hasPartitionErrors = true;
@@ -1090,6 +1488,7 @@ private class FindCoordinatorHandler extends TxnRequestHandler {
         private final FindCoordinatorRequest.Builder builder;
 
         private FindCoordinatorHandler(FindCoordinatorRequest.Builder builder) {
+            super("FindCoordinator");
             this.builder = builder;
         }
 
@@ -1117,27 +1516,30 @@ 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;
                     case TRANSACTION:
                         transactionCoordinator = node;
+
                 }
                 result.done();
+                log.info("Discovered {} coordinator {}", coordinatorType.toString().toLowerCase(Locale.ROOT), node);
             } else if (error == Errors.COORDINATOR_NOT_AVAILABLE) {
                 reenqueue();
             } 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())));
             }
         }
     }
@@ -1146,6 +1548,7 @@ private class EndTxnHandler extends TxnRequestHandler {
         private final EndTxnRequest.Builder builder;
 
         private EndTxnHandler(EndTxnRequest.Builder builder) {
+            super("EndTxn(" + builder.data.committed() + ")");
             this.builder = builder;
         }
 
@@ -1177,12 +1580,16 @@ public void handleResponse(AbstractResponse response) {
                 reenqueue();
             } else if (error == Errors.COORDINATOR_LOAD_IN_PROGRESS || error == Errors.CONCURRENT_TRANSACTIONS) {
                 reenqueue();
-            } else if (error == Errors.INVALID_PRODUCER_EPOCH) {
-                fatalError(error.exception());
+            } else if (error == Errors.INVALID_PRODUCER_EPOCH || error == Errors.PRODUCER_FENCED) {
+                // We could still receive INVALID_PRODUCER_EPOCH from old versioned transaction coordinator,
+                // just treat it the same as PRODUCE_FENCED.
+                fatalError(Errors.PRODUCER_FENCED.exception());
             } else if (error == Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED) {
                 fatalError(error.exception());
             } else if (error == Errors.INVALID_TXN_STATE) {
                 fatalError(error.exception());
+            } else if (error == Errors.UNKNOWN_PRODUCER_ID || error == Errors.INVALID_PRODUCER_ID_MAPPING) {
+                abortableErrorIfPossible(error.exception());
             } else {
                 fatalError(new KafkaException("Unhandled error in EndTxnResponse: " + error.message()));
             }
@@ -1192,11 +1599,15 @@ public void handleResponse(AbstractResponse response) {
     private class AddOffsetsToTxnHandler extends TxnRequestHandler {
         private final AddOffsetsToTxnRequest.Builder builder;
         private final Map offsets;
+        private final ConsumerGroupMetadata groupMetadata;
 
         private AddOffsetsToTxnHandler(AddOffsetsToTxnRequest.Builder builder,
-                                       Map offsets) {
+                                       Map offsets,
+                                       ConsumerGroupMetadata groupMetadata) {
+            super("AddOffsetsToTxn");
             this.builder = builder;
             this.offsets = offsets;
+            this.groupMetadata = groupMetadata;
         }
 
         @Override
@@ -1212,25 +1623,30 @@ Priority priority() {
         @Override
         public void handleResponse(AbstractResponse response) {
             AddOffsetsToTxnResponse addOffsetsToTxnResponse = (AddOffsetsToTxnResponse) response;
-            Errors error = addOffsetsToTxnResponse.error();
+            Errors error = Errors.forCode(addOffsetsToTxnResponse.data().errorCode());
 
             if (error == Errors.NONE) {
-                log.debug("Successfully added partition for consumer group {} to transaction", builder.consumerGroupId());
+                log.debug("Successfully added partition for consumer group {} to transaction", builder.data.groupId());
 
                 // note the result is not completed until the TxnOffsetCommit returns
-                pendingRequests.add(txnOffsetCommitHandler(result, offsets, builder.consumerGroupId()));
+                pendingRequests.add(txnOffsetCommitHandler(result, offsets, groupMetadata));
+
                 transactionStarted = true;
             } else if (error == Errors.COORDINATOR_NOT_AVAILABLE || error == Errors.NOT_COORDINATOR) {
                 lookupCoordinator(FindCoordinatorRequest.CoordinatorType.TRANSACTION, transactionalId);
                 reenqueue();
             } else if (error == Errors.COORDINATOR_LOAD_IN_PROGRESS || error == Errors.CONCURRENT_TRANSACTIONS) {
                 reenqueue();
-            } else if (error == Errors.INVALID_PRODUCER_EPOCH) {
-                fatalError(error.exception());
+            } else if (error == Errors.UNKNOWN_PRODUCER_ID || error == Errors.INVALID_PRODUCER_ID_MAPPING) {
+                abortableErrorIfPossible(error.exception());
+            } else if (error == Errors.INVALID_PRODUCER_EPOCH || error == Errors.PRODUCER_FENCED) {
+                // We could still receive INVALID_PRODUCER_EPOCH from old versioned transaction coordinator,
+                // just treat it the same as PRODUCE_FENCED.
+                fatalError(Errors.PRODUCER_FENCED.exception());
             } else if (error == Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED) {
                 fatalError(error.exception());
             } else if (error == Errors.GROUP_AUTHORIZATION_FAILED) {
-                abortableError(new GroupAuthorizationException(builder.consumerGroupId()));
+                abortableError(GroupAuthorizationException.forGroupId(builder.data.groupId()));
             } else {
                 fatalError(new KafkaException("Unexpected error in AddOffsetsToTxnResponse: " + error.message()));
             }
@@ -1263,57 +1679,69 @@ FindCoordinatorRequest.CoordinatorType coordinatorType() {
 
         @Override
         String coordinatorKey() {
-            return builder.consumerGroupId();
+            return builder.data.groupId();
         }
 
         @Override
         public void handleResponse(AbstractResponse response) {
             TxnOffsetCommitResponse txnOffsetCommitResponse = (TxnOffsetCommitResponse) response;
             boolean coordinatorReloaded = false;
-            boolean hadFailure = false;
             Map errors = txnOffsetCommitResponse.errors();
 
+            log.debug("Received TxnOffsetCommit response for consumer group {}: {}", builder.data.groupId(),
+                    errors);
+
             for (Map.Entry entry : errors.entrySet()) {
                 TopicPartition topicPartition = entry.getKey();
                 Errors error = entry.getValue();
                 if (error == Errors.NONE) {
-                    log.debug("Successfully added offsets {} from consumer group {} to transaction.",
-                            builder.offsets(), builder.consumerGroupId());
                     pendingTxnOffsetCommits.remove(topicPartition);
                 } else if (error == Errors.COORDINATOR_NOT_AVAILABLE
                         || error == Errors.NOT_COORDINATOR
                         || error == Errors.REQUEST_TIMED_OUT) {
-                    hadFailure = true;
                     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) {
-                    hadFailure = true;
+                } 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()));
-                    return;
-                } else if (error == Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED
-                        || error == Errors.INVALID_PRODUCER_EPOCH
-                        || error == Errors.UNSUPPORTED_FOR_MESSAGE_FORMAT) {
+                    abortableError(GroupAuthorizationException.forGroupId(builder.data.groupId()));
+                    break;
+                } else if (error == Errors.FENCED_INSTANCE_ID) {
+                    abortableError(error.exception());
+                    break;
+                } else if (error == Errors.UNKNOWN_MEMBER_ID
+                        || error == Errors.ILLEGAL_GENERATION) {
+                    abortableError(new CommitFailedException("Transaction offset Commit failed " +
+                        "due to consumer group metadata mismatch: " + error.exception().getMessage()));
+                    break;
+                } else if (isFatalException(error)) {
                     fatalError(error.exception());
-                    return;
+                    break;
                 } else {
                     fatalError(new KafkaException("Unexpected error in TxnOffsetCommitResponse: " + error.message()));
-                    return;
+                    break;
                 }
             }
 
-            if (!hadFailure || !result.isSuccessful()) {
-                // all attempted partitions were either successful, or there was a fatal failure.
-                // either way, we are not retrying, so complete the request.
+            if (result.isCompleted()) {
+                pendingTxnOffsetCommits.clear();
+            } else if (pendingTxnOffsetCommits.isEmpty()) {
                 result.done();
-                return;
-            }
-
-            // retry the commits which failed with a retriable error.
-            if (!pendingTxnOffsetCommits.isEmpty())
+            } else {
+                // Retry the commits which failed with a retriable error
                 reenqueue();
+            }
         }
     }
+
+    private boolean isFatalException(Errors error) {
+        return error == Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED
+                   || error == Errors.INVALID_PRODUCER_EPOCH
+                   || error == Errors.PRODUCER_FENCED
+                   || error == Errors.UNSUPPORTED_FOR_MESSAGE_FORMAT;
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/clients/producer/internals/TransactionalRequestResult.java b/clients/src/main/java/org/apache/kafka/clients/producer/internals/TransactionalRequestResult.java
index ff93da872dc97..9881fe8ac4e77 100644
--- a/clients/src/main/java/org/apache/kafka/clients/producer/internals/TransactionalRequestResult.java
+++ b/clients/src/main/java/org/apache/kafka/clients/producer/internals/TransactionalRequestResult.java
@@ -17,25 +17,31 @@
 package org.apache.kafka.clients.producer.internals;
 
 
+import org.apache.kafka.common.errors.InterruptException;
+import org.apache.kafka.common.errors.TimeoutException;
+
+import java.util.Locale;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 
 public final class TransactionalRequestResult {
-    static final TransactionalRequestResult COMPLETE = new TransactionalRequestResult(new CountDownLatch(0));
 
     private final CountDownLatch latch;
     private volatile RuntimeException error = null;
+    private final String operation;
 
-    public TransactionalRequestResult() {
-        this(new CountDownLatch(1));
+    public TransactionalRequestResult(String operation) {
+        this(new CountDownLatch(1), operation);
     }
 
-    private TransactionalRequestResult(CountDownLatch latch) {
+    private TransactionalRequestResult(CountDownLatch latch, String operation) {
         this.latch = latch;
+        this.operation = operation;
     }
 
-    public void setError(RuntimeException error) {
+    public void fail(RuntimeException error) {
         this.error = error;
+        this.latch.countDown();
     }
 
     public void done() {
@@ -58,8 +64,18 @@ public void await() {
             throw error();
     }
 
-    public boolean await(long timeout, TimeUnit unit) throws InterruptedException {
-        return latch.await(timeout, unit);
+    public void await(long timeout, TimeUnit unit) {
+        try {
+            boolean success = latch.await(timeout, unit);
+            if (!isSuccessful()) {
+                throw error();
+            }
+            if (!success) {
+                throw new TimeoutException("Timeout expired after " + timeout + unit.name().toLowerCase(Locale.ROOT) + " while awaiting " + operation);
+            }
+        } catch (InterruptedException e) {
+            throw new InterruptException("Received interrupt while awaiting " + operation, e);
+        }
     }
 
     public RuntimeException error() {
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 0c59f33824d8f..a6bf76310af05 100644
--- a/clients/src/main/java/org/apache/kafka/common/Cluster.java
+++ b/clients/src/main/java/org/apache/kafka/common/Cluster.java
@@ -16,26 +16,28 @@
  */
 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;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
 import java.util.Set;
 
 /**
- * A representation of a subset of the nodes, topics, and partitions in the Kafka cluster.
+ * An immutable representation of a subset of the nodes, topics, and partitions in the Kafka cluster.
  */
 public final class Cluster {
 
     private final boolean isBootstrapConfigured;
     private final List nodes;
     private final Set unauthorizedTopics;
+    private final Set invalidTopics;
     private final Set internalTopics;
     private final Node controller;
     private final Map partitionsByTopicPartition;
@@ -55,7 +57,21 @@ public Cluster(String clusterId,
                    Collection partitions,
                    Set unauthorizedTopics,
                    Set internalTopics) {
-        this(clusterId, false, nodes, partitions, unauthorizedTopics, internalTopics, null);
+        this(clusterId, false, nodes, partitions, unauthorizedTopics, Collections.emptySet(), internalTopics, null);
+    }
+
+    /**
+     * Create a new cluster with the given id, nodes and partitions
+     * @param nodes The nodes in the cluster
+     * @param partitions Information about a subset of the topic-partitions this cluster hosts
+     */
+    public Cluster(String clusterId,
+                   Collection nodes,
+                   Collection partitions,
+                   Set unauthorizedTopics,
+                   Set internalTopics,
+                   Node controller) {
+        this(clusterId, false, nodes, partitions, unauthorizedTopics, Collections.emptySet(), internalTopics, controller);
     }
 
     /**
@@ -67,9 +83,10 @@ public Cluster(String clusterId,
                    Collection nodes,
                    Collection partitions,
                    Set unauthorizedTopics,
+                   Set invalidTopics,
                    Set internalTopics,
                    Node controller) {
-        this(clusterId, false, nodes, partitions, unauthorizedTopics, internalTopics, controller);
+        this(clusterId, false, nodes, partitions, unauthorizedTopics, invalidTopics, internalTopics, controller);
     }
 
     private Cluster(String clusterId,
@@ -77,6 +94,7 @@ private Cluster(String clusterId,
                     Collection nodes,
                     Collection partitions,
                     Set unauthorizedTopics,
+                    Set invalidTopics,
                     Set internalTopics,
                     Node controller) {
         this.isBootstrapConfigured = isBootstrapConfigured;
@@ -85,52 +103,71 @@ private Cluster(String clusterId,
         List copy = new ArrayList<>(nodes);
         Collections.shuffle(copy);
         this.nodes = Collections.unmodifiableList(copy);
-        this.nodesById = new HashMap<>();
-        for (Node node : nodes)
-            this.nodesById.put(node.id(), node);
-
-        // index the partitions by topic/partition for quick lookup
-        this.partitionsByTopicPartition = new HashMap<>(partitions.size());
-        for (PartitionInfo p : partitions)
-            this.partitionsByTopicPartition.put(new TopicPartition(p.topic(), p.partition()), p);
-
-        // index the partitions by topic and node respectively, and make the lists
-        // unmodifiable so we can hand them out in user-facing apis without risk
-        // of the client modifying the contents
-        HashMap> partsForTopic = new HashMap<>();
-        HashMap> partsForNode = new HashMap<>();
-        for (Node n : this.nodes) {
-            partsForNode.put(n.id(), new ArrayList());
+
+        // Index the nodes for quick lookup
+        Map tmpNodesById = new HashMap<>();
+        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<>();
         for (PartitionInfo p : partitions) {
-            if (!partsForTopic.containsKey(p.topic()))
-                partsForTopic.put(p.topic(), new ArrayList());
-            List psTopic = partsForTopic.get(p.topic());
-            psTopic.add(p);
-
-            if (p.leader() != null) {
-                List psNode = Utils.notNull(partsForNode.get(p.leader().id()));
-                psNode.add(p);
-            }
+            tmpPartitionsByTopicPartition.put(new TopicPartition(p.topic(), p.partition()), p);
+            tmpPartitionsByTopic.computeIfAbsent(p.topic(), topic -> new ArrayList<>()).add(p);
+
+            // The leader may not be known
+            if (p.leader() == null || p.leader().isEmpty())
+                continue;
+
+            // If it is known, its node information should be available
+            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()));
         }
-        this.partitionsByTopic = new HashMap<>(partsForTopic.size());
-        this.availablePartitionsByTopic = new HashMap<>(partsForTopic.size());
-        for (Map.Entry> entry : partsForTopic.entrySet()) {
+
+        // 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 partitionList = entry.getValue();
-            this.partitionsByTopic.put(topic, Collections.unmodifiableList(partitionList));
-            List availablePartitions = new ArrayList<>();
-            for (PartitionInfo part : partitionList) {
-                if (part.leader() != null)
-                    availablePartitions.add(part);
+            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;
             }
-            this.availablePartitionsByTopic.put(topic, Collections.unmodifiableList(availablePartitions));
+            tmpAvailablePartitionsByTopic.put(topic, availablePartitionsForTopic);
         }
-        this.partitionsByNode = new HashMap<>(partsForNode.size());
-        for (Map.Entry> entry : partsForNode.entrySet())
-            this.partitionsByNode.put(entry.getKey(), Collections.unmodifiableList(entry.getValue()));
+
+        this.partitionsByTopicPartition = Collections.unmodifiableMap(tmpPartitionsByTopicPartition);
+        this.partitionsByTopic = Collections.unmodifiableMap(tmpPartitionsByTopic);
+        this.availablePartitionsByTopic = Collections.unmodifiableMap(tmpAvailablePartitionsByTopic);
+        this.partitionsByNode = Collections.unmodifiableMap(tmpPartitionsByNode);
 
         this.unauthorizedTopics = Collections.unmodifiableSet(unauthorizedTopics);
+        this.invalidTopics = Collections.unmodifiableSet(invalidTopics);
         this.internalTopics = Collections.unmodifiableSet(internalTopics);
         this.controller = controller;
     }
@@ -139,8 +176,8 @@ private Cluster(String clusterId,
      * Create an empty cluster instance with no nodes and no topic-partitions.
      */
     public static Cluster empty() {
-        return new Cluster(null, new ArrayList(0), new ArrayList(0), Collections.emptySet(),
-                Collections.emptySet(), null);
+        return new Cluster(null, new ArrayList<>(0), new ArrayList<>(0), Collections.emptySet(),
+            Collections.emptySet(), null);
     }
 
     /**
@@ -153,7 +190,8 @@ public static Cluster bootstrap(List addresses) {
         int nodeId = -1;
         for (InetSocketAddress address : addresses)
             nodes.add(new Node(nodeId--, address.getHostString(), address.getPort()));
-        return new Cluster(null, true, nodes, new ArrayList(0), Collections.emptySet(), Collections.emptySet(), null);
+        return new Cluster(null, true, nodes, new ArrayList<>(0),
+            Collections.emptySet(), Collections.emptySet(), Collections.emptySet(), null);
     }
 
     /**
@@ -163,7 +201,8 @@ public Cluster withPartitions(Map partitions) {
         Map combinedPartitions = new HashMap<>(this.partitionsByTopicPartition);
         combinedPartitions.putAll(partitions);
         return new Cluster(clusterResource.clusterId(), this.nodes, combinedPartitions.values(),
-                new HashSet<>(this.unauthorizedTopics), new HashSet<>(this.internalTopics), this.controller);
+                new HashSet<>(this.unauthorizedTopics), new HashSet<>(this.invalidTopics),
+                new HashSet<>(this.internalTopics), this.controller);
     }
 
     /**
@@ -172,7 +211,7 @@ public Cluster withPartitions(Map partitions) {
     public List nodes() {
         return this.nodes;
     }
-    
+
     /**
      * Get the node by the node id (or null if no such node exists)
      * @param id The id of the node
@@ -182,6 +221,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
@@ -198,7 +252,7 @@ public Node leaderFor(TopicPartition topicPartition) {
     /**
      * Get the metadata for the specified partition
      * @param topicPartition The topic and partition to fetch info for
-     * @return The metadata about the given topic and partition
+     * @return The metadata about the given topic and partition, or null if none is found
      */
     public PartitionInfo partition(TopicPartition topicPartition) {
         return partitionsByTopicPartition.get(topicPartition);
@@ -210,12 +264,11 @@ public PartitionInfo partition(TopicPartition topicPartition) {
      * @return A list of partitions
      */
     public List partitionsForTopic(String topic) {
-        List parts = this.partitionsByTopic.get(topic);
-        return (parts == null) ? Collections.emptyList() : parts;
+        return partitionsByTopic.getOrDefault(topic, Collections.emptyList());
     }
 
     /**
-     * Get the number of partitions for the given topic
+     * Get the number of partitions for the given topic.
      * @param topic The topic to get the number of partitions for
      * @return The number of partitions or null if there is no corresponding metadata
      */
@@ -230,8 +283,7 @@ public Integer partitionCountForTopic(String topic) {
      * @return A list of partitions
      */
     public List availablePartitionsForTopic(String topic) {
-        List parts = this.availablePartitionsByTopic.get(topic);
-        return (parts == null) ? Collections.emptyList() : parts;
+        return availablePartitionsByTopic.getOrDefault(topic, Collections.emptyList());
     }
 
     /**
@@ -240,8 +292,7 @@ public List availablePartitionsForTopic(String topic) {
      * @return A list of partitions
      */
     public List partitionsForNode(int nodeId) {
-        List parts = this.partitionsByNode.get(nodeId);
-        return (parts == null) ? Collections.emptyList() : parts;
+        return partitionsByNode.getOrDefault(nodeId, Collections.emptyList());
     }
 
     /**
@@ -249,13 +300,17 @@ public List partitionsForNode(int nodeId) {
      * @return a set of all topics
      */
     public Set topics() {
-        return this.partitionsByTopic.keySet();
+        return partitionsByTopic.keySet();
     }
 
     public Set unauthorizedTopics() {
         return unauthorizedTopics;
     }
 
+    public Set invalidTopics() {
+        return invalidTopics;
+    }
+
     public Set internalTopics() {
         return internalTopics;
     }
@@ -274,7 +329,28 @@ public Node controller() {
 
     @Override
     public String toString() {
-        return "Cluster(id = " + clusterResource.clusterId() + ", nodes = " + this.nodes + ", partitions = " + this.partitionsByTopicPartition.values() + ")";
+        return "Cluster(id = " + clusterResource.clusterId() + ", nodes = " + this.nodes +
+            ", partitions = " + this.partitionsByTopicPartition.values() + ", controller = " + controller + ")";
     }
 
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        Cluster cluster = (Cluster) o;
+        return isBootstrapConfigured == cluster.isBootstrapConfigured &&
+                Objects.equals(nodes, cluster.nodes) &&
+                Objects.equals(unauthorizedTopics, cluster.unauthorizedTopics) &&
+                Objects.equals(invalidTopics, cluster.invalidTopics) &&
+                Objects.equals(internalTopics, cluster.internalTopics) &&
+                Objects.equals(controller, cluster.controller) &&
+                Objects.equals(partitionsByTopicPartition, cluster.partitionsByTopicPartition) &&
+                Objects.equals(clusterResource, cluster.clusterResource);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(isBootstrapConfigured, nodes, unauthorizedTopics, invalidTopics, internalTopics, controller,
+                partitionsByTopicPartition, clusterResource);
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/ClusterResource.java b/clients/src/main/java/org/apache/kafka/common/ClusterResource.java
index c2737c9241805..749f2d124e077 100644
--- a/clients/src/main/java/org/apache/kafka/common/ClusterResource.java
+++ b/clients/src/main/java/org/apache/kafka/common/ClusterResource.java
@@ -17,6 +17,8 @@
 package org.apache.kafka.common;
 
 
+import java.util.Objects;
+
 /**
  * The ClusterResource class encapsulates metadata for a Kafka cluster.
  */
@@ -46,4 +48,17 @@ public String clusterId() {
     public String toString() {
         return "ClusterResource(clusterId=" + clusterId + ")";
     }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        ClusterResource that = (ClusterResource) o;
+        return Objects.equals(clusterId, that.clusterId);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(clusterId);
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/ClusterResourceListener.java b/clients/src/main/java/org/apache/kafka/common/ClusterResourceListener.java
index f8f99ece18384..f1939df29559d 100644
--- a/clients/src/main/java/org/apache/kafka/common/ClusterResourceListener.java
+++ b/clients/src/main/java/org/apache/kafka/common/ClusterResourceListener.java
@@ -26,17 +26,18 @@
  * There will be one invocation of {@link ClusterResourceListener#onUpdate(ClusterResource)} after each metadata response.
  * Note that the cluster id may be null when the Kafka broker version is below 0.10.1.0. If you receive a null cluster id, you can expect it to always be null unless you have a cluster with multiple broker versions which can happen if the cluster is being upgraded while the client is running.
  * 
        
- * {@link org.apache.kafka.clients.producer.ProducerInterceptor} : The {@link ClusterResourceListener#onUpdate(ClusterResource)} method will be invoked after {@link org.apache.kafka.clients.producer.ProducerInterceptor#onSend(ProducerRecord)}
- * but before {@link org.apache.kafka.clients.producer.ProducerInterceptor#onAcknowledgement(RecordMetadata, Exception)} .
+ * {@link org.apache.kafka.clients.producer.ProducerInterceptor} : The {@link ClusterResourceListener#onUpdate(ClusterResource)} method will be invoked after {@link org.apache.kafka.clients.producer.ProducerInterceptor#onSend(org.apache.kafka.clients.producer.ProducerRecord)}
+ * but before {@link org.apache.kafka.clients.producer.ProducerInterceptor#onAcknowledgement(org.apache.kafka.clients.producer.RecordMetadata, Exception)} .
  * 
        
- * {@link org.apache.kafka.clients.consumer.ConsumerInterceptor} : The {@link ClusterResourceListener#onUpdate(ClusterResource)} method will be invoked before {@link org.apache.kafka.clients.consumer.ConsumerInterceptor#onConsume(ConsumerRecords)}
+ * {@link org.apache.kafka.clients.consumer.ConsumerInterceptor} : The {@link ClusterResourceListener#onUpdate(ClusterResource)} method will be invoked before {@link org.apache.kafka.clients.consumer.ConsumerInterceptor#onConsume(org.apache.kafka.clients.consumer.ConsumerRecords)}
  * 
        
  * {@link org.apache.kafka.common.serialization.Serializer} : The {@link ClusterResourceListener#onUpdate(ClusterResource)} method will be invoked before {@link org.apache.kafka.common.serialization.Serializer#serialize(String, Object)}
  * 
        
  * {@link org.apache.kafka.common.serialization.Deserializer} : The {@link ClusterResourceListener#onUpdate(ClusterResource)} method will be invoked before {@link org.apache.kafka.common.serialization.Deserializer#deserialize(String, byte[])}
  * 
        
- * {@link org.apache.kafka.common.metrics.MetricsReporter} : The {@link ClusterResourceListener#onUpdate(ClusterResource)} method will be invoked after first {@link org.apache.kafka.clients.producer.KafkaProducer#send(ProducerRecord)} invocation for Producer metrics reporter
- * and after first {@link org.apache.kafka.clients.consumer.KafkaConsumer#poll(long)} invocation for Consumer metrics reporters. The reporter may receive metric events from the network layer before this method is invoked.
+ * {@link org.apache.kafka.common.metrics.MetricsReporter} : The {@link ClusterResourceListener#onUpdate(ClusterResource)} method will be invoked after first {@link org.apache.kafka.clients.producer.KafkaProducer#send(org.apache.kafka.clients.producer.ProducerRecord)} invocation for Producer metrics reporter
+ * and after first {@link org.apache.kafka.clients.consumer.KafkaConsumer#poll(java.time.Duration)} invocation for Consumer metrics
+ * reporters. The reporter may receive metric events from the network layer before this method is invoked.
  * 
        Broker
        
  * There is a single invocation {@link ClusterResourceListener#onUpdate(ClusterResource)} on broker start-up and the cluster metadata will never change.
  * 
        
diff --git a/clients/src/main/java/org/apache/kafka/common/ConsumerGroupState.java b/clients/src/main/java/org/apache/kafka/common/ConsumerGroupState.java
new file mode 100644
index 0000000000000..36d1a4da36118
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/ConsumerGroupState.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;
+
+import java.util.HashMap;
+
+/**
+ * The consumer group state.
+ */
+public enum ConsumerGroupState {
+    UNKNOWN("Unknown"),
+    PREPARING_REBALANCE("PreparingRebalance"),
+    COMPLETING_REBALANCE("CompletingRebalance"),
+    STABLE("Stable"),
+    DEAD("Dead"),
+    EMPTY("Empty");
+
+    private final static HashMap NAME_TO_ENUM;
+
+    static {
+        NAME_TO_ENUM = new HashMap<>();
+        for (ConsumerGroupState state : ConsumerGroupState.values()) {
+            NAME_TO_ENUM.put(state.name, state);
+        }
+    }
+
+    private final String name;
+
+    ConsumerGroupState(String name) {
+        this.name = name;
+    }
+
+    /**
+     * Parse a string into a consumer group state.
+     */
+    public static ConsumerGroupState parse(String name) {
+        ConsumerGroupState state = NAME_TO_ENUM.get(name);
+        return state == null ? UNKNOWN : state;
+    }
+
+    @Override
+    public String toString() {
+        return name;
+    }
+}
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..55331c5ea9ecd
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/ElectionType.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;
+
+import org.apache.kafka.common.annotation.InterfaceStability;
+
+import java.util.Arrays;
+import java.util.Set;
+
+/**
+ * Options for {@link org.apache.kafka.clients.admin.Admin#electLeaders(ElectionType, Set, org.apache.kafka.clients.admin.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/requests/IsolationLevel.java b/clients/src/main/java/org/apache/kafka/common/IsolationLevel.java
similarity index 96%
rename from clients/src/main/java/org/apache/kafka/common/requests/IsolationLevel.java
rename to clients/src/main/java/org/apache/kafka/common/IsolationLevel.java
index a09b625c050a3..79f0a92954bf1 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/IsolationLevel.java
+++ b/clients/src/main/java/org/apache/kafka/common/IsolationLevel.java
@@ -14,7 +14,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.kafka.common.requests;
+package org.apache.kafka.common;
 
 public enum IsolationLevel {
     READ_UNCOMMITTED((byte) 0), READ_COMMITTED((byte) 1);
diff --git a/clients/src/main/java/org/apache/kafka/common/KafkaFuture.java b/clients/src/main/java/org/apache/kafka/common/KafkaFuture.java
index 23e218137a93a..9cd2e01dc42f6 100644
--- a/clients/src/main/java/org/apache/kafka/common/KafkaFuture.java
+++ b/clients/src/main/java/org/apache/kafka/common/KafkaFuture.java
@@ -35,18 +35,26 @@ public abstract class KafkaFuture implements Future {
     /**
      * A function which takes objects of type A and returns objects of type B.
      */
-    public static abstract class Function {
-        public abstract B apply(A a);
+    public interface BaseFunction {
+        B apply(A a);
     }
 
+    /**
+     * A function which takes objects of type A and returns objects of type B.
+     *
+     * Prefer the functional interface {@link BaseFunction} over the class {@link Function}.  This class is here for
+     * backwards compatibility reasons and might be deprecated/removed in a future release.
+     */
+    public static abstract class Function implements BaseFunction { }
+
     /**
      * A consumer of two different types of object.
      */
-    public static abstract class BiConsumer {
-        public abstract void accept(A a, B b);
+    public interface BiConsumer {
+        void accept(A a, B b);
     }
 
-    private static class AllOfAdapter extends BiConsumer {
+    private static class AllOfAdapter implements BiConsumer {
         private int remainingResponses;
         private KafkaFuture future;
 
@@ -101,11 +109,43 @@ public static KafkaFuture allOf(KafkaFuture... futures) {
     /**
      * Returns a new KafkaFuture that, when this future completes normally, is executed with this
      * futures's result as the argument to the supplied function.
+     *
+     * The function may be invoked by the thread that calls {@code thenApply} or it may be invoked by the thread that
+     * completes the future.
+     */
+    public abstract  KafkaFuture thenApply(BaseFunction function);
+
+    /**
+     * @see KafkaFuture#thenApply(BaseFunction)
+     *
+     * Prefer {@link KafkaFuture#thenApply(BaseFunction)} as this function is here for backwards compatibility reasons
+     * and might be deprecated/removed in a future release.
      */
     public abstract  KafkaFuture thenApply(Function function);
 
-    protected abstract void addWaiter(BiConsumer action);
+    /**
+     * Returns a new KafkaFuture with the same result or exception as this future, that executes the given action
+     * when this future completes.
+     *
+     * When this future is done, the given action is invoked with the result (or null if none) and the exception
+     * (or null if none) of this future as arguments.
+     *
+     * The returned future is completed when the action returns.
+     * The supplied action should not throw an exception. However, if it does, the following rules apply:
+     * if this future completed normally but the supplied action throws an exception, then the returned future completes
+     * exceptionally with the supplied action's exception.
+     * Or, if this future completed exceptionally and the supplied action throws an exception, then the returned future
+     * completes exceptionally with this future's exception.
+     *
+     * The action may be invoked by the thread that calls {@code whenComplete} or it may be invoked by the thread that
+     * completes the future.
+     *
+     * @param action the action to preform
+     * @return the new future
+     */
+    public abstract KafkaFuture whenComplete(BiConsumer action);
 
+    protected abstract void addWaiter(BiConsumer action);
     /**
      * If not already completed, sets the value returned by get() and related methods to the given
      * value.
diff --git a/clients/src/main/java/org/apache/kafka/common/MessageFormatter.java b/clients/src/main/java/org/apache/kafka/common/MessageFormatter.java
new file mode 100644
index 0000000000000..c4a255f315571
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/MessageFormatter.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;
+
+import java.io.Closeable;
+import java.io.PrintStream;
+import java.util.Map;
+import java.util.Properties;
+
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+
+/**
+ * This interface allows to define Formatters that can be used to parse and format records read by a
+ *  Consumer instance for display.
+ * The kafka-console-consumer has built-in support for MessageFormatter, via the --formatter flag.
+ *
+ * Kafka provides a few implementations to display records of internal topics such as __consumer_offsets,
+ * __transaction_state and the MirrorMaker2 topics.
+ *
+ */
+public interface MessageFormatter extends Configurable, Closeable {
+
+    /**
+     * Initialises the MessageFormatter
+     * @param props Properties to configure the formatter
+     * @deprecated Use {@link #configure(Map)} instead, this method is for backward compatibility with the older Formatter interface
+     */
+    @Deprecated
+    default public void init(Properties props) {}
+
+    /**
+     * Configures the MessageFormatter
+     * @param configs Map to configure the formatter
+     */
+    default public void configure(Map configs) {
+        Properties properties = new Properties();
+        properties.putAll(configs);
+        init(properties);
+    }
+
+    /**
+     * Parses and formats a record for display
+     * @param consumerRecord the record to format
+     * @param output the print stream used to output the record
+     */
+    public void writeTo(ConsumerRecord consumerRecord, PrintStream output);
+
+    /**
+     * Closes the formatter
+     */
+    default public void close() {}
+}
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 2136a72b0e4b8..b1ccf302b6599 100644
--- a/clients/src/main/java/org/apache/kafka/common/MetricName.java
+++ b/clients/src/main/java/org/apache/kafka/common/MetricName.java
@@ -17,13 +17,12 @@
 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(...).
  * 
        
- * This class captures the following parameters
+ * This class captures the following parameters:
  * 
          *  name The name of the metric
  *  group logical group name of the metrics to which this metric belongs.
@@ -32,7 +31,7 @@
  * 
        
  * group, tags parameters can be used to create unique metric names while reporting in JMX or any custom reporting.
  * 
        
- * Ex: standard JMX MBean can be constructed like  domainName:type=group,key1=val1,key2=val2
+ * Ex: standard JMX MBean can be constructed like domainName:type=group,key1=val1,key2=val2
  * 
        
  *
  * Usage looks something like this:
@@ -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() {
@@ -106,9 +105,9 @@ public int hashCode() {
             return hash;
         final int prime = 31;
         int result = 1;
-        result = prime * result + ((group == null) ? 0 : group.hashCode());
-        result = prime * result + ((name == null) ? 0 : name.hashCode());
-        result = prime * result + ((tags == null) ? 0 : tags.hashCode());
+        result = prime * result + group.hashCode();
+        result = prime * result + name.hashCode();
+        result = prime * result + tags.hashCode();
         this.hash = result;
         return result;
     }
@@ -122,22 +121,7 @@ public boolean equals(Object obj) {
         if (getClass() != obj.getClass())
             return false;
         MetricName other = (MetricName) obj;
-        if (group == null) {
-            if (other.group != null)
-                return false;
-        } else if (!group.equals(other.group))
-            return false;
-        if (name == null) {
-            if (other.name != null)
-                return false;
-        } else if (!name.equals(other.name))
-            return false;
-        if (tags == null) {
-            if (other.tags != null)
-                return false;
-        } else if (!tags.equals(other.tags))
-            return false;
-        return true;
+        return group.equals(other.group) && name.equals(other.name) && tags.equals(other.tags);
     }
 
     @Override
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 e3ea9950ef159..f9c4ef5f730e6 100644
--- a/clients/src/main/java/org/apache/kafka/common/MetricNameTemplate.java
+++ b/clients/src/main/java/org/apache/kafka/common/MetricNameTemplate.java
@@ -16,56 +16,92 @@
  */
 package org.apache.kafka.common;
 
-import java.util.HashSet;
+import java.util.Collections;
+import java.util.LinkedHashSet;
 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
  * are omitted from the template, but are filled in at runtime with their
- * specified values.
+ * specified values. The order of the tags is maintained, if an ordered set
+ * is provided, so that the mBean names can be compared and sorted lexicographically.
  */
 public class MetricNameTemplate {
     private final String name;
     private final String group;
     private final String description;
-    private Set tags;
+    private LinkedHashSet tags;
 
-    public MetricNameTemplate(String name, String group, String description, Set tags) {
-        this.name = Utils.notNull(name);
-        this.group = Utils.notNull(group);
-        this.description = Utils.notNull(description);
-        this.tags = Utils.notNull(tags);
+    /**
+     * Create a new template. Note that the order of the tags will be preserved if the supplied
+     * {@code tagsNames} set has an order.
+     *
+     * @param name the name of the metric; may not be null
+     * @param group the name of the group; may not be null
+     * @param description the description of the metric; may not be null
+     * @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 = Objects.requireNonNull(name);
+        this.group = Objects.requireNonNull(group);
+        this.description = Objects.requireNonNull(description);
+        this.tags = new LinkedHashSet<>(Objects.requireNonNull(tagsNames));
     }
-    
-    public MetricNameTemplate(String name, String group, String description, String... keys) {
-        this(name, group, description, getTags(keys));
+
+    /**
+     * Create a new template. Note that the order of the tags will be preserved.
+     *
+     * @param name the name of the metric; may not be null
+     * @param group the name of the group; may not be null
+     * @param description the description of the metric; may not be null
+     * @param tagsNames the names of the metric tags in the preferred order; none of the tag names should be null
+     */
+    public MetricNameTemplate(String name, String group, String description, String... tagsNames) {
+        this(name, group, description, getTags(tagsNames));
     }
 
-    private static Set getTags(String... keys) {
-        Set tags = new HashSet();
-        
-        for (int i = 0; i < keys.length; i++)
-            tags.add(keys[i]);
+    private static LinkedHashSet getTags(String... keys) {
+        LinkedHashSet tags = new LinkedHashSet<>();
+
+        Collections.addAll(tags, keys);
 
         return tags;
     }
 
+    /**
+     * Get the name of the metric.
+     *
+     * @return the metric name; never null
+     */
     public String name() {
         return this.name;
     }
 
+    /**
+     * Get the name of the group.
+     *
+     * @return the group name; never null
+     */
     public String group() {
         return this.group;
     }
 
+    /**
+     * Get the description of the metric.
+     *
+     * @return the metric description; never null
+     */
     public String description() {
         return this.description;
     }
 
+    /**
+     * Get the set of tag names for the metric.
+     *
+     * @return the ordered set of tag names; never null but possibly empty
+     */
     public Set tags() {
         return tags;
     }
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 8187369ca3c59..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
  */
@@ -29,12 +31,14 @@ public class Node {
     private final int port;
     private final String rack;
 
+    // Cache hashCode as it is called in performance sensitive parts of the code (e.g. RecordAccumulator.ready)
+    private Integer hash;
+
     public Node(int id, String host, int port) {
         this(id, host, port, null);
     }
 
     public Node(int id, String host, int port, String rack) {
-        super();
         this.id = id;
         this.idString = Integer.toString(id);
         this.host = host;
@@ -100,39 +104,30 @@ public String rack() {
 
     @Override
     public int hashCode() {
-        final int prime = 31;
-        int result = 1;
-        result = prime * result + ((host == null) ? 0 : host.hashCode());
-        result = prime * result + id;
-        result = prime * result + port;
-        result = prime * result + ((rack == null) ? 0 : rack.hashCode());
-        return result;
+        Integer h = this.hash;
+        if (h == null) {
+            int result = 31 + ((host == null) ? 0 : host.hashCode());
+            result = 31 * result + id;
+            result = 31 * result + port;
+            result = 31 * result + ((rack == null) ? 0 : rack.hashCode());
+            this.hash = result;
+            return result;
+        } else {
+            return h;
+        }
     }
 
     @Override
     public boolean equals(Object obj) {
         if (this == obj)
             return true;
-        if (obj == null)
-            return false;
-        if (getClass() != obj.getClass())
+        if (obj == null || getClass() != obj.getClass())
             return false;
         Node other = (Node) obj;
-        if (host == null) {
-            if (other.host != null)
-                return false;
-        } else if (!host.equals(other.host))
-            return false;
-        if (id != other.id)
-            return false;
-        if (port != other.port)
-            return false;
-        if (rack == null) {
-            if (other.rack != null)
-                return false;
-        } else if (!rack.equals(other.rack))
-            return false;
-        return true;
+        return id == other.id &&
+            port == other.port &&
+            Objects.equals(host, other.host) &&
+            Objects.equals(rack, other.rack);
     }
 
     @Override
diff --git a/clients/src/main/java/org/apache/kafka/common/PartitionInfo.java b/clients/src/main/java/org/apache/kafka/common/PartitionInfo.java
index 38e4f6788e90f..7ed54f098a3c8 100644
--- a/clients/src/main/java/org/apache/kafka/common/PartitionInfo.java
+++ b/clients/src/main/java/org/apache/kafka/common/PartitionInfo.java
@@ -20,7 +20,6 @@
  * This is used to describe per-partition state in the MetadataResponse.
  */
 public class PartitionInfo {
-
     private final String topic;
     private final int partition;
     private final Node leader;
@@ -33,7 +32,12 @@ public PartitionInfo(String topic, int partition, Node leader, Node[] replicas,
         this(topic, partition, leader, replicas, inSyncReplicas, new Node[0]);
     }
 
-    public PartitionInfo(String topic, int partition, Node leader, Node[] replicas, Node[] inSyncReplicas, Node[] offlineReplicas) {
+    public PartitionInfo(String topic,
+                         int partition,
+                         Node leader,
+                         Node[] replicas,
+                         Node[] inSyncReplicas,
+                         Node[] offlineReplicas) {
         this.topic = topic;
         this.partition = partition;
         this.leader = leader;
@@ -99,10 +103,12 @@ public String toString() {
     /* Extract the node ids from each item in the array and format for display */
     private String formatNodeIds(Node[] nodes) {
         StringBuilder b = new StringBuilder("[");
-        for (int i = 0; i < nodes.length; i++) {
-            b.append(nodes[i].idString());
-            if (i < nodes.length - 1)
-                b.append(',');
+        if (nodes != null) {
+            for (int i = 0; i < nodes.length; i++) {
+                b.append(nodes[i].idString());
+                if (i < nodes.length - 1)
+                    b.append(',');
+            }
         }
         b.append("]");
         return b.toString();
diff --git a/clients/src/main/java/org/apache/kafka/common/Reconfigurable.java b/clients/src/main/java/org/apache/kafka/common/Reconfigurable.java
new file mode 100644
index 0000000000000..8db9dc27ff35b
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/Reconfigurable.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;
+
+import org.apache.kafka.common.config.ConfigException;
+
+import java.util.Map;
+import java.util.Set;
+
+/**
+ * Interface for reconfigurable classes that support dynamic configuration.
+ */
+public interface Reconfigurable extends Configurable {
+
+    /**
+     * Returns the names of configs that may be reconfigured.
+     */
+    Set reconfigurableConfigs();
+
+    /**
+     * Validates the provided configuration. The provided map contains
+     * all configs including any reconfigurable configs that may be different
+     * from the initial configuration. Reconfiguration will be not performed
+     * if this method throws any exception.
+     * @throws ConfigException if the provided configs are not valid. The exception
+     *         message from ConfigException will be returned to the client in
+     *         the AlterConfigs response.
+     */
+    void validateReconfiguration(Map configs) throws ConfigException;
+
+    /**
+     * Reconfigures this instance with the given key-value pairs. The provided
+     * map contains all configs including any reconfigurable configs that
+     * may have changed since the object was initially configured using
+     * {@link Configurable#configure(Map)}. This method will only be invoked if
+     * the configs have passed validation using {@link #validateReconfiguration(Map)}.
+     */
+    void reconfigure(Map configs);
+
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/TopicPartition.java b/clients/src/main/java/org/apache/kafka/common/TopicPartition.java
index dc79c2e13dc90..2c6add78aab31 100644
--- a/clients/src/main/java/org/apache/kafka/common/TopicPartition.java
+++ b/clients/src/main/java/org/apache/kafka/common/TopicPartition.java
@@ -17,11 +17,13 @@
 package org.apache.kafka.common;
 
 import java.io.Serializable;
+import java.util.Objects;
 
 /**
  * A topic name and partition number
  */
 public final class TopicPartition implements Serializable {
+    private static final long serialVersionUID = -613627415771699627L;
 
     private int hash = 0;
     private final int partition;
@@ -47,7 +49,7 @@ public int hashCode() {
         final int prime = 31;
         int result = 1;
         result = prime * result + partition;
-        result = prime * result + ((topic == null) ? 0 : topic.hashCode());
+        result = prime * result + Objects.hashCode(topic);
         this.hash = result;
         return result;
     }
@@ -61,14 +63,7 @@ public boolean equals(Object obj) {
         if (getClass() != obj.getClass())
             return false;
         TopicPartition other = (TopicPartition) obj;
-        if (partition != other.partition)
-            return false;
-        if (topic == null) {
-            if (other.topic != null)
-                return false;
-        } else if (!topic.equals(other.topic))
-            return false;
-        return true;
+        return partition == other.partition && Objects.equals(topic, other.topic);
     }
 
     @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 be693181aa603..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.
@@ -82,4 +83,26 @@ public String toString() {
         return "(partition=" + partition + ", leader=" + leader + ", replicas=" +
             Utils.join(replicas, ", ") + ", isr=" + Utils.join(isr, ", ") + ")";
     }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+
+        TopicPartitionInfo that = (TopicPartitionInfo) o;
+
+        return partition == that.partition &&
+            Objects.equals(leader, that.leader) &&
+            Objects.equals(replicas, that.replicas) &&
+            Objects.equals(isr, that.isr);
+    }
+
+    @Override
+    public int hashCode() {
+        int result = partition;
+        result = 31 * result + (leader != null ? leader.hashCode() : 0);
+        result = 31 * result + (replicas != null ? replicas.hashCode() : 0);
+        result = 31 * result + (isr != null ? isr.hashCode() : 0);
+        return result;
+    }
 }
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 2a10439587fea..0a7c419c933bd 100644
--- a/clients/src/main/java/org/apache/kafka/common/TopicPartitionReplica.java
+++ b/clients/src/main/java/org/apache/kafka/common/TopicPartitionReplica.java
@@ -17,6 +17,7 @@
 package org.apache.kafka.common;
 
 import java.io.Serializable;
+import java.util.Objects;
 
 
 /**
@@ -30,7 +31,7 @@ public final class TopicPartitionReplica implements Serializable {
     private final String topic;
 
     public TopicPartitionReplica(String topic, int partition, int brokerId) {
-        this.topic = topic;
+        this.topic = Objects.requireNonNull(topic);
         this.partition = partition;
         this.brokerId = brokerId;
     }
@@ -54,7 +55,7 @@ public int hashCode() {
         }
         final int prime = 31;
         int result = 1;
-        result = prime * result + ((topic == null) ? 0 : topic.hashCode());
+        result = prime * result + topic.hashCode();
         result = prime * result + partition;
         result = prime * result + brokerId;
         this.hash = result;
@@ -70,18 +71,7 @@ public boolean equals(Object obj) {
         if (getClass() != obj.getClass())
             return false;
         TopicPartitionReplica other = (TopicPartitionReplica) obj;
-        if (partition != other.partition)
-            return false;
-        if (brokerId != other.brokerId)
-            return false;
-        if (topic == null) {
-            if (other.topic != null) {
-                return false;
-            }
-        } else if (!topic.equals(other.topic)) {
-            return false;
-        }
-        return true;
+        return partition == other.partition && brokerId == other.brokerId && topic.equals(other.topic);
     }
 
     @Override
diff --git a/clients/src/main/java/org/apache/kafka/common/Uuid.java b/clients/src/main/java/org/apache/kafka/common/Uuid.java
new file mode 100644
index 0000000000000..65f56a24abdf7
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/Uuid.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.common;
+
+import java.nio.ByteBuffer;
+import java.util.Base64;
+
+/**
+ * This class defines an immutable universally unique identifier (UUID). It represents a 128-bit value.
+ * More specifically, the random UUIDs generated by this class are variant 2 (Leach-Salz) version 4 UUIDs.
+ * This is the same type of UUID as the ones generated by java.util.UUID. The toString() method prints
+ * using the base64 string encoding. Likewise, the fromString method expects a base64 string encoding.
+ */
+public class Uuid {
+
+    private static final java.util.UUID SENTINEL_ID_INTERNAL = new java.util.UUID(0L, 1L);
+
+    /**
+     * A UUID that represents a null or empty UUID. Will never be returned by the randomUuid method.
+     */
+    public static final Uuid ZERO_UUID = new Uuid(0L, 0L);
+    private static final java.util.UUID ZERO_ID_INTERNAL = new java.util.UUID(0L, 0L);
+
+    private final long mostSignificantBits;
+    private final long leastSignificantBits;
+
+    /**
+     * Constructs a 128-bit type 4 UUID where the first long represents the the most significant 64 bits
+     * and the second long represents the least significant 64 bits.
+     */
+    public Uuid(long mostSigBits, long leastSigBits) {
+        this.mostSignificantBits = mostSigBits;
+        this.leastSignificantBits = leastSigBits;
+    }
+
+    /**
+     * Static factory to retrieve a type 4 (pseudo randomly generated) UUID.
+     */
+    public static Uuid randomUuid() {
+        java.util.UUID uuid = java.util.UUID.randomUUID();
+        while (uuid.equals(SENTINEL_ID_INTERNAL) || uuid.equals(ZERO_ID_INTERNAL)) {
+            uuid = java.util.UUID.randomUUID();
+        }
+        return new Uuid(uuid.getMostSignificantBits(), uuid.getLeastSignificantBits());
+    }
+
+    /**
+     * Returns the most significant bits of the UUID's 128 value.
+     */
+    public long getMostSignificantBits() {
+        return this.mostSignificantBits;
+    }
+
+    /**
+     * Returns the least significant bits of the UUID's 128 value.
+     */
+    public long getLeastSignificantBits() {
+        return this.leastSignificantBits;
+    }
+
+    /**
+     * Returns true iff obj is another Uuid represented by the same two long values.
+     */
+    @Override
+    public boolean equals(Object obj) {
+        if ((null == obj) || (obj.getClass() != this.getClass()))
+            return false;
+        Uuid id = (Uuid) obj;
+        return this.mostSignificantBits == id.mostSignificantBits &&
+                this.leastSignificantBits == id.leastSignificantBits;
+    }
+
+    /**
+     * Returns a hash code for this UUID
+     */
+    @Override
+    public int hashCode() {
+        long xor = mostSignificantBits ^ leastSignificantBits;
+        return (int) (xor >> 32) ^ (int) xor;
+    }
+
+    /**
+     * Returns a base64 string encoding of the UUID.
+     */
+    @Override
+    public String toString() {
+        return Base64.getUrlEncoder().withoutPadding().encodeToString(getBytesFromUuid());
+    }
+
+    /**
+     * Creates a UUID based on a base64 string encoding used in the toString() method.
+     */
+    public static Uuid fromString(String str) {
+        ByteBuffer uuidBytes = ByteBuffer.wrap(Base64.getUrlDecoder().decode(str));
+        return new Uuid(uuidBytes.getLong(), uuidBytes.getLong());
+    }
+
+    private byte[] getBytesFromUuid() {
+        // Extract bytes for uuid which is 128 bits (or 16 bytes) long.
+        ByteBuffer uuidBytes = ByteBuffer.wrap(new byte[16]);
+        uuidBytes.putLong(this.mostSignificantBits);
+        uuidBytes.putLong(this.leastSignificantBits);
+        return uuidBytes.array();
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/acl/AccessControlEntryFilter.java b/clients/src/main/java/org/apache/kafka/common/acl/AccessControlEntryFilter.java
index a95303e8608b6..225e73a47da78 100644
--- a/clients/src/main/java/org/apache/kafka/common/acl/AccessControlEntryFilter.java
+++ b/clients/src/main/java/org/apache/kafka/common/acl/AccessControlEntryFilter.java
@@ -109,9 +109,7 @@ public boolean matches(AccessControlEntry other) {
             return false;
         if ((operation() != AclOperation.ANY) && (!operation().equals(other.operation())))
             return false;
-        if ((permissionType() != AclPermissionType.ANY) && (!permissionType().equals(other.permissionType())))
-            return false;
-        return true;
+        return (permissionType() == AclPermissionType.ANY) || (permissionType().equals(other.permissionType()));
     }
 
     /**
diff --git a/clients/src/main/java/org/apache/kafka/common/acl/AclBinding.java b/clients/src/main/java/org/apache/kafka/common/acl/AclBinding.java
index d264ef1a4f627..67dbfc0545c3b 100644
--- a/clients/src/main/java/org/apache/kafka/common/acl/AclBinding.java
+++ b/clients/src/main/java/org/apache/kafka/common/acl/AclBinding.java
@@ -18,49 +18,61 @@
 package org.apache.kafka.common.acl;
 
 import org.apache.kafka.common.annotation.InterfaceStability;
+import org.apache.kafka.common.resource.PatternType;
 import org.apache.kafka.common.resource.Resource;
+import org.apache.kafka.common.resource.ResourcePattern;
 
 import java.util.Objects;
 
 /**
- * Represents a binding between a resource and an access control entry.
+ * Represents a binding between a resource pattern and an access control entry.
  *
  * The API for this class is still evolving and we may break compatibility in minor releases, if necessary.
  */
 @InterfaceStability.Evolving
 public class AclBinding {
-    private final Resource resource;
+    private final ResourcePattern pattern;
     private final AccessControlEntry entry;
 
+    /**
+     * Create an instance of this class with the provided parameters.
+     *
+     * @param pattern non-null resource pattern.
+     * @param entry non-null entry
+     */
+    public AclBinding(ResourcePattern pattern, AccessControlEntry entry) {
+        this.pattern = Objects.requireNonNull(pattern, "pattern");
+        this.entry = Objects.requireNonNull(entry, "entry");
+    }
+
     /**
      * Create an instance of this class with the provided parameters.
      *
      * @param resource non-null resource
      * @param entry non-null entry
+     * @deprecated Since 2.0. Use {@link #AclBinding(ResourcePattern, AccessControlEntry)}
      */
+    @Deprecated
     public AclBinding(Resource resource, AccessControlEntry entry) {
-        Objects.requireNonNull(resource);
-        this.resource = resource;
-        Objects.requireNonNull(entry);
-        this.entry = entry;
+        this(new ResourcePattern(resource.resourceType(), resource.name(), PatternType.LITERAL), entry);
     }
 
     /**
-     * Return true if this binding has any UNKNOWN components.
+     * @return true if this binding has any UNKNOWN components.
      */
     public boolean isUnknown() {
-        return resource.isUnknown() || entry.isUnknown();
+        return pattern.isUnknown() || entry.isUnknown();
     }
 
     /**
-     * Return the resource for this binding.
+     * @return the resource pattern for this binding.
      */
-    public Resource resource() {
-        return resource;
+    public ResourcePattern pattern() {
+        return pattern;
     }
 
     /**
-     * Return the access control entry for this binding.
+     * @return the access control entry for this binding.
      */
     public final AccessControlEntry entry() {
         return entry;
@@ -70,24 +82,25 @@ public final AccessControlEntry entry() {
      * Create a filter which matches only this AclBinding.
      */
     public AclBindingFilter toFilter() {
-        return new AclBindingFilter(resource.toFilter(), entry.toFilter());
+        return new AclBindingFilter(pattern.toFilter(), entry.toFilter());
     }
 
     @Override
     public String toString() {
-        return "(resource=" + resource + ", entry=" + entry + ")";
+        return "(pattern=" + pattern + ", entry=" + entry + ")";
     }
 
     @Override
     public boolean equals(Object o) {
-        if (!(o instanceof AclBinding))
-            return false;
-        AclBinding other = (AclBinding) o;
-        return resource.equals(other.resource) && entry.equals(other.entry);
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        AclBinding that = (AclBinding) o;
+        return Objects.equals(pattern, that.pattern) &&
+            Objects.equals(entry, that.entry);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(resource, entry);
+        return Objects.hash(pattern, entry);
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/acl/AclBindingFilter.java b/clients/src/main/java/org/apache/kafka/common/acl/AclBindingFilter.java
index 64f16cd746005..3168ec61cfcef 100644
--- a/clients/src/main/java/org/apache/kafka/common/acl/AclBindingFilter.java
+++ b/clients/src/main/java/org/apache/kafka/common/acl/AclBindingFilter.java
@@ -18,8 +18,9 @@
 package org.apache.kafka.common.acl;
 
 import org.apache.kafka.common.annotation.InterfaceStability;
+import org.apache.kafka.common.resource.PatternType;
 import org.apache.kafka.common.resource.ResourceFilter;
-import org.apache.kafka.common.resource.ResourceType;
+import org.apache.kafka.common.resource.ResourcePatternFilter;
 
 import java.util.Objects;
 
@@ -30,45 +31,53 @@
  */
 @InterfaceStability.Evolving
 public class AclBindingFilter {
-    private final ResourceFilter resourceFilter;
+    private final ResourcePatternFilter patternFilter;
     private final AccessControlEntryFilter entryFilter;
 
     /**
      * A filter which matches any ACL binding.
      */
-    public static final AclBindingFilter ANY = new AclBindingFilter(
-        new ResourceFilter(ResourceType.ANY, null),
-        new AccessControlEntryFilter(null, null, AclOperation.ANY, AclPermissionType.ANY));
+    public static final AclBindingFilter ANY = new AclBindingFilter(ResourcePatternFilter.ANY, AccessControlEntryFilter.ANY);
+
+    /**
+     * Create an instance of this filter with the provided parameters.
+     *
+     * @param patternFilter non-null pattern filter
+     * @param entryFilter non-null access control entry filter
+     */
+    public AclBindingFilter(ResourcePatternFilter patternFilter, AccessControlEntryFilter entryFilter) {
+        this.patternFilter = Objects.requireNonNull(patternFilter, "patternFilter");
+        this.entryFilter = Objects.requireNonNull(entryFilter, "entryFilter");
+    }
 
     /**
      * Create an instance of this filter with the provided parameters.
      *
      * @param resourceFilter non-null resource filter
      * @param entryFilter non-null access control entry filter
+     * @deprecated Since 2.0. Use {@link #AclBindingFilter(ResourcePatternFilter, AccessControlEntryFilter)}
      */
+    @Deprecated
     public AclBindingFilter(ResourceFilter resourceFilter, AccessControlEntryFilter entryFilter) {
-        Objects.requireNonNull(resourceFilter);
-        this.resourceFilter = resourceFilter;
-        Objects.requireNonNull(entryFilter);
-        this.entryFilter = entryFilter;
+        this(new ResourcePatternFilter(resourceFilter.resourceType(), resourceFilter.name(), PatternType.LITERAL), entryFilter);
     }
 
     /**
-     * Return true if this filter has any UNKNOWN components.
+     * @return {@code true} if this filter has any UNKNOWN components.
      */
     public boolean isUnknown() {
-        return resourceFilter.isUnknown() || entryFilter.isUnknown();
+        return patternFilter.isUnknown() || entryFilter.isUnknown();
     }
 
     /**
-     * Return the resource filter.
+     * @return the resource pattern filter.
      */
-    public ResourceFilter resourceFilter() {
-        return resourceFilter;
+    public ResourcePatternFilter patternFilter() {
+        return patternFilter;
     }
 
     /**
-     * Return the access control entry filter.
+     * @return the access control entry filter.
      */
     public final AccessControlEntryFilter entryFilter() {
         return entryFilter;
@@ -76,15 +85,16 @@ public final AccessControlEntryFilter entryFilter() {
 
     @Override
     public String toString() {
-        return "(resourceFilter=" + resourceFilter + ", entryFilter=" + entryFilter + ")";
+        return "(patternFilter=" + patternFilter + ", entryFilter=" + entryFilter + ")";
     }
 
     @Override
     public boolean equals(Object o) {
-        if (!(o instanceof AclBindingFilter))
-            return false;
-        AclBindingFilter other = (AclBindingFilter) o;
-        return resourceFilter.equals(other.resourceFilter) && entryFilter.equals(other.entryFilter);
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        AclBindingFilter that = (AclBindingFilter) o;
+        return Objects.equals(patternFilter, that.patternFilter) &&
+            Objects.equals(entryFilter, that.entryFilter);
     }
 
     /**
@@ -92,14 +102,14 @@ public boolean equals(Object o) {
      * there are no ANY or UNKNOWN fields.
      */
     public boolean matchesAtMostOne() {
-        return resourceFilter.matchesAtMostOne() && entryFilter.matchesAtMostOne();
+        return patternFilter.matchesAtMostOne() && entryFilter.matchesAtMostOne();
     }
 
     /**
      * Return a string describing an ANY or UNKNOWN field, or null if there is no such field.
      */
     public String findIndefiniteField() {
-        String indefinite = resourceFilter.findIndefiniteField();
+        String indefinite = patternFilter.findIndefiniteField();
         if (indefinite != null)
             return indefinite;
         return entryFilter.findIndefiniteField();
@@ -109,11 +119,11 @@ public String findIndefiniteField() {
      * Return true if the resource filter matches the binding's resource and the entry filter matches binding's entry.
      */
     public boolean matches(AclBinding binding) {
-        return resourceFilter.matches(binding.resource()) && entryFilter.matches(binding.entry());
+        return patternFilter.matches(binding.pattern()) && entryFilter.matches(binding.entry());
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(resourceFilter, entryFilter);
+        return Objects.hash(patternFilter, entryFilter);
     }
 }
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 3da18cf943254..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
@@ -25,19 +25,19 @@
 /**
  * Represents an operation which an ACL grants or denies permission to perform.
  *
- * Some operations imply other operations.
+ * Some operations imply other operations:
+ * 
          
+ * 
        • ALLOW ALL implies ALLOW everything
+ * 
        • DENY ALL implies DENY everything
  *
- * ALLOW ALL implies ALLOW everything
- * DENY ALL implies DENY everything
+ * 
        • ALLOW READ implies ALLOW DESCRIBE
+ * 
        • ALLOW WRITE implies ALLOW DESCRIBE
+ * 
        • ALLOW DELETE implies ALLOW DESCRIBE
  *
- * ALLOW READ implies ALLOW DESCRIBE
- * ALLOW WRITE implies ALLOW DESCRIBE
- * ALLOW DELETE implies ALLOW DESCRIBE
- *
- * ALLOW ALTER implies ALLOW DESCRIBE
- *
- * ALLOW ALTER_CONFIGS implies ALLOW DESCRIBE_CONFIGS
+ * 
        • ALLOW ALTER implies ALLOW DESCRIBE
  *
+ * 
        • ALLOW ALTER_CONFIGS implies ALLOW DESCRIBE_CONFIGS
+ * 
        
  * The API for this class is still evolving and we may break compatibility in minor releases, if necessary.
  */
 @InterfaceStability.Evolving
@@ -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/cache/LRUCache.java b/clients/src/main/java/org/apache/kafka/common/cache/LRUCache.java
index bdc67ac2732a3..672cb65d66ab6 100644
--- a/clients/src/main/java/org/apache/kafka/common/cache/LRUCache.java
+++ b/clients/src/main/java/org/apache/kafka/common/cache/LRUCache.java
@@ -29,7 +29,7 @@ public LRUCache(final int maxSize) {
         cache = new LinkedHashMap(16, .75f, true) {
             @Override
             protected boolean removeEldestEntry(Map.Entry eldest) {
-                return size() > maxSize;
+                return this.size() > maxSize; // require this. prefix to make lgtm.com happy
             }
         };
     }
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 dc80d98b03dbe..a25c3ea521d0c 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,27 +53,94 @@ 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);
+        this.used = Collections.synchronizedSet(new HashSet<>());
         Map configUpdates = postProcessParsedConfig(Collections.unmodifiableMap(this.values));
         for (Map.Entry update : configUpdates.entrySet()) {
             this.values.put(update.getKey(), update.getValue());
         }
         definition.parse(this.values);
-        this.used = Collections.synchronizedSet(new HashSet());
         this.definition = definition;
         if (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);
+
     }
 
     /**
@@ -133,6 +201,13 @@ public ConfigDef.Type typeOf(String key) {
         return configKey.type;
     }
 
+    public String documentationOf(String key) {
+        ConfigDef.ConfigKey configKey = definition.configKeys().get(key);
+        if (configKey == null)
+            return null;
+        return configKey.documentation;
+    }
+
     public Password getPassword(String key) {
         return (Password) get(key);
     }
@@ -153,6 +228,13 @@ public Map originals() {
         return copy;
     }
 
+    public Map originals(Map configOverrides) {
+        Map copy = new RecordingMap<>();
+        copy.putAll(originals);
+        copy.putAll(configOverrides);
+        return copy;
+    }
+
     /**
      * Get all the original settings, ensuring that all values are of type String.
      * @return the original settings
@@ -176,10 +258,25 @@ public Map originalsStrings() {
      * @return a Map containing the settings with the prefix
      */
     public Map originalsWithPrefix(String prefix) {
+        return originalsWithPrefix(prefix, true);
+    }
+
+    /**
+     * Gets all original settings with the given prefix.
+     *
+     * @param prefix the prefix to use as a filter
+     * @param strip strip the prefix before adding to the output if set true
+     * @return a Map containing the settings with the prefix
+     */
+    public Map originalsWithPrefix(String prefix, boolean strip) {
         Map result = new RecordingMap<>(prefix, false);
         for (Map.Entry entry : originals.entrySet()) {
-            if (entry.getKey().startsWith(prefix) && entry.getKey().length() > prefix.length())
-                result.put(entry.getKey().substring(prefix.length()), entry.getValue());
+            if (entry.getKey().startsWith(prefix) && entry.getKey().length() > prefix.length()) {
+                if (strip)
+                    result.put(entry.getKey().substring(prefix.length()), entry.getValue());
+                else
+                    result.put(entry.getKey(), entry.getValue());
+            }
         }
         return result;
     }
@@ -189,6 +286,16 @@ public Map originalsWithPrefix(String prefix) {
      * put all the remaining keys with the prefix stripped and their parsed values in the result map.
      *
      * This is useful if one wants to allow prefixed configs to override default ones.
+     * 
        
+     * Two forms of prefixes are supported:
+     * 
          
+     *     
        • listener.name.{listenerName}.some.prop: If the provided prefix is `listener.name.{listenerName}.`,
+     *         the key `some.prop` with the value parsed using the definition of `some.prop` is returned.
          • 
+     *     
        • listener.name.{listenerName}.{mechanism}.some.prop: If the provided prefix is `listener.name.{listenerName}.`,
+     *         the key `{mechanism}.some.prop` with the value parsed using the definition of `some.prop` is returned.
+     *          This is used to provide per-mechanism configs for a broker listener (e.g sasl.jaas.config)
          • 
+     * 
        
+     * 
        
      */
     public Map valuesWithPrefixOverride(String prefix) {
         Map result = new RecordingMap<>(values(), prefix, true);
@@ -198,15 +305,57 @@ public Map valuesWithPrefixOverride(String prefix) {
                 ConfigDef.ConfigKey configKey = definition.configKeys().get(keyWithNoPrefix);
                 if (configKey != null)
                     result.put(keyWithNoPrefix, definition.parseValue(configKey, entry.getValue(), true));
+                else {
+                    String keyWithNoSecondaryPrefix = keyWithNoPrefix.substring(keyWithNoPrefix.indexOf('.') + 1);
+                    configKey = definition.configKeys().get(keyWithNoSecondaryPrefix);
+                    if (configKey != null)
+                        result.put(keyWithNoPrefix, definition.parseValue(configKey, entry.getValue(), true));
+                }
             }
         }
         return result;
     }
 
+    /**
+     * If at least one key with {@code prefix} exists, all prefixed values will be parsed and put into map.
+     * If no value with {@code prefix} exists all unprefixed values will be returned.
+     *
+     * This is useful if one wants to allow prefixed configs to override default ones, but wants to use either
+     * only prefixed configs or only regular configs, but not mix them.
+     */
+    public Map valuesWithPrefixAllOrNothing(String prefix) {
+        Map withPrefix = originalsWithPrefix(prefix, true);
+
+        if (withPrefix.isEmpty()) {
+            return new RecordingMap<>(values(), "", true);
+        } else {
+            Map result = new RecordingMap<>(prefix, true);
+
+            for (Map.Entry entry : withPrefix.entrySet()) {
+                ConfigDef.ConfigKey configKey = definition.configKeys().get(entry.getKey());
+                if (configKey != null)
+                    result.put(entry.getKey(), definition.parseValue(configKey, entry.getValue(), true));
+            }
+
+            return result;
+        }
+    }
+
     public Map values() {
         return new RecordingMap<>(values);
     }
 
+    public Map nonInternalValues() {
+        Map nonInternalConfigs = new RecordingMap<>();
+        values.forEach((key, value) -> {
+            ConfigDef.ConfigKey configKey = definition.configKeys().get(key);
+            if (configKey == null || !configKey.internalConfig) {
+                nonInternalConfigs.put(key, value);
+            }
+        });
+        return nonInternalConfigs;
+    }
+
     private void logAll() {
         StringBuilder b = new StringBuilder();
         b.append(getClass().getSimpleName());
@@ -231,6 +380,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.
@@ -240,15 +412,22 @@ public void logUnused() {
      * @return A configured instance of the class
      */
     public  T getConfiguredInstance(String key, Class t) {
+        return getConfiguredInstance(key, t, Collections.emptyMap());
+    }
+
+    /**
+     * Get a configured instance of the give class specified by the given configuration key. If the object implements
+     * Configurable configure it using the configuration.
+     *
+     * @param key The configuration key for the class
+     * @param t The interface the class should implement
+     * @param configOverrides override origin configs
+     * @return A configured instance of the class
+     */
+    public  T getConfiguredInstance(String key, Class t, Map configOverrides) {
         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(configOverrides));
     }
 
     /**
@@ -260,7 +439,7 @@ public  T getConfiguredInstance(String key, Class t) {
      * @return The list of configured instances
      */
     public  List getConfiguredInstances(String key, Class t) {
-        return getConfiguredInstances(key, t, Collections.emptyMap());
+        return getConfiguredInstances(key, t, Collections.emptyMap());
     }
 
     /**
@@ -273,33 +452,134 @@ public  List getConfiguredInstances(String key, Class t) {
      * @return The list of configured instances
      */
     public  List getConfiguredInstances(String key, Class t, Map configOverrides) {
-        List klasses = getList(key);
-        List objects = new ArrayList();
-        if (klasses == null)
+        return getConfiguredInstances(getList(key), t, configOverrides);
+    }
+
+    /**
+     * Get a list of configured instances of the given class specified by the given configuration key. The configuration
+     * may specify either null or an empty string to indicate no configured instances. In both cases, this method
+     * returns an empty list to indicate no configured instances.
+     * @param classNames The list of class names of the instances to create
+     * @param t The interface the class should implement
+     * @param configOverrides Configuration overrides to use.
+     * @return The list of configured instances
+     */
+    public  List getConfiguredInstances(List classNames, Class t, Map configOverrides) {
+        List objects = new ArrayList<>();
+        if (classNames == null)
             return objects;
         Map configPairs = originals();
         configPairs.putAll(configOverrides);
-        for (Object klass : klasses) {
-            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);
+        for (Object klass : classNames) {
+            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;
@@ -360,4 +640,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/ConfigChangeCallback.java b/clients/src/main/java/org/apache/kafka/common/config/ConfigChangeCallback.java
new file mode 100644
index 0000000000000..faa7d3d87074a
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/config/ConfigChangeCallback.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.common.config;
+
+import org.apache.kafka.common.config.provider.ConfigProvider;
+
+/**
+ * A callback passed to {@link ConfigProvider} for subscribing to changes.
+ */
+public interface ConfigChangeCallback {
+
+    /**
+     * Performs an action when configuration data changes.
+     *
+     * @param path the path at which the data resides
+     * @param data the configuration data
+     */
+    void onChange(String path, ConfigData data);
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/config/ConfigData.java b/clients/src/main/java/org/apache/kafka/common/config/ConfigData.java
new file mode 100644
index 0000000000000..8661ee16cbab1
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/config/ConfigData.java
@@ -0,0 +1,68 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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 org.apache.kafka.common.config.provider.ConfigProvider;
+
+import java.util.Map;
+
+/**
+ * Configuration data from a {@link ConfigProvider}.
+ */
+public class ConfigData {
+
+    private final Map data;
+    private final Long ttl;
+
+    /**
+     * Creates a new ConfigData with the given data and TTL (in milliseconds).
+     *
+     * @param data a Map of key-value pairs
+     * @param ttl the time-to-live of the data in milliseconds, or null if there is no TTL
+     */
+    public ConfigData(Map data, Long ttl) {
+        this.data = data;
+        this.ttl = ttl;
+    }
+
+    /**
+     * Creates a new ConfigData with the given data.
+     *
+     * @param data a Map of key-value pairs
+     */
+    public ConfigData(Map data) {
+        this(data, null);
+    }
+
+    /**
+     * Returns the data.
+     *
+     * @return data a Map of key-value pairs
+     */
+    public Map data() {
+        return data;
+    }
+
+    /**
+     * Returns the TTL (in milliseconds).
+     *
+     * @return ttl the time-to-live (in milliseconds) of the data, or null if there is no TTL
+     */
+    public Long ttl() {
+        return ttl;
+    }
+}
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 7b9881fc5a443..0df9335e7c9c4 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,22 +16,25 @@
  */
 package org.apache.kafka.common.config;
 
+import java.util.function.Function;
+import java.util.stream.Collectors;
 import org.apache.kafka.common.config.types.Password;
 import org.apache.kafka.common.utils.Utils;
 
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
-import java.util.Comparator;
 import java.util.HashMap;
 import java.util.HashSet;
-import java.util.Iterator;
 import java.util.LinkedHashMap;
 import java.util.LinkedList;
 import java.util.List;
 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;
 
 /**
  * This class is used for specifying the set of expected configurations. For each configuration, you can specify
@@ -54,10 +57,10 @@
  * 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");
- * 
+ *
  * Map<String, Object> configs = defs.parse(props);
  * // will return "some value"
  * String someConfig = (String) configs.get("config_with_default");
@@ -73,6 +76,9 @@
  * functionality for accessing configs.
  */
 public class ConfigDef {
+
+    private static final Pattern COMMA_WITH_WHITESPACE = Pattern.compile("\\s*,\\s*");
+
     /**
      * A unique Java object which represents the lack of a default value.
      */
@@ -105,6 +111,15 @@ public Set names() {
         return Collections.unmodifiableSet(configKeys.keySet());
     }
 
+    public Map defaultValues() {
+        Map defaultValues = new HashMap<>();
+        for (ConfigKey key : configKeys.values()) {
+            if (key.defaultValue != NO_DEFAULT_VALUE)
+                defaultValues.put(key.name, key.defaultValue);
+        }
+        return defaultValues;
+    }
+
     public ConfigDef define(ConfigKey key) {
         if (configKeys.containsKey(key.name)) {
             throw new ConfigException("Configuration " + key.name + " is defined twice.");
@@ -174,7 +189,7 @@ public ConfigDef define(String name, Type type, Object defaultValue, Validator v
      */
     public ConfigDef define(String name, Type type, Object defaultValue, Validator validator, Importance importance, String documentation,
                             String group, int orderInGroup, Width width, String displayName, Recommender recommender) {
-        return define(name, type, defaultValue, validator, importance, documentation, group, orderInGroup, width, displayName, Collections.emptyList(), recommender);
+        return define(name, type, defaultValue, validator, importance, documentation, group, orderInGroup, width, displayName, Collections.emptyList(), recommender);
     }
 
     /**
@@ -251,7 +266,7 @@ public ConfigDef define(String name, Type type, Object defaultValue, Importance
      */
     public ConfigDef define(String name, Type type, Object defaultValue, Importance importance, String documentation,
                             String group, int orderInGroup, Width width, String displayName, Recommender recommender) {
-        return define(name, type, defaultValue, null, importance, documentation, group, orderInGroup, width, displayName, Collections.emptyList(), recommender);
+        return define(name, type, defaultValue, null, importance, documentation, group, orderInGroup, width, displayName, Collections.emptyList(), recommender);
     }
 
     /**
@@ -324,7 +339,7 @@ public ConfigDef define(String name, Type type, Importance importance, String do
      */
     public ConfigDef define(String name, Type type, Importance importance, String documentation, String group, int orderInGroup,
                             Width width, String displayName, Recommender recommender) {
-        return define(name, type, NO_DEFAULT_VALUE, null, importance, documentation, group, orderInGroup, width, displayName, Collections.emptyList(), recommender);
+        return define(name, type, NO_DEFAULT_VALUE, null, importance, documentation, group, orderInGroup, width, displayName, Collections.emptyList(), recommender);
     }
 
     /**
@@ -586,23 +601,15 @@ private void validate(String name, Map parsed, Map recommendedValues = key.recommender.validValues(name, parsed);
                 List originalRecommendedValues = value.recommendedValues();
                 if (!originalRecommendedValues.isEmpty()) {
                     Set originalRecommendedValueSet = new HashSet<>(originalRecommendedValues);
-                    Iterator it = recommendedValues.iterator();
-                    while (it.hasNext()) {
-                        Object o = it.next();
-                        if (!originalRecommendedValueSet.contains(o)) {
-                            it.remove();
-                        }
-                    }
+                    recommendedValues.removeIf(o -> !originalRecommendedValueSet.contains(o));
                 }
                 value.recommendedValues(recommendedValues);
                 value.visible(key.recommender.visible(name, parsed));
@@ -696,15 +703,22 @@ else if (value instanceof String)
                         if (trimmed.isEmpty())
                             return Collections.emptyList();
                         else
-                            return Arrays.asList(trimmed.split("\\s*,\\s*", -1));
+                            return Arrays.asList(COMMA_WITH_WHITESPACE.split(trimmed, -1));
                     else
                         throw new ConfigException(name, value, "Expected a comma separated list.");
                 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.");
@@ -745,6 +759,30 @@ public static String convertToString(Object parsedValue, Type type) {
         }
     }
 
+    /**
+     * Converts a map of config (key, value) pairs to a map of strings where each value
+     * is converted to a string. This method should be used with care since it stores
+     * actual password values to String. Values from this map should never be used in log entries.
+     */
+    public static  Map convertToStringMapWithPasswordValues(Map configs) {
+        Map result = new HashMap<>();
+        for (Map.Entry entry : configs.entrySet()) {
+            Object value = entry.getValue();
+            String strValue;
+            if (value instanceof Password)
+                strValue = ((Password) value).value();
+            else if (value instanceof List)
+                strValue = convertToString(value, Type.LIST);
+            else if (value instanceof Class)
+                strValue = convertToString(value, Type.CLASS);
+            else
+                strValue = convertToString(value, null);
+            if (strValue != null)
+                result.put(entry.getKey(), strValue);
+        }
+        return result;
+    }
+
     /**
      * The config types
      */
@@ -812,6 +850,11 @@ public static class Range implements Validator {
         private final Number min;
         private final Number max;
 
+        /**
+         *  A numeric range with inclusive upper bound and inclusive lower bound
+         * @param min  the lower bound
+         * @param max  the upper bound
+         */
         private Range(Number min, Number max) {
             this.min = min;
             this.max = max;
@@ -827,7 +870,7 @@ public static Range atLeast(Number min) {
         }
 
         /**
-         * A numeric range that checks both the upper and lower bound
+         * A numeric range that checks both the upper (inclusive) and lower bound
          */
         public static Range between(Number min, Number max) {
             return new Range(min, max);
@@ -844,7 +887,9 @@ public void ensureValid(String name, Object o) {
         }
 
         public String toString() {
-            if (min == null)
+            if (min == null && max == null)
+                return "[...]";
+            else if (min == null)
                 return "[...," + max + "]";
             else if (max == null)
                 return "[" + min + ",...]";
@@ -904,6 +949,105 @@ 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) {
+            if (value == null) {
+                // Pass in the string null to avoid the spotbugs warning
+                throw new ConfigException(name, "null", "entry must be non null");
+            }
+        }
+
+        public String toString() {
+            return "non-null string";
+        }
+    }
+
+    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;
+
+        private CompositeValidator(List validators) {
+            this.validators = Collections.unmodifiableList(validators);
+        }
+
+        public static CompositeValidator of(Validator... validators) {
+            return new CompositeValidator(Arrays.asList(validators));
+        }
+
+        @Override
+        public void ensureValid(String name, Object value) {
+            for (Validator validator: validators) {
+                validator.ensureValid(name, value);
+            }
+        }
+
+        @Override
+        public String toString() {
+            if (validators == null) return "";
+            StringBuilder desc = new StringBuilder();
+            for (Validator v: validators) {
+                if (desc.length() > 0) {
+                    desc.append(',').append(' ');
+                }
+                desc.append(v);
+            }
+            return desc.toString();
+        }
+    }
+
     public static class NonEmptyString implements Validator {
 
         @Override
@@ -920,6 +1064,44 @@ public String toString() {
         }
     }
 
+    public static class NonEmptyStringWithoutControlChars implements Validator {
+
+        public static NonEmptyStringWithoutControlChars nonEmptyStringWithoutControlChars() {
+            return new NonEmptyStringWithoutControlChars();
+        }
+
+        @Override
+        public void ensureValid(String name, Object value) {
+            String s = (String) value;
+
+            if (s == null) {
+                // This can happen during creation of the config object due to no default value being defined for the
+                // name configuration - a missing name parameter is caught when checking for mandatory parameters,
+                // thus we can ok a null value here
+                return;
+            } else if (s.isEmpty()) {
+                throw new ConfigException(name, value, "String may not be empty");
+            }
+
+            // Check name string for illegal characters
+            ArrayList foundIllegalCharacters = new ArrayList<>();
+
+            for (int i = 0; i < s.length(); i++) {
+                if (Character.isISOControl(s.codePointAt(i))) {
+                    foundIllegalCharacters.add(s.codePointAt(i));
+                }
+            }
+
+            if (!foundIllegalCharacters.isEmpty()) {
+                throw new ConfigException(name, value, "String may not contain control sequences but had the following ASCII chars: " + Utils.join(foundIllegalCharacters, ", "));
+            }
+        }
+
+        public String toString() {
+            return "non-empty string without ISO control characters";
+        }
+    }
+
     public static class ConfigKey {
         public final String name;
         public final Type type;
@@ -960,6 +1142,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() {
@@ -981,8 +1167,15 @@ protected String getConfigValue(ConfigKey key, String headerName) {
                     String defaultValueStr = convertToString(key.defaultValue, key.type);
                     if (defaultValueStr.isEmpty())
                         return "\"\"";
-                    else
-                        return defaultValueStr;
+                    else {
+                        String suffix = "";
+                        if (key.name.endsWith(".bytes")) {
+                            suffix = niceMemoryUnits(((Number) key.defaultValue).longValue());
+                        } else if (key.name.endsWith(".ms")) {
+                            suffix = niceTimeUnits(((Number) key.defaultValue).longValue());
+                        }
+                        return defaultValueStr + suffix;
+                    }
                 } else
                     return "";
             case "Valid Values":
@@ -994,17 +1187,85 @@ protected String getConfigValue(ConfigKey key, String headerName) {
         }
     }
 
+    static String niceMemoryUnits(long bytes) {
+        long value = bytes;
+        int i = 0;
+        while (value != 0 && i < 4) {
+            if (value % 1024L == 0) {
+                value /= 1024L;
+                i++;
+            } else {
+                break;
+            }
+        }
+        switch (i) {
+            case 1:
+                return " (" + value + " kibibyte" + (value == 1 ? ")" : "s)");
+            case 2:
+                return " (" + value + " mebibyte" + (value == 1 ? ")" : "s)");
+            case 3:
+                return " (" + value + " gibibyte" + (value == 1 ? ")" : "s)");
+            case 4:
+                return " (" + value + " tebibyte" + (value == 1 ? ")" : "s)");
+            default:
+                return "";
+        }
+    }
+
+    static String niceTimeUnits(long millis) {
+        long value = millis;
+        long[] divisors = {1000, 60, 60, 24};
+        String[] units = {"second", "minute", "hour", "day"};
+        int i = 0;
+        while (value != 0 && i < 4) {
+            if (value % divisors[i] == 0) {
+                value /= divisors[i];
+                i++;
+            } else {
+                break;
+            }
+        }
+        if (i > 0) {
+            return " (" + value + " " + units[i - 1] + (value > 1 ? "s)" : ")");
+        }
+        return "";
+    }
+
     public String toHtmlTable() {
+        return toHtmlTable(Collections.emptyMap());
+    }
+
+    private void addHeader(StringBuilder builder, String headerName) {
+        builder.append("");
+        builder.append(headerName);
+        builder.append("\n");
+    }
+
+    private void addColumnValue(StringBuilder builder, String value) {
+        builder.append("");
+        builder.append(value);
+        builder.append("");
+    }
+
+    /**
+     * Converts this config into an HTML table that can be embedded into docs.
+     * 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
+     */
+    public String toHtmlTable(Map dynamicUpdateModes) {
+        boolean hasUpdateModes = !dynamicUpdateModes.isEmpty();
         List configs = sortedConfigs();
         StringBuilder b = new StringBuilder();
         b.append("\n");
         b.append("\n");
         // print column headers
         for (String headerName : headers()) {
-            b.append("\n");
+            addHeader(b, headerName);
         }
+        if (hasUpdateModes)
+            addHeader(b, "Dynamic Update Mode");
         b.append("\n");
         for (ConfigKey key : configs) {
             if (key.internalConfig) {
@@ -1013,10 +1274,15 @@ public String toHtmlTable() {
             b.append("\n");
             // print column values
             for (String headerName : headers()) {
-                b.append("");
             }
+            if (hasUpdateModes) {
+                String updateMode = dynamicUpdateModes.get(key.name);
+                if (updateMode == null)
+                    updateMode = "read-only";
+                addColumnValue(b, updateMode);
+            }
             b.append("\n");
         }
         b.append("
        ");
-            b.append(headerName);
-            b.append("
        ");
-                b.append(getConfigValue(key, headerName));
+                addColumnValue(b, getConfigValue(key, headerName));
                 b.append("
        ");
@@ -1116,32 +1382,30 @@ private List sortedConfigs() {
         }
 
         List configs = new ArrayList<>(configKeys.values());
-        Collections.sort(configs, new Comparator() {
-            @Override
-            public int compare(ConfigKey k1, ConfigKey k2) {
-                int cmp = k1.group == null
-                        ? (k2.group == null ? 0 : -1)
-                        : (k2.group == null ? 1 : Integer.compare(groupOrd.get(k1.group), groupOrd.get(k2.group)));
-                if (cmp == 0) {
-                    cmp = Integer.compare(k1.orderInGroup, k2.orderInGroup);
-                    if (cmp == 0) {
-                        // first take anything with no default value
-                        if (!k1.hasDefault() && k2.hasDefault()) {
-                            cmp = -1;
-                        } else if (!k2.hasDefault() && k1.hasDefault()) {
-                            cmp = 1;
-                        } else {
-                            cmp = k1.importance.compareTo(k2.importance);
-                            if (cmp == 0) {
-                                return k1.name.compareTo(k2.name);
-                            }
-                        }
-                    }
+        Collections.sort(configs, (k1, k2) -> compare(k1, k2, groupOrd));
+        return configs;
+    }
+
+    private int compare(ConfigKey k1, ConfigKey k2, Map groupOrd) {
+        int cmp = k1.group == null
+            ? (k2.group == null ? 0 : -1)
+            : (k2.group == null ? 1 : Integer.compare(groupOrd.get(k1.group), groupOrd.get(k2.group)));
+        if (cmp == 0) {
+            cmp = Integer.compare(k1.orderInGroup, k2.orderInGroup);
+            if (cmp == 0) {
+                // first take anything with no default value
+                if (!k1.hasDefault() && k2.hasDefault())
+                    cmp = -1;
+                else if (!k2.hasDefault() && k1.hasDefault())
+                    cmp = 1;
+                else {
+                    cmp = k1.importance.compareTo(k2.importance);
+                    if (cmp == 0)
+                        return k1.name.compareTo(k2.name);
                 }
-                return cmp;
             }
-        });
-        return configs;
+        }
+        return cmp;
     }
 
     public void embed(final String keyPrefix, final String groupPrefix, final int startingOrd, final ConfigDef child) {
@@ -1169,11 +1433,15 @@ 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 new ConfigDef.Validator() {
-            @Override
+        return new Validator() {
             public void ensureValid(String name, Object value) {
                 base.ensureValid(name.substring(keyPrefix.length()), value);
             }
+
+            @Override
+            public String toString() {
+                return base.toString();
+            }
         };
     }
 
@@ -1221,4 +1489,81 @@ 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.
+     * @param headerDepth The top level header depth in the generated HTML.
+     * @param idGenerator A function for computing the HTML id attribute in the generated HTML from a given config name.
+     */
+    public String toHtml(int headerDepth, Function idGenerator) {
+        return toHtml(headerDepth, idGenerator, 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) {
+        return toHtml(4, Function.identity(), dynamicUpdateModes);
+    }
+
+    /**
+     * 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 headerDepth The top level header depth in the generated HTML.
+     * @param idGenerator A function for computing the HTML id attribute in the generated HTML from a given config name.
+     * @param dynamicUpdateModes Config name -> update mode mapping.
+     */
+    public String toHtml(int headerDepth, Function idGenerator,
+                         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("
        • \n");
+            b.append(String.format("" +
+                    "%3$s" +
+                    "%n", headerDepth, idGenerator.apply(key.name), key.name));
+            b.append("
          ");
+            b.append(key.documentation.replaceAll("\n", "
          "));
+            b.append("
          \n");
+
+            b.append("" +
+                    "\n");
+            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("
          \n");
+            b.append("
          • \n");
+        }
+        b.append("
        \n");
+        return b.toString();
+    }
+
+    private static void addConfigDetail(StringBuilder builder, String name, String value) {
+        builder.append("" +
+                "" + name + ":" +
+                "" + value + "" +
+                "\n");
+    }
+
 }
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 cd397ad6fc092..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
@@ -17,7 +17,12 @@
 
 package org.apache.kafka.common.config;
 
+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 resources that have configs.
@@ -28,7 +33,25 @@ public final class ConfigResource {
      * Type of resource.
      */
     public enum Type {
-        BROKER, TOPIC, UNKNOWN;
+        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()))
+        );
+
+        private final byte id;
+
+        Type(final byte id) {
+            this.id = id;
+        }
+
+        public byte id() {
+            return id;
+        }
+
+        public static Type forId(final byte id) {
+            return TYPES.getOrDefault(id, UNKNOWN);
+        }
     }
 
     private final Type type;
@@ -61,6 +84,14 @@ public String name() {
         return name;
     }
 
+    /**
+     * Returns true if this is the default resource of a resource type.
+     * Resource name is empty for the default resource.
+     */
+    public boolean isDefault() {
+        return name.isEmpty();
+    }
+
     @Override
     public boolean equals(Object o) {
         if (this == o)
@@ -82,6 +113,6 @@ public int hashCode() {
 
     @Override
     public String toString() {
-        return "ConfigResource{type=" + type + ", name='" + name + "'}";
+        return "ConfigResource(type=" + type + ", name='" + name + "')";
     }
 }
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
new file mode 100644
index 0000000000000..4f078b15a32f3
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/config/ConfigTransformer.java
@@ -0,0 +1,177 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.common.config;
+
+import org.apache.kafka.common.config.provider.ConfigProvider;
+import org.apache.kafka.common.config.provider.FileConfigProvider;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+
+/**
+ * This class wraps a set of {@link ConfigProvider} instances and uses them to perform
+ * transformations.
+ *
+ * 
        The default variable pattern is of the form ${provider:[path:]key},
+ * where the provider corresponds to a {@link ConfigProvider} instance, as passed to
+ * {@link ConfigTransformer#ConfigTransformer(Map)}.  The pattern will extract a set
+ * of paths (which are optional) and keys and then pass them to {@link ConfigProvider#get(String, Set)} to obtain the
+ * values with which to replace the variables.
+ *
+ * 
        For example, if a Map consisting of an entry with a provider name "file" and provider instance
+ * {@link FileConfigProvider} is passed to the {@link ConfigTransformer#ConfigTransformer(Map)}, and a Properties
+ * file with contents
+ * 
        + * fileKey=someValue
+ * 
        
+ * resides at the path "/tmp/properties.txt", then when a configuration Map which has an entry with a key "someKey" and
+ * a value "${file:/tmp/properties.txt:fileKey}" is passed to the {@link #transform(Map)} method, then the transformed
+ * Map will have an entry with key "someKey" and a value "someValue".
+ *
+ * 
        This class only depends on {@link ConfigProvider#get(String, Set)} and does not depend on subscription support
+ * in a {@link ConfigProvider}, such as the {@link ConfigProvider#subscribe(String, Set, ConfigChangeCallback)} and
+ * {@link ConfigProvider#unsubscribe(String, Set, ConfigChangeCallback)} methods.
+ */
+public class ConfigTransformer {
+    public static final Pattern DEFAULT_PATTERN = Pattern.compile("\\$\\{([^}]*?):(([^}]*?):)?([^}]*?)\\}");
+    private static final String EMPTY_PATH = "";
+
+    private final Map configProviders;
+
+    /**
+     * Creates a ConfigTransformer with the default pattern, of the form ${provider:[path:]key}.
+     *
+     * @param configProviders a Map of provider names and {@link ConfigProvider} instances.
+     */
+    public ConfigTransformer(Map configProviders) {
+        this.configProviders = configProviders;
+    }
+
+    /**
+     * Transforms the given configuration data by using the {@link ConfigProvider} instances to
+     * look up values to replace the variables in the pattern.
+     *
+     * @param configs the configuration values to be transformed
+     * @return an instance of {@link ConfigTransformerResult}
+     */
+    public ConfigTransformerResult transform(Map configs) {
+        Map>> keysByProvider = new HashMap<>();
+        Map>> lookupsByProvider = new HashMap<>();
+
+        // Collect the variables from the given configs that need transformation
+        for (Map.Entry config : configs.entrySet()) {
+            if (config.getValue() != null) {
+                List vars = getVars(config.getValue(), DEFAULT_PATTERN);
+                for (ConfigVariable var : vars) {
+                    Map> keysByPath = keysByProvider.computeIfAbsent(var.providerName, k -> new HashMap<>());
+                    Set keys = keysByPath.computeIfAbsent(var.path, k -> new HashSet<>());
+                    keys.add(var.variable);
+                }
+            }
+        }
+
+        // Retrieve requested variables from the ConfigProviders
+        Map ttls = new HashMap<>();
+        for (Map.Entry>> entry : keysByProvider.entrySet()) {
+            String providerName = entry.getKey();
+            ConfigProvider provider = configProviders.get(providerName);
+            Map> keysByPath = entry.getValue();
+            if (provider != null && keysByPath != null) {
+                for (Map.Entry> pathWithKeys : keysByPath.entrySet()) {
+                    String path = pathWithKeys.getKey();
+                    Set keys = new HashSet<>(pathWithKeys.getValue());
+                    ConfigData configData = provider.get(path, keys);
+                    Map data = configData.data();
+                    Long ttl = configData.ttl();
+                    if (ttl != null && ttl >= 0) {
+                        ttls.put(path, ttl);
+                    }
+                    Map> keyValuesByPath =
+                            lookupsByProvider.computeIfAbsent(providerName, k -> new HashMap<>());
+                    keyValuesByPath.put(path, data);
+                }
+            }
+        }
+
+        // Perform the transformations by performing variable replacements
+        Map data = new HashMap<>(configs);
+        for (Map.Entry config : configs.entrySet()) {
+            data.put(config.getKey(), replace(lookupsByProvider, config.getValue(), DEFAULT_PATTERN));
+        }
+        return new ConfigTransformerResult(data, ttls);
+    }
+
+    private static List getVars(String value, Pattern pattern) {
+        List configVars = new ArrayList<>();
+        Matcher matcher = pattern.matcher(value);
+        while (matcher.find()) {
+            configVars.add(new ConfigVariable(matcher));
+        }
+        return configVars;
+    }
+
+    private static String replace(Map>> lookupsByProvider,
+                                  String value,
+                                  Pattern pattern) {
+        if (value == null) {
+            return null;
+        }
+        Matcher matcher = pattern.matcher(value);
+        StringBuilder builder = new StringBuilder();
+        int i = 0;
+        while (matcher.find()) {
+            ConfigVariable configVar = new ConfigVariable(matcher);
+            Map> lookupsByPath = lookupsByProvider.get(configVar.providerName);
+            if (lookupsByPath != null) {
+                Map keyValues = lookupsByPath.get(configVar.path);
+                String replacement = keyValues.get(configVar.variable);
+                builder.append(value, i, matcher.start());
+                if (replacement == null) {
+                    // No replacements will be performed; just return the original value
+                    builder.append(matcher.group(0));
+                } else {
+                    builder.append(replacement);
+                }
+                i = matcher.end();
+            }
+        }
+        builder.append(value, i, value.length());
+        return builder.toString();
+    }
+
+    private static class ConfigVariable {
+        final String providerName;
+        final String path;
+        final String variable;
+
+        ConfigVariable(Matcher matcher) {
+            this.providerName = matcher.group(1);
+            this.path = matcher.group(3) != null ? matcher.group(3) : EMPTY_PATH;
+            this.variable = matcher.group(4);
+        }
+
+        public String toString() {
+            return "(" + providerName + ":" + (path != null ? path + ":" : "") + variable + ")";
+        }
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/config/ConfigTransformerResult.java b/clients/src/main/java/org/apache/kafka/common/config/ConfigTransformerResult.java
new file mode 100644
index 0000000000000..a05669cb6c3ee
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/config/ConfigTransformerResult.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.config;
+
+import org.apache.kafka.common.config.provider.ConfigProvider;
+
+import java.util.Map;
+
+/**
+ * The result of a transformation from {@link ConfigTransformer}.
+ */
+public class ConfigTransformerResult {
+
+    private Map ttls;
+    private Map data;
+
+    /**
+     * Creates a new ConfigTransformerResult with the given data and TTL values for a set of paths.
+     *
+     * @param data a Map of key-value pairs
+     * @param ttls a Map of path and TTL values (in milliseconds)
+     */
+    public ConfigTransformerResult(Map data, Map ttls) {
+        this.data = data;
+        this.ttls = ttls;
+    }
+
+    /**
+     * Returns the transformed data, with variables replaced with corresponding values from the
+     * ConfigProvider instances if found.
+     *
+     * 
        Modifying the transformed data that is returned does not affect the {@link ConfigProvider} nor the
+     * original data that was used as the source of the transformation.
+     *
+     * @return data a Map of key-value pairs
+     */
+    public Map data() {
+        return data;
+    }
+
+    /**
+     * Returns the TTL values (in milliseconds) returned from the ConfigProvider instances for a given set of paths.
+     *
+     * @return data a Map of path and TTL values
+     */
+    public Map ttls() {
+        return ttls;
+    }
+}
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/SaslConfigs.java b/clients/src/main/java/org/apache/kafka/common/config/SaslConfigs.java
index f61b7dd5e6c66..db93ea41647f9 100644
--- a/clients/src/main/java/org/apache/kafka/common/config/SaslConfigs.java
+++ b/clients/src/main/java/org/apache/kafka/common/config/SaslConfigs.java
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.common.config;
 
+import org.apache.kafka.common.config.ConfigDef.Range;
 import org.apache.kafka.common.config.internals.BrokerSecurityConfigs;
 
 import java.util.List;
@@ -49,7 +50,24 @@ public class SaslConfigs {
     public static final String SASL_JAAS_CONFIG = "sasl.jaas.config";
     public static final String SASL_JAAS_CONFIG_DOC = "JAAS login context parameters for SASL connections in the format used by JAAS configuration files. "
         + "JAAS configuration file format is described here. "
-        + "The format for the value is: '  (=)*;'";
+        + "The format for the value is: 'loginModuleClass controlFlag (optionName=optionValue)*;'. For brokers, "
+        + "the config must be prefixed with listener prefix and SASL mechanism name in lower-case. For example, "
+        + "listener.name.sasl_ssl.scram-sha-256.sasl.jaas.config=com.example.ScramLoginModule required;";
+
+    public static final String SASL_CLIENT_CALLBACK_HANDLER_CLASS = "sasl.client.callback.handler.class";
+    public static final String SASL_CLIENT_CALLBACK_HANDLER_CLASS_DOC = "The fully qualified name of a SASL client callback handler class "
+        + "that implements the AuthenticateCallbackHandler interface.";
+
+    public static final String SASL_LOGIN_CALLBACK_HANDLER_CLASS = "sasl.login.callback.handler.class";
+    public static final String SASL_LOGIN_CALLBACK_HANDLER_CLASS_DOC = "The fully qualified name of a SASL login callback handler class "
+            + "that implements the AuthenticateCallbackHandler interface. For brokers, login callback handler config must be prefixed with "
+            + "listener prefix and SASL mechanism name in lower-case. For example, "
+            + "listener.name.sasl_ssl.scram-sha-256.sasl.login.callback.handler.class=com.example.CustomScramLoginCallbackHandler";
+
+    public static final String SASL_LOGIN_CLASS = "sasl.login.class";
+    public static final String SASL_LOGIN_CLASS_DOC = "The fully qualified name of a class that implements the Login interface. "
+        + "For brokers, login config must be prefixed with listener prefix and SASL mechanism name in lower-case. For example, "
+        + "listener.name.sasl_ssl.scram-sha-256.sasl.login.class=com.example.CustomScramLogin";
 
     public static final String SASL_KERBEROS_SERVICE_NAME = "sasl.kerberos.service.name";
     public static final String SASL_KERBEROS_SERVICE_NAME_DOC = "The Kerberos principal name that Kafka runs as. "
@@ -72,6 +90,34 @@ public class SaslConfigs {
     public static final String SASL_KERBEROS_MIN_TIME_BEFORE_RELOGIN_DOC = "Login thread sleep time between refresh attempts.";
     public static final long DEFAULT_KERBEROS_MIN_TIME_BEFORE_RELOGIN = 1 * 60 * 1000L;
 
+    public static final String SASL_LOGIN_REFRESH_WINDOW_FACTOR = "sasl.login.refresh.window.factor";
+    public static final String SASL_LOGIN_REFRESH_WINDOW_FACTOR_DOC = "Login refresh thread will sleep until the specified window factor relative to the"
+            + " credential's lifetime has been reached, at which time it will try to refresh the credential."
+            + " Legal values are between 0.5 (50%) and 1.0 (100%) inclusive; a default value of 0.8 (80%) is used"
+            + " if no value is specified. Currently applies only to OAUTHBEARER.";
+    public static final double DEFAULT_LOGIN_REFRESH_WINDOW_FACTOR = 0.80;
+
+    public static final String SASL_LOGIN_REFRESH_WINDOW_JITTER = "sasl.login.refresh.window.jitter";
+    public static final String SASL_LOGIN_REFRESH_WINDOW_JITTER_DOC = "The maximum amount of random jitter relative to the credential's lifetime"
+            + " that is added to the login refresh thread's sleep time. Legal values are between 0 and 0.25 (25%) inclusive;"
+            + " a default value of 0.05 (5%) is used if no value is specified. Currently applies only to OAUTHBEARER.";
+    public static final double DEFAULT_LOGIN_REFRESH_WINDOW_JITTER = 0.05;
+
+    public static final String SASL_LOGIN_REFRESH_MIN_PERIOD_SECONDS = "sasl.login.refresh.min.period.seconds";
+    public static final String SASL_LOGIN_REFRESH_MIN_PERIOD_SECONDS_DOC = "The desired minimum time for the login refresh thread to wait before refreshing a credential,"
+            + " in seconds. Legal values are between 0 and 900 (15 minutes); a default value of 60 (1 minute) is used if no value is specified.  This value and "
+            + " sasl.login.refresh.buffer.seconds are both ignored if their sum exceeds the remaining lifetime of a credential."
+            + " Currently applies only to OAUTHBEARER.";
+    public static final short DEFAULT_LOGIN_REFRESH_MIN_PERIOD_SECONDS = 60;
+
+    public static final String SASL_LOGIN_REFRESH_BUFFER_SECONDS = "sasl.login.refresh.buffer.seconds";
+    public static final String SASL_LOGIN_REFRESH_BUFFER_SECONDS_DOC = "The amount of buffer time before credential expiration to maintain when refreshing a credential,"
+            + " in seconds. If a refresh would otherwise occur closer to expiration than the number of buffer seconds then the refresh will be moved up to maintain"
+            + " as much of the buffer time as possible. Legal values are between 0 and 3600 (1 hour); a default value of  300 (5 minutes) is used if no value is specified."
+            + " This value and sasl.login.refresh.min.period.seconds are both ignored if their sum exceeds the remaining lifetime of a credential."
+            + " Currently applies only to OAUTHBEARER.";
+    public static final short DEFAULT_LOGIN_REFRESH_BUFFER_SECONDS = 300;
+
     /**
      * @deprecated As of 1.0.0. This field will be removed in a future major release.
      */
@@ -94,7 +140,14 @@ public static void addClientSaslSupport(ConfigDef config) {
                 .define(SaslConfigs.SASL_KERBEROS_TICKET_RENEW_WINDOW_FACTOR, ConfigDef.Type.DOUBLE, SaslConfigs.DEFAULT_KERBEROS_TICKET_RENEW_WINDOW_FACTOR, ConfigDef.Importance.LOW, SaslConfigs.SASL_KERBEROS_TICKET_RENEW_WINDOW_FACTOR_DOC)
                 .define(SaslConfigs.SASL_KERBEROS_TICKET_RENEW_JITTER, ConfigDef.Type.DOUBLE, SaslConfigs.DEFAULT_KERBEROS_TICKET_RENEW_JITTER, ConfigDef.Importance.LOW, SaslConfigs.SASL_KERBEROS_TICKET_RENEW_JITTER_DOC)
                 .define(SaslConfigs.SASL_KERBEROS_MIN_TIME_BEFORE_RELOGIN, ConfigDef.Type.LONG, SaslConfigs.DEFAULT_KERBEROS_MIN_TIME_BEFORE_RELOGIN, ConfigDef.Importance.LOW, SaslConfigs.SASL_KERBEROS_MIN_TIME_BEFORE_RELOGIN_DOC)
+                .define(SaslConfigs.SASL_LOGIN_REFRESH_WINDOW_FACTOR, ConfigDef.Type.DOUBLE, SaslConfigs.DEFAULT_LOGIN_REFRESH_WINDOW_FACTOR, Range.between(0.5, 1.0), ConfigDef.Importance.LOW, SaslConfigs.SASL_LOGIN_REFRESH_WINDOW_FACTOR_DOC)
+                .define(SaslConfigs.SASL_LOGIN_REFRESH_WINDOW_JITTER, ConfigDef.Type.DOUBLE, SaslConfigs.DEFAULT_LOGIN_REFRESH_WINDOW_JITTER, Range.between(0.0, 0.25), ConfigDef.Importance.LOW, SaslConfigs.SASL_LOGIN_REFRESH_WINDOW_JITTER_DOC)
+                .define(SaslConfigs.SASL_LOGIN_REFRESH_MIN_PERIOD_SECONDS, ConfigDef.Type.SHORT, SaslConfigs.DEFAULT_LOGIN_REFRESH_MIN_PERIOD_SECONDS, Range.between(0, 900), ConfigDef.Importance.LOW, SaslConfigs.SASL_LOGIN_REFRESH_MIN_PERIOD_SECONDS_DOC)
+                .define(SaslConfigs.SASL_LOGIN_REFRESH_BUFFER_SECONDS, ConfigDef.Type.SHORT, SaslConfigs.DEFAULT_LOGIN_REFRESH_BUFFER_SECONDS, Range.between(0, 3600), ConfigDef.Importance.LOW, SaslConfigs.SASL_LOGIN_REFRESH_BUFFER_SECONDS_DOC)
                 .define(SaslConfigs.SASL_MECHANISM, ConfigDef.Type.STRING, SaslConfigs.DEFAULT_SASL_MECHANISM, ConfigDef.Importance.MEDIUM, SaslConfigs.SASL_MECHANISM_DOC)
-                .define(SaslConfigs.SASL_JAAS_CONFIG, ConfigDef.Type.PASSWORD, null, ConfigDef.Importance.MEDIUM, SaslConfigs.SASL_JAAS_CONFIG_DOC);
+                .define(SaslConfigs.SASL_JAAS_CONFIG, ConfigDef.Type.PASSWORD, null, ConfigDef.Importance.MEDIUM, SaslConfigs.SASL_JAAS_CONFIG_DOC)
+                .define(SaslConfigs.SASL_CLIENT_CALLBACK_HANDLER_CLASS, ConfigDef.Type.CLASS, null, ConfigDef.Importance.MEDIUM, SaslConfigs.SASL_CLIENT_CALLBACK_HANDLER_CLASS_DOC)
+                .define(SaslConfigs.SASL_LOGIN_CALLBACK_HANDLER_CLASS, ConfigDef.Type.CLASS, null, ConfigDef.Importance.MEDIUM, SaslConfigs.SASL_LOGIN_CALLBACK_HANDLER_CLASS_DOC)
+                .define(SaslConfigs.SASL_LOGIN_CLASS, ConfigDef.Type.CLASS, null, ConfigDef.Importance.MEDIUM, SaslConfigs.SASL_LOGIN_CLASS_DOC);
     }
 }
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 042b051f46b8b..55a58accf5c7d 100644
--- a/clients/src/main/java/org/apache/kafka/common/config/SslConfigs.java
+++ b/clients/src/main/java/org/apache/kafka/common/config/SslConfigs.java
@@ -16,10 +16,14 @@
  */
 package org.apache.kafka.common.config;
 
+import org.apache.kafka.common.config.ConfigDef.Type;
 import org.apache.kafka.common.config.internals.BrokerSecurityConfigs;
+import org.apache.kafka.common.utils.Java;
+import org.apache.kafka.common.utils.Utils;
 
 import javax.net.ssl.KeyManagerFactory;
 import javax.net.ssl.TrustManagerFactory;
+import java.util.Set;
 
 public class SslConfigs {
     /*
@@ -47,11 +51,15 @@ public class SslConfigs {
 
     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. "
-            + "Allowed values in recent JVMs are TLS, TLSv1.1 and TLSv1.2. SSL, SSLv2 and SSLv3 "
-            + "may be supported in older JVMs, but their usage is discouraged due to known security vulnerabilities.";
+        + "The default is 'TLSv1.3' when running with Java 11 or newer, 'TLSv1.2' otherwise. "
+        + "This value should be fine for most use cases. "
+        + "Allowed values in recent JVMs are 'TLSv1.2' and 'TLSv1.3'. 'TLS', 'TLSv1.1', 'SSL', 'SSLv2' and 'SSLv3' "
+        + "may be supported in older JVMs, but their usage is discouraged due to known security vulnerabilities. "
+        + "With the default value for this config and 'ssl.enabled.protocols', clients will downgrade to 'TLSv1.2' if "
+        + "the server does not support 'TLSv1.3'. If this config is set to 'TLSv1.2', clients will not use 'TLSv1.3' even "
+        + "if it is one of the values in ssl.enabled.protocols and the server only supports 'TLSv1.3'.";
 
-    public static final String DEFAULT_SSL_PROTOCOL = "TLS";
+    public static final String DEFAULT_SSL_PROTOCOL;
 
     public static final String SSL_PROVIDER_CONFIG = "ssl.provider";
     public static final String SSL_PROVIDER_DOC = "The name of the security provider used for SSL connections. Default value is the default security provider of the JVM.";
@@ -61,25 +69,53 @@ public class SslConfigs {
             + "By default all the available cipher suites are supported.";
 
     public static final String SSL_ENABLED_PROTOCOLS_CONFIG = "ssl.enabled.protocols";
-    public static final String SSL_ENABLED_PROTOCOLS_DOC = "The list of protocols enabled for SSL connections.";
-    public static final String DEFAULT_SSL_ENABLED_PROTOCOLS = "TLSv1.2,TLSv1.1,TLSv1";
+    public static final String SSL_ENABLED_PROTOCOLS_DOC = "The list of protocols enabled for SSL connections. "
+        + "The default is 'TLSv1.2,TLSv1.3' when running with Java 11 or newer, 'TLSv1.2' otherwise. With the "
+        + "default value for Java 11, clients and servers will prefer TLSv1.3 if both support it and fallback "
+        + "to TLSv1.2 otherwise (assuming both support at least TLSv1.2). This default should be fine for most "
+        + "cases. Also see the config documentation for `ssl.protocol`.";
+    public static final String DEFAULT_SSL_ENABLED_PROTOCOLS;
+
+    static {
+        if (Java.IS_JAVA11_COMPATIBLE) {
+            DEFAULT_SSL_PROTOCOL = "TLSv1.3";
+            DEFAULT_SSL_ENABLED_PROTOCOLS = "TLSv1.2,TLSv1.3";
+        } else {
+            DEFAULT_SSL_PROTOCOL = "TLSv1.2";
+            DEFAULT_SSL_ENABLED_PROTOCOLS = "TLSv1.2";
+        }
+    }
 
     public static final String SSL_KEYSTORE_TYPE_CONFIG = "ssl.keystore.type";
     public static final String SSL_KEYSTORE_TYPE_DOC = "The file format of the key store file. "
             + "This is optional for client.";
     public static final String DEFAULT_SSL_KEYSTORE_TYPE = "JKS";
 
+    public static final String SSL_KEYSTORE_KEY_CONFIG = "ssl.keystore.key";
+    public static final String SSL_KEYSTORE_KEY_DOC = "Private key in the format specified by 'ssl.keystore.type'. "
+        + "Default SSL engine factory supports only PEM format with PKCS#8 keys. If the key is encrypted, "
+        + "key password must be specified using 'ssl.key.password'";
+
+    public static final String SSL_KEYSTORE_CERTIFICATE_CHAIN_CONFIG = "ssl.keystore.certificate.chain";
+    public static final String SSL_KEYSTORE_CERTIFICATE_CHAIN_DOC = "Certificate chain in the format specified by 'ssl.keystore.type'. "
+        + "Default SSL engine factory supports only PEM format with a list of X.509 certificates";
+
+    public static final String SSL_TRUSTSTORE_CERTIFICATES_CONFIG = "ssl.truststore.certificates";
+    public static final String SSL_TRUSTSTORE_CERTIFICATES_DOC = "Trusted certificates in the format specified by 'ssl.truststore.type'. "
+        + "Default SSL engine factory supports only PEM format with X.509 certificates.";
+
     public static final String SSL_KEYSTORE_LOCATION_CONFIG = "ssl.keystore.location";
     public static final String SSL_KEYSTORE_LOCATION_DOC = "The location of the key store file. "
         + "This is optional for client and can be used for two-way authentication for client.";
 
     public static final String SSL_KEYSTORE_PASSWORD_CONFIG = "ssl.keystore.password";
     public static final String SSL_KEYSTORE_PASSWORD_DOC = "The store password for the key store file. "
-        + "This is optional for client and only needed if ssl.keystore.location is configured. ";
+        + "This is optional for client and only needed if 'ssl.keystore.location' is configured. "
+        + " Key store password is not supported for PEM format.";
 
     public static final String SSL_KEY_PASSWORD_CONFIG = "ssl.key.password";
-    public static final String SSL_KEY_PASSWORD_DOC = "The password of the private key in the key store file. "
-            + "This is optional for client.";
+    public static final String SSL_KEY_PASSWORD_DOC = "The password of the private key in the key store file or"
+        + "the PEM key specified in `ssl.keystore.key'. This is required for clients only if two-way authentication is configured.";
 
     public static final String SSL_TRUSTSTORE_TYPE_CONFIG = "ssl.truststore.type";
     public static final String SSL_TRUSTSTORE_TYPE_DOC = "The file format of the trust store file.";
@@ -89,7 +125,9 @@ public class SslConfigs {
     public static final String SSL_TRUSTSTORE_LOCATION_DOC = "The location of the trust store file. ";
 
     public static final String SSL_TRUSTSTORE_PASSWORD_CONFIG = "ssl.truststore.password";
-    public static final String SSL_TRUSTSTORE_PASSWORD_DOC = "The password for the trust store file. If a password is not set access to the truststore is still available, but integrity checking is disabled.";
+    public static final String SSL_TRUSTSTORE_PASSWORD_DOC = "The password for the trust store file. "
+        + "If a password is not set, trust store file configured will still be used, but integrity checking is disabled. "
+        + "Trust store password is not supported for PEM format.";
 
     public static final String SSL_KEYMANAGER_ALGORITHM_CONFIG = "ssl.keymanager.algorithm";
     public static final String SSL_KEYMANAGER_ALGORITHM_DOC = "The algorithm used by key manager factory for SSL connections. "
@@ -103,10 +141,14 @@ public class SslConfigs {
 
     public static final String SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG = "ssl.endpoint.identification.algorithm";
     public static final String SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_DOC = "The endpoint identification algorithm to validate server hostname using server certificate. ";
+    public static final String DEFAULT_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM = "https";
 
     public static final String SSL_SECURE_RANDOM_IMPLEMENTATION_CONFIG = "ssl.secure.random.implementation";
     public static final String SSL_SECURE_RANDOM_IMPLEMENTATION_DOC = "The SecureRandom PRNG implementation to use for SSL cryptography operations. ";
 
+    public static final String SSL_ENGINE_FACTORY_CLASS_CONFIG = "ssl.engine.factory.class";
+    public static final String SSL_ENGINE_FACTORY_CLASS_DOC = "The class of type org.apache.kafka.common.security.auth.SslEngineFactory to provide SSLEngine objects. Default value is org.apache.kafka.common.security.ssl.DefaultSslEngineFactory";
+
     /**
      * @deprecated As of 1.0.0. This field will be removed in a future major release.
      */
@@ -127,12 +169,40 @@ public static void addClientSslSupport(ConfigDef config) {
                 .define(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG, ConfigDef.Type.STRING, null,  ConfigDef.Importance.HIGH, SslConfigs.SSL_KEYSTORE_LOCATION_DOC)
                 .define(SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG, ConfigDef.Type.PASSWORD, null, ConfigDef.Importance.HIGH, SslConfigs.SSL_KEYSTORE_PASSWORD_DOC)
                 .define(SslConfigs.SSL_KEY_PASSWORD_CONFIG, ConfigDef.Type.PASSWORD, null, ConfigDef.Importance.HIGH, SslConfigs.SSL_KEY_PASSWORD_DOC)
+                .define(SslConfigs.SSL_KEYSTORE_KEY_CONFIG, Type.PASSWORD, null,  ConfigDef.Importance.HIGH, SslConfigs.SSL_KEYSTORE_KEY_DOC)
+                .define(SslConfigs.SSL_KEYSTORE_CERTIFICATE_CHAIN_CONFIG, ConfigDef.Type.PASSWORD, null,  ConfigDef.Importance.HIGH, SslConfigs.SSL_KEYSTORE_CERTIFICATE_CHAIN_DOC)
+                .define(SslConfigs.SSL_TRUSTSTORE_CERTIFICATES_CONFIG, ConfigDef.Type.PASSWORD, null,  ConfigDef.Importance.HIGH, SslConfigs.SSL_TRUSTSTORE_CERTIFICATES_DOC)
                 .define(SslConfigs.SSL_TRUSTSTORE_TYPE_CONFIG, ConfigDef.Type.STRING, SslConfigs.DEFAULT_SSL_TRUSTSTORE_TYPE, ConfigDef.Importance.MEDIUM, SslConfigs.SSL_TRUSTSTORE_TYPE_DOC)
                 .define(SslConfigs.SSL_TRUSTSTORE_LOCATION_CONFIG, ConfigDef.Type.STRING, null, ConfigDef.Importance.HIGH, SslConfigs.SSL_TRUSTSTORE_LOCATION_DOC)
                 .define(SslConfigs.SSL_TRUSTSTORE_PASSWORD_CONFIG, ConfigDef.Type.PASSWORD, null, ConfigDef.Importance.HIGH, SslConfigs.SSL_TRUSTSTORE_PASSWORD_DOC)
                 .define(SslConfigs.SSL_KEYMANAGER_ALGORITHM_CONFIG, ConfigDef.Type.STRING, SslConfigs.DEFAULT_SSL_KEYMANGER_ALGORITHM, ConfigDef.Importance.LOW, SslConfigs.SSL_KEYMANAGER_ALGORITHM_DOC)
                 .define(SslConfigs.SSL_TRUSTMANAGER_ALGORITHM_CONFIG, ConfigDef.Type.STRING, SslConfigs.DEFAULT_SSL_TRUSTMANAGER_ALGORITHM, ConfigDef.Importance.LOW, SslConfigs.SSL_TRUSTMANAGER_ALGORITHM_DOC)
-                .define(SslConfigs.SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG, ConfigDef.Type.STRING, null, ConfigDef.Importance.LOW, SslConfigs.SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_DOC)
-                .define(SslConfigs.SSL_SECURE_RANDOM_IMPLEMENTATION_CONFIG, ConfigDef.Type.STRING, null, ConfigDef.Importance.LOW, SslConfigs.SSL_SECURE_RANDOM_IMPLEMENTATION_DOC);
+                .define(SslConfigs.SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG, ConfigDef.Type.STRING, SslConfigs.DEFAULT_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM, ConfigDef.Importance.LOW, SslConfigs.SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_DOC)
+                .define(SslConfigs.SSL_SECURE_RANDOM_IMPLEMENTATION_CONFIG, ConfigDef.Type.STRING, null, ConfigDef.Importance.LOW, SslConfigs.SSL_SECURE_RANDOM_IMPLEMENTATION_DOC)
+                .define(SslConfigs.SSL_ENGINE_FACTORY_CLASS_CONFIG, ConfigDef.Type.CLASS, null, ConfigDef.Importance.LOW, SslConfigs.SSL_ENGINE_FACTORY_CLASS_DOC);
     }
+
+    public static final Set RECONFIGURABLE_CONFIGS = Utils.mkSet(
+            SslConfigs.SSL_KEYSTORE_TYPE_CONFIG,
+            SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG,
+            SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG,
+            SslConfigs.SSL_KEY_PASSWORD_CONFIG,
+            SslConfigs.SSL_TRUSTSTORE_TYPE_CONFIG,
+            SslConfigs.SSL_TRUSTSTORE_LOCATION_CONFIG,
+            SslConfigs.SSL_TRUSTSTORE_PASSWORD_CONFIG,
+            SslConfigs.SSL_KEYSTORE_CERTIFICATE_CHAIN_CONFIG,
+            SslConfigs.SSL_KEYSTORE_KEY_CONFIG,
+            SslConfigs.SSL_TRUSTSTORE_CERTIFICATES_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_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,
+            SslConfigs.SSL_ENGINE_FACTORY_CLASS_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 168dd25359376..7121979e1eb1c 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
@@ -18,11 +18,11 @@
 package org.apache.kafka.common.config;
 
 /**
- * Keys that can be used to configure a topic.  These keys are useful when creating or reconfiguring a
+ * 
        Keys that can be used to configure a topic. These keys are useful when creating or reconfiguring a
  * topic using the AdminClient.
  *
- * The intended pattern is for broker configs to include a `log.` prefix. For example, to set the default broker
- * cleanup policy, one would set log.cleanup.policy instead of cleanup.policy. Unfortunately, there are many cases
+ * 
        The intended pattern is for broker configs to include a `log.` prefix. For example, to set the default broker
+ * cleanup policy, one would set log.cleanup.policy instead of cleanup.policy. Unfortunately, there are many cases
  * where this pattern is not followed.
  */
 // This is a public API, so we should not remove or alter keys without a discussion and a deprecation period.
@@ -73,19 +73,20 @@ public class TopicConfig {
     public static final String RETENTION_MS_DOC = "This configuration controls the maximum time we will retain a " +
         "log before we will discard old log segments to free up space if we are using the " +
         "\"delete\" retention policy. This represents an SLA on how soon consumers must read " +
-        "their data.";
+        "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 " +
-        "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 " +
-        "message format versions, uncompressed records are not grouped into batches and this limit only applies to a " +
-        "single record in that case.
        ";
+    public static final String MAX_MESSAGE_BYTES_DOC =
+        "The largest record batch size allowed by Kafka (after compression if compression is enabled). " +
+        "If this is increased and there are consumers older than 0.10.2, the consumers' fetch " +
+        "size must also be increased so that 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.";
 
     public static final String INDEX_INTERVAL_BYTES_CONFIG = "index.interval.bytes";
     public static final String INDEX_INTERVAL_BYTES_DOCS = "This setting controls how frequently " +
-        "Kafka adds an index entry to it's offset index. The default setting ensures that we index a " +
+        "Kafka adds an index entry to its offset index. The default setting ensures that we index a " +
         "message roughly every 4096 bytes. More indexing allows reads to jump closer to the exact " +
         "position in the log but makes the index larger. You probably don't need to change this.";
 
@@ -104,6 +105,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,13 +116,17 @@ 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";
     public static final String CLEANUP_POLICY_DELETE = "delete";
     public static final String CLEANUP_POLICY_DOC = "A string that is either \"" + CLEANUP_POLICY_DELETE +
-        "\" or \"" + CLEANUP_POLICY_COMPACT + "\". This string designates the retention policy to use on " +
+        "\" or \"" + CLEANUP_POLICY_COMPACT + "\" or both. This string designates the retention policy to use on " +
         "old log segments. The default policy (\"delete\") will discard old segments when their retention " +
         "time or size limit has been reached. The \"compact\" setting will enable log " +
         "compaction on the topic.";
@@ -132,15 +141,15 @@ public class TopicConfig {
         "this configuration specifies the minimum number of replicas that must acknowledge " +
         "a write for the write to be considered successful. If this minimum cannot be met, " +
         "then the producer will raise an exception (either NotEnoughReplicas or " +
-        "NotEnoughReplicasAfterAppend).
        When used together, min.insync.replicas and acks " +
+        "NotEnoughReplicasAfterAppend).
        When used together, min.insync.replicas and acks " +
         "allow you to enforce greater durability guarantees. A typical scenario would be to " +
-        "create a topic with a replication factor of 3, set min.insync.replicas to 2, and " +
-        "produce with acks of \"all\". This will ensure that the producer raises an exception " +
+        "create a topic with a replication factor of 3, set min.insync.replicas to 2, and " +
+        "produce with acks of \"all\". This will ensure that the producer raises an exception " +
         "if a majority of replicas do not receive a write.";
 
     public static final String COMPRESSION_TYPE_CONFIG = "compression.type";
     public static final String COMPRESSION_TYPE_DOC = "Specify the final compression type for a given topic. " +
-        "This configuration accepts the standard compression codecs ('gzip', 'snappy', lz4). It additionally " +
+        "This configuration accepts the standard compression codecs ('gzip', 'snappy', 'lz4', 'zstd'). It additionally " +
         "accepts 'uncompressed' which is equivalent to no compression; and 'producer' which means retain the " +
         "original compression codec set by the producer.";
 
@@ -165,4 +174,11 @@ public class TopicConfig {
         "the timestamp when a broker receives a message and the timestamp specified in the message. If " +
         "message.timestamp.type=CreateTime, a message will be rejected if the difference in timestamp " +
         "exceeds this threshold. This configuration is ignored if message.timestamp.type=LogAppendTime.";
+
+    public static final String MESSAGE_DOWNCONVERSION_ENABLE_CONFIG = "message.downconversion.enable";
+    public static final String MESSAGE_DOWNCONVERSION_ENABLE_DOC = "This configuration controls whether " +
+        "down-conversion of message formats is enabled to satisfy consume requests. When set to false, " +
+        "broker will not perform down-conversion for consumers expecting an older message format. The broker responds " +
+        "with UNSUPPORTED_VERSION error for consume requests from such older clients. This configuration" +
+        "does not apply to any message format conversion that might be required for replication to followers.";
 }
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 18616ec74e57e..3b84908fb154d 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
@@ -33,17 +33,30 @@ public class BrokerSecurityConfigs {
     public static final String SASL_KERBEROS_PRINCIPAL_TO_LOCAL_RULES_CONFIG = "sasl.kerberos.principal.to.local.rules";
     public static final String SSL_CLIENT_AUTH_CONFIG = "ssl.client.auth";
     public static final String SASL_ENABLED_MECHANISMS_CONFIG = "sasl.enabled.mechanisms";
+    public static final String SASL_SERVER_CALLBACK_HANDLER_CLASS = "sasl.server.callback.handler.class";
+    public static final String SSL_PRINCIPAL_MAPPING_RULES_CONFIG = "ssl.principal.mapping.rules";
+    public static final String CONNECTIONS_MAX_REAUTH_MS = "connections.max.reauth.ms";
 
     public static final String PRINCIPAL_BUILDER_CLASS_DOC = "The fully qualified name of a class that implements the " +
             "KafkaPrincipalBuilder interface, which is used to build the KafkaPrincipal object used during " +
             "authorization. This config also supports the deprecated PrincipalBuilder interface which was previously " +
             "used for client authentication over SSL. If no principal builder is defined, the default behavior depends " +
-            "on the security protocol in use. For SSL authentication, the principal name will be the distinguished " +
+            "on the security protocol in use. For SSL authentication,  the principal will be derived using the" +
+            " rules defined by " + SSL_PRINCIPAL_MAPPING_RULES_CONFIG + " applied on the distinguished " +
             "name from the client certificate if one is provided; otherwise, if client authentication is not required, " +
             "the principal name will be ANONYMOUS. For SASL authentication, the principal will be derived using the " +
             "rules defined by " + SASL_KERBEROS_PRINCIPAL_TO_LOCAL_RULES_CONFIG + " if GSSAPI is in use, " +
             "and the SASL authentication ID for other mechanisms. For PLAINTEXT, the principal will be ANONYMOUS.";
 
+    public static final String SSL_PRINCIPAL_MAPPING_RULES_DOC = "A list of rules for mapping from distinguished name" +
+            " from the client certificate to short name. The rules are evaluated in order and the first rule that matches" +
+            " a principal name is used to map it to a short name. Any later rules in the list are ignored. By default," +
+            " distinguished name of the X.500 certificate will be the principal. For more details on the format please" +
+            " 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 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 " +
             "first rule that matches a principal name is used to map it to a short name. Any later rules in the list are " +
@@ -56,15 +69,25 @@ public class BrokerSecurityConfigs {
     public static final String SSL_CLIENT_AUTH_DOC = "Configures kafka broker to request client authentication."
             + " The following settings are common: "
             + " 
          "
-            + " 
        • ssl.client.auth=required If set to required"
-            + " client authentication is required."
+            + " 
        • ssl.client.auth=required If set to required 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.";
+            + " unlike required, 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."
+            + "
        ";
 
     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. "
             + "Only GSSAPI is enabled by default.";
     public static final List DEFAULT_SASL_ENABLED_MECHANISMS = Collections.singletonList(SaslConfigs.GSSAPI_MECHANISM);
 
+    public static final String SASL_SERVER_CALLBACK_HANDLER_CLASS_DOC = "The fully qualified name of a SASL server callback handler "
+            + "class that implements the AuthenticateCallbackHandler interface. Server callback handlers must be prefixed with "
+            + "listener prefix and SASL mechanism name in lower-case. For example, "
+            + "listener.name.sasl_ssl.plain.sasl.server.callback.handler.class=com.example.CustomPlainCallbackHandler.";
+
+    public static final String CONNECTIONS_MAX_REAUTH_MS_DOC = "When explicitly set to a positive number (the default is 0, not a positive number), "
+            + "a session lifetime that will not exceed the configured value will be communicated to v2.2.0 or later clients when they authenticate. "
+            + "The broker will disconnect any such connection that is not re-authenticated within the session lifetime and that is then subsequently "
+            + "used for any purpose other than re-authentication. Configuration names can optionally be prefixed with listener prefix and SASL "
+            + "mechanism name in lower-case. For example, listener.name.sasl_ssl.oauthbearer.connections.max.reauth.ms=3600000";
 }
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
new file mode 100644
index 0000000000000..fe65ddfa30656
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/config/provider/ConfigProvider.java
@@ -0,0 +1,81 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.common.config.provider;
+
+import org.apache.kafka.common.Configurable;
+import org.apache.kafka.common.config.ConfigChangeCallback;
+import org.apache.kafka.common.config.ConfigData;
+
+import java.io.Closeable;
+import java.util.Set;
+
+/**
+ * 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 {
+
+    /**
+     * Retrieves the data at the given path.
+     *
+     * @param path the path where the data resides
+     * @return the configuration data
+     */
+    ConfigData get(String path);
+
+    /**
+     * Retrieves the data with the given keys at the given path.
+     *
+     * @param path the path where the data resides
+     * @param keys the keys whose values will be retrieved
+     * @return the configuration data
+     */
+    ConfigData get(String path, Set keys);
+
+    /**
+     * Subscribes to changes for the given keys at the given path (optional operation).
+     *
+     * @param path the path where the data resides
+     * @param keys the keys whose values will be retrieved
+     * @param callback the callback to invoke upon change
+     * @throws {@link UnsupportedOperationException} if the subscribe operation is not supported
+     */
+    default void subscribe(String path, Set keys, ConfigChangeCallback callback) {
+        throw new UnsupportedOperationException();
+    }
+
+    /**
+     * Unsubscribes to changes for the given keys at the given path (optional operation).
+     *
+     * @param path the path where the data resides
+     * @param keys the keys whose values will be retrieved
+     * @param callback the callback to be unsubscribed from changes
+     * @throws {@link UnsupportedOperationException} if the unsubscribe operation is not supported
+     */
+    default void unsubscribe(String path, Set keys, ConfigChangeCallback callback) {
+        throw new UnsupportedOperationException();
+    }
+
+    /**
+     * Clears all subscribers (optional operation).
+     *
+     * @throws {@link UnsupportedOperationException} if the unsubscribeAll operation is not supported
+     */
+    default void unsubscribeAll() {
+        throw new UnsupportedOperationException();
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/config/provider/DirectoryConfigProvider.java b/clients/src/main/java/org/apache/kafka/common/config/provider/DirectoryConfigProvider.java
new file mode 100644
index 0000000000000..adf2774f09363
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/config/provider/DirectoryConfigProvider.java
@@ -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 org.apache.kafka.common.config.provider;
+
+import org.apache.kafka.common.config.ConfigData;
+import org.apache.kafka.common.config.ConfigException;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.File;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.Map;
+import java.util.Set;
+import java.util.function.Predicate;
+import java.util.stream.Collectors;
+
+import static java.util.Collections.emptyMap;
+
+/**
+ * An implementation of {@link ConfigProvider} based on a directory of files.
+ * Property keys correspond to the names of the regular (i.e. non-directory)
+ * files in a directory given by the path parameter.
+ * Property values are taken from the file contents corresponding to each key.
+ */
+public class DirectoryConfigProvider implements ConfigProvider {
+
+    private static final Logger log = LoggerFactory.getLogger(DirectoryConfigProvider.class);
+
+    @Override
+    public void configure(Map configs) { }
+
+    @Override
+    public void close() throws IOException { }
+
+    /**
+     * Retrieves the data contained in regular files in the directory given by {@code path}.
+     * Non-regular files (such as directories) in the given directory are silently ignored.
+     * @param path the directory where data files reside.
+     * @return the configuration data.
+     */
+    @Override
+    public ConfigData get(String path) {
+        return get(path, Files::isRegularFile);
+    }
+
+    /**
+     * Retrieves the data contained in the regular files named by {@code keys} in the directory given by {@code path}.
+     * Non-regular files (such as directories) in the given directory are silently ignored.
+     * @param path the directory where data files reside.
+     * @param keys the keys whose values will be retrieved.
+     * @return the configuration data.
+     */
+    @Override
+    public ConfigData get(String path, Set keys) {
+        return get(path, pathname ->
+                Files.isRegularFile(pathname)
+                        && keys.contains(pathname.getFileName().toString()));
+    }
+
+    private static ConfigData get(String path, Predicate fileFilter) {
+        Map map = emptyMap();
+        if (path != null && !path.isEmpty()) {
+            Path dir = new File(path).toPath();
+            if (!Files.isDirectory(dir)) {
+                log.warn("The path {} is not a directory", path);
+            } else {
+                try {
+                    map = Files.list(dir)
+                        .filter(fileFilter)
+                        .collect(Collectors.toMap(
+                            p -> p.getFileName().toString(),
+                            p -> read(p)));
+                } catch (IOException e) {
+                    throw new ConfigException("Could not list directory " + dir, e);
+                }
+            }
+        }
+        return new ConfigData(map);
+    }
+
+    private static String read(Path path) {
+        try {
+            return new String(Files.readAllBytes(path), StandardCharsets.UTF_8);
+        } catch (IOException e) {
+            throw new ConfigException("Could not read file " + path + " for property " + path.getFileName(), e);
+        }
+    }
+
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/config/provider/FileConfigProvider.java b/clients/src/main/java/org/apache/kafka/common/config/provider/FileConfigProvider.java
new file mode 100644
index 0000000000000..4e376ecdeed08
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/config/provider/FileConfigProvider.java
@@ -0,0 +1,103 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.common.config.provider;
+
+import org.apache.kafka.common.config.ConfigData;
+import org.apache.kafka.common.config.ConfigException;
+
+import java.io.IOException;
+import java.io.Reader;
+import java.nio.file.Files;
+import java.nio.file.Paths;
+import java.util.Enumeration;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Properties;
+import java.util.Set;
+
+/**
+ * An implementation of {@link ConfigProvider} that represents a Properties file.
+ * All property keys and values are stored as cleartext.
+ */
+public class FileConfigProvider implements ConfigProvider {
+
+    public void configure(Map configs) {
+    }
+
+    /**
+     * Retrieves the data at the given Properties file.
+     *
+     * @param path the file where the data resides
+     * @return the configuration data
+     */
+    public ConfigData get(String path) {
+        Map data = new HashMap<>();
+        if (path == null || path.isEmpty()) {
+            return new ConfigData(data);
+        }
+        try (Reader reader = reader(path)) {
+            Properties properties = new Properties();
+            properties.load(reader);
+            Enumeration keys = properties.keys();
+            while (keys.hasMoreElements()) {
+                String key = keys.nextElement().toString();
+                String value = properties.getProperty(key);
+                if (value != null) {
+                    data.put(key, value);
+                }
+            }
+            return new ConfigData(data);
+        } catch (IOException e) {
+            throw new ConfigException("Could not read properties from file " + path);
+        }
+    }
+
+    /**
+     * Retrieves the data with the given keys at the given Properties file.
+     *
+     * @param path the file where the data resides
+     * @param keys the keys whose values will be retrieved
+     * @return the configuration data
+     */
+    public ConfigData get(String path, Set keys) {
+        Map data = new HashMap<>();
+        if (path == null || path.isEmpty()) {
+            return new ConfigData(data);
+        }
+        try (Reader reader = reader(path)) {
+            Properties properties = new Properties();
+            properties.load(reader);
+            for (String key : keys) {
+                String value = properties.getProperty(key);
+                if (value != null) {
+                    data.put(key, value);
+                }
+            }
+            return new ConfigData(data);
+        } catch (IOException e) {
+            throw new ConfigException("Could not read properties from file " + path);
+        }
+    }
+
+    // visible for testing
+    protected Reader reader(String path) throws IOException {
+        return Files.newBufferedReader(Paths.get(path));
+    }
+
+    public void close() {
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/AuthenticationException.java b/clients/src/main/java/org/apache/kafka/common/errors/AuthenticationException.java
index f6458c6f22ded..7a05eba03f2bc 100644
--- a/clients/src/main/java/org/apache/kafka/common/errors/AuthenticationException.java
+++ b/clients/src/main/java/org/apache/kafka/common/errors/AuthenticationException.java
@@ -40,6 +40,10 @@ public AuthenticationException(String message) {
         super(message);
     }
 
+    public AuthenticationException(Throwable cause) {
+        super(cause);
+    }
+
     public AuthenticationException(String message, Throwable cause) {
         super(message, cause);
     }
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/DelegationTokenAuthorizationException.java b/clients/src/main/java/org/apache/kafka/common/errors/DelegationTokenAuthorizationException.java
new file mode 100644
index 0000000000000..ddc97c64cf38c
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/DelegationTokenAuthorizationException.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;
+
+public class DelegationTokenAuthorizationException extends AuthorizationException {
+
+    private static final long serialVersionUID = 1L;
+
+    public DelegationTokenAuthorizationException(String message) {
+        super(message);
+    }
+
+    public DelegationTokenAuthorizationException(String message, Throwable cause) {
+        super(message, cause);
+    }
+
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/DelegationTokenDisabledException.java b/clients/src/main/java/org/apache/kafka/common/errors/DelegationTokenDisabledException.java
new file mode 100644
index 0000000000000..798611e57ee5d
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/DelegationTokenDisabledException.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;
+
+public class DelegationTokenDisabledException extends ApiException {
+
+    private static final long serialVersionUID = 1L;
+
+    public DelegationTokenDisabledException(String message) {
+        super(message);
+    }
+
+    public DelegationTokenDisabledException(String message, Throwable cause) {
+        super(message, cause);
+    }
+
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/DelegationTokenExpiredException.java b/clients/src/main/java/org/apache/kafka/common/errors/DelegationTokenExpiredException.java
new file mode 100644
index 0000000000000..4dae7f3f493a6
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/DelegationTokenExpiredException.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;
+
+public class DelegationTokenExpiredException extends ApiException {
+
+    private static final long serialVersionUID = 1L;
+
+    public DelegationTokenExpiredException(String message) {
+        super(message);
+    }
+
+    public DelegationTokenExpiredException(String message, Throwable cause) {
+        super(message, cause);
+    }
+
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/DelegationTokenNotFoundException.java b/clients/src/main/java/org/apache/kafka/common/errors/DelegationTokenNotFoundException.java
new file mode 100644
index 0000000000000..5875edfef7ea4
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/DelegationTokenNotFoundException.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;
+
+public class DelegationTokenNotFoundException extends ApiException {
+
+    private static final long serialVersionUID = 1L;
+
+    public DelegationTokenNotFoundException(String message) {
+        super(message);
+    }
+
+    public DelegationTokenNotFoundException(String message, Throwable cause) {
+        super(message, cause);
+    }
+
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/DelegationTokenOwnerMismatchException.java b/clients/src/main/java/org/apache/kafka/common/errors/DelegationTokenOwnerMismatchException.java
new file mode 100644
index 0000000000000..5c8239ebb59a2
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/DelegationTokenOwnerMismatchException.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;
+
+public class DelegationTokenOwnerMismatchException extends ApiException {
+
+    private static final long serialVersionUID = 1L;
+
+    public DelegationTokenOwnerMismatchException(String message) {
+        super(message);
+    }
+
+    public DelegationTokenOwnerMismatchException(String message, Throwable cause) {
+        super(message, cause);
+    }
+
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/DuplicateResourceException.java b/clients/src/main/java/org/apache/kafka/common/errors/DuplicateResourceException.java
new file mode 100644
index 0000000000000..1c0ec43a52c7b
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/DuplicateResourceException.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.common.errors;
+
+/**
+ * Exception thrown due to a request that illegally refers to the same resource twice
+ * (for example, trying to both create and delete the same SCRAM credential for a particular user in a single request).
+ */
+public class DuplicateResourceException extends ApiException {
+
+    private static final long serialVersionUID = 1L;
+
+    private final String resource;
+
+    /**
+     * Constructor
+     *
+     * @param message the exception's message
+     */
+    public DuplicateResourceException(String message) {
+        this(null, message);
+    }
+
+    /**
+     *
+     * @param message the exception's message
+     * @param cause the exception's cause
+     */
+    public DuplicateResourceException(String message, Throwable cause) {
+        this(null, message, cause);
+    }
+
+    /**
+     * Constructor
+     *
+     * @param resource the (potentially null) resource that was referred to twice
+     * @param message the exception's message
+     */
+    public DuplicateResourceException(String resource, String message) {
+        super(message);
+        this.resource = resource;
+    }
+
+    /**
+     * Constructor
+     *
+     * @param resource the (potentially null) resource that was referred to twice
+     * @param message the exception's message
+     * @param cause the exception's cause
+     */
+    public DuplicateResourceException(String resource, String message, Throwable cause) {
+        super(message, cause);
+        this.resource = resource;
+    }
+
+    /**
+     *
+     * @return the (potentially null) resource that was referred to twice
+     */
+    public String resource() {
+        return this.resource;
+    }
+}
\ No newline at end of file
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/FeatureUpdateFailedException.java b/clients/src/main/java/org/apache/kafka/common/errors/FeatureUpdateFailedException.java
new file mode 100644
index 0000000000000..9f5e23d3104da
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/FeatureUpdateFailedException.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 FeatureUpdateFailedException extends ApiException {
+    private static final long serialVersionUID = 1L;
+
+    public FeatureUpdateFailedException(final String message) {
+        super(message);
+    }
+
+    public FeatureUpdateFailedException(final String message, final 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/FencedLeaderEpochException.java b/clients/src/main/java/org/apache/kafka/common/errors/FencedLeaderEpochException.java
new file mode 100644
index 0000000000000..24f0eef93924b
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/FencedLeaderEpochException.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.errors;
+
+/**
+ * The request contained a leader epoch which is smaller than that on the broker that received the
+ * request. This can happen when an operation is attempted before a pending metadata update has been
+ * received. Clients will typically refresh metadata before retrying.
+ */
+public class FencedLeaderEpochException extends InvalidMetadataException {
+    private static final long serialVersionUID = 1L;
+
+    public FencedLeaderEpochException(String message) {
+        super(message);
+    }
+
+    public FencedLeaderEpochException(String message, Throwable cause) {
+        super(message, cause);
+    }
+
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/FetchSessionIdNotFoundException.java b/clients/src/main/java/org/apache/kafka/common/errors/FetchSessionIdNotFoundException.java
new file mode 100644
index 0000000000000..2ce5f740d6719
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/FetchSessionIdNotFoundException.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 FetchSessionIdNotFoundException extends RetriableException {
+    private static final long serialVersionUID = 1L;
+
+    public FetchSessionIdNotFoundException() {
+    }
+
+    public FetchSessionIdNotFoundException(String message) {
+        super(message);
+    }
+}
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/GroupCoordinatorNotAvailableException.java b/clients/src/main/java/org/apache/kafka/common/errors/GroupCoordinatorNotAvailableException.java
deleted file mode 100644
index 03a7719a6f428..0000000000000
--- a/clients/src/main/java/org/apache/kafka/common/errors/GroupCoordinatorNotAvailableException.java
+++ /dev/null
@@ -1,47 +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.errors;
-
-/**
- * The broker returns this error code for consumer metadata requests or offset commit requests if the offsets topic has
- * not yet been created.
- *
- * @deprecated As of Kafka 0.11, this has been replaced by {@link CoordinatorNotAvailableException}
- */
-@Deprecated
-public class GroupCoordinatorNotAvailableException extends RetriableException {
-    public static final GroupCoordinatorNotAvailableException INSTANCE = new GroupCoordinatorNotAvailableException();
-
-    private static final long serialVersionUID = 1L;
-
-    public GroupCoordinatorNotAvailableException() {
-        super();
-    }
-
-    public GroupCoordinatorNotAvailableException(String message) {
-        super(message);
-    }
-
-    public GroupCoordinatorNotAvailableException(String message, Throwable cause) {
-        super(message, cause);
-    }
-
-    public GroupCoordinatorNotAvailableException(Throwable cause) {
-        super(cause);
-    }
-
-}
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/GroupIdNotFoundException.java b/clients/src/main/java/org/apache/kafka/common/errors/GroupIdNotFoundException.java
new file mode 100644
index 0000000000000..a4d509d3a2074
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/GroupIdNotFoundException.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 GroupIdNotFoundException extends ApiException {
+    public GroupIdNotFoundException(String message) {
+        super(message);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/GroupLoadInProgressException.java b/clients/src/main/java/org/apache/kafka/common/errors/GroupLoadInProgressException.java
deleted file mode 100644
index 73daa5f1c3a55..0000000000000
--- a/clients/src/main/java/org/apache/kafka/common/errors/GroupLoadInProgressException.java
+++ /dev/null
@@ -1,46 +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.errors;
-
-/**
- * The broker returns this error code for any coordinator request if it is still loading the metadata (after a leader change
- * for that offsets topic partition) for this group.
- *
- * @deprecated As of Kafka 0.11, this has been replaced by {@link CoordinatorLoadInProgressException}
- */
-@Deprecated
-public class GroupLoadInProgressException extends RetriableException {
-
-    private static final long serialVersionUID = 1L;
-
-    public GroupLoadInProgressException() {
-        super();
-    }
-
-    public GroupLoadInProgressException(String message) {
-        super(message);
-    }
-
-    public GroupLoadInProgressException(String message, Throwable cause) {
-        super(message, cause);
-    }
-
-    public GroupLoadInProgressException(Throwable cause) {
-        super(cause);
-    }
-
-}
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
new file mode 100644
index 0000000000000..85d0c7d25ce85
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/GroupMaxSizeReachedException.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;
+
+/**
+ * Indicates that a consumer group is already at its configured maximum capacity and cannot accommodate more members
+ */
+public class GroupMaxSizeReachedException extends ApiException {
+    private static final long serialVersionUID = 1L;
+
+    public GroupMaxSizeReachedException(String message) {
+        super(message);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/GroupNotEmptyException.java b/clients/src/main/java/org/apache/kafka/common/errors/GroupNotEmptyException.java
new file mode 100644
index 0000000000000..e15b3e6d57f41
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/GroupNotEmptyException.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 GroupNotEmptyException extends ApiException {
+    public GroupNotEmptyException(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/InconsistentVoterSetException.java b/clients/src/main/java/org/apache/kafka/common/errors/InconsistentVoterSetException.java
new file mode 100644
index 0000000000000..8a3667b76dd30
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/InconsistentVoterSetException.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;
+
+public class InconsistentVoterSetException extends ApiException {
+
+    private static final long serialVersionUID = 1;
+
+    public InconsistentVoterSetException(String s) {
+        super(s);
+    }
+
+    public InconsistentVoterSetException(String message, Throwable cause) {
+        super(message, cause);
+    }
+
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/InvalidFetchSessionEpochException.java b/clients/src/main/java/org/apache/kafka/common/errors/InvalidFetchSessionEpochException.java
new file mode 100644
index 0000000000000..3b135c0147d06
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/InvalidFetchSessionEpochException.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 InvalidFetchSessionEpochException extends RetriableException {
+    private static final long serialVersionUID = 1L;
+
+    public InvalidFetchSessionEpochException() {
+    }
+
+    public InvalidFetchSessionEpochException(String message) {
+        super(message);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/InvalidPrincipalTypeException.java b/clients/src/main/java/org/apache/kafka/common/errors/InvalidPrincipalTypeException.java
new file mode 100644
index 0000000000000..a0736e976354f
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/InvalidPrincipalTypeException.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;
+
+public class InvalidPrincipalTypeException extends ApiException {
+
+    private static final long serialVersionUID = 1L;
+
+    public InvalidPrincipalTypeException(String message) {
+        super(message);
+    }
+
+    public InvalidPrincipalTypeException(String message, Throwable cause) {
+        super(message, cause);
+    }
+
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/InvalidProducerEpochException.java b/clients/src/main/java/org/apache/kafka/common/errors/InvalidProducerEpochException.java
new file mode 100644
index 0000000000000..79b82368feb95
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/InvalidProducerEpochException.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.common.errors;
+
+/**
+ * This exception indicates that the produce request sent to the partition leader
+ * contains a non-matching producer epoch. When encountering this exception, user should abort the ongoing transaction
+ * by calling KafkaProducer#abortTransaction which would try to send initPidRequest and reinitialize the producer
+ * under the hood.
+ */
+public class InvalidProducerEpochException extends ApiException {
+
+    private static final long serialVersionUID = 1L;
+
+    public InvalidProducerEpochException(String message) {
+        super(message);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/InvalidTopicException.java b/clients/src/main/java/org/apache/kafka/common/errors/InvalidTopicException.java
index f79e9a7b06719..729ebba8a6f5d 100644
--- a/clients/src/main/java/org/apache/kafka/common/errors/InvalidTopicException.java
+++ b/clients/src/main/java/org/apache/kafka/common/errors/InvalidTopicException.java
@@ -16,6 +16,10 @@
  */
 package org.apache.kafka.common.errors;
 
+import java.util.HashSet;
+import java.util.Set;
+
+
 /**
  * The client has attempted to perform an operation on an invalid topic.
  * For example the topic name is too long, contains invalid characters etc.
@@ -24,23 +28,36 @@
  * @see UnknownTopicOrPartitionException
  */
 public class InvalidTopicException extends ApiException {
-
     private static final long serialVersionUID = 1L;
 
+    private final Set invalidTopics;
+
     public InvalidTopicException() {
         super();
+        invalidTopics = new HashSet<>();
     }
 
     public InvalidTopicException(String message, Throwable cause) {
         super(message, cause);
+        invalidTopics = new HashSet<>();
     }
 
     public InvalidTopicException(String message) {
         super(message);
+        invalidTopics = new HashSet<>();
     }
 
     public InvalidTopicException(Throwable cause) {
         super(cause);
+        invalidTopics = new HashSet<>();
     }
 
+    public InvalidTopicException(Set invalidTopics) {
+        super("Invalid topics: " + invalidTopics);
+        this.invalidTopics = invalidTopics;
+    }
+
+    public Set invalidTopics() {
+        return invalidTopics;
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/InvalidTxnTimeoutException.java b/clients/src/main/java/org/apache/kafka/common/errors/InvalidTxnTimeoutException.java
index c751bc4553d1d..f16af66da7c12 100644
--- a/clients/src/main/java/org/apache/kafka/common/errors/InvalidTxnTimeoutException.java
+++ b/clients/src/main/java/org/apache/kafka/common/errors/InvalidTxnTimeoutException.java
@@ -18,7 +18,7 @@
 
 /**
  * The transaction coordinator returns this error code if the timeout received via the InitProducerIdRequest is larger than
- * the `max.transaction.timeout.ms` config value.
+ * the `transaction.max.timeout.ms` config value.
  */
 public class InvalidTxnTimeoutException extends ApiException {
     private static final long serialVersionUID = 1L;
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/InvalidUpdateVersionException.java b/clients/src/main/java/org/apache/kafka/common/errors/InvalidUpdateVersionException.java
new file mode 100644
index 0000000000000..e41262d59f52d
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/InvalidUpdateVersionException.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 InvalidUpdateVersionException extends ApiException {
+
+    public InvalidUpdateVersionException(String message) {
+        super(message);
+    }
+
+    public InvalidUpdateVersionException(String message, Throwable throwable) {
+        super(message, throwable);
+    }
+
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/ListenerNotFoundException.java b/clients/src/main/java/org/apache/kafka/common/errors/ListenerNotFoundException.java
new file mode 100644
index 0000000000000..82c5d892f49ba
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/ListenerNotFoundException.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.common.errors;
+
+/**
+ * The leader does not have an endpoint corresponding to the listener on which metadata was requested.
+ * This could indicate a broker configuration error or a transient error when listeners are updated
+ * dynamically and client requests are processed before all brokers have updated their listeners.
+ * This is currently used only for missing listeners on leader brokers, but may be used for followers
+ * in future.
+ */
+public class ListenerNotFoundException extends InvalidMetadataException {
+
+    private static final long serialVersionUID = 1L;
+
+    public ListenerNotFoundException(String message) {
+        super(message);
+    }
+
+    public ListenerNotFoundException(String message, Throwable cause) {
+        super(message, cause);
+    }
+
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/MemberIdRequiredException.java b/clients/src/main/java/org/apache/kafka/common/errors/MemberIdRequiredException.java
new file mode 100644
index 0000000000000..55393e0e286af
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/MemberIdRequiredException.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.common.errors;
+
+public class MemberIdRequiredException extends ApiException {
+
+    private static final long serialVersionUID = 1L;
+
+    public MemberIdRequiredException(String message) {
+        super(message);
+    }
+
+    public MemberIdRequiredException(String message, Throwable cause) {
+        super(message, cause);
+    }
+}
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/NotCoordinatorForGroupException.java b/clients/src/main/java/org/apache/kafka/common/errors/NotCoordinatorForGroupException.java
deleted file mode 100644
index cee649502e979..0000000000000
--- a/clients/src/main/java/org/apache/kafka/common/errors/NotCoordinatorForGroupException.java
+++ /dev/null
@@ -1,46 +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.errors;
-
-/**
- * The broker returns this error code if it receives an offset fetch or commit request for a consumer group that it is
- * not a coordinator for.
- *
- * @deprecated As of Kafka 0.11, this has been replaced by {@link NotCoordinatorException}
- */
-@Deprecated
-public class NotCoordinatorForGroupException extends RetriableException {
-
-    private static final long serialVersionUID = 1L;
-
-    public NotCoordinatorForGroupException() {
-        super();
-    }
-
-    public NotCoordinatorForGroupException(String message) {
-        super(message);
-    }
-
-    public NotCoordinatorForGroupException(String message, Throwable cause) {
-        super(message, cause);
-    }
-
-    public NotCoordinatorForGroupException(Throwable cause) {
-        super(cause);
-    }
-
-}
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/NotLeaderForPartitionException.java b/clients/src/main/java/org/apache/kafka/common/errors/NotLeaderForPartitionException.java
index 718277905751d..30efc49dc7584 100644
--- a/clients/src/main/java/org/apache/kafka/common/errors/NotLeaderForPartitionException.java
+++ b/clients/src/main/java/org/apache/kafka/common/errors/NotLeaderForPartitionException.java
@@ -17,8 +17,10 @@
 package org.apache.kafka.common.errors;
 
 /**
- * This server is not the leader for the given partition
+ * This server is not the leader for the given partition.
+ * @deprecated since 2.6. Use {@link NotLeaderOrFollowerException}.
  */
+@Deprecated
 public class NotLeaderForPartitionException extends InvalidMetadataException {
 
     private static final long serialVersionUID = 1L;
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/NotLeaderOrFollowerException.java b/clients/src/main/java/org/apache/kafka/common/errors/NotLeaderOrFollowerException.java
new file mode 100644
index 0000000000000..2db960b738e07
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/NotLeaderOrFollowerException.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.errors;
+
+/**
+ * Broker returns this error if a request could not be processed because the broker is not the leader
+ * or follower for a topic partition. This could be a transient exception during leader elections and
+ * reassignments. For `Produce` and other requests which are intended only for the leader, this exception
+ * indicates that the broker is not the current leader. For consumer `Fetch` requests which may be
+ * satisfied by a leader or follower, this exception indicates that the broker is not a replica
+ * of the topic partition.
+ */
+@SuppressWarnings("deprecation")
+public class NotLeaderOrFollowerException extends NotLeaderForPartitionException {
+
+    private static final long serialVersionUID = 1L;
+
+    public NotLeaderOrFollowerException() {
+        super();
+    }
+
+    public NotLeaderOrFollowerException(String message) {
+        super(message);
+    }
+
+    public NotLeaderOrFollowerException(Throwable cause) {
+        super(cause);
+    }
+
+    public NotLeaderOrFollowerException(String message, Throwable cause) {
+        super(message, cause);
+    }
+
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/OffsetNotAvailableException.java b/clients/src/main/java/org/apache/kafka/common/errors/OffsetNotAvailableException.java
new file mode 100644
index 0000000000000..97de3b3d64e96
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/OffsetNotAvailableException.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;
+
+/**
+ * Indicates that the leader is not able to guarantee monotonically increasing offsets
+ * due to the high watermark lagging behind the epoch start offset after a recent leader election
+ */
+public class OffsetNotAvailableException extends RetriableException {
+    private static final long serialVersionUID = 1L;
+
+    public OffsetNotAvailableException(String message) {
+        super(message);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/PreferredLeaderNotAvailableException.java b/clients/src/main/java/org/apache/kafka/common/errors/PreferredLeaderNotAvailableException.java
new file mode 100644
index 0000000000000..73dfd64ac6d23
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/PreferredLeaderNotAvailableException.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 PreferredLeaderNotAvailableException extends InvalidMetadataException {
+
+    public PreferredLeaderNotAvailableException(String message) {
+        super(message);
+    }
+
+    public PreferredLeaderNotAvailableException(String message, Throwable cause) {
+        super(message, cause);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/PrincipalDeserializationException.java b/clients/src/main/java/org/apache/kafka/common/errors/PrincipalDeserializationException.java
new file mode 100644
index 0000000000000..d0d085802020d
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/PrincipalDeserializationException.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;
+
+/**
+ * Exception used to indicate a kafka principal deserialization failure during request forwarding.
+ */
+public class PrincipalDeserializationException extends ApiException {
+
+    private static final long serialVersionUID = 1L;
+
+    public PrincipalDeserializationException(String message) {
+        super(message);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/ReplicaNotAvailableException.java b/clients/src/main/java/org/apache/kafka/common/errors/ReplicaNotAvailableException.java
index b94d400d8fe31..07971cd57e18a 100644
--- a/clients/src/main/java/org/apache/kafka/common/errors/ReplicaNotAvailableException.java
+++ b/clients/src/main/java/org/apache/kafka/common/errors/ReplicaNotAvailableException.java
@@ -16,7 +16,13 @@
  */
 package org.apache.kafka.common.errors;
 
-public class ReplicaNotAvailableException extends ApiException {
+/**
+ * The replica is not available for the requested topic partition. This may be
+ * a transient exception during reassignments. From version 2.6 onwards, Fetch requests
+ * and other requests intended only for the leader or follower of the topic partition return
+ * {@link NotLeaderOrFollowerException} if the broker is a not a replica of the partition.
+ */
+public class ReplicaNotAvailableException extends InvalidMetadataException {
 
     private static final long serialVersionUID = 1L;
 
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/ResourceNotFoundException.java b/clients/src/main/java/org/apache/kafka/common/errors/ResourceNotFoundException.java
new file mode 100644
index 0000000000000..17dca08248e34
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/ResourceNotFoundException.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.common.errors;
+
+/**
+ * Exception thrown due to a request for a resource that does not exist.
+ */
+public class ResourceNotFoundException extends ApiException {
+
+    private static final long serialVersionUID = 1L;
+
+    private final String resource;
+
+    /**
+     * Constructor
+     *
+     * @param message the exception's message
+     */
+    public ResourceNotFoundException(String message) {
+        this(null, message);
+    }
+
+    /**
+     *
+     * @param message the exception's message
+     * @param cause the exception's cause
+     */
+    public ResourceNotFoundException(String message, Throwable cause) {
+        this(null, message, cause);
+    }
+
+    /**
+     * Constructor
+     *
+     * @param resource the (potentially null) resource that was not found
+     * @param message the exception's message
+     */
+    public ResourceNotFoundException(String resource, String message) {
+        super(message);
+        this.resource = resource;
+    }
+
+    /**
+     * Constructor
+     *
+     * @param resource the (potentially null) resource that was not found
+     * @param message the exception's message
+     * @param cause the exception's cause
+     */
+    public ResourceNotFoundException(String resource, String message, Throwable cause) {
+        super(message, cause);
+        this.resource = resource;
+    }
+
+    /**
+     *
+     * @return the (potentially null) resource that was not found
+     */
+    public String resource() {
+        return this.resource;
+    }
+}
\ No newline at end of file
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/RetriableException.java b/clients/src/main/java/org/apache/kafka/common/errors/RetriableException.java
index 213e4f9fe0844..6d9a76dc739fc 100644
--- a/clients/src/main/java/org/apache/kafka/common/errors/RetriableException.java
+++ b/clients/src/main/java/org/apache/kafka/common/errors/RetriableException.java
@@ -17,7 +17,7 @@
 package org.apache.kafka.common.errors;
 
 /**
- * A retryable exception is a transient exception that if retried may succeed.
+ * A retriable exception is a transient exception that if retried may succeed.
  */
 public abstract class RetriableException extends ApiException {
 
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/StaleBrokerEpochException.java b/clients/src/main/java/org/apache/kafka/common/errors/StaleBrokerEpochException.java
new file mode 100644
index 0000000000000..a5c0b41400757
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/StaleBrokerEpochException.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;
+
+public class StaleBrokerEpochException extends ApiException {
+
+    private static final long serialVersionUID = 1L;
+
+    public StaleBrokerEpochException(String message) {
+        super(message);
+    }
+
+    public StaleBrokerEpochException(String message, Throwable cause) {
+        super(message, cause);
+    }
+
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/ThrottlingQuotaExceededException.java b/clients/src/main/java/org/apache/kafka/common/errors/ThrottlingQuotaExceededException.java
new file mode 100644
index 0000000000000..c4d1350b205dd
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/ThrottlingQuotaExceededException.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.errors;
+
+/**
+ * Exception thrown if an operation on a resource exceeds the throttling quota.
+ */
+public class ThrottlingQuotaExceededException extends RetriableException {
+    private int throttleTimeMs = 0;
+
+    public ThrottlingQuotaExceededException(String message) {
+        super(message);
+    }
+
+    public ThrottlingQuotaExceededException(int throttleTimeMs, String message) {
+        super(message);
+        this.throttleTimeMs = throttleTimeMs;
+    }
+
+    public int throttleTimeMs() {
+        return this.throttleTimeMs;
+    }
+}
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/errors/TopicDeletionDisabledException.java b/clients/src/main/java/org/apache/kafka/common/errors/TopicDeletionDisabledException.java
new file mode 100644
index 0000000000000..41577d2a288e7
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/TopicDeletionDisabledException.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 TopicDeletionDisabledException extends  ApiException {
+    private static final long serialVersionUID = 1L;
+
+    public TopicDeletionDisabledException() {
+    }
+
+    public TopicDeletionDisabledException(String message) {
+        super(message);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/TransactionAbortedException.java b/clients/src/main/java/org/apache/kafka/common/errors/TransactionAbortedException.java
new file mode 100644
index 0000000000000..c394ac5e6a76a
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/TransactionAbortedException.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.common.errors;
+
+/**
+ * This is the Exception thrown when we are aborting any undrained batches during
+ * a transaction which is aborted without any underlying cause - which likely means that the user chose to abort.
+ */
+public class TransactionAbortedException extends ApiException {
+
+    private final static long serialVersionUID = 1L;
+
+    public TransactionAbortedException(String message, Throwable cause) {
+        super(message, cause);
+    }
+
+    public TransactionAbortedException(String message) {
+        super(message);
+    }
+
+    public TransactionAbortedException() {
+        super("Failing batch since transaction was aborted");
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/UnacceptableCredentialException.java b/clients/src/main/java/org/apache/kafka/common/errors/UnacceptableCredentialException.java
new file mode 100644
index 0000000000000..b7cffff6addd8
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/UnacceptableCredentialException.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.errors;
+
+/**
+ * Exception thrown when attempting to define a credential that does not meet the criteria for acceptability
+ * (for example, attempting to create a SCRAM credential with an empty username or password or too few/many iterations).
+ */
+public class UnacceptableCredentialException extends ApiException {
+
+    private static final long serialVersionUID = 1L;
+
+    /**
+     * Constructor
+     *
+     * @param message the exception's message
+     */
+    public UnacceptableCredentialException(String message) {
+        super(message);
+    }
+
+    /**
+     *
+     * @param message the exception's message
+     * @param cause the exception's cause
+     */
+    public UnacceptableCredentialException(String message, Throwable cause) {
+        super(message, cause);
+    }
+}
\ No newline at end of file
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/UnknownLeaderEpochException.java b/clients/src/main/java/org/apache/kafka/common/errors/UnknownLeaderEpochException.java
new file mode 100644
index 0000000000000..3714c364bedba
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/UnknownLeaderEpochException.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.errors;
+
+/**
+ * The request contained a leader epoch which is larger than that on the broker that received the
+ * request. This can happen if the client observes a metadata update before it has been propagated
+ * to all brokers. Clients need not refresh metadata before retrying.
+ */
+public class UnknownLeaderEpochException extends RetriableException {
+    private static final long serialVersionUID = 1L;
+
+    public UnknownLeaderEpochException(String message) {
+        super(message);
+    }
+
+    public UnknownLeaderEpochException(String message, Throwable cause) {
+        super(message, cause);
+    }
+
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/UnstableOffsetCommitException.java b/clients/src/main/java/org/apache/kafka/common/errors/UnstableOffsetCommitException.java
new file mode 100644
index 0000000000000..c89d717e5bbcf
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/UnstableOffsetCommitException.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;
+
+/**
+ * Exception thrown when there are unstable offsets for the requested topic partitions.
+ */
+public class UnstableOffsetCommitException extends RetriableException {
+
+    private static final long serialVersionUID = 1L;
+
+    public UnstableOffsetCommitException(String message) {
+        super(message);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/UnsupportedByAuthenticationException.java b/clients/src/main/java/org/apache/kafka/common/errors/UnsupportedByAuthenticationException.java
new file mode 100644
index 0000000000000..40f357c5e2e0b
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/UnsupportedByAuthenticationException.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.common.errors;
+
+/**
+ * Authentication mechanism does not support the requested function.
+ */
+public class UnsupportedByAuthenticationException extends ApiException {
+    private static final long serialVersionUID = 1L;
+
+    public UnsupportedByAuthenticationException(String message) {
+        super(message);
+    }
+
+    public UnsupportedByAuthenticationException(String message, Throwable cause) {
+        super(message, cause);
+    }
+
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/UnsupportedCompressionTypeException.java b/clients/src/main/java/org/apache/kafka/common/errors/UnsupportedCompressionTypeException.java
new file mode 100644
index 0000000000000..29ffe1b900e23
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/errors/UnsupportedCompressionTypeException.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.errors;
+
+/**
+ * The requesting client does not support the compression type of given partition.
+ */
+public class UnsupportedCompressionTypeException extends ApiException {
+
+    private static final long serialVersionUID = 1L;
+
+    public UnsupportedCompressionTypeException(String message) {
+        super(message);
+    }
+
+    public UnsupportedCompressionTypeException(String message, Throwable cause) {
+        super(message, cause);
+    }
+
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/errors/WakeupException.java b/clients/src/main/java/org/apache/kafka/common/errors/WakeupException.java
index 4726ec13cca94..f8ae8403a5d9f 100644
--- a/clients/src/main/java/org/apache/kafka/common/errors/WakeupException.java
+++ b/clients/src/main/java/org/apache/kafka/common/errors/WakeupException.java
@@ -21,7 +21,7 @@
 /**
  * Exception used to indicate preemption of a blocking operation by an external thread.
  * For example, {@link org.apache.kafka.clients.consumer.KafkaConsumer#wakeup}
- * can be used to break out of an active {@link org.apache.kafka.clients.consumer.KafkaConsumer#poll(long)},
+ * can be used to break out of an active {@link org.apache.kafka.clients.consumer.KafkaConsumer#poll(java.time.Duration)},
  * which would raise an instance of this exception.
  */
 public class WakeupException extends KafkaException {
diff --git a/clients/src/main/java/org/apache/kafka/common/feature/BaseVersionRange.java b/clients/src/main/java/org/apache/kafka/common/feature/BaseVersionRange.java
new file mode 100644
index 0000000000000..2d6ce702e253f
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/feature/BaseVersionRange.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.common.feature;
+
+import static java.util.stream.Collectors.joining;
+
+import java.util.Map;
+import java.util.Objects;
+
+import org.apache.kafka.common.utils.Utils;
+
+/**
+ * Represents an immutable basic version range using 2 attributes: min and max, each of type short.
+ * The min and max attributes need to satisfy 2 rules:
+ *  - they are each expected to be >= 1, as we only consider positive version values to be valid.
+ *  - max should be >= min.
+ *
+ * The class also provides API to convert the version range to a map.
+ * The class allows for configurable labels for the min/max attributes, which can be specialized by
+ * sub-classes (if needed).
+ */
+class BaseVersionRange {
+    // Non-empty label for the min version key, that's used only to convert to/from a map.
+    private final String minKeyLabel;
+
+    // The value of the minimum version.
+    private final short minValue;
+
+    // Non-empty label for the max version key, that's used only to convert to/from a map.
+    private final String maxKeyLabel;
+
+    // The value of the maximum version.
+    private final short maxValue;
+
+    /**
+     * Raises an exception unless the following condition is met:
+     * minValue >= 1 and maxValue >= 1 and maxValue >= minValue.
+     *
+     * @param minKeyLabel   Label for the min version key, that's used only to convert to/from a map.
+     * @param minValue      The minimum version value.
+     * @param maxKeyLabel   Label for the max version key, that's used only to convert to/from a map.
+     * @param maxValue      The maximum version value.
+     *
+     * @throws IllegalArgumentException   If any of the following conditions are true:
+     *                                     - (minValue < 1) OR (maxValue < 1) OR (maxValue < minValue).
+     *                                     - minKeyLabel is empty, OR, minKeyLabel is empty.
+     */
+    protected BaseVersionRange(String minKeyLabel, short minValue, String maxKeyLabel, short maxValue) {
+        if (minValue < 1 || maxValue < 1 || maxValue < minValue) {
+            throw new IllegalArgumentException(
+                String.format(
+                    "Expected minValue >= 1, maxValue >= 1 and maxValue >= minValue, but received" +
+                    " minValue: %d, maxValue: %d", minValue, maxValue));
+        }
+        if (minKeyLabel.isEmpty()) {
+            throw new IllegalArgumentException("Expected minKeyLabel to be non-empty.");
+        }
+        if (maxKeyLabel.isEmpty()) {
+            throw new IllegalArgumentException("Expected maxKeyLabel to be non-empty.");
+        }
+        this.minKeyLabel = minKeyLabel;
+        this.minValue = minValue;
+        this.maxKeyLabel = maxKeyLabel;
+        this.maxValue = maxValue;
+    }
+
+    public short min() {
+        return minValue;
+    }
+
+    public short max() {
+        return maxValue;
+    }
+
+    public String toString() {
+        return String.format(
+            "%s[%s]",
+            this.getClass().getSimpleName(),
+            mapToString(toMap()));
+    }
+
+    public Map toMap() {
+        return Utils.mkMap(Utils.mkEntry(minKeyLabel, min()), Utils.mkEntry(maxKeyLabel, max()));
+    }
+
+    private static String mapToString(final Map map) {
+        return map
+            .entrySet()
+            .stream()
+            .map(entry -> String.format("%s:%d", entry.getKey(), entry.getValue()))
+            .collect(joining(", "));
+    }
+
+    @Override
+    public boolean equals(Object other) {
+        if (this == other) {
+            return true;
+        }
+
+        if (other == null || getClass() != other.getClass()) {
+            return false;
+        }
+
+        final BaseVersionRange that = (BaseVersionRange) other;
+        return Objects.equals(this.minKeyLabel, that.minKeyLabel) &&
+            this.minValue == that.minValue &&
+            Objects.equals(this.maxKeyLabel, that.maxKeyLabel) &&
+            this.maxValue == that.maxValue;
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(minKeyLabel, minValue, maxKeyLabel, maxValue);
+    }
+
+    public static short valueOrThrow(String key, Map versionRangeMap) {
+        final Short value = versionRangeMap.get(key);
+        if (value == null) {
+            throw new IllegalArgumentException(String.format("%s absent in [%s]", key, mapToString(versionRangeMap)));
+        }
+        return value;
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/feature/Features.java b/clients/src/main/java/org/apache/kafka/common/feature/Features.java
new file mode 100644
index 0000000000000..4006d71947fb7
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/feature/Features.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.common.feature;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.stream.Collectors;
+import java.util.Objects;
+
+import static java.util.stream.Collectors.joining;
+
+/**
+ * Represents an immutable dictionary with key being feature name, and value being .
+ * Also provides API to convert the features and their version ranges to/from a map.
+ *
+ * This class can be instantiated only using its factory functions, with the important ones being:
+ * Features.supportedFeatures(...) and Features.finalizedFeatures(...).
+ *
+ * @param  is the type of version range.
+ * @see SupportedVersionRange
+ * @see FinalizedVersionRange
+ */
+public class Features {
+    private final Map features;
+
+    /**
+     * Constructor is made private, as for readability it is preferred the caller uses one of the
+     * static factory functions for instantiation (see below).
+     *
+     * @param features   Map of feature name to a type of VersionRange.
+     */
+    private Features(Map features) {
+        Objects.requireNonNull(features, "Provided features can not be null.");
+        this.features = features;
+    }
+
+    /**
+     * @param features   Map of feature name to SupportedVersionRange.
+     *
+     * @return           Returns a new Features object representing supported features.
+     */
+    public static Features supportedFeatures(Map features) {
+        return new Features<>(features);
+    }
+
+    /**
+     * @param features   Map of feature name to FinalizedVersionRange.
+     *
+     * @return           Returns a new Features object representing finalized features.
+     */
+    public static Features finalizedFeatures(Map features) {
+        return new Features<>(features);
+    }
+
+    // Visible for testing.
+    public static Features emptyFinalizedFeatures() {
+        return new Features<>(new HashMap<>());
+    }
+
+    public static Features emptySupportedFeatures() {
+        return new Features<>(new HashMap<>());
+    }
+
+    public Map features() {
+        return features;
+    }
+
+    public boolean empty() {
+        return features.isEmpty();
+    }
+
+    /**
+     * @param  feature   name of the feature
+     *
+     * @return           the VersionRangeType corresponding to the feature name, or null if the
+     *                   feature is absent
+     */
+    public VersionRangeType get(String feature) {
+        return features.get(feature);
+    }
+
+    public String toString() {
+        return String.format(
+            "Features{%s}",
+            features
+                .entrySet()
+                .stream()
+                .map(entry -> String.format("(%s -> %s)", entry.getKey(), entry.getValue()))
+                .collect(joining(", "))
+        );
+    }
+
+    /**
+     * @return   A map representation of the underlying features. The returned value can be converted
+     *           back to Features using one of the from*FeaturesMap() APIs of this class.
+     */
+    public Map> toMap() {
+        return features.entrySet().stream().collect(
+            Collectors.toMap(
+                Map.Entry::getKey,
+                entry -> entry.getValue().toMap()));
+    }
+
+    /**
+     * An interface that defines behavior to convert from a Map to an object of type BaseVersionRange.
+     */
+    private interface MapToBaseVersionRangeConverter {
+
+        /**
+         * Convert the map representation of an object of type , to an object of type .
+         *
+         * @param  baseVersionRangeMap   the map representation of a BaseVersionRange object.
+         *
+         * @return                       the object of type 
+         */
+        V fromMap(Map baseVersionRangeMap);
+    }
+
+    private static  Features fromFeaturesMap(
+        Map> featuresMap, MapToBaseVersionRangeConverter converter) {
+        return new Features<>(featuresMap.entrySet().stream().collect(
+            Collectors.toMap(
+                Map.Entry::getKey,
+                entry -> converter.fromMap(entry.getValue()))));
+    }
+
+    /**
+     * Converts from a map to Features.
+     *
+     * @param featuresMap  the map representation of a Features object,
+     *                     generated using the toMap() API.
+     *
+     * @return             the Features object
+     */
+    public static Features fromFinalizedFeaturesMap(
+        Map> featuresMap) {
+        return fromFeaturesMap(featuresMap, FinalizedVersionRange::fromMap);
+    }
+
+    /**
+     * Converts from a map to Features.
+     *
+     * @param featuresMap  the map representation of a Features object,
+     *                     generated using the toMap() API.
+     *
+     * @return             the Features object
+     */
+    public static Features fromSupportedFeaturesMap(
+        Map> featuresMap) {
+        return fromFeaturesMap(featuresMap, SupportedVersionRange::fromMap);
+    }
+
+    @Override
+    public boolean equals(Object other) {
+        if (this == other) {
+            return true;
+        }
+        if (!(other instanceof Features)) {
+            return false;
+        }
+
+        final Features that = (Features) other;
+        return Objects.equals(this.features, that.features);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(features);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/feature/FinalizedVersionRange.java b/clients/src/main/java/org/apache/kafka/common/feature/FinalizedVersionRange.java
new file mode 100644
index 0000000000000..27e6440478644
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/feature/FinalizedVersionRange.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 org.apache.kafka.common.feature;
+
+import java.util.Map;
+
+/**
+ * An extended {@link BaseVersionRange} representing the min/max versions for a finalized feature.
+ */
+public class FinalizedVersionRange extends BaseVersionRange {
+    // Label for the min version key, that's used only to convert to/from a map.
+    private static final String MIN_VERSION_LEVEL_KEY_LABEL = "min_version_level";
+
+    // Label for the max version key, that's used only to convert to/from a map.
+    private static final String MAX_VERSION_LEVEL_KEY_LABEL = "max_version_level";
+
+    public FinalizedVersionRange(short minVersionLevel, short maxVersionLevel) {
+        super(MIN_VERSION_LEVEL_KEY_LABEL, minVersionLevel, MAX_VERSION_LEVEL_KEY_LABEL, maxVersionLevel);
+    }
+
+    public static FinalizedVersionRange fromMap(Map versionRangeMap) {
+        return new FinalizedVersionRange(
+            BaseVersionRange.valueOrThrow(MIN_VERSION_LEVEL_KEY_LABEL, versionRangeMap),
+            BaseVersionRange.valueOrThrow(MAX_VERSION_LEVEL_KEY_LABEL, versionRangeMap));
+    }
+
+    /**
+     * Checks if the [min, max] version level range of this object does *NOT* fall within the
+     * [min, max] range of the provided SupportedVersionRange parameter.
+     *
+     * @param supportedVersionRange   the SupportedVersionRange to be checked
+     *
+     * @return                        - true, if the version levels are compatible
+     *                                - false otherwise
+     */
+    public boolean isIncompatibleWith(SupportedVersionRange supportedVersionRange) {
+        return min() < supportedVersionRange.min() || max() > supportedVersionRange.max();
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/feature/SupportedVersionRange.java b/clients/src/main/java/org/apache/kafka/common/feature/SupportedVersionRange.java
new file mode 100644
index 0000000000000..8993014a74b2e
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/feature/SupportedVersionRange.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.feature;
+
+import java.util.Map;
+
+/**
+ * An extended {@link BaseVersionRange} representing the min/max versions for a supported feature.
+ */
+public class SupportedVersionRange extends BaseVersionRange {
+    // Label for the min version key, that's used only to convert to/from a map.
+    private static final String MIN_VERSION_KEY_LABEL = "min_version";
+
+    // Label for the max version key, that's used only to convert to/from a map.
+    private static final String MAX_VERSION_KEY_LABEL = "max_version";
+
+    public SupportedVersionRange(short minVersion, short maxVersion) {
+        super(MIN_VERSION_KEY_LABEL, minVersion, MAX_VERSION_KEY_LABEL, maxVersion);
+    }
+
+    public SupportedVersionRange(short maxVersion) {
+        this((short) 1, maxVersion);
+    }
+
+    public static SupportedVersionRange fromMap(Map versionRangeMap) {
+        return new SupportedVersionRange(
+            BaseVersionRange.valueOrThrow(MIN_VERSION_KEY_LABEL, versionRangeMap),
+            BaseVersionRange.valueOrThrow(MAX_VERSION_KEY_LABEL, versionRangeMap));
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/header/Headers.java b/clients/src/main/java/org/apache/kafka/common/header/Headers.java
index 1796d62944ae5..2353249cebcd3 100644
--- a/clients/src/main/java/org/apache/kafka/common/header/Headers.java
+++ b/clients/src/main/java/org/apache/kafka/common/header/Headers.java
@@ -41,7 +41,7 @@ public interface Headers extends Iterable
         {
      * Removes all headers for the given key returning if the operation succeeded.
      * 
      * @param key to remove all headers for.
-     * @return this instance of the Headers, once the header is added.
+     * @return this instance of the Headers, once the header is removed.
      * @throws IllegalStateException is thrown if headers are in a read-only state.
      */
     Headers remove(String key) throws IllegalStateException;
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..2a29d9da28ece 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
@@ -24,7 +24,8 @@
 import org.apache.kafka.common.utils.Utils;
 
 public class RecordHeader implements Header {
-    private final String key;
+    private ByteBuffer keyBuffer;
+    private String key;
     private ByteBuffer valueBuffer;
     private byte[] value;
 
@@ -34,13 +35,16 @@ public RecordHeader(String key, byte[] value) {
         this.value = value;
     }
 
-    public RecordHeader(String key, ByteBuffer valueBuffer) {
-        Objects.requireNonNull(key, "Null header keys are not permitted");
-        this.key = key;
+    public RecordHeader(ByteBuffer keyBuffer, ByteBuffer valueBuffer) {
+        this.keyBuffer = Objects.requireNonNull(keyBuffer, "Null header keys are not permitted");
         this.valueBuffer = valueBuffer;
     }
     
     public String key() {
+        if (key == null) {
+            key = Utils.utf8(keyBuffer, keyBuffer.remaining());
+            keyBuffer = null;
+        }
         return key;
     }
 
@@ -60,20 +64,20 @@ 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());
     }
 
     @Override
     public int hashCode() {
-        int result = key != null ? key.hashCode() : 0;
+        int result = key().hashCode();
         result = 31 * result + Arrays.hashCode(value());
         return result;
     }
 
     @Override
     public String toString() {
-        return "RecordHeader(key = " + key + ", value = " + Arrays.toString(value()) + ")";
+        return "RecordHeader(key = " + key() + ", value = " + Arrays.toString(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 141c9720c38fc..e6f05c1ab6887 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,19 +16,19 @@
  */
 package org.apache.kafka.common.header.internals;
 
-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.ArrayList;
+import java.util.Arrays;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Objects;
+
 public class RecordHeaders implements Headers {
-    
+
     private final List
         headers;
     private volatile boolean isReadOnly;
 
@@ -37,30 +37,27 @@ public RecordHeaders() {
     }
 
     public RecordHeaders(Header[] headers) {
-        if (headers == null) {
-            this.headers = new ArrayList<>();
-        } else {
-            this.headers = new ArrayList<>(Arrays.asList(headers));
-        }
+        this(headers == null ? null : Arrays.asList(headers));
     }
-    
+
     public RecordHeaders(Iterable
         headers) {
         //Use efficient copy constructor if possible, fallback to iteration otherwise
         if (headers == null) {
             this.headers = new ArrayList<>();
         } else if (headers instanceof RecordHeaders) {
             this.headers = new ArrayList<>(((RecordHeaders) headers).headers);
-        } else if (headers instanceof Collection) {
-            this.headers = new ArrayList<>((Collection
        ) headers);
         } else {
             this.headers = new ArrayList<>();
-            for (Header header : headers)
+            for (Header header : headers) {
+                Objects.requireNonNull(header, "Header cannot be null.");
                 this.headers.add(header);
+            }
         }
     }
 
     @Override
     public Headers add(Header header) throws IllegalStateException {
+        Objects.requireNonNull(header, "Header cannot be null.");
         canWrite();
         headers.add(header);
         return this;
@@ -99,12 +96,7 @@ public Header lastHeader(String key) {
     @Override
     public Iterable
         headers(final String key) {
         checkKey(key);
-        return new Iterable
        () {
-            @Override
-            public Iterator
         iterator() {
-                return new FilterByKeyIterator(headers.iterator(), key);
-            }
-        };
+        return () -> new FilterByKeyIterator(headers.iterator(), key);
     }
 
     @Override
@@ -119,17 +111,15 @@ public void setReadOnly() {
     public Header[] toArray() {
         return headers.isEmpty() ? Record.EMPTY_HEADERS : headers.toArray(new Header[headers.size()]);
     }
-    
+
     private void checkKey(String key) {
-        if (key == null) {
+        if (key == null)
             throw new IllegalArgumentException("key cannot be null.");
-        }
     }
-    
+
     private void canWrite() {
-        if (isReadOnly) {
+        if (isReadOnly)
             throw new IllegalStateException("RecordHeaders has been closed.");
-        }
     }
 
     private Iterator
         closeAware(final Iterator
         original) {
@@ -162,7 +152,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
@@ -177,7 +167,7 @@ public String toString() {
                ", isReadOnly = " + isReadOnly +
                ')';
     }
-    
+
     private static final class FilterByKeyIterator extends AbstractIterator
         {
 
         private final Iterator
         original;
@@ -187,14 +177,13 @@ private FilterByKeyIterator(Iterator
         original, String key) {
             this.original = original;
             this.key = key;
         }
-        
+
         protected Header makeNext() {
             while (true) {
                 if (original.hasNext()) {
                     Header header = original.next();
-                    if (!header.key().equals(key)) {
+                    if (!header.key().equals(key))
                         continue;
-                    }
 
                     return header;
                 }
diff --git a/clients/src/main/java/org/apache/kafka/common/internals/KafkaFutureImpl.java b/clients/src/main/java/org/apache/kafka/common/internals/KafkaFutureImpl.java
index 9ca019b2661ae..c0babe1634ab8 100644
--- a/clients/src/main/java/org/apache/kafka/common/internals/KafkaFutureImpl.java
+++ b/clients/src/main/java/org/apache/kafka/common/internals/KafkaFutureImpl.java
@@ -46,11 +46,11 @@ private static void wrapAndThrow(Throwable t) throws InterruptedException, Execu
         }
     }
 
-    private static class Applicant extends BiConsumer {
-        private final Function function;
+    private static class Applicant implements BiConsumer {
+        private final BaseFunction function;
         private final KafkaFutureImpl future;
 
-        Applicant(Function function, KafkaFutureImpl future) {
+        Applicant(BaseFunction function, KafkaFutureImpl future) {
             this.function = function;
             this.future = future;
         }
@@ -70,7 +70,7 @@ public void accept(A a, Throwable exception) {
         }
     }
 
-    private static class SingleWaiter extends BiConsumer {
+    private static class SingleWaiter implements BiConsumer {
         private R value = null;
         private Throwable exception = null;
         private boolean done = false;
@@ -140,13 +140,62 @@ R await(long timeout, TimeUnit unit)
      * futures's result as the argument to the supplied function.
      */
     @Override
-    public  KafkaFuture thenApply(Function function) {
-        KafkaFutureImpl future = new KafkaFutureImpl();
+    public  KafkaFuture thenApply(BaseFunction function) {
+        KafkaFutureImpl future = new KafkaFutureImpl<>();
         addWaiter(new Applicant<>(function, future));
         return future;
     }
 
+    public  void copyWith(KafkaFuture future, BaseFunction function) {
+        KafkaFutureImpl futureImpl = (KafkaFutureImpl) future;
+        futureImpl.addWaiter(new Applicant<>(function, this));
+    }
+
+    /**
+     * @see KafkaFutureImpl#thenApply(BaseFunction)
+     */
+    @Override
+    public  KafkaFuture thenApply(Function function) {
+        return thenApply((BaseFunction) function);
+    }
+
+    private static class WhenCompleteBiConsumer implements BiConsumer {
+        private final KafkaFutureImpl future;
+        private final BiConsumer biConsumer;
+
+        WhenCompleteBiConsumer(KafkaFutureImpl future, BiConsumer biConsumer) {
+            this.future = future;
+            this.biConsumer = biConsumer;
+        }
+
+        @Override
+        public void accept(T val, Throwable exception) {
+            try {
+                if (exception != null) {
+                    biConsumer.accept(null, exception);
+                } else {
+                    biConsumer.accept(val, null);
+                }
+            } catch (Throwable e) {
+                if (exception == null) {
+                    exception = e;
+                }
+            }
+            if (exception != null) {
+                future.completeExceptionally(exception);
+            } else {
+                future.complete(val);
+            }
+        }
+    }
+
     @Override
+    public KafkaFuture whenComplete(final BiConsumer biConsumer) {
+        final KafkaFutureImpl future = new KafkaFutureImpl<>();
+        addWaiter(new WhenCompleteBiConsumer<>(future, biConsumer));
+        return future;
+    }
+
     protected synchronized void addWaiter(BiConsumer action) {
         if (exception != null) {
             action.accept(null, exception);
@@ -159,7 +208,7 @@ protected synchronized void addWaiter(BiConsumer a
 
     @Override
     public synchronized boolean complete(T newValue) {
-        List> oldWaiters = null;
+        List> oldWaiters;
         synchronized (this) {
             if (done)
                 return false;
@@ -176,7 +225,7 @@ public synchronized boolean complete(T newValue) {
 
     @Override
     public boolean completeExceptionally(Throwable newException) {
-        List> oldWaiters = null;
+        List> oldWaiters;
         synchronized (this) {
             if (done)
                 return false;
@@ -198,9 +247,7 @@ public boolean completeExceptionally(Throwable newException) {
      */
     @Override
     public synchronized boolean cancel(boolean mayInterruptIfRunning) {
-        if (completeExceptionally(new CancellationException()))
-            return true;
-        return exception instanceof CancellationException;
+        return completeExceptionally(new CancellationException()) || exception instanceof CancellationException;
     }
 
     /**
@@ -208,7 +255,7 @@ public synchronized boolean cancel(boolean mayInterruptIfRunning) {
      */
     @Override
     public T get() throws InterruptedException, ExecutionException {
-        SingleWaiter waiter = new SingleWaiter();
+        SingleWaiter waiter = new SingleWaiter<>();
         addWaiter(waiter);
         return waiter.await();
     }
@@ -220,7 +267,7 @@ public T get() throws InterruptedException, ExecutionException {
     @Override
     public T get(long timeout, TimeUnit unit) throws InterruptedException, ExecutionException,
             TimeoutException {
-        SingleWaiter waiter = new SingleWaiter();
+        SingleWaiter waiter = new SingleWaiter<>();
         addWaiter(waiter);
         return waiter.await(timeout, unit);
     }
@@ -243,7 +290,7 @@ public synchronized T getNow(T valueIfAbsent) throws InterruptedException, Execu
      */
     @Override
     public synchronized boolean isCancelled() {
-        return (exception != null) && (exception instanceof CancellationException);
+        return exception instanceof CancellationException;
     }
 
     /**
@@ -261,4 +308,9 @@ public synchronized boolean isCompletedExceptionally() {
     public synchronized boolean isDone() {
         return done;
     }
+
+    @Override
+    public String toString() {
+        return String.format("KafkaFuture{value=%s,exception=%s,done=%b}", value, exception, done);
+    }
 }
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 605372c6ec7e9..e19126d36d32f 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
@@ -19,12 +19,14 @@
 import org.apache.kafka.common.TopicPartition;
 
 import java.util.ArrayList;
-import java.util.HashSet;
+import java.util.Collections;
+import java.util.Iterator;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
+import java.util.function.BiConsumer;
 
 /**
  * This class is a useful building block for doing fetch requests where topic partitions have to be rotated via
@@ -36,10 +38,17 @@
  * topic would "wrap around" and appear twice. However, as partitions are fetched in different orders and partition
  * leadership changes, we will deviate from the optimal. If this turns out to be an issue in practice, we can improve
  * it by tracking the partitions per node or calling `set` every so often.
+ *
+ * Note that this class is not thread-safe with the exception of {@link #size()} which returns the number of
+ * partitions currently tracked.
  */
 public class PartitionStates {
 
     private final LinkedHashMap map = new LinkedHashMap<>();
+    private final Set partitionSetView = Collections.unmodifiableSet(map.keySet());
+
+    /* the number of partitions that are currently assigned available in a thread safe manner */
+    private volatile int size = 0;
 
     public PartitionStates() {}
 
@@ -52,36 +61,41 @@ public void moveToEnd(TopicPartition topicPartition) {
     public void updateAndMoveToEnd(TopicPartition topicPartition, S state) {
         map.remove(topicPartition);
         map.put(topicPartition, state);
+        updateSize();
     }
 
     public void remove(TopicPartition topicPartition) {
         map.remove(topicPartition);
+        updateSize();
     }
 
     /**
-     * Returns the partitions in random order.
+     * Returns an unmodifiable view of the partitions in random order.
+     * changes to this PartitionStates instance will be reflected in this view.
      */
     public Set partitionSet() {
-        return new HashSet<>(map.keySet());
+        return partitionSetView;
     }
 
     public void clear() {
         map.clear();
+        updateSize();
     }
 
     public boolean contains(TopicPartition topicPartition) {
         return map.containsKey(topicPartition);
     }
 
-    /**
-     * Returns the partition states in order.
-     */
-    public List> partitionStates() {
-        List> result = new ArrayList<>();
-        for (Map.Entry entry : map.entrySet()) {
-            result.add(new PartitionState<>(entry.getKey(), entry.getValue()));
-        }
-        return result;
+    public Iterator stateIterator() {
+        return map.values().iterator();
+    }
+
+    public void forEach(BiConsumer biConsumer) {
+        map.forEach(biConsumer);
+    }
+
+    public Map partitionStateMap() {
+        return Collections.unmodifiableMap(map);
     }
 
     /**
@@ -95,8 +109,11 @@ public S stateValue(TopicPartition topicPartition) {
         return map.get(topicPartition);
     }
 
+    /**
+     * Get the number of partitions that are currently being tracked. This is thread-safe.
+     */
     public int size() {
-        return map.size();
+        return size;
     }
 
     /**
@@ -108,16 +125,17 @@ public int size() {
     public void set(Map partitionToState) {
         map.clear();
         update(partitionToState);
+        updateSize();
+    }
+
+    private void updateSize() {
+        size = map.size();
     }
 
     private void update(Map partitionToState) {
         LinkedHashMap> topicToPartitions = new LinkedHashMap<>();
         for (TopicPartition tp : partitionToState.keySet()) {
-            List partitions = topicToPartitions.get(tp.topic());
-            if (partitions == null) {
-                partitions = new ArrayList<>();
-                topicToPartitions.put(tp.topic(), partitions);
-            }
+            List partitions = topicToPartitions.computeIfAbsent(tp.topic(), k -> new ArrayList<>());
             partitions.add(tp);
         }
         for (Map.Entry> entry : topicToPartitions.entrySet()) {
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 041d1c2f420e4..18f8ffe91c66e 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
@@ -77,7 +77,7 @@ protected void bufferToBeReleased(ByteBuffer justReleased) {
     }
 
     @Override
-    public void close() throws Exception {
+    public void close() {
         alive = false;
         gcListenerThread.interrupt();
     }
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 0c49224657cd4..3867091db69f3 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
@@ -16,40 +16,62 @@
  */
 package org.apache.kafka.common.metrics;
 
-import java.lang.management.ManagementFactory;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
+import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.MetricName;
+import org.apache.kafka.common.config.ConfigException;
+import org.apache.kafka.common.utils.ConfigUtils;
+import org.apache.kafka.common.utils.Sanitizer;
+import org.apache.kafka.common.utils.Utils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import javax.management.Attribute;
 import javax.management.AttributeList;
 import javax.management.AttributeNotFoundException;
 import javax.management.DynamicMBean;
-import javax.management.InvalidAttributeValueException;
 import javax.management.JMException;
 import javax.management.MBeanAttributeInfo;
-import javax.management.MBeanException;
 import javax.management.MBeanInfo;
 import javax.management.MBeanServer;
 import javax.management.MalformedObjectNameException;
 import javax.management.ObjectName;
-import javax.management.ReflectionException;
-
-import org.apache.kafka.common.KafkaException;
-import org.apache.kafka.common.MetricName;
-import org.apache.kafka.common.utils.Sanitizer;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
+import java.lang.management.ManagementFactory;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.function.Predicate;
+import java.util.regex.Pattern;
+import java.util.regex.PatternSyntaxException;
 
 /**
  * Register metrics in JMX as dynamic mbeans based on the metric names
  */
 public class JmxReporter implements MetricsReporter {
 
+    public static final String METRICS_CONFIG_PREFIX = "metrics.jmx.";
+
+    public static final String EXCLUDE_CONFIG = METRICS_CONFIG_PREFIX + "exclude";
+    public static final String EXCLUDE_CONFIG_ALIAS = METRICS_CONFIG_PREFIX + "blacklist";
+
+    public static final String INCLUDE_CONFIG = METRICS_CONFIG_PREFIX + "include";
+    public static final String INCLUDE_CONFIG_ALIAS = METRICS_CONFIG_PREFIX + "whitelist";
+
+
+    public static final Set RECONFIGURABLE_CONFIGS = Utils.mkSet(INCLUDE_CONFIG,
+                                                                         INCLUDE_CONFIG_ALIAS,
+                                                                         EXCLUDE_CONFIG,
+                                                                         EXCLUDE_CONFIG_ALIAS);
+
+    public static final String DEFAULT_INCLUDE = ".*";
+    public static final String DEFAULT_EXCLUDE = "";
+
     private static final Logger log = LoggerFactory.getLogger(JmxReporter.class);
     private static final Object LOCK = new Object();
     private String prefix;
-    private final Map mbeans = new HashMap();
+    private final Map mbeans = new HashMap<>();
+    private Predicate mbeanPredicate = s -> true;
 
     public JmxReporter() {
         this("");
@@ -57,55 +79,98 @@ public JmxReporter() {
 
     /**
      * Create a JMX reporter that prefixes all metrics with the given string.
+     *  @deprecated Since 2.6.0. Use {@link JmxReporter#JmxReporter()}
+     *  Initialize JmxReporter with {@link JmxReporter#contextChange(MetricsContext)}
+     *  Populate prefix by adding _namespace/prefix key value pair to {@link MetricsContext}
      */
+    @Deprecated
     public JmxReporter(String prefix) {
-        this.prefix = prefix;
+        this.prefix = prefix != null ? prefix : "";
     }
 
     @Override
-    public void configure(Map configs) {}
+    public void configure(Map configs) {
+        reconfigure(configs);
+    }
+
+    @Override
+    public Set reconfigurableConfigs() {
+        return RECONFIGURABLE_CONFIGS;
+    }
+
+    @Override
+    public void validateReconfiguration(Map configs) throws ConfigException {
+        compilePredicate(configs);
+    }
+
+    @Override
+    public void reconfigure(Map configs) {
+        synchronized (LOCK) {
+            this.mbeanPredicate = JmxReporter.compilePredicate(configs);
+
+            mbeans.forEach((name, mbean) -> {
+                if (mbeanPredicate.test(name)) {
+                    reregister(mbean);
+                } else {
+                    unregister(mbean);
+                }
+            });
+        }
+    }
 
     @Override
     public void init(List metrics) {
         synchronized (LOCK) {
             for (KafkaMetric metric : metrics)
                 addAttribute(metric);
-            for (KafkaMbean mbean : mbeans.values())
-                reregister(mbean);
+
+            mbeans.forEach((name, mbean) -> {
+                if (mbeanPredicate.test(name)) {
+                    reregister(mbean);
+                }
+            });
         }
     }
 
+    public boolean containsMbean(String mbeanName) {
+        return mbeans.containsKey(mbeanName);
+    }
+
     @Override
     public void metricChange(KafkaMetric metric) {
         synchronized (LOCK) {
-            KafkaMbean mbean = addAttribute(metric);
-            reregister(mbean);
+            String mbeanName = addAttribute(metric);
+            if (mbeanName != null && mbeanPredicate.test(mbeanName)) {
+                reregister(mbeans.get(mbeanName));
+            }
         }
     }
 
     @Override
     public void metricRemoval(KafkaMetric metric) {
         synchronized (LOCK) {
-            KafkaMbean mbean = removeAttribute(metric);
+            MetricName metricName = metric.metricName();
+            String mBeanName = getMBeanName(prefix, metricName);
+            KafkaMbean mbean = removeAttribute(metric, mBeanName);
             if (mbean != null) {
-                if (mbean.metrics.isEmpty())
+                if (mbean.metrics.isEmpty()) {
                     unregister(mbean);
-                else
+                    mbeans.remove(mBeanName);
+                } else if (mbeanPredicate.test(mBeanName))
                     reregister(mbean);
             }
         }
     }
 
-    private KafkaMbean removeAttribute(KafkaMetric metric) {
+    private KafkaMbean removeAttribute(KafkaMetric metric, String mBeanName) {
         MetricName metricName = metric.metricName();
-        String mBeanName = getMBeanName(prefix, metricName);
         KafkaMbean mbean = this.mbeans.get(mBeanName);
         if (mbean != null)
             mbean.removeAttribute(metricName.name());
         return mbean;
     }
 
-    private KafkaMbean addAttribute(KafkaMetric metric) {
+    private String addAttribute(KafkaMetric metric) {
         try {
             MetricName metricName = metric.metricName();
             String mBeanName = getMBeanName(prefix, metricName);
@@ -113,7 +178,7 @@ private KafkaMbean addAttribute(KafkaMetric metric) {
                 mbeans.put(mBeanName, new KafkaMbean(mBeanName));
             KafkaMbean mbean = this.mbeans.get(mBeanName);
             mbean.setAttribute(metricName.name(), metric);
-            return mbean;
+            return mBeanName;
         } catch (JMException e) {
             throw new KafkaException("Error creating mbean attribute for metricName :" + metric.metricName(), e);
         }
@@ -169,7 +234,7 @@ private static class KafkaMbean implements DynamicMBean {
         private final ObjectName objectName;
         private final Map metrics;
 
-        public KafkaMbean(String mbeanName) throws MalformedObjectNameException {
+        KafkaMbean(String mbeanName) throws MalformedObjectNameException {
             this.metrics = new HashMap<>();
             this.objectName = new ObjectName(mbeanName);
         }
@@ -178,12 +243,12 @@ public ObjectName name() {
             return objectName;
         }
 
-        public void setAttribute(String name, KafkaMetric metric) {
+        void setAttribute(String name, KafkaMetric metric) {
             this.metrics.put(name, metric);
         }
 
         @Override
-        public Object getAttribute(String name) throws AttributeNotFoundException, MBeanException, ReflectionException {
+        public Object getAttribute(String name) throws AttributeNotFoundException {
             if (this.metrics.containsKey(name))
                 return this.metrics.get(name).metricValue();
             else
@@ -203,7 +268,7 @@ public AttributeList getAttributes(String[] names) {
             return list;
         }
 
-        public KafkaMetric removeAttribute(String name) {
+        KafkaMetric removeAttribute(String name) {
             return this.metrics.remove(name);
         }
 
@@ -226,15 +291,12 @@ public MBeanInfo getMBeanInfo() {
         }
 
         @Override
-        public Object invoke(String name, Object[] params, String[] sig) throws MBeanException, ReflectionException {
+        public Object invoke(String name, Object[] params, String[] sig) {
             throw new UnsupportedOperationException("Set not allowed.");
         }
 
         @Override
-        public void setAttribute(Attribute attribute) throws AttributeNotFoundException,
-                                                     InvalidAttributeValueException,
-                                                     MBeanException,
-                                                     ReflectionException {
+        public void setAttribute(Attribute attribute) {
             throw new UnsupportedOperationException("Set not allowed.");
         }
 
@@ -245,4 +307,50 @@ public AttributeList setAttributes(AttributeList list) {
 
     }
 
+    public static Predicate compilePredicate(Map originalConfig) {
+        Map configs = ConfigUtils.translateDeprecatedConfigs(
+            originalConfig, new String[][]{{INCLUDE_CONFIG, INCLUDE_CONFIG_ALIAS},
+                                           {EXCLUDE_CONFIG, EXCLUDE_CONFIG_ALIAS}});
+        String include = (String) configs.get(INCLUDE_CONFIG);
+        String exclude = (String) configs.get(EXCLUDE_CONFIG);
+
+        if (include == null) {
+            include = DEFAULT_INCLUDE;
+        }
+
+        if (exclude == null) {
+            exclude = DEFAULT_EXCLUDE;
+        }
+
+        try {
+            Pattern includePattern = Pattern.compile(include);
+            Pattern excludePattern = Pattern.compile(exclude);
+
+            return s -> includePattern.matcher(s).matches()
+                        && !excludePattern.matcher(s).matches();
+        } catch (PatternSyntaxException e) {
+            throw new ConfigException("JMX filter for configuration" + METRICS_CONFIG_PREFIX
+                                      + ".(include/exclude) is not a valid regular expression");
+        }
+    }
+
+    @Override
+    public void contextChange(MetricsContext metricsContext) {
+        String namespace = metricsContext.contextLabels().get(MetricsContext.NAMESPACE);
+        Objects.requireNonNull(namespace);
+        synchronized (LOCK) {
+            if (!mbeans.isEmpty()) {
+                throw new IllegalStateException("JMX MetricsContext can only be updated before JMX metrics are created");
+            }
+
+            // prevent prefix from getting reset back to empty for backwards compatibility
+            // with the deprecated JmxReporter(String prefix) constructor, in case contextChange gets called
+            // via one of the Metrics() constructor with a default empty MetricsContext()
+            if (namespace.isEmpty()) {
+                return;
+            }
+
+            prefix = namespace;
+        }
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/metrics/KafkaMetric.java b/clients/src/main/java/org/apache/kafka/common/metrics/KafkaMetric.java
index f04981aee39c9..48999e1057d9d 100644
--- a/clients/src/main/java/org/apache/kafka/common/metrics/KafkaMetric.java
+++ b/clients/src/main/java/org/apache/kafka/common/metrics/KafkaMetric.java
@@ -55,9 +55,7 @@ public MetricName metricName() {
     @Override
     @Deprecated
     public double value() {
-        synchronized (this.lock) {
-            return measurableValue(time.milliseconds());
-        }
+        return measurableValue(time.milliseconds());
     }
 
     @Override
@@ -81,10 +79,12 @@ public Measurable measurable() {
     }
 
     double measurableValue(long timeMs) {
-        if (this.metricValueProvider instanceof Measurable)
-            return ((Measurable) metricValueProvider).measure(config, timeMs);
-        else
-            return 0;
+        synchronized (this.lock) {
+            if (this.metricValueProvider instanceof Measurable)
+                return ((Measurable) metricValueProvider).measure(config, timeMs);
+            else
+                return 0;
+        }
     }
 
     public void config(MetricConfig config) {
diff --git a/clients/src/main/java/org/apache/kafka/common/metrics/KafkaMetricsContext.java b/clients/src/main/java/org/apache/kafka/common/metrics/KafkaMetricsContext.java
new file mode 100644
index 0000000000000..43eb8cb319465
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/metrics/KafkaMetricsContext.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.metrics;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * A implementation of MetricsContext, it encapsulates required metrics context properties for Kafka services and clients
+ */
+public class KafkaMetricsContext implements MetricsContext {
+    /**
+     * Client or Service's contextLabels map.
+     */
+    private final Map contextLabels = new HashMap<>();
+
+    /**
+     * Create a MetricsContext with namespace, no service or client properties
+     * @param namespace value for _namespace key
+     */
+    public KafkaMetricsContext(String namespace) {
+        this(namespace, new HashMap<>());
+    }
+
+    /**
+     * Create a MetricsContext with namespace, service or client properties
+     * @param namespace value for _namespace key
+     * @param contextLabels  contextLabels additional entries to add to the context.
+     *                  values will be converted to string using Object.toString()
+     */
+    public KafkaMetricsContext(String namespace, Map contextLabels) {
+        this.contextLabels.put(MetricsContext.NAMESPACE, namespace);
+        contextLabels.forEach((key, value) -> this.contextLabels.put(key, value != null ? value.toString() : null));
+    }
+
+    @Override
+    public Map contextLabels() {
+        return Collections.unmodifiableMap(contextLabels);
+    }
+
+}
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/MetricConfig.java b/clients/src/main/java/org/apache/kafka/common/metrics/MetricConfig.java
index 1fffee72c168a..7367e966c014e 100644
--- a/clients/src/main/java/org/apache/kafka/common/metrics/MetricConfig.java
+++ b/clients/src/main/java/org/apache/kafka/common/metrics/MetricConfig.java
@@ -33,7 +33,6 @@ public class MetricConfig {
     private Sensor.RecordingLevel recordingLevel;
 
     public MetricConfig() {
-        super();
         this.quota = null;
         this.samples = 2;
         this.eventWindow = Long.MAX_VALUE;
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 81270224e4bcf..52b7794a4c10b 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
@@ -18,9 +18,9 @@
 
 import org.apache.kafka.common.MetricName;
 import org.apache.kafka.common.MetricNameTemplate;
+import org.apache.kafka.common.metrics.internals.MetricsUtils;
 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,14 +33,16 @@
 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;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.ScheduledThreadPoolExecutor;
-import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.TimeUnit;
 
+import static java.util.Collections.emptyList;
+
 /**
  * A registry of sensors and metrics.
  * 
        
@@ -88,7 +90,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);
     }
 
     /**
@@ -96,7 +98,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);
     }
 
 
@@ -106,7 +108,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);
     }
 
     /**
@@ -120,6 +122,18 @@ public Metrics(MetricConfig defaultConfig, List reporters, Time
         this(defaultConfig, reporters, time, false);
     }
 
+    /**
+     * Create a metrics repository with a default config, metric reporters and metric context
+     * Expiration of Sensors is disabled.
+     * @param defaultConfig The default config
+     * @param reporters The metrics reporters
+     * @param time The time instance to use with the metrics
+     * @param metricsContext The metricsContext to initialize metrics reporter with
+     */
+    public Metrics(MetricConfig defaultConfig, List reporters, Time time, MetricsContext metricsContext) {
+        this(defaultConfig, reporters, time, false, metricsContext);
+    }
+
     /**
      * Create a metrics repository with a default config, given metric reporters and the ability to expire eligible sensors
      * @param defaultConfig The default config
@@ -128,36 +142,43 @@ public Metrics(MetricConfig defaultConfig, List reporters, Time
      * @param enableExpiration true if the metrics instance can garbage collect inactive sensors, false otherwise
      */
     public Metrics(MetricConfig defaultConfig, List reporters, Time time, boolean enableExpiration) {
+        this(defaultConfig, reporters, time, enableExpiration, new KafkaMetricsContext(""));
+    }
+
+    /**
+     * Create a metrics repository with a default config, given metric reporters, the ability to expire eligible sensors
+     * and MetricContext
+     * @param defaultConfig The default config
+     * @param reporters The metrics reporters
+     * @param time The time instance to use with the metrics
+     * @param enableExpiration true if the metrics instance can garbage collect inactive sensors, false otherwise
+     * @param metricsContext The metricsContext to initialize metrics reporter with
+     */
+    public Metrics(MetricConfig defaultConfig, List reporters, Time time, boolean enableExpiration,
+                   MetricsContext metricsContext) {
         this.config = defaultConfig;
         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());
+        for (MetricsReporter reporter : reporters) {
+            reporter.contextChange(metricsContext);
+            reporter.init(new ArrayList<>());
+        }
 
         // Create the ThreadPoolExecutor only if expiration of Sensors is enabled.
         if (enableExpiration) {
             this.metricsScheduler = new ScheduledThreadPoolExecutor(1);
             // Creating a daemon thread to not block shutdown
-            this.metricsScheduler.setThreadFactory(new ThreadFactory() {
-                public Thread newThread(Runnable runnable) {
-                    return KafkaThread.daemon("SensorExpiryThread", runnable);
-                }
-            });
+            this.metricsScheduler.setThreadFactory(runnable -> KafkaThread.daemon("SensorExpiryThread", runnable));
             this.metricsScheduler.scheduleAtFixedRate(new ExpireSensorTask(), 30, 30, TimeUnit.SECONDS);
         } else {
             this.metricsScheduler = null;
         }
 
         addMetric(metricName("count", "kafka-metrics-count", "total number of registered metrics"),
-            new Measurable() {
-                @Override
-                public double measure(MetricConfig config, long now) {
-                    return metrics.size();
-                }
-            });
+            (config, now) -> metrics.size());
     }
 
     /**
@@ -184,7 +205,7 @@ public MetricName metricName(String name, String group, String description, Map<
      * @param description A human-readable description to include in the metric
      */
     public MetricName metricName(String name, String group, String description) {
-        return metricName(name, group, description, new HashMap());
+        return metricName(name, group, description, new HashMap<>());
     }
 
     /**
@@ -194,7 +215,7 @@ public MetricName metricName(String name, String group, String description) {
      * @param group       logical group name of the metrics to which this metric belongs
      */
     public MetricName metricName(String name, String group) {
-        return metricName(name, group, "", new HashMap());
+        return metricName(name, group, "", new HashMap<>());
     }
 
     /**
@@ -207,7 +228,7 @@ public MetricName metricName(String name, String group) {
      * @param keyValue      additional key/value attributes of the metric (must come in pairs)
      */
     public MetricName metricName(String name, String group, String description, String... keyValue) {
-        return metricName(name, group, description, getTags(keyValue));
+        return metricName(name, group, description, MetricsUtils.getTags(keyValue));
     }
 
     /**
@@ -222,22 +243,22 @@ public MetricName metricName(String name, String group, Map tags
         return metricName(name, group, "", tags);
     }
 
-    private static Map getTags(String... keyValue) {
-        if ((keyValue.length % 2) != 0)
-            throw new IllegalArgumentException("keyValue needs to be specified in pairs");
-        Map tags = new HashMap();
-
-        for (int i = 0; i < keyValue.length; i += 2)
-            tags.put(keyValue[i], keyValue[i + 1]);
-        return tags;
-    }
-
-    public static String toHtmlTable(String domain, List allMetrics) {
-        Map> beansAndAttributes = new TreeMap>();
+    /**
+     * Use the specified domain and metric name templates to generate an HTML table documenting the metrics. A separate table section
+     * will be generated for each of the MBeans and the associated attributes. The MBean names are lexicographically sorted to
+     * determine the order of these sections. This order is therefore dependent upon the order of the
+     * tags in each {@link MetricNameTemplate}.
+     *
+     * @param domain the domain or prefix for the JMX MBean names; may not be null
+     * @param allMetrics the collection of all {@link MetricNameTemplate} instances each describing one metric; may not be null
+     * @return the string containing the HTML table; never null
+     */
+    public static String toHtmlTable(String domain, Iterable allMetrics) {
+        Map> beansAndAttributes = new TreeMap<>();
     
         try (Metrics metrics = new Metrics()) {
             for (MetricNameTemplate template : allMetrics) {
-                Map tags = new TreeMap();
+                Map tags = new LinkedHashMap<>();
                 for (String s : template.tags()) {
                     tags.put(s, "{" + s + "}");
                 }
@@ -245,7 +266,7 @@ public static String toHtmlTable(String domain, List allMetr
                 MetricName metricName = metrics.metricName(template.name(), template.group(), template.description(), tags);
                 String mBeanName = JmxReporter.getMBeanName(domain, metricName);
                 if (!beansAndAttributes.containsKey(mBeanName)) {
-                    beansAndAttributes.put(mBeanName, new TreeMap());
+                    beansAndAttributes.put(mBeanName, new TreeMap<>());
                 }
                 Map attrAndDesc = beansAndAttributes.get(mBeanName);
                 if (!attrAndDesc.containsKey(template.name())) {
@@ -301,7 +322,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));
     }
 
     /**
@@ -393,15 +414,11 @@ public synchronized Sensor sensor(String name, MetricConfig config, long inactiv
             this.sensors.put(name, s);
             if (parents != null) {
                 for (Sensor parent : parents) {
-                    List children = childrenSensors.get(parent);
-                    if (children == null) {
-                        children = new ArrayList<>();
-                        childrenSensors.put(parent, children);
-                    }
+                    List children = childrenSensors.computeIfAbsent(parent, k -> new ArrayList<>());
                     children.add(s);
                 }
             }
-            log.debug("Added sensor with name {}", name);
+            log.trace("Added sensor with name {}", name);
         }
         return s;
     }
@@ -434,8 +451,11 @@ 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);
+                        }
                     }
                 }
             }
@@ -451,7 +471,7 @@ public void removeSensor(String name) {
      * This is a way to expose existing values as metrics.
      *
      * This method is kept for binary compatibility purposes, it has the same behaviour as
-     * {@link #addMetric(MetricName, MetricValue)}.
+     * {@link #addMetric(MetricName, MetricValueProvider)}.
      *
      * @param metricName The name of the metric
      * @param measurable The measurable that will be measured by this metric
@@ -477,15 +497,16 @@ public void addMetric(MetricName metricName, MetricConfig config, Measurable mea
 
     /**
      * Add a metric to monitor an object that implements MetricValueProvider. This metric won't be associated with any
-     * sensor. This is a way to expose existing values as metrics.
+     * sensor. This is a way to expose existing values as metrics. User is expected to add any additional
+     * synchronization to update and access metric values, if required.
      *
      * @param metricName The name of the metric
      * @param metricValueProvider The metric value provider associated with this metric
      */
     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);
@@ -493,7 +514,8 @@ public void addMetric(MetricName metricName, MetricConfig config, MetricValuePro
 
     /**
      * Add a metric to monitor an object that implements MetricValueProvider. This metric won't be associated with any
-     * sensor. This is a way to expose existing values as metrics.
+     * sensor. This is a way to expose existing values as metrics. User is expected to add any additional
+     * synchronization to update and access metric values, if required.
      *
      * @param metricName The name of the metric
      * @param metricValueProvider The metric value provider associated with this metric
@@ -512,8 +534,14 @@ public void addMetric(MetricName metricName, MetricValueProvider metricValueP
     public synchronized KafkaMetric removeMetric(MetricName metricName) {
         KafkaMetric metric = this.metrics.remove(metricName);
         if (metric != null) {
-            for (MetricsReporter reporter : reporters)
-                reporter.metricRemoval(metric);
+            for (MetricsReporter reporter : reporters) {
+                try {
+                    reporter.metricRemoval(metric);
+                } catch (Exception e) {
+                    log.error("Error when removing metric from " + reporter.getClass().getName(), e);
+                }
+            }
+            log.trace("Removed metric named {}", metricName);
         }
         return metric;
     }
@@ -522,17 +550,32 @@ 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);
     }
 
+    /**
+     * Remove a MetricReporter
+     */
+    public synchronized void removeReporter(MetricsReporter reporter) {
+        if (this.reporters.remove(reporter)) {
+            reporter.close();
+        }
+    }
+
     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.");
         this.metrics.put(metricName, metric);
-        for (MetricsReporter reporter : reporters)
-            reporter.metricChange(metric);
+        for (MetricsReporter reporter : reporters) {
+            try {
+                reporter.metricChange(metric);
+            } catch (Exception e) {
+                log.error("Error when registering metric on " + reporter.getClass().getName(), e);
+            }
+        }
+        log.trace("Registered metric named {}", metricName);
     }
 
     /**
@@ -555,6 +598,7 @@ public KafkaMetric metric(MetricName metricName) {
      * Package private for testing
      */
     class ExpireSensorTask implements Runnable {
+        @Override
         public void run() {
             for (Map.Entry sensorEntry : sensors.entrySet()) {
                 // removeSensor also locks the sensor object. This is fine because synchronized is reentrant
@@ -580,7 +624,7 @@ Map> childrenSensors() {
     }
 
     public MetricName metricInstance(MetricNameTemplate template, String... keyValue) {
-        return metricInstance(template, getTags(keyValue));
+        return metricInstance(template, MetricsUtils.getTags(keyValue));
     }
 
     public MetricName metricInstance(MetricNameTemplate template, Map tags) {
@@ -591,7 +635,7 @@ public MetricName metricInstance(MetricNameTemplate template, Map templateTagKeys = template.tags();
         
         if (!runtimeTagKeys.equals(templateTagKeys)) {
-            throw new IllegalArgumentException("For '" + template.name() + "', runtime-defined metric tags do not match the tags in the template. " + ""
+            throw new IllegalArgumentException("For '" + template.name() + "', runtime-defined metric tags do not match the tags in the template. "
                     + "Runtime = " + runtimeTagKeys.toString() + " Template = " + templateTagKeys.toString());
         }
                 
@@ -612,9 +656,16 @@ public void close() {
                 Thread.currentThread().interrupt();
             }
         }
+        log.info("Metrics scheduler closed");
 
-        for (MetricsReporter reporter : this.reporters)
-            reporter.close();
+        for (MetricsReporter reporter : reporters) {
+            try {
+                log.info("Closing reporter {}", reporter.getClass().getName());
+                reporter.close();
+            } catch (Exception e) {
+                log.error("Error when closing " + reporter.getClass().getName(), e);
+            }
+        }
+        log.info("Metrics reporters closed");
     }
-
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/metrics/MetricsContext.java b/clients/src/main/java/org/apache/kafka/common/metrics/MetricsContext.java
new file mode 100644
index 0000000000000..dcac5c289da6d
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/metrics/MetricsContext.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.metrics;
+
+import org.apache.kafka.common.annotation.InterfaceStability;
+
+import java.util.Map;
+
+/**
+ * MetricsContext encapsulates additional contextLabels about metrics exposed via a
+ * {@link org.apache.kafka.common.metrics.MetricsReporter}
+ *
+ * The contextLabels map provides following information:
+ * - a _namespace field indicating the component exposing metrics
+ *   e.g. kafka.server, kafka.consumer
+ *   {@link JmxReporter} uses this as prefix for mbean names
+ *
+ * - for clients and streams libraries: any freeform fields passed in via
+ *   client properties in the form of `metrics.context.=
+ *
+ * - for kafka brokers: kafka.broker.id, kafka.cluster.id
+ * - for connect workers: connect.kafka.cluster.id, connect.group.id
+ */
+@InterfaceStability.Evolving
+public interface MetricsContext {
+    /* predefined fields */
+    String NAMESPACE = "_namespace"; // metrics namespace, formerly jmx prefix
+
+    /**
+     * Returns the labels for this metrics context.
+     *
+     * @return the map of label keys and values; never null but possibly empty
+     */
+    Map contextLabels();
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/metrics/MetricsReporter.java b/clients/src/main/java/org/apache/kafka/common/metrics/MetricsReporter.java
index 995bdaa6dce3a..75771fb4acedc 100644
--- a/clients/src/main/java/org/apache/kafka/common/metrics/MetricsReporter.java
+++ b/clients/src/main/java/org/apache/kafka/common/metrics/MetricsReporter.java
@@ -16,38 +16,62 @@
  */
 package org.apache.kafka.common.metrics;
 
+import java.util.Collections;
 import java.util.List;
+import java.util.Map;
+import java.util.Set;
 
-import org.apache.kafka.common.Configurable;
+import org.apache.kafka.common.Reconfigurable;
+import org.apache.kafka.common.annotation.InterfaceStability;
+import org.apache.kafka.common.config.ConfigException;
 
 /**
  * A plugin interface to allow things to listen as new metrics are created so they can be reported.
  * 
        
  * 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.
  */
-public interface MetricsReporter extends Configurable {
+public interface MetricsReporter extends Reconfigurable, AutoCloseable {
 
     /**
      * This is called when the reporter is first registered to initially register all existing metrics
      * @param metrics All currently existing metrics
      */
-    public void init(List 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();
 
+    // default methods for backwards compatibility with reporters that only implement Configurable
+    default Set reconfigurableConfigs() {
+        return Collections.emptySet();
+    }
+
+    default void validateReconfiguration(Map configs) throws ConfigException {
+    }
+
+    default void reconfigure(Map configs) {
+    }
+
+    /**
+     * Sets the context labels for the service or library exposing metrics. This will be called before {@link #init(List)} and may be called anytime after that.
+     *
+     * @param metricsContext the metric context
+     */
+    @InterfaceStability.Evolving
+    default void contextChange(MetricsContext metricsContext) {
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/metrics/QuotaViolationException.java b/clients/src/main/java/org/apache/kafka/common/metrics/QuotaViolationException.java
index 4c970f26aea87..7068d3115d795 100644
--- a/clients/src/main/java/org/apache/kafka/common/metrics/QuotaViolationException.java
+++ b/clients/src/main/java/org/apache/kafka/common/metrics/QuotaViolationException.java
@@ -17,7 +17,6 @@
 package org.apache.kafka.common.metrics;
 
 import org.apache.kafka.common.KafkaException;
-import org.apache.kafka.common.MetricName;
 
 /**
  * Thrown when a sensor records a value that causes a metric to go outside the bounds configured as its quota
@@ -25,23 +24,18 @@
 public class QuotaViolationException extends KafkaException {
 
     private static final long serialVersionUID = 1L;
-    private final MetricName metricName;
+    private final KafkaMetric metric;
     private final double value;
     private final double bound;
 
-    public QuotaViolationException(MetricName metricName, double value, double bound) {
-        super(String.format(
-                "'%s' violated quota. Actual: %f, Threshold: %f",
-                metricName,
-                value,
-                bound));
-        this.metricName = metricName;
+    public QuotaViolationException(KafkaMetric metric, double value, double bound) {
+        this.metric = metric;
         this.value = value;
         this.bound = bound;
     }
 
-    public MetricName metricName() {
-        return metricName;
+    public KafkaMetric metric() {
+        return metric;
     }
 
     public double value() {
@@ -51,4 +45,21 @@ public double value() {
     public double bound() {
         return bound;
     }
+
+    @Override
+    public String toString() {
+        return getClass().getName()
+                + ": '"
+                + metric.metricName()
+                + "' violated quota. Actual: "
+                + value
+                + ", Threshold: "
+                + bound;
+    }
+
+    /* avoid the expensive and stack trace for quota violation exceptions */
+    @Override
+    public Throwable fillInStackTrace() {
+        return this;
+    }
 }
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 321fab661cd23..16041e2d66eff 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
@@ -16,19 +16,25 @@
  */
 package org.apache.kafka.common.metrics;
 
+import java.util.function.Supplier;
 import org.apache.kafka.common.MetricName;
 import org.apache.kafka.common.metrics.CompoundStat.NamedMeasurable;
+import org.apache.kafka.common.metrics.stats.TokenBucket;
 import org.apache.kafka.common.utils.Time;
-import org.apache.kafka.common.utils.Utils;
 
 import java.util.ArrayList;
-import java.util.Collections;
 import java.util.HashSet;
+import java.util.LinkedHashMap;
 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;
 
+import static java.util.Arrays.asList;
+import static java.util.Collections.unmodifiableList;
+
 /**
  * A sensor applies a continuous sequence of numerical values to a set of associated metrics. For example a sensor on
  * message size would record a sequence of message sizes using the {@link #record(double)} api and would maintain a set
@@ -39,15 +45,34 @@ public final class Sensor {
     private final Metrics registry;
     private final String name;
     private final Sensor[] parents;
-    private final List stats;
-    private final List metrics;
+    private final List stats;
+    private final Map metrics;
     private final MetricConfig config;
     private final Time time;
     private volatile long lastRecordTime;
     private final long inactiveSensorExpirationTimeMs;
+    private final Object metricLock;
+
+    private static class StatAndConfig {
+        private final Stat stat;
+        private final Supplier configSupplier;
+
+        StatAndConfig(Stat stat, Supplier configSupplier) {
+            this.stat = stat;
+            this.configSupplier = configSupplier;
+        }
+
+        public Stat stat() {
+            return stat;
+        }
+
+        public MetricConfig config() {
+            return configSupplier.get();
+        }
+    }
 
     public enum RecordingLevel {
-        INFO(0, "INFO"), DEBUG(1, "DEBUG");
+        INFO(0, "INFO"), DEBUG(1, "DEBUG"), TRACE(2, "TRACE");
 
         private static final RecordingLevel[] ID_TO_TYPE;
         private static final int MIN_RECORDING_LEVEL_KEY = 0;
@@ -79,8 +104,8 @@ public enum RecordingLevel {
 
         public static RecordingLevel forId(int id) {
             if (id < MIN_RECORDING_LEVEL_KEY || id > MAX_RECORDING_LEVEL_KEY)
-                throw new IllegalArgumentException(String.format("Unexpected RecordLevel id `%s`, it should be between `%s` " +
-                    "and `%s` (inclusive)", id, MIN_RECORDING_LEVEL_KEY, MAX_RECORDING_LEVEL_KEY));
+                throw new IllegalArgumentException(String.format("Unexpected RecordLevel id `%d`, it should be between `%d` " +
+                    "and `%d` (inclusive)", id, MIN_RECORDING_LEVEL_KEY, MAX_RECORDING_LEVEL_KEY));
             return ID_TO_TYPE[id];
         }
 
@@ -90,13 +115,16 @@ public static RecordingLevel forName(String name) {
         }
 
         public boolean shouldRecord(final int configId) {
-            if (configId == DEBUG.id) {
+            if (configId == INFO.id) {
+                return this.id == INFO.id;
+            } else if (configId == DEBUG.id) {
+                return this.id == INFO.id || this.id == DEBUG.id;
+            } else if (configId == TRACE.id) {
                 return true;
             } else {
-                return configId == this.id;
+                throw new IllegalStateException("Did not recognize recording level " + configId);
             }
         }
-
     }
 
     private final RecordingLevel recordingLevel;
@@ -105,16 +133,17 @@ 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 ArrayList<>();
+        this.metrics = new LinkedHashMap<>();
         this.stats = new ArrayList<>();
         this.config = config;
         this.time = time;
         this.inactiveSensorExpirationTimeMs = TimeUnit.MILLISECONDS.convert(inactiveSensorExpirationTimeSeconds, TimeUnit.SECONDS);
         this.lastRecordTime = time.milliseconds();
         this.recordingLevel = recordingLevel;
-        checkForest(new HashSet());
+        this.metricLock = new Object();
+        checkForest(new HashSet<>());
     }
 
     /* Validate that this sensor doesn't end up referencing itself */
@@ -132,13 +161,8 @@ public String name() {
         return this.name;
     }
 
-    /**
-     * Record an occurrence, this is just short-hand for {@link #record(double) record(1.0)}
-     */
-    public void record() {
-        if (shouldRecord()) {
-            record(1.0);
-        }
+    List parents() {
+        return unmodifiableList(asList(parents));
     }
 
     /**
@@ -147,6 +171,16 @@ public void record() {
     public boolean shouldRecord() {
         return this.recordingLevel.shouldRecord(config.recordLevel().id);
     }
+
+    /**
+     * Record an occurrence, this is just short-hand for {@link #record(double) record(1.0)}
+     */
+    public void record() {
+        if (shouldRecord()) {
+            recordInternal(1.0d, time.milliseconds(), true);
+        }
+    }
+
     /**
      * Record a value with this sensor
      * @param value The value to record
@@ -155,7 +189,7 @@ public boolean shouldRecord() {
      */
     public void record(double value) {
         if (shouldRecord()) {
-            record(value, time.milliseconds());
+            recordInternal(value, time.milliseconds(), true);
         }
     }
 
@@ -168,22 +202,40 @@ public void record(double value) {
      *         bound
      */
     public void record(double value, long timeMs) {
-        record(value, timeMs, true);
+        if (shouldRecord()) {
+            recordInternal(value, timeMs, true);
+        }
     }
 
+    /**
+     * Record a value at a known time. This method is slightly faster than {@link #record(double)} since it will reuse
+     * the time stamp.
+     * @param value The value we are recording
+     * @param timeMs The current POSIX time in milliseconds
+     * @param checkQuotas Indicate if quota must be enforced or not
+     * @throws QuotaViolationException if recording this value moves a metric beyond its configured maximum or minimum
+     *         bound
+     */
     public void record(double value, long timeMs, boolean checkQuotas) {
         if (shouldRecord()) {
-            this.lastRecordTime = timeMs;
-            synchronized (this) {
+            recordInternal(value, timeMs, checkQuotas);
+        }
+    }
+
+    private void recordInternal(double value, long timeMs, boolean checkQuotas) {
+        this.lastRecordTime = timeMs;
+        synchronized (this) {
+            synchronized (metricLock()) {
                 // increment all the stats
-                for (Stat stat : this.stats)
-                    stat.record(config, value, timeMs);
-                if (checkQuotas)
-                    checkQuotas(timeMs);
+                for (StatAndConfig statAndConfig : this.stats) {
+                    statAndConfig.stat.record(statAndConfig.config(), value, timeMs);
+                }
             }
-            for (Sensor parent : parents)
-                parent.record(value, timeMs, checkQuotas);
+            if (checkQuotas)
+                checkQuotas(timeMs);
         }
+        for (Sensor parent : parents)
+            parent.record(value, timeMs, checkQuotas);
     }
 
     /**
@@ -194,15 +246,20 @@ public void checkQuotas() {
     }
 
     public void checkQuotas(long timeMs) {
-        for (KafkaMetric metric : this.metrics) {
+        for (KafkaMetric metric : this.metrics.values()) {
             MetricConfig config = metric.config();
             if (config != null) {
                 Quota quota = config.quota();
                 if (quota != null) {
                     double value = metric.measurableValue(timeMs);
-                    if (!quota.acceptable(value)) {
-                        throw new QuotaViolationException(metric.metricName(), value,
-                            quota.bound());
+                    if (metric.measurable() instanceof TokenBucket) {
+                        if (value < 0) {
+                            throw new QuotaViolationException(metric, value, quota.bound());
+                        }
+                    } else {
+                        if (!quota.acceptable(value)) {
+                            throw new QuotaViolationException(metric, value, quota.bound());
+                        }
                     }
                 }
             }
@@ -211,9 +268,11 @@ public void checkQuotas(long timeMs) {
 
     /**
      * Register a compound statistic with this sensor with no config override
+     * @param stat The stat to register
+     * @return true if stat is added to sensor, false if sensor is expired
      */
-    public void add(CompoundStat stat) {
-        add(stat, null);
+    public boolean add(CompoundStat stat) {
+        return add(stat, null);
     }
 
     /**
@@ -221,41 +280,71 @@ public void add(CompoundStat stat) {
      * @param stat The stat to register
      * @param config The configuration for this stat. If null then the stat will use the default configuration for this
      *        sensor.
+     * @return true if stat is added to sensor, false if sensor is expired
      */
-    public synchronized void add(CompoundStat stat, MetricConfig config) {
-        this.stats.add(Utils.notNull(stat));
-        Object lock = new Object();
+    public synchronized boolean add(CompoundStat stat, MetricConfig config) {
+        if (hasExpired())
+            return false;
+
+        final MetricConfig statConfig = config == null ? this.config : config;
+        stats.add(new StatAndConfig(Objects.requireNonNull(stat), () -> statConfig));
+        Object lock = metricLock();
         for (NamedMeasurable m : stat.stats()) {
-            KafkaMetric metric = new KafkaMetric(lock, m.name(), m.stat(), config == null ? this.config : config, time);
-            this.registry.registerMetric(metric);
-            this.metrics.add(metric);
+            final KafkaMetric metric = new KafkaMetric(lock, m.name(), m.stat(), statConfig, time);
+            if (!metrics.containsKey(metric.metricName())) {
+                registry.registerMetric(metric);
+                metrics.put(metric.metricName(), metric);
+            }
         }
+        return true;
     }
 
     /**
      * Register a metric with this sensor
      * @param metricName The name of the metric
      * @param stat The statistic to keep
+     * @return true if metric is added to sensor, false if sensor is expired
      */
-    public void add(MetricName metricName, MeasurableStat stat) {
-        add(metricName, stat, null);
+    public boolean add(MetricName metricName, MeasurableStat stat) {
+        return add(metricName, stat, null);
     }
 
     /**
      * Register a metric with this sensor
+     *
      * @param metricName The name of the metric
-     * @param stat The statistic to keep
-     * @param config A special configuration for this metric. If null use the sensor default configuration.
+     * @param stat       The statistic to keep
+     * @param config     A special configuration for this metric. If null use the sensor default configuration.
+     * @return true if metric is added to sensor, false if sensor is expired
      */
-    public synchronized void add(MetricName metricName, MeasurableStat stat, MetricConfig config) {
-        KafkaMetric metric = new KafkaMetric(new Object(),
-                                             Utils.notNull(metricName),
-                                             Utils.notNull(stat),
-                                             config == null ? this.config : config,
-                                             time);
-        this.registry.registerMetric(metric);
-        this.metrics.add(metric);
-        this.stats.add(stat);
+    public synchronized boolean add(final MetricName metricName, final MeasurableStat stat, final MetricConfig config) {
+        if (hasExpired()) {
+            return false;
+        } else if (metrics.containsKey(metricName)) {
+            return true;
+        } else {
+            final MetricConfig statConfig = config == null ? this.config : config;
+            final KafkaMetric metric = new KafkaMetric(
+                metricLock(),
+                Objects.requireNonNull(metricName),
+                Objects.requireNonNull(stat),
+                statConfig,
+                time
+            );
+            registry.registerMetric(metric);
+            metrics.put(metric.metricName(), metric);
+            stats.add(new StatAndConfig(Objects.requireNonNull(stat), metric::config));
+            return true;
+        }
+    }
+
+    /**
+     * Return if metrics were registered with this sensor.
+     *
+     * @return true if metrics were registered, false otherwise
+     */
+    public synchronized boolean hasMetrics() {
+        return !metrics.isEmpty();
     }
 
     /**
@@ -267,6 +356,30 @@ public boolean hasExpired() {
     }
 
     synchronized List metrics() {
-        return Collections.unmodifiableList(this.metrics);
+        return unmodifiableList(new ArrayList<>(this.metrics.values()));
+    }
+
+    /**
+     * KafkaMetrics of sensors which use SampledStat should be synchronized on the same lock
+     * for sensor record and metric value read to allow concurrent reads and updates. For simplicity,
+     * all sensors are synchronized on this object.
+     * 
        
+     * Sensor object is not used as a lock for reading metric value since metrics reporter is
+     * invoked while holding Sensor and Metrics locks to report addition and removal of metrics
+     * and synchronized reporters may deadlock if Sensor lock is used for reading metrics values.
+     * Note that Sensor object itself is used as a lock to protect the access to stats and metrics
+     * while recording metric values, adding and deleting sensors.
+     * 
        
+     * Locking order (assume all MetricsReporter methods may be synchronized):
+     * 
          
+     *   
        • Sensor#add: Sensor -> Metrics -> MetricsReporter
          • 
+     *   
        • Metrics#removeSensor: Sensor -> Metrics -> MetricsReporter
          • 
+     *   
        • KafkaMetric#metricValue: MetricsReporter -> Sensor#metricLock
          • 
+     *   
        • Sensor#record: Sensor -> Sensor#metricLock
          • 
+     * 
        
+     * 
        
+     */
+    private Object metricLock() {
+        return metricLock;
     }
 }
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/internals/IntGaugeSuite.java b/clients/src/main/java/org/apache/kafka/common/metrics/internals/IntGaugeSuite.java
new file mode 100644
index 0000000000000..cd52759a39069
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/metrics/internals/IntGaugeSuite.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.common.metrics.internals;
+
+import org.apache.kafka.common.MetricName;
+import org.apache.kafka.common.metrics.Gauge;
+import org.apache.kafka.common.metrics.MetricConfig;
+import org.apache.kafka.common.metrics.MetricValueProvider;
+import org.apache.kafka.common.metrics.Metrics;
+import org.slf4j.Logger;
+
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ConcurrentLinkedDeque;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
+import java.util.function.Function;
+
+/**
+ * Manages a suite of integer Gauges.
+ */
+public final class IntGaugeSuite implements AutoCloseable {
+    /**
+     * The log4j logger.
+     */
+    private final Logger log;
+
+    /**
+     * The name of this suite.
+     */
+    private final String suiteName;
+
+    /**
+     * The metrics object to use.
+     */
+    private final Metrics metrics;
+
+    /**
+     * A user-supplied callback which translates keys into unique metric names.
+     */
+    private final Function metricNameCalculator;
+
+    /**
+     * The maximum number of gauges that we will ever create at once.
+     */
+    private final int maxEntries;
+
+    /**
+     * A map from keys to gauges.  Protected by the object monitor.
+     */
+    private final Map gauges;
+
+    /**
+     * The keys of gauges that can be removed, since their value is zero.
+     * Protected by the object monitor.
+     */
+    private final Set removable;
+
+    /**
+     * A lockless list of pending metrics additions and removals.
+     */
+    private final ConcurrentLinkedDeque pending;
+
+    /**
+     * A lock which serializes modifications to metrics.  This lock is not
+     * required to create a new pending operation.
+     */
+    private final Lock modifyMetricsLock;
+
+    /**
+     * True if this suite is closed.  Protected by the object monitor.
+     */
+    private boolean closed;
+
+    /**
+     * A pending metrics addition or removal.
+     */
+    private static class PendingMetricsChange {
+        /**
+         * The name of the metric to add or remove.
+         */
+        private final MetricName metricName;
+
+        /**
+         * In an addition, this field is the MetricValueProvider to add.
+         * In a removal, this field is null.
+         */
+        private final MetricValueProvider provider;
+
+        PendingMetricsChange(MetricName metricName, MetricValueProvider provider) {
+            this.metricName = metricName;
+            this.provider = provider;
+        }
+    }
+
+    /**
+     * The gauge object which we register with the metrics system.
+     */
+    private static class StoredIntGauge implements Gauge {
+        private final MetricName metricName;
+        private int value;
+
+        StoredIntGauge(MetricName metricName) {
+            this.metricName = metricName;
+            this.value = 1;
+        }
+
+        /**
+         * This callback is invoked when the metrics system retrieves the value of this gauge.
+         */
+        @Override
+        public synchronized Integer value(MetricConfig config, long now) {
+            return value;
+        }
+
+        synchronized int increment() {
+            return ++value;
+        }
+
+        synchronized int decrement() {
+            return --value;
+        }
+
+        synchronized int value() {
+            return value;
+        }
+    }
+
+    public IntGaugeSuite(Logger log,
+                         String suiteName,
+                         Metrics metrics,
+                         Function metricNameCalculator,
+                         int maxEntries) {
+        this.log = log;
+        this.suiteName = suiteName;
+        this.metrics = metrics;
+        this.metricNameCalculator = metricNameCalculator;
+        this.maxEntries = maxEntries;
+        this.gauges = new HashMap<>(1);
+        this.removable = new HashSet<>();
+        this.pending = new ConcurrentLinkedDeque<>();
+        this.modifyMetricsLock = new ReentrantLock();
+        this.closed = false;
+        log.trace("{}: created new gauge suite with maxEntries = {}.",
+            suiteName, maxEntries);
+    }
+
+    public void increment(K key) {
+        synchronized (this) {
+            if (closed) {
+                log.warn("{}: Attempted to increment {}, but the GaugeSuite was closed.",
+                    suiteName, key.toString());
+                return;
+            }
+            StoredIntGauge gauge = gauges.get(key);
+            if (gauge != null) {
+                // Fast path: increment the existing counter.
+                if (gauge.increment() > 0) {
+                    removable.remove(key);
+                }
+                return;
+            }
+            if (gauges.size() == maxEntries) {
+                if (removable.isEmpty()) {
+                    log.debug("{}: Attempted to increment {}, but there are already {} entries.",
+                        suiteName, key.toString(), maxEntries);
+                    return;
+                }
+                Iterator iter = removable.iterator();
+                K keyToRemove = iter.next();
+                iter.remove();
+                MetricName metricNameToRemove = gauges.get(keyToRemove).metricName;
+                gauges.remove(keyToRemove);
+                pending.push(new PendingMetricsChange(metricNameToRemove, null));
+                log.trace("{}: Removing the metric {}, which has a value of 0.",
+                    suiteName, keyToRemove.toString());
+            }
+            MetricName metricNameToAdd = metricNameCalculator.apply(key);
+            gauge = new StoredIntGauge(metricNameToAdd);
+            gauges.put(key, gauge);
+            pending.push(new PendingMetricsChange(metricNameToAdd, gauge));
+            log.trace("{}: Adding a new metric {}.", suiteName, key.toString());
+        }
+        // Drop the object monitor and perform any pending metrics additions or removals.
+        performPendingMetricsOperations();
+    }
+
+    /**
+     * Perform pending metrics additions or removals.
+     * It is important to perform them in order.  For example, we don't want to try
+     * to remove a metric that we haven't finished adding yet.
+     */
+    private void performPendingMetricsOperations() {
+        modifyMetricsLock.lock();
+        try {
+            log.trace("{}: entering performPendingMetricsOperations", suiteName);
+            for (PendingMetricsChange change = pending.pollLast();
+                 change != null;
+                 change = pending.pollLast()) {
+                if (change.provider == null) {
+                    if (log.isTraceEnabled()) {
+                        log.trace("{}: removing metric {}", suiteName, change.metricName);
+                    }
+                    metrics.removeMetric(change.metricName);
+                } else {
+                    if (log.isTraceEnabled()) {
+                        log.trace("{}: adding metric {}", suiteName, change.metricName);
+                    }
+                    metrics.addMetric(change.metricName, change.provider);
+                }
+            }
+            log.trace("{}: leaving performPendingMetricsOperations", suiteName);
+        } finally {
+            modifyMetricsLock.unlock();
+        }
+    }
+
+    public synchronized void decrement(K key) {
+        if (closed) {
+            log.warn("{}: Attempted to decrement {}, but the gauge suite was closed.",
+                suiteName, key.toString());
+            return;
+        }
+        StoredIntGauge gauge = gauges.get(key);
+        if (gauge == null) {
+            log.debug("{}: Attempted to decrement {}, but no such metric was registered.",
+                suiteName, key.toString());
+        } else {
+            int cur = gauge.decrement();
+            log.trace("{}: Removed a reference to {}.  {} reference(s) remaining.",
+                suiteName, key.toString(), cur);
+            if (cur <= 0) {
+                removable.add(key);
+            }
+        }
+    }
+
+    @Override
+    public synchronized void close() {
+        if (closed) {
+            log.trace("{}: gauge suite is already closed.", suiteName);
+            return;
+        }
+        closed = true;
+        int prevSize = 0;
+        for (Iterator iter = gauges.values().iterator(); iter.hasNext(); ) {
+            pending.push(new PendingMetricsChange(iter.next().metricName, null));
+            prevSize++;
+            iter.remove();
+        }
+        performPendingMetricsOperations();
+        log.trace("{}: closed {} metric(s).", suiteName, prevSize);
+    }
+
+    /**
+     * Get the maximum number of metrics this suite can create.
+     */
+    public int maxEntries() {
+        return maxEntries;
+    }
+
+    // Visible for testing only.
+    Metrics metrics() {
+        return metrics;
+    }
+
+    /**
+     * Return a map from keys to current reference counts.
+     * Visible for testing only.
+     */
+    synchronized Map values() {
+        HashMap values = new HashMap<>();
+        for (Map.Entry entry : gauges.entrySet()) {
+            values.put(entry.getKey(), entry.getValue().value());
+        }
+        return values;
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/metrics/internals/MetricsUtils.java b/clients/src/main/java/org/apache/kafka/common/metrics/internals/MetricsUtils.java
new file mode 100644
index 0000000000000..edf061f282c50
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/metrics/internals/MetricsUtils.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.metrics.internals;
+
+import org.apache.kafka.common.metrics.Metrics;
+
+import java.util.LinkedHashMap;
+import java.util.Map;
+import java.util.concurrent.TimeUnit;
+
+public class MetricsUtils {
+    /**
+     * Converts the provided time from milliseconds to the requested
+     * time unit.
+     */
+    public static double convert(long timeMs, TimeUnit unit) {
+        switch (unit) {
+            case NANOSECONDS:
+                return timeMs * 1000.0 * 1000.0;
+            case MICROSECONDS:
+                return timeMs * 1000.0;
+            case MILLISECONDS:
+                return timeMs;
+            case SECONDS:
+                return timeMs / 1000.0;
+            case MINUTES:
+                return timeMs / (60.0 * 1000.0);
+            case HOURS:
+                return timeMs / (60.0 * 60.0 * 1000.0);
+            case DAYS:
+                return timeMs / (24.0 * 60.0 * 60.0 * 1000.0);
+            default:
+                throw new IllegalStateException("Unknown unit: " + unit);
+        }
+    }
+
+    /**
+     * Create a set of tags using the supplied key and value pairs. The order of the tags will be kept.
+     *
+     * @param keyValue the key and value pairs for the tags; must be an even number
+     * @return the map of tags that can be supplied to the {@link Metrics} methods; never null
+     */
+    public static Map getTags(String... keyValue) {
+        if ((keyValue.length % 2) != 0)
+            throw new IllegalArgumentException("keyValue needs to be specified in pairs");
+        Map tags = new LinkedHashMap<>(keyValue.length / 2);
+
+        for (int i = 0; i < keyValue.length; i += 2)
+            tags.put(keyValue[i], keyValue[i + 1]);
+        return tags;
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/metrics/stats/Avg.java b/clients/src/main/java/org/apache/kafka/common/metrics/stats/Avg.java
index a09ca5af3e2ca..4e6c3372575fc 100644
--- a/clients/src/main/java/org/apache/kafka/common/metrics/stats/Avg.java
+++ b/clients/src/main/java/org/apache/kafka/common/metrics/stats/Avg.java
@@ -42,7 +42,8 @@ public double combine(List samples, MetricConfig config, long now) {
             total += s.value;
             count += s.eventCount;
         }
-        return count == 0 ? 0 : total / count;
+        return count == 0 ? Double.NaN : total / count;
     }
 
 }
+
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/clients/src/main/java/org/apache/kafka/common/metrics/stats/CumulativeCount.java b/clients/src/main/java/org/apache/kafka/common/metrics/stats/CumulativeCount.java
new file mode 100644
index 0000000000000..85591b5cf726f
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/metrics/stats/CumulativeCount.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.metrics.stats;
+
+import org.apache.kafka.common.metrics.MetricConfig;
+
+/**
+ * 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 extends CumulativeSum {
+    @Override
+    public void record(final MetricConfig config, final double value, final long timeMs) {
+        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/Histogram.java b/clients/src/main/java/org/apache/kafka/common/metrics/stats/Histogram.java
index af2b06407089c..97f91822d8261 100644
--- a/clients/src/main/java/org/apache/kafka/common/metrics/stats/Histogram.java
+++ b/clients/src/main/java/org/apache/kafka/common/metrics/stats/Histogram.java
@@ -16,6 +16,8 @@
  */
 package org.apache.kafka.common.metrics.stats;
 
+import java.util.Arrays;
+
 public class Histogram {
 
     private final BinScheme binScheme;
@@ -55,8 +57,7 @@ public float[] counts() {
     }
 
     public void clear() {
-        for (int i = 0; i < this.hist.length; i++)
-            this.hist[i] = 0.0f;
+        Arrays.fill(this.hist, 0.0f);
         this.count = 0;
     }
 
@@ -156,10 +157,7 @@ public int toBin(double x) {
             if (binNumber < MIN_BIN_NUMBER) {
                 return MIN_BIN_NUMBER;
             }
-            if (binNumber > maxBinNumber) {
-                return maxBinNumber;
-            }
-            return binNumber;
+            return Math.min(binNumber, maxBinNumber);
         }
     }
 
diff --git a/clients/src/main/java/org/apache/kafka/common/metrics/stats/Max.java b/clients/src/main/java/org/apache/kafka/common/metrics/stats/Max.java
index 6d75454a9ccb1..d91bf4099c174 100644
--- a/clients/src/main/java/org/apache/kafka/common/metrics/stats/Max.java
+++ b/clients/src/main/java/org/apache/kafka/common/metrics/stats/Max.java
@@ -37,9 +37,12 @@ protected void update(Sample sample, MetricConfig config, double value, long now
     @Override
     public double combine(List samples, MetricConfig config, long now) {
         double max = Double.NEGATIVE_INFINITY;
-        for (Sample sample : samples)
+        long count = 0;
+        for (Sample sample : samples) {
             max = Math.max(max, sample.value);
-        return max;
+            count += sample.eventCount;
+        }
+        return count == 0 ? Double.NaN : max;
     }
 
 }
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 09263cecae89c..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,45 +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) {
-        this.total = new Total();
+        if (!(rateStat instanceof WindowedSum)) {
+            throw new IllegalArgumentException("Meter is supported only for WindowedCount or WindowedSum.");
+        }
+        this.total = new CumulativeSum();
         this.rate = new Rate(unit, rateStat);
         this.rateMetricName = rateMetricName;
         this.totalMetricName = totalMetricName;
@@ -77,6 +78,8 @@ public List stats() {
     @Override
     public void record(MetricConfig config, double value, long timeMs) {
         rate.record(config, value, timeMs);
-        total.record(config, value, timeMs);
+        // Total metrics with Count stat should record 1.0 (as recorded in the count)
+        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/Min.java b/clients/src/main/java/org/apache/kafka/common/metrics/stats/Min.java
index 7a18a2d0ea44d..3b9925a0d1b45 100644
--- a/clients/src/main/java/org/apache/kafka/common/metrics/stats/Min.java
+++ b/clients/src/main/java/org/apache/kafka/common/metrics/stats/Min.java
@@ -37,9 +37,12 @@ protected void update(Sample sample, MetricConfig config, double value, long now
     @Override
     public double combine(List samples, MetricConfig config, long now) {
         double min = Double.MAX_VALUE;
-        for (Sample sample : samples)
+        long count = 0;
+        for (Sample sample : samples) {
             min = Math.min(min, sample.value);
-        return min;
+            count += sample.eventCount;
+        }
+        return count == 0 ? Double.NaN : min;
     }
 
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/metrics/stats/Percentiles.java b/clients/src/main/java/org/apache/kafka/common/metrics/stats/Percentiles.java
index e970f48242549..4cdc2ce9ce12d 100644
--- a/clients/src/main/java/org/apache/kafka/common/metrics/stats/Percentiles.java
+++ b/clients/src/main/java/org/apache/kafka/common/metrics/stats/Percentiles.java
@@ -20,17 +20,20 @@
 import java.util.List;
 
 import org.apache.kafka.common.metrics.CompoundStat;
-import org.apache.kafka.common.metrics.Measurable;
 import org.apache.kafka.common.metrics.MetricConfig;
 import org.apache.kafka.common.metrics.stats.Histogram.BinScheme;
 import org.apache.kafka.common.metrics.stats.Histogram.ConstantBinScheme;
 import org.apache.kafka.common.metrics.stats.Histogram.LinearBinScheme;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /**
  * A compound stat that reports one or more percentiles
  */
 public class Percentiles extends SampledStat implements CompoundStat {
 
+    private final Logger log = LoggerFactory.getLogger(Percentiles.class);
+
     public enum BucketSizing {
         CONSTANT, LINEAR
     }
@@ -38,6 +41,8 @@ public enum BucketSizing {
     private final int buckets;
     private final Percentile[] percentiles;
     private final BinScheme binScheme;
+    private final double min;
+    private final double max;
 
     public Percentiles(int sizeInBytes, double max, BucketSizing bucketing, Percentile... percentiles) {
         this(sizeInBytes, 0.0, max, bucketing, percentiles);
@@ -47,6 +52,8 @@ public Percentiles(int sizeInBytes, double min, double max, BucketSizing bucketi
         super(0.0);
         this.percentiles = percentiles;
         this.buckets = sizeInBytes / 4;
+        this.min = min;
+        this.max = max;
         if (bucketing == BucketSizing.CONSTANT) {
             this.binScheme = new ConstantBinScheme(buckets, min, max);
         } else if (bucketing == BucketSizing.LINEAR) {
@@ -60,14 +67,13 @@ public Percentiles(int sizeInBytes, double min, double max, BucketSizing bucketi
 
     @Override
     public List stats() {
-        List ms = new ArrayList(this.percentiles.length);
+        List ms = new ArrayList<>(this.percentiles.length);
         for (Percentile percentile : this.percentiles) {
             final double pct = percentile.percentile();
-            ms.add(new NamedMeasurable(percentile.name(), new Measurable() {
-                public double measure(MetricConfig config, long now) {
-                    return value(config, now, pct / 100.0);
-                }
-            }));
+            ms.add(new NamedMeasurable(
+                percentile.name(),
+                (config, now) -> value(config, now, pct / 100.0))
+            );
         }
         return ms;
     }
@@ -105,8 +111,21 @@ protected HistogramSample newSample(long timeMs) {
 
     @Override
     protected void update(Sample sample, MetricConfig config, double value, long timeMs) {
+        final double boundedValue;
+        if (value > max) {
+            log.debug("Received value {} which is greater than max recordable value {}, will be pinned to the max value",
+                     value, max);
+            boundedValue = max;
+        } else if (value < min) {
+            log.debug("Received value {} which is less than min recordable value {}, will be pinned to the min value",
+                     value, min);
+            boundedValue = min;
+        } else {
+            boundedValue = value;
+        }
+
         HistogramSample hist = (HistogramSample) sample;
-        hist.histogram.record(value);
+        hist.histogram.record(boundedValue);
     }
 
     private static class HistogramSample extends SampledStat.Sample {
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..0f3573e4aa5fc 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,13 +16,13 @@
  */
 package org.apache.kafka.common.metrics.stats;
 
-import java.util.List;
 import java.util.Locale;
 import java.util.concurrent.TimeUnit;
 
 import org.apache.kafka.common.metrics.MeasurableStat;
 import org.apache.kafka.common.metrics.MetricConfig;
 
+import static org.apache.kafka.common.metrics.internals.MetricsUtils.convert;
 
 /**
  * The rate of the given quantity. By default this is the total observed over a set of samples from a sampled statistic
@@ -40,7 +40,7 @@ public Rate() {
     }
 
     public Rate(TimeUnit unit) {
-        this(unit, new SampledTotal());
+        this(unit, new WindowedSum());
     }
 
     public Rate(SampledStat stat) {
@@ -64,7 +64,7 @@ public void record(MetricConfig config, double value, long timeMs) {
     @Override
     public double measure(MetricConfig config, long now) {
         double value = stat.measure(config, now);
-        return value / convert(windowSize(config, now));
+        return value / convert(windowSize(config, now), unit);
     }
 
     public long windowSize(MetricConfig config, long now) {
@@ -94,45 +94,10 @@ public long windowSize(MetricConfig config, long now) {
         return totalElapsedTimeMs;
     }
 
-    private double convert(long timeMs) {
-        switch (unit) {
-            case NANOSECONDS:
-                return timeMs * 1000.0 * 1000.0;
-            case MICROSECONDS:
-                return timeMs * 1000.0;
-            case MILLISECONDS:
-                return timeMs;
-            case SECONDS:
-                return timeMs / 1000.0;
-            case MINUTES:
-                return timeMs / (60.0 * 1000.0);
-            case HOURS:
-                return timeMs / (60.0 * 60.0 * 1000.0);
-            case DAYS:
-                return timeMs / (24.0 * 60.0 * 60.0 * 1000.0);
-            default:
-                throw new IllegalStateException("Unknown unit: " + unit);
-        }
-    }
-
-    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/SampledStat.java b/clients/src/main/java/org/apache/kafka/common/metrics/stats/SampledStat.java
index 18dd6f271b076..369709e32e11c 100644
--- a/clients/src/main/java/org/apache/kafka/common/metrics/stats/SampledStat.java
+++ b/clients/src/main/java/org/apache/kafka/common/metrics/stats/SampledStat.java
@@ -40,7 +40,7 @@ public abstract class SampledStat implements MeasurableStat {
 
     public SampledStat(double initialValue) {
         this.initialValue = initialValue;
-        this.samples = new ArrayList(2);
+        this.samples = new ArrayList<>(2);
     }
 
     @Override
diff --git a/clients/src/main/java/org/apache/kafka/common/metrics/stats/SimpleRate.java b/clients/src/main/java/org/apache/kafka/common/metrics/stats/SimpleRate.java
index 69d44a8b2de22..931bd9c35e51f 100644
--- a/clients/src/main/java/org/apache/kafka/common/metrics/stats/SimpleRate.java
+++ b/clients/src/main/java/org/apache/kafka/common/metrics/stats/SimpleRate.java
@@ -34,6 +34,6 @@ public class SimpleRate extends Rate {
     public long windowSize(MetricConfig config, long now) {
         stat.purgeObsoleteSamples(config, now);
         long elapsed = now - stat.oldest(now).lastWindowMs;
-        return elapsed < config.timeWindowMs() ? config.timeWindowMs() : elapsed;
+        return Math.max(elapsed, config.timeWindowMs());
     }
 }
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/TokenBucket.java b/clients/src/main/java/org/apache/kafka/common/metrics/stats/TokenBucket.java
new file mode 100644
index 0000000000000..8416e9715d639
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/metrics/stats/TokenBucket.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.common.metrics.stats;
+
+import java.util.concurrent.TimeUnit;
+import org.apache.kafka.common.metrics.MeasurableStat;
+import org.apache.kafka.common.metrics.MetricConfig;
+import org.apache.kafka.common.metrics.Quota;
+
+import static org.apache.kafka.common.metrics.internals.MetricsUtils.convert;
+
+/**
+ * The {@link TokenBucket} is a {@link MeasurableStat} implementing a token bucket algorithm
+ * that is usable within a {@link org.apache.kafka.common.metrics.Sensor}.
+ *
+ * The {@link Quota#bound()} defined the refill rate of the bucket while the maximum burst or
+ * the maximum number of credits of the bucket is defined by
+ * {@link MetricConfig#samples() * MetricConfig#timeWindowMs() * Quota#bound()}.
+ *
+ * The quota is considered as exhausted when the amount of remaining credits in the bucket
+ * is below zero. The enforcement is done by the {@link org.apache.kafka.common.metrics.Sensor}.
+ *
+ * Token Bucket vs Rate based Quota:
+ * The current sampled rate based quota does not cope well with bursty workloads. The issue is
+ * that a unique and large sample can hold the average above the quota until it is discarded.
+ * Practically, when this happens, one must wait until the sample is expired to bring the rate
+ * below the quota even though less time would be theoretically required. As an examples, let's
+ * imagine that we have:
+ * - Quota (Q)   = 5
+ * - Samples (S) = 100
+ * - Window (W)  = 1s
+ * A burst of 560 brings the average rate (R) to 5.6 (560 / 100). The expected throttle time is
+ * computed as follow: ((R - Q / Q * S * W)) = ((5.6 - 5) / 5 * 100 * 1) = 12 secs. In practice,
+ * the average rate won't go below the quota before the burst is dropped from the samples so one
+ * must wait 100s (S * W).
+ *
+ * The token bucket relies on continuously updated amount of credits. Therefore, it does not
+ * suffers from the above issue. The same example would work as follow:
+ * - Quota (Q) = 5
+ * - Burst (B) = 5 * 1 * 100 = 500 (Q * S * W)
+ * A burst of 560 brings the amount of credits to -60. One must wait 12s (-(-60)/5) to refill the
+ * bucket to zero.
+ */
+public class TokenBucket implements MeasurableStat {
+    private final TimeUnit unit;
+    private double tokens;
+    private long lastUpdateMs;
+
+    public TokenBucket() {
+        this(TimeUnit.SECONDS);
+    }
+
+    public TokenBucket(TimeUnit unit) {
+        this.unit = unit;
+        this.tokens = 0;
+        this.lastUpdateMs = 0;
+    }
+
+    @Override
+    public double measure(final MetricConfig config, final long timeMs) {
+        if (config.quota() == null)
+            return Long.MAX_VALUE;
+        final double quota = config.quota().bound();
+        final double burst = burst(config);
+        refill(quota, burst, timeMs);
+        return this.tokens;
+    }
+
+    @Override
+    public void record(final MetricConfig config, final double value, final long timeMs) {
+        if (config.quota() == null)
+            return;
+        final double quota = config.quota().bound();
+        final double burst = burst(config);
+        refill(quota, burst, timeMs);
+        this.tokens = Math.min(burst, this.tokens - value);
+    }
+
+    private void refill(final double quota, final double burst, final long timeMs) {
+        this.tokens = Math.min(burst, this.tokens + quota * convert(timeMs - lastUpdateMs, unit));
+        this.lastUpdateMs = timeMs;
+    }
+
+    private double burst(final MetricConfig config) {
+        return config.samples() * convert(config.timeWindowMs(), unit) * config.quota().bound();
+    }
+}
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/Authenticator.java b/clients/src/main/java/org/apache/kafka/common/network/Authenticator.java
index 4e2e7273a6815..873c1f90df798 100644
--- a/clients/src/main/java/org/apache/kafka/common/network/Authenticator.java
+++ b/clients/src/main/java/org/apache/kafka/common/network/Authenticator.java
@@ -18,9 +18,11 @@
 
 import org.apache.kafka.common.errors.AuthenticationException;
 import org.apache.kafka.common.security.auth.KafkaPrincipal;
+import org.apache.kafka.common.security.auth.KafkaPrincipalSerde;
 
 import java.io.Closeable;
 import java.io.IOException;
+import java.util.Optional;
 
 /**
  * Authentication for Channel
@@ -37,14 +39,128 @@ public interface Authenticator extends Closeable {
      */
     void authenticate() throws AuthenticationException, IOException;
 
+    /**
+     * Perform any processing related to authentication failure. This is invoked when the channel is about to be closed
+     * because of an {@link AuthenticationException} thrown from a prior {@link #authenticate()} call.
+     * @throws IOException if read/write fails due to an I/O error
+     */
+    default void handleAuthenticationFailure() throws IOException {
+    }
+
     /**
      * Returns Principal using PrincipalBuilder
      */
     KafkaPrincipal principal();
 
+    /**
+     * Returns the serializer/deserializer interface for principal
+     */
+    Optional principalSerde();
+
     /**
      * returns true if authentication is complete otherwise returns false;
      */
     boolean complete();
 
+    /**
+     * Begins re-authentication. Uses transportLayer to read or write tokens as is
+     * done for {@link #authenticate()}. For security protocols PLAINTEXT and SSL,
+     * this is a no-op since re-authentication does not apply/is not supported,
+     * respectively. For SASL_PLAINTEXT and SASL_SSL, this performs a SASL
+     * authentication. Any in-flight responses from prior requests can/will be read
+     * and collected for later processing as required. There must not be partially
+     * written requests; any request queued for writing (for which zero bytes have
+     * been written) remains queued until after re-authentication succeeds.
+     * 
+     * @param reauthenticationContext
+     *            the context in which this re-authentication is occurring. This
+     *            instance is responsible for closing the previous Authenticator
+     *            returned by
+     *            {@link ReauthenticationContext#previousAuthenticator()}.
+     * @throws AuthenticationException
+     *             if authentication fails due to invalid credentials or other
+     *             security configuration errors
+     * @throws IOException
+     *             if read/write fails due to an I/O error
+     */
+    default void reauthenticate(ReauthenticationContext reauthenticationContext) throws IOException {
+        // empty
+    }
+
+    /**
+     * Return the session expiration time, if any, otherwise null. The value is in
+     * nanoseconds as per {@code System.nanoTime()} and is therefore only useful
+     * when compared to such a value -- it's absolute value is meaningless. This
+     * value may be non-null only on the server-side. It represents the time after
+     * which, in the absence of re-authentication, the broker will close the session
+     * if it receives a request unrelated to authentication. We store nanoseconds
+     * here to avoid having to invoke the more expensive {@code milliseconds()} call
+     * on the broker for every request
+     * 
+     * @return the session expiration time, if any, otherwise null
+     */
+    default Long serverSessionExpirationTimeNanos() {
+        return null;
+    }
+
+    /**
+     * Return the time on or after which a client should re-authenticate this
+     * session, if any, otherwise null. The value is in nanoseconds as per
+     * {@code System.nanoTime()} and is therefore only useful when compared to such
+     * a value -- it's absolute value is meaningless. This value may be non-null
+     * only on the client-side. It will be a random time between 85% and 95% of the
+     * full session lifetime to account for latency between client and server and to
+     * avoid re-authentication storms that could be caused by many sessions
+     * re-authenticating simultaneously.
+     * 
+     * @return the time on or after which a client should re-authenticate this
+     *         session, if any, otherwise null
+     */
+    default Long clientSessionReauthenticationTimeNanos() {
+        return null;
+    }
+
+    /**
+     * Return the number of milliseconds that elapsed while re-authenticating this
+     * session from the perspective of this instance, if applicable, otherwise null.
+     * The server-side perspective will yield a lower value than the client-side
+     * perspective of the same re-authentication because the client-side observes an
+     * additional network round-trip.
+     * 
+     * @return the number of milliseconds that elapsed while re-authenticating this
+     *         session from the perspective of this instance, if applicable,
+     *         otherwise null
+     */
+    default Long reauthenticationLatencyMs() {
+        return null;
+    }
+
+    /**
+     * Return the next (always non-null but possibly empty) client-side
+     * {@link NetworkReceive} response that arrived during re-authentication that
+     * is unrelated to re-authentication, if any. These correspond to requests sent
+     * prior to the beginning of re-authentication; the requests were made when the
+     * channel was successfully authenticated, and the responses arrived during the
+     * re-authentication process. The response returned is removed from the authenticator's
+     * queue. Responses of requests sent after completion of re-authentication are
+     * processed only when the authenticator response queue is empty.
+     * 
+     * @return the (always non-null but possibly empty) client-side
+     *         {@link NetworkReceive} response that arrived during
+     *         re-authentication that is unrelated to re-authentication, if any
+     */
+    default Optional pollResponseReceivedDuringReauthentication() {
+        return Optional.empty();
+    }
+    
+    /**
+     * Return true if this is a server-side authenticator and the connected client
+     * has indicated that it supports re-authentication, otherwise false
+     * 
+     * @return true if this is a server-side authenticator and the connected client
+     *         has indicated that it supports re-authentication, otherwise false
+     */
+    default boolean connectedClientSupportsReauthentication() {
+        return false;
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/network/ByteBufferSend.java b/clients/src/main/java/org/apache/kafka/common/network/ByteBufferSend.java
index 23a169f4bc4d9..1402f522cea52 100644
--- a/clients/src/main/java/org/apache/kafka/common/network/ByteBufferSend.java
+++ b/clients/src/main/java/org/apache/kafka/common/network/ByteBufferSend.java
@@ -26,23 +26,22 @@
  */
 public class ByteBufferSend implements Send {
 
-    private final String destination;
-    private final int size;
+    private final long size;
     protected final ByteBuffer[] buffers;
-    private int remaining;
+    private long remaining;
     private boolean pending = false;
 
-    public ByteBufferSend(String destination, ByteBuffer... buffers) {
-        this.destination = destination;
+    public ByteBufferSend(ByteBuffer... buffers) {
         this.buffers = buffers;
         for (ByteBuffer buffer : buffers)
             remaining += buffer.remaining();
         this.size = remaining;
     }
 
-    @Override
-    public String destination() {
-        return destination;
+    public ByteBufferSend(ByteBuffer[] buffers, long size) {
+        this.buffers = buffers;
+        this.size = size;
+        this.remaining = size;
     }
 
     @Override
@@ -64,4 +63,23 @@ public long writeTo(GatheringByteChannel channel) throws IOException {
         pending = TransportLayers.hasPendingWrites(channel);
         return written;
     }
+
+    public long remaining() {
+        return remaining;
+    }
+
+    @Override
+    public String toString() {
+        return "ByteBufferSend(" +
+            ", size=" + size +
+            ", remaining=" + remaining +
+            ", pending=" + pending +
+            ')';
+    }
+
+    public static ByteBufferSend sizePrefixed(ByteBuffer buffer) {
+        ByteBuffer sizeBuffer = ByteBuffer.allocate(4);
+        sizeBuffer.putInt(0, buffer.remaining());
+        return new ByteBufferSend(sizeBuffer, buffer);
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/network/ChannelBuilder.java b/clients/src/main/java/org/apache/kafka/common/network/ChannelBuilder.java
index 54689f37ed858..0cf1d74f85ce6 100644
--- a/clients/src/main/java/org/apache/kafka/common/network/ChannelBuilder.java
+++ b/clients/src/main/java/org/apache/kafka/common/network/ChannelBuilder.java
@@ -16,22 +16,17 @@
  */
 package org.apache.kafka.common.network;
 
-import java.util.Map;
 import java.nio.channels.SelectionKey;
 
 import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.Configurable;
 import org.apache.kafka.common.memory.MemoryPool;
 
 
 /**
  * A ChannelBuilder interface to build Channel based on configs
  */
-public interface ChannelBuilder extends AutoCloseable {
-
-    /**
-     * Configure this class with the given key-value pairs
-     */
-    void configure(Map configs) throws KafkaException;
+public interface ChannelBuilder extends AutoCloseable, Configurable {
 
     /**
      * returns a Channel with TransportLayer and Authenticator configured.
@@ -41,7 +36,8 @@ public interface ChannelBuilder extends AutoCloseable {
      * @param  memoryPool memory pool from which to allocate buffers, or null for none
      * @return KafkaChannel
      */
-    KafkaChannel buildChannel(String id, SelectionKey key, int maxReceiveSize, MemoryPool memoryPool) throws KafkaException;
+    KafkaChannel buildChannel(String id, SelectionKey key, int maxReceiveSize,
+                              MemoryPool memoryPool, ChannelMetadataRegistry metadataRegistry) throws KafkaException;
 
     /**
      * Closes ChannelBuilder
diff --git a/clients/src/main/java/org/apache/kafka/common/network/ChannelBuilders.java b/clients/src/main/java/org/apache/kafka/common/network/ChannelBuilders.java
index 74bd0a0c6f3b8..fce004417f815 100644
--- a/clients/src/main/java/org/apache/kafka/common/network/ChannelBuilders.java
+++ b/clients/src/main/java/org/apache/kafka/common/network/ChannelBuilders.java
@@ -26,12 +26,18 @@
 import org.apache.kafka.common.security.auth.KafkaPrincipalBuilder;
 import org.apache.kafka.common.security.authenticator.CredentialCache;
 import org.apache.kafka.common.security.kerberos.KerberosShortNamer;
+import org.apache.kafka.common.security.ssl.SslPrincipalMapper;
+import org.apache.kafka.common.security.token.delegation.internals.DelegationTokenCache;
+import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.common.utils.Utils;
 
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 
 public class ChannelBuilders {
-
     private ChannelBuilders() { }
 
     /**
@@ -40,17 +46,23 @@ private ChannelBuilders() { }
      * @param config client config
      * @param listenerName the listenerName if contextType is SERVER or null otherwise
      * @param clientSaslMechanism SASL mechanism if mode is CLIENT, ignored otherwise
+     * @param time the time instance
      * @param saslHandshakeRequestEnable flag to enable Sasl handshake requests; disabled only for SASL
      *             inter-broker connections with inter-broker protocol version < 0.10
+     * @param logContext the log context instance
+     *
      * @return the configured `ChannelBuilder`
      * @throws IllegalArgumentException if `mode` invariants described above is not maintained
      */
-    public static ChannelBuilder clientChannelBuilder(SecurityProtocol securityProtocol,
+    public static ChannelBuilder clientChannelBuilder(
+            SecurityProtocol securityProtocol,
             JaasContext.Type contextType,
             AbstractConfig config,
             ListenerName listenerName,
             String clientSaslMechanism,
-            boolean saslHandshakeRequestEnable) {
+            Time time,
+            boolean saslHandshakeRequestEnable,
+            LogContext logContext) {
 
         if (securityProtocol == SecurityProtocol.SASL_PLAINTEXT || securityProtocol == SecurityProtocol.SASL_SSL) {
             if (contextType == null)
@@ -58,23 +70,33 @@ public static ChannelBuilder clientChannelBuilder(SecurityProtocol securityProto
             if (clientSaslMechanism == null)
                 throw new IllegalArgumentException("`clientSaslMechanism` must be non-null in client mode if `securityProtocol` is `" + securityProtocol + "`");
         }
-        return create(securityProtocol, Mode.CLIENT, contextType, config, listenerName, clientSaslMechanism,
-                saslHandshakeRequestEnable, null);
+        return create(securityProtocol, Mode.CLIENT, contextType, config, listenerName, false, clientSaslMechanism,
+                saslHandshakeRequestEnable, null, null, time, logContext);
     }
 
     /**
      * @param listenerName the listenerName
+     * @param isInterBrokerListener whether or not this listener is used for inter-broker requests
      * @param securityProtocol the securityProtocol
      * @param config server config
      * @param credentialCache Credential cache for SASL/SCRAM if SCRAM is enabled
+     * @param tokenCache Delegation token cache
+     * @param time the time instance
+     * @param logContext the log context instance
+     *
      * @return the configured `ChannelBuilder`
      */
     public static ChannelBuilder serverChannelBuilder(ListenerName listenerName,
+                                                      boolean isInterBrokerListener,
                                                       SecurityProtocol securityProtocol,
                                                       AbstractConfig config,
-                                                      CredentialCache credentialCache) {
-        return create(securityProtocol, Mode.SERVER, JaasContext.Type.SERVER, config, listenerName, null,
-                true, credentialCache);
+                                                      CredentialCache credentialCache,
+                                                      DelegationTokenCache tokenCache,
+                                                      Time time,
+                                                      LogContext logContext) {
+        return create(securityProtocol, Mode.SERVER, JaasContext.Type.SERVER, config, listenerName,
+                isInterBrokerListener, null, true, credentialCache,
+                tokenCache, time, logContext);
     }
 
     private static ChannelBuilder create(SecurityProtocol securityProtocol,
@@ -82,30 +104,51 @@ private static ChannelBuilder create(SecurityProtocol securityProtocol,
                                          JaasContext.Type contextType,
                                          AbstractConfig config,
                                          ListenerName listenerName,
+                                         boolean isInterBrokerListener,
                                          String clientSaslMechanism,
                                          boolean saslHandshakeRequestEnable,
-                                         CredentialCache credentialCache) {
-        Map configs;
-        if (listenerName == null)
-            configs = config.values();
-        else
-            configs = config.valuesWithPrefixOverride(listenerName.configPrefix());
+                                         CredentialCache credentialCache,
+                                         DelegationTokenCache tokenCache,
+                                         Time time,
+                                         LogContext logContext) {
+        Map configs = channelBuilderConfigs(config, listenerName);
 
         ChannelBuilder channelBuilder;
         switch (securityProtocol) {
             case SSL:
                 requireNonNullMode(mode, securityProtocol);
-                channelBuilder = new SslChannelBuilder(mode);
+                channelBuilder = new SslChannelBuilder(mode, listenerName, isInterBrokerListener, logContext);
                 break;
             case SASL_SSL:
             case SASL_PLAINTEXT:
                 requireNonNullMode(mode, securityProtocol);
-                JaasContext jaasContext = JaasContext.load(contextType, listenerName, configs);
-                channelBuilder = new SaslChannelBuilder(mode, jaasContext, securityProtocol, listenerName,
-                        clientSaslMechanism, saslHandshakeRequestEnable, credentialCache);
+                Map jaasContexts;
+                if (mode == Mode.SERVER) {
+                    @SuppressWarnings("unchecked")
+                    List enabledMechanisms = (List) configs.get(BrokerSecurityConfigs.SASL_ENABLED_MECHANISMS_CONFIG);
+                    jaasContexts = new HashMap<>(enabledMechanisms.size());
+                    for (String mechanism : enabledMechanisms)
+                        jaasContexts.put(mechanism, JaasContext.loadServerContext(listenerName, mechanism, configs));
+                } else {
+                    // Use server context for inter-broker client connections and client context for other clients
+                    JaasContext jaasContext = contextType == JaasContext.Type.CLIENT ? JaasContext.loadClientContext(configs) :
+                            JaasContext.loadServerContext(listenerName, clientSaslMechanism, configs);
+                    jaasContexts = Collections.singletonMap(clientSaslMechanism, jaasContext);
+                }
+                channelBuilder = new SaslChannelBuilder(mode,
+                        jaasContexts,
+                        securityProtocol,
+                        listenerName,
+                        isInterBrokerListener,
+                        clientSaslMechanism,
+                        saslHandshakeRequestEnable,
+                        credentialCache,
+                        tokenCache,
+                        time,
+                        logContext);
                 break;
             case PLAINTEXT:
-                channelBuilder = new PlaintextChannelBuilder();
+                channelBuilder = new PlaintextChannelBuilder(listenerName);
                 break;
             default:
                 throw new IllegalArgumentException("Unexpected securityProtocol " + securityProtocol);
@@ -115,6 +158,28 @@ private static ChannelBuilder create(SecurityProtocol securityProtocol,
         return channelBuilder;
     }
 
+    /**
+     * @return a mutable RecordingMap. The elements got from RecordingMap are marked as "used".
+     */
+    @SuppressWarnings("unchecked")
+    static Map channelBuilderConfigs(final AbstractConfig config, final ListenerName listenerName) {
+        Map parsedConfigs;
+        if (listenerName == null)
+            parsedConfigs = (Map) config.values();
+        else
+            parsedConfigs = config.valuesWithPrefixOverride(listenerName.configPrefix());
+
+        config.originals().entrySet().stream()
+            .filter(e -> !parsedConfigs.containsKey(e.getKey())) // exclude already parsed configs
+            // exclude already parsed listener prefix configs
+            .filter(e -> !(listenerName != null && e.getKey().startsWith(listenerName.configPrefix()) &&
+                parsedConfigs.containsKey(e.getKey().substring(listenerName.configPrefix().length()))))
+            // exclude keys like `{mechanism}.some.prop` if "listener.name." prefix is present and key `some.prop` exists in parsed configs.
+            .filter(e -> !(listenerName != null && parsedConfigs.containsKey(e.getKey().substring(e.getKey().indexOf('.') + 1))))
+            .forEach(e -> parsedConfigs.put(e.getKey(), e.getValue()));
+        return parsedConfigs;
+    }
+
     private static void requireNonNullMode(Mode mode, SecurityProtocol securityProtocol) {
         if (mode == null)
             throw new IllegalArgumentException("`mode` must be non-null if `securityProtocol` is `" + securityProtocol + "`");
@@ -137,12 +202,13 @@ private static org.apache.kafka.common.security.auth.PrincipalBuilder createPrin
     public static KafkaPrincipalBuilder createPrincipalBuilder(Map configs,
                                                                TransportLayer transportLayer,
                                                                Authenticator authenticator,
-                                                               KerberosShortNamer kerberosShortNamer) {
+                                                               KerberosShortNamer kerberosShortNamer,
+                                                               SslPrincipalMapper sslPrincipalMapper) {
         Class principalBuilderClass = (Class) configs.get(BrokerSecurityConfigs.PRINCIPAL_BUILDER_CLASS_CONFIG);
         final KafkaPrincipalBuilder builder;
 
         if (principalBuilderClass == null || principalBuilderClass == DefaultKafkaPrincipalBuilder.class) {
-            builder = new DefaultKafkaPrincipalBuilder(kerberosShortNamer);
+            builder = new DefaultKafkaPrincipalBuilder(kerberosShortNamer, sslPrincipalMapper);
         } else if (KafkaPrincipalBuilder.class.isAssignableFrom(principalBuilderClass)) {
             builder = (KafkaPrincipalBuilder) Utils.newInstance(principalBuilderClass);
         } else if (org.apache.kafka.common.security.auth.PrincipalBuilder.class.isAssignableFrom(principalBuilderClass)) {
diff --git a/clients/src/main/java/org/apache/kafka/common/network/ChannelMetadataRegistry.java b/clients/src/main/java/org/apache/kafka/common/network/ChannelMetadataRegistry.java
new file mode 100644
index 0000000000000..a3453d881c986
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/network/ChannelMetadataRegistry.java
@@ -0,0 +1,55 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.network;
+
+import java.io.Closeable;
+
+/**
+ * Metadata about a channel is provided in various places in the network stack. This
+ * registry is used as a common place to collect them.
+ */
+public interface ChannelMetadataRegistry extends Closeable {
+
+    /**
+     * Register information about the SSL cipher we are using.
+     *  Re-registering the information will overwrite the previous one.
+     */
+    void registerCipherInformation(CipherInformation cipherInformation);
+
+    /**
+     * Get the currently registered cipher information.
+     */
+    CipherInformation cipherInformation();
+
+    /**
+     * Register information about the client client we are using.
+     * Depending on the clients, the ApiVersionsRequest could be received
+     * multiple times or not at all. Re-registering the information will
+     * overwrite the previous one.
+     */
+    void registerClientInformation(ClientInformation clientInformation);
+
+    /**
+     * Get the currently registered client information.
+     */
+    ClientInformation clientInformation();
+
+    /**
+     * Unregister everything that has been registered and close the registry.
+     */
+    void close();
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/network/ChannelState.java b/clients/src/main/java/org/apache/kafka/common/network/ChannelState.java
index 08ed1a04c94f9..5f6dfb92bd284 100644
--- a/clients/src/main/java/org/apache/kafka/common/network/ChannelState.java
+++ b/clients/src/main/java/org/apache/kafka/common/network/ChannelState.java
@@ -62,7 +62,8 @@ public enum State {
         FAILED_SEND,
         AUTHENTICATION_FAILED,
         LOCAL_CLOSE
-    };
+    }
+
     // AUTHENTICATION_FAILED has a custom exception. For other states,
     // create a reusable `ChannelState` instance per-state.
     public static final ChannelState NOT_CONNECTED = new ChannelState(State.NOT_CONNECTED);
@@ -74,12 +75,20 @@ public enum State {
 
     private final State state;
     private final AuthenticationException exception;
+    private final String remoteAddress;
+
     public ChannelState(State state) {
-        this(state, null);
+        this(state, null, null);
     }
-    public ChannelState(State state, AuthenticationException exception) {
+
+    public ChannelState(State state, String remoteAddress) {
+        this(state, null, remoteAddress);
+    }
+    
+    public ChannelState(State state, AuthenticationException exception, String remoteAddress) {
         this.state = state;
         this.exception = exception;
+        this.remoteAddress = remoteAddress;
     }
 
     public State state() {
@@ -89,4 +98,8 @@ public State state() {
     public AuthenticationException exception() {
         return exception;
     }
+
+    public String remoteAddress() {
+        return remoteAddress;
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/network/CipherInformation.java b/clients/src/main/java/org/apache/kafka/common/network/CipherInformation.java
new file mode 100644
index 0000000000000..d65aeb981c009
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/network/CipherInformation.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.network;
+
+import java.util.Objects;
+
+public class CipherInformation {
+    private final String cipher;
+    private final String protocol;
+
+    public CipherInformation(String cipher, String protocol) {
+        this.cipher = cipher == null || cipher.isEmpty()  ? "unknown" : cipher;
+        this.protocol = protocol == null || protocol.isEmpty()  ? "unknown" : protocol;
+    }
+
+    public String cipher() {
+        return cipher;
+    }
+
+    public String protocol() {
+        return protocol;
+    }
+
+    @Override
+    public String toString() {
+        return "CipherInformation(cipher=" + cipher +
+            ", protocol=" + protocol + ")";
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(cipher, protocol);
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (o == null) {
+            return false;
+        }
+        if (!(o instanceof CipherInformation)) {
+            return false;
+        }
+        CipherInformation other = (CipherInformation) o;
+        return other.cipher.equals(cipher) &&
+            other.protocol.equals(protocol);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/network/ClientInformation.java b/clients/src/main/java/org/apache/kafka/common/network/ClientInformation.java
new file mode 100644
index 0000000000000..cb99a8669e40f
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/network/ClientInformation.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.network;
+
+import java.util.Objects;
+
+public class ClientInformation {
+    public static final String UNKNOWN_NAME_OR_VERSION = "unknown";
+    public static final ClientInformation EMPTY = new ClientInformation(UNKNOWN_NAME_OR_VERSION, UNKNOWN_NAME_OR_VERSION);
+
+    private final String softwareName;
+    private final String softwareVersion;
+
+    public ClientInformation(String softwareName, String softwareVersion) {
+        this.softwareName = softwareName.isEmpty() ? UNKNOWN_NAME_OR_VERSION : softwareName;
+        this.softwareVersion = softwareVersion.isEmpty() ? UNKNOWN_NAME_OR_VERSION : softwareVersion;
+    }
+
+    public String softwareName() {
+        return this.softwareName;
+    }
+
+    public String softwareVersion() {
+        return this.softwareVersion;
+    }
+
+    @Override
+    public String toString() {
+        return "ClientInformation(softwareName=" + softwareName +
+            ", softwareVersion=" + softwareVersion + ")";
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(softwareName, softwareVersion);
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (o == null) {
+            return false;
+        }
+        if (!(o instanceof ClientInformation)) {
+            return false;
+        }
+        ClientInformation other = (ClientInformation) o;
+        return other.softwareName.equals(softwareName) &&
+            other.softwareVersion.equals(softwareVersion);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/network/DefaultChannelMetadataRegistry.java b/clients/src/main/java/org/apache/kafka/common/network/DefaultChannelMetadataRegistry.java
new file mode 100644
index 0000000000000..ae9e9a83a0c2c
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/network/DefaultChannelMetadataRegistry.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.network;
+
+public class DefaultChannelMetadataRegistry implements ChannelMetadataRegistry {
+    private CipherInformation cipherInformation;
+    private ClientInformation clientInformation;
+
+    @Override
+    public void registerCipherInformation(final CipherInformation cipherInformation) {
+        if (this.cipherInformation != null) {
+            this.cipherInformation = cipherInformation;
+        }
+    }
+
+    @Override
+    public CipherInformation cipherInformation() {
+        return this.cipherInformation;
+    }
+
+    @Override
+    public void registerClientInformation(final ClientInformation clientInformation) {
+        this.clientInformation = clientInformation;
+    }
+
+    @Override
+    public ClientInformation clientInformation() {
+        return this.clientInformation;
+    }
+
+    @Override
+    public void close() {
+        this.cipherInformation = null;
+        this.clientInformation = null;
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/network/DelayedResponseAuthenticationException.java b/clients/src/main/java/org/apache/kafka/common/network/DelayedResponseAuthenticationException.java
new file mode 100644
index 0000000000000..8474426c6095c
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/network/DelayedResponseAuthenticationException.java
@@ -0,0 +1,27 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.network;
+
+import org.apache.kafka.common.errors.AuthenticationException;
+
+public class DelayedResponseAuthenticationException extends AuthenticationException {
+    private static final long serialVersionUID = 1L;
+
+    public DelayedResponseAuthenticationException(Throwable cause) {
+        super(cause);
+    }
+}
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 f07035a83b2bd..bc8228009ae7d 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
@@ -17,48 +17,142 @@
 package org.apache.kafka.common.network;
 
 import org.apache.kafka.common.errors.AuthenticationException;
+import org.apache.kafka.common.errors.SslAuthenticationException;
 import org.apache.kafka.common.memory.MemoryPool;
 import org.apache.kafka.common.security.auth.KafkaPrincipal;
+import org.apache.kafka.common.security.auth.KafkaPrincipalSerde;
 import org.apache.kafka.common.utils.Utils;
 
 import java.io.IOException;
 import java.net.InetAddress;
 import java.net.Socket;
+import java.net.SocketAddress;
 import java.nio.channels.SelectionKey;
-import java.util.Objects;
+import java.nio.channels.SocketChannel;
+import java.util.Optional;
+import java.util.function.Supplier;
+
+/**
+ * A Kafka connection either existing on a client (which could be a broker in an
+ * inter-broker scenario) and representing the channel to a remote broker or the
+ * reverse (existing on a broker and representing the channel to a remote
+ * client, which could be a broker in an inter-broker scenario).
+ * 
        
+ * Each instance has the following:
+ * 
          
+ * 
        • a unique ID identifying it in the {@code KafkaClient} instance via which
+ * the connection was made on the client-side or in the instance where it was
+ * accepted on the server-side
          • 
+ * 
        • a reference to the underlying {@link TransportLayer} to allow reading and
+ * writing
          • 
+ * 
        • an {@link Authenticator} that performs the authentication (or
+ * re-authentication, if that feature is enabled and it applies to this
+ * connection) by reading and writing directly from/to the same
+ * {@link TransportLayer}.
          • 
+ * 
        • a {@link MemoryPool} into which responses are read (typically the JVM
+ * heap for clients, though smaller pools can be used for brokers and for
+ * testing out-of-memory scenarios)
          • 
+ * 
        • a {@link NetworkReceive} representing the current incomplete/in-progress
+ * request (from the server-side perspective) or response (from the client-side
+ * perspective) being read, if applicable; or a non-null value that has had no
+ * data read into it yet or a null value if there is no in-progress
+ * request/response (either could be the case)
          • 
+ * 
        • a {@link Send} representing the current request (from the client-side
+ * perspective) or response (from the server-side perspective) that is either
+ * waiting to be sent or partially sent, if applicable, or null
          • 
+ * 
        • a {@link ChannelMuteState} to document if the channel has been muted due
+ * to memory pressure or other reasons
          • 
+ * 
        
+ */
+public class KafkaChannel implements AutoCloseable {
+    private static final long MIN_REAUTH_INTERVAL_ONE_SECOND_NANOS = 1000 * 1000 * 1000;
+
+    /**
+     * Mute States for KafkaChannel:
+     * 
          
+     *   
        •  NOT_MUTED: Channel is not muted. This is the default state. 
          • 
+     *   
        •  MUTED: Channel is muted. Channel must be in this state to be unmuted. 
          • 
+     *   
        •  MUTED_AND_RESPONSE_PENDING: (SocketServer only) Channel is muted and SocketServer has not sent a response
+     *                                    back to the client yet (acks != 0) or is currently waiting to receive a
+     *                                    response from the API layer (acks == 0). 
          • 
+     *   
        •  MUTED_AND_THROTTLED: (SocketServer only) Channel is muted and throttling is in progress due to quota
+     *                             violation. 
          • 
+     *   
        •  MUTED_AND_THROTTLED_AND_RESPONSE_PENDING: (SocketServer only) Channel is muted, throttling is in progress,
+     *                                                  and a response is currently pending. 
          • 
+     * 
        
+     */
+    public enum ChannelMuteState {
+        NOT_MUTED,
+        MUTED,
+        MUTED_AND_RESPONSE_PENDING,
+        MUTED_AND_THROTTLED,
+        MUTED_AND_THROTTLED_AND_RESPONSE_PENDING
+    }
+
+    /** Socket server events that will change the mute state:
+     * 
          
+     *   
        •  REQUEST_RECEIVED: A request has been received from the client. 
          • 
+     *   
        •  RESPONSE_SENT: A response has been sent out to the client (ack != 0) or SocketServer has heard back from
+     *                       the API layer (acks = 0) 
          • 
+     *   
        •  THROTTLE_STARTED: Throttling started due to quota violation. 
          • 
+     *   
        •  THROTTLE_ENDED: Throttling ended. 
          • 
+     * 
        
+     *
+     * Valid transitions on each event are:
+     * 
          
+     *   
        •  REQUEST_RECEIVED: MUTED => MUTED_AND_RESPONSE_PENDING 
          • 
+     *   
        •  RESPONSE_SENT:    MUTED_AND_RESPONSE_PENDING => MUTED, MUTED_AND_THROTTLED_AND_RESPONSE_PENDING => MUTED_AND_THROTTLED 
          • 
+     *   
        •  THROTTLE_STARTED: MUTED_AND_RESPONSE_PENDING => MUTED_AND_THROTTLED_AND_RESPONSE_PENDING 
          • 
+     *   
        •  THROTTLE_ENDED:   MUTED_AND_THROTTLED => MUTED, MUTED_AND_THROTTLED_AND_RESPONSE_PENDING => MUTED_AND_RESPONSE_PENDING 
          • 
+     * 
        
+     */
+    public enum ChannelMuteEvent {
+        REQUEST_RECEIVED,
+        RESPONSE_SENT,
+        THROTTLE_STARTED,
+        THROTTLE_ENDED
+    }
 
-public class KafkaChannel {
     private final String id;
     private final TransportLayer transportLayer;
-    private final Authenticator authenticator;
+    private final Supplier authenticatorCreator;
+    private Authenticator authenticator;
     // Tracks accumulated network thread time. This is updated on the network thread.
     // The values are read and reset after each response is sent.
     private long networkThreadTimeNanos;
     private final int maxReceiveSize;
     private final MemoryPool memoryPool;
+    private final ChannelMetadataRegistry metadataRegistry;
     private NetworkReceive receive;
-    private Send send;
+    private NetworkSend send;
     // Track connection and mute state of channels to enable outstanding requests on channels to be
     // processed after the channel is disconnected.
     private boolean disconnected;
-    private boolean muted;
+    private ChannelMuteState muteState;
     private ChannelState state;
+    private SocketAddress remoteAddress;
+    private int successfulAuthentications;
+    private boolean midWrite;
+    private long lastReauthenticationStartNanos;
 
-    public KafkaChannel(String id, TransportLayer transportLayer, Authenticator authenticator, int maxReceiveSize, MemoryPool memoryPool) throws IOException {
+    public KafkaChannel(String id, TransportLayer transportLayer, Supplier authenticatorCreator,
+                        int maxReceiveSize, MemoryPool memoryPool, ChannelMetadataRegistry metadataRegistry) {
         this.id = id;
         this.transportLayer = transportLayer;
-        this.authenticator = authenticator;
+        this.authenticatorCreator = authenticatorCreator;
+        this.authenticator = authenticatorCreator.get();
         this.networkThreadTimeNanos = 0L;
         this.maxReceiveSize = maxReceiveSize;
         this.memoryPool = memoryPool;
+        this.metadataRegistry = metadataRegistry;
         this.disconnected = false;
-        this.muted = false;
+        this.muteState = ChannelMuteState.NOT_MUTED;
         this.state = ChannelState.NOT_CONNECTED;
     }
 
     public void close() throws IOException {
         this.disconnected = true;
-        Utils.closeAll(transportLayer, authenticator, receive);
+        Utils.closeAll(transportLayer, authenticator, receive, metadataRegistry);
     }
 
     /**
@@ -68,29 +162,47 @@ public KafkaPrincipal principal() {
         return authenticator.principal();
     }
 
+    public Optional principalSerde() {
+        return authenticator.principalSerde();
+    }
+
     /**
      * Does handshake of transportLayer and authentication using configured authenticator.
      * For SSL with client authentication enabled, {@link TransportLayer#handshake()} performs
      * authentication. For SASL, authentication is performed by {@link Authenticator#authenticate()}.
      */
     public void prepare() throws AuthenticationException, IOException {
+        boolean authenticating = false;
         try {
             if (!transportLayer.ready())
                 transportLayer.handshake();
-            if (transportLayer.ready() && !authenticator.complete())
+            if (transportLayer.ready() && !authenticator.complete()) {
+                authenticating = true;
                 authenticator.authenticate();
+            }
         } catch (AuthenticationException e) {
             // Clients are notified of authentication exceptions to enable operations to be terminated
             // without retries. Other errors are handled as network exceptions in Selector.
-            state = new ChannelState(ChannelState.State.AUTHENTICATION_FAILED, e);
+            String remoteDesc = remoteAddress != null ? remoteAddress.toString() : null;
+            state = new ChannelState(ChannelState.State.AUTHENTICATION_FAILED, e, remoteDesc);
+            if (authenticating) {
+                delayCloseOnAuthenticationFailure();
+                throw new DelayedResponseAuthenticationException(e);
+            }
             throw e;
         }
-        if (ready())
+        if (ready()) {
+            ++successfulAuthentications;
             state = ChannelState.READY;
+        }
     }
 
     public void disconnect() {
         disconnected = true;
+        if (state == ChannelState.NOT_CONNECTED && remoteAddress != null) {
+            //if we captured the remote address we can provide more information
+            state = new ChannelState(ChannelState.State.NOT_CONNECTED, remoteAddress.toString());
+        }
         transportLayer.disconnect();
     }
 
@@ -103,9 +215,22 @@ public ChannelState state() {
     }
 
     public boolean finishConnect() throws IOException {
+        //we need to grab remoteAddr before finishConnect() is called otherwise
+        //it becomes inaccessible if the connection was refused.
+        SocketChannel socketChannel = transportLayer.socketChannel();
+        if (socketChannel != null) {
+            remoteAddress = socketChannel.getRemoteAddress();
+        }
         boolean connected = transportLayer.finishConnect();
-        if (connected)
-            state = ready() ? ChannelState.READY : ChannelState.AUTHENTICATE;
+        if (connected) {
+            if (ready()) {
+                state = ChannelState.READY;
+            } else if (remoteAddress != null) {
+                state = new ChannelState(ChannelState.State.AUTHENTICATE, remoteAddress.toString());
+            } else {
+                state = ChannelState.AUTHENTICATE;
+            }
+        }
         return connected;
     }
 
@@ -117,26 +242,102 @@ public String id() {
         return id;
     }
 
+    public SelectionKey selectionKey() {
+        return transportLayer.selectionKey();
+    }
+
     /**
      * externally muting a channel should be done via selector to ensure proper state handling
      */
     void mute() {
-        if (!disconnected)
-            transportLayer.removeInterestOps(SelectionKey.OP_READ);
-        muted = true;
+        if (muteState == ChannelMuteState.NOT_MUTED) {
+            if (!disconnected) transportLayer.removeInterestOps(SelectionKey.OP_READ);
+            muteState = ChannelMuteState.MUTED;
+        }
     }
 
-    void unmute() {
-        if (!disconnected)
-            transportLayer.addInterestOps(SelectionKey.OP_READ);
-        muted = false;
+    /**
+     * Unmute the channel. The channel can be unmuted only if it is in the MUTED state. For other muted states
+     * (MUTED_AND_*), this is a no-op.
+     *
+     * @return Whether or not the channel is in the NOT_MUTED state after the call
+     */
+    boolean maybeUnmute() {
+        if (muteState == ChannelMuteState.MUTED) {
+            if (!disconnected) transportLayer.addInterestOps(SelectionKey.OP_READ);
+            muteState = ChannelMuteState.NOT_MUTED;
+        }
+        return muteState == ChannelMuteState.NOT_MUTED;
+    }
+
+    // Handle the specified channel mute-related event and transition the mute state according to the state machine.
+    public void handleChannelMuteEvent(ChannelMuteEvent event) {
+        boolean stateChanged = false;
+        switch (event) {
+            case REQUEST_RECEIVED:
+                if (muteState == ChannelMuteState.MUTED) {
+                    muteState = ChannelMuteState.MUTED_AND_RESPONSE_PENDING;
+                    stateChanged = true;
+                }
+                break;
+            case RESPONSE_SENT:
+                if (muteState == ChannelMuteState.MUTED_AND_RESPONSE_PENDING) {
+                    muteState = ChannelMuteState.MUTED;
+                    stateChanged = true;
+                }
+                if (muteState == ChannelMuteState.MUTED_AND_THROTTLED_AND_RESPONSE_PENDING) {
+                    muteState = ChannelMuteState.MUTED_AND_THROTTLED;
+                    stateChanged = true;
+                }
+                break;
+            case THROTTLE_STARTED:
+                if (muteState == ChannelMuteState.MUTED_AND_RESPONSE_PENDING) {
+                    muteState = ChannelMuteState.MUTED_AND_THROTTLED_AND_RESPONSE_PENDING;
+                    stateChanged = true;
+                }
+                break;
+            case THROTTLE_ENDED:
+                if (muteState == ChannelMuteState.MUTED_AND_THROTTLED) {
+                    muteState = ChannelMuteState.MUTED;
+                    stateChanged = true;
+                }
+                if (muteState == ChannelMuteState.MUTED_AND_THROTTLED_AND_RESPONSE_PENDING) {
+                    muteState = ChannelMuteState.MUTED_AND_RESPONSE_PENDING;
+                    stateChanged = true;
+                }
+        }
+        if (!stateChanged) {
+            throw new IllegalStateException("Cannot transition from " + muteState.name() + " for " + event.name());
+        }
+    }
+
+    public ChannelMuteState muteState() {
+        return muteState;
+    }
+
+    /**
+     * Delay channel close on authentication failure. This will remove all read/write operations from the channel until
+     * {@link #completeCloseOnAuthenticationFailure()} is called to finish up the channel close.
+     */
+    private void delayCloseOnAuthenticationFailure() {
+        transportLayer.removeInterestOps(SelectionKey.OP_WRITE);
+    }
+
+    /**
+     * Finish up any processing on {@link #prepare()} failure.
+     * @throws IOException
+     */
+    void completeCloseOnAuthenticationFailure() throws IOException {
+        transportLayer.addInterestOps(SelectionKey.OP_WRITE);
+        // Invoke the underlying handler to finish up any processing on authentication failure
+        authenticator.handleAuthenticationFailure();
     }
 
     /**
      * Returns true if this channel has been explicitly muted using {@link KafkaChannel#mute()}
      */
-    public boolean isMute() {
-        return muted;
+    public boolean isMuted() {
+        return muteState != ChannelMuteState.NOT_MUTED;
     }
 
     public boolean isInMutableState() {
@@ -175,39 +376,58 @@ public String socketDescription() {
         return socket.getInetAddress().toString();
     }
 
-    public void setSend(Send send) {
+    public void setSend(NetworkSend send) {
         if (this.send != null)
             throw new IllegalStateException("Attempt to begin a send operation with prior send operation still in progress, connection id is " + id);
         this.send = send;
         this.transportLayer.addInterestOps(SelectionKey.OP_WRITE);
     }
 
-    public NetworkReceive read() throws IOException {
-        NetworkReceive result = null;
+    public NetworkSend maybeCompleteSend() {
+        if (send != null && send.completed()) {
+            midWrite = false;
+            transportLayer.removeInterestOps(SelectionKey.OP_WRITE);
+            NetworkSend result = send;
+            send = null;
+            return result;
+        }
+        return null;
+    }
 
+    public long read() throws IOException {
         if (receive == null) {
             receive = new NetworkReceive(maxReceiveSize, id, memoryPool);
         }
 
-        receive(receive);
-        if (receive.complete()) {
-            receive.payload().rewind();
-            result = receive;
-            receive = null;
-        } else if (receive.requiredMemoryAmountKnown() && !receive.memoryAllocated() && isInMutableState()) {
+        long bytesReceived = receive(this.receive);
+
+        if (this.receive.requiredMemoryAmountKnown() && !this.receive.memoryAllocated() && isInMutableState()) {
             //pool must be out of memory, mute ourselves.
             mute();
         }
-        return result;
+        return bytesReceived;
     }
 
-    public Send write() throws IOException {
-        Send result = null;
-        if (send != null && send(send)) {
-            result = send;
-            send = null;
+    public NetworkReceive currentReceive() {
+        return receive;
+    }
+
+    public NetworkReceive maybeCompleteReceive() {
+        if (receive != null && receive.complete()) {
+            receive.payload().rewind();
+            NetworkReceive result = receive;
+            receive = null;
+            return result;
         }
-        return result;
+        return null;
+    }
+
+    public long write() throws IOException {
+        if (send == null)
+            return 0;
+
+        midWrite = true;
+        return send.writeTo(transportLayer);
     }
 
     /**
@@ -228,15 +448,15 @@ public long getAndResetNetworkThreadTimeNanos() {
     }
 
     private long receive(NetworkReceive receive) throws IOException {
-        return receive.readFrom(transportLayer);
-    }
-
-    private boolean send(Send send) throws IOException {
-        send.writeTo(transportLayer);
-        if (send.completed())
-            transportLayer.removeInterestOps(SelectionKey.OP_WRITE);
-
-        return send.completed();
+        try {
+            return receive.readFrom(transportLayer);
+        } catch (SslAuthenticationException e) {
+            // With TLSv1.3, post-handshake messages may throw SSLExceptions, which are
+            // handled as authentication failures
+            String remoteDesc = remoteAddress != null ? remoteAddress.toString() : null;
+            state = new ChannelState(ChannelState.State.AUTHENTICATION_FAILED, e, remoteDesc);
+            throw e;
+        }
     }
 
     /**
@@ -255,11 +475,200 @@ public boolean equals(Object o) {
             return false;
         }
         KafkaChannel that = (KafkaChannel) o;
-        return Objects.equals(id, that.id);
+        return id.equals(that.id);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(id);
+        return id.hashCode();
+    }
+
+    @Override
+    public String toString() {
+        return super.toString() + " id=" + id;
+    }
+    
+    /**
+     * Return the number of times this instance has successfully authenticated. This
+     * value can only exceed 1 when re-authentication is enabled and it has
+     * succeeded at least once.
+     * 
+     * @return the number of times this instance has successfully authenticated
+     */
+    public int successfulAuthentications() {
+        return successfulAuthentications;
+    }
+
+    /**
+     * If this is a server-side connection that has an expiration time and at least
+     * 1 second has passed since the prior re-authentication (if any) started then
+     * begin the process of re-authenticating the connection and return true,
+     * otherwise return false
+     * 
+     * @param saslHandshakeNetworkReceive
+     *            the mandatory {@link NetworkReceive} containing the
+     *            {@code SaslHandshakeRequest} that has been received on the server
+     *            and that initiates re-authentication.
+     * @param nowNanosSupplier
+     *            {@code Supplier} of the current time. The value must be in
+     *            nanoseconds as per {@code System.nanoTime()} and is therefore only
+     *            useful when compared to such a value -- it's absolute value is
+     *            meaningless.
+     * 
+     * @return true if this is a server-side connection that has an expiration time
+     *         and at least 1 second has passed since the prior re-authentication
+     *         (if any) started to indicate that the re-authentication process has
+     *         begun, otherwise false
+     * @throws AuthenticationException
+     *             if re-authentication fails due to invalid credentials or other
+     *             security configuration errors
+     * @throws IOException
+     *             if read/write fails due to an I/O error
+     * @throws IllegalStateException
+     *             if this channel is not "ready"
+     */
+    public boolean maybeBeginServerReauthentication(NetworkReceive saslHandshakeNetworkReceive,
+            Supplier nowNanosSupplier) throws AuthenticationException, IOException {
+        if (!ready())
+            throw new IllegalStateException(
+                    "KafkaChannel should be \"ready\" when processing SASL Handshake for potential re-authentication");
+        /*
+         * Re-authentication is disabled if there is no session expiration time, in
+         * which case the SASL handshake network receive will be processed normally,
+         * which results in a failure result being sent to the client. Also, no need to
+         * check if we are muted since since we are processing a received packet when we
+         * invoke this.
+         */
+        if (authenticator.serverSessionExpirationTimeNanos() == null)
+            return false;
+        /*
+         * 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();
+        /*
+         * Cannot re-authenticate more than once every second; an attempt to do so will
+         * result in the SASL handshake network receive being processed normally, which
+         * results in a failure result being sent to the client.
+         */
+        if (lastReauthenticationStartNanos != 0
+                && nowNanos - lastReauthenticationStartNanos < MIN_REAUTH_INTERVAL_ONE_SECOND_NANOS)
+            return false;
+        lastReauthenticationStartNanos = nowNanos;
+        swapAuthenticatorsAndBeginReauthentication(
+                new ReauthenticationContext(authenticator, saslHandshakeNetworkReceive, nowNanos));
+        return true;
+    }
+
+    /**
+     * If this is a client-side connection that is not muted, there is no
+     * in-progress write, and there is a session expiration time defined that has
+     * past then begin the process of re-authenticating the connection and return
+     * true, otherwise return false
+     * 
+     * @param nowNanosSupplier
+     *            {@code Supplier} of the current time. The value must be in
+     *            nanoseconds as per {@code System.nanoTime()} and is therefore only
+     *            useful when compared to such a value -- it's absolute value is
+     *            meaningless.
+     * 
+     * @return true if this is a client-side connection that is not muted, there is
+     *         no in-progress write, and there is a session expiration time defined
+     *         that has past to indicate that the re-authentication process has
+     *         begun, otherwise false
+     * @throws AuthenticationException
+     *             if re-authentication fails due to invalid credentials or other
+     *             security configuration errors
+     * @throws IOException
+     *             if read/write fails due to an I/O error
+     * @throws IllegalStateException
+     *             if this channel is not "ready"
+     */
+    public boolean maybeBeginClientReauthentication(Supplier nowNanosSupplier)
+            throws AuthenticationException, IOException {
+        if (!ready())
+            throw new IllegalStateException(
+                    "KafkaChannel should always be \"ready\" when it is checked for possible re-authentication");
+        if (muteState != ChannelMuteState.NOT_MUTED || midWrite
+                || authenticator.clientSessionReauthenticationTimeNanos() == null)
+            return false;
+        /*
+         * 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();
+        if (nowNanos < authenticator.clientSessionReauthenticationTimeNanos())
+            return false;
+        swapAuthenticatorsAndBeginReauthentication(new ReauthenticationContext(authenticator, receive, nowNanos));
+        receive = null;
+        return true;
+    }
+    
+    /**
+     * Return the number of milliseconds that elapsed while re-authenticating this
+     * session from the perspective of this instance, if applicable, otherwise null.
+     * The server-side perspective will yield a lower value than the client-side
+     * perspective of the same re-authentication because the client-side observes an
+     * additional network round-trip.
+     * 
+     * @return the number of milliseconds that elapsed while re-authenticating this
+     *         session from the perspective of this instance, if applicable,
+     *         otherwise null
+     */
+    public Long reauthenticationLatencyMs() {
+        return authenticator.reauthenticationLatencyMs();
+    }
+
+    /**
+     * Return true if this is a server-side channel and the given time is past the
+     * session expiration time, if any, otherwise false
+     * 
+     * @param nowNanos
+     *            the current time in nanoseconds as per {@code System.nanoTime()}
+     * @return true if this is a server-side channel and the given time is past the
+     *         session expiration time, if any, otherwise false
+     */
+    public boolean serverAuthenticationSessionExpired(long nowNanos) {
+        Long serverSessionExpirationTimeNanos = authenticator.serverSessionExpirationTimeNanos();
+        return serverSessionExpirationTimeNanos != null && nowNanos - serverSessionExpirationTimeNanos > 0;
+    }
+    
+    /**
+     * Return the (always non-null but possibly empty) client-side
+     * {@link NetworkReceive} response that arrived during re-authentication but
+     * is unrelated to re-authentication. This corresponds to a request sent
+     * prior to the beginning of re-authentication; the request was made when the
+     * channel was successfully authenticated, and the response arrived during the
+     * re-authentication process.
+     * 
+     * @return client-side {@link NetworkReceive} response that arrived during
+     *         re-authentication that is unrelated to re-authentication. This may
+     *         be empty.
+     */
+    public Optional pollResponseReceivedDuringReauthentication() {
+        return authenticator.pollResponseReceivedDuringReauthentication();
+    }
+    
+    /**
+     * Return true if this is a server-side channel and the connected client has
+     * indicated that it supports re-authentication, otherwise false
+     * 
+     * @return true if this is a server-side channel and the connected client has
+     *         indicated that it supports re-authentication, otherwise false
+     */
+    boolean connectedClientSupportsReauthentication() {
+        return authenticator.connectedClientSupportsReauthentication();
+    }
+
+    private void swapAuthenticatorsAndBeginReauthentication(ReauthenticationContext reauthenticationContext)
+            throws IOException {
+        // it is up to the new authenticator to close the old one
+        // replace with a new one and begin the process of re-authenticating
+        authenticator = authenticatorCreator.get();
+        authenticator.reauthenticate(reauthenticationContext);
+    }
+
+    public ChannelMetadataRegistry channelMetadataRegistry() {
+        return metadataRegistry;
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/network/ListenerName.java b/clients/src/main/java/org/apache/kafka/common/network/ListenerName.java
index 9da4cca97a19f..2decccbc50649 100644
--- a/clients/src/main/java/org/apache/kafka/common/network/ListenerName.java
+++ b/clients/src/main/java/org/apache/kafka/common/network/ListenerName.java
@@ -71,4 +71,12 @@ public String toString() {
     public String configPrefix() {
         return CONFIG_STATIC_PREFIX + "." + value.toLowerCase(Locale.ROOT) + ".";
     }
+
+    public String saslMechanismConfigPrefix(String saslMechanism) {
+        return configPrefix() + saslMechanismPrefix(saslMechanism);
+    }
+
+    public static String saslMechanismPrefix(String saslMechanism) {
+        return saslMechanism.toLowerCase(Locale.ROOT) + ".";
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/network/ListenerReconfigurable.java b/clients/src/main/java/org/apache/kafka/common/network/ListenerReconfigurable.java
new file mode 100644
index 0000000000000..3541212ceff84
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/network/ListenerReconfigurable.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.network;
+
+import org.apache.kafka.common.Reconfigurable;
+
+/**
+ * Interface for reconfigurable entities associated with a listener.
+ */
+public interface ListenerReconfigurable extends Reconfigurable {
+
+    /**
+     * Returns the listener name associated with this reconfigurable. Listener-specific
+     * configs corresponding to this listener name are provided for reconfiguration.
+     */
+    ListenerName listenerName();
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/network/MultiSend.java b/clients/src/main/java/org/apache/kafka/common/network/MultiSend.java
deleted file mode 100644
index f77ff971e8039..0000000000000
--- a/clients/src/main/java/org/apache/kafka/common/network/MultiSend.java
+++ /dev/null
@@ -1,100 +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.network;
-
-import org.apache.kafka.common.KafkaException;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.io.IOException;
-import java.nio.channels.GatheringByteChannel;
-import java.util.Iterator;
-import java.util.List;
-
-/**
- * A set of composite sends, sent one after another
- */
-
-public class MultiSend implements Send {
-
-    private static final Logger log = LoggerFactory.getLogger(MultiSend.class);
-
-    private final String dest;
-    private final Iterator sendsIterator;
-    private final long size;
-
-    private long totalWritten = 0;
-    private Send current;
-
-    public MultiSend(String dest, List sends) {
-        this.dest = dest;
-        this.sendsIterator = sends.iterator();
-        nextSendOrDone();
-        long size = 0;
-        for (Send send : sends)
-            size += send.size();
-        this.size = size;
-    }
-
-    @Override
-    public long size() {
-        return size;
-    }
-
-    @Override
-    public String destination() {
-        return dest;
-    }
-
-    @Override
-    public boolean completed() {
-        return current == null;
-    }
-
-    @Override
-    public long writeTo(GatheringByteChannel channel) throws IOException {
-        if (completed())
-            throw new KafkaException("This operation cannot be invoked on a complete request.");
-
-        int totalWrittenPerCall = 0;
-        boolean sendComplete;
-        do {
-            long written = current.writeTo(channel);
-            totalWrittenPerCall += written;
-            sendComplete = current.completed();
-            if (sendComplete)
-                nextSendOrDone();
-        } while (!completed() && sendComplete);
-
-        totalWritten += totalWrittenPerCall;
-
-        if (completed() && totalWritten != size)
-            log.error("mismatch in sending bytes over socket; expected: " + size + " actual: " + totalWritten);
-
-        log.trace("Bytes written as part of multi-send call: {}, total bytes written so far: {}, expected bytes to write: {}",
-                totalWrittenPerCall, totalWritten, size);
-
-        return totalWrittenPerCall;
-    }
-
-    private void nextSendOrDone() {
-        if (sendsIterator.hasNext())
-            current = sendsIterator.next();
-        else
-            current = null;
-    }
-}
diff --git a/clients/src/main/java/org/apache/kafka/common/network/NetworkReceive.java b/clients/src/main/java/org/apache/kafka/common/network/NetworkReceive.java
index 564fbcd8c217c..5332c8109f360 100644
--- a/clients/src/main/java/org/apache/kafka/common/network/NetworkReceive.java
+++ b/clients/src/main/java/org/apache/kafka/common/network/NetworkReceive.java
@@ -16,14 +16,14 @@
  */
 package org.apache.kafka.common.network;
 
+import org.apache.kafka.common.memory.MemoryPool;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
 import java.io.EOFException;
 import java.io.IOException;
 import java.nio.ByteBuffer;
-import java.nio.channels.ReadableByteChannel;
 import java.nio.channels.ScatteringByteChannel;
-import org.apache.kafka.common.memory.MemoryPool;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
 /**
  * A size delimited Receive that consists of a 4 byte network-ordered size N followed by N bytes of content
@@ -90,33 +90,6 @@ public boolean complete() {
     }
 
     public long readFrom(ScatteringByteChannel channel) throws IOException {
-        return readFromReadableChannel(channel);
-    }
-
-    @Override
-    public boolean requiredMemoryAmountKnown() {
-        return requestedBufferSize != -1;
-    }
-
-    @Override
-    public boolean memoryAllocated() {
-        return buffer != null;
-    }
-
-
-    @Override
-    public void close() throws IOException {
-        if (buffer != null && buffer != EMPTY_BUFFER) {
-            memoryPool.release(buffer);
-            buffer = null;
-        }
-    }
-
-    // Need a method to read from ReadableByteChannel because BlockingChannel requires read with timeout
-    // See: http://stackoverflow.com/questions/2866557/timeout-for-socketchannel-doesnt-work
-    // This can go away after we get rid of BlockingChannel
-    @Deprecated
-    public long readFromReadableChannel(ReadableByteChannel channel) throws IOException {
         int read = 0;
         if (size.hasRemaining()) {
             int bytesRead = channel.read(size);
@@ -151,8 +124,41 @@ public long readFromReadableChannel(ReadableByteChannel channel) throws IOExcept
         return read;
     }
 
+    @Override
+    public boolean requiredMemoryAmountKnown() {
+        return requestedBufferSize != -1;
+    }
+
+    @Override
+    public boolean memoryAllocated() {
+        return buffer != null;
+    }
+
+
+    @Override
+    public void close() throws IOException {
+        if (buffer != null && buffer != EMPTY_BUFFER) {
+            memoryPool.release(buffer);
+            buffer = null;
+        }
+    }
+
     public ByteBuffer payload() {
         return this.buffer;
     }
 
+    public int bytesRead() {
+        if (buffer == null)
+            return size.position();
+        return buffer.position() + size.position();
+    }
+
+    /**
+     * Returns the total size of the receive including payload and size buffer
+     * for use in metrics. This is consistent with {@link NetworkSend#size()}
+     */
+    public int size() {
+        return payload().limit() + size.limit();
+    }
+
 }
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..b0dcc022efb9d 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
@@ -16,26 +16,35 @@
  */
 package org.apache.kafka.common.network;
 
-import java.nio.ByteBuffer;
+import java.io.IOException;
+import java.nio.channels.GatheringByteChannel;
 
-/**
- * A size delimited Send that consists of a 4 byte network-ordered size N followed by N bytes of content
- */
-public class NetworkSend extends ByteBufferSend {
+public class NetworkSend implements Send {
+    private final String destinationId;
+    private final Send send;
+
+    public NetworkSend(String destinationId, Send send) {
+        this.destinationId = destinationId;
+        this.send = send;
+    }
+
+    public String destinationId() {
+        return destinationId;
+    }
 
-    public NetworkSend(String destination, ByteBuffer buffer) {
-        super(destination, sizeDelimit(buffer));
+    @Override
+    public boolean completed() {
+        return send.completed();
     }
 
-    private static ByteBuffer[] sizeDelimit(ByteBuffer buffer) {
-        return new ByteBuffer[] {sizeBuffer(buffer.remaining()), buffer};
+    @Override
+    public long writeTo(GatheringByteChannel channel) throws IOException {
+        return send.writeTo(channel);
     }
 
-    private static ByteBuffer sizeBuffer(int size) {
-        ByteBuffer sizeBuffer = ByteBuffer.allocate(4);
-        sizeBuffer.putInt(size);
-        sizeBuffer.rewind();
-        return sizeBuffer;
+    @Override
+    public long size() {
+        return send.size();
     }
 
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/network/PlaintextChannelBuilder.java b/clients/src/main/java/org/apache/kafka/common/network/PlaintextChannelBuilder.java
index c0d1059aea95a..9369253e5982c 100644
--- a/clients/src/main/java/org/apache/kafka/common/network/PlaintextChannelBuilder.java
+++ b/clients/src/main/java/org/apache/kafka/common/network/PlaintextChannelBuilder.java
@@ -20,6 +20,7 @@
 import org.apache.kafka.common.memory.MemoryPool;
 import org.apache.kafka.common.security.auth.KafkaPrincipal;
 import org.apache.kafka.common.security.auth.KafkaPrincipalBuilder;
+import org.apache.kafka.common.security.auth.KafkaPrincipalSerde;
 import org.apache.kafka.common.security.auth.PlaintextAuthenticationContext;
 import org.apache.kafka.common.utils.Utils;
 import org.slf4j.Logger;
@@ -30,47 +31,79 @@
 import java.net.InetAddress;
 import java.nio.channels.SelectionKey;
 import java.util.Map;
+import java.util.Optional;
+import java.util.function.Supplier;
 
 public class PlaintextChannelBuilder implements ChannelBuilder {
     private static final Logger log = LoggerFactory.getLogger(PlaintextChannelBuilder.class);
+    private final ListenerName listenerName;
     private Map configs;
 
+    /**
+     * Constructs a plaintext channel builder. ListenerName is non-null whenever
+     * it's instantiated in the broker and null otherwise.
+     */
+    public PlaintextChannelBuilder(ListenerName listenerName) {
+        this.listenerName = listenerName;
+    }
+
     public void configure(Map configs) throws KafkaException {
         this.configs = configs;
     }
 
     @Override
-    public KafkaChannel buildChannel(String id, SelectionKey key, int maxReceiveSize, MemoryPool memoryPool) throws KafkaException {
+    public KafkaChannel buildChannel(String id, SelectionKey key, int maxReceiveSize,
+                                     MemoryPool memoryPool, ChannelMetadataRegistry metadataRegistry) throws KafkaException {
         try {
-            PlaintextTransportLayer transportLayer = new PlaintextTransportLayer(key);
-            PlaintextAuthenticator authenticator = new PlaintextAuthenticator(configs, transportLayer);
-            return new KafkaChannel(id, transportLayer, authenticator, maxReceiveSize,
-                    memoryPool != null ? memoryPool : MemoryPool.NONE);
+            PlaintextTransportLayer transportLayer = buildTransportLayer(key);
+            Supplier authenticatorCreator = () -> new PlaintextAuthenticator(configs, transportLayer, listenerName);
+            return buildChannel(id, transportLayer, authenticatorCreator, maxReceiveSize,
+                    memoryPool != null ? memoryPool : MemoryPool.NONE, metadataRegistry);
         } catch (Exception e) {
             log.warn("Failed to create channel due to ", e);
             throw new KafkaException(e);
         }
     }
 
+    // visible for testing
+    KafkaChannel buildChannel(String id, TransportLayer transportLayer, Supplier authenticatorCreator,
+                              int maxReceiveSize, MemoryPool memoryPool, ChannelMetadataRegistry metadataRegistry) {
+        return new KafkaChannel(id, transportLayer, authenticatorCreator, maxReceiveSize, memoryPool, metadataRegistry);
+    }
+
+    protected PlaintextTransportLayer buildTransportLayer(SelectionKey key) throws IOException {
+        return new PlaintextTransportLayer(key);
+    }
+
     @Override
     public void close() {}
 
     private static class PlaintextAuthenticator implements Authenticator {
         private final PlaintextTransportLayer transportLayer;
         private final KafkaPrincipalBuilder principalBuilder;
+        private final ListenerName listenerName;
 
-        private PlaintextAuthenticator(Map configs, PlaintextTransportLayer transportLayer) {
+        private PlaintextAuthenticator(Map configs, PlaintextTransportLayer transportLayer, ListenerName listenerName) {
             this.transportLayer = transportLayer;
-            this.principalBuilder = ChannelBuilders.createPrincipalBuilder(configs, transportLayer, this, null);
+            this.principalBuilder = ChannelBuilders.createPrincipalBuilder(configs, transportLayer, this, null, null);
+            this.listenerName = listenerName;
         }
 
         @Override
-        public void authenticate() throws IOException {}
+        public void authenticate() {}
 
         @Override
         public KafkaPrincipal principal() {
             InetAddress clientAddress = transportLayer.socketChannel().socket().getInetAddress();
-            return principalBuilder.build(new PlaintextAuthenticationContext(clientAddress));
+            // listenerName should only be null in Client mode where principal() should not be called
+            if (listenerName == null)
+                throw new IllegalStateException("Unexpected call to principal() when listenerName is null");
+            return principalBuilder.build(new PlaintextAuthenticationContext(clientAddress, listenerName.value()));
+        }
+
+        @Override
+        public Optional principalSerde() {
+            return principalBuilder instanceof KafkaPrincipalSerde ? Optional.of((KafkaPrincipalSerde) principalBuilder) : Optional.empty();
         }
 
         @Override
diff --git a/clients/src/main/java/org/apache/kafka/common/network/PlaintextTransportLayer.java b/clients/src/main/java/org/apache/kafka/common/network/PlaintextTransportLayer.java
index 11c95655067c8..845b1474f4e31 100644
--- a/clients/src/main/java/org/apache/kafka/common/network/PlaintextTransportLayer.java
+++ b/clients/src/main/java/org/apache/kafka/common/network/PlaintextTransportLayer.java
@@ -63,6 +63,11 @@ public SocketChannel socketChannel() {
         return socketChannel;
     }
 
+    @Override
+    public SelectionKey selectionKey() {
+        return key;
+    }
+
     @Override
     public boolean isOpen() {
         return socketChannel.isOpen();
@@ -73,29 +78,18 @@ public boolean isConnected() {
         return socketChannel.isConnected();
     }
 
-    /**
-     * Closes this channel
-     *
-     * @throws IOException If I/O error occurs
-     */
     @Override
     public void close() throws IOException {
-        try {
-            socketChannel.socket().close();
-            socketChannel.close();
-        } finally {
-            key.attach(null);
-            key.cancel();
-        }
+        socketChannel.socket().close();
+        socketChannel.close();
     }
 
     /**
      * Performs SSL handshake hence is a no-op for the non-secure
      * implementation
-     * @throws IOException
-    */
+     */
     @Override
-    public void handshake() throws IOException {}
+    public void handshake() {}
 
     /**
     * Reads a sequence of bytes from this channel into the given buffer.
@@ -126,7 +120,7 @@ public long read(ByteBuffer[] dsts) throws IOException {
      * @param dsts - The buffers into which bytes are to be transferred
      * @param offset - The offset within the buffer array of the first buffer into which bytes are to be transferred; must be non-negative and no larger than dsts.length.
      * @param length - The maximum number of buffers to be accessed; must be non-negative and no larger than dsts.length - offset
-     * @returns The number of bytes read, possibly zero, or -1 if the channel has reached end-of-stream.
+     * @return The number of bytes read, possibly zero, or -1 if the channel has reached end-of-stream.
      * @throws IOException if some other I/O error occurs
      */
     @Override
@@ -138,7 +132,7 @@ public long read(ByteBuffer[] dsts, int offset, int length) throws IOException {
     * Writes a sequence of bytes to this channel from the given buffer.
     *
     * @param src The buffer from which bytes are to be retrieved
-    * @returns The number of bytes read, possibly zero, or -1 if the channel has reached end-of-stream
+    * @return The number of bytes read, possibly zero, or -1 if the channel has reached end-of-stream
     * @throws IOException If some other I/O error occurs
     */
     @Override
@@ -150,7 +144,7 @@ public int write(ByteBuffer src) throws IOException {
     * Writes a sequence of bytes to this channel from the given buffer.
     *
     * @param srcs The buffer from which bytes are to be retrieved
-    * @returns The number of bytes read, possibly zero, or -1 if the channel has reached end-of-stream
+    * @return The number of bytes read, possibly zero, or -1 if the channel has reached end-of-stream
     * @throws IOException If some other I/O error occurs
     */
     @Override
@@ -185,13 +179,12 @@ public boolean hasPendingWrites() {
      * Returns ANONYMOUS as Principal.
      */
     @Override
-    public Principal peerPrincipal() throws IOException {
+    public Principal peerPrincipal() {
         return principal;
     }
 
     /**
      * Adds the interestOps to selectionKey.
-     * @param ops
      */
     @Override
     public void addInterestOps(int ops) {
@@ -201,7 +194,6 @@ public void addInterestOps(int ops) {
 
     /**
      * Removes the interestOps from selectionKey.
-     * @param ops
      */
     @Override
     public void removeInterestOps(int ops) {
diff --git a/clients/src/main/java/org/apache/kafka/common/network/ReauthenticationContext.java b/clients/src/main/java/org/apache/kafka/common/network/ReauthenticationContext.java
new file mode 100644
index 0000000000000..37e46cbb6ff9c
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/network/ReauthenticationContext.java
@@ -0,0 +1,94 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.network;
+
+import java.util.Objects;
+
+/**
+ * Defines the context in which an {@link Authenticator} is to be created during
+ * a re-authentication.
+ */
+public class ReauthenticationContext {
+    private final NetworkReceive networkReceive;
+    private final Authenticator previousAuthenticator;
+    private final long reauthenticationBeginNanos;
+
+    /**
+     * Constructor
+     * 
+     * @param previousAuthenticator
+     *            the mandatory {@link Authenticator} that was previously used to
+     *            authenticate the channel
+     * @param networkReceive
+     *            the applicable {@link NetworkReceive} instance, if any. For the
+     *            client side this may be a response that has been partially read, a
+     *            non-null instance that has had no data read into it yet, or null;
+     *            if it is non-null then this is the instance that data should
+     *            initially be read into during re-authentication. For the server
+     *            side this is mandatory and it must contain the
+     *            {@code SaslHandshakeRequest} that has been received on the server
+     *            and that initiates re-authentication.
+     * 
+     * @param nowNanos
+     *            the current time. The value is in nanoseconds as per
+     *            {@code System.nanoTime()} and is therefore only useful when
+     *            compared to such a value -- it's absolute value is meaningless.
+     *            This defines the moment when re-authentication begins.
+     */
+    public ReauthenticationContext(Authenticator previousAuthenticator, NetworkReceive networkReceive, long nowNanos) {
+        this.previousAuthenticator = Objects.requireNonNull(previousAuthenticator);
+        this.networkReceive = networkReceive;
+        this.reauthenticationBeginNanos = nowNanos;
+    }
+
+    /**
+     * Return the applicable {@link NetworkReceive} instance, if any. For the client
+     * side this may be a response that has been partially read, a non-null instance
+     * that has had no data read into it yet, or null; if it is non-null then this
+     * is the instance that data should initially be read into during
+     * re-authentication. For the server side this is mandatory and it must contain
+     * the {@code SaslHandshakeRequest} that has been received on the server and
+     * that initiates re-authentication.
+     * 
+     * @return the applicable {@link NetworkReceive} instance, if any
+     */
+    public NetworkReceive networkReceive() {
+        return networkReceive;
+    }
+
+    /**
+     * Return the always non-null {@link Authenticator} that was previously used to
+     * authenticate the channel
+     * 
+     * @return the always non-null {@link Authenticator} that was previously used to
+     *         authenticate the channel
+     */
+    public Authenticator previousAuthenticator() {
+        return previousAuthenticator;
+    }
+
+    /**
+     * Return the time when re-authentication began. The value is in nanoseconds as
+     * per {@code System.nanoTime()} and is therefore only useful when compared to
+     * such a value -- it's absolute value is meaningless.
+     * 
+     * @return the time when re-authentication began
+     */
+    public long reauthenticationBeginNanos() {
+        return reauthenticationBeginNanos;
+    }
+}
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 72a2bc1d91890..f01c4ef113fba 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
@@ -18,114 +18,211 @@
 
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.config.SaslConfigs;
+import org.apache.kafka.common.config.SslConfigs;
 import org.apache.kafka.common.config.internals.BrokerSecurityConfigs;
 import org.apache.kafka.common.memory.MemoryPool;
-import org.apache.kafka.common.security.auth.SecurityProtocol;
 import org.apache.kafka.common.security.JaasContext;
+import org.apache.kafka.common.security.auth.AuthenticateCallbackHandler;
+import org.apache.kafka.common.security.auth.Login;
+import org.apache.kafka.common.security.auth.SecurityProtocol;
 import org.apache.kafka.common.security.authenticator.CredentialCache;
+import org.apache.kafka.common.security.authenticator.DefaultLogin;
 import org.apache.kafka.common.security.authenticator.LoginManager;
 import org.apache.kafka.common.security.authenticator.SaslClientAuthenticator;
+import org.apache.kafka.common.security.authenticator.SaslClientCallbackHandler;
 import org.apache.kafka.common.security.authenticator.SaslServerAuthenticator;
+import org.apache.kafka.common.security.authenticator.SaslServerCallbackHandler;
+import org.apache.kafka.common.security.kerberos.KerberosClientCallbackHandler;
+import org.apache.kafka.common.security.kerberos.KerberosLogin;
+import org.apache.kafka.common.security.kerberos.KerberosName;
 import org.apache.kafka.common.security.kerberos.KerberosShortNamer;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerLoginModule;
+import org.apache.kafka.common.security.oauthbearer.internals.OAuthBearerRefreshingLogin;
+import org.apache.kafka.common.security.oauthbearer.internals.OAuthBearerSaslClientCallbackHandler;
+import org.apache.kafka.common.security.oauthbearer.internals.unsecured.OAuthBearerUnsecuredValidatorCallbackHandler;
+import org.apache.kafka.common.security.plain.internals.PlainSaslServer;
+import org.apache.kafka.common.security.plain.internals.PlainServerCallbackHandler;
+import org.apache.kafka.common.security.scram.ScramCredential;
+import org.apache.kafka.common.security.scram.internals.ScramMechanism;
+import org.apache.kafka.common.security.scram.internals.ScramServerCallbackHandler;
 import org.apache.kafka.common.security.ssl.SslFactory;
-import org.apache.kafka.common.utils.Java;
+import org.apache.kafka.common.security.token.delegation.internals.DelegationTokenCache;
+import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.common.utils.Utils;
+import org.ietf.jgss.GSSContext;
+import org.ietf.jgss.GSSCredential;
+import org.ietf.jgss.GSSException;
+import org.ietf.jgss.GSSManager;
+import org.ietf.jgss.GSSName;
+import org.ietf.jgss.Oid;
 import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
+import javax.security.auth.Subject;
+import javax.security.auth.kerberos.KerberosPrincipal;
 import java.io.IOException;
-import java.lang.reflect.InvocationTargetException;
-import java.lang.reflect.Method;
 import java.net.Socket;
 import java.nio.channels.SelectionKey;
 import java.nio.channels.SocketChannel;
+import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
+import java.util.function.Supplier;
 
-import javax.security.auth.Subject;
-
-public class SaslChannelBuilder implements ChannelBuilder {
-    private static final Logger log = LoggerFactory.getLogger(SaslChannelBuilder.class);
+public class SaslChannelBuilder implements ChannelBuilder, ListenerReconfigurable {
+    static final String GSS_NATIVE_PROP = "sun.security.jgss.native";
 
     private final SecurityProtocol securityProtocol;
     private final ListenerName listenerName;
+    private final boolean isInterBrokerListener;
     private final String clientSaslMechanism;
     private final Mode mode;
-    private final JaasContext jaasContext;
+    private final Map jaasContexts;
     private final boolean handshakeRequestEnable;
     private final CredentialCache credentialCache;
+    private final DelegationTokenCache tokenCache;
+    private final Map loginManagers;
+    private final Map subjects;
 
-    private LoginManager loginManager;
     private SslFactory sslFactory;
     private Map configs;
     private KerberosShortNamer kerberosShortNamer;
+    private Map saslCallbackHandlers;
+    private Map connectionsMaxReauthMsByMechanism;
+    private final Time time;
+    private final LogContext logContext;
+    private final Logger log;
 
     public SaslChannelBuilder(Mode mode,
-                              JaasContext jaasContext,
+                              Map jaasContexts,
                               SecurityProtocol securityProtocol,
                               ListenerName listenerName,
+                              boolean isInterBrokerListener,
                               String clientSaslMechanism,
                               boolean handshakeRequestEnable,
-                              CredentialCache credentialCache) {
+                              CredentialCache credentialCache,
+                              DelegationTokenCache tokenCache,
+                              Time time,
+                              LogContext logContext) {
         this.mode = mode;
-        this.jaasContext = jaasContext;
+        this.jaasContexts = jaasContexts;
+        this.loginManagers = new HashMap<>(jaasContexts.size());
+        this.subjects = new HashMap<>(jaasContexts.size());
         this.securityProtocol = securityProtocol;
         this.listenerName = listenerName;
+        this.isInterBrokerListener = isInterBrokerListener;
         this.handshakeRequestEnable = handshakeRequestEnable;
         this.clientSaslMechanism = clientSaslMechanism;
         this.credentialCache = credentialCache;
+        this.tokenCache = tokenCache;
+        this.saslCallbackHandlers = new HashMap<>();
+        this.connectionsMaxReauthMsByMechanism = new HashMap<>();
+        this.time = time;
+        this.logContext = logContext;
+        this.log = logContext.logger(getClass());
     }
 
+    @SuppressWarnings("unchecked")
     @Override
     public void configure(Map configs) throws KafkaException {
         try {
             this.configs = configs;
-            boolean hasKerberos;
             if (mode == Mode.SERVER) {
-                @SuppressWarnings("unchecked")
-                List enabledMechanisms = (List) this.configs.get(BrokerSecurityConfigs.SASL_ENABLED_MECHANISMS_CONFIG);
-                hasKerberos = enabledMechanisms == null || enabledMechanisms.contains(SaslConfigs.GSSAPI_MECHANISM);
-            } else {
-                hasKerberos = clientSaslMechanism.equals(SaslConfigs.GSSAPI_MECHANISM);
+                createServerCallbackHandlers(configs);
+                createConnectionsMaxReauthMsMap(configs);
+            } else
+                createClientCallbackHandler(configs);
+            for (Map.Entry entry : saslCallbackHandlers.entrySet()) {
+                String mechanism = entry.getKey();
+                entry.getValue().configure(configs, mechanism, jaasContexts.get(mechanism).configurationEntries());
             }
 
-            if (hasKerberos) {
+            Class defaultLoginClass = defaultLoginClass();
+            if (mode == Mode.SERVER && jaasContexts.containsKey(SaslConfigs.GSSAPI_MECHANISM)) {
                 String defaultRealm;
                 try {
                     defaultRealm = defaultKerberosRealm();
                 } catch (Exception ke) {
                     defaultRealm = "";
                 }
-                @SuppressWarnings("unchecked")
                 List principalToLocalRules = (List) configs.get(BrokerSecurityConfigs.SASL_KERBEROS_PRINCIPAL_TO_LOCAL_RULES_CONFIG);
                 if (principalToLocalRules != null)
                     kerberosShortNamer = KerberosShortNamer.fromUnparsedRules(defaultRealm, principalToLocalRules);
             }
-            this.loginManager = LoginManager.acquireLoginManager(jaasContext, hasKerberos, configs);
-
+            for (Map.Entry entry : jaasContexts.entrySet()) {
+                String mechanism = entry.getKey();
+                // With static JAAS configuration, use KerberosLogin if Kerberos is enabled. With dynamic JAAS configuration,
+                // use KerberosLogin only for the LoginContext corresponding to GSSAPI
+                LoginManager loginManager = LoginManager.acquireLoginManager(entry.getValue(), mechanism, defaultLoginClass, configs);
+                loginManagers.put(mechanism, loginManager);
+                Subject subject = loginManager.subject();
+                subjects.put(mechanism, subject);
+                if (mode == Mode.SERVER && mechanism.equals(SaslConfigs.GSSAPI_MECHANISM))
+                    maybeAddNativeGssapiCredentials(subject);
+            }
             if (this.securityProtocol == SecurityProtocol.SASL_SSL) {
                 // Disable SSL client authentication as we are using SASL authentication
-                this.sslFactory = new SslFactory(mode, "none");
+                this.sslFactory = new SslFactory(mode, "none", isInterBrokerListener);
                 this.sslFactory.configure(configs);
             }
-        } catch (Exception e) {
+        } catch (Throwable e) {
             close();
             throw new KafkaException(e);
         }
     }
 
     @Override
-    public KafkaChannel buildChannel(String id, SelectionKey key, int maxReceiveSize, MemoryPool memoryPool) throws KafkaException {
+    public Set reconfigurableConfigs() {
+        return securityProtocol == SecurityProtocol.SASL_SSL ? SslConfigs.RECONFIGURABLE_CONFIGS : Collections.emptySet();
+    }
+
+    @Override
+    public void validateReconfiguration(Map configs) {
+        if (this.securityProtocol == SecurityProtocol.SASL_SSL)
+            sslFactory.validateReconfiguration(configs);
+    }
+
+    @Override
+    public void reconfigure(Map configs) {
+        if (this.securityProtocol == SecurityProtocol.SASL_SSL)
+            sslFactory.reconfigure(configs);
+    }
+
+    @Override
+    public ListenerName listenerName() {
+        return listenerName;
+    }
+
+    @Override
+    public KafkaChannel buildChannel(String id, SelectionKey key, int maxReceiveSize,
+                                     MemoryPool memoryPool, ChannelMetadataRegistry metadataRegistry) throws KafkaException {
         try {
             SocketChannel socketChannel = (SocketChannel) key.channel();
             Socket socket = socketChannel.socket();
-            TransportLayer transportLayer = buildTransportLayer(id, key, socketChannel);
-            Authenticator authenticator;
-            if (mode == Mode.SERVER)
-                authenticator = buildServerAuthenticator(configs, id, transportLayer, loginManager.subject());
-            else
-                authenticator = buildClientAuthenticator(configs, id, socket.getInetAddress().getHostName(),
-                        loginManager.serviceName(), transportLayer, loginManager.subject());
-            return new KafkaChannel(id, transportLayer, authenticator, maxReceiveSize, memoryPool != null ? memoryPool : MemoryPool.NONE);
+            TransportLayer transportLayer = buildTransportLayer(id, key, socketChannel, metadataRegistry);
+            Supplier authenticatorCreator;
+            if (mode == Mode.SERVER) {
+                authenticatorCreator = () -> buildServerAuthenticator(configs,
+                        Collections.unmodifiableMap(saslCallbackHandlers),
+                        id,
+                        transportLayer,
+                        Collections.unmodifiableMap(subjects),
+                        Collections.unmodifiableMap(connectionsMaxReauthMsByMechanism),
+                        metadataRegistry);
+            } else {
+                LoginManager loginManager = loginManagers.get(clientSaslMechanism);
+                authenticatorCreator = () -> buildClientAuthenticator(configs,
+                        saslCallbackHandlers.get(clientSaslMechanism),
+                        id,
+                        socket.getInetAddress().getHostName(),
+                        loginManager.serviceName(),
+                        transportLayer,
+                        subjects.get(clientSaslMechanism));
+            }
+            return new KafkaChannel(id, transportLayer, authenticatorCreator, maxReceiveSize,
+                memoryPool != null ? memoryPool : MemoryPool.NONE, metadataRegistry);
         } catch (Exception e) {
             log.info("Failed to create channel due to ", e);
             throw new KafkaException(e);
@@ -134,58 +231,164 @@ public KafkaChannel buildChannel(String id, SelectionKey key, int maxReceiveSize
 
     @Override
     public void close()  {
-        if (loginManager != null) {
+        for (LoginManager loginManager : loginManagers.values())
             loginManager.release();
-            loginManager = null;
-        }
+        loginManagers.clear();
+        for (AuthenticateCallbackHandler handler : saslCallbackHandlers.values())
+            handler.close();
+        if (sslFactory != null) sslFactory.close();
     }
 
-    private TransportLayer buildTransportLayer(String id, SelectionKey key, SocketChannel socketChannel) throws IOException {
+    // Visible to override for testing
+    protected TransportLayer buildTransportLayer(String id, SelectionKey key, SocketChannel socketChannel,
+                                                 ChannelMetadataRegistry metadataRegistry) throws IOException {
         if (this.securityProtocol == SecurityProtocol.SASL_SSL) {
             return SslTransportLayer.create(id, key,
-                sslFactory.createSslEngine(socketChannel.socket().getInetAddress().getHostName(), socketChannel.socket().getPort()));
+                sslFactory.createSslEngine(socketChannel.socket().getInetAddress().getHostName(),
+                    socketChannel.socket().getPort()),
+                metadataRegistry);
         } else {
             return new PlaintextTransportLayer(key);
         }
     }
 
     // Visible to override for testing
-    protected SaslServerAuthenticator buildServerAuthenticator(Map configs, String id,
-            TransportLayer transportLayer, Subject subject) throws IOException {
-        return new SaslServerAuthenticator(configs, id, jaasContext, subject,
-                kerberosShortNamer, credentialCache, listenerName, securityProtocol, transportLayer);
+    protected SaslServerAuthenticator buildServerAuthenticator(Map configs,
+                                                               Map callbackHandlers,
+                                                               String id,
+                                                               TransportLayer transportLayer,
+                                                               Map subjects,
+                                                               Map connectionsMaxReauthMsByMechanism,
+                                                               ChannelMetadataRegistry metadataRegistry) {
+        return new SaslServerAuthenticator(configs, callbackHandlers, id, subjects,
+                kerberosShortNamer, listenerName, securityProtocol, transportLayer,
+                connectionsMaxReauthMsByMechanism, metadataRegistry, time);
     }
 
     // Visible to override for testing
-    protected SaslClientAuthenticator buildClientAuthenticator(Map configs, String id,
-            String serverHost, String servicePrincipal, TransportLayer transportLayer, Subject subject) throws IOException {
-        return new SaslClientAuthenticator(configs, id, subject, servicePrincipal,
-                serverHost, clientSaslMechanism, handshakeRequestEnable, transportLayer);
+    protected SaslClientAuthenticator buildClientAuthenticator(Map configs,
+                                                               AuthenticateCallbackHandler callbackHandler,
+                                                               String id,
+                                                               String serverHost,
+                                                               String servicePrincipal,
+                                                               TransportLayer transportLayer, Subject subject) {
+        return new SaslClientAuthenticator(configs, callbackHandler, id, subject, servicePrincipal,
+                serverHost, clientSaslMechanism, handshakeRequestEnable, transportLayer, time, logContext);
     }
 
     // Package private for testing
-    LoginManager loginManager() {
-        return loginManager;
+    Map loginManagers() {
+        return loginManagers;
     }
 
-    private static String defaultKerberosRealm() throws ClassNotFoundException, NoSuchMethodException,
-            IllegalArgumentException, IllegalAccessException, InvocationTargetException {
+    private static String defaultKerberosRealm() {
+        // see https://issues.apache.org/jira/browse/HADOOP-10848 for details
+        return new KerberosPrincipal("tmp", 1).getRealm();
+    }
 
-        //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
+    private void createClientCallbackHandler(Map configs) {
+        @SuppressWarnings("unchecked")
+        Class clazz = (Class) configs.get(SaslConfigs.SASL_CLIENT_CALLBACK_HANDLER_CLASS);
+        if (clazz == null)
+            clazz = clientCallbackHandlerClass();
+        AuthenticateCallbackHandler callbackHandler = Utils.newInstance(clazz);
+        saslCallbackHandlers.put(clientSaslMechanism, callbackHandler);
+    }
 
-        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");
+    private void createServerCallbackHandlers(Map configs) {
+        for (String mechanism : jaasContexts.keySet()) {
+            AuthenticateCallbackHandler callbackHandler;
+            String prefix = ListenerName.saslMechanismPrefix(mechanism);
+            @SuppressWarnings("unchecked")
+            Class clazz =
+                    (Class) configs.get(prefix + BrokerSecurityConfigs.SASL_SERVER_CALLBACK_HANDLER_CLASS);
+            if (clazz != null)
+                callbackHandler = Utils.newInstance(clazz);
+            else if (mechanism.equals(PlainSaslServer.PLAIN_MECHANISM))
+                callbackHandler = new PlainServerCallbackHandler();
+            else if (ScramMechanism.isScram(mechanism))
+                callbackHandler = new ScramServerCallbackHandler(credentialCache.cache(mechanism, ScramCredential.class), tokenCache);
+            else if (mechanism.equals(OAuthBearerLoginModule.OAUTHBEARER_MECHANISM))
+                callbackHandler = new OAuthBearerUnsecuredValidatorCallbackHandler();
+            else
+                callbackHandler = new SaslServerCallbackHandler();
+            saslCallbackHandlers.put(mechanism, callbackHandler);
         }
-        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 void createConnectionsMaxReauthMsMap(Map configs) {
+        for (String mechanism : jaasContexts.keySet()) {
+            String prefix = ListenerName.saslMechanismPrefix(mechanism);
+            Long connectionsMaxReauthMs = (Long) configs.get(prefix + BrokerSecurityConfigs.CONNECTIONS_MAX_REAUTH_MS);
+            if (connectionsMaxReauthMs == null)
+                connectionsMaxReauthMs = (Long) configs.get(BrokerSecurityConfigs.CONNECTIONS_MAX_REAUTH_MS);
+            if (connectionsMaxReauthMs != null)
+                connectionsMaxReauthMsByMechanism.put(mechanism, connectionsMaxReauthMs);
+        }
+    }
+
+    protected Class defaultLoginClass() {
+        if (jaasContexts.containsKey(SaslConfigs.GSSAPI_MECHANISM))
+            return KerberosLogin.class;
+        if (OAuthBearerLoginModule.OAUTHBEARER_MECHANISM.equals(clientSaslMechanism))
+            return OAuthBearerRefreshingLogin.class;
+        return DefaultLogin.class;
+    }
+
+    private Class clientCallbackHandlerClass() {
+        switch (clientSaslMechanism) {
+            case SaslConfigs.GSSAPI_MECHANISM:
+                return KerberosClientCallbackHandler.class;
+            case OAuthBearerLoginModule.OAUTHBEARER_MECHANISM:
+                return OAuthBearerSaslClientCallbackHandler.class;
+            default:
+                return SaslClientCallbackHandler.class;
+        }
+    }
+
+    // As described in http://docs.oracle.com/javase/8/docs/technotes/guides/security/jgss/jgss-features.html:
+    // "To enable Java GSS to delegate to the native GSS library and its list of native mechanisms,
+    // set the system property "sun.security.jgss.native" to true"
+    // "In addition, when performing operations as a particular Subject, for example, Subject.doAs(...)
+    // or Subject.doAsPrivileged(...), the to-be-used GSSCredential should be added to Subject's
+    // private credential set. Otherwise, the GSS operations will fail since no credential is found."
+    private void maybeAddNativeGssapiCredentials(Subject subject) {
+        boolean usingNativeJgss = Boolean.getBoolean(GSS_NATIVE_PROP);
+        if (usingNativeJgss && subject.getPrivateCredentials(GSSCredential.class).isEmpty()) {
+
+            final String servicePrincipal = SaslClientAuthenticator.firstPrincipal(subject);
+            KerberosName kerberosName;
+            try {
+                kerberosName = KerberosName.parse(servicePrincipal);
+            } catch (IllegalArgumentException e) {
+                throw new KafkaException("Principal has name with unexpected format " + servicePrincipal);
+            }
+            final String servicePrincipalName = kerberosName.serviceName();
+            final String serviceHostname = kerberosName.hostName();
+
+            try {
+                GSSManager manager = gssManager();
+                // This Oid is used to represent the Kerberos version 5 GSS-API mechanism. It is defined in
+                // RFC 1964.
+                Oid krb5Mechanism = new Oid("1.2.840.113554.1.2.2");
+                GSSName gssName = manager.createName(servicePrincipalName + "@" + serviceHostname, GSSName.NT_HOSTBASED_SERVICE);
+                GSSCredential cred = manager.createCredential(gssName,
+                        GSSContext.INDEFINITE_LIFETIME, krb5Mechanism, GSSCredential.ACCEPT_ONLY);
+                subject.getPrivateCredentials().add(cred);
+                log.info("Configured native GSSAPI private credentials for {}@{}", serviceHostname, serviceHostname);
+            } catch (GSSException ex) {
+                log.warn("Cannot add private credential to subject; clients authentication may fail", ex);
+            }
+        }
+    }
+
+    // Visibility to override for testing
+    protected GSSManager gssManager() {
+        return GSSManager.getInstance();
+    }
+
+    // Visibility for testing
+    protected Subject subject(String saslMechanism) {
+        return subjects.get(saslMechanism);
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/network/Selectable.java b/clients/src/main/java/org/apache/kafka/common/network/Selectable.java
index efb603c8dc361..afdd42e4a9c29 100644
--- a/clients/src/main/java/org/apache/kafka/common/network/Selectable.java
+++ b/clients/src/main/java/org/apache/kafka/common/network/Selectable.java
@@ -19,6 +19,7 @@
 
 import java.io.IOException;
 import java.net.InetSocketAddress;
+import java.util.Collection;
 import java.util.List;
 import java.util.Map;
 
@@ -30,7 +31,7 @@ public interface Selectable {
     /**
      * See {@link #connect(String, InetSocketAddress, int, int) connect()}
      */
-    public static final int USE_DEFAULT_BUFFER_SIZE = -1;
+    int USE_DEFAULT_BUFFER_SIZE = -1;
 
     /**
      * Begin establishing a socket connection to the given address identified by the given address
@@ -40,83 +41,83 @@ public interface Selectable {
      * @param receiveBufferSize The receive buffer for the socket
      * @throws IOException If we cannot begin connecting
      */
-    public void connect(String id, InetSocketAddress address, int sendBufferSize, int receiveBufferSize) throws IOException;
+    void connect(String id, InetSocketAddress address, int sendBufferSize, int receiveBufferSize) throws IOException;
 
     /**
      * Wakeup this selector if it is blocked on I/O
      */
-    public void wakeup();
+    void wakeup();
 
     /**
      * Close this selector
      */
-    public void close();
+    void close();
 
     /**
      * Close the connection identified by the given id
      */
-    public void close(String id);
+    void close(String id);
 
     /**
      * Queue the given request for sending in the subsequent {@link #poll(long) poll()} calls
      * @param send The request to send
      */
-    public void send(Send send);
+    void send(NetworkSend send);
 
     /**
      * Do I/O. Reads, writes, connection establishment, etc.
      * @param timeout The amount of time to block if there is nothing to do
      * @throws IOException
      */
-    public void poll(long timeout) throws IOException;
+    void poll(long timeout) throws IOException;
 
     /**
      * The list of sends that completed on the last {@link #poll(long) poll()} call.
      */
-    public List completedSends();
+    List completedSends();
 
     /**
-     * The list of receives that completed on the last {@link #poll(long) poll()} call.
+     * The collection of receives that completed on the last {@link #poll(long) poll()} call.
      */
-    public List completedReceives();
+    Collection completedReceives();
 
     /**
      * The connections that finished disconnecting on the last {@link #poll(long) poll()}
      * call. Channel state indicates the local channel state at the time of disconnection.
      */
-    public Map disconnected();
+    Map disconnected();
 
     /**
      * The list of connections that completed their connection on the last {@link #poll(long) poll()}
      * call.
      */
-    public List connected();
+    List connected();
 
     /**
      * Disable reads from the given connection
      * @param id The id for the connection
      */
-    public void mute(String id);
+    void mute(String id);
 
     /**
      * Re-enable reads from the given connection
      * @param id The id for the connection
      */
-    public void unmute(String id);
+    void unmute(String id);
 
     /**
      * Disable reads from all connections
      */
-    public void muteAll();
+    void muteAll();
 
     /**
      * Re-enable reads from all connections
      */
-    public void unmuteAll();
+    void unmuteAll();
 
     /**
      * returns true  if a channel is ready
      * @param id The id for the connection
      */
-    public boolean isChannelReady(String id);
+    boolean isChannelReady(String id);
 }
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 09df7b534009e..1e710fdf6a3e3 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
@@ -16,6 +16,24 @@
  */
 package org.apache.kafka.common.network;
 
+import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.MetricName;
+import org.apache.kafka.common.errors.AuthenticationException;
+import org.apache.kafka.common.memory.MemoryPool;
+import org.apache.kafka.common.metrics.Metrics;
+import org.apache.kafka.common.metrics.Sensor;
+import org.apache.kafka.common.metrics.internals.IntGaugeSuite;
+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.Meter;
+import org.apache.kafka.common.metrics.stats.SampledStat;
+import org.apache.kafka.common.metrics.stats.WindowedCount;
+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.net.InetSocketAddress;
 import java.net.Socket;
@@ -23,36 +41,19 @@
 import java.nio.channels.SelectionKey;
 import java.nio.channels.SocketChannel;
 import java.nio.channels.UnresolvedAddressException;
-import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
-import java.util.Deque;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Iterator;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
-
-import org.apache.kafka.common.KafkaException;
-import org.apache.kafka.common.memory.MemoryPool;
-import org.apache.kafka.common.metrics.Measurable;
-import org.apache.kafka.common.metrics.MetricConfig;
-import org.apache.kafka.common.MetricName;
-import org.apache.kafka.common.errors.AuthenticationException;
-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.Meter;
-import org.apache.kafka.common.metrics.stats.SampledStat;
-import org.apache.kafka.common.metrics.stats.Count;
-import org.apache.kafka.common.metrics.stats.Max;
-import org.apache.kafka.common.utils.LogContext;
-import org.apache.kafka.common.utils.Time;
-import org.slf4j.Logger;
+import java.util.concurrent.atomic.AtomicReference;
 
 /**
  * A nioSelector interface for doing non-blocking multi-connection network I/O.
@@ -86,15 +87,27 @@
 public class Selector implements Selectable, AutoCloseable {
 
     public static final long NO_IDLE_TIMEOUT_MS = -1;
+    public static final int NO_FAILED_AUTHENTICATION_DELAY = 0;
+
+    private enum CloseMode {
+        GRACEFUL(true),            // process outstanding buffered receives, notify disconnect
+        NOTIFY_ONLY(true),         // discard any outstanding receives, notify disconnect
+        DISCARD_NO_NOTIFY(false);  // discard any outstanding receives, no disconnect notification
+
+        boolean notifyDisconnect;
+
+        CloseMode(boolean notifyDisconnect) {
+            this.notifyDisconnect = notifyDisconnect;
+        }
+    }
 
     private final Logger log;
     private final java.nio.channels.Selector nioSelector;
     private final Map channels;
     private final Set explicitlyMutedChannels;
     private boolean outOfMemory;
-    private final List completedSends;
-    private final List completedReceives;
-    private final Map> stagedReceives;
+    private final List completedSends;
+    private final LinkedHashMap completedReceives;
     private final Set immediatelyConnectedKeys;
     private final Map closingChannels;
     private Set keysWithBufferedRead;
@@ -107,8 +120,11 @@ public class Selector implements Selectable, AutoCloseable {
     private final int maxReceiveSize;
     private final boolean recordTimePerConnection;
     private final IdleExpiryManager idleExpiryManager;
+    private final LinkedHashMap delayedClosingChannels;
     private final MemoryPool memoryPool;
     private final long lowMemThreshold;
+    private final int failedAuthenticationDelayMs;
+
     //indicates if the previous call to poll was able to make progress in reading already-buffered data.
     //this is used to prevent tight loops when memory is not available to read any more data
     private boolean madeReadProgressLastPoll = true;
@@ -117,6 +133,8 @@ public class Selector implements Selectable, AutoCloseable {
      * Create a new nioSelector
      * @param maxReceiveSize Max size in bytes of a single network receive (use {@link NetworkReceive#UNLIMITED} for no limit)
      * @param connectionMaxIdleMs Max idle connection time (use {@link #NO_IDLE_TIMEOUT_MS} to disable idle timeout)
+     * @param failedAuthenticationDelayMs Minimum time by which failed authentication response and channel close should be delayed by.
+     *                                    Use {@link #NO_FAILED_AUTHENTICATION_DELAY} to disable this delay.
      * @param metrics Registry for Selector metrics
      * @param time Time implementation
      * @param metricGrpPrefix Prefix for the group of metrics registered by Selector
@@ -127,6 +145,7 @@ public class Selector implements Selectable, AutoCloseable {
      */
     public Selector(int maxReceiveSize,
             long connectionMaxIdleMs,
+            int failedAuthenticationDelayMs,
             Metrics metrics,
             Time time,
             String metricGrpPrefix,
@@ -147,37 +166,70 @@ public Selector(int maxReceiveSize,
         this.explicitlyMutedChannels = new HashSet<>();
         this.outOfMemory = false;
         this.completedSends = new ArrayList<>();
-        this.completedReceives = new ArrayList<>();
-        this.stagedReceives = new HashMap<>();
+        this.completedReceives = new LinkedHashMap<>();
         this.immediatelyConnectedKeys = new HashSet<>();
         this.closingChannels = new HashMap<>();
         this.keysWithBufferedRead = new HashSet<>();
         this.connected = new ArrayList<>();
         this.disconnected = new HashMap<>();
         this.failedSends = new ArrayList<>();
+        this.log = logContext.logger(Selector.class);
         this.sensors = new SelectorMetrics(metrics, metricGrpPrefix, metricTags, metricsPerConnection);
         this.channelBuilder = channelBuilder;
         this.recordTimePerConnection = recordTimePerConnection;
         this.idleExpiryManager = connectionMaxIdleMs < 0 ? null : new IdleExpiryManager(time, connectionMaxIdleMs);
         this.memoryPool = memoryPool;
         this.lowMemThreshold = (long) (0.1 * this.memoryPool.size());
-        this.log = logContext.logger(Selector.class);
+        this.failedAuthenticationDelayMs = failedAuthenticationDelayMs;
+        this.delayedClosingChannels = (failedAuthenticationDelayMs > NO_FAILED_AUTHENTICATION_DELAY) ? new LinkedHashMap() : null;
     }
 
     public Selector(int maxReceiveSize,
-            long connectionMaxIdleMs,
-            Metrics metrics,
-            Time time,
-            String metricGrpPrefix,
-            Map metricTags,
-            boolean metricsPerConnection,
-            ChannelBuilder channelBuilder,
-            LogContext logContext) {
-        this(maxReceiveSize, connectionMaxIdleMs, metrics, time, metricGrpPrefix, metricTags, metricsPerConnection, false, channelBuilder, MemoryPool.NONE, logContext);
+                    long connectionMaxIdleMs,
+                    Metrics metrics,
+                    Time time,
+                    String metricGrpPrefix,
+                    Map metricTags,
+                    boolean metricsPerConnection,
+                    boolean recordTimePerConnection,
+                    ChannelBuilder channelBuilder,
+                    MemoryPool memoryPool,
+                    LogContext logContext) {
+        this(maxReceiveSize, connectionMaxIdleMs, NO_FAILED_AUTHENTICATION_DELAY, metrics, time, metricGrpPrefix, metricTags,
+                metricsPerConnection, recordTimePerConnection, channelBuilder, memoryPool, logContext);
+    }
+
+    public Selector(int maxReceiveSize,
+                    long connectionMaxIdleMs,
+                    int failedAuthenticationDelayMs,
+                    Metrics metrics,
+                    Time time,
+                    String metricGrpPrefix,
+                    Map metricTags,
+                    boolean metricsPerConnection,
+                    ChannelBuilder channelBuilder,
+                    LogContext logContext) {
+        this(maxReceiveSize, connectionMaxIdleMs, failedAuthenticationDelayMs, metrics, time, metricGrpPrefix, metricTags, metricsPerConnection, false, channelBuilder, MemoryPool.NONE, logContext);
+    }
+
+    public Selector(int maxReceiveSize,
+                    long connectionMaxIdleMs,
+                    Metrics metrics,
+                    Time time,
+                    String metricGrpPrefix,
+                    Map metricTags,
+                    boolean metricsPerConnection,
+                    ChannelBuilder channelBuilder,
+                    LogContext logContext) {
+        this(maxReceiveSize, connectionMaxIdleMs, NO_FAILED_AUTHENTICATION_DELAY, metrics, time, metricGrpPrefix, metricTags, metricsPerConnection, channelBuilder, logContext);
     }
 
     public Selector(long connectionMaxIdleMS, Metrics metrics, Time time, String metricGrpPrefix, ChannelBuilder channelBuilder, LogContext logContext) {
-        this(NetworkReceive.UNLIMITED, connectionMaxIdleMS, metrics, time, metricGrpPrefix, Collections.emptyMap(), true, channelBuilder, logContext);
+        this(NetworkReceive.UNLIMITED, connectionMaxIdleMS, metrics, time, metricGrpPrefix, Collections.emptyMap(), true, channelBuilder, logContext);
+    }
+
+    public Selector(long connectionMaxIdleMS, int failedAuthenticationDelayMs, Metrics metrics, Time time, String metricGrpPrefix, ChannelBuilder channelBuilder, LogContext logContext) {
+        this(NetworkReceive.UNLIMITED, connectionMaxIdleMS, failedAuthenticationDelayMs, metrics, time, metricGrpPrefix, Collections.emptyMap(), true, channelBuilder, logContext);
     }
 
     /**
@@ -195,12 +247,41 @@ public Selector(long connectionMaxIdleMS, Metrics metrics, Time time, String met
      */
     @Override
     public void connect(String id, InetSocketAddress address, int sendBufferSize, int receiveBufferSize) throws IOException {
-        if (this.channels.containsKey(id))
-            throw new IllegalStateException("There is already a connection for id " + id);
-        if (this.closingChannels.containsKey(id))
-            throw new IllegalStateException("There is already a connection for id " + id + " that is still being closed");
-
+        ensureNotRegistered(id);
         SocketChannel socketChannel = SocketChannel.open();
+        SelectionKey key = null;
+        try {
+            configureSocketChannel(socketChannel, sendBufferSize, receiveBufferSize);
+            boolean connected = doConnect(socketChannel, address);
+            key = registerChannel(id, socketChannel, SelectionKey.OP_CONNECT);
+
+            if (connected) {
+                // OP_CONNECT won't trigger for immediately connected channels
+                log.debug("Immediately connected to node {}", id);
+                immediatelyConnectedKeys.add(key);
+                key.interestOps(0);
+            }
+        } catch (IOException | RuntimeException e) {
+            if (key != null)
+                immediatelyConnectedKeys.remove(key);
+            channels.remove(id);
+            socketChannel.close();
+            throw e;
+        }
+    }
+
+    // Visible to allow test cases to override. In particular, we use this to implement a blocking connect
+    // in order to simulate "immediately connected" sockets.
+    protected boolean doConnect(SocketChannel channel, InetSocketAddress address) throws IOException {
+        try {
+            return channel.connect(address);
+        } catch (UnresolvedAddressException e) {
+            throw new IOException("Can't resolve address: " + address, e);
+        }
+    }
+
+    private void configureSocketChannel(SocketChannel socketChannel, int sendBufferSize, int receiveBufferSize)
+            throws IOException {
         socketChannel.configureBlocking(false);
         Socket socket = socketChannel.socket();
         socket.setKeepAlive(true);
@@ -209,25 +290,6 @@ public void connect(String id, InetSocketAddress address, int sendBufferSize, in
         if (receiveBufferSize != Selectable.USE_DEFAULT_BUFFER_SIZE)
             socket.setReceiveBufferSize(receiveBufferSize);
         socket.setTcpNoDelay(true);
-        boolean connected;
-        try {
-            connected = socketChannel.connect(address);
-        } catch (UnresolvedAddressException e) {
-            socketChannel.close();
-            throw new IOException("Can't resolve address: " + address, e);
-        } catch (IOException e) {
-            socketChannel.close();
-            throw e;
-        }
-        SelectionKey key = socketChannel.register(nioSelector, SelectionKey.OP_CONNECT);
-        KafkaChannel channel = buildChannel(socketChannel, id, key);
-
-        if (connected) {
-            // OP_CONNECT won't trigger for immediately connected channels
-            log.debug("Immediately connected to node {}", channel.id());
-            immediatelyConnectedKeys.add(key);
-            key.interestOps(0);
-        }
     }
 
     /**
@@ -245,19 +307,38 @@ public void connect(String id, InetSocketAddress address, int sendBufferSize, in
      * 
        
      */
     public void register(String id, SocketChannel socketChannel) throws IOException {
+        ensureNotRegistered(id);
+        registerChannel(id, socketChannel, SelectionKey.OP_READ);
+        this.sensors.connectionCreated.record();
+        // Default to empty client information as the ApiVersionsRequest is not
+        // mandatory. In this case, we still want to account for the connection.
+        ChannelMetadataRegistry metadataRegistry = this.channel(id).channelMetadataRegistry();
+        if (metadataRegistry.clientInformation() == null)
+            metadataRegistry.registerClientInformation(ClientInformation.EMPTY);
+    }
+
+    private void ensureNotRegistered(String id) {
         if (this.channels.containsKey(id))
             throw new IllegalStateException("There is already a connection for id " + id);
         if (this.closingChannels.containsKey(id))
             throw new IllegalStateException("There is already a connection for id " + id + " that is still being closed");
+    }
 
-        SelectionKey key = socketChannel.register(nioSelector, SelectionKey.OP_READ);
-        buildChannel(socketChannel, id, key);
+    protected SelectionKey registerChannel(String id, SocketChannel socketChannel, int interestedOps) throws IOException {
+        SelectionKey key = socketChannel.register(nioSelector, interestedOps);
+        KafkaChannel channel = buildAndAttachKafkaChannel(socketChannel, id, key);
+        this.channels.put(id, channel);
+        if (idleExpiryManager != null)
+            idleExpiryManager.update(channel.id(), time.nanoseconds());
+        return key;
     }
 
-    private KafkaChannel buildChannel(SocketChannel socketChannel, String id, SelectionKey key) throws IOException {
-        KafkaChannel channel;
+    private KafkaChannel buildAndAttachKafkaChannel(SocketChannel socketChannel, String id, SelectionKey key) throws IOException {
         try {
-            channel = channelBuilder.buildChannel(id, key, maxReceiveSize, memoryPool);
+            KafkaChannel channel = channelBuilder.buildChannel(id, key, maxReceiveSize, memoryPool,
+                new SelectorChannelMetadataRegistry());
+            key.attach(channel);
+            return channel;
         } catch (Exception e) {
             try {
                 socketChannel.close();
@@ -266,9 +347,6 @@ private KafkaChannel buildChannel(SocketChannel socketChannel, String id, Select
             }
             throw new IOException("Channel could not be created for socket " + socketChannel, e);
         }
-        key.attach(channel);
-        this.channels.put(id, channel);
-        return channel;
     }
 
     /**
@@ -285,23 +363,28 @@ 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);
+        AtomicReference firstException = new AtomicReference<>();
+        Utils.closeAllQuietly(firstException, "release connections",
+                connections.stream().map(id -> (AutoCloseable) () -> close(id)).toArray(AutoCloseable[]::new));
+        // 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.
+        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();
     }
 
     /**
      * Queue the given request for sending in the subsequent {@link #poll(long)} calls
      * @param send The request to send
      */
-    public void send(Send send) {
-        String connectionId = send.destination();
+    public void send(NetworkSend send) {
+        String connectionId = send.destinationId();
         KafkaChannel channel = openOrClosingChannelOrFail(connectionId);
         if (closingChannels.containsKey(connectionId)) {
             // ensure notification via `disconnected`, leave channel in the state in which closing was triggered
@@ -312,9 +395,9 @@ public void send(Send send) {
             } catch (Exception e) {
                 // update the state for consistency, the channel will be discarded after `close`
                 channel.state(ChannelState.FAILED_SEND);
-                // ensure notification via `disconnected`
+                // ensure notification via `disconnected` when `failedSends` are processed in the next poll
                 this.failedSends.add(connectionId);
-                close(channel, false);
+                close(channel, CloseMode.DISCARD_NO_NOTIFY);
                 if (!(e instanceof CancelledKeyException)) {
                     log.error("Unexpected exception during send, closing connection {} and rethrowing exception {}",
                             connectionId, e);
@@ -338,11 +421,10 @@ public void send(Send send) {
      * This requires additional buffers to be maintained as we are reading from network, since the data on the wire is encrypted
      * we won't be able to read exact no.of bytes as kafka protocol requires. We read as many bytes as we can, up to SSLEngine's
      * application buffer size. This means we might be reading additional bytes than the requested size.
-     * If there is no further data to read from socketChannel selector won't invoke that channel and we've have additional bytes
-     * in the buffer. To overcome this issue we added "stagedReceives" map which contains per-channel deque. When we are
-     * reading a channel we read as many responses as we can and store them into "stagedReceives" and pop one response during
-     * the poll to add the completedReceives. If there are any active channels in the "stagedReceives" we set "timeout" to 0
-     * and pop response and add to the completedReceives.
+     * If there is no further data to read from socketChannel selector won't invoke that channel and we have additional bytes
+     * in the buffer. To overcome this issue we added "keysWithBufferedRead" map which tracks channels which have data in the SSL
+     * buffers. If there are channels with buffered data that can by processed, we set "timeout" to 0 and process the data even
+     * if there is no more data to read from the socket.
      *
      * Atmost one entry is added to "completedReceives" for a channel in each poll. This is necessary to guarantee that
      * requests from a channel are processed on the broker in the order they are sent. Since outstanding requests added
@@ -364,7 +446,7 @@ public void poll(long timeout) throws IOException {
 
         boolean dataInBuffers = !keysWithBufferedRead.isEmpty();
 
-        if (hasStagedReceives() || !immediatelyConnectedKeys.isEmpty() || (madeReadProgressLastCall && dataInBuffers))
+        if (!immediatelyConnectedKeys.isEmpty() || (madeReadProgressLastCall && dataInBuffers))
             timeout = 0;
 
         if (!memoryPool.isOutOfMemory() && outOfMemory) {
@@ -372,7 +454,7 @@ public void poll(long timeout) throws IOException {
             log.trace("Broker no longer low on memory - unmuting incoming sockets");
             for (KafkaChannel channel : channels.values()) {
                 if (channel.isInMutableState() && !explicitlyMutedChannels.contains(channel)) {
-                    channel.unmute();
+                    channel.maybeUnmute();
                 }
             }
             outOfMemory = false;
@@ -386,20 +468,22 @@ public void poll(long timeout) throws IOException {
 
         if (numReadyKeys > 0 || !immediatelyConnectedKeys.isEmpty() || dataInBuffers) {
             Set readyKeys = this.nioSelector.selectedKeys();
-            keysWithBufferedRead.removeAll(readyKeys); //so no channel gets polled twice
 
-            //poll from channels that have buffered data (but nothing more from the underlying socket)
-            if (!keysWithBufferedRead.isEmpty()) {
+            // Poll from channels that have buffered data (but nothing more from the underlying socket)
+            if (dataInBuffers) {
+                keysWithBufferedRead.removeAll(readyKeys); //so no channel gets polled twice
                 Set toPoll = keysWithBufferedRead;
                 keysWithBufferedRead = new HashSet<>(); //poll() calls will repopulate if needed
                 pollSelectionKeys(toPoll, false, endSelect);
             }
-            //poll from channels where the underlying socket has more data
-            pollSelectionKeys(readyKeys, false, endSelect);
-            pollSelectionKeys(immediatelyConnectedKeys, true, endSelect);
 
+            // Poll from channels where the underlying socket has more data
+            pollSelectionKeys(readyKeys, false, endSelect);
             // Clear all selected keys so that they are included in the ready count for the next select
             readyKeys.clear();
+
+            pollSelectionKeys(immediatelyConnectedKeys, true, endSelect);
+            immediatelyConnectedKeys.clear();
         } else {
             madeReadProgressLastPoll = true; //no work is also "progress"
         }
@@ -407,13 +491,12 @@ public void poll(long timeout) throws IOException {
         long endIo = time.nanoseconds();
         this.sensors.ioTime.record(endIo - endSelect, time.milliseconds());
 
+        // Close channels that were delayed and are now ready to be closed
+        completeDelayedChannelClose(endIo);
+
         // we use the time at the end of select to ensure that we don't close any connections that
         // have just been processed in pollSelectionKeys
         maybeCloseOldestConnection(endSelect);
-
-        // Add to completedReceives after closing expired connections to avoid removing
-        // channels with completed receives until all staged receives are completed.
-        addToCompletedReceives();
     }
 
     /**
@@ -424,117 +507,204 @@ public void poll(long timeout) throws IOException {
      */
     // package-private for testing
     void pollSelectionKeys(Set selectionKeys,
-                                   boolean isImmediatelyConnected,
-                                   long currentTimeNanos) {
-        Iterator iterator = determineHandlingOrder(selectionKeys).iterator();
-        while (iterator.hasNext()) {
-            SelectionKey key = iterator.next();
+                           boolean isImmediatelyConnected,
+                           long currentTimeNanos) {
+        for (SelectionKey key : determineHandlingOrder(selectionKeys)) {
             KafkaChannel channel = channel(key);
             long channelStartTimeNanos = recordTimePerConnection ? time.nanoseconds() : 0;
+            boolean sendFailed = false;
+            String nodeId = channel.id();
 
             // register all per-connection metrics at once
-            sensors.maybeRegisterConnectionMetrics(channel.id());
+            sensors.maybeRegisterConnectionMetrics(nodeId);
             if (idleExpiryManager != null)
-                idleExpiryManager.update(channel.id(), currentTimeNanos);
+                idleExpiryManager.update(nodeId, currentTimeNanos);
 
             try {
-
                 /* complete any connections that have finished their handshake (either normally or immediately) */
                 if (isImmediatelyConnected || key.isConnectable()) {
                     if (channel.finishConnect()) {
-                        this.connected.add(channel.id());
+                        this.connected.add(nodeId);
                         this.sensors.connectionCreated.record();
+
                         SocketChannel socketChannel = (SocketChannel) key.channel();
                         log.debug("Created socket with SO_RCVBUF = {}, SO_SNDBUF = {}, SO_TIMEOUT = {} to node {}",
                                 socketChannel.socket().getReceiveBufferSize(),
                                 socketChannel.socket().getSendBufferSize(),
                                 socketChannel.socket().getSoTimeout(),
-                                channel.id());
-                    } else
+                                nodeId);
+                    } else {
                         continue;
+                    }
                 }
 
                 /* if channel is not ready finish prepare */
                 if (channel.isConnected() && !channel.ready()) {
-                    try {
-                        channel.prepare();
-                    } catch (AuthenticationException e) {
-                        sensors.failedAuthentication.record();
-                        throw e;
+                    channel.prepare();
+                    if (channel.ready()) {
+                        long readyTimeMs = time.milliseconds();
+                        boolean isReauthentication = channel.successfulAuthentications() > 1;
+                        if (isReauthentication) {
+                            sensors.successfulReauthentication.record(1.0, readyTimeMs);
+                            if (channel.reauthenticationLatencyMs() == null)
+                                log.warn(
+                                    "Should never happen: re-authentication latency for a re-authenticated channel was null; continuing...");
+                            else
+                                sensors.reauthenticationLatency
+                                    .record(channel.reauthenticationLatencyMs().doubleValue(), readyTimeMs);
+                        } else {
+                            sensors.successfulAuthentication.record(1.0, readyTimeMs);
+                            if (!channel.connectedClientSupportsReauthentication())
+                                sensors.successfulAuthenticationNoReauth.record(1.0, readyTimeMs);
+                        }
+                        log.debug("Successfully {}authenticated with {}", isReauthentication ?
+                            "re-" : "", channel.socketDescription());
                     }
-                    if (channel.ready())
-                        sensors.successfulAuthentication.record();
                 }
-
-                attemptRead(key, channel);
+                if (channel.ready() && channel.state() == ChannelState.NOT_CONNECTED)
+                    channel.state(ChannelState.READY);
+                Optional responseReceivedDuringReauthentication = channel.pollResponseReceivedDuringReauthentication();
+                responseReceivedDuringReauthentication.ifPresent(receive -> {
+                    long currentTimeMs = time.milliseconds();
+                    addToCompletedReceives(channel, receive, currentTimeMs);
+                });
+
+                //if channel is ready and has bytes to read from socket or buffer, and has no
+                //previous completed receive then read from it
+                if (channel.ready() && (key.isReadable() || channel.hasBytesBuffered()) && !hasCompletedReceive(channel)
+                        && !explicitlyMutedChannels.contains(channel)) {
+                    attemptRead(channel);
+                }
 
                 if (channel.hasBytesBuffered()) {
                     //this channel has bytes enqueued in intermediary buffers that we could not read
                     //(possibly because no memory). it may be the case that the underlying socket will
                     //not come up in the next poll() and so we need to remember this channel for the
-                    //next poll call otherwise data may be stuck in said buffers forever.
+                    //next poll call otherwise data may be stuck in said buffers forever. If we attempt
+                    //to process buffered data and no progress is made, the channel buffered status is
+                    //cleared to avoid the overhead of checking every time.
                     keysWithBufferedRead.add(key);
                 }
 
                 /* if channel is ready write to any sockets that have space in their buffer and for which we have data */
-                if (channel.ready() && key.isWritable()) {
-                    Send send = channel.write();
-                    if (send != null) {
-                        this.completedSends.add(send);
-                        this.sensors.recordBytesSent(channel.id(), send.size());
-                    }
+
+                long nowNanos = channelStartTimeNanos != 0 ? channelStartTimeNanos : currentTimeNanos;
+                try {
+                    attemptWrite(key, channel, nowNanos);
+                } catch (Exception e) {
+                    sendFailed = true;
+                    throw e;
                 }
 
                 /* cancel any defunct sockets */
                 if (!key.isValid())
-                    close(channel, true);
+                    close(channel, CloseMode.GRACEFUL);
 
             } catch (Exception e) {
                 String desc = channel.socketDescription();
-                if (e instanceof IOException)
+                if (e instanceof IOException) {
                     log.debug("Connection with {} disconnected", desc, e);
-                else if (e instanceof AuthenticationException) // will be logged later as error by clients
-                    log.debug("Connection with {} disconnected due to authentication exception", desc, e);
-                else
+                } else if (e instanceof AuthenticationException) {
+                    boolean isReauthentication = channel.successfulAuthentications() > 0;
+                    if (isReauthentication)
+                        sensors.failedReauthentication.record();
+                    else
+                        sensors.failedAuthentication.record();
+                    String exceptionMessage = e.getMessage();
+                    if (e instanceof DelayedResponseAuthenticationException)
+                        exceptionMessage = e.getCause().getMessage();
+                    log.info("Failed {}authentication with {} ({})", isReauthentication ? "re-" : "",
+                        desc, exceptionMessage);
+                } else {
                     log.warn("Unexpected error from {}; closing connection", desc, e);
-                close(channel, true);
+                }
+
+                if (e instanceof DelayedResponseAuthenticationException)
+                    maybeDelayCloseOnAuthenticationFailure(channel);
+                else
+                    close(channel, sendFailed ? CloseMode.NOTIFY_ONLY : CloseMode.GRACEFUL);
             } finally {
                 maybeRecordTimePerConnection(channel, channelStartTimeNanos);
             }
         }
     }
 
+    private void attemptWrite(SelectionKey key, KafkaChannel channel, long nowNanos) throws IOException {
+        if (channel.hasSend()
+                && channel.ready()
+                && key.isWritable()
+                && !channel.maybeBeginClientReauthentication(() -> nowNanos)) {
+            write(channel);
+        }
+    }
+
+    // package-private for testing
+    void write(KafkaChannel channel) throws IOException {
+        String nodeId = channel.id();
+        long bytesSent = channel.write();
+        NetworkSend send = channel.maybeCompleteSend();
+        // We may complete the send with bytesSent < 1 if `TransportLayer.hasPendingWrites` was true and `channel.write()`
+        // caused the pending writes to be written to the socket channel buffer
+        if (bytesSent > 0 || send != null) {
+            long currentTimeMs = time.milliseconds();
+            if (bytesSent > 0)
+                this.sensors.recordBytesSent(nodeId, bytesSent, currentTimeMs);
+            if (send != null) {
+                this.completedSends.add(send);
+                this.sensors.recordCompletedSend(nodeId, send.size(), currentTimeMs);
+            }
+        }
+    }
+
     private Collection determineHandlingOrder(Set selectionKeys) {
         //it is possible that the iteration order over selectionKeys is the same every invocation.
         //this may cause starvation of reads when memory is low. to address this we shuffle the keys if memory is low.
-        Collection inHandlingOrder;
-
         if (!outOfMemory && memoryPool.availableMemory() < lowMemThreshold) {
-            List temp = new ArrayList<>(selectionKeys);
-            Collections.shuffle(temp);
-            inHandlingOrder = temp;
+            List shuffledKeys = new ArrayList<>(selectionKeys);
+            Collections.shuffle(shuffledKeys);
+            return shuffledKeys;
         } else {
-            inHandlingOrder = selectionKeys;
+            return selectionKeys;
         }
-        return inHandlingOrder;
     }
 
-    private void attemptRead(SelectionKey key, KafkaChannel channel) throws IOException {
-        //if channel is ready and has bytes to read from socket or buffer, and has no
-        //previous receive(s) already staged or otherwise in progress then read from it
-        if (channel.ready() && (key.isReadable() || channel.hasBytesBuffered()) && !hasStagedReceive(channel)
-            && !explicitlyMutedChannels.contains(channel)) {
-            NetworkReceive networkReceive;
-            while ((networkReceive = channel.read()) != null) {
-                madeReadProgressLastPoll = true;
-                addToStagedReceives(channel, networkReceive);
+    private void attemptRead(KafkaChannel channel) throws IOException {
+        String nodeId = channel.id();
+
+        long bytesReceived = channel.read();
+        if (bytesReceived != 0) {
+            long currentTimeMs = time.milliseconds();
+            sensors.recordBytesReceived(nodeId, bytesReceived, currentTimeMs);
+            madeReadProgressLastPoll = true;
+
+            NetworkReceive receive = channel.maybeCompleteReceive();
+            if (receive != null) {
+                addToCompletedReceives(channel, receive, currentTimeMs);
             }
-            if (channel.isMute()) {
-                outOfMemory = true; //channel has muted itself due to memory pressure.
-            } else {
-                madeReadProgressLastPoll = true;
+        }
+        if (channel.isMuted()) {
+            outOfMemory = true; //channel has muted itself due to memory pressure.
+        } else {
+            madeReadProgressLastPoll = true;
+        }
+    }
+
+    private boolean maybeReadFromClosingChannel(KafkaChannel channel) {
+        boolean hasPending;
+        if (channel.state().state() != ChannelState.State.READY)
+            hasPending = false;
+        else if (explicitlyMutedChannels.contains(channel) || hasCompletedReceive(channel))
+            hasPending = true;
+        else {
+            try {
+                attemptRead(channel);
+                hasPending = hasCompletedReceive(channel);
+            } catch (Exception e) {
+                log.trace("Read from closing channel failed, ignoring exception", e);
+                hasPending = false;
             }
         }
+        return hasPending;
     }
 
     // Record time spent in pollSelectionKeys for channel (moved into a method to keep checkstyle happy)
@@ -544,13 +714,13 @@ private void maybeRecordTimePerConnection(KafkaChannel channel, long startTimeNa
     }
 
     @Override
-    public List completedSends() {
+    public List completedSends() {
         return this.completedSends;
     }
 
     @Override
-    public List completedReceives() {
-        return this.completedReceives;
+    public Collection completedReceives() {
+        return this.completedReceives.values();
     }
 
     @Override
@@ -581,8 +751,10 @@ public void unmute(String id) {
     }
 
     private void unmute(KafkaChannel channel) {
-        explicitlyMutedChannels.remove(channel);
-        channel.unmute();
+        // Remove the channel from explicitlyMutedChannels only if the channel has been actually unmuted.
+        if (channel.maybeUnmute()) {
+            explicitlyMutedChannels.remove(channel);
+        }
     }
 
     @Override
@@ -597,6 +769,18 @@ public void unmuteAll() {
             unmute(channel);
     }
 
+    // package-private for testing
+    void completeDelayedChannelClose(long currentTimeNanos) {
+        if (delayedClosingChannels == null)
+            return;
+
+        while (!delayedClosingChannels.isEmpty()) {
+            DelayedAuthenticationFailureClose delayedClose = delayedClosingChannels.values().iterator().next();
+            if (!delayedClose.tryClose(currentTimeNanos))
+                break;
+        }
+    }
+
     private void maybeCloseOldestConnection(long currentTimeNanos) {
         if (idleExpiryManager == null)
             return;
@@ -610,29 +794,59 @@ private void maybeCloseOldestConnection(long currentTimeNanos) {
                     log.trace("About to close the idle connection from {} due to being idle for {} millis",
                             connectionId, (currentTimeNanos - expiredConnection.getValue()) / 1000 / 1000);
                 channel.state(ChannelState.EXPIRED);
-                close(channel, true);
+                close(channel, CloseMode.GRACEFUL);
             }
         }
     }
 
     /**
-     * Clear the results from the prior poll
+     * Clears completed receives. This is used by SocketServer to remove references to
+     * receive buffers after processing completed receives, without waiting for the next
+     * poll().
+     */
+    public void clearCompletedReceives() {
+        this.completedReceives.clear();
+    }
+
+    /**
+     * Clears completed sends. This is used by SocketServer to remove references to
+     * send buffers after processing completed sends, without waiting for the next
+     * poll().
+     */
+    public void clearCompletedSends() {
+        this.completedSends.clear();
+    }
+
+    /**
+     * Clears all the results from the previous poll. This is invoked by Selector at the start of
+     * a poll() when all the results from the previous poll are expected to have been handled.
+     * 
        
+     * SocketServer uses {@link #clearCompletedSends()} and {@link #clearCompletedSends()} to
+     * clear `completedSends` and `completedReceives` as soon as they are processed to avoid
+     * holding onto large request/response buffers from multiple connections longer than necessary.
+     * Clients rely on Selector invoking {@link #clear()} at the start of each poll() since memory usage
+     * is less critical and clearing once-per-poll provides the flexibility to process these results in
+     * any order before the next poll.
      */
     private void clear() {
         this.completedSends.clear();
         this.completedReceives.clear();
         this.connected.clear();
         this.disconnected.clear();
-        // Remove closed channels after all their staged receives have been processed or if a send was requested
+
+        // Remove closed channels after all their buffered receives have been processed or if a send was requested
         for (Iterator> it = closingChannels.entrySet().iterator(); it.hasNext(); ) {
             KafkaChannel channel = it.next().getValue();
-            Deque deque = this.stagedReceives.get(channel);
             boolean sendFailed = failedSends.remove(channel.id());
-            if (deque == null || deque.isEmpty() || sendFailed) {
+            boolean hasPending = false;
+            if (!sendFailed)
+                hasPending = maybeReadFromClosingChannel(channel);
+            if (!hasPending || sendFailed) {
                 doClose(channel, true);
                 it.remove();
             }
         }
+
         for (String channel : this.failedSends)
             this.disconnected.put(channel, ChannelState.FAILED_SEND);
         this.failedSends.clear();
@@ -642,19 +856,17 @@ private void clear() {
     /**
      * Check for data, waiting up to the given timeout.
      *
-     * @param ms Length of time to wait, in milliseconds, which must be non-negative
+     * @param timeoutMs Length of time to wait, in milliseconds, which must be non-negative
      * @return The number of keys ready
-     * @throws IllegalArgumentException
-     * @throws IOException
      */
-    private int select(long ms) throws IOException {
-        if (ms < 0L)
+    private int select(long timeoutMs) throws IOException {
+        if (timeoutMs < 0L)
             throw new IllegalArgumentException("timeout should be >= 0");
 
-        if (ms == 0L)
+        if (timeoutMs == 0L)
             return this.nioSelector.selectNow();
         else
-            return this.nioSelector.select(ms);
+            return this.nioSelector.select(timeoutMs);
     }
 
     /**
@@ -666,7 +878,7 @@ public void close(String id) {
             // There is no disconnect notification for local close, but updating
             // channel state here anyway to avoid confusion.
             channel.state(ChannelState.LOCAL_CLOSE);
-            close(channel, false);
+            close(channel, CloseMode.DISCARD_NO_NOTIFY);
         } else {
             KafkaChannel closingChannel = this.closingChannels.remove(id);
             // Close any closing channel, leave the channel in the state in which closing was triggered
@@ -675,19 +887,35 @@ public void close(String id) {
         }
     }
 
+    private void maybeDelayCloseOnAuthenticationFailure(KafkaChannel channel) {
+        DelayedAuthenticationFailureClose delayedClose = new DelayedAuthenticationFailureClose(channel, failedAuthenticationDelayMs);
+        if (delayedClosingChannels != null)
+            delayedClosingChannels.put(channel.id(), delayedClose);
+        else
+            delayedClose.closeNow();
+    }
+
+    private void handleCloseOnAuthenticationFailure(KafkaChannel channel) {
+        try {
+            channel.completeCloseOnAuthenticationFailure();
+        } catch (Exception e) {
+            log.error("Exception handling close on authentication failure node {}", channel.id(), e);
+        } finally {
+            close(channel, CloseMode.GRACEFUL);
+        }
+    }
+
     /**
      * Begin closing this connection.
+     * If 'closeMode' is `CloseMode.GRACEFUL`, the channel is disconnected here, but outstanding receives
+     * are processed. The channel is closed when there are no outstanding receives or if a send is
+     * requested. For other values of `closeMode`, outstanding receives are discarded and the channel
+     * is closed immediately.
      *
-     * If 'processOutstanding' is true, the channel is disconnected here, but staged receives are
-     * processed. The channel is closed when there are no outstanding receives or if a send
-     * is requested. The channel will be added to disconnect list when it is actually closed.
-     *
-     * If 'processOutstanding' is false, outstanding receives are discarded and the channel is
-     * closed immediately. The channel will not be added to disconnected list and it is the
-     * responsibility of the caller to handle disconnect notifications.
+     * The channel will be added to disconnect list when it is actually closed if `closeMode.notifyDisconnect`
+     * is true.
      */
-    private void close(KafkaChannel channel, boolean processOutstanding) {
-
+    private void close(KafkaChannel channel, CloseMode closeMode) {
         channel.disconnect();
 
         // Ensure that `connected` does not have closed channels. This could happen if `prepare` throws an exception
@@ -700,26 +928,35 @@ private void close(KafkaChannel channel, boolean processOutstanding) {
         // handle close(). When the remote end closes its connection, the channel is retained until
         // a send fails or all outstanding receives are processed. Mute state of disconnected channels
         // are tracked to ensure that requests are processed one-by-one by the broker to preserve ordering.
-        Deque deque = this.stagedReceives.get(channel);
-        if (processOutstanding && deque != null && !deque.isEmpty()) {
-            // stagedReceives will be moved to completedReceives later along with receives from other channels
+        if (closeMode == CloseMode.GRACEFUL && maybeReadFromClosingChannel(channel)) {
             closingChannels.put(channel.id(), channel);
-        } else
-            doClose(channel, processOutstanding);
+            log.debug("Tracking closing connection {} to process outstanding requests", channel.id());
+        } else {
+            doClose(channel, closeMode.notifyDisconnect);
+        }
         this.channels.remove(channel.id());
 
+        if (delayedClosingChannels != null)
+            delayedClosingChannels.remove(channel.id());
+
         if (idleExpiryManager != null)
             idleExpiryManager.remove(channel.id());
     }
 
     private void doClose(KafkaChannel channel, boolean notifyDisconnect) {
+        SelectionKey key = channel.selectionKey();
         try {
+            immediatelyConnectedKeys.remove(key);
+            keysWithBufferedRead.remove(key);
             channel.close();
         } catch (IOException e) {
             log.error("Exception closing connection to node {}:", channel.id(), e);
+        } finally {
+            key.cancel();
+            key.attach(null);
         }
+
         this.sensors.connectionClosed.record();
-        this.stagedReceives.remove(channel);
         this.explicitlyMutedChannels.remove(channel);
         if (notifyDisconnect)
             this.disconnected.put(channel.id(), channel.state());
@@ -766,6 +1003,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
      */
@@ -774,86 +1033,105 @@ private KafkaChannel channel(SelectionKey key) {
     }
 
     /**
-     * Check if given channel has a staged receive
+     * Check if given channel has a completed receive
      */
-    private boolean hasStagedReceive(KafkaChannel channel) {
-        return stagedReceives.containsKey(channel);
+    private boolean hasCompletedReceive(KafkaChannel channel) {
+        return completedReceives.containsKey(channel.id());
     }
 
     /**
-     * check if stagedReceives have unmuted channel
+     * adds a receive to completed receives
      */
-    private boolean hasStagedReceives() {
-        for (KafkaChannel channel : this.stagedReceives.keySet()) {
-            if (!channel.isMute())
-                return true;
-        }
-        return false;
+    private void addToCompletedReceives(KafkaChannel channel, NetworkReceive networkReceive, long currentTimeMs) {
+        if (hasCompletedReceive(channel))
+            throw new IllegalStateException("Attempting to add second completed receive to channel " + channel.id());
+
+        this.completedReceives.put(channel.id(), networkReceive);
+        sensors.recordCompletedReceive(channel.id(), networkReceive.size(), currentTimeMs);
     }
 
+    // only for testing
+    public Set keys() {
+        return new HashSet<>(nioSelector.keys());
+    }
 
-    /**
-     * adds a receive to staged receives
-     */
-    private void addToStagedReceives(KafkaChannel channel, NetworkReceive receive) {
-        if (!stagedReceives.containsKey(channel))
-            stagedReceives.put(channel, new ArrayDeque());
 
-        Deque deque = stagedReceives.get(channel);
-        deque.add(receive);
-    }
+    class SelectorChannelMetadataRegistry implements ChannelMetadataRegistry {
+        private CipherInformation cipherInformation;
+        private ClientInformation clientInformation;
 
-    /**
-     * checks if there are any staged receives and adds to completedReceives
-     */
-    private void addToCompletedReceives() {
-        if (!this.stagedReceives.isEmpty()) {
-            Iterator>> iter = this.stagedReceives.entrySet().iterator();
-            while (iter.hasNext()) {
-                Map.Entry> entry = iter.next();
-                KafkaChannel channel = entry.getKey();
-                if (!explicitlyMutedChannels.contains(channel)) {
-                    Deque deque = entry.getValue();
-                    addToCompletedReceives(channel, deque);
-                    if (deque.isEmpty())
-                        iter.remove();
-                }
+        @Override
+        public void registerCipherInformation(final CipherInformation cipherInformation) {
+            if (this.cipherInformation != null) {
+                if (this.cipherInformation.equals(cipherInformation))
+                    return;
+                sensors.connectionsByCipher.decrement(this.cipherInformation);
             }
+
+            this.cipherInformation = cipherInformation;
+            sensors.connectionsByCipher.increment(cipherInformation);
         }
-    }
 
-    private void addToCompletedReceives(KafkaChannel channel, Deque stagedDeque) {
-        NetworkReceive networkReceive = stagedDeque.poll();
-        this.completedReceives.add(networkReceive);
-        this.sensors.recordBytesReceived(channel.id(), networkReceive.payload().limit());
-    }
+        @Override
+        public CipherInformation cipherInformation() {
+            return cipherInformation;
+        }
 
-    // only for testing
-    public Set keys() {
-        return new HashSet<>(nioSelector.keys());
-    }
+        @Override
+        public void registerClientInformation(final ClientInformation clientInformation) {
+            if (this.clientInformation != null) {
+                if (this.clientInformation.equals(clientInformation))
+                    return;
+                sensors.connectionsByClient.decrement(this.clientInformation);
+            }
 
-    // only for testing
-    public int numStagedReceives(KafkaChannel channel) {
-        Deque deque = stagedReceives.get(channel);
-        return deque == null ? 0 : deque.size();
+            this.clientInformation = clientInformation;
+            sensors.connectionsByClient.increment(clientInformation);
+        }
+
+        @Override
+        public ClientInformation clientInformation() {
+            return clientInformation;
+        }
+
+        @Override
+        public void close() {
+            if (this.cipherInformation != null) {
+                sensors.connectionsByCipher.decrement(this.cipherInformation);
+                this.cipherInformation = null;
+            }
+
+            if (this.clientInformation != null) {
+                sensors.connectionsByClient.decrement(this.clientInformation);
+                this.clientInformation = null;
+            }
+        }
     }
 
-    private class SelectorMetrics {
+    class SelectorMetrics implements AutoCloseable {
         private final Metrics metrics;
-        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;
         public final Sensor successfulAuthentication;
+        public final Sensor successfulReauthentication;
+        public final Sensor successfulAuthenticationNoReauth;
+        public final Sensor reauthenticationLatency;
         public final Sensor failedAuthentication;
+        public final Sensor failedReauthentication;
         public final Sensor bytesTransferred;
         public final Sensor bytesSent;
+        public final Sensor requestsSent;
         public final Sensor bytesReceived;
+        public final Sensor responsesReceived;
         public final Sensor selectTime;
         public final Sensor ioTime;
+        public final IntGaugeSuite connectionsByCipher;
+        public final IntGaugeSuite connectionsByClient;
 
         /* Names of metrics that are not registered through sensors */
         private final List topLevelMetricNames = new ArrayList<>();
@@ -861,10 +1139,10 @@ private class SelectorMetrics {
 
         public SelectorMetrics(Metrics metrics, String metricGrpPrefix, Map metricTags, boolean metricsPerConnection) {
             this.metrics = metrics;
-            this.metricGrpPrefix = metricGrpPrefix;
             this.metricTags = metricTags;
             this.metricsPerConnection = metricsPerConnection;
-            String metricGrpName = metricGrpPrefix + "-metrics";
+            this.metricGrpName = metricGrpPrefix + "-metrics";
+            this.perConnectionMetricGrpName = metricGrpPrefix + "-node-metrics";
             StringBuilder tagsSuffix = new StringBuilder();
 
             for (Map.Entry tag: metricTags.entrySet()) {
@@ -885,33 +1163,62 @@ public SelectorMetrics(Metrics metrics, String metricGrpPrefix, Map(log, "sslCiphers", metrics,
+                cipherInformation -> {
+                    Map tags = new LinkedHashMap<>();
+                    tags.put("cipher", cipherInformation.cipher());
+                    tags.put("protocol", cipherInformation.protocol());
+                    tags.putAll(metricTags);
+                    return metrics.metricName("connections", metricGrpName, "The number of connections with this SSL cipher and protocol.", tags);
+                }, 100);
+
+            this.connectionsByClient = new IntGaugeSuite<>(log, "clients", metrics,
+                clientInformation -> {
+                    Map tags = new LinkedHashMap<>();
+                    tags.put("clientSoftwareName", clientInformation.softwareName());
+                    tags.put("clientSoftwareVersion", clientInformation.softwareVersion());
+                    tags.putAll(metricTags);
+                    return metrics.metricName("connections", metricGrpName, "The number of connections with this client and version.", tags);
+                }, 100);
+
             metricName = metrics.metricName("connection-count", metricGrpName, "The current number of active connections.", metricTags);
             topLevelMetricNames.add(metricName);
-            this.metrics.addMetric(metricName, new Measurable() {
-                public double measure(MetricConfig config, long now) {
-                    return channels.size();
-                }
-            });
+            this.metrics.addMetric(metricName, (config, now) -> channels.size());
         }
 
         private Meter createMeter(Metrics metrics, String groupName, Map metricTags,
@@ -966,56 +1287,78 @@ public void maybeRegisterConnectionMetrics(String connectionId) {
             if (!connectionId.isEmpty() && metricsPerConnection) {
                 // if one sensor of the metrics has been registered for the connection,
                 // then all other sensors should have been registered; and vice versa
-                String nodeRequestName = "node-" + connectionId + ".bytes-sent";
+                String nodeRequestName = "node-" + connectionId + ".requests-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, 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";
+                    String bytesSentName = "node-" + connectionId + ".bytes-sent";
+                    Sensor bytesSent = sensor(bytesSentName);
+                    bytesSent.add(createMeter(metrics, perConnectionMetricGrpName, tags, "outgoing-byte", "outgoing bytes"));
+
+                    String nodeResponseName = "node-" + connectionId + ".responses-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, new WindowedCount(), "response", "responses received"));
+
+                    String bytesReceivedName = "node-" + connectionId + ".bytes-received";
+                    Sensor bytesReceive = sensor(bytesReceivedName);
+                    bytesReceive.add(createMeter(metrics, perConnectionMetricGrpName, tags, "incoming-byte", "incoming bytes"));
 
                     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());
                 }
             }
         }
 
-        public void recordBytesSent(String connectionId, long bytes) {
-            long now = time.milliseconds();
-            this.bytesSent.record(bytes, now);
+        public void recordBytesSent(String connectionId, long bytes, long currentTimeMs) {
+            this.bytesSent.record(bytes, currentTimeMs);
             if (!connectionId.isEmpty()) {
-                String nodeRequestName = "node-" + connectionId + ".bytes-sent";
+                String bytesSentName = "node-" + connectionId + ".bytes-sent";
+                Sensor bytesSent = this.metrics.getSensor(bytesSentName);
+                if (bytesSent != null)
+                    bytesSent.record(bytes, currentTimeMs);
+            }
+        }
+
+        public void recordCompletedSend(String connectionId, long totalBytes, long currentTimeMs) {
+            requestsSent.record(totalBytes, currentTimeMs);
+            if (!connectionId.isEmpty()) {
+                String nodeRequestName = "node-" + connectionId + ".requests-sent";
                 Sensor nodeRequest = this.metrics.getSensor(nodeRequestName);
                 if (nodeRequest != null)
-                    nodeRequest.record(bytes, now);
+                    nodeRequest.record(totalBytes, currentTimeMs);
+            }
+        }
+
+        public void recordBytesReceived(String connectionId, long bytes, long currentTimeMs) {
+            this.bytesReceived.record(bytes, currentTimeMs);
+            if (!connectionId.isEmpty()) {
+                String bytesReceivedName = "node-" + connectionId + ".bytes-received";
+                Sensor bytesReceived = this.metrics.getSensor(bytesReceivedName);
+                if (bytesReceived != null)
+                    bytesReceived.record(bytes, currentTimeMs);
             }
         }
 
-        public void recordBytesReceived(String connection, int bytes) {
-            long now = time.milliseconds();
-            this.bytesReceived.record(bytes, now);
-            if (!connection.isEmpty()) {
-                String nodeRequestName = "node-" + connection + ".bytes-received";
+        public void recordCompletedReceive(String connectionId, long totalBytes, long currentTimeMs) {
+            responsesReceived.record(totalBytes, currentTimeMs);
+            if (!connectionId.isEmpty()) {
+                String nodeRequestName = "node-" + connectionId + ".responses-received";
                 Sensor nodeRequest = this.metrics.getSensor(nodeRequestName);
                 if (nodeRequest != null)
-                    nodeRequest.record(bytes, now);
+                    nodeRequest.record(totalBytes, currentTimeMs);
             }
         }
 
@@ -1024,6 +1367,48 @@ public void close() {
                 metrics.removeMetric(metricName);
             for (Sensor sensor : sensors)
                 metrics.removeSensor(sensor.name());
+            connectionsByCipher.close();
+            connectionsByClient.close();
+        }
+    }
+
+    /**
+     * Encapsulate a channel that must be closed after a specific delay has elapsed due to authentication failure.
+     */
+    private class DelayedAuthenticationFailureClose {
+        private final KafkaChannel channel;
+        private final long endTimeNanos;
+        private boolean closed;
+
+        /**
+         * @param channel The channel whose close is being delayed
+         * @param delayMs The amount of time by which the operation should be delayed
+         */
+        public DelayedAuthenticationFailureClose(KafkaChannel channel, int delayMs) {
+            this.channel = channel;
+            this.endTimeNanos = time.nanoseconds() + (delayMs * 1000L * 1000L);
+            this.closed = false;
+        }
+
+        /**
+         * Try to close this channel if the delay has expired.
+         * @param currentTimeNanos The current time
+         * @return True if the delay has expired and the channel was closed; false otherwise
+         */
+        public final boolean tryClose(long currentTimeNanos) {
+            if (endTimeNanos <= currentTimeNanos)
+                closeNow();
+            return closed;
+        }
+
+        /**
+         * Close the channel now, regardless of whether the delay has expired or not.
+         */
+        public final void closeNow() {
+            if (closed)
+                throw new IllegalStateException("Attempt to close a channel that has already been closed");
+            handleCloseOnAuthenticationFailure(channel);
+            closed = true;
         }
     }
 
@@ -1077,4 +1462,9 @@ boolean isOutOfMemory() {
     boolean isMadeReadProgressLastPoll() {
         return madeReadProgressLastPoll;
     }
+
+    // package-private for testing
+    Map delayedClosingChannels() {
+        return delayedClosingChannels;
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/network/Send.java b/clients/src/main/java/org/apache/kafka/common/network/Send.java
index e6febc8ee864a..4903d55eb8a89 100644
--- a/clients/src/main/java/org/apache/kafka/common/network/Send.java
+++ b/clients/src/main/java/org/apache/kafka/common/network/Send.java
@@ -20,15 +20,10 @@
 import java.nio.channels.GatheringByteChannel;
 
 /**
- * This interface models the in-progress sending of data to a specific destination
+ * This interface models the in-progress sending of data.
  */
 public interface Send {
 
-    /**
-     * The id for the destination of this send
-     */
-    String destination();
-
     /**
      * Is this send complete?
      */
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 9519e586050c4..909009b329f06 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
@@ -17,14 +17,18 @@
 package org.apache.kafka.common.network;
 
 import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.config.SslConfigs;
+import org.apache.kafka.common.config.internals.BrokerSecurityConfigs;
 import org.apache.kafka.common.memory.MemoryPool;
 import org.apache.kafka.common.security.auth.KafkaPrincipal;
 import org.apache.kafka.common.security.auth.KafkaPrincipalBuilder;
+import org.apache.kafka.common.security.auth.KafkaPrincipalSerde;
 import org.apache.kafka.common.security.auth.SslAuthenticationContext;
 import org.apache.kafka.common.security.ssl.SslFactory;
+import org.apache.kafka.common.security.ssl.SslPrincipalMapper;
+import org.apache.kafka.common.utils.LogContext;
 import org.apache.kafka.common.utils.Utils;
 import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
 import java.io.Closeable;
 import java.io.IOException;
@@ -33,34 +37,78 @@
 import java.nio.channels.SelectionKey;
 import java.nio.channels.SocketChannel;
 import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.function.Supplier;
 
-public class SslChannelBuilder implements ChannelBuilder {
-    private static final Logger log = LoggerFactory.getLogger(SslChannelBuilder.class);
+public class SslChannelBuilder implements ChannelBuilder, ListenerReconfigurable {
+    private final ListenerName listenerName;
+    private final boolean isInterBrokerListener;
     private SslFactory sslFactory;
     private Mode mode;
     private Map configs;
+    private SslPrincipalMapper sslPrincipalMapper;
+    private final Logger log;
 
-    public SslChannelBuilder(Mode mode) {
+    /**
+     * 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,
+                             LogContext logContext) {
         this.mode = mode;
+        this.listenerName = listenerName;
+        this.isInterBrokerListener = isInterBrokerListener;
+        this.log = logContext.logger(getClass());
     }
 
     public void configure(Map configs) throws KafkaException {
         try {
             this.configs = configs;
-            this.sslFactory = new SslFactory(mode);
+            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);
             this.sslFactory.configure(this.configs);
+        } catch (KafkaException e) {
+            throw e;
         } catch (Exception e) {
             throw new KafkaException(e);
         }
     }
 
     @Override
-    public KafkaChannel buildChannel(String id, SelectionKey key, int maxReceiveSize, MemoryPool memoryPool) throws KafkaException {
+    public Set reconfigurableConfigs() {
+        return SslConfigs.RECONFIGURABLE_CONFIGS;
+    }
+
+    @Override
+    public void validateReconfiguration(Map configs) {
+        sslFactory.validateReconfiguration(configs);
+    }
+
+    @Override
+    public void reconfigure(Map configs) {
+        sslFactory.reconfigure(configs);
+    }
+
+    @Override
+    public ListenerName listenerName() {
+        return listenerName;
+    }
+
+    @Override
+    public KafkaChannel buildChannel(String id, SelectionKey key, int maxReceiveSize,
+                                     MemoryPool memoryPool, ChannelMetadataRegistry metadataRegistry) throws KafkaException {
         try {
-            SslTransportLayer transportLayer = buildTransportLayer(sslFactory, id, key, peerHost(key));
-            Authenticator authenticator = new SslAuthenticator(configs, transportLayer);
-            return new KafkaChannel(id, transportLayer, authenticator, maxReceiveSize,
-                    memoryPool != null ? memoryPool : MemoryPool.NONE);
+            SslTransportLayer transportLayer = buildTransportLayer(sslFactory, id, key,
+                peerHost(key), metadataRegistry);
+            Supplier authenticatorCreator = () ->
+                new SslAuthenticator(configs, transportLayer, listenerName, sslPrincipalMapper);
+            return new KafkaChannel(id, transportLayer, authenticatorCreator, maxReceiveSize,
+                    memoryPool != null ? memoryPool : MemoryPool.NONE, metadataRegistry);
         } catch (Exception e) {
             log.info("Failed to create channel due to ", e);
             throw new KafkaException(e);
@@ -68,11 +116,15 @@ public KafkaChannel buildChannel(String id, SelectionKey key, int maxReceiveSize
     }
 
     @Override
-    public void close() {}
+    public void close() {
+        if (sslFactory != null) sslFactory.close();
+    }
 
-    protected SslTransportLayer buildTransportLayer(SslFactory sslFactory, String id, SelectionKey key, String host) throws IOException {
+    protected SslTransportLayer buildTransportLayer(SslFactory sslFactory, String id, SelectionKey key,
+                                                    String host, ChannelMetadataRegistry metadataRegistry) throws IOException {
         SocketChannel socketChannel = (SocketChannel) key.channel();
-        return SslTransportLayer.create(id, key, sslFactory.createSslEngine(host, socketChannel.socket().getPort()));
+        return SslTransportLayer.create(id, key, sslFactory.createSslEngine(host, socketChannel.socket().getPort()),
+            metadataRegistry);
     }
 
     /**
@@ -121,16 +173,18 @@ private String peerHost(SelectionKey key) {
     private static class SslAuthenticator implements Authenticator {
         private final SslTransportLayer transportLayer;
         private final KafkaPrincipalBuilder principalBuilder;
+        private final ListenerName listenerName;
 
-        private SslAuthenticator(Map configs, SslTransportLayer transportLayer) {
+        private SslAuthenticator(Map configs, SslTransportLayer transportLayer, ListenerName listenerName, SslPrincipalMapper sslPrincipalMapper) {
             this.transportLayer = transportLayer;
-            this.principalBuilder = ChannelBuilders.createPrincipalBuilder(configs, transportLayer, this, null);
+            this.principalBuilder = ChannelBuilders.createPrincipalBuilder(configs, transportLayer, this, null, sslPrincipalMapper);
+            this.listenerName = listenerName;
         }
         /**
          * No-Op for plaintext authenticator
          */
         @Override
-        public void authenticate() throws IOException {}
+        public void authenticate() {}
 
         /**
          * Constructs Principal using configured principalBuilder.
@@ -139,10 +193,21 @@ public void authenticate() throws IOException {}
         @Override
         public KafkaPrincipal principal() {
             InetAddress clientAddress = transportLayer.socketChannel().socket().getInetAddress();
-            SslAuthenticationContext context = new SslAuthenticationContext(transportLayer.sslSession(), clientAddress);
+            // listenerName should only be null in Client mode where principal() should not be called
+            if (listenerName == null)
+                throw new IllegalStateException("Unexpected call to principal() when listenerName is null");
+            SslAuthenticationContext context = new SslAuthenticationContext(
+                    transportLayer.sslSession(),
+                    clientAddress,
+                    listenerName.value());
             return principalBuilder.build(context);
         }
 
+        @Override
+        public Optional principalSerde() {
+            return principalBuilder instanceof KafkaPrincipalSerde ? Optional.of((KafkaPrincipalSerde) principalBuilder) : Optional.empty();
+        }
+
         @Override
         public void close() throws IOException {
             if (principalBuilder instanceof Closeable)
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 51ebbc1009cf2..b9879ad6da240 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
@@ -31,25 +31,43 @@
 import javax.net.ssl.SSLEngineResult.Status;
 import javax.net.ssl.SSLException;
 import javax.net.ssl.SSLHandshakeException;
-import javax.net.ssl.SSLSession;
+import javax.net.ssl.SSLKeyException;
 import javax.net.ssl.SSLPeerUnverifiedException;
+import javax.net.ssl.SSLProtocolException;
+import javax.net.ssl.SSLSession;
 
 import org.apache.kafka.common.errors.SslAuthenticationException;
 import org.apache.kafka.common.security.auth.KafkaPrincipal;
+import org.apache.kafka.common.utils.ByteUtils;
+import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.common.utils.ByteBufferUnmapper;
 import org.apache.kafka.common.utils.Utils;
 import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
 /*
  * Transport layer for SSL communication
+ *
+ *
+ * TLS v1.3 notes:
+ *   https://tools.ietf.org/html/rfc8446#section-4.6 : Post-Handshake Messages
+ *   "TLS also allows other messages to be sent after the main handshake.
+ *   These messages use a handshake content type and are encrypted under
+ *   the appropriate application traffic key."
  */
 public class SslTransportLayer implements TransportLayer {
-    private static final Logger log = LoggerFactory.getLogger(SslTransportLayer.class);
-
     private enum State {
+        // Initial state
+        NOT_INITIALIZED,
+        // SSLEngine is in handshake mode
         HANDSHAKE,
+        // SSL handshake failed, connection will be terminated
         HANDSHAKE_FAILED,
+        // SSLEngine has completed handshake, post-handshake messages may be pending for TLSv1.3
+        POST_HANDSHAKE,
+        // SSLEngine has completed handshake, any post-handshake messages have been processed for TLSv1.3
+        // For TLSv1.3, we move the channel to READY state when incoming data is processed after handshake
         READY,
+        // Channel is being closed
         CLOSING
     }
 
@@ -57,6 +75,8 @@ private enum State {
     private final SSLEngine sslEngine;
     private final SelectionKey key;
     private final SocketChannel socketChannel;
+    private final ChannelMetadataRegistry metadataRegistry;
+    private final Logger log;
 
     private HandshakeStatus handshakeStatus;
     private SSLEngineResult handshakeResult;
@@ -65,36 +85,39 @@ private enum State {
     private ByteBuffer netReadBuffer;
     private ByteBuffer netWriteBuffer;
     private ByteBuffer appReadBuffer;
-    private ByteBuffer emptyBuf = ByteBuffer.allocate(0);
+    private ByteBuffer fileChannelBuffer;
+    private boolean hasBytesBuffered;
 
-    public static SslTransportLayer create(String channelId, SelectionKey key, SSLEngine sslEngine) throws IOException {
-        SslTransportLayer transportLayer = new SslTransportLayer(channelId, key, sslEngine);
-        transportLayer.startHandshake();
-        return transportLayer;
+    public static SslTransportLayer create(String channelId, SelectionKey key, SSLEngine sslEngine,
+                                           ChannelMetadataRegistry metadataRegistry) throws IOException {
+        return new SslTransportLayer(channelId, key, sslEngine, metadataRegistry);
     }
 
     // Prefer `create`, only use this in tests
-    SslTransportLayer(String channelId, SelectionKey key, SSLEngine sslEngine) throws IOException {
+    SslTransportLayer(String channelId, SelectionKey key, SSLEngine sslEngine,
+                      ChannelMetadataRegistry metadataRegistry) {
         this.channelId = channelId;
         this.key = key;
         this.socketChannel = (SocketChannel) key.channel();
         this.sslEngine = sslEngine;
+        this.state = State.NOT_INITIALIZED;
+        this.metadataRegistry = metadataRegistry;
+
+        final LogContext logContext = new LogContext(String.format("[SslTransportLayer channelId=%s key=%s] ", channelId, key));
+        this.log = logContext.logger(getClass());
     }
 
     // Visible for testing
     protected void startHandshake() throws IOException {
-        if (state != null)
+        if (state != State.NOT_INITIALIZED)
             throw new IllegalStateException("startHandshake() can only be called once, state " + state);
 
         this.netReadBuffer = ByteBuffer.allocate(netReadBufferSize());
         this.netWriteBuffer = ByteBuffer.allocate(netWriteBufferSize());
         this.appReadBuffer = ByteBuffer.allocate(applicationBufferSize());
-        
-        //clear & set netRead & netWrite buffers
-        netWriteBuffer.position(0);
         netWriteBuffer.limit(0);
-        netReadBuffer.position(0);
         netReadBuffer.limit(0);
+
         state = State.HANDSHAKE;
         //initiate handshake
         sslEngine.beginHandshake();
@@ -103,7 +126,7 @@ protected void startHandshake() throws IOException {
 
     @Override
     public boolean ready() {
-        return state == State.READY;
+        return state == State.POST_HANDSHAKE || state == State.READY;
     }
 
     /**
@@ -130,6 +153,11 @@ public SocketChannel socketChannel() {
         return socketChannel;
     }
 
+    @Override
+    public SelectionKey selectionKey() {
+        return key;
+    }
+
     @Override
     public boolean isOpen() {
         return socketChannel.isOpen();
@@ -140,24 +168,24 @@ public boolean isConnected() {
         return socketChannel.isConnected();
     }
 
-
     /**
-    * Sends a SSL close message and closes socketChannel.
+    * Sends an SSL close message and closes socketChannel.
     */
     @Override
     public void close() throws IOException {
+        State prevState = state;
         if (state == State.CLOSING) return;
         state = State.CLOSING;
         sslEngine.closeOutbound();
         try {
-            if (isConnected()) {
+            if (prevState != State.NOT_INITIALIZED && isConnected()) {
                 if (!flush(netWriteBuffer)) {
                     throw new IOException("Remaining data in the network buffer, can't send SSL close message.");
                 }
                 //prep the buffer for the close message
                 netWriteBuffer.clear();
                 //perform the close, since we called sslEngine.closeOutbound
-                SSLEngineResult wrapResult = sslEngine.wrap(emptyBuf, netWriteBuffer);
+                SSLEngineResult wrapResult = sslEngine.wrap(ByteUtils.EMPTY_BUF, netWriteBuffer);
                 //we should be in a close state
                 if (wrapResult.getStatus() != SSLEngineResult.Status.CLOSED) {
                     throw new IOException("Unexpected status returned by SSLEngine.wrap, expected CLOSED, received " +
@@ -167,14 +195,16 @@ public void close() throws IOException {
                 flush(netWriteBuffer);
             }
         } catch (IOException ie) {
-            log.warn("Failed to send SSL Close message ", ie);
+            log.debug("Failed to send SSL Close message", ie);
         } finally {
-            try {
-                socketChannel.socket().close();
-                socketChannel.close();
-            } finally {
-                key.attach(null);
-                key.cancel();
+            socketChannel.socket().close();
+            socketChannel.close();
+            netReadBuffer = null;
+            netWriteBuffer = null;
+            appReadBuffer = null;
+            if (fileChannelBuffer != null) {
+                ByteBufferUnmapper.unmap("fileChannelBuffer", fileChannelBuffer);
+                fileChannelBuffer = null;
             }
         }
     }
@@ -237,34 +267,45 @@ protected boolean flush(ByteBuffer buf) throws IOException {
     */
     @Override
     public void handshake() throws IOException {
-        if (state == State.READY)
+        if (state == State.NOT_INITIALIZED) {
+            try {
+                startHandshake();
+            } catch (SSLException e) {
+                maybeProcessHandshakeFailure(e, false, null);
+            }
+        }
+        if (ready())
             throw renegotiationException();
         if (state == State.CLOSING)
             throw closingException();
 
         int read = 0;
+        boolean readable = key.isReadable();
         try {
             // Read any available bytes before attempting any writes to ensure that handshake failures
             // reported by the peer are processed even if writes fail (since peer closes connection
             // if handshake fails)
-            if (key.isReadable())
+            if (readable)
                 read = readFromSocketChannel();
 
             doHandshake();
+            if (ready())
+                updateBytesBuffered(true);
         } catch (SSLException e) {
-            handshakeFailure(e, true);
+            maybeProcessHandshakeFailure(e, true, null);
         } catch (IOException e) {
             maybeThrowSslAuthenticationException();
 
-            // this exception could be due to a write. If there is data available to unwrap,
-            // process the data so that any SSLExceptions are reported
-            if (handshakeStatus == HandshakeStatus.NEED_UNWRAP && netReadBuffer.position() > 0) {
-                try {
-                    handshakeUnwrap(false);
-                } catch (SSLException e1) {
-                    handshakeFailure(e1, false);
-                }
+            // This exception could be due to a write. If there is data available to unwrap in the buffer, or data available
+            // in the socket channel to read and unwrap, process the data so that any SSL handshake exceptions are reported.
+            try {
+                do {
+                    handshakeUnwrap(false, true);
+                } while (readable && readFromSocketChannel() > 0);
+            } catch (SSLException e1) {
+                maybeProcessHandshakeFailure(e1, false, e);
             }
+
             // If we get here, this is not a handshake failure, throw the original IOException
             throw e;
         }
@@ -276,6 +317,7 @@ public void handshake() throws IOException {
         }
     }
 
+    @SuppressWarnings("fallthrough")
     private void doHandshake() throws IOException {
         boolean read = key.isReadable();
         boolean write = key.isWritable();
@@ -323,7 +365,7 @@ private void doHandshake() throws IOException {
                 log.trace("SSLHandshake NEED_UNWRAP channelId {}, appReadBuffer pos {}, netReadBuffer pos {}, netWriteBuffer pos {}",
                           channelId, appReadBuffer.position(), netReadBuffer.position(), netWriteBuffer.position());
                 do {
-                    handshakeResult = handshakeUnwrap(read);
+                    handshakeResult = handshakeUnwrap(read, false);
                     if (handshakeResult.getStatus() == Status.BUFFER_OVERFLOW) {
                         int currentAppBufferSize = applicationBufferSize();
                         appReadBuffer = Utils.ensureCapacity(appReadBuffer, currentAppBufferSize);
@@ -367,7 +409,7 @@ private void doHandshake() throws IOException {
         }
     }
 
-    private SSLHandshakeException renegotiationException() throws IOException {
+    private SSLHandshakeException renegotiationException() {
         return new SSLHandshakeException("Renegotiation is not supported");
     }
 
@@ -404,11 +446,13 @@ private void handshakeFinished() throws IOException {
             if (netWriteBuffer.hasRemaining())
                 key.interestOps(key.interestOps() | SelectionKey.OP_WRITE);
             else {
-                state = State.READY;
+                state = sslEngine.getSession().getProtocol().equals("TLSv1.3") ? State.POST_HANDSHAKE : State.READY;
                 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());
+                metadataRegistry.registerCipherInformation(
+                    new CipherInformation(session.getCipherSuite(),  session.getProtocol()));
             }
 
             log.trace("SSLHandshake FINISHED channelId {}, appReadBuffer pos {}, netReadBuffer pos {}, netWriteBuffer pos {} ",
@@ -431,7 +475,7 @@ private SSLEngineResult handshakeWrap(boolean doWrite) throws IOException {
         //this should never be called with a network buffer that contains data
         //so we can clear it here.
         netWriteBuffer.clear();
-        SSLEngineResult result = sslEngine.wrap(emptyBuf, netWriteBuffer);
+        SSLEngineResult result = sslEngine.wrap(ByteUtils.EMPTY_BUF, netWriteBuffer);
         //prepare the results to be written
         netWriteBuffer.flip();
         handshakeStatus = result.getHandshakeStatus();
@@ -445,12 +489,13 @@ private SSLEngineResult handshakeWrap(boolean doWrite) throws IOException {
     }
 
     /**
-    * Perform handshake unwrap
-    * @param doRead boolean
-    * @return SSLEngineResult
-    * @throws IOException
-    */
-    private SSLEngineResult handshakeUnwrap(boolean doRead) throws IOException {
+     * Perform handshake unwrap
+     * @param doRead boolean If true, read more from the socket channel
+     * @param ignoreHandshakeStatus If true, continue to unwrap if data available regardless of handshake status
+     * @return SSLEngineResult
+     * @throws IOException
+     */
+    private SSLEngineResult handshakeUnwrap(boolean doRead, boolean ignoreHandshakeStatus) throws IOException {
         log.trace("SSLHandshake handshakeUnwrap {}", channelId);
         SSLEngineResult result;
         int read = 0;
@@ -459,6 +504,7 @@ private SSLEngineResult handshakeUnwrap(boolean doRead) throws IOException {
         boolean cont;
         do {
             //prepare the buffer with the incoming data
+            int position = netReadBuffer.position();
             netReadBuffer.flip();
             result = sslEngine.unwrap(netReadBuffer, appReadBuffer);
             netReadBuffer.compact();
@@ -467,8 +513,9 @@ private SSLEngineResult handshakeUnwrap(boolean doRead) throws IOException {
                 result.getHandshakeStatus() == HandshakeStatus.NEED_TASK) {
                 handshakeStatus = runDelegatedTasks();
             }
-            cont = result.getStatus() == SSLEngineResult.Status.OK &&
-                handshakeStatus == HandshakeStatus.NEED_UNWRAP;
+            cont = (result.getStatus() == SSLEngineResult.Status.OK &&
+                    handshakeStatus == HandshakeStatus.NEED_UNWRAP) ||
+                    (ignoreHandshakeStatus && netReadBuffer.position() != position);
             log.trace("SSLHandshake handshakeUnwrap: handshakeStatus {} status {}", handshakeStatus, result.getStatus());
         } while (netReadBuffer.position() != 0 && cont);
 
@@ -482,7 +529,8 @@ private SSLEngineResult handshakeUnwrap(boolean doRead) throws IOException {
 
 
     /**
-    * Reads a sequence of bytes from this channel into the given buffer.
+    * Reads a sequence of bytes from this channel into the given buffer. Reads as much as possible
+    * until either the dst buffer is full or there is no more data in the socket.
     *
     * @param dst The buffer into which bytes are to be transferred
     * @return The number of bytes read, possible zero or -1 if the channel has reached end-of-stream
@@ -492,7 +540,7 @@ private SSLEngineResult handshakeUnwrap(boolean doRead) throws IOException {
     @Override
     public int read(ByteBuffer dst) throws IOException {
         if (state == State.CLOSING) return -1;
-        else if (state != State.READY) return 0;
+        else if (!ready()) return 0;
 
         //if we have unread decrypted data in appReadBuffer read that into dst buffer.
         int read = 0;
@@ -500,21 +548,43 @@ public int read(ByteBuffer dst) throws IOException {
             read = readFromAppBuffer(dst);
         }
 
-        int netread = 0;
-        if (dst.remaining() > 0) {
+        boolean readFromNetwork = false;
+        boolean isClosed = false;
+        // Each loop reads at most once from the socket.
+        while (dst.remaining() > 0) {
+            int netread = 0;
             netReadBuffer = Utils.ensureCapacity(netReadBuffer, netReadBufferSize());
-            if (netReadBuffer.remaining() > 0)
+            if (netReadBuffer.remaining() > 0) {
                 netread = readFromSocketChannel();
+                if (netread > 0)
+                    readFromNetwork = true;
+            }
 
             while (netReadBuffer.position() > 0) {
                 netReadBuffer.flip();
-                SSLEngineResult unwrapResult = sslEngine.unwrap(netReadBuffer, appReadBuffer);
+                SSLEngineResult unwrapResult;
+                try {
+                    unwrapResult = sslEngine.unwrap(netReadBuffer, appReadBuffer);
+                    if (state == State.POST_HANDSHAKE && appReadBuffer.position() != 0) {
+                        // For TLSv1.3, we have finished processing post-handshake messages since we are now processing data
+                        state = State.READY;
+                    }
+                } catch (SSLException e) {
+                    // For TLSv1.3, handle SSL exceptions while processing post-handshake messages as authentication exceptions
+                    if (state == State.POST_HANDSHAKE) {
+                        state = State.HANDSHAKE_FAILED;
+                        throw new SslAuthenticationException("Failed to process post-handshake messages", e);
+                    } else
+                        throw e;
+                }
                 netReadBuffer.compact();
                 // handle ssl renegotiation.
-                if (unwrapResult.getHandshakeStatus() != HandshakeStatus.NOT_HANDSHAKING && unwrapResult.getStatus() == Status.OK) {
-                    log.trace("Renegotiation requested, but it is not supported, channelId {}, " +
-                        "appReadBuffer pos {}, netReadBuffer pos {}, netWriteBuffer pos {}", channelId,
-                        appReadBuffer.position(), netReadBuffer.position(), netWriteBuffer.position());
+                if (unwrapResult.getHandshakeStatus() != HandshakeStatus.NOT_HANDSHAKING &&
+                        unwrapResult.getHandshakeStatus() != HandshakeStatus.FINISHED &&
+                        unwrapResult.getStatus() == Status.OK) {
+                    log.error("Renegotiation requested, but it is not supported, channelId {}, " +
+                        "appReadBuffer pos {}, netReadBuffer pos {}, netWriteBuffer pos {} handshakeStatus {}", channelId,
+                        appReadBuffer.position(), netReadBuffer.position(), netWriteBuffer.position(), unwrapResult.getHandshakeStatus());
                     throw renegotiationException();
                 }
 
@@ -547,15 +617,20 @@ public int read(ByteBuffer dst) throws IOException {
                     // If data has been read and unwrapped, return the data. Close will be handled on the next poll.
                     if (appReadBuffer.position() == 0 && read == 0)
                         throw new EOFException();
-                    else
+                    else {
+                        isClosed = true;
                         break;
+                    }
                 }
             }
+            if (read == 0 && netread < 0)
+                throw new EOFException("EOF during read");
+            if (netread <= 0 || isClosed)
+                break;
         }
+        updateBytesBuffered(readFromNetwork || read > 0);
         // If data has been read and unwrapped, return the data even if end-of-stream, channel will be closed
         // on a subsequent poll.
-        if (read == 0 && netread < 0)
-            throw new EOFException("EOF during read");
         return read;
     }
 
@@ -608,42 +683,37 @@ public long read(ByteBuffer[] dsts, int offset, int length) throws IOException {
     * Writes a sequence of bytes to this channel from the given buffer.
     *
     * @param src The buffer from which bytes are to be retrieved
-    * @return The number of bytes read, possibly zero, or -1 if the channel has reached end-of-stream
+    * @return The number of bytes read from src, possibly zero, or -1 if the channel has reached end-of-stream
     * @throws IOException If some other I/O error occurs
     */
     @Override
     public int write(ByteBuffer src) throws IOException {
-        int written = 0;
         if (state == State.CLOSING)
             throw closingException();
-        if (state != State.READY)
-            return written;
-
-        if (!flush(netWriteBuffer))
-            return written;
-
-        netWriteBuffer.clear();
-        SSLEngineResult wrapResult = sslEngine.wrap(src, netWriteBuffer);
-        netWriteBuffer.flip();
+        if (!ready())
+            return 0;
 
-        //handle ssl renegotiation
-        if (wrapResult.getHandshakeStatus() != HandshakeStatus.NOT_HANDSHAKING && wrapResult.getStatus() == Status.OK)
-            throw renegotiationException();
-
-        if (wrapResult.getStatus() == Status.OK) {
-            written = wrapResult.bytesConsumed();
-            flush(netWriteBuffer);
-        } else if (wrapResult.getStatus() == Status.BUFFER_OVERFLOW) {
-            int currentNetWriteBufferSize = netWriteBufferSize();
-            netWriteBuffer.compact();
-            netWriteBuffer = Utils.ensureCapacity(netWriteBuffer, currentNetWriteBufferSize);
+        int written = 0;
+        while (flush(netWriteBuffer) && src.hasRemaining()) {
+            netWriteBuffer.clear();
+            SSLEngineResult wrapResult = sslEngine.wrap(src, netWriteBuffer);
             netWriteBuffer.flip();
-            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) {
-            throw new EOFException();
+
+            //handle ssl renegotiation
+            if (wrapResult.getHandshakeStatus() != HandshakeStatus.NOT_HANDSHAKING && wrapResult.getStatus() == Status.OK)
+                throw renegotiationException();
+
+            if (wrapResult.getStatus() == Status.OK) {
+                written += wrapResult.bytesConsumed();
+            } else if (wrapResult.getStatus() == Status.BUFFER_OVERFLOW) {
+                // BUFFER_OVERFLOW means that the last `wrap` call had no effect, so we expand the buffer and try again
+                netWriteBuffer = Utils.ensureCapacity(netWriteBuffer, netWriteBufferSize());
+                netWriteBuffer.position(netWriteBuffer.limit());
+            } else if (wrapResult.getStatus() == Status.BUFFER_UNDERFLOW) {
+                throw new IllegalStateException("SSL BUFFER_UNDERFLOW during write");
+            } else if (wrapResult.getStatus() == Status.CLOSED) {
+                throw new EOFException();
+            }
         }
         return written;
     }
@@ -698,7 +768,7 @@ public long write(ByteBuffer[] srcs) throws IOException {
      * SSLSession's peerPrincipal for the remote host.
      * @return Principal
      */
-    public Principal peerPrincipal() throws IOException {
+    public Principal peerPrincipal() {
         try {
             return sslEngine.getSession().getPeerPrincipal();
         } catch (SSLPeerUnverifiedException se) {
@@ -708,7 +778,7 @@ public Principal peerPrincipal() throws IOException {
     }
 
     /**
-     * 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 {
@@ -723,7 +793,7 @@ public SSLSession sslSession() throws IllegalStateException {
     public void addInterestOps(int ops) {
         if (!key.isValid())
             throw new CancelledKeyException();
-        else if (state != State.READY)
+        else if (!ready())
             throw new IllegalStateException("handshake is not completed");
 
         key.interestOps(key.interestOps() | ops);
@@ -737,7 +807,7 @@ else if (state != State.READY)
     public void removeInterestOps(int ops) {
         if (!key.isValid())
             throw new CancelledKeyException();
-        else if (state != State.READY)
+        else if (!ready())
             throw new IllegalStateException("handshake is not completed");
 
         key.interestOps(key.interestOps() & ~ops);
@@ -771,7 +841,7 @@ private int readFromAppBuffer(ByteBuffer dst) {
     protected int netReadBufferSize() {
         return sslEngine.getSession().getPacketBufferSize();
     }
-    
+
     protected int netWriteBufferSize() {
         return sslEngine.getSession().getPacketBufferSize();
     }
@@ -779,11 +849,16 @@ protected int netWriteBufferSize() {
     protected int applicationBufferSize() {
         return sslEngine.getSession().getApplicationBufferSize();
     }
-    
+
     protected ByteBuffer netReadBuffer() {
         return netReadBuffer;
     }
 
+    // Visibility for testing
+    protected ByteBuffer appReadBuffer() {
+        return appReadBuffer;
+    }
+
     /**
      * SSL exceptions are propagated as authentication failures so that clients can avoid
      * retries and report the failure. If `flush` is true, exceptions are propagated after
@@ -800,8 +875,44 @@ private void handshakeFailure(SSLException sslException, boolean flush) throws I
 
         state = State.HANDSHAKE_FAILED;
         handshakeException = new SslAuthenticationException("SSL handshake failed", sslException);
-        if (!flush || flush(netWriteBuffer))
+
+        // Attempt to flush any outgoing bytes. If flush doesn't complete, delay exception handling until outgoing bytes
+        // are flushed. If write fails because remote end has closed the channel, log the I/O exception and  continue to
+        // handle the handshake failure as an authentication exception.
+        try {
+            if (!flush || flush(netWriteBuffer))
+                throw handshakeException;
+        } catch (IOException e) {
+            log.debug("Failed to flush all bytes before closing channel", e);
             throw handshakeException;
+        }
+    }
+
+    // SSL handshake failures are typically thrown as SSLHandshakeException, SSLProtocolException,
+    // SSLPeerUnverifiedException or SSLKeyException if the cause is known. These exceptions indicate
+    // authentication failures (e.g. configuration errors) which should not be retried. But the SSL engine
+    // may also throw exceptions using the base class SSLException in a few cases:
+    //   a) If there are no matching ciphers or TLS version or the private key is invalid, client will be
+    //      unable to process the server message and an SSLException is thrown:
+    //      javax.net.ssl.SSLException: Unrecognized SSL message, plaintext connection?
+    //   b) If server closes the connection gracefully during handshake, client may receive close_notify
+    //      and and an SSLException is thrown:
+    //      javax.net.ssl.SSLException: Received close_notify during handshake
+    // We want to handle a) as a non-retriable SslAuthenticationException and b) as a retriable IOException.
+    // To do this we need to rely on the exception string. Since it is safer to throw a retriable exception
+    // when we are not sure, we will treat only the first exception string as a handshake exception.
+    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("Received fatal alert: "))
+            handshakeFailure(sslException, flush);
+        else if (ioException == null)
+            throw sslException;
+        else {
+            log.debug("SSLException while unwrapping data after IOException, original IOException will be propagated", sslException);
+            throw ioException;
+        }
     }
 
     // If handshake has already failed, throw the authentication exception.
@@ -817,12 +928,79 @@ public boolean isMute() {
 
     @Override
     public boolean hasBytesBuffered() {
-        return netReadBuffer.position() != 0 || appReadBuffer.position() != 0;
+        return hasBytesBuffered;
+    }
+
+    // Update `hasBytesBuffered` status. If any bytes were read from the network or
+    // if data was returned from read, `hasBytesBuffered` is set to true if any buffered
+    // data is still remaining. If not, `hasBytesBuffered` is set to false since no progress
+    // can be made until more data is available to read from the network.
+    private void updateBytesBuffered(boolean madeProgress) {
+        if (madeProgress)
+            hasBytesBuffered = netReadBuffer.position() != 0 || appReadBuffer.position() != 0;
+        else
+            hasBytesBuffered = false;
     }
 
     @Override
     public long transferFrom(FileChannel fileChannel, long position, long count) throws IOException {
-        return fileChannel.transferTo(position, count, this);
-    }
+        if (state == State.CLOSING)
+            throw closingException();
+        if (state != State.READY)
+            return 0;
+
+        if (!flush(netWriteBuffer))
+            return 0;
+
+        long channelSize = fileChannel.size();
+        if (position > channelSize)
+            return 0;
+        int totalBytesToWrite = (int) Math.min(Math.min(count, channelSize - position), Integer.MAX_VALUE);
+
+        if (fileChannelBuffer == null) {
+            // Pick a size that allows for reasonably efficient disk reads, keeps the memory overhead per connection
+            // manageable and can typically be drained in a single `write` call. The `netWriteBuffer` is typically 16k
+            // and the socket send buffer is 100k by default, so 32k is a good number given the mentioned trade-offs.
+            int transferSize = 32768;
+            // Allocate a direct buffer to avoid one heap to heap buffer copy. SSLEngine copies the source
+            // buffer (fileChannelBuffer) to the destination buffer (netWriteBuffer) and then encrypts in-place.
+            // FileChannel.read() to a heap buffer requires a copy from a direct buffer to a heap buffer, which is not
+            // useful here.
+            fileChannelBuffer = ByteBuffer.allocateDirect(transferSize);
+            // The loop below drains any remaining bytes from the buffer before reading from disk, so we ensure there
+            // are no remaining bytes in the empty buffer
+            fileChannelBuffer.position(fileChannelBuffer.limit());
+        }
 
+        int totalBytesWritten = 0;
+        long pos = position;
+        try {
+            while (totalBytesWritten < totalBytesToWrite) {
+                if (!fileChannelBuffer.hasRemaining()) {
+                    fileChannelBuffer.clear();
+                    int bytesRemaining = totalBytesToWrite - totalBytesWritten;
+                    if (bytesRemaining < fileChannelBuffer.limit())
+                        fileChannelBuffer.limit(bytesRemaining);
+                    int bytesRead = fileChannel.read(fileChannelBuffer, pos);
+                    if (bytesRead <= 0)
+                        break;
+                    fileChannelBuffer.flip();
+                }
+                int networkBytesWritten = write(fileChannelBuffer);
+                totalBytesWritten += networkBytesWritten;
+                // In the case of a partial write we only return the written bytes to the caller. As a result, the
+                // `position` passed in the next `transferFrom` call won't include the bytes remaining in
+                // `fileChannelBuffer`. By draining `fileChannelBuffer` first, we ensure we update `pos` before
+                // we invoke `fileChannel.read`.
+                if (fileChannelBuffer.hasRemaining())
+                    break;
+                pos += networkBytesWritten;
+            }
+            return totalBytesWritten;
+        } catch (IOException e) {
+            if (totalBytesWritten > 0)
+                return totalBytesWritten;
+            throw e;
+        }
+    }
 }
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 23f866b236f88..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
@@ -25,6 +25,7 @@
  */
 import java.io.IOException;
 import java.nio.channels.FileChannel;
+import java.nio.channels.SelectionKey;
 import java.nio.channels.SocketChannel;
 import java.nio.channels.ScatteringByteChannel;
 import java.nio.channels.GatheringByteChannel;
@@ -60,13 +61,16 @@ public interface TransportLayer extends ScatteringByteChannel, GatheringByteChan
      */
     SocketChannel socketChannel();
 
+    /**
+     * Get the underlying selection key
+     */
+    SelectionKey selectionKey();
 
     /**
      * This a no-op for the non-secure PLAINTEXT implementation. For SSL, this performs
      * SSL handshake. The SSL handshake includes client authentication if configured using
-     * {@link org.apache.kafka.common.config.SslConfigsSslConfigs#SSL_CLIENT_AUTH_CONFIG}.
-     * @throws AuthenticationException if handshake fails due to an
-     *         {@link javax.net.ssl.SSLExceptionSSLException}.
+     * {@link org.apache.kafka.common.config.SslConfigs#SSL_CLIENT_AUTH_CONFIG}.
+     * @throws AuthenticationException if handshake fails due to an {@link javax.net.ssl.SSLException}.
      * @throws IOException if read or write fails with an I/O error.
     */
     void handshake() throws AuthenticationException, IOException;
@@ -77,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;
@@ -90,6 +94,7 @@ public interface TransportLayer extends ScatteringByteChannel, GatheringByteChan
 
     /**
      * @return true if channel has bytes to be read in any intermediate buffers
+     * which may be processed without reading additional data from the network.
      */
     boolean hasBytesBuffered();
 
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 cf1bff557331b..21e0b0ffe7d46 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,92 +16,140 @@
  */
 package org.apache.kafka.common.protocol;
 
+import org.apache.kafka.common.message.AddOffsetsToTxnRequestData;
+import org.apache.kafka.common.message.AddOffsetsToTxnResponseData;
+import org.apache.kafka.common.message.AddPartitionsToTxnRequestData;
+import org.apache.kafka.common.message.AddPartitionsToTxnResponseData;
+import org.apache.kafka.common.message.AlterClientQuotasRequestData;
+import org.apache.kafka.common.message.AlterClientQuotasResponseData;
+import org.apache.kafka.common.message.AlterConfigsRequestData;
+import org.apache.kafka.common.message.AlterConfigsResponseData;
+import org.apache.kafka.common.message.AlterIsrRequestData;
+import org.apache.kafka.common.message.AlterIsrResponseData;
+import org.apache.kafka.common.message.AlterPartitionReassignmentsRequestData;
+import org.apache.kafka.common.message.AlterPartitionReassignmentsResponseData;
+import org.apache.kafka.common.message.AlterReplicaLogDirsRequestData;
+import org.apache.kafka.common.message.AlterReplicaLogDirsResponseData;
+import org.apache.kafka.common.message.AlterUserScramCredentialsRequestData;
+import org.apache.kafka.common.message.AlterUserScramCredentialsResponseData;
+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.BeginQuorumEpochRequestData;
+import org.apache.kafka.common.message.BeginQuorumEpochResponseData;
+import org.apache.kafka.common.message.ControlledShutdownRequestData;
+import org.apache.kafka.common.message.ControlledShutdownResponseData;
+import org.apache.kafka.common.message.CreateAclsRequestData;
+import org.apache.kafka.common.message.CreateAclsResponseData;
+import org.apache.kafka.common.message.CreateDelegationTokenRequestData;
+import org.apache.kafka.common.message.CreateDelegationTokenResponseData;
+import org.apache.kafka.common.message.CreatePartitionsRequestData;
+import org.apache.kafka.common.message.CreatePartitionsResponseData;
+import org.apache.kafka.common.message.CreateTopicsRequestData;
+import org.apache.kafka.common.message.CreateTopicsResponseData;
+import org.apache.kafka.common.message.DeleteAclsRequestData;
+import org.apache.kafka.common.message.DeleteAclsResponseData;
+import org.apache.kafka.common.message.DeleteGroupsRequestData;
+import org.apache.kafka.common.message.DeleteGroupsResponseData;
+import org.apache.kafka.common.message.DeleteRecordsRequestData;
+import org.apache.kafka.common.message.DeleteRecordsResponseData;
+import org.apache.kafka.common.message.DeleteTopicsRequestData;
+import org.apache.kafka.common.message.DeleteTopicsResponseData;
+import org.apache.kafka.common.message.DescribeAclsRequestData;
+import org.apache.kafka.common.message.DescribeAclsResponseData;
+import org.apache.kafka.common.message.DescribeClientQuotasRequestData;
+import org.apache.kafka.common.message.DescribeClientQuotasResponseData;
+import org.apache.kafka.common.message.DescribeConfigsRequestData;
+import org.apache.kafka.common.message.DescribeConfigsResponseData;
+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.DescribeLogDirsRequestData;
+import org.apache.kafka.common.message.DescribeLogDirsResponseData;
+import org.apache.kafka.common.message.DescribeQuorumRequestData;
+import org.apache.kafka.common.message.DescribeQuorumResponseData;
+import org.apache.kafka.common.message.DescribeUserScramCredentialsRequestData;
+import org.apache.kafka.common.message.DescribeUserScramCredentialsResponseData;
+import org.apache.kafka.common.message.ElectLeadersRequestData;
+import org.apache.kafka.common.message.ElectLeadersResponseData;
+import org.apache.kafka.common.message.EndQuorumEpochRequestData;
+import org.apache.kafka.common.message.EndQuorumEpochResponseData;
+import org.apache.kafka.common.message.EndTxnRequestData;
+import org.apache.kafka.common.message.EndTxnResponseData;
+import org.apache.kafka.common.message.EnvelopeRequestData;
+import org.apache.kafka.common.message.EnvelopeResponseData;
+import org.apache.kafka.common.message.ExpireDelegationTokenRequestData;
+import org.apache.kafka.common.message.ExpireDelegationTokenResponseData;
+import org.apache.kafka.common.message.FetchRequestData;
+import org.apache.kafka.common.message.FetchResponseData;
+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.ListOffsetRequestData;
+import org.apache.kafka.common.message.ListOffsetResponseData;
+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.OffsetForLeaderEpochRequestData;
+import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData;
+import org.apache.kafka.common.message.ProduceRequestData;
+import org.apache.kafka.common.message.ProduceResponseData;
+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.TxnOffsetCommitRequestData;
+import org.apache.kafka.common.message.TxnOffsetCommitResponseData;
+import org.apache.kafka.common.message.UpdateFeaturesRequestData;
+import org.apache.kafka.common.message.UpdateFeaturesResponseData;
+import org.apache.kafka.common.message.UpdateMetadataRequestData;
+import org.apache.kafka.common.message.UpdateMetadataResponseData;
+import org.apache.kafka.common.message.VoteRequestData;
+import org.apache.kafka.common.message.VoteResponseData;
+import org.apache.kafka.common.message.WriteTxnMarkersRequestData;
+import org.apache.kafka.common.message.WriteTxnMarkersResponseData;
 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 org.apache.kafka.common.record.RecordBatch;
-import org.apache.kafka.common.requests.AddOffsetsToTxnRequest;
-import org.apache.kafka.common.requests.AddOffsetsToTxnResponse;
-import org.apache.kafka.common.requests.AddPartitionsToTxnRequest;
-import org.apache.kafka.common.requests.AddPartitionsToTxnResponse;
-import org.apache.kafka.common.requests.AlterConfigsRequest;
-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.CreatePartitionsRequest;
-import org.apache.kafka.common.requests.CreatePartitionsResponse;
-import org.apache.kafka.common.requests.CreateTopicsRequest;
-import org.apache.kafka.common.requests.CreateTopicsResponse;
-import org.apache.kafka.common.requests.DeleteAclsRequest;
-import org.apache.kafka.common.requests.DeleteAclsResponse;
-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.DescribeGroupsRequest;
-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.EndTxnRequest;
-import org.apache.kafka.common.requests.EndTxnResponse;
-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.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.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.SaslAuthenticateRequest;
-import org.apache.kafka.common.requests.SaslAuthenticateResponse;
-import org.apache.kafka.common.requests.SaslHandshakeRequest;
-import org.apache.kafka.common.requests.SaslHandshakeResponse;
-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;
 
 import java.nio.ByteBuffer;
+import java.util.Arrays;
+import java.util.List;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.stream.Collectors;
 
 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;
 
@@ -109,29 +157,28 @@
  * Identifiers for all the Kafka APIs
  */
 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()),
-    LEAVE_GROUP(13, "LeaveGroup", LeaveGroupRequest.schemaVersions(), LeaveGroupResponse.schemaVersions()),
-    SYNC_GROUP(14, "SyncGroup", SyncGroupRequest.schemaVersions(), SyncGroupResponse.schemaVersions()),
-    DESCRIBE_GROUPS(15, "DescribeGroups", DescribeGroupsRequest.schemaVersions(),
-            DescribeGroupsResponse.schemaVersions()),
-    LIST_GROUPS(16, "ListGroups", ListGroupsRequest.schemaVersions(), ListGroupsResponse.schemaVersions()),
-    SASL_HANDSHAKE(17, "SaslHandshake", SaslHandshakeRequest.schemaVersions(), SaslHandshakeResponse.schemaVersions()),
-    API_VERSIONS(18, "ApiVersions", ApiVersionsRequest.schemaVersions(), ApiVersionsResponse.schemaVersions()) {
+    PRODUCE(0, "Produce", ProduceRequestData.SCHEMAS, ProduceResponseData.SCHEMAS),
+    FETCH(1, "Fetch", FetchRequestData.SCHEMAS, FetchResponseData.SCHEMAS),
+    LIST_OFFSETS(2, "ListOffsets", ListOffsetRequestData.SCHEMAS, ListOffsetResponseData.SCHEMAS),
+    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", SyncGroupRequestData.SCHEMAS, SyncGroupResponseData.SCHEMAS),
+    DESCRIBE_GROUPS(15, "DescribeGroups", DescribeGroupsRequestData.SCHEMAS,
+            DescribeGroupsResponseData.SCHEMAS),
+    LIST_GROUPS(16, "ListGroups", ListGroupsRequestData.SCHEMAS, ListGroupsResponseData.SCHEMAS),
+    SASL_HANDSHAKE(17, "SaslHandshake", SaslHandshakeRequestData.SCHEMAS, SaslHandshakeResponseData.SCHEMAS),
+    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
@@ -140,38 +187,74 @@ public Struct parseResponse(short version, ByteBuffer buffer) {
             return parseResponse(version, buffer, (short) 0);
         }
     },
-    CREATE_TOPICS(19, "CreateTopics", CreateTopicsRequest.schemaVersions(), CreateTopicsResponse.schemaVersions()),
-    DELETE_TOPICS(20, "DeleteTopics", DeleteTopicsRequest.schemaVersions(), DeleteTopicsResponse.schemaVersions()),
-    DELETE_RECORDS(21, "DeleteRecords", DeleteRecordsRequest.schemaVersions(), DeleteRecordsResponse.schemaVersions()),
-    INIT_PRODUCER_ID(22, "InitProducerId", InitProducerIdRequest.schemaVersions(),
-            InitProducerIdResponse.schemaVersions()),
-    OFFSET_FOR_LEADER_EPOCH(23, "OffsetForLeaderEpoch", true, OffsetsForLeaderEpochRequest.schemaVersions(),
-            OffsetsForLeaderEpochResponse.schemaVersions()),
+    CREATE_TOPICS(19, "CreateTopics", CreateTopicsRequestData.SCHEMAS, CreateTopicsResponseData.SCHEMAS, true),
+    DELETE_TOPICS(20, "DeleteTopics", DeleteTopicsRequestData.SCHEMAS, DeleteTopicsResponseData.SCHEMAS, true),
+    DELETE_RECORDS(21, "DeleteRecords", DeleteRecordsRequestData.SCHEMAS, DeleteRecordsResponseData.SCHEMAS),
+    INIT_PRODUCER_ID(22, "InitProducerId", InitProducerIdRequestData.SCHEMAS, InitProducerIdResponseData.SCHEMAS),
+    OFFSET_FOR_LEADER_EPOCH(23, "OffsetForLeaderEpoch", false, OffsetForLeaderEpochRequestData.SCHEMAS,
+        OffsetForLeaderEpochResponseData.SCHEMAS),
     ADD_PARTITIONS_TO_TXN(24, "AddPartitionsToTxn", false, RecordBatch.MAGIC_VALUE_V2,
-            AddPartitionsToTxnRequest.schemaVersions(), AddPartitionsToTxnResponse.schemaVersions()),
-    ADD_OFFSETS_TO_TXN(25, "AddOffsetsToTxn", false, RecordBatch.MAGIC_VALUE_V2, AddOffsetsToTxnRequest.schemaVersions(),
-            AddOffsetsToTxnResponse.schemaVersions()),
-    END_TXN(26, "EndTxn", false, RecordBatch.MAGIC_VALUE_V2, EndTxnRequest.schemaVersions(),
-            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()),
-    DESCRIBE_ACLS(29, "DescribeAcls", DescribeAclsRequest.schemaVersions(), DescribeAclsResponse.schemaVersions()),
-    CREATE_ACLS(30, "CreateAcls", CreateAclsRequest.schemaVersions(), CreateAclsResponse.schemaVersions()),
-    DELETE_ACLS(31, "DeleteAcls", DeleteAclsRequest.schemaVersions(), DeleteAclsResponse.schemaVersions()),
-    DESCRIBE_CONFIGS(32, "DescribeConfigs", DescribeConfigsRequest.schemaVersions(),
-            DescribeConfigsResponse.schemaVersions()),
-    ALTER_CONFIGS(33, "AlterConfigs", AlterConfigsRequest.schemaVersions(),
-            AlterConfigsResponse.schemaVersions()),
-    ALTER_REPLICA_LOG_DIRS(34, "AlterReplicaLogDirs", AlterReplicaLogDirsRequest.schemaVersions(),
-            AlterReplicaLogDirsResponse.schemaVersions()),
-    DESCRIBE_LOG_DIRS(35, "DescribeLogDirs", DescribeLogDirsRequest.schemaVersions(),
-            DescribeLogDirsResponse.schemaVersions()),
-    SASL_AUTHENTICATE(36, "SaslAuthenticate", SaslAuthenticateRequest.schemaVersions(),
-            SaslAuthenticateResponse.schemaVersions()),
-    CREATE_PARTITIONS(37, "CreatePartitions", CreatePartitionsRequest.schemaVersions(),
-            CreatePartitionsResponse.schemaVersions());
+            AddPartitionsToTxnRequestData.SCHEMAS, AddPartitionsToTxnResponseData.SCHEMAS),
+    ADD_OFFSETS_TO_TXN(25, "AddOffsetsToTxn", false, RecordBatch.MAGIC_VALUE_V2, AddOffsetsToTxnRequestData.SCHEMAS,
+            AddOffsetsToTxnResponseData.SCHEMAS),
+    END_TXN(26, "EndTxn", false, RecordBatch.MAGIC_VALUE_V2, EndTxnRequestData.SCHEMAS, EndTxnResponseData.SCHEMAS),
+    WRITE_TXN_MARKERS(27, "WriteTxnMarkers", true, RecordBatch.MAGIC_VALUE_V2, WriteTxnMarkersRequestData.SCHEMAS,
+            WriteTxnMarkersResponseData.SCHEMAS),
+    TXN_OFFSET_COMMIT(28, "TxnOffsetCommit", false, RecordBatch.MAGIC_VALUE_V2, TxnOffsetCommitRequestData.SCHEMAS,
+            TxnOffsetCommitResponseData.SCHEMAS),
+    DESCRIBE_ACLS(29, "DescribeAcls", DescribeAclsRequestData.SCHEMAS, DescribeAclsResponseData.SCHEMAS),
+    CREATE_ACLS(30, "CreateAcls", CreateAclsRequestData.SCHEMAS, CreateAclsResponseData.SCHEMAS, true),
+    DELETE_ACLS(31, "DeleteAcls", DeleteAclsRequestData.SCHEMAS, DeleteAclsResponseData.SCHEMAS, true),
+    DESCRIBE_CONFIGS(32, "DescribeConfigs", DescribeConfigsRequestData.SCHEMAS,
+             DescribeConfigsResponseData.SCHEMAS),
+    ALTER_CONFIGS(33, "AlterConfigs", AlterConfigsRequestData.SCHEMAS,
+            AlterConfigsResponseData.SCHEMAS, true),
+    ALTER_REPLICA_LOG_DIRS(34, "AlterReplicaLogDirs", AlterReplicaLogDirsRequestData.SCHEMAS,
+            AlterReplicaLogDirsResponseData.SCHEMAS),
+    DESCRIBE_LOG_DIRS(35, "DescribeLogDirs", DescribeLogDirsRequestData.SCHEMAS,
+            DescribeLogDirsResponseData.SCHEMAS),
+    SASL_AUTHENTICATE(36, "SaslAuthenticate", SaslAuthenticateRequestData.SCHEMAS,
+            SaslAuthenticateResponseData.SCHEMAS),
+    CREATE_PARTITIONS(37, "CreatePartitions", CreatePartitionsRequestData.SCHEMAS,
+            CreatePartitionsResponseData.SCHEMAS, true),
+    CREATE_DELEGATION_TOKEN(38, "CreateDelegationToken", CreateDelegationTokenRequestData.SCHEMAS,
+            CreateDelegationTokenResponseData.SCHEMAS, true),
+    RENEW_DELEGATION_TOKEN(39, "RenewDelegationToken", RenewDelegationTokenRequestData.SCHEMAS,
+            RenewDelegationTokenResponseData.SCHEMAS, true),
+    EXPIRE_DELEGATION_TOKEN(40, "ExpireDelegationToken", ExpireDelegationTokenRequestData.SCHEMAS,
+            ExpireDelegationTokenResponseData.SCHEMAS, true),
+    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, true),
+    ALTER_PARTITION_REASSIGNMENTS(45, "AlterPartitionReassignments", AlterPartitionReassignmentsRequestData.SCHEMAS,
+            AlterPartitionReassignmentsResponseData.SCHEMAS, true),
+    LIST_PARTITION_REASSIGNMENTS(46, "ListPartitionReassignments", ListPartitionReassignmentsRequestData.SCHEMAS,
+            ListPartitionReassignmentsResponseData.SCHEMAS),
+    OFFSET_DELETE(47, "OffsetDelete", OffsetDeleteRequestData.SCHEMAS, OffsetDeleteResponseData.SCHEMAS),
+    DESCRIBE_CLIENT_QUOTAS(48, "DescribeClientQuotas", DescribeClientQuotasRequestData.SCHEMAS,
+            DescribeClientQuotasResponseData.SCHEMAS),
+    ALTER_CLIENT_QUOTAS(49, "AlterClientQuotas", AlterClientQuotasRequestData.SCHEMAS,
+            AlterClientQuotasResponseData.SCHEMAS, true),
+    DESCRIBE_USER_SCRAM_CREDENTIALS(50, "DescribeUserScramCredentials", DescribeUserScramCredentialsRequestData.SCHEMAS,
+            DescribeUserScramCredentialsResponseData.SCHEMAS),
+    ALTER_USER_SCRAM_CREDENTIALS(51, "AlterUserScramCredentials", AlterUserScramCredentialsRequestData.SCHEMAS,
+            AlterUserScramCredentialsResponseData.SCHEMAS, true),
+    VOTE(52, "Vote", true, false,
+        VoteRequestData.SCHEMAS, VoteResponseData.SCHEMAS),
+    BEGIN_QUORUM_EPOCH(53, "BeginQuorumEpoch", true, false,
+        BeginQuorumEpochRequestData.SCHEMAS, BeginQuorumEpochResponseData.SCHEMAS),
+    END_QUORUM_EPOCH(54, "EndQuorumEpoch", true, false,
+        EndQuorumEpochRequestData.SCHEMAS, EndQuorumEpochResponseData.SCHEMAS),
+    DESCRIBE_QUORUM(55, "DescribeQuorum", true, false,
+        DescribeQuorumRequestData.SCHEMAS, DescribeQuorumResponseData.SCHEMAS),
+    ALTER_ISR(56, "AlterIsr", true, AlterIsrRequestData.SCHEMAS, AlterIsrResponseData.SCHEMAS),
+    UPDATE_FEATURES(57, "UpdateFeatures",
+        UpdateFeaturesRequestData.SCHEMAS, UpdateFeaturesResponseData.SCHEMAS, true),
+    ENVELOPE(58, "Envelope", true, false, EnvelopeRequestData.SCHEMAS, EnvelopeResponseData.SCHEMAS);
 
     private static final ApiKeys[] ID_TO_TYPE;
     private static final int MIN_API_KEY = 0;
@@ -200,6 +283,12 @@ public Struct parseResponse(short version, ByteBuffer buffer) {
     /** indicates the minimum required inter broker magic required to support the API */
     public final byte minRequiredInterBrokerMagic;
 
+    /** indicates whether the API is enabled and should be exposed in ApiVersions **/
+    public final boolean isEnabled;
+
+    /** indicates whether the API is enabled for forwarding **/
+    public final boolean forwardable;
+
     public final Schema[] requestSchemas;
     public final Schema[] responseSchemas;
     public final boolean requiresDelayedAllocation;
@@ -212,14 +301,36 @@ public Struct parseResponse(short version, ByteBuffer buffer) {
         this(id, name, clusterAction, RecordBatch.MAGIC_VALUE_V0, requestSchemas, responseSchemas);
     }
 
+    ApiKeys(int id, String name, Schema[] requestSchemas, Schema[] responseSchemas, boolean forwardable) {
+        this(id, name, false, RecordBatch.MAGIC_VALUE_V0, true, requestSchemas, responseSchemas, forwardable);
+    }
+
+    ApiKeys(int id, String name, boolean clusterAction, boolean isEnabled, Schema[] requestSchemas, Schema[] responseSchemas) {
+        this(id, name, clusterAction, RecordBatch.MAGIC_VALUE_V0, isEnabled, requestSchemas, responseSchemas, false);
+    }
+
     ApiKeys(int id, String name, boolean clusterAction, byte minRequiredInterBrokerMagic,
             Schema[] requestSchemas, Schema[] responseSchemas) {
+        this(id, name, clusterAction, minRequiredInterBrokerMagic, true, requestSchemas, responseSchemas, false);
+    }
+
+    ApiKeys(
+        int id,
+        String name,
+        boolean clusterAction,
+        byte minRequiredInterBrokerMagic,
+        boolean isEnabled,
+        Schema[] requestSchemas,
+        Schema[] responseSchemas,
+        boolean forwardable
+    ) {
         if (id < 0)
             throw new IllegalArgumentException("id must not be negative, id: " + id);
         this.id = (short) id;
         this.name = name;
         this.clusterAction = clusterAction;
         this.minRequiredInterBrokerMagic = minRequiredInterBrokerMagic;
+        this.isEnabled = isEnabled;
 
         if (requestSchemas.length != responseSchemas.length)
             throw new IllegalStateException(requestSchemas.length + " request versions for api " + name
@@ -232,6 +343,13 @@ public Struct parseResponse(short version, ByteBuffer buffer) {
                 throw new IllegalStateException("Response schema for api " + name + " for version " + i + " is null");
         }
 
+        this.requiresDelayedAllocation = forwardable || shouldRetainsBufferReference(requestSchemas);
+        this.requestSchemas = requestSchemas;
+        this.responseSchemas = responseSchemas;
+        this.forwardable = forwardable;
+    }
+
+    private static boolean shouldRetainsBufferReference(Schema[] requestSchemas) {
         boolean requestRetainsBufferReference = false;
         for (Schema requestVersionSchema : requestSchemas) {
             if (retainsBufferReference(requestVersionSchema)) {
@@ -239,9 +357,7 @@ public Struct parseResponse(short version, ByteBuffer buffer) {
                 break;
             }
         }
-        this.requiresDelayedAllocation = requestRetainsBufferReference;
-        this.requestSchemas = requestSchemas;
-        this.responseSchemas = responseSchemas;
+        return requestRetainsBufferReference;
     }
 
     public static ApiKeys forId(int id) {
@@ -302,6 +418,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");
@@ -309,7 +433,7 @@ private static String toHtml() {
         b.append("\n");
         b.append("\n");
         b.append("");
-        for (ApiKeys key : ApiKeys.values()) {
+        for (ApiKeys key : ApiKeys.enabledApis()) {
             b.append("\n");
             b.append(""
+    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(4, (key: String) => "prefix_" + key, Collections.emptyMap())
+    val expectedConfig = "
        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 = {
@@ -112,7 +174,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 4c1f4ae1bf740..a13bedcdd534d 100755
--- a/core/src/test/scala/unit/kafka/log/LogManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/log/LogManagerTest.scala
@@ -18,16 +18,27 @@
 package kafka.log
 
 import java.io._
-import java.util.Properties
+import java.nio.file.Files
+import java.util.{Collections, Properties}
 
-import kafka.common._
+import com.yammer.metrics.core.MetricName
+import kafka.metrics.KafkaYammerMetrics
+import kafka.server.{FetchDataInfo, FetchLogEnd}
 import kafka.server.checkpoints.OffsetCheckpointFile
 import kafka.utils._
-import org.apache.kafka.common.TopicPartition
+import org.apache.directory.api.util.FileUtils
 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 scala.collection.mutable
+import scala.jdk.CollectionConverters._
+import scala.util.{Failure, Try}
 
 class LogManagerTest {
 
@@ -46,37 +57,145 @@ 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)
     // Some tests assign a new LogManager
-    logManager.liveLogDirs.foreach(Utils.delete)
+    if (logManager != null)
+      logManager.liveLogDirs.foreach(Utils.delete)
+  }
+
+  /**
+   * Test that getOrCreateLog on a non-existent log creates a new log and that we can append to the new log.
+   */
+  @Test
+  def testCreateLog(): Unit = {
+    val log = logManager.getOrCreateLog(new TopicPartition(name, 0), () => logConfig)
+    assertEquals(1, logManager.liveLogDirs.size)
+
+    val logFile = new File(logDir, name + "-0")
+    assertTrue(logFile.exists)
+    log.appendAsLeader(TestUtils.singletonRecords("test".getBytes()), leaderEpoch = 0)
+  }
+
+  /**
+   * Tests that all internal futures are completed before LogManager.shutdown() returns to the
+   * caller during error situations.
+   */
+  @Test
+  def testHandlingExceptionsDuringShutdown(): Unit = {
+    // We create two directories logDir1 and logDir2 to help effectively test error handling
+    // during LogManager.shutdown().
+    val logDir1 = TestUtils.tempDir()
+    val logDir2 = TestUtils.tempDir()
+    var logManagerForTest: Option[LogManager] = Option.empty
+    try {
+      logManagerForTest = Some(createLogManager(Seq(logDir1, logDir2)))
+
+      assertEquals(2, logManagerForTest.get.liveLogDirs.size)
+      logManagerForTest.get.startup()
+
+      val log1 = logManagerForTest.get.getOrCreateLog(new TopicPartition(name, 0), () => logConfig)
+      val log2 = logManagerForTest.get.getOrCreateLog(new TopicPartition(name, 1), () => logConfig)
+
+      val logFile1 = new File(logDir1, name + "-0")
+      assertTrue(logFile1.exists)
+      val logFile2 = new File(logDir2, name + "-1")
+      assertTrue(logFile2.exists)
+
+      log1.appendAsLeader(TestUtils.singletonRecords("test1".getBytes()), leaderEpoch = 0)
+      log1.takeProducerSnapshot()
+      log1.appendAsLeader(TestUtils.singletonRecords("test1".getBytes()), leaderEpoch = 0)
+
+      log2.appendAsLeader(TestUtils.singletonRecords("test2".getBytes()), leaderEpoch = 0)
+      log2.takeProducerSnapshot()
+      log2.appendAsLeader(TestUtils.singletonRecords("test2".getBytes()), leaderEpoch = 0)
+
+      // This should cause log1.close() to fail during LogManger shutdown sequence.
+      FileUtils.deleteDirectory(logFile1)
+
+      logManagerForTest.get.shutdown()
+
+      assertFalse(Files.exists(new File(logDir1, Log.CleanShutdownFile).toPath))
+      assertTrue(Files.exists(new File(logDir2, Log.CleanShutdownFile).toPath))
+    } finally {
+      logManagerForTest.foreach(manager => manager.liveLogDirs.foreach(Utils.delete))
+    }
   }
 
   /**
    * Test that getOrCreateLog on a non-existent log creates a new log and that we can append to the new log.
+   * The LogManager is configured with one invalid log directory which should be marked as offline.
    */
   @Test
-  def testCreateLog() {
-    val log = logManager.getOrCreateLog(new TopicPartition(name, 0), logConfig)
+  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)
+    val dirs = Seq(logDir, new File("\u0000"))
+
+    logManager.shutdown()
+    logManager = createLogManager(dirs)
+    logManager.startup()
+
+    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 { invocation =>
+      // 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")
@@ -87,8 +206,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())
@@ -96,7 +215,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))
 
@@ -104,12 +223,17 @@ class LogManagerTest {
     assertEquals("Now there should only be only one segment in the index.", 1, log.numberOfSegments)
     time.sleep(log.config.fileDeleteDelayMs + 1)
 
+    log.logSegments.foreach(s => {
+      s.lazyOffsetIndex.get
+      s.lazyTimeIndex.get
+    })
+
     // there should be a log file, two indexes, one producer snapshot, and the leader epoch checkpoint
     assertEquals("Files should have been deleted", log.numberOfSegments * 4 + 1, log.dir.list.length)
-    assertEquals("Should get empty fetch off new log.", 0, log.readUncommitted(offset+1, 1024).records.sizeInBytes)
+    assertEquals("Should get empty fetch off new log.", 0, readLog(log, offset + 1).records.sizeInBytes)
 
     try {
-      log.readUncommitted(0, 1024)
+      readLog(log, 0)
       fail("Should get exception from fetching earlier.")
     } catch {
       case _: OffsetOutOfRangeException => // This is good.
@@ -122,7 +246,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()
@@ -134,7 +258,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
@@ -142,10 +266,10 @@ class LogManagerTest {
     for (_ <- 0 until numMessages) {
       val set = TestUtils.singletonRecords("test".getBytes())
       val info = log.appendAsLeader(set, leaderEpoch = 0)
-      offset = info.firstOffset
+      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
@@ -154,11 +278,11 @@ class LogManagerTest {
     time.sleep(log.config.fileDeleteDelayMs + 1)
 
     // there should be a log file, two indexes (the txn index is created lazily),
-    // the leader epoch checkpoint and two producer snapshot files (one for the active and previous segments)
-    assertEquals("Files should have been deleted", log.numberOfSegments * 3 + 3, log.dir.list.length)
-    assertEquals("Should get empty fetch off new log.", 0, log.readUncommitted(offset + 1, 1024).records.sizeInBytes)
+    // and a producer snapshot file per segment, and the leader epoch checkpoint.
+    assertEquals("Files should have been deleted", log.numberOfSegments * 4 + 1, log.dir.list.length)
+    assertEquals("Should get empty fetch off new log.", 0, readLog(log, offset + 1).records.sizeInBytes)
     try {
-      log.readUncommitted(0, 1024)
+      readLog(log, 0)
       fail("Should get exception from fetching earlier.")
     } catch {
       case _: OffsetOutOfRangeException => // This is good.
@@ -168,14 +292,27 @@ class LogManagerTest {
   }
 
   /**
-    * Ensures that LogManager only runs on logs with cleanup.policy=delete
+    * Ensures that LogManager doesn't run on logs with cleanup.policy=compact,delete
+    * LogCleaner.CleanerThread handles all logs where compaction is enabled.
+    */
+  @Test
+  def testDoesntCleanLogsWithCompactDeletePolicy(): Unit = {
+    testDoesntCleanLogs(LogConfig.Compact + "," + LogConfig.Delete)
+  }
+
+  /**
+    * Ensures that LogManager doesn't run on logs with cleanup.policy=compact
     * LogCleaner.CleanerThread handles all logs where compaction is enabled.
     */
   @Test
-  def testDoesntCleanLogsWithCompactDeletePolicy() {
+  def testDoesntCleanLogsWithCompactPolicy(): Unit = {
+    testDoesntCleanLogs(LogConfig.Compact)
+  }
+
+  private def testDoesntCleanLogs(policy: String): Unit = {
     val logProps = new Properties()
-    logProps.put(LogConfig.CleanupPolicyProp, LogConfig.Compact + "," + LogConfig.Delete)
-    val log = logManager.getOrCreateLog(new TopicPartition(name, 0), LogConfig.fromProps(logConfig.originals, logProps))
+    logProps.put(LogConfig.CleanupPolicyProp, policy)
+    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())
@@ -196,7 +333,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)
@@ -204,7 +341,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())
@@ -218,7 +355,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(),
@@ -228,7 +365,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)
@@ -239,7 +376,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")
@@ -252,7 +389,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)
   }
 
@@ -260,7 +397,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()
@@ -271,15 +408,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)
@@ -292,7 +429,6 @@ class LogManagerTest {
 
     topicPartitions.zip(logs).foreach { case (tp, log) =>
       assertEquals("Recovery point should equal checkpoint", checkpoints(tp), log.recoveryPoint)
-      assertEquals(Some(log.minSnapshotsOffsetToRetain), log.oldestProducerSnapshotOffset)
     }
   }
 
@@ -304,23 +440,23 @@ 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.index.file.getName
+    val indexName = activeSegment.offsetIndex.file.getName
     val timeIndexName = activeSegment.timeIndex.file.getName
     val txnIndexName = activeSegment.txnIndex.file.getName
     val indexFilesOnDiskBeforeDelete = activeSegment.log.file.getParentFile.listFiles.filter(_.getName.endsWith("index"))
 
-    val removedLog = logManager.asyncDelete(new TopicPartition(name, 0))
+    val removedLog = logManager.asyncDelete(new TopicPartition(name, 0)).get
     val removedSegment = removedLog.activeSegment
-    val indexFilesAfterDelete = Seq(removedSegment.index.file, removedSegment.timeIndex.file,
+    val indexFilesAfterDelete = Seq(removedSegment.lazyOffsetIndex.file, removedSegment.lazyTimeIndex.file,
       removedSegment.txnIndex.file)
 
     assertEquals(new File(removedLog.dir, logName), removedSegment.log.file)
-    assertEquals(new File(removedLog.dir, indexName), removedSegment.index.file)
-    assertEquals(new File(removedLog.dir, timeIndexName), removedSegment.timeIndex.file)
+    assertEquals(new File(removedLog.dir, indexName), removedSegment.lazyOffsetIndex.file)
+    assertEquals(new File(removedLog.dir, timeIndexName), removedSegment.lazyTimeIndex.file)
     assertEquals(new File(removedLog.dir, txnIndexName), removedSegment.txnIndex.file)
 
     // Try to detect the case where a new index type was added and we forgot to update the pointer
@@ -332,5 +468,216 @@ class LogManagerTest {
       assertNotEquals("File reference was not updated in index", fileBeforeDelete.getAbsolutePath,
         fileInIndex.get.getAbsolutePath)
     }
+
+    time.sleep(logManager.InitialTaskDelayMs)
+    assertTrue("Logs deleted too early", logManager.hasLogsToBeDeleted)
+    time.sleep(logManager.currentDefaultConfig.fileDeleteDelayMs - logManager.InitialTaskDelayMs)
+    assertFalse("Logs not deleted", logManager.hasLogsToBeDeleted)
+  }
+
+  @Test
+  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),
+      new TopicPartition("test-a", 2),
+      new TopicPartition("test-b", 0),
+      new TopicPartition("test-b", 1))
+
+    val allLogs = tps.map(logManager.getOrCreateLog(_, () => logConfig))
+    allLogs.foreach { log =>
+      for (_ <- 0 until 50)
+        log.appendAsLeader(TestUtils.singletonRecords("test".getBytes), leaderEpoch = 0)
+      log.flush()
+    }
+
+    logManager.checkpointRecoveryOffsetsInDir(logDir)
+
+    val checkpoints = new OffsetCheckpointFile(new File(logDir, LogManager.RecoveryPointCheckpointFile)).read()
+
+    tps.zip(allLogs).foreach { case (tp, log) =>
+      assertEquals("Recovery point should equal checkpoint", checkpoints(tp), log.recoveryPoint)
+    }
+  }
+
+  private def readLog(log: Log, offset: Long, maxLength: Int = 1024): FetchDataInfo = {
+    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)
+  }
+
+  @Test
+  def testMetricsExistWhenLogIsRecreatedBeforeDeletion(): Unit = {
+    val topicName = "metric-test"
+    def logMetrics: mutable.Set[MetricName] = KafkaYammerMetrics.defaultRegistry.allMetrics.keySet.asScala.
+      filter(metric => metric.getType == "Log" && metric.getScope.contains(topicName))
+
+    val tp = new TopicPartition(topicName, 0)
+    val metricTag = s"topic=${tp.topic},partition=${tp.partition}"
+
+    def verifyMetrics(): Unit = {
+      assertEquals(LogMetricNames.allMetricNames.size, logMetrics.size)
+      logMetrics.foreach { metric =>
+        assertTrue(metric.getMBeanName.contains(metricTag))
+      }
+    }
+
+    // Create the Log and assert that the metrics are present
+    logManager.getOrCreateLog(tp, () => logConfig)
+    verifyMetrics()
+
+    // Trigger the deletion and assert that the metrics have been removed
+    val removedLog = logManager.asyncDelete(tp).get
+    assertTrue(logMetrics.isEmpty)
+
+    // Recreate the Log and assert that the metrics are present
+    logManager.getOrCreateLog(tp, () => logConfig)
+    verifyMetrics()
+
+    // Advance time past the file deletion delay and assert that the removed log has been deleted but the metrics
+    // are still present
+    time.sleep(logConfig.fileDeleteDelayMs + 1)
+    assertTrue(removedLog.logSegments.isEmpty)
+    verifyMetrics()
+  }
+
+  @Test
+  def testMetricsAreRemovedWhenMovingCurrentToFutureLog(): Unit = {
+    val dir1 = TestUtils.tempDir()
+    val dir2 = TestUtils.tempDir()
+    logManager = createLogManager(Seq(dir1, dir2))
+    logManager.startup()
+
+    val topicName = "future-log"
+    def logMetrics: mutable.Set[MetricName] = KafkaYammerMetrics.defaultRegistry.allMetrics.keySet.asScala.
+      filter(metric => metric.getType == "Log" && metric.getScope.contains(topicName))
+
+    val tp = new TopicPartition(topicName, 0)
+    val metricTag = s"topic=${tp.topic},partition=${tp.partition}"
+
+    def verifyMetrics(logCount: Int): Unit = {
+      assertEquals(LogMetricNames.allMetricNames.size * logCount, logMetrics.size)
+      logMetrics.foreach { metric =>
+        assertTrue(metric.getMBeanName.contains(metricTag))
+      }
+    }
+
+    // Create the current and future logs and verify that metrics are present for both current and future logs
+    logManager.maybeUpdatePreferredLogDir(tp, dir1.getAbsolutePath)
+    logManager.getOrCreateLog(tp, () => logConfig)
+    logManager.maybeUpdatePreferredLogDir(tp, dir2.getAbsolutePath)
+    logManager.getOrCreateLog(tp, () => logConfig, isFuture = true)
+    verifyMetrics(2)
+
+    // Replace the current log with the future one and verify that only one set of metrics are present
+    logManager.replaceCurrentWithFutureLog(tp)
+    verifyMetrics(1)
+
+    // Trigger the deletion of the former current directory and verify that one set of metrics is still present
+    time.sleep(logConfig.fileDeleteDelayMs + 1)
+    verifyMetrics(1)
   }
 }
diff --git a/core/src/test/scala/unit/kafka/log/LogSegmentTest.scala b/core/src/test/scala/unit/kafka/log/LogSegmentTest.scala
index cef2bcaeda2ba..93d45e283c8b8 100644
--- a/core/src/test/scala/unit/kafka/log/LogSegmentTest.scala
+++ b/core/src/test/scala/unit/kafka/log/LogSegmentTest.scala
@@ -14,20 +14,24 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
- package kafka.log
+package kafka.log
 
 import java.io.File
 
+import kafka.server.checkpoints.LeaderEpochCheckpoint
+import kafka.server.epoch.EpochEntry
+import kafka.server.epoch.LeaderEpochFileCache
 import kafka.utils.TestUtils
 import kafka.utils.TestUtils.checkEquals
 import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.record._
-import org.apache.kafka.common.utils.{Time, Utils}
+import org.apache.kafka.common.utils.{MockTime, Time, Utils}
 import org.junit.Assert._
 import org.junit.{After, Before, Test}
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 import scala.collection._
+import scala.collection.mutable.ArrayBuffer
 
 class LogSegmentTest {
 
@@ -36,19 +40,10 @@ class LogSegmentTest {
   var logDir: File = _
 
   /* create a segment with the given base offset */
-  def createSegment(offset: Long, indexIntervalBytes: Int = 10): LogSegment = {
-    val msFile = TestUtils.tempFile()
-    val ms = FileRecords.open(msFile)
-    val idxFile = TestUtils.tempFile()
-    val timeIdxFile = TestUtils.tempFile()
-    val txnIdxFile = TestUtils.tempFile()
-    idxFile.delete()
-    timeIdxFile.delete()
-    txnIdxFile.delete()
-    val idx = new OffsetIndex(idxFile, offset, 1000)
-    val timeIdx = new TimeIndex(timeIdxFile, offset, 1500)
-    val txnIndex = new TransactionIndex(offset, txnIdxFile)
-    val seg = new LogSegment(ms, idx, timeIdx, txnIndex, offset, indexIntervalBytes, 0, Time.SYSTEM)
+  def createSegment(offset: Long,
+                    indexIntervalBytes: Int = 10,
+                    time: Time = Time.SYSTEM): LogSegment = {
+    val seg = LogUtils.createSegment(offset, logDir, indexIntervalBytes, time)
     segments += seg
     seg
   }
@@ -65,13 +60,8 @@ class LogSegmentTest {
   }
 
   @After
-  def teardown() {
-    for(seg <- segments) {
-      seg.index.delete()
-      seg.timeIndex.delete()
-      seg.txnIndex.delete()
-      seg.log.delete()
-    }
+  def teardown(): Unit = {
+    segments.foreach(_.close())
     Utils.delete(logDir)
   }
 
@@ -79,9 +69,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)
   }
 
@@ -90,41 +80,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(50, 53, RecordBatch.NO_TIMESTAMP, -1L, ms)
-    val read = seg.read(startOffset = 41, maxSize = 300, maxOffset = None).records
+    seg.append(53, RecordBatch.NO_TIMESTAMP, -1L, ms)
+    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(baseOffset, 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(50, 51, RecordBatch.NO_TIMESTAMP, -1L, ms)
-    val read = seg.read(startOffset = 52, maxSize = 200, maxOffset = None)
+    seg.append(51, RecordBatch.NO_TIMESTAMP, -1L, ms)
+    val read = seg.read(startOffset = 52, maxSize = 200)
     assertNull("Read beyond the last offset in the segment should give null", read)
   }
 
@@ -133,13 +105,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(50, 51, RecordBatch.NO_TIMESTAMP, -1L, ms)
+    seg.append(51, RecordBatch.NO_TIMESTAMP, -1L, ms)
     val ms2 = records(60, "alpha", "beta")
-    seg.append(60, 61, RecordBatch.NO_TIMESTAMP, -1L, ms2)
-    val read = seg.read(startOffset = 55, maxSize = 200, maxOffset = None)
+    seg.append(61, RecordBatch.NO_TIMESTAMP, -1L, ms2)
+    val read = seg.read(startOffset = 55, maxSize = 200)
     checkEquals(ms2.records.iterator, read.records.records.iterator)
   }
 
@@ -148,20 +120,20 @@ 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) {
       val ms1 = records(offset, "hello")
-      seg.append(offset, offset, RecordBatch.NO_TIMESTAMP, -1L, ms1)
+      seg.append(offset, RecordBatch.NO_TIMESTAMP, -1L, ms1)
       val ms2 = records(offset + 1, "hello")
-      seg.append(offset + 1, offset + 1, RecordBatch.NO_TIMESTAMP, -1L, ms2)
+      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
@@ -169,15 +141,56 @@ class LogSegmentTest {
   }
 
   @Test
-  def testReloadLargestTimestampAndNextOffsetAfterTruncation() {
+  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
+
+    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)
+    assertEquals(0, seg.timeIndex.sizeInBytes)
+    assertEquals(0, seg.offsetIndex.sizeInBytes)
+
+    time.sleep(500)
+    reopened.truncateTo(57)
+    assertEquals(0, reopened.timeWaitedForRoll(time.milliseconds(), RecordBatch.NO_TIMESTAMP))
+    assertFalse(reopened.timeIndex.isFull)
+    assertFalse(reopened.offsetIndex.isFull)
+
+    var rollParams = RollParams(maxSegmentMs, maxSegmentBytes = Int.MaxValue, RecordBatch.NO_TIMESTAMP,
+      maxOffsetInMessages = 100L, messagesSize = 1024, time.milliseconds())
+    assertFalse(reopened.shouldRoll(rollParams))
+
+    // The segment should not be rolled even if maxSegmentMs has been exceeded
+    time.sleep(maxSegmentMs + 1)
+    assertEquals(maxSegmentMs + 1, reopened.timeWaitedForRoll(time.milliseconds(), RecordBatch.NO_TIMESTAMP))
+    rollParams = RollParams(maxSegmentMs, maxSegmentBytes = Int.MaxValue, RecordBatch.NO_TIMESTAMP,
+      maxOffsetInMessages = 100L, messagesSize = 1024, time.milliseconds())
+    assertFalse(reopened.shouldRoll(rollParams))
+
+    // But we should still roll the segment if we cannot fit the next offset
+    rollParams = RollParams(maxSegmentMs, maxSegmentBytes = Int.MaxValue, RecordBatch.NO_TIMESTAMP,
+      maxOffsetInMessages = Int.MaxValue.toLong + 200L, messagesSize = 1024, time.milliseconds())
+    assertTrue(reopened.shouldRoll(rollParams))
+  }
+
+  @Test
+  def testReloadLargestTimestampAndNextOffsetAfterTruncation(): Unit = {
     val numMessages = 30
     val seg = createSegment(40, 2 * records(0, "hello").sizeInBytes - 1)
     var offset = 40
     for (_ <- 0 until numMessages) {
-      seg.append(offset, offset, offset, offset, records(offset, "hello"))
+      seg.append(offset, offset, offset, records(offset, "hello"))
       offset += 1
     }
-    assertEquals(offset, seg.nextOffset)
+    assertEquals(offset, seg.readNextOffset)
 
     val expectedNumEntries = numMessages / 2 - 1
     assertEquals(s"Should have $expectedNumEntries time indexes", expectedNumEntries, seg.timeIndex.entries)
@@ -185,32 +198,42 @@ class LogSegmentTest {
     seg.truncateTo(41)
     assertEquals(s"Should have 0 time indexes", 0, seg.timeIndex.entries)
     assertEquals(s"Largest timestamp should be 400", 400L, seg.largestTimestamp)
-    assertEquals(41, seg.nextOffset)
+    assertEquals(41, seg.readNextOffset)
   }
 
   /**
    * 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 seg = createSegment(40)
-    seg.append(40, 41, RecordBatch.NO_TIMESTAMP, -1L, records(40, "hello", "there"))
+    val time = new MockTime
+    val seg = createSegment(40, time = time)
+    seg.append(41, RecordBatch.NO_TIMESTAMP, -1L, records(40, "hello", "there"))
+
+    // If the segment is empty after truncation, the create time should be reset
+    time.sleep(500)
+    assertEquals(500, seg.timeWaitedForRoll(time.milliseconds(), RecordBatch.NO_TIMESTAMP))
+
     seg.truncateTo(0)
-    assertNull("Segment should be empty.", seg.read(0, None, 1024))
-    seg.append(40, 41, RecordBatch.NO_TIMESTAMP, -1L, records(40, "hello", "there"))
+    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, 1024))
+
+    seg.append(41, RecordBatch.NO_TIMESTAMP, -1L, records(40, "hello", "there"))
   }
 
   /**
    * 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
     for (i <- 40 until 50)
-      seg.append(i, i, i * 10, i, records(i, s"msg$i"))
+      seg.append(i, i * 10, i, records(i, s"msg$i"))
 
     assertEquals(490, seg.largestTimestamp)
     // Search for an indexed timestamp
@@ -231,26 +254,40 @@ 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.nextOffset)
-    seg.append(50, 52, RecordBatch.NO_TIMESTAMP, -1L, records(50, "hello", "there", "you"))
-    assertEquals(53, seg.nextOffset)
+    assertEquals(40, seg.readNextOffset)
+    seg.append(52, RecordBatch.NO_TIMESTAMP, -1L, records(50, "hello", "there", "you"))
+    assertEquals(53, seg.readNextOffset)
   }
 
   /**
    * 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.index.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.index.file.getAbsolutePath)
+    assertEquals(indexFile.getAbsolutePath + ".deleted", seg.lazyOffsetIndex.file.getAbsolutePath)
+    assertEquals(timeIndexFile.getAbsolutePath + ".deleted", seg.lazyTimeIndex.file.getAbsolutePath)
     assertTrue(seg.log.file.exists)
-    assertTrue(seg.index.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)
   }
 
   /**
@@ -258,15 +295,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, i, RecordBatch.NO_TIMESTAMP, -1L, records(i, i.toString))
-    val indexFile = seg.index.file
+      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
@@ -280,27 +319,27 @@ class LogSegmentTest {
     val pid2 = 10L
 
     // append transactional records from pid1
-    segment.append(firstOffset = 100L, largestOffset = 101L, largestTimestamp = RecordBatch.NO_TIMESTAMP,
-      shallowOffsetOfMaxTimestamp = 100L, MemoryRecords.withTransactionalRecords(100L, CompressionType.NONE,
+    segment.append(largestOffset = 101L, largestTimestamp = RecordBatch.NO_TIMESTAMP,
+      shallowOffsetOfMaxTimestamp = 100L, records = MemoryRecords.withTransactionalRecords(100L, CompressionType.NONE,
         pid1, producerEpoch, sequence, partitionLeaderEpoch, new SimpleRecord("a".getBytes), new SimpleRecord("b".getBytes)))
 
     // append transactional records from pid2
-    segment.append(firstOffset = 102L, largestOffset = 103L, largestTimestamp = RecordBatch.NO_TIMESTAMP,
-      shallowOffsetOfMaxTimestamp = 102L, MemoryRecords.withTransactionalRecords(102L, CompressionType.NONE,
+    segment.append(largestOffset = 103L, largestTimestamp = RecordBatch.NO_TIMESTAMP,
+      shallowOffsetOfMaxTimestamp = 102L, records = MemoryRecords.withTransactionalRecords(102L, CompressionType.NONE,
         pid2, producerEpoch, sequence, partitionLeaderEpoch, new SimpleRecord("a".getBytes), new SimpleRecord("b".getBytes)))
 
     // append non-transactional records
-    segment.append(firstOffset = 104L, largestOffset = 105L, largestTimestamp = RecordBatch.NO_TIMESTAMP,
-      shallowOffsetOfMaxTimestamp = 104L, MemoryRecords.withRecords(104L, CompressionType.NONE,
+    segment.append(largestOffset = 105L, largestTimestamp = RecordBatch.NO_TIMESTAMP,
+      shallowOffsetOfMaxTimestamp = 104L, records = MemoryRecords.withRecords(104L, CompressionType.NONE,
         partitionLeaderEpoch, new SimpleRecord("a".getBytes), new SimpleRecord("b".getBytes)))
 
     // abort the transaction from pid2 (note LSO should be 100L since the txn from pid1 has not completed)
-    segment.append(firstOffset = 106L, largestOffset = 106L, largestTimestamp = RecordBatch.NO_TIMESTAMP,
-      shallowOffsetOfMaxTimestamp = 106L, endTxnRecords(ControlRecordType.ABORT, pid2, producerEpoch, offset = 106L))
+    segment.append(largestOffset = 106L, largestTimestamp = RecordBatch.NO_TIMESTAMP,
+      shallowOffsetOfMaxTimestamp = 106L, records = endTxnRecords(ControlRecordType.ABORT, pid2, producerEpoch, offset = 106L))
 
     // commit the transaction from pid1
-    segment.append(firstOffset = 107L, largestOffset = 107L, largestTimestamp = RecordBatch.NO_TIMESTAMP,
-      shallowOffsetOfMaxTimestamp = 107L, endTxnRecords(ControlRecordType.COMMIT, pid1, producerEpoch, offset = 107L))
+    segment.append(largestOffset = 107L, largestTimestamp = RecordBatch.NO_TIMESTAMP,
+      shallowOffsetOfMaxTimestamp = 107L, records = endTxnRecords(ControlRecordType.COMMIT, pid1, producerEpoch, offset = 107L))
 
     var stateManager = new ProducerStateManager(topicPartition, logDir)
     segment.recover(stateManager)
@@ -317,8 +356,9 @@ 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 ProducerIdEntry(pid2,
-      mutable.Queue[BatchMetadata](BatchMetadata(10, 10L, 5, RecordBatch.NO_TIMESTAMP)), producerEpoch, 0, Some(75L)))
+    stateManager.loadProducerEntry(new ProducerStateEntry(pid2,
+      mutable.Queue[BatchMetadata](BatchMetadata(10, 10L, 5, RecordBatch.NO_TIMESTAMP)), producerEpoch,
+      0, RecordBatch.NO_TIMESTAMP, Some(75L)))
     segment.recover(stateManager)
     assertEquals(108L, stateManager.mapEndOffset)
 
@@ -331,6 +371,48 @@ class LogSegmentTest {
     assertEquals(100L, abortedTxn.lastStableOffset)
   }
 
+  /**
+   * Create a segment with some data, then recover the segment.
+   * The epoch cache entries should reflect the segment.
+   */
+  @Test
+  def testRecoveryRebuildsEpochCache(): Unit = {
+    val seg = createSegment(0)
+
+    val checkpoint: LeaderEpochCheckpoint = new LeaderEpochCheckpoint {
+      private var epochs = Seq.empty[EpochEntry]
+
+      override def write(epochs: Iterable[EpochEntry]): Unit = {
+        this.epochs = epochs.toVector
+      }
+
+      override def read(): Seq[EpochEntry] = this.epochs
+    }
+
+    val cache = new LeaderEpochFileCache(topicPartition, () => seg.readNextOffset, checkpoint)
+    seg.append(largestOffset = 105L, largestTimestamp = RecordBatch.NO_TIMESTAMP,
+      shallowOffsetOfMaxTimestamp = 104L, records = MemoryRecords.withRecords(104L, CompressionType.NONE, 0,
+        new SimpleRecord("a".getBytes), new SimpleRecord("b".getBytes)))
+
+    seg.append(largestOffset = 107L, largestTimestamp = RecordBatch.NO_TIMESTAMP,
+      shallowOffsetOfMaxTimestamp = 106L, records = MemoryRecords.withRecords(106L, CompressionType.NONE, 1,
+        new SimpleRecord("a".getBytes), new SimpleRecord("b".getBytes)))
+
+    seg.append(largestOffset = 109L, largestTimestamp = RecordBatch.NO_TIMESTAMP,
+      shallowOffsetOfMaxTimestamp = 108L, records = MemoryRecords.withRecords(108L, CompressionType.NONE, 1,
+        new SimpleRecord("a".getBytes), new SimpleRecord("b".getBytes)))
+
+    seg.append(largestOffset = 111L, largestTimestamp = RecordBatch.NO_TIMESTAMP,
+      shallowOffsetOfMaxTimestamp = 110, records = MemoryRecords.withRecords(110L, CompressionType.NONE, 2,
+        new SimpleRecord("a".getBytes), new SimpleRecord("b".getBytes)))
+
+    seg.recover(new ProducerStateManager(topicPartition, logDir), Some(cache))
+    assertEquals(ArrayBuffer(EpochEntry(epoch = 0, startOffset = 104L),
+                             EpochEntry(epoch = 1, startOffset = 106),
+                             EpochEntry(epoch = 2, startOffset = 110)),
+      cache.epochEntries)
+  }
+
   private def endTxnRecords(controlRecordType: ControlRecordType,
                             producerId: Long,
                             producerEpoch: Short,
@@ -347,11 +429,11 @@ 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, i * 10, i, records(i, i.toString))
-    val timeIndexFile = seg.timeIndex.file
+      seg.append(i, i * 10, i, records(i, i.toString))
+    val timeIndexFile = seg.lazyTimeIndex.file
     TestUtils.writeNonsenseToFile(timeIndexFile, 5, timeIndexFile.length.toInt)
     seg.recover(new ProducerStateManager(topicPartition, logDir))
     for(i <- 0 until 100) {
@@ -365,12 +447,12 @@ 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)
-      for(i <- 0 until messagesAppended)
-        seg.append(i, i, RecordBatch.NO_TIMESTAMP, -1L, records(i, i.toString))
+      for (i <- 0 until messagesAppended)
+        seg.append(i, RecordBatch.NO_TIMESTAMP, -1L, records(i, i.toString))
       val offsetToBeginCorruption = TestUtils.random.nextInt(messagesAppended)
       // start corrupting somewhere in the middle of the chosen record all the way to the end
 
@@ -380,14 +462,18 @@ class LogSegmentTest {
       seg.recover(new ProducerStateManager(topicPartition, logDir))
       assertEquals("Should have truncated off bad messages.", (0 until offsetToBeginCorruption).toList,
         seg.log.batches.asScala.map(_.lastOffset).toList)
-      seg.delete()
+      seg.deleteIfExists()
     }
   }
 
-  /* create a segment with   pre allocate */
-  def createSegment(offset: Long, fileAlreadyExists: Boolean, initFileSize: Int, preallocate: Boolean): LogSegment = {
+  private def createSegment(baseOffset: Long, fileAlreadyExists: Boolean, initFileSize: Int, preallocate: Boolean): LogSegment = {
     val tempDir = TestUtils.tempDir()
-    val seg = new LogSegment(tempDir, offset, 10, 1000, 0, Time.SYSTEM, fileAlreadyExists = fileAlreadyExists,
+    val logConfig = LogConfig(Map(
+      LogConfig.IndexIntervalBytesProp -> 10,
+      LogConfig.SegmentIndexBytesProp -> 1000,
+      LogConfig.SegmentJitterMsProp -> 0
+    ).asJava)
+    val seg = LogSegment.open(tempDir, baseOffset, logConfig, Time.SYSTEM, fileAlreadyExists = fileAlreadyExists,
       initFileSize = initFileSize, preallocate = preallocate)
     segments += seg
     seg
@@ -395,27 +481,34 @@ 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(50, 51, RecordBatch.NO_TIMESTAMP, -1L, ms)
+    seg.append(51, RecordBatch.NO_TIMESTAMP, -1L, ms)
     val ms2 = records(60, "alpha", "beta")
-    seg.append(60, 61, RecordBatch.NO_TIMESTAMP, -1L, ms2)
-    val read = seg.read(startOffset = 55, maxSize = 200, maxOffset = None)
+    seg.append(61, RecordBatch.NO_TIMESTAMP, -1L, ms2)
+    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 seg = new LogSegment(tempDir, 40, 10, 1000, 0, Time.SYSTEM, false, 512*1024*1024, true)
+    val logConfig = LogConfig(Map(
+      LogConfig.IndexIntervalBytesProp -> 10,
+      LogConfig.SegmentIndexBytesProp -> 1000,
+      LogConfig.SegmentJitterMsProp -> 0
+    ).asJava)
+
+    val seg = LogSegment.open(tempDir, baseOffset = 40, logConfig, Time.SYSTEM,
+      initFileSize = 512 * 1024 * 1024, preallocate = true)
 
     val ms = records(50, "hello", "there")
-    seg.append(50, 51, RecordBatch.NO_TIMESTAMP, -1L, ms)
+    seg.append(51, RecordBatch.NO_TIMESTAMP, -1L, ms)
     val ms2 = records(60, "alpha", "beta")
-    seg.append(60, 61, RecordBatch.NO_TIMESTAMP, -1L, ms2)
-    val read = seg.read(startOffset = 55, maxSize = 200, maxOffset = None)
+    seg.append(61, RecordBatch.NO_TIMESTAMP, -1L, ms2)
+    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
@@ -425,10 +518,11 @@ class LogSegmentTest {
     //After close, file should be trimmed
     assertEquals(oldSize, seg.log.file.length)
 
-    val segReopen = new LogSegment(tempDir, 40, 10, 1000, 0, Time.SYSTEM, true,  512*1024*1024, true)
+    val segReopen = LogSegment.open(tempDir, baseOffset = 40, logConfig, Time.SYSTEM, fileAlreadyExists = true,
+      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
@@ -439,7 +533,7 @@ class LogSegmentTest {
   }
 
   @Test
-  def shouldTruncateEvenIfOffsetPointsToAGapInTheLog() {
+  def shouldTruncateEvenIfOffsetPointsToAGapInTheLog(): Unit = {
     val seg = createSegment(40)
     val offset = 40
 
@@ -449,17 +543,47 @@ class LogSegmentTest {
 
     //Given two messages with a gap between them (e.g. mid offset compacted away)
     val ms1 = records(offset, "first message")
-    seg.append(offset, offset, RecordBatch.NO_TIMESTAMP, -1L, ms1)
+    seg.append(offset, RecordBatch.NO_TIMESTAMP, -1L, ms1)
     val ms2 = records(offset + 3, "message after gap")
-    seg.append(offset + 3, offset + 3, RecordBatch.NO_TIMESTAMP, -1L, ms2)
+    seg.append(offset + 3, RecordBatch.NO_TIMESTAMP, -1L, ms2)
 
     // When we truncate to an offset without a corresponding log entry
     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)
   }
 
+  @Test
+  def testAppendFromFile(): Unit = {
+    def records(offset: Long, size: Int): MemoryRecords =
+      MemoryRecords.withRecords(RecordBatch.MAGIC_VALUE_V2, offset, CompressionType.NONE, TimestampType.CREATE_TIME,
+        new SimpleRecord(new Array[Byte](size)))
+
+    // create a log file in a separate directory to avoid conflicting with created segments
+    val tempDir = TestUtils.tempDir()
+    val fileRecords = FileRecords.open(Log.logFile(tempDir, 0))
+
+    // Simulate a scenario where we have a single log with an offset range exceeding Int.MaxValue
+    fileRecords.append(records(0, 1024))
+    fileRecords.append(records(500, 1024 * 1024 + 1))
+    val sizeBeforeOverflow = fileRecords.sizeInBytes()
+    fileRecords.append(records(Int.MaxValue + 5L, 1024))
+    val sizeAfterOverflow = fileRecords.sizeInBytes()
+
+    val segment = createSegment(0)
+    val bytesAppended = segment.appendFromFile(fileRecords, 0)
+    assertEquals(sizeBeforeOverflow, bytesAppended)
+    assertEquals(sizeBeforeOverflow, segment.size)
+
+    val overflowSegment = createSegment(Int.MaxValue)
+    val overflowBytesAppended = overflowSegment.appendFromFile(fileRecords, sizeBeforeOverflow)
+    assertEquals(sizeAfterOverflow - sizeBeforeOverflow, overflowBytesAppended)
+    assertEquals(overflowBytesAppended, overflowSegment.size)
+
+    Utils.delete(tempDir)
+  }
+
 }
diff --git a/core/src/test/scala/unit/kafka/log/LogTest.scala b/core/src/test/scala/unit/kafka/log/LogTest.scala
index 1e408eccd8a20..a3aea7a6d0201 100755
--- a/core/src/test/scala/unit/kafka/log/LogTest.scala
+++ b/core/src/test/scala/unit/kafka/log/LogTest.scala
@@ -19,49 +19,59 @@ package kafka.log
 
 import java.io._
 import java.nio.ByteBuffer
-import java.util.Properties
+import java.nio.file.{Files, Paths}
+import java.util.concurrent.{Callable, Executors}
+import java.util.regex.Pattern
+import java.util.{Collections, Optional, Properties}
 
-import org.apache.kafka.common.errors._
-import kafka.common.KafkaException
+import kafka.api.{ApiVersion, KAFKA_0_11_0_IV0}
+import kafka.common.{OffsetsOutOfOrderException, RecordValidationException, UnexpectedAppendOffsetException}
 import kafka.log.Log.DeleteDirSuffix
-import org.junit.Assert._
-import org.junit.{After, Before, Test}
+import kafka.metrics.KafkaYammerMetrics
+import kafka.server.checkpoints.LeaderEpochCheckpointFile
+import kafka.server.epoch.{EpochEntry, LeaderEpochFileCache}
+import kafka.server.{BrokerState, BrokerTopicStats, FetchDataInfo, FetchHighWatermark, FetchIsolation, FetchLogEnd, FetchTxnCommitted, KafkaConfig, LogDirFailureChannel, LogOffsetMetadata}
 import kafka.utils._
-import kafka.server.{BrokerTopicStats, FetchDataInfo, KafkaConfig, LogDirFailureChannel}
-import kafka.server.epoch.{EpochEntry, LeaderEpochCache, LeaderEpochFileCache}
-import org.apache.kafka.common.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
 import org.apache.kafka.common.record.MemoryRecords.RecordFilter.BatchRetention
 import org.apache.kafka.common.record._
 import org.apache.kafka.common.requests.FetchResponse.AbortedTransaction
-import org.apache.kafka.common.requests.IsolationLevel
+import org.apache.kafka.common.requests.{ListOffsetRequest, ListOffsetResponse}
 import org.apache.kafka.common.utils.{Time, Utils}
 import org.easymock.EasyMock
+import org.junit.Assert._
+import org.junit.{After, Before, Test}
+import org.scalatest.Assertions
 
-import scala.collection.JavaConverters._
-import scala.collection.mutable.{ArrayBuffer, ListBuffer}
+import scala.collection.{Iterable, Map, mutable}
+import scala.jdk.CollectionConverters._
+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()
-  var config: KafkaConfig = null
-  val brokerTopicStats = new BrokerTopicStats
+  def metricsKeySet = KafkaYammerMetrics.defaultRegistry.allMetrics.keySet.asScala
 
   @Before
-  def setUp() {
+  def setUp(): Unit = {
     val props = TestUtils.createBrokerConfig(0, "127.0.0.1:1", port = -1)
     config = KafkaConfig.fromProps(props)
   }
 
   @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()
@@ -69,7 +79,323 @@ class LogTest {
   }
 
   @Test
-  def testOffsetFromFile() {
+  def testLogRecoveryIsCalledUponBrokerCrash(): Unit = {
+    // LogManager must realize correctly if the last shutdown was not clean and the logs need
+    // to run recovery while loading upon subsequent broker boot up.
+    val logDir: File = TestUtils.tempDir()
+    val logProps = new Properties()
+    val logConfig = LogConfig(logProps)
+    val logDirs = Seq(logDir)
+    val topicPartition = new TopicPartition("foo", 0)
+    var log: Log = null
+    val time = new MockTime()
+    var cleanShutdownInterceptedValue = false
+    var simulateError = false
+
+    // Create a LogManager with some overridden methods to facilitate interception of clean shutdown
+    // flag and to inject a runtime error
+    def interceptedLogManager(logConfig: LogConfig, logDirs: Seq[File]): LogManager = {
+      new LogManager(logDirs = logDirs.map(_.getAbsoluteFile), initialOfflineDirs = Array.empty[File], topicConfigs = Map(),
+        initialDefaultConfig = logConfig, cleanerConfig = CleanerConfig(enableCleaner = false), recoveryThreadsPerDataDir = 4,
+        flushCheckMs = 1000L, flushRecoveryOffsetCheckpointMs = 10000L, flushStartOffsetCheckpointMs = 10000L,
+        retentionCheckMs = 1000L, maxPidExpirationMs = 60 * 60 * 1000, scheduler = time.scheduler, time = time, brokerState = BrokerState(),
+        brokerTopicStats = new BrokerTopicStats, logDirFailureChannel = new LogDirFailureChannel(logDirs.size)) {
+
+         override def loadLog(logDir: File, hadCleanShutdown: Boolean, recoveryPoints: Map[TopicPartition, Long],
+                     logStartOffsets: Map[TopicPartition, Long]): Log = {
+
+          val topicPartition = Log.parseTopicPartitionName(logDir)
+          val config = topicConfigs.getOrElse(topicPartition.topic, currentDefaultConfig)
+          val logRecoveryPoint = recoveryPoints.getOrElse(topicPartition, 0L)
+          val logStartOffset = logStartOffsets.getOrElse(topicPartition, 0L)
+          val logDirFailureChannel: LogDirFailureChannel = new LogDirFailureChannel(1)
+
+          val producerStateManager = new ProducerStateManager(topicPartition, logDir, maxPidExpirationMs)
+          val log = new Log(logDir, config, logStartOffset, logRecoveryPoint, time.scheduler, brokerTopicStats, time, maxPidExpirationMs,
+            LogManager.ProducerIdExpirationCheckIntervalMs, topicPartition, producerStateManager, logDirFailureChannel, hadCleanShutdown) {
+            override def recoverLog(): Long = {
+              if (simulateError)
+                throw new RuntimeException
+              cleanShutdownInterceptedValue = hadCleanShutdown
+              super.recoverLog()
+            }
+          }
+          log
+
+         }
+
+      }
+    }
+
+    val cleanShutdownFile = new File(logDir, Log.CleanShutdownFile)
+    val logManager: LogManager = interceptedLogManager(logConfig, logDirs)
+    log = logManager.getOrCreateLog(topicPartition, () => logConfig, isNew = true)
+
+    // Load logs after a clean shutdown
+    Files.createFile(cleanShutdownFile.toPath)
+    cleanShutdownInterceptedValue = false
+    logManager.loadLogs()
+    assertTrue("Unexpected value intercepted for clean shutdown flag", cleanShutdownInterceptedValue)
+    assertTrue("Clean shutdown file must not exist after loadLogs has completed", !cleanShutdownFile.exists())
+    // Load logs without clean shutdown file
+    cleanShutdownInterceptedValue = true
+    logManager.loadLogs()
+    assertTrue("Unexpected value intercepted for clean shutdown flag", !cleanShutdownInterceptedValue)
+    assertTrue("Clean shutdown file must not exist after loadLogs has completed", !cleanShutdownFile.exists())
+    // Create clean shutdown file and then simulate error while loading logs such that log loading does not complete.
+    Files.createFile(cleanShutdownFile.toPath)
+    simulateError = true
+    assertThrows[RuntimeException](logManager.loadLogs())
+    assertTrue("Clean shutdown file must not have existed", !cleanShutdownFile.exists())
+    // Do not simulate error on next call to LogManager#loadLogs. LogManager must understand that log had unclean shutdown the last time.
+    simulateError = false
+    cleanShutdownInterceptedValue = true
+    logManager.loadLogs()
+    assertTrue("Unexpected value for clean shutdown flag", !cleanShutdownInterceptedValue)
+  }
+
+  @Test
+  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)
@@ -90,9 +416,9 @@ 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 = createLogConfig(segmentMs = 1 * 60 * 60L)
+    val logConfig = LogTest.createLogConfig(segmentMs = 1 * 60 * 60L)
 
     // create a log
     val log = createLog(logDir, logConfig, maxProducerIdExpirationMs = 24 * 60)
@@ -137,6 +463,52 @@ class LogTest {
     assertEquals("Appending an empty message set should not roll log even if sufficient time has passed.", numSegments, log.numberOfSegments)
   }
 
+  @Test
+  def testRollSegmentThatAlreadyExists(): Unit = {
+    val logConfig = LogTest.createLogConfig(segmentMs = 1 * 60 * 60L)
+
+    // create a log
+    val log = createLog(logDir, logConfig)
+    assertEquals("Log begins with a single empty segment.", 1, log.numberOfSegments)
+
+    // roll active segment with the same base offset of size zero should recreate the segment
+    log.roll(Some(0L))
+    assertEquals("Expect 1 segment after roll() empty segment with base offset.", 1, log.numberOfSegments)
+
+    // should be able to append records to active segment
+    val records = TestUtils.records(
+      List(new SimpleRecord(mockTime.milliseconds, "k1".getBytes, "v1".getBytes)),
+      baseOffset = 0L, partitionLeaderEpoch = 0)
+    log.appendAsFollower(records)
+    assertEquals("Expect one segment.", 1, log.numberOfSegments)
+    assertEquals(0L, log.activeSegment.baseOffset)
+
+    // make sure we can append more records
+    val records2 = TestUtils.records(
+      List(new SimpleRecord(mockTime.milliseconds + 10, "k2".getBytes, "v2".getBytes)),
+      baseOffset = 1L, partitionLeaderEpoch = 0)
+    log.appendAsFollower(records2)
+
+    assertEquals("Expect two records in the log", 2, log.logEndOffset)
+    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()
+    assertEquals("Expect base offset of active segment to be LEO", 2L, log.activeSegment.baseOffset)
+    assertEquals("Expect two segments.", 2, log.numberOfSegments)
+
+    // manually resize offset index to force roll of an empty active segment on next append
+    log.activeSegment.offsetIndex.resize(0)
+    val records3 = TestUtils.records(
+      List(new SimpleRecord(mockTime.milliseconds + 12, "k3".getBytes, "v3".getBytes)),
+      baseOffset = 2L, partitionLeaderEpoch = 0)
+    log.appendAsFollower(records3)
+    assertTrue(log.activeSegment.offsetIndex.maxEntries > 1)
+    assertEquals(2, readLog(log, 2, 1).records.batches.iterator.next().lastOffset)
+    assertEquals("Expect two segments.", 2, log.numberOfSegments)
+  }
+
   @Test(expected = classOf[OutOfOrderSequenceException])
   def testNonSequentialAppend(): Unit = {
     // create a log
@@ -151,11 +523,102 @@ class LogTest {
     log.appendAsLeader(nextRecords, leaderEpoch = 0)
   }
 
+  @Test
+  def testTruncateToEmptySegment(): Unit = {
+    val log = createLog(logDir, LogConfig())
+
+    // Force a segment roll by using a large offset. The first segment will be empty
+    val records = TestUtils.records(List(new SimpleRecord(mockTime.milliseconds, "key".getBytes, "value".getBytes)),
+      baseOffset = Int.MaxValue.toLong + 200)
+    appendAsFollower(log, records)
+    assertEquals(0, log.logSegments.head.size)
+    assertEquals(2, log.logSegments.size)
+
+    // Truncate to an offset before the base offset of the latest segment
+    log.truncateTo(0L)
+    assertEquals(1, log.logSegments.size)
+
+    // Now verify that we can still append to the active segment
+    appendAsFollower(log, TestUtils.records(List(new SimpleRecord(mockTime.milliseconds, "key".getBytes, "value".getBytes)),
+      baseOffset = 100L))
+    assertEquals(1, log.logSegments.size)
+    assertEquals(101L, log.logEndOffset)
+  }
+
+  @Test
+  def testTruncateToEndOffsetClearsEpochCache(): Unit = {
+    val log = createLog(logDir, LogConfig())
+
+    // Seed some initial data in the log
+    val records = TestUtils.records(List(new SimpleRecord("a".getBytes), new SimpleRecord("b".getBytes)),
+      baseOffset = 27)
+    appendAsFollower(log, records, leaderEpoch = 19)
+    assertEquals(Some(EpochEntry(epoch = 19, startOffset = 27)),
+      log.leaderEpochCache.flatMap(_.latestEntry))
+    assertEquals(29, log.logEndOffset)
+
+    def verifyTruncationClearsEpochCache(epoch: Int, truncationOffset: Long): Unit = {
+      // Simulate becoming a leader
+      log.maybeAssignEpochStartOffset(leaderEpoch = epoch, startOffset = log.logEndOffset)
+      assertEquals(Some(EpochEntry(epoch = epoch, startOffset = 29)),
+        log.leaderEpochCache.flatMap(_.latestEntry))
+      assertEquals(29, log.logEndOffset)
+
+      // Now we become the follower and truncate to an offset greater
+      // than or equal to the log end offset. The trivial epoch entry
+      // at the end of the log should be gone
+      log.truncateTo(truncationOffset)
+      assertEquals(Some(EpochEntry(epoch = 19, startOffset = 27)),
+        log.leaderEpochCache.flatMap(_.latestEntry))
+      assertEquals(29, log.logEndOffset)
+    }
+
+    // Truncations greater than or equal to the log end offset should
+    // clear the epoch cache
+    verifyTruncationClearsEpochCache(epoch = 20, truncationOffset = log.logEndOffset)
+    verifyTruncationClearsEpochCache(epoch = 24, truncationOffset = log.logEndOffset + 1)
+  }
+
+  /**
+   * 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
     // snapshot files, and then reloading the log
-    val logConfig = createLogConfig(segmentBytes = 64 * 10)
+    val logConfig = LogTest.createLogConfig(segmentBytes = 64 * 10)
     var log = createLog(logDir, logConfig)
     assertEquals(None, log.oldestProducerSnapshotOffset)
 
@@ -167,7 +630,6 @@ class LogTest {
     val logEndOffset = log.logEndOffset
     log.close()
 
-    val cleanShutdownFile = createCleanShutdownFile()
     deleteProducerSnapshotFiles()
 
     // Reload after clean shutdown
@@ -176,29 +638,193 @@ class LogTest {
     assertEquals(expectedSnapshotOffsets, listProducerSnapshotOffsets)
     log.close()
 
-    Utils.delete(cleanShutdownFile)
     deleteProducerSnapshotFiles()
 
     // Reload after unclean shutdown with recoveryPoint set to log end offset
-    log = createLog(logDir, logConfig, recoveryPoint = logEndOffset)
-    // Note that we don't maintain the guarantee of having a snapshot for the 2 most recent segments in this case
-    expectedSnapshotOffsets = Vector(log.logSegments.last.baseOffset, log.logEndOffset)
+    log = createLog(logDir, logConfig, recoveryPoint = logEndOffset, lastShutdownClean = false)
     assertEquals(expectedSnapshotOffsets, listProducerSnapshotOffsets)
     log.close()
 
     deleteProducerSnapshotFiles()
 
     // Reload after unclean shutdown with recoveryPoint set to 0
-    log = createLog(logDir, logConfig, recoveryPoint = 0L)
-    // Is this working as intended?
+    log = createLog(logDir, logConfig, recoveryPoint = 0L, lastShutdownClean = false)
+    // We progressively create a snapshot for each segment after the recovery point
     expectedSnapshotOffsets = log.logSegments.map(_.baseOffset).tail.toVector :+ log.logEndOffset
     assertEquals(expectedSnapshotOffsets, listProducerSnapshotOffsets)
     log.close()
   }
 
+
   @Test
-  def testProducerSnapshotsRecoveryAfterUncleanShutdown(): Unit = {
-    val logConfig = createLogConfig(segmentBytes = 64 * 10)
+  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, lastShutdownClean = false)
+    assertEquals(secondAppendTimestamp, log.producerStateManager.lastEntry(producerId).get.lastTimestamp)
+    log.close()
+  }
+
+  @Test
+  def testProducerSnapshotsRecoveryAfterUncleanShutdownV1(): Unit = {
+    testProducerSnapshotsRecoveryAfterUncleanShutdown(ApiVersion.minSupportedFor(RecordVersion.V1).version)
+  }
+
+  @Test
+  def testProducerSnapshotsRecoveryAfterUncleanShutdownCurrentMessageFormat(): Unit = {
+    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, ClientRecordDeletion)
+    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, lastShutdownClean = false)
+    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.nonEmpty)
+    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)
     assertEquals(None, log.oldestProducerSnapshotOffset)
 
@@ -215,13 +841,23 @@ 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 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 ++ Set(activeSegmentOffset)
+      expectedSnapshotOffsets ++= log.logSegments.map(_.baseOffset).toVector.takeRight(4) :+ log.logEndOffset
+    }
 
     def createLogWithInterceptedReads(recoveryPoint: Long) = {
       val maxProducerIdExpirationMs = 60 * 60 * 1000
@@ -231,20 +867,19 @@ class LogTest {
       // Intercept all segment read calls
       new Log(logDir, logConfig, logStartOffset = 0, recoveryPoint = recoveryPoint, mockTime.scheduler,
         brokerTopicStats, mockTime, maxProducerIdExpirationMs, LogManager.ProducerIdExpirationCheckIntervalMs,
-        topicPartition, producerStateManager, new LogDirFailureChannel(10)) {
+        topicPartition, producerStateManager, new LogDirFailureChannel(10), hadCleanShutdown = false) {
 
         override def addSegment(segment: LogSegment): LogSegment = {
-          val wrapper = new LogSegment(segment.log, segment.index, segment.timeIndex, segment.txnIndex, segment.baseOffset,
+          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,
-                                 leaderEpochCache: Option[LeaderEpochCache]): Int = {
+                                 leaderEpochCache: Option[LeaderEpochFileCache]): Int = {
               recoveredSegments += this
               super.recover(producerStateManager, leaderEpochCache)
             }
@@ -255,37 +890,31 @@ class LogTest {
     }
 
     // Retain snapshots for the last 2 segments
-    ProducerStateManager.deleteSnapshotsBefore(logDir, segmentOffsets(segmentOffsets.size - 2))
+    log.producerStateManager.deleteSnapshotsBefore(segmentOffsets(segmentOffsets.size - 2))
     log = createLogWithInterceptedReads(offsetForRecoveryPointSegment)
     // We will reload all segments because the recovery point is behind the producer snapshot files (pre KAFKA-5829 behaviour)
-    assertEquals(segOffsetsBeforeRecovery, segmentsWithReads.map(_.baseOffset) -- Seq(activeSegmentOffset))
+    assertEquals(expectedSegmentsWithReads, segmentsWithReads.map(_.baseOffset))
     assertEquals(segOffsetsAfterRecovery, recoveredSegments.map(_.baseOffset))
-    var expectedSnapshotOffsets = segmentOffsets.takeRight(4) :+ log.logEndOffset
-    assertEquals(expectedSnapshotOffsets, listProducerSnapshotOffsets)
+    assertEquals(expectedSnapshotOffsets, listProducerSnapshotOffsets.toSet)
     log.close()
     segmentsWithReads.clear()
     recoveredSegments.clear()
 
     // Only delete snapshots before the base offset of the recovery point segment (post KAFKA-5829 behaviour) to
     // avoid reading all segments
-    ProducerStateManager.deleteSnapshotsBefore(logDir, offsetForRecoveryPointSegment)
+    log.producerStateManager.deleteSnapshotsBefore(offsetForRecoveryPointSegment)
     log = createLogWithInterceptedReads(recoveryPoint = recoveryPoint)
-    assertEquals(Seq(activeSegmentOffset), segmentsWithReads.map(_.baseOffset))
+    assertEquals(Set(activeSegmentOffset), segmentsWithReads.map(_.baseOffset))
     assertEquals(segOffsetsAfterRecovery, recoveredSegments.map(_.baseOffset))
-    expectedSnapshotOffsets = log.logSegments.map(_.baseOffset).toVector.takeRight(4) :+ log.logEndOffset
-    assertEquals(expectedSnapshotOffsets, listProducerSnapshotOffsets)
+    assertEquals(expectedSnapshotOffsets, listProducerSnapshotOffsets.toSet)
 
-    // Verify that we keep 2 snapshot files if we checkpoint the log end offset
-    log.deleteSnapshotsAfterRecoveryPointCheckpoint()
-    expectedSnapshotOffsets = log.logSegments.map(_.baseOffset).toVector.takeRight(2) :+ log.logEndOffset
-    assertEquals(expectedSnapshotOffsets, listProducerSnapshotOffsets)
     log.close()
   }
 
   @Test
   def testSizeForLargeLogs(): Unit = {
     val largeSize = Int.MaxValue.toLong * 2
-    val logSegment = EasyMock.createMock(classOf[LogSegment])
+    val logSegment: LogSegment = EasyMock.createMock(classOf[LogSegment])
 
     EasyMock.expect(logSegment.size).andReturn(Int.MaxValue).anyTimes
     EasyMock.replay(logSegment)
@@ -296,8 +925,8 @@ class LogTest {
   }
 
   @Test
-  def testProducerIdMapOffsetUpdatedForNonIdempotentData() {
-    val logConfig = createLogConfig(segmentBytes = 2048 * 5)
+  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)))
     log.appendAsLeader(records, leaderEpoch = 0)
@@ -307,8 +936,8 @@ class LogTest {
 
   @Test
   def testSkipLoadingIfEmptyProducerStateBeforeTruncation(): Unit = {
-    val stateManager = EasyMock.mock(classOf[ProducerStateManager])
-
+    val stateManager: ProducerStateManager = EasyMock.mock(classOf[ProducerStateManager])
+    EasyMock.expect(stateManager.removeStraySnapshots(EasyMock.anyObject())).anyTimes()
     // Load the log
     EasyMock.expect(stateManager.latestSnapshotOffset).andReturn(None)
 
@@ -340,7 +969,8 @@ class LogTest {
       producerIdExpirationCheckIntervalMs = 30000,
       topicPartition = Log.parseTopicPartitionName(logDir),
       producerStateManager = stateManager,
-      logDirFailureChannel = null)
+      logDirFailureChannel = null,
+      hadCleanShutdown = false)
 
     EasyMock.verify(stateManager)
 
@@ -373,6 +1003,9 @@ class LogTest {
     // We skip directly to updating the map end offset
     stateManager.updateMapEndOffset(1L)
     EasyMock.expectLastCall()
+    // Finally, we take a snapshot
+    stateManager.takeSnapshot()
+    EasyMock.expectLastCall().once()
 
     EasyMock.replay(stateManager)
 
@@ -383,9 +1016,8 @@ class LogTest {
 
   @Test
   def testSkipTruncateAndReloadIfOldMessageFormatAndNoCleanShutdown(): Unit = {
-    val stateManager = EasyMock.mock(classOf[ProducerStateManager])
-
-    EasyMock.expect(stateManager.latestSnapshotOffset).andReturn(None)
+    val stateManager: ProducerStateManager = EasyMock.mock(classOf[ProducerStateManager])
+    EasyMock.expect(stateManager.removeStraySnapshots(EasyMock.anyObject())).anyTimes()
 
     stateManager.updateMapEndOffset(0L)
     EasyMock.expectLastCall().anyTimes()
@@ -393,6 +1025,12 @@ class LogTest {
     stateManager.takeSnapshot()
     EasyMock.expectLastCall().anyTimes()
 
+    EasyMock.expect(stateManager.isEmpty).andReturn(true)
+    EasyMock.expectLastCall().once()
+
+    EasyMock.expect(stateManager.firstUnstableOffset).andReturn(None)
+    EasyMock.expectLastCall().once()
+
     EasyMock.replay(stateManager)
 
     val logProps = new Properties()
@@ -416,9 +1054,8 @@ class LogTest {
 
   @Test
   def testSkipTruncateAndReloadIfOldMessageFormatAndCleanShutdown(): Unit = {
-    val stateManager = EasyMock.mock(classOf[ProducerStateManager])
-
-    EasyMock.expect(stateManager.latestSnapshotOffset).andReturn(None)
+    val stateManager: ProducerStateManager = EasyMock.mock(classOf[ProducerStateManager])
+    EasyMock.expect(stateManager.removeStraySnapshots(EasyMock.anyObject())).anyTimes()
 
     stateManager.updateMapEndOffset(0L)
     EasyMock.expectLastCall().anyTimes()
@@ -426,9 +1063,13 @@ class LogTest {
     stateManager.takeSnapshot()
     EasyMock.expectLastCall().anyTimes()
 
-    EasyMock.replay(stateManager)
+    EasyMock.expect(stateManager.isEmpty).andReturn(true)
+    EasyMock.expectLastCall().once()
 
-    val cleanShutdownFile = createCleanShutdownFile()
+    EasyMock.expect(stateManager.firstUnstableOffset).andReturn(None)
+    EasyMock.expectLastCall().once()
+
+    EasyMock.replay(stateManager)
 
     val logProps = new Properties()
     logProps.put(LogConfig.MessageFormatVersionProp, "0.10.2")
@@ -447,12 +1088,12 @@ class LogTest {
       logDirFailureChannel = null)
 
     EasyMock.verify(stateManager)
-    Utils.delete(cleanShutdownFile)
   }
 
   @Test
   def testSkipTruncateAndReloadIfNewMessageFormatAndCleanShutdown(): Unit = {
-    val stateManager = EasyMock.mock(classOf[ProducerStateManager])
+    val stateManager: ProducerStateManager = EasyMock.mock(classOf[ProducerStateManager])
+    EasyMock.expect(stateManager.removeStraySnapshots(EasyMock.anyObject())).anyTimes()
 
     EasyMock.expect(stateManager.latestSnapshotOffset).andReturn(None)
 
@@ -462,9 +1103,13 @@ class LogTest {
     stateManager.takeSnapshot()
     EasyMock.expectLastCall().anyTimes()
 
-    EasyMock.replay(stateManager)
+    EasyMock.expect(stateManager.isEmpty).andReturn(true)
+    EasyMock.expectLastCall().once()
+
+    EasyMock.expect(stateManager.firstUnstableOffset).andReturn(None)
+    EasyMock.expectLastCall().once()
 
-    val cleanShutdownFile = createCleanShutdownFile()
+    EasyMock.replay(stateManager)
 
     val logProps = new Properties()
     logProps.put(LogConfig.MessageFormatVersionProp, "0.11.0")
@@ -483,12 +1128,11 @@ class LogTest {
       logDirFailureChannel = null)
 
     EasyMock.verify(stateManager)
-    Utils.delete(cleanShutdownFile)
   }
 
   @Test
-  def testRebuildProducerIdMapWithCompactedData() {
-    val logConfig = createLogConfig(segmentBytes = 2048 * 5)
+  def testRebuildProducerIdMapWithCompactedData(): Unit = {
+    val logConfig = LogTest.createLogConfig(segmentBytes = 2048 * 5)
     val log = createLog(logDir, logConfig)
     val pid = 1L
     val epoch = 0.toShort
@@ -497,11 +1141,11 @@ 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)))
-    records.batches.asScala.foreach(_.setPartitionLeaderEpoch(0))
+      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.forEach(_.setPartitionLeaderEpoch(0))
 
     val filtered = ByteBuffer.allocate(2048)
     records.filterTo(new TopicPartition("foo", 0), new RecordFilter {
@@ -515,9 +1159,9 @@ 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)))
-    moreRecords.batches.asScala.foreach(_.setPartitionLeaderEpoch(0))
+      new SimpleRecord(mockTime.milliseconds(), "e".getBytes),
+      new SimpleRecord(mockTime.milliseconds(), "f".getBytes)))
+    moreRecords.batches.forEach(_.setPartitionLeaderEpoch(0))
     log.appendAsFollower(moreRecords)
 
     log.truncateTo(baseOffset + 4)
@@ -530,8 +1174,8 @@ class LogTest {
   }
 
   @Test
-  def testRebuildProducerStateWithEmptyCompactedBatch() {
-    val logConfig = createLogConfig(segmentBytes = 2048 * 5)
+  def testRebuildProducerStateWithEmptyCompactedBatch(): Unit = {
+    val logConfig = LogTest.createLogConfig(segmentBytes = 2048 * 5)
     val log = createLog(logDir, logConfig)
     val pid = 1L
     val epoch = 0.toShort
@@ -540,9 +1184,9 @@ 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)))
-    records.batches.asScala.foreach(_.setPartitionLeaderEpoch(0))
+      new SimpleRecord(mockTime.milliseconds(), "key".getBytes, "a".getBytes),
+      new SimpleRecord(mockTime.milliseconds(), "key".getBytes, "b".getBytes)))
+    records.batches.forEach(_.setPartitionLeaderEpoch(0))
 
     val filtered = ByteBuffer.allocate(2048)
     records.filterTo(new TopicPartition("foo", 0), new RecordFilter {
@@ -556,9 +1200,9 @@ 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)))
-    moreRecords.batches.asScala.foreach(_.setPartitionLeaderEpoch(0))
+      new SimpleRecord(mockTime.milliseconds(), "e".getBytes),
+      new SimpleRecord(mockTime.milliseconds(), "f".getBytes)))
+    moreRecords.batches.forEach(_.setPartitionLeaderEpoch(0))
     log.appendAsFollower(moreRecords)
 
     log.truncateTo(baseOffset + 2)
@@ -571,8 +1215,8 @@ class LogTest {
   }
 
   @Test
-  def testUpdateProducerIdMapWithCompactedData() {
-    val logConfig = createLogConfig(segmentBytes = 2048 * 5)
+  def testUpdateProducerIdMapWithCompactedData(): Unit = {
+    val logConfig = LogTest.createLogConfig(segmentBytes = 2048 * 5)
     val log = createLog(logDir, logConfig)
     val pid = 1L
     val epoch = 0.toShort
@@ -581,11 +1225,11 @@ 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)))
-    records.batches.asScala.foreach(_.setPartitionLeaderEpoch(0))
+      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.forEach(_.setPartitionLeaderEpoch(0))
 
     val filtered = ByteBuffer.allocate(2048)
     records.filterTo(new TopicPartition("foo", 0), new RecordFilter {
@@ -604,8 +1248,8 @@ class LogTest {
   }
 
   @Test
-  def testProducerIdMapTruncateTo() {
-    val logConfig = createLogConfig(segmentBytes = 2048 * 5)
+  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)
     log.appendAsLeader(TestUtils.records(List(new SimpleRecord("b".getBytes))), leaderEpoch = 0)
@@ -619,14 +1263,18 @@ class LogTest {
     assertEquals(2, log.latestProducerStateEndOffset)
 
     log.truncateTo(1)
-    assertEquals(None, log.latestProducerSnapshotOffset)
+    assertEquals(Some(1), log.latestProducerSnapshotOffset)
     assertEquals(1, log.latestProducerStateEndOffset)
+
+    log.truncateTo(0)
+    assertEquals(None, log.latestProducerSnapshotOffset)
+    assertEquals(0, log.latestProducerStateEndOffset)
   }
 
   @Test
-  def testProducerIdMapTruncateToWithNoSnapshots() {
+  def testProducerIdMapTruncateToWithNoSnapshots(): Unit = {
     // This ensures that the upgrade optimization path cannot be hit after initial loading
-    val logConfig = createLogConfig(segmentBytes = 2048 * 5)
+    val logConfig = LogTest.createLogConfig(segmentBytes = 2048 * 5)
     val log = createLog(logDir, logConfig)
     val pid = 1L
     val epoch = 0.toShort
@@ -650,7 +1298,7 @@ class LogTest {
 
   @Test
   def testLoadProducersAfterDeleteRecordsMidSegment(): Unit = {
-    val logConfig = createLogConfig(segmentBytes = 2048 * 5)
+    val logConfig = LogTest.createLogConfig(segmentBytes = 2048 * 5)
     val log = createLog(logDir, logConfig)
     val pid1 = 1L
     val pid2 = 2L
@@ -662,24 +1310,159 @@ class LogTest {
       producerEpoch = epoch, sequence = 0), leaderEpoch = 0)
     assertEquals(2, log.activeProducersWithLastSequence.size)
 
-    log.maybeIncrementLogStartOffset(1L)
+    log.updateHighWatermark(log.logEndOffset)
+    log.maybeIncrementLogStartOffset(1L, ClientRecordDeletion)
 
-    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()
 
-    val reloadedLog = createLog(logDir, logConfig, logStartOffset = 1L)
-    assertEquals(1, reloadedLog.activeProducersWithLastSequence.size)
+    // 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, lastShutdownClean = false)
+    assertEquals(2, reloadedLog.activeProducersWithLastSequence.size)
     val reloadedLastSeqOpt = log.activeProducersWithLastSequence.get(pid2)
     assertEquals(retainedLastSeqOpt, reloadedLastSeqOpt)
   }
 
+  @Test
+  def testRetentionDeletesProducerStateSnapshots(): Unit = {
+    val logConfig = LogTest.createLogConfig(segmentBytes = 2048 * 5, retentionBytes = 0, retentionMs = 1000 * 60, fileDeleteDelayMs = 0)
+    val log = createLog(logDir, logConfig)
+    val pid1 = 1L
+    val epoch = 0.toShort
+
+    log.appendAsLeader(TestUtils.records(List(new SimpleRecord("a".getBytes)), producerId = pid1,
+      producerEpoch = epoch, sequence = 0), leaderEpoch = 0)
+    log.roll()
+    log.appendAsLeader(TestUtils.records(List(new SimpleRecord("b".getBytes)), producerId = pid1,
+      producerEpoch = epoch, sequence = 1), leaderEpoch = 0)
+    log.roll()
+    log.appendAsLeader(TestUtils.records(List(new SimpleRecord("c".getBytes)), producerId = pid1,
+      producerEpoch = epoch, sequence = 2), leaderEpoch = 0)
+
+    log.updateHighWatermark(log.logEndOffset)
+
+    assertEquals(2, ProducerStateManager.listSnapshotFiles(logDir).size)
+    // Sleep to breach the retention period
+    mockTime.sleep(1000 * 60 + 1)
+    log.deleteOldSegments()
+    // Sleep to breach the file delete delay and run scheduled file deletion tasks
+    mockTime.sleep(1)
+    assertEquals("expect a single producer state snapshot remaining", 1, ProducerStateManager.listSnapshotFiles(logDir).size)
+  }
+
+  @Test
+  def testLogStartOffsetMovementDeletesSnapshots(): Unit = {
+    val logConfig = LogTest.createLogConfig(segmentBytes = 2048 * 5, retentionBytes = -1, fileDeleteDelayMs = 0)
+    val log = createLog(logDir, logConfig)
+    val pid1 = 1L
+    val epoch = 0.toShort
+
+    log.appendAsLeader(TestUtils.records(List(new SimpleRecord("a".getBytes)), producerId = pid1,
+      producerEpoch = epoch, sequence = 0), leaderEpoch = 0)
+    log.roll()
+    log.appendAsLeader(TestUtils.records(List(new SimpleRecord("b".getBytes)), producerId = pid1,
+      producerEpoch = epoch, sequence = 1), leaderEpoch = 0)
+    log.roll()
+    log.appendAsLeader(TestUtils.records(List(new SimpleRecord("c".getBytes)), producerId = pid1,
+      producerEpoch = epoch, sequence = 2), leaderEpoch = 0)
+    log.updateHighWatermark(log.logEndOffset)
+    assertEquals(2, ProducerStateManager.listSnapshotFiles(logDir).size)
+
+    // Increment the log start offset to exclude the first two segments.
+    log.maybeIncrementLogStartOffset(log.logEndOffset - 1, ClientRecordDeletion)
+    log.deleteOldSegments()
+    // Sleep to breach the file delete delay and run scheduled file deletion tasks
+    mockTime.sleep(1)
+    assertEquals("expect a single producer state snapshot remaining", 1, ProducerStateManager.listSnapshotFiles(logDir).size)
+  }
+
+  @Test
+  def testCompactionDeletesProducerStateSnapshots(): Unit = {
+    val logConfig = LogTest.createLogConfig(segmentBytes = 2048 * 5, cleanupPolicy = LogConfig.Compact, fileDeleteDelayMs = 0)
+    val log = createLog(logDir, logConfig)
+    val pid1 = 1L
+    val epoch = 0.toShort
+    val cleaner = new Cleaner(id = 0,
+      offsetMap = new FakeOffsetMap(Int.MaxValue),
+      ioBufferSize = 64 * 1024,
+      maxIoBufferSize = 64 * 1024,
+      dupBufferLoadFactor = 0.75,
+      throttler = new Throttler(Double.MaxValue, Long.MaxValue, false, time = mockTime),
+      time = mockTime,
+      checkDone = _ => {})
+
+    log.appendAsLeader(TestUtils.records(List(new SimpleRecord("a".getBytes, "a".getBytes())), producerId = pid1,
+      producerEpoch = epoch, sequence = 0), leaderEpoch = 0)
+    log.roll()
+    log.appendAsLeader(TestUtils.records(List(new SimpleRecord("a".getBytes, "b".getBytes())), producerId = pid1,
+      producerEpoch = epoch, sequence = 1), leaderEpoch = 0)
+    log.roll()
+    log.appendAsLeader(TestUtils.records(List(new SimpleRecord("a".getBytes, "c".getBytes())), producerId = pid1,
+      producerEpoch = epoch, sequence = 2), leaderEpoch = 0)
+    log.updateHighWatermark(log.logEndOffset)
+    assertEquals("expected a snapshot file per segment base offset, except the first segment", log.logSegments.map(_.baseOffset).toSeq.sorted.drop(1), ProducerStateManager.listSnapshotFiles(logDir).map(_.offset).sorted)
+    assertEquals(2, ProducerStateManager.listSnapshotFiles(logDir).size)
+
+    // Clean segments, this should delete everything except the active segment since there only
+    // exists the key "a".
+    cleaner.clean(LogToClean(log.topicPartition, log, 0, log.logEndOffset))
+    log.deleteOldSegments()
+    // Sleep to breach the file delete delay and run scheduled file deletion tasks
+    mockTime.sleep(1)
+    assertEquals("expected a snapshot file per segment base offset, excluding the first", log.logSegments.map(_.baseOffset).toSeq.sorted.drop(1), ProducerStateManager.listSnapshotFiles(logDir).map(_.offset).sorted)
+  }
+
+  /**
+   * After loading the log, producer state is truncated such that there are no producer state snapshot files which
+   * exceed the log end offset. This test verifies that these are removed.
+   */
+  @Test
+  def testLoadingLogDeletesProducerStateSnapshotsPastLogEndOffset(): Unit = {
+    val straySnapshotFile = Log.producerSnapshotFile(logDir, 42).toPath
+    Files.createFile(straySnapshotFile)
+    val logConfig = LogTest.createLogConfig(segmentBytes = 2048 * 5, retentionBytes = -1, fileDeleteDelayMs = 0)
+    createLog(logDir, logConfig)
+    assertEquals("expected producer state snapshots greater than the log end offset to be cleaned up", 0, ProducerStateManager.listSnapshotFiles(logDir).size)
+  }
+
+  @Test
+  def testLoadingLogKeepsLargestStrayProducerStateSnapshot(): Unit = {
+    val logConfig = LogTest.createLogConfig(segmentBytes = 2048 * 5, retentionBytes = 0, retentionMs = 1000 * 60, fileDeleteDelayMs = 0)
+    val log = createLog(logDir, logConfig)
+    val pid1 = 1L
+    val epoch = 0.toShort
+
+    log.appendAsLeader(TestUtils.records(List(new SimpleRecord("a".getBytes)), producerId = pid1, producerEpoch = epoch, sequence = 0), leaderEpoch = 0)
+    log.roll()
+    log.appendAsLeader(TestUtils.records(List(new SimpleRecord("b".getBytes)), producerId = pid1, producerEpoch = epoch, sequence = 1), leaderEpoch = 0)
+    log.roll()
+
+    log.appendAsLeader(TestUtils.records(List(new SimpleRecord("c".getBytes)), producerId = pid1, producerEpoch = epoch, sequence = 2), leaderEpoch = 0)
+    log.appendAsLeader(TestUtils.records(List(new SimpleRecord("d".getBytes)), producerId = pid1, producerEpoch = epoch, sequence = 3), leaderEpoch = 0)
+
+    // Close the log, we should now have 3 segments
+    log.close()
+    assertEquals(log.logSegments.size, 3)
+    // We expect 3 snapshot files, two of which are for the first two segments, the last was written out during log closing.
+    assertEquals(Seq(1, 2, 4), ProducerStateManager.listSnapshotFiles(logDir).map(_.offset).sorted)
+    // Inject a stray snapshot file within the bounds of the log at offset 3, it should be cleaned up after loading the log
+    val straySnapshotFile = Log.producerSnapshotFile(logDir, 3).toPath
+    Files.createFile(straySnapshotFile)
+    assertEquals(Seq(1, 2, 3, 4), ProducerStateManager.listSnapshotFiles(logDir).map(_.offset).sorted)
+
+    createLog(logDir, logConfig, lastShutdownClean = false)
+    // We should clean up the stray producer state snapshot file, but keep the largest snapshot file (4)
+    assertEquals(Seq(1, 2, 4), ProducerStateManager.listSnapshotFiles(logDir).map(_.offset).sorted)
+  }
+
   @Test
   def testLoadProducersAfterDeleteRecordsOnSegment(): Unit = {
-    val logConfig = createLogConfig(segmentBytes = 2048 * 5)
+    val logConfig = LogTest.createLogConfig(segmentBytes = 2048 * 5)
     val log = createLog(logDir, logConfig)
     val pid1 = 1L
     val pid2 = 2L
@@ -694,28 +1477,30 @@ class LogTest {
     assertEquals(2, log.logSegments.size)
     assertEquals(2, log.activeProducersWithLastSequence.size)
 
-    log.maybeIncrementLogStartOffset(1L)
-    log.onHighWatermarkIncremented(log.logEndOffset)
+    log.updateHighWatermark(log.logEndOffset)
+    log.maybeIncrementLogStartOffset(1L, ClientRecordDeletion)
     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()
 
-    val reloadedLog = createLog(logDir, logConfig, logStartOffset = 1L)
+    // After reloading log, producer state should not be regenerated
+    val reloadedLog = createLog(logDir, logConfig, logStartOffset = 1L, lastShutdownClean = false)
     assertEquals(1, reloadedLog.activeProducersWithLastSequence.size)
     val reloadedEntryOpt = log.activeProducersWithLastSequence.get(pid2)
     assertEquals(retainedLastSeqOpt, reloadedEntryOpt)
   }
 
   @Test
-  def testProducerIdMapTruncateFullyAndStartAt() {
+  def testProducerIdMapTruncateFullyAndStartAt(): Unit = {
     val records = TestUtils.singletonRecords("foo".getBytes)
-    val logConfig = createLogConfig(segmentBytes = records.sizeInBytes, retentionBytes = records.sizeInBytes * 2)
+    val logConfig = LogTest.createLogConfig(segmentBytes = records.sizeInBytes, retentionBytes = records.sizeInBytes * 2)
     val log = createLog(logDir, logConfig)
     log.appendAsLeader(records, leaderEpoch = 0)
     log.takeProducerSnapshot()
@@ -735,10 +1520,10 @@ 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 = createLogConfig(segmentBytes = records.sizeInBytes, retentionBytes = records.sizeInBytes * 2)
+    val logConfig = LogTest.createLogConfig(segmentBytes = records.sizeInBytes, retentionBytes = records.sizeInBytes * 2)
     val log = createLog(logDir, logConfig)
     log.appendAsLeader(records, leaderEpoch = 0)
     log.takeProducerSnapshot()
@@ -753,48 +1538,77 @@ 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() {
-    val logConfig = createLogConfig(segmentBytes = 2048 * 5)
+  def testTakeSnapshotOnRollAndDeleteSnapshotOnRecoveryPointCheckpoint(): Unit = {
+    val logConfig = LogTest.createLogConfig(segmentBytes = 2048 * 5)
     val log = createLog(logDir, logConfig)
     log.appendAsLeader(TestUtils.singletonRecords("a".getBytes), leaderEpoch = 0)
-    log.roll(1L)
+    log.roll(Some(1L))
     assertEquals(Some(1L), log.latestProducerSnapshotOffset)
     assertEquals(Some(1L), log.oldestProducerSnapshotOffset)
 
     log.appendAsLeader(TestUtils.singletonRecords("b".getBytes), leaderEpoch = 0)
-    log.roll(2L)
+    log.roll(Some(2L))
     assertEquals(Some(2L), log.latestProducerSnapshotOffset)
     assertEquals(Some(1L), log.oldestProducerSnapshotOffset)
 
     log.appendAsLeader(TestUtils.singletonRecords("c".getBytes), leaderEpoch = 0)
-    log.roll(3L)
+    log.roll(Some(3L))
     assertEquals(Some(3L), log.latestProducerSnapshotOffset)
 
     // roll triggers a flush at the starting offset of the new segment, we should retain all snapshots
     assertEquals(Some(1L), log.oldestProducerSnapshotOffset)
 
-    // retain the snapshots from the active segment and the previous segment, delete the oldest one
-    log.deleteSnapshotsAfterRecoveryPointCheckpoint()
-    assertEquals(Some(2L), log.oldestProducerSnapshotOffset)
-
     // even if we flush within the active segment, the snapshot should remain
     log.appendAsLeader(TestUtils.singletonRecords("baz".getBytes), leaderEpoch = 0)
     log.flush(4L)
     assertEquals(Some(3L), log.latestProducerSnapshotOffset)
-    assertEquals(Some(2L), log.oldestProducerSnapshotOffset)
+  }
+
+  @Test
+  def testProducerSnapshotAfterSegmentRollOnAppend(): Unit = {
+    val producerId = 1L
+    val logConfig = LogTest.createLogConfig(segmentBytes = 1024)
+    val log = createLog(logDir, logConfig)
+
+    log.appendAsLeader(TestUtils.records(Seq(new SimpleRecord(mockTime.milliseconds(), new Array[Byte](512))),
+      producerId = producerId, producerEpoch = 0, sequence = 0),
+      leaderEpoch = 0)
+
+    // The next append should overflow the segment and cause it to roll
+    log.appendAsLeader(TestUtils.records(Seq(new SimpleRecord(mockTime.milliseconds(), new Array[Byte](512))),
+      producerId = producerId, producerEpoch = 0, sequence = 1),
+      leaderEpoch = 0)
+
+    assertEquals(2, log.logSegments.size)
+    assertEquals(1L, log.activeSegment.baseOffset)
+    assertEquals(Some(1L), log.latestProducerSnapshotOffset)
+
+    // Force a reload from the snapshot to check its consistency
+    log.truncateTo(1L)
+
+    assertEquals(2, log.logSegments.size)
+    assertEquals(1L, log.activeSegment.baseOffset)
+    assertTrue(log.activeSegment.log.batches.asScala.isEmpty)
+    assertEquals(Some(1L), log.latestProducerSnapshotOffset)
+
+    val lastEntry = log.producerStateManager.lastEntry(producerId)
+    assertTrue(lastEntry.isDefined)
+    assertEquals(0L, lastEntry.get.firstDataOffset)
+    assertEquals(0L, lastEntry.get.lastDataOffset)
   }
 
   @Test
   def testRebuildTransactionalState(): Unit = {
-    val logConfig = createLogConfig(segmentBytes = 1024 * 1024 * 5)
+    val logConfig = LogTest.createLogConfig(segmentBytes = 1024 * 1024 * 5)
     val log = createLog(logDir, logConfig)
 
     val pid = 137L
@@ -807,17 +1621,16 @@ 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)
 
     log.close()
 
-    val reopenedLog = createLog(logDir, logConfig)
-    reopenedLog.onHighWatermarkIncremented(commitAppendInfo.lastOffset + 1)
+    val reopenedLog = createLog(logDir, logConfig, lastShutdownClean = false)
+    reopenedLog.updateHighWatermark(abortAppendInfo.lastOffset + 1)
     assertEquals(None, reopenedLog.firstUnstableOffset)
   }
 
@@ -825,19 +1638,20 @@ class LogTest {
                             producerId: Long,
                             epoch: Short,
                             offset: Long = 0L,
-                            coordinatorEpoch: Int = 0,
-                            partitionLeaderEpoch: Int = 0): MemoryRecords = {
+                            coordinatorEpoch: Int,
+                            partitionLeaderEpoch: Int = 0,
+                            timestamp: Long): 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
 
     val pid = 23L
-    val logConfig = createLogConfig(segmentBytes = 2048 * 5)
+    val logConfig = LogTest.createLogConfig(segmentBytes = 2048 * 5)
     val log = createLog(logDir, logConfig, maxProducerIdExpirationMs = maxProducerIdExpirationMs,
       producerIdExpirationCheckIntervalMs = producerIdExpirationCheckIntervalMs)
     val records = Seq(new SimpleRecord(mockTime.milliseconds(), "foo".getBytes))
@@ -874,12 +1688,12 @@ class LogTest {
       new SimpleRecord(mockTime.milliseconds, s"key-$seq".getBytes, s"value-$seq".getBytes)
     ), producerId = pid, producerEpoch = epoch, sequence = seq)
     val multiEntryAppendInfo = log.appendAsLeader(createRecords, leaderEpoch = 0)
-    assertEquals("should have appended 3 entries", multiEntryAppendInfo.lastOffset - multiEntryAppendInfo.firstOffset + 1, 3)
+    assertEquals("should have appended 3 entries", multiEntryAppendInfo.lastOffset - multiEntryAppendInfo.firstOffset.get + 1, 3)
 
     // Append a Duplicate of the tail, when the entry at the tail has multiple records.
     val dupMultiEntryAppendInfo = log.appendAsLeader(createRecords, leaderEpoch = 0)
     assertEquals("Somehow appended a duplicate entry with multiple log records to the tail",
-      multiEntryAppendInfo.firstOffset, dupMultiEntryAppendInfo.firstOffset)
+      multiEntryAppendInfo.firstOffset.get, dupMultiEntryAppendInfo.firstOffset.get)
     assertEquals("Somehow appended a duplicate entry with multiple log records to the tail",
       multiEntryAppendInfo.lastOffset, dupMultiEntryAppendInfo.lastOffset)
 
@@ -893,7 +1707,7 @@ class LogTest {
           new SimpleRecord(mockTime.milliseconds, s"key-$seq".getBytes, s"value-$seq".getBytes)),
         producerId = pid, producerEpoch = epoch, sequence = seq - 2)
       log.appendAsLeader(records, leaderEpoch = 0)
-      fail ("Should have received an OutOfOrderSequenceException since we attempted to append a duplicate of a records " +
+      fail("Should have received an OutOfOrderSequenceException since we attempted to append a duplicate of a records " +
         "in the middle of the log.")
     } catch {
       case _: OutOfOrderSequenceException => // Good!
@@ -905,16 +1719,16 @@ class LogTest {
       producerId = pid, producerEpoch = epoch, sequence = 2)
     log.appendAsLeader(duplicateOfFourth, leaderEpoch = 0)
 
-    // Append a Duplicate of an entry older than the last 5 appended batches. This should result in a DuplicateSequenceNumberException.
-     try {
+    // Duplicates at older entries are reported as OutOfOrderSequence errors
+    try {
       val records = TestUtils.records(
         List(new SimpleRecord(mockTime.milliseconds, s"key-1".getBytes, s"value-1".getBytes)),
         producerId = pid, producerEpoch = epoch, sequence = 1)
       log.appendAsLeader(records, leaderEpoch = 0)
-      fail ("Should have received an DuplicateSequenceNumberException since we attempted to append a duplicate of a batch" +
+      fail("Should have received an OutOfOrderSequenceException since we attempted to append a duplicate of a batch " +
         "which is older than the last 5 appended batches.")
     } catch {
-      case _: DuplicateSequenceException => // Good!
+      case _: OutOfOrderSequenceException => // Good!
     }
 
     // Append a duplicate entry with a single records at the tail of the log. This should return the appendInfo of the original entry.
@@ -922,12 +1736,12 @@ class LogTest {
       producerId = pid, producerEpoch = epoch, sequence = seq)
     val origAppendInfo = log.appendAsLeader(createRecordsWithDuplicate, leaderEpoch = 0)
     val newAppendInfo = log.appendAsLeader(createRecordsWithDuplicate, leaderEpoch = 0)
-    assertEquals("Inserted a duplicate records into the log", origAppendInfo.firstOffset, newAppendInfo.firstOffset)
+    assertEquals("Inserted a duplicate records into the log", origAppendInfo.firstOffset.get, newAppendInfo.firstOffset.get)
     assertEquals("Inserted a duplicate records into the log", origAppendInfo.lastOffset, newAppendInfo.lastOffset)
   }
 
   @Test
-  def testMultipleProducerIdsPerMemoryRecord() : Unit = {
+  def testMultipleProducerIdsPerMemoryRecord(): Unit = {
     // create a log
     val log = createLog(logDir, LogConfig())
 
@@ -960,7 +1774,7 @@ class LogTest {
     log.appendAsFollower(memoryRecords)
     log.flush()
 
-    val fetchedData = log.readUncommitted(0, Int.MaxValue)
+    val fetchedData = readLog(log, 0, Int.MaxValue)
 
     val origIterator = memoryRecords.batches.iterator()
     for (batch <- fetchedData.records.batches.asScala) {
@@ -973,8 +1787,8 @@ class LogTest {
   }
 
   @Test
-  def testDuplicateAppendToFollower() : Unit = {
-    val logConfig = createLogConfig(segmentBytes = 1024 * 1024 * 5)
+  def testDuplicateAppendToFollower(): Unit = {
+    val logConfig = LogTest.createLogConfig(segmentBytes = 1024 * 1024 * 5)
     val log = createLog(logDir, logConfig)
     val epoch: Short = 0
     val pid = 1L
@@ -994,8 +1808,8 @@ class LogTest {
   }
 
   @Test
-  def testMultipleProducersWithDuplicatesInSingleAppend() : Unit = {
-    val logConfig = createLogConfig(segmentBytes = 1024 * 1024 * 5)
+  def testMultipleProducersWithDuplicatesInSingleAppend(): Unit = {
+    val logConfig = LogTest.createLogConfig(segmentBytes = 1024 * 1024 * 5)
     val log = createLog(logDir, logConfig)
 
     val pid1 = 1L
@@ -1037,7 +1851,7 @@ class LogTest {
     buffer.flip()
 
     val records = MemoryRecords.readableRecords(buffer)
-    records.batches.asScala.foreach(_.setPartitionLeaderEpoch(0))
+    records.batches.forEach(_.setPartitionLeaderEpoch(0))
 
     // Ensure that batches with duplicates are accepted on the follower.
     assertEquals(0L, log.logEndOffset)
@@ -1045,7 +1859,7 @@ class LogTest {
     assertEquals(5L, log.logEndOffset)
   }
 
-  @Test(expected = classOf[ProducerFencedException])
+  @Test(expected = classOf[InvalidProducerEpochException])
   def testOldProducerEpoch(): Unit = {
     // create a log
     val log = createLog(logDir, LogConfig())
@@ -1060,16 +1874,47 @@ 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.dir).size)
+
+    log.updateHighWatermark(log.logEndOffset)
+    log.maybeIncrementLogStartOffset(2L, ClientRecordDeletion)
+    log.deleteOldSegments() // force retention to kick in so that the snapshot files are cleaned up.
+    mockTime.sleep(logConfig.fileDeleteDelayMs + 1000) // advance the clock so file deletion takes place
+
+    // Deleting records should not remove producer state but should delete snapshots after the file deletion delay.
+    assertEquals(2, log.activeProducersWithLastSequence.size)
+    assertEquals(1, ProducerStateManager.listSnapshotFiles(log.dir).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
-    val logConfig = createLogConfig(segmentMs = 1 * 60 * 60L, segmentJitterMs = maxJitter)
+    val logConfig = LogTest.createLogConfig(segmentMs = 1 * 60 * 60L, segmentJitterMs = maxJitter)
     val log = createLog(logDir, logConfig)
     assertEquals("Log begins with a single empty segment.", 1, log.numberOfSegments)
     log.appendAsLeader(set, leaderEpoch = 0)
@@ -1088,13 +1933,13 @@ 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
     val segmentSize = msgPerSeg * (setSize - 1) // each segment will be 10 messages
     // create a log
-    val logConfig = createLogConfig(segmentBytes = segmentSize)
+    val logConfig = LogTest.createLogConfig(segmentBytes = segmentSize)
     val log = createLog(logDir, logConfig)
     assertEquals("There should be exactly 1 segment.", 1, log.numberOfSegments)
 
@@ -1108,7 +1953,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)
@@ -1118,8 +1963,8 @@ 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() {
-    val logConfig = createLogConfig(segmentBytes = 71)
+  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
 
@@ -1127,23 +1972,23 @@ class LogTest {
       log.appendAsLeader(TestUtils.singletonRecords(value = value), leaderEpoch = 0)
 
     for(i <- values.indices) {
-      val read = log.readUncommitted(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)
       assertEquals("Values not equal", ByteBuffer.wrap(values(i)), actual.value)
     }
     assertEquals("Reading beyond the last message returns nothing.", 0,
-      log.readUncommitted(values.length, 100, None).records.batches.asScala.size)
+      readLog(log, values.length, 100).records.batches.asScala.size)
   }
 
   /**
-   * This test appends a bunch of messages with non-sequential offsets and checks that we can read the correct message
+   * This test appends a bunch of messages with non-sequential offsets and checks that we can an the correct message
    * from any offset less than the logEndOffset including offsets not appended.
    */
   @Test
-  def testAppendAndReadWithNonSequentialOffsets() {
-    val logConfig = createLogConfig(segmentBytes = 72)
+  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
     val records = messageIds.map(id => new SimpleRecord(id.toString.getBytes))
@@ -1153,7 +1998,7 @@ class LogTest {
       log.appendAsFollower(MemoryRecords.withRecords(messageIds(i), CompressionType.NONE, 0, records(i)))
     for(i <- 50 until messageIds.max) {
       val idx = messageIds.indexWhere(_ >= i)
-      val read = log.readUncommitted(i, 100, None).records.records.iterator.next()
+      val read = readLog(log, i, 100).records.records.iterator.next()
       assertEquals("Offset read should match message id.", messageIds(idx), read.offset)
       assertEquals("Message should match appended.", records(idx), new SimpleRecord(read))
     }
@@ -1166,8 +2011,8 @@ class LogTest {
    * first segment has the greatest lower bound on the offset.
    */
   @Test
-  def testReadAtLogGap() {
-    val logConfig = createLogConfig(segmentBytes = 300)
+  def testReadAtLogGap(): Unit = {
+    val logConfig = LogTest.createLogConfig(segmentBytes = 300)
     val log = createLog(logDir, logConfig)
 
     // keep appending until we have two segments with only a single message in the second segment
@@ -1178,20 +2023,20 @@ class LogTest {
     log.logSegments.head.truncateTo(1)
 
     assertEquals("A read should now return the last message in the log", log.logEndOffset - 1,
-      log.readUncommitted(1, 200, None).records.batches.iterator.next().lastOffset)
+      readLog(log, 1, 200).records.batches.iterator.next().lastOffset)
   }
 
   @Test(expected = classOf[KafkaStorageException])
-  def testLogRollAfterLogHandlerClosed() {
-    val logConfig = createLogConfig()
+  def testLogRollAfterLogHandlerClosed(): Unit = {
+    val logConfig = LogTest.createLogConfig()
     val log = createLog(logDir,  logConfig)
     log.closeHandlers()
-    log.roll(1)
+    log.roll(Some(1L))
   }
 
   @Test
-  def testReadWithMinMessage() {
-    val logConfig = createLogConfig(segmentBytes = 72)
+  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
     val records = messageIds.map(id => new SimpleRecord(id.toString.getBytes))
@@ -1203,22 +2048,20 @@ class LogTest {
     for (i <- 50 until messageIds.max) {
       val idx = messageIds.indexWhere(_ >= i)
       val reads = Seq(
-        log.readUncommitted(i, 1, minOneMessage = true),
-        log.readUncommitted(i, 100, minOneMessage = true),
-        log.readUncommitted(i, 100, Some(10000), minOneMessage = true)
+        readLog(log, i, 1),
+        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))
       }
-
-      assertEquals(Seq.empty, log.readUncommitted(i, 1, Some(1), minOneMessage = true).records.batches.asScala.toIndexedSeq)
     }
   }
 
   @Test
-  def testReadWithTooSmallMaxLength() {
-    val logConfig = createLogConfig(segmentBytes = 72)
+  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
     val records = messageIds.map(id => new SimpleRecord(id.toString.getBytes))
@@ -1228,14 +2071,14 @@ class LogTest {
       log.appendAsFollower(MemoryRecords.withRecords(messageIds(i), CompressionType.NONE, 0, records(i)))
 
     for (i <- 50 until messageIds.max) {
-      assertEquals(MemoryRecords.EMPTY, log.readUncommitted(i, 0).records)
+      assertEquals(MemoryRecords.EMPTY, readLog(log, i, maxLength = 0, minOneMessage = false).records)
 
       // we return an incomplete message instead of an empty one for the case below
       // we use this mechanism to tell consumers of the fetch request version 2 and below that the message size is
       // larger than the fetch size
       // in fetch request version 3, we no longer need this as we return oversized messages from the first non-empty
       // partition
-      val fetchInfo = log.readUncommitted(i, 1)
+      val fetchInfo = readLog(log, i, maxLength = 1, minOneMessage = false)
       assertTrue(fetchInfo.firstEntryIncomplete)
       assertTrue(fetchInfo.records.isInstanceOf[FileRecords])
       assertEquals(1, fetchInfo.records.sizeInBytes)
@@ -1249,32 +2092,29 @@ 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 = createLogConfig(segmentBytes = 1024)
+    val logConfig = LogTest.createLogConfig(segmentBytes = 1024)
     val log = createLog(logDir, logConfig)
     log.appendAsLeader(TestUtils.singletonRecords(value = "42".getBytes), leaderEpoch = 0)
 
     assertEquals("Reading at the log end offset should produce 0 byte read.", 0,
-      log.readUncommitted(1025, 1000).records.sizeInBytes)
+      readLog(log, 1025, 1000).records.sizeInBytes)
 
     try {
-      log.readUncommitted(0, 1000)
+      readLog(log, 0, 1000)
       fail("Reading below the log start offset should throw OffsetOutOfRangeException")
     } catch {
       case _: OffsetOutOfRangeException => // This is good.
     }
 
     try {
-      log.readUncommitted(1026, 1000)
+      readLog(log, 1026, 1000)
       fail("Reading at beyond the log end offset should throw OffsetOutOfRangeException")
     } catch {
       case _: OffsetOutOfRangeException => // This is good.
     }
-
-    assertEquals("Reading from below the specified maxOffset should produce 0 byte read.", 0,
-      log.readUncommitted(1025, 1000, Some(1024)).records.sizeInBytes)
   }
 
   /**
@@ -1282,9 +2122,9 @@ 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 = createLogConfig(segmentBytes = 100)
+    val logConfig = LogTest.createLogConfig(segmentBytes = 100)
     val log = createLog(logDir, logConfig)
     val numMessages = 100
     val messageSets = (0 until numMessages).map(i => TestUtils.singletonRecords(value = i.toString.getBytes,
@@ -1295,7 +2135,7 @@ class LogTest {
     /* do successive reads to ensure all our messages are there */
     var offset = 0L
     for(i <- 0 until numMessages) {
-      val messages = log.readUncommitted(offset, 1024*1024).records.batches
+      val messages = readLog(log, offset, 1024*1024).records.batches
       val head = messages.iterator.next()
       assertEquals("Offsets not equal", offset, head.lastOffset)
 
@@ -1306,8 +2146,7 @@ class LogTest {
       assertEquals(s"Timestamps not equal at offset $offset", expected.timestamp, actual.timestamp)
       offset = head.lastOffset + 1
     }
-    val lastRead = log.readUncommitted(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
@@ -1320,16 +2159,16 @@ 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 = createLogConfig(segmentBytes = 110)
+    val logConfig = LogTest.createLogConfig(segmentBytes = 110)
     val log = createLog(logDir, logConfig)
 
     /* append 2 compressed message sets, each with two messages giving offsets 0, 1, 2, 3 */
     log.appendAsLeader(MemoryRecords.withRecords(CompressionType.GZIP, new SimpleRecord("hello".getBytes), new SimpleRecord("there".getBytes)), leaderEpoch = 0)
     log.appendAsLeader(MemoryRecords.withRecords(CompressionType.GZIP, new SimpleRecord("alpha".getBytes), new SimpleRecord("beta".getBytes)), leaderEpoch = 0)
 
-    def read(offset: Int) = log.readUncommitted(offset, 4096).records.records
+    def read(offset: Int) = readLog(log, offset, 4096).records.records
 
     /* we should always get the first message in the compressed set when reading any offset in the set */
     assertEquals("Read at offset 0 should produce 0", 0, read(0).iterator.next().offset)
@@ -1342,11 +2181,11 @@ 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
-      val logConfig = createLogConfig(segmentBytes = 100, retentionMs = 0)
+      val logConfig = LogTest.createLogConfig(segmentBytes = 100, retentionMs = 0)
       val log = createLog(logDir, logConfig)
       for(i <- 0 until messagesToAppend)
         log.appendAsLeader(TestUtils.singletonRecords(value = i.toString.getBytes, timestamp = mockTime.milliseconds - 10), leaderEpoch = 0)
@@ -1355,7 +2194,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)
@@ -1364,7 +2203,7 @@ class LogTest {
       assertEquals("Still no change in the logEndOffset", currOffset, log.logEndOffset)
       assertEquals("Should still be able to append and should get the logEndOffset assigned to the new append",
                    currOffset,
-                   log.appendAsLeader(TestUtils.singletonRecords(value = "hello".getBytes, timestamp = mockTime.milliseconds), leaderEpoch = 0).firstOffset)
+                   log.appendAsLeader(TestUtils.singletonRecords(value = "hello".getBytes, timestamp = mockTime.milliseconds), leaderEpoch = 0).firstOffset.get)
 
       // cleanup the log
       log.delete()
@@ -1376,11 +2215,11 @@ 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
-    val logConfig = createLogConfig(segmentBytes = configSegmentSize)
+    val logConfig = LogTest.createLogConfig(segmentBytes = configSegmentSize)
     val log = createLog(logDir, logConfig)
 
     try {
@@ -1392,7 +2231,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)
@@ -1405,27 +2244,38 @@ class LogTest {
     val messageSetWithKeyedMessage = MemoryRecords.withRecords(CompressionType.NONE, keyedMessage)
     val messageSetWithKeyedMessages = MemoryRecords.withRecords(CompressionType.NONE, keyedMessage, anotherKeyedMessage)
 
-    val logConfig = createLogConfig(cleanupPolicy = LogConfig.Compact)
+    val logConfig = LogTest.createLogConfig(cleanupPolicy = LogConfig.Compact)
     val log = createLog(logDir, logConfig)
 
-    try {
+    val errorMsgPrefix = "Compacted topic cannot accept message without key"
+
+    var e = intercept[RecordValidationException] {
       log.appendAsLeader(messageSetWithUnkeyedMessage, leaderEpoch = 0)
-      fail("Compacted topics cannot accept a message without a key.")
-    } catch {
-      case _: CorruptRecordException => // this is good
     }
-    try {
+    assertTrue(e.invalidException.isInstanceOf[InvalidRecordException])
+    assertEquals(1, e.recordErrors.size)
+    assertEquals(0, e.recordErrors.head.batchIndex)
+    assertTrue(e.recordErrors.head.message.startsWith(errorMsgPrefix))
+
+    e = intercept[RecordValidationException] {
       log.appendAsLeader(messageSetWithOneUnkeyedMessage, leaderEpoch = 0)
-      fail("Compacted topics cannot accept a message without a key.")
-    } catch {
-      case _: CorruptRecordException => // this is good
     }
-    try {
+    assertTrue(e.invalidException.isInstanceOf[InvalidRecordException])
+    assertEquals(1, e.recordErrors.size)
+    assertEquals(0, e.recordErrors.head.batchIndex)
+    assertTrue(e.recordErrors.head.message.startsWith(errorMsgPrefix))
+
+    e = intercept[RecordValidationException] {
       log.appendAsLeader(messageSetWithCompressedUnkeyedMessage, leaderEpoch = 0)
-      fail("Compacted topics cannot accept a message without a key.")
-    } catch {
-      case _: CorruptRecordException => // this is good
     }
+    assertTrue(e.invalidException.isInstanceOf[InvalidRecordException])
+    assertEquals(1, e.recordErrors.size)
+    assertEquals(1, e.recordErrors.head.batchIndex)     // batch index is 1
+    assertTrue(e.recordErrors.head.message.startsWith(errorMsgPrefix))
+
+    // 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)
@@ -1438,7 +2288,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),
@@ -1446,7 +2296,7 @@ class LogTest {
 
     // append messages to log
     val maxMessageSize = second.sizeInBytes - 1
-    val logConfig = createLogConfig(maxMessageBytes = maxMessageSize)
+    val logConfig = LogTest.createLogConfig(maxMessageBytes = maxMessageSize)
     val log = createLog(logDir, logConfig)
 
     // should be able to append the small message
@@ -1459,23 +2309,39 @@ class LogTest {
       case _: RecordTooLargeException => // this is good
     }
   }
+
+  @Test
+  def testMessageSizeCheckInAppendAsFollower(): Unit = {
+    val first = MemoryRecords.withRecords(0, CompressionType.NONE, 0,
+      new SimpleRecord("You".getBytes), new SimpleRecord("bethe".getBytes))
+    val second = MemoryRecords.withRecords(5, CompressionType.NONE, 0,
+      new SimpleRecord("change (I need more bytes)... blah blah blah.".getBytes),
+      new SimpleRecord("More padding boo hoo".getBytes))
+
+    val log = createLog(logDir, LogTest.createLogConfig(maxMessageBytes = second.sizeInBytes - 1))
+
+    log.appendAsFollower(first)
+    // the second record is larger then limit but appendAsFollower does not validate the size.
+    log.appendAsFollower(second)
+  }
+
   /**
    * 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
     val indexInterval = 3 * messageSize
-    val logConfig = createLogConfig(segmentBytes = segmentSize, indexIntervalBytes = indexInterval, segmentIndexBytes = 4096)
+    val logConfig = LogTest.createLogConfig(segmentBytes = segmentSize, indexIntervalBytes = indexInterval, segmentIndexBytes = 4096)
     var log = createLog(logDir, logConfig)
     for(i <- 0 until numMessages)
       log.appendAsLeader(TestUtils.singletonRecords(value = TestUtils.randomBytes(messageSize),
         timestamp = mockTime.milliseconds + i * 10), leaderEpoch = 0)
     assertEquals("After appending %d messages to an empty log, the log end offset should be %d".format(numMessages, numMessages), numMessages, log.logEndOffset)
-    val lastIndexOffset = log.activeSegment.index.lastOffset
-    val numIndexEntries = log.activeSegment.index.entries
+    val lastIndexOffset = log.activeSegment.offsetIndex.lastOffset
+    val numIndexEntries = log.activeSegment.offsetIndex.entries
     val lastOffset = log.logEndOffset
     // After segment is closed, the last entry in the time index should be (largest timestamp -> last offset).
     val lastTimeIndexOffset = log.logEndOffset - 1
@@ -1486,22 +2352,22 @@ 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.index.lastOffset)
-      assertEquals("Should have same number of index entries as before.", numIndexEntries, log.activeSegment.index.entries)
+      assertEquals("Should have same last index offset as before.", lastIndexOffset, log.activeSegment.offsetIndex.lastOffset)
+      assertEquals("Should have same number of index entries as before.", numIndexEntries, log.activeSegment.offsetIndex.entries)
       assertEquals("Should have same last time index timestamp", lastTimeIndexTimestamp, log.activeSegment.timeIndex.lastEntry.timestamp)
       assertEquals("Should have same last time index offset", lastTimeIndexOffset, log.activeSegment.timeIndex.lastEntry.offset)
       assertEquals("Should have same number of time index entries as before.", numTimeIndexEntries, log.activeSegment.timeIndex.entries)
     }
 
-    log = createLog(logDir, logConfig, recoveryPoint = lastOffset)
+    log = createLog(logDir, logConfig, recoveryPoint = lastOffset, lastShutdownClean = false)
     verifyRecoveredLog(log, lastOffset)
     log.close()
 
     // test recovery case
-    log = createLog(logDir, logConfig)
+    log = createLog(logDir, logConfig, lastShutdownClean = false)
     verifyRecoveredLog(log, lastOffset)
     log.close()
   }
@@ -1510,9 +2376,9 @@ 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 = createLogConfig(segmentBytes = 10000, indexIntervalBytes = 1)
+    val logConfig = LogTest.createLogConfig(segmentBytes = 10000, indexIntervalBytes = 1)
     val log = createLog(logDir, logConfig)
 
     val messages = (0 until numMessages).map { i =>
@@ -1529,15 +2395,15 @@ 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 = createLogConfig(segmentBytes = 200, indexIntervalBytes = 1)
+    val logConfig = LogTest.createLogConfig(segmentBytes = 200, indexIntervalBytes = 1)
     var log = createLog(logDir, logConfig)
     for(i <- 0 until numMessages)
       log.appendAsLeader(TestUtils.singletonRecords(value = TestUtils.randomBytes(10), timestamp = mockTime.milliseconds + i * 10), leaderEpoch = 0)
-    val indexFiles = log.logSegments.map(_.index.file)
-    val timeIndexFiles = log.logSegments.map(_.timeIndex.file)
+    val indexFiles = log.logSegments.map(_.lazyOffsetIndex.file)
+    val timeIndexFiles = log.logSegments.map(_.lazyTimeIndex.file)
     log.close()
 
     // delete all the index files
@@ -1545,44 +2411,82 @@ class LogTest {
     timeIndexFiles.foreach(_.delete())
 
     // reopen the log
-    log = createLog(logDir, logConfig)
+    log = createLog(logDir, logConfig, lastShutdownClean = false)
     assertEquals("Should have %d messages when log is reopened".format(numMessages), numMessages, log.logEndOffset)
-    assertTrue("The index should have been rebuilt", log.logSegments.head.index.entries > 0)
+    assertTrue("The index should have been rebuilt", log.logSegments.head.offsetIndex.entries > 0)
     assertTrue("The time index should have been rebuilt", log.logSegments.head.timeIndex.entries > 0)
     for(i <- 0 until numMessages) {
-      assertEquals(i, log.readUncommitted(i, 100, None).records.batches.iterator.next().lastOffset)
+      assertEquals(i, readLog(log, i, 100).records.batches.iterator.next().lastOffset)
       if (i == 0)
-        assertEquals(log.logSegments.head.baseOffset, log.fetchOffsetsByTimestamp(mockTime.milliseconds + i * 10).get.offset)
+        assertEquals(log.logSegments.head.baseOffset, log.fetchOffsetByTimestamp(mockTime.milliseconds + i * 10).get.offset)
       else
-        assertEquals(i, log.fetchOffsetsByTimestamp(mockTime.milliseconds + i * 10).get.offset)
+        assertEquals(i, log.fetchOffsetByTimestamp(mockTime.milliseconds + i * 10).get.offset)
     }
     log.close()
   }
 
+  @Test
+  def testFetchOffsetByTimestampIncludesLeaderEpoch(): Unit = {
+    val logConfig = LogTest.createLogConfig(segmentBytes = 200, indexIntervalBytes = 1)
+    val log = createLog(logDir, logConfig)
+
+    assertEquals(None, log.fetchOffsetByTimestamp(0L))
+
+    val firstTimestamp = mockTime.milliseconds
+    val firstLeaderEpoch = 0
+    log.appendAsLeader(TestUtils.singletonRecords(
+      value = TestUtils.randomBytes(10),
+      timestamp = firstTimestamp),
+      leaderEpoch = firstLeaderEpoch)
+
+    val secondTimestamp = firstTimestamp + 1
+    val secondLeaderEpoch = 1
+    log.appendAsLeader(TestUtils.singletonRecords(
+      value = TestUtils.randomBytes(10),
+      timestamp = secondTimestamp),
+      leaderEpoch = secondLeaderEpoch)
+
+    assertEquals(Some(new TimestampAndOffset(firstTimestamp, 0L, Optional.of(firstLeaderEpoch))),
+      log.fetchOffsetByTimestamp(firstTimestamp))
+    assertEquals(Some(new TimestampAndOffset(secondTimestamp, 1L, Optional.of(secondLeaderEpoch))),
+      log.fetchOffsetByTimestamp(secondTimestamp))
+
+    assertEquals(Some(new TimestampAndOffset(ListOffsetResponse.UNKNOWN_TIMESTAMP, 0L, Optional.of(firstLeaderEpoch))),
+      log.fetchOffsetByTimestamp(ListOffsetRequest.EARLIEST_TIMESTAMP))
+    assertEquals(Some(new TimestampAndOffset(ListOffsetResponse.UNKNOWN_TIMESTAMP, 2L, Optional.of(secondLeaderEpoch))),
+      log.fetchOffsetByTimestamp(ListOffsetRequest.LATEST_TIMESTAMP))
+
+    // The cache can be updated directly after a leader change.
+    // The new latest offset should reflect the updated epoch.
+    log.maybeAssignEpochStartOffset(2, 2L)
+
+    assertEquals(Some(new TimestampAndOffset(ListOffsetResponse.UNKNOWN_TIMESTAMP, 2L, Optional.of(2))),
+      log.fetchOffsetByTimestamp(ListOffsetRequest.LATEST_TIMESTAMP))
+  }
+
   /**
    * 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 = createLogConfig(segmentBytes = segmentSize, indexIntervalBytes = 1, messageFormatVersion = "0.9.0")
+    val logConfig = LogTest.createLogConfig(segmentBytes = segmentSize, indexIntervalBytes = 1, messageFormatVersion = "0.9.0")
     var log = createLog(logDir, logConfig)
-    for(i <- 0 until numMessages)
+    for (i <- 0 until numMessages)
       log.appendAsLeader(TestUtils.singletonRecords(value = TestUtils.randomBytes(10),
         timestamp = mockTime.milliseconds + i * 10, magicValue = RecordBatch.MAGIC_VALUE_V1), leaderEpoch = 0)
-    val timeIndexFiles = log.logSegments.map(_.timeIndex.file)
+    val timeIndexFiles = log.logSegments.map(_.lazyTimeIndex.file)
     log.close()
 
     // Delete the time index.
-    timeIndexFiles.foreach(_.delete())
+    timeIndexFiles.foreach(file => Files.delete(file.toPath))
 
     // The rebuilt time index should be empty
-    log = createLog(logDir, logConfig, recoveryPoint = numMessages + 1)
-    val segArray = log.logSegments.toArray
-    for (i <- segArray.indices.init) {
-      assertEquals("The time index should be empty", 0, segArray(i).timeIndex.entries)
-      assertEquals("The time index file size should be 0", 0, segArray(i).timeIndex.file.length)
+    log = createLog(logDir, logConfig, recoveryPoint = numMessages + 1, lastShutdownClean = false)
+    for (segment <- log.logSegments.init) {
+      assertEquals("The time index should be empty", 0, segment.timeIndex.entries)
+      assertEquals("The time index file size should be 0", 0, segment.lazyTimeIndex.file.length)
     }
   }
 
@@ -1590,15 +2494,15 @@ 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 = createLogConfig(segmentBytes = 200, indexIntervalBytes = 1)
+    val logConfig = LogTest.createLogConfig(segmentBytes = 200, indexIntervalBytes = 1)
     var log = createLog(logDir, logConfig)
     for(i <- 0 until numMessages)
       log.appendAsLeader(TestUtils.singletonRecords(value = TestUtils.randomBytes(10), timestamp = mockTime.milliseconds + i * 10), leaderEpoch = 0)
-    val indexFiles = log.logSegments.map(_.index.file)
-    val timeIndexFiles = log.logSegments.map(_.timeIndex.file)
+    val indexFiles = log.logSegments.map(_.lazyOffsetIndex.file)
+    val timeIndexFiles = log.logSegments.map(_.lazyTimeIndex.file)
     log.close()
 
     // corrupt all the index files
@@ -1615,15 +2519,15 @@ class LogTest {
       bw.close()
     }
 
-    // reopen the log
-    log = createLog(logDir, logConfig, recoveryPoint = 200L)
+    // reopen the log with recovery point=0 so that the segment recovery can be triggered
+    log = createLog(logDir, logConfig, lastShutdownClean = false)
     assertEquals("Should have %d messages when log is reopened".format(numMessages), numMessages, log.logEndOffset)
     for(i <- 0 until numMessages) {
-      assertEquals(i, log.readUncommitted(i, 100, None).records.batches.iterator.next().lastOffset)
+      assertEquals(i, readLog(log, i, 100).records.batches.iterator.next().lastOffset)
       if (i == 0)
-        assertEquals(log.logSegments.head.baseOffset, log.fetchOffsetsByTimestamp(mockTime.milliseconds + i * 10).get.offset)
+        assertEquals(log.logSegments.head.baseOffset, log.fetchOffsetByTimestamp(mockTime.milliseconds + i * 10).get.offset)
       else
-        assertEquals(i, log.fetchOffsetsByTimestamp(mockTime.milliseconds + i * 10).get.offset)
+        assertEquals(i, log.fetchOffsetByTimestamp(mockTime.milliseconds + i * 10).get.offset)
     }
     log.close()
   }
@@ -1632,14 +2536,14 @@ 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
     val segmentSize = msgPerSeg * setSize  // each segment will be 10 messages
 
     // create a log
-    val logConfig = createLogConfig(segmentBytes = segmentSize)
+    val logConfig = LogTest.createLogConfig(segmentBytes = segmentSize)
     val log = createLog(logDir, logConfig)
     assertEquals("There should be exactly 1 segment.", 1, log.numberOfSegments)
 
@@ -1687,11 +2591,11 @@ 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
-    val logConfig = createLogConfig(segmentBytes = segmentSize, indexIntervalBytes = setSize - 1)
+    val logConfig = LogTest.createLogConfig(segmentBytes = segmentSize, indexIntervalBytes = setSize - 1)
     val log = createLog(logDir, logConfig)
     assertEquals("There should be exactly 1 segment.", 1, log.numberOfSegments)
 
@@ -1705,12 +2609,12 @@ class LogTest {
     assertEquals("There should be exactly 2 segment.", 2, log.numberOfSegments)
     val expectedEntries = msgPerSeg - 1
 
-    assertEquals(s"The index of the first segment should have $expectedEntries entries", expectedEntries, log.logSegments.toList.head.index.maxEntries)
+    assertEquals(s"The index of the first segment should have $expectedEntries entries", expectedEntries, log.logSegments.toList.head.offsetIndex.maxEntries)
     assertEquals(s"The time index of the first segment should have $expectedEntries entries", expectedEntries, log.logSegments.toList.head.timeIndex.maxEntries)
 
     log.truncateTo(0)
     assertEquals("There should be exactly 1 segment.", 1, log.numberOfSegments)
-    assertEquals("The index of segment 1 should be resized to maxIndexSize", log.config.maxIndexSize/8, log.logSegments.toList.head.index.maxEntries)
+    assertEquals("The index of segment 1 should be resized to maxIndexSize", log.config.maxIndexSize/8, log.logSegments.toList.head.offsetIndex.maxEntries)
     assertEquals("The time index of segment 1 should be resized to maxIndexSize", log.config.maxIndexSize/12, log.logSegments.toList.head.timeIndex.maxEntries)
 
     mockTime.sleep(msgPerSeg)
@@ -1723,16 +2627,25 @@ 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)
     val bogusTimeIndex2 = Log.timeIndexFile(logDir, 5)
 
+    // The files remain absent until we first access it because we are doing lazy loading for time index and offset index
+    // files but in this test case we need to create these files in order to test we will remove them.
+    bogusIndex2.createNewFile()
+    bogusTimeIndex2.createNewFile()
+
     def createRecords = TestUtils.singletonRecords(value = "test".getBytes, timestamp = mockTime.milliseconds)
-    val logConfig = createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, segmentIndexBytes = 1000, indexIntervalBytes = 1)
+    val logConfig = LogTest.createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, segmentIndexBytes = 1000, indexIntervalBytes = 1)
     val log = createLog(logDir, logConfig)
 
+    // Force the segment to access the index files because we are doing index lazy loading.
+    log.logSegments.toSeq.head.offsetIndex
+    log.logSegments.toSeq.head.timeIndex
+
     assertTrue("The first index file should have been replaced with a larger file", bogusIndex1.length > 0)
     assertTrue("The first time index file should have been replaced with a larger file", bogusTimeIndex1.length > 0)
     assertFalse("The second index file should have been deleted.", bogusIndex2.exists)
@@ -1749,17 +2662,17 @@ 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 = createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, segmentIndexBytes = 1000, indexIntervalBytes = 10000)
+    val logConfig = LogTest.createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, segmentIndexBytes = 1000, indexIntervalBytes = 10000)
     var log = createLog(logDir, logConfig)
 
     // add enough messages to roll over several segments then close and re-open and attempt to truncate
     for (_ <- 0 until 100)
       log.appendAsLeader(createRecords, leaderEpoch = 0)
     log.close()
-    log = createLog(logDir, logConfig)
+    log = createLog(logDir, logConfig, lastShutdownClean = false)
     log.truncateTo(3)
     assertEquals("All but one segment should be deleted.", 1, log.numberOfSegments)
     assertEquals("Log end offset should be 3.", 3, log.logEndOffset)
@@ -1769,10 +2682,10 @@ 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 = createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, segmentIndexBytes = 1000, indexIntervalBytes = 10000,
+    val logConfig = LogTest.createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, segmentIndexBytes = 1000, indexIntervalBytes = 10000,
                                     retentionMs = 999, fileDeleteDelayMs = asyncDeleteMs)
     val log = createLog(logDir, logConfig)
 
@@ -1782,20 +2695,20 @@ class LogTest {
 
     // files should be renamed
     val segments = log.logSegments.toArray
-    val oldFiles = segments.map(_.log.file) ++ segments.map(_.index.file)
+    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)
     assertTrue("All log and index files should end in .deleted", segments.forall(_.log.file.getName.endsWith(Log.DeletedFileSuffix)) &&
-                                                                 segments.forall(_.index.file.getName.endsWith(Log.DeletedFileSuffix)))
+                                                                 segments.forall(_.lazyOffsetIndex.file.getName.endsWith(Log.DeletedFileSuffix)))
     assertTrue("The .deleted files should still be there.", segments.forall(_.log.file.exists) &&
-                                                            segments.forall(_.index.file.exists))
+                                                            segments.forall(_.lazyOffsetIndex.file.exists))
     assertTrue("The original file should be gone.", oldFiles.forall(!_.exists))
 
     // when enough time passes the files should be deleted
-    val deletedFiles = segments.map(_.log.file) ++ segments.map(_.index.file)
+    val deletedFiles = segments.map(_.log.file) ++ segments.map(_.lazyOffsetIndex.file)
     mockTime.sleep(asyncDeleteMs + 1)
     assertTrue("Files should all be gone.", deletedFiles.forall(!_.exists))
   }
@@ -1804,9 +2717,9 @@ 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 = createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, segmentIndexBytes = 1000, retentionMs = 999)
+    val logConfig = LogTest.createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, segmentIndexBytes = 1000, retentionMs = 999)
     var log = createLog(logDir, logConfig)
 
     // append some messages to create some segments
@@ -1814,29 +2727,88 @@ 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)
+    log = createLog(logDir, logConfig, lastShutdownClean = false)
     assertEquals("The deleted segments should be gone.", 1, log.numberOfSegments)
   }
 
   @Test
-  def testAppendMessageWithNullPayload() {
+  def testAppendMessageWithNullPayload(): Unit = {
     val log = createLog(logDir, LogConfig())
     log.appendAsLeader(TestUtils.singletonRecords(value = null), leaderEpoch = 0)
-    val head = log.readUncommitted(0, 4096, None).records.records.iterator.next()
+    val head = readLog(log, 0, 4096).records.records.iterator.next()
     assertEquals(0, head.offset)
     assertTrue("Message payload should be null.", !head.hasValue)
   }
 
-  @Test(expected = classOf[IllegalArgumentException])
-  def testAppendWithOutOfOrderOffsetsThrowsException() {
+  @Test
+  def testAppendWithOutOfOrderOffsetsThrowsException(): Unit = {
+    val log = createLog(logDir, LogConfig())
+
+    val appendOffsets = Seq(0L, 1L, 3L, 2L, 4L)
+    val buffer = ByteBuffer.allocate(512)
+    for (offset <- appendOffsets) {
+      val builder = MemoryRecords.builder(buffer, RecordBatch.MAGIC_VALUE_V2, CompressionType.NONE,
+                                          TimestampType.LOG_APPEND_TIME, offset, mockTime.milliseconds(),
+                                          1L, 0, 0, false, 0)
+      builder.append(new SimpleRecord("key".getBytes, "value".getBytes))
+      builder.close()
+    }
+    buffer.flip()
+    val memoryRecords = MemoryRecords.readableRecords(buffer)
+
+    assertThrows[OffsetsOutOfOrderException] {
+      log.appendAsFollower(memoryRecords)
+    }
+  }
+
+  @Test
+  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))
-    val invalidRecord = MemoryRecords.withRecords(CompressionType.NONE, new SimpleRecord(1.toString.getBytes))
-    log.appendAsFollower(invalidRecord)
+
+    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) {
+      val invalidRecord = MemoryRecords.withRecords(magic, compression, new SimpleRecord(1.toString.getBytes))
+      withClue(s"Magic=$magic, compressionType=$compression") {
+        assertThrows[UnexpectedAppendOffsetException] {
+          log.appendAsFollower(invalidRecord)
+        }
+      }
+    }
+  }
+
+  @Test
+  def testAppendEmptyLogBelowLogStartOffsetThrowsException(): Unit = {
+    createEmptyLogs(logDir, 7)
+    val log = createLog(logDir, LogConfig(), brokerTopicStats = brokerTopicStats)
+    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) {
+      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)
+
+      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)
+        assertEquals(s"Magic=$magic, compressionType=$compression, UnexpectedAppendOffsetException#lastOffset",
+                     firstOffset + 2, exception.lastOffset)
+      }
+    }
   }
 
   @Test
@@ -1847,9 +2819,9 @@ class LogTest {
   }
 
   @Test
-  def testCorruptLog() {
+  def testCorruptLog(): Unit = {
     // append some messages to create some segments
-    val logConfig = createLogConfig(segmentBytes = 1000, indexIntervalBytes = 1, maxMessageBytes = 64 * 1024)
+    val logConfig = LogTest.createLogConfig(segmentBytes = 1000, indexIntervalBytes = 1, maxMessageBytes = 64 * 1024)
     def createRecords = TestUtils.singletonRecords(value = "test".getBytes, timestamp = mockTime.milliseconds)
     val recoveryPoint = 50L
     for (_ <- 0 until 10) {
@@ -1863,11 +2835,11 @@ class LogTest {
       log.close()
 
       // corrupt index and log by appending random bytes
-      TestUtils.appendNonsenseToFile(log.activeSegment.index.file, TestUtils.random.nextInt(1024) + 1)
+      TestUtils.appendNonsenseToFile(log.activeSegment.lazyOffsetIndex.file, TestUtils.random.nextInt(1024) + 1)
       TestUtils.appendNonsenseToFile(log.activeSegment.log.file, TestUtils.random.nextInt(1024) + 1)
 
       // attempt recovery
-      log = createLog(logDir, logConfig, 0L, recoveryPoint)
+      log = createLog(logDir, logConfig, brokerTopicStats, 0L, recoveryPoint, lastShutdownClean = false)
       assertEquals(numMessages, log.logEndOffset)
 
       val recovered = log.logSegments.flatMap(_.log.records.asScala.toList).toList
@@ -1881,52 +2853,442 @@ class LogTest {
         assertEquals(s"Timestamps not equal", expected.timestamp, actual.timestamp)
       }
 
-      Utils.delete(logDir)
-    }
+      Utils.delete(logDir)
+    }
+  }
+
+  @Test
+  def testOverCompactedLogRecovery(): Unit = {
+    // append some messages to create some segments
+    val logConfig = LogTest.createLogConfig(segmentBytes = 1000, indexIntervalBytes = 1, maxMessageBytes = 64 * 1024)
+    val log = createLog(logDir, logConfig)
+    val set1 = MemoryRecords.withRecords(0, CompressionType.NONE, 0, new SimpleRecord("v1".getBytes(), "k1".getBytes()))
+    val set2 = MemoryRecords.withRecords(Integer.MAX_VALUE.toLong + 2, CompressionType.NONE, 0, new SimpleRecord("v3".getBytes(), "k3".getBytes()))
+    val set3 = MemoryRecords.withRecords(Integer.MAX_VALUE.toLong + 3, CompressionType.NONE, 0, new SimpleRecord("v4".getBytes(), "k4".getBytes()))
+    val set4 = MemoryRecords.withRecords(Integer.MAX_VALUE.toLong + 4, CompressionType.NONE, 0, new SimpleRecord("v5".getBytes(), "k5".getBytes()))
+    //Writes into an empty log with baseOffset 0
+    log.appendAsFollower(set1)
+    assertEquals(0L, log.activeSegment.baseOffset)
+    //This write will roll the segment, yielding a new segment with base offset = max(1, Integer.MAX_VALUE+2) = Integer.MAX_VALUE+2
+    log.appendAsFollower(set2)
+    assertEquals(Integer.MAX_VALUE.toLong + 2, log.activeSegment.baseOffset)
+    assertTrue(Log.producerSnapshotFile(logDir, Integer.MAX_VALUE.toLong + 2).exists)
+    //This will go into the existing log
+    log.appendAsFollower(set3)
+    assertEquals(Integer.MAX_VALUE.toLong + 2, log.activeSegment.baseOffset)
+    //This will go into the existing log
+    log.appendAsFollower(set4)
+    assertEquals(Integer.MAX_VALUE.toLong + 2, log.activeSegment.baseOffset)
+    log.close()
+    val indexFiles = logDir.listFiles.filter(file => file.getName.contains(".index"))
+    assertEquals(2, indexFiles.length)
+    for (file <- indexFiles) {
+      val offsetIndex = new OffsetIndex(file, file.getName.replace(".index","").toLong)
+      assertTrue(offsetIndex.lastOffset >= 0)
+      offsetIndex.close()
+    }
+    Utils.delete(logDir)
+  }
+
+  @Test
+  def testWriteLeaderEpochCheckpointAfterDirectoryRename(): 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.latestEpoch)
+
+    // Ensure that after a directory rename, the epoch cache is written to the right location
+    val tp = Log.parseTopicPartitionName(log.dir)
+    log.renameDir(Log.logDeleteDirName(tp))
+    log.appendAsLeader(TestUtils.records(List(new SimpleRecord("foo".getBytes()))), leaderEpoch = 10)
+    assertEquals(Some(10), log.latestEpoch)
+    assertTrue(LeaderEpochCheckpointFile.newFile(log.dir).exists())
+    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)
+    val log = createLog(logDir, logConfig)
+    log.appendAsLeader(TestUtils.records(List(new SimpleRecord("foo".getBytes()))), leaderEpoch = 5)
+    assertEquals(Some(5), log.latestEpoch)
+    log.close()
+
+    // reopen the log with an older message format version and check the cache
+    val downgradedLogConfig = LogTest.createLogConfig(segmentBytes = 1000, indexIntervalBytes = 1,
+      maxMessageBytes = 64 * 1024, messageFormatVersion = kafka.api.KAFKA_0_10_2_IV0.shortVersion)
+    val reopened = createLog(logDir, downgradedLogConfig, lastShutdownClean = false)
+    assertLeaderEpochCacheEmpty(reopened)
+
+    reopened.appendAsLeader(TestUtils.records(List(new SimpleRecord("bar".getBytes())),
+      magicValue = RecordVersion.V1.value), leaderEpoch = 5)
+    assertLeaderEpochCacheEmpty(reopened)
+  }
+
+  @Test
+  def testLeaderEpochCacheClearedAfterDynamicMessageFormatDowngrade(): 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.latestEpoch)
+
+    val downgradedLogConfig = LogTest.createLogConfig(segmentBytes = 1000, indexIntervalBytes = 1,
+      maxMessageBytes = 64 * 1024, messageFormatVersion = kafka.api.KAFKA_0_10_2_IV0.shortVersion)
+    log.updateConfig(downgradedLogConfig)
+    assertLeaderEpochCacheEmpty(log)
+
+    log.appendAsLeader(TestUtils.records(List(new SimpleRecord("bar".getBytes())),
+      magicValue = RecordVersion.V1.value), leaderEpoch = 5)
+    assertLeaderEpochCacheEmpty(log)
+  }
+
+  @Test
+  def testLeaderEpochCacheCreatedAfterMessageFormatUpgrade(): Unit = {
+    val logConfig = LogTest.createLogConfig(segmentBytes = 1000, indexIntervalBytes = 1,
+      maxMessageBytes = 64 * 1024, messageFormatVersion = kafka.api.KAFKA_0_10_2_IV0.shortVersion)
+    val log = createLog(logDir, logConfig)
+    log.appendAsLeader(TestUtils.records(List(new SimpleRecord("bar".getBytes())),
+      magicValue = RecordVersion.V1.value), leaderEpoch = 5)
+    assertLeaderEpochCacheEmpty(log)
+
+    val upgradedLogConfig = LogTest.createLogConfig(segmentBytes = 1000, indexIntervalBytes = 1,
+      maxMessageBytes = 64 * 1024, messageFormatVersion = kafka.api.KAFKA_0_11_0_IV0.shortVersion)
+    log.updateConfig(upgradedLogConfig)
+    log.appendAsLeader(TestUtils.records(List(new SimpleRecord("foo".getBytes()))), leaderEpoch = 5)
+    assertEquals(Some(5), log.latestEpoch)
+  }
+
+  private def assertLeaderEpochCacheEmpty(log: Log): Unit = {
+    assertEquals(None, log.leaderEpochCache)
+    assertEquals(None, log.latestEpoch)
+    assertFalse(LeaderEpochCheckpointFile.newFile(log.dir).exists())
+  }
+
+  @Test
+  def testOverCompactedLogRecoveryMultiRecord(): Unit = {
+    // append some messages to create some segments
+    val logConfig = LogTest.createLogConfig(segmentBytes = 1000, indexIntervalBytes = 1, maxMessageBytes = 64 * 1024)
+    val log = createLog(logDir, logConfig)
+    val set1 = MemoryRecords.withRecords(0, CompressionType.NONE, 0, new SimpleRecord("v1".getBytes(), "k1".getBytes()))
+    val set2 = MemoryRecords.withRecords(Integer.MAX_VALUE.toLong + 2, CompressionType.GZIP, 0,
+      new SimpleRecord("v3".getBytes(), "k3".getBytes()),
+      new SimpleRecord("v4".getBytes(), "k4".getBytes()))
+    val set3 = MemoryRecords.withRecords(Integer.MAX_VALUE.toLong + 4, CompressionType.GZIP, 0,
+      new SimpleRecord("v5".getBytes(), "k5".getBytes()),
+      new SimpleRecord("v6".getBytes(), "k6".getBytes()))
+    val set4 = MemoryRecords.withRecords(Integer.MAX_VALUE.toLong + 6, CompressionType.GZIP, 0,
+      new SimpleRecord("v7".getBytes(), "k7".getBytes()),
+      new SimpleRecord("v8".getBytes(), "k8".getBytes()))
+    //Writes into an empty log with baseOffset 0
+    log.appendAsFollower(set1)
+    assertEquals(0L, log.activeSegment.baseOffset)
+    //This write will roll the segment, yielding a new segment with base offset = max(1, Integer.MAX_VALUE+2) = Integer.MAX_VALUE+2
+    log.appendAsFollower(set2)
+    assertEquals(Integer.MAX_VALUE.toLong + 2, log.activeSegment.baseOffset)
+    assertTrue(Log.producerSnapshotFile(logDir, Integer.MAX_VALUE.toLong + 2).exists)
+    //This will go into the existing log
+    log.appendAsFollower(set3)
+    assertEquals(Integer.MAX_VALUE.toLong + 2, log.activeSegment.baseOffset)
+    //This will go into the existing log
+    log.appendAsFollower(set4)
+    assertEquals(Integer.MAX_VALUE.toLong + 2, log.activeSegment.baseOffset)
+    log.close()
+    val indexFiles = logDir.listFiles.filter(file => file.getName.contains(".index"))
+    assertEquals(2, indexFiles.length)
+    for (file <- indexFiles) {
+      val offsetIndex = new OffsetIndex(file, file.getName.replace(".index","").toLong)
+      assertTrue(offsetIndex.lastOffset >= 0)
+      offsetIndex.close()
+    }
+    Utils.delete(logDir)
+  }
+
+  @Test
+  def testOverCompactedLogRecoveryMultiRecordV1(): Unit = {
+    // append some messages to create some segments
+    val logConfig = LogTest.createLogConfig(segmentBytes = 1000, indexIntervalBytes = 1, maxMessageBytes = 64 * 1024)
+    val log = createLog(logDir, logConfig)
+    val set1 = MemoryRecords.withRecords(RecordBatch.MAGIC_VALUE_V1, 0, CompressionType.NONE,
+      new SimpleRecord("v1".getBytes(), "k1".getBytes()))
+    val set2 = MemoryRecords.withRecords(RecordBatch.MAGIC_VALUE_V1, Integer.MAX_VALUE.toLong + 2, CompressionType.GZIP,
+      new SimpleRecord("v3".getBytes(), "k3".getBytes()),
+      new SimpleRecord("v4".getBytes(), "k4".getBytes()))
+    val set3 = MemoryRecords.withRecords(RecordBatch.MAGIC_VALUE_V1, Integer.MAX_VALUE.toLong + 4, CompressionType.GZIP,
+      new SimpleRecord("v5".getBytes(), "k5".getBytes()),
+      new SimpleRecord("v6".getBytes(), "k6".getBytes()))
+    val set4 = MemoryRecords.withRecords(RecordBatch.MAGIC_VALUE_V1, Integer.MAX_VALUE.toLong + 6, CompressionType.GZIP,
+      new SimpleRecord("v7".getBytes(), "k7".getBytes()),
+      new SimpleRecord("v8".getBytes(), "k8".getBytes()))
+    //Writes into an empty log with baseOffset 0
+    log.appendAsFollower(set1)
+    assertEquals(0L, log.activeSegment.baseOffset)
+    //This write will roll the segment, yielding a new segment with base offset = max(1, 3) = 3
+    log.appendAsFollower(set2)
+    assertEquals(3, log.activeSegment.baseOffset)
+    assertTrue(Log.producerSnapshotFile(logDir, 3).exists)
+    //This will also roll the segment, yielding a new segment with base offset = max(5, Integer.MAX_VALUE+4) = Integer.MAX_VALUE+4
+    log.appendAsFollower(set3)
+    assertEquals(Integer.MAX_VALUE.toLong + 4, log.activeSegment.baseOffset)
+    assertTrue(Log.producerSnapshotFile(logDir, Integer.MAX_VALUE.toLong + 4).exists)
+    //This will go into the existing log
+    log.appendAsFollower(set4)
+    assertEquals(Integer.MAX_VALUE.toLong + 4, log.activeSegment.baseOffset)
+    log.close()
+    val indexFiles = logDir.listFiles.filter(file => file.getName.contains(".index"))
+    assertEquals(3, indexFiles.length)
+    for (file <- indexFiles) {
+      val offsetIndex = new OffsetIndex(file, file.getName.replace(".index","").toLong)
+      assertTrue(offsetIndex.lastOffset >= 0)
+      offsetIndex.close()
+    }
+    Utils.delete(logDir)
+  }
+
+  @Test
+  def testSplitOnOffsetOverflow(): Unit = {
+    // create a log such that one log segment has offsets that overflow, and call the split API on that segment
+    val logConfig = LogTest.createLogConfig(indexIntervalBytes = 1, fileDeleteDelayMs = 1000)
+    val (log, segmentWithOverflow) = createLogWithOffsetOverflow(logConfig)
+    assertTrue("At least one segment must have offset overflow", LogTest.hasOffsetOverflow(log))
+
+    val allRecordsBeforeSplit = LogTest.allRecords(log)
+
+    // split the segment with overflow
+    log.splitOverflowedSegment(segmentWithOverflow)
+
+    // assert we were successfully able to split the segment
+    assertEquals(4, log.numberOfSegments)
+    LogTest.verifyRecordsInLog(log, allRecordsBeforeSplit)
+
+    // verify we do not have offset overflow anymore
+    assertFalse(LogTest.hasOffsetOverflow(log))
+  }
+
+  @Test
+  def testDegenerateSegmentSplit(): Unit = {
+    // This tests a scenario where all of the batches appended to a segment have overflowed.
+    // When we split the overflowed segment, only one new segment will be created.
+
+    val overflowOffset = Int.MaxValue + 1L
+    val batch1 = MemoryRecords.withRecords(overflowOffset, CompressionType.NONE, 0,
+      new SimpleRecord("a".getBytes))
+    val batch2 = MemoryRecords.withRecords(overflowOffset + 1, CompressionType.NONE, 0,
+      new SimpleRecord("b".getBytes))
+
+    testDegenerateSplitSegmentWithOverflow(segmentBaseOffset = 0L, List(batch1, batch2))
+  }
+
+  @Test
+  def testDegenerateSegmentSplitWithOutOfRangeBatchLastOffset(): Unit = {
+    // Degenerate case where the only batch in the segment overflows. In this scenario,
+    // the first offset of the batch is valid, but the last overflows.
+
+    val firstBatchBaseOffset = Int.MaxValue - 1
+    val records = MemoryRecords.withRecords(firstBatchBaseOffset, CompressionType.NONE, 0,
+      new SimpleRecord("a".getBytes),
+      new SimpleRecord("b".getBytes),
+      new SimpleRecord("c".getBytes))
+
+    testDegenerateSplitSegmentWithOverflow(segmentBaseOffset = 0L, List(records))
+  }
+
+  private def testDegenerateSplitSegmentWithOverflow(segmentBaseOffset: Long, records: List[MemoryRecords]): Unit = {
+    val segment = LogTest.rawSegment(logDir, segmentBaseOffset)
+    // Need to create the offset files explicitly to avoid triggering segment recovery to truncate segment.
+    Log.offsetIndexFile(logDir, segmentBaseOffset).createNewFile()
+    Log.timeIndexFile(logDir, segmentBaseOffset).createNewFile()
+    records.foreach(segment.append _)
+    segment.close()
+
+    val logConfig = LogTest.createLogConfig(indexIntervalBytes = 1, fileDeleteDelayMs = 1000)
+    val log = createLog(logDir, logConfig, recoveryPoint = Long.MaxValue)
+
+    val segmentWithOverflow = LogTest.firstOverflowSegment(log).getOrElse {
+      Assertions.fail("Failed to create log with a segment which has overflowed offsets")
+    }
+
+    val allRecordsBeforeSplit = LogTest.allRecords(log)
+    log.splitOverflowedSegment(segmentWithOverflow)
+
+    assertEquals(1, log.numberOfSegments)
+
+    val firstBatchBaseOffset = records.head.batches.asScala.head.baseOffset
+    assertEquals(firstBatchBaseOffset, log.activeSegment.baseOffset)
+    LogTest.verifyRecordsInLog(log, allRecordsBeforeSplit)
+
+    assertFalse(LogTest.hasOffsetOverflow(log))
+  }
+
+  @Test
+  def testRecoveryOfSegmentWithOffsetOverflow(): Unit = {
+    val logConfig = LogTest.createLogConfig(indexIntervalBytes = 1, fileDeleteDelayMs = 1000)
+    val (log, _) = createLogWithOffsetOverflow(logConfig)
+    val expectedKeys = LogTest.keysInLog(log)
+
+    // Run recovery on the log. This should split the segment underneath. Ignore .deleted files as we could have still
+    // have them lying around after the split.
+    val recoveredLog = recoverAndCheck(logConfig, expectedKeys)
+    assertEquals(expectedKeys, LogTest.keysInLog(recoveredLog))
+
+    // Running split again would throw an error
+    for (segment <- recoveredLog.logSegments) {
+      try {
+        log.splitOverflowedSegment(segment)
+        fail()
+      } catch {
+        case _: IllegalArgumentException =>
+      }
+    }
+  }
+
+  @Test
+  def testRecoveryAfterCrashDuringSplitPhase1(): Unit = {
+    val logConfig = LogTest.createLogConfig(indexIntervalBytes = 1, fileDeleteDelayMs = 1000)
+    val (log, segmentWithOverflow) = createLogWithOffsetOverflow(logConfig)
+    val expectedKeys = LogTest.keysInLog(log)
+    val numSegmentsInitial = log.logSegments.size
+
+    // Split the segment
+    val newSegments = log.splitOverflowedSegment(segmentWithOverflow)
+
+    // Simulate recovery just after .cleaned file is created, before rename to .swap. On recovery, existing split
+    // operation is aborted but the recovery process itself kicks off split which should complete.
+    newSegments.reverse.foreach(segment => {
+      segment.changeFileSuffixes("", Log.CleanedFileSuffix)
+      segment.truncateTo(0)
+    })
+    for (file <- logDir.listFiles if file.getName.endsWith(Log.DeletedFileSuffix))
+      Utils.atomicMoveWithFallback(file.toPath, Paths.get(CoreUtils.replaceSuffix(file.getPath, Log.DeletedFileSuffix, "")))
+
+    val recoveredLog = recoverAndCheck(logConfig, expectedKeys)
+    assertEquals(expectedKeys, LogTest.keysInLog(recoveredLog))
+    assertEquals(numSegmentsInitial + 1, recoveredLog.logSegments.size)
+    recoveredLog.close()
+  }
+
+  @Test
+  def testRecoveryAfterCrashDuringSplitPhase2(): Unit = {
+    val logConfig = LogTest.createLogConfig(indexIntervalBytes = 1, fileDeleteDelayMs = 1000)
+    val (log, segmentWithOverflow) = createLogWithOffsetOverflow(logConfig)
+    val expectedKeys = LogTest.keysInLog(log)
+    val numSegmentsInitial = log.logSegments.size
+
+    // Split the segment
+    val newSegments = log.splitOverflowedSegment(segmentWithOverflow)
+
+    // Simulate recovery just after one of the new segments has been renamed to .swap. On recovery, existing split
+    // operation is aborted but the recovery process itself kicks off split which should complete.
+    newSegments.reverse.foreach { segment =>
+      if (segment != newSegments.last)
+        segment.changeFileSuffixes("", Log.CleanedFileSuffix)
+      else
+        segment.changeFileSuffixes("", Log.SwapFileSuffix)
+      segment.truncateTo(0)
+    }
+    for (file <- logDir.listFiles if file.getName.endsWith(Log.DeletedFileSuffix))
+      Utils.atomicMoveWithFallback(file.toPath, Paths.get(CoreUtils.replaceSuffix(file.getPath, Log.DeletedFileSuffix, "")))
+
+    val recoveredLog = recoverAndCheck(logConfig, expectedKeys)
+    assertEquals(expectedKeys, LogTest.keysInLog(recoveredLog))
+    assertEquals(numSegmentsInitial + 1, recoveredLog.logSegments.size)
+    recoveredLog.close()
+  }
+
+  @Test
+  def testRecoveryAfterCrashDuringSplitPhase3(): Unit = {
+    val logConfig = LogTest.createLogConfig(indexIntervalBytes = 1, fileDeleteDelayMs = 1000)
+    val (log, segmentWithOverflow) = createLogWithOffsetOverflow(logConfig)
+    val expectedKeys = LogTest.keysInLog(log)
+    val numSegmentsInitial = log.logSegments.size
+
+    // Split the segment
+    val newSegments = log.splitOverflowedSegment(segmentWithOverflow)
+
+    // Simulate recovery right after all new segments have been renamed to .swap. On recovery, existing split operation
+    // is completed and the old segment must be deleted.
+    newSegments.reverse.foreach(segment => {
+        segment.changeFileSuffixes("", Log.SwapFileSuffix)
+    })
+    for (file <- logDir.listFiles if file.getName.endsWith(Log.DeletedFileSuffix))
+      Utils.atomicMoveWithFallback(file.toPath, Paths.get(CoreUtils.replaceSuffix(file.getPath, Log.DeletedFileSuffix, "")))
+
+    // Truncate the old segment
+    segmentWithOverflow.truncateTo(0)
+
+    val recoveredLog = recoverAndCheck(logConfig, expectedKeys)
+    assertEquals(expectedKeys, LogTest.keysInLog(recoveredLog))
+    assertEquals(numSegmentsInitial + 1, recoveredLog.logSegments.size)
+    log.close()
+  }
+
+  @Test
+  def testRecoveryAfterCrashDuringSplitPhase4(): Unit = {
+    val logConfig = LogTest.createLogConfig(indexIntervalBytes = 1, fileDeleteDelayMs = 1000)
+    val (log, segmentWithOverflow) = createLogWithOffsetOverflow(logConfig)
+    val expectedKeys = LogTest.keysInLog(log)
+    val numSegmentsInitial = log.logSegments.size
+
+    // Split the segment
+    val newSegments = log.splitOverflowedSegment(segmentWithOverflow)
+
+    // Simulate recovery right after all new segments have been renamed to .swap and old segment has been deleted. On
+    // recovery, existing split operation is completed.
+    newSegments.reverse.foreach(_.changeFileSuffixes("", Log.SwapFileSuffix))
+
+    for (file <- logDir.listFiles if file.getName.endsWith(Log.DeletedFileSuffix))
+      Utils.delete(file)
+
+    // Truncate the old segment
+    segmentWithOverflow.truncateTo(0)
+
+    val recoveredLog = recoverAndCheck(logConfig, expectedKeys)
+    assertEquals(expectedKeys, LogTest.keysInLog(recoveredLog))
+    assertEquals(numSegmentsInitial + 1, recoveredLog.logSegments.size)
+    recoveredLog.close()
   }
 
   @Test
-  def testOverCompactedLogRecovery(): Unit = {
-    // append some messages to create some segments
-    val logConfig = createLogConfig(segmentBytes = 1000, indexIntervalBytes = 1, maxMessageBytes = 64 * 1024)
-    val log = createLog(logDir, logConfig)
-    val set1 = MemoryRecords.withRecords(0, CompressionType.NONE, 0, new SimpleRecord("v1".getBytes(), "k1".getBytes()))
-    val set2 = MemoryRecords.withRecords(Integer.MAX_VALUE.toLong + 2, CompressionType.NONE, 0, new SimpleRecord("v3".getBytes(), "k3".getBytes()))
-    val set3 = MemoryRecords.withRecords(Integer.MAX_VALUE.toLong + 3, CompressionType.NONE, 0, new SimpleRecord("v4".getBytes(), "k4".getBytes()))
-    val set4 = MemoryRecords.withRecords(Integer.MAX_VALUE.toLong + 4, CompressionType.NONE, 0, new SimpleRecord("v5".getBytes(), "k5".getBytes()))
-    //Writes into an empty log with baseOffset 0
-    log.appendAsFollower(set1)
-    assertEquals(0L, log.activeSegment.baseOffset)
-    //This write will roll the segment, yielding a new segment with base offset = max(2, 1) = 2
-    log.appendAsFollower(set2)
-    assertEquals(2L, log.activeSegment.baseOffset)
-    assertTrue(Log.producerSnapshotFile(logDir, 2L).exists)
-    //This will also roll the segment, yielding a new segment with base offset = max(3, Integer.MAX_VALUE+3) = Integer.MAX_VALUE+3
-    log.appendAsFollower(set3)
-    assertEquals(Integer.MAX_VALUE.toLong + 3, log.activeSegment.baseOffset)
-    assertTrue(Log.producerSnapshotFile(logDir, Integer.MAX_VALUE.toLong + 3).exists)
-    //This will go into the existing log
-    log.appendAsFollower(set4)
-    assertEquals(Integer.MAX_VALUE.toLong + 3, log.activeSegment.baseOffset)
-    log.close()
-    val indexFiles = logDir.listFiles.filter(file => file.getName.contains(".index"))
-    assertEquals(3, indexFiles.length)
-    for (file <- indexFiles) {
-      val offsetIndex = new OffsetIndex(file, file.getName.replace(".index","").toLong)
-      assertTrue(offsetIndex.lastOffset >= 0)
-      offsetIndex.close()
-    }
-    Utils.delete(logDir)
+  def testRecoveryAfterCrashDuringSplitPhase5(): Unit = {
+    val logConfig = LogTest.createLogConfig(indexIntervalBytes = 1, fileDeleteDelayMs = 1000)
+    val (log, segmentWithOverflow) = createLogWithOffsetOverflow(logConfig)
+    val expectedKeys = LogTest.keysInLog(log)
+    val numSegmentsInitial = log.logSegments.size
+
+    // Split the segment
+    val newSegments = log.splitOverflowedSegment(segmentWithOverflow)
+
+    // Simulate recovery right after one of the new segment has been renamed to .swap and the other to .log. On
+    // recovery, existing split operation is completed.
+    newSegments.last.changeFileSuffixes("", Log.SwapFileSuffix)
+
+    // Truncate the old segment
+    segmentWithOverflow.truncateTo(0)
+
+    val recoveredLog = recoverAndCheck(logConfig, expectedKeys)
+    assertEquals(expectedKeys, LogTest.keysInLog(recoveredLog))
+    assertEquals(numSegmentsInitial + 1, recoveredLog.logSegments.size)
+    recoveredLog.close()
   }
 
   @Test
-  def testCleanShutdownFile() {
+  def testCleanShutdownFile(): Unit = {
     // append some messages to create some segments
-    val logConfig = createLogConfig(segmentBytes = 1000, indexIntervalBytes = 1, maxMessageBytes = 64 * 1024)
+    val logConfig = LogTest.createLogConfig(segmentBytes = 1000, indexIntervalBytes = 1, maxMessageBytes = 64 * 1024)
     def createRecords = TestUtils.singletonRecords(value = "test".getBytes, timestamp = mockTime.milliseconds)
 
-    val cleanShutdownFile = createCleanShutdownFile()
-    assertTrue(".kafka_cleanshutdown must exist", cleanShutdownFile.exists())
     var recoveryPoint = 0L
     // create a log and write some messages to it
     var log = createLog(logDir, logConfig)
@@ -1935,15 +3297,16 @@ class LogTest {
     log.close()
 
     // check if recovery was attempted. Even if the recovery point is 0L, recovery should not be attempted as the
-    // clean shutdown file exists.
+    // clean shutdown file exists. Note: Earlier, Log layer relied on the presence of clean shutdown file to determine the status
+    // of last shutdown. Now, LogManager checks for the presence of this file and immediately deletes the same. It passes
+    // down a clean shutdown flag to the Log layer as log is loaded. Recovery is attempted based on this flag.
     recoveryPoint = log.logEndOffset
     log = createLog(logDir, logConfig)
     assertEquals(recoveryPoint, log.logEndOffset)
-    Utils.delete(cleanShutdownFile)
   }
 
   @Test
-  def testParseTopicPartitionName() {
+  def testParseTopicPartitionName(): Unit = {
     val topic = "test_topic"
     val partition = "143"
     val dir = new File(logDir, topicPartitionName(topic, partition))
@@ -1957,7 +3320,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)))
@@ -1967,7 +3330,7 @@ class LogTest {
   }
 
   @Test
-  def testParseTopicPartitionNameForEmptyName() {
+  def testParseTopicPartitionNameForEmptyName(): Unit = {
     try {
       val dir = new File("")
       Log.parseTopicPartitionName(dir)
@@ -1978,7 +3341,7 @@ class LogTest {
   }
 
   @Test
-  def testParseTopicPartitionNameForNull() {
+  def testParseTopicPartitionNameForNull(): Unit = {
     try {
       val dir: File = null
       Log.parseTopicPartitionName(dir)
@@ -1989,7 +3352,7 @@ class LogTest {
   }
 
   @Test
-  def testParseTopicPartitionNameForMissingSeparator() {
+  def testParseTopicPartitionNameForMissingSeparator(): Unit = {
     val topic = "test_topic"
     val partition = "1999"
     val dir = new File(logDir, topic + partition)
@@ -2010,7 +3373,7 @@ class LogTest {
   }
 
   @Test
-  def testParseTopicPartitionNameForMissingTopic() {
+  def testParseTopicPartitionNameForMissingTopic(): Unit = {
     val topic = ""
     val partition = "1999"
     val dir = new File(logDir, topicPartitionName(topic, partition))
@@ -2032,10 +3395,10 @@ class LogTest {
   }
 
   @Test
-  def testParseTopicPartitionNameForMissingPartition() {
+  def testParseTopicPartitionNameForMissingPartition(): Unit = {
     val topic = "test_topic"
     val partition = ""
-    val dir = new File(logDir + topicPartitionName(topic, partition))
+    val dir = new File(logDir.getPath + topicPartitionName(topic, partition))
     try {
       Log.parseTopicPartitionName(dir)
       fail("KafkaException should have been thrown for dir: " + dir.getCanonicalPath)
@@ -2053,7 +3416,7 @@ class LogTest {
   }
 
   @Test
-  def testParseTopicPartitionNameForInvalidPartition() {
+  def testParseTopicPartitionNameForInvalidPartition(): Unit = {
     val topic = "test_topic"
     val partition = "1999a"
     val dir = new File(logDir, topicPartitionName(topic, partition))
@@ -2074,15 +3437,15 @@ class LogTest {
   }
 
   @Test
-  def testParseTopicPartitionNameForExistingInvalidDir() {
-    val dir1 = new File(logDir + "/non_kafka_dir")
+  def testParseTopicPartitionNameForExistingInvalidDir(): Unit = {
+    val dir1 = new File(logDir.getPath + "/non_kafka_dir")
     try {
       Log.parseTopicPartitionName(dir1)
       fail("KafkaException should have been thrown for dir: " + dir1.getCanonicalPath)
     } catch {
       case _: KafkaException => // should only throw KafkaException
     }
-    val dir2 = new File(logDir + "/non_kafka_dir-delete")
+    val dir2 = new File(logDir.getPath + "/non_kafka_dir-delete")
     try {
       Log.parseTopicPartitionName(dir2)
       fail("KafkaException should have been thrown for dir: " + dir2.getCanonicalPath)
@@ -2095,17 +3458,17 @@ class LogTest {
     topic + "-" + partition
 
   @Test
-  def testDeleteOldSegments() {
+  def testDeleteOldSegments(): Unit = {
     def createRecords = TestUtils.singletonRecords(value = "test".getBytes, timestamp = mockTime.milliseconds - 1000)
-    val logConfig = createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, segmentIndexBytes = 1000, retentionMs = 999)
+    val logConfig = LogTest.createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, segmentIndexBytes = 1000, retentionMs = 999)
     val log = createLog(logDir, logConfig)
 
     // append some messages to create some segments
     for (_ <- 0 until 100)
       log.appendAsLeader(createRecords, leaderEpoch = 0)
 
-    log.leaderEpochCache.assign(0, 40)
-    log.leaderEpochCache.assign(1, 90)
+    log.maybeAssignEpochStartOffset(0, 40)
+    log.maybeAssignEpochStartOffset(1, 90)
 
     // segments are not eligible for deletion if no high watermark has been set
     val numSegments = log.numberOfSegments
@@ -2115,11 +3478,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)
@@ -2128,11 +3491,11 @@ 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)
-    assertEquals("Epoch entry should be the latest epoch and the leo.", EpochEntry(1, 100), epochCache(log).epochEntries().head)
+    assertEquals("Epoch entries should have gone.", 1, epochCache(log).epochEntries.size)
+    assertEquals("Epoch entry should be the latest epoch and the leo.", EpochEntry(1, 100), epochCache(log).epochEntries.head)
 
     // append some messages to create some segments
     for (_ <- 0 until 100)
@@ -2141,123 +3504,123 @@ class LogTest {
     log.delete()
     assertEquals("The number of segments should be 0", 0, log.numberOfSegments)
     assertEquals("The number of deleted segments should be zero.", 0, log.deleteOldSegments())
-    assertEquals("Epoch entries should have gone.", 0, epochCache(log).epochEntries().size)
+    assertEquals("Epoch entries should have gone.", 0, epochCache(log).epochEntries.size)
   }
 
   @Test
-  def testLogDeletionAfterClose() {
+  def testLogDeletionAfterClose(): Unit = {
     def createRecords = TestUtils.singletonRecords(value = "test".getBytes, timestamp = mockTime.milliseconds - 1000)
-    val logConfig = createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, segmentIndexBytes = 1000, retentionMs = 999)
+    val logConfig = LogTest.createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, segmentIndexBytes = 1000, retentionMs = 999)
     val log = createLog(logDir, logConfig)
 
     // append some messages to create some segments
     log.appendAsLeader(createRecords, leaderEpoch = 0)
 
     assertEquals("The deleted segments should be gone.", 1, log.numberOfSegments)
-    assertEquals("Epoch entries should have gone.", 1, epochCache(log).epochEntries().size)
+    assertEquals("Epoch entries should have gone.", 1, epochCache(log).epochEntries.size)
 
     log.close()
     log.delete()
     assertEquals("The number of segments should be 0", 0, log.numberOfSegments)
-    assertEquals("Epoch entries should have gone.", 0, epochCache(log).epochEntries().size)
+    assertEquals("Epoch entries should have gone.", 0, epochCache(log).epochEntries.size)
   }
 
   @Test
-  def testLogDeletionAfterDeleteRecords() {
+  def testLogDeletionAfterDeleteRecords(): Unit = {
     def createRecords = TestUtils.singletonRecords("test".getBytes)
-    val logConfig = createLogConfig(segmentBytes = createRecords.sizeInBytes * 5)
+    val logConfig = LogTest.createLogConfig(segmentBytes = createRecords.sizeInBytes * 5)
     val log = createLog(logDir, logConfig)
 
     for (_ <- 0 until 15)
       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.maybeIncrementLogStartOffset(1, ClientRecordDeletion)
     log.deleteOldSegments()
     assertEquals("should have 3 segments", 3, log.numberOfSegments)
     assertEquals(log.logStartOffset, 1)
 
-    log.maybeIncrementLogStartOffset(6)
+    log.maybeIncrementLogStartOffset(6, ClientRecordDeletion)
     log.deleteOldSegments()
     assertEquals("should have 2 segments", 2, log.numberOfSegments)
     assertEquals(log.logStartOffset, 6)
 
-    log.maybeIncrementLogStartOffset(15)
+    log.maybeIncrementLogStartOffset(15, ClientRecordDeletion)
     log.deleteOldSegments()
     assertEquals("should have 1 segments", 1, log.numberOfSegments)
     assertEquals(log.logStartOffset, 15)
   }
 
   def epochCache(log: Log): LeaderEpochFileCache = {
-    log.leaderEpochCache.asInstanceOf[LeaderEpochFileCache]
+    log.leaderEpochCache.get
   }
 
   @Test
-  def shouldDeleteSizeBasedSegments() {
+  def shouldDeleteSizeBasedSegments(): Unit = {
     def createRecords = TestUtils.singletonRecords("test".getBytes)
-    val logConfig = createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, retentionBytes = createRecords.sizeInBytes * 10)
+    val logConfig = LogTest.createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, retentionBytes = createRecords.sizeInBytes * 10)
     val log = createLog(logDir, logConfig)
 
     // append some messages to create some segments
     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 = createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, retentionBytes = createRecords.sizeInBytes * 15)
+    val logConfig = LogTest.createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, retentionBytes = createRecords.sizeInBytes * 15)
     val log = createLog(logDir, logConfig)
 
     // append some messages to create some segments
     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 = createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, retentionMs = 10000)
+    val logConfig = LogTest.createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, retentionMs = 10000)
     val log = createLog(logDir, logConfig)
 
     // append some messages to create some segments
     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 = createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, retentionMs = 10000000)
+    val logConfig = LogTest.createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, retentionMs = 10000000)
     val log = createLog(logDir, logConfig)
 
     // append some messages to create some segments
     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 = createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, retentionMs = 10000, cleanupPolicy = "compact")
+    val logConfig = LogTest.createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, retentionMs = 10000, cleanupPolicy = "compact")
     val log = createLog(logDir, logConfig)
 
     // append some messages to create some segments
@@ -2268,34 +3631,60 @@ 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 = createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, retentionMs = 10000, cleanupPolicy = "compact,delete")
+    val logConfig = LogTest.createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, retentionMs = 10000, cleanupPolicy = "compact,delete")
     val log = createLog(logDir, logConfig)
 
     // append some messages to create some segments
     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 shouldApplyEpochToMessageOnAppendIfLeader() {
+  def shouldDeleteStartOffsetBreachedSegmentsWhenPolicyDoesNotIncludeDelete(): Unit = {
+    def createRecords = TestUtils.singletonRecords("test".getBytes, key = "test".getBytes, timestamp = 10L)
+    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
+    assertEquals(3, log.logSegments.count(_ => true))
+    log.updateHighWatermark(log.logEndOffset)
+    log.maybeIncrementLogStartOffset(recordsPerSegment, ClientRecordDeletion)
+
+    // 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.updateHighWatermark(log.logEndOffset)
+    log.deleteOldSegments()
+    assertEquals("There should be 2 segments remaining", 2, log.numberOfSegments)
+    assertTrue(log.logSegments.head.baseOffset <= log.logStartOffset)
+    assertTrue(log.logSegments.tail.forall(s => s.baseOffset > log.logStartOffset))
+  }
+
+  @Test
+  def shouldApplyEpochToMessageOnAppendIfLeader(): Unit = {
     val records = (0 until 50).toArray.map(id => new SimpleRecord(id.toString.getBytes))
 
     //Given this partition is on leader epoch 72
     val epoch = 72
     val log = createLog(logDir, LogConfig())
-    log.leaderEpochCache.assign(epoch, records.size)
+    log.maybeAssignEpochStartOffset(epoch, records.size)
 
     //When appending messages as a leader (i.e. assignOffsets = true)
     for (record <- records)
@@ -2306,20 +3695,20 @@ class LogTest {
 
     //Then leader epoch should be set on messages
     for (i <- records.indices) {
-      val read = log.readUncommitted(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))
 
     //Given each message has an offset & epoch, as msgs from leader would
     def recordsForEpoch(i: Int): MemoryRecords = {
       val recs = MemoryRecords.withRecords(messageIds(i), CompressionType.NONE, records(i))
-      recs.batches.asScala.foreach{record =>
+      recs.batches.forEach{record =>
         record.setPartitionLeaderEpoch(42)
         record.setLastOffset(i)
       }
@@ -2332,18 +3721,18 @@ class LogTest {
     for (i <- records.indices)
       log.appendAsFollower(recordsForEpoch(i))
 
-    assertEquals(42, log.leaderEpochCache.asInstanceOf[LeaderEpochFileCache].latestEpoch())
+    assertEquals(Some(42), log.latestEpoch)
   }
 
   @Test
-  def shouldTruncateLeaderEpochsWhenDeletingSegments() {
+  def shouldTruncateLeaderEpochsWhenDeletingSegments(): Unit = {
     def createRecords = TestUtils.singletonRecords("test".getBytes)
-    val logConfig = createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, retentionBytes = createRecords.sizeInBytes * 10)
+    val logConfig = LogTest.createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, retentionBytes = createRecords.sizeInBytes * 10)
     val log = createLog(logDir, logConfig)
     val cache = epochCache(log)
 
     // Given three segments of 5 messages each
-    for (e <- 0 until 15) {
+    for (_ <- 0 until 15) {
       log.appendAsLeader(createRecords, leaderEpoch = 0)
     }
 
@@ -2353,7 +3742,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
@@ -2361,14 +3750,14 @@ class LogTest {
   }
 
   @Test
-  def shouldUpdateOffsetForLeaderEpochsWhenDeletingSegments() {
+  def shouldUpdateOffsetForLeaderEpochsWhenDeletingSegments(): Unit = {
     def createRecords = TestUtils.singletonRecords("test".getBytes)
-    val logConfig = createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, retentionBytes = createRecords.sizeInBytes * 10)
+    val logConfig = LogTest.createLogConfig(segmentBytes = createRecords.sizeInBytes * 5, retentionBytes = createRecords.sizeInBytes * 10)
     val log = createLog(logDir, logConfig)
     val cache = epochCache(log)
 
     // Given three segments of 5 messages each
-    for (e <- 0 until 15) {
+    for (_ <- 0 until 15) {
       log.appendAsLeader(createRecords, leaderEpoch = 0)
     }
 
@@ -2378,28 +3767,33 @@ 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 the first entry should have gone from (0,0) => (0,5)
+    //The first entry should have gone from (0,0) => (0,5)
     assertEquals(ListBuffer(EpochEntry(0, 5), EpochEntry(1, 7), EpochEntry(2, 10)), cache.epochEntries)
   }
 
   @Test
-  def shouldTruncateLeaderEpochFileWhenTruncatingLog() {
-    def createRecords = TestUtils.singletonRecords(value = "test".getBytes, timestamp = mockTime.milliseconds)
-    val logConfig = createLogConfig(segmentBytes = 10 * createRecords.sizeInBytes)
+  def shouldTruncateLeaderEpochCheckpointFileWhenTruncatingLog(): Unit = {
+    def createRecords(startOffset: Long, epoch: Int): MemoryRecords = {
+      TestUtils.records(Seq(new SimpleRecord("value".getBytes)),
+        baseOffset = startOffset, partitionLeaderEpoch = epoch)
+    }
+
+    val logConfig = LogTest.createLogConfig(segmentBytes = 10 * createRecords(0, 0).sizeInBytes)
     val log = createLog(logDir, logConfig)
     val cache = epochCache(log)
 
-    //Given 2 segments, 10 messages per segment
-    for (epoch <- 1 to 20)
-      log.appendAsLeader(createRecords, leaderEpoch = 0)
+    def append(epoch: Int, startOffset: Long, count: Int): Unit = {
+      for (i <- 0 until count)
+        log.appendAsFollower(createRecords(startOffset + i, epoch))
+    }
 
-    //Simulate some leader changes at specific offsets
-    cache.assign(0, 0)
-    cache.assign(1, 10)
-    cache.assign(2, 16)
+    //Given 2 segments, 10 messages per segment
+    append(epoch = 0, startOffset = 0, count = 10)
+    append(epoch = 1, startOffset = 10, count = 6)
+    append(epoch = 2, startOffset = 16, count = 4)
 
     assertEquals(2, log.numberOfSegments)
     assertEquals(20, log.logEndOffset)
@@ -2430,10 +3824,10 @@ 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.
-    */
+   * 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)
@@ -2451,12 +3845,12 @@ class LogTest {
     assertEquals(ListBuffer(EpochEntry(1, 0), EpochEntry(2, 1), EpochEntry(3, 3)), leaderEpochCache.epochEntries)
 
     // deliberately remove some of the epoch entries
-    leaderEpochCache.clearAndFlushLatest(2)
+    leaderEpochCache.truncateFromEnd(2)
     assertNotEquals(ListBuffer(EpochEntry(1, 0), EpochEntry(2, 1), EpochEntry(3, 3)), leaderEpochCache.epochEntries)
     log.close()
 
     // reopen the log and recover from the beginning
-    val recoveredLog = createLog(logDir, LogConfig())
+    val recoveredLog = createLog(logDir, LogConfig(), lastShutdownClean = false)
     val recoveredLeaderEpochCache = epochCache(recoveredLog)
 
     // epoch entries should be recovered
@@ -2465,26 +3859,27 @@ class LogTest {
   }
 
   /**
-    * Wrap a single record log buffer with leader epoch.
-    */
+   * Wrap a single record log buffer with leader epoch.
+   */
   private def singletonRecordsWithLeaderEpoch(value: Array[Byte],
-                                      key: Array[Byte] = null,
-                                      leaderEpoch: Int,
-                                      offset: Long,
-                                      codec: CompressionType = CompressionType.NONE,
-                                      timestamp: Long = RecordBatch.NO_TIMESTAMP,
-                                      magicValue: Byte = RecordBatch.CURRENT_MAGIC_VALUE): MemoryRecords = {
+                                              key: Array[Byte] = null,
+                                              leaderEpoch: Int,
+                                              offset: Long,
+                                              codec: CompressionType = CompressionType.NONE,
+                                              timestamp: Long = RecordBatch.NO_TIMESTAMP,
+                                              magicValue: Byte = RecordBatch.CURRENT_MAGIC_VALUE): MemoryRecords = {
     val records = Seq(new SimpleRecord(timestamp, key, value))
 
     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() {
-    val logConfig = createLogConfig(segmentBytes = 1024 * 1024 * 5)
+  @Test
+  def testFirstUnstableOffsetNoTransactionalData(): Unit = {
+    val logConfig = LogTest.createLogConfig(segmentBytes = 1024 * 1024 * 5)
     val log = createLog(logDir, logConfig)
 
     val records = MemoryRecords.withRecords(CompressionType.NONE,
@@ -2497,8 +3892,8 @@ class LogTest {
   }
 
   @Test
-  def testFirstUnstableOffsetWithTransactionalData() {
-    val logConfig = createLogConfig(segmentBytes = 1024 * 1024 * 5)
+  def testFirstUnstableOffsetWithTransactionalData(): Unit = {
+    val logConfig = LogTest.createLogConfig(segmentBytes = 1024 * 1024 * 5)
     val log = createLog(logDir, logConfig)
 
     val pid = 137L
@@ -2512,7 +3907,7 @@ class LogTest {
       new SimpleRecord("baz".getBytes))
 
     val firstAppendInfo = log.appendAsLeader(records, leaderEpoch = 0)
-    assertEquals(Some(firstAppendInfo.firstOffset), log.firstUnstableOffset.map(_.messageOffset))
+    assertEquals(firstAppendInfo.firstOffset, log.firstUnstableOffset)
 
     // add more transactional records
     seq += 3
@@ -2520,23 +3915,75 @@ class LogTest {
       new SimpleRecord("blah".getBytes)), leaderEpoch = 0)
 
     // LSO should not have changed
-    assertEquals(Some(firstAppendInfo.firstOffset), 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), 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] = () => {
+      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 = () => {
+      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 = createLogConfig(segmentBytes = 1024 * 1024 * 5)
+    val logConfig = LogTest.createLogConfig(segmentBytes = 1024 * 1024 * 5)
     val log = createLog(logDir, logConfig)
     val epoch = 0.toShort
 
@@ -2572,12 +4019,26 @@ 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
   def testFullTransactionIndexRecovery(): Unit = {
-    val logConfig = createLogConfig(segmentBytes = 128 * 5)
+    val logConfig = LogTest.createLogConfig(segmentBytes = 128 * 5)
     val log = createLog(logDir, logConfig)
     val epoch = 0.toShort
 
@@ -2614,21 +4075,21 @@ class LogTest {
 
     // delete all the offset and transaction index files to force recovery
     log.logSegments.foreach { segment =>
-      segment.index.delete()
-      segment.txnIndex.delete()
+      segment.offsetIndex.deleteIfExists()
+      segment.txnIndex.deleteIfExists()
     }
 
     log.close()
 
-    val reloadedLogConfig = createLogConfig(segmentBytes = 1024 * 5)
-    val reloadedLog = createLog(logDir, reloadedLogConfig)
+    val reloadedLogConfig = LogTest.createLogConfig(segmentBytes = 1024 * 5)
+    val reloadedLog = createLog(logDir, reloadedLogConfig, lastShutdownClean = false)
     val abortedTransactions = allAbortedTransactions(reloadedLog)
     assertEquals(List(new AbortedTxn(pid1, 0L, 29L, 8L), new AbortedTxn(pid2, 8L, 74L, 36L)), abortedTransactions)
   }
 
   @Test
   def testRecoverOnlyLastSegment(): Unit = {
-    val logConfig = createLogConfig(segmentBytes = 128 * 5)
+    val logConfig = LogTest.createLogConfig(segmentBytes = 128 * 5)
     val log = createLog(logDir, logConfig)
     val epoch = 0.toShort
 
@@ -2666,20 +4127,20 @@ class LogTest {
     // delete the last offset and transaction index files to force recovery
     val lastSegment = log.logSegments.last
     val recoveryPoint = lastSegment.baseOffset
-    lastSegment.index.delete()
-    lastSegment.txnIndex.delete()
+    lastSegment.offsetIndex.deleteIfExists()
+    lastSegment.txnIndex.deleteIfExists()
 
     log.close()
 
-    val reloadedLogConfig = createLogConfig(segmentBytes = 1024 * 5)
-    val reloadedLog = createLog(logDir, reloadedLogConfig, recoveryPoint = recoveryPoint)
+    val reloadedLogConfig = LogTest.createLogConfig(segmentBytes = 1024 * 5)
+    val reloadedLog = createLog(logDir, reloadedLogConfig, recoveryPoint = recoveryPoint, lastShutdownClean = false)
     val abortedTransactions = allAbortedTransactions(reloadedLog)
     assertEquals(List(new AbortedTxn(pid1, 0L, 29L, 8L), new AbortedTxn(pid2, 8L, 74L, 36L)), abortedTransactions)
   }
 
   @Test
   def testRecoverLastSegmentWithNoSnapshots(): Unit = {
-    val logConfig = createLogConfig(segmentBytes = 128 * 5)
+    val logConfig = LogTest.createLogConfig(segmentBytes = 128 * 5)
     val log = createLog(logDir, logConfig)
     val epoch = 0.toShort
 
@@ -2720,13 +4181,13 @@ class LogTest {
     // the producer state from the start of the log
     val lastSegment = log.logSegments.last
     val recoveryPoint = lastSegment.baseOffset
-    lastSegment.index.delete()
-    lastSegment.txnIndex.delete()
+    lastSegment.offsetIndex.deleteIfExists()
+    lastSegment.txnIndex.deleteIfExists()
 
     log.close()
 
-    val reloadedLogConfig = createLogConfig(segmentBytes = 1024 * 5)
-    val reloadedLog = createLog(logDir, reloadedLogConfig, recoveryPoint = recoveryPoint)
+    val reloadedLogConfig = LogTest.createLogConfig(segmentBytes = 1024 * 5)
+    val reloadedLog = createLog(logDir, reloadedLogConfig, recoveryPoint = recoveryPoint, lastShutdownClean = false)
     val abortedTransactions = allAbortedTransactions(reloadedLog)
     assertEquals(List(new AbortedTxn(pid1, 0L, 29L, 8L), new AbortedTxn(pid2, 8L, 74L, 36L)), abortedTransactions)
   }
@@ -2734,7 +4195,7 @@ class LogTest {
   @Test
   def testTransactionIndexUpdatedThroughReplication(): Unit = {
     val epoch = 0.toShort
-    val logConfig = createLogConfig(segmentBytes = 1024 * 1024 * 5)
+    val logConfig = LogTest.createLogConfig(segmentBytes = 1024 * 1024 * 5)
     val log = createLog(logDir, logConfig)
     val buffer = ByteBuffer.allocate(2048)
 
@@ -2773,14 +4234,59 @@ 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])
   def testZombieCoordinatorFenced(): Unit = {
     val pid = 1L
     val epoch = 0.toShort
-    val logConfig = createLogConfig(segmentBytes = 1024 * 1024 * 5)
+    val logConfig = LogTest.createLogConfig(segmentBytes = 1024 * 1024 * 5)
     val log = createLog(logDir, logConfig)
 
     val append = appendTransactionalAsLeader(log, pid, epoch)
@@ -2795,8 +4301,44 @@ class LogTest {
   }
 
   @Test
-  def testFirstUnstableOffsetDoesNotExceedLogStartOffsetMidSegment(): Unit = {
-    val logConfig = createLogConfig(segmentBytes = 1024 * 1024 * 5)
+  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[InvalidProducerEpochException] {
+      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
     val pid = 1L
@@ -2810,17 +4352,18 @@ class LogTest {
     assertEquals(2, log.logSegments.size)
     appendPid(5)
 
-    assertEquals(Some(0L), log.firstUnstableOffset.map(_.messageOffset))
+    assertEquals(Some(0L), log.firstUnstableOffset)
 
-    log.maybeIncrementLogStartOffset(5L)
+    log.updateHighWatermark(log.logEndOffset)
+    log.maybeIncrementLogStartOffset(5L, ClientRecordDeletion)
 
     // 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 = {
-    val logConfig = createLogConfig(segmentBytes = 1024 * 1024 * 5)
+  def testLastStableOffsetDoesNotExceedLogStartOffsetAfterSegmentDeletion(): Unit = {
+    val logConfig = LogTest.createLogConfig(segmentBytes = 1024 * 1024 * 5)
     val log = createLog(logDir, logConfig)
     val epoch = 0.toShort
     val pid = 1L
@@ -2834,20 +4377,86 @@ class LogTest {
     assertEquals(2, log.logSegments.size)
     appendPid(5)
 
-    assertEquals(Some(0L), log.firstUnstableOffset.map(_.messageOffset))
+    assertEquals(Some(0L), log.firstUnstableOffset)
 
-    log.maybeIncrementLogStartOffset(8L)
-    log.onHighWatermarkIncremented(log.logEndOffset)
+    log.updateHighWatermark(log.logEndOffset)
+    log.maybeIncrementLogStartOffset(8L, ClientRecordDeletion)
+    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.
+    // Note that the second append does not write to the transaction index because the producer
+    // state has already been updated and we do not write index entries for empty transactions.
+    // In the future, we may strengthen the fencing logic so that additional writes to the
+    // log are not possible after an IO error (see KAFKA-10778).
+    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, lastShutdownClean = false)
+    assertEquals(12L, reopenedLog.logEndOffset)
+    assertEquals(1, reopenedLog.activeSegment.txnIndex.allAbortedTxns.size)
+    reopenedLog.updateHighWatermark(12L)
+    assertEquals(None, reopenedLog.firstUnstableOffset)
+  }
+
+  @Test
+  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() {
-    val logConfig = createLogConfig(segmentBytes = 1024 * 1024 * 5)
+  def testLastStableOffsetWithMixedProducerData(): Unit = {
+    val logConfig = LogTest.createLogConfig(segmentBytes = 1024 * 1024 * 5)
     val log = createLog(logDir, logConfig)
 
     // for convenience, both producers share the same epoch
@@ -2863,7 +4472,7 @@ class LogTest {
       new SimpleRecord("a".getBytes),
       new SimpleRecord("b".getBytes),
       new SimpleRecord("c".getBytes)), leaderEpoch = 0)
-    assertEquals(Some(firstAppendInfo.firstOffset), log.firstUnstableOffset.map(_.messageOffset))
+    assertEquals(firstAppendInfo.firstOffset, log.firstUnstableOffset)
 
     // mix in some non-transactional data
     log.appendAsLeader(MemoryRecords.withRecords(CompressionType.NONE,
@@ -2878,27 +4487,25 @@ class LogTest {
       new SimpleRecord("f".getBytes)), leaderEpoch = 0)
 
     // LSO should not have changed
-    assertEquals(Some(firstAppendInfo.firstOffset), 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), 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
@@ -2908,12 +4515,11 @@ class LogTest {
       new SimpleRecord("b".getBytes),
       new SimpleRecord("c".getBytes))
 
-    val logConfig = createLogConfig(segmentBytes = records.sizeInBytes)
+    val logConfig = LogTest.createLogConfig(segmentBytes = records.sizeInBytes)
     val log = createLog(logDir, logConfig)
 
     val firstAppendInfo = log.appendAsLeader(records, leaderEpoch = 0)
-    assertEquals(Some(firstAppendInfo.firstOffset), 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
@@ -2921,70 +4527,33 @@ class LogTest {
       new SimpleRecord("d".getBytes),
       new SimpleRecord("e".getBytes),
       new SimpleRecord("f".getBytes)), leaderEpoch = 0)
-    assertEquals(Some(firstAppendInfo.firstOffset), 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, isolationLevel = IsolationLevel.READ_COMMITTED)
+    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)
   }
 
-  def createLogConfig(segmentMs: Long = Defaults.SegmentMs,
-                      segmentBytes: Int = Defaults.SegmentSize,
-                      retentionMs: Long = Defaults.RetentionMs,
-                      retentionBytes: Long = Defaults.RetentionSize,
-                      segmentJitterMs: Long = Defaults.SegmentJitterMs,
-                      cleanupPolicy: String = Defaults.CleanupPolicy,
-                      maxMessageBytes: Int = Defaults.MaxMessageSize,
-                      indexIntervalBytes: Int = Defaults.IndexInterval,
-                      segmentIndexBytes: Int = Defaults.MaxIndexSize,
-                      messageFormatVersion: String = Defaults.MessageFormatVersion,
-                      fileDeleteDelayMs: Long = Defaults.FileDeleteDelayMs): LogConfig = {
-    val logProps = new Properties()
-
-    logProps.put(LogConfig.SegmentMsProp, segmentMs: java.lang.Long)
-    logProps.put(LogConfig.SegmentBytesProp, segmentBytes: Integer)
-    logProps.put(LogConfig.RetentionMsProp, retentionMs: java.lang.Long)
-    logProps.put(LogConfig.RetentionBytesProp, retentionBytes: java.lang.Long)
-    logProps.put(LogConfig.SegmentJitterMsProp, segmentJitterMs: java.lang.Long)
-    logProps.put(LogConfig.CleanupPolicyProp, cleanupPolicy)
-    logProps.put(LogConfig.MaxMessageBytesProp, maxMessageBytes: Integer)
-    logProps.put(LogConfig.IndexIntervalBytesProp, indexIntervalBytes: Integer)
-    logProps.put(LogConfig.SegmentIndexBytesProp, segmentIndexBytes: Integer)
-    logProps.put(LogConfig.MessageFormatVersionProp, messageFormatVersion)
-    logProps.put(LogConfig.FileDeleteDelayMsProp, fileDeleteDelayMs: java.lang.Long)
-    LogConfig(logProps)
-  }
-
-  def createLog(dir: File,
-                config: LogConfig,
-                logStartOffset: Long = 0L,
-                recoveryPoint: Long = 0L,
-                scheduler: Scheduler = mockTime.scheduler,
-                brokerTopicStats: BrokerTopicStats = brokerTopicStats,
-                time: Time = mockTime,
-                maxProducerIdExpirationMs: Int = 60 * 60 * 1000,
-                producerIdExpirationCheckIntervalMs: Int = LogManager.ProducerIdExpirationCheckIntervalMs): Log = {
-    Log(dir = dir,
-        config = config,
-        logStartOffset = logStartOffset,
-        recoveryPoint = recoveryPoint,
-        scheduler = scheduler,
-        brokerTopicStats = brokerTopicStats,
-        time = time,
-        maxProducerIdExpirationMs = maxProducerIdExpirationMs,
-        producerIdExpirationCheckIntervalMs = producerIdExpirationCheckIntervalMs,
-        logDirFailureChannel = new LogDirFailureChannel(10))
+  @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)
@@ -2993,7 +4562,7 @@ class LogTest {
     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: _*)
@@ -3002,10 +4571,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 = {
@@ -3016,10 +4591,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)
@@ -3030,10 +4609,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 = {
@@ -3045,25 +4629,210 @@ class LogTest {
   }
 
   private def appendAsFollower(log: Log, records: MemoryRecords, leaderEpoch: Int = 0): Unit = {
-    records.batches.asScala.foreach(_.setPartitionLeaderEpoch(leaderEpoch))
+    records.batches.forEach(_.setPartitionLeaderEpoch(leaderEpoch))
     log.appendAsFollower(records)
   }
 
-  private def createCleanShutdownFile(): File = {
-    val parentLogDir = logDir.getParentFile
-    assertTrue("Data directory %s must exist", parentLogDir.isDirectory)
-    val cleanShutdownFile = new File(parentLogDir, Log.CleanShutdownFile)
-    cleanShutdownFile.createNewFile()
-    assertTrue(".kafka_cleanshutdown must exist", cleanShutdownFile.exists())
-    cleanShutdownFile
-  }
-
   private def deleteProducerSnapshotFiles(): Unit = {
     val files = logDir.listFiles.filter(f => f.isFile && f.getName.endsWith(Log.ProducerSnapshotFileSuffix))
     files.foreach(Utils.delete)
   }
 
   private def listProducerSnapshotOffsets: Seq[Long] =
-    ProducerStateManager.listSnapshotFiles(logDir).map(Log.offsetFromFile).sorted
+    ProducerStateManager.listSnapshotFiles(logDir).map(_.offset).sorted
+
+  private def createLog(dir: File,
+                        config: LogConfig,
+                        brokerTopicStats: BrokerTopicStats = brokerTopicStats,
+                        logStartOffset: Long = 0L,
+                        recoveryPoint: Long = 0L,
+                        scheduler: Scheduler = mockTime.scheduler,
+                        time: Time = mockTime,
+                        maxProducerIdExpirationMs: Int = 60 * 60 * 1000,
+                        producerIdExpirationCheckIntervalMs: Int = LogManager.ProducerIdExpirationCheckIntervalMs,
+                        lastShutdownClean: Boolean = true): Log = {
+    LogTest.createLog(dir, config, brokerTopicStats, scheduler, time, logStartOffset, recoveryPoint,
+      maxProducerIdExpirationMs, producerIdExpirationCheckIntervalMs, lastShutdownClean)
+  }
+
+  private def createLogWithOffsetOverflow(logConfig: LogConfig): (Log, LogSegment) = {
+    LogTest.initializeLogDirWithOverflowedSegment(logDir)
+
+    val log = createLog(logDir, logConfig, recoveryPoint = Long.MaxValue)
+    val segmentWithOverflow = LogTest.firstOverflowSegment(log).getOrElse {
+      Assertions.fail("Failed to create log with a segment which has overflowed offsets")
+    }
+
+    (log, segmentWithOverflow)
+  }
+
+  private def recoverAndCheck(config: LogConfig, expectedKeys: Iterable[Long]) = {
+    // method is called only in case of recovery from hard reset
+    LogTest.recoverAndCheck(logDir, config, expectedKeys, brokerTopicStats, mockTime, mockTime.scheduler)
+  }
+
+  private def readLog(log: Log,
+                      startOffset: Long,
+                      maxLength: Int,
+                      isolation: FetchIsolation = FetchLogEnd,
+                      minOneMessage: Boolean = true): FetchDataInfo = {
+    log.read(startOffset, maxLength, isolation, minOneMessage)
+  }
+
+}
+
+object LogTest {
+  def createLogConfig(segmentMs: Long = Defaults.SegmentMs,
+                      segmentBytes: Int = Defaults.SegmentSize,
+                      retentionMs: Long = Defaults.RetentionMs,
+                      retentionBytes: Long = Defaults.RetentionSize,
+                      segmentJitterMs: Long = Defaults.SegmentJitterMs,
+                      cleanupPolicy: String = Defaults.CleanupPolicy,
+                      maxMessageBytes: Int = Defaults.MaxMessageSize,
+                      indexIntervalBytes: Int = Defaults.IndexInterval,
+                      segmentIndexBytes: Int = Defaults.MaxIndexSize,
+                      messageFormatVersion: String = Defaults.MessageFormatVersion,
+                      fileDeleteDelayMs: Long = Defaults.FileDeleteDelayMs): LogConfig = {
+    val logProps = new Properties()
+
+    logProps.put(LogConfig.SegmentMsProp, segmentMs: java.lang.Long)
+    logProps.put(LogConfig.SegmentBytesProp, segmentBytes: Integer)
+    logProps.put(LogConfig.RetentionMsProp, retentionMs: java.lang.Long)
+    logProps.put(LogConfig.RetentionBytesProp, retentionBytes: java.lang.Long)
+    logProps.put(LogConfig.SegmentJitterMsProp, segmentJitterMs: java.lang.Long)
+    logProps.put(LogConfig.CleanupPolicyProp, cleanupPolicy)
+    logProps.put(LogConfig.MaxMessageBytesProp, maxMessageBytes: Integer)
+    logProps.put(LogConfig.IndexIntervalBytesProp, indexIntervalBytes: Integer)
+    logProps.put(LogConfig.SegmentIndexBytesProp, segmentIndexBytes: Integer)
+    logProps.put(LogConfig.MessageFormatVersionProp, messageFormatVersion)
+    logProps.put(LogConfig.FileDeleteDelayMsProp, fileDeleteDelayMs: java.lang.Long)
+    LogConfig(logProps)
+  }
+
+  def createLog(dir: File,
+                config: LogConfig,
+                brokerTopicStats: BrokerTopicStats,
+                scheduler: Scheduler,
+                time: Time,
+                logStartOffset: Long = 0L,
+                recoveryPoint: Long = 0L,
+                maxProducerIdExpirationMs: Int = 60 * 60 * 1000,
+                producerIdExpirationCheckIntervalMs: Int = LogManager.ProducerIdExpirationCheckIntervalMs,
+                lastShutdownClean: Boolean = true): Log = {
+    Log(dir = dir,
+      config = config,
+      logStartOffset = logStartOffset,
+      recoveryPoint = recoveryPoint,
+      scheduler = scheduler,
+      brokerTopicStats = brokerTopicStats,
+      time = time,
+      maxProducerIdExpirationMs = maxProducerIdExpirationMs,
+      producerIdExpirationCheckIntervalMs = producerIdExpirationCheckIntervalMs,
+      logDirFailureChannel = new LogDirFailureChannel(10),
+      lastShutdownClean = lastShutdownClean)
+  }
+
+  /**
+   * Check if the given log contains any segment with records that cause offset overflow.
+   * @param log Log to check
+   * @return true if log contains at least one segment with offset overflow; false otherwise
+   */
+  def hasOffsetOverflow(log: Log): Boolean = firstOverflowSegment(log).isDefined
+
+  def firstOverflowSegment(log: Log): Option[LogSegment] = {
+    def hasOverflow(baseOffset: Long, batch: RecordBatch): Boolean =
+      batch.lastOffset > baseOffset + Int.MaxValue || batch.baseOffset < baseOffset
+
+    for (segment <- log.logSegments) {
+      val overflowBatch = segment.log.batches.asScala.find(batch => hasOverflow(segment.baseOffset, batch))
+      if (overflowBatch.isDefined)
+        return Some(segment)
+    }
+    None
+  }
+
+  private def rawSegment(logDir: File, baseOffset: Long): FileRecords =
+    FileRecords.open(Log.logFile(logDir, baseOffset))
+
+  /**
+   * Initialize the given log directory with a set of segments, one of which will have an
+   * offset which overflows the segment
+   */
+  def initializeLogDirWithOverflowedSegment(logDir: File): Unit = {
+    def writeSampleBatches(baseOffset: Long, segment: FileRecords): Long = {
+      def record(offset: Long) = {
+        val data = offset.toString.getBytes
+        new SimpleRecord(data, data)
+      }
+
+      segment.append(MemoryRecords.withRecords(baseOffset, CompressionType.NONE, 0,
+        record(baseOffset)))
+      segment.append(MemoryRecords.withRecords(baseOffset + 1, CompressionType.NONE, 0,
+        record(baseOffset + 1),
+        record(baseOffset + 2)))
+      segment.append(MemoryRecords.withRecords(baseOffset + Int.MaxValue - 1, CompressionType.NONE, 0,
+        record(baseOffset + Int.MaxValue - 1)))
+      // Need to create the offset files explicitly to avoid triggering segment recovery to truncate segment.
+      Log.offsetIndexFile(logDir, baseOffset).createNewFile()
+      Log.timeIndexFile(logDir, baseOffset).createNewFile()
+      baseOffset + Int.MaxValue
+    }
 
+    def writeNormalSegment(baseOffset: Long): Long = {
+      val segment = rawSegment(logDir, baseOffset)
+      try writeSampleBatches(baseOffset, segment)
+      finally segment.close()
+    }
+
+    def writeOverflowSegment(baseOffset: Long): Long = {
+      val segment = rawSegment(logDir, baseOffset)
+      try {
+        val nextOffset = writeSampleBatches(baseOffset, segment)
+        writeSampleBatches(nextOffset, segment)
+      } finally segment.close()
+    }
+
+    // We create three segments, the second of which contains offsets which overflow
+    var nextOffset = 0L
+    nextOffset = writeNormalSegment(nextOffset)
+    nextOffset = writeOverflowSegment(nextOffset)
+    writeNormalSegment(nextOffset)
+  }
+
+  def allRecords(log: Log): List[Record] = {
+    val recordsFound = ListBuffer[Record]()
+    for (logSegment <- log.logSegments) {
+      for (batch <- logSegment.log.batches.asScala) {
+        recordsFound ++= batch.iterator().asScala
+      }
+    }
+    recordsFound.toList
+  }
+
+  def verifyRecordsInLog(log: Log, expectedRecords: List[Record]): Unit = {
+    assertEquals(expectedRecords, allRecords(log))
+  }
+
+  /* extract all the keys from a log */
+  def keysInLog(log: Log): Iterable[Long] = {
+    for (logSegment <- log.logSegments;
+         batch <- logSegment.log.batches.asScala if !batch.isControlBatch;
+         record <- batch.asScala if record.hasValue && record.hasKey)
+      yield TestUtils.readString(record.key).toLong
+  }
+
+  def recoverAndCheck(logDir: File, config: LogConfig, expectedKeys: Iterable[Long], brokerTopicStats: BrokerTopicStats, time: Time, scheduler: Scheduler): Log = {
+    // Recover log file and check that after recovery, keys are as expected
+    // and all temporary files have been deleted
+    val recoveredLog = createLog(logDir, config, brokerTopicStats, scheduler, time, lastShutdownClean = false)
+    time.sleep(config.fileDeleteDelayMs + 1)
+    for (file <- logDir.listFiles) {
+      assertFalse("Unexpected .deleted file after recovery", file.getName.endsWith(Log.DeletedFileSuffix))
+      assertFalse("Unexpected .cleaned file after recovery", file.getName.endsWith(Log.CleanedFileSuffix))
+      assertFalse("Unexpected .swap file after recovery", file.getName.endsWith(Log.SwapFileSuffix))
+    }
+    assertEquals(expectedKeys, LogTest.keysInLog(recoveredLog))
+    assertFalse(LogTest.hasOffsetOverflow(recoveredLog))
+    recoveredLog
+  }
 }
diff --git a/core/src/test/scala/unit/kafka/log/LogUtils.scala b/core/src/test/scala/unit/kafka/log/LogUtils.scala
new file mode 100644
index 0000000000000..8bf9812b86acd
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/log/LogUtils.scala
@@ -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 kafka.log
+
+import java.io.File
+
+import org.apache.kafka.common.record.FileRecords
+import org.apache.kafka.common.utils.Time
+
+object LogUtils {
+  /**
+    *  Create a segment with the given base offset
+    */
+  def createSegment(offset: Long,
+                    logDir: File,
+                    indexIntervalBytes: Int = 10,
+                    time: Time = Time.SYSTEM): LogSegment = {
+    val ms = FileRecords.open(Log.logFile(logDir, offset))
+    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 131152af4304f..f63d902810025 100644
--- a/core/src/test/scala/unit/kafka/log/LogValidatorTest.scala
+++ b/core/src/test/scala/unit/kafka/log/LogValidatorTest.scala
@@ -17,32 +17,138 @@
 package kafka.log
 
 import java.nio.ByteBuffer
-
-import kafka.common.LongRef
-import kafka.message.{CompressionCodec, DefaultCompressionCodec, GZIPCompressionCodec, NoCompressionCodec, SnappyCompressionCodec}
-import org.apache.kafka.common.errors.{InvalidTimestampException, UnsupportedForMessageFormatException}
+import java.util.concurrent.TimeUnit
+
+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.metrics.KafkaYammerMetrics
+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.{assertThrows, intercept}
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 
 class LogValidatorTest {
 
   val time = Time.SYSTEM
+  val topicPartition = new TopicPartition("topic", 0)
+  val brokerTopicStats = new BrokerTopicStats
+  val metricsKeySet = KafkaYammerMetrics.defaultRegistry.allMetrics.keySet.asScala
+
+  @Test
+  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 testLogAppendTimeNonCompressedV1() {
+  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) {
+  @Test
+  def testLogAppendTimeNonCompressedV2(): Unit = {
+    checkLogAppendTimeNonCompressed(RecordBatch.MAGIC_VALUE_V2)
+  }
+
+  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,
@@ -53,33 +159,35 @@ class LogValidatorTest {
       timestampType = TimestampType.LOG_APPEND_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true)
+      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))
+    validatedRecords.batches.forEach(batch => validateLogAppendTime(now, 1234L, batch))
     assertEquals(s"Max timestamp should be $now", now, validatedResults.maxTimestamp)
-    assertEquals(s"The offset of max timestamp should be 0", 0, validatedResults.shallowOffsetOfMaxTimestamp)
     assertFalse("Message size should not have been changed", validatedResults.messageSizeMaybeChanged)
 
-    verifyRecordsProcessingStats(validatedResults.recordsProcessingStats, numConvertedRecords = 0, records,
+    // we index from last offset in version 2 instead of base offset
+    val expectedMaxTimestampOffset = if (magic >= RecordBatch.MAGIC_VALUE_V2) 2 else 0
+    assertEquals(s"The offset of max timestamp should be $expectedMaxTimestampOffset",
+      expectedMaxTimestampOffset, validatedResults.shallowOffsetOfMaxTimestamp)
+    verifyRecordConversionStats(validatedResults.recordConversionStats, numConvertedRecords = 0, records,
       compressed = false)
   }
 
-  def testLogAppendTimeNonCompressedV2() {
-    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,
@@ -90,37 +198,40 @@ class LogValidatorTest {
       timestampType = TimestampType.LOG_APPEND_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true)
+      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, -1, batch))
+    validatedRecords.batches.forEach(batch => validateLogAppendTime(now, -1, batch))
     assertTrue("MessageSet should still valid", validatedRecords.batches.iterator.next().isValid)
     assertEquals(s"Max timestamp should be $now", now, validatedResults.maxTimestamp)
     assertEquals(s"The offset of max timestamp should be ${records.records.asScala.size - 1}",
       records.records.asScala.size - 1, validatedResults.shallowOffsetOfMaxTimestamp)
     assertTrue("Message size may have been changed", validatedResults.messageSizeMaybeChanged)
 
-    val stats = validatedResults.recordsProcessingStats
-    verifyRecordsProcessingStats(stats, numConvertedRecords = 3, records, compressed = true)
+    val stats = validatedResults.recordConversionStats
+    verifyRecordConversionStats(stats, numConvertedRecords = 3, records, compressed = true)
   }
 
   @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,
@@ -131,33 +242,82 @@ class LogValidatorTest {
       timestampType = TimestampType.LOG_APPEND_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true)
+      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))
+    validatedRecords.batches.forEach(batch => validateLogAppendTime(now, 1234L, batch))
     assertTrue("MessageSet should still valid", validatedRecords.batches.iterator.next().isValid)
     assertEquals(s"Max timestamp should be $now", now, validatedResults.maxTimestamp)
     assertEquals(s"The offset of max timestamp should be ${records.records.asScala.size - 1}",
       records.records.asScala.size - 1, validatedResults.shallowOffsetOfMaxTimestamp)
     assertFalse("Message size should not have been changed", validatedResults.messageSizeMaybeChanged)
 
-    verifyRecordsProcessingStats(validatedResults.recordsProcessingStats, numConvertedRecords = 0, records,
+    verifyRecordConversionStats(validatedResults.recordConversionStats, numConvertedRecords = 0, records,
       compressed = true)
   }
 
   @Test
-  def testLogAppendTimeWithoutRecompressionV2() {
+  def testInvalidOffsetRangeAndRecordCount(): Unit = {
+    // The batch to be written contains 3 records, so the correct lastOffsetDelta is 2
+    validateRecordBatchWithCountOverrides(lastOffsetDelta = 2, count = 3)
+
+    // Count and offset range are inconsistent or invalid
+    assertInvalidBatchCountOverrides(lastOffsetDelta = 0, count = 3)
+    assertInvalidBatchCountOverrides(lastOffsetDelta = 15, count = 3)
+    assertInvalidBatchCountOverrides(lastOffsetDelta = -3, count = 3)
+    assertInvalidBatchCountOverrides(lastOffsetDelta = 2, count = -3)
+    assertInvalidBatchCountOverrides(lastOffsetDelta = 2, count = 6)
+    assertInvalidBatchCountOverrides(lastOffsetDelta = 2, count = 0)
+    assertInvalidBatchCountOverrides(lastOffsetDelta = -3, count = -2)
+
+    // Count and offset range are consistent, but do not match the actual number of records
+    assertInvalidBatchCountOverrides(lastOffsetDelta = 5, count = 6)
+    assertInvalidBatchCountOverrides(lastOffsetDelta = 1, count = 2)
+  }
+
+  private def assertInvalidBatchCountOverrides(lastOffsetDelta: Int, count: Int): Unit = {
+    intercept[InvalidRecordException] {
+      validateRecordBatchWithCountOverrides(lastOffsetDelta, count)
+    }
+  }
+
+  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(),
+      sourceCodec = DefaultCompressionCodec,
+      targetCodec = DefaultCompressionCodec,
+      compactedTopic = false,
+      magic = RecordBatch.MAGIC_VALUE_V2,
+      timestampType = TimestampType.LOG_APPEND_TIME,
+      timestampDiffMaxMs = 1000L,
+      partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats)
+  }
+
+  @Test
+  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)
 
@@ -175,6 +335,7 @@ class LogValidatorTest {
       new SimpleRecord(timestampSeq(2), "beautiful".getBytes))
 
     val validatingResults = LogValidator.validateMessagesAndAssignOffsets(records,
+      topicPartition,
       offsetCounter = new LongRef(0),
       time = time,
       now = System.currentTimeMillis(),
@@ -185,7 +346,9 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = partitionLeaderEpoch,
-      isFromClient = true)
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats)
     val validatedRecords = validatingResults.validatedRecords
 
     var i = 0
@@ -209,12 +372,12 @@ class LogValidatorTest {
     assertEquals(s"Offset of max timestamp should be 1", 1, validatingResults.shallowOffsetOfMaxTimestamp)
     assertFalse("Message size should not have been changed", validatingResults.messageSizeMaybeChanged)
 
-    verifyRecordsProcessingStats(validatingResults.recordsProcessingStats, numConvertedRecords = 0, records,
+    verifyRecordConversionStats(validatingResults.recordConversionStats, numConvertedRecords = 0, records,
       compressed = false)
   }
 
   @Test
-  def testNonCompressedV2() {
+  def testNonCompressedV2(): Unit = {
     checkNonCompressed(RecordBatch.MAGIC_VALUE_V2)
   }
 
@@ -241,6 +404,7 @@ class LogValidatorTest {
       new SimpleRecord(timestampSeq(2), "beautiful".getBytes))
 
     val validatingResults = LogValidator.validateMessagesAndAssignOffsets(records,
+      topicPartition,
       offsetCounter = new LongRef(0),
       time = time,
       now = System.currentTimeMillis(),
@@ -251,7 +415,9 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = partitionLeaderEpoch,
-      isFromClient = true)
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats)
     val validatedRecords = validatingResults.validatedRecords
 
     var i = 0
@@ -274,7 +440,7 @@ class LogValidatorTest {
     assertEquals("Offset of max timestamp should be 2", 2, validatingResults.shallowOffsetOfMaxTimestamp)
     assertTrue("Message size should have been changed", validatingResults.messageSizeMaybeChanged)
 
-    verifyRecordsProcessingStats(validatingResults.recordsProcessingStats, numConvertedRecords = 3, records,
+    verifyRecordConversionStats(validatingResults.recordConversionStats, numConvertedRecords = 3, records,
       compressed = true)
   }
 
@@ -288,9 +454,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(),
@@ -301,7 +468,9 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true)
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats)
     val validatedRecords = validatedResults.validatedRecords
 
     for (batch <- validatedRecords.batches.asScala) {
@@ -318,20 +487,21 @@ class LogValidatorTest {
       validatedRecords.records.asScala.size - 1, validatedResults.shallowOffsetOfMaxTimestamp)
     assertTrue("Message size should have been changed", validatedResults.messageSizeMaybeChanged)
 
-    verifyRecordsProcessingStats(validatedResults.recordsProcessingStats, numConvertedRecords = 3, records,
+    verifyRecordConversionStats(validatedResults.recordConversionStats, numConvertedRecords = 3, records,
       compressed = true)
   }
 
   @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,
@@ -342,7 +512,9 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true)
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats)
     val validatedRecords = validatedResults.validatedRecords
 
     for (batch <- validatedRecords.batches.asScala) {
@@ -359,16 +531,16 @@ class LogValidatorTest {
       validatedRecords.records.asScala.size - 1, validatedResults.shallowOffsetOfMaxTimestamp)
     assertTrue("Message size should have been changed", validatedResults.messageSizeMaybeChanged)
 
-    verifyRecordsProcessingStats(validatedResults.recordsProcessingStats, numConvertedRecords = 3, records,
+    verifyRecordConversionStats(validatedResults.recordConversionStats, numConvertedRecords = 3, records,
       compressed = true)
   }
 
   @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)
 
@@ -386,6 +558,7 @@ class LogValidatorTest {
       new SimpleRecord(timestampSeq(2), "beautiful".getBytes))
 
     val validatedResults = LogValidator.validateMessagesAndAssignOffsets(records,
+      topicPartition,
       offsetCounter = new LongRef(0),
       time = time,
       now = System.currentTimeMillis(),
@@ -396,7 +569,9 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = partitionLeaderEpoch,
-      isFromClient = true)
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats)
     val validatedRecords = validatedResults.validatedRecords
 
     var i = 0
@@ -420,22 +595,23 @@ class LogValidatorTest {
       validatedRecords.records.asScala.size - 1, validatedResults.shallowOffsetOfMaxTimestamp)
     assertFalse("Message size should not have been changed", validatedResults.messageSizeMaybeChanged)
 
-    verifyRecordsProcessingStats(validatedResults.recordsProcessingStats, numConvertedRecords = 0, records,
+    verifyRecordConversionStats(validatedResults.recordConversionStats, numConvertedRecords = 0, records,
       compressed = true)
   }
 
   @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(),
@@ -446,16 +622,19 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true)
+      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(),
@@ -466,16 +645,19 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true)
+      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(),
@@ -486,16 +668,19 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true)
+      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(),
@@ -506,15 +691,18 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true)
+      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(),
@@ -525,15 +713,18 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true).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(),
@@ -544,16 +735,19 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true).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(),
@@ -564,17 +758,20 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 5000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true).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(),
@@ -585,18 +782,21 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 5000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true).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(),
@@ -607,18 +807,21 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 5000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true).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(),
@@ -629,16 +832,19 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 5000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true).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
-    checkOffsets(LogValidator.validateMessagesAndAssignOffsets(records,
+    val validatedResults = LogValidator.validateMessagesAndAssignOffsets(records,
+      topicPartition,
       offsetCounter = new LongRef(offset),
       time = time,
       now = System.currentTimeMillis(),
@@ -649,15 +855,21 @@ class LogValidatorTest {
       timestampType = TimestampType.LOG_APPEND_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true).validatedRecords, offset)
+      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
-    checkOffsets(LogValidator.validateMessagesAndAssignOffsets(records,
+    val validatedResults = LogValidator.validateMessagesAndAssignOffsets(records,
+      topicPartition,
       offsetCounter = new LongRef(offset),
       time = time,
       now = System.currentTimeMillis(),
@@ -668,15 +880,21 @@ class LogValidatorTest {
       timestampType = TimestampType.LOG_APPEND_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true).validatedRecords, offset)
+      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)
-    checkOffsets(LogValidator.validateMessagesAndAssignOffsets(records,
+    val validatedResults = LogValidator.validateMessagesAndAssignOffsets(records,
+      topicPartition,
       offsetCounter = new LongRef(offset),
       time = time,
       now = System.currentTimeMillis(),
@@ -687,15 +905,21 @@ class LogValidatorTest {
       timestampType = TimestampType.LOG_APPEND_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true).validatedRecords, offset)
+      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)
-    checkOffsets(LogValidator.validateMessagesAndAssignOffsets(records,
+    val validatedResults = LogValidator.validateMessagesAndAssignOffsets(records,
+      topicPartition,
       offsetCounter = new LongRef(offset),
       time = time,
       now = System.currentTimeMillis(),
@@ -706,15 +930,21 @@ class LogValidatorTest {
       timestampType = TimestampType.LOG_APPEND_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true).validatedRecords, offset)
+      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(),
@@ -725,15 +955,18 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 5000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true)
+      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(),
@@ -744,7 +977,9 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 5000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = false)
+      origin = AppendOrigin.Coordinator,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats)
     val batches = TestUtils.toList(result.validatedRecords.batches)
     assertEquals(1, batches.size)
     val batch = batches.get(0)
@@ -752,12 +987,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(),
@@ -768,16 +1004,19 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 5000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true).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(),
@@ -788,15 +1027,18 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 5000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true).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(),
@@ -807,15 +1049,18 @@ class LogValidatorTest {
       timestampType = TimestampType.LOG_APPEND_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true).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(),
@@ -826,16 +1071,19 @@ class LogValidatorTest {
       timestampType = TimestampType.LOG_APPEND_TIME,
       timestampDiffMaxMs = 1000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true).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(),
@@ -846,16 +1094,19 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 5000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true).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(),
@@ -866,11 +1117,13 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 5000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true).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
@@ -878,6 +1131,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(),
@@ -888,11 +1142,13 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 5000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true).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
@@ -900,6 +1156,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(),
@@ -910,16 +1167,19 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 5000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true).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(),
@@ -930,16 +1190,19 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 5000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true).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(),
@@ -950,25 +1213,34 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 5000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true).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)
+  @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])
@@ -976,6 +1248,28 @@ class LogValidatorTest {
     testBatchWithoutRecordsNotAllowed(DefaultCompressionCodec, DefaultCompressionCodec)
   }
 
+  @Test(expected = classOf[UnsupportedCompressionTypeException])
+  def testZStdCompressedWithUnavailableIBPVersion(): Unit = {
+    val now = System.currentTimeMillis()
+    // 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,
+      sourceCodec = NoCompressionCodec,
+      targetCodec = ZStdCompressionCodec,
+      compactedTopic = false,
+      magic = RecordBatch.MAGIC_VALUE_V2,
+      timestampType = TimestampType.LOG_APPEND_TIME,
+      timestampDiffMaxMs = 1000L,
+      partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = KAFKA_2_0_IV1,
+      brokerTopicStats = brokerTopicStats)
+  }
+
   @Test(expected = classOf[InvalidRecordException])
   def testUncompressedBatchWithoutRecordsNotAllowed(): Unit = {
     testBatchWithoutRecordsNotAllowed(NoCompressionCodec, NoCompressionCodec)
@@ -986,6 +1280,81 @@ 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, 3)
+  }
+
+  @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)
+    // recordsWithInvalidInnerMagic creates 20 records
+    assertEquals(e.recordErrors.size, 20)
+    e.recordErrors.foreach(assertNotNull(_))
+  }
+
+  @Test
+  def testBatchWithInvalidRecordsAndInvalidTimestamp(): Unit = {
+    val records = (0 until 5).map(id =>
+      LegacyRecord.create(RecordBatch.MAGIC_VALUE_V0, 0L, null, id.toString.getBytes())
+    )
+
+    val buffer = ByteBuffer.allocate(1024)
+    val builder = MemoryRecords.builder(buffer, RecordBatch.MAGIC_VALUE_V1, CompressionType.GZIP,
+      TimestampType.CREATE_TIME, 0L)
+    var offset = 0
+
+    // we want to mix in a record with invalid timestamp range
+    builder.appendUncheckedWithOffset(offset, LegacyRecord.create(RecordBatch.MAGIC_VALUE_V1,
+      1200L, null, "timestamp".getBytes))
+    records.foreach { record =>
+      offset += 30
+      builder.appendUncheckedWithOffset(offset, record)
+    }
+    val invalidOffsetTimestampRecords = builder.build()
+
+    val e = intercept[RecordValidationException] {
+      validateMessages(invalidOffsetTimestampRecords,
+        RecordBatch.MAGIC_VALUE_V0, CompressionType.GZIP, CompressionType.GZIP)
+    }
+    // if there is a mix of both regular InvalidRecordException and InvalidTimestampException,
+    // InvalidTimestampException takes precedence
+    assertTrue(e.invalidException.isInstanceOf[InvalidTimestampException])
+    assertTrue(e.recordErrors.nonEmpty)
+    assertEquals(6, e.recordErrors.size)
+  }
+
   private def testBatchWithoutRecordsNotAllowed(sourceCodec: CompressionCodec, targetCodec: CompressionCodec): Unit = {
     val offset = 1234567
     val (producerId, producerEpoch, baseSequence, isTransactional, partitionLeaderEpoch) =
@@ -997,6 +1366,7 @@ class LogValidatorTest {
     buffer.flip()
     val records = MemoryRecords.readableRecords(buffer)
     LogValidator.validateMessagesAndAssignOffsets(records,
+      topicPartition,
       offsetCounter = new LongRef(offset),
       time = time,
       now = System.currentTimeMillis(),
@@ -1007,7 +1377,9 @@ class LogValidatorTest {
       timestampType = TimestampType.CREATE_TIME,
       timestampDiffMaxMs = 5000L,
       partitionLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH,
-      isFromClient = true)
+      origin = AppendOrigin.Client,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      brokerTopicStats = brokerTopicStats)
   }
 
   private def createRecords(magicValue: Byte,
@@ -1021,8 +1393,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) {
@@ -1031,18 +1430,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
@@ -1062,7 +1480,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)
@@ -1073,12 +1491,14 @@ class LogValidatorTest {
     }
   }
 
-  def verifyRecordsProcessingStats(stats: RecordsProcessingStats, numConvertedRecords: Int, records: MemoryRecords,
-                                   compressed: Boolean): Unit = {
+  def verifyRecordConversionStats(stats: RecordConversionStats, numConvertedRecords: Int, records: MemoryRecords,
+                                  compressed: Boolean): Unit = {
     assertNotNull("Records processing info is null", stats)
     assertEquals(numConvertedRecords, stats.numRecordsConverted)
-    if (numConvertedRecords > 0)
+    if (numConvertedRecords > 0) {
       assertTrue(s"Conversion time not recorded $stats", stats.conversionTimeNanos >= 0)
+      assertTrue(s"Conversion time not valid $stats", stats.conversionTimeNanos <= TimeUnit.MINUTES.toNanos(1))
+    }
     val originalSize = records.sizeInBytes
     val tempBytes = stats.temporaryMemoryBytes
     if (numConvertedRecords > 0 && compressed)
diff --git a/core/src/test/scala/unit/kafka/log/OffsetIndexTest.scala b/core/src/test/scala/unit/kafka/log/OffsetIndexTest.scala
index 8fa3cc196481a..7a4b328ae25ea 100644
--- a/core/src/test/scala/unit/kafka/log/OffsetIndexTest.scala
+++ b/core/src/test/scala/unit/kafka/log/OffsetIndexTest.scala
@@ -24,31 +24,35 @@ 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 kafka.common.InvalidOffsetException
+import org.apache.kafka.common.errors.InvalidOffsetException
 
-class OffsetIndexTest extends JUnitSuite {
+import scala.annotation.nowarn
+
+class OffsetIndexTest {
   
   var idx: OffsetIndex = null
   val maxEntries = 30
+  val baseOffset = 45L
   
   @Before
-  def setup() {
-    this.idx = new OffsetIndex(nonExistentTempFile(), baseOffset = 45L, maxIndexSize = 30 * 8)
+  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()
   }
-  
+
+  @nowarn("cat=deprecation")
   @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
@@ -76,7 +80,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)
@@ -84,9 +88,22 @@ class OffsetIndexTest extends JUnitSuite {
     assertEquals(OffsetPosition(idx.baseOffset, 0), idx.lookup(idx.baseOffset))
     assertEquals(OffsetPosition(idx.baseOffset + idx.maxEntries, idx.maxEntries - 1), idx.lookup(idx.baseOffset + idx.maxEntries))
   }
+
+  @Test
+  def testEntry(): Unit = {
+    for (i <- 0 until idx.maxEntries)
+      idx.append(idx.baseOffset + i + 1, i)
+    for (i <- 0 until idx.maxEntries)
+      assertEquals(OffsetPosition(idx.baseOffset + i + 1, i), idx.entry(i))
+  }
+
+  @Test(expected = classOf[IllegalArgumentException])
+  def testEntryOverflow(): Unit = {
+    idx.entry(0)
+  }
   
   @Test
-  def appendTooMany() {
+  def appendTooMany(): Unit = {
     for(i <- 0 until idx.maxEntries) {
       val offset = idx.baseOffset + i + 1
       idx.append(offset, i)
@@ -95,17 +112,17 @@ class OffsetIndexTest extends JUnitSuite {
   }
   
   @Test(expected = classOf[InvalidOffsetException])
-  def appendOutOfOrder() {
+  def appendOutOfOrder(): Unit = {
     idx.append(51, 0)
     idx.append(50, 1)
   }
 
   @Test
-  def testFetchUpperBoundOffset() {
-    val first = OffsetPosition(0, 0)
-    val second = OffsetPosition(1, 10)
-    val third = OffsetPosition(2, 23)
-    val fourth = OffsetPosition(3, 37)
+  def testFetchUpperBoundOffset(): Unit = {
+    val first = OffsetPosition(baseOffset + 0, 0)
+    val second = OffsetPosition(baseOffset + 1, 10)
+    val third = OffsetPosition(baseOffset + 2, 23)
+    val fourth = OffsetPosition(baseOffset + 3, 37)
 
     assertEquals(None, idx.fetchUpperBoundOffset(first, 5))
 
@@ -123,7 +140,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)
@@ -138,7 +155,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)
@@ -177,8 +194,17 @@ class OffsetIndexTest extends JUnitSuite {
     // mmap should be null after unmap causing lookup to throw a NPE
     intercept[NullPointerException](idx.lookup(1))
   }
+
+  @Test
+  def testSanityLastOffsetEqualToBaseOffset(): Unit = {
+    // Test index sanity for the case where the last offset appended to the index is equal to the base offset
+    val baseOffset = 20L
+    val idx = new OffsetIndex(nonExistentTempFile(), baseOffset = baseOffset, maxIndexSize = 10 * 8)
+    idx.append(baseOffset, 0)
+    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..4ac0e7b1124e2 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)
@@ -83,7 +82,7 @@ object OffsetMapTest {
     val start = System.nanoTime
     val map = test.validateMap(size, load)
     val ellapsedMs = (System.nanoTime - start) / 1000.0 / 1000.0
-    println(map.size + " entries in map of size " + map.slots + " in " + ellapsedMs + " ms")
+    println(s"${map.size} entries in map of size ${map.slots} in $ellapsedMs ms")
     println("Collision rate: %.1f%%".format(100*map.collisionRate))
   }
 }
diff --git a/core/src/test/scala/unit/kafka/log/ProducerStateManagerTest.scala b/core/src/test/scala/unit/kafka/log/ProducerStateManagerTest.scala
index 67b1b15f2dc31..207e129cbccec 100644
--- a/core/src/test/scala/unit/kafka/log/ProducerStateManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/log/ProducerStateManagerTest.scala
@@ -21,19 +21,22 @@ import java.io.File
 import java.nio.ByteBuffer
 import java.nio.channels.FileChannel
 import java.nio.file.StandardOpenOption
+import java.util.Collections
+import java.util.concurrent.atomic.AtomicInteger
 
 import kafka.server.LogOffsetMetadata
 import kafka.utils.TestUtils
 import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.errors._
 import org.apache.kafka.common.internals.Topic
-import org.apache.kafka.common.record.{ControlRecordType, EndTransactionMarker, RecordBatch}
+import org.apache.kafka.common.record._
 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)
@@ -62,8 +65,8 @@ class ProducerStateManagerTest extends JUnitSuite {
     // Second entry for id 0 added
     append(stateManager, producerId, epoch, 1, 0L, 1L)
 
-    // Duplicate sequence number (matches previous sequence number)
-    assertThrows[DuplicateSequenceException] {
+    // Duplicates are checked separately and should result in OutOfOrderSequence if appended
+    assertThrows[OutOfOrderSequenceException] {
       append(stateManager, producerId, epoch, 1, 0L, 1L)
     }
 
@@ -76,17 +79,46 @@ class ProducerStateManagerTest extends JUnitSuite {
     append(stateManager, producerId, (epoch + 1).toShort, 0, 0L, 3L)
 
     // Incorrect epoch
-    assertThrows[ProducerFencedException] {
+    assertThrows[InvalidProducerEpochException] {
       append(stateManager, producerId, epoch, 0, 0L, 4L)
     }
   }
 
+  @Test
+  def testAppendTxnMarkerWithNoProducerState(): Unit = {
+    val producerEpoch = 2.toShort
+    appendEndTxnMarker(stateManager, producerId, producerEpoch, ControlRecordType.COMMIT, offset = 27L)
+
+    val firstEntry = stateManager.lastEntry(producerId).getOrElse(fail("Expected last entry to be defined"))
+    assertEquals(producerEpoch, firstEntry.producerEpoch)
+    assertEquals(producerId, firstEntry.producerId)
+    assertEquals(RecordBatch.NO_SEQUENCE, firstEntry.lastSeq)
+
+    // Fencing should continue to work even if the marker is the only thing left
+    assertThrows[InvalidProducerEpochException] {
+      append(stateManager, producerId, 0.toShort, 0, 0L, 4L)
+    }
+
+    // If the transaction marker is the only thing left in the log, then an attempt to write using a
+    // non-zero sequence number should cause an OutOfOrderSequenceException, so that the producer can reset its state
+    assertThrows[OutOfOrderSequenceException] {
+      append(stateManager, producerId, producerEpoch, 17, 0L, 4L)
+    }
+
+    // The broker should accept the request if the sequence number is reset to 0
+    append(stateManager, producerId, producerEpoch, 0, 39L, 4L)
+    val secondEntry = stateManager.lastEntry(producerId).getOrElse(fail("Expected last entry to be defined"))
+    assertEquals(producerEpoch, secondEntry.producerEpoch)
+    assertEquals(producerId, secondEntry.producerId)
+    assertEquals(0, secondEntry.lastSeq)
+  }
+
   @Test
   def testProducerSequenceWrapAround(): Unit = {
     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)
 
@@ -100,12 +132,31 @@ class ProducerStateManagerTest extends JUnitSuite {
     assertEquals(0, lastEntry.lastSeq)
   }
 
+  @Test
+  def testProducerSequenceWithWrapAroundBatchRecord(): Unit = {
+    val epoch = 15.toShort
+
+    val appendInfo = stateManager.prepareUpdate(producerId, origin = AppendOrigin.Replication)
+    // Sequence number wrap around
+    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)
+
+    val lastEntry = stateManager.lastEntry(producerId).get
+    assertEquals(Int.MaxValue-10, lastEntry.firstSeq)
+    assertEquals(9, lastEntry.lastSeq)
+    assertEquals(2000L, lastEntry.firstDataOffset)
+    assertEquals(2020L, lastEntry.lastDataOffset)
+  }
+
   @Test(expected = classOf[OutOfOrderSequenceException])
   def testProducerSequenceInvalidWrapAround(): Unit = {
     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)
   }
 
@@ -114,7 +165,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)
@@ -124,33 +175,28 @@ 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
-  def testControlRecordBumpsEpoch(): Unit = {
-    val epoch = 0.toShort
-    append(stateManager, producerId, epoch, 0, 0L)
+  def testControlRecordBumpsProducerEpoch(): Unit = {
+    val producerEpoch = 0.toShort
+    append(stateManager, producerId, producerEpoch, 0, 0L)
 
-    val bumpedEpoch = 1.toShort
-    val (completedTxn, lastStableOffset) = appendEndTxnMarker(stateManager, producerId, bumpedEpoch, ControlRecordType.ABORT, 1L)
-    assertEquals(1L, completedTxn.firstOffset)
-    assertEquals(1L, completedTxn.lastOffset)
-    assertEquals(2L, lastStableOffset)
-    assertTrue(completedTxn.isAborted)
-    assertEquals(producerId, completedTxn.producerId)
+    val bumpedProducerEpoch = 1.toShort
+    appendEndTxnMarker(stateManager, producerId, bumpedProducerEpoch, ControlRecordType.ABORT, 1L)
 
     val maybeLastEntry = stateManager.lastEntry(producerId)
     assertTrue(maybeLastEntry.isDefined)
 
     val lastEntry = maybeLastEntry.get
-    assertEquals(bumpedEpoch, lastEntry.producerEpoch)
+    assertEquals(bumpedProducerEpoch, lastEntry.producerEpoch)
     assertEquals(None, lastEntry.currentTxnFirstOffset)
     assertEquals(RecordBatch.NO_SEQUENCE, lastEntry.firstSeq)
     assertEquals(RecordBatch.NO_SEQUENCE, lastEntry.lastSeq)
 
     // should be able to append with the new epoch if we start at sequence 0
-    append(stateManager, producerId, bumpedEpoch, 0, 2L)
+    append(stateManager, producerId, bumpedProducerEpoch, 0, 2L)
     assertEquals(Some(0), stateManager.lastEntry(producerId).map(_.firstSeq))
   }
 
@@ -159,33 +205,155 @@ class ProducerStateManagerTest extends JUnitSuite {
     val producerEpoch = 0.toShort
     val offset = 992342L
     val seq = 0
-    val producerAppendInfo = new ProducerAppendInfo(producerId, ProducerIdEntry.empty(producerId), ValidationType.Full)
-    producerAppendInfo.append(producerEpoch, seq, seq, time.milliseconds(), 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 testSkipEmptyTransactions(): Unit = {
     val producerEpoch = 0.toShort
-    val offset = 992342L
-    val seq = 0
-    val producerAppendInfo = new ProducerAppendInfo(producerId, ProducerIdEntry.empty(producerId), ValidationType.Full)
-    producerAppendInfo.append(producerEpoch, seq, seq, time.milliseconds(), offset, isTransactional = true)
+    val coordinatorEpoch = 27
+    val seq = new AtomicInteger(0)
+
+    def appendEndTxn(
+      recordType: ControlRecordType,
+      offset: Long,
+      appendInfo: ProducerAppendInfo
+    ): Option[CompletedTxn] = {
+      appendInfo.appendEndTxnMarker(new EndTransactionMarker(recordType, coordinatorEpoch),
+        producerEpoch, offset, time.milliseconds())
+    }
 
-    // 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)
+    def appendData(
+      startOffset: Long,
+      endOffset: Long,
+      appendInfo: ProducerAppendInfo
+    ): Unit = {
+      val count = (endOffset - startOffset).toInt
+      appendInfo.appendDataBatch(producerEpoch, seq.get(), seq.addAndGet(count), time.milliseconds(),
+        LogOffsetMetadata(startOffset), endOffset, isTransactional = true)
+      seq.incrementAndGet()
+    }
+
+    // Start one transaction in a separate append
+    val firstAppend = stateManager.prepareUpdate(producerId, origin = AppendOrigin.Client)
+    appendData(16L, 20L, firstAppend)
+    assertEquals(new TxnMetadata(producerId, 16L), firstAppend.startedTransactions.head)
+    stateManager.update(firstAppend)
+    stateManager.onHighWatermarkUpdated(21L)
+    assertEquals(Some(LogOffsetMetadata(16L)), stateManager.firstUnstableOffset)
+
+    // Now do a single append which completes the old transaction, mixes in
+    // some empty transactions, one non-empty complete transaction, and one
+    // incomplete transaction
+    val secondAppend = stateManager.prepareUpdate(producerId, origin = AppendOrigin.Client)
+    val firstCompletedTxn = appendEndTxn(ControlRecordType.COMMIT, 21, secondAppend)
+    assertEquals(Some(CompletedTxn(producerId, 16L, 21, isAborted = false)), firstCompletedTxn)
+    assertEquals(None, appendEndTxn(ControlRecordType.COMMIT, 22, secondAppend))
+    assertEquals(None, appendEndTxn(ControlRecordType.ABORT, 23, secondAppend))
+    appendData(24L, 27L, secondAppend)
+    val secondCompletedTxn = appendEndTxn(ControlRecordType.ABORT, 28L, secondAppend)
+    assertTrue(secondCompletedTxn.isDefined)
+    assertEquals(None, appendEndTxn(ControlRecordType.ABORT, 29L, secondAppend))
+    appendData(30L, 31L, secondAppend)
+
+    assertEquals(2, secondAppend.startedTransactions.size)
+    assertEquals(TxnMetadata(producerId, LogOffsetMetadata(24L)), secondAppend.startedTransactions.head)
+    assertEquals(TxnMetadata(producerId, LogOffsetMetadata(30L)), secondAppend.startedTransactions.last)
+    stateManager.update(secondAppend)
+    stateManager.completeTxn(firstCompletedTxn.get)
+    stateManager.completeTxn(secondCompletedTxn.get)
+    stateManager.onHighWatermarkUpdated(32L)
+    assertEquals(Some(LogOffsetMetadata(30L)), stateManager.firstUnstableOffset)
+  }
+
+  @Test
+  def testLastStableOffsetCompletedTxn(): Unit = {
+    val producerEpoch = 0.toShort
+    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)
+    }
+
+    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))
+
+    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)
+  }
 
-    assertEquals(Some(LogOffsetMetadata(offset)), stateManager.firstUnstableOffset)
+  @Test
+  def testPrepareUpdateDoesNotMutate(): Unit = {
+    val producerEpoch = 0.toShort
+
+    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, 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
+    assertEquals(0, lastEntry.firstSeq)
+    assertEquals(5, lastEntry.lastSeq)
+    assertEquals(20L, lastEntry.lastDataOffset)
+
+    stateManager.update(nextAppendInfo)
+    lastEntry = stateManager.lastEntry(producerId).get
+    assertEquals(0, lastEntry.firstSeq)
+    assertEquals(10, lastEntry.lastSeq)
+    assertEquals(30L, lastEntry.lastDataOffset)
   }
 
   @Test
@@ -195,57 +363,73 @@ 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(), 20L, isTransactional = true)
-    var lastEntry = appendInfo.latestEntry
+    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(0, lastEntry.firstSeq)
+    assertEquals(1, lastEntry.firstSeq)
     assertEquals(5, lastEntry.lastSeq)
-    assertEquals(9L, 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(), 30L, isTransactional = true)
-    lastEntry = appendInfo.latestEntry
+    appendInfo.appendDataBatch(producerEpoch, 6, 10, time.milliseconds(),
+      LogOffsetMetadata(26L), 30L, isTransactional = true)
+    lastEntry = appendInfo.toEntry
     assertEquals(producerEpoch, lastEntry.producerEpoch)
-    assertEquals(0, lastEntry.firstSeq)
+    assertEquals(1, lastEntry.firstSeq)
     assertEquals(10, lastEntry.lastSeq)
-    assertEquals(9L, lastEntry.firstOffset)
+    assertEquals(16L, lastEntry.firstDataOffset)
     assertEquals(30L, lastEntry.lastDataOffset)
     assertEquals(Some(16L), lastEntry.currentTxnFirstOffset)
     assertEquals(List(new TxnMetadata(producerId, 16L)), appendInfo.startedTransactions)
 
     val endTxnMarker = new EndTransactionMarker(ControlRecordType.COMMIT, coordinatorEpoch)
-    val completedTxn = appendInfo.appendEndTxnMarker(endTxnMarker, producerEpoch, 40L, time.milliseconds())
+    val completedTxnOpt = appendInfo.appendEndTxnMarker(endTxnMarker, producerEpoch, 40L, time.milliseconds())
+    assertTrue(completedTxnOpt.isDefined)
+
+    val completedTxn = completedTxnOpt.get
     assertEquals(producerId, completedTxn.producerId)
     assertEquals(16L, completedTxn.firstOffset)
     assertEquals(40L, completedTxn.lastOffset)
     assertFalse(completedTxn.isAborted)
 
-    lastEntry = appendInfo.latestEntry
+    lastEntry = appendInfo.toEntry
     assertEquals(producerEpoch, lastEntry.producerEpoch)
     // verify that appending the transaction marker doesn't affect the metadata of the cached record batches.
-    assertEquals(0, lastEntry.firstSeq)
+    assertEquals(1, lastEntry.firstSeq)
     assertEquals(10, lastEntry.lastSeq)
-    assertEquals(9L, lastEntry.firstOffset)
+    assertEquals(16L, lastEntry.firstDataOffset)
     assertEquals(30L, lastEntry.lastDataOffset)
     assertEquals(coordinatorEpoch, lastEntry.coordinatorEpoch)
     assertEquals(None, lastEntry.currentTxnFirstOffset)
     assertEquals(List(new TxnMetadata(producerId, 16L)), appendInfo.startedTransactions)
   }
 
-  @Test(expected = classOf[OutOfOrderSequenceException])
+  @Test
   def testOutOfSequenceAfterControlRecordEpochBump(): 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)
 
     val bumpedEpoch = 1.toShort
     appendEndTxnMarker(stateManager, producerId, bumpedEpoch, ControlRecordType.ABORT, 1L)
 
     // next append is invalid since we expect the sequence to be reset
-    append(stateManager, producerId, bumpedEpoch, 2, 2L)
+    assertThrows[OutOfOrderSequenceException] {
+      append(stateManager, producerId, bumpedEpoch, 2, 2L, isTransactional = true)
+    }
+
+    assertThrows[OutOfOrderSequenceException] {
+      append(stateManager, producerId, (bumpedEpoch + 1).toShort, 2, 2L, isTransactional = true)
+    }
+
+    // Append with the bumped epoch should be fine if starting from sequence 0
+    append(stateManager, producerId, bumpedEpoch, 0, 0L, isTransactional = true)
+    assertEquals(bumpedEpoch, stateManager.lastEntry(producerId).get.producerEpoch)
+    assertEquals(0, stateManager.lastEntry(producerId).get.lastSeq)
   }
 
   @Test(expected = classOf[InvalidTxnStateException])
@@ -290,20 +474,81 @@ 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(expected = classOf[UnknownProducerIdException])
+  @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
   def testRemoveExpiredPidsOnReload(): Unit = {
     val epoch = 0.toShort
     append(stateManager, producerId, epoch, 0, 0L, 0)
@@ -314,8 +559,12 @@ class ProducerStateManagerTest extends JUnitSuite {
     recoveredMapping.truncateAndReload(0L, 1L, 70000)
 
     // entry added after recovery. The pid should be expired now, and would not exist in the pid mapping. Hence
-    // we should get an out of order sequence exception.
+    // we should accept the append and add the pid back in
     append(recoveredMapping, producerId, epoch, 2, 2L, 70001)
+
+    assertEquals(1, recoveredMapping.activeProducers.size)
+    assertEquals(2, recoveredMapping.activeProducers.head._2.lastSeq)
+    assertEquals(3L, recoveredMapping.mapEndOffset)
   }
 
   @Test
@@ -332,12 +581,12 @@ 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 producerIdEntry = recoveredMapping.activeProducers.get(producerId).head
-    assertEquals(epoch, producerIdEntry.producerEpoch)
-    assertEquals(sequence, producerIdEntry.firstSeq)
-    assertEquals(sequence, producerIdEntry.lastSeq)
+    val producerStateEntry = recoveredMapping.activeProducers.get(producerId).head
+    assertEquals(epoch, producerStateEntry.producerEpoch)
+    assertEquals(sequence, producerStateEntry.firstSeq)
+    assertEquals(sequence, producerStateEntry.lastSeq)
   }
 
   @Test
@@ -348,7 +597,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 =>
@@ -358,7 +607,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)
   }
 
@@ -434,52 +683,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
@@ -490,13 +694,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 entry1 = stateManager.lastEntry(pid1)
+    assertTrue(entry1.isDefined)
+    assertEquals(0, entry1.get.lastSeq)
+    assertEquals(0L, entry1.get.lastDataOffset)
 
-    val entry = stateManager.lastEntry(pid2)
-    assertTrue(entry.isDefined)
-    assertEquals(0, entry.get.lastSeq)
-    assertEquals(1L, entry.get.lastDataOffset)
+    val entry2 = stateManager.lastEntry(pid2)
+    assertTrue(entry2.isDefined)
+    assertEquals(0, entry2.get.lastSeq)
+    assertEquals(1L, entry2.get.lastDataOffset)
   }
 
   @Test
@@ -515,34 +723,20 @@ class ProducerStateManagerTest extends JUnitSuite {
   }
 
   @Test
-  def testStartOffset(): Unit = {
-    val epoch = 0.toShort
-    val pid2 = 2L
-    append(stateManager, pid2, epoch, 0, 0L, 1L)
-    append(stateManager, producerId, epoch, 0, 1L, 2L)
-    append(stateManager, producerId, epoch, 1, 2L, 3L)
-    append(stateManager, producerId, epoch, 2, 3L, 4L)
-    stateManager.takeSnapshot()
-
-    intercept[UnknownProducerIdException] {
-      val recoveredMapping = new ProducerStateManager(partition, logDir, maxPidExpirationMs)
-      recoveredMapping.truncateAndReload(0L, 1L, time.milliseconds)
-      append(recoveredMapping, pid2, epoch, 1, 4L, 5L)
-    }
-  }
-
-  @Test(expected = classOf[UnknownProducerIdException])
-  def testPidExpirationTimeout() {
+  def testPidExpirationTimeout(): Unit = {
     val epoch = 5.toShort
     val sequence = 37
     append(stateManager, producerId, epoch, sequence, 1L)
     time.sleep(maxPidExpirationMs + 1)
     stateManager.removeExpiredProducers(time.milliseconds)
-    append(stateManager, producerId, epoch, sequence + 1, 1L)
+    append(stateManager, producerId, epoch, sequence + 1, 2L)
+    assertEquals(1, stateManager.activeProducers.size)
+    assertEquals(sequence + 1, stateManager.activeProducers.head._2.lastSeq)
+    assertEquals(3L, stateManager.mapEndOffset)
   }
 
   @Test
-  def testFirstUnstableOffset() {
+  def testFirstUnstableOffset(): Unit = {
     val epoch = 5.toShort
     val sequence = 0
 
@@ -576,7 +770,7 @@ class ProducerStateManagerTest extends JUnitSuite {
   }
 
   @Test
-  def testProducersWithOngoingTransactionsDontExpire() {
+  def testProducersWithOngoingTransactionsDontExpire(): Unit = {
     val epoch = 5.toShort
     val sequence = 0
 
@@ -599,12 +793,13 @@ 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])
+  @Test(expected = classOf[InvalidProducerEpochException])
   def testOldEpochForControlRecord(): Unit = {
     val epoch = 5.toShort
     val sequence = 0
@@ -680,6 +875,55 @@ class ProducerStateManagerTest extends JUnitSuite {
     }
   }
 
+  @Test
+  def testAppendEmptyControlBatch(): Unit = {
+    val producerId = 23423L
+    val baseOffset = 15
+
+    val batch: RecordBatch = EasyMock.createMock(classOf[RecordBatch])
+    EasyMock.expect(batch.isControlBatch).andReturn(true).once
+    EasyMock.expect(batch.iterator).andReturn(Collections.emptyIterator[Record]).once
+    EasyMock.replay(batch)
+
+    // Appending the empty control batch should not throw and a new transaction shouldn't be started
+    append(stateManager, producerId, baseOffset, batch, origin = AppendOrigin.Client)
+    assertEquals(None, stateManager.lastEntry(producerId).get.currentTxnFirstOffset)
+  }
+
+  @Test
+  def testRemoveStraySnapshotsKeepCleanShutdownSnapshot(): Unit = {
+    // Test that when stray snapshots are removed, the largest stray snapshot is kept around. This covers the case where
+    // the broker shutdown cleanly and emitted a snapshot file larger than the base offset of the active segment.
+
+    // Create 3 snapshot files at different offsets.
+    Log.producerSnapshotFile(logDir, 5).createNewFile() // not stray
+    Log.producerSnapshotFile(logDir, 2).createNewFile() // stray
+    Log.producerSnapshotFile(logDir, 42).createNewFile() // not stray
+
+    // claim that we only have one segment with a base offset of 5
+    stateManager.removeStraySnapshots(Seq(5))
+
+    // The snapshot file at offset 2 should be considered a stray, but the snapshot at 42 should be kept
+    // around because it is the largest snapshot.
+    assertEquals(Some(42), stateManager.latestSnapshotOffset)
+    assertEquals(Some(5), stateManager.oldestSnapshotOffset)
+    assertEquals(Seq(5, 42), ProducerStateManager.listSnapshotFiles(logDir).map(_.offset).sorted)
+  }
+
+  @Test
+  def testRemoveAllStraySnapshots(): Unit = {
+    // Test that when stray snapshots are removed, we remove only the stray snapshots below the largest segment base offset.
+    // Snapshots associated with an offset in the list of segment base offsets should remain.
+
+    // Create 3 snapshot files at different offsets.
+    Log.producerSnapshotFile(logDir, 5).createNewFile() // stray
+    Log.producerSnapshotFile(logDir, 2).createNewFile() // stray
+    Log.producerSnapshotFile(logDir, 42).createNewFile() // not stray
+
+    stateManager.removeStraySnapshots(Seq(42))
+    assertEquals(Seq(42), ProducerStateManager.listSnapshotFiles(logDir).map(_.offset).sorted)
+  }
+
   private def testLoadFromCorruptSnapshot(makeFileCorrupt: FileChannel => Unit): Unit = {
     val epoch = 0.toShort
     val producerId = 1L
@@ -716,14 +960,14 @@ class ProducerStateManagerTest extends JUnitSuite {
                                  controlType: ControlRecordType,
                                  offset: Long,
                                  coordinatorEpoch: Int = 0,
-                                 timestamp: Long = time.milliseconds()): (CompletedTxn, Long) = {
-    val producerAppendInfo = stateManager.prepareUpdate(producerId, isFromClient = true)
+                                 timestamp: Long = time.milliseconds()): Option[CompletedTxn] = {
+    val producerAppendInfo = stateManager.prepareUpdate(producerId, origin = AppendOrigin.Coordinator)
     val endTxnMarker = new EndTransactionMarker(controlType, coordinatorEpoch)
-    val completedTxn = producerAppendInfo.appendEndTxnMarker(endTxnMarker, producerEpoch, offset, timestamp)
+    val completedTxnOpt = producerAppendInfo.appendEndTxnMarker(endTxnMarker, producerEpoch, offset, timestamp)
     mapping.update(producerAppendInfo)
-    val lastStableOffset = mapping.completeTxn(completedTxn)
+    completedTxnOpt.foreach(mapping.completeTxn)
     mapping.updateMapEndOffset(offset + 1)
-    (completedTxn, lastStableOffset)
+    completedTxnOpt
   }
 
   private def append(stateManager: ProducerStateManager,
@@ -733,14 +977,26 @@ 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, 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)
+  }
+
+  private def append(stateManager: ProducerStateManager,
+                     producerId: Long,
+                     offset: Long,
+                     batch: RecordBatch,
+                     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 c6112a1a007e7..bfe8f56882f15 100644
--- a/core/src/test/scala/unit/kafka/log/TimeIndexTest.scala
+++ b/core/src/test/scala/unit/kafka/log/TimeIndexTest.scala
@@ -19,33 +19,33 @@ package kafka.log
 
 import java.io.File
 
-import kafka.common.InvalidOffsetException
 import kafka.utils.TestUtils
-import org.junit.{Test, After, Before}
-import org.junit.Assert.{assertEquals}
-import org.scalatest.junit.JUnitSuite
+import org.apache.kafka.common.errors.InvalidOffsetException
+import org.junit.{After, Before, Test}
+import org.junit.Assert.assertEquals
+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))
 
@@ -61,7 +61,21 @@ class TimeIndexTest extends JUnitSuite {
   }
 
   @Test
-  def testTruncate() {
+  def testEntry(): Unit = {
+    appendEntries(maxEntries - 1)
+    assertEquals(TimestampOffset(10L, 55L), idx.entry(0))
+    assertEquals(TimestampOffset(20L, 65L), idx.entry(1))
+    assertEquals(TimestampOffset(30L, 75L), idx.entry(2))
+    assertEquals(TimestampOffset(40L, 85L), idx.entry(3))
+  }
+
+  @Test(expected = classOf[IllegalArgumentException])
+  def testEntryOverflow(): Unit = {
+    idx.entry(0)
+  }
+
+  @Test
+  def testTruncate(): Unit = {
     appendEntries(maxEntries - 1)
     idx.truncate()
     assertEquals(0, idx.entries)
@@ -72,7 +86,7 @@ class TimeIndexTest extends JUnitSuite {
   }
 
   @Test
-  def testAppend() {
+  def testAppend(): Unit = {
     appendEntries(maxEntries - 1)
     intercept[IllegalArgumentException] {
       idx.maybeAppend(10000L, 1000L)
@@ -83,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)
   }
@@ -93,5 +107,46 @@ class TimeIndexTest extends JUnitSuite {
     file.delete()
     file
   }
+
+  @Test
+  def testSanityCheck(): Unit = {
+    idx.sanityCheck()
+    appendEntries(5)
+    val firstEntry = idx.entry(0)
+    idx.sanityCheck()
+    idx.close()
+
+    var shouldCorruptOffset = false
+    var shouldCorruptTimestamp = false
+    var shouldCorruptLength = false
+    idx = new TimeIndex(idx.file, baseOffset = baseOffset, maxIndexSize = maxEntries * 12) {
+      override def lastEntry = {
+        val superLastEntry = super.lastEntry
+        val offset = if (shouldCorruptOffset) baseOffset - 1 else superLastEntry.offset
+        val timestamp = if (shouldCorruptTimestamp) firstEntry.timestamp - 1 else superLastEntry.timestamp
+        new TimestampOffset(timestamp, offset)
+      }
+      override def length = {
+        val superLength = super.length
+        if (shouldCorruptLength) superLength - 1 else superLength
+      }
+    }
+
+    shouldCorruptOffset = true
+    intercept[CorruptIndexException](idx.sanityCheck())
+    shouldCorruptOffset = false
+
+    shouldCorruptTimestamp = true
+    intercept[CorruptIndexException](idx.sanityCheck())
+    shouldCorruptTimestamp = false
+
+    shouldCorruptLength = true
+    intercept[CorruptIndexException](idx.sanityCheck())
+    shouldCorruptLength = false
+
+    idx.sanityCheck()
+    idx.close()
+  }
+
 }
 
diff --git a/core/src/test/scala/unit/kafka/log/TransactionIndexTest.scala b/core/src/test/scala/unit/kafka/log/TransactionIndexTest.scala
index 9b90e91d237ca..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 = _
@@ -56,7 +56,7 @@ class TransactionIndexTest extends JUnitSuite {
     assertEquals(abortedTxns ++ List(anotherAbortedTxn), reopenedIndex.allAbortedTxns)
   }
 
-  @Test(expected = classOf[IllegalArgumentException])
+  @Test(expected = classOf[CorruptIndexException])
   def testSanityCheck(): Unit = {
     val abortedTxns = List(
       new AbortedTxn(producerId = 0L, firstOffset = 0, lastOffset = 10, lastStableOffset = 11),
@@ -134,7 +134,7 @@ class TransactionIndexTest extends JUnitSuite {
     index.truncateTo(50)
     assertEquals(abortedTransactions.take(3), index.collectAbortedTxns(0L, 100L).abortedTransactions)
 
-    index.truncate()
+    index.reset()
     assertEquals(List.empty[AbortedTransaction], index.collectAbortedTxns(0L, 100L).abortedTransactions)
   }
 
diff --git a/core/src/test/scala/unit/kafka/message/BaseMessageSetTestCases.scala b/core/src/test/scala/unit/kafka/message/BaseMessageSetTestCases.scala
deleted file mode 100644
index 40581ed20686b..0000000000000
--- a/core/src/test/scala/unit/kafka/message/BaseMessageSetTestCases.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.message
-
-import java.nio.ByteBuffer
-import java.nio.channels.{FileChannel, GatheringByteChannel}
-import java.nio.file.StandardOpenOption
-
-import org.junit.Assert._
-import kafka.utils.TestUtils._
-import org.apache.kafka.common.record.FileRecords
-import org.scalatest.junit.JUnitSuite
-import org.junit.Test
-
-import scala.collection.mutable.ArrayBuffer
-import scala.collection.JavaConverters._
-
-trait BaseMessageSetTestCases extends JUnitSuite {
-
-  private class StubByteChannel(bytesToConsumePerBuffer: Int) extends GatheringByteChannel {
-
-    val data = new ArrayBuffer[Byte]
-
-    def write(srcs: Array[ByteBuffer], offset: Int, length: Int): Long = {
-      srcs.map { src =>
-        val array = new Array[Byte](math.min(bytesToConsumePerBuffer, src.remaining))
-        src.get(array)
-        data ++= array
-        array.length
-      }.sum
-    }
-
-    def write(srcs: Array[ByteBuffer]): Long = write(srcs, 0, srcs.map(_.remaining).sum)
-
-    def write(src: ByteBuffer): Int = write(Array(src)).toInt
-
-    def isOpen: Boolean = true
-
-    def close() {}
-
-  }
-
-
-  val messages = Array(new Message("abcd".getBytes), new Message("efgh".getBytes), new Message("ijkl".getBytes))
-  
-  def createMessageSet(messages: Seq[Message]): MessageSet
-
-  @Test
-  def testWrittenEqualsRead() {
-    val messageSet = createMessageSet(messages)
-    assertEquals(messages.toVector, messageSet.toVector.map(m => m.message))
-  }
-
-  @Test
-  def testIteratorIsConsistent() {
-    val m = createMessageSet(messages)
-    // two iterators over the same set should give the same results
-    checkEquals(m.iterator, m.iterator)
-  }
-
-  @Test
-  def testSizeInBytes() {
-    assertEquals("Empty message set should have 0 bytes.",
-                 0,
-                 createMessageSet(Array[Message]()).sizeInBytes)
-    assertEquals("Predicted size should equal actual size.", 
-                 MessageSet.messageSetSize(messages),
-                 createMessageSet(messages).sizeInBytes)
-  }
-
-  @Test
-  def testWriteTo() {
-    // test empty message set
-    checkWriteToWithMessageSet(createMessageSet(Array[Message]()))
-    checkWriteToWithMessageSet(createMessageSet(messages))
-  }
-
-  /* Tests that writing to a channel that doesn't consume all the bytes in the buffer works correctly */
-  @Test
-  def testWriteToChannelThatConsumesPartially() {
-    val bytesToConsumePerBuffer = 50
-    val messages = (0 until 10).map(_ => new Message(randomString(100).getBytes))
-    val messageSet = createMessageSet(messages)
-    val messageSetSize = messageSet.sizeInBytes
-
-    val channel = new StubByteChannel(bytesToConsumePerBuffer)
-
-    var remaining = messageSetSize
-    var iterations = 0
-    while (remaining > 0) {
-      remaining -= messageSet.asRecords.writeTo(channel, messageSetSize - remaining, remaining).toInt
-      iterations += 1
-    }
-
-    assertEquals((messageSetSize / bytesToConsumePerBuffer) + 1, iterations)
-    checkEquals(new ByteBufferMessageSet(ByteBuffer.wrap(channel.data.toArray)).iterator, messageSet.iterator)
-  }
-
-  def checkWriteToWithMessageSet(messageSet: MessageSet) {
-    checkWriteWithMessageSet(messageSet, messageSet.asRecords.writeTo(_, 0, messageSet.sizeInBytes))
-  }
-
-  def checkWriteWithMessageSet(set: MessageSet, write: GatheringByteChannel => Long) {
-    // do the write twice to ensure the message set is restored to its original state
-    for (_ <- 0 to 1) {
-      val file = tempFile()
-      val channel = FileChannel.open(file.toPath, StandardOpenOption.READ, StandardOpenOption.WRITE)
-      try {
-        val written = write(channel)
-        assertEquals("Expect to write the number of bytes in the set.", set.sizeInBytes, written)
-        val fileRecords = new FileRecords(file, channel, 0, Integer.MAX_VALUE, false)
-        assertEquals(set.asRecords.records.asScala.toVector, fileRecords.records.asScala.toVector)
-        checkEquals(set.asRecords.records.iterator, fileRecords.records.iterator)
-      } finally channel.close()
-    }
-  }
-  
-}
diff --git a/core/src/test/scala/unit/kafka/message/ByteBufferMessageSetTest.scala b/core/src/test/scala/unit/kafka/message/ByteBufferMessageSetTest.scala
deleted file mode 100644
index 00f9dc9f396ab..0000000000000
--- a/core/src/test/scala/unit/kafka/message/ByteBufferMessageSetTest.scala
+++ /dev/null
@@ -1,167 +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.message
-
-import java.nio._
-
-import kafka.utils.TestUtils
-import org.junit.Assert._
-import org.junit.Test
-
-class ByteBufferMessageSetTest extends BaseMessageSetTestCases {
-
-  override def createMessageSet(messages: Seq[Message]): ByteBufferMessageSet = 
-    new ByteBufferMessageSet(NoCompressionCodec, messages: _*)
-
-  @Test
-  def testValidBytes() {
-    {
-      val messages = new ByteBufferMessageSet(NoCompressionCodec, new Message("hello".getBytes), new Message("there".getBytes))
-      val buffer = ByteBuffer.allocate(messages.sizeInBytes + 2)
-      buffer.put(messages.buffer)
-      buffer.putShort(4)
-      val messagesPlus = new ByteBufferMessageSet(buffer)
-      assertEquals("Adding invalid bytes shouldn't change byte count", messages.validBytes, messagesPlus.validBytes)
-    }
-
-    // test valid bytes on empty ByteBufferMessageSet
-    {
-      assertEquals("Valid bytes on an empty ByteBufferMessageSet should return 0", 0,
-        MessageSet.Empty.asInstanceOf[ByteBufferMessageSet].validBytes)
-    }
-  }
-
-  @Test
-  def testValidBytesWithCompression() {
-    val messages = new ByteBufferMessageSet(DefaultCompressionCodec, new Message("hello".getBytes), new Message("there".getBytes))
-    val buffer = ByteBuffer.allocate(messages.sizeInBytes + 2)
-    buffer.put(messages.buffer)
-    buffer.putShort(4)
-    val messagesPlus = new ByteBufferMessageSet(buffer)
-    assertEquals("Adding invalid bytes shouldn't change byte count", messages.validBytes, messagesPlus.validBytes)
-  }
-
-  @Test
-  def testEquals() {
-    var messages = new ByteBufferMessageSet(DefaultCompressionCodec, new Message("hello".getBytes), new Message("there".getBytes))
-    var moreMessages = new ByteBufferMessageSet(DefaultCompressionCodec, new Message("hello".getBytes), new Message("there".getBytes))
-
-    assertTrue(messages.equals(moreMessages))
-
-    messages = new ByteBufferMessageSet(NoCompressionCodec, new Message("hello".getBytes), new Message("there".getBytes))
-    moreMessages = new ByteBufferMessageSet(NoCompressionCodec, new Message("hello".getBytes), new Message("there".getBytes))
-
-    assertTrue(messages.equals(moreMessages))
-  }
-  
-
-  @Test
-  def testIterator() {
-    val messageList = List(
-        new Message("msg1".getBytes),
-        new Message("msg2".getBytes),
-        new Message("msg3".getBytes)
-      )
-
-    // test for uncompressed regular messages
-    {
-      val messageSet = new ByteBufferMessageSet(NoCompressionCodec, messageList: _*)
-      TestUtils.checkEquals[Message](messageList.iterator, TestUtils.getMessageIterator(messageSet.iterator))
-      //make sure ByteBufferMessageSet is re-iterable.
-      TestUtils.checkEquals[Message](messageList.iterator, TestUtils.getMessageIterator(messageSet.iterator))
-
-      //make sure shallow iterator is the same as deep iterator
-      TestUtils.checkEquals[Message](TestUtils.getMessageIterator(messageSet.shallowIterator),
-                                     TestUtils.getMessageIterator(messageSet.iterator))
-    }
-
-    // test for compressed regular messages
-    {
-      val messageSet = new ByteBufferMessageSet(DefaultCompressionCodec, messageList: _*)
-      TestUtils.checkEquals[Message](messageList.iterator, TestUtils.getMessageIterator(messageSet.iterator))
-      //make sure ByteBufferMessageSet is re-iterable.
-      TestUtils.checkEquals[Message](messageList.iterator, TestUtils.getMessageIterator(messageSet.iterator))
-      verifyShallowIterator(messageSet)
-    }
-
-    // test for mixed empty and non-empty messagesets uncompressed
-    {
-      val emptyMessageList : List[Message] = Nil
-      val emptyMessageSet = new ByteBufferMessageSet(NoCompressionCodec, emptyMessageList: _*)
-      val regularMessgeSet = new ByteBufferMessageSet(NoCompressionCodec, messageList: _*)
-      val buffer = ByteBuffer.allocate(emptyMessageSet.buffer.limit() + regularMessgeSet.buffer.limit())
-      buffer.put(emptyMessageSet.buffer)
-      buffer.put(regularMessgeSet.buffer)
-      buffer.rewind
-      val mixedMessageSet = new ByteBufferMessageSet(buffer)
-      TestUtils.checkEquals[Message](messageList.iterator, TestUtils.getMessageIterator(mixedMessageSet.iterator))
-      //make sure ByteBufferMessageSet is re-iterable.
-      TestUtils.checkEquals[Message](messageList.iterator, TestUtils.getMessageIterator(mixedMessageSet.iterator))
-      //make sure shallow iterator is the same as deep iterator
-      TestUtils.checkEquals[Message](TestUtils.getMessageIterator(mixedMessageSet.shallowIterator),
-                                     TestUtils.getMessageIterator(mixedMessageSet.iterator))
-    }
-
-    // test for mixed empty and non-empty messagesets compressed
-    {
-      val emptyMessageList : List[Message] = Nil
-      val emptyMessageSet = new ByteBufferMessageSet(DefaultCompressionCodec, emptyMessageList: _*)
-      val regularMessgeSet = new ByteBufferMessageSet(DefaultCompressionCodec, messageList: _*)
-      val buffer = ByteBuffer.allocate(emptyMessageSet.buffer.limit() + regularMessgeSet.buffer.limit())
-      buffer.put(emptyMessageSet.buffer)
-      buffer.put(regularMessgeSet.buffer)
-      buffer.rewind
-      val mixedMessageSet = new ByteBufferMessageSet(buffer)
-      TestUtils.checkEquals[Message](messageList.iterator, TestUtils.getMessageIterator(mixedMessageSet.iterator))
-      //make sure ByteBufferMessageSet is re-iterable.
-      TestUtils.checkEquals[Message](messageList.iterator, TestUtils.getMessageIterator(mixedMessageSet.iterator))
-      verifyShallowIterator(mixedMessageSet)
-    }
-  }
-
-  @Test
-  def testMessageWithProvidedOffsetSeq() {
-    val offsets = Seq(0L, 2L)
-    val messages = new ByteBufferMessageSet(
-      compressionCodec = NoCompressionCodec,
-      offsetSeq = offsets,
-      new Message("hello".getBytes),
-      new Message("goodbye".getBytes))
-    val iter = messages.iterator
-    assertEquals("first offset should be 0", 0L, iter.next().offset)
-    assertEquals("second offset should be 2", 2L, iter.next().offset)
-  }
-
-  /* check that offsets are assigned based on byte offset from the given base offset */
-  def checkOffsets(messages: ByteBufferMessageSet, baseOffset: Long) {
-    assertTrue("Message set should not be empty", messages.nonEmpty)
-    var offset = baseOffset
-    for(entry <- messages) {
-      assertEquals("Unexpected offset in message set iterator", offset, entry.offset)
-      offset += 1
-    }
-  }
-
-  def verifyShallowIterator(messageSet: ByteBufferMessageSet) {
-    //make sure the offsets returned by a shallow iterator is a subset of that of a deep iterator
-    val shallowOffsets = messageSet.shallowIterator.map(msgAndOff => msgAndOff.offset).toSet
-    val deepOffsets = messageSet.iterator.map(msgAndOff => msgAndOff.offset).toSet
-    assertTrue(shallowOffsets.subsetOf(deepOffsets))
-  }
-
-}
diff --git a/core/src/test/scala/unit/kafka/message/MessageCompressionTest.scala b/core/src/test/scala/unit/kafka/message/MessageCompressionTest.scala
deleted file mode 100644
index ab8046531cf9e..0000000000000
--- a/core/src/test/scala/unit/kafka/message/MessageCompressionTest.scala
+++ /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 kafka.message
-
-import java.io.ByteArrayOutputStream
-import scala.collection._
-import org.junit._
-import org.junit.Assert._
-
-class MessageCompressionTest {
-
-  @Test
-  def testSimpleCompressDecompress() {
-    val codecs = mutable.ArrayBuffer[CompressionCodec](GZIPCompressionCodec)
-    if (isSnappyAvailable)
-      codecs += SnappyCompressionCodec
-    if (isLZ4Available)
-      codecs += LZ4CompressionCodec
-    for (codec <- codecs)
-      testSimpleCompressDecompress(codec)
-  }
-
-  //  A quick test to ensure any growth or increase in compression size is known when upgrading libraries
-  @Test
-  def testCompressSize() {
-    val bytes1k: Array[Byte] = (0 until 1000).map(_.toByte).toArray
-    val bytes2k: Array[Byte] = (1000 until 2000).map(_.toByte).toArray
-    val bytes3k: Array[Byte] = (3000 until 4000).map(_.toByte).toArray
-    val messages: List[Message] = List(new Message(bytes1k, Message.NoTimestamp, Message.MagicValue_V1),
-                                       new Message(bytes2k, Message.NoTimestamp, Message.MagicValue_V1),
-                                       new Message(bytes3k, Message.NoTimestamp, Message.MagicValue_V1))
-
-    testCompressSize(GZIPCompressionCodec, messages, 396)
-
-    if (isSnappyAvailable)
-      testCompressSize(SnappyCompressionCodec, messages, 503)
-
-    if (isLZ4Available)
-      testCompressSize(LZ4CompressionCodec, messages, 387)
-  }
-
-  def testSimpleCompressDecompress(compressionCodec: CompressionCodec) {
-    val messages = List[Message](new Message("hi there".getBytes), new Message("I am fine".getBytes), new Message("I am not so well today".getBytes))
-    val messageSet = new ByteBufferMessageSet(compressionCodec = compressionCodec, messages = messages:_*)
-    assertEquals(compressionCodec, messageSet.shallowIterator.next().message.compressionCodec)
-    val decompressed = messageSet.iterator.map(_.message).toList
-    assertEquals(messages, decompressed)
-  }
-
-  def testCompressSize(compressionCodec: CompressionCodec, messages: List[Message], expectedSize: Int) {
-    val messageSet = new ByteBufferMessageSet(compressionCodec = compressionCodec, messages = messages:_*)
-    assertEquals(s"$compressionCodec size has changed.", expectedSize, messageSet.sizeInBytes)
-  }
-
-  def isSnappyAvailable: Boolean = {
-    try {
-      new org.xerial.snappy.SnappyOutputStream(new ByteArrayOutputStream())
-      true
-    } catch {
-      case _: UnsatisfiedLinkError | _: org.xerial.snappy.SnappyError => false
-    }
-  }
-
-  def isLZ4Available: Boolean = {
-    try {
-      new net.jpountz.lz4.LZ4BlockOutputStream(new ByteArrayOutputStream())
-      true
-    } catch {
-      case _: UnsatisfiedLinkError => false
-    }
-  }
-}
diff --git a/core/src/test/scala/unit/kafka/message/MessageTest.scala b/core/src/test/scala/unit/kafka/message/MessageTest.scala
deleted file mode 100755
index 2390b5b2f425b..0000000000000
--- a/core/src/test/scala/unit/kafka/message/MessageTest.scala
+++ /dev/null
@@ -1,143 +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.message
-
-import java.nio._
-import java.util.HashMap
-
-import org.apache.kafka.common.protocol.Errors
-
-import scala.collection._
-import org.junit.Assert._
-import org.scalatest.junit.JUnitSuite
-import org.junit.{Before, Test}
-import kafka.utils.TestUtils
-import org.apache.kafka.common.utils.ByteUtils
-
-case class MessageTestVal(key: Array[Byte],
-                          payload: Array[Byte],
-                          codec: CompressionCodec,
-                          timestamp: Long,
-                          magicValue: Byte,
-                          message: Message)
-
-class MessageTest extends JUnitSuite {
-
-  var messages = new mutable.ArrayBuffer[MessageTestVal]()
-
-  @Before
-  def setUp(): Unit = {
-    val keys = Array(null, "key".getBytes, "".getBytes)
-    val vals = Array("value".getBytes, "".getBytes, null)
-    val codecs = Array(NoCompressionCodec, GZIPCompressionCodec, SnappyCompressionCodec, LZ4CompressionCodec)
-    val timestamps = Array(Message.NoTimestamp, 0L, 1L)
-    val magicValues = Array(Message.MagicValue_V0, Message.MagicValue_V1)
-    for(k <- keys; v <- vals; codec <- codecs; t <- timestamps; mv <- magicValues) {
-      val timestamp = ensureValid(mv, t)
-      messages += MessageTestVal(k, v, codec, timestamp, mv, new Message(v, k, timestamp, codec, mv))
-    }
-
-    def ensureValid(magicValue: Byte, timestamp: Long): Long =
-      if (magicValue > Message.MagicValue_V0) timestamp else Message.NoTimestamp
-  }
-
-  @Test
-  def testFieldValues(): Unit = {
-    for(v <- messages) {
-      // check payload
-      if(v.payload == null) {
-        assertTrue(v.message.isNull)
-        assertEquals("Payload should be null", null, v.message.payload)
-      } else {
-        TestUtils.checkEquals(ByteBuffer.wrap(v.payload), v.message.payload)
-      }
-      // check timestamp
-      if (v.magicValue > Message.MagicValue_V0)
-        assertEquals("Timestamp should be the same", v.timestamp, v.message.timestamp)
-       else
-        assertEquals("Timestamp should be the NoTimestamp", Message.NoTimestamp, v.message.timestamp)
-
-      // check magic value
-      assertEquals(v.magicValue, v.message.magic)
-      // check key
-      if(v.message.hasKey)
-        TestUtils.checkEquals(ByteBuffer.wrap(v.key), v.message.key)
-      else
-        assertEquals(null, v.message.key)
-      // check compression codec
-      assertEquals(v.codec, v.message.compressionCodec)
-    }
-  }
-
-  @Test
-  def testChecksum() {
-    for(v <- messages) {
-      assertTrue("Auto-computed checksum should be valid", v.message.isValid)
-      // garble checksum
-      val badChecksum: Int = (v.message.checksum + 1 % Int.MaxValue).toInt
-      ByteUtils.writeUnsignedInt(v.message.buffer, Message.CrcOffset, badChecksum)
-      assertFalse("Message with invalid checksum should be invalid", v.message.isValid)
-    }
-  }
-
-  @Test
-  def testEquality() {
-    for (v <- messages) {
-      assertFalse("Should not equal null", v.message.equals(null))
-      assertFalse("Should not equal a random string", v.message.equals("asdf"))
-      assertTrue("Should equal itself", v.message.equals(v.message))
-      val copy = new Message(bytes = v.payload, key = v.key, v.timestamp, codec = v.codec, v.magicValue)
-      assertTrue("Should equal another message with the same content.", v.message.equals(copy))
-    }
-  }
-
-  @Test(expected = classOf[IllegalArgumentException])
-  def testInvalidTimestampAndMagicValueCombination() {
-      new Message("hello".getBytes, 0L, Message.MagicValue_V0)
-  }
-
-  @Test(expected = classOf[IllegalArgumentException])
-  def testInvalidTimestamp() {
-    new Message("hello".getBytes, -3L, Message.MagicValue_V1)
-  }
-
-  @Test(expected = classOf[IllegalArgumentException])
-  def testInvalidMagicByte() {
-    new Message("hello".getBytes, 0L, 2.toByte)
-  }
-
-  @Test
-  def testIsHashable() {
-    // this is silly, but why not
-    val m = new HashMap[Message, Message]()
-    for(v <- messages)
-      m.put(v.message, v.message)
-    for(v <- messages)
-      assertEquals(v.message, m.get(v.message))
-  }
-
-  @Test
-  def testExceptionMapping() {
-    val expected = Errors.CORRUPT_MESSAGE
-    val actual = Errors.forException(new InvalidMessageException())
-
-    assertEquals("InvalidMessageException should map to a corrupt message error", expected, actual)
-  }
-
-}
-
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 c250d850129fb..bac15de13b6dd 100644
--- a/core/src/test/scala/unit/kafka/metrics/MetricsTest.scala
+++ b/core/src/test/scala/unit/kafka/metrics/MetricsTest.scala
@@ -17,25 +17,22 @@
 
 package kafka.metrics
 
+import java.lang.management.ManagementFactory
 import java.util.Properties
-import javax.management.ObjectName
 
-import com.yammer.metrics.Metrics
-import com.yammer.metrics.core.{Meter, MetricPredicate}
+import javax.management.ObjectName
+import com.yammer.metrics.core.MetricPredicate
 import org.junit.Test
 import org.junit.Assert._
 import kafka.integration.KafkaServerTestHarness
 import kafka.server._
-import kafka.serializer._
 import kafka.utils._
-import kafka.utils.TestUtils._
 
 import scala.collection._
-import scala.collection.JavaConverters._
-import scala.util.matching.Regex
-import kafka.consumer.{ConsumerConfig, ZookeeperConsumerConnector}
+import scala.jdk.CollectionConverters._
 import kafka.log.LogConfig
 import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.metrics.JmxReporter
 
 class MetricsTest extends KafkaServerTestHarness with Logging {
   val numNodes = 2
@@ -43,62 +40,78 @@ class MetricsTest extends KafkaServerTestHarness with Logging {
 
   val overridingProps = new Properties
   overridingProps.put(KafkaConfig.NumPartitionsProp, numParts.toString)
+  overridingProps.put(JmxReporter.EXCLUDE_CONFIG, "kafka.server:type=KafkaServer,name=ClusterId")
 
   def generateConfigs =
-    TestUtils.createBrokerConfigs(numNodes, zkConnect, enableDeleteTopic=true).map(KafkaConfig.fromProps(_, overridingProps))
+    TestUtils.createBrokerConfigs(numNodes, zkConnect).map(KafkaConfig.fromProps(_, overridingProps))
 
   val nMessages = 2
 
   @Test
-  @deprecated("This test has been deprecated and it will be removed in a future release", "0.10.0.0")
-  def testMetricsLeak() {
-    val topic = "test-metrics-leak"
-    // create topic topic1 with 1 partition on broker 0
-    createTopic(zkUtils, topic, numPartitions = 1, replicationFactor = 1, servers = servers)
-    // force creation not client's specific metrics.
-    createAndShutdownStep(topic, "group0", "consumer0", "producer0")
-
-    //this assertion is only used for creating the metrics for DelayedFetchMetrics, it should never fail, but should not be removed
-    assertNotNull(DelayedFetchMetrics)
-
-    val countOfStaticMetrics = Metrics.defaultRegistry.allMetrics.keySet.size
-
-    for (i <- 0 to 5) {
-      createAndShutdownStep(topic, "group" + i % 3, "consumer" + i % 2, "producer" + i % 2)
-      assertEquals(countOfStaticMetrics, Metrics.defaultRegistry.allMetrics.keySet.size)
-    }
-  }
-
-  @Test
-  def testMetricsReporterAfterDeletingTopic() {
+  def testMetricsReporterAfterDeletingTopic(): Unit = {
     val topic = "test-topic-metric"
-    adminZkClient.createTopic(topic, 1, 1)
+    createTopic(topic, 1, 1)
     adminZkClient.deleteTopic(topic)
-    TestUtils.verifyTopicDeletion(zkUtils, topic, 1, servers)
+    TestUtils.verifyTopicDeletion(zkClient, topic, 1, servers)
     assertEquals("Topic metrics exists after deleteTopic", Set.empty, topicMetricGroups(topic))
   }
 
   @Test
-  def testBrokerTopicMetricsUnregisteredAfterDeletingTopic() {
+  def testBrokerTopicMetricsUnregisteredAfterDeletingTopic(): Unit = {
     val topic = "test-broker-topic-metric"
-    adminZkClient.createTopic(topic, 2, 1)
+    createTopic(topic, 2, 1)
     // Produce a few messages to create the metrics
     // Don't consume messages as it may cause metrics to be re-created causing the test to fail, see KAFKA-5238
-    TestUtils.produceMessages(servers, topic, nMessages)
+    TestUtils.generateAndProduceMessages(servers, topic, nMessages)
     assertTrue("Topic metrics don't exist", topicMetricGroups(topic).nonEmpty)
     servers.foreach(s => assertNotNull(s.brokerTopicStats.topicStats(topic)))
     adminZkClient.deleteTopic(topic)
-    TestUtils.verifyTopicDeletion(zkUtils, topic, 1, servers)
+    TestUtils.verifyTopicDeletion(zkClient, topic, 1, servers)
     assertEquals("Topic metrics exists after deleteTopic", Set.empty, topicMetricGroups(topic))
   }
 
   @Test
   def testClusterIdMetric(): Unit = {
     // Check if clusterId metric exists.
-    val metrics = Metrics.defaultRegistry.allMetrics
+    val metrics = KafkaYammerMetrics.defaultRegistry.allMetrics
     assertEquals(metrics.keySet.asScala.count(_.getMBeanName == "kafka.server:type=KafkaServer,name=ClusterId"), 1)
   }
 
+  @Test
+  def testJMXFilter(): Unit = {
+    // Check if cluster id metrics is not exposed in JMX
+    assertTrue(ManagementFactory.getPlatformMBeanServer
+                 .isRegistered(new ObjectName("kafka.controller:type=KafkaController,name=ActiveControllerCount")))
+    assertFalse(ManagementFactory.getPlatformMBeanServer
+                  .isRegistered(new ObjectName("kafka.server:type=KafkaServer,name=ClusterId")))
+  }
+
+  @Test
+  def testUpdateJMXFilter(): Unit = {
+    // verify previously exposed metrics are removed and existing matching metrics are added
+    servers.foreach(server => server.kafkaYammerMetrics.reconfigure(
+      Map(JmxReporter.EXCLUDE_CONFIG -> "kafka.controller:type=KafkaController,name=ActiveControllerCount").asJava
+    ))
+    assertFalse(ManagementFactory.getPlatformMBeanServer
+                 .isRegistered(new ObjectName("kafka.controller:type=KafkaController,name=ActiveControllerCount")))
+    assertTrue(ManagementFactory.getPlatformMBeanServer
+                  .isRegistered(new ObjectName("kafka.server:type=KafkaServer,name=ClusterId")))
+  }
+
+  @Test
+  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, 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"
@@ -108,18 +121,6 @@ class MetricsTest extends KafkaServerTestHarness with Logging {
     assert(metric.getMBeanName.endsWith(expectedMBeanName))
   }
 
-  @deprecated("This test has been deprecated and it will be removed in a future release", "0.10.0.0")
-  def createAndShutdownStep(topic: String, group: String, consumerId: String, producerId: String): Unit = {
-    sendMessages(servers, topic, nMessages)
-    // create a consumer
-    val consumerConfig1 = new ConsumerConfig(TestUtils.createConsumerProperties(zkConnect, group, consumerId))
-    val zkConsumerConnector1 = new ZookeeperConsumerConnector(consumerConfig1, true)
-    val topicMessageStreams1 = zkConsumerConnector1.createMessageStreams(Map(topic -> 1), new StringDecoder, new StringDecoder)
-    getMessages(topicMessageStreams1, nMessages)
-
-    zkConsumerConnector1.shutdown()
-  }
-
   @Test
   def testBrokerTopicMetricsBytesInOut(): Unit = {
     val topic = "test-bytes-in-out"
@@ -130,9 +131,9 @@ class MetricsTest extends KafkaServerTestHarness with Logging {
 
     val topicConfig = new Properties
     topicConfig.setProperty(LogConfig.MinInSyncReplicasProp, "2")
-    createTopic(zkUtils, topic, 1, numNodes, servers, topicConfig)
+    createTopic(topic, 1, numNodes, topicConfig)
     // Produce a few messages to create the metrics
-    TestUtils.produceMessages(servers, topic, nMessages)
+    TestUtils.generateAndProduceMessages(servers, topic, nMessages)
 
     // Check the log size for each broker so that we can distinguish between failures caused by replication issues
     // versus failures caused by the metrics
@@ -145,50 +146,69 @@ 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.produceMessages(servers, topic, nMessages)
+    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
   def testControllerMetrics(): Unit = {
-    val metrics = Metrics.defaultRegistry.allMetrics
+    val metrics = KafkaYammerMetrics.defaultRegistry.allMetrics
 
     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)
   }
 
-  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
+  /**
+   * Test that the metrics are created with the right name, testZooKeeperStateChangeRateMetrics
+   * and testZooKeeperSessionStateMetric in ZooKeeperClientTest test the metrics behaviour.
+   */
+  @Test
+  def testSessionExpireListenerMetrics(): Unit = {
+    val metrics = KafkaYammerMetrics.defaultRegistry.allMetrics
+
+    assertEquals(metrics.keySet.asScala.count(_.getMBeanName == "kafka.server:type=SessionExpireListener,name=SessionState"), 1)
+    assertEquals(metrics.keySet.asScala.count(_.getMBeanName == "kafka.server:type=SessionExpireListener,name=ZooKeeperExpiresPerSec"), 1)
+    assertEquals(metrics.keySet.asScala.count(_.getMBeanName == "kafka.server:type=SessionExpireListener,name=ZooKeeperDisconnectsPerSec"), 1)
+  }
+
+  private def topicMetrics(topic: Option[String]): Set[String] = {
+    val metricNames = KafkaYammerMetrics.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)
+    val metricGroups = KafkaYammerMetrics.defaultRegistry.groupedMetrics(MetricPredicate.ALL).keySet.asScala
+    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/ConnectionQuotasTest.scala b/core/src/test/scala/unit/kafka/network/ConnectionQuotasTest.scala
new file mode 100644
index 0000000000000..07283db07617c
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/network/ConnectionQuotasTest.scala
@@ -0,0 +1,943 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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.network
+
+import java.net.InetAddress
+import java.util
+import java.util.concurrent.{Callable, ExecutorService, Executors, TimeUnit}
+import java.util.{Collections, Properties}
+
+import com.yammer.metrics.core.Meter
+import kafka.metrics.KafkaMetricsGroup
+import kafka.network.Processor.ListenerMetricTag
+import kafka.server.{DynamicConfig, KafkaConfig}
+import kafka.utils.{MockTime, TestUtils}
+import org.apache.kafka.common.config.ConfigException
+import org.apache.kafka.common.metrics.internals.MetricsUtils
+import org.apache.kafka.common.metrics.{KafkaMetric, MetricConfig, Metrics}
+import org.apache.kafka.common.network._
+import org.apache.kafka.common.utils.Time
+import org.junit.Assert._
+import org.junit._
+import org.scalatest.Assertions.{assertThrows, intercept}
+
+import scala.jdk.CollectionConverters._
+import scala.collection.{Map, mutable}
+import scala.concurrent.TimeoutException
+
+class ConnectionQuotasTest {
+  private var metrics: Metrics = _
+  private var executor: ExecutorService = _
+  private var connectionQuotas: ConnectionQuotas = _
+  private var time: Time = _
+
+  private val listeners = Map(
+    "EXTERNAL" -> ListenerDesc(new ListenerName("EXTERNAL"), InetAddress.getByName("192.168.1.1")),
+    "ADMIN" -> ListenerDesc(new ListenerName("ADMIN"), InetAddress.getByName("192.168.1.2")),
+    "REPLICATION" -> ListenerDesc(new ListenerName("REPLICATION"), InetAddress.getByName("192.168.1.3")))
+  private val blockedPercentMeters = mutable.Map[String, Meter]()
+  private val knownHost = InetAddress.getByName("192.168.10.0")
+  private val unknownHost = InetAddress.getByName("192.168.2.0")
+
+  private val numQuotaSamples = 2
+  private val quotaWindowSizeSeconds = 1
+  private val eps = 0.01
+
+  case class ListenerDesc(listenerName: ListenerName, defaultIp: InetAddress) {
+    override def toString: String = {
+      s"(listener=${listenerName.value}, client=${defaultIp.getHostAddress})"
+    }
+  }
+
+  def brokerPropsWithDefaultConnectionLimits: Properties = {
+    val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 0)
+    props.put(KafkaConfig.ListenersProp, "EXTERNAL://localhost:0,REPLICATION://localhost:1,ADMIN://localhost:2")
+    // ConnectionQuotas does not limit inter-broker listener even when broker-wide connection limit is reached
+    props.put(KafkaConfig.InterBrokerListenerNameProp, "REPLICATION")
+    props.put(KafkaConfig.ListenerSecurityProtocolMapProp, "EXTERNAL:PLAINTEXT,REPLICATION:PLAINTEXT,ADMIN:PLAINTEXT")
+    props.put(KafkaConfig.NumQuotaSamplesProp, numQuotaSamples.toString)
+    props.put(KafkaConfig.QuotaWindowSizeSecondsProp, quotaWindowSizeSeconds.toString)
+    props
+  }
+
+  private def setupMockTime(): Unit = {
+    // clean up metrics initialized with Time.SYSTEM
+    metrics.close()
+    time = new MockTime()
+    metrics = new Metrics(time)
+  }
+
+  @Before
+  def setUp(): Unit = {
+    // Clean-up any metrics left around by previous tests
+    TestUtils.clearYammerMetrics()
+
+    listeners.keys.foreach { name =>
+        blockedPercentMeters.put(name, KafkaMetricsGroup.newMeter(
+          s"${name}BlockedPercent", "blocked time", TimeUnit.NANOSECONDS, Map(ListenerMetricTag -> name)))
+    }
+    // use system time, because ConnectionQuota causes the current thread to wait with timeout, which waits based on
+    // system time; so using mock time will likely result in test flakiness due to a mixed use of mock and system time
+    time = Time.SYSTEM
+    metrics = new Metrics(new MetricConfig(), Collections.emptyList(), time)
+    executor = Executors.newFixedThreadPool(listeners.size)
+  }
+
+  @After
+  def tearDown(): Unit = {
+    executor.shutdownNow()
+    if (connectionQuotas != null) {
+      connectionQuotas.close()
+    }
+    metrics.close()
+    TestUtils.clearYammerMetrics()
+    blockedPercentMeters.clear()
+  }
+
+  @Test
+  def testFailWhenNoListeners(): Unit = {
+    val config = KafkaConfig.fromProps(brokerPropsWithDefaultConnectionLimits)
+    connectionQuotas = new ConnectionQuotas(config, time, metrics)
+
+    // inc() on a separate thread in case it blocks
+    val listener = listeners("EXTERNAL")
+    executor.submit((() =>
+      intercept[RuntimeException](
+        connectionQuotas.inc(listener.listenerName, listener.defaultIp, blockedPercentMeters("EXTERNAL"))
+      )): Runnable
+    ).get(5, TimeUnit.SECONDS)
+  }
+
+  @Test
+  def testFailDecrementForUnknownIp(): Unit = {
+    val config = KafkaConfig.fromProps(brokerPropsWithDefaultConnectionLimits)
+    connectionQuotas = new ConnectionQuotas(config, time, metrics)
+    addListenersAndVerify(config, connectionQuotas)
+
+    // calling dec() for an IP for which we didn't call inc() should throw an exception
+    intercept[IllegalArgumentException](connectionQuotas.dec(listeners("EXTERNAL").listenerName, unknownHost))
+  }
+
+  @Test
+  def testNoConnectionLimitsByDefault(): Unit = {
+    val config = KafkaConfig.fromProps(brokerPropsWithDefaultConnectionLimits)
+    connectionQuotas = new ConnectionQuotas(config, time, metrics)
+    addListenersAndVerify(config, connectionQuotas)
+
+    // verify there is no limit by accepting 10000 connections as fast as possible
+    val numConnections = 10000
+    val futures = listeners.values.map { listener =>
+      executor.submit((() => acceptConnections(connectionQuotas, listener, numConnections)): Runnable)
+    }
+    futures.foreach(_.get(10, TimeUnit.SECONDS))
+    assertTrue("Expected broker-connection-accept-rate metric to get recorded",
+      metricValue(brokerConnRateMetric())> 0)
+    listeners.values.foreach { listener =>
+      assertEquals(s"Number of connections on $listener:",
+        numConnections, connectionQuotas.get(listener.defaultIp))
+
+      assertTrue(s"Expected connection-accept-rate metric to get recorded for listener $listener",
+        metricValue(listenerConnRateMetric(listener.listenerName.value)) > 0)
+
+      // verify removing one connection
+      connectionQuotas.dec(listener.listenerName, listener.defaultIp)
+      assertEquals(s"Number of connections on $listener:",
+        numConnections - 1, connectionQuotas.get(listener.defaultIp))
+    }
+    // the blocked percent should still be 0, because no limits were reached
+    verifyNoBlockedPercentRecordedOnAllListeners()
+  }
+
+  @Test
+  def testMaxConnectionsPerIp(): Unit = {
+    val maxConnectionsPerIp = 17
+    val props = brokerPropsWithDefaultConnectionLimits
+    props.put(KafkaConfig.MaxConnectionsPerIpProp, maxConnectionsPerIp.toString)
+    val config = KafkaConfig.fromProps(props)
+    connectionQuotas = new ConnectionQuotas(config, time, metrics)
+
+    addListenersAndVerify(config, connectionQuotas)
+
+    val externalListener = listeners("EXTERNAL")
+    executor.submit((() =>
+      acceptConnections(connectionQuotas, externalListener, maxConnectionsPerIp)): Runnable
+    ).get(5, TimeUnit.SECONDS)
+    assertEquals(s"Number of connections on $externalListener:",
+      maxConnectionsPerIp, connectionQuotas.get(externalListener.defaultIp))
+
+    // all subsequent connections will be added to the counters, but inc() will throw TooManyConnectionsException for each
+    executor.submit((() =>
+      acceptConnectionsAboveIpLimit(connectionQuotas, externalListener, 2)): Runnable
+    ).get(5, TimeUnit.SECONDS)
+    assertEquals(s"Number of connections on $externalListener:",
+      maxConnectionsPerIp + 2, connectionQuotas.get(externalListener.defaultIp))
+
+    // connections on the same listener but from a different IP should be accepted
+    executor.submit((() =>
+      acceptConnections(connectionQuotas, externalListener.listenerName, knownHost, maxConnectionsPerIp,
+        0, expectIpThrottle = false)): Runnable
+    ).get(5, TimeUnit.SECONDS)
+
+    // remove two "rejected" connections and remove 2 more connections to free up the space for another 2 connections
+    for (_ <- 0 until 4) connectionQuotas.dec(externalListener.listenerName, externalListener.defaultIp)
+    assertEquals(s"Number of connections on $externalListener:",
+      maxConnectionsPerIp - 2, connectionQuotas.get(externalListener.defaultIp))
+
+    executor.submit((() =>
+      acceptConnections(connectionQuotas, externalListener, 2)): Runnable
+    ).get(5, TimeUnit.SECONDS)
+    assertEquals(s"Number of connections on $externalListener:",
+      maxConnectionsPerIp, connectionQuotas.get(externalListener.defaultIp))
+  }
+
+  @Test
+  def testMaxBrokerWideConnectionLimit(): Unit = {
+    val maxConnections = 800
+    val props = brokerPropsWithDefaultConnectionLimits
+    props.put(KafkaConfig.MaxConnectionsProp, maxConnections.toString)
+    val config = KafkaConfig.fromProps(props)
+    connectionQuotas = new ConnectionQuotas(config, time, metrics)
+
+    addListenersAndVerify(config, connectionQuotas)
+
+    // verify that ConnectionQuota can give all connections to one listener
+    executor.submit((() =>
+      acceptConnections(connectionQuotas, listeners("EXTERNAL"), maxConnections)): Runnable
+    ).get(5, TimeUnit.SECONDS)
+    assertEquals(s"Number of connections on ${listeners("EXTERNAL")}:",
+      maxConnections, connectionQuotas.get(listeners("EXTERNAL").defaultIp))
+
+    // the blocked percent should still be 0, because there should be no wait for a connection slot
+    assertEquals(0, blockedPercentMeters("EXTERNAL").count())
+
+    // the number of connections should be above max for maxConnectionsExceeded to return true
+    assertFalse("Total number of connections is exactly the maximum.",
+      connectionQuotas.maxConnectionsExceeded(listeners("EXTERNAL").listenerName))
+
+    // adding one more connection will block ConnectionQuota.inc()
+    val future = executor.submit((() =>
+      acceptConnections(connectionQuotas, listeners("EXTERNAL"), 1)): Runnable
+    )
+    intercept[TimeoutException](future.get(100, TimeUnit.MILLISECONDS))
+
+    // removing one connection should make the waiting connection to succeed
+    connectionQuotas.dec(listeners("EXTERNAL").listenerName, listeners("EXTERNAL").defaultIp)
+    future.get(1, TimeUnit.SECONDS)
+    assertEquals(s"Number of connections on ${listeners("EXTERNAL")}:",
+      maxConnections, connectionQuotas.get(listeners("EXTERNAL").defaultIp))
+    // metric is recorded in nanoseconds
+    assertTrue("Expected BlockedPercentMeter metric to be recorded",
+      blockedPercentMeters("EXTERNAL").count() > 0)
+
+    // adding inter-broker connections should succeed even when the total number of connections reached the max
+    executor.submit((() =>
+      acceptConnections(connectionQuotas, listeners("REPLICATION"), 1)): Runnable
+    ).get(5, TimeUnit.SECONDS)
+    assertTrue("Expected the number of connections to exceed the maximum.",
+      connectionQuotas.maxConnectionsExceeded(listeners("EXTERNAL").listenerName))
+
+    // adding one more connection on another non-inter-broker will block ConnectionQuota.inc()
+    val future1 = executor.submit((() =>
+      acceptConnections(connectionQuotas, listeners("ADMIN"), 1)): Runnable
+    )
+    intercept[TimeoutException](future1.get(1, TimeUnit.SECONDS))
+
+    // adding inter-broker connection should still succeed, even though a connection from another listener is waiting
+    executor.submit((() =>
+      acceptConnections(connectionQuotas, listeners("REPLICATION"), 1)): Runnable
+    ).get(5, TimeUnit.SECONDS)
+
+    // at this point, we need to remove 3 connections for the waiting connection to succeed
+    // remove 2 first -- should not be enough to accept the waiting connection
+    for (_ <- 0 until 2) connectionQuotas.dec(listeners("EXTERNAL").listenerName, listeners("EXTERNAL").defaultIp)
+    intercept[TimeoutException](future1.get(100, TimeUnit.MILLISECONDS))
+    connectionQuotas.dec(listeners("EXTERNAL").listenerName, listeners("EXTERNAL").defaultIp)
+    future1.get(1, TimeUnit.SECONDS)
+  }
+
+  @Test
+  def testMaxListenerConnectionLimits(): Unit = {
+    val maxConnections = 800
+    // sum of per-listener connection limits is below total connection limit
+    val listenerMaxConnections = 200
+    val props = brokerPropsWithDefaultConnectionLimits
+    props.put(KafkaConfig.MaxConnectionsProp, maxConnections.toString)
+    val config = KafkaConfig.fromProps(props)
+    connectionQuotas = new ConnectionQuotas(config, time, metrics)
+
+    addListenersAndVerify(config, connectionQuotas)
+
+    val listenerConfig = Map(KafkaConfig.MaxConnectionsProp -> listenerMaxConnections.toString).asJava
+    listeners.values.foreach { listener =>
+      connectionQuotas.maxConnectionsPerListener(listener.listenerName).configure(listenerConfig)
+    }
+
+    // verify each listener can create up to max connections configured for that listener
+    val futures = listeners.values.map { listener =>
+      executor.submit((() => acceptConnections(connectionQuotas, listener, listenerMaxConnections)): Runnable)
+    }
+    futures.foreach(_.get(5, TimeUnit.SECONDS))
+    listeners.values.foreach { listener =>
+      assertEquals(s"Number of connections on $listener:",
+        listenerMaxConnections, connectionQuotas.get(listener.defaultIp))
+      assertFalse(s"Total number of connections on $listener should be exactly the maximum.",
+        connectionQuotas.maxConnectionsExceeded(listener.listenerName))
+    }
+
+    // since every listener has exactly the max number of listener connections,
+    // every listener should block on the next connection creation, even the inter-broker listener
+    val overLimitFutures = listeners.values.map { listener =>
+      executor.submit((() => acceptConnections(connectionQuotas, listener, 1)): Runnable)
+    }
+    overLimitFutures.foreach { future =>
+      intercept[TimeoutException](future.get(1, TimeUnit.SECONDS))
+    }
+    listeners.values.foreach { listener =>
+      // free up one connection slot
+      connectionQuotas.dec(listener.listenerName, listener.defaultIp)
+    }
+    // all connections should get added
+    overLimitFutures.foreach(_.get(5, TimeUnit.SECONDS))
+    verifyConnectionCountOnEveryListener(connectionQuotas, listenerMaxConnections)
+  }
+
+  @Test
+  def testBrokerConnectionRateLimitWhenActualRateBelowLimit(): Unit = {
+    val brokerRateLimit = 125
+    // create connections with the total rate < broker-wide quota, and verify there is no throttling
+    val connCreateIntervalMs = 25 // connection creation rate = 40/sec per listener (3 * 40 = 120/sec total)
+    val connectionsPerListener = 200 // should take 5 seconds to create 200 connections with rate = 40/sec
+    val props = brokerPropsWithDefaultConnectionLimits
+    props.put(KafkaConfig.MaxConnectionCreationRateProp, brokerRateLimit.toString)
+    val config = KafkaConfig.fromProps(props)
+    connectionQuotas = new ConnectionQuotas(config, time, metrics)
+
+    addListenersAndVerify(config, connectionQuotas)
+
+    val futures = listeners.values.map { listener =>
+      executor.submit((() => acceptConnections(connectionQuotas, listener, connectionsPerListener, connCreateIntervalMs)): Runnable)
+    }
+    futures.foreach(_.get(10, TimeUnit.SECONDS))
+
+    // the blocked percent should still be 0, because no limits were reached
+    verifyNoBlockedPercentRecordedOnAllListeners()
+    verifyConnectionCountOnEveryListener(connectionQuotas, connectionsPerListener)
+  }
+
+  @Test
+  def testBrokerConnectionRateLimitWhenActualRateAboveLimit(): Unit = {
+    val brokerRateLimit = 90
+    val props = brokerPropsWithDefaultConnectionLimits
+    props.put(KafkaConfig.MaxConnectionCreationRateProp, brokerRateLimit.toString)
+    val config = KafkaConfig.fromProps(props)
+    connectionQuotas = new ConnectionQuotas(config, time, metrics)
+
+    addListenersAndVerify(config, connectionQuotas)
+
+    // each listener creates connections such that the total connection rate > broker-wide quota
+    val connCreateIntervalMs = 10      // connection creation rate = 100
+    val connectionsPerListener = 400
+    val futures = listeners.values.map { listener =>
+      executor.submit((() => acceptConnections(connectionQuotas, listener, connectionsPerListener, connCreateIntervalMs)): Runnable)
+    }
+    futures.foreach(_.get(20, TimeUnit.SECONDS))
+
+    // verify that connections on non-inter-broker listener are throttled
+    verifyOnlyNonInterBrokerListenersBlockedPercentRecorded()
+
+    // expect all connections to be created (no limit on the number of connections)
+    verifyConnectionCountOnEveryListener(connectionQuotas, connectionsPerListener)
+  }
+
+  @Test
+  def testListenerConnectionRateLimitWhenActualRateBelowLimit(): Unit = {
+    val brokerRateLimit = 125
+    val listenerRateLimit = 50
+    val connCreateIntervalMs = 25 // connection creation rate = 40/sec per listener (3 * 40 = 120/sec total)
+    val props = brokerPropsWithDefaultConnectionLimits
+    props.put(KafkaConfig.MaxConnectionCreationRateProp, brokerRateLimit.toString)
+    val config = KafkaConfig.fromProps(props)
+    connectionQuotas = new ConnectionQuotas(config, time, metrics)
+
+    val listenerConfig = Map(KafkaConfig.MaxConnectionCreationRateProp -> listenerRateLimit.toString).asJava
+    addListenersAndVerify(config, listenerConfig, connectionQuotas)
+
+    // create connections with the rate < listener quota on every listener, and verify there is no throttling
+    val connectionsPerListener = 200 // should take 5 seconds to create 200 connections with rate = 40/sec
+    val futures = listeners.values.map { listener =>
+      executor.submit((() => acceptConnections(connectionQuotas, listener, connectionsPerListener, connCreateIntervalMs)): Runnable)
+    }
+    futures.foreach(_.get(10, TimeUnit.SECONDS))
+
+    // the blocked percent should still be 0, because no limits were reached
+    verifyNoBlockedPercentRecordedOnAllListeners()
+
+    verifyConnectionCountOnEveryListener(connectionQuotas, connectionsPerListener)
+  }
+
+  @Test
+  def testListenerConnectionRateLimitWhenActualRateAboveLimit(): Unit = {
+    val brokerRateLimit = 125
+    val listenerRateLimit = 30
+    val connCreateIntervalMs = 25 // connection creation rate = 40/sec per listener (3 * 40 = 120/sec total)
+    val props = brokerPropsWithDefaultConnectionLimits
+    props.put(KafkaConfig.MaxConnectionCreationRateProp, brokerRateLimit.toString)
+    val config = KafkaConfig.fromProps(props)
+    connectionQuotas = new ConnectionQuotas(config, time, metrics)
+
+    val listenerConfig = Map(KafkaConfig.MaxConnectionCreationRateProp -> listenerRateLimit.toString).asJava
+    addListenersAndVerify(config, listenerConfig, connectionQuotas)
+
+    // create connections with the rate > listener quota on every listener
+    // run a bit longer (20 seconds) to also verify the throttle rate
+    val connectionsPerListener = 600 // should take 20 seconds to create 600 connections with rate = 30/sec
+    val futures = listeners.values.map { listener =>
+      executor.submit((() =>
+        // epsilon is set to account for the worst-case where the measurement is taken just before or after the quota window
+        acceptConnectionsAndVerifyRate(connectionQuotas, listener, connectionsPerListener, connCreateIntervalMs, listenerRateLimit, 7)): Runnable)
+    }
+    futures.foreach(_.get(30, TimeUnit.SECONDS))
+
+    // verify that every listener was throttled
+    verifyNonZeroBlockedPercentAndThrottleTimeOnAllListeners()
+
+    // while the connection creation rate was throttled,
+    // expect all connections got created (not limit on the number of connections)
+    verifyConnectionCountOnEveryListener(connectionQuotas, connectionsPerListener)
+  }
+
+  @Test
+  def testIpConnectionRateWhenActualRateBelowLimit(): Unit = {
+    val ipConnectionRateLimit = 30
+    val connCreateIntervalMs = 40 // connection creation rate = 25/sec
+    val props = brokerPropsWithDefaultConnectionLimits
+    val config = KafkaConfig.fromProps(props)
+    // use MockTime for IP connection rate quota tests that don't expect to block
+    setupMockTime()
+    connectionQuotas = new ConnectionQuotas(config, time, metrics)
+    addListenersAndVerify(config, connectionQuotas)
+    val externalListener = listeners("EXTERNAL")
+    connectionQuotas.updateIpConnectionRateQuota(Some(externalListener.defaultIp), Some(ipConnectionRateLimit))
+    val numConnections = 200
+    // create connections with the rate < ip quota and verify there is no throttling
+    acceptConnectionsAndVerifyRate(connectionQuotas, externalListener, numConnections, connCreateIntervalMs,
+      expectedRate = 25, epsilon = 0)
+    assertEquals(s"Number of connections on $externalListener:",
+      numConnections, connectionQuotas.get(externalListener.defaultIp))
+
+    val adminListener = listeners("ADMIN")
+    val unthrottledConnectionCreateInterval = 20 // connection creation rate = 50/s
+    // create connections with an IP with no quota and verify there is no throttling
+    acceptConnectionsAndVerifyRate(connectionQuotas, adminListener, numConnections, unthrottledConnectionCreateInterval,
+      expectedRate = 50, epsilon = 0)
+
+    assertEquals(s"Number of connections on $adminListener:",
+      numConnections, connectionQuotas.get(adminListener.defaultIp))
+
+    // acceptor shouldn't block for IP rate throttling
+    verifyNoBlockedPercentRecordedOnAllListeners()
+    // no IP throttle time should be recorded on any listeners
+    listeners.values.map(_.listenerName).foreach(verifyIpThrottleTimeOnListener(_, expectThrottle = false))
+  }
+
+  @Test
+  def testIpConnectionRateWhenActualRateAboveLimit(): Unit = {
+    val ipConnectionRateLimit = 20
+    val connCreateIntervalMs = 25 // connection creation rate = 40/sec
+    val props = brokerPropsWithDefaultConnectionLimits
+    val config = KafkaConfig.fromProps(props)
+    // use MockTime for IP connection rate quota tests that don't expect to block
+    setupMockTime()
+    connectionQuotas = new ConnectionQuotas(config, time, metrics)
+    addListenersAndVerify(config, connectionQuotas)
+    val externalListener = listeners("EXTERNAL")
+    connectionQuotas.updateIpConnectionRateQuota(Some(externalListener.defaultIp), Some(ipConnectionRateLimit))
+    // create connections with the rate > ip quota
+    val numConnections = 80
+    acceptConnectionsAndVerifyRate(connectionQuotas, externalListener, numConnections, connCreateIntervalMs, ipConnectionRateLimit,
+      1, expectIpThrottle = true)
+    verifyIpThrottleTimeOnListener(externalListener.listenerName, expectThrottle = true)
+
+    // verify that default quota applies to IPs without a quota override
+    connectionQuotas.updateIpConnectionRateQuota(None, Some(ipConnectionRateLimit))
+    val adminListener = listeners("ADMIN")
+    // listener shouldn't have any IP throttle time recorded
+    verifyIpThrottleTimeOnListener(adminListener.listenerName, expectThrottle = false)
+    acceptConnectionsAndVerifyRate(connectionQuotas, adminListener, numConnections, connCreateIntervalMs, ipConnectionRateLimit,
+      1, expectIpThrottle = true)
+    verifyIpThrottleTimeOnListener(adminListener.listenerName, expectThrottle = true)
+
+    // acceptor shouldn't block for IP rate throttling
+    verifyNoBlockedPercentRecordedOnAllListeners()
+    // replication listener shouldn't have any IP throttling recorded
+    verifyIpThrottleTimeOnListener(listeners("REPLICATION").listenerName, expectThrottle = false)
+  }
+
+  @Test
+  def testIpConnectionRateWithListenerConnectionRate(): Unit = {
+    val ipConnectionRateLimit = 25
+    val listenerRateLimit = 35
+    val props = brokerPropsWithDefaultConnectionLimits
+    val config = KafkaConfig.fromProps(props)
+    connectionQuotas = new ConnectionQuotas(config, time, metrics)
+    // with a default per-IP limit of 25 and a listener rate of 30, only one IP should be able to saturate their IP rate
+    // limit, the other IP will hit listener rate limits and block
+    connectionQuotas.updateIpConnectionRateQuota(None, Some(ipConnectionRateLimit))
+    val listenerConfig = Map(KafkaConfig.MaxConnectionCreationRateProp -> listenerRateLimit.toString).asJava
+    addListenersAndVerify(config, listenerConfig, connectionQuotas)
+    val listener = listeners("EXTERNAL").listenerName
+    // use a small number of connections because a longer-running test will have both IPs throttle at different times
+    val numConnections = 35
+    val futures = List(
+      executor.submit((() => acceptConnections(connectionQuotas, listener, knownHost, numConnections,
+        0, true)): Callable[Boolean]),
+      executor.submit((() => acceptConnections(connectionQuotas, listener, unknownHost, numConnections,
+        0, true)): Callable[Boolean])
+    )
+
+    val ipsThrottledResults = futures.map(_.get(3, TimeUnit.SECONDS))
+    val throttledIps = ipsThrottledResults.filter(identity)
+    // at most one IP should get IP throttled before the acceptor blocks on listener quota
+    assertTrue("Expected BlockedPercentMeter metric for EXTERNAL listener to be recorded", blockedPercentMeters("EXTERNAL").count() > 0)
+    assertTrue("Expect at most one IP to get throttled", throttledIps.size < 2)
+  }
+
+  @Test
+  def testRejectedIpConnectionUnrecordedFromConnectionRateQuotas(): Unit = {
+    val config = KafkaConfig.fromProps(brokerPropsWithDefaultConnectionLimits)
+    connectionQuotas = new ConnectionQuotas(config, new MockTime(), metrics)
+    addListenersAndVerify(config, connectionQuotas)
+    val externalListener = listeners("EXTERNAL")
+    val protectedListener = listeners("REPLICATION")
+    connectionQuotas.updateIpConnectionRateQuota(Some(externalListener.defaultIp), Some(0))
+    connectionQuotas.updateIpConnectionRateQuota(Some(protectedListener.defaultIp), Some(0))
+
+    assertThrows[ConnectionThrottledException] {
+      connectionQuotas.inc(externalListener.listenerName, externalListener.defaultIp, blockedPercentMeters("EXTERNAL"))
+    }
+
+    val brokerRateMetric = brokerConnRateMetric()
+    // rejected connection shouldn't be recorded for any of the connection accepted rate metrics
+    assertEquals(0, metricValue(ipConnRateMetric(externalListener.defaultIp.getHostAddress)), eps)
+    assertEquals(0, metricValue(listenerConnRateMetric(externalListener.listenerName.value)), eps)
+    assertEquals(0, metricValue(brokerRateMetric), eps)
+
+    assertThrows[ConnectionThrottledException] {
+      connectionQuotas.inc(protectedListener.listenerName, protectedListener.defaultIp, blockedPercentMeters("REPLICATION"))
+    }
+
+    assertEquals(0, metricValue(ipConnRateMetric(protectedListener.defaultIp.getHostAddress)), eps)
+    assertEquals(0, metricValue(listenerConnRateMetric(protectedListener.listenerName.value)), eps)
+    assertEquals(0, metricValue(brokerRateMetric), eps)
+  }
+
+  @Test
+  def testMaxListenerConnectionListenerMustBeAboveZero(): Unit = {
+    val config = KafkaConfig.fromProps(brokerPropsWithDefaultConnectionLimits)
+    connectionQuotas = new ConnectionQuotas(config, time, metrics)
+
+    connectionQuotas.addListener(config, listeners("EXTERNAL").listenerName)
+
+    val maxListenerConnectionRate = 0
+    val listenerConfig = Map(KafkaConfig.MaxConnectionCreationRateProp -> maxListenerConnectionRate.toString).asJava
+    assertThrows[ConfigException] {
+      connectionQuotas.maxConnectionsPerListener(listeners("EXTERNAL").listenerName).validateReconfiguration(listenerConfig)
+    }
+  }
+
+  @Test
+  def testMaxListenerConnectionRateReconfiguration(): Unit = {
+    val config = KafkaConfig.fromProps(brokerPropsWithDefaultConnectionLimits)
+    connectionQuotas = new ConnectionQuotas(config, time, metrics)
+    connectionQuotas.addListener(config, listeners("EXTERNAL").listenerName)
+
+    val listenerRateLimit = 20
+    val listenerConfig = Map(KafkaConfig.MaxConnectionCreationRateProp -> listenerRateLimit.toString).asJava
+    connectionQuotas.maxConnectionsPerListener(listeners("EXTERNAL").listenerName).configure(listenerConfig)
+
+    // remove connection rate limit
+    connectionQuotas.maxConnectionsPerListener(listeners("EXTERNAL").listenerName).reconfigure(Map.empty.asJava)
+
+    // create connections as fast as possible, will timeout if connections get throttled with previous rate
+    // (50s to create 1000 connections)
+    executor.submit((() =>
+      acceptConnections(connectionQuotas, listeners("EXTERNAL"), 1000)): Runnable
+    ).get(10, TimeUnit.SECONDS)
+    // verify no throttling
+    assertEquals(s"BlockedPercentMeter metric for EXTERNAL listener", 0, blockedPercentMeters("EXTERNAL").count())
+
+    // configure 100 connection/second rate limit
+    val newMaxListenerConnectionRate = 10
+    val newListenerConfig = Map(KafkaConfig.MaxConnectionCreationRateProp -> newMaxListenerConnectionRate.toString).asJava
+    connectionQuotas.maxConnectionsPerListener(listeners("EXTERNAL").listenerName).reconfigure(newListenerConfig)
+
+    // verify rate limit
+    val connectionsPerListener = 200 // should take 20 seconds to create 200 connections with rate = 10/sec
+    executor.submit((() =>
+      acceptConnectionsAndVerifyRate(connectionQuotas, listeners("EXTERNAL"), connectionsPerListener, 5, newMaxListenerConnectionRate, 3)): Runnable
+    ).get(30, TimeUnit.SECONDS)
+    assertTrue("Expected BlockedPercentMeter metric for EXTERNAL listener to be recorded", blockedPercentMeters("EXTERNAL").count() > 0)
+  }
+
+  @Test
+  def testMaxBrokerConnectionRateReconfiguration(): Unit = {
+    val config = KafkaConfig.fromProps(brokerPropsWithDefaultConnectionLimits)
+    connectionQuotas = new ConnectionQuotas(config, time, metrics)
+    connectionQuotas.addListener(config, listeners("EXTERNAL").listenerName)
+
+    addListenersAndVerify(config, connectionQuotas)
+
+    val maxBrokerConnectionRate = 50
+    connectionQuotas.updateBrokerMaxConnectionRate(maxBrokerConnectionRate)
+
+    // create connections with rate = 200 conn/sec (5ms interval), so that connection rate gets throttled
+    val totalConnections = 400
+    executor.submit((() =>
+      // this is a short run, so setting epsilon higher (enough to check that the rate is not unlimited)
+      acceptConnectionsAndVerifyRate(connectionQuotas, listeners("EXTERNAL"), totalConnections, 5, maxBrokerConnectionRate, 20)): Runnable
+    ).get(10, TimeUnit.SECONDS)
+    assertTrue("Expected BlockedPercentMeter metric for EXTERNAL listener to be recorded", blockedPercentMeters("EXTERNAL").count() > 0)
+  }
+
+  @Test
+  def testIpConnectionRateMetricUpdate(): Unit = {
+    val config = KafkaConfig.fromProps(brokerPropsWithDefaultConnectionLimits)
+    connectionQuotas = new ConnectionQuotas(config, time, metrics)
+    connectionQuotas.addListener(config, listeners("EXTERNAL").listenerName)
+    connectionQuotas.addListener(config, listeners("ADMIN").listenerName)
+    val defaultIpRate = 50
+    val defaultOverrideRate = 20
+    val overrideIpRate = 30
+    val externalListener = listeners("EXTERNAL")
+    val adminListener = listeners("ADMIN")
+    // set a non-unlimited default quota so that we create ip rate sensors/metrics
+    connectionQuotas.updateIpConnectionRateQuota(None, Some(defaultIpRate))
+    connectionQuotas.inc(externalListener.listenerName, externalListener.defaultIp, blockedPercentMeters("EXTERNAL"))
+    connectionQuotas.inc(adminListener.listenerName, adminListener.defaultIp, blockedPercentMeters("ADMIN"))
+
+    // both IPs should have the default rate
+    verifyIpConnectionQuota(externalListener.defaultIp, defaultIpRate)
+    verifyIpConnectionQuota(adminListener.defaultIp, defaultIpRate)
+
+    // external listener should have its in-memory quota and metric config updated
+    connectionQuotas.updateIpConnectionRateQuota(Some(externalListener.defaultIp), Some(overrideIpRate))
+    verifyIpConnectionQuota(externalListener.defaultIp, overrideIpRate)
+
+    // update default
+    connectionQuotas.updateIpConnectionRateQuota(None, Some(defaultOverrideRate))
+
+    // external listener IP should not have its quota updated to the new default
+    verifyIpConnectionQuota(externalListener.defaultIp, overrideIpRate)
+    // admin listener IP should have its quota updated with to the new default
+    verifyIpConnectionQuota(adminListener.defaultIp, defaultOverrideRate)
+
+    // remove default connection rate quota
+    connectionQuotas.updateIpConnectionRateQuota(None, None)
+    verifyIpConnectionQuota(adminListener.defaultIp, DynamicConfig.Ip.DefaultConnectionCreationRate)
+    verifyIpConnectionQuota(externalListener.defaultIp, overrideIpRate)
+
+    // remove override for external listener IP
+    connectionQuotas.updateIpConnectionRateQuota(Some(externalListener.defaultIp), None)
+    verifyIpConnectionQuota(externalListener.defaultIp, DynamicConfig.Ip.DefaultConnectionCreationRate)
+  }
+
+  @Test
+  def testEnforcedIpConnectionRateQuotaUpdate(): Unit = {
+    val ipConnectionRateLimit = 20
+    val props = brokerPropsWithDefaultConnectionLimits
+    val config = KafkaConfig.fromProps(props)
+    // use MockTime for IP connection rate quota tests that don't expect to block
+    setupMockTime()
+    connectionQuotas = new ConnectionQuotas(config, time, metrics)
+    addListenersAndVerify(config, connectionQuotas)
+    val externalListener = listeners("EXTERNAL")
+    connectionQuotas.updateIpConnectionRateQuota(Some(externalListener.defaultIp), Some(ipConnectionRateLimit))
+    // create connections with the rate > ip quota
+    val connectionRate = 40
+    assertThrows[ConnectionThrottledException] {
+      acceptConnections(connectionQuotas, externalListener, connectionRate)
+    }
+    assertEquals(s"Number of connections on $externalListener:",
+      ipConnectionRateLimit, connectionQuotas.get(externalListener.defaultIp))
+
+    // increase ip quota, we should accept connections up to the new quota limit
+    val updatedRateLimit = 30
+    connectionQuotas.updateIpConnectionRateQuota(Some(externalListener.defaultIp), Some(updatedRateLimit))
+    assertThrows[ConnectionThrottledException] {
+      acceptConnections(connectionQuotas, externalListener, connectionRate)
+    }
+    assertEquals(s"Number of connections on $externalListener:",
+      updatedRateLimit, connectionQuotas.get(externalListener.defaultIp))
+
+    // remove IP quota, all connections should get accepted
+    connectionQuotas.updateIpConnectionRateQuota(Some(externalListener.defaultIp), None)
+    acceptConnections(connectionQuotas, externalListener, connectionRate)
+    assertEquals(s"Number of connections on $externalListener:",
+      connectionRate + updatedRateLimit, connectionQuotas.get(externalListener.defaultIp))
+
+    // create connections on a different IP,
+    val adminListener = listeners("ADMIN")
+    acceptConnections(connectionQuotas, adminListener, connectionRate)
+    assertEquals(s"Number of connections on $adminListener:",
+      connectionRate, connectionQuotas.get(adminListener.defaultIp))
+
+    // set a default IP quota, verify that quota gets propagated
+    connectionQuotas.updateIpConnectionRateQuota(None, Some(ipConnectionRateLimit))
+    assertThrows[ConnectionThrottledException] {
+      acceptConnections(connectionQuotas, adminListener, connectionRate)
+    }
+    assertEquals(s"Number of connections on $adminListener:",
+      connectionRate + ipConnectionRateLimit, connectionQuotas.get(adminListener.defaultIp))
+
+    // acceptor shouldn't block for IP rate throttling
+    verifyNoBlockedPercentRecordedOnAllListeners()
+  }
+
+  @Test
+  def testNonDefaultConnectionCountLimitAndRateLimit(): Unit = {
+    val brokerRateLimit = 25
+    val maxConnections = 350  // with rate == 25, will run out of connections in 14 seconds
+    val props = brokerPropsWithDefaultConnectionLimits
+    props.put(KafkaConfig.MaxConnectionsProp, maxConnections.toString)
+    props.put(KafkaConfig.MaxConnectionCreationRateProp, brokerRateLimit.toString)
+    val config = KafkaConfig.fromProps(props)
+    connectionQuotas = new ConnectionQuotas(config, time, metrics)
+    connectionQuotas.addListener(config, listeners("EXTERNAL").listenerName)
+
+    addListenersAndVerify(config, connectionQuotas)
+
+    // create connections with rate = 100 conn/sec (10ms interval), so that connection rate gets throttled
+    val listener = listeners("EXTERNAL")
+    executor.submit((() =>
+      acceptConnectionsAndVerifyRate(connectionQuotas, listener, maxConnections, 10, brokerRateLimit, 8)): Runnable
+    ).get(20, TimeUnit.SECONDS)
+    assertTrue("Expected BlockedPercentMeter metric for EXTERNAL listener to be recorded", blockedPercentMeters("EXTERNAL").count() > 0)
+    assertEquals(s"Number of connections on EXTERNAL listener:", maxConnections, connectionQuotas.get(listener.defaultIp))
+
+    // adding one more connection will block ConnectionQuota.inc()
+    val future = executor.submit((() =>
+      acceptConnections(connectionQuotas, listeners("EXTERNAL"), 1)): Runnable
+    )
+    intercept[TimeoutException](future.get(100, TimeUnit.MILLISECONDS))
+
+    // removing one connection should make the waiting connection to succeed
+    connectionQuotas.dec(listener.listenerName, listener.defaultIp)
+    future.get(1, TimeUnit.SECONDS)
+    assertEquals(s"Number of connections on EXTERNAL listener:", maxConnections, connectionQuotas.get(listener.defaultIp))
+  }
+
+  private def addListenersAndVerify(config: KafkaConfig, connectionQuotas: ConnectionQuotas) : Unit = {
+    addListenersAndVerify(config, Map.empty.asJava, connectionQuotas)
+  }
+
+  private def addListenersAndVerify(config: KafkaConfig,
+                                    listenerConfig: util.Map[String, _],
+                                    connectionQuotas: ConnectionQuotas) : Unit = {
+    assertNotNull("Expected broker-connection-accept-rate metric to exist", brokerConnRateMetric())
+
+    // add listeners and verify connection limits not exceeded
+    listeners.foreach { case (name, listener) =>
+      val listenerName = listener.listenerName
+      connectionQuotas.addListener(config, listenerName)
+      connectionQuotas.maxConnectionsPerListener(listenerName).configure(listenerConfig)
+      assertFalse(s"Should not exceed max connection limit on $name listener after initialization",
+        connectionQuotas.maxConnectionsExceeded(listenerName))
+      assertEquals(s"Number of connections on $listener listener:",
+        0, connectionQuotas.get(listener.defaultIp))
+      assertNotNull(s"Expected connection-accept-rate metric to exist for listener ${listenerName.value}",
+        listenerConnRateMetric(listenerName.value))
+      assertEquals(s"Connection acceptance rate metric for listener ${listenerName.value}",
+          0, metricValue(listenerConnRateMetric(listenerName.value)), eps)
+      assertNotNull(s"Expected connection-accept-throttle-time metric to exist for listener ${listenerName.value}",
+        listenerConnThrottleMetric(listenerName.value))
+      assertEquals(s"Listener connection throttle metric for listener ${listenerName.value}",
+        0, metricValue(listenerConnThrottleMetric(listenerName.value)).toLong)
+      assertEquals(s"Ip connection throttle metric for listener ${listenerName.value}",
+        0, metricValue(ipConnThrottleMetric(listenerName.value)).toLong)
+    }
+    verifyNoBlockedPercentRecordedOnAllListeners()
+    assertEquals("Broker-wide connection acceptance rate metric", 0, metricValue(brokerConnRateMetric()), eps)
+  }
+
+  private def verifyNoBlockedPercentRecordedOnAllListeners(): Unit = {
+    blockedPercentMeters.foreach { case (name, meter) =>
+      assertEquals(s"BlockedPercentMeter metric for $name listener", 0, meter.count())
+    }
+  }
+
+  private def verifyNonZeroBlockedPercentAndThrottleTimeOnAllListeners(): Unit = {
+    blockedPercentMeters.foreach { case (name, meter) =>
+      assertTrue(s"Expected BlockedPercentMeter metric for $name listener to be recorded", meter.count() > 0)
+    }
+    listeners.values.foreach { listener =>
+      assertTrue(s"Connection throttle metric for listener ${listener.listenerName.value}",
+        metricValue(listenerConnThrottleMetric(listener.listenerName.value)).toLong > 0)
+    }
+  }
+
+  private def verifyIpThrottleTimeOnListener(listener: ListenerName, expectThrottle: Boolean): Unit = {
+    assertEquals(s"IP connection throttle recorded for listener ${listener.value}", expectThrottle,
+      metricValue(ipConnThrottleMetric(listener.value)).toLong > 0)
+  }
+
+  private def verifyOnlyNonInterBrokerListenersBlockedPercentRecorded(): Unit = {
+    blockedPercentMeters.foreach { case (name, meter) =>
+      name match {
+        case "REPLICATION" =>
+          assertEquals(s"BlockedPercentMeter metric for $name listener", 0, meter.count())
+        case _ =>
+          assertTrue(s"Expected BlockedPercentMeter metric for $name listener to be recorded", meter.count() > 0)
+      }
+    }
+  }
+
+  private def verifyConnectionCountOnEveryListener(connectionQuotas: ConnectionQuotas, expectedConnectionCount: Int): Unit = {
+    listeners.values.foreach { listener =>
+      assertEquals(s"Number of connections on $listener:",
+        expectedConnectionCount, connectionQuotas.get(listener.defaultIp))
+    }
+  }
+
+  private def listenerConnThrottleMetric(listener: String) : KafkaMetric = {
+    val metricName = metrics.metricName(
+      "connection-accept-throttle-time",
+      SocketServer.MetricsGroup,
+      Collections.singletonMap(Processor.ListenerMetricTag, listener))
+    metrics.metric(metricName)
+  }
+
+  private def ipConnThrottleMetric(listener: String): KafkaMetric = {
+    val metricName = metrics.metricName(
+      "ip-connection-accept-throttle-time",
+      SocketServer.MetricsGroup,
+      Collections.singletonMap(Processor.ListenerMetricTag, listener))
+    metrics.metric(metricName)
+  }
+
+  private def listenerConnRateMetric(listener: String) : KafkaMetric = {
+    val metricName = metrics.metricName(
+      "connection-accept-rate",
+      SocketServer.MetricsGroup,
+      Collections.singletonMap(Processor.ListenerMetricTag, listener))
+    metrics.metric(metricName)
+  }
+
+  private def brokerConnRateMetric() : KafkaMetric = {
+    val metricName = metrics.metricName(
+      s"broker-connection-accept-rate",
+      SocketServer.MetricsGroup)
+    metrics.metric(metricName)
+  }
+
+  private def ipConnRateMetric(ip: String): KafkaMetric = {
+    val metricName = metrics.metricName(
+      s"connection-accept-rate",
+      SocketServer.MetricsGroup,
+      Collections.singletonMap("ip", ip))
+    metrics.metric(metricName)
+  }
+
+  private def metricValue(metric: KafkaMetric): Double = {
+    metric.metricValue.asInstanceOf[Double]
+  }
+
+  private def verifyIpConnectionQuota(ip: InetAddress, quota: Int): Unit = {
+    // verify connection quota in-memory rate and metric
+    assertEquals(quota, connectionQuotas.connectionRateForIp(ip))
+    Option(ipConnRateMetric(ip.getHostAddress)) match {
+      case Some(metric) => assertEquals(quota, metric.config.quota.bound, 0.1)
+      case None => fail(s"Expected $ip connection rate metric to be defined")
+    }
+  }
+
+  // this method must be called on a separate thread, because connectionQuotas.inc() may block
+  private def acceptConnections(connectionQuotas: ConnectionQuotas,
+                                listenerDesc: ListenerDesc,
+                                numConnections: Long,
+                                timeIntervalMs: Long = 0L,
+                                expectIpThrottle: Boolean = false) : Unit = {
+    acceptConnections(connectionQuotas, listenerDesc.listenerName, listenerDesc.defaultIp, numConnections,
+      timeIntervalMs, expectIpThrottle)
+  }
+
+  // this method must be called on a separate thread, because connectionQuotas.inc() may block
+  private def acceptConnectionsAndVerifyRate(connectionQuotas: ConnectionQuotas,
+                                             listenerDesc: ListenerDesc,
+                                             numConnections: Long,
+                                             timeIntervalMs: Long,
+                                             expectedRate: Int,
+                                             epsilon: Int,
+                                             expectIpThrottle: Boolean = false) : Unit = {
+    val startTimeMs = time.milliseconds
+    val startNumConnections = connectionQuotas.get(listenerDesc.defaultIp)
+    acceptConnections(connectionQuotas, listenerDesc.listenerName, listenerDesc.defaultIp, numConnections,
+      timeIntervalMs, expectIpThrottle)
+    val elapsedSeconds = MetricsUtils.convert(time.milliseconds - startTimeMs, TimeUnit.SECONDS)
+    val createdConnections = connectionQuotas.get(listenerDesc.defaultIp) - startNumConnections
+    val actualRate = createdConnections.toDouble / elapsedSeconds
+    assertEquals(s"Expected rate ($expectedRate +- $epsilon), but got $actualRate ($createdConnections connections / $elapsedSeconds sec)",
+      expectedRate.toDouble, actualRate, epsilon)
+  }
+
+  /**
+   * This method will "create" connections every 'timeIntervalMs' which translates to 1000/timeIntervalMs connection rate,
+   * as long as the rate is below the connection rate limit. Otherwise, connections will be essentially created as
+   * fast as possible, which would result in the maximum connection creation rate.
+   *
+   * This method must be called on a separate thread, because connectionQuotas.inc() may block
+   */
+  private def acceptConnections(connectionQuotas: ConnectionQuotas,
+                                listenerName: ListenerName,
+                                address: InetAddress,
+                                numConnections: Long,
+                                timeIntervalMs: Long,
+                                expectIpThrottle: Boolean): Boolean = {
+    var nextSendTime = time.milliseconds + timeIntervalMs
+    var ipThrottled = false
+    for (_ <- 0L until numConnections) {
+      // this method may block if broker-wide or listener limit on the number of connections is reached
+      try {
+        connectionQuotas.inc(listenerName, address, blockedPercentMeters(listenerName.value))
+      } catch {
+        case e: ConnectionThrottledException =>
+          if (!expectIpThrottle)
+            throw e
+          ipThrottled = true
+      }
+      val sleepMs = math.max(nextSendTime - time.milliseconds, 0)
+      if (sleepMs > 0)
+        time.sleep(sleepMs)
+
+      nextSendTime = nextSendTime + timeIntervalMs
+    }
+    ipThrottled
+  }
+
+  // this method must be called on a separate thread, because connectionQuotas.inc() may block
+  private def acceptConnectionsAboveIpLimit(connectionQuotas: ConnectionQuotas,
+                                            listenerDesc: ListenerDesc,
+                                            numConnections: Long) : Unit = {
+    val listenerName = listenerDesc.listenerName
+    for (i <- 0L until numConnections) {
+      // this method may block if broker-wide or listener limit is reached
+      intercept[TooManyConnectionsException](
+        connectionQuotas.inc(listenerName, listenerDesc.defaultIp, blockedPercentMeters(listenerName.value))
+      )
+    }
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/network/RequestChannelTest.scala b/core/src/test/scala/unit/kafka/network/RequestChannelTest.scala
new file mode 100644
index 0000000000000..b7b950794fef6
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/network/RequestChannelTest.scala
@@ -0,0 +1,224 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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.network
+
+
+import java.io.IOException
+import java.net.InetAddress
+import java.nio.ByteBuffer
+import java.util.Collections
+
+import com.fasterxml.jackson.databind.ObjectMapper
+import kafka.network
+import org.apache.kafka.clients.admin.AlterConfigOp.OpType
+import org.apache.kafka.common.config.types.Password
+import org.apache.kafka.common.config.{ConfigResource, SaslConfigs, SslConfigs, TopicConfig}
+import org.apache.kafka.common.memory.MemoryPool
+import org.apache.kafka.common.message.IncrementalAlterConfigsRequestData
+import org.apache.kafka.common.message.IncrementalAlterConfigsRequestData._
+import org.apache.kafka.common.network.{ClientInformation, ListenerName}
+import org.apache.kafka.common.requests._
+import org.apache.kafka.common.requests.AlterConfigsRequest._
+import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
+import org.easymock.EasyMock._
+import org.junit.Assert._
+import org.junit._
+
+import scala.jdk.CollectionConverters._
+
+class RequestChannelTest {
+
+  @Test
+  def testAlterRequests(): Unit = {
+
+    val sensitiveValue = "secret"
+    def verifyConfig(resource: ConfigResource, entries: Seq[ConfigEntry], expectedValues: Map[String, String]): Unit = {
+      val alterConfigs = request(new AlterConfigsRequest.Builder(
+          Collections.singletonMap(resource, new Config(entries.asJavaCollection)), true).build())
+
+      val loggableAlterConfigs = alterConfigs.loggableRequest.asInstanceOf[AlterConfigsRequest]
+      val loggedConfig = loggableAlterConfigs.configs.get(resource)
+      assertEquals(expectedValues, toMap(loggedConfig))
+      val alterConfigsDesc = RequestConvertToJson.requestDesc(alterConfigs.header, alterConfigs.requestLog, alterConfigs.isForwarded).toString
+      assertFalse(s"Sensitive config logged $alterConfigsDesc", alterConfigsDesc.contains(sensitiveValue))
+    }
+
+    val brokerResource = new ConfigResource(ConfigResource.Type.BROKER, "1")
+    val keystorePassword = new ConfigEntry(SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG, sensitiveValue)
+    verifyConfig(brokerResource, Seq(keystorePassword), Map(SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG -> Password.HIDDEN))
+
+    val keystoreLocation = new ConfigEntry(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG, "/path/to/keystore")
+    verifyConfig(brokerResource, Seq(keystoreLocation), Map(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG -> "/path/to/keystore"))
+    verifyConfig(brokerResource, Seq(keystoreLocation, keystorePassword),
+      Map(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG -> "/path/to/keystore", SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG -> Password.HIDDEN))
+
+    val listenerKeyPassword = new ConfigEntry(s"listener.name.internal.${SslConfigs.SSL_KEY_PASSWORD_CONFIG}", sensitiveValue)
+    verifyConfig(brokerResource, Seq(listenerKeyPassword), Map(listenerKeyPassword.name -> Password.HIDDEN))
+
+    val listenerKeystore = new ConfigEntry(s"listener.name.internal.${SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG}", "/path/to/keystore")
+    verifyConfig(brokerResource, Seq(listenerKeystore), Map(listenerKeystore.name -> "/path/to/keystore"))
+
+    val plainJaasConfig = new ConfigEntry(s"listener.name.internal.plain.${SaslConfigs.SASL_JAAS_CONFIG}", sensitiveValue)
+    verifyConfig(brokerResource, Seq(plainJaasConfig), Map(plainJaasConfig.name -> Password.HIDDEN))
+
+    val plainLoginCallback = new ConfigEntry(s"listener.name.internal.plain.${SaslConfigs.SASL_LOGIN_CALLBACK_HANDLER_CLASS}", "test.LoginClass")
+    verifyConfig(brokerResource, Seq(plainLoginCallback), Map(plainLoginCallback.name -> plainLoginCallback.value))
+
+    val customConfig = new ConfigEntry("custom.config", sensitiveValue)
+    verifyConfig(brokerResource, Seq(customConfig), Map(customConfig.name -> Password.HIDDEN))
+
+    val topicResource = new ConfigResource(ConfigResource.Type.TOPIC, "testTopic")
+    val compressionType = new ConfigEntry(TopicConfig.COMPRESSION_TYPE_CONFIG, "lz4")
+    verifyConfig(topicResource, Seq(compressionType), Map(TopicConfig.COMPRESSION_TYPE_CONFIG -> "lz4"))
+    verifyConfig(topicResource, Seq(customConfig), Map(customConfig.name -> Password.HIDDEN))
+
+    // Verify empty request
+    val alterConfigs = request(new AlterConfigsRequest.Builder(
+        Collections.emptyMap[ConfigResource, Config], true).build())
+    assertEquals(Collections.emptyMap, alterConfigs.loggableRequest.asInstanceOf[AlterConfigsRequest].configs)
+  }
+
+  @Test
+  def testIncrementalAlterRequests(): Unit = {
+
+    def incrementalAlterConfigs(resource: ConfigResource,
+                                entries: Map[String, String], op: OpType): IncrementalAlterConfigsRequest = {
+      val data = new IncrementalAlterConfigsRequestData()
+      val alterableConfigs = new AlterableConfigCollection()
+      entries.foreach { case (name, value) =>
+        alterableConfigs.add(new AlterableConfig().setName(name).setValue(value).setConfigOperation(op.id))
+      }
+      data.resources.add(new AlterConfigsResource()
+        .setResourceName(resource.name).setResourceType(resource.`type`.id)
+        .setConfigs(alterableConfigs))
+      new IncrementalAlterConfigsRequest.Builder(data).build()
+    }
+
+    val sensitiveValue = "secret"
+    def verifyConfig(resource: ConfigResource,
+                     op: OpType,
+                     entries: Map[String, String],
+                     expectedValues: Map[String, String]): Unit = {
+      val alterConfigs = request(incrementalAlterConfigs(resource, entries, op))
+      val loggableAlterConfigs = alterConfigs.loggableRequest.asInstanceOf[IncrementalAlterConfigsRequest]
+      val loggedConfig = loggableAlterConfigs.data.resources.find(resource.`type`.id, resource.name).configs
+      assertEquals(expectedValues, toMap(loggedConfig))
+      val alterConfigsDesc = RequestConvertToJson.requestDesc(alterConfigs.header, alterConfigs.requestLog, alterConfigs.isForwarded).toString
+      assertFalse(s"Sensitive config logged $alterConfigsDesc", alterConfigsDesc.contains(sensitiveValue))
+    }
+
+    val brokerResource = new ConfigResource(ConfigResource.Type.BROKER, "1")
+    val keystorePassword = Map(SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG -> sensitiveValue)
+    verifyConfig(brokerResource, OpType.SET, keystorePassword, Map(SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG -> Password.HIDDEN))
+
+    val keystoreLocation = Map(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG -> "/path/to/keystore")
+    verifyConfig(brokerResource, OpType.SET, keystoreLocation, keystoreLocation)
+    verifyConfig(brokerResource, OpType.SET, keystoreLocation ++ keystorePassword,
+      Map(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG -> "/path/to/keystore", SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG -> Password.HIDDEN))
+
+    val listenerKeyPassword = Map(s"listener.name.internal.${SslConfigs.SSL_KEY_PASSWORD_CONFIG}" -> sensitiveValue)
+    verifyConfig(brokerResource, OpType.SET, listenerKeyPassword,
+      Map(s"listener.name.internal.${SslConfigs.SSL_KEY_PASSWORD_CONFIG}" -> Password.HIDDEN))
+
+    val listenerKeystore = Map(s"listener.name.internal.${SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG}" -> "/path/to/keystore")
+    verifyConfig(brokerResource, OpType.SET, listenerKeystore, listenerKeystore)
+
+    val plainJaasConfig = Map(s"listener.name.internal.plain.${SaslConfigs.SASL_JAAS_CONFIG}" -> sensitiveValue)
+    verifyConfig(brokerResource, OpType.SET, plainJaasConfig,
+      Map(s"listener.name.internal.plain.${SaslConfigs.SASL_JAAS_CONFIG}" -> Password.HIDDEN))
+
+    val plainLoginCallback = Map(s"listener.name.internal.plain.${SaslConfigs.SASL_LOGIN_CALLBACK_HANDLER_CLASS}" -> "test.LoginClass")
+    verifyConfig(brokerResource, OpType.SET, plainLoginCallback, plainLoginCallback)
+
+    val sslProtocols = Map(SslConfigs.SSL_ENABLED_PROTOCOLS_CONFIG -> "TLSv1.1")
+    verifyConfig(brokerResource, OpType.APPEND, sslProtocols, Map(SslConfigs.SSL_ENABLED_PROTOCOLS_CONFIG -> "TLSv1.1"))
+    verifyConfig(brokerResource, OpType.SUBTRACT, sslProtocols, Map(SslConfigs.SSL_ENABLED_PROTOCOLS_CONFIG -> "TLSv1.1"))
+    val cipherSuites = Map(SslConfigs.SSL_CIPHER_SUITES_CONFIG -> null)
+    verifyConfig(brokerResource, OpType.DELETE, cipherSuites, cipherSuites)
+
+    val customConfig = Map("custom.config" -> sensitiveValue)
+    verifyConfig(brokerResource, OpType.SET, customConfig, Map("custom.config" -> Password.HIDDEN))
+
+    val topicResource = new ConfigResource(ConfigResource.Type.TOPIC, "testTopic")
+    val compressionType = Map(TopicConfig.COMPRESSION_TYPE_CONFIG -> "lz4")
+    verifyConfig(topicResource, OpType.SET, compressionType, compressionType)
+    verifyConfig(topicResource, OpType.SET, customConfig, Map("custom.config" -> Password.HIDDEN))
+  }
+
+  @Test
+  def testNonAlterRequestsNotTransformed(): Unit = {
+    val metadataRequest = request(new MetadataRequest.Builder(List("topic").asJava, true).build())
+    assertSame(metadataRequest.body[MetadataRequest], metadataRequest.loggableRequest)
+  }
+
+  @Test
+  def testJsonRequests(): Unit = {
+    val sensitiveValue = "secret"
+    val resource = new ConfigResource(ConfigResource.Type.BROKER, "1")
+    val keystorePassword = new ConfigEntry(SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG, sensitiveValue)
+    val entries = Seq(keystorePassword)
+
+    val alterConfigs = request(new AlterConfigsRequest.Builder(Collections.singletonMap(resource,
+      new Config(entries.asJavaCollection)), true).build())
+
+    assertTrue(isValidJson(RequestConvertToJson.request(alterConfigs.loggableRequest).toString))
+  }
+
+  private def isValidJson(str: String): Boolean = {
+    try {
+      val mapper = new ObjectMapper
+      mapper.readTree(str)
+      true
+    } catch {
+      case _: IOException => false
+    }
+  }
+
+  def request(req: AbstractRequest): RequestChannel.Request = {
+    val buffer = RequestTestUtils.serializeRequestWithHeader(new RequestHeader(req.apiKey, req.version, "client-id", 1),
+      req)
+    val requestContext = newRequestContext(buffer)
+    new network.RequestChannel.Request(processor = 1,
+      requestContext,
+      startTimeNanos = 0,
+      createNiceMock(classOf[MemoryPool]),
+      buffer,
+      createNiceMock(classOf[RequestChannel.Metrics])
+    )
+  }
+
+  private def newRequestContext(buffer: ByteBuffer): RequestContext = {
+    new RequestContext(
+      RequestHeader.parse(buffer),
+      "connection-id",
+      InetAddress.getLoopbackAddress,
+      new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "user"),
+      ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT),
+      SecurityProtocol.PLAINTEXT,
+      new ClientInformation("name", "version"),
+      false)
+  }
+
+  private def toMap(config: Config): Map[String, String] = {
+    config.entries.asScala.map(e => e.name -> e.value).toMap
+  }
+
+  private def toMap(config: IncrementalAlterConfigsRequestData.AlterableConfigCollection): Map[String, String] = {
+    config.asScala.map(e => e.name -> e.value).toMap
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/network/RequestConvertToJsonTest.scala b/core/src/test/scala/unit/kafka/network/RequestConvertToJsonTest.scala
new file mode 100644
index 0000000000000..5dd7e31496bc7
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/network/RequestConvertToJsonTest.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.network
+
+import java.net.InetAddress
+import java.nio.ByteBuffer
+
+import com.fasterxml.jackson.databind.node.{BooleanNode, DoubleNode, JsonNodeFactory, LongNode, ObjectNode, TextNode}
+import kafka.network
+import kafka.network.RequestConvertToJson.requestHeaderNode
+import org.apache.kafka.common.memory.MemoryPool
+import org.apache.kafka.common.message._
+import org.apache.kafka.common.network.{ClientInformation, ListenerName, NetworkSend}
+import org.junit.Test
+import org.apache.kafka.common.protocol.{ApiKeys, ByteBufferAccessor, ObjectSerializationCache}
+import org.apache.kafka.common.requests._
+import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
+import org.easymock.EasyMock.createNiceMock
+import org.junit.Assert.assertEquals
+
+import scala.collection.mutable.ArrayBuffer
+
+class RequestConvertToJsonTest {
+
+  @Test
+  def testAllRequestTypesHandled(): Unit = {
+    val unhandledKeys = ArrayBuffer[String]()
+    ApiKeys.values().foreach { key => {
+      val version: Short = key.latestVersion()
+      val cache = new ObjectSerializationCache
+      val message = key match {
+        case ApiKeys.DESCRIBE_ACLS =>
+          ApiMessageType.fromApiKey(key.id).newRequest().asInstanceOf[DescribeAclsRequestData]
+            .setPatternTypeFilter(1).setResourceTypeFilter(1).setPermissionType(1).setOperation(1)
+        case _ =>
+          ApiMessageType.fromApiKey(key.id).newRequest()
+      }
+      val messageSize = message.size(cache, version)
+      val bytes = new ByteBufferAccessor(ByteBuffer.allocate(messageSize))
+      message.write(bytes, cache, version)
+      bytes.flip()
+
+      val req = AbstractRequest.parseRequest(key, version, bytes.buffer).request
+      try {
+        RequestConvertToJson.request(req)
+      } catch {
+        case _ : IllegalStateException => unhandledKeys += key.toString
+      }
+    }}
+    assertEquals("Unhandled request keys", ArrayBuffer.empty, unhandledKeys)
+  }
+
+  @Test
+  def testAllResponseTypesHandled(): Unit = {
+    val unhandledKeys = ArrayBuffer[String]()
+    ApiKeys.values().foreach { key => {
+      val version: Short = key.latestVersion()
+      val cache = new ObjectSerializationCache
+      val message = ApiMessageType.fromApiKey(key.id).newResponse()
+      val messageSize = message.size(cache, version)
+      val bytes = new ByteBufferAccessor(ByteBuffer.allocate(messageSize))
+      message.write(bytes, cache, version)
+      bytes.flip()
+      val res = AbstractResponse.parseResponse(key, bytes.buffer, version)
+      try {
+        RequestConvertToJson.response(res, version)
+      } catch {
+        case _ : IllegalStateException => unhandledKeys += key.toString
+      }
+    }}
+    assertEquals("Unhandled response keys", ArrayBuffer.empty, unhandledKeys)
+  }
+
+  @Test
+  def testRequestHeaderNode(): Unit = {
+    val alterIsrRequest = new AlterIsrRequest(new AlterIsrRequestData(), 0)
+    val req = request(alterIsrRequest)
+    val header = req.header
+
+    val expectedNode = RequestHeaderDataJsonConverter.write(header.data, header.headerVersion, false).asInstanceOf[ObjectNode]
+    expectedNode.set("requestApiKeyName", new TextNode(header.apiKey.toString))
+
+    val actualNode = RequestConvertToJson.requestHeaderNode(header)
+
+    assertEquals(expectedNode, actualNode);
+  }
+
+  @Test
+  def testClientInfoNode(): Unit = {
+    val clientInfo = new ClientInformation("name", "1")
+
+    val expectedNode = new ObjectNode(JsonNodeFactory.instance)
+    expectedNode.set("softwareName", new TextNode(clientInfo.softwareName))
+    expectedNode.set("softwareVersion", new TextNode(clientInfo.softwareVersion))
+
+    val actualNode = RequestConvertToJson.clientInfoNode(clientInfo)
+
+    assertEquals(expectedNode, actualNode)
+  }
+
+  @Test
+  def testRequestDesc(): Unit = {
+    val alterIsrRequest = new AlterIsrRequest(new AlterIsrRequestData(), 0)
+    val req = request(alterIsrRequest)
+
+    val expectedNode = new ObjectNode(JsonNodeFactory.instance)
+    expectedNode.set("isForwarded", if (req.isForwarded) BooleanNode.TRUE else BooleanNode.FALSE)
+    expectedNode.set("requestHeader", requestHeaderNode(req.header))
+    expectedNode.set("request", req.requestLog.getOrElse(new TextNode("")))
+
+    val actualNode = RequestConvertToJson.requestDesc(req.header, req.requestLog, req.isForwarded)
+
+    assertEquals(expectedNode, actualNode)
+  }
+
+  @Test
+  def testRequestDescMetrics(): Unit = {
+    val alterIsrRequest = new AlterIsrRequest(new AlterIsrRequestData(), 0)
+    val req = request(alterIsrRequest)
+    val send = new NetworkSend(req.context.connectionId, alterIsrRequest.toSend(req.header))
+    val headerLog = RequestConvertToJson.requestHeaderNode(req.header)
+    val res = new RequestChannel.SendResponse(req, send, Some(headerLog), None)
+
+    val totalTimeMs = 1
+    val requestQueueTimeMs = 2
+    val apiLocalTimeMs = 3
+    val apiRemoteTimeMs = 4
+    val apiThrottleTimeMs = 5
+    val responseQueueTimeMs = 6
+    val responseSendTimeMs = 7
+    val temporaryMemoryBytes = 8
+    val messageConversionsTimeMs = 9
+
+    val expectedNode = RequestConvertToJson.requestDesc(req.header, req.requestLog, req.isForwarded).asInstanceOf[ObjectNode]
+    expectedNode.set("response", res.responseLog.getOrElse(new TextNode("")))
+    expectedNode.set("connection", new TextNode(req.context.connectionId))
+    expectedNode.set("totalTimeMs", new DoubleNode(totalTimeMs))
+    expectedNode.set("requestQueueTimeMs", new DoubleNode(requestQueueTimeMs))
+    expectedNode.set("localTimeMs", new DoubleNode(apiLocalTimeMs))
+    expectedNode.set("remoteTimeMs", new DoubleNode(apiRemoteTimeMs))
+    expectedNode.set("throttleTimeMs", new LongNode(apiThrottleTimeMs))
+    expectedNode.set("responseQueueTimeMs", new DoubleNode(responseQueueTimeMs))
+    expectedNode.set("sendTimeMs", new DoubleNode(responseSendTimeMs))
+    expectedNode.set("securityProtocol", new TextNode(req.context.securityProtocol.toString))
+    expectedNode.set("principal", new TextNode(req.session.principal.toString))
+    expectedNode.set("listener", new TextNode(req.context.listenerName.value))
+    expectedNode.set("clientInformation", RequestConvertToJson.clientInfoNode(req.context.clientInformation))
+    expectedNode.set("temporaryMemoryBytes", new LongNode(temporaryMemoryBytes))
+    expectedNode.set("messageConversionsTime", new DoubleNode(messageConversionsTimeMs))
+
+    val actualNode = RequestConvertToJson.requestDescMetrics(req.header, req.requestLog, res.responseLog, req.context, req.session, req.isForwarded,
+      totalTimeMs, requestQueueTimeMs, apiLocalTimeMs, apiRemoteTimeMs, apiThrottleTimeMs, responseQueueTimeMs,
+      responseSendTimeMs, temporaryMemoryBytes, messageConversionsTimeMs).asInstanceOf[ObjectNode]
+
+    assertEquals(expectedNode, actualNode)
+  }
+
+  def request(req: AbstractRequest): RequestChannel.Request = {
+    val buffer = RequestTestUtils.serializeRequestWithHeader(new RequestHeader(req.apiKey, req.version, "client-id", 1),
+      req)
+    val requestContext = newRequestContext(buffer)
+    new network.RequestChannel.Request(processor = 1,
+      requestContext,
+      startTimeNanos = 0,
+      createNiceMock(classOf[MemoryPool]),
+      buffer,
+      createNiceMock(classOf[RequestChannel.Metrics])
+    )
+  }
+
+  private def newRequestContext(buffer: ByteBuffer): RequestContext = {
+    new RequestContext(
+      RequestHeader.parse(buffer),
+      "connection-id",
+      InetAddress.getLoopbackAddress,
+      new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "user"),
+      ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT),
+      SecurityProtocol.PLAINTEXT,
+      new ClientInformation("name", "version"),
+      false)
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/network/SocketServerTest.scala b/core/src/test/scala/unit/kafka/network/SocketServerTest.scala
index 024e7f9c3c928..adc9355b0616f 100644
--- a/core/src/test/scala/unit/kafka/network/SocketServerTest.scala
+++ b/core/src/test/scala/unit/kafka/network/SocketServerTest.scala
@@ -20,58 +20,85 @@ package kafka.network
 import java.io._
 import java.net._
 import java.nio.ByteBuffer
-import java.nio.channels.SocketChannel
-import java.util.{HashMap, Random}
-import javax.net.ssl._
+import java.nio.channels.{SelectionKey, SocketChannel}
+import java.nio.charset.StandardCharsets
+import java.util
+import java.util.concurrent.{CompletableFuture, ConcurrentLinkedQueue, Executors, TimeUnit}
+import java.util.{Properties, Random}
 
+import com.fasterxml.jackson.databind.node.{JsonNodeFactory, ObjectNode, TextNode}
 import com.yammer.metrics.core.{Gauge, Meter}
-import com.yammer.metrics.{Metrics => YammerMetrics}
-import kafka.network.RequestChannel.SendAction
+import javax.net.ssl._
+import kafka.metrics.KafkaYammerMetrics
 import kafka.security.CredentialProvider
-import kafka.server.KafkaConfig
+import kafka.server.{KafkaConfig, ThrottledChannel}
+import kafka.utils.Implicits._
 import kafka.utils.TestUtils
-import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.memory.MemoryPool
+import org.apache.kafka.common.message.{ProduceRequestData, SaslAuthenticateRequestData, SaslHandshakeRequestData, VoteRequestData}
 import org.apache.kafka.common.metrics.Metrics
-import org.apache.kafka.common.network.{ChannelBuilder, ChannelState, KafkaChannel, ListenerName, NetworkReceive, NetworkSend, Selector, Send}
+import org.apache.kafka.common.network.KafkaChannel.ChannelMuteState
+import org.apache.kafka.common.network.{ClientInformation, _}
 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
+import org.apache.kafka.common.requests._
 import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
-import org.apache.kafka.common.utils.{LogContext, MockTime, Time}
+import org.apache.kafka.common.security.scram.internals.ScramMechanism
+import org.apache.kafka.common.utils.{AppInfoParser, LogContext, MockTime, Time, Utils}
+import org.apache.kafka.test.{TestSslUtils, TestUtils => JTestUtils}
+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.jdk.CollectionConverters._
 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")
   props.put("socket.send.buffer.bytes", "300000")
   props.put("socket.receive.buffer.bytes", "300000")
   props.put("queued.max.requests", "50")
-  props.put("socket.request.max.bytes", "50")
+  props.put("socket.request.max.bytes", "100")
   props.put("max.connections.per.ip", "5")
   props.put("connections.max.idle.ms", "60000")
   val config = KafkaConfig.fromProps(props)
   val metrics = new Metrics
-  val credentialProvider = new CredentialProvider(config.saslEnabledMechanisms)
+  val credentialProvider = new CredentialProvider(ScramMechanism.mechanismNames, null)
   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()
   val sockets = new ArrayBuffer[Socket]
 
-  def sendRequest(socket: Socket, request: Array[Byte], id: Option[Short] = None, flush: Boolean = true) {
+  private val kafkaLogger = org.apache.log4j.LogManager.getLogger("kafka")
+  private var logLevelToRestore: Level = _
+
+  @Before
+  def setUp(): Unit = {
+    // Run the tests with TRACE logging to exercise request logging path
+    logLevelToRestore = kafkaLogger.getLevel
+    kafkaLogger.setLevel(Level.TRACE)
+
+    assertTrue(server.controlPlaneRequestChannelOpt.isEmpty)
+  }
+
+  @After
+  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): Unit = {
     val outgoing = new DataOutputStream(socket.getOutputStream)
     id match {
       case Some(id) =>
@@ -85,6 +112,11 @@ class SocketServerTest extends JUnitSuite {
       outgoing.flush()
   }
 
+  def sendApiRequest(socket: Socket, request: AbstractRequest, header: RequestHeader): Unit = {
+    val serializedBytes = Utils.toArray(RequestTestUtils.serializeRequestWithHeader(header, request))
+    sendRequest(socket, serializedBytes)
+  }
+
   def receiveResponse(socket: Socket): Array[Byte] = {
     val incoming = new DataInputStream(socket.getInputStream)
     val len = incoming.readInt()
@@ -102,35 +134,64 @@ 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) {
-    val byteBuffer = request.body[AbstractRequest].serialize(request.header)
-    byteBuffer.rewind()
+  def processRequest(channel: RequestChannel, request: RequestChannel.Request): Unit = {
+    val byteBuffer = RequestTestUtils.serializeRequestWithHeader(request.header, request.body[AbstractRequest])
+    val send = new NetworkSend(request.context.connectionId, ByteBufferSend.sizePrefixed(byteBuffer))
+    val headerLog = RequestConvertToJson.requestHeaderNode(request.header)
+    channel.sendResponse(new RequestChannel.SendResponse(request, send, Some(headerLog), None))
+  }
 
-    val send = new NetworkSend(request.context.connectionId, byteBuffer)
-    channel.sendResponse(new RequestChannel.Response(request, Some(send), SendAction, None))
+  def processRequestNoOpResponse(channel: RequestChannel, request: RequestChannel.Request): Unit = {
+    channel.sendResponse(new RequestChannel.NoOpResponse(request))
   }
 
-  def connect(s: SocketServer = server, protocol: SecurityProtocol = SecurityProtocol.PLAINTEXT) = {
-    val socket = new Socket("localhost", s.boundPort(ListenerName.forSecurityProtocol(protocol)))
+  def connect(s: SocketServer = server,
+              listenerName: ListenerName = ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT),
+              localAddr: InetAddress = null,
+              port: Int = 0): Socket = {
+    val socket = new Socket("localhost", s.boundPort(listenerName), localAddr, port)
     sockets += socket
     socket
   }
 
+  def sslConnect(s: SocketServer = server): Socket = {
+    val socket = sslClientSocket(s.boundPort(ListenerName.forSecurityProtocol(SecurityProtocol.SSL)))
+    sockets += socket
+    socket
+  }
+
+  private def sslClientSocket(port: Int): Socket = {
+    val sslContext = SSLContext.getInstance(TestSslUtils.DEFAULT_TLS_PROTOCOL_FOR_TESTS)
+    sslContext.init(null, Array(TestUtils.trustAllCerts), new java.security.SecureRandom())
+    val socketFactory = sslContext.getSocketFactory
+    val socket = socketFactory.createSocket("localhost", port)
+    socket.asInstanceOf[SSLSocket].setNeedClientAuth(false)
+    socket
+  }
+
   // Create a client connection, process one request and return (client socket, connectionId)
   def connectAndProcessRequest(s: SocketServer): (Socket, String) = {
-    val socket = connect(s)
+    val securityProtocol = s.dataPlaneAcceptors.asScala.head._1.securityProtocol
+    val socket = securityProtocol match {
+      case SecurityProtocol.PLAINTEXT | SecurityProtocol.SASL_PLAINTEXT =>
+        connect(s)
+      case SecurityProtocol.SSL | SecurityProtocol.SASL_SSL =>
+        sslConnect(s)
+      case _ =>
+        throw new IllegalStateException(s"Unexpected security protocol $securityProtocol")
+    }
     val request = sendAndReceiveRequest(socket, s)
-    processRequest(s.requestChannel, request)
+    processRequest(s.dataPlaneRequestChannel, request)
     (socket, request.context.connectionId)
   }
 
   def sendAndReceiveRequest(socket: Socket, server: SocketServer): RequestChannel.Request = {
-    sendRequest(socket, producerRequestBytes)
-    receiveRequest(server.requestChannel)
+    sendRequest(socket, producerRequestBytes())
+    receiveRequest(server.dataPlaneRequestChannel)
   }
 
   def shutdownServerAndMetrics(server: SocketServer): Unit = {
@@ -138,43 +199,192 @@ class SocketServerTest extends JUnitSuite {
     server.metrics.close()
   }
 
-  @After
-  def tearDown() {
-    shutdownServerAndMetrics(server)
-    sockets.foreach(_.close())
-    sockets.clear()
-  }
-
-  private def producerRequestBytes: Array[Byte] = {
+  private def producerRequestBytes(ack: Short = 0): Array[Byte] = {
     val correlationId = -1
     val clientId = ""
     val ackTimeoutMs = 10000
-    val ack = 0: Short
 
-    val emptyRequest = ProduceRequest.Builder.forCurrentMagic(ack, ackTimeoutMs,
-      new HashMap[TopicPartition, MemoryRecords]()).build()
+    val emptyRequest = requests.ProduceRequest.forCurrentMagic(new ProduceRequestData()
+      .setTopicData(new ProduceRequestData.TopicProduceDataCollection())
+      .setAcks(ack)
+      .setTimeoutMs(ackTimeoutMs)
+      .setTransactionalId(null))
+      .build()
     val emptyHeader = new RequestHeader(ApiKeys.PRODUCE, emptyRequest.version, clientId, correlationId)
-    val byteBuffer = emptyRequest.serialize(emptyHeader)
-    byteBuffer.rewind()
+    Utils.toArray(RequestTestUtils.serializeRequestWithHeader(emptyHeader, emptyRequest))
+  }
 
-    val serializedBytes = new Array[Byte](byteBuffer.remaining)
-    byteBuffer.get(serializedBytes)
-    serializedBytes
+  private def apiVersionRequestBytes(clientId: String, version: Short): Array[Byte] = {
+    val request = new ApiVersionsRequest.Builder().build(version)
+    val header = new RequestHeader(ApiKeys.API_VERSIONS, request.version(), clientId, -1)
+    Utils.toArray(RequestTestUtils.serializeRequestWithHeader(header, request))
   }
 
   @Test
-  def simpleRequest() {
-    val plainSocket = connect(protocol = SecurityProtocol.PLAINTEXT)
-    val serializedBytes = producerRequestBytes
+  def simpleRequest(): Unit = {
+    val plainSocket = connect()
+    val serializedBytes = producerRequestBytes()
 
     // Test PLAINTEXT socket
     sendRequest(plainSocket, serializedBytes)
-    processRequest(server.requestChannel)
+    processRequest(server.dataPlaneRequestChannel)
     assertEquals(serializedBytes.toSeq, receiveResponse(plainSocket).toSeq)
+    verifyAcceptorBlockedPercent("PLAINTEXT", expectBlocked = false)
+  }
+
+
+  private def testClientInformation(version: Short, expectedClientSoftwareName: String,
+                                    expectedClientSoftwareVersion: String): Unit = {
+    val plainSocket = connect()
+    val address = plainSocket.getLocalAddress
+    val clientId = "clientId"
+
+    // Send ApiVersionsRequest - unknown expected
+    sendRequest(plainSocket, apiVersionRequestBytes(clientId, version))
+    var receivedReq = receiveRequest(server.dataPlaneRequestChannel)
+
+    assertEquals(ClientInformation.UNKNOWN_NAME_OR_VERSION, receivedReq.context.clientInformation.softwareName)
+    assertEquals(ClientInformation.UNKNOWN_NAME_OR_VERSION, receivedReq.context.clientInformation.softwareVersion)
+
+    server.dataPlaneRequestChannel.sendResponse(new RequestChannel.NoOpResponse(receivedReq))
+
+    // Send ProduceRequest - client info expected
+    sendRequest(plainSocket, producerRequestBytes())
+    receivedReq = receiveRequest(server.dataPlaneRequestChannel)
+
+    assertEquals(expectedClientSoftwareName, receivedReq.context.clientInformation.softwareName)
+    assertEquals(expectedClientSoftwareVersion, receivedReq.context.clientInformation.softwareVersion)
+
+    server.dataPlaneRequestChannel.sendResponse(new RequestChannel.NoOpResponse(receivedReq))
+
+    // Close the socket
+    plainSocket.setSoLinger(true, 0)
+    plainSocket.close()
+
+    TestUtils.waitUntilTrue(() => server.connectionCount(address) == 0, msg = "Connection not closed")
   }
 
   @Test
-  def tooBigRequestIsRejected() {
+  def testClientInformationWithLatestApiVersionsRequest(): Unit = {
+    testClientInformation(
+      ApiKeys.API_VERSIONS.latestVersion,
+      "apache-kafka-java",
+      AppInfoParser.getVersion
+    )
+  }
+
+  @Test
+  def testClientInformationWithOldestApiVersionsRequest(): Unit = {
+    testClientInformation(
+      ApiKeys.API_VERSIONS.oldestVersion,
+      ClientInformation.UNKNOWN_NAME_OR_VERSION,
+      ClientInformation.UNKNOWN_NAME_OR_VERSION
+    )
+  }
+
+  @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(startProcessingRequests = 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 controlPlaneListenerStarted() = {
+      try {
+        val socket = connect(testableServer, config.controlPlaneListenerName.get, localAddr = InetAddress.getLocalHost)
+        sendAndReceiveControllerRequest(socket, testableServer)
+        true
+      } catch {
+        case _: Throwable => false
+      }
+    }
+
+    def listenerStarted(listenerName: ListenerName) = {
+      try {
+        val socket = connect(testableServer, listenerName, localAddr = InetAddress.getLocalHost)
+        sendAndReceiveRequest(socket, testableServer)
+        true
+      } catch {
+        case _: Throwable => false
+      }
+    }
+
+    try {
+      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((() => testableServer.startProcessingRequests(futures)): Runnable)
+      TestUtils.waitUntilTrue(() => controlPlaneListenerStarted(), "Control plane listener not started")
+      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 testStagedListenerShutdownWhenConnectionQueueIsFull(): 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 connectionQueueSize = 1
+    val testableServer = new TestableSocketServer(config, connectionQueueSize)
+    testableServer.startup(startProcessingRequests = false)
+
+    val socket1 = connect(testableServer, new ListenerName("EXTERNAL"), localAddr = InetAddress.getLocalHost)
+    sendRequest(socket1, producerRequestBytes())
+    val socket2 = connect(testableServer, new ListenerName("EXTERNAL"), localAddr = InetAddress.getLocalHost)
+    sendRequest(socket2, producerRequestBytes())
+
+    testableServer.shutdown()
+  }
+
+  @Test
+  def testDisabledRequestIsRejected(): Unit = {
+    val correlationId = 57
+    val header = new RequestHeader(ApiKeys.VOTE, 0, "", correlationId)
+    val request = new VoteRequest.Builder(new VoteRequestData()).build()
+    val serializedBytes = Utils.toArray(RequestTestUtils.serializeRequestWithHeader(header, request))
+
+    val socket = connect()
+
+    val outgoing = new DataOutputStream(socket.getOutputStream)
+    try {
+      outgoing.writeInt(serializedBytes.length)
+      outgoing.write(serializedBytes)
+      outgoing.flush()
+      receiveResponse(socket)
+    } catch {
+      case _: IOException => // we expect the server to close the socket
+    } finally {
+      outgoing.close()
+    }
+  }
+
+  @Test
+  def tooBigRequestIsRejected(): Unit = {
     val tooManyBytes = new Array[Byte](server.config.socketRequestMaxBytes + 1)
     new Random().nextBytes(tooManyBytes)
     val socket = connect()
@@ -192,28 +402,42 @@ class SocketServerTest extends JUnitSuite {
   }
 
   @Test
-  def testGracefulClose() {
-    val plainSocket = connect(protocol = SecurityProtocol.PLAINTEXT)
-    val serializedBytes = producerRequestBytes
+  def testGracefulClose(): Unit = {
+    val plainSocket = connect()
+    val serializedBytes = producerRequestBytes()
 
     for (_ <- 0 until 10)
       sendRequest(plainSocket, serializedBytes)
     plainSocket.close()
     for (_ <- 0 until 10) {
-      val request = receiveRequest(server.requestChannel)
+      val request = receiveRequest(server.dataPlaneRequestChannel)
       assertNotNull("receiveRequest timed out", request)
-      server.requestChannel.sendResponse(new RequestChannel.Response(request, None, RequestChannel.NoOpAction, None))
+      processRequestNoOpResponse(server.dataPlaneRequestChannel, request)
     }
   }
 
   @Test
-  def testConnectionId() {
-    val sockets = (1 to 5).map(_ => connect(protocol = SecurityProtocol.PLAINTEXT))
-    val serializedBytes = producerRequestBytes
+  def testNoOpAction(): Unit = {
+    val plainSocket = connect()
+    val serializedBytes = producerRequestBytes()
+
+    for (_ <- 0 until 3)
+      sendRequest(plainSocket, serializedBytes)
+    for (_ <- 0 until 3) {
+      val request = receiveRequest(server.dataPlaneRequestChannel)
+      assertNotNull("receiveRequest timed out", request)
+      processRequestNoOpResponse(server.dataPlaneRequestChannel, request)
+    }
+  }
+
+  @Test
+  def testConnectionId(): Unit = {
+    val sockets = (1 to 5).map(_ => connect())
+    val serializedBytes = producerRequestBytes()
 
     val requests = sockets.map{socket =>
       sendRequest(socket, serializedBytes)
-      receiveRequest(server.requestChannel)
+      receiveRequest(server.dataPlaneRequestChannel)
     }
     requests.zipWithIndex.foreach { case (request, i) =>
       val index = request.context.connectionId.split("-").last
@@ -224,44 +448,48 @@ class SocketServerTest extends JUnitSuite {
   }
 
   @Test
-  def testIdleConnection() {
+  def testIdleConnection(): Unit = {
     val idleTimeMs = 60000
     val time = new MockTime()
     props.put(KafkaConfig.ConnectionsMaxIdleMsProp, idleTimeMs.toString)
     val serverMetrics = new Metrics
     val overrideServer = new SocketServer(KafkaConfig.fromProps(props), serverMetrics, time, credentialProvider)
 
-    def openChannel(request: RequestChannel.Request): Option[KafkaChannel] =
-      overrideServer.processor(request.processor).channel(request.context.connectionId)
-    def openOrClosingChannel(request: RequestChannel.Request): Option[KafkaChannel] =
-      overrideServer.processor(request.processor).openOrClosingChannel(request.context.connectionId)
-
     try {
       overrideServer.startup()
-      val serializedBytes = producerRequestBytes
+      val serializedBytes = producerRequestBytes()
+
+      // Connection with no outstanding requests
+      val socket0 = connect(overrideServer)
+      sendRequest(socket0, serializedBytes)
+      val request0 = receiveRequest(overrideServer.dataPlaneRequestChannel)
+      processRequest(overrideServer.dataPlaneRequestChannel, request0)
+      assertTrue("Channel not open", openChannel(request0, overrideServer).nonEmpty)
+      assertEquals(openChannel(request0, overrideServer), openOrClosingChannel(request0, overrideServer))
+      TestUtils.waitUntilTrue(() => !openChannel(request0, overrideServer).get.isMuted, "Failed to unmute channel")
+      time.sleep(idleTimeMs + 1)
+      TestUtils.waitUntilTrue(() => openOrClosingChannel(request0, overrideServer).isEmpty, "Failed to close idle channel")
+      assertTrue("Channel not removed", openChannel(request0, overrideServer).isEmpty)
 
-      // Connection with no staged receives
-      val socket1 = connect(overrideServer, protocol = SecurityProtocol.PLAINTEXT)
+      // Connection with one request being processed (channel is muted), no other in-flight requests
+      val socket1 = connect(overrideServer)
       sendRequest(socket1, serializedBytes)
-      val request1 = receiveRequest(overrideServer.requestChannel)
-      assertTrue("Channel not open", openChannel(request1).nonEmpty)
-      assertEquals(openChannel(request1), openOrClosingChannel(request1))
-
+      val request1 = receiveRequest(overrideServer.dataPlaneRequestChannel)
+      assertTrue("Channel not open", openChannel(request1, overrideServer).nonEmpty)
+      assertEquals(openChannel(request1, overrideServer), openOrClosingChannel(request1, overrideServer))
       time.sleep(idleTimeMs + 1)
-      TestUtils.waitUntilTrue(() => openOrClosingChannel(request1).isEmpty, "Failed to close idle channel")
-      assertTrue("Channel not removed", openChannel(request1).isEmpty)
-      processRequest(overrideServer.requestChannel, request1)
-
-      // Connection with staged receives
-      val socket2 = connect(overrideServer, protocol = SecurityProtocol.PLAINTEXT)
-      val request2 = sendRequestsUntilStagedReceive(overrideServer, socket2, serializedBytes)
+      TestUtils.waitUntilTrue(() => openOrClosingChannel(request1, overrideServer).isEmpty, "Failed to close idle channel")
+      assertTrue("Channel not removed", openChannel(request1, overrideServer).isEmpty)
+      processRequest(overrideServer.dataPlaneRequestChannel, request1)
 
+      // Connection with one request being processed (channel is muted), more in-flight requests
+      val socket2 = connect(overrideServer)
+      val request2 = sendRequestsReceiveOne(overrideServer, socket2, serializedBytes, 3)
       time.sleep(idleTimeMs + 1)
-      TestUtils.waitUntilTrue(() => openChannel(request2).isEmpty, "Failed to close idle channel")
-      TestUtils.waitUntilTrue(() => openOrClosingChannel(request2).nonEmpty, "Channel removed without processing staged receives")
-      processRequest(overrideServer.requestChannel, request2) // this triggers a failed send since channel has been closed
-      TestUtils.waitUntilTrue(() => openOrClosingChannel(request2).isEmpty, "Failed to remove channel with failed sends")
-      assertNull("Received request after failed send", overrideServer.requestChannel.receiveRequest(200))
+      TestUtils.waitUntilTrue(() => openOrClosingChannel(request2, overrideServer).isEmpty, "Failed to close idle channel")
+      assertTrue("Channel not removed", openChannel(request1, overrideServer).isEmpty)
+      processRequest(overrideServer.dataPlaneRequestChannel, request2) // this triggers a failed send since channel has been closed
+      assertNull("Received request on expired channel", overrideServer.dataPlaneRequestChannel.receiveRequest(200))
 
     } finally {
       shutdownServerAndMetrics(overrideServer)
@@ -269,31 +497,32 @@ class SocketServerTest extends JUnitSuite {
   }
 
   @Test
-  def testConnectionIdReuse() {
+  def testConnectionIdReuse(): Unit = {
     val idleTimeMs = 60000
     val time = new MockTime()
     props.put(KafkaConfig.ConnectionsMaxIdleMsProp, idleTimeMs.toString)
-    props.put("listeners", "PLAINTEXT://localhost:0")
+    props ++= sslServerProps
     val serverMetrics = new Metrics
     @volatile var selector: TestableSelector = null
     val overrideConnectionId = "127.0.0.1:1-127.0.0.1:2-0"
     val overrideServer = new SocketServer(KafkaConfig.fromProps(props), serverMetrics, time, credentialProvider) {
-      override def newProcessor(id: Int, connectionQuotas: ConnectionQuotas, listenerName: ListenerName,
-                                protocol: SecurityProtocol, memoryPool: MemoryPool): Processor = {
-        new Processor(id, time, config.socketRequestMaxBytes, requestChannel, connectionQuotas,
-          config.connectionsMaxIdleMs, listenerName, protocol, config, metrics, credentialProvider, memoryPool, new LogContext()) {
+      override def newProcessor(id: Int, requestChannel: RequestChannel, connectionQuotas: ConnectionQuotas, listenerName: ListenerName,
+                                protocol: SecurityProtocol, memoryPool: MemoryPool, isPrivilegedListener: Boolean = false): Processor = {
+        new Processor(id, time, config.socketRequestMaxBytes, dataPlaneRequestChannel, connectionQuotas,
+          config.connectionsMaxIdleMs, config.failedAuthenticationDelayMs, listenerName, protocol, config, metrics,
+          credentialProvider, memoryPool, new LogContext(), isPrivilegedListener = isPrivilegedListener) {
           override protected[network] def connectionId(socket: Socket): String = overrideConnectionId
           override protected[network] def createSelector(channelBuilder: ChannelBuilder): Selector = {
-           val testableSelector = new TestableSelector(config, channelBuilder, time, metrics)
-           selector = testableSelector
-           testableSelector
-        }
+            val testableSelector = new TestableSelector(config, channelBuilder, time, metrics)
+            selector = testableSelector
+            testableSelector
+          }
         }
       }
     }
 
-    def openChannel: Option[KafkaChannel] = overrideServer.processor(0).channel(overrideConnectionId)
-    def openOrClosingChannel: Option[KafkaChannel] = overrideServer.processor(0).openOrClosingChannel(overrideConnectionId)
+    def openChannel: Option[KafkaChannel] = overrideServer.dataPlaneProcessor(0).channel(overrideConnectionId)
+    def openOrClosingChannel: Option[KafkaChannel] = overrideServer.dataPlaneProcessor(0).openOrClosingChannel(overrideConnectionId)
     def connectionCount = overrideServer.connectionCount(InetAddress.getByName("127.0.0.1"))
 
     // Create a client connection and wait for server to register the connection with the selector. For
@@ -301,7 +530,7 @@ class SocketServerTest extends JUnitSuite {
     // only after `register` is processed by the server.
     def connectAndWaitForConnectionRegister(): Socket = {
       val connections = selector.operationCounts(SelectorOperation.Register)
-      val socket = connect(overrideServer)
+      val socket = sslConnect(overrideServer)
       TestUtils.waitUntilTrue(() =>
         selector.operationCounts(SelectorOperation.Register) == connections + 1, "Connection not registered")
       socket
@@ -318,23 +547,21 @@ class SocketServerTest extends JUnitSuite {
       connectAndWaitForConnectionRegister()
       TestUtils.waitUntilTrue(() => connectionCount == 1, "Failed to close channel")
       assertSame(channel1, openChannel.getOrElse(throw new RuntimeException("Channel not found")))
+      socket1.close()
+      TestUtils.waitUntilTrue(() => openChannel.isEmpty, "Channel not closed")
 
-      // Send requests to `channel1` until a receive is staged and advance time beyond idle time so that `channel1` is
-      // closed with staged receives and is in Selector.closingChannels
-      val serializedBytes = producerRequestBytes
-      val request = sendRequestsUntilStagedReceive(overrideServer, socket1, serializedBytes)
-      time.sleep(idleTimeMs + 1)
-      TestUtils.waitUntilTrue(() => openChannel.isEmpty, "Idle channel not closed")
-      TestUtils.waitUntilTrue(() => openOrClosingChannel.isDefined, "Channel removed without processing staged receives")
+      // Create a channel with buffered receive and close remote connection
+      val request = makeChannelWithBufferedRequestsAndCloseRemote(overrideServer, selector)
+      val channel2 = openChannel.getOrElse(throw new RuntimeException("Channel not found"))
 
-      // Create new connection with same id when `channel1` is in Selector.closingChannels
-      // Check that new connection is closed and openOrClosingChannel still contains `channel1`
+      // Create new connection with same id when `channel2` is closing, but still in Selector.channels
+      // Check that new connection is closed and openOrClosingChannel still contains `channel2`
       connectAndWaitForConnectionRegister()
       TestUtils.waitUntilTrue(() => connectionCount == 1, "Failed to close channel")
-      assertSame(channel1, openOrClosingChannel.getOrElse(throw new RuntimeException("Channel not found")))
+      assertSame(channel2, openOrClosingChannel.getOrElse(throw new RuntimeException("Channel not found")))
 
-      // Complete request with failed send so that `channel1` is removed from Selector.closingChannels
-      processRequest(overrideServer.requestChannel, request)
+      // Complete request with failed send so that `channel2` is removed from Selector.channels
+      processRequest(overrideServer.dataPlaneRequestChannel, request)
       TestUtils.waitUntilTrue(() => connectionCount == 0 && openOrClosingChannel.isEmpty, "Failed to remove channel with failed send")
 
       // Check that new connections can be created with the same id since `channel1` is no longer in Selector
@@ -349,54 +576,205 @@ class SocketServerTest extends JUnitSuite {
     }
   }
 
-  private def sendRequestsUntilStagedReceive(server: SocketServer, socket: Socket, requestBytes: Array[Byte]): RequestChannel.Request = {
-    def sendTwoRequestsReceiveOne(): RequestChannel.Request = {
-      sendRequest(socket, requestBytes, flush = false)
-      sendRequest(socket, requestBytes, flush = true)
-      receiveRequest(server.requestChannel)
+  private def makeSocketWithBufferedRequests(server: SocketServer,
+                                             serverSelector: Selector,
+                                             proxyServer: ProxyServer,
+                                             numBufferedRequests: Int = 2): (Socket, RequestChannel.Request) = {
+
+    val requestBytes = producerRequestBytes()
+    val socket = sslClientSocket(proxyServer.localPort)
+    sendRequest(socket, requestBytes)
+    val request1 = receiveRequest(server.dataPlaneRequestChannel)
+
+    val connectionId = request1.context.connectionId
+    val channel = server.dataPlaneProcessor(0).channel(connectionId).getOrElse(throw new IllegalStateException("Channel not found"))
+    val transportLayer: SslTransportLayer = JTestUtils.fieldValue(channel, classOf[KafkaChannel], "transportLayer")
+    val netReadBuffer: ByteBuffer = JTestUtils.fieldValue(transportLayer, classOf[SslTransportLayer], "netReadBuffer")
+
+    proxyServer.enableBuffering(netReadBuffer)
+    (1 to numBufferedRequests).foreach { _ => sendRequest(socket, requestBytes) }
+
+    val keysWithBufferedRead: util.Set[SelectionKey] = JTestUtils.fieldValue(serverSelector, classOf[Selector], "keysWithBufferedRead")
+    keysWithBufferedRead.add(channel.selectionKey)
+    JTestUtils.setFieldValue(transportLayer, "hasBytesBuffered", true)
+
+    (socket, request1)
+  }
+
+  /**
+   * Create a channel with data in SSL buffers and close the remote connection.
+   * The channel should remain open in SocketServer even if it detects that the peer has closed
+   * the connection since there is pending data to be processed.
+   */
+  private def makeChannelWithBufferedRequestsAndCloseRemote(server: SocketServer,
+                                                            serverSelector: Selector,
+                                                            makeClosing: Boolean = false): RequestChannel.Request = {
+
+    val proxyServer = new ProxyServer(server)
+    try {
+      val (socket, request1) = makeSocketWithBufferedRequests(server, serverSelector, proxyServer)
+
+      socket.close()
+      proxyServer.serverConnSocket.close()
+      TestUtils.waitUntilTrue(() => proxyServer.clientConnSocket.isClosed, "Client socket not closed", waitTimeMs = 10000)
+
+      processRequestNoOpResponse(server.dataPlaneRequestChannel, request1)
+      val channel = openOrClosingChannel(request1, server).getOrElse(throw new IllegalStateException("Channel closed too early"))
+      if (makeClosing)
+        serverSelector.asInstanceOf[TestableSelector].pendingClosingChannels.add(channel)
+
+      receiveRequest(server.dataPlaneRequestChannel, timeout = 10000)
+    } finally {
+      proxyServer.close()
     }
-    val (request, hasStagedReceives) = TestUtils.computeUntilTrue(sendTwoRequestsReceiveOne()) { req =>
-      val connectionId = req.context.connectionId
-      val hasStagedReceives = server.processor(0).numStagedReceives(connectionId) > 0
-      if (!hasStagedReceives) {
-        processRequest(server.requestChannel, req)
-        processRequest(server.requestChannel)
+  }
+
+  def sendRequestsReceiveOne(server: SocketServer, socket: Socket, requestBytes: Array[Byte], numRequests: Int): RequestChannel.Request = {
+    (1 to numRequests).foreach(i => sendRequest(socket, requestBytes, flush = i == numRequests))
+    receiveRequest(server.dataPlaneRequestChannel)
+  }
+
+  private def closeSocketWithPendingRequest(server: SocketServer,
+                                            createSocket: () => Socket): RequestChannel.Request = {
+
+    def maybeReceiveRequest(): Option[RequestChannel.Request] = {
+      try {
+        Some(receiveRequest(server.dataPlaneRequestChannel, timeout = 1000))
+      } catch {
+        case e: Exception => None
+      }
+    }
+
+    def closedChannelWithPendingRequest(): Option[RequestChannel.Request] = {
+      val socket = createSocket.apply()
+      val req1 = sendRequestsReceiveOne(server, socket, producerRequestBytes(ack = 0), numRequests = 100)
+      processRequestNoOpResponse(server.dataPlaneRequestChannel, req1)
+      // Set SoLinger to 0 to force a hard disconnect via TCP RST
+      socket.setSoLinger(true, 0)
+      socket.close()
+
+      maybeReceiveRequest().flatMap { req =>
+        processRequestNoOpResponse(server.dataPlaneRequestChannel, req)
+        maybeReceiveRequest()
       }
-      hasStagedReceives
     }
-    assertTrue(s"Receives not staged for ${org.apache.kafka.test.TestUtils.DEFAULT_MAX_WAIT_MS} ms", hasStagedReceives)
+
+    val (request, _) = TestUtils.computeUntilTrue(closedChannelWithPendingRequest()) { req => req.nonEmpty }
+    request.getOrElse(throw new IllegalStateException("Could not create close channel with pending request"))
+  }
+
+  // Prepares test setup for throttled channel tests. throttlingDone controls whether or not throttling has completed
+  // in quota manager.
+  def throttledChannelTestSetUp(socket: Socket, serializedBytes: Array[Byte], noOpResponse: Boolean,
+                                throttlingInProgress: Boolean): RequestChannel.Request = {
+    sendRequest(socket, serializedBytes)
+
+    // Mimic a primitive request handler that fetches the request from RequestChannel and place a response with a
+    // throttled channel.
+    val request = receiveRequest(server.dataPlaneRequestChannel)
+    val byteBuffer = RequestTestUtils.serializeRequestWithHeader(request.header, request.body[AbstractRequest])
+    val send = new NetworkSend(request.context.connectionId, ByteBufferSend.sizePrefixed(byteBuffer))
+    def channelThrottlingCallback(response: RequestChannel.Response): Unit = {
+      server.dataPlaneRequestChannel.sendResponse(response)
+    }
+    val throttledChannel = new ThrottledChannel(request, new MockTime(), 100, channelThrottlingCallback)
+    val headerLog = RequestConvertToJson.requestHeaderNode(request.header)
+    val response =
+      if (!noOpResponse)
+        new RequestChannel.SendResponse(request, send, Some(headerLog), None)
+      else
+        new RequestChannel.NoOpResponse(request)
+    server.dataPlaneRequestChannel.sendResponse(response)
+
+    // Quota manager would call notifyThrottlingDone() on throttling completion. Simulate it if throttleingInProgress is
+    // false.
+    if (!throttlingInProgress)
+      throttledChannel.notifyThrottlingDone()
+
     request
   }
 
+  def openChannel(request: RequestChannel.Request, server: SocketServer = this.server): Option[KafkaChannel] =
+    server.dataPlaneProcessor(0).channel(request.context.connectionId)
+
+  def openOrClosingChannel(request: RequestChannel.Request, server: SocketServer = this.server): Option[KafkaChannel] =
+    server.dataPlaneProcessor(0).openOrClosingChannel(request.context.connectionId)
+
+  @Test
+  def testSendActionResponseWithThrottledChannelWhereThrottlingInProgress(): Unit = {
+    val socket = connect()
+    val serializedBytes = producerRequestBytes()
+    // SendAction with throttling in progress
+    val request = throttledChannelTestSetUp(socket, serializedBytes, false, true)
+
+    // receive response
+    assertEquals(serializedBytes.toSeq, receiveResponse(socket).toSeq)
+    TestUtils.waitUntilTrue(() => openOrClosingChannel(request).exists(c => c.muteState() == ChannelMuteState.MUTED_AND_THROTTLED), "fail")
+    // Channel should still be muted.
+    assertTrue(openOrClosingChannel(request).exists(c => c.isMuted()))
+  }
+
   @Test
-  def testSocketsCloseOnShutdown() {
+  def testSendActionResponseWithThrottledChannelWhereThrottlingAlreadyDone(): Unit = {
+    val socket = connect()
+    val serializedBytes = producerRequestBytes()
+    // SendAction with throttling in progress
+    val request = throttledChannelTestSetUp(socket, serializedBytes, false, false)
+
+    // receive response
+    assertEquals(serializedBytes.toSeq, receiveResponse(socket).toSeq)
+    // Since throttling is already done, the channel can be unmuted after sending out the response.
+    TestUtils.waitUntilTrue(() => openOrClosingChannel(request).exists(c => c.muteState() == ChannelMuteState.NOT_MUTED), "fail")
+    // Channel is now unmuted.
+    assertFalse(openOrClosingChannel(request).exists(c => c.isMuted()))
+  }
+
+  @Test
+  def testNoOpActionResponseWithThrottledChannelWhereThrottlingInProgress(): Unit = {
+    val socket = connect()
+    val serializedBytes = producerRequestBytes()
+    // SendAction with throttling in progress
+    val request = throttledChannelTestSetUp(socket, serializedBytes, true, true)
+
+    TestUtils.waitUntilTrue(() => openOrClosingChannel(request).exists(c => c.muteState() == ChannelMuteState.MUTED_AND_THROTTLED), "fail")
+    // Channel should still be muted.
+    assertTrue(openOrClosingChannel(request).exists(c => c.isMuted()))
+  }
+
+  @Test
+  def testNoOpActionResponseWithThrottledChannelWhereThrottlingAlreadyDone(): Unit = {
+    val socket = connect()
+    val serializedBytes = producerRequestBytes()
+    // SendAction with throttling in progress
+    val request = throttledChannelTestSetUp(socket, serializedBytes, true, false)
+
+    // Since throttling is already done, the channel can be unmuted.
+    TestUtils.waitUntilTrue(() => openOrClosingChannel(request).exists(c => c.muteState() == ChannelMuteState.NOT_MUTED), "fail")
+    // Channel is now unmuted.
+    assertFalse(openOrClosingChannel(request).exists(c => c.isMuted()))
+  }
+
+  @Test
+  def testSocketsCloseOnShutdown(): Unit = {
     // open a connection
-    val plainSocket = connect(protocol = SecurityProtocol.PLAINTEXT)
+    val plainSocket = connect()
     plainSocket.setTcpNoDelay(true)
     val bytes = new Array[Byte](40)
     // send a request first to make sure the connection has been picked up by the socket server
     sendRequest(plainSocket, bytes, Some(0))
-    processRequest(server.requestChannel)
+    processRequest(server.dataPlaneRequestChannel)
     // the following sleep is necessary to reliably detect the connection close when we send data below
     Thread.sleep(200L)
     // make sure the sockets are open
-    server.acceptors.values.foreach(acceptor => assertFalse(acceptor.serverChannel.socket.isClosed))
+    server.dataPlaneAcceptors.asScala.values.foreach(acceptor => assertFalse(acceptor.serverChannel.socket.isClosed))
     // then shutdown the server
     shutdownServerAndMetrics(server)
 
-    val largeChunkOfBytes = new Array[Byte](1000000)
-    // doing a subsequent send should throw an exception as the connection should be closed.
-    // send a large chunk of bytes to trigger a socket flush
-    try {
-      sendRequest(plainSocket, largeChunkOfBytes, Some(0))
-      fail("expected exception when writing to closed plain socket")
-    } catch {
-      case _: IOException => // expected
-    }
+    verifyRemoteConnectionClosed(plainSocket)
   }
 
   @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)
@@ -411,14 +789,51 @@ class SocketServerTest extends JUnitSuite {
     TestUtils.waitUntilTrue(() => server.connectionCount(address) < conns.length,
       "Failed to decrement connection count after close")
     val conn2 = connect()
-    val serializedBytes = producerRequestBytes
+    val serializedBytes = producerRequestBytes()
     sendRequest(conn2, serializedBytes)
-    val request = server.requestChannel.receiveRequest(2000)
+    val request = server.dataPlaneRequestChannel.receiveRequest(2000)
     assertNotNull(request)
   }
 
   @Test
-  def testMaxConnectionsPerIpOverrides() {
+  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"))
+    val server = new SocketServer(KafkaConfig.fromProps(newProps), new Metrics(), Time.SYSTEM, credentialProvider)
+    try {
+      server.startup()
+      // make the maximum allowable number of connections
+      val conns = (0 until 5).map(_ => connect(server))
+      // now try one more (should fail)
+      val conn = connect(server)
+      conn.setSoTimeout(3000)
+      assertEquals(-1, conn.getInputStream.read())
+      conn.close()
+
+      // it should succeed after closing one connection
+      val address = conns.head.getInetAddress
+      conns.head.close()
+      TestUtils.waitUntilTrue(() => server.connectionCount(address) < conns.length,
+        "Failed to decrement connection count after close")
+      val conn2 = connect(server)
+      val serializedBytes = producerRequestBytes()
+      sendRequest(conn2, serializedBytes)
+      val request = server.dataPlaneRequestChannel.receiveRequest(2000)
+      assertNotNull(request)
+
+      // now try to connect from the external facing interface, which should fail
+      val conn3 = connect(s = server, localAddr = InetAddress.getLocalHost)
+      conn3.setSoTimeout(3000)
+      assertEquals(-1, conn3.getInputStream.read())
+      conn3.close()
+    } finally {
+      shutdownServerAndMetrics(server)
+    }
+  }
+
+  @Test
+  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")
@@ -430,9 +845,9 @@ class SocketServerTest extends JUnitSuite {
       val conns = (0 until overrideNum).map(_ => connect(overrideServer))
 
       // it should succeed
-      val serializedBytes = producerRequestBytes
+      val serializedBytes = producerRequestBytes()
       sendRequest(conns.last, serializedBytes)
-      val request = overrideServer.requestChannel.receiveRequest(2000)
+      val request = overrideServer.dataPlaneRequestChannel.receiveRequest(2000)
       assertNotNull(request)
 
       // now try one more (should fail)
@@ -445,17 +860,94 @@ class SocketServerTest extends JUnitSuite {
   }
 
   @Test
-  def testSslSocketServer() {
-    val trustStoreFile = File.createTempFile("truststore", ".jks")
-    val overrideProps = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, interBrokerSecurityProtocol = Some(SecurityProtocol.SSL),
-      trustStoreFile = Some(trustStoreFile))
-    overrideProps.put(KafkaConfig.ListenersProp, "SSL://localhost:0")
+  def testConnectionRatePerIp(): Unit = {
+    val defaultTimeoutMs = 2000
+    val overrideProps = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 0)
+    overrideProps.remove(KafkaConfig.MaxConnectionsPerIpProp)
+    overrideProps.put(KafkaConfig.NumQuotaSamplesProp, String.valueOf(2))
+    val connectionRate = 5
+    val time = new MockTime()
+    val overrideServer = new SocketServer(KafkaConfig.fromProps(overrideProps), new Metrics(), time, credentialProvider)
+    // update the connection rate to 5
+    overrideServer.connectionQuotas.updateIpConnectionRateQuota(None, Some(connectionRate))
+    try {
+      overrideServer.startup()
+      // make the (maximum allowable number + 1) of connections
+      (0 to connectionRate).map(_ => connect(overrideServer))
+
+      val acceptors = overrideServer.dataPlaneAcceptors.asScala.values
+      // waiting for 5 connections got accepted and 1 connection got throttled
+      TestUtils.waitUntilTrue(
+        () => acceptors.foldLeft(0)((accumulator, acceptor) => accumulator + acceptor.throttledSockets.size) == 1,
+        "timeout waiting for 1 connection to get throttled",
+        defaultTimeoutMs)
+
+      // now try one more, so that we can make sure this connection will get throttled
+      var conn = connect(overrideServer)
+      // there should be total 2 connection got throttled now
+      TestUtils.waitUntilTrue(
+        () => acceptors.foldLeft(0)((accumulator, acceptor) => accumulator + acceptor.throttledSockets.size) == 2,
+        "timeout waiting for 2 connection to get throttled",
+        defaultTimeoutMs)
+      // advance time to unthrottle connections
+      time.sleep(defaultTimeoutMs)
+      acceptors.foreach(_.wakeup())
+      // make sure there are no connection got throttled now(and the throttled connections should be closed)
+      TestUtils.waitUntilTrue(() => acceptors.forall(_.throttledSockets.isEmpty),
+        "timeout waiting for connection to be unthrottled",
+        defaultTimeoutMs)
+      // verify the connection is closed now
+      verifyRemoteConnectionClosed(conn)
+
+      // new connection should succeed after previous connection closed, and previous samples have been expired
+      conn = connect(overrideServer)
+      val serializedBytes = producerRequestBytes()
+      sendRequest(conn, serializedBytes)
+      val request = overrideServer.dataPlaneRequestChannel.receiveRequest(defaultTimeoutMs)
+      assertNotNull(request)
+    } finally {
+      shutdownServerAndMetrics(overrideServer)
+    }
+  }
+
+  @Test
+  def testThrottledSocketsClosedOnShutdown(): Unit = {
+    val overrideProps = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 0)
+    overrideProps.remove("max.connections.per.ip")
+    overrideProps.put(KafkaConfig.NumQuotaSamplesProp, String.valueOf(2))
+    val connectionRate = 5
+    val time = new MockTime()
+    val overrideServer = new SocketServer(KafkaConfig.fromProps(overrideProps), new Metrics(), time, credentialProvider)
+    overrideServer.connectionQuotas.updateIpConnectionRateQuota(None, Some(connectionRate))
+    overrideServer.startup()
+    // make the maximum allowable number of connections
+    (0 until connectionRate).map(_ => connect(overrideServer))
+    // now try one more (should get throttled)
+    val conn = connect(overrideServer)
+    // don't advance time so that connection never gets unthrottled
+    shutdownServerAndMetrics(overrideServer)
+    verifyRemoteConnectionClosed(conn)
+  }
 
+  private def verifyRemoteConnectionClosed(connection: Socket): Unit = {
+    val largeChunkOfBytes = new Array[Byte](1000000)
+    // doing a subsequent send should throw an exception as the connection should be closed.
+    // send a large chunk of bytes to trigger a socket flush
+    try {
+      sendRequest(connection, largeChunkOfBytes, Some(0))
+      fail("expected exception when writing to closed plain socket")
+    } catch {
+      case _: IOException => // expected
+    }
+  }
+
+  @Test
+  def testSslSocketServer(): Unit = {
     val serverMetrics = new Metrics
-    val overrideServer = new SocketServer(KafkaConfig.fromProps(overrideProps), serverMetrics, Time.SYSTEM, credentialProvider)
+    val overrideServer = new SocketServer(KafkaConfig.fromProps(sslServerProps), serverMetrics, Time.SYSTEM, credentialProvider)
     try {
       overrideServer.startup()
-      val sslContext = SSLContext.getInstance("TLSv1.2")
+      val sslContext = SSLContext.getInstance(TestSslUtils.DEFAULT_TLS_PROTOCOL_FOR_TESTS)
       sslContext.init(null, Array(TestUtils.trustAllCerts), new java.security.SecureRandom())
       val socketFactory = sslContext.getSocketFactory
       val sslSocket = socketFactory.createSocket("localhost",
@@ -466,17 +958,17 @@ class SocketServerTest extends JUnitSuite {
       val clientId = ""
       val ackTimeoutMs = 10000
       val ack = 0: Short
-      val emptyRequest = ProduceRequest.Builder.forCurrentMagic(ack, ackTimeoutMs,
-        new HashMap[TopicPartition, MemoryRecords]()).build()
+      val emptyRequest = requests.ProduceRequest.forCurrentMagic(new ProduceRequestData()
+        .setTopicData(new ProduceRequestData.TopicProduceDataCollection())
+        .setAcks(ack)
+        .setTimeoutMs(ackTimeoutMs)
+        .setTransactionalId(null))
+        .build()
       val emptyHeader = new RequestHeader(ApiKeys.PRODUCE, emptyRequest.version, clientId, correlationId)
-
-      val byteBuffer = emptyRequest.serialize(emptyHeader)
-      byteBuffer.rewind()
-      val serializedBytes = new Array[Byte](byteBuffer.remaining)
-      byteBuffer.get(serializedBytes)
+      val serializedBytes = Utils.toArray(RequestTestUtils.serializeRequestWithHeader(emptyHeader, emptyRequest))
 
       sendRequest(sslSocket, serializedBytes)
-      processRequest(overrideServer.requestChannel)
+      processRequest(overrideServer.dataPlaneRequestChannel)
       assertEquals(serializedBytes.toSeq, receiveResponse(sslSocket).toSeq)
       sslSocket.close()
     } finally {
@@ -485,25 +977,120 @@ class SocketServerTest extends JUnitSuite {
   }
 
   @Test
-  def testSessionPrincipal() {
+  def testSaslReauthenticationFailureWithKip152SaslAuthenticate(): Unit = {
+    checkSaslReauthenticationFailure(true)
+  }
+
+  @Test
+  def testSaslReauthenticationFailureNoKip152SaslAuthenticate(): Unit = {
+    checkSaslReauthenticationFailure(false)
+  }
+
+  def checkSaslReauthenticationFailure(leverageKip152SaslAuthenticateRequest : Boolean): 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 version : Short = if (leverageKip152SaslAuthenticateRequest) ApiKeys.SASL_HANDSHAKE.latestVersion else 0
+      val saslHandshakeRequest = new SaslHandshakeRequest.Builder(new SaslHandshakeRequestData().setMechanism("PLAIN"))
+        .build(version)
+      val saslHandshakeHeader = new RequestHeader(ApiKeys.SASL_HANDSHAKE, saslHandshakeRequest.version, clientId,
+        correlationId)
+      sendApiRequest(socket, saslHandshakeRequest, saslHandshakeHeader)
+      receiveResponse(socket)
+
+      // now send credentials
+      val authBytes = "admin\u0000admin\u0000admin-secret".getBytes(StandardCharsets.UTF_8)
+      if (leverageKip152SaslAuthenticateRequest) {
+        // send credentials within a SaslAuthenticateRequest
+        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)
+      } else {
+        // send credentials directly, without a SaslAuthenticateRequest
+        sendRequest(socket, authBytes)
+      }
+      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 = requests.ProduceRequest.forCurrentMagic(new ProduceRequestData()
+        .setTopicData(new ProduceRequestData.TopicProduceDataCollection())
+        .setAcks(ack)
+        .setTimeoutMs(ackTimeoutMs)
+        .setTransactionalId(null))
+        .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))
-    assertEquals(KafkaPrincipal.ANONYMOUS, receiveRequest(server.requestChannel).session.principal)
+    assertEquals(KafkaPrincipal.ANONYMOUS, receiveRequest(server.dataPlaneRequestChannel).session.principal)
   }
 
   /* 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 = {
+    // The way we detect a connection close from the client depends on the response size. If it's small, an
+    // IOException ("Connection reset by peer") is thrown when the Selector reads from the socket. If
+    // it's large, an IOException ("Broken pipe") is thrown when the Selector writes to the socket. We test
+    // both paths to ensure they are handled correctly.
+    checkClientDisconnectionUpdatesRequestMetrics(0)
+    checkClientDisconnectionUpdatesRequestMetrics(550000)
+  }
+
+  private def checkClientDisconnectionUpdatesRequestMetrics(responseBufferSize: Int): Unit = {
     val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 0)
     val serverMetrics = new Metrics
     var conn: Socket = null
     val overrideServer = new SocketServer(KafkaConfig.fromProps(props), serverMetrics, Time.SYSTEM, credentialProvider) {
-      override def newProcessor(id: Int, connectionQuotas: ConnectionQuotas, listenerName: ListenerName,
-                                protocol: SecurityProtocol, memoryPool: MemoryPool): Processor = {
-        new Processor(id, time, config.socketRequestMaxBytes, requestChannel, connectionQuotas,
-          config.connectionsMaxIdleMs, listenerName, protocol, config, metrics, credentialProvider, MemoryPool.NONE, new LogContext()) {
-          override protected[network] def sendResponse(response: RequestChannel.Response, responseSend: Send) {
+      override def newProcessor(id: Int, requestChannel: RequestChannel, connectionQuotas: ConnectionQuotas, listenerName: ListenerName,
+                                protocol: SecurityProtocol, memoryPool: MemoryPool, isPrivilegedListener: Boolean = false): Processor = {
+        new Processor(id, time, config.socketRequestMaxBytes, dataPlaneRequestChannel, connectionQuotas,
+          config.connectionsMaxIdleMs, config.failedAuthenticationDelayMs, listenerName, protocol, config, metrics,
+          credentialProvider, MemoryPool.NONE, new LogContext(), isPrivilegedListener = isPrivilegedListener) {
+          override protected[network] def sendResponse(response: RequestChannel.Response, responseSend: Send): Unit = {
             conn.close()
             super.sendResponse(response, responseSend)
           }
@@ -513,23 +1100,20 @@ class SocketServerTest extends JUnitSuite {
     try {
       overrideServer.startup()
       conn = connect(overrideServer)
-      val serializedBytes = producerRequestBytes
+      val serializedBytes = producerRequestBytes()
       sendRequest(conn, serializedBytes)
 
-      val channel = overrideServer.requestChannel
+      val channel = overrideServer.dataPlaneRequestChannel
       val request = receiveRequest(channel)
 
       val requestMetrics = channel.metrics(request.header.apiKey.name)
       def totalTimeHistCount(): Long = requestMetrics.totalTimeHist.count
-      val expectedTotalTimeCount = totalTimeHistCount() + 1
+      val send = new NetworkSend(request.context.connectionId, ByteBufferSend.sizePrefixed(ByteBuffer.allocate(responseBufferSize)))
+      val headerLog = new ObjectNode(JsonNodeFactory.instance)
+      headerLog.set("response", new TextNode("someResponse"))
+      channel.sendResponse(new RequestChannel.SendResponse(request, send, Some(headerLog), None))
 
-      // send a large buffer to ensure that the broker detects the client disconnection while writing to the socket channel.
-      // On Mac OS X, the initial write seems to always succeed and it is able to write up to 102400 bytes on the initial
-      // write. If the buffer is smaller than this, the write is considered complete and the disconnection is not
-      // detected. If the buffer is larger than 102400 bytes, a second write is attempted and it fails with an
-      // IOException.
-      val send = new NetworkSend(request.context.connectionId, ByteBuffer.allocate(550000))
-      channel.sendResponse(new RequestChannel.Response(request, Some(send), SendAction, None))
+      val expectedTotalTimeCount = totalTimeHistCount() + 1
       TestUtils.waitUntilTrue(() => totalTimeHistCount() == expectedTotalTimeCount,
         s"request metrics not updated, expected: $expectedTotalTimeCount, actual: ${totalTimeHistCount()}")
 
@@ -538,26 +1122,61 @@ class SocketServerTest extends JUnitSuite {
     }
   }
 
+  @Test
+  def testClientDisconnectionWithOutstandingReceivesProcessedUntilFailedSend(): Unit = {
+    val serverMetrics = new Metrics
+    @volatile var selector: TestableSelector = null
+    val overrideServer = new SocketServer(KafkaConfig.fromProps(props), serverMetrics, Time.SYSTEM, credentialProvider) {
+      override def newProcessor(id: Int, requestChannel: RequestChannel, connectionQuotas: ConnectionQuotas, listenerName: ListenerName,
+                                protocol: SecurityProtocol, memoryPool: MemoryPool, isPrivilegedListener: Boolean = false): Processor = {
+        new Processor(id, time, config.socketRequestMaxBytes, dataPlaneRequestChannel, connectionQuotas,
+          config.connectionsMaxIdleMs, config.failedAuthenticationDelayMs, listenerName, protocol, config, metrics,
+          credentialProvider, memoryPool, new LogContext(), isPrivilegedListener = isPrivilegedListener) {
+          override protected[network] def createSelector(channelBuilder: ChannelBuilder): Selector = {
+            val testableSelector = new TestableSelector(config, channelBuilder, time, metrics)
+            selector = testableSelector
+            testableSelector
+          }
+        }
+      }
+    }
+
+    try {
+      overrideServer.startup()
+
+      // Create a channel, send some requests and close socket. Receive one pending request after socket was closed.
+      val request = closeSocketWithPendingRequest(overrideServer, () => connect(overrideServer))
+
+      // Complete request with socket exception so that the channel is closed
+      processRequest(overrideServer.dataPlaneRequestChannel, request)
+      TestUtils.waitUntilTrue(() => openOrClosingChannel(request, overrideServer).isEmpty, "Channel not closed after failed send")
+      assertTrue("Unexpected completed send", selector.completedSends.isEmpty)
+    } finally {
+      overrideServer.shutdown()
+      serverMetrics.close()
+    }
+  }
+
   /*
    * Test that we update request metrics if the channel has been removed from the selector when the broker calls
    * `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, "100")
+    props.setProperty(KafkaConfig.ConnectionsMaxIdleMsProp, "110")
     val serverMetrics = new Metrics
     var conn: Socket = null
     val overrideServer = new SocketServer(KafkaConfig.fromProps(props), serverMetrics, Time.SYSTEM, credentialProvider)
     try {
       overrideServer.startup()
       conn = connect(overrideServer)
-      val serializedBytes = producerRequestBytes
+      val serializedBytes = producerRequestBytes()
       sendRequest(conn, serializedBytes)
-      val channel = overrideServer.requestChannel
+      val channel = overrideServer.dataPlaneRequestChannel
       val request = receiveRequest(channel)
 
-      TestUtils.waitUntilTrue(() => overrideServer.processor(request.processor).channel(request.context.connectionId).isEmpty,
+      TestUtils.waitUntilTrue(() => overrideServer.dataPlaneProcessor(request.processor).channel(request.context.connectionId).isEmpty,
         s"Idle connection `${request.context.connectionId}` was not closed by selector")
 
       val requestMetrics = channel.metrics(request.header.apiKey.name)
@@ -577,17 +1196,20 @@ class SocketServerTest extends JUnitSuite {
   @Test
   def testRequestMetricsAfterStop(): Unit = {
     server.stopProcessingRequests()
-
-    server.requestChannel.metrics(ApiKeys.PRODUCE.name).requestRate.mark()
-    server.requestChannel.updateErrorMetrics(ApiKeys.PRODUCE, Map(Errors.NONE -> 1))
-    val nonZeroMeters = Map("kafka.network:type=RequestMetrics,name=RequestsPerSec,request=Produce" -> 1,
-        "kafka.network:type=RequestMetrics,name=ErrorsPerSec,request=Produce,error=NONE" -> 1)
-
-    def requestMetricMeters = YammerMetrics
+    val version = ApiKeys.PRODUCE.latestVersion
+    val version2 = (version - 1).toShort
+    for (_ <- 0 to 1) server.dataPlaneRequestChannel.metrics(ApiKeys.PRODUCE.name).requestRate(version).mark()
+    server.dataPlaneRequestChannel.metrics(ApiKeys.PRODUCE.name).requestRate(version2).mark()
+    assertEquals(2, server.dataPlaneRequestChannel.metrics(ApiKeys.PRODUCE.name).requestRate(version).count())
+    server.dataPlaneRequestChannel.updateErrorMetrics(ApiKeys.PRODUCE, Map(Errors.NONE -> 1))
+    val nonZeroMeters = Map(s"kafka.network:type=RequestMetrics,name=RequestsPerSec,request=Produce,version=$version" -> 2,
+      s"kafka.network:type=RequestMetrics,name=RequestsPerSec,request=Produce,version=$version2" -> 1,
+      "kafka.network:type=RequestMetrics,name=ErrorsPerSec,request=Produce,error=NONE" -> 1)
+
+    def requestMetricMeters = KafkaYammerMetrics
       .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()
@@ -598,12 +1220,12 @@ class SocketServerTest extends JUnitSuite {
   def testMetricCollectionAfterShutdown(): Unit = {
     server.shutdown()
 
-    val nonZeroMetricNamesAndValues = YammerMetrics
+    val nonZeroMetricNamesAndValues = KafkaYammerMetrics
       .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 }
+      .filter { case (_, value) => value != 0.0 && !value.equals(Double.NaN) }
 
     assertEquals(Map.empty, nonZeroMetricNamesAndValues)
   }
@@ -613,14 +1235,14 @@ class SocketServerTest extends JUnitSuite {
     val kafkaMetricNames = metrics.metrics.keySet.asScala.filter(_.tags.asScala.get("listener").nonEmpty)
     assertFalse(kafkaMetricNames.isEmpty)
 
-    val expectedListeners = Set("PLAINTEXT", "TRACE")
+    val expectedListeners = Set("PLAINTEXT")
     kafkaMetricNames.foreach { kafkaMetricName =>
       assertTrue(expectedListeners.contains(kafkaMetricName.tags.get("listener")))
     }
 
     // legacy metrics not tagged
-    val yammerMetricsNames = YammerMetrics.defaultRegistry.allMetrics.asScala
-      .filterKeys(_.getType.equals("Processor"))
+    val yammerMetricsNames = KafkaYammerMetrics.defaultRegistry.allMetrics.asScala
+      .filter { case (k, _) => k.getType.equals("Processor") }
       .collect { case (k, _: Gauge[_]) => k }
     assertFalse(yammerMetricsNames.isEmpty)
 
@@ -641,7 +1263,7 @@ class SocketServerTest extends JUnitSuite {
    */
   @Test
   def configureNewConnectionException(): Unit = {
-    withTestableServer { testableServer =>
+    withTestableServer (testWithServer = { testableServer =>
       val testableSelector = testableServer.testableSelector
 
       testableSelector.updateMinWakeup(2)
@@ -651,7 +1273,7 @@ class SocketServerTest extends JUnitSuite {
       TestUtils.waitUntilTrue(() => testableServer.connectionCount(localAddress) == 1, "Failed channel not removed")
 
       assertProcessorHealthy(testableServer, testableSelector.notFailed(sockets))
-    }
+    })
   }
 
   /**
@@ -666,20 +1288,20 @@ class SocketServerTest extends JUnitSuite {
    */
   @Test
   def processNewResponseException(): Unit = {
-    withTestableServer { testableServer =>
+    withTestableServer (testWithServer = { testableServer =>
       val testableSelector = testableServer.testableSelector
       testableSelector.updateMinWakeup(2)
 
       val sockets = (1 to 2).map(_ => connect(testableServer))
-      sockets.foreach(sendRequest(_, producerRequestBytes))
+      sockets.foreach(sendRequest(_, producerRequestBytes()))
 
       testableServer.testableSelector.addFailure(SelectorOperation.Send)
-      sockets.foreach(_ => processRequest(testableServer.requestChannel))
+      sockets.foreach(_ => processRequest(testableServer.dataPlaneRequestChannel))
       testableSelector.waitForOperations(SelectorOperation.Send, 2)
       testableServer.waitForChannelClose(testableSelector.allFailedChannels.head, locallyClosed = true)
 
       assertProcessorHealthy(testableServer, testableSelector.notFailed(sockets))
-    }
+    })
   }
 
   /**
@@ -689,13 +1311,13 @@ class SocketServerTest extends JUnitSuite {
    */
   @Test
   def sendCancelledKeyException(): Unit = {
-    withTestableServer { testableServer =>
+    withTestableServer (testWithServer = { testableServer =>
       val testableSelector = testableServer.testableSelector
       testableSelector.updateMinWakeup(2)
 
       val sockets = (1 to 2).map(_ => connect(testableServer))
-      sockets.foreach(sendRequest(_, producerRequestBytes))
-      val requestChannel = testableServer.requestChannel
+      sockets.foreach(sendRequest(_, producerRequestBytes()))
+      val requestChannel = testableServer.dataPlaneRequestChannel
 
       val requests = sockets.map(_ => receiveRequest(requestChannel))
       val failedConnectionId = requests(0).context.connectionId
@@ -707,33 +1329,239 @@ class SocketServerTest extends JUnitSuite {
 
       val successfulSocket = if (isSocketConnectionId(failedConnectionId, sockets(0))) sockets(1) else sockets(0)
       assertProcessorHealthy(testableServer, Seq(successfulSocket))
+    })
+  }
+
+  /**
+   * Tests channel send failure handling when send failure is triggered by [[Selector.send]]
+   * to a channel whose peer has closed its connection.
+   */
+  @Test
+  def remoteCloseSendFailure(): Unit = {
+    verifySendFailureAfterRemoteClose(makeClosing = false)
+  }
+
+  /**
+   * Tests channel send failure handling when send failure is triggered by [[Selector.send]]
+   * to a channel whose peer has closed its connection and the channel is in `closingChannels`.
+   */
+  @Test
+  def closingChannelSendFailure(): Unit = {
+    verifySendFailureAfterRemoteClose(makeClosing = true)
+  }
+
+  private def verifySendFailureAfterRemoteClose(makeClosing: Boolean): Unit = {
+    props ++= sslServerProps
+    withTestableServer (testWithServer = { testableServer =>
+      val testableSelector = testableServer.testableSelector
+
+      val serializedBytes = producerRequestBytes()
+      val request = makeChannelWithBufferedRequestsAndCloseRemote(testableServer, testableSelector, makeClosing)
+      val otherSocket = sslConnect(testableServer)
+      sendRequest(otherSocket, serializedBytes)
+
+      processRequest(testableServer.dataPlaneRequestChannel, request)
+      processRequest(testableServer.dataPlaneRequestChannel) // Also process request from other socket
+      testableSelector.waitForOperations(SelectorOperation.Send, 2)
+      testableServer.waitForChannelClose(request.context.connectionId, locallyClosed = false)
+
+      assertProcessorHealthy(testableServer, Seq(otherSocket))
+    })
+  }
+
+  /**
+   * Verifies that all pending buffered receives are processed even if remote connection is closed.
+   * The channel must be closed after pending receives are processed.
+   */
+  @Test
+  def remoteCloseWithBufferedReceives(): Unit = {
+    verifyRemoteCloseWithBufferedReceives(numComplete = 3, hasIncomplete = false)
+  }
+
+  /**
+   * Verifies that channel is closed when remote client closes its connection if there is no
+   * buffered receive.
+   */
+  @Test
+  def remoteCloseWithoutBufferedReceives(): Unit = {
+    verifyRemoteCloseWithBufferedReceives(numComplete = 0, hasIncomplete = false)
+  }
+
+  /**
+   * Verifies that channel is closed when remote client closes its connection if there is a pending
+   * receive that is incomplete.
+   */
+  @Test
+  def remoteCloseWithIncompleteBufferedReceive(): Unit = {
+    verifyRemoteCloseWithBufferedReceives(numComplete = 0, hasIncomplete = true)
+  }
+
+  /**
+   * Verifies that all pending buffered receives are processed even if remote connection is closed.
+   * The channel must be closed after complete receives are processed, even if there is an incomplete
+   * receive remaining in the buffers.
+   */
+  @Test
+  def remoteCloseWithCompleteAndIncompleteBufferedReceives(): Unit = {
+    verifyRemoteCloseWithBufferedReceives(numComplete = 3, hasIncomplete = true)
+  }
+
+  /**
+   * Verifies that pending buffered receives are processed when remote connection is closed
+   * until a response send fails.
+   */
+  @Test
+  def remoteCloseWithBufferedReceivesFailedSend(): Unit = {
+    verifyRemoteCloseWithBufferedReceives(numComplete = 3, hasIncomplete = false, responseRequiredIndex = 1)
+  }
+
+  /**
+   * Verifies that all pending buffered receives are processed for channel in closing state.
+   * The channel must be closed after pending receives are processed.
+   */
+  @Test
+  def closingChannelWithBufferedReceives(): Unit = {
+    verifyRemoteCloseWithBufferedReceives(numComplete = 3, hasIncomplete = false, makeClosing = true)
+  }
+
+  /**
+   * Verifies that all pending buffered receives are processed for channel in closing state.
+   * The channel must be closed after complete receives are processed, even if there is an incomplete
+   * receive remaining in the buffers.
+   */
+  @Test
+  def closingChannelWithCompleteAndIncompleteBufferedReceives(): Unit = {
+    verifyRemoteCloseWithBufferedReceives(numComplete = 3, hasIncomplete = true, makeClosing = false)
+  }
+
+  /**
+   * Verifies that pending buffered receives are processed for a channel in closing state
+   * until a response send fails.
+   */
+  @Test
+  def closingChannelWithBufferedReceivesFailedSend(): Unit = {
+    verifyRemoteCloseWithBufferedReceives(numComplete = 3, hasIncomplete = false, responseRequiredIndex = 1, makeClosing = false)
+  }
+
+  /**
+   * Verifies handling of client disconnections when the server-side channel is in the state
+   * specified using the parameters.
+   *
+   * @param numComplete Number of complete buffered requests
+   * @param hasIncomplete If true, add an additional partial buffered request
+   * @param responseRequiredIndex Index of the buffered request for which a response is sent. Previous requests
+   *                              are completed without a response. If set to -1, all `numComplete` requests
+   *                              are completed without a response.
+   * @param makeClosing If true, put the channel into closing state in the server Selector.
+   */
+  private def verifyRemoteCloseWithBufferedReceives(numComplete: Int,
+                                                    hasIncomplete: Boolean,
+                                                    responseRequiredIndex: Int = -1,
+                                                    makeClosing: Boolean = false): Unit = {
+    props ++= sslServerProps
+
+    // Truncates the last request in the SSL buffers by directly updating the buffers to simulate partial buffered request
+    def truncateBufferedRequest(channel: KafkaChannel): Unit = {
+      val transportLayer: SslTransportLayer = JTestUtils.fieldValue(channel, classOf[KafkaChannel], "transportLayer")
+      val netReadBuffer: ByteBuffer = JTestUtils.fieldValue(transportLayer, classOf[SslTransportLayer], "netReadBuffer")
+      val appReadBuffer: ByteBuffer = JTestUtils.fieldValue(transportLayer, classOf[SslTransportLayer], "appReadBuffer")
+      if (appReadBuffer.position() > 4) {
+        appReadBuffer.position(4)
+        netReadBuffer.position(0)
+      } else {
+        netReadBuffer.position(20)
+      }
     }
+    withTestableServer (testWithServer = { testableServer =>
+      val testableSelector = testableServer.testableSelector
+
+      val proxyServer = new ProxyServer(testableServer)
+      try {
+        // Step 1: Send client requests.
+        //   a) request1 is sent by the client to ProxyServer and this is directly sent to the server. This
+        //      ensures that server-side channel is in muted state until this request is processed in Step 3.
+        //   b) `numComplete` requests are sent and buffered in the server-side channel's SSL buffers
+        //   c) If `hasIncomplete=true`, an extra request is sent and buffered as in b). This will be truncated later
+        //      when previous requests have been processed and only one request is remaining in the SSL buffer,
+        //      making it easy to truncate.
+        val numBufferedRequests = numComplete + (if (hasIncomplete) 1 else 0)
+        val (socket, request1) = makeSocketWithBufferedRequests(testableServer, testableSelector, proxyServer, numBufferedRequests)
+        val channel = openChannel(request1, testableServer).getOrElse(throw new IllegalStateException("Channel closed too early"))
+
+        // Step 2: Close the client-side socket and the proxy socket to the server, triggering close notification in the
+        // server when the client is unmuted in Step 3. Get the channel into its desired closing/buffered state.
+        socket.close()
+        proxyServer.serverConnSocket.close()
+        TestUtils.waitUntilTrue(() => proxyServer.clientConnSocket.isClosed, "Client socket not closed")
+        if (makeClosing)
+          testableSelector.pendingClosingChannels.add(channel)
+        if (numComplete == 0 && hasIncomplete)
+          truncateBufferedRequest(channel)
+
+        // Step 3: Process the first request. Verify that the channel is not removed since the channel
+        // should be retained to process buffered data.
+        processRequestNoOpResponse(testableServer.dataPlaneRequestChannel, request1)
+        assertSame(channel, openOrClosingChannel(request1, testableServer).getOrElse(throw new IllegalStateException("Channel closed too early")))
+
+        // Step 4: Process buffered data. if `responseRequiredIndex>=0`, the channel should be failed and removed when
+        // attempting to send response. Otherwise, the channel should be removed when all completed buffers are processed.
+        // Channel should be closed and removed even if there is a partial buffered request when `hasIncomplete=true`
+        val numRequests = if (responseRequiredIndex >= 0) responseRequiredIndex + 1 else numComplete
+        (0 until numRequests).foreach { i =>
+          val request = receiveRequest(testableServer.dataPlaneRequestChannel)
+          if (i == numComplete - 1 && hasIncomplete)
+            truncateBufferedRequest(channel)
+          if (responseRequiredIndex == i)
+            processRequest(testableServer.dataPlaneRequestChannel, request)
+          else
+            processRequestNoOpResponse(testableServer.dataPlaneRequestChannel, request)
+        }
+        testableServer.waitForChannelClose(channel.id, locallyClosed = false)
+
+        // Verify that SocketServer is healthy
+        val anotherSocket = sslConnect(testableServer)
+        assertProcessorHealthy(testableServer, Seq(anotherSocket))
+      } finally {
+        proxyServer.close()
+      }
+    })
   }
 
   /**
-   * Tests exception handling in [[Processor.processNewResponses]] when [[Selector.send]]
-   * to a channel in closing state throws an exception. Test scenario is similar to
-   * [[SocketServerTest.processNewResponseException]].
+   * Tests idle channel expiry for SSL channels with buffered data. Muted channels are expired
+   * immediately even if there is pending data to be processed. This is consistent with PLAINTEXT where
+   * we expire muted channels even if there is data available on the socket. This scenario occurs if broker
+   * takes longer than idle timeout to process a client request. In this case, typically client would have
+   * expired its connection and would potentially reconnect to retry the request, so immediate expiry enables
+   * the old connection and its associated resources to be freed sooner.
    */
   @Test
-  def closingChannelException(): Unit = {
-    withTestableServer { testableServer =>
+  def idleExpiryWithBufferedReceives(): Unit = {
+    val idleTimeMs = 60000
+    val time = new MockTime()
+    props.put(KafkaConfig.ConnectionsMaxIdleMsProp, idleTimeMs.toString)
+    props ++= sslServerProps
+    val testableServer = new TestableSocketServer(time = time)
+    testableServer.startup()
+
+    assertTrue(testableServer.controlPlaneRequestChannelOpt.isEmpty)
+
+    val proxyServer = new ProxyServer(testableServer)
+    try {
       val testableSelector = testableServer.testableSelector
       testableSelector.updateMinWakeup(2)
 
-      val sockets = (1 to 2).map(_ => connect(testableServer))
-      val serializedBytes = producerRequestBytes
-      val request = sendRequestsUntilStagedReceive(testableServer, sockets(0), serializedBytes)
-      sendRequest(sockets(1), serializedBytes)
-
-      testableSelector.addFailure(SelectorOperation.Send)
-      sockets(0).close()
-      processRequest(testableServer.requestChannel, request)
-      processRequest(testableServer.requestChannel) // Also process request from other channel
-      testableSelector.waitForOperations(SelectorOperation.Send, 2)
-      testableServer.waitForChannelClose(request.context.connectionId, locallyClosed = true)
+      val (socket, request) = makeSocketWithBufferedRequests(testableServer, testableSelector, proxyServer)
+      time.sleep(idleTimeMs + 1)
+      testableServer.waitForChannelClose(request.context.connectionId, locallyClosed = false)
+
+      val otherSocket = sslConnect(testableServer)
+      assertProcessorHealthy(testableServer, Seq(otherSocket))
 
-      assertProcessorHealthy(testableServer, Seq(sockets(1)))
+      socket.close()
+    } finally {
+      proxyServer.close()
+      shutdownServerAndMetrics(testableServer)
     }
   }
 
@@ -749,21 +1577,21 @@ class SocketServerTest extends JUnitSuite {
    */
   @Test
   def processCompletedReceiveException(): Unit = {
-    withTestableServer { testableServer =>
+    withTestableServer (testWithServer = { testableServer =>
       val sockets = (1 to 2).map(_ => connect(testableServer))
       val testableSelector = testableServer.testableSelector
-      val requestChannel = testableServer.requestChannel
+      val requestChannel = testableServer.dataPlaneRequestChannel
 
       testableSelector.cachedCompletedReceives.minPerPoll = 2
       testableSelector.addFailure(SelectorOperation.Mute)
-      sockets.foreach(sendRequest(_, producerRequestBytes))
+      sockets.foreach(sendRequest(_, producerRequestBytes()))
       val requests = sockets.map(_ => receiveRequest(requestChannel))
       testableSelector.waitForOperations(SelectorOperation.Mute, 2)
       testableServer.waitForChannelClose(testableSelector.allFailedChannels.head, locallyClosed = true)
       requests.foreach(processRequest(requestChannel, _))
 
       assertProcessorHealthy(testableServer, testableSelector.notFailed(sockets))
-    }
+    })
   }
 
   /**
@@ -778,18 +1606,18 @@ class SocketServerTest extends JUnitSuite {
    */
   @Test
   def processCompletedSendException(): Unit = {
-    withTestableServer { testableServer =>
+    withTestableServer (testWithServer = { testableServer =>
       val testableSelector = testableServer.testableSelector
       val sockets = (1 to 2).map(_ => connect(testableServer))
       val requests = sockets.map(sendAndReceiveRequest(_, testableServer))
 
       testableSelector.addFailure(SelectorOperation.Unmute)
-      requests.foreach(processRequest(testableServer.requestChannel, _))
+      requests.foreach(processRequest(testableServer.dataPlaneRequestChannel, _))
       testableSelector.waitForOperations(SelectorOperation.Unmute, 2)
       testableServer.waitForChannelClose(testableSelector.allFailedChannels.head, locallyClosed = true)
 
       assertProcessorHealthy(testableServer, testableSelector.notFailed(sockets))
-    }
+    })
   }
 
   /**
@@ -802,7 +1630,7 @@ class SocketServerTest extends JUnitSuite {
    */
   @Test
   def processDisconnectedException(): Unit = {
-    withTestableServer { testableServer =>
+    withTestableServer (testWithServer = { testableServer =>
       val (socket, connectionId) = connectAndProcessRequest(testableServer)
       val testableSelector = testableServer.testableSelector
 
@@ -816,7 +1644,7 @@ class SocketServerTest extends JUnitSuite {
       testableServer.waitForChannelClose(connectionId, locallyClosed = false)
 
       assertProcessorHealthy(testableServer)
-    }
+    })
   }
 
   /**
@@ -824,7 +1652,7 @@ class SocketServerTest extends JUnitSuite {
    */
   @Test
   def pollException(): Unit = {
-    withTestableServer { testableServer =>
+    withTestableServer (testWithServer = { testableServer =>
       val (socket, _) = connectAndProcessRequest(testableServer)
       val testableSelector = testableServer.testableSelector
 
@@ -833,7 +1661,7 @@ class SocketServerTest extends JUnitSuite {
       testableSelector.waitForOperations(SelectorOperation.Poll, 2)
 
       assertProcessorHealthy(testableServer, Seq(socket))
-    }
+    })
   }
 
   /**
@@ -841,34 +1669,178 @@ class SocketServerTest extends JUnitSuite {
    */
   @Test
   def controlThrowable(): Unit = {
-    withTestableServer { testableServer =>
+    withTestableServer (testWithServer = { testableServer =>
       connectAndProcessRequest(testableServer)
       val testableSelector = testableServer.testableSelector
 
       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)
+      assertEquals(1, testableServer.uncaughtExceptions)
+      testableServer.uncaughtExceptions = 0
+    })
+  }
+
+  @Test
+  def testConnectionRateLimit(): Unit = {
+    shutdownServerAndMetrics(server)
+    val numConnections = 5
+    props.put("max.connections.per.ip", numConnections.toString)
+    val testableServer = new TestableSocketServer(KafkaConfig.fromProps(props), connectionQueueSize = 1)
+    testableServer.startup()
+    val testableSelector = testableServer.testableSelector
+    val errors = new mutable.HashSet[String]
+
+    def acceptorStackTraces: scala.collection.Map[Thread, String] = {
+      Thread.getAllStackTraces.asScala.collect {
+        case (thread, stacktraceElement) if thread.getName.contains("kafka-socket-acceptor") =>
+          thread -> stacktraceElement.mkString("\n")
+      }
+    }
+
+    def acceptorBlocked: Boolean = {
+      val stackTraces = acceptorStackTraces
+      if (stackTraces.isEmpty)
+        errors.add(s"Acceptor thread not found, threads=${Thread.getAllStackTraces.keySet}")
+      stackTraces.exists { case (thread, stackTrace) =>
+        thread.getState == Thread.State.WAITING && stackTrace.contains("ArrayBlockingQueue")
+      }
+    }
+
+    def registeredConnectionCount: Int = testableSelector.operationCounts.getOrElse(SelectorOperation.Register, 0)
+
+    try {
+      // Block selector until Acceptor is blocked while connections are pending
+      testableSelector.pollCallback = () => {
+        try {
+          TestUtils.waitUntilTrue(() => errors.nonEmpty || registeredConnectionCount >= numConnections - 1 || acceptorBlocked,
+            "Acceptor not blocked", waitTimeMs = 10000)
+        } catch {
+          case _: Throwable => errors.add(s"Acceptor not blocked: $acceptorStackTraces")
+        }
+      }
+      testableSelector.operationCounts.clear()
+      val sockets = (1 to numConnections).map(_ => connect(testableServer))
+      TestUtils.waitUntilTrue(() => errors.nonEmpty || registeredConnectionCount == numConnections,
+        "Connections not registered", waitTimeMs = 15000)
+      assertEquals(Set.empty, errors)
+      testableSelector.waitForOperations(SelectorOperation.Register, numConnections)
+
+      // In each iteration, SocketServer processes at most connectionQueueSize (1 in this test)
+      // new connections and then does poll() to process data from existing connections. So for
+      // 5 connections, we expect 5 iterations. Since we stop when the 5th connection is processed,
+      // we can safely check that there were atleast 4 polls prior to the 5th connection.
+      val pollCount = testableSelector.operationCounts(SelectorOperation.Poll)
+      assertTrue(s"Connections created too quickly: $pollCount", pollCount >= numConnections - 1)
+      verifyAcceptorBlockedPercent("PLAINTEXT", expectBlocked = true)
+
+      assertProcessorHealthy(testableServer, sockets)
+    } finally {
+      shutdownServerAndMetrics(testableServer)
     }
   }
 
-  private def withTestableServer(testWithServer: TestableSocketServer => Unit): Unit = {
-    props.put("listeners", "PLAINTEXT://localhost:0")
-    val testableServer = new TestableSocketServer
+
+  @Test
+  def testControlPlaneAsPrivilegedListener(): Unit = {
+    val testProps = new Properties
+    testProps ++= props
+    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 controlPlaneSocket = connect(testableServer, config.controlPlaneListenerName.get,
+        localAddr = InetAddress.getLocalHost)
+      val sentRequest = sendAndReceiveControllerRequest(controlPlaneSocket, testableServer)
+      assertTrue(sentRequest.context.fromPrivilegedListener)
+
+      val plainSocket = connect(testableServer, localAddr = InetAddress.getLocalHost)
+      val plainRequest = sendAndReceiveRequest(plainSocket, testableServer)
+      assertFalse(plainRequest.context.fromPrivilegedListener)
+    })
+  }
+
+  @Test
+  def testInterBrokerListenerAsPrivilegedListener(): Unit = {
+    val testProps = new Properties
+    testProps ++= props
+    testProps.put("listeners", "EXTERNAL://localhost:0,INTERNAL://localhost:0")
+    testProps.put("listener.security.protocol.map", "EXTERNAL:PLAINTEXT,INTERNAL:PLAINTEXT")
+    testProps.put("inter.broker.listener.name", "INTERNAL")
+    val config = KafkaConfig.fromProps(testProps)
+    withTestableServer(config, { testableServer =>
+      val interBrokerSocket = connect(testableServer, config.interBrokerListenerName,
+        localAddr = InetAddress.getLocalHost)
+      val sentRequest = sendAndReceiveRequest(interBrokerSocket, testableServer)
+      assertTrue(sentRequest.context.fromPrivilegedListener)
+
+      val externalSocket = connect(testableServer, new ListenerName("EXTERNAL"),
+        localAddr = InetAddress.getLocalHost)
+      val externalRequest = sendAndReceiveRequest(externalSocket, testableServer)
+      assertFalse(externalRequest.context.fromPrivilegedListener)
+    })
+  }
+
+  @Test
+  def testControlPlaneTakePrecedenceOverInterBrokerListenerAsPrivilegedListener(): 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)
+    withTestableServer(config, { testableServer =>
+      val controlPlaneSocket = connect(testableServer, config.controlPlaneListenerName.get,
+        localAddr = InetAddress.getLocalHost)
+      val controlPlaneRequest = sendAndReceiveControllerRequest(controlPlaneSocket, testableServer)
+      assertTrue(controlPlaneRequest.context.fromPrivilegedListener)
+
+      val interBrokerSocket = connect(testableServer, config.interBrokerListenerName,
+        localAddr = InetAddress.getLocalHost)
+      val interBrokerRequest = sendAndReceiveRequest(interBrokerSocket, testableServer)
+      assertFalse(interBrokerRequest.context.fromPrivilegedListener)
+
+      val externalSocket = connect(testableServer, new ListenerName("EXTERNAL"),
+        localAddr = InetAddress.getLocalHost)
+      val externalRequest = sendAndReceiveRequest(externalSocket, testableServer)
+      assertFalse(externalRequest.context.fromPrivilegedListener)
+    })
+  }
+
+  private def sslServerProps: Properties = {
+    val trustStoreFile = File.createTempFile("truststore", ".jks")
+    val sslProps = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, interBrokerSecurityProtocol = Some(SecurityProtocol.SSL),
+      trustStoreFile = Some(trustStoreFile))
+    sslProps.put(KafkaConfig.ListenersProp, "SSL://localhost:0")
+    sslProps
+  }
+
+  private def withTestableServer(config : KafkaConfig = KafkaConfig.fromProps(props),
+                                 testWithServer: TestableSocketServer => Unit): Unit = {
+    val testableServer = new TestableSocketServer(config)
     testableServer.startup()
     try {
-        testWithServer(testableServer)
+      testWithServer(testableServer)
     } finally {
       shutdownServerAndMetrics(testableServer)
+      assertEquals(0, testableServer.uncaughtExceptions)
     }
   }
 
+  def sendAndReceiveControllerRequest(socket: Socket, server: SocketServer): RequestChannel.Request = {
+    sendRequest(socket, producerRequestBytes())
+    receiveRequest(server.controlPlaneRequestChannelOpt.get)
+  }
+
   private def assertProcessorHealthy(testableServer: TestableSocketServer, healthySockets: Seq[Socket] = Seq.empty): Unit = {
     val selector = testableServer.testableSelector
     selector.reset()
-    val requestChannel = testableServer.requestChannel
+    val requestChannel = testableServer.dataPlaneRequestChannel
 
     // Check that existing channels behave as expected
     healthySockets.foreach { socket =>
@@ -880,7 +1852,7 @@ class SocketServerTest extends JUnitSuite {
 
     // Check new channel behaves as expected
     val (socket, connectionId) = connectAndProcessRequest(testableServer)
-    assertArrayEquals(producerRequestBytes, receiveResponse(socket))
+    assertArrayEquals(producerRequestBytes(), receiveResponse(socket))
     assertNotNull("Channel should not have been closed", selector.channel(connectionId))
     assertNull("Channel should not be closing", selector.closingChannel(connectionId))
     socket.close()
@@ -891,22 +1863,45 @@ class SocketServerTest extends JUnitSuite {
   def isSocketConnectionId(connectionId: String, socket: Socket): Boolean =
     connectionId.contains(s":${socket.getLocalPort}-")
 
-  class TestableSocketServer extends SocketServer(KafkaConfig.fromProps(props),
-      new Metrics, Time.SYSTEM, credentialProvider) {
+  private def verifyAcceptorBlockedPercent(listenerName: String, expectBlocked: Boolean): Unit = {
+    val blockedPercentMetricMBeanName = "kafka.network:type=Acceptor,name=AcceptorBlockedPercent,listener=PLAINTEXT"
+    val blockedPercentMetrics = KafkaYammerMetrics.defaultRegistry.allMetrics.asScala.filter { case (k, _) =>
+      k.getMBeanName == blockedPercentMetricMBeanName
+    }.values
+    assertEquals(1, blockedPercentMetrics.size)
+    val blockedPercentMetric = blockedPercentMetrics.head.asInstanceOf[Meter]
+    val blockedPercent = blockedPercentMetric.meanRate
+    if (expectBlocked) {
+      assertTrue(s"Acceptor blocked percent not recorded: $blockedPercent", blockedPercent > 0.0)
+      assertTrue(s"Unexpected blocked percent in acceptor: $blockedPercent", blockedPercent <= 1.0)
+    } else {
+      assertEquals(0.0, blockedPercent, 0.001)
+    }
+  }
+
+  class TestableSocketServer(config : KafkaConfig = KafkaConfig.fromProps(props), val connectionQueueSize: Int = 20,
+                             override val time: Time = Time.SYSTEM) extends SocketServer(config,
+    new Metrics, time, credentialProvider) {
 
     @volatile var selector: Option[TestableSelector] = None
+    @volatile var uncaughtExceptions = 0
 
-    override def newProcessor(id: Int, connectionQuotas: ConnectionQuotas, listenerName: ListenerName,
-                                protocol: SecurityProtocol, memoryPool: MemoryPool): Processor = {
-      new Processor(id, time, config.socketRequestMaxBytes, requestChannel, connectionQuotas,
-        config.connectionsMaxIdleMs, listenerName, protocol, config, metrics, credentialProvider, memoryPool, new
-            LogContext()) {
+    override def newProcessor(id: Int, requestChannel: RequestChannel, connectionQuotas: ConnectionQuotas, listenerName: ListenerName,
+                              protocol: SecurityProtocol, memoryPool: MemoryPool, isPrivilegedListener: Boolean = false): Processor = {
+      new Processor(id, time, config.socketRequestMaxBytes, requestChannel, connectionQuotas, config.connectionsMaxIdleMs,
+        config.failedAuthenticationDelayMs, listenerName, protocol, config, metrics, credentialProvider,
+        memoryPool, new LogContext(), connectionQueueSize, isPrivilegedListener) {
 
         override protected[network] def createSelector(channelBuilder: ChannelBuilder): Selector = {
-           val testableSelector = new TestableSelector(config, channelBuilder, time, metrics)
-           assertEquals(None, selector)
-           selector = Some(testableSelector)
-           testableSelector
+          val testableSelector = new TestableSelector(config, channelBuilder, time, metrics, metricTags.asScala)
+          selector = Some(testableSelector)
+          testableSelector
+        }
+
+        override private[network] def processException(errorMessage: String, throwable: Throwable): Unit = {
+          if (errorMessage.contains("uncaught exception"))
+            uncaughtExceptions += 1
+          super.processException(errorMessage, throwable)
         }
       }
     }
@@ -918,17 +1913,17 @@ class SocketServerTest extends JUnitSuite {
       val selector = testableSelector
       if (locallyClosed) {
         TestUtils.waitUntilTrue(() => selector.allLocallyClosedChannels.contains(connectionId),
-            s"Channel not closed: $connectionId")
+          s"Channel not closed: $connectionId")
         assertTrue("Unexpected disconnect notification", testableSelector.allDisconnectedChannels.isEmpty)
       } else {
         TestUtils.waitUntilTrue(() => selector.allDisconnectedChannels.contains(connectionId),
-            s"Disconnect notification not received: $connectionId")
+          s"Disconnect notification not received: $connectionId")
         assertTrue("Channel closed locally", testableSelector.allLocallyClosedChannels.isEmpty)
       }
       val openCount = selector.allChannels.size - 1 // minus one for the channel just closed above
       TestUtils.waitUntilTrue(() => connectionCount(localAddress) == openCount, "Connection count not decremented")
       TestUtils.waitUntilTrue(() =>
-        processor(0).inflightResponseCount == 0, "Inflight responses not cleared")
+        dataPlaneProcessor(0).inflightResponseCount == 0, "Inflight responses not cleared")
       assertNull("Channel not removed", selector.channel(connectionId))
       assertNull("Closing channel not removed", selector.closingChannel(connectionId))
     }
@@ -946,11 +1941,11 @@ class SocketServerTest extends JUnitSuite {
     case object CloseSelector extends SelectorOperation
   }
 
-  class TestableSelector(config: KafkaConfig, channelBuilder: ChannelBuilder, time: Time, metrics: Metrics)
-        extends Selector(config.socketRequestMaxBytes, config.connectionsMaxIdleMs,
-            metrics, time, "socket-server", new HashMap, false, true, channelBuilder, MemoryPool.NONE, new LogContext()) {
+  class TestableSelector(config: KafkaConfig, channelBuilder: ChannelBuilder, time: Time, metrics: Metrics, metricTags: mutable.Map[String, String] = mutable.Map.empty)
+    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]()
@@ -960,12 +1955,19 @@ class SocketServerTest extends JUnitSuite {
     // Enable data from `Selector.poll()` to be deferred to a subsequent poll() until
     // the number of elements of that type reaches `minPerPoll`. This enables tests to verify
     // that failed processing doesn't impact subsequent processing within the same iteration.
-    class PollData[T] {
+    abstract class PollData[T] {
       var minPerPoll = 1
       val deferredValues = mutable.Buffer[T]()
-      val currentPollValues = mutable.Buffer[T]()
-      def update(newValues: mutable.Buffer[T]): Unit = {
-        if (currentPollValues.nonEmpty || deferredValues.size + newValues.size >= minPerPoll) {
+
+      /**
+       * Process new results and return the results for the current poll if at least
+       * `minPerPoll` results are available including any deferred results. Otherwise
+       * add the provided values to the deferred set and return an empty buffer. This allows
+       * tests to process `minPerPoll` elements as the results of a single poll iteration.
+       */
+      protected def update(newValues: mutable.Buffer[T]): mutable.Buffer[T] = {
+        val currentPollValues = mutable.Buffer[T]()
+        if (deferredValues.size + newValues.size >= minPerPoll) {
           if (deferredValues.nonEmpty) {
             currentPollValues ++= deferredValues
             deferredValues.clear()
@@ -973,19 +1975,56 @@ class SocketServerTest extends JUnitSuite {
           currentPollValues ++= newValues
         } else
           deferredValues ++= newValues
+
+        currentPollValues
+      }
+
+      /**
+       * Process results from the appropriate buffer in Selector and update the buffer to either
+       * defer and return nothing or return all results including previously deferred values.
+       */
+      def updateResults(): Unit
+    }
+
+    class CompletedReceivesPollData(selector: TestableSelector) extends PollData[NetworkReceive] {
+      val completedReceivesMap: util.Map[String, NetworkReceive] = JTestUtils.fieldValue(selector, classOf[Selector], "completedReceives")
+
+      override def updateResults(): Unit = {
+        val currentReceives = update(selector.completedReceives.asScala.toBuffer)
+        completedReceivesMap.clear()
+        currentReceives.foreach { receive =>
+          val channelOpt = Option(selector.channel(receive.source)).orElse(Option(selector.closingChannel(receive.source)))
+          channelOpt.foreach { channel => completedReceivesMap.put(channel.id, receive) }
+        }
+      }
+    }
+
+    class CompletedSendsPollData(selector: TestableSelector) extends PollData[NetworkSend] {
+      override def updateResults(): Unit = {
+        val currentSends = update(selector.completedSends.asScala)
+        selector.completedSends.clear()
+        currentSends.foreach { selector.completedSends.add }
       }
-      def reset(): Unit = {
-        currentPollValues.clear()
+    }
+
+    class DisconnectedPollData(selector: TestableSelector) extends PollData[(String, ChannelState)] {
+      override def updateResults(): Unit = {
+        val currentDisconnected = update(selector.disconnected.asScala.toBuffer)
+        selector.disconnected.clear()
+        currentDisconnected.foreach { case (channelId, state) => selector.disconnected.put(channelId, state) }
       }
     }
-    val cachedCompletedReceives = new PollData[NetworkReceive]()
-    val cachedCompletedSends = new PollData[Send]()
-    val cachedDisconnected = new PollData[(String, ChannelState)]()
+
+    val cachedCompletedReceives = new CompletedReceivesPollData(this)
+    val cachedCompletedSends = new CompletedSendsPollData(this)
+    val cachedDisconnected = new DisconnectedPollData(this)
     val allCachedPollData = Seq(cachedCompletedReceives, cachedCompletedSends, cachedDisconnected)
+    val pendingClosingChannels = new ConcurrentLinkedQueue[KafkaChannel]()
     @volatile var minWakeupCount = 0
     @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"))
     }
@@ -1005,7 +2044,7 @@ class SocketServerTest extends JUnitSuite {
     }
 
     def runOp[T](operation: SelectorOperation, connectionId: Option[String],
-        onFailure: => Unit = {})(code: => T): T = {
+                 onFailure: => Unit = {})(code: => T): T = {
       // If a failure is set on `operation`, throw that exception even if `code` fails
       try code
       finally onOperation(operation, connectionId, onFailure)
@@ -1017,24 +2056,31 @@ class SocketServerTest extends JUnitSuite {
       }
     }
 
-    override def send(s: Send): Unit = {
-      runOp(SelectorOperation.Send, Some(s.destination)) {
+    override def send(s: NetworkSend): Unit = {
+      runOp(SelectorOperation.Send, Some(s.destinationId)) {
         super.send(s)
       }
     }
 
     override def poll(timeout: Long): Unit = {
       try {
-        allCachedPollData.foreach(_.reset)
+        assertEquals(0, super.completedReceives().size)
+        assertEquals(0, super.completedSends().size)
+
+        pollCallback.apply()
+        while (!pendingClosingChannels.isEmpty) {
+          makeClosing(pendingClosingChannels.poll())
+        }
         runOp(SelectorOperation.Poll, None) {
           super.poll(pollTimeoutOverride.getOrElse(timeout))
         }
       } finally {
-        super.channels.asScala.foreach(allChannels += _.id)
+        super.channels.forEach(allChannels += _.id)
         allDisconnectedChannels ++= super.disconnected.asScala.keys
-        cachedCompletedReceives.update(super.completedReceives.asScala)
-        cachedCompletedSends.update(super.completedSends.asScala)
-        cachedDisconnected.update(super.disconnected.asScala.toBuffer)
+
+        cachedCompletedReceives.updateResults()
+        cachedCompletedSends.updateResults()
+        cachedDisconnected.updateResults()
       }
     }
 
@@ -1059,12 +2105,6 @@ class SocketServerTest extends JUnitSuite {
       }
     }
 
-    override def disconnected: java.util.Map[String, ChannelState] = cachedDisconnected.currentPollValues.toMap.asJava
-
-    override def completedSends: java.util.List[Send] = cachedCompletedSends.currentPollValues.asJava
-
-    override def completedReceives: java.util.List[NetworkReceive] = cachedCompletedReceives.currentPollValues.asJava
-
     override def close(id: String): Unit = {
       runOp(SelectorOperation.Close, Some(id)) {
         super.close(id)
@@ -1098,5 +2138,67 @@ class SocketServerTest extends JUnitSuite {
       val failedConnectionId = allFailedChannels.head
       sockets.filterNot(socket => isSocketConnectionId(failedConnectionId, socket))
     }
+
+    private def makeClosing(channel: KafkaChannel): Unit = {
+      val channels: util.Map[String, KafkaChannel] = JTestUtils.fieldValue(this, classOf[Selector], "channels")
+      val closingChannels: util.Map[String, KafkaChannel] = JTestUtils.fieldValue(this, classOf[Selector], "closingChannels")
+      closingChannels.put(channel.id, channel)
+      channels.remove(channel.id)
+    }
+  }
+
+  /**
+   * Proxy server used to intercept connections to SocketServer. This is used for testing SSL channels
+   * with buffered data. A single SSL client is expected to be created by the test using this ProxyServer.
+   * By default, data between the client and the server is simply transferred across to the destination by ProxyServer.
+   * Tests can enable buffering in ProxyServer to directly copy incoming data from the client to the server-side
+   * channel's `netReadBuffer` to simulate scenarios with SSL buffered data.
+   */
+  private class ProxyServer(socketServer: SocketServer) {
+    val serverSocket = new ServerSocket(0)
+    val localPort = serverSocket.getLocalPort
+    val serverConnSocket = new Socket("localhost", socketServer.boundPort(ListenerName.forSecurityProtocol(SecurityProtocol.SSL)))
+    val executor = Executors.newFixedThreadPool(2)
+    @volatile var clientConnSocket: Socket = _
+    @volatile var buffer: Option[ByteBuffer] = None
+
+    executor.submit((() => {
+      try {
+        clientConnSocket = serverSocket.accept()
+        val serverOut = serverConnSocket.getOutputStream
+        val clientIn = clientConnSocket.getInputStream
+        var b: Int = -1
+        while ({b = clientIn.read(); b != -1}) {
+          buffer match {
+            case Some(buf) =>
+              buf.put(b.asInstanceOf[Byte])
+            case None =>
+              serverOut.write(b)
+              serverOut.flush()
+          }
+        }
+      } finally {
+        clientConnSocket.close()
+      }
+    }): Runnable)
+
+    executor.submit((() => {
+      var b: Int = -1
+      val serverIn = serverConnSocket.getInputStream
+      while ({b = serverIn.read(); b != -1}) {
+        clientConnSocket.getOutputStream.write(b)
+      }
+    }): Runnable)
+
+    def enableBuffering(buffer: ByteBuffer): Unit = this.buffer = Some(buffer)
+
+    def close(): Unit = {
+      serverSocket.close()
+      serverConnSocket.close()
+      clientConnSocket.close()
+      executor.shutdownNow()
+      assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS))
+    }
+
   }
 }
diff --git a/core/src/test/scala/unit/kafka/producer/AsyncProducerTest.scala b/core/src/test/scala/unit/kafka/producer/AsyncProducerTest.scala
deleted file mode 100755
index 74f3ad1efea3e..0000000000000
--- a/core/src/test/scala/unit/kafka/producer/AsyncProducerTest.scala
+++ /dev/null
@@ -1,502 +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.producer
-
-import java.util.Properties
-import java.util.concurrent.LinkedBlockingQueue
-
-import org.apache.kafka.common.protocol.Errors
-import org.junit.Assert.{assertEquals, assertTrue}
-import org.easymock.EasyMock
-import org.junit.Test
-import kafka.api._
-import kafka.cluster.BrokerEndPoint
-import kafka.common._
-import kafka.message._
-import kafka.producer.async._
-import kafka.serializer._
-import kafka.server.KafkaConfig
-import kafka.utils.TestUtils._
-
-import scala.collection.Map
-import scala.collection.mutable.ArrayBuffer
-import kafka.utils._
-import org.apache.kafka.common.security.auth.SecurityProtocol
-import org.apache.kafka.common.utils.Time
-
-@deprecated("This test has been deprecated and it will be removed in a future release.", "0.10.0.0")
-class AsyncProducerTest {
-
-  class NegativePartitioner(props: VerifiableProperties = null) extends Partitioner {
-    def partition(data: Any, numPartitions: Int): Int = -1
-  }
-
-  // One of the few cases we can just set a fixed port because the producer is mocked out here since this uses mocks
-  val props = Seq(createBrokerConfig(1, "127.0.0.1:1", port = 65534))
-  val configs = props.map(KafkaConfig.fromProps)
-  val brokerList = configs.map { config =>
-    val endPoint = config.advertisedListeners.find(_.securityProtocol == SecurityProtocol.PLAINTEXT).get
-    org.apache.kafka.common.utils.Utils.formatAddress(endPoint.host, endPoint.port)
-  }.mkString(",")
-
-  @Test
-  def testProducerQueueSize() {
-    // a mock event handler that blocks
-    val mockEventHandler = new EventHandler[String,String] {
-
-      def handle(events: Seq[KeyedMessage[String,String]]) {
-        Thread.sleep(500)
-      }
-
-      def close(): Unit = ()
-    }
-
-    val props = new Properties()
-    props.put("serializer.class", "kafka.serializer.StringEncoder")
-    props.put("metadata.broker.list", brokerList)
-    props.put("producer.type", "async")
-    props.put("queue.buffering.max.messages", "10")
-    props.put("batch.num.messages", "1")
-    props.put("queue.enqueue.timeout.ms", "0")
-
-    val config = new ProducerConfig(props)
-    val produceData = getProduceData(12)
-    val producer = new Producer[String, String](config, mockEventHandler)
-    try {
-      // send all 10 messages, should hit the batch size and then reach broker
-      producer.send(produceData: _*)
-      fail("Queue should be full")
-    }
-    catch {
-      case _: QueueFullException => //expected
-    }finally {
-      producer.close()
-    }
-  }
-
-  @Test
-  def testProduceAfterClosed() {
-    val produceData = getProduceData(10)
-    val producer = createProducer[String, String](
-      brokerList,
-      encoder = classOf[StringEncoder].getName)
-
-    producer.close
-
-    try {
-      producer.send(produceData: _*)
-      fail("should complain that producer is already closed")
-    }
-    catch {
-      case _: ProducerClosedException => //expected
-    }
-  }
-
-  @Test
-  def testBatchSize() {
-    /**
-     *  Send a total of 10 messages with batch size of 5. Expect 2 calls to the handler, one for each batch.
-     */
-    val producerDataList = getProduceData(10)
-    val mockHandler = EasyMock.createStrictMock(classOf[DefaultEventHandler[String,String]])
-    mockHandler.handle(producerDataList.take(5))
-    EasyMock.expectLastCall
-    mockHandler.handle(producerDataList.takeRight(5))
-    EasyMock.expectLastCall
-    EasyMock.replay(mockHandler)
-
-    val queue = new LinkedBlockingQueue[KeyedMessage[String,String]](10)
-    val producerSendThread =
-      new ProducerSendThread[String,String]("thread1", queue, mockHandler, Integer.MAX_VALUE, 5, "")
-    producerSendThread.start()
-
-    for (producerData <- producerDataList)
-      queue.put(producerData)
-
-    producerSendThread.shutdown
-    EasyMock.verify(mockHandler)
-  }
-
-  @Test
-  def testQueueTimeExpired() {
-    /**
-     *  Send a total of 2 messages with batch size of 5 and queue time of 200ms.
-     *  Expect 1 calls to the handler after 200ms.
-     */
-    val producerDataList = getProduceData(2)
-    val mockHandler = EasyMock.createStrictMock(classOf[DefaultEventHandler[String,String]])
-    mockHandler.handle(producerDataList)
-    EasyMock.expectLastCall
-    EasyMock.replay(mockHandler)
-
-    val queueExpirationTime = 200
-    val queue = new LinkedBlockingQueue[KeyedMessage[String,String]](10)
-    val producerSendThread =
-      new ProducerSendThread[String,String]("thread1", queue, mockHandler, queueExpirationTime, 5, "")
-    producerSendThread.start()
-
-    for (producerData <- producerDataList)
-      queue.put(producerData)
-
-    Thread.sleep(queueExpirationTime + 100)
-    EasyMock.verify(mockHandler)
-    producerSendThread.shutdown
-  }
-
-  @Test
-  def testPartitionAndCollateEvents() {
-    val producerDataList = new ArrayBuffer[KeyedMessage[Int,Message]]
-    // use bogus key and partition key override for some messages
-    producerDataList.append(new KeyedMessage[Int,Message]("topic1", key = 0, message = new Message("msg1".getBytes)))
-    producerDataList.append(new KeyedMessage[Int,Message]("topic2", key = -99, partKey = 1, message = new Message("msg2".getBytes)))
-    producerDataList.append(new KeyedMessage[Int,Message]("topic1", key = 2, message = new Message("msg3".getBytes)))
-    producerDataList.append(new KeyedMessage[Int,Message]("topic1", key = -101, partKey = 3, message = new Message("msg4".getBytes)))
-    producerDataList.append(new KeyedMessage[Int,Message]("topic2", key = 4, message = new Message("msg5".getBytes)))
-
-    val props = new Properties()
-    props.put("metadata.broker.list", brokerList)
-    val broker1 = new BrokerEndPoint(0, "localhost", 9092)
-    val broker2 = new BrokerEndPoint(1, "localhost", 9093)
-
-    // form expected partitions metadata
-    val partition1Metadata = new PartitionMetadata(0, Some(broker1), List(broker1, broker2))
-    val partition2Metadata = new PartitionMetadata(1, Some(broker2), List(broker1, broker2))
-    val topic1Metadata = new TopicMetadata("topic1", List(partition1Metadata, partition2Metadata))
-    val topic2Metadata = new TopicMetadata("topic2", List(partition1Metadata, partition2Metadata))
-
-    val topicPartitionInfos = new collection.mutable.HashMap[String, TopicMetadata]
-    topicPartitionInfos.put("topic1", topic1Metadata)
-    topicPartitionInfos.put("topic2", topic2Metadata)
-
-    val intPartitioner = new Partitioner {
-      def partition(key: Any, numPartitions: Int): Int = key.asInstanceOf[Int] % numPartitions
-    }
-    val config = new ProducerConfig(props)
-
-    val producerPool = new ProducerPool(config)
-    val handler = new DefaultEventHandler[Int,String](config,
-                                                      partitioner = intPartitioner,
-                                                      encoder = null.asInstanceOf[Encoder[String]],
-                                                      keyEncoder = new IntEncoder(),
-                                                      producerPool = producerPool,
-                                                      topicPartitionInfos = topicPartitionInfos)
-
-    val topic1Broker1Data =
-      ArrayBuffer[KeyedMessage[Int,Message]](new KeyedMessage[Int,Message]("topic1", 0, new Message("msg1".getBytes)),
-                                             new KeyedMessage[Int,Message]("topic1", 2, new Message("msg3".getBytes)))
-    val topic1Broker2Data = ArrayBuffer[KeyedMessage[Int,Message]](new KeyedMessage[Int,Message]("topic1", -101, 3, new Message("msg4".getBytes)))
-    val topic2Broker1Data = ArrayBuffer[KeyedMessage[Int,Message]](new KeyedMessage[Int,Message]("topic2", 4, new Message("msg5".getBytes)))
-    val topic2Broker2Data = ArrayBuffer[KeyedMessage[Int,Message]](new KeyedMessage[Int,Message]("topic2", -99, 1, new Message("msg2".getBytes)))
-    val expectedResult = Some(Map(
-        0 -> Map(
-              TopicAndPartition("topic1", 0) -> topic1Broker1Data,
-              TopicAndPartition("topic2", 0) -> topic2Broker1Data),
-        1 -> Map(
-              TopicAndPartition("topic1", 1) -> topic1Broker2Data,
-              TopicAndPartition("topic2", 1) -> topic2Broker2Data)
-      ))
-
-    val actualResult = handler.partitionAndCollate(producerDataList)
-    assertEquals(expectedResult, actualResult)
-  }
-
-  @Test
-  def testSerializeEvents() {
-    val produceData = TestUtils.getMsgStrings(5).map(m => new KeyedMessage[String,String]("topic1",m))
-    val props = new Properties()
-    props.put("metadata.broker.list", brokerList)
-    val config = new ProducerConfig(props)
-    // form expected partitions metadata
-    val topic1Metadata = getTopicMetadata("topic1", 0, 0, "localhost", 9092)
-    val topicPartitionInfos = new collection.mutable.HashMap[String, TopicMetadata]
-    topicPartitionInfos.put("topic1", topic1Metadata)
-
-    val producerPool = new ProducerPool(config)
-
-    val handler = new DefaultEventHandler[String,String](config,
-                                                         partitioner = null.asInstanceOf[Partitioner],
-                                                         encoder = new StringEncoder,
-                                                         keyEncoder = new StringEncoder,
-                                                         producerPool = producerPool,
-                                                         topicPartitionInfos = topicPartitionInfos)
-
-    val serializedData = handler.serialize(produceData)
-    val deserializedData = serializedData.map(d => new KeyedMessage[String,String](d.topic, TestUtils.readString(d.message.payload)))
-
-    // Test that the serialize handles seq from a Stream
-    val streamedSerializedData = handler.serialize(Stream(produceData:_*))
-    val deserializedStreamData = streamedSerializedData.map(d => new KeyedMessage[String,String](d.topic, TestUtils.readString(d.message.payload)))
-
-    TestUtils.checkEquals(produceData.iterator, deserializedData.iterator)
-    TestUtils.checkEquals(produceData.iterator, deserializedStreamData.iterator)
-  }
-
-  @Test
-  def testInvalidPartition() {
-    val producerDataList = new ArrayBuffer[KeyedMessage[String,Message]]
-    producerDataList.append(new KeyedMessage[String,Message]("topic1", "key1", new Message("msg1".getBytes)))
-    val props = new Properties()
-    props.put("metadata.broker.list", brokerList)
-    val config = new ProducerConfig(props)
-
-    // form expected partitions metadata
-    val topic1Metadata = getTopicMetadata("topic1", 0, 0, "localhost", 9092)
-
-    val topicPartitionInfos = new collection.mutable.HashMap[String, TopicMetadata]
-    topicPartitionInfos.put("topic1", topic1Metadata)
-
-    val producerPool = new ProducerPool(config)
-
-    val handler = new DefaultEventHandler[String,String](config,
-                                                         partitioner = new NegativePartitioner,
-                                                         encoder = null.asInstanceOf[Encoder[String]],
-                                                         keyEncoder = null.asInstanceOf[Encoder[String]],
-                                                         producerPool = producerPool,
-                                                         topicPartitionInfos = topicPartitionInfos)
-    try {
-      handler.partitionAndCollate(producerDataList)
-    }
-    catch {
-      // should not throw any exception
-      case _: Throwable => fail("Should not throw any exception")
-
-    }
-  }
-
-  @Test
-  def testNoBroker() {
-    val props = new Properties()
-    props.put("metadata.broker.list", brokerList)
-
-    val config = new ProducerConfig(props)
-    // create topic metadata with 0 partitions
-    val topic1Metadata = new TopicMetadata("topic1", Seq.empty)
-
-    val topicPartitionInfos = new collection.mutable.HashMap[String, TopicMetadata]
-    topicPartitionInfos.put("topic1", topic1Metadata)
-
-    val producerPool = new ProducerPool(config)
-
-    val producerDataList = new ArrayBuffer[KeyedMessage[String,String]]
-    producerDataList.append(new KeyedMessage[String,String]("topic1", "msg1"))
-    val handler = new DefaultEventHandler[String,String](config,
-                                                         partitioner = null.asInstanceOf[Partitioner],
-                                                         encoder = new StringEncoder,
-                                                         keyEncoder = new StringEncoder,
-                                                         producerPool = producerPool,
-                                                         topicPartitionInfos = topicPartitionInfos)
-    try {
-      handler.handle(producerDataList)
-      fail("Should fail with FailedToSendMessageException")
-    }
-    catch {
-      case _: FailedToSendMessageException => // we retry on any exception now
-    }
-  }
-
-  @Test
-  def testIncompatibleEncoder() {
-    val props = new Properties()
-    // no need to retry since the send will always fail
-    props.put("message.send.max.retries", "0")
-    val producer= createProducer[String, String](
-      brokerList = brokerList,
-      encoder = classOf[DefaultEncoder].getName,
-      keyEncoder = classOf[DefaultEncoder].getName,
-      producerProps = props)
-
-    try {
-      producer.send(getProduceData(1): _*)
-      fail("Should fail with ClassCastException due to incompatible Encoder")
-    } catch {
-      case _: ClassCastException =>
-    } finally {
-      producer.close()
-    }
-  }
-
-  @Test
-  def testRandomPartitioner() {
-    val props = new Properties()
-    props.put("metadata.broker.list", brokerList)
-    val config = new ProducerConfig(props)
-
-    // create topic metadata with 0 partitions
-    val topic1Metadata = getTopicMetadata("topic1", 0, 0, "localhost", 9092)
-    val topic2Metadata = getTopicMetadata("topic2", 0, 0, "localhost", 9092)
-
-    val topicPartitionInfos = new collection.mutable.HashMap[String, TopicMetadata]
-    topicPartitionInfos.put("topic1", topic1Metadata)
-    topicPartitionInfos.put("topic2", topic2Metadata)
-
-    val producerPool = new ProducerPool(config)
-    val handler = new DefaultEventHandler[String,String](config,
-                                                         partitioner = null.asInstanceOf[Partitioner],
-                                                         encoder = null.asInstanceOf[Encoder[String]],
-                                                         keyEncoder = null.asInstanceOf[Encoder[String]],
-                                                         producerPool = producerPool,
-                                                         topicPartitionInfos = topicPartitionInfos)
-    val producerDataList = new ArrayBuffer[KeyedMessage[String,Message]]
-    producerDataList.append(new KeyedMessage[String,Message]("topic1", new Message("msg1".getBytes)))
-    producerDataList.append(new KeyedMessage[String,Message]("topic2", new Message("msg2".getBytes)))
-    producerDataList.append(new KeyedMessage[String,Message]("topic1", new Message("msg3".getBytes)))
-
-    val partitionedDataOpt = handler.partitionAndCollate(producerDataList)
-    partitionedDataOpt match {
-      case Some(partitionedData) =>
-        for (dataPerBroker <- partitionedData.values) {
-          for (tp <- dataPerBroker.keys)
-            assertTrue(tp.partition == 0)
-        }
-      case None =>
-        fail("Failed to collate requests by topic, partition")
-    }
-  }
-
-  @Test
-  def testFailedSendRetryLogic() {
-    val props = new Properties()
-    props.put("metadata.broker.list", brokerList)
-    props.put("request.required.acks", "1")
-    props.put("serializer.class", classOf[StringEncoder].getName)
-    props.put("key.serializer.class", classOf[NullEncoder[Int]].getName)
-    props.put("producer.num.retries", "3")
-
-    val config = new ProducerConfig(props)
-
-    val topic1 = "topic1"
-    val topic1Metadata = getTopicMetadata(topic1, Array(0, 1), 0, "localhost", 9092)
-    val topicPartitionInfos = new collection.mutable.HashMap[String, TopicMetadata]
-    topicPartitionInfos.put("topic1", topic1Metadata)
-
-    val msgs = TestUtils.getMsgStrings(2)
-
-    import SyncProducerConfig.{DefaultAckTimeoutMs, DefaultClientId}
-
-    // produce request for topic1 and partitions 0 and 1.  Let the first request fail
-    // entirely.  The second request will succeed for partition 1 but fail for partition 0.
-    // On the third try for partition 0, let it succeed.
-    val request1 = TestUtils.produceRequestWithAcks(List(topic1), List(0, 1), messagesToSet(msgs), acks = 1,
-      correlationId = 5, timeout = DefaultAckTimeoutMs, clientId = DefaultClientId)
-    val request2 = TestUtils.produceRequestWithAcks(List(topic1), List(0, 1), messagesToSet(msgs), acks = 1,
-      correlationId = 11, timeout = DefaultAckTimeoutMs, clientId = DefaultClientId)
-    val response1 = ProducerResponse(0,
-      Map((TopicAndPartition("topic1", 0), ProducerResponseStatus(Errors.NOT_LEADER_FOR_PARTITION, 0L)),
-          (TopicAndPartition("topic1", 1), ProducerResponseStatus(Errors.NONE, 0L))))
-    val request3 = TestUtils.produceRequest(topic1, 0, messagesToSet(msgs), acks = 1, correlationId = 15,
-      timeout = DefaultAckTimeoutMs, clientId = DefaultClientId)
-    val response2 = ProducerResponse(0,
-      Map((TopicAndPartition("topic1", 0), ProducerResponseStatus(Errors.NONE, 0L))))
-    val mockSyncProducer = EasyMock.createMock(classOf[SyncProducer])
-    // don't care about config mock
-    val mockConfig = EasyMock.createNiceMock(classOf[SyncProducerConfig])
-    EasyMock.expect(mockSyncProducer.config).andReturn(mockConfig).anyTimes()
-    EasyMock.expect(mockSyncProducer.send(request1)).andThrow(new RuntimeException) // simulate SocketTimeoutException
-    EasyMock.expect(mockSyncProducer.send(request2)).andReturn(response1)
-    EasyMock.expect(mockSyncProducer.send(request3)).andReturn(response2)
-    EasyMock.replay(mockSyncProducer)
-
-    val producerPool = EasyMock.createMock(classOf[ProducerPool])
-    EasyMock.expect(producerPool.getProducer(0)).andReturn(mockSyncProducer).times(3)
-    EasyMock.expect(producerPool.close())
-    EasyMock.replay(producerPool)
-    val time = new Time {
-      override def nanoseconds: Long = 0L
-      override def milliseconds: Long = 0L
-      override def sleep(ms: Long): Unit = {}
-      override def hiResClockMs: Long = 0L
-    }
-    val handler = new DefaultEventHandler(config,
-                                          partitioner = new FixedValuePartitioner(),
-                                          encoder = new StringEncoder(),
-                                          keyEncoder = new NullEncoder[Int](),
-                                          producerPool = producerPool,
-                                          topicPartitionInfos = topicPartitionInfos,
-                                          time = time)
-    val data = msgs.map(m => new KeyedMessage(topic1, 0, m)) ++ msgs.map(m => new KeyedMessage(topic1, 1, m))
-    handler.handle(data)
-    handler.close()
-
-    EasyMock.verify(mockSyncProducer)
-    EasyMock.verify(producerPool)
-  }
-
-  @Test
-  def testJavaProducer() {
-    val topic = "topic1"
-    val msgs = TestUtils.getMsgStrings(5)
-    val scalaProducerData = msgs.map(m => new KeyedMessage[String, String](topic, m))
-    val javaProducerData: java.util.List[KeyedMessage[String, String]] = {
-      import scala.collection.JavaConversions._
-      scalaProducerData
-    }
-
-    val mockScalaProducer = EasyMock.createMock(classOf[kafka.producer.Producer[String, String]])
-    mockScalaProducer.send(scalaProducerData.head)
-    EasyMock.expectLastCall()
-    mockScalaProducer.send(scalaProducerData: _*)
-    EasyMock.expectLastCall()
-    EasyMock.replay(mockScalaProducer)
-
-    val javaProducer = new kafka.javaapi.producer.Producer[String, String](mockScalaProducer)
-    javaProducer.send(javaProducerData.get(0))
-    javaProducer.send(javaProducerData)
-
-    EasyMock.verify(mockScalaProducer)
-  }
-
-  @Test
-  def testInvalidConfiguration() {
-    val props = new Properties()
-    props.put("serializer.class", "kafka.serializer.StringEncoder")
-    props.put("producer.type", "async")
-    try {
-      new ProducerConfig(props)
-      fail("should complain about wrong config")
-    }
-    catch {
-      case _: IllegalArgumentException => //expected
-    }
-  }
-
-  def getProduceData(nEvents: Int): Seq[KeyedMessage[String,String]] = {
-    val producerDataList = new ArrayBuffer[KeyedMessage[String,String]]
-    for (i <- 0 until nEvents)
-      producerDataList.append(new KeyedMessage[String,String]("topic1", null, "msg" + i))
-    producerDataList
-  }
-
-  private def getTopicMetadata(topic: String, partition: Int, brokerId: Int, brokerHost: String, brokerPort: Int): TopicMetadata = {
-    getTopicMetadata(topic, List(partition), brokerId, brokerHost, brokerPort)
-  }
-
-  private def getTopicMetadata(topic: String, partition: Seq[Int], brokerId: Int, brokerHost: String, brokerPort: Int): TopicMetadata = {
-    val broker1 = new BrokerEndPoint(brokerId, brokerHost, brokerPort)
-    new TopicMetadata(topic, partition.map(new PartitionMetadata(_, Some(broker1), List(broker1))))
-  }
-
-  def messagesToSet(messages: Seq[String]): ByteBufferMessageSet = {
-    new ByteBufferMessageSet(NoCompressionCodec, messages.map(m => new Message(m.getBytes, 0L, Message.MagicValue_V1)): _*)
-  }
-
-  def messagesToSet(key: Array[Byte], messages: Seq[Array[Byte]]): ByteBufferMessageSet = {
-    new ByteBufferMessageSet(
-      NoCompressionCodec,
-      messages.map(m => new Message(key = key, bytes = m, timestamp = 0L, magicValue = Message.MagicValue_V1)): _*)
-  }
-}
diff --git a/core/src/test/scala/unit/kafka/producer/ProducerTest.scala b/core/src/test/scala/unit/kafka/producer/ProducerTest.scala
deleted file mode 100755
index 86a0e8626985a..0000000000000
--- a/core/src/test/scala/unit/kafka/producer/ProducerTest.scala
+++ /dev/null
@@ -1,347 +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.producer
-
-import java.nio.ByteBuffer
-import java.util
-import java.util.Properties
-
-import kafka.api.FetchRequestBuilder
-import kafka.common.FailedToSendMessageException
-import kafka.consumer.SimpleConsumer
-import kafka.message.{Message, MessageAndOffset}
-import kafka.serializer.StringEncoder
-import kafka.server.{KafkaConfig, KafkaRequestHandler, KafkaServer}
-import kafka.utils._
-import kafka.zk.ZooKeeperTestHarness
-import org.apache.kafka.common.record.TimestampType
-import org.apache.kafka.common.utils.Time
-import org.apache.log4j.{Level, Logger}
-import org.junit.Assert._
-import org.junit.{After, Before, Test}
-import org.scalatest.exceptions.TestFailedException
-
-@deprecated("This test has been deprecated and it will be removed in a future release.", "0.10.0.0")
-class ProducerTest extends ZooKeeperTestHarness with Logging{
-  private val brokerId1 = 0
-  private val brokerId2 = 1
-  private var server1: KafkaServer = null
-  private var server2: KafkaServer = null
-  private var consumer1: SimpleConsumer = null
-  private var consumer2: SimpleConsumer = null
-  private val requestHandlerLogger = Logger.getLogger(classOf[KafkaRequestHandler])
-  private var servers = List.empty[KafkaServer]
-
-  // Creation of consumers is deferred until they are actually needed. This allows us to kill brokers that use random
-  // ports and then get a consumer instance that will be pointed at the correct port
-  def getConsumer1() = {
-    if (consumer1 == null)
-      consumer1 = new SimpleConsumer("localhost", TestUtils.boundPort(server1), 1000000, 64*1024, "")
-    consumer1
-  }
-
-  def getConsumer2() = {
-    if (consumer2 == null)
-      consumer2 = new SimpleConsumer("localhost", TestUtils.boundPort(server2), 1000000, 64*1024, "")
-    consumer2
-  }
-
-  @Before
-  override def setUp() {
-    super.setUp()
-    // set up 2 brokers with 4 partitions each
-    val props1 = TestUtils.createBrokerConfig(brokerId1, zkConnect, false)
-    props1.put("num.partitions", "4")
-    val config1 = KafkaConfig.fromProps(props1)
-    val props2 = TestUtils.createBrokerConfig(brokerId2, zkConnect, false)
-    props2.put("num.partitions", "4")
-    val config2 = KafkaConfig.fromProps(props2)
-    server1 = TestUtils.createServer(config1)
-    server2 = TestUtils.createServer(config2)
-    servers = List(server1,server2)
-
-    // temporarily set request handler logger to a higher level
-    requestHandlerLogger.setLevel(Level.FATAL)
-  }
-
-  @After
-  override def tearDown() {
-    // restore set request handler logger to a higher level
-    requestHandlerLogger.setLevel(Level.ERROR)
-
-    if (consumer1 != null)
-      consumer1.close()
-    if (consumer2 != null)
-      consumer2.close()
-
-    TestUtils.shutdownServers(Seq(server1, server2))
-    super.tearDown()
-  }
-
-  @Test
-  def testUpdateBrokerPartitionInfo() {
-    val topic = "new-topic"
-    TestUtils.createTopic(zkUtils, topic, numPartitions = 1, replicationFactor = 2, servers = servers)
-
-    val props = new Properties()
-    // no need to retry since the send will always fail
-    props.put("message.send.max.retries", "0")
-    val producer1 = TestUtils.createProducer[String, String](
-        brokerList = "localhost:80,localhost:81",
-        encoder = classOf[StringEncoder].getName,
-        keyEncoder = classOf[StringEncoder].getName,
-        producerProps = props)
-
-    try {
-      producer1.send(new KeyedMessage[String, String](topic, "test", "test1"))
-      fail("Test should fail because the broker list provided are not valid")
-    } catch {
-      case _: FailedToSendMessageException => // this is expected
-    } finally producer1.close()
-
-    val producer2 = TestUtils.createProducer[String, String](
-      brokerList = "localhost:80," + TestUtils.getBrokerListStrFromServers(Seq(server1)),
-      encoder = classOf[StringEncoder].getName,
-      keyEncoder = classOf[StringEncoder].getName)
-
-    try{
-      producer2.send(new KeyedMessage[String, String](topic, "test", "test1"))
-    } catch {
-      case e: Throwable => fail("Should succeed sending the message", e)
-    } finally {
-      producer2.close()
-    }
-
-    val producer3 =  TestUtils.createProducer[String, String](
-      brokerList = TestUtils.getBrokerListStrFromServers(Seq(server1, server2)),
-      encoder = classOf[StringEncoder].getName,
-      keyEncoder = classOf[StringEncoder].getName)
-
-    try{
-      producer3.send(new KeyedMessage[String, String](topic, "test", "test1"))
-    } catch {
-      case e: Throwable => fail("Should succeed sending the message", e)
-    } finally {
-      producer3.close()
-    }
-  }
-
-  @Test
-  def testSendToNewTopic() {
-    val props1 = new util.Properties()
-    props1.put("request.required.acks", "-1")
-
-    val topic = "new-topic"
-    // create topic with 1 partition and await leadership
-    TestUtils.createTopic(zkUtils, topic, numPartitions = 1, replicationFactor = 2, servers = servers)
-
-    val producer1 = TestUtils.createProducer[String, String](
-      brokerList = TestUtils.getBrokerListStrFromServers(Seq(server1, server2)),
-      encoder = classOf[StringEncoder].getName,
-      keyEncoder = classOf[StringEncoder].getName,
-      partitioner = classOf[StaticPartitioner].getName,
-      producerProps = props1)
-    val startTime = System.currentTimeMillis()
-    // Available partition ids should be 0.
-    producer1.send(new KeyedMessage[String, String](topic, "test", "test1"))
-    producer1.send(new KeyedMessage[String, String](topic, "test", "test2"))
-    val endTime = System.currentTimeMillis()
-    // get the leader
-    val leaderOpt = zkUtils.getLeaderForPartition(topic, 0)
-    assertTrue("Leader for topic new-topic partition 0 should exist", leaderOpt.isDefined)
-    val leader = leaderOpt.get
-
-    val messageSet = if(leader == server1.config.brokerId) {
-      val response1 = getConsumer1().fetch(new FetchRequestBuilder().addFetch(topic, 0, 0, 10000).build())
-      response1.messageSet("new-topic", 0).iterator.toBuffer
-    }else {
-      val response2 = getConsumer2().fetch(new FetchRequestBuilder().addFetch(topic, 0, 0, 10000).build())
-      response2.messageSet("new-topic", 0).iterator.toBuffer
-    }
-    assertEquals("Should have fetched 2 messages", 2, messageSet.size)
-    // Message 1
-    assertTrue(ByteBuffer.wrap("test1".getBytes).equals(messageSet.head.message.payload))
-    assertTrue(ByteBuffer.wrap("test".getBytes).equals(messageSet.head.message.key))
-    assertTrue(messageSet.head.message.timestamp >= startTime && messageSet.head.message.timestamp < endTime)
-    assertEquals(TimestampType.CREATE_TIME, messageSet.head.message.timestampType)
-    assertEquals(Message.MagicValue_V1, messageSet.head.message.magic)
-
-    // Message 2
-    assertTrue(ByteBuffer.wrap("test2".getBytes).equals(messageSet(1).message.payload))
-    assertTrue(ByteBuffer.wrap("test".getBytes).equals(messageSet(1).message.key))
-    assertTrue(messageSet(1).message.timestamp >= startTime && messageSet(1).message.timestamp < endTime)
-    assertEquals(TimestampType.CREATE_TIME, messageSet(1).message.timestampType)
-    assertEquals(Message.MagicValue_V1, messageSet(1).message.magic)
-    producer1.close()
-
-    val props2 = new util.Properties()
-    props2.put("request.required.acks", "3")
-    // no need to retry since the send will always fail
-    props2.put("message.send.max.retries", "0")
-
-    try {
-      val producer2 = TestUtils.createProducer[String, String](
-        brokerList = TestUtils.getBrokerListStrFromServers(Seq(server1, server2)),
-        encoder = classOf[StringEncoder].getName,
-        keyEncoder = classOf[StringEncoder].getName,
-        partitioner = classOf[StaticPartitioner].getName,
-        producerProps = props2)
-        producer2.close
-        fail("we don't support request.required.acks greater than 1")
-    }
-    catch {
-      case _: IllegalArgumentException =>  // this is expected
-    }
-  }
-
-
-  @Test
-  def testSendWithDeadBroker() {
-    val props = new Properties()
-    props.put("request.required.acks", "1")
-    // No need to retry since the topic will be created beforehand and normal send will succeed on the first try.
-    // Reducing the retries will save the time on the subsequent failure test.
-    props.put("message.send.max.retries", "0")
-
-    val topic = "new-topic"
-    // create topic
-    TestUtils.createTopic(zkUtils, topic, partitionReplicaAssignment = Map(0->Seq(0), 1->Seq(0), 2->Seq(0), 3->Seq(0)),
-                          servers = servers)
-
-    val producer = TestUtils.createProducer[String, String](
-      brokerList = TestUtils.getBrokerListStrFromServers(Seq(server1, server2)),
-      encoder = classOf[StringEncoder].getName,
-      keyEncoder = classOf[StringEncoder].getName,
-      partitioner = classOf[StaticPartitioner].getName,
-      producerProps = props)
-    val startTime = System.currentTimeMillis()
-    try {
-      // Available partition ids should be 0, 1, 2 and 3, all lead and hosted only
-      // on broker 0
-      producer.send(new KeyedMessage[String, String](topic, "test", "test1"))
-    } catch {
-      case e: Throwable => fail("Unexpected exception: " + e)
-    }
-    val endTime = System.currentTimeMillis()
-    // kill the broker
-    server1.shutdown
-    server1.awaitShutdown()
-
-    try {
-      // These sends should fail since there are no available brokers
-      producer.send(new KeyedMessage[String, String](topic, "test", "test1"))
-      fail("Should fail since no leader exists for the partition.")
-    } catch {
-      case e : TestFailedException => throw e // catch and re-throw the failure message
-      case _: Throwable => // otherwise success
-    }
-
-    // restart server 1
-    server1.startup()
-    TestUtils.waitUntilLeaderIsElectedOrChanged(zkUtils, topic, 0)
-    TestUtils.waitUntilMetadataIsPropagated(servers, topic, 0)
-    TestUtils.waitUntilLeaderIsKnown(servers, topic, 0)
-
-    try {
-      // cross check if broker 1 got the messages
-      val response1 = getConsumer1().fetch(new FetchRequestBuilder().addFetch(topic, 0, 0, 10000).build())
-      val messageSet1 = response1.messageSet(topic, 0).iterator
-      assertTrue("Message set should have 1 message", messageSet1.hasNext)
-      val message = messageSet1.next.message
-      assertTrue(ByteBuffer.wrap("test1".getBytes).equals(message.payload))
-      assertTrue(ByteBuffer.wrap("test".getBytes).equals(message.key))
-      assertTrue(message.timestamp >= startTime && message.timestamp < endTime)
-      assertEquals(TimestampType.CREATE_TIME, message.timestampType)
-      assertEquals(Message.MagicValue_V1, message.magic)
-      assertFalse("Message set should have another message", messageSet1.hasNext)
-    } catch {
-      case e: Exception => fail("Not expected", e)
-    }
-    producer.close
-  }
-
-  @Test
-  def testAsyncSendCanCorrectlyFailWithTimeout() {
-    val topic = "new-topic"
-    // create topics in ZK
-    TestUtils.createTopic(zkUtils, topic, partitionReplicaAssignment = Map(0->Seq(0, 1)), servers = servers)
-
-    val timeoutMs = 500
-    val props = new Properties()
-    props.put("request.timeout.ms", timeoutMs.toString)
-    props.put("request.required.acks", "1")
-    props.put("message.send.max.retries", "0")
-    props.put("client.id","ProducerTest-testAsyncSendCanCorrectlyFailWithTimeout")
-    val producer = TestUtils.createProducer[String, String](
-      brokerList = TestUtils.getBrokerListStrFromServers(Seq(server1, server2)),
-      encoder = classOf[StringEncoder].getName,
-      keyEncoder = classOf[StringEncoder].getName,
-      partitioner = classOf[StaticPartitioner].getName,
-      producerProps = props)
-
-    // do a simple test to make sure plumbing is okay
-    try {
-      // this message should be assigned to partition 0 whose leader is on broker 0
-      producer.send(new KeyedMessage(topic, "test", "test"))
-      // cross check if the broker received the messages
-      // we need the loop because the broker won't return the message until it has been replicated and the producer is
-      // using acks=1
-      var messageSet1: Iterator[MessageAndOffset] = null
-      TestUtils.waitUntilTrue(() => {
-        val response1 = getConsumer1().fetch(new FetchRequestBuilder().addFetch(topic, 0, 0, 10000).build())
-        messageSet1 = response1.messageSet(topic, 0).iterator
-        messageSet1.hasNext
-      }, "Message set should have 1 message")
-      assertEquals(ByteBuffer.wrap("test".getBytes), messageSet1.next.message.payload)
-
-      // stop IO threads and request handling, but leave networking operational
-      // any requests should be accepted and queue up, but not handled
-      server1.requestHandlerPool.shutdown()
-
-      val t1 = Time.SYSTEM.milliseconds
-      try {
-        // this message should be assigned to partition 0 whose leader is on broker 0, but
-        // broker 0 will not respond within timeoutMs millis.
-        producer.send(new KeyedMessage(topic, "test", "test"))
-        fail("Exception should have been thrown")
-      } catch {
-        case _: FailedToSendMessageException => /* success */
-      }
-      val t2 = Time.SYSTEM.milliseconds
-      // make sure we don't wait fewer than timeoutMs
-      assertTrue((t2-t1) >= timeoutMs)
-
-    } finally producer.close()
-  }
-
-  @Test
-  def testSendNullMessage() {
-    val producer = TestUtils.createProducer[String, String](
-      brokerList = TestUtils.getBrokerListStrFromServers(Seq(server1, server2)),
-      encoder = classOf[StringEncoder].getName,
-      keyEncoder = classOf[StringEncoder].getName,
-      partitioner = classOf[StaticPartitioner].getName)
-
-    try {
-      TestUtils.createTopic(zkUtils, "new-topic", 2, 1, servers)
-      producer.send(new KeyedMessage("new-topic", "key", null))
-    } finally {
-      producer.close()
-    }
-  }
-}
diff --git a/core/src/test/scala/unit/kafka/producer/SyncProducerTest.scala b/core/src/test/scala/unit/kafka/producer/SyncProducerTest.scala
deleted file mode 100644
index a0680a2e4d20d..0000000000000
--- a/core/src/test/scala/unit/kafka/producer/SyncProducerTest.scala
+++ /dev/null
@@ -1,254 +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.producer
-
-import java.net.SocketTimeoutException
-import java.util.Properties
-
-import kafka.admin.AdminUtils
-import kafka.api.{ProducerRequest, ProducerResponseStatus}
-import kafka.common.TopicAndPartition
-import kafka.integration.KafkaServerTestHarness
-import kafka.message._
-import kafka.server.KafkaConfig
-import kafka.utils._
-import org.apache.kafka.common.protocol.Errors
-import org.apache.kafka.common.utils.Time
-import org.apache.kafka.common.record.{DefaultRecordBatch, DefaultRecord}
-import org.junit.Test
-import org.junit.Assert._
-
-@deprecated("This test has been deprecated and it will be removed in a future release", "0.10.0.0")
-class SyncProducerTest extends KafkaServerTestHarness {
-  private val messageBytes =  new Array[Byte](2)
-  // turning off controlled shutdown since testProducerCanTimeout() explicitly shuts down request handler pool.
-  def generateConfigs = List(KafkaConfig.fromProps(TestUtils.createBrokerConfigs(1, zkConnect, false).head))
-
-  private def produceRequest(topic: String,
-    partition: Int,
-    message: ByteBufferMessageSet,
-    acks: Int,
-    timeout: Int = SyncProducerConfig.DefaultAckTimeoutMs,
-    correlationId: Int = 0,
-    clientId: String = SyncProducerConfig.DefaultClientId): ProducerRequest = {
-    TestUtils.produceRequest(topic, partition, message, acks, timeout, correlationId, clientId)
-  }
-
-  @Test
-  def testReachableServer() {
-    val server = servers.head
-    val props = TestUtils.getSyncProducerConfig(boundPort(server))
-
-    val producer = new SyncProducer(new SyncProducerConfig(props))
-
-    val firstStart = Time.SYSTEM.milliseconds
-    var response = producer.send(produceRequest("test", 0,
-      new ByteBufferMessageSet(compressionCodec = NoCompressionCodec, messages = new Message(messageBytes)), acks = 1))
-    assertNotNull(response)
-    assertTrue((Time.SYSTEM.milliseconds - firstStart) < 12000)
-
-    val secondStart = Time.SYSTEM.milliseconds
-    response = producer.send(produceRequest("test", 0,
-      new ByteBufferMessageSet(compressionCodec = NoCompressionCodec, messages = new Message(messageBytes)), acks = 1))
-    assertNotNull(response)
-    assertTrue((Time.SYSTEM.milliseconds - secondStart) < 12000)
-
-    response = producer.send(produceRequest("test", 0,
-      new ByteBufferMessageSet(compressionCodec = NoCompressionCodec, messages = new Message(messageBytes)), acks = 1))
-    assertNotNull(response)
-  }
-
-  @Test
-  def testEmptyProduceRequest() {
-    val server = servers.head
-    val props = TestUtils.getSyncProducerConfig(boundPort(server))
-
-
-    val correlationId = 0
-    val clientId = SyncProducerConfig.DefaultClientId
-    val ackTimeoutMs = SyncProducerConfig.DefaultAckTimeoutMs
-    val ack: Short = 1
-    val emptyRequest = new kafka.api.ProducerRequest(correlationId, clientId, ack, ackTimeoutMs, collection.mutable.Map[TopicAndPartition, ByteBufferMessageSet]())
-
-    val producer = new SyncProducer(new SyncProducerConfig(props))
-    val response = producer.send(emptyRequest)
-    assertTrue(response != null)
-    assertTrue(!response.hasError && response.status.isEmpty)
-  }
-
-  @Test
-  def testMessageSizeTooLarge() {
-    val server = servers.head
-    val props = TestUtils.getSyncProducerConfig(boundPort(server))
-
-    val producer = new SyncProducer(new SyncProducerConfig(props))
-    TestUtils.createTopic(zkUtils, "test", numPartitions = 1, replicationFactor = 1, servers = servers)
-
-    val message1 = new Message(new Array[Byte](configs.head.messageMaxBytes + 1))
-    val messageSet1 = new ByteBufferMessageSet(compressionCodec = NoCompressionCodec, messages = message1)
-    val response1 = producer.send(produceRequest("test", 0, messageSet1, acks = 1))
-
-    assertEquals(1, response1.status.count(_._2.error != Errors.NONE))
-    assertEquals(Errors.MESSAGE_TOO_LARGE, response1.status(TopicAndPartition("test", 0)).error)
-    assertEquals(-1L, response1.status(TopicAndPartition("test", 0)).offset)
-
-    val safeSize = configs.head.messageMaxBytes - DefaultRecordBatch.RECORD_BATCH_OVERHEAD - DefaultRecord.MAX_RECORD_OVERHEAD
-    val message2 = new Message(new Array[Byte](safeSize))
-    val messageSet2 = new ByteBufferMessageSet(compressionCodec = NoCompressionCodec, messages = message2)
-    val response2 = producer.send(produceRequest("test", 0, messageSet2, acks = 1))
-
-    assertEquals(1, response1.status.count(_._2.error != Errors.NONE))
-    assertEquals(Errors.NONE, response2.status(TopicAndPartition("test", 0)).error)
-    assertEquals(0, response2.status(TopicAndPartition("test", 0)).offset)
-  }
-
-  @Test
-  def testMessageSizeTooLargeWithAckZero() {
-    val server = servers.head
-    val props = TestUtils.getSyncProducerConfig(boundPort(server))
-
-    props.put("request.required.acks", "0")
-
-    val producer = new SyncProducer(new SyncProducerConfig(props))
-    adminZkClient.createTopic("test", 1, 1)
-    TestUtils.waitUntilLeaderIsElectedOrChanged(zkUtils, "test", 0)
-
-    // This message will be dropped silently since message size too large.
-    producer.send(produceRequest("test", 0,
-      new ByteBufferMessageSet(compressionCodec = NoCompressionCodec, messages = new Message(new Array[Byte](configs.head.messageMaxBytes + 1))), acks = 0))
-
-    // Send another message whose size is large enough to exceed the buffer size so
-    // the socket buffer will be flushed immediately;
-    // this send should fail since the socket has been closed
-    try {
-      producer.send(produceRequest("test", 0,
-        new ByteBufferMessageSet(compressionCodec = NoCompressionCodec, messages = new Message(new Array[Byte](configs.head.messageMaxBytes + 1))), acks = 0))
-    } catch {
-      case _ : java.io.IOException => // success
-    }
-  }
-
-  @Test
-  def testProduceCorrectlyReceivesResponse() {
-    val server = servers.head
-    val props = TestUtils.getSyncProducerConfig(boundPort(server))
-
-    val producer = new SyncProducer(new SyncProducerConfig(props))
-    val messages = new ByteBufferMessageSet(NoCompressionCodec, new Message(messageBytes))
-
-    // #1 - test that we get an error when partition does not belong to broker in response
-    val request = TestUtils.produceRequestWithAcks(Array("topic1", "topic2", "topic3"), Array(0), messages, 1,
-      timeout = SyncProducerConfig.DefaultAckTimeoutMs, clientId = SyncProducerConfig.DefaultClientId)
-    val response = producer.send(request)
-
-    assertNotNull(response)
-    assertEquals(request.correlationId, response.correlationId)
-    assertEquals(3, response.status.size)
-    response.status.values.foreach {
-      case ProducerResponseStatus(error, nextOffset, timestamp) =>
-        assertEquals(Errors.UNKNOWN_TOPIC_OR_PARTITION, error)
-        assertEquals(-1L, nextOffset)
-        assertEquals(Message.NoTimestamp, timestamp)
-    }
-
-    // #2 - test that we get correct offsets when partition is owned by broker
-    adminZkClient.createTopic("topic1", 1, 1)
-    TestUtils.waitUntilLeaderIsElectedOrChanged(zkUtils, "topic1", 0)
-    adminZkClient.createTopic("topic3", 1, 1)
-    TestUtils.waitUntilLeaderIsElectedOrChanged(zkUtils, "topic3", 0)
-
-    val response2 = producer.send(request)
-    assertNotNull(response2)
-    assertEquals(request.correlationId, response2.correlationId)
-    assertEquals(3, response2.status.size)
-
-    // the first and last message should have been accepted by broker
-    assertEquals(Errors.NONE, response2.status(TopicAndPartition("topic1", 0)).error)
-    assertEquals(Errors.NONE, response2.status(TopicAndPartition("topic3", 0)).error)
-    assertEquals(0, response2.status(TopicAndPartition("topic1", 0)).offset)
-    assertEquals(0, response2.status(TopicAndPartition("topic3", 0)).offset)
-
-    // the middle message should have been rejected because broker doesn't lead partition
-    assertEquals(Errors.UNKNOWN_TOPIC_OR_PARTITION,
-                        response2.status(TopicAndPartition("topic2", 0)).error)
-    assertEquals(-1, response2.status(TopicAndPartition("topic2", 0)).offset)
-  }
-
-  @Test
-  def testProducerCanTimeout() {
-    val timeoutMs = 500
-
-    val server = servers.head
-    val props = TestUtils.getSyncProducerConfig(boundPort(server))
-    val producer = new SyncProducer(new SyncProducerConfig(props))
-
-    val messages = new ByteBufferMessageSet(NoCompressionCodec, new Message(messageBytes))
-    val request = produceRequest("topic1", 0, messages, acks = 1)
-
-    // stop IO threads and request handling, but leave networking operational
-    // any requests should be accepted and queue up, but not handled
-    server.requestHandlerPool.shutdown()
-
-    val t1 = Time.SYSTEM.milliseconds
-    try {
-      producer.send(request)
-      fail("Should have received timeout exception since request handling is stopped.")
-    } catch {
-      case _: SocketTimeoutException => /* success */
-    }
-    val t2 = Time.SYSTEM.milliseconds
-    // make sure we don't wait fewer than timeoutMs for a response
-    assertTrue((t2-t1) >= timeoutMs)
-  }
-
-  @Test
-  def testProduceRequestWithNoResponse() {
-    val server = servers.head
-
-    val port = TestUtils.boundPort(server)
-    val props = TestUtils.getSyncProducerConfig(port)
-    val correlationId = 0
-    val clientId = SyncProducerConfig.DefaultClientId
-    val ackTimeoutMs = SyncProducerConfig.DefaultAckTimeoutMs
-    val ack: Short = 0
-    val emptyRequest = new kafka.api.ProducerRequest(correlationId, clientId, ack, ackTimeoutMs, collection.mutable.Map[TopicAndPartition, ByteBufferMessageSet]())
-    val producer = new SyncProducer(new SyncProducerConfig(props))
-    val response = producer.send(emptyRequest)
-    assertTrue(response == null)
-  }
-
-  @Test
-  def testNotEnoughReplicas()  {
-    val topicName = "minisrtest"
-    val server = servers.head
-    val props = TestUtils.getSyncProducerConfig(boundPort(server))
-
-    props.put("request.required.acks", "-1")
-
-    val producer = new SyncProducer(new SyncProducerConfig(props))
-    val topicProps = new Properties()
-    topicProps.put("min.insync.replicas","2")
-    adminZkClient.createTopic(topicName, 1, 1,topicProps)
-    TestUtils.waitUntilLeaderIsElectedOrChanged(zkUtils, topicName, 0)
-
-    val response = producer.send(produceRequest(topicName, 0,
-      new ByteBufferMessageSet(compressionCodec = NoCompressionCodec, messages = new Message(messageBytes)),-1))
-
-    assertEquals(Errors.NOT_ENOUGH_REPLICAS, response.status(TopicAndPartition(topicName, 0)).error)
-  }
-}
diff --git a/core/src/test/scala/unit/kafka/raft/KafkaNetworkChannelTest.scala b/core/src/test/scala/unit/kafka/raft/KafkaNetworkChannelTest.scala
new file mode 100644
index 0000000000000..699e8faa711bb
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/raft/KafkaNetworkChannelTest.scala
@@ -0,0 +1,236 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package kafka.raft
+
+import java.net.InetSocketAddress
+import java.util
+import java.util.Collections
+import java.util.concurrent.atomic.AtomicReference
+
+import org.apache.kafka.clients.MockClient.MockMetadataUpdater
+import org.apache.kafka.clients.{ApiVersion, MockClient, NodeApiVersions}
+import org.apache.kafka.common.message.{BeginQuorumEpochResponseData, EndQuorumEpochResponseData, FetchResponseData, VoteResponseData}
+import org.apache.kafka.common.protocol.{ApiKeys, ApiMessage, Errors}
+import org.apache.kafka.common.requests.{AbstractResponse, BeginQuorumEpochRequest, EndQuorumEpochRequest, VoteRequest, VoteResponse}
+import org.apache.kafka.common.utils.{MockTime, Time}
+import org.apache.kafka.common.{Node, TopicPartition}
+import org.apache.kafka.raft.{RaftRequest, RaftResponse, RaftUtil}
+import org.junit.Assert._
+import org.junit.{Before, Test}
+
+import scala.jdk.CollectionConverters._
+
+class KafkaNetworkChannelTest {
+  import KafkaNetworkChannelTest._
+
+  private val clusterId = "clusterId"
+  private val clientId = "clientId"
+  private val retryBackoffMs = 100
+  private val requestTimeoutMs = 30000
+  private val time = new MockTime()
+  private val client = new MockClient(time, new StubMetadataUpdater)
+  private val topicPartition = new TopicPartition("topic", 0)
+  private val channel = new KafkaNetworkChannel(time, client, clientId, retryBackoffMs, requestTimeoutMs)
+
+  @Before
+  def setupSupportedApis(): Unit = {
+    val supportedApis = RaftApis.map(api => new ApiVersion(api))
+    client.setNodeApiVersions(NodeApiVersions.create(supportedApis.asJava))
+  }
+
+  @Test
+  def testSendToUnknownDestination(): Unit = {
+    val destinationId = 2
+    assertBrokerNotAvailable(destinationId)
+  }
+
+  @Test
+  def testSendToBlackedOutDestination(): Unit = {
+    val destinationId = 2
+    val destinationNode = new Node(destinationId, "127.0.0.1", 9092)
+    channel.updateEndpoint(destinationId, new InetSocketAddress(destinationNode.host, destinationNode.port))
+    client.backoff(destinationNode, 500)
+    assertBrokerNotAvailable(destinationId)
+  }
+
+  @Test
+  def testSendAndDisconnect(): Unit = {
+    val destinationId = 2
+    val destinationNode = new Node(destinationId, "127.0.0.1", 9092)
+    channel.updateEndpoint(destinationId, new InetSocketAddress(destinationNode.host, destinationNode.port))
+
+    for (apiKey <- RaftApis) {
+      val response = KafkaNetworkChannel.buildResponse(buildTestErrorResponse(apiKey, Errors.INVALID_REQUEST))
+      client.prepareResponseFrom(response, destinationNode, true)
+      sendAndAssertErrorResponse(apiKey, destinationId, Errors.BROKER_NOT_AVAILABLE)
+    }
+  }
+
+  @Test
+  def testSendAndFailAuthentication(): Unit = {
+    val destinationId = 2
+    val destinationNode = new Node(destinationId, "127.0.0.1", 9092)
+    channel.updateEndpoint(destinationId, new InetSocketAddress(destinationNode.host, destinationNode.port))
+
+    for (apiKey <- RaftApis) {
+      client.createPendingAuthenticationError(destinationNode, 100)
+      sendAndAssertErrorResponse(apiKey, destinationId, Errors.CLUSTER_AUTHORIZATION_FAILED)
+
+      // reset to clear backoff time
+      client.reset()
+    }
+  }
+
+  private def assertBrokerNotAvailable(destinationId: Int): Unit = {
+    for (apiKey <- RaftApis) {
+      sendAndAssertErrorResponse(apiKey, destinationId, Errors.BROKER_NOT_AVAILABLE)
+    }
+  }
+
+  @Test
+  def testSendAndReceiveOutboundRequest(): Unit = {
+    val destinationId = 2
+    val destinationNode = new Node(destinationId, "127.0.0.1", 9092)
+    channel.updateEndpoint(destinationId, new InetSocketAddress(destinationNode.host, destinationNode.port))
+
+    for (apiKey <- RaftApis) {
+      val expectedError = Errors.INVALID_REQUEST
+      val response = KafkaNetworkChannel.buildResponse(buildTestErrorResponse(apiKey, expectedError))
+      client.prepareResponseFrom(response, destinationNode)
+      sendAndAssertErrorResponse(apiKey, destinationId, expectedError)
+    }
+  }
+
+  @Test
+  def testReceiveAndSendInboundRequest(): Unit = {
+    for (apiKey <- RaftApis) {
+      val request = KafkaNetworkChannel.buildRequest(buildTestRequest(apiKey)).build()
+      val responseRef = new AtomicReference[AbstractResponse]()
+
+      channel.postInboundRequest(request, responseRef.set)
+      val inbound = channel.receive(1000).asScala
+      assertEquals(1, inbound.size)
+
+      val inboundRequest = inbound.head.asInstanceOf[RaftRequest.Inbound]
+      val errorResponse = buildTestErrorResponse(apiKey, Errors.INVALID_REQUEST)
+      val outboundResponse = new RaftResponse.Outbound(inboundRequest.correlationId, errorResponse)
+      channel.send(outboundResponse)
+      channel.receive(1000)
+
+      assertNotNull(responseRef.get)
+      assertEquals(Errors.INVALID_REQUEST, extractError(KafkaNetworkChannel.responseData(responseRef.get)))
+    }
+  }
+
+  private def sendAndAssertErrorResponse(apiKey: ApiKeys,
+                                         destinationId: Int,
+                                         error: Errors): Unit = {
+    val correlationId = channel.newCorrelationId()
+    val createdTimeMs = time.milliseconds()
+    val apiRequest = buildTestRequest(apiKey)
+    val request = new RaftRequest.Outbound(correlationId, apiRequest, destinationId, createdTimeMs)
+
+    channel.send(request)
+    val responses = channel.receive(1000).asScala
+    assertEquals(1, responses.size)
+
+    val response = responses.head.asInstanceOf[RaftResponse.Inbound]
+    assertEquals(destinationId, response.sourceId)
+    assertEquals(correlationId, response.correlationId)
+    assertEquals(apiKey, ApiKeys.forId(response.data.apiKey))
+    assertEquals(error, extractError(response.data))
+  }
+
+  private def buildTestRequest(key: ApiKeys): ApiMessage = {
+    val leaderEpoch = 5
+    val leaderId = 1
+    key match {
+      case ApiKeys.BEGIN_QUORUM_EPOCH =>
+        BeginQuorumEpochRequest.singletonRequest(topicPartition, clusterId, leaderEpoch, leaderId)
+
+      case ApiKeys.END_QUORUM_EPOCH =>
+        EndQuorumEpochRequest.singletonRequest(topicPartition, clusterId, leaderId,
+          leaderEpoch, Collections.singletonList(2))
+
+      case ApiKeys.VOTE =>
+        val lastEpoch = 4
+        VoteRequest.singletonRequest(topicPartition, clusterId, leaderEpoch, leaderId, lastEpoch, 329)
+
+      case ApiKeys.FETCH =>
+        val request = RaftUtil.singletonFetchRequest(topicPartition, fetchPartition => {
+          fetchPartition
+            .setCurrentLeaderEpoch(5)
+            .setFetchOffset(333)
+            .setLastFetchedEpoch(5)
+        })
+        request.setReplicaId(1)
+
+      case _ =>
+        throw new AssertionError(s"Unexpected api $key")
+    }
+  }
+
+  private def buildTestErrorResponse(key: ApiKeys, error: Errors): ApiMessage = {
+    key match {
+      case ApiKeys.BEGIN_QUORUM_EPOCH =>
+        new BeginQuorumEpochResponseData()
+          .setErrorCode(error.code)
+
+      case ApiKeys.END_QUORUM_EPOCH =>
+        new EndQuorumEpochResponseData()
+          .setErrorCode(error.code)
+
+      case ApiKeys.VOTE =>
+        VoteResponse.singletonResponse(error, topicPartition, Errors.NONE, 1, 5, false);
+
+      case ApiKeys.FETCH =>
+        new FetchResponseData()
+          .setErrorCode(error.code)
+
+      case _ =>
+        throw new AssertionError(s"Unexpected api $key")
+    }
+  }
+
+  private def extractError(response: ApiMessage): Errors = {
+    val code = (response: @unchecked) match {
+      case res: BeginQuorumEpochResponseData => res.errorCode
+      case res: EndQuorumEpochResponseData => res.errorCode
+      case res: FetchResponseData => res.errorCode
+      case res: VoteResponseData => res.errorCode
+    }
+    Errors.forCode(code)
+  }
+
+}
+
+object KafkaNetworkChannelTest {
+  val RaftApis = Seq(
+    ApiKeys.VOTE,
+    ApiKeys.BEGIN_QUORUM_EPOCH,
+    ApiKeys.END_QUORUM_EPOCH,
+    ApiKeys.FETCH,
+  )
+
+  private class StubMetadataUpdater extends MockMetadataUpdater {
+    override def fetchNodes(): util.List[Node] = Collections.emptyList()
+
+    override def isUpdateNeeded: Boolean = false
+
+    override def update(time: Time, update: MockClient.MetadataUpdate): Unit = {}
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/security/auth/AclTest.scala b/core/src/test/scala/unit/kafka/security/auth/AclTest.scala
deleted file mode 100644
index dfdd85face30b..0000000000000
--- a/core/src/test/scala/unit/kafka/security/auth/AclTest.scala
+++ /dev/null
@@ -1,45 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package kafka.security.auth
-
-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
-
-class AclTest extends JUnitSuite {
-
-  val AclJson = "{\"version\": 1, \"acls\": [{\"host\": \"host1\",\"permissionType\": \"Deny\",\"operation\": \"READ\", \"principal\": \"User:alice\"  },  " +
-    "{  \"host\":  \"*\" ,  \"permissionType\": \"Allow\",  \"operation\":  \"Read\", \"principal\": \"User:bob\"  },  " +
-    "{  \"host\": \"host1\",  \"permissionType\": \"Deny\",  \"operation\":   \"Read\" ,  \"principal\": \"User:bob\"}  ]}"
-
-  @Test
-  def testAclJsonConversion(): Unit = {
-    val acl1 = new Acl(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "alice"), Deny, "host1" , Read)
-    val acl2 = new Acl(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "bob"), Allow, "*", Read)
-    val acl3 = new Acl(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "bob"), Deny, "host1", Read)
-
-    val acls = Set[Acl](acl1, acl2, acl3)
-    val jsonAcls = Json.encode(Acl.toJsonCompatibleMap(acls))
-
-    Assert.assertEquals(acls, Acl.fromBytes(jsonAcls.getBytes(UTF_8)))
-    Assert.assertEquals(acls, Acl.fromBytes(AclJson.getBytes(UTF_8)))
-  }
-
-}
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..0b35d4af3ce6a 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,9 @@ 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 {
+@deprecated("Scala Authorizer API classes gave been deprecated", "Since 2.5")
+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..a4c885d7e863b 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,10 @@ 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 {
+@deprecated("Scala Authorizer API classes gave been deprecated", "Since 2.5")
+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..68cbbc683b171 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,11 @@ 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 {
+@deprecated("Scala Authorizer API classes gave been deprecated", "Since 2.5")
+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 1e18f1d7bce93..d6342f346b25e 100644
--- a/core/src/test/scala/unit/kafka/security/auth/SimpleAclAuthorizerTest.scala
+++ b/core/src/test/scala/unit/kafka/security/auth/SimpleAclAuthorizerTest.scala
@@ -17,31 +17,50 @@
 package kafka.security.auth
 
 import java.net.InetAddress
+import java.nio.charset.StandardCharsets.UTF_8
 import java.util.UUID
 
+import kafka.api.{ApiVersion, KAFKA_2_0_IV0, KAFKA_2_0_IV1}
 import kafka.network.RequestChannel.Session
-import kafka.security.auth.Acl.WildCardHost
+import kafka.security.auth.Acl.{WildCardHost, WildCardResource}
 import kafka.server.KafkaConfig
 import kafka.utils.TestUtils
-import kafka.zk.ZooKeeperTestHarness
+import kafka.zk.{ZkAclStore, ZooKeeperTestHarness}
+import kafka.zookeeper.{GetChildrenRequest, GetDataRequest, ZooKeeperClient}
+import org.apache.kafka.common.errors.UnsupportedVersionException
+import org.apache.kafka.common.resource.PatternType
+import org.apache.kafka.common.resource.PatternType.{LITERAL, PREFIXED}
 import org.apache.kafka.common.security.auth.KafkaPrincipal
+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 {
 
-  val simpleAclAuthorizer = new SimpleAclAuthorizer
-  val simpleAclAuthorizer2 = new SimpleAclAuthorizer
-  val testPrincipal = Acl.WildCardPrincipal
-  val testHostName = InetAddress.getByName("192.168.0.1")
-  val session = Session(testPrincipal, testHostName)
-  var resource: Resource = null
-  val superUsers = "User:superuser1; User:superuser2"
-  val username = "alice"
-  var config: KafkaConfig = null
+  private val allowReadAcl = Acl(Acl.WildCardPrincipal, Allow, WildCardHost, Read)
+  private val allowWriteAcl = Acl(Acl.WildCardPrincipal, Allow, WildCardHost, Write)
+  private val denyReadAcl = Acl(Acl.WildCardPrincipal, Deny, WildCardHost, Read)
+
+  private val wildCardResource = Resource(Topic, WildCardResource, LITERAL)
+  private val prefixedResource = Resource(Topic, "foo", PREFIXED)
+
+  private val simpleAclAuthorizer = new SimpleAclAuthorizer
+  private val simpleAclAuthorizer2 = new SimpleAclAuthorizer
+  private var resource: Resource = _
+  private val superUsers = "User:superuser1; User:superuser2"
+  private val username = "alice"
+  private val principal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, username)
+  private val session = Session(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() {
+  override def setUp(): Unit = {
     super.setUp()
 
     // Increase maxUpdateRetries to avoid transient failures
@@ -54,18 +73,40 @@ class SimpleAclAuthorizerTest extends ZooKeeperTestHarness {
     config = KafkaConfig.fromProps(props)
     simpleAclAuthorizer.configure(config.originals)
     simpleAclAuthorizer2.configure(config.originals)
-    resource = new Resource(Topic, UUID.randomUUID().toString)
+    resource = Resource(Topic, "foo-" + UUID.randomUUID(), LITERAL)
+
+    zooKeeperClient = new ZooKeeperClient(zkConnect, zkSessionTimeout, zkConnectionTimeout, zkMaxInFlightRequests,
+      Time.SYSTEM, "kafka.test", "SimpleAclAuthorizerTest")
   }
 
   @After
   override def tearDown(): Unit = {
     simpleAclAuthorizer.close()
     simpleAclAuthorizer2.close()
+    zooKeeperClient.close()
     super.tearDown()
   }
 
+  @Test(expected = classOf[IllegalArgumentException])
+  def testAuthorizeThrowsOnNonLiteralResource(): Unit = {
+    simpleAclAuthorizer.authorize(session, Read, Resource(Topic, "something", PREFIXED))
+  }
+
+  @Test
+  def testAuthorizeWithEmptyResourceName(): Unit = {
+    assertFalse(simpleAclAuthorizer.authorize(session, Read, Resource(Group, "", LITERAL)))
+    simpleAclAuthorizer.addAcls(Set[Acl](allowReadAcl), Resource(Group, WildCardResource, LITERAL))
+    assertTrue(simpleAclAuthorizer.authorize(session, Read, Resource(Group, "", LITERAL)))
+  }
+
+  // Authorizing the empty resource is not supported because we create a znode with the resource name.
+  @Test(expected = classOf[IllegalArgumentException])
+  def testEmptyAclThrowsException(): Unit = {
+    simpleAclAuthorizer.addAcls(Set[Acl](allowReadAcl), Resource(Group, "", LITERAL))
+  }
+
   @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")
@@ -116,8 +157,31 @@ class SimpleAclAuthorizerTest extends ZooKeeperTestHarness {
     assertTrue("User3 should have WRITE access from host2", simpleAclAuthorizer.authorize(user3Session, Write, resource))
   }
 
+  /**
+    CustomPrincipals should be compared with their principal type and name
+   */
   @Test
-  def testDenyTakesPrecedence() {
+  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 Acl(user, Deny, host1.getHostAddress, Read)
+    val acl2 = new Acl(user, Allow, host2.getHostAddress, Read)
+    val acls = Set[Acl](acl1, acl2)
+    changeAclAndVerify(Set.empty[Acl], acls, Set.empty[Acl])
+
+    val host1Session = Session(customUserPrincipal, host1)
+    val host2Session = Session(customUserPrincipal, host2)
+
+    assertTrue("User1 should have READ access from host2", simpleAclAuthorizer.authorize(host2Session, Read, resource))
+    assertFalse("User1 should not have READ access from host1 due to denyAcl", simpleAclAuthorizer.authorize(host1Session, Read, resource))
+  }
+
+  @Test
+  def testDenyTakesPrecedence(): Unit = {
     val user = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, username)
     val host = InetAddress.getByName("192.168.2.1")
     val session = Session(user, host)
@@ -132,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])
@@ -142,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])
@@ -154,6 +218,19 @@ class SimpleAclAuthorizerTest extends ZooKeeperTestHarness {
     assertTrue("superuser always has access, no matter what acls.", simpleAclAuthorizer.authorize(session2, Read, resource))
   }
 
+  /**
+    CustomPrincipals should be compared with their principal type and name
+   */
+  @Test
+  def testSuperUserWithCustomPrincipalHasAccess(): Unit = {
+    val denyAllAcl = new Acl(Acl.WildCardPrincipal, Deny, WildCardHost, All)
+    changeAclAndVerify(Set.empty[Acl], Set[Acl](denyAllAcl), Set.empty[Acl])
+
+    val session = Session(new CustomPrincipal(KafkaPrincipal.USER_TYPE, "superuser1"), InetAddress.getByName("192.0.4.4"))
+
+    assertTrue("superuser with custom principal always has access, no matter what acls.", simpleAclAuthorizer.authorize(session, Read, resource))
+  }
+
   @Test
   def testWildCardAcls(): Unit = {
     assertFalse("when acls = [],  authorizer should fail close.", simpleAclAuthorizer.authorize(session, Read, resource))
@@ -161,7 +238,6 @@ class SimpleAclAuthorizerTest extends ZooKeeperTestHarness {
     val user1 = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, username)
     val host1 = InetAddress.getByName("192.168.3.1")
     val readAcl = new Acl(user1, Allow, host1.getHostAddress, Read)
-    val wildCardResource = new Resource(resource.resourceType, Resource.WildCardResource)
 
     val acls = changeAclAndVerify(Set.empty[Acl], Set[Acl](readAcl), Set.empty[Acl], wildCardResource)
 
@@ -180,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")
 
@@ -200,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"
@@ -222,10 +298,10 @@ class SimpleAclAuthorizerTest extends ZooKeeperTestHarness {
     TestUtils.waitUntilTrue(() => Map(resource -> Set(acl3, acl4, acl5)) == simpleAclAuthorizer.getAcls(user2), "changes not propagated in timeout period")
 
     val resourceToAcls = Map[Resource, Set[Acl]](
-      new Resource(Topic, Resource.WildCardResource) -> Set[Acl](new Acl(user2, Allow, WildCardHost, Read)),
-      new Resource(Cluster, Resource.WildCardResource) -> Set[Acl](new Acl(user2, Allow, host1, Read)),
-      new Resource(Group, Resource.WildCardResource) -> acls,
-      new Resource(Group, "test-ConsumerGroup") -> acls
+      new Resource(Topic, Resource.WildCardResource, LITERAL) -> Set[Acl](new Acl(user2, Allow, WildCardHost, Read)),
+      new Resource(Cluster, Resource.WildCardResource, LITERAL) -> Set[Acl](new Acl(user2, Allow, host1, Read)),
+      new Resource(Group, Resource.WildCardResource, LITERAL) -> acls,
+      new Resource(Group, "test-ConsumerGroup", LITERAL) -> acls
     )
 
     resourceToAcls foreach { case (key, value) => changeAclAndVerify(Set.empty[Acl], value, Set.empty[Acl], key) }
@@ -237,28 +313,28 @@ class SimpleAclAuthorizerTest extends ZooKeeperTestHarness {
     //test remove all acls for resource
     simpleAclAuthorizer.removeAcls(resource)
     TestUtils.waitAndVerifyAcls(Set.empty[Acl], simpleAclAuthorizer, resource)
-    assertTrue(!zkClient.resourceExists(resource))
+    assertTrue(!zkClient.resourceExists(resource.toPattern))
 
     //test removing last acl also deletes ZooKeeper path
     acls = changeAclAndVerify(Set.empty[Acl], Set(acl1), Set.empty[Acl])
     changeAclAndVerify(acls, Set.empty[Acl], acls)
-    assertTrue(!zkClient.resourceExists(resource))
+    assertTrue(!zkClient.resourceExists(resource.toPattern))
   }
 
   @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)
     simpleAclAuthorizer.addAcls(acls, resource)
 
     val user2 = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "bob")
-    val resource1 = new Resource(Topic, "test-2")
+    val resource1 = Resource(Topic, "test-2", LITERAL)
     val acl2 = new Acl(user2, Deny, "host3", Read)
     val acls1 = Set[Acl](acl2)
     simpleAclAuthorizer.addAcls(acls1, resource1)
 
-    zkClient.deleteAclChangeNotifications
+    zkClient.deleteAclChangeNotifications()
     val authorizer = new SimpleAclAuthorizer
     try {
       authorizer.configure(config.originals)
@@ -271,8 +347,8 @@ class SimpleAclAuthorizerTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testLocalConcurrentModificationOfResourceAcls() {
-    val commonResource = new Resource(Topic, "test")
+  def testLocalConcurrentModificationOfResourceAcls(): Unit = {
+    val commonResource = Resource(Topic, "test", LITERAL)
 
     val user1 = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, username)
     val acl1 = new Acl(user1, Allow, WildCardHost, Read)
@@ -287,8 +363,8 @@ class SimpleAclAuthorizerTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testDistributedConcurrentModificationOfResourceAcls() {
-    val commonResource = new Resource(Topic, "test")
+  def testDistributedConcurrentModificationOfResourceAcls(): Unit = {
+    val commonResource = Resource(Topic, "test", LITERAL)
 
     val user1 = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, username)
     val acl1 = new Acl(user1, Allow, WildCardHost, Read)
@@ -317,8 +393,8 @@ class SimpleAclAuthorizerTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testHighConcurrencyModificationOfResourceAcls() {
-    val commonResource = new Resource(Topic, "test")
+  def testHighConcurrencyModificationOfResourceAcls(): Unit = {
+    val commonResource = Resource(Topic, "test", LITERAL)
 
     val acls = (0 to 50).map { i =>
       val useri = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, i.toString)
@@ -402,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
@@ -419,6 +495,222 @@ class SimpleAclAuthorizerTest extends ZooKeeperTestHarness {
     TestUtils.waitAndVerifyAcls(Set.empty[Acl], simpleAclAuthorizer2, resource)
   }
 
+  @Test
+  def testAccessAllowedIfAllowAclExistsOnWildcardResource(): Unit = {
+    simpleAclAuthorizer.addAcls(Set[Acl](allowReadAcl), wildCardResource)
+
+    assertTrue(simpleAclAuthorizer.authorize(session, Read, resource))
+  }
+
+  @Test
+  def testDeleteAclOnWildcardResource(): Unit = {
+    simpleAclAuthorizer.addAcls(Set[Acl](allowReadAcl, allowWriteAcl), wildCardResource)
+
+    simpleAclAuthorizer.removeAcls(Set[Acl](allowReadAcl), wildCardResource)
+
+    assertEquals(Set(allowWriteAcl), simpleAclAuthorizer.getAcls(wildCardResource))
+  }
+
+  @Test
+  def testDeleteAllAclOnWildcardResource(): Unit = {
+    simpleAclAuthorizer.addAcls(Set[Acl](allowReadAcl), wildCardResource)
+
+    simpleAclAuthorizer.removeAcls(wildCardResource)
+
+    assertEquals(Map(), simpleAclAuthorizer.getAcls())
+  }
+
+  @Test
+  def testAccessAllowedIfAllowAclExistsOnPrefixedResource(): Unit = {
+    simpleAclAuthorizer.addAcls(Set[Acl](allowReadAcl), prefixedResource)
+
+    assertTrue(simpleAclAuthorizer.authorize(session, Read, resource))
+  }
+
+  @Test
+  def testDeleteAclOnPrefixedResource(): Unit = {
+    simpleAclAuthorizer.addAcls(Set[Acl](allowReadAcl, allowWriteAcl), prefixedResource)
+
+    simpleAclAuthorizer.removeAcls(Set[Acl](allowReadAcl), prefixedResource)
+
+    assertEquals(Set(allowWriteAcl), simpleAclAuthorizer.getAcls(prefixedResource))
+  }
+
+  @Test
+  def testDeleteAllAclOnPrefixedResource(): Unit = {
+    simpleAclAuthorizer.addAcls(Set[Acl](allowReadAcl, allowWriteAcl), prefixedResource)
+
+    simpleAclAuthorizer.removeAcls(prefixedResource)
+
+    assertEquals(Map(), simpleAclAuthorizer.getAcls())
+  }
+
+  @Test
+  def testAddAclsOnLiteralResource(): Unit = {
+    simpleAclAuthorizer.addAcls(Set[Acl](allowReadAcl, allowWriteAcl), resource)
+    simpleAclAuthorizer.addAcls(Set[Acl](allowWriteAcl, denyReadAcl), resource)
+
+    assertEquals(Set(allowReadAcl, allowWriteAcl, denyReadAcl), simpleAclAuthorizer.getAcls(resource))
+    assertEquals(Set(), simpleAclAuthorizer.getAcls(wildCardResource))
+    assertEquals(Set(), simpleAclAuthorizer.getAcls(prefixedResource))
+  }
+
+  @Test
+  def testAddAclsOnWildcardResource(): Unit = {
+    simpleAclAuthorizer.addAcls(Set[Acl](allowReadAcl, allowWriteAcl), wildCardResource)
+    simpleAclAuthorizer.addAcls(Set[Acl](allowWriteAcl, denyReadAcl), wildCardResource)
+
+    assertEquals(Set(allowReadAcl, allowWriteAcl, denyReadAcl), simpleAclAuthorizer.getAcls(wildCardResource))
+    assertEquals(Set(), simpleAclAuthorizer.getAcls(resource))
+    assertEquals(Set(), simpleAclAuthorizer.getAcls(prefixedResource))
+  }
+
+  @Test
+  def testAddAclsOnPrefiexedResource(): Unit = {
+    simpleAclAuthorizer.addAcls(Set[Acl](allowReadAcl, allowWriteAcl), prefixedResource)
+    simpleAclAuthorizer.addAcls(Set[Acl](allowWriteAcl, denyReadAcl), prefixedResource)
+
+    assertEquals(Set(allowReadAcl, allowWriteAcl, denyReadAcl), simpleAclAuthorizer.getAcls(prefixedResource))
+    assertEquals(Set(), simpleAclAuthorizer.getAcls(wildCardResource))
+    assertEquals(Set(), simpleAclAuthorizer.getAcls(resource))
+  }
+
+  @Test
+  def testAuthorizeWithPrefixedResource(): Unit = {
+    simpleAclAuthorizer.addAcls(Set[Acl](denyReadAcl), Resource(Topic, "a_other", LITERAL))
+    simpleAclAuthorizer.addAcls(Set[Acl](denyReadAcl), Resource(Topic, "a_other", PREFIXED))
+    simpleAclAuthorizer.addAcls(Set[Acl](denyReadAcl), Resource(Topic, "foo-" + UUID.randomUUID(), PREFIXED))
+    simpleAclAuthorizer.addAcls(Set[Acl](denyReadAcl), Resource(Topic, "foo-" + UUID.randomUUID(), PREFIXED))
+    simpleAclAuthorizer.addAcls(Set[Acl](denyReadAcl), Resource(Topic, "foo-" + UUID.randomUUID() + "-zzz", PREFIXED))
+    simpleAclAuthorizer.addAcls(Set[Acl](denyReadAcl), Resource(Topic, "fooo-" + UUID.randomUUID(), PREFIXED))
+    simpleAclAuthorizer.addAcls(Set[Acl](denyReadAcl), Resource(Topic, "fo-" + UUID.randomUUID(), PREFIXED))
+    simpleAclAuthorizer.addAcls(Set[Acl](denyReadAcl), Resource(Topic, "fop-" + UUID.randomUUID(), PREFIXED))
+    simpleAclAuthorizer.addAcls(Set[Acl](denyReadAcl), Resource(Topic, "fon-" + UUID.randomUUID(), PREFIXED))
+    simpleAclAuthorizer.addAcls(Set[Acl](denyReadAcl), Resource(Topic, "fon-", PREFIXED))
+    simpleAclAuthorizer.addAcls(Set[Acl](denyReadAcl), Resource(Topic, "z_other", PREFIXED))
+    simpleAclAuthorizer.addAcls(Set[Acl](denyReadAcl), Resource(Topic, "z_other", LITERAL))
+
+    simpleAclAuthorizer.addAcls(Set[Acl](allowReadAcl), prefixedResource)
+
+    assertTrue(simpleAclAuthorizer.authorize(session, Read, resource))
+  }
+
+  @Test
+  def testSingleCharacterResourceAcls(): Unit = {
+    simpleAclAuthorizer.addAcls(Set[Acl](allowReadAcl), Resource(Topic, "f", LITERAL))
+    assertTrue(simpleAclAuthorizer.authorize(session, Read, Resource(Topic, "f", LITERAL)))
+    assertFalse(simpleAclAuthorizer.authorize(session, Read, Resource(Topic, "foo", LITERAL)))
+
+    simpleAclAuthorizer.addAcls(Set[Acl](allowReadAcl), Resource(Topic, "_", PREFIXED))
+    assertTrue(simpleAclAuthorizer.authorize(session, Read, Resource(Topic, "_foo", LITERAL)))
+    assertTrue(simpleAclAuthorizer.authorize(session, Read, Resource(Topic, "_", LITERAL)))
+    assertFalse(simpleAclAuthorizer.authorize(session, Read, Resource(Topic, "foo_", LITERAL)))
+  }
+
+  @Test
+  def testGetAclsPrincipal(): Unit = {
+    val aclOnSpecificPrincipal = new Acl(principal, Allow, WildCardHost, Write)
+    simpleAclAuthorizer.addAcls(Set[Acl](aclOnSpecificPrincipal), resource)
+
+    assertEquals("acl on specific should not be returned for wildcard request",
+      0, simpleAclAuthorizer.getAcls(Acl.WildCardPrincipal).size)
+    assertEquals("acl on specific should be returned for specific request",
+      1, simpleAclAuthorizer.getAcls(principal).size)
+    assertEquals("acl on specific should be returned for different principal instance",
+      1, simpleAclAuthorizer.getAcls(new KafkaPrincipal(principal.getPrincipalType, principal.getName)).size)
+
+    simpleAclAuthorizer.removeAcls(resource)
+    val aclOnWildcardPrincipal = new Acl(Acl.WildCardPrincipal, Allow, WildCardHost, Write)
+    simpleAclAuthorizer.addAcls(Set[Acl](aclOnWildcardPrincipal), resource)
+
+    assertEquals("acl on wildcard should be returned for wildcard request",
+      1, simpleAclAuthorizer.getAcls(Acl.WildCardPrincipal).size)
+    assertEquals("acl on wildcard should not be returned for specific request",
+      0, simpleAclAuthorizer.getAcls(principal).size)
+  }
+
+  @Test(expected = classOf[UnsupportedVersionException])
+  def testThrowsOnAddPrefixedAclIfInterBrokerProtocolVersionTooLow(): Unit = {
+    givenAuthorizerWithProtocolVersion(Option(KAFKA_2_0_IV0))
+    simpleAclAuthorizer.addAcls(Set[Acl](denyReadAcl), Resource(Topic, "z_other", PREFIXED))
+  }
+
+  @Test
+  def testWritesExtendedAclChangeEventIfInterBrokerProtocolNotSet(): Unit = {
+    givenAuthorizerWithProtocolVersion(Option.empty)
+    val resource = Resource(Topic, "z_other", PREFIXED)
+    val expected = new String(ZkAclStore(PREFIXED).changeStore.createChangeNode(resource.toPattern).bytes, UTF_8)
+
+    simpleAclAuthorizer.addAcls(Set[Acl](denyReadAcl), resource)
+
+    val actual = getAclChangeEventAsString(PREFIXED)
+
+    assertEquals(expected, actual)
+  }
+
+  @Test
+  def testWritesExtendedAclChangeEventWhenInterBrokerProtocolAtLeastKafkaV2(): Unit = {
+    givenAuthorizerWithProtocolVersion(Option(KAFKA_2_0_IV1))
+    val resource = Resource(Topic, "z_other", PREFIXED)
+    val expected = new String(ZkAclStore(PREFIXED).changeStore.createChangeNode(resource.toPattern).bytes, UTF_8)
+
+    simpleAclAuthorizer.addAcls(Set[Acl](denyReadAcl), resource)
+
+    val actual = getAclChangeEventAsString(PREFIXED)
+
+    assertEquals(expected, actual)
+  }
+
+  @Test
+  def testWritesLiteralWritesLiteralAclChangeEventWhenInterBrokerProtocolLessThanKafkaV2eralAclChangesForOlderProtocolVersions(): Unit = {
+    givenAuthorizerWithProtocolVersion(Option(KAFKA_2_0_IV0))
+    val resource = Resource(Topic, "z_other", LITERAL)
+    val expected = new String(ZkAclStore(LITERAL).changeStore.createChangeNode(resource.toPattern).bytes, UTF_8)
+
+    simpleAclAuthorizer.addAcls(Set[Acl](denyReadAcl), resource)
+
+    val actual = getAclChangeEventAsString(LITERAL)
+
+    assertEquals(expected, actual)
+  }
+
+  @Test
+  def testWritesLiteralAclChangeEventWhenInterBrokerProtocolIsKafkaV2(): Unit = {
+    givenAuthorizerWithProtocolVersion(Option(KAFKA_2_0_IV1))
+    val resource = Resource(Topic, "z_other", LITERAL)
+    val expected = new String(ZkAclStore(LITERAL).changeStore.createChangeNode(resource.toPattern).bytes, UTF_8)
+
+    simpleAclAuthorizer.addAcls(Set[Acl](denyReadAcl), resource)
+
+    val actual = getAclChangeEventAsString(LITERAL)
+
+    assertEquals(expected, actual)
+  }
+
+  private def givenAuthorizerWithProtocolVersion(protocolVersion: Option[ApiVersion]): Unit = {
+    simpleAclAuthorizer.close()
+
+    val props = TestUtils.createBrokerConfig(0, zkConnect)
+    props.put(SimpleAclAuthorizer.SuperUsersProp, superUsers)
+    protocolVersion.foreach(version => props.put(KafkaConfig.InterBrokerProtocolVersionProp, version.toString))
+
+    config = KafkaConfig.fromProps(props)
+
+    simpleAclAuthorizer.configure(config.originals)
+  }
+
+  private def getAclChangeEventAsString(patternType: PatternType) = {
+    val store = ZkAclStore(patternType)
+    val children = zooKeeperClient.handleRequest(GetChildrenRequest(store.changeStore.aclChangePath, registerWatch = true))
+    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[Acl], addedAcls: Set[Acl], removedAcls: Set[Acl], resource: Resource = resource): Set[Acl] = {
     var acls = originalAcls
 
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 646143c9dd9f6..f530b339e2f4f 100644
--- a/core/src/test/scala/unit/kafka/security/auth/ZkAuthorizationTest.scala
+++ b/core/src/test/scala/unit/kafka/security/auth/ZkAuthorizationTest.scala
@@ -17,24 +17,35 @@
 
 package kafka.security.auth
 
+import java.nio.charset.StandardCharsets
+
 import kafka.admin.ZkSecurityMigrator
-import kafka.utils.{Logging, TestUtils, ZkUtils}
-import kafka.zk.ZooKeeperTestHarness
-import org.apache.kafka.common.KafkaException
+import kafka.utils.{Logging, TestUtils}
+import kafka.zk._
+import org.apache.kafka.common.{KafkaException, TopicPartition, Uuid}
 import org.apache.kafka.common.security.JaasUtils
-import org.apache.zookeeper.data.{ACL}
+import org.apache.zookeeper.data.{ACL, Stat}
 import org.junit.Assert._
 import org.junit.{After, Before, Test}
-import scala.collection.JavaConverters._
-import scala.util.{Try, Success, Failure}
+
+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
+
+import scala.jdk.CollectionConverters._
+import scala.collection.Seq
 
 class ZkAuthorizationTest extends ZooKeeperTestHarness with Logging {
   val jaasFile = kafka.utils.JaasTestUtils.writeJaasContextsToFile(kafka.utils.JaasTestUtils.zkSections)
   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")
@@ -42,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)
@@ -54,15 +65,15 @@ class ZkAuthorizationTest extends ZooKeeperTestHarness with Logging {
    * secure ACLs and authentication with ZooKeeper.
    */
   @Test
-  def testIsZkSecurityEnabled() {
-    assertTrue(JaasUtils.isZkSecurityEnabled())
+  def testIsZkSecurityEnabled(): Unit = {
+    assertTrue(JaasUtils.isZkSaslEnabled())
     Configuration.setConfiguration(null)
     System.clearProperty(JaasUtils.JAVA_LOGIN_CONFIG_PARAM)
-    assertFalse(JaasUtils.isZkSecurityEnabled())
+    assertFalse(JaasUtils.isZkSaslEnabled())
     try {
       Configuration.setConfiguration(null)
       System.setProperty(JaasUtils.JAVA_LOGIN_CONFIG_PARAM, "no-such-file-exists.conf")
-      JaasUtils.isZkSecurityEnabled()
+      JaasUtils.isZkSaslEnabled()
       fail("Should have thrown an exception")
     } catch {
       case _: KafkaException => // Expected
@@ -70,55 +81,78 @@ class ZkAuthorizationTest extends ZooKeeperTestHarness with Logging {
   }
 
   /**
-   * Exercises the code in ZkUtils. The goal is mainly
-   * to verify that the behavior of ZkUtils is correct
+   * Exercises the code in KafkaZkClient. The goal is mainly
+   * to verify that the behavior of KafkaZkClient is correct
    * when isSecure is set to true.
    */
   @Test
-  def testZkUtils() {
-    assertTrue(zkUtils.isSecure)
-    for (path <- zkUtils.persistentZkPaths) {
-      zkUtils.makeSurePersistentPathExists(path)
-      if(!path.equals(ZkUtils.ConsumersPath)) {
-        val aclList = zkUtils.zkConnection.getAcl(path).getKey
-        assertTrue(aclList.size == 2)
-        for (acl: ACL <- aclList.asScala) {
-          assertTrue(TestUtils.isAclSecure(acl, false))
-        }
+  def testKafkaZkClient(): Unit = {
+    assertTrue(zkClient.secure)
+    for (path <- ZkData.PersistentZkPaths) {
+      zkClient.makeSurePersistentPathExists(path)
+      if (ZkData.sensitivePath(path)) {
+        val aclList = zkClient.getAcl(path)
+        assertEquals(s"Unexpected acl list size for $path", 1, aclList.size)
+        for (acl <- aclList)
+          assertTrue(TestUtils.isAclSecure(acl, sensitive = true))
+      } else if (!path.equals(ConsumerPathZNode.path)) {
+        val aclList = zkClient.getAcl(path)
+        assertEquals(s"Unexpected acl list size for $path", 2, aclList.size)
+        for (acl <- aclList)
+          assertTrue(TestUtils.isAclSecure(acl, sensitive = false))
       }
     }
-    // Test that can create: createEphemeralPathExpectConflict
-    zkUtils.createEphemeralPathExpectConflict("/a", "")
-    verify("/a")
-    // Test that can create: createPersistentPath
-    zkUtils.createPersistentPath("/b")
-    verify("/b")
+
+    // Test that creates Ephemeral node
+    val brokerInfo = createBrokerInfo(1, "test.host", 9999, SecurityProtocol.PLAINTEXT)
+    zkClient.registerBroker(brokerInfo)
+    verify(brokerInfo.path)
+
+    // Test that creates persistent nodes
+    val topic1 = "topic1"
+    val topicId = Uuid.randomUuid()
+    val assignment = Map(
+      new TopicPartition(topic1, 0) -> Seq(0, 1),
+      new TopicPartition(topic1, 1) -> Seq(0, 1),
+      new TopicPartition(topic1, 2) -> Seq(1, 2, 3)
+    )
+
+    // create a topic assignment
+    zkClient.createTopicAssignment(topic1, topicId, assignment)
+    verify(TopicZNode.path(topic1))
+
     // Test that can create: createSequentialPersistentPath
-    val seqPath = zkUtils.createSequentialPersistentPath("/c", "")
+    val seqPath = zkClient.createSequentialPersistentPath("/c", "".getBytes(StandardCharsets.UTF_8))
     verify(seqPath)
-    // Test that can update: updateEphemeralPath
-    zkUtils.updateEphemeralPath("/a", "updated")
-    val valueA: String = zkUtils.zkClient.readData("/a")
-    assertTrue(valueA.equals("updated"))
-    // Test that can update: updatePersistentPath
-    zkUtils.updatePersistentPath("/b", "updated")
-    val valueB: String = zkUtils.zkClient.readData("/b")
-    assertTrue(valueB.equals("updated"))
 
-    info("Leaving testZkUtils")
+    // Test that can update Ephemeral node
+    val updatedBrokerInfo = createBrokerInfo(1, "test.host2", 9995, SecurityProtocol.SSL)
+    zkClient.updateBrokerInfo(updatedBrokerInfo)
+    assertEquals(Some(updatedBrokerInfo.broker), zkClient.getBroker(1))
+
+    // Test that can update persistent nodes
+    val updatedAssignment = assignment - new TopicPartition(topic1, 2)
+    zkClient.setTopicAssignment(topic1, topicId,
+      updatedAssignment.map { case (k, v) => k -> ReplicaAssignment(v, List(), List()) })
+    assertEquals(updatedAssignment.size, zkClient.getTopicPartitionCount(topic1).get)
   }
 
+  private def createBrokerInfo(id: Int, host: String, port: Int, securityProtocol: SecurityProtocol,
+                               rack: Option[String] = None): BrokerInfo =
+    BrokerInfo(Broker(id, Seq(new EndPoint(host, port, ListenerName.forSecurityProtocol
+    (securityProtocol), securityProtocol)), rack = rack), ApiVersion.latestVersion, jmxPort = port + 10)
+
   /**
    * Tests the migration tool when making an unsecure
    * cluster secure.
    */
   @Test
-  def testZkMigration() {
-    val unsecureZkUtils = ZkUtils(zkConnect, 6000, 6000, false) 
+  def testZkMigration(): Unit = {
+    val unsecureZkClient = KafkaZkClient(zkConnect, false, 6000, 6000, Int.MaxValue, Time.SYSTEM)
     try {
-      testMigration(zkConnect, unsecureZkUtils, zkUtils)
+      testMigration(zkConnect, unsecureZkClient, zkClient)
     } finally {
-      unsecureZkUtils.close()
+      unsecureZkClient.close()
     }
   }
 
@@ -127,12 +161,12 @@ class ZkAuthorizationTest extends ZooKeeperTestHarness with Logging {
    * cluster unsecure.
    */
   @Test
-  def testZkAntiMigration() {
-    val unsecureZkUtils = ZkUtils(zkConnect, 6000, 6000, false)
+  def testZkAntiMigration(): Unit = {
+    val unsecureZkClient = KafkaZkClient(zkConnect, false, 6000, 6000, Int.MaxValue, Time.SYSTEM)
     try {
-      testMigration(zkConnect, zkUtils, unsecureZkUtils)
+      testMigration(zkConnect, zkClient, unsecureZkClient)
     } finally {
-      unsecureZkUtils.close()
+      unsecureZkClient.close()
     }
   }
 
@@ -140,42 +174,42 @@ class ZkAuthorizationTest extends ZooKeeperTestHarness with Logging {
    * Tests that the persistent paths cannot be deleted.
    */
   @Test
-  def testDelete() {
+  def testDelete(): Unit = {
     info(s"zkConnect string: $zkConnect")
     ZkSecurityMigrator.run(Array("--zookeeper.acl=secure", s"--zookeeper.connect=$zkConnect"))
     deleteAllUnsecure()
   }
 
   /**
-   * Tests that znodes cannot be deleted when the 
+   * Tests that znodes cannot be deleted when the
    * persistent paths have children.
    */
   @Test
-  def testDeleteRecursive() {
+  def testDeleteRecursive(): Unit = {
     info(s"zkConnect string: $zkConnect")
-    for (path <- ZkUtils.SecureZkRootPaths) {
+    for (path <- ZkData.SecureRootPaths) {
       info(s"Creating $path")
-      zkUtils.makeSurePersistentPathExists(path)
-      zkUtils.createPersistentPath(s"$path/fpjwashere", "")
+      zkClient.makeSurePersistentPathExists(path)
+      zkClient.createRecursive(s"$path/fpjwashere", "".getBytes(StandardCharsets.UTF_8))
     }
-    zkUtils.zkConnection.setAcl("/", zkUtils.defaultAcls("/"), -1)
+    zkClient.setAcl("/", zkClient.defaultAcls("/"))
     deleteAllUnsecure()
   }
-  
+
   /**
    * Tests the migration tool when chroot is being used.
    */
   @Test
   def testChroot(): Unit = {
     val zkUrl = zkConnect + "/kafka"
-    zkUtils.createPersistentPath("/kafka")
-    val unsecureZkUtils = ZkUtils(zkUrl, 6000, 6000, false)
-    val secureZkUtils = ZkUtils(zkUrl, 6000, 6000, true)
+    zkClient.createRecursive("/kafka")
+    val unsecureZkClient = KafkaZkClient(zkUrl, false, 6000, 6000, Int.MaxValue, Time.SYSTEM)
+    val secureZkClient = KafkaZkClient(zkUrl, true, 6000, 6000, Int.MaxValue, Time.SYSTEM)
     try {
-      testMigration(zkUrl, unsecureZkUtils, secureZkUtils)
+      testMigration(zkUrl, unsecureZkClient, secureZkClient)
     } finally {
-      unsecureZkUtils.close()
-      secureZkUtils.close()
+      unsecureZkClient.close()
+      secureZkClient.close()
     }
   }
 
@@ -183,62 +217,64 @@ class ZkAuthorizationTest extends ZooKeeperTestHarness with Logging {
    * Exercises the migration tool. It is used in these test cases:
    * testZkMigration, testZkAntiMigration, testChroot.
    */
-  private def testMigration(zkUrl: String, firstZk: ZkUtils, secondZk: ZkUtils) {
+  private def testMigration(zkUrl: String, firstZk: KafkaZkClient, secondZk: KafkaZkClient): Unit = {
     info(s"zkConnect string: $zkUrl")
-    for (path <- ZkUtils.SecureZkRootPaths ++ ZkUtils.SensitiveZkRootPaths) {
+    for (path <- ZkData.SecureRootPaths ++ ZkData.SensitiveRootPaths) {
       info(s"Creating $path")
       firstZk.makeSurePersistentPathExists(path)
       // Create a child for each znode to exercise the recurrent
       // traversal of the data tree
-      firstZk.createPersistentPath(s"$path/fpjwashere", "")
+      firstZk.createRecursive(s"$path/fpjwashere", "".getBytes(StandardCharsets.UTF_8))
     }
     // Getting security option to determine how to verify ACLs.
     // Additionally, we create the consumers znode (not in
     // securePersistentZkPaths) to make sure that we don't
     // add ACLs to it.
     val secureOpt: String =
-      if (secondZk.isSecure) {
-        firstZk.createPersistentPath(ZkUtils.ConsumersPath)
+      if (secondZk.secure) {
+        firstZk.createRecursive(ConsumerPathZNode.path)
         "secure"
       } else {
-        secondZk.createPersistentPath(ZkUtils.ConsumersPath)
+        secondZk.createRecursive(ConsumerPathZNode.path)
         "unsecure"
       }
     ZkSecurityMigrator.run(Array(s"--zookeeper.acl=$secureOpt", s"--zookeeper.connect=$zkUrl"))
     info("Done with migration")
-    for (path <- ZkUtils.SecureZkRootPaths ++ ZkUtils.SensitiveZkRootPaths) {
-      val sensitive = ZkUtils.sensitivePath(path)
-      val listParent = secondZk.zkConnection.getAcl(path).getKey
-      assertTrue(path, isAclCorrect(listParent, secondZk.isSecure, sensitive))
+    for (path <- ZkData.SecureRootPaths ++ ZkData.SensitiveRootPaths) {
+      val sensitive = ZkData.sensitivePath(path)
+      val listParent = secondZk.getAcl(path)
+      assertTrue(path, isAclCorrect(listParent, secondZk.secure, sensitive))
 
       val childPath = path + "/fpjwashere"
-      val listChild = secondZk.zkConnection.getAcl(childPath).getKey
-      assertTrue(childPath, isAclCorrect(listChild, secondZk.isSecure, sensitive))
+      val listChild = secondZk.getAcl(childPath)
+      assertTrue(childPath, isAclCorrect(listChild, secondZk.secure, sensitive))
     }
     // Check consumers path.
-    val consumersAcl = firstZk.zkConnection.getAcl(ZkUtils.ConsumersPath).getKey
-    assertTrue(ZkUtils.ConsumersPath, isAclCorrect(consumersAcl, false, false))
+    val consumersAcl = firstZk.getAcl(ConsumerPathZNode.path)
+    assertTrue(ConsumerPathZNode.path, isAclCorrect(consumersAcl, false, false))
+    assertTrue("/kafka-acl-extended", isAclCorrect(firstZk.getAcl("/kafka-acl-extended"), secondZk.secure,
+      ZkData.sensitivePath(ExtendedAclZNode.path)))
   }
 
   /**
    * Verifies that the path has the appropriate secure ACL.
    */
-  private def verify(path: String): Boolean = {
-    val sensitive = ZkUtils.sensitivePath(path)
-    val list = zkUtils.zkConnection.getAcl(path).getKey
-    list.asScala.forall(TestUtils.isAclSecure(_, sensitive))
+  private def verify(path: String): Unit = {
+    val sensitive = ZkData.sensitivePath(path)
+    val list = zkClient.getAcl(path)
+    assertTrue(list.forall(TestUtils.isAclSecure(_, sensitive)))
   }
 
   /**
    * Verifies ACL.
    */
-  private def isAclCorrect(list: java.util.List[ACL], secure: Boolean, sensitive: Boolean): Boolean = {
+  private def isAclCorrect(list: Seq[ACL], secure: Boolean, sensitive: Boolean): Boolean = {
     val isListSizeCorrect =
       if (secure && !sensitive)
         list.size == 2
       else
         list.size == 1
-    isListSizeCorrect && list.asScala.forall(
+    isListSizeCorrect && list.forall(
       if (secure)
         TestUtils.isAclSecure(_, sensitive)
       else
@@ -251,14 +287,14 @@ class ZkAuthorizationTest extends ZooKeeperTestHarness with Logging {
    * This is used in the testDelete and testDeleteRecursive
    * test cases.
    */
-  private def deleteAllUnsecure() {
+  private def deleteAllUnsecure(): Unit = {
     System.setProperty(JaasUtils.ZK_SASL_CLIENT, "false")
-    val unsecureZkUtils = ZkUtils(zkConnect, 6000, 6000, false)
+    val unsecureZkClient = KafkaZkClient(zkConnect, false, 6000, 6000, Int.MaxValue, Time.SYSTEM)
     val result: Try[Boolean] = {
-      deleteRecursive(unsecureZkUtils, "/")
+      deleteRecursive(unsecureZkClient, "/")
     }
     // Clean up before leaving the test case
-    unsecureZkUtils.close()
+    unsecureZkClient.close()
     System.clearProperty(JaasUtils.ZK_SASL_CLIENT)
     
     // Fail the test if able to delete
@@ -271,13 +307,13 @@ class ZkAuthorizationTest extends ZooKeeperTestHarness with Logging {
   /**
    * Tries to delete znodes recursively
    */
-  private def deleteRecursive(zkUtils: ZkUtils, path: String): Try[Boolean] = {
+  private def deleteRecursive(zkClient: KafkaZkClient, path: String): Try[Boolean] = {
     info(s"Deleting $path")
     var result: Try[Boolean] = Success(true)
-    for (child <- zkUtils.getChildren(path))
+    for (child <- zkClient.getChildren(path))
       result = (path match {
-        case "/" => deleteRecursive(zkUtils, s"/$child")
-        case path => deleteRecursive(zkUtils, s"$path/$child")
+        case "/" => deleteRecursive(zkClient, s"/$child")
+        case path => deleteRecursive(zkClient, s"$path/$child")
       }) match {
         case Success(_) => result
         case Failure(e) => Failure(e)
@@ -288,11 +324,19 @@ class ZkAuthorizationTest extends ZooKeeperTestHarness with Logging {
       // For all other paths, try to delete it
       case path =>
         try {
-          zkUtils.deletePath(path)
+          zkClient.deletePath(path, recursiveDelete = false)
           Failure(new Exception(s"Have been able to delete $path"))
         } catch {
           case _: Exception => result
         }
     }
   }
+
+  @Test
+  def testConsumerOffsetPathAcls(): Unit = {
+    zkClient.makeSurePersistentPathExists(ConsumerPathZNode.path)
+
+    val consumerPathAcls = zkClient.currentZooKeeper.getACL(ConsumerPathZNode.path, new Stat())
+    assertTrue("old consumer znode path acls are not open", consumerPathAcls.asScala.forall(TestUtils.isAclUnsecure))
+  }
 }
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..a5c57b6917397
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/security/authorizer/AclAuthorizerTest.scala
@@ -0,0 +1,1106 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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.io.File
+import java.net.InetAddress
+import java.nio.charset.StandardCharsets.UTF_8
+import java.nio.file.Files
+import java.util.{Collections, UUID}
+import java.util.concurrent.{Executors, Semaphore, TimeUnit}
+
+import kafka.Kafka
+import kafka.api.{ApiVersion, KAFKA_2_0_IV0, KAFKA_2_0_IV1}
+import kafka.security.authorizer.AclEntry.{WildcardHost, WildcardPrincipalString}
+import kafka.server.KafkaConfig
+import kafka.utils.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.ClientInformation
+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.Resource.CLUSTER_NAME
+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.jdk.CollectionConverters._
+import scala.collection.mutable
+
+class AclAuthorizerTest extends ZooKeeperTestHarness {
+
+  private val allowReadAcl = new AccessControlEntry(WildcardPrincipalString, WildcardHost, READ, ALLOW)
+  private val allowWriteAcl = new AccessControlEntry(WildcardPrincipalString, WildcardHost, WRITE, ALLOW)
+  private val denyReadAcl = new AccessControlEntry(WildcardPrincipalString, WildcardHost, READ, DENY)
+
+  private val wildCardResource = new ResourcePattern(TOPIC, WILDCARD_RESOURCE, LITERAL)
+  private val prefixedResource = new ResourcePattern(TOPIC, "foo", PREFIXED)
+  private val clusterResource = new ResourcePattern(CLUSTER, CLUSTER_NAME, LITERAL)
+  private val wildcardPrincipal = JSecurityUtils.parseKafkaPrincipal(WildcardPrincipalString)
+
+  private val 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(WildcardPrincipalString, 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(WildcardPrincipalString, 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(WildcardPrincipalString, 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(WildcardPrincipalString, 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(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(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((() => aclAuthorizer3.configure(config.originals)): Runnable)
+      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(WildcardPrincipalString, 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(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(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(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(resource).bytes, UTF_8)
+
+    addAcls(aclAuthorizer, Set(denyReadAcl), resource)
+
+    val actual = getAclChangeEventAsString(LITERAL)
+
+    assertEquals(expected, actual)
+  }
+
+  @Test
+  def testAuthorizerNoZkConfig(): Unit = {
+    val noTlsProps = Kafka.getPropsFromArgs(Array(prepareDefaultConfig))
+    assertEquals(None, AclAuthorizer.zkClientConfigFromKafkaConfigAndMap(
+      KafkaConfig.fromProps(noTlsProps),
+      mutable.Map(noTlsProps.asInstanceOf[java.util.Map[String, Any]].asScala.toSeq: _*)))
+  }
+
+  @Test
+  def testAuthorizerZkConfigFromKafkaConfigWithDefaults(): Unit = {
+    val props = new java.util.Properties()
+    val kafkaValue = "kafkaValue"
+    val configs = Map("zookeeper.connect" -> "somewhere", // required, otherwise we would omit it
+      KafkaConfig.ZkSslClientEnableProp -> "true",
+      KafkaConfig.ZkClientCnxnSocketProp -> kafkaValue,
+      KafkaConfig.ZkSslKeyStoreLocationProp -> kafkaValue,
+      KafkaConfig.ZkSslKeyStorePasswordProp -> kafkaValue,
+      KafkaConfig.ZkSslKeyStoreTypeProp -> kafkaValue,
+      KafkaConfig.ZkSslTrustStoreLocationProp -> kafkaValue,
+      KafkaConfig.ZkSslTrustStorePasswordProp -> kafkaValue,
+      KafkaConfig.ZkSslTrustStoreTypeProp -> kafkaValue,
+      KafkaConfig.ZkSslEnabledProtocolsProp -> kafkaValue,
+      KafkaConfig.ZkSslCipherSuitesProp -> kafkaValue)
+    configs.foreach{case (key, value) => props.put(key, value.toString) }
+
+    val zkClientConfig = AclAuthorizer.zkClientConfigFromKafkaConfigAndMap(
+      KafkaConfig.fromProps(props), mutable.Map(configs.toSeq: _*))
+    assertTrue(zkClientConfig.isDefined)
+    // confirm we get all the values we expect
+    KafkaConfig.ZkSslConfigToSystemPropertyMap.keys.foreach(prop => prop match {
+      case KafkaConfig.ZkSslClientEnableProp | KafkaConfig.ZkSslEndpointIdentificationAlgorithmProp =>
+        assertEquals("true", KafkaConfig.getZooKeeperClientProperty(zkClientConfig.get, prop).getOrElse(""))
+      case KafkaConfig.ZkSslCrlEnableProp | KafkaConfig.ZkSslOcspEnableProp =>
+        assertEquals("false", KafkaConfig.getZooKeeperClientProperty(zkClientConfig.get, prop).getOrElse(""))
+      case KafkaConfig.ZkSslProtocolProp =>
+        assertEquals("TLSv1.2", KafkaConfig.getZooKeeperClientProperty(zkClientConfig.get, prop).getOrElse(""))
+      case _ => assertEquals(kafkaValue, KafkaConfig.getZooKeeperClientProperty(zkClientConfig.get, prop).getOrElse(""))
+    })
+  }
+
+  @Test
+  def testAuthorizerZkConfigFromKafkaConfig(): Unit = {
+    val props = new java.util.Properties()
+    val kafkaValue = "kafkaValue"
+    val configs = Map("zookeeper.connect" -> "somewhere", // required, otherwise we would omit it
+      KafkaConfig.ZkSslClientEnableProp -> "true",
+      KafkaConfig.ZkClientCnxnSocketProp -> kafkaValue,
+      KafkaConfig.ZkSslKeyStoreLocationProp -> kafkaValue,
+      KafkaConfig.ZkSslKeyStorePasswordProp -> kafkaValue,
+      KafkaConfig.ZkSslKeyStoreTypeProp -> kafkaValue,
+      KafkaConfig.ZkSslTrustStoreLocationProp -> kafkaValue,
+      KafkaConfig.ZkSslTrustStorePasswordProp -> kafkaValue,
+      KafkaConfig.ZkSslTrustStoreTypeProp -> kafkaValue,
+      KafkaConfig.ZkSslProtocolProp -> kafkaValue,
+      KafkaConfig.ZkSslEnabledProtocolsProp -> kafkaValue,
+      KafkaConfig.ZkSslCipherSuitesProp -> kafkaValue,
+      KafkaConfig.ZkSslEndpointIdentificationAlgorithmProp -> "HTTPS",
+      KafkaConfig.ZkSslCrlEnableProp -> "false",
+      KafkaConfig.ZkSslOcspEnableProp -> "false")
+    configs.foreach{case (key, value) => props.put(key, value.toString) }
+
+    val zkClientConfig = AclAuthorizer.zkClientConfigFromKafkaConfigAndMap(
+      KafkaConfig.fromProps(props), mutable.Map(configs.toSeq: _*))
+    assertTrue(zkClientConfig.isDefined)
+    // confirm we get all the values we expect
+    KafkaConfig.ZkSslConfigToSystemPropertyMap.keys.foreach(prop => prop match {
+        case KafkaConfig.ZkSslClientEnableProp | KafkaConfig.ZkSslEndpointIdentificationAlgorithmProp =>
+          assertEquals("true", KafkaConfig.getZooKeeperClientProperty(zkClientConfig.get, prop).getOrElse(""))
+        case KafkaConfig.ZkSslCrlEnableProp | KafkaConfig.ZkSslOcspEnableProp =>
+          assertEquals("false", KafkaConfig.getZooKeeperClientProperty(zkClientConfig.get, prop).getOrElse(""))
+        case _ => assertEquals(kafkaValue, KafkaConfig.getZooKeeperClientProperty(zkClientConfig.get, prop).getOrElse(""))
+      })
+  }
+
+  @Test
+  def testAuthorizerZkConfigFromPrefixOverrides(): Unit = {
+    val props = new java.util.Properties()
+    val kafkaValue = "kafkaValue"
+    val prefixedValue = "prefixedValue"
+    val prefix = "authorizer."
+    val configs = Map("zookeeper.connect" -> "somewhere", // required, otherwise we would omit it
+      KafkaConfig.ZkSslClientEnableProp -> "false",
+      KafkaConfig.ZkClientCnxnSocketProp -> kafkaValue,
+      KafkaConfig.ZkSslKeyStoreLocationProp -> kafkaValue,
+      KafkaConfig.ZkSslKeyStorePasswordProp -> kafkaValue,
+      KafkaConfig.ZkSslKeyStoreTypeProp -> kafkaValue,
+      KafkaConfig.ZkSslTrustStoreLocationProp -> kafkaValue,
+      KafkaConfig.ZkSslTrustStorePasswordProp -> kafkaValue,
+      KafkaConfig.ZkSslTrustStoreTypeProp -> kafkaValue,
+      KafkaConfig.ZkSslProtocolProp -> kafkaValue,
+      KafkaConfig.ZkSslEnabledProtocolsProp -> kafkaValue,
+      KafkaConfig.ZkSslCipherSuitesProp -> kafkaValue,
+      KafkaConfig.ZkSslEndpointIdentificationAlgorithmProp -> "HTTPS",
+      KafkaConfig.ZkSslCrlEnableProp -> "false",
+      KafkaConfig.ZkSslOcspEnableProp -> "false",
+      prefix + KafkaConfig.ZkSslClientEnableProp -> "true",
+      prefix + KafkaConfig.ZkClientCnxnSocketProp -> prefixedValue,
+      prefix + KafkaConfig.ZkSslKeyStoreLocationProp -> prefixedValue,
+      prefix + KafkaConfig.ZkSslKeyStorePasswordProp -> prefixedValue,
+      prefix + KafkaConfig.ZkSslKeyStoreTypeProp -> prefixedValue,
+      prefix + KafkaConfig.ZkSslTrustStoreLocationProp -> prefixedValue,
+      prefix + KafkaConfig.ZkSslTrustStorePasswordProp -> prefixedValue,
+      prefix + KafkaConfig.ZkSslTrustStoreTypeProp -> prefixedValue,
+      prefix + KafkaConfig.ZkSslProtocolProp -> prefixedValue,
+      prefix + KafkaConfig.ZkSslEnabledProtocolsProp -> prefixedValue,
+      prefix + KafkaConfig.ZkSslCipherSuitesProp -> prefixedValue,
+      prefix + KafkaConfig.ZkSslEndpointIdentificationAlgorithmProp -> "",
+      prefix + KafkaConfig.ZkSslCrlEnableProp -> "true",
+      prefix + KafkaConfig.ZkSslOcspEnableProp -> "true")
+    configs.foreach{case (key, value) => props.put(key, value.toString) }
+
+    val zkClientConfig = AclAuthorizer.zkClientConfigFromKafkaConfigAndMap(
+      KafkaConfig.fromProps(props), mutable.Map(configs.toSeq: _*))
+    assertTrue(zkClientConfig.isDefined)
+    // confirm we get all the values we expect
+    KafkaConfig.ZkSslConfigToSystemPropertyMap.keys.foreach(prop => prop match {
+      case KafkaConfig.ZkSslClientEnableProp | KafkaConfig.ZkSslCrlEnableProp | KafkaConfig.ZkSslOcspEnableProp =>
+        assertEquals("true", KafkaConfig.getZooKeeperClientProperty(zkClientConfig.get, prop).getOrElse(""))
+      case KafkaConfig.ZkSslEndpointIdentificationAlgorithmProp =>
+        assertEquals("false", KafkaConfig.getZooKeeperClientProperty(zkClientConfig.get, prop).getOrElse(""))
+      case _ => assertEquals(prefixedValue, KafkaConfig.getZooKeeperClientProperty(zkClientConfig.get, prop).getOrElse(""))
+    })
+  }
+
+  @Test
+  def testCreateDeleteTiming(): Unit = {
+    val literalResource = new ResourcePattern(TOPIC, "foo-" + UUID.randomUUID(), LITERAL)
+    val prefixedResource = new ResourcePattern(TOPIC, "bar-", PREFIXED)
+    val wildcardResource = new ResourcePattern(TOPIC, "*", LITERAL)
+    val ace = new AccessControlEntry(principal.toString, WildcardHost, READ, ALLOW)
+    val updateSemaphore = new Semaphore(1)
+
+    def createAcl(createAuthorizer: AclAuthorizer, resource: ResourcePattern): AclBinding = {
+      val acl = new AclBinding(resource, ace)
+      createAuthorizer.createAcls(requestContext, Collections.singletonList(acl)).asScala
+        .foreach(_.toCompletableFuture.get(15, TimeUnit.SECONDS))
+      acl
+    }
+
+    def deleteAcl(deleteAuthorizer: AclAuthorizer,
+                  resource: ResourcePattern,
+                  deletePatternType: PatternType): List[AclBinding] = {
+
+      val filter = new AclBindingFilter(
+        new ResourcePatternFilter(resource.resourceType(), resource.name(), deletePatternType),
+        AccessControlEntryFilter.ANY)
+      deleteAuthorizer.deleteAcls(requestContext, Collections.singletonList(filter)).asScala
+        .map(_.toCompletableFuture.get(15, TimeUnit.SECONDS))
+        .flatMap(_.aclBindingDeleteResults.asScala)
+        .map(_.aclBinding)
+        .toList
+    }
+
+    def listAcls(authorizer: AclAuthorizer): List[AclBinding] = {
+      authorizer.acls(AclBindingFilter.ANY).asScala.toList
+    }
+
+    def verifyCreateDeleteAcl(deleteAuthorizer: AclAuthorizer,
+                              resource: ResourcePattern,
+                              deletePatternType: PatternType): Unit = {
+      updateSemaphore.acquire()
+      assertEquals(List.empty, listAcls(deleteAuthorizer))
+      val acl = createAcl(aclAuthorizer, resource)
+      val deleted = deleteAcl(deleteAuthorizer, resource, deletePatternType)
+      if (deletePatternType != PatternType.MATCH) {
+        assertEquals(List(acl), deleted)
+      } else {
+        assertEquals(List.empty[AclBinding], deleted)
+      }
+      updateSemaphore.release()
+      if (deletePatternType == PatternType.MATCH) {
+        TestUtils.waitUntilTrue(() => listAcls(deleteAuthorizer).nonEmpty, "ACL not propagated")
+        assertEquals(List(acl), deleteAcl(deleteAuthorizer, resource, deletePatternType))
+      }
+      TestUtils.waitUntilTrue(() => listAcls(deleteAuthorizer).isEmpty, "ACL delete not propagated")
+    }
+
+    val deleteAuthorizer = new AclAuthorizer {
+      override def processAclChangeNotification(resource: ResourcePattern): Unit = {
+        updateSemaphore.acquire()
+        try {
+          super.processAclChangeNotification(resource)
+        } finally {
+          updateSemaphore.release()
+        }
+      }
+    }
+
+    try {
+      deleteAuthorizer.configure(config.originals)
+      List(literalResource, prefixedResource, wildcardResource).foreach { resource =>
+        verifyCreateDeleteAcl(deleteAuthorizer, resource, resource.patternType())
+        verifyCreateDeleteAcl(deleteAuthorizer, resource, PatternType.ANY)
+        verifyCreateDeleteAcl(deleteAuthorizer, resource, PatternType.MATCH)
+      }
+    } finally {
+      deleteAuthorizer.close()
+    }
+  }
+
+  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, registerWatch = true))
+    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, ClientInformation.EMPTY, false)
+  }
+
+  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.ifPresent { 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.ifPresent { e => throw e }
+        result.aclBindingDeleteResults.forEach { r =>
+          r.exception.ifPresent { e => throw e }
+        }
+        !result.aclBindingDeleteResults.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
+  }
+
+  private def prepareDefaultConfig: String =
+    prepareConfig(Array("broker.id=1", "zookeeper.connect=somewhere"))
+
+  private def prepareConfig(lines : Array[String]): String = {
+    val file = File.createTempFile("kafkatest", ".properties")
+    file.deleteOnExit()
+
+    val writer = Files.newOutputStream(file.toPath)
+    try {
+      lines.foreach { l =>
+        writer.write(l.getBytes)
+        writer.write("\n".getBytes)
+      }
+      file.getAbsolutePath
+    } finally writer.close()
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/security/authorizer/AclAuthorizerWithZkSaslTest.scala b/core/src/test/scala/unit/kafka/security/authorizer/AclAuthorizerWithZkSaslTest.scala
new file mode 100644
index 0000000000000..fe9f0f3f1991f
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/security/authorizer/AclAuthorizerWithZkSaslTest.scala
@@ -0,0 +1,186 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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.util
+import java.util.UUID
+import java.util.concurrent.{Executors, TimeUnit}
+
+import javax.security.auth.Subject
+import javax.security.auth.callback.CallbackHandler
+import kafka.api.SaslSetup
+import kafka.security.authorizer.AclEntry.WildcardHost
+import kafka.server.KafkaConfig
+import kafka.utils.JaasTestUtils.{JaasModule, JaasSection}
+import kafka.utils.{JaasTestUtils, TestUtils}
+import kafka.zk.{KafkaZkClient, ZooKeeperTestHarness}
+import kafka.zookeeper.ZooKeeperClient
+import org.apache.kafka.common.acl.{AccessControlEntry, AccessControlEntryFilter, AclBinding, AclBindingFilter}
+import org.apache.kafka.common.acl.AclOperation.{READ, WRITE}
+import org.apache.kafka.common.acl.AclPermissionType.ALLOW
+import org.apache.kafka.common.network.{ClientInformation, ListenerName}
+import org.apache.kafka.common.protocol.ApiKeys
+import org.apache.kafka.common.requests.{RequestContext, RequestHeader}
+import org.apache.kafka.common.resource.PatternType.LITERAL
+import org.apache.kafka.common.resource.ResourcePattern
+import org.apache.kafka.common.resource.ResourceType.TOPIC
+import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
+import org.apache.kafka.test.{TestUtils => JTestUtils}
+import org.apache.zookeeper.server.auth.DigestLoginModule
+import org.junit.Assert.assertEquals
+import org.junit.{After, Before, Test}
+
+import scala.jdk.CollectionConverters._
+import scala.collection.Seq
+
+class AclAuthorizerWithZkSaslTest extends ZooKeeperTestHarness with SaslSetup {
+
+  private val aclAuthorizer = new AclAuthorizer
+  private val aclAuthorizer2 = new AclAuthorizer
+  private val resource: ResourcePattern = new ResourcePattern(TOPIC, "foo-" + UUID.randomUUID(), LITERAL)
+  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 val executor = Executors.newSingleThreadScheduledExecutor
+  private var config: KafkaConfig = _
+
+  @Before
+  override def setUp(): Unit = {
+    // Allow failed clients to avoid server closing the connection before reporting AuthFailed.
+    System.setProperty("zookeeper.allowSaslFailedClients", "true")
+
+    // Configure ZK SASL with TestableDigestLoginModule for clients to inject failures
+    TestableDigestLoginModule.reset()
+    val jaasSections = JaasTestUtils.zkSections
+    val serverJaas = jaasSections.filter(_.contextName == "Server")
+    val clientJaas = jaasSections.filter(_.contextName == "Client")
+      .map(section => new TestableJaasSection(section.contextName, section.modules))
+    startSasl(serverJaas ++ clientJaas)
+
+    // Increase maxUpdateRetries to avoid transient failures
+    aclAuthorizer.maxUpdateRetries = Int.MaxValue
+    aclAuthorizer2.maxUpdateRetries = Int.MaxValue
+
+    super.setUp()
+    config = KafkaConfig.fromProps(TestUtils.createBrokerConfig(0, zkConnect))
+
+    aclAuthorizer.configure(config.originals)
+    aclAuthorizer2.configure(config.originals)
+  }
+
+  @After
+  override def tearDown(): Unit = {
+    System.clearProperty("zookeeper.allowSaslFailedClients")
+    TestableDigestLoginModule.reset()
+    executor.shutdownNow()
+    aclAuthorizer.close()
+    aclAuthorizer2.close()
+    super.tearDown()
+  }
+
+  @Test
+  def testAclUpdateWithSessionExpiration(): Unit = {
+    zkClient(aclAuthorizer).currentZooKeeper.getTestable.injectSessionExpiration()
+    zkClient(aclAuthorizer2).currentZooKeeper.getTestable.injectSessionExpiration()
+    verifyAclUpdate()
+  }
+
+  @Test
+  def testAclUpdateWithAuthFailure(): Unit = {
+    injectTransientAuthenticationFailure()
+    verifyAclUpdate()
+  }
+
+  private def injectTransientAuthenticationFailure(): Unit = {
+    TestableDigestLoginModule.injectInvalidCredentials()
+    zkClient(aclAuthorizer).currentZooKeeper.getTestable.injectSessionExpiration()
+    zkClient(aclAuthorizer2).currentZooKeeper.getTestable.injectSessionExpiration()
+    executor.schedule((() => TestableDigestLoginModule.reset()): Runnable,
+      ZooKeeperClient.RetryBackoffMs * 2, TimeUnit.MILLISECONDS)
+  }
+
+  private def verifyAclUpdate(): Unit = {
+    val allowReadAcl = new AccessControlEntry(principal.toString, WildcardHost, READ, ALLOW)
+    val allowWriteAcl = new AccessControlEntry(principal.toString, WildcardHost, WRITE, ALLOW)
+    val acls = Set(allowReadAcl, allowWriteAcl)
+
+    TestUtils.retry(maxWaitMs = 15000) {
+      try {
+        addAcls(aclAuthorizer, acls, resource)
+      } catch {
+        case _: Exception => // Ignore error and retry
+      }
+      assertEquals(acls, getAcls(aclAuthorizer, resource))
+    }
+    val (acls2, _) = TestUtils.computeUntilTrue(getAcls(aclAuthorizer2, resource)) { _ == acls }
+    assertEquals(acls, acls2)
+  }
+
+  private def zkClient(authorizer: AclAuthorizer): KafkaZkClient = {
+    JTestUtils.fieldValue(authorizer, classOf[AclAuthorizer], "zkClient")
+  }
+
+  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.ifPresent { e => throw e } }
+  }
+
+  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 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, ClientInformation.EMPTY, false)
+  }
+}
+
+object TestableDigestLoginModule {
+  @volatile var injectedPassword: Option[String] = None
+
+  def reset(): Unit = {
+    injectedPassword = None
+  }
+
+  def injectInvalidCredentials(): Unit = {
+    injectedPassword = Some("invalidPassword")
+  }
+}
+
+class TestableDigestLoginModule extends DigestLoginModule {
+  override def initialize(subject: Subject, callbackHandler: CallbackHandler, sharedState: util.Map[String, _], options: util.Map[String, _]): Unit = {
+    super.initialize(subject, callbackHandler, sharedState, options)
+    val injectedPassword = TestableDigestLoginModule.injectedPassword
+    injectedPassword.foreach { newPassword =>
+      val oldPassword = subject.getPrivateCredentials.asScala.head
+      subject.getPrivateCredentials.add(newPassword)
+      subject.getPrivateCredentials.remove(oldPassword)
+    }
+  }
+}
+
+class TestableJaasSection(contextName: String, modules: Seq[JaasModule]) extends JaasSection(contextName, modules) {
+  override def toString: String = {
+    super.toString.replaceFirst(classOf[DigestLoginModule].getName, classOf[TestableDigestLoginModule].getName)
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/security/authorizer/AclEntryTest.scala b/core/src/test/scala/unit/kafka/security/authorizer/AclEntryTest.scala
new file mode 100644
index 0000000000000..80495b03e9b01
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/security/authorizer/AclEntryTest.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.security.authorizer
+
+import java.nio.charset.StandardCharsets.UTF_8
+
+import kafka.utils.Json
+import org.apache.kafka.common.acl.AclOperation.READ
+import org.apache.kafka.common.acl.AclPermissionType.{ALLOW, DENY}
+import org.apache.kafka.common.security.auth.KafkaPrincipal
+import org.junit.{Assert, Test}
+
+import scala.jdk.CollectionConverters._
+
+class AclEntryTest {
+
+  val AclJson = """{"version": 1, "acls": [{"host": "host1","permissionType": "Deny","operation": "READ", "principal": "User:alice"  },
+    {  "host":  "*" ,  "permissionType": "Allow",  "operation":  "Read", "principal": "User:bob"  },
+    {  "host": "host1",  "permissionType": "Deny",  "operation":   "Read" ,  "principal": "User:bob"}]}"""
+
+  @Test
+  def testAclJsonConversion(): Unit = {
+    val acl1 = AclEntry(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "alice"), DENY, "host1" , READ)
+    val acl2 = AclEntry(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "bob"), ALLOW, "*", READ)
+    val acl3 = AclEntry(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "bob"), DENY, "host1", READ)
+
+    val acls = Set[AclEntry](acl1, acl2, acl3)
+
+    Assert.assertEquals(acls, AclEntry.fromBytes(Json.encodeAsBytes(AclEntry.toJsonCompatibleMap(acls).asJava)))
+    Assert.assertEquals(acls, AclEntry.fromBytes(AclJson.getBytes(UTF_8)))
+  }
+}
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
new file mode 100644
index 0000000000000..73fe896875359
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/security/token/delegation/DelegationTokenManagerTest.scala
@@ -0,0 +1,363 @@
+/**
+  * Licensed to the Apache Software Foundation (ASF) under one or more
+  * contributor license agreements.  See the NOTICE file distributed with
+  * this work for additional information regarding copyright ownership.
+  * The ASF licenses this file to You under the Apache License, Version 2.0
+  * (the "License"); you may not use this file except in compliance with
+  * the License.  You may obtain a copy of the License at
+  *
+  *    http://www.apache.org/licenses/LICENSE-2.0
+  *
+  * Unless required by applicable law or agreed to in writing, software
+  * distributed under the License is distributed on an "AS IS" BASIS,
+  * WITHOUT WARRANTIES 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.token.delegation
+
+import java.net.InetAddress
+import java.nio.ByteBuffer
+import java.util.{Base64, Properties}
+
+import kafka.network.RequestChannel.Session
+import kafka.security.authorizer.{AclAuthorizer, AuthorizerUtils}
+import kafka.security.authorizer.AclEntry.WildcardHost
+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.jdk.CollectionConverters._
+import scala.collection.mutable.Buffer
+
+class DelegationTokenManagerTest extends ZooKeeperTestHarness  {
+
+  val time = new MockTime()
+  val owner = SecurityUtils.parseKafkaPrincipal("User:owner")
+  val renewer = List(SecurityUtils.parseKafkaPrincipal("User:renewer1"))
+  val tokenManagers = Buffer[DelegationTokenManager]()
+
+  val secretKey = "secretKey"
+  val maxLifeTimeMsDefault = Defaults.DelegationTokenMaxLifeTimeMsDefault
+  val renewTimeMsDefault = Defaults.DelegationTokenExpiryTimeMsDefault
+  var tokenCache: DelegationTokenCache = null
+  var props: Properties = null
+
+  var createTokenResult: CreateTokenResult = _
+  var error: Errors = Errors.NONE
+  var expiryTimeStamp: Long = 0
+
+  @Before
+  override def setUp(): Unit = {
+    super.setUp()
+    props = TestUtils.createBrokerConfig(0, zkConnect, enableToken = true)
+    props.put(KafkaConfig.SaslEnabledMechanismsProp, ScramMechanism.mechanismNames().asScala.mkString(","))
+    props.put(KafkaConfig.DelegationTokenSecretKeyProp, secretKey)
+    tokenCache = new DelegationTokenCache(ScramMechanism.mechanismNames())
+  }
+
+  @After
+  override def tearDown(): Unit = {
+    tokenManagers.foreach(_.shutdown())
+    super.tearDown()
+  }
+
+  @Test
+  def testTokenRequestsWithDelegationTokenDisabled(): Unit = {
+    val props: Properties = TestUtils.createBrokerConfig(0, zkConnect)
+    val config = KafkaConfig.fromProps(props)
+    val tokenManager = createDelegationTokenManager(config, tokenCache, time, zkClient)
+
+    tokenManager.createToken(owner, renewer, -1, createTokenResultCallBack)
+    assertEquals(Errors.DELEGATION_TOKEN_AUTH_DISABLED, createTokenResult.error)
+    assert(Array[Byte]() sameElements createTokenResult.hmac)
+
+    tokenManager.renewToken(owner, ByteBuffer.wrap("test".getBytes), 1000000, renewResponseCallback)
+    assertEquals(Errors.DELEGATION_TOKEN_AUTH_DISABLED, error)
+
+    tokenManager.expireToken(owner, ByteBuffer.wrap("test".getBytes), 1000000, renewResponseCallback)
+    assertEquals(Errors.DELEGATION_TOKEN_AUTH_DISABLED, error)
+  }
+
+  @Test
+  def testCreateToken(): Unit = {
+    val config = KafkaConfig.fromProps(props)
+    val tokenManager = createDelegationTokenManager(config, tokenCache, time, zkClient)
+    tokenManager.startup()
+
+    tokenManager.createToken(owner, renewer, -1 , createTokenResultCallBack)
+    val issueTime = time.milliseconds
+    val tokenId = createTokenResult.tokenId
+    val password = DelegationTokenManager.createHmac(tokenId, secretKey)
+    assertEquals(CreateTokenResult(issueTime, issueTime + renewTimeMsDefault,  issueTime + maxLifeTimeMsDefault, tokenId, password, Errors.NONE), createTokenResult)
+
+    val token = tokenManager.getToken(tokenId)
+    assertTrue(!token.isEmpty )
+    assertTrue(password sameElements token.get.hmac)
+  }
+
+  @Test
+  def testRenewToken(): Unit = {
+    val config = KafkaConfig.fromProps(props)
+    val tokenManager = createDelegationTokenManager(config, tokenCache, time, zkClient)
+    tokenManager.startup()
+
+    tokenManager.createToken(owner, renewer, -1 , createTokenResultCallBack)
+    val issueTime = time.milliseconds
+    val maxLifeTime = issueTime + maxLifeTimeMsDefault
+    val tokenId = createTokenResult.tokenId
+    val password = DelegationTokenManager.createHmac(tokenId, secretKey)
+    assertEquals(CreateTokenResult(issueTime, issueTime + renewTimeMsDefault,  maxLifeTime, tokenId, password, Errors.NONE), createTokenResult)
+
+    //try renewing non-existing token
+    tokenManager.renewToken(owner, ByteBuffer.wrap("test".getBytes), -1 , renewResponseCallback)
+    assertEquals(Errors.DELEGATION_TOKEN_NOT_FOUND, error)
+
+    //try renew non-owned tokens
+    val unknownOwner = SecurityUtils.parseKafkaPrincipal("User:Unknown")
+    tokenManager.renewToken(unknownOwner, ByteBuffer.wrap(password), -1 , renewResponseCallback)
+    assertEquals(Errors.DELEGATION_TOKEN_OWNER_MISMATCH, error)
+
+    // try renew with default time period
+    time.sleep(24 * 60 * 60 * 1000L)
+    var expectedExpiryStamp = time.milliseconds + renewTimeMsDefault
+    tokenManager.renewToken(owner, ByteBuffer.wrap(password), -1 , renewResponseCallback)
+    assertEquals(expectedExpiryStamp, expiryTimeStamp)
+    assertEquals(Errors.NONE, error)
+
+    // try renew with specific time period
+    time.sleep(24 * 60 * 60 * 1000L)
+    expectedExpiryStamp = time.milliseconds + 1 * 60 * 60 * 1000L
+    tokenManager.renewToken(owner, ByteBuffer.wrap(password), 1 * 60 * 60 * 1000L , renewResponseCallback)
+    assertEquals(expectedExpiryStamp, expiryTimeStamp)
+    assertEquals(Errors.NONE, error)
+
+    //try renewing more than max time period
+    time.sleep( 1 * 60 * 60 * 1000L)
+    tokenManager.renewToken(owner, ByteBuffer.wrap(password), 8 * 24 * 60 * 60 * 1000L, renewResponseCallback)
+    assertEquals(maxLifeTime, expiryTimeStamp)
+    assertEquals(Errors.NONE, error)
+
+    //try renewing expired token
+    time.sleep(8 * 24 * 60 * 60 * 1000L)
+    tokenManager.renewToken(owner, ByteBuffer.wrap(password), -1 , renewResponseCallback)
+    assertEquals(Errors.DELEGATION_TOKEN_EXPIRED, error)
+  }
+
+  @Test
+  def testExpireToken(): Unit = {
+    val config = KafkaConfig.fromProps(props)
+    val tokenManager = createDelegationTokenManager(config, tokenCache, time, zkClient)
+    tokenManager.startup()
+
+    tokenManager.createToken(owner, renewer, -1 , createTokenResultCallBack)
+    val issueTime = time.milliseconds
+    val tokenId = createTokenResult.tokenId
+    val password = DelegationTokenManager.createHmac(tokenId, secretKey)
+    assertEquals(CreateTokenResult(issueTime, issueTime + renewTimeMsDefault,  issueTime + maxLifeTimeMsDefault, tokenId, password, Errors.NONE), createTokenResult)
+
+    //try expire non-existing token
+    tokenManager.expireToken(owner, ByteBuffer.wrap("test".getBytes), -1 , renewResponseCallback)
+    assertEquals(Errors.DELEGATION_TOKEN_NOT_FOUND, error)
+
+    //try expire non-owned tokens
+    val unknownOwner = SecurityUtils.parseKafkaPrincipal("User:Unknown")
+    tokenManager.expireToken(unknownOwner, ByteBuffer.wrap(password), -1 , renewResponseCallback)
+    assertEquals(Errors.DELEGATION_TOKEN_OWNER_MISMATCH, error)
+
+    //try expire token at a timestamp
+    time.sleep(24 * 60 * 60 * 1000L)
+    val expectedExpiryStamp = time.milliseconds + 2 * 60 * 60 * 1000L
+    tokenManager.expireToken(owner, ByteBuffer.wrap(password), 2 * 60 * 60 * 1000L, renewResponseCallback)
+    assertEquals(expectedExpiryStamp, expiryTimeStamp)
+
+    //try expire token immediately
+    time.sleep(1 * 60 * 60 * 1000L)
+    tokenManager.expireToken(owner, ByteBuffer.wrap(password), -1, renewResponseCallback)
+    assert(tokenManager.getToken(tokenId).isEmpty)
+    assertEquals(Errors.NONE, error)
+    assertEquals(time.milliseconds, expiryTimeStamp)
+  }
+
+  @Test
+  def testRemoveTokenHmac():Unit = {
+    val config = KafkaConfig.fromProps(props)
+    val tokenManager = createDelegationTokenManager(config, tokenCache, time, zkClient)
+    tokenManager.startup()
+
+    tokenManager.createToken(owner, renewer, -1 , createTokenResultCallBack)
+    val issueTime = time.milliseconds
+    val tokenId = createTokenResult.tokenId
+    val password = DelegationTokenManager.createHmac(tokenId, secretKey)
+    assertEquals(CreateTokenResult(issueTime, issueTime + renewTimeMsDefault,  issueTime + maxLifeTimeMsDefault, tokenId, password, Errors.NONE), createTokenResult)
+
+    // expire the token immediately
+    tokenManager.expireToken(owner, ByteBuffer.wrap(password), -1, renewResponseCallback)
+
+    val encodedHmac = Base64.getEncoder.encodeToString(password)
+    // check respective hmac map entry is removed for the expired tokenId.
+    val tokenInformation = tokenManager.tokenCache.tokenIdForHmac(encodedHmac)
+    assertNull(tokenInformation)
+
+    //check that the token is removed
+    assert(tokenManager.getToken(tokenId).isEmpty)
+  }
+
+  @Test
+  def testDescribeToken(): Unit = {
+
+    val config = KafkaConfig.fromProps(props)
+
+    val owner1 = SecurityUtils.parseKafkaPrincipal("User:owner1")
+    val owner2 = SecurityUtils.parseKafkaPrincipal("User:owner2")
+    val owner3 = SecurityUtils.parseKafkaPrincipal("User:owner3")
+    val owner4 = SecurityUtils.parseKafkaPrincipal("User:owner4")
+
+    val renewer1 = SecurityUtils.parseKafkaPrincipal("User:renewer1")
+    val renewer2 = SecurityUtils.parseKafkaPrincipal("User:renewer2")
+    val renewer3 = SecurityUtils.parseKafkaPrincipal("User:renewer3")
+    val renewer4 = SecurityUtils.parseKafkaPrincipal("User:renewer4")
+
+    val aclAuthorizer = new AclAuthorizer
+    aclAuthorizer.configure(config.originals)
+
+    var hostSession = new Session(owner1, InetAddress.getByName("192.168.1.1"))
+
+    val tokenManager = createDelegationTokenManager(config, tokenCache, time, zkClient)
+    tokenManager.startup()
+
+    //create tokens
+    tokenManager.createToken(owner1, List(renewer1, renewer2), 1 * 60 * 60 * 1000L, createTokenResultCallBack)
+
+    tokenManager.createToken(owner2, List(renewer3), 1 * 60 * 60 * 1000L, createTokenResultCallBack)
+    val tokenId2 = createTokenResult.tokenId
+
+    tokenManager.createToken(owner3, List(renewer4), 2 * 60 * 60 * 1000L, createTokenResultCallBack)
+    val tokenId3 = createTokenResult.tokenId
+
+    tokenManager.createToken(owner4, List(owner1, renewer4), 2 * 60 * 60 * 1000L, createTokenResultCallBack)
+
+    assert(tokenManager.getAllTokenInformation.size == 4 )
+
+    //get tokens non-exiting owner
+    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, aclAuthorizer, hostSession, owner1, List())
+    assert(tokens.size == 0)
+
+    //get all tokens for owner1
+    tokens = getTokens(tokenManager, aclAuthorizer, hostSession, owner1, List(owner1))
+    assert(tokens.size == 2)
+
+    //get all tokens for owner1
+    tokens = getTokens(tokenManager, aclAuthorizer, hostSession, owner1, null)
+    assert(tokens.size == 2)
+
+    //get all tokens for unknown owner
+    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, 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.ifPresent { e => throw e }
+    }
+
+    //get all tokens for multiple owners (owner1, renewer4) and with permission
+    createAcl(new AclBinding(new ResourcePattern(DELEGATION_TOKEN, tokenId3, LITERAL),
+      new AccessControlEntry(owner1.toString, WildcardHost, DESCRIBE, ALLOW)))
+    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, 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, 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 = 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)
+
+    aclAuthorizer.close()
+  }
+
+  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) = {
+        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)
+    }
+  }
+
+  @Test
+  def testPeriodicTokenExpiry(): Unit = {
+    val config = KafkaConfig.fromProps(props)
+    val tokenManager = createDelegationTokenManager(config, tokenCache, time, zkClient)
+    tokenManager.startup()
+
+    //create tokens
+    tokenManager.createToken(owner, renewer, 1 * 60 * 60 * 1000L, createTokenResultCallBack)
+    tokenManager.createToken(owner, renewer, 1 * 60 * 60 * 1000L, createTokenResultCallBack)
+    tokenManager.createToken(owner, renewer, 2 * 60 * 60 * 1000L, createTokenResultCallBack)
+    tokenManager.createToken(owner, renewer, 2 * 60 * 60 * 1000L, createTokenResultCallBack)
+    assert(tokenManager.getAllTokenInformation.size == 4 )
+
+    time.sleep(2 * 60 * 60 * 1000L)
+    tokenManager.expireTokens()
+    assert(tokenManager.getAllTokenInformation.size == 2 )
+
+  }
+
+  private def createTokenResultCallBack(ret: CreateTokenResult): Unit = {
+    createTokenResult = ret
+  }
+
+  private def renewResponseCallback(ret: Errors, timeStamp: Long): Unit = {
+    error = ret
+    expiryTimeStamp = timeStamp
+  }
+
+  private def createDelegationTokenManager(config: KafkaConfig, tokenCache: DelegationTokenCache,
+                                           time: Time, zkClient: KafkaZkClient): DelegationTokenManager = {
+    val tokenManager = new DelegationTokenManager(config, tokenCache, time, zkClient)
+    tokenManagers += tokenManager
+    tokenManager
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/server/AbstractApiVersionsRequestTest.scala b/core/src/test/scala/unit/kafka/server/AbstractApiVersionsRequestTest.scala
new file mode 100644
index 0000000000000..48d7e214e5344
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/AbstractApiVersionsRequestTest.scala
@@ -0,0 +1,49 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package kafka.server
+
+import org.apache.kafka.common.message.ApiVersionsResponseData.ApiVersionsResponseKey
+import org.apache.kafka.common.protocol.ApiKeys
+import org.apache.kafka.common.requests.{ApiVersionsRequest, ApiVersionsResponse}
+import org.junit.Assert._
+
+import scala.jdk.CollectionConverters._
+
+abstract class AbstractApiVersionsRequestTest extends BaseRequestTest {
+
+  def sendUnsupportedApiVersionRequest(request: ApiVersionsRequest): ApiVersionsResponse = {
+    val overrideHeader = nextRequestHeader(ApiKeys.API_VERSIONS, Short.MaxValue)
+    val socket = connect(anySocketServer)
+    try {
+      sendWithHeader(request, overrideHeader, socket)
+      receive[ApiVersionsResponse](socket, ApiKeys.API_VERSIONS, 0.toShort)
+    } finally socket.close()
+  }
+
+  def validateApiVersionsResponse(apiVersionsResponse: ApiVersionsResponse): Unit = {
+    val enabledPublicApis = ApiKeys.enabledApis()
+    assertEquals("API keys in ApiVersionsResponse must match API keys supported by broker.",
+      enabledPublicApis.size(), 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)
+      assertEquals(s"Received unexpected min version for API key ${actualApiVersion.apiKey}.", expectedApiVersion.minVersion, actualApiVersion.minVersion)
+      assertEquals(s"Received unexpected max version for API key ${actualApiVersion.apiKey}.", expectedApiVersion.maxVersion, actualApiVersion.maxVersion)
+    }
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/server/AbstractCreateTopicsRequestTest.scala b/core/src/test/scala/unit/kafka/server/AbstractCreateTopicsRequestTest.scala
index d89a9df082b1f..10f73a88ae40d 100644
--- a/core/src/test/scala/unit/kafka/server/AbstractCreateTopicsRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/AbstractCreateTopicsRequestTest.scala
@@ -17,46 +17,104 @@
 
 package kafka.server
 
+import java.util
 import java.util.Properties
 
 import kafka.network.SocketServer
 import kafka.utils.TestUtils
-import org.apache.kafka.common.protocol.types.Struct
-import org.apache.kafka.common.protocol.{ApiKeys, Errors}
-import org.apache.kafka.common.requests.{ApiError, CreateTopicsRequest, CreateTopicsResponse, MetadataRequest, MetadataResponse}
+import org.apache.kafka.common.message.CreateTopicsRequestData
+import org.apache.kafka.common.message.CreateTopicsRequestData._
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.requests._
 import org.junit.Assert.{assertEquals, assertFalse, assertNotNull, assertTrue}
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 
-class AbstractCreateTopicsRequestTest extends BaseRequestTest {
+abstract 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],
+                timeout: Integer = 10000,
+                validateOnly: Boolean = false) = {
+    val req = new CreateTopicsRequestData()
+    req.setTimeoutMs(timeout)
+    req.setTopics(new CreatableTopicCollection(topics.asJava.iterator()))
+    req.setValidateOnly(validateOnly)
+    new CreateTopicsRequest.Builder(req).build()
+  }
+
+  def topicReq(name: String,
+               numPartitions: Integer = null,
+               replicationFactor: Integer = null,
+               config: Map[String, String] = null,
+               assignment: Map[Int, Seq[Int]] = null): CreatableTopic = {
+    val topic = new CreatableTopic()
+    topic.setName(name)
+    if (numPartitions != null) {
+      topic.setNumPartitions(numPartitions)
+    } else if (assignment != null) {
+      topic.setNumPartitions(-1)
+    } else {
+      topic.setNumPartitions(1)
+    }
+    if (replicationFactor != null) {
+      topic.setReplicationFactor(replicationFactor.toShort)
+    } else if (assignment != null) {
+      topic.setReplicationFactor((-1).toShort)
+    } else {
+      topic.setReplicationFactor(1.toShort)
+    }
+    if (config != null) {
+      val effectiveConfigs = new CreateableTopicConfigCollection()
+      config.foreach {
+        case (name, value) =>
+          effectiveConfigs.add(new CreateableTopicConfig().setName(name).setValue(value))
+      }
+      topic.setConfigs(effectiveConfigs)
+    }
+    if (assignment != null) {
+      val effectiveAssignments = new CreatableReplicaAssignmentCollection()
+      assignment.foreach {
+        case (partitionIndex, brokerIdList) => {
+          val effectiveAssignment = new CreatableReplicaAssignment()
+          effectiveAssignment.setPartitionIndex(partitionIndex)
+          val brokerIds = new util.ArrayList[java.lang.Integer]()
+          brokerIdList.foreach(brokerId => brokerIds.add(brokerId))
+          effectiveAssignment.setBrokerIds(brokerIds)
+          effectiveAssignments.add(effectiveAssignment)
+        }
+      }
+      topic.setAssignments(effectiveAssignments)
+    }
+    topic
+  }
+
   protected def validateValidCreateTopicsRequests(request: CreateTopicsRequest): Unit = {
     val response = sendCreateTopicRequest(request)
 
-    val error = response.errors.values.asScala.find(_.isFailure)
-    assertTrue(s"There should be no errors, found ${response.errors.asScala}", error.isEmpty)
-
-    request.topics.asScala.foreach { case (topic, details) =>
+    assertTrue(s"There should be no errors, found " +
+      s"${response.errorCounts().keySet().asScala.mkString(", ")},",
+      response.errorCounts().keySet().asScala.find(_.code() > 0).isEmpty)
 
+    request.data.topics.forEach { topic =>
       def verifyMetadata(socketServer: SocketServer) = {
         val metadata = sendMetadataRequest(
-          new MetadataRequest.Builder(List(topic).asJava, true).build()).topicMetadata.asScala
-        val metadataForTopic = metadata.filter(_.topic == topic).head
+          new MetadataRequest.Builder(List(topic.name()).asJava, true).build()).topicMetadata.asScala
+        val metadataForTopic = metadata.filter(_.topic == topic.name()).head
 
-        val partitions = if (!details.replicasAssignments.isEmpty)
-          details.replicasAssignments.size
+        val partitions = if (!topic.assignments().isEmpty)
+          topic.assignments().size
         else
-          details.numPartitions
+          topic.numPartitions
 
-        val replication = if (!details.replicasAssignments.isEmpty)
-          details.replicasAssignments.asScala.head._2.size
+        val replication = if (!topic.assignments().isEmpty)
+          topic.assignments().iterator().next().brokerIds().size()
         else
-          details.replicationFactor
+          topic.replicationFactor
 
-        if (request.validateOnly) {
+        if (request.data.validateOnly) {
           assertNotNull(s"Topic $topic should be created", metadataForTopic)
           assertFalse(s"Error ${metadataForTopic.error} for topic $topic", metadataForTopic.error == Errors.NONE)
           assertTrue("The topic should have no partitions", metadataForTopic.partitionMetadata.isEmpty)
@@ -64,16 +122,26 @@ 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.replicaIds.size)
+          } else {
+            assertEquals("The topic should have the correct replication factor", replication, metadataForTopic.partitionMetadata.asScala.head.replicaIds.size)
+          }
         }
       }
 
       // Verify controller broker has the correct metadata
       verifyMetadata(controllerSocketServer)
-      if (!request.validateOnly) {
+      if (!request.data.validateOnly) {
         // Wait until metadata is propagated and validate non-controller broker has the correct metadata
-        TestUtils.waitUntilMetadataIsPropagated(servers, topic, 0)
+        TestUtils.waitUntilMetadataIsPropagated(servers, topic.name(), 0)
       }
       verifyMetadata(notControllerSocketServer)
     }
@@ -82,44 +150,25 @@ class AbstractCreateTopicsRequestTest extends BaseRequestTest {
   protected def error(error: Errors, errorMessage: Option[String] = None): ApiError =
     new ApiError(error, errorMessage.orNull)
 
-  protected def toStructWithDuplicateFirstTopic(request: CreateTopicsRequest): Struct = {
-    val struct = request.toStruct
-    val topics = struct.getArray("create_topic_requests")
-    val firstTopic = topics(0).asInstanceOf[Struct]
-    val newTopics = firstTopic :: topics.toList
-    struct.set("create_topic_requests", newTopics.toArray)
-    struct
-  }
-
-  protected def addPartitionsAndReplicationFactorToFirstTopic(request: CreateTopicsRequest) = {
-    val struct = request.toStruct
-    val topics = struct.getArray("create_topic_requests")
-    val firstTopic = topics(0).asInstanceOf[Struct]
-    firstTopic.set("num_partitions", 1)
-    firstTopic.set("replication_factor", 1.toShort)
-    new CreateTopicsRequest(struct, request.version)
-  }
-
   protected def validateErrorCreateTopicsRequests(request: CreateTopicsRequest,
                                                   expectedResponse: Map[String, ApiError],
-                                                  checkErrorMessage: Boolean = true,
-                                                  requestStruct: Option[Struct] = None): Unit = {
-    val response = requestStruct.map(sendCreateTopicRequestStruct(_, request.version)).getOrElse(
-      sendCreateTopicRequest(request))
-    val errors = response.errors.asScala
-    assertEquals("The response size should match", expectedResponse.size, response.errors.size)
-
-    expectedResponse.foreach { case (topic, expectedError) =>
-      val expected = expectedResponse(topic)
-      val actual = errors(topic)
-      assertEquals("The response error should match", expected.error, actual.error)
+                                                  checkErrorMessage: Boolean = true): Unit = {
+    val response = sendCreateTopicRequest(request)
+    assertEquals("The response size should match", expectedResponse.size, response.data().topics().size)
+
+    expectedResponse.foreach { case (topicName, expectedError) =>
+      val expected = expectedResponse(topicName)
+      val actual = response.data().topics().find(topicName)
+      if (actual == null) {
+        throw new RuntimeException(s"No response data found for topic $topicName")
+      }
+      assertEquals("The response error should match", expected.error.code(), actual.errorCode())
       if (checkErrorMessage) {
-        assertEquals(expected.message, actual.message)
-        assertEquals(expected.messageWithFallback, actual.messageWithFallback)
+        assertEquals(expected.message, actual.errorMessage())
       }
       // If no error validate topic exists
-      if (expectedError.isSuccess && !request.validateOnly) {
-        validateTopicExists(topic)
+      if (expectedError.isSuccess && !request.data.validateOnly) {
+        validateTopicExists(topicName)
       }
     }
   }
@@ -131,24 +180,13 @@ class AbstractCreateTopicsRequestTest extends BaseRequestTest {
     assertTrue("The topic should be created", metadata.exists(p => p.topic.equals(topic) && p.error == Errors.NONE))
   }
 
-  protected def replicaAssignmentToJava(assignments: Map[Int, List[Int]]) = {
-    assignments.map { case (k, v) => (k: Integer, v.map { i => i: Integer }.asJava) }.asJava
-  }
-
-  protected def sendCreateTopicRequestStruct(requestStruct: Struct, apiVersion: Short,
-                                             socketServer: SocketServer = controllerSocketServer): CreateTopicsResponse = {
-    val response = connectAndSendStruct(requestStruct, ApiKeys.CREATE_TOPICS, apiVersion, socketServer)
-    CreateTopicsResponse.parse(response, apiVersion)
-  }
-
-  protected def sendCreateTopicRequest(request: CreateTopicsRequest, socketServer: SocketServer = controllerSocketServer): CreateTopicsResponse = {
-    val response = connectAndSend(request, ApiKeys.CREATE_TOPICS, socketServer)
-    CreateTopicsResponse.parse(response, request.version)
+  protected def sendCreateTopicRequest(request: CreateTopicsRequest,
+                                       socketServer: SocketServer = controllerSocketServer): CreateTopicsResponse = {
+    connectAndReceive[CreateTopicsResponse](request, socketServer)
   }
 
-  protected def sendMetadataRequest(request: MetadataRequest, destination: SocketServer = anySocketServer): MetadataResponse = {
-    val response = connectAndSend(request, ApiKeys.METADATA, destination = destination)
-    MetadataResponse.parse(response, ApiKeys.METADATA.latestVersion)
+  protected def sendMetadataRequest(request: MetadataRequest): MetadataResponse = {
+    connectAndReceive[MetadataResponse](request)
   }
 
 }
diff --git a/core/src/test/scala/unit/kafka/server/AbstractFetcherManagerTest.scala b/core/src/test/scala/unit/kafka/server/AbstractFetcherManagerTest.scala
new file mode 100644
index 0000000000000..847449b02e3f2
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/AbstractFetcherManagerTest.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.server
+
+import com.yammer.metrics.core.Gauge
+import kafka.cluster.BrokerEndPoint
+import kafka.metrics.KafkaYammerMetrics
+import kafka.utils.TestUtils
+import org.apache.kafka.common.TopicPartition
+import org.easymock.EasyMock
+import org.junit.{Before, Test}
+import org.junit.Assert._
+
+import scala.jdk.CollectionConverters._
+
+class AbstractFetcherManagerTest {
+
+  @Before
+  def cleanMetricRegistry(): Unit = {
+    TestUtils.clearYammerMetrics()
+  }
+
+  private def getMetricValue(name: String): Any = {
+    KafkaYammerMetrics.defaultRegistry.allMetrics.asScala.filter { case (k, _) => k.getName == name }.values.headOption.get.
+      asInstanceOf[Gauge[Int]].value()
+  }
+
+  @Test
+  def testAddAndRemovePartition(): 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 -> initialFetchState)))
+        .andReturn(Set(tp))
+    EasyMock.expect(fetcher.fetchState(tp))
+      .andReturn(Some(PartitionFetchState(fetchOffset, None, leaderEpoch, Truncating, lastFetchedEpoch = None)))
+    EasyMock.expect(fetcher.removePartitions(Set(tp))).andReturn(Map.empty)
+    EasyMock.expect(fetcher.fetchState(tp)).andReturn(None)
+    EasyMock.replay(fetcher)
+
+    fetcherManager.addFetcherForPartitions(Map(tp -> initialFetchState))
+    assertEquals(Some(fetcher), fetcherManager.getFetcher(tp))
+
+    fetcherManager.removeFetcherForPartitions(Set(tp))
+    assertEquals(None, fetcherManager.getFetcher(tp))
+
+    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 -> initialFetchState)))
+        .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 b95f66cba55ca..29ddebc362810 100644
--- a/core/src/test/scala/unit/kafka/server/AbstractFetcherThreadTest.scala
+++ b/core/src/test/scala/unit/kafka/server/AbstractFetcherThreadTest.scala
@@ -17,206 +17,1142 @@
 
 package kafka.server
 
-import AbstractFetcherThread._
-import com.yammer.metrics.Metrics
+import java.nio.ByteBuffer
+import java.util.Optional
+import java.util.concurrent.atomic.AtomicInteger
+
 import kafka.cluster.BrokerEndPoint
-import kafka.server.AbstractFetcherThread.{FetchRequest, PartitionData}
+import kafka.log.LogAppendInfo
+import kafka.message.NoCompressionCodec
+import kafka.metrics.KafkaYammerMetrics
+import kafka.server.AbstractFetcherThread.ReplicaFetch
+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.protocol.Errors
-import org.apache.kafka.common.record.{CompressionType, MemoryRecords, SimpleRecord}
-import org.apache.kafka.common.requests.EpochEndOffset
-import org.junit.Assert.{assertFalse, assertTrue}
+import org.apache.kafka.common.errors.{FencedLeaderEpochException, UnknownLeaderEpochException}
+import org.apache.kafka.common.message.FetchResponseData
+import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData.EpochEndOffset
+import org.apache.kafka.common.protocol.{ApiKeys, Errors}
+import org.apache.kafka.common.record._
+import org.apache.kafka.common.requests.OffsetsForLeaderEpochResponse.{UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET}
+import org.apache.kafka.common.requests.FetchRequest
+import org.apache.kafka.common.utils.Time
+import org.junit.Assert._
 import org.junit.{Before, Test}
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 import scala.collection.{Map, Set, mutable}
+import scala.util.Random
+import org.scalatest.Assertions.assertThrows
+
+import scala.collection.mutable.ArrayBuffer
+import scala.compat.java8.OptionConverters._
 
 class AbstractFetcherThreadTest {
 
+  val truncateOnFetch = true
+  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] = KafkaYammerMetrics.defaultRegistry().allMetrics().asScala.keySet.map(_.getName)
+
+  private def mkBatch(baseOffset: Long, leaderEpoch: Int, records: SimpleRecord*): RecordBatch = {
+    MemoryRecords.withRecords(baseOffset, CompressionType.NONE, leaderEpoch, records: _*)
+      .batches.asScala.head
+  }
+
+  private def initialFetchState(fetchOffset: Long, leaderEpoch: Int): InitialFetchState = {
+    InitialFetchState(leader = new BrokerEndPoint(0, "localhost", 9092),
+      initOffset = fetchOffset, currentLeaderEpoch = leaderEpoch)
   }
 
   @Test
-  def testMetricsRemovedOnShutdown() {
+  def testMetricsRemovedOnShutdown(): Unit = {
     val partition = new TopicPartition("topic", 0)
-    val fetcherThread = new DummyFetcherThread("dummy", "client", new BrokerEndPoint(0, "localhost", 9092))
-
-    fetcherThread.start()
+    val fetcher = new MockFetcherThread
 
     // add one partition to create the consumer lag metric
-    fetcherThread.addPartitions(Map(partition -> 0L))
+    fetcher.setReplicaState(partition, MockFetcherThread.PartitionState(leaderEpoch = 0))
+    fetcher.addPartitions(Map(partition -> initialFetchState(0L, leaderEpoch = 0)))
+    fetcher.setLeaderState(partition, MockFetcherThread.PartitionState(leaderEpoch = 0))
+
+    fetcher.start()
+
+    val brokerTopicStatsMetrics = fetcher.brokerTopicStats.allTopicsStats.metricMap.keySet
+    val fetcherMetrics = Set(FetcherMetrics.BytesPerSec, FetcherMetrics.RequestsPerSec, FetcherMetrics.ConsumerLag)
 
     // wait until all fetcher metrics are present
-    TestUtils.waitUntilTrue(() =>
-      allMetricsNames == Set(FetcherMetrics.BytesPerSec, FetcherMetrics.RequestsPerSec, FetcherMetrics.ConsumerLag),
+    TestUtils.waitUntilTrue(() => allMetricsNames == brokerTopicStatsMetrics ++ fetcherMetrics,
       "Failed waiting for all fetcher metrics to be registered")
 
-    fetcherThread.shutdown()
+    fetcher.shutdown()
 
-    // after shutdown, they should be gone
-    assertTrue(Metrics.defaultRegistry().allMetrics().isEmpty)
+    // verify that all the fetcher metrics are removed and only brokerTopicStats left
+    val metricNames = KafkaYammerMetrics.defaultRegistry().allMetrics().asScala.keySet.map(_.getName).toSet
+    assertTrue(metricNames.intersect(fetcherMetrics).isEmpty)
+    assertEquals(brokerTopicStatsMetrics, metricNames.intersect(brokerTopicStatsMetrics))
   }
 
   @Test
-  def testConsumerLagRemovedWithPartition() {
+  def testConsumerLagRemovedWithPartition(): Unit = {
     val partition = new TopicPartition("topic", 0)
-    val fetcherThread = new DummyFetcherThread("dummy", "client", new BrokerEndPoint(0, "localhost", 9092))
-
-    fetcherThread.start()
+    val fetcher = new MockFetcherThread
 
     // add one partition to create the consumer lag metric
-    fetcherThread.addPartitions(Map(partition -> 0L))
+    fetcher.setReplicaState(partition, MockFetcherThread.PartitionState(leaderEpoch = 0))
+    fetcher.addPartitions(Map(partition -> initialFetchState(0L, leaderEpoch = 0)))
+    fetcher.setLeaderState(partition, MockFetcherThread.PartitionState(leaderEpoch = 0))
 
-    // wait until lag metric is present
-    TestUtils.waitUntilTrue(() => allMetricsNames(FetcherMetrics.ConsumerLag),
-      "Failed waiting for consumer lag metric")
+    fetcher.doWork()
+
+    assertTrue("Failed waiting for consumer lag metric",
+      allMetricsNames(FetcherMetrics.ConsumerLag))
 
     // remove the partition to simulate leader migration
-    fetcherThread.removePartitions(Set(partition))
+    fetcher.removePartitions(Set(partition))
 
     // the lag metric should now be gone
     assertFalse(allMetricsNames(FetcherMetrics.ConsumerLag))
+  }
+
+  @Test
+  def testSimpleFetch(): Unit = {
+    val partition = new TopicPartition("topic", 0)
+    val fetcher = new MockFetcherThread
+
+    fetcher.setReplicaState(partition, MockFetcherThread.PartitionState(leaderEpoch = 0))
+    fetcher.addPartitions(Map(partition -> initialFetchState(0L, leaderEpoch = 0)))
 
-    fetcherThread.shutdown()
+    val batch = mkBatch(baseOffset = 0L, leaderEpoch = 0,
+      new SimpleRecord("a".getBytes), new SimpleRecord("b".getBytes))
+    val leaderState = MockFetcherThread.PartitionState(Seq(batch), leaderEpoch = 0, highWatermark = 2L)
+    fetcher.setLeaderState(partition, leaderState)
+
+    fetcher.doWork()
+
+    val replicaState = fetcher.replicaPartitionState(partition)
+    assertEquals(2L, replicaState.logEndOffset)
+    assertEquals(2L, replicaState.highWatermark)
   }
 
-  private def allMetricsNames = Metrics.defaultRegistry().allMetrics().asScala.keySet.map(_.getName)
+  @Test
+  def testFencedTruncation(): Unit = {
+    val partition = new TopicPartition("topic", 0)
+    val fetcher = new MockFetcherThread
+
+    fetcher.setReplicaState(partition, MockFetcherThread.PartitionState(leaderEpoch = 0))
+    fetcher.addPartitions(Map(partition -> initialFetchState(0L, leaderEpoch = 0)))
 
-  class DummyFetchRequest(val offsets: collection.Map[TopicPartition, Long]) extends FetchRequest {
-    override def isEmpty: Boolean = offsets.isEmpty
+    val batch = mkBatch(baseOffset = 0L, leaderEpoch = 1,
+      new SimpleRecord("a".getBytes), new SimpleRecord("b".getBytes))
+    val leaderState = MockFetcherThread.PartitionState(Seq(batch), leaderEpoch = 1, highWatermark = 2L)
+    fetcher.setLeaderState(partition, leaderState)
 
-    override def offset(topicPartition: TopicPartition): Long = offsets(topicPartition)
+    fetcher.doWork()
+
+    // No progress should be made
+    val replicaState = fetcher.replicaPartitionState(partition)
+    assertEquals(0L, replicaState.logEndOffset)
+    assertEquals(0L, replicaState.highWatermark)
+
+    // After fencing, the fetcher should remove the partition from tracking and mark as failed
+    assertTrue(fetcher.fetchState(partition).isEmpty)
+    assertTrue(failedPartitions.contains(partition))
   }
 
-  class TestPartitionData(records: MemoryRecords = MemoryRecords.EMPTY) extends PartitionData {
-    override def error: Errors = Errors.NONE
+  @Test
+  def testFencedFetch(): Unit = {
+    val partition = new TopicPartition("topic", 0)
+    val fetcher = new MockFetcherThread
+
+    val replicaState = MockFetcherThread.PartitionState(leaderEpoch = 0)
+    fetcher.setReplicaState(partition, replicaState)
+    fetcher.addPartitions(Map(partition -> initialFetchState(0L, leaderEpoch = 0)))
+
+    val batch = mkBatch(baseOffset = 0L, leaderEpoch = 0,
+      new SimpleRecord("a".getBytes),
+      new SimpleRecord("b".getBytes))
+    val leaderState = MockFetcherThread.PartitionState(Seq(batch), leaderEpoch = 0, highWatermark = 2L)
+    fetcher.setLeaderState(partition, leaderState)
+
+    fetcher.doWork()
 
-    override def toRecords: MemoryRecords = records
+    // Verify we have caught up
+    assertEquals(2, replicaState.logEndOffset)
 
-    override def highWatermark: Long = 0L
+    // Bump the epoch on the leader
+    fetcher.leaderPartitionState(partition).leaderEpoch += 1
 
-    override def exception: Option[Throwable] = None
+    fetcher.doWork()
+
+    // After fencing, the fetcher should remove the partition from tracking and mark as failed
+    assertTrue(fetcher.fetchState(partition).isEmpty)
+    assertTrue(failedPartitions.contains(partition))
   }
 
-  class DummyFetcherThread(name: String,
-                           clientId: String,
-                           sourceBroker: BrokerEndPoint,
-                           fetchBackOffMs: Int = 0)
-    extends AbstractFetcherThread(name, clientId, sourceBroker, fetchBackOffMs, isInterruptible = true, includeLogTruncation = false) {
+  @Test
+  def testUnknownLeaderEpochInTruncation(): Unit = {
+    val partition = new TopicPartition("topic", 0)
+    val fetcher = new MockFetcherThread
 
-    type REQ = DummyFetchRequest
-    type PD = PartitionData
+    // The replica's leader epoch is ahead of the leader
+    val replicaState = MockFetcherThread.PartitionState(leaderEpoch = 1)
+    fetcher.setReplicaState(partition, replicaState)
+    fetcher.addPartitions(Map(partition -> initialFetchState(0L, leaderEpoch = 1)), forceTruncation = true)
 
-    override def processPartitionData(topicPartition: TopicPartition,
-                                      fetchOffset: Long,
-                                      partitionData: PartitionData): Unit = {}
+    val batch = mkBatch(baseOffset = 0L, leaderEpoch = 0, new SimpleRecord("a".getBytes))
+    val leaderState = MockFetcherThread.PartitionState(Seq(batch), leaderEpoch = 0, highWatermark = 2L)
+    fetcher.setLeaderState(partition, leaderState)
+
+    fetcher.doWork()
+
+    // Not data has been fetched and the follower is still truncating
+    assertEquals(0, replicaState.logEndOffset)
+    assertEquals(Some(Truncating), fetcher.fetchState(partition).map(_.state))
+
+    // Bump the epoch on the leader
+    fetcher.leaderPartitionState(partition).leaderEpoch += 1
+
+    // Now we can make progress
+    fetcher.doWork()
+
+    assertEquals(1, replicaState.logEndOffset)
+    assertEquals(Some(Fetching), fetcher.fetchState(partition).map(_.state))
+  }
+
+  @Test
+  def testUnknownLeaderEpochWhileFetching(): Unit = {
+    val partition = new TopicPartition("topic", 0)
+    val fetcher = new MockFetcherThread
+
+    // This test is contrived because it shouldn't be possible to to see unknown leader epoch
+    // in the Fetching state as the leader must validate the follower's epoch when it checks
+    // the truncation offset.
+
+    val replicaState = MockFetcherThread.PartitionState(leaderEpoch = 1)
+    fetcher.setReplicaState(partition, replicaState)
+    fetcher.addPartitions(Map(partition -> initialFetchState(0L, leaderEpoch = 1)))
+
+    val leaderState = MockFetcherThread.PartitionState(Seq(
+      mkBatch(baseOffset = 0L, leaderEpoch = 0, new SimpleRecord("a".getBytes)),
+      mkBatch(baseOffset = 1L, leaderEpoch = 0, new SimpleRecord("b".getBytes))
+    ), leaderEpoch = 1, highWatermark = 2L)
+    fetcher.setLeaderState(partition, leaderState)
+
+    fetcher.doWork()
+
+    // We have fetched one batch and gotten out of the truncation phase
+    assertEquals(1, replicaState.logEndOffset)
+    assertEquals(Some(Fetching), fetcher.fetchState(partition).map(_.state))
+
+    // Somehow the leader epoch rewinds
+    fetcher.leaderPartitionState(partition).leaderEpoch = 0
+
+    // We are stuck at the current offset
+    fetcher.doWork()
+    assertEquals(1, replicaState.logEndOffset)
+    assertEquals(Some(Fetching), fetcher.fetchState(partition).map(_.state))
+
+    // After returning to the right epoch, we can continue fetching
+    fetcher.leaderPartitionState(partition).leaderEpoch = 1
+    fetcher.doWork()
+    assertEquals(2, replicaState.logEndOffset)
+    assertEquals(Some(Fetching), fetcher.fetchState(partition).map(_.state))
+  }
+
+  @Test
+  def testTruncation(): Unit = {
+    val partition = new TopicPartition("topic", 0)
+    val fetcher = new MockFetcherThread
+
+    val replicaLog = Seq(
+      mkBatch(baseOffset = 0, leaderEpoch = 0, new SimpleRecord("a".getBytes)),
+      mkBatch(baseOffset = 1, leaderEpoch = 2, new SimpleRecord("b".getBytes)),
+      mkBatch(baseOffset = 2, leaderEpoch = 4, new SimpleRecord("c".getBytes)))
+
+    val replicaState = MockFetcherThread.PartitionState(replicaLog, leaderEpoch = 5, highWatermark = 0L)
+    fetcher.setReplicaState(partition, replicaState)
+    fetcher.addPartitions(Map(partition -> initialFetchState(3L, leaderEpoch = 5)))
+
+    val leaderLog = Seq(
+      mkBatch(baseOffset = 0, leaderEpoch = 1, new SimpleRecord("a".getBytes)),
+      mkBatch(baseOffset = 1, leaderEpoch = 3, new SimpleRecord("b".getBytes)),
+      mkBatch(baseOffset = 2, leaderEpoch = 5, new SimpleRecord("c".getBytes)))
+
+    val leaderState = MockFetcherThread.PartitionState(leaderLog, leaderEpoch = 5, highWatermark = 2L)
+    fetcher.setLeaderState(partition, leaderState)
+
+    TestUtils.waitUntilTrue(() => {
+      fetcher.doWork()
+      fetcher.replicaPartitionState(partition).log == fetcher.leaderPartitionState(partition).log
+    }, "Failed to reconcile leader and follower logs")
+
+    assertEquals(leaderState.logStartOffset, replicaState.logStartOffset)
+    assertEquals(leaderState.logEndOffset, replicaState.logEndOffset)
+    assertEquals(leaderState.highWatermark, replicaState.highWatermark)
+  }
+
+  @Test
+  def testTruncateToHighWatermarkIfLeaderEpochRequestNotSupported(): Unit = {
+    val highWatermark = 2L
+    val partition = new TopicPartition("topic", 0)
+    val fetcher = new MockFetcherThread {
+      override def truncate(topicPartition: TopicPartition, truncationState: OffsetTruncationState): Unit = {
+        assertEquals(highWatermark, truncationState.offset)
+        assertTrue(truncationState.truncationCompleted)
+        super.truncate(topicPartition, truncationState)
+      }
+
+      override def fetchEpochEndOffsets(partitions: Map[TopicPartition, EpochData]): Map[TopicPartition, EpochEndOffset] =
+        throw new UnsupportedOperationException
+
+      override protected val isOffsetForLeaderEpochSupported: Boolean = false
+
+      override protected val isTruncationOnFetchSupported: Boolean = false
+    }
+
+    val replicaLog = Seq(
+      mkBatch(baseOffset = 0, leaderEpoch = 0, new SimpleRecord("a".getBytes)),
+      mkBatch(baseOffset = 1, leaderEpoch = 2, new SimpleRecord("b".getBytes)),
+      mkBatch(baseOffset = 2, leaderEpoch = 4, new SimpleRecord("c".getBytes)))
+
+    val replicaState = MockFetcherThread.PartitionState(replicaLog, leaderEpoch = 5, highWatermark)
+    fetcher.setReplicaState(partition, replicaState)
+    fetcher.addPartitions(Map(partition -> initialFetchState(highWatermark, leaderEpoch = 5)))
+
+    fetcher.doWork()
+
+    assertEquals(highWatermark, replicaState.logEndOffset)
+    assertEquals(highWatermark, fetcher.fetchState(partition).get.fetchOffset)
+    assertTrue(fetcher.fetchState(partition).get.isReadyForFetch)
+  }
+
+  @Test
+  def testTruncateToHighWatermarkIfLeaderEpochInfoNotAvailable(): Unit = {
+    val highWatermark = 2L
+    val partition = new TopicPartition("topic", 0)
+    val fetcher = new MockFetcherThread {
+      override def truncate(topicPartition: TopicPartition, truncationState: OffsetTruncationState): Unit = {
+        assertEquals(highWatermark, truncationState.offset)
+        assertTrue(truncationState.truncationCompleted)
+        super.truncate(topicPartition, truncationState)
+      }
+
+      override def fetchEpochEndOffsets(partitions: Map[TopicPartition, EpochData]): Map[TopicPartition, EpochEndOffset] =
+        throw new UnsupportedOperationException
+
+      override def latestEpoch(topicPartition: TopicPartition): Option[Int] = None
+    }
+
+    val replicaLog = Seq(
+      mkBatch(baseOffset = 0, leaderEpoch = 0, new SimpleRecord("a".getBytes)),
+      mkBatch(baseOffset = 1, leaderEpoch = 2, new SimpleRecord("b".getBytes)),
+      mkBatch(baseOffset = 2, leaderEpoch = 4, new SimpleRecord("c".getBytes)))
+
+    val replicaState = MockFetcherThread.PartitionState(replicaLog, leaderEpoch = 5, highWatermark)
+    fetcher.setReplicaState(partition, replicaState)
+    fetcher.addPartitions(Map(partition -> initialFetchState(highWatermark, leaderEpoch = 5)))
+
+    fetcher.doWork()
+
+    assertEquals(highWatermark, replicaState.logEndOffset)
+    assertEquals(highWatermark, fetcher.fetchState(partition).get.fetchOffset)
+    assertTrue(fetcher.fetchState(partition).get.isReadyForFetch)
+  }
+
+  @Test
+  def testTruncateToHighWatermarkDuringRemovePartitions(): Unit = {
+    val highWatermark = 2L
+    val partition = new TopicPartition("topic", 0)
+    val fetcher = new MockFetcherThread {
+      override def truncateToHighWatermark(partitions: Set[TopicPartition]): Unit = {
+        removePartitions(Set(partition))
+        super.truncateToHighWatermark(partitions)
+      }
+
+      override def latestEpoch(topicPartition: TopicPartition): Option[Int] = None
+    }
+
+    val replicaLog = Seq(
+      mkBatch(baseOffset = 0, leaderEpoch = 0, new SimpleRecord("a".getBytes)),
+      mkBatch(baseOffset = 1, leaderEpoch = 2, new SimpleRecord("b".getBytes)),
+      mkBatch(baseOffset = 2, leaderEpoch = 4, new SimpleRecord("c".getBytes)))
+
+    val replicaState = MockFetcherThread.PartitionState(replicaLog, leaderEpoch = 5, highWatermark)
+    fetcher.setReplicaState(partition, replicaState)
+    fetcher.addPartitions(Map(partition -> initialFetchState(highWatermark, leaderEpoch = 5)))
+
+    fetcher.doWork()
+
+    assertEquals(replicaLog.last.nextOffset(), replicaState.logEndOffset)
+    assertTrue(fetcher.fetchState(partition).isEmpty)
+  }
+
+  @Test
+  def testTruncationSkippedIfNoEpochChange(): Unit = {
+    val partition = new TopicPartition("topic", 0)
+
+    var truncations = 0
+    val fetcher = new MockFetcherThread {
+      override def truncate(topicPartition: TopicPartition, truncationState: OffsetTruncationState): Unit = {
+        truncations += 1
+        super.truncate(topicPartition, truncationState)
+      }
+    }
+
+    val replicaState = MockFetcherThread.PartitionState(leaderEpoch = 5)
+    fetcher.setReplicaState(partition, replicaState)
+    fetcher.addPartitions(Map(partition -> initialFetchState(0L, leaderEpoch = 5)), forceTruncation = true)
+
+    val leaderLog = Seq(
+      mkBatch(baseOffset = 0, leaderEpoch = 1, new SimpleRecord("a".getBytes)),
+      mkBatch(baseOffset = 1, leaderEpoch = 3, new SimpleRecord("b".getBytes)),
+      mkBatch(baseOffset = 2, leaderEpoch = 5, new SimpleRecord("c".getBytes)))
+
+    val leaderState = MockFetcherThread.PartitionState(leaderLog, leaderEpoch = 5, highWatermark = 2L)
+    fetcher.setLeaderState(partition, leaderState)
+
+    // Do one round of truncation
+    fetcher.doWork()
+
+    // We only fetch one record at a time with mock fetcher
+    assertEquals(1, replicaState.logEndOffset)
+    assertEquals(1, truncations)
+
+    // Add partitions again with the same epoch
+    fetcher.addPartitions(Map(partition -> initialFetchState(3L, leaderEpoch = 5)))
+
+    // Verify we did not truncate
+    fetcher.doWork()
+
+    // No truncations occurred and we have fetched another record
+    assertEquals(1, truncations)
+    assertEquals(2, replicaState.logEndOffset)
+  }
+
+  @Test
+  def testFollowerFetchOutOfRangeHigh(): Unit = {
+    val partition = new TopicPartition("topic", 0)
+    val fetcher = new MockFetcherThread()
+
+    val replicaLog = Seq(
+      mkBatch(baseOffset = 0, leaderEpoch = 0, new SimpleRecord("a".getBytes)),
+      mkBatch(baseOffset = 1, leaderEpoch = 2, new SimpleRecord("b".getBytes)),
+      mkBatch(baseOffset = 2, leaderEpoch = 4, new SimpleRecord("c".getBytes)))
+
+    val replicaState = MockFetcherThread.PartitionState(replicaLog, leaderEpoch = 4, highWatermark = 0L)
+    fetcher.setReplicaState(partition, replicaState)
+    fetcher.addPartitions(Map(partition -> initialFetchState(3L, leaderEpoch = 4)))
+
+    val leaderLog = Seq(
+      mkBatch(baseOffset = 0, leaderEpoch = 0, new SimpleRecord("a".getBytes)),
+      mkBatch(baseOffset = 1, leaderEpoch = 2, new SimpleRecord("b".getBytes)),
+      mkBatch(baseOffset = 2, leaderEpoch = 4, new SimpleRecord("c".getBytes)))
+
+    val leaderState = MockFetcherThread.PartitionState(leaderLog, leaderEpoch = 4, highWatermark = 2L)
+    fetcher.setLeaderState(partition, leaderState)
+
+    // initial truncation and verify that the log end offset is updated
+    fetcher.doWork()
+    assertEquals(3L, replicaState.logEndOffset)
+    assertEquals(Option(Fetching), fetcher.fetchState(partition).map(_.state))
+
+    // To hit this case, we have to change the leader log without going through the truncation phase
+    leaderState.log.clear()
+    leaderState.logEndOffset = 0L
+    leaderState.logStartOffset = 0L
+    leaderState.highWatermark = 0L
+
+    fetcher.doWork()
+
+    assertEquals(0L, replicaState.logEndOffset)
+    assertEquals(0L, replicaState.logStartOffset)
+    assertEquals(0L, replicaState.highWatermark)
+  }
+
+  @Test
+  def testFencedOffsetResetAfterOutOfRange(): Unit = {
+    val partition = new TopicPartition("topic", 0)
+    var fetchedEarliestOffset = false
+    val fetcher = new MockFetcherThread() {
+      override protected def fetchEarliestOffsetFromLeader(topicPartition: TopicPartition, leaderEpoch: Int): Long = {
+        fetchedEarliestOffset = true
+        throw new FencedLeaderEpochException(s"Epoch $leaderEpoch is fenced")
+      }
+    }
+
+    val replicaLog = Seq()
+    val replicaState = MockFetcherThread.PartitionState(replicaLog, leaderEpoch = 4, highWatermark = 0L)
+    fetcher.setReplicaState(partition, replicaState)
+    fetcher.addPartitions(Map(partition -> initialFetchState(0L, leaderEpoch = 4)))
+
+    val leaderLog = Seq(
+      mkBatch(baseOffset = 1, leaderEpoch = 2, new SimpleRecord("b".getBytes)),
+      mkBatch(baseOffset = 2, leaderEpoch = 4, new SimpleRecord("c".getBytes)))
+    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 and mark as failed
+    fetcher.doWork()
+    assertEquals(0, replicaState.logEndOffset)
+    assertTrue(fetchedEarliestOffset)
+    assertTrue(fetcher.fetchState(partition).isEmpty)
+    assertTrue(failedPartitions.contains(partition))
+  }
+
+  @Test
+  def testFollowerFetchOutOfRangeLow(): Unit = {
+    val partition = new TopicPartition("topic", 0)
+    val fetcher = new MockFetcherThread
+
+    // The follower begins from an offset which is behind the leader's log start offset
+    val replicaLog = Seq(
+      mkBatch(baseOffset = 0, leaderEpoch = 0, new SimpleRecord("a".getBytes)))
+
+    val replicaState = MockFetcherThread.PartitionState(replicaLog, leaderEpoch = 0, highWatermark = 0L)
+    fetcher.setReplicaState(partition, replicaState)
+    fetcher.addPartitions(Map(partition -> initialFetchState(3L, leaderEpoch = 0)))
+
+    val leaderLog = Seq(
+      mkBatch(baseOffset = 2, leaderEpoch = 4, new SimpleRecord("c".getBytes)))
+
+    val leaderState = MockFetcherThread.PartitionState(leaderLog, leaderEpoch = 0, highWatermark = 2L)
+    fetcher.setLeaderState(partition, leaderState)
 
-    override def handleOffsetOutOfRange(topicPartition: TopicPartition): Long = 0L
+    // initial truncation and verify that the log start offset is updated
+    fetcher.doWork()
+    if (truncateOnFetch) {
+      // Second iteration required here since first iteration is required to
+      // perform initial truncaton based on diverging epoch.
+      fetcher.doWork()
+    }
+    assertEquals(Option(Fetching), fetcher.fetchState(partition).map(_.state))
+    assertEquals(2, replicaState.logStartOffset)
+    assertEquals(List(), replicaState.log.toList)
+
+    TestUtils.waitUntilTrue(() => {
+      fetcher.doWork()
+      fetcher.replicaPartitionState(partition).log == fetcher.leaderPartitionState(partition).log
+    }, "Failed to reconcile leader and follower logs")
+
+    assertEquals(leaderState.logStartOffset, replicaState.logStartOffset)
+    assertEquals(leaderState.logEndOffset, replicaState.logEndOffset)
+    assertEquals(leaderState.highWatermark, replicaState.highWatermark)
+  }
+
+  @Test
+  def testRetryAfterUnknownLeaderEpochInLatestOffsetFetch(): Unit = {
+    val partition = new TopicPartition("topic", 0)
+    val fetcher: MockFetcherThread = new MockFetcherThread {
+      val tries = new AtomicInteger(0)
+      override protected def fetchLatestOffsetFromLeader(topicPartition: TopicPartition, leaderEpoch: Int): Long = {
+        if (tries.getAndIncrement() == 0)
+          throw new UnknownLeaderEpochException("Unexpected leader epoch")
+        super.fetchLatestOffsetFromLeader(topicPartition, leaderEpoch)
+      }
+    }
 
-    override def handlePartitionsWithErrors(partitions: Iterable[TopicPartition]): Unit = {}
+    // The follower begins from an offset which is behind the leader's log start offset
+    val replicaLog = Seq(
+      mkBatch(baseOffset = 0, leaderEpoch = 0, new SimpleRecord("a".getBytes)))
 
-    override protected def fetch(fetchRequest: DummyFetchRequest): Seq[(TopicPartition, TestPartitionData)] =
-      fetchRequest.offsets.mapValues(_ => new TestPartitionData()).toSeq
+    val replicaState = MockFetcherThread.PartitionState(replicaLog, leaderEpoch = 0, highWatermark = 0L)
+    fetcher.setReplicaState(partition, replicaState)
+    fetcher.addPartitions(Map(partition -> initialFetchState(3L, leaderEpoch = 0)))
 
-    override protected def buildFetchRequest(partitionMap: collection.Seq[(TopicPartition, PartitionFetchState)]): ResultWithPartitions[DummyFetchRequest] =
-      ResultWithPartitions(new DummyFetchRequest(partitionMap.map { case (k, v) => (k, v.fetchOffset) }.toMap), Set())
+    val leaderLog = Seq(
+      mkBatch(baseOffset = 2, leaderEpoch = 4, new SimpleRecord("c".getBytes)))
 
-    override def buildLeaderEpochRequest(allPartitions: Seq[(TopicPartition, PartitionFetchState)]): ResultWithPartitions[Map[TopicPartition, Int]] = {
-      ResultWithPartitions(Map(), Set())
+    val leaderState = MockFetcherThread.PartitionState(leaderLog, leaderEpoch = 0, highWatermark = 2L)
+    fetcher.setLeaderState(partition, leaderState)
+
+    // initial truncation and initial error response handling
+    fetcher.doWork()
+    assertEquals(Option(Fetching), fetcher.fetchState(partition).map(_.state))
+
+    TestUtils.waitUntilTrue(() => {
+      fetcher.doWork()
+      fetcher.replicaPartitionState(partition).log == fetcher.leaderPartitionState(partition).log
+    }, "Failed to reconcile leader and follower logs")
+
+    assertEquals(leaderState.logStartOffset, replicaState.logStartOffset)
+    assertEquals(leaderState.logEndOffset, replicaState.logEndOffset)
+    assertEquals(leaderState.highWatermark, replicaState.highWatermark)
+  }
+
+  @Test
+  def testCorruptMessage(): Unit = {
+    val partition = new TopicPartition("topic", 0)
+
+    val fetcher = new MockFetcherThread {
+      var fetchedOnce = false
+      override def fetchFromLeader(fetchRequest: FetchRequest.Builder): Map[TopicPartition, FetchData] = {
+        val fetchedData = super.fetchFromLeader(fetchRequest)
+        if (!fetchedOnce) {
+          val records = fetchedData.head._2.records.asInstanceOf[MemoryRecords]
+          val buffer = records.buffer()
+          buffer.putInt(15, buffer.getInt(15) ^ 23422)
+          buffer.putInt(30, buffer.getInt(30) ^ 93242)
+          fetchedOnce = true
+        }
+        fetchedData
+      }
     }
 
-    override def fetchEpochsFromLeader(partitions: Map[TopicPartition, Int]): Map[TopicPartition, EpochEndOffset] = { Map() }
+    fetcher.setReplicaState(partition, MockFetcherThread.PartitionState(leaderEpoch = 0))
+    fetcher.addPartitions(Map(partition -> initialFetchState(0L, leaderEpoch = 0)))
 
-    override def maybeTruncate(fetchedEpochs: Map[TopicPartition, EpochEndOffset]): ResultWithPartitions[Map[TopicPartition, Long]] = {
-      ResultWithPartitions(Map(), Set())
+    val batch = mkBatch(baseOffset = 0L, leaderEpoch = 0,
+      new SimpleRecord("a".getBytes), new SimpleRecord("b".getBytes))
+    val leaderState = MockFetcherThread.PartitionState(Seq(batch), leaderEpoch = 0, highWatermark = 2L)
+    fetcher.setLeaderState(partition, leaderState)
+
+    fetcher.doWork() // fails with corrupt record
+    fetcher.doWork() // should succeed
+
+    val replicaState = fetcher.replicaPartitionState(partition)
+    assertEquals(2L, replicaState.logEndOffset)
+  }
+
+  @Test
+  def testLeaderEpochChangeDuringFencedFetchEpochsFromLeader(): Unit = {
+    // The leader is on the new epoch when the OffsetsForLeaderEpoch with old epoch is sent, so it
+    // returns the fence error. Validate that response is ignored if the leader epoch changes on
+    // the follower while OffsetsForLeaderEpoch request is in flight, but able to truncate and fetch
+    // in the next of round of "doWork"
+    testLeaderEpochChangeDuringFetchEpochsFromLeader(leaderEpochOnLeader = 1)
+  }
+
+  @Test
+  def testLeaderEpochChangeDuringSuccessfulFetchEpochsFromLeader(): Unit = {
+    // The leader is on the old epoch when the OffsetsForLeaderEpoch with old epoch is sent
+    // and returns the valid response. Validate that response is ignored if the leader epoch changes
+    // on the follower while OffsetsForLeaderEpoch request is in flight, but able to truncate and
+    // fetch once the leader is on the newer epoch (same as follower)
+    testLeaderEpochChangeDuringFetchEpochsFromLeader(leaderEpochOnLeader = 0)
+  }
+
+  private def testLeaderEpochChangeDuringFetchEpochsFromLeader(leaderEpochOnLeader: Int): Unit = {
+    val partition = new TopicPartition("topic", 0)
+    val initialLeaderEpochOnFollower = 0
+    val nextLeaderEpochOnFollower = initialLeaderEpochOnFollower + 1
+
+    val fetcher = new MockFetcherThread {
+      var fetchEpochsFromLeaderOnce = false
+      override def fetchEpochEndOffsets(partitions: Map[TopicPartition, EpochData]): Map[TopicPartition, EpochEndOffset] = {
+        val fetchedEpochs = super.fetchEpochEndOffsets(partitions)
+        if (!fetchEpochsFromLeaderOnce) {
+          // leader epoch changes while fetching epochs from leader
+          removePartitions(Set(partition))
+          setReplicaState(partition, MockFetcherThread.PartitionState(leaderEpoch = nextLeaderEpochOnFollower))
+          addPartitions(Map(partition -> initialFetchState(0L, leaderEpoch = nextLeaderEpochOnFollower)), forceTruncation = true)
+          fetchEpochsFromLeaderOnce = true
+        }
+        fetchedEpochs
+      }
     }
+
+    fetcher.setReplicaState(partition, MockFetcherThread.PartitionState(leaderEpoch = initialLeaderEpochOnFollower))
+    fetcher.addPartitions(Map(partition -> initialFetchState(0L, leaderEpoch = initialLeaderEpochOnFollower)), forceTruncation = true)
+
+    val leaderLog = Seq(
+      mkBatch(baseOffset = 0, leaderEpoch = initialLeaderEpochOnFollower, new SimpleRecord("c".getBytes)))
+    val leaderState = MockFetcherThread.PartitionState(leaderLog, leaderEpochOnLeader, highWatermark = 0L)
+    fetcher.setLeaderState(partition, leaderState)
+
+    // first round of truncation
+    fetcher.doWork()
+
+    // Since leader epoch changed, fetch epochs response is ignored due to partition being in
+    // truncating state with the updated leader epoch
+    assertEquals(Option(Truncating), fetcher.fetchState(partition).map(_.state))
+    assertEquals(Option(nextLeaderEpochOnFollower), fetcher.fetchState(partition).map(_.currentLeaderEpoch))
+
+    if (leaderEpochOnLeader < nextLeaderEpochOnFollower) {
+      fetcher.setLeaderState(
+        partition, MockFetcherThread.PartitionState(leaderLog, nextLeaderEpochOnFollower, highWatermark = 0L))
+    }
+
+    // make sure the fetcher is now able to truncate and fetch
+    fetcher.doWork()
+    assertEquals(fetcher.leaderPartitionState(partition).log, fetcher.replicaPartitionState(partition).log)
+  }
+
+  @Test
+  def testTruncateToEpochEndOffsetsDuringRemovePartitions(): Unit = {
+    val partition = new TopicPartition("topic", 0)
+    val leaderEpochOnLeader = 0
+    val initialLeaderEpochOnFollower = 0
+    val nextLeaderEpochOnFollower = initialLeaderEpochOnFollower + 1
+
+    val fetcher = new MockFetcherThread {
+      override def fetchEpochEndOffsets(partitions: Map[TopicPartition, EpochData]): Map[TopicPartition, EpochEndOffset] = {
+        val fetchedEpochs = super.fetchEpochEndOffsets(partitions)
+        // leader epoch changes while fetching epochs from leader
+        // at the same time, the replica fetcher manager removes the partition
+        removePartitions(Set(partition))
+        setReplicaState(partition, MockFetcherThread.PartitionState(leaderEpoch = nextLeaderEpochOnFollower))
+        fetchedEpochs
+      }
+    }
+
+    fetcher.setReplicaState(partition, MockFetcherThread.PartitionState(leaderEpoch = initialLeaderEpochOnFollower))
+    fetcher.addPartitions(Map(partition -> initialFetchState(0L, leaderEpoch = initialLeaderEpochOnFollower)))
+
+    val leaderLog = Seq(
+      mkBatch(baseOffset = 0, leaderEpoch = initialLeaderEpochOnFollower, new SimpleRecord("c".getBytes)))
+    val leaderState = MockFetcherThread.PartitionState(leaderLog, leaderEpochOnLeader, highWatermark = 0L)
+    fetcher.setLeaderState(partition, leaderState)
+
+    // first round of work
+    fetcher.doWork()
+
+    // since the partition was removed before the fetched endOffsets were filtered against the leader epoch,
+    // we do not expect the partition to be in Truncating state
+    assertEquals(None, fetcher.fetchState(partition).map(_.state))
+    assertEquals(None, fetcher.fetchState(partition).map(_.currentLeaderEpoch))
+
+    fetcher.setLeaderState(
+      partition, MockFetcherThread.PartitionState(leaderLog, nextLeaderEpochOnFollower, highWatermark = 0L))
+
+    // make sure the fetcher is able to continue work
+    fetcher.doWork()
+    assertEquals(ArrayBuffer.empty, fetcher.replicaPartitionState(partition).log)
   }
 
+  @Test
+  def testTruncationThrowsExceptionIfLeaderReturnsPartitionsNotRequestedInFetchEpochs(): Unit = {
+    val partition = new TopicPartition("topic", 0)
+    val fetcher = new MockFetcherThread {
+      override def fetchEpochEndOffsets(partitions: Map[TopicPartition, EpochData]): Map[TopicPartition, EpochEndOffset] = {
+        val unrequestedTp = new TopicPartition("topic2", 0)
+        super.fetchEpochEndOffsets(partitions).toMap + (unrequestedTp -> new EpochEndOffset()
+          .setPartition(unrequestedTp.partition)
+          .setErrorCode(Errors.NONE.code)
+          .setLeaderEpoch(0)
+          .setEndOffset(0))
+      }
+    }
+
+    fetcher.setReplicaState(partition, MockFetcherThread.PartitionState(leaderEpoch = 0))
+    fetcher.addPartitions(Map(partition -> initialFetchState(0L, leaderEpoch = 0)), forceTruncation = true)
+    fetcher.setLeaderState(partition, MockFetcherThread.PartitionState(leaderEpoch = 0))
+
+    // first round of truncation should throw an exception
+    assertThrows[IllegalStateException] {
+      fetcher.doWork()
+    }
+  }
+
+  @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 -> initialFetchState(0L, leaderEpoch = 0)), forceTruncation = true)
+    fetcher.setLeaderState(partition1, MockFetcherThread.PartitionState(leaderEpoch = 0))
+
+    fetcher.setReplicaState(partition2, MockFetcherThread.PartitionState(leaderEpoch = 0))
+    fetcher.addPartitions(Map(partition2 -> initialFetchState(0L, leaderEpoch = 0)), forceTruncation = true)
+    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 -> initialFetchState(0L, leaderEpoch = 1)), forceTruncation = true)
+
+    // partition1 added back
+    assertEquals(Some(Truncating), fetcher.fetchState(partition1).map(_.state))
+    assertFalse(failedPartitions.contains(partition1))
+
+  }
 
   @Test
-  def testFetchRequestCorruptedMessageException() {
+  def testDivergingEpochs(): Unit = {
     val partition = new TopicPartition("topic", 0)
-    val fetcherThread = new CorruptingFetcherThread("test", "client", new BrokerEndPoint(0, "localhost", 9092),
-      fetchBackOffMs = 1)
+    val fetcher = new MockFetcherThread
+
+    val replicaLog = Seq(
+      mkBatch(baseOffset = 0, leaderEpoch = 0, new SimpleRecord("a".getBytes)),
+      mkBatch(baseOffset = 1, leaderEpoch = 2, new SimpleRecord("b".getBytes)),
+      mkBatch(baseOffset = 2, leaderEpoch = 4, new SimpleRecord("c".getBytes)))
 
-    fetcherThread.start()
+    val replicaState = MockFetcherThread.PartitionState(replicaLog, leaderEpoch = 5, highWatermark = 0L)
+    fetcher.setReplicaState(partition, replicaState)
+    fetcher.addPartitions(Map(partition -> initialFetchState(3L, leaderEpoch = 5)))
+    assertEquals(3L, replicaState.logEndOffset)
+    fetcher.verifyLastFetchedEpoch(partition, expectedEpoch = Some(4))
 
-    // Add one partition for fetching
-    fetcherThread.addPartitions(Map(partition -> 0L))
+    val leaderLog = Seq(
+      mkBatch(baseOffset = 0, leaderEpoch = 0, new SimpleRecord("a".getBytes)),
+      mkBatch(baseOffset = 1, leaderEpoch = 2, new SimpleRecord("b".getBytes)),
+      mkBatch(baseOffset = 2, leaderEpoch = 5, new SimpleRecord("d".getBytes)))
 
-    // Wait until fetcherThread finishes the work
-    TestUtils.waitUntilTrue(() => fetcherThread.fetchCount > 3, "Failed waiting for fetcherThread to finish the work")
+    val leaderState = MockFetcherThread.PartitionState(leaderLog, leaderEpoch = 5, highWatermark = 2L)
+    fetcher.setLeaderState(partition, leaderState)
 
-    fetcherThread.shutdown()
+    fetcher.doWork()
+    fetcher.verifyLastFetchedEpoch(partition, Some(2))
 
-    // The fetcherThread should have fetched two normal messages
-    assertTrue(fetcherThread.logEndOffset == 2)
+    TestUtils.waitUntilTrue(() => {
+      fetcher.doWork()
+      fetcher.replicaPartitionState(partition).log == fetcher.leaderPartitionState(partition).log
+    }, "Failed to reconcile leader and follower logs")
+    fetcher.verifyLastFetchedEpoch(partition, Some(5))
   }
 
-  class CorruptingFetcherThread(name: String,
-                                clientId: String,
-                                sourceBroker: BrokerEndPoint,
-                                fetchBackOffMs: Int = 0)
-    extends DummyFetcherThread(name, clientId, sourceBroker, fetchBackOffMs) {
+  object MockFetcherThread {
+    class PartitionState(var log: mutable.Buffer[RecordBatch],
+                         var leaderEpoch: Int,
+                         var logStartOffset: Long,
+                         var logEndOffset: Long,
+                         var highWatermark: Long)
+
+    object PartitionState {
+      def apply(log: Seq[RecordBatch], leaderEpoch: Int, highWatermark: Long): PartitionState = {
+        val logStartOffset = log.headOption.map(_.baseOffset).getOrElse(0L)
+        val logEndOffset = log.lastOption.map(_.nextOffset).getOrElse(0L)
+        new PartitionState(log.toBuffer, leaderEpoch, logStartOffset, logEndOffset, highWatermark)
+      }
+
+      def apply(leaderEpoch: Int): PartitionState = {
+        apply(Seq(), leaderEpoch = leaderEpoch, highWatermark = 0L)
+      }
+    }
+  }
+
+  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()),
+      failedPartitions,
+      brokerTopicStats = new BrokerTopicStats) {
+
+    import MockFetcherThread.PartitionState
+
+    private val replicaPartitionStates = mutable.Map[TopicPartition, PartitionState]()
+    private val leaderPartitionStates = mutable.Map[TopicPartition, PartitionState]()
+    private var latestEpochDefault: Option[Int] = Some(0)
 
-    @volatile var logEndOffset = 0L
-    @volatile var fetchCount = 0
+    def setLeaderState(topicPartition: TopicPartition, state: PartitionState): Unit = {
+      leaderPartitionStates.put(topicPartition, state)
+    }
 
-    private val normalPartitionDataSet = List(
-      new TestPartitionData(MemoryRecords.withRecords(0L, CompressionType.NONE, new SimpleRecord("hello".getBytes()))),
-      new TestPartitionData(MemoryRecords.withRecords(1L, CompressionType.NONE, new SimpleRecord("hello".getBytes())))
-    )
+    def setReplicaState(topicPartition: TopicPartition, state: PartitionState): Unit = {
+      replicaPartitionStates.put(topicPartition, state)
+    }
+
+    def replicaPartitionState(topicPartition: TopicPartition): PartitionState = {
+      replicaPartitionStates.getOrElse(topicPartition,
+        throw new IllegalArgumentException(s"Unknown partition $topicPartition"))
+    }
+
+    def leaderPartitionState(topicPartition: TopicPartition): PartitionState = {
+      leaderPartitionStates.getOrElse(topicPartition,
+        throw new IllegalArgumentException(s"Unknown partition $topicPartition"))
+    }
+
+    def addPartitions(initialFetchStates: Map[TopicPartition, InitialFetchState], forceTruncation: Boolean): Set[TopicPartition] = {
+      latestEpochDefault = if (forceTruncation) None else Some(0)
+      val partitions = super.addPartitions(initialFetchStates)
+      latestEpochDefault = Some(0)
+      partitions
+    }
 
     override def processPartitionData(topicPartition: TopicPartition,
                                       fetchOffset: Long,
-                                      partitionData: PartitionData): Unit = {
+                                      partitionData: FetchData): Option[LogAppendInfo] = {
+      val state = replicaPartitionState(topicPartition)
+
+      if (isTruncationOnFetchSupported && partitionData.divergingEpoch.isPresent) {
+        val divergingEpoch = partitionData.divergingEpoch.get
+        truncateOnFetchResponse(Map(topicPartition -> new EpochEndOffset()
+          .setPartition(topicPartition.partition)
+          .setErrorCode(Errors.NONE.code)
+          .setLeaderEpoch(divergingEpoch.epoch)
+          .setEndOffset(divergingEpoch.endOffset)))
+        return None
+      }
+
       // Throw exception if the fetchOffset does not match the fetcherThread partition state
-      if (fetchOffset != logEndOffset)
-        throw new RuntimeException(
-          "Offset mismatch for partition %s: fetched offset = %d, log end offset = %d."
-            .format(topicPartition, fetchOffset, logEndOffset))
+      if (fetchOffset != state.logEndOffset)
+        throw new RuntimeException(s"Offset mismatch for partition $topicPartition: " +
+          s"fetched offset = $fetchOffset, log end offset = ${state.logEndOffset}.")
 
       // Now check message's crc
-      val records = partitionData.toRecords
-      for (batch <- records.batches.asScala) {
+      val batches = partitionData.records.batches.asScala
+      var maxTimestamp = RecordBatch.NO_TIMESTAMP
+      var offsetOfMaxTimestamp = -1L
+      var lastOffset = state.logEndOffset
+      var lastEpoch: Option[Int] = None
+
+      for (batch <- batches) {
         batch.ensureValid()
-        logEndOffset = batch.nextOffset
+        if (batch.maxTimestamp > maxTimestamp) {
+          maxTimestamp = batch.maxTimestamp
+          offsetOfMaxTimestamp = batch.baseOffset
+        }
+        state.log.append(batch)
+        state.logEndOffset = batch.nextOffset
+        lastOffset = batch.lastOffset
+        lastEpoch = Some(batch.partitionLeaderEpoch)
       }
+
+      state.logStartOffset = partitionData.logStartOffset
+      state.highWatermark = partitionData.highWatermark
+
+      Some(LogAppendInfo(firstOffset = Some(fetchOffset),
+        lastOffset = lastOffset,
+        lastLeaderEpoch = lastEpoch,
+        maxTimestamp = maxTimestamp,
+        offsetOfMaxTimestamp = offsetOfMaxTimestamp,
+        logAppendTime = Time.SYSTEM.milliseconds(),
+        logStartOffset = state.logStartOffset,
+        recordConversionStats = RecordConversionStats.EMPTY,
+        sourceCodec = NoCompressionCodec,
+        targetCodec = NoCompressionCodec,
+        shallowCount = batches.size,
+        validBytes = partitionData.records.sizeInBytes,
+        offsetsMonotonic = true,
+        lastOffsetOfFirstBatch = batches.headOption.map(_.lastOffset).getOrElse(-1)))
     }
 
-    override protected def fetch(fetchRequest: DummyFetchRequest): Seq[(TopicPartition, TestPartitionData)] = {
-      fetchCount += 1
-      // Set the first fetch to get a corrupted message
-      if (fetchCount == 1) {
-        val record = new SimpleRecord("hello".getBytes())
-        val records = MemoryRecords.withRecords(CompressionType.NONE, record)
-        val buffer = records.buffer
+    override def truncate(topicPartition: TopicPartition, truncationState: OffsetTruncationState): Unit = {
+      val state = replicaPartitionState(topicPartition)
+      state.log = state.log.takeWhile { batch =>
+        batch.lastOffset < truncationState.offset
+      }
+      state.logEndOffset = state.log.lastOption.map(_.lastOffset + 1).getOrElse(state.logStartOffset)
+      state.highWatermark = math.min(state.highWatermark, state.logEndOffset)
+    }
+
+    override def truncateFullyAndStartAt(topicPartition: TopicPartition, offset: Long): Unit = {
+      val state = replicaPartitionState(topicPartition)
+      state.log.clear()
+      state.logStartOffset = offset
+      state.logEndOffset = offset
+      state.highWatermark = offset
+    }
 
-        // flip some bits in the message to ensure the crc fails
-        buffer.putInt(15, buffer.getInt(15) ^ 23422)
-        buffer.putInt(30, buffer.getInt(30) ^ 93242)
-        fetchRequest.offsets.mapValues(_ => new TestPartitionData(records)).toSeq
+    override def buildFetch(partitionMap: Map[TopicPartition, PartitionFetchState]): ResultWithPartitions[Option[ReplicaFetch]] = {
+      val fetchData = mutable.Map.empty[TopicPartition, FetchRequest.PartitionData]
+      partitionMap.foreach { case (partition, state) =>
+        if (state.isReadyForFetch) {
+          val replicaState = replicaPartitionState(partition)
+          val lastFetchedEpoch = if (isTruncationOnFetchSupported)
+            state.lastFetchedEpoch.map(_.asInstanceOf[Integer]).asJava
+          else
+            Optional.empty[Integer]
+          fetchData.put(partition, new FetchRequest.PartitionData(state.fetchOffset, replicaState.logStartOffset,
+            1024 * 1024, Optional.of[Integer](state.currentLeaderEpoch), lastFetchedEpoch))
+        }
+      }
+      val fetchRequest = FetchRequest.Builder.forReplica(ApiKeys.FETCH.latestVersion, replicaId, 0, 1, fetchData.asJava)
+      ResultWithPartitions(Some(ReplicaFetch(fetchData.asJava, fetchRequest)), Set.empty)
+    }
+
+    override def latestEpoch(topicPartition: TopicPartition): Option[Int] = {
+      val state = replicaPartitionState(topicPartition)
+      state.log.lastOption.map(_.partitionLeaderEpoch).orElse(latestEpochDefault)
+    }
+
+    override def logStartOffset(topicPartition: TopicPartition): Long = replicaPartitionState(topicPartition).logStartOffset
+
+    override def logEndOffset(topicPartition: TopicPartition): Long = replicaPartitionState(topicPartition).logEndOffset
+
+    override def endOffsetForEpoch(topicPartition: TopicPartition, epoch: Int): Option[OffsetAndEpoch] = {
+      val epochData = new EpochData(Optional.empty[Integer](), epoch)
+      val result = lookupEndOffsetForEpoch(topicPartition, epochData, replicaPartitionState(topicPartition))
+      if (result.endOffset == UNDEFINED_EPOCH_OFFSET)
+        None
+      else
+        Some(OffsetAndEpoch(result.endOffset, result.leaderEpoch))
+    }
+
+    private def checkExpectedLeaderEpoch(expectedEpochOpt: Optional[Integer],
+                                         partitionState: PartitionState): Option[Errors] = {
+      if (expectedEpochOpt.isPresent) {
+        val expectedEpoch = expectedEpochOpt.get
+        if (expectedEpoch < partitionState.leaderEpoch)
+          Some(Errors.FENCED_LEADER_EPOCH)
+        else if (expectedEpoch > partitionState.leaderEpoch)
+          Some(Errors.UNKNOWN_LEADER_EPOCH)
+        else
+          None
       } else {
-        // Then, the following fetches get the normal data
-        fetchRequest.offsets.mapValues(v => normalPartitionDataSet(v.toInt)).toSeq
+        None
       }
     }
 
-    override protected def buildFetchRequest(partitionMap: collection.Seq[(TopicPartition, PartitionFetchState)]): ResultWithPartitions[DummyFetchRequest] = {
-      val requestMap = new mutable.HashMap[TopicPartition, Long]
-      partitionMap.foreach { case (topicPartition, partitionFetchState) =>
-        // Add backoff delay check
-        if (partitionFetchState.isReadyForFetch)
-          requestMap.put(topicPartition, partitionFetchState.fetchOffset)
+    def verifyLastFetchedEpoch(partition: TopicPartition, expectedEpoch: Option[Int]): Unit = {
+      if (isTruncationOnFetchSupported) {
+        assertEquals(Some(Fetching), fetchState(partition).map(_.state))
+        assertEquals(expectedEpoch, fetchState(partition).flatMap(_.lastFetchedEpoch))
+      }
+    }
+
+    private def divergingEpochAndOffset(partition: TopicPartition,
+                                        lastFetchedEpoch: Optional[Integer],
+                                        fetchOffset: Long,
+                                        partitionState: PartitionState): Option[FetchResponseData.EpochEndOffset] = {
+      lastFetchedEpoch.asScala.flatMap { fetchEpoch =>
+        val epochEndOffset = fetchEpochEndOffsets(Map(partition -> new EpochData(Optional.empty[Integer], fetchEpoch)))(partition)
+
+        if (partitionState.log.isEmpty
+            || epochEndOffset.endOffset == UNDEFINED_EPOCH_OFFSET
+            || epochEndOffset.leaderEpoch == UNDEFINED_EPOCH)
+          None
+        else if (epochEndOffset.leaderEpoch < fetchEpoch || epochEndOffset.endOffset < fetchOffset) {
+          Some(new FetchResponseData.EpochEndOffset()
+            .setEpoch(epochEndOffset.leaderEpoch)
+            .setEndOffset(epochEndOffset.endOffset))
+        } else
+          None
+      }
+    }
+
+    private def lookupEndOffsetForEpoch(topicPartition: TopicPartition,
+                                        epochData: EpochData,
+                                        partitionState: PartitionState): EpochEndOffset = {
+      checkExpectedLeaderEpoch(epochData.currentLeaderEpoch, partitionState).foreach { error =>
+        return new EpochEndOffset()
+          .setPartition(topicPartition.partition)
+          .setErrorCode(error.code)
+      }
+
+      var epochLowerBound = UNDEFINED_EPOCH
+      for (batch <- partitionState.log) {
+        if (batch.partitionLeaderEpoch > epochData.leaderEpoch) {
+          // If we don't have the requested epoch, return the next higher entry
+          if (epochLowerBound == UNDEFINED_EPOCH)
+            return new EpochEndOffset()
+              .setPartition(topicPartition.partition)
+              .setErrorCode(Errors.NONE.code)
+              .setLeaderEpoch(batch.partitionLeaderEpoch)
+              .setEndOffset(batch.baseOffset)
+          else
+            return new EpochEndOffset()
+              .setPartition(topicPartition.partition)
+              .setErrorCode(Errors.NONE.code)
+              .setLeaderEpoch(epochLowerBound)
+              .setEndOffset(batch.baseOffset)
+        }
+        epochLowerBound = batch.partitionLeaderEpoch
       }
-      ResultWithPartitions(new DummyFetchRequest(requestMap), Set())
+      new EpochEndOffset()
+        .setPartition(topicPartition.partition)
+        .setErrorCode(Errors.NONE.code)
     }
 
-    override def handlePartitionsWithErrors(partitions: Iterable[TopicPartition]) = delayPartitions(partitions, fetchBackOffMs.toLong)
+    override def fetchEpochEndOffsets(partitions: Map[TopicPartition, EpochData]): Map[TopicPartition, EpochEndOffset] = {
+      val endOffsets = mutable.Map[TopicPartition, EpochEndOffset]()
+      partitions.foreach { case (partition, epochData) =>
+        val leaderState = leaderPartitionState(partition)
+        val epochEndOffset = lookupEndOffsetForEpoch(partition, epochData, leaderState)
+        endOffsets.put(partition, epochEndOffset)
+      }
+      endOffsets
+    }
+
+    override protected val isOffsetForLeaderEpochSupported: Boolean = true
+
+    override protected val isTruncationOnFetchSupported: Boolean = truncateOnFetch
+
+    override def fetchFromLeader(fetchRequest: FetchRequest.Builder): Map[TopicPartition, FetchData] = {
+      fetchRequest.fetchData.asScala.map { case (partition, fetchData) =>
+        val leaderState = leaderPartitionState(partition)
+        val epochCheckError = checkExpectedLeaderEpoch(fetchData.currentLeaderEpoch, leaderState)
+        val divergingEpoch = divergingEpochAndOffset(partition, fetchData.lastFetchedEpoch, fetchData.fetchOffset, leaderState)
+
+        val (error, records) = if (epochCheckError.isDefined) {
+          (epochCheckError.get, MemoryRecords.EMPTY)
+        } else if (fetchData.fetchOffset > leaderState.logEndOffset || fetchData.fetchOffset < leaderState.logStartOffset) {
+          (Errors.OFFSET_OUT_OF_RANGE, MemoryRecords.EMPTY)
+        } else if (divergingEpoch.nonEmpty) {
+          (Errors.NONE, MemoryRecords.EMPTY)
+        } else {
+          // for simplicity, we fetch only one batch at a time
+          val records = leaderState.log.find(_.baseOffset >= fetchData.fetchOffset) match {
+            case Some(batch) =>
+              val buffer = ByteBuffer.allocate(batch.sizeInBytes)
+              batch.writeTo(buffer)
+              buffer.flip()
+              MemoryRecords.readableRecords(buffer)
+
+            case None =>
+              MemoryRecords.EMPTY
+          }
+
+          (Errors.NONE, records)
+        }
+
+        (partition, new FetchData(error, leaderState.highWatermark, leaderState.highWatermark, leaderState.logStartOffset,
+          Optional.empty[Integer], List.empty.asJava, divergingEpoch.asJava, records))
+      }.toMap
+    }
+
+    private def checkLeaderEpochAndThrow(expectedEpoch: Int, partitionState: PartitionState): Unit = {
+      checkExpectedLeaderEpoch(Optional.of[Integer](expectedEpoch), partitionState).foreach { error =>
+        throw error.exception()
+      }
+    }
+
+    override protected def fetchEarliestOffsetFromLeader(topicPartition: TopicPartition, leaderEpoch: Int): Long = {
+      val leaderState = leaderPartitionState(topicPartition)
+      checkLeaderEpochAndThrow(leaderEpoch, leaderState)
+      leaderState.logStartOffset
+    }
+
+    override protected def fetchLatestOffsetFromLeader(topicPartition: TopicPartition, leaderEpoch: Int): Long = {
+      val leaderState = leaderPartitionState(topicPartition)
+      checkLeaderEpochAndThrow(leaderEpoch, leaderState)
+      leaderState.logEndOffset
+    }
 
   }
 
diff --git a/core/src/test/scala/unit/kafka/server/AbstractFetcherThreadWithIbp26Test.scala b/core/src/test/scala/unit/kafka/server/AbstractFetcherThreadWithIbp26Test.scala
new file mode 100644
index 0000000000000..a85246567c50a
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/AbstractFetcherThreadWithIbp26Test.scala
@@ -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 kafka.server
+
+class AbstractFetcherThreadWithIbp26Test extends AbstractFetcherThreadTest {
+
+  override val truncateOnFetch = false
+}
diff --git a/core/src/test/scala/unit/kafka/server/AddPartitionsToTxnRequestServerTest.scala b/core/src/test/scala/unit/kafka/server/AddPartitionsToTxnRequestServerTest.scala
new file mode 100644
index 0000000000000..cb97127438e74
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/AddPartitionsToTxnRequestServerTest.scala
@@ -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 kafka.server
+
+import java.util.Properties
+
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.requests.{AddPartitionsToTxnRequest, AddPartitionsToTxnResponse}
+import org.junit.Assert._
+import org.junit.{Before, Test}
+
+import scala.jdk.CollectionConverters._
+
+class AddPartitionsToTxnRequestServerTest extends BaseRequestTest {
+  private val topic1 = "topic1"
+  val numPartitions = 1
+
+  override def brokerPropertyOverrides(properties: Properties): Unit =
+    properties.put(KafkaConfig.AutoCreateTopicsEnableProp, false.toString)
+
+  @Before
+  override def setUp(): Unit = {
+    super.setUp()
+    createTopic(topic1, numPartitions, servers.size, new Properties())
+  }
+
+  @Test
+  def shouldReceiveOperationNotAttemptedWhenOtherPartitionHasError(): Unit = {
+    // The basic idea is that we have one unknown topic and one created topic. We should get the 'UNKNOWN_TOPIC_OR_PARTITION'
+    // error for the unknown topic and the 'OPERATION_NOT_ATTEMPTED' error for the known and authorized topic.
+    val nonExistentTopic = new TopicPartition("unknownTopic", 0)
+    val createdTopicPartition = new TopicPartition(topic1, 0)
+
+    val transactionalId = "foobar"
+    val producerId = 1000L
+    val producerEpoch: Short = 0
+
+    val request = new AddPartitionsToTxnRequest.Builder(
+      transactionalId,
+      producerId,
+      producerEpoch,
+      List(createdTopicPartition, nonExistentTopic).asJava)
+      .build()
+
+    val leaderId = servers.head.config.brokerId
+    val response = connectAndReceive[AddPartitionsToTxnResponse](request, brokerSocketServer(leaderId))
+
+    assertEquals(2, response.errors.size)
+
+    assertTrue(response.errors.containsKey(createdTopicPartition))
+    assertEquals(Errors.OPERATION_NOT_ATTEMPTED, response.errors.get(createdTopicPartition))
+
+    assertTrue(response.errors.containsKey(nonExistentTopic))
+    assertEquals(Errors.UNKNOWN_TOPIC_OR_PARTITION, response.errors.get(nonExistentTopic))
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/server/AddPartitionsToTxnRequestTest.scala b/core/src/test/scala/unit/kafka/server/AddPartitionsToTxnRequestTest.scala
deleted file mode 100644
index 62dd4c495d552..0000000000000
--- a/core/src/test/scala/unit/kafka/server/AddPartitionsToTxnRequestTest.scala
+++ /dev/null
@@ -1,76 +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.util.Properties
-
-import kafka.utils.TestUtils
-import org.apache.kafka.common.TopicPartition
-import org.apache.kafka.common.protocol.{ApiKeys, Errors}
-import org.apache.kafka.common.requests.{AddPartitionsToTxnRequest, AddPartitionsToTxnResponse}
-import org.junit.{Before, Test}
-import org.junit.Assert._
-
-import scala.collection.JavaConversions._
-
-class AddPartitionsToTxnRequestTest extends BaseRequestTest {
-  private val topic1 = "foobartopic"
-  val numPartitions = 3
-
-  override def propertyOverrides(properties: Properties): Unit =
-    properties.put(KafkaConfig.AutoCreateTopicsEnableProp, false.toString)
-
-  @Before
-  override def setUp(): Unit = {
-    super.setUp()
-    TestUtils.createTopic(zkUtils, topic1, numPartitions, servers.size, servers, new Properties())
-  }
-
-  @Test
-  def shouldReceiveOperationNotAttemptedWhenOtherPartitionHasError(): Unit = {
-    // The basic idea is that we have one unknown topic and one created topic. We should get the 'UNKNOWN_TOPIC_OR_PARTITION'
-    // error for the unknown topic and the 'OPERATION_NOT_ATTEMPTED' error for the known and authorized topic.
-    val nonExistentTopic = new TopicPartition("unknownTopic", 0)
-    val createdTopicPartition = new TopicPartition(topic1, 0)
-
-    val request = createRequest(List(createdTopicPartition, nonExistentTopic))
-    val leaderId = servers.head.config.brokerId
-    val response = sendAddPartitionsRequest(leaderId, request)
-
-    assertEquals(2, response.errors.size)
-
-    assertTrue(response.errors.containsKey(createdTopicPartition))
-    assertEquals(Errors.OPERATION_NOT_ATTEMPTED, response.errors.get(createdTopicPartition))
-
-    assertTrue(response.errors.containsKey(nonExistentTopic))
-    assertEquals(Errors.UNKNOWN_TOPIC_OR_PARTITION, response.errors.get(nonExistentTopic))
-  }
-
-  private def sendAddPartitionsRequest(leaderId: Int, request: AddPartitionsToTxnRequest) : AddPartitionsToTxnResponse = {
-    val response = connectAndSend(request, ApiKeys.ADD_PARTITIONS_TO_TXN, destination = brokerSocketServer(leaderId))
-    AddPartitionsToTxnResponse.parse(response, request.version)
-  }
-
-  private def createRequest(partitions: List[TopicPartition]): AddPartitionsToTxnRequest = {
-    val transactionalId = "foobar"
-    val producerId = 1000L
-    val producerEpoch: Short = 0
-    val builder = new AddPartitionsToTxnRequest.Builder(transactionalId, producerId, producerEpoch, partitions)
-    builder.build()
-  }
-}
diff --git a/core/src/test/scala/unit/kafka/server/AdminManagerTest.scala b/core/src/test/scala/unit/kafka/server/AdminManagerTest.scala
new file mode 100644
index 0000000000000..e97162f640247
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/AdminManagerTest.scala
@@ -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 kafka.server
+
+import kafka.zk.KafkaZkClient
+import org.apache.kafka.common.metrics.Metrics
+import org.easymock.EasyMock
+import kafka.utils.TestUtils
+import org.apache.kafka.common.config.ConfigResource
+import org.apache.kafka.common.message.DescribeConfigsRequestData
+import org.apache.kafka.common.message.DescribeConfigsResponseData
+import org.apache.kafka.common.protocol.Errors
+
+import org.junit.{After, Test}
+import org.junit.Assert.assertEquals
+import org.junit.Assert.assertFalse
+import org.junit.Assert.assertNotNull
+import org.junit.Assert.assertNotEquals
+import java.util.Properties
+
+class AdminManagerTest {
+
+  private val zkClient: KafkaZkClient = EasyMock.createNiceMock(classOf[KafkaZkClient])
+  private val metrics = new Metrics()
+  private val brokerId = 1
+  private val topic = "topic-1"
+  private val metadataCache: MetadataCache = EasyMock.createNiceMock(classOf[MetadataCache])
+
+  @After
+  def tearDown(): Unit = {
+    metrics.close()
+  }
+
+  def createAdminManager(): AdminManager = {
+    val props = TestUtils.createBrokerConfig(brokerId, "zk")
+    new AdminManager(KafkaConfig.fromProps(props), metrics, metadataCache, zkClient)
+  }
+
+  @Test
+  def testDescribeConfigsWithNullConfigurationKeys(): Unit = {
+    EasyMock.expect(zkClient.getEntityConfigs(ConfigType.Topic, topic)).andReturn(TestUtils.createBrokerConfig(brokerId, "zk"))
+    EasyMock.expect(metadataCache.contains(topic)).andReturn(true)
+
+    EasyMock.replay(zkClient, metadataCache)
+
+    val resources = List(new DescribeConfigsRequestData.DescribeConfigsResource()
+      .setResourceName(topic)
+      .setResourceType(ConfigResource.Type.TOPIC.id)
+      .setConfigurationKeys(null))
+    val adminManager = createAdminManager()
+    val results: List[DescribeConfigsResponseData.DescribeConfigsResult] = adminManager.describeConfigs(resources, true, true)
+    assertEquals(Errors.NONE.code, results.head.errorCode())
+    assertFalse("Should return configs", results.head.configs().isEmpty)
+  }
+
+  @Test
+  def testDescribeConfigsWithEmptyConfigurationKeys(): Unit = {
+    EasyMock.expect(zkClient.getEntityConfigs(ConfigType.Topic, topic)).andReturn(TestUtils.createBrokerConfig(brokerId, "zk"))
+    EasyMock.expect(metadataCache.contains(topic)).andReturn(true)
+
+    EasyMock.replay(zkClient, metadataCache)
+
+    val resources = List(new DescribeConfigsRequestData.DescribeConfigsResource()
+      .setResourceName(topic)
+      .setResourceType(ConfigResource.Type.TOPIC.id))
+    val adminManager = createAdminManager()
+    val results: List[DescribeConfigsResponseData.DescribeConfigsResult] = adminManager.describeConfigs(resources, true, true)
+    assertEquals(Errors.NONE.code, results.head.errorCode())
+    assertFalse("Should return configs", results.head.configs().isEmpty)
+  }
+
+  @Test
+  def testDescribeConfigsWithDocumentation(): Unit = {
+    EasyMock.expect(zkClient.getEntityConfigs(ConfigType.Topic, topic)).andReturn(new Properties)
+    EasyMock.expect(zkClient.getEntityConfigs(ConfigType.Broker, brokerId.toString)).andReturn(new Properties)
+    EasyMock.expect(metadataCache.contains(topic)).andReturn(true)
+    EasyMock.replay(zkClient, metadataCache)
+
+    val adminManager = createAdminManager()
+
+    val resources = List(
+      new DescribeConfigsRequestData.DescribeConfigsResource()
+        .setResourceName(topic)
+        .setResourceType(ConfigResource.Type.TOPIC.id),
+      new DescribeConfigsRequestData.DescribeConfigsResource()
+        .setResourceName(brokerId.toString)
+        .setResourceType(ConfigResource.Type.BROKER.id))
+
+    val results: List[DescribeConfigsResponseData.DescribeConfigsResult] = adminManager.describeConfigs(resources, true, true)
+    assertEquals(2, results.size)
+    results.foreach(r => {
+      assertEquals(Errors.NONE.code, r.errorCode)
+      assertFalse("Should return configs", r.configs.isEmpty)
+      r.configs.forEach(c => {
+        assertNotNull(s"Config ${c.name} should have non null documentation", c.documentation)
+        assertNotEquals(s"Config ${c.name} should have non blank documentation", "", c.documentation.trim)
+      })
+    })
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/server/AdvertiseBrokerTest.scala b/core/src/test/scala/unit/kafka/server/AdvertiseBrokerTest.scala
index 1c3efee065c82..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()
   }
@@ -45,7 +45,7 @@ class AdvertiseBrokerTest extends ZooKeeperTestHarness {
     props.put("advertised.port", advertisedPort.toString)
     servers += TestUtils.createServer(KafkaConfig.fromProps(props))
 
-    val brokerInfo = zkUtils.getBrokerInfo(brokerId).get
+    val brokerInfo = zkClient.getBroker(brokerId).get
     assertEquals(1, brokerInfo.endPoints.size)
     val endpoint = brokerInfo.endPoints.head
     assertEquals(advertisedHostName, endpoint.host)
@@ -59,7 +59,7 @@ class AdvertiseBrokerTest extends ZooKeeperTestHarness {
     props.put("advertised.listeners", "PLAINTEXT://routable-listener:3334")
     servers += TestUtils.createServer(KafkaConfig.fromProps(props))
 
-    val brokerInfo = zkUtils.getBrokerInfo(brokerId).get
+    val brokerInfo = zkClient.getBroker(brokerId).get
     assertEquals(1, brokerInfo.endPoints.size)
     val endpoint = brokerInfo.endPoints.head
     assertEquals("routable-listener", endpoint.host)
@@ -76,7 +76,7 @@ class AdvertiseBrokerTest extends ZooKeeperTestHarness {
     props.put("inter.broker.listener.name", "INTERNAL")
     servers += TestUtils.createServer(KafkaConfig.fromProps(props))
 
-    val brokerInfo = zkUtils.getBrokerInfo(brokerId).get
+    val brokerInfo = zkClient.getBroker(brokerId).get
     assertEquals(1, brokerInfo.endPoints.size)
     val endpoint = brokerInfo.endPoints.head
     assertEquals("external-listener", endpoint.host)
diff --git a/core/src/test/scala/unit/kafka/server/AlterIsrManagerTest.scala b/core/src/test/scala/unit/kafka/server/AlterIsrManagerTest.scala
new file mode 100644
index 0000000000000..1f451fef090c1
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/AlterIsrManagerTest.scala
@@ -0,0 +1,303 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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.server
+
+import java.util.Collections
+import java.util.concurrent.atomic.AtomicInteger
+
+import kafka.api.LeaderAndIsr
+import kafka.server.{AlterIsrItem, AlterIsrManager, AlterIsrManagerImpl, BrokerToControllerChannelManager, ControllerRequestCompletionHandler}
+import kafka.utils.{MockScheduler, MockTime}
+import org.apache.kafka.clients.ClientResponse
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.message.AlterIsrResponseData
+import org.apache.kafka.common.metrics.Metrics
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.requests.{AbstractRequest, AlterIsrRequest, AlterIsrResponse}
+import org.easymock.EasyMock
+import org.junit.Assert._
+import org.junit.{Before, Test}
+
+
+class AlterIsrManagerTest {
+
+  val topic = "test-topic"
+  val time = new MockTime
+  val metrics = new Metrics
+  val brokerId = 1
+  val requestTimeout = Long.MaxValue
+
+  var brokerToController: BrokerToControllerChannelManager = _
+
+  val tp0 = new TopicPartition(topic, 0)
+  val tp1 = new TopicPartition(topic, 1)
+  val tp2 = new TopicPartition(topic, 2)
+
+  @Before
+  def setup(): Unit = {
+    brokerToController = EasyMock.createMock(classOf[BrokerToControllerChannelManager])
+  }
+
+  @Test
+  def testBasic(): Unit = {
+    EasyMock.expect(brokerToController.sendRequest(EasyMock.anyObject(), EasyMock.anyObject(), EasyMock.eq(requestTimeout))).once()
+    EasyMock.replay(brokerToController)
+
+    val scheduler = new MockScheduler(time)
+    val alterIsrManager = new AlterIsrManagerImpl(brokerToController, scheduler, time, brokerId, () => 2)
+    alterIsrManager.start()
+    alterIsrManager.enqueue(AlterIsrItem(tp0, new LeaderAndIsr(1, 1, List(1,2,3), 10), _ => {}))
+    time.sleep(50)
+    scheduler.tick()
+
+    EasyMock.verify(brokerToController)
+  }
+
+  @Test
+  def testOverwriteWithinBatch(): Unit = {
+    val capture = EasyMock.newCapture[AbstractRequest.Builder[AlterIsrRequest]]()
+    EasyMock.expect(brokerToController.sendRequest(EasyMock.capture(capture), EasyMock.anyObject(), EasyMock.eq(requestTimeout))).once()
+    EasyMock.replay(brokerToController)
+
+    val scheduler = new MockScheduler(time)
+    val alterIsrManager = new AlterIsrManagerImpl(brokerToController, scheduler, time, brokerId, () => 2)
+    alterIsrManager.start()
+
+    // Only send one ISR update for a given topic+partition
+    assertTrue(alterIsrManager.enqueue(AlterIsrItem(tp0, new LeaderAndIsr(1, 1, List(1,2,3), 10), _ => {})))
+    assertFalse(alterIsrManager.enqueue(AlterIsrItem(tp0, new LeaderAndIsr(1, 1, List(1,2), 10), _ => {})))
+
+    time.sleep(50)
+    scheduler.tick()
+
+    EasyMock.verify(brokerToController)
+
+    val request = capture.getValue.build()
+    assertEquals(request.data().topics().size(), 1)
+    assertEquals(request.data().topics().get(0).partitions().get(0).newIsr().size(), 3)
+  }
+
+  @Test
+  def testSingleBatch(): Unit = {
+    val capture = EasyMock.newCapture[AbstractRequest.Builder[AlterIsrRequest]]()
+    EasyMock.expect(brokerToController.sendRequest(EasyMock.capture(capture), EasyMock.anyObject(), EasyMock.eq(requestTimeout))).once()
+    EasyMock.replay(brokerToController)
+
+    val scheduler = new MockScheduler(time)
+    val alterIsrManager = new AlterIsrManagerImpl(brokerToController, scheduler, time, brokerId, () => 2)
+    alterIsrManager.start()
+
+    for (i <- 0 to 9) {
+      alterIsrManager.enqueue(AlterIsrItem(new TopicPartition(topic, i),
+        new LeaderAndIsr(1, 1, List(1,2,3), 10), _ => {}))
+      time.sleep(1)
+    }
+
+    time.sleep(50)
+    scheduler.tick()
+
+    // This should not be included in the batch
+    alterIsrManager.enqueue(AlterIsrItem(new TopicPartition(topic, 10),
+      new LeaderAndIsr(1, 1, List(1,2,3), 10), _ => {}))
+
+    EasyMock.verify(brokerToController)
+
+    val request = capture.getValue.build()
+    assertEquals(request.data().topics().size(), 1)
+    assertEquals(request.data().topics().get(0).partitions().size(), 10)
+  }
+
+  @Test
+  def testAuthorizationFailed(): Unit = {
+    val isrs = Seq(AlterIsrItem(tp0, new LeaderAndIsr(1, 1, List(1,2,3), 10), _ => { }))
+    val manager = testTopLevelError(isrs, Errors.CLUSTER_AUTHORIZATION_FAILED)
+    // On authz error, we log the exception and keep retrying
+    assertFalse(manager.enqueue(AlterIsrItem(tp0, null, _ => { })))
+  }
+
+  @Test
+  def testStaleBrokerEpoch(): Unit = {
+    val isrs = Seq(AlterIsrItem(tp0, new LeaderAndIsr(1, 1, List(1,2,3), 10), _ => { }))
+    val manager = testTopLevelError(isrs, Errors.STALE_BROKER_EPOCH)
+    // On stale broker epoch, we want to retry, so we don't clear items from the pending map
+    assertFalse(manager.enqueue(AlterIsrItem(tp0, null, _ => { })))
+  }
+
+  @Test
+  def testOtherErrors(): Unit = {
+    val isrs = Seq(AlterIsrItem(tp0, new LeaderAndIsr(1, 1, List(1,2,3), 10), _ => { }))
+    val manager = testTopLevelError(isrs, Errors.UNKNOWN_SERVER_ERROR)
+    // On other unexpected errors, we also want to retry
+    assertFalse(manager.enqueue(AlterIsrItem(tp0, null, _ => { })))
+  }
+
+  def testTopLevelError(isrs: Seq[AlterIsrItem], error: Errors): AlterIsrManager = {
+    val callbackCapture = EasyMock.newCapture[ControllerRequestCompletionHandler]()
+
+    EasyMock.expect(brokerToController.sendRequest(EasyMock.anyObject(), EasyMock.capture(callbackCapture), EasyMock.eq(requestTimeout))).once()
+    EasyMock.replay(brokerToController)
+
+    val scheduler = new MockScheduler(time)
+    val alterIsrManager = new AlterIsrManagerImpl(brokerToController, scheduler, time, brokerId, () => 2)
+    alterIsrManager.start()
+    isrs.foreach(alterIsrManager.enqueue)
+
+    time.sleep(100)
+    scheduler.tick()
+
+    EasyMock.verify(brokerToController)
+
+    val alterIsrResp = new AlterIsrResponse(new AlterIsrResponseData().setErrorCode(error.code))
+    val resp = new ClientResponse(null, null, "", 0L, 0L,
+      false, null, null, alterIsrResp)
+    callbackCapture.getValue.onComplete(resp)
+    alterIsrManager
+  }
+
+  @Test
+  def testPartitionErrors(): Unit = {
+    val errors = Seq(Errors.INVALID_UPDATE_VERSION, Errors.UNKNOWN_TOPIC_OR_PARTITION, Errors.NOT_LEADER_OR_FOLLOWER)
+    errors.foreach(error => {
+      val alterIsrManager = testPartitionError(tp0, error)
+      // Any partition-level error should clear the item from the pending queue allowing for future updates
+      assertTrue(alterIsrManager.enqueue(AlterIsrItem(tp0, null, _ => { })))
+    })
+  }
+
+  def testPartitionError(tp: TopicPartition, error: Errors): AlterIsrManager = {
+    val callbackCapture = EasyMock.newCapture[ControllerRequestCompletionHandler]()
+    EasyMock.reset(brokerToController)
+    EasyMock.expect(brokerToController.sendRequest(EasyMock.anyObject(), EasyMock.capture(callbackCapture), EasyMock.eq(requestTimeout))).once()
+    EasyMock.replay(brokerToController)
+
+    val scheduler = new MockScheduler(time)
+    val alterIsrManager = new AlterIsrManagerImpl(brokerToController, scheduler, time, brokerId, () => 2)
+    alterIsrManager.start()
+
+    var capturedError: Option[Errors] = None
+    val callback = (result:  Either[Errors, LeaderAndIsr]) => {
+      result match {
+        case Left(error: Errors) => capturedError = Some(error)
+        case Right(_) => fail("Should have seen error")
+      }
+    }
+
+    alterIsrManager.enqueue(AlterIsrItem(tp, new LeaderAndIsr(1, 1, List(1,2,3), 10), callback))
+
+    time.sleep(100)
+    scheduler.tick()
+
+    EasyMock.verify(brokerToController)
+
+    val alterIsrResp = new AlterIsrResponse(new AlterIsrResponseData()
+      .setTopics(Collections.singletonList(
+        new AlterIsrResponseData.TopicData()
+          .setName(tp.topic())
+          .setPartitions(Collections.singletonList(
+            new AlterIsrResponseData.PartitionData()
+              .setPartitionIndex(tp.partition())
+              .setErrorCode(error.code))))))
+    val resp = new ClientResponse(null, null, "", 0L, 0L,
+      false, null, null, alterIsrResp)
+    callbackCapture.getValue.onComplete(resp)
+    assertTrue(capturedError.isDefined)
+    assertEquals(capturedError.get, error)
+    alterIsrManager
+  }
+
+  @Test
+  def testOneInFlight(): Unit = {
+    val callbackCapture = EasyMock.newCapture[ControllerRequestCompletionHandler]()
+    EasyMock.reset(brokerToController)
+    EasyMock.expect(brokerToController.sendRequest(EasyMock.anyObject(), EasyMock.capture(callbackCapture), EasyMock.eq(requestTimeout))).once()
+    EasyMock.replay(brokerToController)
+
+    val scheduler = new MockScheduler(time)
+    val alterIsrManager = new AlterIsrManagerImpl(brokerToController, scheduler, time, brokerId, () => 2)
+    alterIsrManager.start()
+    alterIsrManager.enqueue(AlterIsrItem(tp0, new LeaderAndIsr(1, 1, List(1,2,3), 10), _ => {}))
+
+    time.sleep(100)
+    scheduler.tick() // Triggers a request
+
+    // Enqueue more updates
+    alterIsrManager.enqueue(AlterIsrItem(tp1, new LeaderAndIsr(1, 1, List(1,2,3), 10), _ => {}))
+    alterIsrManager.enqueue(AlterIsrItem(tp2, new LeaderAndIsr(1, 1, List(1,2,3), 10), _ => {}))
+
+    time.sleep(100)
+    scheduler.tick() // Trigger the schedule again, but no request this time
+
+    EasyMock.verify(brokerToController)
+
+    // Even an empty response will clear the in-flight
+    val alterIsrResp = new AlterIsrResponse(new AlterIsrResponseData())
+    val resp = new ClientResponse(null, null, "", 0L, 0L,
+      false, null, null, alterIsrResp)
+    callbackCapture.getValue.onComplete(resp)
+
+    EasyMock.reset(brokerToController)
+    EasyMock.expect(brokerToController.sendRequest(EasyMock.anyObject(), EasyMock.capture(callbackCapture), EasyMock.eq(requestTimeout))).once()
+    EasyMock.replay(brokerToController)
+
+    time.sleep(100)
+    scheduler.tick()
+    EasyMock.verify(brokerToController)
+  }
+
+  @Test
+  def testPartitionMissingInResponse(): Unit = {
+    val callbackCapture = EasyMock.newCapture[ControllerRequestCompletionHandler]()
+    EasyMock.reset(brokerToController)
+    EasyMock.expect(brokerToController.sendRequest(EasyMock.anyObject(), EasyMock.capture(callbackCapture), EasyMock.eq(requestTimeout))).once()
+    EasyMock.replay(brokerToController)
+
+    val scheduler = new MockScheduler(time)
+    val alterIsrManager = new AlterIsrManagerImpl(brokerToController, scheduler, time, brokerId, () => 2)
+    alterIsrManager.start()
+
+    val count = new AtomicInteger(0)
+    val callback = (result:  Either[Errors, LeaderAndIsr]) => {
+      count.incrementAndGet()
+      return
+    }
+    alterIsrManager.enqueue(AlterIsrItem(tp0, new LeaderAndIsr(1, 1, List(1,2,3), 10), callback))
+    alterIsrManager.enqueue(AlterIsrItem(tp1, new LeaderAndIsr(1, 1, List(1,2,3), 10), callback))
+    alterIsrManager.enqueue(AlterIsrItem(tp2, new LeaderAndIsr(1, 1, List(1,2,3), 10), callback))
+
+
+    time.sleep(100)
+    scheduler.tick()
+
+    EasyMock.verify(brokerToController)
+
+    // Three partitions were sent, but only one returned
+    val alterIsrResp = new AlterIsrResponse(new AlterIsrResponseData()
+      .setTopics(Collections.singletonList(
+        new AlterIsrResponseData.TopicData()
+          .setName(tp0.topic())
+          .setPartitions(Collections.singletonList(
+            new AlterIsrResponseData.PartitionData()
+              .setPartitionIndex(tp0.partition())
+              .setErrorCode(Errors.UNKNOWN_SERVER_ERROR.code()))))))
+    val resp = new ClientResponse(null, null, "", 0L, 0L,
+      false, null, null, alterIsrResp)
+    callbackCapture.getValue.onComplete(resp)
+
+    assertEquals("Expected all callbacks to run", count.get, 3)
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/server/AlterReplicaLogDirsRequestTest.scala b/core/src/test/scala/unit/kafka/server/AlterReplicaLogDirsRequestTest.scala
index f2dc4a50e31c3..425456f9b6cb4 100644
--- a/core/src/test/scala/unit/kafka/server/AlterReplicaLogDirsRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/AlterReplicaLogDirsRequestTest.scala
@@ -17,29 +17,34 @@
 
 package kafka.server
 
-import kafka.network.SocketServer
-import kafka.utils._
 import java.io.File
 
+import kafka.utils._
 import org.apache.kafka.common.TopicPartition
-import org.apache.kafka.common.protocol.{ApiKeys, Errors}
+import org.apache.kafka.common.message.AlterReplicaLogDirsRequestData
+import org.apache.kafka.common.protocol.Errors
 import org.apache.kafka.common.requests.{AlterReplicaLogDirsRequest, AlterReplicaLogDirsResponse}
 import org.junit.Assert._
 import org.junit.Test
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 import scala.collection.mutable
 import scala.util.Random
 
 class AlterReplicaLogDirsRequestTest extends BaseRequestTest {
-
-  override def numBrokers: Int = 1
-  override def logDirCount: Int = 5
+  override val logDirCount = 5
+  override val brokerCount = 1
 
   val topic = "topic"
 
+  private def findErrorForPartition(response: AlterReplicaLogDirsResponse, tp: TopicPartition): Errors = {
+    Errors.forCode(response.data.results.asScala
+      .find(x => x.topicName == tp.topic).get.partitions.asScala
+      .find(p => p.partitionIndex == tp.partition).get.errorCode)
+  }
+
   @Test
-  def testAlterReplicaLogDirsRequest() {
+  def testAlterReplicaLogDirsRequest(): Unit = {
     val partitionNum = 5
 
     // Alter replica dir before topic creation
@@ -47,14 +52,14 @@ class AlterReplicaLogDirsRequestTest extends BaseRequestTest {
     val partitionDirs1 = (0 until partitionNum).map(partition => new TopicPartition(topic, partition) -> logDir1).toMap
     val alterReplicaLogDirsResponse1 = sendAlterReplicaLogDirsRequest(partitionDirs1)
 
-    // The response should show error REPLICA_NOT_AVAILABLE for all partitions
+    // The response should show error UNKNOWN_TOPIC_OR_PARTITION for all partitions
     (0 until partitionNum).foreach { partition =>
       val tp = new TopicPartition(topic, partition)
-      assertEquals(Errors.REPLICA_NOT_AVAILABLE, alterReplicaLogDirsResponse1.responses().get(tp))
+      assertEquals(Errors.UNKNOWN_TOPIC_OR_PARTITION, findErrorForPartition(alterReplicaLogDirsResponse1, tp))
       assertTrue(servers.head.logManager.getLog(tp).isEmpty)
     }
 
-    TestUtils.createTopic(zkUtils, topic, partitionNum, 1, servers)
+    createTopic(topic, partitionNum, 1)
     (0 until partitionNum).foreach { partition =>
       assertEquals(logDir1, servers.head.logManager.getLog(new TopicPartition(topic, partition)).get.dir.getParent)
     }
@@ -66,7 +71,7 @@ class AlterReplicaLogDirsRequestTest extends BaseRequestTest {
     // The response should succeed for all partitions
     (0 until partitionNum).foreach { partition =>
       val tp = new TopicPartition(topic, partition)
-      assertEquals(Errors.NONE, alterReplicaLogDirsResponse2.responses().get(tp))
+      assertEquals(Errors.NONE, findErrorForPartition(alterReplicaLogDirsResponse2, tp))
       TestUtils.waitUntilTrue(() => {
         logDir2 == servers.head.logManager.getLog(new TopicPartition(topic, partition)).get.dir.getParent
       }, "timed out waiting for replica movement")
@@ -85,18 +90,18 @@ class AlterReplicaLogDirsRequestTest extends BaseRequestTest {
     partitionDirs1.put(new TopicPartition(topic, 0), "invalidDir")
     partitionDirs1.put(new TopicPartition(topic, 1), validDir1)
     val alterReplicaDirResponse1 = sendAlterReplicaLogDirsRequest(partitionDirs1.toMap)
-    assertEquals(Errors.LOG_DIR_NOT_FOUND, alterReplicaDirResponse1.responses().get(new TopicPartition(topic, 0)))
-    assertEquals(Errors.REPLICA_NOT_AVAILABLE, alterReplicaDirResponse1.responses().get(new TopicPartition(topic, 1)))
+    assertEquals(Errors.LOG_DIR_NOT_FOUND, findErrorForPartition(alterReplicaDirResponse1, new TopicPartition(topic, 0)))
+    assertEquals(Errors.UNKNOWN_TOPIC_OR_PARTITION, findErrorForPartition(alterReplicaDirResponse1, new TopicPartition(topic, 1)))
 
-    TestUtils.createTopic(zkUtils, topic, 3, 1, servers)
+    createTopic(topic, 3, 1)
 
     // Test AlterReplicaDirRequest after topic creation
     val partitionDirs2 = mutable.Map.empty[TopicPartition, String]
     partitionDirs2.put(new TopicPartition(topic, 0), "invalidDir")
     partitionDirs2.put(new TopicPartition(topic, 1), validDir2)
     val alterReplicaDirResponse2 = sendAlterReplicaLogDirsRequest(partitionDirs2.toMap)
-    assertEquals(Errors.LOG_DIR_NOT_FOUND, alterReplicaDirResponse2.responses().get(new TopicPartition(topic, 0)))
-    assertEquals(Errors.NONE, alterReplicaDirResponse2.responses().get(new TopicPartition(topic, 1)))
+    assertEquals(Errors.LOG_DIR_NOT_FOUND, findErrorForPartition(alterReplicaDirResponse2, new TopicPartition(topic, 0)))
+    assertEquals(Errors.NONE, findErrorForPartition(alterReplicaDirResponse2, new TopicPartition(topic, 1)))
 
     // Test AlterReplicaDirRequest after topic creation and log directory failure
     servers.head.logDirFailureChannel.maybeAddOfflineLogDir(offlineDir, "", new java.io.IOException())
@@ -106,15 +111,27 @@ class AlterReplicaLogDirsRequestTest extends BaseRequestTest {
     partitionDirs3.put(new TopicPartition(topic, 1), validDir3)
     partitionDirs3.put(new TopicPartition(topic, 2), offlineDir)
     val alterReplicaDirResponse3 = sendAlterReplicaLogDirsRequest(partitionDirs3.toMap)
-    assertEquals(Errors.LOG_DIR_NOT_FOUND, alterReplicaDirResponse3.responses().get(new TopicPartition(topic, 0)))
-    assertEquals(Errors.KAFKA_STORAGE_ERROR, alterReplicaDirResponse3.responses().get(new TopicPartition(topic, 1)))
-    assertEquals(Errors.KAFKA_STORAGE_ERROR, alterReplicaDirResponse3.responses().get(new TopicPartition(topic, 2)))
+    assertEquals(Errors.LOG_DIR_NOT_FOUND, findErrorForPartition(alterReplicaDirResponse3, new TopicPartition(topic, 0)))
+    assertEquals(Errors.KAFKA_STORAGE_ERROR, findErrorForPartition(alterReplicaDirResponse3, new TopicPartition(topic, 1)))
+    assertEquals(Errors.KAFKA_STORAGE_ERROR, findErrorForPartition(alterReplicaDirResponse3, new TopicPartition(topic, 2)))
   }
 
-  private def sendAlterReplicaLogDirsRequest(partitionDirs: Map[TopicPartition, String], socketServer: SocketServer = controllerSocketServer): AlterReplicaLogDirsResponse = {
-    val request = new AlterReplicaLogDirsRequest.Builder(partitionDirs.asJava).build()
-    val response = connectAndSend(request, ApiKeys.ALTER_REPLICA_LOG_DIRS, socketServer)
-    AlterReplicaLogDirsResponse.parse(response, request.version)
+  private def sendAlterReplicaLogDirsRequest(partitionDirs: Map[TopicPartition, String]): AlterReplicaLogDirsResponse = {
+    val logDirs = partitionDirs.groupBy{case (_, dir) => dir}.map{ case(dir, tps) =>
+      new AlterReplicaLogDirsRequestData.AlterReplicaLogDir()
+        .setPath(dir)
+        .setTopics(new AlterReplicaLogDirsRequestData.AlterReplicaLogDirTopicCollection(
+          tps.groupBy { case (tp, _) => tp.topic }
+            .map { case (topic, tpPartitions) =>
+              new AlterReplicaLogDirsRequestData.AlterReplicaLogDirTopic()
+                .setName(topic)
+                .setPartitions(tpPartitions.map{case (tp, _) => tp.partition.asInstanceOf[Integer]}.toList.asJava)
+        }.toList.asJava.iterator))
+    }
+    val data = new AlterReplicaLogDirsRequestData()
+        .setDirs(new AlterReplicaLogDirsRequestData.AlterReplicaLogDirCollection(logDirs.asJava.iterator))
+    val request = new AlterReplicaLogDirsRequest.Builder(data).build()
+    connectAndReceive[AlterReplicaLogDirsResponse](request, destination = controllerSocketServer)
   }
 
 }
diff --git a/core/src/test/scala/unit/kafka/server/AlterUserScramCredentialsRequestTest.scala b/core/src/test/scala/unit/kafka/server/AlterUserScramCredentialsRequestTest.scala
new file mode 100644
index 0000000000000..c6732ba11523f
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/AlterUserScramCredentialsRequestTest.scala
@@ -0,0 +1,433 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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.nio.charset.StandardCharsets
+import java.util
+import java.util.Properties
+
+import kafka.network.SocketServer
+import kafka.security.authorizer.AclAuthorizer
+import org.apache.kafka.clients.admin.ScramMechanism
+import org.apache.kafka.common.message.AlterUserScramCredentialsResponseData.AlterUserScramCredentialsResult
+import org.apache.kafka.common.message.DescribeUserScramCredentialsResponseData.DescribeUserScramCredentialsResult
+import org.apache.kafka.common.message.{AlterUserScramCredentialsRequestData, DescribeUserScramCredentialsRequestData}
+import org.apache.kafka.common.protocol.{ApiKeys, Errors}
+import org.apache.kafka.common.requests.{AlterUserScramCredentialsRequest, AlterUserScramCredentialsResponse, DescribeUserScramCredentialsRequest, DescribeUserScramCredentialsResponse}
+import org.apache.kafka.common.security.auth.{AuthenticationContext, KafkaPrincipal, KafkaPrincipalBuilder}
+import org.apache.kafka.server.authorizer.{Action, AuthorizableRequestContext, AuthorizationResult}
+import org.junit.Assert._
+import org.junit.rules.TestName
+import org.junit.{Rule, Test}
+
+import scala.jdk.CollectionConverters._
+
+/**
+ * Test AlterUserScramCredentialsRequest/Response API for the cases where either no credentials are altered
+ * or failure is expected due to lack of authorization, sending the request to a non-controller broker, or some other issue.
+ * Also tests the Alter and Describe APIs for the case where credentials are successfully altered/described.
+ */
+class AlterUserScramCredentialsRequestTest extends BaseRequestTest {
+  override def brokerPropertyOverrides(properties: Properties): Unit = {
+    properties.put(KafkaConfig.ControlledShutdownEnableProp, "false")
+    properties.put(KafkaConfig.AuthorizerClassNameProp, classOf[AlterCredentialsTest.TestAuthorizer].getName)
+    properties.put(KafkaConfig.PrincipalBuilderClassProp,
+      if (testName.getMethodName.endsWith("NotAuthorized")) {
+        classOf[AlterCredentialsTest.TestPrincipalBuilderReturningUnauthorized].getName
+      } else {
+        classOf[AlterCredentialsTest.TestPrincipalBuilderReturningAuthorized].getName
+      })
+  }
+
+  private val _testName = new TestName
+  @Rule def testName = _testName
+
+  private val saltedPasswordBytes = "saltedPassword".getBytes(StandardCharsets.UTF_8)
+  private val saltBytes = "salt".getBytes(StandardCharsets.UTF_8)
+  private val user1 = "user1"
+  private val user2 = "user2"
+  private val unknownUser = "unknownUser"
+
+  @Test
+  def testAlterNothing(): Unit = {
+    val request = new AlterUserScramCredentialsRequest.Builder(
+      new AlterUserScramCredentialsRequestData()
+        .setDeletions(new util.ArrayList[AlterUserScramCredentialsRequestData.ScramCredentialDeletion])
+        .setUpsertions(new util.ArrayList[AlterUserScramCredentialsRequestData.ScramCredentialUpsertion])).build()
+    val response = sendAlterUserScramCredentialsRequest(request)
+
+    val results = response.data.results
+    assertEquals(0, results.size)
+  }
+
+  @Test
+  def testAlterNothingNotAuthorized(): Unit = {
+    val request = new AlterUserScramCredentialsRequest.Builder(
+      new AlterUserScramCredentialsRequestData()
+        .setDeletions(new util.ArrayList[AlterUserScramCredentialsRequestData.ScramCredentialDeletion])
+        .setUpsertions(new util.ArrayList[AlterUserScramCredentialsRequestData.ScramCredentialUpsertion])).build()
+    val response = sendAlterUserScramCredentialsRequest(request)
+
+    val results = response.data.results
+    assertEquals(0, results.size)
+  }
+
+  @Test
+  def testAlterSomethingNotAuthorized(): Unit = {
+
+    val request = new AlterUserScramCredentialsRequest.Builder(
+      new AlterUserScramCredentialsRequestData()
+        .setDeletions(util.Arrays.asList(new AlterUserScramCredentialsRequestData.ScramCredentialDeletion().setName(user1).setMechanism(ScramMechanism.SCRAM_SHA_256.`type`)))
+        .setUpsertions(util.Arrays.asList(new AlterUserScramCredentialsRequestData.ScramCredentialUpsertion().setName(user2).setMechanism(ScramMechanism.SCRAM_SHA_512.`type`)))).build()
+    val response = sendAlterUserScramCredentialsRequest(request)
+
+    val results = response.data.results
+    assertEquals(2, results.size)
+    checkAllErrorsAlteringCredentials(results, Errors.CLUSTER_AUTHORIZATION_FAILED, "when not authorized")
+  }
+
+  @Test
+  def testAlterSameThingTwice(): Unit = {
+    val deletion1 = new AlterUserScramCredentialsRequestData.ScramCredentialDeletion().setName(user1).setMechanism(ScramMechanism.SCRAM_SHA_256.`type`)
+    val deletion2 = new AlterUserScramCredentialsRequestData.ScramCredentialDeletion().setName(user2).setMechanism(ScramMechanism.SCRAM_SHA_256.`type`)
+    val upsertion1 = new AlterUserScramCredentialsRequestData.ScramCredentialUpsertion().setName(user1).setMechanism(ScramMechanism.SCRAM_SHA_256.`type`)
+      .setIterations(4096).setSalt(saltBytes).setSaltedPassword(saltedPasswordBytes)
+    val upsertion2 = new AlterUserScramCredentialsRequestData.ScramCredentialUpsertion().setName(user2).setMechanism(ScramMechanism.SCRAM_SHA_256.`type`)
+      .setIterations(4096).setSalt(saltBytes).setSaltedPassword(saltedPasswordBytes)
+    val requests = List (
+      new AlterUserScramCredentialsRequest.Builder(
+        new AlterUserScramCredentialsRequestData()
+          .setDeletions(util.Arrays.asList(deletion1, deletion1))
+          .setUpsertions(util.Arrays.asList(upsertion2, upsertion2))).build(),
+      new AlterUserScramCredentialsRequest.Builder(
+        new AlterUserScramCredentialsRequestData()
+          .setDeletions(util.Arrays.asList(deletion1, deletion2))
+          .setUpsertions(util.Arrays.asList(upsertion1, upsertion2))).build(),
+    )
+    requests.foreach(request => {
+      val response = sendAlterUserScramCredentialsRequest(request)
+      val results = response.data.results
+      assertEquals(2, results.size)
+      checkAllErrorsAlteringCredentials(results, Errors.DUPLICATE_RESOURCE, "when altering the same credential twice in a single request")
+    })
+  }
+
+  @Test
+  def testAlterEmptyUser(): Unit = {
+    val deletionEmpty = new AlterUserScramCredentialsRequestData.ScramCredentialDeletion().setName("").setMechanism(ScramMechanism.SCRAM_SHA_256.`type`)
+    val upsertionEmpty = new AlterUserScramCredentialsRequestData.ScramCredentialUpsertion().setName("").setMechanism(ScramMechanism.SCRAM_SHA_256.`type`)
+      .setIterations(4096).setSalt(saltBytes).setSaltedPassword(saltedPasswordBytes)
+    val requests = List (
+      new AlterUserScramCredentialsRequest.Builder(
+        new AlterUserScramCredentialsRequestData()
+          .setDeletions(util.Arrays.asList(deletionEmpty))
+          .setUpsertions(new util.ArrayList[AlterUserScramCredentialsRequestData.ScramCredentialUpsertion])).build(),
+      new AlterUserScramCredentialsRequest.Builder(
+        new AlterUserScramCredentialsRequestData()
+          .setDeletions(new util.ArrayList[AlterUserScramCredentialsRequestData.ScramCredentialDeletion])
+          .setUpsertions(util.Arrays.asList(upsertionEmpty))).build(),
+      new AlterUserScramCredentialsRequest.Builder(
+        new AlterUserScramCredentialsRequestData()
+          .setDeletions(util.Arrays.asList(deletionEmpty, deletionEmpty))
+          .setUpsertions(util.Arrays.asList(upsertionEmpty))).build(),
+    )
+    requests.foreach(request => {
+      val response = sendAlterUserScramCredentialsRequest(request)
+      val results = response.data.results
+      assertEquals(1, results.size)
+      checkAllErrorsAlteringCredentials(results, Errors.UNACCEPTABLE_CREDENTIAL, "when altering an empty user")
+      assertEquals("Username must not be empty", results.get(0).errorMessage)
+    })
+  }
+
+  @Test
+  def testAlterUnknownMechanism(): Unit = {
+    val deletionUnknown1 = new AlterUserScramCredentialsRequestData.ScramCredentialDeletion().setName(user1).setMechanism(ScramMechanism.UNKNOWN.`type`)
+    val deletionValid1 = new AlterUserScramCredentialsRequestData.ScramCredentialDeletion().setName(user1).setMechanism(ScramMechanism.SCRAM_SHA_256.`type`)
+    val deletionUnknown2 = new AlterUserScramCredentialsRequestData.ScramCredentialDeletion().setName(user2).setMechanism(10.toByte)
+    val user3 = "user3"
+    val upsertionUnknown3 = new AlterUserScramCredentialsRequestData.ScramCredentialUpsertion().setName(user3).setMechanism(ScramMechanism.UNKNOWN.`type`)
+      .setIterations(8192).setSalt(saltBytes).setSaltedPassword(saltedPasswordBytes)
+    val upsertionValid3 = new AlterUserScramCredentialsRequestData.ScramCredentialUpsertion().setName(user3).setMechanism(ScramMechanism.SCRAM_SHA_256.`type`)
+      .setIterations(8192).setSalt(saltBytes).setSaltedPassword(saltedPasswordBytes)
+    val user4 = "user4"
+    val upsertionUnknown4 = new AlterUserScramCredentialsRequestData.ScramCredentialUpsertion().setName(user4).setMechanism(10.toByte)
+      .setIterations(8192).setSalt(saltBytes).setSaltedPassword(saltedPasswordBytes)
+    val user5 = "user5"
+    val upsertionUnknown5 = new AlterUserScramCredentialsRequestData.ScramCredentialUpsertion().setName(user5).setMechanism(ScramMechanism.UNKNOWN.`type`)
+      .setIterations(8192).setSalt(saltBytes).setSaltedPassword(saltedPasswordBytes)
+    val request = new AlterUserScramCredentialsRequest.Builder(
+        new AlterUserScramCredentialsRequestData()
+          .setDeletions(util.Arrays.asList(deletionUnknown1, deletionValid1, deletionUnknown2))
+          .setUpsertions(util.Arrays.asList(upsertionUnknown3, upsertionValid3, upsertionUnknown4, upsertionUnknown5))).build()
+    val response = sendAlterUserScramCredentialsRequest(request)
+    val results = response.data.results
+    assertEquals(5, results.size)
+    checkAllErrorsAlteringCredentials(results, Errors.UNSUPPORTED_SASL_MECHANISM, "when altering the credentials with unknown SCRAM mechanisms")
+    results.asScala.foreach(result => assertEquals("Unknown SCRAM mechanism", result.errorMessage))
+  }
+
+  @Test
+  def testAlterTooFewIterations(): Unit = {
+    val upsertionTooFewIterations = new AlterUserScramCredentialsRequestData.ScramCredentialUpsertion().setName(user1)
+      .setMechanism(ScramMechanism.SCRAM_SHA_256.`type`).setIterations(1)
+      .setSalt(saltBytes).setSaltedPassword(saltedPasswordBytes)
+    val request = new AlterUserScramCredentialsRequest.Builder(
+      new AlterUserScramCredentialsRequestData()
+        .setDeletions(util.Collections.emptyList())
+        .setUpsertions(util.Arrays.asList(upsertionTooFewIterations))).build()
+    val response = sendAlterUserScramCredentialsRequest(request)
+    val results = response.data.results
+    assertEquals(1, results.size)
+    checkAllErrorsAlteringCredentials(results, Errors.UNACCEPTABLE_CREDENTIAL, "when altering the credentials with too few iterations")
+    assertEquals("Too few iterations", results.get(0).errorMessage)
+  }
+
+  @Test
+  def testAlterTooManyIterations(): Unit = {
+    val upsertionTooFewIterations = new AlterUserScramCredentialsRequestData.ScramCredentialUpsertion().setName(user1)
+      .setMechanism(ScramMechanism.SCRAM_SHA_256.`type`).setIterations(Integer.MAX_VALUE)
+      .setSalt(saltBytes).setSaltedPassword(saltedPasswordBytes)
+    val request = new AlterUserScramCredentialsRequest.Builder(
+      new AlterUserScramCredentialsRequestData()
+        .setDeletions(util.Collections.emptyList())
+        .setUpsertions(util.Arrays.asList(upsertionTooFewIterations))).build()
+    val response = sendAlterUserScramCredentialsRequest(request)
+    val results = response.data.results
+    assertEquals(1, results.size)
+    checkAllErrorsAlteringCredentials(results, Errors.UNACCEPTABLE_CREDENTIAL, "when altering the credentials with too many iterations")
+    assertEquals("Too many iterations", results.get(0).errorMessage)
+  }
+
+  @Test
+  def testDeleteSomethingThatDoesNotExist(): Unit = {
+    val request = new AlterUserScramCredentialsRequest.Builder(
+      new AlterUserScramCredentialsRequestData()
+        .setDeletions(util.Arrays.asList(new AlterUserScramCredentialsRequestData.ScramCredentialDeletion().setName(user1).setMechanism(ScramMechanism.SCRAM_SHA_256.`type`)))
+        .setUpsertions(new util.ArrayList[AlterUserScramCredentialsRequestData.ScramCredentialUpsertion])).build()
+    val response = sendAlterUserScramCredentialsRequest(request)
+
+    val results = response.data.results
+    assertEquals(1, results.size)
+    checkAllErrorsAlteringCredentials(results, Errors.RESOURCE_NOT_FOUND, "when deleting a non-existing credential")
+  }
+
+  @Test
+  def testAlterNotController(): Unit = {
+    val request = new AlterUserScramCredentialsRequest.Builder(
+      new AlterUserScramCredentialsRequestData()
+        .setDeletions(util.Arrays.asList(new AlterUserScramCredentialsRequestData.ScramCredentialDeletion().setName(user1).setMechanism(ScramMechanism.SCRAM_SHA_256.`type`)))
+        .setUpsertions(util.Arrays.asList(new AlterUserScramCredentialsRequestData.ScramCredentialUpsertion().setName(user2).setMechanism(ScramMechanism.SCRAM_SHA_512.`type`)))).build()
+    val response = sendAlterUserScramCredentialsRequest(request, notControllerSocketServer)
+
+    val results = response.data.results
+    assertEquals(2, results.size)
+    checkAllErrorsAlteringCredentials(results, Errors.NOT_CONTROLLER, "when routed incorrectly to a non-Controller broker")
+  }
+
+  @Test
+  def testAlterAndDescribe(): Unit = {
+    // create a bunch of credentials
+    val request1 = new AlterUserScramCredentialsRequest.Builder(
+      new AlterUserScramCredentialsRequestData()
+        .setUpsertions(util.Arrays.asList(
+          new AlterUserScramCredentialsRequestData.ScramCredentialUpsertion()
+            .setName(user1).setMechanism(ScramMechanism.SCRAM_SHA_256.`type`)
+            .setIterations(4096)
+            .setSalt(saltBytes)
+            .setSaltedPassword(saltedPasswordBytes),
+          new AlterUserScramCredentialsRequestData.ScramCredentialUpsertion()
+            .setName(user1).setMechanism(ScramMechanism.SCRAM_SHA_512.`type`)
+            .setIterations(8192)
+            .setSalt(saltBytes)
+            .setSaltedPassword(saltedPasswordBytes),
+          new AlterUserScramCredentialsRequestData.ScramCredentialUpsertion()
+            .setName(user2).setMechanism(ScramMechanism.SCRAM_SHA_512.`type`)
+            .setIterations(8192)
+            .setSalt(saltBytes)
+            .setSaltedPassword(saltedPasswordBytes),
+        ))).build()
+    val results1 = sendAlterUserScramCredentialsRequest(request1).data.results
+    assertEquals(2, results1.size)
+    checkNoErrorsAlteringCredentials(results1)
+    checkUserAppearsInAlterResults(results1, user1)
+    checkUserAppearsInAlterResults(results1, user2)
+
+    // now describe them all
+    val results2 = describeAllWithNoTopLevelErrorConfirmed().data.results
+    assertEquals(2, results2.size)
+    checkUserHasTwoCredentials(results2, user1)
+    checkForSingleSha512Iterations8192Credential(results2, user2)
+
+    // now describe just one
+    val request3 = new DescribeUserScramCredentialsRequest.Builder(
+      new DescribeUserScramCredentialsRequestData().setUsers(util.Arrays.asList(
+        new DescribeUserScramCredentialsRequestData.UserName().setName(user1)))).build()
+    val response3 = sendDescribeUserScramCredentialsRequest(request3)
+    checkNoTopLevelErrorDescribingCredentials(response3)
+    val results3 = response3.data.results
+    assertEquals(1, results3.size)
+    checkUserHasTwoCredentials(results3, user1)
+
+    // now test per-user errors by describing user1 and an unknown
+    val requestUnknown = new DescribeUserScramCredentialsRequest.Builder(
+      new DescribeUserScramCredentialsRequestData().setUsers(util.Arrays.asList(
+        new DescribeUserScramCredentialsRequestData.UserName().setName(user1),
+        new DescribeUserScramCredentialsRequestData.UserName().setName(unknownUser)))).build()
+    val responseUnknown = sendDescribeUserScramCredentialsRequest(requestUnknown)
+    checkNoTopLevelErrorDescribingCredentials(responseUnknown)
+    val resultsUnknown = responseUnknown.data.results
+    assertEquals(2, resultsUnknown.size)
+    checkUserHasTwoCredentials(resultsUnknown, user1)
+    checkDescribeForError(resultsUnknown, unknownUser, Errors.RESOURCE_NOT_FOUND)
+
+    // now test per-user errors again by describing user1 along with user2 twice
+    val requestDuplicateUser = new DescribeUserScramCredentialsRequest.Builder(
+      new DescribeUserScramCredentialsRequestData().setUsers(util.Arrays.asList(
+        new DescribeUserScramCredentialsRequestData.UserName().setName(user1),
+        new DescribeUserScramCredentialsRequestData.UserName().setName(user2),
+        new DescribeUserScramCredentialsRequestData.UserName().setName(user2)))).build()
+    val responseDuplicateUser = sendDescribeUserScramCredentialsRequest(requestDuplicateUser)
+    checkNoTopLevelErrorDescribingCredentials(responseDuplicateUser)
+    val resultsDuplicateUser = responseDuplicateUser.data.results
+    assertEquals(2, resultsDuplicateUser.size)
+    checkUserHasTwoCredentials(resultsDuplicateUser, user1)
+    checkDescribeForError(resultsDuplicateUser, user2, Errors.DUPLICATE_RESOURCE)
+
+    // now delete a couple of credentials
+    val request4 = new AlterUserScramCredentialsRequest.Builder(
+      new AlterUserScramCredentialsRequestData()
+        .setDeletions(util.Arrays.asList(
+          new AlterUserScramCredentialsRequestData.ScramCredentialDeletion()
+            .setName(user1).setMechanism(ScramMechanism.SCRAM_SHA_256.`type`),
+          new AlterUserScramCredentialsRequestData.ScramCredentialDeletion()
+            .setName(user2).setMechanism(ScramMechanism.SCRAM_SHA_512.`type`),
+        ))).build()
+    val response4 = sendAlterUserScramCredentialsRequest(request4)
+    val results4 = response4.data.results
+    assertEquals(2, results4.size)
+    checkNoErrorsAlteringCredentials(results4)
+    checkUserAppearsInAlterResults(results4, user1)
+    checkUserAppearsInAlterResults(results4, user2)
+
+    // now describe them all, which should just yield 1 credential
+    val results5 = describeAllWithNoTopLevelErrorConfirmed().data.results
+    assertEquals(1, results5.size)
+    checkForSingleSha512Iterations8192Credential(results5, user1)
+
+    // now delete the last one
+    val request6 = new AlterUserScramCredentialsRequest.Builder(
+      new AlterUserScramCredentialsRequestData()
+        .setDeletions(util.Arrays.asList(
+          new AlterUserScramCredentialsRequestData.ScramCredentialDeletion()
+            .setName(user1).setMechanism(ScramMechanism.SCRAM_SHA_512.`type`),
+        ))).build()
+    val results6 = sendAlterUserScramCredentialsRequest(request6).data.results
+    assertEquals(1, results6.size)
+    checkNoErrorsAlteringCredentials(results6)
+    checkUserAppearsInAlterResults(results6, user1)
+
+    // now describe them all, which should yield 0 credentials
+    val results7 = describeAllWithNoTopLevelErrorConfirmed().data.results
+    assertEquals(0, results7.size)
+  }
+
+  private def sendAlterUserScramCredentialsRequest(request: AlterUserScramCredentialsRequest, socketServer: SocketServer = controllerSocketServer): AlterUserScramCredentialsResponse = {
+    connectAndReceive[AlterUserScramCredentialsResponse](request, destination = socketServer)
+  }
+
+  private def sendDescribeUserScramCredentialsRequest(request: DescribeUserScramCredentialsRequest, socketServer: SocketServer = controllerSocketServer): DescribeUserScramCredentialsResponse = {
+    connectAndReceive[DescribeUserScramCredentialsResponse](request, destination = socketServer)
+  }
+
+  private def checkAllErrorsAlteringCredentials(resultsToCheck: util.List[AlterUserScramCredentialsResult], expectedError: Errors, contextMsg: String) = {
+    assertEquals(s"Expected all '${expectedError.name}' errors when altering credentials $contextMsg",
+      0, resultsToCheck.asScala.filterNot(_.errorCode == expectedError.code).size)
+  }
+
+  private def checkNoErrorsAlteringCredentials(resultsToCheck: util.List[AlterUserScramCredentialsResult]) = {
+    assertEquals("Expected no error when altering credentials",
+      0, resultsToCheck.asScala.filterNot(_.errorCode == Errors.NONE.code).size)
+  }
+
+  private def checkUserAppearsInAlterResults(resultsToCheck: util.List[AlterUserScramCredentialsResult], user: String) = {
+    assertTrue(s"Expected result to contain '$user'", resultsToCheck.asScala.exists(_.user == user))
+  }
+
+  private def describeAllWithNoTopLevelErrorConfirmed() = {
+    val response = sendDescribeUserScramCredentialsRequest(
+      new DescribeUserScramCredentialsRequest.Builder(new DescribeUserScramCredentialsRequestData()).build())
+    checkNoTopLevelErrorDescribingCredentials(response)
+    response
+  }
+
+  private def checkNoTopLevelErrorDescribingCredentials(responseToCheck: DescribeUserScramCredentialsResponse) = {
+    assertEquals("Expected no top-level error when describing the credentials", Errors.NONE.code, responseToCheck.data.errorCode)
+  }
+
+  private def checkUserHasTwoCredentials(resultsToCheck: util.List[DescribeUserScramCredentialsResult], user: String) = {
+    assertTrue(s"Expected result to contain '$user' with 2 credentials: $resultsToCheck",
+      resultsToCheck.asScala.exists(result => result.user == user && result.credentialInfos.size == 2 && result.errorCode == Errors.NONE.code))
+    assertTrue(s"Expected result to contain '$user' with SCRAM_SHA_256/4096 and SCRAM_SHA_512/8192 credentials: $resultsToCheck",
+      resultsToCheck.asScala.exists(result => result.user == user && result.credentialInfos.asScala.exists(info =>
+        info.mechanism == ScramMechanism.SCRAM_SHA_256.`type` && info.iterations == 4096)
+        && result.credentialInfos.asScala.exists(info =>
+        info.mechanism == ScramMechanism.SCRAM_SHA_512.`type` && info.iterations == 8192)))
+  }
+
+  private def checkForSingleSha512Iterations8192Credential(resultsToCheck: util.List[DescribeUserScramCredentialsResult], user: String) = {
+    assertTrue(s"Expected result to contain '$user' with 1 credential: $resultsToCheck",
+      resultsToCheck.asScala.exists(result => result.user == user && result.credentialInfos.size == 1 && result.errorCode == Errors.NONE.code))
+    assertTrue(s"Expected result to contain '$user' with SCRAM_SHA_512/8192 credential: $resultsToCheck",
+      resultsToCheck.asScala.exists(result => result.user == user && result.credentialInfos.asScala.exists(info =>
+        info.mechanism == ScramMechanism.SCRAM_SHA_512.`type` && info.iterations == 8192)))
+  }
+
+  private def checkDescribeForError(resultsToCheck: util.List[DescribeUserScramCredentialsResult], user: String, expectedError: Errors) = {
+    assertTrue(s"Expected result to contain '$user' with a ${expectedError.name} error: $resultsToCheck",
+      resultsToCheck.asScala.exists(result => result.user == user && result.credentialInfos.size == 0 && result.errorCode == expectedError.code))
+  }
+}
+
+object AlterCredentialsTest {
+  val UnauthorizedPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "Unauthorized")
+  val AuthorizedPrincipal = KafkaPrincipal.ANONYMOUS
+
+  class TestAuthorizer extends AclAuthorizer {
+    override def authorize(requestContext: AuthorizableRequestContext, actions: util.List[Action]): util.List[AuthorizationResult] = {
+      actions.asScala.map { _ =>
+        if (requestContext.requestType == ApiKeys.ALTER_USER_SCRAM_CREDENTIALS.id && requestContext.principal == UnauthorizedPrincipal)
+          AuthorizationResult.DENIED
+        else
+          AuthorizationResult.ALLOWED
+      }.asJava
+    }
+  }
+
+  class TestPrincipalBuilderReturningAuthorized extends KafkaPrincipalBuilder {
+    override def build(context: AuthenticationContext): KafkaPrincipal = {
+      AuthorizedPrincipal
+    }
+  }
+
+  class TestPrincipalBuilderReturningUnauthorized extends KafkaPrincipalBuilder {
+    override def build(context: AuthenticationContext): KafkaPrincipal = {
+      UnauthorizedPrincipal
+    }
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/server/ApiVersionsRequestTest.scala b/core/src/test/scala/unit/kafka/server/ApiVersionsRequestTest.scala
index 83f7111bb2bbe..916535bb8e3ae 100644
--- a/core/src/test/scala/unit/kafka/server/ApiVersionsRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ApiVersionsRequestTest.scala
@@ -17,46 +17,52 @@
 
 package kafka.server
 
+import org.apache.kafka.common.message.ApiVersionsRequestData
 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
 
-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) {
-      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)
-      assertEquals(s"Received unexpected min version for API key ${actualApiVersion.apiKey}.", expectedApiVersion.minVersion, actualApiVersion.minVersion)
-      assertEquals(s"Received unexpected max version for API key ${actualApiVersion.apiKey}.", expectedApiVersion.maxVersion, actualApiVersion.maxVersion)
-    }
-  }
-}
+class ApiVersionsRequestTest extends AbstractApiVersionsRequestTest {
 
-class ApiVersionsRequestTest extends BaseRequestTest {
+  override def brokerCount: Int = 1
+
+  @Test
+  def testApiVersionsRequest(): Unit = {
+    val request = new ApiVersionsRequest.Builder().build()
+    val apiVersionsResponse = sendApiVersionsRequest(request)
+    validateApiVersionsResponse(apiVersionsResponse)
+  }
 
-  override def numBrokers: Int = 1
+  @Test
+  def testApiVersionsRequestWithUnsupportedVersion(): Unit = {
+    val apiVersionsRequest = new ApiVersionsRequest.Builder().build()
+    val apiVersionsResponse = sendUnsupportedApiVersionRequest(apiVersionsRequest)
+    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())
+  }
 
   @Test
-  def testApiVersionsRequest() {
-    val apiVersionsResponse = sendApiVersionsRequest(new ApiVersionsRequest.Builder().build())
-    ApiVersionsRequestTest.validateApiVersionsResponse(apiVersionsResponse)
+  def testApiVersionsRequestValidationV0(): Unit = {
+    val apiVersionsRequest = new ApiVersionsRequest.Builder().build(0.asInstanceOf[Short])
+    val apiVersionsResponse = sendApiVersionsRequest(apiVersionsRequest)
+    validateApiVersionsResponse(apiVersionsResponse)
   }
 
   @Test
-  def testApiVersionsRequestWithUnsupportedVersion() {
-    val apiVersionsRequest = new ApiVersionsRequest(0)
-    val apiVersionsResponse = sendApiVersionsRequest(apiVersionsRequest, Some(Short.MaxValue), 0)
-    assertEquals(Errors.UNSUPPORTED_VERSION, apiVersionsResponse.error)
+  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)
+    assertEquals(Errors.INVALID_REQUEST.code(), apiVersionsResponse.data.errorCode())
   }
 
-  private def sendApiVersionsRequest(request: ApiVersionsRequest, apiVersion: Option[Short] = None, responseVersion: Short = 1): ApiVersionsResponse = {
-    val response = connectAndSend(request, ApiKeys.API_VERSIONS, apiVersion = apiVersion)
-    ApiVersionsResponse.parse(response, responseVersion)
+  private def sendApiVersionsRequest(request: ApiVersionsRequest): ApiVersionsResponse = {
+    connectAndReceive[ApiVersionsResponse](request)
   }
+
 }
diff --git a/core/src/test/scala/unit/kafka/server/BaseClientQuotaManagerTest.scala b/core/src/test/scala/unit/kafka/server/BaseClientQuotaManagerTest.scala
new file mode 100644
index 0000000000000..423387f336cfd
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/BaseClientQuotaManagerTest.scala
@@ -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 kafka.server
+
+import java.net.InetAddress
+import java.util
+import java.util.Collections
+import kafka.network.RequestChannel
+import kafka.network.RequestChannel.EndThrottlingResponse
+import kafka.network.RequestChannel.Session
+import kafka.network.RequestChannel.StartThrottlingResponse
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.memory.MemoryPool
+import org.apache.kafka.common.metrics.MetricConfig
+import org.apache.kafka.common.metrics.Metrics
+import org.apache.kafka.common.network.ClientInformation
+import org.apache.kafka.common.network.ListenerName
+import org.apache.kafka.common.requests.{AbstractRequest, FetchRequest, RequestContext, RequestHeader, RequestTestUtils}
+import org.apache.kafka.common.requests.FetchRequest.PartitionData
+import org.apache.kafka.common.security.auth.KafkaPrincipal
+import org.apache.kafka.common.security.auth.SecurityProtocol
+import org.apache.kafka.common.utils.MockTime
+import org.easymock.EasyMock
+import org.junit.After
+
+class BaseClientQuotaManagerTest {
+  protected val time = new MockTime
+  protected var numCallbacks: Int = 0
+  protected val metrics = new Metrics(new MetricConfig(), Collections.emptyList(), time)
+
+  @After
+  def tearDown(): Unit = {
+    metrics.close()
+  }
+
+  protected def callback(response: RequestChannel.Response): Unit = {
+    // Count how many times this callback is called for notifyThrottlingDone().
+    (response: @unchecked) match {
+      case _: StartThrottlingResponse =>
+      case _: EndThrottlingResponse => numCallbacks += 1
+    }
+  }
+
+  protected def buildRequest[T <: AbstractRequest](builder: AbstractRequest.Builder[T],
+                                                   listenerName: ListenerName = ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT)): (T, RequestChannel.Request) = {
+
+    val request = builder.build()
+    val buffer = RequestTestUtils.serializeRequestWithHeader(
+      new RequestHeader(builder.apiKey, request.version, "", 0), request)
+    val requestChannelMetrics: RequestChannel.Metrics = EasyMock.createNiceMock(classOf[RequestChannel.Metrics])
+
+    // read the header from the buffer first so that the body can be read next from the Request constructor
+    val header = RequestHeader.parse(buffer)
+    val context = new RequestContext(header, "1", InetAddress.getLocalHost, KafkaPrincipal.ANONYMOUS,
+      listenerName, SecurityProtocol.PLAINTEXT, ClientInformation.EMPTY, false)
+    (request, new RequestChannel.Request(processor = 1, context = context, startTimeNanos =  0, MemoryPool.NONE, buffer,
+      requestChannelMetrics))
+  }
+
+  protected def buildSession(user: String): Session = {
+    val principal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, user)
+    Session(principal, null)
+  }
+
+  protected def maybeRecord(quotaManager: ClientQuotaManager, user: String, clientId: String, value: Double): Int = {
+
+    quotaManager.maybeRecordAndGetThrottleTimeMs(buildSession(user), clientId, value, time.milliseconds)
+  }
+
+  protected def throttle(quotaManager: ClientQuotaManager, user: String, clientId: String, throttleTimeMs: Int,
+                         channelThrottlingCallback: RequestChannel.Response => Unit): Unit = {
+    val (_, request) = buildRequest(FetchRequest.Builder.forConsumer(0, 1000, new util.HashMap[TopicPartition, PartitionData]))
+    quotaManager.throttle(request, throttleTimeMs, channelThrottlingCallback)
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/server/BaseRequestTest.scala b/core/src/test/scala/unit/kafka/server/BaseRequestTest.scala
index f91afd492fede..13cd74e5a3b92 100644
--- a/core/src/test/scala/unit/kafka/server/BaseRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/BaseRequestTest.scala
@@ -21,160 +21,132 @@ import java.io.{DataInputStream, DataOutputStream}
 import java.net.Socket
 import java.nio.ByteBuffer
 import java.util.Properties
-
-import kafka.integration.KafkaServerTestHarness
+import kafka.api.IntegrationTestHarness
 import kafka.network.SocketServer
-import kafka.utils._
+import kafka.utils.NotNothing
 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.security.auth.SecurityProtocol
+import org.apache.kafka.common.requests.{AbstractRequest, AbstractResponse, RequestHeader, RequestTestUtils, ResponseHeader}
+import org.apache.kafka.common.utils.Utils
+
+import scala.annotation.nowarn
+import scala.collection.Seq
+import scala.reflect.ClassTag
 
-abstract class BaseRequestTest extends KafkaServerTestHarness {
+abstract class BaseRequestTest extends IntegrationTestHarness {
   private var correlationId = 0
 
   // If required, set number of brokers
-  protected def numBrokers: Int = 3
-
-  protected def logDirCount: Int = 1
+  override def brokerCount: Int = 3
 
   // If required, override properties by mutating the passed Properties object
-  protected def propertyOverrides(properties: Properties) {}
-
-  def generateConfigs = {
-    val props = TestUtils.createBrokerConfigs(numBrokers, zkConnect,
-      enableControlledShutdown = false, enableDeleteTopic = true,
-      interBrokerSecurityProtocol = Some(securityProtocol),
-      trustStoreFile = trustStoreFile, saslProperties = serverSaslProperties, logDirCount = logDirCount)
-    props.foreach(propertyOverrides)
-    props.map(KafkaConfig.fromProps)
+  protected def brokerPropertyOverrides(properties: Properties): Unit = {}
+
+  override def modifyConfigs(props: Seq[Properties]): Unit = {
+    props.foreach { p =>
+      p.put(KafkaConfig.ControlledShutdownEnableProp, "false")
+      brokerPropertyOverrides(p)
+    }
   }
 
-  def anySocketServer = {
+  def anySocketServer: SocketServer = {
     servers.find { server =>
       val state = server.brokerState.currentState
       state != NotRunning.state && state != BrokerShuttingDown.state
     }.map(_.socketServer).getOrElse(throw new IllegalStateException("No live broker is available"))
   }
 
-  def controllerSocketServer = {
+  def controllerSocketServer: SocketServer = {
     servers.find { server =>
       server.kafkaController.isActive
     }.map(_.socketServer).getOrElse(throw new IllegalStateException("No controller broker is available"))
   }
 
-  def notControllerSocketServer = {
+  def notControllerSocketServer: SocketServer = {
     servers.find { server =>
       !server.kafkaController.isActive
     }.map(_.socketServer).getOrElse(throw new IllegalStateException("No non-controller broker is available"))
   }
 
-  def brokerSocketServer(brokerId: Int) = {
+  def brokerSocketServer(brokerId: Int): SocketServer = {
     servers.find { server =>
       server.config.brokerId == brokerId
     }.map(_.socketServer).getOrElse(throw new IllegalStateException(s"Could not find broker with id $brokerId"))
   }
 
-  def connect(s: SocketServer = anySocketServer, protocol: SecurityProtocol = SecurityProtocol.PLAINTEXT): Socket = {
-    new Socket("localhost", s.boundPort(ListenerName.forSecurityProtocol(protocol)))
+  def connect(socketServer: SocketServer = anySocketServer,
+              listenerName: ListenerName = listenerName): Socket = {
+    new Socket("localhost", socketServer.boundPort(listenerName))
   }
 
-  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)
     outgoing.flush()
   }
 
-  private def receiveResponse(socket: Socket): Array[Byte] = {
+  def receive[T <: AbstractResponse](socket: Socket, apiKey: ApiKeys, version: Short)
+                                    (implicit classTag: ClassTag[T], @nowarn("cat=unused") nn: NotNothing[T]): T = {
     val incoming = new DataInputStream(socket.getInputStream)
     val len = incoming.readInt()
-    val response = new Array[Byte](len)
-    incoming.readFully(response)
-    response
-  }
 
-  def requestAndReceive(socket: Socket, request: Array[Byte]): Array[Byte] = {
-    sendRequest(socket, request)
-    receiveResponse(socket)
-  }
+    val responseBytes = new Array[Byte](len)
+    incoming.readFully(responseBytes)
 
-  /**
-    * @param destination An optional SocketServer ot send the request to. If not set, any available server is used.
-    * @param protocol An optional SecurityProtocol to use. If not set, PLAINTEXT is used.
-    * @return A ByteBuffer containing the response (without the response header)
-    */
-  def connectAndSend(request: AbstractRequest, apiKey: ApiKeys,
-                     destination: SocketServer = anySocketServer,
-                     apiVersion: Option[Short] = None,
-                     protocol: SecurityProtocol = SecurityProtocol.PLAINTEXT): ByteBuffer = {
-    val socket = connect(destination, protocol)
-    try sendAndReceive(request, apiKey, socket, apiVersion)
-    finally socket.close()
-  }
+    val responseBuffer = ByteBuffer.wrap(responseBytes)
+    ResponseHeader.parse(responseBuffer, apiKey.responseHeaderVersion(version))
 
-  /**
-    * @param destination An optional SocketServer ot send the request to. If not set, any available server is used.
-    * @param protocol An optional SecurityProtocol to use. If not set, PLAINTEXT is used.
-    * @return A ByteBuffer containing the response (without the response header).
-    */
-  def connectAndSendStruct(requestStruct: Struct, apiKey: ApiKeys, apiVersion: Short,
-                           destination: SocketServer = anySocketServer,
-                           protocol: SecurityProtocol = SecurityProtocol.PLAINTEXT): ByteBuffer = {
-    val socket = connect(destination, protocol)
-    try sendStructAndReceive(requestStruct, apiKey, socket, apiVersion)
-    finally socket.close()
+    AbstractResponse.parseResponse(apiKey, responseBuffer, version) match {
+      case response: T => response
+      case response =>
+        throw new ClassCastException(s"Expected response with type ${classTag.runtimeClass}, but found ${response.getClass}")
+    }
   }
 
-  /**
-    * 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))
-    val serializedBytes = request.serialize(header).array
-    sendRequest(socket, serializedBytes)
+  def sendAndReceive[T <: AbstractResponse](request: AbstractRequest,
+                                            socket: Socket,
+                                            clientId: String = "client-id",
+                                            correlationId: Option[Int] = None)
+                                           (implicit classTag: ClassTag[T], nn: NotNothing[T]): T = {
+    send(request, socket, clientId, correlationId)
+    receive[T](socket, request.apiKey, request.version)
   }
 
-  /**
-   * Receive response and return a ByteBuffer containing response without the header
-   */
-  def receive(socket: Socket): ByteBuffer = {
-    val response = receiveResponse(socket)
-    skipResponseHeader(response)
+  def connectAndReceive[T <: AbstractResponse](request: AbstractRequest,
+                                               destination: SocketServer = anySocketServer,
+                                               listenerName: ListenerName = listenerName)
+                                              (implicit classTag: ClassTag[T], nn: NotNothing[T]): T = {
+    val socket = connect(destination, listenerName)
+    try sendAndReceive[T](request, socket)
+    finally socket.close()
   }
 
   /**
     * Serializes and sends the request to the given api.
-    * 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 response = receiveResponse(socket)
-    skipResponseHeader(response)
-  }
-
-  /**
-    * Serializes and sends the requestStruct to the given api.
-    * A ByteBuffer containing the response (without the response header) is returned.
     */
-  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 response = requestAndReceive(socket, serializedBytes)
-    skipResponseHeader(response)
+  def send(request: AbstractRequest,
+           socket: Socket,
+           clientId: String = "client-id",
+           correlationId: Option[Int] = None): Unit = {
+    val header = nextRequestHeader(request.apiKey, request.version, clientId, correlationId)
+    sendWithHeader(request, header, socket)
   }
 
-  protected def skipResponseHeader(response: Array[Byte]): ByteBuffer = {
-    val responseBuffer = ByteBuffer.wrap(response)
-    // Parse the header to ensure its valid and move the buffer forward
-    ResponseHeader.parse(responseBuffer)
-    responseBuffer
+  def sendWithHeader(request: AbstractRequest, header: RequestHeader, socket: Socket): Unit = {
+    val serializedBytes = Utils.toArray(RequestTestUtils.serializeRequestWithHeader(header, request))
+    sendRequest(socket, serializedBytes)
   }
 
-  def nextRequestHeader(apiKey: ApiKeys, apiVersion: Short): RequestHeader = {
-    correlationId += 1
-    new RequestHeader(apiKey, apiVersion, "client-id", correlationId)
+  def nextRequestHeader[T <: AbstractResponse](apiKey: ApiKeys,
+                                               apiVersion: Short,
+                                               clientId: String = "client-id",
+                                               correlationIdOpt: Option[Int] = None): RequestHeader = {
+    val correlationId = correlationIdOpt.getOrElse {
+      this.correlationId += 1
+      this.correlationId
+    }
+    new RequestHeader(apiKey, apiVersion, clientId, correlationId)
   }
 
 }
diff --git a/core/src/test/scala/unit/kafka/server/BrokerEpochIntegrationTest.scala b/core/src/test/scala/unit/kafka/server/BrokerEpochIntegrationTest.scala
new file mode 100755
index 0000000000000..560c251e65f41
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/BrokerEpochIntegrationTest.scala
@@ -0,0 +1,280 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package kafka.server
+
+import kafka.api.LeaderAndIsr
+import kafka.cluster.Broker
+import kafka.controller.{ControllerChannelManager, ControllerContext, StateChangeLogger}
+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.StopReplicaRequestData.{StopReplicaPartitionState, StopReplicaTopicState}
+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._
+import org.apache.kafka.common.security.auth.SecurityProtocol
+import org.apache.kafka.common.utils.Time
+import org.junit.Assert._
+import org.junit.{After, Before, Test}
+
+import scala.jdk.CollectionConverters._
+
+class BrokerEpochIntegrationTest extends ZooKeeperTestHarness {
+  val brokerId1 = 0
+  val brokerId2 = 1
+
+  var servers: Seq[KafkaServer] = Seq.empty[KafkaServer]
+
+  @Before
+  override def setUp(): Unit = {
+    super.setUp()
+    val configs = Seq(
+      TestUtils.createBrokerConfig(brokerId1, zkConnect),
+      TestUtils.createBrokerConfig(brokerId2, zkConnect))
+
+    configs.foreach { config =>
+        config.setProperty(KafkaConfig.AutoLeaderRebalanceEnableProp, false.toString)}
+
+    // start both servers
+    servers = configs.map(config => TestUtils.createServer(KafkaConfig.fromProps(config)))
+  }
+
+  @After
+  override def tearDown(): Unit = {
+    TestUtils.shutdownServers(servers)
+    super.tearDown()
+  }
+
+  @Test
+  def testReplicaManagerBrokerEpochMatchesWithZk(): Unit = {
+    val brokerAndEpochs = zkClient.getAllBrokerAndEpochsInCluster
+    assertEquals(brokerAndEpochs.size, servers.size)
+    brokerAndEpochs.foreach {
+      case (broker, epoch) =>
+        val brokerServer = servers.find(e => e.config.brokerId == broker.id)
+        assertTrue(brokerServer.isDefined)
+        assertEquals(epoch, brokerServer.get.kafkaController.brokerEpoch)
+    }
+  }
+
+  @Test
+  def testControllerBrokerEpochCacheMatchesWithZk(): Unit = {
+    val controller = getController
+    val otherBroker = servers.find(e => e.config.brokerId != controller.config.brokerId).get
+
+    // Broker epochs cache matches with zk in steady state
+    checkControllerBrokerEpochsCacheMatchesWithZk(controller.kafkaController.controllerContext)
+
+    // Shutdown a broker and make sure broker epochs cache still matches with zk state
+    otherBroker.shutdown()
+    checkControllerBrokerEpochsCacheMatchesWithZk(controller.kafkaController.controllerContext)
+
+    // Restart a broker and make sure broker epochs cache still matches with zk state
+    otherBroker.startup()
+    checkControllerBrokerEpochsCacheMatchesWithZk(controller.kafkaController.controllerContext)
+  }
+
+  @Test
+  def testControlRequestWithCorrectBrokerEpoch(): Unit = {
+    testControlRequestWithBrokerEpoch(0)
+  }
+
+  @Test
+  def testControlRequestWithStaleBrokerEpoch(): Unit = {
+    testControlRequestWithBrokerEpoch(-1)
+  }
+
+  @Test
+  def testControlRequestWithNewerBrokerEpoch(): Unit = {
+    testControlRequestWithBrokerEpoch(1)
+  }
+
+  private def testControlRequestWithBrokerEpoch(epochInRequestDiffFromCurrentEpoch: Long): Unit = {
+    val tp = new TopicPartition("new-topic", 0)
+
+    // create topic with 1 partition, 2 replicas, one on each broker
+    createTopic(zkClient, tp.topic(), partitionReplicaAssignment = Map(0 -> Seq(brokerId1, brokerId2)), servers = servers)
+
+    val controllerId = 2
+    val controllerEpoch = zkClient.getControllerEpoch.get._1
+
+    val controllerConfig = KafkaConfig.fromProps(TestUtils.createBrokerConfig(controllerId, zkConnect))
+    val securityProtocol = SecurityProtocol.PLAINTEXT
+    val listenerName = ListenerName.forSecurityProtocol(securityProtocol)
+    val brokerAndEpochs = servers.map(s =>
+      (new Broker(s.config.brokerId, "localhost", TestUtils.boundPort(s), listenerName, securityProtocol),
+        s.kafkaController.brokerEpoch)).toMap
+    val nodes = brokerAndEpochs.keys.map(_.node(listenerName))
+
+    val controllerContext = new ControllerContext
+    controllerContext.setLiveBrokers(brokerAndEpochs)
+    val metrics = new Metrics
+    val controllerChannelManager = new ControllerChannelManager(controllerContext, controllerConfig, Time.SYSTEM,
+      metrics, new StateChangeLogger(controllerId, inControllerContext = true, None))
+    controllerChannelManager.startup()
+
+    val broker2 = servers(brokerId2)
+    val epochInRequest = broker2.kafkaController.brokerEpoch + epochInRequestDiffFromCurrentEpoch
+
+    try {
+      // Send LeaderAndIsr request with correct broker epoch
+      {
+        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,
+          partitionStates.asJava, nodes.toSet.asJava)
+
+        if (epochInRequestDiffFromCurrentEpoch < 0) {
+          // stale broker epoch in LEADER_AND_ISR
+          sendAndVerifyStaleBrokerEpochInResponse(controllerChannelManager, requestBuilder)
+        }
+        else {
+          // broker epoch in LEADER_AND_ISR >= current broker epoch
+          sendAndVerifySuccessfulResponse(controllerChannelManager, requestBuilder)
+          TestUtils.waitUntilLeaderIsKnown(Seq(broker2), tp, 10000)
+        }
+      }
+
+      // Send UpdateMetadata request with correct broker epoch
+      {
+        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 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, liveBrokers.asJava)
+
+        if (epochInRequestDiffFromCurrentEpoch < 0) {
+          // stale broker epoch in UPDATE_METADATA
+          sendAndVerifyStaleBrokerEpochInResponse(controllerChannelManager, requestBuilder)
+        }
+        else {
+          // broker epoch in UPDATE_METADATA >= current broker epoch
+          sendAndVerifySuccessfulResponse(controllerChannelManager, requestBuilder)
+          TestUtils.waitUntilMetadataIsPropagated(Seq(broker2), tp.topic, tp.partition, 10000)
+          assertEquals(brokerId2,
+            broker2.metadataCache.getPartitionInfo(tp.topic, tp.partition).get.leader)
+        }
+      }
+
+      // Send StopReplica request with correct broker epoch
+      {
+        val topicStates = Seq(
+          new StopReplicaTopicState()
+            .setTopicName(tp.topic())
+            .setPartitionStates(Seq(new StopReplicaPartitionState()
+              .setPartitionIndex(tp.partition())
+              .setLeaderEpoch(LeaderAndIsr.initialLeaderEpoch + 2)
+              .setDeletePartition(true)).asJava)
+        ).asJava
+        val requestBuilder = new StopReplicaRequest.Builder(
+          ApiKeys.STOP_REPLICA.latestVersion, controllerId, controllerEpoch,
+          epochInRequest, // Correct broker epoch
+          false, topicStates)
+
+        if (epochInRequestDiffFromCurrentEpoch < 0) {
+          // stale broker epoch in STOP_REPLICA
+          sendAndVerifyStaleBrokerEpochInResponse(controllerChannelManager, requestBuilder)
+        } else {
+          // broker epoch in STOP_REPLICA >= current broker epoch
+          sendAndVerifySuccessfulResponse(controllerChannelManager, requestBuilder)
+          assertEquals(HostedPartition.None, broker2.replicaManager.getPartition(tp))
+        }
+      }
+    } finally {
+      controllerChannelManager.shutdown()
+      metrics.close()
+    }
+  }
+
+  private def getController: KafkaServer = {
+    val controllerId = TestUtils.waitUntilControllerElected(zkClient)
+    servers.filter(s => s.config.brokerId == controllerId).head
+  }
+
+  private def checkControllerBrokerEpochsCacheMatchesWithZk(controllerContext: ControllerContext): Unit = {
+    val brokerAndEpochs = zkClient.getAllBrokerAndEpochsInCluster
+    TestUtils.waitUntilTrue(() => {
+      val brokerEpochsInControllerContext = controllerContext.liveBrokerIdAndEpochs
+      if (brokerAndEpochs.size != brokerEpochsInControllerContext.size) false
+      else {
+        brokerAndEpochs.forall {
+          case (broker, epoch) => brokerEpochsInControllerContext.get(broker.id).contains(epoch)
+        }
+      }
+    }, "Broker epoch mismatches")
+  }
+
+  private def sendAndVerifyStaleBrokerEpochInResponse(controllerChannelManager: ControllerChannelManager,
+                                                      builder: AbstractControlRequest.Builder[_ <: AbstractControlRequest]): Unit = {
+    var staleBrokerEpochDetected = false
+    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,
+                                              builder: AbstractControlRequest.Builder[_ <: AbstractControlRequest]): Unit = {
+    @volatile var succeed = false
+    controllerChannelManager.sendRequest(brokerId2, builder, response => {
+        succeed = response.errorCounts().isEmpty ||
+          (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/BrokerFeaturesTest.scala b/core/src/test/scala/unit/kafka/server/BrokerFeaturesTest.scala
new file mode 100644
index 0000000000000..2fe08fb3e1317
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/BrokerFeaturesTest.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 org.apache.kafka.common.feature.{Features, FinalizedVersionRange, SupportedVersionRange}
+import org.junit.Assert.{assertEquals, assertFalse, assertTrue}
+import org.junit.Test
+
+import scala.jdk.CollectionConverters._
+
+class BrokerFeaturesTest {
+
+  @Test
+  def testEmpty(): Unit = {
+    assertTrue(BrokerFeatures.createDefault().supportedFeatures.empty)
+  }
+
+  @Test
+  def testIncompatibilitiesDueToAbsentFeature(): Unit = {
+    val brokerFeatures = BrokerFeatures.createDefault()
+    val supportedFeatures = Features.supportedFeatures(Map[String, SupportedVersionRange](
+      "test_feature_1" -> new SupportedVersionRange(1, 4),
+      "test_feature_2" -> new SupportedVersionRange(1, 3)).asJava)
+    brokerFeatures.setSupportedFeatures(supportedFeatures)
+
+    val compatibleFeatures = Map[String, FinalizedVersionRange](
+      "test_feature_1" -> new FinalizedVersionRange(2, 3))
+    val inCompatibleFeatures = Map[String, FinalizedVersionRange](
+      "test_feature_3" -> new FinalizedVersionRange(3, 4))
+    val features = compatibleFeatures++inCompatibleFeatures
+    val finalizedFeatures = Features.finalizedFeatures(features.asJava)
+
+    assertEquals(
+      Features.finalizedFeatures(inCompatibleFeatures.asJava),
+      brokerFeatures.incompatibleFeatures(finalizedFeatures))
+    assertTrue(BrokerFeatures.hasIncompatibleFeatures(supportedFeatures, finalizedFeatures))
+  }
+
+  @Test
+  def testIncompatibilitiesDueToIncompatibleFeature(): Unit = {
+    val brokerFeatures = BrokerFeatures.createDefault()
+    val supportedFeatures = Features.supportedFeatures(Map[String, SupportedVersionRange](
+      "test_feature_1" -> new SupportedVersionRange(1, 4),
+      "test_feature_2" -> new SupportedVersionRange(1, 3)).asJava)
+    brokerFeatures.setSupportedFeatures(supportedFeatures)
+
+    val compatibleFeatures = Map[String, FinalizedVersionRange](
+      "test_feature_1" -> new FinalizedVersionRange(2, 3))
+    val inCompatibleFeatures = Map[String, FinalizedVersionRange](
+      "test_feature_2" -> new FinalizedVersionRange(1, 4))
+    val features = compatibleFeatures++inCompatibleFeatures
+    val finalizedFeatures = Features.finalizedFeatures(features.asJava)
+
+    assertEquals(
+      Features.finalizedFeatures(inCompatibleFeatures.asJava),
+      brokerFeatures.incompatibleFeatures(finalizedFeatures))
+    assertTrue(BrokerFeatures.hasIncompatibleFeatures(supportedFeatures, finalizedFeatures))
+  }
+
+  @Test
+  def testCompatibleFeatures(): Unit = {
+    val brokerFeatures = BrokerFeatures.createDefault()
+    val supportedFeatures = Features.supportedFeatures(Map[String, SupportedVersionRange](
+      "test_feature_1" -> new SupportedVersionRange(1, 4),
+      "test_feature_2" -> new SupportedVersionRange(1, 3)).asJava)
+    brokerFeatures.setSupportedFeatures(supportedFeatures)
+
+    val compatibleFeatures = Map[String, FinalizedVersionRange](
+      "test_feature_1" -> new FinalizedVersionRange(2, 3),
+      "test_feature_2" -> new FinalizedVersionRange(1, 3))
+    val finalizedFeatures = Features.finalizedFeatures(compatibleFeatures.asJava)
+    assertTrue(brokerFeatures.incompatibleFeatures(finalizedFeatures).empty())
+    assertFalse(BrokerFeatures.hasIncompatibleFeatures(supportedFeatures, finalizedFeatures))
+  }
+
+  @Test
+  def testDefaultFinalizedFeatures(): Unit = {
+    val brokerFeatures = BrokerFeatures.createDefault()
+    val supportedFeatures = Features.supportedFeatures(Map[String, SupportedVersionRange](
+      "test_feature_1" -> new SupportedVersionRange(1, 4),
+      "test_feature_2" -> new SupportedVersionRange(1, 3),
+      "test_feature_3" -> new SupportedVersionRange(3, 7)).asJava)
+    brokerFeatures.setSupportedFeatures(supportedFeatures)
+
+    val expectedFeatures = Map[String, FinalizedVersionRange](
+      "test_feature_1" -> new FinalizedVersionRange(1, 4),
+      "test_feature_2" -> new FinalizedVersionRange(1, 3),
+      "test_feature_3" -> new FinalizedVersionRange(3, 7))
+    assertEquals(Features.finalizedFeatures(expectedFeatures.asJava), brokerFeatures.defaultFinalizedFeatures)
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/server/ClientQuotaManagerTest.scala b/core/src/test/scala/unit/kafka/server/ClientQuotaManagerTest.scala
index 1aabbb3b1d476..8547dd6981712 100644
--- a/core/src/test/scala/unit/kafka/server/ClientQuotaManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ClientQuotaManagerTest.scala
@@ -16,69 +16,62 @@
  */
 package kafka.server
 
-import java.util.Collections
+import java.net.InetAddress
 
-import org.apache.kafka.common.metrics.{MetricConfig, Metrics, Quota}
-import org.apache.kafka.common.utils.{MockTime, Sanitizer}
-import org.junit.Assert.{assertEquals, assertTrue}
-import org.junit.{Before, Test}
+import kafka.network.RequestChannel.Session
+import kafka.server.QuotaType._
+import org.apache.kafka.common.metrics.Quota
+import org.apache.kafka.common.security.auth.KafkaPrincipal
+import org.apache.kafka.common.utils.Sanitizer
 
-class ClientQuotaManagerTest {
-  private val time = new MockTime
+import org.junit.Assert._
+import org.junit.Test
 
-  private val config = ClientQuotaManagerConfig(quotaBytesPerSecondDefault = 500)
+class ClientQuotaManagerTest extends BaseClientQuotaManagerTest {
+  private val config = ClientQuotaManagerConfig(quotaDefault = 500)
 
-  var numCallbacks: Int = 0
-  def callback(delayTimeMs: Int) {
-    numCallbacks += 1
-  }
-
-  @Before
-  def beforeMethod() {
-    numCallbacks = 0
-  }
-
-  private def testQuotaParsing(config: ClientQuotaManagerConfig, client1: UserClient, client2: UserClient, randomClient: UserClient, defaultConfigClient: UserClient) {
-    val clientMetrics = new ClientQuotaManager(config, newMetrics, QuotaType.Produce, time, "")
+  private def testQuotaParsing(config: ClientQuotaManagerConfig, client1: UserClient, client2: UserClient, randomClient: UserClient, defaultConfigClient: UserClient): Unit = {
+    val clientQuotaManager = new ClientQuotaManager(config, metrics, Produce, time, "")
 
     try {
       // Case 1: Update the quota. Assert that the new quota value is returned
-      clientMetrics.updateQuota(client1.configUser, client1.configClientId, client1.sanitizedConfigClientId, Some(new Quota(2000, true)))
-      clientMetrics.updateQuota(client2.configUser, client2.configClientId, client2.sanitizedConfigClientId, Some(new Quota(4000, true)))
+      clientQuotaManager.updateQuota(client1.configUser, client1.configClientId, client1.sanitizedConfigClientId, Some(new Quota(2000, true)))
+      clientQuotaManager.updateQuota(client2.configUser, client2.configClientId, client2.sanitizedConfigClientId, Some(new Quota(4000, true)))
 
-      assertEquals("Default producer quota should be " + config.quotaBytesPerSecondDefault, new Quota(config.quotaBytesPerSecondDefault, true), clientMetrics.quota(randomClient.user, randomClient.clientId))
-      assertEquals("Should return the overridden value (2000)", new Quota(2000, true), clientMetrics.quota(client1.user, client1.clientId))
-      assertEquals("Should return the overridden value (4000)", new Quota(4000, true), clientMetrics.quota(client2.user, client2.clientId))
+      assertEquals("Default producer quota should be " + config.quotaDefault,
+        config.quotaDefault.toDouble, clientQuotaManager.quota(randomClient.user, randomClient.clientId).bound, 0.0)
+      assertEquals("Should return the overridden value (2000)", 2000, clientQuotaManager.quota(client1.user, client1.clientId).bound, 0.0)
+      assertEquals("Should return the overridden value (4000)", 4000, clientQuotaManager.quota(client2.user, client2.clientId).bound, 0.0)
 
       // p1 should be throttled using the overridden quota
-      var throttleTimeMs = clientMetrics.maybeRecordAndThrottle(client1.user, client1.clientId, 2500 * config.numQuotaSamples, this.callback)
+      var throttleTimeMs = maybeRecord(clientQuotaManager, client1.user, client1.clientId, 2500 * config.numQuotaSamples)
       assertTrue(s"throttleTimeMs should be > 0. was $throttleTimeMs", throttleTimeMs > 0)
 
       // Case 2: Change quota again. The quota should be updated within KafkaMetrics as well since the sensor was created.
       // p1 should not longer be throttled after the quota change
-      clientMetrics.updateQuota(client1.configUser, client1.configClientId, client1.sanitizedConfigClientId, Some(new Quota(3000, true)))
-      assertEquals("Should return the newly overridden value (3000)", new Quota(3000, true), clientMetrics.quota(client1.user, client1.clientId))
+      clientQuotaManager.updateQuota(client1.configUser, client1.configClientId, client1.sanitizedConfigClientId, Some(new Quota(3000, true)))
+      assertEquals("Should return the newly overridden value (3000)", 3000, clientQuotaManager.quota(client1.user, client1.clientId).bound, 0.0)
 
-      throttleTimeMs = clientMetrics.maybeRecordAndThrottle(client1.user, client1.clientId, 0, this.callback)
+      throttleTimeMs = maybeRecord(clientQuotaManager, client1.user, client1.clientId, 0)
       assertEquals(s"throttleTimeMs should be 0. was $throttleTimeMs", 0, throttleTimeMs)
 
       // Case 3: Change quota back to default. Should be throttled again
-      clientMetrics.updateQuota(client1.configUser, client1.configClientId, client1.sanitizedConfigClientId, Some(new Quota(500, true)))
-      assertEquals("Should return the default value (500)", new Quota(500, true), clientMetrics.quota(client1.user, client1.clientId))
+      clientQuotaManager.updateQuota(client1.configUser, client1.configClientId, client1.sanitizedConfigClientId, Some(new Quota(500, true)))
+      assertEquals("Should return the default value (500)", 500, clientQuotaManager.quota(client1.user, client1.clientId).bound, 0.0)
 
-      throttleTimeMs = clientMetrics.maybeRecordAndThrottle(client1.user, client1.clientId, 0, this.callback)
+      throttleTimeMs = maybeRecord(clientQuotaManager, client1.user, client1.clientId, 0)
       assertTrue(s"throttleTimeMs should be > 0. was $throttleTimeMs", throttleTimeMs > 0)
 
       // Case 4: Set high default quota, remove p1 quota. p1 should no longer be throttled
-      clientMetrics.updateQuota(client1.configUser, client1.configClientId, client1.sanitizedConfigClientId, None)
-      clientMetrics.updateQuota(defaultConfigClient.configUser, defaultConfigClient.configClientId, defaultConfigClient.sanitizedConfigClientId, Some(new Quota(4000, true)))
-      assertEquals("Should return the newly overridden value (4000)", new Quota(4000, true), clientMetrics.quota(client1.user, client1.clientId))
+      clientQuotaManager.updateQuota(client1.configUser, client1.configClientId, client1.sanitizedConfigClientId, None)
+      clientQuotaManager.updateQuota(defaultConfigClient.configUser, defaultConfigClient.configClientId, defaultConfigClient.sanitizedConfigClientId, Some(new Quota(4000, true)))
+      assertEquals("Should return the newly overridden value (4000)", 4000, clientQuotaManager.quota(client1.user, client1.clientId).bound, 0.0)
 
-      throttleTimeMs = clientMetrics.maybeRecordAndThrottle(client1.user, client1.clientId, 1000 * config.numQuotaSamples, this.callback)
+      throttleTimeMs = maybeRecord(clientQuotaManager, client1.user, client1.clientId, 1000 * config.numQuotaSamples)
       assertEquals(s"throttleTimeMs should be 0. was $throttleTimeMs", 0, throttleTimeMs)
 
     } finally {
-      clientMetrics.shutdown()
+      clientQuotaManager.shutdown()
     }
   }
 
@@ -87,7 +80,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)
@@ -100,12 +93,12 @@ 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)
     val defaultConfigClient = UserClient("", "", Some(ConfigEntityName.Default), None)
-    val config = ClientQuotaManagerConfig(quotaBytesPerSecondDefault = Long.MaxValue)
+    val config = ClientQuotaManagerConfig(quotaDefault = Long.MaxValue)
     testQuotaParsing(config, client1, client2, randomClient, defaultConfigClient)
   }
 
@@ -114,12 +107,12 @@ 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)
     val defaultConfigClient = UserClient("", "", Some(ConfigEntityName.Default), Some(ConfigEntityName.Default))
-    val config = ClientQuotaManagerConfig(quotaBytesPerSecondDefault = Long.MaxValue)
+    val config = ClientQuotaManagerConfig(quotaDefault = Long.MaxValue)
     testQuotaParsing(config, client1, client2, randomClient, defaultConfigClient)
   }
 
@@ -127,7 +120,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)
@@ -139,7 +132,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)
@@ -147,219 +140,239 @@ 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.toDouble, 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.toDouble
+      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, QuotaType.Produce, time, "")
-
-    def checkQuota(user: String, clientId: String, expectedBound: Int, value: Int, expectThrottle: Boolean) {
-      assertEquals(new Quota(expectedBound, true), quotaManager.quota(user, clientId))
-      val throttleTimeMs = quotaManager.maybeRecordAndThrottle(user, clientId, value * config.numQuotaSamples, this.callback)
-      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(quotaDefault = Long.MaxValue, numQuotaSamples = numFullQuotaWindows + 1)
+    val clientQuotaManager = new ClientQuotaManager(nonDefaultConfig, metrics, Fetch, time, "")
+    val userSession = Session(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "userA"), InetAddress.getLocalHost)
+
+    try {
+      // no quota set
+      assertEquals(Double.MaxValue, clientQuotaManager.getMaxValueInQuotaWindow(userSession, "client1"), 0.01)
+
+      // Set default  quota config
+      clientQuotaManager.updateQuota(Some(ConfigEntityName.Default), None, None, Some(new Quota(10, true)))
+      assertEquals(10 * numFullQuotaWindows, clientQuotaManager.getMaxValueInQuotaWindow(userSession, "client1"), 0.01)
+    } finally {
+      clientQuotaManager.shutdown()
     }
+  }
+
+  @Test
+  def testSetAndRemoveDefaultUserQuota(): Unit = {
+    // quotaTypesEnabled will be QuotaTypes.NoQuotas initially
+    val clientQuotaManager = new ClientQuotaManager(ClientQuotaManagerConfig(quotaDefault = Long.MaxValue),
+      metrics, Produce, time, "")
 
     try {
-      quotaManager.updateQuota(Some(ConfigEntityName.Default), None, None, Some(new Quota(1000, true)))
-      quotaManager.updateQuota(None, Some(ConfigEntityName.Default), Some(ConfigEntityName.Default), Some(new Quota(2000, true)))
-      quotaManager.updateQuota(Some(ConfigEntityName.Default), Some(ConfigEntityName.Default), Some(ConfigEntityName.Default), Some(new Quota(3000, true)))
-      quotaManager.updateQuota(Some("userA"), None, None, Some(new Quota(4000, true)))
-      quotaManager.updateQuota(Some("userA"), Some("client1"), Some("client1"), Some(new Quota(5000, true)))
-      quotaManager.updateQuota(Some("userB"), None, None, Some(new Quota(6000, true)))
-      quotaManager.updateQuota(Some("userB"), Some("client1"), Some("client1"), Some(new Quota(7000, true)))
-      quotaManager.updateQuota(Some("userB"), Some(ConfigEntityName.Default), Some(ConfigEntityName.Default), Some(new Quota(8000, true)))
-      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("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("userD", "client1", 3000, 3500, true)  // Default  quota
-      checkQuota("userD", "client2", 3000, 2500, false)
-      checkQuota("userE", "client1", 3000, 2500, false)
+      // no quota set yet, should not throttle
+      checkQuota(clientQuotaManager, "userA", "client1", Long.MaxValue, 1000, 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)
+      // Set default  quota config
+      clientQuotaManager.updateQuota(Some(ConfigEntityName.Default), None, None, Some(new Quota(10, true)))
+      checkQuota(clientQuotaManager, "userA", "client1", 10, 1000, 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)
+      // Remove default  quota config, back to no quotas
+      clientQuotaManager.updateQuota(Some(ConfigEntityName.Default), None, None, None)
+      checkQuota(clientQuotaManager, "userA", "client1", Long.MaxValue, 1000, false)
+    } finally {
+      clientQuotaManager.shutdown()
+    }
+  }
 
-      // 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)
-      quotaManager.updateQuota(Some("userA"), Some("client1"), Some("client1"), None)
-      checkQuota("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)
-      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
+  @Test
+  def testSetAndRemoveUserQuota(): Unit = {
+    // quotaTypesEnabled will be QuotaTypes.NoQuotas initially
+    val clientQuotaManager = new ClientQuotaManager(ClientQuotaManagerConfig(quotaDefault = Long.MaxValue),
+      metrics, Produce, time, "")
+
+    try {
+      // Set  quota config
+      clientQuotaManager.updateQuota(Some("userA"), None, None, Some(new Quota(10, true)))
+      checkQuota(clientQuotaManager, "userA", "client1", 10, 1000, true)
 
+      // Remove  quota config, back to no quotas
+      clientQuotaManager.updateQuota(Some("userA"), None, None, None)
+      checkQuota(clientQuotaManager, "userA", "client1", Long.MaxValue, 1000, false)
     } finally {
-      quotaManager.shutdown()
+      clientQuotaManager.shutdown()
     }
   }
 
   @Test
-  def testQuotaViolation() {
-    val metrics = newMetrics
-    val clientMetrics = new ClientQuotaManager(config, metrics, QuotaType.Produce, time, "")
-    val queueSizeMetric = metrics.metrics().get(metrics.metricName("queue-size", "Produce", ""))
+  def testSetAndRemoveUserClientQuota(): Unit = {
+    // quotaTypesEnabled will be QuotaTypes.NoQuotas initially
+    val clientQuotaManager = new ClientQuotaManager(ClientQuotaManagerConfig(quotaDefault = Long.MaxValue),
+      metrics, Produce, time, "")
+
     try {
-      /* We have 10 second windows. Make sure that there is no quota violation
-       * if we produce under the quota
-       */
-      for (_ <- 0 until 10) {
-        clientMetrics.maybeRecordAndThrottle("ANONYMOUS", "unknown", 400, callback)
-        time.sleep(1000)
-      }
-      assertEquals(10, numCallbacks)
-      assertEquals(0, queueSizeMetric.value().toInt)
+      // Set  quota config
+      clientQuotaManager.updateQuota(Some("userA"), Some("client1"), Some("client1"), Some(new Quota(10, true)))
+      checkQuota(clientQuotaManager, "userA", "client1", 10, 1000, true)
 
-      // Create a spike.
-      // 400*10 + 2000 + 300 = 6300/10.5 = 600 bytes per second.
-      // (600 - quota)/quota*window-size = (600-500)/500*10.5 seconds = 2100
-      // 10.5 seconds because the last window is half complete
-      time.sleep(500)
-      val sleepTime = clientMetrics.maybeRecordAndThrottle("ANONYMOUS", "unknown", 2300, callback)
+      // Remove  quota config, back to no quotas
+      clientQuotaManager.updateQuota(Some("userA"), Some("client1"), Some("client1"), None)
+      checkQuota(clientQuotaManager, "userA", "client1", Long.MaxValue, 1000, false)
+    } finally {
+      clientQuotaManager.shutdown()
+    }
+  }
 
-      assertEquals("Should be throttled", 2100, sleepTime)
-      assertEquals(1, queueSizeMetric.value().toInt)
-      // After a request is delayed, the callback cannot be triggered immediately
-      clientMetrics.throttledRequestReaper.doWork()
-      assertEquals(10, numCallbacks)
-      time.sleep(sleepTime)
+  @Test
+  def testQuotaConfigPrecedence(): Unit = {
+    val clientQuotaManager = new ClientQuotaManager(ClientQuotaManagerConfig(quotaDefault=Long.MaxValue),
+      metrics, Produce, time, "")
 
-      // Callback can only be triggered after the delay time passes
-      clientMetrics.throttledRequestReaper.doWork()
-      assertEquals(0, queueSizeMetric.value().toInt)
-      assertEquals(11, numCallbacks)
+    try {
+      clientQuotaManager.updateQuota(Some(ConfigEntityName.Default), None, None, Some(new Quota(1000, true)))
+      clientQuotaManager.updateQuota(None, Some(ConfigEntityName.Default), Some(ConfigEntityName.Default), Some(new Quota(2000, true)))
+      clientQuotaManager.updateQuota(Some(ConfigEntityName.Default), Some(ConfigEntityName.Default), Some(ConfigEntityName.Default), Some(new Quota(3000, true)))
+      clientQuotaManager.updateQuota(Some("userA"), None, None, Some(new Quota(4000, true)))
+      clientQuotaManager.updateQuota(Some("userA"), Some("client1"), Some("client1"), Some(new Quota(5000, true)))
+      clientQuotaManager.updateQuota(Some("userB"), None, None, Some(new Quota(6000, true)))
+      clientQuotaManager.updateQuota(Some("userB"), Some("client1"), Some("client1"), Some(new Quota(7000, true)))
+      clientQuotaManager.updateQuota(Some("userB"), Some(ConfigEntityName.Default), Some(ConfigEntityName.Default), Some(new Quota(8000, true)))
+      clientQuotaManager.updateQuota(Some("userC"), None, None, Some(new Quota(10000, true)))
+      clientQuotaManager.updateQuota(None, Some("client1"), Some("client1"), Some(new Quota(9000, true)))
+
+      checkQuota(clientQuotaManager, "userA", "client1", 5000, 4500, false) //  quota takes precedence over 
+      checkQuota(clientQuotaManager, "userA", "client2", 4000, 4500, true)  //  quota takes precedence over  and defaults
+      checkQuota(clientQuotaManager, "userA", "client3", 4000, 0, true)     //  quota is shared across clients of user
+      checkQuota(clientQuotaManager, "userA", "client1", 5000, 0, false)    //  is exclusive use, unaffected by other clients
+
+      checkQuota(clientQuotaManager, "userB", "client1", 7000, 8000, true)
+      checkQuota(clientQuotaManager, "userB", "client2", 8000, 7000, false) // Default per-client quota for exclusive use of 
+      checkQuota(clientQuotaManager, "userB", "client3", 8000, 7000, false)
+
+      checkQuota(clientQuotaManager, "userD", "client1", 3000, 3500, true)  // Default  quota
+      checkQuota(clientQuotaManager, "userD", "client2", 3000, 2500, false)
+      checkQuota(clientQuotaManager, "userE", "client1", 3000, 2500, false)
 
-      // Could continue to see delays until the bursty sample disappears
-      for (_ <- 0 until 10) {
-        clientMetrics.maybeRecordAndThrottle("ANONYMOUS", "unknown", 400, callback)
-        time.sleep(1000)
-      }
+      // Remove default  quota config, revert to  default
+      clientQuotaManager.updateQuota(Some(ConfigEntityName.Default), Some(ConfigEntityName.Default), Some(ConfigEntityName.Default), None)
+      checkQuota(clientQuotaManager, "userD", "client1", 1000, 0, false)    // Metrics tags changed, restart counter
+      checkQuota(clientQuotaManager, "userE", "client4", 1000, 1500, true)
+      checkQuota(clientQuotaManager, "userF", "client4", 1000, 800, false)  // Default  quota shared across clients of user
+      checkQuota(clientQuotaManager, "userF", "client5", 1000, 800, true)
+
+      // Remove default  quota config, revert to  default
+      clientQuotaManager.updateQuota(Some(ConfigEntityName.Default), None, None, None)
+      checkQuota(clientQuotaManager, "userF", "client4", 2000, 0, false)  // Default  quota shared across client-id of all users
+      checkQuota(clientQuotaManager, "userF", "client5", 2000, 0, false)
+      checkQuota(clientQuotaManager, "userF", "client5", 2000, 2500, true)
+      checkQuota(clientQuotaManager, "userG", "client5", 2000, 0, true)
+
+      // Update quotas
+      clientQuotaManager.updateQuota(Some("userA"), None, None, Some(new Quota(8000, true)))
+      clientQuotaManager.updateQuota(Some("userA"), Some("client1"), Some("client1"), Some(new Quota(10000, true)))
+      checkQuota(clientQuotaManager, "userA", "client2", 8000, 0, false)
+      checkQuota(clientQuotaManager, "userA", "client2", 8000, 4500, true) // Throttled due to sum of new and earlier values
+      checkQuota(clientQuotaManager, "userA", "client1", 10000, 0, false)
+      checkQuota(clientQuotaManager, "userA", "client1", 10000, 6000, true)
+      clientQuotaManager.updateQuota(Some("userA"), Some("client1"), Some("client1"), None)
+      checkQuota(clientQuotaManager, "userA", "client6", 8000, 0, true)    // Throttled due to shared user quota
+      clientQuotaManager.updateQuota(Some("userA"), Some("client6"), Some("client6"), Some(new Quota(11000, true)))
+      checkQuota(clientQuotaManager, "userA", "client6", 11000, 8500, false)
+      clientQuotaManager.updateQuota(Some("userA"), Some(ConfigEntityName.Default), Some(ConfigEntityName.Default), Some(new Quota(12000, true)))
+      clientQuotaManager.updateQuota(Some("userA"), Some("client6"), Some("client6"), None)
+      checkQuota(clientQuotaManager, "userA", "client6", 12000, 4000, true) // Throttled due to sum of new and earlier values
 
-      assertEquals("Should be unthrottled since bursty sample has rolled over",
-                   0, clientMetrics.maybeRecordAndThrottle("ANONYMOUS", "unknown", 0, callback))
     } finally {
-      clientMetrics.shutdown()
+      clientQuotaManager.shutdown()
     }
   }
 
   @Test
-  def testRequestPercentageQuotaViolation() {
-    val metrics = newMetrics
-    val quotaManager = new ClientRequestQuotaManager(config, metrics, time, "")
-    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
+  def testQuotaViolation(): Unit = {
+    val clientQuotaManager = new ClientQuotaManager(config, metrics, Produce, time, "")
+    val queueSizeMetric = metrics.metrics().get(metrics.metricName("queue-size", "Produce", ""))
     try {
-      /* We have 10 second windows. Make sure that there is no quota violation
-       * if we are under the quota
-       */
+      // We have 10 second windows. Make sure that there is no quota violation
+      // if we produce under the quota
       for (_ <- 0 until 10) {
-        quotaManager.maybeRecordAndThrottle("ANONYMOUS", "test-client", millisToPercent(4), callback)
+        assertEquals(0, maybeRecord(clientQuotaManager, "ANONYMOUS", "unknown", 400))
         time.sleep(1000)
       }
-      assertEquals(10, numCallbacks)
-      assertEquals(0, queueSizeMetric.value().toInt)
+      assertEquals(0, queueSizeMetric.metricValue.asInstanceOf[Double].toInt)
 
       // Create a spike.
-      // quota = 1% (10ms per second)
-      // 4*10 + 67.1 = 107.1/10.5 = 10.2ms per second.
-      // (10.2 - quota)/quota*window-size = (10.2-10)/10*10.5 seconds = 210ms
-      // 10.5 seconds interval because the last window is half complete
+      // 400*10 + 2000 + 300 = 6300/10.5 = 600 bytes per second.
+      // (600 - quota)/quota*window-size = (600-500)/500*10.5 seconds = 2100
+      // 10.5 seconds because the last window is half complete
       time.sleep(500)
-      val throttleTime = quotaManager.maybeRecordAndThrottle("ANONYMOUS", "test-client", millisToPercent(67.1), callback)
+      val throttleTime = maybeRecord(clientQuotaManager, "ANONYMOUS", "unknown", 2300)
 
-      assertEquals("Should be throttled", 210, throttleTime)
-      assertEquals(1, queueSizeMetric.value().toInt)
+      assertEquals("Should be throttled", 2100, throttleTime)
+      throttle(clientQuotaManager, "ANONYMOUS", "unknown", throttleTime, callback)
+      assertEquals(1, queueSizeMetric.metricValue.asInstanceOf[Double].toInt)
       // After a request is delayed, the callback cannot be triggered immediately
-      quotaManager.throttledRequestReaper.doWork()
-      assertEquals(10, numCallbacks)
+      clientQuotaManager.throttledChannelReaper.doWork()
+      assertEquals(0, numCallbacks)
       time.sleep(throttleTime)
 
       // Callback can only be triggered after the delay time passes
-      quotaManager.throttledRequestReaper.doWork()
-      assertEquals(0, queueSizeMetric.value().toInt)
-      assertEquals(11, numCallbacks)
+      clientQuotaManager.throttledChannelReaper.doWork()
+      assertEquals(0, queueSizeMetric.metricValue.asInstanceOf[Double].toInt)
+      assertEquals(1, numCallbacks)
 
       // Could continue to see delays until the bursty sample disappears
-      for (_ <- 0 until 11) {
-        quotaManager.maybeRecordAndThrottle("ANONYMOUS", "test-client", millisToPercent(4), callback)
-        time.sleep(1000)
-      }
-
-      assertEquals("Should be unthrottled since bursty sample has rolled over",
-                   0, quotaManager.maybeRecordAndThrottle("ANONYMOUS", "test-client", 0, callback))
-
-      // Create a very large spike which requires > one quota window to bring within quota
-      assertEquals(1000, quotaManager.maybeRecordAndThrottle("ANONYMOUS", "test-client", millisToPercent(500), callback))
       for (_ <- 0 until 10) {
+        maybeRecord(clientQuotaManager, "ANONYMOUS", "unknown", 400)
         time.sleep(1000)
-        assertEquals(1000, quotaManager.maybeRecordAndThrottle("ANONYMOUS", "test-client", 0, callback))
       }
-      time.sleep(1000)
-      assertEquals("Should be unthrottled since bursty sample has rolled over",
-                   0, quotaManager.maybeRecordAndThrottle("ANONYMOUS", "test-client", 0, callback))
 
+      assertEquals("Should be unthrottled since bursty sample has rolled over",
+                   0, maybeRecord(clientQuotaManager, "ANONYMOUS", "unknown", 0))
     } finally {
-      quotaManager.shutdown()
+      clientQuotaManager.shutdown()
     }
   }
 
   @Test
-  def testExpireThrottleTimeSensor() {
-    val metrics = newMetrics
-    val clientMetrics = new ClientQuotaManager(config, metrics, QuotaType.Produce, time, "")
+  def testExpireThrottleTimeSensor(): Unit = {
+    val clientQuotaManager = new ClientQuotaManager(config, metrics, Produce, time, "")
     try {
-      clientMetrics.maybeRecordAndThrottle("ANONYMOUS", "client1", 100, callback)
+      maybeRecord(clientQuotaManager, "ANONYMOUS", "client1", 100)
       // remove the throttle time sensor
       metrics.removeSensor("ProduceThrottleTime-:client1")
       // should not throw an exception even if the throttle time sensor does not exist.
-      val throttleTime = clientMetrics.maybeRecordAndThrottle("ANONYMOUS", "client1", 10000, callback)
+      val throttleTime = maybeRecord(clientQuotaManager, "ANONYMOUS", "client1", 10000)
       assertTrue("Should be throttled", throttleTime > 0)
       // the sensor should get recreated
       val throttleTimeSensor = metrics.getSensor("ProduceThrottleTime-:client1")
       assertTrue("Throttle time sensor should exist", throttleTimeSensor != null)
+      assertTrue("Throttle time sensor should exist", throttleTimeSensor != null)
     } finally {
-      clientMetrics.shutdown()
+      clientQuotaManager.shutdown()
     }
   }
 
   @Test
-  def testExpireQuotaSensors() {
-    val metrics = newMetrics
-    val clientMetrics = new ClientQuotaManager(config, metrics, QuotaType.Produce, time, "")
+  def testExpireQuotaSensors(): Unit = {
+    val clientQuotaManager = new ClientQuotaManager(config, metrics, Produce, time, "")
     try {
-      clientMetrics.maybeRecordAndThrottle("ANONYMOUS", "client1", 100, callback)
+      maybeRecord(clientQuotaManager, "ANONYMOUS", "client1", 100)
       // remove all the sensors
       metrics.removeSensor("ProduceThrottleTime-:client1")
       metrics.removeSensor("Produce-ANONYMOUS:client1")
       // should not throw an exception
-      val throttleTime = clientMetrics.maybeRecordAndThrottle("ANONYMOUS", "client1", 10000, callback)
+      val throttleTime = maybeRecord(clientQuotaManager, "ANONYMOUS", "client1", 10000)
       assertTrue("Should be throttled", throttleTime > 0)
 
       // all the sensors should get recreated
@@ -369,17 +382,16 @@ class ClientQuotaManagerTest {
       val byteRateSensor = metrics.getSensor("Produce-:client1")
       assertTrue("Byte rate sensor should exist", byteRateSensor != null)
     } finally {
-      clientMetrics.shutdown()
+      clientQuotaManager.shutdown()
     }
   }
 
   @Test
-  def testClientIdNotSanitized() {
-    val metrics = newMetrics
-    val clientMetrics = new ClientQuotaManager(config, metrics, QuotaType.Produce, time, "")
+  def testClientIdNotSanitized(): Unit = {
+    val clientQuotaManager = new ClientQuotaManager(config, metrics, Produce, time, "")
     val clientId = "client@#$%"
     try {
-      clientMetrics.maybeRecordAndThrottle("ANONYMOUS", clientId, 100, callback)
+      maybeRecord(clientQuotaManager, "ANONYMOUS", clientId, 100)
 
       // The metrics should use the raw client ID, even if the reporters internally sanitize them
       val throttleTimeSensor = metrics.getSensor("ProduceThrottleTime-:" + clientId)
@@ -388,14 +400,10 @@ class ClientQuotaManagerTest {
       val byteRateSensor = metrics.getSensor("Produce-:"  + clientId)
       assertTrue("Byte rate sensor should exist", byteRateSensor != null)
     } finally {
-      clientMetrics.shutdown()
+      clientQuotaManager.shutdown()
     }
   }
 
-  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/ClientQuotasRequestTest.scala b/core/src/test/scala/unit/kafka/server/ClientQuotasRequestTest.scala
new file mode 100644
index 0000000000000..b72f3b9121b17
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/ClientQuotasRequestTest.scala
@@ -0,0 +1,604 @@
+/**
+  * Licensed to the Apache Software Foundation (ASF) under one or more
+  * contributor license agreements.  See the NOTICE file distributed with
+  * this work for additional information regarding copyright ownership.
+  * The ASF licenses this file to You under the Apache License, Version 2.0
+  * (the "License"); you may not use this file except in compliance with
+  * the License.  You may obtain a copy of the License at
+  *
+  *    http://www.apache.org/licenses/LICENSE-2.0
+  *
+  * Unless required by applicable law or agreed to in writing, software
+  * distributed under the License is distributed on an "AS IS" BASIS,
+  * WITHOUT WARRANTIES 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.net.InetAddress
+
+import org.apache.kafka.clients.admin.{ScramCredentialInfo, ScramMechanism, UserScramCredentialUpsertion}
+import org.apache.kafka.common.errors.{InvalidRequestException, UnsupportedVersionException}
+import org.apache.kafka.common.internals.KafkaFutureImpl
+import org.apache.kafka.common.quota.{ClientQuotaAlteration, ClientQuotaEntity, ClientQuotaFilter, ClientQuotaFilterComponent}
+import org.apache.kafka.common.requests.{AlterClientQuotasRequest, AlterClientQuotasResponse, DescribeClientQuotasRequest, DescribeClientQuotasResponse}
+import org.junit.Assert._
+import org.junit.Test
+import java.util
+import java.util.concurrent.{ExecutionException, TimeUnit}
+
+import scala.jdk.CollectionConverters._
+
+class ClientQuotasRequestTest extends BaseRequestTest {
+  private val ConsumerByteRateProp = DynamicConfig.Client.ConsumerByteRateOverrideProp
+  private val ProducerByteRateProp = DynamicConfig.Client.ProducerByteRateOverrideProp
+  private val RequestPercentageProp = DynamicConfig.Client.RequestPercentageOverrideProp
+  private val IpConnectionRateProp = DynamicConfig.Ip.IpConnectionRateOverrideProp
+
+  override val brokerCount = 1
+
+  @Test
+  def testAlterClientQuotasRequest(): Unit = {
+
+    val entity = new ClientQuotaEntity(Map((ClientQuotaEntity.USER -> "user"), (ClientQuotaEntity.CLIENT_ID -> "client-id")).asJava)
+
+    // Expect an empty configuration.
+    verifyDescribeEntityQuotas(entity, Map.empty)
+
+    // Add two configuration entries.
+    alterEntityQuotas(entity, Map(
+      (ProducerByteRateProp -> Some(10000.0)),
+      (ConsumerByteRateProp -> Some(20000.0))
+    ), validateOnly = false)
+
+    verifyDescribeEntityQuotas(entity, Map(
+      (ProducerByteRateProp -> 10000.0),
+      (ConsumerByteRateProp -> 20000.0)
+    ))
+
+    // Update an existing entry.
+    alterEntityQuotas(entity, Map(
+      (ProducerByteRateProp -> Some(15000.0))
+    ), validateOnly = false)
+
+    verifyDescribeEntityQuotas(entity, Map(
+      (ProducerByteRateProp -> 15000.0),
+      (ConsumerByteRateProp -> 20000.0)
+    ))
+
+    // Remove an existing configuration entry.
+    alterEntityQuotas(entity, Map(
+      (ProducerByteRateProp -> None)
+    ), validateOnly = false)
+
+    verifyDescribeEntityQuotas(entity, Map(
+      (ConsumerByteRateProp -> 20000.0)
+    ))
+
+    // Remove a non-existent configuration entry.  This should make no changes.
+    alterEntityQuotas(entity, Map(
+      (RequestPercentageProp -> None)
+    ), validateOnly = false)
+
+    verifyDescribeEntityQuotas(entity, Map(
+      (ConsumerByteRateProp -> 20000.0)
+    ))
+
+    // Add back a deleted configuration entry.
+    alterEntityQuotas(entity, Map(
+      (ProducerByteRateProp -> Some(5000.0))
+    ), validateOnly = false)
+
+    verifyDescribeEntityQuotas(entity, Map(
+      (ProducerByteRateProp -> 5000.0),
+      (ConsumerByteRateProp -> 20000.0)
+    ))
+
+    // Perform a mixed update.
+    alterEntityQuotas(entity, Map(
+      (ProducerByteRateProp -> Some(20000.0)),
+      (ConsumerByteRateProp -> None),
+      (RequestPercentageProp -> Some(12.3))
+    ), validateOnly = false)
+
+    verifyDescribeEntityQuotas(entity, Map(
+      (ProducerByteRateProp -> 20000.0),
+      (RequestPercentageProp -> 12.3)
+    ))
+  }
+
+  @Test
+  def testAlterClientQuotasRequestValidateOnly(): Unit = {
+    val entity = new ClientQuotaEntity(Map((ClientQuotaEntity.USER -> "user")).asJava)
+
+    // Set up a configuration.
+    alterEntityQuotas(entity, Map(
+      (ProducerByteRateProp -> Some(20000.0)),
+      (RequestPercentageProp -> Some(23.45))
+    ), validateOnly = false)
+
+    verifyDescribeEntityQuotas(entity, Map(
+      (ProducerByteRateProp -> 20000.0),
+      (RequestPercentageProp -> 23.45)
+    ))
+
+    // Validate-only addition.
+    alterEntityQuotas(entity, Map(
+      (ConsumerByteRateProp -> Some(50000.0))
+    ), validateOnly = true)
+
+    verifyDescribeEntityQuotas(entity, Map(
+      (ProducerByteRateProp -> 20000.0),
+      (RequestPercentageProp -> 23.45)
+    ))
+
+    // Validate-only modification.
+    alterEntityQuotas(entity, Map(
+      (ProducerByteRateProp -> Some(10000.0))
+    ), validateOnly = true)
+
+    verifyDescribeEntityQuotas(entity, Map(
+      (ProducerByteRateProp -> 20000.0),
+      (RequestPercentageProp -> 23.45)
+    ))
+
+    // Validate-only removal.
+    alterEntityQuotas(entity, Map(
+      (RequestPercentageProp -> None)
+    ), validateOnly = true)
+
+    verifyDescribeEntityQuotas(entity, Map(
+      (ProducerByteRateProp -> 20000.0),
+      (RequestPercentageProp -> 23.45)
+    ))
+
+    // Validate-only mixed update.
+    alterEntityQuotas(entity, Map(
+      (ProducerByteRateProp -> Some(10000.0)),
+      (ConsumerByteRateProp -> Some(50000.0)),
+      (RequestPercentageProp -> None)
+    ), validateOnly = true)
+
+    verifyDescribeEntityQuotas(entity, Map(
+      (ProducerByteRateProp -> 20000.0),
+      (RequestPercentageProp -> 23.45)
+    ))
+  }
+
+  @Test
+  def testClientQuotasForScramUsers(): Unit = {
+    val userName = "user"
+
+    val results = createAdminClient().alterUserScramCredentials(util.Arrays.asList(
+      new UserScramCredentialUpsertion(userName, new ScramCredentialInfo(ScramMechanism.SCRAM_SHA_256, 4096), "password")))
+    results.all.get
+
+    val entity = new ClientQuotaEntity(Map(ClientQuotaEntity.USER -> userName).asJava)
+
+    verifyDescribeEntityQuotas(entity, Map.empty)
+
+    alterEntityQuotas(entity, Map(
+      (ProducerByteRateProp -> Some(10000.0)),
+      (ConsumerByteRateProp -> Some(20000.0))
+    ), validateOnly = false)
+
+    verifyDescribeEntityQuotas(entity, Map(
+      (ProducerByteRateProp -> 10000.0),
+      (ConsumerByteRateProp -> 20000.0)
+    ))
+  }
+
+  @Test
+  def testAlterIpQuotasRequest(): Unit = {
+    val knownHost = "1.2.3.4"
+    val unknownHost = "2.3.4.5"
+    val entity = toIpEntity(Some(knownHost))
+    val defaultEntity = toIpEntity(Some(null))
+    val entityFilter = ClientQuotaFilterComponent.ofEntity(ClientQuotaEntity.IP, knownHost)
+    val defaultEntityFilter = ClientQuotaFilterComponent.ofDefaultEntity(ClientQuotaEntity.IP)
+    val allIpEntityFilter = ClientQuotaFilterComponent.ofEntityType(ClientQuotaEntity.IP)
+
+    def verifyIpQuotas(entityFilter: ClientQuotaFilterComponent, expectedMatches: Map[ClientQuotaEntity, Double]): Unit = {
+      val result = describeClientQuotas(ClientQuotaFilter.containsOnly(List(entityFilter).asJava))
+      assertEquals(expectedMatches.keySet, result.asScala.keySet)
+      result.asScala.foreach { case (entity, props) =>
+        assertEquals(Set(IpConnectionRateProp), props.asScala.keySet)
+        assertEquals(expectedMatches(entity), props.get(IpConnectionRateProp))
+        val entityName = entity.entries.get(ClientQuotaEntity.IP)
+        // ClientQuotaEntity with null name maps to default entity
+        val entityIp = if (entityName == null)
+          InetAddress.getByName(unknownHost)
+        else
+          InetAddress.getByName(entityName)
+        assertEquals(expectedMatches(entity), servers.head.socketServer.connectionQuotas.connectionRateForIp(entityIp), 0.01)
+      }
+    }
+
+    // Expect an empty configuration.
+    verifyIpQuotas(allIpEntityFilter, Map.empty)
+
+    // Add a configuration entry.
+    alterEntityQuotas(entity, Map(IpConnectionRateProp -> Some(100.0)), validateOnly = false)
+    verifyIpQuotas(entityFilter, Map(entity -> 100.0))
+
+    // update existing entry
+    alterEntityQuotas(entity, Map(IpConnectionRateProp -> Some(150.0)), validateOnly = false)
+    verifyIpQuotas(entityFilter, Map(entity -> 150.0))
+
+    // update default value
+    alterEntityQuotas(defaultEntity, Map(IpConnectionRateProp -> Some(200.0)), validateOnly = false)
+    verifyIpQuotas(defaultEntityFilter, Map(defaultEntity -> 200.0))
+
+    // describe all IP quotas
+    verifyIpQuotas(allIpEntityFilter, Map(entity -> 150.0, defaultEntity -> 200.0))
+
+    // remove entry
+    alterEntityQuotas(entity, Map(IpConnectionRateProp -> None), validateOnly = false)
+    verifyIpQuotas(entityFilter, Map.empty)
+
+    // remove default value
+    alterEntityQuotas(defaultEntity, Map(IpConnectionRateProp -> None), validateOnly = false)
+    verifyIpQuotas(allIpEntityFilter, Map.empty)
+  }
+
+  @Test(expected = classOf[InvalidRequestException])
+  def testAlterClientQuotasBadUser(): Unit = {
+    val entity = new ClientQuotaEntity(Map((ClientQuotaEntity.USER -> "")).asJava)
+    alterEntityQuotas(entity, Map((RequestPercentageProp -> Some(12.34))), validateOnly = true)
+  }
+
+  @Test(expected = classOf[InvalidRequestException])
+  def testAlterClientQuotasBadClientId(): Unit = {
+    val entity = new ClientQuotaEntity(Map((ClientQuotaEntity.CLIENT_ID -> "")).asJava)
+    alterEntityQuotas(entity, Map((RequestPercentageProp -> Some(12.34))), validateOnly = true)
+  }
+
+  @Test(expected = classOf[InvalidRequestException])
+  def testAlterClientQuotasBadEntityType(): Unit = {
+    val entity = new ClientQuotaEntity(Map(("" -> "name")).asJava)
+    alterEntityQuotas(entity, Map((RequestPercentageProp -> Some(12.34))), validateOnly = true)
+  }
+
+  @Test(expected = classOf[InvalidRequestException])
+  def testAlterClientQuotasEmptyEntity(): Unit = {
+    val entity = new ClientQuotaEntity(Map.empty.asJava)
+    alterEntityQuotas(entity, Map((ProducerByteRateProp -> Some(10000.5))), validateOnly = true)
+  }
+
+  @Test(expected = classOf[InvalidRequestException])
+  def testAlterClientQuotasBadConfigKey(): Unit = {
+    val entity = new ClientQuotaEntity(Map((ClientQuotaEntity.USER -> "user")).asJava)
+    alterEntityQuotas(entity, Map(("bad" -> Some(1.0))), validateOnly = true)
+  }
+
+  @Test(expected = classOf[InvalidRequestException])
+  def testAlterClientQuotasBadConfigValue(): Unit = {
+    val entity = new ClientQuotaEntity(Map((ClientQuotaEntity.USER -> "user")).asJava)
+    alterEntityQuotas(entity, Map((ProducerByteRateProp -> Some(10000.5))), validateOnly = true)
+  }
+
+  private def expectInvalidRequestWithMessage(runnable: => Unit, expectedMessage: String): Unit = {
+    val exception = assertThrows(classOf[InvalidRequestException], () => runnable)
+    assertTrue(s"Expected message $exception to contain $expectedMessage", exception.getMessage.contains(expectedMessage))
+  }
+
+  @Test
+  def testAlterClientQuotasInvalidEntityCombination(): Unit = {
+    val userAndIpEntity = new ClientQuotaEntity(Map(ClientQuotaEntity.USER -> "user", ClientQuotaEntity.IP -> "1.2.3.4").asJava)
+    val clientAndIpEntity = new ClientQuotaEntity(Map(ClientQuotaEntity.CLIENT_ID -> "client", ClientQuotaEntity.IP -> "1.2.3.4").asJava)
+    val expectedExceptionMessage = "Invalid quota entity combination"
+    expectInvalidRequestWithMessage(alterEntityQuotas(userAndIpEntity, Map(RequestPercentageProp -> Some(12.34)),
+      validateOnly = true), expectedExceptionMessage)
+    expectInvalidRequestWithMessage(alterEntityQuotas(clientAndIpEntity, Map(RequestPercentageProp -> Some(12.34)),
+      validateOnly = true), expectedExceptionMessage)
+  }
+
+  @Test
+  def testAlterClientQuotasBadIp(): Unit = {
+    val invalidHostPatternEntity = new ClientQuotaEntity(Map(ClientQuotaEntity.IP -> "abc-123").asJava)
+    val unresolvableHostEntity = new ClientQuotaEntity(Map(ClientQuotaEntity.IP -> "ip").asJava)
+    val expectedExceptionMessage = "not a valid IP"
+    expectInvalidRequestWithMessage(alterEntityQuotas(invalidHostPatternEntity, Map(IpConnectionRateProp -> Some(50.0)),
+      validateOnly = true), expectedExceptionMessage)
+    expectInvalidRequestWithMessage(alterEntityQuotas(unresolvableHostEntity, Map(IpConnectionRateProp -> Some(50.0)),
+      validateOnly = true), expectedExceptionMessage)
+  }
+
+  @Test
+  def testDescribeClientQuotasInvalidFilterCombination(): Unit = {
+    val ipFilterComponent = ClientQuotaFilterComponent.ofEntityType(ClientQuotaEntity.IP)
+    val userFilterComponent = ClientQuotaFilterComponent.ofEntityType(ClientQuotaEntity.USER)
+    val clientIdFilterComponent = ClientQuotaFilterComponent.ofEntityType(ClientQuotaEntity.CLIENT_ID)
+    val expectedExceptionMessage = "Invalid entity filter component combination"
+    expectInvalidRequestWithMessage(describeClientQuotas(ClientQuotaFilter.contains(List(ipFilterComponent, userFilterComponent).asJava)),
+      expectedExceptionMessage)
+    expectInvalidRequestWithMessage(describeClientQuotas(ClientQuotaFilter.contains(List(ipFilterComponent, clientIdFilterComponent).asJava)),
+      expectedExceptionMessage)
+  }
+
+  // Entities to be matched against.
+  private val matchUserClientEntities = List(
+    (Some("user-1"), Some("client-id-1"), 50.50),
+    (Some("user-2"), Some("client-id-1"), 51.51),
+    (Some("user-3"), Some("client-id-2"), 52.52),
+    (Some(null), Some("client-id-1"), 53.53),
+    (Some("user-1"), Some(null), 54.54),
+    (Some("user-3"), Some(null), 55.55),
+    (Some("user-1"), None, 56.56),
+    (Some("user-2"), None, 57.57),
+    (Some("user-3"), None, 58.58),
+    (Some(null), None, 59.59),
+    (None, Some("client-id-2"), 60.60)
+  ).map { case (u, c, v) => (toClientEntity(u, c), v) }
+
+  private val matchIpEntities = List(
+    (Some("1.2.3.4"), 10.0),
+    (Some("2.3.4.5"), 20.0)
+  ).map { case (ip, quota) => (toIpEntity(ip), quota)}
+
+  private def setupDescribeClientQuotasMatchTest() = {
+    val userClientQuotas = matchUserClientEntities.map { case (e, v) =>
+      e -> Map((RequestPercentageProp, Some(v)))
+    }.toMap
+    val ipQuotas = matchIpEntities.map { case (e, v) =>
+      e -> Map((IpConnectionRateProp, Some(v)))
+    }.toMap
+    val result = alterClientQuotas(userClientQuotas ++ ipQuotas, validateOnly = false)
+    (matchUserClientEntities ++ matchIpEntities).foreach(e => result(e._1).get(10, TimeUnit.SECONDS))
+  }
+
+  @Test
+  def testDescribeClientQuotasMatchExact(): Unit = {
+    setupDescribeClientQuotasMatchTest()
+
+    def matchEntity(entity: ClientQuotaEntity) = {
+      val components = entity.entries.asScala.map { case (entityType, entityName) =>
+        entityName match {
+          case null => ClientQuotaFilterComponent.ofDefaultEntity(entityType)
+          case name => ClientQuotaFilterComponent.ofEntity(entityType, name)
+        }
+      }
+      describeClientQuotas(ClientQuotaFilter.containsOnly(components.toList.asJava))
+    }
+
+    // Test exact matches.
+    matchUserClientEntities.foreach { case (e, v) =>
+      val result = matchEntity(e)
+      assertEquals(1, result.size)
+      assertTrue(result.get(e) != null)
+      val value = result.get(e).get(RequestPercentageProp)
+      assertTrue(value != null)
+      assertEquals(value, v, 1e-6)
+    }
+
+    // Entities not contained in `matchEntityList`.
+    val notMatchEntities = List(
+      (Some("user-1"), Some("client-id-2")),
+      (Some("user-3"), Some("client-id-1")),
+      (Some("user-2"), Some(null)),
+      (Some("user-4"), None),
+      (Some(null), Some("client-id-2")),
+      (None, Some("client-id-1")),
+      (None, Some("client-id-3")),
+    ).map { case (u, c) =>
+        new ClientQuotaEntity((u.map((ClientQuotaEntity.USER, _)) ++
+          c.map((ClientQuotaEntity.CLIENT_ID, _))).toMap.asJava)
+    }
+
+    // Verify exact matches of the non-matches returns empty.
+    notMatchEntities.foreach { e =>
+      val result = matchEntity(e)
+      assertEquals(0, result.size)
+    }
+  }
+
+  @Test
+  def testDescribeClientQuotasMatchPartial(): Unit = {
+    setupDescribeClientQuotasMatchTest()
+
+    def testMatchEntities(filter: ClientQuotaFilter, expectedMatchSize: Int, partition: ClientQuotaEntity => Boolean): Unit = {
+      val result = describeClientQuotas(filter)
+      val (expectedMatches, _) = (matchUserClientEntities ++ matchIpEntities).partition(e => partition(e._1))
+      assertEquals(expectedMatchSize, expectedMatches.size)  // for test verification
+      assertEquals(expectedMatchSize, result.size)
+      val expectedMatchesMap = expectedMatches.toMap
+      matchUserClientEntities.foreach { case (entity, expectedValue) =>
+        if (expectedMatchesMap.contains(entity)) {
+          val config = result.get(entity)
+          assertNotNull(config)
+          val value = config.get(RequestPercentageProp)
+          assertNotNull(value)
+          assertEquals(expectedValue, value, 1e-6)
+        } else {
+          assertNull(result.get(entity))
+        }
+      }
+      matchIpEntities.foreach { case (entity, expectedValue) =>
+        if (expectedMatchesMap.contains(entity)) {
+          val config = result.get(entity)
+          assertNotNull(config)
+          val value = config.get(IpConnectionRateProp)
+          assertNotNull(value)
+          assertEquals(expectedValue, value, 1e-6)
+        } else {
+          assertNull(result.get(entity))
+        }
+      }
+    }
+
+    // Match open-ended existing user.
+    testMatchEntities(
+      ClientQuotaFilter.contains(List(ClientQuotaFilterComponent.ofEntity(ClientQuotaEntity.USER, "user-1")).asJava), 3,
+      entity => entity.entries.get(ClientQuotaEntity.USER) == "user-1"
+    )
+
+    // Match open-ended non-existent user.
+    testMatchEntities(
+      ClientQuotaFilter.contains(List(ClientQuotaFilterComponent.ofEntity(ClientQuotaEntity.USER, "unknown")).asJava), 0,
+      entity => false
+    )
+
+    // Match open-ended existing client ID.
+    testMatchEntities(
+      ClientQuotaFilter.contains(List(ClientQuotaFilterComponent.ofEntity(ClientQuotaEntity.CLIENT_ID, "client-id-2")).asJava), 2,
+      entity => entity.entries.get(ClientQuotaEntity.CLIENT_ID) == "client-id-2"
+    )
+
+    // Match open-ended default user.
+    testMatchEntities(
+      ClientQuotaFilter.contains(List(ClientQuotaFilterComponent.ofDefaultEntity(ClientQuotaEntity.USER)).asJava), 2,
+      entity => entity.entries.containsKey(ClientQuotaEntity.USER) && entity.entries.get(ClientQuotaEntity.USER) == null
+    )
+
+    // Match close-ended existing user.
+    testMatchEntities(
+      ClientQuotaFilter.containsOnly(List(ClientQuotaFilterComponent.ofEntity(ClientQuotaEntity.USER, "user-2")).asJava), 1,
+      entity => entity.entries.get(ClientQuotaEntity.USER) == "user-2" && !entity.entries.containsKey(ClientQuotaEntity.CLIENT_ID)
+    )
+
+    // Match close-ended existing client ID that has no matching entity.
+    testMatchEntities(
+      ClientQuotaFilter.containsOnly(List(ClientQuotaFilterComponent.ofEntity(ClientQuotaEntity.CLIENT_ID, "client-id-1")).asJava), 0,
+      entity => false
+    )
+
+    // Match against all entities with the user type in a close-ended match.
+    testMatchEntities(
+      ClientQuotaFilter.containsOnly(List(ClientQuotaFilterComponent.ofEntityType(ClientQuotaEntity.USER)).asJava), 4,
+      entity => entity.entries.containsKey(ClientQuotaEntity.USER) && !entity.entries.containsKey(ClientQuotaEntity.CLIENT_ID)
+    )
+
+    // Match against all entities with the user type in an open-ended match.
+    testMatchEntities(
+      ClientQuotaFilter.contains(List(ClientQuotaFilterComponent.ofEntityType(ClientQuotaEntity.USER)).asJava), 10,
+      entity => entity.entries.containsKey(ClientQuotaEntity.USER)
+    )
+
+    // Match against all entities with the client ID type in a close-ended match.
+    testMatchEntities(
+      ClientQuotaFilter.containsOnly(List(ClientQuotaFilterComponent.ofEntityType(ClientQuotaEntity.CLIENT_ID)).asJava), 1,
+      entity => entity.entries.containsKey(ClientQuotaEntity.CLIENT_ID) && !entity.entries.containsKey(ClientQuotaEntity.USER)
+    )
+
+    // Match against all entities with the client ID type in an open-ended match.
+    testMatchEntities(
+      ClientQuotaFilter.contains(List(ClientQuotaFilterComponent.ofEntityType(ClientQuotaEntity.CLIENT_ID)).asJava),  7,
+      entity => entity.entries.containsKey(ClientQuotaEntity.CLIENT_ID)
+    )
+
+    // Match against all entities with IP type in an open-ended match.
+    testMatchEntities(
+      ClientQuotaFilter.contains(List(ClientQuotaFilterComponent.ofEntityType(ClientQuotaEntity.IP)).asJava), 2,
+      entity => entity.entries.containsKey(ClientQuotaEntity.IP)
+    )
+
+    // Match open-ended empty filter list. This should match all entities.
+    testMatchEntities(ClientQuotaFilter.contains(List.empty.asJava), 13, entity => true)
+
+    // Match close-ended empty filter list. This should match no entities.
+    testMatchEntities(ClientQuotaFilter.containsOnly(List.empty.asJava), 0, entity => false)
+  }
+
+  @Test
+  def testClientQuotasUnsupportedEntityTypes(): Unit = {
+    val entity = new ClientQuotaEntity(Map(("other" -> "name")).asJava)
+    assertThrows(classOf[UnsupportedVersionException], () => verifyDescribeEntityQuotas(entity, Map.empty))
+  }
+
+  @Test
+  def testClientQuotasSanitized(): Unit = {
+    // An entity with name that must be sanitized when writing to Zookeeper.
+    val entity = new ClientQuotaEntity(Map((ClientQuotaEntity.USER -> "user with spaces")).asJava)
+
+    alterEntityQuotas(entity, Map(
+      (ProducerByteRateProp -> Some(20000.0)),
+    ), validateOnly = false)
+
+    verifyDescribeEntityQuotas(entity, Map(
+      (ProducerByteRateProp -> 20000.0),
+    ))
+  }
+
+  @Test
+  def testClientQuotasWithDefaultName(): Unit = {
+    // An entity using the name associated with the default entity name. The entity's name should be sanitized so
+    // that it does not conflict with the default entity name.
+    val entity = new ClientQuotaEntity(Map((ClientQuotaEntity.CLIENT_ID -> ConfigEntityName.Default)).asJava)
+    alterEntityQuotas(entity, Map((ProducerByteRateProp -> Some(20000.0))), validateOnly = false)
+    verifyDescribeEntityQuotas(entity, Map((ProducerByteRateProp -> 20000.0)))
+
+    // This should not match.
+    val result = describeClientQuotas(
+      ClientQuotaFilter.containsOnly(List(ClientQuotaFilterComponent.ofDefaultEntity(ClientQuotaEntity.CLIENT_ID)).asJava))
+    assert(result.isEmpty)
+  }
+
+  private def verifyDescribeEntityQuotas(entity: ClientQuotaEntity, quotas: Map[String, Double]) = {
+    val components = entity.entries.asScala.map { case (entityType, entityName) =>
+      Option(entityName).map{ name => ClientQuotaFilterComponent.ofEntity(entityType, name)}
+        .getOrElse(ClientQuotaFilterComponent.ofDefaultEntity(entityType)
+      )
+    }
+    val describe = describeClientQuotas(ClientQuotaFilter.containsOnly(components.toList.asJava))
+    if (quotas.isEmpty) {
+      assertEquals(0, describe.size)
+    } else {
+      assertEquals(1, describe.size)
+      val configs = describe.get(entity)
+      assertTrue(configs != null)
+      assertEquals(quotas.size, configs.size)
+      quotas.foreach { case (k, v) =>
+        val value = configs.get(k)
+        assertTrue(value != null)
+        assertEquals(v, value, 1e-6)
+      }
+    }
+  }
+
+  private def toClientEntity(user: Option[String], clientId: Option[String]) =
+    new ClientQuotaEntity((user.map((ClientQuotaEntity.USER -> _)) ++ clientId.map((ClientQuotaEntity.CLIENT_ID -> _))).toMap.asJava)
+
+  private def toIpEntity(ip: Option[String]) = new ClientQuotaEntity(ip.map(ClientQuotaEntity.IP -> _).toMap.asJava)
+
+  private def describeClientQuotas(filter: ClientQuotaFilter) = {
+    val result = new KafkaFutureImpl[java.util.Map[ClientQuotaEntity, java.util.Map[String, java.lang.Double]]]
+    sendDescribeClientQuotasRequest(filter).complete(result)
+    try result.get catch {
+      case e: ExecutionException => throw e.getCause
+    }
+  }
+
+  private def sendDescribeClientQuotasRequest(filter: ClientQuotaFilter): DescribeClientQuotasResponse = {
+    val request = new DescribeClientQuotasRequest.Builder(filter).build()
+    connectAndReceive[DescribeClientQuotasResponse](request, destination = controllerSocketServer)
+  }
+
+  private def alterEntityQuotas(entity: ClientQuotaEntity, alter: Map[String, Option[Double]], validateOnly: Boolean) =
+    try alterClientQuotas(Map(entity -> alter), validateOnly).get(entity).get.get(10, TimeUnit.SECONDS) catch {
+      case e: ExecutionException => throw e.getCause
+    }
+
+  private def alterClientQuotas(request: Map[ClientQuotaEntity, Map[String, Option[Double]]], validateOnly: Boolean) = {
+    val entries = request.map { case (entity, alter) =>
+      val ops = alter.map { case (key, value) =>
+        new ClientQuotaAlteration.Op(key, value.map(Double.box).getOrElse(null))
+      }.asJavaCollection
+      new ClientQuotaAlteration(entity, ops)
+    }
+
+    val response = request.map(e => (e._1 -> new KafkaFutureImpl[Void])).asJava
+    sendAlterClientQuotasRequest(entries, validateOnly).complete(response)
+    val result = response.asScala
+    assertEquals(request.size, result.size)
+    request.foreach(e => assertTrue(result.get(e._1).isDefined))
+    result
+  }
+
+  private def sendAlterClientQuotasRequest(entries: Iterable[ClientQuotaAlteration], validateOnly: Boolean): AlterClientQuotasResponse = {
+    val request = new AlterClientQuotasRequest.Builder(entries.asJavaCollection, validateOnly).build()
+    connectAndReceive[AlterClientQuotasResponse](request, destination = controllerSocketServer)
+  }
+
+}
diff --git a/core/src/test/scala/unit/kafka/server/ClientRequestQuotaManagerTest.scala b/core/src/test/scala/unit/kafka/server/ClientRequestQuotaManagerTest.scala
new file mode 100644
index 0000000000000..a3128befc4a0f
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/ClientRequestQuotaManagerTest.scala
@@ -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 kafka.server
+
+import kafka.server.QuotaType.Request
+import org.apache.kafka.common.metrics.Quota
+
+import org.junit.Assert._
+import org.junit.Test
+
+class ClientRequestQuotaManagerTest extends BaseClientQuotaManagerTest {
+  private val config = ClientQuotaManagerConfig()
+
+  @Test
+  def testRequestPercentageQuotaViolation(): Unit = {
+    val clientRequestQuotaManager = new ClientRequestQuotaManager(config, metrics, time, "", None)
+    clientRequestQuotaManager.updateQuota(Some("ANONYMOUS"), Some("test-client"), Some("test-client"), Some(Quota.upperBound(1)))
+    val queueSizeMetric = metrics.metrics().get(metrics.metricName("queue-size", Request.toString, ""))
+    def millisToPercent(millis: Double) = millis * 1000 * 1000 * ClientRequestQuotaManager.NanosToPercentagePerSecond
+    try {
+      // We have 10 second windows. Make sure that there is no quota violation
+      // if we are under the quota
+      for (_ <- 0 until 10) {
+        assertEquals(0, maybeRecord(clientRequestQuotaManager, "ANONYMOUS", "test-client", millisToPercent(4)))
+        time.sleep(1000)
+      }
+      assertEquals(0, queueSizeMetric.metricValue.asInstanceOf[Double].toInt)
+
+      // Create a spike.
+      // quota = 1% (10ms per second)
+      // 4*10 + 67.1 = 107.1/10.5 = 10.2ms per second.
+      // (10.2 - quota)/quota*window-size = (10.2-10)/10*10.5 seconds = 210ms
+      // 10.5 seconds interval because the last window is half complete
+      time.sleep(500)
+      val throttleTime = maybeRecord(clientRequestQuotaManager, "ANONYMOUS", "test-client", millisToPercent(67.1))
+
+      assertEquals("Should be throttled", 210, throttleTime)
+
+      throttle(clientRequestQuotaManager, "ANONYMOUS", "test-client", throttleTime, callback)
+      assertEquals(1, queueSizeMetric.metricValue.asInstanceOf[Double].toInt)
+      // After a request is delayed, the callback cannot be triggered immediately
+      clientRequestQuotaManager.throttledChannelReaper.doWork()
+      assertEquals(0, numCallbacks)
+      time.sleep(throttleTime)
+
+      // Callback can only be triggered after the delay time passes
+      clientRequestQuotaManager.throttledChannelReaper.doWork()
+      assertEquals(0, queueSizeMetric.metricValue.asInstanceOf[Double].toInt)
+      assertEquals(1, numCallbacks)
+
+      // Could continue to see delays until the bursty sample disappears
+      for (_ <- 0 until 11) {
+        maybeRecord(clientRequestQuotaManager, "ANONYMOUS", "test-client", millisToPercent(4))
+        time.sleep(1000)
+      }
+
+      assertEquals("Should be unthrottled since bursty sample has rolled over",
+        0, maybeRecord(clientRequestQuotaManager, "ANONYMOUS", "test-client", 0))
+
+      // Create a very large spike which requires > one quota window to bring within quota
+      assertEquals(1000, maybeRecord(clientRequestQuotaManager, "ANONYMOUS", "test-client", millisToPercent(500)))
+      for (_ <- 0 until 10) {
+        time.sleep(1000)
+        assertEquals(1000, maybeRecord(clientRequestQuotaManager, "ANONYMOUS", "test-client", 0))
+      }
+      time.sleep(1000)
+      assertEquals("Should be unthrottled since bursty sample has rolled over",
+        0, maybeRecord(clientRequestQuotaManager, "ANONYMOUS", "test-client", 0))
+
+    } finally {
+      clientRequestQuotaManager.shutdown()
+    }
+  }
+}
+
diff --git a/core/src/test/scala/unit/kafka/server/ControllerMutationQuotaManagerTest.scala b/core/src/test/scala/unit/kafka/server/ControllerMutationQuotaManagerTest.scala
new file mode 100644
index 0000000000000..08b142426be01
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/ControllerMutationQuotaManagerTest.scala
@@ -0,0 +1,236 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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.server.QuotaType.ControllerMutation
+import org.apache.kafka.common.errors.ThrottlingQuotaExceededException
+import org.apache.kafka.common.metrics.MetricConfig
+import org.apache.kafka.common.metrics.Metrics
+import org.apache.kafka.common.metrics.Quota
+import org.apache.kafka.common.metrics.QuotaViolationException
+import org.apache.kafka.common.metrics.stats.TokenBucket
+import org.apache.kafka.common.utils.MockTime
+import org.junit.Assert._
+import org.junit.Assert.assertEquals
+import org.junit.Assert.assertFalse
+import org.junit.Test
+
+class StrictControllerMutationQuotaTest {
+  @Test
+  def testControllerMutationQuotaViolation(): Unit = {
+    val time = new MockTime(0, System.currentTimeMillis, 0)
+    val metrics = new Metrics(time)
+    val sensor = metrics.sensor("sensor", new MetricConfig()
+      .quota(Quota.upperBound(10))
+      .timeWindow(1, TimeUnit.SECONDS)
+      .samples(10))
+    val metricName = metrics.metricName("rate", "test-group")
+    assertTrue(sensor.add(metricName, new TokenBucket))
+
+    val quota = new StrictControllerMutationQuota(time, sensor)
+    assertFalse(quota.isExceeded)
+
+    // Recording a first value at T to bring the tokens to 10. Value is accepted
+    // because the quota is not exhausted yet.
+    quota.record(90)
+    assertFalse(quota.isExceeded)
+    assertEquals(0, quota.throttleTime)
+
+    // Recording a second value at T to bring the tokens to -80. Value is accepted
+    quota.record(90)
+    assertFalse(quota.isExceeded)
+    assertEquals(0, quota.throttleTime)
+
+    // Recording a third value at T is rejected immediately because there are not
+    // tokens available in the bucket.
+    assertThrows(classOf[ThrottlingQuotaExceededException],
+      () => quota.record(90))
+    assertTrue(quota.isExceeded)
+    assertEquals(8000, quota.throttleTime)
+
+    // Throttle time is adjusted with time
+    time.sleep(5000)
+    assertEquals(3000, quota.throttleTime)
+
+    metrics.close()
+  }
+}
+
+class PermissiveControllerMutationQuotaTest {
+  @Test
+  def testControllerMutationQuotaViolation(): Unit = {
+    val time = new MockTime(0, System.currentTimeMillis, 0)
+    val metrics = new Metrics(time)
+    val sensor = metrics.sensor("sensor", new MetricConfig()
+      .quota(Quota.upperBound(10))
+      .timeWindow(1, TimeUnit.SECONDS)
+      .samples(10))
+    val metricName = metrics.metricName("rate", "test-group")
+    assertTrue(sensor.add(metricName, new TokenBucket))
+
+    val quota = new PermissiveControllerMutationQuota(time, sensor)
+    assertFalse(quota.isExceeded)
+
+    // Recording a first value at T to bring the tokens 10. Value is accepted
+    // because the quota is not exhausted yet.
+    quota.record(90)
+    assertFalse(quota.isExceeded)
+    assertEquals(0, quota.throttleTime)
+
+    // Recording a second value at T to bring the tokens to -80. Value is accepted
+    quota.record(90)
+    assertFalse(quota.isExceeded)
+    assertEquals(8000, quota.throttleTime)
+
+    // Recording a second value at T to bring the tokens to -170. Value is accepted
+    // even though the quota is exhausted.
+    quota.record(90)
+    assertFalse(quota.isExceeded) // quota is never exceeded
+    assertEquals(17000, quota.throttleTime)
+
+    // Throttle time is adjusted with time
+    time.sleep(5000)
+    assertEquals(12000, quota.throttleTime)
+
+    metrics.close()
+  }
+}
+
+class ControllerMutationQuotaManagerTest extends BaseClientQuotaManagerTest {
+  private val User = "ANONYMOUS"
+  private val ClientId = "test-client"
+
+  private val config = ClientQuotaManagerConfig(
+    numQuotaSamples = 10,
+    quotaWindowSizeSeconds = 1
+  )
+
+  private def withQuotaManager(f: ControllerMutationQuotaManager => Unit): Unit = {
+    val quotaManager = new ControllerMutationQuotaManager(config, metrics, time,"", None)
+    try {
+      f(quotaManager)
+    } finally {
+      quotaManager.shutdown()
+    }
+  }
+
+  @Test
+  def testThrottleTime(): Unit = {
+    import ControllerMutationQuotaManager._
+
+    val time = new MockTime(0, System.currentTimeMillis, 0)
+    val metrics = new Metrics(time)
+    val sensor = metrics.sensor("sensor")
+    val metricName = metrics.metricName("tokens", "test-group")
+    sensor.add(metricName, new TokenBucket)
+    val metric = metrics.metric(metricName)
+
+    assertEquals(0, throttleTimeMs(new QuotaViolationException(metric, 0, 10), time.milliseconds()))
+    assertEquals(500, throttleTimeMs(new QuotaViolationException(metric, -5, 10), time.milliseconds()))
+    assertEquals(1000, throttleTimeMs(new QuotaViolationException(metric, -10, 10), time.milliseconds()))
+  }
+
+  @Test
+  def testControllerMutationQuotaViolation(): Unit = {
+    withQuotaManager { quotaManager =>
+      quotaManager.updateQuota(Some(User), Some(ClientId), Some(ClientId),
+        Some(Quota.upperBound(10)))
+      val queueSizeMetric = metrics.metrics().get(
+        metrics.metricName("queue-size", ControllerMutation.toString, ""))
+
+      // Verify that there is no quota violation if we remain under the quota.
+      for (_ <- 0 until 10) {
+        assertEquals(0, maybeRecord(quotaManager, User, ClientId, 10))
+        time.sleep(1000)
+      }
+      assertEquals(0, queueSizeMetric.metricValue.asInstanceOf[Double].toInt)
+
+      // Create a spike worth of 110 mutations.
+      // Current tokens in the bucket = 100
+      // As we use the Strict enforcement, the quota is checked before updating the rate. Hence,
+      // the spike is accepted and no quota violation error is raised.
+      var throttleTime = maybeRecord(quotaManager, User, ClientId, 110)
+      assertEquals("Should not be throttled", 0, throttleTime)
+
+      // Create a spike worth of 110 mutations.
+      // Current tokens in the bucket = 100 - 110 = -10
+      // As the quota is already violated, the spike is rejected immediately without updating the
+      // rate. The client must wait:
+      // 10 / 10 = 1s
+      throttleTime = maybeRecord(quotaManager, User, ClientId, 110)
+      assertEquals("Should be throttled", 1000, throttleTime)
+
+      // Throttle
+      throttle(quotaManager, User, ClientId, throttleTime, callback)
+      assertEquals(1, queueSizeMetric.metricValue.asInstanceOf[Double].toInt)
+
+      // After a request is delayed, the callback cannot be triggered immediately
+      quotaManager.throttledChannelReaper.doWork()
+      assertEquals(0, numCallbacks)
+
+      // Callback can only be triggered after the delay time passes
+      time.sleep(throttleTime)
+      quotaManager.throttledChannelReaper.doWork()
+      assertEquals(0, queueSizeMetric.metricValue.asInstanceOf[Double].toInt)
+      assertEquals(1, numCallbacks)
+
+      // Retry to spike worth of 110 mutations after having waited the required throttle time.
+      // Current tokens in the bucket = 0
+      throttleTime = maybeRecord(quotaManager, User, ClientId, 110)
+      assertEquals("Should be throttled", 0, throttleTime)
+    }
+  }
+
+  @Test
+  def testNewStrictQuotaForReturnsUnboundedQuotaWhenQuotaIsDisabled(): Unit = {
+    withQuotaManager { quotaManager =>
+      assertEquals(UnboundedControllerMutationQuota,
+        quotaManager.newStrictQuotaFor(buildSession(User), ClientId))
+    }
+  }
+
+  @Test
+  def testNewStrictQuotaForReturnsStrictQuotaWhenQuotaIsEnabled(): Unit = {
+    withQuotaManager { quotaManager =>
+      quotaManager.updateQuota(Some(User), Some(ClientId), Some(ClientId),
+        Some(Quota.upperBound(10)))
+      val quota = quotaManager.newStrictQuotaFor(buildSession(User), ClientId)
+      assertTrue(quota.isInstanceOf[StrictControllerMutationQuota])
+
+    }
+  }
+
+  @Test
+  def testNewPermissiveQuotaForReturnsUnboundedQuotaWhenQuotaIsDisabled(): Unit = {
+    withQuotaManager { quotaManager =>
+      assertEquals(UnboundedControllerMutationQuota,
+        quotaManager.newPermissiveQuotaFor(buildSession(User), ClientId))
+    }
+  }
+
+  @Test
+  def testNewPermissiveQuotaForReturnsStrictQuotaWhenQuotaIsEnabled(): Unit = {
+    withQuotaManager { quotaManager =>
+      quotaManager.updateQuota(Some(User), Some(ClientId), Some(ClientId),
+        Some(Quota.upperBound(10)))
+      val quota = quotaManager.newPermissiveQuotaFor(buildSession(User), ClientId)
+      assertTrue(quota.isInstanceOf[PermissiveControllerMutationQuota])
+    }
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/server/ControllerMutationQuotaTest.scala b/core/src/test/scala/unit/kafka/server/ControllerMutationQuotaTest.scala
new file mode 100644
index 0000000000000..7407a984b9029
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/ControllerMutationQuotaTest.scala
@@ -0,0 +1,420 @@
+/**
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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 java.util.concurrent.ExecutionException
+import java.util.concurrent.TimeUnit
+
+import kafka.server.ClientQuotaManager.DefaultTags
+import kafka.utils.TestUtils
+import org.apache.kafka.common.internals.KafkaFutureImpl
+import org.apache.kafka.common.message.CreatePartitionsRequestData
+import org.apache.kafka.common.message.CreatePartitionsRequestData.CreatePartitionsTopic
+import org.apache.kafka.common.message.CreateTopicsRequestData
+import org.apache.kafka.common.message.CreateTopicsRequestData.CreatableTopic
+import org.apache.kafka.common.message.DeleteTopicsRequestData
+import org.apache.kafka.common.metrics.KafkaMetric
+import org.apache.kafka.common.protocol.ApiKeys
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.quota.ClientQuotaAlteration
+import org.apache.kafka.common.quota.ClientQuotaEntity
+import org.apache.kafka.common.requests.AlterClientQuotasRequest
+import org.apache.kafka.common.requests.AlterClientQuotasResponse
+import org.apache.kafka.common.requests.CreatePartitionsRequest
+import org.apache.kafka.common.requests.CreatePartitionsResponse
+import org.apache.kafka.common.requests.CreateTopicsRequest
+import org.apache.kafka.common.requests.CreateTopicsResponse
+import org.apache.kafka.common.requests.DeleteTopicsRequest
+import org.apache.kafka.common.requests.DeleteTopicsResponse
+import org.apache.kafka.common.security.auth.AuthenticationContext
+import org.apache.kafka.common.security.auth.KafkaPrincipal
+import org.apache.kafka.common.security.auth.KafkaPrincipalBuilder
+import org.apache.kafka.test.{TestUtils => JTestUtils}
+import org.junit.Assert.assertEquals
+import org.junit.Assert.assertTrue
+import org.junit.Assert.fail
+import org.junit.Before
+import org.junit.Test
+
+import scala.jdk.CollectionConverters._
+
+object ControllerMutationQuotaTest {
+  // Principal used for all client connections. This is updated by each test.
+  var principal = KafkaPrincipal.ANONYMOUS
+  class TestPrincipalBuilder extends KafkaPrincipalBuilder {
+    override def build(context: AuthenticationContext): KafkaPrincipal = {
+      principal
+    }
+  }
+
+  def asPrincipal(newPrincipal: KafkaPrincipal)(f: => Unit): Unit = {
+    val currentPrincipal = principal
+    principal = newPrincipal
+    try f
+    finally principal = currentPrincipal
+  }
+
+  val ThrottledPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "ThrottledPrincipal")
+  val UnboundedPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "UnboundedPrincipal")
+
+  val StrictCreateTopicsRequestVersion = ApiKeys.CREATE_TOPICS.latestVersion
+  val PermissiveCreateTopicsRequestVersion = 5.toShort
+
+  val StrictDeleteTopicsRequestVersion = ApiKeys.DELETE_TOPICS.latestVersion
+  val PermissiveDeleteTopicsRequestVersion = 4.toShort
+
+  val StrictCreatePartitionsRequestVersion = ApiKeys.CREATE_PARTITIONS.latestVersion
+  val PermissiveCreatePartitionsRequestVersion = 2.toShort
+
+  val Topic1 = "topic-1"
+  val Topic2 = "topic-2"
+  val TopicsWithOnePartition = Map(Topic1 ->  1, Topic2 ->  1)
+  val TopicsWith30Partitions = Map(Topic1 -> 30, Topic2 -> 30)
+  val TopicsWith31Partitions = Map(Topic1 -> 31, Topic2 -> 31)
+
+  val ControllerQuotaSamples = 10
+  val ControllerQuotaWindowSizeSeconds = 1
+  val ControllerMutationRate = 2.0
+}
+
+class ControllerMutationQuotaTest extends BaseRequestTest {
+  import ControllerMutationQuotaTest._
+
+  override def brokerCount: Int = 1
+
+  override def brokerPropertyOverrides(properties: Properties): Unit = {
+    properties.put(KafkaConfig.ControlledShutdownEnableProp, "false")
+    properties.put(KafkaConfig.OffsetsTopicReplicationFactorProp, "1")
+    properties.put(KafkaConfig.OffsetsTopicPartitionsProp, "1")
+    properties.put(KafkaConfig.PrincipalBuilderClassProp,
+      classOf[ControllerMutationQuotaTest.TestPrincipalBuilder].getName)
+    // Specify number of samples and window size.
+    properties.put(KafkaConfig.NumControllerQuotaSamplesProp, ControllerQuotaSamples.toString)
+    properties.put(KafkaConfig.ControllerQuotaWindowSizeSecondsProp, ControllerQuotaWindowSizeSeconds.toString)
+  }
+
+  @Before
+  override def setUp(): Unit = {
+    super.setUp()
+
+    // Define a quota for ThrottledPrincipal
+    defineUserQuota(ThrottledPrincipal.getName, Some(ControllerMutationRate))
+    waitUserQuota(ThrottledPrincipal.getName, ControllerMutationRate)
+  }
+
+  @Test
+  def testSetUnsetQuota(): Unit = {
+    val rate = 1.5
+    val principal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "User")
+    // Default Value
+    waitUserQuota(principal.getName, Long.MaxValue)
+    // Define a new quota
+    defineUserQuota(principal.getName, Some(rate))
+    // Check it
+    waitUserQuota(principal.getName, rate)
+    // Remove it
+    defineUserQuota(principal.getName, None)
+    // Back to the default
+    waitUserQuota(principal.getName, Long.MaxValue)
+  }
+
+  @Test
+  def testQuotaMetric(): Unit = {
+    asPrincipal(ThrottledPrincipal) {
+      // Metric is lazily created
+      assertTrue(quotaMetric(principal.getName).isEmpty)
+
+      // Create a topic to create the metrics
+      val (_, errors) = createTopics(Map("topic" -> 1), StrictDeleteTopicsRequestVersion)
+      assertEquals(Set(Errors.NONE), errors.values.toSet)
+
+      // Metric must be there with the correct config
+      waitQuotaMetric(principal.getName, ControllerMutationRate)
+
+      // Update quota
+      defineUserQuota(ThrottledPrincipal.getName, Some(ControllerMutationRate * 2))
+      waitUserQuota(ThrottledPrincipal.getName, ControllerMutationRate * 2)
+
+      // Metric must be there with the updated config
+      waitQuotaMetric(principal.getName, ControllerMutationRate * 2)
+    }
+  }
+
+  @Test
+  def testStrictCreateTopicsRequest(): Unit = {
+    asPrincipal(ThrottledPrincipal) {
+      // Create two topics worth of 30 partitions each. As we use a strict quota, we
+      // expect one to be created and one to be rejected.
+      // Theoretically, the throttle time should be below or equal to:
+      // -(-10) / 2 = 5s
+      val (throttleTimeMs1, errors1) = createTopics(TopicsWith30Partitions, StrictCreateTopicsRequestVersion)
+      assertThrottleTime(5000, throttleTimeMs1)
+      // Ordering is not guaranteed so we only check the errors
+      assertEquals(Set(Errors.NONE, Errors.THROTTLING_QUOTA_EXCEEDED), errors1.values.toSet)
+
+      // Retry the rejected topic. It should succeed after the throttling delay is passed and the
+      // throttle time should be zero.
+      val rejectedTopicName = errors1.filter(_._2 == Errors.THROTTLING_QUOTA_EXCEEDED).keys.head
+      val rejectedTopicSpec = TopicsWith30Partitions.filter(_._1 == rejectedTopicName)
+      TestUtils.waitUntilTrue(() => {
+        val (throttleTimeMs2, errors2) = createTopics(rejectedTopicSpec, StrictCreateTopicsRequestVersion)
+        throttleTimeMs2 == 0 && errors2 == Map(rejectedTopicName -> Errors.NONE)
+      }, "Failed to create topics after having been throttled")
+    }
+  }
+
+  @Test
+  def testPermissiveCreateTopicsRequest(): Unit = {
+    asPrincipal(ThrottledPrincipal) {
+      // Create two topics worth of 30 partitions each. As we use a permissive quota, we
+      // expect both topics to be created.
+      // Theoretically, the throttle time should be below or equal to:
+      // -(-40) / 2 = 20s
+      val (throttleTimeMs, errors) = createTopics(TopicsWith30Partitions, PermissiveCreateTopicsRequestVersion)
+      assertThrottleTime(20000, throttleTimeMs)
+      assertEquals(Map(Topic1 -> Errors.NONE, Topic2 -> Errors.NONE), errors)
+    }
+  }
+
+  @Test
+  def testUnboundedCreateTopicsRequest(): Unit = {
+    asPrincipal(UnboundedPrincipal) {
+      // Create two topics worth of 30 partitions each. As we use an user without quota, we
+      // expect both topics to be created. The throttle time should be equal to 0.
+      val (throttleTimeMs, errors) = createTopics(TopicsWith30Partitions, StrictCreateTopicsRequestVersion)
+      assertEquals(0, throttleTimeMs)
+      assertEquals(Map(Topic1 -> Errors.NONE, Topic2 -> Errors.NONE), errors)
+    }
+  }
+
+  @Test
+  def testStrictDeleteTopicsRequest(): Unit = {
+    asPrincipal(UnboundedPrincipal) {
+      createTopics(TopicsWith30Partitions, StrictCreateTopicsRequestVersion)
+    }
+
+    asPrincipal(ThrottledPrincipal) {
+      // Delete two topics worth of 30 partitions each. As we use a strict quota, we
+      // expect the first topic to be deleted and the second to be rejected.
+      // Theoretically, the throttle time should be below or equal to:
+      // -(-10) / 2 = 5s
+      val (throttleTimeMs1, errors1) = deleteTopics(TopicsWith30Partitions, StrictDeleteTopicsRequestVersion)
+      assertThrottleTime(5000, throttleTimeMs1)
+      // Ordering is not guaranteed so we only check the errors
+      assertEquals(Set(Errors.NONE, Errors.THROTTLING_QUOTA_EXCEEDED), errors1.values.toSet)
+
+      // Retry the rejected topic. It should succeed after the throttling delay is passed and the
+      // throttle time should be zero.
+      val rejectedTopicName = errors1.filter(_._2 == Errors.THROTTLING_QUOTA_EXCEEDED).keys.head
+      val rejectedTopicSpec = TopicsWith30Partitions.filter(_._1 == rejectedTopicName)
+      TestUtils.waitUntilTrue(() => {
+        val (throttleTimeMs2, errors2) = deleteTopics(rejectedTopicSpec, StrictDeleteTopicsRequestVersion)
+        throttleTimeMs2 == 0 && errors2 == Map(rejectedTopicName -> Errors.NONE)
+      }, "Failed to delete topics after having been throttled")
+    }
+  }
+
+  @Test
+  def testPermissiveDeleteTopicsRequest(): Unit = {
+    asPrincipal(UnboundedPrincipal) {
+      createTopics(TopicsWith30Partitions, StrictCreateTopicsRequestVersion)
+    }
+
+    asPrincipal(ThrottledPrincipal) {
+      // Delete two topics worth of 30 partitions each. As we use a permissive quota, we
+      // expect both topics to be deleted.
+      // Theoretically, the throttle time should be below or equal to:
+      // -(-40) / 2 = 20s
+      val (throttleTimeMs, errors) = deleteTopics(TopicsWith30Partitions, PermissiveDeleteTopicsRequestVersion)
+      assertThrottleTime(20000, throttleTimeMs)
+      assertEquals(Map(Topic1 -> Errors.NONE, Topic2 -> Errors.NONE), errors)
+    }
+  }
+
+  @Test
+  def testUnboundedDeleteTopicsRequest(): Unit = {
+    asPrincipal(UnboundedPrincipal) {
+      createTopics(TopicsWith30Partitions, StrictCreateTopicsRequestVersion)
+
+      // Delete two topics worth of 30 partitions each. As we use an user without quota, we
+      // expect both topics to be deleted. The throttle time should be equal to 0.
+      val (throttleTimeMs, errors) = deleteTopics(TopicsWith30Partitions, StrictDeleteTopicsRequestVersion)
+      assertEquals(0, throttleTimeMs)
+      assertEquals(Map(Topic1 -> Errors.NONE, Topic2 -> Errors.NONE), errors)
+    }
+  }
+
+  @Test
+  def testStrictCreatePartitionsRequest(): Unit = {
+    asPrincipal(UnboundedPrincipal) {
+      createTopics(TopicsWithOnePartition, StrictCreatePartitionsRequestVersion)
+    }
+
+    asPrincipal(ThrottledPrincipal) {
+      // Add 30 partitions to each topic. As we use a strict quota, we
+      // expect the first topic to be extended and the second to be rejected.
+      // Theoretically, the throttle time should be below or equal to:
+      // -(-10) / 2 = 5s
+      val (throttleTimeMs1, errors1) = createPartitions(TopicsWith31Partitions, StrictCreatePartitionsRequestVersion)
+      assertThrottleTime(5000, throttleTimeMs1)
+      // Ordering is not guaranteed so we only check the errors
+      assertEquals(Set(Errors.NONE, Errors.THROTTLING_QUOTA_EXCEEDED), errors1.values.toSet)
+
+      // Retry the rejected topic. It should succeed after the throttling delay is passed and the
+      // throttle time should be zero.
+      val rejectedTopicName = errors1.filter(_._2 == Errors.THROTTLING_QUOTA_EXCEEDED).keys.head
+      val rejectedTopicSpec = TopicsWith30Partitions.filter(_._1 == rejectedTopicName)
+      TestUtils.waitUntilTrue(() => {
+        val (throttleTimeMs2, errors2) = createPartitions(rejectedTopicSpec, StrictCreatePartitionsRequestVersion)
+        throttleTimeMs2 == 0 && errors2 == Map(rejectedTopicName -> Errors.NONE)
+      }, "Failed to create partitions after having been throttled")
+    }
+  }
+
+  @Test
+  def testPermissiveCreatePartitionsRequest(): Unit = {
+    asPrincipal(UnboundedPrincipal) {
+      createTopics(TopicsWithOnePartition, StrictCreatePartitionsRequestVersion)
+    }
+
+    asPrincipal(ThrottledPrincipal) {
+      // Create two topics worth of 30 partitions each. As we use a permissive quota, we
+      // expect both topics to be created.
+      // Theoretically, the throttle time should be below or equal to:
+      // -(-40) / 2 = 20s
+      val (throttleTimeMs, errors) = createPartitions(TopicsWith31Partitions, PermissiveCreatePartitionsRequestVersion)
+      assertThrottleTime(20000, throttleTimeMs)
+      assertEquals(Map(Topic1 -> Errors.NONE, Topic2 -> Errors.NONE), errors)
+    }
+  }
+
+  @Test
+  def testUnboundedCreatePartitionsRequest(): Unit = {
+    asPrincipal(UnboundedPrincipal) {
+      createTopics(TopicsWithOnePartition, StrictCreatePartitionsRequestVersion)
+
+      // Create two topics worth of 30 partitions each. As we use an user without quota, we
+      // expect both topics to be created. The throttle time should be equal to 0.
+      val (throttleTimeMs, errors) = createPartitions(TopicsWith31Partitions, StrictCreatePartitionsRequestVersion)
+      assertEquals(0, throttleTimeMs)
+      assertEquals(Map(Topic1 -> Errors.NONE, Topic2 -> Errors.NONE), errors)
+    }
+  }
+
+  private def assertThrottleTime(max: Int, actual: Int): Unit = {
+    assertTrue(
+      s"Expected a throttle time between 0 and $max but got $actual",
+      (actual >= 0) && (actual <= max))
+  }
+
+  private def createTopics(topics: Map[String, Int], version: Short): (Int, Map[String, Errors]) = {
+    val data = new CreateTopicsRequestData()
+    topics.foreach { case (topic, numPartitions) =>
+      data.topics.add(new CreatableTopic()
+        .setName(topic).setNumPartitions(numPartitions).setReplicationFactor(1))
+    }
+    val request = new CreateTopicsRequest.Builder(data).build(version)
+    val response = connectAndReceive[CreateTopicsResponse](request)
+    response.data.throttleTimeMs -> response.data.topics.asScala
+      .map(topic => topic.name -> Errors.forCode(topic.errorCode)).toMap
+  }
+
+  private def deleteTopics(topics: Map[String, Int], version: Short): (Int, Map[String, Errors]) = {
+    val data = new DeleteTopicsRequestData()
+      .setTimeoutMs(60000)
+      .setTopicNames(topics.keys.toSeq.asJava)
+    val request = new DeleteTopicsRequest.Builder(data).build(version)
+    val response = connectAndReceive[DeleteTopicsResponse](request)
+    response.data.throttleTimeMs -> response.data.responses.asScala
+      .map(topic => topic.name -> Errors.forCode(topic.errorCode)).toMap
+  }
+
+  private def createPartitions(topics: Map[String, Int], version: Short): (Int, Map[String, Errors]) = {
+    val data = new CreatePartitionsRequestData().setTimeoutMs(60000)
+    topics.foreach { case (topic, numPartitions) =>
+      data.topics.add(new CreatePartitionsTopic()
+        .setName(topic).setCount(numPartitions).setAssignments(null))
+    }
+    val request = new CreatePartitionsRequest.Builder(data).build(version)
+    val response = connectAndReceive[CreatePartitionsResponse](request)
+    response.data.throttleTimeMs -> response.data.results.asScala
+      .map(topic => topic.name -> Errors.forCode(topic.errorCode)).toMap
+  }
+
+  private def defineUserQuota(user: String, quota: Option[Double]): Unit = {
+    val entity = new ClientQuotaEntity(Map(ClientQuotaEntity.USER -> user).asJava)
+    val quotas = Map(DynamicConfig.Client.ControllerMutationOverrideProp -> quota)
+
+    try alterClientQuotas(Map(entity -> quotas))(entity).get(10, TimeUnit.SECONDS) catch {
+      case e: ExecutionException => throw e.getCause
+    }
+  }
+
+  private def waitUserQuota(user: String, expectedQuota: Double): Unit = {
+    val quotaManager = servers.head.quotaManagers.controllerMutation
+    var actualQuota = Double.MinValue
+
+    TestUtils.waitUntilTrue(() => {
+      actualQuota = quotaManager.quota(user, "").bound()
+      expectedQuota == actualQuota
+    }, s"Quota of $user is not $expectedQuota but $actualQuota")
+  }
+
+  private def quotaMetric(user: String): Option[KafkaMetric] = {
+    val metrics = servers.head.metrics
+    val metricName = metrics.metricName(
+      "tokens",
+      QuotaType.ControllerMutation.toString,
+      "Tracking remaining tokens in the token bucket per user/client-id",
+      Map(DefaultTags.User -> user, DefaultTags.ClientId -> "").asJava)
+    Option(servers.head.metrics.metric(metricName))
+  }
+
+  private def waitQuotaMetric(user: String, expectedQuota: Double): Unit = {
+    TestUtils.retry(JTestUtils.DEFAULT_MAX_WAIT_MS) {
+      quotaMetric(user) match {
+        case Some(metric) =>
+          val config = metric.config()
+          assertEquals(expectedQuota, config.quota().bound(), 0.1)
+          assertEquals(ControllerQuotaSamples, config.samples())
+          assertEquals(ControllerQuotaWindowSizeSeconds * 1000, config.timeWindowMs())
+
+        case None =>
+          fail(s"Quota metric of $user is not defined")
+      }
+    }
+  }
+
+  private def alterClientQuotas(request: Map[ClientQuotaEntity, Map[String, Option[Double]]]): Map[ClientQuotaEntity, KafkaFutureImpl[Void]] = {
+    val entries = request.map { case (entity, alter) =>
+      val ops = alter.map { case (key, value) =>
+        new ClientQuotaAlteration.Op(key, value.map(Double.box).orNull)
+      }.asJavaCollection
+      new ClientQuotaAlteration(entity, ops)
+    }
+
+    val response = request.map(e => e._1 -> new KafkaFutureImpl[Void]).asJava
+    sendAlterClientQuotasRequest(entries).complete(response)
+    val result = response.asScala
+    assertEquals(request.size, result.size)
+    request.foreach(e => assertTrue(result.get(e._1).isDefined))
+    result.toMap
+  }
+
+  private def sendAlterClientQuotasRequest(entries: Iterable[ClientQuotaAlteration]): AlterClientQuotasResponse = {
+    val request = new AlterClientQuotasRequest.Builder(entries.asJavaCollection, false).build()
+    connectAndReceive[AlterClientQuotasResponse](request, destination = controllerSocketServer)
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/server/CreateTopicsRequestTest.scala b/core/src/test/scala/unit/kafka/server/CreateTopicsRequestTest.scala
index 4ab952032a4cf..245f13f9aca85 100644
--- a/core/src/test/scala/unit/kafka/server/CreateTopicsRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/CreateTopicsRequestTest.scala
@@ -18,71 +18,76 @@
 package kafka.server
 
 import kafka.utils._
+import org.apache.kafka.common.message.CreateTopicsRequestData
+import org.apache.kafka.common.message.CreateTopicsRequestData.CreatableTopicCollection
+import org.apache.kafka.common.protocol.ApiKeys
 import org.apache.kafka.common.protocol.Errors
 import org.apache.kafka.common.requests.CreateTopicsRequest
 import org.junit.Assert._
 import org.junit.Test
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 
 class CreateTopicsRequestTest extends AbstractCreateTopicsRequestTest {
 
   @Test
-  def testValidCreateTopicsRequests() {
-    val timeout = 10000
+  def testValidCreateTopicsRequests(): Unit = {
     // Generated assignments
-    validateValidCreateTopicsRequests(new CreateTopicsRequest.Builder(Map("topic1" -> new CreateTopicsRequest.TopicDetails(1, 1.toShort)).asJava, timeout).build())
-    validateValidCreateTopicsRequests(new CreateTopicsRequest.Builder(Map("topic2" -> new CreateTopicsRequest.TopicDetails(1, 3.toShort)).asJava, timeout).build())
-    val config3 = Map("min.insync.replicas" -> "2").asJava
-    validateValidCreateTopicsRequests(new CreateTopicsRequest.Builder(Map("topic3" -> new CreateTopicsRequest.TopicDetails(5, 2.toShort, config3)).asJava, timeout).build())
+    validateValidCreateTopicsRequests(topicsReq(Seq(topicReq("topic1"))))
+    validateValidCreateTopicsRequests(topicsReq(Seq(topicReq("topic2", replicationFactor = 3))))
+    validateValidCreateTopicsRequests(topicsReq(Seq(topicReq("topic3",
+      numPartitions = 5, replicationFactor = 2, config = Map("min.insync.replicas" -> "2")))))
     // Manual assignments
-    val assignments4 = replicaAssignmentToJava(Map(0 -> List(0)))
-    validateValidCreateTopicsRequests(new CreateTopicsRequest.Builder(Map("topic4" -> new CreateTopicsRequest.TopicDetails(assignments4)).asJava, timeout).build())
-    val assignments5 = replicaAssignmentToJava(Map(0 -> List(0, 1), 1 -> List(1, 0), 2 -> List(1, 2)))
-    val config5 = Map("min.insync.replicas" -> "2").asJava
-    validateValidCreateTopicsRequests(new CreateTopicsRequest.Builder(Map("topic5" -> new CreateTopicsRequest.TopicDetails(assignments5, config5)).asJava, timeout).build())
+    validateValidCreateTopicsRequests(topicsReq(Seq(topicReq("topic4", assignment = Map(0 -> List(0))))))
+    validateValidCreateTopicsRequests(topicsReq(Seq(topicReq("topic5",
+      assignment = Map(0 -> List(0, 1), 1 -> List(1, 0), 2 -> List(1, 2)),
+      config = Map("min.insync.replicas" -> "2")))))
     // Mixed
-    val assignments8 = replicaAssignmentToJava(Map(0 -> List(0, 1), 1 -> List(1, 0), 2 -> List(1, 2)))
-    validateValidCreateTopicsRequests(new CreateTopicsRequest.Builder(Map(
-      "topic6" -> new CreateTopicsRequest.TopicDetails(1, 1.toShort),
-      "topic7" -> new CreateTopicsRequest.TopicDetails(5, 2.toShort),
-      "topic8" -> new CreateTopicsRequest.TopicDetails(assignments8)).asJava, timeout).build()
-    )
-    validateValidCreateTopicsRequests(new CreateTopicsRequest.Builder(Map(
-      "topic9" -> new CreateTopicsRequest.TopicDetails(1, 1.toShort),
-      "topic10" -> new CreateTopicsRequest.TopicDetails(5, 2.toShort),
-      "topic11" -> new CreateTopicsRequest.TopicDetails(assignments8)).asJava, timeout, true).build()
-    )
+    validateValidCreateTopicsRequests(topicsReq(Seq(topicReq("topic6"),
+      topicReq("topic7", numPartitions = 5, replicationFactor = 2),
+      topicReq("topic8", assignment = Map(0 -> List(0, 1), 1 -> List(1, 0), 2 -> List(1, 2))))))
+    validateValidCreateTopicsRequests(topicsReq(Seq(topicReq("topic9"),
+      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() {
-    val timeout = 10000
+  def testErrorCreateTopicsRequests(): Unit = {
     val existingTopic = "existing-topic"
-    TestUtils.createTopic(zkUtils, existingTopic, 1, 1, servers)
-
+    createTopic(existingTopic, 1, 1)
     // Basic
-    validateErrorCreateTopicsRequests(new CreateTopicsRequest.Builder(Map(existingTopic -> new CreateTopicsRequest.TopicDetails(1, 1.toShort)).asJava, timeout).build(),
+    validateErrorCreateTopicsRequests(topicsReq(Seq(topicReq(existingTopic))),
       Map(existingTopic -> error(Errors.TOPIC_ALREADY_EXISTS, Some("Topic 'existing-topic' already exists."))))
-    validateErrorCreateTopicsRequests(new CreateTopicsRequest.Builder(Map("error-partitions" -> new CreateTopicsRequest.TopicDetails(-1, 1.toShort)).asJava, timeout).build(),
+    validateErrorCreateTopicsRequests(topicsReq(Seq(topicReq("error-partitions", numPartitions = -2))),
       Map("error-partitions" -> error(Errors.INVALID_PARTITIONS)), checkErrorMessage = false)
-    validateErrorCreateTopicsRequests(new CreateTopicsRequest.Builder(Map("error-replication" -> new CreateTopicsRequest.TopicDetails(1, (numBrokers + 1).toShort)).asJava, timeout).build(),
+    validateErrorCreateTopicsRequests(topicsReq(Seq(topicReq("error-replication",
+      replicationFactor = brokerCount + 1))),
       Map("error-replication" -> error(Errors.INVALID_REPLICATION_FACTOR)), checkErrorMessage = false)
-    val invalidConfig = Map("not.a.property" -> "error").asJava
-    validateErrorCreateTopicsRequests(new CreateTopicsRequest.Builder(Map("error-config" -> new CreateTopicsRequest.TopicDetails(1, 1.toShort, invalidConfig)).asJava, timeout).build(),
+    validateErrorCreateTopicsRequests(topicsReq(Seq(topicReq("error-config",
+      config=Map("not.a.property" -> "error")))),
       Map("error-config" -> error(Errors.INVALID_CONFIG)), checkErrorMessage = false)
-    val invalidAssignments = replicaAssignmentToJava(Map(0 -> List(0, 1), 1 -> List(0)))
-    validateErrorCreateTopicsRequests(new CreateTopicsRequest.Builder(Map("error-assignment" -> new CreateTopicsRequest.TopicDetails(invalidAssignments)).asJava, timeout).build(),
+    validateErrorCreateTopicsRequests(topicsReq(Seq(topicReq("error-config-value",
+      config=Map("message.format.version" -> "invalid-value")))),
+      Map("error-config-value" -> error(Errors.INVALID_CONFIG)), checkErrorMessage = false)
+    validateErrorCreateTopicsRequests(topicsReq(Seq(topicReq("error-assignment",
+      assignment=Map(0 -> List(0, 1), 1 -> List(0))))),
       Map("error-assignment" -> error(Errors.INVALID_REPLICA_ASSIGNMENT)), checkErrorMessage = false)
 
     // Partial
-    validateErrorCreateTopicsRequests(
-      new CreateTopicsRequest.Builder(Map(
-        existingTopic -> new CreateTopicsRequest.TopicDetails(1, 1.toShort),
-        "partial-partitions" -> new CreateTopicsRequest.TopicDetails(-1, 1.toShort),
-        "partial-replication" -> new CreateTopicsRequest.TopicDetails(1, (numBrokers + 1).toShort),
-        "partial-assignment" -> new CreateTopicsRequest.TopicDetails(invalidAssignments),
-        "partial-none" -> new CreateTopicsRequest.TopicDetails(1, 1.toShort)).asJava, timeout).build(),
+    validateErrorCreateTopicsRequests(topicsReq(Seq(
+      topicReq(existingTopic),
+      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(
         existingTopic -> error(Errors.TOPIC_ALREADY_EXISTS),
         "partial-partitions" -> error(Errors.INVALID_PARTITIONS),
@@ -95,12 +100,15 @@ class CreateTopicsRequestTest extends AbstractCreateTopicsRequestTest {
 
     // Timeout
     // We don't expect a request to ever complete within 1ms. A timeout of 1 ms allows us to test the purgatory timeout logic.
-    validateErrorCreateTopicsRequests(new CreateTopicsRequest.Builder(Map("error-timeout" -> new CreateTopicsRequest.TopicDetails(10, 3.toShort)).asJava, 1).build(),
+    validateErrorCreateTopicsRequests(topicsReq(Seq(
+      topicReq("error-timeout", numPartitions = 10, replicationFactor = 3)), timeout = 1),
       Map("error-timeout" -> error(Errors.REQUEST_TIMED_OUT)), checkErrorMessage = false)
-    validateErrorCreateTopicsRequests(new CreateTopicsRequest.Builder(Map("error-timeout-zero" -> new CreateTopicsRequest.TopicDetails(10, 3.toShort)).asJava, 0).build(),
+    validateErrorCreateTopicsRequests(topicsReq(Seq(
+      topicReq("error-timeout-zero", numPartitions = 10, replicationFactor = 3)), timeout = 0),
       Map("error-timeout-zero" -> error(Errors.REQUEST_TIMED_OUT)), checkErrorMessage = false)
     // Negative timeouts are treated the same as 0
-    validateErrorCreateTopicsRequests(new CreateTopicsRequest.Builder(Map("error-timeout-negative" -> new CreateTopicsRequest.TopicDetails(10, 3.toShort)).asJava, -1).build(),
+    validateErrorCreateTopicsRequests(topicsReq(Seq(
+      topicReq("error-timeout-negative", numPartitions = 10, replicationFactor = 3)), timeout = -1),
       Map("error-timeout-negative" -> error(Errors.REQUEST_TIMED_OUT)), checkErrorMessage = false)
     // The topics should still get created eventually
     TestUtils.waitUntilMetadataIsPropagated(servers, "error-timeout", 0)
@@ -112,55 +120,56 @@ class CreateTopicsRequestTest extends AbstractCreateTopicsRequestTest {
   }
 
   @Test
-  def testInvalidCreateTopicsRequests() {
-    // Duplicate
-    val singleRequest = new CreateTopicsRequest.Builder(Map("duplicate-topic" ->
-        new CreateTopicsRequest.TopicDetails(1, 1.toShort)).asJava, 1000).build()
-    validateErrorCreateTopicsRequests(singleRequest, Map("duplicate-topic" -> error(Errors.INVALID_REQUEST,
-      Some("""Create topics request from client `client-id` contains multiple entries for the following topics: duplicate-topic"""))),
-      requestStruct = Some(toStructWithDuplicateFirstTopic(singleRequest)))
-
-    // Duplicate Partial with validateOnly
-    val doubleRequestValidateOnly = new CreateTopicsRequest.Builder(Map(
-      "duplicate-topic" -> new CreateTopicsRequest.TopicDetails(1, 1.toShort),
-      "other-topic" -> new CreateTopicsRequest.TopicDetails(1, 1.toShort)).asJava, 1000, true).build()
-    validateErrorCreateTopicsRequests(doubleRequestValidateOnly, Map(
-      "duplicate-topic" -> error(Errors.INVALID_REQUEST),
-      "other-topic" -> error(Errors.NONE)), checkErrorMessage = false,
-      requestStruct = Some(toStructWithDuplicateFirstTopic(doubleRequestValidateOnly)))
-
-    // Duplicate Partial
-    val doubleRequest = new CreateTopicsRequest.Builder(Map(
-      "duplicate-topic" -> new CreateTopicsRequest.TopicDetails(1, 1.toShort),
-      "other-topic" -> new CreateTopicsRequest.TopicDetails(1, 1.toShort)).asJava, 1000).build()
-    validateErrorCreateTopicsRequests(doubleRequest, Map(
-      "duplicate-topic" -> error(Errors.INVALID_REQUEST),
-      "other-topic" -> error(Errors.NONE)), checkErrorMessage = false,
-      requestStruct = Some(toStructWithDuplicateFirstTopic(doubleRequest)))
-
+  def testInvalidCreateTopicsRequests(): Unit = {
     // Partitions/ReplicationFactor and ReplicaAssignment
-    val assignments = replicaAssignmentToJava(Map(0 -> List(0)))
-    val assignmentRequest = new CreateTopicsRequest.Builder(Map("bad-args-topic" ->
-        new CreateTopicsRequest.TopicDetails(assignments)).asJava, 1000).build()
-    val badArgumentsRequest = addPartitionsAndReplicationFactorToFirstTopic(assignmentRequest)
-    validateErrorCreateTopicsRequests(badArgumentsRequest, Map("bad-args-topic" -> error(Errors.INVALID_REQUEST)),
-      checkErrorMessage = false)
-
-    // Partitions/ReplicationFactor and ReplicaAssignment with validateOnly
-    val assignmentRequestValidateOnly = new CreateTopicsRequest.Builder(Map("bad-args-topic" ->
-      new CreateTopicsRequest.TopicDetails(assignments)).asJava, 1000, true).build()
-    val badArgumentsRequestValidateOnly = addPartitionsAndReplicationFactorToFirstTopic(assignmentRequestValidateOnly)
-    validateErrorCreateTopicsRequests(badArgumentsRequestValidateOnly, Map("bad-args-topic" -> error(Errors.INVALID_REQUEST)),
-      checkErrorMessage = false)
+    validateErrorCreateTopicsRequests(topicsReq(Seq(
+      topicReq("bad-args-topic", numPartitions = 10, replicationFactor = 3,
+        assignment = Map(0 -> List(0))))),
+      Map("bad-args-topic" -> error(Errors.INVALID_REQUEST)), checkErrorMessage = false)
+
+    validateErrorCreateTopicsRequests(topicsReq(Seq(
+      topicReq("bad-args-topic", numPartitions = 10, replicationFactor = 3,
+        assignment = Map(0 -> List(0)))), validateOnly = true),
+      Map("bad-args-topic" -> error(Errors.INVALID_REQUEST)), checkErrorMessage = false)
   }
 
   @Test
-  def testNotController() {
-    val request = new CreateTopicsRequest.Builder(Map("topic1" -> new CreateTopicsRequest.TopicDetails(1, 1.toShort)).asJava, 1000).build()
-    val response = sendCreateTopicRequest(request, notControllerSocketServer)
-
-    val error = response.errors.asScala.head._2.error
-    assertEquals("Expected controller error when routed incorrectly", Errors.NOT_CONTROLLER, error)
+  def testNotController(): Unit = {
+    val req = topicsReq(Seq(topicReq("topic1")))
+    val response = sendCreateTopicRequest(req, notControllerSocketServer)
+    assertEquals(1, response.errorCounts().get(Errors.NOT_CONTROLLER))
   }
 
+  @Test
+  def testCreateTopicsRequestVersions(): Unit = {
+    for (version <- ApiKeys.CREATE_TOPICS.oldestVersion to ApiKeys.CREATE_TOPICS.latestVersion) {
+      val topic = s"topic_$version"
+      val data = new CreateTopicsRequestData()
+      data.setTimeoutMs(10000)
+      data.setValidateOnly(false)
+      data.setTopics(new CreatableTopicCollection(List(
+        topicReq(topic, numPartitions = 1, replicationFactor = 1,
+          config = Map("min.insync.replicas" -> "2"))
+      ).asJava.iterator()))
+
+      val request = new CreateTopicsRequest.Builder(data).build(version.asInstanceOf[Short])
+      val response = sendCreateTopicRequest(request)
+
+      val topicResponse = response.data.topics.find(topic)
+      assertNotNull(topicResponse)
+      assertEquals(topic, topicResponse.name)
+      assertEquals(Errors.NONE.code, topicResponse.errorCode)
+      if (version >= 5) {
+        assertEquals(1, topicResponse.numPartitions)
+        assertEquals(1, topicResponse.replicationFactor)
+        val config = topicResponse.configs().asScala.find(_.name == "min.insync.replicas")
+        assertTrue(config.isDefined)
+        assertEquals("2", config.get.value)
+      } else {
+        assertEquals(-1, topicResponse.numPartitions)
+        assertEquals(-1, topicResponse.replicationFactor)
+        assertTrue(topicResponse.configs.isEmpty)
+      }
+    }
+  }
 }
diff --git a/core/src/test/scala/unit/kafka/server/CreateTopicsRequestWithForwardingTest.scala b/core/src/test/scala/unit/kafka/server/CreateTopicsRequestWithForwardingTest.scala
new file mode 100644
index 0000000000000..afd8110deca97
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/CreateTopicsRequestWithForwardingTest.scala
@@ -0,0 +1,42 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package kafka.server
+
+import java.util.Properties
+
+import org.apache.kafka.common.protocol.Errors
+import org.junit.Assert.assertEquals
+import org.junit.Test
+
+import scala.jdk.CollectionConverters._
+
+class CreateTopicsRequestWithForwardingTest extends AbstractCreateTopicsRequestTest {
+
+  override def brokerPropertyOverrides(properties: Properties): Unit = {
+    properties.put(KafkaConfig.EnableMetadataQuorumProp, true.toString)
+  }
+
+  @Test
+  def testForwardToController(): Unit = {
+    val req = topicsReq(Seq(topicReq("topic1")))
+    val response = sendCreateTopicRequest(req, notControllerSocketServer)
+    // With forwarding enabled, request could be forwarded to the active controller.
+    assertEquals(Map(Errors.NONE -> 1), response.errorCounts().asScala)
+  }
+
+}
diff --git a/core/src/test/scala/unit/kafka/server/CreateTopicsRequestWithPolicyTest.scala b/core/src/test/scala/unit/kafka/server/CreateTopicsRequestWithPolicyTest.scala
index 36e03af466cd9..e97c04bb768b5 100644
--- a/core/src/test/scala/unit/kafka/server/CreateTopicsRequestWithPolicyTest.scala
+++ b/core/src/test/scala/unit/kafka/server/CreateTopicsRequestWithPolicyTest.scala
@@ -21,86 +21,92 @@ import java.util
 import java.util.Properties
 
 import kafka.log.LogConfig
-import kafka.utils.TestUtils
 import org.apache.kafka.common.errors.PolicyViolationException
 import org.apache.kafka.common.protocol.Errors
-import org.apache.kafka.common.requests.CreateTopicsRequest
 import org.apache.kafka.server.policy.CreateTopicPolicy
 import org.apache.kafka.server.policy.CreateTopicPolicy.RequestMetadata
 import org.junit.Test
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 
 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() {
-    val timeout = 10000
+  def testValidCreateTopicsRequests(): Unit = {
+    validateValidCreateTopicsRequests(topicsReq(Seq(topicReq("topic1",
+      numPartitions = 5))))
 
-    validateValidCreateTopicsRequests(new CreateTopicsRequest.Builder(
-      Map("topic1" -> new CreateTopicsRequest.TopicDetails(5, 1.toShort)).asJava, timeout).build())
+    validateValidCreateTopicsRequests(topicsReq(Seq(topicReq("topic2",
+      numPartitions = 5, replicationFactor = 3)),
+      validateOnly = true))
 
-    validateValidCreateTopicsRequests(new CreateTopicsRequest.Builder(
-      Map("topic2" -> new CreateTopicsRequest.TopicDetails(5, 3.toShort)).asJava, timeout, true).build())
+    validateValidCreateTopicsRequests(topicsReq(Seq(topicReq("topic3",
+      numPartitions = 11, replicationFactor = 2,
+      config = Map(LogConfig.RetentionMsProp -> 4999.toString))),
+      validateOnly = true))
 
-    val configs = Map(LogConfig.RetentionMsProp -> 4999.toString)
-    validateValidCreateTopicsRequests(new CreateTopicsRequest.Builder(
-      Map("topic3" -> new CreateTopicsRequest.TopicDetails(11, 2.toShort, configs.asJava)).asJava, timeout, true).build())
-
-    val assignments = replicaAssignmentToJava(Map(0 -> List(1, 0), 1 -> List(0, 1)))
-    validateValidCreateTopicsRequests(new CreateTopicsRequest.Builder(
-      Map("topic4" -> new CreateTopicsRequest.TopicDetails(assignments)).asJava, timeout).build())
+    validateValidCreateTopicsRequests(topicsReq(Seq(topicReq("topic4",
+      assignment = Map(0 -> List(1, 0), 1 -> List(0, 1))))))
   }
 
   @Test
-  def testErrorCreateTopicsRequests() {
-    val timeout = 10000
+  def testErrorCreateTopicsRequests(): Unit = {
     val existingTopic = "existing-topic"
-    TestUtils.createTopic(zkUtils, existingTopic, 1, 1, servers)
+    createTopic(existingTopic, 1, 1)
 
     // Policy violations
-    validateErrorCreateTopicsRequests(new CreateTopicsRequest.Builder(
-      Map("policy-topic1" -> new CreateTopicsRequest.TopicDetails(4, 1.toShort)).asJava, timeout).build(),
+    validateErrorCreateTopicsRequests(topicsReq(Seq(topicReq("policy-topic1",
+      numPartitions = 4, replicationFactor = 1))),
       Map("policy-topic1" -> error(Errors.POLICY_VIOLATION, Some("Topics should have at least 5 partitions, received 4"))))
 
-    validateErrorCreateTopicsRequests(new CreateTopicsRequest.Builder(
-      Map("policy-topic2" -> new CreateTopicsRequest.TopicDetails(4, 3.toShort)).asJava, timeout, true).build(),
+    validateErrorCreateTopicsRequests(topicsReq(Seq(topicReq("policy-topic2",
+      numPartitions = 4, replicationFactor = 3)), validateOnly = true),
       Map("policy-topic2" -> error(Errors.POLICY_VIOLATION, Some("Topics should have at least 5 partitions, received 4"))))
 
-    val configs = Map(LogConfig.RetentionMsProp -> 5001.toString)
-    validateErrorCreateTopicsRequests(new CreateTopicsRequest.Builder(
-      Map("policy-topic3" -> new CreateTopicsRequest.TopicDetails(11, 2.toShort, configs.asJava)).asJava, timeout, true).build(),
-      Map("policy-topic3" -> error(Errors.POLICY_VIOLATION, Some("RetentionMs should be less than 5000ms if replicationFactor > 5"))))
-
-    validateErrorCreateTopicsRequests(new CreateTopicsRequest.Builder(
-      Map("policy-topic4" -> new CreateTopicsRequest.TopicDetails(11, 3.toShort, Map.empty.asJava)).asJava, timeout, true).build(),
-      Map("policy-topic4" -> error(Errors.POLICY_VIOLATION, Some("RetentionMs should be less than 5000ms if replicationFactor > 5"))))
-
-    val assignments = replicaAssignmentToJava(Map(0 -> List(1), 1 -> List(0)))
-    validateErrorCreateTopicsRequests(new CreateTopicsRequest.Builder(
-      Map("policy-topic5" -> new CreateTopicsRequest.TopicDetails(assignments)).asJava, timeout).build(),
+    validateErrorCreateTopicsRequests(topicsReq(Seq(topicReq("policy-topic3",
+      numPartitions = 11, replicationFactor = 2,
+      config = Map(LogConfig.RetentionMsProp -> 5001.toString))), validateOnly = true),
+      Map("policy-topic3" -> error(Errors.POLICY_VIOLATION,
+        Some("RetentionMs should be less than 5000ms if replicationFactor > 5"))))
+
+    validateErrorCreateTopicsRequests(topicsReq(Seq(topicReq("policy-topic4",
+      numPartitions = 11, replicationFactor = 3,
+      config = Map(LogConfig.RetentionMsProp -> 5001.toString))), validateOnly = true),
+      Map("policy-topic4" -> error(Errors.POLICY_VIOLATION,
+        Some("RetentionMs should be less than 5000ms if replicationFactor > 5"))))
+
+    validateErrorCreateTopicsRequests(topicsReq(Seq(topicReq("policy-topic5",
+      assignment = Map(0 -> List(1), 1 -> List(0)),
+      config = Map(LogConfig.RetentionMsProp -> 5001.toString))), validateOnly = true),
       Map("policy-topic5" -> error(Errors.POLICY_VIOLATION,
         Some("Topic partitions should have at least 2 partitions, received 1 for partition 0"))))
 
     // Check that basic errors still work
-    validateErrorCreateTopicsRequests(new CreateTopicsRequest.Builder(
-      Map(existingTopic -> new CreateTopicsRequest.TopicDetails(5, 1.toShort)).asJava, timeout).build(),
-      Map(existingTopic -> error(Errors.TOPIC_ALREADY_EXISTS, Some("Topic 'existing-topic' already exists."))))
+    validateErrorCreateTopicsRequests(topicsReq(Seq(topicReq(existingTopic,
+      numPartitions = 5, replicationFactor = 1))),
+      Map(existingTopic -> error(Errors.TOPIC_ALREADY_EXISTS,
+        Some("Topic 'existing-topic' already exists."))))
 
-    validateErrorCreateTopicsRequests(new CreateTopicsRequest.Builder(
-      Map("error-replication" -> new CreateTopicsRequest.TopicDetails(10, (numBrokers + 1).toShort)).asJava, timeout, true).build(),
+    validateErrorCreateTopicsRequests(topicsReq(Seq(topicReq("error-replication",
+      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(new CreateTopicsRequest.Builder(
-      Map("error-replication2" -> new CreateTopicsRequest.TopicDetails(10, -1: Short)).asJava, timeout, true).build(),
-      Map("error-replication2" -> error(Errors.INVALID_REPLICATION_FACTOR, Some("Replication factor must be larger than 0."))))
+    validateErrorCreateTopicsRequests(topicsReq(Seq(topicReq("error-replication2",
+      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."))))
   }
 
 }
@@ -140,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 d4d79e554c781..8f481f1c4b582 100644
--- a/core/src/test/scala/unit/kafka/server/DelayedOperationTest.scala
+++ b/core/src/test/scala/unit/kafka/server/DelayedOperationTest.scala
@@ -17,31 +17,76 @@
 
 package kafka.server
 
-import java.util.concurrent.{Executors, Future}
+import java.util.Random
+import java.util.concurrent._
+import java.util.concurrent.atomic.AtomicInteger
 import java.util.concurrent.locks.ReentrantLock
 
 import kafka.utils.CoreUtils.inLock
-
+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.jdk.CollectionConverters._
 
 class DelayedOperationTest {
 
-  var purgatory: DelayedOperationPurgatory[MockDelayedOperation] = null
+  var purgatory: DelayedOperationPurgatory[DelayedOperation] = null
+  var executorService: ExecutorService = null
 
   @Before
-  def setUp() {
-    purgatory = DelayedOperationPurgatory[MockDelayedOperation](purgatoryName = "mock")
+  def setUp(): Unit = {
+    purgatory = DelayedOperationPurgatory[DelayedOperation](purgatoryName = "mock")
   }
 
   @After
-  def tearDown() {
+  def tearDown(): Unit = {
     purgatory.shutdown()
+    if (executorService != null)
+      executorService.shutdown()
+  }
+
+  @Test
+  def testLockInTryCompleteElseWatch(): Unit = {
+    val op = new DelayedOperation(100000L) {
+      override def onExpiration(): Unit = {}
+      override def onComplete(): Unit = {}
+      override def tryComplete(): Boolean = {
+        assertTrue(lock.asInstanceOf[ReentrantLock].isHeldByCurrentThread)
+        false
+      }
+      override def safeTryComplete(): Boolean = {
+        fail("tryCompleteElseWatch should not use safeTryComplete")
+        super.safeTryComplete()
+      }
+    }
+    purgatory.tryCompleteElseWatch(op, Seq("key"))
+  }
+
+  @Test
+  def testSafeTryCompleteOrElse(): Unit = {
+    def op(shouldComplete: Boolean) = new DelayedOperation(100000L) {
+      override def onExpiration(): Unit = {}
+      override def onComplete(): Unit = {}
+      override def tryComplete(): Boolean = {
+        assertTrue(lock.asInstanceOf[ReentrantLock].isHeldByCurrentThread)
+        shouldComplete
+      }
+    }
+    var pass = false
+    assertFalse(op(false).safeTryCompleteOrElse {
+      pass = true
+    })
+    assertTrue(pass)
+    assertTrue(op(true).safeTryCompleteOrElse {
+      fail("this method should NOT be executed")
+    })
   }
 
   @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"))
@@ -58,7 +103,7 @@ class DelayedOperationTest {
   }
 
   @Test
-  def testRequestExpiry() {
+  def testRequestExpiry(): Unit = {
     val expiration = 20L
     val start = Time.SYSTEM.hiResClockMs
     val r1 = new MockDelayedOperation(expiration)
@@ -73,7 +118,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)
@@ -81,17 +188,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")
@@ -105,141 +212,169 @@ 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)
   }
 
+  /**
+    * Test `tryComplete` with multiple threads to verify that there are no timing windows
+    * when completion is not performed even if the thread that makes the operation completable
+    * may not be able to acquire the operation lock. Since it is difficult to test all scenarios,
+    * this test uses random delays with a large number of threads.
+    */
   @Test
-  def testDelayedOperationLock() {
-    verifyDelayedOperationLock(new MockDelayedOperation(100000L), mismatchedLocks = false)
-  }
+  def testTryCompleteWithMultipleThreads(): Unit = {
+    val executor = Executors.newScheduledThreadPool(20)
+    this.executorService = executor
+    val random = new Random
+    val maxDelayMs = 10
+    val completionAttempts = 20
+
+    class TestDelayOperation(index: Int) extends MockDelayedOperation(10000L) {
+      val key = s"key$index"
+      val completionAttemptsRemaining = new AtomicInteger(completionAttempts)
+
+      override def tryComplete(): Boolean = {
+        val shouldComplete = completable
+        Thread.sleep(random.nextInt(maxDelayMs))
+        if (shouldComplete)
+          forceComplete()
+        else
+          false
+      }
+    }
+    val ops = (0 until 100).map { index =>
+      val op = new TestDelayOperation(index)
+      purgatory.tryCompleteElseWatch(op, Seq(op.key))
+      op
+    }
 
-  @Test
-  def testDelayedOperationLockOverride() {
-    def newMockOperation = {
-      val lock = new ReentrantLock
-      new MockDelayedOperation(100000L, Some(lock), Some(lock))
+    def scheduleTryComplete(op: TestDelayOperation, delayMs: Long): Future[_] = {
+      executor.schedule(new Runnable {
+        override def run(): Unit = {
+          if (op.completionAttemptsRemaining.decrementAndGet() == 0)
+            op.completable = true
+          purgatory.checkAndComplete(op.key)
+        }
+      }, delayMs, TimeUnit.MILLISECONDS)
     }
-    verifyDelayedOperationLock(newMockOperation, mismatchedLocks = false)
 
-    verifyDelayedOperationLock(new MockDelayedOperation(100000L, None, Some(new ReentrantLock)),
-        mismatchedLocks = true)
+    (1 to completionAttempts).flatMap { _ =>
+      ops.map { op => scheduleTryComplete(op, random.nextInt(maxDelayMs)) }
+    }.foreach { future => future.get }
+
+    ops.foreach { op => assertTrue("Operation should have completed", op.isCompleted) }
   }
 
-  def verifyDelayedOperationLock(mockDelayedOperation: => MockDelayedOperation, mismatchedLocks: Boolean) {
+  def verifyDelayedOperationLock(mockDelayedOperation: => MockDelayedOperation, mismatchedLocks: Boolean): Unit = {
     val key = "key"
-    val executorService = Executors.newSingleThreadExecutor
-    try {
-      def createDelayedOperations(count: Int): Seq[MockDelayedOperation] = {
-        (1 to count).map { _ =>
-          val op = mockDelayedOperation
-          purgatory.tryCompleteElseWatch(op, Seq(key))
-          assertFalse("Not completable", op.isCompleted)
-          op
-        }
+    executorService = Executors.newSingleThreadExecutor
+    def createDelayedOperations(count: Int): Seq[MockDelayedOperation] = {
+      (1 to count).map { _ =>
+        val op = mockDelayedOperation
+        purgatory.tryCompleteElseWatch(op, Seq(key))
+        assertFalse("Not completable", op.isCompleted)
+        op
       }
+    }
 
-      def createCompletableOperations(count: Int): Seq[MockDelayedOperation] = {
-        (1 to count).map { _ =>
-          val op = mockDelayedOperation
-          op.completable = true
-          op
-        }
+    def createCompletableOperations(count: Int): Seq[MockDelayedOperation] = {
+      (1 to count).map { _ =>
+        val op = mockDelayedOperation
+        op.completable = true
+        op
       }
+    }
 
-      def runOnAnotherThread(fun: => Unit, shouldComplete: Boolean): Future[_] = {
-        val future = executorService.submit(new Runnable {
-          def run() = fun
-        })
-        if (shouldComplete)
-          future.get()
-        else
-          assertFalse("Should not have completed", future.isDone)
-        future
-      }
+    def checkAndComplete(completableOps: Seq[MockDelayedOperation], expectedComplete: Seq[MockDelayedOperation]): Unit = {
+      completableOps.foreach(op => op.completable = true)
+      val completed = purgatory.checkAndComplete(key)
+      assertEquals(expectedComplete.size, completed)
+      expectedComplete.foreach(op => assertTrue("Should have completed", op.isCompleted))
+      val expectedNotComplete = completableOps.toSet -- expectedComplete
+      expectedNotComplete.foreach(op => assertFalse("Should not have completed", op.isCompleted))
+    }
 
-      def checkAndComplete(completableOps: Seq[MockDelayedOperation], expectedComplete: Seq[MockDelayedOperation]): Unit = {
-        completableOps.foreach(op => op.completable = true)
-        val completed = purgatory.checkAndComplete(key)
-        assertEquals(expectedComplete.size, completed)
-        expectedComplete.foreach(op => assertTrue("Should have completed", op.isCompleted))
-        val expectedNotComplete = completableOps.toSet -- expectedComplete
-        expectedNotComplete.foreach(op => assertFalse("Should not have completed", op.isCompleted))
-      }
+    // If locks are free all completable operations should complete
+    var ops = createDelayedOperations(2)
+    checkAndComplete(ops, ops)
 
-      // If locks are free all completable operations should complete
-      var ops = createDelayedOperations(2)
+    // Lock held by current thread, completable operations should complete
+    ops = createDelayedOperations(2)
+    inLock(ops(1).lock) {
       checkAndComplete(ops, ops)
+    }
 
-      // Lock held by current thread, completable operations should complete
-      ops = createDelayedOperations(2)
-      inLock(ops(1).lock) {
-        checkAndComplete(ops, ops)
-      }
+    // Lock held by another thread, should not block, only operations that can be
+    // locked without blocking on the current thread should complete
+    ops = createDelayedOperations(2)
+    runOnAnotherThread(ops(0).lock.lock(), true)
+    try {
+      checkAndComplete(ops, Seq(ops(1)))
+    } finally {
+      runOnAnotherThread(ops(0).lock.unlock(), true)
+      checkAndComplete(Seq(ops(0)), Seq(ops(0)))
+    }
 
-      // Lock held by another thread, should not block, only operations that can be
-      // locked without blocking on the current thread should complete
-      ops = createDelayedOperations(2)
-      runOnAnotherThread(ops(0).lock.lock(), true)
+    // Lock acquired by response callback held by another thread, should not block
+    // if the response lock is used as operation lock, only operations
+    // that can be locked without blocking on the current thread should complete
+    ops = createDelayedOperations(2)
+    ops(0).responseLockOpt.foreach { lock =>
+      runOnAnotherThread(lock.lock(), true)
       try {
-        checkAndComplete(ops, Seq(ops(1)))
-      } finally {
-        runOnAnotherThread(ops(0).lock.unlock(), true)
-        checkAndComplete(Seq(ops(0)), Seq(ops(0)))
-      }
-
-      // Lock acquired by response callback held by another thread, should not block
-      // if the response lock is used as operation lock, only operations
-      // that can be locked without blocking on the current thread should complete
-      ops = createDelayedOperations(2)
-      ops(0).responseLockOpt.foreach { lock =>
-        runOnAnotherThread(lock.lock(), true)
         try {
-          try {
-            checkAndComplete(ops, Seq(ops(1)))
-            assertFalse("Should have failed with mismatched locks", mismatchedLocks)
-          } catch {
-            case e: IllegalStateException =>
-              assertTrue("Should not have failed with valid locks", mismatchedLocks)
-          }
-        } finally {
-          runOnAnotherThread(lock.unlock(), true)
-          checkAndComplete(Seq(ops(0)), Seq(ops(0)))
+          checkAndComplete(ops, Seq(ops(1)))
+          assertFalse("Should have failed with mismatched locks", mismatchedLocks)
+        } catch {
+          case e: IllegalStateException =>
+            assertTrue("Should not have failed with valid locks", mismatchedLocks)
         }
+      } finally {
+        runOnAnotherThread(lock.unlock(), true)
+        checkAndComplete(Seq(ops(0)), Seq(ops(0)))
       }
+    }
 
-      // Immediately completable operations should complete without locking
-      ops = createCompletableOperations(2)
-      ops.foreach { op =>
-        assertTrue("Should have completed", purgatory.tryCompleteElseWatch(op, Seq(key)))
-        assertTrue("Should have completed", op.isCompleted)
-      }
-
-    } finally {
-      executorService.shutdown()
+    // Immediately completable operations should complete without locking
+    ops = createCompletableOperations(2)
+    ops.foreach { op =>
+      assertTrue("Should have completed", purgatory.tryCompleteElseWatch(op, Seq(key)))
+      assertTrue("Should have completed", op.isCompleted)
     }
   }
 
+  private def runOnAnotherThread(fun: => Unit, shouldComplete: Boolean): Future[_] = {
+    val future = executorService.submit(new Runnable {
+      def run() = fun
+    })
+    if (shouldComplete)
+      future.get()
+    else
+      assertFalse("Should not have completed", future.isDone)
+    future
+  }
 
   class MockDelayedOperation(delayMs: Long,
-      lockOpt: Option[ReentrantLock] = None,
-      val responseLockOpt: Option[ReentrantLock] = None) extends DelayedOperation(delayMs, lockOpt) {
+                             lockOpt: Option[ReentrantLock] = None,
+                             val responseLockOpt: Option[ReentrantLock] = None)
+                             extends DelayedOperation(delayMs, lockOpt) {
     var completable = false
 
-    def awaitExpiration() {
+    def awaitExpiration(): Unit = {
       synchronized {
         wait()
       }
@@ -252,11 +387,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
new file mode 100644
index 0000000000000..a5025351e16a7
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/DelegationTokenRequestsOnPlainTextTest.scala
@@ -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 kafka.server
+
+import java.util
+
+import kafka.utils.TestUtils
+import org.apache.kafka.clients.admin.{Admin, AdminClientConfig}
+import org.apache.kafka.common.errors.UnsupportedByAuthenticationException
+import org.junit.{After, Before, Test}
+import org.scalatest.Assertions.intercept
+
+import scala.concurrent.ExecutionException
+
+class DelegationTokenRequestsOnPlainTextTest extends BaseRequestTest {
+  var adminClient: Admin = null
+
+  override def brokerCount = 1
+
+  @Before
+  override def setUp(): Unit = {
+    super.setUp()
+  }
+
+  def createAdminConfig: util.Map[String, Object] = {
+    val config = new util.HashMap[String, Object]
+    config.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList)
+    val securityProps: util.Map[Object, Object] =
+      TestUtils.adminClientSecurityConfigs(securityProtocol, trustStoreFile, clientSaslProperties)
+    securityProps.forEach { (key, value) => config.put(key.asInstanceOf[String], value) }
+    config
+  }
+
+  @Test
+  def testDelegationTokenRequests(): Unit = {
+    adminClient = Admin.create(createAdminConfig)
+
+    val createResult = adminClient.createDelegationToken()
+    intercept[ExecutionException](createResult.delegationToken().get()).getCause.isInstanceOf[UnsupportedByAuthenticationException]
+
+    val describeResult = adminClient.describeDelegationToken()
+    intercept[ExecutionException](describeResult.delegationTokens().get()).getCause.isInstanceOf[UnsupportedByAuthenticationException]
+
+    val renewResult = adminClient.renewDelegationToken("".getBytes())
+    intercept[ExecutionException](renewResult.expiryTimestamp().get()).getCause.isInstanceOf[UnsupportedByAuthenticationException]
+
+    val expireResult = adminClient.expireDelegationToken("".getBytes())
+    intercept[ExecutionException](expireResult.expiryTimestamp().get()).getCause.isInstanceOf[UnsupportedByAuthenticationException]
+  }
+
+
+  @After
+  override def tearDown(): Unit = {
+    if (adminClient != null)
+      adminClient.close()
+    super.tearDown()
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/server/DelegationTokenRequestsTest.scala b/core/src/test/scala/unit/kafka/server/DelegationTokenRequestsTest.scala
new file mode 100644
index 0000000000000..96115d6adfc2b
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/DelegationTokenRequestsTest.scala
@@ -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 kafka.server
+
+import java.util
+
+import kafka.api.{KafkaSasl, SaslSetup}
+import kafka.utils.{JaasTestUtils, TestUtils}
+import org.apache.kafka.clients.admin.{Admin, 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.jdk.CollectionConverters._
+import scala.concurrent.ExecutionException
+
+class DelegationTokenRequestsTest extends BaseRequestTest with SaslSetup {
+  override protected def securityProtocol = SecurityProtocol.SASL_PLAINTEXT
+  private val kafkaClientSaslMechanism = "PLAIN"
+  private val kafkaServerSaslMechanisms = List("PLAIN")
+  protected override val serverSaslProperties = Some(kafkaServerSaslProperties(kafkaServerSaslMechanisms, kafkaClientSaslMechanism))
+  protected override val clientSaslProperties = Some(kafkaClientSaslProperties(kafkaClientSaslMechanism))
+  var adminClient: Admin = null
+
+  override def brokerCount = 1
+
+  @Before
+  override def setUp(): Unit = {
+    startSasl(jaasSections(kafkaServerSaslMechanisms, Some(kafkaClientSaslMechanism), KafkaSasl, JaasTestUtils.KafkaServerContextName))
+    super.setUp()
+  }
+
+  override def generateConfigs = {
+    val props = TestUtils.createBrokerConfigs(brokerCount, zkConnect,
+      enableControlledShutdown = false,
+      interBrokerSecurityProtocol = Some(securityProtocol),
+      trustStoreFile = trustStoreFile, saslProperties = serverSaslProperties, enableToken = true)
+    props.foreach(brokerPropertyOverrides)
+    props.map(KafkaConfig.fromProps)
+  }
+
+  private def createAdminConfig: util.Map[String, Object] = {
+    val config = new util.HashMap[String, Object]
+    config.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList)
+    val securityProps: util.Map[Object, Object] =
+      TestUtils.adminClientSecurityConfigs(securityProtocol, trustStoreFile, clientSaslProperties)
+    securityProps.forEach { (key, value) => config.put(key.asInstanceOf[String], value) }
+    config
+  }
+
+  @Test
+  def testDelegationTokenRequests(): Unit = {
+    adminClient = Admin.create(createAdminConfig)
+
+    // create token1 with renewer1
+    val renewer1 = List(SecurityUtils.parseKafkaPrincipal("User:renewer1")).asJava
+    val createResult1 = adminClient.createDelegationToken(new CreateDelegationTokenOptions().renewers(renewer1))
+    val tokenCreated = createResult1.delegationToken().get()
+
+    //test describe token
+    var tokens = adminClient.describeDelegationToken().delegationTokens().get()
+    assertTrue(tokens.size() == 1)
+    var token1 = tokens.get(0)
+    assertEquals(token1, tokenCreated)
+
+    // create token2 with renewer2
+    val renewer2 = List(SecurityUtils.parseKafkaPrincipal("User:renewer2")).asJava
+    val createResult2 = adminClient.createDelegationToken(new CreateDelegationTokenOptions().renewers(renewer2))
+    val token2 = createResult2.delegationToken().get()
+
+    //get all tokens
+    tokens = adminClient.describeDelegationToken().delegationTokens().get()
+    assertTrue(tokens.size() == 2)
+    assertEquals(Set(token1, token2), tokens.asScala.toSet)
+
+    //get tokens for renewer2
+    tokens = adminClient.describeDelegationToken(new DescribeDelegationTokenOptions().owners(renewer2)).delegationTokens().get()
+    assertTrue(tokens.size() == 1)
+    assertEquals(Set(token2), tokens.asScala.toSet)
+
+    //test renewing tokens
+    val renewResult = adminClient.renewDelegationToken(token1.hmac())
+    var expiryTimestamp = renewResult.expiryTimestamp().get()
+
+    val describeResult = adminClient.describeDelegationToken()
+    val tokenId = token1.tokenInfo().tokenId()
+    token1 = describeResult.delegationTokens().get().asScala.filter(dt => dt.tokenInfo().tokenId() == tokenId).head
+    assertEquals(expiryTimestamp, token1.tokenInfo().expiryTimestamp())
+
+    //test expire tokens
+    val expireResult1 = adminClient.expireDelegationToken(token1.hmac())
+    expiryTimestamp = expireResult1.expiryTimestamp().get()
+
+    val expireResult2 = adminClient.expireDelegationToken(token2.hmac())
+    expiryTimestamp = expireResult2.expiryTimestamp().get()
+
+    tokens = adminClient.describeDelegationToken().delegationTokens().get()
+    assertTrue(tokens.size == 0)
+
+    //create token with invalid principal type
+    val renewer3 = List(SecurityUtils.parseKafkaPrincipal("Group:Renewer3")).asJava
+    val createResult3 = adminClient.createDelegationToken(new CreateDelegationTokenOptions().renewers(renewer3))
+    intercept[ExecutionException](createResult3.delegationToken().get()).getCause.isInstanceOf[InvalidPrincipalTypeException]
+
+    // try describing tokens for unknown owner
+    val unknownOwner = List(SecurityUtils.parseKafkaPrincipal("User:Unknown")).asJava
+    tokens = adminClient.describeDelegationToken(new DescribeDelegationTokenOptions().owners(unknownOwner)).delegationTokens().get()
+    assertTrue(tokens.isEmpty)
+  }
+
+  @After
+  override def tearDown(): Unit = {
+    if (adminClient != null)
+      adminClient.close()
+    super.tearDown()
+    closeSasl()
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/server/DelegationTokenRequestsWithDisableTokenFeatureTest.scala b/core/src/test/scala/unit/kafka/server/DelegationTokenRequestsWithDisableTokenFeatureTest.scala
new file mode 100644
index 0000000000000..12c2138fbc4f7
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/DelegationTokenRequestsWithDisableTokenFeatureTest.scala
@@ -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 kafka.server
+
+import java.util
+
+import kafka.api.{KafkaSasl, SaslSetup}
+import kafka.utils.{JaasTestUtils, TestUtils}
+import org.apache.kafka.clients.admin.{Admin, 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.concurrent.ExecutionException
+
+class DelegationTokenRequestsWithDisableTokenFeatureTest extends BaseRequestTest with SaslSetup {
+  override protected def securityProtocol = SecurityProtocol.SASL_PLAINTEXT
+  private val kafkaClientSaslMechanism = "PLAIN"
+  private val kafkaServerSaslMechanisms = List("PLAIN")
+  protected override val serverSaslProperties = Some(kafkaServerSaslProperties(kafkaServerSaslMechanisms, kafkaClientSaslMechanism))
+  protected override val clientSaslProperties = Some(kafkaClientSaslProperties(kafkaClientSaslMechanism))
+  var adminClient: Admin = null
+
+  override def brokerCount = 1
+
+  @Before
+  override def setUp(): Unit = {
+    startSasl(jaasSections(kafkaServerSaslMechanisms, Some(kafkaClientSaslMechanism), KafkaSasl, JaasTestUtils.KafkaServerContextName))
+    super.setUp()
+  }
+
+  def createAdminConfig: util.Map[String, Object] = {
+    val config = new util.HashMap[String, Object]
+    config.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList)
+    val securityProps: util.Map[Object, Object] =
+      TestUtils.adminClientSecurityConfigs(securityProtocol, trustStoreFile, clientSaslProperties)
+    securityProps.forEach { (key, value) => config.put(key.asInstanceOf[String], value) }
+    config
+  }
+
+  @Test
+  def testDelegationTokenRequests(): Unit = {
+    adminClient = Admin.create(createAdminConfig)
+
+    val createResult = adminClient.createDelegationToken()
+    intercept[ExecutionException](createResult.delegationToken().get()).getCause.isInstanceOf[DelegationTokenDisabledException]
+
+    val describeResult = adminClient.describeDelegationToken()
+    intercept[ExecutionException](describeResult.delegationTokens().get()).getCause.isInstanceOf[DelegationTokenDisabledException]
+
+    val renewResult = adminClient.renewDelegationToken("".getBytes())
+    intercept[ExecutionException](renewResult.expiryTimestamp().get()).getCause.isInstanceOf[DelegationTokenDisabledException]
+
+    val expireResult = adminClient.expireDelegationToken("".getBytes())
+    intercept[ExecutionException](expireResult.expiryTimestamp().get()).getCause.isInstanceOf[DelegationTokenDisabledException]
+  }
+
+  @After
+  override def tearDown(): Unit = {
+    if (adminClient != null)
+      adminClient.close()
+    super.tearDown()
+    closeSasl()
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/server/DeleteTopicsRequestTest.scala b/core/src/test/scala/unit/kafka/server/DeleteTopicsRequestTest.scala
index 881bf8ee790cc..2994fd926e4fe 100644
--- a/core/src/test/scala/unit/kafka/server/DeleteTopicsRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/DeleteTopicsRequestTest.scala
@@ -17,54 +17,65 @@
 
 package kafka.server
 
+import java.util.{Arrays, Collections}
+
 import kafka.network.SocketServer
 import kafka.utils._
-import org.apache.kafka.common.protocol.{ApiKeys, Errors}
+import org.apache.kafka.common.message.DeleteTopicsRequestData
+import org.apache.kafka.common.protocol.Errors
 import org.apache.kafka.common.requests.{DeleteTopicsRequest, DeleteTopicsResponse, MetadataRequest, MetadataResponse}
 import org.junit.Assert._
 import org.junit.Test
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 
 class DeleteTopicsRequestTest extends BaseRequestTest {
 
   @Test
-  def testValidDeleteTopicRequests() {
+  def testValidDeleteTopicRequests(): Unit = {
     val timeout = 10000
     // Single topic
-    TestUtils.createTopic(zkUtils, "topic-1", 1, 1, servers)
-    validateValidDeleteTopicRequests(new DeleteTopicsRequest.Builder(Set("topic-1").asJava, timeout).build())
+    createTopic("topic-1", 1, 1)
+    validateValidDeleteTopicRequests(new DeleteTopicsRequest.Builder(
+        new DeleteTopicsRequestData()
+          .setTopicNames(Arrays.asList("topic-1"))
+          .setTimeoutMs(timeout)).build())
     // Multi topic
-    TestUtils.createTopic(zkUtils, "topic-3", 5, 2, servers)
-    TestUtils.createTopic(zkUtils, "topic-4", 1, 2, servers)
-    validateValidDeleteTopicRequests(new DeleteTopicsRequest.Builder(Set("topic-3", "topic-4").asJava, timeout).build())
+    createTopic("topic-3", 5, 2)
+    createTopic("topic-4", 1, 2)
+    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.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
-    TestUtils.createTopic(zkUtils, "partial-topic-1", 1, 1, servers)
-    validateErrorDeleteTopicRequests(new DeleteTopicsRequest.Builder(Set(
-      "partial-topic-1",
-      "partial-invalid-topic").asJava, timeout).build(),
+    createTopic("partial-topic-1", 1, 1)
+    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
@@ -72,9 +83,12 @@ class DeleteTopicsRequestTest extends BaseRequestTest {
     )
 
     // Timeout
-    TestUtils.createTopic(zkUtils, timeoutTopic, 5, 2, servers)
+    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,28 +112,26 @@ 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 = {
-    val metadata = sendMetadataRequest(new MetadataRequest.
-        Builder(List(topic).asJava, true).build).topicMetadata.asScala
+    val metadata = connectAndReceive[MetadataResponse](new MetadataRequest.Builder(
+      List(topic).asJava, true).build).topicMetadata.asScala
     TestUtils.waitUntilTrue (() => !metadata.exists(p => p.topic.equals(topic) && p.error == Errors.NONE),
       s"The topic $topic should not exist")
   }
 
   private def sendDeleteTopicsRequest(request: DeleteTopicsRequest, socketServer: SocketServer = controllerSocketServer): DeleteTopicsResponse = {
-    val response = connectAndSend(request, ApiKeys.DELETE_TOPICS, socketServer)
-    DeleteTopicsResponse.parse(response, request.version)
+    connectAndReceive[DeleteTopicsResponse](request, destination = socketServer)
   }
 
-  private def sendMetadataRequest(request: MetadataRequest): MetadataResponse = {
-    val response = connectAndSend(request, ApiKeys.METADATA)
-    MetadataResponse.parse(response, request.version)
-  }
 }
diff --git a/core/src/test/scala/unit/kafka/server/DeleteTopicsRequestWithDeletionDisabledTest.scala b/core/src/test/scala/unit/kafka/server/DeleteTopicsRequestWithDeletionDisabledTest.scala
new file mode 100644
index 0000000000000..f9512f66dc59f
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/DeleteTopicsRequestWithDeletionDisabledTest.scala
@@ -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 kafka.server
+
+import java.util.Collections
+
+import kafka.utils._
+import org.apache.kafka.common.message.DeleteTopicsRequestData
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.requests.{DeleteTopicsRequest, DeleteTopicsResponse}
+import org.junit.Assert._
+import org.junit.Test
+
+class DeleteTopicsRequestWithDeletionDisabledTest extends BaseRequestTest {
+
+  override def brokerCount: Int = 1
+
+  override def generateConfigs = {
+    val props = TestUtils.createBrokerConfigs(brokerCount, zkConnect,
+      enableControlledShutdown = false, enableDeleteTopic = false,
+      interBrokerSecurityProtocol = Some(securityProtocol),
+      trustStoreFile = trustStoreFile, saslProperties = serverSaslProperties, logDirCount = logDirCount)
+    props.foreach(brokerPropertyOverrides)
+    props.map(KafkaConfig.fromProps)
+  }
+
+  @Test
+  def testDeleteRecordsRequest(): Unit = {
+    val topic = "topic-1"
+    val request = new DeleteTopicsRequest.Builder(
+        new DeleteTopicsRequestData()
+          .setTopicNames(Collections.singletonList(topic))
+          .setTimeoutMs(1000)).build()
+    val response = sendDeleteTopicsRequest(request)
+    assertEquals(Errors.TOPIC_DELETION_DISABLED.code, response.data.responses.find(topic).errorCode)
+
+    val v2request = new DeleteTopicsRequest.Builder(
+        new DeleteTopicsRequestData()
+        .setTopicNames(Collections.singletonList(topic))
+        .setTimeoutMs(1000)).build(2)
+    val v2response = sendDeleteTopicsRequest(v2request)
+    assertEquals(Errors.INVALID_REQUEST.code, v2response.data.responses.find(topic).errorCode)
+  }
+
+  private def sendDeleteTopicsRequest(request: DeleteTopicsRequest): DeleteTopicsResponse = {
+    connectAndReceive[DeleteTopicsResponse](request, destination = controllerSocketServer)
+  }
+
+}
diff --git a/core/src/test/scala/unit/kafka/server/DescribeLogDirsRequestTest.scala b/core/src/test/scala/unit/kafka/server/DescribeLogDirsRequestTest.scala
index 7f9e63a266d2a..13de19f375dc2 100644
--- a/core/src/test/scala/unit/kafka/server/DescribeLogDirsRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/DescribeLogDirsRequestTest.scala
@@ -17,18 +17,21 @@
 
 package kafka.server
 
+import java.io.File
+
 import kafka.utils._
 import org.apache.kafka.common.TopicPartition
-import org.apache.kafka.common.protocol.{ApiKeys, Errors}
+import org.apache.kafka.common.message.DescribeLogDirsRequestData
+import org.apache.kafka.common.protocol.Errors
 import org.apache.kafka.common.requests._
 import org.junit.Assert._
 import org.junit.Test
-import java.io.File
 
-class DescribeLogDirsRequestTest extends BaseRequestTest {
+import scala.jdk.CollectionConverters._
 
-  override def numBrokers: Int = 1
-  override def logDirCount: Int = 2
+class DescribeLogDirsRequestTest extends BaseRequestTest {
+  override val logDirCount = 2
+  override val brokerCount: Int = 1
 
   val topic = "topic"
   val partitionNum = 2
@@ -40,25 +43,32 @@ class DescribeLogDirsRequestTest extends BaseRequestTest {
     val onlineDir = new File(servers.head.config.logDirs.head).getAbsolutePath
     val offlineDir = new File(servers.head.config.logDirs.tail.head).getAbsolutePath
     servers.head.replicaManager.handleLogDirFailure(offlineDir)
-    TestUtils.createTopic(zkUtils, topic, partitionNum, 1, servers)
-    TestUtils.produceMessages(servers, topic, 10)
+    createTopic(topic, partitionNum, 1)
+    TestUtils.generateAndProduceMessages(servers, topic, 10)
 
-    val request = new DescribeLogDirsRequest.Builder(null).build()
-    val response = connectAndSend(request, ApiKeys.DESCRIBE_LOG_DIRS, controllerSocketServer)
-    val logDirInfos = DescribeLogDirsResponse.parse(response, request.version).logDirInfos()
+    val request = new DescribeLogDirsRequest.Builder(new DescribeLogDirsRequestData().setTopics(null)).build()
+    val response = connectAndReceive[DescribeLogDirsResponse](request, destination = controllerSocketServer)
 
-    assertEquals(logDirCount, logDirInfos.size())
-    assertEquals(Errors.KAFKA_STORAGE_ERROR, logDirInfos.get(offlineDir).error)
-    assertEquals(0, logDirInfos.get(offlineDir).replicaInfos.size())
+    assertEquals(logDirCount, response.data.results.size)
+    val offlineResult = response.data.results.asScala.find(logDirResult => logDirResult.logDir == offlineDir).get
+    assertEquals(Errors.KAFKA_STORAGE_ERROR.code, offlineResult.errorCode)
+    assertEquals(0, offlineResult.topics.asScala.map(t => t.partitions().size()).sum)
 
-    assertEquals(Errors.NONE, logDirInfos.get(onlineDir).error)
-    val replicaInfo0 = logDirInfos.get(onlineDir).replicaInfos.get(tp0)
-    val replicaInfo1 = logDirInfos.get(onlineDir).replicaInfos.get(tp1)
+    val onlineResult = response.data.results.asScala.find(logDirResult => logDirResult.logDir == onlineDir).get
+    assertEquals(Errors.NONE.code, onlineResult.errorCode)
+    val onlinePartitionsMap = onlineResult.topics.asScala.flatMap { topic =>
+      topic.partitions().asScala.map { partitionResult =>
+        new TopicPartition(topic.name, partitionResult.partitionIndex) -> partitionResult
+      }
+    }.toMap
+    val replicaInfo0 = onlinePartitionsMap(tp0)
+    val replicaInfo1 = onlinePartitionsMap(tp1)
     val log0 = servers.head.logManager.getLog(tp0).get
     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)
+    assertEquals(log0.size, replicaInfo0.partitionSize)
+    assertEquals(log1.size, replicaInfo1.partitionSize)
+    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/DescribeUserScramCredentialsRequestTest.scala b/core/src/test/scala/unit/kafka/server/DescribeUserScramCredentialsRequestTest.scala
new file mode 100644
index 0000000000000..299fa6f5b1b2c
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/DescribeUserScramCredentialsRequestTest.scala
@@ -0,0 +1,162 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package kafka.server
+
+import java.util
+import java.util.Properties
+
+import kafka.network.SocketServer
+import kafka.security.authorizer.AclAuthorizer
+import org.apache.kafka.common.message.{DescribeUserScramCredentialsRequestData, DescribeUserScramCredentialsResponseData}
+import org.apache.kafka.common.message.DescribeUserScramCredentialsRequestData.UserName
+import org.apache.kafka.common.protocol.{ApiKeys, Errors}
+import org.apache.kafka.common.requests.{DescribeUserScramCredentialsRequest, DescribeUserScramCredentialsResponse}
+import org.apache.kafka.common.security.auth.{AuthenticationContext, KafkaPrincipal, KafkaPrincipalBuilder}
+import org.apache.kafka.server.authorizer.{Action, AuthorizableRequestContext, AuthorizationResult}
+import org.junit.Assert._
+import org.junit.rules.TestName
+import org.junit.{Rule, Test}
+
+import scala.jdk.CollectionConverters._
+
+/**
+ * Test DescribeUserScramCredentialsRequest/Response API for the cases where no credentials exist
+ * or failure is expected due to lack of authorization, sending the request to a non-controller broker, or some other issue.
+ * Testing the API for the case where there are actually credentials to describe is performed elsewhere.
+ */
+class DescribeUserScramCredentialsRequestTest extends BaseRequestTest {
+  override def brokerPropertyOverrides(properties: Properties): Unit = {
+    properties.put(KafkaConfig.ControlledShutdownEnableProp, "false")
+    properties.put(KafkaConfig.AuthorizerClassNameProp, classOf[DescribeCredentialsTest.TestAuthorizer].getName)
+    properties.put(KafkaConfig.PrincipalBuilderClassProp,
+      if (testName.getMethodName.endsWith("NotAuthorized")) {
+        classOf[DescribeCredentialsTest.TestPrincipalBuilderReturningUnauthorized].getName
+      } else {
+        classOf[DescribeCredentialsTest.TestPrincipalBuilderReturningAuthorized].getName
+      })
+  }
+
+  private val _testName = new TestName
+  @Rule def testName = _testName
+
+  @Test
+  def testDescribeNothing(): Unit = {
+    val request = new DescribeUserScramCredentialsRequest.Builder(
+      new DescribeUserScramCredentialsRequestData()).build()
+    val response = sendDescribeUserScramCredentialsRequest(request)
+
+    val error = response.data.errorCode
+    assertEquals("Expected no error when describing everything and there are no credentials",
+      Errors.NONE.code, error)
+    assertEquals("Expected no credentials when describing everything and there are no credentials",
+      0, response.data.results.size)
+  }
+
+  @Test
+  def testDescribeWithNull(): Unit = {
+    val request = new DescribeUserScramCredentialsRequest.Builder(
+      new DescribeUserScramCredentialsRequestData().setUsers(null)).build()
+    val response = sendDescribeUserScramCredentialsRequest(request)
+
+    val error = response.data.errorCode
+    assertEquals("Expected no error when describing everything and there are no credentials",
+      Errors.NONE.code, error)
+    assertEquals("Expected no credentials when describing everything and there are no credentials",
+      0, response.data.results.size)
+  }
+
+  @Test
+  def testDescribeNotController(): Unit = {
+    val request = new DescribeUserScramCredentialsRequest.Builder(
+      new DescribeUserScramCredentialsRequestData()).build()
+    val response = sendDescribeUserScramCredentialsRequest(request, notControllerSocketServer)
+
+    val error = response.data.errorCode
+    assertEquals("Did not expect controller error when routed to non-controller", Errors.NONE.code, error)
+  }
+
+  @Test
+  def testDescribeNotAuthorized(): Unit = {
+    val request = new DescribeUserScramCredentialsRequest.Builder(
+      new DescribeUserScramCredentialsRequestData()).build()
+    val response = sendDescribeUserScramCredentialsRequest(request)
+
+    val error = response.data.errorCode
+    assertEquals("Expected not authorized error", Errors.CLUSTER_AUTHORIZATION_FAILED.code, error)
+  }
+
+  @Test
+  def testDescribeSameUserTwice(): Unit = {
+    val user = "user1"
+    val userName = new UserName().setName(user)
+    val request = new DescribeUserScramCredentialsRequest.Builder(
+      new DescribeUserScramCredentialsRequestData().setUsers(List(userName, userName).asJava)).build()
+    val response = sendDescribeUserScramCredentialsRequest(request)
+
+    assertEquals("Expected no top-level error", Errors.NONE.code, response.data.errorCode)
+    assertEquals(1, response.data.results.size)
+    val result: DescribeUserScramCredentialsResponseData.DescribeUserScramCredentialsResult = response.data.results.get(0)
+    assertEquals(s"Expected duplicate resource error for $user", Errors.DUPLICATE_RESOURCE.code, result.errorCode)
+    assertEquals(s"Cannot describe SCRAM credentials for the same user twice in a single request: $user", result.errorMessage)
+  }
+
+  @Test
+  def testUnknownUser(): Unit = {
+    val unknownUser = "unknownUser"
+    val request = new DescribeUserScramCredentialsRequest.Builder(
+      new DescribeUserScramCredentialsRequestData().setUsers(List(new UserName().setName(unknownUser)).asJava)).build()
+    val response = sendDescribeUserScramCredentialsRequest(request)
+
+    assertEquals("Expected no top-level error", Errors.NONE.code, response.data.errorCode)
+    assertEquals(1, response.data.results.size)
+    val result: DescribeUserScramCredentialsResponseData.DescribeUserScramCredentialsResult = response.data.results.get(0)
+    assertEquals(s"Expected duplicate resource error for $unknownUser", Errors.RESOURCE_NOT_FOUND.code, result.errorCode)
+    assertEquals(s"Attempt to describe a user credential that does not exist: $unknownUser", result.errorMessage)
+  }
+
+  private def sendDescribeUserScramCredentialsRequest(request: DescribeUserScramCredentialsRequest, socketServer: SocketServer = controllerSocketServer): DescribeUserScramCredentialsResponse = {
+    connectAndReceive[DescribeUserScramCredentialsResponse](request, destination = socketServer)
+  }
+}
+
+object DescribeCredentialsTest {
+  val UnauthorizedPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "Unauthorized")
+  val AuthorizedPrincipal = KafkaPrincipal.ANONYMOUS
+
+  class TestAuthorizer extends AclAuthorizer {
+    override def authorize(requestContext: AuthorizableRequestContext, actions: util.List[Action]): util.List[AuthorizationResult] = {
+      actions.asScala.map { _ =>
+        if (requestContext.requestType == ApiKeys.DESCRIBE_USER_SCRAM_CREDENTIALS.id && requestContext.principal == UnauthorizedPrincipal)
+          AuthorizationResult.DENIED
+        else
+          AuthorizationResult.ALLOWED
+      }.asJava
+    }
+  }
+
+  class TestPrincipalBuilderReturningAuthorized extends KafkaPrincipalBuilder {
+    override def build(context: AuthenticationContext): KafkaPrincipal = {
+      AuthorizedPrincipal
+    }
+  }
+
+  class TestPrincipalBuilderReturningUnauthorized extends KafkaPrincipalBuilder {
+    override def build(context: AuthenticationContext): KafkaPrincipal = {
+      UnauthorizedPrincipal
+    }
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/server/DynamicBrokerConfigTest.scala b/core/src/test/scala/unit/kafka/server/DynamicBrokerConfigTest.scala
new file mode 100755
index 0000000000000..ef115dbce8eb3
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/DynamicBrokerConfigTest.scala
@@ -0,0 +1,450 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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.{lang, util}
+import java.util.Properties
+import java.util.concurrent.CompletionStage
+
+import kafka.utils.TestUtils
+import kafka.zk.KafkaZkClient
+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.Assertions.intercept
+
+import scala.jdk.CollectionConverters._
+import scala.collection.Set
+
+class DynamicBrokerConfigTest {
+
+  @Test
+  def testConfigUpdate(): Unit = {
+    val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
+    val oldKeystore = "oldKs.jks"
+    props.put(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG, oldKeystore)
+    val config = KafkaConfig(props)
+    val dynamicConfig = config.dynamicConfig
+    assertSame(config, dynamicConfig.currentKafkaConfig)
+    assertEquals(oldKeystore, config.values.get(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG))
+    assertEquals(oldKeystore,
+      config.valuesFromThisConfigWithPrefixOverride("listener.name.external.").get(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG))
+    assertEquals(oldKeystore, config.originalsFromThisConfig.get(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG))
+
+    (1 to 2).foreach { i =>
+      val props1 = new Properties
+      val newKeystore = s"ks$i.jks"
+      props1.put(s"listener.name.external.${SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG}", newKeystore)
+      dynamicConfig.updateBrokerConfig(0, props1)
+      assertNotSame(config, dynamicConfig.currentKafkaConfig)
+
+      assertEquals(newKeystore,
+        config.valuesWithPrefixOverride("listener.name.external.").get(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG))
+      assertEquals(newKeystore,
+        config.originalsWithPrefix("listener.name.external.").get(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG))
+      assertEquals(newKeystore,
+        config.valuesWithPrefixOverride("listener.name.external.").get(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG))
+      assertEquals(newKeystore,
+        config.originalsWithPrefix("listener.name.external.").get(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG))
+
+      assertEquals(oldKeystore, config.getString(KafkaConfig.SslKeystoreLocationProp))
+      assertEquals(oldKeystore, config.originals.get(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG))
+      assertEquals(oldKeystore, config.values.get(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG))
+      assertEquals(oldKeystore, config.originalsStrings.get(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG))
+
+      assertEquals(oldKeystore,
+        config.valuesFromThisConfigWithPrefixOverride("listener.name.external.").get(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG))
+      assertEquals(oldKeystore, config.originalsFromThisConfig.get(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG))
+      assertEquals(oldKeystore, config.valuesFromThisConfig.get(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG))
+      assertEquals(oldKeystore, config.originalsFromThisConfig.get(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG))
+      assertEquals(oldKeystore, config.valuesFromThisConfig.get(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG))
+    }
+  }
+
+  @Test
+  def testConfigUpdateWithSomeInvalidConfigs(): Unit = {
+    val origProps = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
+    origProps.put(SslConfigs.SSL_KEYSTORE_TYPE_CONFIG, "JKS")
+    val config = KafkaConfig(origProps)
+
+    val validProps = Map(s"listener.name.external.${SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG}" -> "ks.p12")
+
+    val securityPropsWithoutListenerPrefix = Map(SslConfigs.SSL_KEYSTORE_TYPE_CONFIG -> "PKCS12")
+    verifyConfigUpdateWithInvalidConfig(config, origProps, validProps, securityPropsWithoutListenerPrefix)
+    val nonDynamicProps = Map(KafkaConfig.ZkConnectProp -> "somehost:2181")
+    verifyConfigUpdateWithInvalidConfig(config, origProps, validProps, nonDynamicProps)
+
+    // Test update of configs with invalid type
+    val invalidProps = Map(KafkaConfig.LogCleanerThreadsProp -> "invalid")
+    verifyConfigUpdateWithInvalidConfig(config, origProps, validProps, invalidProps)
+  }
+
+  @Test
+  def testConfigUpdateWithReconfigurableValidationFailure(): Unit = {
+    val origProps = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
+    origProps.put(KafkaConfig.LogCleanerDedupeBufferSizeProp, "100000000")
+    val config = KafkaConfig(origProps)
+    val validProps = Map.empty[String, String]
+    val invalidProps = Map(KafkaConfig.LogCleanerThreadsProp -> "20")
+
+    def validateLogCleanerConfig(configs: util.Map[String, _]): Unit = {
+      val cleanerThreads = configs.get(KafkaConfig.LogCleanerThreadsProp).toString.toInt
+      if (cleanerThreads <=0 || cleanerThreads >= 5)
+        throw new ConfigException(s"Invalid cleaner threads $cleanerThreads")
+    }
+    val reconfigurable = new Reconfigurable {
+      override def configure(configs: util.Map[String, _]): Unit = {}
+      override def reconfigurableConfigs(): util.Set[String] = Set(KafkaConfig.LogCleanerThreadsProp).asJava
+      override def validateReconfiguration(configs: util.Map[String, _]): Unit = validateLogCleanerConfig(configs)
+      override def reconfigure(configs: util.Map[String, _]): Unit = {}
+    }
+    config.dynamicConfig.addReconfigurable(reconfigurable)
+    verifyConfigUpdateWithInvalidConfig(config, origProps, validProps, invalidProps)
+    config.dynamicConfig.removeReconfigurable(reconfigurable)
+
+    val brokerReconfigurable = new BrokerReconfigurable {
+      override def reconfigurableConfigs: collection.Set[String] = Set(KafkaConfig.LogCleanerThreadsProp)
+      override def validateReconfiguration(newConfig: KafkaConfig): Unit = validateLogCleanerConfig(newConfig.originals)
+      override def reconfigure(oldConfig: KafkaConfig, newConfig: KafkaConfig): Unit = {}
+    }
+    config.dynamicConfig.addBrokerReconfigurable(brokerReconfigurable)
+    verifyConfigUpdateWithInvalidConfig(config, origProps, validProps, invalidProps)
+  }
+
+  @Test
+  def testReconfigurableValidation(): Unit = {
+    val origProps = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
+    val config = KafkaConfig(origProps)
+    val invalidReconfigurableProps = Set(KafkaConfig.LogCleanerThreadsProp, KafkaConfig.BrokerIdProp, "some.prop")
+    val validReconfigurableProps = Set(KafkaConfig.LogCleanerThreadsProp, KafkaConfig.LogCleanerDedupeBufferSizeProp, "some.prop")
+
+    def createReconfigurable(configs: Set[String]) = new Reconfigurable {
+      override def configure(configs: util.Map[String, _]): Unit = {}
+      override def reconfigurableConfigs(): util.Set[String] = configs.asJava
+      override def validateReconfiguration(configs: util.Map[String, _]): Unit = {}
+      override def reconfigure(configs: util.Map[String, _]): Unit = {}
+    }
+    intercept[IllegalArgumentException] {
+      config.dynamicConfig.addReconfigurable(createReconfigurable(invalidReconfigurableProps))
+    }
+    config.dynamicConfig.addReconfigurable(createReconfigurable(validReconfigurableProps))
+
+    def createBrokerReconfigurable(configs: Set[String]) = new BrokerReconfigurable {
+      override def reconfigurableConfigs: collection.Set[String] = configs
+      override def validateReconfiguration(newConfig: KafkaConfig): Unit = {}
+      override def reconfigure(oldConfig: KafkaConfig, newConfig: KafkaConfig): Unit = {}
+    }
+    intercept[IllegalArgumentException] {
+      config.dynamicConfig.addBrokerReconfigurable(createBrokerReconfigurable(invalidReconfigurableProps))
+    }
+    config.dynamicConfig.addBrokerReconfigurable(createBrokerReconfigurable(validReconfigurableProps))
+  }
+
+  @Test
+  def testSecurityConfigs(): Unit = {
+    def verifyUpdate(name: String, value: Object): Unit = {
+      verifyConfigUpdate(name, value, perBrokerConfig = true, expectFailure = true)
+      verifyConfigUpdate(s"listener.name.external.$name", value, perBrokerConfig = true, expectFailure = false)
+      verifyConfigUpdate(name, value, perBrokerConfig = false, expectFailure = true)
+      verifyConfigUpdate(s"listener.name.external.$name", value, perBrokerConfig = false, expectFailure = true)
+    }
+
+    verifyUpdate(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG, "ks.jks")
+    verifyUpdate(SslConfigs.SSL_KEYSTORE_TYPE_CONFIG, "JKS")
+    verifyUpdate(SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG, "password")
+    verifyUpdate(SslConfigs.SSL_KEY_PASSWORD_CONFIG, "password")
+  }
+
+  @Test
+  def testConnectionQuota(): Unit = {
+    verifyConfigUpdate(KafkaConfig.MaxConnectionsPerIpProp, "100", perBrokerConfig = true, expectFailure = false)
+    verifyConfigUpdate(KafkaConfig.MaxConnectionsPerIpProp, "100", perBrokerConfig = false, expectFailure = false)
+    //MaxConnectionsPerIpProp can be set to zero only if MaxConnectionsPerIpOverridesProp property is set
+    verifyConfigUpdate(KafkaConfig.MaxConnectionsPerIpProp, "0", perBrokerConfig = false, expectFailure = true)
+
+    verifyConfigUpdate(KafkaConfig.MaxConnectionsPerIpOverridesProp, "hostName1:100,hostName2:0", perBrokerConfig = true,
+      expectFailure = false)
+    verifyConfigUpdate(KafkaConfig.MaxConnectionsPerIpOverridesProp, "hostName1:100,hostName2:0", perBrokerConfig = false,
+      expectFailure = false)
+    //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)
+  }
+
+  @Test
+  def testConnectionRateQuota(): Unit = {
+    verifyConfigUpdate(KafkaConfig.MaxConnectionCreationRateProp, "110", perBrokerConfig = true, expectFailure = false)
+    verifyConfigUpdate(KafkaConfig.MaxConnectionCreationRateProp, "120", perBrokerConfig = false, expectFailure = false)
+    val listenerMaxConnectionsProp = s"listener.name.external.${KafkaConfig.MaxConnectionCreationRateProp}"
+    verifyConfigUpdate(listenerMaxConnectionsProp, "20", perBrokerConfig = true, expectFailure = false)
+    verifyConfigUpdate(listenerMaxConnectionsProp, "30", perBrokerConfig = false, expectFailure = false)
+  }
+
+  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)
+    val props = new Properties
+    props.put(name, value)
+    val oldValue = config.originals.get(name)
+
+    def updateConfig() = {
+      if (perBrokerConfig)
+        config.dynamicConfig.updateBrokerConfig(0, config.dynamicConfig.toPersistentProps(props, perBrokerConfig))
+      else
+        config.dynamicConfig.updateDefaultConfig(props)
+    }
+    if (!expectFailure) {
+      config.dynamicConfig.validate(props, perBrokerConfig)
+      updateConfig()
+      assertEquals(value, config.originals.get(name))
+    } else {
+      try {
+        config.dynamicConfig.validate(props, perBrokerConfig)
+        fail("Invalid config did not fail validation")
+      } catch {
+        case e: Exception => // expected exception
+      }
+      updateConfig()
+      assertEquals(oldValue, config.originals.get(name))
+    }
+  }
+
+  private def verifyConfigUpdateWithInvalidConfig(config: KafkaConfig,
+                                                  origProps: Properties,
+                                                  validProps: Map[String, String],
+                                                  invalidProps: Map[String, String]): Unit = {
+    val props = new Properties
+    validProps.foreach { case (k, v) => props.put(k, v) }
+    invalidProps.foreach { case (k, v) => props.put(k, v) }
+
+    // DynamicBrokerConfig#validate is used by AdminClient to validate the configs provided in
+    // in an AlterConfigs request. Validation should fail with an exception if any of the configs are invalid.
+    try {
+      config.dynamicConfig.validate(props, perBrokerConfig = true)
+      fail("Invalid config did not fail validation")
+    } catch {
+      case e: ConfigException => // expected exception
+    }
+
+    // DynamicBrokerConfig#updateBrokerConfig is used to update configs from ZooKeeper during
+    // startup and when configs are updated in ZK. Update should apply valid configs and ignore
+    // invalid ones.
+    config.dynamicConfig.updateBrokerConfig(0, props)
+    validProps.foreach { case (name, value) => assertEquals(value, config.originals.get(name)) }
+    invalidProps.keySet.foreach { name =>
+      assertEquals(origProps.get(name), config.originals.get(name))
+    }
+  }
+
+  @Test
+  def testPasswordConfigEncryption(): Unit = {
+    val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
+    val configWithoutSecret = KafkaConfig(props)
+    props.put(KafkaConfig.PasswordEncoderSecretProp, "config-encoder-secret")
+    val configWithSecret = KafkaConfig(props)
+    val dynamicProps = new Properties
+    dynamicProps.put(KafkaConfig.SaslJaasConfigProp, "myLoginModule required;")
+
+    try {
+      configWithoutSecret.dynamicConfig.toPersistentProps(dynamicProps, perBrokerConfig = true)
+    } catch {
+      case e: ConfigException => // expected exception
+    }
+    val persistedProps = configWithSecret.dynamicConfig.toPersistentProps(dynamicProps, perBrokerConfig = true)
+    assertFalse("Password not encoded",
+      persistedProps.getProperty(KafkaConfig.SaslJaasConfigProp).contains("myLoginModule"))
+    val decodedProps = configWithSecret.dynamicConfig.fromPersistentProps(persistedProps, perBrokerConfig = true)
+    assertEquals("myLoginModule required;", decodedProps.getProperty(KafkaConfig.SaslJaasConfigProp))
+  }
+
+  @Test
+  def testPasswordConfigEncoderSecretChange(): Unit = {
+    val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
+    props.put(KafkaConfig.SaslJaasConfigProp, "staticLoginModule required;")
+    props.put(KafkaConfig.PasswordEncoderSecretProp, "config-encoder-secret")
+    val config = KafkaConfig(props)
+    val dynamicProps = new Properties
+    dynamicProps.put(KafkaConfig.SaslJaasConfigProp, "dynamicLoginModule required;")
+
+    val persistedProps = config.dynamicConfig.toPersistentProps(dynamicProps, perBrokerConfig = true)
+    assertFalse("Password not encoded",
+      persistedProps.getProperty(KafkaConfig.SaslJaasConfigProp).contains("LoginModule"))
+    config.dynamicConfig.updateBrokerConfig(0, persistedProps)
+    assertEquals("dynamicLoginModule required;", config.values.get(KafkaConfig.SaslJaasConfigProp).asInstanceOf[Password].value)
+
+    // New config with same secret should use the dynamic password config
+    val newConfigWithSameSecret = KafkaConfig(props)
+    newConfigWithSameSecret.dynamicConfig.updateBrokerConfig(0, persistedProps)
+    assertEquals("dynamicLoginModule required;", newConfigWithSameSecret.values.get(KafkaConfig.SaslJaasConfigProp).asInstanceOf[Password].value)
+
+    // New config with new secret should use the dynamic password config if new and old secrets are configured in KafkaConfig
+    props.put(KafkaConfig.PasswordEncoderSecretProp, "new-encoder-secret")
+    props.put(KafkaConfig.PasswordEncoderOldSecretProp, "config-encoder-secret")
+    val newConfigWithNewAndOldSecret = KafkaConfig(props)
+    newConfigWithNewAndOldSecret.dynamicConfig.updateBrokerConfig(0, persistedProps)
+    assertEquals("dynamicLoginModule required;", newConfigWithSameSecret.values.get(KafkaConfig.SaslJaasConfigProp).asInstanceOf[Password].value)
+
+    // New config with new secret alone should revert to static password config since dynamic config cannot be decoded
+    props.put(KafkaConfig.PasswordEncoderSecretProp, "another-new-encoder-secret")
+    val newConfigWithNewSecret = KafkaConfig(props)
+    newConfigWithNewSecret.dynamicConfig.updateBrokerConfig(0, persistedProps)
+    assertEquals("staticLoginModule required;", newConfigWithNewSecret.values.get(KafkaConfig.SaslJaasConfigProp).asInstanceOf[Password].value)
+  }
+
+  @Test
+  def testDynamicListenerConfig(): Unit = {
+    val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 9092)
+    val oldConfig =  KafkaConfig.fromProps(props)
+    val kafkaServer: KafkaServer = EasyMock.createMock(classOf[kafka.server.KafkaServer])
+    EasyMock.expect(kafkaServer.config).andReturn(oldConfig).anyTimes()
+    EasyMock.replay(kafkaServer)
+
+    props.put(KafkaConfig.ListenersProp, "PLAINTEXT://hostname:9092,SASL_PLAINTEXT://hostname:9093")
+    new DynamicListenerConfig(kafkaServer).validateReconfiguration(KafkaConfig(props))
+
+    // it is illegal to update non-reconfiguable configs of existent listeners
+    props.put("listener.name.plaintext.you.should.not.pass", "failure")
+    val dynamicListenerConfig = new DynamicListenerConfig(kafkaServer)
+    assertThrows(classOf[ConfigException], () => dynamicListenerConfig.validateReconfiguration(KafkaConfig(props)))
+  }
+
+  @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"),
+      DynamicBrokerConfig.brokerConfigSynonyms("listener.name.secure.ssl.keystore.type", matchListenerOverride = true))
+    assertEquals(List("listener.name.sasl_ssl.plain.sasl.jaas.config", "sasl.jaas.config"),
+      DynamicBrokerConfig.brokerConfigSynonyms("listener.name.sasl_ssl.plain.sasl.jaas.config", matchListenerOverride = true))
+    assertEquals(List("some.config"),
+      DynamicBrokerConfig.brokerConfigSynonyms("some.config", matchListenerOverride = true))
+    assertEquals(List(KafkaConfig.LogRollTimeMillisProp, KafkaConfig.LogRollTimeHoursProp),
+      DynamicBrokerConfig.brokerConfigSynonyms(KafkaConfig.LogRollTimeMillisProp, matchListenerOverride = true))
+  }
+
+  @Test
+  def testDynamicConfigInitializationWithoutConfigsInZK(): Unit = {
+    val zkClient: KafkaZkClient = EasyMock.createMock(classOf[KafkaZkClient])
+    EasyMock.expect(zkClient.getEntityConfigs(EasyMock.anyString(), EasyMock.anyString())).andReturn(new java.util.Properties()).anyTimes()
+    EasyMock.replay(zkClient)
+
+    val oldConfig =  KafkaConfig.fromProps(TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 9092))
+    val dynamicBrokerConfig = new DynamicBrokerConfig(oldConfig)
+    dynamicBrokerConfig.initialize(zkClient)
+    dynamicBrokerConfig.addBrokerReconfigurable(new TestDynamicThreadPool)
+
+    val newprops = new Properties()
+    newprops.put(KafkaConfig.NumIoThreadsProp, "10")
+    newprops.put(KafkaConfig.BackgroundThreadsProp, "100")
+    dynamicBrokerConfig.updateBrokerConfig(0, newprops)
+  }
+
+  @Test
+  def testImproperConfigsAreRemoved(): Unit = {
+    val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect)
+    val configs = KafkaConfig(props)
+
+    assertEquals(Defaults.MaxConnections, configs.maxConnections)
+    assertEquals(Defaults.MessageMaxBytes, configs.messageMaxBytes)
+
+    var newProps = new Properties()
+    newProps.put(KafkaConfig.MaxConnectionsProp, "9999")
+    newProps.put(KafkaConfig.MessageMaxBytesProp, "2222")
+
+    configs.dynamicConfig.updateDefaultConfig(newProps)
+    assertEquals(9999, configs.maxConnections)
+    assertEquals(2222, configs.messageMaxBytes)
+
+    newProps = new Properties()
+    newProps.put(KafkaConfig.MaxConnectionsProp, "INVALID_INT")
+    newProps.put(KafkaConfig.MessageMaxBytesProp, "1111")
+
+    configs.dynamicConfig.updateDefaultConfig(newProps)
+    // Invalid value should be skipped and reassigned as default value
+    assertEquals(Defaults.MaxConnections, configs.maxConnections)
+    // Even if One property is invalid, the below should get correctly updated.
+    assertEquals(1111, configs.messageMaxBytes)
+  }
+}
+
+class TestDynamicThreadPool() extends BrokerReconfigurable {
+
+  override def reconfigurableConfigs: Set[String] = {
+    DynamicThreadPool.ReconfigurableConfigs
+  }
+
+  override def reconfigure(oldConfig: KafkaConfig, newConfig: KafkaConfig): Unit = {
+    assertEquals(Defaults.NumIoThreads, oldConfig.numIoThreads)
+    assertEquals(Defaults.BackgroundThreads, oldConfig.backgroundThreads)
+
+    assertEquals(10, newConfig.numIoThreads)
+    assertEquals(100, newConfig.backgroundThreads)
+  }
+
+  override def validateReconfiguration(newConfig: KafkaConfig): Unit = {
+    assertEquals(10, newConfig.numIoThreads)
+    assertEquals(100, newConfig.backgroundThreads)
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/server/DynamicConfigChangeTest.scala b/core/src/test/scala/unit/kafka/server/DynamicConfigChangeTest.scala
index 85bd6a11076ff..009db4c1f7fa6 100644
--- a/core/src/test/scala/unit/kafka/server/DynamicConfigChangeTest.scala
+++ b/core/src/test/scala/unit/kafka/server/DynamicConfigChangeTest.scala
@@ -16,26 +16,35 @@
   */
 package kafka.server
 
+import java.net.InetAddress
+import java.nio.charset.StandardCharsets
 import java.util.Properties
+import java.util.concurrent.ExecutionException
 
+import kafka.integration.KafkaServerTestHarness
 import kafka.log.LogConfig._
+import kafka.utils._
 import kafka.server.Constants._
-import org.junit.Assert._
+import kafka.zk.ConfigEntityChangeNotificationZNode
+import org.apache.kafka.clients.CommonClientConfigs
+import org.apache.kafka.clients.admin.{Admin, AlterConfigOp, ConfigEntry}
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.config.ConfigResource
+import org.apache.kafka.common.errors.UnknownTopicOrPartitionException
 import org.apache.kafka.common.metrics.Quota
 import org.easymock.EasyMock
+import org.junit.Assert._
 import org.junit.Test
-import kafka.integration.KafkaServerTestHarness
-import kafka.utils._
-import kafka.admin.AdminOperationException
-import org.apache.kafka.common.TopicPartition
+import org.scalatest.Assertions.assertThrows
 
-import scala.collection.Map
+import scala.collection.{Map, Seq}
+import scala.jdk.CollectionConverters._
 
 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
@@ -43,7 +52,7 @@ class DynamicConfigChangeTest extends KafkaServerTestHarness {
     val tp = new TopicPartition("test", 0)
     val logProps = new Properties()
     logProps.put(FlushMessagesProp, oldVal.toString)
-    adminZkClient.createTopic(tp.topic, 1, 1, logProps)
+    createTopic(tp.topic, 1, 1, logProps)
     TestUtils.retry(10000) {
       val logOpt = this.servers.head.logManager.getLog(tp)
       assertTrue(logOpt.isDefined)
@@ -56,14 +65,41 @@ class DynamicConfigChangeTest extends KafkaServerTestHarness {
     }
   }
 
-  private def testQuotaConfigChange(user: String, clientId: String, rootEntityType: String, configEntityName: String) {
+  @Test
+  def testDynamicTopicConfigChange(): Unit = {
+    val tp = new TopicPartition("test", 0)
+    val oldSegmentSize = 1000
+    val logProps = new Properties()
+    logProps.put(SegmentBytesProp, oldSegmentSize.toString)
+    createTopic(tp.topic, 1, 1, logProps)
+    TestUtils.retry(10000) {
+      val logOpt = this.servers.head.logManager.getLog(tp)
+      assertTrue(logOpt.isDefined)
+      assertEquals(oldSegmentSize, logOpt.get.config.segmentSize)
+    }
+
+    val log = servers.head.logManager.getLog(tp).get
+
+    val newSegmentSize = 2000
+    logProps.put(SegmentBytesProp, newSegmentSize.toString)
+    adminZkClient.changeTopicConfig(tp.topic, logProps)
+    TestUtils.retry(10000) {
+      assertEquals(newSegmentSize, log.config.segmentSize)
+    }
+
+    (1 to 50).foreach(i => TestUtils.produceMessage(servers, tp.topic, i.toString))
+    // Verify that the new config is used for all segments
+    assertTrue("Log segment size change not applied", log.logSegments.forall(_.size > 1000))
+  }
+
+  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()
     props.put(DynamicConfig.Client.ProducerByteRateOverrideProp, "1000")
     props.put(DynamicConfig.Client.ConsumerByteRateOverrideProp, "2000")
 
-    val quotaManagers = servers.head.apis.quotas
+    val quotaManagers = servers.head.dataPlaneRequestProcessor.quotas
     rootEntityType match {
       case ConfigType.Client => adminZkClient.changeClientIdConfig(configEntityName, props)
       case _ => adminZkClient.changeUserOrUserClientIdConfig(configEntityName, props)
@@ -99,37 +135,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()
@@ -147,9 +183,9 @@ class DynamicConfigChangeTest extends KafkaServerTestHarness {
     adminZkClient.changeUserOrUserClientIdConfig("ANONYMOUS/clients/overriddenUserClientId", userClientIdProps)
 
     // Remove config change znodes to force quota initialization only through loading of user/client quotas
-    zkUtils.getChildren(ZkUtils.ConfigChangesPath).foreach { p => zkUtils.deletePath(ZkUtils.ConfigChangesPath + "/" + p) }
+    zkClient.getChildren(ConfigEntityChangeNotificationZNode.path).foreach { p => zkClient.deletePath(ConfigEntityChangeNotificationZNode.path + "/" + p) }
     server.startup()
-    val quotaManagers = server.apis.quotas
+    val quotaManagers = server.dataPlaneRequestProcessor.quotas
 
     assertEquals(Quota.upperBound(1000),  quotaManagers.produce.quota("someuser", "overriddenClientId"))
     assertEquals(Quota.upperBound(2000),  quotaManagers.fetch.quota("someuser", "overriddenClientId"))
@@ -160,15 +196,98 @@ class DynamicConfigChangeTest extends KafkaServerTestHarness {
   }
 
   @Test
-  def testConfigChangeOnNonExistingTopic() {
-    val topic = TestUtils.tempTopic
+  def testIpHandlerUnresolvableAddress(): Unit = {
+    val configHandler = new IpConfigHandler(null)
+    val props: Properties = new Properties()
+    props.put(DynamicConfig.Ip.IpConnectionRateOverrideProp, "1")
+
+    assertThrows[IllegalArgumentException]{
+      configHandler.processConfigChanges("illegal-hostname", props)
+    }
+  }
+
+  @Test
+  def testIpQuotaInitialization(): Unit = {
+    val server = servers.head
+    val ipOverrideProps = new Properties()
+    ipOverrideProps.put(DynamicConfig.Ip.IpConnectionRateOverrideProp, "10")
+    val ipDefaultProps = new Properties()
+    ipDefaultProps.put(DynamicConfig.Ip.IpConnectionRateOverrideProp, "20")
+    server.shutdown()
+
+    adminZkClient.changeIpConfig(ConfigEntityName.Default, ipDefaultProps)
+    adminZkClient.changeIpConfig("1.2.3.4", ipOverrideProps)
+
+    // Remove config change znodes to force quota initialization only through loading of ip quotas
+    zkClient.getChildren(ConfigEntityChangeNotificationZNode.path).foreach { p =>
+      zkClient.deletePath(ConfigEntityChangeNotificationZNode.path + "/" + p)
+    }
+    server.startup()
+
+    val connectionQuotas = server.socketServer.connectionQuotas
+    assertEquals(10L, connectionQuotas.connectionRateForIp(InetAddress.getByName("1.2.3.4")))
+    assertEquals(20L, connectionQuotas.connectionRateForIp(InetAddress.getByName("2.4.6.8")))
+  }
+
+  @Test
+  def testIpQuotaConfigChange(): Unit = {
+    val ipOverrideProps = new Properties()
+    ipOverrideProps.put(DynamicConfig.Ip.IpConnectionRateOverrideProp, "10")
+    val ipDefaultProps = new Properties()
+    ipDefaultProps.put(DynamicConfig.Ip.IpConnectionRateOverrideProp, "20")
+
+    val overrideQuotaIp = InetAddress.getByName("1.2.3.4")
+    val defaultQuotaIp = InetAddress.getByName("2.3.4.5")
+    adminZkClient.changeIpConfig(ConfigEntityName.Default, ipDefaultProps)
+    adminZkClient.changeIpConfig(overrideQuotaIp.getHostAddress, ipOverrideProps)
+
+    val connectionQuotas = servers.head.socketServer.connectionQuotas
+
+    def verifyConnectionQuota(ip: InetAddress, expectedQuota: Integer) = {
+      TestUtils.retry(10000) {
+        val quota = connectionQuotas.connectionRateForIp(ip)
+        assertEquals(s"Unexpected quota for IP $ip", expectedQuota, quota)
+      }
+    }
+
+    verifyConnectionQuota(overrideQuotaIp, 10)
+    verifyConnectionQuota(defaultQuotaIp, 20)
+
+    val emptyProps = new Properties()
+    adminZkClient.changeIpConfig(overrideQuotaIp.getHostAddress, emptyProps)
+    verifyConnectionQuota(overrideQuotaIp, 20)
+
+    adminZkClient.changeIpConfig(ConfigEntityName.Default, emptyProps)
+    verifyConnectionQuota(overrideQuotaIp, DynamicConfig.Ip.DefaultConnectionCreationRate)
+  }
+
+  @Test
+  def testConfigChangeOnNonExistingTopic(): Unit = {
+    val topic = TestUtils.tempTopic()
     try {
       val logProps = new Properties()
       logProps.put(FlushMessagesProp, 10000: java.lang.Integer)
       adminZkClient.changeTopicConfig(topic, logProps)
-      fail("Should fail with AdminOperationException for topic doesn't exist")
+      fail("Should fail with UnknownTopicOrPartitionException for topic doesn't exist")
+    } catch {
+      case _: UnknownTopicOrPartitionException => // expected
+    }
+  }
+
+  @Test
+  def testConfigChangeOnNonExistingTopicWithAdminClient(): Unit = {
+    val topic = TestUtils.tempTopic()
+    val admin = createAdminClient()
+    try {
+      val resource = new ConfigResource(ConfigResource.Type.TOPIC, topic)
+      val op = new AlterConfigOp(new ConfigEntry(FlushMessagesProp, "10000"), AlterConfigOp.OpType.SET)
+      admin.incrementalAlterConfigs(Map(resource -> List(op).asJavaCollection).asJava).all.get
+      fail("Should fail with UnknownTopicOrPartitionException for topic doesn't exist")
     } catch {
-      case _: AdminOperationException => // expected
+      case e: ExecutionException =>
+        assertTrue(e.getCause.isInstanceOf[UnknownTopicOrPartitionException])
+    } finally {
+      admin.close()
     }
   }
 
@@ -180,7 +299,7 @@ class DynamicConfigChangeTest extends KafkaServerTestHarness {
     // Create a mock ConfigHandler to record config changes it is asked to process
     val entityArgument = EasyMock.newCapture[String]
     val propertiesArgument = EasyMock.newCapture[Properties]
-    val handler = EasyMock.createNiceMock(classOf[ConfigHandler])
+    val handler: ConfigHandler = EasyMock.createNiceMock(classOf[ConfigHandler])
     handler.processConfigChanges(
       EasyMock.and(EasyMock.capture(entityArgument), EasyMock.isA(classOf[String])),
       EasyMock.and(EasyMock.capture(propertiesArgument), EasyMock.isA(classOf[Properties])))
@@ -189,12 +308,12 @@ class DynamicConfigChangeTest extends KafkaServerTestHarness {
 
     val configManager = new DynamicConfigManager(zkClient, Map(ConfigType.Topic -> handler))
     // Notifications created using the old TopicConfigManager are ignored.
-    configManager.ConfigChangedNotificationHandler.processNotification("not json")
+    configManager.ConfigChangedNotificationHandler.processNotification("not json".getBytes(StandardCharsets.UTF_8))
 
     // Incorrect Map. No version
     try {
       val jsonMap = Map("v" -> 1, "x" -> 2)
-      configManager.ConfigChangedNotificationHandler.processNotification(Json.encode(jsonMap))
+      configManager.ConfigChangedNotificationHandler.processNotification(Json.encodeAsBytes(jsonMap.asJava))
       fail("Should have thrown an Exception while parsing incorrect notification " + jsonMap)
     }
     catch {
@@ -203,7 +322,7 @@ class DynamicConfigChangeTest extends KafkaServerTestHarness {
     // Version is provided. EntityType is incorrect
     try {
       val jsonMap = Map("version" -> 1, "entity_type" -> "garbage", "entity_name" -> "x")
-      configManager.ConfigChangedNotificationHandler.processNotification(Json.encode(jsonMap))
+      configManager.ConfigChangedNotificationHandler.processNotification(Json.encodeAsBytes(jsonMap.asJava))
       fail("Should have thrown an Exception while parsing incorrect notification " + jsonMap)
     }
     catch {
@@ -213,7 +332,7 @@ class DynamicConfigChangeTest extends KafkaServerTestHarness {
     // EntityName isn't provided
     try {
       val jsonMap = Map("version" -> 1, "entity_type" -> ConfigType.Topic)
-      configManager.ConfigChangedNotificationHandler.processNotification(Json.encode(jsonMap))
+      configManager.ConfigChangedNotificationHandler.processNotification(Json.encodeAsBytes(jsonMap.asJava))
       fail("Should have thrown an Exception while parsing incorrect notification " + jsonMap)
     }
     catch {
@@ -222,7 +341,7 @@ class DynamicConfigChangeTest extends KafkaServerTestHarness {
 
     // Everything is provided
     val jsonMap = Map("version" -> 1, "entity_type" -> ConfigType.Topic, "entity_name" -> "x")
-    configManager.ConfigChangedNotificationHandler.processNotification(Json.encode(jsonMap))
+    configManager.ConfigChangedNotificationHandler.processNotification(Json.encodeAsBytes(jsonMap.asJava))
 
     // Verify that processConfigChanges was only called once
     EasyMock.verify(handler)
@@ -230,7 +349,7 @@ class DynamicConfigChangeTest extends KafkaServerTestHarness {
 
   @Test
   def shouldParseReplicationQuotaProperties(): Unit = {
-    val configHandler: TopicConfigHandler = new TopicConfigHandler(null, null, null)
+    val configHandler: TopicConfigHandler = new TopicConfigHandler(null, null, null, null)
     val props: Properties = new Properties()
 
     //Given
@@ -243,7 +362,7 @@ class DynamicConfigChangeTest extends KafkaServerTestHarness {
 
   @Test
   def shouldParseWildcardReplicationQuotaProperties(): Unit = {
-    val configHandler: TopicConfigHandler = new TopicConfigHandler(null, null, null)
+    val configHandler: TopicConfigHandler = new TopicConfigHandler(null, null, null, null)
     val props: Properties = new Properties()
 
     //Given
@@ -258,7 +377,7 @@ class DynamicConfigChangeTest extends KafkaServerTestHarness {
 
   @Test
   def shouldParseReplicationQuotaReset(): Unit = {
-    val configHandler: TopicConfigHandler = new TopicConfigHandler(null, null, null)
+    val configHandler: TopicConfigHandler = new TopicConfigHandler(null, null, null, null)
     val props: Properties = new Properties()
 
     //Given
@@ -272,8 +391,8 @@ class DynamicConfigChangeTest extends KafkaServerTestHarness {
   }
 
   @Test
-  def shouldParseRegardlessOfWhitespaceAroundValues() {
-    val configHandler: TopicConfigHandler = new TopicConfigHandler(null, null, null)
+  def shouldParseRegardlessOfWhitespaceAroundValues(): Unit = {
+    val configHandler: TopicConfigHandler = new TopicConfigHandler(null, null, null, null)
     assertEquals(AllReplicas, parse(configHandler, "* "))
     assertEquals(Seq(), parse(configHandler, " "))
     assertEquals(Seq(6), parse(configHandler, "6:102"))
@@ -284,4 +403,11 @@ class DynamicConfigChangeTest extends KafkaServerTestHarness {
   def parse(configHandler: TopicConfigHandler, value: String): Seq[Int] = {
     configHandler.parseThrottledPartitions(CoreUtils.propsWith(LeaderReplicationThrottledReplicasProp, value), 102, LeaderReplicationThrottledReplicasProp)
   }
+
+  private def createAdminClient(): Admin = {
+    val props = new Properties()
+    props.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, brokerList)
+    Admin.create(props)
+  }
+
 }
diff --git a/core/src/test/scala/unit/kafka/server/DynamicConfigTest.scala b/core/src/test/scala/unit/kafka/server/DynamicConfigTest.scala
index b2378cfcbbf8c..ce65df9b72566 100644
--- a/core/src/test/scala/unit/kafka/server/DynamicConfigTest.scala
+++ b/core/src/test/scala/unit/kafka/server/DynamicConfigTest.scala
@@ -16,39 +16,50 @@
   */
 package kafka.server
 
+import kafka.admin.AdminOperationException
 import kafka.utils.CoreUtils._
 import kafka.zk.ZooKeeperTestHarness
 import org.apache.kafka.common.config._
 import org.junit.Test
 
-class DynamicConfigTest  extends ZooKeeperTestHarness {
+class DynamicConfigTest extends ZooKeeperTestHarness {
   private final val nonExistentConfig: String = "some.config.that.does.not.exist"
   private final val someValue: String = "some interesting value"
 
   @Test(expected = classOf[IllegalArgumentException])
-  def shouldFailWhenChangingBrokerUnknownConfig() {
-    adminZkClient.changeBrokerConfig(Seq(0), propsWith(nonExistentConfig, someValue))
-  }
-
-  @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"))
   }
-}
\ No newline at end of file
+
+  @Test(expected = classOf[ConfigException])
+  def shouldFailIpConfigsWithInvalidValues(): Unit = {
+    adminZkClient.changeIpConfig("1.2.3.4", propsWith(DynamicConfig.Ip.IpConnectionRateOverrideProp, "-1"))
+  }
+
+  @Test(expected = classOf[AdminOperationException])
+  def shouldFailIpConfigsWithInvalidIpv4Entity(): Unit = {
+    adminZkClient.changeIpConfig("1,1.1.1", propsWith(DynamicConfig.Ip.IpConnectionRateOverrideProp, "2"));
+  }
+
+  @Test(expected = classOf[AdminOperationException])
+  def shouldFailIpConfigsWithBadHost(): Unit = {
+    adminZkClient.changeIpConfig("ip", propsWith(DynamicConfig.Ip.IpConnectionRateOverrideProp, "2"));
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/server/EdgeCaseRequestTest.scala b/core/src/test/scala/unit/kafka/server/EdgeCaseRequestTest.scala
index 7b3011f7a758f..58d86b645cb17 100755
--- a/core/src/test/scala/unit/kafka/server/EdgeCaseRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/EdgeCaseRequestTest.scala
@@ -20,21 +20,24 @@ package kafka.server
 import java.io.{DataInputStream, DataOutputStream}
 import java.net.Socket
 import java.nio.ByteBuffer
+import java.util.Collections
 
 import kafka.integration.KafkaServerTestHarness
 import kafka.network.SocketServer
 import kafka.utils._
-import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.message.ProduceRequestData
 import org.apache.kafka.common.network.ListenerName
 import org.apache.kafka.common.protocol.types.Type
 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, ResponseHeader}
+import org.apache.kafka.common.requests.{ProduceResponse, ResponseHeader}
 import org.apache.kafka.common.security.auth.SecurityProtocol
+import org.apache.kafka.common.utils.ByteUtils
+import org.apache.kafka.common.{TopicPartition, requests}
 import org.junit.Assert._
 import org.junit.Test
 
-import scala.collection.JavaConverters._
+import scala.annotation.nowarn
 
 class EdgeCaseRequestTest extends KafkaServerTestHarness {
 
@@ -50,7 +53,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) =>
@@ -82,12 +85,16 @@ class EdgeCaseRequestTest extends KafkaServerTestHarness {
   }
 
   // Custom header serialization so that protocol assumptions are not forced
-  private def requestHeaderBytes(apiKey: Short, apiVersion: Short, clientId: String = "", correlationId: Int = -1): Array[Byte] = {
+  def requestHeaderBytes(apiKey: Short, apiVersion: Short, clientId: String = "", correlationId: Int = -1): Array[Byte] = {
+    // Check for flex versions, some tests here verify that an invalid apiKey is detected properly, so if -1 is used,
+    // assume the request is not using flex versions.
+    val flexVersion = if (apiKey >= 0) ApiKeys.forId(apiKey).requestHeaderVersion(apiVersion) >= 2 else false
     val size = {
       2 /* apiKey */ +
         2 /* version id */ +
         4 /* correlation id */ +
-        Type.NULLABLE_STRING.sizeOf(clientId) /* client id */
+        Type.NULLABLE_STRING.sizeOf(clientId)  /* client id */ +
+        (if (flexVersion) ByteUtils.sizeOfUnsignedVarint(0) else 0) /* Empty tagged fields for flexible versions */
     }
 
     val buffer = ByteBuffer.allocate(size)
@@ -95,10 +102,11 @@ class EdgeCaseRequestTest extends KafkaServerTestHarness {
     buffer.putShort(apiVersion)
     buffer.putInt(correlationId)
     Type.NULLABLE_STRING.write(buffer, clientId)
+    if (flexVersion) ByteUtils.writeUnsignedVarint(0, buffer)
     buffer.array()
   }
 
-  private def verifyDisconnect(request: Array[Byte]) {
+  private def verifyDisconnect(request: Array[Byte]): Unit = {
     val plainSocket = connect()
     try {
       sendRequest(plainSocket, requestHeaderBytes(-1, 0))
@@ -108,29 +116,40 @@ class EdgeCaseRequestTest extends KafkaServerTestHarness {
     }
   }
 
+  @nowarn("cat=deprecation")
   @Test
-  def testProduceRequestWithNullClientId() {
+  def testProduceRequestWithNullClientId(): Unit = {
     val topic = "topic"
     val topicPartition = new TopicPartition(topic, 0)
     val correlationId = -1
-    TestUtils.createTopic(zkUtils, topic, numPartitions = 1, replicationFactor = 1, servers = servers)
+    createTopic(topic, numPartitions = 1, replicationFactor = 1)
 
     val version = ApiKeys.PRODUCE.latestVersion: Short
-    val serializedBytes = {
-      val headerBytes = requestHeaderBytes(ApiKeys.PRODUCE.id, version, null,
-        correlationId)
-      val records = MemoryRecords.withRecords(CompressionType.NONE, new SimpleRecord("message".getBytes))
-      val request = ProduceRequest.Builder.forCurrentMagic(1, 10000, Map(topicPartition -> records).asJava).build()
-      val byteBuffer = ByteBuffer.allocate(headerBytes.length + request.toStruct.sizeOf)
+    val (serializedBytes, responseHeaderVersion) = {
+      val headerBytes = requestHeaderBytes(ApiKeys.PRODUCE.id, version, "", correlationId)
+      val request = requests.ProduceRequest.forCurrentMagic(new ProduceRequestData()
+        .setTopicData(new ProduceRequestData.TopicProduceDataCollection(
+          Collections.singletonList(new ProduceRequestData.TopicProduceData()
+            .setName(topicPartition.topic()).setPartitionData(Collections.singletonList(
+            new ProduceRequestData.PartitionProduceData()
+              .setIndex(topicPartition.partition())
+              .setRecords(MemoryRecords.withRecords(CompressionType.NONE, new SimpleRecord("message".getBytes))))))
+            .iterator))
+        .setAcks(1.toShort)
+        .setTimeoutMs(10000)
+        .setTransactionalId(null))
+        .build()
+      val bodyBytes = request.serialize
+      val byteBuffer = ByteBuffer.allocate(headerBytes.length + bodyBytes.remaining())
       byteBuffer.put(headerBytes)
-      request.toStruct.writeTo(byteBuffer)
-      byteBuffer.array()
+      byteBuffer.put(bodyBytes)
+      (byteBuffer.array(), request.apiKey.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 +162,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
new file mode 100644
index 0000000000000..a813e19ad01ca
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/FetchRequestDownConversionConfigTest.scala
@@ -0,0 +1,164 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package kafka.server
+
+import java.util
+import java.util.{Optional, Properties}
+
+import kafka.log.LogConfig
+import kafka.utils.TestUtils
+import org.apache.kafka.clients.producer.{KafkaProducer, ProducerRecord}
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.record.MemoryRecords
+import org.apache.kafka.common.requests.{FetchRequest, FetchResponse}
+import org.apache.kafka.common.serialization.StringSerializer
+import org.junit.Assert._
+import org.junit.Test
+
+class FetchRequestDownConversionConfigTest extends BaseRequestTest {
+  private var producer: KafkaProducer[String, String] = null
+  override def brokerCount: Int = 1
+
+  override def setUp(): Unit = {
+    super.setUp()
+    initProducer()
+  }
+
+  override def tearDown(): Unit = {
+    if (producer != null)
+      producer.close()
+    super.tearDown()
+  }
+
+  override protected def brokerPropertyOverrides(properties: Properties): Unit = {
+    super.brokerPropertyOverrides(properties)
+    properties.put(KafkaConfig.LogMessageDownConversionEnableProp, "false")
+  }
+
+  private def initProducer(): Unit = {
+    producer = TestUtils.createProducer(TestUtils.getBrokerListStrFromServers(servers),
+      keySerializer = new StringSerializer, valueSerializer = new StringSerializer)
+  }
+
+  private def createTopics(numTopics: Int, numPartitions: Int,
+                           configs: Map[String, String] = Map.empty, topicSuffixStart: Int = 0): Map[TopicPartition, Int] = {
+    val topics = (0 until numTopics).map(t => s"topic${t + topicSuffixStart}")
+    val topicConfig = new Properties
+    topicConfig.setProperty(LogConfig.MinInSyncReplicasProp, 1.toString)
+    configs.foreach { case (k, v) => topicConfig.setProperty(k, v) }
+    topics.flatMap { topic =>
+      val partitionToLeader = createTopic(topic, numPartitions = numPartitions, replicationFactor = 1,
+        topicConfig = topicConfig)
+      partitionToLeader.map { case (partition, leader) => new TopicPartition(topic, partition) -> leader }
+    }.toMap
+  }
+
+  private def createPartitionMap(maxPartitionBytes: Int, topicPartitions: Seq[TopicPartition],
+                                 offsetMap: Map[TopicPartition, Long] = Map.empty): util.LinkedHashMap[TopicPartition, FetchRequest.PartitionData] = {
+    val partitionMap = new util.LinkedHashMap[TopicPartition, FetchRequest.PartitionData]
+    topicPartitions.foreach { tp =>
+      partitionMap.put(tp, new FetchRequest.PartitionData(offsetMap.getOrElse(tp, 0), 0L,
+        maxPartitionBytes, Optional.empty()))
+    }
+    partitionMap
+  }
+
+  private def sendFetchRequest(leaderId: Int, request: FetchRequest): FetchResponse[MemoryRecords] = {
+    connectAndReceive[FetchResponse[MemoryRecords]](request, destination = brokerSocketServer(leaderId))
+  }
+
+  /**
+   * Tests that fetch request that require down-conversion returns with an error response when down-conversion is disabled on broker.
+   */
+  @Test
+  def testV1FetchWithDownConversionDisabled(): Unit = {
+    val topicMap = createTopics(numTopics = 5, numPartitions = 1)
+    val topicPartitions = topicMap.keySet.toSeq
+    topicPartitions.foreach(tp => producer.send(new ProducerRecord(tp.topic(), "key", "value")).get())
+    val fetchRequest = FetchRequest.Builder.forConsumer(Int.MaxValue, 0, createPartitionMap(1024,
+      topicPartitions)).build(1)
+    val fetchResponse = sendFetchRequest(topicMap.head._2, fetchRequest)
+    topicPartitions.foreach(tp => assertEquals(Errors.UNSUPPORTED_VERSION, fetchResponse.responseData().get(tp).error))
+  }
+
+  /**
+   * Tests that "message.downconversion.enable" has no effect when down-conversion is not required.
+   */
+  @Test
+  def testLatestFetchWithDownConversionDisabled(): Unit = {
+    val topicMap = createTopics(numTopics = 5, numPartitions = 1)
+    val topicPartitions = topicMap.keySet.toSeq
+    topicPartitions.foreach(tp => producer.send(new ProducerRecord(tp.topic(), "key", "value")).get())
+    val fetchRequest = FetchRequest.Builder.forConsumer(Int.MaxValue, 0, createPartitionMap(1024,
+      topicPartitions)).build()
+    val fetchResponse = sendFetchRequest(topicMap.head._2, fetchRequest)
+    topicPartitions.foreach(tp => assertEquals(Errors.NONE, fetchResponse.responseData().get(tp).error))
+  }
+
+  /**
+   * Tests that "message.downconversion.enable" can be set at topic level, and its configuration is obeyed for client
+   * fetch requests.
+   */
+  @Test
+  def testV1FetchWithTopicLevelOverrides(): Unit = {
+    // create topics with default down-conversion configuration (i.e. conversion disabled)
+    val conversionDisabledTopicsMap = createTopics(numTopics = 5, numPartitions = 1, topicSuffixStart = 0)
+    val conversionDisabledTopicPartitions = conversionDisabledTopicsMap.keySet.toSeq
+
+    // create topics with down-conversion configuration enabled
+    val topicConfig = Map(LogConfig.MessageDownConversionEnableProp -> "true")
+    val conversionEnabledTopicsMap = createTopics(numTopics = 5, numPartitions = 1, topicConfig, topicSuffixStart = 5)
+    val conversionEnabledTopicPartitions = conversionEnabledTopicsMap.keySet.toSeq
+
+    val allTopics = conversionDisabledTopicPartitions ++ conversionEnabledTopicPartitions
+    val leaderId = conversionDisabledTopicsMap.head._2
+
+    allTopics.foreach(tp => producer.send(new ProducerRecord(tp.topic(), "key", "value")).get())
+    val fetchRequest = FetchRequest.Builder.forConsumer(Int.MaxValue, 0, createPartitionMap(1024,
+      allTopics)).build(1)
+    val fetchResponse = sendFetchRequest(leaderId, fetchRequest)
+
+    conversionDisabledTopicPartitions.foreach(tp => assertEquals(Errors.UNSUPPORTED_VERSION, fetchResponse.responseData().get(tp).error))
+    conversionEnabledTopicPartitions.foreach(tp => assertEquals(Errors.NONE, fetchResponse.responseData().get(tp).error))
+  }
+
+  /**
+   * Tests that "message.downconversion.enable" has no effect on fetch requests from replicas.
+   */
+  @Test
+  def testV1FetchFromReplica(): Unit = {
+    // create topics with default down-conversion configuration (i.e. conversion disabled)
+    val conversionDisabledTopicsMap = createTopics(numTopics = 5, numPartitions = 1, topicSuffixStart = 0)
+    val conversionDisabledTopicPartitions = conversionDisabledTopicsMap.keySet.toSeq
+
+    // create topics with down-conversion configuration enabled
+    val topicConfig = Map(LogConfig.MessageDownConversionEnableProp -> "true")
+    val conversionEnabledTopicsMap = createTopics(numTopics = 5, numPartitions = 1, topicConfig, topicSuffixStart = 5)
+    val conversionEnabledTopicPartitions = conversionEnabledTopicsMap.keySet.toSeq
+
+    val allTopicPartitions = conversionDisabledTopicPartitions ++ conversionEnabledTopicPartitions
+    val leaderId = conversionDisabledTopicsMap.head._2
+
+    allTopicPartitions.foreach(tp => producer.send(new ProducerRecord(tp.topic(), "key", "value")).get())
+    val fetchRequest = FetchRequest.Builder.forReplica(1, 1, Int.MaxValue, 0,
+      createPartitionMap(1024, allTopicPartitions)).build()
+    val fetchResponse = sendFetchRequest(leaderId, fetchRequest)
+
+    allTopicPartitions.foreach(tp => assertEquals(Errors.NONE, fetchResponse.responseData().get(tp).error))
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/server/FetchRequestMaxBytesTest.scala b/core/src/test/scala/unit/kafka/server/FetchRequestMaxBytesTest.scala
new file mode 100644
index 0000000000000..8203dfcfa559a
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/FetchRequestMaxBytesTest.scala
@@ -0,0 +1,132 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package kafka.server
+
+import java.util.{Optional, Properties}
+
+import kafka.log.LogConfig
+import kafka.utils.TestUtils
+import org.apache.kafka.clients.producer.{KafkaProducer, ProducerRecord}
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.record.MemoryRecords
+import org.apache.kafka.common.requests.FetchRequest.PartitionData
+import org.apache.kafka.common.requests.{FetchRequest, FetchResponse}
+import org.junit.{Assert, Test}
+
+import scala.jdk.CollectionConverters._
+
+/**
+ * This test verifies that the KIP-541 broker-level FetchMaxBytes configuration is honored.
+ */
+class FetchRequestMaxBytesTest extends BaseRequestTest {
+  override def brokerCount: Int = 1
+
+  private var producer: KafkaProducer[Array[Byte], Array[Byte]] = null
+  private val testTopic = "testTopic"
+  private val testTopicPartition = new TopicPartition(testTopic, 0)
+  private val messages = IndexedSeq(
+    multiByteArray(1),
+    multiByteArray(500),
+    multiByteArray(1040),
+    multiByteArray(500),
+    multiByteArray(50))
+
+  private def multiByteArray(length: Int): Array[Byte] = {
+    val array = new Array[Byte](length)
+    array.indices.foreach(i => array(i) = (i % 5).toByte)
+    array
+  }
+
+  private def oneByteArray(value: Byte): Array[Byte] = {
+    val array = new Array[Byte](1)
+    array(0) = value
+    array
+  }
+
+  override def setUp(): Unit = {
+    super.setUp()
+    producer = TestUtils.createProducer(TestUtils.getBrokerListStrFromServers(servers))
+  }
+
+  override def tearDown(): Unit = {
+    if (producer != null)
+      producer.close()
+    super.tearDown()
+  }
+
+  override protected def brokerPropertyOverrides(properties: Properties): Unit = {
+    super.brokerPropertyOverrides(properties)
+    properties.put(KafkaConfig.FetchMaxBytes, "1024")
+  }
+
+  private def createTopics(): Unit = {
+    val topicConfig = new Properties
+    topicConfig.setProperty(LogConfig.MinInSyncReplicasProp, 1.toString)
+    createTopic(testTopic,
+                numPartitions = 1, 
+                replicationFactor = 1,
+                topicConfig = topicConfig)
+    // Produce several messages as single batches.
+    messages.indices.foreach(i => {
+      val record = new ProducerRecord(testTopic, 0, oneByteArray(i.toByte), messages(i))
+      val future = producer.send(record)
+      producer.flush()
+      future.get()
+    })
+  }
+
+  private def sendFetchRequest(leaderId: Int, request: FetchRequest): FetchResponse[MemoryRecords] = {
+    connectAndReceive[FetchResponse[MemoryRecords]](request, destination = brokerSocketServer(leaderId))
+  }
+
+  /**
+   * Tests that each of our fetch requests respects FetchMaxBytes.
+   *
+   * Note that when a single batch is larger than FetchMaxBytes, it will be
+   * returned in full even if this is larger than FetchMaxBytes.  See KIP-74.
+   */
+  @Test
+  def testConsumeMultipleRecords(): Unit = {
+    createTopics()
+
+    expectNextRecords(IndexedSeq(messages(0), messages(1)), 0)
+    expectNextRecords(IndexedSeq(messages(2)), 2)
+    expectNextRecords(IndexedSeq(messages(3), messages(4)), 3)
+  }
+
+  private def expectNextRecords(expected: IndexedSeq[Array[Byte]],
+                                fetchOffset: Long): Unit = {
+    val response = sendFetchRequest(0,
+      FetchRequest.Builder.forConsumer(Int.MaxValue, 0,
+        Map(testTopicPartition ->
+          new PartitionData(fetchOffset, 0, Integer.MAX_VALUE, Optional.empty())).asJava).build(3))
+    val records = response.responseData().get(testTopicPartition).records.records()
+    Assert.assertNotNull(records)
+    val recordsList = records.asScala.toList
+    Assert.assertEquals(expected.size, recordsList.size)
+    recordsList.zipWithIndex.foreach {
+      case (record, i) => {
+        val buffer = record.value().duplicate()
+        val array = new Array[Byte](buffer.remaining())
+        buffer.get(array)
+        Assert.assertArrayEquals(s"expectNextRecords unexpected element ${i}",
+          expected(i), array)
+      }
+    }
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/server/FetchRequestTest.scala b/core/src/test/scala/unit/kafka/server/FetchRequestTest.scala
index 726757d8349cc..97defd7b1b563 100644
--- a/core/src/test/scala/unit/kafka/server/FetchRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/FetchRequestTest.scala
@@ -18,22 +18,23 @@ package kafka.server
 
 import java.io.DataInputStream
 import java.util
-import java.util.Properties
+import java.util.{Optional, Properties}
 
 import kafka.api.KAFKA_0_11_0_IV2
 import kafka.log.LogConfig
+import kafka.message.{GZIPCompressionCodec, ProducerCompressionCodec, ZStdCompressionCodec}
 import kafka.utils.TestUtils
-import kafka.utils.TestUtils._
-import org.apache.kafka.clients.producer.{KafkaProducer, ProducerConfig, ProducerRecord}
-import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.clients.producer.{KafkaProducer, ProducerRecord, RecordMetadata}
 import org.apache.kafka.common.protocol.{ApiKeys, Errors}
-import org.apache.kafka.common.record.{Record, RecordBatch}
-import org.apache.kafka.common.requests.{FetchRequest, FetchResponse}
+import org.apache.kafka.common.record.{MemoryRecords, Record, RecordBatch}
+import org.apache.kafka.common.requests.{FetchRequest, FetchResponse, FetchMetadata => JFetchMetadata}
 import org.apache.kafka.common.serialization.{ByteArraySerializer, StringSerializer}
+import org.apache.kafka.common.{IsolationLevel, TopicPartition}
 import org.junit.Assert._
 import org.junit.Test
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
+import scala.collection.Seq
 import scala.util.Random
 
 /**
@@ -44,7 +45,11 @@ class FetchRequestTest extends BaseRequestTest {
 
   private var producer: KafkaProducer[String, String] = null
 
-  override def tearDown() {
+  override def brokerPropertyOverrides(properties: Properties): Unit = {
+    properties.put(KafkaConfig.FetchMaxBytes, Int.MaxValue.toString)
+  }
+
+  override def tearDown(): Unit = {
     if (producer != null)
       producer.close()
     super.tearDown()
@@ -59,19 +64,19 @@ class FetchRequestTest extends BaseRequestTest {
                                  offsetMap: Map[TopicPartition, Long] = Map.empty): util.LinkedHashMap[TopicPartition, FetchRequest.PartitionData] = {
     val partitionMap = new util.LinkedHashMap[TopicPartition, FetchRequest.PartitionData]
     topicPartitions.foreach { tp =>
-      partitionMap.put(tp, new FetchRequest.PartitionData(offsetMap.getOrElse(tp, 0), 0L, maxPartitionBytes))
+      partitionMap.put(tp, new FetchRequest.PartitionData(offsetMap.getOrElse(tp, 0), 0L, maxPartitionBytes,
+        Optional.empty()))
     }
     partitionMap
   }
 
-  private def sendFetchRequest(leaderId: Int, request: FetchRequest): FetchResponse = {
-    val response = connectAndSend(request, ApiKeys.FETCH, destination = brokerSocketServer(leaderId))
-    FetchResponse.parse(response, request.version)
+  private def sendFetchRequest(leaderId: Int, request: FetchRequest): FetchResponse[MemoryRecords] = {
+    connectAndReceive[FetchResponse[MemoryRecords]](request, destination = brokerSocketServer(leaderId))
   }
 
   private def initProducer(): Unit = {
-    producer = TestUtils.createNewProducer(TestUtils.getBrokerListStrFromServers(servers),
-      retries = 5, keySerializer = new StringSerializer, valueSerializer = new StringSerializer)
+    producer = TestUtils.createProducer(TestUtils.getBrokerListStrFromServers(servers),
+      keySerializer = new StringSerializer, valueSerializer = new StringSerializer)
   }
 
   @Test
@@ -170,8 +175,170 @@ class FetchRequestTest extends BaseRequestTest {
     assertEquals(0, records(partitionData).map(_.sizeInBytes).sum)
   }
 
+  @Test
+  def testFetchRequestV4WithReadCommitted(): Unit = {
+    initProducer()
+    val maxPartitionBytes = 200
+    val (topicPartition, leaderId) = createTopics(numTopics = 1, numPartitions = 1).head
+    producer.send(new ProducerRecord(topicPartition.topic, topicPartition.partition,
+      "key", new String(new Array[Byte](maxPartitionBytes + 1)))).get
+    val fetchRequest = FetchRequest.Builder.forConsumer(Int.MaxValue, 0, createPartitionMap(maxPartitionBytes,
+      Seq(topicPartition))).isolationLevel(IsolationLevel.READ_COMMITTED).build(4)
+    val fetchResponse = sendFetchRequest(leaderId, fetchRequest)
+    val partitionData = fetchResponse.responseData.get(topicPartition)
+    assertEquals(Errors.NONE, partitionData.error)
+    assertTrue(partitionData.lastStableOffset > 0)
+    assertTrue(records(partitionData).map(_.sizeInBytes).sum > 0)
+  }
+
+  @Test
+  def testFetchRequestToNonReplica(): Unit = {
+    val topic = "topic"
+    val partition = 0
+    val topicPartition = new TopicPartition(topic, partition)
+
+    // Create a single-partition topic and find a broker which is not the leader
+    val partitionToLeader = TestUtils.createTopic(zkClient, topic, numPartitions = 1, 1, servers)
+    val leader = partitionToLeader(partition)
+    val nonReplicaOpt = servers.find(_.config.brokerId != leader)
+    assertTrue(nonReplicaOpt.isDefined)
+    val nonReplicaId =  nonReplicaOpt.get.config.brokerId
+
+    // Send the fetch request to the non-replica and verify the error code
+    val fetchRequest = FetchRequest.Builder.forConsumer(Int.MaxValue, 0, createPartitionMap(1024,
+      Seq(topicPartition))).build()
+    val fetchResponse = sendFetchRequest(nonReplicaId, fetchRequest)
+    val partitionData = fetchResponse.responseData.get(topicPartition)
+    assertEquals(Errors.NOT_LEADER_OR_FOLLOWER, partitionData.error)
+  }
+
+  @Test
+  def testLastFetchedEpochValidation(): Unit = {
+    val topic = "topic"
+    val topicPartition = new TopicPartition(topic, 0)
+    val partitionToLeader = TestUtils.createTopic(zkClient, topic, numPartitions = 1, replicationFactor = 3, servers)
+    val firstLeaderId = partitionToLeader(topicPartition.partition)
+    val firstLeaderEpoch = TestUtils.findLeaderEpoch(firstLeaderId, topicPartition, servers)
+
+    initProducer()
+
+    // Write some data in epoch 0
+    val firstEpochResponses = produceData(Seq(topicPartition), 100)
+    val firstEpochEndOffset = firstEpochResponses.lastOption.get.offset + 1
+    // Force a leader change
+    killBroker(firstLeaderId)
+    // Write some more data in epoch 1
+    val secondLeaderId = TestUtils.awaitLeaderChange(servers, topicPartition, firstLeaderId)
+    val secondLeaderEpoch = TestUtils.findLeaderEpoch(secondLeaderId, topicPartition, servers)
+    val secondEpochResponses = produceData(Seq(topicPartition), 100)
+    val secondEpochEndOffset = secondEpochResponses.lastOption.get.offset + 1
+
+    // Build a fetch request in the middle of the second epoch, but with the first epoch
+    val fetchOffset = secondEpochEndOffset + (secondEpochEndOffset - firstEpochEndOffset) / 2
+    val partitionMap = new util.LinkedHashMap[TopicPartition, FetchRequest.PartitionData]
+    partitionMap.put(topicPartition, new FetchRequest.PartitionData(fetchOffset, 0L, 1024,
+      Optional.of(secondLeaderEpoch), Optional.of(firstLeaderEpoch)))
+    val fetchRequest = FetchRequest.Builder.forConsumer(0, 1, partitionMap).build()
+
+    // Validate the expected truncation
+    val fetchResponse = sendFetchRequest(secondLeaderId, fetchRequest)
+    val partitionData = fetchResponse.responseData.get(topicPartition)
+    assertEquals(Errors.NONE, partitionData.error)
+    assertEquals(0L, partitionData.records.sizeInBytes())
+    assertTrue(partitionData.divergingEpoch.isPresent)
+
+    val divergingEpoch = partitionData.divergingEpoch.get()
+    assertEquals(firstLeaderEpoch, divergingEpoch.epoch)
+    assertEquals(firstEpochEndOffset, divergingEpoch.endOffset)
+  }
+
+  @Test
+  def testCurrentEpochValidation(): Unit = {
+    val topic = "topic"
+    val topicPartition = new TopicPartition(topic, 0)
+    val partitionToLeader = TestUtils.createTopic(zkClient, topic, numPartitions = 1, replicationFactor = 3, servers)
+    val firstLeaderId = partitionToLeader(topicPartition.partition)
+
+    def assertResponseErrorForEpoch(error: Errors, brokerId: Int, leaderEpoch: Optional[Integer]): Unit = {
+      val partitionMap = new util.LinkedHashMap[TopicPartition, FetchRequest.PartitionData]
+      partitionMap.put(topicPartition, new FetchRequest.PartitionData(0L, 0L, 1024, leaderEpoch))
+      val fetchRequest = FetchRequest.Builder.forConsumer(0, 1, partitionMap).build()
+      val fetchResponse = sendFetchRequest(brokerId, fetchRequest)
+      val partitionData = fetchResponse.responseData.get(topicPartition)
+      assertEquals(error, partitionData.error)
+    }
+
+    // We need a leader change in order to check epoch fencing since the first epoch is 0 and
+    // -1 is treated as having no epoch at all
+    killBroker(firstLeaderId)
+
+    // Check leader error codes
+    val secondLeaderId = TestUtils.awaitLeaderChange(servers, topicPartition, firstLeaderId)
+    val secondLeaderEpoch = TestUtils.findLeaderEpoch(secondLeaderId, topicPartition, servers)
+    assertResponseErrorForEpoch(Errors.NONE, secondLeaderId, Optional.empty())
+    assertResponseErrorForEpoch(Errors.NONE, secondLeaderId, Optional.of(secondLeaderEpoch))
+    assertResponseErrorForEpoch(Errors.FENCED_LEADER_EPOCH, secondLeaderId, Optional.of(secondLeaderEpoch - 1))
+    assertResponseErrorForEpoch(Errors.UNKNOWN_LEADER_EPOCH, secondLeaderId, Optional.of(secondLeaderEpoch + 1))
+
+    // Check follower error codes
+    val followerId = TestUtils.findFollowerId(topicPartition, servers)
+    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))
+  }
+
+  @Test
+  def testEpochValidationWithinFetchSession(): Unit = {
+    val topic = "topic"
+    val topicPartition = new TopicPartition(topic, 0)
+    val partitionToLeader = TestUtils.createTopic(zkClient, topic, numPartitions = 1, replicationFactor = 3, servers)
+    val firstLeaderId = partitionToLeader(topicPartition.partition)
+
+    // We need a leader change in order to check epoch fencing since the first epoch is 0 and
+    // -1 is treated as having no epoch at all
+    killBroker(firstLeaderId)
+
+    val secondLeaderId = TestUtils.awaitLeaderChange(servers, topicPartition, firstLeaderId)
+    val secondLeaderEpoch = TestUtils.findLeaderEpoch(secondLeaderId, topicPartition, servers)
+    verifyFetchSessionErrors(topicPartition, secondLeaderEpoch, secondLeaderId)
+
+    val followerId = TestUtils.findFollowerId(topicPartition, servers)
+    verifyFetchSessionErrors(topicPartition, secondLeaderEpoch, followerId)
+  }
+
+  private def verifyFetchSessionErrors(topicPartition: TopicPartition,
+                                       leaderEpoch: Int,
+                                       destinationBrokerId: Int): Unit = {
+    val partitionMap = new util.LinkedHashMap[TopicPartition, FetchRequest.PartitionData]
+    partitionMap.put(topicPartition, new FetchRequest.PartitionData(0L, 0L, 1024,
+      Optional.of(leaderEpoch)))
+    val fetchRequest = FetchRequest.Builder.forConsumer(0, 1, partitionMap)
+      .metadata(JFetchMetadata.INITIAL)
+      .build()
+    val fetchResponse = sendFetchRequest(destinationBrokerId, fetchRequest)
+    val sessionId = fetchResponse.sessionId
+
+    def assertResponseErrorForEpoch(expectedError: Errors,
+                                    sessionFetchEpoch: Int,
+                                    leaderEpoch: Optional[Integer]): Unit = {
+      val partitionMap = new util.LinkedHashMap[TopicPartition, FetchRequest.PartitionData]
+      partitionMap.put(topicPartition, new FetchRequest.PartitionData(0L, 0L, 1024, leaderEpoch))
+      val fetchRequest = FetchRequest.Builder.forConsumer(0, 1, partitionMap)
+        .metadata(new JFetchMetadata(sessionId, sessionFetchEpoch))
+        .build()
+      val fetchResponse = sendFetchRequest(destinationBrokerId, fetchRequest)
+      val partitionData = fetchResponse.responseData.get(topicPartition)
+      assertEquals(expectedError, partitionData.error)
+    }
+
+    // We only check errors because we do not expect the partition in the response otherwise
+    assertResponseErrorForEpoch(Errors.FENCED_LEADER_EPOCH, 1, Optional.of(leaderEpoch - 1))
+    assertResponseErrorForEpoch(Errors.UNKNOWN_LEADER_EPOCH, 2, Optional.of(leaderEpoch + 1))
+  }
+
   /**
-   * Tests that down-conversions dont leak memory. Large down conversions are triggered
+   * Tests that down-conversions don't leak memory. Large down conversions are triggered
    * in the server. The client closes its connection after reading partial data when the
    * channel is muted in the server. If buffers are not released this will result in OOM.
    */
@@ -181,11 +348,12 @@ class FetchRequestTest extends BaseRequestTest {
 
     val msgValueLen = 100 * 1000
     val batchSize = 4 * msgValueLen
-    val propsOverride = new Properties
-    propsOverride.put(ProducerConfig.BATCH_SIZE_CONFIG, batchSize.toString)
-    val producer = TestUtils.createNewProducer(TestUtils.getBrokerListStrFromServers(servers),
-      retries = 5, lingerMs = Long.MaxValue,
-      keySerializer = new StringSerializer, valueSerializer = new ByteArraySerializer, props = Some(propsOverride))
+    val producer = TestUtils.createProducer(TestUtils.getBrokerListStrFromServers(servers),
+      lingerMs = Int.MaxValue,
+      deliveryTimeoutMs = Int.MaxValue,
+      batchSize = batchSize,
+      keySerializer = new StringSerializer,
+      valueSerializer = new ByteArraySerializer)
     val bytes = new Array[Byte](msgValueLen)
     val futures = try {
       (0 to 1000).map { _ =>
@@ -198,13 +366,13 @@ class FetchRequestTest extends BaseRequestTest {
     // batch is not complete, but sent when the producer is closed
     futures.foreach(_.get)
 
-    def fetch(version: Short, maxPartitionBytes: Int, closeAfterPartialResponse: Boolean): Option[FetchResponse] = {
+    def fetch(version: Short, maxPartitionBytes: Int, closeAfterPartialResponse: Boolean): Option[FetchResponse[MemoryRecords]] = {
       val fetchRequest = FetchRequest.Builder.forConsumer(Int.MaxValue, 0, createPartitionMap(maxPartitionBytes,
         Seq(topicPartition))).build(version)
 
       val socket = connect(brokerSocketServer(leaderId))
       try {
-        send(fetchRequest, ApiKeys.FETCH, socket)
+        send(fetchRequest, socket)
         if (closeAfterPartialResponse) {
           // read some data to ensure broker has muted this channel and then close socket
           val size = new DataInputStream(socket.getInputStream).readInt()
@@ -215,7 +383,7 @@ class FetchRequestTest extends BaseRequestTest {
               size > maxPartitionBytes - batchSize)
           None
         } else {
-          Some(FetchResponse.parse(receive(socket), version))
+          Some(receive[FetchResponse[MemoryRecords]](socket, ApiKeys.FETCH, version))
         }
       } finally {
         socket.close()
@@ -241,9 +409,12 @@ class FetchRequestTest extends BaseRequestTest {
   @Test
   def testDownConversionFromBatchedToUnbatchedRespectsOffset(): Unit = {
     // Increase linger so that we have control over the batches created
-    producer = TestUtils.createNewProducer(TestUtils.getBrokerListStrFromServers(servers),
-      retries = 5, keySerializer = new StringSerializer, valueSerializer = new StringSerializer,
-      lingerMs = 300 * 1000)
+    producer = TestUtils.createProducer(TestUtils.getBrokerListStrFromServers(servers),
+      retries = 5,
+      keySerializer = new StringSerializer,
+      valueSerializer = new StringSerializer,
+      lingerMs = 30 * 1000,
+      deliveryTimeoutMs = 60 * 1000)
 
     val topicConfig = Map(LogConfig.MessageFormatVersionProp -> KAFKA_0_11_0_IV2.version)
     val (topicPartition, leaderId) = createTopics(numTopics = 1, numPartitions = 1, topicConfig).head
@@ -258,17 +429,32 @@ class FetchRequestTest extends BaseRequestTest {
     secondBatchFutures.foreach(_.get)
 
     def check(fetchOffset: Long, requestVersion: Short, expectedOffset: Long, expectedNumBatches: Int, expectedMagic: Byte): Unit = {
-      val fetchRequest = FetchRequest.Builder.forConsumer(Int.MaxValue, 0, createPartitionMap(Int.MaxValue,
-        Seq(topicPartition), Map(topicPartition -> fetchOffset))).build(requestVersion)
-      val fetchResponse = sendFetchRequest(leaderId, fetchRequest)
-      val partitionData = fetchResponse.responseData.get(topicPartition)
-      assertEquals(Errors.NONE, partitionData.error)
-      assertTrue(partitionData.highWatermark > 0)
-      val batches = partitionData.records.batches.asScala.toBuffer
-      assertEquals(expectedNumBatches, batches.size)
-      val batch = batches.head
-      assertEquals(expectedMagic, batch.magic)
-      assertEquals(expectedOffset, batch.baseOffset)
+      var batchesReceived = 0
+      var currentFetchOffset = fetchOffset
+      var currentExpectedOffset = expectedOffset
+
+      // With KIP-283, we might not receive all batches in a single fetch request so loop through till we have consumed
+      // all batches we are interested in.
+      while (batchesReceived < expectedNumBatches) {
+        val fetchRequest = FetchRequest.Builder.forConsumer(Int.MaxValue, 0, createPartitionMap(Int.MaxValue,
+          Seq(topicPartition), Map(topicPartition -> currentFetchOffset))).build(requestVersion)
+        val fetchResponse = sendFetchRequest(leaderId, fetchRequest)
+
+        // validate response
+        val partitionData = fetchResponse.responseData.get(topicPartition)
+        assertEquals(Errors.NONE, partitionData.error)
+        assertTrue(partitionData.highWatermark > 0)
+        val batches = partitionData.records.batches.asScala.toBuffer
+        val batch = batches.head
+        assertEquals(expectedMagic, batch.magic)
+        assertEquals(currentExpectedOffset, batch.baseOffset)
+
+        currentFetchOffset = batches.last.lastOffset + 1
+        currentExpectedOffset += (batches.last.lastOffset - batches.head.baseOffset + 1)
+        batchesReceived += batches.size
+      }
+
+      assertEquals(expectedNumBatches, batchesReceived)
     }
 
     // down conversion to message format 0, batches of 1 message are returned so we receive the exact offset we requested
@@ -295,11 +481,185 @@ class FetchRequestTest extends BaseRequestTest {
       expectedMagic = RecordBatch.MAGIC_VALUE_V2)
   }
 
-  private def records(partitionData: FetchResponse.PartitionData): Seq[Record] = {
-    partitionData.records.records.asScala.toIndexedSeq
+  /**
+   * Test that when an incremental fetch session contains partitions with an error,
+   * those partitions are returned in all incremental fetch requests.
+   */
+  @Test
+  def testCreateIncrementalFetchWithPartitionsInError(): Unit = {
+    def createFetchRequest(topicPartitions: Seq[TopicPartition],
+                           metadata: JFetchMetadata,
+                           toForget: Seq[TopicPartition]): FetchRequest =
+      FetchRequest.Builder.forConsumer(Int.MaxValue, 0,
+        createPartitionMap(Integer.MAX_VALUE, topicPartitions, Map.empty))
+        .toForget(toForget.asJava)
+        .metadata(metadata)
+        .build()
+    val foo0 = new TopicPartition("foo", 0)
+    val foo1 = new TopicPartition("foo", 1)
+    createTopic("foo", Map(0 -> List(0, 1), 1 -> List(0, 2)))
+    val bar0 = new TopicPartition("bar", 0)
+    val req1 = createFetchRequest(List(foo0, foo1, bar0), JFetchMetadata.INITIAL, Nil)
+    val resp1 = sendFetchRequest(0, req1)
+    assertEquals(Errors.NONE, resp1.error())
+    assertTrue("Expected the broker to create a new incremental fetch session", resp1.sessionId() > 0)
+    debug(s"Test created an incremental fetch session ${resp1.sessionId}")
+    assertTrue(resp1.responseData().containsKey(foo0))
+    assertTrue(resp1.responseData().containsKey(foo1))
+    assertTrue(resp1.responseData().containsKey(bar0))
+    assertEquals(Errors.NONE, resp1.responseData().get(foo0).error)
+    assertEquals(Errors.NONE, resp1.responseData().get(foo1).error)
+    assertEquals(Errors.UNKNOWN_TOPIC_OR_PARTITION, resp1.responseData().get(bar0).error)
+    val req2 = createFetchRequest(Nil, new JFetchMetadata(resp1.sessionId(), 1), Nil)
+    val resp2 = sendFetchRequest(0, req2)
+    assertEquals(Errors.NONE, resp2.error())
+    assertEquals("Expected the broker to continue the incremental fetch session",
+      resp1.sessionId(), resp2.sessionId())
+    assertFalse(resp2.responseData().containsKey(foo0))
+    assertFalse(resp2.responseData().containsKey(foo1))
+    assertTrue(resp2.responseData().containsKey(bar0))
+    assertEquals(Errors.UNKNOWN_TOPIC_OR_PARTITION, resp2.responseData().get(bar0).error)
+    createTopic("bar", Map(0 -> List(0, 1)))
+    val req3 = createFetchRequest(Nil, new JFetchMetadata(resp1.sessionId(), 2), Nil)
+    val resp3 = sendFetchRequest(0, req3)
+    assertEquals(Errors.NONE, resp3.error())
+    assertFalse(resp3.responseData().containsKey(foo0))
+    assertFalse(resp3.responseData().containsKey(foo1))
+    assertTrue(resp3.responseData().containsKey(bar0))
+    assertEquals(Errors.NONE, resp3.responseData().get(bar0).error)
+    val req4 = createFetchRequest(Nil, new JFetchMetadata(resp1.sessionId(), 3), Nil)
+    val resp4 = sendFetchRequest(0, req4)
+    assertEquals(Errors.NONE, resp4.error())
+    assertFalse(resp4.responseData().containsKey(foo0))
+    assertFalse(resp4.responseData().containsKey(foo1))
+    assertFalse(resp4.responseData().containsKey(bar0))
+  }
+
+  @Test
+  def testZStdCompressedTopic(): Unit = {
+    // ZSTD compressed topic
+    val topicConfig = Map(LogConfig.CompressionTypeProp -> ZStdCompressionCodec.name)
+    val (topicPartition, leaderId) = createTopics(numTopics = 1, numPartitions = 1, configs = topicConfig).head
+
+    // Produce messages (v2)
+    producer = TestUtils.createProducer(TestUtils.getBrokerListStrFromServers(servers),
+      keySerializer = new StringSerializer,
+      valueSerializer = new StringSerializer)
+    producer.send(new ProducerRecord(topicPartition.topic, topicPartition.partition,
+      "key1", "value1")).get
+    producer.send(new ProducerRecord(topicPartition.topic, topicPartition.partition,
+      "key2", "value2")).get
+    producer.send(new ProducerRecord(topicPartition.topic, topicPartition.partition,
+      "key3", "value3")).get
+    producer.close()
+
+    // fetch request with version below v10: UNSUPPORTED_COMPRESSION_TYPE error occurs
+    val req0 = new FetchRequest.Builder(0, 9, -1, Int.MaxValue, 0,
+      createPartitionMap(300, Seq(topicPartition), Map.empty))
+      .setMaxBytes(800).build()
+
+    val res0 = sendFetchRequest(leaderId, req0)
+    val data0 = res0.responseData.get(topicPartition)
+    assertEquals(Errors.UNSUPPORTED_COMPRESSION_TYPE, data0.error)
+
+    // fetch request with version 10: works fine!
+    val req1= new FetchRequest.Builder(0, 10, -1, Int.MaxValue, 0,
+      createPartitionMap(300, Seq(topicPartition), Map.empty))
+      .setMaxBytes(800).build()
+    val res1 = sendFetchRequest(leaderId, req1)
+    val data1 = res1.responseData.get(topicPartition)
+    assertEquals(Errors.NONE, data1.error)
+    assertEquals(3, records(data1).size)
+  }
+
+  @Test
+  def testPartitionDataEquals(): Unit = {
+    assertEquals(new FetchRequest.PartitionData(300, 0L, 300, Optional.of(300)),
+    new FetchRequest.PartitionData(300, 0L, 300, Optional.of(300)));
+  }
+
+  @Test
+  def testZStdCompressedRecords(): Unit = {
+    // Producer compressed topic
+    val topicConfig = Map(LogConfig.CompressionTypeProp -> ProducerCompressionCodec.name,
+      LogConfig.MessageFormatVersionProp -> "2.0.0")
+    val (topicPartition, leaderId) = createTopics(numTopics = 1, numPartitions = 1, configs = topicConfig).head
+
+    // Produce GZIP compressed messages (v2)
+    val producer1 = TestUtils.createProducer(TestUtils.getBrokerListStrFromServers(servers),
+      compressionType = GZIPCompressionCodec.name,
+      keySerializer = new StringSerializer,
+      valueSerializer = new StringSerializer)
+    producer1.send(new ProducerRecord(topicPartition.topic, topicPartition.partition,
+      "key1", "value1")).get
+    producer1.close()
+    // Produce ZSTD compressed messages (v2)
+    val producer2 = TestUtils.createProducer(TestUtils.getBrokerListStrFromServers(servers),
+      compressionType = ZStdCompressionCodec.name,
+      keySerializer = new StringSerializer,
+      valueSerializer = new StringSerializer)
+    producer2.send(new ProducerRecord(topicPartition.topic, topicPartition.partition,
+      "key2", "value2")).get
+    producer2.send(new ProducerRecord(topicPartition.topic, topicPartition.partition,
+      "key3", "value3")).get
+    producer2.close()
+
+    // fetch request with fetch version v1 (magic 0):
+    // gzip compressed record is returned with down-conversion.
+    // zstd compressed record raises UNSUPPORTED_COMPRESSION_TYPE error.
+    val req0 = new FetchRequest.Builder(0, 1, -1, Int.MaxValue, 0,
+      createPartitionMap(300, Seq(topicPartition), Map.empty))
+      .setMaxBytes(800)
+      .build()
+
+    val res0 = sendFetchRequest(leaderId, req0)
+    val data0 = res0.responseData.get(topicPartition)
+    assertEquals(Errors.NONE, data0.error)
+    assertEquals(1, records(data0).size)
+
+    val req1 = new FetchRequest.Builder(0, 1, -1, Int.MaxValue, 0,
+      createPartitionMap(300, Seq(topicPartition), Map(topicPartition -> 1L)))
+      .setMaxBytes(800).build()
+
+    val res1 = sendFetchRequest(leaderId, req1)
+    val data1 = res1.responseData.get(topicPartition)
+    assertEquals(Errors.UNSUPPORTED_COMPRESSION_TYPE, data1.error)
+
+    // fetch request with fetch version v3 (magic 1):
+    // gzip compressed record is returned with down-conversion.
+    // zstd compressed record raises UNSUPPORTED_COMPRESSION_TYPE error.
+    val req2 = new FetchRequest.Builder(2, 3, -1, Int.MaxValue, 0,
+      createPartitionMap(300, Seq(topicPartition), Map.empty))
+      .setMaxBytes(800).build()
+
+    val res2 = sendFetchRequest(leaderId, req2)
+    val data2 = res2.responseData.get(topicPartition)
+    assertEquals(Errors.NONE, data2.error)
+    assertEquals(1, records(data2).size)
+
+    val req3 = new FetchRequest.Builder(0, 1, -1, Int.MaxValue, 0,
+      createPartitionMap(300, Seq(topicPartition), Map(topicPartition -> 1L)))
+      .setMaxBytes(800).build()
+
+    val res3 = sendFetchRequest(leaderId, req3)
+    val data3 = res3.responseData.get(topicPartition)
+    assertEquals(Errors.UNSUPPORTED_COMPRESSION_TYPE, data3.error)
+
+    // fetch request with version 10: works fine!
+    val req4= new FetchRequest.Builder(0, 10, -1, Int.MaxValue, 0,
+      createPartitionMap(300, Seq(topicPartition), Map.empty))
+      .setMaxBytes(800).build()
+    val res4 = sendFetchRequest(leaderId, req4)
+    val data4 = res4.responseData.get(topicPartition)
+    assertEquals(Errors.NONE, data4.error)
+    assertEquals(3, records(data4).size)
+  }
+
+  private def records(partitionData: FetchResponse.PartitionData[MemoryRecords]): Seq[Record] = {
+    partitionData.records.records.asScala.toBuffer
   }
 
-  private def checkFetchResponse(expectedPartitions: Seq[TopicPartition], fetchResponse: FetchResponse,
+  private def checkFetchResponse(expectedPartitions: Seq[TopicPartition], fetchResponse: FetchResponse[MemoryRecords],
                                  maxPartitionBytes: Int, maxResponseBytes: Int, numMessagesPerPartition: Int): Unit = {
     assertEquals(expectedPartitions, fetchResponse.responseData.keySet.asScala.toSeq)
     var emptyResponseSeen = false
@@ -314,7 +674,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
@@ -337,18 +697,18 @@ class FetchRequestTest extends BaseRequestTest {
   }
 
   private def createTopics(numTopics: Int, numPartitions: Int, configs: Map[String, String] = Map.empty): Map[TopicPartition, Int] = {
-    val topics = (0 until numPartitions).map(t => s"topic$t")
+    val topics = (0 until numTopics).map(t => s"topic$t")
     val topicConfig = new Properties
     topicConfig.setProperty(LogConfig.MinInSyncReplicasProp, 2.toString)
     configs.foreach { case (k, v) => topicConfig.setProperty(k, v) }
     topics.flatMap { topic =>
-      val partitionToLeader = createTopic(zkUtils, topic, numPartitions = numPartitions, replicationFactor = 2,
-        servers = servers, topicConfig = topicConfig)
+      val partitionToLeader = createTopic(topic, numPartitions = numPartitions, replicationFactor = 2,
+        topicConfig = topicConfig)
       partitionToLeader.map { case (partition, leader) => new TopicPartition(topic, partition) -> leader }
     }.toMap
   }
 
-  private def produceData(topicPartitions: Iterable[TopicPartition], numMessagesPerPartition: Int): Seq[ProducerRecord[String, String]] = {
+  private def produceData(topicPartitions: Iterable[TopicPartition], numMessagesPerPartition: Int): Seq[RecordMetadata] = {
     val records = for {
       tp <- topicPartitions.toSeq
       messageIndex <- 0 until numMessagesPerPartition
@@ -357,7 +717,6 @@ class FetchRequestTest extends BaseRequestTest {
       new ProducerRecord(tp.topic, tp.partition, s"key $suffix", s"value $suffix")
     }
     records.map(producer.send(_).get)
-    records
   }
 
 }
diff --git a/core/src/test/scala/unit/kafka/server/FetchSessionTest.scala b/core/src/test/scala/unit/kafka/server/FetchSessionTest.scala
new file mode 100755
index 0000000000000..9d99fb54152e6
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/FetchSessionTest.scala
@@ -0,0 +1,701 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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
+import java.util.{Collections, Optional}
+
+import kafka.utils.MockTime
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.message.FetchResponseData
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.record.Records
+import org.apache.kafka.common.requests.FetchMetadata.{FINAL_EPOCH, INVALID_SESSION_ID}
+import org.apache.kafka.common.requests.{FetchRequest, FetchResponse, FetchMetadata => JFetchMetadata}
+import org.apache.kafka.common.utils.Utils
+import org.junit.Assert._
+import org.junit.rules.Timeout
+import org.junit.{Rule, Test}
+
+class FetchSessionTest {
+  @Rule
+  def globalTimeout = Timeout.millis(120000)
+
+  @Test
+  def testNewSessionId(): Unit = {
+    val cache = new FetchSessionCache(3, 100)
+    for (i <- 0 to 10000) {
+      val id = cache.newSessionId()
+      assertTrue(id > 0)
+    }
+  }
+
+  def assertCacheContains(cache: FetchSessionCache, sessionIds: Int*) = {
+    var i = 0
+    for (sessionId <- sessionIds) {
+      i = i + 1
+      assertTrue("Missing session " + i + " out of " + sessionIds.size + "(" + sessionId + ")",
+        cache.get(sessionId).isDefined)
+    }
+    assertEquals(sessionIds.size, cache.size)
+  }
+
+  private def dummyCreate(size: Int): FetchSession.CACHE_MAP = {
+    val cacheMap = new FetchSession.CACHE_MAP(size)
+    for (i <- 0 until size) {
+      cacheMap.add(new CachedPartition("test", i))
+    }
+    cacheMap
+  }
+
+  @Test
+  def testSessionCache(): Unit = {
+    val cache = new FetchSessionCache(3, 100)
+    assertEquals(0, cache.size)
+    val id1 = cache.maybeCreateSession(0, false, 10, () => dummyCreate(10))
+    val id2 = cache.maybeCreateSession(10, false, 20, () => dummyCreate(20))
+    val id3 = cache.maybeCreateSession(20, false, 30, () => dummyCreate(30))
+    assertEquals(INVALID_SESSION_ID, cache.maybeCreateSession(30, false, 40, () => dummyCreate(40)))
+    assertEquals(INVALID_SESSION_ID, cache.maybeCreateSession(40, false, 5, () => dummyCreate(5)))
+    assertCacheContains(cache, id1, id2, id3)
+    cache.touch(cache.get(id1).get, 200)
+    val id4 = cache.maybeCreateSession(210, false, 11, () => dummyCreate(11))
+    assertCacheContains(cache, id1, id3, id4)
+    cache.touch(cache.get(id1).get, 400)
+    cache.touch(cache.get(id3).get, 390)
+    cache.touch(cache.get(id4).get, 400)
+    val id5 = cache.maybeCreateSession(410, false, 50, () => dummyCreate(50))
+    assertCacheContains(cache, id3, id4, id5)
+    assertEquals(INVALID_SESSION_ID, cache.maybeCreateSession(410, false, 5, () => dummyCreate(5)))
+    val id6 = cache.maybeCreateSession(410, true, 5, () => dummyCreate(5))
+    assertCacheContains(cache, id3, id5, id6)
+  }
+
+  @Test
+  def testResizeCachedSessions(): Unit = {
+    val cache = new FetchSessionCache(2, 100)
+    assertEquals(0, cache.totalPartitions)
+    assertEquals(0, cache.size)
+    assertEquals(0, cache.evictionsMeter.count)
+    val id1 = cache.maybeCreateSession(0, false, 2, () => dummyCreate(2))
+    assertTrue(id1 > 0)
+    assertCacheContains(cache, id1)
+    val session1 = cache.get(id1).get
+    assertEquals(2, session1.size)
+    assertEquals(2, cache.totalPartitions)
+    assertEquals(1, cache.size)
+    assertEquals(0, cache.evictionsMeter.count)
+    val id2 = cache.maybeCreateSession(0, false, 4, () => dummyCreate(4))
+    val session2 = cache.get(id2).get
+    assertTrue(id2 > 0)
+    assertCacheContains(cache, id1, id2)
+    assertEquals(6, cache.totalPartitions)
+    assertEquals(2, cache.size)
+    assertEquals(0, cache.evictionsMeter.count)
+    cache.touch(session1, 200)
+    cache.touch(session2, 200)
+    val id3 = cache.maybeCreateSession(200, false, 5, () => dummyCreate(5))
+    assertTrue(id3 > 0)
+    assertCacheContains(cache, id2, id3)
+    assertEquals(9, cache.totalPartitions)
+    assertEquals(2, cache.size)
+    assertEquals(1, cache.evictionsMeter.count)
+    cache.remove(id3)
+    assertCacheContains(cache, id2)
+    assertEquals(1, cache.size)
+    assertEquals(1, cache.evictionsMeter.count)
+    assertEquals(4, cache.totalPartitions)
+    val iter = session2.partitionMap.iterator
+    iter.next()
+    iter.remove()
+    assertEquals(3, session2.size)
+    assertEquals(4, session2.cachedSize)
+    cache.touch(session2, session2.lastUsedMs)
+    assertEquals(3, cache.totalPartitions)
+  }
+
+  val EMPTY_PART_LIST = Collections.unmodifiableList(new util.ArrayList[TopicPartition]())
+
+
+  @Test
+  def testCachedLeaderEpoch(): Unit = {
+    val time = new MockTime()
+    val cache = new FetchSessionCache(10, 1000)
+    val fetchManager = new FetchManager(time, cache)
+
+    val tp0 = new TopicPartition("foo", 0)
+    val tp1 = new TopicPartition("foo", 1)
+    val tp2 = new TopicPartition("bar", 1)
+
+    def cachedLeaderEpochs(context: FetchContext): Map[TopicPartition, Optional[Integer]] = {
+      val mapBuilder = Map.newBuilder[TopicPartition, Optional[Integer]]
+      context.foreachPartition((tp, data) => mapBuilder += tp -> data.currentLeaderEpoch)
+      mapBuilder.result()
+    }
+
+    val request1 = new util.LinkedHashMap[TopicPartition, FetchRequest.PartitionData]
+    request1.put(tp0, new FetchRequest.PartitionData(0, 0, 100, Optional.empty()))
+    request1.put(tp1, new FetchRequest.PartitionData(10, 0, 100, Optional.of(1)))
+    request1.put(tp2, new FetchRequest.PartitionData(10, 0, 100, Optional.of(2)))
+
+    val context1 = fetchManager.newContext(JFetchMetadata.INITIAL, request1, EMPTY_PART_LIST, false)
+    val epochs1 = cachedLeaderEpochs(context1)
+    assertEquals(Optional.empty(), epochs1(tp0))
+    assertEquals(Optional.of(1), epochs1(tp1))
+    assertEquals(Optional.of(2), epochs1(tp2))
+
+    val response = new util.LinkedHashMap[TopicPartition, FetchResponse.PartitionData[Records]]
+    response.put(tp0, new FetchResponse.PartitionData(Errors.NONE, 100, 100,
+      100, null, null))
+    response.put(tp1, new FetchResponse.PartitionData(
+      Errors.NONE, 10, 10, 10, null, null))
+    response.put(tp2, new FetchResponse.PartitionData(
+      Errors.NONE, 5, 5, 5, null, null))
+
+    val sessionId = context1.updateAndGenerateResponseData(response).sessionId()
+
+    // With no changes, the cached epochs should remain the same
+    val request2 = new util.LinkedHashMap[TopicPartition, FetchRequest.PartitionData]
+    val context2 = fetchManager.newContext(new JFetchMetadata(sessionId, 1), request2, EMPTY_PART_LIST, false)
+    val epochs2 = cachedLeaderEpochs(context2)
+    assertEquals(Optional.empty(), epochs1(tp0))
+    assertEquals(Optional.of(1), epochs2(tp1))
+    assertEquals(Optional.of(2), epochs2(tp2))
+    context2.updateAndGenerateResponseData(response).sessionId()
+
+    // Now verify we can change the leader epoch and the context is updated
+    val request3 = new util.LinkedHashMap[TopicPartition, FetchRequest.PartitionData]
+    request3.put(tp0, new FetchRequest.PartitionData(0, 0, 100, Optional.of(6)))
+    request3.put(tp1, new FetchRequest.PartitionData(10, 0, 100, Optional.empty()))
+    request3.put(tp2, new FetchRequest.PartitionData(10, 0, 100, Optional.of(3)))
+
+    val context3 = fetchManager.newContext(new JFetchMetadata(sessionId, 2), request3, EMPTY_PART_LIST, false)
+    val epochs3 = cachedLeaderEpochs(context3)
+    assertEquals(Optional.of(6), epochs3(tp0))
+    assertEquals(Optional.empty(), epochs3(tp1))
+    assertEquals(Optional.of(3), epochs3(tp2))
+  }
+
+  @Test
+  def testLastFetchedEpoch(): Unit = {
+    val time = new MockTime()
+    val cache = new FetchSessionCache(10, 1000)
+    val fetchManager = new FetchManager(time, cache)
+
+    val tp0 = new TopicPartition("foo", 0)
+    val tp1 = new TopicPartition("foo", 1)
+    val tp2 = new TopicPartition("bar", 1)
+
+    def cachedLeaderEpochs(context: FetchContext): Map[TopicPartition, Optional[Integer]] = {
+      val mapBuilder = Map.newBuilder[TopicPartition, Optional[Integer]]
+      context.foreachPartition((tp, data) => mapBuilder += tp -> data.currentLeaderEpoch)
+      mapBuilder.result()
+    }
+
+    def cachedLastFetchedEpochs(context: FetchContext): Map[TopicPartition, Optional[Integer]] = {
+      val mapBuilder = Map.newBuilder[TopicPartition, Optional[Integer]]
+      context.foreachPartition((tp, data) => mapBuilder += tp -> data.lastFetchedEpoch)
+      mapBuilder.result()
+    }
+
+    val request1 = new util.LinkedHashMap[TopicPartition, FetchRequest.PartitionData]
+    request1.put(tp0, new FetchRequest.PartitionData(0, 0, 100, Optional.empty[Integer], Optional.empty[Integer]))
+    request1.put(tp1, new FetchRequest.PartitionData(10, 0, 100, Optional.of(1), Optional.empty[Integer]))
+    request1.put(tp2, new FetchRequest.PartitionData(10, 0, 100, Optional.of(2), Optional.of(1)))
+
+    val context1 = fetchManager.newContext(JFetchMetadata.INITIAL, request1, EMPTY_PART_LIST, false)
+    assertEquals(Map(tp0 -> Optional.empty, tp1 -> Optional.of(1), tp2 -> Optional.of(2)),
+      cachedLeaderEpochs(context1))
+    assertEquals(Map(tp0 -> Optional.empty, tp1 -> Optional.empty, tp2 -> Optional.of(1)),
+      cachedLastFetchedEpochs(context1))
+
+    val response = new util.LinkedHashMap[TopicPartition, FetchResponse.PartitionData[Records]]
+    response.put(tp0, new FetchResponse.PartitionData(Errors.NONE, 100, 100, 100, null, null))
+    response.put(tp1, new FetchResponse.PartitionData(Errors.NONE, 10, 10, 10, null, null))
+    response.put(tp2, new FetchResponse.PartitionData(Errors.NONE, 5, 5, 5, null, null))
+
+    val sessionId = context1.updateAndGenerateResponseData(response).sessionId()
+
+    // With no changes, the cached epochs should remain the same
+    val request2 = new util.LinkedHashMap[TopicPartition, FetchRequest.PartitionData]
+    val context2 = fetchManager.newContext(new JFetchMetadata(sessionId, 1), request2, EMPTY_PART_LIST, false)
+    assertEquals(Map(tp0 -> Optional.empty, tp1 -> Optional.of(1), tp2 -> Optional.of(2)), cachedLeaderEpochs(context2))
+    assertEquals(Map(tp0 -> Optional.empty, tp1 -> Optional.empty, tp2 -> Optional.of(1)),
+      cachedLastFetchedEpochs(context2))
+    context2.updateAndGenerateResponseData(response).sessionId()
+
+    // Now verify we can change the leader epoch and the context is updated
+    val request3 = new util.LinkedHashMap[TopicPartition, FetchRequest.PartitionData]
+    request3.put(tp0, new FetchRequest.PartitionData(0, 0, 100, Optional.of(6), Optional.of(5)))
+    request3.put(tp1, new FetchRequest.PartitionData(10, 0, 100, Optional.empty[Integer], Optional.empty[Integer]))
+    request3.put(tp2, new FetchRequest.PartitionData(10, 0, 100, Optional.of(3), Optional.of(3)))
+
+    val context3 = fetchManager.newContext(new JFetchMetadata(sessionId, 2), request3, EMPTY_PART_LIST, false)
+    assertEquals(Map(tp0 -> Optional.of(6), tp1 -> Optional.empty, tp2 -> Optional.of(3)),
+      cachedLeaderEpochs(context3))
+    assertEquals(Map(tp0 -> Optional.of(5), tp1 -> Optional.empty, tp2 -> Optional.of(3)),
+      cachedLastFetchedEpochs(context2))
+  }
+
+  @Test
+  def testFetchRequests(): Unit = {
+    val time = new MockTime()
+    val cache = new FetchSessionCache(10, 1000)
+    val fetchManager = new FetchManager(time, cache)
+
+    // Verify that SESSIONLESS requests get a SessionlessFetchContext
+    val context = fetchManager.newContext(JFetchMetadata.LEGACY,
+        new util.HashMap[TopicPartition, FetchRequest.PartitionData](), EMPTY_PART_LIST, true)
+    assertEquals(classOf[SessionlessFetchContext], context.getClass)
+
+    // Create a new fetch session with a FULL fetch request
+    val reqData2 = new util.LinkedHashMap[TopicPartition, FetchRequest.PartitionData]
+    reqData2.put(new TopicPartition("foo", 0), new FetchRequest.PartitionData(0, 0, 100,
+      Optional.empty()))
+    reqData2.put(new TopicPartition("foo", 1), new FetchRequest.PartitionData(10, 0, 100,
+      Optional.empty()))
+    val context2 = fetchManager.newContext(JFetchMetadata.INITIAL, reqData2, EMPTY_PART_LIST, false)
+    assertEquals(classOf[FullFetchContext], context2.getClass)
+    val reqData2Iter = reqData2.entrySet().iterator()
+    context2.foreachPartition((topicPart, data) => {
+      val entry = reqData2Iter.next()
+      assertEquals(entry.getKey, topicPart)
+      assertEquals(entry.getValue, data)
+    })
+    assertEquals(0, context2.getFetchOffset(new TopicPartition("foo", 0)).get)
+    assertEquals(10, context2.getFetchOffset(new TopicPartition("foo", 1)).get)
+    val respData2 = new util.LinkedHashMap[TopicPartition, FetchResponse.PartitionData[Records]]
+    respData2.put(new TopicPartition("foo", 0), new FetchResponse.PartitionData(
+      Errors.NONE, 100, 100, 100, null, null))
+    respData2.put(new TopicPartition("foo", 1), new FetchResponse.PartitionData(
+      Errors.NONE, 10, 10, 10, null, null))
+    val resp2 = context2.updateAndGenerateResponseData(respData2)
+    assertEquals(Errors.NONE, resp2.error())
+    assertTrue(resp2.sessionId() != INVALID_SESSION_ID)
+    assertEquals(respData2, resp2.responseData())
+
+    // Test trying to create a new session with an invalid epoch
+    val context3 = fetchManager.newContext(
+      new JFetchMetadata(resp2.sessionId(), 5), reqData2, EMPTY_PART_LIST, false)
+    assertEquals(classOf[SessionErrorContext], context3.getClass)
+    assertEquals(Errors.INVALID_FETCH_SESSION_EPOCH,
+      context3.updateAndGenerateResponseData(respData2).error())
+
+    // Test trying to create a new session with a non-existent session id
+    val context4 = fetchManager.newContext(
+      new JFetchMetadata(resp2.sessionId() + 1, 1), reqData2, EMPTY_PART_LIST, false)
+    assertEquals(classOf[SessionErrorContext], context4.getClass)
+    assertEquals(Errors.FETCH_SESSION_ID_NOT_FOUND,
+      context4.updateAndGenerateResponseData(respData2).error())
+
+    // Continue the first fetch session we created.
+    val reqData5 = new util.LinkedHashMap[TopicPartition, FetchRequest.PartitionData]
+    val context5 = fetchManager.newContext(
+      new JFetchMetadata(resp2.sessionId(), 1), reqData5, EMPTY_PART_LIST, false)
+    assertEquals(classOf[IncrementalFetchContext], context5.getClass)
+    val reqData5Iter = reqData2.entrySet().iterator()
+    context5.foreachPartition((topicPart, data) => {
+      val entry = reqData5Iter.next()
+      assertEquals(entry.getKey, topicPart)
+      assertEquals(entry.getValue, data)
+    })
+    assertEquals(10, context5.getFetchOffset(new TopicPartition("foo", 1)).get)
+    val resp5 = context5.updateAndGenerateResponseData(respData2)
+    assertEquals(Errors.NONE, resp5.error())
+    assertEquals(resp2.sessionId(), resp5.sessionId())
+    assertEquals(0, resp5.responseData().size())
+
+    // Test setting an invalid fetch session epoch.
+    val context6 = fetchManager.newContext(
+      new JFetchMetadata(resp2.sessionId(), 5), reqData2, EMPTY_PART_LIST, false)
+    assertEquals(classOf[SessionErrorContext], context6.getClass)
+    assertEquals(Errors.INVALID_FETCH_SESSION_EPOCH,
+      context6.updateAndGenerateResponseData(respData2).error())
+
+    // Test generating a throttled response for the incremental fetch session
+    val reqData7 = new util.LinkedHashMap[TopicPartition, FetchRequest.PartitionData]
+    val context7 = fetchManager.newContext(
+      new JFetchMetadata(resp2.sessionId(), 2), reqData7, EMPTY_PART_LIST, false)
+    val resp7 = context7.getThrottledResponse(100)
+    assertEquals(Errors.NONE, resp7.error())
+    assertEquals(resp2.sessionId(), resp7.sessionId())
+    assertEquals(100, resp7.throttleTimeMs())
+
+    // Close the incremental fetch session.
+    val prevSessionId = resp5.sessionId
+    var nextSessionId = prevSessionId
+    do {
+      val reqData8 = new util.LinkedHashMap[TopicPartition, FetchRequest.PartitionData]
+      reqData8.put(new TopicPartition("bar", 0), new FetchRequest.PartitionData(0, 0, 100,
+        Optional.empty()))
+      reqData8.put(new TopicPartition("bar", 1), new FetchRequest.PartitionData(10, 0, 100,
+        Optional.empty()))
+      val context8 = fetchManager.newContext(
+        new JFetchMetadata(prevSessionId, FINAL_EPOCH), reqData8, EMPTY_PART_LIST, false)
+      assertEquals(classOf[SessionlessFetchContext], context8.getClass)
+      assertEquals(0, cache.size)
+      val respData8 = new util.LinkedHashMap[TopicPartition, FetchResponse.PartitionData[Records]]
+      respData8.put(new TopicPartition("bar", 0),
+        new FetchResponse.PartitionData(Errors.NONE, 100, 100, 100, null, null))
+      respData8.put(new TopicPartition("bar", 1),
+        new FetchResponse.PartitionData(Errors.NONE, 100, 100, 100, null, null))
+      val resp8 = context8.updateAndGenerateResponseData(respData8)
+      assertEquals(Errors.NONE, resp8.error)
+      nextSessionId = resp8.sessionId
+    } while (nextSessionId == prevSessionId)
+  }
+
+  @Test
+  def testIncrementalFetchSession(): Unit = {
+    val time = new MockTime()
+    val cache = new FetchSessionCache(10, 1000)
+    val fetchManager = new FetchManager(time, cache)
+
+    // Create a new fetch session with foo-0 and foo-1
+    val reqData1 = new util.LinkedHashMap[TopicPartition, FetchRequest.PartitionData]
+    reqData1.put(new TopicPartition("foo", 0), new FetchRequest.PartitionData(0, 0, 100,
+      Optional.empty()))
+    reqData1.put(new TopicPartition("foo", 1), new FetchRequest.PartitionData(10, 0, 100,
+      Optional.empty()))
+    val context1 = fetchManager.newContext(JFetchMetadata.INITIAL, reqData1, EMPTY_PART_LIST, false)
+    assertEquals(classOf[FullFetchContext], context1.getClass)
+    val respData1 = new util.LinkedHashMap[TopicPartition, FetchResponse.PartitionData[Records]]
+    respData1.put(new TopicPartition("foo", 0), new FetchResponse.PartitionData(
+      Errors.NONE, 100, 100, 100, null, null))
+    respData1.put(new TopicPartition("foo", 1), new FetchResponse.PartitionData(
+      Errors.NONE, 10, 10, 10, null, null))
+    val resp1 = context1.updateAndGenerateResponseData(respData1)
+    assertEquals(Errors.NONE, resp1.error())
+    assertTrue(resp1.sessionId() != INVALID_SESSION_ID)
+    assertEquals(2, resp1.responseData().size())
+
+    // Create an incremental fetch request that removes foo-0 and adds bar-0
+    val reqData2 = new util.LinkedHashMap[TopicPartition, FetchRequest.PartitionData]
+    reqData2.put(new TopicPartition("bar", 0), new FetchRequest.PartitionData(15, 0, 0,
+      Optional.empty()))
+    val removed2 = new util.ArrayList[TopicPartition]
+    removed2.add(new TopicPartition("foo", 0))
+    val context2 = fetchManager.newContext(
+      new JFetchMetadata(resp1.sessionId(), 1), reqData2, removed2, false)
+    assertEquals(classOf[IncrementalFetchContext], context2.getClass)
+    val parts2 = Set(new TopicPartition("foo", 1), new TopicPartition("bar", 0))
+    val reqData2Iter = parts2.iterator
+    context2.foreachPartition((topicPart, data) => {
+      assertEquals(reqData2Iter.next(), topicPart)
+    })
+    assertEquals(None, context2.getFetchOffset(new TopicPartition("foo", 0)))
+    assertEquals(10, context2.getFetchOffset(new TopicPartition("foo", 1)).get)
+    assertEquals(15, context2.getFetchOffset(new TopicPartition("bar", 0)).get)
+    assertEquals(None, context2.getFetchOffset(new TopicPartition("bar", 2)))
+    val respData2 = new util.LinkedHashMap[TopicPartition, FetchResponse.PartitionData[Records]]
+    respData2.put(new TopicPartition("foo", 1), new FetchResponse.PartitionData(
+      Errors.NONE, 10, 10, 10, null, null))
+    respData2.put(new TopicPartition("bar", 0), new FetchResponse.PartitionData(
+      Errors.NONE, 10, 10, 10, null, null))
+    val resp2 = context2.updateAndGenerateResponseData(respData2)
+    assertEquals(Errors.NONE, resp2.error)
+    assertEquals(1, resp2.responseData.size)
+    assertTrue(resp2.sessionId > 0)
+  }
+
+  @Test
+  def testFetchSessionExpiration(): Unit = {
+    val time = new MockTime()
+    // set maximum entries to 2 to allow for eviction later
+    val cache = new FetchSessionCache(2, 1000)
+    val fetchManager = new FetchManager(time, cache)
+
+    // Create a new fetch session, session 1
+    val session1req = new util.LinkedHashMap[TopicPartition, FetchRequest.PartitionData]
+    session1req.put(new TopicPartition("foo", 0), new FetchRequest.PartitionData(0, 0, 100,
+      Optional.empty()))
+    session1req.put(new TopicPartition("foo", 1), new FetchRequest.PartitionData(10, 0, 100,
+      Optional.empty()))
+    val session1context1 = fetchManager.newContext(JFetchMetadata.INITIAL, session1req, EMPTY_PART_LIST, false)
+    assertEquals(classOf[FullFetchContext], session1context1.getClass)
+    val respData1 = new util.LinkedHashMap[TopicPartition, FetchResponse.PartitionData[Records]]
+    respData1.put(new TopicPartition("foo", 0), new FetchResponse.PartitionData(
+      Errors.NONE, 100, 100, 100, null, null))
+    respData1.put(new TopicPartition("foo", 1), new FetchResponse.PartitionData(
+      Errors.NONE, 10, 10, 10, null, null))
+    val session1resp = session1context1.updateAndGenerateResponseData(respData1)
+    assertEquals(Errors.NONE, session1resp.error())
+    assertTrue(session1resp.sessionId() != INVALID_SESSION_ID)
+    assertEquals(2, session1resp.responseData().size())
+
+    // check session entered into case
+    assertTrue(cache.get(session1resp.sessionId()).isDefined)
+    time.sleep(500)
+
+    // Create a second new fetch session
+    val session2req = new util.LinkedHashMap[TopicPartition, FetchRequest.PartitionData]
+    session2req.put(new TopicPartition("foo", 0), new FetchRequest.PartitionData(0, 0, 100,
+      Optional.empty()))
+    session2req.put(new TopicPartition("foo", 1), new FetchRequest.PartitionData(10, 0, 100,
+      Optional.empty()))
+    val session2context = fetchManager.newContext(JFetchMetadata.INITIAL, session1req, EMPTY_PART_LIST, false)
+    assertEquals(classOf[FullFetchContext], session2context.getClass)
+    val session2RespData = new util.LinkedHashMap[TopicPartition, FetchResponse.PartitionData[Records]]
+    session2RespData.put(new TopicPartition("foo", 0), new FetchResponse.PartitionData(
+      Errors.NONE, 100, 100, 100, null, null))
+    session2RespData.put(new TopicPartition("foo", 1), new FetchResponse.PartitionData(
+      Errors.NONE, 10, 10, 10, null, null))
+    val session2resp = session2context.updateAndGenerateResponseData(respData1)
+    assertEquals(Errors.NONE, session2resp.error())
+    assertTrue(session2resp.sessionId() != INVALID_SESSION_ID)
+    assertEquals(2, session2resp.responseData().size())
+
+    // both newly created entries are present in cache
+    assertTrue(cache.get(session1resp.sessionId()).isDefined)
+    assertTrue(cache.get(session2resp.sessionId()).isDefined)
+    time.sleep(500)
+
+    // Create an incremental fetch request for session 1
+    val context1v2 = fetchManager.newContext(
+      new JFetchMetadata(session1resp.sessionId(), 1),
+      new util.LinkedHashMap[TopicPartition, FetchRequest.PartitionData],
+      new util.ArrayList[TopicPartition], false)
+    assertEquals(classOf[IncrementalFetchContext], context1v2.getClass)
+
+    // total sleep time will now be large enough that fetch session 1 will be evicted if not correctly touched
+    time.sleep(501)
+
+    // create one final session to test that the least recently used entry is evicted
+    // the second session should be evicted because the first session was incrementally fetched
+    // more recently than the second session was created
+    val session3req = new util.LinkedHashMap[TopicPartition, FetchRequest.PartitionData]
+    session3req.put(new TopicPartition("foo", 0), new FetchRequest.PartitionData(0, 0, 100,
+      Optional.empty()))
+    session3req.put(new TopicPartition("foo", 1), new FetchRequest.PartitionData(0, 0, 100,
+      Optional.empty()))
+    val session3context = fetchManager.newContext(JFetchMetadata.INITIAL, session3req, EMPTY_PART_LIST, false)
+    assertEquals(classOf[FullFetchContext], session3context.getClass)
+    val respData3 = new util.LinkedHashMap[TopicPartition, FetchResponse.PartitionData[Records]]
+    respData3.put(new TopicPartition("foo", 0), new FetchResponse.PartitionData(
+      Errors.NONE, 100, 100, 100, null, null))
+    respData3.put(new TopicPartition("foo", 1), new FetchResponse.PartitionData(
+      Errors.NONE, 10, 10, 10, null, null))
+    val session3resp = session3context.updateAndGenerateResponseData(respData3)
+    assertEquals(Errors.NONE, session3resp.error())
+    assertTrue(session3resp.sessionId() != INVALID_SESSION_ID)
+    assertEquals(2, session3resp.responseData().size())
+
+    assertTrue(cache.get(session1resp.sessionId()).isDefined)
+    assertFalse("session 2 should have been evicted by latest session, as session 1 was used more recently",
+      cache.get(session2resp.sessionId()).isDefined)
+    assertTrue(cache.get(session3resp.sessionId()).isDefined)
+  }
+
+  @Test
+  def testPrivilegedSessionHandling(): Unit = {
+    val time = new MockTime()
+    // set maximum entries to 2 to allow for eviction later
+    val cache = new FetchSessionCache(2, 1000)
+    val fetchManager = new FetchManager(time, cache)
+
+    // Create a new fetch session, session 1
+    val session1req = new util.LinkedHashMap[TopicPartition, FetchRequest.PartitionData]
+    session1req.put(new TopicPartition("foo", 0), new FetchRequest.PartitionData(0, 0, 100,
+      Optional.empty()))
+    session1req.put(new TopicPartition("foo", 1), new FetchRequest.PartitionData(10, 0, 100,
+      Optional.empty()))
+    val session1context = fetchManager.newContext(JFetchMetadata.INITIAL, session1req, EMPTY_PART_LIST, true)
+    assertEquals(classOf[FullFetchContext], session1context.getClass)
+    val respData1 = new util.LinkedHashMap[TopicPartition, FetchResponse.PartitionData[Records]]
+    respData1.put(new TopicPartition("foo", 0), new FetchResponse.PartitionData(
+      Errors.NONE, 100, 100, 100, null, null))
+    respData1.put(new TopicPartition("foo", 1), new FetchResponse.PartitionData(
+      Errors.NONE, 10, 10, 10, null, null))
+    val session1resp = session1context.updateAndGenerateResponseData(respData1)
+    assertEquals(Errors.NONE, session1resp.error())
+    assertTrue(session1resp.sessionId() != INVALID_SESSION_ID)
+    assertEquals(2, session1resp.responseData().size())
+    assertEquals(1, cache.size)
+
+    // move time forward to age session 1 a little compared to session 2
+    time.sleep(500)
+
+    // Create a second new fetch session, unprivileged
+    val session2req = new util.LinkedHashMap[TopicPartition, FetchRequest.PartitionData]
+    session2req.put(new TopicPartition("foo", 0), new FetchRequest.PartitionData(0, 0, 100,
+      Optional.empty()))
+    session2req.put(new TopicPartition("foo", 1), new FetchRequest.PartitionData(10, 0, 100,
+      Optional.empty()))
+    val session2context = fetchManager.newContext(JFetchMetadata.INITIAL, session1req, EMPTY_PART_LIST, false)
+    assertEquals(classOf[FullFetchContext], session2context.getClass)
+    val session2RespData = new util.LinkedHashMap[TopicPartition, FetchResponse.PartitionData[Records]]
+    session2RespData.put(new TopicPartition("foo", 0), new FetchResponse.PartitionData(
+      Errors.NONE, 100, 100, 100, null, null))
+    session2RespData.put(new TopicPartition("foo", 1), new FetchResponse.PartitionData(
+      Errors.NONE, 10, 10, 10, null, null))
+    val session2resp = session2context.updateAndGenerateResponseData(respData1)
+    assertEquals(Errors.NONE, session2resp.error())
+    assertTrue(session2resp.sessionId() != INVALID_SESSION_ID)
+    assertEquals(2, session2resp.responseData().size())
+
+    // both newly created entries are present in cache
+    assertTrue(cache.get(session1resp.sessionId()).isDefined)
+    assertTrue(cache.get(session2resp.sessionId()).isDefined)
+    assertEquals(2, cache.size)
+    time.sleep(500)
+
+    // create a session to test session1 privileges mean that session 1 is retained and session 2 is evicted
+    val session3req = new util.LinkedHashMap[TopicPartition, FetchRequest.PartitionData]
+    session3req.put(new TopicPartition("foo", 0), new FetchRequest.PartitionData(0, 0, 100,
+      Optional.empty()))
+    session3req.put(new TopicPartition("foo", 1), new FetchRequest.PartitionData(0, 0, 100,
+      Optional.empty()))
+    val session3context = fetchManager.newContext(JFetchMetadata.INITIAL, session3req, EMPTY_PART_LIST, true)
+    assertEquals(classOf[FullFetchContext], session3context.getClass)
+    val respData3 = new util.LinkedHashMap[TopicPartition, FetchResponse.PartitionData[Records]]
+    respData3.put(new TopicPartition("foo", 0), new FetchResponse.PartitionData(
+      Errors.NONE, 100, 100, 100, null, null))
+    respData3.put(new TopicPartition("foo", 1), new FetchResponse.PartitionData(
+      Errors.NONE, 10, 10, 10, null, null))
+    val session3resp = session3context.updateAndGenerateResponseData(respData3)
+    assertEquals(Errors.NONE, session3resp.error())
+    assertTrue(session3resp.sessionId() != INVALID_SESSION_ID)
+    assertEquals(2, session3resp.responseData().size())
+
+    assertTrue(cache.get(session1resp.sessionId()).isDefined)
+    // even though session 2 is more recent than session 1, and has not reached expiry time, it is less
+    // privileged than session 2, and thus session 3 should be entered and session 2 evicted.
+    assertFalse("session 2 should have been evicted by session 3",
+      cache.get(session2resp.sessionId()).isDefined)
+    assertTrue(cache.get(session3resp.sessionId()).isDefined)
+    assertEquals(2, cache.size)
+
+    time.sleep(501)
+
+    // create a final session to test whether session1 can be evicted due to age even though it is privileged
+    val session4req = new util.LinkedHashMap[TopicPartition, FetchRequest.PartitionData]
+    session4req.put(new TopicPartition("foo", 0), new FetchRequest.PartitionData(0, 0, 100,
+      Optional.empty()))
+    session4req.put(new TopicPartition("foo", 1), new FetchRequest.PartitionData(0, 0, 100,
+      Optional.empty()))
+    val session4context = fetchManager.newContext(JFetchMetadata.INITIAL, session4req, EMPTY_PART_LIST, true)
+    assertEquals(classOf[FullFetchContext], session4context.getClass)
+    val respData4 = new util.LinkedHashMap[TopicPartition, FetchResponse.PartitionData[Records]]
+    respData4.put(new TopicPartition("foo", 0), new FetchResponse.PartitionData(
+      Errors.NONE, 100, 100, 100, null, null))
+    respData4.put(new TopicPartition("foo", 1), new FetchResponse.PartitionData(
+      Errors.NONE, 10, 10, 10, null, null))
+    val session4resp = session3context.updateAndGenerateResponseData(respData4)
+    assertEquals(Errors.NONE, session4resp.error())
+    assertTrue(session4resp.sessionId() != INVALID_SESSION_ID)
+    assertEquals(2, session4resp.responseData().size())
+
+    assertFalse("session 1 should have been evicted by session 4 even though it is privileged as it has hit eviction time",
+      cache.get(session1resp.sessionId()).isDefined)
+    assertTrue(cache.get(session3resp.sessionId()).isDefined)
+    assertTrue(cache.get(session4resp.sessionId()).isDefined)
+    assertEquals(2, cache.size)
+  }
+
+  @Test
+  def testZeroSizeFetchSession(): Unit = {
+    val time = new MockTime()
+    val cache = new FetchSessionCache(10, 1000)
+    val fetchManager = new FetchManager(time, cache)
+
+    // Create a new fetch session with foo-0 and foo-1
+    val reqData1 = new util.LinkedHashMap[TopicPartition, FetchRequest.PartitionData]
+    reqData1.put(new TopicPartition("foo", 0), new FetchRequest.PartitionData(0, 0, 100,
+      Optional.empty()))
+    reqData1.put(new TopicPartition("foo", 1), new FetchRequest.PartitionData(10, 0, 100,
+      Optional.empty()))
+    val context1 = fetchManager.newContext(JFetchMetadata.INITIAL, reqData1, EMPTY_PART_LIST, false)
+    assertEquals(classOf[FullFetchContext], context1.getClass)
+    val respData1 = new util.LinkedHashMap[TopicPartition, FetchResponse.PartitionData[Records]]
+    respData1.put(new TopicPartition("foo", 0), new FetchResponse.PartitionData(
+      Errors.NONE, 100, 100, 100, null, null))
+    respData1.put(new TopicPartition("foo", 1), new FetchResponse.PartitionData(
+      Errors.NONE, 10, 10, 10, null, null))
+    val resp1 = context1.updateAndGenerateResponseData(respData1)
+    assertEquals(Errors.NONE, resp1.error)
+    assertTrue(resp1.sessionId() != INVALID_SESSION_ID)
+    assertEquals(2, resp1.responseData.size)
+
+    // Create an incremental fetch request that removes foo-0 and foo-1
+    // Verify that the previous fetch session was closed.
+    val reqData2 = new util.LinkedHashMap[TopicPartition, FetchRequest.PartitionData]
+    val removed2 = new util.ArrayList[TopicPartition]
+    removed2.add(new TopicPartition("foo", 0))
+    removed2.add(new TopicPartition("foo", 1))
+    val context2 = fetchManager.newContext(
+      new JFetchMetadata(resp1.sessionId, 1), reqData2, removed2, false)
+    assertEquals(classOf[SessionlessFetchContext], context2.getClass)
+    val respData2 = new util.LinkedHashMap[TopicPartition, FetchResponse.PartitionData[Records]]
+    val resp2 = context2.updateAndGenerateResponseData(respData2)
+    assertEquals(INVALID_SESSION_ID, resp2.sessionId)
+    assertTrue(resp2.responseData().isEmpty)
+    assertEquals(0, cache.size)
+  }
+
+  @Test
+  def testDivergingEpoch(): Unit = {
+    val time = new MockTime()
+    val cache = new FetchSessionCache(10, 1000)
+    val fetchManager = new FetchManager(time, cache)
+    val tp1 = new TopicPartition("foo", 1)
+    val tp2 = new TopicPartition("bar", 2)
+
+    val reqData = new util.LinkedHashMap[TopicPartition, FetchRequest.PartitionData]
+    reqData.put(tp1, new FetchRequest.PartitionData(100, 0, 1000, Optional.of(5), Optional.of(4)))
+    reqData.put(tp2, new FetchRequest.PartitionData(100, 0, 1000, Optional.of(5), Optional.of(4)))
+
+    // Full fetch context returns all partitions in the response
+    val context1 = fetchManager.newContext(JFetchMetadata.INITIAL, reqData, EMPTY_PART_LIST, isFollower = false)
+    assertEquals(classOf[FullFetchContext], context1.getClass)
+    val respData = new util.LinkedHashMap[TopicPartition, FetchResponse.PartitionData[Records]]
+    respData.put(tp1, new FetchResponse.PartitionData(Errors.NONE,
+      105, 105, 0, Optional.empty(), Collections.emptyList(), Optional.empty(), null))
+    val divergingEpoch = Optional.of(new FetchResponseData.EpochEndOffset().setEpoch(3).setEndOffset(90))
+    respData.put(tp2, new FetchResponse.PartitionData(Errors.NONE,
+      105, 105, 0, Optional.empty(), Collections.emptyList(), divergingEpoch, null))
+    val resp1 = context1.updateAndGenerateResponseData(respData)
+    assertEquals(Errors.NONE, resp1.error)
+    assertNotEquals(INVALID_SESSION_ID, resp1.sessionId)
+    assertEquals(Utils.mkSet(tp1, tp2), resp1.responseData.keySet)
+
+    // Incremental fetch context returns partitions with divergent epoch even if none
+    // of the other conditions for return are met.
+    val context2 = fetchManager.newContext(new JFetchMetadata(resp1.sessionId, 1), reqData, EMPTY_PART_LIST, isFollower = false)
+    assertEquals(classOf[IncrementalFetchContext], context2.getClass)
+    val resp2 = context2.updateAndGenerateResponseData(respData)
+    assertEquals(Errors.NONE, resp2.error)
+    assertEquals(resp1.sessionId, resp2.sessionId)
+    assertEquals(Collections.singleton(tp2), resp2.responseData.keySet)
+
+    // All partitions with divergent epoch should be returned.
+    respData.put(tp1, new FetchResponse.PartitionData(Errors.NONE,
+      105, 105, 0, Optional.empty(), Collections.emptyList(), divergingEpoch, null))
+    val resp3 = context2.updateAndGenerateResponseData(respData)
+    assertEquals(Errors.NONE, resp3.error)
+    assertEquals(resp1.sessionId, resp3.sessionId)
+    assertEquals(Utils.mkSet(tp1, tp2), resp3.responseData.keySet)
+
+    // Partitions that meet other conditions should be returned regardless of whether
+    // divergingEpoch is set or not.
+    respData.put(tp1, new FetchResponse.PartitionData(Errors.NONE,
+      110, 110, 0, Optional.empty(), Collections.emptyList(), Optional.empty(), null))
+    val resp4 = context2.updateAndGenerateResponseData(respData)
+    assertEquals(Errors.NONE, resp4.error)
+    assertEquals(resp1.sessionId, resp4.sessionId)
+    assertEquals(Utils.mkSet(tp1, tp2), resp4.responseData.keySet)
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/server/FinalizedFeatureCacheTest.scala b/core/src/test/scala/unit/kafka/server/FinalizedFeatureCacheTest.scala
new file mode 100644
index 0000000000000..2f0f70ee0bad3
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/FinalizedFeatureCacheTest.scala
@@ -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 kafka.server
+
+import org.apache.kafka.common.feature.{Features, FinalizedVersionRange, SupportedVersionRange}
+import org.junit.Assert.{assertEquals, assertThrows, assertTrue}
+import org.junit.Test
+
+import scala.jdk.CollectionConverters._
+
+class FinalizedFeatureCacheTest {
+
+  @Test
+  def testEmpty(): Unit = {
+    assertTrue(new FinalizedFeatureCache(BrokerFeatures.createDefault()).get.isEmpty)
+  }
+
+  @Test
+  def testUpdateOrThrowFailedDueToInvalidEpoch(): Unit = {
+    val supportedFeatures = Map[String, SupportedVersionRange](
+      "feature_1" -> new SupportedVersionRange(1, 4))
+    val brokerFeatures = BrokerFeatures.createDefault()
+    brokerFeatures.setSupportedFeatures(Features.supportedFeatures(supportedFeatures.asJava))
+
+    val features = Map[String, FinalizedVersionRange](
+      "feature_1" -> new FinalizedVersionRange(1, 4))
+    val finalizedFeatures = Features.finalizedFeatures(features.asJava)
+
+    val cache = new FinalizedFeatureCache(brokerFeatures)
+    cache.updateOrThrow(finalizedFeatures, 10)
+    assertTrue(cache.get.isDefined)
+    assertEquals(finalizedFeatures, cache.get.get.features)
+    assertEquals(10, cache.get.get.epoch)
+
+    assertThrows(
+      classOf[FeatureCacheUpdateException],
+      () => cache.updateOrThrow(finalizedFeatures, 9))
+
+    // Check that the failed updateOrThrow call did not make any mutations.
+    assertTrue(cache.get.isDefined)
+    assertEquals(finalizedFeatures, cache.get.get.features)
+    assertEquals(10, cache.get.get.epoch)
+  }
+
+  @Test
+  def testUpdateOrThrowFailedDueToInvalidFeatures(): Unit = {
+    val supportedFeatures =
+      Map[String, SupportedVersionRange]("feature_1" -> new SupportedVersionRange(1, 1))
+    val brokerFeatures = BrokerFeatures.createDefault()
+    brokerFeatures.setSupportedFeatures(Features.supportedFeatures(supportedFeatures.asJava))
+
+    val features = Map[String, FinalizedVersionRange](
+      "feature_1" -> new FinalizedVersionRange(1, 2))
+    val finalizedFeatures = Features.finalizedFeatures(features.asJava)
+
+    val cache = new FinalizedFeatureCache(brokerFeatures)
+    assertThrows(
+      classOf[FeatureCacheUpdateException],
+      () => cache.updateOrThrow(finalizedFeatures, 12))
+
+    // Check that the failed updateOrThrow call did not make any mutations.
+    assertTrue(cache.isEmpty)
+  }
+
+  @Test
+  def testUpdateOrThrowSuccess(): Unit = {
+    val supportedFeatures =
+      Map[String, SupportedVersionRange]("feature_1" -> new SupportedVersionRange(1, 4))
+    val brokerFeatures = BrokerFeatures.createDefault()
+    brokerFeatures.setSupportedFeatures(Features.supportedFeatures(supportedFeatures.asJava))
+
+    val features = Map[String, FinalizedVersionRange](
+      "feature_1" -> new FinalizedVersionRange(2, 3))
+    val finalizedFeatures = Features.finalizedFeatures(features.asJava)
+
+    val cache = new FinalizedFeatureCache(brokerFeatures)
+    cache.updateOrThrow(finalizedFeatures, 12)
+    assertTrue(cache.get.isDefined)
+    assertEquals(finalizedFeatures,  cache.get.get.features)
+    assertEquals(12, cache.get.get.epoch)
+  }
+
+  @Test
+  def testClear(): Unit = {
+    val supportedFeatures =
+      Map[String, SupportedVersionRange]("feature_1" -> new SupportedVersionRange(1, 4))
+    val brokerFeatures = BrokerFeatures.createDefault()
+    brokerFeatures.setSupportedFeatures(Features.supportedFeatures(supportedFeatures.asJava))
+
+    val features = Map[String, FinalizedVersionRange](
+      "feature_1" -> new FinalizedVersionRange(2, 3))
+    val finalizedFeatures = Features.finalizedFeatures(features.asJava)
+
+    val cache = new FinalizedFeatureCache(brokerFeatures)
+    cache.updateOrThrow(finalizedFeatures, 12)
+    assertTrue(cache.get.isDefined)
+    assertEquals(finalizedFeatures, cache.get.get.features)
+    assertEquals(12, cache.get.get.epoch)
+
+    cache.clear()
+    assertTrue(cache.isEmpty)
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/server/FinalizedFeatureChangeListenerTest.scala b/core/src/test/scala/unit/kafka/server/FinalizedFeatureChangeListenerTest.scala
new file mode 100644
index 0000000000000..cb8a661317794
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/FinalizedFeatureChangeListenerTest.scala
@@ -0,0 +1,273 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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.{CountDownLatch, TimeoutException}
+
+import kafka.zk.{FeatureZNode, FeatureZNodeStatus, ZkVersion, ZooKeeperTestHarness}
+import kafka.utils.TestUtils
+import org.apache.kafka.common.utils.Exit
+import org.apache.kafka.common.feature.{Features, FinalizedVersionRange, SupportedVersionRange}
+import org.apache.kafka.test.{TestUtils => JTestUtils}
+import org.junit.Assert.{assertEquals, assertFalse, assertNotEquals, assertThrows, assertTrue}
+import org.junit.Test
+
+import scala.jdk.CollectionConverters._
+
+class FinalizedFeatureChangeListenerTest extends ZooKeeperTestHarness {
+
+  private def createBrokerFeatures(): BrokerFeatures = {
+    val supportedFeaturesMap = Map[String, SupportedVersionRange](
+      "feature_1" -> new SupportedVersionRange(1, 4),
+      "feature_2" -> new SupportedVersionRange(1, 3))
+    val brokerFeatures = BrokerFeatures.createDefault()
+    brokerFeatures.setSupportedFeatures(Features.supportedFeatures(supportedFeaturesMap.asJava))
+    brokerFeatures
+  }
+
+  private def createFinalizedFeatures(): FinalizedFeaturesAndEpoch = {
+    val finalizedFeaturesMap = Map[String, FinalizedVersionRange](
+      "feature_1" -> new FinalizedVersionRange(2, 3))
+    val finalizedFeatures = Features.finalizedFeatures(finalizedFeaturesMap.asJava)
+    zkClient.createFeatureZNode(FeatureZNode(FeatureZNodeStatus.Enabled, finalizedFeatures))
+    val (mayBeFeatureZNodeBytes, version) = zkClient.getDataAndVersion(FeatureZNode.path)
+    assertNotEquals(version, ZkVersion.UnknownVersion)
+    assertFalse(mayBeFeatureZNodeBytes.isEmpty)
+    FinalizedFeaturesAndEpoch(finalizedFeatures, version)
+  }
+
+  private def createListener(
+    cache: FinalizedFeatureCache,
+    expectedCacheContent: Option[FinalizedFeaturesAndEpoch]
+  ): FinalizedFeatureChangeListener = {
+    val listener = new FinalizedFeatureChangeListener(cache, zkClient)
+    assertFalse(listener.isListenerInitiated)
+    assertTrue(cache.isEmpty)
+    listener.initOrThrow(15000)
+    assertTrue(listener.isListenerInitiated)
+    if (expectedCacheContent.isDefined) {
+      val mayBeNewCacheContent = cache.get
+      assertFalse(mayBeNewCacheContent.isEmpty)
+      val newCacheContent = mayBeNewCacheContent.get
+      assertEquals(expectedCacheContent.get.features, newCacheContent.features)
+      assertEquals(expectedCacheContent.get.epoch, newCacheContent.epoch)
+    } else {
+      val mayBeNewCacheContent = cache.get
+      assertTrue(mayBeNewCacheContent.isEmpty)
+    }
+    listener
+  }
+
+  /**
+   * Tests that the listener can be initialized, and that it can listen to ZK notifications
+   * successfully from an "Enabled" FeatureZNode (the ZK data has no feature incompatibilities).
+   * Particularly the test checks if multiple notifications can be processed in ZK
+   * (i.e. whether the FeatureZNode watch can be re-established).
+   */
+  @Test
+  def testInitSuccessAndNotificationSuccess(): Unit = {
+    val initialFinalizedFeatures = createFinalizedFeatures()
+    val brokerFeatures = createBrokerFeatures()
+    val cache = new FinalizedFeatureCache(brokerFeatures)
+    val listener = createListener(cache, Some(initialFinalizedFeatures))
+
+    def updateAndCheckCache(finalizedFeatures: Features[FinalizedVersionRange]): Unit = {
+      zkClient.updateFeatureZNode(FeatureZNode(FeatureZNodeStatus.Enabled, finalizedFeatures))
+      val (mayBeFeatureZNodeNewBytes, updatedVersion) = zkClient.getDataAndVersion(FeatureZNode.path)
+      assertNotEquals(updatedVersion, ZkVersion.UnknownVersion)
+      assertFalse(mayBeFeatureZNodeNewBytes.isEmpty)
+      assertTrue(updatedVersion > initialFinalizedFeatures.epoch)
+
+      cache.waitUntilEpochOrThrow(updatedVersion, JTestUtils.DEFAULT_MAX_WAIT_MS)
+      assertEquals(FinalizedFeaturesAndEpoch(finalizedFeatures, updatedVersion), cache.get.get)
+      assertTrue(listener.isListenerInitiated)
+    }
+
+    // Check if the write succeeds and a ZK notification is received that causes the feature cache
+    // to be populated.
+    updateAndCheckCache(
+      Features.finalizedFeatures(
+        Map[String, FinalizedVersionRange](
+        "feature_1" -> new FinalizedVersionRange(2, 4)).asJava))
+    // Check if second write succeeds and a ZK notification is again received that causes the cache
+    // to be populated. This check is needed to verify that the watch on the FeatureZNode was
+    // re-established after the notification was received due to the first write above.
+    updateAndCheckCache(
+      Features.finalizedFeatures(
+        Map[String, FinalizedVersionRange](
+          "feature_1" -> new FinalizedVersionRange(2, 4),
+          "feature_2" -> new FinalizedVersionRange(1, 3)).asJava))
+  }
+
+  /**
+   * Tests that the listener can be initialized, and that it can process FeatureZNode deletion
+   * successfully.
+   */
+  @Test
+  def testFeatureZNodeDeleteNotificationProcessing(): Unit = {
+    val brokerFeatures = createBrokerFeatures()
+    val cache = new FinalizedFeatureCache(brokerFeatures)
+    val initialFinalizedFeatures = createFinalizedFeatures()
+    val listener = createListener(cache, Some(initialFinalizedFeatures))
+
+    zkClient.deleteFeatureZNode()
+    val (mayBeFeatureZNodeDeletedBytes, deletedVersion) = zkClient.getDataAndVersion(FeatureZNode.path)
+    assertEquals(deletedVersion, ZkVersion.UnknownVersion)
+    assertTrue(mayBeFeatureZNodeDeletedBytes.isEmpty)
+    TestUtils.waitUntilTrue(() => {
+      cache.isEmpty
+    }, "Timed out waiting for FinalizedFeatureCache to become empty")
+    assertTrue(listener.isListenerInitiated)
+  }
+
+  /**
+   * Tests that the listener can be initialized, and that it can process disabling of a FeatureZNode
+   * successfully.
+   */
+  @Test
+  def testFeatureZNodeDisablingNotificationProcessing(): Unit = {
+    val brokerFeatures = createBrokerFeatures()
+    val cache = new FinalizedFeatureCache(brokerFeatures)
+    val initialFinalizedFeatures = createFinalizedFeatures()
+
+    val updatedFinalizedFeaturesMap = Map[String, FinalizedVersionRange]()
+    val updatedFinalizedFeatures = Features.finalizedFeatures(updatedFinalizedFeaturesMap.asJava)
+    zkClient.updateFeatureZNode(FeatureZNode(FeatureZNodeStatus.Disabled, updatedFinalizedFeatures))
+    val (mayBeFeatureZNodeNewBytes, updatedVersion) = zkClient.getDataAndVersion(FeatureZNode.path)
+    assertNotEquals(updatedVersion, ZkVersion.UnknownVersion)
+    assertFalse(mayBeFeatureZNodeNewBytes.isEmpty)
+    assertTrue(updatedVersion > initialFinalizedFeatures.epoch)
+    assertTrue(cache.get.isEmpty)
+  }
+
+  /**
+   * Tests that the wait operation on the cache fails (as expected) when an epoch can never be
+   * reached. Also tests that the wait operation on the cache succeeds when an epoch is expected to
+   * be reached.
+   */
+  @Test
+  def testCacheUpdateWaitFailsForUnreachableVersion(): Unit = {
+    val initialFinalizedFeatures = createFinalizedFeatures()
+    val cache = new FinalizedFeatureCache(createBrokerFeatures())
+    val listener = createListener(cache, Some(initialFinalizedFeatures))
+
+    assertThrows(
+      classOf[TimeoutException],
+      () => cache.waitUntilEpochOrThrow(initialFinalizedFeatures.epoch + 1, JTestUtils.DEFAULT_MAX_WAIT_MS))
+
+    val updatedFinalizedFeaturesMap = Map[String, FinalizedVersionRange]()
+    val updatedFinalizedFeatures = Features.finalizedFeatures(updatedFinalizedFeaturesMap.asJava)
+    zkClient.updateFeatureZNode(FeatureZNode(FeatureZNodeStatus.Disabled, updatedFinalizedFeatures))
+    val (mayBeFeatureZNodeNewBytes, updatedVersion) = zkClient.getDataAndVersion(FeatureZNode.path)
+    assertNotEquals(updatedVersion, ZkVersion.UnknownVersion)
+    assertFalse(mayBeFeatureZNodeNewBytes.isEmpty)
+    assertTrue(updatedVersion > initialFinalizedFeatures.epoch)
+
+    assertThrows(
+      classOf[TimeoutException],
+      () => cache.waitUntilEpochOrThrow(updatedVersion, JTestUtils.DEFAULT_MAX_WAIT_MS))
+    assertTrue(cache.get.isEmpty)
+    assertTrue(listener.isListenerInitiated)
+  }
+
+  /**
+   * Tests that the listener initialization fails when it picks up a feature incompatibility from
+   * ZK from an "Enabled" FeatureZNode.
+   */
+  @Test
+  def testInitFailureDueToFeatureIncompatibility(): Unit = {
+    val brokerFeatures = createBrokerFeatures()
+    val cache = new FinalizedFeatureCache(brokerFeatures)
+
+    val incompatibleFinalizedFeaturesMap = Map[String, FinalizedVersionRange](
+      "feature_1" -> new FinalizedVersionRange(2, 5))
+    val incompatibleFinalizedFeatures = Features.finalizedFeatures(incompatibleFinalizedFeaturesMap.asJava)
+    zkClient.createFeatureZNode(FeatureZNode(FeatureZNodeStatus.Enabled, incompatibleFinalizedFeatures))
+    val (mayBeFeatureZNodeBytes, initialVersion) = zkClient.getDataAndVersion(FeatureZNode.path)
+    assertNotEquals(initialVersion, ZkVersion.UnknownVersion)
+    assertFalse(mayBeFeatureZNodeBytes.isEmpty)
+
+    val exitLatch = new CountDownLatch(1)
+    Exit.setExitProcedure((_, _) => exitLatch.countDown())
+    try {
+      val listener = new FinalizedFeatureChangeListener(cache, zkClient)
+      assertFalse(listener.isListenerInitiated)
+      assertTrue(cache.isEmpty)
+      assertThrows(classOf[TimeoutException], () => listener.initOrThrow(5000))
+      exitLatch.await()
+      assertFalse(listener.isListenerInitiated)
+      assertTrue(listener.isListenerDead)
+      assertTrue(cache.isEmpty)
+    } finally {
+      Exit.resetExitProcedure()
+    }
+  }
+
+  /**
+   * Tests that the listener initialization fails when invalid wait time (<= 0) is provided as input.
+   */
+  @Test
+  def testInitFailureDueToInvalidWaitTime(): Unit = {
+    val brokerFeatures = createBrokerFeatures()
+    val cache = new FinalizedFeatureCache(brokerFeatures)
+    val listener = new FinalizedFeatureChangeListener(cache, zkClient)
+    assertThrows(classOf[IllegalArgumentException], () => listener.initOrThrow(0))
+    assertThrows(classOf[IllegalArgumentException], () => listener.initOrThrow(-1))
+  }
+
+  /**
+   * Tests that after successful initialization, the listener fails when it picks up a feature
+   * incompatibility from ZK.
+   */
+  @Test
+  def testNotificationFailureDueToFeatureIncompatibility(): Unit = {
+    val brokerFeatures = createBrokerFeatures()
+    val cache = new FinalizedFeatureCache(brokerFeatures)
+    val initialFinalizedFeatures = createFinalizedFeatures()
+    val listener = createListener(cache, Some(initialFinalizedFeatures))
+
+    val exitLatch = new CountDownLatch(1)
+    Exit.setExitProcedure((_, _) => exitLatch.countDown())
+    val incompatibleFinalizedFeaturesMap = Map[String, FinalizedVersionRange](
+      "feature_1" -> new FinalizedVersionRange(
+        brokerFeatures.supportedFeatures.get("feature_1").min(),
+        (brokerFeatures.supportedFeatures.get("feature_1").max() + 1).asInstanceOf[Short]))
+    val incompatibleFinalizedFeatures = Features.finalizedFeatures(incompatibleFinalizedFeaturesMap.asJava)
+    zkClient.updateFeatureZNode(FeatureZNode(FeatureZNodeStatus.Enabled, incompatibleFinalizedFeatures))
+    val (mayBeFeatureZNodeIncompatibleBytes, updatedVersion) = zkClient.getDataAndVersion(FeatureZNode.path)
+    assertNotEquals(updatedVersion, ZkVersion.UnknownVersion)
+    assertFalse(mayBeFeatureZNodeIncompatibleBytes.isEmpty)
+
+    try {
+      TestUtils.waitUntilTrue(() => {
+        // Make sure the custom exit procedure (defined above) was called.
+        exitLatch.getCount == 0 &&
+        // Make sure the listener is no longer initiated (because, it is dead).
+        !listener.isListenerInitiated &&
+        // Make sure the listener dies after hitting an exception when processing incompatible
+        // features read from ZK.
+        listener.isListenerDead &&
+        // Make sure the cache contents are as expected, and, the incompatible features were not
+        // applied.
+        cache.get.get.equals(initialFinalizedFeatures)
+      }, "Timed out waiting for listener death and FinalizedFeatureCache to be updated")
+    } finally {
+      Exit.resetExitProcedure()
+    }
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/server/ForwardingManagerTest.scala b/core/src/test/scala/unit/kafka/server/ForwardingManagerTest.scala
new file mode 100644
index 0000000000000..5001f671c02f8
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/ForwardingManagerTest.scala
@@ -0,0 +1,157 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package kafka.server
+
+import java.net.InetAddress
+import java.nio.ByteBuffer
+import java.util.Optional
+import kafka.network
+import kafka.network.RequestChannel
+import kafka.utils.MockTime
+import org.apache.kafka.clients.{ClientResponse, RequestCompletionHandler}
+import org.apache.kafka.common.config.{ConfigResource, TopicConfig}
+import org.apache.kafka.common.memory.MemoryPool
+import org.apache.kafka.common.message.AlterConfigsResponseData
+import org.apache.kafka.common.network.{ClientInformation, ListenerName}
+import org.apache.kafka.common.protocol.{ApiKeys, Errors}
+import org.apache.kafka.common.requests.{AbstractRequest, AbstractResponse, AlterConfigsRequest, AlterConfigsResponse, EnvelopeRequest, EnvelopeResponse, RequestContext, RequestHeader, RequestTestUtils}
+import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
+import org.apache.kafka.common.security.authenticator.DefaultKafkaPrincipalBuilder
+import org.junit.Assert._
+import org.junit.Test
+import org.mockito.ArgumentMatchers._
+import org.mockito.{ArgumentMatchers, Mockito}
+
+import scala.jdk.CollectionConverters._
+
+class ForwardingManagerTest {
+  private val brokerToController = Mockito.mock(classOf[BrokerToControllerChannelManager])
+  private val time = new MockTime()
+  private val principalBuilder = new DefaultKafkaPrincipalBuilder(null, null)
+
+  @Test
+  def testResponseCorrelationIdMismatch(): Unit = {
+    val forwardingManager = new ForwardingManager(brokerToController, time, Long.MaxValue)
+    val requestCorrelationId = 27
+    val envelopeCorrelationId = 39
+    val clientPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "client")
+
+    val configResource = new ConfigResource(ConfigResource.Type.TOPIC, "foo")
+    val configs = List(new AlterConfigsRequest.ConfigEntry(TopicConfig.MIN_IN_SYNC_REPLICAS_CONFIG, "1")).asJava
+    val requestBody = new AlterConfigsRequest.Builder(Map(
+      configResource -> new AlterConfigsRequest.Config(configs)
+    ).asJava, false).build()
+    val (requestHeader, requestBuffer) = buildRequest(requestBody, requestCorrelationId)
+    val request = buildRequest(requestHeader, requestBuffer, clientPrincipal)
+
+    val responseBody = new AlterConfigsResponse(new AlterConfigsResponseData())
+    val responseBuffer = RequestTestUtils.serializeResponseWithHeader(responseBody, requestHeader.apiVersion,
+      requestCorrelationId + 1)
+
+    Mockito.when(brokerToController.sendRequest(
+      any(classOf[EnvelopeRequest.Builder]),
+      any(classOf[ControllerRequestCompletionHandler]),
+      ArgumentMatchers.eq(Long.MaxValue)
+    )).thenAnswer(invocation => {
+      val completionHandler = invocation.getArgument[RequestCompletionHandler](1)
+      val response = buildEnvelopeResponse(responseBuffer, envelopeCorrelationId, completionHandler)
+      response.onComplete()
+    })
+
+    var response: AbstractResponse = null
+    forwardingManager.forwardRequest(request, res => response = res)
+
+    assertNotNull(response)
+    assertEquals(Map(Errors.UNKNOWN_SERVER_ERROR -> 1).asJava, response.errorCounts())
+  }
+
+  private def buildEnvelopeResponse(
+    responseBuffer: ByteBuffer,
+    correlationId: Int,
+    completionHandler: RequestCompletionHandler
+  ): ClientResponse = {
+    val envelopeRequestHeader = new RequestHeader(
+      ApiKeys.ENVELOPE,
+      ApiKeys.ENVELOPE.latestVersion(),
+      "clientId",
+      correlationId
+    )
+    val envelopeResponse = new EnvelopeResponse(
+      responseBuffer,
+      Errors.NONE
+    )
+
+    new ClientResponse(
+      envelopeRequestHeader,
+      completionHandler,
+      "1",
+      time.milliseconds(),
+      time.milliseconds(),
+      false,
+      null,
+      null,
+      envelopeResponse
+    )
+  }
+
+  private def buildRequest(
+    body: AbstractRequest,
+    correlationId: Int
+  ): (RequestHeader, ByteBuffer) = {
+    val header = new RequestHeader(
+      body.apiKey,
+      body.version,
+      "clientId",
+      correlationId
+    )
+    val buffer = RequestTestUtils.serializeRequestWithHeader(header, body)
+
+    // Fast-forward buffer to start of the request as `RequestChannel.Request` expects
+    RequestHeader.parse(buffer)
+
+    (header, buffer)
+  }
+
+  private def buildRequest(
+    requestHeader: RequestHeader,
+    requestBuffer: ByteBuffer,
+    principal: KafkaPrincipal
+  ): RequestChannel.Request = {
+    val requestContext = new RequestContext(
+      requestHeader,
+      "1",
+      InetAddress.getLocalHost,
+      principal,
+      new ListenerName("client"),
+      SecurityProtocol.SASL_PLAINTEXT,
+      ClientInformation.EMPTY,
+      false,
+      Optional.of(principalBuilder)
+    )
+
+    new network.RequestChannel.Request(
+      processor = 1,
+      context = requestContext,
+      startTimeNanos = time.nanoseconds(),
+      memoryPool = MemoryPool.NONE,
+      buffer = requestBuffer,
+      metrics = new RequestChannel.Metrics(allowDisabledApis = true),
+      envelope = None
+    )
+  }
+
+}
diff --git a/core/src/test/scala/unit/kafka/server/HighwatermarkPersistenceTest.scala b/core/src/test/scala/unit/kafka/server/HighwatermarkPersistenceTest.scala
index c0871a769d6d7..95c76710c6dc3 100755
--- a/core/src/test/scala/unit/kafka/server/HighwatermarkPersistenceTest.scala
+++ b/core/src/test/scala/unit/kafka/server/HighwatermarkPersistenceTest.scala
@@ -24,18 +24,19 @@ 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 {
 
   val configs = TestUtils.createBrokerConfigs(2, TestUtils.MockZkConnect).map(KafkaConfig.fromProps)
   val topic = "foo"
-  val zkClient = EasyMock.createMock(classOf[KafkaZkClient])
+  val zkClient: KafkaZkClient = EasyMock.createMock(classOf[KafkaZkClient])
   val logManagers = configs map { config =>
     TestUtils.createLogManager(
       logDirs = config.logDirs.map(new File(_)),
@@ -46,109 +47,120 @@ class HighwatermarkPersistenceTest {
     new LogDirFailureChannel(config.logDirs.size)
   }
 
+  val alterIsrManager = TestUtils.createAlterIsrManager()
+
   @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)
 
     // create kafka scheduler
     val scheduler = new KafkaScheduler(2)
-    scheduler.startup
+    scheduler.startup()
     val metrics = new Metrics
     val time = new MockTime
+    val quotaManager = QuotaFactory.instantiate(configs.head, metrics, time, "")
     // create replica manager
     val replicaManager = new ReplicaManager(configs.head, metrics, time, zkClient, scheduler,
-      logManagers.head, new AtomicBoolean(false), QuotaFactory.instantiate(configs.head, metrics, time, ""),
-      new BrokerTopicStats, new MetadataCache(configs.head.brokerId), logDirFailureChannels.head)
+      logManagers.head, new AtomicBoolean(false), quotaManager,
+      new BrokerTopicStats, new MetadataCache(configs.head.brokerId), logDirFailureChannels.head, alterIsrManager)
     replicaManager.startup()
     try {
       replicaManager.checkpointHighWatermarks()
       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),
+        addingReplicas = Seq.empty,
+        removingReplicas = Seq.empty
+      )
+
       replicaManager.checkpointHighWatermarks()
       fooPartition0Hw = hwmFor(replicaManager, topic, 0)
-      assertEquals(leaderReplicaPartition0.highWatermark.messageOffset, fooPartition0Hw)
+      assertEquals(log0.highWatermark, fooPartition0Hw)
       // set the high watermark for local replica
-      partition0.getReplica().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
       replicaManager.shutdown(false)
+      quotaManager.shutdown()
       metrics.close()
       scheduler.shutdown()
     }
   }
 
   @Test
-  def testHighWatermarkPersistenceMultiplePartitions() {
+  def testHighWatermarkPersistenceMultiplePartitions(): Unit = {
     val topic1 = "foo1"
     val topic2 = "foo2"
     // mock zkclient
     EasyMock.replay(zkClient)
     // create kafka scheduler
     val scheduler = new KafkaScheduler(2)
-    scheduler.startup
+    scheduler.startup()
     val metrics = new Metrics
     val time = new MockTime
+    val quotaManager = QuotaFactory.instantiate(configs.head, metrics, time, "")
     // create replica manager
     val replicaManager = new ReplicaManager(configs.head, metrics, time, zkClient,
-      scheduler, logManagers.head, new AtomicBoolean(false), QuotaFactory.instantiate(configs.head, metrics, time, ""),
-      new BrokerTopicStats, new MetadataCache(configs.head.brokerId), logDirFailureChannels.head)
+      scheduler, logManagers.head, new AtomicBoolean(false), quotaManager,
+      new BrokerTopicStats, new MetadataCache(configs.head.brokerId), logDirFailureChannels.head, alterIsrManager)
     replicaManager.startup()
     try {
       replicaManager.checkpointHighWatermarks()
       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.getReplica().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.getReplica().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.getReplica().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)
@@ -160,14 +172,19 @@ class HighwatermarkPersistenceTest {
     } finally {
       // shutdown the replica manager upon test completion
       replicaManager.shutdown(false)
+      quotaManager.shutdown()
       metrics.close()
       scheduler.shutdown()
     }
   }
 
-  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)
+  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 8212ed680c5ba..0000000000000
--- a/core/src/test/scala/unit/kafka/server/ISRExpirationTest.scala
+++ /dev/null
@@ -1,234 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
-*/
-package kafka.server
-
-import java.io.File
-import java.util.Properties
-import java.util.concurrent.atomic.AtomicBoolean
-
-import kafka.cluster.{Partition, Replica}
-import kafka.log.Log
-import kafka.server.epoch.LeaderEpochCache
-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 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 = EasyMock.createMock(classOf[kafka.log.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 (15)
-    for (replica <- partition0.assignedReplicas - leaderReplica)
-      replica.updateLogReadResult(new LogReadResult(info = FetchDataInfo(new LogOffsetMetadata(15L), MemoryRecords.EMPTY),
-                                                    highWatermark = 15L,
-                                                    leaderLogStartOffset = 0L,
-                                                    leaderLogEndOffset = 15L,
-                                                    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(10L), MemoryRecords.EMPTY),
-                                                    highWatermark = 10L,
-                                                    leaderLogStartOffset = 0L,
-                                                    leaderLogEndOffset = 15L,
-                                                    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(11L), MemoryRecords.EMPTY),
-                            highWatermark = 11L,
-                            leaderLogStartOffset = 0L,
-                            leaderLogEndOffset = 15L,
-                            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(15L), MemoryRecords.EMPTY),
-                            highWatermark = 15L,
-                            leaderLogStartOffset = 0L,
-                            leaderLogEndOffset = 15L,
-                            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)
-  }
-
-  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 = EasyMock.createMock(classOf[kafka.log.Log])
-    val cache = EasyMock.createNiceMock(classOf[LeaderEpochCache])
-    EasyMock.expect(log.dir).andReturn(TestUtils.tempDir()).anyTimes()
-    EasyMock.expect(log.leaderEpochCache).andReturn(cache).anyTimes()
-    EasyMock.expect(log.onHighWatermarkIncremented(0L))
-    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..7fc4903e34460
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/IsrExpirationTest.scala
@@ -0,0 +1,252 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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.server.QuotaFactory.QuotaManagers
+import kafka.utils.TestUtils.MockAlterIsrManager
+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 quotaManager: QuotaManagers = null
+  var replicaManager: ReplicaManager = null
+
+  var alterIsrManager: MockAlterIsrManager = _
+
+  @Before
+  def setUp(): Unit = {
+    val logManager: LogManager = EasyMock.createMock(classOf[LogManager])
+    EasyMock.expect(logManager.liveLogDirs).andReturn(Array.empty[File]).anyTimes()
+    EasyMock.replay(logManager)
+
+    alterIsrManager = TestUtils.createAlterIsrManager()
+    quotaManager = QuotaFactory.instantiate(configs.head, metrics, time, "")
+    replicaManager = new ReplicaManager(configs.head, metrics, time, null, null, logManager, new AtomicBoolean(false),
+      quotaManager, new BrokerTopicStats, new MetadataCache(configs.head.brokerId),
+      new LogDirFailureChannel(configs.head.logDirs.size), alterIsrManager)
+  }
+
+  @After
+  def tearDown(): Unit = {
+    Option(replicaManager).foreach(_.shutdown(false))
+    Option(quotaManager).foreach(_.shutdown())
+    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)
+    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)
+
+    // 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)
+    }
+    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)
+    }
+    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)
+
+    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,
+      addingReplicas = Seq.empty,
+      removingReplicas = Seq.empty
+    )
+
+    // set lastCaughtUpTime to current time
+    for (replica <- partition.remoteReplicas)
+      replica.updateFetchState(
+        followerFetchOffsetMetadata = LogOffsetMetadata(0L),
+        followerStartOffset = 0L,
+        followerFetchTimeMs= time.milliseconds,
+        leaderEndOffset = 0L)
+
+    // 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 a51acd04f2730..04e42aa616870 100644
--- a/core/src/test/scala/unit/kafka/server/KafkaApisTest.scala
+++ b/core/src/test/scala/unit/kafka/server/KafkaApisTest.scala
@@ -17,79 +17,120 @@
 
 package kafka.server
 
-import java.lang.{Long => JLong}
 import java.net.InetAddress
+import java.nio.charset.StandardCharsets
 import java.util
+import java.util.Arrays.asList
+import java.util.concurrent.TimeUnit
+import java.util.{Collections, Optional, Properties, Random}
 
-import kafka.api.{ApiVersion, KAFKA_0_10_2_IV0}
-import kafka.cluster.Replica
+import kafka.api.{ApiVersion, KAFKA_0_10_2_IV0, KAFKA_2_2_IV1, LeaderAndIsr}
+import kafka.cluster.{Broker, Partition}
 import kafka.controller.KafkaController
-import kafka.coordinator.group.GroupCoordinator
-import kafka.coordinator.transaction.TransactionCoordinator
-import kafka.log.{Log, TimestampOffset}
+import kafka.coordinator.group.GroupCoordinatorConcurrencyTest.{JoinGroupCallback, SyncGroupCallback}
+import kafka.coordinator.group._
+import kafka.coordinator.transaction.{InitProducerIdResult, TransactionCoordinator}
+import kafka.log.AppendOrigin
 import kafka.network.RequestChannel
-import kafka.security.auth.Authorizer
+import kafka.network.RequestChannel.{CloseConnectionResponse, SendResponse}
 import kafka.server.QuotaFactory.QuotaManagers
-import kafka.utils.{MockTime, TestUtils, ZkUtils}
+import kafka.utils.{MockTime, TestUtils}
 import kafka.zk.KafkaZkClient
-import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.clients.admin.AlterConfigOp.OpType
+import org.apache.kafka.clients.admin.{AlterConfigOp, ConfigEntry}
+import org.apache.kafka.common.acl.AclOperation
+import org.apache.kafka.common.config.ConfigResource
 import org.apache.kafka.common.errors.UnsupportedVersionException
+import org.apache.kafka.common.internals.{KafkaFutureImpl, Topic}
 import org.apache.kafka.common.memory.MemoryPool
+import org.apache.kafka.common.message.CreateTopicsRequestData.{CreatableTopic, CreatableTopicCollection}
+import org.apache.kafka.common.message.JoinGroupRequestData.JoinGroupRequestProtocol
+import org.apache.kafka.common.message.LeaveGroupRequestData.MemberIdentity
+import org.apache.kafka.common.message.ListOffsetRequestData.{ListOffsetPartition, ListOffsetTopic}
+import org.apache.kafka.common.message.MetadataResponseData.MetadataResponseTopic
+import org.apache.kafka.common.message.OffsetDeleteRequestData.{OffsetDeleteRequestPartition, OffsetDeleteRequestTopic, OffsetDeleteRequestTopicCollection}
+import org.apache.kafka.common.message.StopReplicaRequestData.{StopReplicaPartitionState, StopReplicaTopicState}
+import org.apache.kafka.common.message.UpdateMetadataRequestData.{UpdateMetadataBroker, UpdateMetadataEndpoint, UpdateMetadataPartitionState}
+import org.apache.kafka.common.message._
 import org.apache.kafka.common.metrics.Metrics
-import org.apache.kafka.common.network.ListenerName
+import org.apache.kafka.common.network.{ClientInformation, ListenerName}
 import org.apache.kafka.common.protocol.{ApiKeys, Errors}
-import org.apache.kafka.common.record.RecordBatch
+import org.apache.kafka.common.quota.{ClientQuotaAlteration, ClientQuotaEntity}
+import org.apache.kafka.common.record.FileRecords.TimestampAndOffset
+import org.apache.kafka.common.record._
+import org.apache.kafka.common.replica.ClientMetadata
 import org.apache.kafka.common.requests.ProduceResponse.PartitionResponse
 import org.apache.kafka.common.requests.WriteTxnMarkersRequest.TxnMarkerEntry
-import org.apache.kafka.common.requests._
-import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
-import org.apache.kafka.common.utils.Utils
-import org.easymock.{Capture, EasyMock, IAnswer}
-import org.junit.Assert.{assertEquals, assertTrue}
+import org.apache.kafka.common.requests.{FetchMetadata => JFetchMetadata, _}
+import org.apache.kafka.common.resource.{PatternType, Resource, ResourcePattern, ResourceType}
+import org.apache.kafka.common.security.auth.{KafkaPrincipal, KafkaPrincipalSerde, SecurityProtocol}
+import org.apache.kafka.common.utils.ProducerIdAndEpoch
+import org.apache.kafka.common.{IsolationLevel, Node, TopicPartition}
+import org.apache.kafka.server.authorizer.{Action, AuthorizationResult, Authorizer}
+import org.easymock.EasyMock._
+import org.easymock.{Capture, EasyMock, IAnswer, IArgumentMatcher}
+import org.junit.Assert._
 import org.junit.{After, Test}
 
-import scala.collection.JavaConverters._
-import scala.collection.Map
+import scala.annotation.nowarn
+import scala.collection.{Map, Seq, mutable}
+import scala.compat.java8.OptionConverters._
+import scala.jdk.CollectionConverters._
 
 class KafkaApisTest {
 
-  private val requestChannel = EasyMock.createNiceMock(classOf[RequestChannel])
-  private val requestChannelMetrics = EasyMock.createNiceMock(classOf[RequestChannel.Metrics])
-  private val replicaManager = EasyMock.createNiceMock(classOf[ReplicaManager])
-  private val groupCoordinator = EasyMock.createNiceMock(classOf[GroupCoordinator])
-  private val adminManager = EasyMock.createNiceMock(classOf[AdminManager])
-  private val txnCoordinator = EasyMock.createNiceMock(classOf[TransactionCoordinator])
-  private val controller = EasyMock.createNiceMock(classOf[KafkaController])
-  private val zkUtils = EasyMock.createNiceMock(classOf[ZkUtils])
-  private val zkClient = EasyMock.createNiceMock(classOf[KafkaZkClient])
-  private val metadataCache = EasyMock.createNiceMock(classOf[MetadataCache])
+  private val requestChannel: RequestChannel = EasyMock.createNiceMock(classOf[RequestChannel])
+  private val requestChannelMetrics: RequestChannel.Metrics = EasyMock.createNiceMock(classOf[RequestChannel.Metrics])
+  private val replicaManager: ReplicaManager = EasyMock.createNiceMock(classOf[ReplicaManager])
+  private val groupCoordinator: GroupCoordinator = EasyMock.createNiceMock(classOf[GroupCoordinator])
+  private val adminManager: AdminManager = EasyMock.createNiceMock(classOf[AdminManager])
+  private val txnCoordinator: TransactionCoordinator = EasyMock.createNiceMock(classOf[TransactionCoordinator])
+  private val controller: KafkaController = EasyMock.createNiceMock(classOf[KafkaController])
+  private val forwardingManager: ForwardingManager = EasyMock.createNiceMock(classOf[ForwardingManager])
+  private val hostAddress: Array[Byte] = InetAddress.getByName("192.168.1.1").getAddress
+  private val kafkaPrincipalSerde: Option[KafkaPrincipalSerde] = Option(new KafkaPrincipalSerde {
+    override def serialize(principal: KafkaPrincipal): Array[Byte] = null
+    override def deserialize(bytes: Array[Byte]): KafkaPrincipal = null
+  })
+  private val zkClient: KafkaZkClient = EasyMock.createNiceMock(classOf[KafkaZkClient])
   private val metrics = new Metrics()
   private val brokerId = 1
-  private val authorizer: Option[Authorizer] = None
-  private val clientQuotaManager = EasyMock.createNiceMock(classOf[ClientQuotaManager])
-  private val clientRequestQuotaManager = EasyMock.createNiceMock(classOf[ClientRequestQuotaManager])
-  private val replicaQuotaManager = EasyMock.createNiceMock(classOf[ReplicationQuotaManager])
-  private val quotas = QuotaManagers(clientQuotaManager, clientQuotaManager, clientRequestQuotaManager, replicaQuotaManager, replicaQuotaManager, replicaQuotaManager)
+  private var metadataCache: MetadataCache = new MetadataCache(brokerId)
+  private val clientQuotaManager: ClientQuotaManager = EasyMock.createNiceMock(classOf[ClientQuotaManager])
+  private val clientRequestQuotaManager: ClientRequestQuotaManager = EasyMock.createNiceMock(classOf[ClientRequestQuotaManager])
+  private val clientControllerQuotaManager: ControllerMutationQuotaManager = EasyMock.createNiceMock(classOf[ControllerMutationQuotaManager])
+  private val replicaQuotaManager: ReplicationQuotaManager = EasyMock.createNiceMock(classOf[ReplicationQuotaManager])
+  private val quotas = QuotaManagers(clientQuotaManager, clientQuotaManager, clientRequestQuotaManager,
+    clientControllerQuotaManager, replicaQuotaManager, replicaQuotaManager, replicaQuotaManager, None)
+  private val fetchManager: FetchManager = EasyMock.createNiceMock(classOf[FetchManager])
   private val brokerTopicStats = new BrokerTopicStats
   private val clusterId = "clusterId"
   private val time = new MockTime
+  private val clientId = ""
 
   @After
-  def tearDown() {
+  def tearDown(): Unit = {
     quotas.shutdown()
+    TestUtils.clearYammerMetrics()
     metrics.close()
   }
 
-  def createKafkaApis(interBrokerProtocolVersion: ApiVersion = ApiVersion.latestVersion): KafkaApis = {
+  def createKafkaApis(interBrokerProtocolVersion: ApiVersion = ApiVersion.latestVersion,
+                      authorizer: Option[Authorizer] = None,
+                      enableForwarding: Boolean = false): KafkaApis = {
+    val brokerFeatures = BrokerFeatures.createDefault()
+    val cache = new FinalizedFeatureCache(brokerFeatures)
     val properties = TestUtils.createBrokerConfig(brokerId, "zk")
     properties.put(KafkaConfig.InterBrokerProtocolVersionProp, interBrokerProtocolVersion.toString)
     properties.put(KafkaConfig.LogMessageFormatVersionProp, interBrokerProtocolVersion.toString)
+    properties.put(KafkaConfig.EnableMetadataQuorumProp, enableForwarding.toString)
     new KafkaApis(requestChannel,
       replicaManager,
       adminManager,
       groupCoordinator,
       txnCoordinator,
       controller,
+      forwardingManager,
       zkClient,
       brokerId,
       new KafkaConfig(properties),
@@ -97,10 +138,1184 @@ class KafkaApisTest {
       metrics,
       authorizer,
       quotas,
+      fetchManager,
       brokerTopicStats,
       clusterId,
-      time
+      time,
+      null,
+      brokerFeatures,
+      cache)
+  }
+
+  @Test
+  def testAuthorize(): Unit = {
+    val authorizer: Authorizer = EasyMock.niceMock(classOf[Authorizer])
+
+    val operation = AclOperation.WRITE
+    val resourceType = ResourceType.TOPIC
+    val resourceName = "topic-1"
+    val requestHeader = new RequestHeader(ApiKeys.PRODUCE, ApiKeys.PRODUCE.latestVersion,
+      clientId, 0)
+    val requestContext = new RequestContext(requestHeader, "1", InetAddress.getLocalHost,
+      KafkaPrincipal.ANONYMOUS, ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT),
+      SecurityProtocol.PLAINTEXT, ClientInformation.EMPTY, false)
+
+    val expectedActions = Seq(
+      new Action(operation, new ResourcePattern(resourceType, resourceName, PatternType.LITERAL),
+        1, true, true)
+    )
+
+    EasyMock.expect(authorizer.authorize(requestContext, expectedActions.asJava))
+      .andReturn(Seq(AuthorizationResult.ALLOWED).asJava)
+      .once()
+
+    EasyMock.replay(authorizer)
+
+    val result = createKafkaApis(authorizer = Some(authorizer)).authorize(
+      requestContext, operation, resourceType, resourceName)
+
+    verify(authorizer)
+
+    assertEquals(true, result)
+  }
+
+  @Test
+  def testFilterByAuthorized(): Unit = {
+    val authorizer: Authorizer = EasyMock.niceMock(classOf[Authorizer])
+
+    val operation = AclOperation.WRITE
+    val resourceType = ResourceType.TOPIC
+    val resourceName1 = "topic-1"
+    val resourceName2 = "topic-2"
+    val resourceName3 = "topic-3"
+    val requestHeader = new RequestHeader(ApiKeys.PRODUCE, ApiKeys.PRODUCE.latestVersion,
+      clientId, 0)
+    val requestContext = new RequestContext(requestHeader, "1", InetAddress.getLocalHost,
+      KafkaPrincipal.ANONYMOUS, ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT),
+      SecurityProtocol.PLAINTEXT, ClientInformation.EMPTY, false)
+
+    val expectedActions = Seq(
+      new Action(operation, new ResourcePattern(resourceType, resourceName1, PatternType.LITERAL),
+        2, true, true),
+      new Action(operation, new ResourcePattern(resourceType, resourceName2, PatternType.LITERAL),
+        1, true, true),
+      new Action(operation, new ResourcePattern(resourceType, resourceName3, PatternType.LITERAL),
+        1, true, true),
+    )
+
+    EasyMock.expect(authorizer.authorize(
+      EasyMock.eq(requestContext), matchSameElements(expectedActions.asJava)
+    )).andAnswer { () =>
+      val actions = EasyMock.getCurrentArguments.apply(1).asInstanceOf[util.List[Action]].asScala
+      actions.map { action =>
+        if (Set(resourceName1, resourceName3).contains(action.resourcePattern.name))
+          AuthorizationResult.ALLOWED
+        else
+          AuthorizationResult.DENIED
+      }.asJava
+    }.once()
+
+    EasyMock.replay(authorizer)
+
+    val result = createKafkaApis(authorizer = Some(authorizer)).filterByAuthorized(
+      requestContext,
+      operation,
+      resourceType,
+      // Duplicate resource names should not trigger multiple calls to authorize
+      Seq(resourceName1, resourceName2, resourceName1, resourceName3)
+    )(identity)
+
+    verify(authorizer)
+
+    assertEquals(Set(resourceName1, resourceName3), result)
+  }
+
+  /**
+   * Returns true if the elements in both lists are the same irrespective of ordering.
+   */
+  private def matchSameElements[T](list: util.List[T]): util.List[T] = {
+    EasyMock.reportMatcher(new IArgumentMatcher {
+      def matches(argument: Any): Boolean = argument match {
+        case s: util.List[_] => s.asScala.toSet == list.asScala.toSet
+        case _ => false
+      }
+      def appendTo(buffer: StringBuffer): Unit = buffer.append(s"list($list)")
+    })
+    null
+  }
+
+  @Test
+  def testDescribeConfigsWithAuthorizer(): Unit = {
+    val authorizer: Authorizer = EasyMock.niceMock(classOf[Authorizer])
+
+    val operation = AclOperation.DESCRIBE_CONFIGS
+    val resourceType = ResourceType.TOPIC
+    val resourceName = "topic-1"
+    val requestHeader = new RequestHeader(ApiKeys.DESCRIBE_CONFIGS, ApiKeys.DESCRIBE_CONFIGS.latestVersion,
+      clientId, 0)
+
+    val expectedActions = Seq(
+      new Action(operation, new ResourcePattern(resourceType, resourceName, PatternType.LITERAL),
+        1, true, true)
+    )
+
+    // Verify that authorize is only called once
+    EasyMock.expect(authorizer.authorize(anyObject[RequestContext], EasyMock.eq(expectedActions.asJava)))
+      .andReturn(Seq(AuthorizationResult.ALLOWED).asJava)
+      .once()
+
+    expectNoThrottling()
+
+    val configResource = new ConfigResource(ConfigResource.Type.TOPIC, resourceName)
+    EasyMock.expect(adminManager.describeConfigs(anyObject(), EasyMock.eq(true), EasyMock.eq(false)))
+        .andReturn(
+          List(new DescribeConfigsResponseData.DescribeConfigsResult()
+            .setResourceName(configResource.name)
+            .setResourceType(configResource.`type`.id)
+            .setErrorCode(Errors.NONE.code)
+            .setConfigs(Collections.emptyList())))
+
+    EasyMock.replay(replicaManager, clientRequestQuotaManager, requestChannel, authorizer,
+      adminManager)
+
+    val request = buildRequest(new DescribeConfigsRequest.Builder(new DescribeConfigsRequestData()
+      .setIncludeSynonyms(true)
+      .setResources(List(new DescribeConfigsRequestData.DescribeConfigsResource()
+        .setResourceName("topic-1")
+        .setResourceType(ConfigResource.Type.TOPIC.id)).asJava))
+      .build(requestHeader.apiVersion),
+      requestHeader = Option(requestHeader))
+    createKafkaApis(authorizer = Some(authorizer)).handleDescribeConfigsRequest(request)
+
+    verify(authorizer, adminManager)
+  }
+
+  @Test
+  def testEnvelopeRequestHandlingAsController(): Unit = {
+    val authorizer: Authorizer = EasyMock.niceMock(classOf[Authorizer])
+
+    authorizeResource(authorizer, AclOperation.CLUSTER_ACTION, ResourceType.CLUSTER, Resource.CLUSTER_NAME, AuthorizationResult.ALLOWED)
+
+    val operation = AclOperation.ALTER_CONFIGS
+    val resourceName = "topic-1"
+    val requestHeader = new RequestHeader(ApiKeys.ALTER_CONFIGS, ApiKeys.ALTER_CONFIGS.latestVersion,
+      clientId, 0)
+
+    EasyMock.expect(controller.isActive).andReturn(true)
+
+    authorizeResource(authorizer, operation, ResourceType.TOPIC, resourceName, AuthorizationResult.ALLOWED)
+
+    val capturedResponse = expectNoThrottling()
+
+    val configResource = new ConfigResource(ConfigResource.Type.TOPIC, resourceName)
+    EasyMock.expect(adminManager.alterConfigs(anyObject(), EasyMock.eq(false)))
+      .andReturn(Map(configResource -> ApiError.NONE))
+
+    EasyMock.replay(replicaManager, clientRequestQuotaManager, requestChannel, authorizer,
+      adminManager, controller)
+
+    val configs = Map(
+      configResource -> new AlterConfigsRequest.Config(
+        Seq(new AlterConfigsRequest.ConfigEntry("foo", "bar")).asJava))
+    val alterConfigsRequest = new AlterConfigsRequest.Builder(configs.asJava, false).build(requestHeader.apiVersion)
+
+    val request = buildRequestWithEnvelope(alterConfigsRequest, fromPrivilegedListener = true)
+
+    createKafkaApis(authorizer = Some(authorizer), enableForwarding = true).handle(request)
+
+    val envelopeRequest = request.envelope.get.body[EnvelopeRequest]
+    val response = readResponse(envelopeRequest, capturedResponse)
+      .asInstanceOf[EnvelopeResponse]
+
+    assertEquals(Errors.NONE, response.error)
+
+    val innerResponse = AbstractResponse.parseResponse(
+      response.responseData(),
+      requestHeader
+    ).asInstanceOf[AlterConfigsResponse]
+
+    val responseMap = innerResponse.data.responses().asScala.map { resourceResponse =>
+      resourceResponse.resourceName() -> Errors.forCode(resourceResponse.errorCode)
+    }.toMap
+
+    assertEquals(Map(resourceName -> Errors.NONE), responseMap)
+
+    verify(authorizer, controller, adminManager)
+  }
+
+  @Test
+  def testInvalidEnvelopeRequestAsPrimary(): Unit = {
+    val configResource = new ConfigResource(ConfigResource.Type.TOPIC, "name")
+    val requestHeader = new RequestHeader(ApiKeys.ALTER_CONFIGS, ApiKeys.ALTER_CONFIGS.latestVersion,
+      clientId, 0)
+
+    val configs = Map(
+      configResource -> new AlterConfigsRequest.Config(
+        Seq(new AlterConfigsRequest.ConfigEntry("foo", "bar")).asJava))
+    val alterConfigsRequest = new AlterConfigsRequest.Builder(configs.asJava, false).build(requestHeader.apiVersion)
+    val serializedRequestData = RequestTestUtils.serializeRequestWithHeader(requestHeader, alterConfigsRequest)
+
+    val capturedResponse = expectNoThrottling()
+
+    EasyMock.replay(replicaManager, clientRequestQuotaManager, requestChannel, controller)
+
+    val envelopeRequest = new EnvelopeRequest.Builder(serializedRequestData, new Array[Byte](0), hostAddress)
+      .build(ApiKeys.ENVELOPE.latestVersion)
+    val request = buildRequest(envelopeRequest, fromPrivilegedListener = true)
+
+    createKafkaApis(enableForwarding = true).handle(request)
+
+    val response = readResponse(envelopeRequest, capturedResponse)
+      .asInstanceOf[EnvelopeResponse]
+    assertEquals(Errors.UNKNOWN_SERVER_ERROR, response.error())
+  }
+
+  @Test
+  def testInvalidEnvelopeRequestWithNonForwardableAPI(): Unit = {
+    val requestHeader = new RequestHeader(ApiKeys.LEAVE_GROUP, ApiKeys.LEAVE_GROUP.latestVersion,
+      clientId, 0)
+    val leaveGroupRequest = new LeaveGroupRequest.Builder("group",
+      Collections.singletonList(new MemberIdentity())).build(requestHeader.apiVersion)
+    val serializedRequestData = RequestTestUtils.serializeRequestWithHeader(requestHeader, leaveGroupRequest)
+
+    val capturedResponse = expectNoThrottling()
+
+    EasyMock.replay(replicaManager, clientRequestQuotaManager, requestChannel, controller)
+
+    val envelopeHeader = new RequestHeader(ApiKeys.ENVELOPE, ApiKeys.ENVELOPE.latestVersion,
+      clientId, 0)
+
+    val envelopeRequest = new EnvelopeRequest.Builder(serializedRequestData, new Array[Byte](0), hostAddress)
+      .build(envelopeHeader.apiVersion)
+    val request = buildRequestWithEnvelope(leaveGroupRequest, fromPrivilegedListener = true)
+
+    createKafkaApis(enableForwarding = true).handle(request)
+
+    val response = readResponse(envelopeRequest, capturedResponse)
+      .asInstanceOf[EnvelopeResponse]
+    assertEquals(Errors.INVALID_REQUEST, response.error())
+  }
+
+  @Test
+  def testEnvelopeRequestWithNotFromPrivilegedListener(): Unit = {
+    testInvalidEnvelopeRequest(Errors.NONE, fromPrivilegedListener = false, shouldCloseConnection = true)
+  }
+
+  @Test
+  def testEnvelopeRequestNotAuthorized(): Unit = {
+    testInvalidEnvelopeRequest(Errors.CLUSTER_AUTHORIZATION_FAILED,
+      performAuthorize = true, authorizeResult = AuthorizationResult.DENIED)
+  }
+
+  @Test
+  def testEnvelopeRequestNotControllerHandling(): Unit = {
+    testInvalidEnvelopeRequest(Errors.NOT_CONTROLLER, performAuthorize = true, isActiveController = false)
+  }
+
+  private def testInvalidEnvelopeRequest(expectedError: Errors,
+                                         fromPrivilegedListener: Boolean = true,
+                                         shouldCloseConnection: Boolean = false,
+                                         principalSerde: Option[KafkaPrincipalSerde] = kafkaPrincipalSerde,
+                                         performAuthorize: Boolean = false,
+                                         authorizeResult: AuthorizationResult = AuthorizationResult.ALLOWED,
+                                         isActiveController: Boolean = true): Unit = {
+    val authorizer: Authorizer = EasyMock.niceMock(classOf[Authorizer])
+
+    if (performAuthorize) {
+      authorizeResource(authorizer, AclOperation.CLUSTER_ACTION, ResourceType.CLUSTER, Resource.CLUSTER_NAME, authorizeResult)
+    }
+
+    val resourceName = "topic-1"
+    val requestHeader = new RequestHeader(ApiKeys.ALTER_CONFIGS, ApiKeys.ALTER_CONFIGS.latestVersion,
+      clientId, 0)
+
+    EasyMock.expect(controller.isActive).andReturn(isActiveController)
+
+    val capturedResponse = expectNoThrottling()
+
+    val configResource = new ConfigResource(ConfigResource.Type.TOPIC, resourceName)
+
+    EasyMock.replay(replicaManager, clientRequestQuotaManager, requestChannel, authorizer,
+      adminManager, controller)
+
+    val configs = Map(
+      configResource -> new AlterConfigsRequest.Config(
+        Seq(new AlterConfigsRequest.ConfigEntry("foo", "bar")).asJava))
+    val alterConfigsRequest = new AlterConfigsRequest.Builder(configs.asJava, false)
+      .build(requestHeader.apiVersion)
+
+    val request = buildRequestWithEnvelope(alterConfigsRequest,
+      fromPrivilegedListener = fromPrivilegedListener)
+    createKafkaApis(authorizer = Some(authorizer), enableForwarding = true).handle(request)
+
+    if (shouldCloseConnection) {
+      assertTrue(capturedResponse.getValue.isInstanceOf[CloseConnectionResponse])
+    } else {
+      val envelopeRequest = request.envelope.get.body[EnvelopeRequest]
+      val response = readResponse(envelopeRequest, capturedResponse)
+        .asInstanceOf[EnvelopeResponse]
+
+      assertEquals(expectedError, response.error())
+
+      verify(authorizer, adminManager)
+    }
+  }
+
+  @Test
+  def testAlterConfigsWithAuthorizer(): Unit = {
+    val authorizer: Authorizer = EasyMock.niceMock(classOf[Authorizer])
+
+    val authorizedTopic = "authorized-topic"
+    val unauthorizedTopic = "unauthorized-topic"
+    val (authorizedResource, unauthorizedResource) =
+      createConfigsWithAuthorization(authorizer, authorizedTopic, unauthorizedTopic)
+
+    val configs = Map(
+      authorizedResource -> new AlterConfigsRequest.Config(
+        Seq(new AlterConfigsRequest.ConfigEntry("foo", "bar")).asJava),
+      unauthorizedResource -> new AlterConfigsRequest.Config(
+        Seq(new AlterConfigsRequest.ConfigEntry("foo-1", "bar-1")).asJava)
     )
+
+    val topicHeader = new RequestHeader(ApiKeys.ALTER_CONFIGS, ApiKeys.ALTER_CONFIGS.latestVersion,
+      clientId, 0)
+
+    val alterConfigsRequest = new AlterConfigsRequest.Builder(configs.asJava, false)
+      .build(topicHeader.apiVersion)
+    val request = buildRequest(alterConfigsRequest)
+
+    EasyMock.expect(controller.isActive).andReturn(false)
+
+    val capturedResponse = expectNoThrottling()
+
+    EasyMock.expect(adminManager.alterConfigs(anyObject(), EasyMock.eq(false)))
+      .andReturn(Map(authorizedResource -> ApiError.NONE))
+
+    EasyMock.replay(replicaManager, clientRequestQuotaManager, requestChannel, authorizer,
+      adminManager, controller)
+
+    createKafkaApis(authorizer = Some(authorizer)).handleAlterConfigsRequest(request)
+
+    verifyAlterConfigResult(alterConfigsRequest,
+      capturedResponse, Map(authorizedTopic -> Errors.NONE,
+        unauthorizedTopic -> Errors.TOPIC_AUTHORIZATION_FAILED))
+
+    verify(authorizer, adminManager)
+  }
+
+  @Test
+  def testAlterConfigsWithForwarding(): Unit = {
+    val requestBuilder = new AlterConfigsRequest.Builder(Collections.emptyMap(), false)
+    testForwardableAPI(ApiKeys.ALTER_CONFIGS, requestBuilder)
+  }
+
+  private def testForwardableAPI(apiKey: ApiKeys, requestBuilder: AbstractRequest.Builder[_ <: AbstractRequest]): Unit = {
+    val topicHeader = new RequestHeader(apiKey, apiKey.latestVersion,
+      clientId, 0)
+
+    val request = buildRequest(requestBuilder.build(topicHeader.apiVersion))
+
+    EasyMock.expect(controller.isActive).andReturn(false)
+
+    expectNoThrottling()
+
+    EasyMock.expect(forwardingManager.forwardRequest(
+      EasyMock.eq(request),
+      anyObject[AbstractResponse => Unit]()
+    )).once()
+
+    EasyMock.replay(replicaManager, clientRequestQuotaManager, requestChannel, controller, forwardingManager)
+
+    createKafkaApis(enableForwarding = true).handle(request)
+
+    EasyMock.verify(controller, forwardingManager)
+  }
+
+  private def authorizeResource(authorizer: Authorizer,
+                                operation: AclOperation,
+                                resourceType: ResourceType,
+                                resourceName: String,
+                                result: AuthorizationResult,
+                                logIfAllowed: Boolean = true,
+                                logIfDenied: Boolean = true): Unit = {
+    val expectedAuthorizedAction = if (operation == AclOperation.CLUSTER_ACTION)
+      new Action(operation,
+        new ResourcePattern(ResourceType.CLUSTER, Resource.CLUSTER_NAME, PatternType.LITERAL),
+        1, logIfAllowed, logIfDenied)
+    else
+      new Action(operation,
+        new ResourcePattern(resourceType, resourceName, PatternType.LITERAL),
+        1, logIfAllowed, logIfDenied)
+
+    EasyMock.expect(authorizer.authorize(anyObject[RequestContext], EasyMock.eq(Seq(expectedAuthorizedAction).asJava)))
+      .andReturn(Seq(result).asJava)
+      .once()
+  }
+
+  private def verifyAlterConfigResult(alterConfigsRequest: AlterConfigsRequest,
+                                      capturedResponse: Capture[RequestChannel.Response],
+                                      expectedResults: Map[String, Errors]): Unit = {
+    val response = readResponse(alterConfigsRequest, capturedResponse)
+      .asInstanceOf[AlterConfigsResponse]
+    val responseMap = response.data.responses().asScala.map { resourceResponse =>
+      resourceResponse.resourceName() -> Errors.forCode(resourceResponse.errorCode)
+    }.toMap
+
+    assertEquals(expectedResults, responseMap)
+  }
+
+  private def createConfigsWithAuthorization(authorizer: Authorizer,
+                                             authorizedTopic: String,
+                                             unauthorizedTopic: String): (ConfigResource, ConfigResource) = {
+    val authorizedResource = new ConfigResource(ConfigResource.Type.TOPIC, authorizedTopic)
+
+    val unauthorizedResource = new ConfigResource(ConfigResource.Type.TOPIC, unauthorizedTopic)
+
+    createTopicAuthorization(authorizer, AclOperation.ALTER_CONFIGS, authorizedTopic, unauthorizedTopic)
+    (authorizedResource, unauthorizedResource)
+  }
+
+  @Test
+  def testIncrementalAlterConfigsWithAuthorizer(): Unit = {
+    val authorizer: Authorizer = EasyMock.niceMock(classOf[Authorizer])
+
+    val authorizedTopic = "authorized-topic"
+    val unauthorizedTopic = "unauthorized-topic"
+    val (authorizedResource, unauthorizedResource) =
+      createConfigsWithAuthorization(authorizer, authorizedTopic, unauthorizedTopic)
+
+    val requestHeader = new RequestHeader(ApiKeys.INCREMENTAL_ALTER_CONFIGS, ApiKeys.INCREMENTAL_ALTER_CONFIGS.latestVersion, clientId, 0)
+
+    val incrementalAlterConfigsRequest = getIncrementalAlterConfigRequestBuilder(Seq(authorizedResource, unauthorizedResource))
+      .build(requestHeader.apiVersion)
+    val request = buildRequest(incrementalAlterConfigsRequest,
+      fromPrivilegedListener = true, requestHeader = Option(requestHeader))
+
+    EasyMock.expect(controller.isActive).andReturn(true)
+
+    val capturedResponse = expectNoThrottling()
+
+    EasyMock.expect(adminManager.incrementalAlterConfigs(anyObject(), EasyMock.eq(false)))
+      .andReturn(Map(authorizedResource -> ApiError.NONE))
+
+    EasyMock.replay(replicaManager, clientRequestQuotaManager, requestChannel, authorizer,
+      adminManager, controller)
+
+    createKafkaApis(authorizer = Some(authorizer)).handleIncrementalAlterConfigsRequest(request)
+
+    verifyIncrementalAlterConfigResult(incrementalAlterConfigsRequest,
+      capturedResponse, Map(authorizedTopic -> Errors.NONE,
+        unauthorizedTopic -> Errors.TOPIC_AUTHORIZATION_FAILED))
+
+    verify(authorizer, adminManager)
+  }
+
+  @Test
+  def testIncrementalAlterConfigsWithForwarding(): Unit = {
+    val requestBuilder = new IncrementalAlterConfigsRequest.Builder(
+      new IncrementalAlterConfigsRequestData())
+    testForwardableAPI(ApiKeys.INCREMENTAL_ALTER_CONFIGS, requestBuilder)
+  }
+
+  private def getIncrementalAlterConfigRequestBuilder(configResources: Seq[ConfigResource]): IncrementalAlterConfigsRequest.Builder = {
+    val resourceMap = configResources.map(configResource => {
+      configResource -> Set(
+        new AlterConfigOp(new ConfigEntry("foo", "bar"),
+        OpType.forId(configResource.`type`.id))).asJavaCollection
+    }).toMap.asJava
+
+    new IncrementalAlterConfigsRequest.Builder(resourceMap, false)
+  }
+
+  private def verifyIncrementalAlterConfigResult(incrementalAlterConfigsRequest: IncrementalAlterConfigsRequest,
+                                                 capturedResponse: Capture[RequestChannel.Response],
+                                                 expectedResults: Map[String, Errors]): Unit = {
+    val response = readResponse(incrementalAlterConfigsRequest, capturedResponse)
+      .asInstanceOf[IncrementalAlterConfigsResponse]
+    val responseMap = response.data.responses().asScala.map { resourceResponse =>
+      resourceResponse.resourceName() -> Errors.forCode(resourceResponse.errorCode)
+    }.toMap
+
+    assertEquals(expectedResults, responseMap)
+  }
+
+  @Test
+  def testAlterClientQuotasWithAuthorizer(): Unit = {
+    val authorizer: Authorizer = EasyMock.niceMock(classOf[Authorizer])
+
+    authorizeResource(authorizer, AclOperation.ALTER_CONFIGS, ResourceType.CLUSTER,
+      Resource.CLUSTER_NAME, AuthorizationResult.DENIED)
+
+    val quotaEntity = new ClientQuotaEntity(Collections.singletonMap(ClientQuotaEntity.USER, "user"))
+    val quotas = Seq(new ClientQuotaAlteration(quotaEntity, Seq.empty.asJavaCollection))
+
+    val requestHeader = new RequestHeader(ApiKeys.ALTER_CLIENT_QUOTAS, ApiKeys.ALTER_CLIENT_QUOTAS.latestVersion, clientId, 0)
+
+    val alterClientQuotasRequest = new AlterClientQuotasRequest.Builder(quotas.asJavaCollection, false)
+      .build(requestHeader.apiVersion)
+    val request = buildRequest(alterClientQuotasRequest,
+      fromPrivilegedListener = true, requestHeader = Option(requestHeader))
+
+    EasyMock.expect(controller.isActive).andReturn(true)
+
+    val capturedResponse = expectNoThrottling()
+
+    EasyMock.replay(replicaManager, clientRequestQuotaManager, requestChannel, authorizer,
+      adminManager, controller)
+
+    createKafkaApis(authorizer = Some(authorizer)).handleAlterClientQuotasRequest(request)
+
+    verifyAlterClientQuotaResult(alterClientQuotasRequest,
+      capturedResponse, Map(quotaEntity -> Errors.CLUSTER_AUTHORIZATION_FAILED))
+
+    verify(authorizer, adminManager)
+  }
+
+  @Test
+  def testAlterClientQuotasWithForwarding(): Unit = {
+    val requestBuilder = new AlterClientQuotasRequest.Builder(List.empty.asJava, false)
+    testForwardableAPI(ApiKeys.ALTER_CLIENT_QUOTAS, requestBuilder)
+  }
+
+  private def verifyAlterClientQuotaResult(alterClientQuotasRequest: AlterClientQuotasRequest,
+                                           capturedResponse: Capture[RequestChannel.Response],
+                                           expected: Map[ClientQuotaEntity, Errors]): Unit = {
+    val response = readResponse(alterClientQuotasRequest, capturedResponse)
+      .asInstanceOf[AlterClientQuotasResponse]
+    val futures = expected.keys.map(quotaEntity => quotaEntity -> new KafkaFutureImpl[Void]()).toMap
+    response.complete(futures.asJava)
+    futures.foreach {
+      case (entity, future) =>
+        future.whenComplete((_, thrown) =>
+          assertEquals(thrown, expected(entity).exception())
+        ).isDone
+    }
+  }
+
+  @Test
+  def testCreateTopicsWithAuthorizer(): Unit = {
+    val authorizer: Authorizer = EasyMock.niceMock(classOf[Authorizer])
+
+    val authorizedTopic = "authorized-topic"
+    val unauthorizedTopic = "unauthorized-topic"
+
+    authorizeResource(authorizer, AclOperation.CREATE, ResourceType.CLUSTER,
+      Resource.CLUSTER_NAME, AuthorizationResult.DENIED, logIfDenied = false)
+
+    createCombinedTopicAuthorization(authorizer, AclOperation.CREATE,
+      authorizedTopic, unauthorizedTopic)
+
+    createCombinedTopicAuthorization(authorizer, AclOperation.DESCRIBE_CONFIGS,
+      authorizedTopic, unauthorizedTopic, logIfDenied = false)
+
+    val requestHeader = new RequestHeader(ApiKeys.CREATE_TOPICS, ApiKeys.CREATE_TOPICS.latestVersion, clientId, 0)
+
+    EasyMock.expect(controller.isActive).andReturn(true)
+
+    val capturedResponse = expectNoThrottling()
+
+    val topics = new CreateTopicsRequestData.CreatableTopicCollection(2)
+    val topicToCreate = new CreateTopicsRequestData.CreatableTopic()
+      .setName(authorizedTopic)
+    topics.add(topicToCreate)
+
+    val topicToFilter = new CreateTopicsRequestData.CreatableTopic()
+      .setName(unauthorizedTopic)
+    topics.add(topicToFilter)
+
+    val timeout = 10
+    val createTopicsRequest = new CreateTopicsRequest.Builder(
+      new CreateTopicsRequestData()
+        .setTimeoutMs(timeout)
+        .setValidateOnly(false)
+        .setTopics(topics))
+      .build(requestHeader.apiVersion)
+    val request = buildRequest(createTopicsRequest,
+      fromPrivilegedListener = true, requestHeader = Option(requestHeader))
+
+    EasyMock.expect(clientControllerQuotaManager.newQuotaFor(
+      EasyMock.eq(request), EasyMock.eq(6))).andReturn(UnboundedControllerMutationQuota)
+
+    val capturedCallback = EasyMock.newCapture[Map[String, ApiError] => Unit]()
+
+    EasyMock.expect(adminManager.createTopics(
+      EasyMock.eq(timeout),
+      EasyMock.eq(false),
+      EasyMock.eq(Map(authorizedTopic -> topicToCreate)),
+      anyObject(),
+      EasyMock.eq(UnboundedControllerMutationQuota),
+      EasyMock.capture(capturedCallback)))
+
+    EasyMock.replay(replicaManager, clientRequestQuotaManager, clientControllerQuotaManager,
+      requestChannel, authorizer, adminManager, controller)
+
+    createKafkaApis(authorizer = Some(authorizer)).handleCreateTopicsRequest(request)
+
+    capturedCallback.getValue.apply(Map(authorizedTopic -> ApiError.NONE))
+
+    verifyCreateTopicsResult(createTopicsRequest,
+      capturedResponse, Map(authorizedTopic -> Errors.NONE,
+        unauthorizedTopic -> Errors.TOPIC_AUTHORIZATION_FAILED))
+
+    verify(authorizer, adminManager, clientControllerQuotaManager)
+  }
+
+  @Test
+  def testCreateTopicsWithForwarding(): Unit = {
+    val requestBuilder = new CreateTopicsRequest.Builder(
+      new CreateTopicsRequestData().setTopics(
+        new CreatableTopicCollection(Collections.singleton(
+          new CreatableTopic().setName("topic").setNumPartitions(1).
+            setReplicationFactor(1.toShort)).iterator())))
+    testForwardableAPI(ApiKeys.CREATE_TOPICS, requestBuilder)
+  }
+
+  private def createTopicAuthorization(authorizer: Authorizer,
+                                       operation: AclOperation,
+                                       authorizedTopic: String,
+                                       unauthorizedTopic: String,
+                                       logIfAllowed: Boolean = true,
+                                       logIfDenied: Boolean = true): Unit = {
+    authorizeResource(authorizer, operation, ResourceType.TOPIC,
+      authorizedTopic, AuthorizationResult.ALLOWED, logIfAllowed, logIfDenied)
+    authorizeResource(authorizer, operation, ResourceType.TOPIC,
+      unauthorizedTopic, AuthorizationResult.DENIED, logIfAllowed, logIfDenied)
+  }
+
+  private def createCombinedTopicAuthorization(authorizer: Authorizer,
+                                               operation: AclOperation,
+                                               authorizedTopic: String,
+                                               unauthorizedTopic: String,
+                                               logIfAllowed: Boolean = true,
+                                               logIfDenied: Boolean = true): Unit = {
+    val expectedAuthorizedActions = Seq(
+      new Action(operation,
+        new ResourcePattern(ResourceType.TOPIC, authorizedTopic, PatternType.LITERAL),
+        1, logIfAllowed, logIfDenied),
+      new Action(operation,
+        new ResourcePattern(ResourceType.TOPIC, unauthorizedTopic, PatternType.LITERAL),
+        1, logIfAllowed, logIfDenied))
+
+    EasyMock.expect(authorizer.authorize(
+      anyObject[RequestContext], matchSameElements(expectedAuthorizedActions.asJava)
+    )).andAnswer { () =>
+      val actions = EasyMock.getCurrentArguments.apply(1).asInstanceOf[util.List[Action]].asScala
+      actions.map { action =>
+        if (action.resourcePattern().name().equals(authorizedTopic))
+          AuthorizationResult.ALLOWED
+        else
+          AuthorizationResult.DENIED
+      }.asJava
+    }.once()
+  }
+
+  private def verifyCreateTopicsResult(createTopicsRequest: CreateTopicsRequest,
+                                       capturedResponse: Capture[RequestChannel.Response],
+                                       expectedResults: Map[String, Errors]): Unit = {
+    val response = readResponse(createTopicsRequest, capturedResponse)
+      .asInstanceOf[CreateTopicsResponse]
+    val responseMap = response.data.topics().asScala.map { topicResponse =>
+      topicResponse.name() -> Errors.forCode(topicResponse.errorCode)
+    }.toMap
+
+    assertEquals(expectedResults, responseMap)
+  }
+
+  @Test
+  def testCreateAclWithForwarding(): Unit = {
+    val requestBuilder = new CreateAclsRequest.Builder(new CreateAclsRequestData())
+    testForwardableAPI(ApiKeys.CREATE_ACLS, requestBuilder)
+  }
+
+  @Test
+  def testDeleteAclWithForwarding(): Unit = {
+    val requestBuilder = new DeleteAclsRequest.Builder(new DeleteAclsRequestData())
+    testForwardableAPI(ApiKeys.DELETE_ACLS, requestBuilder)
+  }
+
+  @Test
+  def testCreateDelegationTokenWithForwarding(): Unit = {
+    val requestBuilder = new CreateDelegationTokenRequest.Builder(new CreateDelegationTokenRequestData())
+    testForwardableAPI(ApiKeys.CREATE_DELEGATION_TOKEN, requestBuilder)
+  }
+
+  @Test
+  def testRenewDelegationTokenWithForwarding(): Unit = {
+    val requestBuilder = new RenewDelegationTokenRequest.Builder(new RenewDelegationTokenRequestData())
+    testForwardableAPI(ApiKeys.RENEW_DELEGATION_TOKEN, requestBuilder)
+  }
+
+  @Test
+  def testExpireDelegationTokenWithForwarding(): Unit = {
+    val requestBuilder = new ExpireDelegationTokenRequest.Builder(new ExpireDelegationTokenRequestData())
+    testForwardableAPI(ApiKeys.EXPIRE_DELEGATION_TOKEN, requestBuilder)
+  }
+
+  @Test
+  def testAlterPartitionReassignmentsWithForwarding(): Unit = {
+    val requestBuilder = new AlterPartitionReassignmentsRequest.Builder(new AlterPartitionReassignmentsRequestData())
+    testForwardableAPI(ApiKeys.ALTER_PARTITION_REASSIGNMENTS, requestBuilder)
+  }
+
+  @Test
+  def testCreatePartitionsWithForwarding(): Unit = {
+    val requestBuilder = new CreatePartitionsRequest.Builder(new CreatePartitionsRequestData())
+    testForwardableAPI(ApiKeys.CREATE_PARTITIONS, requestBuilder)
+  }
+
+  @Test
+  def testDeleteTopicsWithForwarding(): Unit = {
+    val requestBuilder = new DeleteTopicsRequest.Builder(new DeleteTopicsRequestData())
+    testForwardableAPI(ApiKeys.DELETE_TOPICS, requestBuilder)
+  }
+
+  @Test
+  def testUpdateFeaturesWithForwarding(): Unit = {
+    val requestBuilder = new UpdateFeaturesRequest.Builder(new UpdateFeaturesRequestData())
+    testForwardableAPI(ApiKeys.UPDATE_FEATURES, requestBuilder)
+  }
+
+  @Test
+  def testAlterScramWithForwarding(): Unit = {
+    val requestBuilder = new AlterUserScramCredentialsRequest.Builder(new AlterUserScramCredentialsRequestData())
+    testForwardableAPI(ApiKeys.ALTER_USER_SCRAM_CREDENTIALS, requestBuilder)
+  }
+
+  @Test
+  def testOffsetCommitWithInvalidPartition(): Unit = {
+    val topic = "topic"
+    setupBasicMetadataCache(topic, numPartitions = 1)
+
+    def checkInvalidPartition(invalidPartitionId: Int): Unit = {
+      EasyMock.reset(replicaManager, clientRequestQuotaManager, requestChannel)
+
+      val offsetCommitRequest = 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(""))
+              )
+          ))).build()
+
+      val request = buildRequest(offsetCommitRequest)
+      val capturedResponse = expectNoThrottling()
+      EasyMock.replay(replicaManager, clientRequestQuotaManager, requestChannel)
+      createKafkaApis().handleOffsetCommitRequest(request)
+
+      val response = readResponse(offsetCommitRequest, capturedResponse)
+        .asInstanceOf[OffsetCommitResponse]
+      assertEquals(Errors.UNKNOWN_TOPIC_OR_PARTITION,
+        Errors.forCode(response.data().topics().get(0).partitions().get(0).errorCode()))
+    }
+
+    checkInvalidPartition(-1)
+    checkInvalidPartition(1) // topic has only one partition
+  }
+
+  @Test
+  def testTxnOffsetCommitWithInvalidPartition(): Unit = {
+    val topic = "topic"
+    setupBasicMetadataCache(topic, numPartitions = 1)
+
+    def checkInvalidPartition(invalidPartitionId: Int): Unit = {
+      EasyMock.reset(replicaManager, clientRequestQuotaManager, requestChannel)
+
+      val invalidTopicPartition = new TopicPartition(topic, invalidPartitionId)
+      val partitionOffsetCommitData = new TxnOffsetCommitRequest.CommittedOffset(15L, "", Optional.empty())
+      val offsetCommitRequest = new TxnOffsetCommitRequest.Builder(
+        "txnId",
+        "groupId",
+        15L,
+        0.toShort,
+        Map(invalidTopicPartition -> partitionOffsetCommitData).asJava,
+        false
+      ).build()
+      val request = buildRequest(offsetCommitRequest)
+
+      val capturedResponse = expectNoThrottling()
+      EasyMock.replay(replicaManager, clientRequestQuotaManager, requestChannel)
+      createKafkaApis().handleTxnOffsetCommitRequest(request)
+
+      val response = readResponse(offsetCommitRequest, capturedResponse)
+        .asInstanceOf[TxnOffsetCommitResponse]
+      assertEquals(Errors.UNKNOWN_TOPIC_OR_PARTITION, response.errors().get(invalidTopicPartition))
+    }
+
+    checkInvalidPartition(-1)
+    checkInvalidPartition(1) // topic has only one partition
+  }
+
+  @Test
+  def shouldReplaceCoordinatorNotAvailableWithLoadInProcessInTxnOffsetCommitWithOlderClient(): Unit = {
+    val topic = "topic"
+    setupBasicMetadataCache(topic, numPartitions = 2)
+
+    for (version <- ApiKeys.TXN_OFFSET_COMMIT.oldestVersion to ApiKeys.TXN_OFFSET_COMMIT.latestVersion) {
+      EasyMock.reset(replicaManager, clientRequestQuotaManager, requestChannel, groupCoordinator)
+
+      val topicPartition = new TopicPartition(topic, 1)
+      val capturedResponse: Capture[RequestChannel.Response] = EasyMock.newCapture()
+      val responseCallback: Capture[Map[TopicPartition, Errors] => Unit] = EasyMock.newCapture()
+
+      val partitionOffsetCommitData = new TxnOffsetCommitRequest.CommittedOffset(15L, "", Optional.empty())
+      val groupId = "groupId"
+
+      val producerId = 15L
+      val epoch = 0.toShort
+
+      val offsetCommitRequest = new TxnOffsetCommitRequest.Builder(
+        "txnId",
+        groupId,
+        producerId,
+        epoch,
+        Map(topicPartition -> partitionOffsetCommitData).asJava,
+        false
+      ).build(version.toShort)
+      val request = buildRequest(offsetCommitRequest)
+
+      EasyMock.expect(groupCoordinator.handleTxnCommitOffsets(
+        EasyMock.eq(groupId),
+        EasyMock.eq(producerId),
+        EasyMock.eq(epoch),
+        EasyMock.anyString(),
+        EasyMock.eq(Option.empty),
+        EasyMock.anyInt(),
+        EasyMock.anyObject(),
+        EasyMock.capture(responseCallback)
+      )).andAnswer(
+        () => responseCallback.getValue.apply(Map(topicPartition -> Errors.COORDINATOR_LOAD_IN_PROGRESS)))
+
+    EasyMock.expect(requestChannel.sendResponse(EasyMock.capture(capturedResponse)))
+
+      EasyMock.replay(replicaManager, clientRequestQuotaManager, requestChannel, groupCoordinator)
+
+      createKafkaApis().handleTxnOffsetCommitRequest(request)
+
+      val response = readResponse(offsetCommitRequest, capturedResponse)
+        .asInstanceOf[TxnOffsetCommitResponse]
+
+      if (version < 2) {
+        assertEquals(Errors.COORDINATOR_NOT_AVAILABLE, response.errors().get(topicPartition))
+      } else {
+        assertEquals(Errors.COORDINATOR_LOAD_IN_PROGRESS, response.errors().get(topicPartition))
+      }
+    }
+  }
+
+  @Test
+  def shouldReplaceProducerFencedWithInvalidProducerEpochInInitProducerIdWithOlderClient(): Unit = {
+    val topic = "topic"
+    setupBasicMetadataCache(topic, numPartitions = 2)
+
+    for (version <- ApiKeys.INIT_PRODUCER_ID.oldestVersion to ApiKeys.INIT_PRODUCER_ID.latestVersion) {
+
+      EasyMock.reset(replicaManager, clientRequestQuotaManager, requestChannel, txnCoordinator)
+
+      val capturedResponse: Capture[RequestChannel.Response] = EasyMock.newCapture()
+      val responseCallback: Capture[InitProducerIdResult => Unit] = EasyMock.newCapture()
+
+      val transactionalId = "txnId"
+      val producerId = if (version < 3)
+        RecordBatch.NO_PRODUCER_ID
+      else
+        15
+
+      val epoch = if (version < 3)
+        RecordBatch.NO_PRODUCER_EPOCH
+      else
+        0.toShort
+
+      val txnTimeoutMs = TimeUnit.MINUTES.toMillis(15).toInt
+
+      val initProducerIdRequest = new InitProducerIdRequest.Builder(
+        new InitProducerIdRequestData()
+          .setTransactionalId(transactionalId)
+          .setTransactionTimeoutMs(txnTimeoutMs)
+          .setProducerId(producerId)
+          .setProducerEpoch(epoch)
+      ).build(version.toShort)
+
+      val request = buildRequest(initProducerIdRequest)
+
+      val expectedProducerIdAndEpoch = if (version < 3)
+        Option.empty
+      else
+        Option(new ProducerIdAndEpoch(producerId, epoch))
+
+      EasyMock.expect(txnCoordinator.handleInitProducerId(
+        EasyMock.eq(transactionalId),
+        EasyMock.eq(txnTimeoutMs),
+        EasyMock.eq(expectedProducerIdAndEpoch),
+        EasyMock.capture(responseCallback)
+      )).andAnswer(
+        () => responseCallback.getValue.apply(InitProducerIdResult(producerId, epoch, Errors.PRODUCER_FENCED)))
+
+      EasyMock.expect(requestChannel.sendResponse(EasyMock.capture(capturedResponse)))
+
+      EasyMock.replay(replicaManager, clientRequestQuotaManager, requestChannel, txnCoordinator)
+
+      createKafkaApis().handleInitProducerIdRequest(request)
+
+      val response = readResponse(initProducerIdRequest, capturedResponse)
+        .asInstanceOf[InitProducerIdResponse]
+
+      if (version < 4) {
+        assertEquals(Errors.INVALID_PRODUCER_EPOCH.code, response.data.errorCode)
+      } else {
+        assertEquals(Errors.PRODUCER_FENCED.code, response.data.errorCode)
+      }
+    }
+  }
+
+  @Test
+  def shouldReplaceProducerFencedWithInvalidProducerEpochInAddOffsetToTxnWithOlderClient(): Unit = {
+    val topic = "topic"
+    setupBasicMetadataCache(topic, numPartitions = 2)
+
+    for (version <- ApiKeys.ADD_OFFSETS_TO_TXN.oldestVersion to ApiKeys.ADD_OFFSETS_TO_TXN.latestVersion) {
+
+      EasyMock.reset(replicaManager, clientRequestQuotaManager, requestChannel, groupCoordinator, txnCoordinator)
+
+      val capturedResponse: Capture[RequestChannel.Response] = EasyMock.newCapture()
+      val responseCallback: Capture[Errors => Unit] = EasyMock.newCapture()
+
+      val groupId = "groupId"
+      val transactionalId = "txnId"
+      val producerId = 15L
+      val epoch = 0.toShort
+
+      val addOffsetsToTxnRequest = new AddOffsetsToTxnRequest.Builder(
+        new AddOffsetsToTxnRequestData()
+          .setGroupId(groupId)
+          .setTransactionalId(transactionalId)
+          .setProducerId(producerId)
+          .setProducerEpoch(epoch)
+      ).build(version.toShort)
+      val request = buildRequest(addOffsetsToTxnRequest)
+
+      val partition = 1
+      EasyMock.expect(groupCoordinator.partitionFor(
+        EasyMock.eq(groupId)
+      )).andReturn(partition)
+
+      EasyMock.expect(txnCoordinator.handleAddPartitionsToTransaction(
+        EasyMock.eq(transactionalId),
+        EasyMock.eq(producerId),
+        EasyMock.eq(epoch),
+        EasyMock.eq(Set(new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, partition))),
+        EasyMock.capture(responseCallback)
+      )).andAnswer(
+        () => responseCallback.getValue.apply(Errors.PRODUCER_FENCED))
+
+      EasyMock.expect(requestChannel.sendResponse(EasyMock.capture(capturedResponse)))
+
+      EasyMock.replay(replicaManager, clientRequestQuotaManager, requestChannel, txnCoordinator, groupCoordinator)
+
+      createKafkaApis().handleAddOffsetsToTxnRequest(request)
+
+      val response = readResponse(addOffsetsToTxnRequest, capturedResponse)
+        .asInstanceOf[AddOffsetsToTxnResponse]
+
+      if (version < 2) {
+        assertEquals(Errors.INVALID_PRODUCER_EPOCH.code, response.data.errorCode)
+      } else {
+        assertEquals(Errors.PRODUCER_FENCED.code, response.data.errorCode)
+      }
+    }
+  }
+
+  @Test
+  def shouldReplaceProducerFencedWithInvalidProducerEpochInAddPartitionToTxnWithOlderClient(): Unit = {
+    val topic = "topic"
+    setupBasicMetadataCache(topic, numPartitions = 2)
+
+    for (version <- ApiKeys.ADD_PARTITIONS_TO_TXN.oldestVersion to ApiKeys.ADD_PARTITIONS_TO_TXN.latestVersion) {
+
+      EasyMock.reset(replicaManager, clientRequestQuotaManager, requestChannel, txnCoordinator)
+
+      val capturedResponse: Capture[RequestChannel.Response] = EasyMock.newCapture()
+      val responseCallback: Capture[Errors => Unit] = EasyMock.newCapture()
+
+      val transactionalId = "txnId"
+      val producerId = 15L
+      val epoch = 0.toShort
+
+      val partition = 1
+      val topicPartition = new TopicPartition(topic, partition)
+
+      val addPartitionsToTxnRequest = new AddPartitionsToTxnRequest.Builder(
+        transactionalId,
+        producerId,
+        epoch,
+        Collections.singletonList(topicPartition)
+      ).build(version.toShort)
+      val request = buildRequest(addPartitionsToTxnRequest)
+
+      EasyMock.expect(txnCoordinator.handleAddPartitionsToTransaction(
+        EasyMock.eq(transactionalId),
+        EasyMock.eq(producerId),
+        EasyMock.eq(epoch),
+        EasyMock.eq(Set(topicPartition)),
+
+        EasyMock.capture(responseCallback)
+      )).andAnswer(
+        () => responseCallback.getValue.apply(Errors.PRODUCER_FENCED))
+
+      EasyMock.expect(requestChannel.sendResponse(EasyMock.capture(capturedResponse)))
+
+      EasyMock.replay(replicaManager, clientRequestQuotaManager, requestChannel, txnCoordinator)
+
+      createKafkaApis().handleAddPartitionToTxnRequest(request)
+
+      val response = readResponse(addPartitionsToTxnRequest, capturedResponse)
+        .asInstanceOf[AddPartitionsToTxnResponse]
+
+      if (version < 2) {
+        assertEquals(Collections.singletonMap(topicPartition, Errors.INVALID_PRODUCER_EPOCH), response.errors())
+      } else {
+        assertEquals(Collections.singletonMap(topicPartition, Errors.PRODUCER_FENCED), response.errors())
+      }
+    }
+  }
+
+  @Test
+  def shouldReplaceProducerFencedWithInvalidProducerEpochInEndTxnWithOlderClient(): Unit = {
+    val topic = "topic"
+    setupBasicMetadataCache(topic, numPartitions = 2)
+
+    for (version <- ApiKeys.END_TXN.oldestVersion to ApiKeys.END_TXN.latestVersion) {
+
+      EasyMock.reset(replicaManager, clientRequestQuotaManager, requestChannel, txnCoordinator)
+
+      val capturedResponse: Capture[RequestChannel.Response] = EasyMock.newCapture()
+      val responseCallback: Capture[Errors => Unit]  = EasyMock.newCapture()
+
+      val transactionalId = "txnId"
+      val producerId = 15L
+      val epoch = 0.toShort
+
+      val endTxnRequest = new EndTxnRequest.Builder(
+        new EndTxnRequestData()
+          .setTransactionalId(transactionalId)
+          .setProducerId(producerId)
+          .setProducerEpoch(epoch)
+          .setCommitted(true)
+      ).build(version.toShort)
+      val request = buildRequest(endTxnRequest)
+
+      EasyMock.expect(txnCoordinator.handleEndTransaction(
+        EasyMock.eq(transactionalId),
+        EasyMock.eq(producerId),
+        EasyMock.eq(epoch),
+        EasyMock.eq(TransactionResult.COMMIT),
+        EasyMock.capture(responseCallback)
+      )).andAnswer(
+        () => responseCallback.getValue.apply(Errors.PRODUCER_FENCED))
+
+      EasyMock.expect(requestChannel.sendResponse(EasyMock.capture(capturedResponse)))
+
+      EasyMock.replay(replicaManager, clientRequestQuotaManager, requestChannel, txnCoordinator)
+
+      createKafkaApis().handleEndTxnRequest(request)
+
+      val response = readResponse(endTxnRequest, capturedResponse)
+        .asInstanceOf[EndTxnResponse]
+
+      if (version < 2) {
+        assertEquals(Errors.INVALID_PRODUCER_EPOCH.code, response.data.errorCode)
+      } else {
+        assertEquals(Errors.PRODUCER_FENCED.code, response.data.errorCode)
+      }
+    }
+  }
+
+  @nowarn("cat=deprecation")
+  @Test
+  def shouldReplaceProducerFencedWithInvalidProducerEpochInProduceResponse(): Unit = {
+    val topic = "topic"
+    setupBasicMetadataCache(topic, numPartitions = 2)
+
+    for (version <- ApiKeys.PRODUCE.oldestVersion to ApiKeys.PRODUCE.latestVersion) {
+
+      EasyMock.reset(replicaManager, clientQuotaManager, clientRequestQuotaManager, requestChannel, txnCoordinator)
+
+      val responseCallback: Capture[Map[TopicPartition, PartitionResponse] => Unit] = EasyMock.newCapture()
+
+      val tp = new TopicPartition("topic", 0)
+
+      val produceRequest = ProduceRequest.forCurrentMagic(new ProduceRequestData()
+        .setTopicData(new ProduceRequestData.TopicProduceDataCollection(
+          Collections.singletonList(new ProduceRequestData.TopicProduceData()
+            .setName(tp.topic()).setPartitionData(Collections.singletonList(
+            new ProduceRequestData.PartitionProduceData()
+              .setIndex(tp.partition())
+              .setRecords(MemoryRecords.withRecords(CompressionType.NONE, new SimpleRecord("test".getBytes))))))
+            .iterator))
+        .setAcks(1.toShort)
+        .setTimeoutMs(5000))
+        .build(version.toShort)
+      val request = buildRequest(produceRequest)
+
+      EasyMock.expect(replicaManager.appendRecords(EasyMock.anyLong(),
+        EasyMock.anyShort(),
+        EasyMock.eq(false),
+        EasyMock.eq(AppendOrigin.Client),
+        EasyMock.anyObject(),
+        EasyMock.capture(responseCallback),
+        EasyMock.anyObject(),
+        EasyMock.anyObject())
+      ).andAnswer(() => responseCallback.getValue.apply(Map(tp -> new PartitionResponse(Errors.INVALID_PRODUCER_EPOCH))))
+
+      val capturedResponse = expectNoThrottling()
+      EasyMock.expect(clientQuotaManager.maybeRecordAndGetThrottleTimeMs(
+        anyObject[RequestChannel.Request](), anyDouble, anyLong)).andReturn(0)
+
+      EasyMock.replay(replicaManager, clientQuotaManager, clientRequestQuotaManager, requestChannel, txnCoordinator)
+
+      createKafkaApis().handleProduceRequest(request)
+
+      val response = readResponse(produceRequest, capturedResponse)
+        .asInstanceOf[ProduceResponse]
+
+      assertEquals(1, response.responses().size())
+      for (partitionResponse <- response.responses().asScala) {
+        assertEquals(Errors.INVALID_PRODUCER_EPOCH, partitionResponse._2.error)
+      }
+    }
+  }
+
+  @Test
+  def testAddPartitionsToTxnWithInvalidPartition(): Unit = {
+    val topic = "topic"
+    setupBasicMetadataCache(topic, numPartitions = 1)
+
+    def checkInvalidPartition(invalidPartitionId: Int): Unit = {
+      EasyMock.reset(replicaManager, clientRequestQuotaManager, requestChannel)
+
+      val invalidTopicPartition = new TopicPartition(topic, invalidPartitionId)
+      val addPartitionsToTxnRequest = new AddPartitionsToTxnRequest.Builder(
+        "txnlId", 15L, 0.toShort, List(invalidTopicPartition).asJava
+      ).build()
+      val request = buildRequest(addPartitionsToTxnRequest)
+
+      val capturedResponse = expectNoThrottling()
+      EasyMock.replay(replicaManager, clientRequestQuotaManager, requestChannel)
+      createKafkaApis().handleAddPartitionToTxnRequest(request)
+
+      val response = readResponse(addPartitionsToTxnRequest, capturedResponse)
+        .asInstanceOf[AddPartitionsToTxnResponse]
+      assertEquals(Errors.UNKNOWN_TOPIC_OR_PARTITION, response.errors().get(invalidTopicPartition))
+    }
+
+    checkInvalidPartition(-1)
+    checkInvalidPartition(1) // topic has only one partition
   }
 
   @Test(expected = classOf[UnsupportedVersionException])
@@ -131,7 +1346,7 @@ class KafkaApisTest {
   @Test
   def shouldRespondWithUnsupportedForMessageFormatOnHandleWriteTxnMarkersWhenMagicLowerThanRequired(): Unit = {
     val topicPartition = new TopicPartition("t", 0)
-    val (writeTxnMarkersRequest, request) = createWriteTxnMarkersRequest(Utils.mkList(topicPartition))
+    val (writeTxnMarkersRequest, request) = createWriteTxnMarkersRequest(asList(topicPartition))
     val expectedErrors = Map(topicPartition -> Errors.UNSUPPORTED_FOR_MESSAGE_FORMAT).asJava
     val capturedResponse: Capture[RequestChannel.Response] = EasyMock.newCapture()
 
@@ -142,15 +1357,15 @@ class KafkaApisTest {
 
     createKafkaApis().handleWriteTxnMarkersRequest(request)
 
-    val markersResponse = readResponse(ApiKeys.WRITE_TXN_MARKERS, writeTxnMarkersRequest, capturedResponse)
+    val markersResponse = readResponse(writeTxnMarkersRequest, capturedResponse)
       .asInstanceOf[WriteTxnMarkersResponse]
-    assertEquals(expectedErrors, markersResponse.errors(1))
+    assertEquals(expectedErrors, markersResponse.errorsByProducerId.get(1L))
   }
 
   @Test
   def shouldRespondWithUnknownTopicWhenPartitionIsNotHosted(): Unit = {
     val topicPartition = new TopicPartition("t", 0)
-    val (writeTxnMarkersRequest, request) = createWriteTxnMarkersRequest(Utils.mkList(topicPartition))
+    val (writeTxnMarkersRequest, request) = createWriteTxnMarkersRequest(asList(topicPartition))
     val expectedErrors = Map(topicPartition -> Errors.UNKNOWN_TOPIC_OR_PARTITION).asJava
     val capturedResponse: Capture[RequestChannel.Response] = EasyMock.newCapture()
 
@@ -161,20 +1376,20 @@ class KafkaApisTest {
 
     createKafkaApis().handleWriteTxnMarkersRequest(request)
 
-    val markersResponse = readResponse(ApiKeys.WRITE_TXN_MARKERS, writeTxnMarkersRequest, capturedResponse)
+    val markersResponse = readResponse(writeTxnMarkersRequest, capturedResponse)
       .asInstanceOf[WriteTxnMarkersResponse]
-    assertEquals(expectedErrors, markersResponse.errors(1))
+    assertEquals(expectedErrors, markersResponse.errorsByProducerId.get(1L))
   }
 
   @Test
   def shouldRespondWithUnsupportedMessageFormatForBadPartitionAndNoErrorsForGoodPartition(): Unit = {
     val tp1 = new TopicPartition("t", 0)
     val tp2 = new TopicPartition("t1", 0)
-    val (writeTxnMarkersRequest, request) = createWriteTxnMarkersRequest(Utils.mkList(tp1, tp2))
+    val (writeTxnMarkersRequest, request) = createWriteTxnMarkersRequest(asList(tp1, tp2))
     val expectedErrors = Map(tp1 -> Errors.UNSUPPORTED_FOR_MESSAGE_FORMAT, tp2 -> Errors.NONE).asJava
 
     val capturedResponse: Capture[RequestChannel.Response] = EasyMock.newCapture()
-    val responseCallback: Capture[Map[TopicPartition, PartitionResponse] => Unit]  = EasyMock.newCapture()
+    val responseCallback: Capture[Map[TopicPartition, PartitionResponse] => Unit] = EasyMock.newCapture()
 
     EasyMock.expect(replicaManager.getMagic(tp1))
       .andReturn(Some(RecordBatch.MAGIC_VALUE_V1))
@@ -184,36 +1399,134 @@ 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(),
-      EasyMock.anyObject())).andAnswer(new IAnswer[Unit] {
-      override def answer(): Unit = {
-        responseCallback.getValue.apply(Map(tp2 -> new PartitionResponse(Errors.NONE)))
+      EasyMock.anyObject())
+    ).andAnswer(() => responseCallback.getValue.apply(Map(tp2 -> new PartitionResponse(Errors.NONE))))
+
+    EasyMock.expect(requestChannel.sendResponse(EasyMock.capture(capturedResponse)))
+    EasyMock.replay(replicaManager, replicaQuotaManager, requestChannel)
+
+    createKafkaApis().handleWriteTxnMarkersRequest(request)
+
+    val markersResponse = readResponse(writeTxnMarkersRequest, capturedResponse)
+      .asInstanceOf[WriteTxnMarkersResponse]
+    assertEquals(expectedErrors, markersResponse.errorsByProducerId.get(1L))
+    EasyMock.verify(replicaManager)
+  }
+
+  @Test
+  def shouldResignCoordinatorsIfStopReplicaReceivedWithDeleteFlagAndLeaderEpoch(): Unit = {
+    shouldResignCoordinatorsIfStopReplicaReceivedWithDeleteFlag(
+      LeaderAndIsr.initialLeaderEpoch + 2, deletePartition = true)
+  }
+
+  @Test
+  def shouldResignCoordinatorsIfStopReplicaReceivedWithDeleteFlagAndDeleteSentinel(): Unit = {
+    shouldResignCoordinatorsIfStopReplicaReceivedWithDeleteFlag(
+      LeaderAndIsr.EpochDuringDelete, deletePartition = true)
+  }
+
+  @Test
+  def shouldResignCoordinatorsIfStopReplicaReceivedWithDeleteFlagAndNoEpochSentinel(): Unit = {
+    shouldResignCoordinatorsIfStopReplicaReceivedWithDeleteFlag(
+      LeaderAndIsr.NoEpoch, deletePartition = true)
+  }
+
+  @Test
+  def shouldNotResignCoordinatorsIfStopReplicaReceivedWithoutDeleteFlag(): Unit = {
+    shouldResignCoordinatorsIfStopReplicaReceivedWithDeleteFlag(
+      LeaderAndIsr.initialLeaderEpoch + 2, deletePartition = false)
+  }
+
+  def shouldResignCoordinatorsIfStopReplicaReceivedWithDeleteFlag(leaderEpoch: Int,
+                                                                  deletePartition: Boolean): Unit = {
+    val controllerId = 0
+    val controllerEpoch = 5
+    val brokerEpoch = 230498320L
+
+    val fooPartition = new TopicPartition("foo", 0)
+    val groupMetadataPartition = new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, 0)
+    val txnStatePartition = new TopicPartition(Topic.TRANSACTION_STATE_TOPIC_NAME, 0)
+
+    val topicStates = Seq(
+      new StopReplicaTopicState()
+        .setTopicName(groupMetadataPartition.topic())
+        .setPartitionStates(Seq(new StopReplicaPartitionState()
+          .setPartitionIndex(groupMetadataPartition.partition())
+          .setLeaderEpoch(leaderEpoch)
+          .setDeletePartition(deletePartition)).asJava),
+      new StopReplicaTopicState()
+        .setTopicName(txnStatePartition.topic())
+        .setPartitionStates(Seq(new StopReplicaPartitionState()
+          .setPartitionIndex(txnStatePartition.partition())
+          .setLeaderEpoch(leaderEpoch)
+          .setDeletePartition(deletePartition)).asJava),
+      new StopReplicaTopicState()
+        .setTopicName(fooPartition.topic())
+        .setPartitionStates(Seq(new StopReplicaPartitionState()
+          .setPartitionIndex(fooPartition.partition())
+          .setLeaderEpoch(leaderEpoch)
+          .setDeletePartition(deletePartition)).asJava)
+    ).asJava
+
+    val stopReplicaRequest = new StopReplicaRequest.Builder(
+      ApiKeys.STOP_REPLICA.latestVersion,
+      controllerId,
+      controllerEpoch,
+      brokerEpoch,
+      false,
+      topicStates
+    ).build()
+    val request = buildRequest(stopReplicaRequest)
+
+    EasyMock.expect(replicaManager.stopReplicas(
+      EasyMock.eq(request.context.correlationId),
+      EasyMock.eq(controllerId),
+      EasyMock.eq(controllerEpoch),
+      EasyMock.eq(brokerEpoch),
+      EasyMock.eq(stopReplicaRequest.partitionStates().asScala)
+    )).andReturn(
+      (mutable.Map(
+        groupMetadataPartition -> Errors.NONE,
+        txnStatePartition -> Errors.NONE,
+        fooPartition -> Errors.NONE
+      ), Errors.NONE)
+    )
+    EasyMock.expect(controller.brokerEpoch).andStubReturn(brokerEpoch)
+
+    if (deletePartition) {
+      if (leaderEpoch >= 0) {
+        txnCoordinator.onResignation(txnStatePartition.partition, Some(leaderEpoch))
+      } else {
+        txnCoordinator.onResignation(txnStatePartition.partition, None)
       }
-    })
+      EasyMock.expectLastCall()
+    }
 
-    EasyMock.expect(requestChannel.sendResponse(EasyMock.capture(capturedResponse)))
-    EasyMock.replay(replicaManager, replicaQuotaManager, requestChannel)
+    if (deletePartition) {
+      groupCoordinator.onResignation(groupMetadataPartition.partition)
+      EasyMock.expectLastCall()
+    }
 
-    createKafkaApis().handleWriteTxnMarkersRequest(request)
+    EasyMock.replay(controller, replicaManager, txnCoordinator, groupCoordinator)
 
-    val markersResponse = readResponse(ApiKeys.WRITE_TXN_MARKERS, writeTxnMarkersRequest, capturedResponse)
-      .asInstanceOf[WriteTxnMarkersResponse]
-    assertEquals(expectedErrors, markersResponse.errors(1))
-    EasyMock.verify(replicaManager)
+    createKafkaApis().handleStopReplicaRequest(request)
+
+    EasyMock.verify(txnCoordinator, groupCoordinator)
   }
 
   @Test
   def shouldRespondWithUnknownTopicOrPartitionForBadPartitionAndNoErrorsForGoodPartition(): Unit = {
     val tp1 = new TopicPartition("t", 0)
     val tp2 = new TopicPartition("t1", 0)
-    val (writeTxnMarkersRequest, request) = createWriteTxnMarkersRequest(Utils.mkList(tp1, tp2))
+    val (writeTxnMarkersRequest, request) = createWriteTxnMarkersRequest(asList(tp1, tp2))
     val expectedErrors = Map(tp1 -> Errors.UNKNOWN_TOPIC_OR_PARTITION, tp2 -> Errors.NONE).asJava
 
     val capturedResponse: Capture[RequestChannel.Response] = EasyMock.newCapture()
-    val responseCallback: Capture[Map[TopicPartition, PartitionResponse] => Unit]  = EasyMock.newCapture()
+    val responseCallback: Capture[Map[TopicPartition, PartitionResponse] => Unit] = EasyMock.newCapture()
 
     EasyMock.expect(replicaManager.getMagic(tp1))
       .andReturn(None)
@@ -223,38 +1536,35 @@ 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(),
-      EasyMock.anyObject())).andAnswer(new IAnswer[Unit] {
-      override def answer(): Unit = {
-        responseCallback.getValue.apply(Map(tp2 -> new PartitionResponse(Errors.NONE)))
-      }
-    })
+      EasyMock.anyObject())
+    ).andAnswer(() => responseCallback.getValue.apply(Map(tp2 -> new PartitionResponse(Errors.NONE))))
 
     EasyMock.expect(requestChannel.sendResponse(EasyMock.capture(capturedResponse)))
     EasyMock.replay(replicaManager, replicaQuotaManager, requestChannel)
 
     createKafkaApis().handleWriteTxnMarkersRequest(request)
 
-    val markersResponse = readResponse(ApiKeys.WRITE_TXN_MARKERS, writeTxnMarkersRequest, capturedResponse)
+    val markersResponse = readResponse(writeTxnMarkersRequest, capturedResponse)
       .asInstanceOf[WriteTxnMarkersResponse]
-    assertEquals(expectedErrors, markersResponse.errors(1))
+    assertEquals(expectedErrors, markersResponse.errorsByProducerId.get(1L))
     EasyMock.verify(replicaManager)
   }
 
   @Test
   def shouldAppendToLogOnWriteTxnMarkersWhenCorrectMagicVersion(): Unit = {
     val topicPartition = new TopicPartition("t", 0)
-    val request = createWriteTxnMarkersRequest(Utils.mkList(topicPartition))._2
+    val request = createWriteTxnMarkersRequest(asList(topicPartition))._2
     EasyMock.expect(replicaManager.getMagic(topicPartition))
       .andReturn(Some(RecordBatch.MAGIC_VALUE_V2))
 
     EasyMock.expect(replicaManager.appendRecords(EasyMock.anyLong(),
       EasyMock.anyShort(),
       EasyMock.eq(true),
-      EasyMock.eq(false),
+      EasyMock.eq(AppendOrigin.Coordinator),
       EasyMock.anyObject(),
       EasyMock.anyObject(),
       EasyMock.anyObject(),
@@ -267,173 +1577,1472 @@ class KafkaApisTest {
   }
 
   @Test
-  def testReadUncommittedConsumerListOffsetLimitedAtHighWatermark(): Unit = {
-    testConsumerListOffsetLimit(IsolationLevel.READ_UNCOMMITTED)
+  def testLeaderReplicaIfLocalRaisesFencedLeaderEpoch(): Unit = {
+    testListOffsetFailedGetLeaderReplica(Errors.FENCED_LEADER_EPOCH)
   }
 
   @Test
-  def testReadCommittedConsumerListOffsetLimitedAtLastStableOffset(): Unit = {
-    testConsumerListOffsetLimit(IsolationLevel.READ_COMMITTED)
+  def testLeaderReplicaIfLocalRaisesUnknownLeaderEpoch(): Unit = {
+    testListOffsetFailedGetLeaderReplica(Errors.UNKNOWN_LEADER_EPOCH)
   }
 
-  private def testConsumerListOffsetLimit(isolationLevel: IsolationLevel): Unit = {
-    val tp = new TopicPartition("foo", 0)
-    val timestamp: JLong = time.milliseconds()
-    val limitOffset = 15L
+  @Test
+  def testLeaderReplicaIfLocalRaisesNotLeaderOrFollower(): Unit = {
+    testListOffsetFailedGetLeaderReplica(Errors.NOT_LEADER_OR_FOLLOWER)
+  }
 
-    val capturedResponse = EasyMock.newCapture[RequestChannel.Response]()
-    val capturedThrottleCallback = EasyMock.newCapture[Int => Unit]()
-    val replica = EasyMock.mock(classOf[Replica])
-    val log = EasyMock.mock(classOf[Log])
-    EasyMock.expect(replicaManager.getLeaderReplicaIfLocal(tp)).andReturn(replica)
-    if (isolationLevel == IsolationLevel.READ_UNCOMMITTED)
-      EasyMock.expect(replica.highWatermark).andReturn(LogOffsetMetadata(messageOffset = limitOffset))
-    else
-      EasyMock.expect(replica.lastStableOffset).andReturn(LogOffsetMetadata(messageOffset = limitOffset))
-    EasyMock.expect(replicaManager.getLog(tp)).andReturn(Some(log))
-    EasyMock.expect(log.fetchOffsetsByTimestamp(timestamp)).andReturn(Some(TimestampOffset(timestamp = timestamp, offset = limitOffset)))
-    expectThrottleCallbackAndInvoke(capturedThrottleCallback)
-    EasyMock.expect(requestChannel.sendResponse(EasyMock.capture(capturedResponse)))
-    EasyMock.replay(replicaManager, clientRequestQuotaManager, requestChannel, replica, log)
+  @Test
+  def testLeaderReplicaIfLocalRaisesUnknownTopicOrPartition(): Unit = {
+    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 = new DescribeGroupsRequest.Builder(
+      new DescribeGroupsRequestData().setGroups(List(groupId).asJava)
+    ).build()
+    val request = buildRequest(describeGroupsRequest)
+
+    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(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 testOffsetDelete(): Unit = {
+    val group = "groupId"
+    setupBasicMetadataCache("topic-1", numPartitions = 2)
+    setupBasicMetadataCache("topic-2", numPartitions = 2)
+
+    EasyMock.reset(groupCoordinator, replicaManager, clientRequestQuotaManager, requestChannel)
+
+    val topics = new OffsetDeleteRequestTopicCollection()
+    topics.add(new OffsetDeleteRequestTopic()
+      .setName("topic-1")
+      .setPartitions(Seq(
+        new OffsetDeleteRequestPartition().setPartitionIndex(0),
+        new OffsetDeleteRequestPartition().setPartitionIndex(1)).asJava))
+    topics.add(new OffsetDeleteRequestTopic()
+      .setName("topic-2")
+      .setPartitions(Seq(
+        new OffsetDeleteRequestPartition().setPartitionIndex(0),
+        new OffsetDeleteRequestPartition().setPartitionIndex(1)).asJava))
+
+    val offsetDeleteRequest = new OffsetDeleteRequest.Builder(
+      new OffsetDeleteRequestData()
+        .setGroupId(group)
+        .setTopics(topics)
+    ).build()
+    val request = buildRequest(offsetDeleteRequest)
+
+    val capturedResponse = expectNoThrottling()
+    EasyMock.expect(groupCoordinator.handleDeleteOffsets(
+      EasyMock.eq(group),
+      EasyMock.eq(Seq(
+        new TopicPartition("topic-1", 0),
+        new TopicPartition("topic-1", 1),
+        new TopicPartition("topic-2", 0),
+        new TopicPartition("topic-2", 1)
+      ))
+    )).andReturn((Errors.NONE, Map(
+      new TopicPartition("topic-1", 0) -> Errors.NONE,
+      new TopicPartition("topic-1", 1) -> Errors.NONE,
+      new TopicPartition("topic-2", 0) -> Errors.NONE,
+      new TopicPartition("topic-2", 1) -> Errors.NONE,
+    )))
+
+    EasyMock.replay(groupCoordinator, replicaManager, clientRequestQuotaManager, requestChannel)
+
+    createKafkaApis().handleOffsetDeleteRequest(request)
+
+    val response = readResponse(offsetDeleteRequest, capturedResponse)
+      .asInstanceOf[OffsetDeleteResponse]
+
+    def errorForPartition(topic: String, partition: Int): Errors = {
+      Errors.forCode(response.data.topics.find(topic).partitions.find(partition).errorCode())
+    }
+
+    assertEquals(2, response.data.topics.size)
+    assertEquals(Errors.NONE, errorForPartition("topic-1", 0))
+    assertEquals(Errors.NONE, errorForPartition("topic-1", 1))
+    assertEquals(Errors.NONE, errorForPartition("topic-2", 0))
+    assertEquals(Errors.NONE, errorForPartition("topic-2", 1))
+  }
+
+  @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 = new OffsetDeleteRequest.Builder(
+        new OffsetDeleteRequestData()
+          .setGroupId(group)
+          .setTopics(topics)
+      ).build()
+      val request = buildRequest(offsetDeleteRequest)
+
+      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(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 = new OffsetDeleteRequest.Builder(
+      new OffsetDeleteRequestData()
+        .setGroupId(group)
+    ).build()
+    val request = buildRequest(offsetDeleteRequest)
+
+    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(offsetDeleteRequest, capturedResponse)
+      .asInstanceOf[OffsetDeleteResponse]
+
+    assertEquals(Errors.GROUP_ID_NOT_FOUND, Errors.forCode(response.data.errorCode()))
+  }
 
-    val builder = ListOffsetRequest.Builder.forConsumer(true, isolationLevel)
-      .setTargetTimes(Map(tp -> timestamp).asJava)
-    val (listOffsetRequest, request) = buildRequest(builder)
+  private def testListOffsetFailedGetLeaderReplica(error: Errors): Unit = {
+    val tp = new TopicPartition("foo", 0)
+    val isolationLevel = IsolationLevel.READ_UNCOMMITTED
+    val currentLeaderEpoch = Optional.of[Integer](15)
+
+    EasyMock.expect(replicaManager.fetchOffsetForTimestamp(
+      EasyMock.eq(tp),
+      EasyMock.eq(ListOffsetRequest.EARLIEST_TIMESTAMP),
+      EasyMock.eq(Some(isolationLevel)),
+      EasyMock.eq(currentLeaderEpoch),
+      fetchOnlyFromLeader = EasyMock.eq(true))
+    ).andThrow(error.exception)
+
+    val capturedResponse = expectNoThrottling()
+    EasyMock.replay(replicaManager, clientRequestQuotaManager, requestChannel)
+
+    val targetTimes = List(new ListOffsetTopic()
+      .setName(tp.topic)
+      .setPartitions(List(new ListOffsetPartition()
+        .setPartitionIndex(tp.partition)
+        .setTimestamp(ListOffsetRequest.EARLIEST_TIMESTAMP)
+        .setCurrentLeaderEpoch(currentLeaderEpoch.get)).asJava)).asJava
+    val listOffsetRequest = ListOffsetRequest.Builder.forConsumer(true, isolationLevel)
+      .setTargetTimes(targetTimes).build()
+    val request = buildRequest(listOffsetRequest)
     createKafkaApis().handleListOffsetRequest(request)
 
-    val response = readResponse(ApiKeys.LIST_OFFSETS, listOffsetRequest, capturedResponse).asInstanceOf[ListOffsetResponse]
-    assertTrue(response.responseData.containsKey(tp))
+    val response = readResponse(listOffsetRequest, capturedResponse)
+      .asInstanceOf[ListOffsetResponse]
+    val partitionDataOptional = response.topics.asScala.find(_.name == tp.topic).get
+      .partitions.asScala.find(_.partitionIndex == tp.partition)
+    assertTrue(partitionDataOptional.isDefined)
 
-    val partitionData = response.responseData.get(tp)
-    assertEquals(Errors.NONE, partitionData.error)
+    val partitionData = partitionDataOptional.get
+    assertEquals(error.code, partitionData.errorCode)
     assertEquals(ListOffsetResponse.UNKNOWN_OFFSET, partitionData.offset)
     assertEquals(ListOffsetResponse.UNKNOWN_TIMESTAMP, partitionData.timestamp)
   }
 
   @Test
-  def testReadUncommittedConsumerListOffsetEarliestOffsetEqualsHighWatermark(): Unit = {
-    testConsumerListOffsetEarliestOffsetEqualsLimit(IsolationLevel.READ_UNCOMMITTED)
+  def testReadUncommittedConsumerListOffsetLatest(): Unit = {
+    testConsumerListOffsetLatest(IsolationLevel.READ_UNCOMMITTED)
+  }
+
+  @Test
+  def testReadCommittedConsumerListOffsetLatest(): Unit = {
+    testConsumerListOffsetLatest(IsolationLevel.READ_COMMITTED)
   }
 
+  /**
+   * Verifies that the metadata response is correct if the broker listeners are inconsistent (i.e. one broker has
+   * more listeners than another) and the request is sent on the listener that exists in both brokers.
+   */
   @Test
-  def testReadCommittedConsumerListOffsetEarliestOffsetEqualsLastStableOffset(): Unit = {
-    testConsumerListOffsetEarliestOffsetEqualsLimit(IsolationLevel.READ_COMMITTED)
+  def testMetadataRequestOnSharedListenerWithInconsistentListenersAcrossBrokers(): Unit = {
+    val (plaintextListener, _) = updateMetadataCacheWithInconsistentListeners()
+    val response = sendMetadataRequestWithInconsistentListeners(plaintextListener)
+    assertEquals(Set(0, 1), response.brokers.asScala.map(_.id).toSet)
   }
 
-  private def testConsumerListOffsetEarliestOffsetEqualsLimit(isolationLevel: IsolationLevel): Unit = {
-    val tp = new TopicPartition("foo", 0)
-    val limitOffset = 15L
+  /**
+   * Verifies that the metadata response is correct if the broker listeners are inconsistent (i.e. one broker has
+   * more listeners than another) and the request is sent on the listener that exists in one broker.
+   */
+  @Test
+  def testMetadataRequestOnDistinctListenerWithInconsistentListenersAcrossBrokers(): Unit = {
+    val (_, anotherListener) = updateMetadataCacheWithInconsistentListeners()
+    val response = sendMetadataRequestWithInconsistentListeners(anotherListener)
+    assertEquals(Set(0), response.brokers.asScala.map(_.id).toSet)
+  }
 
-    val capturedResponse = EasyMock.newCapture[RequestChannel.Response]()
-    val capturedThrottleCallback = EasyMock.newCapture[Int => Unit]()
-    val replica = EasyMock.mock(classOf[Replica])
-    val log = EasyMock.mock(classOf[Log])
-    EasyMock.expect(replicaManager.getLeaderReplicaIfLocal(tp)).andReturn(replica)
-    if (isolationLevel == IsolationLevel.READ_UNCOMMITTED)
-      EasyMock.expect(replica.highWatermark).andReturn(LogOffsetMetadata(messageOffset = limitOffset))
-    else
-      EasyMock.expect(replica.lastStableOffset).andReturn(LogOffsetMetadata(messageOffset = limitOffset))
-    EasyMock.expect(replicaManager.getLog(tp)).andReturn(Some(log))
-    EasyMock.expect(log.fetchOffsetsByTimestamp(ListOffsetRequest.EARLIEST_TIMESTAMP))
-      .andReturn(Some(TimestampOffset(timestamp = ListOffsetResponse.UNKNOWN_TIMESTAMP, offset = limitOffset)))
-    expectThrottleCallbackAndInvoke(capturedThrottleCallback)
-    EasyMock.expect(requestChannel.sendResponse(EasyMock.capture(capturedResponse)))
-    EasyMock.replay(replicaManager, clientRequestQuotaManager, requestChannel, replica, log)
 
-    val builder = ListOffsetRequest.Builder.forConsumer(true, isolationLevel)
-      .setTargetTimes(Map(tp -> (ListOffsetRequest.EARLIEST_TIMESTAMP: JLong)).asJava)
-    val (listOffsetRequest, request) = buildRequest(builder)
-    createKafkaApis().handleListOffsetRequest(request)
+  /**
+   * Metadata request to fetch all topics should not result in the followings:
+   * 1) Auto topic creation
+   * 2) UNKNOWN_TOPIC_OR_PARTITION
+   *
+   * This case is testing the case that a topic is being deleted from MetadataCache right after
+   * authorization but before checking in MetadataCache.
+   */
+  @Test
+  def getAllTopicMetadataShouldNotCreateTopicOrReturnUnknownTopicPartition(): Unit = {
+    // Setup: authorizer authorizes 2 topics, but one got deleted in metadata cache
+    metadataCache =
+      EasyMock.partialMockBuilder(classOf[MetadataCache])
+        .withConstructor(classOf[Int])
+        .withArgs(Int.box(brokerId))  // Need to box it for Scala 2.12 and before
+        .addMockedMethod("getAllTopics")
+        .addMockedMethod("getTopicMetadata")
+        .createMock()
+
+    // 2 topics returned for authorization in during handle
+    val topicsReturnedFromMetadataCacheForAuthorization = Set("remaining-topic", "later-deleted-topic")
+    expect(metadataCache.getAllTopics()).andReturn(topicsReturnedFromMetadataCacheForAuthorization).once()
+    // 1 topic is deleted from metadata right at the time between authorization and the next getTopicMetadata() call
+    expect(metadataCache.getTopicMetadata(
+      EasyMock.eq(topicsReturnedFromMetadataCacheForAuthorization),
+      anyObject[ListenerName],
+      anyBoolean,
+      anyBoolean
+    )).andStubReturn(Seq(
+      new MetadataResponseTopic()
+        .setErrorCode(Errors.NONE.code)
+        .setName("remaining-topic")
+        .setIsInternal(false)
+    ))
+
+    EasyMock.replay(metadataCache)
+
+    var createTopicIsCalled: Boolean = false;
+    // Specific mock on zkClient for this use case
+    // Expect it's never called to do auto topic creation
+    expect(zkClient.setOrCreateEntityConfigs(
+      EasyMock.eq(ConfigType.Topic),
+      EasyMock.anyString,
+      EasyMock.anyObject[Properties]
+    )).andStubAnswer(() => {
+      createTopicIsCalled = true
+    })
+    // No need to use
+    expect(zkClient.getAllBrokersInCluster)
+      .andStubReturn(Seq(new Broker(
+        brokerId, "localhost", 9902,
+        ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT), SecurityProtocol.PLAINTEXT
+      )))
+
+    EasyMock.replay(zkClient)
+
+    val (requestListener, _) = updateMetadataCacheWithInconsistentListeners()
+    val response = sendMetadataRequestWithInconsistentListeners(requestListener)
+
+    assertFalse(createTopicIsCalled)
+    assertEquals(List("remaining-topic"), response.topicMetadata().asScala.map { metadata => metadata.topic() })
+    assertTrue(response.topicsByError(Errors.UNKNOWN_TOPIC_OR_PARTITION).isEmpty)
+  }
 
-    val response = readResponse(ApiKeys.LIST_OFFSETS, listOffsetRequest, capturedResponse).asInstanceOf[ListOffsetResponse]
+  /**
+   * Verifies that sending a fetch request with version 9 works correctly when
+   * ReplicaManager.getLogConfig returns None.
+   */
+  @Test
+  def testFetchRequestV9WithNoLogConfig(): Unit = {
+    val tp = new TopicPartition("foo", 0)
+    setupBasicMetadataCache(tp.topic, numPartitions = 1)
+    val hw = 3
+    val timestamp = 1000
+
+    expect(replicaManager.getLogConfig(EasyMock.eq(tp))).andReturn(None)
+
+    replicaManager.fetchMessages(anyLong, anyInt, anyInt, anyInt, anyBoolean,
+      anyObject[Seq[(TopicPartition, FetchRequest.PartitionData)]], anyObject[ReplicaQuota],
+      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 -> FetchPartitionData(Errors.NONE, hw, 0, records,
+          None, None, None, Option.empty, isReassignmentFetch = false)))
+      }
+    })
+
+    val fetchData = Map(tp -> new FetchRequest.PartitionData(0, 0, 1000,
+      Optional.empty())).asJava
+    val fetchMetadata = new JFetchMetadata(0, 0)
+    val fetchContext = new FullFetchContext(time, new FetchSessionCache(1000, 100),
+      fetchMetadata, fetchData, false)
+    expect(fetchManager.newContext(anyObject[JFetchMetadata],
+      anyObject[util.Map[TopicPartition, FetchRequest.PartitionData]],
+      anyObject[util.List[TopicPartition]],
+      anyBoolean)).andReturn(fetchContext)
+
+    val capturedResponse = expectNoThrottling()
+    EasyMock.expect(clientQuotaManager.maybeRecordAndGetThrottleTimeMs(
+      anyObject[RequestChannel.Request](), anyDouble, anyLong)).andReturn(0)
+
+    EasyMock.replay(replicaManager, clientQuotaManager, clientRequestQuotaManager, requestChannel, fetchManager)
+
+    val fetchRequest = new FetchRequest.Builder(9, 9, -1, 100, 0, fetchData)
+      .build()
+    val request = buildRequest(fetchRequest)
+    createKafkaApis().handleFetchRequest(request)
+
+    val response = readResponse(fetchRequest, capturedResponse)
+      .asInstanceOf[FetchResponse[BaseRecords]]
     assertTrue(response.responseData.containsKey(tp))
 
     val partitionData = response.responseData.get(tp)
     assertEquals(Errors.NONE, partitionData.error)
-    assertEquals(limitOffset, partitionData.offset)
-    assertEquals(ListOffsetResponse.UNKNOWN_TIMESTAMP, partitionData.timestamp)
+    assertEquals(hw, partitionData.highWatermark)
+    assertEquals(-1, partitionData.lastStableOffset)
+    assertEquals(0, partitionData.logStartOffset)
+    assertEquals(timestamp,
+      partitionData.records.asInstanceOf[MemoryRecords].batches.iterator.next.maxTimestamp)
+    assertNull(partitionData.abortedTransactions)
   }
 
   @Test
-  def testReadUncommittedConsumerListOffsetLatest(): Unit = {
-    testConsumerListOffsetLatest(IsolationLevel.READ_UNCOMMITTED)
+  def testJoinGroupProtocolsOrder(): Unit = {
+    val protocols = List(
+      ("first", "first".getBytes()),
+      ("second", "second".getBytes())
+    )
+
+    val groupId = "group"
+    val memberId = "member1"
+    val protocolType = "consumer"
+    val rebalanceTimeoutMs = 10
+    val sessionTimeoutMs = 5
+    val capturedProtocols = EasyMock.newCapture[List[(String, Array[Byte])]]()
+
+    EasyMock.expect(groupCoordinator.handleJoinGroup(
+      EasyMock.eq(groupId),
+      EasyMock.eq(memberId),
+      EasyMock.eq(None),
+      EasyMock.eq(true),
+      EasyMock.eq(clientId),
+      EasyMock.eq(InetAddress.getLocalHost.toString),
+      EasyMock.eq(rebalanceTimeoutMs),
+      EasyMock.eq(sessionTimeoutMs),
+      EasyMock.eq(protocolType),
+      EasyMock.capture(capturedProtocols),
+      anyObject()
+    ))
+
+    EasyMock.replay(groupCoordinator)
+
+    createKafkaApis().handleJoinGroupRequest(
+      buildRequest(
+        new JoinGroupRequest.Builder(
+          new JoinGroupRequestData()
+            .setGroupId(groupId)
+            .setMemberId(memberId)
+            .setProtocolType(protocolType)
+            .setRebalanceTimeoutMs(rebalanceTimeoutMs)
+            .setSessionTimeoutMs(sessionTimeoutMs)
+            .setProtocols(new JoinGroupRequestData.JoinGroupRequestProtocolCollection(
+              protocols.map { case (name, protocol) => new JoinGroupRequestProtocol()
+                .setName(name).setMetadata(protocol)
+              }.iterator.asJava))
+        ).build()
+      ))
+
+    EasyMock.verify(groupCoordinator)
+
+    val capturedProtocolsList = capturedProtocols.getValue
+    assertEquals(protocols.size, capturedProtocolsList.size)
+    protocols.zip(capturedProtocolsList).foreach { case ((expectedName, expectedBytes), (name, bytes)) =>
+      assertEquals(expectedName, name)
+      assertArrayEquals(expectedBytes, bytes)
+    }
   }
 
   @Test
-  def testReadCommittedConsumerListOffsetLatest(): Unit = {
-    testConsumerListOffsetLatest(IsolationLevel.READ_COMMITTED)
+  def testJoinGroupWhenAnErrorOccurs(): Unit = {
+    for (version <- ApiKeys.JOIN_GROUP.oldestVersion to ApiKeys.JOIN_GROUP.latestVersion) {
+      testJoinGroupWhenAnErrorOccurs(version.asInstanceOf[Short])
+    }
   }
 
-  private def testConsumerListOffsetLatest(isolationLevel: IsolationLevel): Unit = {
-    val tp = new TopicPartition("foo", 0)
-    val latestOffset = 15L
+  def testJoinGroupWhenAnErrorOccurs(version: Short): Unit = {
+    EasyMock.reset(groupCoordinator, clientRequestQuotaManager, requestChannel)
+
+    val capturedResponse = expectNoThrottling()
+
+    val groupId = "group"
+    val memberId = "member1"
+    val protocolType = "consumer"
+    val rebalanceTimeoutMs = 10
+    val sessionTimeoutMs = 5
+
+    val capturedCallback = EasyMock.newCapture[JoinGroupCallback]()
+
+    EasyMock.expect(groupCoordinator.handleJoinGroup(
+      EasyMock.eq(groupId),
+      EasyMock.eq(memberId),
+      EasyMock.eq(None),
+      EasyMock.eq(if (version >= 4) true else false),
+      EasyMock.eq(clientId),
+      EasyMock.eq(InetAddress.getLocalHost.toString),
+      EasyMock.eq(if (version >= 1) rebalanceTimeoutMs else sessionTimeoutMs),
+      EasyMock.eq(sessionTimeoutMs),
+      EasyMock.eq(protocolType),
+      EasyMock.eq(List.empty),
+      EasyMock.capture(capturedCallback)
+    ))
+
+    val joinGroupRequest = new JoinGroupRequest.Builder(
+      new JoinGroupRequestData()
+        .setGroupId(groupId)
+        .setMemberId(memberId)
+        .setProtocolType(protocolType)
+        .setRebalanceTimeoutMs(rebalanceTimeoutMs)
+        .setSessionTimeoutMs(sessionTimeoutMs)
+    ).build(version)
+
+    val requestChannelRequest = buildRequest(joinGroupRequest)
+
+    EasyMock.replay(groupCoordinator, clientRequestQuotaManager, requestChannel)
+
+    createKafkaApis().handleJoinGroupRequest(requestChannelRequest)
+
+    EasyMock.verify(groupCoordinator)
+
+    capturedCallback.getValue.apply(JoinGroupResult(memberId, Errors.INCONSISTENT_GROUP_PROTOCOL))
+
+    val response = readResponse(joinGroupRequest, capturedResponse)
+      .asInstanceOf[JoinGroupResponse]
+
+    assertEquals(Errors.INCONSISTENT_GROUP_PROTOCOL, response.error)
+    assertEquals(0, response.data.members.size)
+    assertEquals(memberId, response.data.memberId)
+    assertEquals(GroupCoordinator.NoGeneration, response.data.generationId)
+    assertEquals(GroupCoordinator.NoLeader, response.data.leader)
+    assertNull(response.data.protocolType)
+
+    if (version >= 7) {
+      assertNull(response.data.protocolName)
+    } else {
+      assertEquals(GroupCoordinator.NoProtocol, response.data.protocolName)
+    }
+
+    EasyMock.verify(clientRequestQuotaManager, requestChannel)
+  }
 
-    val capturedResponse = EasyMock.newCapture[RequestChannel.Response]()
-    val capturedThrottleCallback = EasyMock.newCapture[Int => Unit]()
-    val replica = EasyMock.mock(classOf[Replica])
-    val log = EasyMock.mock(classOf[Log])
-    EasyMock.expect(replicaManager.getLeaderReplicaIfLocal(tp)).andReturn(replica)
-    if (isolationLevel == IsolationLevel.READ_UNCOMMITTED)
-      EasyMock.expect(replica.highWatermark).andReturn(LogOffsetMetadata(messageOffset = latestOffset))
+  @Test
+  def testJoinGroupProtocolType(): Unit = {
+    for (version <- ApiKeys.JOIN_GROUP.oldestVersion to ApiKeys.JOIN_GROUP.latestVersion) {
+      testJoinGroupProtocolType(version.asInstanceOf[Short])
+    }
+  }
+
+  def testJoinGroupProtocolType(version: Short): Unit = {
+    EasyMock.reset(groupCoordinator, clientRequestQuotaManager, requestChannel)
+
+    val capturedResponse = expectNoThrottling()
+
+    val groupId = "group"
+    val memberId = "member1"
+    val protocolType = "consumer"
+    val protocolName = "range"
+    val rebalanceTimeoutMs = 10
+    val sessionTimeoutMs = 5
+
+    val capturedCallback = EasyMock.newCapture[JoinGroupCallback]()
+
+    EasyMock.expect(groupCoordinator.handleJoinGroup(
+      EasyMock.eq(groupId),
+      EasyMock.eq(memberId),
+      EasyMock.eq(None),
+      EasyMock.eq(if (version >= 4) true else false),
+      EasyMock.eq(clientId),
+      EasyMock.eq(InetAddress.getLocalHost.toString),
+      EasyMock.eq(if (version >= 1) rebalanceTimeoutMs else sessionTimeoutMs),
+      EasyMock.eq(sessionTimeoutMs),
+      EasyMock.eq(protocolType),
+      EasyMock.eq(List.empty),
+      EasyMock.capture(capturedCallback)
+    ))
+
+    val joinGroupRequest = new JoinGroupRequest.Builder(
+      new JoinGroupRequestData()
+        .setGroupId(groupId)
+        .setMemberId(memberId)
+        .setProtocolType(protocolType)
+        .setRebalanceTimeoutMs(rebalanceTimeoutMs)
+        .setSessionTimeoutMs(sessionTimeoutMs)
+    ).build(version)
+
+    val requestChannelRequest = buildRequest(joinGroupRequest)
+
+    EasyMock.replay(groupCoordinator, clientRequestQuotaManager, requestChannel)
+
+    createKafkaApis().handleJoinGroupRequest(requestChannelRequest)
+
+    EasyMock.verify(groupCoordinator)
+
+    capturedCallback.getValue.apply(JoinGroupResult(
+      members = List.empty,
+      memberId = memberId,
+      generationId = 0,
+      protocolType = Some(protocolType),
+      protocolName = Some(protocolName),
+      leaderId = memberId,
+      error = Errors.NONE
+    ))
+
+    val response = readResponse(joinGroupRequest, capturedResponse)
+      .asInstanceOf[JoinGroupResponse]
+
+    assertEquals(Errors.NONE, response.error)
+    assertEquals(0, response.data.members.size)
+    assertEquals(memberId, response.data.memberId)
+    assertEquals(0, response.data.generationId)
+    assertEquals(memberId, response.data.leader)
+    assertEquals(protocolName, response.data.protocolName)
+
+    if (version >= 7) {
+      assertEquals(protocolType, response.data.protocolType)
+    } else {
+      assertNull(response.data.protocolType)
+    }
+
+    EasyMock.verify(clientRequestQuotaManager, requestChannel)
+  }
+
+  @Test
+  def testSyncGroupProtocolTypeAndName(): Unit = {
+    for (version <- ApiKeys.SYNC_GROUP.oldestVersion to ApiKeys.SYNC_GROUP.latestVersion) {
+      testSyncGroupProtocolTypeAndName(version.asInstanceOf[Short])
+    }
+  }
+
+  def testSyncGroupProtocolTypeAndName(version: Short): Unit = {
+    EasyMock.reset(groupCoordinator, clientRequestQuotaManager, requestChannel)
+
+    val capturedResponse = expectNoThrottling()
+
+    val groupId = "group"
+    val memberId = "member1"
+    val protocolType = "consumer"
+    val protocolName = "range"
+
+    val capturedCallback = EasyMock.newCapture[SyncGroupCallback]()
+
+    EasyMock.expect(groupCoordinator.handleSyncGroup(
+      EasyMock.eq(groupId),
+      EasyMock.eq(0),
+      EasyMock.eq(memberId),
+      EasyMock.eq(if (version >= 5) Some(protocolType) else None),
+      EasyMock.eq(if (version >= 5) Some(protocolName) else None),
+      EasyMock.eq(None),
+      EasyMock.eq(Map.empty),
+      EasyMock.capture(capturedCallback)
+    ))
+
+    val syncGroupRequest = new SyncGroupRequest.Builder(
+      new SyncGroupRequestData()
+        .setGroupId(groupId)
+        .setGenerationId(0)
+        .setMemberId(memberId)
+        .setProtocolType(protocolType)
+        .setProtocolName(protocolName)
+    ).build(version)
+
+    val requestChannelRequest = buildRequest(syncGroupRequest)
+
+    EasyMock.replay(groupCoordinator, clientRequestQuotaManager, requestChannel)
+
+    createKafkaApis().handleSyncGroupRequest(requestChannelRequest)
+
+    EasyMock.verify(groupCoordinator)
+
+    capturedCallback.getValue.apply(SyncGroupResult(
+      protocolType = Some(protocolType),
+      protocolName = Some(protocolName),
+      memberAssignment = Array.empty,
+      error = Errors.NONE
+    ))
+
+    val response = readResponse(syncGroupRequest, capturedResponse)
+      .asInstanceOf[SyncGroupResponse]
+
+    assertEquals(Errors.NONE, response.error)
+    assertArrayEquals(Array.empty[Byte], response.data.assignment)
+
+    if (version >= 5) {
+      assertEquals(protocolType, response.data.protocolType)
+    } else {
+      assertNull(response.data.protocolType)
+    }
+
+    EasyMock.verify(clientRequestQuotaManager, requestChannel)
+  }
+
+  @Test
+  def testSyncGroupProtocolTypeAndNameAreMandatorySinceV5(): Unit = {
+    for (version <- ApiKeys.SYNC_GROUP.oldestVersion to ApiKeys.SYNC_GROUP.latestVersion) {
+      testSyncGroupProtocolTypeAndNameAreMandatorySinceV5(version.asInstanceOf[Short])
+    }
+  }
+
+  def testSyncGroupProtocolTypeAndNameAreMandatorySinceV5(version: Short): Unit = {
+    EasyMock.reset(groupCoordinator, clientRequestQuotaManager, requestChannel)
+
+    val capturedResponse = expectNoThrottling()
+
+    val groupId = "group"
+    val memberId = "member1"
+    val protocolType = "consumer"
+    val protocolName = "range"
+
+    val capturedCallback = EasyMock.newCapture[SyncGroupCallback]()
+
+    if (version < 5) {
+      EasyMock.expect(groupCoordinator.handleSyncGroup(
+        EasyMock.eq(groupId),
+        EasyMock.eq(0),
+        EasyMock.eq(memberId),
+        EasyMock.eq(None),
+        EasyMock.eq(None),
+        EasyMock.eq(None),
+        EasyMock.eq(Map.empty),
+        EasyMock.capture(capturedCallback)
+      ))
+    }
+
+    val syncGroupRequest = new SyncGroupRequest.Builder(
+      new SyncGroupRequestData()
+        .setGroupId(groupId)
+        .setGenerationId(0)
+        .setMemberId(memberId)
+    ).build(version)
+
+    val requestChannelRequest = buildRequest(syncGroupRequest)
+
+    EasyMock.replay(groupCoordinator, clientRequestQuotaManager, requestChannel)
+
+    createKafkaApis().handleSyncGroupRequest(requestChannelRequest)
+
+    EasyMock.verify(groupCoordinator)
+
+    if (version < 5) {
+      capturedCallback.getValue.apply(SyncGroupResult(
+        protocolType = Some(protocolType),
+        protocolName = Some(protocolName),
+        memberAssignment = Array.empty,
+        error = Errors.NONE
+      ))
+    }
+
+    val response = readResponse(syncGroupRequest, capturedResponse)
+      .asInstanceOf[SyncGroupResponse]
+
+    if (version < 5) {
+      assertEquals(Errors.NONE, response.error)
+    } else {
+      assertEquals(Errors.INCONSISTENT_GROUP_PROTOCOL, response.error)
+    }
+
+    EasyMock.verify(clientRequestQuotaManager, requestChannel)
+  }
+
+  @Test
+  def rejectJoinGroupRequestWhenStaticMembershipNotSupported(): Unit = {
+    val capturedResponse = expectNoThrottling()
+    EasyMock.replay(clientRequestQuotaManager, requestChannel)
+
+    val joinGroupRequest = new JoinGroupRequest.Builder(
+      new JoinGroupRequestData()
+        .setGroupId("test")
+        .setMemberId("test")
+        .setGroupInstanceId("instanceId")
+        .setProtocolType("consumer")
+        .setProtocols(new JoinGroupRequestData.JoinGroupRequestProtocolCollection)
+    ).build()
+
+    val requestChannelRequest = buildRequest(joinGroupRequest)
+    createKafkaApis(KAFKA_2_2_IV1).handleJoinGroupRequest(requestChannelRequest)
+
+    val response = readResponse(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 = new SyncGroupRequest.Builder(
+      new SyncGroupRequestData()
+        .setGroupId("test")
+        .setMemberId("test")
+        .setGroupInstanceId("instanceId")
+        .setGenerationId(1)
+    ).build()
+
+    val requestChannelRequest = buildRequest(syncGroupRequest)
+    createKafkaApis(KAFKA_2_2_IV1).handleSyncGroupRequest(requestChannelRequest)
+
+    val response = readResponse(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 = new HeartbeatRequest.Builder(
+      new HeartbeatRequestData()
+        .setGroupId("test")
+        .setMemberId("test")
+        .setGroupInstanceId("instanceId")
+        .setGenerationId(1)
+    ).build()
+    val requestChannelRequest = buildRequest(heartbeatRequest)
+    createKafkaApis(KAFKA_2_2_IV1).handleHeartbeatRequest(requestChannelRequest)
+
+    val response = readResponse(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 = 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("")
+            ))
+        ))
+    ).build()
+
+    val requestChannelRequest = buildRequest(offsetCommitRequest)
+    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(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
+      ).build()
+    )
+
+    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
+      ).build()
+    )
+
+    createKafkaApis().handleLeaveGroupRequest(leaveRequest)
+
+    EasyMock.replay(groupCoordinator)
+  }
+
+  @Test
+  def testReassignmentAndReplicationBytesOutRateWhenReassigning(): Unit = {
+    assertReassignmentAndReplicationBytesOutPerSec(true)
+  }
+
+  @Test
+  def testReassignmentAndReplicationBytesOutRateWhenNotReassigning(): Unit = {
+    assertReassignmentAndReplicationBytesOutPerSec(false)
+  }
+
+  private def assertReassignmentAndReplicationBytesOutPerSec(isReassigning: Boolean): Unit = {
+    val leaderEpoch = 0
+    val tp0 = new TopicPartition("tp", 0)
+
+    val fetchData = Collections.singletonMap(tp0, new FetchRequest.PartitionData(0, 0, Int.MaxValue, Optional.of(leaderEpoch)))
+    val fetchFromFollower = buildRequest(new FetchRequest.Builder(
+      ApiKeys.FETCH.oldestVersion(), ApiKeys.FETCH.latestVersion(), 1, 1000, 0, fetchData
+    ).build())
+
+    setupBasicMetadataCache(tp0.topic, numPartitions = 1)
+    val hw = 3
+
+    val records = MemoryRecords.withRecords(CompressionType.NONE,
+      new SimpleRecord(1000, "foo".getBytes(StandardCharsets.UTF_8)))
+    replicaManager.fetchMessages(anyLong, anyInt, anyInt, anyInt, anyBoolean,
+      anyObject[Seq[(TopicPartition, FetchRequest.PartitionData)]], anyObject[ReplicaQuota],
+      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]
+        callback(Seq(tp0 -> FetchPartitionData(Errors.NONE, hw, 0, records,
+          None, None, None, Option.empty, isReassignmentFetch = isReassigning)))
+      }
+    })
+
+    val fetchMetadata = new JFetchMetadata(0, 0)
+    val fetchContext = new FullFetchContext(time, new FetchSessionCache(1000, 100),
+      fetchMetadata, fetchData, true)
+    expect(fetchManager.newContext(anyObject[JFetchMetadata],
+      anyObject[util.Map[TopicPartition, FetchRequest.PartitionData]],
+      anyObject[util.List[TopicPartition]],
+      anyBoolean)).andReturn(fetchContext)
+
+    expect(replicaQuotaManager.record(anyLong()))
+    expect(replicaManager.getLogConfig(EasyMock.eq(tp0))).andReturn(None)
+
+    val partition: Partition = createNiceMock(classOf[Partition])
+    expect(replicaManager.isAddingReplica(anyObject(), anyInt())).andReturn(isReassigning)
+
+    replay(replicaManager, fetchManager, clientQuotaManager, requestChannel, replicaQuotaManager, partition)
+
+    createKafkaApis().handle(fetchFromFollower)
+
+    if (isReassigning)
+      assertEquals(records.sizeInBytes(), brokerTopicStats.allTopicsStats.reassignmentBytesOutPerSec.get.count())
     else
-      EasyMock.expect(replica.lastStableOffset).andReturn(LogOffsetMetadata(messageOffset = latestOffset))
-    expectThrottleCallbackAndInvoke(capturedThrottleCallback)
+      assertEquals(0, brokerTopicStats.allTopicsStats.reassignmentBytesOutPerSec.get.count())
+    assertEquals(records.sizeInBytes(), brokerTopicStats.allTopicsStats.replicationBytesOutRate.get.count())
+
+  }
+
+  @Test
+  def rejectInitProducerIdWhenIdButNotEpochProvided(): Unit = {
+    val capturedResponse = expectNoThrottling()
+    EasyMock.replay(clientRequestQuotaManager, requestChannel)
+
+    val initProducerIdRequest = new InitProducerIdRequest.Builder(
+      new InitProducerIdRequestData()
+        .setTransactionalId("known")
+        .setTransactionTimeoutMs(TimeUnit.MINUTES.toMillis(15).toInt)
+        .setProducerId(10)
+        .setProducerEpoch(RecordBatch.NO_PRODUCER_EPOCH)
+    ).build()
+
+    val requestChannelRequest = buildRequest(initProducerIdRequest)
+    createKafkaApis(KAFKA_2_2_IV1).handleInitProducerIdRequest(requestChannelRequest)
+
+    val response = readResponse(initProducerIdRequest, capturedResponse)
+      .asInstanceOf[InitProducerIdResponse]
+    assertEquals(Errors.INVALID_REQUEST, response.error)
+  }
+
+  @Test
+  def rejectInitProducerIdWhenEpochButNotIdProvided(): Unit = {
+    val capturedResponse = expectNoThrottling()
+    EasyMock.replay(clientRequestQuotaManager, requestChannel)
+
+    val initProducerIdRequest = new InitProducerIdRequest.Builder(
+      new InitProducerIdRequestData()
+        .setTransactionalId("known")
+        .setTransactionTimeoutMs(TimeUnit.MINUTES.toMillis(15).toInt)
+        .setProducerId(RecordBatch.NO_PRODUCER_ID)
+        .setProducerEpoch(2)
+    ).build()
+    val requestChannelRequest = buildRequest(initProducerIdRequest)
+    createKafkaApis(KAFKA_2_2_IV1).handleInitProducerIdRequest(requestChannelRequest)
+
+    val response = readResponse(initProducerIdRequest, capturedResponse).asInstanceOf[InitProducerIdResponse]
+    assertEquals(Errors.INVALID_REQUEST, response.error)
+  }
+
+  @Test
+  def testUpdateMetadataRequestWithCurrentBrokerEpoch(): Unit = {
+    val currentBrokerEpoch = 1239875L
+    testUpdateMetadataRequest(currentBrokerEpoch, currentBrokerEpoch, Errors.NONE)
+  }
+
+  @Test
+  def testUpdateMetadataRequestWithNewerBrokerEpochIsValid(): Unit = {
+    val currentBrokerEpoch = 1239875L
+    testUpdateMetadataRequest(currentBrokerEpoch, currentBrokerEpoch + 1, Errors.NONE)
+  }
+
+  @Test
+  def testUpdateMetadataRequestWithStaleBrokerEpochIsRejected(): Unit = {
+    val currentBrokerEpoch = 1239875L
+    testUpdateMetadataRequest(currentBrokerEpoch, currentBrokerEpoch - 1, Errors.STALE_BROKER_EPOCH)
+  }
+
+  def testUpdateMetadataRequest(currentBrokerEpoch: Long, brokerEpochInRequest: Long, expectedError: Errors): Unit = {
+    val updateMetadataRequest = createBasicMetadataRequest("topicA", 1, brokerEpochInRequest)
+    val request = buildRequest(updateMetadataRequest)
+
+    val capturedResponse: Capture[RequestChannel.Response] = EasyMock.newCapture()
+
+    EasyMock.expect(controller.brokerEpoch).andStubReturn(currentBrokerEpoch)
+    EasyMock.expect(replicaManager.maybeUpdateMetadataCache(
+      EasyMock.eq(request.context.correlationId),
+      EasyMock.anyObject()
+    )).andStubReturn(
+      Seq()
+    )
+
+    EasyMock.expect(requestChannel.sendResponse(EasyMock.capture(capturedResponse)))
+    EasyMock.replay(replicaManager, controller, requestChannel)
+
+    createKafkaApis().handleUpdateMetadataRequest(request)
+    val updateMetadataResponse = readResponse(updateMetadataRequest, capturedResponse)
+      .asInstanceOf[UpdateMetadataResponse]
+    assertEquals(expectedError, updateMetadataResponse.error())
+    EasyMock.verify(replicaManager)
+  }
+
+  @Test
+  def testLeaderAndIsrRequestWithCurrentBrokerEpoch(): Unit = {
+    val currentBrokerEpoch = 1239875L
+    testLeaderAndIsrRequest(currentBrokerEpoch, currentBrokerEpoch, Errors.NONE)
+  }
+
+  @Test
+  def testLeaderAndIsrRequestWithNewerBrokerEpochIsValid(): Unit = {
+    val currentBrokerEpoch = 1239875L
+    testLeaderAndIsrRequest(currentBrokerEpoch, currentBrokerEpoch + 1, Errors.NONE)
+  }
+
+  @Test
+  def testLeaderAndIsrRequestWithStaleBrokerEpochIsRejected(): Unit = {
+    val currentBrokerEpoch = 1239875L
+    testLeaderAndIsrRequest(currentBrokerEpoch, currentBrokerEpoch - 1, Errors.STALE_BROKER_EPOCH)
+  }
+
+  def testLeaderAndIsrRequest(currentBrokerEpoch: Long, brokerEpochInRequest: Long, expectedError: Errors): Unit = {
+    val controllerId = 2
+    val controllerEpoch = 6
+    val capturedResponse: Capture[RequestChannel.Response] = EasyMock.newCapture()
+    val partitionStates = Seq(
+      new LeaderAndIsrRequestData.LeaderAndIsrPartitionState()
+        .setTopicName("topicW")
+        .setPartitionIndex(1)
+        .setControllerEpoch(1)
+        .setLeader(0)
+        .setLeaderEpoch(1)
+        .setIsr(asList(0, 1))
+        .setZkVersion(2)
+        .setReplicas(asList(0, 1, 2))
+        .setIsNew(false)
+    ).asJava
+    val leaderAndIsrRequest = new LeaderAndIsrRequest.Builder(
+      ApiKeys.LEADER_AND_ISR.latestVersion,
+      controllerId,
+      controllerEpoch,
+      brokerEpochInRequest,
+      partitionStates,
+      asList(new Node(0, "host0", 9090), new Node(1, "host1", 9091))
+    ).build()
+    val request = buildRequest(leaderAndIsrRequest)
+    val response = new LeaderAndIsrResponse(new LeaderAndIsrResponseData()
+      .setErrorCode(Errors.NONE.code)
+      .setPartitionErrors(asList()))
+
+    EasyMock.expect(controller.brokerEpoch).andStubReturn(currentBrokerEpoch)
+    EasyMock.expect(replicaManager.becomeLeaderOrFollower(
+      EasyMock.eq(request.context.correlationId),
+      EasyMock.anyObject(),
+      EasyMock.anyObject()
+    )).andStubReturn(
+      response
+    )
+
+    EasyMock.expect(requestChannel.sendResponse(EasyMock.capture(capturedResponse)))
+    EasyMock.replay(replicaManager, controller, requestChannel)
+
+    createKafkaApis().handleLeaderAndIsrRequest(request)
+    val leaderAndIsrResponse = readResponse(leaderAndIsrRequest, capturedResponse)
+      .asInstanceOf[LeaderAndIsrResponse]
+    assertEquals(expectedError, leaderAndIsrResponse.error())
+    EasyMock.verify(replicaManager)
+  }
+
+  @Test
+  def testStopReplicaRequestWithCurrentBrokerEpoch(): Unit = {
+    val currentBrokerEpoch = 1239875L
+    testStopReplicaRequest(currentBrokerEpoch, currentBrokerEpoch, Errors.NONE)
+  }
+
+  @Test
+  def testStopReplicaRequestWithNewerBrokerEpochIsValid(): Unit = {
+    val currentBrokerEpoch = 1239875L
+    testStopReplicaRequest(currentBrokerEpoch, currentBrokerEpoch + 1, Errors.NONE)
+  }
+
+  @Test
+  def testStopReplicaRequestWithStaleBrokerEpochIsRejected(): Unit = {
+    val currentBrokerEpoch = 1239875L
+    testStopReplicaRequest(currentBrokerEpoch, currentBrokerEpoch - 1, Errors.STALE_BROKER_EPOCH)
+  }
+
+  def testStopReplicaRequest(currentBrokerEpoch: Long, brokerEpochInRequest: Long, expectedError: Errors): Unit = {
+    val controllerId = 0
+    val controllerEpoch = 5
+    val capturedResponse: Capture[RequestChannel.Response] = EasyMock.newCapture()
+    val fooPartition = new TopicPartition("foo", 0)
+    val topicStates = Seq(
+      new StopReplicaTopicState()
+        .setTopicName(fooPartition.topic())
+        .setPartitionStates(Seq(new StopReplicaPartitionState()
+          .setPartitionIndex(fooPartition.partition())
+          .setLeaderEpoch(1)
+          .setDeletePartition(false)).asJava)
+    ).asJava
+    val stopReplicaRequest = new StopReplicaRequest.Builder(
+      ApiKeys.STOP_REPLICA.latestVersion,
+      controllerId,
+      controllerEpoch,
+      brokerEpochInRequest,
+      false,
+      topicStates
+    ).build()
+    val request = buildRequest(stopReplicaRequest)
+
+    EasyMock.expect(controller.brokerEpoch).andStubReturn(currentBrokerEpoch)
+    EasyMock.expect(replicaManager.stopReplicas(
+      EasyMock.eq(request.context.correlationId),
+      EasyMock.eq(controllerId),
+      EasyMock.eq(controllerEpoch),
+      EasyMock.eq(brokerEpochInRequest),
+      EasyMock.eq(stopReplicaRequest.partitionStates().asScala)
+    )).andStubReturn(
+      (mutable.Map(
+        fooPartition -> Errors.NONE
+      ), Errors.NONE)
+    )
     EasyMock.expect(requestChannel.sendResponse(EasyMock.capture(capturedResponse)))
-    EasyMock.replay(replicaManager, clientRequestQuotaManager, requestChannel, replica, log)
 
-    val builder = ListOffsetRequest.Builder.forConsumer(true, isolationLevel)
-      .setTargetTimes(Map(tp -> (ListOffsetRequest.LATEST_TIMESTAMP: JLong)).asJava)
-    val (listOffsetRequest, request) = buildRequest(builder)
+    EasyMock.replay(controller, replicaManager, requestChannel)
+
+    createKafkaApis().handleStopReplicaRequest(request)
+    val stopReplicaResponse = readResponse(stopReplicaRequest, capturedResponse)
+      .asInstanceOf[StopReplicaResponse]
+    assertEquals(expectedError, stopReplicaResponse.error())
+    EasyMock.verify(replicaManager)
+  }
+
+  @Test
+  def testListGroupsRequest(): Unit = {
+    val overviews = List(
+      GroupOverview("group1", "protocol1", "Stable"),
+      GroupOverview("group2", "qwerty", "Empty")
+    )
+    val response = listGroupRequest(None, overviews)
+    assertEquals(2, response.data.groups.size)
+    assertEquals("Stable", response.data.groups.get(0).groupState)
+    assertEquals("Empty", response.data.groups.get(1).groupState)
+  }
+
+  @Test
+  def testListGroupsRequestWithState(): Unit = {
+    val overviews = List(
+      GroupOverview("group1", "protocol1", "Stable")
+    )
+    val response = listGroupRequest(Some("Stable"), overviews)
+    assertEquals(1, response.data.groups.size)
+    assertEquals("Stable", response.data.groups.get(0).groupState)
+  }
+
+  private def listGroupRequest(state: Option[String], overviews: List[GroupOverview]): ListGroupsResponse = {
+    EasyMock.reset(groupCoordinator, clientRequestQuotaManager, requestChannel)
+
+    val data = new ListGroupsRequestData()
+    if (state.isDefined)
+      data.setStatesFilter(Collections.singletonList(state.get))
+    val listGroupsRequest = new ListGroupsRequest.Builder(data).build()
+    val requestChannelRequest = buildRequest(listGroupsRequest)
+
+    val capturedResponse = expectNoThrottling()
+    val expectedStates: Set[String] = if (state.isDefined) Set(state.get) else Set()
+    EasyMock.expect(groupCoordinator.handleListGroups(expectedStates))
+      .andReturn((Errors.NONE, overviews))
+    EasyMock.replay(groupCoordinator, clientRequestQuotaManager, requestChannel)
+
+    createKafkaApis().handleListGroupsRequest(requestChannelRequest)
+
+    val response = readResponse(listGroupsRequest, capturedResponse).asInstanceOf[ListGroupsResponse]
+    assertEquals(Errors.NONE.code, response.data.errorCode)
+    response
+  }
+
+  /**
+   * 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 = 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, Seq.empty[UpdateMetadataPartitionState].asJava, brokers.asJava).build()
+    metadataCache.updateMetadata(correlationId = 0, updateMetadataRequest)
+    (plaintextListener, anotherListener)
+  }
+
+  private def sendMetadataRequestWithInconsistentListeners(requestListener: ListenerName): MetadataResponse = {
+    val capturedResponse = expectNoThrottling()
+    EasyMock.replay(clientRequestQuotaManager, requestChannel)
+
+    val metadataRequest = MetadataRequest.Builder.allTopics.build()
+    val requestChannelRequest = buildRequest(metadataRequest, requestListener)
+    createKafkaApis().handleTopicMetadataRequest(requestChannelRequest)
+
+    readResponse(metadataRequest, capturedResponse).asInstanceOf[MetadataResponse]
+  }
+
+  private def testConsumerListOffsetLatest(isolationLevel: IsolationLevel): Unit = {
+    val tp = new TopicPartition("foo", 0)
+    val latestOffset = 15L
+    val currentLeaderEpoch = Optional.empty[Integer]()
+
+    EasyMock.expect(replicaManager.fetchOffsetForTimestamp(
+      EasyMock.eq(tp),
+      EasyMock.eq(ListOffsetRequest.LATEST_TIMESTAMP),
+      EasyMock.eq(Some(isolationLevel)),
+      EasyMock.eq(currentLeaderEpoch),
+      fetchOnlyFromLeader = EasyMock.eq(true))
+    ).andReturn(Some(new TimestampAndOffset(ListOffsetResponse.UNKNOWN_TIMESTAMP, latestOffset, currentLeaderEpoch)))
+
+    val capturedResponse = expectNoThrottling()
+    EasyMock.replay(replicaManager, clientRequestQuotaManager, requestChannel)
+
+    val targetTimes = List(new ListOffsetTopic()
+      .setName(tp.topic)
+      .setPartitions(List(new ListOffsetPartition()
+        .setPartitionIndex(tp.partition)
+        .setTimestamp(ListOffsetRequest.LATEST_TIMESTAMP)).asJava)).asJava
+    val listOffsetRequest = ListOffsetRequest.Builder.forConsumer(true, isolationLevel)
+      .setTargetTimes(targetTimes).build()
+    val request = buildRequest(listOffsetRequest)
     createKafkaApis().handleListOffsetRequest(request)
 
-    val response = readResponse(ApiKeys.LIST_OFFSETS, listOffsetRequest, capturedResponse).asInstanceOf[ListOffsetResponse]
-    assertTrue(response.responseData.containsKey(tp))
+    val response = readResponse(listOffsetRequest, capturedResponse).asInstanceOf[ListOffsetResponse]
+    val partitionDataOptional = response.topics.asScala.find(_.name == tp.topic).get
+      .partitions.asScala.find(_.partitionIndex == tp.partition)
+    assertTrue(partitionDataOptional.isDefined)
 
-    val partitionData = response.responseData.get(tp)
-    assertEquals(Errors.NONE, partitionData.error)
+    val partitionData = partitionDataOptional.get
+    assertEquals(Errors.NONE.code, partitionData.errorCode)
     assertEquals(latestOffset, partitionData.offset)
     assertEquals(ListOffsetResponse.UNKNOWN_TIMESTAMP, partitionData.timestamp)
   }
 
   private def createWriteTxnMarkersRequest(partitions: util.List[TopicPartition]) = {
-    val requestBuilder = new WriteTxnMarkersRequest.Builder(Utils.mkList(
-      new TxnMarkerEntry(1, 1.toShort, 0, TransactionResult.COMMIT, partitions)))
-    buildRequest(requestBuilder)
+    val writeTxnMarkersRequest = new WriteTxnMarkersRequest.Builder(ApiKeys.WRITE_TXN_MARKERS.latestVersion(),
+      asList(new TxnMarkerEntry(1, 1.toShort, 0, TransactionResult.COMMIT, partitions))).build()
+    (writeTxnMarkersRequest, buildRequest(writeTxnMarkersRequest))
+  }
+
+  private def buildRequestWithEnvelope(
+    request: AbstractRequest,
+    fromPrivilegedListener: Boolean,
+    principalSerde: Option[KafkaPrincipalSerde] = kafkaPrincipalSerde
+  ): RequestChannel.Request = {
+    val listenerName = ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT)
+
+    val requestHeader = new RequestHeader(request.apiKey, request.version, clientId, 0)
+    val requestBuffer = RequestTestUtils.serializeRequestWithHeader(requestHeader, request)
+    val requestContext = new RequestContext(requestHeader, "1", InetAddress.getLocalHost,
+      KafkaPrincipal.ANONYMOUS, listenerName, SecurityProtocol.PLAINTEXT, ClientInformation.EMPTY,
+      fromPrivilegedListener)
+
+    val envelopeHeader = new RequestHeader(ApiKeys.ENVELOPE, ApiKeys.ENVELOPE.latestVersion(), clientId, 0)
+    val envelopeBuffer = RequestTestUtils.serializeRequestWithHeader(envelopeHeader, new EnvelopeRequest.Builder(
+      requestBuffer,
+      new Array[Byte](0),
+      InetAddress.getLocalHost.getAddress
+    ).build())
+    val envelopeContext = new RequestContext(envelopeHeader, "1", InetAddress.getLocalHost,
+      KafkaPrincipal.ANONYMOUS, listenerName, SecurityProtocol.PLAINTEXT, ClientInformation.EMPTY,
+      fromPrivilegedListener, principalSerde.asJava)
+
+    RequestHeader.parse(envelopeBuffer)
+    val envelopeRequest = Some(new RequestChannel.Request(
+      processor = 1,
+      context = envelopeContext,
+      startTimeNanos = time.nanoseconds(),
+      memoryPool = MemoryPool.NONE,
+      buffer = envelopeBuffer,
+      metrics = requestChannelMetrics
+    ))
+
+    RequestHeader.parse(requestBuffer)
+    new RequestChannel.Request(
+      processor = 1,
+      context = requestContext,
+      startTimeNanos = time.nanoseconds(),
+      memoryPool = MemoryPool.NONE,
+      buffer = requestBuffer,
+      metrics = requestChannelMetrics,
+      envelope = envelopeRequest
+    )
   }
 
-  private def buildRequest[T <: AbstractRequest](builder: AbstractRequest.Builder[T]): (T, RequestChannel.Request) = {
-    val request = builder.build()
-    val buffer = request.serialize(new RequestHeader(builder.apiKey, request.version, "", 0))
+  private def buildRequest(request: AbstractRequest,
+                           listenerName: ListenerName = ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT),
+                           fromPrivilegedListener: Boolean = false,
+                           requestHeader: Option[RequestHeader] = None): RequestChannel.Request = {
+    val buffer = RequestTestUtils.serializeRequestWithHeader(requestHeader.getOrElse(
+      new RequestHeader(request.apiKey, request.version, clientId, 0)), request)
 
     // read the header from the buffer first so that the body can be read next from the Request constructor
     val header = RequestHeader.parse(buffer)
     val context = new RequestContext(header, "1", InetAddress.getLocalHost, KafkaPrincipal.ANONYMOUS,
-      new ListenerName(""), SecurityProtocol.PLAINTEXT)
-    (request, new RequestChannel.Request(processor = 1, context = context, startTimeNanos =  0,
-      MemoryPool.NONE, buffer, requestChannelMetrics))
+      listenerName, SecurityProtocol.PLAINTEXT, ClientInformation.EMPTY, fromPrivilegedListener,
+      kafkaPrincipalSerde.asJava)
+    new RequestChannel.Request(processor = 1, context = context, startTimeNanos = 0, MemoryPool.NONE, buffer,
+      requestChannelMetrics, envelope = None)
   }
 
-  private def readResponse(api: ApiKeys, request: AbstractRequest, capturedResponse: Capture[RequestChannel.Response]): AbstractResponse = {
-    val send = capturedResponse.getValue.responseSend.get
+  private def readResponse(request: AbstractRequest, capturedResponse: Capture[RequestChannel.Response]) = {
+    val api = request.apiKey
+    val response = capturedResponse.getValue
+    assertTrue(s"Unexpected response type: ${response.getClass}", response.isInstanceOf[SendResponse])
+    val sendResponse = response.asInstanceOf[SendResponse]
+    val send = sendResponse.responseSend
     val channel = new ByteBufferChannel(send.size)
     send.writeTo(channel)
     channel.close()
     channel.buffer.getInt() // read the size
-    ResponseHeader.parse(channel.buffer)
-    val struct = api.responseSchema(request.version).read(channel.buffer)
-    AbstractResponse.parseResponse(api, struct)
+    ResponseHeader.parse(channel.buffer, api.responseHeaderVersion(request.version))
+    AbstractResponse.parseResponse(api, channel.buffer, request.version)
+  }
+
+  private def expectNoThrottling(): Capture[RequestChannel.Response] = {
+    EasyMock.expect(clientRequestQuotaManager.maybeRecordAndGetThrottleTimeMs(EasyMock.anyObject[RequestChannel.Request](),
+      EasyMock.anyObject[Long])).andReturn(0)
+    EasyMock.expect(clientRequestQuotaManager.throttle(EasyMock.anyObject[RequestChannel.Request](), EasyMock.eq(0),
+      EasyMock.anyObject[RequestChannel.Response => Unit]()))
+
+    val capturedResponse = EasyMock.newCapture[RequestChannel.Response]()
+    EasyMock.expect(requestChannel.sendResponse(EasyMock.capture(capturedResponse)))
+    capturedResponse
+  }
+
+  private def createBasicMetadataRequest(topic: String, numPartitions: Int, brokerEpoch: Long): UpdateMetadataRequest = {
+    val replicas = List(0.asInstanceOf[Integer]).asJava
+
+    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 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)
+    new UpdateMetadataRequest.Builder(ApiKeys.UPDATE_METADATA.latestVersion, 0,
+      0, brokerEpoch, partitionStates.asJava, Seq(broker).asJava).build()
   }
 
-  private def expectThrottleCallbackAndInvoke(capturedThrottleCallback: Capture[Int => Unit]): Unit = {
-    EasyMock.expect(clientRequestQuotaManager.maybeRecordAndThrottle(
-      EasyMock.anyObject[RequestChannel.Request](),
-      EasyMock.capture(capturedThrottleCallback)))
-      .andAnswer(new IAnswer[Unit] {
-        override def answer(): Unit = {
-          val callback = capturedThrottleCallback.getValue
-          callback(0)
-        }
-      })
+  private def setupBasicMetadataCache(topic: String, numPartitions: Int): Unit = {
+    val updateMetadataRequest = createBasicMetadataRequest(topic, numPartitions, 0)
+    metadataCache.updateMetadata(correlationId = 0, updateMetadataRequest)
   }
 
+  @Test
+  def testAlterReplicaLogDirs(): Unit = {
+    val data = new AlterReplicaLogDirsRequestData()
+    val dir = new AlterReplicaLogDirsRequestData.AlterReplicaLogDir()
+      .setPath("/foo")
+    dir.topics().add(new AlterReplicaLogDirsRequestData.AlterReplicaLogDirTopic().setName("t0").setPartitions(asList(0, 1, 2)))
+    data.dirs().add(dir)
+    val alterReplicaLogDirsRequest = new AlterReplicaLogDirsRequest.Builder(
+      data
+    ).build()
+    val request = buildRequest(alterReplicaLogDirsRequest)
+
+    EasyMock.reset(replicaManager, clientRequestQuotaManager, requestChannel)
+
+    val capturedResponse = expectNoThrottling()
+    val t0p0 = new TopicPartition("t0", 0)
+    val t0p1 = new TopicPartition("t0", 1)
+    val t0p2 = new TopicPartition("t0", 2)
+    val partitionResults = Map(
+      t0p0 -> Errors.NONE,
+      t0p1 -> Errors.LOG_DIR_NOT_FOUND,
+      t0p2 -> Errors.INVALID_TOPIC_EXCEPTION)
+    EasyMock.expect(replicaManager.alterReplicaLogDirs(EasyMock.eq(Map(
+      t0p0 -> "/foo",
+      t0p1 -> "/foo",
+      t0p2 -> "/foo"))))
+    .andReturn(partitionResults)
+    EasyMock.replay(replicaManager, clientQuotaManager, clientRequestQuotaManager, requestChannel)
+
+    createKafkaApis().handleAlterReplicaLogDirsRequest(request)
+
+    val response = readResponse(alterReplicaLogDirsRequest, capturedResponse)
+      .asInstanceOf[AlterReplicaLogDirsResponse]
+    assertEquals(partitionResults, response.data.results.asScala.flatMap { tr =>
+      tr.partitions().asScala.map { pr =>
+        new TopicPartition(tr.topicName, pr.partitionIndex) -> Errors.forCode(pr.errorCode)
+      }
+    }.toMap)
+    assertEquals(Map(Errors.NONE -> 1,
+      Errors.LOG_DIR_NOT_FOUND -> 1,
+      Errors.INVALID_TOPIC_EXCEPTION -> 1).asJava, response.errorCounts)
+  }
 }
diff --git a/core/src/test/scala/unit/kafka/server/KafkaConfigTest.scala b/core/src/test/scala/unit/kafka/server/KafkaConfigTest.scala
index 9c459d839e3ce..f264cf9ae0fc9 100755
--- a/core/src/test/scala/unit/kafka/server/KafkaConfigTest.scala
+++ b/core/src/test/scala/unit/kafka/server/KafkaConfigTest.scala
@@ -21,11 +21,13 @@ import java.util.Properties
 
 import kafka.api.{ApiVersion, KAFKA_0_8_2}
 import kafka.cluster.EndPoint
+import kafka.log.LogConfig
 import kafka.message._
 import kafka.utils.{CoreUtils, TestUtils}
 import org.apache.kafka.common.config.ConfigException
 import org.apache.kafka.common.metrics.Sensor
 import org.apache.kafka.common.network.ListenerName
+import org.apache.kafka.common.record.Records
 import org.apache.kafka.common.security.auth.SecurityProtocol
 import org.junit.Assert._
 import org.junit.Test
@@ -34,7 +36,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")
 
@@ -43,7 +45,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testLogRetentionTimeMinutesProvided() {
+  def testLogRetentionTimeMinutesProvided(): Unit = {
     val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
     props.put(KafkaConfig.LogRetentionTimeMinutesProp, "30")
 
@@ -52,7 +54,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testLogRetentionTimeMsProvided() {
+  def testLogRetentionTimeMsProvided(): Unit = {
     val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
     props.put(KafkaConfig.LogRetentionTimeMillisProp, "1800000")
 
@@ -61,7 +63,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testLogRetentionTimeNoConfigProvided() {
+  def testLogRetentionTimeNoConfigProvided(): Unit = {
     val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
 
     val cfg = KafkaConfig.fromProps(props)
@@ -69,7 +71,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")
@@ -79,7 +81,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")
@@ -89,7 +91,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)
@@ -143,7 +145,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testAdvertiseDefaults() {
+  def testAdvertiseDefaults(): Unit = {
     val port = "9999"
     val hostName = "fake-host"
 
@@ -159,7 +161,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testAdvertiseConfigured() {
+  def testAdvertiseConfigured(): Unit = {
     val advertisedHostName = "routable-host"
     val advertisedPort = "1234"
 
@@ -176,7 +178,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testAdvertisePortDefault() {
+  def testAdvertisePortDefault(): Unit = {
     val advertisedHostName = "routable-host"
     val port = "9999"
 
@@ -193,7 +195,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testAdvertiseHostNameDefault() {
+  def testAdvertiseHostNameDefault(): Unit = {
     val hostName = "routable-host"
     val advertisedPort = "9999"
 
@@ -210,26 +212,61 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testDuplicateListeners() {
+  def testDuplicateListeners(): Unit = {
     val props = new Properties()
     props.put(KafkaConfig.BrokerIdProp, "1")
     props.put(KafkaConfig.ZkConnectProp, "localhost:2181")
 
     // listeners with duplicate port
-    props.put(KafkaConfig.ListenersProp, "PLAINTEXT://localhost:9091,TRACE://localhost:9091")
-    assertFalse(isValidKafkaConfig(props))
+    props.put(KafkaConfig.ListenersProp, "PLAINTEXT://localhost:9091,SSL://localhost:9091")
+    var caught = intercept[IllegalArgumentException] { KafkaConfig.fromProps(props) }
+    assertTrue(caught.getMessage.contains("Each listener must have a different port"))
 
-    // listeners with duplicate protocol
+    // listeners with duplicate name
     props.put(KafkaConfig.ListenersProp, "PLAINTEXT://localhost:9091,PLAINTEXT://localhost:9092")
-    assertFalse(isValidKafkaConfig(props))
+    caught = intercept[IllegalArgumentException] { KafkaConfig.fromProps(props) }
+    assertTrue(caught.getMessage.contains("Each listener must have a different name"))
 
-    // advertised listeners with duplicate port
-    props.put(KafkaConfig.AdvertisedListenersProp, "PLAINTEXT://localhost:9091,TRACE://localhost:9091")
-    assertFalse(isValidKafkaConfig(props))
+    // advertised listeners can have duplicate ports
+    props.put(KafkaConfig.ListenerSecurityProtocolMapProp, "HOST:SASL_SSL,LB:SASL_SSL")
+    props.put(KafkaConfig.InterBrokerListenerNameProp, "HOST")
+    props.put(KafkaConfig.ListenersProp, "HOST://localhost:9091,LB://localhost:9092")
+    props.put(KafkaConfig.AdvertisedListenersProp, "HOST://localhost:9091,LB://localhost:9091")
+    assertTrue(isValidKafkaConfig(props))
+
+    // but not duplicate names
+    props.put(KafkaConfig.AdvertisedListenersProp, "HOST://localhost:9091,HOST://localhost:9091")
+    caught = intercept[IllegalArgumentException] { KafkaConfig.fromProps(props) }
+    assertTrue(caught.getMessage.contains("Each listener must have a different name"))
+  }
+
+  @Test
+  def testControlPlaneListenerName() = {
+    val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect)
+    props.put("listeners", "PLAINTEXT://localhost:0,CONTROLLER://localhost:5000")
+    props.put("listener.security.protocol.map", "PLAINTEXT:PLAINTEXT,CONTROLLER:SSL")
+    props.put("control.plane.listener.name", "CONTROLLER")
+    assertTrue(isValidKafkaConfig(props))
+
+    val serverConfig = KafkaConfig.fromProps(props)
+    val controlEndpoint = serverConfig.controlPlaneListener.get
+    assertEquals("localhost", controlEndpoint.host)
+    assertEquals(5000, controlEndpoint.port)
+    assertEquals(SecurityProtocol.SSL, controlEndpoint.securityProtocol)
+
+    //advertised listener should contain control-plane listener
+    val advertisedEndpoints = serverConfig.advertisedListeners
+    assertFalse(advertisedEndpoints.filter { endpoint =>
+      endpoint.securityProtocol == controlEndpoint.securityProtocol && endpoint.listenerName.value().equals(controlEndpoint.listenerName.value())
+    }.isEmpty)
+
+    // interBrokerListener name should be different from control-plane listener name
+    val interBrokerListenerName = serverConfig.interBrokerListenerName
+    assertFalse(interBrokerListenerName.value().equals(controlEndpoint.listenerName.value()))
   }
 
   @Test
-  def testBadListenerProtocol() {
+  def testBadListenerProtocol(): Unit = {
     val props = new Properties()
     props.put(KafkaConfig.BrokerIdProp, "1")
     props.put(KafkaConfig.ZkConnectProp, "localhost:2181")
@@ -328,7 +365,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testCaseInsensitiveListenerProtocol() {
+  def testCaseInsensitiveListenerProtocol(): Unit = {
     val props = new Properties()
     props.put(KafkaConfig.BrokerIdProp, "1")
     props.put(KafkaConfig.ZkConnectProp, "localhost:2181")
@@ -343,7 +380,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")
@@ -372,7 +409,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testVersionConfiguration() {
+  def testVersionConfiguration(): Unit = {
     val props = new Properties()
     props.put(KafkaConfig.BrokerIdProp, "1")
     props.put(KafkaConfig.ZkConnectProp, "localhost:2181")
@@ -406,7 +443,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testUncleanLeaderElectionDefault() {
+  def testUncleanLeaderElectionDefault(): Unit = {
     val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
     val serverConfig = KafkaConfig.fromProps(props)
 
@@ -414,7 +451,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)
@@ -423,7 +460,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)
@@ -432,7 +469,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testUncleanElectionInvalid() {
+  def testUncleanElectionInvalid(): Unit = {
     val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
     props.put(KafkaConfig.UncleanLeaderElectionEnableProp, "invalid")
 
@@ -442,7 +479,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testLogRollTimeMsProvided() {
+  def testLogRollTimeMsProvided(): Unit = {
     val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
     props.put(KafkaConfig.LogRollTimeMillisProp, "1800000")
 
@@ -451,7 +488,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")
@@ -461,7 +498,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testLogRollTimeNoConfigProvided() {
+  def testLogRollTimeNoConfigProvided(): Unit = {
     val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
 
     val cfg = KafkaConfig.fromProps(props)
@@ -469,7 +506,7 @@ class KafkaConfigTest {
   }
 
   @Test
-  def testDefaultCompressionType() {
+  def testDefaultCompressionType(): Unit = {
     val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
     val serverConfig = KafkaConfig.fromProps(props)
 
@@ -477,7 +514,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)
@@ -486,7 +523,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] {
@@ -495,7 +532,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)
@@ -505,7 +542,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")
@@ -513,141 +550,189 @@ 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")
-    intercept[IllegalArgumentException] {
+    var caught = intercept[IllegalArgumentException] { KafkaConfig.fromProps(props) }
+    assertTrue(caught.getMessage.contains("No security protocol defined for listener TRACE"))
+
+    props.put(KafkaConfig.ListenerSecurityProtocolMapProp, "PLAINTEXT:PLAINTEXT,TRACE:PLAINTEXT,SSL:SSL")
+    caught = intercept[IllegalArgumentException] { KafkaConfig.fromProps(props) }
+    assertTrue(caught.getMessage.contains("advertised.listeners listener names must be equal to or a subset of the ones defined in listeners"))
+  }
+
+  @Test
+  def testInterBrokerVersionMessageFormatCompatibility(): Unit = {
+    def buildConfig(interBrokerProtocol: ApiVersion, messageFormat: ApiVersion): KafkaConfig = {
+      val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
+      props.put(KafkaConfig.InterBrokerProtocolVersionProp, interBrokerProtocol.version)
+      props.put(KafkaConfig.LogMessageFormatVersionProp, messageFormat.version)
       KafkaConfig.fromProps(props)
     }
+
+    ApiVersion.allVersions.foreach { interBrokerVersion =>
+      ApiVersion.allVersions.foreach { messageFormatVersion =>
+        if (interBrokerVersion.recordVersion.value >= messageFormatVersion.recordVersion.value) {
+          val config = buildConfig(interBrokerVersion, messageFormatVersion)
+          assertEquals(messageFormatVersion, config.logMessageFormatVersion)
+          assertEquals(interBrokerVersion, config.interBrokerProtocolVersion)
+        } else {
+          intercept[IllegalArgumentException] {
+            buildConfig(interBrokerVersion, messageFormatVersion)
+          }
+        }
+      }
+    }
   }
 
   @Test
-  def testFromPropsInvalid() {
-    def getBaseProperties(): Properties = {
+  def testFromPropsInvalid(): Unit = {
+    def baseProperties: Properties = {
       val validRequiredProperties = new Properties()
       validRequiredProperties.put(KafkaConfig.ZkConnectProp, "127.0.0.1:2181")
       validRequiredProperties
     }
     // to ensure a basis is valid - bootstraps all needed validation
-    KafkaConfig.fromProps(getBaseProperties())
+    KafkaConfig.fromProps(baseProperties)
 
-    KafkaConfig.configNames().foreach(name => {
+    KafkaConfig.configNames.foreach { name =>
       name match {
         case KafkaConfig.ZkConnectProp => // ignore string
-        case KafkaConfig.ZkSessionTimeoutMsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
-        case KafkaConfig.ZkConnectionTimeoutMsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
-        case KafkaConfig.ZkSyncTimeMsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
-        case KafkaConfig.ZkEnableSecureAclsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_boolean")
-        case KafkaConfig.ZkMaxInFlightRequestsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "0")
-
-        case KafkaConfig.BrokerIdProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
-        case KafkaConfig.NumNetworkThreadsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "0")
-        case KafkaConfig.NumIoThreadsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "0")
-        case KafkaConfig.BackgroundThreadsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "0")
-        case KafkaConfig.QueuedMaxRequestsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "0")
-        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.ZkSessionTimeoutMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.ZkConnectionTimeoutMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.ZkSyncTimeMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.ZkEnableSecureAclsProp => assertPropertyInvalid(baseProperties, name, "not_a_boolean")
+        case KafkaConfig.ZkMaxInFlightRequestsProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0")
+        case KafkaConfig.ZkSslClientEnableProp => assertPropertyInvalid(baseProperties, name, "not_a_boolean")
+        case KafkaConfig.ZkClientCnxnSocketProp =>  //ignore string
+        case KafkaConfig.ZkSslKeyStoreLocationProp =>  //ignore string
+        case KafkaConfig.ZkSslKeyStorePasswordProp =>  //ignore string
+        case KafkaConfig.ZkSslKeyStoreTypeProp =>  //ignore string
+        case KafkaConfig.ZkSslTrustStoreLocationProp =>  //ignore string
+        case KafkaConfig.ZkSslTrustStorePasswordProp =>  //ignore string
+        case KafkaConfig.ZkSslTrustStoreTypeProp =>  //ignore string
+        case KafkaConfig.ZkSslProtocolProp =>  //ignore string
+        case KafkaConfig.ZkSslEnabledProtocolsProp =>  //ignore string
+        case KafkaConfig.ZkSslCipherSuitesProp =>  //ignore string
+        case KafkaConfig.ZkSslEndpointIdentificationAlgorithmProp => //ignore string
+        case KafkaConfig.ZkSslCrlEnableProp => assertPropertyInvalid(baseProperties, name, "not_a_boolean")
+        case KafkaConfig.ZkSslOcspEnableProp => assertPropertyInvalid(baseProperties, name, "not_a_boolean")
+
+        case KafkaConfig.BrokerIdProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.NumNetworkThreadsProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0")
+        case KafkaConfig.NumIoThreadsProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0")
+        case KafkaConfig.BackgroundThreadsProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0")
+        case KafkaConfig.QueuedMaxRequestsProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0")
+        case KafkaConfig.NumReplicaAlterLogDirsThreadsProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.QueuedMaxBytesProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.RequestTimeoutMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.ConnectionSetupTimeoutMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.ConnectionSetupTimeoutMaxMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
 
         case KafkaConfig.AuthorizerClassNameProp => //ignore string
         case KafkaConfig.CreateTopicPolicyClassNameProp => //ignore string
 
-        case KafkaConfig.PortProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
+        case KafkaConfig.PortProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
         case KafkaConfig.HostNameProp => // ignore string
         case KafkaConfig.AdvertisedHostNameProp => //ignore string
-        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.AdvertisedPortProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.SocketSendBufferBytesProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.SocketReceiveBufferBytesProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
         case KafkaConfig.MaxConnectionsPerIpOverridesProp =>
-          assertPropertyInvalid(getBaseProperties(), name, "127.0.0.1:not_a_number")
-        case KafkaConfig.ConnectionsMaxIdleMsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
+          assertPropertyInvalid(baseProperties, name, "127.0.0.1:not_a_number")
+        case KafkaConfig.ConnectionsMaxIdleMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.FailedAuthenticationDelayMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "-1")
 
-        case KafkaConfig.NumPartitionsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "0")
+        case KafkaConfig.NumPartitionsProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0")
         case KafkaConfig.LogDirsProp => // ignore string
         case KafkaConfig.LogDirProp => // ignore string
-        case KafkaConfig.LogSegmentBytesProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", Message.MinMessageOverhead - 1)
-
-        case KafkaConfig.LogRollTimeMillisProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "0")
-        case KafkaConfig.LogRollTimeHoursProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "0")
-
-        case KafkaConfig.LogRetentionTimeMillisProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "0")
-        case KafkaConfig.LogRetentionTimeMinutesProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "0")
-        case KafkaConfig.LogRetentionTimeHoursProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "0")
-
-        case KafkaConfig.LogRetentionBytesProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
-        case KafkaConfig.LogCleanupIntervalMsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "0")
-        case KafkaConfig.LogCleanupPolicyProp => assertPropertyInvalid(getBaseProperties(), name, "unknown_policy", "0")
-        case KafkaConfig.LogCleanerIoMaxBytesPerSecondProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
-        case KafkaConfig.LogCleanerDedupeBufferSizeProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "1024")
-        case KafkaConfig.LogCleanerDedupeBufferLoadFactorProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
-        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.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")
-        case KafkaConfig.LogFlushSchedulerIntervalMsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
-        case KafkaConfig.LogFlushIntervalMsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
-        case KafkaConfig.LogMessageTimestampDifferenceMaxMsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
-        case KafkaConfig.LogFlushStartOffsetCheckpointIntervalMsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
-        case KafkaConfig.NumRecoveryThreadsPerDataDirProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "0")
-        case KafkaConfig.AutoCreateTopicsEnableProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_boolean", "0")
-        case KafkaConfig.MinInSyncReplicasProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "0")
-        case KafkaConfig.ControllerSocketTimeoutMsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
-        case KafkaConfig.DefaultReplicationFactorProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
-        case KafkaConfig.ReplicaLagTimeMaxMsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
-        case KafkaConfig.ReplicaSocketTimeoutMsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "-2")
-        case KafkaConfig.ReplicaSocketReceiveBufferBytesProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
-        case KafkaConfig.ReplicaFetchMaxBytesProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
-        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.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")
-        case KafkaConfig.ProducerPurgatoryPurgeIntervalRequestsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
-        case KafkaConfig.DeleteRecordsPurgatoryPurgeIntervalRequestsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
-        case KafkaConfig.AutoLeaderRebalanceEnableProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_boolean", "0")
-        case KafkaConfig.LeaderImbalancePerBrokerPercentageProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
-        case KafkaConfig.LeaderImbalanceCheckIntervalSecondsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
-        case KafkaConfig.UncleanLeaderElectionEnableProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_boolean", "0")
-        case KafkaConfig.ControlledShutdownMaxRetriesProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
-        case KafkaConfig.ControlledShutdownRetryBackoffMsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
-        case KafkaConfig.ControlledShutdownEnableProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_boolean", "0")
-        case KafkaConfig.GroupMinSessionTimeoutMsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
-        case KafkaConfig.GroupMaxSessionTimeoutMsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
-        case KafkaConfig.GroupInitialRebalanceDelayMsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
-        case KafkaConfig.OffsetMetadataMaxSizeProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number")
-        case KafkaConfig.OffsetsLoadBufferSizeProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "0")
-        case KafkaConfig.OffsetsTopicReplicationFactorProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "0")
-        case KafkaConfig.OffsetsTopicPartitionsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "0")
-        case KafkaConfig.OffsetsTopicSegmentBytesProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "0")
-        case KafkaConfig.OffsetsTopicCompressionCodecProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "-1")
-        case KafkaConfig.OffsetsRetentionMinutesProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "0")
-        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.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")
-        case KafkaConfig.TransactionsLoadBufferSizeProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "0", "-2")
-        case KafkaConfig.TransactionsTopicPartitionsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "0", "-2")
-        case KafkaConfig.TransactionsTopicSegmentBytesProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "0", "-2")
-        case KafkaConfig.TransactionsTopicReplicationFactorProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "0", "-2")
-        case KafkaConfig.ProducerQuotaBytesPerSecondDefaultProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "0")
-        case KafkaConfig.ConsumerQuotaBytesPerSecondDefaultProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "0")
-        case KafkaConfig.NumQuotaSamplesProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "0")
-        case KafkaConfig.QuotaWindowSizeSecondsProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "0")
-        case KafkaConfig.DeleteTopicEnableProp => assertPropertyInvalid(getBaseProperties(), name, "not_a_boolean", "0")
-
-        case KafkaConfig.MetricNumSamplesProp => assertPropertyInvalid(getBaseProperties, name, "not_a_number", "-1", "0")
-        case KafkaConfig.MetricSampleWindowMsProp => assertPropertyInvalid(getBaseProperties, name, "not_a_number", "-1", "0")
+        case KafkaConfig.LogSegmentBytesProp => assertPropertyInvalid(baseProperties, name, "not_a_number", Records.LOG_OVERHEAD - 1)
+
+        case KafkaConfig.LogRollTimeMillisProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0")
+        case KafkaConfig.LogRollTimeHoursProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0")
+
+        case KafkaConfig.LogRetentionTimeMillisProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0")
+        case KafkaConfig.LogRetentionTimeMinutesProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0")
+        case KafkaConfig.LogRetentionTimeHoursProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0")
+
+        case KafkaConfig.LogRetentionBytesProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.LogCleanupIntervalMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0")
+        case KafkaConfig.LogCleanupPolicyProp => assertPropertyInvalid(baseProperties, name, "unknown_policy", "0")
+        case KafkaConfig.LogCleanerIoMaxBytesPerSecondProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.LogCleanerDedupeBufferSizeProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "1024")
+        case KafkaConfig.LogCleanerDedupeBufferLoadFactorProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.LogCleanerEnableProp => assertPropertyInvalid(baseProperties, name, "not_a_boolean")
+        case KafkaConfig.LogCleanerDeleteRetentionMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.LogCleanerMinCompactionLagMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.LogCleanerMaxCompactionLagMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.LogCleanerMinCleanRatioProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.LogIndexSizeMaxBytesProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "3")
+        case KafkaConfig.LogFlushIntervalMessagesProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0")
+        case KafkaConfig.LogFlushSchedulerIntervalMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.LogFlushIntervalMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.LogMessageTimestampDifferenceMaxMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.LogFlushStartOffsetCheckpointIntervalMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.NumRecoveryThreadsPerDataDirProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0")
+        case KafkaConfig.AutoCreateTopicsEnableProp => assertPropertyInvalid(baseProperties, name, "not_a_boolean", "0")
+        case KafkaConfig.MinInSyncReplicasProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0")
+        case KafkaConfig.ControllerSocketTimeoutMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.DefaultReplicationFactorProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.ReplicaLagTimeMaxMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.ReplicaSocketTimeoutMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "-2")
+        case KafkaConfig.ReplicaSocketReceiveBufferBytesProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.ReplicaFetchMaxBytesProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.ReplicaFetchWaitMaxMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.ReplicaFetchMinBytesProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.ReplicaFetchResponseMaxBytesProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.ReplicaSelectorClassProp => // Ignore string
+        case KafkaConfig.NumReplicaFetchersProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.ReplicaHighWatermarkCheckpointIntervalMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.FetchPurgatoryPurgeIntervalRequestsProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.ProducerPurgatoryPurgeIntervalRequestsProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.DeleteRecordsPurgatoryPurgeIntervalRequestsProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.AutoLeaderRebalanceEnableProp => assertPropertyInvalid(baseProperties, name, "not_a_boolean", "0")
+        case KafkaConfig.LeaderImbalancePerBrokerPercentageProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.LeaderImbalanceCheckIntervalSecondsProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.UncleanLeaderElectionEnableProp => assertPropertyInvalid(baseProperties, name, "not_a_boolean", "0")
+        case KafkaConfig.ControlledShutdownMaxRetriesProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.ControlledShutdownRetryBackoffMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.ControlledShutdownEnableProp => assertPropertyInvalid(baseProperties, name, "not_a_boolean", "0")
+        case KafkaConfig.GroupMinSessionTimeoutMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.GroupMaxSessionTimeoutMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.GroupInitialRebalanceDelayMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.GroupMaxSizeProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0", "-1")
+        case KafkaConfig.OffsetMetadataMaxSizeProp => assertPropertyInvalid(baseProperties, name, "not_a_number")
+        case KafkaConfig.OffsetsLoadBufferSizeProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0")
+        case KafkaConfig.OffsetsTopicReplicationFactorProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0")
+        case KafkaConfig.OffsetsTopicPartitionsProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0")
+        case KafkaConfig.OffsetsTopicSegmentBytesProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0")
+        case KafkaConfig.OffsetsTopicCompressionCodecProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "-1")
+        case KafkaConfig.OffsetsRetentionMinutesProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0")
+        case KafkaConfig.OffsetsRetentionCheckIntervalMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0")
+        case KafkaConfig.OffsetCommitTimeoutMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0")
+        case KafkaConfig.OffsetCommitRequiredAcksProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "-2")
+        case KafkaConfig.TransactionalIdExpirationMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0", "-2")
+        case KafkaConfig.TransactionsMaxTimeoutMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0", "-2")
+        case KafkaConfig.TransactionsTopicMinISRProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0", "-2")
+        case KafkaConfig.TransactionsLoadBufferSizeProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0", "-2")
+        case KafkaConfig.TransactionsTopicPartitionsProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0", "-2")
+        case KafkaConfig.TransactionsTopicSegmentBytesProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0", "-2")
+        case KafkaConfig.TransactionsTopicReplicationFactorProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0", "-2")
+        case KafkaConfig.ProducerQuotaBytesPerSecondDefaultProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0")
+        case KafkaConfig.ConsumerQuotaBytesPerSecondDefaultProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0")
+        case KafkaConfig.NumQuotaSamplesProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0")
+        case KafkaConfig.QuotaWindowSizeSecondsProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0")
+        case KafkaConfig.DeleteTopicEnableProp => assertPropertyInvalid(baseProperties, name, "not_a_boolean", "0")
+
+        case KafkaConfig.MetricNumSamplesProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "-1", "0")
+        case KafkaConfig.MetricSampleWindowMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "-1", "0")
         case KafkaConfig.MetricReporterClassesProp => // ignore string
         case KafkaConfig.MetricRecordingLevelProp => // ignore string
         case KafkaConfig.RackProp => // ignore string
         //SSL Configs
         case KafkaConfig.PrincipalBuilderClassProp =>
+        case KafkaConfig.ConnectionsMaxReauthMsProp =>
         case KafkaConfig.SslProtocolProp => // ignore string
         case KafkaConfig.SslProviderProp => // ignore string
         case KafkaConfig.SslEnabledProtocolsProp =>
@@ -655,28 +740,146 @@ class KafkaConfigTest {
         case KafkaConfig.SslKeystoreLocationProp => // ignore string
         case KafkaConfig.SslKeystorePasswordProp => // ignore string
         case KafkaConfig.SslKeyPasswordProp => // ignore string
+        case KafkaConfig.SslKeystoreCertificateChainProp => // ignore string
+        case KafkaConfig.SslKeystoreKeyProp => // ignore string
         case KafkaConfig.SslTruststoreTypeProp => // ignore string
         case KafkaConfig.SslTruststorePasswordProp => // ignore string
         case KafkaConfig.SslTruststoreLocationProp => // ignore string
+        case KafkaConfig.SslTruststoreCertificatesProp => // ignore string
         case KafkaConfig.SslKeyManagerAlgorithmProp =>
         case KafkaConfig.SslTrustManagerAlgorithmProp =>
         case KafkaConfig.SslClientAuthProp => // ignore string
         case KafkaConfig.SslEndpointIdentificationAlgorithmProp => // ignore string
         case KafkaConfig.SslSecureRandomImplementationProp => // ignore string
         case KafkaConfig.SslCipherSuitesProp => // ignore string
+        case KafkaConfig.SslPrincipalMappingRulesProp => // ignore string
 
         //Sasl Configs
         case KafkaConfig.SaslMechanismInterBrokerProtocolProp => // ignore
         case KafkaConfig.SaslEnabledMechanismsProp =>
+        case KafkaConfig.SaslClientCallbackHandlerClassProp =>
+        case KafkaConfig.SaslServerCallbackHandlerClassProp =>
+        case KafkaConfig.SaslLoginClassProp =>
+        case KafkaConfig.SaslLoginCallbackHandlerClassProp =>
         case KafkaConfig.SaslKerberosServiceNameProp => // ignore string
         case KafkaConfig.SaslKerberosKinitCmdProp =>
         case KafkaConfig.SaslKerberosTicketRenewWindowFactorProp =>
         case KafkaConfig.SaslKerberosTicketRenewJitterProp =>
         case KafkaConfig.SaslKerberosMinTimeBeforeReloginProp =>
         case KafkaConfig.SaslKerberosPrincipalToLocalRulesProp => // ignore string
-        case _ => assertPropertyInvalid(getBaseProperties(), name, "not_a_number", "-1")
+        case KafkaConfig.SaslJaasConfigProp =>
+        case KafkaConfig.SaslLoginRefreshWindowFactorProp =>
+        case KafkaConfig.SaslLoginRefreshWindowJitterProp =>
+        case KafkaConfig.SaslLoginRefreshMinPeriodSecondsProp =>
+        case KafkaConfig.SaslLoginRefreshBufferSecondsProp =>
+
+        // Security config
+        case KafkaConfig.securityProviderClassProp =>
+
+        // Password encoder configs
+        case KafkaConfig.PasswordEncoderSecretProp =>
+        case KafkaConfig.PasswordEncoderOldSecretProp =>
+        case KafkaConfig.PasswordEncoderKeyFactoryAlgorithmProp =>
+        case KafkaConfig.PasswordEncoderCipherAlgorithmProp =>
+        case KafkaConfig.PasswordEncoderKeyLengthProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "-1", "0")
+        case KafkaConfig.PasswordEncoderIterationsProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "-1", "0")
+
+        //delegation token configs
+        case KafkaConfig.DelegationTokenSecretKeyAliasProp => // ignore
+        case KafkaConfig.DelegationTokenSecretKeyProp => // ignore
+        case KafkaConfig.DelegationTokenMaxLifeTimeProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0")
+        case KafkaConfig.DelegationTokenExpiryTimeMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0")
+        case KafkaConfig.DelegationTokenExpiryCheckIntervalMsProp => assertPropertyInvalid(baseProperties, name, "not_a_number", "0")
+
+        //Kafka Yammer metrics reporter configs
+        case KafkaConfig.KafkaMetricsReporterClassesProp => // ignore
+        case KafkaConfig.KafkaMetricsPollingIntervalSecondsProp => //ignore
+
+        case _ => assertPropertyInvalid(baseProperties, name, "not_a_number", "-1")
       }
-    })
+    }
+  }
+
+  @Test
+  def testDynamicLogConfigs(): Unit = {
+    def baseProperties: Properties = {
+      val validRequiredProperties = new Properties()
+      validRequiredProperties.put(KafkaConfig.ZkConnectProp, "127.0.0.1:2181")
+      validRequiredProperties
+    }
+
+    val props = baseProperties
+    val config = KafkaConfig.fromProps(props)
+
+    def assertDynamic(property: String, value: Any, accessor: () => Any): Unit = {
+      val initial = accessor()
+      props.put(property, value.toString)
+      config.updateCurrentConfig(new KafkaConfig(props))
+      assertNotEquals(initial, accessor())
+    }
+
+    // Test dynamic log config values can be correctly passed through via KafkaConfig to LogConfig
+    // Every log config prop must be explicitly accounted for here.
+    // A value other than the default value for this config should be set to ensure that we can check whether
+    // the value is dynamically updatable.
+    LogConfig.TopicConfigSynonyms.foreach { case (logConfig, kafkaConfigProp) =>
+      logConfig match {
+        case LogConfig.CleanupPolicyProp =>
+          assertDynamic(kafkaConfigProp, Defaults.Compact, () => config.logCleanupPolicy)
+        case LogConfig.CompressionTypeProp =>
+          assertDynamic(kafkaConfigProp, "lz4", () => config.compressionType)
+        case LogConfig.SegmentBytesProp =>
+          assertDynamic(kafkaConfigProp, 10000, () => config.logSegmentBytes)
+        case LogConfig.SegmentMsProp =>
+          assertDynamic(kafkaConfigProp, 10001L, () => config.logRollTimeMillis)
+        case LogConfig.DeleteRetentionMsProp =>
+          assertDynamic(kafkaConfigProp, 10002L, () => config.logCleanerDeleteRetentionMs)
+        case LogConfig.FileDeleteDelayMsProp =>
+          assertDynamic(kafkaConfigProp, 10003L, () => config.logDeleteDelayMs)
+        case LogConfig.FlushMessagesProp =>
+          assertDynamic(kafkaConfigProp, 10004L, () => config.logFlushIntervalMessages)
+        case LogConfig.FlushMsProp =>
+          assertDynamic(kafkaConfigProp, 10005L, () => config.logFlushIntervalMs)
+        case LogConfig.MaxCompactionLagMsProp =>
+          assertDynamic(kafkaConfigProp, 10006L, () => config.logCleanerMaxCompactionLagMs)
+        case LogConfig.IndexIntervalBytesProp =>
+          assertDynamic(kafkaConfigProp, 10007, () => config.logIndexIntervalBytes)
+        case LogConfig.MaxMessageBytesProp =>
+          assertDynamic(kafkaConfigProp, 10008, () => config.messageMaxBytes)
+        case LogConfig.MessageDownConversionEnableProp =>
+          assertDynamic(kafkaConfigProp, false, () => config.logMessageDownConversionEnable)
+        case LogConfig.MessageTimestampDifferenceMaxMsProp =>
+          assertDynamic(kafkaConfigProp, 10009, () => config.logMessageTimestampDifferenceMaxMs)
+        case LogConfig.MessageTimestampTypeProp =>
+          assertDynamic(kafkaConfigProp, "LogAppendTime", () => config.logMessageTimestampType.name)
+        case LogConfig.MinCleanableDirtyRatioProp =>
+          assertDynamic(kafkaConfigProp, 0.01, () => config.logCleanerMinCleanRatio)
+        case LogConfig.MinCompactionLagMsProp =>
+          assertDynamic(kafkaConfigProp, 10010L, () => config.logCleanerMinCompactionLagMs)
+        case LogConfig.MinInSyncReplicasProp =>
+          assertDynamic(kafkaConfigProp, 4, () => config.minInSyncReplicas)
+        case LogConfig.PreAllocateEnableProp =>
+          assertDynamic(kafkaConfigProp, true, () => config.logPreAllocateEnable)
+        case LogConfig.RetentionBytesProp =>
+          assertDynamic(kafkaConfigProp, 10011L, () => config.logRetentionBytes)
+        case LogConfig.RetentionMsProp =>
+          assertDynamic(kafkaConfigProp, 10012L, () => config.logRetentionTimeMillis)
+        case LogConfig.SegmentIndexBytesProp =>
+          assertDynamic(kafkaConfigProp, 10013, () => config.logIndexSizeMaxBytes)
+        case LogConfig.SegmentJitterMsProp =>
+          assertDynamic(kafkaConfigProp, 10014L, () => config.logRollTimeJitterMillis)
+        case LogConfig.UncleanLeaderElectionEnableProp =>
+          assertDynamic(kafkaConfigProp, true, () => config.uncleanLeaderElectionEnable)
+        case LogConfig.MessageFormatVersionProp =>
+        // not dynamically updatable
+        case LogConfig.FollowerReplicationThrottledReplicasProp =>
+        // topic only config
+        case LogConfig.LeaderReplicationThrottledReplicasProp =>
+        // topic only config
+        case prop =>
+          fail(prop + " must be explicitly checked for dynamic updatability. Note that LogConfig(s) require that KafkaConfig value lookups are dynamic and not static values.")
+      }
+    }
   }
 
   @Test
@@ -718,17 +921,36 @@ class KafkaConfigTest {
     assertEquals(123L, config.logFlushIntervalMs)
     assertEquals(SnappyCompressionCodec, config.offsetsTopicCompressionCodec)
     assertEquals(Sensor.RecordingLevel.DEBUG.toString, config.metricRecordingLevel)
+    assertEquals(false, config.tokenAuthEnabled)
+    assertEquals(7 * 24 * 60L * 60L * 1000L, config.delegationTokenMaxLifeMs)
+    assertEquals(24 * 60L * 60L * 1000L, config.delegationTokenExpiryTimeMs)
+    assertEquals(1 * 60L * 1000L * 60, config.delegationTokenExpiryCheckIntervalMs)
+
+    defaults.put(KafkaConfig.DelegationTokenSecretKeyProp, "1234567890")
+    val config1 = KafkaConfig.fromProps(defaults)
+    assertEquals(true, config1.tokenAuthEnabled)
   }
 
   @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")
     assertFalse(isValidKafkaConfig(props))
   }
 
-  private def assertPropertyInvalid(validRequiredProps: => Properties, name: String, values: Any*) {
+  @Test
+  def testMaxConnectionsPerIpProp(): Unit = {
+    val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)
+    props.put(KafkaConfig.MaxConnectionsPerIpProp, "0")
+    assertFalse(isValidKafkaConfig(props))
+    props.put(KafkaConfig.MaxConnectionsPerIpOverridesProp, "127.0.0.1:100")
+    assertTrue(isValidKafkaConfig(props))
+    props.put(KafkaConfig.MaxConnectionsPerIpOverridesProp, "127.0.0.0#:100")
+    assertFalse(isValidKafkaConfig(props))
+  }
+
+  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 dfcb4ac279990..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,20 +80,20 @@ 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("kafka.metrics.reporters", "kafka.server.KafkaMetricReporterClusterIdTest$MockKafkaMetricsReporter")
+    props.setProperty(KafkaConfig.KafkaMetricsReporterClassesProp, "kafka.server.KafkaMetricReporterClusterIdTest$MockKafkaMetricsReporter")
     props.setProperty(KafkaConfig.MetricReporterClassesProp, "kafka.server.KafkaMetricReporterClusterIdTest$MockBrokerMetricsReporter")
     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
new file mode 100644
index 0000000000000..4458c74424ff5
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/KafkaMetricReporterExceptionHandlingTest.scala
@@ -0,0 +1,119 @@
+/**
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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.net.Socket
+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.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.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 brokerCount: Int = 1
+
+  override def brokerPropertyOverrides(properties: Properties): Unit = {
+    properties.put(KafkaConfig.MetricReporterClassesProp, classOf[KafkaMetricReporterExceptionHandlingTest.BadReporter].getName + "," + classOf[KafkaMetricReporterExceptionHandlingTest.GoodReporter].getName)
+  }
+
+  @Before
+  override def setUp(): Unit = {
+    super.setUp()
+
+    // need a quota prop to register a "throttle-time" metrics after server startup
+    val quotaProps = new Properties()
+    quotaProps.put(DynamicConfig.Client.RequestPercentageOverrideProp, "0.1")
+    adminZkClient.changeClientIdConfig("", quotaProps)
+  }
+
+  @After
+  override def tearDown(): Unit = {
+    KafkaMetricReporterExceptionHandlingTest.goodReporterRegistered.set(0)
+    KafkaMetricReporterExceptionHandlingTest.badReporterRegistered.set(0)
+    
+    super.tearDown()
+  }
+
+  @Test
+  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 listGroupsRequest = new ListGroupsRequest.Builder(new ListGroupsRequestData).build()
+        val listGroupsResponse = sendAndReceive[ListGroupsResponse](listGroupsRequest, socket)
+        val errors = listGroupsResponse.errorCounts()
+        assertEquals(Collections.singletonMap(Errors.NONE, 1), errors)
+        assertEquals(KafkaMetricReporterExceptionHandlingTest.goodReporterRegistered.get, KafkaMetricReporterExceptionHandlingTest.badReporterRegistered.get)
+        assertTrue(KafkaMetricReporterExceptionHandlingTest.goodReporterRegistered.get > 0)
+      }
+    } finally {
+      socket.close()
+    }
+  }
+}
+
+object KafkaMetricReporterExceptionHandlingTest {
+  var goodReporterRegistered = new AtomicInteger
+  var badReporterRegistered = new AtomicInteger
+
+  class GoodReporter extends MetricsReporter {
+
+    def configure(configs: java.util.Map[String, _]): Unit = {
+    }
+
+    def init(metrics: java.util.List[KafkaMetric]): Unit = {
+    }
+
+    def metricChange(metric: KafkaMetric): Unit = {
+      if (metric.metricName.group == "Request") {
+        goodReporterRegistered.incrementAndGet
+      }
+    }
+
+    def metricRemoval(metric: KafkaMetric): Unit = {
+    }
+
+    def close(): Unit = {
+    }
+  }
+
+  class BadReporter extends GoodReporter {
+
+    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/KafkaMetricsReporterTest.scala b/core/src/test/scala/unit/kafka/server/KafkaMetricsReporterTest.scala
new file mode 100644
index 0000000000000..4e0c9f862afe4
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/KafkaMetricsReporterTest.scala
@@ -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 kafka.server
+
+import java.util
+
+import java.util.concurrent.atomic.AtomicReference
+
+import kafka.utils.{CoreUtils, TestUtils}
+import kafka.zk.ZooKeeperTestHarness
+import org.apache.kafka.common.metrics.{KafkaMetric, MetricsContext, MetricsReporter}
+import org.junit.Assert.{assertEquals}
+import org.junit.{After, Before, Test}
+import org.junit.Assert._
+
+
+object KafkaMetricsReporterTest {
+  val setupError = new AtomicReference[String]("")
+
+  class MockMetricsReporter extends MetricsReporter {
+    def init(metrics: util.List[KafkaMetric]): Unit = {}
+
+    def metricChange(metric: KafkaMetric): Unit = {}
+
+    def metricRemoval(metric: KafkaMetric): Unit = {}
+
+    override def close(): Unit = {}
+
+    override def contextChange(metricsContext: MetricsContext): Unit = {
+      //read jmxPrefix
+
+      MockMetricsReporter.JMXPREFIX.set(metricsContext.contextLabels().get("_namespace").toString)
+      MockMetricsReporter.CLUSTERID.set(metricsContext.contextLabels().get("kafka.cluster.id").toString)
+      MockMetricsReporter.BROKERID.set(metricsContext.contextLabels().get("kafka.broker.id").toString)
+    }
+
+    override def configure(configs: util.Map[String, _]): Unit = {}
+
+  }
+
+  object MockMetricsReporter {
+    val JMXPREFIX: AtomicReference[String] = new AtomicReference[String]
+    val BROKERID : AtomicReference[String] = new AtomicReference[String]
+    val CLUSTERID : AtomicReference[String] = new AtomicReference[String]
+  }
+}
+
+class KafkaMetricsReporterTest extends ZooKeeperTestHarness {
+  var server: KafkaServerStartable = null
+  var config: KafkaConfig = null
+
+  @Before
+  override def setUp(): Unit = {
+    super.setUp()
+    val props = TestUtils.createBrokerConfig(1, zkConnect)
+    props.setProperty(KafkaConfig.MetricReporterClassesProp, "kafka.server.KafkaMetricsReporterTest$MockMetricsReporter")
+    props.setProperty(KafkaConfig.BrokerIdGenerationEnableProp, "true")
+    props.setProperty(KafkaConfig.BrokerIdProp, "-1")
+    config = KafkaConfig.fromProps(props)
+    server = KafkaServerStartable.fromProps(props, threadNamePrefix = Option(this.getClass.getName))
+    server.startup()
+  }
+
+  @Test
+  def testMetricsContextNamespacePresent(): Unit = {
+    assertNotNull(KafkaMetricsReporterTest.MockMetricsReporter.CLUSTERID)
+    assertNotNull(KafkaMetricsReporterTest.MockMetricsReporter.BROKERID)
+    assertNotNull(KafkaMetricsReporterTest.MockMetricsReporter.JMXPREFIX)
+    assertEquals("kafka.server", KafkaMetricsReporterTest.MockMetricsReporter.JMXPREFIX.get())
+
+    server.shutdown()
+    TestUtils.assertNoNonDaemonThreads(this.getClass.getName)
+  }
+
+  @After
+  override def tearDown(): Unit = {
+    server.shutdown()
+    CoreUtils.delete(config.logDirs)
+    super.tearDown()
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/server/KafkaServerTest.scala b/core/src/test/scala/unit/kafka/server/KafkaServerTest.scala
new file mode 100755
index 0000000000000..7342be6928cd3
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/KafkaServerTest.scala
@@ -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 kafka.server
+
+import java.util.Properties
+
+import kafka.utils.TestUtils
+import kafka.zk.ZooKeeperTestHarness
+import org.apache.zookeeper.client.ZKClientConfig
+import org.junit.Test
+import org.junit.Assert.assertEquals
+import org.scalatest.Assertions.intercept
+
+class KafkaServerTest extends ZooKeeperTestHarness {
+
+  @Test
+  def testAlreadyRegisteredAdvertisedListeners(): Unit = {
+    //start a server with a advertised listener
+    val server1 = createServer(1, "myhost", TestUtils.RandomPort)
+
+    //start a server with same advertised listener
+    intercept[IllegalArgumentException] {
+      createServer(2, "myhost", TestUtils.boundPort(server1))
+    }
+
+    //start a server with same host but with different port
+    val server2 = createServer(2, "myhost", TestUtils.RandomPort)
+
+    TestUtils.shutdownServers(Seq(server1, server2))
+  }
+
+  @Test
+  def testCreatesProperZkTlsConfigWhenDisabled(): Unit = {
+    val props = new Properties
+    props.put(KafkaConfig.ZkConnectProp, zkConnect) // required, otherwise we would leave it out
+    props.put(KafkaConfig.ZkSslClientEnableProp, "false")
+    assertEquals(None, KafkaServer.zkClientConfigFromKafkaConfig(KafkaConfig.fromProps(props)))
+  }
+
+  @Test
+  def testCreatesProperZkTlsConfigWithTrueValues(): Unit = {
+    val props = new Properties
+    props.put(KafkaConfig.ZkConnectProp, zkConnect) // required, otherwise we would leave it out
+    // should get correct config for all properties if TLS is enabled
+    val someValue = "some_value"
+    def kafkaConfigValueToSet(kafkaProp: String) : String = kafkaProp match {
+      case KafkaConfig.ZkSslClientEnableProp | KafkaConfig.ZkSslCrlEnableProp | KafkaConfig.ZkSslOcspEnableProp => "true"
+      case KafkaConfig.ZkSslEndpointIdentificationAlgorithmProp => "HTTPS"
+      case _ => someValue
+    }
+    KafkaConfig.ZkSslConfigToSystemPropertyMap.keys.foreach(kafkaProp => props.put(kafkaProp, kafkaConfigValueToSet(kafkaProp)))
+    val zkClientConfig: Option[ZKClientConfig] = KafkaServer.zkClientConfigFromKafkaConfig(KafkaConfig.fromProps(props))
+    // now check to make sure the values were set correctly
+    def zkClientValueToExpect(kafkaProp: String) : String = kafkaProp match {
+      case KafkaConfig.ZkSslClientEnableProp | KafkaConfig.ZkSslCrlEnableProp | KafkaConfig.ZkSslOcspEnableProp => "true"
+      case KafkaConfig.ZkSslEndpointIdentificationAlgorithmProp => "true"
+      case _ => someValue
+    }
+    KafkaConfig.ZkSslConfigToSystemPropertyMap.keys.foreach(kafkaProp =>
+      assertEquals(zkClientValueToExpect(kafkaProp), zkClientConfig.get.getProperty(KafkaConfig.ZkSslConfigToSystemPropertyMap(kafkaProp))))
+  }
+
+  @Test
+  def testCreatesProperZkTlsConfigWithFalseAndListValues(): Unit = {
+    val props = new Properties
+    props.put(KafkaConfig.ZkConnectProp, zkConnect) // required, otherwise we would leave it out
+    // should get correct config for all properties if TLS is enabled
+    val someValue = "some_value"
+    def kafkaConfigValueToSet(kafkaProp: String) : String = kafkaProp match {
+      case KafkaConfig.ZkSslClientEnableProp => "true"
+      case KafkaConfig.ZkSslCrlEnableProp | KafkaConfig.ZkSslOcspEnableProp => "false"
+      case KafkaConfig.ZkSslEndpointIdentificationAlgorithmProp => ""
+      case KafkaConfig.ZkSslEnabledProtocolsProp | KafkaConfig.ZkSslCipherSuitesProp => "A,B"
+      case _ => someValue
+    }
+    KafkaConfig.ZkSslConfigToSystemPropertyMap.keys.foreach(kafkaProp => props.put(kafkaProp, kafkaConfigValueToSet(kafkaProp)))
+    val zkClientConfig: Option[ZKClientConfig] = KafkaServer.zkClientConfigFromKafkaConfig(KafkaConfig.fromProps(props))
+    // now check to make sure the values were set correctly
+    def zkClientValueToExpect(kafkaProp: String) : String = kafkaProp match {
+      case KafkaConfig.ZkSslClientEnableProp => "true"
+      case KafkaConfig.ZkSslCrlEnableProp | KafkaConfig.ZkSslOcspEnableProp => "false"
+      case KafkaConfig.ZkSslEndpointIdentificationAlgorithmProp => "false"
+      case KafkaConfig.ZkSslEnabledProtocolsProp | KafkaConfig.ZkSslCipherSuitesProp => "A,B"
+      case _ => someValue
+    }
+    KafkaConfig.ZkSslConfigToSystemPropertyMap.keys.foreach(kafkaProp =>
+      assertEquals(zkClientValueToExpect(kafkaProp), zkClientConfig.get.getProperty(KafkaConfig.ZkSslConfigToSystemPropertyMap(kafkaProp))))
+  }
+
+  def createServer(nodeId: Int, hostName: String, port: Int): KafkaServer = {
+    val props = TestUtils.createBrokerConfig(nodeId, zkConnect)
+    props.put(KafkaConfig.AdvertisedListenersProp, s"PLAINTEXT://$hostName:$port")
+    val kafkaConfig = KafkaConfig.fromProps(props)
+    TestUtils.createServer(kafkaConfig)
+  }
+
+}
diff --git a/core/src/test/scala/unit/kafka/server/LeaderElectionTest.scala b/core/src/test/scala/unit/kafka/server/LeaderElectionTest.scala
index 56f5f6d9883b9..a3eb5d7572821 100755
--- a/core/src/test/scala/unit/kafka/server/LeaderElectionTest.scala
+++ b/core/src/test/scala/unit/kafka/server/LeaderElectionTest.scala
@@ -19,7 +19,7 @@ package kafka.server
 
 import org.apache.kafka.common.TopicPartition
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 import kafka.api.LeaderAndIsr
 import org.apache.kafka.common.requests._
 import org.junit.Assert._
@@ -28,6 +28,7 @@ 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.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 +45,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 +61,7 @@ class LeaderElectionTest extends ZooKeeperTestHarness {
   }
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     TestUtils.shutdownServers(servers)
     super.tearDown()
   }
@@ -71,55 +72,50 @@ 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(zkUtils, topic, partitionReplicaAssignment = Map(0 -> Seq(0, 1)), servers = servers)(0)
+    val leader1 = createTopic(zkClient, topic, partitionReplicaAssignment = Map(0 -> Seq(0, 1)), servers = servers)(0)
 
-    val leaderEpoch1 = zkUtils.getEpochForPartition(topic, partitionId)
-    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)
+    val leaderEpoch1 = zkClient.getEpochForPartition(new TopicPartition(topic, partitionId)).get
+    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()
-    // check if leader moves to the other server
-    val leader2 = waitUntilLeaderIsElectedOrChanged(zkUtils, topic, partitionId,
-                                                    oldLeaderOpt = if (leader1 == 0) None else Some(leader1))
-    val leaderEpoch2 = zkUtils.getEpochForPartition(topic, partitionId)
-    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()
+    // 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 = Some(leader1))
+    val leaderEpoch2 = zkClient.getEpochForPartition(new TopicPartition(topic, partitionId)).get
+    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(zkUtils, topic, partitionId,
-                                                    oldLeaderOpt = if (leader2 == 1) None else Some(leader2))
-    val leaderEpoch3 = zkUtils.getEpochForPartition(topic, partitionId)
-    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)
+    val leader3 = waitUntilLeaderIsElectedOrChanged(zkClient, topic, partitionId, oldLeaderOpt = Some(leader2))
+    val leaderEpoch3 = zkClient.getEpochForPartition(new TopicPartition(topic, partitionId)).get
+    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
 
     // create topic with 1 partition, 2 replicas, one on each broker
-    val leader1 = createTopic(zkUtils, topic, partitionReplicaAssignment = Map(0 -> Seq(0, 1)), servers = servers)(0)
+    val leader1 = createTopic(zkClient, topic, partitionReplicaAssignment = Map(0 -> Seq(0, 1)), servers = servers)(0)
 
-    val leaderEpoch1 = zkUtils.getEpochForPartition(topic, partitionId)
+    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
@@ -132,28 +128,36 @@ class LeaderElectionTest extends ZooKeeperTestHarness {
     val controllerConfig = KafkaConfig.fromProps(TestUtils.createBrokerConfig(controllerId, zkConnect))
     val securityProtocol = SecurityProtocol.PLAINTEXT
     val listenerName = ListenerName.forSecurityProtocol(securityProtocol)
-    val brokers = servers.map(s => new Broker(s.config.brokerId, "localhost", TestUtils.boundPort(s), listenerName,
-      securityProtocol))
-    val nodes = brokers.map(_.getNode(listenerName))
+    val brokerAndEpochs = servers.map(s =>
+      (new Broker(s.config.brokerId, "localhost", TestUtils.boundPort(s), listenerName, securityProtocol),
+        s.kafkaController.brokerEpoch)).toMap
+    val nodes = brokerAndEpochs.keys.map(_.node(listenerName))
 
     val controllerContext = new ControllerContext
-    controllerContext.liveBrokers = brokers.toSet
+    controllerContext.setLiveBrokers(brokerAndEpochs)
     val metrics = new Metrics
     val controllerChannelManager = new ControllerChannelManager(controllerContext, controllerConfig, Time.SYSTEM,
       metrics, new StateChangeLogger(controllerId, inControllerContext = true, None))
     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, partitionStates.asJava, nodes.toSet.asJava)
+        ApiKeys.LEADER_AND_ISR.latestVersion, controllerId, staleControllerEpoch,
+        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/ListOffsetsRequestTest.scala b/core/src/test/scala/unit/kafka/server/ListOffsetsRequestTest.scala
new file mode 100644
index 0000000000000..ce324c77cd9a2
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/ListOffsetsRequestTest.scala
@@ -0,0 +1,219 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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
+
+import kafka.utils.TestUtils
+import org.apache.kafka.common.message.ListOffsetRequestData.{ListOffsetPartition, ListOffsetTopic}
+import org.apache.kafka.common.protocol.{ApiKeys, Errors}
+import org.apache.kafka.common.requests.{ListOffsetRequest, ListOffsetResponse}
+import org.apache.kafka.common.{IsolationLevel, TopicPartition}
+import org.junit.Assert._
+import org.junit.Test
+
+import scala.jdk.CollectionConverters._
+
+class ListOffsetsRequestTest extends BaseRequestTest {
+
+  val topic = "topic"
+  val partition = new TopicPartition(topic, 0)
+
+  @Test
+  def testListOffsetsErrorCodes(): Unit = {
+    val targetTimes = List(new ListOffsetTopic()
+      .setName(topic)
+      .setPartitions(List(new ListOffsetPartition()
+        .setPartitionIndex(partition.partition)
+        .setTimestamp(ListOffsetRequest.EARLIEST_TIMESTAMP)
+        .setCurrentLeaderEpoch(0)).asJava)).asJava
+
+    val consumerRequest = ListOffsetRequest.Builder
+      .forConsumer(false, IsolationLevel.READ_UNCOMMITTED)
+      .setTargetTimes(targetTimes)
+      .build()
+
+    val replicaRequest = ListOffsetRequest.Builder
+      .forReplica(ApiKeys.LIST_OFFSETS.latestVersion, servers.head.config.brokerId)
+      .setTargetTimes(targetTimes)
+      .build()
+
+    val debugReplicaRequest = ListOffsetRequest.Builder
+      .forReplica(ApiKeys.LIST_OFFSETS.latestVersion, ListOffsetRequest.DEBUGGING_REPLICA_ID)
+      .setTargetTimes(targetTimes)
+      .build()
+
+    // Unknown topic
+    val randomBrokerId = servers.head.config.brokerId
+    assertResponseError(Errors.UNKNOWN_TOPIC_OR_PARTITION, randomBrokerId, consumerRequest)
+    assertResponseError(Errors.UNKNOWN_TOPIC_OR_PARTITION, randomBrokerId, replicaRequest)
+    assertResponseError(Errors.UNKNOWN_TOPIC_OR_PARTITION, randomBrokerId, debugReplicaRequest)
+
+    val partitionToLeader = TestUtils.createTopic(zkClient, topic, numPartitions = 1, replicationFactor = 2, servers)
+    val replicas = zkClient.getReplicasForPartition(partition).toSet
+    val leader = partitionToLeader(partition.partition)
+    val follower = replicas.find(_ != leader).get
+    val nonReplica = servers.map(_.config.brokerId).find(!replicas.contains(_)).get
+
+    // Follower
+    assertResponseError(Errors.NOT_LEADER_OR_FOLLOWER, follower, consumerRequest)
+    assertResponseError(Errors.NOT_LEADER_OR_FOLLOWER, follower, replicaRequest)
+    assertResponseError(Errors.NONE, follower, debugReplicaRequest)
+
+    // Non-replica
+    assertResponseError(Errors.NOT_LEADER_OR_FOLLOWER, nonReplica, consumerRequest)
+    assertResponseError(Errors.NOT_LEADER_OR_FOLLOWER, nonReplica, replicaRequest)
+    assertResponseError(Errors.NOT_LEADER_OR_FOLLOWER, nonReplica, debugReplicaRequest)
+  }
+
+  def assertResponseErrorForEpoch(error: Errors, brokerId: Int, currentLeaderEpoch: Optional[Integer]): Unit = {
+    val listOffsetPartition = new ListOffsetPartition()
+      .setPartitionIndex(partition.partition)
+      .setTimestamp(ListOffsetRequest.EARLIEST_TIMESTAMP)
+    if (currentLeaderEpoch.isPresent)
+      listOffsetPartition.setCurrentLeaderEpoch(currentLeaderEpoch.get)
+    val targetTimes = List(new ListOffsetTopic()
+      .setName(topic)
+      .setPartitions(List(listOffsetPartition).asJava)).asJava
+    val request = ListOffsetRequest.Builder
+      .forConsumer(false, IsolationLevel.READ_UNCOMMITTED)
+      .setTargetTimes(targetTimes)
+      .build()
+    assertResponseError(error, brokerId, request)
+  }
+
+  @Test
+  def testCurrentEpochValidation(): Unit = {
+    val topic = "topic"
+    val topicPartition = new TopicPartition(topic, 0)
+    val partitionToLeader = TestUtils.createTopic(zkClient, topic, numPartitions = 1, replicationFactor = 3, servers)
+    val firstLeaderId = partitionToLeader(topicPartition.partition)
+
+    // We need a leader change in order to check epoch fencing since the first epoch is 0 and
+    // -1 is treated as having no epoch at all
+    killBroker(firstLeaderId)
+
+    // Check leader error codes
+    val secondLeaderId = TestUtils.awaitLeaderChange(servers, topicPartition, firstLeaderId)
+    val secondLeaderEpoch = TestUtils.findLeaderEpoch(secondLeaderId, topicPartition, servers)
+    assertResponseErrorForEpoch(Errors.NONE, secondLeaderId, Optional.empty())
+    assertResponseErrorForEpoch(Errors.NONE, secondLeaderId, Optional.of(secondLeaderEpoch))
+    assertResponseErrorForEpoch(Errors.FENCED_LEADER_EPOCH, secondLeaderId, Optional.of(secondLeaderEpoch - 1))
+    assertResponseErrorForEpoch(Errors.UNKNOWN_LEADER_EPOCH, secondLeaderId, Optional.of(secondLeaderEpoch + 1))
+
+    // Check follower error codes
+    val followerId = TestUtils.findFollowerId(topicPartition, servers)
+    assertResponseErrorForEpoch(Errors.NOT_LEADER_OR_FOLLOWER, followerId, Optional.empty())
+    assertResponseErrorForEpoch(Errors.NOT_LEADER_OR_FOLLOWER, followerId, Optional.of(secondLeaderEpoch))
+    assertResponseErrorForEpoch(Errors.UNKNOWN_LEADER_EPOCH, followerId, Optional.of(secondLeaderEpoch + 1))
+    assertResponseErrorForEpoch(Errors.FENCED_LEADER_EPOCH, followerId, Optional.of(secondLeaderEpoch - 1))
+  }
+
+  // -1 indicate "latest"
+  def fetchOffsetAndEpoch(serverId: Int,
+                          timestamp: Long,
+                          version: Short): (Long, Int) = {
+    val targetTimes = List(new ListOffsetTopic()
+      .setName(topic)
+      .setPartitions(List(new ListOffsetPartition()
+        .setPartitionIndex(partition.partition)
+        .setTimestamp(timestamp)).asJava)).asJava
+
+    val builder = ListOffsetRequest.Builder
+      .forConsumer(false, IsolationLevel.READ_UNCOMMITTED)
+      .setTargetTimes(targetTimes)
+
+    val request = if (version == -1) builder.build() else builder.build(version)
+
+    val response = sendRequest(serverId, request)
+    val partitionData = response.topics.asScala.find(_.name == topic).get
+      .partitions.asScala.find(_.partitionIndex == partition.partition).get
+
+    if (version == 0) {
+      if (partitionData.oldStyleOffsets().isEmpty)
+        (-1, partitionData.leaderEpoch)
+      else
+        (partitionData.oldStyleOffsets().asScala.head, partitionData.leaderEpoch)
+    } else
+      (partitionData.offset, partitionData.leaderEpoch)
+  }
+
+  @Test
+  def testResponseIncludesLeaderEpoch(): Unit = {
+    val partitionToLeader = TestUtils.createTopic(zkClient, topic, numPartitions = 1, replicationFactor = 3, servers)
+    val firstLeaderId = partitionToLeader(partition.partition)
+
+    TestUtils.generateAndProduceMessages(servers, topic, 10)
+
+    assertEquals((0L, 0), fetchOffsetAndEpoch(firstLeaderId, 0L, -1))
+    assertEquals((0L, 0), fetchOffsetAndEpoch(firstLeaderId, ListOffsetRequest.EARLIEST_TIMESTAMP, -1))
+    assertEquals((10L, 0), fetchOffsetAndEpoch(firstLeaderId, ListOffsetRequest.LATEST_TIMESTAMP, -1))
+
+    // Kill the first leader so that we can verify the epoch change when fetching the latest offset
+    killBroker(firstLeaderId)
+    val secondLeaderId = TestUtils.awaitLeaderChange(servers, partition, firstLeaderId)
+    val secondLeaderEpoch = TestUtils.findLeaderEpoch(secondLeaderId, partition, servers)
+
+    // No changes to written data
+    assertEquals((0L, 0), fetchOffsetAndEpoch(secondLeaderId, 0L, -1))
+    assertEquals((0L, 0), fetchOffsetAndEpoch(secondLeaderId, ListOffsetRequest.EARLIEST_TIMESTAMP, -1))
+
+    assertEquals((0L, 0), fetchOffsetAndEpoch(secondLeaderId, 0L, -1))
+    assertEquals((0L, 0), fetchOffsetAndEpoch(secondLeaderId, ListOffsetRequest.EARLIEST_TIMESTAMP, -1))
+
+    // The latest offset reflects the updated epoch
+    assertEquals((10L, secondLeaderEpoch), fetchOffsetAndEpoch(secondLeaderId, ListOffsetRequest.LATEST_TIMESTAMP, -1))
+  }
+
+  @Test
+  def testResponseDefaultOffsetAndLeaderEpochForAllVersions(): Unit = {
+    val partitionToLeader = TestUtils.createTopic(zkClient, topic, numPartitions = 1, replicationFactor = 3, servers)
+    val firstLeaderId = partitionToLeader(partition.partition)
+
+    TestUtils.generateAndProduceMessages(servers, topic, 10)
+
+    for (version <- ApiKeys.LIST_OFFSETS.oldestVersion to ApiKeys.LIST_OFFSETS.latestVersion) {
+      if (version == 0) {
+        assertEquals((-1L, -1), fetchOffsetAndEpoch(firstLeaderId, 0L, version.toShort))
+        assertEquals((0L, -1), fetchOffsetAndEpoch(firstLeaderId, ListOffsetRequest.EARLIEST_TIMESTAMP, version.toShort))
+        assertEquals((10L, -1), fetchOffsetAndEpoch(firstLeaderId, ListOffsetRequest.LATEST_TIMESTAMP, version.toShort))
+      } else if (version >= 1 && version <= 3) {
+        assertEquals((0L, -1), fetchOffsetAndEpoch(firstLeaderId, 0L, version.toShort))
+        assertEquals((0L, -1), fetchOffsetAndEpoch(firstLeaderId, ListOffsetRequest.EARLIEST_TIMESTAMP, version.toShort))
+        assertEquals((10L, -1), fetchOffsetAndEpoch(firstLeaderId, ListOffsetRequest.LATEST_TIMESTAMP, version.toShort))
+      } else if (version >= 4) {
+        assertEquals((0L, 0), fetchOffsetAndEpoch(firstLeaderId, 0L, version.toShort))
+        assertEquals((0L, 0), fetchOffsetAndEpoch(firstLeaderId, ListOffsetRequest.EARLIEST_TIMESTAMP, version.toShort))
+        assertEquals((10L, 0), fetchOffsetAndEpoch(firstLeaderId, ListOffsetRequest.LATEST_TIMESTAMP, version.toShort))
+      }
+    }
+  }
+
+  private def assertResponseError(error: Errors, brokerId: Int, request: ListOffsetRequest): Unit = {
+    val response = sendRequest(brokerId, request)
+    assertEquals(request.topics.size, response.topics.size)
+    response.topics.asScala.foreach { topic =>
+      topic.partitions.asScala.foreach { partition =>
+        assertEquals(error.code, partition.errorCode)
+      }
+    }
+  }
+
+  private def sendRequest(leaderId: Int, request: ListOffsetRequest): ListOffsetResponse = {
+    connectAndReceive[ListOffsetResponse](request, destination = brokerSocketServer(leaderId))
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/server/LogDirFailureTest.scala b/core/src/test/scala/unit/kafka/server/LogDirFailureTest.scala
index bf1ee35fce825..a60e57a1555af 100644
--- a/core/src/test/scala/unit/kafka/server/LogDirFailureTest.scala
+++ b/core/src/test/scala/unit/kafka/server/LogDirFailureTest.scala
@@ -20,51 +20,55 @@ 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.utils.{CoreUtils, Exit, TestUtils, ZkUtils}
+import kafka.utils.TestUtils.{Checkpoint, LogDirFailureType, Roll}
+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.errors.{KafkaStorageException, NotLeaderOrFollowerException}
 import org.apache.kafka.common.utils.Utils
-import org.apache.kafka.common.errors.{KafkaStorageException, NotLeaderForPartitionException}
+import org.junit.Assert.{assertEquals, assertFalse, assertTrue}
 import org.junit.{Before, Test}
-import org.junit.Assert.{assertTrue, assertFalse, assertEquals}
+import org.scalatest.Assertions.fail
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 
 /**
-  * Test whether clients can producer and consume when there is log directory failure
+  * Test whether clients can produce and consume when there is log directory failure
   */
 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
 
-  this.logDirCount = 3
-  this.producerConfig.setProperty(ProducerConfig.RETRIES_CONFIG, "0")
-  this.producerConfig.setProperty(ProducerConfig.RETRY_BACKOFF_MS_CONFIG, "100")
   this.serverConfig.setProperty(KafkaConfig.ReplicaHighWatermarkCheckpointIntervalMsProp, "60000")
   this.serverConfig.setProperty(KafkaConfig.NumReplicaFetchersProp, "1")
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
-    TestUtils.createTopic(zkUtils, topic, partitionNum, serverCount, servers = servers)
+    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)
@@ -73,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)
@@ -93,13 +97,19 @@ class LogDirFailureTest extends IntegrationTestHarness {
   }
 
   @Test
-  def testIOExceptionDuringCheckpoint() {
+  def testProduceErrorFromFailureOnCheckpoint(): Unit = {
+    testProduceErrorsFromLogDirFailureOnLeader(Checkpoint)
+  }
+
+  @Test
+  def testIOExceptionDuringCheckpoint(): Unit = {
     testProduceAfterLogDirFailureOnLeader(Checkpoint)
   }
 
   @Test
-  def testReplicaFetcherThreadAfterLogDirFailureOnFollower() {
-    val producer = producers.head
+  def testReplicaFetcherThreadAfterLogDirFailureOnFollower(): Unit = {
+    this.producerConfig.setProperty(ProducerConfig.RETRIES_CONFIG, "0")
+    val producer = createProducer()
     val partition = new TopicPartition(topic, 0)
 
     val partitionInfo = producer.partitionsFor(topic).asScala.find(_.partition() == 0).get
@@ -112,105 +122,91 @@ 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 = consumers.head
-    subscribeAndWaitForAssignment(topic, consumer)
-    val producer = producers.head
+  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()
+
     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.waitUntilTrue(() => {
-      consumer.poll(0).count() == 1
-    }, "Expected the first message", 3000L)
-
-    // Make log directory of the partition on the leader broker inaccessible by replacing it with a file
-    val replica = leaderServer.replicaManager.getReplicaOrException(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.getReplica(partition).isEmpty)
+    TestUtils.causeLogDirFailure(failureType, leaderServer, partition)
 
-    // The second send() should fail due to either KafkaStorageException or NotLeaderForPartitionException
+    // send() should fail due to either KafkaStorageException or NotLeaderOrFollowerException
     try {
       producer.send(record).get(6000, TimeUnit.MILLISECONDS)
-      fail("send() should fail with either KafkaStorageException or NotLeaderForPartitionException")
+      fail("send() should fail with either KafkaStorageException or NotLeaderOrFollowerException")
     } catch {
       case e: ExecutionException =>
         e.getCause match {
           case t: KafkaStorageException =>
-          case t: NotLeaderForPartitionException => // This may happen if ProduceRequest version <= 3
-          case t: Throwable => fail(s"send() should fail with either KafkaStorageException or NotLeaderForPartitionException instead of ${t.toString}")
+          case t: NotLeaderOrFollowerException => // This may happen if ProduceRequest version <= 3
+          case t: Throwable => fail(s"send() should fail with either KafkaStorageException or NotLeaderOrFollowerException instead of ${t.toString}")
         }
-      case e: Throwable => fail(s"send() should fail with either KafkaStorageException or NotLeaderForPartitionException instead of ${e.toString}")
     }
+  }
+
+  def testProduceAfterLogDirFailureOnLeader(failureType: LogDirFailureType): Unit = {
+    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)
+
+    TestUtils.causeLogDirFailure(failureType, leaderServer, partition)
 
-    // Wait for producer to update metadata for the partition
     TestUtils.waitUntilTrue(() => {
       // ProduceResponse may contain KafkaStorageException and trigger metadata update
       producer.send(record)
       producer.partitionsFor(topic).asScala.find(_.partition() == 0).get.leader().id() != leaderServerId
     }, "Expected new leader for the partition", 6000L)
 
+    // Block on send to ensure that new leader accepts a message.
+    producer.send(record).get(6000L, TimeUnit.MILLISECONDS)
+
     // Consumer should receive some messages
-    TestUtils.waitUntilTrue(() => {
-      consumer.poll(0).count() > 0
-    }, "Expected some messages", 3000L)
+    TestUtils.pollUntilAtLeastNumRecords(consumer, 1)
 
     // There should be no remaining LogDirEventNotification znode
-    assertTrue(zkUtils.getChildrenParentMayNotExist(ZkUtils.LogDirEventNotificationPath).isEmpty)
+    assertTrue(zkClient.getAllLogDirEventNotifications.isEmpty)
 
     // 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 subscribeAndWaitForAssignment(topic: String, consumer: KafkaConsumer[Array[Byte], Array[Byte]]): Unit = {
     consumer.subscribe(Collections.singletonList(topic))
-    TestUtils.waitUntilTrue(() => {
-      consumer.poll(0)
-      !consumer.assignment.isEmpty
-    }, "Expected non-empty assignment")
+    TestUtils.pollUntilTrue(consumer, () => !consumer.assignment.isEmpty, "Expected non-empty assignment")
   }
 
 }
-
-object LogDirFailureTest {
-  sealed trait LogDirFailureType
-  case object Roll extends LogDirFailureType
-  case object Checkpoint extends LogDirFailureType
-}
-
diff --git a/core/src/test/scala/unit/kafka/server/LogOffsetTest.scala b/core/src/test/scala/unit/kafka/server/LogOffsetTest.scala
index b0a7d72f02997..e89a69854ca75 100755
--- a/core/src/test/scala/unit/kafka/server/LogOffsetTest.scala
+++ b/core/src/test/scala/unit/kafka/server/LogOffsetTest.scala
@@ -19,169 +19,151 @@ package kafka.server
 
 import java.io.File
 import java.util.concurrent.atomic.AtomicInteger
-import java.util.{Properties, Random}
-
-import kafka.api.{FetchRequestBuilder, OffsetRequest, PartitionOffsetRequestInfo}
-import kafka.common.TopicAndPartition
-import kafka.consumer.SimpleConsumer
-import kafka.log.{Log, LogSegment}
-import kafka.utils.TestUtils._
-import kafka.utils._
-import kafka.zk.ZooKeeperTestHarness
-import org.apache.kafka.common.TopicPartition
+import java.util.{Optional, Properties, Random}
+
+import kafka.log.{ClientRecordDeletion, Log, LogSegment}
+import kafka.utils.{MockTime, TestUtils}
+import org.apache.kafka.common.message.ListOffsetRequestData.ListOffsetTopic
+import org.apache.kafka.common.message.ListOffsetRequestData.ListOffsetPartition
+import org.apache.kafka.common.message.ListOffsetResponseData.ListOffsetPartitionResponse
+import org.apache.kafka.common.message.ListOffsetResponseData.ListOffsetTopicResponse
 import org.apache.kafka.common.protocol.Errors
-import org.apache.kafka.common.utils.Time
+import org.apache.kafka.common.record.MemoryRecords
+import org.apache.kafka.common.requests.{FetchRequest, FetchResponse, ListOffsetRequest, ListOffsetResponse}
+import org.apache.kafka.common.{IsolationLevel, TopicPartition}
 import org.easymock.{EasyMock, IAnswer}
 import org.junit.Assert._
-import org.junit.{After, Before, Test}
-
-class LogOffsetTest extends ZooKeeperTestHarness {
-  val random = new Random()
-  var logDir: File = null
-  var topicLogDir: File = null
-  var server: KafkaServer = null
-  var logSize: Int = 140
-  var simpleConsumer: SimpleConsumer = null
-  var time: Time = new MockTime()
-
-  @Before
-  override def setUp() {
-    super.setUp()
-    val config: Properties = createBrokerConfig(1)
-    config.put(KafkaConfig.LogMessageTimestampDifferenceMaxMsProp, Long.MaxValue.toString)
-    val logDirPath = config.getProperty("log.dir")
-    logDir = new File(logDirPath)
-    time = new MockTime()
-    server = TestUtils.createServer(KafkaConfig.fromProps(config), time)
-    simpleConsumer = new SimpleConsumer("localhost", TestUtils.boundPort(server), 1000000, 64*1024, "")
-  }
+import org.junit.Test
+
+import scala.jdk.CollectionConverters._
+import scala.collection.mutable.Buffer
+
+class LogOffsetTest extends BaseRequestTest {
+
+  private lazy val time = new MockTime
+
+  override def brokerCount = 1
+
+  protected override def brokerTime(brokerId: Int) = time
 
-  @After
-  override def tearDown() {
-    simpleConsumer.close
-    TestUtils.shutdownServers(Seq(server))
-    super.tearDown()
+  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")
+    props.put("log.retention.check.interval.ms", (5 * 1000 * 60).toString)
+    props.put("log.segment.bytes", "140")
   }
 
+  @deprecated("ListOffsetsRequest V0", since = "")
   @Test
-  def testGetOffsetsForUnknownTopic() {
-    val topicAndPartition = TopicAndPartition("foo", 0)
-    val request = OffsetRequest(
-      Map(topicAndPartition -> PartitionOffsetRequestInfo(OffsetRequest.LatestTime, 10)))
-    val offsetResponse = simpleConsumer.getOffsetsBefore(request)
-    assertEquals(Errors.UNKNOWN_TOPIC_OR_PARTITION,
-                 offsetResponse.partitionErrorAndOffsets(topicAndPartition).error)
+  def testGetOffsetsForUnknownTopic(): Unit = {
+    val topicPartition = new TopicPartition("foo", 0)
+    val request = ListOffsetRequest.Builder.forConsumer(false, IsolationLevel.READ_UNCOMMITTED)
+      .setTargetTimes(buildTargetTimes(topicPartition, ListOffsetRequest.LATEST_TIMESTAMP, 10).asJava).build(0)
+    val response = sendListOffsetsRequest(request)
+    assertEquals(Errors.UNKNOWN_TOPIC_OR_PARTITION.code, findPartition(response.topics.asScala, topicPartition).errorCode)
   }
 
+  @deprecated("ListOffsetsRequest V0", since = "")
   @Test
-  def testGetOffsetsAfterDeleteRecords() {
-    val topicPartition = "kafka-" + 0
-    val topic = topicPartition.split("-").head
-    val part = Integer.valueOf(topicPartition.split("-").last).intValue
+  def testGetOffsetsAfterDeleteRecords(): Unit = {
+    val topic = "kafka-"
+    val topicPartition = new TopicPartition(topic, 0)
 
-    // setup brokers in ZooKeeper as owners of partitions for this test
-    adminZkClient.createTopic(topic, 1, 1)
+    createTopic(topic, 1, 1)
 
     val logManager = server.getLogManager
-    waitUntilTrue(() => logManager.getLog(new TopicPartition(topic, part)).isDefined,
+    TestUtils.waitUntilTrue(() => logManager.getLog(topicPartition).isDefined,
                   "Log for partition [topic,0] should be created")
-    val log = logManager.getLog(new TopicPartition(topic, part)).get
+    val log = logManager.getLog(topicPartition).get
 
     for (_ <- 0 until 20)
       log.appendAsLeader(TestUtils.singletonRecords(value = Integer.toString(42).getBytes()), leaderEpoch = 0)
     log.flush()
 
-    log.onHighWatermarkIncremented(log.logEndOffset)
-    log.maybeIncrementLogStartOffset(3)
+    log.updateHighWatermark(log.logEndOffset)
+    log.maybeIncrementLogStartOffset(3, ClientRecordDeletion)
     log.deleteOldSegments()
 
-    val offsets = server.apis.fetchOffsets(logManager, new TopicPartition(topic, part), OffsetRequest.LatestTime, 15)
+    val offsets = log.legacyFetchOffsetsBefore(ListOffsetRequest.LATEST_TIMESTAMP, 15)
     assertEquals(Seq(20L, 18L, 16L, 14L, 12L, 10L, 8L, 6L, 4L, 3L), offsets)
 
-    waitUntilTrue(() => isLeaderLocalOnBroker(topic, part, server), "Leader should be elected")
-    val topicAndPartition = TopicAndPartition(topic, part)
-    val offsetRequest = OffsetRequest(
-      Map(topicAndPartition -> PartitionOffsetRequestInfo(OffsetRequest.LatestTime, 15)),
-      replicaId = 0)
-    val consumerOffsets =
-      simpleConsumer.getOffsetsBefore(offsetRequest).partitionErrorAndOffsets(topicAndPartition).offsets
+    TestUtils.waitUntilTrue(() => TestUtils.isLeaderLocalOnBroker(topic, topicPartition.partition, server),
+      "Leader should be elected")
+    val request = ListOffsetRequest.Builder.forReplica(0, 0)
+      .setTargetTimes(buildTargetTimes(topicPartition, ListOffsetRequest.LATEST_TIMESTAMP, 15).asJava).build()
+    val consumerOffsets = findPartition(sendListOffsetsRequest(request).topics.asScala, topicPartition).oldStyleOffsets.asScala
     assertEquals(Seq(20L, 18L, 16L, 14L, 12L, 10L, 8L, 6L, 4L, 3L), consumerOffsets)
   }
 
   @Test
-  def testGetOffsetsBeforeLatestTime() {
-    val topicPartition = "kafka-" + 0
-    val topic = topicPartition.split("-").head
-    val part = Integer.valueOf(topicPartition.split("-").last).intValue
+  def testGetOffsetsBeforeLatestTime(): Unit = {
+    val topic = "kafka-"
+    val topicPartition = new TopicPartition(topic, 0)
 
-    // setup brokers in ZooKeeper as owners of partitions for this test
-    adminZkClient.createTopic(topic, 1, 1)
+    createTopic(topic, 1, 1)
 
     val logManager = server.getLogManager
-    waitUntilTrue(() => logManager.getLog(new TopicPartition(topic, part)).isDefined,
-      "Log for partition [topic,0] should be created")
-    val log = logManager.getLog(new TopicPartition(topic, part)).get
+    TestUtils.waitUntilTrue(() => logManager.getLog(topicPartition).isDefined,
+      s"Log for partition $topicPartition should be created")
+    val log = logManager.getLog(topicPartition).get
 
     for (_ <- 0 until 20)
       log.appendAsLeader(TestUtils.singletonRecords(value = Integer.toString(42).getBytes()), leaderEpoch = 0)
     log.flush()
 
-    val offsets = server.apis.fetchOffsets(logManager, new TopicPartition(topic, part), OffsetRequest.LatestTime, 15)
+    val offsets = log.legacyFetchOffsetsBefore(ListOffsetRequest.LATEST_TIMESTAMP, 15)
     assertEquals(Seq(20L, 18L, 16L, 14L, 12L, 10L, 8L, 6L, 4L, 2L, 0L), offsets)
 
-    waitUntilTrue(() => isLeaderLocalOnBroker(topic, part, server), "Leader should be elected")
-    val topicAndPartition = TopicAndPartition(topic, part)
-    val offsetRequest = OffsetRequest(
-      Map(topicAndPartition -> PartitionOffsetRequestInfo(OffsetRequest.LatestTime, 15)),
-      replicaId = 0)
-    val consumerOffsets =
-      simpleConsumer.getOffsetsBefore(offsetRequest).partitionErrorAndOffsets(topicAndPartition).offsets
+    TestUtils.waitUntilTrue(() => TestUtils.isLeaderLocalOnBroker(topic, topicPartition.partition, server),
+      "Leader should be elected")
+    val request = ListOffsetRequest.Builder.forReplica(0, 0)
+      .setTargetTimes(buildTargetTimes(topicPartition, ListOffsetRequest.LATEST_TIMESTAMP, 15).asJava).build()
+    val consumerOffsets = findPartition(sendListOffsetsRequest(request).topics.asScala, topicPartition).oldStyleOffsets.asScala
     assertEquals(Seq(20L, 18L, 16L, 14L, 12L, 10L, 8L, 6L, 4L, 2L, 0L), consumerOffsets)
 
     // try to fetch using latest offset
-    val fetchResponse = simpleConsumer.fetch(
-      new FetchRequestBuilder().addFetch(topic, 0, consumerOffsets.head, 300 * 1024).build())
-    assertFalse(fetchResponse.messageSet(topic, 0).iterator.hasNext)
+    val fetchRequest = FetchRequest.Builder.forConsumer(0, 1,
+      Map(topicPartition -> new FetchRequest.PartitionData(consumerOffsets.head, FetchRequest.INVALID_LOG_START_OFFSET,
+        300 * 1024, Optional.empty())).asJava).build()
+    val fetchResponse = sendFetchRequest(fetchRequest)
+    assertFalse(fetchResponse.responseData.get(topicPartition).records.batches.iterator.hasNext)
   }
 
   @Test
-  def testEmptyLogsGetOffsets() {
-    val topicPartition = "kafka-" + random.nextInt(10)
-    val topicPartitionPath = TestUtils.tempDir().getAbsolutePath + "/" + topicPartition
-    topicLogDir = new File(topicPartitionPath)
+  def testEmptyLogsGetOffsets(): Unit = {
+    val random = new Random
+    val topic = "kafka-"
+    val topicPartition = new TopicPartition(topic, random.nextInt(10))
+    val topicPartitionPath = s"${TestUtils.tempDir().getAbsolutePath}/$topic-${topicPartition.partition}"
+    val topicLogDir = new File(topicPartitionPath)
     topicLogDir.mkdir()
 
-    val topic = topicPartition.split("-").head
-
-    // setup brokers in ZooKeeper as owners of partitions for this test
-    createTopic(zkUtils, topic, numPartitions = 1, replicationFactor = 1, servers = Seq(server))
+    createTopic(topic, numPartitions = 1, replicationFactor = 1)
 
     var offsetChanged = false
     for (_ <- 1 to 14) {
-      val topicAndPartition = TopicAndPartition(topic, 0)
-      val offsetRequest =
-        OffsetRequest(Map(topicAndPartition -> PartitionOffsetRequestInfo(OffsetRequest.EarliestTime, 1)))
-      val consumerOffsets =
-        simpleConsumer.getOffsetsBefore(offsetRequest).partitionErrorAndOffsets(topicAndPartition).offsets
-
-      if(consumerOffsets.head == 1) {
+      val topicPartition = new TopicPartition(topic, 0)
+      val request = ListOffsetRequest.Builder.forReplica(0, 0)
+        .setTargetTimes(buildTargetTimes(topicPartition, ListOffsetRequest.EARLIEST_TIMESTAMP, 1).asJava).build()
+      val consumerOffsets = findPartition(sendListOffsetsRequest(request).topics.asScala, topicPartition).oldStyleOffsets.asScala
+      if (consumerOffsets.head == 1)
         offsetChanged = true
-      }
     }
     assertFalse(offsetChanged)
   }
 
+  @deprecated("legacyFetchOffsetsBefore", since = "")
   @Test
-  def testGetOffsetsBeforeNow() {
-    val topicPartition = "kafka-" + random.nextInt(3)
-    val topic = topicPartition.split("-").head
-    val part = Integer.valueOf(topicPartition.split("-").last).intValue
+  def testGetOffsetsBeforeNow(): Unit = {
+    val random = new Random
+    val topic = "kafka-"
+    val topicPartition = new TopicPartition(topic, random.nextInt(3))
 
-    // setup brokers in ZooKeeper as owners of partitions for this test
-    adminZkClient.createTopic(topic, 3, 1)
+    createTopic(topic, 3, 1)
 
     val logManager = server.getLogManager
-    val log = logManager.getOrCreateLog(new TopicPartition(topic, part), logManager.defaultConfig)
+    val log = logManager.getOrCreateLog(topicPartition, () => logManager.initialDefaultConfig)
 
     for (_ <- 0 until 20)
       log.appendAsLeader(TestUtils.singletonRecords(value = Integer.toString(42).getBytes()), leaderEpoch = 0)
@@ -189,51 +171,50 @@ class LogOffsetTest extends ZooKeeperTestHarness {
 
     val now = time.milliseconds + 30000 // pretend it is the future to avoid race conditions with the fs
 
-    val offsets = server.apis.fetchOffsets(logManager, new TopicPartition(topic, part), now, 15)
+    val offsets = log.legacyFetchOffsetsBefore(now, 15)
     assertEquals(Seq(20L, 18L, 16L, 14L, 12L, 10L, 8L, 6L, 4L, 2L, 0L), offsets)
 
-    waitUntilTrue(() => isLeaderLocalOnBroker(topic, part, server), "Leader should be elected")
-    val topicAndPartition = TopicAndPartition(topic, part)
-    val offsetRequest = OffsetRequest(Map(topicAndPartition -> PartitionOffsetRequestInfo(now, 15)), replicaId = 0)
-    val consumerOffsets =
-      simpleConsumer.getOffsetsBefore(offsetRequest).partitionErrorAndOffsets(topicAndPartition).offsets
+    TestUtils.waitUntilTrue(() => TestUtils.isLeaderLocalOnBroker(topic, topicPartition.partition, server),
+      "Leader should be elected")
+    val request = ListOffsetRequest.Builder.forReplica(0, 0)
+      .setTargetTimes(buildTargetTimes(topicPartition, now, 15).asJava).build()
+    val consumerOffsets = findPartition(sendListOffsetsRequest(request).topics.asScala, topicPartition).oldStyleOffsets.asScala
     assertEquals(Seq(20L, 18L, 16L, 14L, 12L, 10L, 8L, 6L, 4L, 2L, 0L), consumerOffsets)
   }
 
+  @deprecated("legacyFetchOffsetsBefore", since = "")
   @Test
-  def testGetOffsetsBeforeEarliestTime() {
-    val topicPartition = "kafka-" + random.nextInt(3)
-    val topic = topicPartition.split("-").head
-    val part = Integer.valueOf(topicPartition.split("-").last).intValue
+  def testGetOffsetsBeforeEarliestTime(): Unit = {
+    val random = new Random
+    val topic = "kafka-"
+    val topicPartition = new TopicPartition(topic, random.nextInt(3))
 
-    // setup brokers in ZooKeeper as owners of partitions for this test
-    adminZkClient.createTopic(topic, 3, 1)
+    createTopic(topic, 3, 1)
 
     val logManager = server.getLogManager
-    val log = logManager.getOrCreateLog(new TopicPartition(topic, part), logManager.defaultConfig)
+    val log = logManager.getOrCreateLog(topicPartition, () => logManager.initialDefaultConfig)
     for (_ <- 0 until 20)
       log.appendAsLeader(TestUtils.singletonRecords(value = Integer.toString(42).getBytes()), leaderEpoch = 0)
     log.flush()
 
-    val offsets = server.apis.fetchOffsets(logManager, new TopicPartition(topic, part), OffsetRequest.EarliestTime, 10)
+    val offsets = log.legacyFetchOffsetsBefore(ListOffsetRequest.EARLIEST_TIMESTAMP, 10)
 
     assertEquals(Seq(0L), offsets)
 
-    waitUntilTrue(() => isLeaderLocalOnBroker(topic, part, server), "Leader should be elected")
-    val topicAndPartition = TopicAndPartition(topic, part)
-    val offsetRequest =
-      OffsetRequest(Map(topicAndPartition -> PartitionOffsetRequestInfo(OffsetRequest.EarliestTime, 10)))
-    val consumerOffsets =
-      simpleConsumer.getOffsetsBefore(offsetRequest).partitionErrorAndOffsets(topicAndPartition).offsets
+    TestUtils.waitUntilTrue(() => TestUtils.isLeaderLocalOnBroker(topic, topicPartition.partition, server),
+      "Leader should be elected")
+    val request = ListOffsetRequest.Builder.forReplica(0, 0)
+      .setTargetTimes(buildTargetTimes(topicPartition, ListOffsetRequest.EARLIEST_TIMESTAMP, 10).asJava).build()
+    val consumerOffsets = findPartition(sendListOffsetsRequest(request).topics.asScala, topicPartition).oldStyleOffsets.asScala
     assertEquals(Seq(0L), consumerOffsets)
   }
 
   /* We test that `fetchOffsetsBefore` works correctly if `LogSegment.size` changes after each invocation (simulating
    * a race condition) */
   @Test
-  def testFetchOffsetsBeforeWithChangingSegmentSize() {
-    val log = EasyMock.niceMock(classOf[Log])
-    val logSegment = EasyMock.niceMock(classOf[LogSegment])
+  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] {
       private val value = new AtomicInteger(0)
       def answer: Int = value.getAndIncrement()
@@ -242,15 +223,15 @@ class LogOffsetTest extends ZooKeeperTestHarness {
     val logSegments = Seq(logSegment)
     EasyMock.expect(log.logSegments).andStubReturn(logSegments)
     EasyMock.replay(log)
-    server.apis.fetchOffsetsBefore(log, System.currentTimeMillis, 100)
+    log.legacyFetchOffsetsBefore(System.currentTimeMillis, 100)
   }
 
   /* We test that `fetchOffsetsBefore` works correctly if `Log.logSegments` content and size are
    * different (simulating a race condition) */
   @Test
-  def testFetchOffsetsBeforeWithChangingSegments() {
-    val log = EasyMock.niceMock(classOf[Log])
-    val logSegment = EasyMock.niceMock(classOf[LogSegment])
+  def testFetchOffsetsBeforeWithChangingSegments(): Unit = {
+    val log: Log = EasyMock.niceMock(classOf[Log])
+    val logSegment: LogSegment = EasyMock.niceMock(classOf[LogSegment])
     EasyMock.expect(log.logSegments).andStubAnswer {
       new IAnswer[Iterable[LogSegment]] {
         def answer = new Iterable[LogSegment] {
@@ -261,22 +242,32 @@ class LogOffsetTest extends ZooKeeperTestHarness {
     }
     EasyMock.replay(logSegment)
     EasyMock.replay(log)
-    server.apis.fetchOffsetsBefore(log, System.currentTimeMillis, 100)
+    log.legacyFetchOffsetsBefore(System.currentTimeMillis, 100)
   }
 
-  private def createBrokerConfig(nodeId: Int): Properties = {
-    val props = new Properties
-    props.put("broker.id", nodeId.toString)
-    props.put("port", TestUtils.RandomPort.toString())
-    props.put("log.dir", TestUtils.tempDir().getAbsolutePath)
-    props.put("log.flush.interval.messages", "1")
-    props.put("enable.zookeeper", "false")
-    props.put("num.partitions", "20")
-    props.put("log.retention.hours", "10")
-    props.put("log.retention.check.interval.ms", (5*1000*60).toString)
-    props.put("log.segment.bytes", logSize.toString)
-    props.put("zookeeper.connect", zkConnect.toString)
-    props
+  private def server: KafkaServer = servers.head
+
+  private def sendListOffsetsRequest(request: ListOffsetRequest): ListOffsetResponse = {
+    connectAndReceive[ListOffsetResponse](request)
+  }
+
+  private def sendFetchRequest(request: FetchRequest): FetchResponse[MemoryRecords] = {
+    connectAndReceive[FetchResponse[MemoryRecords]](request)
+  }
+
+  private def buildTargetTimes(tp: TopicPartition, timestamp: Long, maxNumOffsets: Int): List[ListOffsetTopic] = {
+    List(new ListOffsetTopic()
+      .setName(tp.topic)
+      .setPartitions(List(new ListOffsetPartition()
+        .setPartitionIndex(tp.partition)
+        .setTimestamp(timestamp)
+        .setMaxNumOffsets(maxNumOffsets)).asJava)
+    )
+  }
+
+  private def findPartition(topics: Buffer[ListOffsetTopicResponse], tp: TopicPartition): ListOffsetPartitionResponse = {
+    topics.find(_.name == tp.topic).get
+      .partitions.asScala.find(_.partitionIndex == tp.partition).get
   }
 
 }
diff --git a/core/src/test/scala/unit/kafka/server/LogRecoveryTest.scala b/core/src/test/scala/unit/kafka/server/LogRecoveryTest.scala
index fa373f5d8bcdc..484e1637c70d8 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
@@ -65,16 +67,15 @@ class LogRecoveryTest extends ZooKeeperTestHarness {
   def updateProducer() = {
     if (producer != null)
       producer.close()
-    producer = TestUtils.createNewProducer(
+    producer = TestUtils.createProducer(
       TestUtils.getBrokerListStrFromServers(servers),
-      retries = 5,
       keySerializer = new IntegerSerializer,
       valueSerializer = new StringSerializer
     )
   }
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
 
     configs = TestUtils.createBrokerConfigs(2, zkConnect, enableControlledShutdown = false).map(KafkaConfig.fromProps(_, overridingProps))
@@ -85,14 +86,14 @@ class LogRecoveryTest extends ZooKeeperTestHarness {
     servers = List(server1, server2)
 
     // create topic with 1 partition, 2 replicas, one on each broker
-    createTopic(zkUtils, topic, partitionReplicaAssignment = Map(0 -> Seq(0,1)), servers = servers)
+    createTopic(zkClient, topic, partitionReplicaAssignment = Map(0 -> Seq(0,1)), servers = servers)
 
     // create the producer
     updateProducer()
   }
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     producer.close()
     TestUtils.shutdownServers(servers)
     super.tearDown()
@@ -105,21 +106,21 @@ class LogRecoveryTest extends ZooKeeperTestHarness {
 
     // give some time for the follower 1 to record leader HW
     TestUtils.waitUntilTrue(() =>
-      server2.replicaManager.getReplica(topicPartition).get.highWatermark.messageOffset == numMessages,
+      server2.replicaManager.localLogOrException(topicPartition).highWatermark == numMessages,
       "Failed to update high watermark for follower after timeout")
 
     servers.foreach(_.replicaManager.checkpointHighWatermarks())
-    val leaderHW = hwFile1.read.getOrElse(topicPartition, 0L)
+    val leaderHW = hwFile1.read().getOrElse(topicPartition, 0L)
     assertEquals(numMessages, leaderHW)
-    val followerHW = hwFile2.read.getOrElse(topicPartition, 0L)
+    val followerHW = hwFile2.read().getOrElse(topicPartition, 0L)
     assertEquals(numMessages, followerHW)
   }
 
   @Test
   def testHWCheckpointWithFailuresSingleLogSegment(): Unit = {
-    var leader = waitUntilLeaderIsElectedOrChanged(zkUtils, topic, partitionId)
+    var leader = waitUntilLeaderIsElectedOrChanged(zkClient, topic, partitionId)
 
-    assertEquals(0L, hwFile1.read.getOrElse(topicPartition, 0L))
+    assertEquals(0L, hwFile1.read().getOrElse(topicPartition, 0L))
 
     sendMessages(1)
     Thread.sleep(1000)
@@ -127,10 +128,10 @@ class LogRecoveryTest extends ZooKeeperTestHarness {
 
     // kill the server hosting the preferred replica
     server1.shutdown()
-    assertEquals(hw, hwFile1.read.getOrElse(topicPartition, 0L))
+    assertEquals(hw, hwFile1.read().getOrElse(topicPartition, 0L))
 
     // check if leader moves to the other server
-    leader = waitUntilLeaderIsElectedOrChanged(zkUtils, topic, partitionId, oldLeaderOpt = Some(leader))
+    leader = waitUntilLeaderIsElectedOrChanged(zkClient, topic, partitionId, oldLeaderOpt = Some(leader))
     assertEquals("Leader must move to broker 1", 1, leader)
 
     // bring the preferred replica back
@@ -138,18 +139,27 @@ class LogRecoveryTest extends ZooKeeperTestHarness {
     // Update producer with new server settings
     updateProducer()
 
-    leader = waitUntilLeaderIsElectedOrChanged(zkUtils, topic, partitionId)
+    leader = waitUntilLeaderIsElectedOrChanged(zkClient, topic, partitionId)
     assertTrue("Leader must remain on broker 1, in case of ZooKeeper session expiration it can move to broker 0",
       leader == 0 || leader == 1)
 
-    assertEquals(hw, hwFile1.read.getOrElse(topicPartition, 0L))
+    assertEquals(hw, hwFile1.read().getOrElse(topicPartition, 0L))
+    /** We plan to shutdown server2 and transfer the leadership to server1.
+      * With unclean leader election turned off, a prerequisite for the successful leadership transition
+      * 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.nonOfflinePartition(topicPartition).get.inSyncReplicaIds.size == 2,
+      "Server 1 is not able to join the ISR after restart")
+
+
     // since server 2 was never shut down, the hw value of 30 is probably not checkpointed to disk yet
     server2.shutdown()
-    assertEquals(hw, hwFile2.read.getOrElse(topicPartition, 0L))
+    assertEquals(hw, hwFile2.read().getOrElse(topicPartition, 0L))
 
     server2.startup()
     updateProducer()
-    leader = waitUntilLeaderIsElectedOrChanged(zkUtils, topic, partitionId, oldLeaderOpt = Some(leader))
+    leader = waitUntilLeaderIsElectedOrChanged(zkClient, topic, partitionId, oldLeaderOpt = Some(leader))
     assertTrue("Leader must remain on broker 0, in case of ZooKeeper session expiration it can move to broker 1",
       leader == 0 || leader == 1)
 
@@ -158,12 +168,12 @@ class LogRecoveryTest extends ZooKeeperTestHarness {
 
     // give some time for follower 1 to record leader HW of 60
     TestUtils.waitUntilTrue(() =>
-      server2.replicaManager.getReplica(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())
-    assertEquals(hw, hwFile1.read.getOrElse(topicPartition, 0L))
-    assertEquals(hw, hwFile2.read.getOrElse(topicPartition, 0L))
+    assertEquals(hw, hwFile1.read().getOrElse(topicPartition, 0L))
+    assertEquals(hw, hwFile2.read().getOrElse(topicPartition, 0L))
   }
 
   @Test
@@ -172,65 +182,65 @@ class LogRecoveryTest extends ZooKeeperTestHarness {
     val hw = 20L
     // give some time for follower 1 to record leader HW of 600
     TestUtils.waitUntilTrue(() =>
-      server2.replicaManager.getReplica(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())
-    val leaderHW = hwFile1.read.getOrElse(topicPartition, 0L)
+    val leaderHW = hwFile1.read().getOrElse(topicPartition, 0L)
     assertEquals(hw, leaderHW)
-    val followerHW = hwFile2.read.getOrElse(topicPartition, 0L)
+    val followerHW = hwFile2.read().getOrElse(topicPartition, 0L)
     assertEquals(hw, followerHW)
   }
 
   @Test
   def testHWCheckpointWithFailuresMultipleLogSegments(): Unit = {
-    var leader = waitUntilLeaderIsElectedOrChanged(zkUtils, topic, partitionId)
+    var leader = waitUntilLeaderIsElectedOrChanged(zkClient, topic, partitionId)
 
     sendMessages(2)
     var hw = 2L
 
     // allow some time for the follower to get the leader HW
     TestUtils.waitUntilTrue(() =>
-      server2.replicaManager.getReplica(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()
     server2.shutdown()
-    assertEquals(hw, hwFile1.read.getOrElse(topicPartition, 0L))
-    assertEquals(hw, hwFile2.read.getOrElse(topicPartition, 0L))
+    assertEquals(hw, hwFile1.read().getOrElse(topicPartition, 0L))
+    assertEquals(hw, hwFile2.read().getOrElse(topicPartition, 0L))
 
     server2.startup()
     updateProducer()
     // check if leader moves to the other server
-    leader = waitUntilLeaderIsElectedOrChanged(zkUtils, topic, partitionId, oldLeaderOpt = Some(leader))
+    leader = waitUntilLeaderIsElectedOrChanged(zkClient, topic, partitionId, oldLeaderOpt = Some(leader))
     assertEquals("Leader must move to broker 1", 1, leader)
 
-    assertEquals(hw, hwFile1.read.getOrElse(topicPartition, 0L))
+    assertEquals(hw, hwFile1.read().getOrElse(topicPartition, 0L))
 
     // bring the preferred replica back
     server1.startup()
     updateProducer()
 
-    assertEquals(hw, hwFile1.read.getOrElse(topicPartition, 0L))
-    assertEquals(hw, hwFile2.read.getOrElse(topicPartition, 0L))
+    assertEquals(hw, hwFile1.read().getOrElse(topicPartition, 0L))
+    assertEquals(hw, hwFile2.read().getOrElse(topicPartition, 0L))
 
     sendMessages(2)
     hw += 2
 
     // allow some time for the follower to create replica
-    TestUtils.waitUntilTrue(() => server1.replicaManager.getReplica(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.getReplica(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())
-    assertEquals(hw, hwFile1.read.getOrElse(topicPartition, 0L))
-    assertEquals(hw, hwFile2.read.getOrElse(topicPartition, 0L))
+    assertEquals(hw, hwFile1.read().getOrElse(topicPartition, 0L))
+    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/MetadataCacheTest.scala b/core/src/test/scala/unit/kafka/server/MetadataCacheTest.scala
index 383c1e238c7a6..1a62065c2110f 100644
--- a/core/src/test/scala/unit/kafka/server/MetadataCacheTest.scala
+++ b/core/src/test/scala/unit/kafka/server/MetadataCacheTest.scala
@@ -19,22 +19,23 @@ package kafka.server
 import java.util
 import util.Arrays.asList
 
-import kafka.common.BrokerEndPointNotAvailableException
-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.record.RecordBatch
 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._
+import scala.jdk.CollectionConverters._
 
 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))
@@ -42,38 +43,72 @@ class MetadataCacheTest {
   }
 
   @Test
-  def getTopicMetadata() {
+  def getTopicMetadata(): Unit = {
     val topic0 = "topic-0"
     val topic1 = "topic-1"
 
-
     val cache = new MetadataCache(1)
 
     val zkVersion = 3
     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,
+    val updateMetadataRequest = new UpdateMetadataRequest.Builder(version, controllerId, controllerEpoch, brokerEpoch,
       partitionStates.asJava, brokers.asJava).build()
-    cache.updateCache(15, updateMetadataRequest)
+    cache.updateMetadata(15, updateMetadataRequest)
 
     for (securityProtocol <- Seq(SecurityProtocol.PLAINTEXT, SecurityProtocol.SSL)) {
       val listenerName = ListenerName.forSecurityProtocol(securityProtocol)
@@ -83,24 +118,22 @@ class MetadataCacheTest {
         assertEquals(1, topicMetadatas.size)
 
         val topicMetadata = topicMetadatas.head
-        assertEquals(Errors.NONE, topicMetadata.error)
-        assertEquals(topic, topicMetadata.topic)
+        assertEquals(Errors.NONE.code, topicMetadata.errorCode)
+        assertEquals(topic, topicMetadata.name)
 
-        val topicPartitionStates = partitionStates.filter { case (tp, _) => tp.topic ==  topic }
-        val partitionMetadatas = topicMetadata.partitionMetadata.asScala.sortBy(_.partition)
+        val topicPartitionStates = partitionStates.filter { ps => ps.topicName == topic }
+        val partitionMetadatas = topicMetadata.partitions.asScala.sortBy(_.partitionIndex)
         assertEquals(s"Unexpected partition count for topic $topic", topicPartitionStates.size, partitionMetadatas.size)
 
         partitionMetadatas.zipWithIndex.foreach { case (partitionMetadata, partitionId) =>
-          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(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)
+          assertEquals(Errors.NONE.code, partitionMetadata.errorCode)
+          assertEquals(partitionId, partitionMetadata.partitionIndex)
+          val partitionState = topicPartitionStates.find(_.partitionIndex == partitionId).getOrElse(
+            Assertions.fail(s"Unable to find partition state for partition $partitionId"))
+          assertEquals(partitionState.leader, partitionMetadata.leaderId)
+          assertEquals(partitionState.leaderEpoch, partitionMetadata.leaderEpoch)
+          assertEquals(partitionState.isr, partitionMetadata.isrNodes)
+          assertEquals(partitionState.replicas, partitionMetadata.replicaNodes)
         }
       }
 
@@ -111,47 +144,110 @@ class MetadataCacheTest {
   }
 
   @Test
-  def getTopicMetadataPartitionLeaderNotAvailable() {
+  def getTopicMetadataPartitionLeaderNotAvailable(): Unit = {
+    val securityProtocol = SecurityProtocol.PLAINTEXT
+    val listenerName = ListenerName.forSecurityProtocol(securityProtocol)
+    val brokers = Seq(new UpdateMetadataBroker()
+      .setId(0)
+      .setEndpoints(Seq(new UpdateMetadataEndpoint()
+        .setHost("foo")
+        .setPort(9092)
+        .setSecurityProtocol(securityProtocol.id)
+        .setListener(listenerName.value)).asJava))
+    val metadataCacheBrokerId = 0
+    // leader is not available. expect LEADER_NOT_AVAILABLE for any metadata version.
+    verifyTopicMetadataPartitionLeaderOrEndpointNotAvailable(metadataCacheBrokerId, brokers, listenerName,
+      leader = 1, Errors.LEADER_NOT_AVAILABLE, errorUnavailableListeners = false)
+    verifyTopicMetadataPartitionLeaderOrEndpointNotAvailable(metadataCacheBrokerId, brokers, listenerName,
+      leader = 1, Errors.LEADER_NOT_AVAILABLE, errorUnavailableListeners = true)
+  }
+
+  @Test
+  def getTopicMetadataPartitionListenerNotAvailableOnLeader(): Unit = {
+    // when listener name is not present in the metadata cache for a broker, getTopicMetadata should
+    // return LEADER_NOT_AVAILABLE or LISTENER_NOT_FOUND errors for old and new versions respectively.
+    val plaintextListenerName = ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT)
+    val sslListenerName = ListenerName.forSecurityProtocol(SecurityProtocol.SSL)
+    val broker0Endpoints = Seq(
+      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))
+    val metadataCacheBrokerId = 0
+    // leader available in cache but listener name not present. expect LISTENER_NOT_FOUND error for new metadata version
+    verifyTopicMetadataPartitionLeaderOrEndpointNotAvailable(metadataCacheBrokerId, brokers, sslListenerName,
+      leader = 1, Errors.LISTENER_NOT_FOUND, errorUnavailableListeners = true)
+    // leader available in cache but listener name not present. expect LEADER_NOT_AVAILABLE error for old metadata version
+    verifyTopicMetadataPartitionLeaderOrEndpointNotAvailable(metadataCacheBrokerId, brokers, sslListenerName,
+      leader = 1, Errors.LEADER_NOT_AVAILABLE, errorUnavailableListeners = false)
+  }
+
+  private def verifyTopicMetadataPartitionLeaderOrEndpointNotAvailable(metadataCacheBrokerId: Int,
+                                                                       brokers: Seq[UpdateMetadataBroker],
+                                                                       listenerName: ListenerName,
+                                                                       leader: Int,
+                                                                       expectedError: Errors,
+                                                                       errorUnavailableListeners: Boolean): Unit = {
     val topic = "topic"
 
-    val cache = new MetadataCache(1)
+    val cache = new MetadataCache(metadataCacheBrokerId)
 
     val zkVersion = 3
     val controllerId = 2
     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 leader = 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,
+    val updateMetadataRequest = new UpdateMetadataRequest.Builder(version, controllerId, controllerEpoch, brokerEpoch,
       partitionStates.asJava, brokers.asJava).build()
-    cache.updateCache(15, updateMetadataRequest)
+    cache.updateMetadata(15, updateMetadataRequest)
 
-    val topicMetadatas = cache.getTopicMetadata(Set(topic), listenerName)
+    val topicMetadatas = cache.getTopicMetadata(Set(topic), listenerName, errorUnavailableListeners = errorUnavailableListeners)
     assertEquals(1, topicMetadatas.size)
 
     val topicMetadata = topicMetadatas.head
-    assertEquals(Errors.NONE, topicMetadata.error)
+    assertEquals(Errors.NONE.code, topicMetadata.errorCode)
 
-    val partitionMetadatas = topicMetadata.partitionMetadata
+    val partitionMetadatas = topicMetadata.partitions
     assertEquals(1, partitionMetadatas.size)
 
     val partitionMetadata = partitionMetadatas.get(0)
-    assertEquals(0, partitionMetadata.partition)
-    assertEquals(Errors.LEADER_NOT_AVAILABLE, partitionMetadata.error)
-    assertTrue(partitionMetadata.isr.isEmpty)
-    assertEquals(1, partitionMetadata.replicas.size)
-    assertEquals(0, partitionMetadata.replicas.get(0).id)
+    assertEquals(0, partitionMetadata.partitionIndex)
+    assertEquals(expectedError.code, partitionMetadata.errorCode)
+    assertFalse(partitionMetadata.isrNodes.isEmpty)
+    assertEquals(List(0), partitionMetadata.replicaNodes.asScala)
   }
 
   @Test
-  def getTopicMetadataReplicaNotAvailable() {
+  def getTopicMetadataReplicaNotAvailable(): Unit = {
     val topic = "topic"
 
     val cache = new MetadataCache(1)
@@ -161,7 +257,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
@@ -169,49 +271,57 @@ 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,
+    val updateMetadataRequest = new UpdateMetadataRequest.Builder(version, controllerId, controllerEpoch, brokerEpoch,
       partitionStates.asJava, brokers.asJava).build()
-    cache.updateCache(15, updateMetadataRequest)
+    cache.updateMetadata(15, updateMetadataRequest)
 
     // Validate errorUnavailableEndpoints = false
     val topicMetadatas = cache.getTopicMetadata(Set(topic), listenerName, errorUnavailableEndpoints = false)
     assertEquals(1, topicMetadatas.size)
 
     val topicMetadata = topicMetadatas.head
-    assertEquals(Errors.NONE, topicMetadata.error)
+    assertEquals(Errors.NONE.code(), topicMetadata.errorCode)
 
-    val partitionMetadatas = topicMetadata.partitionMetadata
+    val partitionMetadatas = topicMetadata.partitions
     assertEquals(1, partitionMetadatas.size)
 
     val partitionMetadata = partitionMetadatas.get(0)
-    assertEquals(0, partitionMetadata.partition)
-    assertEquals(Errors.NONE, partitionMetadata.error)
-    assertEquals(Set(0, 1), partitionMetadata.replicas.asScala.map(_.id).toSet)
-    assertEquals(Set(0), partitionMetadata.isr.asScala.map(_.id).toSet)
+    assertEquals(0, partitionMetadata.partitionIndex)
+    assertEquals(Errors.NONE.code, partitionMetadata.errorCode)
+    assertEquals(Set(0, 1), partitionMetadata.replicaNodes.asScala.toSet)
+    assertEquals(Set(0), partitionMetadata.isrNodes.asScala.toSet)
 
     // Validate errorUnavailableEndpoints = true
     val topicMetadatasWithError = cache.getTopicMetadata(Set(topic), listenerName, errorUnavailableEndpoints = true)
     assertEquals(1, topicMetadatasWithError.size)
 
     val topicMetadataWithError = topicMetadatasWithError.head
-    assertEquals(Errors.NONE, topicMetadataWithError.error)
+    assertEquals(Errors.NONE.code, topicMetadataWithError.errorCode)
 
-    val partitionMetadatasWithError = topicMetadataWithError.partitionMetadata
+    val partitionMetadatasWithError = topicMetadataWithError.partitions()
     assertEquals(1, partitionMetadatasWithError.size)
 
     val partitionMetadataWithError = partitionMetadatasWithError.get(0)
-    assertEquals(0, partitionMetadataWithError.partition)
-    assertEquals(Errors.REPLICA_NOT_AVAILABLE, partitionMetadataWithError.error)
-    assertEquals(Set(0), partitionMetadataWithError.replicas.asScala.map(_.id).toSet)
-    assertEquals(Set(0), partitionMetadataWithError.isr.asScala.map(_.id).toSet)
+    assertEquals(0, partitionMetadataWithError.partitionIndex)
+    assertEquals(Errors.REPLICA_NOT_AVAILABLE.code, partitionMetadataWithError.errorCode)
+    assertEquals(Set(0), partitionMetadataWithError.replicaNodes.asScala.toSet)
+    assertEquals(Set(0), partitionMetadataWithError.isrNodes.asScala.toSet)
   }
 
   @Test
-  def getTopicMetadataIsrNotAvailable() {
+  def getTopicMetadataIsrNotAvailable(): Unit = {
     val topic = "topic"
 
     val cache = new MetadataCache(1)
@@ -221,7 +331,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
@@ -229,106 +346,135 @@ 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,
+    val updateMetadataRequest = new UpdateMetadataRequest.Builder(version, controllerId, controllerEpoch, brokerEpoch,
       partitionStates.asJava, brokers.asJava).build()
-    cache.updateCache(15, updateMetadataRequest)
+    cache.updateMetadata(15, updateMetadataRequest)
 
     // Validate errorUnavailableEndpoints = false
     val topicMetadatas = cache.getTopicMetadata(Set(topic), listenerName, errorUnavailableEndpoints = false)
     assertEquals(1, topicMetadatas.size)
 
     val topicMetadata = topicMetadatas.head
-    assertEquals(Errors.NONE, topicMetadata.error)
+    assertEquals(Errors.NONE.code(), topicMetadata.errorCode)
 
-    val partitionMetadatas = topicMetadata.partitionMetadata
+    val partitionMetadatas = topicMetadata.partitions
     assertEquals(1, partitionMetadatas.size)
 
     val partitionMetadata = partitionMetadatas.get(0)
-    assertEquals(0, partitionMetadata.partition)
-    assertEquals(Errors.NONE, partitionMetadata.error)
-    assertEquals(Set(0), partitionMetadata.replicas.asScala.map(_.id).toSet)
-    assertEquals(Set(0, 1), partitionMetadata.isr.asScala.map(_.id).toSet)
+    assertEquals(0, partitionMetadata.partitionIndex)
+    assertEquals(Errors.NONE.code, partitionMetadata.errorCode)
+    assertEquals(Set(0), partitionMetadata.replicaNodes.asScala.toSet)
+    assertEquals(Set(0, 1), partitionMetadata.isrNodes.asScala.toSet)
 
     // Validate errorUnavailableEndpoints = true
     val topicMetadatasWithError = cache.getTopicMetadata(Set(topic), listenerName, errorUnavailableEndpoints = true)
     assertEquals(1, topicMetadatasWithError.size)
 
     val topicMetadataWithError = topicMetadatasWithError.head
-    assertEquals(Errors.NONE, topicMetadataWithError.error)
+    assertEquals(Errors.NONE.code, topicMetadataWithError.errorCode)
 
-    val partitionMetadatasWithError = topicMetadataWithError.partitionMetadata
+    val partitionMetadatasWithError = topicMetadataWithError.partitions
     assertEquals(1, partitionMetadatasWithError.size)
 
     val partitionMetadataWithError = partitionMetadatasWithError.get(0)
-    assertEquals(0, partitionMetadataWithError.partition)
-    assertEquals(Errors.REPLICA_NOT_AVAILABLE, partitionMetadataWithError.error)
-    assertEquals(Set(0), partitionMetadataWithError.replicas.asScala.map(_.id).toSet)
-    assertEquals(Set(0), partitionMetadataWithError.isr.asScala.map(_.id).toSet)
+    assertEquals(0, partitionMetadataWithError.partitionIndex)
+    assertEquals(Errors.REPLICA_NOT_AVAILABLE.code, partitionMetadataWithError.errorCode)
+    assertEquals(Set(0), partitionMetadataWithError.replicaNodes.asScala.toSet)
+    assertEquals(Set(0), partitionMetadataWithError.isrNodes.asScala.toSet)
   }
 
   @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, partitionStates.asJava,
+    val updateMetadataRequest = new UpdateMetadataRequest.Builder(version, 2, controllerEpoch, brokerEpoch, partitionStates.asJava,
       brokers.asJava).build()
-    cache.updateCache(15, updateMetadataRequest)
-
-    try {
-      val result = cache.getTopicMetadata(Set(topic), ListenerName.forSecurityProtocol(SecurityProtocol.SSL))
-      fail(s"Exception should be thrown by `getTopicMetadata` with non-supported SecurityProtocol, $result was returned instead")
-    }
-    catch {
-      case _: BrokerEndPointNotAvailableException => //expected
-    }
+    cache.updateMetadata(15, updateMetadataRequest)
 
+    val topicMetadata = cache.getTopicMetadata(Set(topic), ListenerName.forSecurityProtocol(SecurityProtocol.SSL))
+    assertEquals(1, topicMetadata.size)
+    assertEquals(1, topicMetadata.head.partitions.size)
+    assertEquals(RecordBatch.NO_PARTITION_LEADER_EPOCH, topicMetadata.head.partitions.get(0).leaderId)
   }
 
   @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, partitionStates.asJava,
+      val updateMetadataRequest = new UpdateMetadataRequest.Builder(version, 2, controllerEpoch, brokerEpoch, partitionStates.asJava,
         brokers.asJava).build()
-      cache.updateCache(15, updateMetadataRequest)
+      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 877d1564980d2..d6b1e7319cdea 100644
--- a/core/src/test/scala/unit/kafka/server/MetadataRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/MetadataRequestTest.scala
@@ -17,40 +17,51 @@
 
 package kafka.server
 
-import java.util.Properties
+import java.util.{Optional, Properties}
 
 import kafka.network.SocketServer
 import kafka.utils.TestUtils
+import org.apache.kafka.common.errors.UnsupportedVersionException
 import org.apache.kafka.common.internals.Topic
-import org.apache.kafka.common.protocol.{ApiKeys, Errors}
+import org.apache.kafka.common.message.MetadataRequestData
+import org.apache.kafka.common.protocol.Errors
 import org.apache.kafka.common.requests.{MetadataRequest, MetadataResponse}
-import org.junit.Assert._
-import org.junit.Test
 import org.apache.kafka.test.TestUtils.isValidClusterId
+import org.junit.Assert._
+import org.junit.{Before, Test}
+import org.scalatest.Assertions.intercept
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
+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)}")
   }
 
+  @Before
+  override def setUp(): Unit = {
+    doSetup(createOffsetsTopic = false)
+  }
+
   @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))
@@ -67,26 +78,26 @@ class MetadataRequestTest extends BaseRequestTest {
     assertNotEquals("Controller id should switch to a new broker", controllerId, controllerId2)
     TestUtils.waitUntilTrue(() => {
       val metadataResponse2 = sendMetadataRequest(MetadataRequest.Builder.allTopics.build(1.toShort))
-      metadataResponse2.controller != null && controllerServer2.apis.brokerId == metadataResponse2.controller.id
+      metadataResponse2.controller != null && controllerServer2.dataPlaneRequestProcessor.brokerId == metadataResponse2.controller.id
     }, "Controller id should match the active controller after failover", 5000)
   }
 
   @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 =>
+    metadataResponse.brokers.forEach { broker =>
       assertEquals("Rack information should match config", s"rack/${broker.id}", broker.rack)
     }
   }
 
   @Test
-  def testIsInternal() {
+  def testIsInternal(): Unit = {
     val internalTopic = Topic.GROUP_METADATA_TOPIC_NAME
     val notInternalTopic = "notInternal"
     // create the topics
-    TestUtils.createTopic(zkUtils, internalTopic, 3, 2, servers)
-    TestUtils.createTopic(zkUtils, notInternalTopic, 3, 2, servers)
+    createTopic(internalTopic, 3, 2)
+    createTopic(notInternalTopic, 3, 2)
 
     val metadataResponse = sendMetadataRequest(MetadataRequest.Builder.allTopics.build(1.toShort))
     assertTrue("Response should have no errors", metadataResponse.errors.isEmpty)
@@ -102,24 +113,23 @@ class MetadataRequestTest extends BaseRequestTest {
   }
 
   @Test
-  def testNoTopicsRequest() {
+  def testNoTopicsRequest(): Unit = {
     // create some topics
-    TestUtils.createTopic(zkUtils, "t1", 3, 2, servers)
-    TestUtils.createTopic(zkUtils, "t2", 3, 2, servers)
+    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)
   }
 
   @Test
   def testAutoTopicCreation(): Unit = {
-    def checkAutoCreatedTopic(existingTopic: String, autoCreatedTopic: String, response: MetadataResponse): Unit = {
-      assertNull(response.errors.get(existingTopic))
+    def checkAutoCreatedTopic(autoCreatedTopic: String, response: MetadataResponse): Unit = {
       assertEquals(Errors.LEADER_NOT_AVAILABLE, response.errors.get(autoCreatedTopic))
-      assertEquals(Some(servers.head.config.numPartitions), zkUtils.getTopicPartitionCount(autoCreatedTopic))
+      assertEquals(Some(servers.head.config.numPartitions), zkClient.getTopicPartitionCount(autoCreatedTopic))
       for (i <- 0 until servers.head.config.numPartitions)
         TestUtils.waitUntilMetadataIsPropagated(servers, autoCreatedTopic, i)
     }
@@ -128,30 +138,82 @@ class MetadataRequestTest extends BaseRequestTest {
     val topic2 = "t2"
     val topic3 = "t3"
     val topic4 = "t4"
-    TestUtils.createTopic(zkUtils, topic1, 1, 1, servers)
+    val topic5 = "t5"
+    createTopic(topic1, numPartitions = 1, replicationFactor = 1)
 
-    val response1 = sendMetadataRequest(new MetadataRequest(Seq(topic1, topic2).asJava, true, ApiKeys.METADATA.latestVersion))
-    checkAutoCreatedTopic(topic1, topic2, response1)
+    val response1 = sendMetadataRequest(new MetadataRequest.Builder(Seq(topic1, topic2).asJava, true).build())
+    assertNull(response1.errors.get(topic1))
+    checkAutoCreatedTopic(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))
-    checkAutoCreatedTopic(topic2, topic3, response2)
+    // The default behavior in old versions of the metadata API is to allow topic creation, so
+    // protocol downgrades should happen gracefully when auto-creation is explicitly requested.
+    val response2 = sendMetadataRequest(new MetadataRequest.Builder(Seq(topic3).asJava, true).build(1))
+    checkAutoCreatedTopic(topic3, response2)
+
+    // V3 doesn't support a configurable allowAutoTopicCreation, so disabling auto-creation is not supported
+    intercept[UnsupportedVersionException] {
+      sendMetadataRequest(new MetadataRequest(requestData(List(topic4), false), 3.toShort))
+    }
 
     // V4 and higher support a configurable allowAutoTopicCreation
-    val response3 = sendMetadataRequest(new MetadataRequest(Seq(topic3, topic4).asJava, false, 4))
-    assertNull(response3.errors.get(topic3))
+    val response3 = sendMetadataRequest(new MetadataRequest.Builder(Seq(topic4, topic5).asJava, false, 4.toShort).build)
     assertEquals(Errors.UNKNOWN_TOPIC_OR_PARTITION, response3.errors.get(topic4))
-    assertEquals(None, zkUtils.getTopicPartitionCount(topic4))
+    assertEquals(Errors.UNKNOWN_TOPIC_OR_PARTITION, response3.errors.get(topic5))
+    assertEquals(None, zkClient.getTopicPartitionCount(topic5))
+  }
+
+  @Test
+  def testAutoCreateTopicWithInvalidReplicationFactor(): Unit = {
+    // Shutdown all but one broker so that the number of brokers is less than the default replication factor
+    servers.tail.foreach(_.shutdown())
+    servers.tail.foreach(_.awaitShutdown())
+
+    val topic1 = "testAutoCreateTopic"
+    val response1 = sendMetadataRequest(new MetadataRequest.Builder(Seq(topic1).asJava, true).build)
+    assertEquals(1, response1.topicMetadata.size)
+    val topicMetadata = response1.topicMetadata.asScala.head
+    assertEquals(Errors.INVALID_REPLICATION_FACTOR, topicMetadata.error)
+    assertEquals(topic1, topicMetadata.topic)
+    assertEquals(0, topicMetadata.partitionMetadata.size)
   }
 
   @Test
-  def testAllTopicsRequest() {
+  def testAutoCreateOfCollidingTopics(): Unit = {
+    val topic1 = "testAutoCreate_Topic"
+    val topic2 = "testAutoCreate.Topic"
+    val response1 = sendMetadataRequest(new MetadataRequest.Builder(Seq(topic1, topic2).asJava, true).build)
+    assertEquals(2, response1.topicMetadata.size)
+    var topicMetadata1 = response1.topicMetadata.asScala.head
+    val topicMetadata2 = response1.topicMetadata.asScala.toSeq(1)
+    assertEquals(Errors.LEADER_NOT_AVAILABLE, topicMetadata1.error)
+    assertEquals(topic1, topicMetadata1.topic)
+    assertEquals(Errors.INVALID_TOPIC_EXCEPTION, topicMetadata2.error)
+    assertEquals(topic2, topicMetadata2.topic)
+
+    TestUtils.waitUntilLeaderIsElectedOrChanged(zkClient, topic1, 0)
+    TestUtils.waitUntilMetadataIsPropagated(servers, topic1, 0)
+
+    // retry the metadata for the first auto created topic
+    val response2 = sendMetadataRequest(new MetadataRequest.Builder(Seq(topic1).asJava, true).build)
+    topicMetadata1 = response2.topicMetadata.asScala.head
+    assertEquals(Errors.NONE, topicMetadata1.error)
+    assertEquals(Seq(Errors.NONE), topicMetadata1.partitionMetadata.asScala.map(_.error))
+    assertEquals(1, topicMetadata1.partitionMetadata.size)
+    val partitionMetadata = topicMetadata1.partitionMetadata.asScala.head
+    assertEquals(0, partitionMetadata.partition)
+    assertEquals(2, partitionMetadata.replicaIds.size)
+    assertTrue(partitionMetadata.leaderId.isPresent)
+    assertTrue(partitionMetadata.leaderId.get >= 0)
+  }
+
+  @Test
+  def testAllTopicsRequest(): Unit = {
     // create some topics
-    TestUtils.createTopic(zkUtils, "t1", 3, 2, servers)
-    TestUtils.createTopic(zkUtils, "t2", 3, 2, servers)
+    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())
 
@@ -167,7 +229,7 @@ class MetadataRequestTest extends BaseRequestTest {
   @Test
   def testPreferredReplica(): Unit = {
     val replicaAssignment = Map(0 -> Seq(1, 2, 0), 1 -> Seq(2, 0, 1))
-    TestUtils.createTopic(zkUtils, "t1", replicaAssignment, servers)
+    createTopic("t1", replicaAssignment)
     // Call controller and one different broker to ensure that metadata propagation works correctly
     val responses = Seq(
       sendMetadataRequest(new MetadataRequest.Builder(Seq("t1").asJava, true).build(), Some(controllerSocketServer)),
@@ -179,64 +241,141 @@ class MetadataRequestTest extends BaseRequestTest {
       assertEquals(Errors.NONE, topicMetadata.error)
       assertEquals("t1", topicMetadata.topic)
       assertEquals(Set(0, 1), topicMetadata.partitionMetadata.asScala.map(_.partition).toSet)
-      topicMetadata.partitionMetadata.asScala.foreach { partitionMetadata =>
+      topicMetadata.partitionMetadata.forEach { partitionMetadata =>
         val assignment = replicaAssignment(partitionMetadata.partition)
-        assertEquals(assignment, partitionMetadata.replicas.asScala.map(_.id))
-        assertEquals(assignment, partitionMetadata.isr.asScala.map(_.id))
-        assertEquals(assignment.head, partitionMetadata.leader.id)
+        assertEquals(assignment, partitionMetadata.replicaIds.asScala)
+        assertEquals(assignment, partitionMetadata.inSyncReplicaIds.asScala)
+        assertEquals(Optional.of(assignment.head), partitionMetadata.leaderId)
       }
     }
   }
 
+  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
 
     // create a topic with 3 replicas
-    TestUtils.createTopic(zkUtils, replicaDownTopic, 1, replicaCount, servers)
+    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).build())
     val partitionMetadata = metadataResponse.topicMetadata.asScala.head.partitionMetadata.asScala.head
     val downNode = servers.find { server =>
-      val serverId = server.apis.brokerId
-      val leaderId = partitionMetadata.leader.id
-      val replicaIds = partitionMetadata.replicas.asScala.map(_.id)
-      serverId != leaderId && replicaIds.contains(serverId)
+      val serverId = server.dataPlaneRequestProcessor.brokerId
+      val leaderId = partitionMetadata.leaderId
+      val replicaIds = partitionMetadata.replicaIds.asScala
+      leaderId.isPresent && leaderId.get() != serverId && replicaIds.contains(serverId)
     }.get
     downNode.shutdown()
 
     TestUtils.waitUntilTrue(() => {
-      val response = sendMetadataRequest(new MetadataRequest(List(replicaDownTopic).asJava, true, 1.toShort))
-      val metadata = response.topicMetadata.asScala.head.partitionMetadata.asScala.head
-      val replica = metadata.replicas.asScala.find(_.id == downNode.apis.brokerId).get
-      replica.host == "" & replica.port == -1
+      val response = sendMetadataRequest(new MetadataRequest.Builder(List(replicaDownTopic).asJava, true).build())
+      !response.brokers.asScala.exists(_.id == downNode.dataPlaneRequestProcessor.brokerId)
     }, "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))
     assertTrue("Response should have one topic", v0MetadataResponse.topicMetadata.size == 1)
     val v0PartitionMetadata = v0MetadataResponse.topicMetadata.asScala.head.partitionMetadata.asScala.head
     assertTrue("PartitionMetadata should have an error", v0PartitionMetadata.error == Errors.REPLICA_NOT_AVAILABLE)
-    assertTrue(s"Response should have ${replicaCount - 1} replicas", v0PartitionMetadata.replicas.size == replicaCount - 1)
+    assertTrue(s"Response should have ${replicaCount - 1} replicas", v0PartitionMetadata.replicaIds.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).build(1))
     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))
     assertEquals("Response should have one topic", 1, v1MetadataResponse.topicMetadata.size)
     val v1PartitionMetadata = v1MetadataResponse.topicMetadata.asScala.head.partitionMetadata.asScala.head
     assertEquals("PartitionMetadata should have no errors", Errors.NONE, v1PartitionMetadata.error)
-    assertEquals(s"Response should have $replicaCount replicas", replicaCount, v1PartitionMetadata.replicas.size)
+    assertEquals(s"Response should have $replicaCount replicas", replicaCount, v1PartitionMetadata.replicaIds.size)
+  }
+
+  @Test
+  def testIsrAfterBrokerShutDownAndJoinsBack(): Unit = {
+    def checkIsr(servers: Seq[KafkaServer], topic: String): Unit = {
+      val activeBrokers = servers.filter(_.brokerState.currentState != NotRunning.state)
+      val expectedIsr = activeBrokers.map(_.config.brokerId).toSet
+
+      // Assert that topic metadata at new brokers is updated correctly
+      activeBrokers.foreach { broker =>
+        var actualIsr = Set.empty[Int]
+        TestUtils.waitUntilTrue(() => {
+          val metadataResponse = sendMetadataRequest(new MetadataRequest.Builder(Seq(topic).asJava, false).build,
+            Some(brokerSocketServer(broker.config.brokerId)))
+          val firstPartitionMetadata = metadataResponse.topicMetadata.asScala.headOption.flatMap(_.partitionMetadata.asScala.headOption)
+          actualIsr = firstPartitionMetadata.map { partitionMetadata =>
+            partitionMetadata.inSyncReplicaIds.asScala.map(Int.unbox).toSet
+          }.getOrElse(Set.empty)
+          expectedIsr == actualIsr
+        }, s"Topic metadata not updated correctly in broker $broker\n" +
+          s"Expected ISR: $expectedIsr \n" +
+          s"Actual ISR : $actualIsr")
+      }
+    }
+
+    val topic = "isr-after-broker-shutdown"
+    val replicaCount = 3
+    createTopic(topic, 1, replicaCount)
+
+    servers.last.shutdown()
+    servers.last.awaitShutdown()
+    servers.last.startup()
+
+    checkIsr(servers, topic)
+  }
+
+  @Test
+  def testAliveBrokersWithNoTopics(): Unit = {
+    def checkMetadata(servers: Seq[KafkaServer], expectedBrokersCount: Int): Unit = {
+      var controllerMetadataResponse: Option[MetadataResponse] = None
+      TestUtils.waitUntilTrue(() => {
+        val metadataResponse = sendMetadataRequest(MetadataRequest.Builder.allTopics.build,
+          Some(controllerSocketServer))
+        controllerMetadataResponse = Some(metadataResponse)
+        metadataResponse.brokers.size == expectedBrokersCount
+      }, s"Expected $expectedBrokersCount brokers, but there are ${controllerMetadataResponse.get.brokers.size} " +
+        "according to the Controller")
+
+      val brokersInController = controllerMetadataResponse.get.brokers.asScala.toSeq.sortBy(_.id)
+
+      // Assert that metadata is propagated correctly
+      servers.filter(_.brokerState.currentState != NotRunning.state).foreach { broker =>
+        TestUtils.waitUntilTrue(() => {
+          val metadataResponse = sendMetadataRequest(MetadataRequest.Builder.allTopics.build,
+            Some(brokerSocketServer(broker.config.brokerId)))
+          val brokers = metadataResponse.brokers.asScala.toSeq.sortBy(_.id)
+          val topicMetadata = metadataResponse.topicMetadata.asScala.toSeq.sortBy(_.topic)
+          brokersInController == brokers && metadataResponse.topicMetadata.asScala.toSeq.sortBy(_.topic) == topicMetadata
+        }, s"Topic metadata not updated correctly")
+      }
+    }
+
+    val serverToShutdown = servers.filterNot(_.kafkaController.isActive).last
+    serverToShutdown.shutdown()
+    serverToShutdown.awaitShutdown()
+    checkMetadata(servers, servers.size - 1)
+
+    serverToShutdown.startup()
+    checkMetadata(servers, servers.size)
   }
 
   private def sendMetadataRequest(request: MetadataRequest, destination: Option[SocketServer] = None): MetadataResponse = {
-    val response = connectAndSend(request, ApiKeys.METADATA, destination = destination.getOrElse(anySocketServer))
-    MetadataResponse.parse(response, request.version)
+    connectAndReceive[MetadataResponse](request, destination = destination.getOrElse(anySocketServer))
   }
+
 }
diff --git a/core/src/test/scala/unit/kafka/server/OffsetCommitTest.scala b/core/src/test/scala/unit/kafka/server/OffsetCommitTest.scala
deleted file mode 100755
index 61e17e308893e..0000000000000
--- a/core/src/test/scala/unit/kafka/server/OffsetCommitTest.scala
+++ /dev/null
@@ -1,331 +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 kafka.api.{GroupCoordinatorRequest, OffsetCommitRequest, OffsetFetchRequest}
-import kafka.consumer.SimpleConsumer
-import kafka.common.{OffsetAndMetadata, OffsetMetadata, OffsetMetadataAndError, TopicAndPartition}
-import kafka.utils._
-import kafka.utils.TestUtils._
-import kafka.zk.ZooKeeperTestHarness
-import org.apache.kafka.common.protocol.Errors
-import org.apache.kafka.common.utils.Time
-import org.junit.{After, Before, Test}
-import org.junit.Assert._
-import java.util.Properties
-import java.io.File
-
-import scala.util.Random
-import scala.collection._
-
-class OffsetCommitTest extends ZooKeeperTestHarness {
-  val random: Random = new Random()
-  val group = "test-group"
-  val retentionCheckInterval: Long = 100L
-  var logDir: File = null
-  var topicLogDir: File = null
-  var server: KafkaServer = null
-  var logSize: Int = 100
-  var simpleConsumer: SimpleConsumer = null
-
-  @Before
-  override def setUp() {
-    super.setUp()
-    val config: Properties = createBrokerConfig(1, zkConnect,  enableDeleteTopic = true)
-    config.setProperty(KafkaConfig.OffsetsTopicReplicationFactorProp, "1")
-    config.setProperty(KafkaConfig.OffsetsRetentionCheckIntervalMsProp, retentionCheckInterval.toString)
-    val logDirPath = config.getProperty("log.dir")
-    logDir = new File(logDirPath)
-    server = TestUtils.createServer(KafkaConfig.fromProps(config), Time.SYSTEM)
-    simpleConsumer = new SimpleConsumer("localhost", TestUtils.boundPort(server), 1000000, 64*1024, "test-client")
-    val consumerMetadataRequest = GroupCoordinatorRequest(group)
-    Stream.continually {
-      val consumerMetadataResponse = simpleConsumer.send(consumerMetadataRequest)
-      consumerMetadataResponse.coordinatorOpt.isDefined
-    }.dropWhile(success => {
-      if (!success) Thread.sleep(1000)
-      !success
-    })
-  }
-
-  @After
-  override def tearDown() {
-    simpleConsumer.close
-    TestUtils.shutdownServers(Seq(server))
-    super.tearDown()
-  }
-
-  @Test
-  def testUpdateOffsets() {
-    val topic = "topic"
-
-    // Commit an offset
-    val topicAndPartition = TopicAndPartition(topic, 0)
-    val expectedReplicaAssignment = Map(0  -> List(1))
-    // create the topic
-    createTopic(zkUtils, topic, partitionReplicaAssignment = expectedReplicaAssignment, servers = Seq(server))
-
-    val commitRequest = OffsetCommitRequest(group, immutable.Map(topicAndPartition -> OffsetAndMetadata(offset = 42L)))
-    val commitResponse = simpleConsumer.commitOffsets(commitRequest)
-
-    assertEquals(Errors.NONE, commitResponse.commitStatus.get(topicAndPartition).get)
-
-    // Fetch it and verify
-    val fetchRequest = OffsetFetchRequest(group, Seq(topicAndPartition))
-    val fetchResponse = simpleConsumer.fetchOffsets(fetchRequest)
-
-    assertEquals(Errors.NONE, fetchResponse.requestInfo.get(topicAndPartition).get.error)
-    assertEquals(OffsetMetadata.NoMetadata, fetchResponse.requestInfo.get(topicAndPartition).get.metadata)
-    assertEquals(42L, fetchResponse.requestInfo.get(topicAndPartition).get.offset)
-
-    // Commit a new offset
-    val commitRequest1 = OffsetCommitRequest(group, immutable.Map(topicAndPartition -> OffsetAndMetadata(
-      offset=100L,
-      metadata="some metadata"
-    )))
-    val commitResponse1 = simpleConsumer.commitOffsets(commitRequest1)
-
-    assertEquals(Errors.NONE, commitResponse1.commitStatus.get(topicAndPartition).get)
-
-    // Fetch it and verify
-    val fetchRequest1 = OffsetFetchRequest(group, Seq(topicAndPartition))
-    val fetchResponse1 = simpleConsumer.fetchOffsets(fetchRequest1)
-
-    assertEquals(Errors.NONE, fetchResponse1.requestInfo.get(topicAndPartition).get.error)
-    assertEquals("some metadata", fetchResponse1.requestInfo.get(topicAndPartition).get.metadata)
-    assertEquals(100L, fetchResponse1.requestInfo.get(topicAndPartition).get.offset)
-
-    // Fetch an unknown topic and verify
-    val unknownTopicAndPartition = TopicAndPartition("unknownTopic", 0)
-    val fetchRequest2 = OffsetFetchRequest(group, Seq(unknownTopicAndPartition))
-    val fetchResponse2 = simpleConsumer.fetchOffsets(fetchRequest2)
-
-    assertEquals(OffsetMetadataAndError.NoOffset, fetchResponse2.requestInfo.get(unknownTopicAndPartition).get)
-    assertEquals(1, fetchResponse2.requestInfo.size)
-  }
-
-  @Test
-  def testCommitAndFetchOffsets() {
-    val topic1 = "topic-1"
-    val topic2 = "topic-2"
-    val topic3 = "topic-3"
-    val topic4 = "topic-4" // Topic that group never consumes
-    val topic5 = "topic-5" // Non-existent topic
-
-    createTopic(zkUtils, topic1, servers = Seq(server), numPartitions = 1)
-    createTopic(zkUtils, topic2, servers = Seq(server), numPartitions = 2)
-    createTopic(zkUtils, topic3, servers = Seq(server), numPartitions = 1)
-    createTopic(zkUtils, topic4, servers = Seq(server), numPartitions = 1)
-
-    val commitRequest = OffsetCommitRequest("test-group", immutable.Map(
-      TopicAndPartition(topic1, 0) -> OffsetAndMetadata(offset=42L, metadata="metadata one"),
-      TopicAndPartition(topic2, 0) -> OffsetAndMetadata(offset=43L, metadata="metadata two"),
-      TopicAndPartition(topic3, 0) -> OffsetAndMetadata(offset=44L, metadata="metadata three"),
-      TopicAndPartition(topic2, 1) -> OffsetAndMetadata(offset=45L)
-    ))
-    val commitResponse = simpleConsumer.commitOffsets(commitRequest)
-    assertEquals(Errors.NONE, commitResponse.commitStatus.get(TopicAndPartition(topic1, 0)).get)
-    assertEquals(Errors.NONE, commitResponse.commitStatus.get(TopicAndPartition(topic2, 0)).get)
-    assertEquals(Errors.NONE, commitResponse.commitStatus.get(TopicAndPartition(topic3, 0)).get)
-    assertEquals(Errors.NONE, commitResponse.commitStatus.get(TopicAndPartition(topic2, 1)).get)
-
-    val fetchRequest = OffsetFetchRequest(group, Seq(
-      TopicAndPartition(topic1, 0),
-      TopicAndPartition(topic2, 0),
-      TopicAndPartition(topic3, 0),
-      TopicAndPartition(topic2, 1),
-      TopicAndPartition(topic3, 1), // An unknown partition
-      TopicAndPartition(topic4, 0), // An unused topic
-      TopicAndPartition(topic5, 0)  // An unknown topic
-    ))
-    val fetchResponse = simpleConsumer.fetchOffsets(fetchRequest)
-
-    assertEquals(Errors.NONE, fetchResponse.requestInfo.get(TopicAndPartition(topic1, 0)).get.error)
-
-    assertEquals(Errors.NONE, fetchResponse.requestInfo.get(TopicAndPartition(topic2, 0)).get.error)
-    assertEquals(Errors.NONE, fetchResponse.requestInfo.get(TopicAndPartition(topic2, 1)).get.error)
-
-    assertEquals(Errors.NONE, fetchResponse.requestInfo.get(TopicAndPartition(topic3, 0)).get.error)
-    assertEquals(Errors.NONE, fetchResponse.requestInfo.get(TopicAndPartition(topic3, 1)).get.error)
-    assertEquals(OffsetMetadataAndError.NoOffset, fetchResponse.requestInfo.get(TopicAndPartition(topic3, 1)).get)
-
-    assertEquals(Errors.NONE, fetchResponse.requestInfo.get(TopicAndPartition(topic4, 0)).get.error)
-    assertEquals(OffsetMetadataAndError.NoOffset, fetchResponse.requestInfo.get(TopicAndPartition(topic4, 0)).get)
-
-    assertEquals(Errors.NONE, fetchResponse.requestInfo.get(TopicAndPartition(topic5, 0)).get.error)
-    assertEquals(OffsetMetadataAndError.NoOffset, fetchResponse.requestInfo.get(TopicAndPartition(topic5, 0)).get)
-
-    assertEquals("metadata one", fetchResponse.requestInfo.get(TopicAndPartition(topic1, 0)).get.metadata)
-    assertEquals("metadata two", fetchResponse.requestInfo.get(TopicAndPartition(topic2, 0)).get.metadata)
-    assertEquals("metadata three", fetchResponse.requestInfo.get(TopicAndPartition(topic3, 0)).get.metadata)
-
-    assertEquals(OffsetMetadata.NoMetadata, fetchResponse.requestInfo.get(TopicAndPartition(topic2, 1)).get.metadata)
-    assertEquals(OffsetMetadata.NoMetadata, fetchResponse.requestInfo.get(TopicAndPartition(topic3, 1)).get.metadata)
-    assertEquals(OffsetMetadata.NoMetadata, fetchResponse.requestInfo.get(TopicAndPartition(topic4, 0)).get.metadata)
-    assertEquals(OffsetMetadata.NoMetadata, fetchResponse.requestInfo.get(TopicAndPartition(topic5, 0)).get.metadata)
-
-    assertEquals(42L, fetchResponse.requestInfo.get(TopicAndPartition(topic1, 0)).get.offset)
-    assertEquals(43L, fetchResponse.requestInfo.get(TopicAndPartition(topic2, 0)).get.offset)
-    assertEquals(44L, fetchResponse.requestInfo.get(TopicAndPartition(topic3, 0)).get.offset)
-    assertEquals(45L, fetchResponse.requestInfo.get(TopicAndPartition(topic2, 1)).get.offset)
-
-    assertEquals(OffsetMetadata.InvalidOffset, fetchResponse.requestInfo.get(TopicAndPartition(topic3, 1)).get.offset)
-    assertEquals(OffsetMetadata.InvalidOffset, fetchResponse.requestInfo.get(TopicAndPartition(topic4, 0)).get.offset)
-    assertEquals(OffsetMetadata.InvalidOffset, fetchResponse.requestInfo.get(TopicAndPartition(topic5, 0)).get.offset)
-  }
-
-  @Test
-  def testLargeMetadataPayload() {
-    val topicAndPartition = TopicAndPartition("large-metadata", 0)
-    val expectedReplicaAssignment = Map(0  -> List(1))
-    createTopic(zkUtils, topicAndPartition.topic, partitionReplicaAssignment = expectedReplicaAssignment,
-                servers = Seq(server))
-
-    val commitRequest = OffsetCommitRequest("test-group", immutable.Map(topicAndPartition -> OffsetAndMetadata(
-      offset=42L,
-      metadata=random.nextString(server.config.offsetMetadataMaxSize)
-    )))
-    val commitResponse = simpleConsumer.commitOffsets(commitRequest)
-
-    assertEquals(Errors.NONE, commitResponse.commitStatus.get(topicAndPartition).get)
-
-    val commitRequest1 = OffsetCommitRequest(group, immutable.Map(topicAndPartition -> OffsetAndMetadata(
-      offset=42L,
-      metadata=random.nextString(server.config.offsetMetadataMaxSize + 1)
-    )))
-    val commitResponse1 = simpleConsumer.commitOffsets(commitRequest1)
-
-    assertEquals(Errors.OFFSET_METADATA_TOO_LARGE, commitResponse1.commitStatus.get(topicAndPartition).get)
-  }
-
-  @Test
-  def testOffsetExpiration() {
-    // set up topic partition
-    val topic = "topic"
-    val topicPartition = TopicAndPartition(topic, 0)
-    createTopic(zkUtils, topic, servers = Seq(server), numPartitions = 1)
-
-    val fetchRequest = OffsetFetchRequest(group, Seq(TopicAndPartition(topic, 0)))
-
-    // v0 version commit request
-    // committed offset should not exist with fetch version 1 since it was stored in ZK
-    val commitRequest0 = OffsetCommitRequest(
-      groupId = group,
-      requestInfo = immutable.Map(topicPartition -> OffsetAndMetadata(1L, "metadata")),
-      versionId = 0
-    )
-    assertEquals(Errors.NONE, simpleConsumer.commitOffsets(commitRequest0).commitStatus.get(topicPartition).get)
-    assertEquals(-1L, simpleConsumer.fetchOffsets(fetchRequest).requestInfo.get(topicPartition).get.offset)
-
-    // committed offset should exist with fetch version 0
-    val offsetFetchReq = OffsetFetchRequest(group, Seq(TopicAndPartition(topic, 0)), versionId = 0)
-    val offsetFetchResp = simpleConsumer.fetchOffsets(offsetFetchReq)
-    assertEquals(1L, offsetFetchResp.requestInfo.get(topicPartition).get.offset)
-
-
-    // v1 version commit request with commit timestamp set to -1
-    // committed offset should not expire
-    val commitRequest1 = OffsetCommitRequest(
-      groupId = group,
-      requestInfo = immutable.Map(topicPartition -> OffsetAndMetadata(2L, "metadata", -1L)),
-      versionId = 1
-    )
-    assertEquals(Errors.NONE, simpleConsumer.commitOffsets(commitRequest1).commitStatus.get(topicPartition).get)
-    Thread.sleep(retentionCheckInterval * 2)
-    assertEquals(2L, simpleConsumer.fetchOffsets(fetchRequest).requestInfo.get(topicPartition).get.offset)
-
-    // v1 version commit request with commit timestamp set to now - two days
-    // committed offset should expire
-    val commitRequest2 = OffsetCommitRequest(
-      groupId = group,
-      requestInfo = immutable.Map(topicPartition -> OffsetAndMetadata(3L, "metadata", Time.SYSTEM.milliseconds - 2*24*60*60*1000L)),
-      versionId = 1
-    )
-    assertEquals(Errors.NONE, simpleConsumer.commitOffsets(commitRequest2).commitStatus.get(topicPartition).get)
-    Thread.sleep(retentionCheckInterval * 2)
-    assertEquals(-1L, simpleConsumer.fetchOffsets(fetchRequest).requestInfo.get(topicPartition).get.offset)
-
-    // v2 version commit request with retention time set to 1 hour
-    // committed offset should not expire
-    val commitRequest3 = OffsetCommitRequest(
-      groupId = group,
-      requestInfo = immutable.Map(topicPartition -> OffsetAndMetadata(4L, "metadata", -1L)),
-      versionId = 2,
-      retentionMs = 1000 * 60 * 60L
-    )
-    assertEquals(Errors.NONE, simpleConsumer.commitOffsets(commitRequest3).commitStatus.get(topicPartition).get)
-    Thread.sleep(retentionCheckInterval * 2)
-    assertEquals(4L, simpleConsumer.fetchOffsets(fetchRequest).requestInfo.get(topicPartition).get.offset)
-
-    // v2 version commit request with retention time set to 0 second
-    // committed offset should expire
-    val commitRequest4 = OffsetCommitRequest(
-      groupId = "test-group",
-      requestInfo = immutable.Map(TopicAndPartition(topic, 0) -> OffsetAndMetadata(5L, "metadata", -1L)),
-      versionId = 2,
-      retentionMs = 0L
-    )
-    assertEquals(Errors.NONE, simpleConsumer.commitOffsets(commitRequest4).commitStatus.get(topicPartition).get)
-    Thread.sleep(retentionCheckInterval * 2)
-    assertEquals(-1L, simpleConsumer.fetchOffsets(fetchRequest).requestInfo.get(topicPartition).get.offset)
-
-  }
-
-  @Test
-  def testNonExistingTopicOffsetCommit() {
-    val topic1 = "topicDoesNotExists"
-    val topic2 = "topic-2"
-
-    createTopic(zkUtils, topic2, servers = Seq(server), numPartitions = 1)
-
-    // Commit an offset
-    val commitRequest = OffsetCommitRequest(group, immutable.Map(
-      TopicAndPartition(topic1, 0) -> OffsetAndMetadata(offset=42L),
-      TopicAndPartition(topic2, 0) -> OffsetAndMetadata(offset=42L)
-    ))
-    val commitResponse = simpleConsumer.commitOffsets(commitRequest)
-
-    assertEquals(Errors.UNKNOWN_TOPIC_OR_PARTITION, commitResponse.commitStatus.get(TopicAndPartition(topic1, 0)).get)
-    assertEquals(Errors.NONE, commitResponse.commitStatus.get(TopicAndPartition(topic2, 0)).get)
-  }
-
-  @Test
-  def testOffsetsDeleteAfterTopicDeletion() {
-    // set up topic partition
-    val topic = "topic"
-    val topicPartition = TopicAndPartition(topic, 0)
-    createTopic(zkUtils, topic, servers = Seq(server), numPartitions = 1)
-
-    val commitRequest = OffsetCommitRequest(group, immutable.Map(topicPartition -> OffsetAndMetadata(offset = 42L)))
-    val commitResponse = simpleConsumer.commitOffsets(commitRequest)
-
-    assertEquals(Errors.NONE, commitResponse.commitStatus.get(topicPartition).get)
-
-    // start topic deletion
-    adminZkClient.deleteTopic(topic)
-    TestUtils.verifyTopicDeletion(zkUtils, topic, 1, Seq(server))
-    Thread.sleep(retentionCheckInterval * 2)
-
-    // check if offsets deleted
-    val fetchRequest = OffsetFetchRequest(group, Seq(TopicAndPartition(topic, 0)))
-    val offsetMetadataAndErrorMap = simpleConsumer.fetchOffsets(fetchRequest)
-    val offsetMetadataAndError = offsetMetadataAndErrorMap.requestInfo(topicPartition)
-    assertEquals(OffsetMetadataAndError.NoOffset, offsetMetadataAndError)
-  }
-
-}
diff --git a/core/src/test/scala/unit/kafka/server/OffsetsForLeaderEpochRequestTest.scala b/core/src/test/scala/unit/kafka/server/OffsetsForLeaderEpochRequestTest.scala
new file mode 100644
index 0000000000000..c04cb5c172995
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/OffsetsForLeaderEpochRequestTest.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.Optional
+
+import kafka.utils.TestUtils
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.protocol.{ApiKeys, Errors}
+import org.apache.kafka.common.requests.{OffsetsForLeaderEpochRequest, OffsetsForLeaderEpochResponse}
+import org.junit.Assert._
+import org.junit.Test
+
+import scala.jdk.CollectionConverters._
+
+class OffsetsForLeaderEpochRequestTest extends BaseRequestTest {
+
+  @Test
+  def testOffsetsForLeaderEpochErrorCodes(): Unit = {
+    val topic = "topic"
+    val partition = new TopicPartition(topic, 0)
+
+    val epochs = Map(partition -> new OffsetsForLeaderEpochRequest.PartitionData(Optional.empty[Integer], 0)).asJava
+    val request = OffsetsForLeaderEpochRequest.Builder.forFollower(
+      ApiKeys.OFFSET_FOR_LEADER_EPOCH.latestVersion, epochs, 1).build()
+
+    // Unknown topic
+    val randomBrokerId = servers.head.config.brokerId
+    assertResponseError(Errors.UNKNOWN_TOPIC_OR_PARTITION, randomBrokerId, request)
+
+    val partitionToLeader = TestUtils.createTopic(zkClient, topic, numPartitions = 1, replicationFactor = 2, servers)
+    val replicas = zkClient.getReplicasForPartition(partition).toSet
+    val leader = partitionToLeader(partition.partition)
+    val follower = replicas.find(_ != leader).get
+    val nonReplica = servers.map(_.config.brokerId).find(!replicas.contains(_)).get
+
+    assertResponseError(Errors.NOT_LEADER_OR_FOLLOWER, follower, request)
+    assertResponseError(Errors.NOT_LEADER_OR_FOLLOWER, nonReplica, request)
+  }
+
+  @Test
+  def testCurrentEpochValidation(): Unit = {
+    val topic = "topic"
+    val topicPartition = new TopicPartition(topic, 0)
+    val partitionToLeader = TestUtils.createTopic(zkClient, topic, numPartitions = 1, replicationFactor = 3, servers)
+    val firstLeaderId = partitionToLeader(topicPartition.partition)
+
+    def assertResponseErrorForEpoch(error: Errors, brokerId: Int, currentLeaderEpoch: Optional[Integer]): Unit = {
+      val epochs = Map(topicPartition -> new OffsetsForLeaderEpochRequest.PartitionData(
+        currentLeaderEpoch, 0)).asJava
+      val request = OffsetsForLeaderEpochRequest.Builder.forFollower(
+        ApiKeys.OFFSET_FOR_LEADER_EPOCH.latestVersion, epochs, 1).build()
+      assertResponseError(error, brokerId, request)
+    }
+
+    // We need a leader change in order to check epoch fencing since the first epoch is 0 and
+    // -1 is treated as having no epoch at all
+    killBroker(firstLeaderId)
+
+    // Check leader error codes
+    val secondLeaderId = TestUtils.awaitLeaderChange(servers, topicPartition, firstLeaderId)
+    val secondLeaderEpoch = TestUtils.findLeaderEpoch(secondLeaderId, topicPartition, servers)
+    assertResponseErrorForEpoch(Errors.NONE, secondLeaderId, Optional.empty())
+    assertResponseErrorForEpoch(Errors.NONE, secondLeaderId, Optional.of(secondLeaderEpoch))
+    assertResponseErrorForEpoch(Errors.FENCED_LEADER_EPOCH, secondLeaderId, Optional.of(secondLeaderEpoch - 1))
+    assertResponseErrorForEpoch(Errors.UNKNOWN_LEADER_EPOCH, secondLeaderId, Optional.of(secondLeaderEpoch + 1))
+
+    // Check follower error codes
+    val followerId = TestUtils.findFollowerId(topicPartition, servers)
+    assertResponseErrorForEpoch(Errors.NOT_LEADER_OR_FOLLOWER, followerId, Optional.empty())
+    assertResponseErrorForEpoch(Errors.NOT_LEADER_OR_FOLLOWER, followerId, Optional.of(secondLeaderEpoch))
+    assertResponseErrorForEpoch(Errors.UNKNOWN_LEADER_EPOCH, followerId, Optional.of(secondLeaderEpoch + 1))
+    assertResponseErrorForEpoch(Errors.FENCED_LEADER_EPOCH, followerId, Optional.of(secondLeaderEpoch - 1))
+  }
+
+  private def assertResponseError(error: Errors, brokerId: Int, request: OffsetsForLeaderEpochRequest): Unit = {
+    val response = sendRequest(brokerId, request)
+    assertEquals(request.data.topics.size, response.data.topics.size)
+    response.data.topics.asScala.foreach { offsetForLeaderTopic =>
+      assertEquals(request.data.topics.find(offsetForLeaderTopic.topic).partitions.size,
+        offsetForLeaderTopic.partitions.size)
+      offsetForLeaderTopic.partitions.asScala.foreach { offsetForLeaderPartition =>
+        assertEquals(error.code(), offsetForLeaderPartition.errorCode())
+      }
+    }
+  }
+
+  private def sendRequest(brokerId: Int, request: OffsetsForLeaderEpochRequest): OffsetsForLeaderEpochResponse = {
+    connectAndReceive[OffsetsForLeaderEpochResponse](request, destination = brokerSocketServer(brokerId))
+  }
+
+}
diff --git a/core/src/test/scala/unit/kafka/server/ProduceRequestTest.scala b/core/src/test/scala/unit/kafka/server/ProduceRequestTest.scala
index 5141ad5b0bce0..ca43d3b8dba05 100644
--- a/core/src/test/scala/unit/kafka/server/ProduceRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ProduceRequestTest.scala
@@ -17,15 +17,24 @@
 
 package kafka.server
 
+import java.nio.ByteBuffer
+import java.util.{Collections, Properties}
+
+import kafka.log.LogConfig
+import kafka.message.ZStdCompressionCodec
+import kafka.metrics.KafkaYammerMetrics
 import kafka.utils.TestUtils
 import org.apache.kafka.common.TopicPartition
-import org.apache.kafka.common.protocol.{ApiKeys, Errors}
-import org.apache.kafka.common.record.{CompressionType, DefaultRecordBatch, MemoryRecords, SimpleRecord}
+import org.apache.kafka.common.message.ProduceRequestData
+import org.apache.kafka.common.protocol.Errors
+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._
+import scala.annotation.nowarn
+import scala.jdk.CollectionConverters._
 
 /**
   * Subclasses of `BaseProduceSendRequestTest` exercise the producer and produce request/response. This class
@@ -33,21 +42,33 @@ import scala.collection.JavaConverters._
   */
 class ProduceRequestTest extends BaseRequestTest {
 
+  val metricsKeySet = KafkaYammerMetrics.defaultRegistry.allMetrics.keySet.asScala
+
+  @nowarn("cat=deprecation")
   @Test
-  def testSimpleProduceRequest() {
+  def testSimpleProduceRequest(): Unit = {
     val (partition, leader) = createTopicAndFindPartitionWithLeader("topic")
 
     def sendAndCheck(memoryRecords: MemoryRecords, expectedOffset: Long): ProduceResponse.PartitionResponse = {
       val topicPartition = new TopicPartition("topic", partition)
-      val partitionRecords = Map(topicPartition -> memoryRecords)
       val produceResponse = sendProduceRequest(leader,
-          ProduceRequest.Builder.forCurrentMagic(-1, 3000, partitionRecords.asJava).build())
+          ProduceRequest.forCurrentMagic(new ProduceRequestData()
+            .setTopicData(new ProduceRequestData.TopicProduceDataCollection(Collections.singletonList(
+              new ProduceRequestData.TopicProduceData()
+                .setName(topicPartition.topic())
+                .setPartitionData(Collections.singletonList(new ProduceRequestData.PartitionProduceData()
+                  .setIndex(topicPartition.partition())
+                  .setRecords(memoryRecords)))).iterator))
+            .setAcks((-1).toShort)
+            .setTimeoutMs(3000)
+            .setTransactionalId(null)).build())
       assertEquals(1, produceResponse.responses.size)
       val (tp, partitionResponse) = produceResponse.responses.asScala.head
       assertEquals(topicPartition, tp)
       assertEquals(Errors.NONE, partitionResponse.error)
       assertEquals(expectedOffset, partitionResponse.baseOffset)
       assertEquals(-1, partitionResponse.logAppendTime)
+      assertTrue(partitionResponse.recordErrors.isEmpty)
       partitionResponse
     }
 
@@ -59,16 +80,94 @@ class ProduceRequestTest extends BaseRequestTest {
       new SimpleRecord(System.currentTimeMillis(), "key2".getBytes, "value2".getBytes)), 1)
   }
 
+  @nowarn("cat=deprecation")
+  @Test
+  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, 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 produceResponse = sendProduceRequest(leader, ProduceRequest.forCurrentMagic(new ProduceRequestData()
+      .setTopicData(new ProduceRequestData.TopicProduceDataCollection(Collections.singletonList(
+        new ProduceRequestData.TopicProduceData()
+          .setName(topicPartition.topic())
+          .setPartitionData(Collections.singletonList(new ProduceRequestData.PartitionProduceData()
+            .setIndex(topicPartition.partition())
+            .setRecords(records)))).iterator))
+      .setAcks((-1).toShort)
+      .setTimeoutMs(3000)
+      .setTransactionalId(null)).build())
+    val (tp, partitionResponse) = produceResponse.responses.asScala.head
+    assertEquals(topicPartition, tp)
+    assertEquals(Errors.INVALID_TIMESTAMP, partitionResponse.error)
+    // there are 3 records with InvalidTimestampException created from inner function createRecords
+    assertEquals(3, partitionResponse.recordErrors.size())
+    assertEquals(0, partitionResponse.recordErrors.get(0).batchIndex)
+    assertEquals(1, partitionResponse.recordErrors.get(1).batchIndex)
+    assertEquals(2, partitionResponse.recordErrors.get(2).batchIndex)
+    for (recordError <- partitionResponse.recordErrors.asScala) {
+      assertNotNull(recordError.message)
+    }
+    assertEquals("One or more records have been rejected due to invalid timestamp", partitionResponse.errorMessage)
+  }
+
+  @nowarn("cat=deprecation")
+  @Test
+  def testProduceToNonReplica(): Unit = {
+    val topic = "topic"
+    val partition = 0
+
+    // Create a single-partition topic and find a broker which is not the leader
+    val partitionToLeader = TestUtils.createTopic(zkClient, topic, numPartitions = 1, 1, servers)
+    val leader = partitionToLeader(partition)
+    val nonReplicaOpt = servers.find(_.config.brokerId != leader)
+    assertTrue(nonReplicaOpt.isDefined)
+    val nonReplicaId =  nonReplicaOpt.get.config.brokerId
+
+    // Send the produce request to the non-replica
+    val records = MemoryRecords.withRecords(CompressionType.NONE, new SimpleRecord("key".getBytes, "value".getBytes))
+    val topicPartition = new TopicPartition("topic", partition)
+    val produceRequest = ProduceRequest.forCurrentMagic(new ProduceRequestData()
+      .setTopicData(new ProduceRequestData.TopicProduceDataCollection(Collections.singletonList(
+        new ProduceRequestData.TopicProduceData()
+          .setName(topicPartition.topic())
+          .setPartitionData(Collections.singletonList(new ProduceRequestData.PartitionProduceData()
+            .setIndex(topicPartition.partition())
+            .setRecords(records)))).iterator))
+      .setAcks((-1).toShort)
+      .setTimeoutMs(3000)
+      .setTransactionalId(null)).build()
+
+    val produceResponse = sendProduceRequest(nonReplicaId, produceRequest)
+    assertEquals(1, produceResponse.responses.size)
+    assertEquals(Errors.NOT_LEADER_OR_FOLLOWER, produceResponse.responses.asScala.head._2.error)
+  }
+
   /* returns a pair of partition id and leader id */
   private def createTopicAndFindPartitionWithLeader(topic: String): (Int, Int) = {
-    val partitionToLeader = TestUtils.createTopic(zkUtils, topic, 3, 2, servers)
+    val partitionToLeader = TestUtils.createTopic(zkClient, topic, 3, 2, servers)
     partitionToLeader.collectFirst {
       case (partition, leader) if leader != -1 => (partition, leader)
     }.getOrElse(fail(s"No leader elected for topic $topic"))
   }
 
+  @nowarn("cat=deprecation")
   @Test
-  def testCorruptLz4ProduceRequest() {
+  def testCorruptLz4ProduceRequest(): Unit = {
     val (partition, leader) = createTopicAndFindPartitionWithLeader("topic")
     val timestamp = 1000000
     val memoryRecords = MemoryRecords.withRecords(CompressionType.LZ4,
@@ -77,20 +176,72 @@ class ProduceRequestTest extends BaseRequestTest {
     val lz4ChecksumOffset = 6
     memoryRecords.buffer.array.update(DefaultRecordBatch.RECORD_BATCH_OVERHEAD + lz4ChecksumOffset, 0)
     val topicPartition = new TopicPartition("topic", partition)
-    val partitionRecords = Map(topicPartition -> memoryRecords)
-    val produceResponse = sendProduceRequest(leader, 
-      ProduceRequest.Builder.forCurrentMagic(-1, 3000, partitionRecords.asJava).build())
+    val produceResponse = sendProduceRequest(leader, ProduceRequest.forCurrentMagic(new ProduceRequestData()
+      .setTopicData(new ProduceRequestData.TopicProduceDataCollection(Collections.singletonList(
+        new ProduceRequestData.TopicProduceData()
+          .setName(topicPartition.topic())
+          .setPartitionData(Collections.singletonList(new ProduceRequestData.PartitionProduceData()
+            .setIndex(topicPartition.partition())
+            .setRecords(memoryRecords)))).iterator))
+      .setAcks((-1).toShort)
+      .setTimeoutMs(3000)
+      .setTransactionalId(null)).build())
     assertEquals(1, produceResponse.responses.size)
     val (tp, partitionResponse) = produceResponse.responses.asScala.head
     assertEquals(topicPartition, tp)
     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)
+  }
+
+  @nowarn("cat=deprecation")
+  @Test
+  def testZSTDProduceRequest(): Unit = {
+    val topic = "topic"
+    val partition = 0
+
+    // Create a single-partition topic compressed with ZSTD
+    val topicConfig = new Properties
+    topicConfig.setProperty(LogConfig.CompressionTypeProp, ZStdCompressionCodec.name)
+    val partitionToLeader = TestUtils.createTopic(zkClient, topic, 1, 1, servers, topicConfig)
+    val leader = partitionToLeader(partition)
+    val memoryRecords = MemoryRecords.withRecords(CompressionType.ZSTD,
+      new SimpleRecord(System.currentTimeMillis(), "key".getBytes, "value".getBytes))
+    val topicPartition = new TopicPartition("topic", partition)
+    val partitionRecords = new ProduceRequestData()
+      .setTopicData(new ProduceRequestData.TopicProduceDataCollection(Collections.singletonList(
+        new ProduceRequestData.TopicProduceData()
+          .setName("topic").setPartitionData(Collections.singletonList(
+            new ProduceRequestData.PartitionProduceData()
+              .setIndex(partition)
+              .setRecords(memoryRecords))))
+        .iterator))
+      .setAcks((-1).toShort)
+      .setTimeoutMs(3000)
+      .setTransactionalId(null)
+
+    // produce request with v7: works fine!
+    val res1 = sendProduceRequest(leader,
+      new ProduceRequest.Builder(7, 7, partitionRecords).build())
+    val (tp1, partitionResponse1) = res1.responses.asScala.head
+    assertEquals(topicPartition, tp1)
+    assertEquals(Errors.NONE, partitionResponse1.error)
+    assertEquals(0, partitionResponse1.baseOffset)
+    assertEquals(-1, partitionResponse1.logAppendTime)
+
+    // produce request with v3: returns Errors.UNSUPPORTED_COMPRESSION_TYPE.
+    val res2 = sendProduceRequest(leader,
+      new ProduceRequest.Builder(3, 3, partitionRecords)
+        .buildUnsafe(3))
+    val (tp2, partitionResponse2) = res2.responses.asScala.head
+    assertEquals(topicPartition, tp2)
+    assertEquals(Errors.UNSUPPORTED_COMPRESSION_TYPE, partitionResponse2.error)
   }
 
   private def sendProduceRequest(leaderId: Int, request: ProduceRequest): ProduceResponse = {
-    val response = connectAndSend(request, ApiKeys.PRODUCE, destination = brokerSocketServer(leaderId))
-    ProduceResponse.parse(response, request.version)
+    connectAndReceive[ProduceResponse](request, destination = brokerSocketServer(leaderId))
   }
 
 }
diff --git a/core/src/test/scala/unit/kafka/server/ReplicaAlterLogDirsThreadTest.scala b/core/src/test/scala/unit/kafka/server/ReplicaAlterLogDirsThreadTest.scala
new file mode 100644
index 0000000000000..fd783273c9bac
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/ReplicaAlterLogDirsThreadTest.scala
@@ -0,0 +1,937 @@
+/**
+  * Licensed to the Apache Software Foundation (ASF) under one or more
+  * contributor license agreements.  See the NOTICE file distributed with
+  * this work for additional information regarding copyright ownership.
+  * The ASF licenses this file to You under the Apache License, Version 2.0
+  * (the "License"); you may not use this file except in compliance with
+  * the License.  You may obtain a copy of the License at
+  *
+  * http://www.apache.org/licenses/LICENSE-2.0
+  *
+  * Unless required by applicable law or agreed to in writing, software
+  * distributed under the License is distributed on an "AS IS" BASIS,
+  * WITHOUT WARRANTIES 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
+
+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.errors.KafkaStorageException
+import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData.EpochEndOffset
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.record.MemoryRecords
+import org.apache.kafka.common.requests.{FetchRequest, OffsetsForLeaderEpochRequest}
+import org.apache.kafka.common.{IsolationLevel, TopicPartition}
+import org.easymock.EasyMock._
+import org.easymock.{Capture, CaptureType, EasyMock, IExpectationSetters}
+import org.junit.Assert._
+import org.junit.Test
+import org.mockito.Mockito.{doNothing, when}
+import org.mockito.{ArgumentCaptor, ArgumentMatchers, Mockito}
+
+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 initialFetchState(fetchOffset: Long, leaderEpoch: Int = 1): InitialFetchState = {
+    InitialFetchState(leader = new BrokerEndPoint(0, "localhost", 9092),
+      initOffset = fetchOffset, currentLeaderEpoch = 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 -> initialFetchState(0L)))
+    assertEquals(Set.empty, addedPartitions)
+    assertEquals(0, thread.partitionCount)
+    assertEquals(None, thread.fetchState(t1p0))
+  }
+
+  @Test
+  def shouldUpdateLeaderEpochAfterFencedEpochError(): Unit = {
+    val brokerId = 1
+    val partitionId = 0
+    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(partition.partitionId).thenReturn(partitionId)
+    when(replicaManager.futureLocalLogOrException(t1p0)).thenReturn(futureLog)
+    when(replicaManager.futureLogExists(t1p0)).thenReturn(true)
+    when(replicaManager.nonOfflinePartition(t1p0)).thenReturn(Some(partition))
+    when(replicaManager.getPartitionOrException(t1p0)).thenReturn(partition)
+
+    when(quotaManager.isQuotaExceeded).thenReturn(false)
+
+    when(partition.lastOffsetForLeaderEpoch(Optional.empty(), leaderEpoch, fetchOnlyFromLeader = false))
+      .thenReturn(new EpochEndOffset()
+        .setPartition(partitionId)
+        .setErrorCode(Errors.NONE.code)
+        .setLeaderEpoch(leaderEpoch)
+        .setEndOffset(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,
+      divergingEpoch = None,
+      lastStableOffset = None,
+      abortedTransactions = None,
+      preferredReadReplica = None,
+      isReassignmentFetch = false)
+    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 -> initialFetchState(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 -> initialFetchState(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,
+      divergingEpoch = None,
+      lastStableOffset = None,
+      abortedTransactions = None,
+      preferredReadReplica = None,
+      isReassignmentFetch = false)
+    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 partitionId = 0
+    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(partition.partitionId).thenReturn(partitionId)
+    when(replicaManager.futureLocalLogOrException(t1p0)).thenReturn(futureLog)
+    when(replicaManager.futureLogExists(t1p0)).thenReturn(true)
+    when(replicaManager.nonOfflinePartition(t1p0)).thenReturn(Some(partition))
+    when(replicaManager.getPartitionOrException(t1p0)).thenReturn(partition)
+
+    when(quotaManager.isQuotaExceeded).thenReturn(false)
+
+    when(partition.lastOffsetForLeaderEpoch(Optional.empty(), leaderEpoch, fetchOnlyFromLeader = false))
+      .thenReturn(new EpochEndOffset()
+        .setPartition(partitionId)
+        .setErrorCode(Errors.NONE.code)
+        .setLeaderEpoch(leaderEpoch)
+        .setEndOffset(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,
+      divergingEpoch = None,
+      lastStableOffset = None,
+      abortedTransactions = None,
+      preferredReadReplica = None,
+      isReassignmentFetch = false)
+    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 -> initialFetchState(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(_ => {
+      callbackCaptor.getValue.apply(Seq((topicPartition, responseData)))
+    })
+  }
+
+  @Test
+  def issuesEpochRequestFromLocalReplica(): Unit = {
+    val config = KafkaConfig.fromProps(TestUtils.createBrokerConfig(1, "localhost:1234"))
+
+    //Setup all dependencies
+
+    val partitionT1p0: Partition = createMock(classOf[Partition])
+    val partitionT1p1: Partition = createMock(classOf[Partition])
+    val replicaManager: ReplicaManager = createMock(classOf[ReplicaManager])
+
+    val partitionT1p0Id = 0
+    val partitionT1p1Id = 1
+    val leaderEpochT1p0 = 2
+    val leaderEpochT1p1 = 5
+    val leoT1p0 = 13
+    val leoT1p1 = 232
+
+    //Stubs
+    expect(partitionT1p0.partitionId).andStubReturn(partitionT1p0Id)
+    expect(partitionT1p0.partitionId).andStubReturn(partitionT1p1Id)
+
+    expect(replicaManager.getPartitionOrException(t1p0))
+      .andStubReturn(partitionT1p0)
+    expect(partitionT1p0.lastOffsetForLeaderEpoch(Optional.empty(), leaderEpochT1p0, fetchOnlyFromLeader = false))
+      .andReturn(new EpochEndOffset()
+        .setPartition(partitionT1p0Id)
+        .setErrorCode(Errors.NONE.code)
+        .setLeaderEpoch(leaderEpochT1p0)
+        .setEndOffset(leoT1p0))
+      .anyTimes()
+
+    expect(replicaManager.getPartitionOrException(t1p1))
+      .andStubReturn(partitionT1p1)
+    expect(partitionT1p1.lastOffsetForLeaderEpoch(Optional.empty(), leaderEpochT1p1, fetchOnlyFromLeader = false))
+      .andReturn(new EpochEndOffset()
+        .setPartition(partitionT1p1Id)
+        .setErrorCode(Errors.NONE.code)
+        .setLeaderEpoch(leaderEpochT1p1)
+        .setEndOffset(leoT1p1))
+      .anyTimes()
+
+    replay(partitionT1p0, partitionT1p1, replicaManager)
+
+    val endPoint = new BrokerEndPoint(0, "localhost", 1000)
+    val thread = new ReplicaAlterLogDirsThread(
+      "alter-logs-dirs-thread-test1",
+      sourceBroker = endPoint,
+      brokerConfig = config,
+      failedPartitions = failedPartitions,
+      replicaMgr = replicaManager,
+      quota = null,
+      brokerTopicStats = null)
+
+    val result = thread.fetchEpochEndOffsets(Map(
+      t1p0 -> new OffsetsForLeaderEpochRequest.PartitionData(Optional.empty(), leaderEpochT1p0),
+      t1p1 -> new OffsetsForLeaderEpochRequest.PartitionData(Optional.empty(), leaderEpochT1p1)))
+
+    val expected = Map(
+      t1p0 -> new EpochEndOffset()
+        .setPartition(t1p0.partition)
+        .setErrorCode(Errors.NONE.code)
+        .setLeaderEpoch(leaderEpochT1p0)
+        .setEndOffset(leoT1p0),
+      t1p1 -> new EpochEndOffset()
+        .setPartition(t1p1.partition)
+        .setErrorCode(Errors.NONE.code)
+        .setLeaderEpoch(leaderEpochT1p1)
+        .setEndOffset(leoT1p1)
+    )
+
+    assertEquals("results from leader epoch request should have offset from local replica",
+                 expected, result)
+  }
+
+  @Test
+  def fetchEpochsFromLeaderShouldHandleExceptionFromGetLocalReplica(): Unit = {
+    val config = KafkaConfig.fromProps(TestUtils.createBrokerConfig(1, "localhost:1234"))
+
+    //Setup all dependencies
+    val partitionT1p0: Partition = createMock(classOf[Partition])
+    val replicaManager: ReplicaManager = createMock(classOf[ReplicaManager])
+
+    val partitionId = 0
+    val leaderEpoch = 2
+    val leo = 13
+
+    //Stubs
+    expect(partitionT1p0.partitionId).andStubReturn(partitionId)
+
+    expect(replicaManager.getPartitionOrException(t1p0))
+      .andStubReturn(partitionT1p0)
+    expect(partitionT1p0.lastOffsetForLeaderEpoch(Optional.empty(), leaderEpoch, fetchOnlyFromLeader = false))
+      .andReturn(new EpochEndOffset()
+        .setPartition(partitionId)
+        .setErrorCode(Errors.NONE.code)
+        .setLeaderEpoch(leaderEpoch)
+        .setEndOffset(leo))
+      .anyTimes()
+
+    expect(replicaManager.getPartitionOrException(t1p1))
+      .andThrow(new KafkaStorageException).once()
+
+    replay(partitionT1p0, replicaManager)
+
+    val endPoint = new BrokerEndPoint(0, "localhost", 1000)
+    val thread = new ReplicaAlterLogDirsThread(
+      "alter-logs-dirs-thread-test1",
+      sourceBroker = endPoint,
+      brokerConfig = config,
+      failedPartitions = failedPartitions,
+      replicaMgr = replicaManager,
+      quota = null,
+      brokerTopicStats = null)
+
+    val result = thread.fetchEpochEndOffsets(Map(
+      t1p0 -> new OffsetsForLeaderEpochRequest.PartitionData(Optional.empty(), leaderEpoch),
+      t1p1 -> new OffsetsForLeaderEpochRequest.PartitionData(Optional.empty(), leaderEpoch)))
+
+    val expected = Map(
+      t1p0 -> new EpochEndOffset()
+        .setPartition(t1p0.partition)
+        .setErrorCode(Errors.NONE.code)
+        .setLeaderEpoch(leaderEpoch)
+        .setEndOffset(leo),
+      t1p1 -> new EpochEndOffset()
+        .setPartition(t1p1.partition)
+        .setErrorCode(Errors.KAFKA_STORAGE_ERROR.code)
+    )
+
+    assertEquals(expected, result)
+  }
+
+  @Test
+  def shouldTruncateToReplicaOffset(): Unit = {
+
+    //Create a capture to track what partitions/offsets are truncated
+    val truncateCaptureT1p0: Capture[Long] = newCapture(CaptureType.ALL)
+    val truncateCaptureT1p1: Capture[Long] = newCapture(CaptureType.ALL)
+
+    // Setup all the dependencies
+    val config = KafkaConfig.fromProps(TestUtils.createBrokerConfig(1, "localhost:1234"))
+    val quotaManager: ReplicationQuotaManager = createNiceMock(classOf[ReplicationQuotaManager])
+    val logManager: LogManager = createMock(classOf[LogManager])
+    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 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])
+    val responseCallback: Capture[Seq[(TopicPartition, FetchPartitionData)] => Unit]  = EasyMock.newCapture()
+
+    val partitionT1p0Id = 0
+    val partitionT1p1Id = 1
+    val leaderEpoch = 2
+    val futureReplicaLEO = 191
+    val replicaT1p0LEO = 190
+    val replicaT1p1LEO = 192
+
+    //Stubs
+    expect(partitionT1p0.partitionId).andStubReturn(partitionT1p0Id)
+    expect(partitionT1p1.partitionId).andStubReturn(partitionT1p1Id)
+
+    expect(replicaManager.getPartitionOrException(t1p0))
+      .andStubReturn(partitionT1p0)
+    expect(replicaManager.getPartitionOrException(t1p1))
+      .andStubReturn(partitionT1p1)
+    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(futureLogT1p0.logEndOffset).andReturn(futureReplicaLEO).anyTimes()
+    expect(futureLogT1p1.logEndOffset).andReturn(futureReplicaLEO).anyTimes()
+
+    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()
+        .setPartition(partitionT1p0Id)
+        .setErrorCode(Errors.NONE.code)
+        .setLeaderEpoch(leaderEpoch)
+        .setEndOffset(replicaT1p0LEO))
+      .anyTimes()
+
+    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()
+        .setPartition(partitionT1p1Id)
+        .setErrorCode(Errors.NONE.code)
+        .setLeaderEpoch(leaderEpoch)
+        .setEndOffset(replicaT1p1LEO))
+      .anyTimes()
+
+    expect(replicaManager.logManager).andReturn(logManager).anyTimes()
+    stubWithFetchMessages(logT1p0, logT1p1, futureLogT1p0, partitionT1p0, replicaManager, responseCallback)
+
+    replay(replicaManager, logManager, quotaManager, partitionT1p0, partitionT1p1, logT1p0, logT1p1, futureLogT1p0, futureLogT1p1)
+
+    //Create the thread
+    val endPoint = new BrokerEndPoint(0, "localhost", 1000)
+    val thread = new ReplicaAlterLogDirsThread(
+      "alter-logs-dirs-thread-test1",
+      sourceBroker = endPoint,
+      brokerConfig = config,
+      failedPartitions = failedPartitions,
+      replicaMgr = replicaManager,
+      quota = quotaManager,
+      brokerTopicStats = null)
+    thread.addPartitions(Map(t1p0 -> initialFetchState(0L), t1p1 -> initialFetchState(0L)))
+
+    //Run it
+    thread.doWork()
+
+    //We should have truncated to the offsets in the response
+    assertEquals(replicaT1p0LEO, truncateCaptureT1p0.getValue)
+    assertEquals(futureReplicaLEO, truncateCaptureT1p1.getValue)
+  }
+
+  @Test
+  def shouldTruncateToEndOffsetOfLargestCommonEpoch(): Unit = {
+
+    //Create a capture to track what partitions/offsets are truncated
+    val truncateToCapture: Capture[Long] = newCapture(CaptureType.ALL)
+
+    // Setup all the dependencies
+    val config = KafkaConfig.fromProps(TestUtils.createBrokerConfig(1, "localhost:1234"))
+    val quotaManager: ReplicationQuotaManager = createNiceMock(classOf[ReplicationQuotaManager])
+    val logManager: LogManager = createMock(classOf[LogManager])
+    val log: Log = createNiceMock(classOf[Log])
+    // one future replica mock because our mocking methods return same values for both future replicas
+    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 partitionId = 0
+    val leaderEpoch = 5
+    val futureReplicaLEO = 195
+    val replicaLEO = 200
+    val replicaEpochEndOffset = 190
+    val futureReplicaEpochEndOffset = 191
+
+    //Stubs
+    expect(partition.partitionId).andStubReturn(partitionId)
+
+    expect(replicaManager.getPartitionOrException(t1p0))
+      .andStubReturn(partition)
+    expect(replicaManager.futureLocalLogOrException(t1p0)).andStubReturn(futureLog)
+    expect(replicaManager.futureLogExists(t1p0)).andStubReturn(true)
+
+    expect(partition.truncateTo(capture(truncateToCapture), EasyMock.eq(true))).anyTimes()
+    expect(futureLog.logEndOffset).andReturn(futureReplicaLEO).anyTimes()
+    expect(futureLog.latestEpoch).andReturn(Some(leaderEpoch)).once()
+    expect(futureLog.latestEpoch).andReturn(Some(leaderEpoch - 2)).times(3)
+
+    // leader replica truncated and fetched new offsets with new leader epoch
+    expect(partition.lastOffsetForLeaderEpoch(Optional.of(1), leaderEpoch, fetchOnlyFromLeader = false))
+      .andReturn(new EpochEndOffset()
+        .setPartition(partitionId)
+        .setErrorCode(Errors.NONE.code)
+        .setLeaderEpoch(leaderEpoch - 1)
+        .setEndOffset(replicaLEO))
+      .anyTimes()
+    // but future replica does not know about this leader epoch, so returns a smaller leader epoch
+    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()
+        .setPartition(partitionId)
+        .setErrorCode(Errors.NONE.code)
+        .setLeaderEpoch(leaderEpoch - 2)
+        .setEndOffset(replicaEpochEndOffset))
+      .anyTimes()
+    expect(futureLog.endOffsetForEpoch(leaderEpoch - 2)).andReturn(
+      Some(OffsetAndEpoch(futureReplicaEpochEndOffset, leaderEpoch - 2))).anyTimes()
+
+    expect(replicaManager.logManager).andReturn(logManager).anyTimes()
+    stubWithFetchMessages(log, null, futureLog, partition, replicaManager, responseCallback)
+
+    replay(replicaManager, logManager, quotaManager, partition, log, futureLog)
+
+    //Create the thread
+    val endPoint = new BrokerEndPoint(0, "localhost", 1000)
+    val thread = new ReplicaAlterLogDirsThread(
+      "alter-logs-dirs-thread-test1",
+      sourceBroker = endPoint,
+      brokerConfig = config,
+      failedPartitions = failedPartitions,
+      replicaMgr = replicaManager,
+      quota = quotaManager,
+      brokerTopicStats = null)
+    thread.addPartitions(Map(t1p0 -> initialFetchState(0L)))
+
+    // First run will result in another offset for leader epoch request
+    thread.doWork()
+    // Second run should actually truncate
+    thread.doWork()
+
+    //We should have truncated to the offsets in the response
+    assertTrue("Expected offset " + replicaEpochEndOffset + " in captured truncation offsets " + truncateToCapture.getValues,
+               truncateToCapture.getValues.asScala.contains(replicaEpochEndOffset))
+  }
+
+  @Test
+  def shouldTruncateToInitialFetchOffsetIfReplicaReturnsUndefinedOffset(): Unit = {
+
+    //Create a capture to track what partitions/offsets are truncated
+    val truncated: Capture[Long] = newCapture(CaptureType.ALL)
+
+    // Setup all the dependencies
+    val config = KafkaConfig.fromProps(TestUtils.createBrokerConfig(1, "localhost:1234"))
+    val quotaManager: ReplicationQuotaManager = createNiceMock(classOf[ReplicationQuotaManager])
+    val logManager: LogManager = createMock(classOf[LogManager])
+    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
+
+    //Stubs
+    expect(replicaManager.getPartitionOrException(t1p0))
+      .andStubReturn(partition)
+    expect(partition.truncateTo(capture(truncated), isFuture = EasyMock.eq(true))).anyTimes()
+    expect(replicaManager.futureLocalLogOrException(t1p0)).andStubReturn(futureLog)
+    expect(replicaManager.futureLogExists(t1p0)).andStubReturn(true)
+
+    expect(replicaManager.logManager).andReturn(logManager).anyTimes()
+
+    // pretend this is a completely new future replica, with no leader epochs recorded
+    expect(futureLog.latestEpoch).andReturn(None).anyTimes()
+
+    stubWithFetchMessages(log, null, futureLog, partition, replicaManager, responseCallback)
+    replay(replicaManager, logManager, quotaManager, partition, log, futureLog)
+
+    //Create the thread
+    val endPoint = new BrokerEndPoint(0, "localhost", 1000)
+    val thread = new ReplicaAlterLogDirsThread(
+      "alter-logs-dirs-thread-test1",
+      sourceBroker = endPoint,
+      brokerConfig = config,
+      failedPartitions = failedPartitions,
+      replicaMgr = replicaManager,
+      quota = quotaManager,
+      brokerTopicStats = null)
+    thread.addPartitions(Map(t1p0 -> initialFetchState(initialFetchOffset)))
+
+    //Run it
+    thread.doWork()
+
+    //We should have truncated to initial fetch offset
+    assertEquals("Expected future replica to truncate to initial fetch offset if replica returns UNDEFINED_EPOCH_OFFSET",
+                 initialFetchOffset, truncated.getValue)
+  }
+
+  @Test
+  def shouldPollIndefinitelyIfReplicaNotAvailable(): Unit = {
+
+    //Create a capture to track what partitions/offsets are truncated
+    val truncated: Capture[Long] = newCapture(CaptureType.ALL)
+
+    // Setup all the dependencies
+    val config = KafkaConfig.fromProps(TestUtils.createBrokerConfig(1, "localhost:1234"))
+    val quotaManager: ReplicationQuotaManager = createNiceMock(classOf[ReplicationQuotaManager])
+    val logManager: LogManager = createMock(classOf[LogManager])
+    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 partitionId = 0
+    val futureReplicaLeaderEpoch = 1
+    val futureReplicaLEO = 290
+    val replicaLEO = 300
+
+    //Stubs
+    expect(partition.partitionId).andStubReturn(partitionId)
+
+    expect(replicaManager.getPartitionOrException(t1p0))
+      .andStubReturn(partition)
+    expect(partition.truncateTo(capture(truncated), isFuture = EasyMock.eq(true))).once()
+
+    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(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))
+      .andReturn(new EpochEndOffset()
+        .setPartition(partitionId)
+        .setErrorCode(Errors.REPLICA_NOT_AVAILABLE.code))
+      .times(3)
+      .andReturn(new EpochEndOffset()
+        .setPartition(partitionId)
+        .setErrorCode(Errors.NONE.code)
+        .setLeaderEpoch(futureReplicaLeaderEpoch)
+        .setEndOffset(replicaLEO))
+
+    expect(replicaManager.logManager).andReturn(logManager).anyTimes()
+    expect(replicaManager.fetchMessages(
+      EasyMock.anyLong(),
+      EasyMock.anyInt(),
+      EasyMock.anyInt(),
+      EasyMock.anyInt(),
+      EasyMock.anyObject(),
+      EasyMock.anyObject(),
+      EasyMock.anyObject(),
+      EasyMock.capture(responseCallback),
+      EasyMock.anyObject(),
+      EasyMock.anyObject())
+    ).andAnswer(() => responseCallback.getValue.apply(Seq.empty[(TopicPartition, FetchPartitionData)])).anyTimes()
+
+    replay(replicaManager, logManager, quotaManager, partition, log, futureLog)
+
+    //Create the thread
+    val endPoint = new BrokerEndPoint(0, "localhost", 1000)
+    val thread = new ReplicaAlterLogDirsThread(
+      "alter-logs-dirs-thread-test1",
+      sourceBroker = endPoint,
+      brokerConfig = config,
+      failedPartitions = failedPartitions,
+      replicaMgr = replicaManager,
+      quota = quotaManager,
+      brokerTopicStats = null)
+    thread.addPartitions(Map(t1p0 -> initialFetchState(0L)))
+
+    // Run thread 3 times (exactly number of times we mock exception for getReplicaOrException)
+    (0 to 2).foreach { _ =>
+      thread.doWork()
+    }
+
+    // Nothing happened since the replica was not available
+    assertEquals(0, truncated.getValues.size())
+
+    // Next time we loop, getReplicaOrException will return replica
+    thread.doWork()
+
+    // Now the final call should have actually done a truncation (to offset futureReplicaLEO)
+    assertEquals(futureReplicaLEO, truncated.getValue)
+  }
+
+  @Test
+  def shouldFetchLeaderEpochOnFirstFetchOnly(): Unit = {
+
+    //Setup all dependencies
+    val config = KafkaConfig.fromProps(TestUtils.createBrokerConfig(1, "localhost:1234"))
+    val quotaManager: ReplicationQuotaManager = createNiceMock(classOf[ReplicationQuotaManager])
+    val logManager: LogManager = createMock(classOf[LogManager])
+    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 partitionId = 0
+    val leaderEpoch = 5
+    val futureReplicaLEO = 190
+    val replicaLEO = 213
+
+    expect(partition.partitionId).andStubReturn(partitionId)
+
+    expect(replicaManager.getPartitionOrException(t1p0))
+        .andStubReturn(partition)
+    expect(partition.lastOffsetForLeaderEpoch(Optional.of(1), leaderEpoch, fetchOnlyFromLeader = false))
+        .andReturn(new EpochEndOffset()
+          .setPartition(partitionId)
+          .setErrorCode(Errors.NONE.code)
+          .setLeaderEpoch(leaderEpoch)
+          .setEndOffset(replicaLEO))
+    expect(partition.truncateTo(futureReplicaLEO, isFuture = true)).once()
+
+    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(log, null, futureLog, partition, replicaManager, responseCallback)
+
+    replay(replicaManager, logManager, quotaManager, partition, log, futureLog)
+
+    //Create the fetcher thread
+    val endPoint = new BrokerEndPoint(0, "localhost", 1000)
+    val thread = new ReplicaAlterLogDirsThread(
+      "alter-logs-dirs-thread-test1",
+      sourceBroker = endPoint,
+      brokerConfig = config,
+      failedPartitions = failedPartitions,
+      replicaMgr = replicaManager,
+      quota = quotaManager,
+      brokerTopicStats = null)
+    thread.addPartitions(Map(t1p0 -> initialFetchState(0L)))
+
+    // loop few times
+    (0 to 3).foreach { _ =>
+      thread.doWork()
+    }
+
+    //Assert that truncate to is called exactly once (despite more loops)
+    verify(partition)
+  }
+
+  @Test
+  def shouldFetchOneReplicaAtATime(): Unit = {
+
+    //Setup all dependencies
+    val config = KafkaConfig.fromProps(TestUtils.createBrokerConfig(1, "localhost:1234"))
+    val quotaManager: ReplicationQuotaManager = createNiceMock(classOf[ReplicationQuotaManager])
+    val logManager: LogManager = createMock(classOf[LogManager])
+    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(replicaManager.logManager).andReturn(logManager).anyTimes()
+    stub(log, null, futureLog, partition, replicaManager)
+
+    replay(replicaManager, logManager, quotaManager, partition, log)
+
+    //Create the fetcher thread
+    val endPoint = new BrokerEndPoint(0, "localhost", 1000)
+    val leaderEpoch = 1
+    val thread = new ReplicaAlterLogDirsThread(
+      "alter-logs-dirs-thread-test1",
+      sourceBroker = endPoint,
+      brokerConfig = config,
+      failedPartitions = failedPartitions,
+      replicaMgr = replicaManager,
+      quota = quotaManager,
+      brokerTopicStats = null)
+    thread.addPartitions(Map(
+      t1p0 -> initialFetchState(0L, leaderEpoch),
+      t1p1 -> initialFetchState(0L, leaderEpoch)))
+
+    val ResultWithPartitions(fetchRequestOpt, partitionsWithError) = thread.buildFetch(Map(
+      t1p0 -> PartitionFetchState(150, None, leaderEpoch, None, state = Fetching, lastFetchedEpoch = None),
+      t1p1 -> PartitionFetchState(160, None, leaderEpoch, None, state = Fetching, lastFetchedEpoch = None)))
+
+    assertTrue(fetchRequestOpt.isDefined)
+    val fetchRequest = fetchRequestOpt.get.fetchRequest
+    assertFalse(fetchRequest.fetchData.isEmpty)
+    assertFalse(partitionsWithError.nonEmpty)
+    val request = fetchRequest.build()
+    assertEquals(0, request.minBytes)
+    val fetchInfos = request.fetchData.asScala.toSeq
+    assertEquals(1, fetchInfos.length)
+    assertEquals("Expected fetch request for first partition", t1p0, fetchInfos.head._1)
+    assertEquals(150, fetchInfos.head._2.fetchOffset)
+  }
+
+  @Test
+  def shouldFetchNonDelayedAndNonTruncatingReplicas(): Unit = {
+
+    //Setup all dependencies
+    val config = KafkaConfig.fromProps(TestUtils.createBrokerConfig(1, "localhost:1234"))
+    val quotaManager: ReplicationQuotaManager = createNiceMock(classOf[ReplicationQuotaManager])
+    val logManager: LogManager = createMock(classOf[LogManager])
+    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
+    val startOffset = 123
+    expect(futureLog.logStartOffset).andReturn(startOffset).anyTimes()
+    expect(replicaManager.logManager).andReturn(logManager).anyTimes()
+    stub(log, null, futureLog, partition, replicaManager)
+
+    replay(replicaManager, logManager, quotaManager, partition, log, futureLog)
+
+    //Create the fetcher thread
+    val endPoint = new BrokerEndPoint(0, "localhost", 1000)
+    val leaderEpoch = 1
+    val thread = new ReplicaAlterLogDirsThread(
+      "alter-logs-dirs-thread-test1",
+      sourceBroker = endPoint,
+      brokerConfig = config,
+      failedPartitions = failedPartitions,
+      replicaMgr = replicaManager,
+      quota = quotaManager,
+      brokerTopicStats = null)
+    thread.addPartitions(Map(
+      t1p0 -> initialFetchState(0L, leaderEpoch),
+      t1p1 -> initialFetchState(0L, leaderEpoch)))
+
+    // one partition is ready and one is truncating
+    val ResultWithPartitions(fetchRequestOpt, partitionsWithError) = thread.buildFetch(Map(
+        t1p0 -> PartitionFetchState(150, None, leaderEpoch, state = Fetching, lastFetchedEpoch = None),
+        t1p1 -> PartitionFetchState(160, None, leaderEpoch, state = Truncating, lastFetchedEpoch = None)))
+
+    assertTrue(fetchRequestOpt.isDefined)
+    val fetchRequest = fetchRequestOpt.get
+    assertFalse(fetchRequest.partitionData.isEmpty)
+    assertFalse(partitionsWithError.nonEmpty)
+    val fetchInfos = fetchRequest.fetchRequest.build().fetchData.asScala.toSeq
+    assertEquals(1, fetchInfos.length)
+    assertEquals("Expected fetch request for non-truncating partition", t1p0, fetchInfos.head._1)
+    assertEquals(150, fetchInfos.head._2.fetchOffset)
+
+    // one partition is ready and one is delayed
+    val ResultWithPartitions(fetchRequest2Opt, partitionsWithError2) = thread.buildFetch(Map(
+        t1p0 -> PartitionFetchState(140, None, leaderEpoch, state = Fetching, lastFetchedEpoch = None),
+        t1p1 -> PartitionFetchState(160, None, leaderEpoch, delay = Some(new DelayedItem(5000)), state = Fetching, lastFetchedEpoch = None)))
+
+    assertTrue(fetchRequest2Opt.isDefined)
+    val fetchRequest2 = fetchRequest2Opt.get
+    assertFalse(fetchRequest2.partitionData.isEmpty)
+    assertFalse(partitionsWithError2.nonEmpty)
+    val fetchInfos2 = fetchRequest2.fetchRequest.build().fetchData.asScala.toSeq
+    assertEquals(1, fetchInfos2.length)
+    assertEquals("Expected fetch request for non-delayed partition", t1p0, fetchInfos2.head._1)
+    assertEquals(140, fetchInfos2.head._2.fetchOffset)
+
+    // both partitions are delayed
+    val ResultWithPartitions(fetchRequest3Opt, partitionsWithError3) = thread.buildFetch(Map(
+        t1p0 -> PartitionFetchState(140, None, leaderEpoch, delay = Some(new DelayedItem(5000)), state = Fetching, lastFetchedEpoch = None),
+        t1p1 -> PartitionFetchState(160, None, leaderEpoch, delay = Some(new DelayedItem(5000)), state = Fetching, lastFetchedEpoch = None)))
+    assertTrue("Expected no fetch requests since all partitions are delayed", fetchRequest3Opt.isEmpty)
+    assertFalse(partitionsWithError3.nonEmpty)
+  }
+
+  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(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(),
+      EasyMock.anyInt(),
+      EasyMock.anyInt(),
+      EasyMock.anyObject(),
+      EasyMock.anyObject(),
+      EasyMock.anyObject(),
+      EasyMock.capture(responseCallback),
+      EasyMock.anyObject(),
+      EasyMock.anyObject())
+    ).andAnswer(() => responseCallback.getValue.apply(Seq.empty[(TopicPartition, FetchPartitionData)])).anyTimes()
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/server/ReplicaFetchTest.scala b/core/src/test/scala/unit/kafka/server/ReplicaFetchTest.scala
index b8c45a5905315..a3a83bc8633f9 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,32 +33,31 @@ 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")
 
     // create a topic and partition and await leadership
     for (topic <- List(topic1,topic2)) {
-      createTopic(zkUtils, topic, numPartitions = 1, replicationFactor = 2, servers = brokers)
+      createTopic(zkClient, topic, numPartitions = 1, replicationFactor = 2, servers = brokers)
     }
 
     // send test messages to leader
-    val producer = TestUtils.createNewProducer(TestUtils.getBrokerListStrFromServers(brokers),
-                                               retries = 5,
+    val producer = TestUtils.createProducer(TestUtils.getBrokerListStrFromServers(brokers),
                                                keySerializer = new StringSerializer,
                                                valueSerializer = new StringSerializer)
     val records = testMessageList1.map(m => new ProducerRecord(topic1, m, m)) ++
@@ -68,9 +69,9 @@ class ReplicaFetchTest extends ZooKeeperTestHarness  {
       var result = true
       for (topic <- List(topic1, topic2)) {
         val tp = new TopicPartition(topic, partition)
-        val expectedOffset = brokers.head.getLogManager().getLog(tp).get.logEndOffset
+        val expectedOffset = brokers.head.getLogManager.getLog(tp).get.logEndOffset
         result = result && expectedOffset > 0 && brokers.forall { item =>
-          expectedOffset == item.getLogManager().getLog(tp).get.logEndOffset
+          expectedOffset == item.getLogManager.getLog(tp).get.logEndOffset
         }
       }
       result
diff --git a/core/src/test/scala/unit/kafka/server/ReplicaFetcherThreadTest.scala b/core/src/test/scala/unit/kafka/server/ReplicaFetcherThreadTest.scala
index 2074044d0e129..51370178c4cc6 100644
--- a/core/src/test/scala/unit/kafka/server/ReplicaFetcherThreadTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ReplicaFetcherThreadTest.scala
@@ -16,25 +16,31 @@
   */
 package kafka.server
 
-import kafka.cluster.{BrokerEndPoint, Replica}
-import kafka.log.LogManager
-import kafka.cluster.Partition
-import kafka.server.epoch.LeaderEpochCache
-import org.apache.kafka.common.requests.EpochEndOffset._
+import java.nio.charset.StandardCharsets
+import java.util.{Collections, Optional}
+
+import kafka.api.{ApiVersion, KAFKA_2_6_IV0}
+import kafka.cluster.{BrokerEndPoint, Partition}
+import kafka.log.{Log, LogAppendInfo, LogManager}
+import kafka.server.QuotaFactory.UnboundedQuota
 import kafka.server.epoch.util.ReplicaFetcherMockBlockingSend
 import kafka.utils.TestUtils
 import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.message.FetchResponseData
+import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData.EpochEndOffset
 import org.apache.kafka.common.metrics.Metrics
-import org.apache.kafka.common.protocol.Errors
 import org.apache.kafka.common.protocol.Errors._
-import org.apache.kafka.common.requests.EpochEndOffset
+import org.apache.kafka.common.protocol.{ApiKeys, Errors}
+import org.apache.kafka.common.record.{CompressionType, MemoryRecords, Records, SimpleRecord}
+import org.apache.kafka.common.requests.OffsetsForLeaderEpochResponse.{UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET}
+import org.apache.kafka.common.requests.{FetchResponse, OffsetsForLeaderEpochRequest}
 import org.apache.kafka.common.utils.SystemTime
 import org.easymock.EasyMock._
 import org.easymock.{Capture, CaptureType}
 import org.junit.Assert._
-import org.junit.Test
+import org.junit.{After, Test}
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 import scala.collection.{Map, mutable}
 
 class ReplicaFetcherThreadTest {
@@ -43,60 +49,173 @@ class ReplicaFetcherThreadTest {
   private val t1p1 = new TopicPartition("topic1", 1)
   private val t2p1 = new TopicPartition("topic2", 1)
 
+  private val brokerEndPoint = new BrokerEndPoint(0, "localhost", 1000)
+  private val failedPartitions = new FailedPartitions
+
+  private def initialFetchState(fetchOffset: Long, leaderEpoch: Int = 1): InitialFetchState = {
+    InitialFetchState(leader = new BrokerEndPoint(0, "localhost", 9092),
+      initOffset = fetchOffset, currentLeaderEpoch = leaderEpoch)
+  }
+
+  @After
+  def cleanup(): Unit = {
+    TestUtils.clearYammerMetrics()
+  }
+
   @Test
-  def shouldNotIssueLeaderEpochRequestIfInterbrokerVersionBelow11(): Unit = {
+  def shouldSendLatestRequestVersionsByDefault(): Unit = {
     val props = TestUtils.createBrokerConfig(1, "localhost:1234")
-    props.put(KafkaConfig.InterBrokerProtocolVersionProp, "0.10.2")
-    props.put(KafkaConfig.LogMessageFormatVersionProp, "0.10.2")
     val config = KafkaConfig.fromProps(props)
-    val endPoint = new BrokerEndPoint(0, "localhost", 1000)
+    val replicaManager: ReplicaManager = mock(classOf[ReplicaManager])
+    expect(replicaManager.brokerTopicStats).andReturn(mock(classOf[BrokerTopicStats]))
+    replay(replicaManager)
     val thread = new ReplicaFetcherThread(
       name = "bob",
       fetcherId = 0,
-      sourceBroker = endPoint,
+      sourceBroker = brokerEndPoint,
       brokerConfig = config,
-      replicaMgr = null,
+      failedPartitions: FailedPartitions,
+      replicaMgr = replicaManager,
       metrics =  new Metrics(),
       time = new SystemTime(),
-      quota = null,
+      quota = UnboundedQuota,
       leaderEndpointBlockingSend = None)
+    assertEquals(ApiKeys.FETCH.latestVersion, thread.fetchRequestVersion)
+    assertEquals(ApiKeys.OFFSET_FOR_LEADER_EPOCH.latestVersion, thread.offsetForLeaderEpochRequestVersion)
+    assertEquals(ApiKeys.LIST_OFFSETS.latestVersion, thread.listOffsetRequestVersion)
+  }
 
-    val result = thread.fetchEpochsFromLeader(Map(t1p0 -> 0, t1p1 -> 0))
+  @Test
+  def testFetchLeaderEpochRequestIfLastEpochDefinedForSomePartitions(): Unit = {
+    val config = kafkaConfigNoTruncateOnFetch
 
-    val expected = Map(
-      t1p0 -> new EpochEndOffset(Errors.NONE, UNDEFINED_EPOCH_OFFSET),
-      t1p1 -> new EpochEndOffset(Errors.NONE, UNDEFINED_EPOCH_OFFSET)
-    )
+    //Setup all dependencies
+    val quota: ReplicationQuotaManager = createNiceMock(classOf[ReplicationQuotaManager])
+    val logManager: LogManager = createMock(classOf[LogManager])
+    val replicaAlterLogDirsManager: ReplicaAlterLogDirsManager = createMock(classOf[ReplicaAlterLogDirsManager])
+    val log: Log = createNiceMock(classOf[Log])
+    val partition: Partition = createMock(classOf[Partition])
+    val replicaManager: ReplicaManager = createMock(classOf[ReplicaManager])
 
-    assertEquals("results from leader epoch request should have undefined offset", expected, result)
+    val leaderEpoch = 5
+
+    //Stubs
+    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()
+    expect(replicaManager.brokerTopicStats).andReturn(mock(classOf[BrokerTopicStats]))
+    stub(partition, replicaManager, log)
+
+    //Expectations
+    expect(partition.truncateTo(anyLong(), anyBoolean())).times(3)
+
+    replay(replicaManager, logManager, quota, partition, log)
+
+    //Define the offsets for the OffsetsForLeaderEpochResponse
+    val offsets = Map(
+      t1p0 -> newOffsetForLeaderPartitionResult(t1p0, leaderEpoch, 1),
+      t1p1 -> newOffsetForLeaderPartitionResult(t1p1, leaderEpoch, 1)).asJava
+
+    //Create the fetcher thread
+    val mockNetwork = new ReplicaFetcherMockBlockingSend(offsets, brokerEndPoint, new SystemTime())
+
+    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(
+      t1p0 -> initialFetchState(0L),
+      t1p1 -> initialFetchState(0L),
+      t2p1 -> initialFetchState(0L)))
+
+    assertPartitionStates(thread, shouldBeReadyForFetch = false, shouldBeTruncatingLog = true, shouldBeDelayed = false)
+    //Loop 1
+    thread.doWork()
+    assertEquals(1, mockNetwork.epochFetchCount)
+    assertEquals(1, mockNetwork.fetchCount)
+
+    assertPartitionStates(thread, shouldBeReadyForFetch = true, shouldBeTruncatingLog = false, shouldBeDelayed = false)
+
+    //Loop 2 we should not fetch epochs
+    thread.doWork()
+    assertEquals(1, mockNetwork.epochFetchCount)
+    assertEquals(2, mockNetwork.fetchCount)
+
+    assertPartitionStates(thread, shouldBeReadyForFetch = true, shouldBeTruncatingLog = false, shouldBeDelayed = false)
+
+    //Loop 3 we should not fetch epochs
+    thread.doWork()
+    assertEquals(1, mockNetwork.epochFetchCount)
+    assertEquals(3, mockNetwork.fetchCount)
+
+    assertPartitionStates(thread, shouldBeReadyForFetch = true, shouldBeTruncatingLog = false, shouldBeDelayed = false)
+
+    //Assert that truncate to is called exactly once (despite two loops)
+    verify(logManager)
+  }
+
+  /**
+    * Assert that all partitions' states are as expected
+    *
+    */
+  def assertPartitionStates(fetcher: AbstractFetcherThread,
+                            shouldBeReadyForFetch: Boolean,
+                            shouldBeTruncatingLog: Boolean,
+                            shouldBeDelayed: Boolean): Unit = {
+    for (tp <- List(t1p0, t1p1, t2p1)) {
+      assertTrue(fetcher.fetchState(tp).isDefined)
+      val fetchState = fetcher.fetchState(tp).get
+
+      assertEquals(
+        s"Partition $tp should${if (!shouldBeReadyForFetch) " NOT" else ""} be ready for fetching",
+        shouldBeReadyForFetch, fetchState.isReadyForFetch)
+
+      assertEquals(
+        s"Partition $tp should${if (!shouldBeTruncatingLog) " NOT" else ""} be truncating its log",
+        shouldBeTruncatingLog, fetchState.isTruncating)
+
+      assertEquals(
+        s"Partition $tp should${if (!shouldBeDelayed) " NOT" else ""} be delayed",
+        shouldBeDelayed, fetchState.isDelayed)
+    }
   }
 
   @Test
   def shouldHandleExceptionFromBlockingSend(): Unit = {
     val props = TestUtils.createBrokerConfig(1, "localhost:1234")
     val config = KafkaConfig.fromProps(props)
-    val endPoint = new BrokerEndPoint(0, "localhost", 1000)
-    val mockBlockingSend = createMock(classOf[BlockingSend])
+    val mockBlockingSend: BlockingSend = createMock(classOf[BlockingSend])
 
     expect(mockBlockingSend.sendRequest(anyObject())).andThrow(new NullPointerException).once()
-    replay(mockBlockingSend)
+    val replicaManager: ReplicaManager = mock(classOf[ReplicaManager])
+    expect(replicaManager.brokerTopicStats).andReturn(mock(classOf[BrokerTopicStats]))
+    replay(mockBlockingSend, replicaManager)
 
     val thread = new ReplicaFetcherThread(
       name = "bob",
       fetcherId = 0,
-      sourceBroker = endPoint,
+      sourceBroker = brokerEndPoint,
       brokerConfig = config,
-      replicaMgr = null,
+      failedPartitions: FailedPartitions,
+      replicaMgr = replicaManager,
       metrics =  new Metrics(),
       time = new SystemTime(),
       quota = null,
       leaderEndpointBlockingSend = Some(mockBlockingSend))
 
-    val result = thread.fetchEpochsFromLeader(Map(t1p0 -> 0, t1p1 -> 0))
+    val result = thread.fetchEpochEndOffsets(Map(
+      t1p0 -> new OffsetsForLeaderEpochRequest.PartitionData(Optional.empty(), 0),
+      t1p1 -> new OffsetsForLeaderEpochRequest.PartitionData(Optional.empty(), 0)))
 
     val expected = Map(
-      t1p0 -> new EpochEndOffset(Errors.UNKNOWN_SERVER_ERROR, UNDEFINED_EPOCH_OFFSET),
-      t1p1 -> new EpochEndOffset(Errors.UNKNOWN_SERVER_ERROR, UNDEFINED_EPOCH_OFFSET)
+      t1p0 -> newOffsetForLeaderPartitionResult(t1p0, Errors.UNKNOWN_SERVER_ERROR, UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET),
+      t1p1 -> newOffsetForLeaderPartitionResult(t1p1, Errors.UNKNOWN_SERVER_ERROR, UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET)
     )
 
     assertEquals("results from leader epoch request should have undefined offset", expected, result)
@@ -104,54 +223,69 @@ class ReplicaFetcherThreadTest {
   }
 
   @Test
-  def shouldFetchLeaderEpochOnFirstFetchOnly(): Unit = {
-    val config = KafkaConfig.fromProps(TestUtils.createBrokerConfig(1, "localhost:1234"))
+  def shouldFetchLeaderEpochOnFirstFetchOnlyIfLeaderEpochKnownToBothIbp26(): Unit = {
+    verifyFetchLeaderEpochOnFirstFetch(KAFKA_2_6_IV0)
+  }
+
+  @Test
+  def shouldNotFetchLeaderEpochOnFirstFetchWithTruncateOnFetch(): Unit = {
+    verifyFetchLeaderEpochOnFirstFetch(ApiVersion.latestVersion, epochFetchCount = 0)
+  }
+
+  private def verifyFetchLeaderEpochOnFirstFetch(ibp: ApiVersion, epochFetchCount: Int = 1): Unit = {
+    val props = TestUtils.createBrokerConfig(1, "localhost:1234")
+    props.setProperty(KafkaConfig.InterBrokerProtocolVersionProp, ibp.version)
+    val config = KafkaConfig.fromProps(props)
 
     //Setup all dependencies
-    val quota = createNiceMock(classOf[ReplicationQuotaManager])
-    val leaderEpochs = createNiceMock(classOf[LeaderEpochCache])
-    val logManager = createMock(classOf[LogManager])
-    val replicaAlterLogDirsManager = createMock(classOf[ReplicaAlterLogDirsManager])
-    val replica = createNiceMock(classOf[Replica])
-    val partition = createMock(classOf[Partition])
-    val replicaManager = createMock(classOf[ReplicaManager])
+    val logManager: LogManager = createMock(classOf[LogManager])
+    val replicaAlterLogDirsManager: ReplicaAlterLogDirsManager = createMock(classOf[ReplicaAlterLogDirsManager])
+    val log: Log = createNiceMock(classOf[Log])
+    val partition: Partition = createMock(classOf[Partition])
+    val replicaManager: ReplicaManager = createMock(classOf[ReplicaManager])
+
+    val leaderEpoch = 5
 
     //Stubs
-    expect(replica.epochs).andReturn(Some(leaderEpochs)).anyTimes()
-    expect(replica.logEndOffset).andReturn(new LogOffsetMetadata(0)).anyTimes()
-    expect(leaderEpochs.latestEpoch).andReturn(5)
+    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)
-
+    expect(replicaManager.brokerTopicStats).andReturn(mock(classOf[BrokerTopicStats]))
+    stub(partition, replicaManager, log)
 
     //Expectations
-    expect(partition.truncateTo(anyLong(), anyBoolean())).once
+    expect(partition.truncateTo(anyLong(), anyBoolean())).times(2)
 
-    replay(leaderEpochs, replicaManager, logManager, quota, replica)
+    replay(replicaManager, logManager, partition, log)
 
     //Define the offsets for the OffsetsForLeaderEpochResponse
-    val offsets = Map(t1p0 -> new EpochEndOffset(1), t1p1 -> new EpochEndOffset(1)).asJava
+    val offsets = Map(
+      t1p0 -> newOffsetForLeaderPartitionResult(t1p0, leaderEpoch, 1),
+      t1p1 -> newOffsetForLeaderPartitionResult(t1p1, leaderEpoch, 1)).asJava
 
     //Create the fetcher thread
-    val endPoint = new BrokerEndPoint(0, "localhost", 1000)
-    val mockNetwork = new ReplicaFetcherMockBlockingSend(offsets, endPoint, new SystemTime())
-    val thread = new ReplicaFetcherThread("bob", 0, endPoint, config, replicaManager, new Metrics(), new SystemTime(), quota, Some(mockNetwork))
-    thread.addPartitions(Map(t1p0 -> 0, t1p1 -> 0))
+    val mockNetwork = new ReplicaFetcherMockBlockingSend(offsets, brokerEndPoint, new SystemTime())
+    val thread = new ReplicaFetcherThread("bob", 0, brokerEndPoint, config, failedPartitions, replicaManager,
+      new Metrics, new SystemTime, UnboundedQuota, Some(mockNetwork))
+    thread.addPartitions(Map(t1p0 -> initialFetchState(0L), t1p1 -> initialFetchState(0L)))
 
     //Loop 1
     thread.doWork()
-    assertEquals(1, mockNetwork.epochFetchCount)
+    assertEquals(epochFetchCount, mockNetwork.epochFetchCount)
     assertEquals(1, mockNetwork.fetchCount)
 
     //Loop 2 we should not fetch epochs
     thread.doWork()
-    assertEquals(1, mockNetwork.epochFetchCount)
+    assertEquals(epochFetchCount, mockNetwork.epochFetchCount)
     assertEquals(2, mockNetwork.fetchCount)
 
     //Loop 3 we should not fetch epochs
     thread.doWork()
-    assertEquals(1, mockNetwork.epochFetchCount)
+    assertEquals(epochFetchCount, mockNetwork.epochFetchCount)
     assertEquals(3, mockNetwork.fetchCount)
 
     //Assert that truncate to is called exactly once (despite two loops)
@@ -165,43 +299,353 @@ class ReplicaFetcherThreadTest {
     val truncateToCapture: Capture[Long] = newCapture(CaptureType.ALL)
 
     // Setup all the dependencies
-    val configs = TestUtils.createBrokerConfigs(1, "localhost:1234").map(KafkaConfig.fromProps)
-    val quota = createNiceMock(classOf[ReplicationQuotaManager])
-    val leaderEpochs = createMock(classOf[LeaderEpochCache])
-    val logManager = createMock(classOf[LogManager])
-    val replicaAlterLogDirsManager = createMock(classOf[ReplicaAlterLogDirsManager])
-    val replica = createNiceMock(classOf[Replica])
-    val partition = createMock(classOf[Partition])
-    val replicaManager = createMock(classOf[ReplicaManager])
+    val config = kafkaConfigNoTruncateOnFetch
+    val quota: ReplicationQuotaManager = createNiceMock(classOf[ReplicationQuotaManager])
+    val logManager: LogManager = createMock(classOf[LogManager])
+    val replicaAlterLogDirsManager: ReplicaAlterLogDirsManager = createMock(classOf[ReplicaAlterLogDirsManager])
+    val log: Log = createNiceMock(classOf[Log])
+    val partition: Partition = createMock(classOf[Partition])
+    val replicaManager: ReplicaManager = createMock(classOf[ReplicaManager])
+
+    val leaderEpoch = 5
+    val initialLEO = 200
+
+    //Stubs
+    expect(partition.truncateTo(capture(truncateToCapture), anyBoolean())).anyTimes()
+    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()
+    expect(replicaManager.brokerTopicStats).andReturn(mock(classOf[BrokerTopicStats]))
+    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 -> newOffsetForLeaderPartitionResult(t1p0, leaderEpoch, 156),
+      t2p1 -> newOffsetForLeaderPartitionResult(t2p1, leaderEpoch, 172)).asJava
+
+    //Create the thread
+    val mockNetwork = new ReplicaFetcherMockBlockingSend(offsetsReply, brokerEndPoint, new SystemTime())
+    val thread = new ReplicaFetcherThread("bob", 0, brokerEndPoint, config, failedPartitions, replicaManager,
+      new Metrics(), new SystemTime(), quota, Some(mockNetwork))
+    thread.addPartitions(Map(t1p0 -> initialFetchState(0L), t2p1 -> initialFetchState(0L)))
+
+    //Run it
+    thread.doWork()
 
+    //We should have truncated to the offsets in the response
+    assertTrue("Expected " + t1p0 + " to truncate to offset 156 (truncation offsets: " + truncateToCapture.getValues + ")",
+               truncateToCapture.getValues.asScala.contains(156))
+    assertTrue("Expected " + t2p1 + " to truncate to offset 172 (truncation offsets: " + truncateToCapture.getValues + ")",
+               truncateToCapture.getValues.asScala.contains(172))
+  }
+
+  @Test
+  def shouldTruncateToOffsetSpecifiedInEpochOffsetResponseIfFollowerHasNoMoreEpochs(): Unit = {
+    // Create a capture to track what partitions/offsets are truncated
+    val truncateToCapture: Capture[Long] = newCapture(CaptureType.ALL)
+
+    // Setup all the dependencies
+    val config = kafkaConfigNoTruncateOnFetch
+    val quota: ReplicationQuotaManager = createNiceMock(classOf[ReplicationQuotaManager])
+    val logManager: LogManager = createMock(classOf[LogManager])
+    val replicaAlterLogDirsManager: ReplicaAlterLogDirsManager = createMock(classOf[ReplicaAlterLogDirsManager])
+    val log: Log = createNiceMock(classOf[Log])
+    val partition: Partition = createMock(classOf[Partition])
+    val replicaManager: ReplicaManager = createMock(classOf[ReplicaManager])
+
+    val leaderEpochAtFollower = 5
+    val leaderEpochAtLeader = 4
     val initialLEO = 200
 
     //Stubs
     expect(partition.truncateTo(capture(truncateToCapture), anyBoolean())).anyTimes()
-    expect(replica.epochs).andReturn(Some(leaderEpochs)).anyTimes()
-    expect(replica.logEndOffset).andReturn(new LogOffsetMetadata(initialLEO)).anyTimes()
-    expect(leaderEpochs.latestEpoch).andReturn(5).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)
+    expect(replicaManager.brokerTopicStats).andReturn(mock(classOf[BrokerTopicStats]))
+    stub(partition, replicaManager, log)
 
-    replay(leaderEpochs, 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(156), t2p1 -> new EpochEndOffset(172)).asJava
+    val offsetsReply = Map(
+      t1p0 -> newOffsetForLeaderPartitionResult(t1p0, leaderEpochAtLeader, 156),
+      t2p1 -> newOffsetForLeaderPartitionResult(t2p1, leaderEpochAtLeader, 202)).asJava
 
     //Create the thread
-    val endPoint = new BrokerEndPoint(0, "localhost", 1000)
-    val mockNetwork = new ReplicaFetcherMockBlockingSend(offsetsReply, endPoint, new SystemTime())
-    val thread = new ReplicaFetcherThread("bob", 0, endPoint, configs(0), replicaManager, new Metrics(), new SystemTime(), quota, Some(mockNetwork))
-    thread.addPartitions(Map(t1p0 -> 0, t2p1 -> 0))
+    val mockNetwork = new ReplicaFetcherMockBlockingSend(offsetsReply, brokerEndPoint, new SystemTime())
+    val thread = new ReplicaFetcherThread("bob", 0, brokerEndPoint, config, failedPartitions,
+      replicaManager, new Metrics(), new SystemTime(), quota, Some(mockNetwork))
+    thread.addPartitions(Map(t1p0 -> initialFetchState(0L), t2p1 -> initialFetchState(0L)))
 
     //Run it
     thread.doWork()
 
     //We should have truncated to the offsets in the response
-    assertTrue(truncateToCapture.getValues.asScala.contains(156))
-    assertTrue(truncateToCapture.getValues.asScala.contains(172))
+    assertTrue("Expected " + t1p0 + " to truncate to offset 156 (truncation offsets: " + truncateToCapture.getValues + ")",
+               truncateToCapture.getValues.asScala.contains(156))
+    assertTrue("Expected " + t2p1 + " to truncate to offset " + initialLEO +
+               " (truncation offsets: " + truncateToCapture.getValues + ")",
+               truncateToCapture.getValues.asScala.contains(initialLEO))
+  }
+
+  @Test
+  def shouldFetchLeaderEpochSecondTimeIfLeaderRepliesWithEpochNotKnownToFollower(): Unit = {
+    // Create a capture to track what partitions/offsets are truncated
+    val truncateToCapture: Capture[Long] = newCapture(CaptureType.ALL)
+
+    val config = kafkaConfigNoTruncateOnFetch
+
+    // Setup all dependencies
+    val quota: ReplicationQuotaManager = createNiceMock(classOf[ReplicationQuotaManager])
+    val logManager: LogManager = createMock(classOf[LogManager])
+    val replicaAlterLogDirsManager: ReplicaAlterLogDirsManager = createMock(classOf[ReplicaAlterLogDirsManager])
+    val log: Log = createNiceMock(classOf[Log])
+    val partition: Partition = createMock(classOf[Partition])
+    val replicaManager: ReplicaManager = createMock(classOf[ReplicaManager])
+
+    val initialLEO = 200
+
+    // Stubs
+    expect(partition.truncateTo(capture(truncateToCapture), anyBoolean())).anyTimes()
+    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(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()
+    expect(replicaManager.brokerTopicStats).andReturn(mock(classOf[BrokerTopicStats]))
+    stub(partition, replicaManager, log)
+
+    replay(replicaManager, logManager, quota, partition, log)
+
+    // Define the offsets for the OffsetsForLeaderEpochResponse
+    val offsets = Map(
+      t1p0 -> newOffsetForLeaderPartitionResult(t1p0, 4, 155),
+      t1p1 -> newOffsetForLeaderPartitionResult(t1p1, 4, 143)).asJava
+
+    // Create the fetcher thread
+    val mockNetwork = new ReplicaFetcherMockBlockingSend(offsets, brokerEndPoint, new SystemTime())
+    val thread = new ReplicaFetcherThread("bob", 0, brokerEndPoint, config, failedPartitions, replicaManager, new Metrics(), new SystemTime(), quota, Some(mockNetwork))
+    thread.addPartitions(Map(t1p0 -> initialFetchState(0L), t1p1 -> initialFetchState(0L)))
+
+    // Loop 1 -- both topic partitions will need to fetch another leader epoch
+    thread.doWork()
+    assertEquals(1, mockNetwork.epochFetchCount)
+    assertEquals(0, mockNetwork.fetchCount)
+
+    // Loop 2 should do the second fetch for both topic partitions because the leader replied with
+    // epoch 4 while follower knows only about epoch 3
+    val nextOffsets = Map(
+      t1p0 -> newOffsetForLeaderPartitionResult(t1p0, 3, 101),
+      t1p1 -> newOffsetForLeaderPartitionResult(t1p1, 3, 102)).asJava
+    mockNetwork.setOffsetsForNextResponse(nextOffsets)
+
+    thread.doWork()
+    assertEquals(2, mockNetwork.epochFetchCount)
+    assertEquals(1, mockNetwork.fetchCount)
+    assertTrue("OffsetsForLeaderEpochRequest version.",
+      mockNetwork.lastUsedOffsetForLeaderEpochVersion >= 3)
+
+    //Loop 3 we should not fetch epochs
+    thread.doWork()
+    assertEquals(2, mockNetwork.epochFetchCount)
+    assertEquals(2, mockNetwork.fetchCount)
+
+
+    //We should have truncated to the offsets in the second response
+    assertTrue("Expected " + t1p1 + " to truncate to offset 102 (truncation offsets: " + truncateToCapture.getValues + ")",
+               truncateToCapture.getValues.asScala.contains(102))
+    assertTrue("Expected " + t1p0 + " to truncate to offset 101 (truncation offsets: " + truncateToCapture.getValues + ")",
+               truncateToCapture.getValues.asScala.contains(101))
+  }
+
+  @Test
+  def shouldTruncateIfLeaderRepliesWithDivergingEpochNotKnownToFollower(): Unit = {
+
+    // Create a capture to track what partitions/offsets are truncated
+    val truncateToCapture: Capture[Long] = newCapture(CaptureType.ALL)
+
+    val config = KafkaConfig.fromProps(TestUtils.createBrokerConfig(1, "localhost:1234"))
+
+    // Setup all dependencies
+    val quota: ReplicationQuotaManager = createNiceMock(classOf[ReplicationQuotaManager])
+    val logManager: LogManager = createMock(classOf[LogManager])
+    val replicaAlterLogDirsManager: ReplicaAlterLogDirsManager = createMock(classOf[ReplicaAlterLogDirsManager])
+    val log: Log = createNiceMock(classOf[Log])
+    val partition: Partition = createNiceMock(classOf[Partition])
+    val replicaManager: ReplicaManager = createMock(classOf[ReplicaManager])
+
+    val initialLEO = 200
+    var latestLogEpoch: Option[Int] = Some(5)
+
+    // Stubs
+    expect(partition.truncateTo(capture(truncateToCapture), anyBoolean())).anyTimes()
+    expect(partition.localLogOrException).andReturn(log).anyTimes()
+    expect(log.highWatermark).andReturn(115).anyTimes()
+    expect(log.latestEpoch).andAnswer(() => latestLogEpoch).anyTimes()
+    expect(log.endOffsetForEpoch(4)).andReturn(Some(OffsetAndEpoch(149, 4))).anyTimes()
+    expect(log.endOffsetForEpoch(3)).andReturn(Some(OffsetAndEpoch(129, 2))).anyTimes()
+    expect(log.endOffsetForEpoch(2)).andReturn(Some(OffsetAndEpoch(119, 1))).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()
+    expect(replicaManager.brokerTopicStats).andReturn(mock(classOf[BrokerTopicStats]))
+    stub(partition, replicaManager, log)
+
+    replay(replicaManager, logManager, quota, partition, log)
+
+    // Create the fetcher thread
+    val mockNetwork = new ReplicaFetcherMockBlockingSend(Collections.emptyMap(), brokerEndPoint, new SystemTime())
+    val thread = new ReplicaFetcherThread("bob", 0, brokerEndPoint, config, failedPartitions, replicaManager, new Metrics(), new SystemTime(), quota, Some(mockNetwork)) {
+      override def processPartitionData(topicPartition: TopicPartition, fetchOffset: Long, partitionData: FetchData): Option[LogAppendInfo] = None
+    }
+    thread.addPartitions(Map(t1p0 -> initialFetchState(initialLEO), t1p1 -> initialFetchState(initialLEO)))
+    val partitions = Set(t1p0, t1p1)
+
+    // Loop 1 -- both topic partitions skip epoch fetch and send fetch request since we can truncate
+    // later based on diverging epochs in fetch response.
+    thread.doWork()
+    assertEquals(0, mockNetwork.epochFetchCount)
+    assertEquals(1, mockNetwork.fetchCount)
+    partitions.foreach { tp => assertEquals(Fetching, thread.fetchState(tp).get.state) }
+
+    def partitionData(divergingEpoch: FetchResponseData.EpochEndOffset): FetchResponse.PartitionData[Records] = {
+      new FetchResponse.PartitionData[Records](
+        Errors.NONE, 0, 0, 0, Optional.empty(), Collections.emptyList(),
+        Optional.of(divergingEpoch), MemoryRecords.EMPTY)
+    }
+
+    // Loop 2 should truncate based on diverging epoch and continue to send fetch requests.
+    mockNetwork.setFetchPartitionDataForNextResponse(Map(
+      t1p0 -> partitionData(new FetchResponseData.EpochEndOffset().setEpoch(4).setEndOffset(140)),
+      t1p1 -> partitionData(new FetchResponseData.EpochEndOffset().setEpoch(4).setEndOffset(141))
+    ))
+    latestLogEpoch = Some(4)
+    thread.doWork()
+    assertEquals(0, mockNetwork.epochFetchCount)
+    assertEquals(2, mockNetwork.fetchCount)
+    assertTrue("Expected " + t1p0 + " to truncate to offset 140 (truncation offsets: " + truncateToCapture.getValues + ")",
+      truncateToCapture.getValues.asScala.contains(140))
+    assertTrue("Expected " + t1p1 + " to truncate to offset 141 (truncation offsets: " + truncateToCapture.getValues + ")",
+      truncateToCapture.getValues.asScala.contains(141))
+    partitions.foreach { tp => assertEquals(Fetching, thread.fetchState(tp).get.state) }
+
+    // Loop 3 should truncate because of diverging epoch. Offset truncation is not complete
+    // because divergent epoch is not known to follower. We truncate and stay in Fetching state.
+    mockNetwork.setFetchPartitionDataForNextResponse(Map(
+      t1p0 -> partitionData(new FetchResponseData.EpochEndOffset().setEpoch(3).setEndOffset(130)),
+      t1p1 -> partitionData(new FetchResponseData.EpochEndOffset().setEpoch(3).setEndOffset(131))
+    ))
+    thread.doWork()
+    assertEquals(0, mockNetwork.epochFetchCount)
+    assertEquals(3, mockNetwork.fetchCount)
+    assertTrue("Expected to truncate to offset 129 (truncation offsets: " + truncateToCapture.getValues + ")",
+      truncateToCapture.getValues.asScala.contains(129))
+    partitions.foreach { tp => assertEquals(Fetching, thread.fetchState(tp).get.state) }
+
+    // Loop 4 should truncate because of diverging epoch. Offset truncation is not complete
+    // because divergent epoch is not known to follower. Last fetched epoch cannot be determined
+    // from the log. We truncate and stay in Fetching state.
+    mockNetwork.setFetchPartitionDataForNextResponse(Map(
+      t1p0 -> partitionData(new FetchResponseData.EpochEndOffset().setEpoch(2).setEndOffset(120)),
+      t1p1 -> partitionData(new FetchResponseData.EpochEndOffset().setEpoch(2).setEndOffset(121))
+    ))
+    latestLogEpoch = None
+    thread.doWork()
+    assertEquals(0, mockNetwork.epochFetchCount)
+    assertEquals(4, mockNetwork.fetchCount)
+    assertTrue("Expected to truncate to offset 119 (truncation offsets: " + truncateToCapture.getValues + ")",
+      truncateToCapture.getValues.asScala.contains(119))
+    partitions.foreach { tp => assertEquals(Fetching, thread.fetchState(tp).get.state) }
+  }
+
+  @Test
+  def shouldUseLeaderEndOffsetIfInterBrokerVersionBelow20(): Unit = {
+
+    // Create a capture to track what partitions/offsets are truncated
+    val truncateToCapture: Capture[Long] = newCapture(CaptureType.ALL)
+
+    val props = TestUtils.createBrokerConfig(1, "localhost:1234")
+    props.put(KafkaConfig.InterBrokerProtocolVersionProp, "0.11.0")
+    val config = KafkaConfig.fromProps(props)
+
+    // Setup all dependencies
+    val quota: ReplicationQuotaManager = createNiceMock(classOf[ReplicationQuotaManager])
+    val logManager: LogManager = createMock(classOf[LogManager])
+    val replicaAlterLogDirsManager: ReplicaAlterLogDirsManager = createMock(classOf[ReplicaAlterLogDirsManager])
+    val log: Log = createNiceMock(classOf[Log])
+    val partition: Partition = createMock(classOf[Partition])
+    val replicaManager: ReplicaManager = createMock(classOf[ReplicaManager])
+
+    val initialLEO = 200
+
+    // Stubs
+    expect(partition.truncateTo(capture(truncateToCapture), anyBoolean())).anyTimes()
+    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(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()
+    expect(replicaManager.brokerTopicStats).andReturn(mock(classOf[BrokerTopicStats]))
+    stub(partition, replicaManager, log)
+
+    replay(replicaManager, logManager, quota, partition, log)
+
+    // Define the offsets for the OffsetsForLeaderEpochResponse with undefined epoch to simulate
+    // older protocol version
+    val offsets = Map(
+      t1p0 -> newOffsetForLeaderPartitionResult(t1p0, UNDEFINED_EPOCH, 155),
+      t1p1 -> newOffsetForLeaderPartitionResult(t1p1, UNDEFINED_EPOCH, 143)).asJava
+
+    // Create the fetcher thread
+    val mockNetwork = new ReplicaFetcherMockBlockingSend(offsets, brokerEndPoint, new SystemTime())
+    val thread = new ReplicaFetcherThread("bob", 0, brokerEndPoint, config, failedPartitions, replicaManager, new Metrics(), new SystemTime(), quota, Some(mockNetwork))
+    thread.addPartitions(Map(t1p0 -> initialFetchState(0L), t1p1 -> initialFetchState(0L)))
+
+    // Loop 1 -- both topic partitions will truncate to leader offset even though they don't know
+    // about leader epoch
+    thread.doWork()
+    assertEquals(1, mockNetwork.epochFetchCount)
+    assertEquals(1, mockNetwork.fetchCount)
+    assertEquals("OffsetsForLeaderEpochRequest version.",
+                 0, mockNetwork.lastUsedOffsetForLeaderEpochVersion)
+
+    //Loop 2 we should not fetch epochs
+    thread.doWork()
+    assertEquals(1, mockNetwork.epochFetchCount)
+    assertEquals(2, mockNetwork.fetchCount)
+
+    //We should have truncated to the offsets in the first response
+    assertTrue("Expected " + t1p0 + " to truncate to offset 155 (truncation offsets: " + truncateToCapture.getValues + ")",
+               truncateToCapture.getValues.asScala.contains(155))
+    assertTrue("Expected " + t1p1 + " to truncate to offset 143 (truncation offsets: " + truncateToCapture.getValues + ")",
+               truncateToCapture.getValues.asScala.contains(143))
   }
 
   @Test
@@ -211,41 +655,40 @@ class ReplicaFetcherThreadTest {
     val truncated: Capture[Long] = newCapture(CaptureType.ALL)
 
     // Setup all the dependencies
-    val configs = TestUtils.createBrokerConfigs(1, "localhost:1234").map(KafkaConfig.fromProps)
-    val quota = createNiceMock(classOf[ReplicationQuotaManager])
-    val leaderEpochs = createNiceMock(classOf[LeaderEpochCache])
-    val logManager = createMock(classOf[LogManager])
-    val replicaAlterLogDirsManager = createMock(classOf[ReplicaAlterLogDirsManager])
-    val replica = createNiceMock(classOf[Replica])
-    val partition = createMock(classOf[Partition])
-    val replicaManager = createMock(classOf[ReplicaManager])
+    val config = kafkaConfigNoTruncateOnFetch
+    val quota: ReplicationQuotaManager = createNiceMock(classOf[ReplicationQuotaManager])
+    val logManager: LogManager = createMock(classOf[LogManager])
+    val replicaAlterLogDirsManager: ReplicaAlterLogDirsManager = createMock(classOf[ReplicaAlterLogDirsManager])
+    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.epochs).andReturn(Some(leaderEpochs)).anyTimes()
-    expect(replica.logEndOffset).andReturn(new LogOffsetMetadata(initialLeo)).anyTimes()
-    expect(leaderEpochs.latestEpoch).andReturn(5)
+    expect(partition.localLogOrException).andReturn(log).anyTimes()
+    expect(log.highWatermark).andReturn(initialFetchOffset).anyTimes()
+    expect(log.latestEpoch).andReturn(Some(5)).times(2)
     expect(replicaManager.logManager).andReturn(logManager).anyTimes()
     expect(replicaManager.replicaAlterLogDirsManager).andReturn(replicaAlterLogDirsManager).anyTimes()
-    stub(replica, partition, replicaManager)
-    replay(leaderEpochs, replicaManager, logManager, quota, replica, partition)
+    expect(replicaManager.brokerTopicStats).andReturn(mock(classOf[BrokerTopicStats]))
+    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_OFFSET)).asJava
+    val offsetsReply = Map(
+      t1p0 -> newOffsetForLeaderPartitionResult(t1p0, UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET)).asJava
 
     //Create the thread
-    val endPoint = new BrokerEndPoint(0, "localhost", 1000)
-    val mockNetwork = new ReplicaFetcherMockBlockingSend(offsetsReply, endPoint, new SystemTime())
-    val thread = new ReplicaFetcherThread("bob", 0, endPoint, configs(0), replicaManager, new Metrics(), new SystemTime(), quota, Some(mockNetwork))
-    thread.addPartitions(Map(t1p0 -> initialFetchOffset))
+    val mockNetwork = new ReplicaFetcherMockBlockingSend(offsetsReply, brokerEndPoint, new SystemTime())
+    val thread = new ReplicaFetcherThread("bob", 0, brokerEndPoint, config, failedPartitions, replicaManager, new Metrics(), new SystemTime(), quota, Some(mockNetwork))
+    thread.addPartitions(Map(t1p0 -> initialFetchState(initialFetchOffset)))
 
     //Run it
     thread.doWork()
 
-    //We should have truncated to the highwatermark for partitino 2 only
+    //We should have truncated to initial fetch offset
     assertEquals(initialFetchOffset, truncated.getValue)
   }
 
@@ -256,40 +699,44 @@ class ReplicaFetcherThreadTest {
     val truncated: Capture[Long] = newCapture(CaptureType.ALL)
 
     // Setup all the dependencies
-    val configs = TestUtils.createBrokerConfigs(1, "localhost:1234").map(KafkaConfig.fromProps)
-    val quota = createNiceMock(classOf[kafka.server.ReplicationQuotaManager])
-    val leaderEpochs = createNiceMock(classOf[LeaderEpochCache])
-    val logManager = createMock(classOf[kafka.log.LogManager])
-    val replicaAlterLogDirsManager = createMock(classOf[ReplicaAlterLogDirsManager])
-    val replica = createNiceMock(classOf[Replica])
-    val partition = createMock(classOf[Partition])
-    val replicaManager = createMock(classOf[kafka.server.ReplicaManager])
-
+    val config = kafkaConfigNoTruncateOnFetch
+    val quota: ReplicationQuotaManager = createNiceMock(classOf[ReplicationQuotaManager])
+    val logManager: LogManager = createMock(classOf[LogManager])
+    val replicaAlterLogDirsManager: ReplicaAlterLogDirsManager = createMock(classOf[ReplicaAlterLogDirsManager])
+    val log: Log = createNiceMock(classOf[Log])
+    val partition: Partition = createMock(classOf[Partition])
+    val replicaManager: ReplicaManager = createMock(classOf[ReplicaManager])
+
+    val leaderEpoch = 5
     val highWaterMark = 100
     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.epochs).andReturn(Some(leaderEpochs)).anyTimes()
-    expect(replica.logEndOffset).andReturn(new LogOffsetMetadata(initialLeo)).anyTimes()
-    expect(leaderEpochs.latestEpoch).andReturn(5)
+    expect(partition.localLogOrException).andReturn(log).anyTimes()
+    expect(log.latestEpoch).andReturn(Some(leaderEpoch)).anyTimes()
+    // this is for the last reply with EpochEndOffset(5, 156)
+    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(leaderEpochs, replicaManager, logManager, quota, replica, partition)
+    expect(replicaManager.brokerTopicStats).andReturn(mock(classOf[BrokerTopicStats]))
+    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(
-      t1p0 -> new EpochEndOffset(NOT_LEADER_FOR_PARTITION, EpochEndOffset.UNDEFINED_EPOCH_OFFSET),
-      t1p1 -> new EpochEndOffset(UNKNOWN_SERVER_ERROR, EpochEndOffset.UNDEFINED_EPOCH_OFFSET)
+      t1p0 -> newOffsetForLeaderPartitionResult(t1p0, NOT_LEADER_OR_FOLLOWER, UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET),
+      t1p1 -> newOffsetForLeaderPartitionResult(t1p1, UNKNOWN_SERVER_ERROR, UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET)
     ).asJava
 
     //Create the thread
-    val endPoint = new BrokerEndPoint(0, "localhost", 1000)
-    val mockNetwork = new ReplicaFetcherMockBlockingSend(offsetsReply, endPoint, new SystemTime())
-    val thread = new ReplicaFetcherThread("bob", 0, endPoint, configs(0), replicaManager, new Metrics(), new SystemTime(), quota, Some(mockNetwork))
-    thread.addPartitions(Map(t1p0 -> 0, t2p1 -> 0))
+    val mockNetwork = new ReplicaFetcherMockBlockingSend(offsetsReply, brokerEndPoint, new SystemTime())
+    val thread = new ReplicaFetcherThread("bob", 0, brokerEndPoint, config, failedPartitions, replicaManager, new Metrics(), new SystemTime(), quota, Some(mockNetwork))
+    thread.addPartitions(Map(t1p0 -> initialFetchState(0L), t1p1 -> initialFetchState(0L)))
 
     //Run thread 3 times
     (0 to 3).foreach { _ =>
@@ -300,7 +747,7 @@ class ReplicaFetcherThreadTest {
     assertEquals(0, truncated.getValues.size())
 
     //New leader elected and replies
-    offsetsReply.put(t1p0, new EpochEndOffset(156))
+    offsetsReply.put(t1p0, newOffsetForLeaderPartitionResult(t1p0, leaderEpoch, 156))
 
     thread.doWork()
 
@@ -310,88 +757,97 @@ class ReplicaFetcherThreadTest {
 
   @Test
   def shouldMovePartitionsOutOfTruncatingLogState(): Unit = {
-    val config = KafkaConfig.fromProps(TestUtils.createBrokerConfig(1, "localhost:1234"))
+    val config = kafkaConfigNoTruncateOnFetch
 
     //Setup all stubs
-    val quota = createNiceMock(classOf[ReplicationQuotaManager])
-    val leaderEpochs = createNiceMock(classOf[LeaderEpochCache])
-    val logManager = createNiceMock(classOf[LogManager])
-    val replicaAlterLogDirsManager = createMock(classOf[ReplicaAlterLogDirsManager])
-    val replica = createNiceMock(classOf[Replica])
-    val partition = createMock(classOf[Partition])
-    val replicaManager = createNiceMock(classOf[ReplicaManager])
+    val quota: ReplicationQuotaManager = createNiceMock(classOf[ReplicationQuotaManager])
+    val logManager: LogManager = createNiceMock(classOf[LogManager])
+    val replicaAlterLogDirsManager: ReplicaAlterLogDirsManager = createMock(classOf[ReplicaAlterLogDirsManager])
+    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.epochs).andReturn(Some(leaderEpochs)).anyTimes()
-    expect(replica.logEndOffset).andReturn(new LogOffsetMetadata(0)).anyTimes()
-    expect(leaderEpochs.latestEpoch).andReturn(5)
+    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(leaderEpochs, replicaManager, logManager, quota, replica)
+    replay(replicaManager, logManager, quota, partition, log)
 
     //Define the offsets for the OffsetsForLeaderEpochResponse
     val offsetsReply = Map(
-      t1p0 -> new EpochEndOffset(1), t1p1 -> new EpochEndOffset(1)
+      t1p0 -> newOffsetForLeaderPartitionResult(t1p0, leaderEpoch, 1),
+      t1p1 -> newOffsetForLeaderPartitionResult(t1p1, leaderEpoch, 1)
     ).asJava
 
     //Create the fetcher thread
-    val endPoint = new BrokerEndPoint(0, "localhost", 1000)
-    val mockNetwork = new ReplicaFetcherMockBlockingSend(offsetsReply, endPoint, new SystemTime())
-    val thread = new ReplicaFetcherThread("bob", 0, endPoint, config, replicaManager, new Metrics(), new SystemTime(), quota, Some(mockNetwork))
+    val mockNetwork = new ReplicaFetcherMockBlockingSend(offsetsReply, brokerEndPoint, new SystemTime())
+    val thread = new ReplicaFetcherThread("bob", 0, brokerEndPoint, config, failedPartitions, replicaManager, new Metrics(), new SystemTime(), quota, Some(mockNetwork))
 
     //When
-    thread.addPartitions(Map(t1p0 -> 0, t1p1 -> 0))
+    thread.addPartitions(Map(t1p0 -> initialFetchState(0L), t1p1 -> initialFetchState(0L)))
 
     //Then all partitions should start in an TruncatingLog state
-    assertTrue(thread.partitionStates.partitionStates().asScala.forall(_.value().truncatingLog))
+    assertEquals(Option(Truncating), thread.fetchState(t1p0).map(_.state))
+    assertEquals(Option(Truncating), thread.fetchState(t1p1).map(_.state))
 
     //When
     thread.doWork()
 
     //Then none should be TruncatingLog anymore
-    assertFalse(thread.partitionStates.partitionStates().asScala.forall(_.value().truncatingLog))
+    assertEquals(Option(Fetching), thread.fetchState(t1p0).map(_.state))
+    assertEquals(Option(Fetching), thread.fetchState(t1p1).map(_.state))
   }
 
   @Test
   def shouldFilterPartitionsMadeLeaderDuringLeaderEpochRequest(): Unit ={
-    val config = KafkaConfig.fromProps(TestUtils.createBrokerConfig(1, "localhost:1234"))
+    val config = kafkaConfigNoTruncateOnFetch
     val truncateToCapture: Capture[Long] = newCapture(CaptureType.ALL)
     val initialLEO = 100
 
     //Setup all stubs
-    val quota = createNiceMock(classOf[ReplicationQuotaManager])
-    val leaderEpochs = createNiceMock(classOf[LeaderEpochCache])
-    val logManager = createNiceMock(classOf[LogManager])
-    val replicaAlterLogDirsManager = createMock(classOf[ReplicaAlterLogDirsManager])
-    val replica = createNiceMock(classOf[Replica])
-    val partition = createMock(classOf[Partition])
-    val replicaManager = createNiceMock(classOf[ReplicaManager])
+    val quota: ReplicationQuotaManager = createNiceMock(classOf[ReplicationQuotaManager])
+    val logManager: LogManager = createNiceMock(classOf[LogManager])
+    val replicaAlterLogDirsManager: ReplicaAlterLogDirsManager = createMock(classOf[ReplicaAlterLogDirsManager])
+    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.epochs).andReturn(Some(leaderEpochs)).anyTimes()
-    expect(replica.logEndOffset).andReturn(new LogOffsetMetadata(initialLEO)).anyTimes()
-    expect(leaderEpochs.latestEpoch).andReturn(5)
+    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(leaderEpochs, replicaManager, logManager, quota, replica, partition)
+    replay(replicaManager, logManager, quota, partition, log)
 
     //Define the offsets for the OffsetsForLeaderEpochResponse
     val offsetsReply = Map(
-      t1p0 -> new EpochEndOffset(52), t1p1 -> new EpochEndOffset(49)
+      t1p0 -> newOffsetForLeaderPartitionResult(t1p0, 5, 52),
+      t1p1 -> newOffsetForLeaderPartitionResult(t1p1, 5, 49)
     ).asJava
 
     //Create the fetcher thread
-    val endPoint = new BrokerEndPoint(0, "localhost", 1000)
-    val mockNetwork = new ReplicaFetcherMockBlockingSend(offsetsReply, endPoint, new SystemTime())
-    val thread = new ReplicaFetcherThread("bob", 0, endPoint, config, replicaManager, new Metrics(), new SystemTime(), quota, Some(mockNetwork))
+    val mockNetwork = new ReplicaFetcherMockBlockingSend(offsetsReply, brokerEndPoint, new SystemTime())
+    val thread = new ReplicaFetcherThread("bob", 0, brokerEndPoint, config, failedPartitions, replicaManager, new Metrics(),
+      new SystemTime(), quota, Some(mockNetwork))
 
     //When
-    thread.addPartitions(Map(t1p0 -> 0, t1p1 -> 0))
+    thread.addPartitions(Map(t1p0 -> initialFetchState(0L), t1p1 -> initialFetchState(0L)))
 
     //When the epoch request is outstanding, remove one of the partitions to simulate a leader change. We do this via a callback passed to the mock thread
     val partitionThatBecameLeader = t1p0
@@ -406,15 +862,131 @@ class ReplicaFetcherThreadTest {
     assertEquals(49, truncateToCapture.getValue)
   }
 
-  def stub(replica: Replica, partition: Partition, replicaManager: ReplicaManager) = {
-    expect(replicaManager.getReplica(t1p0)).andReturn(Some(replica)).anyTimes()
-    expect(replicaManager.getReplicaOrException(t1p0)).andReturn(replica).anyTimes()
-    expect(replicaManager.getPartition(t1p0)).andReturn(Some(partition)).anyTimes()
-    expect(replicaManager.getReplica(t1p1)).andReturn(Some(replica)).anyTimes()
-    expect(replicaManager.getReplicaOrException(t1p1)).andReturn(replica).anyTimes()
-    expect(replicaManager.getPartition(t1p1)).andReturn(Some(partition)).anyTimes()
-    expect(replicaManager.getReplica(t2p1)).andReturn(Some(replica)).anyTimes()
-    expect(replicaManager.getReplicaOrException(t2p1)).andReturn(replica).anyTimes()
-    expect(replicaManager.getPartition(t2p1)).andReturn(Some(partition)).anyTimes()
+  @Test
+  def shouldCatchExceptionFromBlockingSendWhenShuttingDownReplicaFetcherThread(): Unit = {
+    val props = TestUtils.createBrokerConfig(1, "localhost:1234")
+    val config = KafkaConfig.fromProps(props)
+    val mockBlockingSend: BlockingSend = createMock(classOf[BlockingSend])
+
+    expect(mockBlockingSend.initiateClose()).andThrow(new IllegalArgumentException()).once()
+    expect(mockBlockingSend.close()).andThrow(new IllegalStateException()).once()
+    val replicaManager: ReplicaManager = mock(classOf[ReplicaManager])
+    expect(replicaManager.brokerTopicStats).andReturn(mock(classOf[BrokerTopicStats]))
+    replay(mockBlockingSend, replicaManager)
+
+    val thread = new ReplicaFetcherThread(
+      name = "bob",
+      fetcherId = 0,
+      sourceBroker = brokerEndPoint,
+      brokerConfig = config,
+      failedPartitions = failedPartitions,
+      replicaMgr = replicaManager,
+      metrics =  new Metrics(),
+      time = new SystemTime(),
+      quota = null,
+      leaderEndpointBlockingSend = Some(mockBlockingSend))
+    thread.start()
+
+    // Verify that:
+    //   1) IllegalArgumentException thrown by BlockingSend#initiateClose() during `initiateShutdown` is not propagated
+    //   2) BlockingSend.close() is invoked even if BlockingSend#initiateClose() fails
+    //   3) IllegalStateException thrown by BlockingSend.close() during `awaitShutdown` is not propagated
+    thread.initiateShutdown()
+    thread.awaitShutdown()
+    verify(mockBlockingSend)
+  }
+
+  @Test
+  def shouldUpdateReassignmentBytesInMetrics(): Unit = {
+    assertProcessPartitionDataWhen(isReassigning = true)
+  }
+
+  @Test
+  def shouldNotUpdateReassignmentBytesInMetricsWhenNoReassignmentsInProgress(): Unit = {
+    assertProcessPartitionDataWhen(isReassigning = false)
+  }
+
+  private def newOffsetForLeaderPartitionResult(
+   tp: TopicPartition,
+   leaderEpoch: Int,
+   endOffset: Long
+ ): EpochEndOffset = {
+    newOffsetForLeaderPartitionResult(tp, Errors.NONE, leaderEpoch, endOffset)
+  }
+
+  private def newOffsetForLeaderPartitionResult(
+    tp: TopicPartition,
+    error: Errors,
+    leaderEpoch: Int,
+    endOffset: Long
+  ): EpochEndOffset = {
+    new EpochEndOffset()
+      .setPartition(tp.partition)
+      .setErrorCode(error.code)
+      .setLeaderEpoch(leaderEpoch)
+      .setEndOffset(endOffset)
+  }
+
+  private def assertProcessPartitionDataWhen(isReassigning: Boolean): Unit = {
+    val props = TestUtils.createBrokerConfig(1, "localhost:1234")
+    val config = KafkaConfig.fromProps(props)
+
+    val mockBlockingSend: BlockingSend = createNiceMock(classOf[BlockingSend])
+
+    val log: Log = createNiceMock(classOf[Log])
+
+    val partition: Partition = createNiceMock(classOf[Partition])
+    expect(partition.localLogOrException).andReturn(log)
+    expect(partition.isReassigning).andReturn(isReassigning)
+    expect(partition.isAddingLocalReplica).andReturn(isReassigning)
+
+    val replicaManager: ReplicaManager = createNiceMock(classOf[ReplicaManager])
+    expect(replicaManager.nonOfflinePartition(anyObject[TopicPartition])).andReturn(Some(partition))
+    val brokerTopicStats = new BrokerTopicStats
+    expect(replicaManager.brokerTopicStats).andReturn(brokerTopicStats).anyTimes()
+
+    val replicaQuota: ReplicaQuota = createNiceMock(classOf[ReplicaQuota])
+    replay(mockBlockingSend, replicaManager, partition, log, replicaQuota)
+
+    val thread = new ReplicaFetcherThread(
+      name = "bob",
+      fetcherId = 0,
+      sourceBroker = brokerEndPoint,
+      brokerConfig = config,
+      failedPartitions = failedPartitions,
+      replicaMgr = replicaManager,
+      metrics =  new Metrics(),
+      time = new SystemTime(),
+      quota = replicaQuota,
+      leaderEndpointBlockingSend = Some(mockBlockingSend))
+
+    val records = MemoryRecords.withRecords(CompressionType.NONE,
+      new SimpleRecord(1000, "foo".getBytes(StandardCharsets.UTF_8)))
+
+    val partitionData: thread.FetchData = new FetchResponse.PartitionData[Records](
+      Errors.NONE, 0, 0, 0, Optional.empty(), Collections.emptyList(), records)
+    thread.processPartitionData(t1p0, 0, partitionData)
+
+    if (isReassigning)
+      assertEquals(records.sizeInBytes(), brokerTopicStats.allTopicsStats.reassignmentBytesInPerSec.get.count())
+    else
+      assertEquals(0, brokerTopicStats.allTopicsStats.reassignmentBytesInPerSec.get.count())
+
+    assertEquals(records.sizeInBytes(), brokerTopicStats.allTopicsStats.replicationBytesInRate.get.count())
+  }
+
+  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()
+  }
+
+  private def kafkaConfigNoTruncateOnFetch: KafkaConfig = {
+    val props = TestUtils.createBrokerConfig(1, "localhost:1234")
+    props.setProperty(KafkaConfig.InterBrokerProtocolVersionProp, KAFKA_2_6_IV0.version)
+    KafkaConfig.fromProps(props)
   }
-}
\ No newline at end of file
+}
diff --git a/core/src/test/scala/unit/kafka/server/ReplicaManagerQuotasTest.scala b/core/src/test/scala/unit/kafka/server/ReplicaManagerQuotasTest.scala
index 334321ae2307d..25ee84eb7b3f4 100644
--- a/core/src/test/scala/unit/kafka/server/ReplicaManagerQuotasTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ReplicaManagerQuotasTest.scala
@@ -17,11 +17,11 @@
 package kafka.server
 
 import java.io.File
-import java.util.Properties
+import java.util.{Optional, Properties}
 import java.util.concurrent.atomic.AtomicBoolean
 
-import kafka.cluster.Replica
-import kafka.log.Log
+import kafka.cluster.Partition
+import kafka.log.{Log, LogManager, LogOffsetSnapshot}
 import kafka.utils._
 import kafka.zk.KafkaZkClient
 import org.apache.kafka.common.TopicPartition
@@ -30,11 +30,11 @@ import org.apache.kafka.common.record.{CompressionType, MemoryRecords, SimpleRec
 import org.apache.kafka.common.requests.FetchRequest.PartitionData
 import org.easymock.EasyMock
 import EasyMock._
-import org.apache.kafka.common.requests.IsolationLevel
+import kafka.server.QuotaFactory.QuotaManagers
 import org.junit.Assert._
 import org.junit.{After, Test}
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 
 class ReplicaManagerQuotasTest {
   val configs = TestUtils.createBrokerConfigs(2, TestUtils.MockZkConnect).map(KafkaConfig.fromProps(_, new Properties()))
@@ -43,7 +43,10 @@ class ReplicaManagerQuotasTest {
   val record = new SimpleRecord("some-data-in-a-message".getBytes())
   val topicPartition1 = new TopicPartition("test-topic", 1)
   val topicPartition2 = new TopicPartition("test-topic", 2)
-  val fetchInfo = Seq(topicPartition1 -> new PartitionData(0, 0, 100), topicPartition2 -> new PartitionData(0, 0, 100))
+  val fetchInfo = Seq(
+    topicPartition1 -> new PartitionData(0, 0, 100, Optional.empty()),
+    topicPartition2 -> new PartitionData(0, 0, 100, Optional.empty()))
+  var quotaManager: QuotaManagers = _
   var replicaManager: ReplicaManager = _
 
   @Test
@@ -52,19 +55,19 @@ class ReplicaManagerQuotasTest {
     val followerReplicaId = configs.last.brokerId
 
     val quota = mockQuota(1000000)
-    expect(quota.isQuotaExceeded()).andReturn(false).once()
-    expect(quota.isQuotaExceeded()).andReturn(true).once()
+    expect(quota.isQuotaExceeded).andReturn(false).once()
+    expect(quota.isQuotaExceeded).andReturn(true).once()
     replay(quota)
 
     val fetch = replicaManager.readFromLocalLog(
       replicaId = followerReplicaId,
       fetchOnlyFromLeader = true,
-      readOnlyCommitted = true,
+      fetchIsolation = FetchHighWatermark,
       fetchMaxBytes = Int.MaxValue,
       hardMaxBytesLimit = false,
       readPartitionInfo = fetchInfo,
       quota = quota,
-      isolationLevel = IsolationLevel.READ_UNCOMMITTED)
+      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)
 
@@ -78,19 +81,19 @@ class ReplicaManagerQuotasTest {
     val followerReplicaId = configs.last.brokerId
 
     val quota = mockQuota(1000000)
-    expect(quota.isQuotaExceeded()).andReturn(true).once()
-    expect(quota.isQuotaExceeded()).andReturn(true).once()
+    expect(quota.isQuotaExceeded).andReturn(true).once()
+    expect(quota.isQuotaExceeded).andReturn(true).once()
     replay(quota)
 
     val fetch = replicaManager.readFromLocalLog(
       replicaId = followerReplicaId,
       fetchOnlyFromLeader = true,
-      readOnlyCommitted = true,
+      fetchIsolation = FetchHighWatermark,
       fetchMaxBytes = Int.MaxValue,
       hardMaxBytesLimit = false,
       readPartitionInfo = fetchInfo,
       quota = quota,
-      isolationLevel = IsolationLevel.READ_UNCOMMITTED)
+      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,
@@ -103,19 +106,19 @@ class ReplicaManagerQuotasTest {
     val followerReplicaId = configs.last.brokerId
 
     val quota = mockQuota(1000000)
-    expect(quota.isQuotaExceeded()).andReturn(false).once()
-    expect(quota.isQuotaExceeded()).andReturn(false).once()
+    expect(quota.isQuotaExceeded).andReturn(false).once()
+    expect(quota.isQuotaExceeded).andReturn(false).once()
     replay(quota)
 
     val fetch = replicaManager.readFromLocalLog(
       replicaId = followerReplicaId,
       fetchOnlyFromLeader = true,
-      readOnlyCommitted = true,
+      fetchIsolation = FetchHighWatermark,
       fetchMaxBytes = Int.MaxValue,
       hardMaxBytesLimit = false,
       readPartitionInfo = fetchInfo,
       quota = quota,
-      isolationLevel = IsolationLevel.READ_UNCOMMITTED)
+      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,
@@ -128,19 +131,19 @@ class ReplicaManagerQuotasTest {
     val followerReplicaId = configs.last.brokerId
 
     val quota = mockQuota(1000000)
-    expect(quota.isQuotaExceeded()).andReturn(false).once()
-    expect(quota.isQuotaExceeded()).andReturn(true).once()
+    expect(quota.isQuotaExceeded).andReturn(false).once()
+    expect(quota.isQuotaExceeded).andReturn(true).once()
     replay(quota)
 
     val fetch = replicaManager.readFromLocalLog(
       replicaId = followerReplicaId,
       fetchOnlyFromLeader = true,
-      readOnlyCommitted = true,
+      fetchIsolation = FetchHighWatermark,
       fetchMaxBytes = Int.MaxValue,
       hardMaxBytesLimit = false,
       readPartitionInfo = fetchInfo,
       quota = quota,
-      isolationLevel = IsolationLevel.READ_UNCOMMITTED)
+      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)
 
@@ -148,72 +151,127 @@ class ReplicaManagerQuotasTest {
       fetch.find(_._1 == topicPartition2).get._2.info.records.batches.asScala.size)
   }
 
-  def setUpMocks(fetchInfo: Seq[(TopicPartition, PartitionData)], record: SimpleRecord = this.record, bothReplicasInSync: Boolean = false) {
-    val zkClient = EasyMock.createMock(classOf[KafkaZkClient])
-    val scheduler = createNiceMock(classOf[KafkaScheduler])
+  @Test
+  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 = LogOffsetMetadata(messageOffset = 100L, segmentBaseOffset = 0L, relativePositionInSegment = 500)
+      val partition: Partition = EasyMock.createMock(classOf[Partition])
+
+      val offsetSnapshot = LogOffsetSnapshot(
+        logStartOffset = 0L,
+        logEndOffset = endOffsetMetadata,
+        highWatermark = endOffsetMetadata,
+        lastStableOffset = endOffsetMetadata)
+      EasyMock.expect(partition.fetchOffsetSnapshot(Optional.empty(), fetchOnlyFromLeader = true))
+          .andReturn(offsetSnapshot)
+
+      val replicaManager: ReplicaManager = EasyMock.createMock(classOf[ReplicaManager])
+      EasyMock.expect(replicaManager.getPartitionOrException(EasyMock.anyObject[TopicPartition]))
+        .andReturn(partition).anyTimes()
+
+      EasyMock.expect(replicaManager.shouldLeaderThrottle(EasyMock.anyObject[ReplicaQuota], EasyMock.anyObject[Partition], 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(LogOffsetMetadata(messageOffset = 50L, segmentBaseOffset = 0L,
+         relativePositionInSegment = 250), new PartitionData(50, 0, 1, Optional.empty()))
+      val fetchMetadata = FetchMetadata(fetchMinBytes = 1,
+        fetchMaxBytes = 1000,
+        hardMaxBytesLimit = true,
+        fetchOnlyLeader = true,
+        fetchIsolation = FetchLogEnd,
+        isFromFollower = true,
+        replicaId = 1,
+        fetchPartitionStatus = List((tp, fetchPartitionStatus))
+      )
+      new DelayedFetch(delayMs = 600, fetchMetadata = fetchMetadata, replicaManager = replicaManager,
+        quota = null, clientMetadata = None, responseCallback = null) {
+        override def forceComplete(): Boolean = true
+      }
+    }
+
+    assertTrue("In sync replica should complete", setupDelayedFetch(isReplicaInSync = true).tryComplete())
+    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): Unit = {
+    val zkClient: KafkaZkClient = EasyMock.createMock(classOf[KafkaZkClient])
+    val scheduler: KafkaScheduler = createNiceMock(classOf[KafkaScheduler])
 
     //Create log which handles both a regular read and a 0 bytes read
-    val log = createNiceMock(classOf[Log])
+    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(), geq(1), anyObject(), anyObject(),
-      EasyMock.eq(IsolationLevel.READ_UNCOMMITTED))).andReturn(
+    expect(log.read(anyObject(),
+      maxLength = geq(1),
+      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(), EasyMock.eq(0), anyObject(), anyObject(),
-      EasyMock.eq(IsolationLevel.READ_UNCOMMITTED))).andReturn(
+    expect(log.read(anyObject(),
+      maxLength = EasyMock.eq(0),
+      isolation = anyObject(),
+      minOneMessage = anyBoolean())).andReturn(
       FetchDataInfo(
-        new LogOffsetMetadata(0L, 0L, 0),
+        LogOffsetMetadata(0L, 0L, 0),
         MemoryRecords.EMPTY
       )).anyTimes()
     replay(log)
 
     //Create log manager
-    val logManager = createMock(classOf[kafka.log.LogManager])
+    val logManager: LogManager = createMock(classOf[LogManager])
 
     //Return the same log for each partition as it doesn't matter
     expect(logManager.getLog(anyObject(), anyBoolean())).andReturn(Some(log)).anyTimes()
     expect(logManager.liveLogDirs).andReturn(Array.empty[File]).anyTimes()
     replay(logManager)
 
+    val alterIsrManager: AlterIsrManager = createMock(classOf[AlterIsrManager])
+
+    val leaderBrokerId = configs.head.brokerId
+    quotaManager = QuotaFactory.instantiate(configs.head, metrics, time, "")
     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 AtomicBoolean(false), quotaManager,
+      new BrokerTopicStats, new MetadataCache(leaderBrokerId), new LogDirFailureChannel(configs.head.logDirs.size), alterIsrManager)
 
     //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),
+        addingReplicas = Seq.empty,
+        removingReplicas = Seq.empty
+      )
     }
   }
 
   @After
-  def tearDown() {
-    replicaManager.shutdown(false)
+  def tearDown(): Unit = {
+    Option(replicaManager).foreach(_.shutdown(false))
+    Option(quotaManager).foreach(_.shutdown())
     metrics.close()
   }
 
   def mockQuota(bound: Long): ReplicaQuota = {
-    val quota = createMock(classOf[ReplicaQuota])
+    val quota: ReplicaQuota = createMock(classOf[ReplicaQuota])
     expect(quota.isThrottled(anyObject())).andReturn(true).anyTimes()
     quota
   }
diff --git a/core/src/test/scala/unit/kafka/server/ReplicaManagerTest.scala b/core/src/test/scala/unit/kafka/server/ReplicaManagerTest.scala
index ce8868861f78b..f73223e11abcb 100644
--- a/core/src/test/scala/unit/kafka/server/ReplicaManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ReplicaManagerTest.scala
@@ -18,65 +18,93 @@
 package kafka.server
 
 import java.io.File
-import java.util.Properties
-import java.util.concurrent.atomic.AtomicBoolean
-
-import kafka.log.LogConfig
-import kafka.utils.{MockScheduler, MockTime, TestUtils}
-import TestUtils.createBroker
+import java.net.InetAddress
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+import java.util.concurrent.{CountDownLatch, TimeUnit}
+import java.util.{Optional, Properties}
+
+import kafka.api._
+import kafka.log.{AppendOrigin, Log, LogConfig, LogManager, ProducerStateManager}
+import kafka.cluster.{BrokerEndPoint, Partition}
+import kafka.log.LeaderOffsetIncremented
+import kafka.server.QuotaFactory.{QuotaManagers, UnboundedQuota}
+import kafka.server.checkpoints.LazyOffsetCheckpoints
+import kafka.server.checkpoints.OffsetCheckpointFile
+import kafka.server.epoch.util.ReplicaFetcherMockBlockingSend
+import kafka.utils.TestUtils.createBroker
 import kafka.utils.timer.MockTimer
+import kafka.utils.{MockScheduler, MockTime, TestUtils}
 import kafka.zk.KafkaZkClient
-import org.I0Itec.zkclient.ZkClient
+import org.apache.kafka.common.message.LeaderAndIsrRequestData.LeaderAndIsrPartitionState
+import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData.EpochEndOffset
+import org.apache.kafka.common.message.StopReplicaRequestData.StopReplicaPartitionState
 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.{IsolationLevel, LeaderAndIsrRequest}
-import org.apache.kafka.common.requests.ProduceResponse.PartitionResponse
+import org.apache.kafka.common.replica.ClientMetadata
+import org.apache.kafka.common.replica.ClientMetadata.DefaultClientMetadata
 import org.apache.kafka.common.requests.FetchRequest.PartitionData
 import org.apache.kafka.common.requests.FetchResponse.AbortedTransaction
-import org.apache.kafka.common.{Node, TopicPartition}
-import org.apache.zookeeper.data.Stat
+import org.apache.kafka.common.requests.ProduceResponse.PartitionResponse
+import org.apache.kafka.common.requests._
+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.{IsolationLevel, Node, TopicPartition}
 import org.easymock.EasyMock
 import org.junit.Assert._
 import org.junit.{After, Before, Test}
+import org.mockito.Mockito
 
-import scala.collection.JavaConverters._
-import scala.collection.Map
+import scala.collection.mutable
+import scala.jdk.CollectionConverters._
+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 = _
+  var alterIsrManager: AlterIsrManager = _
+  var config: KafkaConfig = _
+  var quotaManager: QuotaManagers = _
+
+  // Constants defined for readability
+  val zkVersion = 0
+  val correlationId = 0
+  var controllerEpoch = 0
+  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)
+
+    val props = TestUtils.createBrokerConfig(1, TestUtils.MockZkConnect)
+    config = KafkaConfig.fromProps(props)
+    alterIsrManager = EasyMock.createMock(classOf[AlterIsrManager])
+    quotaManager = QuotaFactory.instantiate(config, metrics, time, "")
   }
 
   @After
-  def tearDown() {
+  def tearDown(): Unit = {
+    TestUtils.clearYammerMetrics()
+    Option(quotaManager).foreach(_.shutdown())
     metrics.close()
   }
 
   @Test
-  def testHighWaterMarkDirectoryMapping() {
-    val props = TestUtils.createBrokerConfig(1, TestUtils.MockZkConnect)
-    val config = KafkaConfig.fromProps(props)
+  def testHighWaterMarkDirectoryMapping(): Unit = {
     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))
+      new AtomicBoolean(false), quotaManager, new BrokerTopicStats,
+      new MetadataCache(config.brokerId), new LogDirFailureChannel(config.logDirs.size), alterIsrManager)
     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
@@ -85,17 +113,18 @@ 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)
     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))
+      new AtomicBoolean(false), quotaManager, new BrokerTopicStats,
+      new MetadataCache(config.brokerId), new LogDirFailureChannel(config.logDirs.size), alterIsrManager)
     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
@@ -104,13 +133,11 @@ class ReplicaManagerTest {
   }
 
   @Test
-  def testIllegalRequiredAcks() {
-    val props = TestUtils.createBrokerConfig(1, TestUtils.MockZkConnect)
-    val config = KafkaConfig.fromProps(props)
+  def testIllegalRequiredAcks(): Unit = {
     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), Option(this.getClass.getName))
+      new AtomicBoolean(false), quotaManager, new BrokerTopicStats,
+      new MetadataCache(config.brokerId), new LogDirFailureChannel(config.logDirs.size), alterIsrManager, Option(this.getClass.getName))
     try {
       def callback(responseStatus: Map[TopicPartition, PartitionResponse]) = {
         assert(responseStatus.values.head.error == Errors.INVALID_REQUIRED_ACKS)
@@ -119,7 +146,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)
@@ -127,76 +154,163 @@ 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)
     val logProps = new Properties()
     val mockLogMgr = TestUtils.createLogManager(config.logDirs.map(new File(_)), LogConfig(logProps))
     val aliveBrokers = Seq(createBroker(0, "host0", 0), createBroker(1, "host1", 1))
-    val metadataCache = EasyMock.createMock(classOf[MetadataCache])
+    val metadataCache: MetadataCache = EasyMock.createMock(classOf[MetadataCache])
     EasyMock.expect(metadataCache.getAliveBrokers).andReturn(aliveBrokers).anyTimes()
     EasyMock.replay(metadataCache)
     val rm = 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))
+      new AtomicBoolean(false), quotaManager, new BrokerTopicStats,
+      metadataCache, new LogDirFailureChannel(config.logDirs.size), alterIsrManager)
 
     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,
-        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,
+        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.getLeaderReplicaIfLocal(new TopicPartition(topic, 0))
+      rm.getPartitionOrException(new TopicPartition(topic, 0))
+          .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)
+        assertEquals(Errors.NOT_LEADER_OR_FOLLOWER, response.error)
       }
 
-      // Fetch some messages
-      val fetchResult = fetchAsConsumer(rm, new TopicPartition(topic, 0), new PartitionData(0, 0, 100000),
-        minBytes = 100000)
-      assertFalse(fetchResult.isFired)
-
       // Make this replica the follower
-      val leaderAndIsrRequest2 = new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, 0, 0,
-        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,
+        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
+  def testFencedErrorCausedByBecomeLeader(): Unit = {
+    testFencedErrorCausedByBecomeLeader(0)
+    testFencedErrorCausedByBecomeLeader(1)
+    testFencedErrorCausedByBecomeLeader(10)
+  }
+
+  private[this] def testFencedErrorCausedByBecomeLeader(loopEpochChange: Int): Unit = {
+    val replicaManager = setupReplicaManagerWithMockedPurgatories(new MockTimer(time))
+    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,
+        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))
+      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
+    val timer = new MockTimer(time)
     val replicaManager = setupReplicaManagerWithMockedPurgatories(timer)
 
     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,
-        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,
+        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.getLeaderReplicaIfLocal(new TopicPartition(topic, 0))
+      replicaManager.getPartitionOrException(new TopicPartition(topic, 0))
+        .localLogOrException
 
       val producerId = 234L
       val epoch = 5.toShort
@@ -231,22 +345,32 @@ class ReplicaManagerTest {
 
   @Test
   def testReadCommittedFetchLimitedAtLSO(): Unit = {
-    val timer = new MockTimer
+    val timer = new MockTimer(time)
     val replicaManager = setupReplicaManagerWithMockedPurgatories(timer)
 
     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,
-        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,
+        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.getLeaderReplicaIfLocal(new TopicPartition(topic, 0))
-
+      replicaManager.getPartitionOrException(new TopicPartition(topic, 0))
+        .localLogOrException
 
       val producerId = 234L
       val epoch = 5.toShort
@@ -262,12 +386,14 @@ class ReplicaManagerTest {
       }
 
       // fetch as follower to advance the high watermark
-      fetchAsFollower(replicaManager, new TopicPartition(topic, 0), new PartitionData(numRecords, 0, 100000),
+      fetchAsFollower(replicaManager, new TopicPartition(topic, 0),
+        new PartitionData(numRecords, 0, 100000, Optional.empty()),
         isolationLevel = IsolationLevel.READ_UNCOMMITTED)
 
       // fetch should return empty since LSO should be stuck at 0
       var consumerFetchResult = fetchAsConsumer(replicaManager, new TopicPartition(topic, 0),
-        new PartitionData(0, 0, 100000), isolationLevel = IsolationLevel.READ_COMMITTED)
+        new PartitionData(0, 0, 100000, Optional.empty()),
+        isolationLevel = IsolationLevel.READ_COMMITTED)
       var fetchData = consumerFetchResult.assertFired
       assertEquals(Errors.NONE, fetchData.error)
       assertTrue(fetchData.records.batches.asScala.isEmpty)
@@ -275,7 +401,8 @@ class ReplicaManagerTest {
       assertEquals(Some(List.empty[AbortedTransaction]), fetchData.abortedTransactions)
 
       // delayed fetch should timeout and return nothing
-      consumerFetchResult = fetchAsConsumer(replicaManager, new TopicPartition(topic, 0), new PartitionData(0, 0, 100000),
+      consumerFetchResult = fetchAsConsumer(replicaManager, new TopicPartition(topic, 0),
+        new PartitionData(0, 0, 100000, Optional.empty()),
         isolationLevel = IsolationLevel.READ_COMMITTED, minBytes = 1000)
       assertFalse(consumerFetchResult.isFired)
       timer.advanceClock(1001)
@@ -289,12 +416,14 @@ 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
       // none of the data from the transaction should be visible yet
-      consumerFetchResult = fetchAsConsumer(replicaManager, new TopicPartition(topic, 0), new PartitionData(0, 0, 100000),
+      consumerFetchResult = fetchAsConsumer(replicaManager, new TopicPartition(topic, 0),
+        new PartitionData(0, 0, 100000, Optional.empty()),
         isolationLevel = IsolationLevel.READ_COMMITTED)
 
       fetchData = consumerFetchResult.assertFired
@@ -302,11 +431,13 @@ class ReplicaManagerTest {
       assertTrue(fetchData.records.batches.asScala.isEmpty)
 
       // fetch as follower to advance the high watermark
-      fetchAsFollower(replicaManager, new TopicPartition(topic, 0), new PartitionData(numRecords + 1, 0, 100000),
+      fetchAsFollower(replicaManager, new TopicPartition(topic, 0),
+        new PartitionData(numRecords + 1, 0, 100000, Optional.empty()),
         isolationLevel = IsolationLevel.READ_UNCOMMITTED)
 
       // now all of the records should be fetchable
-      consumerFetchResult = fetchAsConsumer(replicaManager, new TopicPartition(topic, 0), new PartitionData(0, 0, 100000),
+      consumerFetchResult = fetchAsConsumer(replicaManager, new TopicPartition(topic, 0),
+        new PartitionData(0, 0, 100000, Optional.empty()),
         isolationLevel = IsolationLevel.READ_COMMITTED)
 
       fetchData = consumerFetchResult.assertFired
@@ -321,20 +452,31 @@ class ReplicaManagerTest {
 
   @Test
   def testDelayedFetchIncludesAbortedTransactions(): Unit = {
-    val timer = new MockTimer
+    val timer = new MockTimer(time)
     val replicaManager = setupReplicaManagerWithMockedPurgatories(timer)
 
     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,
-        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,
+        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.getLeaderReplicaIfLocal(new TopicPartition(topic, 0))
+      replicaManager.getPartitionOrException(new TopicPartition(topic, 0))
+        .localLogOrException
 
       val producerId = 234L
       val epoch = 5.toShort
@@ -352,16 +494,19 @@ 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
-      fetchAsFollower(replicaManager, new TopicPartition(topic, 0), new PartitionData(numRecords + 1, 0, 100000),
+      fetchAsFollower(replicaManager, new TopicPartition(topic, 0),
+        new PartitionData(numRecords + 1, 0, 100000, Optional.empty()),
         isolationLevel = IsolationLevel.READ_UNCOMMITTED)
 
       // Set the minBytes in order force this request to enter purgatory. When it returns, we should still
       // see the newly aborted transaction.
-      val fetchResult = fetchAsConsumer(replicaManager, new TopicPartition(topic, 0), new PartitionData(0, 0, 100000),
+      val fetchResult = fetchAsConsumer(replicaManager, new TopicPartition(topic, 0),
+        new PartitionData(0, 0, 100000, Optional.empty()),
         isolationLevel = IsolationLevel.READ_COMMITTED, minBytes = 10000)
       assertFalse(fetchResult.isFired)
 
@@ -383,20 +528,31 @@ class ReplicaManagerTest {
   }
 
   @Test
-  def testFetchBeyondHighWatermarkReturnEmptyResponse() {
-    val rm = setupReplicaManagerWithMockedPurgatories(new MockTimer, aliveBrokerIds = Seq(0, 1, 2))
+  def testFetchBeyondHighWatermark(): Unit = {
+    val rm = setupReplicaManagerWithMockedPurgatories(new MockTimer(time), 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,
-        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,
+        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.getLeaderReplicaIfLocal(new TopicPartition(topic, 0))
+      rm.getPartitionOrException(new TopicPartition(topic, 0))
+        .localLogOrException
 
       // Append a couple of messages.
       for(i <- 1 to 2) {
@@ -406,14 +562,18 @@ class ReplicaManagerTest {
         }
       }
 
-      // Fetch a message above the high watermark as a follower
-      val followerFetchResult = fetchAsFollower(rm, new TopicPartition(topic, 0), new PartitionData(1, 0, 100000))
+      // Followers are always allowed to fetch above the high watermark
+      val followerFetchResult = fetchAsFollower(rm, new TopicPartition(topic, 0),
+        new PartitionData(1, 0, 100000, Optional.empty()))
       val followerFetchData = followerFetchResult.assertFired
       assertEquals("Should not give an exception", Errors.NONE, followerFetchData.error)
       assertTrue("Should return some data", followerFetchData.records.batches.iterator.hasNext)
 
-      // Fetch a message above the high watermark as a consumer
-      val consumerFetchResult = fetchAsConsumer(rm, new TopicPartition(topic, 0), new PartitionData(1, 0, 100000))
+      // Consumers are not allowed to consume above the high watermark. However, since the
+      // high watermark could be stale at the time of the request, we do not return an out of
+      // range error and instead return an empty record set.
+      val consumerFetchResult = fetchAsConsumer(rm, new TopicPartition(topic, 0),
+        new PartitionData(1, 0, 100000, Optional.empty()))
       val consumerFetchData = consumerFetchResult.assertFired
       assertEquals("Should not give an exception", Errors.NONE, consumerFetchData.error)
       assertEquals("Should return empty response", MemoryRecords.EMPTY, consumerFetchData.records)
@@ -422,26 +582,142 @@ class ReplicaManagerTest {
     }
   }
 
+  @Test
+  def testFollowerStateNotUpdatedIfLogReadFails(): Unit = {
+    val maxFetchBytes = 1024 * 1024
+    val aliveBrokersIds = Seq(0, 1)
+    val leaderEpoch = 5
+    val replicaManager = setupReplicaManagerWithMockedPurgatories(new MockTimer(time), 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,
+        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() {
-    val replicaManager = setupReplicaManagerWithMockedPurgatories(new MockTimer, aliveBrokerIds = Seq(0, 1, 2))
+  def testFetchMessagesWhenNotFollowerForOnePartition(): Unit = {
+    val replicaManager = setupReplicaManagerWithMockedPurgatories(new MockTimer(time), 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,
-        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,
+        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, (_, _) => ())
@@ -463,15 +739,17 @@ class ReplicaManagerTest {
 
         val tp0Status = responseStatusMap.get(tp0)
         assertTrue(tp0Status.isDefined)
-        assertEquals(1, tp0Status.get.highWatermark)
-        assertEquals(None, tp0Status.get.lastStableOffset)
+        // the response contains high watermark on the leader before it is updated based
+        // on this fetch request
+        assertEquals(0, tp0Status.get.highWatermark)
+        assertEquals(Some(0), tp0Status.get.lastStableOffset)
         assertEquals(Errors.NONE, tp0Status.get.error)
         assertTrue(tp0Status.get.records.batches.iterator.hasNext)
 
         val tp1Status = responseStatusMap.get(tp1)
         assertTrue(tp1Status.isDefined)
         assertEquals(0, tp1Status.get.highWatermark)
-        assertEquals(None, tp0Status.get.lastStableOffset)
+        assertEquals(Some(0), tp0Status.get.lastStableOffset)
         assertEquals(Errors.NONE, tp1Status.get.error)
         assertFalse(tp1Status.get.records.batches.iterator.hasNext)
       }
@@ -483,25 +761,829 @@ class ReplicaManagerTest {
         fetchMaxBytes = Int.MaxValue,
         hardMaxBytesLimit = false,
         fetchInfos = Seq(
-          tp0 -> new PartitionData(1, 0, 100000),
-          tp1 -> new PartitionData(1, 0, 100000)),
+          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.getReplica(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.getReplica(tp1)
-      val tp1Replica = replicaManager.getReplica(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)
     }
   }
 
+  @Test
+  def testBecomeFollowerWhenLeaderIsUnchangedButMissedLeaderUpdate(): Unit = {
+    verifyBecomeFollowerWhenLeaderIsUnchangedButMissedLeaderUpdate(new Properties, expectTruncation = false)
+  }
+
+  @Test
+  def testBecomeFollowerWhenLeaderIsUnchangedButMissedLeaderUpdateIbp26(): Unit = {
+    val extraProps = new Properties
+    extraProps.put(KafkaConfig.InterBrokerProtocolVersionProp, KAFKA_2_6_IV0.version)
+    verifyBecomeFollowerWhenLeaderIsUnchangedButMissedLeaderUpdate(extraProps, expectTruncation = true)
+  }
+
+  /**
+   * If a partition becomes a follower and the leader is unchanged it should check for truncation
+   * if the epoch has increased by more than one (which suggests it has missed an update). For
+   * IBP version 2.7 onwards, we don't require this since we can truncate at any time based
+   * on diverging epochs returned in fetch responses.
+   */
+  private def verifyBecomeFollowerWhenLeaderIsUnchangedButMissedLeaderUpdate(extraProps: Properties,
+                                                                             expectTruncation: Boolean): Unit = {
+    val topicPartition = 0
+    val followerBrokerId = 0
+    val leaderBrokerId = 1
+    val controllerId = 0
+    val controllerEpoch = 0
+    var 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, mockLogMgr) = prepareReplicaManagerAndLogManager(new MockTimer(time),
+      topicPartition, leaderEpoch + leaderEpochIncrement, followerBrokerId, leaderBrokerId, countDownLatch,
+      expectTruncation = expectTruncation, localLogOffset = Some(10), extraProps = extraProps)
+
+    // Initialize partition state to follower, with leader = 1, leaderEpoch = 1
+    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,
+      Seq(leaderAndIsrPartitionState(tp, leaderEpoch, leaderBrokerId, aliveBrokerIds)).asJava,
+      Set(new Node(followerBrokerId, "host1", 0),
+        new Node(leaderBrokerId, "host2", 1)).asJava).build()
+    replicaManager.becomeLeaderOrFollower(correlationId, leaderAndIsrRequest0,
+      (_, followers) => assertEquals(followerBrokerId, followers.head.partitionId))
+    assertTrue(countDownLatch.await(1000L, TimeUnit.MILLISECONDS))
+
+    // Truncation should have happened once
+    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(new MockTimer(time),
+      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 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(
+      partition, 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(new MockTimer(time),
+      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,
+      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(new MockTimer(time),
+      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,
+      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
+  def testFollowerFetchWithDefaultSelectorNoForcedHwPropagation(): Unit = {
+    val topicPartition = 0
+    val followerBrokerId = 0
+    val leaderBrokerId = 1
+    val leaderEpoch = 1
+    val leaderEpochIncrement = 2
+    val countDownLatch = new CountDownLatch(1)
+    val timer = new MockTimer(time)
+
+    // Prepare the mocked components for the test
+    val (replicaManager, _) = prepareReplicaManagerAndLogManager(timer,
+      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,
+      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 simpleRecords = Seq(new SimpleRecord("a".getBytes), new SimpleRecord("b".getBytes))
+    val appendResult = appendRecords(replicaManager, tp0,
+      MemoryRecords.withRecords(CompressionType.NONE, simpleRecords.toSeq: _*), AppendOrigin.Client)
+
+    // Increment the hw in the leader by fetching from the last offset
+    val fetchOffset = simpleRecords.size
+    var followerResult = fetchAsFollower(replicaManager, tp0,
+      new PartitionData(fetchOffset, 0, 100000, Optional.empty()),
+      clientMetadata = None)
+    assertTrue(followerResult.isFired)
+    assertEquals(0, followerResult.assertFired.highWatermark)
+
+    assertTrue("Expected producer request to be acked", appendResult.isFired)
+
+    // Fetch from the same offset, no new data is expected and hence the fetch request should
+    // go to the purgatory
+    followerResult = fetchAsFollower(replicaManager, tp0,
+      new PartitionData(fetchOffset, 0, 100000, Optional.empty()),
+      clientMetadata = None, minBytes = 1000)
+    assertFalse("Request completed immediately unexpectedly", followerResult.isFired)
+
+    // Complete the request in the purgatory by advancing the clock
+    timer.advanceClock(1001)
+    assertTrue(followerResult.isFired)
+
+    assertEquals(fetchOffset, followerResult.assertFired.highWatermark)
+  }
+
+  @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(new MockTimer(time),
+      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(new MockTimer(time),
+      topicPartition, leaderEpoch + leaderEpochIncrement, followerBrokerId,
+      leaderBrokerId, countDownLatch, expectTruncation = true)
+    assertFalse(replicaManager.replicaSelectorOpt.isDefined)
+  }
+
+  @Test
+  def testFetchFollowerNotAllowedForOlderClients(): Unit = {
+    val replicaManager = setupReplicaManagerWithMockedPurgatories(new MockTimer(time), 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,
+      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_OR_FOLLOWER, fetchResult.get.error)
+  }
+
+  @Test
+  def testFetchRequestRateMetrics(): Unit = {
+    val mockTimer = new MockTimer(time)
+    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,
+      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, (_, _) => ())
+
+    def assertMetricCount(expected: Int): Unit = {
+      assertEquals(expected, replicaManager.brokerTopicStats.allTopicsStats.totalFetchRequestRate.count)
+      assertEquals(expected, replicaManager.brokerTopicStats.topicStats(topic).totalFetchRequestRate.count)
+    }
+
+    val partitionData = new FetchRequest.PartitionData(0L, 0L, 100,
+      Optional.empty())
+
+    val nonPurgatoryFetchResult = sendConsumerFetch(replicaManager, tp0, partitionData, None, timeout = 0)
+    assertNotNull(nonPurgatoryFetchResult.get)
+    assertEquals(Errors.NONE, nonPurgatoryFetchResult.get.error)
+    assertMetricCount(1)
+
+    val purgatoryFetchResult = sendConsumerFetch(replicaManager, tp0, partitionData, None, timeout = 10)
+    assertNull(purgatoryFetchResult.get)
+    mockTimer.advanceClock(11)
+    assertNotNull(purgatoryFetchResult.get)
+    assertEquals(Errors.NONE, purgatoryFetchResult.get.error)
+    assertMetricCount(2)
+  }
+
+  @Test
+  def testBecomeFollowerWhileOldClientFetchInPurgatory(): Unit = {
+    val mockTimer = new MockTimer(time)
+    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,
+      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,
+      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_OR_FOLLOWER, fetchResult.get.error)
+  }
+
+  @Test
+  def testBecomeFollowerWhileNewClientFetchInPurgatory(): Unit = {
+    val mockTimer = new MockTimer(time)
+    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,
+      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,
+      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(time), 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,
+      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(time)
+    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,
+      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.stopReplicas(2, 0, 0, brokerEpoch,
+      mutable.Map(tp0 -> new StopReplicaPartitionState()
+        .setPartitionIndex(tp0.partition)
+        .setDeletePartition(true)
+        .setLeaderEpoch(LeaderAndIsr.EpochDuringDelete)))
+
+    assertNotNull(fetchResult.get)
+    assertEquals(Errors.NOT_LEADER_OR_FOLLOWER, fetchResult.get.error)
+  }
+
+  @Test
+  def testClearProducePurgatoryOnStopReplica(): Unit = {
+    val mockTimer = new MockTimer(time)
+    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,
+      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.stopReplicas(2, 0, 0, brokerEpoch,
+      mutable.Map(tp0 -> new StopReplicaPartitionState()
+        .setPartitionIndex(tp0.partition)
+        .setDeletePartition(true)
+        .setLeaderEpoch(LeaderAndIsr.EpochDuringDelete)))
+
+    assertNotNull(produceResult.get)
+    assertEquals(Errors.NOT_LEADER_OR_FOLLOWER, 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
+   * 'leaderEpochInLeaderAndIsr' leader epoch for partition 'topicPartition'.
+   */
+  private def prepareReplicaManagerAndLogManager(timer: MockTimer,
+                                                 topicPartition: Int,
+                                                 leaderEpochInLeaderAndIsr: Int,
+                                                 followerBrokerId: Int,
+                                                 leaderBrokerId: Int,
+                                                 countDownLatch: CountDownLatch,
+                                                 expectTruncation: Boolean,
+                                                 localLogOffset: Option[Long] = None,
+                                                 offsetFromLeader: Long = 5,
+                                                 leaderEpochFromLeader: Int = 3,
+                                                 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)
+
+    val mockScheduler = new MockScheduler(time)
+    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"),
+      config = LogConfig(),
+      logStartOffset = 0L,
+      recoveryPoint = 0L,
+      scheduler = mockScheduler,
+      brokerTopicStats = mockBrokerTopicStats,
+      time = time,
+      maxProducerIdExpirationMs = 30000,
+      producerIdExpirationCheckIntervalMs = 30000,
+      topicPartition = new TopicPartition(topic, topicPartition),
+      producerStateManager = new ProducerStateManager(new TopicPartition(topic, topicPartition),
+        new File(new File(config.logDirs.head), s"$topic-$topicPartition"), 30000),
+      logDirFailureChannel = mockLogDirFailureChannel) {
+
+      override def endOffsetForEpoch(leaderEpoch: Int): Option[OffsetAndEpoch] = {
+        assertEquals(leaderEpoch, leaderEpochFromLeader)
+        localLogOffset.map { logOffset =>
+          Some(OffsetAndEpoch(logOffset, leaderEpochFromLeader))
+        }.getOrElse(super.endOffsetForEpoch(leaderEpoch))
+      }
+
+      override def latestEpoch: Option[Int] = Some(leaderEpochFromLeader)
+
+      override def logEndOffsetMetadata: LogOffsetMetadata =
+        localLogOffset.map(LogOffsetMetadata(_)).getOrElse(super.logEndOffsetMetadata)
+
+      override def logEndOffset: Long = localLogOffset.getOrElse(super.logEndOffset)
+    }
+
+    // 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.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(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)
+    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
+    aliveBrokerIds.foreach { brokerId =>
+      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 mockProducePurgatory = new DelayedOperationPurgatory[DelayedProduce](
+      purgatoryName = "Produce", timer, reaperEnabled = false)
+    val mockFetchPurgatory = new DelayedOperationPurgatory[DelayedFetch](
+      purgatoryName = "Fetch", timer, reaperEnabled = false)
+    val mockDeleteRecordsPurgatory = new DelayedOperationPurgatory[DelayedDeleteRecords](
+      purgatoryName = "DeleteRecords", timer, reaperEnabled = false)
+    val mockElectLeaderPurgatory = new DelayedOperationPurgatory[DelayedElectLeader](
+      purgatoryName = "ElectLeader", timer, reaperEnabled = false)
+
+    // Mock network client to show leader offset of 5
+    val blockingSend = new ReplicaFetcherMockBlockingSend(
+      Map(topicPartitionObj -> new EpochEndOffset()
+        .setPartition(topicPartitionObj.partition)
+        .setErrorCode(Errors.NONE.code)
+        .setLeaderEpoch(leaderEpochFromLeader)
+        .setEndOffset(offsetFromLeader)).asJava,
+      BrokerEndPoint(1, "host1" ,1), time)
+    val replicaManager = new ReplicaManager(config, metrics, time, kafkaZkClient, mockScheduler, mockLogMgr,
+      new AtomicBoolean(false), quotaManager, mockBrokerTopicStats,
+      metadataCache, mockLogDirFailureChannel, mockProducePurgatory, mockFetchPurgatory,
+      mockDeleteRecordsPurgatory, mockElectLeaderPurgatory, Option(this.getClass.getName), alterIsrManager) {
+
+      override protected def createReplicaFetcherManager(metrics: Metrics,
+                                                     time: Time,
+                                                     threadNamePrefix: Option[String],
+                                                     replicationQuotaManager: ReplicationQuotaManager): ReplicaFetcherManager = {
+        new ReplicaFetcherManager(config, this, metrics, time, threadNamePrefix, replicationQuotaManager) {
+
+          override def createFetcherThread(fetcherId: Int, sourceBroker: BrokerEndPoint): ReplicaFetcherThread = {
+            new ReplicaFetcherThread(s"ReplicaFetcherThread-$fetcherId", fetcherId,
+              sourceBroker, config, failedPartitions, replicaManager, metrics, time, quotaManager.follower, Some(blockingSend)) {
+
+              override def doWork() = {
+                // In case the thread starts before the partition is added by AbstractFetcherManager,
+                // add it here (it's a no-op if already added)
+                val initialOffset = InitialFetchState(
+                  leader = new BrokerEndPoint(0, "localhost", 9092),
+                  initOffset = 0L, currentLeaderEpoch = leaderEpochInLeaderAndIsr)
+                addPartitions(Map(new TopicPartition(topic, topicPartition) -> initialOffset))
+                super.doWork()
+
+                // Shut the thread down after one iteration to avoid double-counting truncations
+                initiateShutdown()
+                countDownLatch.countDown()
+              }
+            }
+          }
+        }
+      }
+    }
+
+    (replicaManager, mockLogMgr)
+  }
+
+  private def leaderAndIsrPartitionState(topicPartition: TopicPartition,
+                                         leaderEpoch: Int,
+                                         leaderBrokerId: Int,
+                                         aliveBrokerIds: Seq[Integer],
+                                         isNew: Boolean = false): LeaderAndIsrPartitionState = {
+    new LeaderAndIsrPartitionState()
+      .setTopicName(topic)
+      .setPartitionIndex(topicPartition.partition)
+      .setControllerEpoch(controllerEpoch)
+      .setLeader(leaderBrokerId)
+      .setLeaderEpoch(leaderEpoch)
+      .setIsr(aliveBrokerIds.asJava)
+      .setZkVersion(zkVersion)
+      .setReplicas(aliveBrokerIds.asJava)
+      .setIsNew(isNew)
+  }
+
   private class CallbackResult[T] {
     private var value: Option[T] = None
     private var fun: Option[T => Unit] = None
@@ -530,7 +1612,8 @@ class ReplicaManagerTest {
   private def appendRecords(replicaManager: ReplicaManager,
                             partition: TopicPartition,
                             records: MemoryRecords,
-                            isFromClient: Boolean = true): CallbackResult[PartitionResponse] = {
+                            origin: AppendOrigin = AppendOrigin.Client,
+                            requiredAcks: Short = -1): CallbackResult[PartitionResponse] = {
     val result = new CallbackResult[PartitionResponse]()
     def appendCallback(responses: Map[TopicPartition, PartitionResponse]): Unit = {
       val response = responses.get(partition)
@@ -540,9 +1623,9 @@ class ReplicaManagerTest {
 
     replicaManager.appendRecords(
       timeout = 1000,
-      requiredAcks = -1,
+      requiredAcks = requiredAcks,
       internalTopicsAllowed = false,
-      isFromClient = isFromClient,
+      origin = origin,
       entriesPerPartition = Map(partition -> records),
       responseCallback = appendCallback)
 
@@ -553,16 +1636,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,
@@ -570,7 +1655,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)
@@ -586,25 +1672,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 = 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)
@@ -612,11 +1703,477 @@ class ReplicaManagerTest {
       purgatoryName = "Fetch", timer, reaperEnabled = false)
     val mockDeleteRecordsPurgatory = new DelayedOperationPurgatory[DelayedDeleteRecords](
       purgatoryName = "DeleteRecords", 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,
+      new AtomicBoolean(false), quotaManager, new BrokerTopicStats,
       metadataCache, new LogDirFailureChannel(config.logDirs.size), mockProducePurgatory, mockFetchPurgatory,
-      mockDeleteRecordsPurgatory, Option(this.getClass.getName))
+      mockDeleteRecordsPurgatory, mockDelayedElectLeaderPurgatory, Option(this.getClass.getName), alterIsrManager)
+  }
+
+  @Test
+  def testOldLeaderLosesMetricsWhenReassignPartitions(): Unit = {
+    val controllerEpoch = 0
+    val leaderEpoch = 0
+    val leaderEpochIncrement = 1
+    val correlationId = 0
+    val controllerId = 0
+    val mockTopicStats1: BrokerTopicStats = EasyMock.mock(classOf[BrokerTopicStats])
+    val (rm0, rm1) = prepareDifferentReplicaManagers(EasyMock.mock(classOf[BrokerTopicStats]), mockTopicStats1)
+
+    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,
+        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,
+        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 mockTopicStats1: BrokerTopicStats = EasyMock.mock(classOf[BrokerTopicStats])
+    val (rm0, rm1) = prepareDifferentReplicaManagers(EasyMock.mock(classOf[BrokerTopicStats]), mockTopicStats1)
+
+    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,
+        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,
+        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 prepareDifferentReplicaManagers(brokerTopicStats1: BrokerTopicStats,
+                                              brokerTopicStats2: BrokerTopicStats): (ReplicaManager, ReplicaManager) = {
+    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 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), quotaManager,
+      brokerTopicStats1, metadataCache0, new LogDirFailureChannel(config0.logDirs.size), alterIsrManager)
+    val rm1 = new ReplicaManager(config1, metrics, time, kafkaZkClient, new MockScheduler(time), mockLogMgr1,
+      new AtomicBoolean(false), quotaManager,
+      brokerTopicStats2, metadataCache1, new LogDirFailureChannel(config1.logDirs.size), alterIsrManager)
+
+    (rm0, rm1)
+  }
+
+  @Test
+  def testStopReplicaWithStaleControllerEpoch(): Unit = {
+    val mockTimer = new MockTimer(time)
+    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 becomeLeaderRequest = new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, 0, 10, brokerEpoch,
+      Seq(leaderAndIsrPartitionState(tp0, 1, 0, Seq(0, 1), true)).asJava,
+      Set(new Node(0, "host1", 0), new Node(1, "host2", 1)).asJava
+    ).build()
+
+    replicaManager.becomeLeaderOrFollower(1, becomeLeaderRequest, (_, _) => ())
+
+    val partitionStates = Map(tp0 -> new StopReplicaPartitionState()
+      .setPartitionIndex(tp0.partition)
+      .setLeaderEpoch(1)
+      .setDeletePartition(false)
+    )
+
+    val (_, error) = replicaManager.stopReplicas(1, 0, 0, brokerEpoch, partitionStates)
+    assertEquals(Errors.STALE_CONTROLLER_EPOCH, error)
+  }
+
+  @Test
+  def testStopReplicaWithOfflinePartition(): Unit = {
+    val mockTimer = new MockTimer(time)
+    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 becomeLeaderRequest = new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, 0, 0, brokerEpoch,
+      Seq(leaderAndIsrPartitionState(tp0, 1, 0, Seq(0, 1), true)).asJava,
+      Set(new Node(0, "host1", 0), new Node(1, "host2", 1)).asJava
+    ).build()
+
+    replicaManager.becomeLeaderOrFollower(1, becomeLeaderRequest, (_, _) => ())
+    replicaManager.markPartitionOffline(tp0)
+
+    val partitionStates = Map(tp0 -> new StopReplicaPartitionState()
+      .setPartitionIndex(tp0.partition)
+      .setLeaderEpoch(1)
+      .setDeletePartition(false)
+    )
+
+    val (result, error) = replicaManager.stopReplicas(1, 0, 0, brokerEpoch, partitionStates)
+    assertEquals(Errors.NONE, error)
+    assertEquals(Map(tp0 -> Errors.KAFKA_STORAGE_ERROR), result)
+  }
+
+  @Test
+  def testStopReplicaWithInexistentPartition(): Unit = {
+    testStopReplicaWithInexistentPartition(false, false)
+  }
+
+  @Test
+  def testStopReplicaWithInexistentPartitionAndPartitionsDelete(): Unit = {
+    testStopReplicaWithInexistentPartition(true, false)
+  }
+
+  @Test
+  def testStopReplicaWithInexistentPartitionAndPartitionsDeleteAndIOException(): Unit = {
+    testStopReplicaWithInexistentPartition(true, true)
+  }
+
+  private def testStopReplicaWithInexistentPartition(deletePartitions: Boolean, throwIOException: Boolean): Unit = {
+    val mockTimer = new MockTimer(time)
+    val replicaManager = setupReplicaManagerWithMockedPurgatories(mockTimer, aliveBrokerIds = Seq(0, 1))
+
+    val tp0 = new TopicPartition(topic, 0)
+    val log = replicaManager.logManager.getOrCreateLog(tp0, () => LogConfig(), true)
+
+    if (throwIOException) {
+      // Delete the underlying directory to trigger an KafkaStorageException
+      val dir = log.dir.getParentFile
+      Utils.delete(dir)
+      dir.createNewFile()
+    }
+
+    val partitionStates = Map(tp0 -> new StopReplicaPartitionState()
+      .setPartitionIndex(tp0.partition)
+      .setLeaderEpoch(1)
+      .setDeletePartition(deletePartitions)
+    )
+
+    val (result, error) = replicaManager.stopReplicas(1, 0, 0, brokerEpoch, partitionStates)
+    assertEquals(Errors.NONE, error)
+
+    if (throwIOException && deletePartitions) {
+      assertEquals(Map(tp0 -> Errors.KAFKA_STORAGE_ERROR), result)
+      assertTrue(replicaManager.logManager.getLog(tp0).isEmpty)
+    } else if (deletePartitions) {
+      assertEquals(Map(tp0 -> Errors.NONE), result)
+      assertTrue(replicaManager.logManager.getLog(tp0).isEmpty)
+    } else {
+      assertEquals(Map(tp0 -> Errors.NONE), result)
+      assertTrue(replicaManager.logManager.getLog(tp0).isDefined)
+    }
+  }
+
+  @Test
+  def testStopReplicaWithExistingPartitionAndNewerLeaderEpoch(): Unit = {
+    testStopReplicaWithExistingPartition(2, false, false, Errors.NONE)
+  }
+
+  @Test
+  def testStopReplicaWithExistingPartitionAndOlderLeaderEpoch(): Unit = {
+    testStopReplicaWithExistingPartition(0, false, false, Errors.FENCED_LEADER_EPOCH)
+  }
+
+  @Test
+  def testStopReplicaWithExistingPartitionAndEqualLeaderEpoch(): Unit = {
+    testStopReplicaWithExistingPartition(1, false, false, Errors.FENCED_LEADER_EPOCH)
+  }
+
+  @Test
+  def testStopReplicaWithExistingPartitionAndDeleteSentinel(): Unit = {
+    testStopReplicaWithExistingPartition(LeaderAndIsr.EpochDuringDelete, false, false, Errors.NONE)
+  }
+
+  @Test
+  def testStopReplicaWithExistingPartitionAndLeaderEpochNotProvided(): Unit = {
+    testStopReplicaWithExistingPartition(LeaderAndIsr.NoEpoch, false, false, Errors.NONE)
+  }
+
+  @Test
+  def testStopReplicaWithDeletePartitionAndExistingPartitionAndNewerLeaderEpoch(): Unit = {
+    testStopReplicaWithExistingPartition(2, true, false, Errors.NONE)
+  }
+
+  @Test
+  def testStopReplicaWithDeletePartitionAndExistingPartitionAndNewerLeaderEpochAndIOException(): Unit = {
+    testStopReplicaWithExistingPartition(2, true, true, Errors.KAFKA_STORAGE_ERROR)
+  }
+
+  @Test
+  def testStopReplicaWithDeletePartitionAndExistingPartitionAndOlderLeaderEpoch(): Unit = {
+    testStopReplicaWithExistingPartition(0, true, false, Errors.FENCED_LEADER_EPOCH)
+  }
+
+  @Test
+  def testStopReplicaWithDeletePartitionAndExistingPartitionAndEqualLeaderEpoch(): Unit = {
+    testStopReplicaWithExistingPartition(1, true, false, Errors.FENCED_LEADER_EPOCH)
+  }
+
+  @Test
+  def testStopReplicaWithDeletePartitionAndExistingPartitionAndDeleteSentinel(): Unit = {
+    testStopReplicaWithExistingPartition(LeaderAndIsr.EpochDuringDelete, true, false, Errors.NONE)
+  }
+
+  @Test
+  def testStopReplicaWithDeletePartitionAndExistingPartitionAndLeaderEpochNotProvided(): Unit = {
+    testStopReplicaWithExistingPartition(LeaderAndIsr.NoEpoch, true, false, Errors.NONE)
+  }
+
+  private def testStopReplicaWithExistingPartition(leaderEpoch: Int,
+                                                   deletePartition: Boolean,
+                                                   throwIOException: Boolean,
+                                                   expectedOutput: Errors): Unit = {
+    val mockTimer = new MockTimer(time)
+    val replicaManager = setupReplicaManagerWithMockedPurgatories(mockTimer, aliveBrokerIds = Seq(0, 1))
+
+    val tp0 = new TopicPartition(topic, 0)
+    val offsetCheckpoints = new LazyOffsetCheckpoints(replicaManager.highWatermarkCheckpoints)
+    val partition = replicaManager.createPartition(tp0)
+    partition.createLogIfNotExists(isNew = false, isFutureReplica = false, offsetCheckpoints)
+
+    val logDirFailureChannel = new LogDirFailureChannel(replicaManager.config.logDirs.size)
+    val logDir = partition.log.get.parentDirFile
+
+    def readRecoveryPointCheckpoint(): Map[TopicPartition, Long] = {
+      new OffsetCheckpointFile(new File(logDir, LogManager.RecoveryPointCheckpointFile),
+        logDirFailureChannel).read()
+    }
+
+    def readLogStartOffsetCheckpoint(): Map[TopicPartition, Long] = {
+      new OffsetCheckpointFile(new File(logDir, LogManager.LogStartOffsetCheckpointFile),
+        logDirFailureChannel).read()
+    }
+
+    val becomeLeaderRequest = new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, 0, 0, brokerEpoch,
+      Seq(leaderAndIsrPartitionState(tp0, 1, 0, Seq(0, 1), true)).asJava,
+      Set(new Node(0, "host1", 0), new Node(1, "host2", 1)).asJava
+    ).build()
+
+    replicaManager.becomeLeaderOrFollower(1, becomeLeaderRequest, (_, _) => ())
+
+    val batch = TestUtils.records(records = List(
+      new SimpleRecord(10, "k1".getBytes, "v1".getBytes),
+      new SimpleRecord(11, "k2".getBytes, "v2".getBytes)))
+    partition.appendRecordsToLeader(batch, AppendOrigin.Client, requiredAcks = 0)
+    partition.log.get.updateHighWatermark(2L)
+    partition.log.get.maybeIncrementLogStartOffset(1L, LeaderOffsetIncremented)
+    replicaManager.logManager.checkpointLogRecoveryOffsets()
+    replicaManager.logManager.checkpointLogStartOffsets()
+    assertEquals(Some(1L), readRecoveryPointCheckpoint().get(tp0))
+    assertEquals(Some(1L), readLogStartOffsetCheckpoint().get(tp0))
+
+    if (throwIOException) {
+      // Delete the underlying directory to trigger an KafkaStorageException
+      val dir = partition.log.get.dir
+      Utils.delete(dir)
+      dir.createNewFile()
+    }
+
+    val partitionStates = Map(tp0 -> new StopReplicaPartitionState()
+      .setPartitionIndex(tp0.partition)
+      .setLeaderEpoch(leaderEpoch)
+      .setDeletePartition(deletePartition)
+    )
+
+    val (result, error) = replicaManager.stopReplicas(1, 0, 0, brokerEpoch, partitionStates)
+    assertEquals(Errors.NONE, error)
+    assertEquals(Map(tp0 -> expectedOutput), result)
+
+    if (expectedOutput == Errors.NONE && deletePartition) {
+      assertEquals(HostedPartition.None, replicaManager.getPartition(tp0))
+      assertFalse(readRecoveryPointCheckpoint().contains(tp0))
+      assertFalse(readLogStartOffsetCheckpoint().contains(tp0))
+    }
+  }
+
+  @Test
+  def testReplicaNotAvailable(): Unit = {
+
+    def createReplicaManager(): ReplicaManager = {
+      val props = TestUtils.createBrokerConfig(1, TestUtils.MockZkConnect)
+      val config = KafkaConfig.fromProps(props)
+      val mockLogMgr = TestUtils.createLogManager(config.logDirs.map(new File(_)))
+      new ReplicaManager(config, metrics, time, kafkaZkClient, new MockScheduler(time), mockLogMgr,
+        new AtomicBoolean(false), quotaManager, new BrokerTopicStats,
+        new MetadataCache(config.brokerId), new LogDirFailureChannel(config.logDirs.size), alterIsrManager) {
+        override def getPartitionOrException(topicPartition: TopicPartition): Partition = {
+          throw Errors.NOT_LEADER_OR_FOLLOWER.exception()
+        }
+      }
+    }
+
+    val replicaManager = createReplicaManager()
+    try {
+      val tp = new TopicPartition(topic, 0)
+      val dir = replicaManager.logManager.liveLogDirs.head.getAbsolutePath
+      val errors = replicaManager.alterReplicaLogDirs(Map(tp -> dir))
+      assertEquals(Errors.REPLICA_NOT_AVAILABLE, errors(tp))
+    } finally {
+      replicaManager.shutdown(false)
+    }
+  }
 }
diff --git a/core/src/test/scala/unit/kafka/server/ReplicationQuotaManagerTest.scala b/core/src/test/scala/unit/kafka/server/ReplicationQuotaManagerTest.scala
index 54b506d768a6e..c8c508108cdf6 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._
+import scala.jdk.CollectionConverters._
 
 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,14 +49,13 @@ class ReplicationQuotaManagerTest {
 
   @Test
   def shouldExceedQuotaThenReturnBackBelowBoundAsTimePasses(): Unit = {
-    val metrics = newMetrics()
     val quota = new ReplicationQuotaManager(ReplicationQuotaManagerConfig(numQuotaSamples = 10, quotaWindowSizeSeconds = 1), metrics, LeaderReplication, time)
 
     //Given
     quota.updateQuota(new Quota(100, true))
 
     //Quota should not be broken when we start
-    assertFalse(quota.isQuotaExceeded())
+    assertFalse(quota.isQuotaExceeded)
 
     //First window is fixed, so we'll skip it
     time.sleep(1000)
@@ -60,24 +65,24 @@ class ReplicationQuotaManagerTest {
     quota.record(1)
 
     //Then it should not break the quota
-    assertFalse(quota.isQuotaExceeded())
+    assertFalse(quota.isQuotaExceeded)
 
     //When we record half the quota (half way through the window), we still should not break
     quota.record(149) //150B, 1.5s
-    assertFalse(quota.isQuotaExceeded())
+    assertFalse(quota.isQuotaExceeded)
 
     //Add a byte to push over quota
     quota.record(1) //151B, 1.5s
 
     //Then it should break the quota
     assertEquals(151 / 1.5, rate(metrics), 0) //151B, 1.5s
-    assertTrue(quota.isQuotaExceeded())
+    assertTrue(quota.isQuotaExceeded)
 
     //When we sleep for the remaining half the window
     time.sleep(500) //151B, 2s
 
     //Then Our rate should have halved (i.e back down below the quota)
-    assertFalse(quota.isQuotaExceeded())
+    assertFalse(quota.isQuotaExceeded)
     assertEquals(151d / 2, rate(metrics), 0.1) //151B, 2s
 
     //When we sleep for another half a window (now half way through second window)
@@ -86,26 +91,26 @@ class ReplicationQuotaManagerTest {
 
     //Then the rate should be exceeded again
     assertEquals(250 / 2.5, rate(metrics), 0) //250B, 2.5s
-    assertFalse(quota.isQuotaExceeded())
+    assertFalse(quota.isQuotaExceeded)
     quota.record(1)
-    assertTrue(quota.isQuotaExceeded())
+    assertTrue(quota.isQuotaExceeded)
     assertEquals(251 / 2.5, rate(metrics), 0)
 
     //Sleep for 2 more window
     time.sleep(2 * 1000) //so now at 3.5s
-    assertFalse(quota.isQuotaExceeded())
+    assertFalse(quota.isQuotaExceeded)
     assertEquals(251d / 4.5, rate(metrics), 0)
   }
 
   def rate(metrics: Metrics): Double = {
     val metricName = metrics.metricName("byte-rate", LeaderReplication.toString, "Tracking byte-rate for " + LeaderReplication)
-    val leaderThrottledRate = metrics.metrics.asScala(metricName).value()
+    val leaderThrottledRate = metrics.metrics.asScala(metricName).metricValue.asInstanceOf[Double]
     leaderThrottledRate
   }
 
   @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 45a6bddc7d11f..e8518ec229e54 100644
--- a/core/src/test/scala/unit/kafka/server/ReplicationQuotasTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ReplicationQuotasTest.scala
@@ -24,12 +24,14 @@ import kafka.server.KafkaConfig.fromProps
 import kafka.server.QuotaType._
 import kafka.utils.TestUtils._
 import kafka.utils.CoreUtils._
+import kafka.utils.TestUtils
 import kafka.zk.ZooKeeperTestHarness
 import org.apache.kafka.clients.producer.{KafkaProducer, ProducerRecord}
 import org.apache.kafka.common.TopicPartition
 import org.junit.Assert._
 import org.junit.{After, Test}
-import scala.collection.JavaConverters._
+
+import scala.jdk.CollectionConverters._
 
 /**
   * This is the main test which ensure Replication Quotas work correctly.
@@ -41,7 +43,7 @@ import scala.collection.JavaConverters._
   * Anything over 100MB/s tends to fail as this is the non-throttled replication rate
   */
 class ReplicationQuotasTest extends ZooKeeperTestHarness {
-  def percentError(percent: Int, value: Long): Long = Math.round(value * percent / 100)
+  def percentError(percent: Int, value: Long): Long = Math.round(value * percent / 100.0)
 
   val msg100KB = new Array[Byte](100000)
   var brokers: Seq[KafkaServer] = null
@@ -49,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()
@@ -78,7 +80,7 @@ class ReplicationQuotasTest extends ZooKeeperTestHarness {
 
     //Given six partitions, led on nodes 0,1,2,3,4,5 but with followers on node 6,7 (not started yet)
     //And two extra partitions 6,7, which we don't intend on throttling.
-    adminZkClient.createOrUpdateTopicPartitionAssignmentPathInZK(topic, Map(
+    val assignment = Map(
       0 -> Seq(100, 106), //Throttled
       1 -> Seq(101, 106), //Throttled
       2 -> Seq(102, 106), //Throttled
@@ -87,7 +89,8 @@ class ReplicationQuotasTest extends ZooKeeperTestHarness {
       5 -> Seq(105, 107), //Throttled
       6 -> Seq(100, 106), //Not Throttled
       7 -> Seq(101, 107) //Not Throttled
-    ))
+    )
+    TestUtils.createTopic(zkClient, topic, assignment, brokers)
 
     val msg = msg100KB
     val msgCount = 100
@@ -111,7 +114,7 @@ class ReplicationQuotasTest extends ZooKeeperTestHarness {
       adminZkClient.changeTopicConfig(topic, propsWith(FollowerReplicationThrottledReplicasProp, "0:106,1:106,2:106,3:107,4:107,5:107"))
 
     //Add data equally to each partition
-    producer = createNewProducer(getBrokerListStrFromServers(brokers), retries = 5, acks = 1)
+    producer = createProducer(getBrokerListStrFromServers(brokers), acks = 1)
     (0 until msgCount).foreach { _ =>
       (0 to 7).foreach { partition =>
         producer.send(new ProducerRecord(topic, partition, null, msg))
@@ -131,7 +134,7 @@ class ReplicationQuotasTest extends ZooKeeperTestHarness {
 
     //Check that throttled config correctly migrated to the new brokers
     (106 to 107).foreach { brokerId =>
-      assertEquals(throttle, brokerFor(brokerId).quotaManagers.follower.upperBound())
+      assertEquals(throttle, brokerFor(brokerId).quotaManagers.follower.upperBound)
     }
     if (!leaderThrottle) {
       (0 to 2).foreach { partition => assertTrue(brokerFor(106).quotaManagers.follower.isThrottled(tp(partition))) }
@@ -176,7 +179,7 @@ class ReplicationQuotasTest extends ZooKeeperTestHarness {
     val config: Properties = createBrokerConfig(100, zkConnect)
     config.put("log.segment.bytes", (1024 * 1024).toString)
     brokers = Seq(createServer(fromProps(config)))
-    adminZkClient.createOrUpdateTopicPartitionAssignmentPathInZK(topic, Map(0 -> Seq(100, 101)))
+    TestUtils.createTopic(zkClient, topic, Map(0 -> Seq(100, 101)), brokers)
 
     //Write 20MBs and throttle at 5MB/s
     val msg = msg100KB
@@ -191,23 +194,23 @@ class ReplicationQuotasTest extends ZooKeeperTestHarness {
     //Add data
     addData(msgCount, msg)
 
-    val start = System.currentTimeMillis()
-
     //Start the new broker (and hence start replicating)
     debug("Starting new broker")
     brokers = brokers :+ createServer(fromProps(createBrokerConfig(101, zkConnect)))
+    val start = System.currentTimeMillis()
+
     waitForOffsetsToMatch(msgCount, 0, 101)
 
     val throttledTook = System.currentTimeMillis() - start
 
-    assertTrue((s"Throttled replication of ${throttledTook}ms should be > ${expectedDuration * 1000 * 0.9}ms"),
+    assertTrue(s"Throttled replication of ${throttledTook}ms should be > ${expectedDuration * 1000 * 0.9}ms",
       throttledTook > expectedDuration * 1000 * 0.9)
-    assertTrue((s"Throttled replication of ${throttledTook}ms should be < ${expectedDuration * 1500}ms"),
+    assertTrue(s"Throttled replication of ${throttledTook}ms should be < ${expectedDuration * 1500}ms",
       throttledTook < expectedDuration * 1000 * 1.5)
   }
 
   def addData(msgCount: Int, msg: Array[Byte]): Unit = {
-    producer = createNewProducer(getBrokerListStrFromServers(brokers), retries = 5, acks = 0)
+    producer = createProducer(getBrokerListStrFromServers(brokers), acks = 0)
     (0 until msgCount).map(_ => producer.send(new ProducerRecord(topic, msg))).foreach(_.get)
     waitForOffsetsToMatch(msgCount, 0, 100)
   }
@@ -233,6 +236,6 @@ class ReplicationQuotasTest extends ZooKeeperTestHarness {
 
   private def measuredRate(broker: KafkaServer, repType: QuotaType): Double = {
     val metricName = broker.metrics.metricName("byte-rate", repType.toString)
-    broker.metrics.metrics.asScala(metricName).value
+    broker.metrics.metrics.asScala(metricName).metricValue.asInstanceOf[Double]
   }
 }
diff --git a/core/src/test/scala/unit/kafka/server/RequestQuotaTest.scala b/core/src/test/scala/unit/kafka/server/RequestQuotaTest.scala
index ee75933fbefe8..bf22a40d7c5cb 100644
--- a/core/src/test/scala/unit/kafka/server/RequestQuotaTest.scala
+++ b/core/src/test/scala/unit/kafka/server/RequestQuotaTest.scala
@@ -14,52 +14,67 @@
 
 package kafka.server
 
-import java.net.Socket
-import java.nio.ByteBuffer
-import java.util.{Collections, LinkedHashMap, Properties}
+import java.net.InetAddress
+import java.util
 import java.util.concurrent.{Executors, Future, TimeUnit}
+import java.util.{Collections, Optional, Properties}
 
+import kafka.api.LeaderAndIsr
 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.clients.admin.NewPartitions
-import org.apache.kafka.common.acl.{AccessControlEntry, AccessControlEntryFilter, AclBinding, AclBindingFilter, AclOperation, AclPermissionType}
-import org.apache.kafka.common.resource.{ResourceFilter, Resource => AdminResource, ResourceType => AdminResourceType}
-import org.apache.kafka.common.{Node, TopicPartition}
+import org.apache.kafka.common.acl._
+import org.apache.kafka.common.config.ConfigResource
+import org.apache.kafka.common.message.CreatePartitionsRequestData.CreatePartitionsTopic
+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.ListOffsetRequestData.{ListOffsetPartition, ListOffsetTopic}
+import org.apache.kafka.common.message.OffsetForLeaderEpochRequestData.OffsetForLeaderPartition
+import org.apache.kafka.common.message.OffsetForLeaderEpochRequestData.OffsetForLeaderTopic
+import org.apache.kafka.common.message.OffsetForLeaderEpochRequestData.OffsetForLeaderTopicCollection
+import org.apache.kafka.common.message.StopReplicaRequestData.{StopReplicaPartitionState, StopReplicaTopicState}
+import org.apache.kafka.common.message.UpdateMetadataRequestData.{UpdateMetadataBroker, UpdateMetadataEndpoint, UpdateMetadataPartitionState}
+import org.apache.kafka.common.message.{AddOffsetsToTxnRequestData, _}
 import org.apache.kafka.common.metrics.{KafkaMetric, Quota, Sensor}
 import org.apache.kafka.common.network.ListenerName
 import org.apache.kafka.common.protocol.ApiKeys
-import org.apache.kafka.common.protocol.types.Struct
+import org.apache.kafka.common.quota.ClientQuotaFilter
 import org.apache.kafka.common.record._
-import org.apache.kafka.common.requests.CreateAclsRequest.AclCreation
-import org.apache.kafka.common.requests.{Resource => RResource, ResourceType => RResourceType, _}
-import org.apache.kafka.common.security.auth.{AuthenticationContext, KafkaPrincipal, KafkaPrincipalBuilder, SecurityProtocol}
-import org.apache.kafka.common.utils.Sanitizer
+import org.apache.kafka.common.requests._
+import org.apache.kafka.common.resource.{PatternType, ResourceType => AdminResourceType}
+import org.apache.kafka.common.security.auth._
+import org.apache.kafka.common.utils.{Sanitizer, SecurityUtils}
+import org.apache.kafka.common._
+import org.apache.kafka.server.authorizer.{Action, AuthorizableRequestContext, AuthorizationResult}
 import org.junit.Assert._
 import org.junit.{After, Before, Test}
 
-import scala.collection.JavaConverters._
 import scala.collection.mutable.ListBuffer
+import scala.jdk.CollectionConverters._
 
 class RequestQuotaTest extends BaseRequestTest {
 
-  override def numBrokers: Int = 1
+  override def brokerCount: Int = 1
 
   private val topic = "topic-1"
   private val numPartitions = 1
   private val tp = new TopicPartition(topic, 0)
   private val logDir = "logDir"
   private val unthrottledClientId = "unthrottled-client"
+  private val smallQuotaProducerClientId = "small-quota-producer-client"
+  private val smallQuotaConsumerClientId = "small-quota-consumer-client"
   private val brokerId: Integer = 0
-  private var leaderNode: KafkaServer = null
+  private var leaderNode: KafkaServer = _
 
   // Run tests concurrently since a throttle could be up to 1 second because quota percentage allocated is very low
   case class Task(apiKey: ApiKeys, future: Future[_])
   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")
@@ -70,35 +85,51 @@ class RequestQuotaTest extends BaseRequestTest {
   }
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     RequestQuotaTest.principal = KafkaPrincipal.ANONYMOUS
     super.setUp()
 
-    TestUtils.createTopic(zkUtils, topic, numPartitions, 1, servers)
+    createTopic(topic, numPartitions)
     leaderNode = servers.head
 
     // Change default client-id request quota to a small value and a single unthrottledClient with a large quota
     val quotaProps = new Properties()
     quotaProps.put(DynamicConfig.Client.RequestPercentageOverrideProp, "0.01")
+    quotaProps.put(DynamicConfig.Client.ProducerByteRateOverrideProp, "2000")
+    quotaProps.put(DynamicConfig.Client.ConsumerByteRateOverrideProp, "2000")
     adminZkClient.changeClientIdConfig("", quotaProps)
     quotaProps.put(DynamicConfig.Client.RequestPercentageOverrideProp, "2000")
     adminZkClient.changeClientIdConfig(Sanitizer.sanitize(unthrottledClientId), quotaProps)
 
-    TestUtils.retry(10000) {
-      val quotaManager = servers.head.apis.quotas.request
+    // Client ids with small producer and consumer (fetch) quotas. Quota values were picked so that both
+    // producer/consumer and request quotas are violated on the first produce/consume operation, and the delay due to
+    // producer/consumer quota violation will be longer than the delay due to request quota violation.
+    quotaProps.put(DynamicConfig.Client.ProducerByteRateOverrideProp, "1")
+    quotaProps.put(DynamicConfig.Client.RequestPercentageOverrideProp, "0.01")
+    adminZkClient.changeClientIdConfig(Sanitizer.sanitize(smallQuotaProducerClientId), quotaProps)
+    quotaProps.put(DynamicConfig.Client.ConsumerByteRateOverrideProp, "1")
+    quotaProps.put(DynamicConfig.Client.RequestPercentageOverrideProp, "0.01")
+    adminZkClient.changeClientIdConfig(Sanitizer.sanitize(smallQuotaConsumerClientId), quotaProps)
+
+    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))
 
@@ -106,48 +137,65 @@ class RequestQuotaTest extends BaseRequestTest {
   }
 
   @Test
-  def testUnthrottledClient() {
-    for (apiKey <- RequestQuotaTest.ClientActions)
+  def testResponseThrottleTimeWhenBothProduceAndRequestQuotasViolated(): Unit = {
+    submitTest(ApiKeys.PRODUCE, () => checkSmallQuotaProducerRequestThrottleTime())
+    waitAndCheckResults()
+  }
+
+  @Test
+  def testResponseThrottleTimeWhenBothFetchAndRequestQuotasViolated(): Unit = {
+    submitTest(ApiKeys.FETCH, () => checkSmallQuotaConsumerRequestThrottleTime())
+    waitAndCheckResults()
+  }
+
+  @Test
+  def testUnthrottledClient(): Unit = {
+    for (apiKey <- RequestQuotaTest.ClientActions) {
       submitTest(apiKey, () => checkUnthrottledClient(apiKey))
+    }
 
     waitAndCheckResults()
   }
 
   @Test
-  def testExemptRequestTime() {
-    for (apiKey <- RequestQuotaTest.ClusterActions)
+  def testExemptRequestTime(): Unit = {
+    for (apiKey <- RequestQuotaTest.ClusterActions) {
       submitTest(apiKey, () => checkExemptRequestMetric(apiKey))
+    }
 
     waitAndCheckResults()
   }
 
   @Test
-  def testUnauthorizedThrottle() {
+  def testUnauthorizedThrottle(): Unit = {
     RequestQuotaTest.principal = RequestQuotaTest.UnauthorizedPrincipal
 
-    for (apiKey <- ApiKeys.values)
+    for (apiKey <- ApiKeys.enabledApis.asScala) {
       submitTest(apiKey, () => checkUnauthorizedRequestThrottle(apiKey))
+    }
 
     waitAndCheckResults()
   }
 
+  def session(user: String): Session = Session(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, user), null)
+
   private def throttleTimeMetricValue(clientId: String): Double = {
-    val metricName = leaderNode.metrics.metricName("throttle-time",
-                                  QuotaType.Request.toString,
-                                  "",
-                                  "user", "",
-                                  "client-id", clientId)
-    val sensor = leaderNode.quotaManagers.request.getOrCreateQuotaSensors("ANONYMOUS", clientId).throttleTimeSensor
+    throttleTimeMetricValueForQuotaType(clientId, QuotaType.Request)
+  }
+
+  private def throttleTimeMetricValueForQuotaType(clientId: String, quotaType: QuotaType): Double = {
+    val metricName = leaderNode.metrics.metricName("throttle-time", quotaType.toString,
+      "", "user", "", "client-id", clientId)
+    val sensor = leaderNode.quotaManagers.request.getOrCreateQuotaSensors(session("ANONYMOUS"),
+      clientId).throttleTimeSensor
     metricValue(leaderNode.metrics.metrics.get(metricName), sensor)
   }
 
   private def requestTimeMetricValue(clientId: String): Double = {
-    val metricName = leaderNode.metrics.metricName("request-time",
-                                  QuotaType.Request.toString,
-                                  "",
-                                  "user", "",
-                                  "client-id", clientId)
-    val sensor = leaderNode.quotaManagers.request.getOrCreateQuotaSensors("ANONYMOUS", clientId).quotaSensor
+    val metricName = leaderNode.metrics.metricName("request-time", QuotaType.Request.toString,
+      "", "user", "", "client-id", clientId)
+    val sensor = leaderNode.quotaManagers.request.getOrCreateQuotaSensors(session("ANONYMOUS"),
+      clientId).quotaSensor
     metricValue(leaderNode.metrics.metrics.get(metricName), sensor)
   }
 
@@ -158,178 +206,430 @@ class RequestQuotaTest extends BaseRequestTest {
 
   private def metricValue(metric: KafkaMetric, sensor: Sensor): Double = {
     sensor.synchronized {
-      if (metric == null) -1.0 else metric.value
+      if (metric == null) -1.0 else metric.metricValue.asInstanceOf[Double]
     }
   }
 
   private def requestBuilder(apiKey: ApiKeys): AbstractRequest.Builder[_ <: AbstractRequest] = {
     apiKey match {
         case ApiKeys.PRODUCE =>
-          ProduceRequest.Builder.forCurrentMagic(1, 5000,
-            collection.mutable.Map(tp -> MemoryRecords.withRecords(CompressionType.NONE, new SimpleRecord("test".getBytes))).asJava)
+          requests.ProduceRequest.forCurrentMagic(new ProduceRequestData()
+            .setTopicData(new ProduceRequestData.TopicProduceDataCollection(
+              Collections.singletonList(new ProduceRequestData.TopicProduceData()
+                .setName(tp.topic()).setPartitionData(Collections.singletonList(
+                new ProduceRequestData.PartitionProduceData()
+                  .setIndex(tp.partition())
+                  .setRecords(MemoryRecords.withRecords(CompressionType.NONE, new SimpleRecord("test".getBytes))))))
+                .iterator))
+            .setAcks(1.toShort)
+            .setTimeoutMs(5000))
 
         case ApiKeys.FETCH =>
-          val partitionMap = new LinkedHashMap[TopicPartition, FetchRequest.PartitionData]
-          partitionMap.put(tp, new FetchRequest.PartitionData(0, 0, 100))
+          val partitionMap = new util.LinkedHashMap[TopicPartition, FetchRequest.PartitionData]
+          partitionMap.put(tp, new FetchRequest.PartitionData(0, 0, 100, Optional.of(15)))
           FetchRequest.Builder.forConsumer(0, 0, partitionMap)
 
         case ApiKeys.METADATA =>
           new MetadataRequest.Builder(List(topic).asJava, true)
 
         case ApiKeys.LIST_OFFSETS =>
+          val topic = new ListOffsetTopic()
+            .setName(tp.topic)
+            .setPartitions(List(new ListOffsetPartition()
+              .setPartitionIndex(tp.partition)
+              .setTimestamp(0L)
+              .setCurrentLeaderEpoch(15)).asJava)
           ListOffsetRequest.Builder.forConsumer(false, IsolationLevel.READ_UNCOMMITTED)
-            .setTargetTimes(Map(tp -> (0L: java.lang.Long)).asJava)
+            .setTargetTimes(List(topic).asJava)
 
         case ApiKeys.LEADER_AND_ISR =>
-          new LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, brokerId, Int.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,
+            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(brokerId, Int.MaxValue, true, Set(tp).asJava)
+          val topicStates = Seq(
+            new StopReplicaTopicState()
+              .setTopicName(tp.topic())
+              .setPartitionStates(Seq(new StopReplicaPartitionState()
+                .setPartitionIndex(tp.partition())
+                .setLeaderEpoch(LeaderAndIsr.initialLeaderEpoch + 2)
+                .setDeletePartition(true)).asJava)
+          ).asJava
+          new StopReplicaRequest.Builder(ApiKeys.STOP_REPLICA.latestVersion, brokerId,
+            Int.MaxValue, Long.MaxValue, false, topicStates)
 
         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, 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, partitionState, brokers)
 
         case ApiKeys.CONTROLLED_SHUTDOWN =>
-          new ControlledShutdownRequest.Builder(brokerId, 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, "metadata")).asJava).
-            setMemberId("").setGenerationId(1).setRetentionTime(1000)
-
+          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)
+          new OffsetFetchRequest.Builder("test-group", false, List(tp).asJava, false)
 
         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("test-leave-group", "")
+          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(List("test-group").asJava)
+          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("PLAIN")
+          new SaslHandshakeRequest.Builder(new SaslHandshakeRequestData().setMechanism("PLAIN"))
 
         case ApiKeys.SASL_AUTHENTICATE =>
-          new SaslAuthenticateRequest.Builder(ByteBuffer.wrap(new Array[Byte](0)))
+          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(Map("topic-2" -> new CreateTopicsRequest.TopicDetails(1, 1.toShort)).asJava, 0)
+          new CreateTopicsRequest.Builder(
+            new CreateTopicsRequestData().setTopics(
+              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)
+          new DeleteRecordsRequest.Builder(
+            new DeleteRecordsRequestData()
+              .setTimeoutMs(5000)
+              .setTopics(Collections.singletonList(new DeleteRecordsRequestData.DeleteRecordsTopic()
+                .setName(tp.topic())
+                .setPartitions(Collections.singletonList(new DeleteRecordsRequestData.DeleteRecordsPartition()
+                  .setPartitionIndex(tp.partition())
+                  .setOffset(0L))))))
 
         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().add(tp, 0)
+          val epochs = new OffsetForLeaderTopicCollection()
+          epochs.add(new OffsetForLeaderTopic()
+            .setTopic(tp.topic())
+            .setPartitions(List(new OffsetForLeaderPartition()
+              .setPartition(tp.partition())
+              .setLeaderEpoch(0)
+              .setCurrentLeaderEpoch(15)).asJava))
+          OffsetsForLeaderEpochRequest.Builder.forConsumer(epochs)
 
         case ApiKeys.ADD_PARTITIONS_TO_TXN =>
           new AddPartitionsToTxnRequest.Builder("test-transactional-id", 1, 0, List(tp).asJava)
 
         case ApiKeys.ADD_OFFSETS_TO_TXN =>
-          new AddOffsetsToTxnRequest.Builder("test-transactional-id", 1, 0, "test-txn-group")
+          new AddOffsetsToTxnRequest.Builder(new AddOffsetsToTxnRequestData()
+            .setTransactionalId("test-transactional-id")
+            .setProducerId(1)
+            .setProducerEpoch(0)
+            .setGroupId("test-txn-group")
+          )
 
         case ApiKeys.END_TXN =>
-          new EndTxnRequest.Builder("test-transactional-id", 1, 0, TransactionResult.forId(false))
+          new EndTxnRequest.Builder(new EndTxnRequestData()
+            .setTransactionalId("test-transactional-id")
+            .setProducerId(1)
+            .setProducerEpoch(0)
+            .setCommitted(false)
+          )
 
         case ApiKeys.WRITE_TXN_MARKERS =>
-          new WriteTxnMarkersRequest.Builder(List.empty.asJava)
+          new WriteTxnMarkersRequest.Builder(ApiKeys.WRITE_TXN_MARKERS.latestVersion(), 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(
+            "test-transactional-id",
+            "test-txn-group",
+            2,
+            0,
+            Map.empty[TopicPartition, TxnOffsetCommitRequest.CommittedOffset].asJava,
+            false)
 
         case ApiKeys.DESCRIBE_ACLS =>
           new DescribeAclsRequest.Builder(AclBindingFilter.ANY)
 
         case ApiKeys.CREATE_ACLS =>
-          new CreateAclsRequest.Builder(Collections.singletonList(new AclCreation(new AclBinding(
-            new AdminResource(AdminResourceType.TOPIC, "mytopic"),
-            new AccessControlEntry("User:ANONYMOUS", "*", AclOperation.WRITE, AclPermissionType.DENY)))))
-
+          new CreateAclsRequest.Builder(new CreateAclsRequestData().setCreations(Collections.singletonList(
+            new CreateAclsRequestData.AclCreation()
+              .setResourceType(AdminResourceType.TOPIC.code)
+              .setResourceName("mytopic")
+              .setResourcePatternType(PatternType.LITERAL.code)
+              .setPrincipal("User:ANONYMOUS")
+              .setHost("*")
+              .setOperation(AclOperation.WRITE.code)
+              .setPermissionType(AclPermissionType.DENY.code))))
         case ApiKeys.DELETE_ACLS =>
-          new DeleteAclsRequest.Builder(Collections.singletonList(new AclBindingFilter(
-            new ResourceFilter(AdminResourceType.TOPIC, null),
-            new AccessControlEntryFilter("User:ANONYMOUS", "*", AclOperation.ANY, AclPermissionType.DENY))))
-
+          new DeleteAclsRequest.Builder(new DeleteAclsRequestData().setFilters(Collections.singletonList(
+            new DeleteAclsRequestData.DeleteAclsFilter()
+              .setResourceTypeFilter(AdminResourceType.TOPIC.code)
+              .setResourceNameFilter(null)
+              .setPatternTypeFilter(PatternType.LITERAL.code)
+              .setPrincipalFilter("User:ANONYMOUS")
+              .setHostFilter("*")
+              .setOperation(AclOperation.ANY.code)
+              .setPermissionType(AclPermissionType.DENY.code))))
         case ApiKeys.DESCRIBE_CONFIGS =>
-          new DescribeConfigsRequest.Builder(Collections.singleton(new RResource(RResourceType.TOPIC, tp.topic)))
+          new DescribeConfigsRequest.Builder(new DescribeConfigsRequestData()
+            .setResources(Collections.singletonList(new DescribeConfigsRequestData.DescribeConfigsResource()
+              .setResourceType(ConfigResource.Type.TOPIC.id)
+              .setResourceName(tp.topic))))
 
         case ApiKeys.ALTER_CONFIGS =>
           new AlterConfigsRequest.Builder(
-            Collections.singletonMap(new RResource(RResourceType.TOPIC, tp.topic),
+            Collections.singletonMap(new ConfigResource(ConfigResource.Type.TOPIC, tp.topic),
               new AlterConfigsRequest.Config(Collections.singleton(
                 new AlterConfigsRequest.ConfigEntry(LogConfig.MaxMessageBytesProp, "1000000")
               ))), true)
 
         case ApiKeys.ALTER_REPLICA_LOG_DIRS =>
-          new AlterReplicaLogDirsRequest.Builder(Collections.singletonMap(tp, logDir))
+          val dir = new AlterReplicaLogDirsRequestData.AlterReplicaLogDir()
+            .setPath(logDir)
+          dir.topics.add(new AlterReplicaLogDirsRequestData.AlterReplicaLogDirTopic()
+            .setName(tp.topic)
+            .setPartitions(Collections.singletonList(tp.partition)))
+          val data = new AlterReplicaLogDirsRequestData();
+          data.dirs.add(dir)
+          new AlterReplicaLogDirsRequest.Builder(data)
 
         case ApiKeys.DESCRIBE_LOG_DIRS =>
-          new DescribeLogDirsRequest.Builder(Collections.singleton(tp))
+          val data = new DescribeLogDirsRequestData()
+          data.topics.add(new DescribeLogDirsRequestData.DescribableLogDirTopic()
+            .setTopic(tp.topic)
+            .setPartitionIndex(Collections.singletonList(tp.partition)))
+          new DescribeLogDirsRequest.Builder(data)
 
         case ApiKeys.CREATE_PARTITIONS =>
-          new CreatePartitionsRequest.Builder(
-            Collections.singletonMap("topic-2", NewPartitions.increaseTo(1)), 0, false
+          val data = new CreatePartitionsRequestData()
+            .setTimeoutMs(0)
+            .setValidateOnly(false)
+          data.topics().add(new CreatePartitionsTopic().setName("topic-2").setCount(1))
+          new CreatePartitionsRequest.Builder(data)
+
+        case ApiKeys.CREATE_DELEGATION_TOKEN =>
+          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(
+              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(
+              new RenewDelegationTokenRequestData()
+                .setHmac("".getBytes)
+                .setRenewPeriodMs(1000L))
+
+        case ApiKeys.DELETE_GROUPS =>
+          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.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 ApiKeys.DESCRIBE_CLIENT_QUOTAS =>
+          new DescribeClientQuotasRequest.Builder(ClientQuotaFilter.all())
+
+        case ApiKeys.ALTER_CLIENT_QUOTAS =>
+          new AlterClientQuotasRequest.Builder(List.empty.asJava, false)
+
+        case ApiKeys.DESCRIBE_USER_SCRAM_CREDENTIALS =>
+          new DescribeUserScramCredentialsRequest.Builder(new DescribeUserScramCredentialsRequestData())
+
+        case ApiKeys.ALTER_USER_SCRAM_CREDENTIALS =>
+          new AlterUserScramCredentialsRequest.Builder(new AlterUserScramCredentialsRequestData())
+
+        case ApiKeys.VOTE =>
+          new VoteRequest.Builder(VoteRequest.singletonRequest(tp, 1, 2, 0, 10))
+
+        case ApiKeys.BEGIN_QUORUM_EPOCH =>
+          new BeginQuorumEpochRequest.Builder(BeginQuorumEpochRequest.singletonRequest(tp, 2, 5))
+
+        case ApiKeys.END_QUORUM_EPOCH =>
+          new EndQuorumEpochRequest.Builder(EndQuorumEpochRequest.singletonRequest(
+            tp, 10, 5, Collections.singletonList(3)))
+
+        case ApiKeys.ALTER_ISR =>
+          new AlterIsrRequest.Builder(new AlterIsrRequestData())
+
+        case ApiKeys.UPDATE_FEATURES =>
+          new UpdateFeaturesRequest.Builder(new UpdateFeaturesRequestData())
+
+        case ApiKeys.ENVELOPE =>
+          val requestHeader = new RequestHeader(
+            ApiKeys.ALTER_CLIENT_QUOTAS,
+            ApiKeys.ALTER_CLIENT_QUOTAS.latestVersion,
+            "client-id",
+            0
+          )
+          val embedRequestData = RequestTestUtils.serializeRequestWithHeader(requestHeader,
+            new AlterClientQuotasRequest.Builder(List.empty.asJava, false).build())
+          new EnvelopeRequest.Builder(embedRequestData, new Array[Byte](0),
+            InetAddress.getByName("192.168.1.1").getAddress)
+
         case _ =>
           throw new IllegalArgumentException("Unsupported API key " + apiKey)
     }
   }
 
-  private def requestResponse(socket: Socket, clientId: String, correlationId: Int, requestBuilder: AbstractRequest.Builder[_ <: AbstractRequest]): Struct = {
-    val apiKey = requestBuilder.apiKey
-    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)
-    apiKey.parseResponse(request.version, responseBuffer)
-  }
-
   case class Client(clientId: String, apiKey: ApiKeys) {
     var correlationId: Int = 0
-    val builder = requestBuilder(apiKey)
-    def runUntil(until: (Struct) => Boolean): Boolean = {
+    def runUntil(until: AbstractResponse => Boolean): Boolean = {
       val startMs = System.currentTimeMillis
       var done = false
       val socket = connect()
       try {
         while (!done && System.currentTimeMillis < startMs + 10000) {
           correlationId += 1
-          val response = requestResponse(socket, clientId, correlationId, builder)
+          val request = requestBuilder(apiKey).build()
+          val response = sendAndReceive[AbstractResponse](request, socket, clientId, Some(correlationId))
           done = until.apply(response)
         }
       } finally {
@@ -341,123 +641,124 @@ 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)
       } catch {
-        case e: Throwable => {
+        case e: Throwable =>
           error(s"Test failed for api-key ${task.apiKey} with exception $e")
           throw e
-        }
       }
     }
   }
 
-  private def responseThrottleTime(apiKey: ApiKeys, response: Struct): Int = {
-    apiKey match {
-      case ApiKeys.PRODUCE => new ProduceResponse(response).getThrottleTime
-      case ApiKeys.FETCH => new FetchResponse(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.JOIN_GROUP => new JoinGroupResponse(response).throttleTimeMs
-      case ApiKeys.HEARTBEAT => new HeartbeatResponse(response).throttleTimeMs
-      case ApiKeys.LEAVE_GROUP => new LeaveGroupResponse(response).throttleTimeMs
-      case ApiKeys.SYNC_GROUP => new SyncGroupResponse(response).throttleTimeMs
-      case ApiKeys.DESCRIBE_GROUPS => new DescribeGroupsResponse(response).throttleTimeMs
-      case ApiKeys.LIST_GROUPS => new ListGroupsResponse(response).throttleTimeMs
-      case ApiKeys.API_VERSIONS => new ApiVersionsResponse(response).throttleTimeMs
-      case ApiKeys.CREATE_TOPICS => new CreateTopicsResponse(response).throttleTimeMs
-      case ApiKeys.DELETE_TOPICS => new DeleteTopicsResponse(response).throttleTimeMs
-      case ApiKeys.DELETE_RECORDS => new DeleteRecordsResponse(response).throttleTimeMs
-      case ApiKeys.INIT_PRODUCER_ID => new InitProducerIdResponse(response).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.DESCRIBE_ACLS => new DescribeAclsResponse(response).throttleTimeMs
-      case ApiKeys.CREATE_ACLS => new CreateAclsResponse(response).throttleTimeMs
-      case ApiKeys.DELETE_ACLS => new DeleteAclsResponse(response).throttleTimeMs
-      case ApiKeys.DESCRIBE_CONFIGS => new DescribeConfigsResponse(response).throttleTimeMs
-      case ApiKeys.ALTER_CONFIGS => new AlterConfigsResponse(response).throttleTimeMs
-      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 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(_.throttleTimeMs > 0)
 
     assertTrue(s"Response not throttled: $client", throttled)
     assertTrue(s"Throttle time metrics not updated: $client" , throttleTimeMetricValue(clientId) > 0)
   }
 
-  private def checkUnthrottledClient(apiKey: ApiKeys) {
+  private def checkSmallQuotaProducerRequestThrottleTime(): Unit = {
+
+    // Request until throttled using client-id with default small producer quota
+    val smallQuotaProducerClient = Client(smallQuotaProducerClientId, ApiKeys.PRODUCE)
+    val throttled = smallQuotaProducerClient.runUntil(_.throttleTimeMs > 0)
+
+    assertTrue(s"Response not throttled: $smallQuotaProducerClient", throttled)
+    assertTrue(s"Throttle time metrics for produce quota not updated: $smallQuotaProducerClient",
+      throttleTimeMetricValueForQuotaType(smallQuotaProducerClientId, QuotaType.Produce) > 0)
+    assertTrue(s"Throttle time metrics for request quota updated: $smallQuotaProducerClient",
+      throttleTimeMetricValueForQuotaType(smallQuotaProducerClientId, QuotaType.Request).isNaN)
+  }
+
+  private def checkSmallQuotaConsumerRequestThrottleTime(): Unit = {
+
+    // Request until throttled using client-id with default small consumer quota
+    val smallQuotaConsumerClient =   Client(smallQuotaConsumerClientId, ApiKeys.FETCH)
+    val throttled = smallQuotaConsumerClient.runUntil(_.throttleTimeMs > 0)
+
+    assertTrue(s"Response not throttled: $smallQuotaConsumerClientId", throttled)
+    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): Unit = {
 
     // Test that request from client with large quota is not throttled
     val unthrottledClient = Client(unthrottledClientId, apiKey)
-    unthrottledClient.runUntil(response => responseThrottleTime(apiKey, response) <= 0.0)
+    unthrottledClient.runUntil(_.throttleTimeMs <= 0.0)
     assertEquals(1, unthrottledClient.correlationId)
-    assertTrue(s"Client should not have been throttled: $unthrottledClient", throttleTimeMetricValue(unthrottledClientId) <= 0.0)
+    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)
-    val updated = client.runUntil(response => exemptRequestMetricValue > exemptTarget)
+    val updated = client.runUntil(_ => exemptRequestMetricValue > exemptTarget)
 
     assertTrue(s"Exempt-request-time metric not updated: $client", updated)
-    assertTrue(s"Client should not have been throttled: $client", throttleTimeMetricValue(clientId) <= 0.0)
+    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)
+    val throttled = client.runUntil(_ => throttleTimeMetricValue(clientId) > 0.0)
     assertTrue(s"Unauthorized client should have been throttled: $client", throttled)
   }
 }
 
 object RequestQuotaTest {
-  val ClusterActions = ApiKeys.values.toSet.filter(apiKey => apiKey.clusterAction)
+  val ClusterActions = ApiKeys.enabledApis.asScala.filter(_.clusterAction).toSet
   val SaslActions = Set(ApiKeys.SASL_HANDSHAKE, ApiKeys.SASL_AUTHENTICATE)
-  val ClientActions = ApiKeys.values.toSet -- ClusterActions -- SaslActions
+  val ClientActions = ApiKeys.enabledApis.asScala.toSet -- ClusterActions -- SaslActions
 
   val UnauthorizedPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "Unauthorized")
   // 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 {
+  class TestPrincipalBuilder extends KafkaPrincipalBuilder with KafkaPrincipalSerde {
     override def build(context: AuthenticationContext): KafkaPrincipal = {
       principal
     }
+
+    override def serialize(principal: KafkaPrincipal): Array[Byte] = {
+      new Array[Byte](0)
+    }
+
+    override def deserialize(bytes: Array[Byte]): KafkaPrincipal = {
+      principal
+    }
   }
 }
diff --git a/core/src/test/scala/unit/kafka/server/SaslApiVersionsRequestTest.scala b/core/src/test/scala/unit/kafka/server/SaslApiVersionsRequestTest.scala
index 00b9934a72d46..ca276d982d538 100644
--- a/core/src/test/scala/unit/kafka/server/SaslApiVersionsRequestTest.scala
+++ b/core/src/test/scala/unit/kafka/server/SaslApiVersionsRequestTest.scala
@@ -18,24 +18,22 @@ package kafka.server
 
 import java.net.Socket
 import java.util.Collections
-
-import org.apache.kafka.common.protocol.{ApiKeys, Errors}
-import org.apache.kafka.common.requests.{ApiVersionsRequest, ApiVersionsResponse}
-import org.apache.kafka.common.requests.SaslHandshakeRequest
-import org.apache.kafka.common.requests.SaslHandshakeResponse
-import org.junit.{After, Before, Test}
-import org.junit.Assert._
 import kafka.api.{KafkaSasl, SaslSetup}
 import kafka.utils.JaasTestUtils
+import org.apache.kafka.common.message.SaslHandshakeRequestData
+import org.apache.kafka.common.protocol.{ApiKeys, Errors}
+import org.apache.kafka.common.requests.{ApiVersionsRequest, ApiVersionsResponse, SaslHandshakeRequest, SaslHandshakeResponse}
 import org.apache.kafka.common.security.auth.SecurityProtocol
+import org.junit.Assert._
+import org.junit.{After, Before, Test}
 
-class SaslApiVersionsRequestTest extends BaseRequestTest with SaslSetup {
+class SaslApiVersionsRequestTest extends AbstractApiVersionsRequestTest with SaslSetup {
   override protected def securityProtocol = SecurityProtocol.SASL_PLAINTEXT
   private val kafkaClientSaslMechanism = "PLAIN"
   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 = {
@@ -50,55 +48,52 @@ class SaslApiVersionsRequestTest extends BaseRequestTest with SaslSetup {
   }
 
   @Test
-  def testApiVersionsRequestBeforeSaslHandshakeRequest() {
-    val plaintextSocket = connect(protocol = securityProtocol)
+  def testApiVersionsRequestBeforeSaslHandshakeRequest(): Unit = {
+    val socket = connect()
     try {
-      val apiVersionsResponse = sendApiVersionsRequest(plaintextSocket, new ApiVersionsRequest.Builder().build(0))
-      ApiVersionsRequestTest.validateApiVersionsResponse(apiVersionsResponse)
-      sendSaslHandshakeRequestValidateResponse(plaintextSocket)
+      val apiVersionsResponse = sendAndReceive[ApiVersionsResponse](
+        new ApiVersionsRequest.Builder().build(0), socket)
+      validateApiVersionsResponse(apiVersionsResponse)
+      sendSaslHandshakeRequestValidateResponse(socket)
     } finally {
-      plaintextSocket.close()
+      socket.close()
     }
   }
 
   @Test
-  def testApiVersionsRequestAfterSaslHandshakeRequest() {
-    val plaintextSocket = connect(protocol = securityProtocol)
+  def testApiVersionsRequestAfterSaslHandshakeRequest(): Unit = {
+    val socket = connect()
     try {
-      sendSaslHandshakeRequestValidateResponse(plaintextSocket)
-      val response = sendApiVersionsRequest(plaintextSocket, new ApiVersionsRequest.Builder().build(0))
-      assertEquals(Errors.ILLEGAL_SASL_STATE, response.error)
+      sendSaslHandshakeRequestValidateResponse(socket)
+      val response = sendAndReceive[ApiVersionsResponse](
+        new ApiVersionsRequest.Builder().build(0), socket)
+      assertEquals(Errors.ILLEGAL_SASL_STATE.code, response.data.errorCode)
     } finally {
-      plaintextSocket.close()
+      socket.close()
     }
   }
 
   @Test
-  def testApiVersionsRequestWithUnsupportedVersion() {
-    val plaintextSocket = connect(protocol = securityProtocol)
+  def testApiVersionsRequestWithUnsupportedVersion(): Unit = {
+    val socket = connect()
     try {
-      val apiVersionsRequest = new ApiVersionsRequest(0)
-      val apiVersionsResponse = sendApiVersionsRequest(plaintextSocket, apiVersionsRequest, Some(Short.MaxValue))
-      assertEquals(Errors.UNSUPPORTED_VERSION, apiVersionsResponse.error)
-      val apiVersionsResponse2 = sendApiVersionsRequest(plaintextSocket, new ApiVersionsRequest.Builder().build(0))
-      ApiVersionsRequestTest.validateApiVersionsResponse(apiVersionsResponse2)
-      sendSaslHandshakeRequestValidateResponse(plaintextSocket)
+      val apiVersionsRequest = new ApiVersionsRequest.Builder().build(0)
+      val apiVersionsResponse = sendUnsupportedApiVersionRequest(apiVersionsRequest)
+      assertEquals(Errors.UNSUPPORTED_VERSION.code, apiVersionsResponse.data.errorCode)
+      val apiVersionsResponse2 = sendAndReceive[ApiVersionsResponse](
+        new ApiVersionsRequest.Builder().build(0), socket)
+      validateApiVersionsResponse(apiVersionsResponse2)
+      sendSaslHandshakeRequestValidateResponse(socket)
     } finally {
-      plaintextSocket.close()
+      socket.close()
     }
   }
 
-  private def sendApiVersionsRequest(socket: Socket, request: ApiVersionsRequest,
-                                     apiVersion: Option[Short] = None): ApiVersionsResponse = {
-    val response = sendAndReceive(request, ApiKeys.API_VERSIONS, socket, apiVersion)
-    ApiVersionsResponse.parse(response, request.version)
-  }
-
-  private def sendSaslHandshakeRequestValidateResponse(socket: Socket) {
-    val request = new SaslHandshakeRequest("PLAIN")
-    val response = sendAndReceive(request, ApiKeys.SASL_HANDSHAKE, socket)
-    val handshakeResponse = SaslHandshakeResponse.parse(response, request.version)
-    assertEquals(Errors.NONE, handshakeResponse.error)
-    assertEquals(Collections.singletonList("PLAIN"), handshakeResponse.enabledMechanisms)
+  private def sendSaslHandshakeRequestValidateResponse(socket: Socket): Unit = {
+    val request = new SaslHandshakeRequest(new SaslHandshakeRequestData().setMechanism("PLAIN"),
+      ApiKeys.SASL_HANDSHAKE.latestVersion)
+    val response = sendAndReceive[SaslHandshakeResponse](request, socket)
+    assertEquals(Errors.NONE, response.error)
+    assertEquals(Collections.singletonList("PLAIN"), response.enabledMechanisms)
   }
 }
diff --git a/core/src/test/scala/unit/kafka/server/ServerGenerateBrokerIdTest.scala b/core/src/test/scala/unit/kafka/server/ServerGenerateBrokerIdTest.scala
index 0ba133ff5f296..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,16 @@ 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
 
 class ServerGenerateBrokerIdTest extends ZooKeeperTestHarness {
   var props1: Properties = null
@@ -33,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)
@@ -42,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()
@@ -79,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
@@ -118,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()
@@ -135,21 +140,21 @@ 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)
-    intercept[RuntimeException] {
+    val serverB = new KafkaServer(configB, threadNamePrefix = Option(this.getClass.getName))
+    intercept[NodeExistsException] {
       serverB.startup()
     }
     servers = Seq(serverA)
@@ -163,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()
@@ -172,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 = {
@@ -187,12 +192,4 @@ class ServerGenerateBrokerIdTest extends ZooKeeperTestHarness {
     }
     true
   }
-
-  @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/server/ServerGenerateClusterIdTest.scala b/core/src/test/scala/unit/kafka/server/ServerGenerateClusterIdTest.scala
index 1ec80fa66021a..acf8eeba07410 100755
--- a/core/src/test/scala/unit/kafka/server/ServerGenerateClusterIdTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ServerGenerateClusterIdTest.scala
@@ -16,23 +16,33 @@
   */
 package kafka.server
 
+import java.io.File
+
+
+import scala.collection.Seq
 import scala.concurrent._
-import ExecutionContext.Implicits._
 import scala.concurrent.duration._
-import kafka.utils.{TestUtils, ZkUtils}
+import ExecutionContext.Implicits._
+
+import kafka.common.{InconsistentBrokerMetadataException, InconsistentClusterIdException}
+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
 
+
 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 +50,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(zkUtils.pathExists(ZkUtils.ClusterIdPath))
+    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
@@ -61,11 +71,11 @@ class ServerGenerateClusterIdTest extends ZooKeeperTestHarness {
     server1.shutdown()
 
     // Make sure that the cluster id is persistent.
-    assertTrue(zkUtils.pathExists(ZkUtils.ClusterIdPath))
-    assertEquals(zkUtils.getClusterId, Some(clusterIdOnFirstBoot))
+    assertTrue(zkClient.getClusterId.isDefined)
+    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
@@ -74,21 +84,21 @@ class ServerGenerateClusterIdTest extends ZooKeeperTestHarness {
     server1.shutdown()
 
     // Make sure that the cluster id is persistent after multiple reboots.
-    assertTrue(zkUtils.pathExists(ZkUtils.ClusterIdPath))
-    assertEquals(zkUtils.getClusterId, Some(clusterIdOnFirstBoot))
+    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 +117,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 +144,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 bcddd40315b80..436dc9e85d69c 100755
--- a/core/src/test/scala/unit/kafka/server/ServerShutdownTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ServerShutdownTest.scala
@@ -17,20 +17,30 @@
 package kafka.server
 
 import kafka.zk.ZooKeeperTestHarness
-import kafka.consumer.SimpleConsumer
 import kafka.utils.{CoreUtils, TestUtils}
 import kafka.utils.TestUtils._
-import kafka.api.FetchRequestBuilder
-import kafka.message.ByteBufferMessageSet
-import java.io.File
+import java.io.{DataInputStream, File}
+import java.net.ServerSocket
+import java.util.concurrent.{Executors, TimeUnit}
 
+import kafka.cluster.Broker
+import kafka.controller.{ControllerChannelManager, ControllerContext, StateChangeLogger}
 import kafka.log.LogManager
+import kafka.zookeeper.ZooKeeperClientTimeoutException
+import org.apache.kafka.clients.consumer.KafkaConsumer
 import org.apache.kafka.clients.producer.{KafkaProducer, ProducerRecord}
 import org.apache.kafka.common.errors.KafkaStorageException
-import org.apache.kafka.common.serialization.{IntegerSerializer, StringSerializer}
-import org.I0Itec.zkclient.exception.ZkException
+import org.apache.kafka.common.metrics.Metrics
+import org.apache.kafka.common.network.ListenerName
+import org.apache.kafka.common.protocol.ApiKeys
+import org.apache.kafka.common.requests.LeaderAndIsrRequest
+import org.apache.kafka.common.security.auth.SecurityProtocol
+import org.apache.kafka.common.serialization.{IntegerDeserializer, IntegerSerializer, StringDeserializer, StringSerializer}
+import org.apache.kafka.common.utils.Time
 import org.junit.{Before, Test}
 import org.junit.Assert._
+
+import scala.jdk.CollectionConverters._
 import scala.reflect.ClassTag
 
 class ServerShutdownTest extends ZooKeeperTestHarness {
@@ -41,29 +51,36 @@ 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.createNewProducer(
+      TestUtils.createProducer(
         TestUtils.getBrokerListStrFromServers(Seq(server)),
-        retries = 5,
         keySerializer = new IntegerSerializer,
         valueSerializer = new StringSerializer
       )
 
+    def createConsumer(server: KafkaServer): KafkaConsumer[Integer, String] =
+      TestUtils.createConsumer(
+        TestUtils.getBrokerListStrFromServers(Seq(server)),
+        securityProtocol = SecurityProtocol.PLAINTEXT,
+        keyDeserializer = new IntegerDeserializer,
+        valueDeserializer = new StringDeserializer
+      )
+
     var server = new KafkaServer(config, threadNamePrefix = Option(this.getClass.getName))
     server.startup()
     var producer = createProducer(server)
 
     // create topic
-    createTopic(zkUtils, topic, numPartitions = 1, replicationFactor = 1, servers = Seq(server))
+    createTopic(zkClient, topic, numPartitions = 1, replicationFactor = 1, servers = Seq(server))
 
     // send some messages
     sent1.map(value => producer.send(new ProducerRecord(topic, 0, value))).foreach(_.get)
@@ -85,35 +102,27 @@ class ServerShutdownTest extends ZooKeeperTestHarness {
     TestUtils.waitUntilMetadataIsPropagated(Seq(server), topic, 0)
 
     producer = createProducer(server)
-    val consumer = new SimpleConsumer(host, TestUtils.boundPort(server), 1000000, 64*1024, "")
+    val consumer = createConsumer(server)
+    consumer.subscribe(Seq(topic).asJava)
 
-    var fetchedMessage: ByteBufferMessageSet = null
-    while (fetchedMessage == null || fetchedMessage.validBytes == 0) {
-      val fetched = consumer.fetch(new FetchRequestBuilder().addFetch(topic, 0, 0, 10000).maxWait(0).build())
-      fetchedMessage = fetched.messageSet(topic, 0)
-    }
-    assertEquals(sent1, fetchedMessage.map(m => TestUtils.readString(m.message.payload)))
-    val newOffset = fetchedMessage.last.nextOffset
+    val consumerRecords = TestUtils.consumeRecords(consumer, sent1.size)
+    assertEquals(sent1, consumerRecords.map(_.value))
 
     // send some more messages
     sent2.map(value => producer.send(new ProducerRecord(topic, 0, value))).foreach(_.get)
 
-    fetchedMessage = null
-    while (fetchedMessage == null || fetchedMessage.validBytes == 0) {
-      val fetched = consumer.fetch(new FetchRequestBuilder().addFetch(topic, 0, newOffset, 10000).build())
-      fetchedMessage = fetched.messageSet(topic, 0)
-    }
-    assertEquals(sent2, fetchedMessage.map(m => TestUtils.readString(m.message.payload)))
+    val consumerRecords2 = TestUtils.consumeRecords(consumer, sent2.size)
+    assertEquals(sent2, consumerRecords2.map(_.value))
 
     consumer.close()
     producer.close()
     server.shutdown()
     CoreUtils.delete(server.config.logDirs)
-    verifyNonDaemonThreadsStatus
+    verifyNonDaemonThreadsStatus()
   }
 
   @Test
-  def testCleanShutdownWithDeleteTopicEnabled() {
+  def testCleanShutdownWithDeleteTopicEnabled(): Unit = {
     val newProps = TestUtils.createBrokerConfig(0, zkConnect)
     newProps.setProperty("delete.topic.enable", "true")
     val newConfig = KafkaConfig.fromProps(newProps)
@@ -122,22 +131,23 @@ class ServerShutdownTest extends ZooKeeperTestHarness {
     server.shutdown()
     server.awaitShutdown()
     CoreUtils.delete(server.config.logDirs)
-    verifyNonDaemonThreadsStatus
+    verifyNonDaemonThreadsStatus()
   }
 
   @Test
-  def testCleanShutdownAfterFailedStartup() {
+  def testCleanShutdownAfterFailedStartup(): Unit = {
     val newProps = TestUtils.createBrokerConfig(0, zkConnect)
-    newProps.setProperty("zookeeper.connect", "fakehostthatwontresolve:65535")
+    newProps.setProperty(KafkaConfig.ZkConnectionTimeoutMsProp, "50")
+    newProps.setProperty(KafkaConfig.ZkConnectProp, "some.invalid.hostname.foo.bar.local:65535")
     val newConfig = KafkaConfig.fromProps(newProps)
-    verifyCleanShutdownAfterFailedStartup[ZkException](newConfig)
+    verifyCleanShutdownAfterFailedStartup[ZooKeeperClientTimeoutException](newConfig)
   }
 
   @Test
-  def testCleanShutdownAfterFailedStartupDueToCorruptLogs() {
+  def testCleanShutdownAfterFailedStartupDueToCorruptLogs(): Unit = {
     val server = new KafkaServer(config)
     server.startup()
-    createTopic(zkUtils, topic, numPartitions = 1, replicationFactor = 1, servers = Seq(server))
+    createTopic(zkClient, topic, numPartitions = 1, replicationFactor = 1, servers = Seq(server))
     server.shutdown()
     server.awaitShutdown()
     config.logDirs.foreach { dirName =>
@@ -147,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()
@@ -167,31 +177,79 @@ class ServerShutdownTest extends ZooKeeperTestHarness {
       server.awaitShutdown()
     }
     CoreUtils.delete(server.config.logDirs)
-    verifyNonDaemonThreadsStatus
+    verifyNonDaemonThreadsStatus()
   }
 
   private[this] def isNonDaemonKafkaThread(t: Thread): Boolean = {
     !t.isDaemon && t.isAlive && t.getName.startsWith(this.getClass.getName)
   }
 
-  def verifyNonDaemonThreadsStatus() {
-    assertEquals(0, Thread.getAllStackTraces.keySet().toArray
-      .map{ _.asInstanceOf[Thread] }
+  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()
+    server.awaitShutdown()
+    server.shutdown()
+  }
+
+  // Verify that if controller is in the midst of processing a request, shutdown completes
+  // without waiting for request timeout.
+  @Test
+  def testControllerShutdownDuringSend(): Unit = {
+    val securityProtocol = SecurityProtocol.PLAINTEXT
+    val listenerName = ListenerName.forSecurityProtocol(securityProtocol)
+
+    val controllerId = 2
+    val metrics = new Metrics
+    val executor = Executors.newSingleThreadExecutor
+    var serverSocket: ServerSocket = null
+    var controllerChannelManager: ControllerChannelManager = null
+
     try {
-      server.startup()
-      server.shutdown()
-      server.awaitShutdown()
-      server.shutdown()
-      assertTrue(true)
-    }
-    catch{
-      case _: Throwable => fail()
+      // Set up a server to accept a connection and receive one byte from the first request. No response is sent.
+      serverSocket = new ServerSocket(0)
+      val receiveFuture = executor.submit(new Runnable {
+        override def run(): Unit = {
+          val socket = serverSocket.accept()
+          new DataInputStream(socket.getInputStream).readByte()
+        }
+      })
+
+      // Start a ControllerChannelManager
+      val brokerAndEpochs = Map((new Broker(1, "localhost", serverSocket.getLocalPort, listenerName, securityProtocol), 0L))
+      val controllerConfig = KafkaConfig.fromProps(TestUtils.createBrokerConfig(controllerId, zkConnect))
+      val controllerContext = new ControllerContext
+      controllerContext.setLiveBrokers(brokerAndEpochs)
+      controllerChannelManager = new ControllerChannelManager(controllerContext, controllerConfig, Time.SYSTEM,
+        metrics, new StateChangeLogger(controllerId, inControllerContext = true, None))
+      controllerChannelManager.startup()
+
+      // 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, 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
+      val shutdownFuture = executor.submit(new Runnable {
+        override def run(): Unit = controllerChannelManager.shutdown()
+      })
+      shutdownFuture.get(10, TimeUnit.SECONDS)
+
+    } finally {
+      if (serverSocket != null)
+        serverSocket.close()
+      if (controllerChannelManager != null)
+        controllerChannelManager.shutdown()
+      executor.shutdownNow()
+      metrics.close()
     }
   }
 }
diff --git a/core/src/test/scala/unit/kafka/server/ServerStartupTest.scala b/core/src/test/scala/unit/kafka/server/ServerStartupTest.scala
index c764369bf1e59..2448e5058a2d4 100755
--- a/core/src/test/scala/unit/kafka/server/ServerStartupTest.scala
+++ b/core/src/test/scala/unit/kafka/server/ServerStartupTest.scala
@@ -17,8 +17,10 @@
 
 package kafka.server
 
-import kafka.utils.{TestUtils, ZkUtils}
+import kafka.utils.TestUtils
 import kafka.zk.ZooKeeperTestHarness
+import org.apache.kafka.common.KafkaException
+import org.apache.zookeeper.KeeperException.NodeExistsException
 import org.easymock.EasyMock
 import org.junit.Assert._
 import org.junit.{After, Test}
@@ -28,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()
@@ -40,10 +42,10 @@ class ServerStartupTest extends ZooKeeperTestHarness {
     val zookeeperChroot = "/kafka-chroot-for-unittest"
     val props = TestUtils.createBrokerConfig(brokerId, zkConnect)
     val zooKeeperConnect = props.get("zookeeper.connect")
-    props.put("zookeeper.connect", zooKeeperConnect + zookeeperChroot)
+    props.put("zookeeper.connect", zooKeeperConnect.toString + zookeeperChroot)
     server = TestUtils.createServer(KafkaConfig.fromProps(props))
 
-    val pathExists = zkUtils.pathExists(zookeeperChroot)
+    val pathExists = zkClient.pathExists(zookeeperChroot)
     assertTrue(pathExists)
   }
 
@@ -62,7 +64,7 @@ class ServerStartupTest extends ZooKeeperTestHarness {
       TestUtils.createServer(KafkaConfig.fromProps(props2))
       fail("Starting a broker with the same port should fail")
     } catch {
-      case _: RuntimeException => // expected
+      case _: KafkaException => // expected
     }
   }
 
@@ -74,19 +76,19 @@ class ServerStartupTest extends ZooKeeperTestHarness {
     val brokerId = 0
     val props1 = TestUtils.createBrokerConfig(brokerId, zkConnect)
     server = TestUtils.createServer(KafkaConfig.fromProps(props1))
-    val brokerRegistration = zkUtils.readData(ZkUtils.BrokerIdsPath + "/" + brokerId)._1
+    val brokerRegistration = zkClient.getBroker(brokerId).getOrElse(fail("broker doesn't exists"))
 
     val props2 = TestUtils.createBrokerConfig(brokerId, zkConnect)
     try {
       TestUtils.createServer(KafkaConfig.fromProps(props2))
       fail("Registering a broker with a conflicting id should fail")
     } catch {
-      case _: RuntimeException =>
+      case _: NodeExistsException =>
       // this is expected
     }
 
     // broker registration shouldn't change
-    assertEquals(brokerRegistration, zkUtils.readData(ZkUtils.BrokerIdsPath + "/" + brokerId)._1)
+    assertEquals(brokerRegistration, zkClient.getBroker(brokerId).getOrElse(fail("broker doesn't exists")))
   }
 
   @Test
@@ -103,11 +105,11 @@ class ServerStartupTest extends ZooKeeperTestHarness {
   @Test
   def testBrokerStateRunningAfterZK(): Unit = {
     val brokerId = 0
-    val mockBrokerState = EasyMock.niceMock(classOf[kafka.server.BrokerState])
+    val mockBrokerState: BrokerState = EasyMock.niceMock(classOf[BrokerState])
 
     class BrokerStateInterceptor() extends BrokerState {
       override def newState(newState: BrokerStates): Unit = {
-        val brokers = zkUtils.getAllBrokersInCluster()
+        val brokers = zkClient.getAllBrokersInCluster
         assertEquals(1, brokers.size)
         assertEquals(brokerId, brokers.head.id)
       }
diff --git a/core/src/test/scala/unit/kafka/server/SessionExpireListenerTest.scala b/core/src/test/scala/unit/kafka/server/SessionExpireListenerTest.scala
deleted file mode 100644
index fda17c0201c7b..0000000000000
--- a/core/src/test/scala/unit/kafka/server/SessionExpireListenerTest.scala
+++ /dev/null
@@ -1,73 +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 kafka.api.ApiVersion
-import kafka.utils.ZkUtils
-import org.I0Itec.zkclient.ZkClient
-import org.apache.zookeeper.Watcher
-import org.easymock.EasyMock
-import org.junit.{Assert, Test}
-import Assert._
-import com.yammer.metrics.Metrics
-import com.yammer.metrics.core.Meter
-import scala.collection.JavaConverters._
-
-class SessionExpireListenerTest {
-
-  private val brokerId = 1
-
-  private def cleanMetricsRegistry() {
-    val metrics = Metrics.defaultRegistry
-    metrics.allMetrics.keySet.asScala.foreach(metrics.removeMetric)
-  }
-
-  @Test
-  def testSessionExpireListenerMetrics() {
-
-    cleanMetricsRegistry()
-
-    val metrics = Metrics.defaultRegistry
-
-    def checkMeterCount(name: String, expected: Long) {
-      val meter = metrics.allMetrics.asScala.collectFirst {
-        case (metricName, meter: Meter) if metricName.getName == name => meter
-      }.getOrElse(sys.error(s"Unable to find meter with name $name"))
-      assertEquals(s"Unexpected meter count for $name", expected, meter.count)
-    }
-
-    val zkClient = EasyMock.mock(classOf[ZkClient])
-    val zkUtils = ZkUtils(zkClient, isZkSecurityEnabled = false)
-    import Watcher._
-    val healthcheck = new KafkaHealthcheck(brokerId, Seq.empty, zkUtils, None, ApiVersion.latestVersion)
-
-    val expiresPerSecName = "ZooKeeperExpiresPerSec"
-    val disconnectsPerSecName = "ZooKeeperDisconnectsPerSec"
-    checkMeterCount(expiresPerSecName, 0)
-    checkMeterCount(disconnectsPerSecName, 0)
-
-    healthcheck.sessionExpireListener.handleStateChanged(Event.KeeperState.Expired)
-    checkMeterCount(expiresPerSecName, 1)
-    checkMeterCount(disconnectsPerSecName, 0)
-
-    healthcheck.sessionExpireListener.handleStateChanged(Event.KeeperState.Disconnected)
-    checkMeterCount(expiresPerSecName, 1)
-    checkMeterCount(disconnectsPerSecName, 1)
-  }
-
-}
diff --git a/core/src/test/scala/unit/kafka/server/SimpleFetchTest.scala b/core/src/test/scala/unit/kafka/server/SimpleFetchTest.scala
deleted file mode 100644
index 0797b7b3d118d..0000000000000
--- a/core/src/test/scala/unit/kafka/server/SimpleFetchTest.scala
+++ /dev/null
@@ -1,203 +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 kafka.api._
-import kafka.utils._
-import kafka.cluster.Replica
-import kafka.log.Log
-import kafka.server.QuotaFactory.UnboundedQuota
-import kafka.zk.KafkaZkClient
-import org.apache.kafka.common.metrics.Metrics
-import org.apache.kafka.common.requests.FetchRequest.PartitionData
-import org.junit.{After, Before, Test}
-import java.util.Properties
-import java.util.concurrent.atomic.AtomicBoolean
-
-import org.apache.kafka.common.TopicPartition
-import org.apache.kafka.common.record.{CompressionType, MemoryRecords, SimpleRecord}
-import org.apache.kafka.common.requests.IsolationLevel
-import org.easymock.EasyMock
-import org.junit.Assert._
-
-class SimpleFetchTest {
-
-  val replicaLagTimeMaxMs = 100L
-  val replicaFetchWaitMaxMs = 100
-  val replicaLagMaxMessages = 10L
-
-  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))
-
-  // set the replica manager with the partition
-  val time = new MockTime
-  val metrics = new Metrics
-  val leaderLEO = 20L
-  val followerLEO = 15L
-  val partitionHW = 5
-
-  val fetchSize = 100
-  val recordToHW = new SimpleRecord("recordToHW".getBytes())
-  val recordToLEO = new SimpleRecord("recordToLEO".getBytes())
-
-  val topic = "test-topic"
-  val partitionId = 0
-  val topicPartition = new TopicPartition(topic, partitionId)
-
-  val fetchInfo = Seq(topicPartition -> new PartitionData(0, 0, fetchSize))
-
-  var replicaManager: ReplicaManager = _
-
-  @Before
-  def setUp() {
-    // create nice mock since we don't particularly care about zkclient calls
-    val kafkaZkClient = EasyMock.createNiceMock(classOf[KafkaZkClient])
-    EasyMock.replay(kafkaZkClient)
-
-    // create nice mock since we don't particularly care about scheduler calls
-    val scheduler = EasyMock.createNiceMock(classOf[KafkaScheduler])
-    EasyMock.replay(scheduler)
-
-    // create the log which takes read with either HW max offset or none max offset
-    val log = EasyMock.createNiceMock(classOf[Log])
-    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.read(
-      startOffset = 0,
-      maxLength = fetchSize,
-      maxOffset = Some(partitionHW),
-      minOneMessage = true,
-      isolationLevel = IsolationLevel.READ_UNCOMMITTED))
-      .andReturn(FetchDataInfo(
-        new LogOffsetMetadata(0L, 0L, 0),
-        MemoryRecords.withRecords(CompressionType.NONE, recordToHW)
-      )).anyTimes()
-    EasyMock.expect(log.read(
-      startOffset = 0,
-      maxLength = fetchSize,
-      maxOffset = None,
-      minOneMessage = true,
-      isolationLevel = IsolationLevel.READ_UNCOMMITTED))
-      .andReturn(FetchDataInfo(
-        new LogOffsetMetadata(0L, 0L, 0),
-        MemoryRecords.withRecords(CompressionType.NONE, recordToLEO)
-      )).anyTimes()
-    EasyMock.replay(log)
-
-    // create the log manager that is aware of this mock log
-    val logManager = EasyMock.createMock(classOf[kafka.log.LogManager])
-    EasyMock.expect(logManager.getLog(topicPartition, false)).andReturn(Some(log)).anyTimes()
-    EasyMock.expect(logManager.liveLogDirs).andReturn(Array.empty[File]).anyTimes()
-    EasyMock.replay(logManager)
-
-    // create the replica manager
-    replicaManager = new ReplicaManager(configs.head, metrics, time, kafkaZkClient, scheduler, logManager,
-      new AtomicBoolean(false), QuotaFactory.instantiate(configs.head, metrics, time, ""), new BrokerTopicStats,
-      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))
-
-    // 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)
-
-    // 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
-  }
-
-  @After
-  def tearDown() {
-    replicaManager.shutdown(false)
-    metrics.close()
-  }
-
-  /**
-   * The scenario for this test is that there is one topic that has one partition
-   * with one leader replica on broker "0" and one follower replica on broker "1"
-   * inside the replica manager's metadata.
-   *
-   * The leader replica on "0" has HW of "5" and LEO of "20".  The follower on
-   * broker "1" has a local replica with a HW matching the leader's ("5") and
-   * LEO of "15", meaning it's not in-sync but is still in ISR (hasn't yet expired from ISR).
-   *
-   * When a fetch operation with read committed data turned on is received, the replica manager
-   * should only return data up to the HW of the partition; when a fetch operation with read
-   * committed data turned off is received, the replica manager could return data up to the LEO
-   * of the local leader replica's log.
-   *
-   * This test also verifies counts of fetch requests recorded by the ReplicaManager
-   */
-  @Test
-  def testReadFromLog() {
-    val brokerTopicStats = new BrokerTopicStats
-    val initialTopicCount = brokerTopicStats.topicStats(topic).totalFetchRequestRate.count()
-    val initialAllTopicsCount = brokerTopicStats.allTopicsStats.totalFetchRequestRate.count()
-
-    val readCommittedRecords = replicaManager.readFromLocalLog(
-      replicaId = Request.OrdinaryConsumerId,
-      fetchOnlyFromLeader = true,
-      readOnlyCommitted = true,
-      fetchMaxBytes = Int.MaxValue,
-      hardMaxBytesLimit = false,
-      readPartitionInfo = fetchInfo,
-      quota = UnboundedQuota,
-      isolationLevel = IsolationLevel.READ_UNCOMMITTED).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))
-
-    val readAllRecords = replicaManager.readFromLocalLog(
-      replicaId = Request.OrdinaryConsumerId,
-      fetchOnlyFromLeader = true,
-      readOnlyCommitted = false,
-      fetchMaxBytes = Int.MaxValue,
-      hardMaxBytesLimit = false,
-      readPartitionInfo = fetchInfo,
-      quota = UnboundedQuota,
-      isolationLevel = IsolationLevel.READ_UNCOMMITTED).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,
-      new SimpleRecord(firstRecord))
-
-    assertEquals("Counts should increment after fetch", initialTopicCount+2, brokerTopicStats.topicStats(topic).totalFetchRequestRate.count())
-    assertEquals("Counts should increment after fetch", initialAllTopicsCount+2, brokerTopicStats.allTopicsStats.totalFetchRequestRate.count())
-  }
-}
diff --git a/core/src/test/scala/unit/kafka/server/StopReplicaRequestTest.scala b/core/src/test/scala/unit/kafka/server/StopReplicaRequestTest.scala
new file mode 100644
index 0000000000000..cfb6cbf2bed65
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/StopReplicaRequestTest.scala
@@ -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 kafka.server
+
+import kafka.api.LeaderAndIsr
+import kafka.utils._
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.message.StopReplicaRequestData.{StopReplicaPartitionState, StopReplicaTopicState}
+import org.apache.kafka.common.protocol.ApiKeys
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.requests._
+import org.junit.Assert._
+import org.junit.Test
+
+import scala.jdk.CollectionConverters._
+import scala.collection.Seq
+
+class StopReplicaRequestTest extends BaseRequestTest {
+  override val logDirCount = 2
+  override val brokerCount: Int = 1
+
+  val topic = "topic"
+  val partitionNum = 2
+  val tp0 = new TopicPartition(topic, 0)
+  val tp1 = new TopicPartition(topic, 1)
+
+  @Test
+  def testStopReplicaRequest(): Unit = {
+    createTopic(topic, partitionNum, 1)
+    TestUtils.generateAndProduceMessages(servers, topic, 10)
+
+    val server = servers.head
+    val offlineDir = server.logManager.getLog(tp1).get.dir.getParent
+    server.replicaManager.handleLogDirFailure(offlineDir, sendZkNotification = false)
+
+    val topicStates = Seq(
+      new StopReplicaTopicState()
+        .setTopicName(tp0.topic())
+        .setPartitionStates(Seq(new StopReplicaPartitionState()
+          .setPartitionIndex(tp0.partition())
+          .setLeaderEpoch(LeaderAndIsr.initialLeaderEpoch + 2)
+          .setDeletePartition(true)).asJava),
+      new StopReplicaTopicState()
+        .setTopicName(tp1.topic())
+        .setPartitionStates(Seq(new StopReplicaPartitionState()
+          .setPartitionIndex(tp1.partition())
+          .setLeaderEpoch(LeaderAndIsr.initialLeaderEpoch + 2)
+          .setDeletePartition(true)).asJava)
+    ).asJava
+
+    for (_ <- 1 to 2) {
+      val request1 = new StopReplicaRequest.Builder(ApiKeys.STOP_REPLICA.latestVersion,
+        server.config.brokerId, server.replicaManager.controllerEpoch, server.kafkaController.brokerEpoch,
+        false, topicStates).build()
+      val response1 = connectAndReceive[StopReplicaResponse](request1, destination = controllerSocketServer)
+      val partitionErrors1 = response1.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
new file mode 100644
index 0000000000000..ad8da634138a5
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/ThrottledChannelExpirationTest.scala
@@ -0,0 +1,126 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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.net.InetAddress
+import java.util
+import java.util.Collections
+import java.util.concurrent.{DelayQueue, TimeUnit}
+import kafka.network.RequestChannel
+import kafka.network.RequestChannel.{EndThrottlingResponse, Response, StartThrottlingResponse}
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.memory.MemoryPool
+import org.apache.kafka.common.metrics.MetricConfig
+import org.apache.kafka.common.network.ClientInformation
+import org.apache.kafka.common.network.ListenerName
+import org.apache.kafka.common.requests.FetchRequest.PartitionData
+import org.apache.kafka.common.requests.{AbstractRequest, FetchRequest, RequestContext, RequestHeader, RequestTestUtils}
+import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
+import org.apache.kafka.common.utils.MockTime
+import org.easymock.EasyMock
+import org.junit.{Assert, Before, Test}
+
+class ThrottledChannelExpirationTest {
+  private val time = new MockTime
+  private var numCallbacksForStartThrottling: Int = 0
+  private var numCallbacksForEndThrottling: Int = 0
+  private val metrics = new org.apache.kafka.common.metrics.Metrics(new MetricConfig(),
+                                                                    Collections.emptyList(),
+                                                                    time)
+  private val request = buildRequest(FetchRequest.Builder.forConsumer(0, 1000, new util.HashMap[TopicPartition, PartitionData]))._2
+
+  private def buildRequest[T <: AbstractRequest](builder: AbstractRequest.Builder[T],
+                                                 listenerName: ListenerName = ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT)): (T, RequestChannel.Request) = {
+
+    val request = builder.build()
+    val buffer = RequestTestUtils.serializeRequestWithHeader(
+      new RequestHeader(builder.apiKey, request.version, "", 0), request)
+    val requestChannelMetrics: RequestChannel.Metrics = EasyMock.createNiceMock(classOf[RequestChannel.Metrics])
+
+    // read the header from the buffer first so that the body can be read next from the Request constructor
+    val header = RequestHeader.parse(buffer)
+    val context = new RequestContext(header, "1", InetAddress.getLocalHost, KafkaPrincipal.ANONYMOUS,
+      listenerName, SecurityProtocol.PLAINTEXT, ClientInformation.EMPTY, false)
+    (request, new RequestChannel.Request(processor = 1, context = context, startTimeNanos =  0, MemoryPool.NONE, buffer,
+      requestChannelMetrics))
+  }
+
+  def callback(response: Response): Unit = {
+    (response: @unchecked) match {
+      case _: StartThrottlingResponse => numCallbacksForStartThrottling += 1
+      case _: EndThrottlingResponse => numCallbacksForEndThrottling += 1
+    }
+  }
+
+  @Before
+  def beforeMethod(): Unit = {
+    numCallbacksForStartThrottling = 0
+    numCallbacksForEndThrottling = 0
+  }
+
+  @Test
+  def testCallbackInvocationAfterExpiration(): Unit = {
+    val clientMetrics = new ClientQuotaManager(ClientQuotaManagerConfig(), metrics, QuotaType.Produce, time, "")
+
+    val delayQueue = new DelayQueue[ThrottledChannel]()
+    val reaper = new clientMetrics.ThrottledChannelReaper(delayQueue, "")
+    try {
+      // Add 4 elements to the queue out of order. Add 2 elements with the same expire timestamp.
+      val channel1 = new ThrottledChannel(request, time, 10, callback)
+      val channel2 = new ThrottledChannel(request, time, 30, callback)
+      val channel3 = new ThrottledChannel(request, time, 30, callback)
+      val channel4 = new ThrottledChannel(request, time, 20, callback)
+      delayQueue.add(channel1)
+      delayQueue.add(channel2)
+      delayQueue.add(channel3)
+      delayQueue.add(channel4)
+      Assert.assertEquals(4, numCallbacksForStartThrottling)
+
+      for(itr <- 1 to 3) {
+        time.sleep(10)
+        reaper.doWork()
+        Assert.assertEquals(itr, numCallbacksForEndThrottling)
+      }
+      reaper.doWork()
+      Assert.assertEquals(4, numCallbacksForEndThrottling)
+      Assert.assertEquals(0, delayQueue.size())
+      reaper.doWork()
+      Assert.assertEquals(4, numCallbacksForEndThrottling)
+    } finally {
+      clientMetrics.shutdown()
+    }
+  }
+
+  @Test
+  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)
+    Assert.assertEquals(10, t1.throttleTimeMs)
+    Assert.assertEquals(20, t2.throttleTimeMs)
+    Assert.assertEquals(20, t3.throttleTimeMs)
+
+    for(itr <- 0 to 2) {
+      Assert.assertEquals(10 - 10*itr, t1.getDelay(TimeUnit.MILLISECONDS))
+      Assert.assertEquals(20 - 10*itr, t2.getDelay(TimeUnit.MILLISECONDS))
+      Assert.assertEquals(20 - 10*itr, t3.getDelay(TimeUnit.MILLISECONDS))
+      time.sleep(10)
+    }
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/server/ThrottledResponseExpirationTest.scala b/core/src/test/scala/unit/kafka/server/ThrottledResponseExpirationTest.scala
deleted file mode 100644
index d7634c0c0bf74..0000000000000
--- a/core/src/test/scala/unit/kafka/server/ThrottledResponseExpirationTest.scala
+++ /dev/null
@@ -1,89 +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.util.Collections
-import java.util.concurrent.{DelayQueue, TimeUnit}
-
-import org.apache.kafka.common.metrics.MetricConfig
-import org.apache.kafka.common.utils.MockTime
-import org.junit.{Assert, Before, Test}
-
-class ThrottledResponseExpirationTest {
-  private val time = new MockTime
-  private var numCallbacks: Int = 0
-  private val metrics = new org.apache.kafka.common.metrics.Metrics(new MetricConfig(),
-                                                                    Collections.emptyList(),
-                                                                    time)
-
-  def callback(delayTimeMs: Int) {
-    numCallbacks += 1
-  }
-
-  @Before
-  def beforeMethod() {
-    numCallbacks = 0
-  }
-
-  @Test
-  def testExpire() {
-    val clientMetrics = new ClientQuotaManager(ClientQuotaManagerConfig(), metrics, QuotaType.Produce, time, "")
-
-    val delayQueue = new DelayQueue[ThrottledResponse]()
-    val reaper = new clientMetrics.ThrottledRequestReaper(delayQueue, "")
-    try {
-      // Add 4 elements to the queue out of order. Add 2 elements with the same expire timestamp
-      delayQueue.add(new ThrottledResponse(time, 10, callback))
-      delayQueue.add(new ThrottledResponse(time, 30, callback))
-      delayQueue.add(new ThrottledResponse(time, 30, callback))
-      delayQueue.add(new ThrottledResponse(time, 20, callback))
-
-      for(itr <- 1 to 3) {
-        time.sleep(10)
-        reaper.doWork()
-        Assert.assertEquals(itr, numCallbacks)
-
-      }
-      reaper.doWork()
-      Assert.assertEquals(4, numCallbacks)
-      Assert.assertEquals(0, delayQueue.size())
-      reaper.doWork()
-      Assert.assertEquals(4, numCallbacks)
-    } finally {
-      clientMetrics.shutdown()
-    }
-  }
-
-  @Test
-  def testThrottledRequest() {
-    val t1: ThrottledResponse = new ThrottledResponse(time, 10, callback)
-    val t2: ThrottledResponse = new ThrottledResponse(time, 20, callback)
-    val t3: ThrottledResponse = new ThrottledResponse(time, 20, callback)
-    Assert.assertEquals(10, t1.throttleTimeMs)
-    Assert.assertEquals(20, t2.throttleTimeMs)
-    Assert.assertEquals(20, t3.throttleTimeMs)
-
-    for(itr <- 0 to 2) {
-      Assert.assertEquals(10 - 10*itr, t1.getDelay(TimeUnit.MILLISECONDS))
-      Assert.assertEquals(20 - 10*itr, t2.getDelay(TimeUnit.MILLISECONDS))
-      Assert.assertEquals(20 - 10*itr, t3.getDelay(TimeUnit.MILLISECONDS))
-      time.sleep(10)
-    }
-  }
-}
diff --git a/core/src/test/scala/unit/kafka/server/UpdateFeaturesTest.scala b/core/src/test/scala/unit/kafka/server/UpdateFeaturesTest.scala
new file mode 100644
index 0000000000000..f0099090912fb
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/UpdateFeaturesTest.scala
@@ -0,0 +1,580 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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 java.util.concurrent.ExecutionException
+
+import kafka.api.KAFKA_2_7_IV0
+import kafka.utils.TestUtils
+import kafka.zk.{FeatureZNode, FeatureZNodeStatus, ZkVersion}
+import kafka.utils.TestUtils.waitUntilTrue
+import org.apache.kafka.clients.admin.{Admin, FeatureUpdate, UpdateFeaturesOptions, UpdateFeaturesResult}
+import org.apache.kafka.common.errors.InvalidRequestException
+import org.apache.kafka.common.feature.FinalizedVersionRange
+import org.apache.kafka.common.feature.{Features, SupportedVersionRange}
+import org.apache.kafka.common.message.UpdateFeaturesRequestData
+import org.apache.kafka.common.message.UpdateFeaturesRequestData.FeatureUpdateKeyCollection
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.requests.{UpdateFeaturesRequest, UpdateFeaturesResponse}
+import org.apache.kafka.common.utils.Utils
+import org.junit.Test
+import org.junit.Assert.{assertEquals, assertFalse, assertNotEquals, assertNotNull, assertTrue}
+import org.scalatest.Assertions.intercept
+
+import scala.jdk.CollectionConverters._
+import scala.reflect.ClassTag
+import scala.util.matching.Regex
+
+class UpdateFeaturesTest extends BaseRequestTest {
+
+  override def brokerCount = 3
+
+  override def brokerPropertyOverrides(props: Properties): Unit = {
+    props.put(KafkaConfig.InterBrokerProtocolVersionProp, KAFKA_2_7_IV0.toString)
+  }
+
+  private def defaultSupportedFeatures(): Features[SupportedVersionRange] = {
+    Features.supportedFeatures(Utils.mkMap(Utils.mkEntry("feature_1", new SupportedVersionRange(1, 3))))
+  }
+
+  private def defaultFinalizedFeatures(): Features[FinalizedVersionRange] = {
+    Features.finalizedFeatures(Utils.mkMap(Utils.mkEntry("feature_1", new FinalizedVersionRange(1, 2))))
+  }
+
+  private def updateSupportedFeatures(
+    features: Features[SupportedVersionRange], targetServers: Set[KafkaServer]): Unit = {
+    targetServers.foreach(s => {
+      s.brokerFeatures.setSupportedFeatures(features)
+      s.zkClient.updateBrokerInfo(s.createBrokerInfo)
+    })
+
+    // Wait until updates to all BrokerZNode supported features propagate to the controller.
+    val brokerIds = targetServers.map(s => s.config.brokerId)
+    waitUntilTrue(
+      () => servers.exists(s => {
+        if (s.kafkaController.isActive) {
+          s.kafkaController.controllerContext.liveOrShuttingDownBrokers
+            .filter(b => brokerIds.contains(b.id))
+            .forall(b => {
+              b.features.equals(features)
+            })
+        } else {
+          false
+        }
+      }),
+      "Controller did not get broker updates")
+  }
+
+  private def updateSupportedFeaturesInAllBrokers(features: Features[SupportedVersionRange]): Unit = {
+    updateSupportedFeatures(features, Set[KafkaServer]() ++ servers)
+  }
+
+  private def updateFeatureZNode(features: Features[FinalizedVersionRange]): Int = {
+    val server = serverForId(0).get
+    val newNode = new FeatureZNode(FeatureZNodeStatus.Enabled, features)
+    val newVersion = server.zkClient.updateFeatureZNode(newNode)
+    servers.foreach(s => {
+      s.featureCache.waitUntilEpochOrThrow(newVersion, s.config.zkConnectionTimeoutMs)
+    })
+    newVersion
+  }
+
+  private def getFeatureZNode(): FeatureZNode = {
+    val (mayBeFeatureZNodeBytes, version) = serverForId(0).get.zkClient.getDataAndVersion(FeatureZNode.path)
+    assertNotEquals(version, ZkVersion.UnknownVersion)
+    FeatureZNode.decode(mayBeFeatureZNodeBytes.get)
+  }
+
+  private def finalizedFeatures(features: java.util.Map[String, org.apache.kafka.clients.admin.FinalizedVersionRange]): Features[FinalizedVersionRange] = {
+    Features.finalizedFeatures(features.asScala.map {
+      case(name, versionRange) =>
+        (name, new FinalizedVersionRange(versionRange.minVersionLevel(), versionRange.maxVersionLevel()))
+    }.asJava)
+  }
+
+  private def supportedFeatures(features: java.util.Map[String, org.apache.kafka.clients.admin.SupportedVersionRange]): Features[SupportedVersionRange] = {
+    Features.supportedFeatures(features.asScala.map {
+      case(name, versionRange) =>
+        (name, new SupportedVersionRange(versionRange.minVersion(), versionRange.maxVersion()))
+    }.asJava)
+  }
+
+  private def checkFeatures(client: Admin,
+                            expectedNode: FeatureZNode,
+                            expectedFinalizedFeatures: Features[FinalizedVersionRange],
+                            expectedFinalizedFeaturesEpoch: Long,
+                            expectedSupportedFeatures: Features[SupportedVersionRange]): Unit = {
+    assertEquals(expectedNode, getFeatureZNode())
+    val featureMetadata = client.describeFeatures.featureMetadata.get
+    assertEquals(expectedFinalizedFeatures, finalizedFeatures(featureMetadata.finalizedFeatures))
+    assertEquals(expectedSupportedFeatures, supportedFeatures(featureMetadata.supportedFeatures))
+    assertEquals(Optional.of(expectedFinalizedFeaturesEpoch), featureMetadata.finalizedFeaturesEpoch)
+  }
+
+  private def checkException[ExceptionType <: Throwable](result: UpdateFeaturesResult,
+                                                         featureExceptionMsgPatterns: Map[String, Regex])
+                                                        (implicit tag: ClassTag[ExceptionType]): Unit = {
+    featureExceptionMsgPatterns.foreach {
+      case (feature, exceptionMsgPattern) =>
+        val exception = intercept[ExecutionException] {
+          result.values().get(feature).get()
+        }
+        val cause = exception.getCause
+        assertNotNull(cause)
+        assertEquals(cause.getClass, tag.runtimeClass)
+        assertTrue(s"Received unexpected error message: ${cause.getMessage}",
+                   exceptionMsgPattern.findFirstIn(cause.getMessage).isDefined)
+    }
+  }
+
+  /**
+   * Tests whether an invalid feature update does not get processed on the server as expected,
+   * and raises the ExceptionType on the client side as expected.
+   *
+   * @param feature               the feature to be updated
+   * @param invalidUpdate         the invalid feature update to be sent in the
+   *                              updateFeatures request to the server
+   * @param exceptionMsgPattern   a pattern for the expected exception message
+   */
+  private def testWithInvalidFeatureUpdate[ExceptionType <: Throwable](feature: String,
+                                                                       invalidUpdate: FeatureUpdate,
+                                                                       exceptionMsgPattern: Regex)
+                                                                      (implicit tag: ClassTag[ExceptionType]): Unit = {
+    TestUtils.waitUntilControllerElected(zkClient)
+
+    updateSupportedFeaturesInAllBrokers(defaultSupportedFeatures())
+    val versionBefore = updateFeatureZNode(defaultFinalizedFeatures())
+    val adminClient = createAdminClient()
+    val nodeBefore = getFeatureZNode()
+
+    val result = adminClient.updateFeatures(Utils.mkMap(Utils.mkEntry(feature, invalidUpdate)), new UpdateFeaturesOptions())
+
+    checkException[ExceptionType](result, Map(feature -> exceptionMsgPattern))
+    checkFeatures(
+      adminClient,
+      nodeBefore,
+      defaultFinalizedFeatures(),
+      versionBefore,
+      defaultSupportedFeatures())
+  }
+
+  /**
+   * Tests that an UpdateFeatures request sent to a non-Controller node fails as expected.
+   */
+  @Test
+  def testShouldFailRequestIfNotController(): Unit = {
+    TestUtils.waitUntilControllerElected(zkClient)
+
+    updateSupportedFeaturesInAllBrokers(defaultSupportedFeatures())
+    val versionBefore = updateFeatureZNode(defaultFinalizedFeatures())
+
+    val nodeBefore = getFeatureZNode()
+    val validUpdates = new FeatureUpdateKeyCollection()
+    val validUpdate = new UpdateFeaturesRequestData.FeatureUpdateKey();
+    validUpdate.setFeature("feature_1");
+    validUpdate.setMaxVersionLevel(defaultSupportedFeatures().get("feature_1").max())
+    validUpdate.setAllowDowngrade(false)
+    validUpdates.add(validUpdate)
+
+    val response = connectAndReceive[UpdateFeaturesResponse](
+      new UpdateFeaturesRequest.Builder(new UpdateFeaturesRequestData().setFeatureUpdates(validUpdates)).build(),
+      notControllerSocketServer)
+
+    assertEquals(Errors.NOT_CONTROLLER, Errors.forCode(response.data.errorCode()))
+    assertNotNull(response.data.errorMessage())
+    assertEquals(0, response.data.results.size)
+    checkFeatures(
+      createAdminClient(),
+      nodeBefore,
+      defaultFinalizedFeatures(),
+      versionBefore,
+      defaultSupportedFeatures())
+  }
+
+  /**
+   * Tests that an UpdateFeatures request fails in the Controller, when, for a feature the
+   * allowDowngrade flag is not set during a downgrade request.
+   */
+  @Test
+  def testShouldFailRequestWhenDowngradeFlagIsNotSetDuringDowngrade(): Unit = {
+    val targetMaxVersionLevel = (defaultFinalizedFeatures().get("feature_1").max() - 1).asInstanceOf[Short]
+    testWithInvalidFeatureUpdate[InvalidRequestException](
+      "feature_1",
+      new FeatureUpdate(targetMaxVersionLevel,false),
+      ".*Can not downgrade finalized feature.*allowDowngrade.*".r)
+  }
+
+  /**
+   * Tests that an UpdateFeatures request fails in the Controller, when, for a feature the downgrade
+   * is attempted to a max version level higher than the existing max version level.
+   */
+  @Test
+  def testShouldFailRequestWhenDowngradeToHigherVersionLevelIsAttempted(): Unit = {
+    val targetMaxVersionLevel = (defaultFinalizedFeatures().get("feature_1").max() + 1).asInstanceOf[Short]
+    testWithInvalidFeatureUpdate[InvalidRequestException](
+      "feature_1",
+      new FeatureUpdate(targetMaxVersionLevel, true),
+      ".*When the allowDowngrade flag set in the request, the provided maxVersionLevel:3.*existing maxVersionLevel:2.*".r)
+  }
+
+  /**
+   * Tests that an UpdateFeatures request fails in the Controller, when, a feature deletion is
+   * attempted without setting the allowDowngrade flag.
+   */
+  @Test
+  def testShouldFailRequestInServerWhenDowngradeFlagIsNotSetDuringDeletion(): Unit = {
+    TestUtils.waitUntilControllerElected(zkClient)
+
+    updateSupportedFeaturesInAllBrokers(defaultSupportedFeatures())
+    val versionBefore = updateFeatureZNode(defaultFinalizedFeatures())
+
+    val adminClient = createAdminClient()
+    val nodeBefore = getFeatureZNode()
+
+    val invalidUpdates
+      = new UpdateFeaturesRequestData.FeatureUpdateKeyCollection();
+    val invalidUpdate = new UpdateFeaturesRequestData.FeatureUpdateKey();
+    invalidUpdate.setFeature("feature_1")
+    invalidUpdate.setMaxVersionLevel(0)
+    invalidUpdate.setAllowDowngrade(false)
+    invalidUpdates.add(invalidUpdate);
+    val requestData = new UpdateFeaturesRequestData()
+    requestData.setFeatureUpdates(invalidUpdates);
+
+    val response = connectAndReceive[UpdateFeaturesResponse](
+      new UpdateFeaturesRequest.Builder(new UpdateFeaturesRequestData().setFeatureUpdates(invalidUpdates)).build(),
+      controllerSocketServer)
+
+    assertEquals(1, response.data().results().size())
+    val result = response.data.results.asScala.head
+    assertEquals("feature_1", result.feature)
+    assertEquals(Errors.INVALID_REQUEST, Errors.forCode(result.errorCode))
+    assertNotNull(result.errorMessage)
+    assertFalse(result.errorMessage.isEmpty)
+    val exceptionMsgPattern = ".*Can not provide maxVersionLevel: 0 less than 1.*allowDowngrade.*".r
+    assertTrue(result.errorMessage, exceptionMsgPattern.findFirstIn(result.errorMessage).isDefined)
+    checkFeatures(
+      adminClient,
+      nodeBefore,
+      defaultFinalizedFeatures(),
+      versionBefore,
+      defaultSupportedFeatures())
+  }
+
+  /**
+   * Tests that an UpdateFeatures request fails in the Controller, when, a feature version level
+   * upgrade is attempted for a non-existing feature.
+   */
+  @Test
+  def testShouldFailRequestDuringDeletionOfNonExistingFeature(): Unit = {
+    testWithInvalidFeatureUpdate[InvalidRequestException](
+      "feature_non_existing",
+      new FeatureUpdate(3, true),
+      ".*Could not apply finalized feature update because the provided feature is not supported.*".r)
+  }
+
+  /**
+   * Tests that an UpdateFeatures request fails in the Controller, when, a feature version level
+   * upgrade is attempted to a version level same as the existing max version level.
+   */
+  @Test
+  def testShouldFailRequestWhenUpgradingToSameVersionLevel(): Unit = {
+    val targetMaxVersionLevel = defaultFinalizedFeatures().get("feature_1").max()
+    testWithInvalidFeatureUpdate[InvalidRequestException](
+      "feature_1",
+      new FeatureUpdate(targetMaxVersionLevel, false),
+      ".*Can not upgrade a finalized feature.*to the same value.*".r)
+  }
+
+  private def testShouldFailRequestDuringBrokerMaxVersionLevelIncompatibility(
+    featureName: String,
+    supportedVersionRange: SupportedVersionRange,
+    initialFinalizedVersionRange: Option[FinalizedVersionRange]
+  ): Unit = {
+    TestUtils.waitUntilControllerElected(zkClient)
+
+    val controller = servers.filter { server => server.kafkaController.isActive}.head
+    val nonControllerServers = servers.filter { server => !server.kafkaController.isActive}
+    // We setup the supported features on the broker such that 1/3 of the brokers does not
+    // support an expected feature version, while 2/3 brokers support the expected feature
+    // version.
+    val brokersWithVersionIncompatibility = Set[KafkaServer](nonControllerServers.head)
+    val versionCompatibleBrokers = Set[KafkaServer](nonControllerServers(1), controller)
+
+    val supportedFeatures = Features.supportedFeatures(Utils.mkMap(Utils.mkEntry(featureName, supportedVersionRange)))
+    updateSupportedFeatures(supportedFeatures, versionCompatibleBrokers)
+
+    val unsupportedMaxVersion = (supportedVersionRange.max() - 1).asInstanceOf[Short]
+    val supportedFeaturesWithVersionIncompatibility = Features.supportedFeatures(
+      Utils.mkMap(
+        Utils.mkEntry("feature_1",
+          new SupportedVersionRange(
+            supportedVersionRange.min(),
+            unsupportedMaxVersion))))
+    updateSupportedFeatures(supportedFeaturesWithVersionIncompatibility, brokersWithVersionIncompatibility)
+
+    val initialFinalizedFeatures = initialFinalizedVersionRange.map(
+      versionRange => Features.finalizedFeatures(Utils.mkMap(Utils.mkEntry(featureName, versionRange)))
+    ).getOrElse(Features.emptyFinalizedFeatures())
+    val versionBefore = updateFeatureZNode(initialFinalizedFeatures)
+
+    val invalidUpdate = new FeatureUpdate(supportedVersionRange.max(), false)
+    val nodeBefore = getFeatureZNode()
+    val adminClient = createAdminClient()
+    val result = adminClient.updateFeatures(
+      Utils.mkMap(Utils.mkEntry("feature_1", invalidUpdate)),
+      new UpdateFeaturesOptions())
+
+    checkException[InvalidRequestException](result, Map("feature_1" -> ".*brokers.*incompatible.*".r))
+    checkFeatures(
+      adminClient,
+      nodeBefore,
+      initialFinalizedFeatures,
+      versionBefore,
+      supportedFeatures)
+  }
+
+  /**
+   * Tests that an UpdateFeatures request fails in the Controller, when for an existing finalized
+   * feature, a version level upgrade introduces a version incompatibility with existing supported
+   * features.
+   */
+  @Test
+  def testShouldFailRequestDuringBrokerMaxVersionLevelIncompatibilityForExistingFinalizedFeature(): Unit = {
+    val feature = "feature_1"
+    testShouldFailRequestDuringBrokerMaxVersionLevelIncompatibility(
+      feature,
+      defaultSupportedFeatures().get(feature),
+      Some(defaultFinalizedFeatures().get(feature)))
+  }
+
+  /**
+   * Tests that an UpdateFeatures request fails in the Controller, when for a non-existing finalized
+   * feature, a version level upgrade introduces a version incompatibility with existing supported
+   * features.
+   */
+  @Test
+  def testShouldFailRequestDuringBrokerMaxVersionLevelIncompatibilityWithNoExistingFinalizedFeature(): Unit = {
+    val feature = "feature_1"
+    testShouldFailRequestDuringBrokerMaxVersionLevelIncompatibility(
+      feature,
+      defaultSupportedFeatures().get(feature),
+      Option.empty)
+  }
+
+  /**
+   * Tests that an UpdateFeatures request succeeds in the Controller, when, there are no existing
+   * finalized features in FeatureZNode when the test starts.
+   */
+  @Test
+  def testSuccessfulFeatureUpgradeAndWithNoExistingFinalizedFeatures(): Unit = {
+    TestUtils.waitUntilControllerElected(zkClient)
+
+    val supportedFeatures =
+      Features.supportedFeatures(
+        Utils.mkMap(
+          Utils.mkEntry("feature_1", new SupportedVersionRange(1, 3)),
+          Utils.mkEntry("feature_2", new SupportedVersionRange(2, 5))))
+    updateSupportedFeaturesInAllBrokers(supportedFeatures)
+    val versionBefore = updateFeatureZNode(Features.emptyFinalizedFeatures())
+
+    val targetFinalizedFeatures = Features.finalizedFeatures(
+      Utils.mkMap(
+        Utils.mkEntry("feature_1", new FinalizedVersionRange(1, 3)),
+        Utils.mkEntry("feature_2", new FinalizedVersionRange(2, 3))))
+    val update1 = new FeatureUpdate(targetFinalizedFeatures.get("feature_1").max(), false)
+    val update2 = new FeatureUpdate(targetFinalizedFeatures.get("feature_2").max(), false)
+
+    val adminClient = createAdminClient()
+    adminClient.updateFeatures(
+      Utils.mkMap(Utils.mkEntry("feature_1", update1), Utils.mkEntry("feature_2", update2)),
+      new UpdateFeaturesOptions()
+    ).all().get()
+
+    checkFeatures(
+      adminClient,
+      new FeatureZNode(FeatureZNodeStatus.Enabled, targetFinalizedFeatures),
+      targetFinalizedFeatures,
+      versionBefore + 1,
+      supportedFeatures)
+  }
+
+  /**
+   * Tests that an UpdateFeatures request succeeds in the Controller, when, the request contains
+   * both a valid feature version level upgrade as well as a downgrade request.
+   */
+  @Test
+  def testSuccessfulFeatureUpgradeAndDowngrade(): Unit = {
+    TestUtils.waitUntilControllerElected(zkClient)
+
+    val supportedFeatures = Features.supportedFeatures(
+      Utils.mkMap(
+        Utils.mkEntry("feature_1", new SupportedVersionRange(1, 3)),
+        Utils.mkEntry("feature_2", new SupportedVersionRange(2, 5))))
+    updateSupportedFeaturesInAllBrokers(supportedFeatures)
+    val initialFinalizedFeatures = Features.finalizedFeatures(
+      Utils.mkMap(
+        Utils.mkEntry("feature_1", new FinalizedVersionRange(1, 2)),
+        Utils.mkEntry("feature_2", new FinalizedVersionRange(2, 4))))
+    val versionBefore = updateFeatureZNode(initialFinalizedFeatures)
+
+    // Below we aim to do the following:
+    // - Valid upgrade of feature_1 maxVersionLevel from 2 to 3
+    // - Valid downgrade of feature_2 maxVersionLevel from 4 to 3
+    val targetFinalizedFeatures = Features.finalizedFeatures(
+      Utils.mkMap(
+        Utils.mkEntry("feature_1", new FinalizedVersionRange(1, 3)),
+        Utils.mkEntry("feature_2", new FinalizedVersionRange(2, 3))))
+    val update1 = new FeatureUpdate(targetFinalizedFeatures.get("feature_1").max(), false)
+    val update2 = new FeatureUpdate(targetFinalizedFeatures.get("feature_2").max(), true)
+
+    val adminClient = createAdminClient()
+    adminClient.updateFeatures(
+      Utils.mkMap(Utils.mkEntry("feature_1", update1), Utils.mkEntry("feature_2", update2)),
+      new UpdateFeaturesOptions()
+    ).all().get()
+
+    checkFeatures(
+      adminClient,
+      new FeatureZNode(FeatureZNodeStatus.Enabled, targetFinalizedFeatures),
+      targetFinalizedFeatures,
+      versionBefore + 1,
+      supportedFeatures)
+  }
+
+  /**
+   * Tests that an UpdateFeatures request succeeds partially in the Controller, when, the request
+   * contains a valid feature version level upgrade and an invalid version level downgrade.
+   * i.e. expect the upgrade operation to succeed, and the downgrade operation to fail.
+   */
+  @Test
+  def testPartialSuccessDuringValidFeatureUpgradeAndInvalidDowngrade(): Unit = {
+    TestUtils.waitUntilControllerElected(zkClient)
+
+    val supportedFeatures = Features.supportedFeatures(
+      Utils.mkMap(
+        Utils.mkEntry("feature_1", new SupportedVersionRange(1, 3)),
+        Utils.mkEntry("feature_2", new SupportedVersionRange(2, 5))))
+    updateSupportedFeaturesInAllBrokers(supportedFeatures)
+    val initialFinalizedFeatures = Features.finalizedFeatures(
+      Utils.mkMap(
+        Utils.mkEntry("feature_1", new FinalizedVersionRange(1, 2)),
+        Utils.mkEntry("feature_2", new FinalizedVersionRange(2, 4))))
+    val versionBefore = updateFeatureZNode(initialFinalizedFeatures)
+
+    // Below we aim to do the following:
+    // - Valid upgrade of feature_1 maxVersionLevel from 2 to 3
+    // - Invalid downgrade of feature_2 maxVersionLevel from 4 to 3
+    //   (because we intentionally do not set the allowDowngrade flag)
+    val targetFinalizedFeatures = Features.finalizedFeatures(
+      Utils.mkMap(
+        Utils.mkEntry("feature_1", new FinalizedVersionRange(1, 3)),
+        Utils.mkEntry("feature_2", new FinalizedVersionRange(2, 3))))
+    val validUpdate = new FeatureUpdate(targetFinalizedFeatures.get("feature_1").max(), false)
+    val invalidUpdate = new FeatureUpdate(targetFinalizedFeatures.get("feature_2").max(), false)
+
+    val adminClient = createAdminClient()
+    val result = adminClient.updateFeatures(
+      Utils.mkMap(Utils.mkEntry("feature_1", validUpdate), Utils.mkEntry("feature_2", invalidUpdate)),
+      new UpdateFeaturesOptions())
+
+    // Expect update for "feature_1" to have succeeded.
+    result.values().get("feature_1").get()
+    // Expect update for "feature_2" to have failed.
+    checkException[InvalidRequestException](
+      result, Map("feature_2" -> ".*Can not downgrade finalized feature.*allowDowngrade.*".r))
+    val expectedFeatures = Features.finalizedFeatures(
+      Utils.mkMap(
+        Utils.mkEntry("feature_1", targetFinalizedFeatures.get("feature_1")),
+        Utils.mkEntry("feature_2", initialFinalizedFeatures.get("feature_2"))))
+    checkFeatures(
+      adminClient,
+      FeatureZNode(FeatureZNodeStatus.Enabled, expectedFeatures),
+      expectedFeatures,
+      versionBefore + 1,
+      supportedFeatures)
+  }
+
+  /**
+   * Tests that an UpdateFeatures request succeeds partially in the Controller, when, the request
+   * contains an invalid feature version level upgrade and a valid version level downgrade.
+   * i.e. expect the downgrade operation to succeed, and the upgrade operation to fail.
+   */
+  @Test
+  def testPartialSuccessDuringInvalidFeatureUpgradeAndValidDowngrade(): Unit = {
+    TestUtils.waitUntilControllerElected(zkClient)
+
+    val controller = servers.filter { server => server.kafkaController.isActive}.head
+    val nonControllerServers = servers.filter { server => !server.kafkaController.isActive}
+    // We setup the supported features on the broker such that 1/3 of the brokers does not
+    // support an expected feature version, while 2/3 brokers support the expected feature
+    // version.
+    val brokersWithVersionIncompatibility = Set[KafkaServer](nonControllerServers.head)
+    val versionCompatibleBrokers = Set[KafkaServer](nonControllerServers(1), controller)
+
+    val supportedFeatures = Features.supportedFeatures(
+      Utils.mkMap(
+        Utils.mkEntry("feature_1", new SupportedVersionRange(1, 3)),
+        Utils.mkEntry("feature_2", new SupportedVersionRange(2, 5))))
+    updateSupportedFeatures(supportedFeatures, versionCompatibleBrokers)
+
+    val supportedFeaturesWithVersionIncompatibility = Features.supportedFeatures(
+      Utils.mkMap(
+        Utils.mkEntry("feature_1", new SupportedVersionRange(1, 2)),
+        Utils.mkEntry("feature_2", supportedFeatures.get("feature_2"))))
+    updateSupportedFeatures(supportedFeaturesWithVersionIncompatibility, brokersWithVersionIncompatibility)
+
+    val initialFinalizedFeatures = Features.finalizedFeatures(
+      Utils.mkMap(
+        Utils.mkEntry("feature_1", new FinalizedVersionRange(1, 2)),
+        Utils.mkEntry("feature_2", new FinalizedVersionRange(2, 4))))
+    val versionBefore = updateFeatureZNode(initialFinalizedFeatures)
+
+    // Below we aim to do the following:
+    // - Invalid upgrade of feature_1 maxVersionLevel from 2 to 3
+    //   (because one of the brokers does not support the max version: 3)
+    // - Valid downgrade of feature_2 maxVersionLevel from 4 to 3
+    val targetFinalizedFeatures = Features.finalizedFeatures(
+      Utils.mkMap(
+        Utils.mkEntry("feature_1", new FinalizedVersionRange(1, 3)),
+        Utils.mkEntry("feature_2", new FinalizedVersionRange(2, 3))))
+    val invalidUpdate = new FeatureUpdate(targetFinalizedFeatures.get("feature_1").max(), false)
+    val validUpdate = new FeatureUpdate(targetFinalizedFeatures.get("feature_2").max(), true)
+
+    val adminClient = createAdminClient()
+    val result = adminClient.updateFeatures(
+      Utils.mkMap(Utils.mkEntry("feature_1", invalidUpdate), Utils.mkEntry("feature_2", validUpdate)),
+      new UpdateFeaturesOptions())
+
+    // Expect update for "feature_2" to have succeeded.
+    result.values().get("feature_2").get()
+    // Expect update for "feature_1" to have failed.
+    checkException[InvalidRequestException](result, Map("feature_1" -> ".*brokers.*incompatible.*".r))
+    val expectedFeatures = Features.finalizedFeatures(
+      Utils.mkMap(
+        Utils.mkEntry("feature_1", initialFinalizedFeatures.get("feature_1")),
+        Utils.mkEntry("feature_2", targetFinalizedFeatures.get("feature_2"))))
+    checkFeatures(
+      adminClient,
+      FeatureZNode(FeatureZNodeStatus.Enabled, expectedFeatures),
+      expectedFeatures,
+      versionBefore + 1,
+      supportedFeatures)
+  }
+}
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 e7c6a9785bce6..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,10 +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 ={
@@ -69,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 a44f900aa43c6..cc8eb097a9831 100644
--- a/core/src/test/scala/unit/kafka/server/epoch/EpochDrivenReplicationProtocolAcceptanceTest.scala
+++ b/core/src/test/scala/unit/kafka/server/epoch/EpochDrivenReplicationProtocolAcceptanceTest.scala
@@ -20,7 +20,7 @@ package kafka.server.epoch
 import java.io.{File, RandomAccessFile}
 import java.util.Properties
 
-import kafka.api.KAFKA_0_11_0_IV1
+import kafka.api.ApiVersion
 import kafka.log.Log
 import kafka.server.KafkaConfig._
 import kafka.server.{KafkaConfig, KafkaServer}
@@ -29,19 +29,20 @@ import kafka.utils.{CoreUtils, Logging, TestUtils}
 import kafka.utils.TestUtils._
 import kafka.zk.ZooKeeperTestHarness
 import org.apache.kafka.clients.consumer.{ConsumerConfig, KafkaConsumer}
-import org.apache.kafka.clients.producer.{KafkaProducer, ProducerConfig, ProducerRecord}
+import org.apache.kafka.clients.producer.{KafkaProducer, ProducerRecord}
 import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.record.RecordBatch
-import org.apache.kafka.common.serialization.Deserializer
+import org.apache.kafka.common.serialization.ByteArrayDeserializer
 import org.junit.Assert.{assertEquals, assertTrue}
 import org.junit.{After, Before, Test}
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 import scala.collection.mutable.{ListBuffer => Buffer}
+import scala.collection.Seq
 
 /**
   * These tests were written to assert the addition of leader epochs to the replication protocol fix the problems
-  * described in KIP-101. There is a boolean KIP_101_ENABLED which can be toggled to demonstrate the tests failing in the pre-KIP-101 case
+  * described in KIP-101.
   *
   * https://cwiki.apache.org/confluence/display/KAFKA/KIP-101+-+Alter+Replication+Protocol+to+use+Leader+Epoch+rather+than+High+Watermark+for+Truncation
   *
@@ -49,6 +50,8 @@ import scala.collection.mutable.{ListBuffer => Buffer}
   */
 class EpochDrivenReplicationProtocolAcceptanceTest extends ZooKeeperTestHarness with Logging {
 
+  // Set this to KAFKA_0_11_0_IV1 to demonstrate the tests failing in the pre-KIP-101 case
+  val apiVersion = ApiVersion.latestVersion
   val topic = "topic1"
   val msg = new Array[Byte](1000)
   val msgBigger = new Array[Byte](10000)
@@ -56,15 +59,13 @@ class EpochDrivenReplicationProtocolAcceptanceTest extends ZooKeeperTestHarness
   var producer: KafkaProducer[Array[Byte], Array[Byte]] = null
   var consumer: KafkaConsumer[Array[Byte], Array[Byte]] = null
 
-  val KIP_101_ENABLED = true
-
   @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()
@@ -77,35 +78,35 @@ class EpochDrivenReplicationProtocolAcceptanceTest extends ZooKeeperTestHarness
     brokers = (100 to 101).map(createBroker(_))
 
     //A single partition topic with 2 replicas
-    adminZkClient.createOrUpdateTopicPartitionAssignmentPathInZK(topic, Map(0 -> Seq(100, 101)))
-    producer = createProducer()
+    TestUtils.createTopic(zkClient, topic, Map(0 -> Seq(100, 101)), brokers)
+    producer = createProducer
     val tp = new TopicPartition(topic, 0)
 
     //When one record is written to the leader
     producer.send(new ProducerRecord(topic, 0, null, msg)).get
 
     //The message should have epoch 0 stamped onto it in both leader and follower
-    assertEquals(0, latestRecord(leader).partitionLeaderEpoch())
-    assertEquals(0, latestRecord(follower).partitionLeaderEpoch())
+    assertEquals(0, latestRecord(leader).partitionLeaderEpoch)
+    assertEquals(0, latestRecord(follower).partitionLeaderEpoch)
 
     //Both leader and follower should have recorded Epoch 0 at Offset 0
-    assertEquals(Buffer(EpochEntry(0, 0)), epochCache(leader).epochEntries())
-    assertEquals(Buffer(EpochEntry(0, 0)), epochCache(follower).epochEntries())
+    assertEquals(Buffer(EpochEntry(0, 0)), epochCache(leader).epochEntries)
+    assertEquals(Buffer(EpochEntry(0, 0)), epochCache(follower).epochEntries)
 
     //Bounce the follower
     bounce(follower)
     awaitISR(tp)
 
     //Nothing happens yet as we haven't sent any new messages.
-    assertEquals(Buffer(EpochEntry(0, 0)), epochCache(leader).epochEntries())
-    assertEquals(Buffer(EpochEntry(0, 0)), epochCache(follower).epochEntries())
+    assertEquals(Buffer(EpochEntry(0, 0), EpochEntry(1, 1)), epochCache(leader).epochEntries)
+    assertEquals(Buffer(EpochEntry(0, 0)), epochCache(follower).epochEntries)
 
     //Send a message
     producer.send(new ProducerRecord(topic, 0, null, msg)).get
 
     //Epoch1 should now propagate to the follower with the written message
-    assertEquals(Buffer(EpochEntry(0, 0), EpochEntry(1, 1)), epochCache(leader).epochEntries())
-    assertEquals(Buffer(EpochEntry(0, 0), EpochEntry(1, 1)), epochCache(follower).epochEntries())
+    assertEquals(Buffer(EpochEntry(0, 0), EpochEntry(1, 1)), epochCache(leader).epochEntries)
+    assertEquals(Buffer(EpochEntry(0, 0), EpochEntry(1, 1)), epochCache(follower).epochEntries)
 
     //The new message should have epoch 1 stamped
     assertEquals(1, latestRecord(leader).partitionLeaderEpoch())
@@ -116,8 +117,8 @@ class EpochDrivenReplicationProtocolAcceptanceTest extends ZooKeeperTestHarness
     awaitISR(tp)
 
     //Epochs 2 should be added to the leader, but not on the follower (yet), as there has been no replication.
-    assertEquals(Buffer(EpochEntry(0, 0), EpochEntry(1, 1)), epochCache(leader).epochEntries())
-    assertEquals(Buffer(EpochEntry(0, 0), EpochEntry(1, 1)), epochCache(follower).epochEntries())
+    assertEquals(Buffer(EpochEntry(0, 0), EpochEntry(1, 1), EpochEntry(2, 2)), epochCache(leader).epochEntries)
+    assertEquals(Buffer(EpochEntry(0, 0), EpochEntry(1, 1)), epochCache(follower).epochEntries)
 
     //Send a message
     producer.send(new ProducerRecord(topic, 0, null, msg)).get
@@ -127,55 +128,53 @@ class EpochDrivenReplicationProtocolAcceptanceTest extends ZooKeeperTestHarness
     assertEquals(2, latestRecord(follower).partitionLeaderEpoch())
 
     //The leader epoch files should now match on leader and follower
-    assertEquals(Buffer(EpochEntry(0, 0), EpochEntry(1, 1), EpochEntry(2, 2)), epochCache(leader).epochEntries())
-    assertEquals(Buffer(EpochEntry(0, 0), EpochEntry(1, 1), EpochEntry(2, 2)), epochCache(follower).epochEntries())
+    assertEquals(Buffer(EpochEntry(0, 0), EpochEntry(1, 1), EpochEntry(2, 2)), epochCache(leader).epochEntries)
+    assertEquals(Buffer(EpochEntry(0, 0), EpochEntry(1, 1), EpochEntry(2, 2)), epochCache(follower).epochEntries)
   }
 
   @Test
   def shouldNotAllowDivergentLogs(): Unit = {
-
     //Given two brokers
     brokers = (100 to 101).map { id => createServer(fromProps(createBrokerConfig(id, zkConnect))) }
+    val broker100 = brokers(0)
+    val broker101 = brokers(1)
 
     //A single partition topic with 2 replicas
-    adminZkClient.createOrUpdateTopicPartitionAssignmentPathInZK(topic, Map(
-      0 -> Seq(100, 101)
-    ))
-    producer = createProducer()
+    TestUtils.createTopic(zkClient, topic, Map(0 -> Seq(100, 101)), brokers)
+    producer = createProducer
 
-    //Write 10 messages
+    //Write 10 messages (ensure they are not batched so we can truncate in the middle below)
     (0 until 10).foreach { i =>
-      producer.send(new ProducerRecord(topic, 0, null, msg))
-      producer.flush()
+      producer.send(new ProducerRecord(topic, 0, s"$i".getBytes, msg)).get()
     }
 
-    //Stop the brokers
-    brokers.foreach { b => b.shutdown() }
+    //Stop the brokers (broker 101 first so that 100 is the leader)
+    broker101.shutdown()
+    broker100.shutdown()
 
     //Delete the clean shutdown file to simulate crash
-    new File(brokers(0).config.logDirs(0), Log.CleanShutdownFile).delete()
+    new File(broker100.config.logDirs.head, Log.CleanShutdownFile).delete()
 
     //Delete 5 messages from the leader's log on 100
-    deleteMessagesFromLogFile(5 * msg.length, brokers(0), 0)
+    deleteMessagesFromLogFile(5 * msg.length, broker100, 0)
 
     //Restart broker 100
-    brokers(0).startup()
+    broker100.startup()
 
-    //Bounce the producer (this is required, although I'm unsure as to why?)
+    //Bounce the producer (this is required since the broker uses a random port)
     producer.close()
-    producer = createProducer()
+    producer = createProducer
 
-    //Write ten larger messages (so we can easily distinguish between messages written in the two phases)
-    (0 until 10).foreach { _ =>
-      producer.send(new ProducerRecord(topic, 0, null, msgBigger))
-      producer.flush()
-    }
+    //Write ten additional messages
+    (11 until 20).map { i =>
+      producer.send(new ProducerRecord(topic, 0, s"$i".getBytes, msg))
+    }.foreach(_.get())
 
-    //Start broker 101
-    brokers(1).startup()
+    //Start broker 101 (we expect it to truncate to match broker 100's log)
+    broker101.startup()
 
     //Wait for replication to resync
-    waitForLogsToMatch(brokers(0), brokers(1))
+    waitForLogsToMatch(broker100, broker101)
 
     assertEquals("Log files should match Broker0 vs Broker 1", getLogFile(brokers(0), 0).length, getLogFile(brokers(1), 0).length)
   }
@@ -188,10 +187,8 @@ class EpochDrivenReplicationProtocolAcceptanceTest extends ZooKeeperTestHarness
     brokers = (100 to 101).map(createBroker(_))
 
     //A single partition topic with 2 replicas
-    adminZkClient.createOrUpdateTopicPartitionAssignmentPathInZK(topic, Map(
-      0 -> Seq(100, 101)
-    ))
-    producer = bufferingProducer()
+    TestUtils.createTopic(zkClient, topic, Map(0 -> Seq(100, 101)), brokers)
+    producer = createBufferingProducer
 
     //Write 100 messages
     (0 until 100).foreach { i =>
@@ -213,7 +210,7 @@ class EpochDrivenReplicationProtocolAcceptanceTest extends ZooKeeperTestHarness
 
     //Bounce the producer (this is required, although I'm unsure as to why?)
     producer.close()
-    producer = bufferingProducer()
+    producer = createBufferingProducer
 
     //Write two large batches of messages. This will ensure that the LeO of the follower's log aligns with the middle
     //of the a compressed message set in the leader (which, when forwarded, will result in offsets going backwards)
@@ -242,7 +239,7 @@ class EpochDrivenReplicationProtocolAcceptanceTest extends ZooKeeperTestHarness
 
     //Search to see if we have non-monotonic offsets in the log
     startConsumer()
-    val records = consumer.poll(1000).asScala
+    val records = TestUtils.pollUntilAtLeastNumRecords(consumer, 100)
     var prevOffset = -1L
     records.foreach { r =>
       assertTrue(s"Offset $prevOffset came before ${r.offset} ", r.offset > prevOffset)
@@ -265,8 +262,8 @@ class EpochDrivenReplicationProtocolAcceptanceTest extends ZooKeeperTestHarness
     brokers = (100 to 101).map(createBroker(_))
 
     //A single partition topic with 2 replicas
-    adminZkClient.createOrUpdateTopicPartitionAssignmentPathInZK(topic, Map(0 -> Seq(100, 101)))
-    producer = createProducer()
+    TestUtils.createTopic(zkClient, topic, Map(0 -> Seq(100, 101)), brokers)
+    producer = createProducer
 
     //Kick off with a single record
     producer.send(new ProducerRecord(topic, 0, null, msg)).get
@@ -274,7 +271,7 @@ class EpochDrivenReplicationProtocolAcceptanceTest extends ZooKeeperTestHarness
 
     //Now invoke the fast leader change bug
     (0 until 5).foreach { i =>
-      val leaderId = zkUtils.getLeaderForPartition(topic, 0).get
+      val leaderId = zkClient.getLeaderForPartition(new TopicPartition(topic, 0)).get
       val leader = brokers.filter(_.config.brokerId == leaderId)(0)
       val follower = brokers.filter(_.config.brokerId != leaderId)(0)
 
@@ -298,10 +295,90 @@ class EpochDrivenReplicationProtocolAcceptanceTest extends ZooKeeperTestHarness
     }
   }
 
+  @Test
+  def logsShouldNotDivergeOnUncleanLeaderElections(): Unit = {
+
+    // Given two brokers, unclean leader election is enabled
+    brokers = (100 to 101).map(createBroker(_, enableUncleanLeaderElection = true))
+
+    // A single partition topic with 2 replicas, min.isr = 1
+    TestUtils.createTopic(zkClient, topic, Map(0 -> Seq(100, 101)), brokers,
+      CoreUtils.propsWith((KafkaConfig.MinInSyncReplicasProp, "1")))
+
+    producer = TestUtils.createProducer(getBrokerListStrFromServers(brokers), acks = 1)
+
+    // Write one message while both brokers are up
+    (0 until 1).foreach { i =>
+      producer.send(new ProducerRecord(topic, 0, null, msg))
+      producer.flush()}
+
+    // Since we use producer with acks = 1, make sure that logs match for the first epoch
+    waitForLogsToMatch(brokers(0), brokers(1))
+
+    // shutdown broker 100
+    brokers(0).shutdown()
+
+    //Write 1 message
+    (0 until 1).foreach { i =>
+      producer.send(new ProducerRecord(topic, 0, null, msg))
+      producer.flush()}
+
+    brokers(1).shutdown()
+    brokers(0).startup()
+
+    //Bounce the producer (this is required, probably because the broker port changes on restart?)
+    producer.close()
+    producer = TestUtils.createProducer(getBrokerListStrFromServers(brokers), acks = 1)
+
+    //Write 3 messages
+    (0 until 3).foreach { i =>
+      producer.send(new ProducerRecord(topic, 0, null, msgBigger))
+      producer.flush()}
+
+    brokers(0).shutdown()
+    brokers(1).startup()
+
+    //Bounce the producer (this is required, probably because the broker port changes on restart?)
+    producer.close()
+    producer = TestUtils.createProducer(getBrokerListStrFromServers(brokers), acks = 1)
+
+    //Write 1 message
+    (0 until 1).foreach { i =>
+      producer.send(new ProducerRecord(topic, 0, null, msg))
+      producer.flush()}
+
+    brokers(1).shutdown()
+    brokers(0).startup()
+
+    //Bounce the producer (this is required, probably because the broker port changes on restart?)
+    producer.close()
+    producer = TestUtils.createProducer(getBrokerListStrFromServers(brokers), acks = 1)
+
+    //Write 2 messages
+    (0 until 2).foreach { i =>
+      producer.send(new ProducerRecord(topic, 0, null, msgBigger))
+      producer.flush()}
+
+    printSegments()
+
+    brokers(1).startup()
+
+    waitForLogsToMatch(brokers(0), brokers(1))
+    printSegments()
+
+    def crcSeq(broker: KafkaServer, partition: Int = 0): Seq[Long] = {
+      val batches = getLog(broker, partition).activeSegment.read(0, Integer.MAX_VALUE)
+        .records.batches().asScala.toSeq
+      batches.map(_.checksum)
+    }
+    assertTrue(s"Logs on Broker 100 and Broker 101 should match",
+               crcSeq(brokers(0)) == crcSeq(brokers(1)))
+  }
+
   private def log(leader: KafkaServer, follower: KafkaServer): Unit = {
     info(s"Bounce complete for follower ${follower.config.brokerId}")
-    info(s"Leader: leo${leader.config.brokerId}: " + getLog(leader, 0).logEndOffset + " cache: " + epochCache(leader).epochEntries())
-    info(s"Follower: leo${follower.config.brokerId}: " + getLog(follower, 0).logEndOffset + " cache: " + epochCache(follower).epochEntries())
+    info(s"Leader: leo${leader.config.brokerId}: " + getLog(leader, 0).logEndOffset + " cache: " + epochCache(leader).epochEntries)
+    info(s"Follower: leo${follower.config.brokerId}: " + getLog(follower, 0).logEndOffset + " cache: " + epochCache(follower).epochEntries)
   }
 
   private def waitForLogsToMatch(b1: KafkaServer, b2: KafkaServer, partition: Int = 0): Unit = {
@@ -320,7 +397,7 @@ class EpochDrivenReplicationProtocolAcceptanceTest extends ZooKeeperTestHarness
     consumerConfig.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, getBrokerListStrFromServers(brokers))
     consumerConfig.put(ConsumerConfig.FETCH_MAX_BYTES_CONFIG, String.valueOf(getLogFile(brokers(1), 0).length() * 2))
     consumerConfig.put(ConsumerConfig.MAX_PARTITION_FETCH_BYTES_CONFIG, String.valueOf(getLogFile(brokers(1), 0).length() * 2))
-    consumer = new KafkaConsumer(consumerConfig, new StubDeserializer, new StubDeserializer)
+    consumer = new KafkaConsumer(consumerConfig, new ByteArrayDeserializer, new ByteArrayDeserializer)
     consumer.assign(List(new TopicPartition(topic, 0)).asJava)
     consumer.seek(new TopicPartition(topic, 0), 0)
     consumer
@@ -333,12 +410,12 @@ class EpochDrivenReplicationProtocolAcceptanceTest extends ZooKeeperTestHarness
     writable.close()
   }
 
-  private def bufferingProducer(): KafkaProducer[Array[Byte], Array[Byte]] = {
-    createNewProducer(getBrokerListStrFromServers(brokers), retries = 5, acks = -1, lingerMs = 10000,
-      props = Option(CoreUtils.propsWith(
-        (ProducerConfig.BATCH_SIZE_CONFIG, String.valueOf(msg.length * 1000))
-        , (ProducerConfig.COMPRESSION_TYPE_CONFIG, "snappy")
-      )))
+  private def createBufferingProducer: KafkaProducer[Array[Byte], Array[Byte]] = {
+    TestUtils.createProducer(getBrokerListStrFromServers(brokers),
+      acks = -1,
+      lingerMs = 10000,
+      batchSize = msg.length * 1000,
+      compressionType = "snappy")
   }
 
   private def getLogFile(broker: KafkaServer, partition: Int): File = {
@@ -355,54 +432,43 @@ class EpochDrivenReplicationProtocolAcceptanceTest extends ZooKeeperTestHarness
     follower.shutdown()
     follower.startup()
     producer.close()
-    producer = createProducer() //TODO not sure why we need to recreate the producer, but it doesn't reconnect if we don't
+    producer = createProducer //TODO not sure why we need to recreate the producer, but it doesn't reconnect if we don't
   }
 
-  private def epochCache(broker: KafkaServer): LeaderEpochFileCache = {
-    getLog(broker, 0).leaderEpochCache.asInstanceOf[LeaderEpochFileCache]
-  }
+  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")
   }
 
-  private def createProducer(): KafkaProducer[Array[Byte], Array[Byte]] = {
-    createNewProducer(getBrokerListStrFromServers(brokers), retries = 5, acks = -1)
+  private def createProducer: KafkaProducer[Array[Byte], Array[Byte]] = {
+    TestUtils.createProducer(getBrokerListStrFromServers(brokers), acks = -1)
   }
 
-  private def leader(): KafkaServer = {
+  private def leader: KafkaServer = {
     assertEquals(2, brokers.size)
-    val leaderId = zkUtils.getLeaderForPartition(topic, 0).get
-    brokers.filter(_.config.brokerId == leaderId)(0)
+    val leaderId = zkClient.getLeaderForPartition(new TopicPartition(topic, 0)).get
+    brokers.filter(_.config.brokerId == leaderId).head
   }
 
-  private def follower(): KafkaServer = {
+  private def follower: KafkaServer = {
     assertEquals(2, brokers.size)
-    val leader = zkUtils.getLeaderForPartition(topic, 0).get
-    brokers.filter(_.config.brokerId != leader)(0)
+    val leader = zkClient.getLeaderForPartition(new TopicPartition(topic, 0)).get
+    brokers.filter(_.config.brokerId != leader).head
   }
 
-  private def createBroker(id: Int): KafkaServer = {
+  private def createBroker(id: Int, enableUncleanLeaderElection: Boolean = false): KafkaServer = {
     val config = createBrokerConfig(id, zkConnect)
-    if(!KIP_101_ENABLED) {
-      config.setProperty(KafkaConfig.InterBrokerProtocolVersionProp, KAFKA_0_11_0_IV1.version)
-      config.setProperty(KafkaConfig.LogMessageFormatVersionProp, KAFKA_0_11_0_IV1.version)
-    }
+    config.setProperty(KafkaConfig.InterBrokerProtocolVersionProp, apiVersion.version)
+    config.setProperty(KafkaConfig.LogMessageFormatVersionProp, apiVersion.version)
+    config.setProperty(KafkaConfig.UncleanLeaderElectionEnableProp, enableUncleanLeaderElection.toString)
     createServer(fromProps(config))
   }
-
-  private class StubDeserializer extends Deserializer[Array[Byte]] {
-    override def configure(configs: java.util.Map[String, _], isKey: Boolean): Unit = {}
-
-    override def deserialize(topic: String, data: Array[Byte]): Array[Byte] = { data }
-
-    override def close(): Unit = {}
-  }
 }
diff --git a/core/src/test/scala/unit/kafka/server/epoch/EpochDrivenReplicationProtocolAcceptanceWithIbp26Test.scala b/core/src/test/scala/unit/kafka/server/epoch/EpochDrivenReplicationProtocolAcceptanceWithIbp26Test.scala
new file mode 100644
index 0000000000000..2ad4776bb2ca3
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/server/epoch/EpochDrivenReplicationProtocolAcceptanceWithIbp26Test.scala
@@ -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 kafka.server.epoch
+
+import kafka.api.KAFKA_2_6_IV0
+
+/**
+ * With IBP 2.7 onwards, we truncate based on diverging epochs returned in fetch responses.
+ * EpochDrivenReplicationProtocolAcceptanceTest tests epochs with latest version. This test
+ * verifies that we handle older IBP versions with truncation on leader/follower change correctly.
+ */
+class EpochDrivenReplicationProtocolAcceptanceWithIbp26Test extends EpochDrivenReplicationProtocolAcceptanceTest {
+  override val apiVersion = KAFKA_2_6_IV0
+}
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 8460fe4c54c7c..8d91c6c3870d0 100644
--- a/core/src/test/scala/unit/kafka/server/epoch/LeaderEpochFileCacheTest.scala
+++ b/core/src/test/scala/unit/kafka/server/epoch/LeaderEpochFileCacheTest.scala
@@ -16,199 +16,203 @@
   */
 
 package kafka.server.epoch
+
 import java.io.File
 
-import kafka.server.LogOffsetMetadata
+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
+import org.apache.kafka.common.requests.OffsetsForLeaderEpochResponse.{UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET}
 import org.apache.kafka.common.TopicPartition
 import org.junit.Assert._
-import org.junit.{Before, Test}
-
-import scala.collection.mutable.ListBuffer
+import org.junit.Test
 
 /**
   * Unit test for the LeaderEpochFileCache.
   */
 class LeaderEpochFileCacheTest {
   val tp = new TopicPartition("TestTopic", 5)
-  var checkpoint: LeaderEpochCheckpoint = _
+  private var logEndOffset = 0L
+  private val checkpoint: LeaderEpochCheckpoint = new LeaderEpochCheckpoint {
+    private var epochs: Seq[EpochEntry] = Seq()
+    override def write(epochs: Iterable[EpochEntry]): Unit = this.epochs = epochs.toSeq
+    override def read(): Seq[EpochEntry] = this.epochs
+  }
+  private val cache = new LeaderEpochFileCache(tp, () => logEndOffset, checkpoint)
 
   @Test
-  def shouldAddEpochAndMessageOffsetToCache() = {
-    var leo = 0
-    def leoFinder() = new LogOffsetMetadata(leo)
+  def testPreviousEpoch(): Unit = {
+    assertEquals(None, cache.previousEpoch)
 
-    //Given
-    val cache = new LeaderEpochFileCache(tp, () => leoFinder, checkpoint)
+    cache.assign(epoch = 2, startOffset = 10)
+    assertEquals(None, cache.previousEpoch)
+
+    cache.assign(epoch = 4, startOffset = 15)
+    assertEquals(Some(2), cache.previousEpoch)
+
+    cache.assign(epoch = 10, startOffset = 20)
+    assertEquals(Some(4), cache.previousEpoch)
+
+    cache.truncateFromEnd(18)
+    assertEquals(Some(2), cache.previousEpoch)
+  }
 
+  @Test
+  def shouldAddEpochAndMessageOffsetToCache() = {
     //When
-    cache.assign(epoch = 2, offset = 10)
-    leo = 11
+    cache.assign(epoch = 2, startOffset = 10)
+    logEndOffset = 11
 
     //Then
-    assertEquals(2, cache.latestEpoch())
-    assertEquals(EpochEntry(2, 10), cache.epochEntries()(0))
-    assertEquals(11, cache.endOffsetFor(2)) //should match leo
+    assertEquals(Some(2), cache.latestEpoch)
+    assertEquals(EpochEntry(2, 10), cache.epochEntries(0))
+    assertEquals((2, logEndOffset), cache.endOffsetFor(2)) //should match logEndOffset
   }
 
   @Test
   def shouldReturnLogEndOffsetIfLatestEpochRequested() = {
-    var leo = 0
-    def leoFinder() = new LogOffsetMetadata(leo)
-
-    //Given
-    val cache = new LeaderEpochFileCache(tp, () => leoFinder, checkpoint)
-
     //When just one epoch
-    cache.assign(epoch = 2, offset = 11)
-    cache.assign(epoch = 2, offset = 12)
-    leo = 14
+    cache.assign(epoch = 2, startOffset = 11)
+    cache.assign(epoch = 2, startOffset = 12)
+    logEndOffset = 14
 
     //Then
-    assertEquals(14, cache.endOffsetFor(2))
+    assertEquals((2, logEndOffset), cache.endOffsetFor(2))
   }
 
   @Test
   def shouldReturnUndefinedOffsetIfUndefinedEpochRequested() = {
-    def leoFinder() = new LogOffsetMetadata(0)
+    val expectedEpochEndOffset = (UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET)
 
-    //Given cache with some data on leader
-    val cache = new LeaderEpochFileCache(tp, () => leoFinder, checkpoint)
-    cache.assign(epoch = 2, offset = 11)
-    cache.assign(epoch = 3, offset = 12)
+    // assign couple of epochs
+    cache.assign(epoch = 2, startOffset = 11)
+    cache.assign(epoch = 3, startOffset = 12)
 
     //When (say a bootstraping follower) sends request for UNDEFINED_EPOCH
-    val offsetFor = cache.endOffsetFor(UNDEFINED_EPOCH)
+    val epochAndOffsetFor = cache.endOffsetFor(UNDEFINED_EPOCH)
 
     //Then
-    assertEquals(UNDEFINED_EPOCH_OFFSET, offsetFor)
+    assertEquals("Expected undefined epoch and offset if undefined epoch requested. Cache not empty.",
+                 expectedEpochEndOffset, epochAndOffsetFor)
   }
 
   @Test
   def shouldNotOverwriteLogEndOffsetForALeaderEpochOnceItHasBeenAssigned() = {
-    var leo = 0
-    def leoFinder() = new LogOffsetMetadata(leo)
-
     //Given
-    leo = 9
-    val cache = new LeaderEpochFileCache(tp, () => leoFinder, checkpoint)
+    logEndOffset = 9
 
-    cache.assign(2, leo)
+    cache.assign(2, logEndOffset)
 
     //When called again later
     cache.assign(2, 10)
 
     //Then the offset should NOT have been updated
-    assertEquals(leo, cache.epochEntries()(0).startOffset)
+    assertEquals(logEndOffset, cache.epochEntries(0).startOffset)
+    assertEquals(ListBuffer(EpochEntry(2, 9)), cache.epochEntries)
   }
 
   @Test
-  def shouldAllowLeaderEpochToChangeEvenIfOffsetDoesNot() = {
-    def leoFinder() = new LogOffsetMetadata(0)
-
+  def shouldEnforceMonotonicallyIncreasingStartOffsets() = {
     //Given
-    val cache = new LeaderEpochFileCache(tp, () => leoFinder, checkpoint)
     cache.assign(2, 9)
 
     //When update epoch new epoch but same offset
     cache.assign(3, 9)
 
     //Then epoch should have been updated
-    assertEquals(ListBuffer(EpochEntry(2, 9), EpochEntry(3, 9)), cache.epochEntries())
+    assertEquals(ListBuffer(EpochEntry(3, 9)), cache.epochEntries)
   }
   
   @Test
   def shouldNotOverwriteOffsetForALeaderEpochOnceItHasBeenAssigned() = {
-    //Given
-    val cache = new LeaderEpochFileCache(tp, () => new LogOffsetMetadata(0), checkpoint)
     cache.assign(2, 6)
 
     //When called again later with a greater offset
     cache.assign(2, 10)
 
     //Then later update should have been ignored
-    assertEquals(6, cache.epochEntries()(0).startOffset)
+    assertEquals(6, cache.epochEntries(0).startOffset)
   }
 
   @Test
-  def shouldReturnUnsupportedIfNoEpochRecorded(){
-    def leoFinder() = new LogOffsetMetadata(0)
-
-    //Given
-    val cache = new LeaderEpochFileCache(tp, () => leoFinder, checkpoint)
-
+  def shouldReturnUnsupportedIfNoEpochRecorded(): Unit = {
     //Then
-    assertEquals(UNDEFINED_EPOCH_OFFSET, cache.endOffsetFor(0))
+    assertEquals((UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET), cache.endOffsetFor(0))
   }
 
   @Test
-  def shouldReturnUnsupportedIfRequestedEpochLessThanFirstEpoch(){
-    def leoFinder() = new LogOffsetMetadata(0)
+  def shouldReturnUnsupportedIfNoEpochRecordedAndUndefinedEpochRequested(): Unit = {
+    logEndOffset = 73
 
-    //Given
-    val cache = new LeaderEpochFileCache(tp, () => leoFinder, checkpoint)
+    //When (say a follower on older message format version) sends request for UNDEFINED_EPOCH
+    val offsetFor = cache.endOffsetFor(UNDEFINED_EPOCH)
 
-    cache.assign(epoch = 5, offset = 11)
-    cache.assign(epoch = 6, offset = 12)
-    cache.assign(epoch = 7, offset = 13)
+    //Then
+    assertEquals("Expected undefined epoch and offset if undefined epoch requested. Empty cache.",
+                 (UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET), offsetFor)
+  }
+
+  @Test
+  def shouldReturnFirstEpochIfRequestedEpochLessThanFirstEpoch(): Unit = {
+    cache.assign(epoch = 5, startOffset = 11)
+    cache.assign(epoch = 6, startOffset = 12)
+    cache.assign(epoch = 7, startOffset = 13)
 
     //When
-    val offset = cache.endOffsetFor(5 - 1)
+    val epochAndOffset = cache.endOffsetFor(4)
 
     //Then
-    assertEquals(UNDEFINED_EPOCH_OFFSET, offset)
+    assertEquals((4, 11), epochAndOffset)
   }
 
   @Test
-  def shouldGetFirstOffsetOfSubsequentEpochWhenOffsetRequestedForPreviousEpoch() = {
-    var leo = 0
-    def leoFinder() = new LogOffsetMetadata(leo)
+  def shouldTruncateIfMatchingEpochButEarlierStartingOffset(): Unit = {
+    cache.assign(epoch = 5, startOffset = 11)
+    cache.assign(epoch = 6, startOffset = 12)
+    cache.assign(epoch = 7, startOffset = 13)
 
-    //Given
-    val cache = new LeaderEpochFileCache(tp, () => leoFinder, checkpoint)
+    // epoch 7 starts at an earlier offset
+    cache.assign(epoch = 7, startOffset = 12)
+
+    assertEquals((5, 12), cache.endOffsetFor(5))
+    assertEquals((5, 12), cache.endOffsetFor(6))
+  }
 
+  @Test
+  def shouldGetFirstOffsetOfSubsequentEpochWhenOffsetRequestedForPreviousEpoch() = {
     //When several epochs
-    cache.assign(epoch = 1, offset = 11)
-    cache.assign(epoch = 1, offset = 12)
-    cache.assign(epoch = 2, offset = 13)
-    cache.assign(epoch = 2, offset = 14)
-    cache.assign(epoch = 3, offset = 15)
-    cache.assign(epoch = 3, offset = 16)
-    leo = 17
+    cache.assign(epoch = 1, startOffset = 11)
+    cache.assign(epoch = 1, startOffset = 12)
+    cache.assign(epoch = 2, startOffset = 13)
+    cache.assign(epoch = 2, startOffset = 14)
+    cache.assign(epoch = 3, startOffset = 15)
+    cache.assign(epoch = 3, startOffset = 16)
+    logEndOffset = 17
 
     //Then get the start offset of the next epoch
-    assertEquals(15, cache.endOffsetFor(2))
+    assertEquals((2, 15), cache.endOffsetFor(2))
   }
 
   @Test
-  def shouldReturnNextAvailableEpochIfThereIsNoExactEpochForTheOneRequested(){
-    def leoFinder() = new LogOffsetMetadata(0)
-
-    //Given
-    val cache = new LeaderEpochFileCache(tp, () => leoFinder, checkpoint)
-
+  def shouldReturnNextAvailableEpochIfThereIsNoExactEpochForTheOneRequested(): Unit = {
     //When
-    cache.assign(epoch = 0, offset = 10)
-    cache.assign(epoch = 2, offset = 13)
-    cache.assign(epoch = 4, offset = 17)
+    cache.assign(epoch = 0, startOffset = 10)
+    cache.assign(epoch = 2, startOffset = 13)
+    cache.assign(epoch = 4, startOffset = 17)
 
     //Then
-    assertEquals(13, cache.endOffsetFor(requestedEpoch = 1))
-    assertEquals(17, cache.endOffsetFor(requestedEpoch = 2))
+    assertEquals((0, 13), cache.endOffsetFor(requestedEpoch = 1))
+    assertEquals((2, 17), cache.endOffsetFor(requestedEpoch = 2))
+    assertEquals((2, 17), cache.endOffsetFor(requestedEpoch = 3))
   }
 
   @Test
   def shouldNotUpdateEpochAndStartOffsetIfItDidNotChange() = {
-    def leoFinder() = new LogOffsetMetadata(0)
-
-    //Given
-    val cache = new LeaderEpochFileCache(tp, () => leoFinder, checkpoint)
-
     //When
-    cache.assign(epoch = 2, offset = 6)
-    cache.assign(epoch = 2, offset = 7)
+    cache.assign(epoch = 2, startOffset = 6)
+    cache.assign(epoch = 2, startOffset = 7)
 
     //Then
     assertEquals(1, cache.epochEntries.size)
@@ -217,50 +221,39 @@ class LeaderEpochFileCacheTest {
 
   @Test
   def shouldReturnInvalidOffsetIfEpochIsRequestedWhichIsNotCurrentlyTracked(): Unit = {
-    val leo = 100
-    def leoFinder() = new LogOffsetMetadata(leo)
-
-    //Given
-    val cache = new LeaderEpochFileCache(tp, () => leoFinder, checkpoint)
+    logEndOffset = 100
 
     //When
-    cache.assign(epoch = 2, offset = 100)
+    cache.assign(epoch = 2, startOffset = 100)
 
     //Then
-    assertEquals(UNDEFINED_EPOCH_OFFSET, cache.endOffsetFor(3))
+    assertEquals((UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET), cache.endOffsetFor(3))
   }
 
   @Test
   def shouldSupportEpochsThatDoNotStartFromZero(): Unit = {
-    var leo = 0
-    def leoFinder() = new LogOffsetMetadata(leo)
-
-    //Given
-    val cache = new LeaderEpochFileCache(tp, () => leoFinder, checkpoint)
-
     //When
-    cache.assign(epoch = 2, offset = 6)
-    leo = 7
+    cache.assign(epoch = 2, startOffset = 6)
+    logEndOffset = 7
 
     //Then
-    assertEquals(leo, cache.endOffsetFor(2))
+    assertEquals((2, logEndOffset), cache.endOffsetFor(2))
     assertEquals(1, cache.epochEntries.size)
-    assertEquals(EpochEntry(2, 6), cache.epochEntries()(0))
+    assertEquals(EpochEntry(2, 6), cache.epochEntries(0))
   }
 
   @Test
-  def shouldPersistEpochsBetweenInstances(){
-    def leoFinder() = new LogOffsetMetadata(0)
+  def shouldPersistEpochsBetweenInstances(): Unit = {
     val checkpointPath = TestUtils.tempFile().getAbsolutePath
-    checkpoint = new LeaderEpochCheckpointFile(new File(checkpointPath))
+    val checkpoint = new LeaderEpochCheckpointFile(new File(checkpointPath))
 
     //Given
-    val cache = new LeaderEpochFileCache(tp, () => leoFinder, checkpoint)
-    cache.assign(epoch = 2, offset = 6)
+    val cache = new LeaderEpochFileCache(tp, () => logEndOffset, checkpoint)
+    cache.assign(epoch = 2, startOffset = 6)
 
     //When
     val checkpoint2 = new LeaderEpochCheckpointFile(new File(checkpointPath))
-    val cache2 = new LeaderEpochFileCache(tp, () => leoFinder, checkpoint2)
+    val cache2 = new LeaderEpochFileCache(tp, () => logEndOffset, checkpoint2)
 
     //Then
     assertEquals(1, cache2.epochEntries.size)
@@ -268,140 +261,118 @@ class LeaderEpochFileCacheTest {
   }
 
   @Test
-  def shouldNotLetEpochGoBackwardsEvenIfMessageEpochsDo(): Unit = {
-    var leo = 0
-    def leoFinder() = new LogOffsetMetadata(leo)
-
-    val cache = new LeaderEpochFileCache(tp, () => leoFinder, checkpoint)
-
+  def shouldEnforceMonotonicallyIncreasingEpochs(): Unit = {
     //Given
-    cache.assign(epoch = 1, offset = 5); leo = 6
-    cache.assign(epoch = 2, offset = 6); leo = 7
-
-    //When we update an epoch in the past with an earlier offset
-    cache.assign(epoch = 1, offset = 7); leo = 8
+    cache.assign(epoch = 1, startOffset = 5); logEndOffset = 6
+    cache.assign(epoch = 2, startOffset = 6); logEndOffset = 7
 
-    //Then epoch should not be changed
-    assertEquals(2, cache.latestEpoch())
+    //When we update an epoch in the past with a different offset, the log has already reached
+    //an inconsistent state. Our options are either to raise an error, ignore the new append,
+    //or truncate the cached epochs to the point of conflict. We take this latter approach in
+    //order to guarantee that epochs and offsets in the cache increase monotonically, which makes
+    //the search logic simpler to reason about.
+    cache.assign(epoch = 1, startOffset = 7); logEndOffset = 8
 
-    //Then end offset for epoch 1 shouldn't have changed
-    assertEquals(6, cache.endOffsetFor(1))
+    //Then later epochs will be removed
+    assertEquals(Some(1), cache.latestEpoch)
 
-    //Then end offset for epoch 2 has to be the offset of the epoch 1 message (I can't thing of a better option)
-    assertEquals(8, cache.endOffsetFor(2))
+    //Then end offset for epoch 1 will have changed
+    assertEquals((1, 8), cache.endOffsetFor(1))
 
-    //Epoch history shouldn't have changed
-    assertEquals(EpochEntry(1, 5), cache.epochEntries()(0))
-    assertEquals(EpochEntry(2, 6), cache.epochEntries()(1))
+    //Then end offset for epoch 2 is now undefined
+    assertEquals((UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET), cache.endOffsetFor(2))
+    assertEquals(EpochEntry(1, 7), cache.epochEntries(0))
   }
 
   @Test
-  def shouldNotLetOffsetsGoBackwardsEvenIfEpochsProgress() = {
-    def leoFinder() = new LogOffsetMetadata(0)
-
-    //Given
-    val cache = new LeaderEpochFileCache(tp, () => leoFinder, checkpoint)
-
+  def shouldEnforceOffsetsIncreaseMonotonically() = {
     //When epoch goes forward but offset goes backwards
-    cache.assign(epoch = 2, offset = 6)
-    cache.assign(epoch = 3, offset = 5)
+    cache.assign(epoch = 2, startOffset = 6)
+    cache.assign(epoch = 3, startOffset = 5)
 
-    //Then latter assign should be ignored
-    assertEquals(EpochEntry(2, 6), cache.epochEntries.toList(0))
+    //The last assignment wins and the conflicting one is removed from the log
+    assertEquals(EpochEntry(3, 5), cache.epochEntries.toList(0))
   }
 
   @Test
   def shouldIncreaseAndTrackEpochsAsLeadersChangeManyTimes(): Unit = {
-    var leo = 0
-    def leoFinder() = new LogOffsetMetadata(leo)
-
     //Given
-    val cache = new LeaderEpochFileCache(tp, () => leoFinder, checkpoint)
-    cache.assign(epoch = 0, offset = 0) //leo=0
+    cache.assign(epoch = 0, startOffset = 0) //logEndOffset=0
 
     //When
-    cache.assign(epoch = 1, offset = 0) //leo=0
+    cache.assign(epoch = 1, startOffset = 0) //logEndOffset=0
 
     //Then epoch should go up
-    assertEquals(1, cache.latestEpoch())
+    assertEquals(Some(1), cache.latestEpoch)
     //offset for 1 should still be 0
-    assertEquals(0, cache.endOffsetFor(1))
+    assertEquals((1, 0), cache.endOffsetFor(1))
     //offset for epoch 0 should still be 0
-    assertEquals(0, cache.endOffsetFor(0))
+    assertEquals((0, 0), cache.endOffsetFor(0))
 
     //When we write 5 messages as epoch 1
-    leo = 5
+    logEndOffset = 5
 
-    //Then end offset for epoch(1) should be leo => 5
-    assertEquals(5, cache.endOffsetFor(1))
+    //Then end offset for epoch(1) should be logEndOffset => 5
+    assertEquals((1, 5), cache.endOffsetFor(1))
     //Epoch 0 should still be at offset 0
-    assertEquals(0, cache.endOffsetFor(0))
+    assertEquals((0, 0), cache.endOffsetFor(0))
 
     //When
-    cache.assign(epoch = 2, offset = 5) //leo=5
+    cache.assign(epoch = 2, startOffset = 5) //logEndOffset=5
 
-    leo = 10 //write another 5 messages
+    logEndOffset = 10 //write another 5 messages
 
-    //Then end offset for epoch(2) should be leo => 10
-    assertEquals(10, cache.endOffsetFor(2))
+    //Then end offset for epoch(2) should be logEndOffset => 10
+    assertEquals((2, 10), cache.endOffsetFor(2))
 
     //end offset for epoch(1) should be the start offset of epoch(2) => 5
-    assertEquals(5, cache.endOffsetFor(1))
+    assertEquals((1, 5), cache.endOffsetFor(1))
 
     //epoch (0) should still be 0
-    assertEquals(0, cache.endOffsetFor(0))
+    assertEquals((0, 0), cache.endOffsetFor(0))
   }
 
   @Test
   def shouldIncreaseAndTrackEpochsAsFollowerReceivesManyMessages(): Unit = {
-    var leo = 0
-    def leoFinder() = new LogOffsetMetadata(leo)
-
-    //When new
-    val cache = new LeaderEpochFileCache(tp, () => leoFinder, checkpoint)
-
     //When Messages come in
-    cache.assign(epoch = 0, offset = 0); leo = 1
-    cache.assign(epoch = 0, offset = 1); leo = 2
-    cache.assign(epoch = 0, offset = 2); leo = 3
+    cache.assign(epoch = 0, startOffset = 0); logEndOffset = 1
+    cache.assign(epoch = 0, startOffset = 1); logEndOffset = 2
+    cache.assign(epoch = 0, startOffset = 2); logEndOffset = 3
 
     //Then epoch should stay, offsets should grow
-    assertEquals(0, cache.latestEpoch())
-    assertEquals(leo, cache.endOffsetFor(0))
+    assertEquals(Some(0), cache.latestEpoch)
+    assertEquals((0, logEndOffset), cache.endOffsetFor(0))
 
     //When messages arrive with greater epoch
-    cache.assign(epoch = 1, offset = 3); leo = 4
-    cache.assign(epoch = 1, offset = 4); leo = 5
-    cache.assign(epoch = 1, offset = 5); leo = 6
+    cache.assign(epoch = 1, startOffset = 3); logEndOffset = 4
+    cache.assign(epoch = 1, startOffset = 4); logEndOffset = 5
+    cache.assign(epoch = 1, startOffset = 5); logEndOffset = 6
 
-    assertEquals(1, cache.latestEpoch())
-    assertEquals(leo, cache.endOffsetFor(1))
+    assertEquals(Some(1), cache.latestEpoch)
+    assertEquals((1, logEndOffset), cache.endOffsetFor(1))
 
     //When
-    cache.assign(epoch = 2, offset = 6); leo = 7
-    cache.assign(epoch = 2, offset = 7); leo = 8
-    cache.assign(epoch = 2, offset = 8); leo = 9
+    cache.assign(epoch = 2, startOffset = 6); logEndOffset = 7
+    cache.assign(epoch = 2, startOffset = 7); logEndOffset = 8
+    cache.assign(epoch = 2, startOffset = 8); logEndOffset = 9
 
-    assertEquals(2, cache.latestEpoch())
-    assertEquals(leo, cache.endOffsetFor(2))
+    assertEquals(Some(2), cache.latestEpoch)
+    assertEquals((2, logEndOffset), cache.endOffsetFor(2))
 
     //Older epochs should return the start offset of the first message in the subsequent epoch.
-    assertEquals(3, cache.endOffsetFor(0))
-    assertEquals(6, cache.endOffsetFor(1))
+    assertEquals((0, 3), cache.endOffsetFor(0))
+    assertEquals((1, 6), cache.endOffsetFor(1))
   }
 
   @Test
   def shouldDropEntriesOnEpochBoundaryWhenRemovingLatestEntries(): Unit = {
-    def leoFinder() = new LogOffsetMetadata(0)
-
     //Given
-    val cache = new LeaderEpochFileCache(tp, () => leoFinder, checkpoint)
-    cache.assign(epoch = 2, offset = 6)
-    cache.assign(epoch = 3, offset = 8)
-    cache.assign(epoch = 4, offset = 11)
+    cache.assign(epoch = 2, startOffset = 6)
+    cache.assign(epoch = 3, startOffset = 8)
+    cache.assign(epoch = 4, startOffset = 11)
 
     //When clear latest on epoch boundary
-    cache.clearAndFlushLatest(offset = 8)
+    cache.truncateFromEnd(endOffset = 8)
 
     //Then should remove two latest epochs (remove is inclusive)
     assertEquals(ListBuffer(EpochEntry(2, 6)), cache.epochEntries)
@@ -409,16 +380,13 @@ class LeaderEpochFileCacheTest {
 
   @Test
   def shouldPreserveResetOffsetOnClearEarliestIfOneExists(): Unit = {
-    def leoFinder() = new LogOffsetMetadata(0)
-
     //Given
-    val cache = new LeaderEpochFileCache(tp, () => leoFinder, checkpoint)
-    cache.assign(epoch = 2, offset = 6)
-    cache.assign(epoch = 3, offset = 8)
-    cache.assign(epoch = 4, offset = 11)
+    cache.assign(epoch = 2, startOffset = 6)
+    cache.assign(epoch = 3, startOffset = 8)
+    cache.assign(epoch = 4, startOffset = 11)
 
     //When reset to offset ON epoch boundary
-    cache.clearAndFlushEarliest(offset = 8)
+    cache.truncateFromStart(startOffset = 8)
 
     //Then should preserve (3, 8)
     assertEquals(ListBuffer(EpochEntry(3, 8), EpochEntry(4, 11)), cache.epochEntries)
@@ -426,16 +394,13 @@ class LeaderEpochFileCacheTest {
 
   @Test
   def shouldUpdateSavedOffsetWhenOffsetToClearToIsBetweenEpochs(): Unit = {
-    def leoFinder() = new LogOffsetMetadata(0)
-
     //Given
-    val cache = new LeaderEpochFileCache(tp, () => leoFinder, checkpoint)
-    cache.assign(epoch = 2, offset = 6)
-    cache.assign(epoch = 3, offset = 8)
-    cache.assign(epoch = 4, offset = 11)
+    cache.assign(epoch = 2, startOffset = 6)
+    cache.assign(epoch = 3, startOffset = 8)
+    cache.assign(epoch = 4, startOffset = 11)
 
     //When reset to offset BETWEEN epoch boundaries
-    cache.clearAndFlushEarliest(offset = 9)
+    cache.truncateFromStart(startOffset = 9)
 
     //Then we should retain epoch 3, but update it's offset to 9 as 8 has been removed
     assertEquals(ListBuffer(EpochEntry(3, 9), EpochEntry(4, 11)), cache.epochEntries)
@@ -443,16 +408,13 @@ class LeaderEpochFileCacheTest {
 
   @Test
   def shouldNotClearAnythingIfOffsetToEarly(): Unit = {
-    def leoFinder() = new LogOffsetMetadata(0)
-
     //Given
-    val cache = new LeaderEpochFileCache(tp, () => leoFinder, checkpoint)
-    cache.assign(epoch = 2, offset = 6)
-    cache.assign(epoch = 3, offset = 8)
-    cache.assign(epoch = 4, offset = 11)
+    cache.assign(epoch = 2, startOffset = 6)
+    cache.assign(epoch = 3, startOffset = 8)
+    cache.assign(epoch = 4, startOffset = 11)
 
     //When reset to offset before first epoch offset
-    cache.clearAndFlushEarliest(offset = 1)
+    cache.truncateFromStart(startOffset = 1)
 
     //Then nothing should change
     assertEquals(ListBuffer(EpochEntry(2, 6),EpochEntry(3, 8), EpochEntry(4, 11)), cache.epochEntries)
@@ -460,16 +422,13 @@ class LeaderEpochFileCacheTest {
 
   @Test
   def shouldNotClearAnythingIfOffsetToFirstOffset(): Unit = {
-    def leoFinder() = new LogOffsetMetadata(0)
-
     //Given
-    val cache = new LeaderEpochFileCache(tp, () => leoFinder, checkpoint)
-    cache.assign(epoch = 2, offset = 6)
-    cache.assign(epoch = 3, offset = 8)
-    cache.assign(epoch = 4, offset = 11)
+    cache.assign(epoch = 2, startOffset = 6)
+    cache.assign(epoch = 3, startOffset = 8)
+    cache.assign(epoch = 4, startOffset = 11)
 
     //When reset to offset on earliest epoch boundary
-    cache.clearAndFlushEarliest(offset = 6)
+    cache.truncateFromStart(startOffset = 6)
 
     //Then nothing should change
     assertEquals(ListBuffer(EpochEntry(2, 6),EpochEntry(3, 8), EpochEntry(4, 11)), cache.epochEntries)
@@ -477,16 +436,13 @@ class LeaderEpochFileCacheTest {
 
   @Test
   def shouldRetainLatestEpochOnClearAllEarliest(): Unit = {
-    def leoFinder() = new LogOffsetMetadata(0)
-
     //Given
-    val cache = new LeaderEpochFileCache(tp, () => leoFinder, checkpoint)
-    cache.assign(epoch = 2, offset = 6)
-    cache.assign(epoch = 3, offset = 8)
-    cache.assign(epoch = 4, offset = 11)
+    cache.assign(epoch = 2, startOffset = 6)
+    cache.assign(epoch = 3, startOffset = 8)
+    cache.assign(epoch = 4, startOffset = 11)
 
     //When
-    cache.clearAndFlushEarliest(offset = 11)
+    cache.truncateFromStart(startOffset = 11)
 
     //Then retain the last
     assertEquals(ListBuffer(EpochEntry(4, 11)), cache.epochEntries)
@@ -494,16 +450,13 @@ class LeaderEpochFileCacheTest {
 
   @Test
   def shouldUpdateOffsetBetweenEpochBoundariesOnClearEarliest(): Unit = {
-    def leoFinder() = new LogOffsetMetadata(0)
-
     //Given
-    val cache = new LeaderEpochFileCache(tp, () => leoFinder, checkpoint)
-    cache.assign(epoch = 2, offset = 6)
-    cache.assign(epoch = 3, offset = 8)
-    cache.assign(epoch = 4, offset = 11)
+    cache.assign(epoch = 2, startOffset = 6)
+    cache.assign(epoch = 3, startOffset = 8)
+    cache.assign(epoch = 4, startOffset = 11)
 
-    //When we clear from a postition between offset 8 & offset 11
-    cache.clearAndFlushEarliest(offset = 9)
+    //When we clear from a position between offset 8 & offset 11
+    cache.truncateFromStart(startOffset = 9)
 
     //Then we should update the middle epoch entry's offset
     assertEquals(ListBuffer(EpochEntry(3, 9), EpochEntry(4, 11)), cache.epochEntries)
@@ -511,33 +464,27 @@ class LeaderEpochFileCacheTest {
 
   @Test
   def shouldUpdateOffsetBetweenEpochBoundariesOnClearEarliest2(): Unit = {
-    def leoFinder() = new LogOffsetMetadata(0)
-
     //Given
-    val cache = new LeaderEpochFileCache(tp, () => leoFinder, checkpoint)
-    cache.assign(epoch = 0, offset = 0)
-    cache.assign(epoch = 1, offset = 7)
-    cache.assign(epoch = 2, offset = 10)
+    cache.assign(epoch = 0, startOffset = 0)
+    cache.assign(epoch = 1, startOffset = 7)
+    cache.assign(epoch = 2, startOffset = 10)
 
-    //When we clear from a postition between offset 0 & offset 7
-    cache.clearAndFlushEarliest(offset = 5)
+    //When we clear from a position between offset 0 & offset 7
+    cache.truncateFromStart(startOffset = 5)
 
-    //Then we should keeep epoch 0 but update the offset appropriately
+    //Then we should keep epoch 0 but update the offset appropriately
     assertEquals(ListBuffer(EpochEntry(0,5), EpochEntry(1, 7), EpochEntry(2, 10)), cache.epochEntries)
   }
 
   @Test
   def shouldRetainLatestEpochOnClearAllEarliestAndUpdateItsOffset(): Unit = {
-    def leoFinder() = new LogOffsetMetadata(0)
-
     //Given
-    val cache = new LeaderEpochFileCache(tp, () => leoFinder, checkpoint)
-    cache.assign(epoch = 2, offset = 6)
-    cache.assign(epoch = 3, offset = 8)
-    cache.assign(epoch = 4, offset = 11)
+    cache.assign(epoch = 2, startOffset = 6)
+    cache.assign(epoch = 3, startOffset = 8)
+    cache.assign(epoch = 4, startOffset = 11)
 
     //When reset to offset beyond last epoch
-    cache.clearAndFlushEarliest(offset = 15)
+    cache.truncateFromStart(startOffset = 15)
 
     //Then update the last
     assertEquals(ListBuffer(EpochEntry(4, 15)), cache.epochEntries)
@@ -545,51 +492,42 @@ class LeaderEpochFileCacheTest {
 
   @Test
   def shouldDropEntriesBetweenEpochBoundaryWhenRemovingNewest(): Unit = {
-    def leoFinder() = new LogOffsetMetadata(0)
-
     //Given
-    val cache = new LeaderEpochFileCache(tp, () => leoFinder, checkpoint)
-    cache.assign(epoch = 2, offset = 6)
-    cache.assign(epoch = 3, offset = 8)
-    cache.assign(epoch = 4, offset = 11)
+    cache.assign(epoch = 2, startOffset = 6)
+    cache.assign(epoch = 3, startOffset = 8)
+    cache.assign(epoch = 4, startOffset = 11)
 
     //When reset to offset BETWEEN epoch boundaries
-    cache.clearAndFlushLatest(offset = 9)
+    cache.truncateFromEnd(endOffset = 9)
 
     //Then should keep the preceding epochs
-    assertEquals(3, cache.latestEpoch())
+    assertEquals(Some(3), cache.latestEpoch)
     assertEquals(ListBuffer(EpochEntry(2, 6), EpochEntry(3, 8)), cache.epochEntries)
   }
 
   @Test
   def shouldClearAllEntries(): Unit = {
-    def leoFinder() = new LogOffsetMetadata(0)
-
     //Given
-    val cache = new LeaderEpochFileCache(tp, () => leoFinder, checkpoint)
-    cache.assign(epoch = 2, offset = 6)
-    cache.assign(epoch = 3, offset = 8)
-    cache.assign(epoch = 4, offset = 11)
+    cache.assign(epoch = 2, startOffset = 6)
+    cache.assign(epoch = 3, startOffset = 8)
+    cache.assign(epoch = 4, startOffset = 11)
 
-    //When 
+    //When
     cache.clearAndFlush()
 
-    //Then 
+    //Then
     assertEquals(0, cache.epochEntries.size)
   }
 
   @Test
   def shouldNotResetEpochHistoryHeadIfUndefinedPassed(): Unit = {
-    def leoFinder() = new LogOffsetMetadata(0)
-
     //Given
-    val cache = new LeaderEpochFileCache(tp, () => leoFinder, checkpoint)
-    cache.assign(epoch = 2, offset = 6)
-    cache.assign(epoch = 3, offset = 8)
-    cache.assign(epoch = 4, offset = 11)
+    cache.assign(epoch = 2, startOffset = 6)
+    cache.assign(epoch = 3, startOffset = 8)
+    cache.assign(epoch = 4, startOffset = 11)
 
     //When reset to offset on epoch boundary
-    cache.clearAndFlushLatest(offset = UNDEFINED_EPOCH_OFFSET)
+    cache.truncateFromStart(startOffset = UNDEFINED_EPOCH_OFFSET)
 
     //Then should do nothing
     assertEquals(3, cache.epochEntries.size)
@@ -597,16 +535,13 @@ class LeaderEpochFileCacheTest {
 
   @Test
   def shouldNotResetEpochHistoryTailIfUndefinedPassed(): Unit = {
-    def leoFinder() = new LogOffsetMetadata(0)
-
     //Given
-    val cache = new LeaderEpochFileCache(tp, () => leoFinder, checkpoint)
-    cache.assign(epoch = 2, offset = 6)
-    cache.assign(epoch = 3, offset = 8)
-    cache.assign(epoch = 4, offset = 11)
+    cache.assign(epoch = 2, startOffset = 6)
+    cache.assign(epoch = 3, startOffset = 8)
+    cache.assign(epoch = 4, startOffset = 11)
 
     //When reset to offset on epoch boundary
-    cache.clearAndFlushEarliest(offset = UNDEFINED_EPOCH_OFFSET)
+    cache.truncateFromEnd(endOffset = UNDEFINED_EPOCH_OFFSET)
 
     //Then should do nothing
     assertEquals(3, cache.epochEntries.size)
@@ -614,54 +549,26 @@ class LeaderEpochFileCacheTest {
 
   @Test
   def shouldFetchLatestEpochOfEmptyCache(): Unit = {
-    //Given
-    def leoFinder() = new LogOffsetMetadata(0)
-
-    //When
-    val cache = new LeaderEpochFileCache(tp, () => leoFinder, checkpoint)
-
     //Then
-    assertEquals(-1, cache.latestEpoch)
+    assertEquals(None, cache.latestEpoch)
   }
 
   @Test
   def shouldFetchEndOffsetOfEmptyCache(): Unit = {
-    //Given
-    def leoFinder() = new LogOffsetMetadata(0)
-
-    //When
-    val cache = new LeaderEpochFileCache(tp, () => leoFinder, checkpoint)
-
     //Then
-    assertEquals(-1, cache.endOffsetFor(7))
+    assertEquals((UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET), cache.endOffsetFor(7))
   }
 
   @Test
   def shouldClearEarliestOnEmptyCache(): Unit = {
-    //Given
-    def leoFinder() = new LogOffsetMetadata(0)
-
-    //When
-    val cache = new LeaderEpochFileCache(tp, () => leoFinder, checkpoint)
-
     //Then
-    cache.clearAndFlushEarliest(7)
+    cache.truncateFromStart(7)
   }
 
   @Test
   def shouldClearLatestOnEmptyCache(): Unit = {
-    //Given
-    def leoFinder() = new LogOffsetMetadata(0)
-
-    //When
-    val cache = new LeaderEpochFileCache(tp, () => leoFinder, checkpoint)
-
     //Then
-    cache.clearAndFlushLatest(7)
+    cache.truncateFromEnd(7)
   }
 
-  @Before
-  def setUp() {
-    checkpoint = new LeaderEpochCheckpointFile(TestUtils.tempFile())
-  }
-}
\ No newline at end of file
+}
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 eaafb2a371f4e..8a49aac987b6a 100644
--- a/core/src/test/scala/unit/kafka/server/epoch/LeaderEpochIntegrationTest.scala
+++ b/core/src/test/scala/unit/kafka/server/epoch/LeaderEpochIntegrationTest.scala
@@ -16,7 +16,7 @@
   */
 package kafka.server.epoch
 
-import java.util.{Map => JMap}
+import java.util.Optional
 
 import kafka.server.KafkaConfig._
 import kafka.server.{BlockingSend, KafkaServer, ReplicaFetcherBlockingSend}
@@ -26,20 +26,22 @@ import kafka.zk.ZooKeeperTestHarness
 import org.apache.kafka.clients.producer.{KafkaProducer, ProducerRecord}
 import org.apache.kafka.common.metrics.Metrics
 import org.apache.kafka.common.protocol.Errors._
-import org.apache.kafka.common.requests.EpochEndOffset._
 import org.apache.kafka.common.serialization.StringSerializer
 import org.apache.kafka.common.utils.{LogContext, SystemTime}
 import org.apache.kafka.common.TopicPartition
-
+import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData.EpochEndOffset
+import org.apache.kafka.common.protocol.ApiKeys
+import org.apache.kafka.common.requests.{OffsetsForLeaderEpochRequest, OffsetsForLeaderEpochResponse}
+import org.apache.kafka.common.requests.OffsetsForLeaderEpochResponse.UNDEFINED_EPOCH_OFFSET
 import org.junit.Assert._
 import org.junit.{After, Test}
-import org.apache.kafka.common.requests.{EpochEndOffset, OffsetsForLeaderEpochRequest, OffsetsForLeaderEpochResponse}
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 import scala.collection.Map
+import scala.collection.mutable.ListBuffer
 
 class LeaderEpochIntegrationTest extends ZooKeeperTestHarness with Logging {
-  var brokers: Seq[KafkaServer] = null
+  var brokers: ListBuffer[KafkaServer] = ListBuffer()
   val topic1 = "foo"
   val topic2 = "bar"
   val t1p0 = new TopicPartition(topic1, 0)
@@ -51,7 +53,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)
@@ -59,12 +61,12 @@ class LeaderEpochIntegrationTest extends ZooKeeperTestHarness with Logging {
   }
 
   @Test
-  def shouldAddCurrentLeaderEpochToMessagesAsTheyAreWrittenToLeader() {
-    brokers = (0 to 1).map { id => createServer(fromProps(createBrokerConfig(id, zkConnect))) }
+  def shouldAddCurrentLeaderEpochToMessagesAsTheyAreWrittenToLeader(): Unit = {
+    brokers ++= (0 to 1).map { id => createServer(fromProps(createBrokerConfig(id, zkConnect))) }
 
     // Given two topics with replication of a single partition
     for (topic <- List(topic1, topic2)) {
-      createTopic(zkUtils, topic, Map(0 -> Seq(0, 1)), servers = brokers)
+      createTopic(zkClient, topic, Map(0 -> Seq(0, 1)), servers = brokers)
     }
 
     // When we send four messages
@@ -94,17 +96,16 @@ class LeaderEpochIntegrationTest extends ZooKeeperTestHarness with Logging {
   def shouldSendLeaderEpochRequestAndGetAResponse(): Unit = {
 
     //3 brokers, put partition on 100/101 and then pretend to be 102
-    brokers = (100 to 102).map { id => createServer(fromProps(createBrokerConfig(id, zkConnect))) }
-    adminZkClient.createOrUpdateTopicPartitionAssignmentPathInZK(topic1, Map(
-      0 -> Seq(100),
-      1 -> Seq(101)
-    ))
-    adminZkClient.createOrUpdateTopicPartitionAssignmentPathInZK(topic2, Map(
-      0 -> Seq(100)
-    ))
+    brokers ++= (100 to 102).map { id => createServer(fromProps(createBrokerConfig(id, zkConnect))) }
+
+    val assignment1 = Map(0 -> Seq(100), 1 -> Seq(101))
+    TestUtils.createTopic(zkClient, topic1, assignment1, brokers)
+
+    val assignment2 = Map(0 -> Seq(100))
+    TestUtils.createTopic(zkClient, topic2, assignment2, brokers)
 
     //Send messages equally to the two partitions, then half as many to a third
-    producer = createNewProducer(getBrokerListStrFromServers(brokers), retries = 5, acks = -1)
+    producer = createProducer(getBrokerListStrFromServers(brokers), acks = -1)
     (0 until 10).foreach { _ =>
       producer.send(new ProducerRecord(topic1, 0, null, "IHeartLogs".getBytes))
     }
@@ -127,8 +128,7 @@ class LeaderEpochIntegrationTest extends ZooKeeperTestHarness with Logging {
     assertEquals(30, offsetsForEpochs(t2p0).endOffset)
 
     //And should get no leader for partition error from t1p1 (as it's not on broker 0)
-    assertTrue(offsetsForEpochs(t1p1).hasError)
-    assertEquals(UNKNOWN_TOPIC_OR_PARTITION, offsetsForEpochs(t1p1).error)
+    assertEquals(NOT_LEADER_OR_FOLLOWER.code, offsetsForEpochs(t1p1).errorCode)
     assertEquals(UNDEFINED_EPOCH_OFFSET, offsetsForEpochs(t1p1).endOffset)
 
     //Repointing to broker 1 we should get the correct offset for t1p1
@@ -139,22 +139,26 @@ class LeaderEpochIntegrationTest extends ZooKeeperTestHarness with Logging {
 
   @Test
   def shouldIncreaseLeaderEpochBetweenLeaderRestarts(): Unit = {
-
     //Setup: we are only interested in the single partition on broker 101
-    brokers = Seq(100, 101).map { id => createServer(fromProps(createBrokerConfig(id, zkConnect))) }
-    def leo() = brokers(1).replicaManager.getReplica(tp).get.logEndOffset.messageOffset
-    adminZkClient.createOrUpdateTopicPartitionAssignmentPathInZK(tp.topic, Map(tp.partition -> Seq(101)))
-    producer = createNewProducer(getBrokerListStrFromServers(brokers), retries = 10, acks = -1)
+    brokers += createServer(fromProps(createBrokerConfig(100, zkConnect)))
+    assertEquals(100, TestUtils.waitUntilControllerElected(zkClient))
+
+    brokers += createServer(fromProps(createBrokerConfig(101, zkConnect)))
+
+    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)
 
     //1. Given a single message
     producer.send(new ProducerRecord(tp.topic, tp.partition, null, "IHeartLogs".getBytes)).get
     var fetcher = new TestFetcherThread(sender(brokers(0), brokers(1)))
 
     //Then epoch should be 0 and leo: 1
-    var offset = fetcher.leaderOffsetsFor(Map(tp -> 0))(tp).endOffset()
-    assertEquals(1, offset)
-    assertEquals(leo(), offset)
-
+    var epochEndOffset = fetcher.leaderOffsetsFor(Map(tp -> 0))(tp)
+    assertEquals(0, epochEndOffset.leaderEpoch)
+    assertEquals(1, epochEndOffset.endOffset)
+    assertEquals(1, leo())
 
     //2. When broker is bounced
     brokers(1).shutdown()
@@ -163,15 +167,23 @@ class LeaderEpochIntegrationTest extends ZooKeeperTestHarness with Logging {
     producer.send(new ProducerRecord(tp.topic, tp.partition, null, "IHeartLogs".getBytes)).get
     fetcher = new TestFetcherThread(sender(brokers(0), brokers(1)))
 
-
     //Then epoch 0 should still be the start offset of epoch 1
-    offset = fetcher.leaderOffsetsFor(Map(tp -> 0))(tp).endOffset()
-    assertEquals(1, offset)
-
-    //Then epoch 2 should be the leo (NB: The leader epoch goes up in factors of 2 - This is because we have to first change leader to -1 and then change it again to the live replica)
-    assertEquals(2, fetcher.leaderOffsetsFor(Map(tp -> 2))(tp).endOffset())
-    assertEquals(leo(), fetcher.leaderOffsetsFor(Map(tp -> 2))(tp).endOffset())
-
+    epochEndOffset = fetcher.leaderOffsetsFor(Map(tp -> 0))(tp)
+    assertEquals(1, epochEndOffset.endOffset)
+    assertEquals(0, epochEndOffset.leaderEpoch)
+
+    //No data written in epoch 1
+    epochEndOffset = fetcher.leaderOffsetsFor(Map(tp -> 1))(tp)
+    assertEquals(0, epochEndOffset.leaderEpoch)
+    assertEquals(1, epochEndOffset.endOffset)
+
+    //Then epoch 2 should be the leo (NB: The leader epoch goes up in factors of 2 -
+    //This is because we have to first change leader to -1 and then change it again to the live replica)
+    //Note that the expected leader changes depend on the controller being on broker 100, which is not restarted
+    epochEndOffset = fetcher.leaderOffsetsFor(Map(tp -> 2))(tp)
+    assertEquals(2, epochEndOffset.leaderEpoch)
+    assertEquals(2, epochEndOffset.endOffset)
+    assertEquals(2, leo())
 
     //3. When broker is bounced again
     brokers(1).shutdown()
@@ -180,7 +192,6 @@ class LeaderEpochIntegrationTest extends ZooKeeperTestHarness with Logging {
     producer.send(new ProducerRecord(tp.topic, tp.partition, null, "IHeartLogs".getBytes)).get
     fetcher = new TestFetcherThread(sender(brokers(0), brokers(1)))
 
-
     //Then Epoch 0 should still map to offset 1
     assertEquals(1, fetcher.leaderOffsetsFor(Map(tp -> 0))(tp).endOffset())
 
@@ -214,16 +225,13 @@ class LeaderEpochIntegrationTest extends ZooKeeperTestHarness with Logging {
   }
 
   private def sender(from: KafkaServer, to: KafkaServer): BlockingSend = {
-    val endPoint = from.metadataCache.getAliveBrokers.find(_.id == to.config.brokerId).get.getBrokerEndPoint(from.config.interBrokerListenerName)
+    val endPoint = from.metadataCache.getAliveBrokers.find(_.id == to.config.brokerId).get.brokerEndPoint(from.config.interBrokerListenerName)
     new ReplicaFetcherBlockingSend(endPoint, from.config, new Metrics(), new SystemTime(), 42, "TestFetcher", new LogContext())
   }
 
   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")
   }
 
@@ -231,13 +239,13 @@ class LeaderEpochIntegrationTest extends ZooKeeperTestHarness with Logging {
     var result = true
     for (topic <- List(topic1, topic2)) {
       val tp = new TopicPartition(topic, 0)
-      val leo = broker.getLogManager().getLog(tp).get.logEndOffset
+      val leo = broker.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) {
+        broker.getLogManager.getLog(tp).get.logSegments.iterator.forall { segment =>
+          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()
             )
@@ -251,7 +259,8 @@ class LeaderEpochIntegrationTest extends ZooKeeperTestHarness with Logging {
   private def sendFourMessagesToEachTopic() = {
     val testMessageList1 = List("test1", "test2", "test3", "test4")
     val testMessageList2 = List("test5", "test6", "test7", "test8")
-    val producer = TestUtils.createNewProducer(TestUtils.getBrokerListStrFromServers(brokers), retries = 5, keySerializer = new StringSerializer, valueSerializer = new StringSerializer)
+    val producer = TestUtils.createProducer(TestUtils.getBrokerListStrFromServers(brokers),
+      keySerializer = new StringSerializer, valueSerializer = new StringSerializer)
     val records =
       testMessageList1.map(m => new ProducerRecord(topic1, m, m)) ++
         testMessageList2.map(m => new ProducerRecord(topic2, m, m))
@@ -265,13 +274,18 @@ class LeaderEpochIntegrationTest extends ZooKeeperTestHarness with Logging {
   private[epoch] class TestFetcherThread(sender: BlockingSend) extends Logging {
 
     def leaderOffsetsFor(partitions: Map[TopicPartition, Int]): Map[TopicPartition, EpochEndOffset] = {
-      val request = new OffsetsForLeaderEpochRequest.Builder(toJavaFormat(partitions))
-      val response = sender.sendRequest(request)
-      response.responseBody.asInstanceOf[OffsetsForLeaderEpochResponse].responses.asScala
-    }
+      val partitionData = partitions.map { case (k, v) =>
+        k -> new OffsetsForLeaderEpochRequest.PartitionData(Optional.empty(), v)
+      }
 
-    def toJavaFormat(partitions: Map[TopicPartition, Int]): JMap[TopicPartition, Integer] = {
-      partitions.map { case (tp, epoch) => tp -> epoch.asInstanceOf[Integer] }.toMap.asJava
+      val request = OffsetsForLeaderEpochRequest.Builder.forFollower(
+        ApiKeys.OFFSET_FOR_LEADER_EPOCH.latestVersion, partitionData.asJava, 1)
+      val response = sender.sendRequest(request)
+      response.responseBody.asInstanceOf[OffsetsForLeaderEpochResponse].data.topics.asScala.flatMap { topic =>
+        topic.partitions.asScala.map { partition =>
+          new TopicPartition(topic.topic, partition.partition) -> partition
+        }
+      }.toMap
     }
   }
 }
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 1c01d622438eb..af597f19067d8 100644
--- a/core/src/test/scala/unit/kafka/server/epoch/OffsetsForLeaderEpochTest.scala
+++ b/core/src/test/scala/unit/kafka/server/epoch/OffsetsForLeaderEpochTest.scala
@@ -19,99 +19,149 @@ package kafka.server.epoch
 import java.io.File
 import java.util.concurrent.atomic.AtomicBoolean
 
-import kafka.cluster.Replica
+import kafka.log.{Log, LogManager}
+import kafka.server.QuotaFactory.QuotaManagers
 import kafka.server._
 import kafka.utils.{MockTime, TestUtils}
 import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.message.OffsetForLeaderEpochRequestData.{OffsetForLeaderTopic, OffsetForLeaderPartition}
+import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData.{OffsetForLeaderTopicResult, EpochEndOffset}
 import org.apache.kafka.common.metrics.Metrics
 import org.apache.kafka.common.protocol.Errors
-import org.apache.kafka.common.requests.EpochEndOffset
-import org.apache.kafka.common.requests.EpochEndOffset._
+import org.apache.kafka.common.record.RecordBatch
+import org.apache.kafka.common.requests.OffsetsForLeaderEpochResponse.{UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET}
 import org.easymock.EasyMock._
 import org.junit.Assert._
-import org.junit.Test
+import org.junit.{After, Before, Test}
 
+import scala.jdk.CollectionConverters._
 
 class OffsetsForLeaderEpochTest {
   private val config = TestUtils.createBrokerConfigs(1, TestUtils.MockZkConnect).map(KafkaConfig.fromProps).head
   private val time = new MockTime
   private val metrics = new Metrics
+  private val alterIsrManager = TestUtils.createAlterIsrManager()
   private val tp = new TopicPartition("topic", 1)
+  private var replicaManager: ReplicaManager = _
+  private var quotaManager: QuotaManagers = _
+
+  @Before
+  def setUp(): Unit = {
+    quotaManager = QuotaFactory.instantiate(config, metrics, time, "")
+  }
 
   @Test
   def shouldGetEpochsFromReplica(): Unit = {
     //Given
-    val offset = 42
+    val offsetAndEpoch = OffsetAndEpoch(42L, 5)
     val epochRequested: Integer = 5
-    val request = Map(tp -> epochRequested)
+    val request = Seq(newOffsetForLeaderTopic(tp, RecordBatch.NO_PARTITION_LEADER_EPOCH, epochRequested))
 
     //Stubs
-    val mockLog = createNiceMock(classOf[kafka.log.Log])
-    val mockCache = createNiceMock(classOf[kafka.server.epoch.LeaderEpochCache])
-    val logManager = createNiceMock(classOf[kafka.log.LogManager])
-    expect(mockCache.endOffsetFor(epochRequested)).andReturn(offset)
-    expect(mockLog.leaderEpochCache).andReturn(mockCache).anyTimes()
+    val mockLog: Log = createNiceMock(classOf[Log])
+    val logManager: LogManager = createNiceMock(classOf[LogManager])
+    expect(mockLog.endOffsetForEpoch(epochRequested)).andReturn(Some(offsetAndEpoch))
     expect(logManager.liveLogDirs).andReturn(Array.empty[File]).anyTimes()
-    replay(mockCache, mockLog, logManager)
+    replay(mockLog, logManager)
 
     // create a replica manager with 1 partition that has 1 replica
-    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)
+    replicaManager = new ReplicaManager(config, metrics, time, null, null, logManager, new AtomicBoolean(false),
+      quotaManager, new BrokerTopicStats,
+      new MetadataCache(config.brokerId), new LogDirFailureChannel(config.logDirs.size), alterIsrManager)
+    val partition = replicaManager.createPartition(tp)
+    partition.setLog(mockLog, isFutureLog = false)
     partition.leaderReplicaIdOpt = Some(config.brokerId)
 
     //When
     val response = replicaManager.lastOffsetForLeaderEpoch(request)
 
     //Then
-    assertEquals(new EpochEndOffset(Errors.NONE, offset), response(tp))
+    assertEquals(
+      Seq(newOffsetForLeaderTopicResult(tp, Errors.NONE, offsetAndEpoch.leaderEpoch, offsetAndEpoch.offset)),
+      response)
   }
 
   @Test
   def shouldReturnNoLeaderForPartitionIfThrown(): Unit = {
-    val logManager = createNiceMock(classOf[kafka.log.LogManager])
+    val logManager: LogManager = createNiceMock(classOf[LogManager])
     expect(logManager.liveLogDirs).andReturn(Array.empty[File]).anyTimes()
     replay(logManager)
 
     //create a replica manager with 1 partition that has 0 replica
-    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 = new ReplicaManager(config, metrics, time, null, null, logManager, new AtomicBoolean(false),
+      quotaManager, new BrokerTopicStats,
+      new MetadataCache(config.brokerId), new LogDirFailureChannel(config.logDirs.size), alterIsrManager)
+    replicaManager.createPartition(tp)
 
     //Given
     val epochRequested: Integer = 5
-    val request = Map(tp -> epochRequested)
+    val request = Seq(newOffsetForLeaderTopic(tp, RecordBatch.NO_PARTITION_LEADER_EPOCH, epochRequested))
 
     //When
     val response = replicaManager.lastOffsetForLeaderEpoch(request)
 
     //Then
-    assertEquals(new EpochEndOffset(Errors.NOT_LEADER_FOR_PARTITION, UNDEFINED_EPOCH_OFFSET), response(tp))
+    assertEquals(
+      Seq(newOffsetForLeaderTopicResult(tp, Errors.NOT_LEADER_OR_FOLLOWER, UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET)),
+      response)
   }
 
   @Test
   def shouldReturnUnknownTopicOrPartitionIfThrown(): Unit = {
-    val logManager = createNiceMock(classOf[kafka.log.LogManager])
+    val logManager: LogManager = createNiceMock(classOf[LogManager])
     expect(logManager.liveLogDirs).andReturn(Array.empty[File]).anyTimes()
     replay(logManager)
 
     //create a replica manager with 0 partition
-    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 = new ReplicaManager(config, metrics, time, null, null, logManager, new AtomicBoolean(false),
+      quotaManager, new BrokerTopicStats,
+      new MetadataCache(config.brokerId), new LogDirFailureChannel(config.logDirs.size), alterIsrManager)
 
     //Given
     val epochRequested: Integer = 5
-    val request = Map(tp -> epochRequested)
+    val request = Seq(newOffsetForLeaderTopic(tp, RecordBatch.NO_PARTITION_LEADER_EPOCH, epochRequested))
 
     //When
     val response = replicaManager.lastOffsetForLeaderEpoch(request)
 
     //Then
-    assertEquals(new EpochEndOffset(Errors.UNKNOWN_TOPIC_OR_PARTITION, UNDEFINED_EPOCH_OFFSET), response(tp))
+    assertEquals(
+      Seq(newOffsetForLeaderTopicResult(tp, Errors.UNKNOWN_TOPIC_OR_PARTITION, UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET)),
+      response)
+  }
+
+  @After
+  def tearDown(): Unit = {
+    Option(replicaManager).foreach(_.shutdown(checkpointHW = false))
+    Option(quotaManager).foreach(_.shutdown())
+    metrics.close()
+  }
+
+  private def newOffsetForLeaderTopic(
+    tp: TopicPartition,
+    currentLeaderEpoch: Int,
+    leaderEpoch: Int
+  ): OffsetForLeaderTopic = {
+    new OffsetForLeaderTopic()
+      .setTopic(tp.topic)
+      .setPartitions(List(new OffsetForLeaderPartition()
+        .setPartition(tp.partition)
+        .setCurrentLeaderEpoch(currentLeaderEpoch)
+        .setLeaderEpoch(leaderEpoch)).asJava)
+  }
+
+  private def newOffsetForLeaderTopicResult(
+     tp: TopicPartition,
+     error: Errors,
+     leaderEpoch: Int,
+     endOffset: Long
+  ): OffsetForLeaderTopicResult = {
+    new OffsetForLeaderTopicResult()
+      .setTopic(tp.topic)
+      .setPartitions(List(new EpochEndOffset()
+        .setPartition(tp.partition)
+        .setErrorCode(error.code)
+        .setLeaderEpoch(leaderEpoch)
+        .setEndOffset(endOffset)).asJava)
   }
-}
\ No newline at end of file
+}
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 0692afb6591c1..c03bf69881d5e 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,30 +16,66 @@
   */
 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.{ClientRequest, ClientResponse, MockClient}
-import org.apache.kafka.common.{Node, TopicPartition}
-import org.apache.kafka.common.protocol.ApiKeys
+import org.apache.kafka.clients.MockClient.MockMetadataUpdater
+import org.apache.kafka.clients.{ClientRequest, ClientResponse, MockClient, NetworkClientUtils}
+import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData
+import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData.{OffsetForLeaderTopicResult, EpochEndOffset}
+import org.apache.kafka.common.protocol.{ApiKeys, Errors}
+import org.apache.kafka.common.record.Records
 import org.apache.kafka.common.requests.AbstractRequest.Builder
-import org.apache.kafka.common.requests.FetchResponse.PartitionData
-import org.apache.kafka.common.requests.{AbstractRequest, EpochEndOffset, FetchResponse, OffsetsForLeaderEpochResponse}
+import org.apache.kafka.common.requests.{AbstractRequest, FetchResponse, OffsetsForLeaderEpochResponse, FetchMetadata => JFetchMetadata}
 import org.apache.kafka.common.utils.{SystemTime, Time}
+import org.apache.kafka.common.{Node, TopicPartition}
+
+import scala.collection.Map
 
 /**
   * Stub network client used for testing the ReplicaFetcher, wraps the MockClient used for consumer testing
+  *
+  * The common case is that there is only one OFFSET_FOR_LEADER_EPOCH request/response. So, the
+  * response to OFFSET_FOR_LEADER_EPOCH is 'offsets' map. If the test needs to set another round of
+  * 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 {
-  private val client = new MockClient(new SystemTime)
+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
+    override def update(time: Time, update: MockClient.MetadataUpdate): Unit = {}
+  })
   var fetchCount = 0
   var epochFetchCount = 0
+  var lastUsedOffsetForLeaderEpochVersion = -1
   var callback: Option[() => Unit] = None
+  var currentOffsets: util.Map[TopicPartition, EpochEndOffset] = offsets
+  var fetchPartitionData: Map[TopicPartition, FetchResponse.PartitionData[Records]] = Map.empty
+  private val sourceNode = new Node(sourceBroker.id, sourceBroker.host, sourceBroker.port)
 
-  def setEpochRequestCallback(postEpochFunction: () => Unit){
+  def setEpochRequestCallback(postEpochFunction: () => Unit): Unit = {
     callback = Some(postEpochFunction)
   }
 
+  def setOffsetsForNextResponse(newOffsets: util.Map[TopicPartition, EpochEndOffset]): Unit = {
+    currentOffsets = newOffsets
+  }
+
+  def setFetchPartitionDataForNextResponse(partitionData: Map[TopicPartition, FetchResponse.PartitionData[Records]]): Unit = {
+    fetchPartitionData = partitionData
+  }
+
   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)
@@ -50,28 +86,47 @@ class ReplicaFetcherMockBlockingSend(offsets: java.util.Map[TopicPartition, Epoc
       case ApiKeys.OFFSET_FOR_LEADER_EPOCH =>
         callback.foreach(_.apply())
         epochFetchCount += 1
-        new OffsetsForLeaderEpochResponse(offsets)
+        lastUsedOffsetForLeaderEpochVersion = requestBuilder.latestAllowedVersion()
+
+        val data = new OffsetForLeaderEpochResponseData()
+        currentOffsets.forEach((tp, offsetForLeaderPartition) => {
+          var topic = data.topics.find(tp.topic)
+          if (topic == null) {
+            topic = new OffsetForLeaderTopicResult()
+              .setTopic(tp.topic)
+            data.topics.add(topic)
+          }
+          topic.partitions.add(offsetForLeaderPartition)
+        })
+
+        new OffsetsForLeaderEpochResponse(data)
 
       case ApiKeys.FETCH =>
         fetchCount += 1
-        new FetchResponse(new java.util.LinkedHashMap[TopicPartition, PartitionData], 0)
+        val partitionData = new util.LinkedHashMap[TopicPartition, FetchResponse.PartitionData[Records]]
+        fetchPartitionData.foreach { case (tp, data) => partitionData.put(tp, data) }
+        fetchPartitionData = Map.empty
+        new FetchResponse(Errors.NONE, partitionData, 0,
+          if (partitionData.isEmpty) JFetchMetadata.INVALID_SESSION_ID else 1)
 
       case _ =>
         throw new UnsupportedOperationException
     }
 
     //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)
   }
 
+  override def initiateClose(): Unit = {}
+
   override def close(): Unit = {}
-}
\ No newline at end of file
+}
diff --git a/core/src/test/scala/unit/kafka/tools/ConsoleConsumerTest.scala b/core/src/test/scala/unit/kafka/tools/ConsoleConsumerTest.scala
index 364f6b3b93e3b..8387201f07398 100644
--- a/core/src/test/scala/unit/kafka/tools/ConsoleConsumerTest.scala
+++ b/core/src/test/scala/unit/kafka/tools/ConsoleConsumerTest.scala
@@ -17,118 +17,132 @@
 
 package kafka.tools
 
-import java.io.{PrintStream, FileOutputStream}
+import java.io.{ByteArrayOutputStream, PrintStream}
+import java.nio.file.Files
+import java.util.{HashMap, Map => JMap}
 
-import kafka.common.MessageFormatter
-import kafka.consumer.{BaseConsumer, BaseConsumerRecord}
+import kafka.tools.ConsoleConsumer.ConsumerWrapper
 import kafka.utils.{Exit, TestUtils}
+import org.apache.kafka.clients.consumer.{ConsumerRecord, MockConsumer, OffsetResetStrategy}
+import org.apache.kafka.common.{MessageFormatter, TopicPartition}
+import org.apache.kafka.common.record.TimestampType
 import org.apache.kafka.clients.consumer.ConsumerConfig
-import org.easymock.EasyMock
+import org.apache.kafka.test.MockDeserializer
+import org.mockito.Mockito._
+import org.mockito.ArgumentMatchers
+import ArgumentMatchers._
 import org.junit.Assert._
-import org.junit.Test
+import org.junit.{Before, Test}
+
+import scala.jdk.CollectionConverters._
 
 class ConsoleConsumerTest {
 
+  @Before
+  def setup(): Unit = {
+    ConsoleConsumer.messageCount = 0
+  }
+
   @Test
-  def shouldLimitReadsToMaxMessageLimit() {
-    //Mocks
-    val consumer = EasyMock.createNiceMock(classOf[BaseConsumer])
-    val formatter = EasyMock.createNiceMock(classOf[MessageFormatter])
+  def shouldResetUnConsumedOffsetsBeforeExit(): Unit = {
+    val topic = "test"
+    val maxMessages: Int = 123
+    val totalMessages: Int = 700
+    val startOffset: java.lang.Long = 0L
 
-    //Stubs
-    val record = new BaseConsumerRecord(topic = "foo", partition = 1, offset = 1, key = Array[Byte](), value = Array[Byte]())
+    val mockConsumer = new MockConsumer[Array[Byte], Array[Byte]](OffsetResetStrategy.EARLIEST)
+    val tp1 = new TopicPartition(topic, 0)
+    val tp2 = new TopicPartition(topic, 1)
 
-    //Expectations
-    val messageLimit: Int = 10
-    EasyMock.expect(formatter.writeTo(EasyMock.anyObject(), EasyMock.anyObject())).times(messageLimit)
-    EasyMock.expect(consumer.receive()).andReturn(record).times(messageLimit)
+    val consumer = new ConsumerWrapper(Some(topic), None, None, None, mockConsumer)
 
-    EasyMock.replay(consumer)
-    EasyMock.replay(formatter)
+    mockConsumer.rebalance(List(tp1, tp2).asJava)
+    mockConsumer.updateBeginningOffsets(Map(tp1 -> startOffset, tp2 -> startOffset).asJava)
 
-    //Test
-    ConsoleConsumer.process(messageLimit, formatter, consumer, System.out, true)
+    0 until totalMessages foreach { i =>
+      // add all records, each partition should have half of `totalMessages`
+      mockConsumer.addRecord(new ConsumerRecord[Array[Byte], Array[Byte]](topic, i % 2, i / 2, "key".getBytes, "value".getBytes))
+    }
+
+    val formatter = mock(classOf[MessageFormatter])
+
+    ConsoleConsumer.process(maxMessages, formatter, consumer, System.out, skipMessageOnError = false)
+    assertEquals(totalMessages, mockConsumer.position(tp1) + mockConsumer.position(tp2))
+
+    consumer.resetUnconsumedOffsets()
+    assertEquals(maxMessages, mockConsumer.position(tp1) + mockConsumer.position(tp2))
+
+    verify(formatter, times(maxMessages)).writeTo(any(), any())
   }
 
   @Test
-  def shouldStopWhenOutputCheckErrorFails() {
-    //Mocks
-    val consumer = EasyMock.createNiceMock(classOf[BaseConsumer])
-    val formatter = EasyMock.createNiceMock(classOf[MessageFormatter])
-    val printStream = EasyMock.createNiceMock(classOf[PrintStream])
-
-    //Stubs
-    val record = new BaseConsumerRecord(topic = "foo", partition = 1, offset = 1, key = Array[Byte](), value = Array[Byte]())
-
-    //Expectations
-    EasyMock.expect(consumer.receive()).andReturn(record)
-    EasyMock.expect(formatter.writeTo(EasyMock.anyObject(), EasyMock.eq(printStream)))
-    //Simulate an error on System.out after the first record has been printed
-    EasyMock.expect(printStream.checkError()).andReturn(true)
+  def shouldLimitReadsToMaxMessageLimit(): Unit = {
+    val consumer = mock(classOf[ConsumerWrapper])
+    val formatter = mock(classOf[MessageFormatter])
+    val record = new ConsumerRecord("foo", 1, 1, Array[Byte](), Array[Byte]())
 
-    EasyMock.replay(consumer)
-    EasyMock.replay(formatter)
-    EasyMock.replay(printStream)
+    val messageLimit: Int = 10
+    when(consumer.receive()).thenReturn(record)
 
-    //Test
-    ConsoleConsumer.process(-1, formatter, consumer, printStream, true)
+    ConsoleConsumer.process(messageLimit, formatter, consumer, System.out, true)
 
-    //Verify
-    EasyMock.verify(consumer, formatter, printStream)
+    verify(consumer, times(messageLimit)).receive()
+    verify(formatter, times(messageLimit)).writeTo(any(), any())
+
+    consumer.cleanup()
   }
 
   @Test
-  def shouldParseValidOldConsumerValidConfig() {
-    //Given
-    val args: Array[String] = Array(
-      "--zookeeper", "localhost:2181",
-      "--topic", "test",
-      "--from-beginning")
+  def shouldStopWhenOutputCheckErrorFails(): Unit = {
+    val consumer = mock(classOf[ConsumerWrapper])
+    val formatter = mock(classOf[MessageFormatter])
+    val printStream = mock(classOf[PrintStream])
 
-    //When
-    val config = new ConsoleConsumer.ConsumerConfig(args)
+    val record = new ConsumerRecord("foo", 1, 1, Array[Byte](), Array[Byte]())
 
-    //Then
-    assertTrue(config.useOldConsumer)
-    assertEquals("localhost:2181", config.zkConnectionStr)
-    assertEquals("test", config.topicArg)
-    assertEquals(true, config.fromBeginning)
+    when(consumer.receive()).thenReturn(record)
+    //Simulate an error on System.out after the first record has been printed
+    when(printStream.checkError()).thenReturn(true)
+
+    ConsoleConsumer.process(-1, formatter, consumer, printStream, true)
+
+    verify(formatter).writeTo(any(), ArgumentMatchers.eq(printStream))
+    verify(consumer).receive()
+    verify(printStream).checkError()
+
+    consumer.cleanup()
   }
 
   @Test
-  def shouldParseValidNewConsumerValidConfig() {
+  def shouldParseValidConsumerValidConfig(): Unit = {
     //Given
     val args: Array[String] = Array(
       "--bootstrap-server", "localhost:9092",
       "--topic", "test",
-      "--from-beginning",
-      "--new-consumer") //new
+      "--from-beginning")
 
     //When
     val config = new ConsoleConsumer.ConsumerConfig(args)
 
     //Then
-    assertFalse(config.useOldConsumer)
     assertEquals("localhost:9092", config.bootstrapServer)
     assertEquals("test", config.topicArg)
     assertEquals(true, config.fromBeginning)
   }
 
   @Test
-  def shouldParseValidNewSimpleConsumerValidConfigWithNumericOffset(): Unit = {
+  def shouldParseValidSimpleConsumerValidConfigWithNumericOffset(): Unit = {
     //Given
     val args: Array[String] = Array(
       "--bootstrap-server", "localhost:9092",
       "--topic", "test",
       "--partition", "0",
-      "--offset", "3",
-      "--new-consumer") //new
+      "--offset", "3")
 
     //When
     val config = new ConsoleConsumer.ConsumerConfig(args)
 
     //Then
-    assertFalse(config.useOldConsumer)
     assertEquals("localhost:9092", config.bootstrapServer)
     assertEquals("test", config.topicArg)
     assertEquals(0, config.partitionArg.get)
@@ -137,37 +151,39 @@ class ConsoleConsumerTest {
 
   }
 
-  @Test
-  def testDefaultConsumer() {
+  @Test(expected = classOf[IllegalArgumentException])
+  def shouldExitOnUnrecognizedNewConsumerOption(): Unit = {
+    Exit.setExitProcedure((_, message) => throw new IllegalArgumentException(message.orNull))
+
     //Given
     val args: Array[String] = Array(
+      "--new-consumer",
       "--bootstrap-server", "localhost:9092",
       "--topic", "test",
       "--from-beginning")
 
     //When
-    val config = new ConsoleConsumer.ConsumerConfig(args)
-
-    //Then
-    assertFalse(config.useOldConsumer)
+    try {
+      new ConsoleConsumer.ConsumerConfig(args)
+    } finally {
+      Exit.resetExitProcedure()
+    }
   }
 
   @Test
-  def shouldParseValidNewSimpleConsumerValidConfigWithStringOffset() {
+  def shouldParseValidSimpleConsumerValidConfigWithStringOffset(): Unit = {
     //Given
     val args: Array[String] = Array(
       "--bootstrap-server", "localhost:9092",
       "--topic", "test",
       "--partition", "0",
       "--offset", "LatEst",
-      "--new-consumer", //new
       "--property", "print.value=false")
 
     //When
     val config = new ConsoleConsumer.ConsumerConfig(args)
 
     //Then
-    assertFalse(config.useOldConsumer)
     assertEquals("localhost:9092", config.bootstrapServer)
     assertEquals("test", config.topicArg)
     assertEquals(0, config.partitionArg.get)
@@ -177,67 +193,7 @@ class ConsoleConsumerTest {
   }
 
   @Test
-  def shouldParseValidOldConsumerConfigWithAutoOffsetResetSmallest() {
-    //Given
-    val args: Array[String] = Array(
-      "--zookeeper", "localhost:2181",
-      "--topic", "test",
-      "--consumer-property", "auto.offset.reset=smallest")
-
-    //When
-    val config = new ConsoleConsumer.ConsumerConfig(args)
-    val consumerProperties = ConsoleConsumer.getOldConsumerProps(config)
-
-    //Then
-    assertTrue(config.useOldConsumer)
-    assertEquals("localhost:2181", config.zkConnectionStr)
-    assertEquals("test", config.topicArg)
-    assertEquals(false, config.fromBeginning)
-    assertEquals("smallest", consumerProperties.getProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG))
-  }
-
-  @Test
-  def shouldParseValidOldConsumerConfigWithAutoOffsetResetLargest() {
-    //Given
-    val args: Array[String] = Array(
-      "--zookeeper", "localhost:2181",
-      "--topic", "test",
-      "--consumer-property", "auto.offset.reset=largest")
-
-    //When
-    val config = new ConsoleConsumer.ConsumerConfig(args)
-    val consumerProperties = ConsoleConsumer.getOldConsumerProps(config)
-
-    //Then
-    assertTrue(config.useOldConsumer)
-    assertEquals("localhost:2181", config.zkConnectionStr)
-    assertEquals("test", config.topicArg)
-    assertEquals(false, config.fromBeginning)
-    assertEquals("largest", consumerProperties.getProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG))
-  }
-
-  @Test
-  def shouldSetAutoResetToSmallestWhenFromBeginningConfigured() {
-    //Given
-    val args = Array(
-      "--zookeeper", "localhost:2181",
-      "--topic", "test",
-      "--from-beginning")
-
-    //When
-    val config = new ConsoleConsumer.ConsumerConfig(args)
-    val consumerProperties = ConsoleConsumer.getOldConsumerProps(config)
-
-    //Then
-    assertTrue(config.useOldConsumer)
-    assertEquals("localhost:2181", config.zkConnectionStr)
-    assertEquals("test", config.topicArg)
-    assertEquals(true, config.fromBeginning)
-    assertEquals("smallest", consumerProperties.getProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG))
-  }
-
-  @Test
-  def shouldParseValidNewConsumerConfigWithAutoOffsetResetLatest() {
+  def shouldParseValidConsumerConfigWithAutoOffsetResetLatest(): Unit = {
     //Given
     val args: Array[String] = Array(
       "--bootstrap-server", "localhost:9092",
@@ -246,10 +202,9 @@ class ConsoleConsumerTest {
 
     //When
     val config = new ConsoleConsumer.ConsumerConfig(args)
-    val consumerProperties = ConsoleConsumer.getNewConsumerProps(config)
+    val consumerProperties = ConsoleConsumer.consumerProps(config)
 
     //Then
-    assertFalse(config.useOldConsumer)
     assertEquals("localhost:9092", config.bootstrapServer)
     assertEquals("test", config.topicArg)
     assertEquals(false, config.fromBeginning)
@@ -257,7 +212,7 @@ class ConsoleConsumerTest {
   }
 
   @Test
-  def shouldParseValidNewConsumerConfigWithAutoOffsetResetEarliest() {
+  def shouldParseValidConsumerConfigWithAutoOffsetResetEarliest(): Unit = {
     //Given
     val args: Array[String] = Array(
       "--bootstrap-server", "localhost:9092",
@@ -266,10 +221,9 @@ class ConsoleConsumerTest {
 
     //When
     val config = new ConsoleConsumer.ConsumerConfig(args)
-    val consumerProperties = ConsoleConsumer.getNewConsumerProps(config)
+    val consumerProperties = ConsoleConsumer.consumerProps(config)
 
     //Then
-    assertFalse(config.useOldConsumer)
     assertEquals("localhost:9092", config.bootstrapServer)
     assertEquals("test", config.topicArg)
     assertEquals(false, config.fromBeginning)
@@ -277,7 +231,7 @@ class ConsoleConsumerTest {
   }
 
   @Test
-  def shouldParseValidNewConsumerConfigWithAutoOffsetResetAndMatchingFromBeginning() {
+  def shouldParseValidConsumerConfigWithAutoOffsetResetAndMatchingFromBeginning(): Unit = {
     //Given
     val args: Array[String] = Array(
       "--bootstrap-server", "localhost:9092",
@@ -287,10 +241,9 @@ class ConsoleConsumerTest {
 
     //When
     val config = new ConsoleConsumer.ConsumerConfig(args)
-    val consumerProperties = ConsoleConsumer.getNewConsumerProps(config)
+    val consumerProperties = ConsoleConsumer.consumerProps(config)
 
     //Then
-    assertFalse(config.useOldConsumer)
     assertEquals("localhost:9092", config.bootstrapServer)
     assertEquals("test", config.topicArg)
     assertEquals(true, config.fromBeginning)
@@ -298,7 +251,7 @@ class ConsoleConsumerTest {
   }
 
   @Test
-  def shouldParseValidNewConsumerConfigWithNoOffsetReset() {
+  def shouldParseValidConsumerConfigWithNoOffsetReset(): Unit = {
     //Given
     val args: Array[String] = Array(
       "--bootstrap-server", "localhost:9092",
@@ -306,10 +259,9 @@ class ConsoleConsumerTest {
 
     //When
     val config = new ConsoleConsumer.ConsumerConfig(args)
-    val consumerProperties = ConsoleConsumer.getNewConsumerProps(config)
+    val consumerProperties = ConsoleConsumer.consumerProps(config)
 
     //Then
-    assertFalse(config.useOldConsumer)
     assertEquals("localhost:9092", config.bootstrapServer)
     assertEquals("test", config.topicArg)
     assertEquals(false, config.fromBeginning)
@@ -317,10 +269,7 @@ class ConsoleConsumerTest {
   }
 
   @Test(expected = classOf[IllegalArgumentException])
-  def shouldExitOnInvalidConfigWithAutoOffsetResetAndConflictingFromBeginningNewConsumer() {
-
-    // Override exit procedure to throw an exception instead of exiting, so we can catch the exit
-    // properly for this test case
+  def shouldExitOnInvalidConfigWithAutoOffsetResetAndConflictingFromBeginning(): Unit = {
     Exit.setExitProcedure((_, message) => throw new IllegalArgumentException(message.orNull))
 
     //Given
@@ -332,42 +281,16 @@ class ConsoleConsumerTest {
 
     try {
       val config = new ConsoleConsumer.ConsumerConfig(args)
-      ConsoleConsumer.getNewConsumerProps(config)
-    } finally {
-      Exit.resetExitProcedure()
-    }
-
-    fail("Expected consumer property construction to fail due to inconsistent reset options")
-  }
-
-  @Test(expected = classOf[IllegalArgumentException])
-  def shouldExitOnInvalidConfigWithAutoOffsetResetAndConflictingFromBeginningOldConsumer() {
-
-    // Override exit procedure to throw an exception instead of exiting, so we can catch the exit
-    // properly for this test case
-    Exit.setExitProcedure((_, message) => throw new IllegalArgumentException(message.orNull))
-
-    //Given
-    val args: Array[String] = Array(
-      "--zookeeper", "localhost:2181",
-      "--topic", "test",
-      "--consumer-property", "auto.offset.reset=largest",
-      "--from-beginning")
-
-    try {
-      val config = new ConsoleConsumer.ConsumerConfig(args)
-      ConsoleConsumer.getOldConsumerProps(config)
+      ConsoleConsumer.consumerProps(config)
     } finally {
       Exit.resetExitProcedure()
     }
-
-    fail("Expected consumer property construction to fail due to inconsistent reset options")
   }
 
   @Test
-  def shouldParseConfigsFromFile() {
+  def shouldParseConfigsFromFile(): Unit = {
     val propsFile = TestUtils.tempFile()
-    val propsStream = new FileOutputStream(propsFile)
+    val propsStream = Files.newOutputStream(propsFile.toPath)
     propsStream.write("request.timeout.ms=1000\n".getBytes())
     propsStream.write("group.id=group1".getBytes())
     propsStream.close()
@@ -384,12 +307,12 @@ class ConsoleConsumerTest {
   }
 
   @Test
-  def groupIdsProvidedInDifferentPlacesMustMatch() {
+  def groupIdsProvidedInDifferentPlacesMustMatch(): Unit = {
     Exit.setExitProcedure((_, message) => throw new IllegalArgumentException(message.orNull))
 
     // different in all three places
     var propsFile = TestUtils.tempFile()
-    var propsStream = new FileOutputStream(propsFile)
+    var propsStream = Files.newOutputStream(propsFile.toPath)
     propsStream.write("group.id=group-from-file".getBytes())
     propsStream.close()
     var args: Array[String] = Array(
@@ -409,7 +332,7 @@ class ConsoleConsumerTest {
 
     // the same in all three places
     propsFile = TestUtils.tempFile()
-    propsStream = new FileOutputStream(propsFile)
+    propsStream = Files.newOutputStream(propsFile.toPath)
     propsStream.write("group.id=test-group".getBytes())
     propsStream.close()
     args = Array(
@@ -421,12 +344,12 @@ class ConsoleConsumerTest {
     )
 
     var config = new ConsoleConsumer.ConsumerConfig(args)
-    var props = ConsoleConsumer.getNewConsumerProps(config)
+    var props = ConsoleConsumer.consumerProps(config)
     assertEquals("test-group", props.getProperty("group.id"))
 
     // different via --consumer-property and --consumer.config
     propsFile = TestUtils.tempFile()
-    propsStream = new FileOutputStream(propsFile)
+    propsStream = Files.newOutputStream(propsFile.toPath)
     propsStream.write("group.id=group-from-file".getBytes())
     propsStream.close()
     args = Array(
@@ -460,7 +383,7 @@ class ConsoleConsumerTest {
 
     // different via --group and --consumer.config
     propsFile = TestUtils.tempFile()
-    propsStream = new FileOutputStream(propsFile)
+    propsStream = Files.newOutputStream(propsFile.toPath)
     propsStream.write("group.id=group-from-file".getBytes())
     propsStream.close()
     args = Array(
@@ -485,9 +408,164 @@ class ConsoleConsumerTest {
     )
 
     config = new ConsoleConsumer.ConsumerConfig(args)
-    props = ConsoleConsumer.getNewConsumerProps(config)
+    props = ConsoleConsumer.consumerProps(config)
     assertEquals("group-from-arguments", props.getProperty("group.id"))
 
     Exit.resetExitProcedure()
   }
+
+  @Test
+  def testCustomPropertyShouldBePassedToConfigureMethod(): Unit = {
+    val args = Array(
+      "--bootstrap-server", "localhost:9092",
+      "--topic", "test",
+      "--property", "print.key=true",
+      "--property", "key.deserializer=org.apache.kafka.test.MockDeserializer",
+      "--property", "key.deserializer.my-props=abc"
+    )
+    val config = new ConsoleConsumer.ConsumerConfig(args)
+    assertTrue(config.formatter.isInstanceOf[DefaultMessageFormatter])
+    assertTrue(config.formatterArgs.containsKey("key.deserializer.my-props"))
+    val formatter = config.formatter.asInstanceOf[DefaultMessageFormatter]
+    assertTrue(formatter.keyDeserializer.get.isInstanceOf[MockDeserializer])
+    assertEquals(1, formatter.keyDeserializer.get.asInstanceOf[MockDeserializer].configs.size)
+    assertEquals("abc", formatter.keyDeserializer.get.asInstanceOf[MockDeserializer].configs.get("my-props"))
+    assertTrue(formatter.keyDeserializer.get.asInstanceOf[MockDeserializer].isKey)
+  }
+
+  @Test
+  def shouldParseGroupIdFromBeginningGivenTogether(): Unit = {
+    // Start from earliest
+    var args: Array[String] = Array(
+      "--bootstrap-server", "localhost:9092",
+      "--topic", "test",
+      "--group", "test-group",
+      "--from-beginning")
+
+    var config = new ConsoleConsumer.ConsumerConfig(args)
+    assertEquals("localhost:9092", config.bootstrapServer)
+    assertEquals("test", config.topicArg)
+    assertEquals(-2, config.offsetArg)
+    assertEquals(true, config.fromBeginning)
+
+    // Start from latest
+    args = Array(
+      "--bootstrap-server", "localhost:9092",
+      "--topic", "test",
+      "--group", "test-group"
+    )
+
+    config = new ConsoleConsumer.ConsumerConfig(args)
+    assertEquals("localhost:9092", config.bootstrapServer)
+    assertEquals("test", config.topicArg)
+    assertEquals(-1, config.offsetArg)
+    assertEquals(false, config.fromBeginning)
+  }
+
+  @Test(expected = classOf[IllegalArgumentException])
+  def shouldExitOnGroupIdAndPartitionGivenTogether(): Unit = {
+    Exit.setExitProcedure((_, message) => throw new IllegalArgumentException(message.orNull))
+    //Given
+    val args: Array[String] = Array(
+      "--bootstrap-server", "localhost:9092",
+      "--topic", "test",
+      "--group", "test-group",
+      "--partition", "0")
+
+    //When
+    try {
+      new ConsoleConsumer.ConsumerConfig(args)
+    } finally {
+      Exit.resetExitProcedure()
+    }
+  }
+
+  @Test(expected = classOf[IllegalArgumentException])
+  def shouldExitOnOffsetWithoutPartition(): Unit = {
+    Exit.setExitProcedure((_, message) => throw new IllegalArgumentException(message.orNull))
+    //Given
+    val args: Array[String] = Array(
+      "--bootstrap-server", "localhost:9092",
+      "--topic", "test",
+      "--offset", "10")
+
+    //When
+    try {
+      new ConsoleConsumer.ConsumerConfig(args)
+    } finally {
+      Exit.resetExitProcedure()
+    }
+  }
+
+  @Test
+  def testDefaultMessageFormatter(): Unit = {
+    val record = new ConsumerRecord("topic", 0, 123, "key".getBytes, "value".getBytes)
+    val formatter = new DefaultMessageFormatter()
+    val configs: JMap[String, String] = new HashMap()
+
+    formatter.configure(configs)
+    var out = new ByteArrayOutputStream()
+    formatter.writeTo(record, new PrintStream(out))
+    assertEquals("value\n", out.toString)
+
+    configs.put("print.key", "true")
+    formatter.configure(configs)
+    out = new ByteArrayOutputStream()
+    formatter.writeTo(record, new PrintStream(out))
+    assertEquals("key\tvalue\n", out.toString)
+
+    configs.put("print.partition", "true")
+    formatter.configure(configs)
+    out = new ByteArrayOutputStream()
+    formatter.writeTo(record, new PrintStream(out))
+    assertEquals("Partition:0\tkey\tvalue\n", out.toString)
+
+    configs.put("print.timestamp", "true")
+    formatter.configure(configs)
+    out = new ByteArrayOutputStream()
+    formatter.writeTo(record, new PrintStream(out))
+    assertEquals("NO_TIMESTAMP\tPartition:0\tkey\tvalue\n", out.toString)
+
+    configs.put("print.offset", "true")
+    formatter.configure(configs)
+    out = new ByteArrayOutputStream()
+    formatter.writeTo(record, new PrintStream(out))
+    assertEquals("NO_TIMESTAMP\tPartition:0\tOffset:123\tkey\tvalue\n", out.toString)
+
+    out = new ByteArrayOutputStream()
+    val record2 = new ConsumerRecord("topic", 0, 123, 123L, TimestampType.CREATE_TIME, 321L, -1, -1, "key".getBytes, "value".getBytes)
+    formatter.writeTo(record2, new PrintStream(out))
+    assertEquals("CreateTime:123\tPartition:0\tOffset:123\tkey\tvalue\n", out.toString)
+    formatter.close()
+  }
+
+  @Test
+  def testNoOpMessageFormatter(): Unit = {
+    val record = new ConsumerRecord("topic", 0, 123, "key".getBytes, "value".getBytes)
+    val formatter = new NoOpMessageFormatter()
+
+    formatter.configure(new HashMap())
+    val out = new ByteArrayOutputStream()
+    formatter.writeTo(record, new PrintStream(out))
+    assertEquals("", out.toString)
+  }
+
+  @Test
+  def testChecksumMessageFormatter(): Unit = {
+    val record = new ConsumerRecord("topic", 0, 123, "key".getBytes, "value".getBytes)
+    val formatter = new ChecksumMessageFormatter()
+    val configs: JMap[String, String] = new HashMap()
+
+    formatter.configure(configs)
+    var out = new ByteArrayOutputStream()
+    formatter.writeTo(record, new PrintStream(out))
+    assertEquals("checksum:-1\n", out.toString)
+
+    configs.put("topic", "topic1")
+    formatter.configure(configs)
+    out = new ByteArrayOutputStream()
+    formatter.writeTo(record, new PrintStream(out))
+    assertEquals("topic1:checksum:-1\n", out.toString)
+  }
+
 }
diff --git a/core/src/test/scala/unit/kafka/tools/ConsoleProducerTest.scala b/core/src/test/scala/unit/kafka/tools/ConsoleProducerTest.scala
index da80c0d41e25c..a636c32370c11 100644
--- a/core/src/test/scala/unit/kafka/tools/ConsoleProducerTest.scala
+++ b/core/src/test/scala/unit/kafka/tools/ConsoleProducerTest.scala
@@ -17,14 +17,17 @@
 
 package kafka.tools
 
-import kafka.producer.ProducerConfig
+import java.util
+
 import ConsoleProducer.LineMessageReader
-import org.apache.kafka.clients.producer.KafkaProducer
+import org.apache.kafka.clients.producer.ProducerConfig
 import org.junit.{Assert, Test}
+import Assert.assertEquals
+import kafka.utils.Exit
 
 class ConsoleProducerTest {
 
-  val validArgs: Array[String] = Array(
+  val brokerListValidArgs: Array[String] = Array(
     "--broker-list",
     "localhost:1001,localhost:1002",
     "--topic",
@@ -34,45 +37,85 @@ class ConsoleProducerTest {
     "--property",
     "key.separator=#"
   )
-
+  val bootstrapServerValidArgs: Array[String] = Array(
+    "--bootstrap-server",
+    "localhost:1003,localhost:1004",
+    "--topic",
+    "t3",
+    "--property",
+    "parse.key=true",
+    "--property",
+    "key.separator=#"
+  )
   val invalidArgs: Array[String] = Array(
     "--t", // not a valid argument
     "t3"
   )
+  val bootstrapServerOverride: Array[String] = Array(
+    "--broker-list",
+    "localhost:1001",
+    "--bootstrap-server",
+    "localhost:1002",
+    "--topic",
+    "t3",
+  )
+  val clientIdOverride: Array[String] = Array(
+    "--broker-list",
+    "localhost:1001",
+    "--topic",
+    "t3",
+    "--producer-property",
+    "client.id=producer-1"
+  )
 
   @Test
-  def testValidConfigsNewProducer() {
-    val config = new ConsoleProducer.ProducerConfig(validArgs)
-    // New ProducerConfig constructor is package private, so we can't call it directly
-    // Creating new Producer to validate instead
-    val producer = new KafkaProducer(ConsoleProducer.getNewProducerProps(config))
-    producer.close()
+  def testValidConfigsBrokerList(): Unit = {
+    val config = new ConsoleProducer.ProducerConfig(brokerListValidArgs)
+    val producerConfig = new ProducerConfig(ConsoleProducer.producerProps(config))
+    assertEquals(util.Arrays.asList("localhost:1001", "localhost:1002"),
+      producerConfig.getList(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG))
   }
 
   @Test
-  @deprecated("This test has been deprecated and it will be removed in a future release.", "0.10.0.0")
-  def testValidConfigsOldProducer() {
-    val config = new ConsoleProducer.ProducerConfig(validArgs)
-    new ProducerConfig(ConsoleProducer.getOldProducerProps(config))
+  def testValidConfigsBootstrapServer(): Unit = {
+    val config = new ConsoleProducer.ProducerConfig(bootstrapServerValidArgs)
+    val producerConfig = new ProducerConfig(ConsoleProducer.producerProps(config))
+    assertEquals(util.Arrays.asList("localhost:1003", "localhost:1004"),
+      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()
     }
   }
 
   @Test
   def testParseKeyProp(): Unit = {
-    val config = new ConsoleProducer.ProducerConfig(validArgs)
-    val reader = Class.forName(config.readerClass).newInstance().asInstanceOf[LineMessageReader]
+    val config = new ConsoleProducer.ProducerConfig(brokerListValidArgs)
+    val reader = Class.forName(config.readerClass).getDeclaredConstructor().newInstance().asInstanceOf[LineMessageReader]
     reader.init(System.in,ConsoleProducer.getReaderProps(config))
     assert(reader.keySeparator == "#")
     assert(reader.parseKey)
   }
 
+  @Test
+  def testBootstrapServerOverride(): Unit = {
+    val config = new ConsoleProducer.ProducerConfig(bootstrapServerOverride)
+    val producerConfig = new ProducerConfig(ConsoleProducer.producerProps(config))
+    assertEquals(util.Arrays.asList("localhost:1002"),
+      producerConfig.getList(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG))
+  }
+
+  @Test
+  def testClientIdOverride(): Unit = {
+    val config = new ConsoleProducer.ProducerConfig(clientIdOverride)
+    val producerConfig = new ProducerConfig(ConsoleProducer.producerProps(config))
+    assertEquals("producer-1",
+      producerConfig.getString(ProducerConfig.CLIENT_ID_CONFIG))
+  }
 }
diff --git a/core/src/test/scala/unit/kafka/tools/ConsumerPerformanceTest.scala b/core/src/test/scala/unit/kafka/tools/ConsumerPerformanceTest.scala
index bafe8ed136bb1..94b732d2fd57d 100644
--- a/core/src/test/scala/unit/kafka/tools/ConsumerPerformanceTest.scala
+++ b/core/src/test/scala/unit/kafka/tools/ConsumerPerformanceTest.scala
@@ -20,6 +20,7 @@ package kafka.tools
 import java.io.ByteArrayOutputStream
 import java.text.SimpleDateFormat
 
+import kafka.utils.Exit
 import org.junit.Assert.assertEquals
 import org.junit.Test
 
@@ -30,24 +31,93 @@ class ConsumerPerformanceTest {
 
   @Test
   def testDetailedHeaderMatchBody(): Unit = {
-    testHeaderMatchContent(detailed = true, useOldConsumer = false, 2,
-      () => ConsumerPerformance.printNewConsumerProgress(1, 1024 * 1024, 0, 1, 0, 0, 1, dateFormat, 1L))
-    testHeaderMatchContent(detailed = true, useOldConsumer = true, 4,
-      () => ConsumerPerformance.printOldConsumerProgress(1, 1024 * 1024, 0, 1, 0, 0, 1,
-      dateFormat))
+    testHeaderMatchContent(detailed = true, 2,
+      () => ConsumerPerformance.printConsumerProgress(1, 1024 * 1024, 0, 1, 0, 0, 1, dateFormat, 1L))
   }
 
   @Test
   def testNonDetailedHeaderMatchBody(): Unit = {
-    testHeaderMatchContent(detailed = false, useOldConsumer = false, 2, () => println(s"${dateFormat.format(System.currentTimeMillis)}, " +
+    testHeaderMatchContent(detailed = false, 2, () => println(s"${dateFormat.format(System.currentTimeMillis)}, " +
       s"${dateFormat.format(System.currentTimeMillis)}, 1.0, 1.0, 1, 1.0, 1, 1, 1.1, 1.1"))
-    testHeaderMatchContent(detailed = false, useOldConsumer = true, 4, () => println(s"${dateFormat.format(System.currentTimeMillis)}, " +
-      s"${dateFormat.format(System.currentTimeMillis)}, 1.0, 1.0, 1, 1.0"))
   }
 
-  private def testHeaderMatchContent(detailed: Boolean, useOldConsumer: Boolean, expectedOutputLineCount: Int, fun: () => Unit): Unit = {
+  @Test
+  def testConfigBrokerList(): Unit = {
+    //Given
+    val args: Array[String] = Array(
+      "--broker-list", "localhost:9092",
+      "--topic", "test",
+      "--messages", "10"
+    )
+
+    //When
+    val config = new ConsumerPerformance.ConsumerPerfConfig(args)
+
+    //Then
+    assertEquals("localhost:9092", config.brokerHostsAndPorts)
+    assertEquals("test", config.topic)
+    assertEquals(10, config.numMessages)
+  }
+
+  @Test
+  def testConfigBootStrapServer(): Unit = {
+    //Given
+    val args: Array[String] = Array(
+      "--bootstrap-server", "localhost:9092",
+      "--topic", "test",
+      "--messages", "10",
+      "--print-metrics"
+    )
+
+    //When
+    val config = new ConsumerPerformance.ConsumerPerfConfig(args)
+
+    //Then
+    assertEquals("localhost:9092", config.brokerHostsAndPorts)
+    assertEquals("test", config.topic)
+    assertEquals(10, config.numMessages)
+  }
+
+  @Test
+  def testBrokerListOverride(): Unit = {
+    //Given
+    val args: Array[String] = Array(
+      "--broker-list", "localhost:9094",
+      "--bootstrap-server", "localhost:9092",
+      "--topic", "test",
+      "--messages", "10"
+    )
+
+    //When
+    val config = new ConsumerPerformance.ConsumerPerfConfig(args)
+
+    //Then
+    assertEquals("localhost:9092", config.brokerHostsAndPorts)
+    assertEquals("test", config.topic)
+    assertEquals(10, config.numMessages)
+  }
+
+  @Test(expected = classOf[IllegalArgumentException])
+  def testConfigWithUnrecognizedOption(): Unit = {
+    Exit.setExitProcedure((_, message) => throw new IllegalArgumentException(message.orNull))
+    //Given
+    val args: Array[String] = Array(
+      "--broker-list", "localhost:9092",
+      "--topic", "test",
+      "--messages", "10",
+      "--new-consumer"
+    )
+    try {
+      //When
+      new ConsumerPerformance.ConsumerPerfConfig(args)
+    } finally {
+      Exit.resetExitProcedure()
+    }
+  }
+
+  private def testHeaderMatchContent(detailed: Boolean, expectedOutputLineCount: Int, fun: () => Unit): Unit = {
     Console.withOut(outContent) {
-      ConsumerPerformance.printHeader(detailed, useOldConsumer)
+      ConsumerPerformance.printHeader(detailed)
       fun()
 
       val contents = outContent.toString.split("\n")
diff --git a/core/src/test/scala/unit/kafka/tools/DumpLogSegmentsTest.scala b/core/src/test/scala/unit/kafka/tools/DumpLogSegmentsTest.scala
new file mode 100644
index 0000000000000..0f166080a051a
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/tools/DumpLogSegmentsTest.scala
@@ -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 kafka.tools
+
+import java.io.{ByteArrayOutputStream, File}
+import java.util.Properties
+
+import kafka.log.{Log, LogConfig, LogManager}
+import kafka.server.{BrokerTopicStats, LogDirFailureChannel}
+import kafka.tools.DumpLogSegments.TimeIndexDumpErrors
+import kafka.utils.{MockTime, TestUtils}
+import org.apache.kafka.common.record.{CompressionType, MemoryRecords, SimpleRecord}
+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()
+  val logDir = TestUtils.randomPartitionLogDir(tmpDir)
+  val segmentName = "00000000000000000000"
+  val logFilePath = s"$logDir/$segmentName.log"
+  val indexFilePath = s"$logDir/$segmentName.index"
+  val timeIndexFilePath = s"$logDir/$segmentName.timeindex"
+  val time = new MockTime(0, 0)
+
+  val batches = new ArrayBuffer[BatchInfo]
+  var log: Log = _
+
+  @Before
+  def setUp(): Unit = {
+    val props = new Properties
+    props.setProperty(LogConfig.IndexIntervalBytesProp, "128")
+    log = Log(logDir, LogConfig(props), logStartOffset = 0L, recoveryPoint = 0L, scheduler = time.scheduler,
+      time = time, brokerTopicStats = new BrokerTopicStats, maxProducerIdExpirationMs = 60 * 60 * 1000,
+      producerIdExpirationCheckIntervalMs = LogManager.ProducerIdExpirationCheckIntervalMs,
+      logDirFailureChannel = new LogDirFailureChannel(10))
+
+    val now = System.currentTimeMillis()
+    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()
+  }
+
+  @After
+  def tearDown(): Unit = {
+    log.close()
+    Utils.delete(tmpDir)
+  }
+
+  @Test
+  def testPrintDataLog(): Unit = {
+
+    def verifyRecordsInOutput(checkKeysAndValues: Boolean, args: Array[String]): Unit = {
+      def isBatch(index: Int): Boolean = {
+        var i = 0
+        batches.zipWithIndex.foreach { case (batch, batchIndex) =>
+          if (i == index)
+            return true
+
+          i += 1
+
+          batch.records.indices.foreach { recordIndex =>
+            if (i == index)
+              return false
+            i += 1
+          }
+        }
+        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(_.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)) {
+          assertTrue(s"Not a valid batch-level message record: $line", line.startsWith(s"baseOffset: $offset lastOffset: "))
+          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
+        }
+      }
+    }
+
+    def verifyNoRecordsInOutput(args: Array[String]): Unit = {
+      val output = runDumpLogSegments(args)
+      assertFalse(s"Data should not have been printed: $output", output.matches("(?s).*offset: [0-9]* isvalid.*"))
+    }
+
+    // Verify that records are printed with --print-data-log even if --deep-iteration is not specified
+    verifyRecordsInOutput(true, Array("--print-data-log", "--files", logFilePath))
+    // Verify that records are printed with --print-data-log if --deep-iteration is also specified
+    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(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(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(false, Array("--deep-iteration", "--files", logFilePath))
+
+    // Verify that records are not printed by default
+    verifyNoRecordsInOutput(Array("--files", logFilePath))
+  }
+
+  @Test
+  def testDumpIndexMismatches(): Unit = {
+    val offsetMismatches = mutable.Map[String, List[(Long, Long)]]()
+    DumpLogSegments.dumpIndex(new File(indexFilePath), indexSanityOnly = false, verifyOnly = true, offsetMismatches,
+      Int.MaxValue)
+    assertEquals(Map.empty, offsetMismatches)
+  }
+
+  @Test
+  def testDumpTimeIndexErrors(): Unit = {
+    val errors = new TimeIndexDumpErrors
+    DumpLogSegments.dumpTimeIndex(new File(timeIndexFilePath), indexSanityOnly = false, verifyOnly = true, errors,
+      Int.MaxValue)
+    assertEquals(Map.empty, errors.misMatchesForTimeIndexFilesMap)
+    assertEquals(Map.empty, errors.outOfOrderTimestamp)
+    assertEquals(Map.empty, errors.shallowOffsetNotFound)
+  }
+
+  private def runDumpLogSegments(args: Array[String]): String = {
+    val outContent = new ByteArrayOutputStream
+    Console.withOut(outContent) {
+      DumpLogSegments.main(args)
+    }
+    outContent.toString
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/tools/MirrorMakerTest.scala b/core/src/test/scala/unit/kafka/tools/MirrorMakerTest.scala
index a3abd4c10e060..413b5ba9aab53 100644
--- a/core/src/test/scala/unit/kafka/tools/MirrorMakerTest.scala
+++ b/core/src/test/scala/unit/kafka/tools/MirrorMakerTest.scala
@@ -19,14 +19,18 @@ package kafka.tools
 
 import kafka.consumer.BaseConsumerRecord
 import org.apache.kafka.common.record.{RecordBatch, TimestampType}
-import scala.collection.JavaConverters._
+
+import scala.jdk.CollectionConverters._
 import org.junit.Assert._
 import org.junit.Test
 
+import scala.annotation.nowarn
+
+@nowarn("cat=deprecation")
 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 +46,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 +62,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 50023f8046407..2977a2b210c3c 100644
--- a/core/src/test/scala/unit/kafka/utils/CommandLineUtilsTest.scala
+++ b/core/src/test/scala/unit/kafka/utils/CommandLineUtilsTest.scala
@@ -17,6 +17,9 @@
 
 package kafka.utils
 
+import java.util.Properties
+
+import joptsimple.{OptionParser, OptionSpec}
 import org.junit.Assert._
 import org.junit.Test
 
@@ -24,32 +27,197 @@ class CommandLineUtilsTest {
 
 
   @Test(expected = classOf[java.lang.IllegalArgumentException])
-  def testParseEmptyArg() {
+  def testParseEmptyArg(): Unit = {
     val argArray = Array("my.empty.property=")
-    CommandLineUtils.parseKeyValueArgs(argArray, false)
+
+    CommandLineUtils.parseKeyValueArgs(argArray, acceptMissingValue = false)
   }
 
+  @Test(expected = classOf[java.lang.IllegalArgumentException])
+  def testParseEmptyArgWithNoDelimiter(): Unit = {
+    val argArray = Array("my.empty.property")
+
+    CommandLineUtils.parseKeyValueArgs(argArray, acceptMissingValue = false)
+  }
 
   @Test
-  def testParseEmptyArgAsValid() {
-    val argArray = Array("my.empty.property=")
+  def testParseEmptyArgAsValid(): Unit = {
+    val argArray = Array("my.empty.property=", "my.empty.property1")
     val props = CommandLineUtils.parseKeyValueArgs(argArray)
-    assertEquals("Value of a key with missing value should be an empty string",props.getProperty("my.empty.property"),"")
+
+    assertEquals("Value of a key with missing value should be an empty string", props.getProperty("my.empty.property"), "")
+    assertEquals("Value of a key with missing value with no delimiter should be an empty string", props.getProperty("my.empty.property1"), "")
   }
 
   @Test
-  def testParseSingleArg() {
+  def testParseSingleArg(): Unit = {
     val argArray = Array("my.property=value")
     val props = CommandLineUtils.parseKeyValueArgs(argArray)
-    assertEquals("Value of a single property should be 'value' ",props.getProperty("my.property"),"value")
+
+    assertEquals("Value of a single property should be 'value' ", props.getProperty("my.property"), "value")
   }
 
   @Test
-  def testParseArgs() {
+  def testParseArgs(): Unit = {
     val argArray = Array("first.property=first","second.property=second")
-    val props = CommandLineUtils.parseKeyValueArgs(argArray, false)
-    assertEquals("Value of first property should be 'first'",props.getProperty("first.property"),"first")
-    assertEquals("Value of second property should be 'second'",props.getProperty("second.property"),"second")
+    val props = CommandLineUtils.parseKeyValueArgs(argArray)
+
+    assertEquals("Value of first property should be 'first'", props.getProperty("first.property"), "first")
+    assertEquals("Value of second property should be 'second'", props.getProperty("second.property"), "second")
+  }
+
+  @Test
+  def testParseArgsWithMultipleDelimiters(): Unit = {
+    val argArray = Array("first.property==first", "second.property=second=", "third.property=thi=rd")
+    val props = CommandLineUtils.parseKeyValueArgs(argArray)
+
+    assertEquals("Value of first property should be '=first'", props.getProperty("first.property"), "=first")
+    assertEquals("Value of second property should be 'second='", props.getProperty("second.property"), "second=")
+    assertEquals("Value of second property should be 'thi=rd'", props.getProperty("third.property"), "thi=rd")
+  }
+
+  val props = new Properties()
+  val parser = new OptionParser(false)
+  var stringOpt : OptionSpec[String] = _
+  var intOpt : OptionSpec[java.lang.Integer] = _
+  var stringOptOptionalArg : OptionSpec[String] = _
+  var intOptOptionalArg : OptionSpec[java.lang.Integer] = _
+  var stringOptOptionalArgNoDefault : OptionSpec[String] = _
+  var intOptOptionalArgNoDefault : OptionSpec[java.lang.Integer] = _
+
+  def setUpOptions(): Unit = {
+    stringOpt = parser.accepts("str")
+      .withRequiredArg
+      .ofType(classOf[String])
+      .defaultsTo("default-string")
+    intOpt = parser.accepts("int")
+      .withRequiredArg()
+      .ofType(classOf[java.lang.Integer])
+      .defaultsTo(100)
+    stringOptOptionalArg = parser.accepts("str-opt")
+      .withOptionalArg
+      .ofType(classOf[String])
+      .defaultsTo("default-string-2")
+    intOptOptionalArg = parser.accepts("int-opt")
+      .withOptionalArg
+      .ofType(classOf[java.lang.Integer])
+      .defaultsTo(200)
+    stringOptOptionalArgNoDefault = parser.accepts("str-opt-nodef")
+      .withOptionalArg
+      .ofType(classOf[String])
+    intOptOptionalArgNoDefault = parser.accepts("int-opt-nodef")
+      .withOptionalArg
+      .ofType(classOf[java.lang.Integer])
+  }
+
+  @Test
+  def testMaybeMergeOptionsOverwriteExisting(): Unit = {
+    setUpOptions()
+
+    props.put("skey", "existing-string")
+    props.put("ikey", "300")
+    props.put("sokey", "existing-string-2")
+    props.put("iokey", "400")
+    props.put("sondkey", "existing-string-3")
+    props.put("iondkey", "500")
+
+    val options = parser.parse(
+      "--str", "some-string",
+      "--int", "600",
+      "--str-opt", "some-string-2",
+      "--int-opt", "700",
+      "--str-opt-nodef", "some-string-3",
+      "--int-opt-nodef", "800"
+    )
+
+    CommandLineUtils.maybeMergeOptions(props, "skey", options, stringOpt)
+    CommandLineUtils.maybeMergeOptions(props, "ikey", options, intOpt)
+    CommandLineUtils.maybeMergeOptions(props, "sokey", options, stringOptOptionalArg)
+    CommandLineUtils.maybeMergeOptions(props, "iokey", options, intOptOptionalArg)
+    CommandLineUtils.maybeMergeOptions(props, "sondkey", options, stringOptOptionalArgNoDefault)
+    CommandLineUtils.maybeMergeOptions(props, "iondkey", options, intOptOptionalArgNoDefault)
+
+    assertEquals("some-string", props.get("skey"))
+    assertEquals("600", props.get("ikey"))
+    assertEquals("some-string-2", props.get("sokey"))
+    assertEquals("700", props.get("iokey"))
+    assertEquals("some-string-3", props.get("sondkey"))
+    assertEquals("800", props.get("iondkey"))
+  }
+
+  @Test
+  def testMaybeMergeOptionsDefaultOverwriteExisting(): Unit = {
+    setUpOptions()
+
+    props.put("sokey", "existing-string")
+    props.put("iokey", "300")
+    props.put("sondkey", "existing-string-2")
+    props.put("iondkey", "400")
+
+    val options = parser.parse(
+      "--str-opt",
+      "--int-opt",
+      "--str-opt-nodef",
+      "--int-opt-nodef"
+    )
+
+    CommandLineUtils.maybeMergeOptions(props, "sokey", options, stringOptOptionalArg)
+    CommandLineUtils.maybeMergeOptions(props, "iokey", options, intOptOptionalArg)
+    CommandLineUtils.maybeMergeOptions(props, "sondkey", options, stringOptOptionalArgNoDefault)
+    CommandLineUtils.maybeMergeOptions(props, "iondkey", options, intOptOptionalArgNoDefault)
+
+    assertEquals("default-string-2", props.get("sokey"))
+    assertEquals("200", props.get("iokey"))
+    assertNull(props.get("sondkey"))
+    assertNull(props.get("iondkey"))
+  }
+
+  @Test
+  def testMaybeMergeOptionsDefaultValueIfNotExist(): Unit = {
+    setUpOptions()
+
+    val options = parser.parse()
+
+    CommandLineUtils.maybeMergeOptions(props, "skey", options, stringOpt)
+    CommandLineUtils.maybeMergeOptions(props, "ikey", options, intOpt)
+    CommandLineUtils.maybeMergeOptions(props, "sokey", options, stringOptOptionalArg)
+    CommandLineUtils.maybeMergeOptions(props, "iokey", options, intOptOptionalArg)
+    CommandLineUtils.maybeMergeOptions(props, "sondkey", options, stringOptOptionalArgNoDefault)
+    CommandLineUtils.maybeMergeOptions(props, "iondkey", options, intOptOptionalArgNoDefault)
+
+    assertEquals("default-string", props.get("skey"))
+    assertEquals("100", props.get("ikey"))
+    assertEquals("default-string-2", props.get("sokey"))
+    assertEquals("200", props.get("iokey"))
+    assertNull(props.get("sondkey"))
+    assertNull(props.get("iondkey"))
   }
 
+  @Test
+  def testMaybeMergeOptionsNotOverwriteExisting(): Unit = {
+    setUpOptions()
+
+    props.put("skey", "existing-string")
+    props.put("ikey", "300")
+    props.put("sokey", "existing-string-2")
+    props.put("iokey", "400")
+    props.put("sondkey", "existing-string-3")
+    props.put("iondkey", "500")
+
+    val options = parser.parse()
+
+    CommandLineUtils.maybeMergeOptions(props, "skey", options, stringOpt)
+    CommandLineUtils.maybeMergeOptions(props, "ikey", options, intOpt)
+    CommandLineUtils.maybeMergeOptions(props, "sokey", options, stringOptOptionalArg)
+    CommandLineUtils.maybeMergeOptions(props, "iokey", options, intOptOptionalArg)
+    CommandLineUtils.maybeMergeOptions(props, "sondkey", options, stringOptOptionalArgNoDefault)
+    CommandLineUtils.maybeMergeOptions(props, "iondkey", options, intOptOptionalArgNoDefault)
+
+    assertEquals("existing-string", props.get("skey"))
+    assertEquals("300", props.get("ikey"))
+    assertEquals("existing-string-2", props.get("sokey"))
+    assertEquals("400", props.get("iokey"))
+    assertEquals("existing-string-3", props.get("sondkey"))
+    assertEquals("500", props.get("iondkey"))
+  }
 }
diff --git a/core/src/test/scala/unit/kafka/utils/CoreUtilsTest.scala b/core/src/test/scala/unit/kafka/utils/CoreUtilsTest.scala
index 9ff47dd91aa86..c6bcd46496bdb 100755
--- a/core/src/test/scala/unit/kafka/utils/CoreUtilsTest.scala
+++ b/core/src/test/scala/unit/kafka/utils/CoreUtilsTest.scala
@@ -17,43 +17,93 @@
 
 package kafka.utils
 
-import java.util.{Arrays, UUID}
+import java.util.{Arrays, Base64, UUID}
 import java.util.concurrent.{ConcurrentHashMap, Executors, TimeUnit}
 import java.util.concurrent.atomic.AtomicInteger
 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.common.KafkaException
 import kafka.utils.CoreUtils.inLock
+import org.apache.kafka.common.KafkaException
 import org.junit.Test
-import org.apache.kafka.common.utils.{Base64, Utils}
+import org.apache.kafka.common.utils.Utils
 import org.slf4j.event.Level
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
+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)
   }
 
   @Test
-  def testCircularIterator() {
+  def testTryAll(): Unit = {
+    case class TestException(key: String) extends Exception
+
+    val recorded = mutable.Map.empty[String, Either[TestException, String]]
+    def recordingFunction(v: Either[TestException, String]): Unit = {
+      val key = v match {
+        case Right(key) => key
+        case Left(e) => e.key
+      }
+      recorded(key) = v
+    }
+
+    CoreUtils.tryAll(Seq(
+      () => recordingFunction(Right("valid-0")),
+      () => recordingFunction(Left(new TestException("exception-1"))),
+      () => recordingFunction(Right("valid-2")),
+      () => recordingFunction(Left(new TestException("exception-3")))
+    ))
+    var expected = Map(
+      "valid-0" -> Right("valid-0"),
+      "exception-1" -> Left(TestException("exception-1")),
+      "valid-2" -> Right("valid-2"),
+      "exception-3" -> Left(TestException("exception-3"))
+    )
+    assertEquals(expected, recorded)
+
+    recorded.clear()
+    CoreUtils.tryAll(Seq(
+      () => recordingFunction(Right("valid-0")),
+      () => recordingFunction(Right("valid-1"))
+    ))
+    expected = Map(
+      "valid-0" -> Right("valid-0"),
+      "valid-1" -> Right("valid-1")
+    )
+    assertEquals(expected, recorded)
+
+    recorded.clear()
+    CoreUtils.tryAll(Seq(
+      () => recordingFunction(Left(new TestException("exception-0"))),
+      () => recordingFunction(Left(new TestException("exception-1")))
+    ))
+    expected = Map(
+      "exception-0" -> Left(TestException("exception-0")),
+      "exception-1" -> Left(TestException("exception-1"))
+    )
+    assertEquals(expected, recorded)
+  }
+
+  @Test
+  def testCircularIterator(): Unit = {
     val l = List(1, 2)
     val itl = CoreUtils.circularIterator(l)
     assertEquals(1, itl.next())
     assertEquals(2, itl.next())
     assertEquals(1, itl.next())
     assertEquals(2, itl.next())
-    assertFalse(itl.hasDefiniteSize)
+    assertFalse(itl.isEmpty)
 
     val s = Set(1, 2)
     val its = CoreUtils.circularIterator(s)
@@ -65,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))))
@@ -73,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))
@@ -82,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", ""))
@@ -90,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) {
@@ -100,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)
@@ -113,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]
@@ -148,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)
@@ -159,17 +209,17 @@ 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.urlEncoderNoPadding.encodeToString(CoreUtils.getBytesFromUuid(UUID.fromString(
+    val clusterId1 = Base64.getUrlEncoder.withoutPadding.encodeToString(CoreUtils.getBytesFromUuid(UUID.fromString(
       "a149b4a3-06e1-4b49-a8cb-8a9c4a59fa46")))
     assertEquals(clusterId1, "oUm0owbhS0moy4qcSln6Rg")
     assertEquals(clusterId1.length, 22)
     assertTrue(clusterIdPattern.matcher(clusterId1).matches())
 
     // Test a UUID that has + or / characters in base64 encoding [d418ec02-277e-4853-81e6-afe30259daec ->(base64)-> 1BjsAid+SFOB5q/jAlna7A==]
-    val clusterId2 = Base64.urlEncoderNoPadding.encodeToString(CoreUtils.getBytesFromUuid(UUID.fromString(
+    val clusterId2 = Base64.getUrlEncoder.withoutPadding.encodeToString(CoreUtils.getBytesFromUuid(UUID.fromString(
       "d418ec02-277e-4853-81e6-afe30259daec")))
     assertEquals(clusterId2, "1BjsAid-SFOB5q_jAlna7A")
     assertEquals(clusterId2.length, 22)
@@ -177,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/IteratorTemplateTest.scala b/core/src/test/scala/unit/kafka/utils/IteratorTemplateTest.scala
deleted file mode 100644
index 67e64d57441db..0000000000000
--- a/core/src/test/scala/unit/kafka/utils/IteratorTemplateTest.scala
+++ /dev/null
@@ -1,57 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package kafka.utils
-
-import org.junit.Assert._
-import org.scalatest.Assertions
-import org.junit.Test
-
-class IteratorTemplateTest extends Assertions {
-  
-  val lst = 0 until 10
-  val iterator = new IteratorTemplate[Int]() {
-    var i = 0
-    override def makeNext() = {
-      if(i >= lst.size) {
-        allDone()
-      } else {
-        val item = lst(i)
-        i += 1
-        item
-      }
-    }
-  }
-
-  @Test
-  def testIterator() {
-    for(i <- 0 until 10) {
-      assertEquals("We should have an item to read.", true, iterator.hasNext)
-      assertEquals("Checking again shouldn't change anything.", true, iterator.hasNext)
-      assertEquals("Peeking at the item should show the right thing.", i, iterator.peek)
-      assertEquals("Peeking again shouldn't change anything", i, iterator.peek)
-      assertEquals("Getting the item should give the right thing.", i, iterator.next)
-    }
-    assertEquals("All gone!", false, iterator.hasNext)
-    intercept[NoSuchElementException] {
-      iterator.peek
-    }
-    intercept[NoSuchElementException] {
-      iterator.next
-    }
-  }
-  
-}
diff --git a/core/src/test/scala/unit/kafka/utils/JaasTestUtils.scala b/core/src/test/scala/unit/kafka/utils/JaasTestUtils.scala
index 1f9ccf3152b2b..3d257bae8e9cb 100644
--- a/core/src/test/scala/unit/kafka/utils/JaasTestUtils.scala
+++ b/core/src/test/scala/unit/kafka/utils/JaasTestUtils.scala
@@ -16,9 +16,12 @@
  */
 package kafka.utils
 
-import java.io.{File, BufferedWriter, FileWriter}
+import java.io.{BufferedWriter, File, FileWriter}
 import java.util.Properties
+
+import scala.collection.Seq
 import kafka.server.KafkaConfig
+import org.apache.kafka.clients.admin.ScramMechanism
 import org.apache.kafka.common.utils.Java
 
 object JaasTestUtils {
@@ -72,14 +75,26 @@ object JaasTestUtils {
 
   case class ScramLoginModule(username: String,
                               password: String,
-                              debug: Boolean = false) extends JaasModule {
+                              debug: Boolean = false,
+                              tokenProps: Map[String, String] = Map.empty) extends JaasModule {
 
     def name = "org.apache.kafka.common.security.scram.ScramLoginModule"
 
     def entries: Map[String, String] = Map(
       "username" -> username,
       "password" -> password
+    ) ++ tokenProps.map { case (name, value) => name -> value }
+  }
+
+  case class OAuthBearerLoginModule(username: String,
+                              debug: Boolean = false) extends JaasModule {
+
+    def name = "org.apache.kafka.common.security.oauthbearer.OAuthBearerLoginModule"
+
+    def entries: Map[String, String] = Map(
+      "unsecuredLoginStringClaim_sub" -> username
     )
+
   }
 
   sealed trait JaasModule {
@@ -113,7 +128,7 @@ object JaasTestUtils {
   val KafkaServerContextName = "KafkaServer"
   val KafkaServerPrincipalUnqualifiedName = "kafka"
   private val KafkaServerPrincipal = KafkaServerPrincipalUnqualifiedName + "/localhost@EXAMPLE.COM"
-  private val KafkaClientContextName = "KafkaClient"
+  val KafkaClientContextName = "KafkaClient"
   val KafkaClientPrincipalUnqualifiedName = "client"
   private val KafkaClientPrincipal = KafkaClientPrincipalUnqualifiedName + "@EXAMPLE.COM"
   val KafkaClientPrincipalUnqualifiedName2 = "client2"
@@ -133,13 +148,14 @@ object JaasTestUtils {
   val KafkaScramAdmin = "scram-admin"
   val KafkaScramAdminPassword = "scram-admin-secret"
 
+  val KafkaOAuthBearerUser = "oauthbearer-user"
+  val KafkaOAuthBearerUser2 = "oauthbearer-user2"
+  val KafkaOAuthBearerAdmin = "oauthbearer-admin"
+
   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))
@@ -153,9 +169,30 @@ object JaasTestUtils {
     jaasFile
   }
 
+  // Returns a SASL/SCRAM configuration using credentials for the given user and password
+  def scramClientLoginModule(mechanism: String, scramUser: String, scramPassword: String): String = {
+    if (ScramMechanism.fromMechanismName(mechanism) == ScramMechanism.UNKNOWN) {
+      throw new IllegalArgumentException("Unsupported SCRAM mechanism " + mechanism)
+    }
+    ScramLoginModule(
+      scramUser,
+      scramPassword
+    ).toString
+  }
+
   // Returns the dynamic configuration, using credentials for user #1
-  def clientLoginModule(mechanism: String, keytabLocation: Option[File]): String =
-    kafkaClientModule(mechanism, keytabLocation, KafkaClientPrincipal, KafkaPlainUser, KafkaPlainPassword, KafkaScramUser, KafkaScramPassword).toString
+  def clientLoginModule(mechanism: String, keytabLocation: Option[File], serviceName: String = serviceName): String =
+    kafkaClientModule(mechanism, keytabLocation, KafkaClientPrincipal, KafkaPlainUser, KafkaPlainPassword, KafkaScramUser, KafkaScramPassword, KafkaOAuthBearerUser, serviceName).toString
+
+  def tokenClientLoginModule(tokenId: String, password: String): String = {
+    ScramLoginModule(
+      tokenId,
+      password,
+      debug = false,
+      Map(
+        "tokenauth" -> "true"
+      )).toString
+  }
 
   def zkSections: Seq[JaasSection] = Seq(
     JaasSection(ZkServerContextName, Seq(ZkDigestModule(debug = false,
@@ -184,21 +221,28 @@ object JaasTestUtils {
             KafkaPlainUser -> KafkaPlainPassword,
             KafkaPlainUser2 -> KafkaPlainPassword2
           ))
-      case "SCRAM-SHA-256" | "SCRAM-SHA-512" =>
-        ScramLoginModule(
-          KafkaScramAdmin,
-          KafkaScramAdminPassword,
-          debug = false)
-      case mechanism => throw new IllegalArgumentException("Unsupported server mechanism " + mechanism)
+      case "OAUTHBEARER" =>
+        OAuthBearerLoginModule(KafkaOAuthBearerAdmin)
+      case mechanism => {
+        if (ScramMechanism.fromMechanismName(mechanism) != ScramMechanism.UNKNOWN) {
+          ScramLoginModule(
+            KafkaScramAdmin,
+            KafkaScramAdminPassword,
+            debug = false)
+        } else {
+          throw new IllegalArgumentException("Unsupported server mechanism " + mechanism)
+        }
+      }
     }
     JaasSection(contextName, modules)
   }
 
   // consider refactoring if more mechanisms are added
-  private def kafkaClientModule(mechanism: String, 
+  private def kafkaClientModule(mechanism: String,
       keytabLocation: Option[File], clientPrincipal: String,
       plainUser: String, plainPassword: String, 
-      scramUser: String, scramPassword: String): JaasModule = {
+      scramUser: String, scramPassword: String,
+      oauthBearerUser: String, serviceName: String = serviceName): JaasModule = {
     mechanism match {
       case "GSSAPI" =>
         Krb5LoginModule(
@@ -214,12 +258,20 @@ object JaasTestUtils {
           plainUser,
           plainPassword
         )
-      case "SCRAM-SHA-256" | "SCRAM-SHA-512" =>
-        ScramLoginModule(
-          scramUser,
-          scramPassword
+      case "OAUTHBEARER" =>
+        OAuthBearerLoginModule(
+          oauthBearerUser
         )
-      case mechanism => throw new IllegalArgumentException("Unsupported client mechanism " + mechanism)
+      case mechanism => {
+        if (ScramMechanism.fromMechanismName(mechanism) != ScramMechanism.UNKNOWN) {
+          ScramLoginModule(
+            scramUser,
+            scramPassword
+          )
+        } else {
+          throw new IllegalArgumentException("Unsupported client mechanism " + mechanism)
+        }
+      }
     }
   }
 
@@ -228,13 +280,13 @@ object JaasTestUtils {
    */
   def kafkaClientSection(mechanism: Option[String], keytabLocation: Option[File]): JaasSection = {
     JaasSection(KafkaClientContextName, mechanism.map(m =>
-      kafkaClientModule(m, keytabLocation, KafkaClientPrincipal2, KafkaPlainUser2, KafkaPlainPassword2, KafkaScramUser2, KafkaScramPassword2)).toSeq)
+      kafkaClientModule(m, keytabLocation, KafkaClientPrincipal2, KafkaPlainUser2, KafkaPlainPassword2, KafkaScramUser2, KafkaScramPassword2, KafkaOAuthBearerUser2)).toSeq)
   }
 
   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 8bba50bb46ca0..60d0234057a58 100644
--- a/core/src/test/scala/unit/kafka/utils/JsonTest.scala
+++ b/core/src/test/scala/unit/kafka/utils/JsonTest.scala
@@ -16,17 +16,28 @@
  */
 package kafka.utils
 
-import org.junit.Assert._
-import org.junit.Test
+import java.nio.charset.StandardCharsets
+
+import com.fasterxml.jackson.annotation.JsonProperty
+import com.fasterxml.jackson.core.JsonParseException
 import com.fasterxml.jackson.databind.JsonNode
 import com.fasterxml.jackson.databind.node._
+import kafka.utils.JsonTest.TestObject
 import kafka.utils.json.JsonValue
-import scala.collection.JavaConverters._
+import org.junit.Assert._
+import org.junit.Test
+
+import scala.jdk.CollectionConverters._
+import scala.collection.Map
+
+object JsonTest {
+  case class TestObject(@JsonProperty("foo") foo: String, @JsonProperty("bar") bar: Int)
+}
 
 class JsonTest {
 
   @Test
-  def testJsonParse() {
+  def testJsonParse(): Unit = {
     val jnf = JsonNodeFactory.instance
 
     assertEquals(Json.parseFull("{}"), Some(JsonValue(new ObjectNode(jnf))))
@@ -42,25 +53,73 @@ class JsonTest {
     val arrayNode = new ArrayNode(jnf)
     Vector(1, 2, 3).map(new IntNode(_)).foreach(arrayNode.add)
     assertEquals(Json.parseFull("[1, 2, 3]"), Some(JsonValue(arrayNode)))
+
+    // Test with encoder that properly escapes backslash and quotes
+    val map = Map("foo1" -> """bar1\,bar2""", "foo2" -> """\bar""").asJava
+    val encoded = Json.encodeAsString(map)
+    val decoded = Json.parseFull(encoded)
+    assertEquals(Json.parseFull("""{"foo1":"bar1\\,bar2", "foo2":"\\bar"}"""), decoded)
+  }
+
+  @Test
+  def testEncodeAsString(): Unit = {
+    assertEquals("null", Json.encodeAsString(null))
+    assertEquals("1", Json.encodeAsString(1))
+    assertEquals("1", Json.encodeAsString(1L))
+    assertEquals("1", Json.encodeAsString(1.toByte))
+    assertEquals("1", Json.encodeAsString(1.toShort))
+    assertEquals("1.0", Json.encodeAsString(1.0))
+    assertEquals(""""str"""", Json.encodeAsString("str"))
+    assertEquals("true", Json.encodeAsString(true))
+    assertEquals("false", Json.encodeAsString(false))
+    assertEquals("[]", Json.encodeAsString(Seq().asJava))
+    assertEquals("[null]", Json.encodeAsString(Seq(null).asJava))
+    assertEquals("[1,2,3]", Json.encodeAsString(Seq(1,2,3).asJava))
+    assertEquals("""[1,"2",[3],null]""", Json.encodeAsString(Seq(1,"2",Seq(3).asJava,null).asJava))
+    assertEquals("{}", Json.encodeAsString(Map().asJava))
+    assertEquals("""{"a":1,"b":2,"c":null}""", Json.encodeAsString(Map("a" -> 1, "b" -> 2, "c" -> null).asJava))
+    assertEquals("""{"a":[1,2],"c":[3,4]}""", Json.encodeAsString(Map("a" -> Seq(1,2).asJava, "c" -> Seq(3,4).asJava).asJava))
+    assertEquals("""{"a":[1,2],"b":[3,4],"c":null}""", Json.encodeAsString(Map("a" -> Seq(1,2).asJava, "b" -> Seq(3,4).asJava, "c" -> null).asJava))
+    assertEquals(""""str1\\,str2"""", Json.encodeAsString("""str1\,str2"""))
+    assertEquals(""""\"quoted\""""", Json.encodeAsString(""""quoted""""))
+  }
+
+  @Test
+  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))
+    assertEquals("1", new String(Json.encodeAsBytes(1.toByte), StandardCharsets.UTF_8))
+    assertEquals("1", new String(Json.encodeAsBytes(1.toShort), StandardCharsets.UTF_8))
+    assertEquals("1.0", new String(Json.encodeAsBytes(1.0), StandardCharsets.UTF_8))
+    assertEquals(""""str"""",  new String(Json.encodeAsBytes("str"), StandardCharsets.UTF_8))
+    assertEquals("true", new String(Json.encodeAsBytes(true), StandardCharsets.UTF_8))
+    assertEquals("false", new String(Json.encodeAsBytes(false), StandardCharsets.UTF_8))
+    assertEquals("[]", new String(Json.encodeAsBytes(Seq().asJava), StandardCharsets.UTF_8))
+    assertEquals("[null]", new String(Json.encodeAsBytes(Seq(null).asJava), StandardCharsets.UTF_8))
+    assertEquals("[1,2,3]", new String(Json.encodeAsBytes(Seq(1,2,3).asJava), StandardCharsets.UTF_8))
+    assertEquals("""[1,"2",[3],null]""", new String(Json.encodeAsBytes(Seq(1,"2",Seq(3).asJava,null).asJava), StandardCharsets.UTF_8))
+    assertEquals("{}", new String(Json.encodeAsBytes(Map().asJava), StandardCharsets.UTF_8))
+    assertEquals("""{"a":1,"b":2,"c":null}""", new String(Json.encodeAsBytes(Map("a" -> 1, "b" -> 2, "c" -> null).asJava), StandardCharsets.UTF_8))
+    assertEquals("""{"a":[1,2],"c":[3,4]}""", new String(Json.encodeAsBytes(Map("a" -> Seq(1,2).asJava, "c" -> Seq(3,4).asJava).asJava), StandardCharsets.UTF_8))
+    assertEquals("""{"a":[1,2],"b":[3,4],"c":null}""", new String(Json.encodeAsBytes(Map("a" -> Seq(1,2).asJava, "b" -> Seq(3,4).asJava, "c" -> null).asJava), StandardCharsets.UTF_8))
+    assertEquals(""""str1\\,str2"""", new String(Json.encodeAsBytes("""str1\,str2"""), StandardCharsets.UTF_8))
+    assertEquals(""""\"quoted\""""", new String(Json.encodeAsBytes(""""quoted""""), StandardCharsets.UTF_8))
+  }
+
+  @Test
+  def testParseTo() = {
+    val foo = "baz"
+    val bar = 1
+
+    val result = Json.parseStringAs[TestObject](s"""{"foo": "$foo", "bar": $bar}""")
+
+    assertEquals(Right(TestObject(foo, bar)), result)
   }
 
   @Test
-  def testJsonEncoding() {
-    assertEquals("null", Json.encode(null))
-    assertEquals("1", Json.encode(1))
-    assertEquals("1", Json.encode(1L))
-    assertEquals("1", Json.encode(1.toByte))
-    assertEquals("1", Json.encode(1.toShort))
-    assertEquals("1.0", Json.encode(1.0))
-    assertEquals("\"str\"", Json.encode("str"))
-    assertEquals("true", Json.encode(true))
-    assertEquals("false", Json.encode(false))
-    assertEquals("[]", Json.encode(Seq()))
-    assertEquals("[1,2,3]", Json.encode(Seq(1,2,3)))
-    assertEquals("[1,\"2\",[3]]", Json.encode(Seq(1,"2",Seq(3))))
-    assertEquals("{}", Json.encode(Map()))
-    assertEquals("{\"a\":1,\"b\":2}", Json.encode(Map("a" -> 1, "b" -> 2)))
-    assertEquals("{\"a\":[1,2],\"c\":[3,4]}", Json.encode(Map("a" -> Seq(1,2), "c" -> Seq(3,4))))
+  def testParseToWithInvalidJson() = {
+    val result = Json.parseStringAs[TestObject]("{invalid json}")
+    assertEquals(Left(classOf[JsonParseException]), result.left.map(_.getClass))
   }
-  
 }
diff --git a/core/src/test/scala/unit/kafka/utils/LogCaptureAppender.scala b/core/src/test/scala/unit/kafka/utils/LogCaptureAppender.scala
new file mode 100644
index 0000000000000..2d071452829ff
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/utils/LogCaptureAppender.scala
@@ -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 kafka.utils
+
+import org.apache.log4j.{AppenderSkeleton, Level, Logger}
+import org.apache.log4j.spi.LoggingEvent
+
+import scala.collection.mutable.ListBuffer
+
+class LogCaptureAppender extends AppenderSkeleton {
+  private val events: ListBuffer[LoggingEvent] = ListBuffer.empty
+
+  override protected def append(event: LoggingEvent): Unit = {
+    events.synchronized {
+      events += event
+    }
+  }
+
+  def getMessages: ListBuffer[LoggingEvent] = {
+    events.synchronized {
+      return events.clone()
+    }
+  }
+
+  override def close(): Unit = {
+    events.synchronized {
+      events.clear()
+    }
+  }
+
+  override def requiresLayout: Boolean = false
+}
+
+object LogCaptureAppender {
+  def createAndRegister(): LogCaptureAppender = {
+    val logCaptureAppender: LogCaptureAppender = new LogCaptureAppender
+    Logger.getRootLogger.addAppender(logCaptureAppender)
+    logCaptureAppender
+  }
+
+  def setClassLoggerLevel(clazz: Class[_], logLevel: Level): Level = {
+    val logger = Logger.getLogger(clazz)
+    val previousLevel = logger.getLevel
+    Logger.getLogger(clazz).setLevel(logLevel)
+    previousLevel
+  }
+
+  def unregister(logCaptureAppender: LogCaptureAppender): Unit = {
+    Logger.getRootLogger.removeAppender(logCaptureAppender)
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/utils/MockScheduler.scala b/core/src/test/scala/unit/kafka/utils/MockScheduler.scala
index c5f383c5b1bde..de3dfd7a03add 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
 
@@ -35,19 +35,20 @@ import org.apache.kafka.common.utils.Time
  * Incrementing the time to the exact next execution time of a task will result in that task executing (it as if execution itself takes no time).
  */
 class MockScheduler(val time: Time) extends Scheduler {
-  
+
   /* a priority queue of tasks ordered by next execution time */
-  var tasks = new PriorityQueue[MockTask]()
+  private val tasks = new PriorityQueue[MockTask]()
   
   def isStarted = true
 
-  def startup() {}
+  def startup(): Unit = {}
   
-  def shutdown() {
-    this synchronized {
-      tasks.foreach(_.fun())
-      tasks.clear()
-    }
+  def shutdown(): Unit = {
+    var currTask: Option[MockTask] = None
+    do {
+      currTask = poll(_ => true)
+      currTask.foreach(_.fun())
+    } while (currTask.nonEmpty)
   }
   
   /**
@@ -55,32 +56,53 @@ 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() {
-    this synchronized {
-      val now = time.milliseconds
-      while(tasks.nonEmpty && tasks.head.nextExecution <= now) {
-        /* pop and execute the task with the lowest next execution time */
-        val curr = tasks.dequeue
+  def tick(): Unit = {
+    val now = time.milliseconds
+    var currTask: Option[MockTask] = None
+    /* pop and execute the task with the lowest next execution time if ready */
+    do {
+      currTask = poll(_.nextExecution <= now)
+      currTask.foreach { curr =>
         curr.fun()
         /* if the task is periodic, reschedule it and re-enqueue */
         if(curr.periodic) {
           curr.nextExecution += curr.period
-          this.tasks += curr
+          add(curr)
         }
       }
+    } while (currTask.nonEmpty)
+  }
+
+  def schedule(name: String, fun: () => Unit, delay: Long = 0, period: Long = -1, unit: TimeUnit = TimeUnit.MILLISECONDS): ScheduledFuture[Unit] = {
+    val task = MockTask(name, fun, time.milliseconds + delay, period = period, time=time)
+    add(task)
+    tick()
+    task
+  }
+
+  def clear(): Unit = {
+    this synchronized {
+      tasks.clear()
     }
   }
-  
-  def schedule(name: String, fun: () => Unit, delay: Long = 0, period: Long = -1, unit: TimeUnit = TimeUnit.MILLISECONDS) {
+
+  private def poll(predicate: MockTask => Boolean): Option[MockTask] = {
     this synchronized {
-      tasks += MockTask(name, fun, time.milliseconds + delay, period = period)
-      tick()
+      if (tasks.nonEmpty && predicate.apply(tasks.head))
+        Some(tasks.dequeue())
+      else
+        None
+    }
+  }
+
+  private def add(task: MockTask): Unit = {
+    this synchronized {
+      tasks += task
     }
   }
-  
 }
 
-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)
@@ -90,4 +112,38 @@ 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] = (x, y) => 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/PasswordEncoderTest.scala b/core/src/test/scala/unit/kafka/utils/PasswordEncoderTest.scala
new file mode 100755
index 0000000000000..e24a518751bd6
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/utils/PasswordEncoderTest.scala
@@ -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 kafka.utils
+
+
+import javax.crypto.SecretKeyFactory
+
+import kafka.server.Defaults
+import org.apache.kafka.common.config.ConfigException
+import org.apache.kafka.common.config.types.Password
+import org.junit.Assert._
+import org.junit.Test
+
+class PasswordEncoderTest {
+
+  @Test
+  def testEncodeDecode(): Unit = {
+    val encoder = new PasswordEncoder(new Password("password-encoder-secret"),
+      None,
+      Defaults.PasswordEncoderCipherAlgorithm,
+      Defaults.PasswordEncoderKeyLength,
+      Defaults.PasswordEncoderIterations)
+    val password = "test-password"
+    val encoded = encoder.encode(new Password(password))
+    val encodedMap = CoreUtils.parseCsvMap(encoded)
+    assertEquals("4096", encodedMap(PasswordEncoder.IterationsProp))
+    assertEquals("128", encodedMap(PasswordEncoder.KeyLengthProp))
+    val defaultKeyFactoryAlgorithm = try {
+      SecretKeyFactory.getInstance("PBKDF2WithHmacSHA512")
+      "PBKDF2WithHmacSHA512"
+    } catch {
+      case _: Exception => "PBKDF2WithHmacSHA1"
+    }
+    assertEquals(defaultKeyFactoryAlgorithm, encodedMap(PasswordEncoder.KeyFactoryAlgorithmProp))
+    assertEquals("AES/CBC/PKCS5Padding", encodedMap(PasswordEncoder.CipherAlgorithmProp))
+
+    verifyEncodedPassword(encoder, password, encoded)
+  }
+
+  @Test
+  def testEncoderConfigChange(): Unit = {
+    val encoder = new PasswordEncoder(new Password("password-encoder-secret"),
+      Some("PBKDF2WithHmacSHA1"),
+      "DES/CBC/PKCS5Padding",
+      64,
+      1024)
+    val password = "test-password"
+    val encoded = encoder.encode(new Password(password))
+    val encodedMap = CoreUtils.parseCsvMap(encoded)
+    assertEquals("1024", encodedMap(PasswordEncoder.IterationsProp))
+    assertEquals("64", encodedMap(PasswordEncoder.KeyLengthProp))
+    assertEquals("PBKDF2WithHmacSHA1", encodedMap(PasswordEncoder.KeyFactoryAlgorithmProp))
+    assertEquals("DES/CBC/PKCS5Padding", encodedMap(PasswordEncoder.CipherAlgorithmProp))
+
+    // Test that decoding works even if PasswordEncoder algorithm, iterations etc. are altered
+    val decoder = new PasswordEncoder(new Password("password-encoder-secret"),
+      Some("PBKDF2WithHmacSHA1"),
+      "AES/CBC/PKCS5Padding",
+      128,
+      2048)
+    assertEquals(password, decoder.decode(encoded).value)
+
+    // Test that decoding fails if secret is altered
+    val decoder2 = new PasswordEncoder(new Password("secret-2"),
+      Some("PBKDF2WithHmacSHA1"),
+      "AES/CBC/PKCS5Padding",
+      128,
+      1024)
+    try {
+      decoder2.decode(encoded)
+    } catch {
+      case e: ConfigException => // expected exception
+    }
+  }
+
+  @Test
+  def testEncodeDecodeAlgorithms(): Unit = {
+
+    def verifyEncodeDecode(keyFactoryAlg: Option[String], cipherAlg: String, keyLength: Int): Unit = {
+      val encoder = new PasswordEncoder(new Password("password-encoder-secret"),
+        keyFactoryAlg,
+        cipherAlg,
+        keyLength,
+        Defaults.PasswordEncoderIterations)
+      val password = "test-password"
+      val encoded = encoder.encode(new Password(password))
+      verifyEncodedPassword(encoder, password, encoded)
+    }
+
+    verifyEncodeDecode(keyFactoryAlg = None, "DES/CBC/PKCS5Padding", keyLength = 64)
+    verifyEncodeDecode(keyFactoryAlg = None, "DESede/CBC/PKCS5Padding", keyLength = 192)
+    verifyEncodeDecode(keyFactoryAlg = None, "AES/CBC/PKCS5Padding", keyLength = 128)
+    verifyEncodeDecode(keyFactoryAlg = None, "AES/CFB/PKCS5Padding", keyLength = 128)
+    verifyEncodeDecode(keyFactoryAlg = None, "AES/OFB/PKCS5Padding", keyLength = 128)
+    verifyEncodeDecode(keyFactoryAlg = Some("PBKDF2WithHmacSHA1"), Defaults.PasswordEncoderCipherAlgorithm, keyLength = 128)
+    verifyEncodeDecode(keyFactoryAlg = None, "AES/GCM/NoPadding", keyLength = 128)
+    verifyEncodeDecode(keyFactoryAlg = Some("PBKDF2WithHmacSHA256"), Defaults.PasswordEncoderCipherAlgorithm, keyLength = 128)
+    verifyEncodeDecode(keyFactoryAlg = Some("PBKDF2WithHmacSHA512"), Defaults.PasswordEncoderCipherAlgorithm, keyLength = 128)
+  }
+
+  private def verifyEncodedPassword(encoder: PasswordEncoder, password: String, encoded: String): Unit = {
+    val encodedMap = CoreUtils.parseCsvMap(encoded)
+    assertEquals(password.length.toString, encodedMap(PasswordEncoder.PasswordLengthProp))
+    assertNotNull("Invalid salt", encoder.base64Decode(encodedMap("salt")))
+    assertNotNull("Invalid encoding parameters", encoder.base64Decode(encodedMap(PasswordEncoder.InitializationVectorProp)))
+    assertNotNull("Invalid encoded password", encoder.base64Decode(encodedMap(PasswordEncoder.EncyrptedPasswordProp)))
+    assertEquals(password, encoder.decode(encoded).value)
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/utils/PoolTest.scala b/core/src/test/scala/unit/kafka/utils/PoolTest.scala
new file mode 100644
index 0000000000000..74751806a11ef
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/utils/PoolTest.scala
@@ -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 kafka.utils
+
+import org.junit.Assert.assertEquals
+import org.junit.Test
+
+
+class PoolTest {
+  @Test
+  def testRemoveAll(): Unit = {
+    val pool = new Pool[Int, String]
+    pool.put(1, "1")
+    pool.put(2, "2")
+    pool.put(3, "3")
+
+    assertEquals(3, pool.size)
+
+    pool.removeAll(Seq(1, 2))
+    assertEquals(1, pool.size)
+    assertEquals("3", pool.get(3))
+    pool.removeAll(Seq(3))
+    assertEquals(0, pool.size)
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/utils/QuotaUtilsTest.scala b/core/src/test/scala/unit/kafka/utils/QuotaUtilsTest.scala
new file mode 100755
index 0000000000000..6d1b960fd0d6d
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/utils/QuotaUtilsTest.scala
@@ -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 kafka.utils
+
+import java.util.concurrent.TimeUnit
+
+import org.apache.kafka.common.MetricName
+import org.apache.kafka.common.metrics.{KafkaMetric, MetricConfig, Quota, QuotaViolationException}
+import org.apache.kafka.common.metrics.stats.{Rate, Value}
+
+import scala.jdk.CollectionConverters._
+import org.junit.Assert._
+import org.junit.Test
+import org.scalatest.Assertions.assertThrows
+
+class QuotaUtilsTest {
+
+  private val time = new MockTime
+  private val numSamples = 10
+  private val sampleWindowSec = 1
+  private val maxThrottleTimeMs = 500
+  private val metricName = new MetricName("test-metric", "groupA", "testA", Map.empty.asJava)
+
+  @Test
+  def testThrottleTimeObservedRateEqualsQuota(): Unit = {
+    val numSamples = 10
+    val observedValue = 16.5
+
+    assertEquals(0, throttleTime(observedValue, observedValue, numSamples))
+
+    // should be independent of window size
+    assertEquals(0, throttleTime(observedValue, observedValue, numSamples + 1))
+  }
+
+  @Test
+  def testThrottleTimeObservedRateBelowQuota(): Unit = {
+    val observedValue = 16.5
+    val quota = 20.4
+    assertTrue(throttleTime(observedValue, quota, numSamples) < 0)
+
+    // should be independent of window size
+    assertTrue(throttleTime(observedValue, quota, numSamples + 1) < 0)
+  }
+
+  @Test
+  def testThrottleTimeObservedRateAboveQuota(): Unit = {
+    val quota = 50.0
+    val observedValue = 100.0
+    assertEquals(2000, throttleTime(observedValue, quota, 3))
+  }
+
+  @Test
+  def testBoundedThrottleTimeObservedRateEqualsQuota(): Unit = {
+    val observedValue = 18.2
+    assertEquals(0, boundedThrottleTime(observedValue, observedValue, numSamples, maxThrottleTimeMs))
+
+    // should be independent of window size
+    assertEquals(0, boundedThrottleTime(observedValue, observedValue, numSamples + 1, maxThrottleTimeMs))
+  }
+
+  @Test
+  def testBoundedThrottleTimeObservedRateBelowQuota(): Unit = {
+    val observedValue = 16.5
+    val quota = 22.4
+
+    assertTrue(boundedThrottleTime(observedValue, quota, numSamples, maxThrottleTimeMs) < 0)
+
+    // should be independent of window size
+    assertTrue(boundedThrottleTime(observedValue, quota, numSamples + 1, maxThrottleTimeMs) < 0)
+  }
+
+  @Test
+  def testBoundedThrottleTimeObservedRateAboveQuotaBelowLimit(): Unit = {
+    val quota = 50.0
+    val observedValue = 55.0
+    assertEquals(100, boundedThrottleTime(observedValue, quota, 2, maxThrottleTimeMs))
+  }
+
+  @Test
+  def testBoundedThrottleTimeObservedRateAboveQuotaAboveLimit(): Unit = {
+    val quota = 50.0
+    val observedValue = 100.0
+    assertEquals(maxThrottleTimeMs, boundedThrottleTime(observedValue, quota, numSamples, maxThrottleTimeMs))
+  }
+
+  @Test
+  def testThrottleTimeThrowsExceptionIfProvidedNonRateMetric(): Unit = {
+    val testMetric = new KafkaMetric(new Object(), metricName, new Value(), new MetricConfig, time);
+
+    assertThrows[IllegalArgumentException] {
+      QuotaUtils.throttleTime(new QuotaViolationException(testMetric, 10.0, 20.0), time.milliseconds)
+    }
+  }
+
+  @Test
+  def testBoundedThrottleTimeThrowsExceptionIfProvidedNonRateMetric(): Unit = {
+    val testMetric = new KafkaMetric(new Object(), metricName, new Value(), new MetricConfig, time);
+
+    assertThrows[IllegalArgumentException] {
+      QuotaUtils.boundedThrottleTime(new QuotaViolationException(testMetric, 10.0, 20.0), maxThrottleTimeMs, time.milliseconds)
+    }
+  }
+
+  // the `metric` passed into the returned QuotaViolationException will return windowSize = 'numSamples' - 1
+  private def quotaViolationException(observedValue: Double, quota: Double, numSamples: Int): QuotaViolationException = {
+    val metricConfig = new MetricConfig()
+      .timeWindow(sampleWindowSec, TimeUnit.SECONDS)
+      .samples(numSamples)
+      .quota(new Quota(quota, true))
+    val metric = new KafkaMetric(new Object(), metricName, new Rate(), metricConfig, time)
+    new QuotaViolationException(metric, observedValue, quota)
+  }
+
+  private def throttleTime(observedValue: Double, quota: Double, numSamples: Int): Long = {
+    val e = quotaViolationException(observedValue, quota, numSamples)
+    QuotaUtils.throttleTime(e, time.milliseconds)
+  }
+
+  private def boundedThrottleTime(observedValue: Double, quota: Double, numSamples: Int, maxThrottleTime: Long): Long = {
+    val e = quotaViolationException(observedValue, quota, numSamples)
+    QuotaUtils.boundedThrottleTime(e, maxThrottleTime, time.milliseconds)
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/utils/ReplicationUtilsTest.scala b/core/src/test/scala/unit/kafka/utils/ReplicationUtilsTest.scala
index ebab756e9b614..e8de8d9ecfb08 100644
--- a/core/src/test/scala/unit/kafka/utils/ReplicationUtilsTest.scala
+++ b/core/src/test/scala/unit/kafka/utils/ReplicationUtilsTest.scala
@@ -17,90 +17,79 @@
 
 package kafka.utils
 
-import kafka.controller.LeaderIsrAndControllerEpoch
-import kafka.server.{KafkaConfig, ReplicaFetcherManager}
+import kafka.server.{KafkaConfig, ReplicaFetcherManager, ReplicaManager}
 import kafka.api.LeaderAndIsr
-import kafka.zk.ZooKeeperTestHarness
+import kafka.controller.LeaderIsrAndControllerEpoch
+import kafka.log.{Log, LogManager}
+import kafka.zk._
 import org.apache.kafka.common.TopicPartition
 import org.junit.Assert._
 import org.junit.{Before, Test}
 import org.easymock.EasyMock
 
-
 class ReplicationUtilsTest extends ZooKeeperTestHarness {
-  val topic = "my-topic-test"
-  val partitionId = 0
-  val brokerId = 1
-  val leaderEpoch = 1
-  val controllerEpoch = 1
-  val zkVersion = 1
-  val topicPath = "/brokers/topics/my-topic-test/partitions/0/state"
-  val topicData = Json.encode(Map("controller_epoch" -> 1, "leader" -> 1,
-    "versions" -> 1, "leader_epoch" -> 1,"isr" -> List(1,2)))
-  val topicDataVersionMismatch = Json.encode(Map("controller_epoch" -> 1, "leader" -> 1,
-    "versions" -> 2, "leader_epoch" -> 1,"isr" -> List(1,2)))
-  val topicDataMismatch = Json.encode(Map("controller_epoch" -> 1, "leader" -> 1,
-    "versions" -> 2, "leader_epoch" -> 2,"isr" -> List(1,2)))
-
-  val topicDataLeaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(LeaderAndIsr(1,leaderEpoch,List(1,2),0), controllerEpoch)
+  private val zkVersion = 1
+  private val topic = "my-topic-test"
+  private val partition = 0
+  private val leader = 1
+  private val leaderEpoch = 1
+  private val controllerEpoch = 1
+  private val isr = List(1, 2)
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
-    zkUtils.createPersistentPath(topicPath, topicData)
+    zkClient.makeSurePersistentPathExists(TopicZNode.path(topic))
+    val topicPartition = new TopicPartition(topic, partition)
+    val leaderAndIsr = LeaderAndIsr(leader, leaderEpoch, isr, 1)
+    val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(leaderAndIsr, controllerEpoch)
+    zkClient.createTopicPartitionStatesRaw(Map(topicPartition -> leaderIsrAndControllerEpoch), ZkVersion.MatchAnyVersion)
   }
 
   @Test
-  def testUpdateLeaderAndIsr() {
+  def testUpdateLeaderAndIsr(): Unit = {
     val configs = TestUtils.createBrokerConfigs(1, zkConnect).map(KafkaConfig.fromProps)
-    val log = EasyMock.createMock(classOf[kafka.log.Log])
+    val log: Log = EasyMock.createMock(classOf[Log])
     EasyMock.expect(log.logEndOffset).andReturn(20).anyTimes()
     EasyMock.expect(log)
     EasyMock.replay(log)
 
-    val logManager = EasyMock.createMock(classOf[kafka.log.LogManager])
-    EasyMock.expect(logManager.getLog(new TopicPartition(topic, partitionId), false)).andReturn(Some(log)).anyTimes()
+    val logManager: LogManager = EasyMock.createMock(classOf[LogManager])
+    EasyMock.expect(logManager.getLog(new TopicPartition(topic, partition), false)).andReturn(Some(log)).anyTimes()
     EasyMock.replay(logManager)
 
-    val replicaManager = EasyMock.createMock(classOf[kafka.server.ReplicaManager])
+    val replicaManager: ReplicaManager = EasyMock.createMock(classOf[ReplicaManager])
     EasyMock.expect(replicaManager.config).andReturn(configs.head)
     EasyMock.expect(replicaManager.logManager).andReturn(logManager)
     EasyMock.expect(replicaManager.replicaFetcherManager).andReturn(EasyMock.createMock(classOf[ReplicaFetcherManager]))
     EasyMock.expect(replicaManager.zkClient).andReturn(zkClient)
     EasyMock.replay(replicaManager)
 
-    zkUtils.makeSurePersistentPathExists(ZkUtils.IsrChangeNotificationPath)
+    zkClient.makeSurePersistentPathExists(IsrChangeNotificationZNode.path)
 
-    val replicas = List(0,1)
+    val replicas = List(0, 1)
 
     // regular update
-    val newLeaderAndIsr1 = new LeaderAndIsr(brokerId, leaderEpoch, replicas, 0)
-    val (updateSucceeded1,newZkVersion1) = ReplicationUtils.updateLeaderAndIsr(zkClient,
-      "my-topic-test", partitionId, newLeaderAndIsr1, controllerEpoch, 0)
+    val newLeaderAndIsr1 = new LeaderAndIsr(leader, leaderEpoch, replicas, 0)
+    val (updateSucceeded1, newZkVersion1) = ReplicationUtils.updateLeaderAndIsr(zkClient,
+      new TopicPartition(topic, partition), newLeaderAndIsr1, controllerEpoch)
     assertTrue(updateSucceeded1)
     assertEquals(newZkVersion1, 1)
 
     // mismatched zkVersion with the same data
-    val newLeaderAndIsr2 = new LeaderAndIsr(brokerId, leaderEpoch, replicas, zkVersion + 1)
-    val (updateSucceeded2,newZkVersion2) = ReplicationUtils.updateLeaderAndIsr(zkClient,
-      "my-topic-test", partitionId, newLeaderAndIsr2, controllerEpoch, zkVersion + 1)
+    val newLeaderAndIsr2 = new LeaderAndIsr(leader, leaderEpoch, replicas, zkVersion + 1)
+    val (updateSucceeded2, newZkVersion2) = ReplicationUtils.updateLeaderAndIsr(zkClient,
+      new TopicPartition(topic, partition), newLeaderAndIsr2, controllerEpoch)
     assertTrue(updateSucceeded2)
     // returns true with existing zkVersion
-    assertEquals(newZkVersion2,1)
+    assertEquals(newZkVersion2, 1)
 
     // mismatched zkVersion and leaderEpoch
-    val newLeaderAndIsr3 = new LeaderAndIsr(brokerId, leaderEpoch + 1, replicas, zkVersion + 1)
-    val (updateSucceeded3,newZkVersion3) = ReplicationUtils.updateLeaderAndIsr(zkClient,
-      "my-topic-test", partitionId, newLeaderAndIsr3, controllerEpoch, zkVersion + 1)
+    val newLeaderAndIsr3 = new LeaderAndIsr(leader, leaderEpoch + 1, replicas, zkVersion + 1)
+    val (updateSucceeded3, newZkVersion3) = ReplicationUtils.updateLeaderAndIsr(zkClient,
+      new TopicPartition(topic, partition), newLeaderAndIsr3, controllerEpoch)
     assertFalse(updateSucceeded3)
-    assertEquals(newZkVersion3,-1)
-  }
-
-  @Test
-  def testGetLeaderIsrAndEpochForPartition() {
-    val leaderIsrAndControllerEpoch = ReplicationUtils.getLeaderIsrAndEpochForPartition(zkUtils, topic, partitionId)
-    assertEquals(topicDataLeaderIsrAndControllerEpoch, leaderIsrAndControllerEpoch.get)
-    assertEquals(None, ReplicationUtils.getLeaderIsrAndEpochForPartition(zkUtils, topic, partitionId + 1))
+    assertEquals(newZkVersion3, -1)
   }
 
 }
diff --git a/core/src/test/scala/unit/kafka/utils/SchedulerTest.scala b/core/src/test/scala/unit/kafka/utils/SchedulerTest.scala
index dbb818c046948..42241dee71cfd 100644
--- a/core/src/test/scala/unit/kafka/utils/SchedulerTest.scala
+++ b/core/src/test/scala/unit/kafka/utils/SchedulerTest.scala
@@ -16,10 +16,15 @@
  */
 package kafka.utils
 
-import org.junit.Assert._
+import java.util.Properties
 import java.util.concurrent.atomic._
-import org.junit.{Test, After, Before}
+import java.util.concurrent.{CountDownLatch, Executors, TimeUnit}
+
+import kafka.log.{Log, LogConfig, LogManager, ProducerStateManager}
+import kafka.server.{BrokerTopicStats, LogDirFailureChannel}
 import kafka.utils.TestUtils.retry
+import org.junit.Assert._
+import org.junit.{After, Before, Test}
 
 class SchedulerTest {
 
@@ -29,17 +34,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 +58,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 +72,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 +89,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 +97,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 +112,55 @@ class SchedulerTest {
     mockTime.sleep(1)
     assertEquals(2, counter1.get())
   }
-}
\ No newline at end of file
+
+  @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)))
+  }
+
+  /**
+   * Verify that scheduler lock is not held when invoking task method, allowing new tasks to be scheduled
+   * when another is being executed. This is required to avoid deadlocks when:
+   *   a) Thread1 executes a task which attempts to acquire LockA
+   *   b) Thread2 holding LockA attempts to schedule a new task
+   */
+  @Test(timeout = 15000)
+  def testMockSchedulerLocking(): Unit = {
+    val initLatch = new CountDownLatch(1)
+    val completionLatch = new CountDownLatch(2)
+    val taskLatches = List(new CountDownLatch(1), new CountDownLatch(1))
+    def scheduledTask(taskLatch: CountDownLatch): Unit = {
+      initLatch.countDown()
+      assertTrue("Timed out waiting for latch", taskLatch.await(30, TimeUnit.SECONDS))
+      completionLatch.countDown()
+    }
+    mockTime.scheduler.schedule("test1", () => scheduledTask(taskLatches.head), delay=1)
+    val tickExecutor = Executors.newSingleThreadScheduledExecutor()
+    try {
+      tickExecutor.scheduleWithFixedDelay(() => mockTime.sleep(1), 0, 1, TimeUnit.MILLISECONDS)
+
+      // wait for first task to execute and then schedule the next task while the first one is running
+      assertTrue(initLatch.await(10, TimeUnit.SECONDS))
+      mockTime.scheduler.schedule("test2", () => scheduledTask(taskLatches(1)), delay = 1)
+
+      taskLatches.foreach(_.countDown())
+      assertTrue("Tasks did not complete", completionLatch.await(10, TimeUnit.SECONDS))
+
+    } finally {
+      tickExecutor.shutdownNow()
+    }
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/utils/ShutdownableThreadTest.scala b/core/src/test/scala/unit/kafka/utils/ShutdownableThreadTest.scala
index 31f7efb5ed3b1..672ac87a35b4a 100644
--- a/core/src/test/scala/unit/kafka/utils/ShutdownableThreadTest.scala
+++ b/core/src/test/scala/unit/kafka/utils/ShutdownableThreadTest.scala
@@ -38,7 +38,7 @@ class ShutdownableThreadTest {
     }
     val latch = new CountDownLatch(1)
     val thread = new ShutdownableThread("shutdownable-thread-test") {
-      override def doWork: Unit = {
+      override def doWork(): Unit = {
         latch.countDown()
         throw new FatalExitError
       }
diff --git a/core/src/test/scala/unit/kafka/utils/TestUtils.scala b/core/src/test/scala/unit/kafka/utils/TestUtils.scala
index 360b9dc383182..5368c20ce5403 100755
--- a/core/src/test/scala/unit/kafka/utils/TestUtils.scala
+++ b/core/src/test/scala/unit/kafka/utils/TestUtils.scala
@@ -14,53 +14,63 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-
 package kafka.utils
 
 import java.io._
 import java.nio._
 import java.nio.channels._
 import java.nio.charset.{Charset, StandardCharsets}
+import java.nio.file.{Files, StandardOpenOption}
 import java.security.cert.X509Certificate
-import java.util.{Collections, Properties}
-import java.util.concurrent.{Callable, Executors, TimeUnit}
-import javax.net.ssl.X509TrustManager
+import java.time.Duration
+import java.util.concurrent.atomic.AtomicInteger
+import java.util.{Arrays, Collections, Properties}
+import java.util.concurrent.{Callable, ExecutionException, Executors, TimeUnit}
 
-import kafka.admin.AdminUtils
+import javax.net.ssl.X509TrustManager
 import kafka.api._
-import kafka.cluster.{Broker, EndPoint}
-import kafka.common.TopicAndPartition
-import kafka.consumer.{ConsumerConfig, ConsumerTimeoutException, KafkaStream}
+import kafka.cluster.{Broker, EndPoint, IsrChangeListener}
 import kafka.log._
-import kafka.message._
-import kafka.producer._
-import kafka.security.auth.{Acl, Authorizer, Resource}
-import kafka.serializer.{DefaultEncoder, Encoder, StringEncoder}
+import kafka.security.auth.{Acl, Resource, Authorizer => LegacyAuthorizer}
 import kafka.server._
 import kafka.server.checkpoints.OffsetCheckpointFile
-import ZkUtils._
-import Implicits._
+import com.yammer.metrics.core.Meter
+import kafka.controller.LeaderIsrAndControllerEpoch
+import kafka.metrics.KafkaYammerMetrics
+import kafka.utils.Implicits._
+import kafka.zk._
 import org.apache.kafka.clients.CommonClientConfigs
-import org.apache.kafka.clients.consumer.{ConsumerRecord, KafkaConsumer, OffsetAndMetadata, RangeAssignor}
+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.TopicPartition
+import org.apache.kafka.common.acl.{AccessControlEntry, AccessControlEntryFilter, AclBinding, AclBindingFilter}
+import org.apache.kafka.common.config.ConfigResource
+import org.apache.kafka.common.errors.{KafkaStorageException, 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.quota.{ClientQuotaAlteration, ClientQuotaEntity}
 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.{ByteArraySerializer, Serializer}
-import org.apache.kafka.common.utils.Time
+import org.apache.kafka.common.serialization.{ByteArrayDeserializer, ByteArraySerializer, Deserializer, IntegerSerializer, Serializer}
+import org.apache.kafka.common.utils.{Time, Utils}
 import org.apache.kafka.common.utils.Utils._
+import org.apache.kafka.common.{KafkaFuture, TopicPartition}
+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.jdk.CollectionConverters._
 import scala.collection.mutable.{ArrayBuffer, ListBuffer}
-import scala.util.Try
+import scala.collection.{Map, Seq, mutable}
+import scala.concurrent.duration.FiniteDuration
+import scala.concurrent.{Await, ExecutionContext, Future}
 
 /**
  * Utility functions to help with testing
@@ -72,15 +82,25 @@ object TestUtils extends Logging {
   /* 0 gives a random port; you can then retrieve the assigned port from the Socket object. */
   val RandomPort = 0
 
+  /* Incorrect broker port which can used by kafka clients in tests. This port should not be used
+   by any other service and hence we use a reserved port. */
+  val IncorrectBrokerPort = 225
+
   /** Port to use for unit tests that mock/don't require a real ZK server. */
   val MockZkPort = 1
   /** ZooKeeper connection string to use for unit tests that mock/don't require a real ZK server. */
   val MockZkConnect = "127.0.0.1:" + MockZkPort
+  // CN in SSL certificates - this is used for endpoint validation when enabled
+  val SslCertificateCn = "localhost"
 
   private val transactionStatusKey = "transactionStatus"
   private val committedValue : Array[Byte] = "committed".getBytes(StandardCharsets.UTF_8)
   private val abortedValue : Array[Byte] = "aborted".getBytes(StandardCharsets.UTF_8)
 
+  sealed trait LogDirFailureType
+  case object Roll extends LogDirFailureType
+  case object Checkpoint extends LogDirFailureType
+
   /**
    * Create a temporary directory
    */
@@ -120,7 +140,8 @@ object TestUtils extends Logging {
   /**
    * Create a temporary file and return an open file channel for this file
    */
-  def tempChannel(): FileChannel = new RandomAccessFile(tempFile(), "rw").getChannel()
+  def tempChannel(): FileChannel =
+    FileChannel.open(tempFile().toPath, StandardOpenOption.READ, StandardOpenOption.WRITE)
 
   /**
    * Create a kafka server instance with appropriate test settings
@@ -129,7 +150,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
   }
@@ -140,6 +169,11 @@ object TestUtils extends Logging {
   def createBroker(id: Int, host: String, port: Int, securityProtocol: SecurityProtocol = SecurityProtocol.PLAINTEXT): Broker =
     new Broker(id, host, port, ListenerName.forSecurityProtocol(securityProtocol), securityProtocol)
 
+  def createBrokerAndEpoch(id: Int, host: String, port: Int, securityProtocol: SecurityProtocol = SecurityProtocol.PLAINTEXT,
+                           epoch: Long = 0): (Broker, Long) = {
+    (new Broker(id, host, port, ListenerName.forSecurityProtocol(securityProtocol), securityProtocol), epoch)
+  }
+
   /**
    * Create a test config for the provided parameters.
    *
@@ -148,7 +182,7 @@ object TestUtils extends Logging {
   def createBrokerConfigs(numConfigs: Int,
     zkConnect: String,
     enableControlledShutdown: Boolean = true,
-    enableDeleteTopic: Boolean = false,
+    enableDeleteTopic: Boolean = true,
     interBrokerSecurityProtocol: Option[SecurityProtocol] = None,
     trustStoreFile: Option[File] = None,
     saslProperties: Option[Properties] = None,
@@ -157,11 +191,15 @@ object TestUtils extends Logging {
     enableSaslPlaintext: Boolean = false,
     enableSaslSsl: Boolean = false,
     rackInfo: Map[Int, String] = Map(),
-    logDirCount: Int = 1): Seq[Properties] = {
+    logDirCount: Int = 1,
+    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)
+        enableSaslPlaintext = enableSaslPlaintext, enableSaslSsl = enableSaslSsl, rack = rackInfo.get(node), logDirCount = logDirCount, enableToken = enableToken,
+        numPartitions = numPartitions, defaultReplicationFactor = defaultReplicationFactor)
     }
   }
 
@@ -184,11 +222,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))
   }
 
   /**
@@ -199,7 +241,7 @@ object TestUtils extends Logging {
   def createBrokerConfig(nodeId: Int,
                          zkConnect: String,
                          enableControlledShutdown: Boolean = true,
-                         enableDeleteTopic: Boolean = false,
+                         enableDeleteTopic: Boolean = true,
                          port: Int = RandomPort,
                          interBrokerSecurityProtocol: Option[SecurityProtocol] = None,
                          trustStoreFile: Option[File] = None,
@@ -212,7 +254,10 @@ object TestUtils extends Logging {
                          enableSaslSsl: Boolean = false,
                          saslSslPort: Int = RandomPort,
                          rack: Option[String] = None,
-                         logDirCount: Int = 1): Properties = {
+                         logDirCount: Int = 1,
+                         enableToken: Boolean = false,
+                         numPartitions: Int = 1,
+                         defaultReplicationFactor: Short = 1): Properties = {
     def shouldEnable(protocol: SecurityProtocol) = interBrokerSecurityProtocol.fold(false)(_ == protocol)
 
     val protocolAndPorts = ArrayBuffer[(SecurityProtocol, Int)]()
@@ -269,6 +314,12 @@ object TestUtils extends Logging {
       props.put(KafkaConfig.InterBrokerSecurityProtocolProp, protocol.name)
     }
 
+    if (enableToken)
+      props.put(KafkaConfig.DelegationTokenSecretKeyProp, "secretkey")
+
+    props.put(KafkaConfig.NumPartitionsProp, numPartitions.toString)
+    props.put(KafkaConfig.DefaultReplicationFactorProp, defaultReplicationFactor.toString)
+
     props
   }
 
@@ -277,18 +328,29 @@ object TestUtils extends Logging {
    * Wait until the leader is elected and the metadata is propagated to all brokers.
    * Return the leader for each partition.
    */
-  def createTopic(zkUtils: ZkUtils,
+  def createTopic(zkClient: KafkaZkClient,
                   topic: String,
                   numPartitions: Int = 1,
                   replicationFactor: Int = 1,
                   servers: Seq[KafkaServer],
                   topicConfig: Properties = new Properties): scala.collection.immutable.Map[Int, Int] = {
+    val adminZkClient = new AdminZkClient(zkClient)
     // create topic
-    AdminUtils.createTopic(zkUtils, topic, numPartitions, replicationFactor, topicConfig)
+    waitUntilTrue( () => {
+      var hasSessionExpirationException = false
+      try {
+        adminZkClient.createTopic(topic, numPartitions, replicationFactor, topicConfig)
+      } catch {
+        case _: SessionExpiredException => hasSessionExpirationException = true
+        case e: Throwable => throw e // let other exceptions propagate
+      }
+      !hasSessionExpirationException},
+      s"Can't create topic $topic")
+
     // wait until the update metadata request for new topic reaches all servers
     (0 until numPartitions).map { i =>
       TestUtils.waitUntilMetadataIsPropagated(servers, topic, i)
-      i -> TestUtils.waitUntilLeaderIsElectedOrChanged(zkUtils, topic, i)
+      i -> TestUtils.waitUntilLeaderIsElectedOrChanged(zkClient, topic, i)
     }.toMap
   }
 
@@ -297,14 +359,40 @@ object TestUtils extends Logging {
    * Wait until the leader is elected and the metadata is propagated to all brokers.
    * Return the leader for each partition.
    */
-  def createTopic(zkUtils: ZkUtils, topic: String, partitionReplicaAssignment: collection.Map[Int, Seq[Int]],
+  def createTopic(zkClient: KafkaZkClient,
+                  topic: String,
+                  partitionReplicaAssignment: collection.Map[Int, Seq[Int]],
                   servers: Seq[KafkaServer]): scala.collection.immutable.Map[Int, Int] = {
+    createTopic(zkClient, topic, partitionReplicaAssignment, servers, new Properties())
+  }
+
+  /**
+   * Create a topic in ZooKeeper using a customized replica assignment.
+   * Wait until the leader is elected and the metadata is propagated to all brokers.
+   * Return the leader for each partition.
+   */
+  def createTopic(zkClient: KafkaZkClient,
+                  topic: String,
+                  partitionReplicaAssignment: collection.Map[Int, Seq[Int]],
+                  servers: Seq[KafkaServer],
+                  topicConfig: Properties): scala.collection.immutable.Map[Int, Int] = {
+    val adminZkClient = new AdminZkClient(zkClient)
     // create topic
-    AdminUtils.createOrUpdateTopicPartitionAssignmentPathInZK(zkUtils, topic, partitionReplicaAssignment)
+    waitUntilTrue( () => {
+      var hasSessionExpirationException = false
+      try {
+        adminZkClient.createTopicWithAssignment(topic, topicConfig, partitionReplicaAssignment)
+      } catch {
+        case _: SessionExpiredException => hasSessionExpirationException = true
+        case e: Throwable => throw e // let other exceptions propagate
+      }
+      !hasSessionExpirationException},
+      s"Can't create topic $topic")
+
     // wait until the update metadata request for new topic reaches all servers
-    partitionReplicaAssignment.keySet.map { case i =>
+    partitionReplicaAssignment.keySet.map { i =>
       TestUtils.waitUntilMetadataIsPropagated(servers, topic, i)
-      i -> TestUtils.waitUntilLeaderIsElectedOrChanged(zkUtils, topic, i)
+      i -> TestUtils.waitUntilLeaderIsElectedOrChanged(zkClient, topic, i)
     }.toMap
   }
 
@@ -312,40 +400,15 @@ object TestUtils extends Logging {
     * Create the consumer offsets/group metadata topic and wait until the leader is elected and metadata is propagated
     * to all brokers.
     */
-  def createOffsetsTopic(zkUtils: ZkUtils, servers: Seq[KafkaServer]): Unit = {
+  def createOffsetsTopic(zkClient: KafkaZkClient, servers: Seq[KafkaServer]): Unit = {
     val server = servers.head
-    createTopic(zkUtils, Topic.GROUP_METADATA_TOPIC_NAME,
+    createTopic(zkClient, Topic.GROUP_METADATA_TOPIC_NAME,
       server.config.getInt(KafkaConfig.OffsetsTopicPartitionsProp),
       server.config.getShort(KafkaConfig.OffsetsTopicReplicationFactorProp).toInt,
       servers,
       server.groupCoordinator.offsetsTopicConfigs)
   }
 
-  /**
-   * Create a test config for a consumer
-   */
-  def createConsumerProperties(zkConnect: String, groupId: String, consumerId: String,
-                               consumerTimeout: Long = -1): Properties = {
-    val props = new Properties
-    props.put("zookeeper.connect", zkConnect)
-    props.put("group.id", groupId)
-    props.put("consumer.id", consumerId)
-    props.put("consumer.timeout.ms", consumerTimeout.toString)
-    props.put("zookeeper.session.timeout.ms", "6000")
-    props.put("zookeeper.sync.time.ms", "200")
-    props.put("auto.commit.interval.ms", "1000")
-    props.put("rebalance.max.retries", "4")
-    props.put("auto.offset.reset", "smallest")
-    props.put("num.consumer.fetchers", "2")
-
-    props
-  }
-
-  /**
-   * 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.
    */
@@ -369,10 +432,11 @@ object TestUtils extends Logging {
               producerId: Long = RecordBatch.NO_PRODUCER_ID,
               producerEpoch: Short = RecordBatch.NO_PRODUCER_EPOCH,
               sequence: Int = RecordBatch.NO_SEQUENCE,
-              baseOffset: Long = 0L): MemoryRecords = {
+              baseOffset: Long = 0L,
+              partitionLeaderEpoch: Int = RecordBatch.NO_PARTITION_LEADER_EPOCH): MemoryRecords = {
     val buf = ByteBuffer.allocate(DefaultRecordBatch.sizeInBytes(records.asJava))
     val builder = MemoryRecords.builder(buf, magicValue, codec, TimestampType.CREATE_TIME, baseOffset,
-      System.currentTimeMillis, producerId, producerEpoch, sequence)
+      System.currentTimeMillis, producerId, producerEpoch, sequence, false, partitionLeaderEpoch)
     records.foreach(builder.append)
     builder.build()
   }
@@ -395,7 +459,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))
@@ -405,18 +469,18 @@ 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
-      assertEquals(expected.next, actual.next)
+      assertEquals(expected.next(), actual.next())
     }
 
     // check if the expected iterator is longer
     if (expected.hasNext) {
       var length1 = length
       while (expected.hasNext) {
-        expected.next
+        expected.next()
         length1 += 1
       }
       assertFalse("Iterators have uneven length-- first has more: "+length1 + " > " + length, true)
@@ -426,7 +490,7 @@ object TestUtils extends Logging {
     if (actual.hasNext) {
       var length2 = length
       while (actual.hasNext) {
-        actual.next
+        actual.next()
         length2 += 1
       }
       assertFalse("Iterators have uneven length-- second has more: "+length2 + " > " + length, true)
@@ -437,11 +501,11 @@ 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
-      s1.next
+      n += 1
+      s1.next()
     }
     assertEquals(expectedLength, n)
   }
@@ -450,7 +514,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)
@@ -466,7 +530,7 @@ object TestUtils extends Logging {
         while (true) {
           if (cur == null) {
             if (topIterator.hasNext)
-              cur = topIterator.next
+              cur = topIterator.next()
             else
               return false
           }
@@ -474,11 +538,11 @@ object TestUtils extends Logging {
             return true
           cur = null
         }
-        // should never reach her
+        // should never reach here
         throw new RuntimeException("should not reach here")
       }
 
-      def next() : T = cur.next
+      def next() : T = cur.next()
     }
   }
 
@@ -497,36 +561,16 @@ object TestUtils extends Logging {
     builder.toString
   }
 
-  /**
-   * Create a producer with a few pre-configured properties.
-   * If certain properties need to be overridden, they can be provided in producerProps.
-   */
-  @deprecated("This method has been deprecated and it will be removed in a future release.", "0.10.0.0")
-  def createProducer[K, V](brokerList: String,
-                           encoder: String = classOf[DefaultEncoder].getName,
-                           keyEncoder: String = classOf[DefaultEncoder].getName,
-                           partitioner: String = classOf[DefaultPartitioner].getName,
-                           producerProps: Properties = null): Producer[K, V] = {
-    val props: Properties = getProducerConfig(brokerList)
-
-    //override any explicitly specified properties
-    if (producerProps != null)
-      props ++= producerProps
-
-    props.put("serializer.class", encoder)
-    props.put("key.serializer.class", keyEncoder)
-    props.put("partitioner.class", partitioner)
-    new Producer[K, V](new kafka.producer.ProducerConfig(props))
-  }
-
-  private def securityConfigs(mode: Mode,
-                              securityProtocol: SecurityProtocol,
-                              trustStoreFile: Option[File],
-                              certAlias: String,
-                              saslProperties: Option[Properties]): Properties = {
+  def securityConfigs(mode: Mode,
+                      securityProtocol: SecurityProtocol,
+                      trustStoreFile: Option[File],
+                      certAlias: String,
+                      certCn: String,
+                      saslProperties: Option[Properties],
+                      tlsProtocol: String = TestSslUtils.DEFAULT_TLS_PROTOCOL_FOR_TESTS): Properties = {
     val props = new Properties
     if (usesSslTransportLayer(securityProtocol))
-      props ++= sslConfigs(mode, securityProtocol == SecurityProtocol.SSL, trustStoreFile, certAlias)
+      props ++= sslConfigs(mode, securityProtocol == SecurityProtocol.SSL, trustStoreFile, certAlias, certCn, tlsProtocol)
 
     if (usesSaslAuthentication(securityProtocol))
       props ++= JaasTestUtils.saslConfigs(saslProperties)
@@ -534,54 +578,43 @@ object TestUtils extends Logging {
     props
   }
 
-  def producerSecurityConfigs(securityProtocol: SecurityProtocol, trustStoreFile: Option[File], saslProperties: Option[Properties]): Properties =
-    securityConfigs(Mode.CLIENT, securityProtocol, trustStoreFile, "producer", saslProperties)
+  def producerSecurityConfigs(securityProtocol: SecurityProtocol,
+                              trustStoreFile: Option[File],
+                              saslProperties: Option[Properties]): Properties =
+    securityConfigs(Mode.CLIENT, securityProtocol, trustStoreFile, "producer", SslCertificateCn, saslProperties)
 
   /**
    * Create a (new) producer with a few pre-configured properties.
    */
-  def createNewProducer[K, V](brokerList: String,
-                        acks: Int = -1,
-                        maxBlockMs: Long = 60 * 1000L,
-                        bufferSize: Long = 1024L * 1024L,
-                        retries: Int = 0,
-                        lingerMs: Long = 0,
-                        requestTimeoutMs: Long = 10 * 1024L,
-                        securityProtocol: SecurityProtocol = SecurityProtocol.PLAINTEXT,
-                        trustStoreFile: Option[File] = None,
-                        saslProperties: Option[Properties] = None,
-                        keySerializer: Serializer[K] = new ByteArraySerializer,
-                        valueSerializer: Serializer[V] = new ByteArraySerializer,
-                        props: Option[Properties] = None): KafkaProducer[K, V] = {
-
-    val producerProps = props.getOrElse(new Properties)
+  def createProducer[K, V](brokerList: String,
+                           acks: Int = -1,
+                           maxBlockMs: Long = 60 * 1000L,
+                           bufferSize: Long = 1024L * 1024L,
+                           retries: Int = Int.MaxValue,
+                           deliveryTimeoutMs: Int = 30 * 1000,
+                           lingerMs: Int = 0,
+                           batchSize: Int = 16384,
+                           compressionType: String = "none",
+                           requestTimeoutMs: Int = 20 * 1000,
+                           securityProtocol: SecurityProtocol = SecurityProtocol.PLAINTEXT,
+                           trustStoreFile: Option[File] = None,
+                           saslProperties: Option[Properties] = None,
+                           keySerializer: Serializer[K] = new ByteArraySerializer,
+                           valueSerializer: Serializer[V] = new ByteArraySerializer,
+                           enableIdempotence: Boolean = false): KafkaProducer[K, V] = {
+    val producerProps = new Properties
     producerProps.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList)
     producerProps.put(ProducerConfig.ACKS_CONFIG, acks.toString)
     producerProps.put(ProducerConfig.MAX_BLOCK_MS_CONFIG, maxBlockMs.toString)
     producerProps.put(ProducerConfig.BUFFER_MEMORY_CONFIG, bufferSize.toString)
     producerProps.put(ProducerConfig.RETRIES_CONFIG, retries.toString)
+    producerProps.put(ProducerConfig.DELIVERY_TIMEOUT_MS_CONFIG, deliveryTimeoutMs.toString)
     producerProps.put(ProducerConfig.REQUEST_TIMEOUT_MS_CONFIG, requestTimeoutMs.toString)
-
-    /* Only use these if not already set */
-    val defaultProps = Map(
-      ProducerConfig.RETRY_BACKOFF_MS_CONFIG -> "100",
-      ProducerConfig.RECONNECT_BACKOFF_MS_CONFIG -> "200",
-      ProducerConfig.LINGER_MS_CONFIG -> lingerMs.toString
-    )
-
-    defaultProps.foreach { case (key, value) =>
-      if (!producerProps.containsKey(key)) producerProps.put(key, value)
-    }
-
-    /*
-     * It uses CommonClientConfigs.SECURITY_PROTOCOL_CONFIG to determine whether
-     * securityConfigs has been invoked already. For example, we need to
-     * invoke it before this call in IntegrationTestHarness, otherwise the
-     * SSL client auth fails.
-     */
-    if (!producerProps.containsKey(CommonClientConfigs.SECURITY_PROTOCOL_CONFIG))
-      producerProps ++= producerSecurityConfigs(securityProtocol, trustStoreFile, saslProperties)
-
+    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.ENABLE_IDEMPOTENCE_CONFIG, enableIdempotence.toString)
+    producerProps ++= producerSecurityConfigs(securityProtocol, trustStoreFile, saslProperties)
     new KafkaProducer[K, V](producerProps, keySerializer, valueSerializer)
   }
 
@@ -596,119 +629,54 @@ object TestUtils extends Logging {
   }
 
   def consumerSecurityConfigs(securityProtocol: SecurityProtocol, trustStoreFile: Option[File], saslProperties: Option[Properties]): Properties =
-    securityConfigs(Mode.CLIENT, securityProtocol, trustStoreFile, "consumer", saslProperties)
+    securityConfigs(Mode.CLIENT, securityProtocol, trustStoreFile, "consumer", SslCertificateCn, saslProperties)
 
   def adminClientSecurityConfigs(securityProtocol: SecurityProtocol, trustStoreFile: Option[File], saslProperties: Option[Properties]): Properties =
-    securityConfigs(Mode.CLIENT, securityProtocol, trustStoreFile, "admin-client", saslProperties)
+    securityConfigs(Mode.CLIENT, securityProtocol, trustStoreFile, "admin-client", SslCertificateCn, saslProperties)
 
   /**
-   * Create a new consumer with a few pre-configured properties.
+   * Create a consumer with a few pre-configured properties.
    */
-  def createNewConsumer(brokerList: String,
-                        groupId: String = "group",
-                        autoOffsetReset: String = "earliest",
-                        partitionFetchSize: Long = 4096L,
-                        partitionAssignmentStrategy: String = classOf[RangeAssignor].getName,
-                        sessionTimeout: Int = 30000,
-                        securityProtocol: SecurityProtocol,
-                        trustStoreFile: Option[File] = None,
-                        saslProperties: Option[Properties] = None,
-                        props: Option[Properties] = None) : KafkaConsumer[Array[Byte],Array[Byte]] = {
-    import org.apache.kafka.clients.consumer.ConsumerConfig
-
-    val consumerProps = props.getOrElse(new Properties())
+  def createConsumer[K, V](brokerList: String,
+                           groupId: String = "group",
+                           autoOffsetReset: String = "earliest",
+                           enableAutoCommit: Boolean = true,
+                           readCommitted: Boolean = false,
+                           maxPollRecords: Int = 500,
+                           securityProtocol: SecurityProtocol = SecurityProtocol.PLAINTEXT,
+                           trustStoreFile: Option[File] = None,
+                           saslProperties: Option[Properties] = None,
+                           keyDeserializer: Deserializer[K] = new ByteArrayDeserializer,
+                           valueDeserializer: Deserializer[V] = new ByteArrayDeserializer): KafkaConsumer[K, V] = {
+    val consumerProps = new Properties
     consumerProps.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList)
     consumerProps.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, autoOffsetReset)
-    consumerProps.put(ConsumerConfig.MAX_PARTITION_FETCH_BYTES_CONFIG, partitionFetchSize.toString)
-
-    val defaultProps = Map(
-      ConsumerConfig.RETRY_BACKOFF_MS_CONFIG -> "100",
-      ConsumerConfig.RECONNECT_BACKOFF_MS_CONFIG -> "200",
-      ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG -> "org.apache.kafka.common.serialization.ByteArrayDeserializer",
-      ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG -> "org.apache.kafka.common.serialization.ByteArrayDeserializer",
-      ConsumerConfig.PARTITION_ASSIGNMENT_STRATEGY_CONFIG -> partitionAssignmentStrategy,
-      ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG -> sessionTimeout.toString,
-      ConsumerConfig.GROUP_ID_CONFIG -> groupId)
-
-    defaultProps.foreach { case (key, value) =>
-      if (!consumerProps.containsKey(key)) consumerProps.put(key, value)
-    }
-
-    /*
-     * It uses CommonClientConfigs.SECURITY_PROTOCOL_CONFIG to determine whether
-     * securityConfigs has been invoked already. For example, we need to
-     * invoke it before this call in IntegrationTestHarness, otherwise the
-     * SSL client auth fails.
-     */
-    if(!consumerProps.containsKey(CommonClientConfigs.SECURITY_PROTOCOL_CONFIG))
-      consumerProps ++= consumerSecurityConfigs(securityProtocol, trustStoreFile, saslProperties)
-
-    new KafkaConsumer[Array[Byte],Array[Byte]](consumerProps)
-  }
-
-  /**
-   * Create a default producer config properties map with the given metadata broker list
-   */
-  def getProducerConfig(brokerList: String): Properties = {
-    val props = new Properties()
-    props.put("metadata.broker.list", brokerList)
-    props.put("message.send.max.retries", "5")
-    props.put("retry.backoff.ms", "1000")
-    props.put("request.timeout.ms", "2000")
-    props.put("request.required.acks", "-1")
-    props.put("send.buffer.bytes", "65536")
-
-    props
-  }
-
-  @deprecated("This method has been deprecated and will be removed in a future release", "0.11.0.0")
-  def getSyncProducerConfig(port: Int): Properties = {
-    val props = new Properties()
-    props.put("host", "localhost")
-    props.put("port", port.toString)
-    props.put("request.timeout.ms", "10000")
-    props.put("request.required.acks", "1")
-    props.put("serializer.class", classOf[StringEncoder].getName)
-    props
-  }
-
-
-  @deprecated("This method has been deprecated and will be removed in a future release.", "0.11.0.0")
-  def updateConsumerOffset(config : ConsumerConfig, path : String, offset : Long) = {
-    val zkUtils = ZkUtils(config.zkConnect, config.zkSessionTimeoutMs, config.zkConnectionTimeoutMs, false)
-    zkUtils.updatePersistentPath(path, offset.toString)
-    zkUtils.close()
-
-  }
-
-  def getMessageIterator(iter: Iterator[MessageAndOffset]): Iterator[Message] = {
-    new IteratorTemplate[Message] {
-      override def makeNext(): Message = {
-        if (iter.hasNext)
-          iter.next.message
-        else
-          allDone()
-      }
-    }
+    consumerProps.put(ConsumerConfig.GROUP_ID_CONFIG, groupId)
+    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 ++= consumerSecurityConfigs(securityProtocol, trustStoreFile, saslProperties)
+    new KafkaConsumer[K, V](consumerProps, keyDeserializer, valueDeserializer)
   }
 
-  def createBrokersInZk(zkUtils: ZkUtils, ids: Seq[Int]): Seq[Broker] =
-    createBrokersInZk(ids.map(kafka.admin.BrokerMetadata(_, None)), zkUtils)
+  def createBrokersInZk(zkClient: KafkaZkClient, ids: Seq[Int]): Seq[Broker] =
+    createBrokersInZk(ids.map(kafka.admin.BrokerMetadata(_, None)), zkClient)
 
-  def createBrokersInZk(brokerMetadatas: Seq[kafka.admin.BrokerMetadata], zkUtils: ZkUtils): Seq[Broker] = {
+  def createBrokersInZk(brokerMetadatas: Seq[kafka.admin.BrokerMetadata], zkClient: KafkaZkClient): Seq[Broker] = {
+    zkClient.makeSurePersistentPathExists(BrokerIdsZNode.path)
     val brokers = brokerMetadatas.map { b =>
       val protocol = SecurityProtocol.PLAINTEXT
       val listenerName = ListenerName.forSecurityProtocol(protocol)
       Broker(b.id, Seq(EndPoint("localhost", 6667, listenerName, protocol)), b.rack)
     }
-    brokers.foreach(b => zkUtils.registerBrokerInZk(b.id, "localhost", 6667, b.endPoints, jmxPort = -1,
-      rack = b.rack, ApiVersion.latestVersion))
+    brokers.foreach(b => zkClient.registerBroker(BrokerInfo(Broker(b.id, b.endPoints, rack = b.rack),
+      ApiVersion.latestVersion, jmxPort = -1)))
     brokers
   }
 
-  def deleteBrokersInZk(zkUtils: ZkUtils, ids: Seq[Int]): Seq[Broker] = {
+  def deleteBrokersInZk(zkClient: KafkaZkClient, ids: Seq[Int]): Seq[Broker] = {
     val brokers = ids.map(createBroker(_, "localhost", 6667, SecurityProtocol.PLAINTEXT))
-    brokers.foreach(b => zkUtils.deletePath(ZkUtils.BrokerIdsPath + "/" + b))
+    ids.foreach(b => zkClient.deletePath(BrokerIdsZNode.path + "/" + b))
     brokers
   }
 
@@ -719,52 +687,18 @@ object TestUtils extends Logging {
     buffer
   }
 
-  /**
-   * Create a wired format request based on simple basic information
-   */
-  @deprecated("This method has been deprecated and it will be removed in a future release", "0.10.0.0")
-  def produceRequest(topic: String,
-                     partition: Int,
-                     message: ByteBufferMessageSet,
-                     acks: Int,
-                     timeout: Int,
-                     correlationId: Int = 0,
-                     clientId: String): ProducerRequest = {
-    produceRequestWithAcks(Seq(topic), Seq(partition), message, acks, timeout, correlationId, clientId)
-  }
-
-  @deprecated("This method has been deprecated and it will be removed in a future release", "0.10.0.0")
-  def produceRequestWithAcks(topics: Seq[String],
-                             partitions: Seq[Int],
-                             message: ByteBufferMessageSet,
-                             acks: Int,
-                             timeout: Int,
-                             correlationId: Int = 0,
-                             clientId: String): ProducerRequest = {
-    val data = topics.flatMap(topic =>
-      partitions.map(partition => (TopicAndPartition(topic,  partition), message))
-    )
-    new ProducerRequest(correlationId, clientId, acks.toShort, timeout, collection.mutable.Map(data:_*))
-  }
-
-  def makeLeaderForPartition(zkUtils: ZkUtils,
+  def makeLeaderForPartition(zkClient: KafkaZkClient,
                              topic: String,
                              leaderPerPartitionMap: scala.collection.immutable.Map[Int, Int],
-                             controllerEpoch: Int) {
-    leaderPerPartitionMap.foreach { case (partition, leader) =>
-      try {
-        val newLeaderAndIsr = zkUtils.getLeaderAndIsrForPartition(topic, partition)
-          .map(_.newLeader(leader))
-          .getOrElse(LeaderAndIsr(leader, List(leader)))
-
-        zkUtils.updatePersistentPath(
-          getTopicPartitionLeaderAndIsrPath(topic, partition),
-          zkUtils.leaderAndIsrZkData(newLeaderAndIsr, controllerEpoch)
-        )
-      } catch {
-        case oe: Throwable => error(s"Error while electing leader for partition [$topic,$partition]", oe)
-      }
+                             controllerEpoch: Int): Unit = {
+    val newLeaderIsrAndControllerEpochs = leaderPerPartitionMap.map { case (partition, leader) =>
+      val topicPartition = new TopicPartition(topic, partition)
+      val newLeaderAndIsr = zkClient.getTopicPartitionState(topicPartition)
+        .map(_.leaderAndIsr.newLeader(leader))
+        .getOrElse(LeaderAndIsr(leader, List(leader)))
+      topicPartition -> LeaderIsrAndControllerEpoch(newLeaderAndIsr, controllerEpoch)
     }
+    zkClient.setTopicPartitionStatesRaw(newLeaderIsrAndControllerEpochs, ZkVersion.MatchAnyVersion)
   }
 
   /**
@@ -776,7 +710,7 @@ object TestUtils extends Logging {
    *         LeaderDuringDelete).
    * @throws AssertionError if the expected condition is not true within the timeout.
    */
-  def waitUntilLeaderIsElectedOrChanged(zkUtils: ZkUtils, topic: String, partition: Int, timeoutMs: Long = 30000L,
+  def waitUntilLeaderIsElectedOrChanged(zkClient: KafkaZkClient, topic: String, partition: Int, timeoutMs: Long = 30000L,
                                         oldLeaderOpt: Option[Int] = None, newLeaderOpt: Option[Int] = None): Int = {
     require(!(oldLeaderOpt.isDefined && newLeaderOpt.isDefined), "Can't define both the old and the new leader")
     val startTime = System.currentTimeMillis()
@@ -789,7 +723,7 @@ object TestUtils extends Logging {
     var electedOrChangedLeader: Option[Int] = None
     while (electedOrChangedLeader.isEmpty && System.currentTimeMillis() < startTime + timeoutMs) {
       // check if leader is elected
-      leader = zkUtils.getLeaderForPartition(topic, partition)
+      leader = zkClient.getLeaderForPartition(topicPartition)
       leader match {
         case Some(l) => (newLeaderOpt, oldLeaderOpt) match {
           case (Some(newLeader), _) if newLeader == l =>
@@ -827,7 +761,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) {
@@ -848,19 +782,73 @@ object TestUtils extends Logging {
     }
   }
 
+  def pollUntilTrue(consumer: Consumer[_, _],
+                    action: () => Boolean,
+                    msg: => String,
+                    waitTimeMs: Long = JTestUtils.DEFAULT_MAX_WAIT_MS): Unit = {
+    waitUntilTrue(() => {
+      consumer.poll(Duration.ofMillis(100))
+      action()
+    }, msg = msg, pause = 0L, waitTimeMs = waitTimeMs)
+  }
+
+  def pollRecordsUntilTrue[K, V](consumer: Consumer[K, V],
+                                 action: ConsumerRecords[K, V] => Boolean,
+                                 msg: => String,
+                                 waitTimeMs: Long = JTestUtils.DEFAULT_MAX_WAIT_MS): Unit = {
+    waitUntilTrue(() => {
+      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 until the given condition is true or throw an exception if the given wait time elapses.
+   * 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.
+    *
+    * @param condition condition to check
+    * @param msg error message
+    * @param waitTimeMs maximum time to wait and retest the condition before failing the test
+    * @param pause delay between condition checks
+    */
   def waitUntilTrue(condition: () => Boolean, msg: => String,
-                    waitTime: Long = JTestUtils.DEFAULT_MAX_WAIT_MS, pause: Long = 100L): Unit = {
+                    waitTimeMs: Long = JTestUtils.DEFAULT_MAX_WAIT_MS, pause: Long = 100L): Unit = {
     val startTime = System.currentTimeMillis()
     while (true) {
       if (condition())
         return
-      if (System.currentTimeMillis() > startTime + waitTime)
+      if (System.currentTimeMillis() > startTime + waitTimeMs)
         fail(msg)
-      Thread.sleep(waitTime.min(pause))
+      Thread.sleep(waitTimeMs.min(pause))
     }
+
     // should never hit here
     throw new RuntimeException("unexpected error")
   }
@@ -873,7 +861,7 @@ object TestUtils extends Logging {
     * This method is useful in cases where `waitUntilTrue` makes it awkward to provide good error messages.
     */
   def computeUntilTrue[T](compute: => T, waitTime: Long = JTestUtils.DEFAULT_MAX_WAIT_MS, pause: Long = 100L)(
-                    predicate: T => Boolean): (T, Boolean) = {
+                          predicate: T => Boolean): (T, Boolean) = {
     val startTime = System.currentTimeMillis()
     while (true) {
       val result = compute
@@ -888,15 +876,29 @@ 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 createRequestByteBuffer(request: RequestOrResponse): ByteBuffer = {
-    val byteBuffer = ByteBuffer.allocate(request.sizeInBytes + 2)
-    byteBuffer.putShort(request.requestId.get)
-    request.writeTo(byteBuffer)
-    byteBuffer.rewind()
-    byteBuffer
+  def findLeaderEpoch(brokerId: Int,
+                      topicPartition: TopicPartition,
+                      servers: Iterable[KafkaServer]): Int = {
+    val leaderServer = servers.find(_.config.brokerId == brokerId)
+    val leaderPartition = leaderServer.flatMap(_.replicaManager.nonOfflinePartition(topicPartition))
+      .getOrElse(fail(s"Failed to find expected replica on broker $brokerId"))
+    leaderPartition.getLeaderEpoch
+  }
+
+  def findFollowerId(topicPartition: TopicPartition,
+                     servers: Iterable[KafkaServer]): Int = {
+    val followerOpt = servers.find { server =>
+      server.replicaManager.nonOfflinePartition(topicPartition) match {
+        case Some(partition) => !partition.leaderReplicaIdOpt.contains(server.config.brokerId)
+        case None => false
+      }
+    }
+    followerOpt
+      .map(_.config.brokerId)
+      .getOrElse(fail(s"Unable to locate follower for $topicPartition"))
   }
 
   /**
@@ -908,8 +910,8 @@ 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 =>
-      expectedBrokerIds == server.apis.metadataCache.getAliveBrokers.map(_.id).toSet
+    waitUntilTrue(() => servers.forall(server =>
+      expectedBrokerIds == server.dataPlaneRequestProcessor.metadataCache.getAliveBrokers.map(_.id).toSet
     ), "Timed out waiting for broker metadata to propagate to all servers", timeout)
   }
 
@@ -926,43 +928,59 @@ 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.apis.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),
-      waitTime = timeout)
+      waitTimeMs = timeout)
 
     leader
   }
 
-  def waitUntilControllerElected(zkUtils: ZkUtils, timeout: Long = JTestUtils.DEFAULT_MAX_WAIT_MS): Int = {
-    var controllerIdTry: Try[Int] = null
-    TestUtils.waitUntilTrue(() => {
-      controllerIdTry = Try { zkUtils.getController() }
-      controllerIdTry.isSuccess
-    }, s"Controller not elected after $timeout ms", waitTime = timeout)
-    controllerIdTry.get
-  }
-
-  def waitUntilLeaderIsKnown(servers: Seq[KafkaServer], topic: String, partition: Int,
-                             timeout: Long = JTestUtils.DEFAULT_MAX_WAIT_MS): Unit = {
-    val tp = new TopicPartition(topic, partition)
-    TestUtils.waitUntilTrue(() =>
-      servers.exists { server =>
-        server.replicaManager.getPartition(tp).exists(_.leaderReplicaIfLocal.isDefined)
-      }, s"Partition $tp leaders not made yet after $timeout ms", waitTime = timeout
-    )
+  def waitUntilControllerElected(zkClient: KafkaZkClient, timeout: Long = JTestUtils.DEFAULT_MAX_WAIT_MS): Int = {
+    val (controllerId, _) = TestUtils.computeUntilTrue(zkClient.getControllerId, waitTime = timeout)(_.isDefined)
+    controllerId.getOrElse(fail(s"Controller not elected after $timeout ms"))
+  }
+
+  def awaitLeaderChange(servers: Seq[KafkaServer],
+                        tp: TopicPartition,
+                        oldLeader: Int,
+                        timeout: Long = JTestUtils.DEFAULT_MAX_WAIT_MS): Int = {
+    def newLeaderExists: Option[Int] = {
+      servers.find { server =>
+        server.config.brokerId != oldLeader &&
+          server.replicaManager.nonOfflinePartition(tp).exists(_.leaderLogIfLocal.isDefined)
+      }.map(_.config.brokerId)
+    }
+
+    waitUntilTrue(() => newLeaderExists.isDefined,
+      s"Did not observe leader change for partition $tp after $timeout ms", waitTimeMs = timeout)
+
+    newLeaderExists.get
+  }
+
+  def waitUntilLeaderIsKnown(servers: Seq[KafkaServer],
+                             tp: TopicPartition,
+                             timeout: Long = JTestUtils.DEFAULT_MAX_WAIT_MS): Int = {
+    def leaderIfExists: Option[Int] = {
+      servers.find { server =>
+        server.replicaManager.nonOfflinePartition(tp).exists(_.leaderLogIfLocal.isDefined)
+      }.map(_.config.brokerId)
+    }
+
+    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)
@@ -970,52 +988,58 @@ object TestUtils extends Logging {
     file.close()
   }
 
-  def appendNonsenseToFile(fileName: File, size: Int) {
-    val file = new FileOutputStream(fileName, true)
-    for (_ <- 0 until size)
-      file.write(random.nextInt(255))
-    file.close()
+  def appendNonsenseToFile(file: File, size: Int): Unit = {
+    val outputStream = Files.newOutputStream(file.toPath(), StandardOpenOption.APPEND)
+    try {
+      for (_ <- 0 until size)
+        outputStream.write(random.nextInt(255))
+    } finally outputStream.close()
   }
 
-  def checkForPhantomInSyncReplicas(zkUtils: ZkUtils, topic: String, partitionToBeReassigned: Int, assignedReplicas: Seq[Int]) {
-    val inSyncReplicas = zkUtils.getInSyncReplicasForPartition(topic, partitionToBeReassigned)
+  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.toSet -- assignedReplicas.toSet
+    val phantomInSyncReplicas = inSyncReplicas.get.toSet -- assignedReplicas.toSet
     assertTrue("All in sync replicas %s must be in the assigned replica list %s".format(inSyncReplicas, assignedReplicas),
       phantomInSyncReplicas.isEmpty)
   }
 
-  def ensureNoUnderReplicatedPartitions(zkUtils: ZkUtils, topic: String, partitionToBeReassigned: Int, assignedReplicas: Seq[Int],
-                                                servers: Seq[KafkaServer]) {
-    TestUtils.waitUntilTrue(() => {
-        val inSyncReplicas = zkUtils.getInSyncReplicasForPartition(topic, partitionToBeReassigned)
-        inSyncReplicas.size == assignedReplicas.size
+  def ensureNoUnderReplicatedPartitions(zkClient: KafkaZkClient, topic: String, partitionToBeReassigned: Int, assignedReplicas: Seq[Int],
+                                                servers: Seq[KafkaServer]): Unit = {
+    val topicPartition = new TopicPartition(topic, partitionToBeReassigned)
+    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(() => {
-        leader = zkUtils.getLeaderForPartition(topic, partitionToBeReassigned)
+    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 checkIfReassignPartitionPathExists(zkUtils: ZkUtils): Boolean = {
-    zkUtils.pathExists(ZkUtils.ReassignPartitionsPath)
-  }
-
-  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)
     }
     assertEquals(0, threadCount)
   }
 
+  def allThreadStackTraces(): String = {
+    Thread.getAllStackTraces.asScala.map { case (thread, stackTrace) =>
+      thread.getName + "\n\t" + stackTrace.toList.map(_.toString).mkString("\n\t")
+    }.mkString("\n")
+  }
+
   /**
    * Create new LogManager instance with default configuration for testing
    */
@@ -1023,12 +1047,12 @@ object TestUtils extends Logging {
                        defaultConfig: LogConfig = LogConfig(),
                        cleanerConfig: CleanerConfig = CleanerConfig(enableCleaner = false),
                        time: MockTime = new MockTime()): LogManager = {
-    new LogManager(logDirs = logDirs,
+    new LogManager(logDirs = logDirs.map(_.getAbsoluteFile),
                    initialOfflineDirs = Array.empty[File],
                    topicConfigs = Map(),
-                   defaultConfig = defaultConfig,
+                   initialDefaultConfig = defaultConfig,
                    cleanerConfig = cleanerConfig,
-                   ioThreads = 4,
+                   recoveryThreadsPerDataDir = 4,
                    flushCheckMs = 1000L,
                    flushRecoveryOffsetCheckpointMs = 10000L,
                    flushStartOffsetCheckpointMs = 10000L,
@@ -1041,159 +1065,118 @@ object TestUtils extends Logging {
                    logDirFailureChannel = new LogDirFailureChannel(logDirs.size))
   }
 
-  @deprecated("This method has been deprecated and it will be removed in a future release.", "0.10.0.0")
-  def sendMessages(servers: Seq[KafkaServer],
-                   topic: String,
-                   numMessages: Int,
-                   partition: Int = -1,
-                   compression: CompressionCodec = NoCompressionCodec): List[String] = {
-    val header = "test-%d".format(partition)
-    val props = new Properties()
-    props.put("compression.codec", compression.codec.toString)
-    val ms = 0.until(numMessages).map(x => header + "-" + x)
-
-    // Specific Partition
-    if (partition >= 0) {
-      val producer: Producer[Int, String] =
-        createProducer(TestUtils.getBrokerListStrFromServers(servers),
-          encoder = classOf[StringEncoder].getName,
-          keyEncoder = classOf[IntEncoder].getName,
-          partitioner = classOf[FixedValuePartitioner].getName,
-          producerProps = props)
-
-      producer.send(ms.map(m => new KeyedMessage[Int, String](topic, partition, m)): _*)
-      debug("Sent %d messages for partition [%s,%d]".format(ms.size, topic, partition))
-      producer.close()
-      ms.toList
-    } else {
-      // Use topic as the key to determine partition
-      val producer: Producer[String, String] = createProducer(
-        TestUtils.getBrokerListStrFromServers(servers),
-        encoder = classOf[StringEncoder].getName,
-        keyEncoder = classOf[StringEncoder].getName,
-        partitioner = classOf[DefaultPartitioner].getName,
-        producerProps = props)
-      producer.send(ms.map(m => new KeyedMessage[String, String](topic, topic, m)): _*)
-      producer.close()
-      debug("Sent %d messages for topic [%s]".format(ms.size, topic))
-      ms.toList
+  class MockAlterIsrManager extends AlterIsrManager {
+    val isrUpdates: mutable.Queue[AlterIsrItem] = new mutable.Queue[AlterIsrItem]()
+
+    override def enqueue(alterIsrItem: AlterIsrItem): Boolean = {
+      isrUpdates += alterIsrItem
+      true
+    }
+
+    override def clearPending(topicPartition: TopicPartition): Unit = {
+      isrUpdates.clear()
     }
+
+    override def start(): Unit = { }
   }
 
-  def produceMessages(servers: Seq[KafkaServer],
-                      topic: String,
-                      numMessages: Int,
-                      acks: Int = -1,
-                      valueBytes: Int = -1): Seq[Array[Byte]] = {
-
-    val producer = createNewProducer(
-      TestUtils.getBrokerListStrFromServers(servers),
-      retries = 5,
-      requestTimeoutMs = 2000,
-      acks = acks
-    )
+  def createAlterIsrManager(): MockAlterIsrManager = {
+    new MockAlterIsrManager()
+  }
 
-    val values = (0 until numMessages).map(x => valueBytes match {
-      case -1 => s"test-$x".getBytes
-      case _ => new Array[Byte](valueBytes)
-    })
+  class MockIsrChangeListener extends IsrChangeListener {
+    val expands: AtomicInteger = new AtomicInteger(0)
+    val shrinks: AtomicInteger = new AtomicInteger(0)
+    val failures: AtomicInteger = new AtomicInteger(0)
 
-    val futures = values.map { value =>
-      producer.send(new ProducerRecord(topic, value))
-    }
-    futures.foreach(_.get)
-    producer.close()
+    override def markExpand(): Unit = expands.incrementAndGet()
 
-    debug(s"Sent ${values.size} messages for topic [$topic]")
+    override def markShrink(): Unit = shrinks.incrementAndGet()
 
-    values
+    override def markFailed(): Unit = failures.incrementAndGet()
+
+    def reset(): Unit = {
+      expands.set(0)
+      shrinks.set(0)
+      failures.set(0)
+    }
   }
 
-  def produceMessage(servers: Seq[KafkaServer], topic: String, message: String) {
-    val producer = createNewProducer(
-      TestUtils.getBrokerListStrFromServers(servers),
-      retries = 5,
-      requestTimeoutMs = 2000
-    )
-    producer.send(new ProducerRecord(topic, topic.getBytes, message.getBytes)).get
-    producer.close()
+  def createIsrChangeListener(): MockIsrChangeListener = {
+    new MockIsrChangeListener()
   }
 
-  /**
-   * Consume all messages (or a specific number of messages)
-   *
-   * @param topicMessageStreams the Topic Message Streams
-   * @param nMessagesPerThread an optional field to specify the exact number of messages to be returned.
-   *                           ConsumerTimeoutException will be thrown if there are no messages to be consumed.
-   *                           If not specified, then all available messages will be consumed, and no exception is thrown.
-   * @return the list of messages consumed.
-   */
-  @deprecated("This method has been deprecated and will be removed in a future release.", "0.11.0.0")
-  def getMessages(topicMessageStreams: Map[String, List[KafkaStream[String, String]]],
-                     nMessagesPerThread: Int = -1): List[String] = {
-
-    var messages: List[String] = Nil
-    val shouldGetAllMessages = nMessagesPerThread < 0
-    for (messageStreams <- topicMessageStreams.values) {
-      for (messageStream <- messageStreams) {
-        val iterator = messageStream.iterator()
-        try {
-          var i = 0
-          while ((shouldGetAllMessages && iterator.hasNext()) || (i < nMessagesPerThread)) {
-            assertTrue(iterator.hasNext)
-            val message = iterator.next().message // will throw a timeout exception if the message isn't there
-            messages ::= message
-            debug("received message: " + message)
-            i += 1
-          }
-        } catch {
-          case e: ConsumerTimeoutException =>
-            if (shouldGetAllMessages) {
-              // swallow the exception
-              debug("consumer timed out after receiving " + messages.length + " message(s).")
-            } else {
-              throw e
-            }
-        }
-      }
+  def produceMessages(servers: Seq[KafkaServer],
+                      records: Seq[ProducerRecord[Array[Byte], Array[Byte]]],
+                      acks: Int = -1): Unit = {
+    val producer = createProducer(TestUtils.getBrokerListStrFromServers(servers), acks = acks)
+    try {
+      val futures = records.map(producer.send)
+      futures.foreach(_.get)
+    } finally {
+      producer.close()
+    }
+
+    val topics = records.map(_.topic).distinct
+    debug(s"Sent ${records.size} messages for topics ${topics.mkString(",")}")
+  }
+
+  def generateAndProduceMessages(servers: Seq[KafkaServer],
+                                 topic: String,
+                                 numMessages: Int,
+                                 acks: Int = -1): Seq[String] = {
+    val values = (0 until numMessages).map(x =>  s"test-$x")
+    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
+  }
 
-    messages.reverse
+  def produceMessage(servers: Seq[KafkaServer], topic: String, message: String,
+                     deliveryTimeoutMs: Int = 30 * 1000, requestTimeoutMs: Int = 20 * 1000): Unit = {
+    val producer = createProducer(TestUtils.getBrokerListStrFromServers(servers),
+      deliveryTimeoutMs = deliveryTimeoutMs, requestTimeoutMs = requestTimeoutMs)
+    try {
+      producer.send(new ProducerRecord(topic, topic.getBytes, message.getBytes)).get
+    } finally {
+      producer.close()
+    }
   }
 
-  def verifyTopicDeletion(zkUtils: ZkUtils, 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(() => !zkUtils.isTopicMarkedForDeletion(topic),
-      "Admin path /admin/delete_topic/%s path not deleted even after a replica is restarted".format(topic))
-    TestUtils.waitUntilTrue(() => !zkUtils.pathExists(getTopicPath(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)
           }
@@ -1202,6 +1185,31 @@ object TestUtils extends Logging {
     ), "Failed to hard-delete the delete directory")
   }
 
+
+  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)
+  }
+
   /**
    * Translate the given buffer into a string
    *
@@ -1220,15 +1228,23 @@ object TestUtils extends Logging {
     copy
   }
 
-  def sslConfigs(mode: Mode, clientCert: Boolean, trustStoreFile: Option[File], certAlias: String): Properties = {
+  def sslConfigs(mode: Mode, clientCert: Boolean, trustStoreFile: Option[File], certAlias: String,
+                 certCn: String = SslCertificateCn,
+                 tlsProtocol: String = TestSslUtils.DEFAULT_TLS_PROTOCOL_FOR_TESTS): Properties = {
     val trustStore = trustStoreFile.getOrElse {
       throw new Exception("SSL enabled but no trustStoreFile provided")
     }
 
-    val sslConfigs = TestSslUtils.createSslConfig(clientCert, true, mode, trustStore, certAlias)
+    val sslConfigs = new TestSslUtils.SslConfigsBuilder(mode)
+      .useClientCert(clientCert)
+      .createNewTrustStore(trustStore)
+      .certAlias(certAlias)
+      .cn(certCn)
+      .tlsProtocol(tlsProtocol)
+      .build()
 
     val sslProps = new Properties()
-    sslConfigs.asScala.foreach { case (k, v) => sslProps.put(k, v) }
+    sslConfigs.forEach { (k, v) => sslProps.put(k, v) }
     sslProps
   }
 
@@ -1238,17 +1254,33 @@ 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[Acl], authorizer: Authorizer, resource: Resource) = {
-    TestUtils.waitUntilTrue(() => authorizer.getAcls(resource) == expected,
-      s"expected acls $expected but got ${authorizer.getAcls(resource)}", waitTime = JTestUtils.DEFAULT_MAX_WAIT_MS)
+  def waitAndVerifyAcls(expected: Set[AccessControlEntry],
+                        authorizer: JAuthorizer,
+                        resource: ResourcePattern,
+                        accessControlEntryFilter: AccessControlEntryFilter = AccessControlEntryFilter.ANY): Unit = {
+    val newLine = scala.util.Properties.lineSeparator
+
+    val filter = new AclBindingFilter(resource.toFilter, accessControlEntryFilter)
+    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)}")
+  }
+
+  @deprecated("Use org.apache.kafka.server.authorizer.Authorizer", "Since 2.5")
+  def waitAndVerifyAcls(expected: Set[Acl], authorizer: LegacyAuthorizer, resource: Resource): Unit = {
+    val newLine = scala.util.Properties.lineSeparator
+
+    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)}")
   }
 
   /**
@@ -1275,30 +1307,30 @@ object TestUtils extends Logging {
     }
   }
 
-  private def secureZkPaths(zkUtils: ZkUtils): Seq[String] = {
+  private def secureZkPaths(zkClient: KafkaZkClient): Seq[String] = {
     def subPaths(path: String): Seq[String] = {
-      if (zkUtils.pathExists(path))
-        path +: zkUtils.getChildren(path).map(c => path + "/" + c).flatMap(subPaths)
+      if (zkClient.pathExists(path))
+        path +: zkClient.getChildren(path).map(c => path + "/" + c).flatMap(subPaths)
       else
         Seq.empty
     }
-    val topLevelPaths = ZkUtils.SecureZkRootPaths ++ ZkUtils.SensitiveZkRootPaths
+    val topLevelPaths = ZkData.SecureRootPaths ++ ZkData.SensitiveRootPaths
     topLevelPaths.flatMap(subPaths)
   }
 
   /**
    * Verifies that all secure paths in ZK are created with the expected ACL.
    */
-  def verifySecureZkAcls(zkUtils: ZkUtils, usersWithAccess: Int) {
-    secureZkPaths(zkUtils).foreach(path => {
-      if (zkUtils.pathExists(path)) {
-        val sensitive = ZkUtils.sensitivePath(path)
+  def verifySecureZkAcls(zkClient: KafkaZkClient, usersWithAccess: Int): Unit = {
+    secureZkPaths(zkClient).foreach(path => {
+      if (zkClient.pathExists(path)) {
+        val sensitive = ZkData.sensitivePath(path)
         // usersWithAccess have ALL access to path. For paths that are
         // not sensitive, world has READ access.
         val aclCount = if (sensitive) usersWithAccess else usersWithAccess + 1
-        val acls = zkUtils.zkConnection.getAcl(path).getKey
+        val acls = zkClient.getAcl(path)
         assertEquals(s"Invalid ACLs for $path $acls", aclCount, acls.size)
-        acls.asScala.foreach(acl => isAclSecure(acl, sensitive))
+        acls.foreach(acl => isAclSecure(acl, sensitive))
       }
     })
   }
@@ -1307,12 +1339,12 @@ 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(zkUtils: ZkUtils) {
-    secureZkPaths(zkUtils).foreach(path => {
-      if (zkUtils.pathExists(path)) {
-        val acls = zkUtils.zkConnection.getAcl(path).getKey
+  def verifyUnsecureZkAcls(zkClient: KafkaZkClient): Unit = {
+    secureZkPaths(zkClient).foreach(path => {
+      if (zkClient.pathExists(path)) {
+        val acls = zkClient.getAcl(path)
         assertEquals(s"Invalid ACLs for $path $acls", 1, acls.size)
-        acls.asScala.foreach(isAclUnsecure)
+        acls.foreach(isAclUnsecure)
       }
     })
   }
@@ -1322,9 +1354,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")
     }
 
@@ -1354,31 +1386,44 @@ object TestUtils extends Logging {
       threadPool.shutdownNow()
     }
     assertTrue(s"$message failed with exception(s) $exceptions", exceptions.isEmpty)
-
   }
 
   def consumeTopicRecords[K, V](servers: Seq[KafkaServer],
                                 topic: String,
                                 numMessages: Int,
+                                groupId: String = "group",
                                 securityProtocol: SecurityProtocol = SecurityProtocol.PLAINTEXT,
                                 trustStoreFile: Option[File] = None,
                                 waitTime: Long = JTestUtils.DEFAULT_MAX_WAIT_MS): Seq[ConsumerRecord[Array[Byte], Array[Byte]]] = {
-    val consumer = createNewConsumer(TestUtils.getBrokerListStrFromServers(servers, securityProtocol),
-      securityProtocol = securityProtocol, trustStoreFile = trustStoreFile)
+    val consumer = createConsumer(TestUtils.getBrokerListStrFromServers(servers, securityProtocol),
+      groupId = groupId,
+      securityProtocol = securityProtocol,
+      trustStoreFile = trustStoreFile)
     try {
       consumer.subscribe(Collections.singleton(topic))
       consumeRecords(consumer, numMessages, waitTime)
     } finally consumer.close()
   }
 
-  def consumeRecords[K, V](consumer: KafkaConsumer[K, V], numMessages: Int,
-                           waitTime: Long = JTestUtils.DEFAULT_MAX_WAIT_MS): Seq[ConsumerRecord[K, V]] = {
+  def pollUntilAtLeastNumRecords[K, V](consumer: Consumer[K, V],
+                                       numRecords: Int,
+                                       waitTimeMs: Long = JTestUtils.DEFAULT_MAX_WAIT_MS): Seq[ConsumerRecord[K, V]] = {
     val records = new ArrayBuffer[ConsumerRecord[K, V]]()
-    waitUntilTrue(() => {
-      records ++= consumer.poll(50).asScala
-      records.size >= numMessages
-    }, s"Consumed ${records.size} records until timeout instead of the expected $numMessages records", waitTime)
-    assertEquals("Consumed more records than expected", numMessages, records.size)
+    def pollAction(polledRecords: ConsumerRecords[K, V]): Boolean = {
+      records ++= polledRecords.asScala
+      records.size >= numRecords
+    }
+    pollRecordsUntilTrue(consumer, pollAction,
+      waitTimeMs = waitTimeMs,
+      msg = s"Consumed ${records.size} records before timeout instead of the expected $numRecords records")
+    records
+  }
+
+  def consumeRecords[K, V](consumer: Consumer[K, V],
+                           numRecords: Int,
+                           waitTimeMs: Long = JTestUtils.DEFAULT_MAX_WAIT_MS): Seq[ConsumerRecord[K, V]] = {
+    val records = pollUntilAtLeastNumRecords(consumer, numRecords, waitTimeMs)
+    assertEquals("Consumed more records than expected", numRecords, records.size)
     records
   }
 
@@ -1390,33 +1435,44 @@ object TestUtils extends Logging {
     *
     * @return All the records consumed by the consumer within the specified duration.
     */
-  def consumeRecordsFor[K, V](consumer: KafkaConsumer[K, V], duration: Long): Seq[ConsumerRecord[K, V]] = {
+  def consumeRecordsFor[K, V](consumer: KafkaConsumer[K, V], duration: Long = JTestUtils.DEFAULT_MAX_WAIT_MS): Seq[ConsumerRecord[K, V]] = {
     val startTime = System.currentTimeMillis()
     val records = new ArrayBuffer[ConsumerRecord[K, V]]()
     waitUntilTrue(() => {
-      records ++= consumer.poll(50).asScala
+      records ++= consumer.poll(Duration.ofMillis(50)).asScala
       System.currentTimeMillis() - startTime > duration
     }, s"The timeout $duration was greater than the maximum wait time.")
     records
   }
 
-  def createTransactionalProducer(transactionalId: String, servers: Seq[KafkaServer], batchSize: Int = 16384,
-                                  transactionTimeoutMs: Long = 60000) = {
+  def createTransactionalProducer(transactionalId: String,
+                                  servers: Seq[KafkaServer],
+                                  batchSize: Int = 16384,
+                                  transactionTimeoutMs: Long = 60000,
+                                  maxBlockMs: Long = 60000,
+                                  deliveryTimeoutMs: Int = 120000,
+                                  requestTimeoutMs: Int = 30000,
+                                  maxInFlight: Int = 5): KafkaProducer[Array[Byte], Array[Byte]] = {
     val props = new Properties()
+    props.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, TestUtils.getBrokerListStrFromServers(servers))
+    props.put(ProducerConfig.ACKS_CONFIG, "all")
+    props.put(ProducerConfig.BATCH_SIZE_CONFIG, batchSize.toString)
     props.put(ProducerConfig.TRANSACTIONAL_ID_CONFIG, transactionalId)
-    props.put(ProducerConfig.MAX_IN_FLIGHT_REQUESTS_PER_CONNECTION, "5")
     props.put(ProducerConfig.ENABLE_IDEMPOTENCE_CONFIG, "true")
-    props.put(ProducerConfig.BATCH_SIZE_CONFIG, batchSize.toString)
     props.put(ProducerConfig.TRANSACTION_TIMEOUT_CONFIG, transactionTimeoutMs.toString)
-    TestUtils.createNewProducer(TestUtils.getBrokerListStrFromServers(servers), retries = Integer.MAX_VALUE, acks = -1, props = Some(props))
+    props.put(ProducerConfig.MAX_BLOCK_MS_CONFIG, maxBlockMs.toString)
+    props.put(ProducerConfig.DELIVERY_TIMEOUT_MS_CONFIG, deliveryTimeoutMs.toString)
+    props.put(ProducerConfig.REQUEST_TIMEOUT_MS_CONFIG, requestTimeoutMs.toString)
+    props.put(ProducerConfig.MAX_IN_FLIGHT_REQUESTS_PER_CONNECTION, maxInFlight.toString)
+    new KafkaProducer[Array[Byte], Array[Byte]](props, new ByteArraySerializer, new ByteArraySerializer)
   }
 
   // Seeds the given topic with records with keys and values in the range [0..numRecords)
   def seedTopicWithNumberedRecords(topic: String, numRecords: Int, servers: Seq[KafkaServer]): Unit = {
     val props = new Properties()
     props.put(ProducerConfig.ENABLE_IDEMPOTENCE_CONFIG, "true")
-    val producer = TestUtils.createNewProducer(TestUtils.getBrokerListStrFromServers(servers),
-      retries = Integer.MAX_VALUE, acks = -1, props = Some(props))
+    props.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, TestUtils.getBrokerListStrFromServers(servers))
+    val producer = new KafkaProducer[Array[Byte], Array[Byte]](props, new ByteArraySerializer, new ByteArraySerializer)
     try {
       for (i <- 0 until numRecords) {
         producer.send(new ProducerRecord[Array[Byte], Array[Byte]](topic, asBytes(i.toString), asBytes(i.toString)))
@@ -1449,74 +1505,308 @@ object TestUtils extends Logging {
     asString(record.value)
   }
 
-  def producerRecordWithExpectedTransactionStatus(topic: String, key: Array[Byte], value: Array[Byte],
-                                                  willBeCommitted: Boolean) : ProducerRecord[Array[Byte], Array[Byte]] = {
+  def producerRecordWithExpectedTransactionStatus(topic: String, partition: Integer, key: Array[Byte], value: Array[Byte], willBeCommitted: Boolean): ProducerRecord[Array[Byte], Array[Byte]] = {
     val header = new Header {override def key() = transactionStatusKey
       override def value() = if (willBeCommitted)
         committedValue
       else
         abortedValue
     }
-    new ProducerRecord[Array[Byte], Array[Byte]](topic, null, key, value, Collections.singleton(header))
+    new ProducerRecord[Array[Byte], Array[Byte]](topic, partition, key, value, Collections.singleton(header))
   }
 
-  def producerRecordWithExpectedTransactionStatus(topic: String, key: String, value: String,
-                                                  willBeCommitted: Boolean) : ProducerRecord[Array[Byte], Array[Byte]] = {
-    producerRecordWithExpectedTransactionStatus(topic, asBytes(key), asBytes(value), willBeCommitted)
+  def producerRecordWithExpectedTransactionStatus(topic: String, partition: Integer, key: String, value: String, willBeCommitted: Boolean): ProducerRecord[Array[Byte], Array[Byte]] = {
+    producerRecordWithExpectedTransactionStatus(topic, partition, asBytes(key), asBytes(value), willBeCommitted)
   }
 
   // Collect the current positions for all partition in the consumers current assignment.
   def consumerPositions(consumer: KafkaConsumer[Array[Byte], Array[Byte]]) : Map[TopicPartition, OffsetAndMetadata]  = {
     val offsetsToCommit = new mutable.HashMap[TopicPartition, OffsetAndMetadata]()
-    consumer.assignment.asScala.foreach { topicPartition =>
+    consumer.assignment.forEach { topicPartition =>
       offsetsToCommit.put(topicPartition, new OffsetAndMetadata(consumer.position(topicPartition)))
     }
     offsetsToCommit.toMap
   }
 
-  def pollUntilAtLeastNumRecords(consumer: KafkaConsumer[Array[Byte], Array[Byte]], numRecords: Int): Seq[ConsumerRecord[Array[Byte], Array[Byte]]] = {
-    val records = new ArrayBuffer[ConsumerRecord[Array[Byte], Array[Byte]]]()
-    TestUtils.waitUntilTrue(() => {
-      records ++= consumer.poll(50).asScala
-      records.size >= numRecords
-    }, s"Consumed ${records.size} records until timeout, but expected $numRecords records.")
-    records
-  }
+  def resetToCommittedPositions(consumer: KafkaConsumer[Array[Byte], Array[Byte]]): Unit = {
+    val committed = consumer.committed(consumer.assignment).asScala.filter(_._2 != null).map { case (k, v) => k -> v.offset }
 
-  def resetToCommittedPositions(consumer: KafkaConsumer[Array[Byte], Array[Byte]]) = {
-    consumer.assignment.asScala.foreach { case(topicPartition) =>
-      val offset = consumer.committed(topicPartition)
-      if (offset != null)
-        consumer.seek(topicPartition, offset.offset)
+    consumer.assignment.forEach { topicPartition =>
+      if (committed.contains(topicPartition))
+        consumer.seek(topicPartition, committed(topicPartition))
       else
         consumer.seekToBeginning(Collections.singletonList(topicPartition))
     }
   }
 
+  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 alterClientQuotas(adminClient: Admin, request: Map[ClientQuotaEntity, Map[String, Option[Double]]]): AlterClientQuotasResult = {
+    val entries = request.map { case (entity, alter) =>
+      val ops = alter.map { case (key, value) =>
+        new ClientQuotaAlteration.Op(key, value.map(Double.box).getOrElse(null))
+      }.asJavaCollection
+      new ClientQuotaAlteration(entity, ops)
+    }.asJavaCollection
+    adminClient.alterClientQuotas(entries)
+  }
+
+  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"
+    )
+  }
+
+  def waitForReplicasAssigned(client: Admin, partition: TopicPartition, brokerIds: Seq[Int]): Unit = {
+    TestUtils.waitUntilTrue(
+      () => {
+        val description = client.describeTopics(Set(partition.topic).asJava).all.get.asScala
+        val replicas = description
+          .values
+          .flatMap(_.partitions.asScala.flatMap(_.replicas.asScala))
+          .map(_.id)
+          .toSeq
+
+        brokerIds == replicas
+      },
+      s"Expected brokers $brokerIds to be the replicas for $partition"
+    )
+  }
+
   /**
    * Capture the console output during the execution of the provided function.
    */
   def grabConsoleOutput(f: => Unit) : String = {
     val out = new ByteArrayOutputStream
     try scala.Console.withOut(out)(f)
-    finally scala.Console.out.flush
+    finally scala.Console.out.flush()
     out.toString
   }
 
-}
+  /**
+   * Capture the console error during the execution of the provided function.
+   */
+  def grabConsoleError(f: => Unit) : String = {
+    val err = new ByteArrayOutputStream
+    try scala.Console.withErr(err)(f)
+    finally scala.Console.err.flush()
+    err.toString
+  }
 
-class IntEncoder(props: VerifiableProperties = null) extends Encoder[Int] {
-  override def toBytes(n: Int) = n.toString.getBytes
-}
+  /**
+   * Capture both the console output and console error during the execution of the provided function.
+   */
+  def grabConsoleOutputAndError(f: => Unit) : (String, String) = {
+    val out = new ByteArrayOutputStream
+    val err = new ByteArrayOutputStream
+    try scala.Console.withOut(out)(scala.Console.withErr(err)(f))
+    finally {
+      scala.Console.out.flush()
+      scala.Console.err.flush()
+    }
+    (out.toString, err.toString)
+  }
 
-@deprecated("This class is deprecated and it will be removed in a future release.", "0.10.0.0")
-class StaticPartitioner(props: VerifiableProperties = null) extends Partitioner {
-  def partition(data: Any, numPartitions: Int): Int = {
-    data.asInstanceOf[String].length % numPartitions
+  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
+        assertTrue("Expected an exception of type " + clazz.getName + "; got type " +
+            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)))
+    }
+  }
+
+  def totalMetricValue(server: KafkaServer, metricName: String): Long = {
+    val allMetrics = server.metrics.metrics
+    val total = allMetrics.values().asScala.filter(_.metricName().name() == metricName)
+      .foldLeft(0.0)((total, metric) => total + metric.metricValue.asInstanceOf[Double])
+    total.toLong
+  }
+
+  def meterCount(metricName: String): Long = {
+    KafkaYammerMetrics.defaultRegistry.allMetrics.asScala
+      .filter { case (k, _) => k.getMBeanName.endsWith(metricName) }
+      .values
+      .headOption
+      .getOrElse(fail(s"Unable to find metric $metricName"))
+      .asInstanceOf[Meter]
+      .count
+  }
+
+  def clearYammerMetrics(): Unit = {
+    for (metricName <- KafkaYammerMetrics.defaultRegistry.allMetrics.keySet.asScala)
+      KafkaYammerMetrics.defaultRegistry.removeMetric(metricName)
+  }
+
+  def stringifyTopicPartitions(partitions: Set[TopicPartition]): String = {
+    Json.encodeAsString(Map("partitions" ->
+      partitions.map(tp => Map("topic" -> tp.topic, "partition" -> tp.partition).asJava).asJava).asJava)
+  }
+
+  def resource[R <: AutoCloseable, A](resource: R)(func: R => A): A = {
+    try {
+      func(resource)
+    } finally {
+      resource.close()
+    }
+  }
+
+  /**
+   * Set broker replication quotas and enable throttling for a set of partitions. This
+   * will override any previous replication quotas, but will leave the throttling status
+   * of other partitions unaffected.
+   */
+  def setReplicationThrottleForPartitions(admin: Admin,
+                                          brokerIds: Seq[Int],
+                                          partitions: Set[TopicPartition],
+                                          throttleBytes: Int): Unit = {
+    throttleAllBrokersReplication(admin, brokerIds, throttleBytes)
+    assignThrottledPartitionReplicas(admin, partitions.map(_ -> brokerIds).toMap)
+  }
+
+  /**
+   * Remove a set of throttled partitions and reset the overall replication quota.
+   */
+  def removeReplicationThrottleForPartitions(admin: Admin,
+                                             brokerIds: Seq[Int],
+                                             partitions: Set[TopicPartition]): Unit = {
+    removePartitionReplicaThrottles(admin, partitions)
+    resetBrokersThrottle(admin, brokerIds)
+  }
+
+  /**
+    * 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(",")
+
+  def waitForAllReassignmentsToComplete(adminClient: Admin, pause: Long = 100L): Unit = {
+    waitUntilTrue(() => adminClient.listPartitionReassignments().reassignments().get().isEmpty,
+      s"There still are ongoing reassignments", pause = pause)
+  }
+
+  def addAndVerifyAcls(server: KafkaServer, acls: Set[AccessControlEntry], resource: ResourcePattern): Unit = {
+    val authorizer = server.dataPlaneRequestProcessor.authorizer.get
+    val aclBindings = acls.map { acl => new AclBinding(resource, acl) }
+    authorizer.createAcls(null, aclBindings.toList.asJava).asScala
+      .map(_.toCompletableFuture.get)
+      .foreach { result =>
+        result.exception.ifPresent { e => throw e }
+      }
+    val aclFilter = new AclBindingFilter(resource.toFilter, AccessControlEntryFilter.ANY)
+    waitAndVerifyAcls(
+      authorizer.acls(aclFilter).asScala.map(_.entry).toSet ++ acls,
+      authorizer, resource)
   }
-}
 
-@deprecated("This class has been deprecated and it will be removed in a future release.", "0.10.0.0")
-class FixedValuePartitioner(props: VerifiableProperties = null) extends Partitioner {
-  def partition(data: Any, numPartitions: Int): Int = data.asInstanceOf[Int]
 }
diff --git a/core/src/test/scala/unit/kafka/utils/ThrottlerTest.scala b/core/src/test/scala/unit/kafka/utils/ThrottlerTest.scala
new file mode 100755
index 0000000000000..a0b44c296a555
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/utils/ThrottlerTest.scala
@@ -0,0 +1,62 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package kafka.utils
+
+import org.apache.kafka.common.utils.MockTime
+import org.junit.Test
+import org.junit.Assert.{assertTrue, assertEquals}
+
+
+class ThrottlerTest {
+  @Test
+  def testThrottleDesiredRate(): Unit = {
+    val throttleCheckIntervalMs = 100
+    val desiredCountPerSec = 1000.0
+    val desiredCountPerInterval = desiredCountPerSec * throttleCheckIntervalMs / 1000.0
+
+    val mockTime = new MockTime()
+    val throttler = new Throttler(desiredRatePerSec = desiredCountPerSec,
+                                  checkIntervalMs = throttleCheckIntervalMs,
+                                  time = mockTime)
+
+    // Observe desiredCountPerInterval at t1
+    val t1 = mockTime.milliseconds()
+    throttler.maybeThrottle(desiredCountPerInterval)
+    assertEquals(t1, mockTime.milliseconds())
+
+    // Observe desiredCountPerInterval at t1 + throttleCheckIntervalMs + 1,
+    mockTime.sleep(throttleCheckIntervalMs + 1)
+    throttler.maybeThrottle(desiredCountPerInterval)
+    val t2 = mockTime.milliseconds()
+    assertTrue(t2 >= t1 + 2 * throttleCheckIntervalMs)
+
+    // Observe desiredCountPerInterval at t2
+    throttler.maybeThrottle(desiredCountPerInterval)
+    assertEquals(t2, mockTime.milliseconds())
+
+    // Observe desiredCountPerInterval at t2 + throttleCheckIntervalMs + 1
+    mockTime.sleep(throttleCheckIntervalMs + 1)
+    throttler.maybeThrottle(desiredCountPerInterval)
+    val t3 = mockTime.milliseconds()
+    assertTrue(t3 >= t2 + 2 * throttleCheckIntervalMs)
+
+    val elapsedTimeMs = t3 - t1
+    val actualCountPerSec = 4 * desiredCountPerInterval * 1000 / elapsedTimeMs
+    assertTrue(actualCountPerSec <= desiredCountPerSec)
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/utils/TopicFilterTest.scala b/core/src/test/scala/unit/kafka/utils/TopicFilterTest.scala
new file mode 100644
index 0000000000000..b89f58c28d114
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/utils/TopicFilterTest.scala
@@ -0,0 +1,49 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package kafka.utils
+
+import org.apache.kafka.common.internals.Topic
+import org.junit.Assert._
+import org.junit.Test
+
+class TopicFilterTest {
+
+  @Test
+  def testIncludeLists(): Unit = {
+
+    val topicFilter1 = IncludeList("yes1,yes2")
+    assertTrue(topicFilter1.isTopicAllowed("yes2", excludeInternalTopics = true))
+    assertTrue(topicFilter1.isTopicAllowed("yes2", excludeInternalTopics = false))
+    assertFalse(topicFilter1.isTopicAllowed("no1", excludeInternalTopics = true))
+    assertFalse(topicFilter1.isTopicAllowed("no1", excludeInternalTopics = false))
+
+    val topicFilter2 = IncludeList(".+")
+    assertTrue(topicFilter2.isTopicAllowed("alltopics", excludeInternalTopics = true))
+    assertFalse(topicFilter2.isTopicAllowed(Topic.GROUP_METADATA_TOPIC_NAME, excludeInternalTopics = true))
+    assertTrue(topicFilter2.isTopicAllowed(Topic.GROUP_METADATA_TOPIC_NAME, excludeInternalTopics = false))
+
+    val topicFilter3 = IncludeList("included-topic.+")
+    assertTrue(topicFilter3.isTopicAllowed("included-topic1", excludeInternalTopics = true))
+    assertFalse(topicFilter3.isTopicAllowed("no1", excludeInternalTopics = true))
+
+    val topicFilter4 = IncludeList("test-(?!bad\\b)[\\w]+")
+    assertTrue(topicFilter4.isTopicAllowed("test-good", excludeInternalTopics = true))
+    assertFalse(topicFilter4.isTopicAllowed("test-bad", 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 1ad37fcabafc3..0000000000000
--- a/core/src/test/scala/unit/kafka/utils/ZkUtilsTest.scala
+++ /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 kafka.utils
-
-import kafka.common.TopicAndPartition
-import kafka.zk.ZooKeeperTestHarness
-import org.junit.Assert._
-import org.junit.Test
-
-class ZkUtilsTest extends ZooKeeperTestHarness {
-
-  val path = "/path"
-
-  @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())
-  }
-}
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 b12d0f396b84e..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._
@@ -25,7 +27,7 @@ import kafka.utils.Json
 
 class JsonValueTest {
 
-  val json = """
+  private val json = """
     |{
     |  "boolean": false,
     |  "int": 1234,
@@ -66,7 +68,7 @@ class JsonValueTest {
   }
 
   @Test
-  def testAsJsonObject: Unit = {
+  def testAsJsonObject(): Unit = {
     val parsed = parse(json).asJsonObject
     val obj = parsed("object")
     assertEquals(obj, obj.asJsonObject)
@@ -74,14 +76,14 @@ class JsonValueTest {
   }
 
   @Test
-  def testAsJsonObjectOption: Unit = {
+  def testAsJsonObjectOption(): Unit = {
     val parsed = parse(json).asJsonObject
     assertTrue(parsed("object").asJsonObjectOption.isDefined)
     assertEquals(None, parsed("array").asJsonObjectOption)
   }
 
   @Test
-  def testAsJsonArray: Unit = {
+  def testAsJsonArray(): Unit = {
     val parsed = parse(json).asJsonObject
     val array = parsed("array")
     assertEquals(array, array.asJsonArray)
@@ -89,28 +91,28 @@ class JsonValueTest {
   }
 
   @Test
-  def testAsJsonArrayOption: Unit = {
+  def testAsJsonArrayOption(): Unit = {
     val parsed = parse(json).asJsonObject
     assertTrue(parsed("array").asJsonArrayOption.isDefined)
     assertEquals(None, parsed("object").asJsonArrayOption)
   }
 
   @Test
-  def testJsonObjectGet: Unit = {
+  def testJsonObjectGet(): Unit = {
     val parsed = parse(json).asJsonObject
     assertEquals(Some(parse("""{"a":true,"b":false}""")), parsed.get("object"))
     assertEquals(None, parsed.get("aaaaa"))
   }
 
   @Test
-  def testJsonObjectApply: Unit = {
+  def testJsonObjectApply(): Unit = {
     val parsed = parse(json).asJsonObject
     assertEquals(parse("""{"a":true,"b":false}"""), parsed("object"))
     assertThrow[JsonMappingException](parsed("aaaaaaaa"))
   }
 
   @Test
-  def testJsonObjectIterator: Unit = {
+  def testJsonObjectIterator(): Unit = {
     assertEquals(
       Vector("a" -> parse("true"), "b" -> parse("false")),
       parse(json).asJsonObject("object").asJsonObject.iterator.toVector
@@ -118,12 +120,12 @@ class JsonValueTest {
   }
 
   @Test
-  def testJsonArrayIterator: Unit = {
+  def testJsonArrayIterator(): Unit = {
     assertEquals(Vector("4.0", "11.1", "44.5").map(parse), parse(json).asJsonObject("array").asJsonArray.iterator.toVector)
   }
 
   @Test
-  def testJsonValueEquals: Unit = {
+  def testJsonValueEquals(): Unit = {
 
     assertEquals(parse(json), parse(json))
 
@@ -139,24 +141,24 @@ class JsonValueTest {
   }
 
   @Test
-  def testJsonValueHashCode: Unit = {
+  def testJsonValueHashCode(): Unit = {
     assertEquals(new ObjectMapper().readTree(json).hashCode, parse(json).hashCode)
   }
 
   @Test
-  def testJsonValueToString: Unit = {
+  def testJsonValueToString(): Unit = {
     val js = """{"boolean":false,"int":1234,"array":[4.0,11.1,44.5],"object":{"a":true,"b":false}}"""
     assertEquals(js, parse(js).toString)
   }
 
   @Test
-  def testDecodeBoolean: Unit = {
+  def testDecodeBoolean(): Unit = {
     assertTo[Boolean](false, _("boolean"))
     assertToFails[Boolean](_("int"))
   }
 
   @Test
-  def testDecodeString: Unit = {
+  def testDecodeString(): Unit = {
     assertTo[String]("string", _("string"))
     assertTo[String]("123", _("number_as_string"))
     assertToFails[String](_("int"))
@@ -164,20 +166,20 @@ class JsonValueTest {
   }
 
   @Test
-  def testDecodeInt: Unit = {
+  def testDecodeInt(): Unit = {
     assertTo[Int](1234, _("int"))
     assertToFails[Int](_("long"))
   }
 
   @Test
-  def testDecodeLong: Unit = {
+  def testDecodeLong(): Unit = {
     assertTo[Long](3000000000L, _("long"))
     assertTo[Long](1234, _("int"))
     assertToFails[Long](_("string"))
   }
 
   @Test
-  def testDecodeDouble: Unit = {
+  def testDecodeDouble(): Unit = {
     assertTo[Double](16.244355, _("double"))
     assertTo[Double](1234.0, _("int"))
     assertTo[Double](3000000000L, _("long"))
@@ -185,7 +187,7 @@ class JsonValueTest {
   }
 
   @Test
-  def testDecodeSeq: Unit = {
+  def testDecodeSeq(): Unit = {
     assertTo[Seq[Double]](Seq(4.0, 11.1, 44.5), _("array"))
     assertToFails[Seq[Double]](_("string"))
     assertToFails[Seq[Double]](_("object"))
@@ -193,7 +195,7 @@ class JsonValueTest {
   }
 
   @Test
-  def testDecodeMap: Unit = {
+  def testDecodeMap(): Unit = {
     assertTo[Map[String, Boolean]](Map("a" -> true, "b" -> false), _("object"))
     assertToFails[Map[String, Int]](_("object"))
     assertToFails[Map[String, String]](_("object"))
@@ -201,7 +203,7 @@ class JsonValueTest {
   }
 
   @Test
-  def testDecodeOption: Unit = {
+  def testDecodeOption(): Unit = {
     assertTo[Option[Int]](None, _("null"))
     assertTo[Option[Int]](Some(1234), _("int"))
     assertToFails[Option[String]](_("int"))
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 e4ac4fa1ec70c..819954a6d75a0 100644
--- a/core/src/test/scala/unit/kafka/utils/timer/MockTimer.scala
+++ b/core/src/test/scala/unit/kafka/utils/timer/MockTimer.scala
@@ -20,16 +20,18 @@ import kafka.utils.MockTime
 
 import scala.collection.mutable
 
-class MockTimer extends Timer {
+class MockTimer(val time: MockTime = new MockTime) 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
-      taskQueue.enqueue(new TimerTaskEntry(timerTask, timerTask.delayMs + time.milliseconds))
+    else {
+      taskQueue synchronized {
+        taskQueue.enqueue(new TimerTaskEntry(timerTask, timerTask.delayMs + time.milliseconds))
+      }
+    }
   }
 
   def advanceClock(timeoutMs: Long): Boolean = {
@@ -38,15 +40,25 @@ class MockTimer extends Timer {
     var executed = false
     val now = time.milliseconds
 
-    while (taskQueue.nonEmpty && now > taskQueue.head.expirationMs) {
-      val taskEntry = taskQueue.dequeue()
-      if (!taskEntry.cancelled) {
-        val task = taskEntry.timerTask
-        task.run()
-        executed = true
+    var hasMore = true
+    while (hasMore) {
+      hasMore = false
+      val head = taskQueue synchronized {
+        if (taskQueue.nonEmpty && now > taskQueue.head.expirationMs) {
+          val entry = Some(taskQueue.dequeue())
+          hasMore = taskQueue.nonEmpty
+          entry
+        } else
+          None
+      }
+      head.foreach { taskEntry =>
+        if (!taskEntry.cancelled) {
+          val task = taskEntry.timerTask
+          task.run()
+          executed = true
+        }
       }
     }
-
     executed
   }
 
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 993dcf2497e9c..6d7f9de01ced3 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._
@@ -28,50 +29,70 @@ import kafka.utils.TestUtils._
 import kafka.utils.{Logging, TestUtils}
 import kafka.zk.{AdminZkClient, KafkaZkClient, ZooKeeperTestHarness}
 import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.config.TopicConfig
 import org.apache.kafka.common.errors.{InvalidReplicaAssignmentException, InvalidTopicException, TopicExistsException}
 import org.apache.kafka.common.metrics.Quota
+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.jdk.CollectionConverters._
+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(zkUtils, brokers)
+    TestUtils.createBrokersInZk(zkClient, brokers)
+
+    val topicConfig = new Properties()
 
     // duplicate brokers
     intercept[InvalidReplicaAssignmentException] {
-      adminZkClient.createOrUpdateTopicPartitionAssignmentPathInZK("test", Map(0->Seq(0,0)))
+      adminZkClient.createTopicWithAssignment("test", topicConfig, Map(0->Seq(0,0)))
     }
 
     // inconsistent replication factor
     intercept[InvalidReplicaAssignmentException] {
-      adminZkClient.createOrUpdateTopicPartitionAssignmentPathInZK("test", Map(0->Seq(0,1), 1->Seq(0)))
+      adminZkClient.createTopicWithAssignment("test", topicConfig, Map(0->Seq(0,1), 1->Seq(0)))
+    }
+
+    // partitions should be 0-based
+    intercept[InvalidReplicaAssignmentException] {
+      adminZkClient.createTopicWithAssignment("test", topicConfig, Map(1->Seq(1,2), 2->Seq(1,2)))
+    }
+
+    // partitions should be 0-based and consecutive
+    intercept[InvalidReplicaAssignmentException] {
+      adminZkClient.createTopicWithAssignment("test", topicConfig, Map(0->Seq(1,2), 0->Seq(1,2), 3->Seq(1,2)))
+    }
+
+    // partitions should be 0-based and consecutive
+    intercept[InvalidReplicaAssignmentException] {
+      adminZkClient.createTopicWithAssignment("test", topicConfig, Map(-1->Seq(1,2), 1->Seq(1,2), 2->Seq(1,2), 4->Seq(1,2)))
     }
 
     // good assignment
     val assignment = Map(0 -> List(0, 1, 2),
                          1 -> List(1, 2, 3))
-    adminZkClient.createOrUpdateTopicPartitionAssignmentPathInZK("test", assignment)
+    adminZkClient.createTopicWithAssignment("test", topicConfig, assignment)
     val found = zkClient.getPartitionAssignmentForTopics(Set("test"))
-    assertEquals(assignment, found("test"))
+    assertEquals(assignment.map { case (k, v) => k -> ReplicaAssignment(v, List(), List()) }, found("test"))
   }
 
   @Test
-  def testTopicCreationInZK() {
+  def testTopicCreationInZK(): Unit = {
     val expectedReplicaAssignment = Map(
       0  -> List(0, 1, 2),
       1  -> List(1, 2, 3),
@@ -101,11 +122,12 @@ class AdminZkClientTest extends ZooKeeperTestHarness with Logging with RackAware
       11 -> 1
     )
     val topic = "test"
-    TestUtils.createBrokersInZk(zkUtils, List(0, 1, 2, 3, 4))
+    val topicConfig = new Properties()
+    TestUtils.createBrokersInZk(zkClient, List(0, 1, 2, 3, 4))
     // create the topic
-    adminZkClient.createOrUpdateTopicPartitionAssignmentPathInZK(topic, expectedReplicaAssignment)
+    adminZkClient.createTopicWithAssignment(topic, topicConfig, expectedReplicaAssignment)
     // create leaders for all partitions
-    TestUtils.makeLeaderForPartition(zkUtils, topic, leaderForPartitionMap, 1)
+    TestUtils.makeLeaderForPartition(zkClient, topic, leaderForPartitionMap, 1)
     val actualReplicaMap = leaderForPartitionMap.keys.map(p => p -> zkClient.getReplicasForPartition(new TopicPartition(topic, p))).toMap
     assertEquals(expectedReplicaAssignment.size, actualReplicaMap.size)
     for(i <- 0 until actualReplicaMap.size)
@@ -113,15 +135,15 @@ class AdminZkClientTest extends ZooKeeperTestHarness with Logging with RackAware
 
     intercept[TopicExistsException] {
       // shouldn't be able to create a topic that already exists
-      adminZkClient.createOrUpdateTopicPartitionAssignmentPathInZK(topic, expectedReplicaAssignment)
+      adminZkClient.createTopicWithAssignment(topic, topicConfig, expectedReplicaAssignment)
     }
   }
 
   @Test
-  def testTopicCreationWithCollision() {
+  def testTopicCreationWithCollision(): Unit = {
     val topic = "test.topic"
     val collidingTopic = "test_topic"
-    TestUtils.createBrokersInZk(zkUtils, List(0, 1, 2, 3, 4))
+    TestUtils.createBrokersInZk(zkClient, List(0, 1, 2, 3, 4))
     // create the topic
     adminZkClient.createTopic(topic, 3, 1)
 
@@ -132,19 +154,42 @@ class AdminZkClientTest extends ZooKeeperTestHarness with Logging with RackAware
   }
 
   @Test
-  def testConcurrentTopicCreation() {
+  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 = EasyMock.createNiceMock(classOf[KafkaZkClient])
+    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(false)).andReturn(Set("some.topic", topic, "some.other.topic"))
     EasyMock.replay(zkMock)
     val adminZkClient = new AdminZkClient(zkMock)
 
     intercept[TopicExistsException] {
-      adminZkClient.validateCreateOrUpdateTopic(topic, Map.empty, new Properties, update = false)
+      adminZkClient.validateTopicCreate(topic, Map.empty, new Properties)
+    }
+  }
+
+  @Test
+  def testConcurrentTopicCreation(): Unit = {
+    val topic = "test-concurrent-topic-creation"
+    TestUtils.createBrokersInZk(zkClient, List(0, 1, 2, 3, 4))
+    val props = new Properties
+    props.setProperty(TopicConfig.MIN_IN_SYNC_REPLICAS_CONFIG, "2")
+    def createTopic(): Unit = {
+      try adminZkClient.createTopic(topic, 3, 1, props)
+      catch { case _: TopicExistsException => () }
+      val (_, partitionAssignment) = zkClient.getPartitionAssignmentForTopics(Set(topic)).head
+      assertEquals(3, partitionAssignment.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)
     }
+
+    TestUtils.assertConcurrent("Concurrent topic creation failed", Seq(() => createTopic(), () => createTopic()),
+      JTestUtils.DEFAULT_MAX_WAIT_MS.toInt)
   }
 
   /**
@@ -152,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)))
@@ -167,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 == "")
@@ -225,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)
@@ -270,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")
@@ -280,23 +325,23 @@ class AdminZkClientTest extends ZooKeeperTestHarness with Logging with RackAware
     zkClient.setOrCreateEntityConfigs(ConfigType.Client, clientId, props)
 
     val configInZk: Map[String, Properties] = adminZkClient.fetchAllEntityConfigs(ConfigType.Client)
-    assertEquals("Must have 1 overriden client config", 1, configInZk.size)
+    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.apis.quotas.produce.quota("ANONYMOUS", clientId))
-    assertEquals(new Quota(2000, true), server.apis.quotas.fetch.quota("ANONYMOUS", clientId))
+    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() {
+  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")
     val brokerMetadatas = toBrokerMetadata(rackInfo, brokersWithoutRack = brokerList.filterNot(rackInfo.keySet))
-    TestUtils.createBrokersInZk(brokerMetadatas, zkUtils)
+    TestUtils.createBrokersInZk(brokerMetadatas, zkClient)
 
     val processedMetadatas1 = adminZkClient.getBrokerMetadatas(RackAwareMode.Disabled)
     assertEquals(brokerList, processedMetadatas1.map(_.id))
diff --git a/core/src/test/scala/unit/kafka/zk/EmbeddedZookeeper.scala b/core/src/test/scala/unit/kafka/zk/EmbeddedZookeeper.scala
index d4a829dfe053a..28b592eaf7af8 100755
--- a/core/src/test/scala/unit/kafka/zk/EmbeddedZookeeper.scala
+++ b/core/src/test/scala/unit/kafka/zk/EmbeddedZookeeper.scala
@@ -38,7 +38,9 @@ class EmbeddedZookeeper() extends Logging {
 
   val snapshotDir = TestUtils.tempDir()
   val logDir = TestUtils.tempDir()
-  val tickTime = 500
+  val tickTime = 800 // allow a maxSessionTimeout of 20 * 800ms = 16 secs
+
+  System.setProperty("zookeeper.forceSync", "no")  //disable fsync to ZK txn log in tests to avoid timeout
   val zookeeper = new ZooKeeperServer(snapshotDir, logDir, tickTime)
   val factory = new NIOServerCnxnFactory()
   private val addr = new InetSocketAddress("127.0.0.1", TestUtils.RandomPort)
@@ -46,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 = {
@@ -58,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 28d84308d0686..7bcfbfd94a07e 100644
--- a/core/src/test/scala/unit/kafka/zk/KafkaZkClientTest.scala
+++ b/core/src/test/scala/unit/kafka/zk/KafkaZkClientTest.scala
@@ -16,23 +16,110 @@
 */
 package kafka.zk
 
-import java.util.{Properties, UUID}
+import java.util.{Collections, Properties}
+import java.nio.charset.StandardCharsets.UTF_8
+import java.util.concurrent.{CountDownLatch, TimeUnit}
 
+import kafka.api.{ApiVersion, LeaderAndIsr}
+import kafka.cluster.{Broker, EndPoint}
 import kafka.log.LogConfig
-import kafka.security.auth._
-import kafka.server.ConfigType
-import org.apache.kafka.common.TopicPartition
-import org.apache.kafka.common.security.auth.KafkaPrincipal
-import org.junit.Assert.{assertEquals, assertFalse, assertTrue}
-import org.junit.Test
+import kafka.server.{ConfigType, KafkaConfig}
+import kafka.utils.CoreUtils
+import org.apache.kafka.common.{TopicPartition, Uuid}
+import org.apache.kafka.common.network.ListenerName
+import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
+import org.apache.kafka.common.security.token.delegation.TokenInformation
+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.jdk.CollectionConverters._
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.{Seq, mutable}
+import scala.util.Random
+import kafka.controller.{LeaderIsrAndControllerEpoch, ReplicaAssignment}
+import kafka.security.authorizer.AclEntry
+import kafka.zk.KafkaZkClient.UpdateLeaderAndIsrResult
+import kafka.zookeeper._
+import org.apache.kafka.common.acl.AclOperation.READ
+import org.apache.kafka.common.acl.AclPermissionType.{ALLOW, DENY}
+import org.apache.kafka.common.errors.ControllerMovedException
+import org.apache.kafka.common.feature.{Features, SupportedVersionRange}
+import org.apache.kafka.common.feature.Features._
+import org.apache.kafka.common.resource.ResourcePattern
+import org.apache.kafka.common.resource.ResourceType.{GROUP, TOPIC}
+import org.apache.kafka.common.security.JaasUtils
+import org.apache.zookeeper.ZooDefs
+import org.apache.zookeeper.client.ZKClientConfig
+import org.apache.zookeeper.data.Stat
 
 class KafkaZkClientTest extends ZooKeeperTestHarness {
 
   private val group = "my-group"
+  private val topic1 = "topic1"
+  private val topic2 = "topic2"
+  private val topicIds = Map(topic1 -> Uuid.randomUuid(), topic2 -> Uuid.randomUuid())
+
+  val topicPartition10 = new TopicPartition(topic1, 0)
+  val topicPartition11 = new TopicPartition(topic1, 1)
+  val topicPartition20 = new TopicPartition(topic2, 0)
+  val topicPartitions10_11 = Seq(topicPartition10, topicPartition11)
+  val controllerEpochZkVersion = 0
+
+  var otherZkClient: KafkaZkClient = _
+  var expiredSessionZkClient: ExpiredKafkaZkClient = _
+
+  @Before
+  override def setUp(): Unit = {
+    super.setUp()
+    zkClient.createControllerEpochRaw(1)
+    otherZkClient = KafkaZkClient(zkConnect, zkAclsEnabled.getOrElse(JaasUtils.isZkSaslEnabled), zkSessionTimeout,
+      zkConnectionTimeout, zkMaxInFlightRequests, Time.SYSTEM)
+    expiredSessionZkClient = ExpiredKafkaZkClient(zkConnect, zkAclsEnabled.getOrElse(JaasUtils.isZkSaslEnabled),
+      zkSessionTimeout, zkConnectionTimeout, zkMaxInFlightRequests, Time.SYSTEM)
+  }
+
+  @After
+  override def tearDown(): Unit = {
+    if (otherZkClient != null)
+      otherZkClient.close()
+    zkClient.deletePath(ControllerEpochZNode.path)
+    if (expiredSessionZkClient != null)
+      expiredSessionZkClient.close()
+    super.tearDown()
+  }
+
   private val topicPartition = new TopicPartition("topic", 0)
 
   @Test
-  def testSetAndGetConsumerOffset() {
+  def testConnectionViaNettyClient(): Unit = {
+    // Confirm that we can explicitly set client connection configuration, which is necessary for TLS.
+    // TLS connectivity itself is tested in system tests rather than here to avoid having to add TLS support
+    // to kafka.zk.EmbeddedZoopeeper
+    val clientConfig = new ZKClientConfig()
+    val propKey = KafkaConfig.ZkClientCnxnSocketProp
+    val propVal = "org.apache.zookeeper.ClientCnxnSocketNetty"
+    KafkaConfig.setZooKeeperClientProperty(clientConfig, propKey, propVal)
+    val client = KafkaZkClient(zkConnect, zkAclsEnabled.getOrElse(JaasUtils.isZkSaslEnabled), zkSessionTimeout,
+      zkConnectionTimeout, zkMaxInFlightRequests, Time.SYSTEM, zkClientConfig = Some(clientConfig))
+    try {
+      assertEquals(Some(propVal), KafkaConfig.getZooKeeperClientProperty(client.currentZooKeeper.getClientConfig, propKey))
+      // For a sanity check, make sure a bad client connection socket class name generates an exception
+      val badClientConfig = new ZKClientConfig()
+      KafkaConfig.setZooKeeperClientProperty(badClientConfig, propKey, propVal + "BadClassName")
+      intercept[Exception] {
+        KafkaZkClient(zkConnect, zkAclsEnabled.getOrElse(JaasUtils.isZkSaslEnabled), zkSessionTimeout,
+          zkConnectionTimeout, zkMaxInFlightRequests, Time.SYSTEM, zkClientConfig = Some(badClientConfig))
+      }
+    } finally {
+      client.close()
+    }
+  }
+
+  @Test
+  def testSetAndGetConsumerOffset(): Unit = {
     val offset = 123L
     // None if no committed offsets
     assertTrue(zkClient.getConsumerOffset(group, topicPartition).isEmpty)
@@ -45,28 +132,41 @@ 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")
     assertTrue(zkClient.pathExists("/delete/some/random/path"))
-    zkClient.deleteRecursive("/delete")
-    assertFalse(zkClient.pathExists("/delete/some/random/path"))
-    assertFalse(zkClient.pathExists("/delete/some/random"))
-    assertFalse(zkClient.pathExists("/delete/some"))
+    assertTrue(zkClient.deleteRecursive("/delete"))
     assertFalse(zkClient.pathExists("/delete"))
 
     intercept[IllegalArgumentException](zkClient.deleteRecursive("delete-invalid-path"))
   }
 
   @Test
-  def testCreateRecursive() {
+  def testDeleteRecursiveWithControllerEpochVersionCheck(): Unit = {
+    assertFalse(zkClient.deleteRecursive("/delete/does-not-exist", controllerEpochZkVersion))
+
+    zkClient.createRecursive("/delete/some/random/path")
+    assertTrue(zkClient.pathExists("/delete/some/random/path"))
+    intercept[ControllerMovedException](
+      zkClient.deleteRecursive("/delete", controllerEpochZkVersion + 1))
+
+    assertTrue(zkClient.deleteRecursive("/delete", controllerEpochZkVersion))
+    assertFalse(zkClient.pathExists("/delete"))
+
+    intercept[IllegalArgumentException](zkClient.deleteRecursive(
+      "delete-invalid-path", controllerEpochZkVersion))
+  }
+
+  @Test
+  def testCreateRecursive(): Unit = {
     zkClient.createRecursive("/create-newrootpath")
     assertTrue(zkClient.pathExists("/create-newrootpath"))
 
@@ -78,11 +178,11 @@ class KafkaZkClientTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testTopicAssignmentMethods() {
-    val topic1 = "topic1"
-    val topic2 = "topic2"
+  def testTopicAssignmentMethods(): Unit = {
+    assertTrue(zkClient.getAllTopicsInCluster().isEmpty)
 
     // test with non-existing topic
+    assertFalse(zkClient.topicExists(topic1))
     assertTrue(zkClient.getTopicPartitionCount(topic1).isEmpty)
     assertTrue(zkClient.getPartitionAssignmentForTopics(Set(topic1)).isEmpty)
     assertTrue(zkClient.getPartitionsForTopics(Set(topic1)).isEmpty)
@@ -95,22 +195,25 @@ class KafkaZkClientTest extends ZooKeeperTestHarness {
     )
 
     // create a topic assignment
-    zkClient.createTopicAssignment(topic1, assignment)
+    zkClient.createTopicAssignment(topic1, topicIds(topic1), assignment)
+
+    assertTrue(zkClient.topicExists(topic1))
 
     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)
-    assertEquals(expectedAssignment, zkClient.getPartitionAssignmentForTopics(Set(topic1)).get(topic1).get)
-    assertEquals(Set(0, 1, 2), zkClient.getPartitionsForTopics(Set(topic1)).get(topic1).get.toSet)
+    assertEquals(expectedAssignment, zkClient.getPartitionAssignmentForTopics(Set(topic1))(topic1))
+    assertEquals(Set(0, 1, 2), zkClient.getPartitionsForTopics(Set(topic1))(topic1).toSet)
     assertEquals(Set(1, 2, 3), zkClient.getReplicasForPartition(new TopicPartition(topic1, 2)).toSet)
 
     val updatedAssignment = assignment - new TopicPartition(topic1, 2)
 
-    zkClient.setTopicAssignment(topic1, updatedAssignment)
+    zkClient.setTopicAssignment(topic1, topicIds(topic1), updatedAssignment.map {
+      case (k, v) => k -> ReplicaAssignment(v, List(), List()) })
     assertEquals(updatedAssignment.size, zkClient.getTopicPartitionCount(topic1).get)
 
     // add second topic
@@ -119,60 +222,257 @@ class KafkaZkClientTest extends ZooKeeperTestHarness {
       new TopicPartition(topic2, 1) -> Seq(0, 1)
     )
 
-    zkClient.createTopicAssignment(topic2, secondAssignment)
+    zkClient.createTopicAssignment(topic2, topicIds(topic2), secondAssignment)
+
+    assertEquals(Set(topic1, topic2), zkClient.getAllTopicsInCluster())
+  }
+
+  @Test
+  def testGetAllTopicsInClusterTriggersWatch(): Unit = {
+    zkClient.createTopLevelPaths()
+    val latch = registerChildChangeHandler(1)
+
+    // Listing all the topics and register the watch
+    assertTrue(zkClient.getAllTopicsInCluster(true).isEmpty)
+
+    // Verifies that listing all topics without registering the watch does
+    // not interfere with the previous registered watcher
+    assertTrue(zkClient.getAllTopicsInCluster(false).isEmpty)
+
+    zkClient.createTopicAssignment(topic1, topicIds(topic1), Map.empty)
+
+    assertTrue("Failed to receive watch notification",
+      latch.await(5, TimeUnit.SECONDS))
+
+    assertTrue(zkClient.topicExists(topic1))
+  }
+
+  @Test
+  def testGetAllTopicsInClusterDoesNotTriggerWatch(): Unit = {
+    zkClient.createTopLevelPaths()
+    val latch = registerChildChangeHandler(1)
+
+    // Listing all the topics and don't register the watch
+    assertTrue(zkClient.getAllTopicsInCluster(false).isEmpty)
+
+    zkClient.createTopicAssignment(topic1, topicIds(topic1), Map.empty)
+
+    assertFalse("Received watch notification",
+      latch.await(100, TimeUnit.MILLISECONDS))
+
+    assertTrue(zkClient.topicExists(topic1))
+  }
+
+  private def registerChildChangeHandler(count: Int): CountDownLatch = {
+    val znodeChildChangeHandlerCountDownLatch = new CountDownLatch(1)
+    val znodeChildChangeHandler = new ZNodeChildChangeHandler {
+      override val path: String = TopicsZNode.path
 
-    assertEquals(Set(topic1, topic2), zkClient.getAllTopicsInCluster.toSet)
+      override def handleChildChange(): Unit = {
+        znodeChildChangeHandlerCountDownLatch.countDown()
+      }
+    }
+    zkClient.registerZNodeChildChangeHandler(znodeChildChangeHandler)
+    znodeChildChangeHandlerCountDownLatch
   }
 
   @Test
-  def testGetDataAndVersion() {
+  def testGetDataAndVersion(): Unit = {
     val path = "/testpath"
 
     // test with non-existing path
-    var dataAndVersion = zkClient.getDataAndVersion(path)
-    assertTrue(dataAndVersion._1.isEmpty)
-    assertEquals(-1, dataAndVersion._2)
+    val (data0, version0) = zkClient.getDataAndVersion(path)
+    assertTrue(data0.isEmpty)
+    assertEquals(ZkVersion.UnknownVersion, version0)
 
     // create a test path
     zkClient.createRecursive(path)
-    zkClient.conditionalUpdatePath(path, "version1", 0)
+    zkClient.conditionalUpdatePath(path, "version1".getBytes(UTF_8), 0)
 
     // test with existing path
-    dataAndVersion = zkClient.getDataAndVersion(path)
-    assertEquals("version1", dataAndVersion._1.get)
-    assertEquals(1, dataAndVersion._2)
-
-    zkClient.conditionalUpdatePath(path, "version2", 1)
-    dataAndVersion = zkClient.getDataAndVersion(path)
-    assertEquals("version2", dataAndVersion._1.get)
-    assertEquals(2, dataAndVersion._2)
+    val (data1, version1) = zkClient.getDataAndVersion(path)
+    assertEquals("version1", new String(data1.get, UTF_8))
+    assertEquals(1, version1)
+
+    zkClient.conditionalUpdatePath(path, "version2".getBytes(UTF_8), 1)
+    val (data2, version2) = zkClient.getDataAndVersion(path)
+    assertEquals("version2", new String(data2.get, UTF_8))
+    assertEquals(2, version2)
   }
 
   @Test
-  def testConditionalUpdatePath() {
+  def testConditionalUpdatePath(): Unit = {
     val path = "/testconditionalpath"
 
     // test with non-existing path
-    var statusAndVersion = zkClient.conditionalUpdatePath(path, "version0", 0)
+    var statusAndVersion = zkClient.conditionalUpdatePath(path, "version0".getBytes(UTF_8), 0)
     assertFalse(statusAndVersion._1)
-    assertEquals(-1, statusAndVersion._2)
+    assertEquals(ZkVersion.UnknownVersion, statusAndVersion._2)
 
     // create path
     zkClient.createRecursive(path)
 
     // test with valid expected version
-    statusAndVersion = zkClient.conditionalUpdatePath(path, "version1", 0)
+    statusAndVersion = zkClient.conditionalUpdatePath(path, "version1".getBytes(UTF_8), 0)
     assertTrue(statusAndVersion._1)
     assertEquals(1, statusAndVersion._2)
 
     // test with invalid expected version
-    statusAndVersion = zkClient.conditionalUpdatePath(path, "version2", 2)
+    statusAndVersion = zkClient.conditionalUpdatePath(path, "version2".getBytes(UTF_8), 2)
     assertFalse(statusAndVersion._1)
-    assertEquals(-1, statusAndVersion._2)
+    assertEquals(ZkVersion.UnknownVersion, statusAndVersion._2)
+  }
+
+  @Test
+  def testCreateSequentialPersistentPath(): Unit = {
+    val path = "/testpath"
+    zkClient.createRecursive(path)
+
+    var result = zkClient.createSequentialPersistentPath(path + "/sequence_", null)
+    assertEquals(s"$path/sequence_0000000000", result)
+    assertTrue(zkClient.pathExists(s"$path/sequence_0000000000"))
+    assertEquals(None, dataAsString(s"$path/sequence_0000000000"))
+
+    result = zkClient.createSequentialPersistentPath(path + "/sequence_", "some value".getBytes(UTF_8))
+    assertEquals(s"$path/sequence_0000000001", result)
+    assertTrue(zkClient.pathExists(s"$path/sequence_0000000001"))
+    assertEquals(Some("some value"), dataAsString(s"$path/sequence_0000000001"))
   }
 
   @Test
-  def testSetGetAndDeletePartitionReassignment() {
+  def testPropagateIsrChanges(): Unit = {
+    zkClient.createRecursive("/isr_change_notification")
+
+    zkClient.propagateIsrChanges(Set(new TopicPartition("topic-a", 0), new TopicPartition("topic-b", 0)))
+    var expectedPath = "/isr_change_notification/isr_change_0000000000"
+    assertTrue(zkClient.pathExists(expectedPath))
+    assertEquals(Some("""{"version":1,"partitions":[{"topic":"topic-a","partition":0},{"topic":"topic-b","partition":0}]}"""),
+      dataAsString(expectedPath))
+
+    zkClient.propagateIsrChanges(Set(new TopicPartition("topic-b", 0)))
+    expectedPath = "/isr_change_notification/isr_change_0000000001"
+    assertTrue(zkClient.pathExists(expectedPath))
+    assertEquals(Some("""{"version":1,"partitions":[{"topic":"topic-b","partition":0}]}"""), dataAsString(expectedPath))
+  }
+
+  @Test
+  def testIsrChangeNotificationGetters(): Unit = {
+    assertEquals("Failed for non existing parent ZK node", Seq.empty, zkClient.getAllIsrChangeNotifications)
+    assertEquals("Failed for non existing parent ZK node", Seq.empty, zkClient.getPartitionsFromIsrChangeNotifications(Seq("0000000000")))
+
+    zkClient.createRecursive("/isr_change_notification")
+
+    zkClient.propagateIsrChanges(Set(topicPartition10, topicPartition11))
+    zkClient.propagateIsrChanges(Set(topicPartition10))
+
+    assertEquals(Set("0000000000", "0000000001"), zkClient.getAllIsrChangeNotifications.toSet)
+
+    // A partition can have multiple notifications
+    assertEquals(Seq(topicPartition10, topicPartition11, topicPartition10),
+      zkClient.getPartitionsFromIsrChangeNotifications(Seq("0000000000", "0000000001")))
+  }
+
+  @Test
+  def testIsrChangeNotificationsDeletion(): Unit = {
+    // Should not fail even if parent node does not exist
+    zkClient.deleteIsrChangeNotifications(Seq("0000000000"), controllerEpochZkVersion)
+
+    zkClient.createRecursive("/isr_change_notification")
+
+    zkClient.propagateIsrChanges(Set(topicPartition10, topicPartition11))
+    zkClient.propagateIsrChanges(Set(topicPartition10))
+    zkClient.propagateIsrChanges(Set(topicPartition11))
+
+    // Should throw exception if the controllerEpochZkVersion does not match
+    intercept[ControllerMovedException](zkClient.deleteIsrChangeNotifications(Seq("0000000001"), controllerEpochZkVersion + 1))
+    // Delete should not succeed
+    assertEquals(Set("0000000000", "0000000001", "0000000002"), zkClient.getAllIsrChangeNotifications.toSet)
+
+    zkClient.deleteIsrChangeNotifications(Seq("0000000001"), controllerEpochZkVersion)
+    // Should not fail if called on a non-existent notification
+    zkClient.deleteIsrChangeNotifications(Seq("0000000001"), controllerEpochZkVersion)
+
+    assertEquals(Set("0000000000", "0000000002"), zkClient.getAllIsrChangeNotifications.toSet)
+    zkClient.deleteIsrChangeNotifications(controllerEpochZkVersion)
+    assertEquals(Seq.empty, zkClient.getAllIsrChangeNotifications)
+  }
+
+  @Test
+  def testPropagateLogDir(): Unit = {
+    zkClient.createRecursive("/log_dir_event_notification")
+
+    val brokerId = 3
+
+    zkClient.propagateLogDirEvent(brokerId)
+    var expectedPath = "/log_dir_event_notification/log_dir_event_0000000000"
+    assertTrue(zkClient.pathExists(expectedPath))
+    assertEquals(Some("""{"version":1,"broker":3,"event":1}"""), dataAsString(expectedPath))
+
+    zkClient.propagateLogDirEvent(brokerId)
+    expectedPath = "/log_dir_event_notification/log_dir_event_0000000001"
+    assertTrue(zkClient.pathExists(expectedPath))
+    assertEquals(Some("""{"version":1,"broker":3,"event":1}"""), dataAsString(expectedPath))
+
+    val anotherBrokerId = 4
+    zkClient.propagateLogDirEvent(anotherBrokerId)
+    expectedPath = "/log_dir_event_notification/log_dir_event_0000000002"
+    assertTrue(zkClient.pathExists(expectedPath))
+    assertEquals(Some("""{"version":1,"broker":4,"event":1}"""), dataAsString(expectedPath))
+  }
+
+  @Test
+  def testLogDirGetters(): Unit = {
+    assertEquals("getAllLogDirEventNotifications failed for non existing parent ZK node",
+      Seq.empty, zkClient.getAllLogDirEventNotifications)
+    assertEquals("getBrokerIdsFromLogDirEvents failed for non existing parent ZK node",
+      Seq.empty, zkClient.getBrokerIdsFromLogDirEvents(Seq("0000000000")))
+
+    zkClient.createRecursive("/log_dir_event_notification")
+
+    val brokerId = 3
+    zkClient.propagateLogDirEvent(brokerId)
+
+    assertEquals(Seq(3), zkClient.getBrokerIdsFromLogDirEvents(Seq("0000000000")))
+
+    zkClient.propagateLogDirEvent(brokerId)
+
+    val anotherBrokerId = 4
+    zkClient.propagateLogDirEvent(anotherBrokerId)
+
+    val notifications012 = Seq("0000000000", "0000000001", "0000000002")
+    assertEquals(notifications012.toSet, zkClient.getAllLogDirEventNotifications.toSet)
+    assertEquals(Seq(3, 3, 4), zkClient.getBrokerIdsFromLogDirEvents(notifications012))
+  }
+
+  @Test
+  def testLogDirEventNotificationsDeletion(): Unit = {
+    // Should not fail even if parent node does not exist
+    zkClient.deleteLogDirEventNotifications(Seq("0000000000", "0000000002"), controllerEpochZkVersion)
+
+    zkClient.createRecursive("/log_dir_event_notification")
+
+    val brokerId = 3
+    val anotherBrokerId = 4
+
+    zkClient.propagateLogDirEvent(brokerId)
+    zkClient.propagateLogDirEvent(brokerId)
+    zkClient.propagateLogDirEvent(anotherBrokerId)
+
+    intercept[ControllerMovedException](zkClient.deleteLogDirEventNotifications(Seq("0000000000", "0000000002"), controllerEpochZkVersion + 1))
+    assertEquals(Seq("0000000000", "0000000001", "0000000002"), zkClient.getAllLogDirEventNotifications)
+
+    zkClient.deleteLogDirEventNotifications(Seq("0000000000", "0000000002"), controllerEpochZkVersion)
+
+    assertEquals(Seq("0000000001"), zkClient.getAllLogDirEventNotifications)
+
+    zkClient.propagateLogDirEvent(anotherBrokerId)
+
+    zkClient.deleteLogDirEventNotifications(controllerEpochZkVersion)
+    assertEquals(Seq.empty, zkClient.getAllLogDirEventNotifications)
+  }
+
+  @Test
+  def testSetGetAndDeletePartitionReassignment(): Unit = {
     zkClient.createRecursive(AdminZNode.path)
 
     assertEquals(Map.empty, zkClient.getPartitionReassignment)
@@ -183,14 +483,18 @@ class KafkaZkClientTest extends ZooKeeperTestHarness {
       new TopicPartition("topic_b", 0) -> Seq(4, 5),
       new TopicPartition("topic_c", 0) -> Seq(5, 3)
     )
-    zkClient.setOrCreatePartitionReassignment(reassignment)
+
+    // Should throw ControllerMovedException if the controller epoch zkVersion does not match
+    intercept[ControllerMovedException](zkClient.setOrCreatePartitionReassignment(reassignment, controllerEpochZkVersion + 1))
+
+    zkClient.setOrCreatePartitionReassignment(reassignment, controllerEpochZkVersion)
     assertEquals(reassignment, zkClient.getPartitionReassignment)
 
-    val updatedReassingment = reassignment - new TopicPartition("topic_b", 0)
-    zkClient.setOrCreatePartitionReassignment(updatedReassingment)
-    assertEquals(updatedReassingment, zkClient.getPartitionReassignment)
+    val updatedReassignment = reassignment - new TopicPartition("topic_b", 0)
+    zkClient.setOrCreatePartitionReassignment(updatedReassignment, controllerEpochZkVersion)
+    assertEquals(updatedReassignment, zkClient.getPartitionReassignment)
 
-    zkClient.deletePartitionReassignment()
+    zkClient.deletePartitionReassignment(controllerEpochZkVersion)
     assertEquals(Map.empty, zkClient.getPartitionReassignment)
 
     zkClient.createPartitionReassignment(reassignment)
@@ -198,31 +502,31 @@ class KafkaZkClientTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testGetDataAndStat() {
+  def testGetDataAndStat(): Unit = {
     val path = "/testpath"
 
     // test with non-existing path
-    var dataAndVersion = zkClient.getDataAndStat(path)
-    assertTrue(dataAndVersion._1.isEmpty)
-    assertEquals(0, dataAndVersion._2.getVersion)
+    val (data0, version0) = zkClient.getDataAndStat(path)
+    assertTrue(data0.isEmpty)
+    assertEquals(0, version0.getVersion)
 
     // create a test path
     zkClient.createRecursive(path)
-    zkClient.conditionalUpdatePath(path, "version1", 0)
+    zkClient.conditionalUpdatePath(path, "version1".getBytes(UTF_8), 0)
 
     // test with existing path
-    dataAndVersion = zkClient.getDataAndStat(path)
-    assertEquals("version1", dataAndVersion._1.get)
-    assertEquals(1, dataAndVersion._2.getVersion)
-
-    zkClient.conditionalUpdatePath(path, "version2", 1)
-    dataAndVersion = zkClient.getDataAndStat(path)
-    assertEquals("version2", dataAndVersion._1.get)
-    assertEquals(2, dataAndVersion._2.getVersion)
+    val (data1, version1) = zkClient.getDataAndStat(path)
+    assertEquals("version1", new String(data1.get, UTF_8))
+    assertEquals(1, version1.getVersion)
+
+    zkClient.conditionalUpdatePath(path, "version2".getBytes(UTF_8), 1)
+    val (data2, version2) = zkClient.getDataAndStat(path)
+    assertEquals("version2", new String(data2.get, UTF_8))
+    assertEquals(2, version2.getVersion)
   }
 
   @Test
-  def testGetChildren() {
+  def testGetChildren(): Unit = {
     val path = "/testpath"
 
     // test with non-existing path
@@ -240,81 +544,111 @@ class KafkaZkClientTest extends ZooKeeperTestHarness {
   }
 
   @Test
-  def testAclManagementMethods() {
-
-    assertFalse(zkClient.pathExists(AclZNode.path))
-    assertFalse(zkClient.pathExists(AclChangeNotificationZNode.path))
-    ResourceType.values.foreach(resource => assertFalse(zkClient.pathExists(ResourceTypeZNode.path(resource.name))))
+  def testAclManagementMethods(): Unit = {
+    ZkAclStore.stores.foreach(store => {
+      assertFalse(zkClient.pathExists(store.aclPath))
+      assertFalse(zkClient.pathExists(store.changeStore.aclChangePath))
+      AclEntry.ResourceTypes.foreach(resource => assertFalse(zkClient.pathExists(store.path(resource))))
+    })
 
     // create acl paths
-    zkClient.createAclPaths
+    zkClient.createAclPaths()
+
+    ZkAclStore.stores.foreach(store => {
+      assertTrue(zkClient.pathExists(store.aclPath))
+      assertTrue(zkClient.pathExists(store.changeStore.aclChangePath))
+      AclEntry.ResourceTypes.foreach(resource => assertTrue(zkClient.pathExists(store.path(resource))))
 
-    assertTrue(zkClient.pathExists(AclZNode.path))
-    assertTrue(zkClient.pathExists(AclChangeNotificationZNode.path))
-    ResourceType.values.foreach(resource => assertTrue(zkClient.pathExists(ResourceTypeZNode.path(resource.name))))
+      val resource1 = new ResourcePattern(TOPIC, Uuid.randomUuid().toString, store.patternType)
+      val resource2 = new ResourcePattern(TOPIC, Uuid.randomUuid().toString, store.patternType)
 
-    val resource1 = new Resource(Topic, UUID.randomUUID().toString)
-    val resource2 = new Resource(Topic, UUID.randomUUID().toString)
+      // try getting acls for non-existing resource
+      var versionedAcls = zkClient.getVersionedAclsForResource(resource1)
+      assertTrue(versionedAcls.acls.isEmpty)
+      assertEquals(ZkVersion.UnknownVersion, versionedAcls.zkVersion)
+      assertFalse(zkClient.resourceExists(resource1))
 
-    // try getting acls for non-existing resource
-    var versionedAcls = zkClient.getVersionedAclsForResource(resource1)
-    assertTrue(versionedAcls.acls.isEmpty)
-    assertEquals(-1, versionedAcls.zkVersion)
-    assertFalse(zkClient.resourceExists(resource1))
 
+      val acl1 = AclEntry(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "alice"), DENY, "host1" , READ)
+      val acl2 = AclEntry(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "bob"), ALLOW, "*", READ)
+      val acl3 = AclEntry(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "bob"), DENY, "host1", READ)
 
-    val acl1 = new Acl(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "alice"), Deny, "host1" , Read)
-    val acl2 = new Acl(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "bob"), Allow, "*", Read)
-    val acl3 = new Acl(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "bob"), Deny, "host1", Read)
+      // Conditional set should fail if path not created
+      assertFalse(zkClient.conditionalSetAclsForResource(resource1, Set(acl1, acl3), 0)._1)
 
-    //create acls for resources
-    zkClient.conditionalSetOrCreateAclsForResource(resource1, Set(acl1, acl2), 0)
-    zkClient.conditionalSetOrCreateAclsForResource(resource2, Set(acl1, acl3), 0)
+      //create acls for resources
+      assertTrue(zkClient.createAclsForResourceIfNotExists(resource1, Set(acl1, acl2))._1)
+      assertTrue(zkClient.createAclsForResourceIfNotExists(resource2, Set(acl1, acl3))._1)
 
-    versionedAcls = zkClient.getVersionedAclsForResource(resource1)
-    assertEquals(Set(acl1, acl2), versionedAcls.acls)
-    assertEquals(0, versionedAcls.zkVersion)
-    assertTrue(zkClient.resourceExists(resource1))
+      // Create should fail if path already exists
+      assertFalse(zkClient.createAclsForResourceIfNotExists(resource2, Set(acl1, acl3))._1)
 
-    //update acls for resource
-    zkClient.conditionalSetOrCreateAclsForResource(resource1, Set(acl1, acl3), 0)
+      versionedAcls = zkClient.getVersionedAclsForResource(resource1)
+      assertEquals(Set(acl1, acl2), versionedAcls.acls)
+      assertEquals(0, versionedAcls.zkVersion)
+      assertTrue(zkClient.resourceExists(resource1))
 
-    versionedAcls = zkClient.getVersionedAclsForResource(resource1)
-    assertEquals(Set(acl1, acl3), versionedAcls.acls)
-    assertEquals(1, versionedAcls.zkVersion)
+      //update acls for resource
+      assertTrue(zkClient.conditionalSetAclsForResource(resource1, Set(acl1, acl3), 0)._1)
 
-    //get resource Types
-    assertTrue(ResourceType.values.map( rt => rt.name).toSet == zkClient.getResourceTypes().toSet)
+      versionedAcls = zkClient.getVersionedAclsForResource(resource1)
+      assertEquals(Set(acl1, acl3), versionedAcls.acls)
+      assertEquals(1, versionedAcls.zkVersion)
 
-    //get resource name
-    val resourceNames = zkClient.getResourceNames(Topic.name)
-    assertEquals(2, resourceNames.size)
-    assertTrue(Set(resource1.name,resource2.name) == resourceNames.toSet)
+      //get resource Types
+      assertEquals(AclEntry.ResourceTypes.map(SecurityUtils.resourceTypeName), zkClient.getResourceTypes(store.patternType).toSet)
 
-    //delete resource
-    assertTrue(zkClient.deleteResource(resource1))
-    assertFalse(zkClient.resourceExists(resource1))
+      //get resource name
+      val resourceNames = zkClient.getResourceNames(store.patternType, TOPIC)
+      assertEquals(2, resourceNames.size)
+      assertTrue(Set(resource1.name,resource2.name) == resourceNames.toSet)
 
-    //delete with invalid expected zk version
-    assertFalse(zkClient.conditionalDelete(resource2, 10))
-    //delete with valid expected zk version
-    assertTrue(zkClient.conditionalDelete(resource2, 0))
+      //delete resource
+      assertTrue(zkClient.deleteResource(resource1))
+      assertFalse(zkClient.resourceExists(resource1))
 
+      //delete with invalid expected zk version
+      assertFalse(zkClient.conditionalDelete(resource2, 10))
+      //delete with valid expected zk version
+      assertTrue(zkClient.conditionalDelete(resource2, 0))
 
-    zkClient.createAclChangeNotification("resource1")
-    zkClient.createAclChangeNotification("resource2")
+      zkClient.createAclChangeNotification(new ResourcePattern(GROUP, "resource1", store.patternType))
+      zkClient.createAclChangeNotification(new ResourcePattern(TOPIC, "resource2", store.patternType))
 
-    assertEquals(2, zkClient.getChildren(AclChangeNotificationZNode.path).size)
+      assertEquals(2, zkClient.getChildren(store.changeStore.aclChangePath).size)
 
-    zkClient.deleteAclChangeNotifications()
-    assertTrue(zkClient.getChildren(AclChangeNotificationZNode.path).isEmpty)
+      zkClient.deleteAclChangeNotifications()
+      assertTrue(zkClient.getChildren(store.changeStore.aclChangePath).isEmpty)
+    })
   }
 
   @Test
-  def testDeleteTopicPathMethods() {
-    val topic1 = "topic1"
-    val topic2 = "topic2"
+  def testDeletePath(): Unit = {
+    val path = "/a/b/c"
+    zkClient.createRecursive(path)
+    zkClient.deletePath(path)
+    assertFalse(zkClient.pathExists(path))
+
+    zkClient.createRecursive(path)
+    zkClient.deletePath("/a")
+    assertFalse(zkClient.pathExists(path))
+
+    zkClient.createRecursive(path)
+    zkClient.deletePath(path, recursiveDelete =  false)
+    assertFalse(zkClient.pathExists(path))
+    assertTrue(zkClient.pathExists("/a/b"))
+  }
 
+  @Test
+  def testDeleteTopicZNode(): Unit = {
+    zkClient.deleteTopicZNode(topic1, controllerEpochZkVersion)
+    zkClient.createRecursive(TopicZNode.path(topic1))
+    zkClient.deleteTopicZNode(topic1, controllerEpochZkVersion)
+    assertFalse(zkClient.pathExists(TopicZNode.path(topic1)))
+  }
+
+  @Test
+  def testDeleteTopicPathMethods(): Unit = {
     assertFalse(zkClient.isTopicMarkedForDeletion(topic1))
     assertTrue(zkClient.getTopicDeletions.isEmpty)
 
@@ -324,21 +658,32 @@ class KafkaZkClientTest extends ZooKeeperTestHarness {
     assertTrue(zkClient.isTopicMarkedForDeletion(topic1))
     assertEquals(Set(topic1, topic2), zkClient.getTopicDeletions.toSet)
 
-    zkClient.deleteTopicDeletions(Seq(topic1, topic2))
+    intercept[ControllerMovedException](zkClient.deleteTopicDeletions(Seq(topic1, topic2), controllerEpochZkVersion + 1))
+    assertEquals(Set(topic1, topic2), zkClient.getTopicDeletions.toSet)
+
+    zkClient.deleteTopicDeletions(Seq(topic1, topic2), controllerEpochZkVersion)
     assertTrue(zkClient.getTopicDeletions.isEmpty)
   }
 
+  private def assertPathExistenceAndData(expectedPath: String, data: String): Unit = {
+    assertTrue(zkClient.pathExists(expectedPath))
+    assertEquals(Some(data), dataAsString(expectedPath))
+   }
+
   @Test
-  def testEntityConfigManagementMethods() {
-    val topic1 = "topic1"
-    val topic2 = "topic2"
+  def testCreateTokenChangeNotification(): Unit = {
+    intercept[NoNodeException] {
+      zkClient.createTokenChangeNotification("delegationToken")
+    }
+    zkClient.createDelegationTokenPaths()
 
-    assertTrue(zkClient.getEntityConfigs(ConfigType.Topic, topic1).isEmpty)
+    zkClient.createTokenChangeNotification("delegationToken")
+    assertPathExistenceAndData("/delegation_token/token_changes/token_change_0000000000", "delegationToken")
+  }
 
-    val logProps = new Properties()
-    logProps.put(LogConfig.SegmentBytesProp, "1024")
-    logProps.put(LogConfig.SegmentIndexBytesProp, "1024")
-    logProps.put(LogConfig.CleanupPolicyProp, LogConfig.Compact)
+  @Test
+  def testEntityConfigManagementMethods(): Unit = {
+    assertTrue(zkClient.getEntityConfigs(ConfigType.Topic, topic1).isEmpty)
 
     zkClient.setOrCreateEntityConfigs(ConfigType.Topic, topic1, logProps)
     assertEquals(logProps, zkClient.getEntityConfigs(ConfigType.Topic, topic1))
@@ -350,7 +695,647 @@ class KafkaZkClientTest extends ZooKeeperTestHarness {
     zkClient.setOrCreateEntityConfigs(ConfigType.Topic, topic2, logProps)
     assertEquals(Set(topic1, topic2), zkClient.getAllEntitiesWithConfig(ConfigType.Topic).toSet)
 
-    zkClient.deleteTopicConfigs(Seq(topic1, topic2))
+    zkClient.deleteTopicConfigs(Seq(topic1, topic2), controllerEpochZkVersion)
     assertTrue(zkClient.getEntityConfigs(ConfigType.Topic, topic1).isEmpty)
   }
+
+  @Test
+  def testCreateConfigChangeNotification(): Unit = {
+    assertFalse(zkClient.pathExists(ConfigEntityChangeNotificationZNode.path))
+
+    // The parent path is created if needed
+    zkClient.createConfigChangeNotification(ConfigEntityZNode.path(ConfigType.Topic, topic1))
+    assertPathExistenceAndData(
+      "/config/changes/config_change_0000000000",
+      """{"version":2,"entity_path":"/config/topics/topic1"}""")
+
+    // Creation does not fail if the parent path exists
+    zkClient.createConfigChangeNotification(ConfigEntityZNode.path(ConfigType.Topic, topic2))
+    assertPathExistenceAndData(
+      "/config/changes/config_change_0000000001",
+      """{"version":2,"entity_path":"/config/topics/topic2"}""")
+  }
+
+  private def createLogProps(bytesProp: Int): Properties = {
+    val logProps = new Properties()
+    logProps.put(LogConfig.SegmentBytesProp, bytesProp.toString)
+    logProps.put(LogConfig.SegmentIndexBytesProp, bytesProp.toString)
+    logProps.put(LogConfig.CleanupPolicyProp, LogConfig.Compact)
+    logProps
+  }
+
+  private val logProps = createLogProps(1024)
+
+  @Test
+  def testGetLogConfigs(): Unit = {
+    val emptyConfig = LogConfig(Collections.emptyMap())
+    assertEquals("Non existent config, no defaults",
+      (Map(topic1 -> emptyConfig), Map.empty),
+      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(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(Set(topic1, topic2), Collections.emptyMap()))
+
+    val logProps1WithMoreValues = createLogProps(1024)
+    logProps1WithMoreValues.put(LogConfig.SegmentJitterMsProp, "100")
+    logProps1WithMoreValues.put(LogConfig.SegmentBytesProp, "1024")
+
+    assertEquals("Config with defaults",
+      (Map(topic1 -> LogConfig(logProps1WithMoreValues)), Map.empty),
+      zkClient.getLogConfigs(Set(topic1),
+        Map[String, AnyRef](LogConfig.SegmentJitterMsProp -> "100", LogConfig.SegmentBytesProp -> "128").asJava))
+  }
+
+  private def createBrokerInfo(id: Int, host: String, port: Int, securityProtocol: SecurityProtocol,
+                               rack: Option[String] = None,
+                               features: Features[SupportedVersionRange] = emptySupportedFeatures): BrokerInfo =
+    BrokerInfo(
+      Broker(
+        id,
+        Seq(new EndPoint(host, port, ListenerName.forSecurityProtocol(securityProtocol), securityProtocol)),
+        rack = rack,
+        features = features),
+      ApiVersion.latestVersion, jmxPort = port + 10)
+
+  @Test
+  def testRegisterBrokerInfo(): Unit = {
+    zkClient.createTopLevelPaths()
+
+    val brokerInfo = createBrokerInfo(
+      1, "test.host", 9999, SecurityProtocol.PLAINTEXT,
+      rack = None,
+      features = Features.supportedFeatures(
+        Map[String, SupportedVersionRange](
+          "feature1" -> new SupportedVersionRange(1, 2)).asJava))
+    val differentBrokerInfoWithSameId = createBrokerInfo(
+      1, "test.host2", 9995, SecurityProtocol.SSL,
+      features = Features.supportedFeatures(
+        Map[String, SupportedVersionRange](
+          "feature2" -> new SupportedVersionRange(4, 7)).asJava))
+
+    zkClient.registerBroker(brokerInfo)
+    assertEquals(Some(brokerInfo.broker), zkClient.getBroker(1))
+    assertEquals("Other ZK clients can read broker info", Some(brokerInfo.broker), otherZkClient.getBroker(1))
+
+    // Node exists, owned by current session - no error, no update
+    zkClient.registerBroker(differentBrokerInfoWithSameId)
+    assertEquals(Some(brokerInfo.broker), zkClient.getBroker(1))
+
+    // Other client tries to register broker with same id causes failure, info is not changed in ZK
+    intercept[NodeExistsException] {
+      otherZkClient.registerBroker(differentBrokerInfoWithSameId)
+    }
+    assertEquals(Some(brokerInfo.broker), zkClient.getBroker(1))
+  }
+
+  @Test
+  def testRetryRegisterBrokerInfo(): Unit = {
+    val brokerId = 5
+    val brokerPort = 9999
+    val brokerHost = "test.host"
+    val expiredBrokerInfo = createBrokerInfo(brokerId, brokerHost, brokerPort, SecurityProtocol.PLAINTEXT)
+    expiredSessionZkClient.createTopLevelPaths()
+
+    // Register the broker, for the first time
+    expiredSessionZkClient.registerBroker(expiredBrokerInfo)
+    assertEquals(Some(expiredBrokerInfo.broker), expiredSessionZkClient.getBroker(brokerId))
+    val originalCzxid = expiredSessionZkClient.getPathCzxid(BrokerIdZNode.path(brokerId))
+
+    // Here, the node exists already, when trying to register under a different session id,
+    // the node will be deleted and created again using the new session id.
+    expiredSessionZkClient.registerBroker(expiredBrokerInfo)
+
+    // The broker info should be the same, no error should be raised
+    assertEquals(Some(expiredBrokerInfo.broker), expiredSessionZkClient.getBroker(brokerId))
+    val newCzxid = expiredSessionZkClient.getPathCzxid(BrokerIdZNode.path(brokerId))
+
+    assertNotEquals("The Czxid of original ephemeral znode should be different " +
+      "from the new ephemeral znode Czxid", originalCzxid, newCzxid)
+  }
+
+  @Test
+  def testGetBrokerMethods(): Unit = {
+    zkClient.createTopLevelPaths()
+
+    assertEquals(Seq.empty,zkClient.getAllBrokersInCluster)
+    assertEquals(Seq.empty, zkClient.getSortedBrokerList)
+    assertEquals(None, zkClient.getBroker(0))
+
+    val brokerInfo0 = createBrokerInfo(
+      0, "test.host0", 9998, SecurityProtocol.PLAINTEXT,
+      features = Features.supportedFeatures(
+        Map[String, SupportedVersionRange](
+          "feature1" -> new SupportedVersionRange(1, 2)).asJava))
+    val brokerInfo1 = createBrokerInfo(
+      1, "test.host1", 9999, SecurityProtocol.SSL,
+      features = Features.supportedFeatures(
+        Map[String, SupportedVersionRange](
+          "feature2" -> new SupportedVersionRange(3, 6)).asJava))
+
+    zkClient.registerBroker(brokerInfo1)
+    otherZkClient.registerBroker(brokerInfo0)
+
+    assertEquals(Seq(0, 1), zkClient.getSortedBrokerList)
+    assertEquals(
+      Seq(brokerInfo0.broker, brokerInfo1.broker),
+      zkClient.getAllBrokersInCluster
+    )
+    assertEquals(Some(brokerInfo0.broker), zkClient.getBroker(0))
+  }
+
+  @Test
+  def testUpdateBrokerInfo(): Unit = {
+    zkClient.createTopLevelPaths()
+
+    // Updating info of a broker not existing in ZK fails
+    val originalBrokerInfo = createBrokerInfo(1, "test.host", 9999, SecurityProtocol.PLAINTEXT)
+    intercept[NoNodeException]{
+      zkClient.updateBrokerInfo(originalBrokerInfo)
+    }
+
+    zkClient.registerBroker(originalBrokerInfo)
+
+    val updatedBrokerInfo = createBrokerInfo(1, "test.host2", 9995, SecurityProtocol.SSL)
+    zkClient.updateBrokerInfo(updatedBrokerInfo)
+    assertEquals(Some(updatedBrokerInfo.broker), zkClient.getBroker(1))
+
+    // Other ZK clients can update info
+    otherZkClient.updateBrokerInfo(originalBrokerInfo)
+    assertEquals(Some(originalBrokerInfo.broker), otherZkClient.getBroker(1))
+  }
+
+  private def statWithVersion(version: Int): Stat = {
+    val stat = new Stat(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0)
+    stat.setVersion(version)
+    stat
+  }
+
+  private def leaderIsrAndControllerEpochs(state: Int, zkVersion: Int): Map[TopicPartition, LeaderIsrAndControllerEpoch] =
+    Map(
+      topicPartition10 -> LeaderIsrAndControllerEpoch(
+        LeaderAndIsr(leader = 1, leaderEpoch = state, isr = List(2 + state, 3 + state), zkVersion = zkVersion),
+        controllerEpoch = 4),
+      topicPartition11 -> LeaderIsrAndControllerEpoch(
+        LeaderAndIsr(leader = 0, leaderEpoch = state + 1, isr = List(1 + state, 2 + state), zkVersion = zkVersion),
+        controllerEpoch = 4))
+
+  val initialLeaderIsrAndControllerEpochs: Map[TopicPartition, LeaderIsrAndControllerEpoch] =
+    leaderIsrAndControllerEpochs(0, 0)
+
+  val initialLeaderIsrs: Map[TopicPartition, LeaderAndIsr] =
+    initialLeaderIsrAndControllerEpochs.map { case (k, v) => k -> v.leaderAndIsr }
+
+  private def leaderIsrs(state: Int, zkVersion: Int): Map[TopicPartition, LeaderAndIsr] =
+    leaderIsrAndControllerEpochs(state, zkVersion).map { case (k, v) => k -> v.leaderAndIsr }
+
+  private def checkUpdateLeaderAndIsrResult(
+                  expectedSuccessfulPartitions: Map[TopicPartition, LeaderAndIsr],
+                  expectedPartitionsToRetry: Seq[TopicPartition],
+                  expectedFailedPartitions: Map[TopicPartition, (Class[_], String)],
+                  actualUpdateLeaderAndIsrResult: UpdateLeaderAndIsrResult): Unit = {
+    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, successfulPartitions)
+  }
+
+  @Test
+  def testTopicAssignments(): Unit = {
+    val topicId = Uuid.randomUuid()
+    assertEquals(0, zkClient.getPartitionAssignmentForTopics(Set(topicPartition.topic())).size)
+    zkClient.createTopicAssignment(topicPartition.topic(), topicId,
+      Map(topicPartition -> Seq()))
+
+    val expectedAssignment = ReplicaAssignment(Seq(1,2,3), Seq(1), Seq(3))
+    val response = zkClient.setTopicAssignmentRaw(topicPartition.topic(), topicId,
+      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
+  def testUpdateLeaderAndIsr(): Unit = {
+    zkClient.createRecursive(TopicZNode.path(topic1))
+
+    // Non-existing topicPartitions
+    checkUpdateLeaderAndIsrResult(
+        Map.empty,
+        mutable.ArrayBuffer.empty,
+      Map(
+        topicPartition10 -> (classOf[NoNodeException], "KeeperErrorCode = NoNode for /brokers/topics/topic1/partitions/0/state"),
+        topicPartition11 -> (classOf[NoNodeException], "KeeperErrorCode = NoNode for /brokers/topics/topic1/partitions/1/state")),
+      zkClient.updateLeaderAndIsr(initialLeaderIsrs, controllerEpoch = 4, controllerEpochZkVersion))
+
+    zkClient.createTopicPartitionStatesRaw(initialLeaderIsrAndControllerEpochs, controllerEpochZkVersion)
+
+    // Mismatch controller epoch zkVersion
+    intercept[ControllerMovedException](zkClient.updateLeaderAndIsr(initialLeaderIsrs, controllerEpoch = 4, controllerEpochZkVersion + 1))
+
+    // successful updates
+    checkUpdateLeaderAndIsrResult(
+      leaderIsrs(state = 1, zkVersion = 1),
+      mutable.ArrayBuffer.empty,
+      Map.empty,
+      zkClient.updateLeaderAndIsr(leaderIsrs(state = 1, zkVersion = 0),controllerEpoch = 4, controllerEpochZkVersion))
+
+    // Try to update with wrong ZK version
+    checkUpdateLeaderAndIsrResult(
+      Map.empty,
+      ArrayBuffer(topicPartition10, topicPartition11),
+      Map.empty,
+      zkClient.updateLeaderAndIsr(leaderIsrs(state = 1, zkVersion = 0),controllerEpoch = 4, controllerEpochZkVersion))
+
+    // Trigger successful, to be retried and failed partitions in same call
+    val mixedState = Map(
+      topicPartition10 -> LeaderAndIsr(leader = 1, leaderEpoch = 2, isr = List(4, 5), zkVersion = 1),
+      topicPartition11 -> LeaderAndIsr(leader = 0, leaderEpoch = 2, isr = List(3, 4), zkVersion = 0),
+      topicPartition20 -> LeaderAndIsr(leader = 0, leaderEpoch = 2, isr = List(3, 4), zkVersion = 0))
+
+    checkUpdateLeaderAndIsrResult(
+      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")),
+      zkClient.updateLeaderAndIsr(mixedState, controllerEpoch = 4, controllerEpochZkVersion))
+  }
+
+  private def checkGetDataResponse(
+      leaderIsrAndControllerEpochs: Map[TopicPartition,LeaderIsrAndControllerEpoch],
+      topicPartition: TopicPartition,
+      response: GetDataResponse): Unit = {
+    val zkVersion = leaderIsrAndControllerEpochs(topicPartition).leaderAndIsr.zkVersion
+    assertEquals(Code.OK, response.resultCode)
+    assertEquals(TopicPartitionStateZNode.path(topicPartition), response.path)
+    assertEquals(Some(topicPartition), response.ctx)
+    assertEquals(
+      Some(leaderIsrAndControllerEpochs(topicPartition)),
+      TopicPartitionStateZNode.decode(response.data, statWithVersion(zkVersion)))
+  }
+
+  private def eraseMetadata(response: CreateResponse): CreateResponse =
+    response.copy(metadata = ResponseMetadata(0, 0))
+
+  @Test
+  def testGetTopicsAndPartitions(): Unit = {
+    assertTrue(zkClient.getAllTopicsInCluster().isEmpty)
+    assertTrue(zkClient.getAllPartitions.isEmpty)
+
+    zkClient.createRecursive(TopicZNode.path(topic1))
+    zkClient.createRecursive(TopicZNode.path(topic2))
+    assertEquals(Set(topic1, topic2), zkClient.getAllTopicsInCluster())
+
+    assertTrue(zkClient.getAllPartitions.isEmpty)
+
+    zkClient.createTopicPartitionStatesRaw(initialLeaderIsrAndControllerEpochs, controllerEpochZkVersion)
+    assertEquals(Set(topicPartition10, topicPartition11), zkClient.getAllPartitions)
+  }
+
+  @Test
+  def testCreateAndGetTopicPartitionStatesRaw(): Unit = {
+    zkClient.createRecursive(TopicZNode.path(topic1))
+
+    // Mismatch controller epoch zkVersion
+    intercept[ControllerMovedException](zkClient.createTopicPartitionStatesRaw(initialLeaderIsrAndControllerEpochs, controllerEpochZkVersion + 1))
+
+    assertEquals(
+      Seq(
+        CreateResponse(Code.OK, TopicPartitionStateZNode.path(topicPartition10), Some(topicPartition10),
+          TopicPartitionStateZNode.path(topicPartition10), ResponseMetadata(0, 0)),
+        CreateResponse(Code.OK, TopicPartitionStateZNode.path(topicPartition11), Some(topicPartition11),
+          TopicPartitionStateZNode.path(topicPartition11), ResponseMetadata(0, 0))),
+      zkClient.createTopicPartitionStatesRaw(initialLeaderIsrAndControllerEpochs, controllerEpochZkVersion)
+        .map(eraseMetadata).toList)
+
+    val getResponses = zkClient.getTopicPartitionStatesRaw(topicPartitions10_11)
+    assertEquals(2, getResponses.size)
+    topicPartitions10_11.zip(getResponses) foreach {case (tp, r) => checkGetDataResponse(initialLeaderIsrAndControllerEpochs, tp, r)}
+
+    // Trying to create existing topicPartition states fails
+    assertEquals(
+      Seq(
+        CreateResponse(Code.NODEEXISTS, TopicPartitionStateZNode.path(topicPartition10), Some(topicPartition10), null, ResponseMetadata(0, 0)),
+        CreateResponse(Code.NODEEXISTS, TopicPartitionStateZNode.path(topicPartition11), Some(topicPartition11), null, ResponseMetadata(0, 0))),
+      zkClient.createTopicPartitionStatesRaw(initialLeaderIsrAndControllerEpochs, controllerEpochZkVersion).map(eraseMetadata).toList)
+  }
+
+  @Test
+  def testSetTopicPartitionStatesRaw(): Unit = {
+
+    def expectedSetDataResponses(topicPartitions: TopicPartition*)(resultCode: Code, stat: Stat) =
+      topicPartitions.map { topicPartition =>
+        SetDataResponse(resultCode, TopicPartitionStateZNode.path(topicPartition),
+          Some(topicPartition), stat, ResponseMetadata(0, 0))
+      }
+
+    zkClient.createRecursive(TopicZNode.path(topic1))
+
+    // Trying to set non-existing topicPartition's data results in NONODE responses
+    assertEquals(
+      expectedSetDataResponses(topicPartition10, topicPartition11)(Code.NONODE, null),
+      zkClient.setTopicPartitionStatesRaw(initialLeaderIsrAndControllerEpochs, controllerEpochZkVersion).map {
+        _.copy(metadata = ResponseMetadata(0, 0))}.toList)
+
+    zkClient.createTopicPartitionStatesRaw(initialLeaderIsrAndControllerEpochs, controllerEpochZkVersion)
+
+    assertEquals(
+      expectedSetDataResponses(topicPartition10, topicPartition11)(Code.OK, statWithVersion(1)),
+      zkClient.setTopicPartitionStatesRaw(leaderIsrAndControllerEpochs(state = 1, zkVersion = 0), controllerEpochZkVersion).map {
+        eraseMetadataAndStat}.toList)
+
+    // Mismatch controller epoch zkVersion
+    intercept[ControllerMovedException](zkClient.setTopicPartitionStatesRaw(leaderIsrAndControllerEpochs(state = 1, zkVersion = 0), controllerEpochZkVersion + 1))
+
+    val getResponses = zkClient.getTopicPartitionStatesRaw(topicPartitions10_11)
+    assertEquals(2, getResponses.size)
+    topicPartitions10_11.zip(getResponses) foreach {case (tp, r) => checkGetDataResponse(leaderIsrAndControllerEpochs(state = 1, zkVersion = 0), tp, r)}
+
+    // Other ZK client can also write the state of a partition
+    assertEquals(
+      expectedSetDataResponses(topicPartition10, topicPartition11)(Code.OK, statWithVersion(2)),
+      otherZkClient.setTopicPartitionStatesRaw(leaderIsrAndControllerEpochs(state = 2, zkVersion = 1), controllerEpochZkVersion).map {
+        eraseMetadataAndStat}.toList)
+  }
+
+  @Test
+  def testReassignPartitionsInProgress(): Unit = {
+    assertFalse(zkClient.reassignPartitionsInProgress)
+    zkClient.createRecursive(ReassignPartitionsZNode.path)
+    assertTrue(zkClient.reassignPartitionsInProgress)
+  }
+
+  @Test
+  def testGetTopicPartitionStates(): Unit = {
+    assertEquals(None, zkClient.getTopicPartitionState(topicPartition10))
+    assertEquals(None, zkClient.getLeaderForPartition(topicPartition10))
+
+    zkClient.createRecursive(TopicZNode.path(topic1))
+
+    zkClient.createTopicPartitionStatesRaw(initialLeaderIsrAndControllerEpochs, controllerEpochZkVersion)
+    assertEquals(
+      initialLeaderIsrAndControllerEpochs,
+      zkClient.getTopicPartitionStates(Seq(topicPartition10, topicPartition11))
+    )
+
+    assertEquals(
+      Some(initialLeaderIsrAndControllerEpochs(topicPartition10)),
+      zkClient.getTopicPartitionState(topicPartition10)
+    )
+
+    assertEquals(Some(1), zkClient.getLeaderForPartition(topicPartition10))
+
+    val notExistingPartition = new TopicPartition(topic1, 2)
+    assertTrue(zkClient.getTopicPartitionStates(Seq(notExistingPartition)).isEmpty)
+    assertEquals(
+      Map(topicPartition10 -> initialLeaderIsrAndControllerEpochs(topicPartition10)),
+      zkClient.getTopicPartitionStates(Seq(topicPartition10, notExistingPartition))
+    )
+
+    assertEquals(None, zkClient.getTopicPartitionState(notExistingPartition))
+    assertEquals(None, zkClient.getLeaderForPartition(notExistingPartition))
+
+  }
+
+  private def eraseMetadataAndStat(response: SetDataResponse): SetDataResponse = {
+    val stat = if (response.stat != null) statWithVersion(response.stat.getVersion) else null
+    response.copy(metadata = ResponseMetadata(0, 0), stat = stat)
+  }
+
+  @Test
+  def testControllerEpochMethods(): Unit = {
+    zkClient.deletePath(ControllerEpochZNode.path)
+
+    assertEquals(None, zkClient.getControllerEpoch)
+
+    assertEquals("Setting non existing nodes should return NONODE results",
+      SetDataResponse(Code.NONODE, ControllerEpochZNode.path, None, null, ResponseMetadata(0, 0)),
+      eraseMetadataAndStat(zkClient.setControllerEpochRaw(1, 0)))
+
+    assertEquals("Creating non existing nodes is OK",
+      CreateResponse(Code.OK, ControllerEpochZNode.path, None, ControllerEpochZNode.path, ResponseMetadata(0, 0)),
+      eraseMetadata(zkClient.createControllerEpochRaw(0)))
+    assertEquals(0, zkClient.getControllerEpoch.get._1)
+
+    assertEquals("Attemt to create existing nodes should return NODEEXISTS",
+      CreateResponse(Code.NODEEXISTS, ControllerEpochZNode.path, None, null, ResponseMetadata(0, 0)),
+      eraseMetadata(zkClient.createControllerEpochRaw(0)))
+
+    assertEquals("Updating existing nodes is OK",
+      SetDataResponse(Code.OK, ControllerEpochZNode.path, None, statWithVersion(1), ResponseMetadata(0, 0)),
+      eraseMetadataAndStat(zkClient.setControllerEpochRaw(1, 0)))
+    assertEquals(1, zkClient.getControllerEpoch.get._1)
+
+    assertEquals("Updating with wrong ZK version returns BADVERSION",
+      SetDataResponse(Code.BADVERSION, ControllerEpochZNode.path, None, null, ResponseMetadata(0, 0)),
+      eraseMetadataAndStat(zkClient.setControllerEpochRaw(1, 0)))
+  }
+
+  @Test
+  def testControllerManagementMethods(): Unit = {
+    // No controller
+    assertEquals(None, zkClient.getControllerId)
+    // Create controller
+    val (_, newEpochZkVersion) = zkClient.registerControllerAndIncrementControllerEpoch(controllerId = 1)
+    assertEquals(Some(1), zkClient.getControllerId)
+    zkClient.deleteController(newEpochZkVersion)
+    assertEquals(None, zkClient.getControllerId)
+  }
+
+  @Test
+  def testZNodeChangeHandlerForDataChange(): Unit = {
+    val mockPath = "/foo"
+
+    val znodeChangeHandlerCountDownLatch = new CountDownLatch(1)
+    val zNodeChangeHandler = new ZNodeChangeHandler {
+      override def handleCreation(): Unit = {
+        znodeChangeHandlerCountDownLatch.countDown()
+      }
+
+      override val path: String = mockPath
+    }
+
+    zkClient.registerZNodeChangeHandlerAndCheckExistence(zNodeChangeHandler)
+    zkClient.createRecursive(mockPath)
+    assertTrue("Failed to receive create notification", znodeChangeHandlerCountDownLatch.await(5, TimeUnit.SECONDS))
+  }
+
+  @Test
+  def testClusterIdMethods(): Unit = {
+    val clusterId = CoreUtils.generateUuidAsBase64()
+
+    zkClient.createOrGetClusterId(clusterId)
+    assertEquals(clusterId, zkClient.getClusterId.getOrElse(fail("No cluster id found")))
+   }
+
+  @Test
+  def testBrokerSequenceIdMethods(): Unit = {
+    val sequenceId = zkClient.generateBrokerSequenceId()
+    assertEquals(sequenceId + 1, zkClient.generateBrokerSequenceId())
+  }
+
+  @Test
+  def testCreateTopLevelPaths(): Unit = {
+    zkClient.createTopLevelPaths()
+
+    ZkData.PersistentZkPaths.foreach(path => assertTrue(zkClient.pathExists(path)))
+  }
+
+  @Test
+  def testPreferredReplicaElectionMethods(): Unit = {
+
+    assertTrue(zkClient.getPreferredReplicaElection.isEmpty)
+
+    val electionPartitions = Set(new TopicPartition(topic1, 0), new TopicPartition(topic1, 1))
+
+    zkClient.createPreferredReplicaElection(electionPartitions)
+    assertEquals(electionPartitions, zkClient.getPreferredReplicaElection)
+
+    intercept[NodeExistsException] {
+      zkClient.createPreferredReplicaElection(electionPartitions)
+    }
+
+    // Mismatch controller epoch zkVersion
+    intercept[ControllerMovedException](zkClient.deletePreferredReplicaElection(controllerEpochZkVersion + 1))
+    assertEquals(electionPartitions, zkClient.getPreferredReplicaElection)
+
+    zkClient.deletePreferredReplicaElection(controllerEpochZkVersion)
+    assertTrue(zkClient.getPreferredReplicaElection.isEmpty)
+  }
+
+  private def dataAsString(path: String): Option[String] = {
+    val (data, _) = zkClient.getDataAndStat(path)
+    data.map(new String(_, UTF_8))
+  }
+
+  @Test
+  def testDelegationTokenMethods(): Unit = {
+    assertFalse(zkClient.pathExists(DelegationTokensZNode.path))
+    assertFalse(zkClient.pathExists(DelegationTokenChangeNotificationZNode.path))
+
+    zkClient.createDelegationTokenPaths()
+    assertTrue(zkClient.pathExists(DelegationTokensZNode.path))
+    assertTrue(zkClient.pathExists(DelegationTokenChangeNotificationZNode.path))
+
+    val tokenId = "token1"
+    val owner = SecurityUtils.parseKafkaPrincipal("User:owner1")
+    val renewers = List(SecurityUtils.parseKafkaPrincipal("User:renewer1"), SecurityUtils.parseKafkaPrincipal("User:renewer1"))
+
+    val tokenInfo = new TokenInformation(tokenId, owner, renewers.asJava,
+      System.currentTimeMillis(), System.currentTimeMillis(), System.currentTimeMillis())
+    val bytes = new Array[Byte](20)
+    Random.nextBytes(bytes)
+    val token = new org.apache.kafka.common.security.token.delegation.DelegationToken(tokenInfo, bytes)
+
+    // test non-existent token
+    assertTrue(zkClient.getDelegationTokenInfo(tokenId).isEmpty)
+    assertFalse(zkClient.deleteDelegationToken(tokenId))
+
+    // create a token
+    zkClient.setOrCreateDelegationToken(token)
+
+    //get created token
+    assertEquals(tokenInfo, zkClient.getDelegationTokenInfo(tokenId).get)
+
+    //update expiryTime
+    tokenInfo.setExpiryTimestamp(System.currentTimeMillis())
+    zkClient.setOrCreateDelegationToken(token)
+
+    //test updated token
+    assertEquals(tokenInfo, zkClient.getDelegationTokenInfo(tokenId).get)
+
+    //test deleting token
+    assertTrue(zkClient.deleteDelegationToken(tokenId))
+    assertEquals(None, zkClient.getDelegationTokenInfo(tokenId))
+  }
+
+  @Test
+  def testConsumerOffsetPath(): Unit = {
+    def getConsumersOffsetsZkPath(consumerGroup: String, topic: String, partition: Int): String = {
+      s"/consumers/$consumerGroup/offsets/$topic/$partition"
+    }
+
+    val consumerGroup = "test-group"
+    val topic = "test-topic"
+    val partition = 2
+
+    val expectedConsumerGroupOffsetsPath = getConsumersOffsetsZkPath(consumerGroup, topic, partition)
+    val actualConsumerGroupOffsetsPath = ConsumerOffset.path(consumerGroup, topic, partition)
+
+    assertEquals(expectedConsumerGroupOffsetsPath, actualConsumerGroupOffsetsPath)
+  }
+
+  @Test
+  def testAclMethods(): Unit = {
+    val mockPath = "/foo"
+
+    intercept[NoNodeException] {
+      zkClient.getAcl(mockPath)
+    }
+
+    intercept[NoNodeException] {
+      zkClient.setAcl(mockPath, ZooDefs.Ids.OPEN_ACL_UNSAFE.asScala)
+    }
+
+    zkClient.createRecursive(mockPath)
+
+    zkClient.setAcl(mockPath, ZooDefs.Ids.READ_ACL_UNSAFE.asScala)
+
+    assertEquals(ZooDefs.Ids.READ_ACL_UNSAFE.asScala, zkClient.getAcl(mockPath))
+  }
+
+  class ExpiredKafkaZkClient private (zooKeeperClient: ZooKeeperClient, isSecure: Boolean, time: Time)
+    extends KafkaZkClient(zooKeeperClient, isSecure, time) {
+    // Overwriting this method from the parent class to force the client to re-register the Broker.
+    override def shouldReCreateEphemeralZNode(ephemeralOwnerId: Long): Boolean = {
+      true
+    }
+
+    def getPathCzxid(path: String): Long = {
+      val getDataRequest = GetDataRequest(path)
+      val getDataResponse = retryRequestUntilConnected(getDataRequest)
+
+      getDataResponse.stat.getCzxid
+    }
+  }
+
+  private object ExpiredKafkaZkClient {
+    def apply(connectString: String,
+              isSecure: Boolean,
+              sessionTimeoutMs: Int,
+              connectionTimeoutMs: Int,
+              maxInFlightRequests: Int,
+              time: Time,
+              metricGroup: String = "kafka.server",
+              metricType: String = "SessionExpireListener") = {
+      val zooKeeperClient = new ZooKeeperClient(connectString, sessionTimeoutMs, connectionTimeoutMs, maxInFlightRequests,
+        time, metricGroup, metricType)
+      new ExpiredKafkaZkClient(zooKeeperClient, isSecure, time)
+    }
+  }
 }
diff --git a/core/src/test/scala/unit/kafka/zk/ReassignPartitionsZNodeTest.scala b/core/src/test/scala/unit/kafka/zk/ReassignPartitionsZNodeTest.scala
new file mode 100644
index 0000000000000..5b9e3e4d95f77
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/zk/ReassignPartitionsZNodeTest.scala
@@ -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 kafka.zk
+
+import java.nio.charset.StandardCharsets
+
+import com.fasterxml.jackson.core.JsonProcessingException
+import org.apache.kafka.common.TopicPartition
+import org.junit.Assert._
+import org.junit.Test
+import org.scalatest.Assertions
+
+class ReassignPartitionsZNodeTest {
+
+  private val topic = "foo"
+  private val partition1 = 0
+  private val replica1 = 1
+  private val replica2 = 2
+
+  private val reassignPartitionData = Map(new TopicPartition(topic, partition1) -> Seq(replica1, replica2))
+  private val reassignmentJson = """{"version":1,"partitions":[{"topic":"foo","partition":0,"replicas":[1,2]}]}"""
+
+  @Test
+  def testEncode(): Unit = {
+    val encodedJsonString = new String(ReassignPartitionsZNode.encode(reassignPartitionData), StandardCharsets.UTF_8)
+    assertEquals(reassignmentJson, encodedJsonString)
+  }
+
+  @Test
+  def testDecodeInvalidJson(): Unit = {
+    val result = ReassignPartitionsZNode.decode("invalid json".getBytes)
+    val exception = result.left.getOrElse(Assertions.fail(s"decode should have failed, result $result"))
+    assertTrue(exception.isInstanceOf[JsonProcessingException])
+  }
+
+  @Test
+  def testDecodeValidJson(): Unit = {
+    val result = ReassignPartitionsZNode.decode(reassignmentJson.getBytes)
+    val replicas = result.map(assignmentMap => assignmentMap(new TopicPartition(topic, partition1)))
+    assertEquals(Right(Seq(replica1, replica2)), replicas)
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/zk/ZKEphemeralTest.scala b/core/src/test/scala/unit/kafka/zk/ZKEphemeralTest.scala
deleted file mode 100644
index 6280d97b23ce5..0000000000000
--- a/core/src/test/scala/unit/kafka/zk/ZKEphemeralTest.scala
+++ /dev/null
@@ -1,176 +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.zk
-
-import java.lang.Iterable
-import javax.security.auth.login.Configuration
-
-import scala.collection.JavaConverters._
-
-import kafka.consumer.ConsumerConfig
-import kafka.utils.ZkUtils
-import kafka.utils.ZKCheckedEphemeral
-import kafka.utils.TestUtils
-import org.apache.kafka.common.security.JaasUtils
-import org.apache.zookeeper.CreateMode
-import org.apache.zookeeper.WatchedEvent
-import org.apache.zookeeper.Watcher
-import org.apache.zookeeper.ZooDefs.Ids
-import org.I0Itec.zkclient.exception.ZkNodeExistsException
-import org.junit.{After, Before, Test, Assert}
-import org.junit.runners.Parameterized
-import org.junit.runners.Parameterized.Parameters
-import org.junit.runner.RunWith
-
-object ZKEphemeralTest {
-
-  @Parameters
-  def enableSecurityOptions: Iterable[Array[java.lang.Boolean]] =
-    Seq[Array[java.lang.Boolean]](Array(true), Array(false)).asJava
-
-}
-
-@RunWith(value = classOf[Parameterized])
-class ZKEphemeralTest(val secure: Boolean) extends ZooKeeperTestHarness {
-  val jaasFile = kafka.utils.JaasTestUtils.writeJaasContextsToFile(kafka.utils.JaasTestUtils.zkSections)
-  val authProvider = "zookeeper.authProvider.1"
-  var zkSessionTimeoutMs = 1000
-  
-  @Before
-  override def setUp() {
-    if (secure) {
-      System.setProperty(JaasUtils.JAVA_LOGIN_CONFIG_PARAM, jaasFile.getAbsolutePath)
-      Configuration.setConfiguration(null)
-      System.setProperty(authProvider, "org.apache.zookeeper.server.auth.SASLAuthenticationProvider")
-      if (!JaasUtils.isZkSecurityEnabled)
-        fail("Secure access not enabled")
-    }
-    super.setUp
-  }
-  
-  @After
-  override def tearDown() {
-    super.tearDown
-    System.clearProperty(JaasUtils.JAVA_LOGIN_CONFIG_PARAM)
-    System.clearProperty(authProvider)
-    Configuration.setConfiguration(null)
-  }
-  
-  @Test
-  def testEphemeralNodeCleanup(): Unit = {
-    val config = new ConsumerConfig(TestUtils.createConsumerProperties(zkConnect, "test", "1"))
-    var zkUtils = ZkUtils(zkConnect, zkSessionTimeoutMs, config.zkConnectionTimeoutMs, JaasUtils.isZkSecurityEnabled())
-
-    try {
-      zkUtils.createEphemeralPathExpectConflict("/tmp/zktest", "node created")
-    } catch {                       
-      case _: Exception =>
-    }
-
-    var testData: String = null
-    testData = zkUtils.readData("/tmp/zktest")._1
-    Assert.assertNotNull(testData)
-    zkUtils.close
-    zkUtils = ZkUtils(zkConnect, zkSessionTimeoutMs, config.zkConnectionTimeoutMs, JaasUtils.isZkSecurityEnabled())
-    val nodeExists = zkUtils.pathExists("/tmp/zktest")
-    Assert.assertFalse(nodeExists)
-    zkUtils.close()
-  }
-
-  /*****
-   ***** Tests for ZkWatchedEphemeral
-   *****/
-  
-  /**
-   * Tests basic creation
-   */
-  @Test
-  def testZkWatchedEphemeral(): Unit = {
-    testCreation("/zwe-test")
-    testCreation("/zwe-test-parent/zwe-test")
-  }
- 
-  private def testCreation(path: String) {
-    val zk = zkUtils.zkConnection.getZookeeper
-    val zwe = new ZKCheckedEphemeral(path, "", zk, JaasUtils.isZkSecurityEnabled())
-    var created = false
-
-    zk.exists(path, new Watcher() {
-      def process(event: WatchedEvent) {
-        if(event.getType == Watcher.Event.EventType.NodeCreated) {
-          created = true
-        }
-      }
-    })
-    zwe.create()
-    // Waits until the znode is created
-    TestUtils.waitUntilTrue(() => zkUtils.pathExists(path),
-                            s"Znode $path wasn't created")
-  }
-
-  /**
-   * Tests that it fails in the presence of an overlapping
-   * session.
-   */
-  @Test
-  def testOverlappingSessions(): Unit = {
-    val path = "/zwe-test"
-    val zk1 = zkUtils.zkConnection.getZookeeper
-
-    //Creates a second session
-    val (zkClient2, zkConnection2) = ZkUtils.createZkClientAndConnection(zkConnect, zkSessionTimeoutMs, zkConnectionTimeout)
-    val zk2 = zkConnection2.getZookeeper
-    val zwe = new ZKCheckedEphemeral(path, "", zk2, JaasUtils.isZkSecurityEnabled())
-
-    // Creates znode for path in the first session
-    zk1.create(path, Array[Byte](), Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL)
-    
-    //Bootstraps the ZKWatchedEphemeral object
-    val gotException =
-      try {
-        zwe.create()
-        false
-      } catch {
-        case _: ZkNodeExistsException => true
-      }
-    Assert.assertTrue(gotException)
-    zkClient2.close()
-  }
-  
-  /**
-   * Tests if succeeds with znode from the same session
-   */
-  @Test
-  def testSameSession(): Unit = {
-    val path = "/zwe-test"
-    val zk = zkUtils.zkConnection.getZookeeper
-    // Creates znode for path in the first session
-    zk.create(path, Array[Byte](), Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL)
-    
-    val zwe = new ZKCheckedEphemeral(path, "", zk, JaasUtils.isZkSecurityEnabled())
-    //Bootstraps the ZKWatchedEphemeral object
-    val gotException =
-      try {
-        zwe.create()
-        false
-      } catch {
-        case _: ZkNodeExistsException => true
-      }
-    Assert.assertFalse(gotException)
-  }
-}
diff --git a/core/src/test/scala/unit/kafka/zk/ZKPathTest.scala b/core/src/test/scala/unit/kafka/zk/ZKPathTest.scala
deleted file mode 100644
index 0cf836df68199..0000000000000
--- a/core/src/test/scala/unit/kafka/zk/ZKPathTest.scala
+++ /dev/null
@@ -1,130 +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.zk
-
-import kafka.consumer.ConsumerConfig
-import kafka.utils.{TestUtils, ZkUtils}
-import org.apache.kafka.common.config.ConfigException
-import org.junit.Assert._
-import org.junit.Test
-
-class ZKPathTest extends ZooKeeperTestHarness {
-
-  val path = "/some_dir"
-  val zkSessionTimeoutMs = 1000
-  def zkConnectWithInvalidRoot: String = zkConnect + "/ghost"
-
-  @Test
-  def testCreatePersistentPathThrowsException(): Unit = {
-    val config = new ConsumerConfig(TestUtils.createConsumerProperties(zkConnectWithInvalidRoot,
-      "test", "1"))
-    val zkUtils = ZkUtils(zkConnectWithInvalidRoot, zkSessionTimeoutMs,
-      config.zkConnectionTimeoutMs, false)
-    try {
-      zkUtils.zkPath.resetNamespaceCheckedState
-      zkUtils.createPersistentPath(path)
-      fail("Failed to throw ConfigException for missing ZooKeeper root node")
-    } catch {
-      case _: ConfigException =>
-    }
-    zkUtils.close()
-  }
-
-  @Test
-  def testCreatePersistentPath(): Unit = {
-    val config = new ConsumerConfig(TestUtils.createConsumerProperties(zkConnect, "test", "1"))
-    val zkUtils = ZkUtils(zkConnect, zkSessionTimeoutMs, config.zkConnectionTimeoutMs, false)
-    zkUtils.zkPath.resetNamespaceCheckedState
-    zkUtils.createPersistentPath(path)
-    assertTrue("Failed to create persistent path", zkUtils.pathExists(path))
-    zkUtils.close()
-  }
-
-  @Test
-  def testMakeSurePersistsPathExistsThrowsException(): Unit = {
-    val config = new ConsumerConfig(TestUtils.createConsumerProperties(zkConnectWithInvalidRoot, "test", "1"))
-    val zkUtils = ZkUtils(zkConnectWithInvalidRoot, zkSessionTimeoutMs, config.zkConnectionTimeoutMs, false)
-    try {
-      zkUtils.zkPath.resetNamespaceCheckedState
-      zkUtils.makeSurePersistentPathExists(path)
-      fail("Failed to throw ConfigException for missing ZooKeeper root node")
-    } catch {
-      case _: ConfigException =>
-    }
-    zkUtils.close()
-  }
-
-  @Test
-  def testMakeSurePersistsPathExists(): Unit = {
-    val config = new ConsumerConfig(TestUtils.createConsumerProperties(zkConnect, "test", "1"))
-    val zkUtils = ZkUtils(zkConnect, zkSessionTimeoutMs, config.zkConnectionTimeoutMs, false)
-    zkUtils.zkPath.resetNamespaceCheckedState
-    zkUtils.makeSurePersistentPathExists(path)
-    assertTrue("Failed to create persistent path", zkUtils.pathExists(path))
-    zkUtils.close()
-  }
-
-  @Test
-  def testCreateEphemeralPathThrowsException(): Unit = {
-    val config = new ConsumerConfig(TestUtils.createConsumerProperties(zkConnectWithInvalidRoot, "test", "1"))
-    val zkUtils = ZkUtils(zkConnectWithInvalidRoot, zkSessionTimeoutMs, config.zkConnectionTimeoutMs, false)
-    try {
-      zkUtils.zkPath.resetNamespaceCheckedState
-      zkUtils.createEphemeralPathExpectConflict(path, "somedata")
-      fail("Failed to throw ConfigException for missing ZooKeeper root node")
-    } catch {
-      case _: ConfigException =>
-    }
-    zkUtils.close()
-  }
-
-  @Test
-  def testCreateEphemeralPathExists(): Unit = {
-    val config = new ConsumerConfig(TestUtils.createConsumerProperties(zkConnect, "test", "1"))
-    val zkUtils = ZkUtils(zkConnect, zkSessionTimeoutMs, config.zkConnectionTimeoutMs, false)
-    zkUtils.zkPath.resetNamespaceCheckedState
-    zkUtils.createEphemeralPathExpectConflict(path, "somedata")
-    assertTrue("Failed to create ephemeral path", zkUtils.pathExists(path))
-    zkUtils.close()
-  }
-
-  @Test
-  def testCreatePersistentSequentialThrowsException(): Unit = {
-    val config = new ConsumerConfig(TestUtils.createConsumerProperties(zkConnectWithInvalidRoot,
-      "test", "1"))
-    val zkUtils = ZkUtils(zkConnectWithInvalidRoot, zkSessionTimeoutMs, config.zkConnectionTimeoutMs, false)
-    try {
-      zkUtils.zkPath.resetNamespaceCheckedState
-      zkUtils.createSequentialPersistentPath(path)
-      fail("Failed to throw ConfigException for missing ZooKeeper root node")
-    } catch {
-      case _: ConfigException =>
-    }
-    zkUtils.close()
-  }
-
-  @Test
-  def testCreatePersistentSequentialExists(): Unit = {
-    val config = new ConsumerConfig(TestUtils.createConsumerProperties(zkConnect, "test", "1"))
-    val zkUtils = ZkUtils(zkConnect, zkSessionTimeoutMs, config.zkConnectionTimeoutMs, false)
-    zkUtils.zkPath.resetNamespaceCheckedState
-    val actualPath = zkUtils.createSequentialPersistentPath(path)
-    assertTrue("Failed to create persistent path", zkUtils.pathExists(actualPath))
-    zkUtils.close()
-  }
-}
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 4966b10b50f05..dfd5e002b569c 100755
--- a/core/src/test/scala/unit/kafka/zk/ZooKeeperTestHarness.scala
+++ b/core/src/test/scala/unit/kafka/zk/ZooKeeperTestHarness.scala
@@ -19,32 +19,31 @@ package kafka.zk
 
 import javax.security.auth.login.Configuration
 
-import kafka.utils.{CoreUtils, Logging, TestUtils, ZkUtils}
+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
 
 import scala.collection.Set
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 import org.apache.kafka.clients.producer.KafkaProducer
 import org.apache.kafka.clients.consumer.internals.AbstractCoordinator
 import kafka.controller.ControllerEventManager
-import kafka.zookeeper.ZooKeeperClient
+import org.apache.kafka.clients.admin.AdminClientUnitTestEnv
+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 = 6000
+  val zkSessionTimeout = 15000 // Allows us to avoid ZK session expiration due to GC up to 2/3 * 15000ms = 10 secs
   val zkMaxInFlightRequests = Int.MaxValue
 
-  protected val zkAclsEnabled: Option[Boolean] = None
+  protected def zkAclsEnabled: Option[Boolean] = None
 
-  var zkUtils: ZkUtils = null
-  var zooKeeperClient: ZooKeeperClient = null
   var zkClient: KafkaZkClient = null
   var adminZkClient: AdminZkClient = null
 
@@ -54,29 +53,37 @@ 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()
-    zkUtils = ZkUtils(zkConnect, zkSessionTimeout, zkConnectionTimeout, zkAclsEnabled.getOrElse(JaasUtils.isZkSecurityEnabled()))
-
-    zooKeeperClient = new ZooKeeperClient(zkConnect, zkSessionTimeout, zkConnectionTimeout, zkMaxInFlightRequests)
-    zkClient = new KafkaZkClient(zooKeeperClient, zkAclsEnabled.getOrElse(JaasUtils.isZkSecurityEnabled()))
+    zkClient = KafkaZkClient(zkConnect, zkAclsEnabled.getOrElse(JaasUtils.isZkSaslEnabled), zkSessionTimeout,
+      zkConnectionTimeout, zkMaxInFlightRequests, Time.SYSTEM)
     adminZkClient = new AdminZkClient(zkClient)
   }
 
   @After
-  def tearDown() {
-    if (zkUtils != null)
-     CoreUtils.swallow(zkUtils.close(), this)
+  def tearDown(): Unit = {
     if (zkClient != null)
      zkClient.close()
     if (zookeeper != null)
       CoreUtils.swallow(zookeeper.shutdown(), this)
     Configuration.setConfiguration(null)
   }
+
+  // Trigger session expiry by reusing the session id in another client
+  def createZooKeeperClientToTriggerSessionExpiry(zooKeeper: ZooKeeper): ZooKeeper = {
+    val dummyWatcher = new Watcher {
+      override def process(event: WatchedEvent): Unit = {}
+    }
+    val anotherZkClient = new ZooKeeper(zkConnect, 1000, dummyWatcher,
+      zooKeeper.getSessionId,
+      zooKeeper.getSessionPasswd)
+    assertNull(anotherZkClient.exists("/nonexistent", false)) // Make sure new client works
+    anotherZkClient
+  }
 }
 
 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
@@ -84,8 +91,9 @@ object ZooKeeperTestHarness {
   // which reset static JAAS configuration.
   val unexpectedThreadNames = Set(ControllerEventManager.ControllerEventThreadName,
                                   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.
@@ -93,7 +101,7 @@ object ZooKeeperTestHarness {
    * which is true for core tests where this harness is used.
    */
   @BeforeClass
-  def setUpClass() {
+  def setUpClass(): Unit = {
     verifyNoUnexpectedThreads("@BeforeClass")
   }
 
@@ -101,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")
   }
 
@@ -109,11 +117,15 @@ 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)))
     }
-    assertTrue(s"Found unexpected threads during $context, allThreads=$threads", noUnexpected)
+    assertTrue(
+      s"Found unexpected threads during $context, allThreads=$threads, " +
+        s"unexpected=${threads.filterNot(t => unexpectedThreadNames.forall(s => !t.contains(s)))}",
+      noUnexpected
+    )
   }
 }
diff --git a/core/src/test/scala/unit/kafka/zookeeper/ZooKeeperClientTest.scala b/core/src/test/scala/unit/kafka/zookeeper/ZooKeeperClientTest.scala
index ea6a47518a6ed..a6200889cc397 100644
--- a/core/src/test/scala/unit/kafka/zookeeper/ZooKeeperClientTest.scala
+++ b/core/src/test/scala/unit/kafka/zookeeper/ZooKeeperClientTest.scala
@@ -16,56 +16,128 @@
  */
 package kafka.zookeeper
 
-import java.net.UnknownHostException
 import java.nio.charset.StandardCharsets
 import java.util.UUID
-import java.util.concurrent.atomic.AtomicBoolean
-import java.util.concurrent.{ArrayBlockingQueue, CountDownLatch, TimeUnit}
-import javax.security.auth.login.Configuration
+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.core.{Gauge, Meter, MetricName}
+import kafka.server.KafkaConfig
+import kafka.metrics.KafkaYammerMetrics
 import kafka.zk.ZooKeeperTestHarness
 import org.apache.kafka.common.security.JaasUtils
-import org.apache.zookeeper.KeeperException.Code
-import org.apache.zookeeper.{CreateMode, ZooDefs}
-import org.junit.Assert.{assertArrayEquals, assertEquals, assertTrue}
-import org.junit.{After, Test}
+import org.apache.kafka.common.utils.Time
+import org.apache.zookeeper.KeeperException.{Code, NoNodeException}
+import org.apache.zookeeper.Watcher.Event.{EventType, KeeperState}
+import org.apache.zookeeper.ZooKeeper.States
+import org.apache.zookeeper.client.ZKClientConfig
+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.jdk.CollectionConverters._
 
 class ZooKeeperClientTest extends ZooKeeperTestHarness {
   private val mockPath = "/foo"
+  private val time = Time.SYSTEM
+
+  private var zooKeeperClient: ZooKeeperClient = _
+
+  @Before
+  override def setUp(): Unit = {
+    ZooKeeperTestHarness.verifyNoUnexpectedThreads("@Before")
+    cleanMetricsRegistry()
+    super.setUp()
+    zooKeeperClient = new ZooKeeperClient(zkConnect, zkSessionTimeout, zkConnectionTimeout, zkMaxInFlightRequests,
+      Time.SYSTEM, "testMetricGroup", "testMetricType")
+  }
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
+    if (zooKeeperClient != null)
+      zooKeeperClient.close()
     super.tearDown()
     System.clearProperty(JaasUtils.JAVA_LOGIN_CONFIG_PARAM)
-    Configuration.setConfiguration(null)
+    ZooKeeperTestHarness.verifyNoUnexpectedThreads("@After")
   }
 
-  @Test(expected = classOf[UnknownHostException])
+  @Test
   def testUnresolvableConnectString(): Unit = {
-    new ZooKeeperClient("some.invalid.hostname.foo.bar.local", -1, -1, Int.MaxValue)
+    try {
+      new ZooKeeperClient("some.invalid.hostname.foo.bar.local", zkSessionTimeout, connectionTimeoutMs = 10,
+        Int.MaxValue, time, "testMetricGroup", "testMetricType")
+    } catch {
+      case e: ZooKeeperClientTimeoutException =>
+        assertEquals("ZooKeeper client threads still running", Set.empty,  runningZkSendThreads)
+    }
   }
 
+  private def runningZkSendThreads: collection.Set[String] = Thread.getAllStackTraces.keySet.asScala
+    .filter(_.isAlive)
+    .map(_.getName)
+    .filter(t => t.contains("SendThread()"))
+
   @Test(expected = classOf[ZooKeeperClientTimeoutException])
   def testConnectionTimeout(): Unit = {
     zookeeper.shutdown()
-    new ZooKeeperClient(zkConnect, zkSessionTimeout, connectionTimeoutMs = 100, Int.MaxValue)
+    new ZooKeeperClient(zkConnect, zkSessionTimeout, connectionTimeoutMs = 10, Int.MaxValue, time, "testMetricGroup",
+      "testMetricType").close()
   }
 
   @Test
   def testConnection(): Unit = {
-    new ZooKeeperClient(zkConnect, zkSessionTimeout, zkConnectionTimeout, Int.MaxValue)
+    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
+  def testConnectionViaNettyClient(): Unit = {
+    // Confirm that we can explicitly set client connection configuration, which is necessary for TLS.
+    // TLS connectivity itself is tested in system tests rather than here to avoid having to add TLS support
+    // to kafka.zk.EmbeddedZoopeeper
+    val clientConfig = new ZKClientConfig()
+    val propKey = KafkaConfig.ZkClientCnxnSocketProp
+    val propVal = "org.apache.zookeeper.ClientCnxnSocketNetty"
+    KafkaConfig.setZooKeeperClientProperty(clientConfig, propKey, propVal)
+    val client = new ZooKeeperClient(zkConnect, zkSessionTimeout, zkConnectionTimeout, Int.MaxValue, time, "testMetricGroup",
+      "testMetricType", None, Some(clientConfig))
+    try {
+      assertEquals(Some(propVal), KafkaConfig.getZooKeeperClientProperty(client.getClientConfig, propKey))
+      // For a sanity check, make sure a bad client connection socket class name generates an exception
+      val badClientConfig = new ZKClientConfig()
+      KafkaConfig.setZooKeeperClientProperty(badClientConfig, propKey, propVal + "BadClassName")
+      intercept[Exception] {
+          new ZooKeeperClient(zkConnect, zkSessionTimeout, zkConnectionTimeout, Int.MaxValue, time, "testMetricGroup",
+            "testMetricType", None, Some(badClientConfig))
+      }
+    } finally {
+      client.close()
+    }
   }
 
   @Test
   def testDeleteNonExistentZNode(): Unit = {
     val deleteResponse = zooKeeperClient.handleRequest(DeleteRequest(mockPath, -1))
     assertEquals("Response code should be NONODE", Code.NONODE, deleteResponse.resultCode)
+    intercept[NoNodeException] {
+      deleteResponse.maybeThrow()
+    }
   }
 
   @Test
   def testDeleteExistingZNode(): Unit = {
-    import scala.collection.JavaConverters._
-    val createResponse = zooKeeperClient.handleRequest(CreateRequest(mockPath, Array.empty[Byte], ZooDefs.Ids.OPEN_ACL_UNSAFE.asScala, CreateMode.PERSISTENT))
+    val createResponse = zooKeeperClient.handleRequest(CreateRequest(mockPath, Array.empty[Byte],
+      ZooDefs.Ids.OPEN_ACL_UNSAFE.asScala, CreateMode.PERSISTENT))
     assertEquals("Response code for create should be OK", Code.OK, createResponse.resultCode)
     val deleteResponse = zooKeeperClient.handleRequest(DeleteRequest(mockPath, -1))
     assertEquals("Response code for delete should be OK", Code.OK, deleteResponse.resultCode)
@@ -79,8 +151,8 @@ class ZooKeeperClientTest extends ZooKeeperTestHarness {
 
   @Test
   def testExistsExistingZNode(): Unit = {
-    import scala.collection.JavaConverters._
-    val createResponse = zooKeeperClient.handleRequest(CreateRequest(mockPath, Array.empty[Byte], ZooDefs.Ids.OPEN_ACL_UNSAFE.asScala, CreateMode.PERSISTENT))
+    val createResponse = zooKeeperClient.handleRequest(CreateRequest(mockPath, Array.empty[Byte],
+      ZooDefs.Ids.OPEN_ACL_UNSAFE.asScala, CreateMode.PERSISTENT))
     assertEquals("Response code for create should be OK", Code.OK, createResponse.resultCode)
     val existsResponse = zooKeeperClient.handleRequest(ExistsRequest(mockPath))
     assertEquals("Response code for exists should be OK", Code.OK, existsResponse.resultCode)
@@ -94,7 +166,6 @@ class ZooKeeperClientTest extends ZooKeeperTestHarness {
 
   @Test
   def testGetDataExistingZNode(): Unit = {
-    import scala.collection.JavaConverters._
     val data = bytes
     val createResponse = zooKeeperClient.handleRequest(CreateRequest(mockPath, data, ZooDefs.Ids.OPEN_ACL_UNSAFE.asScala,
       CreateMode.PERSISTENT))
@@ -112,7 +183,6 @@ class ZooKeeperClientTest extends ZooKeeperTestHarness {
 
   @Test
   def testSetDataExistingZNode(): Unit = {
-    import scala.collection.JavaConverters._
     val data = bytes
     val createResponse = zooKeeperClient.handleRequest(CreateRequest(mockPath, Array.empty[Byte],
       ZooDefs.Ids.OPEN_ACL_UNSAFE.asScala, CreateMode.PERSISTENT))
@@ -132,7 +202,6 @@ class ZooKeeperClientTest extends ZooKeeperTestHarness {
 
   @Test
   def testGetAclExistingZNode(): Unit = {
-    import scala.collection.JavaConverters._
     val createResponse = zooKeeperClient.handleRequest(CreateRequest(mockPath, Array.empty[Byte], ZooDefs.Ids.OPEN_ACL_UNSAFE.asScala, CreateMode.PERSISTENT))
     assertEquals("Response code for create should be OK", Code.OK, createResponse.resultCode)
     val getAclResponse = zooKeeperClient.handleRequest(GetAclRequest(mockPath))
@@ -142,31 +211,28 @@ class ZooKeeperClientTest extends ZooKeeperTestHarness {
 
   @Test
   def testSetAclNonExistentZNode(): Unit = {
-    import scala.collection.JavaConverters._
     val setAclResponse = zooKeeperClient.handleRequest(SetAclRequest(mockPath, ZooDefs.Ids.OPEN_ACL_UNSAFE.asScala, -1))
     assertEquals("Response code should be NONODE", Code.NONODE, setAclResponse.resultCode)
   }
 
   @Test
   def testGetChildrenNonExistentZNode(): Unit = {
-    val getChildrenResponse = zooKeeperClient.handleRequest(GetChildrenRequest(mockPath))
+    val getChildrenResponse = zooKeeperClient.handleRequest(GetChildrenRequest(mockPath, registerWatch = true))
     assertEquals("Response code should be NONODE", Code.NONODE, getChildrenResponse.resultCode)
   }
 
   @Test
   def testGetChildrenExistingZNode(): Unit = {
-    import scala.collection.JavaConverters._
     val createResponse = zooKeeperClient.handleRequest(CreateRequest(mockPath, Array.empty[Byte],
       ZooDefs.Ids.OPEN_ACL_UNSAFE.asScala, CreateMode.PERSISTENT))
     assertEquals("Response code for create should be OK", Code.OK, createResponse.resultCode)
-    val getChildrenResponse = zooKeeperClient.handleRequest(GetChildrenRequest(mockPath))
+    val getChildrenResponse = zooKeeperClient.handleRequest(GetChildrenRequest(mockPath, registerWatch = true))
     assertEquals("Response code for getChildren should be OK", Code.OK, getChildrenResponse.resultCode)
     assertEquals("getChildren should return no children", Seq.empty[String], getChildrenResponse.children)
   }
 
   @Test
   def testGetChildrenExistingZNodeWithChildren(): Unit = {
-    import scala.collection.JavaConverters._
     val child1 = "child1"
     val child2 = "child2"
     val child1Path = mockPath + "/" + child1
@@ -181,14 +247,13 @@ class ZooKeeperClientTest extends ZooKeeperTestHarness {
       ZooDefs.Ids.OPEN_ACL_UNSAFE.asScala, CreateMode.PERSISTENT))
     assertEquals("Response code for create child2 should be OK", Code.OK, createResponseChild2.resultCode)
 
-    val getChildrenResponse = zooKeeperClient.handleRequest(GetChildrenRequest(mockPath))
+    val getChildrenResponse = zooKeeperClient.handleRequest(GetChildrenRequest(mockPath, registerWatch = true))
     assertEquals("Response code for getChildren should be OK", Code.OK, getChildrenResponse.resultCode)
     assertEquals("getChildren should return two children", Seq(child1, child2), getChildrenResponse.children.sorted)
   }
 
   @Test
   def testPipelinedGetData(): Unit = {
-    import scala.collection.JavaConverters._
     val createRequests = (1 to 3).map(x => CreateRequest("/" + x, (x * 2).toString.getBytes, ZooDefs.Ids.OPEN_ACL_UNSAFE.asScala, CreateMode.PERSISTENT))
     val createResponses = createRequests.map(zooKeeperClient.handleRequest)
     createResponses.foreach(createResponse => assertEquals("Response code for create should be OK", Code.OK, createResponse.resultCode))
@@ -204,7 +269,6 @@ class ZooKeeperClientTest extends ZooKeeperTestHarness {
 
   @Test
   def testMixedPipeline(): Unit = {
-    import scala.collection.JavaConverters._
     val createResponse = zooKeeperClient.handleRequest(CreateRequest(mockPath, Array.empty[Byte],
       ZooDefs.Ids.OPEN_ACL_UNSAFE.asScala, CreateMode.PERSISTENT))
     assertEquals("Response code for create should be OK", Code.OK, createResponse.resultCode)
@@ -218,7 +282,6 @@ class ZooKeeperClientTest extends ZooKeeperTestHarness {
 
   @Test
   def testZNodeChangeHandlerForCreation(): Unit = {
-    import scala.collection.JavaConverters._
     val znodeChangeHandlerCountDownLatch = new CountDownLatch(1)
     val zNodeChangeHandler = new ZNodeChangeHandler {
       override def handleCreation(): Unit = {
@@ -238,7 +301,6 @@ class ZooKeeperClientTest extends ZooKeeperTestHarness {
 
   @Test
   def testZNodeChangeHandlerForDeletion(): Unit = {
-    import scala.collection.JavaConverters._
     val znodeChangeHandlerCountDownLatch = new CountDownLatch(1)
     val zNodeChangeHandler = new ZNodeChangeHandler {
       override def handleDeletion(): Unit = {
@@ -260,7 +322,6 @@ class ZooKeeperClientTest extends ZooKeeperTestHarness {
 
   @Test
   def testZNodeChangeHandlerForDataChange(): Unit = {
-    import scala.collection.JavaConverters._
     val znodeChangeHandlerCountDownLatch = new CountDownLatch(1)
     val zNodeChangeHandler = new ZNodeChangeHandler {
       override def handleDataChange(): Unit = {
@@ -280,9 +341,86 @@ class ZooKeeperClientTest extends ZooKeeperTestHarness {
     assertTrue("Failed to receive data change notification", znodeChangeHandlerCountDownLatch.await(5, TimeUnit.SECONDS))
   }
 
+  @Test
+  def testBlockOnRequestCompletionFromStateChangeHandler(): Unit = {
+    // This tests the scenario exposed by KAFKA-6879 in which the expiration callback awaits
+    // completion of a request which is handled by another thread
+
+    val latch = new CountDownLatch(1)
+    val stateChangeHandler = new StateChangeHandler {
+      override val name = this.getClass.getName
+      override def beforeInitializingSession(): Unit = {
+        latch.await()
+      }
+    }
+
+    zooKeeperClient.close()
+    zooKeeperClient = new ZooKeeperClient(zkConnect, zkSessionTimeout, zkConnectionTimeout, Int.MaxValue, time,
+      "testMetricGroup", "testMetricType")
+    zooKeeperClient.registerStateChangeHandler(stateChangeHandler)
+
+    val requestThread = new Thread() {
+      override def run(): Unit = {
+        try
+          zooKeeperClient.handleRequest(CreateRequest(mockPath, Array.empty[Byte],
+            ZooDefs.Ids.OPEN_ACL_UNSAFE.asScala, CreateMode.PERSISTENT))
+        finally
+          latch.countDown()
+      }
+    }
+
+    val reinitializeThread = new Thread() {
+      override def run(): Unit = {
+        zooKeeperClient.forceReinitialize()
+      }
+    }
+
+    reinitializeThread.start()
+
+    // sleep briefly before starting the request thread so that the initialization
+    // thread is blocking on the latch
+    Thread.sleep(100)
+    requestThread.start()
+
+    reinitializeThread.join()
+    requestThread.join()
+  }
+
+  @Test
+  def testExceptionInBeforeInitializingSession(): Unit = {
+    val faultyHandler = new StateChangeHandler {
+      override val name = this.getClass.getName
+      override def beforeInitializingSession(): Unit = {
+        throw new RuntimeException()
+      }
+    }
+
+    val goodCalls = new AtomicInteger(0)
+    val goodHandler = new StateChangeHandler {
+      override val name = this.getClass.getName
+      override def beforeInitializingSession(): Unit = {
+        goodCalls.incrementAndGet()
+      }
+    }
+
+    zooKeeperClient.close()
+    zooKeeperClient = new ZooKeeperClient(zkConnect, zkSessionTimeout, zkConnectionTimeout, Int.MaxValue, time,
+      "testMetricGroup", "testMetricType")
+    zooKeeperClient.registerStateChangeHandler(faultyHandler)
+    zooKeeperClient.registerStateChangeHandler(goodHandler)
+
+    zooKeeperClient.forceReinitialize()
+
+    assertEquals(1, goodCalls.get)
+
+    // Client should be usable even if the callback throws an error
+    val createResponse = zooKeeperClient.handleRequest(CreateRequest(mockPath, Array.empty[Byte],
+      ZooDefs.Ids.OPEN_ACL_UNSAFE.asScala, CreateMode.PERSISTENT))
+    assertEquals("Response code for create should be OK", Code.OK, createResponse.resultCode)
+  }
+
   @Test
   def testZNodeChildChangeHandlerForChildChange(): Unit = {
-    import scala.collection.JavaConverters._
     val zNodeChildChangeHandlerCountDownLatch = new CountDownLatch(1)
     val zNodeChildChangeHandler = new ZNodeChildChangeHandler {
       override def handleChildChange(): Unit = {
@@ -293,14 +431,42 @@ class ZooKeeperClientTest extends ZooKeeperTestHarness {
 
     val child1 = "child1"
     val child1Path = mockPath + "/" + child1
-    val createResponse = zooKeeperClient.handleRequest(CreateRequest(mockPath, Array.empty[Byte], ZooDefs.Ids.OPEN_ACL_UNSAFE.asScala, CreateMode.PERSISTENT))
+    val createResponse = zooKeeperClient.handleRequest(
+      CreateRequest(mockPath, Array.empty[Byte], ZooDefs.Ids.OPEN_ACL_UNSAFE.asScala, CreateMode.PERSISTENT))
+    assertEquals("Response code for create should be OK", Code.OK, createResponse.resultCode)
+    zooKeeperClient.registerZNodeChildChangeHandler(zNodeChildChangeHandler)
+    val getChildrenResponse = zooKeeperClient.handleRequest(GetChildrenRequest(mockPath, registerWatch = true))
+    assertEquals("Response code for getChildren should be OK", Code.OK, getChildrenResponse.resultCode)
+    val createResponseChild1 = zooKeeperClient.handleRequest(
+      CreateRequest(child1Path, Array.empty[Byte], ZooDefs.Ids.OPEN_ACL_UNSAFE.asScala, CreateMode.PERSISTENT))
+    assertEquals("Response code for create child1 should be OK", Code.OK, createResponseChild1.resultCode)
+    assertTrue("Failed to receive child change notification",
+      zNodeChildChangeHandlerCountDownLatch.await(5, TimeUnit.SECONDS))
+  }
+
+  @Test
+  def testZNodeChildChangeHandlerForChildChangeNotTriggered(): Unit = {
+    val zNodeChildChangeHandlerCountDownLatch = new CountDownLatch(1)
+    val zNodeChildChangeHandler = new ZNodeChildChangeHandler {
+      override def handleChildChange(): Unit = {
+        zNodeChildChangeHandlerCountDownLatch.countDown()
+      }
+      override val path: String = mockPath
+    }
+
+    val child1 = "child1"
+    val child1Path = mockPath + "/" + child1
+    val createResponse = zooKeeperClient.handleRequest(
+      CreateRequest(mockPath, Array.empty[Byte], ZooDefs.Ids.OPEN_ACL_UNSAFE.asScala, CreateMode.PERSISTENT))
     assertEquals("Response code for create should be OK", Code.OK, createResponse.resultCode)
     zooKeeperClient.registerZNodeChildChangeHandler(zNodeChildChangeHandler)
-    val getChildrenResponse = zooKeeperClient.handleRequest(GetChildrenRequest(mockPath))
+    val getChildrenResponse = zooKeeperClient.handleRequest(GetChildrenRequest(mockPath, registerWatch = false))
     assertEquals("Response code for getChildren should be OK", Code.OK, getChildrenResponse.resultCode)
-    val createResponseChild1 = zooKeeperClient.handleRequest(CreateRequest(child1Path, Array.empty[Byte], ZooDefs.Ids.OPEN_ACL_UNSAFE.asScala, CreateMode.PERSISTENT))
+    val createResponseChild1 = zooKeeperClient.handleRequest(
+      CreateRequest(child1Path, Array.empty[Byte], ZooDefs.Ids.OPEN_ACL_UNSAFE.asScala, CreateMode.PERSISTENT))
     assertEquals("Response code for create child1 should be OK", Code.OK, createResponseChild1.resultCode)
-    assertTrue("Failed to receive child change notification", zNodeChildChangeHandlerCountDownLatch.await(5, TimeUnit.SECONDS))
+    assertFalse("Child change notification received",
+      zNodeChildChangeHandlerCountDownLatch.await(100, TimeUnit.MILLISECONDS))
   }
 
   @Test
@@ -315,44 +481,241 @@ class ZooKeeperClientTest extends ZooKeeperTestHarness {
       }
     }
 
-    val zooKeeperClient = new ZooKeeperClient(zkConnect, zkSessionTimeout, zkConnectionTimeout, Int.MaxValue)
-    zooKeeperClient.registerStateChangeHandler(stateChangeHandler)
-    zooKeeperClient.reinitialize()
+    val zooKeeperClient = new ZooKeeperClient(zkConnect, zkSessionTimeout, zkConnectionTimeout, Int.MaxValue, time,
+      "testMetricGroup", "testMetricType")
+    try {
+      zooKeeperClient.registerStateChangeHandler(stateChangeHandler)
+      zooKeeperClient.forceReinitialize()
 
-    assertTrue("Failed to receive auth failed notification", stateChangeHandlerCountDownLatch.await(5, TimeUnit.SECONDS))
+      assertTrue("Failed to receive auth failed notification", stateChangeHandlerCountDownLatch.await(5, TimeUnit.SECONDS))
+    } finally zooKeeperClient.close()
   }
 
   @Test
   def testConnectionLossRequestTermination(): Unit = {
     val batchSize = 10
-    val zooKeeperClient = new ZooKeeperClient(zkConnect, zkSessionTimeout, zkConnectionTimeout, 2)
+    val zooKeeperClient = new ZooKeeperClient(zkConnect, zkSessionTimeout, zkConnectionTimeout, 2, time,
+      "testGroupType", "testGroupName")
     zookeeper.shutdown()
-    val requests = (1 to batchSize).map(i => GetDataRequest(s"/$i"))
-    val countDownLatch = new CountDownLatch(1)
-    val running = new AtomicBoolean(true)
-    val unexpectedResponses = new ArrayBlockingQueue[GetDataResponse](batchSize)
-    val requestThread = new Thread {
-      override def run(): Unit = {
-        while (running.get()) {
-          val responses = zooKeeperClient.handleRequests(requests)
-          val suffix = responses.dropWhile(response => response.resultCode != Code.CONNECTIONLOSS)
-          if (!suffix.forall(response => response.resultCode == Code.CONNECTIONLOSS))
-            responses.foreach(unexpectedResponses.add)
-          if (!unexpectedResponses.isEmpty || suffix.nonEmpty)
-            running.set(false)
+    try {
+      val requests = (1 to batchSize).map(i => GetDataRequest(s"/$i"))
+      val countDownLatch = new CountDownLatch(1)
+      val running = new AtomicBoolean(true)
+      val unexpectedResponses = new ArrayBlockingQueue[GetDataResponse](batchSize)
+      val requestThread = new Thread {
+        override def run(): Unit = {
+          while (running.get()) {
+            val responses = zooKeeperClient.handleRequests(requests)
+            val suffix = responses.dropWhile(response => response.resultCode != Code.CONNECTIONLOSS)
+            if (!suffix.forall(response => response.resultCode == Code.CONNECTIONLOSS))
+              responses.foreach(unexpectedResponses.add)
+            if (!unexpectedResponses.isEmpty || suffix.nonEmpty)
+              running.set(false)
+          }
+          countDownLatch.countDown()
         }
-        countDownLatch.countDown()
+      }
+      requestThread.start()
+      val requestThreadTerminated = countDownLatch.await(30, TimeUnit.SECONDS)
+      if (!requestThreadTerminated) {
+        running.set(false)
+        requestThread.join(5000)
+        fail("Failed to receive a CONNECTIONLOSS response code after zookeeper has shutdown.")
+      } else if (!unexpectedResponses.isEmpty) {
+        fail(s"Received an unexpected non-CONNECTIONLOSS response code after a CONNECTIONLOSS response code from a single batch: $unexpectedResponses")
+      }
+    } finally zooKeeperClient.close()
+  }
+
+  /**
+    * Tests that if session expiry notification is received while a thread is processing requests,
+    * session expiry is handled and the request thread completes with responses to all requests,
+    * even though some requests may fail due to session expiry or disconnection.
+    *
+    * Sequence of events on different threads:
+    *   Request thread:
+    *       - Sends `maxInflightRequests` requests (these may complete before session is expired)
+    *   Main thread:
+    *       - Waits for at least one request to be processed (this should succeed)
+    *       - Expires session by creating new client with same session id
+    *       - Unblocks another `maxInflightRequests` requests before and after new client is closed (these may fail)
+    *   ZooKeeperClient Event thread:
+    *       - Delivers responses and session expiry (no ordering guarantee between these, both are processed asynchronously)
+    *   Response executor thread:
+    *       - Blocks subsequent sends by delaying response until session expiry is processed
+    *   ZooKeeperClient Session Expiry Handler:
+    *       - Unblocks subsequent sends
+    *   Main thread:
+    *       - Waits for all sends to complete. The requests sent after session expiry processing should succeed.
+    */
+  @Test
+  def testSessionExpiry(): Unit = {
+    val maxInflightRequests = 2
+    val responseExecutor = Executors.newSingleThreadExecutor
+    val sendSemaphore = new Semaphore(0)
+    val sendCompleteSemaphore = new Semaphore(0)
+    val sendSize = maxInflightRequests * 5
+    @volatile var resultCodes: Seq[Code] = null
+    val stateChanges = new ConcurrentLinkedQueue[String]()
+    val zooKeeperClient = new ZooKeeperClient(zkConnect, zkSessionTimeout, zkConnectionTimeout, maxInflightRequests,
+      time, "testGroupType", "testGroupName") {
+      override def send[Req <: AsyncRequest](request: Req)(processResponse: Req#Response => Unit): Unit = {
+        super.send(request)( response => {
+          responseExecutor.submit(new Runnable {
+            override def run(): Unit = {
+              sendCompleteSemaphore.release()
+              sendSemaphore.acquire()
+              processResponse(response)
+            }
+          })
+        })
       }
     }
-    requestThread.start()
-    val requestThreadTerminated = countDownLatch.await(30, TimeUnit.SECONDS)
-    if (!requestThreadTerminated) {
-      running.set(false)
-      requestThread.join(5000)
-      fail("Failed to receive a CONNECTIONLOSS response code after zookeeper has shutdown.")
-    } else if (!unexpectedResponses.isEmpty) {
-      fail(s"Received an unexpected non-CONNECTIONLOSS response code after a CONNECTIONLOSS response code from a single batch: $unexpectedResponses")
+    try {
+      zooKeeperClient.registerStateChangeHandler(new StateChangeHandler {
+        override val name: String ="test-state-change-handler"
+        override def afterInitializingSession(): Unit = {
+          verifyHandlerThread()
+          stateChanges.add("afterInitializingSession")
+        }
+        override def beforeInitializingSession(): Unit = {
+          verifyHandlerThread()
+          stateChanges.add("beforeInitializingSession")
+          sendSemaphore.release(sendSize) // Resume remaining sends
+        }
+        private def verifyHandlerThread(): Unit = {
+          val threadName = Thread.currentThread.getName
+          assertTrue(s"Unexpected thread + $threadName", threadName.startsWith(zooKeeperClient.reinitializeScheduler.threadNamePrefix))
+        }
+      })
+
+      val requestThread = new Thread {
+        override def run(): Unit = {
+          val requests = (1 to sendSize).map(i => GetDataRequest(s"/$i"))
+          resultCodes = zooKeeperClient.handleRequests(requests).map(_.resultCode)
+        }
+      }
+      requestThread.start()
+      sendCompleteSemaphore.acquire() // Wait for request thread to start processing requests
+
+      val anotherZkClient = createZooKeeperClientToTriggerSessionExpiry(zooKeeperClient.currentZooKeeper)
+      sendSemaphore.release(maxInflightRequests) // Resume a few more sends which may fail
+      anotherZkClient.close()
+      sendSemaphore.release(maxInflightRequests) // Resume a few more sends which may fail
+
+      requestThread.join(10000)
+      if (requestThread.isAlive) {
+        requestThread.interrupt()
+        fail("Request thread did not complete")
+      }
+      assertEquals(Seq("beforeInitializingSession", "afterInitializingSession"), stateChanges.asScala.toSeq)
+
+      assertEquals(resultCodes.size, sendSize)
+      val connectionLostCount = resultCodes.count(_ == Code.CONNECTIONLOSS)
+      assertTrue(s"Unexpected connection lost requests $resultCodes", connectionLostCount <= maxInflightRequests)
+      val expiredCount = resultCodes.count(_ == Code.SESSIONEXPIRED)
+      assertTrue(s"Unexpected session expired requests $resultCodes", expiredCount <= maxInflightRequests)
+      assertTrue(s"No connection lost or expired requests $resultCodes", connectionLostCount + expiredCount > 0)
+      assertEquals(Code.NONODE, resultCodes.head)
+      assertEquals(Code.NONODE, resultCodes.last)
+      assertTrue(s"Unexpected result code $resultCodes",
+        resultCodes.filterNot(Set(Code.NONODE, Code.SESSIONEXPIRED, Code.CONNECTIONLOSS).contains).isEmpty)
+
+    } finally {
+      zooKeeperClient.close()
+      responseExecutor.shutdownNow()
     }
+    assertFalse("Expiry executor not shutdown", zooKeeperClient.reinitializeScheduler.isStarted)
+  }
+
+  @Test
+  def testSessionExpiryDuringClose(): Unit = {
+    val semaphore = new Semaphore(0)
+    val closeExecutor = Executors.newSingleThreadExecutor
+    try {
+      zooKeeperClient.reinitializeScheduler.schedule("test", () => semaphore.acquireUninterruptibly(),
+        delay = 0, period = -1, TimeUnit.SECONDS)
+      zooKeeperClient.scheduleReinitialize("session-expired", "Session expired.", delayMs = 0L)
+      val closeFuture = closeExecutor.submit(new Runnable {
+        override def run(): Unit = {
+          zooKeeperClient.close()
+        }
+      })
+      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.reinitializeScheduler.isStarted)
+    } finally {
+      closeExecutor.shutdownNow()
+    }
+  }
+
+  @Test
+  def testReinitializeAfterAuthFailure(): Unit = {
+    val sessionInitializedCountDownLatch = new CountDownLatch(1)
+    val changeHandler = new StateChangeHandler {
+      override val name = this.getClass.getName
+      override def beforeInitializingSession(): Unit = {
+        sessionInitializedCountDownLatch.countDown()
+      }
+    }
+
+    zooKeeperClient.close()
+    zooKeeperClient = new ZooKeeperClient(zkConnect, zkSessionTimeout, zkConnectionTimeout, Int.MaxValue, time,
+      "testMetricGroup", "testMetricType")
+    zooKeeperClient.registerStateChangeHandler(changeHandler)
+
+    zooKeeperClient.ZooKeeperClientWatcher.process(new WatchedEvent(EventType.None, KeeperState.AuthFailed, null))
+    assertTrue("Failed to receive session initializing notification", sessionInitializedCountDownLatch.await(5, TimeUnit.SECONDS))
+  }
+
+  def isExpectedMetricName(metricName: MetricName, name: String): Boolean =
+    metricName.getName == name && metricName.getGroup == "testMetricGroup" && metricName.getType == "testMetricType"
+
+  @Test
+  def testZooKeeperStateChangeRateMetrics(): Unit = {
+    def checkMeterCount(name: String, expected: Long): Unit = {
+      val meter = KafkaYammerMetrics.defaultRegistry.allMetrics.asScala.collectFirst {
+        case (metricName, meter: Meter) if isExpectedMetricName(metricName, name) => meter
+      }.getOrElse(sys.error(s"Unable to find meter with name $name"))
+      assertEquals(s"Unexpected meter count for $name", expected, meter.count)
+    }
+
+    val expiresPerSecName = "ZooKeeperExpiresPerSec"
+    val disconnectsPerSecName = "ZooKeeperDisconnectsPerSec"
+    checkMeterCount(expiresPerSecName, 0)
+    checkMeterCount(disconnectsPerSecName, 0)
+
+    zooKeeperClient.ZooKeeperClientWatcher.process(new WatchedEvent(EventType.None, KeeperState.Expired, null))
+    checkMeterCount(expiresPerSecName, 1)
+    checkMeterCount(disconnectsPerSecName, 0)
+
+    zooKeeperClient.ZooKeeperClientWatcher.process(new WatchedEvent(EventType.None, KeeperState.Disconnected, null))
+    checkMeterCount(expiresPerSecName, 1)
+    checkMeterCount(disconnectsPerSecName, 1)
+  }
+
+  @Test
+  def testZooKeeperSessionStateMetric(): Unit = {
+    def gaugeValue(name: String): Option[String] = {
+      KafkaYammerMetrics.defaultRegistry.allMetrics.asScala.collectFirst {
+        case (metricName, gauge: Gauge[_]) if isExpectedMetricName(metricName, name) => gauge.value.asInstanceOf[String]
+      }
+    }
+
+    assertEquals(Some(States.CONNECTED.toString), gaugeValue("SessionState"))
+    assertEquals(States.CONNECTED, zooKeeperClient.connectionState)
+
+    zooKeeperClient.close()
+
+    assertEquals(None, gaugeValue("SessionState"))
+    assertEquals(States.CLOSED, zooKeeperClient.connectionState)
+  }
+
+  private def cleanMetricsRegistry(): Unit = {
+    val metrics = KafkaYammerMetrics.defaultRegistry
+    metrics.allMetrics.keySet.forEach(metrics.removeMetric)
   }
 
   private def bytes = UUID.randomUUID().toString.getBytes(StandardCharsets.UTF_8)
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 97771861883dd..94d5f3e0da3e9 100644
--- a/docs/api.html
+++ b/docs/api.html
@@ -21,12 +21,12 @@
 	
      • The Consumer API allows applications to read streams of data from topics in the Kafka cluster.
 	
      • The Streams API allows transforming streams of data from input topics to output topics.
 	
      • 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.
-	
      • The AdminClient API allows managing and inspecting topics, brokers, and other Kafka objects.
+	
      • 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.
 
-	
        2.1 Producer API
        
+	
        2.1 Producer API
        
 
 	The Producer API allows applications to send streams of data to topics in the Kafka cluster.
 	
        
@@ -35,15 +35,13 @@ 
        2.1 Producer API
        
 	
        
 	To use the producer, you can use the following maven dependency:
 
-	
        -		<dependency>
+	
        		<dependency>
 			<groupId>org.apache.kafka</groupId>
 			<artifactId>kafka-clients</artifactId>
 			<version>{{fullDotVersion}}</version>
-		</dependency>
-	
        
+		</dependency>
        
 
-	
        2.2 Consumer API
        
+	
        2.2 Consumer API
        
 
 	The Consumer API allows applications to read streams of data from topics in the Kafka cluster.
 	
        
@@ -51,15 +49,13 @@ 
        2.2 Consumer API
        
 	javadocs.
 	
        
 	To use the consumer, you can use the following maven dependency:
-	
        -		<dependency>
+	
        		<dependency>
 			<groupId>org.apache.kafka</groupId>
 			<artifactId>kafka-clients</artifactId>
 			<version>{{fullDotVersion}}</version>
-		</dependency>
-	
        
+		</dependency>
        
 
-	
        2.3 Streams API
        
+	
        2.3 Streams API
        
 
 	The Streams API allows transforming streams of data from input topics to output topics.
 	
        
@@ -70,15 +66,24 @@ 
        2.3 Streams API
        
 	
        
 	To use Kafka Streams you can use the following maven dependency:
 
-	
        -		<dependency>
+	
        		<dependency>
 			<groupId>org.apache.kafka</groupId>
 			<artifactId>kafka-streams</artifactId>
 			<version>{{fullDotVersion}}</version>
-		</dependency>
-	
        
+		</dependency>
        
 
-	
        2.4 Connect API
        
+	
        
+	When using Scala you may optionally include the kafka-streams-scala library.  Additional documentation on using the Kafka Streams DSL for Scala is available in the developer guide.
+	
        
+	To use Kafka Streams DSL for Scala for Scala {{scalaVersion}} you can use the following maven dependency:
+
+	
        		<dependency>
+			<groupId>org.apache.kafka</groupId>
+			<artifactId>kafka-streams-scala_{{scalaVersion}}</artifactId>
+			<version>{{fullDotVersion}}</version>
+		</dependency>
        
+
+	
        2.4 Connect API
        
 
 	The Connect API allows implementing connectors that continually pull from some source data system into Kafka or push from Kafka into some sink data system.
 	
        
@@ -87,27 +92,19 @@ 
        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:
-	
        -		<dependency>
+	To use the Admin API, add the following Maven dependency:
+	
        		<dependency>
 			<groupId>org.apache.kafka</groupId>
 			<artifactId>kafka-clients</artifactId>
 			<version>{{fullDotVersion}}</version>
-		</dependency>
-	
        
-	For more information about the AdminClient APIs, see the javadoc.
+		</dependency>
        
+	For more information about the Admin APIs, see the javadoc.
 	
        
 
-	
        2.6 Legacy APIs
        
-
-	
        
-	A more limited legacy producer and consumer api is also included in Kafka. These old Scala APIs are deprecated and only still available for compatibility purposes. Information on them can be found here 
-	here.
-	
        
 
 
 
        
diff --git a/docs/configuration.html b/docs/configuration.html
index 3f962e9854140..34340a944d36a 100644
--- a/docs/configuration.html
+++ b/docs/configuration.html
@@ -18,7 +18,7 @@
 
diff --git a/docs/connect.html b/docs/connect.html
index b910cf57d4d7f..ef4cfe9668c7a 100644
--- a/docs/connect.html
+++ b/docs/connect.html
@@ -32,7 +32,7 @@ 
        8.1 Overview
        
 
     
        8.2 User Guide
        
 
-    
        The quickstart provides a brief example of how to run a standalone version of Kafka Connect. This section describes how to configure, run, and manage Kafka Connect in more detail.
        
+    
        The quickstart provides a brief example of how to run a standalone version of Kafka Connect. This section describes how to configure, run, and manage Kafka Connect in more detail.
        
 
     
        Running Kafka Connect
        
 
@@ -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 exeption 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).
        
 
@@ -165,19 +167,22 @@ 
        Transformations
 
     
        You can see that the lines we've read are now part of a JSON map, and there is an extra field with the static value we specified. This is just one example of what you can do with transformations.
        
+    
+    
        Included transformations
        
 
     
        Several widely-applicable data and routing transformations are included with Kafka Connect:
        
 
     
          
         
        • InsertField - Add a field using either static data or record metadata
          • 
         
        • ReplaceField - Filter or rename fields
          • 
-        
        • MaskField - Replace field with valid null value for the type (0, empty string, etc)
          • 
-        
        • ValueToKey
          • 
+        
        • MaskField - Replace field with valid null value for the type (0, empty string, etc) or custom replacement (non-empty string or numeric value only)
          • 
+        
        • ValueToKey - Replace the record key with a new key formed from a subset of fields in the record value
          • 
         
        • HoistField - Wrap the entire event as a single field inside a Struct or a Map
          • 
         
        • ExtractField - Extract a specific field from Struct and Map and include only this field in results
          • 
         
        • SetSchemaMetadata - modify the schema name or version
          • 
         
        • TimestampRouter - Modify the topic of a record based on original topic and timestamp. Useful when using a sink that needs to write to different tables or indexes based on timestamps
          • 
         
        • RegexRouter - modify the topic of a record based on original topic, replacement string and a regular expression
          • 
+        
        • Filter - Removes messages from all further processing. This is used with a predicate to selectively filter certain messages.
          • 
     
        
 
     
        Details on how to configure each transformation are listed below:
        
@@ -185,9 +190,111 @@ 
        Transformations
 
+
+    
        Predicates
        
+
+    
        Transformations can be configured with predicates so that the transformation is applied only to messages which satisfy some condition. In particular, when combined with the Filter transformation predicates can be used to selectively filter out certain messages.
        
+
+    
        Predicates are specified in the connector configuration.
        
+
+    
          
+        
        • predicates - Set of aliases for the predicates to be applied to some of the transformations.
          • 
+        
        • predicates.$alias.type - Fully qualified class name for the predicate.
          • 
+        
        • predicates.$alias.$predicateSpecificConfig - Configuration properties for the predicate.
          • 
+    
        
+
+    
        All transformations have the implicit config properties predicate and negate. A predicular predicate is associated with a transformation by setting the transformation's predicate config to the predicate's alias. The predicate's value can be reversed using the negate configuration property.
        
+
+    
        For example, suppose you have a source connector which produces messages to many different topics and you want to:
        
+    
          
+        
        • filter out the messages in the 'foo' topic entirely
          • 
+        
        • apply the ExtractField transformation with the field name 'other_field' to records in all topics except the topic 'bar'
          • 
+    
        
+
+    
        To do this we need first to filter out the records destined for the topic 'foo'. The Filter transformation removes records from further processing, and can use the TopicNameMatches predicate to apply the transformation only to records in topics which match a certain regular expression. TopicNameMatches's only configuration property is pattern which is a Java regular expression for matching against the topic name. The configuration would look like this:
        
+
+    
        +        transforms=Filter
+        transforms.Filter.type=org.apache.kafka.connect.transforms.Filter
+        transforms.Filter.predicate=IsFoo
+
+        predicates=IsFoo
+        predicates.IsFoo.type=org.apache.kafka.connect.transforms.predicates.TopicNameMatches
+        predicates.IsFoo.pattern=foo
+    
        
+        
+    
        Next we need to apply ExtractField only when the topic name of the record is not 'bar'. We can't just use TopicNameMatches directly, because that would apply the transformation to matching topic names, not topic names which do not match. The transformation's implicit negate config properties allows us to invert the set of records which a predicate matches. Adding the configuration for this to the previous example we arrive at:
        
+
+    
        +        transforms=Filter,Extract
+        transforms.Filter.type=org.apache.kafka.connect.transforms.Filter
+        transforms.Filter.predicate=IsFoo
+
+        transforms.Extract.type=org.apache.kafka.connect.transforms.ExtractField$Key
+        transforms.Extract.field=other_field
+        transforms.Extract.predicate=IsBar
+        transforms.Extract.negate=true
+
+        predicates=IsFoo,IsBar
+        predicates.IsFoo.type=org.apache.kafka.connect.transforms.predicates.TopicNameMatches
+        predicates.IsFoo.pattern=foo
+
+        predicates.IsBar.type=org.apache.kafka.connect.transforms.predicates.TopicNameMatches
+        predicates.IsBar.pattern=bar
+    
        
+
+    
        Kafka Connect includes the following predicates:
        
+
+    
          
+        
        • TopicNameMatches - matches records in a topic with a name matching a particular Java regular expression.
          • 
+        
        • HasHeaderKey - matches records which have a header with the given key.
          • 
+        
        • RecordIsTombstone - matches tombstone records, that is records with a null value.
          • 
+    
        
+
+    
        Details on how to configure each predicate are listed below:
        
+    
+    
+
+
     
        REST API
        
 
-    
        Since Kafka Connect is intended to be run as a service, it also provides a REST API for managing connectors. By default, this service runs on port 8083. The following are the currently supported endpoints:
        
+    
        Since Kafka Connect is intended to be run as a service, it also provides a REST API for managing connectors. The REST API server can be configured using the listeners configuration option.
+        This field should contain a list of listeners in the following format: protocol://host:port,protocol2://host2:port2. Currently supported protocols are http and https.
+        For example:
        
+
+    
        +        listeners=http://localhost:8080,https://localhost:8443
+    
        
+
+    
        By default, if no listeners are specified, the REST server runs on port 8083 using the HTTP protocol. When using HTTPS, the configuration has to include the SSL configuration.
+    By default, it will use the ssl.* settings. In case it is needed to use different configuration for the REST API than for connecting to Kafka brokers, the fields can be prefixed with listeners.https.
+        When using the prefix, only the prefixed options will be used and the ssl.* options without the prefix will be ignored. Following fields can be used to configure HTTPS for the REST API:
        
+
+    
          
+        
        • ssl.keystore.location
          • 
+        
        • ssl.keystore.password
          • 
+        
        • ssl.keystore.type
          • 
+        
        • ssl.key.password
          • 
+        
        • ssl.truststore.location
          • 
+        
        • ssl.truststore.password
          • 
+        
        • ssl.truststore.type
          • 
+        
        • ssl.enabled.protocols
          • 
+        
        • ssl.provider
          • 
+        
        • ssl.protocol
          • 
+        
        • ssl.cipher.suites
          • 
+        
        • ssl.keymanager.algorithm
          • 
+        
        • ssl.secure.random.implementation
          • 
+        
        • ssl.trustmanager.algorithm
          • 
+        
        • ssl.endpoint.identification.algorithm
          • 
+        
        • ssl.client.auth
          • 
+    
        
+
+    
        The REST API is used not only by users to monitor / manage Kafka Connect. It is also used for the Kafka Connect cross-cluster communication. Requests received on the follower nodes REST API will be forwarded to the leader node REST API.
+    In case the URI under which is given host reachable is different from the URI which it listens on, the configuration options rest.advertised.host.name, rest.advertised.port and rest.advertised.listener
+    can be used to change the URI which will be used by the follower nodes to connect with the leader. When using both HTTP and HTTPS listeners, the rest.advertised.listener option can be also used to define which listener
+        will be used for the cross-cluster communication. When using HTTPS for communication between nodes, the same ssl.* or listeners.https options will be used to configure the HTTPS client.
        
+
+    
        The following are the currently supported REST API endpoints:
        
 
     
          
         
        • GET /connectors - return a list of active connectors
          • 
@@ -203,6 +310,8 @@ 
          REST API
          
         
        • POST /connectors/{name}/restart - restart a connector (typically because it has failed)
          • 
         
        • POST /connectors/{name}/tasks/{taskId}/restart - restart an individual task (typically because it has failed)
          • 
         
        • DELETE /connectors/{name} - delete a connector, halting all tasks and deleting its configuration
          • 
+        
        • GET /connectors/{name}/topics - get the set of topics that a specific connector is using since the connector was created or since a request to reset its set of active topics was issued
          • 
+        
        • PUT /connectors/{name}/topics/reset - send a request to empty the set of active topics of a connector
          • 
     
        
 
     
        Kafka Connect also provides a REST API for getting information about connector plugins:
        
@@ -212,6 +321,54 @@ 
        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.
+    
        
+
+    
        Error Reporting in Connect
        
+
+    
        Kafka Connect provides error reporting to handle errors encountered along various stages of processing. By default, any error encountered during conversion or within transformations will cause the connector to fail. Each connector configuration can also enable tolerating such errors by skipping them, optionally writing each error and the details of the failed operation and problematic record (with various levels of detail) to the Connect application log. These mechanisms also capture errors when a sink connector is processing the messages consumed from its Kafka topics, and all of the errors can be written to a configurable "dead letter queue" (DLQ) Kafka topic.
        
+
+    
        To report errors within a connector's converter, transforms, or within the sink connector itself to the log, set errors.log.enable=true in the connector configuration to log details of each error and problem record's topic, partition, and offset. For additional debugging purposes, set errors.log.include.messages=true to also log the problem record key, value, and headers to the log (note this may log sensitive information).
        
+
+    
        To report errors within a connector's converter, transforms, or within the sink connector itself to a dead letter queue topic, set errors.deadletterqueue.topic.name, and optionally errors.deadletterqueue.context.headers.enable=true.
        
+
+    
        By default connectors exhibit "fail fast" behavior immediately upon an error or exception. This is equivalent to adding the following configuration properties with their defaults to a connector configuration:
        
+
+    
        +        # disable retries on failure
+        errors.retry.timeout=0
+
+        # do not log the error and their contexts
+        errors.log.enable=false
+
+        # do not record errors in a dead letter queue topic
+        errors.deadletterqueue.topic.name=
+
+        # Fail on first error
+        errors.tolerance=none
+    
        
+
+    
        These and other related connector configuration properties can be changed to provide different behavior. For example, the following configuration properties can be added to a connector configuration to setup error handling with multiple retries, logging to the application logs and the my-connector-errors Kafka topic, and tolerating all errors by reporting them rather than failing the connector task:
        
+
+    
        +        # retry for at most 10 minutes times waiting up to 30 seconds between consecutive failures
+        errors.retry.timeout=600000
+        errors.retry.delay.max.ms=30000
+
+        # log error context along with application logs, but do not include configs and messages
+        errors.log.enable=true
+        errors.log.include.messages=false
+
+        # produce error context into the Kafka topic
+        errors.deadletterqueue.topic.name=my-connector-errors
+
+        # Tolerate all errors.
+        errors.tolerance=all
+    
        
+
     
        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.
        
@@ -263,7 +420,7 @@ 
        Connector
     }
     
 
-    
        We will define the FileStreamSourceTask class below. Next, we add some standard lifecycle methods, start() and stop()
        :
+    
        We will define the FileStreamSourceTask class below. Next, we add some standard lifecycle methods, start() and stop():
        
 
     
             @Override
@@ -358,7 +515,7 @@ 
        Task Example - Sourc
     }
     
        
 
-    
        Again, we've omitted some details, but we can see the important steps: the poll() method is going to be called repeatedly, and for each call it will loop trying to read records from the file. For each line it reads, it also tracks the file offset. It uses this information to create an output SourceRecord with four pieces of information: the source partition (there is only one, the single file being read), source offset (byte offset in the file), output topic name, and output value (the line, and we include a schema indicating this value will always be a string). Other variants of the SourceRecord constructor can also include a specific output partition and a key.
        
+    
        Again, we've omitted some details, but we can see the important steps: the poll() method is going to be called repeatedly, and for each call it will loop trying to read records from the file. For each line it reads, it also tracks the file offset. It uses this information to create an output SourceRecord with four pieces of information: the source partition (there is only one, the single file being read), source offset (byte offset in the file), output topic name, and output value (the line, and we include a schema indicating this value will always be a string). Other variants of the SourceRecord constructor can also include a specific output partition, a key, and headers.
        
 
     
        Note that this implementation uses the normal Java InputStream interface and may sleep if data is not available. This is acceptable because Kafka Connect provides each task with a dedicated thread. While task implementations have to conform to the basic poll() interface, they have a lot of flexibility in how they are implemented. In this case, an NIO-based implementation would be more efficient, but this simple approach works, is quick to implement, and is compatible with older versions of Java.
        
 
@@ -378,11 +535,47 @@ 
        Sink Tasks
        
         }
     
 
-    
        The SinkTask documentation contains full details, but this interface is nearly as simple as the SourceTask. The put() method should contain most of the implementation, accepting sets of SinkRecords, performing any required translation, and storing them in the destination system. This method does not need to ensure the data has been fully written to the destination system before returning. In fact, in many cases internal buffering will be useful so an entire batch of records can be sent at once, reducing the overhead of inserting events into the downstream data store. The SinkRecords contain essentially the same information as SourceRecords: Kafka topic, partition, offset and the event key and value.
        
+    
        The SinkTask documentation contains full details, but this interface is nearly as simple as the SourceTask. The put() method should contain most of the implementation, accepting sets of SinkRecords, performing any required translation, and storing them in the destination system. This method does not need to ensure the data has been fully written to the destination system before returning. In fact, in many cases internal buffering will be useful so an entire batch of records can be sent at once, reducing the overhead of inserting events into the downstream data store. The SinkRecords contain essentially the same information as SourceRecords: Kafka topic, partition, offset, the event key and value, and optional headers.
        
 
     
        The flush() method is used during the offset commit process, which allows tasks to recover from failures and resume from a safe point such that no events will be missed. The method should push any outstanding data to the destination system and then block until the write has been acknowledged. The offsets parameter can often be ignored, but is useful in some cases where implementations want to store offset information in the destination store to provide exactly-once
     delivery. For example, an HDFS connector could do this and use atomic move operations to make sure the flush() operation atomically commits the data and offsets to a final location in HDFS.
        
 
+    
        Errant Record Reporter
        
+
+    
        When error reporting is enabled for a connector, the connector can use an ErrantRecordReporter to report problems with individual records sent to a sink connector. The following example shows how a connector's SinkTask subclass might obtain and use the ErrantRecordReporter, safely handling a null reporter when the DLQ is not enabled or when the connector is installed in an older Connect runtime that doesn't have this reporter feature:
        
+
+    
        +        private ErrantRecordReporter reporter;
+
+        @Override
+        public void start(Map<String, String> props) {
+            ...
+            try {
+                reporter = context.errantRecordReporter(); // may be null if DLQ not enabled
+            } catch (NoSuchMethodException | NoClassDefFoundError e) {
+                // Will occur in Connect runtimes earlier than 2.6
+                reporter = null;
+            }
+        }
+
+        @Override
+        public void put(Collection<SinkRecord> records) {
+            for (SinkRecord record: records) {
+                try {
+                    // attempt to process and send record to data sink
+                    process(record);
+                } catch(Exception e) {
+                    if (reporter != null) {
+                        // Send errant record to error reporter
+                        reporter.report(record, e);
+                    } else {
+                        // There's no error reporter, so fail
+                        throw new ConnectException("Failed on record", e);
+                    }
+                }
+            }
+        }
+    
        
 
     
        Resuming from Previous Offsets
        
 
@@ -451,7 +644,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:
     
        
 
     
        @@ -501,12 +726,17 @@ 
        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.
        • 
     
 
     
        
     In most cases, connector and task states will match, though they may be different for short periods of time when changes are occurring or if tasks have failed. For example, when a connector is first started, there may be a noticeable delay before the connector and its tasks have all transitioned to the RUNNING state. States will also diverge when tasks fail since Connect does not automatically restart failed tasks. To restart a connector/task manually, you can use the restart APIs listed above. Note that if you try to restart a task while a rebalance is taking place, Connect will return a 409 (Conflict) status code. You can retry after the rebalance completes, but it might not be necessary since rebalances effectively restart all the connectors and tasks in the cluster.
     
        
 
+    
        
+    Starting with 2.5.0, Kafka Connect uses the status.storage.topic to also store information related to the topics that each connector is using. Connect Workers use these per-connector topic status updates to respond to requests to the REST endpoint GET /connectors/{name}/topics by returning the set of topic names that a connector is using. A request to the REST endpoint PUT /connectors/{name}/topics/reset resets the set of active topics for a connector and allows a new set to be populated, based on the connector's latest pattern of topic usage. Upon connector deletion, the set of the connector's active topics is also deleted. Topic tracking is enabled by default but can be disabled by setting topic.tracking.enable=false. If you want to disallow requests to reset the active topics of connectors during runtime, set the Worker property topic.tracking.allow.reset=false.
+    
        
+
     
        
     It's sometimes useful to temporarily stop the message processing of a connector. For example, if the remote system is undergoing maintenance, it would be preferable for source connectors to stop polling it for new data instead of filling logs with exception spam. For this use case, Connect offers a pause/resume API. While a source connector is paused, Connect will stop polling it for additional records. While a sink connector is paused, Connect will stop pushing new messages to it. The pause state is persistent, so even if you restart the cluster, the connector will not begin message processing again until the task has been resumed. Note that there may be a delay before all of a connector's tasks have transitioned to the PAUSED state since it may take time for them to finish whatever processing they were in the middle of when being paused. Additionally, failed tasks will not transition to the PAUSED state until they have been restarted.
     
        
diff --git a/docs/design.html b/docs/design.html
index 564df386db781..a1e438ee86bfd 100644
--- a/docs/design.html
+++ b/docs/design.html
@@ -16,7 +16,7 @@
 -->
 
 
 
 
        
diff --git a/docs/introduction.html b/docs/introduction.html
index 7f4c3e214365c..da79386a97ac0 100644
--- a/docs/introduction.html
+++ b/docs/introduction.html
@@ -18,197 +18,203 @@
 
 
 
 
 
        
diff --git a/docs/js/templateData.js b/docs/js/templateData.js
index 28f0eec129cd6..378b8bd5a6af9 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": "11",
-    "dotVersion": "1.1",
-    "fullDotVersion": "1.1.0",
-    "scalaVersion": "2.11"
+    "version": "28",
+    "dotVersion": "2.8",
+    "fullDotVersion": "2.8.0",
+    "scalaVersion": "2.13"
 };
diff --git a/docs/migration.html b/docs/migration.html
index 08a62715d675a..95fc87ffacafe 100644
--- a/docs/migration.html
+++ b/docs/migration.html
@@ -16,11 +16,11 @@
 -->
 
 
-
        Migrating from 0.7.x to 0.8
        
+
        Migrating from 0.7.x to 0.8
        
 
 0.8 is our first (and hopefully last) release with a non-backwards-compatible wire protocol, ZooKeeper     layout, and on-disk data format. This was a chance for us to clean up a lot of cruft and start fresh. This means performing a no-downtime upgrade is more painful than normal—you cannot just swap in the new code in-place.
 
-
        Migration Steps
        
+
        Migration Steps
        
 
 
          
     
        1. Setup a new cluster running 0.8.
diff --git a/docs/ops.html b/docs/ops.html
index 92077df81edf3..3c74273d25120 100644
--- a/docs/ops.html
+++ b/docs/ops.html
@@ -14,23 +14,20 @@
  See the License for the specific language governing permissions and
  limitations under the License.
 -->
-
+
+
+
+
          
diff --git a/docs/quickstart-zookeeper.html b/docs/quickstart-zookeeper.html
new file mode 100644
index 0000000000000..30cc832fccc18
--- /dev/null
+++ b/docs/quickstart-zookeeper.html
@@ -0,0 +1,277 @@
+
+
+
+
+
+
+
          
diff --git a/docs/quickstart.html b/docs/quickstart.html
deleted file mode 100644
index 063fec076b77d..0000000000000
--- a/docs/quickstart.html
+++ /dev/null
@@ -1,300 +0,0 @@
-
-
-
-
-
-
-
          
diff --git a/docs/security.html b/docs/security.html
index 8c1bda9f54d08..b22a165770304 100644
--- a/docs/security.html
+++ b/docs/security.html
@@ -16,7 +16,7 @@
 -->
 
 
 
 
          
diff --git a/docs/streams/architecture.html b/docs/streams/architecture.html
index 0dbb1dc97ff70..67594c232340b 100644
--- a/docs/streams/architecture.html
+++ b/docs/streams/architecture.html
@@ -19,6 +19,19 @@
 
 
 
 
 
-
          
+
          
     
     
          
-        
+        
         
         
          
     
          
diff --git a/docs/streams/core-concepts.html b/docs/streams/core-concepts.html
index dbc3f215f2cff..ddb37eac2c19e 100644
--- a/docs/streams/core-concepts.html
+++ b/docs/streams/core-concepts.html
@@ -20,16 +20,18 @@
 
 
 
 
-
          
+
          
     
     
          
-        
+        
         
         
          
     
          
diff --git a/docs/streams/developer-guide.html b/docs/streams/developer-guide.html
deleted file mode 100644
index 5730f5330e198..0000000000000
--- a/docs/streams/developer-guide.html
+++ /dev/null
@@ -1,3026 +0,0 @@
-
-
-
-
-
-
-
-
-
          
-    
-    
          
-        
-        
-        
          
-    
          
-
          
-
-
\ No newline at end of file
diff --git a/docs/streams/developer-guide/app-reset-tool.html b/docs/streams/developer-guide/app-reset-tool.html
new file mode 100644
index 0000000000000..1b90af0cb801b
--- /dev/null
+++ b/docs/streams/developer-guide/app-reset-tool.html
@@ -0,0 +1,198 @@
+
+
+
+
+
+
+
+
+
          
+  
+  
          
+    
+    
+    
          
+  
          
+
          
+
+
diff --git a/docs/streams/developer-guide/config-streams.html b/docs/streams/developer-guide/config-streams.html
new file mode 100644
index 0000000000000..6395644efb48e
--- /dev/null
+++ b/docs/streams/developer-guide/config-streams.html
@@ -0,0 +1,1032 @@
+
+
+
+
+
+
+                
+                
+                    
          
+                    
+                    
          
+                    
+                    
+                
          
+                    
          
+                    
          
+                    
+                    
diff --git a/docs/streams/developer-guide/datatypes.html b/docs/streams/developer-guide/datatypes.html
new file mode 100644
index 0000000000000..e6930bc31ce83
--- /dev/null
+++ b/docs/streams/developer-guide/datatypes.html
@@ -0,0 +1,235 @@
+
+
+
+
+
+
+
+
+
          
+  
+  
          
+    
+    
+    
          
+  
          
+
          
+
+
diff --git a/docs/streams/developer-guide/dsl-api.html b/docs/streams/developer-guide/dsl-api.html
new file mode 100644
index 0000000000000..cd734dea55d20
--- /dev/null
+++ b/docs/streams/developer-guide/dsl-api.html
@@ -0,0 +1,3908 @@
+
+
+
+
+
+
+
+
+
          
+  
+  
          
+    
+    
+    
          
+  
          
+
          
+
+
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..a3b2a53a5558d
--- /dev/null
+++ b/docs/streams/developer-guide/dsl-topology-naming.html
@@ -0,0 +1,350 @@
+
+
+
+
+
+
+
+
+
          
+	
+	
          
+		
+		
+		
          
+	
          
+
          
+
+
+
+
+
+
diff --git a/docs/streams/developer-guide/index.html b/docs/streams/developer-guide/index.html
new file mode 100644
index 0000000000000..19f638ece0daf
--- /dev/null
+++ b/docs/streams/developer-guide/index.html
@@ -0,0 +1,106 @@
+
+
+
+
+
+
+
+
+
          
+    
+    
          
+        
+        
+        
          
+    
          
+
          
+
+
diff --git a/docs/streams/developer-guide/interactive-queries.html b/docs/streams/developer-guide/interactive-queries.html
new file mode 100644
index 0000000000000..ad890e0494563
--- /dev/null
+++ b/docs/streams/developer-guide/interactive-queries.html
@@ -0,0 +1,518 @@
+
+
+
+
+
+
+                
+                
+                    
          
+                    
+                    
          
+                    
+                    
+                
          
+                    
          
+                    
          
+                    
+                    
diff --git a/docs/streams/developer-guide/manage-topics.html b/docs/streams/developer-guide/manage-topics.html
new file mode 100644
index 0000000000000..d65e375d74955
--- /dev/null
+++ b/docs/streams/developer-guide/manage-topics.html
@@ -0,0 +1,128 @@
+
+
+
+
+
+
+
+
+
          
+  
+  
          
+    
+    
+    
          
+  
          
+
          
+
+
diff --git a/docs/streams/developer-guide/memory-mgmt.html b/docs/streams/developer-guide/memory-mgmt.html
new file mode 100644
index 0000000000000..0070705a961e6
--- /dev/null
+++ b/docs/streams/developer-guide/memory-mgmt.html
@@ -0,0 +1,288 @@
+
+
+
+
+
+
+
+
+
          
+  
+  
          
+    
+    
+    
          
+  
          
+
          
+
+
diff --git a/docs/streams/developer-guide/processor-api.html b/docs/streams/developer-guide/processor-api.html
new file mode 100644
index 0000000000000..9cabac030b90c
--- /dev/null
+++ b/docs/streams/developer-guide/processor-api.html
@@ -0,0 +1,517 @@
+
+
+
+
+
+
+                
+                
+                    
          
+                    
+                    
          
+                    
+                    
+                
          
+                    
          
+                    
          
+                    
+                    
diff --git a/docs/streams/developer-guide/running-app.html b/docs/streams/developer-guide/running-app.html
new file mode 100644
index 0000000000000..87ee8f0a9da30
--- /dev/null
+++ b/docs/streams/developer-guide/running-app.html
@@ -0,0 +1,189 @@
+
+
+
+
+
+
+                
+                
+                    
          
+                    
+                    
          
+                    
+                    
+                
          
+                    
          
+                    
          
+                    
+                    
\ No newline at end of file
diff --git a/docs/streams/developer-guide/security.html b/docs/streams/developer-guide/security.html
new file mode 100644
index 0000000000000..05de0794d63a9
--- /dev/null
+++ b/docs/streams/developer-guide/security.html
@@ -0,0 +1,193 @@
+
+
+
+
+
+
+                
+                
+                    
          
+                    
+                    
          
+                    
+                    
+                
          
+                    
          
+                    
          
+                    
+                    
diff --git a/docs/streams/developer-guide/testing.html b/docs/streams/developer-guide/testing.html
new file mode 100644
index 0000000000000..79ae6e117bb50
--- /dev/null
+++ b/docs/streams/developer-guide/testing.html
@@ -0,0 +1,407 @@
+
+
+
+
+
+
+
+
+
          
+    
+    
          
+        
+        
+        
          
+    
          
+
          
+
+
diff --git a/docs/streams/developer-guide/write-streams.html b/docs/streams/developer-guide/write-streams.html
new file mode 100644
index 0000000000000..720b0c376756c
--- /dev/null
+++ b/docs/streams/developer-guide/write-streams.html
@@ -0,0 +1,255 @@
+
+
+
+
+
+
+
+
+
          
+  
+  
          
+    
+    
+    
          
+  
          
+
          
+
+
diff --git a/docs/streams/index.html b/docs/streams/index.html
index ab72c871a5885..3d84bbfee7802 100644
--- a/docs/streams/index.html
+++ b/docs/streams/index.html
@@ -15,20 +15,25 @@
 
+
 
 
 
-
          
+
          
   
   
          
-    
+    
     
     
          
@@ -308,52 +306,52 @@ 
          Hello Kafka Streams
          
 
 
\ No newline at end of file
+
diff --git a/docs/streams/quickstart.html b/docs/streams/quickstart.html
index 314bce3020c82..2cc48ef0897f6 100644
--- a/docs/streams/quickstart.html
+++ b/docs/streams/quickstart.html
@@ -18,18 +18,20 @@
 
 
 
 
 
@@ -625,13 +571,13 @@ 
          Writing a th
 
 
 
-
          
+
          
     
     
          
-        
+        
         
         
          
     
          
diff --git a/docs/streams/upgrade-guide.html b/docs/streams/upgrade-guide.html
index 297405802ca92..f24d065f57412 100644
--- a/docs/streams/upgrade-guide.html
+++ b/docs/streams/upgrade-guide.html
@@ -18,33 +18,654 @@
 
 
 
 
 
 
-
          
+
          
     
     
          
-        
+        
         
         
          
     
          
diff --git a/docs/toc.html b/docs/toc.html
index e7d939e72b776..b8c977a88985b 100644
--- a/docs/toc.html
+++ b/docs/toc.html
@@ -35,8 +35,7 @@
                 
        2. 2.2 Consumer API
                 
        3. 2.3 Streams API
                 
        4. 2.4 Connect API
-                
        5. 2.5 AdminClient API
-                
        6. 2.6 Legacy APIs
+                
        7. 2.5 Admin API
             
         
          8. 
         
        8. 3. Configuration
@@ -45,10 +44,6 @@
                 
        9. 3.2 Topic Configs
                 
        10. 3.3 Producer Configs
                 
        11. 3.4 Consumer Configs
-                    
                 
        12. 3.5 Kafka Connect Configs
                 
        13. 3.6 Kafka Streams Configs
                 
        14. 3.7 AdminClient Configs
@@ -106,6 +101,15 @@
                         
        15. Ext4 Notes
                     
                 
        16. 6.6 Monitoring
+                    
                 
        17. 6.7 ZooKeeper
                     
         
          18. 
         
        18. 8. Kafka Connect
diff --git a/docs/upgrade.html b/docs/upgrade.html
index 27c70811bad0d..d3147003bd28b 100644
--- a/docs/upgrade.html
+++ b/docs/upgrade.html
@@ -19,6 +19,829 @@
 
 
        NameKey
        ");
             b.append("" + key.name + "");
@@ -332,7 +456,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);
             }
         };
@@ -340,4 +465,10 @@ public void visit(Type field) {
         return hasBuffer.get();
     }
 
+    public static List enabledApis() {
+        return Arrays.stream(values())
+            .filter(api -> api.isEnabled)
+            .collect(Collectors.toList());
+    }
+
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/protocol/ApiMessage.java b/clients/src/main/java/org/apache/kafka/common/protocol/ApiMessage.java
new file mode 100644
index 0000000000000..4f175653bf416
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/protocol/ApiMessage.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.protocol;
+
+/**
+ * A Message which is part of the top-level Kafka API.
+ */
+public interface ApiMessage extends Message {
+    /**
+     * Returns the API key of this message, or -1 if there is none.
+     */
+    short apiKey();
+}
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
new file mode 100644
index 0000000000000..3c5c30973165f
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/protocol/ByteBufferAccessor.java
@@ -0,0 +1,143 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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 org.apache.kafka.common.utils.ByteUtils;
+
+import java.nio.ByteBuffer;
+
+public class ByteBufferAccessor implements Readable, Writable {
+    private final ByteBuffer buf;
+
+    public ByteBufferAccessor(ByteBuffer buf) {
+        this.buf = buf;
+    }
+
+    @Override
+    public byte readByte() {
+        return buf.get();
+    }
+
+    @Override
+    public short readShort() {
+        return buf.getShort();
+    }
+
+    @Override
+    public int readInt() {
+        return buf.getInt();
+    }
+
+    @Override
+    public long readLong() {
+        return buf.getLong();
+    }
+
+    @Override
+    public double readDouble() {
+        return ByteUtils.readDouble(buf);
+    }
+
+    @Override
+    public void readArray(byte[] arr) {
+        buf.get(arr);
+    }
+
+    @Override
+    public int readUnsignedVarint() {
+        return ByteUtils.readUnsignedVarint(buf);
+    }
+
+    @Override
+    public ByteBuffer readByteBuffer(int length) {
+        ByteBuffer res = buf.slice();
+        res.limit(length);
+
+        buf.position(buf.position() + length);
+
+        return res;
+    }
+
+    @Override
+    public void writeByte(byte val) {
+        buf.put(val);
+    }
+
+    @Override
+    public void writeShort(short val) {
+        buf.putShort(val);
+    }
+
+    @Override
+    public void writeInt(int val) {
+        buf.putInt(val);
+    }
+
+    @Override
+    public void writeLong(long val) {
+        buf.putLong(val);
+    }
+
+    @Override
+    public void writeDouble(double val) {
+        ByteUtils.writeDouble(val, buf);
+    }
+
+    @Override
+    public void writeByteArray(byte[] arr) {
+        buf.put(arr);
+    }
+
+    @Override
+    public void writeUnsignedVarint(int i) {
+        ByteUtils.writeUnsignedVarint(i, buf);
+    }
+
+    @Override
+    public void writeByteBuffer(ByteBuffer src) {
+        buf.put(src.duplicate());
+    }
+
+    @Override
+    public void writeVarint(int i) {
+        ByteUtils.writeVarint(i, buf);
+    }
+
+    @Override
+    public void writeVarlong(long i) {
+        ByteUtils.writeVarlong(i, buf);
+    }
+
+    @Override
+    public int readVarint() {
+        return ByteUtils.readVarint(buf);
+    }
+
+    @Override
+    public long readVarlong() {
+        return ByteUtils.readVarlong(buf);
+    }
+
+    public void flip() {
+        buf.flip();
+    }
+
+    public ByteBuffer buffer() {
+        return 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 e1d188453d13c..122d4da602324 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
@@ -26,28 +26,10 @@ public class CommonFields {
     public static final Field.Int32 PARTITION_ID = new Field.Int32("partition", "Topic partition id");
     public static final Field.Int16 ERROR_CODE = new Field.Int16("error_code", "Response error code");
     public static final Field.NullableStr ERROR_MESSAGE = new Field.NullableStr("error_message", "Response error message");
-
-    // 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.");
-    public static final Field.NullableStr NULLABLE_TRANSACTIONAL_ID = new Field.NullableStr("transactional_id",
-            "The transactional id or null if the producer is not transactional");
-    public static final Field.Int64 PRODUCER_ID = new Field.Int64("producer_id", "Current producer id in use by the transactional id.");
-    public static final Field.Int16 PRODUCER_EPOCH = new Field.Int16("producer_epoch", "Current epoch associated with the producer id.");
-
-    // ACL APIs
-    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.Str PRINCIPAL = new Field.Str("principal", "The ACL principal");
-    public static final Field.NullableStr PRINCIPAL_FILTER = new Field.NullableStr("principal", "The ACL principal filter");
-    public static final Field.Str HOST = new Field.Str("host", "The ACL host");
-    public static final Field.NullableStr HOST_FILTER = new Field.NullableStr("host", "The ACL host filter");
-    public static final Field.Int8 OPERATION = new Field.Int8("operation", "The ACL operation");
-    public static final Field.Int8 PERMISSION_TYPE = new Field.Int8("permission_type", "The ACL permission type");
+    public static final Field.Int32 LEADER_EPOCH = new Field.Int32("leader_epoch", "The leader epoch");
+    public static final Field.Int32 CURRENT_LEADER_EPOCH = new Field.Int32("current_leader_epoch",
+            "The current leader epoch, if provided, is used to fence consumers/replicas with old metadata. " +
+                    "If the epoch provided by the client is larger than the current epoch known to the broker, then " +
+                    "the UNKNOWN_LEADER_EPOCH error code will be returned. If the provided epoch is smaller, then " +
+                    "the FENCED_LEADER_EPOCH error code will be returned.");
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/protocol/DataInputStreamReadable.java b/clients/src/main/java/org/apache/kafka/common/protocol/DataInputStreamReadable.java
new file mode 100644
index 0000000000000..93c6c597d7102
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/protocol/DataInputStreamReadable.java
@@ -0,0 +1,130 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.common.protocol;
+
+import org.apache.kafka.common.utils.ByteUtils;
+
+import java.io.Closeable;
+import java.io.DataInputStream;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+
+public class DataInputStreamReadable implements Readable, Closeable {
+    protected final DataInputStream input;
+
+    public DataInputStreamReadable(DataInputStream input) {
+        this.input = input;
+    }
+
+    @Override
+    public byte readByte() {
+        try {
+            return input.readByte();
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    public short readShort() {
+        try {
+            return input.readShort();
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    public int readInt() {
+        try {
+            return input.readInt();
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    public long readLong() {
+        try {
+            return input.readLong();
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    public double readDouble() {
+        try {
+            return input.readDouble();
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    public void readArray(byte[] arr) {
+        try {
+            input.readFully(arr);
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    public int readUnsignedVarint() {
+        try {
+            return ByteUtils.readUnsignedVarint(input);
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    public ByteBuffer readByteBuffer(int length) {
+        byte[] arr = new byte[length];
+        readArray(arr);
+        return ByteBuffer.wrap(arr);
+    }
+
+    @Override
+    public int readVarint() {
+        try {
+            return ByteUtils.readVarint(input);
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    public long readVarlong() {
+        try {
+            return ByteUtils.readVarlong(input);
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    public void close() {
+        try {
+            input.close();
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/protocol/DataOutputStreamWritable.java b/clients/src/main/java/org/apache/kafka/common/protocol/DataOutputStreamWritable.java
new file mode 100644
index 0000000000000..f484016a5300a
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/protocol/DataOutputStreamWritable.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.protocol;
+
+import org.apache.kafka.common.utils.ByteUtils;
+import org.apache.kafka.common.utils.Utils;
+
+import java.io.Closeable;
+import java.io.DataOutputStream;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+
+public class DataOutputStreamWritable implements Writable, Closeable {
+    protected final DataOutputStream out;
+
+    public DataOutputStreamWritable(DataOutputStream out) {
+        this.out = out;
+    }
+
+    @Override
+    public void writeByte(byte val) {
+        try {
+            out.writeByte(val);
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    public void writeShort(short val) {
+        try {
+            out.writeShort(val);
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    public void writeInt(int val) {
+        try {
+            out.writeInt(val);
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    public void writeLong(long val) {
+        try {
+            out.writeLong(val);
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    public void writeDouble(double val) {
+        try {
+            out.writeDouble(val);
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    public void writeByteArray(byte[] arr) {
+        try {
+            out.write(arr);
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    public void writeUnsignedVarint(int i) {
+        try {
+            ByteUtils.writeUnsignedVarint(i, out);
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    public void writeByteBuffer(ByteBuffer buf) {
+        try {
+            if (buf.hasArray()) {
+                out.write(buf.array(), buf.position(), buf.limit());
+            } else {
+                byte[] bytes = Utils.toArray(buf);
+                out.write(bytes);
+            }
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    public void writeVarint(int i) {
+        try {
+            ByteUtils.writeVarint(i, out);
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    public void writeVarlong(long i) {
+        try {
+            ByteUtils.writeVarlong(i, out);
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    public void flush() {
+        try {
+            out.flush();
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    public void close() {
+        try {
+            out.close();
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+}
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 d9370549305d8..5332369f0f9be 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,8 +16,8 @@
  */
 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.SaslAuthenticationException;
 import org.apache.kafka.common.errors.BrokerNotAvailableException;
 import org.apache.kafka.common.errors.ClusterAuthorizationException;
 import org.apache.kafka.common.errors.ConcurrentTransactionsException;
@@ -25,18 +25,37 @@
 import org.apache.kafka.common.errors.CoordinatorLoadInProgressException;
 import org.apache.kafka.common.errors.CoordinatorNotAvailableException;
 import org.apache.kafka.common.errors.CorruptRecordException;
-import org.apache.kafka.common.errors.LogDirNotFoundException;
+import org.apache.kafka.common.errors.DelegationTokenAuthorizationException;
+import org.apache.kafka.common.errors.DelegationTokenDisabledException;
+import org.apache.kafka.common.errors.DelegationTokenExpiredException;
+import org.apache.kafka.common.errors.DelegationTokenNotFoundException;
+import org.apache.kafka.common.errors.DelegationTokenOwnerMismatchException;
+import org.apache.kafka.common.errors.DuplicateResourceException;
 import org.apache.kafka.common.errors.DuplicateSequenceException;
+import org.apache.kafka.common.errors.ElectionNotNeededException;
+import org.apache.kafka.common.errors.EligibleLeadersNotAvailableException;
+import org.apache.kafka.common.errors.FencedInstanceIdException;
+import org.apache.kafka.common.errors.FeatureUpdateFailedException;
+import org.apache.kafka.common.errors.FencedLeaderEpochException;
+import org.apache.kafka.common.errors.InvalidUpdateVersionException;
+import org.apache.kafka.common.errors.FetchSessionIdNotFoundException;
 import org.apache.kafka.common.errors.GroupAuthorizationException;
+import org.apache.kafka.common.errors.GroupIdNotFoundException;
+import org.apache.kafka.common.errors.GroupMaxSizeReachedException;
+import org.apache.kafka.common.errors.GroupNotEmptyException;
+import org.apache.kafka.common.errors.GroupSubscribedToTopicException;
 import org.apache.kafka.common.errors.IllegalGenerationException;
 import org.apache.kafka.common.errors.IllegalSaslStateException;
 import org.apache.kafka.common.errors.InconsistentGroupProtocolException;
+import org.apache.kafka.common.errors.InconsistentVoterSetException;
 import org.apache.kafka.common.errors.InvalidCommitOffsetSizeException;
 import org.apache.kafka.common.errors.InvalidConfigurationException;
+import org.apache.kafka.common.errors.InvalidFetchSessionEpochException;
 import org.apache.kafka.common.errors.InvalidFetchSizeException;
 import org.apache.kafka.common.errors.InvalidGroupIdException;
 import org.apache.kafka.common.errors.InvalidPartitionsException;
 import org.apache.kafka.common.errors.InvalidPidMappingException;
+import org.apache.kafka.common.errors.InvalidPrincipalTypeException;
 import org.apache.kafka.common.errors.InvalidReplicaAssignmentException;
 import org.apache.kafka.common.errors.InvalidReplicationFactorException;
 import org.apache.kafka.common.errors.InvalidRequestException;
@@ -48,42 +67,61 @@
 import org.apache.kafka.common.errors.InvalidTxnTimeoutException;
 import org.apache.kafka.common.errors.KafkaStorageException;
 import org.apache.kafka.common.errors.LeaderNotAvailableException;
+import org.apache.kafka.common.errors.ListenerNotFoundException;
+import org.apache.kafka.common.errors.LogDirNotFoundException;
+import org.apache.kafka.common.errors.MemberIdRequiredException;
 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;
 import org.apache.kafka.common.errors.NotEnoughReplicasException;
-import org.apache.kafka.common.errors.NotLeaderForPartitionException;
+import org.apache.kafka.common.errors.NotLeaderOrFollowerException;
 import org.apache.kafka.common.errors.OffsetMetadataTooLarge;
+import org.apache.kafka.common.errors.OffsetNotAvailableException;
 import org.apache.kafka.common.errors.OffsetOutOfRangeException;
 import org.apache.kafka.common.errors.OperationNotAttemptedException;
 import org.apache.kafka.common.errors.OutOfOrderSequenceException;
 import org.apache.kafka.common.errors.PolicyViolationException;
+import org.apache.kafka.common.errors.PreferredLeaderNotAvailableException;
+import org.apache.kafka.common.errors.PrincipalDeserializationException;
 import org.apache.kafka.common.errors.ProducerFencedException;
 import org.apache.kafka.common.errors.ReassignmentInProgressException;
 import org.apache.kafka.common.errors.RebalanceInProgressException;
 import org.apache.kafka.common.errors.RecordBatchTooLargeException;
 import org.apache.kafka.common.errors.RecordTooLargeException;
 import org.apache.kafka.common.errors.ReplicaNotAvailableException;
+import org.apache.kafka.common.errors.ResourceNotFoundException;
 import org.apache.kafka.common.errors.RetriableException;
+import org.apache.kafka.common.errors.SaslAuthenticationException;
 import org.apache.kafka.common.errors.SecurityDisabledException;
+import org.apache.kafka.common.errors.StaleBrokerEpochException;
+import org.apache.kafka.common.errors.ThrottlingQuotaExceededException;
 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.TopicExistsException;
-import org.apache.kafka.common.errors.TransactionalIdAuthorizationException;
 import org.apache.kafka.common.errors.TransactionCoordinatorFencedException;
+import org.apache.kafka.common.errors.TransactionalIdAuthorizationException;
+import org.apache.kafka.common.errors.UnacceptableCredentialException;
+import org.apache.kafka.common.errors.UnknownLeaderEpochException;
 import org.apache.kafka.common.errors.UnknownMemberIdException;
 import org.apache.kafka.common.errors.UnknownProducerIdException;
 import org.apache.kafka.common.errors.UnknownServerException;
 import org.apache.kafka.common.errors.UnknownTopicOrPartitionException;
+import org.apache.kafka.common.errors.UnstableOffsetCommitException;
+import org.apache.kafka.common.errors.UnsupportedByAuthenticationException;
+import org.apache.kafka.common.errors.UnsupportedCompressionTypeException;
 import org.apache.kafka.common.errors.UnsupportedForMessageFormatException;
 import org.apache.kafka.common.errors.UnsupportedSaslMechanismException;
 import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.errors.InvalidProducerEpochException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.util.HashMap;
 import java.util.Map;
+import java.util.function.Function;
 
 /**
  * This class contains all the client-server errors--those errors that must be sent from the server to the client. These
@@ -97,459 +135,213 @@
  * Do not add exceptions that occur only on the client or only on the server here.
  */
 public enum Errors {
-    UNKNOWN_SERVER_ERROR(-1, "The server experienced an unexpected error when processing the request",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new UnknownServerException(message);
-            }
-        }),
-    NONE(0, null,
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return null;
-            }
-        }),
+    UNKNOWN_SERVER_ERROR(-1, "The server experienced an unexpected error when processing the request.",
+            UnknownServerException::new),
+    NONE(0, null, message -> null),
     OFFSET_OUT_OF_RANGE(1, "The requested offset is not within the range of offsets maintained by the server.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new OffsetOutOfRangeException(message);
-            }
-        }),
-    CORRUPT_MESSAGE(2, "This message has failed its CRC checksum, exceeds the valid size, or is otherwise corrupt.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new CorruptRecordException(message);
-            }
-        }),
+            OffsetOutOfRangeException::new),
+    CORRUPT_MESSAGE(2, "This message has failed its CRC checksum, exceeds the valid size, has a null key for a compacted topic, or is otherwise corrupt.",
+            CorruptRecordException::new),
     UNKNOWN_TOPIC_OR_PARTITION(3, "This server does not host this topic-partition.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new UnknownTopicOrPartitionException(message);
-            }
-        }),
+            UnknownTopicOrPartitionException::new),
     INVALID_FETCH_SIZE(4, "The requested fetch size is invalid.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new InvalidFetchSizeException(message);
-            }
-        }),
+            InvalidFetchSizeException::new),
     LEADER_NOT_AVAILABLE(5, "There is no leader for this topic-partition as we are in the middle of a leadership election.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new LeaderNotAvailableException(message);
-            }
-        }),
-    NOT_LEADER_FOR_PARTITION(6, "This server is not the leader for that topic-partition.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new NotLeaderForPartitionException(message);
-            }
-        }),
+            LeaderNotAvailableException::new),
+    NOT_LEADER_OR_FOLLOWER(6, "For requests intended only for the leader, this error indicates that the broker is not the current leader. " +
+            "For requests intended for any replica, this error indicates that the broker is not a replica of the topic partition.",
+            NotLeaderOrFollowerException::new),
     REQUEST_TIMED_OUT(7, "The request timed out.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new TimeoutException(message);
-            }
-        }),
+            TimeoutException::new),
     BROKER_NOT_AVAILABLE(8, "The broker is not available.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new BrokerNotAvailableException(message);
-            }
-        }),
-    REPLICA_NOT_AVAILABLE(9, "The replica is not available for the requested topic-partition",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new ReplicaNotAvailableException(message);
-            }
-        }),
+            BrokerNotAvailableException::new),
+    REPLICA_NOT_AVAILABLE(9, "The replica is not available for the requested topic-partition. Produce/Fetch requests and other requests " +
+            "intended only for the leader or follower return NOT_LEADER_OR_FOLLOWER if the broker is not a replica of the topic-partition.",
+            ReplicaNotAvailableException::new),
     MESSAGE_TOO_LARGE(10, "The request included a message larger than the max message size the server will accept.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new RecordTooLargeException(message);
-            }
-        }),
+            RecordTooLargeException::new),
     STALE_CONTROLLER_EPOCH(11, "The controller moved to another broker.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new ControllerMovedException(message);
-            }
-        }),
+            ControllerMovedException::new),
     OFFSET_METADATA_TOO_LARGE(12, "The metadata field of the offset request was too large.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new OffsetMetadataTooLarge(message);
-            }
-        }),
+            OffsetMetadataTooLarge::new),
     NETWORK_EXCEPTION(13, "The server disconnected before a response was received.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new NetworkException(message);
-            }
-        }),
+            NetworkException::new),
     COORDINATOR_LOAD_IN_PROGRESS(14, "The coordinator is loading and hence can't process requests.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new CoordinatorLoadInProgressException(message);
-            }
-        }),
+            CoordinatorLoadInProgressException::new),
     COORDINATOR_NOT_AVAILABLE(15, "The coordinator is not available.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new CoordinatorNotAvailableException(message);
-            }
-        }),
+            CoordinatorNotAvailableException::new),
     NOT_COORDINATOR(16, "This is not the correct coordinator.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new NotCoordinatorException(message);
-            }
-        }),
+            NotCoordinatorException::new),
     INVALID_TOPIC_EXCEPTION(17, "The request attempted to perform an operation on an invalid topic.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new InvalidTopicException(message);
-            }
-        }),
+            InvalidTopicException::new),
     RECORD_LIST_TOO_LARGE(18, "The request included message batch larger than the configured segment size on the server.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new RecordBatchTooLargeException(message);
-            }
-        }),
+            RecordBatchTooLargeException::new),
     NOT_ENOUGH_REPLICAS(19, "Messages are rejected since there are fewer in-sync replicas than required.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new NotEnoughReplicasException(message);
-            }
-        }),
+            NotEnoughReplicasException::new),
     NOT_ENOUGH_REPLICAS_AFTER_APPEND(20, "Messages are written to the log, but to fewer in-sync replicas than required.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new NotEnoughReplicasAfterAppendException(message);
-            }
-        }),
+            NotEnoughReplicasAfterAppendException::new),
     INVALID_REQUIRED_ACKS(21, "Produce request specified an invalid value for required acks.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new InvalidRequiredAcksException(message);
-            }
-        }),
+            InvalidRequiredAcksException::new),
     ILLEGAL_GENERATION(22, "Specified group generation id is not valid.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new IllegalGenerationException(message);
-            }
-        }),
+            IllegalGenerationException::new),
     INCONSISTENT_GROUP_PROTOCOL(23,
-            "The group member's supported protocols are incompatible with those of existing members" +
-                " or first group member tried to join with empty protocol type or empty protocol list.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new InconsistentGroupProtocolException(message);
-            }
-        }),
-    INVALID_GROUP_ID(24, "The configured groupId is invalid",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new InvalidGroupIdException(message);
-            }
-        }),
+            "The group member's supported protocols are incompatible with those of existing members " +
+            "or first group member tried to join with empty protocol type or empty protocol list.",
+            InconsistentGroupProtocolException::new),
+    INVALID_GROUP_ID(24, "The configured groupId is invalid.",
+            InvalidGroupIdException::new),
     UNKNOWN_MEMBER_ID(25, "The coordinator is not aware of this member.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new UnknownMemberIdException(message);
-            }
-        }),
+            UnknownMemberIdException::new),
     INVALID_SESSION_TIMEOUT(26,
             "The session timeout is not within the range allowed by the broker " +
             "(as configured by group.min.session.timeout.ms and group.max.session.timeout.ms).",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new InvalidSessionTimeoutException(message);
-            }
-        }),
+            InvalidSessionTimeoutException::new),
     REBALANCE_IN_PROGRESS(27, "The group is rebalancing, so a rejoin is needed.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new RebalanceInProgressException(message);
-            }
-        }),
-    INVALID_COMMIT_OFFSET_SIZE(28, "The committing offset data size is not valid",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new InvalidCommitOffsetSizeException(message);
-            }
-        }),
-    TOPIC_AUTHORIZATION_FAILED(29, "Topic authorization failed.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new TopicAuthorizationException(message);
-            }
-        }),
-    GROUP_AUTHORIZATION_FAILED(30, "Group authorization failed.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new GroupAuthorizationException(message);
-            }
-        }),
+            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),
     CLUSTER_AUTHORIZATION_FAILED(31, "Cluster authorization failed.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new ClusterAuthorizationException(message);
-            }
-        }),
+            ClusterAuthorizationException::new),
     INVALID_TIMESTAMP(32, "The timestamp of the message is out of acceptable range.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new InvalidTimestampException(message);
-            }
-        }),
+            InvalidTimestampException::new),
     UNSUPPORTED_SASL_MECHANISM(33, "The broker does not support the requested SASL mechanism.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new UnsupportedSaslMechanismException(message);
-            }
-        }),
+            UnsupportedSaslMechanismException::new),
     ILLEGAL_SASL_STATE(34, "Request is not valid given the current SASL state.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new IllegalSaslStateException(message);
-            }
-        }),
+            IllegalSaslStateException::new),
     UNSUPPORTED_VERSION(35, "The version of API is not supported.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new UnsupportedVersionException(message);
-            }
-        }),
+            UnsupportedVersionException::new),
     TOPIC_ALREADY_EXISTS(36, "Topic with this name already exists.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new TopicExistsException(message);
-            }
-        }),
-    INVALID_PARTITIONS(37, "Number of partitions is invalid.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new InvalidPartitionsException(message);
-            }
-        }),
-    INVALID_REPLICATION_FACTOR(38, "Replication-factor is invalid.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new InvalidReplicationFactorException(message);
-            }
-        }),
+            TopicExistsException::new),
+    INVALID_PARTITIONS(37, "Number of partitions is below 1.",
+            InvalidPartitionsException::new),
+    INVALID_REPLICATION_FACTOR(38, "Replication factor is below 1 or larger than the number of available brokers.",
+            InvalidReplicationFactorException::new),
     INVALID_REPLICA_ASSIGNMENT(39, "Replica assignment is invalid.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new InvalidReplicaAssignmentException(message);
-            }
-        }),
+            InvalidReplicaAssignmentException::new),
     INVALID_CONFIG(40, "Configuration is invalid.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new InvalidConfigurationException(message);
-            }
-        }),
+            InvalidConfigurationException::new),
     NOT_CONTROLLER(41, "This is not the correct controller for this cluster.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new NotControllerException(message);
-            }
-        }),
+            NotControllerException::new),
     INVALID_REQUEST(42, "This most likely occurs because of a request being malformed by the " +
-                "client library or the message was sent to an incompatible broker. See the broker logs " +
-                "for more details.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new InvalidRequestException(message);
-            }
-        }),
+            "client library or the message was sent to an incompatible broker. See the broker logs " +
+            "for more details.",
+            InvalidRequestException::new),
     UNSUPPORTED_FOR_MESSAGE_FORMAT(43, "The message format version on the broker does not support the request.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new UnsupportedForMessageFormatException(message);
-            }
-        }),
+            UnsupportedForMessageFormatException::new),
     POLICY_VIOLATION(44, "Request parameters do not satisfy the configured policy.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new PolicyViolationException(message);
-            }
-        }),
-    OUT_OF_ORDER_SEQUENCE_NUMBER(45, "The broker received an out of order sequence number",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new OutOfOrderSequenceException(message);
-            }
-        }),
-    DUPLICATE_SEQUENCE_NUMBER(46, "The broker received a duplicate sequence number",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new DuplicateSequenceException(message);
-            }
-        }),
-    INVALID_PRODUCER_EPOCH(47, "Producer attempted an operation with an old epoch. Either there is a newer producer " +
-            "with the same transactionalId, or the producer's transaction has been expired by the broker.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new ProducerFencedException(message);
-            }
-        }),
-    INVALID_TXN_STATE(48, "The producer attempted a transactional operation in an invalid state",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new InvalidTxnStateException(message);
-            }
-        }),
+            PolicyViolationException::new),
+    OUT_OF_ORDER_SEQUENCE_NUMBER(45, "The broker received an out of order sequence number.",
+            OutOfOrderSequenceException::new),
+    DUPLICATE_SEQUENCE_NUMBER(46, "The broker received a duplicate sequence number.",
+            DuplicateSequenceException::new),
+    INVALID_PRODUCER_EPOCH(47, "Producer attempted to produce with an old epoch.",
+            InvalidProducerEpochException::new),
+    INVALID_TXN_STATE(48, "The producer attempted a transactional operation in an invalid state.",
+            InvalidTxnStateException::new),
     INVALID_PRODUCER_ID_MAPPING(49, "The producer attempted to use a producer id which is not currently assigned to " +
-            "its transactional id",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new InvalidPidMappingException(message);
-            }
-        }),
+            "its transactional id.",
+            InvalidPidMappingException::new),
     INVALID_TRANSACTION_TIMEOUT(50, "The transaction timeout is larger than the maximum value allowed by " +
-                "the broker (as configured by max.transaction.timeout.ms).",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new InvalidTxnTimeoutException(message);
-            }
-        }),
+            "the broker (as configured by transaction.max.timeout.ms).",
+            InvalidTxnTimeoutException::new),
     CONCURRENT_TRANSACTIONS(51, "The producer attempted to update a transaction " +
-                "while another concurrent operation on the same transaction was ongoing",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new ConcurrentTransactionsException(message);
-            }
-        }),
+            "while another concurrent operation on the same transaction was ongoing.",
+            ConcurrentTransactionsException::new),
     TRANSACTION_COORDINATOR_FENCED(52, "Indicates that the transaction coordinator sending a WriteTxnMarker " +
-            "is no longer the current coordinator for a given producer",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new TransactionCoordinatorFencedException(message);
-            }
-        }),
-    TRANSACTIONAL_ID_AUTHORIZATION_FAILED(53, "Transactional Id authorization failed",
-                                                  new ApiExceptionBuilder() {
-        @Override
-        public ApiException build(String message) {
-            return new TransactionalIdAuthorizationException(message);
-        }
-    }),
-    SECURITY_DISABLED(54, "Security features are disabled.", new ApiExceptionBuilder() {
-        @Override
-        public ApiException build(String message) {
-            return new SecurityDisabledException(message);
-        }
-    }),
-    OPERATION_NOT_ATTEMPTED(55, "The broker did not attempt to execute this operation. This may happen for batched RPCs " +
-            "where some operations in the batch failed, causing the broker to respond without trying the rest.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new OperationNotAttemptedException(message);
-            }
-    }),
+            "is no longer the current coordinator for a given producer.",
+            TransactionCoordinatorFencedException::new),
+    TRANSACTIONAL_ID_AUTHORIZATION_FAILED(53, "Transactional Id authorization failed.",
+            TransactionalIdAuthorizationException::new),
+    SECURITY_DISABLED(54, "Security features are disabled.",
+            SecurityDisabledException::new),
+    OPERATION_NOT_ATTEMPTED(55, "The broker did not attempt to execute this operation. This may happen for " +
+            "batched RPCs where some operations in the batch failed, causing the broker to respond without " +
+            "trying the rest.",
+            OperationNotAttemptedException::new),
     KAFKA_STORAGE_ERROR(56, "Disk error when trying to access log file on the disk.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new KafkaStorageException(message);
-            }
-    }),
+            KafkaStorageException::new),
     LOG_DIR_NOT_FOUND(57, "The user-specified log directory is not found in the broker config.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new LogDirNotFoundException(message);
-            }
-    }),
+            LogDirNotFoundException::new),
     SASL_AUTHENTICATION_FAILED(58, "SASL Authentication failed.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new SaslAuthenticationException(message);
-            }
-    }),
+            SaslAuthenticationException::new),
     UNKNOWN_PRODUCER_ID(59, "This exception is raised by the broker if it could not locate the producer metadata " +
             "associated with the producerId in question. This could happen if, for instance, the producer's records " +
             "were deleted because their retention time had elapsed. Once the last records of the producerId are " +
             "removed, the producer's metadata is removed from the broker, and future appends by the producer will " +
             "return this exception.",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new UnknownProducerIdException(message);
-            }
-    }),
-    REASSIGNMENT_IN_PROGRESS(60, "A partition reassignment is in progress",
-        new ApiExceptionBuilder() {
-            @Override
-            public ApiException build(String message) {
-                return new ReassignmentInProgressException(message);
-            }
-    });
-
-    private interface ApiExceptionBuilder {
-        ApiException build(String message);
-    }
+            UnknownProducerIdException::new),
+    REASSIGNMENT_IN_PROGRESS(60, "A partition reassignment is in progress.",
+            ReassignmentInProgressException::new),
+    DELEGATION_TOKEN_AUTH_DISABLED(61, "Delegation Token feature is not enabled.",
+            DelegationTokenDisabledException::new),
+    DELEGATION_TOKEN_NOT_FOUND(62, "Delegation Token is not found on server.",
+            DelegationTokenNotFoundException::new),
+    DELEGATION_TOKEN_OWNER_MISMATCH(63, "Specified Principal is not valid Owner/Renewer.",
+            DelegationTokenOwnerMismatchException::new),
+    DELEGATION_TOKEN_REQUEST_NOT_ALLOWED(64, "Delegation Token requests are not allowed on PLAINTEXT/1-way SSL " +
+            "channels and on delegation token authenticated channels.",
+            UnsupportedByAuthenticationException::new),
+    DELEGATION_TOKEN_AUTHORIZATION_FAILED(65, "Delegation Token authorization failed.",
+            DelegationTokenAuthorizationException::new),
+    DELEGATION_TOKEN_EXPIRED(66, "Delegation Token is expired.",
+            DelegationTokenExpiredException::new),
+    INVALID_PRINCIPAL_TYPE(67, "Supplied principalType is not supported.",
+            InvalidPrincipalTypeException::new),
+    NON_EMPTY_GROUP(68, "The group is not empty.",
+            GroupNotEmptyException::new),
+    GROUP_ID_NOT_FOUND(69, "The group id does not exist.",
+            GroupIdNotFoundException::new),
+    FETCH_SESSION_ID_NOT_FOUND(70, "The fetch session ID was not found.",
+            FetchSessionIdNotFoundException::new),
+    INVALID_FETCH_SESSION_EPOCH(71, "The fetch session epoch is invalid.",
+            InvalidFetchSessionEpochException::new),
+    LISTENER_NOT_FOUND(72, "There is no listener on the leader broker that matches the listener on which " +
+            "metadata request was processed.",
+            ListenerNotFoundException::new),
+    TOPIC_DELETION_DISABLED(73, "Topic deletion is disabled.",
+            TopicDeletionDisabledException::new),
+    FENCED_LEADER_EPOCH(74, "The leader epoch in the request is older than the epoch on the broker.",
+            FencedLeaderEpochException::new),
+    UNKNOWN_LEADER_EPOCH(75, "The leader epoch in the request is newer than the epoch on the broker.",
+            UnknownLeaderEpochException::new),
+    UNSUPPORTED_COMPRESSION_TYPE(76, "The requesting client does not support the compression type of given partition.",
+            UnsupportedCompressionTypeException::new),
+    STALE_BROKER_EPOCH(77, "Broker epoch has changed.",
+            StaleBrokerEpochException::new),
+    OFFSET_NOT_AVAILABLE(78, "The leader high watermark has not caught up from a recent leader " +
+            "election so the offsets cannot be guaranteed to be monotonically increasing.",
+            OffsetNotAvailableException::new),
+    MEMBER_ID_REQUIRED(79, "The group member needs to have a valid member id before actually entering a consumer group.",
+            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),
+    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 will be rejected.", InvalidRecordException::new),
+    UNSTABLE_OFFSET_COMMIT(88, "There are unstable offsets that need to be cleared.", UnstableOffsetCommitException::new),
+    THROTTLING_QUOTA_EXCEEDED(89, "The throttling quota has been exceeded.", ThrottlingQuotaExceededException::new),
+    PRODUCER_FENCED(90, "There is a newer producer with the same transactionalId " +
+            "which fences the current one.", ProducerFencedException::new),
+    RESOURCE_NOT_FOUND(91, "A request illegally referred to a resource that does not exist.", ResourceNotFoundException::new),
+    DUPLICATE_RESOURCE(92, "A request illegally referred to the same resource twice.", DuplicateResourceException::new),
+    UNACCEPTABLE_CREDENTIAL(93, "Requested credential would not meet criteria for acceptability.", UnacceptableCredentialException::new),
+    INCONSISTENT_VOTER_SET(94, "Indicates that the either the sender or recipient of a " +
+            "voter-only request is not one of the expected voters", InconsistentVoterSetException::new),
+    INVALID_UPDATE_VERSION(95, "The given update version was invalid.", InvalidUpdateVersionException::new),
+    FEATURE_UPDATE_FAILED(96, "Unable to update finalized features due to an unexpected server error.", FeatureUpdateFailedException::new),
+    PRINCIPAL_DESERIALIZATION_FAILURE(97, "Request principal deserialization failed during forwarding. " +
+         "This indicates an internal error on the broker cluster security setup.", PrincipalDeserializationException::new);
 
     private static final Logger log = LoggerFactory.getLogger(Errors.class);
 
@@ -558,20 +350,23 @@ private interface ApiExceptionBuilder {
 
     static {
         for (Errors error : Errors.values()) {
-            codeToError.put(error.code(), error);
+            if (codeToError.put(error.code(), error) != null)
+                throw new ExceptionInInitializerError("Code " + error.code() + " for error " +
+                        error + " has already been used");
+
             if (error.exception != null)
                 classToError.put(error.exception.getClass(), error);
         }
     }
 
     private final short code;
-    private final ApiExceptionBuilder builder;
+    private final Function builder;
     private final ApiException exception;
 
-    Errors(int code, String defaultExceptionString, ApiExceptionBuilder builder) {
+    Errors(int code, String defaultExceptionString, Function builder) {
         this.code = (short) code;
         this.builder = builder;
-        this.exception = builder.build(defaultExceptionString);
+        this.exception = builder.apply(defaultExceptionString);
     }
 
     /**
@@ -593,7 +388,7 @@ public ApiException exception(String message) {
             return exception;
         }
         // Return an exception with the given error message.
-        return builder.build(message);
+        return builder.apply(message);
     }
 
     /**
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
new file mode 100644
index 0000000000000..75641d3c0e2d8
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/protocol/Message.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.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.
+ */
+public interface Message {
+    /**
+     * Returns the lowest supported API key of this message, inclusive.
+     */
+    short lowestSupportedVersion();
+
+    /**
+     * Returns the highest supported API key of this message, inclusive.
+     */
+    short highestSupportedVersion();
+
+    /**
+     * 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.
+     */
+    default int size(ObjectSerializationCache cache, short version) {
+        MessageSizeAccumulator size = new MessageSizeAccumulator();
+        addSize(size, cache, version);
+        return size.totalSize();
+    }
+
+    /**
+     * Add the size of this message to an accumulator.
+     *
+     * @param size          The size accumulator to add to
+     * @param cache         The serialization size cache to populate.
+     * @param version       The version to use.
+     */
+    void addSize(MessageSizeAccumulator size, ObjectSerializationCache cache, short version);
+
+    /**
+     * 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, ObjectSerializationCache cache, short version);
+
+    /**
+     * Reads this message from the given Readable.  This will overwrite all
+     * relevant fields with information from the byte buffer.
+     *
+     * @param readable      The source readable.
+     * @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 read(Readable readable, short version);
+
+    /**
+     * Reads this message from a Struct object.  This will overwrite all
+     * relevant fields with information from the Struct.
+     *
+     * @param struct        The source struct.
+     * @param version       The version to use.
+     *
+     * @throws {@see org.apache.kafka.common.errors.UnsupportedVersionException}
+     *                      If the specified struct can't be processed with the
+     *                      specified message version.
+     */
+    void fromStruct(Struct struct, short version);
+
+    /**
+     * Writes out this message to a Struct.
+     *
+     * @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.
+     */
+    Struct toStruct(short version);
+
+    /**
+     * Returns a list of tagged fields which this software can't understand.
+     *
+     * @return              The raw tagged fields.
+     */
+    List unknownTaggedFields();
+
+    /**
+     * Make a deep copy of the message.
+     *
+     * @return              A copy of the message which does not share any mutable fields.
+     */
+    Message duplicate();
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/protocol/MessageSizeAccumulator.java b/clients/src/main/java/org/apache/kafka/common/protocol/MessageSizeAccumulator.java
new file mode 100644
index 0000000000000..dac007ec5551f
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/protocol/MessageSizeAccumulator.java
@@ -0,0 +1,68 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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;
+
+/**
+ * Helper class which facilitates zero-copy network transmission. See {@link SendBuilder}.
+ */
+public class MessageSizeAccumulator {
+    private int totalSize = 0;
+    private int zeroCopySize = 0;
+
+    /**
+     * Get the total size of the message.
+     *
+     * @return total size in bytes
+     */
+    public int totalSize() {
+        return totalSize;
+    }
+
+    /**
+     * Size excluding zero copy fields as specified by {@link #zeroCopySize}. This is typically the size of the byte
+     * buffer used to serialize messages.
+     */
+    public int sizeExcludingZeroCopy() {
+        return totalSize - zeroCopySize;
+    }
+
+    /**
+     * Get the total "zero-copy" size of the message. This is the summed
+     * total of all fields which have either have a type of 'bytes' with
+     * 'zeroCopy' enabled, or a type of 'records'
+     *
+     * @return total size of zero-copy data in the message
+     */
+    public int zeroCopySize() {
+        return zeroCopySize;
+    }
+
+    public void addZeroCopyBytes(int size) {
+        zeroCopySize += size;
+        totalSize += size;
+    }
+
+    public void addBytes(int size) {
+        totalSize += size;
+    }
+
+    public void add(MessageSizeAccumulator size) {
+        this.totalSize += size.totalSize;
+        this.zeroCopySize += size.zeroCopySize;
+    }
+
+}
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
new file mode 100644
index 0000000000000..d3812bbf5bcdc
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/protocol/MessageUtil.java
@@ -0,0 +1,211 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.kafka.common.protocol;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import org.apache.kafka.common.protocol.types.RawTaggedField;
+import org.apache.kafka.common.utils.Utils;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Arrays;
+import java.util.Iterator;
+import java.util.List;
+import java.util.UUID;
+
+public final class MessageUtil {
+    public static final UUID ZERO_UUID = new UUID(0L, 0L);
+
+    /**
+     * Copy a byte buffer into an array.  This will not affect the buffer's
+     * position or mark.
+     */
+    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 arr;
+    }
+
+    public static String deepToString(Iterator iter) {
+        StringBuilder bld = new StringBuilder("[");
+        String prefix = "";
+        while (iter.hasNext()) {
+            Object object = iter.next();
+            bld.append(prefix);
+            bld.append(object.toString());
+            prefix = ", ";
+        }
+        bld.append("]");
+        return bld.toString();
+    }
+
+    public static byte jsonNodeToByte(JsonNode node, String about) {
+        int value = jsonNodeToInt(node, about);
+        if (value > Byte.MAX_VALUE) {
+            if (value <= 256) {
+                // It's more traditional to refer to bytes as unsigned,
+                // so we support that here.
+                value -= 128;
+            } else {
+                throw new RuntimeException(about + ": value " + value +
+                    " does not fit in an 8-bit signed integer.");
+            }
+        }
+        if (value < Byte.MIN_VALUE) {
+            throw new RuntimeException(about + ": value " + value +
+                " does not fit in an 8-bit signed integer.");
+        }
+        return (byte) value;
+    }
+
+    public static short jsonNodeToShort(JsonNode node, String about) {
+        int value = jsonNodeToInt(node, about);
+        if ((value < Short.MIN_VALUE) || (value > Short.MAX_VALUE)) {
+            throw new RuntimeException(about + ": value " + value +
+                " does not fit in a 16-bit signed integer.");
+        }
+        return (short) value;
+    }
+
+    public static int jsonNodeToInt(JsonNode node, String about) {
+        if (node.isInt()) {
+            return node.asInt();
+        }
+        if (node.isTextual()) {
+            throw new NumberFormatException(about + ": expected an integer or " +
+                "string type, but got " + node.getNodeType());
+        }
+        String text = node.asText();
+        if (text.startsWith("0x")) {
+            try {
+                return Integer.parseInt(text.substring(2), 16);
+            } catch (NumberFormatException e) {
+                throw new NumberFormatException(about + ": failed to " +
+                    "parse hexadecimal number: " + e.getMessage());
+            }
+        } else {
+            try {
+                return Integer.parseInt(text);
+            } catch (NumberFormatException e) {
+                throw new NumberFormatException(about + ": failed to " +
+                    "parse number: " + e.getMessage());
+            }
+        }
+    }
+
+    public static long jsonNodeToLong(JsonNode node, String about) {
+        if (node.isLong()) {
+            return node.asLong();
+        }
+        if (node.isTextual()) {
+            throw new NumberFormatException(about + ": expected an integer or " +
+                "string type, but got " + node.getNodeType());
+        }
+        String text = node.asText();
+        if (text.startsWith("0x")) {
+            try {
+                return Long.parseLong(text.substring(2), 16);
+            } catch (NumberFormatException e) {
+                throw new NumberFormatException(about + ": failed to " +
+                    "parse hexadecimal number: " + e.getMessage());
+            }
+        } else {
+            try {
+                return Long.parseLong(text);
+            } catch (NumberFormatException e) {
+                throw new NumberFormatException(about + ": failed to " +
+                    "parse number: " + e.getMessage());
+            }
+        }
+    }
+
+    public static byte[] jsonNodeToBinary(JsonNode node, String about) {
+        if (!node.isBinary()) {
+            throw new RuntimeException(about + ": expected Base64-encoded binary data.");
+        }
+        try {
+            byte[] value = node.binaryValue();
+            return value;
+        } catch (IOException e) {
+            throw new RuntimeException(about + ": unable to retrieve Base64-encoded binary data", e);
+        }
+    }
+
+    public static double jsonNodeToDouble(JsonNode node, String about) {
+        if (!node.isFloatingPointNumber()) {
+            throw new NumberFormatException(about + ": expected a floating point " +
+                "type, but got " + node.getNodeType());
+        }
+        return node.asDouble();
+    }
+
+    public static byte[] duplicate(byte[] array) {
+        if (array == null)
+            return null;
+        return Arrays.copyOf(array, array.length);
+    }
+
+    /**
+     * Compare two RawTaggedFields lists.
+     * A null list is equivalent to an empty one in this context.
+     */
+    public static boolean compareRawTaggedFields(List first,
+                                                 List second) {
+        if (first == null) {
+            return second == null || second.isEmpty();
+        } else if (second == null) {
+            return first.isEmpty();
+        } else {
+            return first.equals(second);
+        }
+    }
+
+    public static ByteBuffer toByteBuffer(final Message message, final short version) {
+        ObjectSerializationCache cache = new ObjectSerializationCache();
+        int messageSize = message.size(cache, version);
+        ByteBufferAccessor bytes = new ByteBufferAccessor(ByteBuffer.allocate(messageSize));
+        message.write(bytes, cache, version);
+        bytes.flip();
+        return bytes.buffer();
+    }
+
+    public static ByteBuffer toVersionPrefixedByteBuffer(final short version, final Message message) {
+        ObjectSerializationCache cache = new ObjectSerializationCache();
+        int messageSize = message.size(cache, version);
+        ByteBufferAccessor bytes = new ByteBufferAccessor(ByteBuffer.allocate(messageSize + 2));
+        bytes.writeShort(version);
+        message.write(bytes, cache, version);
+        bytes.flip();
+        return bytes.buffer();
+    }
+
+    public static byte[] toVersionPrefixedBytes(final short version, final Message message) {
+        ByteBuffer buffer = toVersionPrefixedByteBuffer(version, message);
+        // take the inner array directly if it is full with data
+        if (buffer.hasArray() &&
+                buffer.arrayOffset() == 0 &&
+                buffer.position() == 0 &&
+                buffer.limit() == buffer.array().length) return buffer.array();
+        else return Utils.toArray(buffer);
+    }
+}
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..208b0567737ee
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/protocol/ObjectSerializationCache.java
@@ -0,0 +1,55 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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, Integer size) {
+        map.put(o, size);
+    }
+
+    public Integer getArraySizeInBytes(Object o) {
+        return (Integer) map.get(o);
+    }
+
+    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..6e7d49380f9af 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,19 +119,21 @@ 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 (ApiKeys key : ApiKeys.values()) {
+        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.enabledApis()) {
             // Key
             b.append("
        ");
             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
new file mode 100644
index 0000000000000..8bf9be8883bfc
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/protocol/Readable.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.common.protocol;
+
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.common.protocol.types.RawTaggedField;
+import org.apache.kafka.common.record.MemoryRecords;
+
+import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.List;
+
+public interface Readable {
+    byte readByte();
+    short readShort();
+    int readInt();
+    long readLong();
+    double readDouble();
+    void readArray(byte[] arr);
+    int readUnsignedVarint();
+    ByteBuffer readByteBuffer(int length);
+    int readVarint();
+    long readVarlong();
+
+    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;
+    }
+
+    default MemoryRecords readRecords(int length) {
+        if (length < 0) {
+            // no records
+            return null;
+        } else {
+            ByteBuffer recordsBuffer = readByteBuffer(length);
+            return MemoryRecords.readableRecords(recordsBuffer);
+        }
+    }
+
+    /**
+     * Read a UUID with the most significant digits first.
+     */
+    default Uuid readUuid() {
+        return new Uuid(readLong(), readLong());
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/protocol/SendBuilder.java b/clients/src/main/java/org/apache/kafka/common/protocol/SendBuilder.java
new file mode 100644
index 0000000000000..b76e0d8dea833
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/protocol/SendBuilder.java
@@ -0,0 +1,224 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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 org.apache.kafka.common.network.ByteBufferSend;
+import org.apache.kafka.common.network.Send;
+import org.apache.kafka.common.record.BaseRecords;
+import org.apache.kafka.common.record.MemoryRecords;
+import org.apache.kafka.common.record.MultiRecordsSend;
+import org.apache.kafka.common.requests.RequestHeader;
+import org.apache.kafka.common.requests.ResponseHeader;
+import org.apache.kafka.common.utils.ByteUtils;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayDeque;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Queue;
+
+/**
+ * This class provides a way to build {@link Send} objects for network transmission
+ * from generated {@link org.apache.kafka.common.protocol.ApiMessage} types without
+ * allocating new space for "zero-copy" fields (see {@link #writeByteBuffer(ByteBuffer)}
+ * and {@link #writeRecords(BaseRecords)}).
+ *
+ * See {@link org.apache.kafka.common.requests.EnvelopeRequest#toSend(RequestHeader)}
+ * for example usage.
+ */
+public class SendBuilder implements Writable {
+    private final ByteBuffer buffer;
+
+    private final Queue sends = new ArrayDeque<>(1);
+    private long sizeOfSends = 0;
+
+    private final List buffers = new ArrayList<>();
+    private long sizeOfBuffers = 0;
+
+    SendBuilder(int size) {
+        this.buffer = ByteBuffer.allocate(size);
+        this.buffer.mark();
+    }
+
+    @Override
+    public void writeByte(byte val) {
+        buffer.put(val);
+    }
+
+    @Override
+    public void writeShort(short val) {
+        buffer.putShort(val);
+    }
+
+    @Override
+    public void writeInt(int val) {
+        buffer.putInt(val);
+    }
+
+    @Override
+    public void writeLong(long val) {
+        buffer.putLong(val);
+    }
+
+    @Override
+    public void writeDouble(double val) {
+        buffer.putDouble(val);
+    }
+
+    @Override
+    public void writeByteArray(byte[] arr) {
+        buffer.put(arr);
+    }
+
+    @Override
+    public void writeUnsignedVarint(int i) {
+        ByteUtils.writeUnsignedVarint(i, buffer);
+    }
+
+    /**
+     * Write a byte buffer. The reference to the underlying buffer will
+     * be retained in the result of {@link #build()}.
+     *
+     * @param buf the buffer to write
+     */
+    @Override
+    public void writeByteBuffer(ByteBuffer buf) {
+        flushPendingBuffer();
+        addBuffer(buf.duplicate());
+    }
+
+    @Override
+    public void writeVarint(int i) {
+        ByteUtils.writeVarint(i, buffer);
+    }
+
+    @Override
+    public void writeVarlong(long i) {
+        ByteUtils.writeVarlong(i, buffer);
+    }
+
+    private void addBuffer(ByteBuffer buffer) {
+        buffers.add(buffer);
+        sizeOfBuffers += buffer.remaining();
+    }
+
+    private void addSend(Send send) {
+        sends.add(send);
+        sizeOfSends += send.size();
+    }
+
+    private void clearBuffers() {
+        buffers.clear();
+        sizeOfBuffers = 0;
+    }
+
+    /**
+     * Write a record set. The underlying record data will be retained
+     * in the result of {@link #build()}. See {@link BaseRecords#toSend()}.
+     *
+     * @param records the records to write
+     */
+    @Override
+    public void writeRecords(BaseRecords records) {
+        if (records instanceof MemoryRecords) {
+            flushPendingBuffer();
+            addBuffer(((MemoryRecords) records).buffer());
+        } else {
+            flushPendingSend();
+            addSend(records.toSend());
+        }
+    }
+
+    private void flushPendingSend() {
+        flushPendingBuffer();
+        if (!buffers.isEmpty()) {
+            ByteBuffer[] byteBufferArray = buffers.toArray(new ByteBuffer[0]);
+            addSend(new ByteBufferSend(byteBufferArray, sizeOfBuffers));
+            clearBuffers();
+        }
+    }
+
+    private void flushPendingBuffer() {
+        int latestPosition = buffer.position();
+        buffer.reset();
+
+        if (latestPosition > buffer.position()) {
+            buffer.limit(latestPosition);
+            addBuffer(buffer.slice());
+
+            buffer.position(latestPosition);
+            buffer.limit(buffer.capacity());
+            buffer.mark();
+        }
+    }
+
+    public Send build() {
+        flushPendingSend();
+
+        if (sends.size() == 1) {
+            return sends.poll();
+        } else {
+            return new MultiRecordsSend(sends, sizeOfSends);
+        }
+    }
+
+    public static Send buildRequestSend(
+        RequestHeader header,
+        Message apiRequest
+    ) {
+        return buildSend(
+            header.data(),
+            header.headerVersion(),
+            apiRequest,
+            header.apiVersion()
+        );
+    }
+
+    public static Send buildResponseSend(
+        ResponseHeader header,
+        Message apiResponse,
+        short apiVersion
+    ) {
+        return buildSend(
+            header.data(),
+            header.headerVersion(),
+            apiResponse,
+            apiVersion
+        );
+    }
+
+    private static Send buildSend(
+        Message header,
+        short headerVersion,
+        Message apiMessage,
+        short apiVersion
+    ) {
+        ObjectSerializationCache serializationCache = new ObjectSerializationCache();
+
+        MessageSizeAccumulator messageSize = new MessageSizeAccumulator();
+        header.addSize(messageSize, serializationCache, headerVersion);
+        apiMessage.addSize(messageSize, serializationCache, apiVersion);
+
+        SendBuilder builder = new SendBuilder(messageSize.sizeExcludingZeroCopy() + 4);
+        builder.writeInt(messageSize.totalSize());
+        header.write(builder, serializationCache, headerVersion);
+        apiMessage.write(builder, serializationCache, apiVersion);
+
+        return builder.build();
+    }
+
+}
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
new file mode 100644
index 0000000000000..1ca59d04e0fec
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/protocol/Writable.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.protocol;
+
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.common.record.BaseRecords;
+import org.apache.kafka.common.record.MemoryRecords;
+
+import java.nio.ByteBuffer;
+
+public interface Writable {
+    void writeByte(byte val);
+    void writeShort(short val);
+    void writeInt(int val);
+    void writeLong(long val);
+    void writeDouble(double val);
+    void writeByteArray(byte[] arr);
+    void writeUnsignedVarint(int i);
+    void writeByteBuffer(ByteBuffer buf);
+    void writeVarint(int i);
+    void writeVarlong(long i);
+
+    default void writeRecords(BaseRecords records) {
+        if (records instanceof MemoryRecords) {
+            MemoryRecords memRecords = (MemoryRecords) records;
+            writeByteBuffer(memRecords.buffer());
+        } else {
+            throw new UnsupportedOperationException("Unsupported record type " + records.getClass());
+        }
+    }
+
+    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 4213ecd08f054..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
@@ -16,12 +16,17 @@
  */
 package org.apache.kafka.common.protocol.types;
 
+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
  */
-public class ArrayOf extends Type {
+public class ArrayOf extends DocumentedType {
+
+    private static final String ARRAY_TYPE_NAME = "ARRAY";
 
     private final Type type;
     private final boolean nullable;
@@ -87,13 +92,14 @@ public int sizeOf(Object o) {
         return size;
     }
 
-    public Type type() {
-        return type;
+    @Override
+    public Optional arrayElementType() {
+        return Optional.of(type);
     }
 
     @Override
     public String toString() {
-        return "ARRAY(" + type + ")";
+        return ARRAY_TYPE_NAME + "(" + type + ")";
     }
 
     @Override
@@ -110,4 +116,18 @@ public Object[] validate(Object item) {
             throw new SchemaException("Not an Object[].");
         }
     }
+
+    @Override
+    public String typeName() {
+        return 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 is given as an " + INT32 + ". Then N instances of type T follow. " +
+                "A null array is represented with a length of -1. " +
+                "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/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 ec217f5bd0b46..7d5c3f5d3d4bc 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
@@ -50,6 +50,9 @@ public static class Int8 extends Field {
         public Int8(String name, String docString) {
             super(name, Type.INT8, docString, false, null);
         }
+        public Int8(String name, String docString, byte defaultValue) {
+            super(name, Type.INT8, docString, true, defaultValue);
+        }
     }
 
     public static class Int32 extends Field {
@@ -72,21 +75,116 @@ 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);
         }
     }
 
+    public static class Float64 extends Field {
+        public Float64(String name, String docString) {
+            super(name, Type.FLOAT64, docString, false, null);
+        }
+
+        public Float64(String name, String docString, double defaultValue) {
+            super(name, Type.FLOAT64, docString, true, defaultValue);
+        }
+    }
+
     public static class Str extends Field {
         public Str(String name, String docString) {
             super(name, Type.STRING, docString, false, null);
         }
     }
 
+    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);
+        }
+    }
+
+    public static class Array extends Field {
+        public Array(String name, Type elementType, String docString) {
+            super(name, new ArrayOf(elementType), docString, false, null);
+        }
+    }
+
+    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;
+
+        public ComplexArray(String name, String docString) {
+            this.name = name;
+            this.docString = docString;
+        }
+
+        public Field withFields(Field... fields) {
+            Schema elementType = new Schema(fields);
+            return new Field(name, new ArrayOf(elementType), docString, false, null);
+        }
+
+        public Field nullableWithFields(Field... fields) {
+            Schema elementType = new Schema(fields);
+            return new Field(name, ArrayOf.nullable(elementType), docString, false, null);
+        }
+
+        public Field withFields(String docStringOverride, Field... fields) {
+            Schema elementType = new Schema(fields);
+            return new Field(name, new ArrayOf(elementType), docStringOverride, false, null);
+        }
+    }
 }
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 faa1540b499cd..aa6ffbef7718a 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.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,9 +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/SchemaException.java b/clients/src/main/java/org/apache/kafka/common/protocol/types/SchemaException.java
index 48a3bdc5f777f..8bbab3222d55e 100644
--- a/clients/src/main/java/org/apache/kafka/common/protocol/types/SchemaException.java
+++ b/clients/src/main/java/org/apache/kafka/common/protocol/types/SchemaException.java
@@ -29,4 +29,7 @@ public SchemaException(String message) {
         super(message);
     }
 
+    public SchemaException(String message, Throwable cause) {
+        super(message, cause);
+    }
 }
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 b825201ae3b6b..e6e4a6fdc2ce0 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
@@ -16,10 +16,12 @@
  */
 package org.apache.kafka.common.protocol.types;
 
-import org.apache.kafka.common.record.Records;
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.common.record.BaseRecords;
 
 import java.nio.ByteBuffer;
 import java.util.Arrays;
+import java.util.Objects;
 
 /**
  * A record that can be serialized and deserialized according to a pre-defined schema
@@ -87,10 +89,18 @@ 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);
     }
 
+    public Double get(Field.Float64 field) {
+        return getDouble(field.name);
+    }
+
     public String get(Field.Str field) {
         return getString(field.name);
     }
@@ -99,18 +109,54 @@ public String get(Field.NullableStr field) {
         return getString(field.name);
     }
 
+    public Boolean get(Field.Bool field) {
+        return getBoolean(field.name);
+    }
+
+    public Object[] get(Field.Array field) {
+        return getArray(field.name);
+    }
+
+    public Object[] get(Field.ComplexArray field) {
+        return getArray(field.name);
+    }
+
     public Long getOrElse(Field.Int64 field, long alternative) {
         if (hasField(field.name))
             return getLong(field.name);
         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);
+        return alternative;
+    }
+
+    public Byte getOrElse(Field.Int8 field, byte alternative) {
+        if (hasField(field.name))
+            return getByte(field.name);
+        return alternative;
+    }
+
     public Integer getOrElse(Field.Int32 field, int alternative) {
         if (hasField(field.name))
             return getInt(field.name);
         return alternative;
     }
 
+    public Double getOrElse(Field.Float64 field, double alternative) {
+        if (hasField(field.name))
+            return getDouble(field.name);
+        return alternative;
+    }
+
     public String getOrElse(Field.NullableStr field, String alternative) {
         if (hasField(field.name))
             return getString(field.name);
@@ -123,6 +169,24 @@ public String getOrElse(Field.Str field, String alternative) {
         return alternative;
     }
 
+    public boolean getOrElse(Field.Bool field, boolean alternative) {
+        if (hasField(field.name))
+            return getBoolean(field.name);
+        return alternative;
+    }
+
+    public Object[] getOrEmpty(Field.Array field) {
+        if (hasField(field.name))
+            return getArray(field.name);
+        return new Object[0];
+    }
+
+    public Object[] getOrEmpty(Field.ComplexArray field) {
+        if (hasField(field.name))
+            return getArray(field.name);
+        return new Object[0];
+    }
+
     /**
      * Get the record value for the field with the given name by doing a hash table lookup (slower!)
      *
@@ -150,6 +214,10 @@ public boolean hasField(Field def) {
         return schema.get(def.name) != null;
     }
 
+    public boolean hasField(Field.ComplexArray def) {
+        return schema.get(def.name) != null;
+    }
+
     public Struct getStruct(BoundField field) {
         return (Struct) get(field);
     }
@@ -166,8 +234,8 @@ public byte getByte(String name) {
         return (Byte) get(name);
     }
 
-    public Records getRecords(String name) {
-        return (Records) get(name);
+    public BaseRecords getRecords(String name) {
+        return (BaseRecords) get(name);
     }
 
     public Short getShort(BoundField field) {
@@ -198,6 +266,22 @@ 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 Double getDouble(BoundField field) {
+        return (Double) get(field);
+    }
+
+    public Double getDouble(String name) {
+        return (Double) get(name);
+    }
+
     public Object[] getArray(BoundField field) {
         return (Object[]) get(field);
     }
@@ -236,6 +320,17 @@ public ByteBuffer getBytes(String name) {
         return (ByteBuffer) result;
     }
 
+    public byte[] getByteArray(String name) {
+        Object result = get(name);
+        if (result instanceof byte[])
+            return (byte[]) result;
+        ByteBuffer buf = (ByteBuffer) result;
+        byte[] arr = new byte[buf.remaining()];
+        buf.get(arr);
+        buf.flip();
+        return arr;
+    }
+
     /**
      * Set the given field to the specified value
      *
@@ -284,12 +379,49 @@ 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);
     }
 
+    public Struct set(Field.Float64 def, double value) {
+        return set(def.name, value);
+    }
+
+    public Struct set(Field.Bool def, boolean value) {
+        return set(def.name, value);
+    }
+
+    public Struct set(Field.Array def, Object[] value) {
+        return set(def.name, value);
+    }
+
+    public Struct set(Field.ComplexArray def, Object[] value) {
+        return set(def.name, value);
+    }
+
+    public Struct setByteArray(String name, byte[] value) {
+        ByteBuffer buf = value == null ? null : ByteBuffer.wrap(value);
+        return set(name, buf);
+    }
+
+    public Struct setIfExists(Field.Array def, Object[] value) {
+        return setIfExists(def.name, value);
+    }
+
+    public Struct setIfExists(Field.ComplexArray def, Object[] value) {
+        return setIfExists(def.name, value);
+    }
+
     public Struct setIfExists(Field def, Object value) {
-        BoundField field = this.schema.get(def.name);
+        return setIfExists(def.name, value);
+    }
+
+    public Struct setIfExists(String fieldName, Object value) {
+        BoundField field = this.schema.get(fieldName);
         if (field != null)
             this.values[field.index] = value;
         return this;
@@ -308,9 +440,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);
         }
@@ -327,6 +458,14 @@ public Struct instance(String field) {
         return instance(schema.get(field));
     }
 
+    public Struct instance(Field field) {
+        return instance(schema.get(field.name));
+    }
+
+    public Struct instance(Field.ComplexArray field) {
+        return instance(schema.get(field.name));
+    }
+
     /**
      * Empty all the values from this record
      */
@@ -354,6 +493,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)
@@ -377,7 +517,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++) {
@@ -401,7 +541,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)
@@ -431,17 +571,16 @@ 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;
         }
         return true;
     }
-
 }
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 57d31f459fbf1..43a77fc29c3cc 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
@@ -16,13 +16,14 @@
  */
 package org.apache.kafka.common.protocol.types;
 
-import org.apache.kafka.common.record.FileRecords;
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.common.record.BaseRecords;
 import org.apache.kafka.common.record.MemoryRecords;
-import org.apache.kafka.common.record.Records;
 import org.apache.kafka.common.utils.ByteUtils;
 import org.apache.kafka.common.utils.Utils;
 
 import java.nio.ByteBuffer;
+import java.util.Optional;
 
 /**
  * A serializable type
@@ -63,6 +64,43 @@ 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.
+     */
+    public static abstract class DocumentedType extends Type {
+
+        /**
+         * Short name of the type to identify it in documentation;
+         * @return the name of the type
+         */
+        public abstract String typeName();
+
+        /**
+         * Documentation of the Type.
+         *
+         * @return details about valid values, representation
+         */
+        public abstract String documentation();
+
+        @Override
+        public String toString() {
+            return typeName();
+        }
+    }
     /**
      * The Boolean type represents a boolean value in a byte by using
      * the value of 0 to represent false, and 1 to represent true.
@@ -70,7 +108,7 @@ public boolean isNullable() {
      * If for some reason a value that is not 0 or 1 is read,
      * then any non-zero value will return true.
      */
-    public static final Type BOOLEAN = new Type() {
+    public static final DocumentedType BOOLEAN = new DocumentedType() {
         @Override
         public void write(ByteBuffer buffer, Object o) {
             if ((Boolean) o)
@@ -91,7 +129,7 @@ public int sizeOf(Object o) {
         }
 
         @Override
-        public String toString() {
+        public String typeName() {
             return "BOOLEAN";
         }
 
@@ -102,9 +140,16 @@ public Boolean validate(Object item) {
             else
                 throw new SchemaException(item + " is not a Boolean.");
         }
+
+        @Override
+        public String documentation() {
+            return "Represents a boolean value in a byte. " +
+                    "Values 0 and 1 are used to represent false and true respectively. " +
+                    "When reading a boolean value, any non-zero value is considered true.";
+        }
     };
 
-    public static final Type INT8 = new Type() {
+    public static final DocumentedType INT8 = new DocumentedType() {
         @Override
         public void write(ByteBuffer buffer, Object o) {
             buffer.put((Byte) o);
@@ -121,7 +166,7 @@ public int sizeOf(Object o) {
         }
 
         @Override
-        public String toString() {
+        public String typeName() {
             return "INT8";
         }
 
@@ -132,9 +177,14 @@ public Byte validate(Object item) {
             else
                 throw new SchemaException(item + " is not a Byte.");
         }
+
+        @Override
+        public String documentation() {
+            return "Represents an integer between -27 and 27-1 inclusive.";
+        }
     };
 
-    public static final Type INT16 = new Type() {
+    public static final DocumentedType INT16 = new DocumentedType() {
         @Override
         public void write(ByteBuffer buffer, Object o) {
             buffer.putShort((Short) o);
@@ -151,7 +201,7 @@ public int sizeOf(Object o) {
         }
 
         @Override
-        public String toString() {
+        public String typeName() {
             return "INT16";
         }
 
@@ -162,9 +212,15 @@ public Short validate(Object item) {
             else
                 throw new SchemaException(item + " is not a Short.");
         }
+
+        @Override
+        public String documentation() {
+            return "Represents an integer between -215 and 215-1 inclusive. " +
+                    "The values are encoded using two bytes in network byte order (big-endian).";
+        }
     };
 
-    public static final Type INT32 = new Type() {
+    public static final DocumentedType INT32 = new DocumentedType() {
         @Override
         public void write(ByteBuffer buffer, Object o) {
             buffer.putInt((Integer) o);
@@ -181,7 +237,7 @@ public int sizeOf(Object o) {
         }
 
         @Override
-        public String toString() {
+        public String typeName() {
             return "INT32";
         }
 
@@ -192,9 +248,15 @@ public Integer validate(Object item) {
             else
                 throw new SchemaException(item + " is not an Integer.");
         }
+
+        @Override
+        public String documentation() {
+            return "Represents an integer between -231 and 231-1 inclusive. " +
+                    "The values are encoded using four bytes in network byte order (big-endian).";
+        }
     };
 
-    public static final Type UNSIGNED_INT32 = new Type() {
+    public static final DocumentedType UNSIGNED_INT32 = new DocumentedType() {
         @Override
         public void write(ByteBuffer buffer, Object o) {
             ByteUtils.writeUnsignedInt(buffer, (long) o);
@@ -211,7 +273,7 @@ public int sizeOf(Object o) {
         }
 
         @Override
-        public String toString() {
+        public String typeName() {
             return "UINT32";
         }
 
@@ -222,9 +284,15 @@ public Long validate(Object item) {
             else
                 throw new SchemaException(item + " is not a Long.");
         }
+
+        @Override
+        public String documentation() {
+            return "Represents an integer between 0 and 232-1 inclusive. " +
+                    "The values are encoded using four bytes in network byte order (big-endian).";
+        }
     };
 
-    public static final Type INT64 = new Type() {
+    public static final DocumentedType INT64 = new DocumentedType() {
         @Override
         public void write(ByteBuffer buffer, Object o) {
             buffer.putLong((Long) o);
@@ -241,7 +309,7 @@ public int sizeOf(Object o) {
         }
 
         @Override
-        public String toString() {
+        public String typeName() {
             return "INT64";
         }
 
@@ -252,9 +320,89 @@ public Long validate(Object item) {
             else
                 throw new SchemaException(item + " is not a Long.");
         }
+
+        @Override
+        public String documentation() {
+            return "Represents an integer between -263 and 263-1 inclusive. " +
+                    "The values are encoded using eight bytes in network byte order (big-endian).";
+        }
+    };
+
+    public static final DocumentedType UUID = new DocumentedType() {
+        @Override
+        public void write(ByteBuffer buffer, Object o) {
+            final Uuid uuid = (Uuid) o;
+            buffer.putLong(uuid.getMostSignificantBits());
+            buffer.putLong(uuid.getLeastSignificantBits());
+        }
+
+        @Override
+        public Object read(ByteBuffer buffer) {
+            return new 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 type 4 immutable universally unique identifier (Uuid). " +
+                    "The values are encoded using sixteen bytes in network byte order (big-endian).";
+        }
+    };
+
+    public static final DocumentedType FLOAT64 = new DocumentedType() {
+        @Override
+        public void write(ByteBuffer buffer, Object o) {
+            ByteUtils.writeDouble((Double) o, buffer);
+        }
+
+        @Override
+        public Object read(ByteBuffer buffer) {
+            return ByteUtils.readDouble(buffer);
+        }
+
+        @Override
+        public int sizeOf(Object o) {
+            return 8;
+        }
+
+        @Override
+        public String typeName() {
+            return "FLOAT64";
+        }
+
+        @Override
+        public Double validate(Object item) {
+            if (item instanceof Double)
+                return (Double) item;
+            else
+                throw new SchemaException(item + " is not a Double.");
+        }
+
+        @Override
+        public String documentation() {
+            return "Represents a double-precision 64-bit format IEEE 754 value. " +
+                    "The values are encoded using eight bytes in network byte order (big-endian).";
+        }
     };
 
-    public static final Type STRING = new Type() {
+    public static final DocumentedType STRING = new DocumentedType() {
         @Override
         public void write(ByteBuffer buffer, Object o) {
             byte[] bytes = Utils.utf8((String) o);
@@ -282,7 +430,7 @@ public int sizeOf(Object o) {
         }
 
         @Override
-        public String toString() {
+        public String typeName() {
             return "STRING";
         }
 
@@ -293,9 +441,66 @@ public String validate(Object item) {
             else
                 throw new SchemaException(item + " is not a String.");
         }
+
+        @Override
+        public String documentation() {
+            return "Represents a sequence of characters. First the length N is given as an " + INT16 +
+                    ". Then N bytes follow which are the UTF-8 encoding of the character sequence. " +
+                    "Length must not be negative.";
+        }
+    };
+
+    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 Type NULLABLE_STRING = new Type() {
+    public static final DocumentedType NULLABLE_STRING = new DocumentedType() {
         @Override
         public boolean isNullable() {
             return true;
@@ -336,7 +541,7 @@ public int sizeOf(Object o) {
         }
 
         @Override
-        public String toString() {
+        public String typeName() {
             return "NULLABLE_STRING";
         }
 
@@ -350,9 +555,84 @@ public String validate(Object item) {
             else
                 throw new SchemaException(item + " is not a String.");
         }
+
+        @Override
+        public String documentation() {
+            return "Represents a sequence of characters or null. For non-null strings, first the length N is given as an " + INT16 +
+                    ". Then N bytes follow which are the UTF-8 encoding of the character sequence. " +
+                    "A null value is encoded with length of -1 and there are no following bytes.";
+        }
     };
 
-    public static final Type BYTES = new Type() {
+    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) {
             ByteBuffer arg = (ByteBuffer) o;
@@ -383,7 +663,7 @@ public int sizeOf(Object o) {
         }
 
         @Override
-        public String toString() {
+        public String typeName() {
             return "BYTES";
         }
 
@@ -394,9 +674,66 @@ public ByteBuffer validate(Object 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 is given as an " + INT32 +
+                    ". Then N bytes follow.";
+        }
+    };
+
+    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 Type NULLABLE_BYTES = new Type() {
+    public static final DocumentedType NULLABLE_BYTES = new DocumentedType() {
         @Override
         public boolean isNullable() {
             return true;
@@ -440,7 +777,7 @@ public int sizeOf(Object o) {
         }
 
         @Override
-        public String toString() {
+        public String typeName() {
             return "NULLABLE_BYTES";
         }
 
@@ -454,9 +791,15 @@ public ByteBuffer validate(Object item) {
 
             throw new SchemaException(item + " is not a java.nio.ByteBuffer.");
         }
+
+        @Override
+        public String documentation() {
+            return "Represents a raw sequence of bytes or null. For non-null values, first the length N is given as an " + INT32 +
+                    ". Then N bytes follow. A null value is encoded with length of -1 and there are no following bytes.";
+        }
     };
 
-    public static final Type RECORDS = new Type() {
+    public static final DocumentedType COMPACT_NULLABLE_BYTES = new DocumentedType() {
         @Override
         public boolean isNullable() {
             return true;
@@ -464,16 +807,153 @@ public boolean isNullable() {
 
         @Override
         public void write(ByteBuffer buffer, Object o) {
-            if (o instanceof FileRecords)
-                throw new IllegalArgumentException("FileRecords must be written to the channel directly");
-            MemoryRecords records = (MemoryRecords) o;
-            NULLABLE_BYTES.write(buffer, records.buffer().duplicate());
+            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 Records read(ByteBuffer buffer) {
+        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 COMPACT_RECORDS = new DocumentedType() {
+        @Override
+        public boolean isNullable() {
+            return true;
+        }
+
+        @Override
+        public void write(ByteBuffer buffer, Object o) {
+            if (o == null) {
+                COMPACT_NULLABLE_BYTES.write(buffer, null);
+            } else if (o instanceof MemoryRecords) {
+                MemoryRecords records = (MemoryRecords) o;
+                COMPACT_NULLABLE_BYTES.write(buffer, records.buffer().duplicate());
+            } else {
+                throw new IllegalArgumentException("Unexpected record type: " + o.getClass());
+            }
+        }
+
+        @Override
+        public MemoryRecords read(ByteBuffer buffer) {
+            ByteBuffer recordsBuffer = (ByteBuffer) COMPACT_NULLABLE_BYTES.read(buffer);
+            if (recordsBuffer == null) {
+                return null;
+            } else {
+                return MemoryRecords.readableRecords(recordsBuffer);
+            }
+        }
+
+        @Override
+        public int sizeOf(Object o) {
+            if (o == null) {
+                return 1;
+            }
+
+            BaseRecords records = (BaseRecords) o;
+            int recordsSize = records.sizeInBytes();
+            return ByteUtils.sizeOfUnsignedVarint(recordsSize + 1) + recordsSize;
+        }
+
+        @Override
+        public String typeName() {
+            return "COMPACT_RECORDS";
+        }
+
+        @Override
+        public BaseRecords validate(Object item) {
+            if (item == null)
+                return null;
+
+            if (item instanceof BaseRecords)
+                return (BaseRecords) item;
+
+            throw new SchemaException(item + " is not an instance of " + BaseRecords.class.getName());
+        }
+
+        @Override
+        public String documentation() {
+            return "Represents a sequence of Kafka records as " + COMPACT_NULLABLE_BYTES + ". " +
+                "For a detailed description of records see " +
+                "Message Sets.";
+        }
+    };
+
+    public static final DocumentedType RECORDS = new DocumentedType() {
+        @Override
+        public boolean isNullable() {
+            return true;
+        }
+
+        @Override
+        public void write(ByteBuffer buffer, Object o) {
+            if (o == null) {
+                NULLABLE_BYTES.write(buffer, null);
+            } else if (o instanceof MemoryRecords) {
+                MemoryRecords records = (MemoryRecords) o;
+                NULLABLE_BYTES.write(buffer, records.buffer().duplicate());
+            } else {
+                throw new IllegalArgumentException("Unexpected record type: " + o.getClass());
+            }
+        }
+
+        @Override
+        public MemoryRecords read(ByteBuffer buffer) {
             ByteBuffer recordsBuffer = (ByteBuffer) NULLABLE_BYTES.read(buffer);
-            return MemoryRecords.readableRecords(recordsBuffer);
+            if (recordsBuffer == null) {
+                return null;
+            } else {
+                return MemoryRecords.readableRecords(recordsBuffer);
+            }
         }
 
         @Override
@@ -481,28 +961,35 @@ public int sizeOf(Object o) {
             if (o == null)
                 return 4;
 
-            Records records = (Records) o;
+            BaseRecords records = (BaseRecords) o;
             return 4 + records.sizeInBytes();
         }
 
         @Override
-        public String toString() {
+        public String typeName() {
             return "RECORDS";
         }
 
         @Override
-        public Records validate(Object item) {
+        public BaseRecords validate(Object item) {
             if (item == null)
                 return null;
 
-            if (item instanceof Records)
-                return (Records) item;
+            if (item instanceof BaseRecords)
+                return (BaseRecords) item;
 
-            throw new SchemaException(item + " is not an instance of " + Records.class.getName());
+            throw new SchemaException(item + " is not an instance of " + BaseRecords.class.getName());
+        }
+
+        @Override
+        public String documentation() {
+            return "Represents a sequence of Kafka records as " + NULLABLE_BYTES + ". " +
+                    "For a detailed description of records see " +
+                    "Message Sets.";
         }
     };
 
-    public static final Type VARINT = new Type() {
+    public static final DocumentedType VARINT = new DocumentedType() {
         @Override
         public void write(ByteBuffer buffer, Object o) {
             ByteUtils.writeVarint((Integer) o, buffer);
@@ -520,7 +1007,7 @@ public Integer validate(Object item) {
             throw new SchemaException(item + " is not an integer");
         }
 
-        public String toString() {
+        public String typeName() {
             return "VARINT";
         }
 
@@ -528,9 +1015,16 @@ public String toString() {
         public int sizeOf(Object o) {
             return ByteUtils.sizeOfVarint((Integer) o);
         }
+
+        @Override
+        public String documentation() {
+            return "Represents an integer between -231 and 231-1 inclusive. " +
+                    "Encoding follows the variable-length zig-zag encoding from " +
+                    "  Google Protocol Buffers.";
+        }
     };
 
-    public static final Type VARLONG = new Type() {
+    public static final DocumentedType VARLONG = new DocumentedType() {
         @Override
         public void write(ByteBuffer buffer, Object o) {
             ByteUtils.writeVarlong((Long) o, buffer);
@@ -548,7 +1042,7 @@ public Long validate(Object item) {
             throw new SchemaException(item + " is not a long");
         }
 
-        public String toString() {
+        public String typeName() {
             return "VARLONG";
         }
 
@@ -556,6 +1050,43 @@ public String toString() {
         public int sizeOf(Object o) {
             return ByteUtils.sizeOfVarlong((Long) o);
         }
+
+        @Override
+        public String documentation() {
+            return "Represents an integer between -263 and 263-1 inclusive. " +
+                    "Encoding follows the variable-length zig-zag encoding from " +
+                    "  Google Protocol Buffers.";
+        }
     };
 
+    private static String toHtml() {
+        DocumentedType[] types = {
+            BOOLEAN, INT8, INT16, INT32, INT64,
+            UNSIGNED_INT32, VARINT, VARLONG, UUID, FLOAT64,
+            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("");
+        b.append("\n");
+        b.append("\n");
+        b.append("\n");
+        for (DocumentedType type : types) {
+            b.append("");
+            b.append("");
+            b.append("");
+            b.append("\n");
+        }
+        b.append("
        TypeDescription
        ");
+            b.append(type.typeName());
+            b.append("");
+            b.append(type.documentation());
+            b.append("
        \n");
+        return b.toString();
+    }
+
+    public static void main(String[] args) {
+        System.out.println(toHtml());
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/quota/ClientQuotaAlteration.java b/clients/src/main/java/org/apache/kafka/common/quota/ClientQuotaAlteration.java
new file mode 100644
index 0000000000000..88670ce8ddc05
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/quota/ClientQuotaAlteration.java
@@ -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 org.apache.kafka.common.quota;
+
+import java.util.Collection;
+import java.util.Objects;
+
+/**
+ * Describes a configuration alteration to be made to a client quota entity.
+ */
+public class ClientQuotaAlteration {
+
+    public static class Op {
+        private final String key;
+        private final Double value;
+
+        /**
+         * @param key the quota type to alter
+         * @param value if set then the existing value is updated,
+         *              otherwise if null, the existing value is cleared
+         */
+        public Op(String key, Double value) {
+            this.key = key;
+            this.value = value;
+        }
+
+        /**
+         * @return the quota type to alter
+         */
+        public String key() {
+            return this.key;
+        }
+
+        /**
+         * @return if set then the existing value is updated,
+         *         otherwise if null, the existing value is cleared
+         */
+        public Double value() {
+            return this.value;
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (o == null || getClass() != o.getClass()) return false;
+            Op that = (Op) o;
+            return Objects.equals(key, that.key) && Objects.equals(value, that.value);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(key, value);
+        }
+
+        @Override
+        public String toString() {
+            return "ClientQuotaAlteration.Op(key=" + key + ", value=" + value + ")";
+        }
+    }
+
+    private final ClientQuotaEntity entity;
+    private final Collection ops;
+
+    /**
+     * @param entity the entity whose config will be modified
+     * @param ops the alteration to perform
+     */
+    public ClientQuotaAlteration(ClientQuotaEntity entity, Collection ops) {
+        this.entity = entity;
+        this.ops = ops;
+    }
+
+    /**
+     * @return the entity whose config will be modified
+     */
+    public ClientQuotaEntity entity() {
+        return this.entity;
+    }
+
+    /**
+     * @return the alteration to perform
+     */
+    public Collection ops() {
+        return this.ops;
+    }
+
+    @Override
+    public String toString() {
+        return "ClientQuotaAlteration(entity=" + entity + ", ops=" + ops + ")";
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/quota/ClientQuotaEntity.java b/clients/src/main/java/org/apache/kafka/common/quota/ClientQuotaEntity.java
new file mode 100644
index 0000000000000..dfad50917a926
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/quota/ClientQuotaEntity.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.quota;
+
+import java.util.Map;
+import java.util.Objects;
+
+/**
+ * Describes a client quota entity, which is a mapping of entity types to their names.
+ */
+public class ClientQuotaEntity {
+
+    private final Map entries;
+
+    /**
+     * The type of an entity entry.
+     */
+    public static final String USER = "user";
+    public static final String CLIENT_ID = "client-id";
+    public static final String IP = "ip";
+
+    /**
+     * Constructs a quota entity for the given types and names. If a name is null,
+     * then it is mapped to the built-in default entity name.
+     *
+     * @param entries maps entity type to its name
+     */
+    public ClientQuotaEntity(Map entries) {
+        this.entries = entries;
+    }
+
+    /**
+     * @return map of entity type to its name
+     */
+    public Map entries() {
+        return this.entries;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        ClientQuotaEntity that = (ClientQuotaEntity) o;
+        return Objects.equals(entries, that.entries);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(entries);
+    }
+
+    @Override
+    public String toString() {
+        return "ClientQuotaEntity(entries=" + entries + ")";
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/quota/ClientQuotaFilter.java b/clients/src/main/java/org/apache/kafka/common/quota/ClientQuotaFilter.java
new file mode 100644
index 0000000000000..e8a6a7290d480
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/quota/ClientQuotaFilter.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.common.quota;
+
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Objects;
+
+/**
+ * Describes a client quota entity filter.
+ */
+public class ClientQuotaFilter {
+
+    private final Collection components;
+    private final boolean strict;
+
+    /**
+     * A filter to be applied to matching client quotas.
+     *
+     * @param components the components to filter on
+     * @param strict whether the filter only includes specified components
+     */
+    private ClientQuotaFilter(Collection components, boolean strict) {
+        this.components = components;
+        this.strict = strict;
+    }
+
+    /**
+     * Constructs and returns a quota filter that matches all provided components. Matching entities
+     * with entity types that are not specified by a component will also be included in the result.
+     *
+     * @param components the components for the filter
+     */
+    public static ClientQuotaFilter contains(Collection components) {
+        return new ClientQuotaFilter(components, false);
+    }
+
+    /**
+     * Constructs and returns a quota filter that matches all provided components. Matching entities
+     * with entity types that are not specified by a component will *not* be included in the result.
+     *
+     * @param components the components for the filter
+     */
+    public static ClientQuotaFilter containsOnly(Collection components) {
+        return new ClientQuotaFilter(components, true);
+    }
+
+    /**
+     * Constructs and returns a quota filter that matches all configured entities.
+     */
+    public static ClientQuotaFilter all() {
+        return new ClientQuotaFilter(Collections.emptyList(), false);
+    }
+
+    /**
+     * @return the filter's components
+     */
+    public Collection components() {
+        return this.components;
+    }
+
+    /**
+     * @return whether the filter is strict, i.e. only includes specified components
+     */
+    public boolean strict() {
+        return this.strict;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        ClientQuotaFilter that = (ClientQuotaFilter) o;
+        return Objects.equals(components, that.components) && Objects.equals(strict, that.strict);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(components, strict);
+    }
+
+    @Override
+    public String toString() {
+        return "ClientQuotaFilter(components=" + components + ", strict=" + strict + ")";
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/quota/ClientQuotaFilterComponent.java b/clients/src/main/java/org/apache/kafka/common/quota/ClientQuotaFilterComponent.java
new file mode 100644
index 0000000000000..9fc8f72fa9f23
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/quota/ClientQuotaFilterComponent.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.quota;
+
+import java.util.Objects;
+import java.util.Optional;
+
+/**
+ * Describes a component for applying a client quota filter.
+ */
+public class ClientQuotaFilterComponent {
+
+    private final String entityType;
+    private final Optional match;
+
+    /**
+     * A filter to be applied.
+     *
+     * @param entityType the entity type the filter component applies to
+     * @param match if present, the name that's matched exactly
+     *              if empty, matches the default name
+     *              if null, matches any specified name
+     */
+    private ClientQuotaFilterComponent(String entityType, Optional match) {
+        this.entityType = Objects.requireNonNull(entityType);
+        this.match = match;
+    }
+
+    /**
+     * Constructs and returns a filter component that exactly matches the provided entity
+     * name for the entity type.
+     *
+     * @param entityType the entity type the filter component applies to
+     * @param entityName the entity name that's matched exactly
+     */
+    public static ClientQuotaFilterComponent ofEntity(String entityType, String entityName) {
+        return new ClientQuotaFilterComponent(entityType, Optional.of(Objects.requireNonNull(entityName)));
+    }
+
+    /**
+     * Constructs and returns a filter component that matches the built-in default entity name
+     * for the entity type.
+     *
+     * @param entityType the entity type the filter component applies to
+     */
+    public static ClientQuotaFilterComponent ofDefaultEntity(String entityType) {
+        return new ClientQuotaFilterComponent(entityType, Optional.empty());
+    }
+
+    /**
+     * Constructs and returns a filter component that matches any specified name for the
+     * entity type.
+     *
+     * @param entityType the entity type the filter component applies to
+     */
+    public static ClientQuotaFilterComponent ofEntityType(String entityType) {
+        return new ClientQuotaFilterComponent(entityType, null);
+    }
+
+    /**
+     * @return the component's entity type
+     */
+    public String entityType() {
+        return this.entityType;
+    }
+
+    /**
+     * @return the optional match string, where:
+     *         if present, the name that's matched exactly
+     *         if empty, matches the default name
+     *         if null, matches any specified name
+     */
+    public Optional match() {
+        return this.match;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        ClientQuotaFilterComponent that = (ClientQuotaFilterComponent) o;
+        return Objects.equals(that.entityType, entityType) && Objects.equals(that.match, match);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(entityType, match);
+    }
+
+    @Override
+    public String toString() {
+        return "ClientQuotaFilterComponent(entityType=" + entityType + ", match=" + match + ")";
+    }
+}
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 6ac073dd29801..83637640af49d 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;
@@ -29,10 +30,10 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.nio.ByteBuffer;
-import java.nio.channels.FileChannel;
 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;
@@ -227,7 +228,7 @@ 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);
 
@@ -322,6 +323,8 @@ private DeepRecordsIterator(AbstractLegacyRecordBatch wrapperEntry,
                 throw new InvalidRecordException("Invalid wrapper magic found in legacy deep record iterator " + wrapperMagic);
 
             CompressionType compressionType = wrapperRecord.compressionType();
+            if (compressionType == CompressionType.ZSTD)
+                throw new InvalidRecordException("Invalid wrapper compressionType found in legacy deep record iterator " + wrapperMagic);
             ByteBuffer wrapperValue = wrapperRecord.value();
             if (wrapperValue == null)
                 throw new InvalidRecordException("Found invalid compressed record set with null value (magic = " +
@@ -438,13 +441,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;
         }
     }
@@ -501,6 +504,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());
@@ -515,7 +528,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
@@ -528,10 +541,10 @@ static class LegacyFileChannelRecordBatch extends FileLogInputStream.FileChannel
 
         LegacyFileChannelRecordBatch(long offset,
                                      byte magic,
-                                     FileChannel channel,
+                                     FileRecords fileRecords,
                                      int position,
                                      int batchSize) {
-            super(offset, magic, channel, position, batchSize);
+            super(offset, magic, fileRecords, position, batchSize);
         }
 
         @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 2452798d48551..0f55a508fca9f 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
@@ -18,22 +18,14 @@
 
 import org.apache.kafka.common.header.Header;
 import org.apache.kafka.common.utils.AbstractIterator;
-import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.common.utils.Utils;
 
 import java.nio.ByteBuffer;
-import java.util.ArrayList;
 import java.util.Iterator;
-import java.util.List;
 
 public abstract class AbstractRecords implements Records {
 
-    private final Iterable records = new Iterable() {
-        @Override
-        public Iterator iterator() {
-            return recordsIterator();
-        }
-    };
+    private final Iterable records = this::recordsIterator;
 
     @Override
     public boolean hasMatchingMagic(byte magic) {
@@ -51,90 +43,13 @@ public boolean hasCompatibleMagic(byte magic) {
         return true;
     }
 
-    /**
-     * Down convert batches to the provided message format version. The first offset parameter is only relevant in the
-     * conversion from uncompressed v2 or higher to v1 or lower. The reason is that uncompressed records in v0 and v1
-     * are not batched (put another way, each batch always has 1 record).
-     *
-     * If a client requests records in v1 format starting from the middle of an uncompressed batch in v2 format, we
-     * need to drop records from the batch during the conversion. Some versions of librdkafka rely on this for
-     * correctness.
-     *
-     * The temporaryMemoryBytes computation assumes that the batches are not loaded into the heap
-     * (via classes like FileChannelRecordBatch) before this method is called. This is the case in the broker (we
-     * only load records into the heap when down converting), but it's not for the producer. However, down converting
-     * in the producer is very uncommon and the extra complexity to handle that case is not worth it.
-     */
-    protected ConvertedRecords downConvert(Iterable batches, byte toMagic,
-            long firstOffset, Time time) {
-        // maintain the batch along with the decompressed records to avoid the need to decompress again
-        List recordBatchAndRecordsList = new ArrayList<>();
-        int totalSizeEstimate = 0;
-        long startNanos = time.nanoseconds();
-
-        for (RecordBatch batch : batches) {
-            if (toMagic < RecordBatch.MAGIC_VALUE_V2 && batch.isControlBatch())
-                continue;
-
-            if (batch.magic() <= toMagic) {
-                totalSizeEstimate += batch.sizeInBytes();
-                recordBatchAndRecordsList.add(new RecordBatchAndRecords(batch, null, null));
-            } else {
-                List records = new ArrayList<>();
-                for (Record record : batch) {
-                    // See the method javadoc for an explanation
-                    if (toMagic > RecordBatch.MAGIC_VALUE_V1 || batch.isCompressed() || record.offset() >= firstOffset)
-                        records.add(record);
-                }
-                if (records.isEmpty())
-                    continue;
-                final long baseOffset;
-                if (batch.magic() >= RecordBatch.MAGIC_VALUE_V2 && toMagic >= RecordBatch.MAGIC_VALUE_V2)
-                    baseOffset = batch.baseOffset();
-                else
-                    baseOffset = records.get(0).offset();
-                totalSizeEstimate += estimateSizeInBytes(toMagic, baseOffset, batch.compressionType(), records);
-                recordBatchAndRecordsList.add(new RecordBatchAndRecords(batch, records, baseOffset));
-            }
-        }
-
-        ByteBuffer buffer = ByteBuffer.allocate(totalSizeEstimate);
-        long temporaryMemoryBytes = 0;
-        int numRecordsConverted = 0;
-        for (RecordBatchAndRecords recordBatchAndRecords : recordBatchAndRecordsList) {
-            temporaryMemoryBytes += recordBatchAndRecords.batch.sizeInBytes();
-            if (recordBatchAndRecords.batch.magic() <= toMagic) {
-                recordBatchAndRecords.batch.writeTo(buffer);
-            } else {
-                MemoryRecordsBuilder builder = convertRecordBatch(toMagic, buffer, recordBatchAndRecords);
-                buffer = builder.buffer();
-                temporaryMemoryBytes += builder.uncompressedBytesWritten();
-                numRecordsConverted += builder.numRecords();
-            }
-        }
+    public RecordBatch firstBatch() {
+        Iterator iterator = batches().iterator();
 
-        buffer.flip();
-        RecordsProcessingStats stats = new RecordsProcessingStats(temporaryMemoryBytes, numRecordsConverted,
-                time.nanoseconds() - startNanos);
-        return new ConvertedRecords<>(MemoryRecords.readableRecords(buffer), stats);
-    }
+        if (!iterator.hasNext())
+            return null;
 
-    /**
-     * Return a buffer containing the converted record batches. The returned buffer may not be the same as the received
-     * one (e.g. it may require expansion).
-     */
-    private MemoryRecordsBuilder convertRecordBatch(byte magic, ByteBuffer buffer, RecordBatchAndRecords recordBatchAndRecords) {
-        RecordBatch batch = recordBatchAndRecords.batch;
-        final TimestampType timestampType = batch.timestampType();
-        long logAppendTime = timestampType == TimestampType.LOG_APPEND_TIME ? batch.maxTimestamp() : RecordBatch.NO_TIMESTAMP;
-
-        MemoryRecordsBuilder builder = MemoryRecords.builder(buffer, magic, batch.compressionType(),
-                timestampType, recordBatchAndRecords.baseOffset, logAppendTime);
-        for (Record record : recordBatchAndRecords.records)
-            builder.append(record);
-
-        builder.close();
-        return builder;
+        return iterator.next();
     }
 
     /**
@@ -146,6 +61,11 @@ public Iterable records() {
         return records;
     }
 
+    @Override
+    public DefaultRecordsSend toSend() {
+        return new DefaultRecordsSend(this);
+    }
+
     private Iterator recordsIterator() {
         return new AbstractIterator() {
             private final Iterator batches = batches().iterator();
@@ -236,16 +156,5 @@ public static int recordBatchHeaderSizeInBytes(byte magic, CompressionType compr
         }
     }
 
-    private static class RecordBatchAndRecords {
-        private final RecordBatch batch;
-        private final List records;
-        private final Long baseOffset;
-
-        private RecordBatchAndRecords(RecordBatch batch, List records, Long baseOffset) {
-            this.batch = batch;
-            this.records = records;
-            this.baseOffset = baseOffset;
-        }
-    }
 
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/record/BaseRecords.java b/clients/src/main/java/org/apache/kafka/common/record/BaseRecords.java
new file mode 100644
index 0000000000000..49e316f301c5d
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/record/BaseRecords.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.record;
+
+/**
+ * Base interface for accessing records which could be contained in the log, or an in-memory materialization of log records.
+ */
+public interface BaseRecords {
+    /**
+     * The size of these records in bytes.
+     * @return The size in bytes of the records
+     */
+    int sizeInBytes();
+
+    /**
+     * Encapsulate this {@link BaseRecords} object into {@link RecordsSend}
+     * @return Initialized {@link RecordsSend} object
+     */
+    RecordsSend toSend();
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/record/BufferSupplier.java b/clients/src/main/java/org/apache/kafka/common/record/BufferSupplier.java
index 2e09f7d1a2cc3..1a6c92c712fba 100644
--- a/clients/src/main/java/org/apache/kafka/common/record/BufferSupplier.java
+++ b/clients/src/main/java/org/apache/kafka/common/record/BufferSupplier.java
@@ -93,4 +93,36 @@ public void close() {
             bufferMap.clear();
         }
     }
+
+    /**
+     * Simple buffer supplier for single-threaded usage. It caches a single buffer, which grows
+     * monotonically as needed to fulfill the allocation request.
+     */
+    public static class GrowableBufferSupplier extends BufferSupplier {
+        private ByteBuffer cachedBuffer;
+
+        @Override
+        public ByteBuffer get(int minCapacity) {
+            if (cachedBuffer != null && cachedBuffer.capacity() >= minCapacity) {
+                ByteBuffer res = cachedBuffer;
+                cachedBuffer = null;
+                return res;
+            } else {
+                cachedBuffer = null;
+                return ByteBuffer.allocate(minCapacity);
+            }
+        }
+
+        @Override
+        public void release(ByteBuffer buffer) {
+            buffer.clear();
+            cachedBuffer = buffer;
+        }
+
+        @Override
+        public void close() {
+            cachedBuffer = null;
+        }
+    }
+
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/record/ByteBufferLogInputStream.java b/clients/src/main/java/org/apache/kafka/common/record/ByteBufferLogInputStream.java
index bd54e40a457ac..572abd8e812c7 100644
--- a/clients/src/main/java/org/apache/kafka/common/record/ByteBufferLogInputStream.java
+++ b/clients/src/main/java/org/apache/kafka/common/record/ByteBufferLogInputStream.java
@@ -18,9 +18,9 @@
 
 import org.apache.kafka.common.errors.CorruptRecordException;
 
-import java.io.IOException;
 import java.nio.ByteBuffer;
 
+import static org.apache.kafka.common.record.Records.HEADER_SIZE_UP_TO_MAGIC;
 import static org.apache.kafka.common.record.Records.LOG_OVERHEAD;
 import static org.apache.kafka.common.record.Records.MAGIC_OFFSET;
 import static org.apache.kafka.common.record.Records.SIZE_OFFSET;
@@ -38,20 +38,11 @@ class ByteBufferLogInputStream implements LogInputStream {
         this.maxMessageSize = maxMessageSize;
     }
 
-    public MutableRecordBatch nextBatch() throws IOException {
+    public MutableRecordBatch nextBatch() {
         int remaining = buffer.remaining();
-        if (remaining < LOG_OVERHEAD)
-            return null;
-
-        int recordSize = buffer.getInt(buffer.position() + SIZE_OFFSET);
-        // V0 has the smallest overhead, stricter checking is done later
-        if (recordSize < LegacyRecord.RECORD_OVERHEAD_V0)
-            throw new CorruptRecordException(String.format("Record size is less than the minimum record overhead (%d)", LegacyRecord.RECORD_OVERHEAD_V0));
-        if (recordSize > maxMessageSize)
-            throw new CorruptRecordException(String.format("Record size exceeds the largest allowable message size (%d).", maxMessageSize));
 
-        int batchSize = recordSize + LOG_OVERHEAD;
-        if (remaining < batchSize)
+        Integer batchSize = nextBatchSize();
+        if (batchSize == null || remaining < batchSize)
             return null;
 
         byte magic = buffer.get(buffer.position() + MAGIC_OFFSET);
@@ -60,13 +51,38 @@ public MutableRecordBatch nextBatch() throws IOException {
         batchSlice.limit(batchSize);
         buffer.position(buffer.position() + batchSize);
 
-        if (magic < 0 || magic > RecordBatch.CURRENT_MAGIC_VALUE)
-            throw new CorruptRecordException("Invalid magic found in record: " + magic);
-
         if (magic > RecordBatch.MAGIC_VALUE_V1)
             return new DefaultRecordBatch(batchSlice);
         else
             return new AbstractLegacyRecordBatch.ByteBufferLegacyRecordBatch(batchSlice);
     }
 
+    /**
+     * Validates the header of the next batch and returns batch size.
+     * @return next batch size including LOG_OVERHEAD if buffer contains header up to
+     *         magic byte, null otherwise
+     * @throws CorruptRecordException if record size or magic is invalid
+     */
+    Integer nextBatchSize() throws CorruptRecordException {
+        int remaining = buffer.remaining();
+        if (remaining < LOG_OVERHEAD)
+            return null;
+        int recordSize = buffer.getInt(buffer.position() + SIZE_OFFSET);
+        // V0 has the smallest overhead, stricter checking is done later
+        if (recordSize < LegacyRecord.RECORD_OVERHEAD_V0)
+            throw new CorruptRecordException(String.format("Record size %d is less than the minimum record overhead (%d)",
+                    recordSize, LegacyRecord.RECORD_OVERHEAD_V0));
+        if (recordSize > maxMessageSize)
+            throw new CorruptRecordException(String.format("Record size %d exceeds the largest allowable message size (%d).",
+                    recordSize, maxMessageSize));
+
+        if (remaining < HEADER_SIZE_UP_TO_MAGIC)
+            return null;
+
+        byte magic = buffer.get(buffer.position() + MAGIC_OFFSET);
+        if (magic < 0 || magic > RecordBatch.CURRENT_MAGIC_VALUE)
+            throw new CorruptRecordException("Invalid magic found in record: " + magic);
+
+        return recordSize + LOG_OVERHEAD;
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/record/CompressionRatioEstimator.java b/clients/src/main/java/org/apache/kafka/common/record/CompressionRatioEstimator.java
index 7f117849cd79c..264525b380920 100644
--- a/clients/src/main/java/org/apache/kafka/common/record/CompressionRatioEstimator.java
+++ b/clients/src/main/java/org/apache/kafka/common/record/CompressionRatioEstimator.java
@@ -46,7 +46,7 @@ public static float updateEstimation(String topic, CompressionType type, float o
             if (observedRatio > currentEstimation)
                 compressionRatioForTopic[type.id] = Math.max(currentEstimation + COMPRESSION_RATIO_DETERIORATE_STEP, observedRatio);
             else if (observedRatio < currentEstimation) {
-                compressionRatioForTopic[type.id] = currentEstimation - COMPRESSION_RATIO_IMPROVING_STEP;
+                compressionRatioForTopic[type.id] = Math.max(currentEstimation - COMPRESSION_RATIO_IMPROVING_STEP, observedRatio);
             }
         }
         return compressionRatioForTopic[type.id];
@@ -108,4 +108,4 @@ private static float[] initialCompressionRatio() {
         }
         return compressionRatio;
     }
-}
\ No newline at end of file
+}
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 16d6e01f097c8..cf6a026d0a526 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
@@ -20,6 +20,8 @@
 import org.apache.kafka.common.utils.ByteBufferInputStream;
 import org.apache.kafka.common.utils.ByteBufferOutputStream;
 
+import java.io.BufferedInputStream;
+import java.io.BufferedOutputStream;
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.lang.invoke.MethodHandle;
@@ -49,8 +51,10 @@ public InputStream wrapForInput(ByteBuffer buffer, byte messageVersion, BufferSu
         @Override
         public OutputStream wrapForOutput(ByteBufferOutputStream buffer, byte messageVersion) {
             try {
-                // GZIPOutputStream has a default buffer size of 512 bytes, which is too small
-                return new GZIPOutputStream(buffer, 8 * 1024);
+                // Set input buffer (uncompressed) to 16 KB (none by default) and output buffer (compressed) to
+                // 8 KB (0.5 KB by default) to ensure reasonable performance in cases where the caller passes a small
+                // number of bytes to write (potentially a single byte)
+                return new BufferedOutputStream(new GZIPOutputStream(buffer, 8 * 1024), 16 * 1024);
             } catch (Exception e) {
                 throw new KafkaException(e);
             }
@@ -59,7 +63,11 @@ public OutputStream wrapForOutput(ByteBufferOutputStream buffer, byte messageVer
         @Override
         public InputStream wrapForInput(ByteBuffer buffer, byte messageVersion, BufferSupplier decompressionBufferSupplier) {
             try {
-                return new GZIPInputStream(new ByteBufferInputStream(buffer));
+                // Set output buffer (uncompressed) to 16 KB (none by default) and input buffer (compressed) to
+                // 8 KB (0.5 KB by default) to ensure reasonable performance in cases where the caller reads a small
+                // number of bytes (potentially a single byte)
+                return new BufferedInputStream(new GZIPInputStream(new ByteBufferInputStream(buffer), 8 * 1024),
+                        16 * 1024);
             } catch (Exception e) {
                 throw new KafkaException(e);
             }
@@ -105,6 +113,30 @@ public InputStream wrapForInput(ByteBuffer inputBuffer, byte messageVersion, Buf
                 throw new KafkaException(e);
             }
         }
+    },
+
+    ZSTD(4, "zstd", 1.0f) {
+        @Override
+        public OutputStream wrapForOutput(ByteBufferOutputStream buffer, byte messageVersion) {
+            try {
+                // Set input buffer (uncompressed) to 16 KB (none by default) to ensure reasonable performance
+                // in cases where the caller passes a small number of bytes to write (potentially a single byte)
+                return new BufferedOutputStream((OutputStream) ZstdConstructors.OUTPUT.invoke(buffer), 16 * 1024);
+            } catch (Throwable e) {
+                throw new KafkaException(e);
+            }
+        }
+
+        @Override
+        public InputStream wrapForInput(ByteBuffer buffer, byte messageVersion, BufferSupplier decompressionBufferSupplier) {
+            try {
+                // Set output buffer (uncompressed) to 16 KB (none by default) to ensure reasonable performance
+                // in cases where the caller reads a small number of bytes (potentially a single byte)
+                return new BufferedInputStream((InputStream) ZstdConstructors.INPUT.invoke(new ByteBufferInputStream(buffer)), 16 * 1024);
+            } catch (Throwable e) {
+                throw new KafkaException(e);
+            }
+        }
     };
 
     public final int id;
@@ -120,7 +152,7 @@ public InputStream wrapForInput(ByteBuffer inputBuffer, byte messageVersion, Buf
     /**
      * Wrap bufferStream with an OutputStream that will compress data with this CompressionType.
      *
-     * Note: Unlike {@link #wrapForInput}, {@link #wrapForOutput} cannot take {@#link ByteBuffer}s directly.
+     * Note: Unlike {@link #wrapForInput}, {@link #wrapForOutput} cannot take {@link ByteBuffer}s directly.
      * Currently, {@link MemoryRecordsBuilder#writeDefaultBatchHeader()} and {@link MemoryRecordsBuilder#writeLegacyCompressedWrapperHeader()}
      * write to the underlying buffer in the given {@link ByteBufferOutputStream} after the compressed data has been written.
      * In the event that the buffer needs to be expanded while writing the data, access to the underlying buffer needs to be preserved.
@@ -148,6 +180,8 @@ public static CompressionType forId(int id) {
                 return SNAPPY;
             case 3:
                 return LZ4;
+            case 4:
+                return ZSTD;
             default:
                 throw new IllegalArgumentException("Unknown compression type id: " + id);
         }
@@ -162,6 +196,8 @@ else if (SNAPPY.name.equals(name))
             return SNAPPY;
         else if (LZ4.name.equals(name))
             return LZ4;
+        else if (ZSTD.name.equals(name))
+            return ZSTD;
         else
             throw new IllegalArgumentException("Unknown compression name: " + name);
     }
@@ -169,7 +205,7 @@ else if (LZ4.name.equals(name))
     // We should only have a runtime dependency on compression algorithms in case the native libraries don't support
     // some platforms.
     //
-    // For Snappy, we dynamically load the classes and rely on the initialization-on-demand holder idiom to ensure
+    // For Snappy and Zstd, we dynamically load the classes and rely on the initialization-on-demand holder idiom to ensure
     // they're only loaded if used.
     //
     // For LZ4 we are using org.apache.kafka classes, which should always be in the classpath, and would not trigger
@@ -182,6 +218,13 @@ private static class SnappyConstructors {
                 MethodType.methodType(void.class, OutputStream.class));
     }
 
+    private static class ZstdConstructors {
+        static final MethodHandle INPUT = findConstructor("com.github.luben.zstd.ZstdInputStream",
+            MethodType.methodType(void.class, InputStream.class));
+        static final MethodHandle OUTPUT = findConstructor("com.github.luben.zstd.ZstdOutputStream",
+            MethodType.methodType(void.class, OutputStream.class));
+    }
+
     private static MethodHandle findConstructor(String className, MethodType methodType) {
         try {
             return MethodHandles.publicLookup().findConstructor(Class.forName(className), methodType);
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..091b83d9defff 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;
@@ -42,6 +43,9 @@
 public enum ControlRecordType {
     ABORT((short) 0),
     COMMIT((short) 1),
+    // Raft quorum related control messages.
+    QUORUM_REASSIGNMENT((short) 2),
+    LEADER_CHANGE((short) 3),
 
     // UNKNOWN is used to indicate a control type which the client is not aware of and should be ignored
     UNKNOWN((short) -1);
@@ -92,6 +96,10 @@ public static ControlRecordType fromTypeId(short typeId) {
                 return ABORT;
             case 1:
                 return COMMIT;
+            case 2:
+                return QUORUM_REASSIGNMENT;
+            case 3:
+                return LEADER_CHANGE;
             default:
                 return UNKNOWN;
         }
diff --git a/clients/src/main/java/org/apache/kafka/common/record/ControlRecordUtils.java b/clients/src/main/java/org/apache/kafka/common/record/ControlRecordUtils.java
new file mode 100644
index 0000000000000..fb086dc5f1b45
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/record/ControlRecordUtils.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.record;
+
+import org.apache.kafka.common.message.LeaderChangeMessage;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+
+import java.nio.ByteBuffer;
+
+/**
+ * Utility class for easy interaction with control records.
+ */
+public class ControlRecordUtils {
+
+    public static final short LEADER_CHANGE_SCHEMA_VERSION = new LeaderChangeMessage().highestSupportedVersion();
+
+    public static LeaderChangeMessage deserializeLeaderChangeMessage(Record record) {
+        ControlRecordType recordType = ControlRecordType.parse(record.key());
+        if (recordType != ControlRecordType.LEADER_CHANGE) {
+            throw new IllegalArgumentException(
+                "Expected LEADER_CHANGE control record type(3), but found " + recordType.toString());
+        }
+        return deserializeLeaderChangeMessage(record.value().duplicate());
+    }
+
+    public static LeaderChangeMessage deserializeLeaderChangeMessage(ByteBuffer data) {
+        ByteBufferAccessor byteBufferAccessor = new ByteBufferAccessor(data.duplicate());
+        return new LeaderChangeMessage(byteBufferAccessor, LEADER_CHANGE_SCHEMA_VERSION);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/record/ConvertedRecords.java b/clients/src/main/java/org/apache/kafka/common/record/ConvertedRecords.java
index fe37c48976fbf..d9150e5044db6 100644
--- a/clients/src/main/java/org/apache/kafka/common/record/ConvertedRecords.java
+++ b/clients/src/main/java/org/apache/kafka/common/record/ConvertedRecords.java
@@ -19,18 +19,18 @@
 public class ConvertedRecords {
 
     private final T records;
-    private final RecordsProcessingStats recordsProcessingStats;
+    private final RecordConversionStats recordConversionStats;
 
-    public ConvertedRecords(T records, RecordsProcessingStats recordsProcessingStats) {
+    public ConvertedRecords(T records, RecordConversionStats recordConversionStats) {
         this.records = records;
-        this.recordsProcessingStats = recordsProcessingStats;
+        this.recordConversionStats = recordConversionStats;
     }
 
     public T records() {
         return records;
     }
 
-    public RecordsProcessingStats recordsProcessingStats() {
-        return recordsProcessingStats;
+    public RecordConversionStats recordConversionStats() {
+        return recordConversionStats;
     }
 }
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..d85f1000bc1c9 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);
@@ -369,6 +373,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++) {
@@ -376,7 +535,8 @@ private static Header[] readHeaders(ByteBuffer buffer, int numHeaders) {
             if (headerKeySize < 0)
                 throw new InvalidRecordException("Invalid negative header key size " + headerKeySize);
 
-            String headerKey = Utils.utf8(buffer, headerKeySize);
+            ByteBuffer headerKeyBuffer = buffer.slice();
+            headerKeyBuffer.limit(headerKeySize);
             buffer.position(buffer.position() + headerKeySize);
 
             ByteBuffer headerValue = null;
@@ -387,7 +547,7 @@ private static Header[] readHeaders(ByteBuffer buffer, int numHeaders) {
                 buffer.position(buffer.position() + headerValueSize);
             }
 
-            headers[i] = new RecordHeader(headerKey, headerValue);
+            headers[i] = new RecordHeader(headerKeyBuffer, headerValue);
         }
 
         return headers;
@@ -416,17 +576,16 @@ private static int sizeOfBodyInBytes(int offsetDelta,
                                          ByteBuffer key,
                                          ByteBuffer value,
                                          Header[] headers) {
-
         int keySize = key == null ? -1 : key.remaining();
         int valueSize = value == null ? -1 : value.remaining();
         return sizeOfBodyInBytes(offsetDelta, timestampDelta, keySize, valueSize, headers);
     }
 
-    private static int sizeOfBodyInBytes(int offsetDelta,
-                                         long timestampDelta,
-                                         int keySize,
-                                         int valueSize,
-                                         Header[] headers) {
+    public static int sizeOfBodyInBytes(int offsetDelta,
+                                        long timestampDelta,
+                                        int keySize,
+                                        int valueSize,
+                                        Header[] headers) {
         int size = 1; // always one byte for attributes
         size += ByteUtils.sizeOfVarint(offsetDelta);
         size += ByteUtils.sizeOfVarlong(timestampDelta);
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 ff8d3b990ba0f..33709c0387963 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;
@@ -28,12 +30,12 @@
 import java.io.IOException;
 import java.nio.BufferUnderflowException;
 import java.nio.ByteBuffer;
-import java.nio.channels.FileChannel;
 import java.util.ArrayList;
 import java.util.Collections;
 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;
 
@@ -106,7 +108,7 @@ public class DefaultRecordBatch extends AbstractRecordBatch implements MutableRe
     static final int CRC_LENGTH = 4;
     static final int ATTRIBUTES_OFFSET = CRC_OFFSET + CRC_LENGTH;
     static final int ATTRIBUTE_LENGTH = 2;
-    static final int LAST_OFFSET_DELTA_OFFSET = ATTRIBUTES_OFFSET + ATTRIBUTE_LENGTH;
+    public static final int LAST_OFFSET_DELTA_OFFSET = ATTRIBUTES_OFFSET + ATTRIBUTE_LENGTH;
     static final int LAST_OFFSET_DELTA_LENGTH = 4;
     static final int FIRST_TIMESTAMP_OFFSET = LAST_OFFSET_DELTA_OFFSET + LAST_OFFSET_DELTA_LENGTH;
     static final int FIRST_TIMESTAMP_LENGTH = 8;
@@ -118,7 +120,7 @@ public class DefaultRecordBatch extends AbstractRecordBatch implements MutableRe
     static final int PRODUCER_EPOCH_LENGTH = 2;
     static final int BASE_SEQUENCE_OFFSET = PRODUCER_EPOCH_OFFSET + PRODUCER_EPOCH_LENGTH;
     static final int BASE_SEQUENCE_LENGTH = 4;
-    static final int RECORDS_COUNT_OFFSET = BASE_SEQUENCE_OFFSET + BASE_SEQUENCE_LENGTH;
+    public static final int RECORDS_COUNT_OFFSET = BASE_SEQUENCE_OFFSET + BASE_SEQUENCE_LENGTH;
     static final int RECORDS_COUNT_LENGTH = 4;
     static final int RECORDS_OFFSET = RECORDS_COUNT_OFFSET + RECORDS_COUNT_LENGTH;
     public static final int RECORD_BATCH_OVERHEAD = RECORDS_OFFSET;
@@ -128,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) {
@@ -142,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() + ")");
     }
 
@@ -251,42 +255,33 @@ public int partitionLeaderEpoch() {
         return buffer.getInt(PARTITION_LEADER_EPOCH_OFFSET);
     }
 
-    private CloseableIterator compressedIterator(BufferSupplier bufferSupplier) {
+    public DataInputStream recordInputStream(BufferSupplier bufferSupplier) {
         final ByteBuffer buffer = this.buffer.duplicate();
         buffer.position(RECORDS_OFFSET);
-        final DataInputStream inputStream = new DataInputStream(compressionType().wrapForInput(buffer, magic(),
-                bufferSupplier));
+        return new DataInputStream(compressionType().wrapForInput(buffer, magic(), 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);
-                }
-            }
+    private CloseableIterator compressedIterator(BufferSupplier bufferSupplier, boolean skipKeyValue) {
+        final DataInputStream inputStream = recordInputStream(bufferSupplier);
 
-            @Override
-            protected boolean ensureNoneRemaining() {
-                try {
-                    return inputStream.read() == -1;
-                } catch (IOException e) {
-                    throw new KafkaException("Error checking for remaining bytes after reading batch", 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
-            public void close() {
-                try {
-                    inputStream.close();
-                } catch (IOException e) {
-                    throw new KafkaException("Failed to close record stream", 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);
                 }
-            }
-        };
+            };
+        } 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() {
@@ -321,7 +316,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());
@@ -329,10 +324,29 @@ 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 streamingIterator(BufferSupplier bufferSupplier) {
         if (isCompressed())
-            return compressedIterator(bufferSupplier);
+            return compressedIterator(bufferSupplier, false);
         else
             return uncompressedIterator();
     }
@@ -387,7 +401,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
@@ -429,22 +443,22 @@ public static void writeEmptyHeader(ByteBuffer buffer,
                 producerEpoch, baseSequence, isTransactional, isControlRecord, partitionLeaderEpoch, 0);
     }
 
-    static void writeHeader(ByteBuffer buffer,
-                            long baseOffset,
-                            int lastOffsetDelta,
-                            int sizeInBytes,
-                            byte magic,
-                            CompressionType compressionType,
-                            TimestampType timestampType,
-                            long firstTimestamp,
-                            long maxTimestamp,
-                            long producerId,
-                            short epoch,
-                            int sequence,
-                            boolean isTransactional,
-                            boolean isControlBatch,
-                            int partitionLeaderEpoch,
-                            int numRecords) {
+    public static void writeHeader(ByteBuffer buffer,
+                                   long baseOffset,
+                                   int lastOffsetDelta,
+                                   int sizeInBytes,
+                                   byte magic,
+                                   CompressionType compressionType,
+                                   TimestampType timestampType,
+                                   long firstTimestamp,
+                                   long maxTimestamp,
+                                   long producerId,
+                                   short epoch,
+                                   int sequence,
+                                   boolean isTransactional,
+                                   boolean isControlBatch,
+                                   int partitionLeaderEpoch,
+                                   int numRecords) {
         if (magic < RecordBatch.CURRENT_MAGIC_VALUE)
             throw new IllegalArgumentException("Invalid magic value " + magic);
         if (firstTimestamp < 0 && firstTimestamp != NO_TIMESTAMP)
@@ -473,6 +487,8 @@ static void writeHeader(ByteBuffer buffer,
     @Override
     public String toString() {
         return "RecordBatch(magic=" + magic() + ", offsets=[" + baseOffset() + ", " + lastOffset() + "], " +
+                "sequence=[" + baseSequence() + ", " + lastSequence() + "], " +
+                "isTransactional=" + isTransactional() + ", isControlBatch=" + isControlBatch() + ", " +
                 "compression=" + compressionType() + ", timestampType=" + timestampType() + ", crc=" + checksum() + ")";
     }
 
@@ -523,10 +539,16 @@ static int estimateBatchSizeUpperBound(ByteBuffer key, ByteBuffer value, Header[
         return RECORD_BATCH_OVERHEAD + DefaultRecord.recordSizeUpperBound(key, value, headers);
     }
 
-    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 sequence, int decrement) {
+        if (sequence < decrement)
+            return Integer.MAX_VALUE - (decrement - sequence) + 1;
+        return sequence - decrement;
     }
 
     private abstract class RecordIterator implements CloseableIterator {
@@ -537,7 +559,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();
@@ -582,14 +604,54 @@ 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,
                                       byte magic,
-                                      FileChannel channel,
+                                      FileRecords fileRecords,
                                       int position,
                                       int batchSize) {
-            super(offset, magic, channel, position, batchSize);
+            super(offset, magic, fileRecords, position, batchSize);
         }
 
         @Override
diff --git a/clients/src/main/java/org/apache/kafka/common/record/DefaultRecordsSend.java b/clients/src/main/java/org/apache/kafka/common/record/DefaultRecordsSend.java
new file mode 100644
index 0000000000000..82be5d1623338
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/record/DefaultRecordsSend.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.record;
+
+import java.io.IOException;
+import java.nio.channels.GatheringByteChannel;
+
+public class DefaultRecordsSend extends RecordsSend {
+    public DefaultRecordsSend(Records records) {
+        this(records, records.sizeInBytes());
+    }
+
+    public DefaultRecordsSend(Records records, int maxBytesToWrite) {
+        super(records, maxBytesToWrite);
+    }
+
+    @Override
+    protected long writeTo(GatheringByteChannel channel, long previouslyWritten, int remaining) throws IOException {
+        return records().writeTo(channel, previouslyWritten, remaining);
+    }
+}
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 75eb1b3e35438..15c09dea32c8d 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;
@@ -40,26 +41,27 @@
 public class FileLogInputStream implements LogInputStream {
     private int position;
     private final int end;
-    private final FileChannel channel;
+    private final FileRecords fileRecords;
     private final ByteBuffer logHeaderBuffer = ByteBuffer.allocate(HEADER_SIZE_UP_TO_MAGIC);
 
     /**
      * Create a new log input stream over the FileChannel
-     * @param channel Underlying FileChannel
+     * @param records Underlying FileRecords instance
      * @param start Position in the file channel to start from
      * @param end Position in the file channel not to read past
      */
-    FileLogInputStream(FileChannel channel,
+    FileLogInputStream(FileRecords records,
                        int start,
                        int end) {
-        this.channel = channel;
+        this.fileRecords = records;
         this.position = start;
         this.end = end;
     }
 
     @Override
     public FileChannelRecordBatch nextBatch() throws IOException {
-        if (position + HEADER_SIZE_UP_TO_MAGIC >= end)
+        FileChannel channel = fileRecords.channel();
+        if (position >= end - HEADER_SIZE_UP_TO_MAGIC)
             return null;
 
         logHeaderBuffer.rewind();
@@ -71,18 +73,19 @@ public FileChannelRecordBatch nextBatch() throws IOException {
 
         // V0 has the smallest overhead, stricter checking is done later
         if (size < LegacyRecord.RECORD_OVERHEAD_V0)
-            throw new CorruptRecordException(String.format("Record size is smaller than minimum record overhead (%d).", LegacyRecord.RECORD_OVERHEAD_V0));
+            throw new CorruptRecordException(String.format("Found record size %d smaller than minimum record " +
+                            "overhead (%d) in file %s.", size, LegacyRecord.RECORD_OVERHEAD_V0, fileRecords.file()));
 
-        if (position + LOG_OVERHEAD + size > end)
+        if (position > end - LOG_OVERHEAD - size)
             return null;
 
         byte magic = logHeaderBuffer.get(MAGIC_OFFSET);
         final FileChannelRecordBatch batch;
 
         if (magic < RecordBatch.MAGIC_VALUE_V2)
-            batch = new LegacyFileChannelRecordBatch(offset, magic, channel, position, size);
+            batch = new LegacyFileChannelRecordBatch(offset, magic, fileRecords, position, size);
         else
-            batch = new DefaultFileChannelRecordBatch(offset, magic, channel, position, size);
+            batch = new DefaultFileChannelRecordBatch(offset, magic, fileRecords, position, size);
 
         position += batch.sizeInBytes();
         return batch;
@@ -96,7 +99,7 @@ public FileChannelRecordBatch nextBatch() throws IOException {
     public abstract static class FileChannelRecordBatch extends AbstractRecordBatch {
         protected final long offset;
         protected final byte magic;
-        protected final FileChannel channel;
+        protected final FileRecords fileRecords;
         protected final int position;
         protected final int batchSize;
 
@@ -105,12 +108,12 @@ public abstract static class FileChannelRecordBatch extends AbstractRecordBatch
 
         FileChannelRecordBatch(long offset,
                                byte magic,
-                               FileChannel channel,
+                               FileRecords fileRecords,
                                int position,
                                int batchSize) {
             this.offset = offset;
             this.magic = magic;
-            this.channel = channel;
+            this.fileRecords = fileRecords;
             this.position = position;
             this.batchSize = batchSize;
         }
@@ -171,14 +174,14 @@ public int sizeInBytes() {
 
         @Override
         public void writeTo(ByteBuffer buffer) {
+            FileChannel channel = fileRecords.channel();
             try {
                 int limit = buffer.limit();
                 buffer.limit(buffer.position() + sizeInBytes());
                 Utils.readFully(channel, buffer, position);
                 buffer.limit(limit);
             } catch (IOException e) {
-                throw new KafkaException("Failed to read record batch at position " + position + " from file channel " +
-                        channel, e);
+                throw new KafkaException("Failed to read record batch at position " + position + " from " + fileRecords, e);
             }
         }
 
@@ -205,13 +208,14 @@ protected RecordBatch loadBatchHeader() {
         }
 
         private RecordBatch loadBatchWithSize(int size, String description) {
+            FileChannel channel = fileRecords.channel();
             try {
                 ByteBuffer buffer = ByteBuffer.allocate(size);
                 Utils.readFullyOrFail(channel, buffer, position, description);
                 buffer.rewind();
                 return toMemoryRecordBatch(buffer);
             } catch (IOException e) {
-                throw new KafkaException(e);
+                throw new KafkaException("Failed to load record batch at position " + position + " from " + fileRecords, e);
             }
         }
 
@@ -224,15 +228,20 @@ public boolean equals(Object o) {
 
             FileChannelRecordBatch that = (FileChannelRecordBatch) o;
 
+            FileChannel channel = fileRecords == null ? null : fileRecords.channel();
+            FileChannel thatChannel = that.fileRecords == null ? null : that.fileRecords.channel();
+
             return offset == that.offset &&
                     position == that.position &&
                     batchSize == that.batchSize &&
-                    (channel == null ? that.channel == null : channel.equals(that.channel));
+                    Objects.equals(channel, thatChannel);
         }
 
         @Override
         public int hashCode() {
-            int result = (int) (offset ^ (offset >>> 32));
+            FileChannel channel = fileRecords == null ? null : fileRecords.channel();
+
+            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 36591a0ce2d55..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
@@ -19,18 +19,21 @@
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.network.TransportLayer;
 import org.apache.kafka.common.record.FileLogInputStream.FileChannelRecordBatch;
+import org.apache.kafka.common.utils.AbstractIterator;
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.common.utils.Utils;
 
 import java.io.Closeable;
 import java.io.File;
-import java.io.FileInputStream;
 import java.io.IOException;
 import java.io.RandomAccessFile;
 import java.nio.ByteBuffer;
 import java.nio.channels.FileChannel;
 import java.nio.channels.GatheringByteChannel;
-import java.util.Iterator;
+import java.nio.file.Files;
+import java.nio.file.StandardOpenOption;
+import java.util.Objects;
+import java.util.Optional;
 import java.util.concurrent.atomic.AtomicInteger;
 
 /**
@@ -53,11 +56,11 @@ public class FileRecords extends AbstractRecords implements Closeable {
      * The {@code FileRecords.open} methods should be used instead of this constructor whenever possible.
      * The constructor is visible for tests.
      */
-    public FileRecords(File file,
-                       FileChannel channel,
-                       int start,
-                       int end,
-                       boolean isSlice) throws IOException {
+    FileRecords(File file,
+                FileChannel channel,
+                int start,
+                int end,
+                boolean isSlice) throws IOException {
         this.file = file;
         this.channel = channel;
         this.start = start;
@@ -69,6 +72,10 @@ public FileRecords(File file,
             // don't check the file size if this is just a slice view
             size.set(end - start);
         } else {
+            if (channel.size() > Integer.MAX_VALUE)
+                throw new KafkaException("The size of segment " + file + " (" + channel.size() +
+                        ") is larger than the maximum allowed segment size of " + Integer.MAX_VALUE);
+
             int limit = Math.min((int) channel.size(), end);
             size.set(limit - start);
 
@@ -107,14 +114,12 @@ public FileChannel channel() {
      *
      * @param buffer The buffer to write the batches to
      * @param position Position in the buffer to read from
-     * @return The same buffer
      * @throws IOException If an I/O error occurs, see {@link FileChannel#read(ByteBuffer, long)} for details on the
      * possible exceptions
      */
-    public ByteBuffer readInto(ByteBuffer buffer, int position) throws IOException {
+    public void readInto(ByteBuffer buffer, int position) throws IOException {
         Utils.readFully(channel, buffer, position + this.start);
         buffer.flip();
-        return buffer;
     }
 
     /**
@@ -129,25 +134,36 @@ public ByteBuffer readInto(ByteBuffer buffer, int position) throws IOException {
      * @param size The number of bytes after the start position to include
      * @return A sliced wrapper on this message set limited based on the given position and size
      */
-    public FileRecords read(int position, int size) throws IOException {
+    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);
+            throw new IllegalArgumentException("Invalid position: " + position + " in read from " + this);
+        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);
+            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);
     }
 
     /**
-     * Append log batches to the buffer
+     * Append a set of records to the file. This method is not thread-safe and must be
+     * protected with a lock.
+     *
      * @param records The records to append
      * @return the number of bytes written to the underlying file
      */
     public int append(MemoryRecords records) throws IOException {
+        if (records.sizeInBytes() > Integer.MAX_VALUE - size.get())
+            throw new IllegalArgumentException("Append of size " + records.sizeInBytes() +
+                    " bytes is too large for segment with current file position at " + size.get());
+
         int written = records.writeFullyTo(channel);
         size.getAndAdd(written);
         return written;
@@ -178,11 +194,13 @@ public void closeHandlers() throws IOException {
 
     /**
      * Delete this message set from the filesystem
-     * @return True iff this message set was deleted.
+     * @throws IOException if deletion fails due to an I/O error
+     * @return  {@code true} if the file was deleted by this method; {@code false} if the file could not be deleted
+     *          because it did not exist
      */
-    public boolean delete() {
+    public boolean deleteIfExists() throws IOException {
         Utils.closeQuietly(channel, "FileChannel");
-        return file.delete();
+        return Files.deleteIfExists(file.toPath());
     }
 
     /**
@@ -193,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());
     }
 
     /**
@@ -225,7 +243,7 @@ public void renameTo(File f) throws IOException {
     public int truncateTo(int targetSize) throws IOException {
         int originalSize = sizeInBytes();
         if (targetSize > originalSize || targetSize < 0)
-            throw new KafkaException("Attempt to truncate log segment to " + targetSize + " bytes failed, " +
+            throw new KafkaException("Attempt to truncate log segment " + file + " to " + targetSize + " bytes failed, " +
                     " size of this log segment is " + originalSize + " bytes.");
         if (targetSize < (int) channel.size()) {
             channel.truncate(targetSize);
@@ -236,8 +254,8 @@ public int truncateTo(int targetSize) throws IOException {
 
     @Override
     public ConvertedRecords downConvert(byte toMagic, long firstOffset, Time time) {
-        ConvertedRecords convertedRecords = downConvert(batches, toMagic, firstOffset, time);
-        if (convertedRecords.recordsProcessingStats().numRecordsConverted() == 0) {
+        ConvertedRecords convertedRecords = RecordsUtil.downConvert(batches, toMagic, firstOffset, time);
+        if (convertedRecords.recordConversionStats().numRecordsConverted() == 0) {
             // This indicates that the message is too large, which means that the buffer is not large
             // enough to hold a full record batch. We just return all the bytes in this instance.
             // Even though the record batch does not have the right format version, we expect old clients
@@ -245,7 +263,7 @@ public ConvertedRecords downConvert(byte toMagic, long firstO
             // are not enough available bytes in the response to read it fully. Note that this is
             // only possible prior to KIP-74, after which the broker was changed to always return at least
             // one full record batch, even if it requires exceeding the max fetch size requested by the client.
-            return new ConvertedRecords<>(this, RecordsProcessingStats.EMPTY);
+            return new ConvertedRecords<>(this, RecordConversionStats.EMPTY);
         } else {
             return convertedRecords;
         }
@@ -307,7 +325,8 @@ public TimestampAndOffset searchForTimestamp(long targetTimestamp, int startingP
                 for (Record record : batch) {
                     long timestamp = record.timestamp();
                     if (timestamp >= targetTimestamp && record.offset() >= startingOffset)
-                        return new TimestampAndOffset(timestamp, record.offset());
+                        return new TimestampAndOffset(timestamp, record.offset(),
+                                maybeLeaderEpoch(batch.partitionLeaderEpoch()));
                 }
             }
         }
@@ -322,15 +341,23 @@ public TimestampAndOffset searchForTimestamp(long targetTimestamp, int startingP
     public TimestampAndOffset largestTimestampAfter(int startingPosition) {
         long maxTimestamp = RecordBatch.NO_TIMESTAMP;
         long offsetOfMaxTimestamp = -1L;
+        int leaderEpochOfMaxTimestamp = RecordBatch.NO_PARTITION_LEADER_EPOCH;
 
         for (RecordBatch batch : batchesFrom(startingPosition)) {
             long timestamp = batch.maxTimestamp();
             if (timestamp > maxTimestamp) {
                 maxTimestamp = timestamp;
                 offsetOfMaxTimestamp = batch.lastOffset();
+                leaderEpochOfMaxTimestamp = batch.partitionLeaderEpoch();
             }
         }
-        return new TimestampAndOffset(maxTimestamp, offsetOfMaxTimestamp);
+        return new TimestampAndOffset(maxTimestamp, offsetOfMaxTimestamp,
+                maybeLeaderEpoch(leaderEpochOfMaxTimestamp));
+    }
+
+    private Optional maybeLeaderEpoch(int leaderEpoch) {
+        return leaderEpoch == RecordBatch.NO_PARTITION_LEADER_EPOCH ?
+                Optional.empty() : Optional.of(leaderEpoch);
     }
 
     /**
@@ -344,22 +371,38 @@ public Iterable batches() {
         return batches;
     }
 
-    private Iterable batchesFrom(final int start) {
-        return new Iterable() {
-            @Override
-            public Iterator iterator() {
-                return batchIterator(start);
-            }
-        };
+    @Override
+    public String toString() {
+        return "FileRecords(size=" + sizeInBytes() +
+                ", file=" + file +
+                ", start=" + start +
+                ", end=" + end +
+                ")";
     }
 
-    private Iterator batchIterator(int start) {
+    /**
+     * Get an iterator over the record batches in the file, starting at a specific position. This is similar to
+     * {@link #batches()} except that callers specify a particular position to start reading the batches from. This
+     * method must be used with caution: the start position passed in must be a known start of a batch.
+     * @param start The position to start record iteration from; must be a known position for start of a batch
+     * @return An iterator over batches starting from {@code start}
+     */
+    public Iterable batchesFrom(final int start) {
+        return () -> batchIterator(start);
+    }
+
+    @Override
+    public AbstractIterator batchIterator() {
+        return batchIterator(start);
+    }
+
+    private AbstractIterator batchIterator(int start) {
         final int end;
         if (isSlice)
             end = this.end;
         else
             end = this.sizeInBytes();
-        FileLogInputStream inputStream = new FileLogInputStream(channel, start, end);
+        FileLogInputStream inputStream = new FileLogInputStream(this, start, end);
         return new RecordBatchIterator<>(inputStream);
     }
 
@@ -404,19 +447,16 @@ private static FileChannel openChannel(File file,
                                            int initFileSize,
                                            boolean preallocate) throws IOException {
         if (mutable) {
-            if (fileAlreadyExists) {
-                return new RandomAccessFile(file, "rw").getChannel();
+            if (fileAlreadyExists || !preallocate) {
+                return FileChannel.open(file.toPath(), StandardOpenOption.CREATE, StandardOpenOption.READ,
+                        StandardOpenOption.WRITE);
             } else {
-                if (preallocate) {
-                    RandomAccessFile randomAccessFile = new RandomAccessFile(file, "rw");
-                    randomAccessFile.setLength(initFileSize);
-                    return randomAccessFile.getChannel();
-                } else {
-                    return new RandomAccessFile(file, "rw").getChannel();
-                }
+                RandomAccessFile randomAccessFile = new RandomAccessFile(file, "rw");
+                randomAccessFile.setLength(initFileSize);
+                return randomAccessFile.getChannel();
             }
         } else {
-            return new FileInputStream(file).getChannel();
+            return FileChannel.open(file.toPath());
         }
     }
 
@@ -448,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;
@@ -467,28 +507,27 @@ public String toString() {
     public static class TimestampAndOffset {
         public final long timestamp;
         public final long offset;
+        public final Optional leaderEpoch;
 
-        public TimestampAndOffset(long timestamp, long offset) {
+        public TimestampAndOffset(long timestamp, long offset, Optional leaderEpoch) {
             this.timestamp = timestamp;
             this.offset = offset;
+            this.leaderEpoch = leaderEpoch;
         }
 
         @Override
         public boolean equals(Object o) {
             if (this == o) return true;
             if (o == null || getClass() != o.getClass()) return false;
-
             TimestampAndOffset that = (TimestampAndOffset) o;
-
-            if (timestamp != that.timestamp) return false;
-            return offset == that.offset;
+            return timestamp == that.timestamp &&
+                    offset == that.offset &&
+                    Objects.equals(leaderEpoch, that.leaderEpoch);
         }
 
         @Override
         public int hashCode() {
-            int result = (int) (timestamp ^ (timestamp >>> 32));
-            result = 31 * result + (int) (offset ^ (offset >>> 32));
-            return result;
+            return Objects.hash(timestamp, offset, leaderEpoch);
         }
 
         @Override
@@ -496,8 +535,8 @@ public String toString() {
             return "TimestampAndOffset(" +
                     "timestamp=" + timestamp +
                     ", offset=" + offset +
+                    ", leaderEpoch=" + leaderEpoch +
                     ')';
         }
     }
-
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/record/KafkaLZ4BlockInputStream.java b/clients/src/main/java/org/apache/kafka/common/record/KafkaLZ4BlockInputStream.java
index 56f10581db409..850b1e96e55f2 100644
--- a/clients/src/main/java/org/apache/kafka/common/record/KafkaLZ4BlockInputStream.java
+++ b/clients/src/main/java/org/apache/kafka/common/record/KafkaLZ4BlockInputStream.java
@@ -264,12 +264,12 @@ public long skip(long n) throws IOException {
     }
 
     @Override
-    public int available() throws IOException {
+    public int available() {
         return decompressedBuffer == null ? 0 : decompressedBuffer.remaining();
     }
 
     @Override
-    public void close() throws IOException {
+    public void close() {
         bufferSupplier.release(decompressionBuffer);
     }
 
@@ -279,7 +279,7 @@ public void mark(int readlimit) {
     }
 
     @Override
-    public void reset() throws IOException {
+    public void reset() {
         throw new RuntimeException("reset not supported");
     }
 
diff --git a/clients/src/main/java/org/apache/kafka/common/record/KafkaLZ4BlockOutputStream.java b/clients/src/main/java/org/apache/kafka/common/record/KafkaLZ4BlockOutputStream.java
index 8cfc37be826cb..591ab1693646c 100644
--- a/clients/src/main/java/org/apache/kafka/common/record/KafkaLZ4BlockOutputStream.java
+++ b/clients/src/main/java/org/apache/kafka/common/record/KafkaLZ4BlockOutputStream.java
@@ -16,7 +16,6 @@
  */
 package org.apache.kafka.common.record;
 
-import java.io.FilterOutputStream;
 import java.io.IOException;
 import java.io.OutputStream;
 
@@ -34,7 +33,7 @@
  *
  * This class is not thread-safe.
  */
-public final class KafkaLZ4BlockOutputStream extends FilterOutputStream {
+public final class KafkaLZ4BlockOutputStream extends OutputStream {
 
     public static final int MAGIC = 0x184D2204;
     public static final int LZ4_MAX_HEADER_LENGTH = 19;
@@ -52,9 +51,10 @@ public final class KafkaLZ4BlockOutputStream extends FilterOutputStream {
     private final boolean useBrokenFlagDescriptorChecksum;
     private final FLG flg;
     private final BD bd;
-    private final byte[] buffer;
-    private final byte[] compressedBuffer;
     private final int maxBlockSize;
+    private OutputStream out;
+    private byte[] buffer;
+    private byte[] compressedBuffer;
     private int bufferOffset;
     private boolean finished;
 
@@ -71,7 +71,7 @@ public final class KafkaLZ4BlockOutputStream extends FilterOutputStream {
      * @throws IOException
      */
     public KafkaLZ4BlockOutputStream(OutputStream out, int blockSize, boolean blockChecksum, boolean useBrokenFlagDescriptorChecksum) throws IOException {
-        super(out);
+        this.out = out;
         compressor = LZ4Factory.fastestInstance().fastCompressor();
         checksum = XXHashFactory.fastestInstance().hash32();
         this.useBrokenFlagDescriptorChecksum = useBrokenFlagDescriptorChecksum;
@@ -204,7 +204,6 @@ private void writeBlock() throws IOException {
     private void writeEndMark() throws IOException {
         ByteUtils.writeUnsignedIntLE(out, 0);
         // TODO implement content checksum, update flg.validate()
-        finished = true;
     }
 
     @Override
@@ -259,15 +258,26 @@ private void ensureNotFinished() {
 
     @Override
     public void close() throws IOException {
-        if (!finished) {
-            // basically flush the buffer writing the last block
-            writeBlock();
-            // write the end block and finish the stream
-            writeEndMark();
-        }
-        if (out != null) {
-            out.close();
-            out = null;
+        try {
+            if (!finished) {
+                // basically flush the buffer writing the last block
+                writeBlock();
+                // write the end block
+                writeEndMark();
+            }
+        } finally {
+            try {
+                if (out != null) {
+                    try (OutputStream outStream = out) {
+                        outStream.flush();
+                    }
+                }
+            } finally {
+                out = null;
+                buffer = null;
+                compressedBuffer = null;
+                finished = true;
+            }
         }
     }
 
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
new file mode 100644
index 0000000000000..c3598501f5091
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/record/LazyDownConversionRecords.java
@@ -0,0 +1,189 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.TopicPartition;
+import org.apache.kafka.common.utils.AbstractIterator;
+import org.apache.kafka.common.utils.Time;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Objects;
+
+/**
+ * Encapsulation for holding records that require down-conversion in a lazy, chunked manner (KIP-283). See
+ * {@link LazyDownConversionRecordsSend} for the actual chunked send implementation.
+ */
+public class LazyDownConversionRecords implements BaseRecords {
+    private final TopicPartition topicPartition;
+    private final Records records;
+    private final byte toMagic;
+    private final long firstOffset;
+    private ConvertedRecords firstConvertedBatch;
+    private final int sizeInBytes;
+    private final Time time;
+
+    /**
+     * @param topicPartition The topic-partition to which records belong
+     * @param records Records to lazily down-convert
+     * @param toMagic Magic version to down-convert to
+     * @param firstOffset The starting offset for down-converted records. This only impacts some cases. See
+     *                    {@link RecordsUtil#downConvert(Iterable, byte, long, Time)} for an explanation.
+     * @param time The time instance to use
+     *
+     * @throws org.apache.kafka.common.errors.UnsupportedCompressionTypeException If the first batch to down-convert
+     *    has a compression type which we do not support down-conversion for.
+     */
+    public LazyDownConversionRecords(TopicPartition topicPartition, Records records, byte toMagic, long firstOffset, Time time) {
+        this.topicPartition = Objects.requireNonNull(topicPartition);
+        this.records = Objects.requireNonNull(records);
+        this.toMagic = toMagic;
+        this.firstOffset = firstOffset;
+        this.time = Objects.requireNonNull(time);
+
+        // Kafka consumers expect at least one full batch of messages for every topic-partition. To guarantee this, we
+        // need to make sure that we are able to accommodate one full batch of down-converted messages. The way we achieve
+        // this is by having sizeInBytes method factor in the size of the first down-converted batch and return at least
+        // its size.
+        java.util.Iterator> it = iterator(0);
+        if (it.hasNext()) {
+            firstConvertedBatch = it.next();
+            sizeInBytes = Math.max(records.sizeInBytes(), firstConvertedBatch.records().sizeInBytes());
+        } else {
+            // If there are no messages we got after down-conversion, make sure we are able to send at least an overflow
+            // message to the consumer. Typically, the consumer would need to increase the fetch size in such cases.
+            firstConvertedBatch = null;
+            sizeInBytes = LazyDownConversionRecordsSend.MIN_OVERFLOW_MESSAGE_LENGTH;
+        }
+    }
+
+    @Override
+    public int sizeInBytes() {
+        return sizeInBytes;
+    }
+
+    @Override
+    public LazyDownConversionRecordsSend toSend() {
+        return new LazyDownConversionRecordsSend(this);
+    }
+
+    public TopicPartition topicPartition() {
+        return topicPartition;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (o instanceof LazyDownConversionRecords) {
+            LazyDownConversionRecords that = (LazyDownConversionRecords) o;
+            return toMagic == that.toMagic &&
+                    firstOffset == that.firstOffset &&
+                    topicPartition.equals(that.topicPartition) &&
+                    records.equals(that.records);
+        }
+        return false;
+    }
+
+    @Override
+    public int hashCode() {
+        int result = toMagic;
+        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.
+        ConvertedRecords firstBatch = firstConvertedBatch;
+        firstConvertedBatch = null;
+        return new Iterator(records, maximumReadSize, firstBatch);
+    }
+
+    /**
+     * Implementation for being able to iterate over down-converted records. Goal of this implementation is to keep
+     * it as memory-efficient as possible by not having to maintain all down-converted records in-memory. Maintains
+     * a view into batches of down-converted records.
+     */
+    private class Iterator extends AbstractIterator> {
+        private final AbstractIterator batchIterator;
+        private final long maximumReadSize;
+        private ConvertedRecords firstConvertedBatch;
+
+        /**
+         * @param recordsToDownConvert Records that require down-conversion
+         * @param maximumReadSize Maximum possible size of underlying records that will be down-converted in each call to
+         *                        {@link #makeNext()}. This is a soft limit as {@link #makeNext()} will always convert
+         *                        and return at least one full message batch.
+         */
+        private Iterator(Records recordsToDownConvert, long maximumReadSize, ConvertedRecords firstConvertedBatch) {
+            this.batchIterator = recordsToDownConvert.batchIterator();
+            this.maximumReadSize = maximumReadSize;
+            this.firstConvertedBatch = firstConvertedBatch;
+            // If we already have the first down-converted batch, advance the underlying records iterator to next batch
+            if (firstConvertedBatch != null)
+                this.batchIterator.next();
+        }
+
+        /**
+         * Make next set of down-converted records
+         * @return Down-converted records
+         */
+        @Override
+        protected ConvertedRecords makeNext() {
+            // If we have cached the first down-converted batch, return that now
+            if (firstConvertedBatch != null) {
+                ConvertedRecords convertedBatch = firstConvertedBatch;
+                firstConvertedBatch = null;
+                return convertedBatch;
+            }
+
+            while (batchIterator.hasNext()) {
+                final List batches = new ArrayList<>();
+                boolean isFirstBatch = true;
+                long sizeSoFar = 0;
+
+                // Figure out batches we should down-convert based on the size constraints
+                while (batchIterator.hasNext() &&
+                        (isFirstBatch || (batchIterator.peek().sizeInBytes() + sizeSoFar) <= maximumReadSize)) {
+                    RecordBatch currentBatch = batchIterator.next();
+                    batches.add(currentBatch);
+                    sizeSoFar += currentBatch.sizeInBytes();
+                    isFirstBatch = false;
+                }
+
+                ConvertedRecords convertedRecords = RecordsUtil.downConvert(batches, toMagic, firstOffset, time);
+                // During conversion, it is possible that we drop certain batches because they do not have an equivalent
+                // representation in the message format we want to convert to. For example, V0 and V1 message formats
+                // have no notion of transaction markers which were introduced in V2 so they get dropped during conversion.
+                // We return converted records only when we have at least one valid batch of messages after conversion.
+                if (convertedRecords.records().sizeInBytes() > 0)
+                    return convertedRecords;
+            }
+            return allDone();
+        }
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/record/LazyDownConversionRecordsSend.java b/clients/src/main/java/org/apache/kafka/common/record/LazyDownConversionRecordsSend.java
new file mode 100644
index 0000000000000..34c3d69979097
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/record/LazyDownConversionRecordsSend.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.record;
+
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.errors.UnsupportedCompressionTypeException;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.channels.GatheringByteChannel;
+import java.util.Iterator;
+
+/**
+ * Encapsulation for {@link RecordsSend} for {@link LazyDownConversionRecords}. Records are down-converted in batches and
+ * on-demand when {@link #writeTo} method is called.
+ */
+public final class LazyDownConversionRecordsSend extends RecordsSend {
+    private static final Logger log = LoggerFactory.getLogger(LazyDownConversionRecordsSend.class);
+    private static final int MAX_READ_SIZE = 128 * 1024;
+    static final int MIN_OVERFLOW_MESSAGE_LENGTH = Records.LOG_OVERHEAD;
+
+    private RecordConversionStats recordConversionStats;
+    private RecordsSend convertedRecordsWriter;
+    private Iterator> convertedRecordsIterator;
+
+    public LazyDownConversionRecordsSend(LazyDownConversionRecords records) {
+        super(records, records.sizeInBytes());
+        convertedRecordsWriter = null;
+        recordConversionStats = new RecordConversionStats();
+        convertedRecordsIterator = records().iterator(MAX_READ_SIZE);
+    }
+
+    private MemoryRecords buildOverflowBatch(int remaining) {
+        // We do not have any records left to down-convert. Construct an overflow message for the length remaining.
+        // This message will be ignored by the consumer because its length will be past the length of maximum
+        // possible response size.
+        // DefaultRecordBatch =>
+        //      BaseOffset => Int64
+        //      Length => Int32
+        //      ...
+        ByteBuffer overflowMessageBatch = ByteBuffer.allocate(
+                Math.max(MIN_OVERFLOW_MESSAGE_LENGTH, Math.min(remaining + 1, MAX_READ_SIZE)));
+        overflowMessageBatch.putLong(-1L);
+
+        // Fill in the length of the overflow batch. A valid batch must be at least as long as the minimum batch
+        // overhead.
+        overflowMessageBatch.putInt(Math.max(remaining + 1, DefaultRecordBatch.RECORD_BATCH_OVERHEAD));
+        log.debug("Constructed overflow message batch for partition {} with length={}", topicPartition(), remaining);
+        return MemoryRecords.readableRecords(overflowMessageBatch);
+    }
+
+    @Override
+    public long writeTo(GatheringByteChannel channel, long previouslyWritten, int remaining) throws IOException {
+        if (convertedRecordsWriter == null || convertedRecordsWriter.completed()) {
+            MemoryRecords convertedRecords;
+
+            try {
+                // Check if we have more chunks left to down-convert
+                if (convertedRecordsIterator.hasNext()) {
+                    // Get next chunk of down-converted messages
+                    ConvertedRecords recordsAndStats = convertedRecordsIterator.next();
+                    convertedRecords = (MemoryRecords) recordsAndStats.records();
+                    recordConversionStats.add(recordsAndStats.recordConversionStats());
+                    log.debug("Down-converted records for partition {} with length={}", topicPartition(), convertedRecords.sizeInBytes());
+                } else {
+                    convertedRecords = buildOverflowBatch(remaining);
+                }
+            } catch (UnsupportedCompressionTypeException e) {
+                // We have encountered a compression type which does not support down-conversion (e.g. zstd).
+                // Since we have already sent at least one batch and we have committed to the fetch size, we
+                // send an overflow batch. The consumer will read the first few records and then fetch from the
+                // offset of the batch which has the unsupported compression type. At that time, we will
+                // send back the UNSUPPORTED_COMPRESSION_TYPE erro which will allow the consumer to fail gracefully.
+                convertedRecords = buildOverflowBatch(remaining);
+            }
+
+            convertedRecordsWriter = new DefaultRecordsSend(convertedRecords, Math.min(convertedRecords.sizeInBytes(), remaining));
+        }
+        return convertedRecordsWriter.writeTo(channel);
+    }
+
+    public RecordConversionStats recordConversionStats() {
+        return recordConversionStats;
+    }
+
+    public TopicPartition topicPartition() {
+        return records().topicPartition();
+    }
+}
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 2a25aade9beef..72f39e9a73085 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
@@ -17,7 +17,10 @@
 package org.apache.kafka.common.record;
 
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.errors.CorruptRecordException;
+import org.apache.kafka.common.message.LeaderChangeMessage;
 import org.apache.kafka.common.record.MemoryRecords.RecordFilter.BatchRetention;
+import org.apache.kafka.common.utils.AbstractIterator;
 import org.apache.kafka.common.utils.ByteBufferOutputStream;
 import org.apache.kafka.common.utils.CloseableIterator;
 import org.apache.kafka.common.utils.Time;
@@ -29,7 +32,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;
 
@@ -44,12 +46,7 @@ public class MemoryRecords extends AbstractRecords {
 
     private final ByteBuffer buffer;
 
-    private final Iterable batches = new Iterable() {
-        @Override
-        public Iterator iterator() {
-            return new RecordBatchIterator<>(new ByteBufferLogInputStream(buffer.duplicate(), Integer.MAX_VALUE));
-        }
-    };
+    private final Iterable batches = this::batchIterator;
 
     private int validBytes = -1;
 
@@ -113,7 +110,24 @@ public int validBytes() {
 
     @Override
     public ConvertedRecords downConvert(byte toMagic, long firstOffset, Time time) {
-        return downConvert(batches(), toMagic, firstOffset, time);
+        return RecordsUtil.downConvert(batches(), toMagic, firstOffset, time);
+    }
+
+    @Override
+    public AbstractIterator batchIterator() {
+        return new RecordBatchIterator<>(new ByteBufferLogInputStream(buffer.duplicate(), Integer.MAX_VALUE));
+    }
+
+    /**
+     * Validates the header of the first batch and returns batch size.
+     * @return first batch size including LOG_OVERHEAD if buffer contains header up to
+     *         magic byte, null otherwise
+     * @throws CorruptRecordException if record size or magic is invalid
+     */
+    public Integer firstBatchSize() {
+        if (buffer.remaining() < HEADER_SIZE_UP_TO_MAGIC)
+            return null;
+        return new ByteBufferLogInputStream(buffer, Integer.MAX_VALUE).nextBatchSize();
     }
 
     /**
@@ -140,20 +154,14 @@ public FilterResult filterTo(TopicPartition partition, RecordFilter filter, Byte
     private static FilterResult filterTo(TopicPartition partition, Iterable batches,
                                          RecordFilter filter, ByteBuffer destinationBuffer, int maxRecordBatchSize,
                                          BufferSupplier decompressionBufferSupplier) {
-        long maxTimestamp = RecordBatch.NO_TIMESTAMP;
-        long maxOffset = -1L;
-        long shallowOffsetOfMaxTimestamp = -1L;
-        int messagesRead = 0;
-        int bytesRead = 0;
-        int messagesRetained = 0;
-        int bytesRetained = 0;
-
+        FilterResult filterResult = new FilterResult(destinationBuffer);
         ByteBufferOutputStream bufferOutputStream = new ByteBufferOutputStream(destinationBuffer);
 
         for (MutableRecordBatch batch : batches) {
-            bytesRead += batch.sizeInBytes();
-
+            long maxOffset = -1L;
             BatchRetention batchRetention = filter.checkBatchRetention(batch);
+            filterResult.bytesRead += batch.sizeInBytes();
+
             if (batchRetention == BatchRetention.DELETE)
                 continue;
 
@@ -169,7 +177,7 @@ private static FilterResult filterTo(TopicPartition partition, Iterable iterator = batch.streamingIterator(decompressionBufferSupplier)) {
                 while (iterator.hasNext()) {
                     Record record = iterator.next();
-                    messagesRead += 1;
+                    filterResult.messagesRead += 1;
 
                     if (filter.shouldRetainRecord(batch, record)) {
                         // Check for log corruption due to KAFKA-4298. If we find it, make sure that we overwrite
@@ -190,20 +198,11 @@ private static FilterResult filterTo(TopicPartition partition, Iterable maxTimestamp) {
-                        maxTimestamp = batch.maxTimestamp();
-                        shallowOffsetOfMaxTimestamp = batch.lastOffset();
-                    }
+                    filterResult.updateRetainedBatchMetadata(batch, retainedRecords.size(), false);
                 } else {
                     MemoryRecordsBuilder builder = buildRetainedRecordsInto(batch, retainedRecords, bufferOutputStream);
                     MemoryRecords records = builder.build();
                     int filteredBatchSize = records.sizeInBytes();
-
-                    messagesRetained += retainedRecords.size();
-                    bytesRetained += filteredBatchSize;
-
                     if (filteredBatchSize > batch.sizeInBytes() && filteredBatchSize > maxRecordBatchSize)
                         log.warn("Record batch from {} with last offset {} exceeded max record batch size {} after cleaning " +
                                         "(new size is {}). Consumers with version earlier than 0.10.1.0 may need to " +
@@ -211,10 +210,8 @@ private static FilterResult filterTo(TopicPartition partition, Iterable maxTimestamp) {
-                        maxTimestamp = info.maxTimestamp;
-                        shallowOffsetOfMaxTimestamp = info.shallowOffsetOfMaxTimestamp;
-                    }
+                    filterResult.updateRetainedBatchMetadata(info.maxTimestamp, info.shallowOffsetOfMaxTimestamp,
+                            maxOffset, retainedRecords.size(), filteredBatchSize);
                 }
             } else if (batchRetention == BatchRetention.RETAIN_EMPTY) {
                 if (batchMagic < RecordBatch.MAGIC_VALUE_V2)
@@ -225,18 +222,19 @@ private static FilterResult filterTo(TopicPartition partition, Iterable batches() {
 
     @Override
     public String toString() {
-        Iterator iter = records().iterator();
-        StringBuilder builder = new StringBuilder();
-        builder.append('[');
-        while (iter.hasNext()) {
-            Record record = iter.next();
-            builder.append('(');
-            builder.append("record=");
-            builder.append(record);
-            builder.append(")");
-            if (iter.hasNext())
-                builder.append(", ");
-        }
-        builder.append(']');
-        return builder.toString();
+        return "MemoryRecords(size=" + sizeInBytes() +
+                ", buffer=" + buffer +
+                ")";
     }
 
     @Override
@@ -335,31 +322,76 @@ public enum BatchRetention {
     }
 
     public static class FilterResult {
-        public final ByteBuffer output;
-        public final int messagesRead;
-        public final int bytesRead;
-        public final int messagesRetained;
-        public final int bytesRetained;
-        public final long maxOffset;
-        public final long maxTimestamp;
-        public final long shallowOffsetOfMaxTimestamp;
-
-        public FilterResult(ByteBuffer output,
-                            int messagesRead,
-                            int bytesRead,
-                            int messagesRetained,
-                            int bytesRetained,
-                            long maxOffset,
-                            long maxTimestamp,
-                            long shallowOffsetOfMaxTimestamp) {
-            this.output = output;
-            this.messagesRead = messagesRead;
-            this.bytesRead = bytesRead;
-            this.messagesRetained = messagesRetained;
-            this.bytesRetained = bytesRetained;
-            this.maxOffset = maxOffset;
-            this.maxTimestamp = maxTimestamp;
-            this.shallowOffsetOfMaxTimestamp = shallowOffsetOfMaxTimestamp;
+        private ByteBuffer outputBuffer;
+        private int messagesRead = 0;
+        // Note that `bytesRead` should contain only bytes from batches that have been processed, i.e. bytes from
+        // `messagesRead` and any discarded batches.
+        private int bytesRead = 0;
+        private int messagesRetained = 0;
+        private int bytesRetained = 0;
+        private long maxOffset = -1L;
+        private long maxTimestamp = RecordBatch.NO_TIMESTAMP;
+        private long shallowOffsetOfMaxTimestamp = -1L;
+
+        private FilterResult(ByteBuffer outputBuffer) {
+            this.outputBuffer = outputBuffer;
+        }
+
+        private void updateRetainedBatchMetadata(MutableRecordBatch retainedBatch, int numMessagesInBatch, boolean headerOnly) {
+            int bytesRetained = headerOnly ? DefaultRecordBatch.RECORD_BATCH_OVERHEAD : retainedBatch.sizeInBytes();
+            updateRetainedBatchMetadata(retainedBatch.maxTimestamp(), retainedBatch.lastOffset(),
+                    retainedBatch.lastOffset(), numMessagesInBatch, bytesRetained);
+        }
+
+        private void updateRetainedBatchMetadata(long maxTimestamp, long shallowOffsetOfMaxTimestamp, long maxOffset,
+                                                int messagesRetained, int bytesRetained) {
+            validateBatchMetadata(maxTimestamp, shallowOffsetOfMaxTimestamp, maxOffset);
+            if (maxTimestamp > this.maxTimestamp) {
+                this.maxTimestamp = maxTimestamp;
+                this.shallowOffsetOfMaxTimestamp = shallowOffsetOfMaxTimestamp;
+            }
+            this.maxOffset = Math.max(maxOffset, this.maxOffset);
+            this.messagesRetained += messagesRetained;
+            this.bytesRetained += bytesRetained;
+        }
+
+        private void validateBatchMetadata(long maxTimestamp, long shallowOffsetOfMaxTimestamp, long maxOffset) {
+            if (maxTimestamp != RecordBatch.NO_TIMESTAMP && shallowOffsetOfMaxTimestamp < 0)
+                throw new IllegalArgumentException("shallowOffset undefined for maximum timestamp " + maxTimestamp);
+            if (maxOffset < 0)
+                throw new IllegalArgumentException("maxOffset undefined");
+        }
+
+        public ByteBuffer outputBuffer() {
+            return outputBuffer;
+        }
+
+        public int messagesRead() {
+            return messagesRead;
+        }
+
+        public int bytesRead() {
+            return bytesRead;
+        }
+
+        public int messagesRetained() {
+            return messagesRetained;
+        }
+
+        public int bytesRetained() {
+            return bytesRetained;
+        }
+
+        public long maxOffset() {
+            return maxOffset;
+        }
+
+        public long maxTimestamp() {
+            return maxTimestamp;
+        }
+
+        public long shallowOffsetOfMaxTimestamp() {
+            return shallowOffsetOfMaxTimestamp;
         }
     }
 
@@ -495,6 +527,10 @@ public static MemoryRecords withRecords(long initialOffset, CompressionType comp
                 records);
     }
 
+    public static MemoryRecords withRecords(byte magic, long initialOffset, CompressionType compressionType, SimpleRecord... records) {
+        return withRecords(magic, initialOffset, compressionType, TimestampType.CREATE_TIME, records);
+    }
+
     public static MemoryRecords withRecords(long initialOffset, CompressionType compressionType, Integer partitionLeaderEpoch, SimpleRecord... records) {
         return withRecords(RecordBatch.CURRENT_MAGIC_VALUE, initialOffset, compressionType, TimestampType.CREATE_TIME, RecordBatch.NO_PRODUCER_ID,
                 RecordBatch.NO_PRODUCER_EPOCH, RecordBatch.NO_SEQUENCE, partitionLeaderEpoch, false, records);
@@ -550,7 +586,7 @@ public static MemoryRecords withRecords(byte magic, long initialOffset, Compress
     public static MemoryRecords withRecords(byte magic, long initialOffset, CompressionType compressionType,
                                             TimestampType timestampType, long producerId, short producerEpoch,
                                             int baseSequence, int partitionLeaderEpoch, boolean isTransactional,
-                                            SimpleRecord ... records) {
+                                            SimpleRecord... records) {
         if (records.length == 0)
             return MemoryRecords.EMPTY;
         int sizeEstimate = AbstractRecords.estimateSizeInBytes(magic, compressionType, Arrays.asList(records));
@@ -593,13 +629,35 @@ public static void writeEndTransactionalMarker(ByteBuffer buffer, long initialOf
                                                    int partitionLeaderEpoch, long producerId, short producerEpoch,
                                                    EndTransactionMarker marker) {
         boolean isTransactional = true;
-        boolean isControlBatch = true;
         MemoryRecordsBuilder builder = new MemoryRecordsBuilder(buffer, RecordBatch.CURRENT_MAGIC_VALUE, CompressionType.NONE,
                 TimestampType.CREATE_TIME, initialOffset, timestamp, producerId, producerEpoch,
-                RecordBatch.NO_SEQUENCE, isTransactional, isControlBatch, partitionLeaderEpoch,
+                RecordBatch.NO_SEQUENCE, isTransactional, true, partitionLeaderEpoch,
                 buffer.capacity());
         builder.appendEndTxnMarker(timestamp, marker);
         builder.close();
     }
 
+    public static MemoryRecords withLeaderChangeMessage(long timestamp, int leaderEpoch, LeaderChangeMessage leaderChangeMessage) {
+        // To avoid calling message toStruct multiple times, we supply a fixed message size
+        // for leader change, as it happens rare and the buffer could still grow if not sufficient in
+        // certain edge cases.
+        ByteBuffer buffer = ByteBuffer.allocate(256);
+        writeLeaderChangeMessage(buffer, 0L, timestamp, leaderEpoch, leaderChangeMessage);
+        buffer.flip();
+        return MemoryRecords.readableRecords(buffer);
+    }
+
+    private static void writeLeaderChangeMessage(ByteBuffer buffer,
+                                                 long initialOffset,
+                                                 long timestamp,
+                                                 int leaderEpoch,
+                                                 LeaderChangeMessage leaderChangeMessage) {
+        MemoryRecordsBuilder builder = new MemoryRecordsBuilder(buffer, RecordBatch.CURRENT_MAGIC_VALUE, CompressionType.NONE,
+            TimestampType.CREATE_TIME, initialOffset, timestamp,
+            RecordBatch.NO_PRODUCER_ID, RecordBatch.NO_PRODUCER_EPOCH, RecordBatch.NO_SEQUENCE,
+            false, true, leaderEpoch, buffer.capacity());
+        builder.appendLeaderChangeMessage(timestamp, leaderChangeMessage);
+        builder.close();
+    }
+
 }
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 a9b57ac22df09..480f4dc7a343b 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
@@ -18,11 +18,16 @@
 
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.header.Header;
+import org.apache.kafka.common.message.LeaderChangeMessage;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.ObjectSerializationCache;
 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;
+import java.io.OutputStream;
 import java.nio.ByteBuffer;
 
 import static org.apache.kafka.common.utils.Utils.wrapNullable;
@@ -36,13 +41,17 @@
  * and the builder is closed (e.g. the Producer), it's important to call `closeForRecordAppends` when the former happens.
  * This will release resources like compression buffers that can be relatively large (64 KB for LZ4).
  */
-public class MemoryRecordsBuilder {
+public class MemoryRecordsBuilder implements AutoCloseable {
     private static final float COMPRESSION_RATE_ESTIMATION_FACTOR = 1.05f;
+    private static final DataOutputStream CLOSED_STREAM = new DataOutputStream(new OutputStream() {
+        @Override
+        public void write(int b) {
+            throw new IllegalStateException("MemoryRecordsBuilder is closed for record appends");
+        }
+    });
 
     private final TimestampType timestampType;
     private final CompressionType compressionType;
-    // Used to append records, may compress data on the fly
-    private final DataOutputStream appendStream;
     // Used to hold a reference to the underlying ByteBuffer so that we can write the record batch header and access
     // the written bytes. ByteBufferOutputStream allocates a new ByteBuffer if the existing one is not large enough,
     // so it's not safe to hold a direct reference to the underlying ByteBuffer.
@@ -60,7 +69,8 @@ public class MemoryRecordsBuilder {
     // from previous batches before appending any records.
     private float estimatedCompressionRatio = 1.0F;
 
-    private boolean appendStreamIsClosed = false;
+    // Used to append records, may compress data on the fly
+    private DataOutputStream appendStream;
     private boolean isTransactional;
     private long producerId;
     private short producerEpoch;
@@ -96,6 +106,8 @@ public MemoryRecordsBuilder(ByteBufferOutputStream bufferStream,
                 throw new IllegalArgumentException("Transactional records are not supported for magic " + magic);
             if (isControlBatch)
                 throw new IllegalArgumentException("Control records are not supported for magic " + magic);
+            if (compressionType == CompressionType.ZSTD)
+                throw new IllegalArgumentException("ZStandard compression is not supported for magic " + magic);
         }
 
         this.magic = magic;
@@ -265,12 +277,13 @@ public void overrideLastOffset(long lastOffset) {
      * possible to update the RecordBatch header.
      */
     public void closeForRecordAppends() {
-        if (!appendStreamIsClosed) {
+        if (appendStream != CLOSED_STREAM) {
             try {
                 appendStream.close();
-                appendStreamIsClosed = true;
             } catch (IOException e) {
                 throw new KafkaException(e);
+            } finally {
+                appendStream = CLOSED_STREAM;
             }
         }
     }
@@ -409,7 +422,7 @@ private Long appendWithOffset(long offset, boolean isControlRecord, long timesta
                 appendDefaultRecord(offset, timestamp, key, value, headers);
                 return null;
             } else {
-                return appendLegacyRecord(offset, timestamp, key, value);
+                return appendLegacyRecord(offset, timestamp, key, value, magic);
             }
         } catch (IOException e) {
             throw new KafkaException("I/O exception when writing to the append stream, closing", e);
@@ -476,6 +489,24 @@ public Long appendWithOffset(long offset, SimpleRecord record) {
         return appendWithOffset(offset, record.timestamp(), record.key(), record.value(), record.headers());
     }
 
+    /**
+     * Append a control record at the given offset. The control record type must be known or
+     * this method will raise an error.
+     *
+     * @param offset The absolute offset of the record in the log buffer
+     * @param record The record to append
+     * @return CRC of the record or null if record-level CRC is not supported for the message format
+     */
+    public Long appendControlRecordWithOffset(long offset, SimpleRecord record) {
+        short typeId = ControlRecordType.parseTypeId(record.key());
+        ControlRecordType type = ControlRecordType.fromTypeId(typeId);
+        if (type == ControlRecordType.UNKNOWN)
+            throw new IllegalArgumentException("Cannot append record with unknown control record type " + typeId);
+
+        return appendWithOffset(offset, true, record.timestamp(),
+            record.key(), record.value(), record.headers());
+    }
+
     /**
      * Append a new record at the next sequential offset.
      * @param timestamp The record timestamp
@@ -558,6 +589,23 @@ public Long appendEndTxnMarker(long timestamp, EndTransactionMarker marker) {
         return appendControlRecord(timestamp, marker.controlType(), value);
     }
 
+    /**
+     * Return CRC of the record or null if record-level CRC is not supported for the message format
+     */
+    public Long appendLeaderChangeMessage(long timestamp, LeaderChangeMessage leaderChangeMessage) {
+        if (partitionLeaderEpoch == RecordBatch.NO_PARTITION_LEADER_EPOCH) {
+            throw new IllegalArgumentException("Partition leader epoch must be valid, but get " + partitionLeaderEpoch);
+        }
+
+        ObjectSerializationCache cache = new ObjectSerializationCache();
+        int size = leaderChangeMessage.size(cache, ControlRecordUtils.LEADER_CHANGE_SCHEMA_VERSION);
+        ByteBuffer buffer = ByteBuffer.allocate(size);
+        ByteBufferAccessor accessor = new ByteBufferAccessor(buffer);
+        leaderChangeMessage.write(accessor, cache, ControlRecordUtils.LEADER_CHANGE_SCHEMA_VERSION);
+        buffer.flip();
+        return appendControlRecord(timestamp, ControlRecordType.LEADER_CHANGE, buffer);
+    }
+
     /**
      * Add a legacy record without doing offset/magic validation (this should only be used in testing).
      * @param offset The offset of the record
@@ -578,6 +626,35 @@ 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 a record at the next sequential offset.
      * @param record the record to add
@@ -623,7 +700,7 @@ private void appendDefaultRecord(long offset, long timestamp, ByteBuffer key, By
         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;
@@ -663,7 +740,7 @@ private void recordWritten(long offset, long timestamp, int size) {
     }
 
     private void ensureOpenForRecordAppend() {
-        if (appendStreamIsClosed)
+        if (appendStream == CLOSED_STREAM)
             throw new IllegalStateException("Tried to append a record, but MemoryRecordsBuilder is closed for record appends");
     }
 
@@ -738,7 +815,7 @@ public boolean isClosed() {
     public boolean isFull() {
         // note that the write limit is respected only after the first record is added which ensures we can always
         // create non-empty batches (this is used to disable batching when the producer's batch size is set to 0).
-        return appendStreamIsClosed || (this.numRecords > 0 && this.writeLimit <= estimatedBytesWritten());
+        return appendStream == CLOSED_STREAM || (this.numRecords > 0 && this.writeLimit <= estimatedBytesWritten());
     }
 
     /**
diff --git a/clients/src/main/java/org/apache/kafka/common/record/MultiRecordsSend.java b/clients/src/main/java/org/apache/kafka/common/record/MultiRecordsSend.java
new file mode 100644
index 0000000000000..8b3200d8cd814
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/record/MultiRecordsSend.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.common.record;
+
+import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.network.Send;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.nio.channels.GatheringByteChannel;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Queue;
+
+/**
+ * A set of composite sends with nested {@link RecordsSend}, sent one after another
+ */
+public class MultiRecordsSend implements Send {
+    private static final Logger log = LoggerFactory.getLogger(MultiRecordsSend.class);
+
+    private final Queue sendQueue;
+    private final long size;
+    private Map recordConversionStats;
+
+    private long totalWritten = 0;
+    private Send current;
+
+    /**
+     * Construct a MultiRecordsSend from a queue of Send objects. The queue will be consumed as the MultiRecordsSend
+     * progresses (on completion, it will be empty).
+     */
+    public MultiRecordsSend(Queue sends) {
+        this.sendQueue = sends;
+
+        long size = 0;
+        for (Send send : sends)
+            size += send.size();
+        this.size = size;
+
+        this.current = sendQueue.poll();
+    }
+
+    public MultiRecordsSend(Queue sends, long size) {
+        this.sendQueue = sends;
+        this.size = size;
+        this.current = sendQueue.poll();
+    }
+
+    @Override
+    public long size() {
+        return size;
+    }
+
+    @Override
+    public boolean completed() {
+        return current == null;
+    }
+
+    // Visible for testing
+    int numResidentSends() {
+        int count = 0;
+        if (current != null)
+            count += 1;
+        count += sendQueue.size();
+        return count;
+    }
+
+    @Override
+    public long writeTo(GatheringByteChannel channel) throws IOException {
+        if (completed())
+            throw new KafkaException("This operation cannot be invoked on a complete request.");
+
+        int totalWrittenPerCall = 0;
+        boolean sendComplete;
+        do {
+            long written = current.writeTo(channel);
+            totalWrittenPerCall += written;
+            sendComplete = current.completed();
+            if (sendComplete) {
+                updateRecordConversionStats(current);
+                current = sendQueue.poll();
+            }
+        } while (!completed() && sendComplete);
+
+        totalWritten += totalWrittenPerCall;
+
+        if (completed() && totalWritten != size)
+            log.error("mismatch in sending bytes over socket; expected: {} actual: {}", size, totalWritten);
+
+        log.trace("Bytes written as part of multi-send call: {}, total bytes written so far: {}, expected bytes to write: {}",
+                totalWrittenPerCall, totalWritten, size);
+
+        return totalWrittenPerCall;
+    }
+
+    /**
+     * Get any statistics that were recorded as part of executing this {@link MultiRecordsSend}.
+     * @return Records processing statistics (could be null if no statistics were collected)
+     */
+    public Map recordConversionStats() {
+        return recordConversionStats;
+    }
+
+    @Override
+    public String toString() {
+        return "MultiRecordsSend(" +
+            "size=" + size +
+            ", totalWritten=" + totalWritten +
+            ')';
+    }
+
+    private void updateRecordConversionStats(Send completedSend) {
+        // The underlying send might have accumulated statistics that need to be recorded. For example,
+        // LazyDownConversionRecordsSend accumulates statistics related to the number of bytes down-converted, the amount
+        // of temporary memory used for down-conversion, etc. Pull out any such statistics from the underlying send
+        // and fold it up appropriately.
+        if (completedSend instanceof LazyDownConversionRecordsSend) {
+            if (recordConversionStats == null)
+                recordConversionStats = new HashMap<>();
+            LazyDownConversionRecordsSend lazyRecordsSend = (LazyDownConversionRecordsSend) completedSend;
+            recordConversionStats.put(lazyRecordsSend.topicPartition(), lazyRecordsSend.recordConversionStats());
+        }
+    }
+}
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/RecordConversionStats.java b/clients/src/main/java/org/apache/kafka/common/record/RecordConversionStats.java
new file mode 100644
index 0000000000000..4f0bca527fb97
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/record/RecordConversionStats.java
@@ -0,0 +1,70 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.common.record;
+
+public class RecordConversionStats {
+
+    public static final RecordConversionStats EMPTY = new RecordConversionStats();
+
+    private long temporaryMemoryBytes;
+    private int numRecordsConverted;
+    private long conversionTimeNanos;
+
+    public RecordConversionStats(long temporaryMemoryBytes, int numRecordsConverted, long conversionTimeNanos) {
+        this.temporaryMemoryBytes = temporaryMemoryBytes;
+        this.numRecordsConverted = numRecordsConverted;
+        this.conversionTimeNanos = conversionTimeNanos;
+    }
+
+    public RecordConversionStats() {
+        this(0, 0, 0);
+    }
+
+    public void add(RecordConversionStats stats) {
+        temporaryMemoryBytes += stats.temporaryMemoryBytes;
+        numRecordsConverted += stats.numRecordsConverted;
+        conversionTimeNanos += stats.conversionTimeNanos;
+    }
+
+    /**
+     * Returns the number of temporary memory bytes allocated to process the records.
+     * This size depends on whether the records need decompression and/or conversion:
+     * 
          
+     *   
        • Non compressed, no conversion: zero
          • 
+     *   
        • Non compressed, with conversion: size of the converted buffer
          • 
+     *   
        • Compressed, no conversion: size of the original buffer after decompression
          • 
+     *   
        • Compressed, with conversion: size of the original buffer after decompression + size of the converted buffer uncompressed
          • 
+     * 
        
+     */
+    public long temporaryMemoryBytes() {
+        return temporaryMemoryBytes;
+    }
+
+    public int numRecordsConverted() {
+        return numRecordsConverted;
+    }
+
+    public long conversionTimeNanos() {
+        return conversionTimeNanos;
+    }
+
+    @Override
+    public String toString() {
+        return String.format("RecordConversionStats(temporaryMemoryBytes=%d, numRecordsConverted=%d, conversionTimeNanos=%d)",
+                temporaryMemoryBytes, numRecordsConverted, conversionTimeNanos);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/record/RecordVersion.java b/clients/src/main/java/org/apache/kafka/common/record/RecordVersion.java
new file mode 100644
index 0000000000000..8406d5331c8a8
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/record/RecordVersion.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.record;
+
+/**
+ * Defines the record format versions supported by Kafka.
+ *
+ * For historical reasons, the record format version is also known as `magic` and `message format version`. Note that
+ * the version actually applies to the {@link RecordBatch} (instead of the {@link Record}). Finally, the
+ * `message.format.version` topic config confusingly expects an ApiVersion instead of a RecordVersion.
+ */
+public enum RecordVersion {
+    V0(0), V1(1), V2(2);
+
+    private static final RecordVersion[] VALUES = values();
+
+    public final byte value;
+
+    RecordVersion(int value) {
+        this.value = (byte) value;
+    }
+
+    /**
+     * Check whether this version precedes another version.
+     *
+     * @return true only if the magic value is less than the other's
+     */
+    public boolean precedes(RecordVersion other) {
+        return this.value < other.value;
+    }
+
+    public static RecordVersion lookup(byte value) {
+        if (value < 0 || value >= VALUES.length)
+            throw new IllegalArgumentException("Unknown record version: " + value);
+        return VALUES[value];
+    }
+
+    public static RecordVersion current() {
+        return V2;
+    }
+
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/record/Records.java b/clients/src/main/java/org/apache/kafka/common/record/Records.java
index 19152bae26790..23607b4617190 100644
--- a/clients/src/main/java/org/apache/kafka/common/record/Records.java
+++ b/clients/src/main/java/org/apache/kafka/common/record/Records.java
@@ -16,10 +16,13 @@
  */
 package org.apache.kafka.common.record;
 
+import org.apache.kafka.common.utils.AbstractIterator;
+import org.apache.kafka.common.utils.Time;
+
 import java.io.IOException;
 import java.nio.channels.GatheringByteChannel;
+import java.util.Iterator;
 
-import org.apache.kafka.common.utils.Time;
 
 /**
  * Interface for accessing the records contained in a log. The log itself is represented as a sequence of record
@@ -28,20 +31,19 @@
  * For magic versions 1 and below, each batch consists of an 8 byte offset, a 4 byte record size, and a "shallow"
  * {@link Record record}. If the batch is not compressed, then each batch will have only the shallow record contained
  * inside it. If it is compressed, the batch contains "deep" records, which are packed into the value field of the
- * shallow record. To iterate over the shallow batches, use {@link #batches()}; for the deep records, use
- * {@link #records()}. Note that the deep iterator handles both compressed and non-compressed batches: if the batch is
- * not compressed, the shallow record is returned; otherwise, the shallow batch is decompressed and the deep records
- * are returned.
+ * shallow record. To iterate over the shallow batches, use {@link Records#batches()}; for the deep records, use
+ * {@link Records#records()}. Note that the deep iterator handles both compressed and non-compressed batches:
+ * if the batch is not compressed, the shallow record is returned; otherwise, the shallow batch is decompressed and the
+ * deep records are returned.
  *
  * For magic version 2, every batch contains 1 or more log record, regardless of compression. You can iterate
- * over the batches directly using {@link #batches()}. Records can be iterated either directly from an individual
- * batch or through {@link #records()}. Just as in previous versions, iterating over the records typically involves
+ * over the batches directly using {@link Records#batches()}. Records can be iterated either directly from an individual
+ * batch or through {@link Records#records()}. Just as in previous versions, iterating over the records typically involves
  * decompression and should therefore be used with caution.
  *
  * See {@link MemoryRecords} for the in-memory representation and {@link FileRecords} for the on-disk representation.
  */
-public interface Records {
-
+public interface Records extends BaseRecords {
     int OFFSET_OFFSET = 0;
     int OFFSET_LENGTH = 8;
     int SIZE_OFFSET = OFFSET_OFFSET + OFFSET_LENGTH;
@@ -54,12 +56,6 @@ public interface Records {
     int MAGIC_LENGTH = 1;
     int HEADER_SIZE_UP_TO_MAGIC = MAGIC_OFFSET + MAGIC_LENGTH;
 
-    /**
-     * The size of these records in bytes.
-     * @return The size in bytes of the records
-     */
-    int sizeInBytes();
-
     /**
      * Attempts to write the contents of this buffer to a channel.
      * @param channel The channel to write to
@@ -79,6 +75,13 @@ public interface Records {
      */
     Iterable batches();
 
+    /**
+     * Get an iterator over the record batches. This is similar to {@link #batches()} but returns an {@link AbstractIterator}
+     * instead of {@link Iterator}, so that clients can use methods like {@link AbstractIterator#peek() peek}.
+     * @return An iterator over the record batches of the log
+     */
+    AbstractIterator batchIterator();
+
     /**
      * Check whether all batches in this buffer have a certain magic value.
      * @param magic The magic value to check
@@ -99,7 +102,7 @@ public interface Records {
      * deep iteration since all of the deep records must also be converted to the desired format.
      * @param toMagic The magic value to convert to
      * @param firstOffset The starting offset for returned records. This only impacts some cases. See
-     *                    {@link AbstractRecords#downConvert(Iterable, byte, long, Time) for an explanation.
+     *                    {@link RecordsUtil#downConvert(Iterable, byte, long, Time)} for an explanation.
      * @param time instance used for reporting stats
      * @return A ConvertedRecords instance which may or may not contain the same instance in its records field.
      */
diff --git a/clients/src/main/java/org/apache/kafka/common/record/RecordsProcessingStats.java b/clients/src/main/java/org/apache/kafka/common/record/RecordsProcessingStats.java
deleted file mode 100644
index e104bc8189f7e..0000000000000
--- a/clients/src/main/java/org/apache/kafka/common/record/RecordsProcessingStats.java
+++ /dev/null
@@ -1,60 +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.record;
-
-public class RecordsProcessingStats {
-
-    public static final RecordsProcessingStats EMPTY = new RecordsProcessingStats(0L, 0, -1);
-
-    private final long temporaryMemoryBytes;
-    private final int numRecordsConverted;
-    private final long conversionTimeNanos;
-
-    public RecordsProcessingStats(long temporaryMemoryBytes, int numRecordsConverted, long conversionTimeNanos) {
-        this.temporaryMemoryBytes = temporaryMemoryBytes;
-        this.numRecordsConverted = numRecordsConverted;
-        this.conversionTimeNanos = conversionTimeNanos;
-    }
-
-    /**
-     * Returns the number of temporary memory bytes allocated to process the records.
-     * This size depends on whether the records need decompression and/or conversion:
-     * 
          
-     *   
        • Non compressed, no conversion: zero
          • 
-     *   
        • Non compressed, with conversion: size of the converted buffer
          • 
-     *   
        • Compressed, no conversion: size of the original buffer after decompression
          • 
-     *   
        • Compressed, with conversion: size of the original buffer after decompression + size of the converted buffer uncompressed
          • 
-     * 
        
-     */
-    public long temporaryMemoryBytes() {
-        return temporaryMemoryBytes;
-    }
-
-    public int numRecordsConverted() {
-        return numRecordsConverted;
-    }
-
-    public long conversionTimeNanos() {
-        return conversionTimeNanos;
-    }
-
-    @Override
-    public String toString() {
-        return String.format("RecordsProcessingStats(temporaryMemoryBytes=%d, numRecordsConverted=%d, conversionTimeNanos=%d)",
-                temporaryMemoryBytes, numRecordsConverted, conversionTimeNanos);
-    }
-}
diff --git a/clients/src/main/java/org/apache/kafka/common/record/RecordsSend.java b/clients/src/main/java/org/apache/kafka/common/record/RecordsSend.java
new file mode 100644
index 0000000000000..588dc4273aeac
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/record/RecordsSend.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 org.apache.kafka.common.record;
+
+import org.apache.kafka.common.network.Send;
+import org.apache.kafka.common.network.TransportLayers;
+
+import java.io.EOFException;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.channels.GatheringByteChannel;
+
+public abstract class RecordsSend implements Send {
+    private static final ByteBuffer EMPTY_BYTE_BUFFER = ByteBuffer.allocate(0);
+
+    private final T records;
+    private final int maxBytesToWrite;
+    private int remaining;
+    private boolean pending = false;
+
+    protected RecordsSend(T records, int maxBytesToWrite) {
+        this.records = records;
+        this.maxBytesToWrite = maxBytesToWrite;
+        this.remaining = maxBytesToWrite;
+    }
+
+    @Override
+    public boolean completed() {
+        return remaining <= 0 && !pending;
+    }
+
+    @Override
+    public final long writeTo(GatheringByteChannel channel) throws IOException {
+        long written = 0;
+
+        if (remaining > 0) {
+            written = writeTo(channel, size() - remaining, remaining);
+            if (written < 0)
+                throw new EOFException("Wrote negative bytes to channel. This shouldn't happen.");
+            remaining -= written;
+        }
+
+        pending = TransportLayers.hasPendingWrites(channel);
+        if (remaining <= 0 && pending)
+            channel.write(EMPTY_BYTE_BUFFER);
+
+        return written;
+    }
+
+    @Override
+    public long size() {
+        return maxBytesToWrite;
+    }
+
+    protected T records() {
+        return records;
+    }
+
+    /**
+     * Write records up to `remaining` bytes to `channel`. The implementation is allowed to be stateful. The contract
+     * from the caller is that the first invocation will be with `previouslyWritten` equal to 0, and `remaining` equal to
+     * the to maximum bytes we want to write the to `channel`. `previouslyWritten` and `remaining` will be adjusted
+     * appropriately for every subsequent invocation. See {@link #writeTo} for example expected usage.
+     * @param channel The channel to write to
+     * @param previouslyWritten Bytes written in previous calls to {@link #writeTo(GatheringByteChannel, long, int)}; 0 if being called for the first time
+     * @param remaining Number of bytes remaining to be written
+     * @return The number of bytes actually written
+     * @throws IOException For any IO errors
+     */
+    protected abstract long writeTo(GatheringByteChannel channel, long previouslyWritten, int remaining) throws IOException;
+}
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
new file mode 100644
index 0000000000000..423d1e1656f86
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/record/RecordsUtil.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.common.record;
+
+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;
+import java.util.List;
+
+public class RecordsUtil {
+    /**
+     * Down convert batches to the provided message format version. The first offset parameter is only relevant in the
+     * conversion from uncompressed v2 or higher to v1 or lower. The reason is that uncompressed records in v0 and v1
+     * are not batched (put another way, each batch always has 1 record).
+     *
+     * If a client requests records in v1 format starting from the middle of an uncompressed batch in v2 format, we
+     * need to drop records from the batch during the conversion. Some versions of librdkafka rely on this for
+     * correctness.
+     *
+     * The temporaryMemoryBytes computation assumes that the batches are not loaded into the heap
+     * (via classes like FileChannelRecordBatch) before this method is called. This is the case in the broker (we
+     * only load records into the heap when down converting), but it's not for the producer. However, down converting
+     * in the producer is very uncommon and the extra complexity to handle that case is not worth it.
+     */
+    protected static ConvertedRecords downConvert(Iterable batches, byte toMagic,
+                                                                 long firstOffset, Time time) {
+        // maintain the batch along with the decompressed records to avoid the need to decompress again
+        List recordBatchAndRecordsList = new ArrayList<>();
+        int totalSizeEstimate = 0;
+        long startNanos = time.nanoseconds();
+
+        for (RecordBatch batch : batches) {
+            if (toMagic < RecordBatch.MAGIC_VALUE_V2) {
+                if (batch.isControlBatch())
+                    continue;
+
+                if (batch.compressionType() == CompressionType.ZSTD)
+                    throw new UnsupportedCompressionTypeException("Down-conversion of zstandard-compressed batches " +
+                        "is not supported");
+            }
+
+            if (batch.magic() <= toMagic) {
+                totalSizeEstimate += batch.sizeInBytes();
+                recordBatchAndRecordsList.add(new RecordBatchAndRecords(batch, null, null));
+            } else {
+                List records = new ArrayList<>();
+                for (Record record : batch) {
+                    // See the method javadoc for an explanation
+                    if (toMagic > RecordBatch.MAGIC_VALUE_V1 || batch.isCompressed() || record.offset() >= firstOffset)
+                        records.add(record);
+                }
+                if (records.isEmpty())
+                    continue;
+                final long baseOffset;
+                if (batch.magic() >= RecordBatch.MAGIC_VALUE_V2 && toMagic >= RecordBatch.MAGIC_VALUE_V2)
+                    baseOffset = batch.baseOffset();
+                else
+                    baseOffset = records.get(0).offset();
+                totalSizeEstimate += AbstractRecords.estimateSizeInBytes(toMagic, baseOffset, batch.compressionType(), records);
+                recordBatchAndRecordsList.add(new RecordBatchAndRecords(batch, records, baseOffset));
+            }
+        }
+
+        ByteBuffer buffer = ByteBuffer.allocate(totalSizeEstimate);
+        long temporaryMemoryBytes = 0;
+        int numRecordsConverted = 0;
+
+        for (RecordBatchAndRecords recordBatchAndRecords : recordBatchAndRecordsList) {
+            temporaryMemoryBytes += recordBatchAndRecords.batch.sizeInBytes();
+            if (recordBatchAndRecords.batch.magic() <= toMagic) {
+                buffer = Utils.ensureCapacity(buffer, buffer.position() + recordBatchAndRecords.batch.sizeInBytes());
+                recordBatchAndRecords.batch.writeTo(buffer);
+            } else {
+                MemoryRecordsBuilder builder = convertRecordBatch(toMagic, buffer, recordBatchAndRecords);
+                buffer = builder.buffer();
+                temporaryMemoryBytes += builder.uncompressedBytesWritten();
+                numRecordsConverted += builder.numRecords();
+            }
+        }
+
+        buffer.flip();
+        RecordConversionStats stats = new RecordConversionStats(temporaryMemoryBytes, numRecordsConverted,
+                time.nanoseconds() - startNanos);
+        return new ConvertedRecords<>(MemoryRecords.readableRecords(buffer), stats);
+    }
+
+    /**
+     * Return a buffer containing the converted record batches. The returned buffer may not be the same as the received
+     * one (e.g. it may require expansion).
+     */
+    private static MemoryRecordsBuilder convertRecordBatch(byte magic, ByteBuffer buffer, RecordBatchAndRecords recordBatchAndRecords) {
+        RecordBatch batch = recordBatchAndRecords.batch;
+        final TimestampType timestampType = batch.timestampType();
+        long logAppendTime = timestampType == TimestampType.LOG_APPEND_TIME ? batch.maxTimestamp() : RecordBatch.NO_TIMESTAMP;
+
+        MemoryRecordsBuilder builder = MemoryRecords.builder(buffer, magic, batch.compressionType(),
+                timestampType, recordBatchAndRecords.baseOffset, logAppendTime);
+        for (Record record : recordBatchAndRecords.records) {
+            // Down-convert this record. Ignore headers when down-converting to V0 and V1 since they are not supported
+            if (magic > RecordBatch.MAGIC_VALUE_V1)
+                builder.append(record);
+            else
+                builder.appendWithOffset(record.offset(), record.timestamp(), record.key(), record.value());
+        }
+
+        builder.close();
+        return builder;
+    }
+
+
+    private static class RecordBatchAndRecords {
+        private final RecordBatch batch;
+        private final List records;
+        private final Long baseOffset;
+
+        private RecordBatchAndRecords(RecordBatch batch, List records, Long baseOffset) {
+            this.batch = batch;
+            this.records = records;
+            this.baseOffset = baseOffset;
+        }
+    }
+
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/record/SimpleRecord.java b/clients/src/main/java/org/apache/kafka/common/record/SimpleRecord.java
index fd361c417f3dd..9b8c87f79351f 100644
--- a/clients/src/main/java/org/apache/kafka/common/record/SimpleRecord.java
+++ b/clients/src/main/java/org/apache/kafka/common/record/SimpleRecord.java
@@ -61,6 +61,10 @@ public SimpleRecord(byte[] value) {
         this(RecordBatch.NO_TIMESTAMP, null, value);
     }
 
+    public SimpleRecord(ByteBuffer value) {
+        this(RecordBatch.NO_TIMESTAMP, null, value);
+    }
+
     public SimpleRecord(byte[] key, byte[] value) {
         this(RecordBatch.NO_TIMESTAMP, key, value);
     }
@@ -94,8 +98,8 @@ public boolean equals(Object o) {
 
         SimpleRecord that = (SimpleRecord) o;
         return timestamp == that.timestamp &&
-                (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);
     }
 
@@ -103,7 +107,7 @@ public boolean equals(Object o) {
     public int hashCode() {
         int result = key != null ? key.hashCode() : 0;
         result = 31 * result + (value != null ? value.hashCode() : 0);
-        result = 31 * result + (int) (timestamp ^ (timestamp >>> 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
new file mode 100644
index 0000000000000..dc4a1e21e8dd9
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/AbstractControlRequest.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.requests;
+
+import org.apache.kafka.common.protocol.ApiKeys;
+
+// 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;
+
+    public static abstract class Builder extends AbstractRequest.Builder {
+        protected final int controllerId;
+        protected final int controllerEpoch;
+        protected final long brokerEpoch;
+
+        protected Builder(ApiKeys api, short version, int controllerId, int controllerEpoch, long brokerEpoch) {
+            super(api, version);
+            this.controllerId = controllerId;
+            this.controllerEpoch = controllerEpoch;
+            this.brokerEpoch = brokerEpoch;
+        }
+
+    }
+
+    protected AbstractControlRequest(ApiKeys api, short version) {
+        super(api, version);
+    }
+
+    public abstract int controllerId();
+
+    public abstract int controllerEpoch();
+
+    public abstract long brokerEpoch();
+
+}
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 5a1c4f49926d9..28988895c074d 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
@@ -16,16 +16,19 @@
  */
 package org.apache.kafka.common.requests;
 
-import org.apache.kafka.common.network.NetworkSend;
+import org.apache.kafka.common.errors.UnsupportedVersionException;
 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.Struct;
+import org.apache.kafka.common.protocol.Message;
+import org.apache.kafka.common.protocol.MessageUtil;
+import org.apache.kafka.common.protocol.ObjectSerializationCache;
+import org.apache.kafka.common.protocol.SendBuilder;
 
 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;
@@ -75,9 +78,13 @@ public T build() {
     }
 
     private final short version;
+    private final ApiKeys apiKey;
 
-    public AbstractRequest(short version) {
+    public AbstractRequest(ApiKeys apiKey, short version) {
+        if (!apiKey.isVersionSupported(version))
+            throw new UnsupportedVersionException("The " + apiKey + " protocol does not support version " + version);
         this.version = version;
+        this.apiKey = apiKey;
     }
 
     /**
@@ -87,21 +94,28 @@ public short version() {
         return version;
     }
 
-    public Send toSend(String destination, RequestHeader header) {
-        return new NetworkSend(destination, serialize(header));
+    public ApiKeys apiKey() {
+        return apiKey;
     }
 
-    /**
-     * Use with care, typically {@link #toSend(String, RequestHeader)} should be used instead.
-     */
-    public ByteBuffer serialize(RequestHeader header) {
-        return serialize(header.toStruct(), toStruct());
+    public final Send toSend(RequestHeader header) {
+        return SendBuilder.buildRequestSend(header, data());
+    }
+
+    public abstract Message data();
+
+    // Visible for testing
+    public final ByteBuffer serialize() {
+        return MessageUtil.toByteBuffer(data(), version);
     }
 
-    protected abstract Struct toStruct();
+    // Visible for testing
+    final int sizeInBytes() {
+        return data().size(new ObjectSerializationCache(), version);
+    }
 
     public String toString(boolean verbose) {
-        return toStruct().toString();
+        return data().toString();
     }
 
     @Override
@@ -136,84 +150,131 @@ public Map errorCounts(Throwable e) {
     /**
      * Factory method for getting a request object based on ApiKey ID and a version
      */
-    public static AbstractRequest parseRequest(ApiKeys apiKey, short apiVersion, Struct struct) {
+    public static RequestAndSize parseRequest(ApiKeys apiKey, short apiVersion, ByteBuffer buffer) {
+        int bufferSize = buffer.remaining();
+        return new RequestAndSize(doParseRequest(apiKey, apiVersion, buffer), bufferSize);
+    }
+
+    private static AbstractRequest doParseRequest(ApiKeys apiKey, short apiVersion, ByteBuffer buffer) {
         switch (apiKey) {
             case PRODUCE:
-                return new ProduceRequest(struct, apiVersion);
+                return ProduceRequest.parse(buffer, apiVersion);
             case FETCH:
-                return new FetchRequest(struct, apiVersion);
+                return FetchRequest.parse(buffer, apiVersion);
             case LIST_OFFSETS:
-                return new ListOffsetRequest(struct, apiVersion);
+                return ListOffsetRequest.parse(buffer, apiVersion);
             case METADATA:
-                return new MetadataRequest(struct, apiVersion);
+                return MetadataRequest.parse(buffer, apiVersion);
             case OFFSET_COMMIT:
-                return new OffsetCommitRequest(struct, apiVersion);
+                return OffsetCommitRequest.parse(buffer, apiVersion);
             case OFFSET_FETCH:
-                return new OffsetFetchRequest(struct, apiVersion);
+                return OffsetFetchRequest.parse(buffer, apiVersion);
             case FIND_COORDINATOR:
-                return new FindCoordinatorRequest(struct, apiVersion);
+                return FindCoordinatorRequest.parse(buffer, apiVersion);
             case JOIN_GROUP:
-                return new JoinGroupRequest(struct, apiVersion);
+                return JoinGroupRequest.parse(buffer, apiVersion);
             case HEARTBEAT:
-                return new HeartbeatRequest(struct, apiVersion);
+                return HeartbeatRequest.parse(buffer, apiVersion);
             case LEAVE_GROUP:
-                return new LeaveGroupRequest(struct, apiVersion);
+                return LeaveGroupRequest.parse(buffer, apiVersion);
             case SYNC_GROUP:
-                return new SyncGroupRequest(struct, apiVersion);
+                return SyncGroupRequest.parse(buffer, apiVersion);
             case STOP_REPLICA:
-                return new StopReplicaRequest(struct, apiVersion);
+                return StopReplicaRequest.parse(buffer, apiVersion);
             case CONTROLLED_SHUTDOWN:
-                return new ControlledShutdownRequest(struct, apiVersion);
+                return ControlledShutdownRequest.parse(buffer, apiVersion);
             case UPDATE_METADATA:
-                return new UpdateMetadataRequest(struct, apiVersion);
+                return UpdateMetadataRequest.parse(buffer, apiVersion);
             case LEADER_AND_ISR:
-                return new LeaderAndIsrRequest(struct, apiVersion);
+                return LeaderAndIsrRequest.parse(buffer, apiVersion);
             case DESCRIBE_GROUPS:
-                return new DescribeGroupsRequest(struct, apiVersion);
+                return DescribeGroupsRequest.parse(buffer, apiVersion);
             case LIST_GROUPS:
-                return new ListGroupsRequest(struct, apiVersion);
+                return ListGroupsRequest.parse(buffer, apiVersion);
             case SASL_HANDSHAKE:
-                return new SaslHandshakeRequest(struct, apiVersion);
+                return SaslHandshakeRequest.parse(buffer, apiVersion);
             case API_VERSIONS:
-                return new ApiVersionsRequest(struct, apiVersion);
+                return ApiVersionsRequest.parse(buffer, apiVersion);
             case CREATE_TOPICS:
-                return new CreateTopicsRequest(struct, apiVersion);
+                return CreateTopicsRequest.parse(buffer, apiVersion);
             case DELETE_TOPICS:
-                return new DeleteTopicsRequest(struct, apiVersion);
+                return DeleteTopicsRequest.parse(buffer, apiVersion);
             case DELETE_RECORDS:
-                return new DeleteRecordsRequest(struct, apiVersion);
+                return DeleteRecordsRequest.parse(buffer, apiVersion);
             case INIT_PRODUCER_ID:
-                return new InitProducerIdRequest(struct, apiVersion);
+                return InitProducerIdRequest.parse(buffer, apiVersion);
             case OFFSET_FOR_LEADER_EPOCH:
-                return new OffsetsForLeaderEpochRequest(struct, apiVersion);
+                return OffsetsForLeaderEpochRequest.parse(buffer, apiVersion);
             case ADD_PARTITIONS_TO_TXN:
-                return new AddPartitionsToTxnRequest(struct, apiVersion);
+                return AddPartitionsToTxnRequest.parse(buffer, apiVersion);
             case ADD_OFFSETS_TO_TXN:
-                return new AddOffsetsToTxnRequest(struct, apiVersion);
+                return AddOffsetsToTxnRequest.parse(buffer, apiVersion);
             case END_TXN:
-                return new EndTxnRequest(struct, apiVersion);
+                return EndTxnRequest.parse(buffer, apiVersion);
             case WRITE_TXN_MARKERS:
-                return new WriteTxnMarkersRequest(struct, apiVersion);
+                return WriteTxnMarkersRequest.parse(buffer, apiVersion);
             case TXN_OFFSET_COMMIT:
-                return new TxnOffsetCommitRequest(struct, apiVersion);
+                return TxnOffsetCommitRequest.parse(buffer, apiVersion);
             case DESCRIBE_ACLS:
-                return new DescribeAclsRequest(struct, apiVersion);
+                return DescribeAclsRequest.parse(buffer, apiVersion);
             case CREATE_ACLS:
-                return new CreateAclsRequest(struct, apiVersion);
+                return CreateAclsRequest.parse(buffer, apiVersion);
             case DELETE_ACLS:
-                return new DeleteAclsRequest(struct, apiVersion);
+                return DeleteAclsRequest.parse(buffer, apiVersion);
             case DESCRIBE_CONFIGS:
-                return new DescribeConfigsRequest(struct, apiVersion);
+                return DescribeConfigsRequest.parse(buffer, apiVersion);
             case ALTER_CONFIGS:
-                return new AlterConfigsRequest(struct, apiVersion);
+                return AlterConfigsRequest.parse(buffer, apiVersion);
             case ALTER_REPLICA_LOG_DIRS:
-                return new AlterReplicaLogDirsRequest(struct, apiVersion);
+                return AlterReplicaLogDirsRequest.parse(buffer, apiVersion);
             case DESCRIBE_LOG_DIRS:
-                return new DescribeLogDirsRequest(struct, apiVersion);
+                return DescribeLogDirsRequest.parse(buffer, apiVersion);
             case SASL_AUTHENTICATE:
-                return new SaslAuthenticateRequest(struct, apiVersion);
+                return SaslAuthenticateRequest.parse(buffer, apiVersion);
             case CREATE_PARTITIONS:
-                return new CreatePartitionsRequest(struct, apiVersion);
+                return CreatePartitionsRequest.parse(buffer, apiVersion);
+            case CREATE_DELEGATION_TOKEN:
+                return CreateDelegationTokenRequest.parse(buffer, apiVersion);
+            case RENEW_DELEGATION_TOKEN:
+                return RenewDelegationTokenRequest.parse(buffer, apiVersion);
+            case EXPIRE_DELEGATION_TOKEN:
+                return ExpireDelegationTokenRequest.parse(buffer, apiVersion);
+            case DESCRIBE_DELEGATION_TOKEN:
+                return DescribeDelegationTokenRequest.parse(buffer, apiVersion);
+            case DELETE_GROUPS:
+                return DeleteGroupsRequest.parse(buffer, apiVersion);
+            case ELECT_LEADERS:
+                return ElectLeadersRequest.parse(buffer, apiVersion);
+            case INCREMENTAL_ALTER_CONFIGS:
+                return IncrementalAlterConfigsRequest.parse(buffer, apiVersion);
+            case ALTER_PARTITION_REASSIGNMENTS:
+                return AlterPartitionReassignmentsRequest.parse(buffer, apiVersion);
+            case LIST_PARTITION_REASSIGNMENTS:
+                return ListPartitionReassignmentsRequest.parse(buffer, apiVersion);
+            case OFFSET_DELETE:
+                return OffsetDeleteRequest.parse(buffer, apiVersion);
+            case DESCRIBE_CLIENT_QUOTAS:
+                return DescribeClientQuotasRequest.parse(buffer, apiVersion);
+            case ALTER_CLIENT_QUOTAS:
+                return AlterClientQuotasRequest.parse(buffer, apiVersion);
+            case DESCRIBE_USER_SCRAM_CREDENTIALS:
+                return DescribeUserScramCredentialsRequest.parse(buffer, apiVersion);
+            case ALTER_USER_SCRAM_CREDENTIALS:
+                return AlterUserScramCredentialsRequest.parse(buffer, apiVersion);
+            case VOTE:
+                return VoteRequest.parse(buffer, apiVersion);
+            case BEGIN_QUORUM_EPOCH:
+                return BeginQuorumEpochRequest.parse(buffer, apiVersion);
+            case END_QUORUM_EPOCH:
+                return EndQuorumEpochRequest.parse(buffer, apiVersion);
+            case DESCRIBE_QUORUM:
+                return DescribeQuorumRequest.parse(buffer, apiVersion);
+            case ALTER_ISR:
+                return AlterIsrRequest.parse(buffer, apiVersion);
+            case UPDATE_FEATURES:
+                return UpdateFeaturesRequest.parse(buffer, apiVersion);
+            case ENVELOPE:
+                return EnvelopeRequest.parse(buffer, 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..5d655bd0075dd 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,4 @@
  */
 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 6294af4cf7214..617ddbb1c132d 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
@@ -16,40 +16,64 @@
  */
 package org.apache.kafka.common.requests;
 
-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.Struct;
+import org.apache.kafka.common.protocol.Message;
+import org.apache.kafka.common.protocol.MessageUtil;
+import org.apache.kafka.common.protocol.SendBuilder;
 
 import java.nio.ByteBuffer;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
-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));
+    private final ApiKeys apiKey;
+
+    protected AbstractResponse(ApiKeys apiKey) {
+        this.apiKey = apiKey;
+    }
+
+    public final Send toSend(ResponseHeader header, short version) {
+        return SendBuilder.buildResponseSend(header, data(), version);
     }
 
     /**
-     * Visible for testing, typically {@link #toSend(String, ResponseHeader, short)} should be used instead.
+     * Serializes header and body without prefixing with size (unlike `toSend`, which does include a size prefix).
      */
-    public ByteBuffer serialize(short version, ResponseHeader responseHeader) {
-        return serialize(responseHeader.toStruct(), toStruct(version));
+    final ByteBuffer serializeWithHeader(ResponseHeader header, short version) {
+        return RequestUtils.serialize(header.data(), header.headerVersion(), data(), version);
     }
 
+    // Visible for testing
+    final ByteBuffer serialize(short version) {
+        return MessageUtil.toByteBuffer(data(), version);
+    }
+
+    /**
+     * The number of each type of error in the response, including {@link Errors#NONE} and top-level errors as well as
+     * more specifically scoped errors (such as topic or partition-level errors).
+     * @return A count of errors.
+     */
     public abstract Map errorCounts();
 
     protected Map errorCounts(Errors error) {
         return Collections.singletonMap(error, 1);
     }
 
-    protected Map errorCounts(Map errors) {
+    protected Map errorCounts(Stream errors) {
+        return errors.collect(Collectors.groupingBy(e -> e, Collectors.summingInt(e -> 1)));
+    }
+
+    protected Map errorCounts(Collection errors) {
         Map errorCounts = new HashMap<>();
-        for (Errors error : errors.values())
+        for (Errors error : errors)
             updateErrorCounts(errorCounts, error);
         return errorCounts;
     }
@@ -62,97 +86,174 @@ protected Map apiErrorCounts(Map errors) {
     }
 
     protected void updateErrorCounts(Map errorCounts, Errors error) {
-        Integer count = errorCounts.get(error);
-        errorCounts.put(error, count == null ? 1 : count + 1);
+        Integer count = errorCounts.getOrDefault(error, 0);
+        errorCounts.put(error, count + 1);
     }
 
-    protected abstract Struct toStruct(short version);
+    protected abstract Message data();
+
+    /**
+     * Parse a response from the provided buffer. The buffer is expected to hold both
+     * the {@link ResponseHeader} as well as the response payload.
+     */
+    public static AbstractResponse parseResponse(ByteBuffer buffer, RequestHeader requestHeader) {
+        ApiKeys apiKey = requestHeader.apiKey();
+        short apiVersion = requestHeader.apiVersion();
+
+        ResponseHeader responseHeader = ResponseHeader.parse(buffer, apiKey.responseHeaderVersion(apiVersion));
+
+        if (requestHeader.correlationId() != responseHeader.correlationId()) {
+            throw new CorrelationIdMismatchException("Correlation id for response ("
+                + responseHeader.correlationId() + ") does not match request ("
+                + requestHeader.correlationId() + "), request header: " + requestHeader,
+                requestHeader.correlationId(), responseHeader.correlationId());
+        }
+
+        return AbstractResponse.parseResponse(apiKey, buffer, apiVersion);
+    }
 
-    public static AbstractResponse parseResponse(ApiKeys apiKey, Struct struct) {
+    public static AbstractResponse parseResponse(ApiKeys apiKey, ByteBuffer responseBuffer, short version) {
         switch (apiKey) {
             case PRODUCE:
-                return new ProduceResponse(struct);
+                return ProduceResponse.parse(responseBuffer, version);
             case FETCH:
-                return new FetchResponse(struct);
+                return FetchResponse.parse(responseBuffer, version);
             case LIST_OFFSETS:
-                return new ListOffsetResponse(struct);
+                return ListOffsetResponse.parse(responseBuffer, version);
             case METADATA:
-                return new MetadataResponse(struct);
+                return MetadataResponse.parse(responseBuffer, version);
             case OFFSET_COMMIT:
-                return new OffsetCommitResponse(struct);
+                return OffsetCommitResponse.parse(responseBuffer, version);
             case OFFSET_FETCH:
-                return new OffsetFetchResponse(struct);
+                return OffsetFetchResponse.parse(responseBuffer, version);
             case FIND_COORDINATOR:
-                return new FindCoordinatorResponse(struct);
+                return FindCoordinatorResponse.parse(responseBuffer, version);
             case JOIN_GROUP:
-                return new JoinGroupResponse(struct);
+                return JoinGroupResponse.parse(responseBuffer, version);
             case HEARTBEAT:
-                return new HeartbeatResponse(struct);
+                return HeartbeatResponse.parse(responseBuffer, version);
             case LEAVE_GROUP:
-                return new LeaveGroupResponse(struct);
+                return LeaveGroupResponse.parse(responseBuffer, version);
             case SYNC_GROUP:
-                return new SyncGroupResponse(struct);
+                return SyncGroupResponse.parse(responseBuffer, version);
             case STOP_REPLICA:
-                return new StopReplicaResponse(struct);
+                return StopReplicaResponse.parse(responseBuffer, version);
             case CONTROLLED_SHUTDOWN:
-                return new ControlledShutdownResponse(struct);
+                return ControlledShutdownResponse.parse(responseBuffer, version);
             case UPDATE_METADATA:
-                return new UpdateMetadataResponse(struct);
+                return UpdateMetadataResponse.parse(responseBuffer, version);
             case LEADER_AND_ISR:
-                return new LeaderAndIsrResponse(struct);
+                return LeaderAndIsrResponse.parse(responseBuffer, version);
             case DESCRIBE_GROUPS:
-                return new DescribeGroupsResponse(struct);
+                return DescribeGroupsResponse.parse(responseBuffer, version);
             case LIST_GROUPS:
-                return new ListGroupsResponse(struct);
+                return ListGroupsResponse.parse(responseBuffer, version);
             case SASL_HANDSHAKE:
-                return new SaslHandshakeResponse(struct);
+                return SaslHandshakeResponse.parse(responseBuffer, version);
             case API_VERSIONS:
-                return new ApiVersionsResponse(struct);
+                return ApiVersionsResponse.parse(responseBuffer, version);
             case CREATE_TOPICS:
-                return new CreateTopicsResponse(struct);
+                return CreateTopicsResponse.parse(responseBuffer, version);
             case DELETE_TOPICS:
-                return new DeleteTopicsResponse(struct);
+                return DeleteTopicsResponse.parse(responseBuffer, version);
             case DELETE_RECORDS:
-                return new DeleteRecordsResponse(struct);
+                return DeleteRecordsResponse.parse(responseBuffer, version);
             case INIT_PRODUCER_ID:
-                return new InitProducerIdResponse(struct);
+                return InitProducerIdResponse.parse(responseBuffer, version);
             case OFFSET_FOR_LEADER_EPOCH:
-                return new OffsetsForLeaderEpochResponse(struct);
+                return OffsetsForLeaderEpochResponse.parse(responseBuffer, version);
             case ADD_PARTITIONS_TO_TXN:
-                return new AddPartitionsToTxnResponse(struct);
+                return AddPartitionsToTxnResponse.parse(responseBuffer, version);
             case ADD_OFFSETS_TO_TXN:
-                return new AddOffsetsToTxnResponse(struct);
+                return AddOffsetsToTxnResponse.parse(responseBuffer, version);
             case END_TXN:
-                return new EndTxnResponse(struct);
+                return EndTxnResponse.parse(responseBuffer, version);
             case WRITE_TXN_MARKERS:
-                return new WriteTxnMarkersResponse(struct);
+                return WriteTxnMarkersResponse.parse(responseBuffer, version);
             case TXN_OFFSET_COMMIT:
-                return new TxnOffsetCommitResponse(struct);
+                return TxnOffsetCommitResponse.parse(responseBuffer, version);
             case DESCRIBE_ACLS:
-                return new DescribeAclsResponse(struct);
+                return DescribeAclsResponse.parse(responseBuffer, version);
             case CREATE_ACLS:
-                return new CreateAclsResponse(struct);
+                return CreateAclsResponse.parse(responseBuffer, version);
             case DELETE_ACLS:
-                return new DeleteAclsResponse(struct);
+                return DeleteAclsResponse.parse(responseBuffer, version);
             case DESCRIBE_CONFIGS:
-                return new DescribeConfigsResponse(struct);
+                return DescribeConfigsResponse.parse(responseBuffer, version);
             case ALTER_CONFIGS:
-                return new AlterConfigsResponse(struct);
+                return AlterConfigsResponse.parse(responseBuffer, version);
             case ALTER_REPLICA_LOG_DIRS:
-                return new AlterReplicaLogDirsResponse(struct);
+                return AlterReplicaLogDirsResponse.parse(responseBuffer, version);
             case DESCRIBE_LOG_DIRS:
-                return new DescribeLogDirsResponse(struct);
+                return DescribeLogDirsResponse.parse(responseBuffer, version);
             case SASL_AUTHENTICATE:
-                return new SaslAuthenticateResponse(struct);
+                return SaslAuthenticateResponse.parse(responseBuffer, version);
             case CREATE_PARTITIONS:
-                return new CreatePartitionsResponse(struct);
+                return CreatePartitionsResponse.parse(responseBuffer, version);
+            case CREATE_DELEGATION_TOKEN:
+                return CreateDelegationTokenResponse.parse(responseBuffer, version);
+            case RENEW_DELEGATION_TOKEN:
+                return RenewDelegationTokenResponse.parse(responseBuffer, version);
+            case EXPIRE_DELEGATION_TOKEN:
+                return ExpireDelegationTokenResponse.parse(responseBuffer, version);
+            case DESCRIBE_DELEGATION_TOKEN:
+                return DescribeDelegationTokenResponse.parse(responseBuffer, version);
+            case DELETE_GROUPS:
+                return DeleteGroupsResponse.parse(responseBuffer, version);
+            case ELECT_LEADERS:
+                return ElectLeadersResponse.parse(responseBuffer, version);
+            case INCREMENTAL_ALTER_CONFIGS:
+                return IncrementalAlterConfigsResponse.parse(responseBuffer, version);
+            case ALTER_PARTITION_REASSIGNMENTS:
+                return AlterPartitionReassignmentsResponse.parse(responseBuffer, version);
+            case LIST_PARTITION_REASSIGNMENTS:
+                return ListPartitionReassignmentsResponse.parse(responseBuffer, version);
+            case OFFSET_DELETE:
+                return OffsetDeleteResponse.parse(responseBuffer, version);
+            case DESCRIBE_CLIENT_QUOTAS:
+                return DescribeClientQuotasResponse.parse(responseBuffer, version);
+            case ALTER_CLIENT_QUOTAS:
+                return AlterClientQuotasResponse.parse(responseBuffer, version);
+            case DESCRIBE_USER_SCRAM_CREDENTIALS:
+                return DescribeUserScramCredentialsResponse.parse(responseBuffer, version);
+            case ALTER_USER_SCRAM_CREDENTIALS:
+                return AlterUserScramCredentialsResponse.parse(responseBuffer, version);
+            case VOTE:
+                return VoteResponse.parse(responseBuffer, version);
+            case BEGIN_QUORUM_EPOCH:
+                return BeginQuorumEpochResponse.parse(responseBuffer, version);
+            case END_QUORUM_EPOCH:
+                return EndQuorumEpochResponse.parse(responseBuffer, version);
+            case DESCRIBE_QUORUM:
+                return DescribeQuorumResponse.parse(responseBuffer, version);
+            case ALTER_ISR:
+                return AlterIsrResponse.parse(responseBuffer, version);
+            case UPDATE_FEATURES:
+                return UpdateFeaturesResponse.parse(responseBuffer, version);
+            case ENVELOPE:
+                return EnvelopeResponse.parse(responseBuffer, version);
             default:
                 throw new AssertionError(String.format("ApiKey %s is not currently handled in `parseResponse`, the " +
                         "code should be updated to do so.", apiKey));
         }
     }
 
-    public String toString(short version) {
-        return toStruct(version).toString();
+    /**
+     * Returns whether or not client should throttle upon receiving a response of the specified version with a non-zero
+     * throttle time. Client-side throttling is needed when communicating with a newer version of broker which, on
+     * quota violation, sends out responses before throttling.
+     */
+    public boolean shouldClientThrottle(short version) {
+        return false;
+    }
+
+    public ApiKeys apiKey() {
+        return apiKey;
+    }
+
+    public abstract int throttleTimeMs();
+
+    public String toString() {
+        return data().toString();
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/AddOffsetsToTxnRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/AddOffsetsToTxnRequest.java
index f6a1722112d00..1e5f9862178bf 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/AddOffsetsToTxnRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/AddOffsetsToTxnRequest.java
@@ -16,119 +16,55 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.message.AddOffsetsToTxnRequestData;
+import org.apache.kafka.common.message.AddOffsetsToTxnResponseData;
 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.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.CommonFields.PRODUCER_EPOCH;
-import static org.apache.kafka.common.protocol.CommonFields.PRODUCER_ID;
-import static org.apache.kafka.common.protocol.CommonFields.TRANSACTIONAL_ID;
-
 public class AddOffsetsToTxnRequest extends AbstractRequest {
-    private static final Schema ADD_OFFSETS_TO_TXN_REQUEST_V0 = new Schema(
-            TRANSACTIONAL_ID,
-            PRODUCER_ID,
-            PRODUCER_EPOCH,
-            GROUP_ID);
 
-    public static Schema[] schemaVersions() {
-        return new Schema[]{ADD_OFFSETS_TO_TXN_REQUEST_V0};
-    }
+    private final AddOffsetsToTxnRequestData data;
 
     public static class Builder extends AbstractRequest.Builder {
-        private final String transactionalId;
-        private final long producerId;
-        private final short producerEpoch;
-        private final String consumerGroupId;
+        public AddOffsetsToTxnRequestData data;
 
-        public Builder(String transactionalId, long producerId, short producerEpoch, String consumerGroupId) {
+        public Builder(AddOffsetsToTxnRequestData data) {
             super(ApiKeys.ADD_OFFSETS_TO_TXN);
-            this.transactionalId = transactionalId;
-            this.producerId = producerId;
-            this.producerEpoch = producerEpoch;
-            this.consumerGroupId = consumerGroupId;
-        }
-
-        public String consumerGroupId() {
-            return consumerGroupId;
+            this.data = data;
         }
 
         @Override
         public AddOffsetsToTxnRequest build(short version) {
-            return new AddOffsetsToTxnRequest(version, transactionalId, producerId, producerEpoch, consumerGroupId);
+            return new AddOffsetsToTxnRequest(data, version);
         }
 
         @Override
         public String toString() {
-            StringBuilder bld = new StringBuilder();
-            bld.append("(type=AddOffsetsToTxnRequest").
-                    append(", transactionalId=").append(transactionalId).
-                    append(", producerId=").append(producerId).
-                    append(", producerEpoch=").append(producerEpoch).
-                    append(", consumerGroupId=").append(consumerGroupId).
-                    append(")");
-            return bld.toString();
+            return data.toString();
         }
     }
 
-    private final String transactionalId;
-    private final long producerId;
-    private final short producerEpoch;
-    private final String consumerGroupId;
-
-    private AddOffsetsToTxnRequest(short version, String transactionalId, long producerId, short producerEpoch, String consumerGroupId) {
-        super(version);
-        this.transactionalId = transactionalId;
-        this.producerId = producerId;
-        this.producerEpoch = producerEpoch;
-        this.consumerGroupId = consumerGroupId;
-    }
-
-    public AddOffsetsToTxnRequest(Struct struct, short version) {
-        super(version);
-        this.transactionalId = struct.get(TRANSACTIONAL_ID);
-        this.producerId = struct.get(PRODUCER_ID);
-        this.producerEpoch = struct.get(PRODUCER_EPOCH);
-        this.consumerGroupId = struct.get(GROUP_ID);
-    }
-
-    public String transactionalId() {
-        return transactionalId;
-    }
-
-    public long producerId() {
-        return producerId;
-    }
-
-    public short producerEpoch() {
-        return producerEpoch;
-    }
-
-    public String consumerGroupId() {
-        return consumerGroupId;
+    public AddOffsetsToTxnRequest(AddOffsetsToTxnRequestData data, short version) {
+        super(ApiKeys.ADD_OFFSETS_TO_TXN, version);
+        this.data = data;
     }
 
     @Override
-    protected Struct toStruct() {
-        Struct struct = new Struct(ApiKeys.ADD_OFFSETS_TO_TXN.requestSchema(version()));
-        struct.set(TRANSACTIONAL_ID, transactionalId);
-        struct.set(PRODUCER_ID, producerId);
-        struct.set(PRODUCER_EPOCH, producerEpoch);
-        struct.set(GROUP_ID, consumerGroupId);
-        return struct;
+    public AddOffsetsToTxnRequestData data() {
+        return data;
     }
 
     @Override
     public AddOffsetsToTxnResponse getErrorResponse(int throttleTimeMs, Throwable e) {
-        return new AddOffsetsToTxnResponse(throttleTimeMs, Errors.forException(e));
+        return new AddOffsetsToTxnResponse(new AddOffsetsToTxnResponseData()
+                                               .setErrorCode(Errors.forException(e).code())
+                                               .setThrottleTimeMs(throttleTimeMs));
     }
 
     public static AddOffsetsToTxnRequest parse(ByteBuffer buffer, short version) {
-        return new AddOffsetsToTxnRequest(ApiKeys.ADD_OFFSETS_TO_TXN.parseRequest(version, buffer), version);
+        return new AddOffsetsToTxnRequest(new AddOffsetsToTxnRequestData(new ByteBufferAccessor(buffer), version), version);
     }
-
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/AddOffsetsToTxnResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/AddOffsetsToTxnResponse.java
index 69b8ad4feaeb4..ce9a6cf7d6063 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/AddOffsetsToTxnResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/AddOffsetsToTxnResponse.java
@@ -16,80 +16,62 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.message.AddOffsetsToTxnResponseData;
 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.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;
-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#INVALID_PRODUCER_ID_MAPPING}
+ *   - {@link Errors#INVALID_PRODUCER_EPOCH} // for version <=1
+ *   - {@link Errors#PRODUCER_FENCED}
+ *   - {@link Errors#INVALID_TXN_STATE}
+ *   - {@link Errors#GROUP_AUTHORIZATION_FAILED}
+ *   - {@link Errors#TRANSACTIONAL_ID_AUTHORIZATION_FAILED}
+ */
 public class AddOffsetsToTxnResponse extends AbstractResponse {
-    private static final Schema ADD_OFFSETS_TO_TXN_RESPONSE_V0 = new Schema(
-            THROTTLE_TIME_MS,
-            ERROR_CODE);
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{ADD_OFFSETS_TO_TXN_RESPONSE_V0};
-    }
 
-    // Possible error codes:
-    //   NotCoordinator
-    //   CoordinatorNotAvailable
-    //   CoordinatorLoadInProgress
-    //   InvalidProducerIdMapping
-    //   InvalidProducerEpoch
-    //   InvalidTxnState
-    //   GroupAuthorizationFailed
-    //   TransactionalIdAuthorizationFailed
+    private final AddOffsetsToTxnResponseData data;
 
-    private final Errors error;
-    private final int throttleTimeMs;
-
-    public AddOffsetsToTxnResponse(int throttleTimeMs, Errors error) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.error = error;
-    }
-
-    public AddOffsetsToTxnResponse(Struct struct) {
-        this.throttleTimeMs = struct.get(THROTTLE_TIME_MS);
-        this.error = Errors.forCode(struct.get(ERROR_CODE));
+    public AddOffsetsToTxnResponse(AddOffsetsToTxnResponseData data) {
+        super(ApiKeys.ADD_OFFSETS_TO_TXN);
+        this.data = data;
     }
 
-    public int throttleTimeMs() {
-        return throttleTimeMs;
-    }
-
-    public Errors error() {
-        return error;
+    @Override
+    public Map errorCounts() {
+        return errorCounts(Errors.forCode(data.errorCode()));
     }
 
     @Override
-    public Map errorCounts() {
-        return errorCounts(error);
+    public int throttleTimeMs() {
+        return data.throttleTimeMs();
     }
 
     @Override
-    protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.ADD_OFFSETS_TO_TXN.responseSchema(version));
-        struct.set(THROTTLE_TIME_MS, throttleTimeMs);
-        struct.set(ERROR_CODE, error.code());
-        return struct;
+    public AddOffsetsToTxnResponseData data() {
+        return data;
     }
 
     public static AddOffsetsToTxnResponse parse(ByteBuffer buffer, short version) {
-        return new AddOffsetsToTxnResponse(ApiKeys.ADD_OFFSETS_TO_TXN.parseResponse(version, buffer));
+        return new AddOffsetsToTxnResponse(new AddOffsetsToTxnResponseData(new ByteBufferAccessor(buffer), version));
     }
 
     @Override
     public String toString() {
-        return "AddOffsetsToTxnResponse(" +
-                "error=" + error +
-                ", throttleTimeMs=" + throttleTimeMs +
-                ')';
+        return data.toString();
     }
 
+    @Override
+    public boolean shouldClientThrottle(short version) {
+        return version >= 1;
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/AddPartitionsToTxnRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/AddPartitionsToTxnRequest.java
index 0ca32beed1d25..1034c0f7adc55 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/AddPartitionsToTxnRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/AddPartitionsToTxnRequest.java
@@ -17,13 +17,12 @@
 package org.apache.kafka.common.requests;
 
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.message.AddPartitionsToTxnRequestData;
+import org.apache.kafka.common.message.AddPartitionsToTxnRequestData.AddPartitionsToTxnTopic;
+import org.apache.kafka.common.message.AddPartitionsToTxnRequestData.AddPartitionsToTxnTopicCollection;
 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.Struct;
-import org.apache.kafka.common.utils.CollectionUtils;
 
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
@@ -31,145 +30,102 @@
 import java.util.List;
 import java.util.Map;
 
-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 static org.apache.kafka.common.protocol.types.Type.INT32;
-
 public class AddPartitionsToTxnRequest extends AbstractRequest {
-    private static final String TOPICS_KEY_NAME = "topics";
-    private static final String PARTITIONS_KEY_NAME = "partitions";
-
-    private static final Schema ADD_PARTITIONS_TO_TXN_REQUEST_V0 = new Schema(
-            TRANSACTIONAL_ID,
-            PRODUCER_ID,
-            PRODUCER_EPOCH,
-            new Field(TOPICS_KEY_NAME, new ArrayOf(new Schema(
-                    TOPIC_NAME,
-                    new Field(PARTITIONS_KEY_NAME, new ArrayOf(INT32)))),
-                    "The partitions to add to the transaction."));
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{ADD_PARTITIONS_TO_TXN_REQUEST_V0};
-    }
+
+    private final AddPartitionsToTxnRequestData data;
+
+    private List cachedPartitions = null;
 
     public static class Builder extends AbstractRequest.Builder {
-        private final String transactionalId;
-        private final long producerId;
-        private final short producerEpoch;
-        private final List partitions;
+        public final AddPartitionsToTxnRequestData data;
+
+        public Builder(final AddPartitionsToTxnRequestData data) {
+            super(ApiKeys.ADD_PARTITIONS_TO_TXN);
+            this.data = data;
+        }
 
-        public Builder(String transactionalId, long producerId, short producerEpoch, List partitions) {
+        public Builder(final String transactionalId,
+                       final long producerId,
+                       final short producerEpoch,
+                       final List partitions) {
             super(ApiKeys.ADD_PARTITIONS_TO_TXN);
-            this.transactionalId = transactionalId;
-            this.producerId = producerId;
-            this.producerEpoch = producerEpoch;
-            this.partitions = partitions;
+
+            Map> partitionMap = new HashMap<>();
+            for (TopicPartition topicPartition : partitions) {
+                String topicName = topicPartition.topic();
+
+                partitionMap.compute(topicName, (key, subPartitions) -> {
+                    if (subPartitions == null) {
+                        subPartitions = new ArrayList<>();
+                    }
+                    subPartitions.add(topicPartition.partition());
+                    return subPartitions;
+                });
+            }
+
+            AddPartitionsToTxnTopicCollection topics = new AddPartitionsToTxnTopicCollection();
+            for (Map.Entry> partitionEntry : partitionMap.entrySet()) {
+                topics.add(new AddPartitionsToTxnTopic()
+                               .setName(partitionEntry.getKey())
+                               .setPartitions(partitionEntry.getValue()));
+            }
+
+            this.data = new AddPartitionsToTxnRequestData()
+                            .setTransactionalId(transactionalId)
+                            .setProducerId(producerId)
+                            .setProducerEpoch(producerEpoch)
+                            .setTopics(topics);
         }
 
         @Override
         public AddPartitionsToTxnRequest build(short version) {
-            return new AddPartitionsToTxnRequest(version, transactionalId, producerId, producerEpoch, partitions);
+            return new AddPartitionsToTxnRequest(data, version);
         }
 
-        public List partitions() {
+        static List getPartitions(AddPartitionsToTxnRequestData data) {
+            List partitions = new ArrayList<>();
+            for (AddPartitionsToTxnTopic topicCollection : data.topics()) {
+                for (Integer partition : topicCollection.partitions()) {
+                    partitions.add(new TopicPartition(topicCollection.name(), partition));
+                }
+            }
             return partitions;
         }
 
         @Override
         public String toString() {
-            StringBuilder bld = new StringBuilder();
-            bld.append("(type=AddPartitionsToTxnRequest").
-                    append(", transactionalId=").append(transactionalId).
-                    append(", producerId=").append(producerId).
-                    append(", producerEpoch=").append(producerEpoch).
-                    append(", partitions=").append(partitions).
-                    append(")");
-            return bld.toString();
+            return data.toString();
         }
     }
 
-    private final String transactionalId;
-    private final long producerId;
-    private final short producerEpoch;
-    private final List partitions;
-
-    private AddPartitionsToTxnRequest(short version, String transactionalId, long producerId, short producerEpoch,
-                                      List partitions) {
-        super(version);
-        this.transactionalId = transactionalId;
-        this.producerId = producerId;
-        this.producerEpoch = producerEpoch;
-        this.partitions = partitions;
-    }
-
-    public AddPartitionsToTxnRequest(Struct struct, short version) {
-        super(version);
-        this.transactionalId = struct.get(TRANSACTIONAL_ID);
-        this.producerId = struct.get(PRODUCER_ID);
-        this.producerEpoch = struct.get(PRODUCER_EPOCH);
-
-        List partitions = new ArrayList<>();
-        Object[] topicPartitionsArray = struct.getArray(TOPICS_KEY_NAME);
-        for (Object topicPartitionObj : topicPartitionsArray) {
-            Struct topicPartitionStruct = (Struct) topicPartitionObj;
-            String topic = topicPartitionStruct.get(TOPIC_NAME);
-            for (Object partitionObj : topicPartitionStruct.getArray(PARTITIONS_KEY_NAME)) {
-                partitions.add(new TopicPartition(topic, (Integer) partitionObj));
-            }
-        }
-        this.partitions = partitions;
-    }
-
-    public String transactionalId() {
-        return transactionalId;
-    }
-
-    public long producerId() {
-        return producerId;
-    }
-
-    public short producerEpoch() {
-        return producerEpoch;
+    public AddPartitionsToTxnRequest(final AddPartitionsToTxnRequestData data, short version) {
+        super(ApiKeys.ADD_PARTITIONS_TO_TXN, version);
+        this.data = data;
     }
 
     public List partitions() {
-        return partitions;
+        if (cachedPartitions != null) {
+            return cachedPartitions;
+        }
+        cachedPartitions = Builder.getPartitions(data);
+        return cachedPartitions;
     }
 
     @Override
-    protected Struct toStruct() {
-        Struct struct = new Struct(ApiKeys.ADD_PARTITIONS_TO_TXN.requestSchema(version()));
-        struct.set(TRANSACTIONAL_ID, transactionalId);
-        struct.set(PRODUCER_ID, producerId);
-        struct.set(PRODUCER_EPOCH, producerEpoch);
-
-        Map> mappedPartitions = CollectionUtils.groupDataByTopic(partitions);
-        Object[] partitionsArray = new Object[mappedPartitions.size()];
-        int i = 0;
-        for (Map.Entry> topicAndPartitions : mappedPartitions.entrySet()) {
-            Struct topicPartitionsStruct = struct.instance(TOPICS_KEY_NAME);
-            topicPartitionsStruct.set(TOPIC_NAME, topicAndPartitions.getKey());
-            topicPartitionsStruct.set(PARTITIONS_KEY_NAME, topicAndPartitions.getValue().toArray());
-            partitionsArray[i++] = topicPartitionsStruct;
-        }
-
-        struct.set(TOPICS_KEY_NAME, partitionsArray);
-        return struct;
+    public AddPartitionsToTxnRequestData data() {
+        return data;
     }
 
     @Override
     public AddPartitionsToTxnResponse getErrorResponse(int throttleTimeMs, Throwable e) {
         final HashMap errors = new HashMap<>();
-        for (TopicPartition partition : partitions) {
+        for (TopicPartition partition : partitions()) {
             errors.put(partition, Errors.forException(e));
         }
         return new AddPartitionsToTxnResponse(throttleTimeMs, errors);
     }
 
     public static AddPartitionsToTxnRequest parse(ByteBuffer buffer, short version) {
-        return new AddPartitionsToTxnRequest(ApiKeys.ADD_PARTITIONS_TO_TXN.parseRequest(version, buffer), version);
+        return new AddPartitionsToTxnRequest(new AddPartitionsToTxnRequestData(new ByteBufferAccessor(buffer), version), version);
     }
-
 }
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 44724499636e0..57b2a5a5d7c08 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
@@ -17,123 +17,121 @@
 package org.apache.kafka.common.requests;
 
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.message.AddPartitionsToTxnResponseData;
+import org.apache.kafka.common.message.AddPartitionsToTxnResponseData.AddPartitionsToTxnPartitionResult;
+import org.apache.kafka.common.message.AddPartitionsToTxnResponseData.AddPartitionsToTxnPartitionResultCollection;
+import org.apache.kafka.common.message.AddPartitionsToTxnResponseData.AddPartitionsToTxnTopicResult;
+import org.apache.kafka.common.message.AddPartitionsToTxnResponseData.AddPartitionsToTxnTopicResultCollection;
 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.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 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:
+ *
+ *   - {@link Errors#NOT_COORDINATOR}
+ *   - {@link Errors#COORDINATOR_NOT_AVAILABLE}
+ *   - {@link Errors#COORDINATOR_LOAD_IN_PROGRESS}
+ *   - {@link Errors#INVALID_TXN_STATE}
+ *   - {@link Errors#INVALID_PRODUCER_ID_MAPPING}
+ *   - {@link Errors#INVALID_PRODUCER_EPOCH} // for version <=1
+ *   - {@link Errors#PRODUCER_FENCED}
+ *   - {@link Errors#TOPIC_AUTHORIZATION_FAILED}
+ *   - {@link Errors#TRANSACTIONAL_ID_AUTHORIZATION_FAILED}
+ *   - {@link Errors#UNKNOWN_TOPIC_OR_PARTITION}
+ */
 public class AddPartitionsToTxnResponse extends AbstractResponse {
-    private static final String ERRORS_KEY_NAME = "errors";
-    private static final String PARTITION_ERRORS = "partition_errors";
-
-    private static final Schema ADD_PARTITIONS_TO_TXN_RESPONSE_V0 = new Schema(
-            THROTTLE_TIME_MS,
-            new Field(ERRORS_KEY_NAME, new ArrayOf(new Schema(
-                    TOPIC_NAME,
-                    new Field(PARTITION_ERRORS, new ArrayOf(new Schema(
-                            PARTITION_ID,
-                            ERROR_CODE)))))));
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{ADD_PARTITIONS_TO_TXN_RESPONSE_V0};
-    }
 
-    private final int throttleTimeMs;
-
-    // Possible error codes:
-    //   NotCoordinator
-    //   CoordinatorNotAvailable
-    //   CoordinatorLoadInProgress
-    //   InvalidTxnState
-    //   InvalidProducerIdMapping
-    //   TopicAuthorizationFailed
-    //   InvalidProducerEpoch
-    //   UnknownTopicOrPartition
-    //   TopicAuthorizationFailed
-    //   TransactionalIdAuthorizationFailed
-    private final Map errors;
+    private final AddPartitionsToTxnResponseData data;
 
-    public AddPartitionsToTxnResponse(int throttleTimeMs, Map errors) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.errors = errors;
+    private Map cachedErrorsMap = null;
+
+    public AddPartitionsToTxnResponse(AddPartitionsToTxnResponseData data) {
+        super(ApiKeys.ADD_PARTITIONS_TO_TXN);
+        this.data = data;
     }
 
-    public AddPartitionsToTxnResponse(Struct struct) {
-        this.throttleTimeMs = struct.get(THROTTLE_TIME_MS);
-        errors = new HashMap<>();
-        for (Object topic : struct.getArray(ERRORS_KEY_NAME)) {
-            Struct topicStruct = (Struct) topic;
-            final String topicName = topicStruct.get(TOPIC_NAME);
-            for (Object partition : topicStruct.getArray(PARTITION_ERRORS)) {
-                Struct partitionStruct = (Struct) partition;
-                TopicPartition topicPartition = new TopicPartition(topicName, partitionStruct.get(PARTITION_ID));
-                errors.put(topicPartition, Errors.forCode(partitionStruct.get(ERROR_CODE)));
-            }
+    public AddPartitionsToTxnResponse(int throttleTimeMs, Map errors) {
+        super(ApiKeys.ADD_PARTITIONS_TO_TXN);
+
+        Map resultMap = new HashMap<>();
+
+        for (Map.Entry entry : errors.entrySet()) {
+            TopicPartition topicPartition = entry.getKey();
+            String topicName = topicPartition.topic();
+
+            AddPartitionsToTxnPartitionResult partitionResult =
+                new AddPartitionsToTxnPartitionResult()
+                    .setErrorCode(entry.getValue().code())
+                    .setPartitionIndex(topicPartition.partition());
+
+            AddPartitionsToTxnPartitionResultCollection partitionResultCollection = resultMap.getOrDefault(
+                topicName, new AddPartitionsToTxnPartitionResultCollection()
+            );
+
+            partitionResultCollection.add(partitionResult);
+            resultMap.put(topicName, partitionResultCollection);
         }
+
+        AddPartitionsToTxnTopicResultCollection topicCollection = new AddPartitionsToTxnTopicResultCollection();
+        for (Map.Entry entry : resultMap.entrySet()) {
+            topicCollection.add(new AddPartitionsToTxnTopicResult()
+                                    .setName(entry.getKey())
+                                    .setResults(entry.getValue()));
+        }
+
+        this.data = new AddPartitionsToTxnResponseData()
+                        .setThrottleTimeMs(throttleTimeMs)
+                        .setResults(topicCollection);
     }
 
+    @Override
     public int throttleTimeMs() {
-        return throttleTimeMs;
+        return data.throttleTimeMs();
     }
 
     public Map errors() {
-        return errors;
+        if (cachedErrorsMap != null) {
+            return cachedErrorsMap;
+        }
+
+        cachedErrorsMap = new HashMap<>();
+
+        for (AddPartitionsToTxnTopicResult topicResult : this.data.results()) {
+            for (AddPartitionsToTxnPartitionResult partitionResult : topicResult.results()) {
+                cachedErrorsMap.put(new TopicPartition(
+                        topicResult.name(), partitionResult.partitionIndex()),
+                    Errors.forCode(partitionResult.errorCode()));
+            }
+        }
+        return cachedErrorsMap;
     }
 
     @Override
     public Map errorCounts() {
-        return errorCounts(errors);
+        return errorCounts(errors().values());
     }
 
     @Override
-    protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.ADD_PARTITIONS_TO_TXN.responseSchema(version));
-        struct.set(THROTTLE_TIME_MS, throttleTimeMs);
-
-        Map> errorsByTopic = CollectionUtils.groupDataByTopic(errors);
-        List topics = new ArrayList<>(errorsByTopic.size());
-        for (Map.Entry> entry : errorsByTopic.entrySet()) {
-            Struct topicErrorCodes = struct.instance(ERRORS_KEY_NAME);
-            topicErrorCodes.set(TOPIC_NAME, entry.getKey());
-            List partitionArray = new ArrayList<>();
-            for (Map.Entry partitionErrors : entry.getValue().entrySet()) {
-                final Struct partitionData = topicErrorCodes.instance(PARTITION_ERRORS)
-                        .set(PARTITION_ID, partitionErrors.getKey())
-                        .set(ERROR_CODE, partitionErrors.getValue().code());
-                partitionArray.add(partitionData);
-
-            }
-            topicErrorCodes.set(PARTITION_ERRORS, partitionArray.toArray());
-            topics.add(topicErrorCodes);
-        }
-        struct.set(ERRORS_KEY_NAME, topics.toArray());
-        return struct;
+    public AddPartitionsToTxnResponseData data() {
+        return data;
     }
 
     public static AddPartitionsToTxnResponse parse(ByteBuffer buffer, short version) {
-        return new AddPartitionsToTxnResponse(ApiKeys.ADD_PARTITIONS_TO_TXN.parseResponse(version, buffer));
+        return new AddPartitionsToTxnResponse(new AddPartitionsToTxnResponseData(new ByteBufferAccessor(buffer), version));
     }
 
     @Override
     public String toString() {
-        return "AddPartitionsToTxnResponse(" +
-                "errors=" + errors +
-                ", throttleTimeMs=" + throttleTimeMs +
-                ')';
+        return data.toString();
     }
 
+    @Override
+    public boolean shouldClientThrottle(short version) {
+        return version >= 1;
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/AlterClientQuotasRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/AlterClientQuotasRequest.java
new file mode 100644
index 0000000000000..d03c2671a290f
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/AlterClientQuotasRequest.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.requests;
+
+import org.apache.kafka.common.message.AlterClientQuotasRequestData;
+import org.apache.kafka.common.message.AlterClientQuotasRequestData.EntityData;
+import org.apache.kafka.common.message.AlterClientQuotasRequestData.EntryData;
+import org.apache.kafka.common.message.AlterClientQuotasRequestData.OpData;
+import org.apache.kafka.common.message.AlterClientQuotasResponseData;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.quota.ClientQuotaAlteration;
+import org.apache.kafka.common.quota.ClientQuotaEntity;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+public class AlterClientQuotasRequest extends AbstractRequest {
+
+    public static class Builder extends AbstractRequest.Builder {
+
+        private final AlterClientQuotasRequestData data;
+
+        public Builder(Collection entries, boolean validateOnly) {
+            super(ApiKeys.ALTER_CLIENT_QUOTAS);
+
+            List entryData = new ArrayList<>(entries.size());
+            for (ClientQuotaAlteration entry : entries) {
+                List entityData = new ArrayList<>(entry.entity().entries().size());
+                for (Map.Entry entityEntries : entry.entity().entries().entrySet()) {
+                    entityData.add(new EntityData()
+                            .setEntityType(entityEntries.getKey())
+                            .setEntityName(entityEntries.getValue()));
+                }
+
+                List opData = new ArrayList<>(entry.ops().size());
+                for (ClientQuotaAlteration.Op op : entry.ops()) {
+                    opData.add(new OpData()
+                            .setKey(op.key())
+                            .setValue(op.value() == null ? 0.0 : op.value())
+                            .setRemove(op.value() == null));
+                }
+
+                entryData.add(new EntryData()
+                        .setEntity(entityData)
+                        .setOps(opData));
+            }
+
+            this.data = new AlterClientQuotasRequestData()
+                    .setEntries(entryData)
+                    .setValidateOnly(validateOnly);
+        }
+
+        @Override
+        public AlterClientQuotasRequest build(short version) {
+            return new AlterClientQuotasRequest(data, version);
+        }
+
+        @Override
+        public String toString() {
+            return data.toString();
+        }
+    }
+
+    private final AlterClientQuotasRequestData data;
+
+    public AlterClientQuotasRequest(AlterClientQuotasRequestData data, short version) {
+        super(ApiKeys.ALTER_CLIENT_QUOTAS, version);
+        this.data = data;
+    }
+
+    public List entries() {
+        List entries = new ArrayList<>(data.entries().size());
+        for (EntryData entryData : data.entries()) {
+            Map entity = new HashMap<>(entryData.entity().size());
+            for (EntityData entityData : entryData.entity()) {
+                entity.put(entityData.entityType(), entityData.entityName());
+            }
+
+            List ops = new ArrayList<>(entryData.ops().size());
+            for (OpData opData : entryData.ops()) {
+                Double value = opData.remove() ? null : opData.value();
+                ops.add(new ClientQuotaAlteration.Op(opData.key(), value));
+            }
+
+            entries.add(new ClientQuotaAlteration(new ClientQuotaEntity(entity), ops));
+        }
+        return entries;
+    }
+
+    public boolean validateOnly() {
+        return data.validateOnly();
+    }
+
+    @Override
+    public AlterClientQuotasRequestData data() {
+        return data;
+    }
+
+    @Override
+    public AlterClientQuotasResponse getErrorResponse(int throttleTimeMs, Throwable e) {
+        List responseEntries = new ArrayList<>();
+        for (EntryData entryData : data.entries()) {
+            List responseEntities = new ArrayList<>();
+            for (EntityData entityData : entryData.entity()) {
+                responseEntities.add(new AlterClientQuotasResponseData.EntityData()
+                    .setEntityType(entityData.entityType())
+                    .setEntityName(entityData.entityName()));
+            }
+            responseEntries.add(new AlterClientQuotasResponseData.EntryData().setEntity(responseEntities));
+        }
+        AlterClientQuotasResponseData responseData = new AlterClientQuotasResponseData()
+                .setThrottleTimeMs(throttleTimeMs)
+                .setEntries(responseEntries);
+        return new AlterClientQuotasResponse(responseData);
+    }
+
+    public static AlterClientQuotasRequest parse(ByteBuffer buffer, short version) {
+        return new AlterClientQuotasRequest(new AlterClientQuotasRequestData(new ByteBufferAccessor(buffer), version), version);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/AlterClientQuotasResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/AlterClientQuotasResponse.java
new file mode 100644
index 0000000000000..fcacc5d95ef07
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/AlterClientQuotasResponse.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.requests;
+
+import org.apache.kafka.common.internals.KafkaFutureImpl;
+import org.apache.kafka.common.message.AlterClientQuotasResponseData;
+import org.apache.kafka.common.message.AlterClientQuotasResponseData.EntityData;
+import org.apache.kafka.common.message.AlterClientQuotasResponseData.EntryData;
+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.quota.ClientQuotaEntity;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+public class AlterClientQuotasResponse extends AbstractResponse {
+
+    private final AlterClientQuotasResponseData data;
+
+    public AlterClientQuotasResponse(AlterClientQuotasResponseData data) {
+        super(ApiKeys.ALTER_CLIENT_QUOTAS);
+        this.data = data;
+    }
+
+    public void complete(Map> futures) {
+        for (EntryData entryData : data.entries()) {
+            Map entityEntries = new HashMap<>(entryData.entity().size());
+            for (EntityData entityData : entryData.entity()) {
+                entityEntries.put(entityData.entityType(), entityData.entityName());
+            }
+            ClientQuotaEntity entity = new ClientQuotaEntity(entityEntries);
+
+            KafkaFutureImpl future = futures.get(entity);
+            if (future == null) {
+                throw new IllegalArgumentException("Future map must contain entity " + entity);
+            }
+
+            Errors error = Errors.forCode(entryData.errorCode());
+            if (error == Errors.NONE) {
+                future.complete(null);
+            } else {
+                future.completeExceptionally(error.exception(entryData.errorMessage()));
+            }
+        }
+    }
+
+    @Override
+    public int throttleTimeMs() {
+        return data.throttleTimeMs();
+    }
+
+    @Override
+    public Map errorCounts() {
+        Map counts = new HashMap<>();
+        data.entries().forEach(entry ->
+            updateErrorCounts(counts, Errors.forCode(entry.errorCode()))
+        );
+        return counts;
+    }
+
+    @Override
+    public AlterClientQuotasResponseData data() {
+        return data;
+    }
+
+    private static List toEntityData(ClientQuotaEntity entity) {
+        List entityData = new ArrayList<>(entity.entries().size());
+        for (Map.Entry entry : entity.entries().entrySet()) {
+            entityData.add(new AlterClientQuotasResponseData.EntityData()
+                    .setEntityType(entry.getKey())
+                    .setEntityName(entry.getValue()));
+        }
+        return entityData;
+    }
+
+    public static AlterClientQuotasResponse parse(ByteBuffer buffer, short version) {
+        return new AlterClientQuotasResponse(new AlterClientQuotasResponseData(new ByteBufferAccessor(buffer), version));
+    }
+
+    public static AlterClientQuotasResponse fromQuotaEntities(Map result, int throttleTimeMs) {
+        List entries = new ArrayList<>(result.size());
+        for (Map.Entry entry : result.entrySet()) {
+            ApiError e = entry.getValue();
+            entries.add(new EntryData()
+                    .setErrorCode(e.error().code())
+                    .setErrorMessage(e.message())
+                    .setEntity(toEntityData(entry.getKey())));
+        }
+
+        return new AlterClientQuotasResponse(new AlterClientQuotasResponseData()
+            .setThrottleTimeMs(throttleTimeMs)
+            .setEntries(entries));
+    }
+
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/AlterConfigsRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/AlterConfigsRequest.java
index 14b39ae9daceb..b4d35d52ae35f 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/AlterConfigsRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/AlterConfigsRequest.java
@@ -17,58 +17,25 @@
 
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.config.ConfigResource;
+import org.apache.kafka.common.message.AlterConfigsRequestData;
+import org.apache.kafka.common.message.AlterConfigsResponseData;
 import org.apache.kafka.common.protocol.ApiKeys;
-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.protocol.ByteBufferAccessor;
 
 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.types.Type.BOOLEAN;
-import static org.apache.kafka.common.protocol.types.Type.INT8;
-import static org.apache.kafka.common.protocol.types.Type.NULLABLE_STRING;
-import static org.apache.kafka.common.protocol.types.Type.STRING;
+import java.util.Objects;
+import java.util.stream.Collectors;
 
 public class AlterConfigsRequest extends AbstractRequest {
 
-    private static final String RESOURCES_KEY_NAME = "resources";
-    private static final String RESOURCE_TYPE_KEY_NAME = "resource_type";
-    private static final String RESOURCE_NAME_KEY_NAME = "resource_name";
-    private static final String VALIDATE_ONLY_KEY_NAME = "validate_only";
-
-    private static final String CONFIG_ENTRIES_KEY_NAME = "config_entries";
-    private static final String CONFIG_NAME = "config_name";
-    private static final String CONFIG_VALUE = "config_value";
-
-    private static final Schema CONFIG_ENTRY = new Schema(
-            new Field(CONFIG_NAME, STRING, "Configuration name"),
-            new Field(CONFIG_VALUE, NULLABLE_STRING, "Configuration value"));
-
-    private static final Schema ALTER_CONFIGS_REQUEST_RESOURCE_V0 = new Schema(
-            new Field(RESOURCE_TYPE_KEY_NAME, INT8),
-            new Field(RESOURCE_NAME_KEY_NAME, STRING),
-            new Field(CONFIG_ENTRIES_KEY_NAME, new ArrayOf(CONFIG_ENTRY)));
-
-    private static final Schema ALTER_CONFIGS_REQUEST_V0 = new Schema(
-            new Field(RESOURCES_KEY_NAME, new ArrayOf(ALTER_CONFIGS_REQUEST_RESOURCE_V0),
-                    "An array of resources to update with the provided configs."),
-            new Field(VALIDATE_ONLY_KEY_NAME, BOOLEAN));
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {ALTER_CONFIGS_REQUEST_V0};
-    }
-
     public static class Config {
         private final Collection entries;
 
         public Config(Collection entries) {
-            this.entries = entries;
+            this.entries = Objects.requireNonNull(entries, "entries");
         }
 
         public Collection entries() {
@@ -81,8 +48,8 @@ public static class ConfigEntry {
         private final String value;
 
         public ConfigEntry(String name, String value) {
-            this.name = name;
-            this.value = value;
+            this.name = Objects.requireNonNull(name, "name");
+            this.value = Objects.requireNonNull(value, "value");
         }
 
         public String name() {
@@ -95,111 +62,77 @@ public String value() {
 
     }
 
-    public static class Builder extends AbstractRequest.Builder {
+    public static class Builder extends AbstractRequest.Builder {
 
-        private final Map configs;
-        private final boolean validateOnly;
+        private final AlterConfigsRequestData data = new AlterConfigsRequestData();
 
-        public Builder(Map configs, boolean validateOnly) {
+        public Builder(Map configs, boolean validateOnly) {
             super(ApiKeys.ALTER_CONFIGS);
-            this.configs = configs;
-            this.validateOnly = validateOnly;
+            Objects.requireNonNull(configs, "configs");
+            for (Map.Entry entry : configs.entrySet()) {
+                AlterConfigsRequestData.AlterConfigsResource resource =
+                    new AlterConfigsRequestData.AlterConfigsResource()
+                        .setResourceName(entry.getKey().name())
+                        .setResourceType(entry.getKey().type().id());
+                for (ConfigEntry x : entry.getValue().entries) {
+                    resource.configs().add(new AlterConfigsRequestData.AlterableConfig()
+                                               .setName(x.name())
+                                               .setValue(x.value()));
+                }
+                this.data.resources().add(resource);
+            }
+            this.data.setValidateOnly(validateOnly);
         }
 
         @Override
         public AlterConfigsRequest build(short version) {
-            return new AlterConfigsRequest(version, configs, validateOnly);
+            return new AlterConfigsRequest(data, version);
         }
     }
 
-    private final Map configs;
-    private final boolean validateOnly;
+    private final AlterConfigsRequestData data;
 
-    public AlterConfigsRequest(short version, Map configs, boolean validateOnly) {
-        super(version);
-        this.configs = configs;
-        this.validateOnly = validateOnly;
+    public AlterConfigsRequest(AlterConfigsRequestData data, short version) {
+        super(ApiKeys.ALTER_CONFIGS, version);
+        this.data = data;
     }
 
-    public AlterConfigsRequest(Struct struct, short version) {
-        super(version);
-        validateOnly = struct.getBoolean(VALIDATE_ONLY_KEY_NAME);
-        Object[] resourcesArray = struct.getArray(RESOURCES_KEY_NAME);
-        configs = new HashMap<>(resourcesArray.length);
-        for (Object resourcesObj : resourcesArray) {
-            Struct resourcesStruct = (Struct) resourcesObj;
-
-            ResourceType resourceType = ResourceType.forId(resourcesStruct.getByte(RESOURCE_TYPE_KEY_NAME));
-            String resourceName = resourcesStruct.getString(RESOURCE_NAME_KEY_NAME);
-            Resource resource = new Resource(resourceType, resourceName);
-
-            Object[] configEntriesArray = resourcesStruct.getArray(CONFIG_ENTRIES_KEY_NAME);
-            List configEntries = new ArrayList<>(configEntriesArray.length);
-            for (Object configEntriesObj: configEntriesArray) {
-                Struct configEntriesStruct = (Struct) configEntriesObj;
-                String configName = configEntriesStruct.getString(CONFIG_NAME);
-                String configValue = configEntriesStruct.getString(CONFIG_VALUE);
-                configEntries.add(new ConfigEntry(configName, configValue));
-            }
-            Config config = new Config(configEntries);
-            configs.put(resource, config);
-        }
-    }
-
-    public Map configs() {
-        return configs;
+    public Map configs() {
+        return data.resources().stream().collect(Collectors.toMap(
+            resource -> new ConfigResource(
+                    ConfigResource.Type.forId(resource.resourceType()),
+                    resource.resourceName()),
+            resource -> new Config(resource.configs().stream()
+                    .map(entry -> new ConfigEntry(entry.name(), entry.value()))
+                    .collect(Collectors.toList()))));
     }
 
     public boolean validateOnly() {
-        return validateOnly;
+        return data.validateOnly();
     }
 
     @Override
-    protected Struct toStruct() {
-        Struct struct = new Struct(ApiKeys.ALTER_CONFIGS.requestSchema(version()));
-        struct.set(VALIDATE_ONLY_KEY_NAME, validateOnly);
-        List resourceStructs = new ArrayList<>(configs.size());
-        for (Map.Entry entry : configs.entrySet()) {
-            Struct resourceStruct = struct.instance(RESOURCES_KEY_NAME);
-
-            Resource resource = entry.getKey();
-            resourceStruct.set(RESOURCE_TYPE_KEY_NAME, resource.type().id());
-            resourceStruct.set(RESOURCE_NAME_KEY_NAME, resource.name());
-
-            Config config = entry.getValue();
-            List configEntryStructs = new ArrayList<>(config.entries.size());
-            for (ConfigEntry configEntry : config.entries) {
-                Struct configEntriesStruct = resourceStruct.instance(CONFIG_ENTRIES_KEY_NAME);
-                configEntriesStruct.set(CONFIG_NAME, configEntry.name);
-                configEntriesStruct.set(CONFIG_VALUE, configEntry.value);
-                configEntryStructs.add(configEntriesStruct);
-            }
-            resourceStruct.set(CONFIG_ENTRIES_KEY_NAME, configEntryStructs.toArray(new Struct[0]));
-
-            resourceStructs.add(resourceStruct);
-        }
-        struct.set(RESOURCES_KEY_NAME, resourceStructs.toArray(new Struct[0]));
-        return struct;
+    public AlterConfigsRequestData data() {
+        return data;
     }
 
     @Override
     public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
-        short version = version();
-        switch (version) {
-            case 0:
-                ApiError error = ApiError.fromThrowable(e);
-                Map errors = new HashMap<>(configs.size());
-                for (Resource resource : configs.keySet())
-                    errors.put(resource, error);
-                return new AlterConfigsResponse(throttleTimeMs, errors);
-            default:
-                throw new IllegalArgumentException(String.format("Version %d is not valid. Valid versions for %s are 0 to %d",
-                        version, this.getClass().getSimpleName(), ApiKeys.ALTER_CONFIGS.latestVersion()));
+        ApiError error = ApiError.fromThrowable(e);
+        AlterConfigsResponseData data = new AlterConfigsResponseData()
+                .setThrottleTimeMs(throttleTimeMs);
+        for (AlterConfigsRequestData.AlterConfigsResource resource : this.data.resources()) {
+            data.responses().add(new AlterConfigsResponseData.AlterConfigsResourceResponse()
+                    .setResourceType(resource.resourceType())
+                    .setResourceName(resource.resourceName())
+                    .setErrorMessage(error.message())
+                    .setErrorCode(error.error().code()));
         }
+        return new AlterConfigsResponse(data);
+
     }
 
     public static AlterConfigsRequest parse(ByteBuffer buffer, short version) {
-        return new AlterConfigsRequest(ApiKeys.ALTER_CONFIGS.parseRequest(version, buffer), version);
+        return new AlterConfigsRequest(new AlterConfigsRequestData(new ByteBufferAccessor(buffer), version), version);
     }
-
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/AlterConfigsResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/AlterConfigsResponse.java
index f292ef6731b0e..1115f06ee80a9 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/AlterConfigsResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/AlterConfigsResponse.java
@@ -17,99 +17,55 @@
 
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.config.ConfigResource;
+import org.apache.kafka.common.message.AlterConfigsResponseData;
 import org.apache.kafka.common.protocol.ApiKeys;
-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.ByteBufferAccessor;
 import org.apache.kafka.common.protocol.Errors;
-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.ERROR_MESSAGE;
-import static org.apache.kafka.common.protocol.CommonFields.THROTTLE_TIME_MS;
-import static org.apache.kafka.common.protocol.types.Type.INT8;
-import static org.apache.kafka.common.protocol.types.Type.STRING;
+import java.util.stream.Collectors;
 
 public class AlterConfigsResponse extends AbstractResponse {
 
-    private static final String RESOURCES_KEY_NAME = "resources";
-    private static final String RESOURCE_TYPE_KEY_NAME = "resource_type";
-    private static final String RESOURCE_NAME_KEY_NAME = "resource_name";
-
-    private static final Schema ALTER_CONFIGS_RESPONSE_ENTITY_V0 = new Schema(
-            ERROR_CODE,
-            ERROR_MESSAGE,
-            new Field(RESOURCE_TYPE_KEY_NAME, INT8),
-            new Field(RESOURCE_NAME_KEY_NAME, STRING));
-
-    private static final Schema ALTER_CONFIGS_RESPONSE_V0 = new Schema(
-            THROTTLE_TIME_MS,
-            new Field(RESOURCES_KEY_NAME, new ArrayOf(ALTER_CONFIGS_RESPONSE_ENTITY_V0)));
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{ALTER_CONFIGS_RESPONSE_V0};
-    }
-
-    private final int throttleTimeMs;
-    private final Map errors;
+    private final AlterConfigsResponseData data;
 
-    public AlterConfigsResponse(int throttleTimeMs, Map errors) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.errors = errors;
-
-    }
-
-    public AlterConfigsResponse(Struct struct) {
-        throttleTimeMs = struct.get(THROTTLE_TIME_MS);
-        Object[] resourcesArray = struct.getArray(RESOURCES_KEY_NAME);
-        errors = new HashMap<>(resourcesArray.length);
-        for (Object resourceObj : resourcesArray) {
-            Struct resourceStruct = (Struct) resourceObj;
-            ApiError error = new ApiError(resourceStruct);
-            ResourceType resourceType = ResourceType.forId(resourceStruct.getByte(RESOURCE_TYPE_KEY_NAME));
-            String resourceName = resourceStruct.getString(RESOURCE_NAME_KEY_NAME);
-            errors.put(new Resource(resourceType, resourceName), error);
-        }
+    public AlterConfigsResponse(AlterConfigsResponseData data) {
+        super(ApiKeys.ALTER_CONFIGS);
+        this.data = data;
     }
 
-    public Map errors() {
-        return errors;
+    public Map errors() {
+        return data.responses().stream().collect(Collectors.toMap(
+            response -> new ConfigResource(
+                    ConfigResource.Type.forId(response.resourceType()),
+                    response.resourceName()),
+            response -> new ApiError(Errors.forCode(response.errorCode()), response.errorMessage())
+        ));
     }
 
     @Override
     public Map errorCounts() {
-        return apiErrorCounts(errors);
+        return apiErrorCounts(errors());
     }
 
+    @Override
     public int throttleTimeMs() {
-        return throttleTimeMs;
+        return data.throttleTimeMs();
     }
 
     @Override
-    protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.ALTER_CONFIGS.responseSchema(version));
-        struct.set(THROTTLE_TIME_MS, throttleTimeMs);
-        List resourceStructs = new ArrayList<>(errors.size());
-        for (Map.Entry entry : errors.entrySet()) {
-            Struct resourceStruct = struct.instance(RESOURCES_KEY_NAME);
-            Resource resource = entry.getKey();
-            entry.getValue().write(resourceStruct);
-            resourceStruct.set(RESOURCE_TYPE_KEY_NAME, resource.type().id());
-            resourceStruct.set(RESOURCE_NAME_KEY_NAME, resource.name());
-            resourceStructs.add(resourceStruct);
-        }
-        struct.set(RESOURCES_KEY_NAME, resourceStructs.toArray(new Struct[0]));
-        return struct;
+    public AlterConfigsResponseData data() {
+        return data;
     }
 
     public static AlterConfigsResponse parse(ByteBuffer buffer, short version) {
-        return new AlterConfigsResponse(ApiKeys.ALTER_CONFIGS.parseResponse(version, buffer));
+        return new AlterConfigsResponse(new AlterConfigsResponseData(new ByteBufferAccessor(buffer), version));
     }
 
+    @Override
+    public boolean shouldClientThrottle(short version) {
+        return version >= 1;
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/AlterIsrRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/AlterIsrRequest.java
new file mode 100644
index 0000000000000..516c2ce76aa85
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/AlterIsrRequest.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.AlterIsrRequestData;
+import org.apache.kafka.common.message.AlterIsrResponseData;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.Errors;
+
+import java.nio.ByteBuffer;
+
+public class AlterIsrRequest extends AbstractRequest {
+
+    private final AlterIsrRequestData data;
+
+    public AlterIsrRequest(AlterIsrRequestData data, short apiVersion) {
+        super(ApiKeys.ALTER_ISR, apiVersion);
+        this.data = data;
+    }
+
+    @Override
+    public AlterIsrRequestData data() {
+        return data;
+    }
+
+    /**
+     * Get an error response for a request with specified throttle time in the response if applicable
+     */
+    @Override
+    public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
+        return new AlterIsrResponse(new AlterIsrResponseData()
+                .setThrottleTimeMs(throttleTimeMs)
+                .setErrorCode(Errors.forException(e).code()));
+    }
+
+    public static AlterIsrRequest parse(ByteBuffer buffer, short version) {
+        return new AlterIsrRequest(new AlterIsrRequestData(new ByteBufferAccessor(buffer), version), version);
+    }
+
+    public static class Builder extends AbstractRequest.Builder {
+
+        private final AlterIsrRequestData data;
+
+        public Builder(AlterIsrRequestData data) {
+            super(ApiKeys.ALTER_ISR);
+            this.data = data;
+        }
+
+        @Override
+        public AlterIsrRequest build(short version) {
+            return new AlterIsrRequest(data, version);
+        }
+
+        @Override
+        public String toString() {
+            return data.toString();
+        }
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/AlterIsrResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/AlterIsrResponse.java
new file mode 100644
index 0000000000000..c3106ed94cbde
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/AlterIsrResponse.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.requests;
+
+import org.apache.kafka.common.message.AlterIsrResponseData;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.Errors;
+
+import java.nio.ByteBuffer;
+import java.util.HashMap;
+import java.util.Map;
+
+public class AlterIsrResponse extends AbstractResponse {
+
+    private final AlterIsrResponseData data;
+
+    public AlterIsrResponse(AlterIsrResponseData data) {
+        super(ApiKeys.ALTER_ISR);
+        this.data = data;
+    }
+
+    @Override
+    public AlterIsrResponseData data() {
+        return data;
+    }
+
+    @Override
+    public Map errorCounts() {
+        Map counts = new HashMap<>();
+        updateErrorCounts(counts, Errors.forCode(data.errorCode()));
+        data.topics().forEach(topicResponse -> topicResponse.partitions().forEach(partitionResponse -> {
+            updateErrorCounts(counts, Errors.forCode(partitionResponse.errorCode()));
+        }));
+        return counts;
+    }
+
+    @Override
+    public int throttleTimeMs() {
+        return data.throttleTimeMs();
+    }
+
+    public static AlterIsrResponse parse(ByteBuffer buffer, short version) {
+        return new AlterIsrResponse(new AlterIsrResponseData(new ByteBufferAccessor(buffer), version));
+    }
+}
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..2d289cc1497e1
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/AlterPartitionReassignmentsRequest.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.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.ByteBufferAccessor;
+
+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 AlterPartitionReassignmentsRequest(AlterPartitionReassignmentsRequestData data, short version) {
+        super(ApiKeys.ALTER_PARTITION_REASSIGNMENTS, version);
+        this.data = data;
+    }
+
+    public static AlterPartitionReassignmentsRequest parse(ByteBuffer buffer, short version) {
+        return new AlterPartitionReassignmentsRequest(new AlterPartitionReassignmentsRequestData(
+            new ByteBufferAccessor(buffer), version), version);
+    }
+
+    public AlterPartitionReassignmentsRequestData data() {
+        return data;
+    }
+
+    @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/AlterPartitionReassignmentsResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/AlterPartitionReassignmentsResponse.java
new file mode 100644
index 0000000000000..ab166b812718c
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/AlterPartitionReassignmentsResponse.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.common.requests;
+
+import org.apache.kafka.common.message.AlterPartitionReassignmentsResponseData;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.Errors;
+
+import java.nio.ByteBuffer;
+import java.util.HashMap;
+import java.util.Map;
+
+public class AlterPartitionReassignmentsResponse extends AbstractResponse {
+
+    private final AlterPartitionReassignmentsResponseData data;
+
+    public AlterPartitionReassignmentsResponse(AlterPartitionReassignmentsResponseData data) {
+        super(ApiKeys.ALTER_PARTITION_REASSIGNMENTS);
+        this.data = data;
+    }
+
+    public static AlterPartitionReassignmentsResponse parse(ByteBuffer buffer, short version) {
+        return new AlterPartitionReassignmentsResponse(
+            new AlterPartitionReassignmentsResponseData(new ByteBufferAccessor(buffer), version));
+    }
+
+    @Override
+    public AlterPartitionReassignmentsResponseData 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<>();
+        updateErrorCounts(counts, Errors.forCode(data.errorCode()));
+
+        data.responses().forEach(topicResponse ->
+            topicResponse.partitions().forEach(partitionResponse ->
+                updateErrorCounts(counts, Errors.forCode(partitionResponse.errorCode()))
+        ));
+        return counts;
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/AlterReplicaLogDirsRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/AlterReplicaLogDirsRequest.java
index ba217598efb2a..68a87e6bf407e 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/AlterReplicaLogDirsRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/AlterReplicaLogDirsRequest.java
@@ -17,149 +17,78 @@
 
 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.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.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.stream.Collectors;
 
-import static org.apache.kafka.common.protocol.CommonFields.TOPIC_NAME;
-import static org.apache.kafka.common.protocol.types.Type.INT32;
-import static org.apache.kafka.common.protocol.types.Type.STRING;
+import org.apache.kafka.common.message.AlterReplicaLogDirsRequestData;
+import org.apache.kafka.common.message.AlterReplicaLogDirsResponseData;
+import org.apache.kafka.common.message.AlterReplicaLogDirsResponseData.AlterReplicaLogDirTopicResult;
 
 public class AlterReplicaLogDirsRequest extends AbstractRequest {
 
-    // request level key names
-    private static final String LOG_DIRS_KEY_NAME = "log_dirs";
-
-    // log dir level key names
-    private static final String LOG_DIR_KEY_NAME = "log_dir";
-    private static final String TOPICS_KEY_NAME = "topics";
-
-    // topic level key names
-    private static final String PARTITIONS_KEY_NAME = "partitions";
-
-    private static final Schema ALTER_REPLICA_LOG_DIRS_REQUEST_V0 = new Schema(
-            new Field("log_dirs", new ArrayOf(new Schema(
-                    new Field("log_dir", STRING, "The absolute log directory path."),
-                    new Field("topics", new ArrayOf(new Schema(
-                            TOPIC_NAME,
-                            new Field("partitions", new ArrayOf(INT32), "List of partition ids of the topic."))))))));
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{ALTER_REPLICA_LOG_DIRS_REQUEST_V0};
-    }
-
-    private final Map partitionDirs;
+    private final AlterReplicaLogDirsRequestData data;
 
     public static class Builder extends AbstractRequest.Builder {
-        private final Map partitionDirs;
+        private final AlterReplicaLogDirsRequestData data;
 
-        public Builder(Map partitionDirs) {
+        public Builder(AlterReplicaLogDirsRequestData data) {
             super(ApiKeys.ALTER_REPLICA_LOG_DIRS);
-            this.partitionDirs = partitionDirs;
+            this.data = data;
         }
 
         @Override
         public AlterReplicaLogDirsRequest build(short version) {
-            return new AlterReplicaLogDirsRequest(partitionDirs, version);
+            return new AlterReplicaLogDirsRequest(data, version);
         }
 
         @Override
         public String toString() {
-            StringBuilder builder = new StringBuilder();
-            builder.append("(type=AlterReplicaLogDirsRequest")
-                .append(", partitionDirs=")
-                .append(partitionDirs)
-                .append(")");
-            return builder.toString();
-        }
-    }
-
-    public AlterReplicaLogDirsRequest(Struct struct, short version) {
-        super(version);
-        partitionDirs = new HashMap<>();
-        for (Object logDirStructObj : struct.getArray(LOG_DIRS_KEY_NAME)) {
-            Struct logDirStruct = (Struct) logDirStructObj;
-            String logDir = logDirStruct.getString(LOG_DIR_KEY_NAME);
-            for (Object topicStructObj : logDirStruct.getArray(TOPICS_KEY_NAME)) {
-                Struct topicStruct = (Struct) topicStructObj;
-                String topic = topicStruct.get(TOPIC_NAME);
-                for (Object partitionObj : topicStruct.getArray(PARTITIONS_KEY_NAME)) {
-                    int partition = (Integer) partitionObj;
-                    partitionDirs.put(new TopicPartition(topic, partition), logDir);
-                }
-            }
+            return data.toString();
         }
     }
 
-    public AlterReplicaLogDirsRequest(Map partitionDirs, short version) {
-        super(version);
-        this.partitionDirs = partitionDirs;
+    public AlterReplicaLogDirsRequest(AlterReplicaLogDirsRequestData data, short version) {
+        super(ApiKeys.ALTER_REPLICA_LOG_DIRS, version);
+        this.data = data;
     }
 
     @Override
-    protected Struct toStruct() {
-        Map> dirPartitions = new HashMap<>();
-        for (Map.Entry entry: partitionDirs.entrySet()) {
-            if (!dirPartitions.containsKey(entry.getValue()))
-                dirPartitions.put(entry.getValue(), new ArrayList());
-            dirPartitions.get(entry.getValue()).add(entry.getKey());
-        }
-
-        Struct struct = new Struct(ApiKeys.ALTER_REPLICA_LOG_DIRS.requestSchema(version()));
-        List logDirStructArray = new ArrayList<>();
-        for (Map.Entry> logDirEntry: dirPartitions.entrySet()) {
-            Struct logDirStruct = struct.instance(LOG_DIRS_KEY_NAME);
-            logDirStruct.set(LOG_DIR_KEY_NAME, logDirEntry.getKey());
-
-            List topicStructArray = new ArrayList<>();
-            for (Map.Entry> topicEntry: CollectionUtils.groupDataByTopic(logDirEntry.getValue()).entrySet()) {
-                Struct topicStruct = logDirStruct.instance(TOPICS_KEY_NAME);
-                topicStruct.set(TOPIC_NAME, topicEntry.getKey());
-                topicStruct.set(PARTITIONS_KEY_NAME, topicEntry.getValue().toArray());
-                topicStructArray.add(topicStruct);
-            }
-            logDirStruct.set(TOPICS_KEY_NAME, topicStructArray.toArray());
-            logDirStructArray.add(logDirStruct);
-        }
-        struct.set(LOG_DIRS_KEY_NAME, logDirStructArray.toArray());
-        return struct;
+    public AlterReplicaLogDirsRequestData data() {
+        return data;
     }
 
-    @Override
-    public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
-        Map responseMap = new HashMap<>();
-
-        for (Map.Entry entry : partitionDirs.entrySet()) {
-            responseMap.put(entry.getKey(), Errors.forException(e));
-        }
-
-        short versionId = version();
-        switch (versionId) {
-            case 0:
-                return new AlterReplicaLogDirsResponse(throttleTimeMs, responseMap);
-            default:
-                throw new IllegalArgumentException(
-                    String.format("Version %d is not valid. Valid versions for %s are 0 to %d", versionId,
-                        this.getClass().getSimpleName(), ApiKeys.ALTER_REPLICA_LOG_DIRS.latestVersion()));
-        }
+    public AlterReplicaLogDirsResponse getErrorResponse(int throttleTimeMs, Throwable e) {
+        AlterReplicaLogDirsResponseData data = new AlterReplicaLogDirsResponseData();
+        data.setResults(this.data.dirs().stream().flatMap(alterDir ->
+            alterDir.topics().stream().map(topic ->
+                new AlterReplicaLogDirTopicResult()
+                    .setTopicName(topic.name())
+                    .setPartitions(topic.partitions().stream().map(partitionId ->
+                        new AlterReplicaLogDirsResponseData.AlterReplicaLogDirPartitionResult()
+                            .setErrorCode(Errors.forException(e).code())
+                            .setPartitionIndex(partitionId)).collect(Collectors.toList())))).collect(Collectors.toList()));
+        return new AlterReplicaLogDirsResponse(data.setThrottleTimeMs(throttleTimeMs));
     }
 
     public Map partitionDirs() {
-        return partitionDirs;
+        Map result = new HashMap<>();
+        data.dirs().forEach(alterDir ->
+            alterDir.topics().forEach(topic ->
+                topic.partitions().forEach(partition ->
+                    result.put(new TopicPartition(topic.name(), partition.intValue()), alterDir.path())))
+        );
+        return result;
     }
 
     public static AlterReplicaLogDirsRequest parse(ByteBuffer buffer, short version) {
-        return new AlterReplicaLogDirsRequest(ApiKeys.ALTER_REPLICA_LOG_DIRS.parseRequest(version, buffer), version);
+        return new AlterReplicaLogDirsRequest(new AlterReplicaLogDirsRequestData(new ByteBufferAccessor(buffer), version), version);
     }
 }
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 f8d15466d62e9..afa658d1e150f 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
@@ -17,119 +17,57 @@
 
 package org.apache.kafka.common.requests;
 
-import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.message.AlterReplicaLogDirsResponseData;
 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.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 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:
+ *
+ * {@link Errors#LOG_DIR_NOT_FOUND}
+ * {@link Errors#KAFKA_STORAGE_ERROR}
+ * {@link Errors#REPLICA_NOT_AVAILABLE}
+ * {@link Errors#UNKNOWN_SERVER_ERROR}
+ */
 public class AlterReplicaLogDirsResponse extends AbstractResponse {
 
-    // request level key names
-    private static final String TOPICS_KEY_NAME = "topics";
-
-    // topic level key names
-    private static final String PARTITIONS_KEY_NAME = "partitions";
-
-    private static final Schema ALTER_REPLICA_LOG_DIRS_RESPONSE_V0 = new Schema(
-            THROTTLE_TIME_MS,
-            new Field(TOPICS_KEY_NAME, new ArrayOf(new Schema(
-                    TOPIC_NAME,
-                    new Field(PARTITIONS_KEY_NAME, new ArrayOf(new Schema(
-                            PARTITION_ID,
-                            ERROR_CODE)))))));
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{ALTER_REPLICA_LOG_DIRS_RESPONSE_V0};
-    }
-
-    /**
-     * Possible error code:
-     *
-     * LOG_DIR_NOT_FOUND (57)
-     * KAFKA_STORAGE_ERROR (56)
-     * REPLICA_NOT_AVAILABLE (9)
-     * UNKNOWN (-1)
-     */
-    private final Map responses;
-    private final int throttleTimeMs;
+    private final AlterReplicaLogDirsResponseData data;
 
-    public AlterReplicaLogDirsResponse(Struct struct) {
-        throttleTimeMs = struct.get(THROTTLE_TIME_MS);
-        responses = new HashMap<>();
-        for (Object topicStructObj : struct.getArray(TOPICS_KEY_NAME)) {
-            Struct topicStruct = (Struct) topicStructObj;
-            String topic = topicStruct.get(TOPIC_NAME);
-            for (Object partitionStructObj : topicStruct.getArray(PARTITIONS_KEY_NAME)) {
-                Struct partitionStruct = (Struct) partitionStructObj;
-                int partition = partitionStruct.get(PARTITION_ID);
-                Errors error = Errors.forCode(partitionStruct.get(ERROR_CODE));
-                responses.put(new TopicPartition(topic, partition), error);
-            }
-        }
-    }
-
-    /**
-     * Constructor for version 0.
-     */
-    public AlterReplicaLogDirsResponse(int throttleTimeMs, Map responses) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.responses = responses;
+    public AlterReplicaLogDirsResponse(AlterReplicaLogDirsResponseData data) {
+        super(ApiKeys.ALTER_REPLICA_LOG_DIRS);
+        this.data = data;
     }
 
     @Override
-    protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.ALTER_REPLICA_LOG_DIRS.responseSchema(version));
-        struct.set(THROTTLE_TIME_MS, throttleTimeMs);
-        Map> responsesByTopic = CollectionUtils.groupDataByTopic(responses);
-        List topicStructArray = new ArrayList<>();
-        for (Map.Entry> responsesByTopicEntry : responsesByTopic.entrySet()) {
-            Struct topicStruct = struct.instance(TOPICS_KEY_NAME);
-            topicStruct.set(TOPIC_NAME, responsesByTopicEntry.getKey());
-            List partitionStructArray = new ArrayList<>();
-            for (Map.Entry responsesByPartitionEntry : responsesByTopicEntry.getValue().entrySet()) {
-                Struct partitionStruct = topicStruct.instance(PARTITIONS_KEY_NAME);
-                Errors response = responsesByPartitionEntry.getValue();
-                partitionStruct.set(PARTITION_ID, responsesByPartitionEntry.getKey());
-                partitionStruct.set(ERROR_CODE, response.code());
-                partitionStructArray.add(partitionStruct);
-            }
-            topicStruct.set(PARTITIONS_KEY_NAME, partitionStructArray.toArray());
-            topicStructArray.add(topicStruct);
-        }
-        struct.set(TOPICS_KEY_NAME, topicStructArray.toArray());
-        return struct;
+    public AlterReplicaLogDirsResponseData data() {
+        return data;
     }
 
+    @Override
     public int throttleTimeMs() {
-        return throttleTimeMs;
-    }
-
-    public Map responses() {
-        return this.responses;
+        return data.throttleTimeMs();
     }
 
     @Override
     public Map errorCounts() {
-        return errorCounts(responses);
+        Map errorCounts = new HashMap<>();
+        data.results().forEach(topicResult ->
+            topicResult.partitions().forEach(partitionResult ->
+                updateErrorCounts(errorCounts, Errors.forCode(partitionResult.errorCode()))));
+        return errorCounts;
     }
 
     public static AlterReplicaLogDirsResponse parse(ByteBuffer buffer, short version) {
-        return new AlterReplicaLogDirsResponse(ApiKeys.ALTER_REPLICA_LOG_DIRS.responseSchema(version).read(buffer));
+        return new AlterReplicaLogDirsResponse(new AlterReplicaLogDirsResponseData(new ByteBufferAccessor(buffer), version));
+    }
+
+    @Override
+    public boolean shouldClientThrottle(short version) {
+        return version >= 1;
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/AlterUserScramCredentialsRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/AlterUserScramCredentialsRequest.java
new file mode 100644
index 0000000000000..1ca7ea77aa422
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/AlterUserScramCredentialsRequest.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 org.apache.kafka.common.requests;
+
+import org.apache.kafka.common.message.AlterUserScramCredentialsRequestData;
+import org.apache.kafka.common.message.AlterUserScramCredentialsResponseData;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+
+import java.nio.ByteBuffer;
+import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+public class AlterUserScramCredentialsRequest extends AbstractRequest {
+
+    public static class Builder extends AbstractRequest.Builder {
+        private final AlterUserScramCredentialsRequestData data;
+
+        public Builder(AlterUserScramCredentialsRequestData data) {
+            super(ApiKeys.ALTER_USER_SCRAM_CREDENTIALS);
+            this.data = data;
+        }
+
+        @Override
+        public AlterUserScramCredentialsRequest build(short version) {
+            return new AlterUserScramCredentialsRequest(data, version);
+        }
+
+        @Override
+        public String toString() {
+            return data.toString();
+        }
+    }
+
+    private final AlterUserScramCredentialsRequestData data;
+
+    private AlterUserScramCredentialsRequest(AlterUserScramCredentialsRequestData data, short version) {
+        super(ApiKeys.ALTER_USER_SCRAM_CREDENTIALS, version);
+        this.data = data;
+    }
+
+    public static AlterUserScramCredentialsRequest parse(ByteBuffer buffer, short version) {
+        return new AlterUserScramCredentialsRequest(new AlterUserScramCredentialsRequestData(new ByteBufferAccessor(buffer), version), version);
+    }
+
+    @Override
+    public AlterUserScramCredentialsRequestData data() {
+        return data;
+    }
+
+    @Override
+    public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
+        ApiError apiError = ApiError.fromThrowable(e);
+        short errorCode = apiError.error().code();
+        String errorMessage = apiError.message();
+        Set users = Stream.concat(
+                this.data.deletions().stream().map(deletion -> deletion.name()),
+                this.data.upsertions().stream().map(upsertion -> upsertion.name()))
+                .collect(Collectors.toSet());
+        List results =
+                users.stream().sorted().map(user ->
+                        new AlterUserScramCredentialsResponseData.AlterUserScramCredentialsResult()
+                                .setUser(user)
+                                .setErrorCode(errorCode)
+                                .setErrorMessage(errorMessage))
+                        .collect(Collectors.toList());
+        return new AlterUserScramCredentialsResponse(new AlterUserScramCredentialsResponseData().setResults(results));
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/AlterUserScramCredentialsResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/AlterUserScramCredentialsResponse.java
new file mode 100644
index 0000000000000..97c0b7d17b204
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/AlterUserScramCredentialsResponse.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.common.requests;
+
+import org.apache.kafka.common.message.AlterUserScramCredentialsResponseData;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.Errors;
+
+import java.nio.ByteBuffer;
+import java.util.Map;
+
+public class AlterUserScramCredentialsResponse extends AbstractResponse {
+
+    private final AlterUserScramCredentialsResponseData data;
+
+    public AlterUserScramCredentialsResponse(AlterUserScramCredentialsResponseData responseData) {
+        super(ApiKeys.ALTER_USER_SCRAM_CREDENTIALS);
+        this.data = responseData;
+    }
+
+    @Override
+    public AlterUserScramCredentialsResponseData data() {
+        return data;
+    }
+
+    @Override
+    public boolean shouldClientThrottle(short version) {
+        return true;
+    }
+
+    @Override
+    public int throttleTimeMs() {
+        return data.throttleTimeMs();
+    }
+
+    @Override
+    public Map errorCounts() {
+        return errorCounts(data.results().stream().map(r -> Errors.forCode(r.errorCode())));
+    }
+
+    public static AlterUserScramCredentialsResponse parse(ByteBuffer buffer, short version) {
+        return new AlterUserScramCredentialsResponse(new AlterUserScramCredentialsResponseData(new ByteBufferAccessor(buffer), 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..38712adbd55f8 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
@@ -19,10 +19,8 @@
 
 import org.apache.kafka.common.errors.ApiException;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.Struct;
 
-import static org.apache.kafka.common.protocol.CommonFields.ERROR_CODE;
-import static org.apache.kafka.common.protocol.CommonFields.ERROR_MESSAGE;
+import java.util.Objects;
 
 /**
  * Encapsulates an error code (via the Errors enum) and an optional message. Generally, the optional message is only
@@ -38,16 +36,15 @@ public class ApiError {
     private final String message;
 
     public static ApiError fromThrowable(Throwable t) {
-        // Avoid populating the error message if it's a generic one
+        // Avoid populating the error message if it's a generic one. Also don't populate error
+        // message for UNKNOWN_SERVER_ERROR to ensure we don't leak sensitive information.
         Errors error = Errors.forException(t);
-        String message = error.message().equals(t.getMessage()) ? null : t.getMessage();
+        String message = error == Errors.UNKNOWN_SERVER_ERROR || error.message().equals(t.getMessage()) ? null : t.getMessage();
         return new ApiError(error, message);
     }
 
-    public ApiError(Struct struct) {
-        error = Errors.forCode(struct.get(ERROR_CODE));
-        // In some cases, the error message field was introduced in newer version
-        message = struct.getOrElse(ERROR_MESSAGE, null);
+    public ApiError(Errors error) {
+        this(error, error.message());
     }
 
     public ApiError(Errors error, String message) {
@@ -55,10 +52,9 @@ public ApiError(Errors error, String message) {
         this.message = message;
     }
 
-    public void write(Struct struct) {
-        struct.set(ERROR_CODE, error.code());
-        if (error != Errors.NONE)
-            struct.setIfExists(ERROR_MESSAGE, message);
+    public ApiError(short code, String message) {
+        this.error = Errors.forCode(code);
+        this.message = message;
     }
 
     public boolean is(Errors error) {
@@ -98,6 +94,21 @@ public ApiException exception() {
         return error.exception(message);
     }
 
+    @Override
+    public int hashCode() {
+        return Objects.hash(error, message);
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (!(o instanceof ApiError)) {
+            return false;
+        }
+        ApiError other = (ApiError) o;
+        return Objects.equals(error, other.error) &&
+            Objects.equals(message, other.message);
+    }
+
     @Override
     public String toString() {
         return "ApiError(error=" + error + ", 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 22daf6ce2c01c..946a9d9da7bed 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,25 +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.ByteBufferAccessor;
 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;
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{API_VERSIONS_REQUEST_V0, API_VERSIONS_REQUEST_V1};
-    }
 
     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);
@@ -46,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);
+    private final ApiVersionsRequestData data;
+
+    public ApiVersionsRequest(ApiVersionsRequestData data, short version) {
+        this(data, version, null);
     }
 
-    public ApiVersionsRequest(short version, Short unsupportedRequestVersion) {
-        super(version);
+    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
@@ -72,35 +78,49 @@ public ApiVersionsRequest(short version, Short unsupportedRequestVersion) {
         this.unsupportedRequestVersion = unsupportedRequestVersion;
     }
 
-    public ApiVersionsRequest(Struct struct, short version) {
-        this(version, null);
-    }
-
     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()));
+    public ApiVersionsRequestData data() {
+        return data;
     }
 
     @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:
-                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) {
-        return new ApiVersionsRequest(ApiKeys.API_VERSIONS.parseRequest(version, buffer), version);
+        return new ApiVersionsRequest(new ApiVersionsRequestData(new ByteBufferAccessor(buffer), version), 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 753db9e94bd86..4d7397ad96df1 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,180 +16,165 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.feature.Features;
+import org.apache.kafka.common.feature.FinalizedVersionRange;
+import org.apache.kafka.common.feature.SupportedVersionRange;
+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.ApiVersionsResponseData.FinalizedFeatureKey;
+import org.apache.kafka.common.message.ApiVersionsResponseData.FinalizedFeatureKeyCollection;
+import org.apache.kafka.common.message.ApiVersionsResponseData.SupportedFeatureKey;
+import org.apache.kafka.common.message.ApiVersionsResponseData.SupportedFeatureKeyCollection;
 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.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);
-
-    // 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};
-    }
 
-    /**
-     * 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 static final long UNKNOWN_FINALIZED_FEATURES_EPOCH = -1L;
 
-        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 +
-                    ")";
-        }
-    }
+    private final ApiVersionsResponseData data;
 
-    public ApiVersionsResponse(Errors error, List apiVersions) {
-        this(DEFAULT_THROTTLE_TIME, error, apiVersions);
+    public ApiVersionsResponse(ApiVersionsResponseData data) {
+        super(ApiKeys.API_VERSIONS);
+        this.data = data;
     }
 
-    public ApiVersionsResponse(int throttleTimeMs, Errors error, List apiVersions) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.error = error;
-        this.apiKeyToApiVersion = buildApiKeyToApiVersion(apiVersions);
+    @Override
+    public ApiVersionsResponseData data() {
+        return data;
     }
 
-    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 ApiVersionsResponseKey apiVersion(short apiKey) {
+        return data.apiKeys().find(apiKey);
     }
 
     @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;
-    }
-
-    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 Map errorCounts() {
+        return errorCounts(Errors.forCode(this.data.errorCode()));
     }
 
+    @Override
     public int throttleTimeMs() {
-        return throttleTimeMs;
-    }
-
-    public Collection apiVersions() {
-        return apiKeyToApiVersion.values();
-    }
-
-    public ApiVersion apiVersion(short apiKey) {
-        return apiKeyToApiVersion.get(apiKey);
-    }
-
-    public Errors error() {
-        return error;
+        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 which means that the version received by this
+        // method is not necessarily 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);
-        }
-        return tempApiIdToApiVersion;
+    public static ApiVersionsResponse createApiVersionsResponse(final int throttleTimeMs, final byte minMagic) {
+        return createApiVersionsResponse(throttleTimeMs, minMagic, Features.emptySupportedFeatures(),
+            Features.emptyFinalizedFeatures(), UNKNOWN_FINALIZED_FEATURES_EPOCH);
     }
 
-    public static ApiVersionsResponse createApiVersionsResponse(int throttleTimeMs, final byte minMagic) {
-        List versionList = new ArrayList<>();
-        for (ApiKeys apiKey : ApiKeys.values()) {
+    private static ApiVersionsResponse createApiVersionsResponse(
+        final int throttleTimeMs,
+        final byte minMagic,
+        final Features latestSupportedFeatures,
+        final Features finalizedFeatures,
+        final long finalizedFeaturesEpoch) {
+        return new ApiVersionsResponse(
+            createApiVersionsResponseData(
+                throttleTimeMs,
+                Errors.NONE,
+                defaultApiKeys(minMagic),
+                latestSupportedFeatures,
+                finalizedFeatures,
+                finalizedFeaturesEpoch));
+    }
+
+    public static ApiVersionsResponseKeyCollection defaultApiKeys(final byte minMagic) {
+        ApiVersionsResponseKeyCollection apiKeys = new ApiVersionsResponseKeyCollection();
+        for (ApiKeys apiKey : ApiKeys.enabledApis()) {
             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);
-    }
+        return apiKeys;
+    }
+
+    public static ApiVersionsResponseData createApiVersionsResponseData(
+        final int throttleTimeMs,
+        final Errors error,
+        final ApiVersionsResponseKeyCollection apiKeys,
+        final Features latestSupportedFeatures,
+        final Features finalizedFeatures,
+        final long finalizedFeaturesEpoch
+    ) {
+        final ApiVersionsResponseData data = new ApiVersionsResponseData();
+        data.setThrottleTimeMs(throttleTimeMs);
+        data.setErrorCode(error.code());
+        data.setApiKeys(apiKeys);
+        data.setSupportedFeatures(createSupportedFeatureKeys(latestSupportedFeatures));
+        data.setFinalizedFeatures(createFinalizedFeatureKeys(finalizedFeatures));
+        data.setFinalizedFeaturesEpoch(finalizedFeaturesEpoch);
+
+        return data;
+    }
+
+    private static SupportedFeatureKeyCollection createSupportedFeatureKeys(
+        Features latestSupportedFeatures) {
+        SupportedFeatureKeyCollection converted = new SupportedFeatureKeyCollection();
+        for (Map.Entry feature : latestSupportedFeatures.features().entrySet()) {
+            final SupportedFeatureKey key = new SupportedFeatureKey();
+            final SupportedVersionRange versionRange = feature.getValue();
+            key.setName(feature.getKey());
+            key.setMinVersion(versionRange.min());
+            key.setMaxVersion(versionRange.max());
+            converted.add(key);
+        }
 
-    public static ApiVersionsResponse defaultApiVersionsResponse() {
-        if (defaultApiVersionsResponse == null)
-            defaultApiVersionsResponse = createApiVersionsResponse(DEFAULT_THROTTLE_TIME, RecordBatch.CURRENT_MAGIC_VALUE);
-        return defaultApiVersionsResponse;
+        return converted;
     }
 
+    private static FinalizedFeatureKeyCollection createFinalizedFeatureKeys(
+        Features finalizedFeatures) {
+        FinalizedFeatureKeyCollection converted = new FinalizedFeatureKeyCollection();
+        for (Map.Entry feature : finalizedFeatures.features().entrySet()) {
+            final FinalizedFeatureKey key = new FinalizedFeatureKey();
+            final FinalizedVersionRange versionLevelRange = feature.getValue();
+            key.setName(feature.getKey());
+            key.setMinVersionLevel(versionLevelRange.min());
+            key.setMaxVersionLevel(versionLevelRange.max());
+            converted.add(key);
+        }
+
+        return converted;
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/BasePartitionState.java b/clients/src/main/java/org/apache/kafka/common/requests/BasePartitionState.java
deleted file mode 100644
index c9a461093caaf..0000000000000
--- a/clients/src/main/java/org/apache/kafka/common/requests/BasePartitionState.java
+++ /dev/null
@@ -1,40 +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 java.util.List;
-
-// This class contains the common fields shared between LeaderAndIsrRequest.PartitionState and UpdateMetadataRequest.PartitionState
-public class BasePartitionState {
-
-    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;
-    }
-
-}
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/BeginQuorumEpochRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/BeginQuorumEpochRequest.java
new file mode 100644
index 0000000000000..0794fb460959a
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/BeginQuorumEpochRequest.java
@@ -0,0 +1,94 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.BeginQuorumEpochRequestData;
+import org.apache.kafka.common.message.BeginQuorumEpochResponseData;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.Errors;
+
+import java.nio.ByteBuffer;
+import java.util.Collections;
+
+public class BeginQuorumEpochRequest extends AbstractRequest {
+    public static class Builder extends AbstractRequest.Builder {
+        private final BeginQuorumEpochRequestData data;
+
+        public Builder(BeginQuorumEpochRequestData data) {
+            super(ApiKeys.BEGIN_QUORUM_EPOCH);
+            this.data = data;
+        }
+
+        @Override
+        public BeginQuorumEpochRequest build(short version) {
+            return new BeginQuorumEpochRequest(data, version);
+        }
+
+        @Override
+        public String toString() {
+            return data.toString();
+        }
+    }
+
+    private final BeginQuorumEpochRequestData data;
+
+    private BeginQuorumEpochRequest(BeginQuorumEpochRequestData data, short version) {
+        super(ApiKeys.BEGIN_QUORUM_EPOCH, version);
+        this.data = data;
+    }
+
+    @Override
+    public BeginQuorumEpochRequestData data() {
+        return data;
+    }
+
+    @Override
+    public BeginQuorumEpochResponse getErrorResponse(int throttleTimeMs, Throwable e) {
+        return new BeginQuorumEpochResponse(new BeginQuorumEpochResponseData()
+            .setErrorCode(Errors.forException(e).code()));
+    }
+
+    public static BeginQuorumEpochRequest parse(ByteBuffer buffer, short version) {
+        return new BeginQuorumEpochRequest(new BeginQuorumEpochRequestData(new ByteBufferAccessor(buffer), version), version);
+    }
+
+    public static BeginQuorumEpochRequestData singletonRequest(TopicPartition topicPartition,
+                                                               int leaderEpoch,
+                                                               int leaderId) {
+        return singletonRequest(topicPartition, null, leaderEpoch, leaderId);
+    }
+
+    public static BeginQuorumEpochRequestData singletonRequest(TopicPartition topicPartition,
+                                                               String clusterId,
+                                                               int leaderEpoch,
+                                                               int leaderId) {
+        return new BeginQuorumEpochRequestData()
+                   .setClusterId(clusterId)
+                   .setTopics(Collections.singletonList(
+                       new BeginQuorumEpochRequestData.TopicData()
+                           .setTopicName(topicPartition.topic())
+                           .setPartitions(Collections.singletonList(
+                               new BeginQuorumEpochRequestData.PartitionData()
+                                   .setPartitionIndex(topicPartition.partition())
+                                   .setLeaderEpoch(leaderEpoch)
+                                   .setLeaderId(leaderId))))
+                   );
+    }
+
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/BeginQuorumEpochResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/BeginQuorumEpochResponse.java
new file mode 100644
index 0000000000000..c8c0328c93a15
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/BeginQuorumEpochResponse.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.TopicPartition;
+import org.apache.kafka.common.message.BeginQuorumEpochResponseData;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.Errors;
+
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * Possible error codes.
+ *
+ * Top level errors:
+ * - {@link Errors#CLUSTER_AUTHORIZATION_FAILED}
+ * - {@link Errors#BROKER_NOT_AVAILABLE}
+ *
+ * Partition level errors:
+ * - {@link Errors#FENCED_LEADER_EPOCH}
+ * - {@link Errors#INVALID_REQUEST}
+ * - {@link Errors#INCONSISTENT_VOTER_SET}
+ * - {@link Errors#UNKNOWN_TOPIC_OR_PARTITION}
+ */
+public class BeginQuorumEpochResponse extends AbstractResponse {
+    private final BeginQuorumEpochResponseData data;
+
+    public BeginQuorumEpochResponse(BeginQuorumEpochResponseData data) {
+        super(ApiKeys.BEGIN_QUORUM_EPOCH);
+        this.data = data;
+    }
+
+    public static BeginQuorumEpochResponseData singletonResponse(
+        Errors topLevelError,
+        TopicPartition topicPartition,
+        Errors partitionLevelError,
+        int leaderEpoch,
+        int leaderId
+    ) {
+        return new BeginQuorumEpochResponseData()
+                   .setErrorCode(topLevelError.code())
+                   .setTopics(Collections.singletonList(
+                       new BeginQuorumEpochResponseData.TopicData()
+                           .setTopicName(topicPartition.topic())
+                           .setPartitions(Collections.singletonList(
+                               new BeginQuorumEpochResponseData.PartitionData()
+                                   .setErrorCode(partitionLevelError.code())
+                                   .setLeaderId(leaderId)
+                                   .setLeaderEpoch(leaderEpoch)
+                           )))
+                   );
+    }
+
+    @Override
+    public Map errorCounts() {
+        Map errors = new HashMap<>();
+
+        errors.put(Errors.forCode(data.errorCode()), 1);
+
+        for (BeginQuorumEpochResponseData.TopicData topicResponse : data.topics()) {
+            for (BeginQuorumEpochResponseData.PartitionData partitionResponse : topicResponse.partitions()) {
+                errors.compute(Errors.forCode(partitionResponse.errorCode()),
+                    (error, count) -> count == null ? 1 : count + 1);
+            }
+        }
+        return errors;
+    }
+
+    @Override
+    public BeginQuorumEpochResponseData data() {
+        return data;
+    }
+
+    @Override
+    public int throttleTimeMs() {
+        return DEFAULT_THROTTLE_TIME;
+    }
+
+    public static BeginQuorumEpochResponse parse(ByteBuffer buffer, short version) {
+        return new BeginQuorumEpochResponse(new BeginQuorumEpochResponseData(new ByteBufferAccessor(buffer), version));
+    }
+
+}
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 e6e873429c071..088c351302fdd 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,89 +16,57 @@
  */
 package org.apache.kafka.common.requests;
 
-import org.apache.kafka.common.TopicPartition;
+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.ByteBufferAccessor;
 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;
 
 public class ControlledShutdownRequest extends AbstractRequest {
-    private static final String BROKER_ID_KEY_NAME = "broker_id";
-
-    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;
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {CONTROLLED_SHUTDOWN_REQUEST_V0, CONTROLLED_SHUTDOWN_REQUEST_V1};
-    }
 
     public static class Builder extends AbstractRequest.Builder {
-        private final int brokerId;
 
-        public Builder(int brokerId, short desiredVersion) {
+        private final ControlledShutdownRequestData data;
+
+        public Builder(ControlledShutdownRequestData data, short desiredVersion) {
             super(ApiKeys.CONTROLLED_SHUTDOWN, desiredVersion);
-            this.brokerId = brokerId;
+            this.data = data;
         }
 
         @Override
         public ControlledShutdownRequest build(short version) {
-            return new ControlledShutdownRequest(brokerId, version);
+            return new ControlledShutdownRequest(data, version);
         }
 
         @Override
         public String toString() {
-            StringBuilder bld = new StringBuilder();
-            bld.append("(type=ControlledShutdownRequest").
-                append(", brokerId=").append(brokerId).
-                append(")");
-            return bld.toString();
+            return data.toString();
         }
     }
-    private final int brokerId;
 
-    private ControlledShutdownRequest(int brokerId, short version) {
-        super(version);
-        this.brokerId = brokerId;
-    }
+    private final ControlledShutdownRequestData data;
 
-    public ControlledShutdownRequest(Struct struct, short version) {
-        super(version);
-        brokerId = struct.getInt(BROKER_ID_KEY_NAME);
+    private ControlledShutdownRequest(ControlledShutdownRequestData data, short version) {
+        super(ApiKeys.CONTROLLED_SHUTDOWN, version);
+        this.data = data;
     }
 
     @Override
-    public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
-        short versionId = version();
-        switch (versionId) {
-            case 0:
-            case 1:
-                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 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) {
-        return new ControlledShutdownRequest(
-                ApiKeys.CONTROLLED_SHUTDOWN.parseRequest(version, buffer), version);
+        return new ControlledShutdownRequest(new ControlledShutdownRequestData(new ByteBufferAccessor(buffer), version),
+            version);
     }
 
     @Override
-    protected Struct toStruct() {
-        Struct struct = new Struct(ApiKeys.CONTROLLED_SHUTDOWN.requestSchema(version()));
-        struct.set(BROKER_ID_KEY_NAME, brokerId);
-        return struct;
+    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 f9852dd472d84..73b6a50268379 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,43 +17,19 @@
 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.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.Struct;
 
 import java.nio.ByteBuffer;
-import java.util.ArrayList;
-import java.util.HashSet;
-import java.util.List;
 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;
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{CONTROLLED_SHUTDOWN_RESPONSE_V0, CONTROLLED_SHUTDOWN_RESPONSE_V1};
-    }
-
     /**
      * Possible error codes:
      *
@@ -61,58 +37,47 @@ public static Schema[] schemaVersions() {
      * BROKER_NOT_AVAILABLE(8)
      * STALE_CONTROLLER_EPOCH(11)
      */
-    private final Errors error;
-
-    private final Set partitionsRemaining;
+    private final ControlledShutdownResponseData data;
 
-    public ControlledShutdownResponse(Errors error, Set partitionsRemaining) {
-        this.error = error;
-        this.partitionsRemaining = partitionsRemaining;
-    }
-
-    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(ControlledShutdownResponseData data) {
+        super(ApiKeys.CONTROLLED_SHUTDOWN);
+        this.data = data;
     }
 
     public Errors error() {
-        return error;
+        return Errors.forCode(data.errorCode());
     }
 
     @Override
     public Map errorCounts() {
-        return errorCounts(error);
+        return errorCounts(error());
     }
 
-    public Set partitionsRemaining() {
-        return partitionsRemaining;
+    @Override
+    public int throttleTimeMs() {
+        return DEFAULT_THROTTLE_TIME;
     }
 
     public static ControlledShutdownResponse parse(ByteBuffer buffer, short version) {
-        return new ControlledShutdownResponse(ApiKeys.CONTROLLED_SHUTDOWN.parseResponse(version, buffer));
+        return new ControlledShutdownResponse(new ControlledShutdownResponseData(new ByteBufferAccessor(buffer), version));
     }
 
     @Override
-    protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.CONTROLLED_SHUTDOWN.responseSchema(version));
-        struct.set(ERROR_CODE, error.code());
-
-        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/CorrelationIdMismatchException.java b/clients/src/main/java/org/apache/kafka/common/requests/CorrelationIdMismatchException.java
new file mode 100644
index 0000000000000..2610a27ea41a2
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/CorrelationIdMismatchException.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.requests;
+
+/**
+ * Raised if the correlationId in a response header does not match
+ * the expected value from the request header.
+ */
+public class CorrelationIdMismatchException extends IllegalStateException {
+    private final int requestCorrelationId;
+    private final int responseCorrelationId;
+
+    public CorrelationIdMismatchException(
+        String message,
+        int requestCorrelationId,
+        int responseCorrelationId
+    ) {
+        super(message);
+        this.requestCorrelationId = requestCorrelationId;
+        this.responseCorrelationId = responseCorrelationId;
+    }
+
+    public int requestCorrelationId() {
+        return requestCorrelationId;
+    }
+
+    public int responseCorrelationId() {
+        return responseCorrelationId;
+    }
+
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/CreateAclsRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/CreateAclsRequest.java
index d281b3ba4afb3..29df8326bc5a6 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/CreateAclsRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/CreateAclsRequest.java
@@ -19,137 +19,119 @@
 
 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.errors.UnsupportedVersionException;
+import org.apache.kafka.common.message.CreateAclsRequestData;
+import org.apache.kafka.common.message.CreateAclsRequestData.AclCreation;
+import org.apache.kafka.common.message.CreateAclsResponseData;
+import org.apache.kafka.common.message.CreateAclsResponseData.AclCreationResult;
 import org.apache.kafka.common.protocol.ApiKeys;
-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.resource.Resource;
-import org.apache.kafka.common.utils.Utils;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.resource.PatternType;
+import org.apache.kafka.common.resource.ResourcePattern;
+import org.apache.kafka.common.resource.ResourceType;
 
 import java.nio.ByteBuffer;
-import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 
-import static org.apache.kafka.common.protocol.CommonFields.HOST;
-import static org.apache.kafka.common.protocol.CommonFields.OPERATION;
-import static org.apache.kafka.common.protocol.CommonFields.PERMISSION_TYPE;
-import static org.apache.kafka.common.protocol.CommonFields.PRINCIPAL;
-import static org.apache.kafka.common.protocol.CommonFields.RESOURCE_NAME;
-import static org.apache.kafka.common.protocol.CommonFields.RESOURCE_TYPE;
-
 public class CreateAclsRequest extends AbstractRequest {
-    private final static String CREATIONS_KEY_NAME = "creations";
-
-    private static final Schema CREATE_ACLS_REQUEST_V0 = new Schema(
-            new Field(CREATIONS_KEY_NAME, new ArrayOf(new Schema(
-                    RESOURCE_TYPE,
-                    RESOURCE_NAME,
-                    PRINCIPAL,
-                    HOST,
-                    OPERATION,
-                    PERMISSION_TYPE))));
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{CREATE_ACLS_REQUEST_V0};
-    }
-
-    public static class AclCreation {
-        private final AclBinding acl;
-
-        public AclCreation(AclBinding acl) {
-            this.acl = acl;
-        }
-
-        static AclCreation fromStruct(Struct struct) {
-            Resource resource = RequestUtils.resourceFromStructFields(struct);
-            AccessControlEntry entry = RequestUtils.aceFromStructFields(struct);
-            return new AclCreation(new AclBinding(resource, entry));
-        }
-
-        public AclBinding acl() {
-            return acl;
-        }
-
-        void setStructFields(Struct struct) {
-            RequestUtils.resourceSetStructFields(acl.resource(), struct);
-            RequestUtils.aceSetStructFields(acl.entry(), struct);
-        }
-
-        @Override
-        public String toString() {
-            return "(acl=" + acl + ")";
-        }
-    }
 
     public static class Builder extends AbstractRequest.Builder {
-        private final List creations;
+        private final CreateAclsRequestData data;
 
-        public Builder(List creations) {
+        public Builder(CreateAclsRequestData data) {
             super(ApiKeys.CREATE_ACLS);
-            this.creations = creations;
+            this.data = data;
         }
 
         @Override
         public CreateAclsRequest build(short version) {
-            return new CreateAclsRequest(version, creations);
+            return new CreateAclsRequest(data, version);
         }
 
         @Override
         public String toString() {
-            return "(type=CreateAclsRequest, creations=" + Utils.join(creations, ", ") + ")";
+            return data.toString();
         }
     }
 
-    private final List aclCreations;
+    private final CreateAclsRequestData data;
 
-    CreateAclsRequest(short version, List aclCreations) {
-        super(version);
-        this.aclCreations = aclCreations;
+    CreateAclsRequest(CreateAclsRequestData data, short version) {
+        super(ApiKeys.CREATE_ACLS, version);
+        validate(data);
+        this.data = data;
     }
 
-    public CreateAclsRequest(Struct struct, short version) {
-        super(version);
-        this.aclCreations = new ArrayList<>();
-        for (Object creationStructObj : struct.getArray(CREATIONS_KEY_NAME)) {
-            Struct creationStruct = (Struct) creationStructObj;
-            aclCreations.add(AclCreation.fromStruct(creationStruct));
-        }
+    public List aclCreations() {
+        return data.creations();
     }
 
     @Override
-    protected Struct toStruct() {
-        Struct struct = new Struct(ApiKeys.CREATE_ACLS.requestSchema(version()));
-        List requests = new ArrayList<>();
-        for (AclCreation creation : aclCreations) {
-            Struct creationStruct = struct.instance(CREATIONS_KEY_NAME);
-            creation.setStructFields(creationStruct);
-            requests.add(creationStruct);
-        }
-        struct.set(CREATIONS_KEY_NAME, requests.toArray());
-        return struct;
-    }
-
-    public List aclCreations() {
-        return aclCreations;
+    public CreateAclsRequestData data() {
+        return data;
     }
 
     @Override
     public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable throwable) {
-        short versionId = version();
-        switch (versionId) {
-            case 0:
-                List responses = new ArrayList<>();
-                for (int i = 0; i < aclCreations.size(); i++)
-                    responses.add(new CreateAclsResponse.AclCreationResponse(ApiError.fromThrowable(throwable)));
-                return new CreateAclsResponse(throttleTimeMs, 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.CREATE_ACLS.latestVersion()));
-        }
+        CreateAclsResponseData.AclCreationResult result = CreateAclsRequest.aclResult(throwable);
+        List results = Collections.nCopies(data.creations().size(), result);
+        return new CreateAclsResponse(new CreateAclsResponseData()
+            .setThrottleTimeMs(throttleTimeMs)
+            .setResults(results));
     }
 
     public static CreateAclsRequest parse(ByteBuffer buffer, short version) {
-        return new CreateAclsRequest(ApiKeys.CREATE_ACLS.parseRequest(version, buffer), version);
+        return new CreateAclsRequest(new CreateAclsRequestData(new ByteBufferAccessor(buffer), version), version);
+    }
+
+    private void validate(CreateAclsRequestData data) {
+        if (version() == 0) {
+            final boolean unsupported = data.creations().stream().anyMatch(creation ->
+                creation.resourcePatternType() != PatternType.LITERAL.code());
+            if (unsupported)
+                throw new UnsupportedVersionException("Version 0 only supports literal resource pattern types");
+        }
+
+        final boolean unknown = data.creations().stream().anyMatch(creation ->
+            creation.resourcePatternType() == PatternType.UNKNOWN.code()
+                || creation.resourceType() == ResourceType.UNKNOWN.code()
+                || creation.permissionType() == AclPermissionType.UNKNOWN.code()
+                || creation.operation() == AclOperation.UNKNOWN.code());
+        if (unknown)
+            throw new IllegalArgumentException("CreatableAcls contain unknown elements: " + data.creations());
+    }
+
+    public static AclBinding aclBinding(AclCreation acl) {
+        ResourcePattern pattern = new ResourcePattern(
+            ResourceType.fromCode(acl.resourceType()),
+            acl.resourceName(),
+            PatternType.fromCode(acl.resourcePatternType()));
+        AccessControlEntry entry = new AccessControlEntry(
+            acl.principal(),
+            acl.host(),
+            AclOperation.fromCode(acl.operation()),
+            AclPermissionType.fromCode(acl.permissionType()));
+        return new AclBinding(pattern, entry);
+    }
+
+    public static AclCreation aclCreation(AclBinding binding) {
+        return new AclCreation()
+            .setHost(binding.entry().host())
+            .setOperation(binding.entry().operation().code())
+            .setPermissionType(binding.entry().permissionType().code())
+            .setPrincipal(binding.entry().principal())
+            .setResourceName(binding.pattern().name())
+            .setResourceType(binding.pattern().resourceType().code())
+            .setResourcePatternType(binding.pattern().patternType().code());
+    }
+
+    private static AclCreationResult aclResult(Throwable throwable) {
+        ApiError apiError = ApiError.fromThrowable(throwable);
+        return new AclCreationResult()
+            .setErrorCode(apiError.error().code())
+            .setErrorMessage(apiError.message());
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/CreateAclsResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/CreateAclsResponse.java
index 6c3579893cef5..8bc6643f9de01 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/CreateAclsResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/CreateAclsResponse.java
@@ -16,103 +16,48 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.message.CreateAclsResponseData;
 import org.apache.kafka.common.protocol.ApiKeys;
-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.ByteBufferAccessor;
 import org.apache.kafka.common.protocol.Errors;
-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.ERROR_MESSAGE;
-import static org.apache.kafka.common.protocol.CommonFields.THROTTLE_TIME_MS;
-
 public class CreateAclsResponse extends AbstractResponse {
-    private final static String CREATION_RESPONSES_KEY_NAME = "creation_responses";
-
-    private static final Schema CREATE_ACLS_RESPONSE_V0 = new Schema(
-            THROTTLE_TIME_MS,
-            new Field(CREATION_RESPONSES_KEY_NAME, new ArrayOf(new Schema(
-                    ERROR_CODE,
-                    ERROR_MESSAGE))));
+    private final CreateAclsResponseData data;
 
-    public static Schema[] schemaVersions() {
-        return new Schema[]{CREATE_ACLS_RESPONSE_V0};
-    }
-
-    public static class AclCreationResponse {
-        private final ApiError error;
-
-        public AclCreationResponse(ApiError error) {
-            this.error = error;
-        }
-
-        public ApiError error() {
-            return error;
-        }
-
-        @Override
-        public String toString() {
-            return "(" + error + ")";
-        }
-    }
-
-    private final int throttleTimeMs;
-
-    private final List aclCreationResponses;
-
-    public CreateAclsResponse(int throttleTimeMs, List aclCreationResponses) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.aclCreationResponses = aclCreationResponses;
-    }
-
-    public CreateAclsResponse(Struct struct) {
-        this.throttleTimeMs = struct.get(THROTTLE_TIME_MS);
-        this.aclCreationResponses = new ArrayList<>();
-        for (Object responseStructObj : struct.getArray(CREATION_RESPONSES_KEY_NAME)) {
-            Struct responseStruct = (Struct) responseStructObj;
-            ApiError error = new ApiError(responseStruct);
-            this.aclCreationResponses.add(new AclCreationResponse(error));
-        }
+    public CreateAclsResponse(CreateAclsResponseData data) {
+        super(ApiKeys.CREATE_ACLS);
+        this.data = data;
     }
 
     @Override
-    protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.CREATE_ACLS.responseSchema(version));
-        struct.set(THROTTLE_TIME_MS, throttleTimeMs);
-        List responseStructs = new ArrayList<>();
-        for (AclCreationResponse response : aclCreationResponses) {
-            Struct responseStruct = struct.instance(CREATION_RESPONSES_KEY_NAME);
-            response.error.write(responseStruct);
-            responseStructs.add(responseStruct);
-        }
-        struct.set(CREATION_RESPONSES_KEY_NAME, responseStructs.toArray());
-        return struct;
+    public CreateAclsResponseData data() {
+        return data;
     }
 
+    @Override
     public int throttleTimeMs() {
-        return throttleTimeMs;
+        return data.throttleTimeMs();
     }
 
-    public List aclCreationResponses() {
-        return aclCreationResponses;
+    public List results() {
+        return data.results();
     }
 
     @Override
     public Map errorCounts() {
-        Map errorCounts = new HashMap<>();
-        for (AclCreationResponse response : aclCreationResponses)
-            updateErrorCounts(errorCounts, response.error.error());
-        return errorCounts;
+        return errorCounts(results().stream().map(r -> Errors.forCode(r.errorCode())));
     }
 
     public static CreateAclsResponse parse(ByteBuffer buffer, short version) {
-        return new CreateAclsResponse(ApiKeys.CREATE_ACLS.responseSchema(version).read(buffer));
+        return new CreateAclsResponse(new CreateAclsResponseData(new ByteBufferAccessor(buffer), version));
+    }
+
+    @Override
+    public boolean shouldClientThrottle(short version) {
+        return version >= 1;
     }
 }
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
new file mode 100644
index 0000000000000..1fee1b71eb3a4
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/CreateDelegationTokenRequest.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.common.requests;
+
+import org.apache.kafka.common.message.CreateDelegationTokenRequestData;
+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.security.auth.KafkaPrincipal;
+
+import java.nio.ByteBuffer;
+
+public class CreateDelegationTokenRequest extends AbstractRequest {
+
+    private final CreateDelegationTokenRequestData data;
+
+    private CreateDelegationTokenRequest(CreateDelegationTokenRequestData data, short version) {
+        super(ApiKeys.CREATE_DELEGATION_TOKEN, version);
+        this.data = data;
+    }
+
+    public static CreateDelegationTokenRequest parse(ByteBuffer buffer, short version) {
+        return new CreateDelegationTokenRequest(new CreateDelegationTokenRequestData(new ByteBufferAccessor(buffer), version),
+            version);
+    }
+
+    @Override
+    public CreateDelegationTokenRequestData data() {
+        return data;
+    }
+
+    @Override
+    public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
+        return CreateDelegationTokenResponse.prepareResponse(throttleTimeMs, Errors.forException(e), KafkaPrincipal.ANONYMOUS);
+    }
+
+    public static class Builder extends AbstractRequest.Builder {
+        private final CreateDelegationTokenRequestData data;
+
+        public Builder(CreateDelegationTokenRequestData data) {
+            super(ApiKeys.CREATE_DELEGATION_TOKEN);
+            this.data = data;
+        }
+
+        @Override
+        public CreateDelegationTokenRequest build(short version) {
+            return new CreateDelegationTokenRequest(data, version);
+        }
+
+        @Override
+        public String 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
new file mode 100644
index 0000000000000..b679a30c8dd5c
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/CreateDelegationTokenResponse.java
@@ -0,0 +1,94 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.CreateDelegationTokenResponseData;
+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.security.auth.KafkaPrincipal;
+
+import java.nio.ByteBuffer;
+import java.util.Map;
+
+public class CreateDelegationTokenResponse extends AbstractResponse {
+
+    private final CreateDelegationTokenResponseData data;
+
+    public CreateDelegationTokenResponse(CreateDelegationTokenResponseData data) {
+        super(ApiKeys.CREATE_DELEGATION_TOKEN);
+        this.data = data;
+    }
+
+    public static CreateDelegationTokenResponse parse(ByteBuffer buffer, short version) {
+        return new CreateDelegationTokenResponse(
+            new CreateDelegationTokenResponseData(new ByteBufferAccessor(buffer), version));
+    }
+
+    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 static CreateDelegationTokenResponse prepareResponse(int throttleTimeMs, Errors error, KafkaPrincipal owner) {
+        return prepareResponse(throttleTimeMs, error, owner, -1, -1, -1, "", ByteBuffer.wrap(new byte[] {}));
+    }
+
+    @Override
+    public CreateDelegationTokenResponseData data() {
+        return data;
+    }
+
+    @Override
+    public Map errorCounts() {
+        return errorCounts(error());
+    }
+
+    @Override
+    public int throttleTimeMs() {
+        return data.throttleTimeMs();
+    }
+
+    public Errors error() {
+        return Errors.forCode(data.errorCode());
+    }
+
+    public boolean hasError() {
+        return error() != Errors.NONE;
+    }
+
+    @Override
+    public boolean shouldClientThrottle(short version) {
+        return version >= 1;
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/CreatePartitionsRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/CreatePartitionsRequest.java
index 41c5327189ca4..d371bbb216995 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/CreatePartitionsRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/CreatePartitionsRequest.java
@@ -17,198 +17,66 @@
 
 package org.apache.kafka.common.requests;
 
-import org.apache.kafka.clients.admin.NewPartitions;
+import org.apache.kafka.common.message.CreatePartitionsRequestData;
+import org.apache.kafka.common.message.CreatePartitionsRequestData.CreatePartitionsTopic;
+import org.apache.kafka.common.message.CreatePartitionsResponseData;
+import org.apache.kafka.common.message.CreatePartitionsResponseData.CreatePartitionsTopicResult;
 import org.apache.kafka.common.protocol.ApiKeys;
-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.protocol.ByteBufferAccessor;
 
 import java.nio.ByteBuffer;
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
-
-import static org.apache.kafka.common.protocol.CommonFields.TOPIC_NAME;
-import static org.apache.kafka.common.protocol.types.Type.BOOLEAN;
-import static org.apache.kafka.common.protocol.types.Type.INT32;
 
 public class CreatePartitionsRequest extends AbstractRequest {
 
-    private static final String TOPIC_PARTITIONS_KEY_NAME = "topic_partitions";
-    private static final String NEW_PARTITIONS_KEY_NAME = "new_partitions";
-    private static final String COUNT_KEY_NAME = "count";
-    private static final String ASSIGNMENT_KEY_NAME = "assignment";
-    private static final String TIMEOUT_KEY_NAME = "timeout";
-    private static final String VALIDATE_ONLY_KEY_NAME = "validate_only";
-
-    private static final Schema CREATE_PARTITIONS_REQUEST_V0 = new Schema(
-            new Field(TOPIC_PARTITIONS_KEY_NAME, new ArrayOf(
-                    new Schema(
-                            TOPIC_NAME,
-                            new Field(NEW_PARTITIONS_KEY_NAME, new Schema(
-                                    new Field(COUNT_KEY_NAME, INT32, "The new partition count."),
-                                    new Field(ASSIGNMENT_KEY_NAME, ArrayOf.nullable(new ArrayOf(INT32)),
-                                            "The assigned brokers.")
-                            )))),
-                    "List of topic and the corresponding new partitions."),
-            new Field(TIMEOUT_KEY_NAME, INT32, "The time in ms to wait for the partitions to be created."),
-            new Field(VALIDATE_ONLY_KEY_NAME, BOOLEAN,
-                    "If true then validate the request, but don't actually increase the number of partitions."));
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{CREATE_PARTITIONS_REQUEST_V0};
-    }
-
-    // It is an error for duplicate topics to be present in the request,
-    // so track duplicates here to allow KafkaApis to report per-topic errors.
-    private final Set duplicates;
-    private final Map newPartitions;
-    private final int timeout;
-    private final boolean validateOnly;
+    private final CreatePartitionsRequestData data;
 
     public static class Builder extends AbstractRequest.Builder {
 
-        private final Map newPartitions;
-        private final int timeout;
-        private final boolean validateOnly;
+        private final CreatePartitionsRequestData data;
 
-        public Builder(Map newPartitions, int timeout, boolean validateOnly) {
+        public Builder(CreatePartitionsRequestData data) {
             super(ApiKeys.CREATE_PARTITIONS);
-            this.newPartitions = newPartitions;
-            this.timeout = timeout;
-            this.validateOnly = validateOnly;
+            this.data = data;
         }
 
         @Override
         public CreatePartitionsRequest build(short version) {
-            return new CreatePartitionsRequest(newPartitions, timeout, validateOnly, version);
+            return new CreatePartitionsRequest(data, version);
         }
 
         @Override
         public String toString() {
-            StringBuilder bld = new StringBuilder();
-            bld.append("(type=CreatePartitionsRequest").
-                    append(", newPartitions=").append(newPartitions).
-                    append(", timeout=").append(timeout).
-                    append(", validateOnly=").append(validateOnly).
-                    append(")");
-            return bld.toString();
-        }
-    }
-
-    CreatePartitionsRequest(Map newPartitions, int timeout, boolean validateOnly, short apiVersion) {
-        super(apiVersion);
-        this.newPartitions = newPartitions;
-        this.duplicates = Collections.emptySet();
-        this.timeout = timeout;
-        this.validateOnly = validateOnly;
-    }
-
-    public CreatePartitionsRequest(Struct struct, short apiVersion) {
-        super(apiVersion);
-        Object[] topicCountArray = struct.getArray(TOPIC_PARTITIONS_KEY_NAME);
-        Map counts = new HashMap<>(topicCountArray.length);
-        Set dupes = new HashSet<>();
-        for (Object topicPartitionCountObj : topicCountArray) {
-            Struct topicPartitionCountStruct = (Struct) topicPartitionCountObj;
-            String topic = topicPartitionCountStruct.get(TOPIC_NAME);
-            Struct partitionCountStruct = topicPartitionCountStruct.getStruct(NEW_PARTITIONS_KEY_NAME);
-            int count = partitionCountStruct.getInt(COUNT_KEY_NAME);
-            Object[] assignmentsArray = partitionCountStruct.getArray(ASSIGNMENT_KEY_NAME);
-            NewPartitions newPartition;
-            if (assignmentsArray != null) {
-                List> assignments = new ArrayList<>(assignmentsArray.length);
-                for (Object replicas : assignmentsArray) {
-                    Object[] replicasArray = (Object[]) replicas;
-                    List replicasList = new ArrayList<>(replicasArray.length);
-                    assignments.add(replicasList);
-                    for (Object broker : replicasArray) {
-                        replicasList.add((Integer) broker);
-                    }
-                }
-                newPartition = NewPartitions.increaseTo(count, assignments);
-            } else {
-                newPartition = NewPartitions.increaseTo(count);
-            }
-            NewPartitions dupe = counts.put(topic, newPartition);
-            if (dupe != null) {
-                dupes.add(topic);
-            }
+            return data.toString();
         }
-        this.newPartitions = counts;
-        this.duplicates = dupes;
-        this.timeout = struct.getInt(TIMEOUT_KEY_NAME);
-        this.validateOnly = struct.getBoolean(VALIDATE_ONLY_KEY_NAME);
     }
 
-    public Set duplicates() {
-        return duplicates;
-    }
-
-    public Map newPartitions() {
-        return newPartitions;
-    }
-
-    public int timeout() {
-        return timeout;
-    }
-
-    public boolean validateOnly() {
-        return validateOnly;
+    CreatePartitionsRequest(CreatePartitionsRequestData data, short apiVersion) {
+        super(ApiKeys.CREATE_PARTITIONS, apiVersion);
+        this.data = data;
     }
 
     @Override
-    protected Struct toStruct() {
-        Struct struct = new Struct(ApiKeys.CREATE_PARTITIONS.requestSchema(version()));
-        List topicPartitionsList = new ArrayList<>();
-        for (Map.Entry topicPartitionCount : this.newPartitions.entrySet()) {
-            Struct topicPartitionCountStruct = struct.instance(TOPIC_PARTITIONS_KEY_NAME);
-            topicPartitionCountStruct.set(TOPIC_NAME, topicPartitionCount.getKey());
-            NewPartitions count = topicPartitionCount.getValue();
-            Struct partitionCountStruct = topicPartitionCountStruct.instance(NEW_PARTITIONS_KEY_NAME);
-            partitionCountStruct.set(COUNT_KEY_NAME, count.totalCount());
-            Object[][] assignments = null;
-            if (count.assignments() != null) {
-                assignments = new Object[count.assignments().size()][];
-                int i = 0;
-                for (List partitionAssignment : count.assignments()) {
-                    assignments[i] = partitionAssignment.toArray(new Object[0]);
-                    i++;
-                }
-            }
-            partitionCountStruct.set(ASSIGNMENT_KEY_NAME, assignments);
-            topicPartitionCountStruct.set(NEW_PARTITIONS_KEY_NAME, partitionCountStruct);
-            topicPartitionsList.add(topicPartitionCountStruct);
-        }
-        struct.set(TOPIC_PARTITIONS_KEY_NAME, topicPartitionsList.toArray(new Object[0]));
-        struct.set(TIMEOUT_KEY_NAME, this.timeout);
-        struct.set(VALIDATE_ONLY_KEY_NAME, this.validateOnly);
-        return struct;
+    public CreatePartitionsRequestData data() {
+        return data;
     }
 
     @Override
     public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
-        Map topicErrors = new HashMap<>();
-        for (String topic : newPartitions.keySet()) {
-            topicErrors.put(topic, ApiError.fromThrowable(e));
-        }
-
-        short versionId = version();
-        switch (versionId) {
-            case 0:
-                return new CreatePartitionsResponse(throttleTimeMs, topicErrors);
-            default:
-                throw new IllegalArgumentException(String.format("Version %d is not valid. Valid versions for %s are 0 to %d",
-                        versionId, this.getClass().getSimpleName(), ApiKeys.CREATE_PARTITIONS.latestVersion()));
+        CreatePartitionsResponseData response = new CreatePartitionsResponseData();
+        response.setThrottleTimeMs(throttleTimeMs);
+
+        ApiError apiError = ApiError.fromThrowable(e);
+        for (CreatePartitionsTopic topic : data.topics()) {
+            response.results().add(new CreatePartitionsTopicResult()
+                    .setName(topic.name())
+                    .setErrorCode(apiError.error().code())
+                    .setErrorMessage(apiError.message())
+            );
         }
+        return new CreatePartitionsResponse(response);
     }
 
     public static CreatePartitionsRequest parse(ByteBuffer buffer, short version) {
-        return new CreatePartitionsRequest(ApiKeys.CREATE_PARTITIONS.parseRequest(version, buffer), version);
+        return new CreatePartitionsRequest(new CreatePartitionsRequestData(new ByteBufferAccessor(buffer), version), version);
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/CreatePartitionsResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/CreatePartitionsResponse.java
index 73a6b56566657..e59ac981f112a 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/CreatePartitionsResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/CreatePartitionsResponse.java
@@ -17,95 +17,49 @@
 
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.message.CreatePartitionsResponseData;
 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.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.ERROR_MESSAGE;
-import static org.apache.kafka.common.protocol.CommonFields.THROTTLE_TIME_MS;
-import static org.apache.kafka.common.protocol.CommonFields.TOPIC_NAME;
-
-
 public class CreatePartitionsResponse extends AbstractResponse {
 
-    private static final String TOPIC_ERRORS_KEY_NAME = "topic_errors";
-
-    private static final Schema CREATE_PARTITIONS_RESPONSE_V0 = new Schema(
-            THROTTLE_TIME_MS,
-            new Field(TOPIC_ERRORS_KEY_NAME, new ArrayOf(
-                    new Schema(
-                            TOPIC_NAME,
-                            ERROR_CODE,
-                            ERROR_MESSAGE
-                    )), "Per topic results for the create partitions request")
-    );
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{CREATE_PARTITIONS_RESPONSE_V0};
-    }
-
-    private final int throttleTimeMs;
-    private final Map errors;
+    private final CreatePartitionsResponseData data;
 
-    public CreatePartitionsResponse(int throttleTimeMs, Map errors) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.errors = errors;
+    public CreatePartitionsResponse(CreatePartitionsResponseData data) {
+        super(ApiKeys.CREATE_PARTITIONS);
+        this.data = data;
     }
 
-    public CreatePartitionsResponse(Struct struct) {
-        super();
-        Object[] topicErrorsArray = struct.getArray(TOPIC_ERRORS_KEY_NAME);
-        Map errors = new HashMap<>(topicErrorsArray.length);
-        for (Object topicErrorObj : topicErrorsArray) {
-            Struct topicErrorStruct = (Struct) topicErrorObj;
-            String topic = topicErrorStruct.get(TOPIC_NAME);
-            ApiError error = new ApiError(topicErrorStruct);
-            errors.put(topic, error);
-        }
-        this.throttleTimeMs = struct.get(THROTTLE_TIME_MS);
-        this.errors = errors;
+    @Override
+    public CreatePartitionsResponseData data() {
+        return data;
     }
 
     @Override
-    protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.CREATE_PARTITIONS.responseSchema(version));
-        List topicErrors = new ArrayList<>(errors.size());
-        for (Map.Entry error : errors.entrySet()) {
-            Struct errorStruct = struct.instance(TOPIC_ERRORS_KEY_NAME);
-            errorStruct.set(TOPIC_NAME, error.getKey());
-            error.getValue().write(errorStruct);
-            topicErrors.add(errorStruct);
-        }
-        struct.set(THROTTLE_TIME_MS, throttleTimeMs);
-        struct.set(TOPIC_ERRORS_KEY_NAME, topicErrors.toArray(new Object[topicErrors.size()]));
-        return struct;
+    public Map errorCounts() {
+        Map counts = new HashMap<>();
+        data.results().forEach(result ->
+            updateErrorCounts(counts, Errors.forCode(result.errorCode()))
+        );
+        return counts;
     }
 
-    public Map errors() {
-        return errors;
+    public static CreatePartitionsResponse parse(ByteBuffer buffer, short version) {
+        return new CreatePartitionsResponse(new CreatePartitionsResponseData(new ByteBufferAccessor(buffer), version));
     }
 
     @Override
-    public Map errorCounts() {
-        return apiErrorCounts(errors);
+    public boolean shouldClientThrottle(short version) {
+        return version >= 1;
     }
 
+    @Override
     public int throttleTimeMs() {
-        return throttleTimeMs;
+        return data.throttleTimeMs();
     }
-
-    public static CreatePartitionsResponse parse(ByteBuffer buffer, short version) {
-        return new CreatePartitionsResponse(ApiKeys.CREATE_PARTITIONS.parseResponse(version, buffer));
-    }
-
 }
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 3a1de51f8f224..9a1032b8c0955 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,332 +16,89 @@
  */
 package org.apache.kafka.common.requests;
 
-import org.apache.kafka.common.errors.UnsupportedVersionException;
-import org.apache.kafka.common.protocol.ApiKeys;
-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.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.BOOLEAN;
-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;
-import static org.apache.kafka.common.protocol.types.Type.STRING;
+import java.util.stream.Collectors;
+import org.apache.kafka.common.errors.UnsupportedVersionException;
+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 org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
 
 public class CreateTopicsRequest extends AbstractRequest {
-    private static final String REQUESTS_KEY_NAME = "create_topic_requests";
-
-    private static final String TIMEOUT_KEY_NAME = "timeout";
-    private static final String VALIDATE_ONLY_KEY_NAME = "validate_only";
-    private static final String NUM_PARTITIONS_KEY_NAME = "num_partitions";
-    private static final String REPLICATION_FACTOR_KEY_NAME = "replication_factor";
-    private static final String REPLICA_ASSIGNMENT_KEY_NAME = "replica_assignment";
-    private static final String REPLICA_ASSIGNMENT_REPLICAS_KEY_NAME = "replicas";
-
-    private static final String CONFIG_NAME_KEY_NAME = "config_name";
-    private static final String CONFIG_VALUE_KEY_NAME = "config_value";
-    private static final String CONFIG_ENTRIES_KEY_NAME = "config_entries";
-
-    private static final Schema CONFIG_ENTRY = new Schema(
-            new Field(CONFIG_NAME_KEY_NAME, STRING, "Configuration name"),
-            new Field(CONFIG_VALUE_KEY_NAME, NULLABLE_STRING, "Configuration value"));
-
-    private static final Schema PARTITION_REPLICA_ASSIGNMENT_ENTRY = new Schema(
-            PARTITION_ID,
-            new Field(REPLICA_ASSIGNMENT_REPLICAS_KEY_NAME, new ArrayOf(INT32), "The set of all nodes that should " +
-                    "host this partition. The first replica in the list is the preferred leader."));
-
-    private static final Schema SINGLE_CREATE_TOPIC_REQUEST_V0 = new Schema(
-            TOPIC_NAME,
-            new Field(NUM_PARTITIONS_KEY_NAME, INT32, "Number of partitions to be created. -1 indicates unset."),
-            new Field(REPLICATION_FACTOR_KEY_NAME, INT16, "Replication factor for the topic. -1 indicates unset."),
-            new Field(REPLICA_ASSIGNMENT_KEY_NAME, new ArrayOf(PARTITION_REPLICA_ASSIGNMENT_ENTRY),
-                    "Replica assignment among kafka brokers for this topic partitions. If this is set num_partitions " +
-                            "and replication_factor must be unset."),
-            new Field(CONFIG_ENTRIES_KEY_NAME, new ArrayOf(CONFIG_ENTRY), "Topic level configuration for topic to be set."));
-
-    private static final Schema SINGLE_CREATE_TOPIC_REQUEST_V1 = SINGLE_CREATE_TOPIC_REQUEST_V0;
-
-    private static final Schema CREATE_TOPICS_REQUEST_V0 = new Schema(
-            new Field(REQUESTS_KEY_NAME, new ArrayOf(SINGLE_CREATE_TOPIC_REQUEST_V0),
-                    "An array of single topic creation requests. Can not have multiple entries for the same topic."),
-            new Field(TIMEOUT_KEY_NAME, INT32, "The time in ms to wait for a topic to be completely created on the " +
-                    "controller node. Values <= 0 will trigger topic creation and return immediately"));
-
-    private static final Schema CREATE_TOPICS_REQUEST_V1 = new Schema(
-            new Field(REQUESTS_KEY_NAME, new ArrayOf(SINGLE_CREATE_TOPIC_REQUEST_V1), "An array of single " +
-                    "topic creation requests. Can not have multiple entries for the same topic."),
-            new Field(TIMEOUT_KEY_NAME, INT32, "The time in ms to wait for a topic to be completely created on the " +
-                    "controller node. Values <= 0 will trigger topic creation and return immediately"),
-            new Field(VALIDATE_ONLY_KEY_NAME, BOOLEAN, "If this is true, the request will be validated, but the " +
-                    "topic won't be created."));
-
-    /* v2 request is the same as v1. Throttle time has been added to the response */
-    private static final Schema CREATE_TOPICS_REQUEST_V2 = CREATE_TOPICS_REQUEST_V1;
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{CREATE_TOPICS_REQUEST_V0, CREATE_TOPICS_REQUEST_V1, CREATE_TOPICS_REQUEST_V2};
-    }
-
-    public static final class TopicDetails {
-        public final int numPartitions;
-        public final short replicationFactor;
-        public final Map> replicasAssignments;
-        public final Map configs;
-
-        private TopicDetails(int numPartitions,
-                             short replicationFactor,
-                             Map> replicasAssignments,
-                             Map configs) {
-            this.numPartitions = numPartitions;
-            this.replicationFactor = replicationFactor;
-            this.replicasAssignments = replicasAssignments;
-            this.configs = configs;
-        }
-
-        public TopicDetails(int partitions,
-                            short replicationFactor,
-                            Map configs) {
-            this(partitions, replicationFactor, Collections.>emptyMap(), configs);
-        }
-
-        public TopicDetails(int partitions,
-                            short replicationFactor) {
-            this(partitions, replicationFactor, Collections.emptyMap());
-        }
-
-        public TopicDetails(Map> replicasAssignments,
-                            Map configs) {
-            this(NO_NUM_PARTITIONS, NO_REPLICATION_FACTOR, replicasAssignments, configs);
-        }
-
-        public TopicDetails(Map> replicasAssignments) {
-            this(replicasAssignments, Collections.emptyMap());
-        }
-
-        @Override
-        public String toString() {
-            StringBuilder bld = new StringBuilder();
-            bld.append("(numPartitions=").append(numPartitions).
-                    append(", replicationFactor=").append(replicationFactor).
-                    append(", replicasAssignments=").append(replicasAssignments).
-                    append(", configs=").append(configs).
-                    append(")");
-            return bld.toString();
-        }
-    }
-
     public static class Builder extends AbstractRequest.Builder {
-        private final Map topics;
-        private final int timeout;
-        private final boolean validateOnly; // introduced in V1
-
-        public Builder(Map topics, int timeout) {
-            this(topics, timeout, false);
-        }
+        private final CreateTopicsRequestData data;
 
-        public Builder(Map topics, int timeout, boolean validateOnly) {
+        public Builder(CreateTopicsRequestData data) {
             super(ApiKeys.CREATE_TOPICS);
-            this.topics = topics;
-            this.timeout = timeout;
-            this.validateOnly = validateOnly;
+            this.data = data;
         }
 
         @Override
         public CreateTopicsRequest build(short version) {
-            if (validateOnly && version == 0)
+            if (data.validateOnly() && version == 0)
                 throw new UnsupportedVersionException("validateOnly is not supported in version 0 of " +
                         "CreateTopicsRequest");
-            return new CreateTopicsRequest(topics, timeout, validateOnly, version);
+
+            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);
         }
 
         @Override
         public String toString() {
-            StringBuilder bld = new StringBuilder();
-            bld.append("(type=CreateTopicsRequest").
-                append(", topics=").append(topics).
-                append(", timeout=").append(timeout).
-                append(", validateOnly=").append(validateOnly).
-                append(")");
-            return bld.toString();
+            return data.toString();
         }
     }
 
-    private final Map topics;
-    private final Integer timeout;
-    private final boolean validateOnly; // introduced in V1
-
-    // Set to handle special case where 2 requests for the same topic exist on the wire.
-    // This allows the broker to return an error code for these topics.
-    private final Set duplicateTopics;
+    private final CreateTopicsRequestData data;
 
     public static final int NO_NUM_PARTITIONS = -1;
     public static final short NO_REPLICATION_FACTOR = -1;
 
-    private CreateTopicsRequest(Map topics, Integer timeout, boolean validateOnly, short version) {
-        super(version);
-        this.topics = topics;
-        this.timeout = timeout;
-        this.validateOnly = validateOnly;
-        this.duplicateTopics = Collections.emptySet();
+    public CreateTopicsRequest(CreateTopicsRequestData data, short version) {
+        super(ApiKeys.CREATE_TOPICS, version);
+        this.data = data;
     }
 
-    public CreateTopicsRequest(Struct struct, short version) {
-        super(version);
-
-        Object[] requestStructs = struct.getArray(REQUESTS_KEY_NAME);
-        Map topics = new HashMap<>();
-        Set duplicateTopics = new HashSet<>();
-
-        for (Object requestStructObj : requestStructs) {
-            Struct singleRequestStruct = (Struct) requestStructObj;
-            String topic = singleRequestStruct.get(TOPIC_NAME);
-
-            if (topics.containsKey(topic))
-                duplicateTopics.add(topic);
-
-            int numPartitions = singleRequestStruct.getInt(NUM_PARTITIONS_KEY_NAME);
-            short replicationFactor = singleRequestStruct.getShort(REPLICATION_FACTOR_KEY_NAME);
-
-            //replica assignment
-            Object[] assignmentsArray = singleRequestStruct.getArray(REPLICA_ASSIGNMENT_KEY_NAME);
-            Map> partitionReplicaAssignments = new HashMap<>(assignmentsArray.length);
-            for (Object assignmentStructObj : assignmentsArray) {
-                Struct assignmentStruct = (Struct) assignmentStructObj;
-
-                Integer partitionId = assignmentStruct.get(PARTITION_ID);
-
-                Object[] replicasArray = assignmentStruct.getArray(REPLICA_ASSIGNMENT_REPLICAS_KEY_NAME);
-                List replicas = new ArrayList<>(replicasArray.length);
-                for (Object replica : replicasArray) {
-                    replicas.add((Integer) replica);
-                }
-
-                partitionReplicaAssignments.put(partitionId, replicas);
-            }
-
-            Object[] configArray = singleRequestStruct.getArray(CONFIG_ENTRIES_KEY_NAME);
-            Map configs = new HashMap<>(configArray.length);
-            for (Object configStructObj : configArray) {
-                Struct configStruct = (Struct) configStructObj;
-
-                String key = configStruct.getString(CONFIG_NAME_KEY_NAME);
-                String value = configStruct.getString(CONFIG_VALUE_KEY_NAME);
-
-                configs.put(key, value);
-            }
-
-            TopicDetails args = new TopicDetails(numPartitions, replicationFactor, partitionReplicaAssignments, configs);
-
-            topics.put(topic, args);
-        }
-
-        this.topics = topics;
-        this.timeout = struct.getInt(TIMEOUT_KEY_NAME);
-        if (struct.hasField(VALIDATE_ONLY_KEY_NAME))
-            this.validateOnly = struct.getBoolean(VALIDATE_ONLY_KEY_NAME);
-        else
-            this.validateOnly = false;
-        this.duplicateTopics = duplicateTopics;
+    @Override
+    public CreateTopicsRequestData data() {
+        return data;
     }
 
     @Override
     public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
-        Map topicErrors = new HashMap<>();
-        for (String topic : topics.keySet()) {
-            topicErrors.put(topic, ApiError.fromThrowable(e));
+        CreateTopicsResponseData response = new CreateTopicsResponseData();
+        if (version() >= 2) {
+            response.setThrottleTimeMs(throttleTimeMs);
         }
-
-        short versionId = version();
-        switch (versionId) {
-            case 0:
-            case 1:
-                return new CreateTopicsResponse(topicErrors);
-            case 2:
-                return new CreateTopicsResponse(throttleTimeMs, topicErrors);
-            default:
-                throw new IllegalArgumentException(String.format("Version %d is not valid. Valid versions for %s are 0 to %d",
-                    versionId, this.getClass().getSimpleName(), ApiKeys.CREATE_TOPICS.latestVersion()));
+        ApiError apiError = ApiError.fromThrowable(e);
+        for (CreatableTopic topic : data.topics()) {
+            response.topics().add(new CreatableTopicResult().
+                setName(topic.name()).
+                setErrorCode(apiError.error().code()).
+                setErrorMessage(apiError.message()));
         }
-    }
-
-    public Map topics() {
-        return this.topics;
-    }
-
-    public int timeout() {
-        return this.timeout;
-    }
-
-    public boolean validateOnly() {
-        return validateOnly;
-    }
-
-    public Set duplicateTopics() {
-        return this.duplicateTopics;
+        return new CreateTopicsResponse(response);
     }
 
     public static CreateTopicsRequest parse(ByteBuffer buffer, short version) {
-        return new CreateTopicsRequest(ApiKeys.CREATE_TOPICS.parseRequest(version, buffer), version);
-    }
-
-    /**
-     * Visible for testing.
-     */
-    @Override
-    public Struct toStruct() {
-        short version = version();
-        Struct struct = new Struct(ApiKeys.CREATE_TOPICS.requestSchema(version));
-
-        List createTopicRequestStructs = new ArrayList<>(topics.size());
-        for (Map.Entry entry : topics.entrySet()) {
-
-            Struct singleRequestStruct = struct.instance(REQUESTS_KEY_NAME);
-            String topic = entry.getKey();
-            TopicDetails args = entry.getValue();
-
-            singleRequestStruct.set(TOPIC_NAME, topic);
-            singleRequestStruct.set(NUM_PARTITIONS_KEY_NAME, args.numPartitions);
-            singleRequestStruct.set(REPLICATION_FACTOR_KEY_NAME, args.replicationFactor);
-
-            // replica assignment
-            List replicaAssignmentsStructs = new ArrayList<>(args.replicasAssignments.size());
-            for (Map.Entry> partitionReplicaAssignment : args.replicasAssignments.entrySet()) {
-                Struct replicaAssignmentStruct = singleRequestStruct.instance(REPLICA_ASSIGNMENT_KEY_NAME);
-                replicaAssignmentStruct.set(PARTITION_ID, partitionReplicaAssignment.getKey());
-                replicaAssignmentStruct.set(REPLICA_ASSIGNMENT_REPLICAS_KEY_NAME, partitionReplicaAssignment.getValue().toArray());
-                replicaAssignmentsStructs.add(replicaAssignmentStruct);
-            }
-            singleRequestStruct.set(REPLICA_ASSIGNMENT_KEY_NAME, replicaAssignmentsStructs.toArray());
-
-            // configs
-            List configsStructs = new ArrayList<>(args.configs.size());
-            for (Map.Entry configEntry : args.configs.entrySet()) {
-                Struct configStruct = singleRequestStruct.instance(CONFIG_ENTRIES_KEY_NAME);
-                configStruct.set(CONFIG_NAME_KEY_NAME, configEntry.getKey());
-                configStruct.set(CONFIG_VALUE_KEY_NAME, configEntry.getValue());
-                configsStructs.add(configStruct);
-            }
-            singleRequestStruct.set(CONFIG_ENTRIES_KEY_NAME, configsStructs.toArray());
-            createTopicRequestStructs.add(singleRequestStruct);
-        }
-        struct.set(REQUESTS_KEY_NAME, createTopicRequestStructs.toArray());
-        struct.set(TIMEOUT_KEY_NAME, timeout);
-        if (version >= 1)
-            struct.set(VALIDATE_ONLY_KEY_NAME, validateOnly);
-        return struct;
-
+        return new CreateTopicsRequest(new CreateTopicsRequestData(new ByteBufferAccessor(buffer), version), version);
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/CreateTopicsResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/CreateTopicsResponse.java
index 1f73edac1da19..dd0627742587c 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/CreateTopicsResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/CreateTopicsResponse.java
@@ -14,60 +14,25 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.kafka.common.requests;
 
+package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.message.CreateTopicsResponseData;
 import org.apache.kafka.common.protocol.ApiKeys;
-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.ByteBufferAccessor;
 import org.apache.kafka.common.protocol.Errors;
-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.ERROR_MESSAGE;
-import static org.apache.kafka.common.protocol.CommonFields.THROTTLE_TIME_MS;
-import static org.apache.kafka.common.protocol.CommonFields.TOPIC_NAME;
-
 public class CreateTopicsResponse extends AbstractResponse {
-    private static final String TOPIC_ERRORS_KEY_NAME = "topic_errors";
-
-    private static final Schema TOPIC_ERROR_CODE = new Schema(
-            TOPIC_NAME,
-            ERROR_CODE);
-
-    // Improves on TOPIC_ERROR_CODE by adding an error_message to complement the error_code
-    private static final Schema TOPIC_ERROR = new Schema(
-            TOPIC_NAME,
-            ERROR_CODE,
-            ERROR_MESSAGE);
-
-    private static final Schema CREATE_TOPICS_RESPONSE_V0 = new Schema(
-            new Field(TOPIC_ERRORS_KEY_NAME, new ArrayOf(TOPIC_ERROR_CODE), "An array of per topic error codes."));
-
-    private static final Schema CREATE_TOPICS_RESPONSE_V1 = new Schema(
-            new Field(TOPIC_ERRORS_KEY_NAME, new ArrayOf(TOPIC_ERROR), "An array of per topic errors."));
-
-    private static final Schema CREATE_TOPICS_RESPONSE_V2 = new Schema(
-            THROTTLE_TIME_MS,
-            new Field(TOPIC_ERRORS_KEY_NAME, new ArrayOf(TOPIC_ERROR), "An array of per topic errors."));
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{CREATE_TOPICS_RESPONSE_V0, CREATE_TOPICS_RESPONSE_V1, CREATE_TOPICS_RESPONSE_V2};
-    }
-
     /**
      * Possible error codes:
      *
      * REQUEST_TIMED_OUT(7)
      * INVALID_TOPIC_EXCEPTION(17)
-     * CLUSTER_AUTHORIZATION_FAILED(31)
+     * TOPIC_AUTHORIZATION_FAILED(29)
      * TOPIC_ALREADY_EXISTS(36)
      * INVALID_PARTITIONS(37)
      * INVALID_REPLICATION_FACTOR(38)
@@ -75,63 +40,41 @@ public static Schema[] schemaVersions() {
      * INVALID_CONFIG(40)
      * NOT_CONTROLLER(41)
      * INVALID_REQUEST(42)
+     * POLICY_VIOLATION(44)
      */
 
-    private final Map errors;
-    private final int throttleTimeMs;
+    private final CreateTopicsResponseData data;
 
-    public CreateTopicsResponse(Map errors) {
-        this(DEFAULT_THROTTLE_TIME, errors);
-    }
-
-    public CreateTopicsResponse(int throttleTimeMs, Map errors) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.errors = errors;
-    }
-
-    public CreateTopicsResponse(Struct struct) {
-        Object[] topicErrorStructs = struct.getArray(TOPIC_ERRORS_KEY_NAME);
-        Map errors = new HashMap<>();
-        for (Object topicErrorStructObj : topicErrorStructs) {
-            Struct topicErrorStruct = (Struct) topicErrorStructObj;
-            String topic = topicErrorStruct.get(TOPIC_NAME);
-            errors.put(topic, new ApiError(topicErrorStruct));
-        }
-
-        this.throttleTimeMs = struct.getOrElse(THROTTLE_TIME_MS, DEFAULT_THROTTLE_TIME);
-        this.errors = errors;
+    public CreateTopicsResponse(CreateTopicsResponseData data) {
+        super(ApiKeys.CREATE_TOPICS);
+        this.data = data;
     }
 
     @Override
-    protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.CREATE_TOPICS.responseSchema(version));
-        struct.setIfExists(THROTTLE_TIME_MS, throttleTimeMs);
-
-        List topicErrorsStructs = new ArrayList<>(errors.size());
-        for (Map.Entry topicError : errors.entrySet()) {
-            Struct topicErrorsStruct = struct.instance(TOPIC_ERRORS_KEY_NAME);
-            topicErrorsStruct.set(TOPIC_NAME, topicError.getKey());
-            topicError.getValue().write(topicErrorsStruct);
-            topicErrorsStructs.add(topicErrorsStruct);
-        }
-        struct.set(TOPIC_ERRORS_KEY_NAME, topicErrorsStructs.toArray());
-        return struct;
+    public CreateTopicsResponseData data() {
+        return data;
     }
 
+    @Override
     public int throttleTimeMs() {
-        return throttleTimeMs;
-    }
-
-    public Map errors() {
-        return errors;
+        return data.throttleTimeMs();
     }
 
     @Override
     public Map errorCounts() {
-        return apiErrorCounts(errors);
+        HashMap counts = new HashMap<>();
+        data.topics().forEach(result ->
+            updateErrorCounts(counts, Errors.forCode(result.errorCode()))
+        );
+        return counts;
     }
 
     public static CreateTopicsResponse parse(ByteBuffer buffer, short version) {
-        return new CreateTopicsResponse(ApiKeys.CREATE_TOPICS.responseSchema(version).read(buffer));
+        return new CreateTopicsResponse(new CreateTopicsResponseData(new ByteBufferAccessor(buffer), version));
+    }
+
+    @Override
+    public boolean shouldClientThrottle(short version) {
+        return version >= 3;
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/DeleteAclsRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/DeleteAclsRequest.java
index 2d50ea65cda38..98fd6589b2280 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/DeleteAclsRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/DeleteAclsRequest.java
@@ -16,118 +16,131 @@
  */
 package org.apache.kafka.common.requests;
 
+import java.util.Collections;
+import java.util.stream.Collectors;
 import org.apache.kafka.common.acl.AccessControlEntryFilter;
 import org.apache.kafka.common.acl.AclBindingFilter;
+import org.apache.kafka.common.acl.AclOperation;
+import org.apache.kafka.common.acl.AclPermissionType;
+import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.message.DeleteAclsRequestData;
+import org.apache.kafka.common.message.DeleteAclsRequestData.DeleteAclsFilter;
+import org.apache.kafka.common.message.DeleteAclsResponseData;
+import org.apache.kafka.common.message.DeleteAclsResponseData.DeleteAclsFilterResult;
 import org.apache.kafka.common.protocol.ApiKeys;
-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.resource.ResourceFilter;
-import org.apache.kafka.common.utils.Utils;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.resource.PatternType;
+import org.apache.kafka.common.resource.ResourcePatternFilter;
+import org.apache.kafka.common.resource.ResourceType;
 
 import java.nio.ByteBuffer;
-import java.util.ArrayList;
-import java.util.Collections;
 import java.util.List;
 
 import static org.apache.kafka.common.protocol.ApiKeys.DELETE_ACLS;
-import static org.apache.kafka.common.protocol.CommonFields.HOST_FILTER;
-import static org.apache.kafka.common.protocol.CommonFields.OPERATION;
-import static org.apache.kafka.common.protocol.CommonFields.PERMISSION_TYPE;
-import static org.apache.kafka.common.protocol.CommonFields.PRINCIPAL_FILTER;
-import static org.apache.kafka.common.protocol.CommonFields.RESOURCE_NAME_FILTER;
-import static org.apache.kafka.common.protocol.CommonFields.RESOURCE_TYPE;
 
 public class DeleteAclsRequest extends AbstractRequest {
-    private final static String FILTERS = "filters";
-
-    private static final Schema DELETE_ACLS_REQUEST_V0 = new Schema(
-            new Field(FILTERS, new ArrayOf(new Schema(
-                    RESOURCE_TYPE,
-                    RESOURCE_NAME_FILTER,
-                    PRINCIPAL_FILTER,
-                    HOST_FILTER,
-                    OPERATION,
-                    PERMISSION_TYPE))));
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{DELETE_ACLS_REQUEST_V0};
-    }
-
     public static class Builder extends AbstractRequest.Builder {
-        private final List filters;
+        private final DeleteAclsRequestData data;
 
-        public Builder(List filters) {
+        public Builder(DeleteAclsRequestData data) {
             super(DELETE_ACLS);
-            this.filters = filters;
+            this.data = data;
         }
 
         @Override
         public DeleteAclsRequest build(short version) {
-            return new DeleteAclsRequest(version, filters);
+            return new DeleteAclsRequest(data, version);
         }
 
         @Override
         public String toString() {
-            return "(type=DeleteAclsRequest, filters=" + Utils.join(filters, ", ") + ")";
+            return data.toString();
         }
+
     }
 
-    private final List filters;
+    private final DeleteAclsRequestData data;
 
-    DeleteAclsRequest(short version, List filters) {
-        super(version);
-        this.filters = filters;
+    private DeleteAclsRequest(DeleteAclsRequestData data, short version) {
+        super(ApiKeys.DELETE_ACLS, version);
+        this.data = data;
+        normalizeAndValidate();
     }
 
-    public DeleteAclsRequest(Struct struct, short version) {
-        super(version);
-        this.filters = new ArrayList<>();
-        for (Object filterStructObj : struct.getArray(FILTERS)) {
-            Struct filterStruct = (Struct) filterStructObj;
-            ResourceFilter resourceFilter = RequestUtils.resourceFilterFromStructFields(filterStruct);
-            AccessControlEntryFilter aceFilter = RequestUtils.aceFilterFromStructFields(filterStruct);
-            this.filters.add(new AclBindingFilter(resourceFilter, aceFilter));
+    private void normalizeAndValidate() {
+        if (version() == 0) {
+            for (DeleteAclsRequestData.DeleteAclsFilter filter : data.filters()) {
+                PatternType patternType = PatternType.fromCode(filter.patternTypeFilter());
+
+                // On older brokers, no pattern types existed except LITERAL (effectively). So even though ANY is not
+                // directly supported on those brokers, we can get the same effect as ANY by setting the pattern type
+                // to LITERAL. Note that the wildcard `*` is considered `LITERAL` for compatibility reasons.
+                if (patternType == PatternType.ANY)
+                    filter.setPatternTypeFilter(PatternType.LITERAL.code());
+                else if (patternType != PatternType.LITERAL)
+                    throw new UnsupportedVersionException("Version 0 does not support pattern type " +
+                            patternType + " (only LITERAL and ANY are supported)");
+            }
+        }
+
+        final boolean unknown = data.filters().stream().anyMatch(filter ->
+                filter.patternTypeFilter() == PatternType.UNKNOWN.code()
+                        || filter.resourceTypeFilter() == ResourceType.UNKNOWN.code()
+                        || filter.operation() == AclOperation.UNKNOWN.code()
+                        || filter.permissionType() == AclPermissionType.UNKNOWN.code()
+        );
+
+        if (unknown) {
+            throw new IllegalArgumentException("Filters contain UNKNOWN elements, filters: " + data.filters());
         }
     }
 
     public List filters() {
-        return filters;
+        return data.filters().stream().map(DeleteAclsRequest::aclBindingFilter).collect(Collectors.toList());
     }
 
     @Override
-    protected Struct toStruct() {
-        Struct struct = new Struct(DELETE_ACLS.requestSchema(version()));
-        List filterStructs = new ArrayList<>();
-        for (AclBindingFilter filter : filters) {
-            Struct filterStruct = struct.instance(FILTERS);
-            RequestUtils.resourceFilterSetStructFields(filter.resourceFilter(), filterStruct);
-            RequestUtils.aceFilterSetStructFields(filter.entryFilter(), filterStruct);
-            filterStructs.add(filterStruct);
-        }
-        struct.set(FILTERS, filterStructs.toArray());
-        return struct;
+    public DeleteAclsRequestData data() {
+        return data;
     }
 
     @Override
     public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable throwable) {
-        short versionId = version();
-        switch (versionId) {
-            case 0:
-                List responses = new ArrayList<>();
-                for (int i = 0; i < filters.size(); i++) {
-                    responses.add(new DeleteAclsResponse.AclFilterResponse(
-                        ApiError.fromThrowable(throwable), Collections.emptySet()));
-                }
-                return new DeleteAclsResponse(throttleTimeMs, 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.DELETE_ACLS.latestVersion()));
-        }
+        ApiError apiError = ApiError.fromThrowable(throwable);
+        List filterResults = Collections.nCopies(data.filters().size(),
+            new DeleteAclsResponseData.DeleteAclsFilterResult()
+                .setErrorCode(apiError.error().code())
+                .setErrorMessage(apiError.message()));
+        return new DeleteAclsResponse(new DeleteAclsResponseData()
+            .setThrottleTimeMs(throttleTimeMs)
+            .setFilterResults(filterResults), version());
     }
 
     public static DeleteAclsRequest parse(ByteBuffer buffer, short version) {
-        return new DeleteAclsRequest(DELETE_ACLS.parseRequest(version, buffer), version);
+        return new DeleteAclsRequest(new DeleteAclsRequestData(new ByteBufferAccessor(buffer), version), version);
+    }
+
+    public static DeleteAclsFilter deleteAclsFilter(AclBindingFilter filter) {
+        return new DeleteAclsFilter()
+            .setResourceNameFilter(filter.patternFilter().name())
+            .setResourceTypeFilter(filter.patternFilter().resourceType().code())
+            .setPatternTypeFilter(filter.patternFilter().patternType().code())
+            .setHostFilter(filter.entryFilter().host())
+            .setOperation(filter.entryFilter().operation().code())
+            .setPermissionType(filter.entryFilter().permissionType().code())
+            .setPrincipalFilter(filter.entryFilter().principal());
+    }
+
+    private static AclBindingFilter aclBindingFilter(DeleteAclsFilter filter) {
+        ResourcePatternFilter patternFilter = new ResourcePatternFilter(
+            ResourceType.fromCode(filter.resourceTypeFilter()),
+            filter.resourceNameFilter(),
+            PatternType.fromCode(filter.patternTypeFilter()));
+        AccessControlEntryFilter entryFilter = new AccessControlEntryFilter(
+            filter.principalFilter(),
+            filter.hostFilter(),
+            AclOperation.fromCode(filter.operation()),
+            AclPermissionType.fromCode(filter.permissionType()));
+        return new AclBindingFilter(patternFilter, entryFilter);
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/DeleteAclsResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/DeleteAclsResponse.java
index 7ae25da2c62e5..7482953a00d64 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/DeleteAclsResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/DeleteAclsResponse.java
@@ -19,186 +19,125 @@
 import org.apache.kafka.common.acl.AccessControlEntry;
 import org.apache.kafka.common.acl.AclBinding;
 import org.apache.kafka.common.protocol.ApiKeys;
-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.ByteBufferAccessor;
+import org.apache.kafka.common.resource.ResourcePattern;
+import org.apache.kafka.common.acl.AclOperation;
+import org.apache.kafka.common.acl.AclPermissionType;
+import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.message.DeleteAclsResponseData;
+import org.apache.kafka.common.message.DeleteAclsResponseData.DeleteAclsFilterResult;
+import org.apache.kafka.common.message.DeleteAclsResponseData.DeleteAclsMatchingAcl;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.Struct;
-import org.apache.kafka.common.resource.Resource;
-import org.apache.kafka.common.utils.Utils;
+import org.apache.kafka.common.resource.PatternType;
+import org.apache.kafka.common.resource.ResourceType;
+import org.apache.kafka.server.authorizer.AclDeleteResult;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 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.ERROR_MESSAGE;
-import static org.apache.kafka.common.protocol.CommonFields.HOST;
-import static org.apache.kafka.common.protocol.CommonFields.OPERATION;
-import static org.apache.kafka.common.protocol.CommonFields.PERMISSION_TYPE;
-import static org.apache.kafka.common.protocol.CommonFields.PRINCIPAL;
-import static org.apache.kafka.common.protocol.CommonFields.RESOURCE_NAME;
-import static org.apache.kafka.common.protocol.CommonFields.RESOURCE_TYPE;
-import static org.apache.kafka.common.protocol.CommonFields.THROTTLE_TIME_MS;
+import java.util.stream.Collectors;
 
 public class DeleteAclsResponse extends AbstractResponse {
     public static final Logger log = LoggerFactory.getLogger(DeleteAclsResponse.class);
-    private final static String FILTER_RESPONSES_KEY_NAME = "filter_responses";
-    private final static String MATCHING_ACLS_KEY_NAME = "matching_acls";
-
-    private static final Schema MATCHING_ACL = new Schema(
-            ERROR_CODE,
-            ERROR_MESSAGE,
-            RESOURCE_TYPE,
-            RESOURCE_NAME,
-            PRINCIPAL,
-            HOST,
-            OPERATION,
-            PERMISSION_TYPE);
-
-    private static final Schema DELETE_ACLS_RESPONSE_V0 = new Schema(
-            THROTTLE_TIME_MS,
-            new Field(FILTER_RESPONSES_KEY_NAME,
-                    new ArrayOf(new Schema(
-                            ERROR_CODE,
-                            ERROR_MESSAGE,
-                            new Field(MATCHING_ACLS_KEY_NAME, new ArrayOf(MATCHING_ACL), "The matching ACLs")))));
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{DELETE_ACLS_RESPONSE_V0};
-    }
-
-    public static class AclDeletionResult {
-        private final ApiError error;
-        private final AclBinding acl;
-
-        public AclDeletionResult(ApiError error, AclBinding acl) {
-            this.error = error;
-            this.acl = acl;
-        }
 
-        public AclDeletionResult(AclBinding acl) {
-            this(ApiError.NONE, acl);
-        }
-
-        public ApiError error() {
-            return error;
-        }
-
-        public AclBinding acl() {
-            return acl;
-        }
+    private final DeleteAclsResponseData data;
 
-        @Override
-        public String toString() {
-            return "(error=" + error + ", acl=" + acl + ")";
-        }
+    public DeleteAclsResponse(DeleteAclsResponseData data, short version) {
+        super(ApiKeys.DELETE_ACLS);
+        this.data = data;
+        validate(version);
     }
 
-    public static class AclFilterResponse {
-        private final ApiError error;
-        private final Collection deletions;
-
-        public AclFilterResponse(ApiError error, Collection deletions) {
-            this.error = error;
-            this.deletions = deletions;
-        }
-
-        public AclFilterResponse(Collection deletions) {
-            this(ApiError.NONE, deletions);
-        }
-
-        public ApiError error() {
-            return error;
-        }
-
-        public Collection deletions() {
-            return deletions;
-        }
+    @Override
+    public DeleteAclsResponseData data() {
+        return data;
+    }
 
-        @Override
-        public String toString() {
-            return "(error=" + error + ", deletions=" + Utils.join(deletions, ",") + ")";
-        }
+    @Override
+    public int throttleTimeMs() {
+        return data.throttleTimeMs();
     }
 
-    private final int throttleTimeMs;
+    public List filterResults() {
+        return data.filterResults();
+    }
 
-    private final List responses;
+    @Override
+    public Map errorCounts() {
+        return errorCounts(filterResults().stream().map(r -> Errors.forCode(r.errorCode())));
+    }
 
-    public DeleteAclsResponse(int throttleTimeMs, List responses) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.responses = responses;
+    public static DeleteAclsResponse parse(ByteBuffer buffer, short version) {
+        return new DeleteAclsResponse(new DeleteAclsResponseData(new ByteBufferAccessor(buffer), version), version);
     }
 
-    public DeleteAclsResponse(Struct struct) {
-        this.throttleTimeMs = struct.get(THROTTLE_TIME_MS);
-        this.responses = new ArrayList<>();
-        for (Object responseStructObj : struct.getArray(FILTER_RESPONSES_KEY_NAME)) {
-            Struct responseStruct = (Struct) responseStructObj;
-            ApiError error = new ApiError(responseStruct);
-            List deletions = new ArrayList<>();
-            for (Object matchingAclStructObj : responseStruct.getArray(MATCHING_ACLS_KEY_NAME)) {
-                Struct matchingAclStruct = (Struct) matchingAclStructObj;
-                ApiError matchError = new ApiError(matchingAclStruct);
-                AccessControlEntry entry = RequestUtils.aceFromStructFields(matchingAclStruct);
-                Resource resource = RequestUtils.resourceFromStructFields(matchingAclStruct);
-                deletions.add(new AclDeletionResult(matchError, new AclBinding(resource, entry)));
-            }
-            this.responses.add(new AclFilterResponse(error, deletions));
-        }
+    public String toString() {
+        return data.toString();
     }
 
     @Override
-    protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.DELETE_ACLS.responseSchema(version));
-        struct.set(THROTTLE_TIME_MS, throttleTimeMs);
-        List responseStructs = new ArrayList<>();
-        for (AclFilterResponse response : responses) {
-            Struct responseStruct = struct.instance(FILTER_RESPONSES_KEY_NAME);
-            response.error.write(responseStruct);
-            List deletionStructs = new ArrayList<>();
-            for (AclDeletionResult deletion : response.deletions()) {
-                Struct deletionStruct = responseStruct.instance(MATCHING_ACLS_KEY_NAME);
-                deletion.error.write(deletionStruct);
-                RequestUtils.resourceSetStructFields(deletion.acl().resource(), deletionStruct);
-                RequestUtils.aceSetStructFields(deletion.acl().entry(), deletionStruct);
-                deletionStructs.add(deletionStruct);
-            }
-            responseStruct.set(MATCHING_ACLS_KEY_NAME, deletionStructs.toArray(new Struct[0]));
-            responseStructs.add(responseStruct);
-        }
-        struct.set(FILTER_RESPONSES_KEY_NAME, responseStructs.toArray());
-        return struct;
+    public boolean shouldClientThrottle(short version) {
+        return version >= 1;
     }
 
-    public int throttleTimeMs() {
-        return throttleTimeMs;
+    private void validate(short version) {
+        if (version == 0) {
+            final boolean unsupported = filterResults().stream()
+                .flatMap(r -> r.matchingAcls().stream())
+                .anyMatch(matchingAcl -> matchingAcl.patternType() != PatternType.LITERAL.code());
+            if (unsupported)
+                throw new UnsupportedVersionException("Version 0 only supports literal resource pattern types");
+        }
+
+        final boolean unknown = filterResults().stream()
+                .flatMap(r -> r.matchingAcls().stream())
+                .anyMatch(matchingAcl -> matchingAcl.patternType() == PatternType.UNKNOWN.code()
+                    || matchingAcl.resourceType() == ResourceType.UNKNOWN.code()
+                    || matchingAcl.permissionType() == AclPermissionType.UNKNOWN.code()
+                    || matchingAcl.operation() == AclOperation.UNKNOWN.code());
+        if (unknown)
+            throw new IllegalArgumentException("DeleteAclsMatchingAcls contain UNKNOWN elements");
     }
 
-    public List responses() {
-        return responses;
+    public static DeleteAclsFilterResult filterResult(AclDeleteResult result) {
+        ApiError error = result.exception().map(e -> ApiError.fromThrowable(e)).orElse(ApiError.NONE);
+        List matchingAcls = result.aclBindingDeleteResults().stream()
+            .map(DeleteAclsResponse::matchingAcl)
+            .collect(Collectors.toList());
+        return new DeleteAclsFilterResult()
+            .setErrorCode(error.error().code())
+            .setErrorMessage(error.message())
+            .setMatchingAcls(matchingAcls);
     }
 
-    @Override
-    public Map errorCounts() {
-        Map errorCounts = new HashMap<>();
-        for (AclFilterResponse response : responses)
-            updateErrorCounts(errorCounts, response.error.error());
-        return errorCounts;
+    private static DeleteAclsMatchingAcl matchingAcl(AclDeleteResult.AclBindingDeleteResult result) {
+        ApiError error = result.exception().map(e -> ApiError.fromThrowable(e)).orElse(ApiError.NONE);
+        AclBinding acl = result.aclBinding();
+        return matchingAcl(acl, error);
     }
 
-    public static DeleteAclsResponse parse(ByteBuffer buffer, short version) {
-        return new DeleteAclsResponse(ApiKeys.DELETE_ACLS.responseSchema(version).read(buffer));
+    // Visible for testing
+    public static DeleteAclsMatchingAcl matchingAcl(AclBinding acl, ApiError error) {
+        return new DeleteAclsMatchingAcl()
+            .setErrorCode(error.error().code())
+            .setErrorMessage(error.message())
+            .setResourceName(acl.pattern().name())
+            .setResourceType(acl.pattern().resourceType().code())
+            .setPatternType(acl.pattern().patternType().code())
+            .setHost(acl.entry().host())
+            .setOperation(acl.entry().operation().code())
+            .setPermissionType(acl.entry().permissionType().code())
+            .setPrincipal(acl.entry().principal());
     }
 
-    public String toString() {
-        return "(responses=" + Utils.join(responses, ",") + ")";
+    public static AclBinding aclBinding(DeleteAclsMatchingAcl matchingAcl) {
+        ResourcePattern resourcePattern = new ResourcePattern(ResourceType.fromCode(matchingAcl.resourceType()),
+            matchingAcl.resourceName(), PatternType.fromCode(matchingAcl.patternType()));
+        AccessControlEntry accessControlEntry = new AccessControlEntry(matchingAcl.principal(), matchingAcl.host(),
+            AclOperation.fromCode(matchingAcl.operation()), AclPermissionType.fromCode(matchingAcl.permissionType()));
+        return new AclBinding(resourcePattern, accessControlEntry);
     }
 
 }
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
new file mode 100644
index 0000000000000..87d6deedc12c4
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/DeleteGroupsRequest.java
@@ -0,0 +1,81 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.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.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.Errors;
+
+import java.nio.ByteBuffer;
+
+public class DeleteGroupsRequest extends AbstractRequest {
+    public static class Builder extends AbstractRequest.Builder {
+        private final DeleteGroupsRequestData data;
+
+        public Builder(DeleteGroupsRequestData data) {
+            super(ApiKeys.DELETE_GROUPS);
+            this.data = data;
+        }
+
+        @Override
+        public DeleteGroupsRequest build(short version) {
+            return new DeleteGroupsRequest(data, version);
+        }
+
+        @Override
+        public String toString() {
+            return data.toString();
+        }
+    }
+
+    private final DeleteGroupsRequestData data;
+
+    public DeleteGroupsRequest(DeleteGroupsRequestData data, short version) {
+        super(ApiKeys.DELETE_GROUPS, version);
+        this.data = data;
+    }
+
+    @Override
+    public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
+        Errors error = Errors.forException(e);
+        DeletableGroupResultCollection groupResults = new DeletableGroupResultCollection();
+        for (String groupId : data.groupsNames()) {
+            groupResults.add(new DeletableGroupResult()
+                                 .setGroupId(groupId)
+                                 .setErrorCode(error.code()));
+        }
+
+        return new DeleteGroupsResponse(
+            new DeleteGroupsResponseData()
+                .setResults(groupResults)
+                .setThrottleTimeMs(throttleTimeMs)
+        );
+    }
+
+    public static DeleteGroupsRequest parse(ByteBuffer buffer, short version) {
+        return new DeleteGroupsRequest(new DeleteGroupsRequestData(new ByteBufferAccessor(buffer), version), version);
+    }
+
+    @Override
+    public DeleteGroupsRequestData data() {
+        return data;
+    }
+}
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
new file mode 100644
index 0000000000000..4cbffda422138
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/DeleteGroupsResponse.java
@@ -0,0 +1,92 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.Errors;
+
+import java.nio.ByteBuffer;
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * 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 final DeleteGroupsResponseData data;
+
+    public DeleteGroupsResponse(DeleteGroupsResponseData data) {
+        super(ApiKeys.DELETE_GROUPS);
+        this.data = data;
+    }
+
+    @Override
+    public DeleteGroupsResponseData data() {
+        return data;
+    }
+
+    public Map errors() {
+        Map errorMap = new HashMap<>();
+        for (DeletableGroupResult result : data.results()) {
+            errorMap.put(result.groupId(), Errors.forCode(result.errorCode()));
+        }
+        return errorMap;
+    }
+
+    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() {
+        Map counts = new HashMap<>();
+        data.results().forEach(result ->
+            updateErrorCounts(counts, Errors.forCode(result.errorCode()))
+        );
+        return counts;
+    }
+
+    public static DeleteGroupsResponse parse(ByteBuffer buffer, short version) {
+        return new DeleteGroupsResponse(new DeleteGroupsResponseData(new ByteBufferAccessor(buffer), version));
+    }
+
+    @Override
+    public int throttleTimeMs() {
+        return data.throttleTimeMs();
+    }
+
+    @Override
+    public boolean shouldClientThrottle(short version) {
+        return version >= 1;
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/DeleteRecordsRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/DeleteRecordsRequest.java
index 6467e06a2189c..a4f62e18f675a 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/DeleteRecordsRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/DeleteRecordsRequest.java
@@ -17,158 +17,69 @@
 
 package org.apache.kafka.common.requests;
 
-import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.message.DeleteRecordsRequestData;
+import org.apache.kafka.common.message.DeleteRecordsRequestData.DeleteRecordsTopic;
+import org.apache.kafka.common.message.DeleteRecordsResponseData;
+import org.apache.kafka.common.message.DeleteRecordsResponseData.DeleteRecordsTopicResult;
 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.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 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 org.apache.kafka.common.protocol.types.Type.INT64;
 
 public class DeleteRecordsRequest extends AbstractRequest {
 
     public static final long HIGH_WATERMARK = -1L;
 
-    // request level key names
-    private static final String TOPICS_KEY_NAME = "topics";
-    private static final String TIMEOUT_KEY_NAME = "timeout";
-
-    // topic level key names
-    private static final String PARTITIONS_KEY_NAME = "partitions";
-
-    // partition level key names
-    private static final String OFFSET_KEY_NAME = "offset";
-
-
-    private static final Schema DELETE_RECORDS_REQUEST_PARTITION_V0 = new Schema(
-            PARTITION_ID,
-            new Field(OFFSET_KEY_NAME, INT64, "The offset before which the messages will be deleted. -1 means high-watermark for the partition."));
-
-    private static final Schema DELETE_RECORDS_REQUEST_TOPIC_V0 = new Schema(
-            TOPIC_NAME,
-            new Field(PARTITIONS_KEY_NAME, new ArrayOf(DELETE_RECORDS_REQUEST_PARTITION_V0)));
-
-    private static final Schema DELETE_RECORDS_REQUEST_V0 = new Schema(
-            new Field(TOPICS_KEY_NAME, new ArrayOf(DELETE_RECORDS_REQUEST_TOPIC_V0)),
-            new Field(TIMEOUT_KEY_NAME, INT32, "The maximum time to await a response in ms."));
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{DELETE_RECORDS_REQUEST_V0};
-    }
-
-    private final int timeout;
-    private final Map partitionOffsets;
+    private final DeleteRecordsRequestData data;
 
     public static class Builder extends AbstractRequest.Builder {
-        private final int timeout;
-        private final Map partitionOffsets;
+        private DeleteRecordsRequestData data;
 
-        public Builder(int timeout, Map partitionOffsets) {
+        public Builder(DeleteRecordsRequestData data) {
             super(ApiKeys.DELETE_RECORDS);
-            this.timeout = timeout;
-            this.partitionOffsets = partitionOffsets;
+            this.data = data;
         }
 
         @Override
         public DeleteRecordsRequest build(short version) {
-            return new DeleteRecordsRequest(timeout, partitionOffsets, version);
+            return new DeleteRecordsRequest(data, version);
         }
 
         @Override
         public String toString() {
-            StringBuilder builder = new StringBuilder();
-            builder.append("(type=DeleteRecordsRequest")
-                   .append(", timeout=").append(timeout)
-                   .append(", partitionOffsets=(").append(partitionOffsets)
-                   .append("))");
-            return builder.toString();
+            return data.toString();
         }
     }
 
-
-    public DeleteRecordsRequest(Struct struct, short version) {
-        super(version);
-        partitionOffsets = new HashMap<>();
-        for (Object topicStructObj : struct.getArray(TOPICS_KEY_NAME)) {
-            Struct topicStruct = (Struct) topicStructObj;
-            String topic = topicStruct.get(TOPIC_NAME);
-            for (Object partitionStructObj : topicStruct.getArray(PARTITIONS_KEY_NAME)) {
-                Struct partitionStruct = (Struct) partitionStructObj;
-                int partition = partitionStruct.get(PARTITION_ID);
-                long offset = partitionStruct.getLong(OFFSET_KEY_NAME);
-                partitionOffsets.put(new TopicPartition(topic, partition), offset);
-            }
-        }
-        timeout = struct.getInt(TIMEOUT_KEY_NAME);
+    private DeleteRecordsRequest(DeleteRecordsRequestData data, short version) {
+        super(ApiKeys.DELETE_RECORDS, version);
+        this.data = data;
     }
 
-    public DeleteRecordsRequest(int timeout, Map partitionOffsets, short version) {
-        super(version);
-        this.timeout = timeout;
-        this.partitionOffsets = partitionOffsets;
-    }
     @Override
-    protected Struct toStruct() {
-        Struct struct = new Struct(ApiKeys.DELETE_RECORDS.requestSchema(version()));
-        Map> offsetsByTopic = CollectionUtils.groupDataByTopic(partitionOffsets);
-        struct.set(TIMEOUT_KEY_NAME, timeout);
-        List topicStructArray = new ArrayList<>();
-        for (Map.Entry> offsetsByTopicEntry : offsetsByTopic.entrySet()) {
-            Struct topicStruct = struct.instance(TOPICS_KEY_NAME);
-            topicStruct.set(TOPIC_NAME, offsetsByTopicEntry.getKey());
-            List partitionStructArray = new ArrayList<>();
-            for (Map.Entry offsetsByPartitionEntry : offsetsByTopicEntry.getValue().entrySet()) {
-                Struct partitionStruct = topicStruct.instance(PARTITIONS_KEY_NAME);
-                partitionStruct.set(PARTITION_ID, offsetsByPartitionEntry.getKey());
-                partitionStruct.set(OFFSET_KEY_NAME, offsetsByPartitionEntry.getValue());
-                partitionStructArray.add(partitionStruct);
-            }
-            topicStruct.set(PARTITIONS_KEY_NAME, partitionStructArray.toArray());
-            topicStructArray.add(topicStruct);
-        }
-        struct.set(TOPICS_KEY_NAME, topicStructArray.toArray());
-        return struct;
+    public DeleteRecordsRequestData data() {
+        return data;
     }
 
     @Override
     public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
-        Map responseMap = new HashMap<>();
-
-        for (Map.Entry entry : partitionOffsets.entrySet()) {
-            responseMap.put(entry.getKey(), new DeleteRecordsResponse.PartitionResponse(DeleteRecordsResponse.INVALID_LOW_WATERMARK, Errors.forException(e)));
-        }
-
-        short versionId = version();
-        switch (versionId) {
-            case 0:
-                return new DeleteRecordsResponse(throttleTimeMs, responseMap);
-            default:
-                throw new IllegalArgumentException(String.format("Version %d is not valid. Valid versions for %s are 0 to %d",
-                    versionId, this.getClass().getSimpleName(), ApiKeys.DELETE_RECORDS.latestVersion()));
+        DeleteRecordsResponseData result = new DeleteRecordsResponseData().setThrottleTimeMs(throttleTimeMs);
+        short errorCode = Errors.forException(e).code();
+        for (DeleteRecordsTopic topic : data.topics()) {
+            DeleteRecordsTopicResult topicResult = new DeleteRecordsTopicResult().setName(topic.name());
+            result.topics().add(topicResult);
+            for (DeleteRecordsRequestData.DeleteRecordsPartition partition : topic.partitions()) {
+                topicResult.partitions().add(new DeleteRecordsResponseData.DeleteRecordsPartitionResult()
+                        .setPartitionIndex(partition.partitionIndex())
+                        .setErrorCode(errorCode)
+                        .setLowWatermark(DeleteRecordsResponse.INVALID_LOW_WATERMARK));
+            }
         }
-    }
-
-    public int timeout() {
-        return timeout;
-    }
-
-    public Map partitionOffsets() {
-        return partitionOffsets;
+        return new DeleteRecordsResponse(result);
     }
 
     public static DeleteRecordsRequest parse(ByteBuffer buffer, short version) {
-        return new DeleteRecordsRequest(ApiKeys.DELETE_RECORDS.parseRequest(version, buffer), version);
+        return new DeleteRecordsRequest(new DeleteRecordsRequestData(new ByteBufferAccessor(buffer), version), version);
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/DeleteRecordsResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/DeleteRecordsResponse.java
index 4b4551738efd3..b090543faddfb 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/DeleteRecordsResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/DeleteRecordsResponse.java
@@ -17,158 +17,62 @@
 
 package org.apache.kafka.common.requests;
 
-import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.message.DeleteRecordsResponseData;
 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.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 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.protocol.types.Type.INT64;
-
 public class DeleteRecordsResponse extends AbstractResponse {
 
     public static final long INVALID_LOW_WATERMARK = -1L;
-
-    // request level key names
-    private static final String TOPICS_KEY_NAME = "topics";
-
-    // topic level key names
-    private static final String PARTITIONS_KEY_NAME = "partitions";
-
-    // partition level key names
-    private static final String LOW_WATERMARK_KEY_NAME = "low_watermark";
-
-    private static final Schema DELETE_RECORDS_RESPONSE_PARTITION_V0 = new Schema(
-            PARTITION_ID,
-            new Field(LOW_WATERMARK_KEY_NAME, INT64, "Smallest available offset of all live replicas"),
-            ERROR_CODE);
-
-    private static final Schema DELETE_RECORDS_RESPONSE_TOPIC_V0 = new Schema(
-            TOPIC_NAME,
-            new Field(PARTITIONS_KEY_NAME, new ArrayOf(DELETE_RECORDS_RESPONSE_PARTITION_V0)));
-
-    private static final Schema DELETE_RECORDS_RESPONSE_V0 = new Schema(
-            THROTTLE_TIME_MS,
-            new Field(TOPICS_KEY_NAME, new ArrayOf(DELETE_RECORDS_RESPONSE_TOPIC_V0)));
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{DELETE_RECORDS_RESPONSE_V0};
-    }
-
-    private final int throttleTimeMs;
-    private final Map responses;
+    private final DeleteRecordsResponseData data;
 
     /**
      * Possible error code:
      *
      * OFFSET_OUT_OF_RANGE (1)
      * UNKNOWN_TOPIC_OR_PARTITION (3)
-     * NOT_LEADER_FOR_PARTITION (6)
+     * NOT_LEADER_OR_FOLLOWER (6)
      * REQUEST_TIMED_OUT (7)
      * UNKNOWN (-1)
      */
 
-    public static final class PartitionResponse {
-        public long lowWatermark;
-        public Errors error;
-
-        public PartitionResponse(long lowWatermark, Errors error) {
-            this.lowWatermark = lowWatermark;
-            this.error = error;
-        }
-
-        @Override
-        public String toString() {
-            StringBuilder builder = new StringBuilder();
-            builder.append('{')
-                   .append(",low_watermark: ")
-                   .append(lowWatermark)
-                   .append("error: ")
-                   .append(error.toString())
-                   .append('}');
-            return builder.toString();
-        }
-    }
-
-    public DeleteRecordsResponse(Struct struct) {
-        this.throttleTimeMs = struct.getOrElse(THROTTLE_TIME_MS, DEFAULT_THROTTLE_TIME);
-        responses = new HashMap<>();
-        for (Object topicStructObj : struct.getArray(TOPICS_KEY_NAME)) {
-            Struct topicStruct = (Struct) topicStructObj;
-            String topic = topicStruct.get(TOPIC_NAME);
-            for (Object partitionStructObj : topicStruct.getArray(PARTITIONS_KEY_NAME)) {
-                Struct partitionStruct = (Struct) partitionStructObj;
-                int partition = partitionStruct.get(PARTITION_ID);
-                long lowWatermark = partitionStruct.getLong(LOW_WATERMARK_KEY_NAME);
-                Errors error = Errors.forCode(partitionStruct.get(ERROR_CODE));
-                responses.put(new TopicPartition(topic, partition), new PartitionResponse(lowWatermark, error));
-            }
-        }
-    }
-
-    /**
-     * Constructor for version 0.
-     */
-    public DeleteRecordsResponse(int throttleTimeMs, Map responses) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.responses = responses;
+    public DeleteRecordsResponse(DeleteRecordsResponseData data) {
+        super(ApiKeys.DELETE_RECORDS);
+        this.data = data;
     }
 
     @Override
-    protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.DELETE_RECORDS.responseSchema(version));
-        struct.setIfExists(THROTTLE_TIME_MS, throttleTimeMs);
-        Map> responsesByTopic = CollectionUtils.groupDataByTopic(responses);
-        List topicStructArray = new ArrayList<>();
-        for (Map.Entry> responsesByTopicEntry : responsesByTopic.entrySet()) {
-            Struct topicStruct = struct.instance(TOPICS_KEY_NAME);
-            topicStruct.set(TOPIC_NAME, responsesByTopicEntry.getKey());
-            List partitionStructArray = new ArrayList<>();
-            for (Map.Entry responsesByPartitionEntry : responsesByTopicEntry.getValue().entrySet()) {
-                Struct partitionStruct = topicStruct.instance(PARTITIONS_KEY_NAME);
-                PartitionResponse response = responsesByPartitionEntry.getValue();
-                partitionStruct.set(PARTITION_ID, responsesByPartitionEntry.getKey());
-                partitionStruct.set(LOW_WATERMARK_KEY_NAME, response.lowWatermark);
-                partitionStruct.set(ERROR_CODE, response.error.code());
-                partitionStructArray.add(partitionStruct);
-            }
-            topicStruct.set(PARTITIONS_KEY_NAME, partitionStructArray.toArray());
-            topicStructArray.add(topicStruct);
-        }
-        struct.set(TOPICS_KEY_NAME, topicStructArray.toArray());
-        return struct;
+    public DeleteRecordsResponseData data() {
+        return data;
     }
 
+    @Override
     public int throttleTimeMs() {
-        return throttleTimeMs;
-    }
-
-    public Map responses() {
-        return this.responses;
+        return data.throttleTimeMs();
     }
 
     @Override
     public Map errorCounts() {
         Map errorCounts = new HashMap<>();
-        for (PartitionResponse response : responses.values())
-            updateErrorCounts(errorCounts, response.error);
+        data.topics().forEach(topicResponses ->
+            topicResponses.partitions().forEach(response ->
+                updateErrorCounts(errorCounts, Errors.forCode(response.errorCode()))
+            )
+        );
         return errorCounts;
     }
 
     public static DeleteRecordsResponse parse(ByteBuffer buffer, short version) {
-        return new DeleteRecordsResponse(ApiKeys.DELETE_RECORDS.responseSchema(version).read(buffer));
+        return new DeleteRecordsResponse(new DeleteRecordsResponseData(new ByteBufferAccessor(buffer), version));
+    }
+
+    @Override
+    public boolean shouldClientThrottle(short version) {
+        return version >= 1;
     }
 }
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 4696b503147ac..dfd2e72df03d5 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,121 +16,64 @@
  */
 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 org.apache.kafka.common.protocol.ByteBufferAccessor;
 
 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;
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{DELETE_TOPICS_REQUEST_V0, DELETE_TOPICS_REQUEST_V1};
-    }
-
-    private final Set topics;
-    private final Integer timeout;
 
     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) {
-        super(version);
-        this.topics = topics;
-        this.timeout = timeout;
-    }
-
-    public DeleteTopicsRequest(Struct struct, short version) {
-        super(version);
-        Object[] topicsArray = struct.getArray(TOPICS_KEY_NAME);
-        Set topics = new HashSet<>(topicsArray.length);
-        for (Object topic : topicsArray)
-            topics.add((String) topic);
+    private DeleteTopicsRequestData data;
 
-        this.topics = topics;
-        this.timeout = struct.getInt(TIMEOUT_KEY_NAME);
+    private DeleteTopicsRequest(DeleteTopicsRequestData data, short version) {
+        super(ApiKeys.DELETE_TOPICS, version);
+        this.data = data;
     }
 
     @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;
+    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:
-                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) {
-        return new DeleteTopicsRequest(ApiKeys.DELETE_TOPICS.parseRequest(version, buffer), version);
+        return new DeleteTopicsRequest(new DeleteTopicsRequestData(new ByteBufferAccessor(buffer), version), 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 d13e4412b6c4d..2090c4fd2e2ee 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,42 +16,17 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.message.DeleteTopicsResponseData;
 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.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."));
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{DELETE_TOPICS_RESPONSE_V0, DELETE_TOPICS_RESPONSE_V1};
-    }
-
 
     /**
      * Possible error codes:
@@ -60,62 +35,41 @@ public static Schema[] schemaVersions() {
      * INVALID_TOPIC_EXCEPTION(17)
      * TOPIC_AUTHORIZATION_FAILED(29)
      * NOT_CONTROLLER(41)
+     * INVALID_REQUEST(42)
+     * TOPIC_DELETION_DISABLED(73)
      */
-    private final Map errors;
-    private final int throttleTimeMs;
-
-    public DeleteTopicsResponse(Map errors) {
-        this(DEFAULT_THROTTLE_TIME, errors);
-    }
+    private final DeleteTopicsResponseData 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(DeleteTopicsResponseData data) {
+        super(ApiKeys.DELETE_TOPICS);
+        this.data = data;
     }
 
     @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;
-    }
-
     public int throttleTimeMs() {
-        return throttleTimeMs;
+        return data.throttleTimeMs();
     }
 
-    public Map errors() {
-        return errors;
+    @Override
+    public DeleteTopicsResponseData data() {
+        return data;
     }
 
     @Override
     public Map errorCounts() {
-        return errorCounts(errors);
+        HashMap counts = new HashMap<>();
+        data.responses().forEach(result ->
+            updateErrorCounts(counts, Errors.forCode(result.errorCode()))
+        );
+        return counts;
     }
 
     public static DeleteTopicsResponse parse(ByteBuffer buffer, short version) {
-        return new DeleteTopicsResponse(ApiKeys.DELETE_TOPICS.responseSchema(version).read(buffer));
+        return new DeleteTopicsResponse(new DeleteTopicsResponseData(new ByteBufferAccessor(buffer), version));
+    }
+
+    @Override
+    public boolean shouldClientThrottle(short version) {
+        return version >= 2;
     }
 }
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 1bacac7c7f52c..1ddf5bf99fc89 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
@@ -17,95 +17,108 @@
 package org.apache.kafka.common.requests;
 
 import org.apache.kafka.common.acl.AccessControlEntryFilter;
-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.acl.AclPermissionType;
+import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.message.DescribeAclsRequestData;
+import org.apache.kafka.common.message.DescribeAclsResponseData;
 import org.apache.kafka.common.protocol.ApiKeys;
-import org.apache.kafka.common.protocol.types.Schema;
-import org.apache.kafka.common.protocol.types.Struct;
-import org.apache.kafka.common.resource.ResourceFilter;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.resource.PatternType;
+import org.apache.kafka.common.resource.ResourcePatternFilter;
+import org.apache.kafka.common.resource.ResourceType;
 
 import java.nio.ByteBuffer;
-import java.util.Collections;
-
-import static org.apache.kafka.common.protocol.CommonFields.HOST_FILTER;
-import static org.apache.kafka.common.protocol.CommonFields.OPERATION;
-import static org.apache.kafka.common.protocol.CommonFields.PERMISSION_TYPE;
-import static org.apache.kafka.common.protocol.CommonFields.PRINCIPAL_FILTER;
-import static org.apache.kafka.common.protocol.CommonFields.RESOURCE_NAME_FILTER;
-import static org.apache.kafka.common.protocol.CommonFields.RESOURCE_TYPE;
 
 public class DescribeAclsRequest extends AbstractRequest {
-    private static final Schema DESCRIBE_ACLS_REQUEST_V0 = new Schema(
-            RESOURCE_TYPE,
-            RESOURCE_NAME_FILTER,
-            PRINCIPAL_FILTER,
-            HOST_FILTER,
-            OPERATION,
-            PERMISSION_TYPE);
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{DESCRIBE_ACLS_REQUEST_V0};
-    }
 
     public static class Builder extends AbstractRequest.Builder {
-        private final AclBindingFilter filter;
+        private final DescribeAclsRequestData data;
 
         public Builder(AclBindingFilter filter) {
             super(ApiKeys.DESCRIBE_ACLS);
-            this.filter = filter;
+            ResourcePatternFilter patternFilter = filter.patternFilter();
+            AccessControlEntryFilter entryFilter = filter.entryFilter();
+            data = new DescribeAclsRequestData()
+                .setHostFilter(entryFilter.host())
+                .setOperation(entryFilter.operation().code())
+                .setPermissionType(entryFilter.permissionType().code())
+                .setPrincipalFilter(entryFilter.principal())
+                .setResourceNameFilter(patternFilter.name())
+                .setPatternTypeFilter(patternFilter.patternType().code())
+                .setResourceTypeFilter(patternFilter.resourceType().code());
         }
 
         @Override
         public DescribeAclsRequest build(short version) {
-            return new DescribeAclsRequest(filter, version);
+            return new DescribeAclsRequest(data, version);
         }
 
         @Override
         public String toString() {
-            return "(type=DescribeAclsRequest, filter=" + filter + ")";
+            return data.toString();
         }
     }
 
-    private final AclBindingFilter filter;
+    private final DescribeAclsRequestData data;
 
-    DescribeAclsRequest(AclBindingFilter filter, short version) {
-        super(version);
-        this.filter = filter;
+    private DescribeAclsRequest(DescribeAclsRequestData data, short version) {
+        super(ApiKeys.DESCRIBE_ACLS, version);
+        this.data = data;
+        normalizeAndValidate(version);
     }
 
-    public DescribeAclsRequest(Struct struct, short version) {
-        super(version);
-        ResourceFilter resourceFilter = RequestUtils.resourceFilterFromStructFields(struct);
-        AccessControlEntryFilter entryFilter = RequestUtils.aceFilterFromStructFields(struct);
-        this.filter = new AclBindingFilter(resourceFilter, entryFilter);
+    private void normalizeAndValidate(short version) {
+        if (version == 0) {
+            PatternType patternType = PatternType.fromCode(data.patternTypeFilter());
+            // On older brokers, no pattern types existed except LITERAL (effectively). So even though ANY is not
+            // directly supported on those brokers, we can get the same effect as ANY by setting the pattern type
+            // to LITERAL. Note that the wildcard `*` is considered `LITERAL` for compatibility reasons.
+            if (patternType == PatternType.ANY)
+                data.setPatternTypeFilter(PatternType.LITERAL.code());
+            else if (patternType != PatternType.LITERAL)
+                throw new UnsupportedVersionException("Version 0 only supports literal resource pattern types");
+        }
+
+        if (data.patternTypeFilter() == PatternType.UNKNOWN.code()
+                || data.resourceTypeFilter() == ResourceType.UNKNOWN.code()
+                || data.permissionType() == AclPermissionType.UNKNOWN.code()
+                || data.operation() == AclOperation.UNKNOWN.code()) {
+            throw new IllegalArgumentException("DescribeAclsRequest contains UNKNOWN elements: " + data);
+        }
     }
 
     @Override
-    protected Struct toStruct() {
-        Struct struct = new Struct(ApiKeys.DESCRIBE_ACLS.requestSchema(version()));
-        RequestUtils.resourceFilterSetStructFields(filter.resourceFilter(), struct);
-        RequestUtils.aceFilterSetStructFields(filter.entryFilter(), struct);
-        return struct;
+    public DescribeAclsRequestData data() {
+        return data;
     }
 
     @Override
     public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable throwable) {
-        short versionId = version();
-        switch (versionId) {
-            case 0:
-                return new DescribeAclsResponse(throttleTimeMs, ApiError.fromThrowable(throwable),
-                        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.DESCRIBE_ACLS.latestVersion()));
-        }
+        ApiError error = ApiError.fromThrowable(throwable);
+        DescribeAclsResponseData response = new DescribeAclsResponseData()
+            .setThrottleTimeMs(throttleTimeMs)
+            .setErrorCode(error.error().code())
+            .setErrorMessage(error.message());
+        return new DescribeAclsResponse(response, version());
     }
 
     public static DescribeAclsRequest parse(ByteBuffer buffer, short version) {
-        return new DescribeAclsRequest(ApiKeys.DESCRIBE_ACLS.parseRequest(version, buffer), version);
+        return new DescribeAclsRequest(new DescribeAclsRequestData(new ByteBufferAccessor(buffer), version), version);
     }
 
     public AclBindingFilter filter() {
-        return filter;
+        ResourcePatternFilter rpf = new ResourcePatternFilter(
+                ResourceType.fromCode(data.resourceTypeFilter()),
+                data.resourceNameFilter(),
+                PatternType.fromCode(data.patternTypeFilter()));
+        AccessControlEntryFilter acef =  new AccessControlEntryFilter(
+                data.principalFilter(),
+                data.hostFilter(),
+                AclOperation.fromCode(data.operation()),
+                AclPermissionType.fromCode(data.permissionType()));
+        return new AclBindingFilter(rpf, acef);
     }
+
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/DescribeAclsResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/DescribeAclsResponse.java
index a21230bfc39e3..c4190e65640ea 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/DescribeAclsResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/DescribeAclsResponse.java
@@ -20,12 +20,11 @@
 import org.apache.kafka.common.acl.AccessControlEntry;
 import org.apache.kafka.common.acl.AclBinding;
 import org.apache.kafka.common.protocol.ApiKeys;
-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.ByteBufferAccessor;
+import org.apache.kafka.common.resource.PatternType;
+import org.apache.kafka.common.resource.ResourcePattern;
+import org.apache.kafka.common.errors.UnsupportedVersionException;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.Struct;
-import org.apache.kafka.common.resource.Resource;
 
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
@@ -33,117 +32,130 @@
 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.ERROR_MESSAGE;
-import static org.apache.kafka.common.protocol.CommonFields.HOST;
-import static org.apache.kafka.common.protocol.CommonFields.OPERATION;
-import static org.apache.kafka.common.protocol.CommonFields.PERMISSION_TYPE;
-import static org.apache.kafka.common.protocol.CommonFields.PRINCIPAL;
-import static org.apache.kafka.common.protocol.CommonFields.RESOURCE_NAME;
-import static org.apache.kafka.common.protocol.CommonFields.RESOURCE_TYPE;
-import static org.apache.kafka.common.protocol.CommonFields.THROTTLE_TIME_MS;
+import java.util.Map.Entry;
+import java.util.Optional;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+import org.apache.kafka.common.acl.AclOperation;
+import org.apache.kafka.common.acl.AclPermissionType;
+import org.apache.kafka.common.message.DescribeAclsResponseData;
+import org.apache.kafka.common.message.DescribeAclsResponseData.AclDescription;
+import org.apache.kafka.common.message.DescribeAclsResponseData.DescribeAclsResource;
+import org.apache.kafka.common.resource.ResourceType;
 
 public class DescribeAclsResponse extends AbstractResponse {
-    private final static String RESOURCES_KEY_NAME = "resources";
-    private final static String ACLS_KEY_NAME = "acls";
-
-    private static final Schema DESCRIBE_ACLS_RESOURCE = new Schema(
-            RESOURCE_TYPE,
-            RESOURCE_NAME,
-            new Field(ACLS_KEY_NAME, new ArrayOf(new Schema(
-                    PRINCIPAL,
-                    HOST,
-                    OPERATION,
-                    PERMISSION_TYPE))));
-
-    private static final Schema DESCRIBE_ACLS_RESPONSE_V0 = new Schema(
-            THROTTLE_TIME_MS,
-            ERROR_CODE,
-            ERROR_MESSAGE,
-            new Field(RESOURCES_KEY_NAME, new ArrayOf(DESCRIBE_ACLS_RESOURCE), "The resources and their associated ACLs."));
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{DESCRIBE_ACLS_RESPONSE_V0};
-    }
 
-    private final int throttleTimeMs;
-    private final ApiError error;
-    private final Collection acls;
+    private final DescribeAclsResponseData data;
 
-    public DescribeAclsResponse(int throttleTimeMs, ApiError error, Collection acls) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.error = error;
-        this.acls = acls;
+    public DescribeAclsResponse(DescribeAclsResponseData data, short version) {
+        super(ApiKeys.DESCRIBE_ACLS);
+        this.data = data;
+        validate(Optional.of(version));
     }
 
-    public DescribeAclsResponse(Struct struct) {
-        this.throttleTimeMs = struct.get(THROTTLE_TIME_MS);
-        this.error = new ApiError(struct);
-        this.acls = new ArrayList<>();
-        for (Object resourceStructObj : struct.getArray(RESOURCES_KEY_NAME)) {
-            Struct resourceStruct = (Struct) resourceStructObj;
-            Resource resource = RequestUtils.resourceFromStructFields(resourceStruct);
-            for (Object aclDataStructObj : resourceStruct.getArray(ACLS_KEY_NAME)) {
-                Struct aclDataStruct = (Struct) aclDataStructObj;
-                AccessControlEntry entry = RequestUtils.aceFromStructFields(aclDataStruct);
-                this.acls.add(new AclBinding(resource, entry));
-            }
-        }
+    // Skips version validation, visible for testing
+    DescribeAclsResponse(DescribeAclsResponseData data) {
+        super(ApiKeys.DESCRIBE_ACLS);
+        this.data = data;
+        validate(Optional.empty());
     }
 
     @Override
-    protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.DESCRIBE_ACLS.responseSchema(version));
-        struct.set(THROTTLE_TIME_MS, throttleTimeMs);
-        error.write(struct);
-
-        Map> resourceToData = new HashMap<>();
-        for (AclBinding acl : acls) {
-            List entry = resourceToData.get(acl.resource());
-            if (entry == null) {
-                entry = new ArrayList<>();
-                resourceToData.put(acl.resource(), entry);
-            }
-            entry.add(acl.entry());
-        }
-
-        List resourceStructs = new ArrayList<>();
-        for (Map.Entry> tuple : resourceToData.entrySet()) {
-            Resource resource = tuple.getKey();
-            Struct resourceStruct = struct.instance(RESOURCES_KEY_NAME);
-            RequestUtils.resourceSetStructFields(resource, resourceStruct);
-            List dataStructs = new ArrayList<>();
-            for (AccessControlEntry entry : tuple.getValue()) {
-                Struct dataStruct = resourceStruct.instance(ACLS_KEY_NAME);
-                RequestUtils.aceSetStructFields(entry, dataStruct);
-                dataStructs.add(dataStruct);
-            }
-            resourceStruct.set(ACLS_KEY_NAME, dataStructs.toArray());
-            resourceStructs.add(resourceStruct);
-        }
-        struct.set(RESOURCES_KEY_NAME, resourceStructs.toArray());
-        return struct;
+    public DescribeAclsResponseData data() {
+        return data;
     }
 
+    @Override
     public int throttleTimeMs() {
-        return throttleTimeMs;
+        return data.throttleTimeMs();
     }
 
     public ApiError error() {
-        return error;
+        return new ApiError(Errors.forCode(data.errorCode()), data.errorMessage());
     }
 
     @Override
     public Map errorCounts() {
-        return errorCounts(error.error());
+        return errorCounts(Errors.forCode(data.errorCode()));
     }
 
-    public Collection acls() {
-        return acls;
+    public List acls() {
+        return data.resources();
     }
 
     public static DescribeAclsResponse parse(ByteBuffer buffer, short version) {
-        return new DescribeAclsResponse(ApiKeys.DESCRIBE_ACLS.responseSchema(version).read(buffer));
+        return new DescribeAclsResponse(new DescribeAclsResponseData(new ByteBufferAccessor(buffer), version), version);
+    }
+
+    @Override
+    public boolean shouldClientThrottle(short version) {
+        return version >= 1;
+    }
+
+    private void validate(Optional version) {
+        if (version.isPresent() && version.get() == 0) {
+            final boolean unsupported = acls().stream()
+                .anyMatch(acl -> acl.patternType() != PatternType.LITERAL.code());
+            if (unsupported) {
+                throw new UnsupportedVersionException("Version 0 only supports literal resource pattern types");
+            }
+        }
+
+        for (DescribeAclsResource resource : acls()) {
+            if (resource.patternType() == PatternType.UNKNOWN.code() || resource.resourceType() == ResourceType.UNKNOWN.code())
+                throw new IllegalArgumentException("Contain UNKNOWN elements");
+            for (AclDescription acl : resource.acls()) {
+                if (acl.operation() == AclOperation.UNKNOWN.code() || acl.permissionType() == AclPermissionType.UNKNOWN.code()) {
+                    throw new IllegalArgumentException("Contain UNKNOWN elements");
+                }
+            }
+        }
+    }
+
+    private static Stream aclBindings(DescribeAclsResource resource) {
+        return resource.acls().stream().map(acl -> {
+            ResourcePattern pattern = new ResourcePattern(
+                    ResourceType.fromCode(resource.resourceType()),
+                    resource.resourceName(),
+                    PatternType.fromCode(resource.patternType()));
+            AccessControlEntry entry = new AccessControlEntry(
+                    acl.principal(),
+                    acl.host(),
+                    AclOperation.fromCode(acl.operation()),
+                    AclPermissionType.fromCode(acl.permissionType()));
+            return new AclBinding(pattern, entry);
+        });
+    }
+
+    public static List aclBindings(List resources) {
+        return resources.stream().flatMap(DescribeAclsResponse::aclBindings).collect(Collectors.toList());
+    }
+
+    public static List aclsResources(Collection acls) {
+        Map> patternToEntries = new HashMap<>();
+        for (AclBinding acl : acls) {
+            patternToEntries.computeIfAbsent(acl.pattern(), v -> new ArrayList<>()).add(acl.entry());
+        }
+        List resources = new ArrayList<>(patternToEntries.size());
+        for (Entry> entry : patternToEntries.entrySet()) {
+            ResourcePattern key = entry.getKey();
+            List aclDescriptions = new ArrayList<>();
+            for (AccessControlEntry ace : entry.getValue()) {
+                AclDescription ad = new AclDescription()
+                    .setHost(ace.host())
+                    .setOperation(ace.operation().code())
+                    .setPermissionType(ace.permissionType().code())
+                    .setPrincipal(ace.principal());
+                aclDescriptions.add(ad);
+            }
+            DescribeAclsResource dar = new DescribeAclsResource()
+                .setResourceName(key.name())
+                .setPatternType(key.patternType().code())
+                .setResourceType(key.resourceType().code())
+                .setAcls(aclDescriptions);
+            resources.add(dar);
+        }
+        return resources;
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/DescribeClientQuotasRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/DescribeClientQuotasRequest.java
new file mode 100644
index 0000000000000..68f4ecf9c0702
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/DescribeClientQuotasRequest.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.message.DescribeClientQuotasRequestData;
+import org.apache.kafka.common.message.DescribeClientQuotasRequestData.ComponentData;
+import org.apache.kafka.common.message.DescribeClientQuotasResponseData;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.quota.ClientQuotaFilter;
+import org.apache.kafka.common.quota.ClientQuotaFilterComponent;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.List;
+
+public class DescribeClientQuotasRequest extends AbstractRequest {
+    // These values must not change.
+    private static final byte MATCH_TYPE_EXACT = 0;
+    private static final byte MATCH_TYPE_DEFAULT = 1;
+    private static final byte MATCH_TYPE_SPECIFIED = 2;
+
+    public static class Builder extends AbstractRequest.Builder {
+
+        private final DescribeClientQuotasRequestData data;
+
+        public Builder(ClientQuotaFilter filter) {
+            super(ApiKeys.DESCRIBE_CLIENT_QUOTAS);
+
+            List componentData = new ArrayList<>(filter.components().size());
+            for (ClientQuotaFilterComponent component : filter.components()) {
+                ComponentData fd = new ComponentData().setEntityType(component.entityType());
+                if (component.match() == null) {
+                    fd.setMatchType(MATCH_TYPE_SPECIFIED);
+                    fd.setMatch(null);
+                } else if (component.match().isPresent()) {
+                    fd.setMatchType(MATCH_TYPE_EXACT);
+                    fd.setMatch(component.match().get());
+                } else {
+                    fd.setMatchType(MATCH_TYPE_DEFAULT);
+                    fd.setMatch(null);
+                }
+                componentData.add(fd);
+            }
+            this.data = new DescribeClientQuotasRequestData()
+                .setComponents(componentData)
+                .setStrict(filter.strict());
+        }
+
+        @Override
+        public DescribeClientQuotasRequest build(short version) {
+            return new DescribeClientQuotasRequest(data, version);
+        }
+
+        @Override
+        public String toString() {
+            return data.toString();
+        }
+    }
+
+    private final DescribeClientQuotasRequestData data;
+
+    public DescribeClientQuotasRequest(DescribeClientQuotasRequestData data, short version) {
+        super(ApiKeys.DESCRIBE_CLIENT_QUOTAS, version);
+        this.data = data;
+    }
+
+    public ClientQuotaFilter filter() {
+        List components = new ArrayList<>(data.components().size());
+        for (ComponentData componentData : data.components()) {
+            ClientQuotaFilterComponent component;
+            switch (componentData.matchType()) {
+                case MATCH_TYPE_EXACT:
+                    component = ClientQuotaFilterComponent.ofEntity(componentData.entityType(), componentData.match());
+                    break;
+                case MATCH_TYPE_DEFAULT:
+                    component = ClientQuotaFilterComponent.ofDefaultEntity(componentData.entityType());
+                    break;
+                case MATCH_TYPE_SPECIFIED:
+                    component = ClientQuotaFilterComponent.ofEntityType(componentData.entityType());
+                    break;
+                default:
+                    throw new IllegalArgumentException("Unexpected match type: " + componentData.matchType());
+            }
+            components.add(component);
+        }
+        if (data.strict()) {
+            return ClientQuotaFilter.containsOnly(components);
+        } else {
+            return ClientQuotaFilter.contains(components);
+        }
+    }
+
+    @Override
+    public DescribeClientQuotasRequestData data() {
+        return data;
+    }
+
+    @Override
+    public DescribeClientQuotasResponse getErrorResponse(int throttleTimeMs, Throwable e) {
+        ApiError error = ApiError.fromThrowable(e);
+        return new DescribeClientQuotasResponse(new DescribeClientQuotasResponseData()
+            .setThrottleTimeMs(throttleTimeMs)
+            .setErrorCode(error.error().code())
+            .setErrorMessage(error.message())
+            .setEntries(null));
+    }
+
+    public static DescribeClientQuotasRequest parse(ByteBuffer buffer, short version) {
+        return new DescribeClientQuotasRequest(new DescribeClientQuotasRequestData(new ByteBufferAccessor(buffer), version),
+            version);
+    }
+
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/DescribeClientQuotasResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/DescribeClientQuotasResponse.java
new file mode 100644
index 0000000000000..3a81e21dac65e
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/DescribeClientQuotasResponse.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.common.requests;
+
+import org.apache.kafka.common.internals.KafkaFutureImpl;
+import org.apache.kafka.common.message.DescribeClientQuotasResponseData;
+import org.apache.kafka.common.message.DescribeClientQuotasResponseData.EntityData;
+import org.apache.kafka.common.message.DescribeClientQuotasResponseData.EntryData;
+import org.apache.kafka.common.message.DescribeClientQuotasResponseData.ValueData;
+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.quota.ClientQuotaEntity;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+public class DescribeClientQuotasResponse extends AbstractResponse {
+
+    private final DescribeClientQuotasResponseData data;
+
+    public DescribeClientQuotasResponse(DescribeClientQuotasResponseData data) {
+        super(ApiKeys.DESCRIBE_CLIENT_QUOTAS);
+        this.data = data;
+    }
+
+    public void complete(KafkaFutureImpl>> future) {
+        Errors error = Errors.forCode(data.errorCode());
+        if (error != Errors.NONE) {
+            future.completeExceptionally(error.exception(data.errorMessage()));
+            return;
+        }
+
+        Map> result = new HashMap<>(data.entries().size());
+        for (EntryData entries : data.entries()) {
+            Map entity = new HashMap<>(entries.entity().size());
+            for (EntityData entityData : entries.entity()) {
+                entity.put(entityData.entityType(), entityData.entityName());
+            }
+
+            Map values = new HashMap<>(entries.values().size());
+            for (ValueData valueData : entries.values()) {
+                values.put(valueData.key(), valueData.value());
+            }
+
+            result.put(new ClientQuotaEntity(entity), values);
+        }
+        future.complete(result);
+    }
+
+    @Override
+    public int throttleTimeMs() {
+        return data.throttleTimeMs();
+    }
+
+    @Override
+    public DescribeClientQuotasResponseData data() {
+        return data;
+    }
+
+    @Override
+    public Map errorCounts() {
+        return errorCounts(Errors.forCode(data.errorCode()));
+    }
+
+    public static DescribeClientQuotasResponse parse(ByteBuffer buffer, short version) {
+        return new DescribeClientQuotasResponse(new DescribeClientQuotasResponseData(new ByteBufferAccessor(buffer), version));
+    }
+
+    public static DescribeClientQuotasResponse fromQuotaEntities(Map> entities,
+                                                                 int throttleTimeMs) {
+        List entries = new ArrayList<>(entities.size());
+        for (Map.Entry> entry : entities.entrySet()) {
+            ClientQuotaEntity quotaEntity = entry.getKey();
+            List entityData = new ArrayList<>(quotaEntity.entries().size());
+            for (Map.Entry entityEntry : quotaEntity.entries().entrySet()) {
+                entityData.add(new EntityData()
+                        .setEntityType(entityEntry.getKey())
+                        .setEntityName(entityEntry.getValue()));
+            }
+
+            Map quotaValues = entry.getValue();
+            List valueData = new ArrayList<>(quotaValues.size());
+            for (Map.Entry valuesEntry : entry.getValue().entrySet()) {
+                valueData.add(new ValueData()
+                        .setKey(valuesEntry.getKey())
+                        .setValue(valuesEntry.getValue()));
+            }
+
+            entries.add(new EntryData()
+                    .setEntity(entityData)
+                    .setValues(valueData));
+        }
+
+        return new DescribeClientQuotasResponse(new DescribeClientQuotasResponseData()
+            .setThrottleTimeMs(throttleTimeMs)
+            .setErrorCode((short) 0)
+            .setErrorMessage(null)
+            .setEntries(entries));
+    }
+
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/DescribeConfigsRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/DescribeConfigsRequest.java
index 74e25f4763a1c..d612ca80949a0 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/DescribeConfigsRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/DescribeConfigsRequest.java
@@ -16,145 +16,58 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.message.DescribeConfigsRequestData;
+import org.apache.kafka.common.message.DescribeConfigsResponseData;
 import org.apache.kafka.common.protocol.ApiKeys;
-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.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.Errors;
 
 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.protocol.types.Type.INT8;
-import static org.apache.kafka.common.protocol.types.Type.STRING;
+import java.util.stream.Collectors;
 
 public class DescribeConfigsRequest extends AbstractRequest {
 
-    private static final String RESOURCES_KEY_NAME = "resources";
-    private static final String RESOURCE_TYPE_KEY_NAME = "resource_type";
-    private static final String RESOURCE_NAME_KEY_NAME = "resource_name";
-    private static final String CONFIG_NAMES_KEY_NAME = "config_names";
-
-    private static final Schema DESCRIBE_CONFIGS_REQUEST_RESOURCE_V0 = new Schema(
-            new Field(RESOURCE_TYPE_KEY_NAME, INT8),
-            new Field(RESOURCE_NAME_KEY_NAME, STRING),
-            new Field(CONFIG_NAMES_KEY_NAME, ArrayOf.nullable(STRING)));
-
-    private static final Schema DESCRIBE_CONFIGS_REQUEST_V0 = new Schema(
-            new Field(RESOURCES_KEY_NAME, new ArrayOf(DESCRIBE_CONFIGS_REQUEST_RESOURCE_V0), "An array of config resources to be returned."));
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{DESCRIBE_CONFIGS_REQUEST_V0};
-    }
-
-    public static class Builder extends AbstractRequest.Builder {
-        private final Map> resourceToConfigNames;
+    public static class Builder extends AbstractRequest.Builder {
+        private final DescribeConfigsRequestData data;
 
-        public Builder(Map> resourceToConfigNames) {
+        public Builder(DescribeConfigsRequestData data) {
             super(ApiKeys.DESCRIBE_CONFIGS);
-            this.resourceToConfigNames = resourceToConfigNames;
-        }
-
-        public Builder(Collection resources) {
-            this(toResourceToConfigNames(resources));
-        }
-
-        private static Map> toResourceToConfigNames(Collection resources) {
-            Map> result = new HashMap<>(resources.size());
-            for (Resource resource : resources)
-                result.put(resource, null);
-            return result;
+            this.data = data;
         }
 
         @Override
         public DescribeConfigsRequest build(short version) {
-            return new DescribeConfigsRequest(version, resourceToConfigNames);
-        }
-    }
-
-    private final Map> resourceToConfigNames;
-
-    public DescribeConfigsRequest(short version, Map> resourceToConfigNames) {
-        super(version);
-        this.resourceToConfigNames = resourceToConfigNames;
-
-    }
-
-    public DescribeConfigsRequest(Struct struct, short version) {
-        super(version);
-        Object[] resourcesArray = struct.getArray(RESOURCES_KEY_NAME);
-        resourceToConfigNames = new HashMap<>(resourcesArray.length);
-        for (Object resourceObj : resourcesArray) {
-            Struct resourceStruct = (Struct) resourceObj;
-            ResourceType resourceType = ResourceType.forId(resourceStruct.getByte(RESOURCE_TYPE_KEY_NAME));
-            String resourceName = resourceStruct.getString(RESOURCE_NAME_KEY_NAME);
-
-            Object[] configNamesArray = resourceStruct.getArray(CONFIG_NAMES_KEY_NAME);
-            List configNames = null;
-            if (configNamesArray != null) {
-                configNames = new ArrayList<>(configNamesArray.length);
-                for (Object configNameObj : configNamesArray)
-                    configNames.add((String) configNameObj);
-            }
-
-            resourceToConfigNames.put(new Resource(resourceType, resourceName), configNames);
+            return new DescribeConfigsRequest(data, version);
         }
     }
 
-    public Collection resources() {
-        return resourceToConfigNames.keySet();
-    }
+    private final DescribeConfigsRequestData data;
 
-    /**
-     * Return null if all config names should be returned.
-     */
-    public Collection configNames(Resource resource) {
-        return resourceToConfigNames.get(resource);
+    public DescribeConfigsRequest(DescribeConfigsRequestData data, short version) {
+        super(ApiKeys.DESCRIBE_CONFIGS, version);
+        this.data = data;
     }
 
     @Override
-    protected Struct toStruct() {
-        Struct struct = new Struct(ApiKeys.DESCRIBE_CONFIGS.requestSchema(version()));
-        List resourceStructs = new ArrayList<>(resources().size());
-        for (Map.Entry> entry : resourceToConfigNames.entrySet()) {
-            Resource resource = entry.getKey();
-            Struct resourceStruct = struct.instance(RESOURCES_KEY_NAME);
-            resourceStruct.set(RESOURCE_TYPE_KEY_NAME, resource.type().id());
-            resourceStruct.set(RESOURCE_NAME_KEY_NAME, resource.name());
-
-            String[] configNames = entry.getValue() == null ? null : entry.getValue().toArray(new String[0]);
-            resourceStruct.set(CONFIG_NAMES_KEY_NAME, configNames);
-
-            resourceStructs.add(resourceStruct);
-        }
-        struct.set(RESOURCES_KEY_NAME, resourceStructs.toArray(new Struct[0]));
-        return struct;
+    public DescribeConfigsRequestData data() {
+        return data;
     }
 
     @Override
     public DescribeConfigsResponse getErrorResponse(int throttleTimeMs, Throwable e) {
-        short version = version();
-        switch (version) {
-            case 0:
-                ApiError error = ApiError.fromThrowable(e);
-                Map errors = new HashMap<>(resources().size());
-                DescribeConfigsResponse.Config config = new DescribeConfigsResponse.Config(error,
-                        Collections.emptyList());
-                for (Resource resource : resources())
-                    errors.put(resource, config);
-                return new DescribeConfigsResponse(throttleTimeMs, errors);
-            default:
-                throw new IllegalArgumentException(String.format("Version %d is not valid. Valid versions for %s are 0 to %d",
-                        version, this.getClass().getSimpleName(), ApiKeys.DESCRIBE_CONFIGS.latestVersion()));
-        }
+        Errors error = Errors.forException(e);
+        return new DescribeConfigsResponse(new DescribeConfigsResponseData()
+                .setThrottleTimeMs(throttleTimeMs)
+                .setResults(data.resources().stream().map(result -> {
+                    return new DescribeConfigsResponseData.DescribeConfigsResult().setErrorCode(error.code())
+                            .setErrorMessage(error.message())
+                            .setResourceName(result.resourceName())
+                            .setResourceType(result.resourceType());
+                }).collect(Collectors.toList())
+        ));
     }
 
     public static DescribeConfigsRequest parse(ByteBuffer buffer, short version) {
-        return new DescribeConfigsRequest(ApiKeys.DESCRIBE_CONFIGS.parseRequest(version, buffer), version);
+        return new DescribeConfigsRequest(new DescribeConfigsRequestData(new ByteBufferAccessor(buffer), version), version);
     }
 }
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 91bf30e2311f4..aa7a713e8ab4e 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
@@ -17,70 +17,29 @@
 
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.config.ConfigResource;
+import org.apache.kafka.common.message.DescribeConfigsResponseData;
 import org.apache.kafka.common.protocol.ApiKeys;
-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.ByteBufferAccessor;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.protocol.types.Struct;
 
 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.ERROR_MESSAGE;
-import static org.apache.kafka.common.protocol.CommonFields.THROTTLE_TIME_MS;
-import static org.apache.kafka.common.protocol.types.Type.BOOLEAN;
-import static org.apache.kafka.common.protocol.types.Type.INT8;
-import static org.apache.kafka.common.protocol.types.Type.NULLABLE_STRING;
-import static org.apache.kafka.common.protocol.types.Type.STRING;
+import java.util.Objects;
+import java.util.function.Function;
+import java.util.stream.Collectors;
 
 public class DescribeConfigsResponse extends AbstractResponse {
 
-    private static final String RESOURCES_KEY_NAME = "resources";
-
-    private static final String RESOURCE_TYPE_KEY_NAME = "resource_type";
-    private static final String RESOURCE_NAME_KEY_NAME = "resource_name";
-
-    private static final String CONFIG_ENTRIES_KEY_NAME = "config_entries";
-
-    private static final String CONFIG_NAME_KEY_NAME = "config_name";
-    private static final String CONFIG_VALUE_KEY_NAME = "config_value";
-    private static final String IS_SENSITIVE_KEY_NAME = "is_sensitive";
-    private static final String IS_DEFAULT_KEY_NAME = "is_default";
-    private static final String READ_ONLY_KEY_NAME = "read_only";
-
-    private static final Schema DESCRIBE_CONFIGS_RESPONSE_ENTITY_V0 = new Schema(
-            ERROR_CODE,
-            ERROR_MESSAGE,
-            new Field(RESOURCE_TYPE_KEY_NAME, INT8),
-            new Field(RESOURCE_NAME_KEY_NAME, STRING),
-            new Field(CONFIG_ENTRIES_KEY_NAME, new ArrayOf(new Schema(
-                    new Field(CONFIG_NAME_KEY_NAME, STRING),
-                    new Field(CONFIG_VALUE_KEY_NAME, NULLABLE_STRING),
-                    new Field(READ_ONLY_KEY_NAME, BOOLEAN),
-                    new Field(IS_DEFAULT_KEY_NAME, BOOLEAN),
-                    new Field(IS_SENSITIVE_KEY_NAME, BOOLEAN)))));
-
-    private static final Schema DESCRIBE_CONFIGS_RESPONSE_V0 = new Schema(
-            THROTTLE_TIME_MS,
-            new Field(RESOURCES_KEY_NAME, new ArrayOf(DESCRIBE_CONFIGS_RESPONSE_ENTITY_V0)));
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{DESCRIBE_CONFIGS_RESPONSE_V0};
-    }
-
     public static class Config {
         private final ApiError error;
         private final Collection entries;
 
         public Config(ApiError error, Collection entries) {
-            this.error = error;
-            this.entries = entries;
+            this.error = Objects.requireNonNull(error, "error");
+            this.entries = Objects.requireNonNull(entries, "entries");
         }
 
         public ApiError error() {
@@ -96,15 +55,28 @@ public static class ConfigEntry {
         private final String name;
         private final String value;
         private final boolean isSensitive;
-        private final boolean isDefault;
+        private final ConfigSource source;
         private final boolean readOnly;
+        private final Collection synonyms;
+        private final ConfigType type;
+        private final String documentation;
 
-        public ConfigEntry(String name, String value, boolean isSensitive, boolean isDefault, boolean readOnly) {
-            this.name = name;
+        public ConfigEntry(String name, String value, ConfigSource source, boolean isSensitive, boolean readOnly,
+            Collection synonyms) {
+            this(name, value, source, isSensitive, readOnly, synonyms, ConfigType.UNKNOWN, null);
+        }
+
+        public ConfigEntry(String name, String value, ConfigSource source, boolean isSensitive, boolean readOnly,
+                           Collection synonyms, ConfigType type, String documentation) {
+
+            this.name = Objects.requireNonNull(name, "name");
             this.value = value;
+            this.source = Objects.requireNonNull(source, "source");
             this.isSensitive = isSensitive;
-            this.isDefault = isDefault;
             this.readOnly = readOnly;
+            this.synonyms = Objects.requireNonNull(synonyms, "synonyms");
+            this.type = type;
+            this.documentation = documentation;
         }
 
         public String name() {
@@ -119,106 +91,185 @@ public boolean isSensitive() {
             return isSensitive;
         }
 
-        public boolean isDefault() {
-            return isDefault;
+        public ConfigSource source() {
+            return source;
         }
 
         public boolean isReadOnly() {
             return readOnly;
         }
+
+        public Collection synonyms() {
+            return synonyms;
+        }
+
+        public ConfigType type() {
+            return type;
+        }
+
+        public String documentation() {
+            return documentation;
+        }
+    }
+
+    public enum ConfigSource {
+        UNKNOWN((byte) 0, org.apache.kafka.clients.admin.ConfigEntry.ConfigSource.UNKNOWN),
+        TOPIC_CONFIG((byte) 1, org.apache.kafka.clients.admin.ConfigEntry.ConfigSource.DYNAMIC_TOPIC_CONFIG),
+        DYNAMIC_BROKER_CONFIG((byte) 2, org.apache.kafka.clients.admin.ConfigEntry.ConfigSource.DYNAMIC_BROKER_CONFIG),
+        DYNAMIC_DEFAULT_BROKER_CONFIG((byte) 3, org.apache.kafka.clients.admin.ConfigEntry.ConfigSource.DYNAMIC_DEFAULT_BROKER_CONFIG),
+        STATIC_BROKER_CONFIG((byte) 4, org.apache.kafka.clients.admin.ConfigEntry.ConfigSource.STATIC_BROKER_CONFIG),
+        DEFAULT_CONFIG((byte) 5, org.apache.kafka.clients.admin.ConfigEntry.ConfigSource.DEFAULT_CONFIG),
+        DYNAMIC_BROKER_LOGGER_CONFIG((byte) 6, org.apache.kafka.clients.admin.ConfigEntry.ConfigSource.DYNAMIC_BROKER_LOGGER_CONFIG);
+
+        final byte id;
+        private final org.apache.kafka.clients.admin.ConfigEntry.ConfigSource source;
+        private static final ConfigSource[] VALUES = values();
+
+        ConfigSource(byte id, org.apache.kafka.clients.admin.ConfigEntry.ConfigSource source) {
+            this.id = id;
+            this.source = source;
+        }
+
+        public byte id() {
+            return id;
+        }
+
+        public static ConfigSource forId(byte id) {
+            if (id < 0)
+                throw new IllegalArgumentException("id should be positive, id: " + id);
+            if (id >= VALUES.length)
+                return UNKNOWN;
+            return VALUES[id];
+        }
+
+        public org.apache.kafka.clients.admin.ConfigEntry.ConfigSource source() {
+            return source;
+        }
     }
 
-    private final int throttleTimeMs;
-    private final Map configs;
+    public enum ConfigType {
+        UNKNOWN((byte) 0, org.apache.kafka.clients.admin.ConfigEntry.ConfigType.UNKNOWN),
+        BOOLEAN((byte) 1, org.apache.kafka.clients.admin.ConfigEntry.ConfigType.BOOLEAN),
+        STRING((byte) 2, org.apache.kafka.clients.admin.ConfigEntry.ConfigType.STRING),
+        INT((byte) 3, org.apache.kafka.clients.admin.ConfigEntry.ConfigType.INT),
+        SHORT((byte) 4, org.apache.kafka.clients.admin.ConfigEntry.ConfigType.SHORT),
+        LONG((byte) 5, org.apache.kafka.clients.admin.ConfigEntry.ConfigType.LONG),
+        DOUBLE((byte) 6, org.apache.kafka.clients.admin.ConfigEntry.ConfigType.DOUBLE),
+        LIST((byte) 7, org.apache.kafka.clients.admin.ConfigEntry.ConfigType.LIST),
+        CLASS((byte) 8, org.apache.kafka.clients.admin.ConfigEntry.ConfigType.CLASS),
+        PASSWORD((byte) 9, org.apache.kafka.clients.admin.ConfigEntry.ConfigType.PASSWORD);
+
+        final byte id;
+        final org.apache.kafka.clients.admin.ConfigEntry.ConfigType type;
+        private static final ConfigType[] VALUES = values();
+
+        ConfigType(byte id, org.apache.kafka.clients.admin.ConfigEntry.ConfigType type) {
+            this.id = id;
+            this.type = type;
+        }
+
+        public byte id() {
+            return id;
+        }
 
-    public DescribeConfigsResponse(int throttleTimeMs, Map configs) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.configs = configs;
+        public static ConfigType forId(byte id) {
+            if (id < 0)
+                throw new IllegalArgumentException("id should be positive, id: " + id);
+            if (id >= VALUES.length)
+                return UNKNOWN;
+            return VALUES[id];
+        }
+
+        public org.apache.kafka.clients.admin.ConfigEntry.ConfigType type() {
+            return type;
+        }
     }
 
-    public DescribeConfigsResponse(Struct struct) {
-        throttleTimeMs = struct.get(THROTTLE_TIME_MS);
-        Object[] resourcesArray = struct.getArray(RESOURCES_KEY_NAME);
-        configs = new HashMap<>(resourcesArray.length);
-        for (Object resourceObj : resourcesArray) {
-            Struct resourceStruct = (Struct) resourceObj;
-
-            ApiError error = new ApiError(resourceStruct);
-            ResourceType resourceType = ResourceType.forId(resourceStruct.getByte(RESOURCE_TYPE_KEY_NAME));
-            String resourceName = resourceStruct.getString(RESOURCE_NAME_KEY_NAME);
-            Resource resource = new Resource(resourceType, resourceName);
-
-            Object[] configEntriesArray = resourceStruct.getArray(CONFIG_ENTRIES_KEY_NAME);
-            List configEntries = new ArrayList<>(configEntriesArray.length);
-            for (Object configEntriesObj: configEntriesArray) {
-                Struct configEntriesStruct = (Struct) configEntriesObj;
-                String configName = configEntriesStruct.getString(CONFIG_NAME_KEY_NAME);
-                String configValue = configEntriesStruct.getString(CONFIG_VALUE_KEY_NAME);
-                boolean isSensitive = configEntriesStruct.getBoolean(IS_SENSITIVE_KEY_NAME);
-                boolean isDefault = configEntriesStruct.getBoolean(IS_DEFAULT_KEY_NAME);
-                boolean readOnly = configEntriesStruct.getBoolean(READ_ONLY_KEY_NAME);
-                configEntries.add(new ConfigEntry(configName, configValue, isSensitive, isDefault, readOnly));
-            }
-            Config config = new Config(error, configEntries);
-            configs.put(resource, config);
+    public static class ConfigSynonym {
+        private final String name;
+        private final String value;
+        private final ConfigSource source;
+
+        public ConfigSynonym(String name, String value, ConfigSource source) {
+            this.name = Objects.requireNonNull(name, "name");
+            this.value = value;
+            this.source = Objects.requireNonNull(source, "source");
+        }
+
+        public String name() {
+            return name;
+        }
+        public String value() {
+            return value;
         }
+        public ConfigSource source() {
+            return source;
+        }
+    }
+
+    public Map resultMap() {
+        return data().results().stream().collect(Collectors.toMap(
+            configsResult ->
+                    new ConfigResource(ConfigResource.Type.forId(configsResult.resourceType()),
+                            configsResult.resourceName()),
+            Function.identity()));
     }
 
-    public Map configs() {
-        return configs;
+    private final DescribeConfigsResponseData data;
+
+    public DescribeConfigsResponse(DescribeConfigsResponseData data) {
+        super(ApiKeys.DESCRIBE_CONFIGS);
+        this.data = data;
     }
 
-    public Config config(Resource resource) {
-        return configs.get(resource);
+    // This constructor should only be used after deserialization, it has special handling for version 0
+    private DescribeConfigsResponse(DescribeConfigsResponseData data, short version) {
+        super(ApiKeys.DESCRIBE_CONFIGS);
+        this.data = data;
+        if (version == 0) {
+            for (DescribeConfigsResponseData.DescribeConfigsResult result : data.results()) {
+                for (DescribeConfigsResponseData.DescribeConfigsResourceResult config : result.configs()) {
+                    if (config.isDefault()) {
+                        config.setConfigSource(ConfigSource.DEFAULT_CONFIG.id);
+                    } else {
+                        if (result.resourceType() == ConfigResource.Type.BROKER.id()) {
+                            config.setConfigSource(ConfigSource.STATIC_BROKER_CONFIG.id);
+                        } else if (result.resourceType() == ConfigResource.Type.TOPIC.id()) {
+                            config.setConfigSource(ConfigSource.TOPIC_CONFIG.id);
+                        } else {
+                            config.setConfigSource(ConfigSource.UNKNOWN.id);
+                        }
+                    }
+                }
+            }
+        }
     }
 
+    @Override
+    public DescribeConfigsResponseData data() {
+        return data;
+    }
+
+    @Override
     public int throttleTimeMs() {
-        return throttleTimeMs;
+        return data.throttleTimeMs();
     }
 
     @Override
     public Map errorCounts() {
         Map errorCounts = new HashMap<>();
-        for (Config response : configs.values())
-            updateErrorCounts(errorCounts, response.error.error());
+        data.results().forEach(response ->
+            updateErrorCounts(errorCounts, Errors.forCode(response.errorCode()))
+        );
         return errorCounts;
     }
 
-    @Override
-    protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.DESCRIBE_CONFIGS.responseSchema(version));
-        struct.set(THROTTLE_TIME_MS, throttleTimeMs);
-        List resourceStructs = new ArrayList<>(configs.size());
-        for (Map.Entry entry : configs.entrySet()) {
-            Struct resourceStruct = struct.instance(RESOURCES_KEY_NAME);
-
-            Resource resource = entry.getKey();
-            resourceStruct.set(RESOURCE_TYPE_KEY_NAME, resource.type().id());
-            resourceStruct.set(RESOURCE_NAME_KEY_NAME, resource.name());
-
-            Config config = entry.getValue();
-            config.error.write(resourceStruct);
-
-            List configEntryStructs = new ArrayList<>(config.entries.size());
-            for (ConfigEntry configEntry : config.entries) {
-                Struct configEntriesStruct = resourceStruct.instance(CONFIG_ENTRIES_KEY_NAME);
-                configEntriesStruct.set(CONFIG_NAME_KEY_NAME, configEntry.name);
-                configEntriesStruct.set(CONFIG_VALUE_KEY_NAME, configEntry.value);
-                configEntriesStruct.set(IS_SENSITIVE_KEY_NAME, configEntry.isSensitive);
-                configEntriesStruct.set(IS_DEFAULT_KEY_NAME, configEntry.isDefault);
-                configEntriesStruct.set(READ_ONLY_KEY_NAME, configEntry.readOnly);
-                configEntryStructs.add(configEntriesStruct);
-            }
-            resourceStruct.set(CONFIG_ENTRIES_KEY_NAME, configEntryStructs.toArray(new Struct[0]));
-            
-            resourceStructs.add(resourceStruct);
-        }
-        struct.set(RESOURCES_KEY_NAME, resourceStructs.toArray(new Struct[0]));
-        return struct;
-    }
-
     public static DescribeConfigsResponse parse(ByteBuffer buffer, short version) {
-        return new DescribeConfigsResponse(ApiKeys.DESCRIBE_CONFIGS.parseResponse(version, buffer));
+        return new DescribeConfigsResponse(new DescribeConfigsResponseData(new ByteBufferAccessor(buffer), version), version);
     }
 
+    @Override
+    public boolean shouldClientThrottle(short version) {
+        return version >= 2;
+    }
 }
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
new file mode 100644
index 0000000000000..9bf59e844a6c3
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/DescribeDelegationTokenRequest.java
@@ -0,0 +1,81 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.DescribeDelegationTokenRequestData;
+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.security.auth.KafkaPrincipal;
+
+import java.nio.ByteBuffer;
+import java.util.List;
+import java.util.stream.Collectors;
+
+public class DescribeDelegationTokenRequest extends AbstractRequest {
+
+    public static class Builder extends AbstractRequest.Builder {
+        private final DescribeDelegationTokenRequestData data;
+
+        public Builder(List owners) {
+            super(ApiKeys.DESCRIBE_DELEGATION_TOKEN);
+            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(data, version);
+        }
+
+        @Override
+        public String toString() {
+            return data.toString();
+        }
+    }
+
+    private final DescribeDelegationTokenRequestData data;
+
+    public DescribeDelegationTokenRequest(DescribeDelegationTokenRequestData data, short version) {
+        super(ApiKeys.DESCRIBE_DELEGATION_TOKEN, version);
+        this.data = data;
+    }
+
+    @Override
+    public DescribeDelegationTokenRequestData data() {
+        return data;
+    }
+
+    public boolean ownersListEmpty() {
+        return data.owners() != null && data.owners().isEmpty();
+    }
+
+    @Override
+    public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
+        return new DescribeDelegationTokenResponse(throttleTimeMs, Errors.forException(e));
+    }
+
+    public static DescribeDelegationTokenRequest parse(ByteBuffer buffer, short version) {
+        return new DescribeDelegationTokenRequest(new DescribeDelegationTokenRequestData(
+            new ByteBufferAccessor(buffer), version), version);
+    }
+}
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
new file mode 100644
index 0000000000000..4a2162f53aaef
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/DescribeDelegationTokenResponse.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.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.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.Errors;
+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 java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+public class DescribeDelegationTokenResponse extends AbstractResponse {
+
+    private final DescribeDelegationTokenResponseData data;
+
+    public DescribeDelegationTokenResponse(int throttleTimeMs, Errors error, List tokens) {
+        super(ApiKeys.DESCRIBE_DELEGATION_TOKEN);
+        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(DescribeDelegationTokenResponseData data) {
+        super(ApiKeys.DESCRIBE_DELEGATION_TOKEN);
+        this.data = data;
+    }
+
+    public static DescribeDelegationTokenResponse parse(ByteBuffer buffer, short version) {
+        return new DescribeDelegationTokenResponse(new DescribeDelegationTokenResponseData(
+            new ByteBufferAccessor(buffer), version));
+    }
+
+    @Override
+    public Map errorCounts() {
+        return errorCounts(error());
+    }
+
+    @Override
+    public DescribeDelegationTokenResponseData data() {
+        return data;
+    }
+
+    @Override
+    public int throttleTimeMs() {
+        return data.throttleTimeMs();
+    }
+
+    public Errors error() {
+        return Errors.forCode(data.errorCode());
+    }
+
+    public List 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 error() != Errors.NONE;
+    }
+
+    @Override
+    public boolean shouldClientThrottle(short version) {
+        return version >= 1;
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/DescribeGroupsRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/DescribeGroupsRequest.java
index 56117dad53051..eff5bb9fff1ec 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/DescribeGroupsRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/DescribeGroupsRequest.java
@@ -16,95 +16,57 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.message.DescribeGroupsRequestData;
 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.Struct;
-import org.apache.kafka.common.utils.Utils;
 
 import java.nio.ByteBuffer;
-import java.util.ArrayList;
-import java.util.List;
 
-import static org.apache.kafka.common.protocol.types.Type.STRING;
+import static org.apache.kafka.common.requests.AbstractResponse.DEFAULT_THROTTLE_TIME;
 
 public class DescribeGroupsRequest extends AbstractRequest {
-    private static final String GROUP_IDS_KEY_NAME = "group_ids";
-
-    /* Describe group api */
-    private static final Schema DESCRIBE_GROUPS_REQUEST_V0 = new Schema(
-            new Field(GROUP_IDS_KEY_NAME, new ArrayOf(STRING), "List of groupIds to request metadata for (an " +
-                    "empty groupId array will return empty group metadata)."));
-
-    /* v1 request is the same as v0. Throttle time has been added to response */
-    private static final Schema DESCRIBE_GROUPS_REQUEST_V1 = DESCRIBE_GROUPS_REQUEST_V0;
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{DESCRIBE_GROUPS_REQUEST_V0, DESCRIBE_GROUPS_REQUEST_V1};
-    }
-
     public static class Builder extends AbstractRequest.Builder {
-        private final List groupIds;
+        private final DescribeGroupsRequestData data;
 
-        public Builder(List groupIds) {
+        public Builder(DescribeGroupsRequestData data) {
             super(ApiKeys.DESCRIBE_GROUPS);
-            this.groupIds = groupIds;
+            this.data = data;
         }
 
         @Override
         public DescribeGroupsRequest build(short version) {
-            return new DescribeGroupsRequest(this.groupIds, version);
+            return new DescribeGroupsRequest(data, version);
         }
 
         @Override
         public String toString() {
-            return "(type=DescribeGroupsRequest, groupIds=(" + Utils.join(groupIds, ",") + "))";
+            return data.toString();
         }
     }
 
-    private final List groupIds;
+    private final DescribeGroupsRequestData data;
 
-    private DescribeGroupsRequest(List groupIds, short version) {
-        super(version);
-        this.groupIds = groupIds;
-    }
-
-    public DescribeGroupsRequest(Struct struct, short version) {
-        super(version);
-        this.groupIds = new ArrayList<>();
-        for (Object groupId : struct.getArray(GROUP_IDS_KEY_NAME))
-            this.groupIds.add((String) groupId);
-    }
-
-    public List groupIds() {
-        return groupIds;
+    private DescribeGroupsRequest(DescribeGroupsRequestData data, short version) {
+        super(ApiKeys.DESCRIBE_GROUPS, version);
+        this.data = data;
     }
 
     @Override
-    protected Struct toStruct() {
-        Struct struct = new Struct(ApiKeys.DESCRIBE_GROUPS.requestSchema(version()));
-        struct.set(GROUP_IDS_KEY_NAME, groupIds.toArray());
-        return struct;
+    public DescribeGroupsRequestData data() {
+        return data;
     }
 
     @Override
     public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
-        short version = version();
-        switch (version) {
-            case 0:
-                return DescribeGroupsResponse.fromError(Errors.forException(e), groupIds);
-            case 1:
-                return DescribeGroupsResponse.fromError(throttleTimeMs, Errors.forException(e), groupIds);
-
-            default:
-                throw new IllegalArgumentException(String.format("Version %d is not valid. Valid versions for %s are 0 to %d",
-                        version, this.getClass().getSimpleName(), ApiKeys.DESCRIBE_GROUPS.latestVersion()));
+        if (version() == 0) {
+            return DescribeGroupsResponse.fromError(DEFAULT_THROTTLE_TIME, Errors.forException(e), data.groups());
+        } else {
+            return DescribeGroupsResponse.fromError(throttleTimeMs, Errors.forException(e), data.groups());
         }
     }
 
     public static DescribeGroupsRequest parse(ByteBuffer buffer, short version) {
-        return new DescribeGroupsRequest(ApiKeys.DESCRIBE_GROUPS.parseRequest(version, buffer), version);
+        return new DescribeGroupsRequest(new DescribeGroupsRequestData(new ByteBufferAccessor(buffer), version), version);
     }
 }
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 174b14bc7442b..360caf01e468e 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
@@ -16,74 +16,24 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.message.DescribeGroupsResponseData;
+import org.apache.kafka.common.message.DescribeGroupsResponseData.DescribedGroup;
+import org.apache.kafka.common.message.DescribeGroupsResponseData.DescribedGroupMember;
 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.Struct;
+import org.apache.kafka.common.utils.Utils;
 
 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.GROUP_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;
+import java.util.Set;
 
 public class DescribeGroupsResponse extends AbstractResponse {
 
-    private static final String GROUPS_KEY_NAME = "groups";
-
-    private static final String GROUP_STATE_KEY_NAME = "state";
-    private static final String PROTOCOL_TYPE_KEY_NAME = "protocol_type";
-    private static final String PROTOCOL_KEY_NAME = "protocol";
-
-    private static final String MEMBERS_KEY_NAME = "members";
-    private static final String CLIENT_ID_KEY_NAME = "client_id";
-    private static final String CLIENT_HOST_KEY_NAME = "client_host";
-    private static final String MEMBER_METADATA_KEY_NAME = "member_metadata";
-    private static final String MEMBER_ASSIGNMENT_KEY_NAME = "member_assignment";
-
-    private static final Schema DESCRIBE_GROUPS_RESPONSE_MEMBER_V0 = new Schema(
-            MEMBER_ID,
-            new Field(CLIENT_ID_KEY_NAME, STRING, "The client id used in the member's latest join group request"),
-            new Field(CLIENT_HOST_KEY_NAME, STRING, "The client host used in the request session corresponding to the " +
-                    "member's join group."),
-            new Field(MEMBER_METADATA_KEY_NAME, BYTES, "The metadata corresponding to the current group protocol in " +
-                    "use (will only be present if the group is stable)."),
-            new Field(MEMBER_ASSIGNMENT_KEY_NAME, BYTES, "The current assignment provided by the group leader " +
-                    "(will only be present if the group is stable)."));
-
-    private static final Schema DESCRIBE_GROUPS_RESPONSE_GROUP_METADATA_V0 = new Schema(
-            ERROR_CODE,
-            GROUP_ID,
-            new Field(GROUP_STATE_KEY_NAME, STRING, "The current state of the group (one of: Dead, Stable, CompletingRebalance, " +
-                    "PreparingRebalance, or empty if there is no active group)"),
-            new Field(PROTOCOL_TYPE_KEY_NAME, STRING, "The current group protocol type (will be empty if there is no active group)"),
-            new Field(PROTOCOL_KEY_NAME, STRING, "The current group protocol (only provided if the group is Stable)"),
-            new Field(MEMBERS_KEY_NAME, new ArrayOf(DESCRIBE_GROUPS_RESPONSE_MEMBER_V0), "Current group members " +
-                    "(only provided if the group is not Dead)"));
-
-    private static final Schema DESCRIBE_GROUPS_RESPONSE_V0 = new Schema(
-            new Field(GROUPS_KEY_NAME, new ArrayOf(DESCRIBE_GROUPS_RESPONSE_GROUP_METADATA_V0)));
-    private static final Schema DESCRIBE_GROUPS_RESPONSE_V1 = new Schema(
-            THROTTLE_TIME_MS,
-            new Field(GROUPS_KEY_NAME, new ArrayOf(DESCRIBE_GROUPS_RESPONSE_GROUP_METADATA_V0)));
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {DESCRIBE_GROUPS_RESPONSE_V0, DESCRIBE_GROUPS_RESPONSE_V1};
-    }
-
-    public static final String UNKNOWN_STATE = "";
-    public static final String UNKNOWN_PROTOCOL_TYPE = "";
-    public static final String UNKNOWN_PROTOCOL = "";
+    public static final int AUTHORIZED_OPERATIONS_OMITTED = Integer.MIN_VALUE;
 
     /**
      * Possible per-group error codes:
@@ -94,195 +44,108 @@ public static Schema[] schemaVersions() {
      * AUTHORIZATION_FAILED (29)
      */
 
-    private final Map groups;
-    private final int throttleTimeMs;
+    private final DescribeGroupsResponseData data;
 
-    public DescribeGroupsResponse(Map groups) {
-        this(DEFAULT_THROTTLE_TIME, groups);
+    public DescribeGroupsResponse(DescribeGroupsResponseData data) {
+        super(ApiKeys.DESCRIBE_GROUPS);
+        this.data = data;
     }
 
-    public DescribeGroupsResponse(int throttleTimeMs, Map groups) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.groups = groups;
+    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)
+            .setMemberMetadata(metadata);
     }
 
-    public DescribeGroupsResponse(Struct struct) {
-        this.throttleTimeMs = struct.getOrElse(THROTTLE_TIME_MS, DEFAULT_THROTTLE_TIME);
-        this.groups = new HashMap<>();
-        for (Object groupObj : struct.getArray(GROUPS_KEY_NAME)) {
-            Struct groupStruct = (Struct) groupObj;
+    public static DescribedGroup groupMetadata(
+        final String groupId,
+        final Errors error,
+        final String state,
+        final String protocolType,
+        final String protocol,
+        final List members,
+        final Set authorizedOperations) {
+        DescribedGroup groupMetadata = new DescribedGroup();
+        groupMetadata.setGroupId(groupId)
+            .setErrorCode(error.code())
+            .setGroupState(state)
+            .setProtocolType(protocolType)
+            .setProtocolData(protocol)
+            .setMembers(members)
+            .setAuthorizedOperations(Utils.to32BitField(authorizedOperations));
+        return  groupMetadata;
+    }
 
-            String groupId = groupStruct.get(GROUP_ID);
-            Errors error = Errors.forCode(groupStruct.get(ERROR_CODE));
-            String state = groupStruct.getString(GROUP_STATE_KEY_NAME);
-            String protocolType = groupStruct.getString(PROTOCOL_TYPE_KEY_NAME);
-            String protocol = groupStruct.getString(PROTOCOL_KEY_NAME);
+    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 groupMetadata = new DescribedGroup();
+        groupMetadata.setGroupId(groupId)
+            .setErrorCode(error.code())
+            .setGroupState(state)
+            .setProtocolType(protocolType)
+            .setProtocolData(protocol)
+            .setMembers(members)
+            .setAuthorizedOperations(authorizedOperations);
+        return  groupMetadata;
+    }
 
-            List members = new ArrayList<>();
-            for (Object memberObj : groupStruct.getArray(MEMBERS_KEY_NAME)) {
-                Struct memberStruct = (Struct) memberObj;
-                String memberId = memberStruct.get(MEMBER_ID);
-                String clientId = memberStruct.getString(CLIENT_ID_KEY_NAME);
-                String clientHost = memberStruct.getString(CLIENT_HOST_KEY_NAME);
-                ByteBuffer memberMetadata = memberStruct.getBytes(MEMBER_METADATA_KEY_NAME);
-                ByteBuffer memberAssignment = memberStruct.getBytes(MEMBER_ASSIGNMENT_KEY_NAME);
-                members.add(new GroupMember(memberId, clientId, clientHost,
-                        memberMetadata, memberAssignment));
-            }
-            this.groups.put(groupId, new GroupMetadata(error, state, protocolType, protocol, members));
-        }
+    @Override
+    public DescribeGroupsResponseData data() {
+        return data;
     }
 
+    @Override
     public int throttleTimeMs() {
-        return throttleTimeMs;
+        return data.throttleTimeMs();
     }
 
-    public Map groups() {
-        return groups;
-    }
+    public static final String UNKNOWN_STATE = "";
+    public static final String UNKNOWN_PROTOCOL_TYPE = "";
+    public static final String UNKNOWN_PROTOCOL = "";
 
     @Override
     public Map errorCounts() {
         Map errorCounts = new HashMap<>();
-        for (GroupMetadata response : groups.values())
-            updateErrorCounts(errorCounts, response.error);
+        data.groups().forEach(describedGroup ->
+            updateErrorCounts(errorCounts, Errors.forCode(describedGroup.errorCode())));
         return errorCounts;
     }
 
-    public static class GroupMetadata {
-        private final Errors error;
-        private final String state;
-        private final String protocolType;
-        private final String protocol;
-        private final List members;
-
-        public GroupMetadata(Errors error,
-                             String state,
-                             String protocolType,
-                             String protocol,
-                             List members) {
-            this.error = error;
-            this.state = state;
-            this.protocolType = protocolType;
-            this.protocol = protocol;
-            this.members = members;
-        }
-
-        public Errors error() {
-            return error;
-        }
-
-        public String state() {
-            return state;
-        }
-
-        public String protocolType() {
-            return protocolType;
-        }
-
-        public String protocol() {
-            return protocol;
-        }
-
-        public List members() {
-            return members;
-        }
-
-        public static GroupMetadata forError(Errors error) {
-            return new DescribeGroupsResponse.GroupMetadata(
-                    error,
-                    DescribeGroupsResponse.UNKNOWN_STATE,
-                    DescribeGroupsResponse.UNKNOWN_PROTOCOL_TYPE,
-                    DescribeGroupsResponse.UNKNOWN_PROTOCOL,
-                    Collections.emptyList());
-        }
-    }
-
-    public static class GroupMember {
-        private final String memberId;
-        private final String clientId;
-        private final String clientHost;
-        private final ByteBuffer memberMetadata;
-        private final ByteBuffer memberAssignment;
-
-        public GroupMember(String memberId,
-                           String clientId,
-                           String clientHost,
-                           ByteBuffer memberMetadata,
-                           ByteBuffer memberAssignment) {
-            this.memberId = memberId;
-            this.clientId = clientId;
-            this.clientHost = clientHost;
-            this.memberMetadata = memberMetadata;
-            this.memberAssignment = memberAssignment;
-        }
-
-        public String memberId() {
-            return memberId;
-        }
-
-        public String clientId() {
-            return clientId;
-        }
-
-        public String clientHost() {
-            return clientHost;
-        }
-
-        public ByteBuffer memberMetadata() {
-            return memberMetadata;
-        }
-
-        public ByteBuffer memberAssignment() {
-            return memberAssignment;
-        }
-    }
-
-    public static DescribeGroupsResponse fromError(Errors error, List groupIds) {
-        return fromError(DEFAULT_THROTTLE_TIME, error, groupIds);
+    public static DescribedGroup forError(String groupId, Errors error) {
+        return groupMetadata(groupId, error, DescribeGroupsResponse.UNKNOWN_STATE, DescribeGroupsResponse.UNKNOWN_PROTOCOL_TYPE,
+                DescribeGroupsResponse.UNKNOWN_PROTOCOL, Collections.emptyList(), AUTHORIZED_OPERATIONS_OMITTED);
     }
 
     public static DescribeGroupsResponse fromError(int throttleTimeMs, Errors error, List groupIds) {
-        GroupMetadata errorMetadata = GroupMetadata.forError(error);
-        Map groups = new HashMap<>();
+        DescribeGroupsResponseData describeGroupsResponseData = new DescribeGroupsResponseData();
+        describeGroupsResponseData.setThrottleTimeMs(throttleTimeMs);
         for (String groupId : groupIds)
-            groups.put(groupId, errorMetadata);
-        return new DescribeGroupsResponse(throttleTimeMs, groups);
+            describeGroupsResponseData.groups().add(DescribeGroupsResponse.forError(groupId, error));
+        return new DescribeGroupsResponse(describeGroupsResponseData);
     }
 
-    @Override
-    protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.DESCRIBE_GROUPS.responseSchema(version));
-        struct.setIfExists(THROTTLE_TIME_MS, throttleTimeMs);
-
-        List groupStructs = new ArrayList<>();
-        for (Map.Entry groupEntry : groups.entrySet()) {
-            Struct groupStruct = struct.instance(GROUPS_KEY_NAME);
-            GroupMetadata group = groupEntry.getValue();
-            groupStruct.set(GROUP_ID, groupEntry.getKey());
-            groupStruct.set(ERROR_CODE, group.error.code());
-            groupStruct.set(GROUP_STATE_KEY_NAME, group.state);
-            groupStruct.set(PROTOCOL_TYPE_KEY_NAME, group.protocolType);
-            groupStruct.set(PROTOCOL_KEY_NAME, group.protocol);
-            List membersList = new ArrayList<>();
-            for (GroupMember member : group.members) {
-                Struct memberStruct = groupStruct.instance(MEMBERS_KEY_NAME);
-                memberStruct.set(MEMBER_ID, member.memberId);
-                memberStruct.set(CLIENT_ID_KEY_NAME, member.clientId);
-                memberStruct.set(CLIENT_HOST_KEY_NAME, member.clientHost);
-                memberStruct.set(MEMBER_METADATA_KEY_NAME, member.memberMetadata);
-                memberStruct.set(MEMBER_ASSIGNMENT_KEY_NAME, member.memberAssignment);
-                membersList.add(memberStruct);
-            }
-            groupStruct.set(MEMBERS_KEY_NAME, membersList.toArray());
-            groupStructs.add(groupStruct);
-        }
-        struct.set(GROUPS_KEY_NAME, groupStructs.toArray());
-
-        return struct;
+    public static DescribeGroupsResponse parse(ByteBuffer buffer, short version) {
+        return new DescribeGroupsResponse(new DescribeGroupsResponseData(new ByteBufferAccessor(buffer), version));
     }
 
-    public static DescribeGroupsResponse parse(ByteBuffer buffer, short version) {
-        return new DescribeGroupsResponse(ApiKeys.DESCRIBE_GROUPS.parseResponse(version, buffer));
+    @Override
+    public boolean shouldClientThrottle(short version) {
+        return version >= 2;
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/DescribeLogDirsRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/DescribeLogDirsRequest.java
index 5f35c436ec7a5..05ca4f0d59391 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/DescribeLogDirsRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/DescribeLogDirsRequest.java
@@ -17,143 +17,56 @@
 
 package org.apache.kafka.common.requests;
 
-import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.message.DescribeLogDirsRequestData;
+import org.apache.kafka.common.message.DescribeLogDirsResponseData;
 import org.apache.kafka.common.protocol.ApiKeys;
-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.requests.DescribeLogDirsResponse.LogDirInfo;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
 
 import java.nio.ByteBuffer;
-import java.util.ArrayList;
-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.TOPIC_NAME;
-import static org.apache.kafka.common.protocol.types.Type.INT32;
 
 public class DescribeLogDirsRequest extends AbstractRequest {
 
-    // request level key names
-    private static final String TOPICS_KEY_NAME = "topics";
-
-    // topic level key names
-    private static final String PARTITIONS_KEY_NAME = "partitions";
-
-    private static final Schema DESCRIBE_LOG_DIRS_REQUEST_V0 = new Schema(
-            new Field(TOPICS_KEY_NAME, ArrayOf.nullable(new Schema(
-                    TOPIC_NAME,
-                    new Field(PARTITIONS_KEY_NAME, new ArrayOf(INT32), "List of partition ids of the topic.")))));
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{DESCRIBE_LOG_DIRS_REQUEST_V0};
-    }
-
-    private final Set topicPartitions;
+    private final DescribeLogDirsRequestData data;
 
     public static class Builder extends AbstractRequest.Builder {
-        private final Set topicPartitions;
+        private final DescribeLogDirsRequestData data;
 
-        // topicPartitions == null indicates requesting all partitions, and an empty list indicates requesting no partitions.
-        public Builder(Set partitions) {
+        public Builder(DescribeLogDirsRequestData data) {
             super(ApiKeys.DESCRIBE_LOG_DIRS);
-            this.topicPartitions = partitions;
+            this.data = data;
         }
 
         @Override
         public DescribeLogDirsRequest build(short version) {
-            return new DescribeLogDirsRequest(topicPartitions, version);
+            return new DescribeLogDirsRequest(data, version);
         }
 
         @Override
         public String toString() {
-            StringBuilder builder = new StringBuilder();
-            builder.append("(type=DescribeLogDirsRequest")
-                .append(", topicPartitions=")
-                .append(topicPartitions)
-                .append(")");
-            return builder.toString();
-        }
-    }
-
-    public DescribeLogDirsRequest(Struct struct, short version) {
-        super(version);
-
-        if (struct.getArray(TOPICS_KEY_NAME) == null) {
-            topicPartitions = null;
-        } else {
-            topicPartitions = new HashSet<>();
-            for (Object topicStructObj : struct.getArray(TOPICS_KEY_NAME)) {
-                Struct topicStruct = (Struct) topicStructObj;
-                String topic = topicStruct.get(TOPIC_NAME);
-                for (Object partitionObj : topicStruct.getArray(PARTITIONS_KEY_NAME)) {
-                    int partition = (Integer) partitionObj;
-                    topicPartitions.add(new TopicPartition(topic, partition));
-                }
-            }
+            return data.toString();
         }
     }
 
-    // topicPartitions == null indicates requesting all partitions, and an empty list indicates requesting no partitions.
-    public DescribeLogDirsRequest(Set topicPartitions, short version) {
-        super(version);
-        this.topicPartitions = topicPartitions;
+    public DescribeLogDirsRequest(DescribeLogDirsRequestData data, short version) {
+        super(ApiKeys.DESCRIBE_LOG_DIRS, version);
+        this.data = data;
     }
 
     @Override
-    protected Struct toStruct() {
-        Struct struct = new Struct(ApiKeys.DESCRIBE_LOG_DIRS.requestSchema(version()));
-        if (topicPartitions == null) {
-            struct.set(TOPICS_KEY_NAME, null);
-            return struct;
-        }
-
-        Map> partitionsByTopic = new HashMap<>();
-        for (TopicPartition tp : topicPartitions) {
-            if (!partitionsByTopic.containsKey(tp.topic())) {
-                partitionsByTopic.put(tp.topic(), new ArrayList());
-            }
-            partitionsByTopic.get(tp.topic()).add(tp.partition());
-        }
-
-        List topicStructArray = new ArrayList<>();
-        for (Map.Entry> partitionsByTopicEntry : partitionsByTopic.entrySet()) {
-            Struct topicStruct = struct.instance(TOPICS_KEY_NAME);
-            topicStruct.set(TOPIC_NAME, partitionsByTopicEntry.getKey());
-            topicStruct.set(PARTITIONS_KEY_NAME, partitionsByTopicEntry.getValue().toArray());
-            topicStructArray.add(topicStruct);
-        }
-        struct.set(TOPICS_KEY_NAME, topicStructArray.toArray());
-
-        return struct;
+    public DescribeLogDirsRequestData data() {
+        return data;
     }
 
     @Override
     public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
-        short versionId = version();
-        switch (versionId) {
-            case 0:
-                return new DescribeLogDirsResponse(throttleTimeMs, new HashMap());
-            default:
-                throw new IllegalArgumentException(
-                    String.format("Version %d is not valid. Valid versions for %s are 0 to %d", versionId,
-                        this.getClass().getSimpleName(), ApiKeys.DESCRIBE_LOG_DIRS.latestVersion()));
-        }
+        return new DescribeLogDirsResponse(new DescribeLogDirsResponseData().setThrottleTimeMs(throttleTimeMs));
     }
 
     public boolean isAllTopicPartitions() {
-        return topicPartitions == null;
-    }
-
-    public Set topicPartitions() {
-        return topicPartitions;
+        return data.topics() == null;
     }
 
     public static DescribeLogDirsRequest parse(ByteBuffer buffer, short version) {
-        return new DescribeLogDirsRequest(ApiKeys.DESCRIBE_LOG_DIRS.parseRequest(version, buffer), version);
+        return new DescribeLogDirsRequest(new DescribeLogDirsRequestData(new ByteBufferAccessor(buffer), version), version);
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/DescribeLogDirsResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/DescribeLogDirsResponse.java
index a242240e0fe9d..cd1326be9d1f9 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/DescribeLogDirsResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/DescribeLogDirsResponse.java
@@ -18,172 +18,63 @@
 package org.apache.kafka.common.requests;
 
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.message.DescribeLogDirsResponseData;
 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.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 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.protocol.types.Type.BOOLEAN;
-import static org.apache.kafka.common.protocol.types.Type.INT64;
-import static org.apache.kafka.common.protocol.types.Type.STRING;
-
 
 public class DescribeLogDirsResponse extends AbstractResponse {
 
     public static final long INVALID_OFFSET_LAG = -1L;
 
-    // request level key names
-    private static final String LOG_DIRS_KEY_NAME = "log_dirs";
-
-    // dir level key names
-    private static final String LOG_DIR_KEY_NAME = "log_dir";
-    private static final String TOPICS_KEY_NAME = "topics";
-
-    // topic level key names
-    private static final String PARTITIONS_KEY_NAME = "partitions";
-
-    // partition level key names
-    private static final String SIZE_KEY_NAME = "size";
-    private static final String OFFSET_LAG_KEY_NAME = "offset_lag";
-    private static final String IS_FUTURE_KEY_NAME = "is_future";
-
-    private static final Schema DESCRIBE_LOG_DIRS_RESPONSE_V0 = new Schema(
-            THROTTLE_TIME_MS,
-            new Field(LOG_DIRS_KEY_NAME, new ArrayOf(new Schema(
-                    ERROR_CODE,
-                    new Field(LOG_DIR_KEY_NAME, STRING, "The absolute log directory path."),
-                    new Field(TOPICS_KEY_NAME, new ArrayOf(new Schema(
-                            TOPIC_NAME,
-                            new Field(PARTITIONS_KEY_NAME, new ArrayOf(new Schema(
-                                    PARTITION_ID,
-                                    new Field(SIZE_KEY_NAME, INT64, "The size of the log segments of the partition in bytes."),
-                                    new Field(OFFSET_LAG_KEY_NAME, INT64, "The lag of the log's LEO w.r.t. partition's HW " +
-                                            "(if it is the current log for the partition) or current replica's LEO " +
-                                            "(if it is the future log for the partition)"),
-                                    new Field(IS_FUTURE_KEY_NAME, BOOLEAN, "True if this log is created by " +
-                                            "AlterReplicaLogDirsRequest and will replace the current log of the replica " +
-                                            "in the future.")))))))))));
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{DESCRIBE_LOG_DIRS_RESPONSE_V0};
-    }
-
-    private final int throttleTimeMs;
-    private final Map logDirInfos;
-
-    public DescribeLogDirsResponse(Struct struct) {
-        throttleTimeMs = struct.get(THROTTLE_TIME_MS);
-        logDirInfos = new HashMap<>();
-
-        for (Object logDirStructObj : struct.getArray(LOG_DIRS_KEY_NAME)) {
-            Struct logDirStruct = (Struct) logDirStructObj;
-            Errors error = Errors.forCode(logDirStruct.get(ERROR_CODE));
-            String logDir = logDirStruct.getString(LOG_DIR_KEY_NAME);
-            Map replicaInfos = new HashMap<>();
-
-            for (Object topicStructObj : logDirStruct.getArray(TOPICS_KEY_NAME)) {
-                Struct topicStruct = (Struct) topicStructObj;
-                String topic = topicStruct.get(TOPIC_NAME);
-
-                for (Object partitionStructObj : topicStruct.getArray(PARTITIONS_KEY_NAME)) {
-                    Struct partitionStruct = (Struct) partitionStructObj;
-                    int partition = partitionStruct.get(PARTITION_ID);
-                    long size = partitionStruct.getLong(SIZE_KEY_NAME);
-                    long offsetLag = partitionStruct.getLong(OFFSET_LAG_KEY_NAME);
-                    boolean isFuture = partitionStruct.getBoolean(IS_FUTURE_KEY_NAME);
-                    ReplicaInfo replicaInfo = new ReplicaInfo(size, offsetLag, isFuture);
-                    replicaInfos.put(new TopicPartition(topic, partition), replicaInfo);
-                }
-            }
-
-            logDirInfos.put(logDir, new LogDirInfo(error, replicaInfos));
-        }
-    }
+    private final DescribeLogDirsResponseData data;
 
-    /**
-     * Constructor for version 0.
-     */
-    public DescribeLogDirsResponse(int throttleTimeMs, Map logDirInfos) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.logDirInfos = logDirInfos;
+    public DescribeLogDirsResponse(DescribeLogDirsResponseData data) {
+        super(ApiKeys.DESCRIBE_LOG_DIRS);
+        this.data = data;
     }
 
     @Override
-    protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.DESCRIBE_LOG_DIRS.responseSchema(version));
-        struct.set(THROTTLE_TIME_MS, throttleTimeMs);
-        List logDirStructArray = new ArrayList<>();
-        for (Map.Entry logDirInfosEntry : logDirInfos.entrySet()) {
-            LogDirInfo logDirInfo = logDirInfosEntry.getValue();
-            Struct logDirStruct = struct.instance(LOG_DIRS_KEY_NAME);
-            logDirStruct.set(ERROR_CODE, logDirInfo.error.code());
-            logDirStruct.set(LOG_DIR_KEY_NAME, logDirInfosEntry.getKey());
-
-            Map> replicaInfosByTopic = CollectionUtils.groupDataByTopic(logDirInfo.replicaInfos);
-            List topicStructArray = new ArrayList<>();
-            for (Map.Entry> replicaInfosByTopicEntry : replicaInfosByTopic.entrySet()) {
-                Struct topicStruct = logDirStruct.instance(TOPICS_KEY_NAME);
-                topicStruct.set(TOPIC_NAME, replicaInfosByTopicEntry.getKey());
-                List partitionStructArray = new ArrayList<>();
-
-                for (Map.Entry replicaInfosByPartitionEntry : replicaInfosByTopicEntry.getValue().entrySet()) {
-                    Struct partitionStruct = topicStruct.instance(PARTITIONS_KEY_NAME);
-                    ReplicaInfo replicaInfo = replicaInfosByPartitionEntry.getValue();
-                    partitionStruct.set(PARTITION_ID, replicaInfosByPartitionEntry.getKey());
-                    partitionStruct.set(SIZE_KEY_NAME, replicaInfo.size);
-                    partitionStruct.set(OFFSET_LAG_KEY_NAME, replicaInfo.offsetLag);
-                    partitionStruct.set(IS_FUTURE_KEY_NAME, replicaInfo.isFuture);
-                    partitionStructArray.add(partitionStruct);
-                }
-                topicStruct.set(PARTITIONS_KEY_NAME, partitionStructArray.toArray());
-                topicStructArray.add(topicStruct);
-            }
-            logDirStruct.set(TOPICS_KEY_NAME, topicStructArray.toArray());
-            logDirStructArray.add(logDirStruct);
-        }
-        struct.set(LOG_DIRS_KEY_NAME, logDirStructArray.toArray());
-        return struct;
+    public DescribeLogDirsResponseData data() {
+        return data;
     }
 
+    @Override
     public int throttleTimeMs() {
-        return throttleTimeMs;
+        return data.throttleTimeMs();
     }
 
     @Override
     public Map errorCounts() {
         Map errorCounts = new HashMap<>();
-        for (LogDirInfo logDirInfo : logDirInfos.values())
-            updateErrorCounts(errorCounts, logDirInfo.error);
+        data.results().forEach(result -> {
+            updateErrorCounts(errorCounts, Errors.forCode(result.errorCode()));
+        });
         return errorCounts;
     }
 
-    public Map logDirInfos() {
-        return logDirInfos;
-    }
-
     public static DescribeLogDirsResponse parse(ByteBuffer buffer, short version) {
-        return new DescribeLogDirsResponse(ApiKeys.DESCRIBE_LOG_DIRS.responseSchema(version).read(buffer));
+        return new DescribeLogDirsResponse(new DescribeLogDirsResponseData(new ByteBufferAccessor(buffer), version));
     }
 
+    // Note this class is part of the public API, reachable from Admin.describeLogDirs()
     /**
      * Possible error code:
      *
      * KAFKA_STORAGE_ERROR (56)
      * UNKNOWN (-1)
+     *
+     * @deprecated Deprecated Since Kafka 2.7.
+     * Use {@link org.apache.kafka.clients.admin.DescribeLogDirsResult#descriptions()}
+     * and {@link org.apache.kafka.clients.admin.DescribeLogDirsResult#allDescriptions()} to access the replacement
+     * class {@link org.apache.kafka.clients.admin.LogDirDescription}.
      */
+    @Deprecated
     static public class LogDirInfo {
         public final Errors error;
         public final Map replicaInfos;
@@ -192,8 +83,28 @@ public LogDirInfo(Errors error, Map replicaInfos) {
             this.error = error;
             this.replicaInfos = replicaInfos;
         }
+
+        @Override
+        public String toString() {
+            StringBuilder builder = new StringBuilder();
+            builder.append("(error=")
+                    .append(error)
+                    .append(", replicas=")
+                    .append(replicaInfos)
+                    .append(")");
+            return builder.toString();
+        }
     }
 
+    // Note this class is part of the public API, reachable from Admin.describeLogDirs()
+
+    /**
+     * @deprecated Deprecated Since Kafka 2.7.
+     * Use {@link org.apache.kafka.clients.admin.DescribeLogDirsResult#descriptions()}
+     * and {@link org.apache.kafka.clients.admin.DescribeLogDirsResult#allDescriptions()} to access the replacement
+     * class {@link org.apache.kafka.clients.admin.ReplicaInfo}.
+     */
+    @Deprecated
     static public class ReplicaInfo {
 
         public final long size;
@@ -219,4 +130,9 @@ public String toString() {
             return builder.toString();
         }
     }
+
+    @Override
+    public boolean shouldClientThrottle(short version) {
+        return version >= 1;
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/DescribeQuorumRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/DescribeQuorumRequest.java
new file mode 100644
index 0000000000000..acdb11c664420
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/DescribeQuorumRequest.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.requests;
+
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.message.DescribeQuorumRequestData;
+import org.apache.kafka.common.message.DescribeQuorumResponseData;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.Errors;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.stream.Collectors;
+
+public class DescribeQuorumRequest extends AbstractRequest {
+    public static class Builder extends AbstractRequest.Builder {
+        private final DescribeQuorumRequestData data;
+
+        public Builder(DescribeQuorumRequestData data) {
+            super(ApiKeys.DESCRIBE_QUORUM);
+            this.data = data;
+        }
+
+        @Override
+        public DescribeQuorumRequest build(short version) {
+            return new DescribeQuorumRequest(data, version);
+        }
+
+        @Override
+        public String toString() {
+            return data.toString();
+        }
+    }
+
+    private final DescribeQuorumRequestData data;
+
+    private DescribeQuorumRequest(DescribeQuorumRequestData data, short version) {
+        super(ApiKeys.DESCRIBE_QUORUM, version);
+        this.data = data;
+    }
+
+    public static DescribeQuorumRequest parse(ByteBuffer buffer, short version) {
+        return new DescribeQuorumRequest(new DescribeQuorumRequestData(new ByteBufferAccessor(buffer), version), version);
+    }
+
+    public static DescribeQuorumRequestData singletonRequest(TopicPartition topicPartition) {
+        return new DescribeQuorumRequestData()
+            .setTopics(Collections.singletonList(
+                new DescribeQuorumRequestData.TopicData()
+                    .setTopicName(topicPartition.topic())
+                    .setPartitions(Collections.singletonList(
+                        new DescribeQuorumRequestData.PartitionData()
+                            .setPartitionIndex(topicPartition.partition()))
+            )));
+    }
+
+    @Override
+    public DescribeQuorumRequestData data() {
+        return data;
+    }
+
+    @Override
+    public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
+        return new DescribeQuorumResponse(getTopLevelErrorResponse(Errors.forException(e)));
+    }
+
+    public static DescribeQuorumResponseData getPartitionLevelErrorResponse(DescribeQuorumRequestData data, Errors error) {
+        short errorCode = error.code();
+
+        List topicResponses = new ArrayList<>();
+        for (DescribeQuorumRequestData.TopicData topic : data.topics()) {
+            topicResponses.add(
+                new DescribeQuorumResponseData.TopicData()
+                    .setTopicName(topic.topicName())
+                    .setPartitions(topic.partitions().stream().map(
+                        requestPartition -> new DescribeQuorumResponseData.PartitionData()
+                                                .setPartitionIndex(requestPartition.partitionIndex())
+                                                .setErrorCode(errorCode)
+                    ).collect(Collectors.toList())));
+        }
+
+        return new DescribeQuorumResponseData().setTopics(topicResponses);
+    }
+
+    public static DescribeQuorumResponseData getTopLevelErrorResponse(Errors topLevelError) {
+        return new DescribeQuorumResponseData().setErrorCode(topLevelError.code());
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/DescribeQuorumResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/DescribeQuorumResponse.java
new file mode 100644
index 0000000000000..cbf945b70409a
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/DescribeQuorumResponse.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.common.requests;
+
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.message.DescribeQuorumResponseData;
+import org.apache.kafka.common.message.DescribeQuorumResponseData.ReplicaState;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.Errors;
+
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Possible error codes.
+ *
+ * Top level errors:
+ * - {@link Errors#CLUSTER_AUTHORIZATION_FAILED}
+ * - {@link Errors#BROKER_NOT_AVAILABLE}
+ *
+ * Partition level errors:
+ * - {@link Errors#INVALID_REQUEST}
+ * - {@link Errors#UNKNOWN_TOPIC_OR_PARTITION}
+ */
+public class DescribeQuorumResponse extends AbstractResponse {
+    private final DescribeQuorumResponseData data;
+
+    public DescribeQuorumResponse(DescribeQuorumResponseData data) {
+        super(ApiKeys.DESCRIBE_QUORUM);
+        this.data = data;
+    }
+
+    @Override
+    public Map errorCounts() {
+        Map errors = new HashMap<>();
+
+        errors.put(Errors.forCode(data.errorCode()), 1);
+
+        for (DescribeQuorumResponseData.TopicData topicResponse : data.topics()) {
+            for (DescribeQuorumResponseData.PartitionData partitionResponse : topicResponse.partitions()) {
+                updateErrorCounts(errors, Errors.forCode(partitionResponse.errorCode()));
+            }
+        }
+        return errors;
+    }
+
+    @Override
+    public DescribeQuorumResponseData data() {
+        return data;
+    }
+
+    @Override
+    public int throttleTimeMs() {
+        return DEFAULT_THROTTLE_TIME;
+    }
+
+    public static DescribeQuorumResponseData singletonResponse(TopicPartition topicPartition,
+                                                               int leaderId,
+                                                               int leaderEpoch,
+                                                               long highWatermark,
+                                                               List voterStates,
+                                                               List observerStates) {
+        return new DescribeQuorumResponseData()
+            .setTopics(Collections.singletonList(new DescribeQuorumResponseData.TopicData()
+                .setTopicName(topicPartition.topic())
+                .setPartitions(Collections.singletonList(new DescribeQuorumResponseData.PartitionData()
+                    .setErrorCode(Errors.NONE.code())
+                    .setLeaderId(leaderId)
+                    .setLeaderEpoch(leaderEpoch)
+                    .setHighWatermark(highWatermark)
+                    .setCurrentVoters(voterStates)
+                    .setObservers(observerStates)))));
+    }
+
+    public static DescribeQuorumResponse parse(ByteBuffer buffer, short version) {
+        return new DescribeQuorumResponse(new DescribeQuorumResponseData(new ByteBufferAccessor(buffer), version));
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/DescribeUserScramCredentialsRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/DescribeUserScramCredentialsRequest.java
new file mode 100644
index 0000000000000..0142e5ab6f592
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/DescribeUserScramCredentialsRequest.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.common.requests;
+
+import org.apache.kafka.common.message.DescribeUserScramCredentialsRequestData;
+import org.apache.kafka.common.message.DescribeUserScramCredentialsResponseData;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+
+import java.nio.ByteBuffer;
+
+public class DescribeUserScramCredentialsRequest extends AbstractRequest {
+
+    public static class Builder extends AbstractRequest.Builder {
+        private final DescribeUserScramCredentialsRequestData data;
+
+        public Builder(DescribeUserScramCredentialsRequestData data) {
+            super(ApiKeys.DESCRIBE_USER_SCRAM_CREDENTIALS);
+            this.data = data;
+        }
+
+        @Override
+        public DescribeUserScramCredentialsRequest build(short version) {
+            return new DescribeUserScramCredentialsRequest(data, version);
+        }
+
+        @Override
+        public String toString() {
+            return data.toString();
+        }
+    }
+
+    private final DescribeUserScramCredentialsRequestData data;
+    private final short version;
+
+    private DescribeUserScramCredentialsRequest(DescribeUserScramCredentialsRequestData data, short version) {
+        super(ApiKeys.DESCRIBE_USER_SCRAM_CREDENTIALS, version);
+        this.data = data;
+        this.version = version;
+    }
+
+    public static DescribeUserScramCredentialsRequest parse(ByteBuffer buffer, short version) {
+        return new DescribeUserScramCredentialsRequest(new DescribeUserScramCredentialsRequestData(
+                new ByteBufferAccessor(buffer), version), version);
+    }
+
+    @Override
+    public DescribeUserScramCredentialsRequestData data() {
+        return data;
+    }
+
+    @Override
+    public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
+        ApiError apiError = ApiError.fromThrowable(e);
+        return new DescribeUserScramCredentialsResponse(new DescribeUserScramCredentialsResponseData()
+                .setErrorCode(apiError.error().code())
+                .setErrorMessage(apiError.message()));
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/DescribeUserScramCredentialsResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/DescribeUserScramCredentialsResponse.java
new file mode 100644
index 0000000000000..001cefae41a6b
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/DescribeUserScramCredentialsResponse.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.common.requests;
+
+import org.apache.kafka.common.message.DescribeUserScramCredentialsResponseData;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.Errors;
+
+import java.nio.ByteBuffer;
+import java.util.Map;
+
+public class DescribeUserScramCredentialsResponse extends AbstractResponse {
+
+    private final DescribeUserScramCredentialsResponseData data;
+
+    public DescribeUserScramCredentialsResponse(DescribeUserScramCredentialsResponseData responseData) {
+        super(ApiKeys.DESCRIBE_USER_SCRAM_CREDENTIALS);
+        this.data = responseData;
+    }
+
+    @Override
+    public DescribeUserScramCredentialsResponseData data() {
+        return data;
+    }
+
+    @Override
+    public boolean shouldClientThrottle(short version) {
+        return true;
+    }
+
+    @Override
+    public int throttleTimeMs() {
+        return data.throttleTimeMs();
+    }
+
+    @Override
+    public Map errorCounts() {
+        return errorCounts(data.results().stream().map(r -> Errors.forCode(r.errorCode())));
+    }
+
+    public static DescribeUserScramCredentialsResponse parse(ByteBuffer buffer, short version) {
+        return new DescribeUserScramCredentialsResponse(new DescribeUserScramCredentialsResponseData(new ByteBufferAccessor(buffer), version));
+    }
+}
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..92f6b45eed59d
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/ElectLeadersRequest.java
@@ -0,0 +1,125 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.MessageUtil;
+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;
+    }
+
+    @Override
+    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(new ElectLeadersRequestData(new ByteBufferAccessor(buffer), version), 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..88d4d19fc021e
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/ElectLeadersResponse.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.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;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.message.ElectLeadersResponseData.ReplicaElectionResult;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.Errors;
+
+public class ElectLeadersResponse extends AbstractResponse {
+
+    private final ElectLeadersResponseData data;
+
+    public ElectLeadersResponse(ElectLeadersResponseData data) {
+        super(ApiKeys.ELECT_LEADERS);
+        this.data = data;
+    }
+
+    public ElectLeadersResponse(
+            int throttleTimeMs,
+            short errorCode,
+            List electionResults,
+            short version) {
+        super(ApiKeys.ELECT_LEADERS);
+        this.data = new ElectLeadersResponseData();
+        data.setThrottleTimeMs(throttleTimeMs);
+        if (version >= 1)
+            data.setErrorCode(errorCode);
+        data.setReplicaElectionResults(electionResults);
+    }
+
+    @Override
+    public ElectLeadersResponseData data() {
+        return data;
+    }
+
+    @Override
+    public int throttleTimeMs() {
+        return data.throttleTimeMs();
+    }
+
+    @Override
+    public Map errorCounts() {
+        HashMap counts = new HashMap<>();
+        updateErrorCounts(counts, Errors.forCode(data.errorCode()));
+        data.replicaElectionResults().forEach(result ->
+            result.partitionResult().forEach(partitionResult ->
+                updateErrorCounts(counts, Errors.forCode(partitionResult.errorCode()))
+            )
+        );
+        return counts;
+    }
+
+    public static ElectLeadersResponse parse(ByteBuffer buffer, short version) {
+        return new ElectLeadersResponse(new ElectLeadersResponseData(new ByteBufferAccessor(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/EndQuorumEpochRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/EndQuorumEpochRequest.java
new file mode 100644
index 0000000000000..136bc54d50b2e
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/EndQuorumEpochRequest.java
@@ -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 org.apache.kafka.common.requests;
+
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.message.EndQuorumEpochRequestData;
+import org.apache.kafka.common.message.EndQuorumEpochResponseData;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.Errors;
+
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.List;
+
+public class EndQuorumEpochRequest extends AbstractRequest {
+    public static class Builder extends AbstractRequest.Builder {
+        private final EndQuorumEpochRequestData data;
+
+        public Builder(EndQuorumEpochRequestData data) {
+            super(ApiKeys.END_QUORUM_EPOCH);
+            this.data = data;
+        }
+
+        @Override
+        public EndQuorumEpochRequest build(short version) {
+            return new EndQuorumEpochRequest(data, version);
+        }
+
+        @Override
+        public String toString() {
+            return data.toString();
+        }
+    }
+
+    private final EndQuorumEpochRequestData data;
+
+    private EndQuorumEpochRequest(EndQuorumEpochRequestData data, short version) {
+        super(ApiKeys.END_QUORUM_EPOCH, version);
+        this.data = data;
+    }
+
+    @Override
+    public EndQuorumEpochRequestData data() {
+        return data;
+    }
+
+    @Override
+    public EndQuorumEpochResponse getErrorResponse(int throttleTimeMs, Throwable e) {
+        return new EndQuorumEpochResponse(new EndQuorumEpochResponseData()
+            .setErrorCode(Errors.forException(e).code()));
+    }
+
+    public static EndQuorumEpochRequest parse(ByteBuffer buffer, short version) {
+        return new EndQuorumEpochRequest(new EndQuorumEpochRequestData(new ByteBufferAccessor(buffer), version), version);
+    }
+
+    public static EndQuorumEpochRequestData singletonRequest(TopicPartition topicPartition,
+                                                             int leaderEpoch,
+                                                             int leaderId,
+                                                             List preferredSuccessors) {
+        return singletonRequest(topicPartition, null, leaderEpoch, leaderId, preferredSuccessors);
+    }
+
+    public static EndQuorumEpochRequestData singletonRequest(TopicPartition topicPartition,
+                                                             String clusterId,
+                                                             int leaderEpoch,
+                                                             int leaderId,
+                                                             List preferredSuccessors) {
+        return new EndQuorumEpochRequestData()
+                   .setClusterId(clusterId)
+                   .setTopics(Collections.singletonList(
+                       new EndQuorumEpochRequestData.TopicData()
+                           .setTopicName(topicPartition.topic())
+                           .setPartitions(Collections.singletonList(
+                               new EndQuorumEpochRequestData.PartitionData()
+                                   .setPartitionIndex(topicPartition.partition())
+                                   .setLeaderEpoch(leaderEpoch)
+                                   .setLeaderId(leaderId)
+                                   .setPreferredSuccessors(preferredSuccessors))))
+                   );
+    }
+
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/EndQuorumEpochResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/EndQuorumEpochResponse.java
new file mode 100644
index 0000000000000..ac2c0c5c9d5c3
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/EndQuorumEpochResponse.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.common.requests;
+
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.message.EndQuorumEpochResponseData;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.Errors;
+
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * Possible error codes.
+ *
+ * Top level errors:
+ * - {@link Errors#CLUSTER_AUTHORIZATION_FAILED}
+ * - {@link Errors#BROKER_NOT_AVAILABLE}
+ *
+ * Partition level errors:
+ * - {@link Errors#FENCED_LEADER_EPOCH}
+ * - {@link Errors#INVALID_REQUEST}
+ * - {@link Errors#INCONSISTENT_VOTER_SET}
+ * - {@link Errors#UNKNOWN_TOPIC_OR_PARTITION}
+ */
+public class EndQuorumEpochResponse extends AbstractResponse {
+    private final EndQuorumEpochResponseData data;
+
+    public EndQuorumEpochResponse(EndQuorumEpochResponseData data) {
+        super(ApiKeys.END_QUORUM_EPOCH);
+        this.data = data;
+    }
+
+    @Override
+    public Map errorCounts() {
+        Map errors = new HashMap<>();
+
+        errors.put(Errors.forCode(data.errorCode()), 1);
+
+        for (EndQuorumEpochResponseData.TopicData topicResponse : data.topics()) {
+            for (EndQuorumEpochResponseData.PartitionData partitionResponse : topicResponse.partitions()) {
+                updateErrorCounts(errors, Errors.forCode(partitionResponse.errorCode()));
+            }
+        }
+        return errors;
+    }
+
+    @Override
+    public EndQuorumEpochResponseData data() {
+        return data;
+    }
+
+    @Override
+    public int throttleTimeMs() {
+        return DEFAULT_THROTTLE_TIME;
+    }
+
+    public static EndQuorumEpochResponseData singletonResponse(
+        Errors topLevelError,
+        TopicPartition topicPartition,
+        Errors partitionLevelError,
+        int leaderEpoch,
+        int leaderId
+    ) {
+        return new EndQuorumEpochResponseData()
+                   .setErrorCode(topLevelError.code())
+                   .setTopics(Collections.singletonList(
+                       new EndQuorumEpochResponseData.TopicData()
+                           .setTopicName(topicPartition.topic())
+                           .setPartitions(Collections.singletonList(
+                               new EndQuorumEpochResponseData.PartitionData()
+                                   .setErrorCode(partitionLevelError.code())
+                                   .setLeaderId(leaderId)
+                                   .setLeaderEpoch(leaderEpoch)
+                           )))
+                   );
+    }
+
+    public static EndQuorumEpochResponse parse(ByteBuffer buffer, short version) {
+        return new EndQuorumEpochResponse(new EndQuorumEpochResponseData(new ByteBufferAccessor(buffer), version));
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/EndTxnRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/EndTxnRequest.java
index 9118d6a5a086d..c9ea98005fd09 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/EndTxnRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/EndTxnRequest.java
@@ -16,122 +16,63 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.message.EndTxnRequestData;
+import org.apache.kafka.common.message.EndTxnResponseData;
 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.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.PRODUCER_EPOCH;
-import static org.apache.kafka.common.protocol.CommonFields.PRODUCER_ID;
-import static org.apache.kafka.common.protocol.CommonFields.TRANSACTIONAL_ID;
-import static org.apache.kafka.common.protocol.types.Type.BOOLEAN;
-
 public class EndTxnRequest extends AbstractRequest {
-    private static final String TRANSACTION_RESULT_KEY_NAME = "transaction_result";
-
-    private static final Schema END_TXN_REQUEST_V0 = new Schema(
-            TRANSACTIONAL_ID,
-            PRODUCER_ID,
-            PRODUCER_EPOCH,
-            new Field(TRANSACTION_RESULT_KEY_NAME, BOOLEAN, "The result of the transaction (0 = ABORT, 1 = COMMIT)"));
 
-    public static Schema[] schemaVersions() {
-        return new Schema[]{END_TXN_REQUEST_V0};
-    }
+    private final EndTxnRequestData data;
 
     public static class Builder extends AbstractRequest.Builder {
-        private final String transactionalId;
-        private final long producerId;
-        private final short producerEpoch;
-        private final TransactionResult result;
+        public final EndTxnRequestData data;
 
-        public Builder(String transactionalId, long producerId, short producerEpoch, TransactionResult result) {
+        public Builder(EndTxnRequestData data) {
             super(ApiKeys.END_TXN);
-            this.transactionalId = transactionalId;
-            this.producerId = producerId;
-            this.producerEpoch = producerEpoch;
-            this.result = result;
-        }
-
-        public TransactionResult result() {
-            return result;
+            this.data = data;
         }
 
         @Override
         public EndTxnRequest build(short version) {
-            return new EndTxnRequest(version, transactionalId, producerId, producerEpoch, result);
+            return new EndTxnRequest(data, version);
         }
 
         @Override
         public String toString() {
-            StringBuilder bld = new StringBuilder();
-            bld.append("(type=EndTxnRequest").
-                    append(", transactionalId=").append(transactionalId).
-                    append(", producerId=").append(producerId).
-                    append(", producerEpoch=").append(producerEpoch).
-                    append(", result=").append(result).
-                    append(")");
-            return bld.toString();
+            return data.toString();
         }
     }
 
-    private final String transactionalId;
-    private final long producerId;
-    private final short producerEpoch;
-    private final TransactionResult result;
-
-    private EndTxnRequest(short version, String transactionalId, long producerId, short producerEpoch, TransactionResult result) {
-        super(version);
-        this.transactionalId = transactionalId;
-        this.producerId = producerId;
-        this.producerEpoch = producerEpoch;
-        this.result = result;
-    }
-
-    public EndTxnRequest(Struct struct, short version) {
-        super(version);
-        this.transactionalId = struct.get(TRANSACTIONAL_ID);
-        this.producerId = struct.get(PRODUCER_ID);
-        this.producerEpoch = struct.get(PRODUCER_EPOCH);
-        this.result = TransactionResult.forId(struct.getBoolean(TRANSACTION_RESULT_KEY_NAME));
+    private EndTxnRequest(EndTxnRequestData data, short version) {
+        super(ApiKeys.END_TXN, version);
+        this.data = data;
     }
 
-    public String transactionalId() {
-        return transactionalId;
-    }
-
-    public long producerId() {
-        return producerId;
-    }
-
-    public short producerEpoch() {
-        return producerEpoch;
-    }
-
-    public TransactionResult command() {
-        return result;
+    public TransactionResult result() {
+        if (data.committed())
+            return TransactionResult.COMMIT;
+        else
+            return TransactionResult.ABORT;
     }
 
     @Override
-    protected Struct toStruct() {
-        Struct struct = new Struct(ApiKeys.END_TXN.requestSchema(version()));
-        struct.set(TRANSACTIONAL_ID, transactionalId);
-        struct.set(PRODUCER_ID, producerId);
-        struct.set(PRODUCER_EPOCH, producerEpoch);
-        struct.set(TRANSACTION_RESULT_KEY_NAME, result.id);
-        return struct;
+    public EndTxnRequestData data() {
+        return data;
     }
 
     @Override
     public EndTxnResponse getErrorResponse(int throttleTimeMs, Throwable e) {
-        return new EndTxnResponse(throttleTimeMs, Errors.forException(e));
+        return new EndTxnResponse(new EndTxnResponseData()
+                                      .setErrorCode(Errors.forException(e).code())
+                                      .setThrottleTimeMs(throttleTimeMs)
+        );
     }
 
     public static EndTxnRequest parse(ByteBuffer buffer, short version) {
-        return new EndTxnRequest(ApiKeys.END_TXN.parseRequest(version, buffer), version);
+        return new EndTxnRequest(new EndTxnRequestData(new ByteBufferAccessor(buffer), version), version);
     }
-
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/EndTxnResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/EndTxnResponse.java
index a0a453d2f582a..029e7d0ce5909 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/EndTxnResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/EndTxnResponse.java
@@ -16,78 +16,66 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.message.EndTxnResponseData;
 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.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;
-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#INVALID_TXN_STATE}
+ *   - {@link Errors#INVALID_PRODUCER_ID_MAPPING}
+ *   - {@link Errors#INVALID_PRODUCER_EPOCH} // for version <=1
+ *   - {@link Errors#PRODUCER_FENCED}
+ *   - {@link Errors#TRANSACTIONAL_ID_AUTHORIZATION_FAILED}
+ */
 public class EndTxnResponse extends AbstractResponse {
-    private static final Schema END_TXN_RESPONSE_V0 = new Schema(
-            THROTTLE_TIME_MS,
-            ERROR_CODE);
 
-    public static Schema[] schemaVersions() {
-        return new Schema[]{END_TXN_RESPONSE_V0};
-    }
+    private final EndTxnResponseData data;
 
-    // Possible error codes:
-    //   NotCoordinator
-    //   CoordinatorNotAvailable
-    //   CoordinatorLoadInProgress
-    //   InvalidTxnState
-    //   InvalidProducerIdMapping
-    //   InvalidProducerEpoch
-    //   TransactionalIdAuthorizationFailed
-
-    private final Errors error;
-    private final int throttleTimeMs;
-
-    public EndTxnResponse(int throttleTimeMs, Errors error) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.error = error;
-    }
-
-    public EndTxnResponse(Struct struct) {
-        this.throttleTimeMs = struct.get(THROTTLE_TIME_MS);
-        this.error = Errors.forCode(struct.get(ERROR_CODE));
+    public EndTxnResponse(EndTxnResponseData data) {
+        super(ApiKeys.END_TXN);
+        this.data = data;
     }
 
+    @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 errorCounts(error());
     }
 
     @Override
-    protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.END_TXN.responseSchema(version));
-        struct.set(THROTTLE_TIME_MS, throttleTimeMs);
-        struct.set(ERROR_CODE, error.code());
-        return struct;
+    public EndTxnResponseData data() {
+        return data;
     }
 
     public static EndTxnResponse parse(ByteBuffer buffer, short version) {
-        return new EndTxnResponse(ApiKeys.END_TXN.parseResponse(version, buffer));
+        return new EndTxnResponse(new EndTxnResponseData(new ByteBufferAccessor(buffer), version));
     }
 
     @Override
     public String toString() {
-        return "EndTxnResponse(" +
-                "error=" + error +
-                ", throttleTimeMs=" + throttleTimeMs +
-                ')';
+        return data.toString();
+    }
+
+    @Override
+    public boolean shouldClientThrottle(short version) {
+        return version >= 1;
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/EnvelopeRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/EnvelopeRequest.java
new file mode 100644
index 0000000000000..5e8d3faea40ee
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/EnvelopeRequest.java
@@ -0,0 +1,87 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.EnvelopeRequestData;
+import org.apache.kafka.common.message.EnvelopeResponseData;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.Errors;
+
+import java.nio.ByteBuffer;
+
+public class EnvelopeRequest extends AbstractRequest {
+
+    public static class Builder extends AbstractRequest.Builder {
+
+        private final EnvelopeRequestData data;
+
+        public Builder(ByteBuffer requestData,
+                       byte[] serializedPrincipal,
+                       byte[] clientAddress) {
+            super(ApiKeys.ENVELOPE);
+            this.data = new EnvelopeRequestData()
+                            .setRequestData(requestData)
+                            .setRequestPrincipal(serializedPrincipal)
+                            .setClientHostAddress(clientAddress);
+        }
+
+        @Override
+        public EnvelopeRequest build(short version) {
+            return new EnvelopeRequest(data, version);
+        }
+
+        @Override
+        public String toString() {
+            return data.toString();
+        }
+    }
+
+    private final EnvelopeRequestData data;
+
+    public EnvelopeRequest(EnvelopeRequestData data, short version) {
+        super(ApiKeys.ENVELOPE, version);
+        this.data = data;
+    }
+
+    public ByteBuffer requestData() {
+        return data.requestData();
+    }
+
+    public byte[] clientAddress() {
+        return data.clientHostAddress();
+    }
+
+    public byte[] requestPrincipal() {
+        return data.requestPrincipal();
+    }
+
+    @Override
+    public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
+        return new EnvelopeResponse(new EnvelopeResponseData()
+                                        .setErrorCode(Errors.forException(e).code()));
+    }
+
+    public static EnvelopeRequest parse(ByteBuffer buffer, short version) {
+        return new EnvelopeRequest(new EnvelopeRequestData(new ByteBufferAccessor(buffer), version), version);
+    }
+
+    @Override
+    public EnvelopeRequestData data() {
+        return data;
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/EnvelopeResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/EnvelopeResponse.java
new file mode 100644
index 0000000000000..529f616bb26fc
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/EnvelopeResponse.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.common.requests;
+
+import org.apache.kafka.common.message.EnvelopeResponseData;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.Errors;
+
+import java.nio.ByteBuffer;
+import java.util.Map;
+
+public class EnvelopeResponse extends AbstractResponse {
+
+    private final EnvelopeResponseData data;
+
+    public EnvelopeResponse(ByteBuffer responseData, Errors error) {
+        super(ApiKeys.ENVELOPE);
+        this.data = new EnvelopeResponseData()
+                        .setResponseData(responseData)
+                        .setErrorCode(error.code());
+    }
+
+    public EnvelopeResponse(Errors error) {
+        this(null, error);
+    }
+
+    public EnvelopeResponse(EnvelopeResponseData data) {
+        super(ApiKeys.ENVELOPE);
+        this.data = data;
+    }
+
+    public ByteBuffer responseData() {
+        return data.responseData();
+    }
+
+    @Override
+    public Map errorCounts() {
+        return errorCounts(error());
+    }
+
+    public Errors error() {
+        return Errors.forCode(data.errorCode());
+    }
+
+    @Override
+    public EnvelopeResponseData data() {
+        return data;
+    }
+
+    @Override
+    public int throttleTimeMs() {
+        return DEFAULT_THROTTLE_TIME;
+    }
+
+    public static EnvelopeResponse parse(ByteBuffer buffer, short version) {
+        return new EnvelopeResponse(new EnvelopeResponseData(new ByteBufferAccessor(buffer), version));
+    }
+
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/EpochEndOffset.java b/clients/src/main/java/org/apache/kafka/common/requests/EpochEndOffset.java
deleted file mode 100644
index 0965e3612d8aa..0000000000000
--- a/clients/src/main/java/org/apache/kafka/common/requests/EpochEndOffset.java
+++ /dev/null
@@ -1,81 +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.protocol.Errors;
-
-import static org.apache.kafka.common.record.RecordBatch.NO_PARTITION_LEADER_EPOCH;
-
-/**
- * The offset, fetched from a leader, for a particular partition.
- */
-
-public class EpochEndOffset {
-    public static final long UNDEFINED_EPOCH_OFFSET = NO_PARTITION_LEADER_EPOCH;
-    public static final int UNDEFINED_EPOCH = -1;
-
-    private Errors error;
-    private long endOffset;
-
-    public EpochEndOffset(Errors error, long endOffset) {
-        this.error = error;
-        this.endOffset = endOffset;
-    }
-
-    public EpochEndOffset(long endOffset) {
-        this.error = Errors.NONE;
-        this.endOffset = endOffset;
-    }
-
-    public Errors error() {
-        return error;
-    }
-
-    public boolean hasError() {
-        return error != Errors.NONE;
-    }
-
-    public long endOffset() {
-        return endOffset;
-    }
-
-    @Override
-    public String toString() {
-        return "EpochEndOffset{" +
-                "error=" + error +
-                ", endOffset=" + endOffset +
-                '}';
-    }
-
-    @Override
-    public boolean equals(Object o) {
-        if (this == o) return true;
-        if (o == null || getClass() != o.getClass()) return false;
-
-        EpochEndOffset that = (EpochEndOffset) o;
-
-        if (error != that.error) return false;
-        return endOffset == that.endOffset;
-    }
-
-    @Override
-    public int hashCode() {
-        int result = (int) error.code();
-        result = 31 * result + (int) (endOffset ^ (endOffset >>> 32));
-        return result;
-    }
-}
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
new file mode 100644
index 0000000000000..85b02382baf89
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/ExpireDelegationTokenRequest.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.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.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.Errors;
+
+public class ExpireDelegationTokenRequest extends AbstractRequest {
+
+    private final ExpireDelegationTokenRequestData data;
+
+    private ExpireDelegationTokenRequest(ExpireDelegationTokenRequestData data, short version) {
+        super(ApiKeys.EXPIRE_DELEGATION_TOKEN, version);
+        this.data = data;
+    }
+
+    public static ExpireDelegationTokenRequest parse(ByteBuffer buffer, short version) {
+        return new ExpireDelegationTokenRequest(
+            new ExpireDelegationTokenRequestData(new ByteBufferAccessor(buffer), version), version);
+    }
+
+    @Override
+    public ExpireDelegationTokenRequestData data() {
+        return data;
+    }
+
+    @Override
+    public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
+        return new ExpireDelegationTokenResponse(
+                new ExpireDelegationTokenResponseData()
+                    .setErrorCode(Errors.forException(e).code())
+                    .setThrottleTimeMs(throttleTimeMs));
+    }
+
+    public ByteBuffer hmac() {
+        return ByteBuffer.wrap(data.hmac());
+    }
+
+    public long expiryTimePeriod() {
+        return data.expiryTimePeriodMs();
+    }
+
+    public static class Builder extends AbstractRequest.Builder {
+        private final ExpireDelegationTokenRequestData data;
+
+        public Builder(ExpireDelegationTokenRequestData data) {
+            super(ApiKeys.EXPIRE_DELEGATION_TOKEN);
+            this.data = data;
+        }
+
+        @Override
+        public ExpireDelegationTokenRequest build(short version) {
+            return new ExpireDelegationTokenRequest(data, version);
+        }
+
+        @Override
+        public String 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
new file mode 100644
index 0000000000000..163ee78d0adcd
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/ExpireDelegationTokenResponse.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.requests;
+
+import java.nio.ByteBuffer;
+import java.util.Map;
+
+import org.apache.kafka.common.message.ExpireDelegationTokenResponseData;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.Errors;
+
+public class ExpireDelegationTokenResponse extends AbstractResponse {
+
+    private final ExpireDelegationTokenResponseData data;
+
+    public ExpireDelegationTokenResponse(ExpireDelegationTokenResponseData data) {
+        super(ApiKeys.EXPIRE_DELEGATION_TOKEN);
+        this.data = data;
+    }
+
+    public static ExpireDelegationTokenResponse parse(ByteBuffer buffer, short version) {
+        return new ExpireDelegationTokenResponse(new ExpireDelegationTokenResponseData(new ByteBufferAccessor(buffer),
+            version));
+    }
+
+    public Errors error() {
+        return Errors.forCode(data.errorCode());
+    }
+
+    public long expiryTimestamp() {
+        return data.expiryTimestampMs();
+    }
+
+    @Override
+    public Map errorCounts() {
+        return errorCounts(error());
+    }
+
+    @Override
+    public ExpireDelegationTokenResponseData data() {
+        return data;
+    }
+
+    @Override
+    public int throttleTimeMs() {
+        return data.throttleTimeMs();
+    }
+
+    public boolean hasError() {
+        return error() != Errors.NONE;
+    }
+
+    @Override
+    public boolean shouldClientThrottle(short version) {
+        return version >= 1;
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/FetchMetadata.java b/clients/src/main/java/org/apache/kafka/common/requests/FetchMetadata.java
new file mode 100644
index 0000000000000..feb6953f9dafd
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/FetchMetadata.java
@@ -0,0 +1,154 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.Objects;
+
+public class FetchMetadata {
+    public static final Logger log = LoggerFactory.getLogger(FetchMetadata.class);
+
+    /**
+     * The session ID used by clients with no session.
+     */
+    public static final int INVALID_SESSION_ID = 0;
+
+    /**
+     * The first epoch.  When used in a fetch request, indicates that the client
+     * wants to create or recreate a session.
+     */
+    public static final int INITIAL_EPOCH = 0;
+
+    /**
+     * An invalid epoch.  When used in a fetch request, indicates that the client
+     * wants to close any existing session, and not create a new one.
+     */
+    public static final int FINAL_EPOCH = -1;
+
+    /**
+     * The FetchMetadata that is used when initializing a new FetchSessionHandler.
+     */
+    public static final FetchMetadata INITIAL = new FetchMetadata(INVALID_SESSION_ID, INITIAL_EPOCH);
+
+    /**
+     * The FetchMetadata that is implicitly used for handling older FetchRequests that
+     * don't include fetch metadata.
+     */
+    public static final FetchMetadata LEGACY = new FetchMetadata(INVALID_SESSION_ID, FINAL_EPOCH);
+
+    /**
+     * Returns the next epoch.
+     *
+     * @param prevEpoch The previous epoch.
+     * @return          The next epoch.
+     */
+    public static int nextEpoch(int prevEpoch) {
+        if (prevEpoch < 0) {
+            // The next epoch after FINAL_EPOCH is always FINAL_EPOCH itself.
+            return FINAL_EPOCH;
+        } else if (prevEpoch == Integer.MAX_VALUE) {
+            return 1;
+        } else {
+            return prevEpoch + 1;
+        }
+    }
+
+    /**
+     * The fetch session ID.
+     */
+    private final int sessionId;
+
+    /**
+     * The fetch session epoch.
+     */
+    private final int epoch;
+
+    public FetchMetadata(int sessionId, int epoch) {
+        this.sessionId = sessionId;
+        this.epoch = epoch;
+    }
+
+    /**
+     * Returns true if this is a full fetch request.
+     */
+    public boolean isFull() {
+        return (this.epoch == INITIAL_EPOCH) || (this.epoch == FINAL_EPOCH);
+    }
+
+    public int sessionId() {
+        return sessionId;
+    }
+
+    public int epoch() {
+        return epoch;
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(sessionId, epoch);
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        FetchMetadata that = (FetchMetadata) o;
+        return sessionId == that.sessionId && epoch == that.epoch;
+    }
+
+    /**
+     * Return the metadata for the next error response.
+     */
+    public FetchMetadata nextCloseExisting() {
+        return new FetchMetadata(sessionId, INITIAL_EPOCH);
+    }
+
+    /**
+     * Return the metadata for the next full fetch request.
+     */
+    public static FetchMetadata newIncremental(int sessionId) {
+        return new FetchMetadata(sessionId, nextEpoch(INITIAL_EPOCH));
+    }
+
+    /**
+     * Return the metadata for the next incremental response.
+     */
+    public FetchMetadata nextIncremental() {
+        return new FetchMetadata(sessionId, nextEpoch(epoch));
+    }
+
+    @Override
+    public String toString() {
+        StringBuilder bld = new StringBuilder();
+        if (sessionId == INVALID_SESSION_ID) {
+            bld.append("(sessionId=INVALID, ");
+        } else {
+            bld.append("(sessionId=").append(sessionId).append(", ");
+        }
+        if (epoch == INITIAL_EPOCH) {
+            bld.append("epoch=INITIAL)");
+        } else if (epoch == FINAL_EPOCH) {
+            bld.append("epoch=FINAL)");
+        } else {
+            bld.append("epoch=").append(epoch).append(")");
+        }
+        return bld.toString();
+    }
+}
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 18425d0360ba8..950b03efe14fa 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
@@ -16,162 +16,134 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.IsolationLevel;
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.message.FetchRequestData;
 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.Struct;
 import org.apache.kafka.common.record.MemoryRecords;
+import org.apache.kafka.common.record.RecordBatch;
+import org.apache.kafka.common.utils.Utils;
 
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Iterator;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
-
-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 org.apache.kafka.common.protocol.types.Type.INT64;
-import static org.apache.kafka.common.protocol.types.Type.INT8;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.stream.Collectors;
 
 public class FetchRequest extends AbstractRequest {
+
     public static final int CONSUMER_REPLICA_ID = -1;
-    private static final String REPLICA_ID_KEY_NAME = "replica_id";
-    private static final String MAX_WAIT_KEY_NAME = "max_wait_time";
-    private static final String MIN_BYTES_KEY_NAME = "min_bytes";
-    private static final String ISOLATION_LEVEL_KEY_NAME = "isolation_level";
-    private static final String TOPICS_KEY_NAME = "topics";
-
-    // request and partition level name
-    private static final String MAX_BYTES_KEY_NAME = "max_bytes";
-
-    // topic level field names
-    private static final String PARTITIONS_KEY_NAME = "partitions";
-
-    // partition level field names
-    private static final String FETCH_OFFSET_KEY_NAME = "fetch_offset";
-    private static final String LOG_START_OFFSET_KEY_NAME = "log_start_offset";
-
-    private static final Schema FETCH_REQUEST_PARTITION_V0 = new Schema(
-            PARTITION_ID,
-            new Field(FETCH_OFFSET_KEY_NAME, INT64, "Message offset."),
-            new Field(MAX_BYTES_KEY_NAME, INT32, "Maximum bytes to fetch."));
-
-    // FETCH_REQUEST_PARTITION_V5 added log_start_offset field - the earliest available offset of partition data that can be consumed.
-    private static final Schema FETCH_REQUEST_PARTITION_V5 = new Schema(
-            PARTITION_ID,
-            new Field(FETCH_OFFSET_KEY_NAME, INT64, "Message offset."),
-            new Field(LOG_START_OFFSET_KEY_NAME, INT64, "Earliest available offset of the follower replica. " +
-                            "The field is only used when request is sent by follower. "),
-            new Field(MAX_BYTES_KEY_NAME, INT32, "Maximum bytes to fetch."));
-
-    private static final Schema FETCH_REQUEST_TOPIC_V0 = new Schema(
-            TOPIC_NAME,
-            new Field(PARTITIONS_KEY_NAME, new ArrayOf(FETCH_REQUEST_PARTITION_V0), "Partitions to fetch."));
-
-    private static final Schema FETCH_REQUEST_TOPIC_V5 = new Schema(
-            TOPIC_NAME,
-            new Field(PARTITIONS_KEY_NAME, new ArrayOf(FETCH_REQUEST_PARTITION_V5), "Partitions to fetch."));
-
-    private static final Schema FETCH_REQUEST_V0 = new Schema(
-            new Field(REPLICA_ID_KEY_NAME, INT32, "Broker id of the follower. For normal consumers, use -1."),
-            new Field(MAX_WAIT_KEY_NAME, INT32, "Maximum time in ms to wait for the response."),
-            new Field(MIN_BYTES_KEY_NAME, INT32, "Minimum bytes to accumulate in the response."),
-            new Field(TOPICS_KEY_NAME, new ArrayOf(FETCH_REQUEST_TOPIC_V0), "Topics to fetch."));
-
-    // The V1 Fetch Request body is the same as V0.
-    // Only the version number is incremented to indicate a newer client
-    private static final Schema FETCH_REQUEST_V1 = FETCH_REQUEST_V0;
-    // The V2 Fetch Request body is the same as V1.
-    // Only the version number is incremented to indicate the client support message format V1 which uses
-    // relative offset and has timestamp.
-    private static final Schema FETCH_REQUEST_V2 = FETCH_REQUEST_V1;
-    // Fetch Request V3 added top level max_bytes field - the total size of partition data to accumulate in response.
-    // The partition ordering is now relevant - partitions will be processed in order they appear in request.
-    private static final Schema FETCH_REQUEST_V3 = new Schema(
-            new Field(REPLICA_ID_KEY_NAME, INT32, "Broker id of the follower. For normal consumers, use -1."),
-            new Field(MAX_WAIT_KEY_NAME, INT32, "Maximum time in ms to wait for the response."),
-            new Field(MIN_BYTES_KEY_NAME, INT32, "Minimum bytes to accumulate in the response."),
-            new Field(MAX_BYTES_KEY_NAME, INT32, "Maximum bytes to accumulate in the response. Note that this is not an absolute maximum, " +
-                    "if the first message in the first non-empty partition of the fetch is larger than this " +
-                    "value, the message will still be returned to ensure that progress can be made."),
-            new Field(TOPICS_KEY_NAME, new ArrayOf(FETCH_REQUEST_TOPIC_V0), "Topics to fetch in the order provided."));
-
-    // The V4 Fetch Request adds the fetch isolation level and exposes magic v2 (via the response).
-    private static final Schema FETCH_REQUEST_V4 = new Schema(
-            new Field(REPLICA_ID_KEY_NAME, INT32, "Broker id of the follower. For normal consumers, use -1."),
-            new Field(MAX_WAIT_KEY_NAME, INT32, "Maximum time in ms to wait for the response."),
-            new Field(MIN_BYTES_KEY_NAME, INT32, "Minimum bytes to accumulate in the response."),
-            new Field(MAX_BYTES_KEY_NAME, INT32, "Maximum bytes to accumulate in the response. Note that this is not an absolute maximum, " +
-                    "if the first message in the first non-empty partition of the fetch is larger than this " +
-                    "value, the message will still be returned to ensure that progress can be made."),
-            new Field(ISOLATION_LEVEL_KEY_NAME, INT8, "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"),
-            new Field(TOPICS_KEY_NAME, new ArrayOf(FETCH_REQUEST_TOPIC_V0), "Topics to fetch in the order provided."));
-
-    // FETCH_REQUEST_V5 added a per-partition log_start_offset field - the earliest available offset of partition data that can be consumed.
-    private static final Schema FETCH_REQUEST_V5 = new Schema(
-            new Field(REPLICA_ID_KEY_NAME, INT32, "Broker id of the follower. For normal consumers, use -1."),
-            new Field(MAX_WAIT_KEY_NAME, INT32, "Maximum time in ms to wait for the response."),
-            new Field(MIN_BYTES_KEY_NAME, INT32, "Minimum bytes to accumulate in the response."),
-            new Field(MAX_BYTES_KEY_NAME, INT32, "Maximum bytes to accumulate in the response. Note that this is not an absolute maximum, " +
-                    "if the first message in the first non-empty partition of the fetch is larger than this " +
-                    "value, the message will still be returned to ensure that progress can be made."),
-            new Field(ISOLATION_LEVEL_KEY_NAME, INT8, "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"),
-            new Field(TOPICS_KEY_NAME, new ArrayOf(FETCH_REQUEST_TOPIC_V5), "Topics to fetch in the order provided."));
-
-    /**
-     * The body of FETCH_REQUEST_V6 is the same as FETCH_REQUEST_V5.
-     * The version number is bumped up to indicate that the client supports KafkaStorageException.
-     * The KafkaStorageException will be translated to NotLeaderForPartitionException in the response if version <= 5
-     */
-    private static final Schema FETCH_REQUEST_V6 = FETCH_REQUEST_V5;
-
-    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};
-    };
 
     // default values for older versions where a request level limit did not exist
     public static final int DEFAULT_RESPONSE_MAX_BYTES = Integer.MAX_VALUE;
     public static final long INVALID_LOG_START_OFFSET = -1L;
 
-    private final int replicaId;
-    private final int maxWait;
-    private final int minBytes;
-    private final int maxBytes;
-    private final IsolationLevel isolationLevel;
-    private final LinkedHashMap fetchData;
+    private final FetchRequestData data;
+
+    // These are immutable read-only structures derived from FetchRequestData
+    private final Map fetchData;
+    private final List toForget;
+    private final FetchMetadata metadata;
 
     public static final class PartitionData {
         public final long fetchOffset;
         public final long logStartOffset;
         public final int maxBytes;
+        public final Optional currentLeaderEpoch;
+        public final Optional lastFetchedEpoch;
+
+        public PartitionData(
+            long fetchOffset,
+            long logStartOffset,
+            int maxBytes,
+            Optional currentLeaderEpoch
+        ) {
+            this(fetchOffset, logStartOffset, maxBytes, currentLeaderEpoch, Optional.empty());
+        }
 
-        public PartitionData(long fetchOffset, long logStartOffset, int maxBytes) {
+        public PartitionData(
+            long fetchOffset,
+            long logStartOffset,
+            int maxBytes,
+            Optional currentLeaderEpoch,
+            Optional lastFetchedEpoch
+        ) {
             this.fetchOffset = fetchOffset;
             this.logStartOffset = logStartOffset;
             this.maxBytes = maxBytes;
+            this.currentLeaderEpoch = currentLeaderEpoch;
+            this.lastFetchedEpoch = lastFetchedEpoch;
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (o == null || getClass() != o.getClass()) return false;
+            PartitionData that = (PartitionData) o;
+            return fetchOffset == that.fetchOffset &&
+                logStartOffset == that.logStartOffset &&
+                maxBytes == that.maxBytes &&
+                Objects.equals(currentLeaderEpoch, that.currentLeaderEpoch) &&
+                Objects.equals(lastFetchedEpoch, that.lastFetchedEpoch);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(fetchOffset, logStartOffset, maxBytes, currentLeaderEpoch, lastFetchedEpoch);
         }
 
         @Override
         public String toString() {
-            return "(offset=" + fetchOffset + ", logStartOffset=" + logStartOffset + ", maxBytes=" + maxBytes + ")";
+            return "PartitionData(" +
+                "fetchOffset=" + fetchOffset +
+                ", logStartOffset=" + logStartOffset +
+                ", maxBytes=" + maxBytes +
+                ", currentLeaderEpoch=" + currentLeaderEpoch +
+                ", lastFetchedEpoch=" + lastFetchedEpoch +
+                ')';
+        }
+    }
+
+    private Optional optionalEpoch(int rawEpochValue) {
+        if (rawEpochValue < 0) {
+            return Optional.empty();
+        } else {
+            return Optional.of(rawEpochValue);
         }
     }
 
+    private Map toPartitionDataMap(List fetchableTopics) {
+        Map result = new LinkedHashMap<>();
+        fetchableTopics.forEach(fetchTopic -> fetchTopic.partitions().forEach(fetchPartition -> {
+            result.put(new TopicPartition(fetchTopic.topic(), fetchPartition.partition()),
+                new PartitionData(
+                    fetchPartition.fetchOffset(),
+                    fetchPartition.logStartOffset(),
+                    fetchPartition.partitionMaxBytes(),
+                    optionalEpoch(fetchPartition.currentLeaderEpoch()),
+                    optionalEpoch(fetchPartition.lastFetchedEpoch())
+                ));
+        }));
+        return Collections.unmodifiableMap(result);
+    }
+
+    private List toForgottenTopicList(List forgottenTopics) {
+        List result = new ArrayList<>();
+        forgottenTopics.forEach(forgottenTopic ->
+            forgottenTopic.partitions().forEach(partitionId ->
+                result.add(new TopicPartition(forgottenTopic.topic(), partitionId))
+            )
+        );
+        return result;
+    }
+
     static final class TopicAndPartitionData {
         public final String topic;
         public final LinkedHashMap partitions;
@@ -181,9 +153,10 @@ public TopicAndPartitionData(String topic) {
             this.partitions = new LinkedHashMap<>();
         }
 
-        public static  List> batchByTopic(LinkedHashMap data) {
+        public static  List> batchByTopic(Iterator> iter) {
             List> topics = new ArrayList<>();
-            for (Map.Entry topicEntry : data.entrySet()) {
+            while (iter.hasNext()) {
+                Map.Entry topicEntry = iter.next();
                 String topic = topicEntry.getKey().topic();
                 int partition = topicEntry.getKey().partition();
                 T partitionData = topicEntry.getValue();
@@ -199,37 +172,48 @@ public static class Builder extends AbstractRequest.Builder {
         private final int maxWait;
         private final int minBytes;
         private final int replicaId;
-        private final LinkedHashMap fetchData;
-        private final IsolationLevel isolationLevel;
+        private final Map fetchData;
+        private IsolationLevel isolationLevel = IsolationLevel.READ_UNCOMMITTED;
         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, LinkedHashMap fetchData) {
-            return forConsumer(maxWait, minBytes, fetchData, IsolationLevel.READ_UNCOMMITTED);
-        }
-
-        public static Builder forConsumer(int maxWait, int minBytes, LinkedHashMap fetchData,
-                                          IsolationLevel isolationLevel) {
-            return new Builder(ApiKeys.FETCH.oldestVersion(), ApiKeys.FETCH.latestVersion(), CONSUMER_REPLICA_ID,
-                    maxWait, minBytes, fetchData, isolationLevel);
+        public static Builder forConsumer(int maxWait, int minBytes, Map fetchData) {
+            return new Builder(ApiKeys.FETCH.oldestVersion(), ApiKeys.FETCH.latestVersion(),
+                CONSUMER_REPLICA_ID, maxWait, minBytes, fetchData);
         }
 
         public static Builder forReplica(short allowedVersion, int replicaId, int maxWait, int minBytes,
-                                         LinkedHashMap fetchData) {
-            return new Builder(allowedVersion, allowedVersion, replicaId, maxWait, minBytes, fetchData,
-                    IsolationLevel.READ_UNCOMMITTED);
+                                         Map fetchData) {
+            return new Builder(allowedVersion, allowedVersion, replicaId, maxWait, minBytes, fetchData);
         }
 
-        private Builder(short minVersion, short maxVersion, int replicaId, int maxWait, int minBytes,
-                        LinkedHashMap fetchData, IsolationLevel isolationLevel) {
+        public Builder(short minVersion, short maxVersion, int replicaId, int maxWait, int minBytes,
+                        Map fetchData) {
             super(ApiKeys.FETCH, minVersion, maxVersion);
             this.replicaId = replicaId;
             this.maxWait = maxWait;
             this.minBytes = minBytes;
             this.fetchData = fetchData;
+        }
+
+        public Builder isolationLevel(IsolationLevel isolationLevel) {
             this.isolationLevel = isolationLevel;
+            return this;
+        }
+
+        public Builder metadata(FetchMetadata metadata) {
+            this.metadata = metadata;
+            return this;
         }
 
-        public LinkedHashMap fetchData() {
+        public Builder rackId(String rackId) {
+            this.rackId = rackId;
+            return this;
+        }
+
+        public Map fetchData() {
             return this.fetchData;
         }
 
@@ -238,13 +222,68 @@ public Builder setMaxBytes(int maxBytes) {
             return this;
         }
 
+        public List toForget() {
+            return toForget;
+        }
+
+        public Builder toForget(List toForget) {
+            this.toForget = toForget;
+            return this;
+        }
+
         @Override
         public FetchRequest build(short version) {
             if (version < 3) {
                 maxBytes = DEFAULT_RESPONSE_MAX_BYTES;
             }
 
-            return new FetchRequest(version, replicaId, maxWait, minBytes, maxBytes, fetchData, isolationLevel);
+            FetchRequestData fetchRequestData = new FetchRequestData();
+            fetchRequestData.setReplicaId(replicaId);
+            fetchRequestData.setMaxWaitMs(maxWait);
+            fetchRequestData.setMinBytes(minBytes);
+            fetchRequestData.setMaxBytes(maxBytes);
+            fetchRequestData.setIsolationLevel(isolationLevel.id());
+            fetchRequestData.setForgottenTopicsData(new ArrayList<>());
+            toForget.stream()
+                .collect(Collectors.groupingBy(TopicPartition::topic, LinkedHashMap::new, Collectors.toList()))
+                .forEach((topic, partitions) ->
+                    fetchRequestData.forgottenTopicsData().add(new FetchRequestData.ForgottenTopic()
+                        .setTopic(topic)
+                        .setPartitions(partitions.stream().map(TopicPartition::partition).collect(Collectors.toList())))
+                );
+            fetchRequestData.setTopics(new ArrayList<>());
+
+            // We collect the partitions in a single FetchTopic only if they appear sequentially in the fetchData
+            FetchRequestData.FetchTopic fetchTopic = null;
+            for (Map.Entry entry : fetchData.entrySet()) {
+                TopicPartition topicPartition = entry.getKey();
+                PartitionData partitionData = entry.getValue();
+
+                if (fetchTopic == null || !topicPartition.topic().equals(fetchTopic.topic())) {
+                    fetchTopic = new FetchRequestData.FetchTopic()
+                       .setTopic(topicPartition.topic())
+                       .setPartitions(new ArrayList<>());
+                    fetchRequestData.topics().add(fetchTopic);
+                }
+
+                FetchRequestData.FetchPartition fetchPartition = new FetchRequestData.FetchPartition()
+                    .setPartition(topicPartition.partition())
+                    .setCurrentLeaderEpoch(partitionData.currentLeaderEpoch.orElse(RecordBatch.NO_PARTITION_LEADER_EPOCH))
+                    .setLastFetchedEpoch(partitionData.lastFetchedEpoch.orElse(RecordBatch.NO_PARTITION_LEADER_EPOCH))
+                    .setFetchOffset(partitionData.fetchOffset)
+                    .setLogStartOffset(partitionData.logStartOffset)
+                    .setPartitionMaxBytes(partitionData.maxBytes);
+
+                fetchTopic.partitions().add(fetchPartition);
+            }
+
+            if (metadata != null) {
+                fetchRequestData.setSessionEpoch(metadata.epoch());
+                fetchRequestData.setSessionId(metadata.sessionId());
+            }
+            fetchRequestData.setRackId(rackId);
+
+            return new FetchRequest(fetchRequestData, version);
         }
 
         @Override
@@ -257,131 +296,87 @@ public String toString() {
                     append(", maxBytes=").append(maxBytes).
                     append(", fetchData=").append(fetchData).
                     append(", isolationLevel=").append(isolationLevel).
+                    append(", toForget=").append(Utils.join(toForget, ", ")).
+                    append(", metadata=").append(metadata).
+                    append(", rackId=").append(rackId).
                     append(")");
             return bld.toString();
         }
     }
 
-    private FetchRequest(short version, int replicaId, int maxWait, int minBytes, int maxBytes,
-                         LinkedHashMap fetchData, IsolationLevel isolationLevel) {
-        super(version);
-        this.replicaId = replicaId;
-        this.maxWait = maxWait;
-        this.minBytes = minBytes;
-        this.maxBytes = maxBytes;
-        this.fetchData = fetchData;
-        this.isolationLevel = isolationLevel;
-    }
-
-    public FetchRequest(Struct struct, short version) {
-        super(version);
-        replicaId = struct.getInt(REPLICA_ID_KEY_NAME);
-        maxWait = struct.getInt(MAX_WAIT_KEY_NAME);
-        minBytes = struct.getInt(MIN_BYTES_KEY_NAME);
-        if (struct.hasField(MAX_BYTES_KEY_NAME))
-            maxBytes = struct.getInt(MAX_BYTES_KEY_NAME);
-        else
-            maxBytes = DEFAULT_RESPONSE_MAX_BYTES;
-
-        if (struct.hasField(ISOLATION_LEVEL_KEY_NAME))
-            isolationLevel = IsolationLevel.forId(struct.getByte(ISOLATION_LEVEL_KEY_NAME));
-        else
-            isolationLevel = IsolationLevel.READ_UNCOMMITTED;
-
-        fetchData = new LinkedHashMap<>();
-        for (Object topicResponseObj : struct.getArray(TOPICS_KEY_NAME)) {
-            Struct topicResponse = (Struct) topicResponseObj;
-            String topic = topicResponse.get(TOPIC_NAME);
-            for (Object partitionResponseObj : topicResponse.getArray(PARTITIONS_KEY_NAME)) {
-                Struct partitionResponse = (Struct) partitionResponseObj;
-                int partition = partitionResponse.get(PARTITION_ID);
-                long offset = partitionResponse.getLong(FETCH_OFFSET_KEY_NAME);
-                int maxBytes = partitionResponse.getInt(MAX_BYTES_KEY_NAME);
-                long logStartOffset = partitionResponse.hasField(LOG_START_OFFSET_KEY_NAME) ?
-                    partitionResponse.getLong(LOG_START_OFFSET_KEY_NAME) : INVALID_LOG_START_OFFSET;
-                PartitionData partitionData = new PartitionData(offset, logStartOffset, maxBytes);
-                fetchData.put(new TopicPartition(topic, partition), partitionData);
-            }
-        }
+    public FetchRequest(FetchRequestData fetchRequestData, short version) {
+        super(ApiKeys.FETCH, version);
+        this.data = fetchRequestData;
+        this.fetchData = toPartitionDataMap(fetchRequestData.topics());
+        this.toForget = toForgottenTopicList(fetchRequestData.forgottenTopicsData());
+        this.metadata = new FetchMetadata(fetchRequestData.sessionId(), fetchRequestData.sessionEpoch());
     }
 
     @Override
     public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
-        LinkedHashMap responseData = new LinkedHashMap<>();
-
-        for (Map.Entry entry: fetchData.entrySet()) {
-            FetchResponse.PartitionData partitionResponse = new FetchResponse.PartitionData(Errors.forException(e),
-                FetchResponse.INVALID_HIGHWATERMARK, FetchResponse.INVALID_LAST_STABLE_OFFSET, FetchResponse.INVALID_LOG_START_OFFSET,
-                null, MemoryRecords.EMPTY);
+        // The error is indicated in two ways: by setting the same error code in all partitions, and by
+        // setting the top-level error code.  The form where we set the same error code in all partitions
+        // is needed in order to maintain backwards compatibility with older versions of the protocol
+        // in which there was no top-level error code. Note that for incremental fetch responses, there
+        // may not be any partitions at all in the response.  For this reason, the top-level error code
+        // is essential for them.
+        Errors error = Errors.forException(e);
+        LinkedHashMap> responseData = new LinkedHashMap<>();
+        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, Optional.empty(), null, MemoryRecords.EMPTY);
             responseData.put(entry.getKey(), partitionResponse);
         }
-        return new FetchResponse(responseData, throttleTimeMs);
+        return new FetchResponse<>(error, responseData, throttleTimeMs, data.sessionId());
     }
 
     public int replicaId() {
-        return replicaId;
+        return data.replicaId();
     }
 
     public int maxWait() {
-        return maxWait;
+        return data.maxWaitMs();
     }
 
     public int minBytes() {
-        return minBytes;
+        return data.minBytes();
     }
 
     public int maxBytes() {
-        return maxBytes;
+        return data.maxBytes();
     }
 
     public Map fetchData() {
         return fetchData;
     }
 
+    public List toForget() {
+        return toForget;
+    }
+
     public boolean isFromFollower() {
-        return replicaId >= 0;
+        return replicaId() >= 0;
     }
 
     public IsolationLevel isolationLevel() {
-        return isolationLevel;
+        return IsolationLevel.forId(data.isolationLevel());
+    }
+
+    public FetchMetadata metadata() {
+        return metadata;
+    }
+
+    public String rackId() {
+        return data.rackId();
     }
 
     public static FetchRequest parse(ByteBuffer buffer, short version) {
-        return new FetchRequest(ApiKeys.FETCH.parseRequest(version, buffer), version);
+        return new FetchRequest(new FetchRequestData(new ByteBufferAccessor(buffer), version), version);
     }
 
     @Override
-    protected Struct toStruct() {
-        Struct struct = new Struct(ApiKeys.FETCH.requestSchema(version()));
-        List> topicsData = TopicAndPartitionData.batchByTopic(fetchData);
-
-        struct.set(REPLICA_ID_KEY_NAME, replicaId);
-        struct.set(MAX_WAIT_KEY_NAME, maxWait);
-        struct.set(MIN_BYTES_KEY_NAME, minBytes);
-        if (struct.hasField(MAX_BYTES_KEY_NAME))
-            struct.set(MAX_BYTES_KEY_NAME, maxBytes);
-        if (struct.hasField(ISOLATION_LEVEL_KEY_NAME))
-            struct.set(ISOLATION_LEVEL_KEY_NAME, isolationLevel.id());
-
-        List topicArray = new ArrayList<>();
-        for (TopicAndPartitionData topicEntry : topicsData) {
-            Struct topicData = struct.instance(TOPICS_KEY_NAME);
-            topicData.set(TOPIC_NAME, topicEntry.topic);
-            List partitionArray = new ArrayList<>();
-            for (Map.Entry partitionEntry : topicEntry.partitions.entrySet()) {
-                PartitionData fetchPartitionData = partitionEntry.getValue();
-                Struct partitionData = topicData.instance(PARTITIONS_KEY_NAME);
-                partitionData.set(PARTITION_ID, partitionEntry.getKey());
-                partitionData.set(FETCH_OFFSET_KEY_NAME, fetchPartitionData.fetchOffset);
-                if (partitionData.hasField(LOG_START_OFFSET_KEY_NAME))
-                    partitionData.set(LOG_START_OFFSET_KEY_NAME, fetchPartitionData.logStartOffset);
-                partitionData.set(MAX_BYTES_KEY_NAME, fetchPartitionData.maxBytes);
-                partitionArray.add(partitionData);
-            }
-            topicData.set(PARTITIONS_KEY_NAME, partitionArray.toArray());
-            topicArray.add(topicData);
-        }
-        struct.set(TOPICS_KEY_NAME, topicArray.toArray());
-        return struct;
+    public FetchRequestData data() {
+        return data;
     }
 }
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 0d09027802ab8..30b462bd21933 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
@@ -17,159 +17,60 @@
 package org.apache.kafka.common.requests;
 
 import org.apache.kafka.common.TopicPartition;
-import org.apache.kafka.common.network.ByteBufferSend;
-import org.apache.kafka.common.network.MultiSend;
-import org.apache.kafka.common.network.Send;
+import org.apache.kafka.common.message.FetchResponseData;
 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.Struct;
-import org.apache.kafka.common.record.Records;
+import org.apache.kafka.common.protocol.ObjectSerializationCache;
+import org.apache.kafka.common.record.BaseRecords;
+import org.apache.kafka.common.record.MemoryRecords;
 
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.Iterator;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
+import java.util.stream.Collectors;
 
-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.protocol.types.Type.INT64;
-import static org.apache.kafka.common.protocol.types.Type.RECORDS;
-import static org.apache.kafka.common.protocol.types.Type.STRING;
+import static org.apache.kafka.common.requests.FetchMetadata.INVALID_SESSION_ID;
 
 /**
  * This wrapper supports all versions of the Fetch API
+ *
+ * Possible error codes:
+ *
+ * - {@link Errors#OFFSET_OUT_OF_RANGE} If the fetch offset is out of range for a requested partition
+ * - {@link Errors#TOPIC_AUTHORIZATION_FAILED} If the user does not have READ access to a requested topic
+ * - {@link Errors#REPLICA_NOT_AVAILABLE} If the request is received by a broker with version < 2.6 which is not a replica
+ * - {@link Errors#NOT_LEADER_OR_FOLLOWER} If the broker is not a leader or follower and either the provided leader epoch
+ *     matches the known leader epoch on the broker or is empty
+ * - {@link Errors#FENCED_LEADER_EPOCH} If the epoch is lower than the broker's epoch
+ * - {@link Errors#UNKNOWN_LEADER_EPOCH} If the epoch is larger than the broker's epoch
+ * - {@link Errors#UNKNOWN_TOPIC_OR_PARTITION} If the broker does not have metadata for a topic or partition
+ * - {@link Errors#KAFKA_STORAGE_ERROR} If the log directory for one of the requested partitions is offline
+ * - {@link Errors#UNSUPPORTED_COMPRESSION_TYPE} If a fetched topic is using a compression type which is
+ *     not supported by the fetch request version
+ * - {@link Errors#CORRUPT_MESSAGE} If corrupt message encountered, e.g. when the broker scans the log to find
+ *     the fetch offset after the index lookup
+ * - {@link Errors#UNKNOWN_SERVER_ERROR} For any unexpected errors
  */
-public class FetchResponse extends AbstractResponse {
-
-    private static final String RESPONSES_KEY_NAME = "responses";
-
-    // topic level field names
-    private static final String PARTITIONS_KEY_NAME = "partition_responses";
-
-    // partition level field names
-    private static final String PARTITION_HEADER_KEY_NAME = "partition_header";
-    private static final String HIGH_WATERMARK_KEY_NAME = "high_watermark";
-    private static final String LAST_STABLE_OFFSET_KEY_NAME = "last_stable_offset";
-    private static final String LOG_START_OFFSET_KEY_NAME = "log_start_offset";
-    private static final String ABORTED_TRANSACTIONS_KEY_NAME = "aborted_transactions";
-    private static final String RECORD_SET_KEY_NAME = "record_set";
-
-    // aborted transaction field names
-    private static final String PRODUCER_ID_KEY_NAME = "producer_id";
-    private static final String FIRST_OFFSET_KEY_NAME = "first_offset";
-
-    private static final Schema FETCH_RESPONSE_PARTITION_HEADER_V0 = new Schema(
-            PARTITION_ID,
-            ERROR_CODE,
-            new Field(HIGH_WATERMARK_KEY_NAME, INT64, "Last committed offset."));
-    private static final Schema FETCH_RESPONSE_PARTITION_V0 = new Schema(
-            new Field(PARTITION_HEADER_KEY_NAME, FETCH_RESPONSE_PARTITION_HEADER_V0),
-            new Field(RECORD_SET_KEY_NAME, RECORDS));
-
-    private static final Schema FETCH_RESPONSE_TOPIC_V0 = new Schema(
-            TOPIC_NAME,
-            new Field(PARTITIONS_KEY_NAME, new ArrayOf(FETCH_RESPONSE_PARTITION_V0)));
-
-    private static final Schema FETCH_RESPONSE_V0 = new Schema(
-            new Field(RESPONSES_KEY_NAME, new ArrayOf(FETCH_RESPONSE_TOPIC_V0)));
-
-    private static final Schema FETCH_RESPONSE_V1 = new Schema(
-            THROTTLE_TIME_MS,
-            new Field(RESPONSES_KEY_NAME, new ArrayOf(FETCH_RESPONSE_TOPIC_V0)));
-    // Even though fetch response v2 has the same protocol as v1, the record set in the response is different. In v1,
-    // record set only includes messages of v0 (magic byte 0). In v2, record set can include messages of v0 and v1
-    // (magic byte 0 and 1). For details, see Records, RecordBatch and Record.
-    private static final Schema FETCH_RESPONSE_V2 = FETCH_RESPONSE_V1;
-
-    // The partition ordering is now relevant - partitions will be processed in order they appear in request.
-    private static final Schema FETCH_RESPONSE_V3 = FETCH_RESPONSE_V2;
-
-    // The v4 Fetch Response adds features for transactional consumption (the aborted transaction list and the
-    // last stable offset). It also exposes messages with magic v2 (along with older formats).
-    private static final Schema FETCH_RESPONSE_ABORTED_TRANSACTION_V4 = new Schema(
-            new Field(PRODUCER_ID_KEY_NAME, INT64, "The producer id associated with the aborted transactions"),
-            new Field(FIRST_OFFSET_KEY_NAME, INT64, "The first offset in the aborted transaction"));
-
-    private static final Schema FETCH_RESPONSE_ABORTED_TRANSACTION_V5 = FETCH_RESPONSE_ABORTED_TRANSACTION_V4;
-
-    private static final Schema FETCH_RESPONSE_PARTITION_HEADER_V4 = new Schema(
-            PARTITION_ID,
-            ERROR_CODE,
-            new Field(HIGH_WATERMARK_KEY_NAME, INT64, "Last committed offset."),
-            new Field(LAST_STABLE_OFFSET_KEY_NAME, INT64, "The last stable offset (or LSO) of the partition. This is the last offset such that the state " +
-                    "of all transactional records prior to this offset have been decided (ABORTED or COMMITTED)"),
-            new Field(ABORTED_TRANSACTIONS_KEY_NAME, ArrayOf.nullable(FETCH_RESPONSE_ABORTED_TRANSACTION_V4)));
-
-    // FETCH_RESPONSE_PARTITION_HEADER_V5 added log_start_offset field - the earliest available offset of partition data that can be consumed.
-    private static final Schema FETCH_RESPONSE_PARTITION_HEADER_V5 = new Schema(
-            PARTITION_ID,
-            ERROR_CODE,
-            new Field(HIGH_WATERMARK_KEY_NAME, INT64, "Last committed offset."),
-            new Field(LAST_STABLE_OFFSET_KEY_NAME, INT64, "The last stable offset (or LSO) of the partition. This is the last offset such that the state " +
-                    "of all transactional records prior to this offset have been decided (ABORTED or COMMITTED)"),
-            new Field(LOG_START_OFFSET_KEY_NAME, INT64, "Earliest available offset."),
-            new Field(ABORTED_TRANSACTIONS_KEY_NAME, ArrayOf.nullable(FETCH_RESPONSE_ABORTED_TRANSACTION_V5)));
-
-    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));
-
-    private static final Schema FETCH_RESPONSE_PARTITION_V5 = new Schema(
-            new Field(PARTITION_HEADER_KEY_NAME, FETCH_RESPONSE_PARTITION_HEADER_V5),
-            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)));
-
-    private static final Schema FETCH_RESPONSE_TOPIC_V5 = new Schema(
-            TOPIC_NAME,
-            new Field(PARTITIONS_KEY_NAME, new ArrayOf(FETCH_RESPONSE_PARTITION_V5)));
-
-    private static final Schema FETCH_RESPONSE_V4 = new Schema(
-            THROTTLE_TIME_MS,
-            new Field(RESPONSES_KEY_NAME, new ArrayOf(FETCH_RESPONSE_TOPIC_V4)));
-
-    private static final Schema FETCH_RESPONSE_V5 = new Schema(
-            THROTTLE_TIME_MS,
-            new Field(RESPONSES_KEY_NAME, new ArrayOf(FETCH_RESPONSE_TOPIC_V5)));
-
-    /**
-     * The body of FETCH_RESPONSE_V6 is the same as FETCH_RESPONSE_V5.
-     * The version number is bumped up to indicate that the client supports KafkaStorageException.
-     * The KafkaStorageException will be translated to NotLeaderForPartitionException in the response if version <= 5
-     */
-    private static final Schema FETCH_RESPONSE_V6 = FETCH_RESPONSE_V5;
-
-    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};
-    }
-
+public class FetchResponse extends AbstractResponse {
 
     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;
 
-    /**
-     * Possible error codes:
-     *
-     *  OFFSET_OUT_OF_RANGE (1)
-     *  UNKNOWN_TOPIC_OR_PARTITION (3)
-     *  NOT_LEADER_FOR_PARTITION (6)
-     *  REPLICA_NOT_AVAILABLE (9)
-     *  UNKNOWN (-1)
-     */
+    private final FetchResponseData data;
+    private final LinkedHashMap> responseDataMap;
 
-    private final LinkedHashMap responseData;
-    private final int throttleTimeMs;
+    @Override
+    public FetchResponseData data() {
+        return data;
+    }
 
     public static final class AbortedTransaction {
         public final long producerId;
@@ -194,8 +95,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;
         }
 
@@ -203,28 +104,91 @@ public int hashCode() {
         public String toString() {
             return "(producerId=" + producerId + ", firstOffset=" + firstOffset + ")";
         }
+
+        static AbortedTransaction fromMessage(FetchResponseData.AbortedTransaction abortedTransaction) {
+            return new AbortedTransaction(abortedTransaction.producerId(), abortedTransaction.firstOffset());
+        }
     }
 
-    public static final class PartitionData {
-        public final Errors error;
-        public final long highWatermark;
-        public final long lastStableOffset;
-        public final long logStartOffset;
-        public final List abortedTransactions;
-        public final Records records;
+    public static final class PartitionData {
+        private final FetchResponseData.FetchablePartitionResponse partitionResponse;
+
+        // Derived fields
+        private final Optional preferredReplica;
+        private final List abortedTransactions;
+        private final Errors error;
+
+        private PartitionData(FetchResponseData.FetchablePartitionResponse partitionResponse) {
+            // We partially construct FetchablePartitionResponse since we don't know the partition ID at this point
+            // When we convert the PartitionData (and other fields) into FetchResponseData down in toMessage, we
+            // set the partition IDs.
+            this.partitionResponse = partitionResponse;
+            this.preferredReplica = Optional.of(partitionResponse.preferredReadReplica())
+                .filter(replicaId -> replicaId != INVALID_PREFERRED_REPLICA_ID);
+
+            if (partitionResponse.abortedTransactions() == null) {
+                this.abortedTransactions = null;
+            } else {
+                this.abortedTransactions = partitionResponse.abortedTransactions().stream()
+                    .map(AbortedTransaction::fromMessage)
+                    .collect(Collectors.toList());
+            }
+
+            this.error = Errors.forCode(partitionResponse.errorCode());
+        }
 
         public PartitionData(Errors error,
                              long highWatermark,
                              long lastStableOffset,
                              long logStartOffset,
+                             Optional preferredReadReplica,
                              List abortedTransactions,
-                             Records records) {
-            this.error = error;
-            this.highWatermark = highWatermark;
-            this.lastStableOffset = lastStableOffset;
-            this.logStartOffset = logStartOffset;
+                             Optional divergingEpoch,
+                             T records) {
+            this.preferredReplica = preferredReadReplica;
             this.abortedTransactions = abortedTransactions;
-            this.records = records;
+            this.error = error;
+
+            FetchResponseData.FetchablePartitionResponse partitionResponse =
+                new FetchResponseData.FetchablePartitionResponse();
+            partitionResponse.setErrorCode(error.code())
+                .setHighWatermark(highWatermark)
+                .setLastStableOffset(lastStableOffset)
+                .setLogStartOffset(logStartOffset);
+            if (abortedTransactions != null) {
+                partitionResponse.setAbortedTransactions(abortedTransactions.stream().map(
+                    aborted -> new FetchResponseData.AbortedTransaction()
+                        .setProducerId(aborted.producerId)
+                        .setFirstOffset(aborted.firstOffset))
+                    .collect(Collectors.toList()));
+            } else {
+                partitionResponse.setAbortedTransactions(null);
+            }
+            partitionResponse.setPreferredReadReplica(preferredReadReplica.orElse(INVALID_PREFERRED_REPLICA_ID));
+            partitionResponse.setRecordSet(records);
+            divergingEpoch.ifPresent(partitionResponse::setDivergingEpoch);
+
+            this.partitionResponse = partitionResponse;
+        }
+
+        public PartitionData(Errors error,
+                             long highWatermark,
+                             long lastStableOffset,
+                             long logStartOffset,
+                             Optional preferredReadReplica,
+                             List abortedTransactions,
+                             T records) {
+            this(error, highWatermark, lastStableOffset, logStartOffset, preferredReadReplica,
+                abortedTransactions, Optional.empty(), records);
+        }
+
+        public PartitionData(Errors error,
+                             long highWatermark,
+                             long lastStableOffset,
+                             long logStartOffset,
+                             List abortedTransactions,
+                             T records) {
+            this(error, highWatermark, lastStableOffset, logStartOffset, Optional.empty(), abortedTransactions, records);
         }
 
         @Override
@@ -236,243 +200,179 @@ public boolean equals(Object o) {
 
             PartitionData that = (PartitionData) o;
 
-            return error == that.error &&
-                    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));
+            return this.partitionResponse.equals(that.partitionResponse);
         }
 
         @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 + (abortedTransactions != null ? abortedTransactions.hashCode() : 0);
-            result = 31 * result + (records != null ? records.hashCode() : 0);
-            return result;
+            return this.partitionResponse.hashCode();
         }
 
         @Override
         public String toString() {
-            return "(error=" + error +
-                    ", highWaterMark=" + highWatermark +
-                    ", lastStableOffset = " + lastStableOffset +
-                    ", logStartOffset = " + logStartOffset +
-                    ", abortedTransactions = " + abortedTransactions +
-                    ", recordsSizeInBytes=" + records.sizeInBytes() + ")";
+            return "(error=" + error() +
+                    ", highWaterMark=" + highWatermark() +
+                    ", lastStableOffset = " + lastStableOffset() +
+                    ", logStartOffset = " + logStartOffset() +
+                    ", preferredReadReplica = " + preferredReadReplica().map(Object::toString).orElse("absent") +
+                    ", abortedTransactions = " + abortedTransactions() +
+                    ", divergingEpoch =" + divergingEpoch() +
+                    ", recordsSizeInBytes=" + records().sizeInBytes() + ")";
+        }
+
+        public Errors error() {
+            return error;
+        }
+
+        public long highWatermark() {
+            return partitionResponse.highWatermark();
+        }
+
+        public long lastStableOffset() {
+            return partitionResponse.lastStableOffset();
+        }
+
+        public long logStartOffset() {
+            return partitionResponse.logStartOffset();
+        }
+
+        public Optional preferredReadReplica() {
+            return preferredReplica;
+        }
+
+        public List abortedTransactions() {
+            return abortedTransactions;
         }
 
+        public Optional divergingEpoch() {
+            FetchResponseData.EpochEndOffset epochEndOffset = partitionResponse.divergingEpoch();
+            if (epochEndOffset.epoch() < 0) {
+                return Optional.empty();
+            } else {
+                return Optional.of(epochEndOffset);
+            }
+        }
+
+        @SuppressWarnings("unchecked")
+        public T records() {
+            return (T) partitionResponse.recordSet();
+        }
     }
 
     /**
-     * Constructor for all versions.
-     *
      * From version 3 or later, the entries in `responseData` should be in the same order as the entries in
      * `FetchRequest.fetchData`.
      *
-     * @param responseData fetched data grouped by topic-partition
-     * @param throttleTimeMs Time in milliseconds the response was throttled
+     * @param error             The top-level error code.
+     * @param responseData      The fetched data grouped by partition.
+     * @param throttleTimeMs    The time in milliseconds that the response was throttled
+     * @param sessionId         The fetch session id.
      */
-    public FetchResponse(LinkedHashMap responseData, int throttleTimeMs) {
-        this.responseData = responseData;
-        this.throttleTimeMs = throttleTimeMs;
+    public FetchResponse(Errors error,
+                         LinkedHashMap> responseData,
+                         int throttleTimeMs,
+                         int sessionId) {
+        super(ApiKeys.FETCH);
+        this.data = toMessage(throttleTimeMs, error, responseData.entrySet().iterator(), sessionId);
+        this.responseDataMap = responseData;
     }
 
-    public FetchResponse(Struct struct) {
-        LinkedHashMap responseData = new LinkedHashMap<>();
-        for (Object topicResponseObj : struct.getArray(RESPONSES_KEY_NAME)) {
-            Struct topicResponse = (Struct) topicResponseObj;
-            String topic = topicResponse.get(TOPIC_NAME);
-            for (Object partitionResponseObj : topicResponse.getArray(PARTITIONS_KEY_NAME)) {
-                Struct partitionResponse = (Struct) partitionResponseObj;
-                Struct partitionResponseHeader = partitionResponse.getStruct(PARTITION_HEADER_KEY_NAME);
-                int partition = partitionResponseHeader.get(PARTITION_ID);
-                Errors error = Errors.forCode(partitionResponseHeader.get(ERROR_CODE));
-                long highWatermark = partitionResponseHeader.getLong(HIGH_WATERMARK_KEY_NAME);
-                long lastStableOffset = INVALID_LAST_STABLE_OFFSET;
-                if (partitionResponseHeader.hasField(LAST_STABLE_OFFSET_KEY_NAME))
-                    lastStableOffset = partitionResponseHeader.getLong(LAST_STABLE_OFFSET_KEY_NAME);
-                long logStartOffset = INVALID_LOG_START_OFFSET;
-                if (partitionResponseHeader.hasField(LOG_START_OFFSET_KEY_NAME))
-                    logStartOffset = partitionResponseHeader.getLong(LOG_START_OFFSET_KEY_NAME);
-
-                Records records = partitionResponse.getRecords(RECORD_SET_KEY_NAME);
-
-                List abortedTransactions = null;
-                if (partitionResponseHeader.hasField(ABORTED_TRANSACTIONS_KEY_NAME)) {
-                    Object[] abortedTransactionsArray = partitionResponseHeader.getArray(ABORTED_TRANSACTIONS_KEY_NAME);
-                    if (abortedTransactionsArray != null) {
-                        abortedTransactions = new ArrayList<>(abortedTransactionsArray.length);
-                        for (Object abortedTransactionObj : abortedTransactionsArray) {
-                            Struct abortedTransactionStruct = (Struct) abortedTransactionObj;
-                            long producerId = abortedTransactionStruct.getLong(PRODUCER_ID_KEY_NAME);
-                            long firstOffset = abortedTransactionStruct.getLong(FIRST_OFFSET_KEY_NAME);
-                            abortedTransactions.add(new AbortedTransaction(producerId, firstOffset));
-                        }
-                    }
-                }
-
-                PartitionData partitionData = new PartitionData(error, highWatermark, lastStableOffset, logStartOffset,
-                        abortedTransactions, records);
-                responseData.put(new TopicPartition(topic, partition), partitionData);
-            }
-        }
-        this.responseData = responseData;
-        this.throttleTimeMs = struct.getOrElse(THROTTLE_TIME_MS, DEFAULT_THROTTLE_TIME);
+    public FetchResponse(FetchResponseData fetchResponseData) {
+        super(ApiKeys.FETCH);
+        this.data = fetchResponseData;
+        this.responseDataMap = toResponseDataMap(fetchResponseData);
     }
 
-    @Override
-    public Struct toStruct(short version) {
-        return toStruct(version, responseData, throttleTimeMs);
+    public Errors error() {
+        return Errors.forCode(data.errorCode());
     }
 
-    @Override
-    protected Send toSend(String dest, ResponseHeader responseHeader, short apiVersion) {
-        Struct responseHeaderStruct = responseHeader.toStruct();
-        Struct responseBodyStruct = toStruct(apiVersion);
-
-        // write the total size and the response header
-        ByteBuffer buffer = ByteBuffer.allocate(responseHeaderStruct.sizeOf() + 4);
-        buffer.putInt(responseHeaderStruct.sizeOf() + responseBodyStruct.sizeOf());
-        responseHeaderStruct.writeTo(buffer);
-        buffer.rewind();
-
-        List sends = new ArrayList<>();
-        sends.add(new ByteBufferSend(dest, buffer));
-        addResponseData(responseBodyStruct, throttleTimeMs, dest, sends);
-        return new MultiSend(dest, sends);
+    public LinkedHashMap> responseData() {
+        return responseDataMap;
     }
 
-    public LinkedHashMap responseData() {
-        return responseData;
+    @Override
+    public int throttleTimeMs() {
+        return data.throttleTimeMs();
     }
 
-    public int throttleTimeMs() {
-        return this.throttleTimeMs;
+    public int sessionId() {
+        return data.sessionId();
     }
 
     @Override
     public Map errorCounts() {
         Map errorCounts = new HashMap<>();
-        for (PartitionData response : responseData.values())
-            updateErrorCounts(errorCounts, response.error);
+        updateErrorCounts(errorCounts, error());
+        responseDataMap.values().forEach(response ->
+            updateErrorCounts(errorCounts, response.error())
+        );
         return errorCounts;
     }
 
-    public static FetchResponse parse(ByteBuffer buffer, short version) {
-        return new FetchResponse(ApiKeys.FETCH.responseSchema(version).read(buffer));
+    public static FetchResponse parse(ByteBuffer buffer, short version) {
+        return new FetchResponse<>(new FetchResponseData(new ByteBufferAccessor(buffer), version));
     }
 
-    private static void addResponseData(Struct struct, int throttleTimeMs, String dest, List sends) {
-        Object[] allTopicData = struct.getArray(RESPONSES_KEY_NAME);
-
-        if (struct.hasField(THROTTLE_TIME_MS)) {
-            ByteBuffer buffer = ByteBuffer.allocate(8);
-            buffer.putInt(throttleTimeMs);
-            buffer.putInt(allTopicData.length);
-            buffer.rewind();
-            sends.add(new ByteBufferSend(dest, buffer));
-        } else {
-            ByteBuffer buffer = ByteBuffer.allocate(4);
-            buffer.putInt(allTopicData.length);
-            buffer.rewind();
-            sends.add(new ByteBufferSend(dest, buffer));
-        }
-
-        for (Object topicData : allTopicData)
-            addTopicData(dest, sends, (Struct) topicData);
+    @SuppressWarnings("unchecked")
+    private static  LinkedHashMap> toResponseDataMap(
+            FetchResponseData message) {
+        LinkedHashMap> responseMap = new LinkedHashMap<>();
+        message.responses().forEach(topicResponse -> {
+            topicResponse.partitionResponses().forEach(partitionResponse -> {
+                TopicPartition tp = new TopicPartition(topicResponse.topic(), partitionResponse.partition());
+                PartitionData partitionData = new PartitionData<>(partitionResponse);
+                responseMap.put(tp, partitionData);
+            });
+        });
+        return responseMap;
     }
 
-    private static void addTopicData(String dest, List sends, Struct topicData) {
-        String topic = topicData.get(TOPIC_NAME);
-        Object[] allPartitionData = topicData.getArray(PARTITIONS_KEY_NAME);
-
-        // include the topic header and the count for the number of partitions
-        ByteBuffer buffer = ByteBuffer.allocate(STRING.sizeOf(topic) + 4);
-        STRING.write(buffer, topic);
-        buffer.putInt(allPartitionData.length);
-        buffer.rewind();
-        sends.add(new ByteBufferSend(dest, buffer));
-
-        for (Object partitionData : allPartitionData)
-            addPartitionData(dest, sends, (Struct) partitionData);
+    private static  FetchResponseData toMessage(int throttleTimeMs, Errors error,
+                                                                       Iterator>> partIterator,
+                                                                       int sessionId) {
+        FetchResponseData message = new FetchResponseData();
+        message.setThrottleTimeMs(throttleTimeMs);
+        message.setErrorCode(error.code());
+        message.setSessionId(sessionId);
+
+        List topicResponseList = new ArrayList<>();
+        List>> topicsData =
+                FetchRequest.TopicAndPartitionData.batchByTopic(partIterator);
+        topicsData.forEach(partitionDataTopicAndPartitionData -> {
+            List partitionResponses = new ArrayList<>();
+            partitionDataTopicAndPartitionData.partitions.forEach((partitionId, partitionData) -> {
+                // Since PartitionData alone doesn't know the partition ID, we set it here
+                partitionData.partitionResponse.setPartition(partitionId);
+                partitionResponses.add(partitionData.partitionResponse);
+            });
+            topicResponseList.add(new FetchResponseData.FetchableTopicResponse()
+                .setTopic(partitionDataTopicAndPartitionData.topic)
+                .setPartitionResponses(partitionResponses));
+        });
+
+        message.setResponses(topicResponseList);
+        return message;
     }
 
-    private static void addPartitionData(String dest, List sends, Struct partitionData) {
-        Struct header = partitionData.getStruct(PARTITION_HEADER_KEY_NAME);
-        Records records = partitionData.getRecords(RECORD_SET_KEY_NAME);
-
-        // include the partition header and the size of the record set
-        ByteBuffer buffer = ByteBuffer.allocate(header.sizeOf() + 4);
-        header.writeTo(buffer);
-        buffer.putInt(records.sizeInBytes());
-        buffer.rewind();
-        sends.add(new ByteBufferSend(dest, buffer));
-
-        // finally the send for the record set itself
-        sends.add(new RecordsSend(dest, records));
-    }
-
-    private static Struct toStruct(short version, LinkedHashMap responseData, int throttleTimeMs) {
-        Struct struct = new Struct(ApiKeys.FETCH.responseSchema(version));
-        List> topicsData = FetchRequest.TopicAndPartitionData.batchByTopic(responseData);
-        List topicArray = new ArrayList<>();
-        for (FetchRequest.TopicAndPartitionData topicEntry: topicsData) {
-            Struct topicData = struct.instance(RESPONSES_KEY_NAME);
-            topicData.set(TOPIC_NAME, topicEntry.topic);
-            List partitionArray = new ArrayList<>();
-            for (Map.Entry partitionEntry : topicEntry.partitions.entrySet()) {
-                PartitionData fetchPartitionData = partitionEntry.getValue();
-                short errorCode = fetchPartitionData.error.code();
-                // If consumer sends FetchRequest V5 or earlier, the client library is not guaranteed to recognize the error code
-                // for KafkaStorageException. In this case the client library will translate KafkaStorageException to
-                // UnknownServerException which is not retriable. We can ensure that consumer will update metadata and retry
-                // by converting the KafkaStorageException to NotLeaderForPartitionException in the response if FetchRequest version <= 5
-                if (errorCode == Errors.KAFKA_STORAGE_ERROR.code() && version <= 5)
-                    errorCode = Errors.NOT_LEADER_FOR_PARTITION.code();
-                Struct partitionData = topicData.instance(PARTITIONS_KEY_NAME);
-                Struct partitionDataHeader = partitionData.instance(PARTITION_HEADER_KEY_NAME);
-                partitionDataHeader.set(PARTITION_ID, partitionEntry.getKey());
-                partitionDataHeader.set(ERROR_CODE, errorCode);
-                partitionDataHeader.set(HIGH_WATERMARK_KEY_NAME, fetchPartitionData.highWatermark);
-
-                if (partitionDataHeader.hasField(LAST_STABLE_OFFSET_KEY_NAME)) {
-                    partitionDataHeader.set(LAST_STABLE_OFFSET_KEY_NAME, fetchPartitionData.lastStableOffset);
-
-                    if (fetchPartitionData.abortedTransactions == null) {
-                        partitionDataHeader.set(ABORTED_TRANSACTIONS_KEY_NAME, null);
-                    } else {
-                        List abortedTransactionStructs = new ArrayList<>(fetchPartitionData.abortedTransactions.size());
-                        for (AbortedTransaction abortedTransaction : fetchPartitionData.abortedTransactions) {
-                            Struct abortedTransactionStruct = partitionDataHeader.instance(ABORTED_TRANSACTIONS_KEY_NAME);
-                            abortedTransactionStruct.set(PRODUCER_ID_KEY_NAME, abortedTransaction.producerId);
-                            abortedTransactionStruct.set(FIRST_OFFSET_KEY_NAME, abortedTransaction.firstOffset);
-                            abortedTransactionStructs.add(abortedTransactionStruct);
-                        }
-                        partitionDataHeader.set(ABORTED_TRANSACTIONS_KEY_NAME, abortedTransactionStructs.toArray());
-                    }
-                }
-                if (partitionDataHeader.hasField(LOG_START_OFFSET_KEY_NAME))
-                    partitionDataHeader.set(LOG_START_OFFSET_KEY_NAME, fetchPartitionData.logStartOffset);
-
-                partitionData.set(PARTITION_HEADER_KEY_NAME, partitionDataHeader);
-                partitionData.set(RECORD_SET_KEY_NAME, fetchPartitionData.records);
-                partitionArray.add(partitionData);
-            }
-            topicData.set(PARTITIONS_KEY_NAME, partitionArray.toArray());
-            topicArray.add(topicData);
-        }
-        struct.set(RESPONSES_KEY_NAME, topicArray.toArray());
-        struct.setIfExists(THROTTLE_TIME_MS, throttleTimeMs);
-
-        return struct;
+    /**
+     * Convenience method to find the size of a response.
+     *
+     * @param version       The version of the response to use.
+     * @param partIterator  The partition iterator.
+     * @return              The response size in bytes.
+     */
+    public static  int sizeOf(short version,
+                                                     Iterator>> partIterator) {
+        // Since the throttleTimeMs and metadata field sizes are constant and fixed, we can
+        // use arbitrary values here without affecting the result.
+        FetchResponseData data = toMessage(0, Errors.NONE, partIterator, INVALID_SESSION_ID);
+        ObjectSerializationCache cache = new ObjectSerializationCache();
+        return 4 + data.size(cache, version);
     }
 
-    public static int sizeOf(short version, LinkedHashMap responseData) {
-        return 4 + toStruct(version, responseData, 0).sizeOf();
+    @Override
+    public boolean shouldClientThrottle(short version) {
+        return version >= 8;
     }
-
 }
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 9bfc968d55faa..dc512a554e246 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,132 +18,68 @@
 
 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.ByteBufferAccessor;
 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)"));
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {FIND_COORDINATOR_REQUEST_V0, FIND_COORDINATOR_REQUEST_V1};
-    }
 
     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();
         }
-    }
-
-    private final String coordinatorKey;
-    private final CoordinatorType coordinatorType;
 
-    private FindCoordinatorRequest(CoordinatorType coordinatorType, String coordinatorKey, short version) {
-        super(version);
-        this.coordinatorType = coordinatorType;
-        this.coordinatorKey = coordinatorKey;
+        public FindCoordinatorRequestData data() {
+            return data;
+        }
     }
 
-    public FindCoordinatorRequest(Struct struct, short version) {
-        super(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);
+    private final FindCoordinatorRequestData data;
+
+    private FindCoordinatorRequest(FindCoordinatorRequestData data, short version) {
+        super(ApiKeys.FIND_COORDINATOR, version);
+        this.data = data;
     }
 
     @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:
-                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) {
-        return new FindCoordinatorRequest(ApiKeys.FIND_COORDINATOR.parseRequest(version, buffer), version);
+        return new FindCoordinatorRequest(new FindCoordinatorRequestData(new ByteBufferAccessor(buffer), version),
+            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;
+    public FindCoordinatorRequestData data() {
+        return data;
     }
 
     public enum CoordinatorType {
@@ -155,6 +91,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 bddc41f770ae4..11f3d48887900 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,129 +17,82 @@
 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.ByteBufferAccessor;
 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.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"));
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {FIND_COORDINATOR_RESPONSE_V0, FIND_COORDINATOR_RESPONSE_V1};
-    }
 
     /**
      * Possible error codes:
      *
+     * COORDINATOR_LOAD_IN_PROGRESS (14)
      * COORDINATOR_NOT_AVAILABLE (15)
-     * NOT_COORDINATOR (16)
      * GROUP_AUTHORIZATION_FAILED (30)
+     * INVALID_REQUEST (42)
+     * 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) {
+        super(ApiKeys.FIND_COORDINATOR);
+        this.data = data;
     }
 
-    public FindCoordinatorResponse(int throttleTimeMs, Errors error, Node node) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.error = error;
-        this.node = node;
-        this.errorMessage = null;
+    @Override
+    public FindCoordinatorResponseData data() {
+        return data;
     }
 
-    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);
-
-        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 error() != Errors.NONE;
     }
 
     public Errors error() {
-        return error;
+        return Errors.forCode(data.errorCode());
     }
 
     @Override
     public Map errorCounts() {
-        return errorCounts(error);
+        return errorCounts(error());
     }
 
-    public Node node() {
-        return node;
+    public static FindCoordinatorResponse parse(ByteBuffer buffer, short version) {
+        return new FindCoordinatorResponse(new FindCoordinatorResponseData(new ByteBufferAccessor(buffer), version));
     }
 
     @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;
+    public String toString() {
+        return data.toString();
     }
 
-    public static FindCoordinatorResponse parse(ByteBuffer buffer, short version) {
-        return new FindCoordinatorResponse(ApiKeys.FIND_COORDINATOR.responseSchema(version).read(buffer));
+    @Override
+    public boolean shouldClientThrottle(short version) {
+        return version >= 2;
     }
 
-    @Override
-    public String toString() {
-        return "FindCoordinatorResponse(" +
-                "throttleTimeMs=" + throttleTimeMs +
-                ", errorMessage='" + errorMessage + '\'' +
-                ", error=" + error +
-                ", node=" + node +
-                ')';
+    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 7d84918d3f4c8..482e61a255a8e 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,113 +16,63 @@
  */
 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.ByteBufferAccessor;
 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;
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {HEARTBEAT_REQUEST_V0, HEARTBEAT_REQUEST_V1};
-    }
 
     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;
-
-    private HeartbeatRequest(String groupId, int groupGenerationId, String memberId, short version) {
-        super(version);
-        this.groupId = groupId;
-        this.groupGenerationId = groupGenerationId;
-        this.memberId = memberId;
-    }
+    private final HeartbeatRequestData data;
 
-    public HeartbeatRequest(Struct struct, short version) {
-        super(version);
-        groupId = struct.get(GROUP_ID);
-        groupGenerationId = struct.get(GENERATION_ID);
-        memberId = struct.get(MEMBER_ID);
+    private HeartbeatRequest(HeartbeatRequestData data, short version) {
+        super(ApiKeys.HEARTBEAT, version);
+        this.data = data;
     }
 
     @Override
     public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
-        short versionId = version();
-        switch (versionId) {
-            case 0:
-                return new HeartbeatResponse(Errors.forException(e));
-            case 1:
-                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) {
-        return new HeartbeatRequest(ApiKeys.HEARTBEAT.parseRequest(version, buffer), version);
+        return new HeartbeatRequest(new HeartbeatRequestData(new ByteBufferAccessor(buffer), version), 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;
+    public HeartbeatRequestData data() {
+        return data;
     }
 }
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 b52a9936872c9..eb402fcbab9f7 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,29 +16,16 @@
  */
 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.ByteBufferAccessor;
 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;
-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);
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {HEARTBEAT_RESPONSE_V0, HEARTBEAT_RESPONSE_V1};
-    }
-
     /**
      * Possible error codes:
      *
@@ -49,45 +36,38 @@ 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(Struct struct) {
-        this.throttleTimeMs = struct.getOrElse(THROTTLE_TIME_MS, DEFAULT_THROTTLE_TIME);
-        error = Errors.forCode(struct.get(ERROR_CODE));
+    public HeartbeatResponse(HeartbeatResponseData data) {
+        super(ApiKeys.HEARTBEAT);
+        this.data = data;
     }
 
+    @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 errorCounts(error());
     }
 
     @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;
+    public HeartbeatResponseData data() {
+        return data;
     }
 
     public static HeartbeatResponse parse(ByteBuffer buffer, short version) {
-        return new HeartbeatResponse(ApiKeys.HEARTBEAT.parseResponse(version, buffer));
+        return new HeartbeatResponse(new HeartbeatResponseData(new ByteBufferAccessor(buffer), version));
+    }
+
+    @Override
+    public boolean shouldClientThrottle(short version) {
+        return version >= 2;
     }
 }
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..2bc591410f9a9
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/IncrementalAlterConfigsRequest.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.requests;
+
+import org.apache.kafka.clients.admin.AlterConfigOp;
+import org.apache.kafka.common.config.ConfigResource;
+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.ByteBufferAccessor;
+
+import java.nio.ByteBuffer;
+import java.util.Collection;
+import java.util.Map;
+
+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;
+        }
+
+        public Builder(final Collection resources,
+                       final Map> configs,
+                       final boolean validateOnly) {
+            super(ApiKeys.INCREMENTAL_ALTER_CONFIGS);
+            this.data = new IncrementalAlterConfigsRequestData()
+                            .setValidateOnly(validateOnly);
+            for (ConfigResource resource : resources) {
+                IncrementalAlterConfigsRequestData.AlterableConfigCollection alterableConfigSet =
+                    new IncrementalAlterConfigsRequestData.AlterableConfigCollection();
+                for (AlterConfigOp configEntry : configs.get(resource))
+                    alterableConfigSet.add(new IncrementalAlterConfigsRequestData.AlterableConfig()
+                                               .setName(configEntry.configEntry().name())
+                                               .setValue(configEntry.configEntry().value())
+                                               .setConfigOperation(configEntry.opType().id()));
+                IncrementalAlterConfigsRequestData.AlterConfigsResource alterConfigsResource = new IncrementalAlterConfigsRequestData.AlterConfigsResource();
+                alterConfigsResource.setResourceType(resource.type().id())
+                    .setResourceName(resource.name()).setConfigs(alterableConfigSet);
+                data.resources().add(alterConfigsResource);
+            }
+        }
+
+        public Builder(final Map> configs,
+                       final boolean validateOnly) {
+            this(configs.keySet(), configs, validateOnly);
+        }
+
+        @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;
+    }
+
+    public static IncrementalAlterConfigsRequest parse(ByteBuffer buffer, short version) {
+        return new IncrementalAlterConfigsRequest(new IncrementalAlterConfigsRequestData(
+            new ByteBufferAccessor(buffer), version), version);
+    }
+
+    @Override
+    public IncrementalAlterConfigsRequestData data() {
+        return data;
+    }
+
+    @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..b5887de9b4b75
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/IncrementalAlterConfigsResponse.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.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.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.Errors;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+public class IncrementalAlterConfigsResponse extends AbstractResponse {
+
+    public IncrementalAlterConfigsResponse(final int requestThrottleMs,
+                                           final Map results) {
+        super(ApiKeys.INCREMENTAL_ALTER_CONFIGS);
+        final List newResults = new ArrayList<>(results.size());
+        results.forEach(
+            (resource, error) -> newResults.add(
+                new AlterConfigsResourceResponse()
+                    .setErrorCode(error.error().code())
+                    .setErrorMessage(error.message())
+                    .setResourceName(resource.name())
+                    .setResourceType(resource.type().id()))
+        );
+
+        this.data = new IncrementalAlterConfigsResponseData()
+                        .setResponses(newResults)
+                        .setThrottleTimeMs(requestThrottleMs);
+    }
+
+    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) {
+        super(ApiKeys.INCREMENTAL_ALTER_CONFIGS);
+        this.data = data;
+    }
+
+    @Override
+    public IncrementalAlterConfigsResponseData data() {
+        return data;
+    }
+
+    @Override
+    public Map errorCounts() {
+        HashMap counts = new HashMap<>();
+        data.responses().forEach(response ->
+            updateErrorCounts(counts, Errors.forCode(response.errorCode()))
+        );
+        return counts;
+    }
+
+    @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(new IncrementalAlterConfigsResponseData(
+            new ByteBufferAccessor(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 6c659ff207aa8..5c24b41b351df 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,101 +16,65 @@
  */
 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.ByteBufferAccessor;
 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."));
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{INIT_PRODUCER_ID_REQUEST_V0};
-    }
-
-    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) {
-        super(version);
-        this.transactionalId = struct.get(NULLABLE_TRANSACTIONAL_ID);
-        this.transactionTimeoutMs = struct.getInt(TRANSACTION_TIMEOUT_KEY_NAME);
-    }
+    private final InitProducerIdRequestData data;
 
-    private InitProducerIdRequest(short version, String transactionalId, int transactionTimeoutMs) {
-        super(version);
-        this.transactionalId = transactionalId;
-        this.transactionTimeoutMs = transactionTimeoutMs;
+    private InitProducerIdRequest(InitProducerIdRequestData data, short version) {
+        super(ApiKeys.INIT_PRODUCER_ID, version);
+        this.data = data;
     }
 
     @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;
+        return new InitProducerIdRequest(new InitProducerIdRequestData(new ByteBufferAccessor(buffer), version), version);
     }
 
     @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;
+    public InitProducerIdRequestData data() {
+        return data;
     }
 
 }
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 7a988ca806a36..f8451d7863b3f 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,103 +16,63 @@
  */
 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.ByteBufferAccessor;
 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}
+ *   - {@link Errors#INVALID_PRODUCER_EPOCH} // for version <=3
+ *   - {@link Errors#PRODUCER_FENCED}
+ */
 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);
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{INIT_PRODUCER_ID_RESPONSE_V0};
-    }
+    private final InitProducerIdResponseData data;
 
-    private final int throttleTimeMs;
-    private final Errors error;
-    private final long producerId;
-    private final short epoch;
-
-    public InitProducerIdResponse(int throttleTimeMs, Errors error, long producerId, short epoch) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.error = error;
-        this.producerId = producerId;
-        this.epoch = epoch;
-    }
-
-    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(InitProducerIdResponseData data) {
+        super(ApiKeys.INIT_PRODUCER_ID);
+        this.data = data;
     }
 
+    @Override
     public int throttleTimeMs() {
-        return throttleTimeMs;
+        return data.throttleTimeMs();
     }
 
-    public long producerId() {
-        return producerId;
-    }
-
-    public Errors error() {
-        return error;
+    @Override
+    public Map errorCounts() {
+        return errorCounts(Errors.forCode(data.errorCode()));
     }
 
     @Override
-    public Map errorCounts() {
-        return errorCounts(error);
+    public InitProducerIdResponseData data() {
+        return data;
     }
 
-    public short epoch() {
-        return epoch;
+    public static InitProducerIdResponse parse(ByteBuffer buffer, short version) {
+        return new InitProducerIdResponse(new InitProducerIdResponseData(new ByteBufferAccessor(buffer), version));
     }
 
     @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;
+    public String toString() {
+        return data.toString();
     }
 
-    public static InitProducerIdResponse parse(ByteBuffer buffer, short version) {
-        return new InitProducerIdResponse(ApiKeys.INIT_PRODUCER_ID.parseResponse(version, buffer));
+    public Errors error() {
+        return Errors.forCode(data.errorCode());
     }
 
     @Override
-    public String toString() {
-        return "InitProducerIdResponse(" +
-                "error=" + error +
-                ", producerId=" + producerId +
-                ", producerEpoch=" + epoch +
-                ", throttleTimeMs=" + throttleTimeMs +
-                ')';
+    public boolean shouldClientThrottle(short version) {
+        return version >= 1;
     }
 }
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 a7b62a98660f2..35155a085b6f7 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,254 +16,122 @@
  */
 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.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.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;
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {JOIN_GROUP_REQUEST_V0, JOIN_GROUP_REQUEST_V1, JOIN_GROUP_REQUEST_V2};
-    }
-
-    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(version);
-        this.groupId = groupId;
-        this.sessionTimeout = sessionTimeout;
-        this.rebalanceTimeout = rebalanceTimeout;
-        this.memberId = memberId;
-        this.protocolType = protocolType;
-        this.groupProtocols = groupProtocols;
-    }
-
-    public JoinGroupRequest(Struct struct, short versionId) {
-        super(versionId);
+    private final JoinGroupRequestData data;
 
-        groupId = struct.get(GROUP_ID);
-        sessionTimeout = struct.getInt(SESSION_TIMEOUT_KEY_NAME);
-
-        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;
-
-        memberId = struct.get(MEMBER_ID);
-        protocolType = struct.getString(PROTOCOL_TYPE_KEY_NAME);
-
-        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 static final String UNKNOWN_MEMBER_ID = "";
+    public static final int UNKNOWN_GENERATION_ID = -1;
+    public static final String UNKNOWN_PROTOCOL_NAME = "";
+
+    private static final int MAX_GROUP_INSTANCE_ID_LENGTH = 249;
+
+    /**
+     * 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 '-'");
     }
 
-    @Override
-    public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
-        short versionId = version();
-        switch (versionId) {
-            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());
-            case 2:
-                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());
-
-            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()));
+    /**
+     * 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;
     }
 
-    public String groupId() {
-        return groupId;
-    }
-
-    public int sessionTimeout() {
-        return sessionTimeout;
+    public JoinGroupRequest(JoinGroupRequestData data, short version) {
+        super(ApiKeys.JOIN_GROUP, version);
+        this.data = data;
+        maybeOverrideRebalanceTimeout(version);
     }
 
-    public int rebalanceTimeout() {
-        return rebalanceTimeout;
+    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 String memberId() {
-        return memberId;
+    @Override
+    public JoinGroupRequestData data() {
+        return data;
     }
 
-    public List groupProtocols() {
-        return groupProtocols;
-    }
+    @Override
+    public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
+        JoinGroupResponseData data = new JoinGroupResponseData()
+            .setThrottleTimeMs(throttleTimeMs)
+            .setErrorCode(Errors.forException(e).code())
+            .setGenerationId(UNKNOWN_GENERATION_ID)
+            .setProtocolName(UNKNOWN_PROTOCOL_NAME)
+            .setLeader(UNKNOWN_MEMBER_ID)
+            .setMemberId(UNKNOWN_MEMBER_ID)
+            .setMembers(Collections.emptyList());
+
+        if (version() >= 7)
+            data.setProtocolName(null);
+        else
+            data.setProtocolName(UNKNOWN_PROTOCOL_NAME);
 
-    public String protocolType() {
-        return protocolType;
+        return new JoinGroupResponse(data);
     }
 
     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 new JoinGroupRequest(new JoinGroupRequestData(new ByteBufferAccessor(buffer), version), 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 4bcd6e6d65fd6..336c82462a21d 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,203 +16,57 @@
  */
 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.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.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.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 final JoinGroupResponseData data;
 
-    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)));
-
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {JOIN_GROUP_RESPONSE_V0, JOIN_GROUP_RESPONSE_V1, JOIN_GROUP_RESPONSE_V2};
+    public JoinGroupResponse(JoinGroupResponseData data) {
+        super(ApiKeys.JOIN_GROUP);
+        this.data = 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)
-     */
-
-    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;
+    @Override
+    public JoinGroupResponseData data() {
+        return 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);
+    public boolean isLeader() {
+        return data.memberId().equals(data.leader());
     }
 
+    @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 int generationId() {
-        return generationId;
-    }
-
-    public String groupProtocol() {
-        return groupProtocol;
-    }
-
-    public String memberId() {
-        return memberId;
-    }
-
-    public String leaderId() {
-        return leaderId;
-    }
-
-    public boolean isLeader() {
-        return memberId.equals(leaderId);
-    }
-
-    public Map members() {
-        return members;
+        return errorCounts(Errors.forCode(data.errorCode()));
     }
 
     public static JoinGroupResponse parse(ByteBuffer buffer, short version) {
-        return new JoinGroupResponse(ApiKeys.JOIN_GROUP.parseResponse(version, buffer));
+        return new JoinGroupResponse(new JoinGroupResponseData(new ByteBufferAccessor(buffer), version));
     }
 
     @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;
+    public String toString() {
+        return data.toString();
     }
 
     @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(), ",")) + ")";
+    public boolean shouldClientThrottle(short version) {
+        return version >= 3;
     }
 }
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 4fa5337188489..833e0255336d2 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,112 +17,75 @@
 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.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.Struct;
+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 java.util.stream.Collectors;
 
-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.BOOLEAN;
-import static org.apache.kafka.common.protocol.types.Type.INT32;
-import static org.apache.kafka.common.protocol.types.Type.STRING;
+public class LeaderAndIsrRequest extends AbstractControlRequest {
 
-public class LeaderAndIsrRequest extends AbstractRequest {
-    private static final String CONTROLLER_ID_KEY_NAME = "controller_id";
-    private static final String CONTROLLER_EPOCH_KEY_NAME = "controller_epoch";
-    private static final String PARTITION_STATES_KEY_NAME = "partition_states";
-    private static final String LIVE_LEADERS_KEY_NAME = "live_leaders";
+    public static class Builder extends AbstractControlRequest.Builder {
 
-    // partition_states key names
-    private static final String LEADER_KEY_NAME = "leader";
-    private static final String LEADER_EPOCH_KEY_NAME = "leader_epoch";
-    private static final String ISR_KEY_NAME = "isr";
-    private static final String ZK_VERSION_KEY_NAME = "zk_version";
-    private static final String REPLICAS_KEY_NAME = "replicas";
-    private static final String IS_NEW_KEY_NAME = "is_new";
+        private final List partitionStates;
+        private final Collection liveLeaders;
 
-    // live_leaders key names
-    private static final String END_POINT_ID_KEY_NAME = "id";
-    private static final String HOST_KEY_NAME = "host";
-    private static final String PORT_KEY_NAME = "port";
-
-    private static final Schema LEADER_AND_ISR_REQUEST_PARTITION_STATE_V0 = new Schema(
-            TOPIC_NAME,
-            PARTITION_ID,
-            new Field(CONTROLLER_EPOCH_KEY_NAME, INT32, "The controller epoch."),
-            new Field(LEADER_KEY_NAME, INT32, "The broker id for the leader."),
-            new Field(LEADER_EPOCH_KEY_NAME, INT32, "The leader epoch."),
-            new Field(ISR_KEY_NAME, new ArrayOf(INT32), "The in sync replica ids."),
-            new Field(ZK_VERSION_KEY_NAME, INT32, "The ZK version."),
-            new Field(REPLICAS_KEY_NAME, new ArrayOf(INT32), "The replica ids."));
-
-    // LEADER_AND_ISR_REQUEST_PARTITION_STATE_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_PARTITION_STATE_V1 = new Schema(
-            TOPIC_NAME,
-            PARTITION_ID,
-            new Field(CONTROLLER_EPOCH_KEY_NAME, INT32, "The controller epoch."),
-            new Field(LEADER_KEY_NAME, INT32, "The broker id for the leader."),
-            new Field(LEADER_EPOCH_KEY_NAME, INT32, "The leader epoch."),
-            new Field(ISR_KEY_NAME, new ArrayOf(INT32), "The in sync replica ids."),
-            new Field(ZK_VERSION_KEY_NAME, INT32, "The ZK version."),
-            new Field(REPLICAS_KEY_NAME, new ArrayOf(INT32), "The replica ids."),
-            new Field(IS_NEW_KEY_NAME, BOOLEAN, "Whether the replica should have existed on the broker or not"));
-
-    private static final Schema LEADER_AND_ISR_REQUEST_LIVE_LEADER_V0 = new Schema(
-            new Field(END_POINT_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 LEADER_AND_ISR_REQUEST_V0 = new Schema(
-            new Field(CONTROLLER_ID_KEY_NAME, INT32, "The controller id."),
-            new Field(CONTROLLER_EPOCH_KEY_NAME, INT32, "The controller epoch."),
-            new Field(PARTITION_STATES_KEY_NAME, new ArrayOf(LEADER_AND_ISR_REQUEST_PARTITION_STATE_V0)),
-            new Field(LIVE_LEADERS_KEY_NAME, new ArrayOf(LEADER_AND_ISR_REQUEST_LIVE_LEADER_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(
-            new Field(CONTROLLER_ID_KEY_NAME, INT32, "The controller id."),
-            new Field(CONTROLLER_EPOCH_KEY_NAME, INT32, "The controller epoch."),
-            new Field(PARTITION_STATES_KEY_NAME, new ArrayOf(LEADER_AND_ISR_REQUEST_PARTITION_STATE_V1)),
-            new Field(LIVE_LEADERS_KEY_NAME, new ArrayOf(LEADER_AND_ISR_REQUEST_LIVE_LEADER_V0)));
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{LEADER_AND_ISR_REQUEST_V0, LEADER_AND_ISR_REQUEST_V1};
-    }
-
-    public static class Builder extends AbstractRequest.Builder {
-        private final int controllerId;
-        private final int controllerEpoch;
-        private final Map partitionStates;
-        private final Set liveLeaders;
-
-        public Builder(short version, int controllerId, int controllerEpoch,
-                       Map partitionStates, Set liveLeaders) {
-            super(ApiKeys.LEADER_AND_ISR, version);
-            this.controllerId = controllerId;
-            this.controllerEpoch = controllerEpoch;
+        public Builder(short version, int controllerId, int controllerEpoch, long brokerEpoch,
+                       List partitionStates, Collection liveLeaders) {
+            super(ApiKeys.LEADER_AND_ISR, version, controllerId, controllerEpoch, brokerEpoch);
             this.partitionStates = partitionStates;
             this.liveLeaders = liveLeaders;
         }
 
         @Override
         public LeaderAndIsrRequest build(short version) {
-            return new LeaderAndIsrRequest(controllerId, controllerEpoch, 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)
+                .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
@@ -131,174 +94,84 @@ public String toString() {
             bld.append("(type=LeaderAndIsRequest")
                 .append(", controllerId=").append(controllerId)
                 .append(", controllerEpoch=").append(controllerEpoch)
+                .append(", brokerEpoch=").append(brokerEpoch)
                 .append(", partitionStates=").append(partitionStates)
                 .append(", liveLeaders=(").append(Utils.join(liveLeaders, ", ")).append(")")
                 .append(")");
             return bld.toString();
-        }
-    }
-
-    private final int controllerId;
-    private final int controllerEpoch;
-    private final Map partitionStates;
-    private final Set liveLeaders;
-
-    private LeaderAndIsrRequest(int controllerId, int controllerEpoch, Map partitionStates,
-                                Set liveLeaders, short version) {
-        super(version);
-        this.controllerId = controllerId;
-        this.controllerEpoch = controllerEpoch;
-        this.partitionStates = partitionStates;
-        this.liveLeaders = liveLeaders;
-    }
-
-    public LeaderAndIsrRequest(Struct struct, short version) {
-        super(version);
 
-        Map partitionStates = new HashMap<>();
-        for (Object partitionStateDataObj : struct.getArray(PARTITION_STATES_KEY_NAME)) {
-            Struct partitionStateData = (Struct) partitionStateDataObj;
-            String topic = partitionStateData.get(TOPIC_NAME);
-            int partition = partitionStateData.get(PARTITION_ID);
-            int controllerEpoch = partitionStateData.getInt(CONTROLLER_EPOCH_KEY_NAME);
-            int leader = partitionStateData.getInt(LEADER_KEY_NAME);
-            int leaderEpoch = partitionStateData.getInt(LEADER_EPOCH_KEY_NAME);
-
-            Object[] isrArray = partitionStateData.getArray(ISR_KEY_NAME);
-            List isr = new ArrayList<>(isrArray.length);
-            for (Object r : isrArray)
-                isr.add((Integer) r);
-
-            int zkVersion = partitionStateData.getInt(ZK_VERSION_KEY_NAME);
-
-            Object[] replicasArray = partitionStateData.getArray(REPLICAS_KEY_NAME);
-            List replicas = new ArrayList<>(replicasArray.length);
-            for (Object r : replicasArray)
-                replicas.add((Integer) r);
-            boolean isNew = partitionStateData.hasField(IS_NEW_KEY_NAME) ? partitionStateData.getBoolean(IS_NEW_KEY_NAME) : false;
-
-            PartitionState partitionState = new PartitionState(controllerEpoch, leader, leaderEpoch, isr, zkVersion, replicas, isNew);
-            partitionStates.put(new TopicPartition(topic, partition), partitionState);
         }
-
-        Set leaders = new HashSet<>();
-        for (Object leadersDataObj : struct.getArray(LIVE_LEADERS_KEY_NAME)) {
-            Struct leadersData = (Struct) leadersDataObj;
-            int id = leadersData.getInt(END_POINT_ID_KEY_NAME);
-            String host = leadersData.getString(HOST_KEY_NAME);
-            int port = leadersData.getInt(PORT_KEY_NAME);
-            leaders.add(new Node(id, host, port));
-        }
-
-        controllerId = struct.getInt(CONTROLLER_ID_KEY_NAME);
-        controllerEpoch = struct.getInt(CONTROLLER_EPOCH_KEY_NAME);
-        this.partitionStates = partitionStates;
-        this.liveLeaders = leaders;
     }
 
-    @Override
-    protected Struct toStruct() {
-        short version = version();
-        Struct struct = new Struct(ApiKeys.LEADER_AND_ISR.requestSchema(version));
-        struct.set(CONTROLLER_ID_KEY_NAME, controllerId);
-        struct.set(CONTROLLER_EPOCH_KEY_NAME, controllerEpoch);
+    private final LeaderAndIsrRequestData data;
 
-        List partitionStatesData = new ArrayList<>(partitionStates.size());
-        for (Map.Entry entry : partitionStates.entrySet()) {
-            Struct partitionStateData = struct.instance(PARTITION_STATES_KEY_NAME);
-            TopicPartition topicPartition = entry.getKey();
-            partitionStateData.set(TOPIC_NAME, topicPartition.topic());
-            partitionStateData.set(PARTITION_ID, topicPartition.partition());
-            PartitionState partitionState = entry.getValue();
-            partitionStateData.set(CONTROLLER_EPOCH_KEY_NAME, partitionState.basePartitionState.controllerEpoch);
-            partitionStateData.set(LEADER_KEY_NAME, partitionState.basePartitionState.leader);
-            partitionStateData.set(LEADER_EPOCH_KEY_NAME, partitionState.basePartitionState.leaderEpoch);
-            partitionStateData.set(ISR_KEY_NAME, partitionState.basePartitionState.isr.toArray());
-            partitionStateData.set(ZK_VERSION_KEY_NAME, partitionState.basePartitionState.zkVersion);
-            partitionStateData.set(REPLICAS_KEY_NAME, partitionState.basePartitionState.replicas.toArray());
-            if (partitionStateData.hasField(IS_NEW_KEY_NAME))
-                partitionStateData.set(IS_NEW_KEY_NAME, partitionState.isNew);
-            partitionStatesData.add(partitionStateData);
-        }
-        struct.set(PARTITION_STATES_KEY_NAME, partitionStatesData.toArray());
+    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();
+    }
 
-        List leadersData = new ArrayList<>(liveLeaders.size());
-        for (Node leader : liveLeaders) {
-            Struct leaderData = struct.instance(LIVE_LEADERS_KEY_NAME);
-            leaderData.set(END_POINT_ID_KEY_NAME, leader.id());
-            leaderData.set(HOST_KEY_NAME, leader.host());
-            leaderData.set(PORT_KEY_NAME, leader.port());
-            leadersData.add(leaderData);
+    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());
+                }
+            }
         }
-        struct.set(LIVE_LEADERS_KEY_NAME, leadersData.toArray());
-        return struct;
     }
 
     @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:
-                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;
+    public Iterable partitionStates() {
+        if (version() >= 2)
+            return () -> new FlattenedIterator<>(data.topicStates().iterator(),
+                topicState -> topicState.partitionStates().iterator());
+        return data.ungroupedPartitionStates();
     }
 
-    public static LeaderAndIsrRequest parse(ByteBuffer buffer, short version) {
-        return new LeaderAndIsrRequest(ApiKeys.LEADER_AND_ISR.parseRequest(version, buffer), version);
+    public List liveLeaders() {
+        return Collections.unmodifiableList(data.liveLeaders());
     }
 
-    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;
-        }
-
-        @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 + ")";
-        }
+    @Override
+    public LeaderAndIsrRequestData data() {
+        return data;
     }
 
+    public static LeaderAndIsrRequest parse(ByteBuffer buffer, short version) {
+        return new LeaderAndIsrRequest(new LeaderAndIsrRequestData(new ByteBufferAccessor(buffer), version), 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 c21f9a783b3b8..974dde84ec4c7 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,107 +16,69 @@
  */
 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.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.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;
-
 public class LeaderAndIsrResponse extends AbstractResponse {
-    private static final String PARTITIONS_KEY_NAME = "partitions";
-
-    private static final Schema LEADER_AND_ISR_RESPONSE_PARTITION_V0 = new Schema(
-            TOPIC_NAME,
-            PARTITION_ID,
-            ERROR_CODE);
-    private static final Schema LEADER_AND_ISR_RESPONSE_V0 = new Schema(
-            ERROR_CODE,
-            new Field(PARTITIONS_KEY_NAME, new ArrayOf(LEADER_AND_ISR_RESPONSE_PARTITION_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;
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{LEADER_AND_ISR_RESPONSE_V0, LEADER_AND_ISR_RESPONSE_V1};
-    }
 
     /**
      * Possible error code:
      *
      * STALE_CONTROLLER_EPOCH (11)
+     * STALE_BROKER_EPOCH (77)
      */
-    private final Errors error;
-
-    private final Map responses;
-
-    public LeaderAndIsrResponse(Errors error, Map responses) {
-        this.responses = responses;
-        this.error = error;
-    }
-
-    public LeaderAndIsrResponse(Struct struct) {
-        responses = new HashMap<>();
-        for (Object responseDataObj : struct.getArray(PARTITIONS_KEY_NAME)) {
-            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);
-        }
+    private final LeaderAndIsrResponseData data;
 
-        error = Errors.forCode(struct.get(ERROR_CODE));
+    public LeaderAndIsrResponse(LeaderAndIsrResponseData data) {
+        super(ApiKeys.LEADER_AND_ISR);
+        this.data = data;
     }
 
-    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() + 1);
+        Map errors = errorCounts(data.partitionErrors().stream().map(l -> Errors.forCode(l.errorCode())));
+        // Top level error
+        updateErrorCounts(errors, Errors.NONE);
+        return errors;
+    }
+
+    @Override
+    public int throttleTimeMs() {
+        return DEFAULT_THROTTLE_TIME;
     }
 
     public static LeaderAndIsrResponse parse(ByteBuffer buffer, short version) {
-        return new LeaderAndIsrResponse(ApiKeys.LEADER_AND_ISR.parseResponse(version, buffer));
+        return new LeaderAndIsrResponse(new LeaderAndIsrResponseData(new ByteBufferAccessor(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_KEY_NAME);
-            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_KEY_NAME, responseDatas.toArray());
-        struct.set(ERROR_CODE, error.code());
+    public LeaderAndIsrResponseData data() {
+        return data;
+    }
 
-        return struct;
+    @Override
+    public String toString() {
+        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 b0d0ad6d685bb..8ce95350fc9bd 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,100 +16,101 @@
  */
 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.ByteBufferAccessor;
 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.protocol.MessageUtil;
 
 import java.nio.ByteBuffer;
-
-import static org.apache.kafka.common.protocol.CommonFields.GROUP_ID;
-import static org.apache.kafka.common.protocol.CommonFields.MEMBER_ID;
+import java.util.Collections;
+import java.util.List;
 
 public class LeaveGroupRequest extends AbstractRequest {
-    private static final Schema LEAVE_GROUP_REQUEST_V0 = new Schema(
-            GROUP_ID,
-            MEMBER_ID);
-
-    /* v1 request is the same as v0. Throttle time has been added to response */
-    private static final Schema LEAVE_GROUP_REQUEST_V1 = LEAVE_GROUP_REQUEST_V0;
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {LEAVE_GROUP_REQUEST_V0, LEAVE_GROUP_REQUEST_V1};
-    }
 
     public static class Builder extends AbstractRequest.Builder {
         private final String groupId;
-        private final String memberId;
+        private final List members;
 
-        public Builder(String groupId, String memberId) {
-            super(ApiKeys.LEAVE_GROUP);
+        public Builder(String groupId, List members) {
+            this(groupId, members, ApiKeys.LEAVE_GROUP.oldestVersion(), ApiKeys.LEAVE_GROUP.latestVersion());
+        }
+
+        Builder(String groupId, List members, short oldestVersion, short latestVersion) {
+            super(ApiKeys.LEAVE_GROUP, oldestVersion, latestVersion);
             this.groupId = groupId;
-            this.memberId = memberId;
+            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) {
-            return new LeaveGroupRequest(groupId, memberId, 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() {
-            StringBuilder bld = new StringBuilder();
-            bld.append("(type=LeaveGroupRequest").
-                append(", groupId=").append(groupId).
-                append(", memberId=").append(memberId).
-                append(")");
-            return bld.toString();
+            return "(type=LeaveGroupRequest" +
+                       ", groupId=" + groupId +
+                       ", members=" + MessageUtil.deepToString(members.iterator()) +
+                       ")";
         }
     }
+    private final LeaveGroupRequestData data;
 
-    private final String groupId;
-    private final String memberId;
+    private LeaveGroupRequest(LeaveGroupRequestData data, short version) {
+        super(ApiKeys.LEAVE_GROUP, version);
+        this.data = data;
+    }
 
-    private LeaveGroupRequest(String groupId, String memberId, short version) {
-        super(version);
-        this.groupId = groupId;
-        this.memberId = memberId;
+    @Override
+    public LeaveGroupRequestData data() {
+        return data;
     }
 
-    public LeaveGroupRequest(Struct struct, short version) {
-        super(version);
-        groupId = struct.get(GROUP_ID);
-        memberId = struct.get(MEMBER_ID);
+    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) {
-        short versionId = version();
-        switch (versionId) {
-            case 0:
-                return new LeaveGroupResponse(Errors.forException(e));
-            case 1:
-                return new LeaveGroupResponse(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.LEAVE_GROUP.latestVersion()));
-        }
-    }
-
-    public String groupId() {
-        return groupId;
-    }
+        LeaveGroupResponseData responseData = new LeaveGroupResponseData()
+                                                  .setErrorCode(Errors.forException(e).code());
 
-    public String memberId() {
-        return memberId;
+        if (version() >= 1) {
+            responseData.setThrottleTimeMs(throttleTimeMs);
+        }
+        return new LeaveGroupResponse(responseData);
     }
 
     public static LeaveGroupRequest parse(ByteBuffer buffer, short version) {
-        return new LeaveGroupRequest(ApiKeys.LEAVE_GROUP.parseRequest(version, buffer), version);
-    }
-
-    @Override
-    protected Struct toStruct() {
-        Struct struct = new Struct(ApiKeys.LEAVE_GROUP.requestSchema(version()));
-        struct.set(GROUP_ID, groupId);
-        struct.set(MEMBER_ID, memberId);
-        return struct;
+        return new LeaveGroupRequest(new LeaveGroupRequestData(new ByteBufferAccessor(buffer), version), 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 f8682ecde7137..9a59139f4e77c 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
@@ -16,78 +16,138 @@
  */
 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.ByteBufferAccessor;
 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.HashMap;
+import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 
-import static org.apache.kafka.common.protocol.CommonFields.ERROR_CODE;
-import static org.apache.kafka.common.protocol.CommonFields.THROTTLE_TIME_MS;
-
+/**
+ * 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 static final Schema LEAVE_GROUP_RESPONSE_V0 = new Schema(
-            ERROR_CODE);
-    private static final Schema LEAVE_GROUP_RESPONSE_V1 = new Schema(
-            THROTTLE_TIME_MS,
-            ERROR_CODE);
+    private final LeaveGroupResponseData data;
+
+    public LeaveGroupResponse(LeaveGroupResponseData data) {
+        super(ApiKeys.LEAVE_GROUP);
+        this.data = data;
+    }
 
-    public static Schema[] schemaVersions() {
-        return new Schema[] {LEAVE_GROUP_RESPONSE_V0, LEAVE_GROUP_RESPONSE_V1};
+    public LeaveGroupResponse(List memberResponses,
+                              Errors topLevelError,
+                              final int throttleTimeMs,
+                              final short version) {
+        super(ApiKeys.LEAVE_GROUP);
+        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);
+        }
     }
 
-    /**
-     * Possible error code:
-     *
-     * GROUP_LOAD_IN_PROGRESS (14)
-     * CONSUMER_COORDINATOR_NOT_AVAILABLE (15)
-     * NOT_COORDINATOR_FOR_CONSUMER (16)
-     * UNKNOWN_CONSUMER_ID (25)
-     * GROUP_AUTHORIZATION_FAILED (30)
-     */
-    private final Errors error;
-    private final int throttleTimeMs;
-
-    public LeaveGroupResponse(Errors error) {
-        this(DEFAULT_THROTTLE_TIME, error);
+    @Override
+    public int throttleTimeMs() {
+        return data.throttleTimeMs();
     }
 
-    public LeaveGroupResponse(int throttleTimeMs, Errors error) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.error = error;
+    public List memberResponses() {
+        return data.members();
     }
 
-    public LeaveGroupResponse(Struct struct) {
-        this.throttleTimeMs = struct.getOrElse(THROTTLE_TIME_MS, DEFAULT_THROTTLE_TIME);
-        this.error = Errors.forCode(struct.get(ERROR_CODE));
+    public Errors error() {
+        return getError(Errors.forCode(data.errorCode()), data.members());
     }
 
-    public int throttleTimeMs() {
-        return throttleTimeMs;
+    public Errors topLevelError() {
+        return Errors.forCode(data.errorCode());
     }
 
-    public Errors error() {
-        return error;
+    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 errorCounts(error);
+        Map combinedErrorCounts = new HashMap<>();
+        // Top level error.
+        updateErrorCounts(combinedErrorCounts, Errors.forCode(data.errorCode()));
+
+        // Member level error.
+        data.members().forEach(memberResponse -> {
+            updateErrorCounts(combinedErrorCounts, Errors.forCode(memberResponse.errorCode()));
+        });
+        return combinedErrorCounts;
     }
 
     @Override
-    public Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.LEAVE_GROUP.responseSchema(version));
-        struct.setIfExists(THROTTLE_TIME_MS, throttleTimeMs);
-        struct.set(ERROR_CODE, error.code());
-        return struct;
+    public LeaveGroupResponseData data() {
+        return data;
     }
 
-    public static LeaveGroupResponse parse(ByteBuffer buffer, short versionId) {
-        return new LeaveGroupResponse(ApiKeys.LEAVE_GROUP.parseResponse(versionId, buffer));
+    public static LeaveGroupResponse parse(ByteBuffer buffer, short version) {
+        return new LeaveGroupResponse(new LeaveGroupResponseData(new ByteBufferAccessor(buffer), version));
     }
 
+    @Override
+    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);
+    }
+
+    @Override
+    public String toString() {
+        return data.toString();
+    }
 }
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 f279b4c618f46..ab7ec61f6c343 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,70 +16,73 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.errors.UnsupportedVersionException;
+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.ByteBufferAccessor;
 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();
+    public static class Builder extends AbstractRequest.Builder {
 
-    /* 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;
+        private final ListGroupsRequestData data;
 
-    public static Schema[] schemaVersions() {
-        return new Schema[] {LIST_GROUPS_REQUEST_V0, LIST_GROUPS_REQUEST_V1};
-    }
-
-    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);
+            if (!data.statesFilter().isEmpty() && version < 4) {
+                throw new UnsupportedVersionException("The broker only supports ListGroups " +
+                        "v" + version + ", but we need v4 or newer to request groups by states.");
+            }
+            return new ListGroupsRequest(data, version);
         }
 
         @Override
         public String toString() {
-            return "(type=ListGroupsRequest)";
+            return data.toString();
         }
     }
 
-    public ListGroupsRequest(short version) {
-        super(version);
-    }
+    private final ListGroupsRequestData data;
 
-    public ListGroupsRequest(Struct struct, short versionId) {
-        super(versionId);
+    public ListGroupsRequest(ListGroupsRequestData data, short version) {
+        super(ApiKeys.LIST_GROUPS, version);
+        this.data = data;
     }
 
     @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:
-                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) {
-        return new ListGroupsRequest(ApiKeys.LIST_GROUPS.parseRequest(version, buffer), version);
+        return new ListGroupsRequest(new ListGroupsRequestData(new ByteBufferAccessor(buffer), version), version);
     }
 
     @Override
-    protected Struct toStruct() {
-        return new Struct(ApiKeys.LIST_GROUPS.requestSchema(version()));
+    public ListGroupsRequestData data() {
+        return data;
     }
 }
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 9c82ae06bd637..270c43c0568ad 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,130 +16,44 @@
  */
 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.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.Struct;
 
 import java.nio.ByteBuffer;
-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.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)));
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {LIST_GROUPS_RESPONSE_V0, LIST_GROUPS_RESPONSE_V1};
-    }
-
-    /**
-     * Possible error codes:
-     *
-     * COORDINATOR_NOT_AVAILABLE (15)
-     * AUTHORIZATION_FAILED (29)
-     */
-
-    private final Errors error;
-    private final int throttleTimeMs;
-    private final List groups;
+    private final ListGroupsResponseData data;
 
-    public ListGroupsResponse(Errors error, List groups) {
-        this(DEFAULT_THROTTLE_TIME, error, groups);
+    public ListGroupsResponse(ListGroupsResponseData data) {
+        super(ApiKeys.LIST_GROUPS);
+        this.data = data;
     }
 
-    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));
-        }
+    @Override
+    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);
+        return errorCounts(Errors.forCode(data.errorCode()));
     }
 
-    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;
-        }
-
+    public static ListGroupsResponse parse(ByteBuffer buffer, short version) {
+        return new ListGroupsResponse(new ListGroupsResponseData(new ByteBufferAccessor(buffer), version));
     }
 
     @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;
+    public boolean shouldClientThrottle(short version) {
+        return version >= 2;
     }
-
-    public static ListGroupsResponse parse(ByteBuffer buffer, short version) {
-        return new ListGroupsResponse(ApiKeys.LIST_GROUPS.parseResponse(version, buffer));
-    }
-
 }
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 98f53bdfc44f0..9ebb06bdd6511 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
@@ -16,15 +16,6 @@
  */
 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.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.CollectionUtils;
-
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -34,11 +25,17 @@
 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 org.apache.kafka.common.protocol.types.Type.INT64;
-import static org.apache.kafka.common.protocol.types.Type.INT8;
+import org.apache.kafka.common.IsolationLevel;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.message.ListOffsetRequestData;
+import org.apache.kafka.common.message.ListOffsetRequestData.ListOffsetPartition;
+import org.apache.kafka.common.message.ListOffsetRequestData.ListOffsetTopic;
+import org.apache.kafka.common.message.ListOffsetResponseData;
+import org.apache.kafka.common.message.ListOffsetResponseData.ListOffsetPartitionResponse;
+import org.apache.kafka.common.message.ListOffsetResponseData.ListOffsetTopicResponse;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.Errors;
 
 public class ListOffsetRequest extends AbstractRequest {
     public static final long EARLIEST_TIMESTAMP = -2L;
@@ -47,65 +44,11 @@ public class ListOffsetRequest extends AbstractRequest {
     public static final int CONSUMER_REPLICA_ID = -1;
     public static final int DEBUGGING_REPLICA_ID = -2;
 
-    private static final String REPLICA_ID_KEY_NAME = "replica_id";
-    private static final String ISOLATION_LEVEL_KEY_NAME = "isolation_level";
-    private static final String TOPICS_KEY_NAME = "topics";
-
-    // topic level field names
-    private static final String PARTITIONS_KEY_NAME = "partitions";
-
-    // partition level field names
-    private static final String TIMESTAMP_KEY_NAME = "timestamp";
-    private static final String MAX_NUM_OFFSETS_KEY_NAME = "max_num_offsets";
-
-    private static final Schema LIST_OFFSET_REQUEST_PARTITION_V0 = new Schema(
-            PARTITION_ID,
-            new Field(TIMESTAMP_KEY_NAME, INT64, "Timestamp."),
-            new Field(MAX_NUM_OFFSETS_KEY_NAME, INT32, "Maximum offsets to return."));
-    private static final Schema LIST_OFFSET_REQUEST_PARTITION_V1 = new Schema(
-            PARTITION_ID,
-            new Field(TIMESTAMP_KEY_NAME, INT64, "The target timestamp for the partition."));
-
-    private static final Schema LIST_OFFSET_REQUEST_TOPIC_V0 = new Schema(
-            TOPIC_NAME,
-            new Field(PARTITIONS_KEY_NAME, new ArrayOf(LIST_OFFSET_REQUEST_PARTITION_V0), "Partitions to list offset."));
-    private static final Schema LIST_OFFSET_REQUEST_TOPIC_V1 = new Schema(
-            TOPIC_NAME,
-            new Field(PARTITIONS_KEY_NAME, new ArrayOf(LIST_OFFSET_REQUEST_PARTITION_V1), "Partitions to list offset."));
-
-    private static final Schema LIST_OFFSET_REQUEST_V0 = new Schema(
-            new Field(REPLICA_ID_KEY_NAME, INT32, "Broker id of the follower. For normal consumers, use -1."),
-            new Field(TOPICS_KEY_NAME, new ArrayOf(LIST_OFFSET_REQUEST_TOPIC_V0), "Topics to list offsets."));
-    private static final Schema LIST_OFFSET_REQUEST_V1 = new Schema(
-            new Field(REPLICA_ID_KEY_NAME, INT32, "Broker id of the follower. For normal consumers, use -1."),
-            new Field(TOPICS_KEY_NAME, new ArrayOf(LIST_OFFSET_REQUEST_TOPIC_V1), "Topics to list offsets."));
-
-    private static final Schema LIST_OFFSET_REQUEST_V2 = new Schema(
-            new Field(REPLICA_ID_KEY_NAME, INT32, "Broker id of the follower. For normal consumers, use -1."),
-            new Field(ISOLATION_LEVEL_KEY_NAME, INT8, "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"),
-            new Field(TOPICS_KEY_NAME, new ArrayOf(LIST_OFFSET_REQUEST_TOPIC_V1), "Topics to list offsets."));;
-
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {LIST_OFFSET_REQUEST_V0, LIST_OFFSET_REQUEST_V1, LIST_OFFSET_REQUEST_V2};
-    }
-
-    private final int replicaId;
-    private final IsolationLevel isolationLevel;
-    private final Map offsetData;
-    private final Map partitionTimestamps;
+    private final ListOffsetRequestData data;
     private final Set duplicatePartitions;
 
     public static class Builder extends AbstractRequest.Builder {
-        private final int replicaId;
-        private final IsolationLevel isolationLevel;
-        private Map offsetData = null;
-        private Map partitionTimestamps = null;
+        private final ListOffsetRequestData data;
 
         public static Builder forReplica(short allowedVersion, int replicaId) {
             return new Builder((short) 0, allowedVersion, replicaId, IsolationLevel.READ_UNCOMMITTED);
@@ -120,180 +63,95 @@ else if (requireTimestamp)
             return new Builder(minVersion, ApiKeys.LIST_OFFSETS.latestVersion(), CONSUMER_REPLICA_ID, isolationLevel);
         }
 
-        private Builder(short oldestAllowedVersion, short latestAllowedVersion, int replicaId, IsolationLevel isolationLevel) {
+        private Builder(short oldestAllowedVersion,
+                        short latestAllowedVersion,
+                        int replicaId,
+                        IsolationLevel isolationLevel) {
             super(ApiKeys.LIST_OFFSETS, oldestAllowedVersion, latestAllowedVersion);
-            this.replicaId = replicaId;
-            this.isolationLevel = isolationLevel;
+            data = new ListOffsetRequestData()
+                      .setIsolationLevel(isolationLevel.id())
+                      .setReplicaId(replicaId);
         }
 
-        public Builder setOffsetData(Map offsetData) {
-            this.offsetData = offsetData;
-            return this;
-        }
-
-        public Builder setTargetTimes(Map partitionTimestamps) {
-            this.partitionTimestamps = partitionTimestamps;
+        public Builder setTargetTimes(List topics) {
+            data.setTopics(topics);
             return this;
         }
 
         @Override
         public ListOffsetRequest build(short version) {
-            if (version == 0) {
-                if (offsetData == null) {
-                    if (partitionTimestamps == null) {
-                        throw new RuntimeException("Must set partitionTimestamps or offsetData when creating a v0 " +
-                            "ListOffsetRequest");
-                    } else {
-                        offsetData = new HashMap<>();
-                        for (Map.Entry entry: partitionTimestamps.entrySet()) {
-                            offsetData.put(entry.getKey(),
-                                    new PartitionData(entry.getValue(), 1));
-                        }
-                        this.partitionTimestamps = null;
-                    }
-                }
-            } else {
-                if (offsetData != null) {
-                    throw new RuntimeException("Cannot create a v" + version + " ListOffsetRequest with v0 " +
-                        "PartitionData.");
-                } else if (partitionTimestamps == null) {
-                    throw new RuntimeException("Must set partitionTimestamps when creating a v" +
-                            version + " ListOffsetRequest");
-                }
-            }
-            Map m = (version == 0) ?  offsetData : partitionTimestamps;
-            return new ListOffsetRequest(replicaId, m, isolationLevel, version);
-        }
-
-        @Override
-        public String toString() {
-            StringBuilder bld = new StringBuilder();
-            bld.append("(type=ListOffsetRequest")
-               .append(", replicaId=").append(replicaId);
-            if (offsetData != null) {
-                bld.append(", offsetData=").append(offsetData);
-            }
-            if (partitionTimestamps != null) {
-                bld.append(", partitionTimestamps=").append(partitionTimestamps);
-            }
-            bld.append(", isolationLevel=").append(isolationLevel);
-            bld.append(")");
-            return bld.toString();
-        }
-    }
-
-    /**
-     * This class is only used by ListOffsetRequest v0 which has been deprecated.
-     */
-    @Deprecated
-    public static final class PartitionData {
-        public final long timestamp;
-        public final int maxNumOffsets;
-
-        public PartitionData(long timestamp, int maxNumOffsets) {
-            this.timestamp = timestamp;
-            this.maxNumOffsets = maxNumOffsets;
+            return new ListOffsetRequest(data, version);
         }
 
         @Override
         public String toString() {
-            StringBuilder bld = new StringBuilder();
-            bld.append("{timestamp: ").append(timestamp).
-                append(", maxNumOffsets: ").append(maxNumOffsets).
-                append("}");
-            return bld.toString();
+            return data.toString();
         }
     }
 
     /**
      * Private constructor with a specified version.
      */
-    @SuppressWarnings("unchecked")
-    private ListOffsetRequest(int replicaId, Map targetTimes, IsolationLevel isolationLevel, short version) {
-        super(version);
-        this.replicaId = replicaId;
-        this.isolationLevel = isolationLevel;
-        this.offsetData = version == 0 ? (Map) targetTimes : null;
-        this.partitionTimestamps = version >= 1 ? (Map) targetTimes : null;
-        this.duplicatePartitions = Collections.emptySet();
-    }
-
-    public ListOffsetRequest(Struct struct, short version) {
-        super(version);
-        Set duplicatePartitions = new HashSet<>();
-        replicaId = struct.getInt(REPLICA_ID_KEY_NAME);
-        isolationLevel = struct.hasField(ISOLATION_LEVEL_KEY_NAME) ?
-                IsolationLevel.forId(struct.getByte(ISOLATION_LEVEL_KEY_NAME)) :
-                IsolationLevel.READ_UNCOMMITTED;
-        offsetData = new HashMap<>();
-        partitionTimestamps = new HashMap<>();
-        for (Object topicResponseObj : struct.getArray(TOPICS_KEY_NAME)) {
-            Struct topicResponse = (Struct) topicResponseObj;
-            String topic = topicResponse.get(TOPIC_NAME);
-            for (Object partitionResponseObj : topicResponse.getArray(PARTITIONS_KEY_NAME)) {
-                Struct partitionResponse = (Struct) partitionResponseObj;
-                int partition = partitionResponse.get(PARTITION_ID);
-                long timestamp = partitionResponse.getLong(TIMESTAMP_KEY_NAME);
-                TopicPartition tp = new TopicPartition(topic, partition);
-                if (partitionResponse.hasField(MAX_NUM_OFFSETS_KEY_NAME)) {
-                    int maxNumOffsets = partitionResponse.getInt(MAX_NUM_OFFSETS_KEY_NAME);
-                    PartitionData partitionData = new PartitionData(timestamp, maxNumOffsets);
-                    offsetData.put(tp, partitionData);
-                } else {
-                    if (partitionTimestamps.put(tp, timestamp) != null)
-                        duplicatePartitions.add(tp);
+    private ListOffsetRequest(ListOffsetRequestData data, short version) {
+        super(ApiKeys.LIST_OFFSETS, version);
+        this.data = data;
+        duplicatePartitions = new HashSet<>();
+        Set partitions = new HashSet<>();
+        for (ListOffsetTopic topic : data.topics()) {
+            for (ListOffsetPartition partition : topic.partitions()) {
+                TopicPartition tp = new TopicPartition(topic.name(), partition.partitionIndex());
+                if (!partitions.add(tp)) {
+                    duplicatePartitions.add(tp);
                 }
             }
         }
-        this.duplicatePartitions = duplicatePartitions;
     }
 
     @Override
-    @SuppressWarnings("deprecation")
     public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
-        Map responseData = new HashMap<>();
-
         short versionId = version();
-        if (versionId == 0) {
-            for (Map.Entry entry : offsetData.entrySet()) {
-                ListOffsetResponse.PartitionData partitionResponse = new ListOffsetResponse.PartitionData(
-                        Errors.forException(e), Collections.emptyList());
-                responseData.put(entry.getKey(), partitionResponse);
-            }
-        } else {
-            for (Map.Entry entry : partitionTimestamps.entrySet()) {
-                ListOffsetResponse.PartitionData partitionResponse = new ListOffsetResponse.PartitionData(
-                        Errors.forException(e), -1L, -1L);
-                responseData.put(entry.getKey(), partitionResponse);
+        short errorCode = Errors.forException(e).code();
+
+        List responses = new ArrayList<>();
+        for (ListOffsetTopic topic : data.topics()) {
+            ListOffsetTopicResponse topicResponse = new ListOffsetTopicResponse().setName(topic.name());
+            List partitions = new ArrayList<>();
+            for (ListOffsetPartition partition : topic.partitions()) {
+                ListOffsetPartitionResponse partitionResponse = new ListOffsetPartitionResponse()
+                        .setErrorCode(errorCode)
+                        .setPartitionIndex(partition.partitionIndex());
+                if (versionId == 0) {
+                    partitionResponse.setOldStyleOffsets(Collections.emptyList());
+                } else {
+                    partitionResponse.setOffset(ListOffsetResponse.UNKNOWN_OFFSET)
+                                     .setTimestamp(ListOffsetResponse.UNKNOWN_TIMESTAMP);
+                }
+                partitions.add(partitionResponse);
             }
+            topicResponse.setPartitions(partitions);
+            responses.add(topicResponse);
         }
+        ListOffsetResponseData responseData = new ListOffsetResponseData()
+                .setThrottleTimeMs(throttleTimeMs)
+                .setTopics(responses);
+        return new ListOffsetResponse(responseData);
+    }
 
-        switch (versionId) {
-            case 0:
-            case 1:
-            case 2:
-                return new ListOffsetResponse(throttleTimeMs, 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.LIST_OFFSETS.latestVersion()));
-        }
+    @Override
+    public ListOffsetRequestData data() {
+        return data;
     }
 
     public int replicaId() {
-        return replicaId;
+        return data.replicaId();
     }
 
     public IsolationLevel isolationLevel() {
-        return isolationLevel;
+        return IsolationLevel.forId(data.isolationLevel());
     }
 
-    @Deprecated
-    public Map offsetData() {
-        return offsetData;
-    }
-
-    public Map partitionTimestamps() {
-        return partitionTimestamps;
+    public List topics() {
+        return data.topics();
     }
 
     public Set duplicatePartitions() {
@@ -301,46 +159,25 @@ public Set duplicatePartitions() {
     }
 
     public static ListOffsetRequest parse(ByteBuffer buffer, short version) {
-        return new ListOffsetRequest(ApiKeys.LIST_OFFSETS.parseRequest(version, buffer), version);
+        return new ListOffsetRequest(new ListOffsetRequestData(new ByteBufferAccessor(buffer), version), version);
     }
 
-    @Override
-    protected Struct toStruct() {
-        short version = version();
-        Struct struct = new Struct(ApiKeys.LIST_OFFSETS.requestSchema(version));
-
-        Map targetTimes = partitionTimestamps == null ? offsetData : partitionTimestamps;
-        Map> topicsData = CollectionUtils.groupDataByTopic(targetTimes);
-
-        struct.set(REPLICA_ID_KEY_NAME, replicaId);
-
-        if (struct.hasField(ISOLATION_LEVEL_KEY_NAME))
-            struct.set(ISOLATION_LEVEL_KEY_NAME, isolationLevel.id());
-        List topicArray = new ArrayList<>();
-        for (Map.Entry> topicEntry: topicsData.entrySet()) {
-            Struct topicData = struct.instance(TOPICS_KEY_NAME);
-            topicData.set(TOPIC_NAME, topicEntry.getKey());
-            List partitionArray = new ArrayList<>();
-            for (Map.Entry partitionEntry : topicEntry.getValue().entrySet()) {
-                if (version == 0) {
-                    PartitionData offsetPartitionData = (PartitionData) partitionEntry.getValue();
-                    Struct partitionData = topicData.instance(PARTITIONS_KEY_NAME);
-                    partitionData.set(PARTITION_ID, partitionEntry.getKey());
-                    partitionData.set(TIMESTAMP_KEY_NAME, offsetPartitionData.timestamp);
-                    partitionData.set(MAX_NUM_OFFSETS_KEY_NAME, offsetPartitionData.maxNumOffsets);
-                    partitionArray.add(partitionData);
-                } else {
-                    Long timestamp = (Long) partitionEntry.getValue();
-                    Struct partitionData = topicData.instance(PARTITIONS_KEY_NAME);
-                    partitionData.set(PARTITION_ID, partitionEntry.getKey());
-                    partitionData.set(TIMESTAMP_KEY_NAME, timestamp);
-                    partitionArray.add(partitionData);
-                }
-            }
-            topicData.set(PARTITIONS_KEY_NAME, partitionArray.toArray());
-            topicArray.add(topicData);
+    public static List toListOffsetTopics(Map timestampsToSearch) {
+        Map topics = new HashMap<>();
+        for (Map.Entry entry : timestampsToSearch.entrySet()) {
+            TopicPartition tp = entry.getKey();
+            ListOffsetTopic topic = topics.computeIfAbsent(tp.topic(), k -> new ListOffsetTopic().setName(tp.topic()));
+            topic.partitions().add(entry.getValue());
         }
-        struct.set(TOPICS_KEY_NAME, topicArray.toArray());
-        return struct;
+        return new ArrayList<>(topics.values());
+    }
+
+    public static ListOffsetTopic singletonRequestData(String topic, int partitionIndex, long timestamp, int maxNumOffsets) {
+        return new ListOffsetTopic()
+                .setName(topic)
+                .setPartitions(Collections.singletonList(new ListOffsetPartition()
+                        .setPartitionIndex(partitionIndex)
+                        .setTimestamp(timestamp)
+                        .setMaxNumOffsets(maxNumOffsets)));
     }
 }
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 13f2dfb71066d..ca57382d9c0ca 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
@@ -16,232 +16,96 @@
  */
 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.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.CollectionUtils;
-import org.apache.kafka.common.utils.Utils;
-
 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.THROTTLE_TIME_MS;
-import static org.apache.kafka.common.protocol.CommonFields.TOPIC_NAME;
-import static org.apache.kafka.common.protocol.types.Type.INT64;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.message.ListOffsetResponseData;
+import org.apache.kafka.common.message.ListOffsetResponseData.ListOffsetPartitionResponse;
+import org.apache.kafka.common.message.ListOffsetResponseData.ListOffsetTopicResponse;
+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.record.RecordBatch;
 
+/**
+ * Possible error codes:
+ *
+ * - {@link Errors#UNSUPPORTED_FOR_MESSAGE_FORMAT} If the message format does not support lookup by timestamp
+ * - {@link Errors#TOPIC_AUTHORIZATION_FAILED} If the user does not have DESCRIBE access to a requested topic
+ * - {@link Errors#REPLICA_NOT_AVAILABLE} If the request is received by a broker with version < 2.6 which is not a replica
+ * - {@link Errors#NOT_LEADER_OR_FOLLOWER} If the broker is not a leader or follower and either the provided leader epoch
+ *     matches the known leader epoch on the broker or is empty
+ * - {@link Errors#FENCED_LEADER_EPOCH} If the epoch is lower than the broker's epoch
+ * - {@link Errors#UNKNOWN_LEADER_EPOCH} If the epoch is larger than the broker's epoch
+ * - {@link Errors#UNKNOWN_TOPIC_OR_PARTITION} If the broker does not have metadata for a topic or partition
+ * - {@link Errors#KAFKA_STORAGE_ERROR} If the log directory for one of the requested partitions is offline
+ * - {@link Errors#UNKNOWN_SERVER_ERROR} For any unexpected errors
+ * - {@link Errors#LEADER_NOT_AVAILABLE} The leader's HW has not caught up after recent election (v4 protocol)
+ * - {@link Errors#OFFSET_NOT_AVAILABLE} The leader's HW has not caught up after recent election (v5+ protocol)
+ */
 public class ListOffsetResponse extends AbstractResponse {
     public static final long UNKNOWN_TIMESTAMP = -1L;
     public static final long UNKNOWN_OFFSET = -1L;
+    public static final int UNKNOWN_EPOCH = RecordBatch.NO_PARTITION_LEADER_EPOCH;
 
-    private static final String RESPONSES_KEY_NAME = "responses";
-
-    // topic level field names
-    private static final String PARTITIONS_KEY_NAME = "partition_responses";
-
-    /**
-     * Possible error code:
-     *
-     *  UNKNOWN_TOPIC_OR_PARTITION (3)
-     *  NOT_LEADER_FOR_PARTITION (6)
-     *  UNSUPPORTED_FOR_MESSAGE_FORMAT (43)
-     *  UNKNOWN (-1)
-     */
-
-    // This key is only used by ListOffsetResponse v0
-    @Deprecated
-    private static final String OFFSETS_KEY_NAME = "offsets";
-    private static final String TIMESTAMP_KEY_NAME = "timestamp";
-    private static final String OFFSET_KEY_NAME = "offset";
-
-    private static final Schema LIST_OFFSET_RESPONSE_PARTITION_V0 = new Schema(
-            PARTITION_ID,
-            ERROR_CODE,
-            new Field(OFFSETS_KEY_NAME, new ArrayOf(INT64), "A list of offsets."));
-
-    private static final Schema LIST_OFFSET_RESPONSE_PARTITION_V1 = new Schema(
-            PARTITION_ID,
-            ERROR_CODE,
-            new Field(TIMESTAMP_KEY_NAME, INT64, "The timestamp associated with the returned offset"),
-            new Field(OFFSET_KEY_NAME, INT64, "offset found"));
-
-    private static final Schema LIST_OFFSET_RESPONSE_TOPIC_V0 = new Schema(
-            TOPIC_NAME,
-            new Field(PARTITIONS_KEY_NAME, new ArrayOf(LIST_OFFSET_RESPONSE_PARTITION_V0)));
-
-    private static final Schema LIST_OFFSET_RESPONSE_TOPIC_V1 = new Schema(
-            TOPIC_NAME,
-            new Field(PARTITIONS_KEY_NAME, new ArrayOf(LIST_OFFSET_RESPONSE_PARTITION_V1)));
-
-    private static final Schema LIST_OFFSET_RESPONSE_V0 = new Schema(
-            new Field(RESPONSES_KEY_NAME, new ArrayOf(LIST_OFFSET_RESPONSE_TOPIC_V0)));
-
-    private static final Schema LIST_OFFSET_RESPONSE_V1 = new Schema(
-            new Field(RESPONSES_KEY_NAME, new ArrayOf(LIST_OFFSET_RESPONSE_TOPIC_V1)));
-    private static final Schema LIST_OFFSET_RESPONSE_V2 = new Schema(
-            THROTTLE_TIME_MS,
-            new Field(RESPONSES_KEY_NAME, new ArrayOf(LIST_OFFSET_RESPONSE_TOPIC_V1)));
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {LIST_OFFSET_RESPONSE_V0, LIST_OFFSET_RESPONSE_V1, LIST_OFFSET_RESPONSE_V2};
-    }
-
-    public static final class PartitionData {
-        public final Errors error;
-        // The offsets list is only used in ListOffsetResponse v0.
-        @Deprecated
-        public final List offsets;
-        public final Long timestamp;
-        public final Long offset;
-
-        /**
-         * Constructor for ListOffsetResponse v0
-         */
-        @Deprecated
-        public PartitionData(Errors error, List offsets) {
-            this.error = error;
-            this.offsets = offsets;
-            this.timestamp = null;
-            this.offset = null;
-        }
-
-        /**
-         * Constructor for ListOffsetResponse v1
-         */
-        public PartitionData(Errors error, long timestamp, long offset) {
-            this.error = error;
-            this.timestamp = timestamp;
-            this.offset = offset;
-            this.offsets = null;
-        }
-
-        @Override
-        public String toString() {
-            StringBuilder bld = new StringBuilder();
-            bld.append("PartitionData{").
-                append("errorCode: ").append((int) error.code()).
-                append(", timestamp: ").append(timestamp).
-                append(", offset: ").append(offset).
-                append(", offsets: ");
-            if (offsets == null) {
-                bld.append(offsets);
-            } else {
-                bld.append("[").append(Utils.join(this.offsets, ",")).append("]");
-            }
-            bld.append("}");
-            return bld.toString();
-        }
-    }
-
-    private final int throttleTimeMs;
-    private final Map responseData;
+    private final ListOffsetResponseData data;
 
-    /**
-     * Constructor for all versions without throttle time
-     */
-    public ListOffsetResponse(Map responseData) {
-        this(DEFAULT_THROTTLE_TIME, responseData);
+    public ListOffsetResponse(ListOffsetResponseData data) {
+        super(ApiKeys.LIST_OFFSETS);
+        this.data = data;
     }
 
-    public ListOffsetResponse(int throttleTimeMs, Map responseData) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.responseData = responseData;
-    }
-
-    public ListOffsetResponse(Struct struct) {
-        this.throttleTimeMs = struct.getOrElse(THROTTLE_TIME_MS, DEFAULT_THROTTLE_TIME);
-        responseData = new HashMap<>();
-        for (Object topicResponseObj : struct.getArray(RESPONSES_KEY_NAME)) {
-            Struct topicResponse = (Struct) topicResponseObj;
-            String topic = topicResponse.get(TOPIC_NAME);
-            for (Object partitionResponseObj : topicResponse.getArray(PARTITIONS_KEY_NAME)) {
-                Struct partitionResponse = (Struct) partitionResponseObj;
-                int partition = partitionResponse.get(PARTITION_ID);
-                Errors error = Errors.forCode(partitionResponse.get(ERROR_CODE));
-                PartitionData partitionData;
-                if (partitionResponse.hasField(OFFSETS_KEY_NAME)) {
-                    Object[] offsets = partitionResponse.getArray(OFFSETS_KEY_NAME);
-                    List offsetsList = new ArrayList<>();
-                    for (Object offset : offsets)
-                        offsetsList.add((Long) offset);
-                    partitionData = new PartitionData(error, offsetsList);
-                } else {
-                    long timestamp = partitionResponse.getLong(TIMESTAMP_KEY_NAME);
-                    long offset = partitionResponse.getLong(OFFSET_KEY_NAME);
-                    partitionData = new PartitionData(error, timestamp, offset);
-                }
-                responseData.put(new TopicPartition(topic, partition), partitionData);
-            }
-        }
+    @Override
+    public int throttleTimeMs() {
+        return data.throttleTimeMs();
     }
 
-    public int throttleTimeMs() {
-        return throttleTimeMs;
+    @Override
+    public ListOffsetResponseData data() {
+        return data;
     }
 
-    public Map responseData() {
-        return responseData;
+    public List topics() {
+        return data.topics();
     }
 
     @Override
     public Map errorCounts() {
         Map errorCounts = new HashMap<>();
-        for (PartitionData response : responseData.values())
-            updateErrorCounts(errorCounts, response.error);
+        topics().forEach(topic ->
+            topic.partitions().forEach(partition ->
+                updateErrorCounts(errorCounts, Errors.forCode(partition.errorCode()))
+            )
+        );
         return errorCounts;
     }
 
     public static ListOffsetResponse parse(ByteBuffer buffer, short version) {
-        return new ListOffsetResponse(ApiKeys.LIST_OFFSETS.parseResponse(version, buffer));
+        return new ListOffsetResponse(new ListOffsetResponseData(new ByteBufferAccessor(buffer), version));
     }
 
     @Override
-    protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.LIST_OFFSETS.responseSchema(version));
-        struct.setIfExists(THROTTLE_TIME_MS, throttleTimeMs);
-        Map> topicsData = CollectionUtils.groupDataByTopic(responseData);
-
-        List topicArray = new ArrayList<>();
-        for (Map.Entry> topicEntry: topicsData.entrySet()) {
-            Struct topicData = struct.instance(RESPONSES_KEY_NAME);
-            topicData.set(TOPIC_NAME, topicEntry.getKey());
-            List partitionArray = new ArrayList<>();
-            for (Map.Entry partitionEntry : topicEntry.getValue().entrySet()) {
-                PartitionData offsetPartitionData = partitionEntry.getValue();
-                Struct partitionData = topicData.instance(PARTITIONS_KEY_NAME);
-                partitionData.set(PARTITION_ID, partitionEntry.getKey());
-                partitionData.set(ERROR_CODE, offsetPartitionData.error.code());
-                if (version == 0)
-                    partitionData.set(OFFSETS_KEY_NAME, offsetPartitionData.offsets.toArray());
-                else {
-                    partitionData.set(TIMESTAMP_KEY_NAME, offsetPartitionData.timestamp);
-                    partitionData.set(OFFSET_KEY_NAME, offsetPartitionData.offset);
-                }
-                partitionArray.add(partitionData);
-            }
-            topicData.set(PARTITIONS_KEY_NAME, partitionArray.toArray());
-            topicArray.add(topicData);
-        }
-        struct.set(RESPONSES_KEY_NAME, topicArray.toArray());
-
-        return struct;
+    public String toString() {
+        return data.toString();
     }
 
     @Override
-    public String toString() {
-        StringBuilder bld = new StringBuilder();
-        bld.append("(type=ListOffsetResponse")
-            .append(", throttleTimeMs=").append(throttleTimeMs)
-            .append(", responseData=").append(responseData)
-            .append(")");
-        return bld.toString();
+    public boolean shouldClientThrottle(short version) {
+        return version >= 3;
+    }
+
+    public static ListOffsetTopicResponse singletonListOffsetTopicResponse(TopicPartition tp, Errors error, long timestamp, long offset, int epoch) {
+        return new ListOffsetTopicResponse()
+                 .setName(tp.topic())
+                 .setPartitions(Collections.singletonList(new ListOffsetPartitionResponse()
+                         .setPartitionIndex(tp.partition())
+                         .setErrorCode(error.code())
+                         .setTimestamp(timestamp)
+                         .setOffset(offset)
+                         .setLeaderEpoch(epoch)));
     }
 }
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..03affd11b6706
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/ListPartitionReassignmentsRequest.java
@@ -0,0 +1,93 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.ByteBufferAccessor;
+
+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 ListPartitionReassignmentsRequest(ListPartitionReassignmentsRequestData data, short version) {
+        super(ApiKeys.LIST_PARTITION_REASSIGNMENTS, version);
+        this.data = data;
+    }
+
+    public static ListPartitionReassignmentsRequest parse(ByteBuffer buffer, short version) {
+        return new ListPartitionReassignmentsRequest(new ListPartitionReassignmentsRequestData(
+            new ByteBufferAccessor(buffer), version), version);
+    }
+
+    @Override
+    public ListPartitionReassignmentsRequestData data() {
+        return data;
+    }
+
+    @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..4a890e8b50cd1
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/ListPartitionReassignmentsResponse.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.message.ListPartitionReassignmentsResponseData;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.Errors;
+
+import java.nio.ByteBuffer;
+import java.util.Map;
+
+public class ListPartitionReassignmentsResponse extends AbstractResponse {
+
+    private final ListPartitionReassignmentsResponseData data;
+
+    public ListPartitionReassignmentsResponse(ListPartitionReassignmentsResponseData responseData) {
+        super(ApiKeys.LIST_PARTITION_REASSIGNMENTS);
+        this.data = responseData;
+    }
+
+    public static ListPartitionReassignmentsResponse parse(ByteBuffer buffer, short version) {
+        return new ListPartitionReassignmentsResponse(new ListPartitionReassignmentsResponseData(
+            new ByteBufferAccessor(buffer), version));
+    }
+
+    @Override
+    public ListPartitionReassignmentsResponseData data() {
+        return data;
+    }
+
+    @Override
+    public boolean shouldClientThrottle(short version) {
+        return true;
+    }
+
+    @Override
+    public int throttleTimeMs() {
+        return data.throttleTimeMs();
+    }
+
+    @Override
+    public Map errorCounts() {
+        return errorCounts(Errors.forCode(data.errorCode()));
+    }
+}
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 934b0ed06d153..816f600061568 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
@@ -16,192 +16,147 @@
  */
 package org.apache.kafka.common.requests;
 
-import org.apache.kafka.common.Node;
 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.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.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.BOOLEAN;
-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 String ALLOW_AUTO_TOPIC_CREATION_KEY_NAME = "allow_auto_topic_creation";
-
-    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;
+    public static class Builder extends AbstractRequest.Builder {
+        private static final MetadataRequestData ALL_TOPICS_REQUEST_DATA = new MetadataRequestData().
+            setTopics(null).setAllowAutoTopicCreation(true);
 
-    /* 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;
+        private final MetadataRequestData data;
 
-    /* The v4 metadata request has an additional field for allowing auto topic creation. The response is the same as v3. */
-    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."),
-            new Field(ALLOW_AUTO_TOPIC_CREATION_KEY_NAME, BOOLEAN, "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."));
+        public Builder(MetadataRequestData data) {
+            super(ApiKeys.METADATA);
+            this.data = data;
+        }
 
-    /* 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;
+        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};
-    }
+        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;
-
-    /**
-     * 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) {
-        super(version);
-        this.topics = topics;
-        this.allowAutoTopicCreation = allowAutoTopicCreation;
+    private final MetadataRequestData data;
+
+    public MetadataRequest(MetadataRequestData data, short version) {
+        super(ApiKeys.METADATA, version);
+        this.data = data;
     }
 
-    public MetadataRequest(Struct struct, short version) {
-        super(version);
-        Object[] topicArray = struct.getArray(TOPICS_KEY_NAME);
-        if (topicArray != null) {
-            topics = new ArrayList<>();
-            for (Object topicObj: topicArray) {
-                topics.add((String) topicObj);
-            }
-        } else {
-            topics = null;
-        }
-        if (struct.hasField(ALLOW_AUTO_TOPIC_CREATION_KEY_NAME))
-            allowAutoTopicCreation = struct.getBoolean(ALLOW_AUTO_TOPIC_CREATION_KEY_NAME);
-        else
-            allowAutoTopicCreation = true;
+    @Override
+    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();
-        switch (versionId) {
-            case 0:
-            case 1:
-            case 2:
-                return new MetadataResponse(Collections.emptyList(), null, MetadataResponse.NO_CONTROLLER_ID, topicMetadatas);
-            case 3:
-            case 4:
-            case 5:
-                return new MetadataResponse(throttleTimeMs, Collections.emptyList(), null, MetadataResponse.NO_CONTROLLER_ID, topicMetadatas);
-            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()));
-        }
+        responseData.setThrottleTimeMs(throttleTimeMs);
+        return new MetadataResponse(responseData, true);
     }
 
     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);
+        return new MetadataRequest(new MetadataRequestData(new ByteBufferAccessor(buffer), version), version);
     }
 
-    @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());
-        if (struct.hasField(ALLOW_AUTO_TOPIC_CREATION_KEY_NAME))
-            struct.set(ALLOW_AUTO_TOPIC_CREATION_KEY_NAME, allowAutoTopicCreation);
-        return struct;
+    public static List convertToMetadataRequestTopic(final Collection topics) {
+        return topics.stream().map(topic -> new MetadataRequestTopic()
+            .setName(topic))
+            .collect(Collectors.toList());
     }
 }
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 99c4ffb5fb5b6..097c26e687dea 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,266 +19,68 @@
 import org.apache.kafka.common.Cluster;
 import org.apache.kafka.common.Node;
 import org.apache.kafka.common.PartitionInfo;
-import org.apache.kafka.common.errors.InvalidMetadataException;
+import org.apache.kafka.common.TopicPartition;
+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.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.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.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 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.protocol.types.Type.BOOLEAN;
-import static org.apache.kafka.common.protocol.types.Type.INT32;
-import static org.apache.kafka.common.protocol.types.Type.NULLABLE_STRING;
-import static org.apache.kafka.common.protocol.types.Type.STRING;
-
+import java.util.function.Function;
+import java.util.stream.Collectors;
+
+/**
+ * Possible topic-level error codes:
+ *  UnknownTopic (3)
+ *  LeaderNotAvailable (5)
+ *  InvalidTopic (17)
+ *  TopicAuthorizationFailed (29)
+
+ * Possible partition-level error codes:
+ *  LeaderNotAvailable (5)
+ *  ReplicaNotAvailable (9)
+ */
 public class MetadataResponse extends AbstractResponse {
-    private static final String BROKERS_KEY_NAME = "brokers";
-    private static final String TOPIC_METADATA_KEY_NAME = "topic_metadata";
-
-    // broker 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 String RACK_KEY_NAME = "rack";
-
-    private static final String CONTROLLER_ID_KEY_NAME = "controller_id";
     public static final int NO_CONTROLLER_ID = -1;
+    public static final int NO_LEADER_ID = -1;
+    public static final int AUTHORIZED_OPERATIONS_OMITTED = Integer.MIN_VALUE;
 
-    private static final String CLUSTER_ID_KEY_NAME = "cluster_id";
-
-    /**
-     * Possible error codes:
-     *
-     * UnknownTopic (3)
-     * LeaderNotAvailable (5)
-     * InvalidTopic (17)
-     * TopicAuthorizationFailed (29)
-     */
-
-    private static final String IS_INTERNAL_KEY_NAME = "is_internal";
-    private static final String PARTITION_METADATA_KEY_NAME = "partition_metadata";
-
-    /**
-     * Possible error codes:
-     *
-     * LeaderNotAvailable (5)
-     * ReplicaNotAvailable (9)
-     */
-
-    private static final String LEADER_KEY_NAME = "leader";
-    private static final String REPLICAS_KEY_NAME = "replicas";
-    private static final String ISR_KEY_NAME = "isr";
-    private static final String OFFLINE_REPLICAS_KEY_NAME = "offline_replicas";
-
-    private static final Schema METADATA_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 PARTITION_METADATA_V0 = new Schema(
-            ERROR_CODE,
-            PARTITION_ID,
-            new Field(LEADER_KEY_NAME, INT32, "The id of the broker acting as leader for this partition."),
-            new Field(REPLICAS_KEY_NAME, new ArrayOf(INT32), "The set of all nodes that host this partition."),
-            new Field(ISR_KEY_NAME, new ArrayOf(INT32), "The set of nodes that are in sync with the leader for this partition."));
-
-    private static final Schema TOPIC_METADATA_V0 = new Schema(
-            ERROR_CODE,
-            TOPIC_NAME,
-            new Field(PARTITION_METADATA_KEY_NAME, new ArrayOf(PARTITION_METADATA_V0), "Metadata for each partition of the topic."));
-
-    private static final Schema METADATA_RESPONSE_V0 = new Schema(
-            new Field(BROKERS_KEY_NAME, new ArrayOf(METADATA_BROKER_V0), "Host and port information for all brokers."),
-            new Field(TOPIC_METADATA_KEY_NAME, new ArrayOf(TOPIC_METADATA_V0)));
-
-    private static final Schema METADATA_BROKER_V1 = 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."),
-            new Field(RACK_KEY_NAME, NULLABLE_STRING, "The rack of the broker."));
-
-    private static final Schema PARTITION_METADATA_V1 = PARTITION_METADATA_V0;
-
-    // PARTITION_METADATA_V2 added a per-partition offline_replicas field. This field specifies the list of replicas that are offline.
-    private static final Schema PARTITION_METADATA_V2 = new Schema(
-            ERROR_CODE,
-            PARTITION_ID,
-            new Field(LEADER_KEY_NAME, INT32, "The id of the broker acting as leader for this partition."),
-            new Field(REPLICAS_KEY_NAME, new ArrayOf(INT32), "The set of all nodes that host this partition."),
-            new Field(ISR_KEY_NAME, new ArrayOf(INT32), "The set of nodes that are in sync with the leader for this partition."),
-            new Field(OFFLINE_REPLICAS_KEY_NAME, new ArrayOf(INT32), "The set of offline replicas of this partition."));
-
-    private static final Schema TOPIC_METADATA_V1 = new Schema(
-            ERROR_CODE,
-            TOPIC_NAME,
-            new Field(IS_INTERNAL_KEY_NAME, BOOLEAN, "Indicates if the topic is considered a Kafka internal topic"),
-            new Field(PARTITION_METADATA_KEY_NAME, new ArrayOf(PARTITION_METADATA_V1), "Metadata for each partition of the topic."));
-
-    // TOPIC_METADATA_V2 added a per-partition offline_replicas field. This field specifies the list of replicas that are offline.
-    private static final Schema TOPIC_METADATA_V2 = new Schema(
-            ERROR_CODE,
-            TOPIC_NAME,
-            new Field(IS_INTERNAL_KEY_NAME, BOOLEAN, "Indicates if the topic is considered a Kafka internal topic"),
-            new Field(PARTITION_METADATA_KEY_NAME, new ArrayOf(PARTITION_METADATA_V2), "Metadata for each partition of the topic."));
-
-    private static final Schema METADATA_RESPONSE_V1 = new Schema(
-            new Field(BROKERS_KEY_NAME, new ArrayOf(METADATA_BROKER_V1), "Host and port information for all brokers."),
-            new Field(CONTROLLER_ID_KEY_NAME, INT32, "The broker id of the controller broker."),
-            new Field(TOPIC_METADATA_KEY_NAME, new ArrayOf(TOPIC_METADATA_V1)));
-
-    private static final Schema METADATA_RESPONSE_V2 = new Schema(
-            new Field(BROKERS_KEY_NAME, new ArrayOf(METADATA_BROKER_V1), "Host and port information for all brokers."),
-            new Field(CLUSTER_ID_KEY_NAME, NULLABLE_STRING, "The cluster id that this broker belongs to."),
-            new Field(CONTROLLER_ID_KEY_NAME, INT32, "The broker id of the controller broker."),
-            new Field(TOPIC_METADATA_KEY_NAME, new ArrayOf(TOPIC_METADATA_V1)));
-
-    private static final Schema METADATA_RESPONSE_V3 = new Schema(
-            THROTTLE_TIME_MS,
-            new Field(BROKERS_KEY_NAME, new ArrayOf(METADATA_BROKER_V1), "Host and port information for all brokers."),
-            new Field(CLUSTER_ID_KEY_NAME, NULLABLE_STRING, "The cluster id that this broker belongs to."),
-            new Field(CONTROLLER_ID_KEY_NAME, INT32, "The broker id of the controller broker."),
-            new Field(TOPIC_METADATA_KEY_NAME, new ArrayOf(TOPIC_METADATA_V1)));
-
-    private static final Schema METADATA_RESPONSE_V4 = METADATA_RESPONSE_V3;
-
-    // METADATA_RESPONSE_V5 added a per-partition offline_replicas field. This field specifies the list of replicas that are offline.
-    private static final Schema METADATA_RESPONSE_V5 = new Schema(
-            THROTTLE_TIME_MS,
-            new Field(BROKERS_KEY_NAME, new ArrayOf(METADATA_BROKER_V1), "Host and port information for all brokers."),
-            new Field(CLUSTER_ID_KEY_NAME, NULLABLE_STRING, "The cluster id that this broker belongs to."),
-            new Field(CONTROLLER_ID_KEY_NAME, INT32, "The broker id of the controller broker."),
-            new Field(TOPIC_METADATA_KEY_NAME, new ArrayOf(TOPIC_METADATA_V2)));
-
-    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};
-    }
-
-    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, short version) {
+        this(data, hasReliableLeaderEpochs(version));
     }
 
-    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;
+    MetadataResponse(MetadataResponseData data, boolean hasReliableLeaderEpochs) {
+        super(ApiKeys.METADATA);
+        this.data = data;
+        this.hasReliableLeaderEpochs = hasReliableLeaderEpochs;
     }
 
-    public MetadataResponse(Struct struct) {
-        this.throttleTimeMs = struct.getOrElse(THROTTLE_TIME_MS, DEFAULT_THROTTLE_TIME);
-        Map brokers = new HashMap<>();
-        Object[] brokerStructs = (Object[]) struct.get(BROKERS_KEY_NAME);
-        for (Object brokerStruct : brokerStructs) {
-            Struct broker = (Struct) brokerStruct;
-            int nodeId = broker.getInt(NODE_ID_KEY_NAME);
-            String host = broker.getString(HOST_KEY_NAME);
-            int port = broker.getInt(PORT_KEY_NAME);
-            // 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.hasField(RACK_KEY_NAME) ? broker.getString(RACK_KEY_NAME) : 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 = NO_CONTROLLER_ID;
-        if (struct.hasField(CONTROLLER_ID_KEY_NAME))
-            controllerId = struct.getInt(CONTROLLER_ID_KEY_NAME);
-
-        // This field only exists in v2+
-        if (struct.hasField(CLUSTER_ID_KEY_NAME)) {
-            this.clusterId = struct.getString(CLUSTER_ID_KEY_NAME);
-        } else {
-            this.clusterId = null;
-        }
-
-        List topicMetadata = new ArrayList<>();
-        Object[] topicInfos = (Object[]) struct.get(TOPIC_METADATA_KEY_NAME);
-        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.hasField(IS_INTERNAL_KEY_NAME) ? topicInfo.getBoolean(IS_INTERNAL_KEY_NAME) : false;
-
-            List partitionMetadata = new ArrayList<>();
-
-            Object[] partitionInfos = (Object[]) topicInfo.get(PARTITION_METADATA_KEY_NAME);
-            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.getInt(LEADER_KEY_NAME);
-                Node leaderNode = leader == -1 ? null : brokers.get(leader);
-
-                Object[] replicas = (Object[]) partitionInfo.get(REPLICAS_KEY_NAME);
-                List replicaNodes = convertToNodes(brokers, replicas);
-
-                Object[] isr = (Object[]) partitionInfo.get(ISR_KEY_NAME);
-                List isrNodes = convertToNodes(brokers, isr);
-
-                Object[] offlineReplicas = partitionInfo.hasField(OFFLINE_REPLICAS_KEY_NAME) ?
-                    (Object[]) partitionInfo.get(OFFLINE_REPLICAS_KEY_NAME) : new Object[0];
-                List offlineNodes = convertToNodes(brokers, offlineReplicas);
-
-                partitionMetadata.add(new PartitionMetadata(partitionError, partition, leaderNode, replicaNodes, isrNodes, offlineNodes));
-            }
-
-            topicMetadata.add(new TopicMetadata(topicError, topic, isInternal, partitionMetadata));
-        }
-
-        this.brokers = brokers.values();
-        this.controller = getControllerNode(controllerId, brokers.values());
-        this.topicMetadata = topicMetadata;
-    }
-
-    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 Node getControllerNode(int controllerId, Collection brokers) {
-        for (Node broker : brokers) {
-            if (broker.id() == controllerId)
-                return broker;
-        }
-        return null;
+    @Override
+    public MetadataResponseData data() {
+        return data;
     }
 
+    @Override
     public int throttleTimeMs() {
-        return throttleTimeMs;
+        return data.throttleTimeMs();
     }
 
     /**
@@ -287,9 +89,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;
     }
@@ -297,8 +99,10 @@ public Map errors() {
     @Override
     public Map errorCounts() {
         Map errorCounts = new HashMap<>();
-        for (TopicMetadata metadata : topicMetadata)
-            updateErrorCounts(errorCounts, metadata.error);
+        data.topics().forEach(metadata -> {
+            metadata.partitions().forEach(p -> updateErrorCounts(errorCounts, Errors.forCode(p.errorCode())));
+            updateErrorCounts(errorCounts, Errors.forCode(metadata.errorCode()));
+        });
         return errorCounts;
     }
 
@@ -307,36 +111,13 @@ 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;
     }
 
-    /**
-     * Returns the set of topics with an error indicating invalid metadata
-     * and topics with any partition whose error indicates invalid metadata.
-     * This includes all non-existent topics specified in the metadata request
-     * and any topic returned with one or more partitions whose leader is not known.
-     */
-    public Set unavailableTopics() {
-        Set invalidMetadataTopics = new HashSet<>();
-        for (TopicMetadata topicMetadata : this.topicMetadata) {
-            if (topicMetadata.error.exception() instanceof InvalidMetadataException)
-                invalidMetadataTopics.add(topicMetadata.topic);
-            else {
-                for (PartitionMetadata partitionMetadata : topicMetadata.partitionMetadata) {
-                    if (partitionMetadata.error.exception() instanceof InvalidMetadataException) {
-                        invalidMetadataTopics.add(topicMetadata.topic);
-                        break;
-                    }
-                }
-            }
-        }
-        return invalidMetadataTopics;
-    }
-
     /**
      * Get a snapshot of the cluster metadata from this response
      * @return the cluster snapshot
@@ -344,23 +125,64 @@ public Set unavailableTopics() {
     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);
-                for (PartitionMetadata partitionMetadata : metadata.partitionMetadata)
-                    partitions.add(new PartitionInfo(
-                            metadata.topic,
-                            partitionMetadata.partition,
-                            partitionMetadata.leader,
-                            partitionMetadata.replicas.toArray(new Node[0]),
-                            partitionMetadata.isr.toArray(new Node[0]),
-                            partitionMetadata.offlineReplicas.toArray(new Node[0])));
+                for (PartitionMetadata partitionMetadata : metadata.partitionMetadata) {
+                    partitions.add(toPartitionInfo(partitionMetadata, holder().brokers));
+                }
             }
         }
+        return new Cluster(data.clusterId(), brokers(), partitions, topicsByError(Errors.TOPIC_AUTHORIZATION_FAILED),
+                topicsByError(Errors.INVALID_TOPIC_EXCEPTION), internalTopics, controller());
+    }
+
+    public static PartitionInfo toPartitionInfo(PartitionMetadata metadata, Map nodesById) {
+        return new PartitionInfo(metadata.topic(),
+                metadata.partition(),
+                metadata.leaderId.map(nodesById::get).orElse(null),
+                convertToNodeArray(metadata.replicaIds, nodesById),
+                convertToNodeArray(metadata.inSyncReplicaIds, nodesById),
+                convertToNodeArray(metadata.offlineReplicaIds, nodesById));
+    }
+
+    private static Node[] convertToNodeArray(List replicaIds, Map nodesById) {
+        return replicaIds.stream().map(replicaId -> {
+            Node node = nodesById.get(replicaId);
+            if (node == null)
+                return new Node(replicaId, "", -1);
+            return node;
+        }).toArray(Node[]::new);
+    }
+
+    /**
+     * 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());
+    }
+
+    /**
+     * Returns a 32-bit bitfield to represent authorized operations for this cluster.
+     */
+    public int clusterAuthorizedOperations() {
+        return data.clusterAuthorizedOperations();
+    }
 
-        return new Cluster(this.clusterId, this.brokers, partitions, topicsByError(Errors.TOPIC_AUTHORIZATION_FAILED),
-                internalTopics, this.controller);
+    private Holder holder() {
+        if (holder == null) {
+            synchronized (data) {
+                if (holder == null)
+                    holder = new Holder(data);
+            }
+        }
+        return holder;
     }
 
     /**
@@ -368,7 +190,11 @@ public Cluster cluster() {
      * @return the brokers
      */
     public Collection brokers() {
-        return brokers;
+        return holder().brokers.values();
+    }
+
+    public Map brokersById() {
+        return holder().brokers;
     }
 
     /**
@@ -376,7 +202,7 @@ public Collection brokers() {
      * @return the topicMetadata
      */
     public Collection topicMetadata() {
-        return topicMetadata;
+        return holder().topicMetadata;
     }
 
     /**
@@ -384,7 +210,7 @@ public Collection topicMetadata() {
      * @return the controller node or null if it doesn't exist
      */
     public Node controller() {
-        return controller;
+        return holder().controller;
     }
 
     /**
@@ -392,11 +218,34 @@ 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;
+    }
+
+    // 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.
+    private static boolean hasReliableLeaderEpochs(short version) {
+        return version >= 9;
     }
 
     public static MetadataResponse parse(ByteBuffer buffer, short version) {
-        return new MetadataResponse(ApiKeys.METADATA.parseResponse(version, buffer));
+        return new MetadataResponse(new MetadataResponseData(new ByteBufferAccessor(buffer), version),
+            hasReliableLeaderEpochs(version));
     }
 
     public static class TopicMetadata {
@@ -404,15 +253,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, AUTHORIZED_OPERATIONS_OMITTED);
         }
 
         public Errors error() {
@@ -431,128 +290,185 @@ public List partitionMetadata() {
             return partitionMetadata;
         }
 
-    }
+        public void authorizedOperations(int authorizedOperations) {
+            this.authorizedOperations = authorizedOperations;
+        }
 
-    // This is used to describe per-partition state in the MetadataResponse
-    public static class PartitionMetadata {
-        private final Errors error;
-        private final int partition;
-        private final Node leader;
-        private final List replicas;
-        private final List isr;
-        private final List offlineReplicas;
+        public int authorizedOperations() {
+            return authorizedOperations;
+        }
 
-        public PartitionMetadata(Errors error,
-                                 int partition,
-                                 Node leader,
-                                 List replicas,
-                                 List isr,
-                                 List offlineReplicas) {
-            this.error = error;
-            this.partition = partition;
-            this.leader = leader;
-            this.replicas = replicas;
-            this.isr = isr;
-            this.offlineReplicas = offlineReplicas;
+        @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);
         }
 
-        public Errors error() {
-            return error;
+        @Override
+        public int hashCode() {
+            return Objects.hash(error, topic, isInternal, partitionMetadata, authorizedOperations);
         }
 
-        public int partition() {
-            return partition;
+        @Override
+        public String toString() {
+            return "TopicMetadata{" +
+                "error=" + error +
+                ", topic='" + topic + '\'' +
+                ", isInternal=" + isInternal +
+                ", partitionMetadata=" + partitionMetadata +
+                ", authorizedOperations=" + authorizedOperations +
+                '}';
         }
+    }
 
-        public Node leader() {
-            return leader;
+    // This is used to describe per-partition state in the MetadataResponse
+    public static class PartitionMetadata {
+        public final TopicPartition topicPartition;
+        public final Errors error;
+        public final Optional leaderId;
+        public final Optional leaderEpoch;
+        public final List replicaIds;
+        public final List inSyncReplicaIds;
+        public final List offlineReplicaIds;
+
+        public PartitionMetadata(Errors error,
+                                 TopicPartition topicPartition,
+                                 Optional leaderId,
+                                 Optional leaderEpoch,
+                                 List replicaIds,
+                                 List inSyncReplicaIds,
+                                 List offlineReplicaIds) {
+            this.error = error;
+            this.topicPartition = topicPartition;
+            this.leaderId = leaderId;
+            this.leaderEpoch = leaderEpoch;
+            this.replicaIds = replicaIds;
+            this.inSyncReplicaIds = inSyncReplicaIds;
+            this.offlineReplicaIds = offlineReplicaIds;
         }
 
-        public List replicas() {
-            return replicas;
+        public int partition() {
+            return topicPartition.partition();
         }
 
-        public List isr() {
-            return isr;
+        public String topic() {
+            return topicPartition.topic();
         }
 
-        public List offlineReplicas() {
-            return offlineReplicas;
+        public PartitionMetadata withoutLeaderEpoch() {
+            return new PartitionMetadata(error,
+                    topicPartition,
+                    leaderId,
+                    Optional.empty(),
+                    replicaIds,
+                    inSyncReplicaIds,
+                    offlineReplicaIds);
         }
 
         @Override
         public String toString() {
-            return "(type=PartitionMetadata," +
-                    ", error=" + error +
-                    ", partition=" + partition +
-                    ", leader=" + leader +
-                    ", replicas=" + Utils.join(replicas, ",") +
-                    ", isr=" + Utils.join(isr, ",") +
-                    ", offlineReplicas=" + Utils.join(offlineReplicas, ",") + ')';
+            return "PartitionMetadata(" +
+                    "error=" + error +
+                    ", partition=" + topicPartition +
+                    ", leader=" + leaderId +
+                    ", leaderEpoch=" + leaderEpoch +
+                    ", replicas=" + Utils.join(replicaIds, ",") +
+                    ", isr=" + Utils.join(inSyncReplicaIds, ",") +
+                    ", offlineReplicas=" + Utils.join(offlineReplicaIds, ",") + ')';
         }
     }
 
-    @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_KEY_NAME);
-            broker.set(NODE_ID_KEY_NAME, node.id());
-            broker.set(HOST_KEY_NAME, node.host());
-            broker.set(PORT_KEY_NAME, node.port());
-            // This field only exists in v1+
-            if (broker.hasField(RACK_KEY_NAME))
-                broker.set(RACK_KEY_NAME, 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 = Collections.unmodifiableMap(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, Function.identity()));
         }
-        struct.set(BROKERS_KEY_NAME, brokerArray.toArray());
-
-        // This field only exists in v1+
-        if (struct.hasField(CONTROLLER_ID_KEY_NAME))
-            struct.set(CONTROLLER_ID_KEY_NAME, controller == null ? NO_CONTROLLER_ID : controller.id());
-
-        // This field only exists in v2+
-        if (struct.hasField(CLUSTER_ID_KEY_NAME))
-            struct.set(CLUSTER_ID_KEY_NAME, clusterId);
-
-        List topicMetadataArray = new ArrayList<>(topicMetadata.size());
-        for (TopicMetadata metadata : topicMetadata) {
-            Struct topicData = struct.instance(TOPIC_METADATA_KEY_NAME);
-            topicData.set(TOPIC_NAME, metadata.topic);
-            topicData.set(ERROR_CODE, metadata.error.code());
-            // This field only exists in v1+
-            if (topicData.hasField(IS_INTERNAL_KEY_NAME))
-                topicData.set(IS_INTERNAL_KEY_NAME, metadata.isInternal());
-
-            List partitionMetadataArray = new ArrayList<>(metadata.partitionMetadata.size());
-            for (PartitionMetadata partitionMetadata : metadata.partitionMetadata()) {
-                Struct partitionData = topicData.instance(PARTITION_METADATA_KEY_NAME);
-                partitionData.set(ERROR_CODE, partitionMetadata.error.code());
-                partitionData.set(PARTITION_ID, partitionMetadata.partition);
-                partitionData.set(LEADER_KEY_NAME, partitionMetadata.leader.id());
-                ArrayList replicas = new ArrayList<>(partitionMetadata.replicas.size());
-                for (Node node : partitionMetadata.replicas)
-                    replicas.add(node.id());
-                partitionData.set(REPLICAS_KEY_NAME, replicas.toArray());
-                ArrayList isr = new ArrayList<>(partitionMetadata.isr.size());
-                for (Node node : partitionMetadata.isr)
-                    isr.add(node.id());
-                partitionData.set(ISR_KEY_NAME, isr.toArray());
-                if (partitionData.hasField(OFFLINE_REPLICAS_KEY_NAME)) {
-                    ArrayList offlineReplicas = new ArrayList<>(partitionMetadata.offlineReplicas.size());
-                    for (Node node : partitionMetadata.offlineReplicas)
-                        offlineReplicas.add(node.id());
-                    partitionData.set(OFFLINE_REPLICAS_KEY_NAME, 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 leaderId = partitionMetadata.leaderId();
+                    Optional leaderIdOpt = leaderId < 0 ? Optional.empty() : Optional.of(leaderId);
+
+                    Optional leaderEpoch = RequestUtils.getLeaderEpoch(partitionMetadata.leaderEpoch());
+                    TopicPartition topicPartition = new TopicPartition(topic, partitionIndex);
+                    partitionMetadataList.add(new PartitionMetadata(partitionError, topicPartition, leaderIdOpt,
+                            leaderEpoch, partitionMetadata.replicaNodes(), partitionMetadata.isrNodes(),
+                            partitionMetadata.offlineReplicas()));
                 }
-                partitionMetadataArray.add(partitionData);
 
+                topicMetadataList.add(new TopicMetadata(topicError, topic, isInternal, partitionMetadataList,
+                        topicMetadata.topicAuthorizedOperations()));
             }
-            topicData.set(PARTITION_METADATA_KEY_NAME, partitionMetadataArray.toArray());
-            topicMetadataArray.add(topicData);
+            return topicMetadataList;
         }
-        struct.set(TOPIC_METADATA_KEY_NAME, topicMetadataArray.toArray());
-        return struct;
+
+    }
+
+    public static MetadataResponse prepareResponse(short version,
+                                                   int throttleTimeMs,
+                                                   Collection brokers,
+                                                   String clusterId,
+                                                   int controllerId,
+                                                   List topics,
+                                                   int clusterAuthorizedOperations) {
+        return prepareResponse(hasReliableLeaderEpochs(version), throttleTimeMs, brokers, clusterId, controllerId,
+                topics, clusterAuthorizedOperations);
+    }
+
+    // Visible for testing
+    public static MetadataResponse prepareResponse(boolean hasReliableEpoch,
+                                                   int throttleTimeMs,
+                                                   Collection brokers,
+                                                   String clusterId,
+                                                   int controllerId,
+                                                   List topics,
+                                                   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);
+
+        topics.forEach(topicMetadata -> responseData.topics().add(topicMetadata));
+        return new MetadataResponse(responseData, hasReliableEpoch);
+    }
+
+    @Override
+    public boolean shouldClientThrottle(short version) {
+        return version >= 6;
     }
 }
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 4686c3bc2580f..9869da5d254b9 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
@@ -18,13 +18,14 @@
 
 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.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.Struct;
-import org.apache.kafka.common.utils.CollectionUtils;
 
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
@@ -32,312 +33,92 @@
 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.CommonFields.PARTITION_ID;
-import static org.apache.kafka.common.protocol.CommonFields.TOPIC_NAME;
-import static org.apache.kafka.common.protocol.types.Type.INT64;
-import static org.apache.kafka.common.protocol.types.Type.NULLABLE_STRING;
-
-/**
- * This wrapper supports both v0 and v1 of OffsetCommitRequest.
- */
 public class OffsetCommitRequest extends AbstractRequest {
-    private static final String TOPICS_KEY_NAME = "topics";
-    private static final String RETENTION_TIME_KEY_NAME = "retention_time";
-
-    // topic level field names
-    private static final String PARTITIONS_KEY_NAME = "partitions";
-
-    // partition level field names
-    private static final String COMMIT_OFFSET_KEY_NAME = "offset";
-    private static final String METADATA_KEY_NAME = "metadata";
-
-    @Deprecated
-    private static final String TIMESTAMP_KEY_NAME = "timestamp";         // for v0, v1
-
-    /* Offset commit api */
-    private static final Schema OFFSET_COMMIT_REQUEST_PARTITION_V0 = new Schema(
-            PARTITION_ID,
-            new Field(COMMIT_OFFSET_KEY_NAME, INT64, "Message offset to be committed."),
-            new Field(METADATA_KEY_NAME, NULLABLE_STRING, "Any associated metadata the client wants to keep."));
-
-    private static final Schema OFFSET_COMMIT_REQUEST_PARTITION_V1 = new Schema(
-            PARTITION_ID,
-            new Field(COMMIT_OFFSET_KEY_NAME, INT64, "Message offset to be committed."),
-            new Field(TIMESTAMP_KEY_NAME, INT64, "Timestamp of the commit"),
-            new Field(METADATA_KEY_NAME, NULLABLE_STRING, "Any associated metadata the client wants to keep."));
-
-    private static final Schema OFFSET_COMMIT_REQUEST_PARTITION_V2 = new Schema(
-            PARTITION_ID,
-            new Field(COMMIT_OFFSET_KEY_NAME, INT64, "Message offset to be committed."),
-            new Field(METADATA_KEY_NAME, NULLABLE_STRING, "Any associated metadata the client wants to keep."));
-
-    private static final Schema OFFSET_COMMIT_REQUEST_TOPIC_V0 = new Schema(
-            TOPIC_NAME,
-            new Field(PARTITIONS_KEY_NAME, new ArrayOf(OFFSET_COMMIT_REQUEST_PARTITION_V0), "Partitions to commit offsets."));
-
-    private static final Schema OFFSET_COMMIT_REQUEST_TOPIC_V1 = new Schema(
-            TOPIC_NAME,
-            new Field(PARTITIONS_KEY_NAME, new ArrayOf(OFFSET_COMMIT_REQUEST_PARTITION_V1), "Partitions to commit offsets."));
-
-    private static final Schema OFFSET_COMMIT_REQUEST_TOPIC_V2 = new Schema(
-            TOPIC_NAME,
-            new Field(PARTITIONS_KEY_NAME, new ArrayOf(OFFSET_COMMIT_REQUEST_PARTITION_V2), "Partitions to commit offsets."));
-
-    private static final Schema OFFSET_COMMIT_REQUEST_V0 = new Schema(
-            GROUP_ID,
-            new Field(TOPICS_KEY_NAME, new ArrayOf(OFFSET_COMMIT_REQUEST_TOPIC_V0), "Topics to commit offsets."));
-
-    private static final Schema OFFSET_COMMIT_REQUEST_V1 = new Schema(
-            GROUP_ID,
-            GENERATION_ID,
-            MEMBER_ID,
-            new Field(TOPICS_KEY_NAME, new ArrayOf(OFFSET_COMMIT_REQUEST_TOPIC_V1), "Topics to commit offsets."));
-
-    private static final Schema OFFSET_COMMIT_REQUEST_V2 = new Schema(
-            GROUP_ID,
-            GENERATION_ID,
-            MEMBER_ID,
-            new Field(RETENTION_TIME_KEY_NAME, INT64, "Time period in ms to retain the offset."),
-            new Field(TOPICS_KEY_NAME, new ArrayOf(OFFSET_COMMIT_REQUEST_TOPIC_V2), "Topics to commit offsets."));
-
-    /* v3 request is same as v2. Throttle time has been added to response */
-    private static final Schema OFFSET_COMMIT_REQUEST_V3 = OFFSET_COMMIT_REQUEST_V2;
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {OFFSET_COMMIT_REQUEST_V0, OFFSET_COMMIT_REQUEST_V1, OFFSET_COMMIT_REQUEST_V2,
-            OFFSET_COMMIT_REQUEST_V3};
-    }
-
     // default values for the current version
     public static final int DEFAULT_GENERATION_ID = -1;
     public static final String DEFAULT_MEMBER_ID = "";
     public static final long DEFAULT_RETENTION_TIME = -1L;
 
-    // default values for old versions,
-    // will be removed after these versions are deprecated
-    @Deprecated
+    // default values for old versions, will be removed after these versions are no longer supported
     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;
-
-        @Deprecated
-        public PartitionData(long offset, long timestamp, String metadata) {
-            this.offset = offset;
-            this.timestamp = timestamp;
-            this.metadata = metadata;
-        }
-
-        public PartitionData(long offset, String metadata) {
-            this(offset, DEFAULT_TIMESTAMP, metadata);
-        }
-
-        @Override
-        public String toString() {
-            StringBuilder bld = new StringBuilder();
-            bld.append("(timestamp=").append(timestamp).
-                append(", offset=").append(offset).
-                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;
-        private long retentionTime = DEFAULT_RETENTION_TIME;
-
-        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;
-        }
-
-        public Builder setGenerationId(int generationId) {
-            this.generationId = generationId;
-            return this;
-        }
+        private final OffsetCommitRequestData data;
 
-        public Builder setRetentionTime(long retentionTime) {
-            this.retentionTime = retentionTime;
-            return this;
+        public Builder(OffsetCommitRequestData data) {
+            super(ApiKeys.OFFSET_COMMIT);
+            this.data = data;
         }
 
         @Override
         public OffsetCommitRequest build(short version) {
-            switch (version) {
-                case 0:
-                    return new OffsetCommitRequest(groupId, DEFAULT_GENERATION_ID, DEFAULT_MEMBER_ID,
-                            DEFAULT_RETENTION_TIME, offsetData, version);
-                case 1:
-                case 2:
-                case 3:
-                    long retentionTime = version == 1 ? DEFAULT_RETENTION_TIME : this.retentionTime;
-                    return new OffsetCommitRequest(groupId, generationId, memberId, retentionTime, offsetData, version);
-                default:
-                    throw new UnsupportedVersionException("Unsupported version " + 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(", retentionTime=").append(retentionTime).
-                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) {
-        super(version);
-        this.groupId = groupId;
-        this.generationId = generationId;
-        this.memberId = memberId;
-        this.retentionTime = retentionTime;
-        this.offsetData = offsetData;
+    public OffsetCommitRequest(OffsetCommitRequestData data, short version) {
+        super(ApiKeys.OFFSET_COMMIT, version);
+        this.data = data;
     }
 
-    public OffsetCommitRequest(Struct struct, short versionId) {
-        super(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);
-
-        // This field only exists in v2
-        if (struct.hasField(RETENTION_TIME_KEY_NAME))
-            retentionTime = struct.getLong(RETENTION_TIME_KEY_NAME);
-        else
-            retentionTime = DEFAULT_RETENTION_TIME;
+    @Override
+    public OffsetCommitRequestData data() {
+        return data;
+    }
 
-        offsetData = new HashMap<>();
-        for (Object topicDataObj : struct.getArray(TOPICS_KEY_NAME)) {
-            Struct topicData = (Struct) topicDataObj;
-            String topic = topicData.get(TOPIC_NAME);
-            for (Object partitionDataObj : topicData.getArray(PARTITIONS_KEY_NAME)) {
-                Struct partitionDataStruct = (Struct) partitionDataObj;
-                int partition = partitionDataStruct.get(PARTITION_ID);
-                long offset = partitionDataStruct.getLong(COMMIT_OFFSET_KEY_NAME);
-                String metadata = partitionDataStruct.getString(METADATA_KEY_NAME);
-                PartitionData partitionOffset;
-                // This field only exists in v1
-                if (partitionDataStruct.hasField(TIMESTAMP_KEY_NAME)) {
-                    long timestamp = partitionDataStruct.getLong(TIMESTAMP_KEY_NAME);
-                    partitionOffset = new PartitionData(offset, timestamp, metadata);
-                } else {
-                    partitionOffset = new PartitionData(offset, 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.groupDataByTopic(offsetData);
-        List topicArray = new ArrayList<>();
-        for (Map.Entry> topicEntry: topicsData.entrySet()) {
-            Struct topicData = struct.instance(TOPICS_KEY_NAME);
-            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_KEY_NAME);
-                partitionData.set(PARTITION_ID, partitionEntry.getKey());
-                partitionData.set(COMMIT_OFFSET_KEY_NAME, fetchPartitionData.offset);
-                // Only for v1
-                if (partitionData.hasField(TIMESTAMP_KEY_NAME))
-                    partitionData.set(TIMESTAMP_KEY_NAME, fetchPartitionData.timestamp);
-                partitionData.set(METADATA_KEY_NAME, 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_KEY_NAME, partitionArray.toArray());
-            topicArray.add(topicData);
+            responseTopicData.add(new OffsetCommitResponseTopic()
+                    .setName(entry.name())
+                    .setPartitions(responsePartitions)
+            );
         }
-        struct.set(TOPICS_KEY_NAME, topicArray.toArray());
-        struct.setIfExists(GENERATION_ID, generationId);
-        struct.setIfExists(MEMBER_ID, memberId);
-        if (struct.hasField(RETENTION_TIME_KEY_NAME))
-            struct.set(RETENTION_TIME_KEY_NAME, 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));
-        }
-
-        short versionId = version();
-        switch (versionId) {
-            case 0:
-            case 1:
-            case 2:
-                return new OffsetCommitResponse(responseData);
-            case 3:
-                return new OffsetCommitResponse(throttleTimeMs, 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.OFFSET_COMMIT.latestVersion()));
-        }
-    }
-
-    public String groupId() {
-        return groupId;
-    }
-
-    public int generationId() {
-        return generationId;
-    }
-
-    public String memberId() {
-        return memberId;
-    }
-
-    public long retentionTime() {
-        return retentionTime;
-    }
-
-    public Map offsetData() {
-        return offsetData;
+    public OffsetCommitResponse getErrorResponse(int throttleTimeMs, Throwable e) {
+        List
+                responseTopicData = getErrorResponseTopics(data.topics(), Errors.forException(e));
+        return new OffsetCommitResponse(new OffsetCommitResponseData()
+                .setTopics(responseTopicData)
+                .setThrottleTimeMs(throttleTimeMs));
     }
 
     public static OffsetCommitRequest parse(ByteBuffer buffer, short version) {
-        Schema schema = ApiKeys.OFFSET_COMMIT.requestSchema(version);
-        return new OffsetCommitRequest(schema.read(buffer), version);
+        return new OffsetCommitRequest(new OffsetCommitRequestData(new ByteBufferAccessor(buffer), version), 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 b439034e1af41..2ed0e312983ce 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,141 +17,98 @@
 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.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.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 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:
+ *
+ *   - {@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 String RESPONSES_KEY_NAME = "responses";
-
-    // topic level fields
-    private static final String PARTITIONS_KEY_NAME = "partition_responses";
-
-    /**
-     * 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)
-     */
-
-    private static final Schema OFFSET_COMMIT_RESPONSE_PARTITION_V0 = new Schema(
-            PARTITION_ID,
-            ERROR_CODE);
-
-    private static final Schema OFFSET_COMMIT_RESPONSE_TOPIC_V0 = new Schema(
-            TOPIC_NAME,
-            new Field(PARTITIONS_KEY_NAME, new ArrayOf(OFFSET_COMMIT_RESPONSE_PARTITION_V0)));
-
-    private static final Schema OFFSET_COMMIT_RESPONSE_V0 = new Schema(
-            new Field(RESPONSES_KEY_NAME, new ArrayOf(OFFSET_COMMIT_RESPONSE_TOPIC_V0)));
-
-
-    /* The response types for V0, V1 and V2 of OFFSET_COMMIT_REQUEST are the same. */
-    private static final Schema OFFSET_COMMIT_RESPONSE_V1 = OFFSET_COMMIT_RESPONSE_V0;
-    private static final Schema OFFSET_COMMIT_RESPONSE_V2 = OFFSET_COMMIT_RESPONSE_V0;
-
-    private static final Schema OFFSET_COMMIT_RESPONSE_V3 = new Schema(
-            THROTTLE_TIME_MS,
-            new Field(RESPONSES_KEY_NAME, new ArrayOf(OFFSET_COMMIT_RESPONSE_TOPIC_V0)));
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {OFFSET_COMMIT_RESPONSE_V0, OFFSET_COMMIT_RESPONSE_V1, OFFSET_COMMIT_RESPONSE_V2,
-            OFFSET_COMMIT_RESPONSE_V3};
-    }
+    private final OffsetCommitResponseData data;
 
-    private final Map responseData;
-    private final int throttleTimeMs;
-
-    public OffsetCommitResponse(Map responseData) {
-        this(DEFAULT_THROTTLE_TIME, responseData);
+    public OffsetCommitResponse(OffsetCommitResponseData data) {
+        super(ApiKeys.OFFSET_COMMIT);
+        this.data = data;
     }
 
-    public OffsetCommitResponse(int throttleTimeMs, Map responseData) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.responseData = responseData;
-    }
+    public OffsetCommitResponse(int requestThrottleMs, Map responseData) {
+        super(ApiKeys.OFFSET_COMMIT);
+        Map
+                responseTopicDataMap = new HashMap<>();
 
-    public OffsetCommitResponse(Struct struct) {
-        this.throttleTimeMs = struct.getOrElse(THROTTLE_TIME_MS, DEFAULT_THROTTLE_TIME);
-        responseData = new HashMap<>();
-        for (Object topicResponseObj : struct.getArray(RESPONSES_KEY_NAME)) {
-            Struct topicResponse = (Struct) topicResponseObj;
-            String topic = topicResponse.get(TOPIC_NAME);
-            for (Object partitionResponseObj : topicResponse.getArray(PARTITIONS_KEY_NAME)) {
-                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);
-            }
-        }
-    }
+        for (Map.Entry entry : responseData.entrySet()) {
+            TopicPartition topicPartition = entry.getKey();
+            String topicName = topicPartition.topic();
 
-    @Override
-    public Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.OFFSET_COMMIT.responseSchema(version));
-        struct.setIfExists(THROTTLE_TIME_MS, throttleTimeMs);
-
-        Map> topicsData = CollectionUtils.groupDataByTopic(responseData);
-        List topicArray = new ArrayList<>();
-        for (Map.Entry> entries: topicsData.entrySet()) {
-            Struct topicData = struct.instance(RESPONSES_KEY_NAME);
-            topicData.set(TOPIC_NAME, entries.getKey());
-            List partitionArray = new ArrayList<>();
-            for (Map.Entry partitionEntry : entries.getValue().entrySet()) {
-                Struct partitionData = topicData.instance(PARTITIONS_KEY_NAME);
-                partitionData.set(PARTITION_ID, partitionEntry.getKey());
-                partitionData.set(ERROR_CODE, partitionEntry.getValue().code());
-                partitionArray.add(partitionData);
-            }
-            topicData.set(PARTITIONS_KEY_NAME, partitionArray.toArray());
-            topicArray.add(topicData);
+            OffsetCommitResponseTopic topic = responseTopicDataMap.getOrDefault(
+                topicName, new OffsetCommitResponseTopic().setName(topicName));
+
+            topic.partitions().add(new OffsetCommitResponsePartition()
+                                       .setErrorCode(entry.getValue().code())
+                                       .setPartitionIndex(topicPartition.partition()));
+            responseTopicDataMap.put(topicName, topic);
         }
-        struct.set(RESPONSES_KEY_NAME, topicArray.toArray());
 
-        return struct;
+        data = new OffsetCommitResponseData()
+                .setTopics(new ArrayList<>(responseTopicDataMap.values()))
+                .setThrottleTimeMs(requestThrottleMs);
     }
 
-    public int throttleTimeMs() {
-        return throttleTimeMs;
+    public OffsetCommitResponse(Map responseData) {
+        this(DEFAULT_THROTTLE_TIME, responseData);
     }
 
-    public Map responseData() {
-        return responseData;
+    @Override
+    public OffsetCommitResponseData data() {
+        return data;
     }
 
     @Override
     public Map errorCounts() {
-        return errorCounts(responseData);
+        return errorCounts(data.topics().stream().flatMap(topicResult ->
+                topicResult.partitions().stream().map(partitionResult ->
+                        Errors.forCode(partitionResult.errorCode()))));
     }
 
     public static OffsetCommitResponse parse(ByteBuffer buffer, short version) {
-        return new OffsetCommitResponse(ApiKeys.OFFSET_COMMIT.parseResponse(version, buffer));
+        return new OffsetCommitResponse(new OffsetCommitResponseData(new ByteBufferAccessor(buffer), version));
+    }
+
+    @Override
+    public String toString() {
+        return data.toString();
     }
 
+    @Override
+    public int throttleTimeMs() {
+        return data.throttleTimeMs();
+    }
+
+    @Override
+    public boolean shouldClientThrottle(short version) {
+        return version >= 4;
+    }
 }
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..28b763d520f00
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/OffsetDeleteRequest.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.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.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.Errors;
+
+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();
+        }
+    }
+
+    private final OffsetDeleteRequestData data;
+
+    public OffsetDeleteRequest(OffsetDeleteRequestData data, short version) {
+        super(ApiKeys.OFFSET_DELETE, version);
+        this.data = data;
+    }
+
+    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(new OffsetDeleteRequestData(new ByteBufferAccessor(buffer), version), version);
+    }
+
+    @Override
+    public OffsetDeleteRequestData data() {
+        return data;
+    }
+}
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..79f6f4e6d3495
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/OffsetDeleteResponse.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.requests;
+
+import org.apache.kafka.common.message.OffsetDeleteResponseData;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.Errors;
+
+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 {
+
+    private final OffsetDeleteResponseData data;
+
+    public OffsetDeleteResponse(OffsetDeleteResponseData data) {
+        super(ApiKeys.OFFSET_DELETE);
+        this.data = data;
+    }
+
+    @Override
+    public OffsetDeleteResponseData data() {
+        return data;
+    }
+
+    @Override
+    public Map errorCounts() {
+        Map counts = new HashMap<>();
+        updateErrorCounts(counts, Errors.forCode(data.errorCode()));
+        data.topics().forEach(topic ->
+            topic.partitions().forEach(partition ->
+                updateErrorCounts(counts, Errors.forCode(partition.errorCode()))
+            )
+        );
+        return counts;
+    }
+
+    public static OffsetDeleteResponse parse(ByteBuffer buffer, short version) {
+        return new OffsetDeleteResponse(new OffsetDeleteResponseData(new ByteBufferAccessor(buffer), version));
+    }
+
+    @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 0db1c50ca57b5..c35d479d64dba 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,141 +18,119 @@
 
 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.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.Struct;
-import org.apache.kafka.common.utils.CollectionUtils;
-import org.apache.kafka.common.utils.Utils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 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.GROUP_ID;
-import static org.apache.kafka.common.protocol.CommonFields.PARTITION_ID;
-import static org.apache.kafka.common.protocol.CommonFields.TOPIC_NAME;
+import java.util.Optional;
 
 public class OffsetFetchRequest extends AbstractRequest {
-    private static final String TOPICS_KEY_NAME = "topics";
-
-    // topic level field names
-    private static final String PARTITIONS_KEY_NAME = "partitions";
-
-    /*
-     * 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 Schema OFFSET_FETCH_REQUEST_PARTITION_V0 = new Schema(
-            PARTITION_ID);
-
-    private static final Schema OFFSET_FETCH_REQUEST_TOPIC_V0 = new Schema(
-            TOPIC_NAME,
-            new Field(PARTITIONS_KEY_NAME, new ArrayOf(OFFSET_FETCH_REQUEST_PARTITION_V0), "Partitions to fetch offsets."));
-
-    private static final Schema OFFSET_FETCH_REQUEST_V0 = new Schema(
-            GROUP_ID,
-            new Field(TOPICS_KEY_NAME, new ArrayOf(OFFSET_FETCH_REQUEST_TOPIC_V0), "Topics to fetch offsets."));
-
-    private static final Schema OFFSET_FETCH_REQUEST_V1 = OFFSET_FETCH_REQUEST_V0;
-
-    private static final Schema OFFSET_FETCH_REQUEST_V2 = new Schema(
-            GROUP_ID,
-            new Field(TOPICS_KEY_NAME, ArrayOf.nullable(OFFSET_FETCH_REQUEST_TOPIC_V0), "Topics to fetch offsets. If the " +
-                    "topic array is null fetch offsets for all topics."));
-
-    /* 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;
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {OFFSET_FETCH_REQUEST_V0, OFFSET_FETCH_REQUEST_V1, OFFSET_FETCH_REQUEST_V2,
-            OFFSET_FETCH_REQUEST_V3};
-    }
+
+    private static final Logger log = LoggerFactory.getLogger(OffsetFetchRequest.class);
+
+    private static final List ALL_TOPIC_PARTITIONS = null;
+    private 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 Builder(String groupId, List partitions) {
+        public final OffsetFetchRequestData data;
+        private final boolean throwOnFetchStableOffsetsUnsupported;
+
+        public Builder(String groupId,
+                       boolean requireStable,
+                       List partitions,
+                       boolean throwOnFetchStableOffsetsUnsupported) {
             super(ApiKeys.OFFSET_FETCH);
-            this.groupId = groupId;
-            this.partitions = partitions;
-        }
 
-        public static Builder allTopicPartitions(String groupId) {
-            return new Builder(groupId, ALL_TOPIC_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)
+                            .setRequireStable(requireStable)
+                            .setTopics(topics);
+            this.throwOnFetchStableOffsetsUnsupported = throwOnFetchStableOffsetsUnsupported;
         }
 
-        public boolean isAllTopicPartitions() {
-            return this.partitions == ALL_TOPIC_PARTITIONS;
+        boolean isAllTopicPartitions() {
+            return this.data.topics() == ALL_TOPIC_PARTITIONS;
         }
 
         @Override
         public OffsetFetchRequest build(short version) {
-            if (isAllTopicPartitions() && version < 2)
+            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);
+                    "v" + version + ", but we need v2 or newer to request all topic partitions.");
+            }
+
+            if (data.requireStable() && version < 7) {
+                if (throwOnFetchStableOffsetsUnsupported) {
+                    throw new UnsupportedVersionException("Broker unexpectedly " +
+                        "doesn't support requireStable flag on version " + version);
+                } else {
+                    log.trace("Fallback the requireStable flag to false as broker " +
+                                  "only supports OffsetFetchRequest version {}. Need " +
+                                  "v7 or newer to enable this feature", version);
+
+                    return new OffsetFetchRequest(data.setRequireStable(false), 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 static OffsetFetchRequest forAllPartitions(String groupId) {
-        return new OffsetFetchRequest.Builder(groupId, null).build((short) 2);
+    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;
     }
 
-    // v0, v1, and v2 have the same fields.
-    private OffsetFetchRequest(String groupId, List partitions, short version) {
-        super(version);
-        this.groupId = groupId;
-        this.partitions = partitions;
+    public String groupId() {
+        return data.groupId();
     }
 
-    public OffsetFetchRequest(Struct struct, short version) {
-        super(version);
-
-        Object[] topicArray = struct.getArray(TOPICS_KEY_NAME);
-        if (topicArray != null) {
-            partitions = new ArrayList<>();
-            for (Object topicResponseObj : struct.getArray(TOPICS_KEY_NAME)) {
-                Struct topicResponse = (Struct) topicResponseObj;
-                String topic = topicResponse.get(TOPIC_NAME);
-                for (Object partitionResponseObj : topicResponse.getArray(PARTITIONS_KEY_NAME)) {
-                    Struct partitionResponse = (Struct) partitionResponseObj;
-                    int partition = partitionResponse.get(PARTITION_ID);
-                    partitions.add(new TopicPartition(topic, partition));
-                }
-            }
-        } else
-            partitions = null;
-
+    public boolean requireStable() {
+        return data.requireStable();
+    }
 
-        groupId = struct.get(GROUP_ID);
+    private OffsetFetchRequest(OffsetFetchRequestData data, short version) {
+        super(ApiKeys.OFFSET_FETCH, version);
+        this.data = data;
     }
 
     public OffsetFetchResponse getErrorResponse(Errors error) {
@@ -160,28 +138,26 @@ public OffsetFetchResponse getErrorResponse(Errors error) {
     }
 
     public OffsetFetchResponse getErrorResponse(int throttleTimeMs, Errors error) {
-        short versionId = version();
-
         Map responsePartitions = new HashMap<>();
-        if (versionId < 2) {
-            for (TopicPartition partition : this.partitions) {
-                responsePartitions.put(partition, new OffsetFetchResponse.PartitionData(
-                        OffsetFetchResponse.INVALID_OFFSET,
-                        OffsetFetchResponse.NO_METADATA,
-                        error));
+        if (version() < 2) {
+            OffsetFetchResponse.PartitionData partitionError = new OffsetFetchResponse.PartitionData(
+                    OffsetFetchResponse.INVALID_OFFSET,
+                    Optional.empty(),
+                    OffsetFetchResponse.NO_METADATA,
+                    error);
+
+            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:
-                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);
         }
     }
 
@@ -190,46 +166,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);
+        return new OffsetFetchRequest(new OffsetFetchRequestData(new ByteBufferAccessor(buffer), version), 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.groupDataByTopic(partitions);
-
-            List topicArray = new ArrayList<>();
-            for (Map.Entry> entries : topicsData.entrySet()) {
-                Struct topicData = struct.instance(TOPICS_KEY_NAME);
-                topicData.set(TOPIC_NAME, entries.getKey());
-                List partitionArray = new ArrayList<>();
-                for (Integer partitionId : entries.getValue()) {
-                    Struct partitionData = topicData.instance(PARTITIONS_KEY_NAME);
-                    partitionData.set(PARTITION_ID, partitionId);
-                    partitionArray.add(partitionData);
-                }
-                topicData.set(PARTITIONS_KEY_NAME, partitionArray.toArray());
-                topicArray.add(topicData);
-            }
-            struct.set(TOPICS_KEY_NAME, topicArray.toArray());
-        } else
-            struct.set(TOPICS_KEY_NAME, null);
-
-        return struct;
+    public OffsetFetchRequestData data() {
+        return data;
     }
 }
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 e398442e847dc..594eb0e5bcf93 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,103 +17,67 @@
 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.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.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.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.protocol.types.Type.INT64;
-import static org.apache.kafka.common.protocol.types.Type.NULLABLE_STRING;
+import static org.apache.kafka.common.record.RecordBatch.NO_PARTITION_LEADER_EPOCH;
 
+/**
+ * Possible error codes:
+ *
+ * - Partition errors:
+ *   - {@link Errors#UNKNOWN_TOPIC_OR_PARTITION}
+ *   - {@link Errors#TOPIC_AUTHORIZATION_FAILED}
+ *   - {@link Errors#UNSTABLE_OFFSET_COMMIT}
+ *
+ * - Group or coordinator errors:
+ *   - {@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 String RESPONSES_KEY_NAME = "responses";
-
-    // topic level fields
-    private static final String PARTITIONS_KEY_NAME = "partition_responses";
-
-    // partition level fields
-    private static final String COMMIT_OFFSET_KEY_NAME = "offset";
-    private static final String METADATA_KEY_NAME = "metadata";
-
-    private static final Schema OFFSET_FETCH_RESPONSE_PARTITION_V0 = new Schema(
-            PARTITION_ID,
-            new Field(COMMIT_OFFSET_KEY_NAME, INT64, "Last committed message offset."),
-            new Field(METADATA_KEY_NAME, NULLABLE_STRING, "Any associated metadata the client wants to keep."),
-            ERROR_CODE);
-
-    private static final Schema OFFSET_FETCH_RESPONSE_TOPIC_V0 = new Schema(
-            TOPIC_NAME,
-            new Field(PARTITIONS_KEY_NAME, new ArrayOf(OFFSET_FETCH_RESPONSE_PARTITION_V0)));
-
-    private static final Schema OFFSET_FETCH_RESPONSE_V0 = new Schema(
-            new Field(RESPONSES_KEY_NAME, new ArrayOf(OFFSET_FETCH_RESPONSE_TOPIC_V0)));
-
-    private static final Schema OFFSET_FETCH_RESPONSE_V1 = OFFSET_FETCH_RESPONSE_V0;
-
-    private static final Schema OFFSET_FETCH_RESPONSE_V2 = new Schema(
-            new Field(RESPONSES_KEY_NAME, new ArrayOf(OFFSET_FETCH_RESPONSE_TOPIC_V0)),
-            ERROR_CODE);
-
-    /* v3 request is the same as v2. Throttle time has been added to v3 response */
-    private static final Schema OFFSET_FETCH_RESPONSE_V3 = new Schema(
-            THROTTLE_TIME_MS,
-            new Field(RESPONSES_KEY_NAME, new ArrayOf(OFFSET_FETCH_RESPONSE_TOPIC_V0)),
-            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};
-    }
-
     public static final long INVALID_OFFSET = -1L;
     public static final String NO_METADATA = "";
-    public static final PartitionData UNKNOWN_PARTITION = new PartitionData(INVALID_OFFSET, NO_METADATA,
-            Errors.UNKNOWN_TOPIC_OR_PARTITION);
-    public static final PartitionData UNAUTHORIZED_PARTITION = new PartitionData(INVALID_OFFSET, NO_METADATA,
-            Errors.TOPIC_AUTHORIZATION_FAILED);
-
-    /**
-     * Possible error codes:
-     *
-     * - Partition errors:
-     *   - UNKNOWN_TOPIC_OR_PARTITION (3)
-     *
-     * - Group or coordinator errors:
-     *   - COORDINATOR_LOAD_IN_PROGRESS (14)
-     *   - COORDINATOR_NOT_AVAILABLE (15)
-     *   - NOT_COORDINATOR (16)
-     *   - GROUP_AUTHORIZATION_FAILED (30)
-     */
-
-    private static final List PARTITION_ERRORS = Collections.singletonList(Errors.UNKNOWN_TOPIC_OR_PARTITION);
-
-    private final Map responseData;
+    public static final PartitionData UNKNOWN_PARTITION = new PartitionData(INVALID_OFFSET,
+                                                                            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);
+    private static final List PARTITION_ERRORS = Arrays.asList(
+        Errors.UNKNOWN_TOPIC_OR_PARTITION, Errors.TOPIC_AUTHORIZATION_FAILED);
+
+    private final OffsetFetchResponseData data;
     private final Errors error;
-    private final int throttleTimeMs;
 
     public static final class PartitionData {
         public final long offset;
         public final String metadata;
         public final Errors error;
+        public final Optional leaderEpoch;
 
-        public PartitionData(long offset, String metadata, Errors error) {
+        public PartitionData(long offset,
+                             Optional leaderEpoch,
+                             String metadata,
+                             Errors error) {
             this.offset = offset;
+            this.leaderEpoch = leaderEpoch;
             this.metadata = metadata;
             this.error = error;
         }
@@ -121,6 +85,32 @@ public PartitionData(long offset, String metadata, Errors error) {
         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 Objects.equals(this.offset, otherPartition.offset)
+                   && Objects.equals(this.leaderEpoch, otherPartition.leaderEpoch)
+                   && Objects.equals(this.metadata, otherPartition.metadata)
+                   && Objects.equals(this.error, 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);
+        }
     }
 
     /**
@@ -139,99 +129,102 @@ 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;
+        super(ApiKeys.OFFSET_FETCH);
+        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.getArray(RESPONSES_KEY_NAME)) {
-            Struct topicResponse = (Struct) topicResponseObj;
-            String topic = topicResponse.get(TOPIC_NAME);
-            for (Object partitionResponseObj : topicResponse.getArray(PARTITIONS_KEY_NAME)) {
-                Struct partitionResponse = (Struct) partitionResponseObj;
-                int partition = partitionResponse.get(PARTITION_ID);
-                long offset = partitionResponse.getLong(COMMIT_OFFSET_KEY_NAME);
-                String metadata = partitionResponse.getString(METADATA_KEY_NAME);
-                Errors error = Errors.forCode(partitionResponse.get(ERROR_CODE));
-                if (error != Errors.NONE && !PARTITION_ERRORS.contains(error))
-                    topLevelError = error;
-                PartitionData partitionData = new PartitionData(offset, metadata, error);
-                this.responseData.put(new TopicPartition(topic, partition), partitionData);
-            }
-        }
-
+    public OffsetFetchResponse(OffsetFetchResponseData data, short version) {
+        super(ApiKeys.OFFSET_FETCH);
+        this.data = data;
         // for version 2 and later use the top-level error code (in ERROR_CODE_KEY_NAME) from the response.
         // 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;
+        this.error = version >= 2 ? Errors.forCode(data.errorCode()) : topLevelError(data);
     }
 
-    public void maybeThrowFirstPartitionError() {
-        Collection partitionsData = this.responseData.values();
-        for (PartitionData data : partitionsData) {
-            if (data.hasError())
-                throw data.error.exception();
+    private static Errors topLevelError(OffsetFetchResponseData data) {
+        for (OffsetFetchResponseTopic topic : data.topics()) {
+            for (OffsetFetchResponsePartition partition : topic.partitions()) {
+                Errors partitionError = Errors.forCode(partition.errorCode());
+                if (partitionError != Errors.NONE && !PARTITION_ERRORS.contains(partitionError)) {
+                    return partitionError;
+                }
+            }
         }
+        return Errors.NONE;
     }
 
+    @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
     public Map errorCounts() {
-        return errorCounts(error);
+        Map counts = new HashMap<>();
+        updateErrorCounts(counts, error);
+        data.topics().forEach(topic ->
+                topic.partitions().forEach(partition ->
+                        updateErrorCounts(counts, Errors.forCode(partition.errorCode()))));
+        return counts;
     }
 
     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(new OffsetFetchResponseData(new ByteBufferAccessor(buffer), version), 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.groupDataByTopic(responseData);
-        List topicArray = new ArrayList<>();
-        for (Map.Entry> entries : topicsData.entrySet()) {
-            Struct topicData = struct.instance(RESPONSES_KEY_NAME);
-            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_KEY_NAME);
-                partitionData.set(PARTITION_ID, partitionEntry.getKey());
-                partitionData.set(COMMIT_OFFSET_KEY_NAME, fetchPartitionData.offset);
-                partitionData.set(METADATA_KEY_NAME, fetchPartitionData.metadata);
-                partitionData.set(ERROR_CODE, fetchPartitionData.error.code());
-                partitionArray.add(partitionData);
-            }
-            topicData.set(PARTITIONS_KEY_NAME, partitionArray.toArray());
-            topicArray.add(topicData);
-        }
-        struct.set(RESPONSES_KEY_NAME, topicArray.toArray());
-
-        if (version > 1)
-            struct.set(ERROR_CODE, this.error.code());
+    public OffsetFetchResponseData data() {
+        return data;
+    }
 
-        return struct;
+    @Override
+    public boolean shouldClientThrottle(short version) {
+        return version >= 4;
     }
 }
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 d0585bed6d552..727c7087c439a 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,141 +17,156 @@
 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.OffsetForLeaderEpochRequestData;
+import org.apache.kafka.common.message.OffsetForLeaderEpochRequestData.OffsetForLeaderPartition;
+import org.apache.kafka.common.message.OffsetForLeaderEpochRequestData.OffsetForLeaderTopic;
+import org.apache.kafka.common.message.OffsetForLeaderEpochRequestData.OffsetForLeaderTopicCollection;
+import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData;
+import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData.EpochEndOffset;
+import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData.OffsetForLeaderTopicResult;
 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.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.Optional;
 
-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 org.apache.kafka.common.requests.OffsetsForLeaderEpochResponse.UNDEFINED_EPOCH;
+import static org.apache.kafka.common.requests.OffsetsForLeaderEpochResponse.UNDEFINED_EPOCH_OFFSET;
 
 public class OffsetsForLeaderEpochRequest extends AbstractRequest {
-    private static final String TOPICS_KEY_NAME = "topics";
-    private static final String PARTITIONS_KEY_NAME = "partitions";
-    private static final String LEADER_EPOCH = "leader_epoch";
-
-    /* Offsets for Leader Epoch api */
-    private static final Schema OFFSET_FOR_LEADER_EPOCH_REQUEST_PARTITION_V0 = new Schema(
-            PARTITION_ID,
-            new Field(LEADER_EPOCH, INT32, "The epoch"));
-    private static final Schema OFFSET_FOR_LEADER_EPOCH_REQUEST_TOPIC_V0 = new Schema(
-            TOPIC_NAME,
-            new Field(PARTITIONS_KEY_NAME, new ArrayOf(OFFSET_FOR_LEADER_EPOCH_REQUEST_PARTITION_V0)));
-    private static final Schema OFFSET_FOR_LEADER_EPOCH_REQUEST_V0 = new Schema(
-            new Field(TOPICS_KEY_NAME, new ArrayOf(OFFSET_FOR_LEADER_EPOCH_REQUEST_TOPIC_V0), "An array of topics to get epochs for"));
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{OFFSET_FOR_LEADER_EPOCH_REQUEST_V0};
-    }
 
-    private Map epochsByPartition;
+    /**
+     * Sentinel replica_id value to indicate a regular consumer rather than another broker
+     */
+    public static final int CONSUMER_REPLICA_ID = -1;
 
-    public Map epochsByTopicPartition() {
-        return epochsByPartition;
-    }
+    /**
+     * 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 final OffsetForLeaderEpochRequestData data;
 
     public static class Builder extends AbstractRequest.Builder {
-        private Map epochsByPartition = new HashMap<>();
+        private final OffsetForLeaderEpochRequestData data;
 
-        public Builder() {
-            super(ApiKeys.OFFSET_FOR_LEADER_EPOCH);
+        Builder(short oldestAllowedVersion, short latestAllowedVersion, OffsetForLeaderEpochRequestData data) {
+            super(ApiKeys.OFFSET_FOR_LEADER_EPOCH, oldestAllowedVersion, latestAllowedVersion);
+            this.data = data;
         }
 
-        public Builder(Map epochsByPartition) {
-            super(ApiKeys.OFFSET_FOR_LEADER_EPOCH);
-            this.epochsByPartition = epochsByPartition;
+        public static Builder forConsumer(OffsetForLeaderTopicCollection 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.
+            OffsetForLeaderEpochRequestData data = new OffsetForLeaderEpochRequestData();
+            data.setReplicaId(CONSUMER_REPLICA_ID);
+            data.setTopics(epochsByPartition);
+            return new Builder((short) 3, ApiKeys.OFFSET_FOR_LEADER_EPOCH.latestVersion(), data);
         }
 
-        public Builder add(TopicPartition topicPartition, Integer epoch) {
-            epochsByPartition.put(topicPartition, epoch);
-            return this;
+        public static Builder forFollower(short version, Map epochsByPartition, int replicaId) {
+            OffsetForLeaderEpochRequestData data = new OffsetForLeaderEpochRequestData();
+            data.setReplicaId(replicaId);
+
+            epochsByPartition.forEach((partitionKey, partitionValue) -> {
+                OffsetForLeaderTopic topic = data.topics().find(partitionKey.topic());
+                if (topic == null) {
+                    topic = new OffsetForLeaderTopic().setTopic(partitionKey.topic());
+                    data.topics().add(topic);
+                }
+                topic.partitions().add(new OffsetForLeaderPartition()
+                    .setPartition(partitionKey.partition())
+                    .setLeaderEpoch(partitionValue.leaderEpoch)
+                    .setCurrentLeaderEpoch(partitionValue.currentLeaderEpoch
+                        .orElse(RecordBatch.NO_PARTITION_LEADER_EPOCH))
+                );
+            });
+            return new Builder(version, version, data);
         }
 
         @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);
 
-        public static OffsetsForLeaderEpochRequest parse(ByteBuffer buffer, short version) {
-            return new OffsetsForLeaderEpochRequest(ApiKeys.OFFSET_FOR_LEADER_EPOCH.parseRequest(version, buffer), version);
+            return new OffsetsForLeaderEpochRequest(data, version);
         }
 
         @Override
         public String toString() {
-            StringBuilder bld = new StringBuilder();
-            bld.append("(type=OffsetForLeaderEpochRequest, ").
-                    append("epochsByTopic=").append(epochsByPartition).
-                    append(")");
-            return bld.toString();
+            return data.toString();
         }
     }
 
-    public OffsetsForLeaderEpochRequest(Map epochsByPartition, short version) {
-        super(version);
-        this.epochsByPartition = epochsByPartition;
+    public OffsetsForLeaderEpochRequest(OffsetForLeaderEpochRequestData data, short version) {
+        super(ApiKeys.OFFSET_FOR_LEADER_EPOCH, version);
+        this.data = data;
     }
 
-    public OffsetsForLeaderEpochRequest(Struct struct, short version) {
-        super(version);
-        epochsByPartition = new HashMap<>();
-        for (Object topicAndEpochsObj : struct.getArray(TOPICS_KEY_NAME)) {
-            Struct topicAndEpochs = (Struct) topicAndEpochsObj;
-            String topic = topicAndEpochs.get(TOPIC_NAME);
-            for (Object partitionAndEpochObj : topicAndEpochs.getArray(PARTITIONS_KEY_NAME)) {
-                Struct partitionAndEpoch = (Struct) partitionAndEpochObj;
-                int partitionId = partitionAndEpoch.get(PARTITION_ID);
-                int epoch = partitionAndEpoch.getInt(LEADER_EPOCH);
-                TopicPartition tp = new TopicPartition(topic, partitionId);
-                epochsByPartition.put(tp, epoch);
-            }
-        }
+    @Override
+    public OffsetForLeaderEpochRequestData data() {
+        return data;
     }
 
-    public static OffsetsForLeaderEpochRequest parse(ByteBuffer buffer, short versionId) {
-        return new OffsetsForLeaderEpochRequest(ApiKeys.OFFSET_FOR_LEADER_EPOCH.parseRequest(versionId, buffer), versionId);
+    public int replicaId() {
+        return data.replicaId();
     }
 
-    @Override
-    protected Struct toStruct() {
-        Struct requestStruct = new Struct(ApiKeys.OFFSET_FOR_LEADER_EPOCH.requestSchema(version()));
-
-        Map> topicsToPartitionEpochs = CollectionUtils.groupDataByTopic(epochsByPartition);
-
-        List topics = new ArrayList<>();
-        for (Map.Entry> topicToEpochs : topicsToPartitionEpochs.entrySet()) {
-            Struct topicsStruct = requestStruct.instance(TOPICS_KEY_NAME);
-            topicsStruct.set(TOPIC_NAME, topicToEpochs.getKey());
-            List partitions = new ArrayList<>();
-            for (Map.Entry partitionEpoch : topicToEpochs.getValue().entrySet()) {
-                Struct partitionStruct = topicsStruct.instance(PARTITIONS_KEY_NAME);
-                partitionStruct.set(PARTITION_ID, partitionEpoch.getKey());
-                partitionStruct.set(LEADER_EPOCH, partitionEpoch.getValue());
-                partitions.add(partitionStruct);
-            }
-            topicsStruct.set(PARTITIONS_KEY_NAME, partitions.toArray());
-            topics.add(topicsStruct);
-        }
-        requestStruct.set(TOPICS_KEY_NAME, topics.toArray());
-        return requestStruct;
+    public static OffsetsForLeaderEpochRequest parse(ByteBuffer buffer, short version) {
+        return new OffsetsForLeaderEpochRequest(new OffsetForLeaderEpochRequestData(new ByteBufferAccessor(buffer), version), version);
     }
 
     @Override
     public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
         Errors error = Errors.forException(e);
-        Map errorResponse = new HashMap<>();
-        for (TopicPartition tp : epochsByPartition.keySet()) {
-            errorResponse.put(tp, new EpochEndOffset(error, EpochEndOffset.UNDEFINED_EPOCH_OFFSET));
+
+        OffsetForLeaderEpochResponseData responseData = new OffsetForLeaderEpochResponseData();
+        data.topics().forEach(topic -> {
+            OffsetForLeaderTopicResult topicData = new OffsetForLeaderTopicResult()
+                .setTopic(topic.topic());
+            topic.partitions().forEach(partition ->
+                topicData.partitions().add(new EpochEndOffset()
+                    .setPartition(partition.partition())
+                    .setErrorCode(error.code())
+                    .setLeaderEpoch(UNDEFINED_EPOCH)
+                    .setEndOffset(UNDEFINED_EPOCH_OFFSET)));
+            responseData.topics().add(topicData);
+        });
+
+        return new OffsetsForLeaderEpochResponse(responseData);
+    }
+
+    public static class PartitionData {
+        public final Optional currentLeaderEpoch;
+        public final int leaderEpoch;
+
+        public PartitionData(Optional currentLeaderEpoch, int leaderEpoch) {
+            this.currentLeaderEpoch = currentLeaderEpoch;
+            this.leaderEpoch = leaderEpoch;
         }
-        return new OffsetsForLeaderEpochResponse(errorResponse);
+
+        @Override
+        public String toString() {
+            StringBuilder bld = new StringBuilder();
+            bld.append("(currentLeaderEpoch=").append(currentLeaderEpoch).
+                append(", leaderEpoch=").append(leaderEpoch).
+                append(")");
+            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 4a91533938d84..893d5a2af20a0 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
@@ -16,107 +16,64 @@
  */
 package org.apache.kafka.common.requests;
 
-import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData;
 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.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 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 static org.apache.kafka.common.protocol.types.Type.INT64;
+import static org.apache.kafka.common.record.RecordBatch.NO_PARTITION_LEADER_EPOCH;
 
+/**
+ * Possible error codes:
+ * - {@link Errors#TOPIC_AUTHORIZATION_FAILED} If the user does not have DESCRIBE access to a requested topic
+ * - {@link Errors#REPLICA_NOT_AVAILABLE} If the request is received by a broker with version < 2.6 which is not a replica
+ * - {@link Errors#NOT_LEADER_OR_FOLLOWER} If the broker is not a leader or follower and either the provided leader epoch
+ *     matches the known leader epoch on the broker or is empty
+ * - {@link Errors#FENCED_LEADER_EPOCH} If the epoch is lower than the broker's epoch
+ * - {@link Errors#UNKNOWN_LEADER_EPOCH} If the epoch is larger than the broker's epoch
+ * - {@link Errors#UNKNOWN_TOPIC_OR_PARTITION} If the broker does not have metadata for a topic or partition
+ * - {@link Errors#KAFKA_STORAGE_ERROR} If the log directory for one of the requested partitions is offline
+ * - {@link Errors#UNKNOWN_SERVER_ERROR} For any unexpected errors
+ */
 public class OffsetsForLeaderEpochResponse extends AbstractResponse {
-    private static final String TOPICS_KEY_NAME = "topics";
-    private static final String PARTITIONS_KEY_NAME = "partitions";
-    private static final String END_OFFSET_KEY_NAME = "end_offset";
-
-    private static final Schema OFFSET_FOR_LEADER_EPOCH_RESPONSE_PARTITION_V0 = new Schema(
-            ERROR_CODE,
-            PARTITION_ID,
-            new Field(END_OFFSET_KEY_NAME, INT64, "The end offset"));
-    private static final Schema OFFSET_FOR_LEADER_EPOCH_RESPONSE_TOPIC_V0 = new Schema(
-            TOPIC_NAME,
-            new Field(PARTITIONS_KEY_NAME, new ArrayOf(OFFSET_FOR_LEADER_EPOCH_RESPONSE_PARTITION_V0)));
-    private static final Schema OFFSET_FOR_LEADER_EPOCH_RESPONSE_V0 = new Schema(
-            new Field(TOPICS_KEY_NAME, new ArrayOf(OFFSET_FOR_LEADER_EPOCH_RESPONSE_TOPIC_V0),
-                    "An array of topics for which we have leader offsets for some requested Partition Leader Epoch"));
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{OFFSET_FOR_LEADER_EPOCH_RESPONSE_V0};
-    }
-
-    private Map epochEndOffsetsByPartition;
+    public static final long UNDEFINED_EPOCH_OFFSET = NO_PARTITION_LEADER_EPOCH;
+    public static final int UNDEFINED_EPOCH = NO_PARTITION_LEADER_EPOCH;
 
-    public OffsetsForLeaderEpochResponse(Struct struct) {
-        epochEndOffsetsByPartition = new HashMap<>();
-        for (Object topicAndEpocsObj : struct.getArray(TOPICS_KEY_NAME)) {
-            Struct topicAndEpochs = (Struct) topicAndEpocsObj;
-            String topic = topicAndEpochs.get(TOPIC_NAME);
-            for (Object partitionAndEpochObj : topicAndEpochs.getArray(PARTITIONS_KEY_NAME)) {
-                Struct partitionAndEpoch = (Struct) partitionAndEpochObj;
-                Errors error = Errors.forCode(partitionAndEpoch.get(ERROR_CODE));
-                int partitionId = partitionAndEpoch.get(PARTITION_ID);
-                TopicPartition tp = new TopicPartition(topic, partitionId);
-                long endOffset = partitionAndEpoch.getLong(END_OFFSET_KEY_NAME);
-                epochEndOffsetsByPartition.put(tp, new EpochEndOffset(error, endOffset));
-            }
-        }
-    }
+    private final OffsetForLeaderEpochResponseData data;
 
-    public OffsetsForLeaderEpochResponse(Map epochsByTopic) {
-        this.epochEndOffsetsByPartition = epochsByTopic;
+    public OffsetsForLeaderEpochResponse(OffsetForLeaderEpochResponseData data) {
+        super(ApiKeys.OFFSET_FOR_LEADER_EPOCH);
+        this.data = data;
     }
 
-    public Map responses() {
-        return epochEndOffsetsByPartition;
+    @Override
+    public OffsetForLeaderEpochResponseData data() {
+        return data;
     }
 
     @Override
     public Map errorCounts() {
         Map errorCounts = new HashMap<>();
-        for (EpochEndOffset response : epochEndOffsetsByPartition.values())
-            updateErrorCounts(errorCounts, response.error());
+        data.topics().forEach(topic ->
+            topic.partitions().forEach(partition ->
+                updateErrorCounts(errorCounts, Errors.forCode(partition.errorCode()))));
         return errorCounts;
     }
 
-    public static OffsetsForLeaderEpochResponse parse(ByteBuffer buffer, short versionId) {
-        return new OffsetsForLeaderEpochResponse(ApiKeys.OFFSET_FOR_LEADER_EPOCH.responseSchema(versionId).read(buffer));
+    public int throttleTimeMs() {
+        return data.throttleTimeMs();
     }
 
-    @Override
-    protected Struct toStruct(short version) {
-        Struct responseStruct = new Struct(ApiKeys.OFFSET_FOR_LEADER_EPOCH.responseSchema(version));
-
-        Map> endOffsetsByTopic = CollectionUtils.groupDataByTopic(epochEndOffsetsByPartition);
+    public static OffsetsForLeaderEpochResponse parse(ByteBuffer buffer, short version) {
+        return new OffsetsForLeaderEpochResponse(new OffsetForLeaderEpochResponseData(new ByteBufferAccessor(buffer), version));
+    }
 
-        List topics = new ArrayList<>(endOffsetsByTopic.size());
-        for (Map.Entry> topicToPartitionEpochs : endOffsetsByTopic.entrySet()) {
-            Struct topicStruct = responseStruct.instance(TOPICS_KEY_NAME);
-            topicStruct.set(TOPIC_NAME, topicToPartitionEpochs.getKey());
-            Map partitionEpochs = topicToPartitionEpochs.getValue();
-            List partitions = new ArrayList<>();
-            for (Map.Entry partitionEndOffset : partitionEpochs.entrySet()) {
-                Struct partitionStruct = topicStruct.instance(PARTITIONS_KEY_NAME);
-                partitionStruct.set(ERROR_CODE, partitionEndOffset.getValue().error().code());
-                partitionStruct.set(PARTITION_ID, partitionEndOffset.getKey());
-                partitionStruct.set(END_OFFSET_KEY_NAME, partitionEndOffset.getValue().endOffset());
-                partitions.add(partitionStruct);
-            }
-            topicStruct.set(PARTITIONS_KEY_NAME, partitions.toArray());
-            topics.add(topicStruct);
-        }
-        responseStruct.set(TOPICS_KEY_NAME, topics.toArray());
-        return responseStruct;
+    @Override
+    public String toString() {
+        return data.toString();
     }
 }
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 91e3aebcf3f8d..758631a1d87aa 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,276 +16,147 @@
  */
 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.message.ProduceRequestData;
+import org.apache.kafka.common.message.ProduceResponseData;
 import org.apache.kafka.common.protocol.ApiKeys;
-import org.apache.kafka.common.protocol.CommonFields;
+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.Struct;
-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.BaseRecords;
+import org.apache.kafka.common.record.CompressionType;
 import org.apache.kafka.common.record.RecordBatch;
-import org.apache.kafka.common.utils.CollectionUtils;
+import org.apache.kafka.common.record.Records;
 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.Iterator;
-import java.util.List;
 import java.util.Map;
+import java.util.stream.Collectors;
 
-import static org.apache.kafka.common.protocol.CommonFields.NULLABLE_TRANSACTIONAL_ID;
-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.INT16;
-import static org.apache.kafka.common.protocol.types.Type.INT32;
-import static org.apache.kafka.common.protocol.types.Type.RECORDS;
+import static org.apache.kafka.common.requests.ProduceResponse.INVALID_OFFSET;
 
 public class ProduceRequest extends AbstractRequest {
-    private static final String ACKS_KEY_NAME = "acks";
-    private static final String TIMEOUT_KEY_NAME = "timeout";
-    private static final String TOPIC_DATA_KEY_NAME = "topic_data";
-
-    // topic level field names
-    private static final String PARTITION_DATA_KEY_NAME = "data";
-
-    // partition level field names
-    private static final String RECORD_SET_KEY_NAME = "record_set";
-
-
-    private static final Schema TOPIC_PRODUCE_DATA_V0 = new Schema(
-            TOPIC_NAME,
-            new Field(PARTITION_DATA_KEY_NAME, new ArrayOf(new Schema(
-                    PARTITION_ID,
-                    new Field(RECORD_SET_KEY_NAME, RECORDS)))));
-
-    private static final Schema PRODUCE_REQUEST_V0 = new Schema(
-            new Field(ACKS_KEY_NAME, INT16, "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."),
-            new Field(TIMEOUT_KEY_NAME, INT32, "The time to await a response in ms."),
-            new Field(TOPIC_DATA_KEY_NAME, new ArrayOf(TOPIC_PRODUCE_DATA_V0)));
-
-    /**
-     * The body of PRODUCE_REQUEST_V1 is the same as PRODUCE_REQUEST_V0.
-     * The version number is bumped up to indicate that the client supports quota throttle time field in the response.
-     */
-    private static final Schema PRODUCE_REQUEST_V1 = PRODUCE_REQUEST_V0;
-    /**
-     * The body of PRODUCE_REQUEST_V2 is the same as PRODUCE_REQUEST_V1.
-     * The version number is bumped up to indicate that message format V1 is used which has relative offset and
-     * timestamp.
-     */
-    private static final Schema PRODUCE_REQUEST_V2 = PRODUCE_REQUEST_V1;
-
-    // Produce request V3 adds the transactional id which is used for authorization when attempting to write
-    // transactional data. This version also adds support for message format V2.
-    private static final Schema PRODUCE_REQUEST_V3 = new Schema(
-            CommonFields.NULLABLE_TRANSACTIONAL_ID,
-            new Field(ACKS_KEY_NAME, INT16, "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."),
-            new Field(TIMEOUT_KEY_NAME, INT32, "The time to await a response in ms."),
-            new Field(TOPIC_DATA_KEY_NAME, new ArrayOf(TOPIC_PRODUCE_DATA_V0)));
-
-    /**
-     * The body of PRODUCE_REQUEST_V4 is the same as PRODUCE_REQUEST_V3.
-     * The version number is bumped up to indicate that the client supports KafkaStorageException.
-     * The KafkaStorageException will be translated to NotLeaderForPartitionException in the response if version <= 3
-     */
-    private static final Schema PRODUCE_REQUEST_V4 = PRODUCE_REQUEST_V3;
-
-    /**
-     * The body of the PRODUCE_REQUEST_V5 is the same as PRODUCE_REQUEST_V4.
-     * The version number is bumped since the PRODUCE_RESPONSE_V5 includes an additional partition level
-     * field: the log_start_offset.
-     */
-    private static final Schema PRODUCE_REQUEST_V5 = PRODUCE_REQUEST_V4;
 
+    public static Builder forMagic(byte magic, ProduceRequestData data) {
+        // Message format upgrades correspond with a bump in the produce request version. Older
+        // message format versions are generally not supported by the produce request versions
+        // following the bump.
+
+        final short minVersion;
+        final short maxVersion;
+        if (magic < RecordBatch.MAGIC_VALUE_V2) {
+            minVersion = 2;
+            maxVersion = 2;
+        } else {
+            minVersion = 3;
+            maxVersion = ApiKeys.PRODUCE.latestVersion();
+        }
+        return new Builder(minVersion, maxVersion, data);
+    }
 
-    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};
+    public static Builder forCurrentMagic(ProduceRequestData data) {
+        return forMagic(RecordBatch.CURRENT_MAGIC_VALUE, data);
     }
 
     public static class Builder extends AbstractRequest.Builder {
-        private final short acks;
-        private final int timeout;
-        private final Map partitionRecords;
-        private final String transactionalId;
-
-        public static Builder forCurrentMagic(short acks,
-                                              int timeout,
-                                              Map partitionRecords) {
-            return forMagic(RecordBatch.CURRENT_MAGIC_VALUE, acks, timeout, partitionRecords, null);
-        }
+        private final ProduceRequestData data;
 
-        public static Builder forMagic(byte magic,
-                                       short acks,
-                                       int timeout,
-                                       Map partitionRecords,
-                                       String transactionalId) {
-            // Message format upgrades correspond with a bump in the produce request version. Older
-            // message format versions are generally not supported by the produce request versions
-            // following the bump.
-
-            final short minVersion;
-            final short maxVersion;
-            if (magic < RecordBatch.MAGIC_VALUE_V2) {
-                minVersion = 2;
-                maxVersion = 2;
-            } else {
-                minVersion = 3;
-                maxVersion = ApiKeys.PRODUCE.latestVersion();
-            }
-            return new Builder(minVersion, maxVersion, acks, timeout, partitionRecords, transactionalId);
-        }
-
-        private Builder(short minVersion,
-                        short maxVersion,
-                        short acks,
-                        int timeout,
-                        Map partitionRecords,
-                        String transactionalId) {
+        public Builder(short minVersion,
+                       short maxVersion,
+                       ProduceRequestData data) {
             super(ApiKeys.PRODUCE, minVersion, maxVersion);
-            this.acks = acks;
-            this.timeout = timeout;
-            this.partitionRecords = partitionRecords;
-            this.transactionalId = transactionalId;
+            this.data = data;
         }
 
         @Override
         public ProduceRequest build(short version) {
-            return new ProduceRequest(version, acks, timeout, partitionRecords, transactionalId);
+            return build(version, true);
+        }
+
+        // Visible for testing only
+        public ProduceRequest buildUnsafe(short version) {
+            return build(version, false);
+        }
+
+        private ProduceRequest build(short version, boolean validate) {
+            if (validate) {
+                // Validate the given records first
+                data.topicData().forEach(tpd ->
+                        tpd.partitionData().forEach(partitionProduceData ->
+                                ProduceRequest.validateRecords(version, partitionProduceData.records())));
+            }
+            return new ProduceRequest(data, version);
         }
 
         @Override
         public String toString() {
             StringBuilder bld = new StringBuilder();
             bld.append("(type=ProduceRequest")
-                    .append(", acks=").append(acks)
-                    .append(", timeout=").append(timeout)
-                    .append(", partitionRecords=(").append(partitionRecords)
-                    .append("), transactionalId='").append(transactionalId != null ? transactionalId : "")
+                    .append(", acks=").append(data.acks())
+                    .append(", timeout=").append(data.timeoutMs())
+                    .append(", partitionRecords=(").append(data.topicData().stream().flatMap(d -> d.partitionData().stream()).collect(Collectors.toList()))
+                    .append("), transactionalId='").append(data.transactionalId() != null ? data.transactionalId() : "")
                     .append("'");
             return bld.toString();
         }
     }
 
+    /**
+     * We have to copy acks, timeout, transactionalId and partitionSizes from data since data maybe reset to eliminate
+     * the reference to ByteBuffer but those metadata are still useful.
+     */
     private final short acks;
     private final int timeout;
     private final String transactionalId;
-
-    private final Map partitionSizes;
-
     // This is set to null by `clearPartitionRecords` to prevent unnecessary memory retention when a produce request is
     // put in the purgatory (due to client throttling, it can take a while before the response is sent).
     // Care should be taken in methods that use this field.
-    private volatile Map partitionRecords;
-    private boolean transactional = false;
-    private boolean idempotent = false;
-
-    private ProduceRequest(short version, short acks, int timeout, Map partitionRecords, String transactionalId) {
-        super(version);
-        this.acks = acks;
-        this.timeout = timeout;
-
-        this.transactionalId = transactionalId;
-        this.partitionRecords = partitionRecords;
-        this.partitionSizes = createPartitionSizes(partitionRecords);
-
-        for (MemoryRecords records : partitionRecords.values())
-            validateRecords(version, records);
+    private volatile ProduceRequestData data;
+    // the partitionSizes is lazily initialized since it is used by server-side in production.
+    private volatile Map partitionSizes;
+
+    public ProduceRequest(ProduceRequestData produceRequestData, short version) {
+        super(ApiKeys.PRODUCE, version);
+        this.data = produceRequestData;
+        this.acks = data.acks();
+        this.timeout = data.timeoutMs();
+        this.transactionalId = data.transactionalId();
     }
 
-    private static Map createPartitionSizes(Map partitionRecords) {
-        Map result = new HashMap<>(partitionRecords.size());
-        for (Map.Entry entry : partitionRecords.entrySet())
-            result.put(entry.getKey(), entry.getValue().sizeInBytes());
-        return result;
-    }
-
-    public ProduceRequest(Struct struct, short version) {
-        super(version);
-        partitionRecords = new HashMap<>();
-        for (Object topicDataObj : struct.getArray(TOPIC_DATA_KEY_NAME)) {
-            Struct topicData = (Struct) topicDataObj;
-            String topic = topicData.get(TOPIC_NAME);
-            for (Object partitionResponseObj : topicData.getArray(PARTITION_DATA_KEY_NAME)) {
-                Struct partitionResponse = (Struct) partitionResponseObj;
-                int partition = partitionResponse.get(PARTITION_ID);
-                MemoryRecords records = (MemoryRecords) partitionResponse.getRecords(RECORD_SET_KEY_NAME);
-                validateRecords(version, records);
-                partitionRecords.put(new TopicPartition(topic, partition), records);
+    // visible for testing
+    Map partitionSizes() {
+        if (partitionSizes == null) {
+            // this method may be called by different thread (see the comment on data)
+            synchronized (this) {
+                if (partitionSizes == null) {
+                    partitionSizes = new HashMap<>();
+                    data.topicData().forEach(topicData ->
+                        topicData.partitionData().forEach(partitionData ->
+                            partitionSizes.compute(new TopicPartition(topicData.name(), partitionData.index()),
+                                (ignored, previousValue) ->
+                                    partitionData.records().sizeInBytes() + (previousValue == null ? 0 : previousValue))
+                        )
+                    );
+                }
             }
         }
-        partitionSizes = createPartitionSizes(partitionRecords);
-        acks = struct.getShort(ACKS_KEY_NAME);
-        timeout = struct.getInt(TIMEOUT_KEY_NAME);
-        transactionalId = struct.getOrElse(NULLABLE_TRANSACTIONAL_ID, null);
-    }
-
-    private void validateRecords(short version, MemoryRecords records) {
-        if (version >= 3) {
-            Iterator iterator = records.batches().iterator();
-            if (!iterator.hasNext())
-                throw new InvalidRecordException("Produce requests with version " + version + " must have at least " +
-                        "one record batch");
-
-            MutableRecordBatch entry = iterator.next();
-            if (entry.magic() != RecordBatch.MAGIC_VALUE_V2)
-                throw new InvalidRecordException("Produce requests with version " + version + " are only allowed to " +
-                        "contain record batches with magic version 2");
-
-            if (iterator.hasNext())
-                throw new InvalidRecordException("Produce requests with version " + version + " are only allowed to " +
-                        "contain exactly one record batch");
-            idempotent = entry.hasProducerId();
-            transactional = entry.isTransactional();
-        }
-
-        // Note that we do not do similar validation for older versions to ensure compatibility with
-        // clients which send the wrong magic version in the wrong version of the produce request. The broker
-        // did not do this validation before, so we maintain that behavior here.
+        return partitionSizes;
     }
 
     /**
-     * Visible for testing.
+     * @return data or IllegalStateException if the data is removed (to prevent unnecessary memory retention).
      */
     @Override
-    public Struct toStruct() {
+    public ProduceRequestData data() {
         // Store it in a local variable to protect against concurrent updates
-        Map partitionRecords = partitionRecordsOrFail();
-        short version = version();
-        Struct struct = new Struct(ApiKeys.PRODUCE.requestSchema(version));
-        Map> recordsByTopic = CollectionUtils.groupDataByTopic(partitionRecords);
-        struct.set(ACKS_KEY_NAME, acks);
-        struct.set(TIMEOUT_KEY_NAME, timeout);
-        struct.setIfExists(NULLABLE_TRANSACTIONAL_ID, transactionalId);
-
-        List topicDatas = new ArrayList<>(recordsByTopic.size());
-        for (Map.Entry> topicEntry : recordsByTopic.entrySet()) {
-            Struct topicData = struct.instance(TOPIC_DATA_KEY_NAME);
-            topicData.set(TOPIC_NAME, topicEntry.getKey());
-            List partitionArray = new ArrayList<>();
-            for (Map.Entry partitionEntry : topicEntry.getValue().entrySet()) {
-                MemoryRecords records = partitionEntry.getValue();
-                Struct part = topicData.instance(PARTITION_DATA_KEY_NAME)
-                        .set(PARTITION_ID, partitionEntry.getKey())
-                        .set(RECORD_SET_KEY_NAME, records);
-                partitionArray.add(part);
-            }
-            topicData.set(PARTITION_DATA_KEY_NAME, partitionArray.toArray());
-            topicDatas.add(topicData);
-        }
-        struct.set(TOPIC_DATA_KEY_NAME, topicDatas.toArray());
-        return struct;
+        ProduceRequestData tmp = data;
+        if (tmp == null)
+            throw new IllegalStateException("The partition records are no longer available because clearPartitionRecords() has been invoked.");
+        return tmp;
     }
 
     @Override
@@ -296,9 +167,9 @@ public String toString(boolean verbose) {
                 .append(",timeout=").append(timeout);
 
         if (verbose)
-            bld.append(",partitionSizes=").append(Utils.mkString(partitionSizes, "[", "]", "=", ","));
+            bld.append(",partitionSizes=").append(Utils.mkString(partitionSizes(), "[", "]", "=", ","));
         else
-            bld.append(",numPartitions=").append(partitionSizes.size());
+            bld.append(",numPartitions=").append(partitionSizes().size());
 
         bld.append("}");
         return bld.toString();
@@ -307,39 +178,31 @@ public String toString(boolean verbose) {
     @Override
     public ProduceResponse getErrorResponse(int throttleTimeMs, Throwable e) {
         /* In case the producer doesn't actually want any response */
-        if (acks == 0)
-            return null;
-
-        Errors error = Errors.forException(e);
-        Map responseMap = new HashMap<>();
-        ProduceResponse.PartitionResponse partitionResponse = new ProduceResponse.PartitionResponse(error);
-
-        for (TopicPartition tp : partitions())
-            responseMap.put(tp, partitionResponse);
-
-        short versionId = version();
-        switch (versionId) {
-            case 0:
-            case 1:
-            case 2:
-            case 3:
-            case 4:
-            case 5:
-                return new ProduceResponse(responseMap, 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.PRODUCE.latestVersion()));
-        }
+        if (acks == 0) return null;
+        ApiError apiError = ApiError.fromThrowable(e);
+        ProduceResponseData data = new ProduceResponseData().setThrottleTimeMs(throttleTimeMs);
+        partitionSizes().forEach((tp, ignored) -> {
+            ProduceResponseData.TopicProduceResponse tpr = data.responses().find(tp.topic());
+            if (tpr == null) {
+                tpr = new ProduceResponseData.TopicProduceResponse().setName(tp.topic());
+                data.responses().add(tpr);
+            }
+            tpr.partitionResponses().add(new ProduceResponseData.PartitionProduceResponse()
+                    .setIndex(tp.partition())
+                    .setRecordErrors(Collections.emptyList())
+                    .setBaseOffset(INVALID_OFFSET)
+                    .setLogAppendTimeMs(RecordBatch.NO_TIMESTAMP)
+                    .setLogStartOffset(INVALID_OFFSET)
+                    .setErrorMessage(apiError.message())
+                    .setErrorCode(apiError.error().code()));
+        });
+        return new ProduceResponse(data);
     }
 
     @Override
     public Map errorCounts(Throwable e) {
         Errors error = Errors.forException(e);
-        return Collections.singletonMap(error, partitions().size());
-    }
-
-    private Collection partitions() {
-        return partitionSizes.keySet();
+        return Collections.singletonMap(error, partitionSizes().size());
     }
 
     public short acks() {
@@ -354,35 +217,50 @@ public String transactionalId() {
         return transactionalId;
     }
 
-    public boolean isTransactional() {
-        return transactional;
-    }
-
-    public boolean isIdempotent() {
-        return idempotent;
+    public void clearPartitionRecords() {
+        // lazily initialize partitionSizes.
+        partitionSizes();
+        data = null;
     }
 
-    /**
-     * Returns the partition records or throws IllegalStateException if clearPartitionRecords() has been invoked.
-     */
-    public Map partitionRecordsOrFail() {
-        // Store it in a local variable to protect against concurrent updates
-        Map partitionRecords = this.partitionRecords;
-        if (partitionRecords == null)
-            throw new IllegalStateException("The partition records are no longer available because " +
-                    "clearPartitionRecords() has been invoked.");
-        return partitionRecords;
-    }
+    public static void validateRecords(short version, BaseRecords baseRecords) {
+        if (version >= 3) {
+            if (baseRecords instanceof Records) {
+                Records records = (Records) baseRecords;
+                Iterator iterator = records.batches().iterator();
+                if (!iterator.hasNext())
+                    throw new InvalidRecordException("Produce requests with version " + version + " must have at least " +
+                            "one record batch");
+
+                RecordBatch entry = iterator.next();
+                if (entry.magic() != RecordBatch.MAGIC_VALUE_V2)
+                    throw new InvalidRecordException("Produce requests with version " + version + " are only allowed to " +
+                            "contain record batches with magic version 2");
+                if (version < 7 && entry.compressionType() == CompressionType.ZSTD) {
+                    throw new UnsupportedCompressionTypeException("Produce requests with version " + version + " are not allowed to " +
+                            "use ZStandard compression");
+                }
+
+                if (iterator.hasNext())
+                    throw new InvalidRecordException("Produce requests with version " + version + " are only allowed to " +
+                            "contain exactly one record batch");
+            }
+        }
 
-    public void clearPartitionRecords() {
-        partitionRecords = null;
+        // Note that we do not do similar validation for older versions to ensure compatibility with
+        // clients which send the wrong magic version in the wrong version of the produce request. The broker
+        // did not do this validation before, so we maintain that behavior here.
     }
 
     public static ProduceRequest parse(ByteBuffer buffer, short version) {
-        return new ProduceRequest(ApiKeys.PRODUCE.parseRequest(version, buffer), version);
+        return new ProduceRequest(new ProduceRequestData(new ByteBufferAccessor(buffer), version), version);
     }
 
     public static byte requiredMagicForVersion(short produceRequestVersion) {
+        if (produceRequestVersion < ApiKeys.PRODUCE.oldestVersion() || produceRequestVersion > ApiKeys.PRODUCE.latestVersion())
+            throw new IllegalArgumentException("Magic value to use for produce request version " +
+                    produceRequestVersion + " is not known");
+
         switch (produceRequestVersion) {
             case 0:
             case 1:
@@ -391,18 +269,8 @@ public static byte requiredMagicForVersion(short produceRequestVersion) {
             case 2:
                 return RecordBatch.MAGIC_VALUE_V1;
 
-            case 3:
-            case 4:
-            case 5:
-                return RecordBatch.MAGIC_VALUE_V2;
-
             default:
-                // raise an exception if the version has not been explicitly added to this method.
-                // this ensures that we cannot accidentally use the wrong magic value if we forget
-                // to update this method on a bump to the produce request version.
-                throw new IllegalArgumentException("Magic value to use for produce request version " +
-                        produceRequestVersion + " is not known");
+                return RecordBatch.MAGIC_VALUE_V2;
         }
     }
-
 }
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 afedc9d6e88d9..b94e48b54a1c6 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
@@ -17,141 +17,56 @@
 package org.apache.kafka.common.requests;
 
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.message.ProduceResponseData;
 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.Struct;
 import org.apache.kafka.common.record.RecordBatch;
-import org.apache.kafka.common.utils.CollectionUtils;
 
 import java.nio.ByteBuffer;
-import java.util.ArrayList;
+import java.util.AbstractMap;
+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.THROTTLE_TIME_MS;
-import static org.apache.kafka.common.protocol.CommonFields.TOPIC_NAME;
-import static org.apache.kafka.common.protocol.types.Type.INT64;
+import java.util.Objects;
+import java.util.stream.Collectors;
 
 /**
- * This wrapper supports both v0 and v1 of ProduceResponse.
+ * This wrapper supports both v0 and v8 of ProduceResponse.
+ *
+ * Possible error code:
+ *
+ * {@link Errors#CORRUPT_MESSAGE}
+ * {@link Errors#UNKNOWN_TOPIC_OR_PARTITION}
+ * {@link Errors#NOT_LEADER_OR_FOLLOWER}
+ * {@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}
  */
 public class ProduceResponse extends AbstractResponse {
-
-    private static final String RESPONSES_KEY_NAME = "responses";
-
-    // topic level field names
-    private static final String PARTITION_RESPONSES_KEY_NAME = "partition_responses";
-
     public static final long INVALID_OFFSET = -1L;
+    private final ProduceResponseData data;
 
-    /**
-     * 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)
-     */
-
-    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 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 Schema PRODUCE_RESPONSE_V0 = 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))))))));
-
-    private static final Schema PRODUCE_RESPONSE_V1 = 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))))))),
-            THROTTLE_TIME_MS);
-
-    /**
-     * PRODUCE_RESPONSE_V2 added a timestamp field in the per partition response status.
-     * The timestamp is log append time if the topic is configured to use log append time. Or it is NoTimestamp when create
-     * time is used for the topic.
-     */
-    private static final Schema PRODUCE_RESPONSE_V2 = 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."))))))),
-            THROTTLE_TIME_MS);
-
-    private static final Schema PRODUCE_RESPONSE_V3 = PRODUCE_RESPONSE_V2;
-
-    /**
-     * The body of PRODUCE_RESPONSE_V4 is the same as PRODUCE_RESPONSE_V3.
-     * The version number is bumped up to indicate that the client supports KafkaStorageException.
-     * The KafkaStorageException will be translated to NotLeaderForPartitionException in the response if version <= 3
-     */
-    private static final Schema PRODUCE_RESPONSE_V4 = PRODUCE_RESPONSE_V3;
-
-
-    /**
-     * Add in the log_start_offset field to the partition response to filter out spurious OutOfOrderSequencExceptions
-     * on the client.
-     */
-    public static final Schema PRODUCE_RESPONSE_V5 = 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)))))),
-            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};
+    public ProduceResponse(ProduceResponseData produceResponseData) {
+        super(ApiKeys.PRODUCE);
+        this.data = produceResponseData;
     }
 
-    private final Map responses;
-    private final int throttleTime;
-
     /**
      * Constructor for Version 0
      * @param responses Produced data grouped by topic-partition
      */
+    @Deprecated
     public ProduceResponse(Map responses) {
         this(responses, DEFAULT_THROTTLE_TIME);
     }
@@ -159,85 +74,78 @@ public ProduceResponse(Map responses) {
     /**
      * Constructor for the latest version
      * @param responses Produced data grouped by topic-partition
-     * @param throttleTime Time in milliseconds the response was throttled
+     * @param throttleTimeMs Time in milliseconds the response was throttled
      */
-    public ProduceResponse(Map responses, int throttleTime) {
-        this.responses = responses;
-        this.throttleTime = throttleTime;
+    @Deprecated
+    public ProduceResponse(Map responses, int throttleTimeMs) {
+        this(toData(responses, throttleTimeMs));
     }
 
-    /**
-     * Constructor from a {@link Struct}.
-     */
-    public ProduceResponse(Struct struct) {
-        responses = new HashMap<>();
-        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 logStartOffset = partRespStruct.getOrElse(LOG_START_OFFSET_FIELD, INVALID_OFFSET);
-                TopicPartition tp = new TopicPartition(topic, partition);
-                responses.put(tp, new PartitionResponse(error, offset, logAppendTime, logStartOffset));
+    private static ProduceResponseData toData(Map responses, int throttleTimeMs) {
+        ProduceResponseData data = new ProduceResponseData().setThrottleTimeMs(throttleTimeMs);
+        responses.forEach((tp, response) -> {
+            ProduceResponseData.TopicProduceResponse tpr = data.responses().find(tp.topic());
+            if (tpr == null) {
+                tpr = new ProduceResponseData.TopicProduceResponse().setName(tp.topic());
+                data.responses().add(tpr);
             }
-        }
-        this.throttleTime = struct.getOrElse(THROTTLE_TIME_MS, DEFAULT_THROTTLE_TIME);
+            tpr.partitionResponses()
+                .add(new ProduceResponseData.PartitionProduceResponse()
+                    .setIndex(tp.partition())
+                    .setBaseOffset(response.baseOffset)
+                    .setLogStartOffset(response.logStartOffset)
+                    .setLogAppendTimeMs(response.logAppendTime)
+                    .setErrorMessage(response.errorMessage)
+                    .setErrorCode(response.error.code())
+                    .setRecordErrors(response.recordErrors
+                        .stream()
+                        .map(e -> new ProduceResponseData.BatchIndexAndErrorMessage()
+                            .setBatchIndex(e.batchIndex)
+                            .setBatchIndexErrorMessage(e.message))
+                        .collect(Collectors.toList())));
+        });
+        return data;
     }
 
     @Override
-    protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.PRODUCE.responseSchema(version));
-
-        Map> responseByTopic = CollectionUtils.groupDataByTopic(responses);
-        List topicDatas = new ArrayList<>(responseByTopic.size());
-        for (Map.Entry> entry : responseByTopic.entrySet()) {
-            Struct topicData = struct.instance(RESPONSES_KEY_NAME);
-            topicData.set(TOPIC_NAME, entry.getKey());
-            List partitionArray = new ArrayList<>();
-            for (Map.Entry partitionEntry : entry.getValue().entrySet()) {
-                PartitionResponse part = partitionEntry.getValue();
-                short errorCode = part.error.code();
-                // If producer sends ProduceRequest V3 or earlier, the client library is not guaranteed to recognize the error code
-                // for KafkaStorageException. In this case the client library will translate KafkaStorageException to
-                // UnknownServerException which is not retriable. We can ensure that producer will update metadata and retry
-                // by converting the KafkaStorageException to NotLeaderForPartitionException in the response if ProduceRequest version <= 3
-                if (errorCode == Errors.KAFKA_STORAGE_ERROR.code() && version <= 3)
-                    errorCode = Errors.NOT_LEADER_FOR_PARTITION.code();
-                Struct partStruct = topicData.instance(PARTITION_RESPONSES_KEY_NAME)
-                        .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_START_OFFSET_FIELD, part.logStartOffset);
-                partitionArray.add(partStruct);
-            }
-            topicData.set(PARTITION_RESPONSES_KEY_NAME, partitionArray.toArray());
-            topicDatas.add(topicData);
-        }
-        struct.set(RESPONSES_KEY_NAME, topicDatas.toArray());
-        struct.setIfExists(THROTTLE_TIME_MS, throttleTime);
-
-        return struct;
+    public ProduceResponseData data() {
+        return this.data;
     }
 
+    /**
+     * this method is used by testing only.
+     * refactor the tests which are using this method and then remove this method from production code.
+     * https://issues.apache.org/jira/browse/KAFKA-10697
+     */
+    @Deprecated
     public Map responses() {
-        return this.responses;
+        return data.responses()
+                .stream()
+                .flatMap(t -> t.partitionResponses()
+                        .stream()
+                        .map(p -> new AbstractMap.SimpleEntry<>(new TopicPartition(t.name(), p.index()),
+                                new PartitionResponse(
+                                        Errors.forCode(p.errorCode()),
+                                        p.baseOffset(),
+                                        p.logAppendTimeMs(),
+                                        p.logStartOffset(),
+                                        p.recordErrors()
+                                                .stream()
+                                                .map(e -> new RecordError(e.batchIndex(), e.batchIndexErrorMessage()))
+                                                .collect(Collectors.toList()),
+                                        p.errorMessage()))))
+                .collect(Collectors.toMap(AbstractMap.SimpleEntry::getKey, AbstractMap.SimpleEntry::getValue));
     }
 
-    public int getThrottleTime() {
-        return this.throttleTime;
+    @Override
+    public int throttleTimeMs() {
+        return this.data.throttleTimeMs();
     }
 
     @Override
     public Map errorCounts() {
         Map errorCounts = new HashMap<>();
-        for (PartitionResponse response : responses.values())
-            updateErrorCounts(errorCounts, response.error);
+        data.responses().forEach(t -> t.partitionResponses().forEach(p -> updateErrorCounts(errorCounts, Errors.forCode(p.errorCode()))));
         return errorCounts;
     }
 
@@ -246,16 +154,50 @@ 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, String errorMessage) {
+            this(error, INVALID_OFFSET, RecordBatch.NO_TIMESTAMP, INVALID_OFFSET, Collections.emptyList(), errorMessage);
+        }
+
         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
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (o == null || getClass() != o.getClass()) return false;
+            PartitionResponse that = (PartitionResponse) o;
+            return baseOffset == that.baseOffset &&
+                    logAppendTime == that.logAppendTime &&
+                    logStartOffset == that.logStartOffset &&
+                    error == that.error &&
+                    Objects.equals(recordErrors, that.recordErrors) &&
+                    Objects.equals(errorMessage, that.errorMessage);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(error, baseOffset, logAppendTime, logStartOffset, recordErrors, errorMessage);
         }
 
         @Override
@@ -270,13 +212,54 @@ 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;
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (o == null || getClass() != o.getClass()) return false;
+            RecordError that = (RecordError) o;
+            return batchIndex == that.batchIndex &&
+                    Objects.equals(message, that.message);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(batchIndex, message);
+        }
+    }
+
     public static ProduceResponse parse(ByteBuffer buffer, short version) {
-        return new ProduceResponse(ApiKeys.PRODUCE.responseSchema(version).read(buffer));
+        return new ProduceResponse(new ProduceResponseData(new ByteBufferAccessor(buffer), version));
     }
 
+    @Override
+    public boolean shouldClientThrottle(short version) {
+        return version >= 6;
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/RecordsSend.java b/clients/src/main/java/org/apache/kafka/common/requests/RecordsSend.java
deleted file mode 100644
index 6608e9b121823..0000000000000
--- a/clients/src/main/java/org/apache/kafka/common/requests/RecordsSend.java
+++ /dev/null
@@ -1,74 +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.network.Send;
-import org.apache.kafka.common.network.TransportLayers;
-import org.apache.kafka.common.record.Records;
-
-import java.io.EOFException;
-import java.io.IOException;
-import java.nio.ByteBuffer;
-import java.nio.channels.GatheringByteChannel;
-
-public class RecordsSend implements Send {
-    private static final ByteBuffer EMPTY_BYTE_BUFFER = ByteBuffer.allocate(0);
-
-    private final String destination;
-    private final Records records;
-    private int remaining;
-    private boolean pending = false;
-
-    public RecordsSend(String destination, Records records) {
-        this.destination = destination;
-        this.records = records;
-        this.remaining = records.sizeInBytes();
-    }
-
-    @Override
-    public String destination() {
-        return destination;
-    }
-
-    @Override
-    public boolean completed() {
-        return remaining <= 0 && !pending;
-    }
-
-    @Override
-    public long writeTo(GatheringByteChannel channel) throws IOException {
-        long written = 0;
-
-        if (remaining > 0) {
-            written = records.writeTo(channel, size() - remaining, remaining);
-            if (written < 0)
-                throw new EOFException("Wrote negative bytes to channel. This shouldn't happen.");
-            remaining -= written;
-        }
-
-        pending = TransportLayers.hasPendingWrites(channel);
-        if (remaining <= 0 && pending)
-            channel.write(EMPTY_BYTE_BUFFER);
-
-        return written;
-    }
-
-    @Override
-    public long size() {
-        return records.sizeInBytes();
-    }
-}
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
new file mode 100644
index 0000000000000..91a9f968d9d48
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/RenewDelegationTokenRequest.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.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.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.Errors;
+
+public class RenewDelegationTokenRequest extends AbstractRequest {
+
+    private final RenewDelegationTokenRequestData data;
+
+    public RenewDelegationTokenRequest(RenewDelegationTokenRequestData data, short version) {
+        super(ApiKeys.RENEW_DELEGATION_TOKEN, version);
+        this.data = data;
+    }
+
+    public static RenewDelegationTokenRequest parse(ByteBuffer buffer, short version) {
+        return new RenewDelegationTokenRequest(new RenewDelegationTokenRequestData(
+            new ByteBufferAccessor(buffer), version), version);
+    }
+
+    @Override
+    public RenewDelegationTokenRequestData data() {
+        return data;
+    }
+
+    @Override
+    public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
+        return new RenewDelegationTokenResponse(
+                new RenewDelegationTokenResponseData()
+                    .setThrottleTimeMs(throttleTimeMs)
+                    .setErrorCode(Errors.forException(e).code()));
+    }
+
+    public static class Builder extends AbstractRequest.Builder {
+        private final RenewDelegationTokenRequestData data;
+
+        public Builder(RenewDelegationTokenRequestData data) {
+            super(ApiKeys.RENEW_DELEGATION_TOKEN);
+            this.data = data;
+        }
+
+        @Override
+        public RenewDelegationTokenRequest build(short version) {
+            return new RenewDelegationTokenRequest(data, version);
+        }
+
+        @Override
+        public String 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
new file mode 100644
index 0000000000000..30708ff038c25
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/RenewDelegationTokenResponse.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.requests;
+
+import java.nio.ByteBuffer;
+import java.util.Map;
+
+import org.apache.kafka.common.message.RenewDelegationTokenResponseData;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.Errors;
+
+public class RenewDelegationTokenResponse extends AbstractResponse {
+
+    private final RenewDelegationTokenResponseData data;
+
+    public RenewDelegationTokenResponse(RenewDelegationTokenResponseData data) {
+        super(ApiKeys.RENEW_DELEGATION_TOKEN);
+        this.data = data;
+    }
+
+    public static RenewDelegationTokenResponse parse(ByteBuffer buffer, short version) {
+        return new RenewDelegationTokenResponse(new RenewDelegationTokenResponseData(
+            new ByteBufferAccessor(buffer), version));
+    }
+
+    @Override
+    public Map errorCounts() {
+        return errorCounts(error());
+    }
+
+    @Override
+    public RenewDelegationTokenResponseData data() {
+        return data;
+    }
+
+    @Override
+    public int throttleTimeMs() {
+        return data.throttleTimeMs();
+    }
+
+    public Errors error() {
+        return Errors.forCode(data.errorCode());
+    }
+
+    public long expiryTimestamp() {
+        return data.expiryTimestampMs();
+    }
+
+    public boolean hasError() {
+        return error() != Errors.NONE;
+    }
+
+    @Override
+    public boolean shouldClientThrottle(short version) {
+        return version >= 1;
+    }
+}
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 232c18aadd13b..225db375370d5 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,52 +17,82 @@
 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.ClientInformation;
 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.types.Struct;
 import org.apache.kafka.common.security.auth.KafkaPrincipal;
+import org.apache.kafka.common.security.auth.KafkaPrincipalSerde;
 import org.apache.kafka.common.security.auth.SecurityProtocol;
+import org.apache.kafka.server.authorizer.AuthorizableRequestContext;
 
 import java.net.InetAddress;
 import java.nio.ByteBuffer;
+import java.util.Optional;
 
 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;
     public final KafkaPrincipal principal;
     public final ListenerName listenerName;
     public final SecurityProtocol securityProtocol;
+    public final ClientInformation clientInformation;
+    public final boolean fromPrivilegedListener;
+    public final Optional principalSerde;
 
     public RequestContext(RequestHeader header,
                           String connectionId,
                           InetAddress clientAddress,
                           KafkaPrincipal principal,
                           ListenerName listenerName,
-                          SecurityProtocol securityProtocol) {
+                          SecurityProtocol securityProtocol,
+                          ClientInformation clientInformation,
+                          boolean fromPrivilegedListener) {
+        this(header,
+            connectionId,
+            clientAddress,
+            principal,
+            listenerName,
+            securityProtocol,
+            clientInformation,
+            fromPrivilegedListener,
+            Optional.empty());
+    }
+
+    public RequestContext(RequestHeader header,
+                          String connectionId,
+                          InetAddress clientAddress,
+                          KafkaPrincipal principal,
+                          ListenerName listenerName,
+                          SecurityProtocol securityProtocol,
+                          ClientInformation clientInformation,
+                          boolean fromPrivilegedListener,
+                          Optional principalSerde) {
         this.header = header;
         this.connectionId = connectionId;
         this.clientAddress = clientAddress;
         this.principal = principal;
         this.listenerName = listenerName;
         this.securityProtocol = securityProtocol;
+        this.clientInformation = clientInformation;
+        this.fromPrivilegedListener = fromPrivilegedListener;
+        this.principalSerde = principalSerde;
     }
 
     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();
             try {
                 short apiVersion = header.apiVersion();
-                Struct struct = apiKey.parseRequest(apiVersion, buffer);
-                AbstractRequest body = AbstractRequest.parseRequest(apiKey, apiVersion, struct);
-                return new RequestAndSize(body, struct.sizeOf());
+                return AbstractRequest.parseRequest(apiKey, apiVersion, buffer);
             } catch (Throwable ex) {
                 throw new InvalidRequestException("Error getting request for apiKey: " + apiKey +
                         ", apiVersion: " + header.apiVersion() +
@@ -73,19 +103,76 @@ public RequestAndSize parseRequest(ByteBuffer buffer) {
         }
     }
 
-    public Send buildResponse(AbstractResponse body) {
-        ResponseHeader responseHeader = header.toResponseHeader();
-        short version = header.apiVersion();
+    /**
+     * Build a {@link Send} for direct transmission of the provided response
+     * over the network.
+     */
+    public Send buildResponseSend(AbstractResponse body) {
+        return body.toSend(header.toResponseHeader(), apiVersion());
+    }
+
+    /**
+     * Serialize a response into a {@link ByteBuffer}. This is used when the response
+     * will be encapsulated in an {@link EnvelopeResponse}. The buffer will contain
+     * both the serialized {@link ResponseHeader} as well as the bytes from the response.
+     * There is no `size` prefix unlike the output from {@link #buildResponseSend(AbstractResponse)}.
+     *
+     * Note that envelope requests are reserved only for APIs which have set the
+     * {@link ApiKeys#forwardable} flag. Notably the `Fetch` API cannot be forwarded,
+     * so we do not lose the benefit of "zero copy" transfers from disk.
+     */
+    public ByteBuffer buildResponseEnvelopePayload(AbstractResponse body) {
+        return body.serializeWithHeader(header.toResponseHeader(), apiVersion());
+    }
+
+    private boolean isUnsupportedApiVersionsRequest() {
+        return header.apiKey() == API_VERSIONS && !API_VERSIONS.isVersionSupported(header.apiVersion());
+    }
 
+    public short apiVersion() {
         // Use v0 when serializing an unhandled ApiVersion response
         if (isUnsupportedApiVersionsRequest())
-            version = 0;
+            return 0;
+        return header.apiVersion();
+    }
 
-        return body.toSend(connectionId, responseHeader, version);
+    @Override
+    public String listenerName() {
+        return listenerName.value();
     }
 
-    private boolean isUnsupportedApiVersionsRequest() {
-        return header.apiKey() == API_VERSIONS && !API_VERSIONS.isVersionSupported(header.apiVersion());
+    @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..29539491d1469 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,97 @@
 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.types.Struct;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.ObjectSerializationCache;
 
 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 RequestHeader(ApiKeys apiKey, short version, String clientId, int correlation) {
-        this.apiKey = requireNonNull(apiKey);
-        this.apiVersion = version;
-        this.clientId = clientId;
-        this.correlationId = correlation;
+public class RequestHeader implements AbstractRequestResponse {
+    private final RequestHeaderData data;
+    private final short headerVersion;
+
+    public RequestHeader(ApiKeys requestApiKey, short requestVersion, String clientId, int correlationId) {
+        this(new RequestHeaderData().
+                setRequestApiKey(requestApiKey.id).
+                setRequestApiVersion(requestVersion).
+                setClientId(clientId).
+                setCorrelationId(correlationId),
+            requestApiKey.requestHeaderVersion(requestVersion));
     }
 
-    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;
+    public RequestHeader(RequestHeaderData data, short headerVersion) {
+        this.data = data;
+        this.headerVersion = 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 void write(ByteBuffer buffer, ObjectSerializationCache serializationCache) {
+        data.write(new ByteBufferAccessor(buffer), serializationCache, headerVersion);
+    }
+
+    public int size(ObjectSerializationCache serializationCache) {
+        return data.size(serializationCache, headerVersion);
     }
 
     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();
+            // We derive the header version from the request api version, so we read that first.
+            // The request api version is part of `RequestHeaderData`, so we reset the buffer position after the read.
+            int position = buffer.position();
+            apiKey = buffer.getShort();
             short apiVersion = buffer.getShort();
-            Schema schema = schema(apiKey, apiVersion);
-            buffer.rewind();
-            return new RequestHeader(schema.read(buffer));
-        } catch (InvalidRequestException e) {
-            throw e;
-        } catch (Throwable  ex) {
+            short headerVersion = ApiKeys.forId(apiKey).requestHeaderVersion(apiVersion);
+            buffer.position(position);
+            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 +115,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 103246cac5064..929eb2fd2a0cb 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
@@ -16,78 +16,81 @@
  */
 package org.apache.kafka.common.requests;
 
-import org.apache.kafka.common.acl.AccessControlEntry;
-import org.apache.kafka.common.acl.AccessControlEntryFilter;
-import org.apache.kafka.common.acl.AclOperation;
-import org.apache.kafka.common.acl.AclPermissionType;
-import org.apache.kafka.common.protocol.types.Struct;
-import org.apache.kafka.common.resource.Resource;
-import org.apache.kafka.common.resource.ResourceFilter;
-import org.apache.kafka.common.resource.ResourceType;
+import org.apache.kafka.common.message.ProduceRequestData;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.Message;
+import org.apache.kafka.common.protocol.ObjectSerializationCache;
+import org.apache.kafka.common.record.BaseRecords;
+import org.apache.kafka.common.record.RecordBatch;
+import org.apache.kafka.common.record.Records;
 
-import static org.apache.kafka.common.protocol.CommonFields.HOST;
-import static org.apache.kafka.common.protocol.CommonFields.HOST_FILTER;
-import static org.apache.kafka.common.protocol.CommonFields.OPERATION;
-import static org.apache.kafka.common.protocol.CommonFields.PERMISSION_TYPE;
-import static org.apache.kafka.common.protocol.CommonFields.PRINCIPAL;
-import static org.apache.kafka.common.protocol.CommonFields.PRINCIPAL_FILTER;
-import static org.apache.kafka.common.protocol.CommonFields.RESOURCE_NAME;
-import static org.apache.kafka.common.protocol.CommonFields.RESOURCE_NAME_FILTER;
-import static org.apache.kafka.common.protocol.CommonFields.RESOURCE_TYPE;
+import java.nio.ByteBuffer;
+import java.util.AbstractMap;
+import java.util.Iterator;
+import java.util.Optional;
 
-class RequestUtils {
+public final class RequestUtils {
 
-    static Resource resourceFromStructFields(Struct struct) {
-        byte resourceType = struct.get(RESOURCE_TYPE);
-        String name = struct.get(RESOURCE_NAME);
-        return new Resource(ResourceType.fromCode(resourceType), name);
-    }
+    private RequestUtils() {}
 
-    static void resourceSetStructFields(Resource resource, Struct struct) {
-        struct.set(RESOURCE_TYPE, resource.resourceType().code());
-        struct.set(RESOURCE_NAME, resource.name());
+    static Optional getLeaderEpoch(int leaderEpoch) {
+        return leaderEpoch == RecordBatch.NO_PARTITION_LEADER_EPOCH ?
+            Optional.empty() : Optional.of(leaderEpoch);
     }
 
-    static ResourceFilter resourceFilterFromStructFields(Struct struct) {
-        byte resourceType = struct.get(RESOURCE_TYPE);
-        String name = struct.get(RESOURCE_NAME_FILTER);
-        return new ResourceFilter(ResourceType.fromCode(resourceType), name);
+    // visible for testing
+    public static boolean hasIdempotentRecords(ProduceRequest request) {
+        return flags(request).getKey();
     }
 
-    static void resourceFilterSetStructFields(ResourceFilter resourceFilter, Struct struct) {
-        struct.set(RESOURCE_TYPE, resourceFilter.resourceType().code());
-        struct.set(RESOURCE_NAME_FILTER, resourceFilter.name());
+    // visible for testing
+    public static boolean hasTransactionalRecords(ProduceRequest request) {
+        return flags(request).getValue();
     }
 
-    static AccessControlEntry aceFromStructFields(Struct struct) {
-        String principal = struct.get(PRINCIPAL);
-        String host = struct.get(HOST);
-        byte operation = struct.get(OPERATION);
-        byte permissionType = struct.get(PERMISSION_TYPE);
-        return new AccessControlEntry(principal, host, AclOperation.fromCode(operation),
-            AclPermissionType.fromCode(permissionType));
+    /**
+     * Get both hasIdempotentRecords flag and hasTransactionalRecords flag from produce request.
+     * Noted that we find all flags at once to avoid duplicate loop and record batch construction.
+     * @return first flag is "hasIdempotentRecords" and another is "hasTransactionalRecords"
+     */
+    public static AbstractMap.SimpleEntry flags(ProduceRequest request) {
+        boolean hasIdempotentRecords = false;
+        boolean hasTransactionalRecords = false;
+        for (ProduceRequestData.TopicProduceData tpd : request.data().topicData()) {
+            for (ProduceRequestData.PartitionProduceData ppd : tpd.partitionData()) {
+                BaseRecords records = ppd.records();
+                if (records instanceof Records) {
+                    Iterator iterator = ((Records) records).batches().iterator();
+                    if (iterator.hasNext()) {
+                        RecordBatch batch = iterator.next();
+                        hasIdempotentRecords = hasIdempotentRecords || batch.hasProducerId();
+                        hasTransactionalRecords = hasTransactionalRecords || batch.isTransactional();
+                    }
+                }
+                // return early
+                if (hasIdempotentRecords && hasTransactionalRecords)
+                    return new AbstractMap.SimpleEntry<>(true, true);
+            }
+        }
+        return new AbstractMap.SimpleEntry<>(hasIdempotentRecords, hasTransactionalRecords);
     }
 
-    static void aceSetStructFields(AccessControlEntry data, Struct struct) {
-        struct.set(PRINCIPAL, data.principal());
-        struct.set(HOST, data.host());
-        struct.set(OPERATION, data.operation().code());
-        struct.set(PERMISSION_TYPE, data.permissionType().code());
-    }
+    public static ByteBuffer serialize(
+        Message header,
+        short headerVersion,
+        Message apiMessage,
+        short apiVersion
+    ) {
+        ObjectSerializationCache cache = new ObjectSerializationCache();
 
-    static AccessControlEntryFilter aceFilterFromStructFields(Struct struct) {
-        String principal = struct.get(PRINCIPAL_FILTER);
-        String host = struct.get(HOST_FILTER);
-        byte operation = struct.get(OPERATION);
-        byte permissionType = struct.get(PERMISSION_TYPE);
-        return new AccessControlEntryFilter(principal, host, AclOperation.fromCode(operation),
-            AclPermissionType.fromCode(permissionType));
-    }
+        int headerSize = header.size(cache, headerVersion);
+        int messageSize = apiMessage.size(cache, apiVersion);
+        ByteBufferAccessor writable = new ByteBufferAccessor(ByteBuffer.allocate(headerSize + messageSize));
+
+        header.write(writable, cache, headerVersion);
+        apiMessage.write(writable, cache, apiVersion);
 
-    static void aceFilterSetStructFields(AccessControlEntryFilter filter, Struct struct) {
-        struct.set(PRINCIPAL_FILTER, filter.principal());
-        struct.set(HOST_FILTER, filter.host());
-        struct.set(OPERATION, filter.operation().code());
-        struct.set(PERMISSION_TYPE, filter.permissionType().code());
+        writable.flip();
+        return writable.buffer();
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/Resource.java b/clients/src/main/java/org/apache/kafka/common/requests/Resource.java
deleted file mode 100644
index 6a360a5900a61..0000000000000
--- a/clients/src/main/java/org/apache/kafka/common/requests/Resource.java
+++ /dev/null
@@ -1,60 +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;
-
-public final class Resource {
-    private final ResourceType type;
-    private final String name;
-
-    public Resource(ResourceType type, String name) {
-        this.type = type;
-        this.name = name;
-    }
-
-    public ResourceType type() {
-        return type;
-    }
-
-    public String name() {
-        return name;
-    }
-
-    @Override
-    public boolean equals(Object o) {
-        if (this == o)
-            return true;
-        if (o == null || getClass() != o.getClass())
-            return false;
-
-        Resource resource = (Resource) o;
-
-        return type == resource.type && name.equals(resource.name);
-    }
-
-    @Override
-    public int hashCode() {
-        int result = type.hashCode();
-        result = 31 * result + name.hashCode();
-        return result;
-    }
-
-    @Override
-    public String toString() {
-        return "Resource(type=" + type + ", name='" + name + "'}";
-    }
-}
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/ResourceType.java b/clients/src/main/java/org/apache/kafka/common/requests/ResourceType.java
deleted file mode 100644
index 2c117727f93d6..0000000000000
--- a/clients/src/main/java/org/apache/kafka/common/requests/ResourceType.java
+++ /dev/null
@@ -1,42 +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;
-
-public enum ResourceType {
-    UNKNOWN((byte) 0), ANY((byte) 1), TOPIC((byte) 2), GROUP((byte) 3), BROKER((byte) 4);
-
-    private static final ResourceType[] VALUES = values();
-
-    private final byte id;
-
-    ResourceType(byte id) {
-        this.id = id;
-    }
-
-    public byte id() {
-        return id;
-    }
-
-    public static ResourceType forId(byte id) {
-        if (id < 0)
-            throw new IllegalArgumentException("id should be positive, id: " + id);
-        if (id >= VALUES.length)
-            return UNKNOWN;
-        return VALUES[id];
-    }
-}
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..d45e28dc1d81d 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,49 +16,74 @@
  */
 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.protocol.types.Struct;
+import org.apache.kafka.common.message.ResponseHeaderData;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.ObjectSerializationCache;
 
 import java.nio.ByteBuffer;
-
-import static org.apache.kafka.common.protocol.types.Type.INT32;
+import java.util.Objects;
 
 /**
  * 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;
+
+    public ResponseHeader(int correlationId, short headerVersion) {
+        this(new ResponseHeaderData().setCorrelationId(correlationId), headerVersion);
+    }
 
-    private final int correlationId;
+    public ResponseHeader(ResponseHeaderData data, short headerVersion) {
+        this.data = data;
+        this.headerVersion = headerVersion;
+    }
 
-    public ResponseHeader(Struct struct) {
-        correlationId = struct.getInt(CORRELATION_KEY_FIELD);
+    public int size(ObjectSerializationCache serializationCache) {
+        return data().size(serializationCache, headerVersion);
     }
 
-    public ResponseHeader(int correlationId) {
-        this.correlationId = correlationId;
+    public int correlationId() {
+        return this.data.correlationId();
     }
 
-    public int sizeOf() {
-        return toStruct().sizeOf();
+    public short headerVersion() {
+        return headerVersion;
     }
 
-    public Struct toStruct() {
-        Struct struct = new Struct(SCHEMA);
-        struct.set(CORRELATION_KEY_FIELD, correlationId);
-        return struct;
+    public ResponseHeaderData data() {
+        return data;
     }
 
-    public int correlationId() {
-        return correlationId;
+    public void write(ByteBuffer buffer, ObjectSerializationCache serializationCache) {
+        data.write(new ByteBufferAccessor(buffer), serializationCache, headerVersion);
     }
 
-    public static ResponseHeader parse(ByteBuffer buffer) {
-        return new ResponseHeader(SCHEMA.read(buffer));
+    @Override
+    public String toString() {
+        return "ResponseHeader("
+            + "correlationId=" + data.correlationId()
+            + ", headerVersion=" + headerVersion
+            + ")";
     }
 
+    public static ResponseHeader parse(ByteBuffer buffer, short headerVersion) {
+        return new ResponseHeader(
+            new ResponseHeaderData(new ByteBufferAccessor(buffer), headerVersion),
+                headerVersion);
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        ResponseHeader that = (ResponseHeader) o;
+        return headerVersion == that.headerVersion &&
+            Objects.equals(data, that.data);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(data, headerVersion);
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/SaslAuthenticateRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/SaslAuthenticateRequest.java
index 74d31a688de66..e2080ce841bfc 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/SaslAuthenticateRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/SaslAuthenticateRequest.java
@@ -16,15 +16,13 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.message.SaslAuthenticateRequestData;
+import org.apache.kafka.common.message.SaslAuthenticateResponseData;
 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.protocol.ByteBufferAccessor;
 
 import java.nio.ByteBuffer;
 
-import static org.apache.kafka.common.protocol.types.Type.BYTES;
 
 /**
  * Request from SASL client containing client SASL authentication token as defined by the
@@ -35,77 +33,50 @@
  * brokers will send SaslHandshake request v0 followed by SASL tokens without the Kafka request headers.
  */
 public class SaslAuthenticateRequest extends AbstractRequest {
-    private static final String SASL_AUTH_BYTES_KEY_NAME = "sasl_auth_bytes";
-
-    private static final Schema SASL_AUTHENTICATE_REQUEST_V0 = new Schema(
-            new Field(SASL_AUTH_BYTES_KEY_NAME, BYTES, "SASL authentication bytes from client as defined by the SASL mechanism."));
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{SASL_AUTHENTICATE_REQUEST_V0};
-    }
-
-    private final ByteBuffer saslAuthBytes;
 
     public static class Builder extends AbstractRequest.Builder {
-        private final ByteBuffer saslAuthBytes;
+        private final SaslAuthenticateRequestData data;
 
-        public Builder(ByteBuffer saslAuthBytes) {
+        public Builder(SaslAuthenticateRequestData data) {
             super(ApiKeys.SASL_AUTHENTICATE);
-            this.saslAuthBytes = saslAuthBytes;
+            this.data = data;
         }
 
         @Override
         public SaslAuthenticateRequest build(short version) {
-            return new SaslAuthenticateRequest(saslAuthBytes, version);
+            return new SaslAuthenticateRequest(data, version);
         }
 
         @Override
         public String toString() {
-            StringBuilder bld = new StringBuilder();
-            bld.append("(type=SaslAuthenticateRequest)");
-            return bld.toString();
+            return "(type=SaslAuthenticateRequest)";
         }
     }
 
-    public SaslAuthenticateRequest(ByteBuffer saslAuthBytes) {
-        this(saslAuthBytes, ApiKeys.SASL_AUTHENTICATE.latestVersion());
-    }
-
-    public SaslAuthenticateRequest(ByteBuffer saslAuthBytes, short version) {
-        super(version);
-        this.saslAuthBytes = saslAuthBytes;
-    }
+    private final SaslAuthenticateRequestData data;
 
-    public SaslAuthenticateRequest(Struct struct, short version) {
-        super(version);
-        saslAuthBytes = struct.getBytes(SASL_AUTH_BYTES_KEY_NAME);
+    public SaslAuthenticateRequest(SaslAuthenticateRequestData data, short version) {
+        super(ApiKeys.SASL_AUTHENTICATE, version);
+        this.data = data;
     }
 
-    public ByteBuffer saslAuthBytes() {
-        return saslAuthBytes;
+    @Override
+    public SaslAuthenticateRequestData data() {
+        return data;
     }
 
     @Override
     public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
-        short versionId = version();
-        switch (versionId) {
-            case 0:
-                return new SaslAuthenticateResponse(Errors.forException(e), e.getMessage());
-            default:
-                throw new IllegalArgumentException(String.format("Version %d is not valid. Valid versions for %s are 0 to %d",
-                        versionId, this.getClass().getSimpleName(), ApiKeys.SASL_AUTHENTICATE.latestVersion()));
-        }
+        ApiError apiError = ApiError.fromThrowable(e);
+        SaslAuthenticateResponseData response = new SaslAuthenticateResponseData()
+                .setErrorCode(apiError.error().code())
+                .setErrorMessage(apiError.message());
+        return new SaslAuthenticateResponse(response);
     }
 
     public static SaslAuthenticateRequest parse(ByteBuffer buffer, short version) {
-        return new SaslAuthenticateRequest(ApiKeys.SASL_AUTHENTICATE.parseRequest(version, buffer), version);
-    }
-
-    @Override
-    protected Struct toStruct() {
-        Struct struct = new Struct(ApiKeys.SASL_AUTHENTICATE.requestSchema(version()));
-        struct.set(SASL_AUTH_BYTES_KEY_NAME, saslAuthBytes);
-        return struct;
+        return new SaslAuthenticateRequest(new SaslAuthenticateRequestData(new ByteBufferAccessor(buffer), version),
+            version);
     }
 }
 
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/SaslAuthenticateResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/SaslAuthenticateResponse.java
index e2441825b57c4..bd12d3d4ae7bb 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/SaslAuthenticateResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/SaslAuthenticateResponse.java
@@ -16,91 +16,63 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.message.SaslAuthenticateResponseData;
 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.Field;
-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;
-import static org.apache.kafka.common.protocol.CommonFields.ERROR_MESSAGE;
-import static org.apache.kafka.common.protocol.types.Type.BYTES;
-
-
 /**
  * Response from SASL server which for a SASL challenge as defined by the SASL protocol
  * for the mechanism configured for the client.
  */
 public class SaslAuthenticateResponse extends AbstractResponse {
-    private static final String SASL_AUTH_BYTES_KEY_NAME = "sasl_auth_bytes";
 
-    private static final Schema SASL_AUTHENTICATE_RESPONSE_V0 = new Schema(
-            ERROR_CODE,
-            ERROR_MESSAGE,
-            new Field(SASL_AUTH_BYTES_KEY_NAME, BYTES, "SASL authentication bytes from server as defined by the SASL mechanism."));
+    private final SaslAuthenticateResponseData data;
 
-    public static Schema[] schemaVersions() {
-        return new Schema[]{SASL_AUTHENTICATE_RESPONSE_V0};
+    public SaslAuthenticateResponse(SaslAuthenticateResponseData data) {
+        super(ApiKeys.SASL_AUTHENTICATE);
+        this.data = data;
     }
 
-    private static final ByteBuffer EMPTY_BUFFER = ByteBuffer.allocate(0);
-
-    private final ByteBuffer saslAuthBytes;
-
     /**
      * Possible error codes:
      *   SASL_AUTHENTICATION_FAILED(57) : Authentication failed
      */
-    private final Errors error;
-    private final String errorMessage;
-
-    public SaslAuthenticateResponse(Errors error, String errorMessage) {
-        this(error, errorMessage, EMPTY_BUFFER);
-    }
-
-    public SaslAuthenticateResponse(Errors error, String errorMessage, ByteBuffer saslAuthBytes) {
-        this.error = error;
-        this.errorMessage = errorMessage;
-        this.saslAuthBytes = saslAuthBytes;
+    public Errors error() {
+        return Errors.forCode(data.errorCode());
     }
 
-    public SaslAuthenticateResponse(Struct struct) {
-        error = Errors.forCode(struct.get(ERROR_CODE));
-        errorMessage = struct.get(ERROR_MESSAGE);
-        saslAuthBytes = struct.getBytes(SASL_AUTH_BYTES_KEY_NAME);
+    @Override
+    public Map errorCounts() {
+        return errorCounts(Errors.forCode(data.errorCode()));
     }
 
-    public Errors error() {
-        return error;
+    public String errorMessage() {
+        return data.errorMessage();
     }
 
-    public String errorMessage() {
-        return errorMessage;
+    public long sessionLifetimeMs() {
+        return data.sessionLifetimeMs();
     }
 
-    public ByteBuffer saslAuthBytes() {
-        return saslAuthBytes;
+    public byte[] saslAuthBytes() {
+        return data.authBytes();
     }
 
     @Override
-    public Map errorCounts() {
-        return errorCounts(error);
+    public int throttleTimeMs() {
+        return DEFAULT_THROTTLE_TIME;
     }
 
     @Override
-    public Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.SASL_AUTHENTICATE.responseSchema(version));
-        struct.set(ERROR_CODE, error.code());
-        struct.set(ERROR_MESSAGE, errorMessage);
-        struct.set(SASL_AUTH_BYTES_KEY_NAME, saslAuthBytes);
-        return struct;
+    public SaslAuthenticateResponseData data() {
+        return data;
     }
 
     public static SaslAuthenticateResponse parse(ByteBuffer buffer, short version) {
-        return new SaslAuthenticateResponse(ApiKeys.SASL_AUTHENTICATE.parseResponse(version, buffer));
+        return new SaslAuthenticateResponse(new SaslAuthenticateResponseData(new ByteBufferAccessor(buffer), version));
     }
 }
-
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/SaslHandshakeRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/SaslHandshakeRequest.java
index a06a4db870ce1..09d3a87c8e727 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/SaslHandshakeRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/SaslHandshakeRequest.java
@@ -16,17 +16,13 @@
  */
 package org.apache.kafka.common.requests;
 
+
+import org.apache.kafka.common.message.SaslHandshakeRequestData;
+import org.apache.kafka.common.message.SaslHandshakeResponseData;
 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.protocol.ByteBufferAccessor;
 
 import java.nio.ByteBuffer;
-import java.util.Collections;
-import java.util.List;
-
-import static org.apache.kafka.common.protocol.types.Type.STRING;
 
 /**
  * Request from SASL client containing client SASL mechanism.
@@ -38,84 +34,47 @@
  * making it easy to distinguish from a GSSAPI packet.
  */
 public class SaslHandshakeRequest extends AbstractRequest {
-    private static final String MECHANISM_KEY_NAME = "mechanism";
-
-    private static final Schema SASL_HANDSHAKE_REQUEST_V0 = new Schema(
-            new Field("mechanism", STRING, "SASL Mechanism chosen by the client."));
-
-    // SASL_HANDSHAKE_REQUEST_V1 added to support SASL_AUTHENTICATE request to improve diagnostics
-    private static final Schema SASL_HANDSHAKE_REQUEST_V1 = SASL_HANDSHAKE_REQUEST_V0;
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{SASL_HANDSHAKE_REQUEST_V0, SASL_HANDSHAKE_REQUEST_V1};
-    }
-
-    private final String mechanism;
 
     public static class Builder extends AbstractRequest.Builder {
-        private final String mechanism;
+        private final SaslHandshakeRequestData data;
 
-        public Builder(String mechanism) {
+        public Builder(SaslHandshakeRequestData data) {
             super(ApiKeys.SASL_HANDSHAKE);
-            this.mechanism = mechanism;
+            this.data = data;
         }
 
         @Override
         public SaslHandshakeRequest build(short version) {
-            return new SaslHandshakeRequest(mechanism, version);
+            return new SaslHandshakeRequest(data, version);
         }
 
         @Override
         public String toString() {
-            StringBuilder bld = new StringBuilder();
-            bld.append("(type=SaslHandshakeRequest").
-                append(", mechanism=").append(mechanism).
-                append(")");
-            return bld.toString();
+            return data.toString();
         }
     }
 
-    public SaslHandshakeRequest(String mechanism) {
-        this(mechanism, ApiKeys.SASL_HANDSHAKE.latestVersion());
-    }
+    private final SaslHandshakeRequestData data;
 
-    public SaslHandshakeRequest(String mechanism, short version) {
-        super(version);
-        this.mechanism = mechanism;
+    public SaslHandshakeRequest(SaslHandshakeRequestData data, short version) {
+        super(ApiKeys.SASL_HANDSHAKE, version);
+        this.data = data;
     }
 
-    public SaslHandshakeRequest(Struct struct, short version) {
-        super(version);
-        mechanism = struct.getString(MECHANISM_KEY_NAME);
-    }
-
-    public String mechanism() {
-        return mechanism;
+    @Override
+    public SaslHandshakeRequestData data() {
+        return data;
     }
 
     @Override
     public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
-        short versionId = version();
-        switch (versionId) {
-            case 0:
-            case 1:
-                List enabledMechanisms = Collections.emptyList();
-                return new SaslHandshakeResponse(Errors.forException(e), enabledMechanisms);
-            default:
-                throw new IllegalArgumentException(String.format("Version %d is not valid. Valid versions for %s are 0 to %d",
-                        versionId, this.getClass().getSimpleName(), ApiKeys.SASL_HANDSHAKE.latestVersion()));
-        }
+        SaslHandshakeResponseData response = new SaslHandshakeResponseData();
+        response.setErrorCode(ApiError.fromThrowable(e).error().code());
+        return new SaslHandshakeResponse(response);
     }
 
     public static SaslHandshakeRequest parse(ByteBuffer buffer, short version) {
-        return new SaslHandshakeRequest(ApiKeys.SASL_HANDSHAKE.parseRequest(version, buffer), version);
-    }
-
-    @Override
-    protected Struct toStruct() {
-        Struct struct = new Struct(ApiKeys.SASL_HANDSHAKE.requestSchema(version()));
-        struct.set(MECHANISM_KEY_NAME, mechanism);
-        return struct;
+        return new SaslHandshakeRequest(new SaslHandshakeRequestData(new ByteBufferAccessor(buffer), version), version);
     }
 }
 
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/SaslHandshakeResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/SaslHandshakeResponse.java
index 252c62dd23065..63c047a06196b 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/SaslHandshakeResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/SaslHandshakeResponse.java
@@ -16,85 +16,57 @@
  */
 package org.apache.kafka.common.requests;
 
+import org.apache.kafka.common.message.SaslHandshakeResponseData;
 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.Struct;
-import org.apache.kafka.common.protocol.types.Type;
 
 import java.nio.ByteBuffer;
-import java.util.ArrayList;
-import java.util.Collection;
 import java.util.List;
 import java.util.Map;
 
-import static org.apache.kafka.common.protocol.CommonFields.ERROR_CODE;
-
-
 /**
  * Response from SASL server which indicates if the client-chosen mechanism is enabled in the server.
  * For error responses, the list of enabled mechanisms is included in the response.
  */
 public class SaslHandshakeResponse extends AbstractResponse {
-    private static final String ENABLED_MECHANISMS_KEY_NAME = "enabled_mechanisms";
-
-    private static final Schema SASL_HANDSHAKE_RESPONSE_V0 = new Schema(
-            ERROR_CODE,
-            new Field(ENABLED_MECHANISMS_KEY_NAME, new ArrayOf(Type.STRING), "Array of mechanisms enabled in the server."));
-
-    private static final Schema SASL_HANDSHAKE_RESPONSE_V1 = SASL_HANDSHAKE_RESPONSE_V0;
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{SASL_HANDSHAKE_RESPONSE_V0, SASL_HANDSHAKE_RESPONSE_V1};
-    }
 
-    /**
-     * Possible error codes:
-     *   UNSUPPORTED_SASL_MECHANISM(33): Client mechanism not enabled in server
-     *   ILLEGAL_SASL_STATE(34) : Invalid request during SASL handshake
-     */
-    private final Errors error;
-    private final List enabledMechanisms;
+    private final SaslHandshakeResponseData data;
 
-    public SaslHandshakeResponse(Errors error, Collection enabledMechanisms) {
-        this.error = error;
-        this.enabledMechanisms = new ArrayList<>(enabledMechanisms);
-    }
-
-    public SaslHandshakeResponse(Struct struct) {
-        error = Errors.forCode(struct.get(ERROR_CODE));
-        Object[] mechanisms = struct.getArray(ENABLED_MECHANISMS_KEY_NAME);
-        ArrayList enabledMechanisms = new ArrayList<>();
-        for (Object mechanism : mechanisms)
-            enabledMechanisms.add((String) mechanism);
-        this.enabledMechanisms = enabledMechanisms;
+    public SaslHandshakeResponse(SaslHandshakeResponseData data) {
+        super(ApiKeys.SASL_HANDSHAKE);
+        this.data = data;
     }
 
+    /*
+    * Possible error codes:
+    *   UNSUPPORTED_SASL_MECHANISM(33): Client mechanism not enabled in server
+    *   ILLEGAL_SASL_STATE(34) : Invalid request during SASL handshake
+    */
     public Errors error() {
-        return error;
+        return Errors.forCode(data.errorCode());
     }
 
     @Override
     public Map errorCounts() {
-        return errorCounts(error);
+        return errorCounts(Errors.forCode(data.errorCode()));
+    }
+
+    @Override
+    public int throttleTimeMs() {
+        return DEFAULT_THROTTLE_TIME;
     }
 
     @Override
-    public Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.SASL_HANDSHAKE.responseSchema(version));
-        struct.set(ERROR_CODE, error.code());
-        struct.set(ENABLED_MECHANISMS_KEY_NAME, enabledMechanisms.toArray());
-        return struct;
+    public SaslHandshakeResponseData data() {
+        return data;
     }
 
     public List enabledMechanisms() {
-        return enabledMechanisms;
+        return data.mechanisms();
     }
 
     public static SaslHandshakeResponse parse(ByteBuffer buffer, short version) {
-        return new SaslHandshakeResponse(ApiKeys.SASL_HANDSHAKE.parseResponse(version, buffer));
+        return new SaslHandshakeResponse(new SaslHandshakeResponseData(new ByteBufferAccessor(buffer), version));
     }
 }
-
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 d79c938d7941b..4326aaffd8731 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,65 +17,69 @@
 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.StopReplicaPartitionState;
+import org.apache.kafka.common.message.StopReplicaRequestData.StopReplicaPartitionV0;
+import org.apache.kafka.common.message.StopReplicaRequestData.StopReplicaTopicV1;
+import org.apache.kafka.common.message.StopReplicaRequestData.StopReplicaTopicState;
+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.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.Struct;
+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.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.BOOLEAN;
-import static org.apache.kafka.common.protocol.types.Type.INT32;
-
-public class StopReplicaRequest extends AbstractRequest {
-    private static final String CONTROLLER_ID_KEY_NAME = "controller_id";
-    private static final String CONTROLLER_EPOCH_KEY_NAME = "controller_epoch";
-    private static final String DELETE_PARTITIONS_KEY_NAME = "delete_partitions";
-    private static final String PARTITIONS_KEY_NAME = "partitions";
-
-    private static final Schema STOP_REPLICA_REQUEST_PARTITION_V0 = new Schema(
-            TOPIC_NAME,
-            PARTITION_ID);
-    private static final Schema STOP_REPLICA_REQUEST_V0 = new Schema(
-            new Field(CONTROLLER_ID_KEY_NAME, INT32, "The controller id."),
-            new Field(CONTROLLER_EPOCH_KEY_NAME, INT32, "The controller epoch."),
-            new Field(DELETE_PARTITIONS_KEY_NAME, BOOLEAN, "Boolean which indicates if replica's partitions must be deleted."),
-            new Field(PARTITIONS_KEY_NAME, new ArrayOf(STOP_REPLICA_REQUEST_PARTITION_V0)));
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {STOP_REPLICA_REQUEST_V0};
-    }
+import java.util.stream.Collectors;
+
+public class StopReplicaRequest extends AbstractControlRequest {
 
-    public static class Builder extends AbstractRequest.Builder {
-        private final int controllerId;
-        private final int controllerEpoch;
+    public static class Builder extends AbstractControlRequest.Builder {
         private final boolean deletePartitions;
-        private final Set partitions;
+        private final List topicStates;
 
-        public Builder(int controllerId, int controllerEpoch, boolean deletePartitions,
-                       Set partitions) {
-            super(ApiKeys.STOP_REPLICA);
-            this.controllerId = controllerId;
-            this.controllerEpoch = controllerEpoch;
+        public Builder(short version, int controllerId, int controllerEpoch, long brokerEpoch,
+                       boolean deletePartitions, List topicStates) {
+            super(ApiKeys.STOP_REPLICA, version, controllerId, controllerEpoch, brokerEpoch);
             this.deletePartitions = deletePartitions;
-            this.partitions = partitions;
+            this.topicStates = topicStates;
         }
 
-        @Override
         public StopReplicaRequest build(short version) {
-            return new StopReplicaRequest(controllerId, controllerEpoch,
-                    deletePartitions, partitions, version);
+            StopReplicaRequestData data = new StopReplicaRequestData()
+                .setControllerId(controllerId)
+                .setControllerEpoch(controllerEpoch)
+                .setBrokerEpoch(brokerEpoch);
+
+            if (version >= 3) {
+                data.setTopicStates(topicStates);
+            } else if (version >= 1) {
+                data.setDeletePartitions(deletePartitions);
+                List topics = topicStates.stream().map(topic ->
+                    new StopReplicaTopicV1()
+                        .setName(topic.topicName())
+                        .setPartitionIndexes(topic.partitionStates().stream()
+                            .map(StopReplicaPartitionState::partitionIndex)
+                            .collect(Collectors.toList())))
+                    .collect(Collectors.toList());
+                data.setTopics(topics);
+            } else {
+                data.setDeletePartitions(deletePartitions);
+                List partitions = topicStates.stream().flatMap(topic ->
+                    topic.partitionStates().stream().map(partition ->
+                        new StopReplicaPartitionV0()
+                            .setTopicName(topic.topicName())
+                            .setPartitionIndex(partition.partitionIndex())))
+                    .collect(Collectors.toList());
+                data.setUngroupedPartitions(partitions);
+            }
+
+            return new StopReplicaRequest(data, version);
         }
 
         @Override
@@ -84,99 +88,130 @@ public String toString() {
             bld.append("(type=StopReplicaRequest").
                 append(", controllerId=").append(controllerId).
                 append(", controllerEpoch=").append(controllerEpoch).
+                append(", brokerEpoch=").append(brokerEpoch).
                 append(", deletePartitions=").append(deletePartitions).
-                append(", partitions=").append(Utils.join(partitions, ",")).
+                append(", topicStates=").append(Utils.join(topicStates, ",")).
                 append(")");
             return bld.toString();
         }
     }
 
-    private final int controllerId;
-    private final int controllerEpoch;
-    private final boolean deletePartitions;
-    private final Set partitions;
-
-    private StopReplicaRequest(int controllerId, int controllerEpoch, boolean deletePartitions,
-                               Set partitions, short version) {
-        super(version);
-        this.controllerId = controllerId;
-        this.controllerEpoch = controllerEpoch;
-        this.deletePartitions = deletePartitions;
-        this.partitions = partitions;
-    }
-
-    public StopReplicaRequest(Struct struct, short version) {
-        super(version);
+    private final StopReplicaRequestData data;
 
-        partitions = new HashSet<>();
-        for (Object partitionDataObj : struct.getArray(PARTITIONS_KEY_NAME)) {
-            Struct partitionData = (Struct) partitionDataObj;
-            String topic = partitionData.get(TOPIC_NAME);
-            int partition = partitionData.get(PARTITION_ID);
-            partitions.add(new TopicPartition(topic, partition));
-        }
-
-        controllerId = struct.getInt(CONTROLLER_ID_KEY_NAME);
-        controllerEpoch = struct.getInt(CONTROLLER_EPOCH_KEY_NAME);
-        deletePartitions = struct.getBoolean(DELETE_PARTITIONS_KEY_NAME);
+    private StopReplicaRequest(StopReplicaRequestData data, short version) {
+        super(ApiKeys.STOP_REPLICA, version);
+        this.data = data;
     }
 
     @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 (StopReplicaTopicState topic : topicStates()) {
+            for (StopReplicaPartitionState partition : topic.partitionStates()) {
+                partitions.add(new StopReplicaPartitionError()
+                    .setTopicName(topic.topicName())
+                    .setPartitionIndex(partition.partitionIndex())
+                    .setErrorCode(error.code()));
+            }
         }
+        data.setPartitionErrors(partitions);
 
-        short versionId = version();
-        switch (versionId) {
-            case 0:
-                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()));
+        return new StopReplicaResponse(data);
+    }
+
+    /**
+     * 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 topicStates() {
+        if (version() < 1) {
+            Map topicStates = new HashMap<>();
+            for (StopReplicaPartitionV0 partition : data.ungroupedPartitions()) {
+                StopReplicaTopicState topicState = topicStates.computeIfAbsent(partition.topicName(),
+                    topic -> new StopReplicaTopicState().setTopicName(topic));
+                topicState.partitionStates().add(new StopReplicaPartitionState()
+                    .setPartitionIndex(partition.partitionIndex())
+                    .setDeletePartition(data.deletePartitions()));
+            }
+            return topicStates.values();
+        } else if (version() < 3) {
+            return () -> new MappedIterator<>(data.topics().iterator(), topic ->
+                new StopReplicaTopicState()
+                    .setTopicName(topic.name())
+                    .setPartitionStates(topic.partitionIndexes().stream()
+                        .map(partition -> new StopReplicaPartitionState()
+                            .setPartitionIndex(partition)
+                            .setDeletePartition(data.deletePartitions()))
+                        .collect(Collectors.toList())));
+        } else {
+            return data.topicStates();
         }
     }
 
-    public int controllerId() {
-        return controllerId;
+    public Map partitionStates() {
+        Map partitionStates = new HashMap<>();
+
+        if (version() < 1) {
+            for (StopReplicaPartitionV0 partition : data.ungroupedPartitions()) {
+                partitionStates.put(
+                    new TopicPartition(partition.topicName(), partition.partitionIndex()),
+                    new StopReplicaPartitionState()
+                        .setPartitionIndex(partition.partitionIndex())
+                        .setDeletePartition(data.deletePartitions()));
+            }
+        } else if (version() < 3) {
+            for (StopReplicaTopicV1 topic : data.topics()) {
+                for (Integer partitionIndex : topic.partitionIndexes()) {
+                    partitionStates.put(
+                        new TopicPartition(topic.name(), partitionIndex),
+                        new StopReplicaPartitionState()
+                            .setPartitionIndex(partitionIndex)
+                            .setDeletePartition(data.deletePartitions()));
+                }
+            }
+        } else {
+            for (StopReplicaTopicState topicState : data.topicStates()) {
+                for (StopReplicaPartitionState partitionState: topicState.partitionStates()) {
+                    partitionStates.put(
+                        new TopicPartition(topicState.topicName(), partitionState.partitionIndex()),
+                        partitionState);
+                }
+            }
+        }
+
+        return partitionStates;
     }
 
-    public int controllerEpoch() {
-        return controllerEpoch;
+    @Override
+    public int controllerId() {
+        return data.controllerId();
     }
 
-    public boolean deletePartitions() {
-        return deletePartitions;
+    @Override
+    public int controllerEpoch() {
+        return data.controllerEpoch();
     }
 
-    public Set partitions() {
-        return partitions;
+    @Override
+    public long brokerEpoch() {
+        return data.brokerEpoch();
     }
 
     public static StopReplicaRequest parse(ByteBuffer buffer, short version) {
-        return new StopReplicaRequest(ApiKeys.STOP_REPLICA.parseRequest(version, buffer), version);
+        return new StopReplicaRequest(new StopReplicaRequestData(new ByteBufferAccessor(buffer), version), version);
     }
 
     @Override
-    protected Struct toStruct() {
-        Struct struct = new Struct(ApiKeys.STOP_REPLICA.requestSchema(version()));
-
-        struct.set(CONTROLLER_ID_KEY_NAME, controllerId);
-        struct.set(CONTROLLER_EPOCH_KEY_NAME, controllerEpoch);
-        struct.set(DELETE_PARTITIONS_KEY_NAME, deletePartitions);
-
-        List partitionDatas = new ArrayList<>(partitions.size());
-        for (TopicPartition partition : partitions) {
-            Struct partitionData = struct.instance(PARTITIONS_KEY_NAME);
-            partitionData.set(TOPIC_NAME, partition.topic());
-            partitionData.set(PARTITION_ID, partition.partition());
-            partitionDatas.add(partitionData);
-        }
-
-        struct.set(PARTITIONS_KEY_NAME, partitionDatas.toArray());
-        return struct;
+    public StopReplicaRequestData data() {
+        return data;
     }
 }
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 777416d175836..10ab153f440b9 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,103 +16,68 @@
  */
 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.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.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;
-
 public class StopReplicaResponse extends AbstractResponse {
-    private static final String PARTITIONS_KEY_NAME = "partitions";
-
-    private static final Schema STOP_REPLICA_RESPONSE_PARTITION_V0 = new Schema(
-            TOPIC_NAME,
-            PARTITION_ID,
-            ERROR_CODE);
-    private static final Schema STOP_REPLICA_RESPONSE_V0 = new Schema(
-            ERROR_CODE,
-            new Field(PARTITIONS_KEY_NAME, new ArrayOf(STOP_REPLICA_RESPONSE_PARTITION_V0)));
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {STOP_REPLICA_RESPONSE_V0};
-    }
-
-    private final Map responses;
 
     /**
      * Possible error code:
-     *
-     * STALE_CONTROLLER_EPOCH (11)
+     *  - {@link Errors#STALE_CONTROLLER_EPOCH}
+     *  - {@link Errors#STALE_BROKER_EPOCH}
+     *  - {@link Errors#FENCED_LEADER_EPOCH}
+     *  - {@link Errors#KAFKA_STORAGE_ERROR}
      */
-    private final Errors error;
+    private final StopReplicaResponseData data;
 
-    public StopReplicaResponse(Errors error, Map responses) {
-        this.responses = responses;
-        this.error = error;
+    public StopReplicaResponse(StopReplicaResponseData data) {
+        super(ApiKeys.STOP_REPLICA);
+        this.data = data;
     }
 
-    public StopReplicaResponse(Struct struct) {
-        responses = new HashMap<>();
-        for (Object responseDataObj : struct.getArray(PARTITIONS_KEY_NAME)) {
-            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 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() + 1);
+        Map errors = errorCounts(data.partitionErrors().stream().map(p -> Errors.forCode(p.errorCode())));
+        updateErrorCounts(errors, Errors.forCode(data.errorCode())); // top level error
+        return errors;
     }
 
     public static StopReplicaResponse parse(ByteBuffer buffer, short version) {
-        return new StopReplicaResponse(ApiKeys.STOP_REPLICA.parseResponse(version, buffer));
+        return new StopReplicaResponse(new StopReplicaResponseData(new ByteBufferAccessor(buffer), version));
     }
 
     @Override
-    protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.STOP_REPLICA.responseSchema(version));
+    public int throttleTimeMs() {
+        return DEFAULT_THROTTLE_TIME;
+    }
 
-        List responseDatas = new ArrayList<>(responses.size());
-        for (Map.Entry response : responses.entrySet()) {
-            Struct partitionData = struct.instance(PARTITIONS_KEY_NAME);
-            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);
-        }
+    @Override
+    public StopReplicaResponseData data() {
+        return data;
+    }
 
-        struct.set(PARTITIONS_KEY_NAME, responseDatas.toArray());
-        struct.set(ERROR_CODE, error.code());
-        return struct;
+    @Override
+    public String toString() {
+        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 14ed2625d5a6a..8242b71a03fbd 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,162 +16,83 @@
  */
 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.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.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;
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {SYNC_GROUP_REQUEST_V0, SYNC_GROUP_REQUEST_V1};
-    }
 
     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) {
-        super(version);
-        this.groupId = groupId;
-        this.generationId = generationId;
-        this.memberId = memberId;
-        this.groupAssignment = groupAssignment;
-    }
-
-    public SyncGroupRequest(Struct struct, short version) {
-        super(version);
-        this.groupId = struct.get(GROUP_ID);
-        this.generationId = struct.get(GENERATION_ID);
-        this.memberId = struct.get(MEMBER_ID);
 
-        groupAssignment = new HashMap<>();
+    private final SyncGroupRequestData data;
 
-        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);
-        }
+    public SyncGroupRequest(SyncGroupRequestData data, short version) {
+        super(ApiKeys.SYNC_GROUP, version);
+        this.data = data;
     }
 
     @Override
     public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
-        short versionId = version();
-        switch (versionId) {
-            case 0:
-                return new SyncGroupResponse(
-                        Errors.forException(e),
-                        ByteBuffer.wrap(new byte[]{}));
-            case 1:
-                return new SyncGroupResponse(
-                        throttleTimeMs,
-                        Errors.forException(e),
-                        ByteBuffer.wrap(new byte[]{}));
-            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()));
-        }
+        return new SyncGroupResponse(new SyncGroupResponseData()
+                .setErrorCode(Errors.forException(e).code())
+                .setAssignment(new byte[0])
+                .setThrottleTimeMs(throttleTimeMs));
     }
 
-    public String groupId() {
-        return groupId;
-    }
-
-    public int generationId() {
-        return generationId;
-    }
-
-    public Map groupAssignment() {
-        return groupAssignment;
+    public Map groupAssignments() {
+        Map groupAssignments = new HashMap<>();
+        for (SyncGroupRequestData.SyncGroupRequestAssignment assignment : data.assignments()) {
+            groupAssignments.put(assignment.memberId(), ByteBuffer.wrap(assignment.assignment()));
+        }
+        return groupAssignments;
     }
 
-    public String memberId() {
-        return memberId;
+    /**
+     * ProtocolType and ProtocolName are mandatory since version 5. This methods verifies that
+     * they are defined for version 5 or higher, or returns true otherwise for older versions.
+     */
+    public boolean areMandatoryProtocolTypeAndNamePresent() {
+        if (version() >= 5)
+            return data.protocolType() != null && data.protocolName() != null;
+        else
+            return true;
     }
 
     public static SyncGroupRequest parse(ByteBuffer buffer, short version) {
-        return new SyncGroupRequest(ApiKeys.SYNC_GROUP.parseRequest(version, buffer), version);
+        return new SyncGroupRequest(new SyncGroupRequestData(new ByteBufferAccessor(buffer), version), 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;
+    public SyncGroupRequestData data() {
+        return data;
     }
 }
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 77f951251fc63..822a3e78b9949 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,93 +16,53 @@
  */
 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.ByteBufferAccessor;
 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.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));
-
-    public static Schema[] schemaVersions() {
-        return new Schema[] {SYNC_GROUP_RESPONSE_V0, SYNC_GROUP_RESPONSE_V1};
-    }
 
-    /**
-     * Possible error codes:
-     *
-     * GROUP_COORDINATOR_NOT_AVAILABLE (15)
-     * NOT_COORDINATOR (16)
-     * ILLEGAL_GENERATION (22)
-     * UNKNOWN_MEMBER_ID (25)
-     * REBALANCE_IN_PROGRESS (27)
-     * GROUP_AUTHORIZATION_FAILED (30)
-     */
-
-    private final Errors error;
-    private final int throttleTimeMs;
-    private final ByteBuffer memberState;
-
-    public SyncGroupResponse(Errors error, ByteBuffer memberState) {
-        this(DEFAULT_THROTTLE_TIME, error, memberState);
-    }
+    private final SyncGroupResponseData data;
 
-    public SyncGroupResponse(int throttleTimeMs, Errors error, ByteBuffer memberState) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.error = error;
-        this.memberState = memberState;
-    }
-
-    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(SyncGroupResponseData data) {
+        super(ApiKeys.SYNC_GROUP);
+        this.data = data;
     }
 
+    @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 errorCounts(Errors.forCode(data.errorCode()));
     }
 
-    public ByteBuffer memberAssignment() {
-        return memberState;
+    @Override
+    public SyncGroupResponseData data() {
+        return data;
     }
 
     @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;
+    public String toString() {
+        return data.toString();
     }
 
     public static SyncGroupResponse parse(ByteBuffer buffer, short version) {
-        return new SyncGroupResponse(ApiKeys.SYNC_GROUP.parseResponse(version, buffer));
+        return new SyncGroupResponse(new SyncGroupResponseData(new ByteBufferAccessor(buffer), version));
     }
 
+    @Override
+    public boolean shouldClientThrottle(short version) {
+        return version >= 2;
+    }
 }
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 9787c2df0a897..e96f81ae7e2f9 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,228 @@
 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.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.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.Struct;
-import org.apache.kafka.common.utils.CollectionUtils;
+import org.apache.kafka.common.record.RecordBatch;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 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.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 static org.apache.kafka.common.protocol.types.Type.INT64;
-import static org.apache.kafka.common.protocol.types.Type.NULLABLE_STRING;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.stream.Collectors;
 
 public class TxnOffsetCommitRequest extends AbstractRequest {
-    private static final String TOPICS_KEY_NAME = "topics";
-    private static final String PARTITIONS_KEY_NAME = "partitions";
-    private static final String OFFSET_KEY_NAME = "offset";
-    private static final String METADATA_KEY_NAME = "metadata";
-
-    private static final Schema TXN_OFFSET_COMMIT_PARTITION_OFFSET_METADATA_REQUEST_V0 = new Schema(
-            PARTITION_ID,
-            new Field(OFFSET_KEY_NAME, INT64),
-            new Field(METADATA_KEY_NAME, NULLABLE_STRING));
-
-    private static final Schema TXN_OFFSET_COMMIT_REQUEST_V0 = new Schema(
-            TRANSACTIONAL_ID,
-            GROUP_ID,
-            PRODUCER_ID,
-            PRODUCER_EPOCH,
-            new Field(TOPICS_KEY_NAME, new ArrayOf(new Schema(
-                    TOPIC_NAME,
-                    new Field(PARTITIONS_KEY_NAME, new ArrayOf(TXN_OFFSET_COMMIT_PARTITION_OFFSET_METADATA_REQUEST_V0)))),
-                    "The partitions to write markers for."));
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{TXN_OFFSET_COMMIT_REQUEST_V0};
-    }
+
+    private static final Logger log = LoggerFactory.getLogger(TxnOffsetCommitRequest.class);
+
+    private 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;
+
+        private final boolean autoDowngrade;
+
+        public Builder(final String transactionalId,
+                       final String consumerGroupId,
+                       final long producerId,
+                       final short producerEpoch,
+                       final Map pendingTxnOffsetCommits,
+                       final boolean autoDowngrade) {
+            this(transactionalId,
+                consumerGroupId,
+                producerId,
+                producerEpoch,
+                pendingTxnOffsetCommits,
+                JoinGroupRequest.UNKNOWN_MEMBER_ID,
+                JoinGroupRequest.UNKNOWN_GENERATION_ID,
+                Optional.empty(),
+                autoDowngrade);
         }
 
-        public Map offsets() {
-            return offsets;
+        public Builder(final String transactionalId,
+                       final String consumerGroupId,
+                       final long producerId,
+                       final short producerEpoch,
+                       final Map pendingTxnOffsetCommits,
+                       final String memberId,
+                       final int generationId,
+                       final Optional groupInstanceId,
+                       final boolean autoDowngrade) {
+            super(ApiKeys.TXN_OFFSET_COMMIT);
+            this.data = new TxnOffsetCommitRequestData()
+                            .setTransactionalId(transactionalId)
+                            .setGroupId(consumerGroupId)
+                            .setProducerId(producerId)
+                            .setProducerEpoch(producerEpoch)
+                            .setTopics(getTopics(pendingTxnOffsetCommits))
+                            .setMemberId(memberId)
+                            .setGenerationId(generationId)
+                            .setGroupInstanceId(groupInstanceId.orElse(null));
+            this.autoDowngrade = autoDowngrade;
         }
 
         @Override
         public TxnOffsetCommitRequest build(short version) {
-            return new TxnOffsetCommitRequest(version, transactionalId, consumerGroupId, producerId, producerEpoch, offsets);
+            if (version < 3 && groupMetadataSet()) {
+                if (autoDowngrade) {
+                    log.trace("Downgrade the request by resetting group metadata fields: " +
+                                  "[member.id:{}, generation.id:{}, group.instance.id:{}], because broker " +
+                                  "only supports TxnOffsetCommit version {}. Need " +
+                                  "v3 or newer to enable this feature",
+                        data.memberId(), data.generationId(), data.groupInstanceId(), version);
+
+                    data.setGenerationId(JoinGroupRequest.UNKNOWN_GENERATION_ID)
+                        .setMemberId(JoinGroupRequest.UNKNOWN_MEMBER_ID)
+                        .setGroupInstanceId(null);
+                } else {
+                    throw new UnsupportedVersionException("Broker unexpectedly " +
+                        "doesn't support group metadata commit API on version " + version);
+                }
+            }
+            return new TxnOffsetCommitRequest(data, version);
+        }
+
+        private boolean groupMetadataSet() {
+            return !data.memberId().equals(JoinGroupRequest.UNKNOWN_MEMBER_ID) ||
+                       data.generationId() != JoinGroupRequest.UNKNOWN_GENERATION_ID ||
+                       data.groupInstanceId() != null;
         }
 
         @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) {
-        super(version);
-        this.transactionalId = transactionalId;
-        this.consumerGroupId = consumerGroupId;
-        this.producerId = producerId;
-        this.producerEpoch = producerEpoch;
-        this.offsets = offsets;
+    public TxnOffsetCommitRequest(TxnOffsetCommitRequestData data, short version) {
+        super(ApiKeys.TXN_OFFSET_COMMIT, version);
+        this.data = data;
     }
 
-    public TxnOffsetCommitRequest(Struct struct, short version) {
-        super(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);
-
-        Map offsets = new HashMap<>();
-        Object[] topicPartitionsArray = struct.getArray(TOPICS_KEY_NAME);
-        for (Object topicPartitionObj : topicPartitionsArray) {
-            Struct topicPartitionStruct = (Struct) topicPartitionObj;
-            String topic = topicPartitionStruct.get(TOPIC_NAME);
-            for (Object partitionObj : topicPartitionStruct.getArray(PARTITIONS_KEY_NAME)) {
-                Struct partitionStruct = (Struct) partitionObj;
-                TopicPartition partition = new TopicPartition(topic, partitionStruct.get(PARTITION_ID));
-                long offset = partitionStruct.getLong(OFFSET_KEY_NAME);
-                String metadata = partitionStruct.getString(METADATA_KEY_NAME);
-                offsets.put(partition, new CommittedOffset(offset, metadata));
+    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;
-    }
-
-    public String transactionalId() {
-        return transactionalId;
-    }
-
-    public String consumerGroupId() {
-        return consumerGroupId;
-    }
-
-    public long producerId() {
-        return producerId;
+        return offsetMap;
     }
 
-    public short producerEpoch() {
-        return producerEpoch;
+    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());
     }
 
-    public Map offsets() {
-        return offsets;
+    @Override
+    public TxnOffsetCommitRequestData data() {
+        return data;
     }
 
-    @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.groupDataByTopic(offsets);
-        Object[] partitionsArray = new Object[mappedPartitionOffsets.size()];
-        int i = 0;
-        for (Map.Entry> topicAndPartitions : mappedPartitionOffsets.entrySet()) {
-            Struct topicPartitionsStruct = struct.instance(TOPICS_KEY_NAME);
-            topicPartitionsStruct.set(TOPIC_NAME, topicAndPartitions.getKey());
-
-            Map partitionOffsets = topicAndPartitions.getValue();
-            Object[] partitionOffsetsArray = new Object[partitionOffsets.size()];
-            int j = 0;
-            for (Map.Entry partitionOffset : partitionOffsets.entrySet()) {
-                Struct partitionOffsetStruct = topicPartitionsStruct.instance(PARTITIONS_KEY_NAME);
-                partitionOffsetStruct.set(PARTITION_ID, partitionOffset.getKey());
-                CommittedOffset committedOffset = partitionOffset.getValue();
-                partitionOffsetStruct.set(OFFSET_KEY_NAME, committedOffset.offset);
-                partitionOffsetStruct.set(METADATA_KEY_NAME, committedOffset.metadata);
-                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_KEY_NAME, partitionOffsetsArray);
-            partitionsArray[i++] = topicPartitionsStruct;
+            responseTopicData.add(new TxnOffsetCommitResponseTopic()
+                                      .setName(entry.name())
+                                      .setPartitions(responsePartitions)
+            );
         }
-
-        struct.set(TOPICS_KEY_NAME, 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) {
-        return new TxnOffsetCommitRequest(ApiKeys.TXN_OFFSET_COMMIT.parseRequest(version, buffer), version);
+        return new TxnOffsetCommitRequest(new TxnOffsetCommitRequestData(
+            new ByteBufferAccessor(buffer), version), version);
     }
 
     public static class CommittedOffset {
-        private final long offset;
-        private final String metadata;
+        public final long offset;
+        public final String metadata;
+        public final Optional leaderEpoch;
 
-        public CommittedOffset(long offset, String metadata) {
+        public CommittedOffset(long offset, String metadata, Optional leaderEpoch) {
             this.offset = offset;
             this.metadata = metadata;
+            this.leaderEpoch = leaderEpoch;
         }
 
         @Override
         public String toString() {
             return "CommittedOffset(" +
                     "offset=" + offset +
+                    ", leaderEpoch=" + leaderEpoch +
                     ", metadata='" + metadata + "')";
         }
 
-        public long offset() {
-            return offset;
+        @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);
         }
 
-        public String metadata() {
-            return 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 ff0f8cefeefeb..b4de54741e6d0 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,129 +17,105 @@
 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.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.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:
+ *
+ *   - {@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}
+ *   - {@link Errors#UNKNOWN_MEMBER_ID}
+ *   - {@link Errors#FENCED_INSTANCE_ID}
+ *   - {@link Errors#ILLEGAL_GENERATION}
+ */
 public class TxnOffsetCommitResponse extends AbstractResponse {
-    private static final String TOPICS_KEY_NAME = "topics";
-    private static final String PARTITIONS_KEY_NAME = "partitions";
-
-    private static final Schema TXN_OFFSET_COMMIT_PARTITION_ERROR_RESPONSE_V0 = new Schema(
-            PARTITION_ID,
-            ERROR_CODE);
-
-    private static final Schema TXN_OFFSET_COMMIT_RESPONSE_V0 = new Schema(
-            THROTTLE_TIME_MS,
-            new Field(TOPICS_KEY_NAME, new ArrayOf(new Schema(
-                    TOPIC_NAME,
-                    new Field(PARTITIONS_KEY_NAME, new ArrayOf(TXN_OFFSET_COMMIT_PARTITION_ERROR_RESPONSE_V0)))),
-                    "Errors per partition from writing markers."));
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{TXN_OFFSET_COMMIT_RESPONSE_V0};
-    }
 
-    // Possible error codes:
-    //   InvalidProducerEpoch
-    //   NotCoordinator
-    //   CoordinatorNotAvailable
-    //   CoordinatorLoadInProgress
-    //   OffsetMetadataTooLarge
-    //   GroupAuthorizationFailed
-    //   InvalidCommitOffsetSize
-    //   TransactionalIdAuthorizationFailed
-    //   RequestTimedOut
-
-    private final Map errors;
-    private final int throttleTimeMs;
-
-    public TxnOffsetCommitResponse(int throttleTimeMs, Map errors) {
-        this.throttleTimeMs = throttleTimeMs;
-        this.errors = errors;
-    }
+    private final TxnOffsetCommitResponseData data;
 
-    public TxnOffsetCommitResponse(Struct struct) {
-        this.throttleTimeMs = struct.get(THROTTLE_TIME_MS);
-        Map errors = new HashMap<>();
-        Object[] topicPartitionsArray = struct.getArray(TOPICS_KEY_NAME);
-        for (Object topicPartitionObj : topicPartitionsArray) {
-            Struct topicPartitionStruct = (Struct) topicPartitionObj;
-            String topic = topicPartitionStruct.get(TOPIC_NAME);
-            for (Object partitionObj : topicPartitionStruct.getArray(PARTITIONS_KEY_NAME)) {
-                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);
-            }
-        }
-        this.errors = errors;
+    public TxnOffsetCommitResponse(TxnOffsetCommitResponseData data) {
+        super(ApiKeys.TXN_OFFSET_COMMIT);
+        this.data = data;
     }
 
-    @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.groupDataByTopic(errors);
-        Object[] partitionsArray = new Object[mappedPartitions.size()];
-        int i = 0;
-        for (Map.Entry> topicAndPartitions : mappedPartitions.entrySet()) {
-            Struct topicPartitionsStruct = struct.instance(TOPICS_KEY_NAME);
-            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_KEY_NAME);
-                partitionAndErrorStruct.set(PARTITION_ID, partitionAndError.getKey());
-                partitionAndErrorStruct.set(ERROR_CODE, partitionAndError.getValue().code());
-                partitionAndErrorsArray[j++] = partitionAndErrorStruct;
-            }
-            topicPartitionsStruct.set(PARTITIONS_KEY_NAME, partitionAndErrorsArray);
-            partitionsArray[i++] = topicPartitionsStruct;
+    public TxnOffsetCommitResponse(int requestThrottleMs, Map responseData) {
+        super(ApiKeys.TXN_OFFSET_COMMIT);
+        Map responseTopicDataMap = new HashMap<>();
+
+        for (Map.Entry entry : responseData.entrySet()) {
+            TopicPartition topicPartition = entry.getKey();
+            String topicName = topicPartition.topic();
+
+            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);
         }
 
-        struct.set(TOPICS_KEY_NAME, partitionsArray);
-        return struct;
+        data = new TxnOffsetCommitResponseData()
+                   .setTopics(new ArrayList<>(responseTopicDataMap.values()))
+                   .setThrottleTimeMs(requestThrottleMs);
     }
 
-    public int throttleTimeMs() {
-        return throttleTimeMs;
+    @Override
+    public TxnOffsetCommitResponseData data() {
+        return data;
     }
 
-    public Map errors() {
-        return errors;
+    @Override
+    public int throttleTimeMs() {
+        return data.throttleTimeMs();
     }
 
     @Override
     public Map errorCounts() {
-        return errorCounts(errors);
+        return errorCounts(data.topics().stream().flatMap(topic ->
+                topic.partitions().stream().map(partition ->
+                        Errors.forCode(partition.errorCode()))));
+    }
+
+    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(new TxnOffsetCommitResponseData(new ByteBufferAccessor(buffer), version));
     }
 
     @Override
     public String toString() {
-        return "TxnOffsetCommitResponse(" +
-                "errors=" + errors +
-                ", throttleTimeMs=" + throttleTimeMs +
-                ')';
+        return data.toString();
     }
 
+    @Override
+    public boolean shouldClientThrottle(short version) {
+        return version >= 1;
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/UpdateFeaturesRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/UpdateFeaturesRequest.java
new file mode 100644
index 0000000000000..fc1892475001d
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/UpdateFeaturesRequest.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.requests;
+
+import java.nio.ByteBuffer;
+import org.apache.kafka.common.message.UpdateFeaturesRequestData.FeatureUpdateKey;
+import org.apache.kafka.common.message.UpdateFeaturesResponseData;
+import org.apache.kafka.common.message.UpdateFeaturesRequestData;
+import org.apache.kafka.common.message.UpdateFeaturesResponseData.UpdatableFeatureResult;
+import org.apache.kafka.common.message.UpdateFeaturesResponseData.UpdatableFeatureResultCollection;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+
+public class UpdateFeaturesRequest extends AbstractRequest {
+
+    public static class Builder extends AbstractRequest.Builder {
+
+        private final UpdateFeaturesRequestData data;
+
+        public Builder(UpdateFeaturesRequestData data) {
+            super(ApiKeys.UPDATE_FEATURES);
+            this.data = data;
+        }
+
+        @Override
+        public UpdateFeaturesRequest build(short version) {
+            return new UpdateFeaturesRequest(data, version);
+        }
+
+        @Override
+        public String toString() {
+            return data.toString();
+        }
+    }
+
+    private final UpdateFeaturesRequestData data;
+
+    public UpdateFeaturesRequest(UpdateFeaturesRequestData data, short version) {
+        super(ApiKeys.UPDATE_FEATURES, version);
+        this.data = data;
+    }
+
+    @Override
+    public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
+        final ApiError apiError = ApiError.fromThrowable(e);
+        final UpdatableFeatureResultCollection results = new UpdatableFeatureResultCollection();
+        for (FeatureUpdateKey update : this.data.featureUpdates().valuesSet()) {
+            final UpdatableFeatureResult result = new UpdatableFeatureResult()
+                .setFeature(update.feature())
+                .setErrorCode(apiError.error().code())
+                .setErrorMessage(apiError.message());
+            results.add(result);
+        }
+        final UpdateFeaturesResponseData responseData = new UpdateFeaturesResponseData()
+            .setThrottleTimeMs(throttleTimeMs)
+            .setResults(results);
+        return new UpdateFeaturesResponse(responseData);
+    }
+
+    @Override
+    public UpdateFeaturesRequestData data() {
+        return data;
+    }
+
+    public static UpdateFeaturesRequest parse(ByteBuffer buffer, short version) {
+        return new UpdateFeaturesRequest(new UpdateFeaturesRequestData(new ByteBufferAccessor(buffer), version), version);
+    }
+
+    public static boolean isDeleteRequest(UpdateFeaturesRequestData.FeatureUpdateKey update) {
+        return update.maxVersionLevel() < 1 && update.allowDowngrade();
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/UpdateFeaturesResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/UpdateFeaturesResponse.java
new file mode 100644
index 0000000000000..028f8855ce5c7
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/UpdateFeaturesResponse.java
@@ -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 org.apache.kafka.common.requests;
+
+import java.nio.ByteBuffer;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+import org.apache.kafka.common.message.UpdateFeaturesResponseData;
+import org.apache.kafka.common.message.UpdateFeaturesResponseData.UpdatableFeatureResult;
+import org.apache.kafka.common.message.UpdateFeaturesResponseData.UpdatableFeatureResultCollection;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.Errors;
+
+
+/**
+ * Possible error codes:
+ *
+ *   - {@link Errors#CLUSTER_AUTHORIZATION_FAILED}
+ *   - {@link Errors#NOT_CONTROLLER}
+ *   - {@link Errors#INVALID_REQUEST}
+ *   - {@link Errors#FEATURE_UPDATE_FAILED}
+ */
+public class UpdateFeaturesResponse extends AbstractResponse {
+
+    private final UpdateFeaturesResponseData data;
+
+    public UpdateFeaturesResponse(UpdateFeaturesResponseData data) {
+        super(ApiKeys.UPDATE_FEATURES);
+        this.data = data;
+    }
+
+    public Map errors() {
+        return data.results().valuesSet().stream().collect(
+            Collectors.toMap(
+                result -> result.feature(),
+                result -> new ApiError(Errors.forCode(result.errorCode()), result.errorMessage())));
+    }
+
+    @Override
+    public Map errorCounts() {
+        return apiErrorCounts(errors());
+    }
+
+    @Override
+    public int throttleTimeMs() {
+        return data.throttleTimeMs();
+    }
+
+    @Override
+    public String toString() {
+        return data.toString();
+    }
+
+    @Override
+    public UpdateFeaturesResponseData data() {
+        return data;
+    }
+
+    public static UpdateFeaturesResponse parse(ByteBuffer buffer, short version) {
+        return new UpdateFeaturesResponse(new UpdateFeaturesResponseData(new ByteBufferAccessor(buffer), version));
+    }
+
+    public static UpdateFeaturesResponse createWithErrors(ApiError topLevelError, Map updateErrors, int throttleTimeMs) {
+        final UpdatableFeatureResultCollection results = new UpdatableFeatureResultCollection();
+        for (final Map.Entry updateError : updateErrors.entrySet()) {
+            final String feature = updateError.getKey();
+            final ApiError error = updateError.getValue();
+            final UpdatableFeatureResult result = new UpdatableFeatureResult();
+            result.setFeature(feature)
+                .setErrorCode(error.error().code())
+                .setErrorMessage(error.message());
+            results.add(result);
+        }
+        final UpdateFeaturesResponseData responseData = new UpdateFeaturesResponseData()
+            .setThrottleTimeMs(throttleTimeMs)
+            .setErrorCode(topLevelError.error().code())
+            .setErrorMessage(topLevelError.message())
+            .setResults(results)
+            .setThrottleTimeMs(throttleTimeMs);
+        return new UpdateFeaturesResponse(responseData);
+    }
+}
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 da48e9fa85b02..2d443b9b6a20b 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,183 +16,91 @@
  */
 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.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.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.Struct;
 import org.apache.kafka.common.security.auth.SecurityProtocol;
+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.INT16;
-import static org.apache.kafka.common.protocol.types.Type.INT32;
-import static org.apache.kafka.common.protocol.types.Type.NULLABLE_STRING;
-import static org.apache.kafka.common.protocol.types.Type.STRING;
+import static java.util.Collections.singletonList;
 
-public class UpdateMetadataRequest extends AbstractRequest {
+public class UpdateMetadataRequest extends AbstractControlRequest {
 
-    private static final String CONTROLLER_ID_KEY_NAME = "controller_id";
-    private static final String CONTROLLER_EPOCH_KEY_NAME = "controller_epoch";
-    private static final String PARTITION_STATES_KEY_NAME = "partition_states";
-    private static final String LIVE_BROKERS_KEY_NAME = "live_brokers";
+    public static class Builder extends AbstractControlRequest.Builder {
+        private final List partitionStates;
+        private final List liveBrokers;
 
-    // PartitionState key names
-    private static final String LEADER_KEY_NAME = "leader";
-    private static final String LEADER_EPOCH_KEY_NAME = "leader_epoch";
-    private static final String ISR_KEY_NAME = "isr";
-    private static final String ZK_VERSION_KEY_NAME = "zk_version";
-    private static final String REPLICAS_KEY_NAME = "replicas";
-    private static final String OFFLINE_REPLICAS_KEY_NAME = "offline_replicas";
-
-    // Broker key names
-    private static final String BROKER_ID_KEY_NAME = "id";
-    private static final String ENDPOINTS_KEY_NAME = "end_points";
-    private static final String RACK_KEY_NAME = "rack";
-
-    // EndPoint key names
-    private static final String HOST_KEY_NAME = "host";
-    private static final String PORT_KEY_NAME = "port";
-    private static final String LISTENER_NAME_KEY_NAME = "listener_name";
-    private static final String SECURITY_PROTOCOL_TYPE_KEY_NAME = "security_protocol_type";
-
-    private static final Schema UPDATE_METADATA_REQUEST_PARTITION_STATE_V0 = new Schema(
-                    TOPIC_NAME,
-                    PARTITION_ID,
-                    new Field(CONTROLLER_EPOCH_KEY_NAME, INT32, "The controller epoch."),
-                    new Field(LEADER_KEY_NAME, INT32, "The broker id for the leader."),
-                    new Field(LEADER_EPOCH_KEY_NAME, INT32, "The leader epoch."),
-                    new Field(ISR_KEY_NAME, new ArrayOf(INT32), "The in sync replica ids."),
-                    new Field(ZK_VERSION_KEY_NAME, INT32, "The ZK version."),
-                    new Field(REPLICAS_KEY_NAME, new ArrayOf(INT32), "The replica ids."));
-
-    private static final Schema UPDATE_METADATA_REQUEST_BROKER_V0 = new Schema(
-            new Field(BROKER_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 UPDATE_METADATA_REQUEST_V0 = new Schema(
-            new Field(CONTROLLER_ID_KEY_NAME, INT32, "The controller id."),
-            new Field(CONTROLLER_EPOCH_KEY_NAME, INT32, "The controller epoch."),
-            new Field(PARTITION_STATES_KEY_NAME, new ArrayOf(UPDATE_METADATA_REQUEST_PARTITION_STATE_V0)),
-            new Field(LIVE_BROKERS_KEY_NAME, new ArrayOf(UPDATE_METADATA_REQUEST_BROKER_V0)));
-
-    private static final Schema UPDATE_METADATA_REQUEST_PARTITION_STATE_V1 = UPDATE_METADATA_REQUEST_PARTITION_STATE_V0;
-
-    // for some reason, V1 sends `port` before `host` while V0 sends `host` before `port
-    private static final Schema UPDATE_METADATA_REQUEST_END_POINT_V1 = new Schema(
-            new Field(PORT_KEY_NAME, INT32, "The port on which the broker accepts requests."),
-            new Field(HOST_KEY_NAME, STRING, "The hostname of the broker."),
-            new Field(SECURITY_PROTOCOL_TYPE_KEY_NAME, INT16, "The security protocol type."));
-
-    private static final Schema UPDATE_METADATA_REQUEST_BROKER_V1 = new Schema(
-            new Field(BROKER_ID_KEY_NAME, INT32, "The broker id."),
-            new Field(ENDPOINTS_KEY_NAME, new ArrayOf(UPDATE_METADATA_REQUEST_END_POINT_V1)));
-
-    private static final Schema UPDATE_METADATA_REQUEST_V1 = new Schema(
-            new Field(CONTROLLER_ID_KEY_NAME, INT32, "The controller id."),
-            new Field(CONTROLLER_EPOCH_KEY_NAME, INT32, "The controller epoch."),
-            new Field(PARTITION_STATES_KEY_NAME, new ArrayOf(UPDATE_METADATA_REQUEST_PARTITION_STATE_V1)),
-            new Field(LIVE_BROKERS_KEY_NAME, new ArrayOf(UPDATE_METADATA_REQUEST_BROKER_V1)));
-
-    private static final Schema UPDATE_METADATA_REQUEST_PARTITION_STATE_V2 = UPDATE_METADATA_REQUEST_PARTITION_STATE_V1;
-
-    private static final Schema UPDATE_METADATA_REQUEST_END_POINT_V2 = UPDATE_METADATA_REQUEST_END_POINT_V1;
-
-    private static final Schema UPDATE_METADATA_REQUEST_BROKER_V2 = new Schema(
-            new Field(BROKER_ID_KEY_NAME, INT32, "The broker id."),
-            new Field(ENDPOINTS_KEY_NAME, new ArrayOf(UPDATE_METADATA_REQUEST_END_POINT_V2)),
-            new Field(RACK_KEY_NAME, NULLABLE_STRING, "The rack"));
-
-    private static final Schema UPDATE_METADATA_REQUEST_V2 = new Schema(
-            new Field(CONTROLLER_ID_KEY_NAME, INT32, "The controller id."),
-            new Field(CONTROLLER_EPOCH_KEY_NAME, INT32, "The controller epoch."),
-            new Field(PARTITION_STATES_KEY_NAME, new ArrayOf(UPDATE_METADATA_REQUEST_PARTITION_STATE_V2)),
-            new Field(LIVE_BROKERS_KEY_NAME, new ArrayOf(UPDATE_METADATA_REQUEST_BROKER_V2)));
-
-    private static final Schema UPDATE_METADATA_REQUEST_PARTITION_STATE_V3 = UPDATE_METADATA_REQUEST_PARTITION_STATE_V2;
-
-    // UPDATE_METADATA_REQUEST_PARTITION_STATE_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_PARTITION_STATE_V4 = new Schema(
-            TOPIC_NAME,
-            PARTITION_ID,
-            new Field(CONTROLLER_EPOCH_KEY_NAME, INT32, "The controller epoch."),
-            new Field(LEADER_KEY_NAME, INT32, "The broker id for the leader."),
-            new Field(LEADER_EPOCH_KEY_NAME, INT32, "The leader epoch."),
-            new Field(ISR_KEY_NAME, new ArrayOf(INT32), "The in sync replica ids."),
-            new Field(ZK_VERSION_KEY_NAME, INT32, "The ZK version."),
-            new Field(REPLICAS_KEY_NAME, new ArrayOf(INT32), "The replica ids."),
-            new Field(OFFLINE_REPLICAS_KEY_NAME, new ArrayOf(INT32), "The offline replica ids"));
-
-    private static final Schema UPDATE_METADATA_REQUEST_END_POINT_V3 = new Schema(
-            new Field(PORT_KEY_NAME, INT32, "The port on which the broker accepts requests."),
-            new Field(HOST_KEY_NAME, STRING, "The hostname of the broker."),
-            new Field(LISTENER_NAME_KEY_NAME, STRING, "The listener name."),
-            new Field(SECURITY_PROTOCOL_TYPE_KEY_NAME, INT16, "The security protocol type."));
-
-    private static final Schema UPDATE_METADATA_REQUEST_BROKER_V3 = new Schema(
-            new Field(BROKER_ID_KEY_NAME, INT32, "The broker id."),
-            new Field(ENDPOINTS_KEY_NAME, new ArrayOf(UPDATE_METADATA_REQUEST_END_POINT_V3)),
-            new Field(RACK_KEY_NAME, NULLABLE_STRING, "The rack"));
-
-    private static final Schema UPDATE_METADATA_REQUEST_V3 = new Schema(
-            new Field(CONTROLLER_ID_KEY_NAME, INT32, "The controller id."),
-            new Field(CONTROLLER_EPOCH_KEY_NAME, INT32, "The controller epoch."),
-            new Field(PARTITION_STATES_KEY_NAME, new ArrayOf(UPDATE_METADATA_REQUEST_PARTITION_STATE_V3)),
-            new Field(LIVE_BROKERS_KEY_NAME, new ArrayOf(UPDATE_METADATA_REQUEST_BROKER_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(
-            new Field(CONTROLLER_ID_KEY_NAME, INT32, "The controller id."),
-            new Field(CONTROLLER_EPOCH_KEY_NAME, INT32, "The controller epoch."),
-            new Field(PARTITION_STATES_KEY_NAME, new ArrayOf(UPDATE_METADATA_REQUEST_PARTITION_STATE_V4)),
-            new Field(LIVE_BROKERS_KEY_NAME, new ArrayOf(UPDATE_METADATA_REQUEST_BROKER_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};
-    }
-
-    public static class Builder extends AbstractRequest.Builder {
-        private final int controllerId;
-        private final int controllerEpoch;
-        private final Map partitionStates;
-        private final Set liveBrokers;
-
-        public Builder(short version, int controllerId, int controllerEpoch,
-                       Map partitionStates, Set liveBrokers) {
-            super(ApiKeys.UPDATE_METADATA, version);
-            this.controllerId = controllerId;
-            this.controllerEpoch = controllerEpoch;
+        public Builder(short version, int controllerId, int controllerEpoch, long brokerEpoch,
+                       List partitionStates, List liveBrokers) {
+            super(ApiKeys.UPDATE_METADATA, version, controllerId, controllerEpoch, brokerEpoch);
             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");
+            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, partitionStates, liveBrokers);
+
+            UpdateMetadataRequestData data = new UpdateMetadataRequestData()
+                    .setControllerId(controllerId)
+                    .setControllerEpoch(controllerEpoch)
+                    .setBrokerEpoch(brokerEpoch)
+                    .setLiveBrokers(liveBrokers);
+
+            if (version >= 5) {
+                Map topicStatesMap = groupByTopic(partitionStates);
+                data.setTopicStates(new ArrayList<>(topicStatesMap.values()));
+            } else {
+                data.setUngroupedPartitionStates(partitionStates);
+            }
+
+            return new UpdateMetadataRequest(data, version);
+        }
+
+        private static Map groupByTopic(List partitionStates) {
+            Map topicStates = new HashMap<>();
+            for (UpdateMetadataPartitionState partition : partitionStates) {
+                // We don't null out the topic name in UpdateMetadataPartitionState 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
@@ -201,6 +109,7 @@ public String toString() {
             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(")");
@@ -208,254 +117,95 @@ public String 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;
-        }
-
-        @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()) + ")";
-        }
-    }
-
-    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;
-        }
-
-        @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 static final class EndPoint {
-        public final String host;
-        public final int port;
-        public final SecurityProtocol securityProtocol;
-        public final ListenerName listenerName; // introduced in V3
-
-        public EndPoint(String host, int port, SecurityProtocol securityProtocol, ListenerName listenerName) {
-            this.host = host;
-            this.port = port;
-            this.securityProtocol = securityProtocol;
-            this.listenerName = listenerName;
-        }
-
-        @Override
-        public String toString() {
-            return "(host=" + host + ", port=" + port + ", listenerName=" + listenerName +
-                    ", securityProtocol=" + securityProtocol + ")";
-        }
-    }
-
-    private final int controllerId;
-    private final int controllerEpoch;
-    private final Map partitionStates;
-    private final Set liveBrokers;
+    private final UpdateMetadataRequestData data;
 
-    private UpdateMetadataRequest(short version, int controllerId, int controllerEpoch,
-                                  Map partitionStates, Set liveBrokers) {
-        super(version);
-        this.controllerId = controllerId;
-        this.controllerEpoch = controllerEpoch;
-        this.partitionStates = partitionStates;
-        this.liveBrokers = liveBrokers;
+    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();
     }
 
-    public UpdateMetadataRequest(Struct struct, short versionId) {
-        super(versionId);
-        Map partitionStates = new HashMap<>();
-        for (Object partitionStateDataObj : struct.getArray(PARTITION_STATES_KEY_NAME)) {
-            Struct partitionStateData = (Struct) partitionStateDataObj;
-            String topic = partitionStateData.get(TOPIC_NAME);
-            int partition = partitionStateData.get(PARTITION_ID);
-            int controllerEpoch = partitionStateData.getInt(CONTROLLER_EPOCH_KEY_NAME);
-            int leader = partitionStateData.getInt(LEADER_KEY_NAME);
-            int leaderEpoch = partitionStateData.getInt(LEADER_EPOCH_KEY_NAME);
-
-            Object[] isrArray = partitionStateData.getArray(ISR_KEY_NAME);
-            List isr = new ArrayList<>(isrArray.length);
-            for (Object r : isrArray)
-                isr.add((Integer) r);
-
-            int zkVersion = partitionStateData.getInt(ZK_VERSION_KEY_NAME);
-
-            Object[] replicasArray = partitionStateData.getArray(REPLICAS_KEY_NAME);
-            List replicas = new ArrayList<>(replicasArray.length);
-            for (Object r : replicasArray)
-                replicas.add((Integer) r);
-
-            List offlineReplicas = new ArrayList<>();
-            if (partitionStateData.hasField(OFFLINE_REPLICAS_KEY_NAME)) {
-                Object[] offlineReplicasArray = partitionStateData.getArray(OFFLINE_REPLICAS_KEY_NAME);
-                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));
+                    }
+                }
             }
-
-            PartitionState partitionState =
-                new PartitionState(controllerEpoch, leader, leaderEpoch, isr, zkVersion, replicas, offlineReplicas);
-            partitionStates.put(new TopicPartition(topic, partition), partitionState);
         }
 
-        Set liveBrokers = new HashSet<>();
-
-        for (Object brokerDataObj : struct.getArray(LIVE_BROKERS_KEY_NAME)) {
-            Struct brokerData = (Struct) brokerDataObj;
-            int brokerId = brokerData.getInt(BROKER_ID_KEY_NAME);
-
-            // V0
-            if (brokerData.hasField(HOST_KEY_NAME)) {
-                String host = brokerData.getString(HOST_KEY_NAME);
-                int port = brokerData.getInt(PORT_KEY_NAME);
-                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.getArray(ENDPOINTS_KEY_NAME)) {
-                    Struct endPointData = (Struct) endPointDataObj;
-                    int port = endPointData.getInt(PORT_KEY_NAME);
-                    String host = endPointData.getString(HOST_KEY_NAME);
-                    short protocolTypeId = endPointData.getShort(SECURITY_PROTOCOL_TYPE_KEY_NAME);
-                    SecurityProtocol securityProtocol = SecurityProtocol.forId(protocolTypeId);
-                    String listenerName;
-                    if (endPointData.hasField(LISTENER_NAME_KEY_NAME)) // V3
-                        listenerName = endPointData.getString(LISTENER_NAME_KEY_NAME);
-                    else
-                        listenerName = securityProtocol.name;
-                    endPoints.add(new EndPoint(host, port, securityProtocol, new ListenerName(listenerName)));
+        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());
                 }
-                String rack = null;
-                if (brokerData.hasField(RACK_KEY_NAME)) { // V2
-                    rack = brokerData.getString(RACK_KEY_NAME);
-                }
-                liveBrokers.add(new Broker(brokerId, endPoints, rack));
             }
         }
-        controllerId = struct.getInt(CONTROLLER_ID_KEY_NAME);
-        controllerEpoch = struct.getInt(CONTROLLER_EPOCH_KEY_NAME);
-        this.partitionStates = partitionStates;
-        this.liveBrokers = liveBrokers;
     }
 
-    @Override
-    protected Struct toStruct() {
-        short version = version();
-        Struct struct = new Struct(ApiKeys.UPDATE_METADATA.requestSchema(version));
-        struct.set(CONTROLLER_ID_KEY_NAME, controllerId);
-        struct.set(CONTROLLER_EPOCH_KEY_NAME, controllerEpoch);
-
-        List partitionStatesData = new ArrayList<>(partitionStates.size());
-        for (Map.Entry entry : partitionStates.entrySet()) {
-            Struct partitionStateData = struct.instance(PARTITION_STATES_KEY_NAME);
-            TopicPartition topicPartition = entry.getKey();
-            partitionStateData.set(TOPIC_NAME, topicPartition.topic());
-            partitionStateData.set(PARTITION_ID, topicPartition.partition());
-            PartitionState partitionState = entry.getValue();
-            partitionStateData.set(CONTROLLER_EPOCH_KEY_NAME, partitionState.basePartitionState.controllerEpoch);
-            partitionStateData.set(LEADER_KEY_NAME, partitionState.basePartitionState.leader);
-            partitionStateData.set(LEADER_EPOCH_KEY_NAME, partitionState.basePartitionState.leaderEpoch);
-            partitionStateData.set(ISR_KEY_NAME, partitionState.basePartitionState.isr.toArray());
-            partitionStateData.set(ZK_VERSION_KEY_NAME, partitionState.basePartitionState.zkVersion);
-            partitionStateData.set(REPLICAS_KEY_NAME, partitionState.basePartitionState.replicas.toArray());
-            if (partitionStateData.hasField(OFFLINE_REPLICAS_KEY_NAME))
-              partitionStateData.set(OFFLINE_REPLICAS_KEY_NAME, partitionState.offlineReplicas.toArray());
-            partitionStatesData.add(partitionStateData);
-        }
-        struct.set(PARTITION_STATES_KEY_NAME, partitionStatesData.toArray());
-
-        List brokersData = new ArrayList<>(liveBrokers.size());
-        for (Broker broker : liveBrokers) {
-            Struct brokerData = struct.instance(LIVE_BROKERS_KEY_NAME);
-            brokerData.set(BROKER_ID_KEY_NAME, broker.id);
-
-            if (version == 0) {
-                EndPoint endPoint = broker.endPoints.get(0);
-                brokerData.set(HOST_KEY_NAME, endPoint.host);
-                brokerData.set(PORT_KEY_NAME, endPoint.port);
-            } else {
-                List endPointsData = new ArrayList<>(broker.endPoints.size());
-                for (EndPoint endPoint : broker.endPoints) {
-                    Struct endPointData = brokerData.instance(ENDPOINTS_KEY_NAME);
-                    endPointData.set(PORT_KEY_NAME, endPoint.port);
-                    endPointData.set(HOST_KEY_NAME, endPoint.host);
-                    endPointData.set(SECURITY_PROTOCOL_TYPE_KEY_NAME, endPoint.securityProtocol.id);
-                    if (version >= 3)
-                        endPointData.set(LISTENER_NAME_KEY_NAME, endPoint.listenerName.value());
-                    endPointsData.add(endPointData);
-
-                }
-                brokerData.set(ENDPOINTS_KEY_NAME, endPointsData.toArray());
-                if (version >= 2) {
-                    brokerData.set(RACK_KEY_NAME, broker.rack);
-                }
-            }
+    private static String listenerNameFromSecurityProtocol(UpdateMetadataEndpoint endpoint) {
+        SecurityProtocol securityProtocol = SecurityProtocol.forId(endpoint.securityProtocol());
+        return ListenerName.forSecurityProtocol(securityProtocol).value();
+    }
 
-            brokersData.add(brokerData);
-        }
-        struct.set(LIVE_BROKERS_KEY_NAME, brokersData.toArray());
+    @Override
+    public int controllerId() {
+        return data.controllerId();
+    }
 
-        return struct;
+    @Override
+    public int controllerEpoch() {
+        return data.controllerEpoch();
     }
 
     @Override
-    public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
-        short versionId = version();
-        if (versionId <= 4)
-            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 long brokerEpoch() {
+        return data.brokerEpoch();
     }
 
-    public int controllerId() {
-        return controllerId;
+    @Override
+    public UpdateMetadataResponse getErrorResponse(int throttleTimeMs, Throwable e) {
+        UpdateMetadataResponseData data = new UpdateMetadataResponseData()
+                .setErrorCode(Errors.forException(e).code());
+        return new UpdateMetadataResponse(data);
     }
 
-    public int controllerEpoch() {
-        return controllerEpoch;
+    public Iterable partitionStates() {
+        if (version() >= 5) {
+            return () -> new FlattenedIterator<>(data.topicStates().iterator(),
+                topicState -> topicState.partitionStates().iterator());
+        }
+        return data.ungroupedPartitionStates();
     }
 
-    public Map partitionStates() {
-        return partitionStates;
+    public List liveBrokers() {
+        return data.liveBrokers();
     }
 
-    public Set liveBrokers() {
-        return liveBrokers;
+    @Override
+    public UpdateMetadataRequestData data() {
+        return data;
     }
 
     public static UpdateMetadataRequest parse(ByteBuffer buffer, short version) {
-        return new UpdateMetadataRequest(ApiKeys.UPDATE_METADATA.parseRequest(version, buffer), version);
+        return new UpdateMetadataRequest(new UpdateMetadataRequestData(new ByteBufferAccessor(buffer), version), 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 4c21cdeaccc36..cc7749a47242c 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,60 +16,44 @@
  */
 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.ByteBufferAccessor;
 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;
-
-    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};
-    }
 
-    /**
-     * Possible error code:
-     *
-     * STALE_CONTROLLER_EPOCH (11)
-     */
-    private final Errors error;
-
-    public UpdateMetadataResponse(Errors error) {
-        this.error = error;
-    }
+    private final UpdateMetadataResponseData data;
 
-    public UpdateMetadataResponse(Struct struct) {
-        error = Errors.forCode(struct.get(ERROR_CODE));
+    public UpdateMetadataResponse(UpdateMetadataResponseData data) {
+        super(ApiKeys.UPDATE_METADATA);
+        this.data = data;
     }
 
     public Errors error() {
-        return error;
+        return Errors.forCode(data.errorCode());
     }
 
     @Override
     public Map errorCounts() {
-        return errorCounts(error);
+        return errorCounts(error());
+    }
+
+    @Override
+    public int throttleTimeMs() {
+        return DEFAULT_THROTTLE_TIME;
     }
 
     public static UpdateMetadataResponse parse(ByteBuffer buffer, short version) {
-        return new UpdateMetadataResponse(ApiKeys.UPDATE_METADATA.parseResponse(version, buffer));
+        return new UpdateMetadataResponse(new UpdateMetadataResponseData(new ByteBufferAccessor(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;
+    public UpdateMetadataResponseData data() {
+        return data;
     }
+
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/VoteRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/VoteRequest.java
new file mode 100644
index 0000000000000..8fba2f085d566
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/VoteRequest.java
@@ -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 org.apache.kafka.common.requests;
+
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.message.VoteRequestData;
+import org.apache.kafka.common.message.VoteResponseData;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.Errors;
+
+import java.nio.ByteBuffer;
+import java.util.Collections;
+
+public class VoteRequest extends AbstractRequest {
+
+    public static class Builder extends AbstractRequest.Builder {
+        private final VoteRequestData data;
+
+        public Builder(VoteRequestData data) {
+            super(ApiKeys.VOTE);
+            this.data = data;
+        }
+
+        @Override
+        public VoteRequest build(short version) {
+            return new VoteRequest(data, version);
+        }
+
+        @Override
+        public String toString() {
+            return data.toString();
+        }
+    }
+
+    private final VoteRequestData data;
+
+    private VoteRequest(VoteRequestData data, short version) {
+        super(ApiKeys.VOTE, version);
+        this.data = data;
+    }
+
+    @Override
+    public VoteRequestData data() {
+        return data;
+    }
+
+    @Override
+    public AbstractResponse getErrorResponse(int throttleTimeMs, Throwable e) {
+        return new VoteResponse(new VoteResponseData()
+            .setErrorCode(Errors.forException(e).code()));
+    }
+
+    public static VoteRequest parse(ByteBuffer buffer, short version) {
+        return new VoteRequest(new VoteRequestData(new ByteBufferAccessor(buffer), version), version);
+    }
+
+    public static VoteRequestData singletonRequest(TopicPartition topicPartition,
+                                                   int candidateEpoch,
+                                                   int candidateId,
+                                                   int lastEpoch,
+                                                   long lastEpochEndOffset) {
+        return singletonRequest(topicPartition,
+            null,
+            candidateEpoch,
+            candidateId,
+            lastEpoch,
+            lastEpochEndOffset);
+    }
+
+    public static VoteRequestData singletonRequest(TopicPartition topicPartition,
+                                                   String clusterId,
+                                                   int candidateEpoch,
+                                                   int candidateId,
+                                                   int lastEpoch,
+                                                   long lastEpochEndOffset) {
+        return new VoteRequestData()
+                   .setClusterId(clusterId)
+                   .setTopics(Collections.singletonList(
+                       new VoteRequestData.TopicData()
+                           .setTopicName(topicPartition.topic())
+                           .setPartitions(Collections.singletonList(
+                               new VoteRequestData.PartitionData()
+                                   .setPartitionIndex(topicPartition.partition())
+                                   .setCandidateEpoch(candidateEpoch)
+                                   .setCandidateId(candidateId)
+                                   .setLastOffsetEpoch(lastEpoch)
+                                   .setLastOffset(lastEpochEndOffset))
+                           )));
+    }
+
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/VoteResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/VoteResponse.java
new file mode 100644
index 0000000000000..51991adcf0cbb
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/requests/VoteResponse.java
@@ -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 org.apache.kafka.common.requests;
+
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.message.VoteResponseData;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.Errors;
+
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * Possible error codes.
+ *
+ * Top level errors:
+ * - {@link Errors#CLUSTER_AUTHORIZATION_FAILED}
+ * - {@link Errors#BROKER_NOT_AVAILABLE}
+ *
+ * Partition level errors:
+ * - {@link Errors#FENCED_LEADER_EPOCH}
+ * - {@link Errors#INVALID_REQUEST}
+ * - {@link Errors#INCONSISTENT_VOTER_SET}
+ * - {@link Errors#UNKNOWN_TOPIC_OR_PARTITION}
+ */
+public class VoteResponse extends AbstractResponse {
+    private final VoteResponseData data;
+
+    public VoteResponse(VoteResponseData data) {
+        super(ApiKeys.VOTE);
+        this.data = data;
+    }
+
+    public static VoteResponseData singletonResponse(Errors topLevelError,
+                                                     TopicPartition topicPartition,
+                                                     Errors partitionLevelError,
+                                                     int leaderEpoch,
+                                                     int leaderId,
+                                                     boolean voteGranted) {
+        return new VoteResponseData()
+            .setErrorCode(topLevelError.code())
+            .setTopics(Collections.singletonList(
+                new VoteResponseData.TopicData()
+                    .setTopicName(topicPartition.topic())
+                    .setPartitions(Collections.singletonList(
+                        new VoteResponseData.PartitionData()
+                            .setErrorCode(partitionLevelError.code())
+                            .setLeaderId(leaderId)
+                            .setLeaderEpoch(leaderEpoch)
+                            .setVoteGranted(voteGranted)))));
+    }
+
+    @Override
+    public Map errorCounts() {
+        Map errors = new HashMap<>();
+
+        errors.put(Errors.forCode(data.errorCode()), 1);
+
+        for (VoteResponseData.TopicData topicResponse : data.topics()) {
+            for (VoteResponseData.PartitionData partitionResponse : topicResponse.partitions()) {
+                updateErrorCounts(errors, Errors.forCode(partitionResponse.errorCode()));
+            }
+        }
+        return errors;
+    }
+
+    @Override
+    public VoteResponseData data() {
+        return data;
+    }
+
+    @Override
+    public int throttleTimeMs() {
+        return DEFAULT_THROTTLE_TIME;
+    }
+
+    public static VoteResponse parse(ByteBuffer buffer, short version) {
+        return new VoteResponse(new VoteResponseData(new ByteBufferAccessor(buffer), version));
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/WriteTxnMarkersRequest.java b/clients/src/main/java/org/apache/kafka/common/requests/WriteTxnMarkersRequest.java
index 3f7a0c9fc9b52..64a3df49c7b39 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/WriteTxnMarkersRequest.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/WriteTxnMarkersRequest.java
@@ -17,13 +17,12 @@
 package org.apache.kafka.common.requests;
 
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.message.WriteTxnMarkersRequestData;
+import org.apache.kafka.common.message.WriteTxnMarkersRequestData.WritableTxnMarker;
+import org.apache.kafka.common.message.WriteTxnMarkersRequestData.WritableTxnMarkerTopic;
 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.Struct;
-import org.apache.kafka.common.utils.CollectionUtils;
 
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
@@ -32,40 +31,7 @@
 import java.util.Map;
 import java.util.Objects;
 
-import static org.apache.kafka.common.protocol.CommonFields.TOPIC_NAME;
-import static org.apache.kafka.common.protocol.types.Type.BOOLEAN;
-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.INT64;
-
 public class WriteTxnMarkersRequest extends AbstractRequest {
-    private static final String COORDINATOR_EPOCH_KEY_NAME = "coordinator_epoch";
-    private static final String TXN_MARKERS_KEY_NAME = "transaction_markers";
-
-    private static final String PRODUCER_ID_KEY_NAME = "producer_id";
-    private static final String PRODUCER_EPOCH_KEY_NAME = "producer_epoch";
-    private static final String TRANSACTION_RESULT_KEY_NAME = "transaction_result";
-    private static final String TOPICS_KEY_NAME = "topics";
-    private static final String PARTITIONS_KEY_NAME = "partitions";
-
-    private static final Schema WRITE_TXN_MARKERS_ENTRY_V0 = new Schema(
-            new Field(PRODUCER_ID_KEY_NAME, INT64, "Current producer id in use by the transactional id."),
-            new Field(PRODUCER_EPOCH_KEY_NAME, INT16, "Current epoch associated with the producer id."),
-            new Field(TRANSACTION_RESULT_KEY_NAME, BOOLEAN, "The result of the transaction to write to the " +
-                    "partitions (false = ABORT, true = COMMIT)."),
-            new Field(TOPICS_KEY_NAME, new ArrayOf(new Schema(
-                    TOPIC_NAME,
-                    new Field(PARTITIONS_KEY_NAME, new ArrayOf(INT32)))), "The partitions to write markers for."),
-            new Field(COORDINATOR_EPOCH_KEY_NAME, INT32, "Epoch associated with the transaction state partition " +
-                    "hosted by this transaction coordinator"));
-
-    private static final Schema WRITE_TXN_MARKERS_REQUEST_V0 = new Schema(
-            new Field(TXN_MARKERS_KEY_NAME, new ArrayOf(WRITE_TXN_MARKERS_ENTRY_V0), "The transaction markers to " +
-                    "be written."));
-
-    public static Schema[] schemaVersions() {
-        return new Schema[]{WRITE_TXN_MARKERS_REQUEST_V0};
-    }
 
     public static class TxnMarkerEntry {
         private final long producerId;
@@ -106,16 +72,15 @@ public List partitions() {
             return partitions;
         }
 
-
         @Override
         public String toString() {
             return "TxnMarkerEntry{" +
-                    "producerId=" + producerId +
-                    ", producerEpoch=" + producerEpoch +
-                    ", coordinatorEpoch=" + coordinatorEpoch +
-                    ", result=" + result +
-                    ", partitions=" + partitions +
-                    '}';
+                       "producerId=" + producerId +
+                       ", producerEpoch=" + producerEpoch +
+                       ", coordinatorEpoch=" + coordinatorEpoch +
+                       ", result=" + result +
+                       ", partitions=" + partitions +
+                       '}';
         }
 
         @Override
@@ -124,10 +89,10 @@ public boolean equals(final Object o) {
             if (o == null || getClass() != o.getClass()) return false;
             final TxnMarkerEntry that = (TxnMarkerEntry) o;
             return producerId == that.producerId &&
-                    producerEpoch == that.producerEpoch &&
-                    coordinatorEpoch == that.coordinatorEpoch &&
-                    result == that.result &&
-                    Objects.equals(partitions, that.partitions);
+                       producerEpoch == that.producerEpoch &&
+                       coordinatorEpoch == that.coordinatorEpoch &&
+                       result == that.result &&
+                       Objects.equals(partitions, that.partitions);
         }
 
         @Override
@@ -137,108 +102,90 @@ public int hashCode() {
     }
 
     public static class Builder extends AbstractRequest.Builder {
-        private final List markers;
 
-        public Builder(List markers) {
-            super(ApiKeys.WRITE_TXN_MARKERS);
-            this.markers = markers;
+        public final WriteTxnMarkersRequestData data;
+
+        public Builder(short version, final List markers) {
+            super(ApiKeys.WRITE_TXN_MARKERS, version);
+            List dataMarkers = new ArrayList<>();
+            for (TxnMarkerEntry marker : markers) {
+                final Map topicMap = new HashMap<>();
+                for (TopicPartition topicPartition : marker.partitions) {
+                    WritableTxnMarkerTopic topic = topicMap.getOrDefault(topicPartition.topic(),
+                                                                         new WritableTxnMarkerTopic()
+                                                                             .setName(topicPartition.topic()));
+                    topic.partitionIndexes().add(topicPartition.partition());
+                    topicMap.put(topicPartition.topic(), topic);
+                }
+
+                dataMarkers.add(new WritableTxnMarker()
+                                    .setProducerId(marker.producerId)
+                                    .setProducerEpoch(marker.producerEpoch)
+                                    .setCoordinatorEpoch(marker.coordinatorEpoch)
+                                    .setTransactionResult(marker.transactionResult().id)
+                                    .setTopics(new ArrayList<>(topicMap.values())));
+            }
+            this.data = new WriteTxnMarkersRequestData().setMarkers(dataMarkers);
         }
 
         @Override
         public WriteTxnMarkersRequest build(short version) {
-            return new WriteTxnMarkersRequest(version, markers);
+            return new WriteTxnMarkersRequest(data, version);
         }
     }
 
-    private final List markers;
+    private final WriteTxnMarkersRequestData data;
 
-    private WriteTxnMarkersRequest(short version, List markers) {
-        super(version);
-
-        this.markers = markers;
+    private WriteTxnMarkersRequest(WriteTxnMarkersRequestData data, short version) {
+        super(ApiKeys.WRITE_TXN_MARKERS, version);
+        this.data = data;
     }
 
-    public WriteTxnMarkersRequest(Struct struct, short version) {
-        super(version);
-        List markers = new ArrayList<>();
-        Object[] markersArray = struct.getArray(TXN_MARKERS_KEY_NAME);
-        for (Object markerObj : markersArray) {
-            Struct markerStruct = (Struct) markerObj;
+    @Override
+    public WriteTxnMarkersRequestData data() {
+        return data;
+    }
 
-            long producerId = markerStruct.getLong(PRODUCER_ID_KEY_NAME);
-            short producerEpoch = markerStruct.getShort(PRODUCER_EPOCH_KEY_NAME);
-            int coordinatorEpoch = markerStruct.getInt(COORDINATOR_EPOCH_KEY_NAME);
-            TransactionResult result = TransactionResult.forId(markerStruct.getBoolean(TRANSACTION_RESULT_KEY_NAME));
+    @Override
+    public WriteTxnMarkersResponse getErrorResponse(int throttleTimeMs, Throwable e) {
+        Errors error = Errors.forException(e);
 
-            List partitions = new ArrayList<>();
-            Object[] topicPartitionsArray = markerStruct.getArray(TOPICS_KEY_NAME);
-            for (Object topicPartitionObj : topicPartitionsArray) {
-                Struct topicPartitionStruct = (Struct) topicPartitionObj;
-                String topic = topicPartitionStruct.get(TOPIC_NAME);
-                for (Object partitionObj : topicPartitionStruct.getArray(PARTITIONS_KEY_NAME)) {
-                    partitions.add(new TopicPartition(topic, (Integer) partitionObj));
+        final Map> errors = new HashMap<>(data.markers().size());
+        for (WritableTxnMarker markerEntry : data.markers()) {
+            Map errorsPerPartition = new HashMap<>();
+            for (WritableTxnMarkerTopic topic : markerEntry.topics()) {
+                for (Integer partitionIdx : topic.partitionIndexes()) {
+                    errorsPerPartition.put(new TopicPartition(topic.name(), partitionIdx), error);
                 }
             }
-
-            markers.add(new TxnMarkerEntry(producerId, producerEpoch, coordinatorEpoch, result, partitions));
+            errors.put(markerEntry.producerId(), errorsPerPartition);
         }
 
-        this.markers = markers;
+        return new WriteTxnMarkersResponse(errors);
     }
 
-
     public List markers() {
-        return markers;
-    }
-
-    @Override
-    protected Struct toStruct() {
-        Struct struct = new Struct(ApiKeys.WRITE_TXN_MARKERS.requestSchema(version()));
-
-        Object[] markersArray = new Object[markers.size()];
-        int i = 0;
-        for (TxnMarkerEntry entry : markers) {
-            Struct markerStruct = struct.instance(TXN_MARKERS_KEY_NAME);
-            markerStruct.set(PRODUCER_ID_KEY_NAME, entry.producerId);
-            markerStruct.set(PRODUCER_EPOCH_KEY_NAME, entry.producerEpoch);
-            markerStruct.set(COORDINATOR_EPOCH_KEY_NAME, entry.coordinatorEpoch);
-            markerStruct.set(TRANSACTION_RESULT_KEY_NAME, entry.result.id);
-
-            Map> mappedPartitions = CollectionUtils.groupDataByTopic(entry.partitions);
-            Object[] partitionsArray = new Object[mappedPartitions.size()];
-            int j = 0;
-            for (Map.Entry> topicAndPartitions : mappedPartitions.entrySet()) {
-                Struct topicPartitionsStruct = markerStruct.instance(TOPICS_KEY_NAME);
-                topicPartitionsStruct.set(TOPIC_NAME, topicAndPartitions.getKey());
-                topicPartitionsStruct.set(PARTITIONS_KEY_NAME, topicAndPartitions.getValue().toArray());
-                partitionsArray[j++] = topicPartitionsStruct;
+        List markers = new ArrayList<>();
+        for (WritableTxnMarker markerEntry : data.markers()) {
+            List topicPartitions = new ArrayList<>();
+            for (WritableTxnMarkerTopic topic : markerEntry.topics()) {
+                for (Integer partitionIdx : topic.partitionIndexes()) {
+                    topicPartitions.add(new TopicPartition(topic.name(), partitionIdx));
+                }
             }
-            markerStruct.set(TOPICS_KEY_NAME, partitionsArray);
-            markersArray[i++] = markerStruct;
-        }
-        struct.set(TXN_MARKERS_KEY_NAME, markersArray);
-
-        return struct;
-    }
-
-    @Override
-    public WriteTxnMarkersResponse getErrorResponse(int throttleTimeMs, Throwable e) {
-        Errors error = Errors.forException(e);
-
-        Map> errors = new HashMap<>(markers.size());
-        for (TxnMarkerEntry entry : markers) {
-            Map errorsPerPartition = new HashMap<>(entry.partitions.size());
-            for (TopicPartition partition : entry.partitions)
-                errorsPerPartition.put(partition, error);
-
-            errors.put(entry.producerId, errorsPerPartition);
+            markers.add(new TxnMarkerEntry(
+                markerEntry.producerId(),
+                markerEntry.producerEpoch(),
+                markerEntry.coordinatorEpoch(),
+                TransactionResult.forId(markerEntry.transactionResult()),
+                topicPartitions)
+            );
         }
-
-        return new WriteTxnMarkersResponse(errors);
+        return markers;
     }
 
     public static WriteTxnMarkersRequest parse(ByteBuffer buffer, short version) {
-        return new WriteTxnMarkersRequest(ApiKeys.WRITE_TXN_MARKERS.parseRequest(version, buffer), version);
+        return new WriteTxnMarkersRequest(new WriteTxnMarkersRequestData(new ByteBufferAccessor(buffer), version), version);
     }
 
     @Override
@@ -246,11 +193,11 @@ public boolean equals(final Object o) {
         if (this == o) return true;
         if (o == null || getClass() != o.getClass()) return false;
         final WriteTxnMarkersRequest that = (WriteTxnMarkersRequest) o;
-        return Objects.equals(markers, that.markers);
+        return Objects.equals(this.data, that.data);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(markers);
+        return Objects.hash(this.data);
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/requests/WriteTxnMarkersResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/WriteTxnMarkersResponse.java
index f4bf157adaf51..fd2a834d24569 100644
--- a/clients/src/main/java/org/apache/kafka/common/requests/WriteTxnMarkersResponse.java
+++ b/clients/src/main/java/org/apache/kafka/common/requests/WriteTxnMarkersResponse.java
@@ -17,151 +17,110 @@
 package org.apache.kafka.common.requests;
 
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.message.WriteTxnMarkersResponseData;
+import org.apache.kafka.common.message.WriteTxnMarkersResponseData.WritableTxnMarkerPartitionResult;
+import org.apache.kafka.common.message.WriteTxnMarkersResponseData.WritableTxnMarkerResult;
+import org.apache.kafka.common.message.WriteTxnMarkersResponseData.WritableTxnMarkerTopicResult;
 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.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 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 static org.apache.kafka.common.protocol.types.Type.INT64;
-
+/**
+ * Possible error codes:
+ *
+ *   - {@link Errors#CORRUPT_MESSAGE}
+ *   - {@link Errors#INVALID_PRODUCER_EPOCH}
+ *   - {@link Errors#UNKNOWN_TOPIC_OR_PARTITION}
+ *   - {@link Errors#NOT_LEADER_OR_FOLLOWER}
+ *   - {@link Errors#MESSAGE_TOO_LARGE}
+ *   - {@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#TRANSACTION_COORDINATOR_FENCED}
+ *   - {@link Errors#REQUEST_TIMED_OUT}
+ *   - {@link Errors#CLUSTER_AUTHORIZATION_FAILED}
+ */
 public class WriteTxnMarkersResponse extends AbstractResponse {
-    private static final String TXN_MARKERS_KEY_NAME = "transaction_markers";
-
-    private static final String PRODUCER_ID_KEY_NAME = "producer_id";
-    private static final String TOPICS_KEY_NAME = "topics";
-    private static final String PARTITIONS_KEY_NAME = "partitions";
-
-    private static final Schema WRITE_TXN_MARKERS_PARTITION_ERROR_RESPONSE_V0 = new Schema(
-            PARTITION_ID,
-            ERROR_CODE);
-
-    private static final Schema WRITE_TXN_MARKERS_ENTRY_V0 = new Schema(
-            new Field(PRODUCER_ID_KEY_NAME, INT64, "Current producer id in use by the transactional id."),
-            new Field(TOPICS_KEY_NAME, new ArrayOf(new Schema(
-                    TOPIC_NAME,
-                    new Field(PARTITIONS_KEY_NAME, new ArrayOf(WRITE_TXN_MARKERS_PARTITION_ERROR_RESPONSE_V0)))),
-                    "Errors per partition from writing markers."));
-
-    private static final Schema WRITE_TXN_MARKERS_RESPONSE_V0 = new Schema(
-            new Field(TXN_MARKERS_KEY_NAME, new ArrayOf(WRITE_TXN_MARKERS_ENTRY_V0), "Errors per partition from " +
-                    "writing markers."));
 
-    public static Schema[] schemaVersions() {
-        return new Schema[]{WRITE_TXN_MARKERS_RESPONSE_V0};
-    }
-
-    // Possible error codes:
-    //   CorruptRecord
-    //   InvalidProducerEpoch
-    //   UnknownTopicOrPartition
-    //   NotLeaderForPartition
-    //   MessageTooLarge
-    //   RecordListTooLarge
-    //   NotEnoughReplicas
-    //   NotEnoughReplicasAfterAppend
-    //   InvalidRequiredAcks
-    //   TransactionCoordinatorFenced
-    //   RequestTimeout
-    //   ClusterAuthorizationFailed
-
-    private final Map> errors;
+    private final WriteTxnMarkersResponseData data;
 
     public WriteTxnMarkersResponse(Map> errors) {
-        this.errors = errors;
-    }
-
-    public WriteTxnMarkersResponse(Struct struct) {
-        Map> errors = new HashMap<>();
-
-        Object[] responseArray = struct.getArray(TXN_MARKERS_KEY_NAME);
-        for (Object responseObj : responseArray) {
-            Struct responseStruct = (Struct) responseObj;
-
-            long producerId = responseStruct.getLong(PRODUCER_ID_KEY_NAME);
-
-            Map errorPerPartition = new HashMap<>();
-            Object[] topicPartitionsArray = responseStruct.getArray(TOPICS_KEY_NAME);
-            for (Object topicPartitionObj : topicPartitionsArray) {
-                Struct topicPartitionStruct = (Struct) topicPartitionObj;
-                String topic = topicPartitionStruct.get(TOPIC_NAME);
-                for (Object partitionObj : topicPartitionStruct.getArray(PARTITIONS_KEY_NAME)) {
-                    Struct partitionStruct = (Struct) partitionObj;
-                    Integer partition = partitionStruct.get(PARTITION_ID);
-                    Errors error = Errors.forCode(partitionStruct.get(ERROR_CODE));
-                    errorPerPartition.put(new TopicPartition(topic, partition), error);
-                }
+        super(ApiKeys.WRITE_TXN_MARKERS);
+        List markers = new ArrayList<>();
+        for (Map.Entry> markerEntry : errors.entrySet()) {
+            Map responseTopicDataMap = new HashMap<>();
+            for (Map.Entry topicEntry : markerEntry.getValue().entrySet()) {
+                TopicPartition topicPartition = topicEntry.getKey();
+                String topicName = topicPartition.topic();
+
+                WritableTxnMarkerTopicResult topic =
+                    responseTopicDataMap.getOrDefault(topicName, new WritableTxnMarkerTopicResult().setName(topicName));
+                topic.partitions().add(new WritableTxnMarkerPartitionResult()
+                                           .setErrorCode(topicEntry.getValue().code())
+                                           .setPartitionIndex(topicPartition.partition())
+                );
+                responseTopicDataMap.put(topicName, topic);
             }
-            errors.put(producerId, errorPerPartition);
+
+            markers.add(new WritableTxnMarkerResult()
+                            .setProducerId(markerEntry.getKey())
+                            .setTopics(new ArrayList<>(responseTopicDataMap.values()))
+            );
         }
+        this.data = new WriteTxnMarkersResponseData()
+                        .setMarkers(markers);
+    }
 
-        this.errors = errors;
+    public WriteTxnMarkersResponse(WriteTxnMarkersResponseData data) {
+        super(ApiKeys.WRITE_TXN_MARKERS);
+        this.data = data;
     }
 
     @Override
-    protected Struct toStruct(short version) {
-        Struct struct = new Struct(ApiKeys.WRITE_TXN_MARKERS.responseSchema(version));
-
-        Object[] responsesArray = new Object[errors.size()];
-        int k = 0;
-        for (Map.Entry> responseEntry : errors.entrySet()) {
-            Struct responseStruct = struct.instance(TXN_MARKERS_KEY_NAME);
-            responseStruct.set(PRODUCER_ID_KEY_NAME, responseEntry.getKey());
-
-            Map partitionAndErrors = responseEntry.getValue();
-            Map> mappedPartitions = CollectionUtils.groupDataByTopic(partitionAndErrors);
-            Object[] partitionsArray = new Object[mappedPartitions.size()];
-            int i = 0;
-            for (Map.Entry> topicAndPartitions : mappedPartitions.entrySet()) {
-                Struct topicPartitionsStruct = responseStruct.instance(TOPICS_KEY_NAME);
-                topicPartitionsStruct.set(TOPIC_NAME, topicAndPartitions.getKey());
-                Map partitionIdAndErrors = topicAndPartitions.getValue();
+    public WriteTxnMarkersResponseData data() {
+        return data;
+    }
 
-                Object[] partitionAndErrorsArray = new Object[partitionIdAndErrors.size()];
-                int j = 0;
-                for (Map.Entry partitionAndError : partitionIdAndErrors.entrySet()) {
-                    Struct partitionAndErrorStruct = topicPartitionsStruct.instance(PARTITIONS_KEY_NAME);
-                    partitionAndErrorStruct.set(PARTITION_ID, partitionAndError.getKey());
-                    partitionAndErrorStruct.set(ERROR_CODE, partitionAndError.getValue().code());
-                    partitionAndErrorsArray[j++] = partitionAndErrorStruct;
+    public Map> errorsByProducerId() {
+        Map> errors = new HashMap<>();
+        for (WritableTxnMarkerResult marker : data.markers()) {
+            Map topicPartitionErrorsMap = new HashMap<>();
+            for (WritableTxnMarkerTopicResult topic : marker.topics()) {
+                for (WritableTxnMarkerPartitionResult partitionResult : topic.partitions()) {
+                    topicPartitionErrorsMap.put(new TopicPartition(topic.name(), partitionResult.partitionIndex()),
+                            Errors.forCode(partitionResult.errorCode()));
                 }
-                topicPartitionsStruct.set(PARTITIONS_KEY_NAME, partitionAndErrorsArray);
-                partitionsArray[i++] = topicPartitionsStruct;
             }
-            responseStruct.set(TOPICS_KEY_NAME, partitionsArray);
-
-            responsesArray[k++] = responseStruct;
+            errors.put(marker.producerId(), topicPartitionErrorsMap);
         }
-
-        struct.set(TXN_MARKERS_KEY_NAME, responsesArray);
-        return struct;
+        return errors;
     }
 
-    public Map errors(long producerId) {
-        return errors.get(producerId);
+    @Override
+    public int throttleTimeMs() {
+        return DEFAULT_THROTTLE_TIME;
     }
 
     @Override
     public Map errorCounts() {
         Map errorCounts = new HashMap<>();
-        for (Map allErrors : errors.values()) {
-            for (Errors error : allErrors.values())
-                updateErrorCounts(errorCounts, error);
+        for (WritableTxnMarkerResult marker : data.markers()) {
+            for (WritableTxnMarkerTopicResult topic : marker.topics()) {
+                for (WritableTxnMarkerPartitionResult partitionResult : topic.partitions())
+                    updateErrorCounts(errorCounts, Errors.forCode(partitionResult.errorCode()));
+            }
         }
         return errorCounts;
     }
 
     public static WriteTxnMarkersResponse parse(ByteBuffer buffer, short version) {
-        return new WriteTxnMarkersResponse(ApiKeys.WRITE_TXN_MARKERS.parseResponse(version, buffer));
+        return new WriteTxnMarkersResponse(new WriteTxnMarkersResponseData(new ByteBufferAccessor(buffer), version));
     }
-
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/resource/PatternType.java b/clients/src/main/java/org/apache/kafka/common/resource/PatternType.java
new file mode 100644
index 0000000000000..0c05a0b0466b4
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/resource/PatternType.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.resource;
+
+import org.apache.kafka.common.annotation.InterfaceStability;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.Map;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+
+/**
+ * Resource pattern type.
+ */
+@InterfaceStability.Evolving
+public enum PatternType {
+    /**
+     * Represents any PatternType which this client cannot understand, perhaps because this client is too old.
+     */
+    UNKNOWN((byte) 0),
+
+    /**
+     * In a filter, matches any resource pattern type.
+     */
+    ANY((byte) 1),
+
+    /**
+     * In a filter, will perform pattern matching.
+     *
+     * e.g. Given a filter of {@code ResourcePatternFilter(TOPIC, "payments.received", MATCH)`}, the filter match
+     * any {@link ResourcePattern} that matches topic 'payments.received'. This might include:
+     * 
          
+     *     
        • A Literal pattern with the same type and name, e.g. {@code ResourcePattern(TOPIC, "payments.received", LITERAL)}
          • 
+     *     
        • A Wildcard pattern with the same type, e.g. {@code ResourcePattern(TOPIC, "*", LITERAL)}
          • 
+     *     
        • A Prefixed pattern with the same type and where the name is a matching prefix, e.g. {@code ResourcePattern(TOPIC, "payments.", PREFIXED)}
          • 
+     * 
        
+     */
+    MATCH((byte) 2),
+
+    /**
+     * A literal resource name.
+     *
+     * A literal name defines the full name of a resource, e.g. topic with name 'foo', or group with name 'bob'.
+     *
+     * The special wildcard character {@code *} can be used to represent a resource with any name.
+     */
+    LITERAL((byte) 3),
+
+    /**
+     * A prefixed resource name.
+     *
+     * A prefixed name defines a prefix for a resource, e.g. topics with names that start with 'foo'.
+     */
+    PREFIXED((byte) 4);
+
+    private final static Map CODE_TO_VALUE =
+        Collections.unmodifiableMap(
+            Arrays.stream(PatternType.values())
+                .collect(Collectors.toMap(PatternType::code, Function.identity()))
+        );
+
+    private final static Map NAME_TO_VALUE =
+        Collections.unmodifiableMap(
+            Arrays.stream(PatternType.values())
+                .collect(Collectors.toMap(PatternType::name, Function.identity()))
+        );
+
+    private final byte code;
+
+    PatternType(byte code) {
+        this.code = code;
+    }
+
+    /**
+     * @return the code of this resource.
+     */
+    public byte code() {
+        return code;
+    }
+
+    /**
+     * @return whether this resource pattern type is UNKNOWN.
+     */
+    public boolean isUnknown() {
+        return this == UNKNOWN;
+    }
+
+    /**
+     * @return whether this resource pattern type is a concrete type, rather than UNKNOWN or one of the filter types.
+     */
+    public boolean isSpecific() {
+        return this != UNKNOWN && this != ANY && this != MATCH;
+    }
+
+    /**
+     * Return the PatternType with the provided code or {@link #UNKNOWN} if one cannot be found.
+     */
+    public static PatternType fromCode(byte code) {
+        return CODE_TO_VALUE.getOrDefault(code, UNKNOWN);
+    }
+
+    /**
+     * Return the PatternType with the provided name or {@link #UNKNOWN} if one cannot be found.
+     */
+    public static PatternType fromString(String name) {
+        return NAME_TO_VALUE.getOrDefault(name, UNKNOWN);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/resource/ResourceFilter.java b/clients/src/main/java/org/apache/kafka/common/resource/ResourceFilter.java
index 0a4611f9874b2..2ea9032e61941 100644
--- a/clients/src/main/java/org/apache/kafka/common/resource/ResourceFilter.java
+++ b/clients/src/main/java/org/apache/kafka/common/resource/ResourceFilter.java
@@ -93,9 +93,7 @@ public int hashCode() {
     public boolean matches(Resource other) {
         if ((name != null) && (!name.equals(other.name())))
             return false;
-        if ((resourceType != ResourceType.ANY) && (!resourceType.equals(other.resourceType())))
-            return false;
-        return true;
+        return (resourceType == ResourceType.ANY) || (resourceType.equals(other.resourceType()));
     }
 
     /**
diff --git a/clients/src/main/java/org/apache/kafka/common/resource/ResourcePattern.java b/clients/src/main/java/org/apache/kafka/common/resource/ResourcePattern.java
new file mode 100644
index 0000000000000..2b7504f70a598
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/resource/ResourcePattern.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.common.resource;
+
+import org.apache.kafka.common.annotation.InterfaceStability;
+
+import java.util.Objects;
+
+/**
+ * Represents a pattern that is used by ACLs to match zero or more
+ * {@link org.apache.kafka.common.resource.Resource Resources}.
+ *
+ * The API for this class is still evolving and we may break compatibility in minor releases, if necessary.
+ */
+@InterfaceStability.Evolving
+public class ResourcePattern {
+    /**
+     * A special literal resource name that corresponds to 'all resources of a certain type'.
+     */
+    public static final String WILDCARD_RESOURCE = "*";
+
+    private final ResourceType resourceType;
+    private final String name;
+    private final PatternType patternType;
+
+   /**
+     * Create a pattern using the supplied parameters.
+     *
+     * @param resourceType non-null, specific, resource type
+     * @param name non-null resource name, which can be the {@link #WILDCARD_RESOURCE}.
+     * @param patternType non-null, specific, resource pattern type, which controls how the pattern will match resource names.
+     */
+    public ResourcePattern(ResourceType resourceType, String name, PatternType patternType) {
+        this.resourceType = Objects.requireNonNull(resourceType, "resourceType");
+        this.name = Objects.requireNonNull(name, "name");
+        this.patternType = Objects.requireNonNull(patternType, "patternType");
+
+        if (resourceType == ResourceType.ANY) {
+            throw new IllegalArgumentException("resourceType must not be ANY");
+        }
+
+        if (patternType == PatternType.MATCH || patternType == PatternType.ANY) {
+            throw new IllegalArgumentException("patternType must not be " + patternType);
+        }
+    }
+
+    /**
+     * @return the specific resource type this pattern matches
+     */
+    public ResourceType resourceType() {
+        return resourceType;
+    }
+
+    /**
+     * @return the resource name.
+     */
+    public String name() {
+        return name;
+    }
+
+    /**
+     * @return the resource pattern type.
+     */
+    public PatternType patternType() {
+        return patternType;
+    }
+
+    /**
+     * @return a filter which matches only this pattern.
+     */
+    public ResourcePatternFilter toFilter() {
+        return new ResourcePatternFilter(resourceType, name, patternType);
+    }
+
+    @Override
+    public String toString() {
+        return "ResourcePattern(resourceType=" + resourceType + ", name=" + ((name == null) ? "" : name) + ", patternType=" + patternType + ")";
+    }
+
+    /**
+     * @return {@code true} if this Resource has any UNKNOWN components.
+     */
+    public boolean isUnknown() {
+        return resourceType.isUnknown() || patternType.isUnknown();
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o)
+            return true;
+        if (o == null || getClass() != o.getClass())
+            return false;
+
+        final ResourcePattern resource = (ResourcePattern) o;
+        return resourceType == resource.resourceType &&
+            Objects.equals(name, resource.name) &&
+            patternType == resource.patternType;
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(resourceType, name, patternType);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/resource/ResourcePatternFilter.java b/clients/src/main/java/org/apache/kafka/common/resource/ResourcePatternFilter.java
new file mode 100644
index 0000000000000..6f511c9a345d6
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/resource/ResourcePatternFilter.java
@@ -0,0 +1,170 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.resource;
+
+import org.apache.kafka.common.annotation.InterfaceStability;
+
+import java.util.Objects;
+
+import static org.apache.kafka.common.resource.ResourcePattern.WILDCARD_RESOURCE;
+
+/**
+ * Represents a filter that can match {@link ResourcePattern}.
+ * 
        
+ * The API for this class is still evolving and we may break compatibility in minor releases, if necessary.
+ */
+@InterfaceStability.Evolving
+public class ResourcePatternFilter {
+    /**
+     * Matches any resource pattern.
+     */
+    public static final ResourcePatternFilter ANY = new ResourcePatternFilter(ResourceType.ANY, null, PatternType.ANY);
+
+    private final ResourceType resourceType;
+    private final String name;
+    private final PatternType patternType;
+
+    /**
+     * Create a filter using the supplied parameters.
+     *
+     * @param resourceType non-null resource type.
+     *                     If {@link ResourceType#ANY}, the filter will ignore the resource type of the pattern.
+     *                     If any other resource type, the filter will match only patterns with the same type.
+     * @param name         resource name or {@code null}.
+     *                     If {@code null}, the filter will ignore the name of resources.
+     *                     If {@link ResourcePattern#WILDCARD_RESOURCE}, will match only wildcard patterns.
+     * @param patternType  non-null resource pattern type.
+     *                     If {@link PatternType#ANY}, the filter will match patterns regardless of pattern type.
+     *                     If {@link PatternType#MATCH}, the filter will match patterns that would match the supplied
+     *                     {@code name}, including a matching prefixed and wildcards patterns.
+     *                     If any other resource pattern type, the filter will match only patterns with the same type.
+     */
+    public ResourcePatternFilter(ResourceType resourceType, String name, PatternType patternType) {
+        this.resourceType = Objects.requireNonNull(resourceType, "resourceType");
+        this.name = name;
+        this.patternType = Objects.requireNonNull(patternType, "patternType");
+    }
+
+    /**
+     * @return {@code true} if this filter has any UNKNOWN components.
+     */
+    public boolean isUnknown() {
+        return resourceType.isUnknown() || patternType.isUnknown();
+    }
+
+    /**
+     * @return the specific resource type this pattern matches
+     */
+    public ResourceType resourceType() {
+        return resourceType;
+    }
+
+    /**
+     * @return the resource name.
+     */
+    public String name() {
+        return name;
+    }
+
+    /**
+     * @return the resource pattern type.
+     */
+    public PatternType patternType() {
+        return patternType;
+    }
+
+    /**
+     * @return {@code true} if this filter matches the given pattern.
+     */
+    public boolean matches(ResourcePattern pattern) {
+        if (!resourceType.equals(ResourceType.ANY) && !resourceType.equals(pattern.resourceType())) {
+            return false;
+        }
+
+        if (!patternType.equals(PatternType.ANY) && !patternType.equals(PatternType.MATCH) && !patternType.equals(pattern.patternType())) {
+            return false;
+        }
+
+        if (name == null) {
+            return true;
+        }
+
+        if (patternType.equals(PatternType.ANY) || patternType.equals(pattern.patternType())) {
+            return name.equals(pattern.name());
+        }
+
+        switch (pattern.patternType()) {
+            case LITERAL:
+                return name.equals(pattern.name()) || pattern.name().equals(WILDCARD_RESOURCE);
+
+            case PREFIXED:
+                return name.startsWith(pattern.name());
+
+            default:
+                throw new IllegalArgumentException("Unsupported PatternType: " + pattern.patternType());
+        }
+    }
+
+    /**
+     * @return {@code true} if this filter could only match one pattern.
+     * In other words, if there are no ANY or UNKNOWN fields.
+     */
+    public boolean matchesAtMostOne() {
+        return findIndefiniteField() == null;
+    }
+
+    /**
+     * @return a string describing any ANY or UNKNOWN field, or null if there is no such field.
+     */
+    public String findIndefiniteField() {
+        if (resourceType == ResourceType.ANY)
+            return "Resource type is ANY.";
+        if (resourceType == ResourceType.UNKNOWN)
+            return "Resource type is UNKNOWN.";
+        if (name == null)
+            return "Resource name is NULL.";
+        if (patternType == PatternType.MATCH)
+            return "Resource pattern type is MATCH.";
+        if (patternType == PatternType.UNKNOWN)
+            return "Resource pattern type is UNKNOWN.";
+        return null;
+    }
+
+    @Override
+    public String toString() {
+        return "ResourcePattern(resourceType=" + resourceType + ", name=" + ((name == null) ? "" : name) + ", patternType=" + patternType + ")";
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o)
+            return true;
+        if (o == null || getClass() != o.getClass())
+            return false;
+
+        final ResourcePatternFilter resource = (ResourcePatternFilter) o;
+        return resourceType == resource.resourceType &&
+            Objects.equals(name, resource.name) &&
+            patternType == resource.patternType;
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(resourceType, name, patternType);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/resource/ResourceType.java b/clients/src/main/java/org/apache/kafka/common/resource/ResourceType.java
index d83382d21d26f..2ce653fbeb2ec 100644
--- a/clients/src/main/java/org/apache/kafka/common/resource/ResourceType.java
+++ b/clients/src/main/java/org/apache/kafka/common/resource/ResourceType.java
@@ -58,7 +58,12 @@ public enum ResourceType {
     /**
      * A transactional ID.
      */
-    TRANSACTIONAL_ID((byte) 5);
+    TRANSACTIONAL_ID((byte) 5),
+
+    /**
+     * A token ID.
+     */
+    DELEGATION_TOKEN((byte) 6);
 
     private final static HashMap CODE_TO_VALUE = new HashMap<>();
 
diff --git a/clients/src/main/java/org/apache/kafka/common/security/JaasConfig.java b/clients/src/main/java/org/apache/kafka/common/security/JaasConfig.java
index 24bdac2378728..5e837a69c1602 100644
--- a/clients/src/main/java/org/apache/kafka/common/security/JaasConfig.java
+++ b/clients/src/main/java/org/apache/kafka/common/security/JaasConfig.java
@@ -81,6 +81,9 @@ public AppConfigurationEntry[] getAppConfigurationEntry(String name) {
     }
 
     private LoginModuleControlFlag loginModuleControlFlag(String flag) {
+        if (flag == null)
+            throw new IllegalArgumentException("Login module control flag is not available in the JAAS config");
+
         LoginModuleControlFlag controlFlag;
         switch (flag.toUpperCase(Locale.ROOT)) {
             case "REQUIRED":
diff --git a/clients/src/main/java/org/apache/kafka/common/security/JaasContext.java b/clients/src/main/java/org/apache/kafka/common/security/JaasContext.java
index a13acd2dc6d4a..48216a8a90c4c 100644
--- a/clients/src/main/java/org/apache/kafka/common/security/JaasContext.java
+++ b/clients/src/main/java/org/apache/kafka/common/security/JaasContext.java
@@ -33,7 +33,7 @@
 
 public class JaasContext {
 
-    private static final Logger LOG = LoggerFactory.getLogger(JaasUtils.class);
+    private static final Logger LOG = LoggerFactory.getLogger(JaasContext.class);
 
     private static final String GLOBAL_CONTEXT_NAME_SERVER = "KafkaServer";
     private static final String GLOBAL_CONTEXT_NAME_CLIENT = "KafkaClient";
@@ -41,57 +41,57 @@ public class JaasContext {
     /**
      * Returns an instance of this class.
      *
-     * For contextType SERVER, the context will contain the default Configuration and the context name will be one of:
+     * The context will contain the configuration specified by the JAAS configuration property
+     * {@link SaslConfigs#SASL_JAAS_CONFIG} with prefix `listener.name.{listenerName}.{mechanism}.`
+     * with listenerName and mechanism in lower case. The context `KafkaServer` will be returned
+     * with a single login context entry loaded from the property.
+     * 
        
+     * If the property is not defined, the context will contain the default Configuration and
+     * the context name will be one of:
+     * 
          
+     *   
        1. Lowercased listener name followed by a period and the string `KafkaServer`
          2. 
+     *   
        2. The string `KafkaServer`
          3. 
+     *  
        
+     * If both are valid entries in the default JAAS configuration, the first option is chosen.
+     * 
        
      *
-     * 1. Lowercased listener name followed by a period and the string `KafkaServer`
-     * 2. The string `KafkaServer`
-     *
-     * If both are valid entries in the JAAS configuration, the first option is chosen.
+     * @throws IllegalArgumentException if listenerName or mechanism is not defined.
+     */
+    public static JaasContext loadServerContext(ListenerName listenerName, String mechanism, Map configs) {
+        if (listenerName == null)
+            throw new IllegalArgumentException("listenerName should not be null for SERVER");
+        if (mechanism == null)
+            throw new IllegalArgumentException("mechanism should not be null for SERVER");
+        String listenerContextName = listenerName.value().toLowerCase(Locale.ROOT) + "." + GLOBAL_CONTEXT_NAME_SERVER;
+        Password dynamicJaasConfig = (Password) configs.get(mechanism.toLowerCase(Locale.ROOT) + "." + SaslConfigs.SASL_JAAS_CONFIG);
+        if (dynamicJaasConfig == null && configs.get(SaslConfigs.SASL_JAAS_CONFIG) != null)
+            LOG.warn("Server config {} should be prefixed with SASL mechanism name, ignoring config", SaslConfigs.SASL_JAAS_CONFIG);
+        return load(Type.SERVER, listenerContextName, GLOBAL_CONTEXT_NAME_SERVER, dynamicJaasConfig);
+    }
+
+    /**
+     * Returns an instance of this class.
      *
-     * For contextType CLIENT, if JAAS configuration property @link SaslConfigs#SASL_JAAS_CONFIG} is specified,
+     * If JAAS configuration property @link SaslConfigs#SASL_JAAS_CONFIG} is specified,
      * the configuration object is created by parsing the property value. Otherwise, the default Configuration
      * is returned. The context name is always `KafkaClient`.
      *
-     * @throws IllegalArgumentException if JAAS configuration property is specified for contextType SERVER, if
-     * listenerName is not defined for contextType SERVER of if listenerName is defined for contextType CLIENT.
      */
-    public static JaasContext load(JaasContext.Type contextType, ListenerName listenerName,
-                                   Map configs) {
-        String listenerContextName;
-        String globalContextName;
-        switch (contextType) {
-            case CLIENT:
-                if (listenerName != null)
-                    throw new IllegalArgumentException("listenerName should be null for CLIENT");
-                globalContextName = GLOBAL_CONTEXT_NAME_CLIENT;
-                listenerContextName = null;
-                break;
-            case SERVER:
-                if (listenerName == null)
-                    throw new IllegalArgumentException("listenerName should not be null for SERVER");
-                globalContextName = GLOBAL_CONTEXT_NAME_SERVER;
-                listenerContextName = listenerName.value().toLowerCase(Locale.ROOT) + "." + GLOBAL_CONTEXT_NAME_SERVER;
-                break;
-            default:
-                throw new IllegalArgumentException("Unexpected context type " + contextType);
-        }
-        return load(contextType, listenerContextName, globalContextName, configs);
+    public static JaasContext loadClientContext(Map configs) {
+        Password dynamicJaasConfig = (Password) configs.get(SaslConfigs.SASL_JAAS_CONFIG);
+        return load(JaasContext.Type.CLIENT, null, GLOBAL_CONTEXT_NAME_CLIENT, dynamicJaasConfig);
     }
 
     static JaasContext load(JaasContext.Type contextType, String listenerContextName,
-                            String globalContextName, Map configs) {
-        Password jaasConfigArgs = (Password) configs.get(SaslConfigs.SASL_JAAS_CONFIG);
-        if (jaasConfigArgs != null) {
-            if (contextType == JaasContext.Type.SERVER)
-                throw new IllegalArgumentException("JAAS config property not supported for server");
-            else {
-                JaasConfig jaasConfig = new JaasConfig(globalContextName, jaasConfigArgs.value());
-                AppConfigurationEntry[] clientModules = jaasConfig.getAppConfigurationEntry(globalContextName);
-                int numModules = clientModules == null ? 0 : clientModules.length;
-                if (numModules != 1)
-                    throw new IllegalArgumentException("JAAS config property contains " + numModules + " login modules, should be 1 module");
-                return new JaasContext(globalContextName, contextType, jaasConfig);
-            }
+                            String globalContextName, Password dynamicJaasConfig) {
+        if (dynamicJaasConfig != null) {
+            JaasConfig jaasConfig = new JaasConfig(globalContextName, dynamicJaasConfig.value());
+            AppConfigurationEntry[] contextModules = jaasConfig.getAppConfigurationEntry(globalContextName);
+            if (contextModules == null || contextModules.length == 0)
+                throw new IllegalArgumentException("JAAS config property does not contain any login modules");
+            else if (contextModules.length != 1)
+                throw new IllegalArgumentException("JAAS config property contains " + contextModules.length + " login modules, should be 1 module");
+            return new JaasContext(globalContextName, contextType, jaasConfig, dynamicJaasConfig);
         } else
             return defaultContext(contextType, listenerContextName, globalContextName);
     }
@@ -131,21 +131,22 @@ private static JaasContext defaultContext(JaasContext.Type contextType, String l
             throw new IllegalArgumentException(errorMessage);
         }
 
-        return new JaasContext(contextName, contextType, jaasConfig);
+        return new JaasContext(contextName, contextType, jaasConfig, null);
     }
 
     /**
      * The type of the SASL login context, it should be SERVER for the broker and CLIENT for the clients (consumer, producer,
      * etc.). This is used to validate behaviour (e.g. some functionality is only available in the broker or clients).
      */
-    public enum Type { CLIENT, SERVER; }
+    public enum Type { CLIENT, SERVER }
 
     private final String name;
     private final Type type;
     private final Configuration configuration;
     private final List configurationEntries;
+    private final Password dynamicJaasConfig;
 
-    public JaasContext(String name, Type type, Configuration configuration) {
+    public JaasContext(String name, Type type, Configuration configuration, Password dynamicJaasConfig) {
         this.name = name;
         this.type = type;
         this.configuration = configuration;
@@ -153,6 +154,7 @@ public JaasContext(String name, Type type, Configuration configuration) {
         if (entries == null)
             throw new IllegalArgumentException("Could not find a '" + name + "' entry in this JAAS configuration.");
         this.configurationEntries = Collections.unmodifiableList(new ArrayList<>(Arrays.asList(entries)));
+        this.dynamicJaasConfig = dynamicJaasConfig;
     }
 
     public String name() {
@@ -171,11 +173,15 @@ public List configurationEntries() {
         return configurationEntries;
     }
 
+    public Password dynamicJaasConfig() {
+        return dynamicJaasConfig;
+    }
+
     /**
      * Returns the configuration option for key from this context.
      * If login module name is specified, return option value only from that module.
      */
-    public String configEntryOption(String key, String loginModuleName) {
+    public static String configEntryOption(List configurationEntries, String key, String loginModuleName) {
         for (AppConfigurationEntry entry : configurationEntries) {
             if (loginModuleName != null && !loginModuleName.equals(entry.getLoginModuleName()))
                 continue;
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 08f41c3349be5..baff5633a3f6d 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,13 +16,13 @@
  */
 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;
 
-public class JaasUtils {
+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,25 +31,51 @@ public class JaasUtils {
     public static final String ZK_SASL_CLIENT = "zookeeper.sasl.client";
     public static final String ZK_LOGIN_CONTEXT_NAME_KEY = "zookeeper.sasl.clientconfig";
 
-    public static boolean isZkSecurityEnabled() {
-        boolean zkSaslEnabled = Boolean.parseBoolean(System.getProperty(ZK_SASL_CLIENT, "true"));
-        String zkLoginContextName = System.getProperty(ZK_LOGIN_CONTEXT_NAME_KEY, "Client");
+    private static final String DEFAULT_ZK_LOGIN_CONTEXT_NAME = "Client";
+    private static final String DEFAULT_ZK_SASL_CLIENT = "true";
+
+    private JaasUtils() {}
 
-        boolean isSecurityEnabled;
+    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 isZkSaslEnabled() {
+        // Technically a client must also check if TLS mutual authentication has been configured,
+        // but we will leave that up to the client code to determine since direct connectivity to ZooKeeper
+        // has been deprecated in many clients and we don't wish to re-introduce a ZooKeeper jar dependency here.
+        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);
+
+        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/AuthenticateCallbackHandler.java b/clients/src/main/java/org/apache/kafka/common/security/auth/AuthenticateCallbackHandler.java
new file mode 100644
index 0000000000000..8951d3a589367
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/auth/AuthenticateCallbackHandler.java
@@ -0,0 +1,62 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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 java.util.List;
+import java.util.Map;
+
+import javax.security.auth.callback.CallbackHandler;
+import javax.security.auth.login.AppConfigurationEntry;
+
+/*
+ * Callback handler for SASL-based authentication
+ */
+public interface AuthenticateCallbackHandler extends CallbackHandler {
+
+    /**
+     * Configures this callback handler for the specified SASL mechanism.
+     *
+     * @param configs Key-value pairs containing the parsed configuration options of
+     *        the client or broker. Note that these are the Kafka configuration options
+     *        and not the JAAS configuration options. JAAS config options may be obtained
+     *        from `jaasConfigEntries` for callbacks which obtain some configs from the
+     *        JAAS configuration. For configs that may be specified as both Kafka config
+     *        as well as JAAS config (e.g. sasl.kerberos.service.name), the configuration
+     *        is treated as invalid if conflicting values are provided.
+     * @param saslMechanism Negotiated SASL mechanism. For clients, this is the SASL
+     *        mechanism configured for the client. For brokers, this is the mechanism
+     *        negotiated with the client and is one of the mechanisms enabled on the broker.
+     * @param jaasConfigEntries JAAS configuration entries from the JAAS login context.
+     *        This list contains a single entry for clients and may contain more than
+     *        one entry for brokers if multiple mechanisms are enabled on a listener using
+     *        static JAAS configuration where there is no mapping between mechanisms and
+     *        login module entries. In this case, callback handlers can use the login module in
+     *        `jaasConfigEntries` to identify the entry corresponding to `saslMechanism`.
+     *        Alternatively, dynamic JAAS configuration option
+     *        {@link org.apache.kafka.common.config.SaslConfigs#SASL_JAAS_CONFIG} may be
+     *        configured on brokers with listener and mechanism prefix, in which case
+     *        only the configuration entry corresponding to `saslMechanism` will be provided
+     *        in `jaasConfigEntries`.
+     */
+    void configure(Map configs, String saslMechanism, List jaasConfigEntries);
+
+    /**
+     * Closes this instance.
+     */
+    void close();
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/auth/AuthenticationContext.java b/clients/src/main/java/org/apache/kafka/common/security/auth/AuthenticationContext.java
index b8c08477027cf..a8abea858960d 100644
--- a/clients/src/main/java/org/apache/kafka/common/security/auth/AuthenticationContext.java
+++ b/clients/src/main/java/org/apache/kafka/common/security/auth/AuthenticationContext.java
@@ -18,9 +18,11 @@
 
 import java.net.InetAddress;
 
+
 /**
  * An object representing contextual information from the authentication session. See
- * {@link SaslAuthenticationContext} and {@link SslAuthenticationContext}.
+ * {@link PlaintextAuthenticationContext}, {@link SaslAuthenticationContext}
+ * and {@link SslAuthenticationContext}. This class is only used in the broker.
  */
 public interface AuthenticationContext {
     /**
@@ -32,4 +34,9 @@ public interface AuthenticationContext {
      * Address of the authenticated client
      */
     InetAddress clientAddress();
+
+    /**
+     * Name of the listener used for the connection
+     */
+    String listenerName();
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/security/auth/KafkaPrincipal.java b/clients/src/main/java/org/apache/kafka/common/security/auth/KafkaPrincipal.java
index ed3c9563d2e2e..1a8669ac475c5 100644
--- a/clients/src/main/java/org/apache/kafka/common/security/auth/KafkaPrincipal.java
+++ b/clients/src/main/java/org/apache/kafka/common/security/auth/KafkaPrincipal.java
@@ -23,23 +23,24 @@
 import static java.util.Objects.requireNonNull;
 
 /**
- * Principals in Kafka are defined by a type and a name. The principal type will always be "User"
+ * 
        Principals in Kafka are defined by a type and a name. The principal type will always be "User"
  * for the simple authorizer that is enabled by default, but custom authorizers can leverage different
  * principal types (such as to enable group or role-based ACLs). The {@link KafkaPrincipalBuilder} interface
  * is used when you need to derive a different principal type from the authentication context, or when
  * you need to represent relations between different principals. For example, you could extend
  * {@link KafkaPrincipal} in order to link a user principal to one or more role principals.
  *
- * For custom extensions of {@link KafkaPrincipal}, there two key points to keep in mind:
- *
- * 1. To be compatible with the ACL APIs provided by Kafka (including the command line tool), each ACL
+ * 
        For custom extensions of {@link KafkaPrincipal}, there two key points to keep in mind:
+ * 
          
+ * 
        1. To be compatible with the ACL APIs provided by Kafka (including the command line tool), each ACL
  *    can only represent a permission granted to a single principal (consisting of a principal type and name).
  *    It is possible to use richer ACL semantics, but you must implement your own mechanisms for adding
  *    and removing ACLs.
- * 2. In general, {@link KafkaPrincipal} extensions are only useful when the corresponding Authorizer
+ * 
        2. In general, {@link KafkaPrincipal} extensions are only useful when the corresponding Authorizer
  *    is also aware of the extension. If you have a {@link KafkaPrincipalBuilder} which derives user groups
  *    from the authentication context (e.g. from an SSL client certificate), then you need a custom
  *    authorizer which is capable of using the additional group information.
+ * 
        
  */
 public class KafkaPrincipal implements Principal {
     public static final String USER_TYPE = "User";
@@ -47,10 +48,16 @@ public class KafkaPrincipal implements Principal {
 
     private final String principalType;
     private final String name;
+    private volatile boolean tokenAuthenticated;
 
     public KafkaPrincipal(String principalType, String name) {
+        this(principalType, name, false);
+    }
+
+    public KafkaPrincipal(String principalType, String name, boolean tokenAuthenticated) {
         this.principalType = requireNonNull(principalType, "Principal type cannot be null");
         this.name = requireNonNull(name, "Principal name cannot be null");
+        this.tokenAuthenticated = tokenAuthenticated;
     }
 
     /**
@@ -83,8 +90,8 @@ public boolean equals(Object o) {
 
     @Override
     public int hashCode() {
-        int result = principalType.hashCode();
-        result = 31 * result + name.hashCode();
+        int result = principalType != null ? principalType.hashCode() : 0;
+        result = 31 * result + (name != null ? name.hashCode() : 0);
         return result;
     }
 
@@ -97,4 +104,12 @@ public String getPrincipalType() {
         return principalType;
     }
 
+    public void tokenAuthenticated(boolean tokenAuthenticated) {
+        this.tokenAuthenticated = tokenAuthenticated;
+    }
+
+    public boolean tokenAuthenticated() {
+        return tokenAuthenticated;
+    }
 }
+
diff --git a/clients/src/main/java/org/apache/kafka/common/security/auth/KafkaPrincipalSerde.java b/clients/src/main/java/org/apache/kafka/common/security/auth/KafkaPrincipalSerde.java
new file mode 100644
index 0000000000000..02dbdb639171f
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/auth/KafkaPrincipalSerde.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.errors.SerializationException;
+
+/**
+ * Serializer/Deserializer interface for {@link KafkaPrincipal} for the the purpose of inter-broker forwarding.
+ * Any serialization/deserialization failure should raise a {@link SerializationException} to be consistent.
+ */
+public interface KafkaPrincipalSerde {
+
+    /**
+     * Serialize a {@link KafkaPrincipal} into byte array.
+     *
+     * @param principal principal to be serialized
+     * @return serialized bytes
+     * @throws SerializationException
+     */
+    byte[] serialize(KafkaPrincipal principal) throws SerializationException;
+
+    /**
+     * Deserialize a {@link KafkaPrincipal} from byte array.
+     * @param bytes byte array to be deserialized
+     * @return the deserialized principal
+     * @throws SerializationException
+     */
+    KafkaPrincipal deserialize(byte[] bytes) throws SerializationException;
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/auth/Login.java b/clients/src/main/java/org/apache/kafka/common/security/auth/Login.java
new file mode 100644
index 0000000000000..eda5e7a225af1
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/auth/Login.java
@@ -0,0 +1,68 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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 java.util.Map;
+
+import javax.security.auth.Subject;
+import javax.security.auth.login.Configuration;
+import javax.security.auth.login.LoginContext;
+import javax.security.auth.login.LoginException;
+
+/**
+ * Login interface for authentication.
+ */
+public interface Login {
+
+    /**
+     * Configures this login instance.
+     * @param configs Key-value pairs containing the parsed configuration options of
+     *        the client or broker. Note that these are the Kafka configuration options
+     *        and not the JAAS configuration options. The JAAS options may be obtained
+     *        from `jaasConfiguration`.
+     * @param contextName JAAS context name for this login which may be used to obtain
+     *        the login context from `jaasConfiguration`.
+     * @param jaasConfiguration JAAS configuration containing the login context named
+     *        `contextName`. If static JAAS configuration is used, this `Configuration`
+     *         may also contain other login contexts.
+     * @param loginCallbackHandler Login callback handler instance to use for this Login.
+     *        Login callback handler class may be configured using
+     *        {@link org.apache.kafka.common.config.SaslConfigs#SASL_LOGIN_CALLBACK_HANDLER_CLASS}.
+     */
+    void configure(Map configs, String contextName, Configuration jaasConfiguration,
+                   AuthenticateCallbackHandler loginCallbackHandler);
+
+    /**
+     * Performs login for each login module specified for the login context of this instance.
+     */
+    LoginContext login() throws LoginException;
+
+    /**
+     * Returns the authenticated subject of this login context.
+     */
+    Subject subject();
+
+    /**
+     * Returns the service name to be used for SASL.
+     */
+    String serviceName();
+
+    /**
+     * Closes this instance.
+     */
+    void close();
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/auth/PlaintextAuthenticationContext.java b/clients/src/main/java/org/apache/kafka/common/security/auth/PlaintextAuthenticationContext.java
index bc14d36aaf540..a111f21ec1302 100644
--- a/clients/src/main/java/org/apache/kafka/common/security/auth/PlaintextAuthenticationContext.java
+++ b/clients/src/main/java/org/apache/kafka/common/security/auth/PlaintextAuthenticationContext.java
@@ -20,9 +20,11 @@
 
 public class PlaintextAuthenticationContext implements AuthenticationContext {
     private final InetAddress clientAddress;
+    private final String listenerName;
 
-    public PlaintextAuthenticationContext(InetAddress clientAddress) {
+    public PlaintextAuthenticationContext(InetAddress clientAddress, String listenerName) {
         this.clientAddress = clientAddress;
+        this.listenerName = listenerName;
     }
 
     @Override
@@ -35,4 +37,9 @@ public InetAddress clientAddress() {
         return clientAddress;
     }
 
+    @Override
+    public String listenerName() {
+        return listenerName;
+    }
+
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/security/auth/SaslAuthenticationContext.java b/clients/src/main/java/org/apache/kafka/common/security/auth/SaslAuthenticationContext.java
index 89e606377d9eb..719d041770d13 100644
--- a/clients/src/main/java/org/apache/kafka/common/security/auth/SaslAuthenticationContext.java
+++ b/clients/src/main/java/org/apache/kafka/common/security/auth/SaslAuthenticationContext.java
@@ -17,17 +17,20 @@
 package org.apache.kafka.common.security.auth;
 
 import javax.security.sasl.SaslServer;
+
 import java.net.InetAddress;
 
 public class SaslAuthenticationContext implements AuthenticationContext {
     private final SaslServer server;
     private final SecurityProtocol securityProtocol;
     private final InetAddress clientAddress;
+    private final String listenerName;
 
-    public SaslAuthenticationContext(SaslServer server, SecurityProtocol securityProtocol, InetAddress clientAddress) {
+    public SaslAuthenticationContext(SaslServer server, SecurityProtocol securityProtocol, InetAddress clientAddress, String listenerName) {
         this.server = server;
         this.securityProtocol = securityProtocol;
         this.clientAddress = clientAddress;
+        this.listenerName = listenerName;
     }
 
     public SaslServer server() {
@@ -43,4 +46,9 @@ public SecurityProtocol securityProtocol() {
     public InetAddress clientAddress() {
         return clientAddress;
     }
+
+    @Override
+    public String listenerName() {
+        return listenerName;
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/security/auth/SaslExtensions.java b/clients/src/main/java/org/apache/kafka/common/security/auth/SaslExtensions.java
new file mode 100644
index 0000000000000..c129f1ec400f7
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/auth/SaslExtensions.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.security.auth;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * A simple immutable value object class holding customizable SASL extensions
+ */
+public class SaslExtensions {
+    /**
+     * An "empty" instance indicating no SASL extensions
+     */
+    public static final SaslExtensions NO_SASL_EXTENSIONS = new SaslExtensions(Collections.emptyMap());
+    private final Map extensionsMap;
+
+    public SaslExtensions(Map extensionsMap) {
+        this.extensionsMap = Collections.unmodifiableMap(new HashMap<>(extensionsMap));
+    }
+
+    /**
+     * Returns an immutable map of the extension names and their values
+     */
+    public Map map() {
+        return extensionsMap;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        return extensionsMap.equals(((SaslExtensions) o).extensionsMap);
+    }
+
+    @Override
+    public String toString() {
+        return extensionsMap.toString();
+    }
+
+    @Override
+    public int hashCode() {
+        return extensionsMap.hashCode();
+    }
+
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/auth/SaslExtensionsCallback.java b/clients/src/main/java/org/apache/kafka/common/security/auth/SaslExtensionsCallback.java
new file mode 100644
index 0000000000000..c5bd449e0cc08
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/auth/SaslExtensionsCallback.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.security.auth;
+
+import java.util.Objects;
+
+import javax.security.auth.callback.Callback;
+
+/**
+ * Optional callback used for SASL mechanisms if any extensions need to be set
+ * in the SASL exchange.
+ */
+public class SaslExtensionsCallback implements Callback {
+    private SaslExtensions extensions = SaslExtensions.NO_SASL_EXTENSIONS;
+
+    /**
+     * Returns always non-null {@link SaslExtensions} consisting of the extension
+     * names and values that are sent by the client to the server in the initial
+     * client SASL authentication message. The default value is
+     * {@link SaslExtensions#NO_SASL_EXTENSIONS} so that if this callback is
+     * unhandled the client will see a non-null value.
+     */
+    public SaslExtensions extensions() {
+        return extensions;
+    }
+
+    /**
+     * Sets the SASL extensions on this callback.
+     * 
+     * @param extensions
+     *            the mandatory extensions to set
+     */
+    public void extensions(SaslExtensions extensions) {
+        this.extensions = Objects.requireNonNull(extensions, "extensions must not be null");
+    }
+}
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/auth/SslAuthenticationContext.java b/clients/src/main/java/org/apache/kafka/common/security/auth/SslAuthenticationContext.java
index d87a8920b479e..88819f91cf82f 100644
--- a/clients/src/main/java/org/apache/kafka/common/security/auth/SslAuthenticationContext.java
+++ b/clients/src/main/java/org/apache/kafka/common/security/auth/SslAuthenticationContext.java
@@ -22,10 +22,12 @@
 public class SslAuthenticationContext implements AuthenticationContext {
     private final SSLSession session;
     private final InetAddress clientAddress;
+    private final String listenerName;
 
-    public SslAuthenticationContext(SSLSession session, InetAddress clientAddress) {
+    public SslAuthenticationContext(SSLSession session, InetAddress clientAddress, String listenerName) {
         this.session = session;
         this.clientAddress = clientAddress;
+        this.listenerName = listenerName;
     }
 
     public SSLSession session() {
@@ -41,4 +43,9 @@ public SecurityProtocol securityProtocol() {
     public InetAddress clientAddress() {
         return clientAddress;
     }
+
+    @Override
+    public String listenerName() {
+        return listenerName;
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/security/auth/SslEngineFactory.java b/clients/src/main/java/org/apache/kafka/common/security/auth/SslEngineFactory.java
new file mode 100644
index 0000000000000..586017d55f226
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/auth/SslEngineFactory.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.common.security.auth;
+
+import org.apache.kafka.common.Configurable;
+
+import javax.net.ssl.SSLEngine;
+import java.io.Closeable;
+import java.security.KeyStore;
+import java.util.Map;
+import java.util.Set;
+
+/**
+ * Plugin interface for allowing creation of SSLEngine object in a custom way.
+ * For example, you can use this to customize loading your key material and trust material needed for SSLContext.
+ * This is complementary to the existing Java Security Provider mechanism which allows the entire provider
+ * to be replaced with a custom provider. In scenarios where only the configuration mechanism for SSL engines
+ * need to be updated, this interface provides a convenient method for overriding the default implementation.
+ */
+public interface SslEngineFactory extends Configurable, Closeable {
+
+    /**
+     * Creates a new SSLEngine object to be used by the client.
+     *
+     * @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.
+     */
+    SSLEngine createClientSslEngine(String peerHost, int peerPort, String endpointIdentification);
+
+    /**
+     * Creates a new SSLEngine object to be used by the server.
+     *
+     * @param peerHost               The peer host to use. This is a hint and not used for validation.
+     * @param peerPort               The peer port to use. This is a hint and not used for validation.
+     * @return The new SSLEngine.
+     */
+    SSLEngine createServerSslEngine(String peerHost, int peerPort);
+
+    /**
+     * Returns true if SSLEngine needs to be rebuilt. This method will be called when reconfiguration is triggered on
+     * the SslFactory used to create SSL engines. Based on the new configs provided in nextConfigs, this method
+     * will decide whether underlying SSLEngine object needs to be rebuilt. If this method returns true, the
+     * SslFactory will create a new instance of this object with nextConfigs and run other
+     * checks before deciding to use the new object for new incoming connection requests. Existing connections
+     * are not impacted by this and will not see any changes done as part of reconfiguration.
+     * 
        
+     * For example, if the implementation depends on file-based key material, it can check if the file was updated
+     * compared to the previous/last-loaded timestamp and return true.
+     * 
        
+     *
+     * @param nextConfigs       The new configuration we want to use.
+     * @return                  True only if the underlying SSLEngine object should be rebuilt.
+     */
+    boolean shouldBeRebuilt(Map nextConfigs);
+
+    /**
+     * Returns the names of configs that may be reconfigured.
+     * @return Names of configuration options that are dynamically reconfigurable.
+     */
+    Set reconfigurableConfigs();
+
+    /**
+     * Returns keystore configured for this factory.
+     * @return The keystore for this factory or null if a keystore is not configured.
+     */
+    KeyStore keystore();
+
+    /**
+     * Returns truststore configured for this factory.
+     * @return The truststore for this factory or null if a truststore is not configured.
+     */
+    KeyStore truststore();
+}
\ No newline at end of file
diff --git a/clients/src/main/java/org/apache/kafka/common/security/authenticator/AbstractLogin.java b/clients/src/main/java/org/apache/kafka/common/security/authenticator/AbstractLogin.java
index 643f859e8298c..7e1350864e438 100644
--- a/clients/src/main/java/org/apache/kafka/common/security/authenticator/AbstractLogin.java
+++ b/clients/src/main/java/org/apache/kafka/common/security/authenticator/AbstractLogin.java
@@ -16,20 +16,23 @@
  */
 package org.apache.kafka.common.security.authenticator;
 
+import javax.security.auth.login.AppConfigurationEntry;
+import javax.security.auth.login.Configuration;
 import javax.security.auth.login.LoginContext;
 import javax.security.auth.login.LoginException;
 import javax.security.sasl.RealmCallback;
 import javax.security.auth.callback.Callback;
-import javax.security.auth.callback.CallbackHandler;
 import javax.security.auth.callback.NameCallback;
 import javax.security.auth.callback.PasswordCallback;
 import javax.security.auth.callback.UnsupportedCallbackException;
 import javax.security.auth.Subject;
 
-import org.apache.kafka.common.security.JaasContext;
+import org.apache.kafka.common.security.auth.AuthenticateCallbackHandler;
+import org.apache.kafka.common.security.auth.Login;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.util.List;
 import java.util.Map;
 
 /**
@@ -38,17 +41,22 @@
 public abstract class AbstractLogin implements Login {
     private static final Logger log = LoggerFactory.getLogger(AbstractLogin.class);
 
-    private JaasContext jaasContext;
+    private String contextName;
+    private Configuration configuration;
     private LoginContext loginContext;
+    private AuthenticateCallbackHandler loginCallbackHandler;
 
     @Override
-    public void configure(Map configs, JaasContext jaasContext) {
-        this.jaasContext = jaasContext;
+    public void configure(Map configs, String contextName, Configuration configuration,
+                          AuthenticateCallbackHandler loginCallbackHandler) {
+        this.contextName = contextName;
+        this.configuration = configuration;
+        this.loginCallbackHandler = loginCallbackHandler;
     }
 
     @Override
     public LoginContext login() throws LoginException {
-        loginContext = new LoginContext(jaasContext.name(), null, new LoginCallbackHandler(), jaasContext.configuration());
+        loginContext = new LoginContext(contextName, null, loginCallbackHandler, configuration);
         loginContext.login();
         log.info("Successfully logged in.");
         return loginContext;
@@ -59,8 +67,12 @@ public Subject subject() {
         return loginContext.getSubject();
     }
 
-    protected JaasContext jaasContext() {
-        return jaasContext;
+    protected String contextName() {
+        return contextName;
+    }
+
+    protected Configuration configuration() {
+        return configuration;
     }
 
     /**
@@ -70,7 +82,11 @@ protected JaasContext jaasContext() {
      * callback handlers which require additional user input.
      *
      */
-    public static class LoginCallbackHandler implements CallbackHandler {
+    public static class DefaultLoginCallbackHandler implements AuthenticateCallbackHandler {
+
+        @Override
+        public void configure(Map configs, String saslMechanism, List jaasConfigEntries) {
+        }
 
         @Override
         public void handle(Callback[] callbacks) throws UnsupportedCallbackException {
@@ -90,6 +106,10 @@ public void handle(Callback[] callbacks) throws UnsupportedCallbackException {
                 }
             }
         }
+
+        @Override
+        public void close() {
+        }
     }
 }
 
diff --git a/clients/src/main/java/org/apache/kafka/common/security/authenticator/AuthCallbackHandler.java b/clients/src/main/java/org/apache/kafka/common/security/authenticator/AuthCallbackHandler.java
deleted file mode 100644
index d517162a1762c..0000000000000
--- a/clients/src/main/java/org/apache/kafka/common/security/authenticator/AuthCallbackHandler.java
+++ /dev/null
@@ -1,45 +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.security.authenticator;
-
-import java.util.Map;
-
-import org.apache.kafka.common.network.Mode;
-
-import javax.security.auth.Subject;
-import javax.security.auth.callback.CallbackHandler;
-
-/*
- * Callback handler for SASL-based authentication
- */
-public interface AuthCallbackHandler extends CallbackHandler {
-
-    /**
-     * Configures this callback handler.
-     *
-     * @param configs Configuration
-     * @param mode The mode that indicates if this is a client or server connection
-     * @param subject Subject from login context
-     * @param saslMechanism Negotiated SASL mechanism
-     */
-    void configure(Map configs, Mode mode, Subject subject, String saslMechanism);
-
-    /**
-     * Closes this instance.
-     */
-    void close();
-}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/authenticator/CredentialCache.java b/clients/src/main/java/org/apache/kafka/common/security/authenticator/CredentialCache.java
index aa39101753fe0..ecf3ea9d4a6f1 100644
--- a/clients/src/main/java/org/apache/kafka/common/security/authenticator/CredentialCache.java
+++ b/clients/src/main/java/org/apache/kafka/common/security/authenticator/CredentialCache.java
@@ -16,18 +16,17 @@
  */
 package org.apache.kafka.common.security.authenticator;
 
-import java.util.HashMap;
-import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 
 public class CredentialCache {
 
-    private final Map> cacheMap = new HashMap<>();
+    private final ConcurrentHashMap> cacheMap = new ConcurrentHashMap<>();
 
     public  Cache createCache(String mechanism, Class credentialClass) {
-        Cache cache = new Cache(credentialClass);
-        cacheMap.put(mechanism, cache);
-        return cache;
+        Cache cache = new Cache<>(credentialClass);
+        @SuppressWarnings("unchecked")
+        Cache oldCache = (Cache) cacheMap.putIfAbsent(mechanism, cache);
+        return oldCache == null ? cache : oldCache;
     }
 
     @SuppressWarnings("unchecked")
@@ -66,4 +65,4 @@ public Class credentialClass() {
             return credentialClass;
         }
     }
-}
\ No newline at end of file
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/authenticator/DefaultKafkaPrincipalBuilder.java b/clients/src/main/java/org/apache/kafka/common/security/authenticator/DefaultKafkaPrincipalBuilder.java
index 30b0a3e0980f0..c459316eca63c 100644
--- a/clients/src/main/java/org/apache/kafka/common/security/authenticator/DefaultKafkaPrincipalBuilder.java
+++ b/clients/src/main/java/org/apache/kafka/common/security/authenticator/DefaultKafkaPrincipalBuilder.java
@@ -16,13 +16,19 @@
  */
 package org.apache.kafka.common.security.authenticator;
 
+import javax.security.auth.x500.X500Principal;
+
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.config.SaslConfigs;
+import org.apache.kafka.common.errors.SerializationException;
+import org.apache.kafka.common.message.DefaultPrincipalData;
 import org.apache.kafka.common.network.Authenticator;
 import org.apache.kafka.common.network.TransportLayer;
+import org.apache.kafka.common.protocol.types.Struct;
 import org.apache.kafka.common.security.auth.AuthenticationContext;
 import org.apache.kafka.common.security.auth.KafkaPrincipal;
 import org.apache.kafka.common.security.auth.KafkaPrincipalBuilder;
+import org.apache.kafka.common.security.auth.KafkaPrincipalSerde;
 import org.apache.kafka.common.security.auth.PlaintextAuthenticationContext;
 import org.apache.kafka.common.security.auth.SaslAuthenticationContext;
 import org.apache.kafka.common.security.auth.SslAuthenticationContext;
@@ -32,8 +38,11 @@
 import javax.net.ssl.SSLPeerUnverifiedException;
 import javax.net.ssl.SSLSession;
 import javax.security.sasl.SaslServer;
+import org.apache.kafka.common.security.ssl.SslPrincipalMapper;
+
 import java.io.Closeable;
 import java.io.IOException;
+import java.nio.ByteBuffer;
 import java.security.Principal;
 
 import static java.util.Objects.requireNonNull;
@@ -48,13 +57,14 @@
  * of {@link KafkaPrincipalBuilder}, there is no default no-arg constructor since this class
  * must adapt implementations of the older {@link org.apache.kafka.common.security.auth.PrincipalBuilder} interface.
  */
-public class DefaultKafkaPrincipalBuilder implements KafkaPrincipalBuilder, Closeable {
+public class DefaultKafkaPrincipalBuilder implements KafkaPrincipalBuilder, KafkaPrincipalSerde, Closeable {
     // Use FQN to avoid import deprecation warnings
     @SuppressWarnings("deprecation")
     private final org.apache.kafka.common.security.auth.PrincipalBuilder oldPrincipalBuilder;
     private final Authenticator authenticator;
     private final TransportLayer transportLayer;
     private final KerberosShortNamer kerberosShortNamer;
+    private final SslPrincipalMapper sslPrincipalMapper;
 
     /**
      * Construct a new instance which wraps an instance of the older {@link org.apache.kafka.common.security.auth.PrincipalBuilder}.
@@ -73,27 +83,31 @@ public static DefaultKafkaPrincipalBuilder fromOldPrincipalBuilder(Authenticator
                 requireNonNull(authenticator),
                 requireNonNull(transportLayer),
                 requireNonNull(oldPrincipalBuilder),
-                kerberosShortNamer);
+                kerberosShortNamer,
+                null);
     }
 
     @SuppressWarnings("deprecation")
     private DefaultKafkaPrincipalBuilder(Authenticator authenticator,
                                          TransportLayer transportLayer,
                                          org.apache.kafka.common.security.auth.PrincipalBuilder oldPrincipalBuilder,
-                                         KerberosShortNamer kerberosShortNamer) {
+                                         KerberosShortNamer kerberosShortNamer,
+                                         SslPrincipalMapper sslPrincipalMapper) {
         this.authenticator = authenticator;
         this.transportLayer = transportLayer;
         this.oldPrincipalBuilder = oldPrincipalBuilder;
         this.kerberosShortNamer = kerberosShortNamer;
+        this.sslPrincipalMapper =  sslPrincipalMapper;
     }
 
     /**
      * Construct a new instance.
      *
      * @param kerberosShortNamer Kerberos name rewrite rules or null if none have been configured
+     * @param sslPrincipalMapper SSL Principal mapper or null if none have been configured
      */
-    public DefaultKafkaPrincipalBuilder(KerberosShortNamer kerberosShortNamer) {
-        this(null, null, null, kerberosShortNamer);
+    public DefaultKafkaPrincipalBuilder(KerberosShortNamer kerberosShortNamer, SslPrincipalMapper sslPrincipalMapper) {
+        this(null, null, null, kerberosShortNamer, sslPrincipalMapper);
     }
 
     @Override
@@ -110,7 +124,7 @@ public KafkaPrincipal build(AuthenticationContext context) {
                 return convertToKafkaPrincipal(oldPrincipalBuilder.buildPrincipal(transportLayer, authenticator));
 
             try {
-                return convertToKafkaPrincipal(sslSession.getPeerPrincipal());
+                return applySslPrincipalMapper(sslSession.getPeerPrincipal());
             } catch (SSLPeerUnverifiedException se) {
                 return KafkaPrincipal.ANONYMOUS;
             }
@@ -136,14 +150,53 @@ private KafkaPrincipal applyKerberosShortNamer(String authorizationId) {
         }
     }
 
+    private KafkaPrincipal applySslPrincipalMapper(Principal principal) {
+        try {
+            if (!(principal instanceof X500Principal) || principal == KafkaPrincipal.ANONYMOUS) {
+                return new KafkaPrincipal(KafkaPrincipal.USER_TYPE, principal.getName());
+            } else {
+                return new KafkaPrincipal(KafkaPrincipal.USER_TYPE, sslPrincipalMapper.getName(principal.getName()));
+            }
+        } catch (IOException e) {
+            throw new KafkaException("Failed to map name for '" + principal.getName() +
+                    "' based on SSL principal mapping rules.", e);
+        }
+    }
+
     private KafkaPrincipal convertToKafkaPrincipal(Principal principal) {
         return new KafkaPrincipal(KafkaPrincipal.USER_TYPE, principal.getName());
     }
 
+    @Override
+    public byte[] serialize(KafkaPrincipal principal) {
+        DefaultPrincipalData data = new DefaultPrincipalData()
+                                        .setType(principal.getPrincipalType())
+                                        .setName(principal.getName())
+                                        .setTokenAuthenticated(principal.tokenAuthenticated());
+        Struct dataStruct = data.toStruct(DefaultPrincipalData.HIGHEST_SUPPORTED_VERSION);
+        ByteBuffer buffer = ByteBuffer.allocate(2 + dataStruct.sizeOf());
+        buffer.putShort(DefaultPrincipalData.HIGHEST_SUPPORTED_VERSION);
+        dataStruct.writeTo(buffer);
+        return buffer.array();
+    }
+
+    @Override
+    public KafkaPrincipal deserialize(byte[] bytes) {
+        ByteBuffer buffer = ByteBuffer.wrap(bytes);
+        short version = buffer.getShort();
+        if (version < 0 || version >= DefaultPrincipalData.SCHEMAS.length) {
+            throw new SerializationException("Invalid principal data version " + version);
+        }
+
+        DefaultPrincipalData data = new DefaultPrincipalData(
+            DefaultPrincipalData.SCHEMAS[version].read(buffer),
+            version);
+        return new KafkaPrincipal(data.type(), data.name(), data.tokenAuthenticated());
+    }
+
     @Override
     public void close() {
         if (oldPrincipalBuilder != null)
             oldPrincipalBuilder.close();
     }
-
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/security/authenticator/Login.java b/clients/src/main/java/org/apache/kafka/common/security/authenticator/Login.java
deleted file mode 100644
index b41d1b2572fcc..0000000000000
--- a/clients/src/main/java/org/apache/kafka/common/security/authenticator/Login.java
+++ /dev/null
@@ -1,57 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements. See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License. You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.kafka.common.security.authenticator;
-
-import org.apache.kafka.common.security.JaasContext;
-
-import java.util.Map;
-
-import javax.security.auth.Subject;
-import javax.security.auth.login.LoginContext;
-import javax.security.auth.login.LoginException;
-
-/**
- * Login interface for authentication.
- */
-public interface Login {
-
-    /**
-     * Configures this login instance.
-     */
-    void configure(Map configs, JaasContext jaasContext);
-
-    /**
-     * Performs login for each login module specified for the login context of this instance.
-     */
-    LoginContext login() throws LoginException;
-
-    /**
-     * Returns the authenticated subject of this login context.
-     */
-    Subject subject();
-
-    /**
-     * Returns the service name to be used for SASL.
-     */
-    String serviceName();
-
-    /**
-     * Closes this instance.
-     */
-    void close();
-}
-
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 a576e37f4632c..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
@@ -16,40 +16,49 @@
  */
 package org.apache.kafka.common.security.authenticator;
 
-import javax.security.auth.Subject;
-import javax.security.auth.login.LoginException;
-
-import java.io.IOException;
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.Map;
-
+import org.apache.kafka.common.config.ConfigException;
 import org.apache.kafka.common.config.SaslConfigs;
 import org.apache.kafka.common.config.types.Password;
+import org.apache.kafka.common.network.ListenerName;
 import org.apache.kafka.common.security.JaasContext;
-import org.apache.kafka.common.security.kerberos.KerberosLogin;
+import org.apache.kafka.common.security.auth.AuthenticateCallbackHandler;
+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;
 
+import javax.security.auth.Subject;
+import javax.security.auth.login.LoginException;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Objects;
+
 public class LoginManager {
 
     private static final Logger LOGGER = LoggerFactory.getLogger(LoginManager.class);
 
     // static configs (broker or client)
-    private static final Map STATIC_INSTANCES = new HashMap<>();
+    private static final Map, LoginManager> STATIC_INSTANCES = new HashMap<>();
 
-    // dynamic configs (client-only)
-    private static final Map DYNAMIC_INSTANCES = new HashMap<>();
+    // dynamic configs (broker or client)
+    private static final Map, LoginManager> DYNAMIC_INSTANCES = new HashMap<>();
 
     private final Login login;
-    private final Object cacheKey;
+    private final LoginMetadata loginMetadata;
+    private final AuthenticateCallbackHandler loginCallbackHandler;
     private int refCount;
 
-    private LoginManager(JaasContext jaasContext, boolean hasKerberos, Map configs,
-                         Password jaasConfigValue) throws IOException, LoginException {
-        this.cacheKey = jaasConfigValue != null ? jaasConfigValue : jaasContext.name();
-        login = hasKerberos ? new KerberosLogin() : new DefaultLogin();
-        login.configure(configs, jaasContext);
+    private LoginManager(JaasContext jaasContext, String saslMechanism, Map configs,
+                         LoginMetadata loginMetadata) throws LoginException {
+        this.loginMetadata = loginMetadata;
+        this.login = Utils.newInstance(loginMetadata.loginClass);
+        loginCallbackHandler = Utils.newInstance(loginMetadata.loginCallbackClass);
+        loginCallbackHandler.configure(configs, saslMechanism, jaasContext.configurationEntries());
+        login.configure(configs, jaasContext.name(), jaasContext.configuration(), loginCallbackHandler);
         login.login();
     }
 
@@ -57,32 +66,54 @@ private LoginManager(JaasContext jaasContext, boolean hasKerberos, Map configs) throws IOException, LoginException {
+    public static LoginManager acquireLoginManager(JaasContext jaasContext, String saslMechanism,
+                                                   Class defaultLoginClass,
+                                                   Map configs) throws LoginException {
+        Class loginClass = configuredClassOrDefault(configs, jaasContext,
+                saslMechanism, SaslConfigs.SASL_LOGIN_CLASS, defaultLoginClass);
+        Class defaultLoginCallbackHandlerClass = OAuthBearerLoginModule.OAUTHBEARER_MECHANISM
+                .equals(saslMechanism) ? OAuthBearerUnsecuredLoginCallbackHandler.class
+                        : AbstractLogin.DefaultLoginCallbackHandler.class;
+        Class loginCallbackClass = configuredClassOrDefault(configs, jaasContext,
+                saslMechanism, SaslConfigs.SASL_LOGIN_CALLBACK_HANDLER_CLASS, defaultLoginCallbackHandlerClass);
         synchronized (LoginManager.class) {
-            // SASL_JAAS_CONFIG is only supported by clients
             LoginManager loginManager;
-            Password jaasConfigValue = (Password) configs.get(SaslConfigs.SASL_JAAS_CONFIG);
-            if (jaasContext.type() == JaasContext.Type.CLIENT && jaasConfigValue != null) {
-                loginManager = DYNAMIC_INSTANCES.get(jaasConfigValue);
+            Password jaasConfigValue = jaasContext.dynamicJaasConfig();
+            if (jaasConfigValue != null) {
+                LoginMetadata loginMetadata = new LoginMetadata<>(jaasConfigValue, loginClass, loginCallbackClass);
+                loginManager = DYNAMIC_INSTANCES.get(loginMetadata);
                 if (loginManager == null) {
-                    loginManager = new LoginManager(jaasContext, hasKerberos, configs, jaasConfigValue);
-                    DYNAMIC_INSTANCES.put(jaasConfigValue, loginManager);
+                    loginManager = new LoginManager(jaasContext, saslMechanism, configs, loginMetadata);
+                    DYNAMIC_INSTANCES.put(loginMetadata, loginManager);
                 }
             } else {
-                loginManager = STATIC_INSTANCES.get(jaasContext.name());
+                LoginMetadata loginMetadata = new LoginMetadata<>(jaasContext.name(), loginClass, loginCallbackClass);
+                loginManager = STATIC_INSTANCES.get(loginMetadata);
                 if (loginManager == null) {
-                    loginManager = new LoginManager(jaasContext, hasKerberos, configs, jaasConfigValue);
-                    STATIC_INSTANCES.put(jaasContext.name(), loginManager);
+                    loginManager = new LoginManager(jaasContext, saslMechanism, configs, loginMetadata);
+                    STATIC_INSTANCES.put(loginMetadata, loginManager);
                 }
             }
+            SecurityUtils.addConfiguredSecurityProviders(configs);
             return loginManager.acquire();
         }
     }
@@ -95,6 +126,11 @@ public String serviceName() {
         return login.serviceName();
     }
 
+    // Only for testing
+    Object cacheKey() {
+        return loginMetadata.configInfo;
+    }
+
     private LoginManager acquire() {
         ++refCount;
         LOGGER.trace("{} acquired", this);
@@ -109,12 +145,13 @@ public void release() {
             if (refCount == 0)
                 throw new IllegalStateException("release() called on disposed " + this);
             else if (refCount == 1) {
-                if (cacheKey instanceof Password) {
-                    DYNAMIC_INSTANCES.remove(cacheKey);
+                if (loginMetadata.configInfo instanceof Password) {
+                    DYNAMIC_INSTANCES.remove(loginMetadata);
                 } else {
-                    STATIC_INSTANCES.remove(cacheKey);
+                    STATIC_INSTANCES.remove(loginMetadata);
                 }
                 login.close();
+                loginCallbackHandler.close();
             }
             --refCount;
             LOGGER.trace("{} released", this);
@@ -132,10 +169,57 @@ public String toString() {
     /* Should only be used in tests. */
     public static void closeAll() {
         synchronized (LoginManager.class) {
-            for (String key : new ArrayList<>(STATIC_INSTANCES.keySet()))
+            for (LoginMetadata key : new ArrayList<>(STATIC_INSTANCES.keySet()))
                 STATIC_INSTANCES.remove(key).login.close();
-            for (Password key : new ArrayList<>(DYNAMIC_INSTANCES.keySet()))
+            for (LoginMetadata key : new ArrayList<>(DYNAMIC_INSTANCES.keySet()))
                 DYNAMIC_INSTANCES.remove(key).login.close();
         }
     }
+
+    private static  Class configuredClassOrDefault(Map configs,
+                                                     JaasContext jaasContext,
+                                                     String saslMechanism,
+                                                     String configName,
+                                                     Class defaultClass) {
+        String prefix  = jaasContext.type() == JaasContext.Type.SERVER ? ListenerName.saslMechanismPrefix(saslMechanism) : "";
+        @SuppressWarnings("unchecked")
+        Class clazz = (Class) configs.get(prefix + configName);
+        if (clazz != null && jaasContext.configurationEntries().size() != 1) {
+            String errorMessage = configName + " cannot be specified with multiple login modules in the JAAS context. " +
+                    SaslConfigs.SASL_JAAS_CONFIG + " must be configured to override mechanism-specific configs.";
+            throw new ConfigException(errorMessage);
+        }
+        if (clazz == null)
+            clazz = defaultClass;
+        return clazz;
+    }
+
+    private static class LoginMetadata {
+        final T configInfo;
+        final Class loginClass;
+        final Class loginCallbackClass;
+
+        LoginMetadata(T configInfo, Class loginClass,
+                      Class loginCallbackClass) {
+            this.configInfo = configInfo;
+            this.loginClass = loginClass;
+            this.loginCallbackClass = loginCallbackClass;
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(configInfo, loginClass, loginCallbackClass);
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (o == null || getClass() != o.getClass()) return false;
+
+            LoginMetadata loginMetadata = (LoginMetadata) o;
+            return Objects.equals(configInfo, loginMetadata.configInfo) &&
+                   Objects.equals(loginClass, loginMetadata.loginClass) &&
+                   Objects.equals(loginCallbackClass, loginMetadata.loginCallbackClass);
+        }
+    }
 }
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 8b0116563d8be..916585017a791 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,10 +23,14 @@
 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;
-import org.apache.kafka.common.network.Mode;
+import org.apache.kafka.common.network.ByteBufferSend;
 import org.apache.kafka.common.network.NetworkReceive;
-import org.apache.kafka.common.network.NetworkSend;
+import org.apache.kafka.common.network.ReauthenticationContext;
 import org.apache.kafka.common.network.Send;
 import org.apache.kafka.common.network.TransportLayer;
 import org.apache.kafka.common.protocol.ApiKeys;
@@ -35,48 +39,102 @@
 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;
 import org.apache.kafka.common.requests.SaslHandshakeRequest;
 import org.apache.kafka.common.requests.SaslHandshakeResponse;
+import org.apache.kafka.common.security.auth.AuthenticateCallbackHandler;
 import org.apache.kafka.common.security.auth.KafkaPrincipal;
+import org.apache.kafka.common.security.auth.KafkaPrincipalSerde;
+import org.apache.kafka.common.security.kerberos.KerberosError;
+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 org.slf4j.LoggerFactory;
 
 import javax.security.auth.Subject;
 import javax.security.sasl.Sasl;
 import javax.security.sasl.SaslClient;
 import javax.security.sasl.SaslException;
 import java.io.IOException;
+import java.nio.BufferUnderflowException;
 import java.nio.ByteBuffer;
 import java.nio.channels.SelectionKey;
 import java.security.Principal;
 import java.security.PrivilegedActionException;
 import java.security.PrivilegedExceptionAction;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Iterator;
+import java.util.List;
 import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Random;
 import java.util.Set;
 
 public class SaslClientAuthenticator implements Authenticator {
-
+    /**
+     * The internal state transitions for initial authentication of a channel are
+     * declared in order, starting with {@link #SEND_APIVERSIONS_REQUEST} and ending
+     * in either {@link #COMPLETE} or {@link #FAILED}.
+     * 
        
+     * Re-authentication of a channel starts with the state
+     * {@link #REAUTH_PROCESS_ORIG_APIVERSIONS_RESPONSE} and then flows to
+     * {@link #REAUTH_SEND_HANDSHAKE_REQUEST} followed by
+     * {@link #REAUTH_RECEIVE_HANDSHAKE_OR_OTHER_RESPONSE} and then
+     * {@link #REAUTH_INITIAL}; after that the flow joins the authentication flow
+     * at the {@link #INTERMEDIATE} state and ends at either {@link #COMPLETE} or
+     * {@link #FAILED}.
+     */
     public enum SaslState {
-        SEND_APIVERSIONS_REQUEST,     // Initial state: client sends ApiVersionsRequest in this state
-        RECEIVE_APIVERSIONS_RESPONSE, // Awaiting ApiVersionsResponse from server
-        SEND_HANDSHAKE_REQUEST,       // Received ApiVersionsResponse, send SaslHandshake request
-        RECEIVE_HANDSHAKE_RESPONSE,   // Awaiting SaslHandshake request from server
-        INITIAL,                      // Initial state starting SASL token exchange for configured mechanism, send first token
-        INTERMEDIATE,                 // Intermediate state during SASL token exchange, process challenges and send responses
-        CLIENT_COMPLETE,              // Sent response to last challenge. If using SaslAuthenticate, wait for authentication status from server, else COMPLETE
-        COMPLETE,                     // Authentication sequence complete. If using SaslAuthenticate, this state implies successful authentication.
-        FAILED                        // Failed authentication due to an error at some stage
+        SEND_APIVERSIONS_REQUEST,                   // Initial state for authentication: client sends ApiVersionsRequest in this state when authenticating
+        RECEIVE_APIVERSIONS_RESPONSE,               // Awaiting ApiVersionsResponse from server
+        SEND_HANDSHAKE_REQUEST,                     // Received ApiVersionsResponse, send SaslHandshake request
+        RECEIVE_HANDSHAKE_RESPONSE,                 // Awaiting SaslHandshake response from server when authenticating
+        INITIAL,                                    // Initial authentication state starting SASL token exchange for configured mechanism, send first token
+        INTERMEDIATE,                               // Intermediate state during SASL token exchange, process challenges and send responses
+        CLIENT_COMPLETE,                            // Sent response to last challenge. If using SaslAuthenticate, wait for authentication status from server, else COMPLETE
+        COMPLETE,                                   // Authentication sequence complete. If using SaslAuthenticate, this state implies successful authentication.
+        FAILED,                                     // Failed authentication due to an error at some stage
+        REAUTH_PROCESS_ORIG_APIVERSIONS_RESPONSE,   // Initial state for re-authentication: process ApiVersionsResponse from original authentication
+        REAUTH_SEND_HANDSHAKE_REQUEST,              // Processed original ApiVersionsResponse, send SaslHandshake request as part of re-authentication
+        REAUTH_RECEIVE_HANDSHAKE_OR_OTHER_RESPONSE, // Awaiting SaslHandshake response from server when re-authenticating, and may receive other, in-flight responses sent prior to start of re-authentication as well
+        REAUTH_INITIAL,                             // Initial re-authentication state starting SASL token exchange for configured mechanism, send first token
     }
 
-    private static final Logger LOG = LoggerFactory.getLogger(SaslClientAuthenticator.class);
     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;
@@ -87,7 +145,10 @@ public enum SaslState {
     private final SaslClient saslClient;
     private final Map configs;
     private final String clientPrincipalName;
-    private final AuthCallbackHandler callbackHandler;
+    private final AuthenticateCallbackHandler callbackHandler;
+    private final Time time;
+    private final Logger log;
+    private final ReauthInfo reauthInfo;
 
     // buffers used in `authenticate`
     private NetworkReceive netInBuffer;
@@ -103,24 +164,33 @@ public enum SaslState {
     private RequestHeader currentRequestHeader;
     // Version of SaslAuthenticate request/responses
     private short saslAuthenticateVersion;
+    // Version of SaslHandshake request/responses
+    private short saslHandshakeVersion;
 
     public SaslClientAuthenticator(Map configs,
+                                   AuthenticateCallbackHandler callbackHandler,
                                    String node,
                                    Subject subject,
                                    String servicePrincipal,
                                    String host,
                                    String mechanism,
                                    boolean handshakeRequestEnable,
-                                   TransportLayer transportLayer) throws IOException {
+                                   TransportLayer transportLayer,
+                                   Time time,
+                                   LogContext logContext) {
         this.node = node;
         this.subject = subject;
+        this.callbackHandler = callbackHandler;
         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;
+        this.time = time;
+        this.log = logContext.logger(getClass());
+        this.reauthInfo = new ReauthInfo();
 
         try {
             setSaslState(handshakeRequestEnable ? SaslState.SEND_APIVERSIONS_REQUEST : SaslState.INITIAL);
@@ -133,24 +203,24 @@ public SaslClientAuthenticator(Map configs,
             else
                 this.clientPrincipalName = null;
 
-            callbackHandler = new SaslClientCallbackHandler();
-            callbackHandler.configure(configs, Mode.CLIENT, subject, mechanism);
-
             saslClient = createSaslClient();
         } catch (Exception e) {
             throw new SaslAuthenticationException("Failed to configure SaslClientAuthenticator", e);
         }
     }
 
-    private SaslClient createSaslClient() {
+    // visible for testing
+    SaslClient createSaslClient() {
         try {
-            return Subject.doAs(subject, new PrivilegedExceptionAction() {
-                public SaslClient run() throws SaslException {
-                    String[] mechs = {mechanism};
-                    LOG.debug("Creating SaslClient: client={};service={};serviceHostname={};mechs={}",
-                        clientPrincipalName, servicePrincipal, host, Arrays.toString(mechs));
-                    return Sasl.createSaslClient(mechs, clientPrincipalName, servicePrincipal, host, configs, callbackHandler);
+            return Subject.doAs(subject, (PrivilegedExceptionAction) () -> {
+                String[] mechs = {mechanism};
+                log.debug("Creating SaslClient: client={};service={};serviceHostname={};mechs={}",
+                    clientPrincipalName, servicePrincipal, host, Arrays.toString(mechs));
+                SaslClient retvalSaslClient = Sasl.createSaslClient(mechs, clientPrincipalName, servicePrincipal, host, configs, callbackHandler);
+                if (retvalSaslClient == null) {
+                    throw new SaslAuthenticationException("Failed to create SaslClient with mechanism " + mechanism);
                 }
+                return retvalSaslClient;
             });
         } catch (PrivilegedActionException e) {
             throw new SaslAuthenticationException("Failed to create SaslClient with mechanism " + mechanism, e.getCause());
@@ -164,16 +234,16 @@ public SaslClient run() throws SaslException {
      * The messages are sent and received as size delimited bytes that consists of a 4 byte network-ordered size N
      * followed by N bytes representing the opaque payload.
      */
+    @SuppressWarnings("fallthrough")
     public void authenticate() throws IOException {
-        short saslHandshakeVersion = 0;
         if (netOutBuffer != null && !flushNetOutBufferAndUpdateInterestOps())
             return;
 
         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);
-                send(apiVersionsRequest.toSend(node, nextRequestHeader(ApiKeys.API_VERSIONS, apiVersionsRequest.version())));
+                ApiVersionsRequest apiVersionsRequest = new ApiVersionsRequest.Builder().build((short) 0);
+                send(apiVersionsRequest.toSend(nextRequestHeader(ApiKeys.API_VERSIONS, apiVersionsRequest.version())));
                 setSaslState(SaslState.RECEIVE_APIVERSIONS_RESPONSE);
                 break;
             case RECEIVE_APIVERSIONS_RESPONSE:
@@ -181,16 +251,13 @@ public void authenticate() throws IOException {
                 if (apiVersionsResponse == null)
                     break;
                 else {
-                    saslHandshakeVersion = apiVersionsResponse.apiVersion(ApiKeys.SASL_HANDSHAKE.id).maxVersion;
-                    ApiVersion authenticateVersion = apiVersionsResponse.apiVersion(ApiKeys.SASL_AUTHENTICATE.id);
-                    if (authenticateVersion != null)
-                        saslAuthenticateVersion((short) Math.min(authenticateVersion.maxVersion, ApiKeys.SASL_AUTHENTICATE.latestVersion()));
+                    setSaslAuthenticateAndHandshakeVersions(apiVersionsResponse);
+                    reauthInfo.apiVersionsResponseReceivedFromBroker = apiVersionsResponse;
                     setSaslState(SaslState.SEND_HANDSHAKE_REQUEST);
                     // Fall through to send handshake request with the latest supported version
                 }
             case SEND_HANDSHAKE_REQUEST:
-                SaslHandshakeRequest handshakeRequest = createSaslHandshakeRequest(saslHandshakeVersion);
-                send(handshakeRequest.toSend(node, nextRequestHeader(ApiKeys.SASL_HANDSHAKE, handshakeRequest.version())));
+                sendHandshakeRequest(saslHandshakeVersion);
                 setSaslState(SaslState.RECEIVE_HANDSHAKE_RESPONSE);
                 break;
             case RECEIVE_HANDSHAKE_RESPONSE:
@@ -203,7 +270,32 @@ public void authenticate() throws IOException {
                     // Fall through and start SASL authentication using the configured client mechanism
                 }
             case INITIAL:
-                sendSaslClientToken(new byte[0], true);
+                sendInitialToken();
+                setSaslState(SaslState.INTERMEDIATE);
+                break;
+            case REAUTH_PROCESS_ORIG_APIVERSIONS_RESPONSE:
+                setSaslAuthenticateAndHandshakeVersions(reauthInfo.apiVersionsResponseFromOriginalAuthentication);
+                setSaslState(SaslState.REAUTH_SEND_HANDSHAKE_REQUEST); // Will set immediately
+                // Fall through to send handshake request with the latest supported version
+            case REAUTH_SEND_HANDSHAKE_REQUEST:
+                sendHandshakeRequest(saslHandshakeVersion);
+                setSaslState(SaslState.REAUTH_RECEIVE_HANDSHAKE_OR_OTHER_RESPONSE);
+                break;
+            case REAUTH_RECEIVE_HANDSHAKE_OR_OTHER_RESPONSE:
+                handshakeResponse = (SaslHandshakeResponse) receiveKafkaResponse();
+                if (handshakeResponse == null)
+                    break;
+                handleSaslHandshakeResponse(handshakeResponse);
+                setSaslState(SaslState.REAUTH_INITIAL); // Will set immediately
+                /*
+                 * Fall through and start SASL authentication using the configured client
+                 * mechanism. Note that we have to either fall through or add a loop to enter
+                 * the switch statement again. We will fall through to avoid adding the loop and
+                 * therefore minimize the changes to authentication-related code due to the
+                 * changes related to re-authentication.
+                 */
+            case REAUTH_INITIAL:
+                sendInitialToken();
                 setSaslState(SaslState.INTERMEDIATE);
                 break;
             case INTERMEDIATE:
@@ -231,20 +323,83 @@ public void authenticate() throws IOException {
         }
     }
 
+    private void sendHandshakeRequest(short version) throws IOException {
+        SaslHandshakeRequest handshakeRequest = createSaslHandshakeRequest(version);
+        send(handshakeRequest.toSend(nextRequestHeader(ApiKeys.SASL_HANDSHAKE, handshakeRequest.version())));
+    }
+
+    private void sendInitialToken() throws IOException {
+        sendSaslClientToken(new byte[0], true);
+    }
+
+    @Override
+    public void reauthenticate(ReauthenticationContext reauthenticationContext) throws IOException {
+        SaslClientAuthenticator previousSaslClientAuthenticator = (SaslClientAuthenticator) Objects
+                .requireNonNull(reauthenticationContext).previousAuthenticator();
+        ApiVersionsResponse apiVersionsResponseFromOriginalAuthentication = previousSaslClientAuthenticator.reauthInfo
+                .apiVersionsResponse();
+        previousSaslClientAuthenticator.close();
+        reauthInfo.reauthenticating(apiVersionsResponseFromOriginalAuthentication,
+                reauthenticationContext.reauthenticationBeginNanos());
+        NetworkReceive netInBufferFromChannel = reauthenticationContext.networkReceive();
+        netInBuffer = netInBufferFromChannel;
+        setSaslState(SaslState.REAUTH_PROCESS_ORIG_APIVERSIONS_RESPONSE); // Will set immediately
+        authenticate();
+    }
+
+    @Override
+    public Optional pollResponseReceivedDuringReauthentication() {
+        return reauthInfo.pollResponseReceivedDuringReauthentication();
+    }
+
+    @Override
+    public Long clientSessionReauthenticationTimeNanos() {
+        return reauthInfo.clientSessionReauthenticationTimeNanos;
+    }
+
+    @Override
+    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;
     }
 
     // Visible to override for testing
     protected SaslHandshakeRequest createSaslHandshakeRequest(short version) {
-        return new SaslHandshakeRequest.Builder(mechanism).build(version);
+        return new SaslHandshakeRequest.Builder(
+                new SaslHandshakeRequestData().setMechanism(mechanism)).build(version);
     }
 
     // Visible to override for testing
-    protected void saslAuthenticateVersion(short version) {
-        this.saslAuthenticateVersion = version;
+    protected void setSaslAuthenticateAndHandshakeVersions(ApiVersionsResponse apiVersionsResponse) {
+        ApiVersionsResponseKey authenticateVersion = apiVersionsResponse.apiVersion(ApiKeys.SASL_AUTHENTICATE.id);
+        if (authenticateVersion != null) {
+            this.saslAuthenticateVersion = (short) Math.min(authenticateVersion.maxVersion(),
+                    ApiKeys.SASL_AUTHENTICATE.latestVersion());
+        }
+        ApiVersionsResponseKey handshakeVersion = apiVersionsResponse.apiVersion(ApiKeys.SASL_HANDSHAKE.id);
+        if (handshakeVersion != null) {
+            this.saslHandshakeVersion = (short) Math.min(handshakeVersion.maxVersion(),
+                    ApiKeys.SASL_HANDSHAKE.latestVersion());
+        }
     }
 
     private void setSaslState(SaslState saslState) {
@@ -253,9 +408,18 @@ private void setSaslState(SaslState saslState) {
         else {
             this.pendingSaslState = null;
             this.saslState = saslState;
-            LOG.debug("Set SASL client state to {}", saslState);
-            if (saslState == SaslState.COMPLETE)
-                transportLayer.removeInterestOps(SelectionKey.OP_WRITE);
+            log.debug("Set SASL client state to {}", saslState);
+            if (saslState == SaslState.COMPLETE) {
+                reauthInfo.setAuthenticationEndAndSessionReauthenticationTimes(time.nanoseconds());
+                if (!reauthInfo.reauthenticating())
+                    transportLayer.removeInterestOps(SelectionKey.OP_WRITE);
+                else
+                    /*
+                     * Re-authentication is triggered by a write, so we have to make sure that
+                     * pending write is actually sent.
+                     */
+                    transportLayer.addInterestOps(SelectionKey.OP_WRITE);
+            }
         }
     }
 
@@ -269,11 +433,16 @@ private boolean sendSaslClientToken(byte[] serverToken, boolean isInitial) throw
             byte[] saslToken = createSaslToken(serverToken, isInitial);
             if (saslToken != null) {
                 ByteBuffer tokenBuf = ByteBuffer.wrap(saslToken);
-                if (saslAuthenticateVersion != DISABLE_KAFKA_SASL_AUTHENTICATE_HEADER) {
-                    SaslAuthenticateRequest request = new SaslAuthenticateRequest.Builder(tokenBuf).build(saslAuthenticateVersion);
-                    tokenBuf = request.serialize(nextRequestHeader(ApiKeys.SASL_AUTHENTICATE, saslAuthenticateVersion));
+                Send send;
+                if (saslAuthenticateVersion == DISABLE_KAFKA_SASL_AUTHENTICATE_HEADER) {
+                    send = ByteBufferSend.sizePrefixed(tokenBuf);
+                } else {
+                    SaslAuthenticateRequestData data = new SaslAuthenticateRequestData()
+                            .setAuthBytes(tokenBuf.array());
+                    SaslAuthenticateRequest request = new SaslAuthenticateRequest.Builder(data).build(saslAuthenticateVersion);
+                    send = request.toSend(nextRequestHeader(ApiKeys.SASL_AUTHENTICATE, saslAuthenticateVersion));
                 }
-                send(new NetworkSend(node, tokenBuf));
+                send(send);
                 return true;
             }
         }
@@ -318,6 +487,11 @@ public KafkaPrincipal principal() {
         return new KafkaPrincipal(KafkaPrincipal.USER_TYPE, clientPrincipalName);
     }
 
+    @Override
+    public Optional principalSerde() {
+        return Optional.empty();
+    }
+
     public boolean complete() {
         return saslState == SaslState.COMPLETE;
     }
@@ -325,8 +499,6 @@ public boolean complete() {
     public void close() throws IOException {
         if (saslClient != null)
             saslClient.dispose();
-        if (callbackHandler != null)
-            callbackHandler.close();
     }
 
     private byte[] receiveToken() throws IOException {
@@ -341,7 +513,10 @@ private byte[] receiveToken() throws IOException {
                     String errMsg = response.errorMessage();
                     throw errMsg == null ? error.exception() : error.exception(errMsg);
                 }
-                return Utils.readBytes(response.saslAuthBytes());
+                long sessionLifetimeMs = response.sessionLifetimeMs();
+                if (sessionLifetimeMs > 0L)
+                    reauthInfo.positiveSessionLifetimeMs = sessionLifetimeMs;
+                return Utils.copyArray(response.saslAuthBytes());
             } else
                 return null;
         }
@@ -356,27 +531,29 @@ private byte[] createSaslToken(final byte[] saslToken, boolean isInitial) throws
             if (isInitial && !saslClient.hasInitialResponse())
                 return saslToken;
             else
-                return Subject.doAs(subject, new PrivilegedExceptionAction() {
-                    public byte[] run() throws SaslException {
-                        return saslClient.evaluateChallenge(saslToken);
-                    }
-                });
+                return Subject.doAs(subject, (PrivilegedExceptionAction) () -> saslClient.evaluateChallenge(saslToken));
         } catch (PrivilegedActionException e) {
             String error = "An error: (" + e + ") occurred when evaluating SASL token received from the Kafka Broker.";
+            KerberosError kerberosError = KerberosError.fromException(e);
             // Try to provide hints to use about what went wrong so they can fix their configuration.
-            // TODO: introspect about e: look for GSS information.
-            final String unknownServerErrorText =
-                "(Mechanism level: Server not found in Kerberos database (7) - UNKNOWN_SERVER)";
-            if (e.toString().contains(unknownServerErrorText)) {
+            if (kerberosError == KerberosError.SERVER_NOT_FOUND) {
                 error += " This may be caused by Java's being unable to resolve the Kafka Broker's" +
                     " hostname correctly. You may want to try to adding" +
                     " '-Dsun.net.spi.nameservice.provider.1=dns,sun' to your client's JVMFLAGS environment." +
                     " Users must configure FQDN of kafka brokers when authenticating using SASL and" +
                     " `socketChannel.socket().getInetAddress().getHostName()` must match the hostname in `principal/hostname@realm`";
             }
-            error += " Kafka Client will go to AUTHENTICATION_FAILED state.";
             //Unwrap the SaslException inside `PrivilegedActionException`
-            throw new SaslAuthenticationException(error, e.getCause());
+            Throwable cause = e.getCause();
+            // Treat transient Kerberos errors as non-fatal SaslExceptions that are processed as I/O exceptions
+            // and all other failures as fatal SaslAuthenticationException.
+            if ((kerberosError != null && kerberosError.retriable()) || (kerberosError == null && KerberosError.isRetriableClientGssException(e))) {
+                error += " Kafka Client will retry.";
+                throw new SaslException(error, cause);
+            } else {
+                error += " Kafka Client will go to AUTHENTICATION_FAILED state.";
+                throw new SaslAuthenticationException(error, cause);
+            }
         }
     }
 
@@ -388,6 +565,9 @@ private boolean flushNetOutBuffer() throws IOException {
     }
 
     private AbstractResponse receiveKafkaResponse() throws IOException {
+        if (netInBuffer == null)
+            netInBuffer = new NetworkReceive(node);
+        NetworkReceive receive = netInBuffer;
         try {
             byte[] responseBytes = receiveResponseOrToken();
             if (responseBytes == null)
@@ -397,8 +577,21 @@ private AbstractResponse receiveKafkaResponse() throws IOException {
                 currentRequestHeader = null;
                 return response;
             }
-        } catch (SchemaException | IllegalArgumentException e) {
-            LOG.debug("Invalid SASL mechanism response, server may be expecting only GSSAPI tokens");
+        } catch (BufferUnderflowException | SchemaException | IllegalArgumentException e) {
+            /*
+             * Account for the fact that during re-authentication there may be responses
+             * arriving for requests that were sent in the past.
+             */
+            if (reauthInfo.reauthenticating()) {
+                /*
+                 * It didn't match the current request header, so it must be unrelated to
+                 * re-authentication. Save it so it can be processed later.
+                 */
+                receive.payload().rewind();
+                reauthInfo.pendingAuthenticatedReceives.add(receive);
+                return null;
+            }
+            log.debug("Invalid SASL mechanism response, server may be expecting only GSSAPI tokens");
             setSaslState(SaslState.FAILED);
             throw new IllegalSaslStateException("Invalid SASL mechanism response, server may be expecting a different protocol", e);
         }
@@ -429,7 +622,7 @@ private void handleSaslHandshakeResponse(SaslHandshakeResponse response) {
      *     During Kerberos re-login, principal is reset on Subject. An exception is
      *     thrown so that the connection is retried after any configured backoff.
      */
-    static final String firstPrincipal(Subject subject) {
+    public static String firstPrincipal(Subject subject) {
         Set principals = subject.getPrincipals();
         synchronized (principals) {
             Iterator iterator = principals.iterator();
@@ -439,4 +632,80 @@ static final String firstPrincipal(Subject subject) {
                 throw new KafkaException("Principal could not be determined from Subject, this may be a transient failure due to Kerberos re-login");
         }
     }
+
+    /**
+     * Information related to re-authentication
+     */
+    private class ReauthInfo {
+        public ApiVersionsResponse apiVersionsResponseFromOriginalAuthentication;
+        public long reauthenticationBeginNanos;
+        public List pendingAuthenticatedReceives = new ArrayList<>();
+        public ApiVersionsResponse apiVersionsResponseReceivedFromBroker;
+        public Long positiveSessionLifetimeMs;
+        public long authenticationEndNanos;
+        public Long clientSessionReauthenticationTimeNanos;
+
+        public void reauthenticating(ApiVersionsResponse apiVersionsResponseFromOriginalAuthentication,
+                                     long reauthenticationBeginNanos) {
+            this.apiVersionsResponseFromOriginalAuthentication = Objects
+                    .requireNonNull(apiVersionsResponseFromOriginalAuthentication);
+            this.reauthenticationBeginNanos = reauthenticationBeginNanos;
+        }
+
+        public boolean reauthenticating() {
+            return apiVersionsResponseFromOriginalAuthentication != null;
+        }
+
+        public ApiVersionsResponse apiVersionsResponse() {
+            return reauthenticating() ? apiVersionsResponseFromOriginalAuthentication
+                    : apiVersionsResponseReceivedFromBroker;
+        }
+
+        /**
+         * Return the (always non-null but possibly empty) NetworkReceive response that
+         * arrived during re-authentication that is unrelated to re-authentication, if
+         * any. This corresponds to a request sent prior to the beginning of
+         * re-authentication; the request was made when the channel was successfully
+         * authenticated, and the response arrived during the re-authentication
+         * process.
+         * 
+         * @return the (always non-null but possibly empty) NetworkReceive response
+         *         that arrived during re-authentication that is unrelated to
+         *         re-authentication, if any
+         */
+        public Optional pollResponseReceivedDuringReauthentication() {
+            if (pendingAuthenticatedReceives.isEmpty())
+                return Optional.empty();
+            return Optional.of(pendingAuthenticatedReceives.remove(0));
+        }
+
+        public void setAuthenticationEndAndSessionReauthenticationTimes(long nowNanos) {
+            authenticationEndNanos = nowNanos;
+            long sessionLifetimeMsToUse = 0;
+            if (positiveSessionLifetimeMs != null) {
+                // pick a random percentage between 85% and 95% for session re-authentication
+                double pctWindowFactorToTakeNetworkLatencyAndClockDriftIntoAccount = 0.85;
+                double pctWindowJitterToAvoidReauthenticationStormAcrossManyChannelsSimultaneously = 0.10;
+                double pctToUse = pctWindowFactorToTakeNetworkLatencyAndClockDriftIntoAccount + RNG.nextDouble()
+                        * pctWindowJitterToAvoidReauthenticationStormAcrossManyChannelsSimultaneously;
+                sessionLifetimeMsToUse = (long) (positiveSessionLifetimeMs * pctToUse);
+                clientSessionReauthenticationTimeNanos = authenticationEndNanos + 1000 * 1000 * sessionLifetimeMsToUse;
+                log.debug(
+                        "Finished {} with session expiration in {} ms and session re-authentication on or after {} ms",
+                        authenticationOrReauthenticationText(), positiveSessionLifetimeMs, sessionLifetimeMsToUse);
+            } else
+                log.debug("Finished {} with no session expiration and no session re-authentication",
+                        authenticationOrReauthenticationText());
+        }
+
+        public Long reauthenticationLatencyMs() {
+            return reauthenticating()
+                    ? Math.round((authenticationEndNanos - reauthenticationBeginNanos) / 1000.0 / 1000.0)
+                    : null;
+        }
+
+        private String authenticationOrReauthenticationText() {
+            return reauthenticating() ? "re-authentication" : "authentication";
+        }
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/security/authenticator/SaslClientCallbackHandler.java b/clients/src/main/java/org/apache/kafka/common/security/authenticator/SaslClientCallbackHandler.java
index 4756387bba8a8..e141bb652af3e 100644
--- a/clients/src/main/java/org/apache/kafka/common/security/authenticator/SaslClientCallbackHandler.java
+++ b/clients/src/main/java/org/apache/kafka/common/security/authenticator/SaslClientCallbackHandler.java
@@ -16,6 +16,8 @@
  */
 package org.apache.kafka.common.security.authenticator;
 
+import java.security.AccessController;
+import java.util.List;
 import java.util.Map;
 
 import javax.security.auth.Subject;
@@ -23,51 +25,51 @@
 import javax.security.auth.callback.NameCallback;
 import javax.security.auth.callback.PasswordCallback;
 import javax.security.auth.callback.UnsupportedCallbackException;
+import javax.security.auth.login.AppConfigurationEntry;
 import javax.security.sasl.AuthorizeCallback;
 import javax.security.sasl.RealmCallback;
 
 import org.apache.kafka.common.config.SaslConfigs;
-import org.apache.kafka.common.network.Mode;
+import org.apache.kafka.common.security.auth.SaslExtensionsCallback;
+import org.apache.kafka.common.security.auth.AuthenticateCallbackHandler;
+import org.apache.kafka.common.security.auth.SaslExtensions;
+import org.apache.kafka.common.security.scram.ScramExtensionsCallback;
+import org.apache.kafka.common.security.scram.internals.ScramMechanism;
 
 /**
- * Callback handler for Sasl clients. The callbacks required for the SASL mechanism
+ * Default callback handler for Sasl clients. The callbacks required for the SASL mechanism
  * configured for the client should be supported by this callback handler. See
  * Java SASL API
  * for the list of SASL callback handlers required for each SASL mechanism.
+ *
+ * For adding custom SASL extensions, a {@link SaslExtensions} may be added to the subject's public credentials
  */
-public class SaslClientCallbackHandler implements AuthCallbackHandler {
+public class SaslClientCallbackHandler implements AuthenticateCallbackHandler {
 
-    private boolean isKerberos;
-    private Subject subject;
+    private String mechanism;
 
     @Override
-    public void configure(Map configs, Mode mode, Subject subject, String mechanism) {
-        this.isKerberos = mechanism.equals(SaslConfigs.GSSAPI_MECHANISM);
-        this.subject = subject;
+    public void configure(Map configs, String saslMechanism, List jaasConfigEntries) {
+        this.mechanism  = saslMechanism;
     }
 
     @Override
     public void handle(Callback[] callbacks) throws UnsupportedCallbackException {
+        Subject subject = Subject.getSubject(AccessController.getContext());
         for (Callback callback : callbacks) {
             if (callback instanceof NameCallback) {
                 NameCallback nc = (NameCallback) callback;
-                if (!isKerberos && subject != null && !subject.getPublicCredentials(String.class).isEmpty()) {
+                if (subject != null && !subject.getPublicCredentials(String.class).isEmpty()) {
                     nc.setName(subject.getPublicCredentials(String.class).iterator().next());
                 } else
                     nc.setName(nc.getDefaultName());
             } else if (callback instanceof PasswordCallback) {
-                if (!isKerberos && subject != null && !subject.getPrivateCredentials(String.class).isEmpty()) {
+                if (subject != null && !subject.getPrivateCredentials(String.class).isEmpty()) {
                     char[] password = subject.getPrivateCredentials(String.class).iterator().next().toCharArray();
                     ((PasswordCallback) callback).setPassword(password);
                 } else {
                     String errorMessage = "Could not login: the client is being asked for a password, but the Kafka" +
                              " client code does not currently support obtaining a password from the user.";
-                    if (isKerberos) {
-                        errorMessage += " Make sure -Djava.security.auth.login.config property passed to JVM and" +
-                             " the client is configured to use a ticket cache (using" +
-                             " the JAAS configuration setting 'useTicketCache=true)'. Make sure you are using" +
-                             " FQDN of the Kafka broker you are trying to connect to.";
-                    }
                     throw new UnsupportedCallbackException(callback, errorMessage);
                 }
             } else if (callback instanceof RealmCallback) {
@@ -80,7 +82,19 @@ public void handle(Callback[] callbacks) throws UnsupportedCallbackException {
                 ac.setAuthorized(authId.equals(authzId));
                 if (ac.isAuthorized())
                     ac.setAuthorizedID(authzId);
-            } else {
+            } else if (callback instanceof ScramExtensionsCallback) {
+                if (ScramMechanism.isScram(mechanism) && subject != null && !subject.getPublicCredentials(Map.class).isEmpty()) {
+                    @SuppressWarnings("unchecked")
+                    Map extensions = (Map) subject.getPublicCredentials(Map.class).iterator().next();
+                    ((ScramExtensionsCallback) callback).extensions(extensions);
+                }
+            } else if (callback instanceof SaslExtensionsCallback) {
+                if (!SaslConfigs.GSSAPI_MECHANISM.equals(mechanism) &&
+                        subject != null && !subject.getPublicCredentials(SaslExtensions.class).isEmpty()) {
+                    SaslExtensions extensions = subject.getPublicCredentials(SaslExtensions.class).iterator().next();
+                    ((SaslExtensionsCallback) callback).extensions(extensions);
+                }
+            }  else {
                 throw new UnsupportedCallbackException(callback, "Unrecognized SASL ClientCallback");
             }
         }
diff --git a/clients/src/main/java/org/apache/kafka/common/security/authenticator/SaslInternalConfigs.java b/clients/src/main/java/org/apache/kafka/common/security/authenticator/SaslInternalConfigs.java
new file mode 100644
index 0000000000000..c1793ebc3192b
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/authenticator/SaslInternalConfigs.java
@@ -0,0 +1,41 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.common.security.authenticator;
+
+import org.apache.kafka.common.config.internals.BrokerSecurityConfigs;
+
+public class SaslInternalConfigs {
+    /**
+     * The server (broker) specifies a positive session length in milliseconds to a
+     * SASL client when {@link BrokerSecurityConfigs#CONNECTIONS_MAX_REAUTH_MS} is
+     * positive as per KIP
+     * 368: Allow SASL Connections to Periodically Re-Authenticate. The session
+     * length is the minimum of the configured value and any session length implied
+     * by the credential presented during authentication. The lifetime defined by
+     * the credential, in terms of milliseconds since the epoch, is available via a
+     * negotiated property on the SASL Server instance, and that value can be
+     * converted to a session length by subtracting the time at which authentication
+     * occurred. This variable defines the negotiated property key that is used to
+     * communicate the credential lifetime in milliseconds since the epoch.
+     */
+    public static final String CREDENTIAL_LIFETIME_MS_SASL_NEGOTIATED_PROPERTY_KEY = "CREDENTIAL.LIFETIME.MS";
+
+    private SaslInternalConfigs() {
+        // empty
+    }
+}
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 1a2aaf497b9ca..744b941af4d27 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
@@ -25,17 +25,20 @@
 import org.apache.kafka.common.errors.SaslAuthenticationException;
 import org.apache.kafka.common.errors.UnsupportedSaslMechanismException;
 import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.message.SaslAuthenticateResponseData;
+import org.apache.kafka.common.message.SaslHandshakeResponseData;
 import org.apache.kafka.common.network.Authenticator;
+import org.apache.kafka.common.network.ByteBufferSend;
 import org.apache.kafka.common.network.ChannelBuilders;
+import org.apache.kafka.common.network.ChannelMetadataRegistry;
+import org.apache.kafka.common.network.ClientInformation;
 import org.apache.kafka.common.network.ListenerName;
-import org.apache.kafka.common.network.Mode;
 import org.apache.kafka.common.network.NetworkReceive;
-import org.apache.kafka.common.network.NetworkSend;
+import org.apache.kafka.common.network.ReauthenticationContext;
 import org.apache.kafka.common.network.Send;
 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.security.auth.SecurityProtocol;
 import org.apache.kafka.common.requests.AbstractResponse;
 import org.apache.kafka.common.requests.ApiVersionsRequest;
 import org.apache.kafka.common.requests.ApiVersionsResponse;
@@ -46,22 +49,19 @@
 import org.apache.kafka.common.requests.SaslAuthenticateResponse;
 import org.apache.kafka.common.requests.SaslHandshakeRequest;
 import org.apache.kafka.common.requests.SaslHandshakeResponse;
-import org.apache.kafka.common.security.JaasContext;
+import org.apache.kafka.common.security.auth.AuthenticateCallbackHandler;
 import org.apache.kafka.common.security.auth.KafkaPrincipal;
 import org.apache.kafka.common.security.auth.KafkaPrincipalBuilder;
+import org.apache.kafka.common.security.auth.KafkaPrincipalSerde;
 import org.apache.kafka.common.security.auth.SaslAuthenticationContext;
+import org.apache.kafka.common.security.auth.SecurityProtocol;
+import org.apache.kafka.common.security.kerberos.KerberosError;
 import org.apache.kafka.common.security.kerberos.KerberosName;
 import org.apache.kafka.common.security.kerberos.KerberosShortNamer;
-import org.apache.kafka.common.security.scram.ScramCredential;
-import org.apache.kafka.common.security.scram.ScramMechanism;
-import org.apache.kafka.common.security.scram.ScramServerCallbackHandler;
+import org.apache.kafka.common.security.scram.ScramLoginModule;
+import org.apache.kafka.common.security.scram.internals.ScramMechanism;
+import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.common.utils.Utils;
-import org.ietf.jgss.GSSContext;
-import org.ietf.jgss.GSSCredential;
-import org.ietf.jgss.GSSException;
-import org.ietf.jgss.GSSManager;
-import org.ietf.jgss.GSSName;
-import org.ietf.jgss.Oid;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -76,37 +76,55 @@
 import java.nio.channels.SelectionKey;
 import java.security.PrivilegedActionException;
 import java.security.PrivilegedExceptionAction;
+import java.util.ArrayList;
+import java.util.Date;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
+import java.util.Objects;
+import java.util.Optional;
 
 public class SaslServerAuthenticator implements Authenticator {
-
     // GSSAPI limits requests to 64K, but we allow a bit extra for custom SASL mechanisms
     static final int MAX_RECEIVE_SIZE = 524288;
     private static final Logger LOG = LoggerFactory.getLogger(SaslServerAuthenticator.class);
-    private static final ByteBuffer EMPTY_BUFFER = ByteBuffer.allocate(0);
 
+    /**
+     * The internal state transitions for initial authentication of a channel on the
+     * server side are declared in order, starting with {@link #INITIAL_REQUEST} and
+     * ending in either {@link #COMPLETE} or {@link #FAILED}.
+     * 
        
+     * Re-authentication of a channel on the server side starts with the state
+     * {@link #REAUTH_PROCESS_HANDSHAKE}. It may then flow to
+     * {@link #REAUTH_BAD_MECHANISM} before a transition to {@link #FAILED} if
+     * re-authentication is attempted with a mechanism different than the original
+     * one; otherwise it joins the authentication flow at the {@link #AUTHENTICATE}
+     * state and likewise ends at either {@link #COMPLETE} or {@link #FAILED}.
+     */
     private enum SaslState {
-        INITIAL_REQUEST,               // May be GSSAPI token, SaslHandshake or ApiVersions
+        INITIAL_REQUEST,               // May be GSSAPI token, SaslHandshake or ApiVersions for authentication
         HANDSHAKE_OR_VERSIONS_REQUEST, // May be SaslHandshake or ApiVersions
         HANDSHAKE_REQUEST,             // After an ApiVersions request, next request must be SaslHandshake
         AUTHENTICATE,                  // Authentication tokens (SaslHandshake v1 and above indicate SaslAuthenticate headers)
         COMPLETE,                      // Authentication completed successfully
-        FAILED                         // Authentication failed
+        FAILED,                        // Authentication failed
+        REAUTH_PROCESS_HANDSHAKE,      // Initial state for re-authentication, processes SASL handshake request
+        REAUTH_BAD_MECHANISM,          // When re-authentication requested with wrong mechanism, generate exception
     }
 
     private final SecurityProtocol securityProtocol;
     private final ListenerName listenerName;
     private final String connectionId;
-    private final JaasContext jaasContext;
-    private final Subject subject;
-    private final CredentialCache credentialCache;
+    private final Map subjects;
     private final TransportLayer transportLayer;
-    private final Set enabledMechanisms;
+    private final List enabledMechanisms;
     private final Map configs;
     private final KafkaPrincipalBuilder principalBuilder;
+    private final Map callbackHandlers;
+    private final Map connectionsMaxReauthMsByMechanism;
+    private final Time time;
+    private final ReauthInfo reauthInfo;
+    private final ChannelMetadataRegistry metadataRegistry;
 
     // Current SASL state
     private SaslState saslState = SaslState.INITIAL_REQUEST;
@@ -116,71 +134,77 @@ private enum SaslState {
     private AuthenticationException pendingException = null;
     private SaslServer saslServer;
     private String saslMechanism;
-    private AuthCallbackHandler callbackHandler;
 
     // buffers used in `authenticate`
     private NetworkReceive netInBuffer;
     private Send netOutBuffer;
+    private Send authenticationFailureSend = null;
     // flag indicating if sasl tokens are sent as Kafka SaslAuthenticate request/responses
     private boolean enableKafkaSaslAuthenticateHeaders;
 
     public SaslServerAuthenticator(Map configs,
+                                   Map callbackHandlers,
                                    String connectionId,
-                                   JaasContext jaasContext,
-                                   Subject subject,
+                                   Map subjects,
                                    KerberosShortNamer kerberosNameParser,
-                                   CredentialCache credentialCache,
                                    ListenerName listenerName,
                                    SecurityProtocol securityProtocol,
-                                   TransportLayer transportLayer) throws IOException {
-        if (subject == null)
-            throw new IllegalArgumentException("subject cannot be null");
+                                   TransportLayer transportLayer,
+                                   Map connectionsMaxReauthMsByMechanism,
+                                   ChannelMetadataRegistry metadataRegistry,
+                                   Time time) {
+        this.callbackHandlers = callbackHandlers;
         this.connectionId = connectionId;
-        this.jaasContext = jaasContext;
-        this.subject = subject;
-        this.credentialCache = credentialCache;
+        this.subjects = subjects;
         this.listenerName = listenerName;
         this.securityProtocol = securityProtocol;
         this.enableKafkaSaslAuthenticateHeaders = false;
-
         this.transportLayer = transportLayer;
+        this.connectionsMaxReauthMsByMechanism = connectionsMaxReauthMsByMechanism;
+        this.time = time;
+        this.reauthInfo = new ReauthInfo();
+        this.metadataRegistry = metadataRegistry;
 
         this.configs = configs;
         @SuppressWarnings("unchecked")
         List enabledMechanisms = (List) this.configs.get(BrokerSecurityConfigs.SASL_ENABLED_MECHANISMS_CONFIG);
         if (enabledMechanisms == null || enabledMechanisms.isEmpty())
             throw new IllegalArgumentException("No SASL mechanisms are enabled");
-        this.enabledMechanisms = new HashSet<>(enabledMechanisms);
+        this.enabledMechanisms = new ArrayList<>(new HashSet<>(enabledMechanisms));
+        for (String mechanism : this.enabledMechanisms) {
+            if (!callbackHandlers.containsKey(mechanism))
+                throw new IllegalArgumentException("Callback handler not specified for SASL mechanism " + mechanism);
+            if (!subjects.containsKey(mechanism))
+                throw new IllegalArgumentException("Subject cannot be null for SASL mechanism " + mechanism);
+            LOG.debug("{} for mechanism={}: {}", BrokerSecurityConfigs.CONNECTIONS_MAX_REAUTH_MS, mechanism,
+                    connectionsMaxReauthMsByMechanism.get(mechanism));
+        }
 
         // Note that the old principal builder does not support SASL, so we do not need to pass the
         // authenticator or the transport layer
-        this.principalBuilder = ChannelBuilders.createPrincipalBuilder(configs, null, null, kerberosNameParser);
+        this.principalBuilder = ChannelBuilders.createPrincipalBuilder(configs, null, null, kerberosNameParser, null);
     }
 
     private void createSaslServer(String mechanism) throws IOException {
         this.saslMechanism = mechanism;
-        if (!ScramMechanism.isScram(mechanism))
-            callbackHandler = new SaslServerCallbackHandler(jaasContext);
-        else
-            callbackHandler = new ScramServerCallbackHandler(credentialCache.cache(mechanism, ScramCredential.class));
-        callbackHandler.configure(configs, Mode.SERVER, subject, saslMechanism);
+        Subject subject = subjects.get(mechanism);
+        final AuthenticateCallbackHandler callbackHandler = callbackHandlers.get(mechanism);
         if (mechanism.equals(SaslConfigs.GSSAPI_MECHANISM)) {
             saslServer = createSaslKerberosServer(callbackHandler, configs, subject);
         } else {
             try {
-                saslServer = Subject.doAs(subject, new PrivilegedExceptionAction() {
-                    public SaslServer run() throws SaslException {
-                        return Sasl.createSaslServer(saslMechanism, "kafka", serverAddress().getHostName(),
-                                configs, callbackHandler);
-                    }
-                });
+                saslServer = Subject.doAs(subject, (PrivilegedExceptionAction) () ->
+                    Sasl.createSaslServer(saslMechanism, "kafka", serverAddress().getHostName(), configs, callbackHandler));
+                if (saslServer == null) {
+                    throw new SaslException("Kafka Server failed to create a SaslServer to interact with a client during session authentication with server mechanism " + saslMechanism);
+                }
             } catch (PrivilegedActionException e) {
-                throw new SaslException("Kafka Server failed to create a SaslServer to interact with a client during session authentication", e.getCause());
+                throw new SaslException("Kafka Server failed to create a SaslServer to interact with a client during session authentication with server mechanism " + saslMechanism, e.getCause());
             }
         }
     }
 
-    private SaslServer createSaslKerberosServer(final AuthCallbackHandler saslServerCallbackHandler, final Map configs, Subject subject) throws IOException {
+    private SaslServer createSaslKerberosServer(final AuthenticateCallbackHandler saslServerCallbackHandler, final Map configs, Subject subject) throws IOException {
         // server is using a JAAS-authenticated subject: determine service principal name and hostname from kafka server's subject.
         final String servicePrincipal = SaslClientAuthenticator.firstPrincipal(subject);
         KerberosName kerberosName;
@@ -194,33 +218,9 @@ private SaslServer createSaslKerberosServer(final AuthCallbackHandler saslServer
 
         LOG.debug("Creating SaslServer for {} with mechanism {}", kerberosName, saslMechanism);
 
-        // As described in http://docs.oracle.com/javase/8/docs/technotes/guides/security/jgss/jgss-features.html:
-        // "To enable Java GSS to delegate to the native GSS library and its list of native mechanisms,
-        // set the system property "sun.security.jgss.native" to true"
-        // "In addition, when performing operations as a particular Subject, for example, Subject.doAs(...)
-        // or Subject.doAsPrivileged(...), the to-be-used GSSCredential should be added to Subject's
-        // private credential set. Otherwise, the GSS operations will fail since no credential is found."
-        boolean usingNativeJgss = Boolean.getBoolean("sun.security.jgss.native");
-        if (usingNativeJgss) {
-            try {
-                GSSManager manager = GSSManager.getInstance();
-                // This Oid is used to represent the Kerberos version 5 GSS-API mechanism. It is defined in
-                // RFC 1964.
-                Oid krb5Mechanism = new Oid("1.2.840.113554.1.2.2");
-                GSSName gssName = manager.createName(servicePrincipalName + "@" + serviceHostname, GSSName.NT_HOSTBASED_SERVICE);
-                GSSCredential cred = manager.createCredential(gssName, GSSContext.INDEFINITE_LIFETIME, krb5Mechanism, GSSCredential.ACCEPT_ONLY);
-                subject.getPrivateCredentials().add(cred);
-            } catch (GSSException ex) {
-                LOG.warn("Cannot add private credential to subject; clients authentication may fail", ex);
-            }
-        }
-
         try {
-            return Subject.doAs(subject, new PrivilegedExceptionAction() {
-                public SaslServer run() throws SaslException {
-                    return Sasl.createSaslServer(saslMechanism, servicePrincipalName, serviceHostname, configs, saslServerCallbackHandler);
-                }
-            });
+            return Subject.doAs(subject, (PrivilegedExceptionAction) () ->
+                    Sasl.createSaslServer(saslMechanism, servicePrincipalName, serviceHostname, configs, saslServerCallbackHandler));
         } catch (PrivilegedActionException e) {
             throw new SaslException("Kafka Server failed to create a SaslServer to interact with a client during session authentication", e.getCause());
         }
@@ -233,63 +233,77 @@ public SaslServer run() throws SaslException {
      * The messages are sent and received as size delimited bytes that consists of a 4 byte network-ordered size N
      * followed by N bytes representing the opaque payload.
      */
+    @SuppressWarnings("fallthrough")
     @Override
     public void authenticate() throws IOException {
-        if (netOutBuffer != null && !flushNetOutBufferAndUpdateInterestOps())
-            return;
+        if (saslState != SaslState.REAUTH_PROCESS_HANDSHAKE) {
+            if (netOutBuffer != null && !flushNetOutBufferAndUpdateInterestOps())
+                return;
 
-        if (saslServer != null && saslServer.isComplete()) {
-            setSaslState(SaslState.COMPLETE);
-            return;
-        }
-
-        if (netInBuffer == null) netInBuffer = new NetworkReceive(MAX_RECEIVE_SIZE, connectionId);
+            if (saslServer != null && saslServer.isComplete()) {
+                setSaslState(SaslState.COMPLETE);
+                return;
+            }
 
-        netInBuffer.readFrom(transportLayer);
+            // allocate on heap (as opposed to any socket server memory pool)
+            if (netInBuffer == null) netInBuffer = new NetworkReceive(MAX_RECEIVE_SIZE, connectionId);
 
-        if (netInBuffer.complete()) {
+            netInBuffer.readFrom(transportLayer);
+            if (!netInBuffer.complete())
+                return;
             netInBuffer.payload().rewind();
-            byte[] clientToken = new byte[netInBuffer.payload().remaining()];
-            netInBuffer.payload().get(clientToken, 0, clientToken.length);
-            netInBuffer = null; // reset the networkReceive as we read all the data.
-            try {
-                switch (saslState) {
-                    case HANDSHAKE_OR_VERSIONS_REQUEST:
-                    case HANDSHAKE_REQUEST:
-                        handleKafkaRequest(clientToken);
-                        break;
-                    case INITIAL_REQUEST:
-                        if (handleKafkaRequest(clientToken))
-                            break;
-                        // For default GSSAPI, fall through to authenticate using the client token as the first GSSAPI packet.
-                        // This is required for interoperability with 0.9.0.x clients which do not send handshake request
-                    case AUTHENTICATE:
-                        handleSaslToken(clientToken);
-                        // When the authentication exchange is complete and no more tokens are expected from the client,
-                        // update SASL state. Current SASL state will be updated when outgoing writes to the client complete.
-                        if (saslServer.isComplete())
-                            setSaslState(SaslState.COMPLETE);
-                        break;
-                    default:
+        }
+        byte[] clientToken = new byte[netInBuffer.payload().remaining()];
+        netInBuffer.payload().get(clientToken, 0, clientToken.length);
+        netInBuffer = null; // reset the networkReceive as we read all the data.
+        try {
+            switch (saslState) {
+                case REAUTH_PROCESS_HANDSHAKE:
+                case HANDSHAKE_OR_VERSIONS_REQUEST:
+                case HANDSHAKE_REQUEST:
+                    handleKafkaRequest(clientToken);
+                    break;
+                case REAUTH_BAD_MECHANISM:
+                    throw new SaslAuthenticationException(reauthInfo.badMechanismErrorMessage);
+                case INITIAL_REQUEST:
+                    if (handleKafkaRequest(clientToken))
                         break;
-                }
-            } catch (SaslException | AuthenticationException e) {
-                // Exception will be propagated after response is sent to client
-                AuthenticationException authException = (e instanceof AuthenticationException) ?
-                        (AuthenticationException) e : new AuthenticationException("SASL authentication failed", e);
-                setSaslState(SaslState.FAILED, authException);
-            } catch (Exception e) {
-                // In the case of IOExceptions and other unexpected exceptions, fail immediately
-                saslState = SaslState.FAILED;
-                throw e;
+                    // For default GSSAPI, fall through to authenticate using the client token as the first GSSAPI packet.
+                    // This is required for interoperability with 0.9.0.x clients which do not send handshake request
+                case AUTHENTICATE:
+                    handleSaslToken(clientToken);
+                    // When the authentication exchange is complete and no more tokens are expected from the client,
+                    // update SASL state. Current SASL state will be updated when outgoing writes to the client complete.
+                    if (saslServer.isComplete())
+                        setSaslState(SaslState.COMPLETE);
+                    break;
+                default:
+                    break;
             }
+        } catch (AuthenticationException e) {
+            // Exception will be propagated after response is sent to client
+            setSaslState(SaslState.FAILED, e);
+        } catch (Exception e) {
+            // In the case of IOExceptions and other unexpected exceptions, fail immediately
+            saslState = SaslState.FAILED;
+            LOG.debug("Failed during {}: {}", reauthInfo.authenticationOrReauthenticationText(), e.getMessage());
+            throw e;
         }
     }
 
     @Override
     public KafkaPrincipal principal() {
-        SaslAuthenticationContext context = new SaslAuthenticationContext(saslServer, securityProtocol, clientAddress());
-        return principalBuilder.build(context);
+        SaslAuthenticationContext context = new SaslAuthenticationContext(saslServer, securityProtocol, clientAddress(), listenerName.value());
+        KafkaPrincipal principal = principalBuilder.build(context);
+        if (ScramMechanism.isScram(saslMechanism) && Boolean.parseBoolean((String) saslServer.getNegotiatedProperty(ScramLoginModule.TOKEN_AUTH_CONFIG))) {
+            principal.tokenAuthenticated(true);
+        }
+        return principal;
+    }
+
+    @Override
+    public Optional principalSerde() {
+        return principalBuilder instanceof KafkaPrincipalSerde ? Optional.of((KafkaPrincipalSerde) principalBuilder) : Optional.empty();
     }
 
     @Override
@@ -297,27 +311,62 @@ public boolean complete() {
         return saslState == SaslState.COMPLETE;
     }
 
+    @Override
+    public void handleAuthenticationFailure() throws IOException {
+        sendAuthenticationFailureResponse();
+    }
+
     @Override
     public void close() throws IOException {
         if (principalBuilder instanceof Closeable)
             Utils.closeQuietly((Closeable) principalBuilder, "principal builder");
         if (saslServer != null)
             saslServer.dispose();
-        if (callbackHandler != null)
-            callbackHandler.close();
     }
 
-    private void setSaslState(SaslState saslState) throws IOException {
+    @Override
+    public void reauthenticate(ReauthenticationContext reauthenticationContext) throws IOException {
+        NetworkReceive saslHandshakeReceive = reauthenticationContext.networkReceive();
+        if (saslHandshakeReceive == null)
+            throw new IllegalArgumentException(
+                    "Invalid saslHandshakeReceive in server-side re-authentication context: null");
+        SaslServerAuthenticator previousSaslServerAuthenticator = (SaslServerAuthenticator) reauthenticationContext.previousAuthenticator();
+        reauthInfo.reauthenticating(previousSaslServerAuthenticator.saslMechanism,
+                previousSaslServerAuthenticator.principal(), reauthenticationContext.reauthenticationBeginNanos());
+        previousSaslServerAuthenticator.close();
+        netInBuffer = saslHandshakeReceive;
+        LOG.debug("Beginning re-authentication: {}", this);
+        netInBuffer.payload().rewind();
+        setSaslState(SaslState.REAUTH_PROCESS_HANDSHAKE);
+        authenticate();
+    }
+
+    @Override
+    public Long serverSessionExpirationTimeNanos() {
+        return reauthInfo.sessionExpirationTimeNanos;
+    }
+
+    @Override
+    public Long reauthenticationLatencyMs() {
+        return reauthInfo.reauthenticationLatencyMs();
+    }
+
+    @Override
+    public boolean connectedClientSupportsReauthentication() {
+        return reauthInfo.connectedClientSupportsReauthentication;
+    }
+
+    private void setSaslState(SaslState saslState) {
         setSaslState(saslState, null);
     }
 
-    private void setSaslState(SaslState saslState, AuthenticationException exception) throws IOException {
+    private void setSaslState(SaslState saslState, AuthenticationException exception) {
         if (netOutBuffer != null && !netOutBuffer.completed()) {
             pendingSaslState = saslState;
             pendingException = exception;
         } else {
             this.saslState = saslState;
-            LOG.debug("Set SASL server state to {}", saslState);
+            LOG.debug("Set SASL server state to {} during {}", saslState, reauthInfo.authenticationOrReauthenticationText());
             this.pendingSaslState = null;
             this.pendingException = null;
             if (exception != null)
@@ -353,8 +402,13 @@ private InetAddress clientAddress() {
     private void handleSaslToken(byte[] clientToken) throws IOException {
         if (!enableKafkaSaslAuthenticateHeaders) {
             byte[] response = saslServer.evaluateResponse(clientToken);
+            if (saslServer.isComplete()) {
+                reauthInfo.calcCompletionTimesAndReturnSessionLifetimeMs();
+                if (reauthInfo.reauthenticating())
+                    reauthInfo.ensurePrincipalUnchanged(principal());
+            }
             if (response != null) {
-                netOutBuffer = new NetworkSend(connectionId, ByteBuffer.wrap(response));
+                netOutBuffer = ByteBufferSend.sizePrefixed(ByteBuffer.wrap(response));
                 flushNetOutBufferAndUpdateInterestOps();
             }
         } else {
@@ -363,11 +417,11 @@ private void handleSaslToken(byte[] clientToken) throws IOException {
             ApiKeys apiKey = header.apiKey();
             short version = header.apiVersion();
             RequestContext requestContext = new RequestContext(header, connectionId, clientAddress(),
-                    KafkaPrincipal.ANONYMOUS, listenerName, securityProtocol);
+                    KafkaPrincipal.ANONYMOUS, listenerName, securityProtocol, ClientInformation.EMPTY, false);
             RequestAndSize requestAndSize = requestContext.parseRequest(requestBuffer);
             if (apiKey != ApiKeys.SASL_AUTHENTICATE) {
                 IllegalSaslStateException e = new IllegalSaslStateException("Unexpected Kafka request of type " + apiKey + " during SASL authentication.");
-                sendKafkaResponse(requestContext, requestAndSize.request.getErrorResponse(e));
+                buildResponseOnAuthenticateFailure(requestContext, requestAndSize.request.getErrorResponse(e));
                 throw e;
             }
             if (!apiKey.isVersionSupported(version)) {
@@ -375,19 +429,53 @@ private void handleSaslToken(byte[] clientToken) throws IOException {
                 // This should not normally occur since clients typically check supported versions using ApiVersionsRequest
                 throw new UnsupportedVersionException("Version " + version + " is not supported for apiKey " + apiKey);
             }
+            /*
+             * The client sends multiple SASL_AUTHENTICATE requests, and the client is known
+             * to support the required version if any one of them indicates it supports that
+             * version.
+             */
+            if (!reauthInfo.connectedClientSupportsReauthentication)
+                reauthInfo.connectedClientSupportsReauthentication = version > 0;
             SaslAuthenticateRequest saslAuthenticateRequest = (SaslAuthenticateRequest) requestAndSize.request;
 
             try {
-                byte[] responseToken = saslServer.evaluateResponse(Utils.readBytes(saslAuthenticateRequest.saslAuthBytes()));
+                byte[] responseToken = saslServer.evaluateResponse(
+                        Utils.copyArray(saslAuthenticateRequest.data().authBytes()));
+                if (reauthInfo.reauthenticating() && saslServer.isComplete())
+                    reauthInfo.ensurePrincipalUnchanged(principal());
                 // For versions with SASL_AUTHENTICATE header, send a response to SASL_AUTHENTICATE request even if token is empty.
-                ByteBuffer responseBuf = responseToken == null ? EMPTY_BUFFER : ByteBuffer.wrap(responseToken);
-                sendKafkaResponse(requestContext, new SaslAuthenticateResponse(Errors.NONE, null, responseBuf));
-            } catch (SaslAuthenticationException | SaslException e) {
-                String errorMessage = e instanceof SaslAuthenticationException ? e.getMessage() :
-                    "Authentication failed due to invalid credentials with SASL mechanism " + saslMechanism;
-                sendKafkaResponse(requestContext, new SaslAuthenticateResponse(Errors.SASL_AUTHENTICATION_FAILED,
-                        errorMessage));
+                byte[] responseBytes = responseToken == null ? new byte[0] : responseToken;
+                long sessionLifetimeMs = !saslServer.isComplete() ? 0L
+                        : reauthInfo.calcCompletionTimesAndReturnSessionLifetimeMs();
+                sendKafkaResponse(requestContext, new SaslAuthenticateResponse(
+                        new SaslAuthenticateResponseData()
+                        .setErrorCode(Errors.NONE.code())
+                        .setAuthBytes(responseBytes)
+                        .setSessionLifetimeMs(sessionLifetimeMs)));
+            } catch (SaslAuthenticationException e) {
+                buildResponseOnAuthenticateFailure(requestContext,
+                        new SaslAuthenticateResponse(
+                                new SaslAuthenticateResponseData()
+                                .setErrorCode(Errors.SASL_AUTHENTICATION_FAILED.code())
+                                .setErrorMessage(e.getMessage())));
                 throw e;
+            } catch (SaslException e) {
+                KerberosError kerberosError = KerberosError.fromException(e);
+                if (kerberosError != null && kerberosError.retriable()) {
+                    // Handle retriable Kerberos exceptions as I/O exceptions rather than authentication exceptions
+                    throw e;
+                } else {
+                    // DO NOT include error message from the `SaslException` in the client response since it may
+                    // contain sensitive data like the existence of the user.
+                    String errorMessage = "Authentication failed during "
+                            + reauthInfo.authenticationOrReauthenticationText()
+                            + " due to invalid credentials with SASL mechanism " + saslMechanism;
+                    buildResponseOnAuthenticateFailure(requestContext, new SaslAuthenticateResponse(
+                            new SaslAuthenticateResponseData()
+                            .setErrorCode(Errors.SASL_AUTHENTICATION_FAILED.code())
+                            .setErrorMessage(errorMessage)));
+                    throw new SaslAuthenticationException(errorMessage, e);
+                }
             }
         }
     }
@@ -411,11 +499,11 @@ private boolean handleKafkaRequest(byte[] requestBytes) throws IOException, Auth
             if (apiKey != ApiKeys.API_VERSIONS && apiKey != ApiKeys.SASL_HANDSHAKE)
                 throw new IllegalSaslStateException("Unexpected Kafka request of type " + apiKey + " during SASL handshake.");
 
-            LOG.debug("Handling Kafka request {}", apiKey);
+            LOG.debug("Handling Kafka request {} during {}", apiKey, reauthInfo.authenticationOrReauthenticationText());
 
 
             RequestContext requestContext = new RequestContext(header, connectionId, clientAddress(),
-                    KafkaPrincipal.ANONYMOUS, listenerName, securityProtocol);
+                    KafkaPrincipal.ANONYMOUS, listenerName, securityProtocol, ClientInformation.EMPTY, false);
             RequestAndSize requestAndSize = requestContext.parseRequest(requestBuffer);
             if (apiKey == ApiKeys.API_VERSIONS)
                 handleApiVersionsRequest(requestContext, (ApiVersionsRequest) requestAndSize.request);
@@ -443,7 +531,8 @@ private boolean handleKafkaRequest(byte[] requestBytes) throws IOException, Auth
             } else
                 throw e;
         }
-        if (clientMechanism != null) {
+        if (clientMechanism != null && (!reauthInfo.reauthenticating()
+                || reauthInfo.saslMechanismUnchanged(clientMechanism))) {
             createSaslServer(clientMechanism);
             setSaslState(SaslState.AUTHENTICATE);
         }
@@ -451,24 +540,26 @@ private boolean handleKafkaRequest(byte[] requestBytes) throws IOException, Auth
     }
 
     private String handleHandshakeRequest(RequestContext context, SaslHandshakeRequest handshakeRequest) throws IOException, UnsupportedSaslMechanismException {
-        String clientMechanism = handshakeRequest.mechanism();
+        String clientMechanism = handshakeRequest.data().mechanism();
         short version = context.header.apiVersion();
         if (version >= 1)
             this.enableKafkaSaslAuthenticateHeaders(true);
         if (enabledMechanisms.contains(clientMechanism)) {
             LOG.debug("Using SASL mechanism '{}' provided by client", clientMechanism);
-            sendKafkaResponse(context, new SaslHandshakeResponse(Errors.NONE, enabledMechanisms));
+            sendKafkaResponse(context, new SaslHandshakeResponse(
+                    new SaslHandshakeResponseData().setErrorCode(Errors.NONE.code()).setMechanisms(enabledMechanisms)));
             return clientMechanism;
         } else {
             LOG.debug("SASL mechanism '{}' requested by client is not supported", clientMechanism);
-            sendKafkaResponse(context, new SaslHandshakeResponse(Errors.UNSUPPORTED_SASL_MECHANISM, enabledMechanisms));
+            buildResponseOnAuthenticateFailure(context, new SaslHandshakeResponse(
+                    new SaslHandshakeResponseData().setErrorCode(Errors.UNSUPPORTED_SASL_MECHANISM.code()).setMechanisms(enabledMechanisms)));
             throw new UnsupportedSaslMechanismException("Unsupported SASL mechanism " + clientMechanism);
         }
     }
 
     // Visible to override for testing
     protected ApiVersionsResponse apiVersionsResponse() {
-        return ApiVersionsResponse.defaultApiVersionsResponse();
+        return ApiVersionsResponse.DEFAULT_API_VERSIONS_RESPONSE;
     }
 
     // Visible to override for testing
@@ -482,18 +573,150 @@ 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 {
+            metadataRegistry.registerClientInformation(new ClientInformation(apiVersionsRequest.data().clientSoftwareName(),
+                apiVersionsRequest.data().clientSoftwareVersion()));
             sendKafkaResponse(context, apiVersionsResponse());
             setSaslState(SaslState.HANDSHAKE_REQUEST);
         }
     }
 
+    /**
+     * Build a {@link Send} response on {@link #authenticate()} failure. The actual response is sent out when
+     * {@link #sendAuthenticationFailureResponse()} is called.
+     */
+    private void buildResponseOnAuthenticateFailure(RequestContext context, AbstractResponse response) {
+        authenticationFailureSend = context.buildResponseSend(response);
+    }
+
+    /**
+     * Send any authentication failure response that may have been previously built.
+     */
+    private void sendAuthenticationFailureResponse() throws IOException {
+        if (authenticationFailureSend == null)
+            return;
+        sendKafkaResponse(authenticationFailureSend);
+        authenticationFailureSend = null;
+    }
+
     private void sendKafkaResponse(RequestContext context, AbstractResponse response) throws IOException {
-        sendKafkaResponse(context.buildResponse(response));
+        sendKafkaResponse(context.buildResponseSend(response));
     }
 
     private void sendKafkaResponse(Send send) throws IOException {
         netOutBuffer = send;
         flushNetOutBufferAndUpdateInterestOps();
     }
+
+    /**
+     * Information related to re-authentication
+     */
+    private class ReauthInfo {
+        public String previousSaslMechanism;
+        public KafkaPrincipal previousKafkaPrincipal;
+        public long reauthenticationBeginNanos;
+        public Long sessionExpirationTimeNanos;
+        public boolean connectedClientSupportsReauthentication;
+        public long authenticationEndNanos;
+        public String badMechanismErrorMessage;
+
+        public void reauthenticating(String previousSaslMechanism, KafkaPrincipal previousKafkaPrincipal,
+                long reauthenticationBeginNanos) {
+            this.previousSaslMechanism = Objects.requireNonNull(previousSaslMechanism);
+            this.previousKafkaPrincipal = Objects.requireNonNull(previousKafkaPrincipal);
+            this.reauthenticationBeginNanos = reauthenticationBeginNanos;
+        }
+
+        public boolean reauthenticating() {
+            return previousSaslMechanism != null;
+        }
+
+        public String authenticationOrReauthenticationText() {
+            return reauthenticating() ? "re-authentication" : "authentication";
+        }
+
+        public void ensurePrincipalUnchanged(KafkaPrincipal reauthenticatedKafkaPrincipal) throws SaslAuthenticationException {
+            if (!previousKafkaPrincipal.equals(reauthenticatedKafkaPrincipal)) {
+                throw new SaslAuthenticationException(String.format(
+                        "Cannot change principals during re-authentication from %s.%s: %s.%s",
+                        previousKafkaPrincipal.getPrincipalType(), previousKafkaPrincipal.getName(),
+                        reauthenticatedKafkaPrincipal.getPrincipalType(), reauthenticatedKafkaPrincipal.getName()));
+            }
+        }
+
+        /*
+         * We define the REAUTH_BAD_MECHANISM state because the failed re-authentication
+         * metric does not get updated if we send back an error immediately upon the
+         * start of re-authentication.
+         */
+        public boolean saslMechanismUnchanged(String clientMechanism) {
+            if (previousSaslMechanism.equals(clientMechanism))
+                return true;
+            badMechanismErrorMessage = String.format(
+                    "SASL mechanism '%s' requested by client is not supported for re-authentication of mechanism '%s'",
+                    clientMechanism, previousSaslMechanism);
+            LOG.debug(badMechanismErrorMessage);
+            setSaslState(SaslState.REAUTH_BAD_MECHANISM);
+            return false;
+        }
+
+        private long calcCompletionTimesAndReturnSessionLifetimeMs() {
+            long retvalSessionLifetimeMs = 0L;
+            long authenticationEndMs = time.milliseconds();
+            authenticationEndNanos = time.nanoseconds();
+            Long credentialExpirationMs = (Long) saslServer
+                    .getNegotiatedProperty(SaslInternalConfigs.CREDENTIAL_LIFETIME_MS_SASL_NEGOTIATED_PROPERTY_KEY);
+            Long connectionsMaxReauthMs = connectionsMaxReauthMsByMechanism.get(saslMechanism);
+            if (credentialExpirationMs != null || connectionsMaxReauthMs != null) {
+                if (credentialExpirationMs == null)
+                    retvalSessionLifetimeMs = zeroIfNegative(connectionsMaxReauthMs);
+                else if (connectionsMaxReauthMs == null)
+                    retvalSessionLifetimeMs = zeroIfNegative(credentialExpirationMs - authenticationEndMs);
+                else
+                    retvalSessionLifetimeMs = zeroIfNegative(
+                            Math.min(credentialExpirationMs - authenticationEndMs, connectionsMaxReauthMs));
+                if (retvalSessionLifetimeMs > 0L)
+                    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),
+                            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),
+                            credentialExpirationMs - authenticationEndMs);
+            } else {
+                if (sessionExpirationTimeNanos != null)
+                    LOG.debug(
+                            "Authentication complete; session max lifetime from broker config={} ms, no credential expiration; session expiration = {} ({} ms), sending {} ms to client",
+                            connectionsMaxReauthMs, new Date(authenticationEndMs + retvalSessionLifetimeMs),
+                            retvalSessionLifetimeMs, retvalSessionLifetimeMs);
+                else
+                    LOG.debug(
+                            "Authentication complete; session max lifetime from broker config={} ms, no credential expiration; no session expiration, sending 0 ms to client",
+                            connectionsMaxReauthMs);
+            }
+            return retvalSessionLifetimeMs;
+        }
+
+        public Long reauthenticationLatencyMs() {
+            if (!reauthenticating())
+                return null;
+            // record at least 1 ms if there is some latency
+            long latencyNanos = authenticationEndNanos - reauthenticationBeginNanos;
+            return latencyNanos == 0L ? 0L : Math.max(1L, Math.round(latencyNanos / 1000.0 / 1000.0));
+        }
+
+        private long zeroIfNegative(long value) {
+            return Math.max(0L, value);
+        }        
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/security/authenticator/SaslServerCallbackHandler.java b/clients/src/main/java/org/apache/kafka/common/security/authenticator/SaslServerCallbackHandler.java
index 7d5372db9771d..d3d43cbfa26d2 100644
--- a/clients/src/main/java/org/apache/kafka/common/security/authenticator/SaslServerCallbackHandler.java
+++ b/clients/src/main/java/org/apache/kafka/common/security/authenticator/SaslServerCallbackHandler.java
@@ -16,51 +16,46 @@
  */
 package org.apache.kafka.common.security.authenticator;
 
-import java.io.IOException;
+import java.util.List;
 import java.util.Map;
 
-import org.apache.kafka.common.security.JaasContext;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import javax.security.auth.Subject;
 import javax.security.auth.callback.Callback;
 import javax.security.auth.callback.UnsupportedCallbackException;
+import javax.security.auth.login.AppConfigurationEntry;
 import javax.security.sasl.AuthorizeCallback;
 import javax.security.sasl.RealmCallback;
 
-import org.apache.kafka.common.network.Mode;
+import org.apache.kafka.common.config.SaslConfigs;
+import org.apache.kafka.common.security.auth.AuthenticateCallbackHandler;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /**
- * Callback handler for Sasl servers. The callbacks required for all the SASL
+ * Default callback handler for Sasl servers. The callbacks required for all the SASL
  * mechanisms enabled in the server should be supported by this callback handler. See
  * Java SASL API
  * for the list of SASL callback handlers required for each SASL mechanism.
  */
-public class SaslServerCallbackHandler implements AuthCallbackHandler {
+public class SaslServerCallbackHandler implements AuthenticateCallbackHandler {
     private static final Logger LOG = LoggerFactory.getLogger(SaslServerCallbackHandler.class);
-    private final JaasContext jaasContext;
 
-    public SaslServerCallbackHandler(JaasContext jaasContext) throws IOException {
-        this.jaasContext = jaasContext;
-    }
+    private String mechanism;
 
     @Override
-    public void configure(Map configs, Mode mode, Subject subject, String saslMechanism) {
-    }
-
-    public JaasContext jaasContext() {
-        return jaasContext;
+    public void configure(Map configs, String mechanism, List jaasConfigEntries) {
+        this.mechanism = mechanism;
     }
 
     @Override
     public void handle(Callback[] callbacks) throws UnsupportedCallbackException {
         for (Callback callback : callbacks) {
-            if (callback instanceof RealmCallback) {
+            if (callback instanceof RealmCallback)
                 handleRealmCallback((RealmCallback) callback);
-            } else if (callback instanceof AuthorizeCallback) {
+            else if (callback instanceof AuthorizeCallback && mechanism.equals(SaslConfigs.GSSAPI_MECHANISM))
                 handleAuthorizeCallback((AuthorizeCallback) callback);
-            }
+            else
+                throw new UnsupportedCallbackException(callback);
         }
     }
 
diff --git a/clients/src/main/java/org/apache/kafka/common/security/kerberos/KerberosClientCallbackHandler.java b/clients/src/main/java/org/apache/kafka/common/security/kerberos/KerberosClientCallbackHandler.java
new file mode 100644
index 0000000000000..fa9cad261c589
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/kerberos/KerberosClientCallbackHandler.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.common.security.kerberos;
+
+import org.apache.kafka.common.config.SaslConfigs;
+import org.apache.kafka.common.security.auth.AuthenticateCallbackHandler;
+
+import javax.security.auth.callback.Callback;
+import javax.security.auth.callback.NameCallback;
+import javax.security.auth.callback.PasswordCallback;
+import javax.security.auth.callback.UnsupportedCallbackException;
+import javax.security.auth.login.AppConfigurationEntry;
+import javax.security.sasl.AuthorizeCallback;
+import javax.security.sasl.RealmCallback;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Callback handler for SASL/GSSAPI clients.
+ */
+public class KerberosClientCallbackHandler implements AuthenticateCallbackHandler {
+
+    @Override
+    public void configure(Map configs, String saslMechanism, List jaasConfigEntries) {
+        if (!saslMechanism.equals(SaslConfigs.GSSAPI_MECHANISM))
+            throw new IllegalStateException("Kerberos callback handler should only be used with GSSAPI");
+    }
+
+    @Override
+    public void handle(Callback[] callbacks) throws UnsupportedCallbackException {
+        for (Callback callback : callbacks) {
+            if (callback instanceof NameCallback) {
+                NameCallback nc = (NameCallback) callback;
+                nc.setName(nc.getDefaultName());
+            } else if (callback instanceof PasswordCallback) {
+                String errorMessage = "Could not login: the client is being asked for a password, but the Kafka" +
+                             " client code does not currently support obtaining a password from the user.";
+                errorMessage += " Make sure -Djava.security.auth.login.config property passed to JVM and" +
+                             " the client is configured to use a ticket cache (using" +
+                             " the JAAS configuration setting 'useTicketCache=true)'. Make sure you are using" +
+                             " FQDN of the Kafka broker you are trying to connect to.";
+                throw new UnsupportedCallbackException(callback, errorMessage);
+            } else if (callback instanceof RealmCallback) {
+                RealmCallback rc = (RealmCallback) callback;
+                rc.setText(rc.getDefaultText());
+            } else if (callback instanceof AuthorizeCallback) {
+                AuthorizeCallback ac = (AuthorizeCallback) callback;
+                String authId = ac.getAuthenticationID();
+                String authzId = ac.getAuthorizationID();
+                ac.setAuthorized(authId.equals(authzId));
+                if (ac.isAuthorized())
+                    ac.setAuthorizedID(authzId);
+            }  else {
+                throw new UnsupportedCallbackException(callback, "Unrecognized SASL ClientCallback");
+            }
+        }
+    }
+
+    @Override
+    public void close() {
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/kerberos/KerberosError.java b/clients/src/main/java/org/apache/kafka/common/security/kerberos/KerberosError.java
new file mode 100644
index 0000000000000..4b8e8e0d67ba3
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/kerberos/KerberosError.java
@@ -0,0 +1,131 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.common.security.kerberos;
+
+import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.security.authenticator.SaslClientAuthenticator;
+import org.apache.kafka.common.utils.Java;
+import org.ietf.jgss.GSSException;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.security.sasl.SaslClient;
+import java.lang.reflect.Method;
+
+/**
+ * Kerberos exceptions that may require special handling. The standard Kerberos error codes
+ * for these errors are retrieved using KrbException#errorCode() from the underlying Kerberos
+ * exception thrown during {@link SaslClient#evaluateChallenge(byte[])}.
+ */
+public enum KerberosError {
+    // (Mechanism level: Server not found in Kerberos database (7) - UNKNOWN_SERVER)
+    // This is retriable, but included here to add extra logging for this case.
+    SERVER_NOT_FOUND(7, false),
+    // (Mechanism level: Client not yet valid - try again later (21))
+    CLIENT_NOT_YET_VALID(21, true),
+    // (Mechanism level: Ticket not yet valid (33) - Ticket not yet valid)])
+    // This could be a small timing window.
+    TICKET_NOT_YET_VALID(33, true),
+    // (Mechanism level: Request is a replay (34) - Request is a replay)
+    // Replay detection used to prevent DoS attacks can result in false positives, so retry on error.
+    REPLAY(34, true);
+
+    private static final Logger log = LoggerFactory.getLogger(SaslClientAuthenticator.class);
+    private static final Class KRB_EXCEPTION_CLASS;
+    private static final Method KRB_EXCEPTION_RETURN_CODE_METHOD;
+
+    static {
+        try {
+            // different IBM JDKs versions include different security implementations
+            if (Java.isIbmJdk() && canLoad("com.ibm.security.krb5.KrbException")) {
+                KRB_EXCEPTION_CLASS = Class.forName("com.ibm.security.krb5.KrbException");
+            } else if (Java.isIbmJdk() && canLoad("com.ibm.security.krb5.internal.KrbException")) {
+                KRB_EXCEPTION_CLASS = Class.forName("com.ibm.security.krb5.internal.KrbException");
+            } else {
+                KRB_EXCEPTION_CLASS = Class.forName("sun.security.krb5.KrbException");
+            }
+            KRB_EXCEPTION_RETURN_CODE_METHOD = KRB_EXCEPTION_CLASS.getMethod("returnCode");
+        } catch (Exception e) {
+            throw new KafkaException("Kerberos exceptions could not be initialized", e);
+        }
+    }
+
+    private static boolean canLoad(String clazz) {
+        try {
+            Class.forName(clazz);
+            return true;
+        } catch (Exception e) {
+            return false;
+        }
+    }
+
+    private final int errorCode;
+    private final boolean retriable;
+
+    KerberosError(int errorCode, boolean retriable) {
+        this.errorCode = errorCode;
+        this.retriable = retriable;
+    }
+
+    public boolean retriable() {
+        return retriable;
+    }
+
+    public static KerberosError fromException(Exception exception) {
+        Throwable cause = exception.getCause();
+        while (cause != null && !KRB_EXCEPTION_CLASS.isInstance(cause)) {
+            cause = cause.getCause();
+        }
+        if (cause == null)
+            return null;
+        else {
+            try {
+                Integer errorCode = (Integer) KRB_EXCEPTION_RETURN_CODE_METHOD.invoke(cause);
+                return fromErrorCode(errorCode);
+            } catch (Exception e) {
+                log.trace("Kerberos return code could not be determined from {} due to {}", exception, e);
+                return null;
+            }
+        }
+    }
+
+    private static KerberosError fromErrorCode(int errorCode) {
+        for (KerberosError error : values()) {
+            if (error.errorCode == errorCode)
+                return error;
+        }
+        return null;
+    }
+
+    /**
+     * Returns true if the exception should be handled as a transient failure on clients.
+     * We handle GSSException.NO_CRED as retriable on the client-side since this may
+     * occur during re-login if a clients attempts to authentication after logout, but
+     * before the subsequent login.
+     */
+    public static boolean isRetriableClientGssException(Exception exception) {
+        Throwable cause = exception.getCause();
+        while (cause != null && !(cause instanceof GSSException)) {
+            cause = cause.getCause();
+        }
+        if (cause != null) {
+            GSSException gssException = (GSSException) cause;
+            return gssException.getMajor() == GSSException.NO_CRED;
+        }
+        return false;
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/kerberos/KerberosLogin.java b/clients/src/main/java/org/apache/kafka/common/security/kerberos/KerberosLogin.java
index a0aad542971ab..9626da846df38 100644
--- a/clients/src/main/java/org/apache/kafka/common/security/kerberos/KerberosLogin.java
+++ b/clients/src/main/java/org/apache/kafka/common/security/kerberos/KerberosLogin.java
@@ -18,6 +18,7 @@
 
 import javax.security.auth.kerberos.KerberosPrincipal;
 import javax.security.auth.login.AppConfigurationEntry;
+import javax.security.auth.login.Configuration;
 import javax.security.auth.login.LoginContext;
 import javax.security.auth.login.LoginException;
 import javax.security.auth.kerberos.KerberosTicket;
@@ -25,6 +26,7 @@
 
 import org.apache.kafka.common.security.JaasContext;
 import org.apache.kafka.common.security.JaasUtils;
+import org.apache.kafka.common.security.auth.AuthenticateCallbackHandler;
 import org.apache.kafka.common.security.authenticator.AbstractLogin;
 import org.apache.kafka.common.config.SaslConfigs;
 import org.apache.kafka.common.utils.KafkaThread;
@@ -33,11 +35,12 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.util.Arrays;
 import java.util.Date;
 import java.util.List;
+import java.util.Map;
 import java.util.Random;
 import java.util.Set;
-import java.util.Map;
 
 /**
  * This class is responsible for refreshing Kerberos credentials for
@@ -78,13 +81,15 @@ public class KerberosLogin extends AbstractLogin {
     private String serviceName;
     private long lastLogin;
 
-    public void configure(Map configs, JaasContext jaasContext) {
-        super.configure(configs, jaasContext);
+    @Override
+    public void configure(Map configs, String contextName, Configuration configuration,
+                          AuthenticateCallbackHandler callbackHandler) {
+        super.configure(configs, contextName, configuration, callbackHandler);
         this.ticketRenewWindowFactor = (Double) configs.get(SaslConfigs.SASL_KERBEROS_TICKET_RENEW_WINDOW_FACTOR);
         this.ticketRenewJitter = (Double) configs.get(SaslConfigs.SASL_KERBEROS_TICKET_RENEW_JITTER);
         this.minTimeBeforeRelogin = (Long) configs.get(SaslConfigs.SASL_KERBEROS_MIN_TIME_BEFORE_RELOGIN);
         this.kinitCmd = (String) configs.get(SaslConfigs.SASL_KERBEROS_KINIT_CMD);
-        this.serviceName = getServiceName(configs, jaasContext);
+        this.serviceName = getServiceName(configs, contextName, configuration);
     }
 
     /**
@@ -99,13 +104,13 @@ public LoginContext login() throws LoginException {
         subject = loginContext.getSubject();
         isKrbTicket = !subject.getPrivateCredentials(KerberosTicket.class).isEmpty();
 
-        List entries = jaasContext().configurationEntries();
-        if (entries.isEmpty()) {
+        AppConfigurationEntry[] entries = configuration().getAppConfigurationEntry(contextName());
+        if (entries.length == 0) {
             isUsingTicketCache = false;
             principal = null;
         } else {
             // there will only be a single entry
-            AppConfigurationEntry entry = entries.get(0);
+            AppConfigurationEntry entry = entries[0];
             if (entry.getOptions().get("useTicketCache") != null) {
                 String val = (String) entry.getOptions().get("useTicketCache");
                 isUsingTicketCache = val.equals("true");
@@ -129,128 +134,127 @@ public LoginContext login() throws LoginException {
         // TGT's existing expiry date and the configured minTimeBeforeRelogin. For testing and development,
         // you can decrease the interval of expiration of tickets (for example, to 3 minutes) by running:
         //  "modprinc -maxlife 3mins " in kadmin.
-        t = KafkaThread.daemon(String.format("kafka-kerberos-refresh-thread-%s", principal), new Runnable() {
-            public void run() {
-                log.info("[Principal={}]: TGT refresh thread started.", principal);
-                while (true) {  // renewal thread's main loop. if it exits from here, thread will exit.
-                    KerberosTicket tgt = getTGT();
-                    long now = currentWallTime();
-                    long nextRefresh;
-                    Date nextRefreshDate;
-                    if (tgt == null) {
-                        nextRefresh = now + minTimeBeforeRelogin;
-                        nextRefreshDate = new Date(nextRefresh);
-                        log.warn("[Principal={}]: No TGT found: will try again at {}", principal, nextRefreshDate);
+        t = KafkaThread.daemon(String.format("kafka-kerberos-refresh-thread-%s", principal), () -> {
+            log.info("[Principal={}]: TGT refresh thread started.", principal);
+            while (true) {  // renewal thread's main loop. if it exits from here, thread will exit.
+                KerberosTicket tgt = getTGT();
+                long now = currentWallTime();
+                long nextRefresh;
+                Date nextRefreshDate;
+                if (tgt == null) {
+                    nextRefresh = now + minTimeBeforeRelogin;
+                    nextRefreshDate = new Date(nextRefresh);
+                    log.warn("[Principal={}]: No TGT found: will try again at {}", principal, nextRefreshDate);
+                } else {
+                    nextRefresh = getRefreshTime(tgt);
+                    long expiry = tgt.getEndTime().getTime();
+                    Date expiryDate = new Date(expiry);
+                    if (isUsingTicketCache && tgt.getRenewTill() != null && tgt.getRenewTill().getTime() < expiry) {
+                        log.warn("The TGT cannot be renewed beyond the next expiry date: {}." +
+                            "This process will not be able to authenticate new SASL connections after that " +
+                            "time (for example, it will not be able to authenticate a new connection with a Kafka " +
+                            "Broker).  Ask your system administrator to either increase the " +
+                            "'renew until' time by doing : 'modprinc -maxrenewlife {} ' within " +
+                            "kadmin, or instead, to generate a keytab for {}. Because the TGT's " +
+                            "expiry cannot be further extended by refreshing, exiting refresh thread now.",
+                            expiryDate, principal, principal);
+                        return;
+                    }
+                    // determine how long to sleep from looking at ticket's expiry.
+                    // We should not allow the ticket to expire, but we should take into consideration
+                    // minTimeBeforeRelogin. Will not sleep less than minTimeBeforeRelogin, unless doing so
+                    // would cause ticket expiration.
+                    if ((nextRefresh > expiry) || (now + minTimeBeforeRelogin > expiry)) {
+                        // expiry is before next scheduled refresh).
+                        log.info("[Principal={}]: Refreshing now because expiry is before next scheduled refresh time.", principal);
+                        nextRefresh = now;
                     } else {
-                        nextRefresh = getRefreshTime(tgt);
-                        long expiry = tgt.getEndTime().getTime();
-                        Date expiryDate = new Date(expiry);
-                        if (isUsingTicketCache && tgt.getRenewTill() != null && tgt.getRenewTill().getTime() < expiry) {
-                            log.warn("The TGT cannot be renewed beyond the next expiry date: {}." +
-                                    "This process will not be able to authenticate new SASL connections after that " +
-                                    "time (for example, it will not be able to authenticate a new connection with a Kafka " +
-                                    "Broker).  Ask your system administrator to either increase the " +
-                                    "'renew until' time by doing : 'modprinc -maxrenewlife {} ' within " +
-                                    "kadmin, or instead, to generate a keytab for {}. Because the TGT's " +
-                                    "expiry cannot be further extended by refreshing, exiting refresh thread now.",
-                                    expiryDate, principal, principal);
-                            return;
-                        }
-                        // determine how long to sleep from looking at ticket's expiry.
-                        // We should not allow the ticket to expire, but we should take into consideration
-                        // minTimeBeforeRelogin. Will not sleep less than minTimeBeforeRelogin, unless doing so
-                        // would cause ticket expiration.
-                        if ((nextRefresh > expiry) || (now + minTimeBeforeRelogin > expiry)) {
-                            // expiry is before next scheduled refresh).
-                            log.info("[Principal={}]: Refreshing now because expiry is before next scheduled refresh time.", principal);
-                            nextRefresh = now;
-                        } else {
-                            if (nextRefresh < (now + minTimeBeforeRelogin)) {
-                                // next scheduled refresh is sooner than (now + MIN_TIME_BEFORE_LOGIN).
-                                Date until = new Date(nextRefresh);
-                                Date newUntil = new Date(now + minTimeBeforeRelogin);
-                                log.warn("[Principal={}]: TGT refresh thread time adjusted from {} to {} since the former is sooner " +
-                                        "than the minimum refresh interval ({} seconds) from now.",
-                                        principal, until, newUntil, minTimeBeforeRelogin / 1000);
-                            }
-                            nextRefresh = Math.max(nextRefresh, now + minTimeBeforeRelogin);
-                        }
-                        nextRefreshDate = new Date(nextRefresh);
-                        if (nextRefresh > expiry) {
-                            log.error("[Principal={}]: Next refresh: {} is later than expiry {}. This may indicate a clock skew problem." +
-                                    "Check that this host and the KDC hosts' clocks are in sync. Exiting refresh thread.",
-                                    principal, nextRefreshDate, expiryDate);
-                            return;
+                        if (nextRefresh < (now + minTimeBeforeRelogin)) {
+                            // next scheduled refresh is sooner than (now + MIN_TIME_BEFORE_LOGIN).
+                            Date until = new Date(nextRefresh);
+                            Date newUntil = new Date(now + minTimeBeforeRelogin);
+                            log.warn("[Principal={}]: TGT refresh thread time adjusted from {} to {} since the former is sooner " +
+                                "than the minimum refresh interval ({} seconds) from now.",
+                                principal, until, newUntil, minTimeBeforeRelogin / 1000);
                         }
+                        nextRefresh = Math.max(nextRefresh, now + minTimeBeforeRelogin);
                     }
-                    if (now < nextRefresh) {
-                        Date until = new Date(nextRefresh);
-                        log.info("[Principal={}]: TGT refresh sleeping until: {}", principal, until);
-                        try {
-                            Thread.sleep(nextRefresh - now);
-                        } catch (InterruptedException ie) {
-                            log.warn("[Principal={}]: TGT renewal thread has been interrupted and will exit.", principal);
-                            return;
-                        }
-                    } else {
-                        log.error("[Principal={}]: NextRefresh: {} is in the past: exiting refresh thread. Check"
-                                + " clock sync between this host and KDC - (KDC's clock is likely ahead of this host)."
-                                + " Manual intervention will be required for this client to successfully authenticate."
-                                + " Exiting refresh thread.", principal, nextRefreshDate);
+                    nextRefreshDate = new Date(nextRefresh);
+                    if (nextRefresh > expiry) {
+                        log.error("[Principal={}]: Next refresh: {} is later than expiry {}. This may indicate a clock skew problem." +
+                            "Check that this host and the KDC hosts' clocks are in sync. Exiting refresh thread.",
+                            principal, nextRefreshDate, expiryDate);
                         return;
                     }
-                    if (isUsingTicketCache) {
-                        String kinitArgs = "-R";
-                        int retry = 1;
-                        while (retry >= 0) {
-                            try {
-                                log.debug("[Principal={}]: Running ticket cache refresh command: {} {}", principal, kinitCmd, kinitArgs);
-                                Shell.execCommand(kinitCmd, kinitArgs);
-                                break;
-                            } catch (Exception e) {
-                                if (retry > 0) {
-                                    --retry;
-                                    // sleep for 10 seconds
-                                    try {
-                                        Thread.sleep(10 * 1000);
-                                    } catch (InterruptedException ie) {
-                                        log.error("[Principal={}]: Interrupted while renewing TGT, exiting Login thread", principal);
-                                        return;
-                                    }
-                                } else {
-                                    log.warn("[Principal={}]: Could not renew TGT due to problem running shell command: '{} {}'; " +
-                                            "exception was: %s. Exiting refresh thread.", 
-                                            principal, kinitCmd, kinitArgs, e, e);
+                }
+                if (now < nextRefresh) {
+                    Date until = new Date(nextRefresh);
+                    log.info("[Principal={}]: TGT refresh sleeping until: {}", principal, until);
+                    try {
+                        Thread.sleep(nextRefresh - now);
+                    } catch (InterruptedException ie) {
+                        log.warn("[Principal={}]: TGT renewal thread has been interrupted and will exit.", principal);
+                        return;
+                    }
+                } else {
+                    log.error("[Principal={}]: NextRefresh: {} is in the past: exiting refresh thread. Check"
+                        + " clock sync between this host and KDC - (KDC's clock is likely ahead of this host)."
+                        + " Manual intervention will be required for this client to successfully authenticate."
+                        + " Exiting refresh thread.", principal, nextRefreshDate);
+                    return;
+                }
+                if (isUsingTicketCache) {
+                    String kinitArgs = "-R";
+                    int retry = 1;
+                    while (retry >= 0) {
+                        try {
+                            log.debug("[Principal={}]: Running ticket cache refresh command: {} {}", principal, kinitCmd, kinitArgs);
+                            Shell.execCommand(kinitCmd, kinitArgs);
+                            break;
+                        } catch (Exception e) {
+                            if (retry > 0) {
+                                log.warn("[Principal={}]: Error when trying to renew with TicketCache, but will retry ", principal, e);
+                                --retry;
+                                // sleep for 10 seconds
+                                try {
+                                    Thread.sleep(10 * 1000);
+                                } catch (InterruptedException ie) {
+                                    log.error("[Principal={}]: Interrupted while renewing TGT, exiting Login thread", principal);
                                     return;
                                 }
+                            } else {
+                                log.warn("[Principal={}]: Could not renew TGT due to problem running shell command: '{} {}'. " +
+                                    "Exiting refresh thread.", principal, kinitCmd, kinitArgs, e);
+                                return;
                             }
                         }
                     }
-                    try {
-                        int retry = 1;
-                        while (retry >= 0) {
-                            try {
-                                reLogin();
-                                break;
-                            } catch (LoginException le) {
-                                if (retry > 0) {
-                                    --retry;
-                                    // sleep for 10 seconds.
-                                    try {
-                                        Thread.sleep(10 * 1000);
-                                    } catch (InterruptedException e) {
-                                        log.error("[Principal={}]: Interrupted during login retry after LoginException:", principal, le);
-                                        throw le;
-                                    }
-                                } else {
-                                    log.error("[Principal={}]: Could not refresh TGT.", principal, le);
+                }
+                try {
+                    int retry = 1;
+                    while (retry >= 0) {
+                        try {
+                            reLogin();
+                            break;
+                        } catch (LoginException le) {
+                            if (retry > 0) {
+                                log.warn("[Principal={}]: Error when trying to re-Login, but will retry ", principal, le);
+                                --retry;
+                                // sleep for 10 seconds.
+                                try {
+                                    Thread.sleep(10 * 1000);
+                                } catch (InterruptedException e) {
+                                    log.error("[Principal={}]: Interrupted during login retry after LoginException:", principal, le);
+                                    throw le;
                                 }
+                            } else {
+                                log.error("[Principal={}]: Could not refresh TGT.", principal, le);
                             }
                         }
-                    } catch (LoginException le) {
-                        log.error("[Principal={}]: Failed to refresh TGT: refresh thread exiting now.", principal, le);
-                        return;
                     }
+                } catch (LoginException le) {
+                    log.error("[Principal={}]: Failed to refresh TGT: refresh thread exiting now.", principal, le);
+                    return;
                 }
             }
         });
@@ -281,8 +285,9 @@ public String serviceName() {
         return serviceName;
     }
 
-    private static String getServiceName(Map configs, JaasContext jaasContext) {
-        String jaasServiceName = jaasContext.configEntryOption(JaasUtils.SERVICE_NAME, null);
+    private static String getServiceName(Map configs, String contextName, Configuration configuration) {
+        List configEntries = Arrays.asList(configuration.getAppConfigurationEntry(contextName));
+        String jaasServiceName = JaasContext.configEntryOption(configEntries, JaasUtils.SERVICE_NAME, null);
         String configServiceName = (String) configs.get(SaslConfigs.SASL_KERBEROS_SERVICE_NAME);
         if (jaasServiceName != null && configServiceName != null && !jaasServiceName.equals(configServiceName)) {
             String message = String.format("Conflicting serviceName values found in JAAS and Kafka configs " +
@@ -314,7 +319,7 @@ private long getRefreshTime(KerberosTicket tgt) {
             return proposedRefresh;
     }
 
-    private synchronized KerberosTicket getTGT() {
+    private KerberosTicket getTGT() {
         Set tickets = subject.getPrivateCredentials(KerberosTicket.class);
         for (KerberosTicket ticket : tickets) {
             KerberosPrincipal server = ticket.getServer();
@@ -341,7 +346,7 @@ private boolean hasSufficientTimeElapsed() {
      * Re-login a principal. This method assumes that {@link #login()} has happened already.
      * @throws javax.security.auth.login.LoginException on a failure
      */
-    private synchronized void reLogin() throws LoginException {
+    protected void reLogin() throws LoginException {
         if (!isKrbTicket) {
             return;
         }
@@ -358,15 +363,20 @@ private synchronized void reLogin() throws LoginException {
             //clear up the kerberos state. But the tokens are not cleared! As per
             //the Java kerberos login module code, only the kerberos credentials
             //are cleared
-            loginContext.logout();
+            logout();
             //login and also update the subject field of this instance to
             //have the new credentials (pass it to the LoginContext constructor)
-            loginContext = new LoginContext(jaasContext().name(), subject, null, jaasContext().configuration());
+            loginContext = new LoginContext(contextName(), subject, null, configuration());
             log.info("Initiating re-login for {}", principal);
             loginContext.login();
         }
     }
 
+    // Visibility to override for testing
+    protected void logout() throws LoginException {
+        loginContext.logout();
+    }
+
     private long currentElapsedTime() {
         return time.hiResClockMs();
     }
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 40a9a1fad6e54..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
@@ -17,6 +17,7 @@
 package org.apache.kafka.common.security.kerberos;
 
 import java.io.IOException;
+import java.util.Locale;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
@@ -44,6 +45,8 @@ class KerberosRule {
     private final Pattern fromPattern;
     private final String toPattern;
     private final boolean repeat;
+    private final boolean toLowerCase;
+    private final boolean toUpperCase;
 
     KerberosRule(String defaultRealm) {
         this.defaultRealm = defaultRealm;
@@ -54,10 +57,12 @@ class KerberosRule {
         fromPattern = null;
         toPattern = null;
         repeat = false;
+        toLowerCase = false;
+        toUpperCase = false;
     }
 
     KerberosRule(String defaultRealm, int numOfComponents, String format, String match, String fromPattern,
-                 String toPattern, boolean repeat) {
+                 String toPattern, boolean repeat, boolean toLowerCase, boolean toUpperCase) {
         this.defaultRealm = defaultRealm;
         isDefault = false;
         this.numOfComponents = numOfComponents;
@@ -67,6 +72,8 @@ class KerberosRule {
                 fromPattern == null ? null : Pattern.compile(fromPattern);
         this.toPattern = toPattern;
         this.repeat = repeat;
+        this.toLowerCase = toLowerCase;
+        this.toUpperCase = toUpperCase;
     }
 
     @Override
@@ -95,13 +102,19 @@ public String toString() {
                     buf.append('g');
                 }
             }
+            if (toLowerCase) {
+                buf.append("/L");
+            }
+            if (toUpperCase) {
+                buf.append("/U");
+            }
         }
         return buf.toString();
     }
 
     /**
-     * Replace the numbered parameters of the form $n where n is from 1 to
-     * the length of params. Normal text is copied directly and $n is replaced
+     * Replace the numbered parameters of the form $n where n is from 0 to
+     * the length of params - 1. Normal text is copied directly and $n is replaced
      * by the corresponding parameter.
      * @param format the string to replace parameters again
      * @param params the list of parameters
@@ -119,7 +132,7 @@ static String replaceParameters(String format,
             if (paramNum != null) {
                 try {
                     int num = Integer.parseInt(paramNum);
-                    if (num < 0 || num > params.length) {
+                    if (num < 0 || num >= params.length) {
                         throw new BadFormatString("index " + num + " from " + format +
                                 " is outside of the valid range 0 to " +
                                 (params.length - 1));
@@ -182,6 +195,12 @@ String apply(String[] params) throws IOException {
         if (result != null && NON_SIMPLE_PATTERN.matcher(result).find()) {
             throw new NoMatchingRule("Non-simple name " + result + " after auth_to_local rule " + this);
         }
+        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 1db984a6cb744..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)?)?))");
+    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;
@@ -60,12 +60,14 @@ private static List parseRules(String defaultRealm, List r
                 result.add(new KerberosRule(defaultRealm));
             } else {
                 result.add(new KerberosRule(defaultRealm,
-                        Integer.parseInt(matcher.group(4)),
-                        matcher.group(5),
-                        matcher.group(7),
-                        matcher.group(9),
+                        Integer.parseInt(matcher.group(5)),
+                        matcher.group(6),
+                        matcher.group(8),
                         matcher.group(10),
-                        "g".equals(matcher.group(11))));
+                        matcher.group(11),
+                        "g".equals(matcher.group(12)),
+                        "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
new file mode 100644
index 0000000000000..eab208baa6f2b
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerExtensionsValidatorCallback.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.security.oauthbearer;
+
+import org.apache.kafka.common.security.auth.SaslExtensions;
+
+import javax.security.auth.callback.Callback;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Objects;
+
+import static org.apache.kafka.common.utils.CollectionUtils.subtractMap;
+
+/**
+ * A {@code Callback} for use by the {@code SaslServer} implementation when it
+ * needs to validate the SASL extensions for the OAUTHBEARER mechanism
+ * Callback handlers should use the {@link #valid(String)}
+ * method to communicate valid extensions back to the SASL server.
+ * Callback handlers should use the
+ * {@link #error(String, String)} method to communicate validation errors back to
+ * the SASL Server.
+ * As per RFC-7628 (https://tools.ietf.org/html/rfc7628#section-3.1), unknown extensions must be ignored by the server.
+ * The callback handler implementation should simply ignore unknown extensions,
+ * not calling {@link #error(String, String)} nor {@link #valid(String)}.
+ * Callback handlers should communicate other problems by raising an {@code IOException}.
+ * 
        
+ * The OAuth bearer token is provided in the callback for better context in extension validation.
+ * It is very important that token validation is done in its own {@link OAuthBearerValidatorCallback}
+ * irregardless of provided extensions, as they are inherently insecure.
+ */
+public class OAuthBearerExtensionsValidatorCallback implements Callback {
+    private final OAuthBearerToken token;
+    private final SaslExtensions inputExtensions;
+    private final Map validatedExtensions = new HashMap<>();
+    private final Map invalidExtensions = new HashMap<>();
+
+    public OAuthBearerExtensionsValidatorCallback(OAuthBearerToken token, SaslExtensions extensions) {
+        this.token = Objects.requireNonNull(token);
+        this.inputExtensions = Objects.requireNonNull(extensions);
+    }
+
+    /**
+     * @return {@link OAuthBearerToken} the OAuth bearer token of the client
+     */
+    public OAuthBearerToken token() {
+        return token;
+    }
+
+    /**
+     * @return {@link SaslExtensions} consisting of the unvalidated extension names and values that were sent by the client
+     */
+    public SaslExtensions inputExtensions() {
+        return inputExtensions;
+    }
+
+    /**
+     * @return an unmodifiable {@link Map} consisting of the validated and recognized by the server extension names and values
+     */
+    public Map validatedExtensions() {
+        return Collections.unmodifiableMap(validatedExtensions);
+    }
+
+    /**
+     * @return An immutable {@link Map} consisting of the name->error messages of extensions which failed validation
+     */
+    public Map invalidExtensions() {
+        return Collections.unmodifiableMap(invalidExtensions);
+    }
+
+    /**
+     * @return An immutable {@link Map} consisting of the extensions that have neither been validated nor invalidated
+     */
+    public Map ignoredExtensions() {
+        return Collections.unmodifiableMap(subtractMap(subtractMap(inputExtensions.map(), invalidExtensions), validatedExtensions));
+    }
+
+    /**
+     * Validates a specific extension in the original {@code inputExtensions} map
+     * @param extensionName - the name of the extension which was validated
+     */
+    public void valid(String extensionName) {
+        if (!inputExtensions.map().containsKey(extensionName))
+            throw new IllegalArgumentException(String.format("Extension %s was not found in the original extensions", extensionName));
+        validatedExtensions.put(extensionName, inputExtensions.map().get(extensionName));
+    }
+    /**
+     * Set the error value for a specific extension key-value pair if validation has failed
+     *
+     * @param invalidExtensionName
+     *            the mandatory extension name which caused the validation failure
+     * @param errorMessage
+     *            error message describing why the validation failed
+     */
+    public void error(String invalidExtensionName, String errorMessage) {
+        if (Objects.requireNonNull(invalidExtensionName).isEmpty())
+            throw new IllegalArgumentException("extension name must not be empty");
+        this.invalidExtensions.put(invalidExtensionName, errorMessage);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerLoginModule.java b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerLoginModule.java
new file mode 100644
index 0000000000000..e7976b5550605
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerLoginModule.java
@@ -0,0 +1,428 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.oauthbearer;
+
+import java.io.IOException;
+import java.util.Collections;
+import java.util.Iterator;
+import java.util.Map;
+import java.util.Objects;
+
+import javax.security.auth.Subject;
+import javax.security.auth.callback.Callback;
+import javax.security.auth.callback.CallbackHandler;
+import javax.security.auth.callback.UnsupportedCallbackException;
+import javax.security.auth.login.LoginException;
+import javax.security.auth.spi.LoginModule;
+
+import org.apache.kafka.common.config.SaslConfigs;
+import org.apache.kafka.common.security.auth.AuthenticateCallbackHandler;
+import org.apache.kafka.common.security.auth.Login;
+import org.apache.kafka.common.security.auth.SaslExtensionsCallback;
+import org.apache.kafka.common.security.auth.SaslExtensions;
+import org.apache.kafka.common.security.oauthbearer.internals.OAuthBearerSaslClientProvider;
+import org.apache.kafka.common.security.oauthbearer.internals.OAuthBearerSaslServerProvider;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * The {@code LoginModule} for the SASL/OAUTHBEARER mechanism. When a client
+ * (whether a non-broker client or a broker when SASL/OAUTHBEARER is the
+ * inter-broker protocol) connects to Kafka the {@code OAuthBearerLoginModule}
+ * instance asks its configured {@link AuthenticateCallbackHandler}
+ * implementation to handle an instance of {@link OAuthBearerTokenCallback} and
+ * return an instance of {@link OAuthBearerToken}. A default, builtin
+ * {@link AuthenticateCallbackHandler} implementation creates an unsecured token
+ * as defined by these JAAS module options:
+ * 
        
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
        JAAS Module Option for Unsecured Token RetrievalDocumentation
        {@code unsecuredLoginStringClaim_="value"}Creates a {@code String} claim with the given name and value. Any valid
+ * claim name can be specified except '{@code iat}' and '{@code exp}' (these are
+ * automatically generated).
        {@code unsecuredLoginNumberClaim_="value"}Creates a {@code Number} claim with the given name and value. Any valid
+ * claim name can be specified except '{@code iat}' and '{@code exp}' (these are
+ * automatically generated).
        {@code unsecuredLoginListClaim_="value"}Creates a {@code String List} claim with the given name and values parsed
+ * from the given value where the first character is taken as the delimiter. For
+ * example: {@code unsecuredLoginListClaim_fubar="|value1|value2"}. Any valid
+ * claim name can be specified except '{@code iat}' and '{@code exp}' (these are
+ * automatically generated).
        {@code unsecuredLoginPrincipalClaimName}Set to a custom claim name if you wish the name of the {@code String}
+ * claim holding the principal name to be something other than
+ * '{@code sub}'.
        {@code unsecuredLoginLifetimeSeconds}Set to an integer value if the token expiration is to be set to something
+ * other than the default value of 3600 seconds (which is 1 hour). The
+ * '{@code exp}' claim will be set to reflect the expiration time.
        {@code unsecuredLoginScopeClaimName}Set to a custom claim name if you wish the name of the {@code String} or
+ * {@code String List} claim holding any token scope to be something other than
+ * '{@code scope}'.
        
+ * 
        
+ * 
        
+ * You can also add custom unsecured SASL extensions when using the default, builtin {@link AuthenticateCallbackHandler}
+ * implementation through using the configurable option {@code unsecuredLoginExtension_}. Note that there
+ * are validations for the key/values in order to conform to the SASL/OAUTHBEARER standard
+ * (https://tools.ietf.org/html/rfc7628#section-3.1), including the reserved key at
+ * {@link org.apache.kafka.common.security.oauthbearer.internals.OAuthBearerClientInitialResponse#AUTH_KEY}.
+ * The {@code OAuthBearerLoginModule} instance also asks its configured {@link AuthenticateCallbackHandler}
+ * implementation to handle an instance of {@link SaslExtensionsCallback} and return an instance of {@link SaslExtensions}.
+ * The configured callback handler does not need to handle this callback, though -- any {@code UnsupportedCallbackException}
+ * that is thrown is ignored, and no SASL extensions will be associated with the login.
+ * 
        
+ * Production use cases will require writing an implementation of
+ * {@link AuthenticateCallbackHandler} that can handle an instance of
+ * {@link OAuthBearerTokenCallback} and declaring it via either the
+ * {@code sasl.login.callback.handler.class} configuration option for a
+ * non-broker client or via the
+ * {@code listener.name.sasl_ssl.oauthbearer.sasl.login.callback.handler.class}
+ * configuration option for brokers (when SASL/OAUTHBEARER is the inter-broker
+ * protocol).
+ * 
        
+ * This class stores the retrieved {@link OAuthBearerToken} in the
+ * {@code Subject}'s private credentials where the {@code SaslClient} can
+ * retrieve it. An appropriate, builtin {@code SaslClient} implementation is
+ * automatically used and configured such that it can perform that retrieval.
+ * 
        
+ * Here is a typical, basic JAAS configuration for a client leveraging unsecured
+ * SASL/OAUTHBEARER authentication:
+ *
+ * 
        + * KafkaClient {
+ *      org.apache.kafka.common.security.oauthbearer.OAuthBearerLoginModule Required
+ *      unsecuredLoginStringClaim_sub="thePrincipalName";
+ * };
+ * 
        
+ *
+ * An implementation of the {@link Login} interface specific to the
+ * {@code OAUTHBEARER} mechanism is automatically applied; it periodically
+ * refreshes any token before it expires so that the client can continue to make
+ * connections to brokers. The parameters that impact how the refresh algorithm
+ * operates are specified as part of the producer/consumer/broker configuration
+ * and are as follows. See the documentation for these properties elsewhere for
+ * details.
+ * 
        
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
        Producer/Consumer/Broker Configuration Property
        {@code sasl.login.refresh.window.factor}
        {@code sasl.login.refresh.window.jitter}
        {@code sasl.login.refresh.min.period.seconds}
        {@code sasl.login.refresh.min.buffer.seconds}
        
+ * 
        
+ * When a broker accepts a SASL/OAUTHBEARER connection the instance of the
+ * builtin {@code SaslServer} implementation asks its configured
+ * {@link AuthenticateCallbackHandler} implementation to handle an instance of
+ * {@link OAuthBearerValidatorCallback} constructed with the OAuth 2 Bearer
+ * Token's compact serialization and return an instance of
+ * {@link OAuthBearerToken} if the value validates. A default, builtin
+ * {@link AuthenticateCallbackHandler} implementation validates an unsecured
+ * token as defined by these JAAS module options:
+ * 
        
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
        JAAS Module Option for Unsecured Token ValidationDocumentation
        {@code unsecuredValidatorPrincipalClaimName="value"}Set to a non-empty value if you wish a particular {@code String} claim
+ * holding a principal name to be checked for existence; the default is to check
+ * for the existence of the '{@code sub}' claim.
        {@code unsecuredValidatorScopeClaimName="value"}Set to a custom claim name if you wish the name of the {@code String} or
+ * {@code String List} claim holding any token scope to be something other than
+ * '{@code scope}'.
        {@code unsecuredValidatorRequiredScope="value"}Set to a space-delimited list of scope values if you wish the
+ * {@code String/String List} claim holding the token scope to be checked to
+ * make sure it contains certain values.
        {@code unsecuredValidatorAllowableClockSkewMs="value"}Set to a positive integer value if you wish to allow up to some number of
+ * positive milliseconds of clock skew (the default is 0).
        
+ * 
        
+ * Here is a typical, basic JAAS configuration for a broker leveraging unsecured
+ * SASL/OAUTHBEARER validation:
+ *
+ * 
        + * KafkaServer {
+ *      org.apache.kafka.common.security.oauthbearer.OAuthBearerLoginModule Required
+ *      unsecuredLoginStringClaim_sub="thePrincipalName";
+ * };
+ * 
        
+ *
+ * Production use cases will require writing an implementation of
+ * {@link AuthenticateCallbackHandler} that can handle an instance of
+ * {@link OAuthBearerValidatorCallback} and declaring it via the
+ * {@code listener.name.sasl_ssl.oauthbearer.sasl.server.callback.handler.class}
+ * broker configuration option.
+ * 
        
+ * The builtin {@code SaslServer} implementation for SASL/OAUTHBEARER in Kafka
+ * makes the instance of {@link OAuthBearerToken} available upon successful
+ * authentication via the negotiated property "{@code OAUTHBEARER.token}"; the
+ * token could be used in a custom authorizer (to authorize based on JWT claims
+ * rather than ACLs, for example).
+ * 
        
+ * This implementation's {@code logout()} method will logout the specific token
+ * that this instance logged in if it's {@code Subject} instance is shared
+ * across multiple {@code LoginContext}s and there happen to be multiple tokens
+ * on the {@code Subject}. This functionality is useful because it means a new
+ * token with a longer lifetime can be created before a soon-to-expire token is
+ * actually logged out. Otherwise, if multiple simultaneous tokens were not
+ * supported like this, the soon-to-be expired token would have to be logged out
+ * first, and then if the new token could not be retrieved (maybe the
+ * authorization server is temporarily unavailable, for example) the client
+ * would be left without a token and would be unable to create new connections.
+ * Better to mitigate this possibility by leaving the existing token (which
+ * still has some lifetime left) in place until a new replacement token is
+ * actually retrieved. This implementation supports this.
+ *
+ * @see SaslConfigs#SASL_LOGIN_REFRESH_WINDOW_FACTOR_DOC
+ * @see SaslConfigs#SASL_LOGIN_REFRESH_WINDOW_JITTER_DOC
+ * @see SaslConfigs#SASL_LOGIN_REFRESH_MIN_PERIOD_SECONDS_DOC
+ * @see SaslConfigs#SASL_LOGIN_REFRESH_BUFFER_SECONDS_DOC
+ */
+public class OAuthBearerLoginModule implements LoginModule {
+
+    /**
+     * Login state transitions:
+     *   Initial state: NOT_LOGGED_IN
+     *   login()      : NOT_LOGGED_IN => LOGGED_IN_NOT_COMMITTED
+     *   commit()     : LOGGED_IN_NOT_COMMITTED => COMMITTED
+     *   abort()      : LOGGED_IN_NOT_COMMITTED => NOT_LOGGED_IN
+     *   logout()     : Any state => NOT_LOGGED_IN
+     */
+    private enum LoginState {
+        NOT_LOGGED_IN,
+        LOGGED_IN_NOT_COMMITTED,
+        COMMITTED
+    }
+
+    /**
+     * The SASL Mechanism name for OAuth 2: {@code OAUTHBEARER}
+     */
+    public static final String OAUTHBEARER_MECHANISM = "OAUTHBEARER";
+    private static final Logger log = LoggerFactory.getLogger(OAuthBearerLoginModule.class);
+    private static final SaslExtensions EMPTY_EXTENSIONS = new SaslExtensions(Collections.emptyMap());
+    private Subject subject = null;
+    private AuthenticateCallbackHandler callbackHandler = null;
+    private OAuthBearerToken tokenRequiringCommit = null;
+    private OAuthBearerToken myCommittedToken = null;
+    private SaslExtensions extensionsRequiringCommit = null;
+    private SaslExtensions myCommittedExtensions = null;
+    private LoginState loginState;
+
+    static {
+        OAuthBearerSaslClientProvider.initialize(); // not part of public API
+        OAuthBearerSaslServerProvider.initialize(); // not part of public API
+    }
+
+    @Override
+    public void initialize(Subject subject, CallbackHandler callbackHandler, Map sharedState,
+            Map options) {
+        this.subject = Objects.requireNonNull(subject);
+        if (!(Objects.requireNonNull(callbackHandler) instanceof AuthenticateCallbackHandler))
+            throw new IllegalArgumentException(String.format("Callback handler must be castable to %s: %s",
+                    AuthenticateCallbackHandler.class.getName(), callbackHandler.getClass().getName()));
+        this.callbackHandler = (AuthenticateCallbackHandler) callbackHandler;
+    }
+
+    @Override
+    public boolean login() throws LoginException {
+        if (loginState == LoginState.LOGGED_IN_NOT_COMMITTED) {
+            if (tokenRequiringCommit != null)
+                throw new IllegalStateException(String.format(
+                    "Already have an uncommitted token with private credential token count=%d", committedTokenCount()));
+            else
+                throw new IllegalStateException("Already logged in without a token");
+        }
+        if (loginState == LoginState.COMMITTED) {
+            if (myCommittedToken != null)
+                throw new IllegalStateException(String.format(
+                    "Already have a committed token with private credential token count=%d; must login on another login context or logout here first before reusing the same login context",
+                    committedTokenCount()));
+            else
+                throw new IllegalStateException("Login has already been committed without a token");
+        }
+
+        identifyToken();
+        if (tokenRequiringCommit != null)
+            identifyExtensions();
+        else
+            log.debug("Logged in without a token, this login cannot be used to establish client connections");
+
+        loginState = LoginState.LOGGED_IN_NOT_COMMITTED;
+        log.debug("Login succeeded; invoke commit() to commit it; current committed token count={}",
+                committedTokenCount());
+        return true;
+    }
+
+    private void identifyToken() throws LoginException {
+        OAuthBearerTokenCallback tokenCallback = new OAuthBearerTokenCallback();
+        try {
+            callbackHandler.handle(new Callback[] {tokenCallback});
+        } catch (IOException | UnsupportedCallbackException e) {
+            log.error(e.getMessage(), e);
+            throw new LoginException("An internal error occurred while retrieving token from callback handler");
+        }
+
+        tokenRequiringCommit = tokenCallback.token();
+        if (tokenCallback.errorCode() != null) {
+            log.info("Login failed: {} : {} (URI={})", tokenCallback.errorCode(), tokenCallback.errorDescription(),
+                    tokenCallback.errorUri());
+            throw new LoginException(tokenCallback.errorDescription());
+        }
+    }
+
+    /**
+     * Attaches SASL extensions to the Subject
+     */
+    private void identifyExtensions() throws LoginException {
+        SaslExtensionsCallback extensionsCallback = new SaslExtensionsCallback();
+        try {
+            callbackHandler.handle(new Callback[] {extensionsCallback});
+            extensionsRequiringCommit = extensionsCallback.extensions();
+        } catch (IOException e) {
+            log.error(e.getMessage(), e);
+            throw new LoginException("An internal error occurred while retrieving SASL extensions from callback handler");
+        } catch (UnsupportedCallbackException e) {
+            extensionsRequiringCommit = EMPTY_EXTENSIONS;
+            log.debug("CallbackHandler {} does not support SASL extensions. No extensions will be added", callbackHandler.getClass().getName());
+        }
+        if (extensionsRequiringCommit ==  null) {
+            log.error("SASL Extensions cannot be null. Check whether your callback handler is explicitly setting them as null.");
+            throw new LoginException("Extensions cannot be null.");
+        }
+    }
+
+    @Override
+    public boolean logout() {
+        if (loginState == LoginState.LOGGED_IN_NOT_COMMITTED)
+            throw new IllegalStateException(
+                    "Cannot call logout() immediately after login(); need to first invoke commit() or abort()");
+        if (loginState != LoginState.COMMITTED) {
+            log.debug("Nothing here to log out");
+            return false;
+        }
+        if (myCommittedToken != null) {
+            log.trace("Logging out my token; current committed token count = {}", committedTokenCount());
+            for (Iterator iterator = subject.getPrivateCredentials().iterator(); iterator.hasNext(); ) {
+                Object privateCredential = iterator.next();
+                if (privateCredential == myCommittedToken) {
+                    iterator.remove();
+                    myCommittedToken = null;
+                    break;
+                }
+            }
+            log.debug("Done logging out my token; committed token count is now {}", committedTokenCount());
+        } else
+            log.debug("No tokens to logout for this login");
+
+        if (myCommittedExtensions != null) {
+            log.trace("Logging out my extensions");
+            if (subject.getPublicCredentials().removeIf(e -> myCommittedExtensions == e))
+                myCommittedExtensions = null;
+            log.debug("Done logging out my extensions");
+        } else
+            log.debug("No extensions to logout for this login");
+
+        loginState = LoginState.NOT_LOGGED_IN;
+        return true;
+    }
+
+    @Override
+    public boolean commit() {
+        if (loginState != LoginState.LOGGED_IN_NOT_COMMITTED) {
+            log.debug("Nothing here to commit");
+            return false;
+        }
+
+        if (tokenRequiringCommit != null) {
+            log.trace("Committing my token; current committed token count = {}", committedTokenCount());
+            subject.getPrivateCredentials().add(tokenRequiringCommit);
+            myCommittedToken = tokenRequiringCommit;
+            tokenRequiringCommit = null;
+            log.debug("Done committing my token; committed token count is now {}", committedTokenCount());
+        } else
+            log.debug("No tokens to commit, this login cannot be used to establish client connections");
+
+        if (extensionsRequiringCommit != null) {
+            subject.getPublicCredentials().add(extensionsRequiringCommit);
+            myCommittedExtensions = extensionsRequiringCommit;
+            extensionsRequiringCommit = null;
+        }
+
+        loginState = LoginState.COMMITTED;
+        return true;
+    }
+
+    @Override
+    public boolean abort() {
+        if (loginState == LoginState.LOGGED_IN_NOT_COMMITTED) {
+            log.debug("Login aborted");
+            tokenRequiringCommit = null;
+            extensionsRequiringCommit = null;
+            loginState = LoginState.NOT_LOGGED_IN;
+            return true;
+        }
+        log.debug("Nothing here to abort");
+        return false;
+    }
+
+    private int committedTokenCount() {
+        return subject.getPrivateCredentials(OAuthBearerToken.class).size();
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerToken.java b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerToken.java
new file mode 100644
index 0000000000000..ee443ede7a2f2
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerToken.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.security.oauthbearer;
+
+import java.util.Set;
+
+import org.apache.kafka.common.annotation.InterfaceStability;
+
+/**
+ * The b64token value as defined in
+ * RFC 6750 Section
+ * 2.1 along with the token's specific scope and lifetime and principal
+ * name.
+ * 
        
+ * A network request would be required to re-hydrate an opaque token, and that
+ * could result in (for example) an {@code IOException}, but retrievers for
+ * various attributes ({@link #scope()}, {@link #lifetimeMs()}, etc.) declare no
+ * exceptions. Therefore, if a network request is required for any of these
+ * retriever methods, that request could be performed at construction time so
+ * that the various attributes can be reliably provided thereafter. For example,
+ * a constructor might declare {@code throws IOException} in such a case.
+ * Alternatively, the retrievers could throw unchecked exceptions.
+ * 
        
+ * This interface was introduced in 2.0.0 and, while it feels stable, it could
+ * evolve. We will try to evolve the API in a compatible manner (easier now that
+ * Java 7 and its lack of default methods doesn't have to be supported), 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.
+ * 
+ * @see RFC 6749
+ *      Section 1.4 and
+ *      RFC 6750
+ *      Section 2.1
+ */
+@InterfaceStability.Evolving
+public interface OAuthBearerToken {
+    /**
+     * The b64token value as defined in
+     * RFC 6750 Section
+     * 2.1
+     * 
+     * @return b64token value as defined in
+     *         RFC 6750
+     *         Section 2.1
+     */
+    String value();
+
+    /**
+     * The token's scope of access, as per
+     * RFC 6749 Section
+     * 1.4
+     * 
+     * @return the token's (always non-null but potentially empty) scope of access,
+     *         as per RFC
+     *         6749 Section 1.4. Note that all values in the returned set will
+     *         be trimmed of preceding and trailing whitespace, and the result will
+     *         never contain the empty string.
+     */
+    Set scope();
+
+    /**
+     * The token's lifetime, expressed as the number of milliseconds since the
+     * epoch, as per RFC
+     * 6749 Section 1.4
+     * 
+     * @return the token'slifetime, expressed as the number of milliseconds since
+     *         the epoch, as per
+     *         RFC 6749
+     *         Section 1.4.
+     */
+    long lifetimeMs();
+
+    /**
+     * The name of the principal to which this credential applies
+     * 
+     * @return the always non-null/non-empty principal name
+     */
+    String principalName();
+
+    /**
+     * When the credential became valid, in terms of the number of milliseconds
+     * since the epoch, if known, otherwise null. An expiring credential may not
+     * necessarily indicate when it was created -- just when it expires -- so we
+     * need to support a null return value here.
+     * 
+     * @return the time when the credential became valid, in terms of the number of
+     *         milliseconds since the epoch, if known, otherwise null
+     */
+    Long startTimeMs();
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerTokenCallback.java b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerTokenCallback.java
new file mode 100644
index 0000000000000..3f4f269606ab2
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerTokenCallback.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.common.security.oauthbearer;
+
+import java.util.Objects;
+
+import javax.security.auth.callback.Callback;
+
+import org.apache.kafka.common.annotation.InterfaceStability;
+
+/**
+ * A {@code Callback} for use by the {@code SaslClient} and {@code Login}
+ * implementations when they require an OAuth 2 bearer token. Callback handlers
+ * should use the {@link #error(String, String, String)} method to communicate
+ * errors returned by the authorization server as per
+ * RFC 6749: The OAuth
+ * 2.0 Authorization Framework. Callback handlers should communicate other
+ * problems by raising an {@code IOException}.
+ * 
        
+ * This class was introduced in 2.0.0 and, while it feels stable, it could
+ * evolve. 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 class OAuthBearerTokenCallback implements Callback {
+    private OAuthBearerToken token = null;
+    private String errorCode = null;
+    private String errorDescription = null;
+    private String errorUri = null;
+
+    /**
+     * Return the (potentially null) token
+     * 
+     * @return the (potentially null) token
+     */
+    public OAuthBearerToken token() {
+        return token;
+    }
+
+    /**
+     * Return the optional (but always non-empty if not null) error code as per
+     * RFC 6749: The OAuth
+     * 2.0 Authorization Framework.
+     * 
+     * @return the optional (but always non-empty if not null) error code
+     */
+    public String errorCode() {
+        return errorCode;
+    }
+
+    /**
+     * Return the (potentially null) error description as per
+     * RFC 6749: The OAuth
+     * 2.0 Authorization Framework.
+     * 
+     * @return the (potentially null) error description
+     */
+    public String errorDescription() {
+        return errorDescription;
+    }
+
+    /**
+     * Return the (potentially null) error URI as per
+     * RFC 6749: The OAuth
+     * 2.0 Authorization Framework.
+     * 
+     * @return the (potentially null) error URI
+     */
+    public String errorUri() {
+        return errorUri;
+    }
+
+    /**
+     * Set the token. All error-related values are cleared.
+     * 
+     * @param token
+     *            the optional token to set
+     */
+    public void token(OAuthBearerToken token) {
+        this.token = token;
+        this.errorCode = null;
+        this.errorDescription = null;
+        this.errorUri = null;
+    }
+
+    /**
+     * Set the error values as per
+     * RFC 6749: The OAuth
+     * 2.0 Authorization Framework. Any token is cleared.
+     * 
+     * @param errorCode
+     *            the mandatory error code to set
+     * @param errorDescription
+     *            the optional error description to set
+     * @param errorUri
+     *            the optional error URI to set
+     */
+    public void error(String errorCode, String errorDescription, String errorUri) {
+        if (Objects.requireNonNull(errorCode).isEmpty())
+            throw new IllegalArgumentException("error code must not be empty");
+        this.errorCode = errorCode;
+        this.errorDescription = errorDescription;
+        this.errorUri = errorUri;
+        this.token = null;
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerValidatorCallback.java b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerValidatorCallback.java
new file mode 100644
index 0000000000000..36bcf089ac322
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerValidatorCallback.java
@@ -0,0 +1,154 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.oauthbearer;
+
+import java.util.Objects;
+
+import javax.security.auth.callback.Callback;
+
+import org.apache.kafka.common.annotation.InterfaceStability;
+
+/**
+ * A {@code Callback} for use by the {@code SaslServer} implementation when it
+ * needs to provide an OAuth 2 bearer token compact serialization for
+ * validation. Callback handlers should use the
+ * {@link #error(String, String, String)} method to communicate errors back to
+ * the SASL Client as per
+ * RFC 6749: The OAuth
+ * 2.0 Authorization Framework and the IANA
+ * OAuth Extensions Error Registry. Callback handlers should communicate
+ * other problems by raising an {@code IOException}.
+ * 
        
+ * This class was introduced in 2.0.0 and, while it feels stable, it could
+ * evolve. 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 class OAuthBearerValidatorCallback implements Callback {
+    private final String tokenValue;
+    private OAuthBearerToken token = null;
+    private String errorStatus = null;
+    private String errorScope = null;
+    private String errorOpenIDConfiguration = null;
+
+    /**
+     * Constructor
+     * 
+     * @param tokenValue
+     *            the mandatory/non-blank token value
+     */
+    public OAuthBearerValidatorCallback(String tokenValue) {
+        if (Objects.requireNonNull(tokenValue).isEmpty())
+            throw new IllegalArgumentException("token value must not be empty");
+        this.tokenValue = tokenValue;
+    }
+
+    /**
+     * Return the (always non-null) token value
+     * 
+     * @return the (always non-null) token value
+     */
+    public String tokenValue() {
+        return tokenValue;
+    }
+
+    /**
+     * Return the (potentially null) token
+     * 
+     * @return the (potentially null) token
+     */
+    public OAuthBearerToken token() {
+        return token;
+    }
+
+    /**
+     * Return the (potentially null) error status value as per
+     * RFC 7628: A Set
+     * of Simple Authentication and Security Layer (SASL) Mechanisms for OAuth
+     * and the IANA
+     * OAuth Extensions Error Registry.
+     * 
+     * @return the (potentially null) error status value
+     */
+    public String errorStatus() {
+        return errorStatus;
+    }
+
+    /**
+     * Return the (potentially null) error scope value as per
+     * RFC 7628: A Set
+     * of Simple Authentication and Security Layer (SASL) Mechanisms for OAuth.
+     * 
+     * @return the (potentially null) error scope value
+     */
+    public String errorScope() {
+        return errorScope;
+    }
+
+    /**
+     * Return the (potentially null) error openid-configuration value as per
+     * RFC 7628: A Set
+     * of Simple Authentication and Security Layer (SASL) Mechanisms for OAuth.
+     * 
+     * @return the (potentially null) error openid-configuration value
+     */
+    public String errorOpenIDConfiguration() {
+        return errorOpenIDConfiguration;
+    }
+
+    /**
+     * Set the token. The token value is unchanged and is expected to match the
+     * provided token's value. All error values are cleared.
+     * 
+     * @param token
+     *            the mandatory token to set
+     */
+    public void token(OAuthBearerToken token) {
+        this.token = Objects.requireNonNull(token);
+        this.errorStatus = null;
+        this.errorScope = null;
+        this.errorOpenIDConfiguration = null;
+    }
+
+    /**
+     * Set the error values as per
+     * RFC 7628: A Set
+     * of Simple Authentication and Security Layer (SASL) Mechanisms for OAuth.
+     * Any token is cleared.
+     * 
+     * @param errorStatus
+     *            the mandatory error status value from the IANA
+     *            OAuth Extensions Error Registry to set
+     * @param errorScope
+     *            the optional error scope value to set
+     * @param errorOpenIDConfiguration
+     *            the optional error openid-configuration value to set
+     */
+    public void error(String errorStatus, String errorScope, String errorOpenIDConfiguration) {
+        if (Objects.requireNonNull(errorStatus).isEmpty())
+            throw new IllegalArgumentException("error status must not be empty");
+        this.errorStatus = errorStatus;
+        this.errorScope = errorScope;
+        this.errorOpenIDConfiguration = errorOpenIDConfiguration;
+        this.token = null;
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/OAuthBearerClientInitialResponse.java b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/OAuthBearerClientInitialResponse.java
new file mode 100644
index 0000000000000..a356f0da3ddb9
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/OAuthBearerClientInitialResponse.java
@@ -0,0 +1,189 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.oauthbearer.internals;
+
+import org.apache.kafka.common.security.auth.SaslExtensions;
+import org.apache.kafka.common.utils.Utils;
+
+import javax.security.sasl.SaslException;
+import java.nio.charset.StandardCharsets;
+import java.util.Map;
+import java.util.Objects;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+
+public class OAuthBearerClientInitialResponse {
+    static final String SEPARATOR = "\u0001";
+
+    private static final String SASLNAME = "(?:[\\x01-\\x7F&&[^=,]]|=2C|=3D)+";
+    private static final String KEY = "[A-Za-z]+";
+    private static final String VALUE = "[\\x21-\\x7E \t\r\n]+";
+
+    private static final String KVPAIRS = String.format("(%s=%s%s)*", KEY, VALUE, SEPARATOR);
+    private static final Pattern AUTH_PATTERN = Pattern.compile("(?[\\w]+)[ ]+(?[-_\\.a-zA-Z0-9]+)");
+    private static final Pattern CLIENT_INITIAL_RESPONSE_PATTERN = Pattern.compile(
+            String.format("n,(a=(?%s))?,%s(?%s)%s", SASLNAME, SEPARATOR, KVPAIRS, SEPARATOR));
+    public static final String AUTH_KEY = "auth";
+
+    private final String tokenValue;
+    private final String authorizationId;
+    private SaslExtensions saslExtensions;
+
+    public static final Pattern EXTENSION_KEY_PATTERN = Pattern.compile(KEY);
+    public static final Pattern EXTENSION_VALUE_PATTERN = Pattern.compile(VALUE);
+
+    public OAuthBearerClientInitialResponse(byte[] response) throws SaslException {
+        String responseMsg = new String(response, StandardCharsets.UTF_8);
+        Matcher matcher = CLIENT_INITIAL_RESPONSE_PATTERN.matcher(responseMsg);
+        if (!matcher.matches())
+            throw new SaslException("Invalid OAUTHBEARER client first message");
+        String authzid = matcher.group("authzid");
+        this.authorizationId = authzid == null ? "" : authzid;
+        String kvPairs = matcher.group("kvpairs");
+        Map properties = Utils.parseMap(kvPairs, "=", SEPARATOR);
+        String auth = properties.get(AUTH_KEY);
+        if (auth == null)
+            throw new SaslException("Invalid OAUTHBEARER client first message: 'auth' not specified");
+        properties.remove(AUTH_KEY);
+        SaslExtensions extensions = new SaslExtensions(properties);
+        validateExtensions(extensions);
+        this.saslExtensions = extensions;
+
+        Matcher authMatcher = AUTH_PATTERN.matcher(auth);
+        if (!authMatcher.matches())
+            throw new SaslException("Invalid OAUTHBEARER client first message: invalid 'auth' format");
+        if (!"bearer".equalsIgnoreCase(authMatcher.group("scheme"))) {
+            String msg = String.format("Invalid scheme in OAUTHBEARER client first message: %s",
+                    matcher.group("scheme"));
+            throw new SaslException(msg);
+        }
+        this.tokenValue = authMatcher.group("token");
+    }
+
+    /**
+     * Constructor
+     * 
+     * @param tokenValue
+     *            the mandatory token value
+     * @param extensions
+     *            the optional extensions
+     * @throws SaslException
+     *             if any extension name or value fails to conform to the required
+     *             regular expression as defined by the specification, or if the
+     *             reserved {@code auth} appears as a key
+     */
+    public OAuthBearerClientInitialResponse(String tokenValue, SaslExtensions extensions) throws SaslException {
+        this(tokenValue, "", extensions);
+    }
+
+    /**
+     * Constructor
+     * 
+     * @param tokenValue
+     *            the mandatory token value
+     * @param authorizationId
+     *            the optional authorization ID
+     * @param extensions
+     *            the optional extensions
+     * @throws SaslException
+     *             if any extension name or value fails to conform to the required
+     *             regular expression as defined by the specification, or if the
+     *             reserved {@code auth} appears as a key
+     */
+    public OAuthBearerClientInitialResponse(String tokenValue, String authorizationId, SaslExtensions extensions) throws SaslException {
+        this.tokenValue = Objects.requireNonNull(tokenValue, "token value must not be null");
+        this.authorizationId = authorizationId == null ? "" : authorizationId;
+        validateExtensions(extensions);
+        this.saslExtensions = extensions != null ? extensions : SaslExtensions.NO_SASL_EXTENSIONS;
+    }
+
+    /**
+     * Return the always non-null extensions
+     * 
+     * @return the always non-null extensions
+     */
+    public SaslExtensions extensions() {
+        return saslExtensions;
+    }
+
+    public byte[] toBytes() {
+        String authzid = authorizationId.isEmpty() ? "" : "a=" + authorizationId;
+        String extensions = extensionsMessage();
+        if (extensions.length() > 0)
+            extensions = SEPARATOR + extensions;
+
+        String message = String.format("n,%s,%sauth=Bearer %s%s%s%s", authzid,
+                SEPARATOR, tokenValue, extensions, SEPARATOR, SEPARATOR);
+
+        return message.getBytes(StandardCharsets.UTF_8);
+    }
+
+    /**
+     * Return the always non-null token value
+     * 
+     * @return the always non-null toklen value
+     */
+    public String tokenValue() {
+        return tokenValue;
+    }
+
+    /**
+     * Return the always non-null authorization ID
+     * 
+     * @return the always non-null authorization ID
+     */
+    public String authorizationId() {
+        return authorizationId;
+    }
+
+    /**
+     * Validates that the given extensions conform to the standard. They should also not contain the reserve key name {@link OAuthBearerClientInitialResponse#AUTH_KEY}
+     *
+     * @param extensions
+     *            optional extensions to validate
+     * @throws SaslException
+     *             if any extension name or value fails to conform to the required
+     *             regular expression as defined by the specification, or if the
+     *             reserved {@code auth} appears as a key
+     *
+     * @see RFC 7628,
+     *  Section 3.1
+     */
+    public static void validateExtensions(SaslExtensions extensions) throws SaslException {
+        if (extensions == null)
+            return;
+        if (extensions.map().containsKey(OAuthBearerClientInitialResponse.AUTH_KEY))
+            throw new SaslException("Extension name " + OAuthBearerClientInitialResponse.AUTH_KEY + " is invalid");
+
+        for (Map.Entry entry : extensions.map().entrySet()) {
+            String extensionName = entry.getKey();
+            String extensionValue = entry.getValue();
+
+            if (!EXTENSION_KEY_PATTERN.matcher(extensionName).matches())
+                throw new SaslException("Extension name " + extensionName + " is invalid");
+            if (!EXTENSION_VALUE_PATTERN.matcher(extensionValue).matches())
+                throw new SaslException("Extension value (" + extensionValue + ") for extension " + extensionName + " is invalid");
+        }
+    }
+
+    /**
+     * Converts the SASLExtensions to an OAuth protocol-friendly string
+     */
+    private String extensionsMessage() {
+        return Utils.mkString(saslExtensions.map(), "", "", "=", SEPARATOR);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/OAuthBearerRefreshingLogin.java b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/OAuthBearerRefreshingLogin.java
new file mode 100644
index 0000000000000..4adbe394d0874
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/OAuthBearerRefreshingLogin.java
@@ -0,0 +1,153 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.oauthbearer.internals;
+
+import java.util.Map;
+import java.util.Set;
+
+import javax.security.auth.Subject;
+import javax.security.auth.login.Configuration;
+import javax.security.auth.login.LoginContext;
+import javax.security.auth.login.LoginException;
+
+import org.apache.kafka.common.config.SaslConfigs;
+import org.apache.kafka.common.security.auth.AuthenticateCallbackHandler;
+import org.apache.kafka.common.security.auth.Login;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerLoginModule;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerToken;
+import org.apache.kafka.common.security.oauthbearer.internals.expiring.ExpiringCredential;
+import org.apache.kafka.common.security.oauthbearer.internals.expiring.ExpiringCredentialRefreshConfig;
+import org.apache.kafka.common.security.oauthbearer.internals.expiring.ExpiringCredentialRefreshingLogin;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * This class is responsible for refreshing logins for both Kafka client and
+ * server when the credential is an OAuth 2 bearer token communicated over
+ * SASL/OAUTHBEARER. An OAuth 2 bearer token has a limited lifetime, and an
+ * instance of this class periodically refreshes it so that the client can
+ * create new connections to brokers on an ongoing basis.
+ * 
        
+ * This class does not need to be explicitly set via the
+ * {@code sasl.login.class} client configuration property or the
+ * {@code listener.name.sasl_[plaintext|ssl].oauthbearer.sasl.login.class}
+ * broker configuration property when the SASL mechanism is OAUTHBEARER; it is
+ * automatically set by default in that case.
+ * 
        
+ * The parameters that impact how the refresh algorithm operates are specified
+ * as part of the producer/consumer/broker configuration and are as follows. See
+ * the documentation for these properties elsewhere for details.
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
+ * 
        Producer/Consumer/Broker Configuration Property
        {@code sasl.login.refresh.window.factor}
        {@code sasl.login.refresh.window.jitter}
        {@code sasl.login.refresh.min.period.seconds}
        {@code sasl.login.refresh.min.buffer.seconds}
        
+ * 
+ * @see OAuthBearerLoginModule
+ * @see SaslConfigs#SASL_LOGIN_REFRESH_WINDOW_FACTOR_DOC
+ * @see SaslConfigs#SASL_LOGIN_REFRESH_WINDOW_JITTER_DOC
+ * @see SaslConfigs#SASL_LOGIN_REFRESH_MIN_PERIOD_SECONDS_DOC
+ * @see SaslConfigs#SASL_LOGIN_REFRESH_BUFFER_SECONDS_DOC
+ */
+public class OAuthBearerRefreshingLogin implements Login {
+    private static final Logger log = LoggerFactory.getLogger(OAuthBearerRefreshingLogin.class);
+    private ExpiringCredentialRefreshingLogin expiringCredentialRefreshingLogin = null;
+
+    @Override
+    public void configure(Map configs, String contextName, Configuration configuration,
+            AuthenticateCallbackHandler loginCallbackHandler) {
+        /*
+         * Specify this class as the one to synchronize on so that only one OAuth 2
+         * Bearer Token is refreshed at a given time. Specify null if we don't mind
+         * multiple simultaneously refreshes. Refreshes happen on the order of minutes
+         * rather than seconds or milliseconds, and there are typically minutes of
+         * lifetime remaining when the refresh occurs, so serializing them seems
+         * reasonable.
+         */
+        Class classToSynchronizeOnPriorToRefresh = OAuthBearerRefreshingLogin.class;
+        expiringCredentialRefreshingLogin = new ExpiringCredentialRefreshingLogin(contextName, configuration,
+                new ExpiringCredentialRefreshConfig(configs, true), loginCallbackHandler,
+                classToSynchronizeOnPriorToRefresh) {
+            @Override
+            public ExpiringCredential expiringCredential() {
+                Set privateCredentialTokens = expiringCredentialRefreshingLogin.subject()
+                        .getPrivateCredentials(OAuthBearerToken.class);
+                if (privateCredentialTokens.isEmpty())
+                    return null;
+                final OAuthBearerToken token = privateCredentialTokens.iterator().next();
+                if (log.isDebugEnabled())
+                    log.debug("Found expiring credential with principal '{}'.", token.principalName());
+                return new ExpiringCredential() {
+                    @Override
+                    public String principalName() {
+                        return token.principalName();
+                    }
+
+                    @Override
+                    public Long startTimeMs() {
+                        return token.startTimeMs();
+                    }
+
+                    @Override
+                    public long expireTimeMs() {
+                        return token.lifetimeMs();
+                    }
+
+                    @Override
+                    public Long absoluteLastRefreshTimeMs() {
+                        return null;
+                    }
+                };
+            }
+        };
+    }
+
+    @Override
+    public void close() {
+        if (expiringCredentialRefreshingLogin != null)
+            expiringCredentialRefreshingLogin.close();
+    }
+
+    @Override
+    public Subject subject() {
+        return expiringCredentialRefreshingLogin != null ? expiringCredentialRefreshingLogin.subject() : null;
+    }
+
+    @Override
+    public String serviceName() {
+        return expiringCredentialRefreshingLogin != null ? expiringCredentialRefreshingLogin.serviceName() : null;
+    }
+
+    @Override
+    public synchronized LoginContext login() throws LoginException {
+        if (expiringCredentialRefreshingLogin != null)
+            return expiringCredentialRefreshingLogin.login();
+        throw new LoginException("Login was not configured properly");
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/OAuthBearerSaslClient.java b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/OAuthBearerSaslClient.java
new file mode 100644
index 0000000000000..e32e16d5a0928
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/OAuthBearerSaslClient.java
@@ -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 org.apache.kafka.common.security.oauthbearer.internals;
+
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
+import java.util.Map;
+import java.util.Objects;
+
+import javax.security.auth.callback.Callback;
+import javax.security.auth.callback.CallbackHandler;
+import javax.security.auth.callback.UnsupportedCallbackException;
+import javax.security.sasl.SaslClient;
+import javax.security.sasl.SaslClientFactory;
+import javax.security.sasl.SaslException;
+
+import org.apache.kafka.common.errors.IllegalSaslStateException;
+import org.apache.kafka.common.security.auth.SaslExtensions;
+import org.apache.kafka.common.security.auth.SaslExtensionsCallback;
+import org.apache.kafka.common.security.auth.AuthenticateCallbackHandler;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerLoginModule;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerToken;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerTokenCallback;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * {@code SaslClient} implementation for SASL/OAUTHBEARER in Kafka. This
+ * implementation requires an instance of {@code AuthenticateCallbackHandler}
+ * that can handle an instance of {@link OAuthBearerTokenCallback} and return
+ * the {@link OAuthBearerToken} generated by the {@code login()} event on the
+ * {@code LoginContext}. Said handler can also optionally handle an instance of {@link SaslExtensionsCallback}
+ * to return any extensions generated by the {@code login()} event on the {@code LoginContext}.
+ *
+ * @see RFC 6750,
+ *      Section 2.1
+ *
+ */
+public class OAuthBearerSaslClient implements SaslClient {
+    static final byte BYTE_CONTROL_A = (byte) 0x01;
+    private static final Logger log = LoggerFactory.getLogger(OAuthBearerSaslClient.class);
+    private final CallbackHandler callbackHandler;
+
+    enum State {
+        SEND_CLIENT_FIRST_MESSAGE, RECEIVE_SERVER_FIRST_MESSAGE, RECEIVE_SERVER_MESSAGE_AFTER_FAILURE, COMPLETE, FAILED
+    }
+
+    private State state;
+
+    public OAuthBearerSaslClient(AuthenticateCallbackHandler callbackHandler) {
+        this.callbackHandler = Objects.requireNonNull(callbackHandler);
+        setState(State.SEND_CLIENT_FIRST_MESSAGE);
+    }
+
+    public CallbackHandler callbackHandler() {
+        return callbackHandler;
+    }
+
+    @Override
+    public String getMechanismName() {
+        return OAuthBearerLoginModule.OAUTHBEARER_MECHANISM;
+    }
+
+    @Override
+    public boolean hasInitialResponse() {
+        return true;
+    }
+
+    @Override
+    public byte[] evaluateChallenge(byte[] challenge) throws SaslException {
+        try {
+            OAuthBearerTokenCallback callback = new OAuthBearerTokenCallback();
+            switch (state) {
+                case SEND_CLIENT_FIRST_MESSAGE:
+                    if (challenge != null && challenge.length != 0)
+                        throw new SaslException("Expected empty challenge");
+                    callbackHandler().handle(new Callback[] {callback});
+                    SaslExtensions extensions = retrieveCustomExtensions();
+
+                    setState(State.RECEIVE_SERVER_FIRST_MESSAGE);
+
+                    return new OAuthBearerClientInitialResponse(callback.token().value(), extensions).toBytes();
+                case RECEIVE_SERVER_FIRST_MESSAGE:
+                    if (challenge != null && challenge.length != 0) {
+                        String jsonErrorResponse = new String(challenge, StandardCharsets.UTF_8);
+                        if (log.isDebugEnabled())
+                            log.debug("Sending %%x01 response to server after receiving an error: {}",
+                                    jsonErrorResponse);
+                        setState(State.RECEIVE_SERVER_MESSAGE_AFTER_FAILURE);
+                        return new byte[] {BYTE_CONTROL_A};
+                    }
+                    callbackHandler().handle(new Callback[] {callback});
+                    if (log.isDebugEnabled())
+                        log.debug("Successfully authenticated as {}", callback.token().principalName());
+                    setState(State.COMPLETE);
+                    return null;
+                default:
+                    throw new IllegalSaslStateException("Unexpected challenge in Sasl client state " + state);
+            }
+        } catch (SaslException e) {
+            setState(State.FAILED);
+            throw e;
+        } catch (IOException | UnsupportedCallbackException e) {
+            setState(State.FAILED);
+            throw new SaslException(e.getMessage(), e);
+        }
+    }
+
+    @Override
+    public boolean isComplete() {
+        return state == State.COMPLETE;
+    }
+
+    @Override
+    public byte[] unwrap(byte[] incoming, int offset, int len) {
+        if (!isComplete())
+            throw new IllegalStateException("Authentication exchange has not completed");
+        return Arrays.copyOfRange(incoming, offset, offset + len);
+    }
+
+    @Override
+    public byte[] wrap(byte[] outgoing, int offset, int len) {
+        if (!isComplete())
+            throw new IllegalStateException("Authentication exchange has not completed");
+        return Arrays.copyOfRange(outgoing, offset, offset + len);
+    }
+
+    @Override
+    public Object getNegotiatedProperty(String propName) {
+        if (!isComplete())
+            throw new IllegalStateException("Authentication exchange has not completed");
+        return null;
+    }
+
+    @Override
+    public void dispose() {
+    }
+
+    private void setState(State state) {
+        log.debug("Setting SASL/{} client state to {}", OAuthBearerLoginModule.OAUTHBEARER_MECHANISM, state);
+        this.state = state;
+    }
+
+    private SaslExtensions retrieveCustomExtensions() throws SaslException {
+        SaslExtensionsCallback extensionsCallback = new SaslExtensionsCallback();
+        try {
+            callbackHandler().handle(new Callback[] {extensionsCallback});
+        } catch (UnsupportedCallbackException e) {
+            log.debug("Extensions callback is not supported by client callback handler {}, no extensions will be added",
+                    callbackHandler());
+        } catch (Exception e) {
+            throw new SaslException("SASL extensions could not be obtained", e);
+        }
+
+        return extensionsCallback.extensions();
+    }
+
+    public static class OAuthBearerSaslClientFactory implements SaslClientFactory {
+        @Override
+        public SaslClient createSaslClient(String[] mechanisms, String authorizationId, String protocol,
+                String serverName, Map props, CallbackHandler callbackHandler) {
+            String[] mechanismNamesCompatibleWithPolicy = getMechanismNames(props);
+            for (String mechanism : mechanisms) {
+                for (int i = 0; i < mechanismNamesCompatibleWithPolicy.length; i++) {
+                    if (mechanismNamesCompatibleWithPolicy[i].equals(mechanism)) {
+                        if (!(Objects.requireNonNull(callbackHandler) instanceof AuthenticateCallbackHandler))
+                            throw new IllegalArgumentException(String.format(
+                                    "Callback handler must be castable to %s: %s",
+                                    AuthenticateCallbackHandler.class.getName(), callbackHandler.getClass().getName()));
+                        return new OAuthBearerSaslClient((AuthenticateCallbackHandler) callbackHandler);
+                    }
+                }
+            }
+            return null;
+        }
+
+        @Override
+        public String[] getMechanismNames(Map props) {
+            return OAuthBearerSaslServer.mechanismNamesCompatibleWithPolicy(props);
+        }
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/OAuthBearerSaslClientCallbackHandler.java b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/OAuthBearerSaslClientCallbackHandler.java
new file mode 100644
index 0000000000000..bca55be13bb8c
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/OAuthBearerSaslClientCallbackHandler.java
@@ -0,0 +1,143 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.oauthbearer.internals;
+
+import java.io.IOException;
+import java.security.AccessController;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.Date;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.SortedSet;
+import java.util.TreeSet;
+
+import javax.security.auth.Subject;
+import javax.security.auth.callback.Callback;
+import javax.security.auth.callback.UnsupportedCallbackException;
+import javax.security.auth.login.AppConfigurationEntry;
+
+import org.apache.kafka.common.security.auth.SaslExtensionsCallback;
+import org.apache.kafka.common.security.auth.AuthenticateCallbackHandler;
+import org.apache.kafka.common.security.auth.SaslExtensions;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerLoginModule;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerToken;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerTokenCallback;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * An implementation of {@code AuthenticateCallbackHandler} that recognizes
+ * {@link OAuthBearerTokenCallback} and retrieves OAuth 2 Bearer Token that was
+ * created when the {@code OAuthBearerLoginModule} logged in by looking for an
+ * instance of {@link OAuthBearerToken} in the {@code Subject}'s private
+ * credentials. This class also recognizes {@link SaslExtensionsCallback} and retrieves any SASL extensions that were
+ * created when the {@code OAuthBearerLoginModule} logged in by looking for an instance of {@link SaslExtensions}
+ * in the {@code Subject}'s public credentials
+ * 
        
+ * Use of this class is configured automatically and does not need to be
+ * explicitly set via the {@code sasl.client.callback.handler.class}
+ * configuration property.
+ */
+public class OAuthBearerSaslClientCallbackHandler implements AuthenticateCallbackHandler {
+    private static final Logger log = LoggerFactory.getLogger(OAuthBearerSaslClientCallbackHandler.class);
+    private boolean configured = false;
+
+    /**
+     * Return true if this instance has been configured, otherwise false
+     *
+     * @return true if this instance has been configured, otherwise false
+     */
+    public boolean configured() {
+        return configured;
+    }
+
+    @Override
+    public void configure(Map configs, String saslMechanism, List jaasConfigEntries) {
+        if (!OAuthBearerLoginModule.OAUTHBEARER_MECHANISM.equals(saslMechanism))
+            throw new IllegalArgumentException(String.format("Unexpected SASL mechanism: %s", saslMechanism));
+        configured = true;
+    }
+
+    @Override
+    public void handle(Callback[] callbacks) throws IOException, UnsupportedCallbackException {
+        if (!configured())
+            throw new IllegalStateException("Callback handler not configured");
+        for (Callback callback : callbacks) {
+            if (callback instanceof OAuthBearerTokenCallback)
+                handleCallback((OAuthBearerTokenCallback) callback);
+            else if (callback instanceof SaslExtensionsCallback)
+                handleCallback((SaslExtensionsCallback) callback, Subject.getSubject(AccessController.getContext()));
+            else
+                throw new UnsupportedCallbackException(callback);
+        }
+    }
+
+    @Override
+    public void close() {
+        // empty
+    }
+
+    private void handleCallback(OAuthBearerTokenCallback callback) throws IOException {
+        if (callback.token() != null)
+            throw new IllegalArgumentException("Callback had a token already");
+        Subject subject = Subject.getSubject(AccessController.getContext());
+        Set privateCredentials = subject != null
+            ? subject.getPrivateCredentials(OAuthBearerToken.class)
+            : Collections.emptySet();
+        if (privateCredentials.size() == 0)
+            throw new IOException("No OAuth Bearer tokens in Subject's private credentials");
+        if (privateCredentials.size() == 1)
+            callback.token(privateCredentials.iterator().next());
+        else {
+            /*
+             * There a very small window of time upon token refresh (on the order of milliseconds)
+             * where both an old and a new token appear on the Subject's private credentials.
+             * Rather than implement a lock to eliminate this window, we will deal with it by
+             * checking for the existence of multiple tokens and choosing the one that has the
+             * longest lifetime.  It is also possible that a bug could cause multiple tokens to
+             * exist (e.g. KAFKA-7902), so dealing with the unlikely possibility that occurs
+             * during normal operation also allows us to deal more robustly with potential bugs.
+             */
+            SortedSet sortedByLifetime =
+                new TreeSet<>(
+                    new Comparator() {
+                        @Override
+                        public int compare(OAuthBearerToken o1, OAuthBearerToken o2) {
+                            return Long.compare(o1.lifetimeMs(), o2.lifetimeMs());
+                        }
+                    });
+            sortedByLifetime.addAll(privateCredentials);
+            log.warn("Found {} OAuth Bearer tokens in Subject's private credentials; the oldest expires at {}, will use the newest, which expires at {}",
+                sortedByLifetime.size(),
+                new Date(sortedByLifetime.first().lifetimeMs()),
+                new Date(sortedByLifetime.last().lifetimeMs()));
+            callback.token(sortedByLifetime.last());
+        }
+    }
+
+    /**
+     * Attaches the first {@link SaslExtensions} found in the public credentials of the Subject
+     */
+    private static void handleCallback(SaslExtensionsCallback extensionsCallback, Subject subject) {
+        if (subject != null && !subject.getPublicCredentials(SaslExtensions.class).isEmpty()) {
+            SaslExtensions extensions = subject.getPublicCredentials(SaslExtensions.class).iterator().next();
+            extensionsCallback.extensions(extensions);
+        }
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/OAuthBearerSaslClientProvider.java b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/OAuthBearerSaslClientProvider.java
new file mode 100644
index 0000000000000..08777ef593e5b
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/OAuthBearerSaslClientProvider.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.oauthbearer.internals;
+
+import java.security.Provider;
+import java.security.Security;
+
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerLoginModule;
+import org.apache.kafka.common.security.oauthbearer.internals.OAuthBearerSaslClient.OAuthBearerSaslClientFactory;
+
+public class OAuthBearerSaslClientProvider extends Provider {
+    private static final long serialVersionUID = 1L;
+
+    protected OAuthBearerSaslClientProvider() {
+        super("SASL/OAUTHBEARER Client Provider", 1.0, "SASL/OAUTHBEARER Client Provider for Kafka");
+        put("SaslClientFactory." + OAuthBearerLoginModule.OAUTHBEARER_MECHANISM,
+                OAuthBearerSaslClientFactory.class.getName());
+    }
+
+    public static void initialize() {
+        Security.addProvider(new OAuthBearerSaslClientProvider());
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/OAuthBearerSaslServer.java b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/OAuthBearerSaslServer.java
new file mode 100644
index 0000000000000..8735f49d07e15
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/OAuthBearerSaslServer.java
@@ -0,0 +1,245 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.oauthbearer.internals;
+
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
+import java.util.Map;
+import java.util.Objects;
+
+import javax.security.auth.callback.Callback;
+import javax.security.auth.callback.CallbackHandler;
+import javax.security.auth.callback.UnsupportedCallbackException;
+import javax.security.sasl.Sasl;
+import javax.security.sasl.SaslException;
+import javax.security.sasl.SaslServer;
+import javax.security.sasl.SaslServerFactory;
+
+import org.apache.kafka.common.errors.SaslAuthenticationException;
+import org.apache.kafka.common.security.auth.SaslExtensions;
+import org.apache.kafka.common.security.authenticator.SaslInternalConfigs;
+import org.apache.kafka.common.security.auth.AuthenticateCallbackHandler;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerExtensionsValidatorCallback;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerLoginModule;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerToken;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerValidatorCallback;
+import org.apache.kafka.common.utils.Utils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * {@code SaslServer} implementation for SASL/OAUTHBEARER in Kafka. An instance
+ * of {@link OAuthBearerToken} is available upon successful authentication via
+ * the negotiated property "{@code OAUTHBEARER.token}"; the token could be used
+ * in a custom authorizer (to authorize based on JWT claims rather than ACLs,
+ * for example).
+ */
+public class OAuthBearerSaslServer implements SaslServer {
+
+    private static final Logger log = LoggerFactory.getLogger(OAuthBearerSaslServer.class);
+    private static final String NEGOTIATED_PROPERTY_KEY_TOKEN = OAuthBearerLoginModule.OAUTHBEARER_MECHANISM + ".token";
+    private static final String INTERNAL_ERROR_ON_SERVER = "Authentication could not be performed due to an internal error on the server";
+
+    private final AuthenticateCallbackHandler callbackHandler;
+
+    private boolean complete;
+    private OAuthBearerToken tokenForNegotiatedProperty = null;
+    private String errorMessage = null;
+    private SaslExtensions extensions;
+
+    public OAuthBearerSaslServer(CallbackHandler callbackHandler) {
+        if (!(Objects.requireNonNull(callbackHandler) instanceof AuthenticateCallbackHandler))
+            throw new IllegalArgumentException(String.format("Callback handler must be castable to %s: %s",
+                    AuthenticateCallbackHandler.class.getName(), callbackHandler.getClass().getName()));
+        this.callbackHandler = (AuthenticateCallbackHandler) callbackHandler;
+    }
+
+    /**
+     * @throws SaslAuthenticationException
+     *             if access token cannot be validated
+     *             
        
+     *             Note: This method may throw
+     *             {@link SaslAuthenticationException} to provide custom error
+     *             messages to clients. But care should be taken to avoid including
+     *             any information in the exception message that should not be
+     *             leaked to unauthenticated clients. It may be safer to throw
+     *             {@link SaslException} in some cases so that a standard error
+     *             message is returned to clients.
+     *             
        
+     */
+    @Override
+    public byte[] evaluateResponse(byte[] response) throws SaslException, SaslAuthenticationException {
+        if (response.length == 1 && response[0] == OAuthBearerSaslClient.BYTE_CONTROL_A && errorMessage != null) {
+            log.debug("Received %x01 response from client after it received our error");
+            throw new SaslAuthenticationException(errorMessage);
+        }
+        errorMessage = null;
+
+        OAuthBearerClientInitialResponse clientResponse;
+        try {
+            clientResponse = new OAuthBearerClientInitialResponse(response);
+        } catch (SaslException e) {
+            log.debug(e.getMessage());
+            throw e;
+        }
+
+        return process(clientResponse.tokenValue(), clientResponse.authorizationId(), clientResponse.extensions());
+    }
+
+    @Override
+    public String getAuthorizationID() {
+        if (!complete)
+            throw new IllegalStateException("Authentication exchange has not completed");
+        return tokenForNegotiatedProperty.principalName();
+    }
+
+    @Override
+    public String getMechanismName() {
+        return OAuthBearerLoginModule.OAUTHBEARER_MECHANISM;
+    }
+
+    @Override
+    public Object getNegotiatedProperty(String propName) {
+        if (!complete)
+            throw new IllegalStateException("Authentication exchange has not completed");
+        if (NEGOTIATED_PROPERTY_KEY_TOKEN.equals(propName))
+            return tokenForNegotiatedProperty;
+        if (SaslInternalConfigs.CREDENTIAL_LIFETIME_MS_SASL_NEGOTIATED_PROPERTY_KEY.equals(propName))
+            return tokenForNegotiatedProperty.lifetimeMs();
+        return extensions.map().get(propName);
+    }
+
+    @Override
+    public boolean isComplete() {
+        return complete;
+    }
+
+    @Override
+    public byte[] unwrap(byte[] incoming, int offset, int len) {
+        if (!complete)
+            throw new IllegalStateException("Authentication exchange has not completed");
+        return Arrays.copyOfRange(incoming, offset, offset + len);
+    }
+
+    @Override
+    public byte[] wrap(byte[] outgoing, int offset, int len) {
+        if (!complete)
+            throw new IllegalStateException("Authentication exchange has not completed");
+        return Arrays.copyOfRange(outgoing, offset, offset + len);
+    }
+
+    @Override
+    public void dispose() {
+        complete = false;
+        tokenForNegotiatedProperty = null;
+        extensions = null;
+    }
+
+    private byte[] process(String tokenValue, String authorizationId, SaslExtensions extensions) throws SaslException {
+        OAuthBearerValidatorCallback callback = new OAuthBearerValidatorCallback(tokenValue);
+        try {
+            callbackHandler.handle(new Callback[] {callback});
+        } catch (IOException | UnsupportedCallbackException e) {
+            handleCallbackError(e);
+        }
+        OAuthBearerToken token = callback.token();
+        if (token == null) {
+            errorMessage = jsonErrorResponse(callback.errorStatus(), callback.errorScope(),
+                    callback.errorOpenIDConfiguration());
+            log.debug(errorMessage);
+            return errorMessage.getBytes(StandardCharsets.UTF_8);
+        }
+        /*
+         * We support the client specifying an authorization ID as per the SASL
+         * specification, but it must match the principal name if it is specified.
+         */
+        if (!authorizationId.isEmpty() && !authorizationId.equals(token.principalName()))
+            throw new SaslAuthenticationException(String.format(
+                    "Authentication failed: Client requested an authorization id (%s) that is different from the token's principal name (%s)",
+                    authorizationId, token.principalName()));
+
+        Map validExtensions = processExtensions(token, extensions);
+
+        tokenForNegotiatedProperty = token;
+        this.extensions = new SaslExtensions(validExtensions);
+        complete = true;
+        log.debug("Successfully authenticate User={}", token.principalName());
+        return new byte[0];
+    }
+
+    private Map processExtensions(OAuthBearerToken token, SaslExtensions extensions) throws SaslException {
+        OAuthBearerExtensionsValidatorCallback extensionsCallback = new OAuthBearerExtensionsValidatorCallback(token, extensions);
+        try {
+            callbackHandler.handle(new Callback[] {extensionsCallback});
+        } catch (UnsupportedCallbackException e) {
+            // backwards compatibility - no extensions will be added
+        } catch (IOException e) {
+            handleCallbackError(e);
+        }
+        if (!extensionsCallback.invalidExtensions().isEmpty()) {
+            String errorMessage = String.format("Authentication failed: %d extensions are invalid! They are: %s",
+                    extensionsCallback.invalidExtensions().size(),
+                    Utils.mkString(extensionsCallback.invalidExtensions(), "", "", ": ", "; "));
+            log.debug(errorMessage);
+            throw new SaslAuthenticationException(errorMessage);
+        }
+
+        return extensionsCallback.validatedExtensions();
+    }
+
+    private static String jsonErrorResponse(String errorStatus, String errorScope, String errorOpenIDConfiguration) {
+        String jsonErrorResponse = String.format("{\"status\":\"%s\"", errorStatus);
+        if (errorScope != null)
+            jsonErrorResponse = String.format("%s, \"scope\":\"%s\"", jsonErrorResponse, errorScope);
+        if (errorOpenIDConfiguration != null)
+            jsonErrorResponse = String.format("%s, \"openid-configuration\":\"%s\"", jsonErrorResponse,
+                    errorOpenIDConfiguration);
+        jsonErrorResponse = String.format("%s}", jsonErrorResponse);
+        return jsonErrorResponse;
+    }
+
+    private void handleCallbackError(Exception e) throws SaslException {
+        String msg = String.format("%s: %s", INTERNAL_ERROR_ON_SERVER, e.getMessage());
+        log.debug(msg, e);
+        throw new SaslException(msg);
+    }
+
+    public static String[] mechanismNamesCompatibleWithPolicy(Map props) {
+        return props != null && "true".equals(String.valueOf(props.get(Sasl.POLICY_NOPLAINTEXT))) ? new String[] {}
+                : new String[] {OAuthBearerLoginModule.OAUTHBEARER_MECHANISM};
+    }
+
+    public static class OAuthBearerSaslServerFactory implements SaslServerFactory {
+        @Override
+        public SaslServer createSaslServer(String mechanism, String protocol, String serverName, Map props,
+                CallbackHandler callbackHandler) {
+            String[] mechanismNamesCompatibleWithPolicy = getMechanismNames(props);
+            for (int i = 0; i < mechanismNamesCompatibleWithPolicy.length; i++) {
+                if (mechanismNamesCompatibleWithPolicy[i].equals(mechanism)) {
+                    return new OAuthBearerSaslServer(callbackHandler);
+                }
+            }
+            return null;
+        }
+
+        @Override
+        public String[] getMechanismNames(Map props) {
+            return OAuthBearerSaslServer.mechanismNamesCompatibleWithPolicy(props);
+        }
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/OAuthBearerSaslServerProvider.java b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/OAuthBearerSaslServerProvider.java
new file mode 100644
index 0000000000000..2e179ce94c57e
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/OAuthBearerSaslServerProvider.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.oauthbearer.internals;
+
+import java.security.Provider;
+import java.security.Security;
+
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerLoginModule;
+import org.apache.kafka.common.security.oauthbearer.internals.OAuthBearerSaslServer.OAuthBearerSaslServerFactory;
+
+public class OAuthBearerSaslServerProvider extends Provider {
+    private static final long serialVersionUID = 1L;
+
+    protected OAuthBearerSaslServerProvider() {
+        super("SASL/OAUTHBEARER Server Provider", 1.0, "SASL/OAUTHBEARER Server Provider for Kafka");
+        put("SaslServerFactory." + OAuthBearerLoginModule.OAUTHBEARER_MECHANISM,
+                OAuthBearerSaslServerFactory.class.getName());
+    }
+
+    public static void initialize() {
+        Security.addProvider(new OAuthBearerSaslServerProvider());
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/expiring/ExpiringCredential.java b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/expiring/ExpiringCredential.java
new file mode 100644
index 0000000000000..1bfa4b25c4e92
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/expiring/ExpiringCredential.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.security.oauthbearer.internals.expiring;
+
+/**
+ * A credential that expires and that can potentially be refreshed
+ * 
+ * @see ExpiringCredentialRefreshingLogin
+ */
+public interface ExpiringCredential {
+    /**
+     * The name of the principal to which this credential applies (used only for
+     * logging)
+     * 
+     * @return the always non-null/non-empty principal name
+     */
+    String principalName();
+
+    /**
+     * When the credential became valid, in terms of the number of milliseconds
+     * since the epoch, if known, otherwise null. An expiring credential may not
+     * necessarily indicate when it was created -- just when it expires -- so we
+     * need to support a null return value here.
+     * 
+     * @return the time when the credential became valid, in terms of the number of
+     *         milliseconds since the epoch, if known, otherwise null
+     */
+    Long startTimeMs();
+
+    /**
+     * When the credential expires, in terms of the number of milliseconds since the
+     * epoch. All expiring credentials by definition must indicate their expiration
+     * time -- thus, unlike other methods, we do not support a null return value
+     * here.
+     * 
+     * @return the time when the credential expires, in terms of the number of
+     *         milliseconds since the epoch
+     */
+    long expireTimeMs();
+
+    /**
+     * The point after which the credential can no longer be refreshed, in terms of
+     * the number of milliseconds since the epoch, if any, otherwise null. Some
+     * expiring credentials can be refreshed over and over again without limit, so
+     * we support a null return value here.
+     * 
+     * @return the point after which the credential can no longer be refreshed, in
+     *         terms of the number of milliseconds since the epoch, if any,
+     *         otherwise null
+     */
+    Long absoluteLastRefreshTimeMs();
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/expiring/ExpiringCredentialRefreshConfig.java b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/expiring/ExpiringCredentialRefreshConfig.java
new file mode 100644
index 0000000000000..1df69f756a4e0
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/expiring/ExpiringCredentialRefreshConfig.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.security.oauthbearer.internals.expiring;
+
+import java.util.Map;
+import java.util.Objects;
+
+import org.apache.kafka.common.config.SaslConfigs;
+
+/**
+ * Immutable refresh-related configuration for expiring credentials that can be
+ * parsed from a producer/consumer/broker config.
+ */
+public class ExpiringCredentialRefreshConfig {
+    private final double loginRefreshWindowFactor;
+    private final double loginRefreshWindowJitter;
+    private final short loginRefreshMinPeriodSeconds;
+    private final short loginRefreshBufferSeconds;
+    private final boolean loginRefreshReloginAllowedBeforeLogout;
+
+    /**
+     * Constructor based on producer/consumer/broker configs and the indicated value
+     * for whether or not client relogin is allowed before logout
+     * 
+     * @param configs
+     *            the mandatory (but possibly empty) producer/consumer/broker
+     *            configs upon which to build this instance
+     * @param clientReloginAllowedBeforeLogout
+     *            if the {@code LoginModule} and {@code SaslClient} implementations
+     *            support multiple simultaneous login contexts on a single
+     *            {@code Subject} at the same time. If true, then upon refresh,
+     *            logout will only be invoked on the original {@code LoginContext}
+     *            after a new one successfully logs in. This can be helpful if the
+     *            original credential still has some lifetime left when an attempt
+     *            to refresh the credential fails; the client will still be able to
+     *            create new connections as long as the original credential remains
+     *            valid. Otherwise, if logout is immediately invoked prior to
+     *            relogin, a relogin failure leaves the client without the ability
+     *            to connect until relogin does in fact succeed.
+     */
+    public ExpiringCredentialRefreshConfig(Map configs, boolean clientReloginAllowedBeforeLogout) {
+        Objects.requireNonNull(configs);
+        this.loginRefreshWindowFactor = (Double) configs.get(SaslConfigs.SASL_LOGIN_REFRESH_WINDOW_FACTOR);
+        this.loginRefreshWindowJitter = (Double) configs.get(SaslConfigs.SASL_LOGIN_REFRESH_WINDOW_JITTER);
+        this.loginRefreshMinPeriodSeconds = (Short) configs.get(SaslConfigs.SASL_LOGIN_REFRESH_MIN_PERIOD_SECONDS);
+        this.loginRefreshBufferSeconds = (Short) configs.get(SaslConfigs.SASL_LOGIN_REFRESH_BUFFER_SECONDS);
+        this.loginRefreshReloginAllowedBeforeLogout = clientReloginAllowedBeforeLogout;
+    }
+
+    /**
+     * Background login refresh thread will sleep until the specified window factor
+     * relative to the credential's total lifetime has been reached, at which time
+     * it will try to refresh the credential.
+     * 
+     * @return the login refresh window factor
+     */
+    public double loginRefreshWindowFactor() {
+        return loginRefreshWindowFactor;
+    }
+
+    /**
+     * Amount of random jitter added to the background login refresh thread's sleep
+     * time.
+     * 
+     * @return the login refresh window jitter
+     */
+    public double loginRefreshWindowJitter() {
+        return loginRefreshWindowJitter;
+    }
+
+    /**
+     * The desired minimum time between checks by the background login refresh
+     * thread, in seconds
+     * 
+     * @return the desired minimum refresh period, in seconds
+     */
+    public short loginRefreshMinPeriodSeconds() {
+        return loginRefreshMinPeriodSeconds;
+    }
+
+    /**
+     * The amount of buffer time before expiration to maintain when refreshing. If a
+     * refresh is scheduled to occur closer to expiration than the number of seconds
+     * defined here then the refresh will be moved up to maintain as much of the
+     * desired buffer as possible.
+     * 
+     * @return the refresh buffer, in seconds
+     */
+    public short loginRefreshBufferSeconds() {
+        return loginRefreshBufferSeconds;
+    }
+
+    /**
+     * If the LoginModule and SaslClient implementations support multiple
+     * simultaneous login contexts on a single Subject at the same time. If true,
+     * then upon refresh, logout will only be invoked on the original LoginContext
+     * after a new one successfully logs in. This can be helpful if the original
+     * credential still has some lifetime left when an attempt to refresh the
+     * credential fails; the client will still be able to create new connections as
+     * long as the original credential remains valid. Otherwise, if logout is
+     * immediately invoked prior to relogin, a relogin failure leaves the client
+     * without the ability to connect until relogin does in fact succeed.
+     * 
+     * @return true if relogin is allowed prior to discarding an existing
+     *         (presumably unexpired) credential, otherwise false
+     */
+    public boolean loginRefreshReloginAllowedBeforeLogout() {
+        return loginRefreshReloginAllowedBeforeLogout;
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/expiring/ExpiringCredentialRefreshingLogin.java b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/expiring/ExpiringCredentialRefreshingLogin.java
new file mode 100644
index 0000000000000..ab2f303e09124
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/expiring/ExpiringCredentialRefreshingLogin.java
@@ -0,0 +1,444 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.oauthbearer.internals.expiring;
+
+import java.util.Date;
+import java.util.Objects;
+import java.util.Random;
+
+import javax.security.auth.Subject;
+import javax.security.auth.login.Configuration;
+import javax.security.auth.login.LoginContext;
+import javax.security.auth.login.LoginException;
+
+import org.apache.kafka.common.security.auth.AuthenticateCallbackHandler;
+import org.apache.kafka.common.security.auth.Login;
+import org.apache.kafka.common.utils.KafkaThread;
+import org.apache.kafka.common.utils.Time;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * This class is responsible for refreshing logins for both Kafka client and
+ * server when the login is a type that has a limited lifetime/will expire. The
+ * credentials for the login must implement {@link ExpiringCredential}.
+ */
+public abstract class ExpiringCredentialRefreshingLogin implements AutoCloseable {
+    /**
+     * Class that can be overridden for testing
+     */
+    static class LoginContextFactory {
+        public LoginContext createLoginContext(ExpiringCredentialRefreshingLogin expiringCredentialRefreshingLogin)
+                throws LoginException {
+            return new LoginContext(expiringCredentialRefreshingLogin.contextName(),
+                    expiringCredentialRefreshingLogin.subject(), expiringCredentialRefreshingLogin.callbackHandler(),
+                    expiringCredentialRefreshingLogin.configuration());
+        }
+
+        public void refresherThreadStarted() {
+            // empty
+        }
+
+        public void refresherThreadDone() {
+            // empty
+        }
+    }
+
+    private static class ExitRefresherThreadDueToIllegalStateException extends Exception {
+        private static final long serialVersionUID = -6108495378411920380L;
+
+        public ExitRefresherThreadDueToIllegalStateException(String message) {
+            super(message);
+        }
+    }
+
+    private class Refresher implements Runnable {
+        @Override
+        public void run() {
+            log.info("[Principal={}]: Expiring credential re-login thread started.", principalLogText());
+            while (true) {
+                /*
+                 * Refresh thread's main loop. Each expiring credential lives for one iteration
+                 * of the loop. Thread will exit if the loop exits from here.
+                 */
+                long nowMs = currentMs();
+                Long nextRefreshMs = refreshMs(nowMs);
+                if (nextRefreshMs == null) {
+                    loginContextFactory.refresherThreadDone();
+                    return;
+                }
+                // safety check motivated by KAFKA-7945,
+                // should generally never happen except due to a bug
+                if (nextRefreshMs.longValue() < nowMs) {
+                    log.warn("[Principal={}]: Expiring credential re-login sleep time was calculated to be in the past! Will explicitly adjust. ({})", principalLogText(),
+                            new Date(nextRefreshMs));
+                    nextRefreshMs = Long.valueOf(nowMs + 10 * 1000); // refresh in 10 seconds
+                }
+                log.info("[Principal={}]: Expiring credential re-login sleeping until: {}", principalLogText(),
+                        new Date(nextRefreshMs));
+                time.sleep(nextRefreshMs - nowMs);
+                if (Thread.currentThread().isInterrupted()) {
+                    log.info("[Principal={}]: Expiring credential re-login thread has been interrupted and will exit.",
+                            principalLogText());
+                    loginContextFactory.refresherThreadDone();
+                    return;
+                }
+                while (true) {
+                    /*
+                     * Perform a re-login over and over again with some intervening delay
+                     * unless/until either the refresh succeeds or we are interrupted.
+                     */
+                    try {
+                        reLogin();
+                        break; // success
+                    } catch (ExitRefresherThreadDueToIllegalStateException e) {
+                        log.error(e.getMessage(), e);
+                        loginContextFactory.refresherThreadDone();
+                        return;
+                    } catch (LoginException loginException) {
+                        log.warn(String.format(
+                                "[Principal=%s]: LoginException during login retry; will sleep %d seconds before trying again.",
+                                principalLogText(), DELAY_SECONDS_BEFORE_NEXT_RETRY_WHEN_RELOGIN_FAILS),
+                                loginException);
+                        // Sleep and allow loop to run/try again unless interrupted
+                        time.sleep(DELAY_SECONDS_BEFORE_NEXT_RETRY_WHEN_RELOGIN_FAILS * 1000);
+                        if (Thread.currentThread().isInterrupted()) {
+                            log.error(
+                                    "[Principal={}]: Interrupted while trying to perform a subsequent expiring credential re-login after one or more initial re-login failures: re-login thread exiting now: {}",
+                                    principalLogText(), String.valueOf(loginException.getMessage()));
+                            loginContextFactory.refresherThreadDone();
+                            return;
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    private static final Logger log = LoggerFactory.getLogger(ExpiringCredentialRefreshingLogin.class);
+    private static final long DELAY_SECONDS_BEFORE_NEXT_RETRY_WHEN_RELOGIN_FAILS = 10L;
+    private static final Random RNG = new Random();
+    private final Time time;
+    private Thread refresherThread;
+
+    private final LoginContextFactory loginContextFactory;
+    private final String contextName;
+    private final Configuration configuration;
+    private final ExpiringCredentialRefreshConfig expiringCredentialRefreshConfig;
+    private final AuthenticateCallbackHandler callbackHandler;
+
+    // mark volatile due to existence of public subject() method
+    private volatile Subject subject = null;
+    private boolean hasExpiringCredential = false;
+    private String principalName = null;
+    private LoginContext loginContext = null;
+    private ExpiringCredential expiringCredential = null;
+    private final Class mandatoryClassToSynchronizeOnPriorToRefresh;
+
+    public ExpiringCredentialRefreshingLogin(String contextName, Configuration configuration,
+            ExpiringCredentialRefreshConfig expiringCredentialRefreshConfig,
+            AuthenticateCallbackHandler callbackHandler, Class mandatoryClassToSynchronizeOnPriorToRefresh) {
+        this(contextName, configuration, expiringCredentialRefreshConfig, callbackHandler,
+                mandatoryClassToSynchronizeOnPriorToRefresh, new LoginContextFactory(), Time.SYSTEM);
+    }
+
+    public ExpiringCredentialRefreshingLogin(String contextName, Configuration configuration,
+            ExpiringCredentialRefreshConfig expiringCredentialRefreshConfig,
+            AuthenticateCallbackHandler callbackHandler, Class mandatoryClassToSynchronizeOnPriorToRefresh,
+            LoginContextFactory loginContextFactory, Time time) {
+        this.contextName = Objects.requireNonNull(contextName);
+        this.configuration = Objects.requireNonNull(configuration);
+        this.expiringCredentialRefreshConfig = Objects.requireNonNull(expiringCredentialRefreshConfig);
+        this.callbackHandler = callbackHandler;
+        this.mandatoryClassToSynchronizeOnPriorToRefresh = Objects
+                .requireNonNull(mandatoryClassToSynchronizeOnPriorToRefresh);
+        this.loginContextFactory = loginContextFactory;
+        this.time = Objects.requireNonNull(time);
+    }
+
+    public Subject subject() {
+        return subject; // field requires volatile keyword
+    }
+
+    public String contextName() {
+        return contextName;
+    }
+
+    public Configuration configuration() {
+        return configuration;
+    }
+
+    public AuthenticateCallbackHandler callbackHandler() {
+        return callbackHandler;
+    }
+
+    public String serviceName() {
+        return "kafka";
+    }
+
+    /**
+     * Performs login for each login module specified for the login context of this
+     * instance and starts the thread used to periodically re-login.
+     * 
        
+     * The synchronized keyword is not necessary because an implementation of
+     * {@link Login} will delegate to this code (e.g. OAuthBearerRefreshingLogin},
+     * and the {@code login()} method on the delegating class will itself be
+     * synchronized if necessary.
+     */
+    public LoginContext login() throws LoginException {
+        LoginContext tmpLoginContext = loginContextFactory.createLoginContext(this);
+        tmpLoginContext.login();
+        log.info("Successfully logged in.");
+        loginContext = tmpLoginContext;
+        subject = loginContext.getSubject();
+        expiringCredential = expiringCredential();
+        hasExpiringCredential = expiringCredential != null;
+        if (!hasExpiringCredential) {
+            // do not bother with re-logins.
+            log.debug("No Expiring Credential");
+            principalName = null;
+            refresherThread = null;
+            return loginContext;
+        }
+
+        principalName = expiringCredential.principalName();
+
+        // Check for a clock skew problem
+        long expireTimeMs = expiringCredential.expireTimeMs();
+        long nowMs = currentMs();
+        if (nowMs > expireTimeMs) {
+            log.error(
+                    "[Principal={}]: Current clock: {} is later than expiry {}. This may indicate a clock skew problem."
+                            + " Check that this host's and remote host's clocks are in sync. Not starting refresh thread."
+                            + " This process is likely unable to authenticate SASL connections (for example, it is unlikely"
+                            + " to be able to authenticate a connection with a Kafka Broker).",
+                    principalLogText(), new Date(nowMs), new Date(expireTimeMs));
+            return loginContext;
+        }
+
+        if (log.isDebugEnabled())
+            log.debug("[Principal={}]: It is an expiring credential", principalLogText());
+
+        /*
+         * Re-login periodically. How often is determined by the expiration date of the
+         * credential and refresh-related configuration values.
+         */
+        refresherThread = KafkaThread.daemon(String.format("kafka-expiring-relogin-thread-%s", principalName),
+                new Refresher());
+        refresherThread.start();
+        loginContextFactory.refresherThreadStarted();
+        return loginContext;
+    }
+
+    public void close() {
+        if (refresherThread != null && refresherThread.isAlive()) {
+            refresherThread.interrupt();
+            try {
+                refresherThread.join();
+            } catch (InterruptedException e) {
+                log.warn("[Principal={}]: Interrupted while waiting for re-login thread to shutdown.",
+                        principalLogText(), e);
+                Thread.currentThread().interrupt();
+            }
+        }
+    }
+
+    public abstract ExpiringCredential expiringCredential();
+
+    /**
+     * Determine when to sleep until before performing a refresh
+     * 
+     * @param relativeToMs
+     *            the point (in terms of number of milliseconds since the epoch) at
+     *            which to perform the calculation
+     * @return null if no refresh should occur, otherwise the time to sleep until
+     *         (in terms of the number of milliseconds since the epoch) before
+     *         performing a refresh
+     */
+    private Long refreshMs(long relativeToMs) {
+        if (expiringCredential == null) {
+            /*
+             * Re-login failed because our login() invocation did not generate a credential
+             * but also did not generate an exception. Try logging in again after some delay
+             * (it seems likely to be a bug, but it doesn't hurt to keep trying to refresh).
+             */
+            long retvalNextRefreshMs = relativeToMs + DELAY_SECONDS_BEFORE_NEXT_RETRY_WHEN_RELOGIN_FAILS * 1000L;
+            log.warn("[Principal={}]: No Expiring credential found: will try again at {}", principalLogText(),
+                    new Date(retvalNextRefreshMs));
+            return retvalNextRefreshMs;
+        }
+        long expireTimeMs = expiringCredential.expireTimeMs();
+        if (relativeToMs > expireTimeMs) {
+            boolean logoutRequiredBeforeLoggingBackIn = isLogoutRequiredBeforeLoggingBackIn();
+            if (logoutRequiredBeforeLoggingBackIn) {
+                log.error(
+                        "[Principal={}]: Current clock: {} is later than expiry {}. This may indicate a clock skew problem."
+                                + " Check that this host's and remote host's clocks are in sync. Exiting refresh thread.",
+                        principalLogText(), new Date(relativeToMs), new Date(expireTimeMs));
+                return null;
+            } else {
+                /*
+                 * Since the current soon-to-expire credential isn't logged out until we have a
+                 * new credential with a refreshed lifetime, it is possible that the current
+                 * credential could expire if the re-login continually fails over and over again
+                 * making us unable to get the new credential. Therefore keep trying rather than
+                 * exiting.
+                 */
+                long retvalNextRefreshMs = relativeToMs + DELAY_SECONDS_BEFORE_NEXT_RETRY_WHEN_RELOGIN_FAILS * 1000L;
+                log.warn("[Principal={}]: Expiring credential already expired at {}: will try to refresh again at {}",
+                        principalLogText(), new Date(expireTimeMs), new Date(retvalNextRefreshMs));
+                return retvalNextRefreshMs;
+            }
+        }
+        Long absoluteLastRefreshTimeMs = expiringCredential.absoluteLastRefreshTimeMs();
+        if (absoluteLastRefreshTimeMs != null && absoluteLastRefreshTimeMs.longValue() < expireTimeMs) {
+            log.warn("[Principal={}]: Expiring credential refresh thread exiting because the"
+                    + " expiring credential's current expiration time ({}) exceeds the latest possible refresh time ({})."
+                    + " This process will not be able to authenticate new SASL connections after that"
+                    + " time (for example, it will not be able to authenticate a new connection with a Kafka Broker).",
+                    principalLogText(), new Date(expireTimeMs), new Date(absoluteLastRefreshTimeMs.longValue()));
+            return null;
+        }
+        Long optionalStartTime = expiringCredential.startTimeMs();
+        long startMs = optionalStartTime != null ? optionalStartTime.longValue() : relativeToMs;
+        log.info("[Principal={}]: Expiring credential valid from {} to {}", expiringCredential.principalName(),
+                new java.util.Date(startMs), new java.util.Date(expireTimeMs));
+
+        double pct = expiringCredentialRefreshConfig.loginRefreshWindowFactor()
+                + (expiringCredentialRefreshConfig.loginRefreshWindowJitter() * RNG.nextDouble());
+        /*
+         * Ignore buffer times if the credential's remaining lifetime is less than their
+         * sum.
+         */
+        long refreshMinPeriodSeconds = expiringCredentialRefreshConfig.loginRefreshMinPeriodSeconds();
+        long clientRefreshBufferSeconds = expiringCredentialRefreshConfig.loginRefreshBufferSeconds();
+        if (relativeToMs + 1000L * (refreshMinPeriodSeconds + clientRefreshBufferSeconds) > expireTimeMs) {
+            long retvalRefreshMs = relativeToMs + (long) ((expireTimeMs - relativeToMs) * pct);
+            log.warn(
+                    "[Principal={}]: Expiring credential expires at {}, so buffer times of {} and {} seconds"
+                            + " at the front and back, respectively, cannot be accommodated.  We will refresh at {}.",
+                    principalLogText(), new Date(expireTimeMs), refreshMinPeriodSeconds, clientRefreshBufferSeconds,
+                    new Date(retvalRefreshMs));
+            return retvalRefreshMs;
+        }
+        long proposedRefreshMs = startMs + (long) ((expireTimeMs - startMs) * pct);
+        // Don't let it violate the requested end buffer time
+        long beginningOfEndBufferTimeMs = expireTimeMs - clientRefreshBufferSeconds * 1000;
+        if (proposedRefreshMs > beginningOfEndBufferTimeMs) {
+            log.info(
+                    "[Principal={}]: Proposed refresh time of {} extends into the desired buffer time of {} seconds before expiration, so refresh it at the desired buffer begin point, at {}",
+                    expiringCredential.principalName(), new Date(proposedRefreshMs), clientRefreshBufferSeconds,
+                    new Date(beginningOfEndBufferTimeMs));
+            return beginningOfEndBufferTimeMs;
+        }
+        // Don't let it violate the minimum refresh period
+        long endOfMinRefreshBufferTime = relativeToMs + 1000 * refreshMinPeriodSeconds;
+        if (proposedRefreshMs < endOfMinRefreshBufferTime) {
+            log.info(
+                    "[Principal={}]: Expiring credential re-login thread time adjusted from {} to {} since the former is sooner "
+                            + "than the minimum refresh interval ({} seconds from now).",
+                    principalLogText(), new Date(proposedRefreshMs), new Date(endOfMinRefreshBufferTime),
+                    refreshMinPeriodSeconds);
+            return endOfMinRefreshBufferTime;
+        }
+        // Proposed refresh time doesn't violate any constraints
+        return proposedRefreshMs;
+    }
+
+    private void reLogin() throws LoginException, ExitRefresherThreadDueToIllegalStateException {
+        synchronized (mandatoryClassToSynchronizeOnPriorToRefresh) {
+            // Only perform one refresh of a particular type at a time
+            boolean logoutRequiredBeforeLoggingBackIn = isLogoutRequiredBeforeLoggingBackIn();
+            if (hasExpiringCredential && logoutRequiredBeforeLoggingBackIn) {
+                String principalLogTextPriorToLogout = principalLogText();
+                log.info("Initiating logout for {}", principalLogTextPriorToLogout);
+                loginContext.logout();
+                // Make absolutely sure we were logged out
+                expiringCredential = expiringCredential();
+                hasExpiringCredential = expiringCredential != null;
+                if (hasExpiringCredential)
+                    // We can't force the removal because we don't know how to do it, so abort
+                    throw new ExitRefresherThreadDueToIllegalStateException(String.format(
+                            "Subject's private credentials still contains an instance of %s even though logout() was invoked; exiting refresh thread",
+                            expiringCredential.getClass().getName()));
+            }
+            /*
+             * Perform a login, making note of any credential that might need a logout()
+             * afterwards
+             */
+            ExpiringCredential optionalCredentialToLogout = expiringCredential;
+            LoginContext optionalLoginContextToLogout = loginContext;
+            boolean cleanLogin = false; // remember to restore the original if necessary
+            try {
+                loginContext = loginContextFactory.createLoginContext(ExpiringCredentialRefreshingLogin.this);
+                log.info("Initiating re-login for {}, logout() still needs to be called on a previous login = {}",
+                        principalName, optionalCredentialToLogout != null);
+                loginContext.login();
+                cleanLogin = true; // no need to restore the original
+                // Perform a logout() on any original credential if necessary
+                if (optionalCredentialToLogout != null)
+                    optionalLoginContextToLogout.logout();
+            } finally {
+                if (!cleanLogin)
+                    // restore the original
+                    loginContext = optionalLoginContextToLogout;
+            }
+            /*
+             * Get the new credential and make sure it is not any old one that required a
+             * logout() after the login()
+             */
+            expiringCredential = expiringCredential();
+            hasExpiringCredential = expiringCredential != null;
+            if (!hasExpiringCredential) {
+                /*
+                 * Re-login has failed because our login() invocation has not generated a
+                 * credential but has also not generated an exception. We won't exit here;
+                 * instead we will allow login retries in case we can somehow fix the issue (it
+                 * seems likely to be a bug, but it doesn't hurt to keep trying to refresh).
+                 */
+                log.error("No Expiring Credential after a supposedly-successful re-login");
+                principalName = null;
+            } else {
+                if (expiringCredential == optionalCredentialToLogout)
+                    /*
+                     * The login() didn't identify a new credential; we still have the old one. We
+                     * don't know how to fix this, so abort.
+                     */
+                    throw new ExitRefresherThreadDueToIllegalStateException(String.format(
+                            "Subject's private credentials still contains the previous, soon-to-expire instance of %s even though login() followed by logout() was invoked; exiting refresh thread",
+                            expiringCredential.getClass().getName()));
+                principalName = expiringCredential.principalName();
+                if (log.isDebugEnabled())
+                    log.debug("[Principal={}]: It is an expiring credential after re-login as expected",
+                            principalLogText());
+            }
+        }
+    }
+
+    private String principalLogText() {
+        return expiringCredential == null ? principalName
+                : expiringCredential.getClass().getSimpleName() + ":" + principalName;
+    }
+
+    private long currentMs() {
+        return time.milliseconds();
+    }
+
+    private boolean isLogoutRequiredBeforeLoggingBackIn() {
+        return !expiringCredentialRefreshConfig.loginRefreshReloginAllowedBeforeLogout();
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerConfigException.java b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerConfigException.java
new file mode 100644
index 0000000000000..3dcdcd00afc57
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerConfigException.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.security.oauthbearer.internals.unsecured;
+
+import org.apache.kafka.common.KafkaException;
+
+/**
+ * Exception thrown when there is a problem with the configuration (an invalid
+ * option in a JAAS config, for example).
+ */
+public class OAuthBearerConfigException extends KafkaException {
+    private static final long serialVersionUID = -8056105648062343518L;
+
+    public OAuthBearerConfigException(String s) {
+        super(s);
+    }
+
+    public OAuthBearerConfigException(String message, Throwable cause) {
+        super(message, cause);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerIllegalTokenException.java b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerIllegalTokenException.java
new file mode 100644
index 0000000000000..78859003f7e42
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerIllegalTokenException.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 org.apache.kafka.common.security.oauthbearer.internals.unsecured;
+
+import java.util.Objects;
+
+import org.apache.kafka.common.KafkaException;
+
+/**
+ * Exception thrown when token validation fails due to a problem with the token
+ * itself (as opposed to a missing remote resource or a configuration problem)
+ */
+public class OAuthBearerIllegalTokenException extends KafkaException {
+    private static final long serialVersionUID = -5275276640051316350L;
+    private final OAuthBearerValidationResult reason;
+
+    /**
+     * Constructor
+     * 
+     * @param reason
+     *            the mandatory reason for the validation failure; it must indicate
+     *            failure
+     */
+    public OAuthBearerIllegalTokenException(OAuthBearerValidationResult reason) {
+        super(Objects.requireNonNull(reason).failureDescription());
+        if (reason.success())
+            throw new IllegalArgumentException("The reason indicates success; it must instead indicate failure");
+        this.reason = reason;
+    }
+
+    /**
+     * Return the (always non-null) reason for the validation failure
+     * 
+     * @return the reason for the validation failure
+     */
+    public OAuthBearerValidationResult reason() {
+        return reason;
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerScopeUtils.java b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerScopeUtils.java
new file mode 100644
index 0000000000000..7cae41a804887
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerScopeUtils.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.security.oauthbearer.internals.unsecured;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Objects;
+import java.util.regex.Pattern;
+
+/**
+ * Utility class for help dealing with
+ * Access Token
+ * Scopes
+ */
+public class OAuthBearerScopeUtils {
+    private static final Pattern INDIVIDUAL_SCOPE_ITEM_PATTERN = Pattern.compile("[\\x23-\\x5B\\x5D-\\x7E\\x21]+");
+
+    /**
+     * Return true if the given value meets the definition of a valid scope item as
+     * per RFC 6749
+     * Section 3.3, otherwise false
+     * 
+     * @param scopeItem
+     *            the mandatory scope item to check for validity
+     * @return true if the given value meets the definition of a valid scope item,
+     *         otherwise false
+     */
+    public static boolean isValidScopeItem(String scopeItem) {
+        return INDIVIDUAL_SCOPE_ITEM_PATTERN.matcher(Objects.requireNonNull(scopeItem)).matches();
+    }
+
+    /**
+     * Convert a space-delimited list of scope values (for example,
+     * "scope1 scope2") to a List containing the individual elements
+     * ("scope1" and "scope2")
+     * 
+     * @param spaceDelimitedScope
+     *            the mandatory (but possibly empty) space-delimited scope values,
+     *            each of which must be valid according to
+     *            {@link #isValidScopeItem(String)}
+     * @return the list of the given (possibly empty) space-delimited values
+     * @throws OAuthBearerConfigException
+     *             if any of the individual scope values are malformed/illegal
+     */
+    public static List parseScope(String spaceDelimitedScope) throws OAuthBearerConfigException {
+        List retval = new ArrayList<>();
+        for (String individualScopeItem : Objects.requireNonNull(spaceDelimitedScope).split(" ")) {
+            if (!individualScopeItem.isEmpty()) {
+                if (!isValidScopeItem(individualScopeItem))
+                    throw new OAuthBearerConfigException(String.format("Invalid scope value: %s", individualScopeItem));
+                retval.add(individualScopeItem);
+            }
+        }
+        return Collections.unmodifiableList(retval);
+    }
+
+    private OAuthBearerScopeUtils() {
+        // empty
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerUnsecuredJws.java b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerUnsecuredJws.java
new file mode 100644
index 0000000000000..b5b301650ef4f
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerUnsecuredJws.java
@@ -0,0 +1,369 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.oauthbearer.internals.unsecured;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Base64;
+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.Map.Entry;
+import java.util.Objects;
+import java.util.Set;
+
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerToken;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.node.JsonNodeType;
+
+/**
+ * A simple unsecured JWS implementation. The '{@code nbf}' claim is ignored if
+ * it is given because the related logic is not required for Kafka testing and
+ * development purposes.
+ *
+ * @see RFC 7515
+ */
+public class OAuthBearerUnsecuredJws implements OAuthBearerToken {
+    private final String compactSerialization;
+    private final List splits;
+    private final Map header;
+    private final String principalClaimName;
+    private final String scopeClaimName;
+    private final Map claims;
+    private final Set scope;
+    private final long lifetime;
+    private final String principalName;
+    private final Long startTimeMs;
+
+    /**
+     * Constructor with the given principal and scope claim names
+     *
+     * @param compactSerialization
+     *            the compact serialization to parse as an unsecured JWS
+     * @param principalClaimName
+     *            the required principal claim name
+     * @param scopeClaimName
+     *            the required scope claim name
+     * @throws OAuthBearerIllegalTokenException
+     *             if the compact serialization is not a valid unsecured JWS
+     *             (meaning it did not have 3 dot-separated Base64URL sections
+     *             without an empty digital signature; or the header or claims
+     *             either are not valid Base 64 URL encoded values or are not JSON
+     *             after decoding; or the mandatory '{@code alg}' header value is
+     *             not "{@code none}")
+     */
+    public OAuthBearerUnsecuredJws(String compactSerialization, String principalClaimName, String scopeClaimName)
+            throws OAuthBearerIllegalTokenException {
+        this.compactSerialization = Objects.requireNonNull(compactSerialization);
+        if (compactSerialization.contains(".."))
+            throw new OAuthBearerIllegalTokenException(
+                    OAuthBearerValidationResult.newFailure("Malformed compact serialization contains '..'"));
+        this.splits = extractCompactSerializationSplits();
+        this.header = toMap(splits().get(0));
+        String claimsSplit = splits.get(1);
+        this.claims = toMap(claimsSplit);
+        String alg = Objects.requireNonNull(header().get("alg"), "JWS header must have an Algorithm value").toString();
+        if (!"none".equals(alg))
+            throw new OAuthBearerIllegalTokenException(
+                    OAuthBearerValidationResult.newFailure("Unsecured JWS must have 'none' for an algorithm"));
+        String digitalSignatureSplit = splits.get(2);
+        if (!digitalSignatureSplit.isEmpty())
+            throw new OAuthBearerIllegalTokenException(
+                    OAuthBearerValidationResult.newFailure("Unsecured JWS must not contain a digital signature"));
+        this.principalClaimName = Objects.requireNonNull(principalClaimName).trim();
+        if (this.principalClaimName.isEmpty())
+            throw new IllegalArgumentException("Must specify a non-blank principal claim name");
+        this.scopeClaimName = Objects.requireNonNull(scopeClaimName).trim();
+        if (this.scopeClaimName.isEmpty())
+            throw new IllegalArgumentException("Must specify a non-blank scope claim name");
+        this.scope = calculateScope();
+        Number expirationTimeSeconds = expirationTime();
+        if (expirationTimeSeconds == null)
+            throw new OAuthBearerIllegalTokenException(
+                    OAuthBearerValidationResult.newFailure("No expiration time in JWT"));
+        lifetime = convertClaimTimeInSecondsToMs(expirationTimeSeconds);
+        String principalName = claim(this.principalClaimName, String.class);
+        if (principalName == null || principalName.trim().isEmpty())
+            throw new OAuthBearerIllegalTokenException(OAuthBearerValidationResult
+                    .newFailure("No principal name in JWT claim: " + this.principalClaimName));
+        this.principalName = principalName;
+        this.startTimeMs = calculateStartTimeMs();
+    }
+
+    @Override
+    public String value() {
+        return compactSerialization;
+    }
+
+    /**
+     * Return the 3 or 5 dot-separated sections of the JWT compact serialization
+     *
+     * @return the 3 or 5 dot-separated sections of the JWT compact serialization
+     */
+    public List splits() {
+        return splits;
+    }
+
+    /**
+     * Return the JOSE Header as a {@code Map}
+     *
+     * @return the JOSE header
+     */
+    public Map header() {
+        return header;
+    }
+
+    @Override
+    public String principalName() {
+        return principalName;
+    }
+
+    @Override
+    public Long startTimeMs() {
+        return startTimeMs;
+    }
+
+    @Override
+    public long lifetimeMs() {
+        return lifetime;
+    }
+
+    @Override
+    public Set scope() throws OAuthBearerIllegalTokenException {
+        return scope;
+    }
+
+    /**
+     * Return the JWT Claim Set as a {@code Map}
+     *
+     * @return the (always non-null but possibly empty) claims
+     */
+    public Map claims() {
+        return claims;
+    }
+
+    /**
+     * Return the (always non-null/non-empty) principal claim name
+     *
+     * @return the (always non-null/non-empty) principal claim name
+     */
+    public String principalClaimName() {
+        return principalClaimName;
+    }
+
+    /**
+     * Return the (always non-null/non-empty) scope claim name
+     *
+     * @return the (always non-null/non-empty) scope claim name
+     */
+    public String scopeClaimName() {
+        return scopeClaimName;
+    }
+
+    /**
+     * Indicate if the claim exists and is the given type
+     *
+     * @param claimName
+     *            the mandatory JWT claim name
+     * @param type
+     *            the mandatory type, which should either be String.class,
+     *            Number.class, or List.class
+     * @return true if the claim exists and is the given type, otherwise false
+     */
+    public boolean isClaimType(String claimName, Class type) {
+        Object value = rawClaim(claimName);
+        Objects.requireNonNull(type);
+        if (value == null)
+            return false;
+        if (type == String.class && value instanceof String)
+            return true;
+        if (type == Number.class && value instanceof Number)
+            return true;
+        return type == List.class && value instanceof List;
+    }
+
+    /**
+     * Extract a claim of the given type
+     *
+     * @param claimName
+     *            the mandatory JWT claim name
+     * @param type
+     *            the mandatory type, which must either be String.class,
+     *            Number.class, or List.class
+     * @return the claim if it exists, otherwise null
+     * @throws OAuthBearerIllegalTokenException
+     *             if the claim exists but is not the given type
+     */
+    public  T claim(String claimName, Class type) throws OAuthBearerIllegalTokenException {
+        Object value = rawClaim(claimName);
+        try {
+            return Objects.requireNonNull(type).cast(value);
+        } catch (ClassCastException e) {
+            throw new OAuthBearerIllegalTokenException(
+                    OAuthBearerValidationResult.newFailure(String.format("The '%s' claim was not of type %s: %s",
+                            claimName, type.getSimpleName(), value.getClass().getSimpleName())));
+        }
+    }
+
+    /**
+     * Extract a claim in its raw form
+     *
+     * @param claimName
+     *            the mandatory JWT claim name
+     * @return the raw claim value, if it exists, otherwise null
+     */
+    public Object rawClaim(String claimName) {
+        return claims().get(Objects.requireNonNull(claimName));
+    }
+
+    /**
+     * Return the
+     * Expiration
+     * Time claim
+     *
+     * @return the Expiration
+     *         Time claim if available, otherwise null
+     * @throws OAuthBearerIllegalTokenException
+     *             if the claim value is the incorrect type
+     */
+    public Number expirationTime() throws OAuthBearerIllegalTokenException {
+        return claim("exp", Number.class);
+    }
+
+    /**
+     * Return the Issued
+     * At claim
+     *
+     * @return the
+     *         Issued
+     *         At claim if available, otherwise null
+     * @throws OAuthBearerIllegalTokenException
+     *             if the claim value is the incorrect type
+     */
+    public Number issuedAt() throws OAuthBearerIllegalTokenException {
+        return claim("iat", Number.class);
+    }
+
+    /**
+     * Return the
+     * Subject claim
+     *
+     * @return the Subject claim
+     *         if available, otherwise null
+     * @throws OAuthBearerIllegalTokenException
+     *             if the claim value is the incorrect type
+     */
+    public String subject() throws OAuthBearerIllegalTokenException {
+        return claim("sub", String.class);
+    }
+
+    /**
+     * Decode the given Base64URL-encoded value, parse the resulting JSON as a JSON
+     * object, and return the map of member names to their values (each value being
+     * represented as either a String, a Number, or a List of Strings).
+     *
+     * @param split
+     *            the value to decode and parse
+     * @return the map of JSON member names to their String, Number, or String List
+     *         value
+     * @throws OAuthBearerIllegalTokenException
+     *             if the given Base64URL-encoded value cannot be decoded or parsed
+     */
+    public static Map toMap(String split) throws OAuthBearerIllegalTokenException {
+        Map retval = new HashMap<>();
+        try {
+            byte[] decode = Base64.getDecoder().decode(split);
+            JsonNode jsonNode = new ObjectMapper().readTree(decode);
+            if (jsonNode == null)
+                throw new OAuthBearerIllegalTokenException(OAuthBearerValidationResult.newFailure("malformed JSON"));
+            for (Iterator> iterator = jsonNode.fields(); iterator.hasNext();) {
+                Entry entry = iterator.next();
+                retval.put(entry.getKey(), convert(entry.getValue()));
+            }
+            return Collections.unmodifiableMap(retval);
+        } catch (IllegalArgumentException e) {
+            // potentially thrown by java.util.Base64.Decoder implementations
+            throw new OAuthBearerIllegalTokenException(
+                    OAuthBearerValidationResult.newFailure("malformed Base64 URL encoded value"));
+        } catch (IOException e) {
+            throw new OAuthBearerIllegalTokenException(OAuthBearerValidationResult.newFailure("malformed JSON"));
+        }
+    }
+
+    private List extractCompactSerializationSplits() {
+        List tmpSplits = new ArrayList<>(Arrays.asList(compactSerialization.split("\\.")));
+        if (compactSerialization.endsWith("."))
+            tmpSplits.add("");
+        if (tmpSplits.size() != 3)
+            throw new OAuthBearerIllegalTokenException(OAuthBearerValidationResult.newFailure(
+                    "Unsecured JWS compact serializations must have 3 dot-separated Base64URL-encoded values"));
+        return Collections.unmodifiableList(tmpSplits);
+    }
+
+    private static Object convert(JsonNode value) {
+        if (value.isArray()) {
+            List retvalList = new ArrayList<>();
+            for (JsonNode arrayElement : value)
+                retvalList.add(arrayElement.asText());
+            return retvalList;
+        }
+        return value.getNodeType() == JsonNodeType.NUMBER ? value.numberValue() : value.asText();
+    }
+
+    private Long calculateStartTimeMs() throws OAuthBearerIllegalTokenException {
+        Number issuedAtSeconds = claim("iat", Number.class);
+        return issuedAtSeconds == null ? null : convertClaimTimeInSecondsToMs(issuedAtSeconds);
+    }
+
+    private static long convertClaimTimeInSecondsToMs(Number claimValue) {
+        return Math.round(claimValue.doubleValue() * 1000);
+    }
+
+    private Set calculateScope() {
+        String scopeClaimName = scopeClaimName();
+        if (isClaimType(scopeClaimName, String.class)) {
+            String scopeClaimValue = claim(scopeClaimName, String.class);
+            if (scopeClaimValue.trim().isEmpty())
+                return Collections.emptySet();
+            else {
+                Set retval = new HashSet<>();
+                retval.add(scopeClaimValue.trim());
+                return Collections.unmodifiableSet(retval);
+            }
+        }
+        List scopeClaimValue = claim(scopeClaimName, List.class);
+        if (scopeClaimValue == null || scopeClaimValue.isEmpty())
+            return Collections.emptySet();
+        @SuppressWarnings("unchecked")
+        List stringList = (List) scopeClaimValue;
+        Set retval = new HashSet<>();
+        for (String scope : stringList) {
+            if (scope != null && !scope.trim().isEmpty()) {
+                retval.add(scope.trim());
+            }
+        }
+        return Collections.unmodifiableSet(retval);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerUnsecuredLoginCallbackHandler.java b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerUnsecuredLoginCallbackHandler.java
new file mode 100644
index 0000000000000..e7a4f2cc798d1
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerUnsecuredLoginCallbackHandler.java
@@ -0,0 +1,346 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.oauthbearer.internals.unsecured;
+
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
+import java.util.Base64;
+import java.util.Base64.Encoder;
+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.Set;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+
+import javax.security.auth.callback.Callback;
+import javax.security.auth.callback.UnsupportedCallbackException;
+import javax.security.auth.login.AppConfigurationEntry;
+import javax.security.sasl.SaslException;
+
+import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.config.ConfigException;
+import org.apache.kafka.common.security.auth.AuthenticateCallbackHandler;
+import org.apache.kafka.common.security.auth.SaslExtensionsCallback;
+import org.apache.kafka.common.security.auth.SaslExtensions;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerLoginModule;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerTokenCallback;
+import org.apache.kafka.common.security.oauthbearer.internals.OAuthBearerClientInitialResponse;
+import org.apache.kafka.common.utils.Time;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@code CallbackHandler} that recognizes {@link OAuthBearerTokenCallback}
+ * to return an unsecured OAuth 2 bearer token and {@link SaslExtensionsCallback} to return SASL extensions
+ * 
        
+ * Claims and their values on the returned token can be specified using
+ * {@code unsecuredLoginStringClaim_},
+ * {@code unsecuredLoginNumberClaim_}, and
+ * {@code unsecuredLoginListClaim_} options. The first character of
+ * the value is taken as the delimiter for list claims. You may define any claim
+ * name and value except '{@code iat}' and '{@code exp}', both of which are
+ * calculated automatically.
+ * 
        
+ * 
        
+ * You can also add custom unsecured SASL extensions using
+ * {@code unsecuredLoginExtension_}. Extension keys and values are subject to regex validation.
+ * The extension key must also not be equal to the reserved key {@link OAuthBearerClientInitialResponse#AUTH_KEY}
+ * 
        
+ * This implementation also accepts the following options:
+ * 
          
+ * 
        • {@code unsecuredLoginPrincipalClaimName} set to a custom claim name if
+ * you wish the name of the String claim holding the principal name to be
+ * something other than '{@code sub}'.
          • 
+ * 
        • {@code unsecuredLoginLifetimeSeconds} set to an integer value if the
+ * token expiration is to be set to something other than the default value of
+ * 3600 seconds (which is 1 hour). The '{@code exp}' claim reflects the
+ * expiration time.
          • 
+ * 
        • {@code unsecuredLoginScopeClaimName} set to a custom claim name if you
+ * wish the name of the String or String List claim holding any token scope to
+ * be something other than '{@code scope}'
          • 
+ * 
        
+ * For example:
+ *
+ * 
        + * KafkaClient {
+ *      org.apache.kafka.common.security.oauthbearer.OAuthBearerLoginModule Required
+ *      unsecuredLoginStringClaim_sub="thePrincipalName"
+ *      unsecuredLoginListClaim_scope="|scopeValue1|scopeValue2"
+ *      unsecuredLoginLifetimeSeconds="60"
+ *      unsecuredLoginExtension_traceId="123";
+ * };
+ * 
        
+ *
+ * This class is the default when the SASL mechanism is OAUTHBEARER and no value
+ * is explicitly set via either the {@code sasl.login.callback.handler.class}
+ * client configuration property or the
+ * {@code listener.name.sasl_[plaintext|ssl].oauthbearer.sasl.login.callback.handler.class}
+ * broker configuration property.
+ */
+public class OAuthBearerUnsecuredLoginCallbackHandler implements AuthenticateCallbackHandler {
+    private final Logger log = LoggerFactory.getLogger(OAuthBearerUnsecuredLoginCallbackHandler.class);
+    private static final String OPTION_PREFIX = "unsecuredLogin";
+    private static final String PRINCIPAL_CLAIM_NAME_OPTION = OPTION_PREFIX + "PrincipalClaimName";
+    private static final String LIFETIME_SECONDS_OPTION = OPTION_PREFIX + "LifetimeSeconds";
+    private static final String SCOPE_CLAIM_NAME_OPTION = OPTION_PREFIX + "ScopeClaimName";
+    private static final Set RESERVED_CLAIMS = Collections
+            .unmodifiableSet(new HashSet<>(Arrays.asList("iat", "exp")));
+    private static final String DEFAULT_PRINCIPAL_CLAIM_NAME = "sub";
+    private static final String DEFAULT_LIFETIME_SECONDS_ONE_HOUR = "3600";
+    private static final String DEFAULT_SCOPE_CLAIM_NAME = "scope";
+    private static final String STRING_CLAIM_PREFIX = OPTION_PREFIX + "StringClaim_";
+    private static final String NUMBER_CLAIM_PREFIX = OPTION_PREFIX + "NumberClaim_";
+    private static final String LIST_CLAIM_PREFIX = OPTION_PREFIX + "ListClaim_";
+    private static final String EXTENSION_PREFIX = OPTION_PREFIX + "Extension_";
+    private static final String QUOTE = "\"";
+    private Time time = Time.SYSTEM;
+    private Map moduleOptions = null;
+    private boolean configured = false;
+
+    private static final Pattern DOUBLEQUOTE = Pattern.compile("\"", Pattern.LITERAL);
+
+    private static final Pattern BACKSLASH = Pattern.compile("\\", Pattern.LITERAL);
+
+    /**
+     * For testing
+     *
+     * @param time
+     *            the mandatory time to set
+     */
+    void time(Time time) {
+        this.time = Objects.requireNonNull(time);
+    }
+
+    /**
+     * Return true if this instance has been configured, otherwise false
+     *
+     * @return true if this instance has been configured, otherwise false
+     */
+    public boolean configured() {
+        return configured;
+    }
+
+    @SuppressWarnings("unchecked")
+    @Override
+    public void configure(Map configs, String saslMechanism, List jaasConfigEntries) {
+        if (!OAuthBearerLoginModule.OAUTHBEARER_MECHANISM.equals(saslMechanism))
+            throw new IllegalArgumentException(String.format("Unexpected SASL mechanism: %s", saslMechanism));
+        if (Objects.requireNonNull(jaasConfigEntries).size() != 1 || jaasConfigEntries.get(0) == null)
+            throw new IllegalArgumentException(
+                    String.format("Must supply exactly 1 non-null JAAS mechanism configuration (size was %d)",
+                            jaasConfigEntries.size()));
+        this.moduleOptions = Collections.unmodifiableMap((Map) jaasConfigEntries.get(0).getOptions());
+        configured = true;
+    }
+
+    @Override
+    public void handle(Callback[] callbacks) throws IOException, UnsupportedCallbackException {
+        if (!configured())
+            throw new IllegalStateException("Callback handler not configured");
+        for (Callback callback : callbacks) {
+            if (callback instanceof OAuthBearerTokenCallback)
+                try {
+                    handleTokenCallback((OAuthBearerTokenCallback) callback);
+                } catch (KafkaException e) {
+                    throw new IOException(e.getMessage(), e);
+                }
+            else if (callback instanceof SaslExtensionsCallback)
+                try {
+                    handleExtensionsCallback((SaslExtensionsCallback) callback);
+                } catch (KafkaException e) {
+                    throw new IOException(e.getMessage(), e);
+                }
+            else
+                throw new UnsupportedCallbackException(callback);
+        }
+    }
+
+    @Override
+    public void close() {
+        // empty
+    }
+
+    private void handleTokenCallback(OAuthBearerTokenCallback callback) {
+        if (callback.token() != null)
+            throw new IllegalArgumentException("Callback had a token already");
+        if (moduleOptions.isEmpty()) {
+            log.debug("Token not provided, this login cannot be used to establish client connections");
+            callback.token(null);
+            return;
+        }
+        if (moduleOptions.keySet().stream().noneMatch(name -> !name.startsWith(EXTENSION_PREFIX))) {
+            throw new OAuthBearerConfigException("Extensions provided in login context without a token");
+        }
+        String principalClaimNameValue = optionValue(PRINCIPAL_CLAIM_NAME_OPTION);
+        String principalClaimName = principalClaimNameValue != null && !principalClaimNameValue.trim().isEmpty()
+                ? principalClaimNameValue.trim()
+                : DEFAULT_PRINCIPAL_CLAIM_NAME;
+        String scopeClaimNameValue = optionValue(SCOPE_CLAIM_NAME_OPTION);
+        String scopeClaimName = scopeClaimNameValue != null && !scopeClaimNameValue.trim().isEmpty()
+                ? scopeClaimNameValue.trim()
+                : DEFAULT_SCOPE_CLAIM_NAME;
+        String headerJson = "{" + claimOrHeaderJsonText("alg", "none") + "}";
+        String lifetimeSecondsValueToUse = optionValue(LIFETIME_SECONDS_OPTION, DEFAULT_LIFETIME_SECONDS_ONE_HOUR);
+        String claimsJson;
+        try {
+            claimsJson = String.format("{%s,%s%s}", expClaimText(Long.parseLong(lifetimeSecondsValueToUse)),
+                    claimOrHeaderJsonText("iat", time.milliseconds() / 1000.0),
+                    commaPrependedStringNumberAndListClaimsJsonText());
+        } catch (NumberFormatException e) {
+            throw new OAuthBearerConfigException(e.getMessage());
+        }
+        try {
+            Encoder urlEncoderNoPadding = Base64.getUrlEncoder().withoutPadding();
+            OAuthBearerUnsecuredJws jws = new OAuthBearerUnsecuredJws(
+                    String.format("%s.%s.",
+                            urlEncoderNoPadding.encodeToString(headerJson.getBytes(StandardCharsets.UTF_8)),
+                            urlEncoderNoPadding.encodeToString(claimsJson.getBytes(StandardCharsets.UTF_8))),
+                    principalClaimName, scopeClaimName);
+            log.info("Retrieved token with principal {}", jws.principalName());
+            callback.token(jws);
+        } catch (OAuthBearerIllegalTokenException e) {
+            // occurs if the principal claim doesn't exist or has an empty value
+            throw new OAuthBearerConfigException(e.getMessage(), e);
+        }
+    }
+
+    /**
+     *  Add and validate all the configured extensions.
+     *  Token keys, apart from passing regex validation, must not be equal to the reserved key {@link OAuthBearerClientInitialResponse#AUTH_KEY}
+     */
+    private void handleExtensionsCallback(SaslExtensionsCallback callback) {
+        Map extensions = new HashMap<>();
+        for (Map.Entry configEntry : this.moduleOptions.entrySet()) {
+            String key = configEntry.getKey();
+            if (!key.startsWith(EXTENSION_PREFIX))
+                continue;
+
+            extensions.put(key.substring(EXTENSION_PREFIX.length()), configEntry.getValue());
+        }
+
+        SaslExtensions saslExtensions = new SaslExtensions(extensions);
+        try {
+            OAuthBearerClientInitialResponse.validateExtensions(saslExtensions);
+        } catch (SaslException e) {
+            throw new ConfigException(e.getMessage());
+        }
+
+        callback.extensions(saslExtensions);
+    }
+
+    private String commaPrependedStringNumberAndListClaimsJsonText() throws OAuthBearerConfigException {
+        StringBuilder sb = new StringBuilder();
+        for (String key : moduleOptions.keySet()) {
+            if (key.startsWith(STRING_CLAIM_PREFIX) && key.length() > STRING_CLAIM_PREFIX.length())
+                sb.append(',').append(claimOrHeaderJsonText(
+                        confirmNotReservedClaimName(key.substring(STRING_CLAIM_PREFIX.length())), optionValue(key)));
+            else if (key.startsWith(NUMBER_CLAIM_PREFIX) && key.length() > NUMBER_CLAIM_PREFIX.length())
+                sb.append(',')
+                        .append(claimOrHeaderJsonText(
+                                confirmNotReservedClaimName(key.substring(NUMBER_CLAIM_PREFIX.length())),
+                                Double.valueOf(optionValue(key))));
+            else if (key.startsWith(LIST_CLAIM_PREFIX) && key.length() > LIST_CLAIM_PREFIX.length())
+                sb.append(',')
+                        .append(claimOrHeaderJsonArrayText(
+                                confirmNotReservedClaimName(key.substring(LIST_CLAIM_PREFIX.length())),
+                                listJsonText(optionValue(key))));
+        }
+        return sb.toString();
+    }
+
+    private String confirmNotReservedClaimName(String claimName) throws OAuthBearerConfigException {
+        if (RESERVED_CLAIMS.contains(claimName))
+            throw new OAuthBearerConfigException(String.format("Cannot explicitly set the '%s' claim", claimName));
+        return claimName;
+    }
+
+    private String listJsonText(String value) {
+        if (value.isEmpty() || value.length() <= 1)
+            return "[]";
+        String delimiter;
+        String unescapedDelimiterChar = value.substring(0, 1);
+        switch (unescapedDelimiterChar) {
+            case "\\":
+            case ".":
+            case "[":
+            case "(":
+            case "{":
+            case "|":
+            case "^":
+            case "$":
+                delimiter = "\\" + unescapedDelimiterChar;
+                break;
+            default:
+                delimiter = unescapedDelimiterChar;
+                break;
+        }
+        String listText = value.substring(1);
+        String[] elements = listText.split(delimiter);
+        StringBuilder sb = new StringBuilder();
+        for (String element : elements) {
+            sb.append(sb.length() == 0 ? '[' : ',');
+            sb.append('"').append(escape(element)).append('"');
+        }
+        if (listText.startsWith(unescapedDelimiterChar) || listText.endsWith(unescapedDelimiterChar)
+                || listText.contains(unescapedDelimiterChar + unescapedDelimiterChar))
+            sb.append(",\"\"");
+        return sb.append(']').toString();
+    }
+
+    private String optionValue(String key) {
+        return optionValue(key, null);
+    }
+
+    private String optionValue(String key, String defaultValue) {
+        String explicitValue = option(key);
+        return explicitValue != null ? explicitValue : defaultValue;
+    }
+
+    private String option(String key) {
+        if (!configured)
+            throw new IllegalStateException("Callback handler not configured");
+        return moduleOptions.get(Objects.requireNonNull(key));
+    }
+
+    private String claimOrHeaderJsonText(String claimName, Number claimValue) {
+        return QUOTE + escape(claimName) + QUOTE + ":" + claimValue;
+    }
+
+    private String claimOrHeaderJsonText(String claimName, String claimValue) {
+        return QUOTE + escape(claimName) + QUOTE + ":" + QUOTE + escape(claimValue) + QUOTE;
+    }
+
+    private String claimOrHeaderJsonArrayText(String claimName, String escapedClaimValue) {
+        if (!escapedClaimValue.startsWith("[") || !escapedClaimValue.endsWith("]"))
+            throw new IllegalArgumentException(String.format("Illegal JSON array: %s", escapedClaimValue));
+        return QUOTE + escape(claimName) + QUOTE + ":" + escapedClaimValue;
+    }
+
+    private String escape(String jsonStringValue) {
+        String replace1 = DOUBLEQUOTE.matcher(jsonStringValue).replaceAll(Matcher.quoteReplacement("\\\""));
+        return BACKSLASH.matcher(replace1).replaceAll(Matcher.quoteReplacement("\\\\"));
+    }
+
+    private String expClaimText(long lifetimeSeconds) {
+        return claimOrHeaderJsonText("exp", time.milliseconds() / 1000.0 + lifetimeSeconds);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerUnsecuredValidatorCallbackHandler.java b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerUnsecuredValidatorCallbackHandler.java
new file mode 100644
index 0000000000000..54e289c2c5621
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerUnsecuredValidatorCallbackHandler.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.security.oauthbearer.internals.unsecured;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+
+import javax.security.auth.callback.Callback;
+import javax.security.auth.callback.UnsupportedCallbackException;
+import javax.security.auth.login.AppConfigurationEntry;
+
+import org.apache.kafka.common.security.auth.AuthenticateCallbackHandler;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerExtensionsValidatorCallback;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerLoginModule;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerValidatorCallback;
+import org.apache.kafka.common.utils.Time;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@code CallbackHandler} that recognizes
+ * {@link OAuthBearerValidatorCallback} and validates an unsecured OAuth 2
+ * bearer token. It requires there to be an "exp" (Expiration Time)
+ * claim of type Number. If "iat" (Issued At) or
+ * "nbf" (Not Before) claims are present each must be a number that
+ * precedes the Expiration Time claim, and if both are present the Not Before
+ * claim must not precede the Issued At claim. It also accepts the following
+ * options, none of which are required:
+ * 
          
+ * 
        • {@code unsecuredValidatorPrincipalClaimName} set to a non-empty value if
+ * you wish a particular String claim holding a principal name to be checked for
+ * existence; the default is to check for the existence of the '{@code sub}'
+ * claim
          • 
+ * 
        • {@code unsecuredValidatorScopeClaimName} set to a custom claim name if
+ * you wish the name of the String or String List claim holding any token scope
+ * to be something other than '{@code scope}'
          • 
+ * 
        • {@code unsecuredValidatorRequiredScope} set to a space-delimited list of
+ * scope values if you wish the String/String List claim holding the token scope
+ * to be checked to make sure it contains certain values
          • 
+ * 
        • {@code unsecuredValidatorAllowableClockSkewMs} set to a positive integer
+ * value if you wish to allow up to some number of positive milliseconds of
+ * clock skew (the default is 0)
          • 
+ * 
            
+ * For example:
+ *
+ * 
            + * KafkaServer {
+ *      org.apache.kafka.common.security.oauthbearer.OAuthBearerLoginModule Required
+ *      unsecuredLoginStringClaim_sub="thePrincipalName"
+ *      unsecuredLoginListClaim_scope=",KAFKA_BROKER,LOGIN_TO_KAFKA"
+ *      unsecuredValidatorRequiredScope="LOGIN_TO_KAFKA"
+ *      unsecuredValidatorAllowableClockSkewMs="3000";
+ * };
+ * 
            
+ * It also recognizes {@link OAuthBearerExtensionsValidatorCallback} and validates every extension passed to it.
+ *
+ * This class is the default when the SASL mechanism is OAUTHBEARER and no value
+ * is explicitly set via the
+ * {@code listener.name.sasl_[plaintext|ssl].oauthbearer.sasl.server.callback.handler.class}
+ * broker configuration property.
+ * It is worth noting that this class is not suitable for production use due to the use of unsecured JWT tokens and
+ * validation of every given extension.
+ */
+public class OAuthBearerUnsecuredValidatorCallbackHandler implements AuthenticateCallbackHandler {
+    private static final Logger log = LoggerFactory.getLogger(OAuthBearerUnsecuredValidatorCallbackHandler.class);
+    private static final String OPTION_PREFIX = "unsecuredValidator";
+    private static final String PRINCIPAL_CLAIM_NAME_OPTION = OPTION_PREFIX + "PrincipalClaimName";
+    private static final String SCOPE_CLAIM_NAME_OPTION = OPTION_PREFIX + "ScopeClaimName";
+    private static final String REQUIRED_SCOPE_OPTION = OPTION_PREFIX + "RequiredScope";
+    private static final String ALLOWABLE_CLOCK_SKEW_MILLIS_OPTION = OPTION_PREFIX + "AllowableClockSkewMs";
+    private Time time = Time.SYSTEM;
+    private Map moduleOptions = null;
+    private boolean configured = false;
+
+    /**
+     * For testing
+     *
+     * @param time
+     *            the mandatory time to set
+     */
+    void time(Time time) {
+        this.time = Objects.requireNonNull(time);
+    }
+
+    /**
+     * Return true if this instance has been configured, otherwise false
+     *
+     * @return true if this instance has been configured, otherwise false
+     */
+    public boolean configured() {
+        return configured;
+    }
+
+    @SuppressWarnings("unchecked")
+    @Override
+    public void configure(Map configs, String saslMechanism, List jaasConfigEntries) {
+        if (!OAuthBearerLoginModule.OAUTHBEARER_MECHANISM.equals(saslMechanism))
+            throw new IllegalArgumentException(String.format("Unexpected SASL mechanism: %s", saslMechanism));
+        if (Objects.requireNonNull(jaasConfigEntries).size() != 1 || jaasConfigEntries.get(0) == null)
+            throw new IllegalArgumentException(
+                    String.format("Must supply exactly 1 non-null JAAS mechanism configuration (size was %d)",
+                            jaasConfigEntries.size()));
+        final Map unmodifiableModuleOptions = Collections
+                .unmodifiableMap((Map) jaasConfigEntries.get(0).getOptions());
+        this.moduleOptions = unmodifiableModuleOptions;
+        configured = true;
+    }
+
+    @Override
+    public void handle(Callback[] callbacks) throws UnsupportedCallbackException {
+        if (!configured())
+            throw new IllegalStateException("Callback handler not configured");
+        for (Callback callback : callbacks) {
+            if (callback instanceof OAuthBearerValidatorCallback) {
+                OAuthBearerValidatorCallback validationCallback = (OAuthBearerValidatorCallback) callback;
+                try {
+                    handleCallback(validationCallback);
+                } catch (OAuthBearerIllegalTokenException e) {
+                    OAuthBearerValidationResult failureReason = e.reason();
+                    String failureScope = failureReason.failureScope();
+                    validationCallback.error(failureScope != null ? "insufficient_scope" : "invalid_token",
+                            failureScope, failureReason.failureOpenIdConfig());
+                }
+            } else if (callback instanceof OAuthBearerExtensionsValidatorCallback) {
+                OAuthBearerExtensionsValidatorCallback extensionsCallback = (OAuthBearerExtensionsValidatorCallback) callback;
+                extensionsCallback.inputExtensions().map().forEach((extensionName, v) -> extensionsCallback.valid(extensionName));
+            } else
+                throw new UnsupportedCallbackException(callback);
+        }
+    }
+
+    @Override
+    public void close() {
+        // empty
+    }
+
+    private void handleCallback(OAuthBearerValidatorCallback callback) {
+        String tokenValue = callback.tokenValue();
+        if (tokenValue == null)
+            throw new IllegalArgumentException("Callback missing required token value");
+        String principalClaimName = principalClaimName();
+        String scopeClaimName = scopeClaimName();
+        List requiredScope = requiredScope();
+        int allowableClockSkewMs = allowableClockSkewMs();
+        OAuthBearerUnsecuredJws unsecuredJwt = new OAuthBearerUnsecuredJws(tokenValue, principalClaimName,
+                scopeClaimName);
+        long now = time.milliseconds();
+        OAuthBearerValidationUtils
+                .validateClaimForExistenceAndType(unsecuredJwt, true, principalClaimName, String.class)
+                .throwExceptionIfFailed();
+        OAuthBearerValidationUtils.validateIssuedAt(unsecuredJwt, false, now, allowableClockSkewMs)
+                .throwExceptionIfFailed();
+        OAuthBearerValidationUtils.validateExpirationTime(unsecuredJwt, now, allowableClockSkewMs)
+                .throwExceptionIfFailed();
+        OAuthBearerValidationUtils.validateTimeConsistency(unsecuredJwt).throwExceptionIfFailed();
+        OAuthBearerValidationUtils.validateScope(unsecuredJwt, requiredScope).throwExceptionIfFailed();
+        log.info("Successfully validated token with principal {}: {}", unsecuredJwt.principalName(),
+                unsecuredJwt.claims().toString());
+        callback.token(unsecuredJwt);
+    }
+
+    private String principalClaimName() {
+        String principalClaimNameValue = option(PRINCIPAL_CLAIM_NAME_OPTION);
+        String principalClaimName = principalClaimNameValue != null && !principalClaimNameValue.trim().isEmpty()
+                ? principalClaimNameValue.trim()
+                : "sub";
+        return principalClaimName;
+    }
+
+    private String scopeClaimName() {
+        String scopeClaimNameValue = option(SCOPE_CLAIM_NAME_OPTION);
+        String scopeClaimName = scopeClaimNameValue != null && !scopeClaimNameValue.trim().isEmpty()
+                ? scopeClaimNameValue.trim()
+                : "scope";
+        return scopeClaimName;
+    }
+
+    private List requiredScope() {
+        String requiredSpaceDelimitedScope = option(REQUIRED_SCOPE_OPTION);
+        List requiredScope = requiredSpaceDelimitedScope == null || requiredSpaceDelimitedScope.trim().isEmpty()
+            ? Collections.emptyList()
+            : OAuthBearerScopeUtils.parseScope(requiredSpaceDelimitedScope.trim());
+        return requiredScope;
+    }
+
+    private int allowableClockSkewMs() {
+        String allowableClockSkewMsValue = option(ALLOWABLE_CLOCK_SKEW_MILLIS_OPTION);
+        int allowableClockSkewMs = 0;
+        try {
+            allowableClockSkewMs = allowableClockSkewMsValue == null || allowableClockSkewMsValue.trim().isEmpty() ? 0
+                    : Integer.parseInt(allowableClockSkewMsValue.trim());
+        } catch (NumberFormatException e) {
+            throw new OAuthBearerConfigException(e.getMessage(), e);
+        }
+        if (allowableClockSkewMs < 0) {
+            throw new OAuthBearerConfigException(
+                    String.format("Allowable clock skew millis must not be negative: %s", allowableClockSkewMsValue));
+        }
+        return allowableClockSkewMs;
+    }
+
+    private String option(String key) {
+        if (!configured)
+            throw new IllegalStateException("Callback handler not configured");
+        return moduleOptions.get(Objects.requireNonNull(key));
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerValidationResult.java b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerValidationResult.java
new file mode 100644
index 0000000000000..2806b4dddcdcf
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerValidationResult.java
@@ -0,0 +1,126 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.oauthbearer.internals.unsecured;
+
+import java.io.Serializable;
+
+/**
+ * The result of some kind of token validation
+ */
+public class OAuthBearerValidationResult implements Serializable {
+    private static final long serialVersionUID = 5774669940899777373L;
+    private final boolean success;
+    private final String failureDescription;
+    private final String failureScope;
+    private final String failureOpenIdConfig;
+
+    /**
+     * Return an instance indicating success
+     * 
+     * @return an instance indicating success
+     */
+    public static OAuthBearerValidationResult newSuccess() {
+        return new OAuthBearerValidationResult(true, null, null, null);
+    }
+
+    /**
+     * Return a new validation failure instance
+     * 
+     * @param failureDescription
+     *            optional description of the failure
+     * @return a new validation failure instance
+     */
+    public static OAuthBearerValidationResult newFailure(String failureDescription) {
+        return newFailure(failureDescription, null, null);
+    }
+
+    /**
+     * Return a new validation failure instance
+     * 
+     * @param failureDescription
+     *            optional description of the failure
+     * @param failureScope
+     *            optional scope to be reported with the failure
+     * @param failureOpenIdConfig
+     *            optional OpenID Connect configuration to be reported with the
+     *            failure
+     * @return a new validation failure instance
+     */
+    public static OAuthBearerValidationResult newFailure(String failureDescription, String failureScope,
+            String failureOpenIdConfig) {
+        return new OAuthBearerValidationResult(false, failureDescription, failureScope, failureOpenIdConfig);
+    }
+
+    private OAuthBearerValidationResult(boolean success, String failureDescription, String failureScope,
+            String failureOpenIdConfig) {
+        if (success && (failureScope != null || failureOpenIdConfig != null))
+            throw new IllegalArgumentException("success was indicated but failure scope/OpenIdConfig were provided");
+        this.success = success;
+        this.failureDescription = failureDescription;
+        this.failureScope = failureScope;
+        this.failureOpenIdConfig = failureOpenIdConfig;
+    }
+
+    /**
+     * Return true if this instance indicates success, otherwise false
+     * 
+     * @return true if this instance indicates success, otherwise false
+     */
+    public boolean success() {
+        return success;
+    }
+
+    /**
+     * Return the (potentially null) descriptive message for the failure
+     * 
+     * @return the (potentially null) descriptive message for the failure
+     */
+    public String failureDescription() {
+        return failureDescription;
+    }
+
+    /**
+     * Return the (potentially null) scope to be reported with the failure
+     * 
+     * @return the (potentially null) scope to be reported with the failure
+     */
+    public String failureScope() {
+        return failureScope;
+    }
+
+    /**
+     * Return the (potentially null) OpenID Connect configuration to be reported
+     * with the failure
+     * 
+     * @return the (potentially null) OpenID Connect configuration to be reported
+     *         with the failure
+     */
+    public String failureOpenIdConfig() {
+        return failureOpenIdConfig;
+    }
+
+    /**
+     * Raise an exception if this instance indicates failure, otherwise do nothing
+     * 
+     * @throws OAuthBearerIllegalTokenException
+     *             if this instance indicates failure
+     */
+    public void throwExceptionIfFailed() throws OAuthBearerIllegalTokenException {
+        if (!success())
+            throw new OAuthBearerIllegalTokenException(this);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerValidationUtils.java b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerValidationUtils.java
new file mode 100644
index 0000000000000..ce1b62b82aa7e
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerValidationUtils.java
@@ -0,0 +1,200 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.oauthbearer.internals.unsecured;
+
+import java.util.List;
+import java.util.Objects;
+import java.util.Set;
+
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerToken;
+
+public class OAuthBearerValidationUtils {
+    /**
+     * Validate the given claim for existence and type. It can be required to exist
+     * in the given claims, and if it exists it must be one of the types indicated
+     * 
+     * @param jwt
+     *            the mandatory JWT to which the validation will be applied
+     * @param required
+     *            true if the claim is required to exist
+     * @param claimName
+     *            the required claim name identifying the claim to be checked
+     * @param allowedTypes
+     *            one or more of {@code String.class}, {@code Number.class}, and
+     *            {@code List.class} identifying the type(s) that the claim value is
+     *            allowed to be if it exists
+     * @return the result of the validation
+     */
+    public static OAuthBearerValidationResult validateClaimForExistenceAndType(OAuthBearerUnsecuredJws jwt,
+            boolean required, String claimName, Class... allowedTypes) {
+        Object rawClaim = Objects.requireNonNull(jwt).rawClaim(Objects.requireNonNull(claimName));
+        if (rawClaim == null)
+            return required
+                    ? OAuthBearerValidationResult.newFailure(String.format("Required claim missing: %s", claimName))
+                    : OAuthBearerValidationResult.newSuccess();
+        for (Class allowedType : allowedTypes) {
+            if (allowedType != null && allowedType.isAssignableFrom(rawClaim.getClass()))
+                return OAuthBearerValidationResult.newSuccess();
+        }
+        return OAuthBearerValidationResult.newFailure(String.format("The %s claim had the incorrect type: %s",
+                claimName, rawClaim.getClass().getSimpleName()));
+    }
+
+    /**
+     * Validate the 'iat' (Issued At) claim. It can be required to exist in the
+     * given claims, and if it exists it must be a (potentially fractional) number
+     * of seconds since the epoch defining when the JWT was issued; it is a
+     * validation error if the Issued At time is after the time at which the check
+     * is being done (plus any allowable clock skew).
+     * 
+     * @param jwt
+     *            the mandatory JWT to which the validation will be applied
+     * @param required
+     *            true if the claim is required to exist
+     * @param whenCheckTimeMs
+     *            the time relative to which the validation is to occur
+     * @param allowableClockSkewMs
+     *            non-negative number to take into account some potential clock skew
+     * @return the result of the validation
+     * @throws OAuthBearerConfigException
+     *             if the given allowable clock skew is negative
+     */
+    public static OAuthBearerValidationResult validateIssuedAt(OAuthBearerUnsecuredJws jwt, boolean required,
+            long whenCheckTimeMs, int allowableClockSkewMs) throws OAuthBearerConfigException {
+        Number value;
+        try {
+            value = Objects.requireNonNull(jwt).issuedAt();
+        } catch (OAuthBearerIllegalTokenException e) {
+            return e.reason();
+        }
+        boolean exists = value != null;
+        if (!exists)
+            return doesNotExistResult(required, "iat");
+        double doubleValue = value.doubleValue();
+        return 1000 * doubleValue > whenCheckTimeMs + confirmNonNegative(allowableClockSkewMs)
+                ? OAuthBearerValidationResult.newFailure(String.format(
+                        "The Issued At value (%f seconds) was after the indicated time (%d ms) plus allowable clock skew (%d ms)",
+                        doubleValue, whenCheckTimeMs, allowableClockSkewMs))
+                : OAuthBearerValidationResult.newSuccess();
+    }
+
+    /**
+     * Validate the 'exp' (Expiration Time) claim. It must exist and it must be a
+     * (potentially fractional) number of seconds defining the point at which the
+     * JWT expires. It is a validation error if the time at which the check is being
+     * done (minus any allowable clock skew) is on or after the Expiration Time
+     * time.
+     * 
+     * @param jwt
+     *            the mandatory JWT to which the validation will be applied
+     * @param whenCheckTimeMs
+     *            the time relative to which the validation is to occur
+     * @param allowableClockSkewMs
+     *            non-negative number to take into account some potential clock skew
+     * @return the result of the validation
+     * @throws OAuthBearerConfigException
+     *             if the given allowable clock skew is negative
+     */
+    public static OAuthBearerValidationResult validateExpirationTime(OAuthBearerUnsecuredJws jwt, long whenCheckTimeMs,
+            int allowableClockSkewMs) throws OAuthBearerConfigException {
+        Number value;
+        try {
+            value = Objects.requireNonNull(jwt).expirationTime();
+        } catch (OAuthBearerIllegalTokenException e) {
+            return e.reason();
+        }
+        boolean exists = value != null;
+        if (!exists)
+            return doesNotExistResult(true, "exp");
+        double doubleValue = value.doubleValue();
+        return whenCheckTimeMs - confirmNonNegative(allowableClockSkewMs) >= 1000 * doubleValue
+                ? OAuthBearerValidationResult.newFailure(String.format(
+                        "The indicated time (%d ms) minus allowable clock skew (%d ms) was on or after the Expiration Time value (%f seconds)",
+                        whenCheckTimeMs, allowableClockSkewMs, doubleValue))
+                : OAuthBearerValidationResult.newSuccess();
+    }
+
+    /**
+     * Validate the 'iat' (Issued At) and 'exp' (Expiration Time) claims for
+     * internal consistency. The following must be true if both claims exist:
+     * 
+     * 
            +     * exp > iat
+     * 
            
+     * 
+     * @param jwt
+     *            the mandatory JWT to which the validation will be applied
+     * @return the result of the validation
+     */
+    public static OAuthBearerValidationResult validateTimeConsistency(OAuthBearerUnsecuredJws jwt) {
+        Number issuedAt;
+        Number expirationTime;
+        try {
+            issuedAt = Objects.requireNonNull(jwt).issuedAt();
+            expirationTime = jwt.expirationTime();
+        } catch (OAuthBearerIllegalTokenException e) {
+            return e.reason();
+        }
+        if (expirationTime != null && issuedAt != null && expirationTime.doubleValue() <= issuedAt.doubleValue())
+            return OAuthBearerValidationResult.newFailure(
+                    String.format("The Expiration Time time (%f seconds) was not after the Issued At time (%f seconds)",
+                            expirationTime.doubleValue(), issuedAt.doubleValue()));
+        return OAuthBearerValidationResult.newSuccess();
+    }
+
+    /**
+     * Validate the given token's scope against the required scope. Every required
+     * scope element (if any) must exist in the provided token's scope for the
+     * validation to succeed.
+     * 
+     * @param token
+     *            the required token for which the scope will to validate
+     * @param requiredScope
+     *            the optional required scope against which the given token's scope
+     *            will be validated
+     * @return the result of the validation
+     */
+    public static OAuthBearerValidationResult validateScope(OAuthBearerToken token, List requiredScope) {
+        final Set tokenScope = token.scope();
+        if (requiredScope == null || requiredScope.isEmpty())
+            return OAuthBearerValidationResult.newSuccess();
+        for (String requiredScopeElement : requiredScope) {
+            if (!tokenScope.contains(requiredScopeElement))
+                return OAuthBearerValidationResult.newFailure(String.format(
+                        "The provided scope (%s) was mising a required scope (%s).  All required scope elements: %s",
+                        String.valueOf(tokenScope), requiredScopeElement, requiredScope.toString()),
+                        requiredScope.toString(), null);
+        }
+        return OAuthBearerValidationResult.newSuccess();
+    }
+
+    private static int confirmNonNegative(int allowableClockSkewMs) throws OAuthBearerConfigException {
+        if (allowableClockSkewMs < 0)
+            throw new OAuthBearerConfigException(
+                    String.format("Allowable clock skew must not be negative: %d", allowableClockSkewMs));
+        return allowableClockSkewMs;
+    }
+
+    private static OAuthBearerValidationResult doesNotExistResult(boolean required, String claimName) {
+        return required ? OAuthBearerValidationResult.newFailure(String.format("Required claim missing: %s", claimName))
+                : OAuthBearerValidationResult.newSuccess();
+    }
+
+    private OAuthBearerValidationUtils() {
+        // empty
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/plain/PlainAuthenticateCallback.java b/clients/src/main/java/org/apache/kafka/common/security/plain/PlainAuthenticateCallback.java
new file mode 100644
index 0000000000000..7f42645e487fb
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/plain/PlainAuthenticateCallback.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.security.plain;
+
+import javax.security.auth.callback.Callback;
+
+/*
+ * Authentication callback for SASL/PLAIN authentication. Callback handler must
+ * set authenticated flag to true if the client provided password in the callback
+ * matches the expected password.
+ */
+public class PlainAuthenticateCallback implements Callback {
+    private final char[] password;
+    private boolean authenticated;
+
+    /**
+     * Creates a callback with the password provided by the client
+     * @param password The password provided by the client during SASL/PLAIN authentication
+     */
+    public PlainAuthenticateCallback(char[] password) {
+        this.password = password;
+    }
+
+    /**
+     * Returns the password provided by the client during SASL/PLAIN authentication
+     */
+    public char[] password() {
+        return password;
+    }
+
+    /**
+     * Returns true if client password matches expected password, false otherwise.
+     * This state is set the server-side callback handler.
+     */
+    public boolean authenticated() {
+        return this.authenticated;
+    }
+
+    /**
+     * Sets the authenticated state. This is set by the server-side callback handler
+     * by matching the client provided password with the expected password.
+     *
+     * @param authenticated true indicates successful authentication
+     */
+    public void authenticated(boolean authenticated) {
+        this.authenticated = authenticated;
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/plain/PlainLoginModule.java b/clients/src/main/java/org/apache/kafka/common/security/plain/PlainLoginModule.java
index c8b29fc69a9cc..40851681ba63f 100644
--- a/clients/src/main/java/org/apache/kafka/common/security/plain/PlainLoginModule.java
+++ b/clients/src/main/java/org/apache/kafka/common/security/plain/PlainLoginModule.java
@@ -16,11 +16,12 @@
  */
 package org.apache.kafka.common.security.plain;
 
+import org.apache.kafka.common.security.plain.internals.PlainSaslServerProvider;
+
 import java.util.Map;
 
 import javax.security.auth.Subject;
 import javax.security.auth.callback.CallbackHandler;
-import javax.security.auth.login.LoginException;
 import javax.security.auth.spi.LoginModule;
 
 public class PlainLoginModule implements LoginModule {
@@ -43,22 +44,22 @@ public void initialize(Subject subject, CallbackHandler callbackHandler, Map
- * Valid users with passwords are specified in the Jaas configuration file. Each user
- * is specified with user_ as key and  as value. This is consistent
- * with Zookeeper Digest-MD5 implementation.
- * 
            
- * To avoid storing clear passwords on disk or to integrate with external authentication
- * servers in production systems, this module can be replaced with a different implementation.
- *
- */
-public class PlainSaslServer implements SaslServer {
-
-    public static final String PLAIN_MECHANISM = "PLAIN";
-    private static final String JAAS_USER_PREFIX = "user_";
-
-    private final JaasContext jaasContext;
-
-    private boolean complete;
-    private String authorizationId;
-
-    public PlainSaslServer(JaasContext jaasContext) {
-        this.jaasContext = jaasContext;
-    }
-
-    /**
-     * @throws SaslAuthenticationException if username/password combination is invalid or if the requested
-     *         authorization id is not the same as username.
-     * 
            
-     * Note: This method may throw {@link SaslAuthenticationException} to provide custom error messages
-     * to clients. But care should be taken to avoid including any information in the exception message that
-     * should not be leaked to unauthenticated clients. It may be safer to throw {@link SaslException} in
-     * some cases so that a standard error message is returned to clients.
-     * 
            
-     */
-    @Override
-    public byte[] evaluateResponse(byte[] response) throws SaslException, SaslAuthenticationException {
-        /*
-         * Message format (from https://tools.ietf.org/html/rfc4616):
-         *
-         * message   = [authzid] UTF8NUL authcid UTF8NUL passwd
-         * authcid   = 1*SAFE ; MUST accept up to 255 octets
-         * authzid   = 1*SAFE ; MUST accept up to 255 octets
-         * passwd    = 1*SAFE ; MUST accept up to 255 octets
-         * UTF8NUL   = %x00 ; UTF-8 encoded NUL character
-         *
-         * SAFE      = UTF1 / UTF2 / UTF3 / UTF4
-         *                ;; any UTF-8 encoded Unicode character except NUL
-         */
-
-        String[] tokens;
-        try {
-            tokens = new String(response, "UTF-8").split("\u0000");
-        } catch (UnsupportedEncodingException e) {
-            throw new SaslException("UTF-8 encoding not supported", e);
-        }
-        if (tokens.length != 3)
-            throw new SaslException("Invalid SASL/PLAIN response: expected 3 tokens, got " + tokens.length);
-        String authorizationIdFromClient = tokens[0];
-        String username = tokens[1];
-        String password = tokens[2];
-
-        if (username.isEmpty()) {
-            throw new SaslException("Authentication failed: username not specified");
-        }
-        if (password.isEmpty()) {
-            throw new SaslException("Authentication failed: password not specified");
-        }
-
-        String expectedPassword = jaasContext.configEntryOption(JAAS_USER_PREFIX + username,
-                PlainLoginModule.class.getName());
-        if (!password.equals(expectedPassword)) {
-            throw new SaslAuthenticationException("Authentication failed: Invalid username or password");
-        }
-
-        if (!authorizationIdFromClient.isEmpty() && !authorizationIdFromClient.equals(username))
-            throw new SaslAuthenticationException("Authentication failed: Client requested an authorization id that is different from username");
-
-        this.authorizationId = username;
-
-        complete = true;
-        return new byte[0];
-    }
-
-    @Override
-    public String getAuthorizationID() {
-        if (!complete)
-            throw new IllegalStateException("Authentication exchange has not completed");
-        return authorizationId;
-    }
-
-    @Override
-    public String getMechanismName() {
-        return PLAIN_MECHANISM;
-    }
-
-    @Override
-    public Object getNegotiatedProperty(String propName) {
-        if (!complete)
-            throw new IllegalStateException("Authentication exchange has not completed");
-        return null;
-    }
-
-    @Override
-    public boolean isComplete() {
-        return complete;
-    }
-
-    @Override
-    public byte[] unwrap(byte[] incoming, int offset, int len) throws SaslException {
-        if (!complete)
-            throw new IllegalStateException("Authentication exchange has not completed");
-        return Arrays.copyOfRange(incoming, offset, offset + len);
-    }
-
-    @Override
-    public byte[] wrap(byte[] outgoing, int offset, int len) throws SaslException {
-        if (!complete)
-            throw new IllegalStateException("Authentication exchange has not completed");
-        return Arrays.copyOfRange(outgoing, offset, offset + len);
-    }
-
-    @Override
-    public void dispose() throws SaslException {
-    }
-
-    public static class PlainSaslServerFactory implements SaslServerFactory {
-
-        @Override
-        public SaslServer createSaslServer(String mechanism, String protocol, String serverName, Map props, CallbackHandler cbh)
-            throws SaslException {
-
-            if (!PLAIN_MECHANISM.equals(mechanism))
-                throw new SaslException(String.format("Mechanism \'%s\' is not supported. Only PLAIN is supported.", mechanism));
-
-            if (!(cbh instanceof SaslServerCallbackHandler))
-                throw new SaslException("CallbackHandler must be of type SaslServerCallbackHandler, but it is: " + cbh.getClass());
-
-            return new PlainSaslServer(((SaslServerCallbackHandler) cbh).jaasContext());
-        }
-
-        @Override
-        public String[] getMechanismNames(Map props) {
-            if (props == null) return new String[]{PLAIN_MECHANISM};
-            String noPlainText = (String) props.get(Sasl.POLICY_NOPLAINTEXT);
-            if ("true".equals(noPlainText))
-                return new String[]{};
-            else
-                return new String[]{PLAIN_MECHANISM};
-        }
-    }
-}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/plain/internals/PlainSaslServer.java b/clients/src/main/java/org/apache/kafka/common/security/plain/internals/PlainSaslServer.java
new file mode 100644
index 0000000000000..7a65fe67acd30
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/plain/internals/PlainSaslServer.java
@@ -0,0 +1,201 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.common.security.plain.internals;
+
+import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
+
+import javax.security.auth.callback.Callback;
+import javax.security.auth.callback.CallbackHandler;
+import javax.security.auth.callback.NameCallback;
+import javax.security.sasl.Sasl;
+import javax.security.sasl.SaslException;
+import javax.security.sasl.SaslServer;
+import javax.security.sasl.SaslServerFactory;
+
+import org.apache.kafka.common.errors.SaslAuthenticationException;
+import org.apache.kafka.common.security.plain.PlainAuthenticateCallback;
+
+/**
+ * Simple SaslServer implementation for SASL/PLAIN. In order to make this implementation
+ * fully pluggable, authentication of username/password is fully contained within the
+ * server implementation.
+ * 
            
+ * Valid users with passwords are specified in the Jaas configuration file. Each user
+ * is specified with user_ as key and  as value. This is consistent
+ * with Zookeeper Digest-MD5 implementation.
+ * 
            
+ * To avoid storing clear passwords on disk or to integrate with external authentication
+ * servers in production systems, this module can be replaced with a different implementation.
+ *
+ */
+public class PlainSaslServer implements SaslServer {
+
+    public static final String PLAIN_MECHANISM = "PLAIN";
+
+    private final CallbackHandler callbackHandler;
+    private boolean complete;
+    private String authorizationId;
+
+    public PlainSaslServer(CallbackHandler callbackHandler) {
+        this.callbackHandler = callbackHandler;
+    }
+
+    /**
+     * @throws SaslAuthenticationException if username/password combination is invalid or if the requested
+     *         authorization id is not the same as username.
+     * 
            
+     * Note: This method may throw {@link SaslAuthenticationException} to provide custom error messages
+     * to clients. But care should be taken to avoid including any information in the exception message that
+     * should not be leaked to unauthenticated clients. It may be safer to throw {@link SaslException} in
+     * some cases so that a standard error message is returned to clients.
+     * 
            
+     */
+    @Override
+    public byte[] evaluateResponse(byte[] responseBytes) throws SaslAuthenticationException {
+        /*
+         * Message format (from https://tools.ietf.org/html/rfc4616):
+         *
+         * message   = [authzid] UTF8NUL authcid UTF8NUL passwd
+         * authcid   = 1*SAFE ; MUST accept up to 255 octets
+         * authzid   = 1*SAFE ; MUST accept up to 255 octets
+         * passwd    = 1*SAFE ; MUST accept up to 255 octets
+         * UTF8NUL   = %x00 ; UTF-8 encoded NUL character
+         *
+         * SAFE      = UTF1 / UTF2 / UTF3 / UTF4
+         *                ;; any UTF-8 encoded Unicode character except NUL
+         */
+
+        String response = new String(responseBytes, StandardCharsets.UTF_8);
+        List tokens = extractTokens(response);
+        String authorizationIdFromClient = tokens.get(0);
+        String username = tokens.get(1);
+        String password = tokens.get(2);
+
+        if (username.isEmpty()) {
+            throw new SaslAuthenticationException("Authentication failed: username not specified");
+        }
+        if (password.isEmpty()) {
+            throw new SaslAuthenticationException("Authentication failed: password not specified");
+        }
+
+        NameCallback nameCallback = new NameCallback("username", username);
+        PlainAuthenticateCallback authenticateCallback = new PlainAuthenticateCallback(password.toCharArray());
+        try {
+            callbackHandler.handle(new Callback[]{nameCallback, authenticateCallback});
+        } catch (Throwable e) {
+            throw new SaslAuthenticationException("Authentication failed: credentials for user could not be verified", e);
+        }
+        if (!authenticateCallback.authenticated())
+            throw new SaslAuthenticationException("Authentication failed: Invalid username or password");
+        if (!authorizationIdFromClient.isEmpty() && !authorizationIdFromClient.equals(username))
+            throw new SaslAuthenticationException("Authentication failed: Client requested an authorization id that is different from username");
+
+        this.authorizationId = username;
+
+        complete = true;
+        return new byte[0];
+    }
+
+    private List extractTokens(String string) {
+        List tokens = new ArrayList<>();
+        int startIndex = 0;
+        for (int i = 0; i < 4; ++i) {
+            int endIndex = string.indexOf("\u0000", startIndex);
+            if (endIndex == -1) {
+                tokens.add(string.substring(startIndex));
+                break;
+            }
+            tokens.add(string.substring(startIndex, endIndex));
+            startIndex = endIndex + 1;
+        }
+
+        if (tokens.size() != 3)
+            throw new SaslAuthenticationException("Invalid SASL/PLAIN response: expected 3 tokens, got " +
+                tokens.size());
+
+        return tokens;
+    }
+
+    @Override
+    public String getAuthorizationID() {
+        if (!complete)
+            throw new IllegalStateException("Authentication exchange has not completed");
+        return authorizationId;
+    }
+
+    @Override
+    public String getMechanismName() {
+        return PLAIN_MECHANISM;
+    }
+
+    @Override
+    public Object getNegotiatedProperty(String propName) {
+        if (!complete)
+            throw new IllegalStateException("Authentication exchange has not completed");
+        return null;
+    }
+
+    @Override
+    public boolean isComplete() {
+        return complete;
+    }
+
+    @Override
+    public byte[] unwrap(byte[] incoming, int offset, int len) {
+        if (!complete)
+            throw new IllegalStateException("Authentication exchange has not completed");
+        return Arrays.copyOfRange(incoming, offset, offset + len);
+    }
+
+    @Override
+    public byte[] wrap(byte[] outgoing, int offset, int len) {
+        if (!complete)
+            throw new IllegalStateException("Authentication exchange has not completed");
+        return Arrays.copyOfRange(outgoing, offset, offset + len);
+    }
+
+    @Override
+    public void dispose() {
+    }
+
+    public static class PlainSaslServerFactory implements SaslServerFactory {
+
+        @Override
+        public SaslServer createSaslServer(String mechanism, String protocol, String serverName, Map props, CallbackHandler cbh)
+            throws SaslException {
+
+            if (!PLAIN_MECHANISM.equals(mechanism))
+                throw new SaslException(String.format("Mechanism \'%s\' is not supported. Only PLAIN is supported.", mechanism));
+
+            return new PlainSaslServer(cbh);
+        }
+
+        @Override
+        public String[] getMechanismNames(Map props) {
+            if (props == null) return new String[]{PLAIN_MECHANISM};
+            String noPlainText = (String) props.get(Sasl.POLICY_NOPLAINTEXT);
+            if ("true".equals(noPlainText))
+                return new String[]{};
+            else
+                return new String[]{PLAIN_MECHANISM};
+        }
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/plain/PlainSaslServerProvider.java b/clients/src/main/java/org/apache/kafka/common/security/plain/internals/PlainSaslServerProvider.java
similarity index 89%
rename from clients/src/main/java/org/apache/kafka/common/security/plain/PlainSaslServerProvider.java
rename to clients/src/main/java/org/apache/kafka/common/security/plain/internals/PlainSaslServerProvider.java
index ae1424417f8fb..33c3dc5623db7 100644
--- a/clients/src/main/java/org/apache/kafka/common/security/plain/PlainSaslServerProvider.java
+++ b/clients/src/main/java/org/apache/kafka/common/security/plain/internals/PlainSaslServerProvider.java
@@ -14,12 +14,12 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.kafka.common.security.plain;
+package org.apache.kafka.common.security.plain.internals;
 
 import java.security.Provider;
 import java.security.Security;
 
-import org.apache.kafka.common.security.plain.PlainSaslServer.PlainSaslServerFactory;
+import org.apache.kafka.common.security.plain.internals.PlainSaslServer.PlainSaslServerFactory;
 
 public class PlainSaslServerProvider extends Provider {
 
diff --git a/clients/src/main/java/org/apache/kafka/common/security/plain/internals/PlainServerCallbackHandler.java b/clients/src/main/java/org/apache/kafka/common/security/plain/internals/PlainServerCallbackHandler.java
new file mode 100644
index 0000000000000..842f986abd7e6
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/plain/internals/PlainServerCallbackHandler.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.common.security.plain.internals;
+
+import org.apache.kafka.common.security.JaasContext;
+import org.apache.kafka.common.security.auth.AuthenticateCallbackHandler;
+import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.security.plain.PlainAuthenticateCallback;
+import org.apache.kafka.common.security.plain.PlainLoginModule;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
+
+import javax.security.auth.callback.Callback;
+import javax.security.auth.callback.NameCallback;
+import javax.security.auth.callback.UnsupportedCallbackException;
+import javax.security.auth.login.AppConfigurationEntry;
+
+public class PlainServerCallbackHandler implements AuthenticateCallbackHandler {
+
+    private static final String JAAS_USER_PREFIX = "user_";
+    private List jaasConfigEntries;
+
+    @Override
+    public void configure(Map configs, String mechanism, List jaasConfigEntries) {
+        this.jaasConfigEntries = jaasConfigEntries;
+    }
+
+    @Override
+    public void handle(Callback[] callbacks) throws IOException, UnsupportedCallbackException {
+        String username = null;
+        for (Callback callback: callbacks) {
+            if (callback instanceof NameCallback)
+                username = ((NameCallback) callback).getDefaultName();
+            else if (callback instanceof PlainAuthenticateCallback) {
+                PlainAuthenticateCallback plainCallback = (PlainAuthenticateCallback) callback;
+                boolean authenticated = authenticate(username, plainCallback.password());
+                plainCallback.authenticated(authenticated);
+            } else
+                throw new UnsupportedCallbackException(callback);
+        }
+    }
+
+    protected boolean authenticate(String username, char[] password) throws IOException {
+        if (username == null)
+            return false;
+        else {
+            String expectedPassword = JaasContext.configEntryOption(jaasConfigEntries,
+                    JAAS_USER_PREFIX + username,
+                    PlainLoginModule.class.getName());
+            return expectedPassword != null && Arrays.equals(password, expectedPassword.toCharArray());
+        }
+    }
+
+    @Override
+    public void close() throws KafkaException {
+    }
+
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/scram/ScramCredential.java b/clients/src/main/java/org/apache/kafka/common/security/scram/ScramCredential.java
index 09ff0aacd3c84..dfbfef15b4f74 100644
--- a/clients/src/main/java/org/apache/kafka/common/security/scram/ScramCredential.java
+++ b/clients/src/main/java/org/apache/kafka/common/security/scram/ScramCredential.java
@@ -16,6 +16,11 @@
  */
 package org.apache.kafka.common.security.scram;
 
+/**
+ * SCRAM credential class that encapsulates the credential data persisted for each user that is
+ * accessible to the server. See RFC rfc5802
+ * for details.
+ */
 public class ScramCredential {
 
     private final byte[] salt;
@@ -23,6 +28,9 @@ public class ScramCredential {
     private final byte[] storedKey;
     private final int iterations;
 
+    /**
+     * Constructs a new credential.
+     */
     public ScramCredential(byte[] salt, byte[] storedKey, byte[] serverKey, int iterations) {
         this.salt = salt;
         this.serverKey = serverKey;
@@ -30,18 +38,30 @@ public ScramCredential(byte[] salt, byte[] storedKey, byte[] serverKey, int iter
         this.iterations = iterations;
     }
 
+    /**
+     * Returns the salt used to process this credential using the SCRAM algorithm.
+     */
     public byte[] salt() {
         return salt;
     }
 
+    /**
+     * Server key computed from the client password using the SCRAM algorithm.
+     */
     public byte[] serverKey() {
         return serverKey;
     }
 
+    /**
+     * Stored key computed from the client password using the SCRAM algorithm.
+     */
     public byte[] storedKey() {
         return storedKey;
     }
 
+    /**
+     * Number of iterations used to process this credential using the SCRAM algorithm.
+     */
     public int iterations() {
         return iterations;
     }
diff --git a/clients/src/main/java/org/apache/kafka/common/security/scram/ScramCredentialCallback.java b/clients/src/main/java/org/apache/kafka/common/security/scram/ScramCredentialCallback.java
index 931210a0e0702..d5988cbeb8905 100644
--- a/clients/src/main/java/org/apache/kafka/common/security/scram/ScramCredentialCallback.java
+++ b/clients/src/main/java/org/apache/kafka/common/security/scram/ScramCredentialCallback.java
@@ -18,14 +18,23 @@
 
 import javax.security.auth.callback.Callback;
 
+/**
+ * Callback used for SCRAM mechanisms.
+ */
 public class ScramCredentialCallback implements Callback {
     private ScramCredential scramCredential;
 
-    public ScramCredential scramCredential() {
-        return scramCredential;
-    }
-
+    /**
+     * Sets the SCRAM credential for this instance.
+     */
     public void scramCredential(ScramCredential scramCredential) {
         this.scramCredential = scramCredential;
     }
-}
\ No newline at end of file
+
+    /**
+     * Returns the SCRAM credential if set on this instance.
+     */
+    public ScramCredential scramCredential() {
+        return scramCredential;
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/scram/ScramExtensionsCallback.java b/clients/src/main/java/org/apache/kafka/common/security/scram/ScramExtensionsCallback.java
new file mode 100644
index 0000000000000..b83c94e04bcc4
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/scram/ScramExtensionsCallback.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.common.security.scram;
+
+import javax.security.auth.callback.Callback;
+import java.util.Collections;
+import java.util.Map;
+
+
+/**
+ * Optional callback used for SCRAM mechanisms if any extensions need to be set
+ * in the SASL/SCRAM exchange.
+ */
+public class ScramExtensionsCallback implements Callback {
+    private Map extensions = Collections.emptyMap();
+
+    /**
+     * Returns map of the extension names and values that are sent by the client to
+     * the server in the initial client SCRAM authentication message.
+     * Default is an empty unmodifiable map.
+     */
+    public Map extensions() {
+        return extensions;
+    }
+
+    /**
+     * Sets the SCRAM extensions on this callback. Maps passed in should be unmodifiable
+     */
+    public void extensions(Map extensions) {
+        this.extensions = extensions;
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/scram/ScramFormatter.java b/clients/src/main/java/org/apache/kafka/common/security/scram/ScramFormatter.java
deleted file mode 100644
index 406c28599b3bb..0000000000000
--- a/clients/src/main/java/org/apache/kafka/common/security/scram/ScramFormatter.java
+++ /dev/null
@@ -1,167 +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.security.scram;
-
-import java.math.BigInteger;
-import java.nio.charset.StandardCharsets;
-import java.security.InvalidKeyException;
-import java.security.MessageDigest;
-import java.security.NoSuchAlgorithmException;
-import java.security.SecureRandom;
-
-import javax.crypto.Mac;
-import javax.crypto.spec.SecretKeySpec;
-
-import org.apache.kafka.common.KafkaException;
-import org.apache.kafka.common.security.scram.ScramMessages.ClientFinalMessage;
-import org.apache.kafka.common.security.scram.ScramMessages.ClientFirstMessage;
-import org.apache.kafka.common.security.scram.ScramMessages.ServerFirstMessage;
-
-/**
- * Scram message salt and hash functions defined in RFC 5802.
- */
-public class ScramFormatter {
-
-    private final MessageDigest messageDigest;
-    private final Mac mac;
-    private final SecureRandom random;
-
-    public ScramFormatter(ScramMechanism mechanism) throws NoSuchAlgorithmException {
-        this.messageDigest = MessageDigest.getInstance(mechanism.hashAlgorithm());
-        this.mac = Mac.getInstance(mechanism.macAlgorithm());
-        this.random = new SecureRandom();
-    }
-
-    public byte[] hmac(byte[] key, byte[] bytes) throws InvalidKeyException {
-        mac.init(new SecretKeySpec(key, mac.getAlgorithm()));
-        return mac.doFinal(bytes);
-    }
-
-    public byte[] hash(byte[] str) {
-        return messageDigest.digest(str);
-    }
-
-    public byte[] xor(byte[] first, byte[] second) {
-        if (first.length != second.length)
-            throw new IllegalArgumentException("Argument arrays must be of the same length");
-        byte[] result = new byte[first.length];
-        for (int i = 0; i < result.length; i++)
-            result[i] = (byte) (first[i] ^ second[i]);
-        return result;
-    }
-
-    public byte[] hi(byte[] str, byte[] salt, int iterations) throws InvalidKeyException {
-        mac.init(new SecretKeySpec(str, mac.getAlgorithm()));
-        mac.update(salt);
-        byte[] u1 = mac.doFinal(new byte[]{0, 0, 0, 1});
-        byte[] prev = u1;
-        byte[] result = u1;
-        for (int i = 2; i <= iterations; i++) {
-            byte[] ui = hmac(str, prev);
-            result = xor(result, ui);
-            prev = ui;
-        }
-        return result;
-    }
-
-    public byte[] normalize(String str) {
-        return toBytes(str);
-    }
-
-    public byte[] saltedPassword(String password, byte[] salt, int iterations) throws InvalidKeyException {
-        return hi(normalize(password), salt, iterations);
-    }
-
-    public byte[] clientKey(byte[] saltedPassword) throws InvalidKeyException {
-        return hmac(saltedPassword, toBytes("Client Key"));
-    }
-
-    public byte[] storedKey(byte[] clientKey) {
-        return hash(clientKey);
-    }
-
-    public String saslName(String username) {
-        return username.replace("=", "=3D").replace(",", "=2C");
-    }
-
-    public String username(String saslName) {
-        String username = saslName.replace("=2C", ",");
-        if (username.replace("=3D", "").indexOf('=') >= 0)
-            throw new IllegalArgumentException("Invalid username: " + saslName);
-        return username.replace("=3D", "=");
-    }
-
-    public String authMessage(String clientFirstMessageBare, String serverFirstMessage, String clientFinalMessageWithoutProof) {
-        return clientFirstMessageBare + "," + serverFirstMessage + "," + clientFinalMessageWithoutProof;
-    }
-
-    public byte[] clientSignature(byte[] storedKey, ClientFirstMessage clientFirstMessage, ServerFirstMessage serverFirstMessage, ClientFinalMessage clientFinalMessage) throws InvalidKeyException {
-        byte[] authMessage = authMessage(clientFirstMessage, serverFirstMessage, clientFinalMessage);
-        return hmac(storedKey, authMessage);
-    }
-
-    public byte[] clientProof(byte[] saltedPassword, ClientFirstMessage clientFirstMessage, ServerFirstMessage serverFirstMessage, ClientFinalMessage clientFinalMessage) throws InvalidKeyException {
-        byte[] clientKey = clientKey(saltedPassword);
-        byte[] storedKey = hash(clientKey);
-        byte[] clientSignature = hmac(storedKey, authMessage(clientFirstMessage, serverFirstMessage, clientFinalMessage));
-        return xor(clientKey, clientSignature);
-    }
-
-    private byte[] authMessage(ClientFirstMessage clientFirstMessage, ServerFirstMessage serverFirstMessage, ClientFinalMessage clientFinalMessage) {
-        return toBytes(authMessage(clientFirstMessage.clientFirstMessageBare(),
-                serverFirstMessage.toMessage(),
-                clientFinalMessage.clientFinalMessageWithoutProof()));
-    }
-
-    public byte[] storedKey(byte[] clientSignature, byte[] clientProof) {
-        return hash(xor(clientSignature, clientProof));
-    }
-
-    public byte[] serverKey(byte[] saltedPassword) throws InvalidKeyException {
-        return hmac(saltedPassword, toBytes("Server Key"));
-    }
-
-    public byte[] serverSignature(byte[] serverKey, ClientFirstMessage clientFirstMessage, ServerFirstMessage serverFirstMessage, ClientFinalMessage clientFinalMessage) throws InvalidKeyException {
-        byte[] authMessage = authMessage(clientFirstMessage, serverFirstMessage, clientFinalMessage);
-        return hmac(serverKey, authMessage);
-    }
-
-    public String secureRandomString() {
-        return new BigInteger(130, random).toString(Character.MAX_RADIX);
-    }
-
-    public byte[] secureRandomBytes() {
-        return toBytes(secureRandomString());
-    }
-
-    public byte[] toBytes(String str) {
-        return str.getBytes(StandardCharsets.UTF_8);
-    }
-
-    public ScramCredential generateCredential(String password, int iterations) {
-        try {
-            byte[] salt = secureRandomBytes();
-            byte[] saltedPassword = saltedPassword(password, salt, iterations);
-            byte[] clientKey = clientKey(saltedPassword);
-            byte[] storedKey = storedKey(clientKey);
-            byte[] serverKey = serverKey(saltedPassword);
-            return new ScramCredential(salt, storedKey, serverKey, iterations);
-        } catch (InvalidKeyException e) {
-            throw new KafkaException("Could not create credential", e);
-        }
-    }
-}
\ No newline at end of file
diff --git a/clients/src/main/java/org/apache/kafka/common/security/scram/ScramLoginModule.java b/clients/src/main/java/org/apache/kafka/common/security/scram/ScramLoginModule.java
index c087a3215d782..104b3fc04ad6c 100644
--- a/clients/src/main/java/org/apache/kafka/common/security/scram/ScramLoginModule.java
+++ b/clients/src/main/java/org/apache/kafka/common/security/scram/ScramLoginModule.java
@@ -16,17 +16,21 @@
  */
 package org.apache.kafka.common.security.scram;
 
+import org.apache.kafka.common.security.scram.internals.ScramSaslClientProvider;
+import org.apache.kafka.common.security.scram.internals.ScramSaslServerProvider;
+
+import java.util.Collections;
 import java.util.Map;
 
 import javax.security.auth.Subject;
 import javax.security.auth.callback.CallbackHandler;
-import javax.security.auth.login.LoginException;
 import javax.security.auth.spi.LoginModule;
 
 public class ScramLoginModule implements LoginModule {
 
     private static final String USERNAME_CONFIG = "username";
     private static final String PASSWORD_CONFIG = "password";
+    public static final String TOKEN_AUTH_CONFIG = "tokenauth";
 
     static {
         ScramSaslClientProvider.initialize();
@@ -41,25 +45,31 @@ public void initialize(Subject subject, CallbackHandler callbackHandler, Map scramExtensions = Collections.singletonMap(TOKEN_AUTH_CONFIG, "true");
+            subject.getPublicCredentials().add(scramExtensions);
+        }
     }
 
     @Override
-    public boolean login() throws LoginException {
+    public boolean login() {
         return true;
     }
 
     @Override
-    public boolean logout() throws LoginException {
+    public boolean logout() {
         return true;
     }
 
     @Override
-    public boolean commit() throws LoginException {
+    public boolean commit() {
         return true;
     }
 
     @Override
-    public boolean abort() throws LoginException {
+    public boolean abort() {
         return false;
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/security/scram/ScramMechanism.java b/clients/src/main/java/org/apache/kafka/common/security/scram/ScramMechanism.java
deleted file mode 100644
index d8c0c6d3d603e..0000000000000
--- a/clients/src/main/java/org/apache/kafka/common/security/scram/ScramMechanism.java
+++ /dev/null
@@ -1,77 +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.security.scram;
-
-import java.util.Collection;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.Map;
-
-public enum ScramMechanism {
-
-    SCRAM_SHA_256("SHA-256", "HmacSHA256", 4096),
-    SCRAM_SHA_512("SHA-512", "HmacSHA512", 4096);
-
-    private final String mechanismName;
-    private final String hashAlgorithm;
-    private final String macAlgorithm;
-    private final int minIterations;
-
-    private static final Map MECHANISMS_MAP;
-
-    static {
-        Map map = new HashMap<>();
-        for (ScramMechanism mech : values())
-            map.put(mech.mechanismName, mech);
-        MECHANISMS_MAP = Collections.unmodifiableMap(map);
-    }
-
-    ScramMechanism(String hashAlgorithm, String macAlgorithm, int minIterations) {
-        this.mechanismName = "SCRAM-" + hashAlgorithm;
-        this.hashAlgorithm = hashAlgorithm;
-        this.macAlgorithm = macAlgorithm;
-        this.minIterations = minIterations;
-    }
-
-    public final String mechanismName() {
-        return mechanismName;
-    }
-
-    public String hashAlgorithm() {
-        return hashAlgorithm;
-    }
-
-    public String macAlgorithm() {
-        return macAlgorithm;
-    }
-
-    public int minIterations() {
-        return minIterations;
-    }
-
-    public static ScramMechanism forMechanismName(String mechanismName) {
-        return MECHANISMS_MAP.get(mechanismName);
-    }
-
-    public static Collection mechanismNames() {
-        return MECHANISMS_MAP.keySet();
-    }
-
-    public static boolean isScram(String mechanismName) {
-        return MECHANISMS_MAP.containsKey(mechanismName);
-    }
-}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/scram/ScramSaslServer.java b/clients/src/main/java/org/apache/kafka/common/security/scram/ScramSaslServer.java
deleted file mode 100644
index b250e155ae1b8..0000000000000
--- a/clients/src/main/java/org/apache/kafka/common/security/scram/ScramSaslServer.java
+++ /dev/null
@@ -1,230 +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.security.scram;
-
-import java.io.IOException;
-import java.security.InvalidKeyException;
-import java.security.NoSuchAlgorithmException;
-import java.util.Arrays;
-import java.util.Collection;
-import java.util.Map;
-
-import javax.security.auth.callback.Callback;
-import javax.security.auth.callback.CallbackHandler;
-import javax.security.auth.callback.NameCallback;
-import javax.security.auth.callback.UnsupportedCallbackException;
-import javax.security.sasl.SaslException;
-import javax.security.sasl.SaslServer;
-import javax.security.sasl.SaslServerFactory;
-
-import org.apache.kafka.common.errors.IllegalSaslStateException;
-import org.apache.kafka.common.errors.SaslAuthenticationException;
-import org.apache.kafka.common.security.scram.ScramMessages.ClientFinalMessage;
-import org.apache.kafka.common.security.scram.ScramMessages.ClientFirstMessage;
-import org.apache.kafka.common.security.scram.ScramMessages.ServerFinalMessage;
-import org.apache.kafka.common.security.scram.ScramMessages.ServerFirstMessage;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-/**
- * SaslServer implementation for SASL/SCRAM. This server is configured with a callback
- * handler for integration with a credential manager. Kafka brokers provide callbacks
- * based on a Zookeeper-based password store.
- *
- * @see RFC 5802
- */
-public class ScramSaslServer implements SaslServer {
-
-    private static final Logger log = LoggerFactory.getLogger(ScramSaslServer.class);
-
-    enum State {
-        RECEIVE_CLIENT_FIRST_MESSAGE,
-        RECEIVE_CLIENT_FINAL_MESSAGE,
-        COMPLETE,
-        FAILED
-    };
-
-    private final ScramMechanism mechanism;
-    private final ScramFormatter formatter;
-    private final CallbackHandler callbackHandler;
-    private State state;
-    private String username;
-    private ClientFirstMessage clientFirstMessage;
-    private ServerFirstMessage serverFirstMessage;
-    private ScramCredential scramCredential;
-
-    public ScramSaslServer(ScramMechanism mechanism, Map props, CallbackHandler callbackHandler) throws NoSuchAlgorithmException {
-        this.mechanism = mechanism;
-        this.formatter = new ScramFormatter(mechanism);
-        this.callbackHandler = callbackHandler;
-        setState(State.RECEIVE_CLIENT_FIRST_MESSAGE);
-    }
-
-    /**
-     * @throws SaslAuthenticationException if the requested authorization id is not the same as username.
-     * 
            
-     * Note: This method may throw {@link SaslAuthenticationException} to provide custom error messages
-     * to clients. But care should be taken to avoid including any information in the exception message that
-     * should not be leaked to unauthenticated clients. It may be safer to throw {@link SaslException} in
-     * most cases so that a standard error message is returned to clients.
-     * 
            
-     */
-    @Override
-    public byte[] evaluateResponse(byte[] response) throws SaslException, SaslAuthenticationException {
-        try {
-            switch (state) {
-                case RECEIVE_CLIENT_FIRST_MESSAGE:
-                    this.clientFirstMessage = new ClientFirstMessage(response);
-                    String serverNonce = formatter.secureRandomString();
-                    try {
-                        String saslName = clientFirstMessage.saslName();
-                        this.username = formatter.username(saslName);
-                        NameCallback nameCallback = new NameCallback("username", username);
-                        ScramCredentialCallback credentialCallback = new ScramCredentialCallback();
-                        callbackHandler.handle(new Callback[]{nameCallback, credentialCallback});
-                        this.scramCredential = credentialCallback.scramCredential();
-                        if (scramCredential == null)
-                            throw new SaslException("Authentication failed: Invalid user credentials");
-                        String authorizationIdFromClient = clientFirstMessage.authorizationId();
-                        if (!authorizationIdFromClient.isEmpty() && !authorizationIdFromClient.equals(username))
-                            throw new SaslAuthenticationException("Authentication failed: Client requested an authorization id that is different from username");
-
-                        if (scramCredential.iterations() < mechanism.minIterations())
-                            throw new SaslException("Iterations " + scramCredential.iterations() +  " is less than the minimum " + mechanism.minIterations() + " for " + mechanism);
-                        this.serverFirstMessage = new ServerFirstMessage(clientFirstMessage.nonce(),
-                                serverNonce,
-                                scramCredential.salt(),
-                                scramCredential.iterations());
-                        setState(State.RECEIVE_CLIENT_FINAL_MESSAGE);
-                        return serverFirstMessage.toBytes();
-                    } catch (IOException | NumberFormatException | UnsupportedCallbackException e) {
-                        throw new SaslException("Authentication failed: Credentials could not be obtained", e);
-                    }
-
-                case RECEIVE_CLIENT_FINAL_MESSAGE:
-                    try {
-                        ClientFinalMessage clientFinalMessage = new ClientFinalMessage(response);
-                        verifyClientProof(clientFinalMessage);
-                        byte[] serverKey = scramCredential.serverKey();
-                        byte[] serverSignature = formatter.serverSignature(serverKey, clientFirstMessage, serverFirstMessage, clientFinalMessage);
-                        ServerFinalMessage serverFinalMessage = new ServerFinalMessage(null, serverSignature);
-                        clearCredentials();
-                        setState(State.COMPLETE);
-                        return serverFinalMessage.toBytes();
-                    } catch (InvalidKeyException e) {
-                        throw new SaslException("Authentication failed: Invalid client final message", e);
-                    }
-
-                default:
-                    throw new IllegalSaslStateException("Unexpected challenge in Sasl server state " + state);
-            }
-        } catch (SaslException e) {
-            clearCredentials();
-            setState(State.FAILED);
-            throw e;
-        }
-    }
-
-    @Override
-    public String getAuthorizationID() {
-        if (!isComplete())
-            throw new IllegalStateException("Authentication exchange has not completed");
-        return username;
-    }
-
-    @Override
-    public String getMechanismName() {
-        return mechanism.mechanismName();
-    }
-
-    @Override
-    public Object getNegotiatedProperty(String propName) {
-        if (!isComplete())
-            throw new IllegalStateException("Authentication exchange has not completed");
-        return null;
-    }
-
-    @Override
-    public boolean isComplete() {
-        return state == State.COMPLETE;
-    }
-
-    @Override
-    public byte[] unwrap(byte[] incoming, int offset, int len) throws SaslException {
-        if (!isComplete())
-            throw new IllegalStateException("Authentication exchange has not completed");
-        return Arrays.copyOfRange(incoming, offset, offset + len);
-    }
-
-    @Override
-    public byte[] wrap(byte[] outgoing, int offset, int len) throws SaslException {
-        if (!isComplete())
-            throw new IllegalStateException("Authentication exchange has not completed");
-        return Arrays.copyOfRange(outgoing, offset, offset + len);
-    }
-
-    @Override
-    public void dispose() throws SaslException {
-    }
-
-    private void setState(State state) {
-        log.debug("Setting SASL/{} server state to {}", mechanism, state);
-        this.state = state;
-    }
-
-    private void verifyClientProof(ClientFinalMessage clientFinalMessage) throws SaslException {
-        try {
-            byte[] expectedStoredKey = scramCredential.storedKey();
-            byte[] clientSignature = formatter.clientSignature(expectedStoredKey, clientFirstMessage, serverFirstMessage, clientFinalMessage);
-            byte[] computedStoredKey = formatter.storedKey(clientSignature, clientFinalMessage.proof());
-            if (!Arrays.equals(computedStoredKey, expectedStoredKey))
-                throw new SaslException("Invalid client credentials");
-        } catch (InvalidKeyException e) {
-            throw new SaslException("Sasl client verification failed", e);
-        }
-    }
-
-    private void clearCredentials() {
-        scramCredential = null;
-        clientFirstMessage = null;
-        serverFirstMessage = null;
-    }
-
-    public static class ScramSaslServerFactory implements SaslServerFactory {
-
-        @Override
-        public SaslServer createSaslServer(String mechanism, String protocol, String serverName, Map props, CallbackHandler cbh)
-            throws SaslException {
-
-            if (!ScramMechanism.isScram(mechanism)) {
-                throw new SaslException(String.format("Requested mechanism '%s' is not supported. Supported mechanisms are '%s'.",
-                        mechanism, ScramMechanism.mechanismNames()));
-            }
-            try {
-                return new ScramSaslServer(ScramMechanism.forMechanismName(mechanism), props, cbh);
-            } catch (NoSuchAlgorithmException e) {
-                throw new SaslException("Hash algorithm not supported for mechanism " + mechanism, e);
-            }
-        }
-
-        @Override
-        public String[] getMechanismNames(Map props) {
-            Collection mechanisms = ScramMechanism.mechanismNames();
-            return mechanisms.toArray(new String[mechanisms.size()]);
-        }
-    }
-}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/scram/ScramServerCallbackHandler.java b/clients/src/main/java/org/apache/kafka/common/security/scram/ScramServerCallbackHandler.java
deleted file mode 100644
index d3b245d7ba37f..0000000000000
--- a/clients/src/main/java/org/apache/kafka/common/security/scram/ScramServerCallbackHandler.java
+++ /dev/null
@@ -1,59 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements. See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License. You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.kafka.common.security.scram;
-
-import java.io.IOException;
-import java.util.Map;
-
-import javax.security.auth.Subject;
-import javax.security.auth.callback.Callback;
-import javax.security.auth.callback.NameCallback;
-import javax.security.auth.callback.UnsupportedCallbackException;
-
-import org.apache.kafka.common.network.Mode;
-import org.apache.kafka.common.security.authenticator.AuthCallbackHandler;
-import org.apache.kafka.common.security.authenticator.CredentialCache;
-
-public class ScramServerCallbackHandler implements AuthCallbackHandler {
-
-    private final CredentialCache.Cache credentialCache;
-
-    public ScramServerCallbackHandler(CredentialCache.Cache credentialCache) {
-        this.credentialCache = credentialCache;
-    }
-
-    @Override
-    public void handle(Callback[] callbacks) throws IOException, UnsupportedCallbackException {
-        String username = null;
-        for (Callback callback : callbacks) {
-            if (callback instanceof NameCallback)
-                username = ((NameCallback) callback).getDefaultName();
-            else if (callback instanceof ScramCredentialCallback)
-                ((ScramCredentialCallback) callback).scramCredential(credentialCache.get(username));
-            else
-                throw new UnsupportedCallbackException(callback);
-        }
-    }
-
-    @Override
-    public void configure(Map configs, Mode mode, Subject subject, String saslMechanism) {
-    }
-
-    @Override
-    public void close() {
-    }
-}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/scram/ScramCredentialUtils.java b/clients/src/main/java/org/apache/kafka/common/security/scram/internals/ScramCredentialUtils.java
similarity index 79%
rename from clients/src/main/java/org/apache/kafka/common/security/scram/ScramCredentialUtils.java
rename to clients/src/main/java/org/apache/kafka/common/security/scram/internals/ScramCredentialUtils.java
index 55b0651e8f048..0ce51a5ea16e4 100644
--- a/clients/src/main/java/org/apache/kafka/common/security/scram/ScramCredentialUtils.java
+++ b/clients/src/main/java/org/apache/kafka/common/security/scram/internals/ScramCredentialUtils.java
@@ -14,13 +14,14 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.kafka.common.security.scram;
+package org.apache.kafka.common.security.scram.internals;
 
+import java.util.Base64;
 import java.util.Collection;
 import java.util.Properties;
 
 import org.apache.kafka.common.security.authenticator.CredentialCache;
-import org.apache.kafka.common.utils.Base64;
+import org.apache.kafka.common.security.scram.ScramCredential;
 
 /**
  * SCRAM Credential persistence utility functions. Implements format conversion used
@@ -31,20 +32,22 @@
  * 
  *
  */
-public class ScramCredentialUtils {
+public final class ScramCredentialUtils {
     private static final String SALT = "salt";
     private static final String STORED_KEY = "stored_key";
     private static final String SERVER_KEY = "server_key";
     private static final String ITERATIONS = "iterations";
 
+    private ScramCredentialUtils() {}
+
     public static String credentialToString(ScramCredential credential) {
         return String.format("%s=%s,%s=%s,%s=%s,%s=%d",
                SALT,
-               Base64.encoder().encodeToString(credential.salt()),
+               Base64.getEncoder().encodeToString(credential.salt()),
                STORED_KEY,
-                Base64.encoder().encodeToString(credential.storedKey()),
+               Base64.getEncoder().encodeToString(credential.storedKey()),
                SERVER_KEY,
-                Base64.encoder().encodeToString(credential.serverKey()),
+               Base64.getEncoder().encodeToString(credential.serverKey()),
                ITERATIONS,
                credential.iterations());
     }
@@ -55,9 +58,9 @@ public static ScramCredential credentialFromString(String str) {
                 !props.containsKey(SERVER_KEY) || !props.containsKey(ITERATIONS)) {
             throw new IllegalArgumentException("Credentials not valid: " + str);
         }
-        byte[] salt = Base64.decoder().decode(props.getProperty(SALT));
-        byte[] storedKey = Base64.decoder().decode(props.getProperty(STORED_KEY));
-        byte[] serverKey = Base64.decoder().decode(props.getProperty(SERVER_KEY));
+        byte[] salt = Base64.getDecoder().decode(props.getProperty(SALT));
+        byte[] storedKey = Base64.getDecoder().decode(props.getProperty(STORED_KEY));
+        byte[] serverKey = Base64.getDecoder().decode(props.getProperty(SERVER_KEY));
         int iterations = Integer.parseInt(props.getProperty(ITERATIONS));
         return new ScramCredential(salt, storedKey, serverKey, iterations);
     }
@@ -74,9 +77,9 @@ private static Properties toProps(String str) {
         return props;
     }
 
-    public static void createCache(CredentialCache cache, Collection enabledMechanisms) {
+    public static void createCache(CredentialCache cache, Collection mechanisms) {
         for (String mechanism : ScramMechanism.mechanismNames()) {
-            if (enabledMechanisms.contains(mechanism))
+            if (mechanisms.contains(mechanism))
                 cache.createCache(mechanism, ScramCredential.class);
         }
     }
diff --git a/clients/src/main/java/org/apache/kafka/common/security/scram/internals/ScramExtensions.java b/clients/src/main/java/org/apache/kafka/common/security/scram/internals/ScramExtensions.java
new file mode 100644
index 0000000000000..439bcfe14f218
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/scram/internals/ScramExtensions.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.scram.internals;
+
+import org.apache.kafka.common.security.auth.SaslExtensions;
+import org.apache.kafka.common.security.scram.ScramLoginModule;
+import org.apache.kafka.common.utils.Utils;
+
+import java.util.Collections;
+import java.util.Map;
+
+public class ScramExtensions extends SaslExtensions {
+
+    public ScramExtensions() {
+        this(Collections.emptyMap());
+    }
+
+    public ScramExtensions(String extensions) {
+        this(Utils.parseMap(extensions, "=", ","));
+    }
+
+    public ScramExtensions(Map extensionMap) {
+        super(extensionMap);
+    }
+
+    public boolean tokenAuthenticated() {
+        return Boolean.parseBoolean(map().get(ScramLoginModule.TOKEN_AUTH_CONFIG));
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/scram/internals/ScramFormatter.java b/clients/src/main/java/org/apache/kafka/common/security/scram/internals/ScramFormatter.java
new file mode 100644
index 0000000000000..4c03fa158b80e
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/scram/internals/ScramFormatter.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.common.security.scram.internals;
+
+import java.math.BigInteger;
+import java.nio.charset.StandardCharsets;
+import java.security.InvalidKeyException;
+import java.security.MessageDigest;
+import java.security.NoSuchAlgorithmException;
+import java.security.SecureRandom;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+
+import javax.crypto.Mac;
+import javax.crypto.spec.SecretKeySpec;
+
+import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.security.scram.ScramCredential;
+import org.apache.kafka.common.security.scram.internals.ScramMessages.ClientFinalMessage;
+import org.apache.kafka.common.security.scram.internals.ScramMessages.ClientFirstMessage;
+import org.apache.kafka.common.security.scram.internals.ScramMessages.ServerFirstMessage;
+
+/**
+ * Scram message salt and hash functions defined in RFC 5802.
+ */
+public class ScramFormatter {
+
+    private static final Pattern EQUAL = Pattern.compile("=", Pattern.LITERAL);
+    private static final Pattern COMMA = Pattern.compile(",", Pattern.LITERAL);
+    private static final Pattern EQUAL_TWO_C = Pattern.compile("=2C", Pattern.LITERAL);
+    private static final Pattern EQUAL_THREE_D = Pattern.compile("=3D", Pattern.LITERAL);
+
+    private final MessageDigest messageDigest;
+    private final Mac mac;
+    private final SecureRandom random;
+
+    public ScramFormatter(ScramMechanism mechanism) throws NoSuchAlgorithmException {
+        this.messageDigest = MessageDigest.getInstance(mechanism.hashAlgorithm());
+        this.mac = Mac.getInstance(mechanism.macAlgorithm());
+        this.random = new SecureRandom();
+    }
+
+    public byte[] hmac(byte[] key, byte[] bytes) throws InvalidKeyException {
+        mac.init(new SecretKeySpec(key, mac.getAlgorithm()));
+        return mac.doFinal(bytes);
+    }
+
+    public byte[] hash(byte[] str) {
+        return messageDigest.digest(str);
+    }
+
+    public static byte[] xor(byte[] first, byte[] second) {
+        if (first.length != second.length)
+            throw new IllegalArgumentException("Argument arrays must be of the same length");
+        byte[] result = new byte[first.length];
+        for (int i = 0; i < result.length; i++)
+            result[i] = (byte) (first[i] ^ second[i]);
+        return result;
+    }
+
+    public byte[] hi(byte[] str, byte[] salt, int iterations) throws InvalidKeyException {
+        mac.init(new SecretKeySpec(str, mac.getAlgorithm()));
+        mac.update(salt);
+        byte[] u1 = mac.doFinal(new byte[]{0, 0, 0, 1});
+        byte[] prev = u1;
+        byte[] result = u1;
+        for (int i = 2; i <= iterations; i++) {
+            byte[] ui = hmac(str, prev);
+            result = xor(result, ui);
+            prev = ui;
+        }
+        return result;
+    }
+
+    public static byte[] normalize(String str) {
+        return toBytes(str);
+    }
+
+    public byte[] saltedPassword(String password, byte[] salt, int iterations) throws InvalidKeyException {
+        return hi(normalize(password), salt, iterations);
+    }
+
+    public byte[] clientKey(byte[] saltedPassword) throws InvalidKeyException {
+        return hmac(saltedPassword, toBytes("Client Key"));
+    }
+
+    public byte[] storedKey(byte[] clientKey) {
+        return hash(clientKey);
+    }
+
+    public static String saslName(String username) {
+        String replace1 = EQUAL.matcher(username).replaceAll(Matcher.quoteReplacement("=3D"));
+        return COMMA.matcher(replace1).replaceAll(Matcher.quoteReplacement("=2C"));
+    }
+
+    public static String username(String saslName) {
+        String username = EQUAL_TWO_C.matcher(saslName).replaceAll(Matcher.quoteReplacement(","));
+        if (EQUAL_THREE_D.matcher(username).replaceAll(Matcher.quoteReplacement("")).indexOf('=') >= 0) {
+            throw new IllegalArgumentException("Invalid username: " + saslName);
+        }
+        return EQUAL_THREE_D.matcher(username).replaceAll(Matcher.quoteReplacement("="));
+    }
+
+    public static String authMessage(String clientFirstMessageBare, String serverFirstMessage, String clientFinalMessageWithoutProof) {
+        return clientFirstMessageBare + "," + serverFirstMessage + "," + clientFinalMessageWithoutProof;
+    }
+
+    public byte[] clientSignature(byte[] storedKey, ClientFirstMessage clientFirstMessage, ServerFirstMessage serverFirstMessage, ClientFinalMessage clientFinalMessage) throws InvalidKeyException {
+        byte[] authMessage = authMessage(clientFirstMessage, serverFirstMessage, clientFinalMessage);
+        return hmac(storedKey, authMessage);
+    }
+
+    public byte[] clientProof(byte[] saltedPassword, ClientFirstMessage clientFirstMessage, ServerFirstMessage serverFirstMessage, ClientFinalMessage clientFinalMessage) throws InvalidKeyException {
+        byte[] clientKey = clientKey(saltedPassword);
+        byte[] storedKey = hash(clientKey);
+        byte[] clientSignature = hmac(storedKey, authMessage(clientFirstMessage, serverFirstMessage, clientFinalMessage));
+        return xor(clientKey, clientSignature);
+    }
+
+    private byte[] authMessage(ClientFirstMessage clientFirstMessage, ServerFirstMessage serverFirstMessage, ClientFinalMessage clientFinalMessage) {
+        return toBytes(authMessage(clientFirstMessage.clientFirstMessageBare(),
+                serverFirstMessage.toMessage(),
+                clientFinalMessage.clientFinalMessageWithoutProof()));
+    }
+
+    public byte[] storedKey(byte[] clientSignature, byte[] clientProof) {
+        return hash(xor(clientSignature, clientProof));
+    }
+
+    public byte[] serverKey(byte[] saltedPassword) throws InvalidKeyException {
+        return hmac(saltedPassword, toBytes("Server Key"));
+    }
+
+    public byte[] serverSignature(byte[] serverKey, ClientFirstMessage clientFirstMessage, ServerFirstMessage serverFirstMessage, ClientFinalMessage clientFinalMessage) throws InvalidKeyException {
+        byte[] authMessage = authMessage(clientFirstMessage, serverFirstMessage, clientFinalMessage);
+        return hmac(serverKey, authMessage);
+    }
+
+    public String secureRandomString() {
+        return secureRandomString(random);
+    }
+
+    public static String secureRandomString(SecureRandom random) {
+        return new BigInteger(130, random).toString(Character.MAX_RADIX);
+    }
+
+    public byte[] secureRandomBytes() {
+        return secureRandomBytes(random);
+    }
+
+    public static byte[] secureRandomBytes(SecureRandom random) {
+        return toBytes(secureRandomString(random));
+    }
+
+    public static byte[] toBytes(String str) {
+        return str.getBytes(StandardCharsets.UTF_8);
+    }
+
+    public ScramCredential generateCredential(String password, int iterations) {
+        try {
+            byte[] salt = secureRandomBytes();
+            byte[] saltedPassword = saltedPassword(password, salt, iterations);
+            return generateCredential(salt, saltedPassword, iterations);
+        } catch (InvalidKeyException e) {
+            throw new KafkaException("Could not create credential", e);
+        }
+    }
+
+    public ScramCredential generateCredential(byte[] salt, byte[] saltedPassword, int iterations) {
+        try {
+            byte[] clientKey = clientKey(saltedPassword);
+            byte[] storedKey = storedKey(clientKey);
+            byte[] serverKey = serverKey(saltedPassword);
+            return new ScramCredential(salt, storedKey, serverKey, iterations);
+        } catch (InvalidKeyException e) {
+            throw new KafkaException("Could not create credential", e);
+        }
+    }
+}
\ No newline at end of file
diff --git a/clients/src/main/java/org/apache/kafka/common/security/scram/internals/ScramMechanism.java b/clients/src/main/java/org/apache/kafka/common/security/scram/internals/ScramMechanism.java
new file mode 100644
index 0000000000000..9f6e69d9c5e05
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/scram/internals/ScramMechanism.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.common.security.scram.internals;
+
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+public enum ScramMechanism {
+
+    SCRAM_SHA_256("SHA-256", "HmacSHA256", 4096),
+    SCRAM_SHA_512("SHA-512", "HmacSHA512", 4096);
+
+    private final String mechanismName;
+    private final String hashAlgorithm;
+    private final String macAlgorithm;
+    private final int minIterations;
+
+    private static final Map MECHANISMS_MAP;
+
+    static {
+        Map map = new HashMap<>();
+        for (ScramMechanism mech : values())
+            map.put(mech.mechanismName, mech);
+        MECHANISMS_MAP = Collections.unmodifiableMap(map);
+    }
+
+    ScramMechanism(String hashAlgorithm, String macAlgorithm, int minIterations) {
+        this.mechanismName = "SCRAM-" + hashAlgorithm;
+        this.hashAlgorithm = hashAlgorithm;
+        this.macAlgorithm = macAlgorithm;
+        this.minIterations = minIterations;
+    }
+
+    public final String mechanismName() {
+        return mechanismName;
+    }
+
+    public String hashAlgorithm() {
+        return hashAlgorithm;
+    }
+
+    public String macAlgorithm() {
+        return macAlgorithm;
+    }
+
+    public int minIterations() {
+        return minIterations;
+    }
+
+    public static ScramMechanism forMechanismName(String mechanismName) {
+        return MECHANISMS_MAP.get(mechanismName);
+    }
+
+    public static Collection mechanismNames() {
+        return MECHANISMS_MAP.keySet();
+    }
+
+    public static boolean isScram(String mechanismName) {
+        return MECHANISMS_MAP.containsKey(mechanismName);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/scram/ScramMessages.java b/clients/src/main/java/org/apache/kafka/common/security/scram/internals/ScramMessages.java
similarity index 85%
rename from clients/src/main/java/org/apache/kafka/common/security/scram/ScramMessages.java
rename to clients/src/main/java/org/apache/kafka/common/security/scram/internals/ScramMessages.java
index 1ad7266e7bf4d..05512962906a9 100644
--- a/clients/src/main/java/org/apache/kafka/common/security/scram/ScramMessages.java
+++ b/clients/src/main/java/org/apache/kafka/common/security/scram/internals/ScramMessages.java
@@ -14,11 +14,13 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.kafka.common.security.scram;
+package org.apache.kafka.common.security.scram.internals;
 
-import org.apache.kafka.common.utils.Base64;
+import org.apache.kafka.common.utils.Utils;
 
 import java.nio.charset.StandardCharsets;
+import java.util.Base64;
+import java.util.Map;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
@@ -64,16 +66,18 @@ protected String toMessage(byte[] messageBytes) {
      */
     public static class ClientFirstMessage extends AbstractScramMessage {
         private static final Pattern PATTERN = Pattern.compile(String.format(
-                "n,(a=(?%s))?,%sn=(?%s),r=(?%s)%s",
+                "n,(a=(?%s))?,%sn=(?%s),r=(?%s)(?%s)",
                 SASLNAME,
                 RESERVED,
                 SASLNAME,
                 PRINTABLE,
                 EXTENSIONS));
 
+
         private final String saslName;
         private final String nonce;
         private final String authorizationId;
+        private final ScramExtensions extensions;
         public ClientFirstMessage(byte[] messageBytes) throws SaslException {
             String message = toMessage(messageBytes);
             Matcher matcher = PATTERN.matcher(message);
@@ -83,10 +87,14 @@ public ClientFirstMessage(byte[] messageBytes) throws SaslException {
             this.authorizationId = authzid != null ? authzid : "";
             this.saslName = matcher.group("saslname");
             this.nonce = matcher.group("nonce");
+            String extString = matcher.group("extensions");
+
+            this.extensions = extString.startsWith(",") ? new ScramExtensions(extString.substring(1)) : new ScramExtensions();
         }
-        public ClientFirstMessage(String saslName, String nonce) {
+        public ClientFirstMessage(String saslName, String nonce, Map extensions) {
             this.saslName = saslName;
             this.nonce = nonce;
+            this.extensions = new ScramExtensions(extensions);
             this.authorizationId = ""; // Optional authzid not specified in gs2-header
         }
         public String saslName() {
@@ -101,8 +109,17 @@ public String authorizationId() {
         public String gs2Header() {
             return "n," + authorizationId + ",";
         }
+        public ScramExtensions extensions() {
+            return extensions;
+        }
+
         public String clientFirstMessageBare() {
-            return String.format("n=%s,r=%s", saslName, nonce);
+            String extensionStr = Utils.mkString(extensions.map(), "", "", "=", ",");
+
+            if (extensionStr.isEmpty())
+                return String.format("n=%s,r=%s", saslName, nonce);
+            else
+                return String.format("n=%s,r=%s,%s", saslName, nonce, extensionStr);
         }
         String toMessage() {
             return gs2Header() + clientFirstMessageBare();
@@ -141,7 +158,7 @@ public ServerFirstMessage(byte[] messageBytes) throws SaslException {
             }
             this.nonce = matcher.group("nonce");
             String salt = matcher.group("salt");
-            this.salt = Base64.decoder().decode(salt);
+            this.salt = Base64.getDecoder().decode(salt);
         }
         public ServerFirstMessage(String clientNonce, String serverNonce, byte[] salt, int iterations) {
             this.nonce = clientNonce + serverNonce;
@@ -158,7 +175,7 @@ public int iterations() {
             return iterations;
         }
         String toMessage() {
-            return String.format("r=%s,s=%s,i=%d", nonce, Base64.encoder().encodeToString(salt), iterations);
+            return String.format("r=%s,s=%s,i=%d", nonce, Base64.getEncoder().encodeToString(salt), iterations);
         }
     }
     /**
@@ -185,9 +202,9 @@ public ClientFinalMessage(byte[] messageBytes) throws SaslException {
             if (!matcher.matches())
                 throw new SaslException("Invalid SCRAM client final message format: " + message);
 
-            this.channelBinding = Base64.decoder().decode(matcher.group("channel"));
+            this.channelBinding = Base64.getDecoder().decode(matcher.group("channel"));
             this.nonce = matcher.group("nonce");
-            this.proof = Base64.decoder().decode(matcher.group("proof"));
+            this.proof = Base64.getDecoder().decode(matcher.group("proof"));
         }
         public ClientFinalMessage(byte[] channelBinding, String nonce) {
             this.channelBinding = channelBinding;
@@ -207,13 +224,13 @@ public void proof(byte[] proof) {
         }
         public String clientFinalMessageWithoutProof() {
             return String.format("c=%s,r=%s",
-                    Base64.encoder().encodeToString(channelBinding),
+                    Base64.getEncoder().encodeToString(channelBinding),
                     nonce);
         }
         String toMessage() {
             return String.format("%s,p=%s",
                     clientFinalMessageWithoutProof(),
-                    Base64.encoder().encodeToString(proof));
+                    Base64.getEncoder().encodeToString(proof));
         }
     }
     /**
@@ -244,7 +261,7 @@ public ServerFinalMessage(byte[] messageBytes) throws SaslException {
                 // ignore
             }
             if (error == null) {
-                this.serverSignature = Base64.decoder().decode(matcher.group("signature"));
+                this.serverSignature = Base64.getDecoder().decode(matcher.group("signature"));
                 this.error = null;
             } else {
                 this.serverSignature = null;
@@ -265,7 +282,7 @@ String toMessage() {
             if (error != null)
                 return "e=" + error;
             else
-                return "v=" + Base64.encoder().encodeToString(serverSignature);
+                return "v=" + Base64.getEncoder().encodeToString(serverSignature);
         }
     }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/security/scram/ScramSaslClient.java b/clients/src/main/java/org/apache/kafka/common/security/scram/internals/ScramSaslClient.java
similarity index 85%
rename from clients/src/main/java/org/apache/kafka/common/security/scram/ScramSaslClient.java
rename to clients/src/main/java/org/apache/kafka/common/security/scram/internals/ScramSaslClient.java
index e5b0f84744b51..6ee2e31d90f09 100644
--- a/clients/src/main/java/org/apache/kafka/common/security/scram/ScramSaslClient.java
+++ b/clients/src/main/java/org/apache/kafka/common/security/scram/internals/ScramSaslClient.java
@@ -14,9 +14,8 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.kafka.common.security.scram;
+package org.apache.kafka.common.security.scram.internals;
 
-import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.security.InvalidKeyException;
 import java.security.NoSuchAlgorithmException;
@@ -34,9 +33,10 @@
 import javax.security.sasl.SaslException;
 
 import org.apache.kafka.common.errors.IllegalSaslStateException;
-import org.apache.kafka.common.security.scram.ScramMessages.ClientFinalMessage;
-import org.apache.kafka.common.security.scram.ScramMessages.ServerFinalMessage;
-import org.apache.kafka.common.security.scram.ScramMessages.ServerFirstMessage;
+import org.apache.kafka.common.security.scram.ScramExtensionsCallback;
+import org.apache.kafka.common.security.scram.internals.ScramMessages.ClientFinalMessage;
+import org.apache.kafka.common.security.scram.internals.ScramMessages.ServerFinalMessage;
+import org.apache.kafka.common.security.scram.internals.ScramMessages.ServerFirstMessage;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -59,7 +59,7 @@ enum State {
         RECEIVE_SERVER_FINAL_MESSAGE,
         COMPLETE,
         FAILED
-    };
+    }
 
     private final ScramMechanism mechanism;
     private final CallbackHandler callbackHandler;
@@ -97,14 +97,24 @@ public byte[] evaluateChallenge(byte[] challenge) throws SaslException {
                         throw new SaslException("Expected empty challenge");
                     clientNonce = formatter.secureRandomString();
                     NameCallback nameCallback = new NameCallback("Name:");
+                    ScramExtensionsCallback extensionsCallback = new ScramExtensionsCallback();
+
                     try {
                         callbackHandler.handle(new Callback[]{nameCallback});
-                    } catch (IOException | UnsupportedCallbackException e) {
-                        throw new SaslException("User name could not be obtained", e);
+                        try {
+                            callbackHandler.handle(new Callback[]{extensionsCallback});
+                        } catch (UnsupportedCallbackException e) {
+                            log.debug("Extensions callback is not supported by client callback handler {}, no extensions will be added",
+                                    callbackHandler);
+                        }
+                    } catch (Throwable e) {
+                        throw new SaslException("User name or extensions could not be obtained", e);
                     }
+
                     String username = nameCallback.getName();
-                    String saslName = formatter.saslName(username);
-                    this.clientFirstMessage = new ScramMessages.ClientFirstMessage(saslName, clientNonce);
+                    String saslName = ScramFormatter.saslName(username);
+                    Map extensions = extensionsCallback.extensions();
+                    this.clientFirstMessage = new ScramMessages.ClientFirstMessage(saslName, clientNonce, extensions);
                     setState(State.RECEIVE_SERVER_FIRST_MESSAGE);
                     return clientFirstMessage.toBytes();
 
@@ -117,7 +127,7 @@ public byte[] evaluateChallenge(byte[] challenge) throws SaslException {
                     PasswordCallback passwordCallback = new PasswordCallback("Password:", false);
                     try {
                         callbackHandler.handle(new Callback[]{passwordCallback});
-                    } catch (IOException | UnsupportedCallbackException e) {
+                    } catch (Throwable e) {
                         throw new SaslException("User name could not be obtained", e);
                     }
                     this.clientFinalMessage = handleServerFirstMessage(passwordCallback.getPassword());
@@ -147,14 +157,14 @@ public boolean isComplete() {
     }
 
     @Override
-    public byte[] unwrap(byte[] incoming, int offset, int len) throws SaslException {
+    public byte[] unwrap(byte[] incoming, int offset, int len) {
         if (!isComplete())
             throw new IllegalStateException("Authentication exchange has not completed");
         return Arrays.copyOfRange(incoming, offset, offset + len);
     }
 
     @Override
-    public byte[] wrap(byte[] outgoing, int offset, int len) throws SaslException {
+    public byte[] wrap(byte[] outgoing, int offset, int len) {
         if (!isComplete())
             throw new IllegalStateException("Authentication exchange has not completed");
         return Arrays.copyOfRange(outgoing, offset, offset + len);
@@ -168,7 +178,7 @@ public Object getNegotiatedProperty(String propName) {
     }
 
     @Override
-    public void dispose() throws SaslException {
+    public void dispose() {
     }
 
     private void setState(State state) {
@@ -178,7 +188,7 @@ private void setState(State state) {
 
     private ClientFinalMessage handleServerFirstMessage(char[] password) throws SaslException {
         try {
-            byte[] passwordBytes = formatter.normalize(new String(password));
+            byte[] passwordBytes = ScramFormatter.normalize(new String(password));
             this.saltedPassword = formatter.hi(passwordBytes, serverFirstMessage.salt(), serverFirstMessage.iterations());
 
             ClientFinalMessage clientFinalMessage = new ClientFinalMessage("n,,".getBytes(StandardCharsets.UTF_8), serverFirstMessage.nonce());
diff --git a/clients/src/main/java/org/apache/kafka/common/security/scram/ScramSaslClientProvider.java b/clients/src/main/java/org/apache/kafka/common/security/scram/internals/ScramSaslClientProvider.java
similarity index 90%
rename from clients/src/main/java/org/apache/kafka/common/security/scram/ScramSaslClientProvider.java
rename to clients/src/main/java/org/apache/kafka/common/security/scram/internals/ScramSaslClientProvider.java
index d389f044a9967..8c5b85a1a9277 100644
--- a/clients/src/main/java/org/apache/kafka/common/security/scram/ScramSaslClientProvider.java
+++ b/clients/src/main/java/org/apache/kafka/common/security/scram/internals/ScramSaslClientProvider.java
@@ -14,12 +14,12 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.kafka.common.security.scram;
+package org.apache.kafka.common.security.scram.internals;
 
 import java.security.Provider;
 import java.security.Security;
 
-import org.apache.kafka.common.security.scram.ScramSaslClient.ScramSaslClientFactory;
+import org.apache.kafka.common.security.scram.internals.ScramSaslClient.ScramSaslClientFactory;
 
 public class ScramSaslClientProvider extends Provider {
 
diff --git a/clients/src/main/java/org/apache/kafka/common/security/scram/internals/ScramSaslServer.java b/clients/src/main/java/org/apache/kafka/common/security/scram/internals/ScramSaslServer.java
new file mode 100644
index 0000000000000..f9bce80360e33
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/scram/internals/ScramSaslServer.java
@@ -0,0 +1,265 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.scram.internals;
+
+import java.security.InvalidKeyException;
+import java.security.NoSuchAlgorithmException;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Map;
+import java.util.Set;
+
+import javax.security.auth.callback.Callback;
+import javax.security.auth.callback.CallbackHandler;
+import javax.security.auth.callback.NameCallback;
+import javax.security.sasl.SaslException;
+import javax.security.sasl.SaslServer;
+import javax.security.sasl.SaslServerFactory;
+
+import org.apache.kafka.common.errors.AuthenticationException;
+import org.apache.kafka.common.errors.IllegalSaslStateException;
+import org.apache.kafka.common.errors.SaslAuthenticationException;
+import org.apache.kafka.common.security.authenticator.SaslInternalConfigs;
+import org.apache.kafka.common.security.scram.ScramCredential;
+import org.apache.kafka.common.security.scram.ScramCredentialCallback;
+import org.apache.kafka.common.security.scram.ScramLoginModule;
+import org.apache.kafka.common.security.scram.internals.ScramMessages.ClientFinalMessage;
+import org.apache.kafka.common.security.scram.internals.ScramMessages.ClientFirstMessage;
+import org.apache.kafka.common.security.scram.internals.ScramMessages.ServerFinalMessage;
+import org.apache.kafka.common.security.scram.internals.ScramMessages.ServerFirstMessage;
+import org.apache.kafka.common.security.token.delegation.internals.DelegationTokenCredentialCallback;
+import org.apache.kafka.common.utils.Utils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * SaslServer implementation for SASL/SCRAM. This server is configured with a callback
+ * handler for integration with a credential manager. Kafka brokers provide callbacks
+ * based on a Zookeeper-based password store.
+ *
+ * @see RFC 5802
+ */
+public class ScramSaslServer implements SaslServer {
+
+    private static final Logger log = LoggerFactory.getLogger(ScramSaslServer.class);
+    private static final Set SUPPORTED_EXTENSIONS = Utils.mkSet(ScramLoginModule.TOKEN_AUTH_CONFIG);
+
+    enum State {
+        RECEIVE_CLIENT_FIRST_MESSAGE,
+        RECEIVE_CLIENT_FINAL_MESSAGE,
+        COMPLETE,
+        FAILED
+    }
+
+    private final ScramMechanism mechanism;
+    private final ScramFormatter formatter;
+    private final CallbackHandler callbackHandler;
+    private State state;
+    private String username;
+    private ClientFirstMessage clientFirstMessage;
+    private ServerFirstMessage serverFirstMessage;
+    private ScramExtensions scramExtensions;
+    private ScramCredential scramCredential;
+    private String authorizationId;
+    private Long tokenExpiryTimestamp;
+
+    public ScramSaslServer(ScramMechanism mechanism, Map props, CallbackHandler callbackHandler) throws NoSuchAlgorithmException {
+        this.mechanism = mechanism;
+        this.formatter = new ScramFormatter(mechanism);
+        this.callbackHandler = callbackHandler;
+        setState(State.RECEIVE_CLIENT_FIRST_MESSAGE);
+    }
+
+    /**
+     * @throws SaslAuthenticationException if the requested authorization id is not the same as username.
+     * 
            
+     * Note: This method may throw {@link SaslAuthenticationException} to provide custom error messages
+     * to clients. But care should be taken to avoid including any information in the exception message that
+     * should not be leaked to unauthenticated clients. It may be safer to throw {@link SaslException} in
+     * most cases so that a standard error message is returned to clients.
+     * 
            
+     */
+    @Override
+    public byte[] evaluateResponse(byte[] response) throws SaslException, SaslAuthenticationException {
+        try {
+            switch (state) {
+                case RECEIVE_CLIENT_FIRST_MESSAGE:
+                    this.clientFirstMessage = new ClientFirstMessage(response);
+                    this.scramExtensions = clientFirstMessage.extensions();
+                    if (!SUPPORTED_EXTENSIONS.containsAll(scramExtensions.map().keySet())) {
+                        log.debug("Unsupported extensions will be ignored, supported {}, provided {}",
+                                SUPPORTED_EXTENSIONS, scramExtensions.map().keySet());
+                    }
+                    String serverNonce = formatter.secureRandomString();
+                    try {
+                        String saslName = clientFirstMessage.saslName();
+                        this.username = ScramFormatter.username(saslName);
+                        NameCallback nameCallback = new NameCallback("username", username);
+                        ScramCredentialCallback credentialCallback;
+                        if (scramExtensions.tokenAuthenticated()) {
+                            DelegationTokenCredentialCallback tokenCallback = new DelegationTokenCredentialCallback();
+                            credentialCallback = tokenCallback;
+                            callbackHandler.handle(new Callback[]{nameCallback, tokenCallback});
+                            if (tokenCallback.tokenOwner() == null)
+                                throw new SaslException("Token Authentication failed: Invalid tokenId : " + username);
+                            this.authorizationId = tokenCallback.tokenOwner();
+                            this.tokenExpiryTimestamp = tokenCallback.tokenExpiryTimestamp();
+                        } else {
+                            credentialCallback = new ScramCredentialCallback();
+                            callbackHandler.handle(new Callback[]{nameCallback, credentialCallback});
+                            this.authorizationId = username;
+                            this.tokenExpiryTimestamp = null;
+                        }
+                        this.scramCredential = credentialCallback.scramCredential();
+                        if (scramCredential == null)
+                            throw new SaslException("Authentication failed: Invalid user credentials");
+                        String authorizationIdFromClient = clientFirstMessage.authorizationId();
+                        if (!authorizationIdFromClient.isEmpty() && !authorizationIdFromClient.equals(username))
+                            throw new SaslAuthenticationException("Authentication failed: Client requested an authorization id that is different from username");
+
+                        if (scramCredential.iterations() < mechanism.minIterations())
+                            throw new SaslException("Iterations " + scramCredential.iterations() +  " is less than the minimum " + mechanism.minIterations() + " for " + mechanism);
+                        this.serverFirstMessage = new ServerFirstMessage(clientFirstMessage.nonce(),
+                                serverNonce,
+                                scramCredential.salt(),
+                                scramCredential.iterations());
+                        setState(State.RECEIVE_CLIENT_FINAL_MESSAGE);
+                        return serverFirstMessage.toBytes();
+                    } catch (SaslException | AuthenticationException e) {
+                        throw e;
+                    } catch (Throwable e) {
+                        throw new SaslException("Authentication failed: Credentials could not be obtained", e);
+                    }
+
+                case RECEIVE_CLIENT_FINAL_MESSAGE:
+                    try {
+                        ClientFinalMessage clientFinalMessage = new ClientFinalMessage(response);
+                        verifyClientProof(clientFinalMessage);
+                        byte[] serverKey = scramCredential.serverKey();
+                        byte[] serverSignature = formatter.serverSignature(serverKey, clientFirstMessage, serverFirstMessage, clientFinalMessage);
+                        ServerFinalMessage serverFinalMessage = new ServerFinalMessage(null, serverSignature);
+                        clearCredentials();
+                        setState(State.COMPLETE);
+                        return serverFinalMessage.toBytes();
+                    } catch (InvalidKeyException e) {
+                        throw new SaslException("Authentication failed: Invalid client final message", e);
+                    }
+
+                default:
+                    throw new IllegalSaslStateException("Unexpected challenge in Sasl server state " + state);
+            }
+        } catch (SaslException | AuthenticationException e) {
+            clearCredentials();
+            setState(State.FAILED);
+            throw e;
+        }
+    }
+
+    @Override
+    public String getAuthorizationID() {
+        if (!isComplete())
+            throw new IllegalStateException("Authentication exchange has not completed");
+        return authorizationId;
+    }
+
+    @Override
+    public String getMechanismName() {
+        return mechanism.mechanismName();
+    }
+
+    @Override
+    public Object getNegotiatedProperty(String propName) {
+        if (!isComplete())
+            throw new IllegalStateException("Authentication exchange has not completed");
+        if (SaslInternalConfigs.CREDENTIAL_LIFETIME_MS_SASL_NEGOTIATED_PROPERTY_KEY.equals(propName))
+            return tokenExpiryTimestamp; // will be null if token not used
+        if (SUPPORTED_EXTENSIONS.contains(propName))
+            return scramExtensions.map().get(propName);
+        else
+            return null;
+    }
+
+    @Override
+    public boolean isComplete() {
+        return state == State.COMPLETE;
+    }
+
+    @Override
+    public byte[] unwrap(byte[] incoming, int offset, int len) {
+        if (!isComplete())
+            throw new IllegalStateException("Authentication exchange has not completed");
+        return Arrays.copyOfRange(incoming, offset, offset + len);
+    }
+
+    @Override
+    public byte[] wrap(byte[] outgoing, int offset, int len) {
+        if (!isComplete())
+            throw new IllegalStateException("Authentication exchange has not completed");
+        return Arrays.copyOfRange(outgoing, offset, offset + len);
+    }
+
+    @Override
+    public void dispose() {
+    }
+
+    private void setState(State state) {
+        log.debug("Setting SASL/{} server state to {}", mechanism, state);
+        this.state = state;
+    }
+
+    private void verifyClientProof(ClientFinalMessage clientFinalMessage) throws SaslException {
+        try {
+            byte[] expectedStoredKey = scramCredential.storedKey();
+            byte[] clientSignature = formatter.clientSignature(expectedStoredKey, clientFirstMessage, serverFirstMessage, clientFinalMessage);
+            byte[] computedStoredKey = formatter.storedKey(clientSignature, clientFinalMessage.proof());
+            if (!Arrays.equals(computedStoredKey, expectedStoredKey))
+                throw new SaslException("Invalid client credentials");
+        } catch (InvalidKeyException e) {
+            throw new SaslException("Sasl client verification failed", e);
+        }
+    }
+
+    private void clearCredentials() {
+        scramCredential = null;
+        clientFirstMessage = null;
+        serverFirstMessage = null;
+    }
+
+    public static class ScramSaslServerFactory implements SaslServerFactory {
+
+        @Override
+        public SaslServer createSaslServer(String mechanism, String protocol, String serverName, Map props, CallbackHandler cbh)
+            throws SaslException {
+
+            if (!ScramMechanism.isScram(mechanism)) {
+                throw new SaslException(String.format("Requested mechanism '%s' is not supported. Supported mechanisms are '%s'.",
+                        mechanism, ScramMechanism.mechanismNames()));
+            }
+            try {
+                return new ScramSaslServer(ScramMechanism.forMechanismName(mechanism), props, cbh);
+            } catch (NoSuchAlgorithmException e) {
+                throw new SaslException("Hash algorithm not supported for mechanism " + mechanism, e);
+            }
+        }
+
+        @Override
+        public String[] getMechanismNames(Map props) {
+            Collection mechanisms = ScramMechanism.mechanismNames();
+            return mechanisms.toArray(new String[mechanisms.size()]);
+        }
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/scram/ScramSaslServerProvider.java b/clients/src/main/java/org/apache/kafka/common/security/scram/internals/ScramSaslServerProvider.java
similarity index 90%
rename from clients/src/main/java/org/apache/kafka/common/security/scram/ScramSaslServerProvider.java
rename to clients/src/main/java/org/apache/kafka/common/security/scram/internals/ScramSaslServerProvider.java
index 9f2a6b3d25621..6a868600722ee 100644
--- a/clients/src/main/java/org/apache/kafka/common/security/scram/ScramSaslServerProvider.java
+++ b/clients/src/main/java/org/apache/kafka/common/security/scram/internals/ScramSaslServerProvider.java
@@ -14,12 +14,12 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.kafka.common.security.scram;
+package org.apache.kafka.common.security.scram.internals;
 
 import java.security.Provider;
 import java.security.Security;
 
-import org.apache.kafka.common.security.scram.ScramSaslServer.ScramSaslServerFactory;
+import org.apache.kafka.common.security.scram.internals.ScramSaslServer.ScramSaslServerFactory;
 
 public class ScramSaslServerProvider extends Provider {
 
diff --git a/clients/src/main/java/org/apache/kafka/common/security/scram/internals/ScramServerCallbackHandler.java b/clients/src/main/java/org/apache/kafka/common/security/scram/internals/ScramServerCallbackHandler.java
new file mode 100644
index 0000000000000..1af38e9d41801
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/scram/internals/ScramServerCallbackHandler.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.common.security.scram.internals;
+
+import java.util.List;
+import java.util.Map;
+
+import javax.security.auth.callback.Callback;
+import javax.security.auth.callback.NameCallback;
+import javax.security.auth.callback.UnsupportedCallbackException;
+import javax.security.auth.login.AppConfigurationEntry;
+
+import org.apache.kafka.common.security.auth.AuthenticateCallbackHandler;
+import org.apache.kafka.common.security.authenticator.CredentialCache;
+import org.apache.kafka.common.security.scram.ScramCredential;
+import org.apache.kafka.common.security.scram.ScramCredentialCallback;
+import org.apache.kafka.common.security.token.delegation.TokenInformation;
+import org.apache.kafka.common.security.token.delegation.internals.DelegationTokenCache;
+import org.apache.kafka.common.security.token.delegation.internals.DelegationTokenCredentialCallback;
+
+public class ScramServerCallbackHandler implements AuthenticateCallbackHandler {
+
+    private final CredentialCache.Cache credentialCache;
+    private final DelegationTokenCache tokenCache;
+    private String saslMechanism;
+
+    public ScramServerCallbackHandler(CredentialCache.Cache credentialCache,
+                                      DelegationTokenCache tokenCache) {
+        this.credentialCache = credentialCache;
+        this.tokenCache = tokenCache;
+    }
+
+    @Override
+    public void configure(Map configs, String mechanism, List jaasConfigEntries) {
+        this.saslMechanism = mechanism;
+    }
+
+    @Override
+    public void handle(Callback[] callbacks) throws UnsupportedCallbackException {
+        String username = null;
+        for (Callback callback : callbacks) {
+            if (callback instanceof NameCallback)
+                username = ((NameCallback) callback).getDefaultName();
+            else if (callback instanceof DelegationTokenCredentialCallback) {
+                DelegationTokenCredentialCallback tokenCallback = (DelegationTokenCredentialCallback) callback;
+                tokenCallback.scramCredential(tokenCache.credential(saslMechanism, username));
+                tokenCallback.tokenOwner(tokenCache.owner(username));
+                TokenInformation tokenInfo = tokenCache.token(username);
+                if (tokenInfo != null)
+                    tokenCallback.tokenExpiryTimestamp(tokenInfo.expiryTimestamp());
+            } else if (callback instanceof ScramCredentialCallback) {
+                ScramCredentialCallback sc = (ScramCredentialCallback) callback;
+                sc.scramCredential(credentialCache.get(username));
+            } else
+                throw new UnsupportedCallbackException(callback);
+        }
+    }
+
+    @Override
+    public void close() {
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/security/ssl/DefaultSslEngineFactory.java b/clients/src/main/java/org/apache/kafka/common/security/ssl/DefaultSslEngineFactory.java
new file mode 100644
index 0000000000000..1528a4a8e0b8d
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/ssl/DefaultSslEngineFactory.java
@@ -0,0 +1,584 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.config.types.Password;
+import org.apache.kafka.common.errors.InvalidConfigurationException;
+import org.apache.kafka.common.network.Mode;
+import org.apache.kafka.common.security.auth.SslEngineFactory;
+import org.apache.kafka.common.utils.SecurityUtils;
+import org.apache.kafka.common.utils.Utils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.file.Files;
+import java.nio.file.Paths;
+import java.security.GeneralSecurityException;
+import java.security.Key;
+import java.security.KeyFactory;
+import java.security.KeyStore;
+import java.security.PrivateKey;
+import java.security.SecureRandom;
+import java.security.cert.Certificate;
+import java.security.cert.CertificateFactory;
+import java.security.spec.InvalidKeySpecException;
+import java.security.spec.PKCS8EncodedKeySpec;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Base64;
+import java.util.Collections;
+import java.util.Date;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+import java.util.stream.Collectors;
+
+import javax.crypto.Cipher;
+import javax.crypto.EncryptedPrivateKeyInfo;
+import javax.crypto.SecretKey;
+import javax.crypto.SecretKeyFactory;
+import javax.crypto.spec.PBEKeySpec;
+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;
+
+public final class DefaultSslEngineFactory implements SslEngineFactory {
+
+    private static final Logger log = LoggerFactory.getLogger(DefaultSslEngineFactory.class);
+    public static final String PEM_TYPE = "PEM";
+
+    private Map configs;
+    private String protocol;
+    private String provider;
+    private String kmfAlgorithm;
+    private String tmfAlgorithm;
+    private SecurityStore keystore;
+    private SecurityStore truststore;
+    private String[] cipherSuites;
+    private String[] enabledProtocols;
+    private SecureRandom secureRandomImplementation;
+    private SSLContext sslContext;
+    private SslClientAuth sslClientAuth;
+
+
+    @Override
+    public SSLEngine createClientSslEngine(String peerHost, int peerPort, String endpointIdentification) {
+        return createSslEngine(Mode.CLIENT, peerHost, peerPort, endpointIdentification);
+    }
+
+    @Override
+    public SSLEngine createServerSslEngine(String peerHost, int peerPort) {
+        return createSslEngine(Mode.SERVER, peerHost, peerPort, null);
+    }
+
+    @Override
+    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;
+    }
+
+    @Override
+    public Set reconfigurableConfigs() {
+        return SslConfigs.RECONFIGURABLE_CONFIGS;
+    }
+
+    @Override
+    public KeyStore keystore() {
+        return this.keystore != null ? this.keystore.get() : null;
+    }
+
+    @Override
+    public KeyStore truststore() {
+        return this.truststore != null ? this.truststore.get() : null;
+    }
+
+    @SuppressWarnings("unchecked")
+    @Override
+    public void configure(Map configs) {
+        this.configs = Collections.unmodifiableMap(configs);
+        this.protocol = (String) configs.get(SslConfigs.SSL_PROTOCOL_CONFIG);
+        this.provider = (String) configs.get(SslConfigs.SSL_PROVIDER_CONFIG);
+        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[0]);
+        } 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[0]);
+        } 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),
+                (Password) configs.get(SslConfigs.SSL_KEYSTORE_KEY_CONFIG),
+                (Password) configs.get(SslConfigs.SSL_KEYSTORE_CERTIFICATE_CHAIN_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),
+                (Password) configs.get(SslConfigs.SSL_TRUSTSTORE_CERTIFICATES_CONFIG));
+
+        this.sslContext = createSSLContext(keystore, truststore);
+    }
+
+    @Override
+    public void close() {
+        this.sslContext = null;
+    }
+
+    //For Test only
+    public SSLContext sslContext() {
+        return this.sslContext;
+    }
+
+    private 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;
+    }
+    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(Enum::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(SecurityStore keystore, SecurityStore truststore) {
+        try {
+            SSLContext sslContext;
+            if (provider != null)
+                sslContext = SSLContext.getInstance(protocol, provider);
+            else
+                sslContext = SSLContext.getInstance(protocol);
+
+            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) {
+                    kmf.init(keystore.get(), keystore.keyPassword());
+                } 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.get();
+            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);
+        }
+    }
+
+    // Visibility to override for testing
+    protected SecurityStore createKeystore(String type, String path, Password password, Password keyPassword, Password privateKey, Password certificateChain) {
+        if (privateKey != null) {
+            if (!PEM_TYPE.equals(type))
+                throw new InvalidConfigurationException("SSL private key can be specified only for PEM, but key store type is " + type + ".");
+            else if (certificateChain == null)
+                throw new InvalidConfigurationException("SSL private key is specified, but certificate chain is not specified.");
+            else if (path != null)
+                throw new InvalidConfigurationException("Both SSL key store location and separate private key are specified.");
+            else if (password != null)
+                throw new InvalidConfigurationException("SSL key store password cannot be specified with PEM format, only key password may be specified.");
+            else
+                return new PemStore(certificateChain, privateKey, keyPassword);
+        } else if (certificateChain != null) {
+            throw new InvalidConfigurationException("SSL certificate chain is specified, but private key is not specified");
+        } else if (PEM_TYPE.equals(type) && path != null) {
+            if (password != null)
+                throw new InvalidConfigurationException("SSL key store password cannot be specified with PEM format, only key password may be specified");
+            else if (keyPassword == null)
+                throw new InvalidConfigurationException("SSL PEM key store is specified, but key password is not specified.");
+            else
+                return new FileBasedPemStore(path, keyPassword, true);
+        } else if (path == null && password != null) {
+            throw new InvalidConfigurationException("SSL key store is not specified, but key store password is specified.");
+        } else if (path != null && password == null) {
+            throw new InvalidConfigurationException("SSL key store is specified, but key store password is not specified.");
+        } else if (path != null && password != null) {
+            return new FileBasedStore(type, path, password, keyPassword, true);
+        } 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, Password trustStoreCerts) {
+        if (trustStoreCerts != null) {
+            if (!PEM_TYPE.equals(type))
+                throw new InvalidConfigurationException("SSL trust store certs can be specified only for PEM, but trust store type is " + type + ".");
+            else if (path != null)
+                throw new InvalidConfigurationException("Both SSL trust store location and separate trust certificates are specified.");
+            else if (password != null)
+                throw new InvalidConfigurationException("SSL trust store password cannot be specified for PEM format.");
+            else
+                return new PemStore(trustStoreCerts);
+        } else if (PEM_TYPE.equals(type) && path != null) {
+            if (password != null)
+                throw new InvalidConfigurationException("SSL trust store password cannot be specified for PEM format.");
+            else
+                return new FileBasedPemStore(path, null, false);
+        } else if (path == null && password != null) {
+            throw new InvalidConfigurationException("SSL trust store is not specified, but trust store password is specified.");
+        } else if (path != null) {
+            return new FileBasedStore(type, path, password, null, false);
+        } else
+            return null;
+    }
+
+    static interface SecurityStore {
+        KeyStore get();
+        char[] keyPassword();
+        boolean modified();
+    }
+
+    // package access for testing
+    static class FileBasedStore implements SecurityStore {
+        private final String type;
+        protected final String path;
+        private final Password password;
+        protected final Password keyPassword;
+        private final Long fileLastModifiedMs;
+        private final KeyStore keyStore;
+
+        FileBasedStore(String type, String path, Password password, Password keyPassword, boolean isKeyStore) {
+            Objects.requireNonNull(type, "type must not be null");
+            this.type = type;
+            this.path = path;
+            this.password = password;
+            this.keyPassword = keyPassword;
+            fileLastModifiedMs = lastModifiedMs(path);
+            this.keyStore = load(isKeyStore);
+        }
+
+        @Override
+        public KeyStore get() {
+            return keyStore;
+        }
+
+        @Override
+        public char[] keyPassword() {
+            Password passwd = keyPassword != null ? keyPassword : password;
+            return passwd == null ? null : passwd.value().toCharArray();
+        }
+
+        /**
+         * 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)
+         */
+        protected KeyStore load(boolean isKeyStore) {
+            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;
+            }
+        }
+
+        public 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)) + ")";
+        }
+    }
+
+    static class FileBasedPemStore extends FileBasedStore {
+        FileBasedPemStore(String path, Password keyPassword, boolean isKeyStore) {
+            super(PEM_TYPE, path, null, keyPassword, isKeyStore);
+        }
+
+        @Override
+        protected KeyStore load(boolean isKeyStore) {
+            try {
+                Password storeContents = new Password(Utils.readFileAsString(path));
+                PemStore pemStore = isKeyStore ? new PemStore(storeContents, storeContents, keyPassword) :
+                    new PemStore(storeContents);
+                return pemStore.keyStore;
+            } catch (Exception e) {
+                throw new InvalidConfigurationException("Failed to load PEM SSL keystore " + path, e);
+            }
+        }
+    }
+
+    static class PemStore implements SecurityStore {
+        private static final PemParser CERTIFICATE_PARSER = new PemParser("CERTIFICATE");
+        private static final PemParser PRIVATE_KEY_PARSER = new PemParser("PRIVATE KEY");
+        private static final List KEY_FACTORIES = Arrays.asList(
+                keyFactory("RSA"),
+                keyFactory("DSA"),
+                keyFactory("EC")
+        );
+
+        private final char[] keyPassword;
+        private final KeyStore keyStore;
+
+        PemStore(Password certificateChain, Password privateKey, Password keyPassword) {
+            this.keyPassword = keyPassword == null ? null : keyPassword.value().toCharArray();
+            keyStore = createKeyStoreFromPem(privateKey.value(), certificateChain.value(), this.keyPassword);
+        }
+
+        PemStore(Password trustStoreCerts) {
+            this.keyPassword = null;
+            keyStore = createTrustStoreFromPem(trustStoreCerts.value());
+        }
+
+        @Override
+        public KeyStore get() {
+            return keyStore;
+        }
+
+        @Override
+        public char[] keyPassword() {
+            return keyPassword;
+        }
+
+        @Override
+        public boolean modified() {
+            return false;
+        }
+
+        private KeyStore createKeyStoreFromPem(String privateKeyPem, String certChainPem, char[] keyPassword) {
+            try {
+                KeyStore ks = KeyStore.getInstance("PKCS12");
+                ks.load(null, null);
+                Key key = privateKey(privateKeyPem, keyPassword);
+                Certificate[] certChain = certs(certChainPem);
+                ks.setKeyEntry("kafka", key, keyPassword, certChain);
+                return ks;
+            } catch (Exception e) {
+                throw new InvalidConfigurationException("Invalid PEM keystore configs", e);
+            }
+        }
+
+        private KeyStore createTrustStoreFromPem(String trustedCertsPem) {
+            try {
+                KeyStore ts = KeyStore.getInstance("PKCS12");
+                ts.load(null, null);
+                Certificate[] certs = certs(trustedCertsPem);
+                for (int i = 0; i < certs.length; i++) {
+                    ts.setCertificateEntry("kafka" + i, certs[i]);
+                }
+                return ts;
+            } catch (InvalidConfigurationException e) {
+                throw e;
+            } catch (Exception e) {
+                throw new InvalidConfigurationException("Invalid PEM keystore configs", e);
+            }
+        }
+
+        private Certificate[] certs(String pem) throws GeneralSecurityException {
+            List certEntries = CERTIFICATE_PARSER.pemEntries(pem);
+            if (certEntries.isEmpty())
+                throw new InvalidConfigurationException("At least one certificate expected, but none found");
+
+            Certificate[] certs = new Certificate[certEntries.size()];
+            for (int i = 0; i < certs.length; i++) {
+                certs[i] = CertificateFactory.getInstance("X.509")
+                    .generateCertificate(new ByteArrayInputStream(certEntries.get(i)));
+            }
+            return certs;
+        }
+
+        private PrivateKey privateKey(String pem, char[] keyPassword) throws Exception {
+            List keyEntries = PRIVATE_KEY_PARSER.pemEntries(pem);
+            if (keyEntries.isEmpty())
+                throw new InvalidConfigurationException("Private key not provided");
+            if (keyEntries.size() != 1)
+                throw new InvalidConfigurationException("Expected one private key, but found " + keyEntries.size());
+
+            byte[] keyBytes = keyEntries.get(0);
+            PKCS8EncodedKeySpec keySpec;
+            if (keyPassword == null) {
+                keySpec = new PKCS8EncodedKeySpec(keyBytes);
+            } else {
+                EncryptedPrivateKeyInfo keyInfo = new EncryptedPrivateKeyInfo(keyBytes);
+                String algorithm = keyInfo.getAlgName();
+                SecretKeyFactory keyFactory = SecretKeyFactory.getInstance(algorithm);
+                SecretKey pbeKey = keyFactory.generateSecret(new PBEKeySpec(keyPassword));
+                Cipher cipher = Cipher.getInstance(algorithm);
+                cipher.init(Cipher.DECRYPT_MODE, pbeKey, keyInfo.getAlgParameters());
+                keySpec = keyInfo.getKeySpec(cipher);
+            }
+
+            InvalidKeySpecException firstException = null;
+            for (KeyFactory factory : KEY_FACTORIES) {
+                try {
+                    return factory.generatePrivate(keySpec);
+                } catch (InvalidKeySpecException e) {
+                    if (firstException == null)
+                        firstException = e;
+                }
+            }
+            throw new InvalidConfigurationException("Private key could not be loaded", firstException);
+        }
+
+        private static KeyFactory keyFactory(String algorithm) {
+            try {
+                return KeyFactory.getInstance(algorithm);
+            } catch (Exception e) {
+                throw new InvalidConfigurationException("Could not create key factory for algorithm " + algorithm, e);
+            }
+        }
+    }
+
+    /**
+     * Parser to process certificate/private key entries from PEM files
+     * Examples:
+     *   -----BEGIN CERTIFICATE-----
+     *   Base64 cert
+     *   -----END CERTIFICATE-----
+     *
+     *   -----BEGIN ENCRYPTED PRIVATE KEY-----
+     *   Base64 private key
+     *   -----END ENCRYPTED PRIVATE KEY-----
+     *   Additional data may be included before headers, so we match all entries within the PEM.
+     */
+    static class PemParser {
+        private final String name;
+        private final Pattern pattern;
+
+        PemParser(String name) {
+            this.name = name;
+            String beginOrEndFormat = "-+%s\\s*.*%s[^-]*-+\\s+";
+            String nameIgnoreSpace = name.replace(" ", "\\s+");
+
+            String encodingParams = "\\s*[^\\r\\n]*:[^\\r\\n]*[\\r\\n]+";
+            String base64Pattern = "([a-zA-Z0-9/+=\\s]*)";
+            String patternStr =  String.format(beginOrEndFormat, "BEGIN", nameIgnoreSpace) +
+                String.format("(?:%s)*", encodingParams) +
+                base64Pattern +
+                String.format(beginOrEndFormat, "END", nameIgnoreSpace);
+            pattern = Pattern.compile(patternStr);
+        }
+
+        private List pemEntries(String pem) {
+            Matcher matcher = pattern.matcher(pem + "\n"); // allow last newline to be omitted in value
+            List  entries = new ArrayList<>();
+            while (matcher.find()) {
+                String base64Str = matcher.group(1).replaceAll("\\s", "");
+                entries.add(Base64.getDecoder().decode(base64Str));
+            }
+            if (entries.isEmpty())
+                throw new InvalidConfigurationException("No matching " + name + " entries in PEM file");
+            return entries;
+        }
+    }
+}
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 6b2fe74219edc..91b23b61a9397 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
@@ -16,215 +16,419 @@
  */
 package org.apache.kafka.common.security.ssl;
 
-import org.apache.kafka.common.Configurable;
 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.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.security.auth.SslEngineFactory;
+import org.apache.kafka.common.utils.Utils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
-import java.io.FileInputStream;
-import java.io.IOException;
+import javax.net.ssl.SSLEngine;
+import javax.net.ssl.SSLEngineResult;
+import javax.net.ssl.SSLException;
+import java.io.Closeable;
+import java.nio.ByteBuffer;
+import java.security.cert.Certificate;
+import java.security.cert.X509Certificate;
 import java.security.GeneralSecurityException;
 import java.security.KeyStore;
-import java.security.SecureRandom;
+import java.security.Principal;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Enumeration;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Set;
+import java.util.HashSet;
 
-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;
-
-public class SslFactory implements Configurable {
+public class SslFactory implements Reconfigurable, Closeable {
+    private static final Logger log = LoggerFactory.getLogger(SslFactory.class);
 
     private final Mode mode;
     private final String clientAuthConfigOverride;
-
-    private String protocol;
-    private String provider;
-    private String kmfAlgorithm;
-    private String tmfAlgorithm;
-    private SecurityStore keystore = null;
-    private Password keyPassword;
-    private SecurityStore truststore;
-    private String[] cipherSuites;
-    private String[] enabledProtocols;
+    private final boolean keystoreVerifiableUsingTruststore;
     private String endpointIdentification;
-    private SecureRandom secureRandomImplementation;
-    private SSLContext sslContext;
-    private boolean needClientAuth;
-    private boolean wantClientAuth;
+    private SslEngineFactory sslEngineFactory;
+    private Map sslEngineFactoryConfig;
 
     public SslFactory(Mode mode) {
-        this(mode, null);
+        this(mode, null, false);
     }
 
-    public SslFactory(Mode mode, String clientAuthConfigOverride) {
+    /**
+     * 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;
     }
 
+    @SuppressWarnings("unchecked")
     @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)
-            this.cipherSuites = cipherSuitesList.toArray(new String[cipherSuitesList.size()]);
-
-        @SuppressWarnings("unchecked")
-        List enabledProtocolsList = (List) configs.get(SslConfigs.SSL_ENABLED_PROTOCOLS_CONFIG);
-        if (enabledProtocolsList != null)
-            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;
+        if (sslEngineFactory != null) {
+            throw new IllegalStateException("SslFactory was already configured.");
+        }
+        this.endpointIdentification = (String) configs.get(SslConfigs.SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG);
 
-        String secureRandomImplementation = (String) configs.get(SslConfigs.SSL_SECURE_RANDOM_IMPLEMENTATION_CONFIG);
-        if (secureRandomImplementation != null) {
+        // The input map must be a mutable RecordingMap in production.
+        Map nextConfigs = (Map) configs;
+        if (clientAuthConfigOverride != null) {
+            nextConfigs.put(BrokerSecurityConfigs.SSL_CLIENT_AUTH_CONFIG, clientAuthConfigOverride);
+        }
+        SslEngineFactory builder = instantiateSslEngineFactory(nextConfigs);
+        if (keystoreVerifiableUsingTruststore) {
             try {
-                this.secureRandomImplementation = SecureRandom.getInstance(secureRandomImplementation);
-            } catch (GeneralSecurityException e) {
-                throw new KafkaException(e);
+                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.sslEngineFactory = builder;
+    }
 
-        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.kmfAlgorithm = (String) configs.get(SslConfigs.SSL_KEYMANAGER_ALGORITHM_CONFIG);
-        this.tmfAlgorithm = (String) configs.get(SslConfigs.SSL_TRUSTMANAGER_ALGORITHM_CONFIG);
+    @Override
+    public Set reconfigurableConfigs() {
+        return sslEngineFactory.reconfigurableConfigs();
+    }
 
-        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));
+    @Override
+    public void validateReconfiguration(Map newConfigs) {
+        createNewSslEngineFactory(newConfigs);
+    }
 
-        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();
-        } catch (Exception e) {
-            throw new KafkaException(e);
+    @Override
+    public void reconfigure(Map newConfigs) throws KafkaException {
+        SslEngineFactory newSslEngineFactory = createNewSslEngineFactory(newConfigs);
+        if (newSslEngineFactory != this.sslEngineFactory) {
+            Utils.closeQuietly(this.sslEngineFactory, "close stale ssl engine factory");
+            this.sslEngineFactory = newSslEngineFactory;
+            log.info("Created new {} SSL engine builder with keystore {} truststore {}", mode,
+                    newSslEngineFactory.keystore(), newSslEngineFactory.truststore());
         }
     }
 
-
-    private SSLContext createSSLContext() 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 = this.keyPassword != null ? this.keyPassword : keystore.password;
-            kmf.init(ks, keyPassword.value().toCharArray());
-            keyManagers = kmf.getKeyManagers();
+    private SslEngineFactory instantiateSslEngineFactory(Map configs) {
+        @SuppressWarnings("unchecked")
+        Class sslEngineFactoryClass =
+                (Class) configs.get(SslConfigs.SSL_ENGINE_FACTORY_CLASS_CONFIG);
+        SslEngineFactory sslEngineFactory;
+        if (sslEngineFactoryClass == null) {
+            sslEngineFactory = new DefaultSslEngineFactory();
+        } else {
+            sslEngineFactory = Utils.newInstance(sslEngineFactoryClass);
         }
+        sslEngineFactory.configure(configs);
+        this.sslEngineFactoryConfig = configs;
+        return sslEngineFactory;
+    }
 
-        String tmfAlgorithm = this.tmfAlgorithm != null ? this.tmfAlgorithm : TrustManagerFactory.getDefaultAlgorithm();
-        TrustManagerFactory tmf = TrustManagerFactory.getInstance(tmfAlgorithm);
-        KeyStore ts = truststore == null ? null : truststore.load();
-        tmf.init(ts);
+    private SslEngineFactory createNewSslEngineFactory(Map newConfigs) {
+        if (sslEngineFactory == null) {
+            throw new IllegalStateException("SslFactory has not been configured.");
+        }
+        Map nextConfigs = new HashMap<>(sslEngineFactoryConfig);
+        copyMapEntries(nextConfigs, newConfigs, reconfigurableConfigs());
+        if (clientAuthConfigOverride != null) {
+            nextConfigs.put(BrokerSecurityConfigs.SSL_CLIENT_AUTH_CONFIG, clientAuthConfigOverride);
+        }
+        if (!sslEngineFactory.shouldBeRebuilt(nextConfigs)) {
+            return sslEngineFactory;
+        }
+        try {
+            SslEngineFactory newSslEngineFactory = instantiateSslEngineFactory(nextConfigs);
+            if (sslEngineFactory.keystore() == null) {
+                if (newSslEngineFactory.keystore() != null) {
+                    throw new ConfigException("Cannot add SSL keystore to an existing listener for " +
+                            "which no keystore was configured.");
+                }
+            } else {
+                if (newSslEngineFactory.keystore() == null) {
+                    throw new ConfigException("Cannot remove the SSL keystore from an existing listener for " +
+                            "which a keystore was configured.");
+                }
 
-        sslContext.init(keyManagers, tmf.getTrustManagers(), this.secureRandomImplementation);
-        return sslContext;
+                CertificateEntries.ensureCompatible(newSslEngineFactory.keystore(), sslEngineFactory.keystore());
+            }
+            if (sslEngineFactory.truststore() == null && newSslEngineFactory.truststore() != null) {
+                throw new ConfigException("Cannot add SSL truststore to an existing listener for which no " +
+                        "truststore was configured.");
+            }
+            if (keystoreVerifiableUsingTruststore) {
+                if (sslEngineFactory.truststore() != null || sslEngineFactory.keystore() != null) {
+                    SslEngineValidator.validate(sslEngineFactory, newSslEngineFactory);
+                }
+            }
+            return newSslEngineFactory;
+        } 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);
+        }
     }
 
     public SSLEngine createSslEngine(String peerHost, int peerPort) {
-        SSLEngine sslEngine = sslContext.createSSLEngine(peerHost, peerPort);
-        if (cipherSuites != null) sslEngine.setEnabledCipherSuites(cipherSuites);
-        if (enabledProtocols != null) sslEngine.setEnabledProtocols(enabledProtocols);
-
-        // SSLParameters#setEndpointIdentificationAlgorithm enables endpoint validation
-        // only in client mode. Hence, validation is enabled only for clients.
+        if (sslEngineFactory == null) {
+            throw new IllegalStateException("SslFactory has not been configured.");
+        }
         if (mode == Mode.SERVER) {
-            sslEngine.setUseClientMode(false);
-            if (needClientAuth)
-                sslEngine.setNeedClientAuth(needClientAuth);
-            else
-                sslEngine.setWantClientAuth(wantClientAuth);
+            return sslEngineFactory.createServerSslEngine(peerHost, peerPort);
         } else {
-            sslEngine.setUseClientMode(true);
-            SSLParameters sslParams = sslEngine.getSSLParameters();
-            sslParams.setEndpointIdentificationAlgorithm(endpointIdentification);
-            sslEngine.setSSLParameters(sslParams);
+            return sslEngineFactory.createClientSslEngine(peerHost, peerPort, endpointIdentification);
         }
-        return sslEngine;
+    }
+
+    public SslEngineFactory sslEngineFactory() {
+        return sslEngineFactory;
     }
 
     /**
-     * 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 static  void copyMapEntries(Map destMap,
+                                              Map srcMap,
+                                              Set keySet) {
+        for (K k : keySet) {
+            copyMapEntry(destMap, srcMap, k);
+        }
     }
 
-    private void 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) {
-            this.keystore = new SecurityStore(type, path, password);
-            this.keyPassword = keyPassword;
+    /**
+     * 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));
         }
     }
 
-    private void 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) {
-            this.truststore = new SecurityStore(type, path, password);
+    @Override
+    public void close() {
+        Utils.closeQuietly(sslEngineFactory, "close engine factory");
+    }
+
+    static class CertificateEntries {
+        private final String alias;
+        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(alias, (X509Certificate) cert));
+            }
+            return entries;
+        }
+
+        static void ensureCompatible(KeyStore newKeystore, KeyStore oldKeystore) throws GeneralSecurityException {
+            List newEntries = CertificateEntries.create(newKeystore);
+            List oldEntries = CertificateEntries.create(oldKeystore);
+            if (newEntries.size() != oldEntries.size()) {
+                throw new ConfigException(String.format("Keystore entries do not match, existing store contains %d entries, new store contains %d entries",
+                    oldEntries.size(), newEntries.size()));
+            }
+            for (int i = 0; i < newEntries.size(); i++) {
+                CertificateEntries newEntry = newEntries.get(i);
+                CertificateEntries oldEntry = oldEntries.get(i);
+                if (!Objects.equals(newEntry.subjectPrincipal, oldEntry.subjectPrincipal)) {
+                    throw new ConfigException(String.format("Keystore DistinguishedName does not match: " +
+                        " existing={alias=%s, DN=%s}, new={alias=%s, DN=%s}",
+                        oldEntry.alias, oldEntry.subjectPrincipal, newEntry.alias, newEntry.subjectPrincipal));
+                }
+                if (!newEntry.subjectAltNames.containsAll(oldEntry.subjectAltNames)) {
+                    throw new ConfigException(String.format("Keystore SubjectAltNames do not match: " +
+                            " existing={alias=%s, SAN=%s}, new={alias=%s, SAN=%s}",
+                        oldEntry.alias, oldEntry.subjectAltNames, newEntry.alias, newEntry.subjectAltNames));
+                }
+            }
+        }
+
+        CertificateEntries(String alias, X509Certificate cert) throws GeneralSecurityException {
+            this.alias = alias;
+            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;
         }
     }
 
-    private static class SecurityStore {
-        private final String type;
-        private final String path;
-        private final Password password;
+    /**
+     * Validator used to verify dynamic update of keystore used in inter-broker communication.
+     * The validator checks that a successful handshake can be performed using the keystore and
+     * truststore configured on this SslFactory.
+     */
+    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(SslEngineFactory oldEngineBuilder,
+                             SslEngineFactory newEngineBuilder) throws SSLException {
+            validate(createSslEngineForValidation(oldEngineBuilder, Mode.SERVER),
+                    createSslEngineForValidation(newEngineBuilder, Mode.CLIENT));
+            validate(createSslEngineForValidation(newEngineBuilder, Mode.SERVER),
+                    createSslEngineForValidation(oldEngineBuilder, Mode.CLIENT));
+        }
 
-        private SecurityStore(String type, String path, Password password) {
-            Objects.requireNonNull(type, "type must not be null");
-            this.type = type;
-            this.path = path;
-            this.password = password;
+        private static SSLEngine createSslEngineForValidation(SslEngineFactory sslEngineFactory, Mode mode) {
+            // Use empty hostname, disable hostname verification
+            if (mode == Mode.SERVER) {
+                return sslEngineFactory.createServerSslEngine("", 0);
+            } else {
+                return sslEngineFactory.createClientSslEngine("", 0, "");
+            }
         }
 
-        private KeyStore load() throws GeneralSecurityException, IOException {
-            FileInputStream in = null;
+        static void validate(SSLEngine clientEngine, SSLEngine serverEngine) throws SSLException {
+            SslEngineValidator clientValidator = new SslEngineValidator(clientEngine);
+            SslEngineValidator serverValidator = new SslEngineValidator(serverEngine);
             try {
-                KeyStore ks = KeyStore.getInstance(type);
-                in = new FileInputStream(path);
-                // 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;
+                clientValidator.beginHandshake();
+                serverValidator.beginHandshake();
+                while (!serverValidator.complete() || !clientValidator.complete()) {
+                    clientValidator.handshake(serverValidator);
+                    serverValidator.handshake(clientValidator);
+                }
             } finally {
-                if (in != null) in.close();
+                clientValidator.close();
+                serverValidator.close();
             }
         }
-    }
 
-}
+        private SslEngineValidator(SSLEngine engine) {
+            this.sslEngine = engine;
+            appBuffer = ByteBuffer.allocate(sslEngine.getSession().getApplicationBufferSize());
+            netBuffer = ByteBuffer.allocate(sslEngine.getSession().getPacketBufferSize());
+        }
+
+        void beginHandshake() throws SSLException {
+            sslEngine.beginHandshake();
+        }
+        void handshake(SslEngineValidator peerValidator) throws SSLException {
+            SSLEngineResult.HandshakeStatus handshakeStatus = sslEngine.getHandshakeStatus();
+            while (true) {
+                switch (handshakeStatus) {
+                    case NEED_WRAP:
+                        if (netBuffer.position() != 0) // Wait for peer to consume previously wrapped data
+                            return;
+                        handshakeResult = sslEngine.wrap(EMPTY_BUF, netBuffer);
+                        switch (handshakeResult.getStatus()) {
+                            case OK: break;
+                            case BUFFER_OVERFLOW:
+                                netBuffer.compact();
+                                netBuffer = Utils.ensureCapacity(netBuffer, sslEngine.getSession().getPacketBufferSize());
+                                netBuffer.flip();
+                                break;
+                            case BUFFER_UNDERFLOW:
+                            case CLOSED:
+                            default:
+                                throw new SSLException("Unexpected handshake status: " + handshakeResult.getStatus());
+                        }
+                        return;
+                    case NEED_UNWRAP:
+                        if (peerValidator.netBuffer.position() == 0) // no data to unwrap, return to process peer
+                            return;
+                        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;
+                            case BUFFER_OVERFLOW:
+                                appBuffer = Utils.ensureCapacity(appBuffer, sslEngine.getSession().getApplicationBufferSize());
+                                break;
+                            case BUFFER_UNDERFLOW:
+                                netBuffer = Utils.ensureCapacity(netBuffer, sslEngine.getSession().getPacketBufferSize());
+                                break;
+                            case CLOSED:
+                            default:
+                                throw new SSLException("Unexpected handshake status: " + handshakeResult.getStatus());
+                        }
+                        break;
+                    case NEED_TASK:
+                        sslEngine.getDelegatedTask().run();
+                        handshakeStatus = sslEngine.getHandshakeStatus();
+                        break;
+                    case FINISHED:
+                        return;
+                    case NOT_HANDSHAKING:
+                        if (handshakeResult.getHandshakeStatus() != SSLEngineResult.HandshakeStatus.FINISHED)
+                            throw new SSLException("Did not finish handshake");
+                        return;
+                    default:
+                        throw new IllegalStateException("Unexpected handshake status " + handshakeStatus);
+                }
+            }
+        }
+
+        boolean complete() {
+            return sslEngine.getHandshakeStatus() == SSLEngineResult.HandshakeStatus.FINISHED ||
+                    sslEngine.getHandshakeStatus() == SSLEngineResult.HandshakeStatus.NOT_HANDSHAKING;
+        }
+
+        void close() {
+            sslEngine.closeOutbound();
+            try {
+                sslEngine.closeInbound();
+            } catch (Exception e) {
+                // ignore
+            }
+        }
+    }
+}
\ No newline at end of file
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
new file mode 100644
index 0000000000000..33da9642d97b2
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/ssl/SslPrincipalMapper.java
@@ -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 org.apache.kafka.common.security.ssl;
+
+import java.io.IOException;
+import java.util.List;
+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 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(String sslPrincipalMappingRules) {
+        this.rules = parseRules(splitRules(sslPrincipalMappingRules));
+    }
+
+    public static SslPrincipalMapper fromRules(String sslPrincipalMappingRules) {
+        return new SslPrincipalMapper(sslPrincipalMappingRules);
+    }
+
+    private static List splitRules(String sslPrincipalMappingRules) {
+        if (sslPrincipalMappingRules == null) {
+            sslPrincipalMappingRules = DEFAULT_SSL_PRINCIPAL_MAPPING_RULES;
+        }
+
+        List result = new ArrayList<>();
+        Matcher matcher = RULE_SPLITTER.matcher(sslPrincipalMappingRules.trim());
+        while (matcher.find()) {
+            result.add(matcher.group(1));
+        }
+
+        return result;
+    }
+
+    private static List parseRules(List rules) {
+        List result = new ArrayList<>();
+        for (String rule : rules) {
+            Matcher matcher = RULE_PARSER.matcher(rule);
+            if (!matcher.lookingAt()) {
+                throw new IllegalArgumentException("Invalid rule: " + rule);
+            }
+            if (rule.length() != matcher.end()) {
+                throw new IllegalArgumentException("Invalid rule: `" + rule + "`, unmatched substring: `" + rule.substring(matcher.end()) + "`");
+            }
+
+            // empty rules are ignored
+            if (matcher.group(1) != null) {
+                result.add(new Rule());
+            } 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;
+    }
+
+    public String getName(String distinguishedName) throws IOException {
+        for (Rule r : rules) {
+            String principalName = r.apply(distinguishedName);
+            if (principalName != null) {
+                return principalName;
+            }
+        }
+        throw new NoMatchingRule("No rules apply to " + distinguishedName + ", rules " + rules);
+    }
+
+    @Override
+    public String toString() {
+        return "SslPrincipalMapper(rules = " + rules + ")";
+    }
+
+    public static class NoMatchingRule extends IOException {
+        NoMatchingRule(String msg) {
+            super(msg);
+        }
+    }
+
+    private static class Rule {
+        private static final Pattern BACK_REFERENCE_PATTERN = Pattern.compile("\\$(\\d+)");
+
+        private final boolean isDefault;
+        private final Pattern pattern;
+        private final String replacement;
+        private final boolean toLowerCase;
+        private final boolean toUpperCase;
+
+        Rule() {
+            isDefault = true;
+            pattern = null;
+            replacement = null;
+            toLowerCase = false;
+            toUpperCase = false;
+        }
+
+        Rule(String pattern, String replacement, boolean toLowerCase, boolean toUpperCase) {
+            isDefault = false;
+            this.pattern = pattern == null ? null : Pattern.compile(pattern);
+            this.replacement = replacement;
+            this.toLowerCase = toLowerCase;
+            this.toUpperCase = toUpperCase;
+        }
+
+        String apply(String distinguishedName) {
+            if (isDefault) {
+                return distinguishedName;
+            }
+
+            String result = null;
+            final Matcher m = pattern.matcher(distinguishedName);
+
+            if (m.matches()) {
+                result = distinguishedName.replaceAll(pattern.pattern(), escapeLiteralBackReferences(replacement, m.groupCount()));
+            }
+
+            if (toLowerCase && result != null) {
+                result = result.toLowerCase(Locale.ENGLISH);
+            } else if (toUpperCase & result != null) {
+                result = result.toUpperCase(Locale.ENGLISH);
+            }
+
+            return result;
+        }
+
+        //If we find a back reference that is not valid, then we will treat it as a literal string. For example, if we have 3 capturing
+        //groups and the Replacement Value has the value is "$1@$4", then we want to treat the $4 as a literal "$4", rather
+        //than attempting to use it as a back reference.
+        //This method was taken from Apache Nifi project : org.apache.nifi.authorization.util.IdentityMappingUtil
+        private String escapeLiteralBackReferences(final String unescaped, final int numCapturingGroups) {
+            if (numCapturingGroups == 0) {
+                return unescaped;
+            }
+
+            String value = unescaped;
+            final Matcher backRefMatcher = BACK_REFERENCE_PATTERN.matcher(value);
+            while (backRefMatcher.find()) {
+                final String backRefNum = backRefMatcher.group(1);
+                if (backRefNum.startsWith("0")) {
+                    continue;
+                }
+                int backRefIndex = Integer.parseInt(backRefNum);
+
+
+                // if we have a replacement value like $123, and we have less than 123 capturing groups, then
+                // we want to truncate the 3 and use capturing group 12; if we have less than 12 capturing groups,
+                // then we want to truncate the 2 and use capturing group 1; if we don't have a capturing group then
+                // we want to truncate the 1 and get 0.
+                while (backRefIndex > numCapturingGroups && backRefIndex >= 10) {
+                    backRefIndex /= 10;
+                }
+
+                if (backRefIndex > numCapturingGroups) {
+                    final StringBuilder sb = new StringBuilder(value.length() + 1);
+                    final int groupStart = backRefMatcher.start(1);
+
+                    sb.append(value.substring(0, groupStart - 1));
+                    sb.append("\\");
+                    sb.append(value.substring(groupStart - 1));
+                    value = sb.toString();
+                }
+            }
+
+            return value;
+        }
+
+        @Override
+        public String toString() {
+            StringBuilder buf = new StringBuilder();
+            if (isDefault) {
+                buf.append("DEFAULT");
+            } else {
+                buf.append("RULE:");
+                if (pattern != null) {
+                    buf.append(pattern);
+                }
+                if (replacement != null) {
+                    buf.append("/");
+                    buf.append(replacement);
+                }
+                if (toLowerCase) {
+                    buf.append("/L");
+                } else if (toUpperCase) {
+                    buf.append("/U");
+                }
+            }
+            return buf.toString();
+        }
+
+    }
+}
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
new file mode 100644
index 0000000000000..b389a199a9201
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/token/delegation/DelegationToken.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.security.token.delegation;
+
+import org.apache.kafka.common.annotation.InterfaceStability;
+
+import java.util.Arrays;
+import java.util.Base64;
+import java.util.Objects;
+
+/**
+ * A class representing a delegation token.
+ *
+ */
+@InterfaceStability.Evolving
+public class DelegationToken {
+    private TokenInformation tokenInformation;
+    private byte[] hmac;
+
+    public DelegationToken(TokenInformation tokenInformation, byte[] hmac) {
+        this.tokenInformation = tokenInformation;
+        this.hmac = hmac;
+    }
+
+    public TokenInformation tokenInfo() {
+        return tokenInformation;
+    }
+
+    public byte[] hmac() {
+        return hmac;
+    }
+
+    public String hmacAsBase64String() {
+        return Base64.getEncoder().encodeToString(hmac);
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+
+        DelegationToken token = (DelegationToken) o;
+
+        return Objects.equals(tokenInformation, token.tokenInformation) && Arrays.equals(hmac, token.hmac);
+    }
+
+    @Override
+    public int hashCode() {
+        int result = tokenInformation != null ? tokenInformation.hashCode() : 0;
+        result = 31 * result + Arrays.hashCode(hmac);
+        return result;
+    }
+
+    @Override
+    public String toString() {
+        return "DelegationToken{" +
+            "tokenInformation=" + tokenInformation +
+            ", hmac=[*******]" +
+            '}';
+    }
+}
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
new file mode 100644
index 0000000000000..9903eb51b235f
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/token/delegation/TokenInformation.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.security.token.delegation;
+
+import org.apache.kafka.common.annotation.InterfaceStability;
+import org.apache.kafka.common.security.auth.KafkaPrincipal;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Objects;
+
+/**
+ * A class representing a delegation token details.
+ *
+ */
+@InterfaceStability.Evolving
+public class TokenInformation {
+
+    private KafkaPrincipal owner;
+    private Collection renewers;
+    private long issueTimestamp;
+    private long maxTimestamp;
+    private long expiryTimestamp;
+    private String tokenId;
+
+    public TokenInformation(String tokenId, KafkaPrincipal owner, Collection renewers,
+                            long issueTimestamp, long maxTimestamp, long expiryTimestamp) {
+        this.tokenId = tokenId;
+        this.owner = owner;
+        this.renewers = renewers;
+        this.issueTimestamp =  issueTimestamp;
+        this.maxTimestamp =  maxTimestamp;
+        this.expiryTimestamp =  expiryTimestamp;
+    }
+
+    public KafkaPrincipal owner() {
+        return owner;
+    }
+
+    public String ownerAsString() {
+        return owner.toString();
+    }
+
+    public Collection renewers() {
+        return renewers;
+    }
+
+    public Collection renewersAsString() {
+        Collection renewerList = new ArrayList<>();
+        for (KafkaPrincipal renewer : renewers) {
+            renewerList.add(renewer.toString());
+        }
+        return renewerList;
+    }
+
+    public long issueTimestamp() {
+        return issueTimestamp;
+    }
+
+    public long expiryTimestamp() {
+        return expiryTimestamp;
+    }
+
+    public void setExpiryTimestamp(long expiryTimestamp) {
+        this.expiryTimestamp = expiryTimestamp;
+    }
+
+    public String tokenId() {
+        return tokenId;
+    }
+
+    public long maxTimestamp() {
+        return maxTimestamp;
+    }
+
+    public boolean ownerOrRenewer(KafkaPrincipal principal) {
+        return owner.equals(principal) || renewers.contains(principal);
+    }
+
+    @Override
+    public String toString() {
+        return "TokenInformation{" +
+            "owner=" + owner +
+            ", renewers=" + renewers +
+            ", issueTimestamp=" + issueTimestamp +
+            ", maxTimestamp=" + maxTimestamp +
+            ", expiryTimestamp=" + expiryTimestamp +
+            ", tokenId='" + tokenId + '\'' +
+            '}';
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+
+        TokenInformation that = (TokenInformation) o;
+
+        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 + 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/security/token/delegation/internals/DelegationTokenCache.java b/clients/src/main/java/org/apache/kafka/common/security/token/delegation/internals/DelegationTokenCache.java
new file mode 100644
index 0000000000000..9cc913f57502e
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/token/delegation/internals/DelegationTokenCache.java
@@ -0,0 +1,125 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.token.delegation.internals;
+
+import org.apache.kafka.common.security.authenticator.CredentialCache;
+import org.apache.kafka.common.security.scram.ScramCredential;
+import org.apache.kafka.common.security.scram.internals.ScramCredentialUtils;
+import org.apache.kafka.common.security.scram.internals.ScramMechanism;
+import org.apache.kafka.common.security.token.delegation.DelegationToken;
+import org.apache.kafka.common.security.token.delegation.TokenInformation;
+
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+public class DelegationTokenCache {
+
+    private CredentialCache credentialCache = new CredentialCache();
+
+    //Cache to hold all the tokens
+    private Map tokenCache = new ConcurrentHashMap<>();
+
+    //Cache to hold hmac->tokenId mapping. This is required for renew, expire requests
+    private Map hmacTokenIdCache = new ConcurrentHashMap<>();
+
+    //Cache to hold tokenId->hmac mapping. This is required for removing entry from hmacTokenIdCache using tokenId.
+    private Map tokenIdHmacCache = new ConcurrentHashMap<>();
+
+    public DelegationTokenCache(Collection scramMechanisms) {
+        //Create caches for scramMechanisms
+        ScramCredentialUtils.createCache(credentialCache, scramMechanisms);
+    }
+
+    public ScramCredential credential(String mechanism, String tokenId) {
+        CredentialCache.Cache cache = credentialCache.cache(mechanism, ScramCredential.class);
+        return cache == null ? null : cache.get(tokenId);
+    }
+
+    public String owner(String tokenId) {
+        TokenInformation tokenInfo = tokenCache.get(tokenId);
+        return tokenInfo == null ? null : tokenInfo.owner().getName();
+    }
+
+    public void updateCache(DelegationToken token, Map scramCredentialMap) {
+        //Update TokenCache
+        String tokenId =  token.tokenInfo().tokenId();
+        addToken(tokenId, token.tokenInfo());
+        String hmac = token.hmacAsBase64String();
+        //Update Scram Credentials
+        updateCredentials(tokenId, scramCredentialMap);
+        //Update hmac-id cache
+        hmacTokenIdCache.put(hmac, tokenId);
+        tokenIdHmacCache.put(tokenId, hmac);
+    }
+
+    public void removeCache(String tokenId) {
+        removeToken(tokenId);
+        updateCredentials(tokenId, new HashMap<>());
+    }
+
+    public String tokenIdForHmac(String base64hmac) {
+        return hmacTokenIdCache.get(base64hmac);
+    }
+
+    public TokenInformation tokenForHmac(String base64hmac) {
+        String tokenId = hmacTokenIdCache.get(base64hmac);
+        return tokenId == null ? null : tokenCache.get(tokenId);
+    }
+
+    public TokenInformation addToken(String tokenId, TokenInformation tokenInfo) {
+        return tokenCache.put(tokenId, tokenInfo);
+    }
+
+    public void removeToken(String tokenId) {
+        TokenInformation tokenInfo = tokenCache.remove(tokenId);
+        if (tokenInfo != null) {
+            String hmac = tokenIdHmacCache.remove(tokenInfo.tokenId());
+            if (hmac != null) {
+                hmacTokenIdCache.remove(hmac);
+            }
+        }
+    }
+
+    public Collection tokens() {
+        return tokenCache.values();
+    }
+
+    public TokenInformation token(String tokenId) {
+        return tokenCache.get(tokenId);
+    }
+
+    public CredentialCache.Cache credentialCache(String mechanism) {
+        return credentialCache.cache(mechanism, ScramCredential.class);
+    }
+
+    private void updateCredentials(String tokenId, Map scramCredentialMap) {
+        for (String mechanism : ScramMechanism.mechanismNames()) {
+            CredentialCache.Cache cache = credentialCache.cache(mechanism, ScramCredential.class);
+            if (cache != null) {
+                ScramCredential credential = scramCredentialMap.get(mechanism);
+                if (credential == null) {
+                    cache.remove(tokenId);
+                } else {
+                    cache.put(tokenId, credential);
+                }
+            }
+        }
+    }
+}
\ No newline at end of file
diff --git a/clients/src/main/java/org/apache/kafka/common/security/token/delegation/internals/DelegationTokenCredentialCallback.java b/clients/src/main/java/org/apache/kafka/common/security/token/delegation/internals/DelegationTokenCredentialCallback.java
new file mode 100644
index 0000000000000..5d9eee986a1bd
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/security/token/delegation/internals/DelegationTokenCredentialCallback.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.common.security.token.delegation.internals;
+
+import org.apache.kafka.common.security.scram.ScramCredentialCallback;
+
+public class DelegationTokenCredentialCallback extends ScramCredentialCallback {
+    private String tokenOwner;
+    private Long tokenExpiryTimestamp;
+
+    public void tokenOwner(String tokenOwner) {
+        this.tokenOwner = tokenOwner;
+    }
+
+    public String tokenOwner() {
+        return tokenOwner;
+    }
+    
+    public void tokenExpiryTimestamp(Long tokenExpiryTimestamp) {
+        this.tokenExpiryTimestamp = tokenExpiryTimestamp;
+    }
+
+    public Long tokenExpiryTimestamp() {
+        return tokenExpiryTimestamp;
+    }
+}
\ No newline at end of file
diff --git a/clients/src/main/java/org/apache/kafka/common/serialization/ByteArrayDeserializer.java b/clients/src/main/java/org/apache/kafka/common/serialization/ByteArrayDeserializer.java
index 267211576b6d0..1147f45534f26 100644
--- a/clients/src/main/java/org/apache/kafka/common/serialization/ByteArrayDeserializer.java
+++ b/clients/src/main/java/org/apache/kafka/common/serialization/ByteArrayDeserializer.java
@@ -16,22 +16,10 @@
  */
 package org.apache.kafka.common.serialization;
 
-import java.util.Map;
-
 public class ByteArrayDeserializer implements Deserializer {
 
-    @Override
-    public void configure(Map configs, boolean isKey) {
-        // nothing to do
-    }
-
     @Override
     public byte[] deserialize(String topic, byte[] data) {
         return data;
     }
-
-    @Override
-    public void close() {
-        // nothing to do
-    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/serialization/ByteArraySerializer.java b/clients/src/main/java/org/apache/kafka/common/serialization/ByteArraySerializer.java
index d069e9495e6ec..6bebaa6531fc1 100644
--- a/clients/src/main/java/org/apache/kafka/common/serialization/ByteArraySerializer.java
+++ b/clients/src/main/java/org/apache/kafka/common/serialization/ByteArraySerializer.java
@@ -16,22 +16,9 @@
  */
 package org.apache.kafka.common.serialization;
 
-import java.util.Map;
-
 public class ByteArraySerializer implements Serializer {
-
-    @Override
-    public void configure(Map configs, boolean isKey) {
-        // nothing to do
-    }
-
     @Override
     public byte[] serialize(String topic, byte[] data) {
         return data;
     }
-
-    @Override
-    public void close() {
-        // nothing to do
-    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/serialization/ByteBufferDeserializer.java b/clients/src/main/java/org/apache/kafka/common/serialization/ByteBufferDeserializer.java
index d41f03c6675ee..0dfcf5f26c328 100644
--- a/clients/src/main/java/org/apache/kafka/common/serialization/ByteBufferDeserializer.java
+++ b/clients/src/main/java/org/apache/kafka/common/serialization/ByteBufferDeserializer.java
@@ -17,22 +17,12 @@
 package org.apache.kafka.common.serialization;
 
 import java.nio.ByteBuffer;
-import java.util.Map;
 
 public class ByteBufferDeserializer implements Deserializer {
-
-    public void configure(Map configs, boolean isKey) {
-        // nothing to do
-    }
-
     public ByteBuffer deserialize(String topic, byte[] data) {
         if (data == null)
             return null;
 
         return ByteBuffer.wrap(data);
     }
-
-    public void close() {
-        // nothing to do
-    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/serialization/ByteBufferSerializer.java b/clients/src/main/java/org/apache/kafka/common/serialization/ByteBufferSerializer.java
index c8c369272ddc3..9fb12544e0fe1 100644
--- a/clients/src/main/java/org/apache/kafka/common/serialization/ByteBufferSerializer.java
+++ b/clients/src/main/java/org/apache/kafka/common/serialization/ByteBufferSerializer.java
@@ -17,14 +17,8 @@
 package org.apache.kafka.common.serialization;
 
 import java.nio.ByteBuffer;
-import java.util.Map;
 
 public class ByteBufferSerializer implements Serializer {
-
-    public void configure(Map configs, boolean isKey) {
-        // nothing to do
-    }
-
     public byte[] serialize(String topic, ByteBuffer data) {
         if (data == null)
             return null;
@@ -43,8 +37,4 @@ public byte[] serialize(String topic, ByteBuffer data) {
         data.rewind();
         return ret;
     }
-
-    public void close() {
-        // nothing to do
-    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/serialization/BytesDeserializer.java b/clients/src/main/java/org/apache/kafka/common/serialization/BytesDeserializer.java
index 66b07eb58419e..1350dca21dd01 100644
--- a/clients/src/main/java/org/apache/kafka/common/serialization/BytesDeserializer.java
+++ b/clients/src/main/java/org/apache/kafka/common/serialization/BytesDeserializer.java
@@ -18,22 +18,11 @@
 
 import org.apache.kafka.common.utils.Bytes;
 
-import java.util.Map;
-
 public class BytesDeserializer implements Deserializer {
-
-    public void configure(Map configs, boolean isKey) {
-        // nothing to do
-    }
-
     public Bytes deserialize(String topic, byte[] data) {
         if (data == null)
             return null;
 
         return new Bytes(data);
     }
-
-    public void close() {
-        // nothing to do
-    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/serialization/BytesSerializer.java b/clients/src/main/java/org/apache/kafka/common/serialization/BytesSerializer.java
index 0dc4476d46d32..62ea6ec321fd3 100644
--- a/clients/src/main/java/org/apache/kafka/common/serialization/BytesSerializer.java
+++ b/clients/src/main/java/org/apache/kafka/common/serialization/BytesSerializer.java
@@ -18,23 +18,12 @@
 
 import org.apache.kafka.common.utils.Bytes;
 
-import java.util.Map;
-
 public class BytesSerializer implements Serializer {
-
-    public void configure(Map configs, boolean isKey) {
-        // nothing to do
-    }
-
     public byte[] serialize(String topic, Bytes data) {
         if (data == null)
             return null;
 
         return data.get();
     }
-
-    public void close() {
-        // nothing to do
-    }
 }
 
diff --git a/clients/src/main/java/org/apache/kafka/common/serialization/Deserializer.java b/clients/src/main/java/org/apache/kafka/common/serialization/Deserializer.java
index f9eb398cbee23..eb56485abce3c 100644
--- a/clients/src/main/java/org/apache/kafka/common/serialization/Deserializer.java
+++ b/clients/src/main/java/org/apache/kafka/common/serialization/Deserializer.java
@@ -16,6 +16,8 @@
  */
 package org.apache.kafka.common.serialization;
 
+import org.apache.kafka.common.header.Headers;
+
 import java.io.Closeable;
 import java.util.Map;
 
@@ -35,7 +37,9 @@ public interface Deserializer extends Closeable {
      * @param configs configs in key/value pairs
      * @param isKey whether is for key or value
      */
-    void configure(Map configs, boolean isKey);
+    default void configure(Map configs, boolean isKey) {
+        // intentionally left blank
+    }
 
     /**
      * Deserialize a record value from a byte array into a value or object.
@@ -45,6 +49,24 @@ public interface Deserializer extends Closeable {
      */
     T deserialize(String topic, byte[] data);
 
+    /**
+     * Deserialize a record value from a byte array into a value or object.
+     * @param topic topic associated with the data
+     * @param headers headers associated with the record; may be empty.
+     * @param data serialized bytes; may be null; implementations are recommended to handle null by returning a value or null rather than throwing an exception.
+     * @return deserialized typed data; may be null
+     */
+    default T deserialize(String topic, Headers headers, byte[] data) {
+        return deserialize(topic, data);
+    }
+
+    /**
+     * Close this deserializer.
+     * 
            
+     * This method must be idempotent as it may be called multiple times.
+     */
     @Override
-    void close();
+    default void close() {
+        // intentionally left blank
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/serialization/DoubleDeserializer.java b/clients/src/main/java/org/apache/kafka/common/serialization/DoubleDeserializer.java
index 24f6007cb35fc..0fa1cce4d7409 100644
--- a/clients/src/main/java/org/apache/kafka/common/serialization/DoubleDeserializer.java
+++ b/clients/src/main/java/org/apache/kafka/common/serialization/DoubleDeserializer.java
@@ -18,15 +18,8 @@
 
 import org.apache.kafka.common.errors.SerializationException;
 
-import java.util.Map;
-
 public class DoubleDeserializer implements Deserializer {
 
-    @Override
-    public void configure(Map configs, boolean isKey) {
-        // nothing to do
-    }
-
     @Override
     public Double deserialize(String topic, byte[] data) {
         if (data == null)
@@ -42,9 +35,4 @@ public Double deserialize(String topic, byte[] data) {
         }
         return Double.longBitsToDouble(value);
     }
-
-    @Override
-    public void close() {
-        // nothing to do
-    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/serialization/DoubleSerializer.java b/clients/src/main/java/org/apache/kafka/common/serialization/DoubleSerializer.java
index 7dd4edc3b62ab..99781b53d0e01 100644
--- a/clients/src/main/java/org/apache/kafka/common/serialization/DoubleSerializer.java
+++ b/clients/src/main/java/org/apache/kafka/common/serialization/DoubleSerializer.java
@@ -16,15 +16,7 @@
  */
 package org.apache.kafka.common.serialization;
 
-import java.util.Map;
-
 public class DoubleSerializer implements Serializer {
-
-    @Override
-    public void configure(Map configs, boolean isKey) {
-        // nothing to do
-    }
-
     @Override
     public byte[] serialize(String topic, Double data) {
         if (data == null)
@@ -42,9 +34,4 @@ public byte[] serialize(String topic, Double data) {
             (byte) bits
         };
     }
-
-    @Override
-    public void close() {
-        // nothing to do
-    }
 }
\ No newline at end of file
diff --git a/clients/src/main/java/org/apache/kafka/common/serialization/ExtendedDeserializer.java b/clients/src/main/java/org/apache/kafka/common/serialization/ExtendedDeserializer.java
index 2d5be4a148d31..2f4a012fc8f06 100644
--- a/clients/src/main/java/org/apache/kafka/common/serialization/ExtendedDeserializer.java
+++ b/clients/src/main/java/org/apache/kafka/common/serialization/ExtendedDeserializer.java
@@ -29,7 +29,9 @@
  *
  * A class that implements this interface is expected to have a constructor with no parameters.
  * @param 
+ * @deprecated This class has been deprecated and will be removed in a future release. Please use {@link Deserializer} instead.
  */
+@Deprecated
 public interface ExtendedDeserializer extends Deserializer {
 
     /**
diff --git a/clients/src/main/java/org/apache/kafka/common/serialization/ExtendedSerializer.java b/clients/src/main/java/org/apache/kafka/common/serialization/ExtendedSerializer.java
index 14fbb47b1dbf3..8c949807a03a5 100644
--- a/clients/src/main/java/org/apache/kafka/common/serialization/ExtendedSerializer.java
+++ b/clients/src/main/java/org/apache/kafka/common/serialization/ExtendedSerializer.java
@@ -29,7 +29,9 @@
  *
  * A class that implements this interface is expected to have a constructor with no parameters.
  * @param 
+ * @deprecated This class has been deprecated and will be removed in a future release. Please use {@link Serializer} instead.
  */
+@Deprecated
 public interface ExtendedSerializer extends Serializer {
 
     /**
diff --git a/clients/src/main/java/org/apache/kafka/common/serialization/FloatDeserializer.java b/clients/src/main/java/org/apache/kafka/common/serialization/FloatDeserializer.java
index 3834ce20b079b..09031779426c7 100644
--- a/clients/src/main/java/org/apache/kafka/common/serialization/FloatDeserializer.java
+++ b/clients/src/main/java/org/apache/kafka/common/serialization/FloatDeserializer.java
@@ -18,15 +18,7 @@
 
 import org.apache.kafka.common.errors.SerializationException;
 
-import java.util.Map;
-
 public class FloatDeserializer implements Deserializer {
-
-    @Override
-    public void configure(final Map configs, final boolean isKey) {
-        // nothing to do
-    }
-
     @Override
     public Float deserialize(final String topic, final byte[] data) {
         if (data == null)
@@ -42,10 +34,4 @@ public Float deserialize(final String topic, final byte[] data) {
         }
         return Float.intBitsToFloat(value);
     }
-
-    @Override
-    public void close() {
-        // nothing to do
-    }
-
 }
\ No newline at end of file
diff --git a/clients/src/main/java/org/apache/kafka/common/serialization/FloatSerializer.java b/clients/src/main/java/org/apache/kafka/common/serialization/FloatSerializer.java
index 6eb766dcd4254..aa72d43a91eca 100644
--- a/clients/src/main/java/org/apache/kafka/common/serialization/FloatSerializer.java
+++ b/clients/src/main/java/org/apache/kafka/common/serialization/FloatSerializer.java
@@ -16,15 +16,7 @@
  */
 package org.apache.kafka.common.serialization;
 
-import java.util.Map;
-
 public class FloatSerializer implements Serializer {
-
-    @Override
-    public void configure(final Map configs, final boolean isKey) {
-        // nothing to do
-    }
-
     @Override
     public byte[] serialize(final String topic, final Float data) {
         if (data == null)
@@ -38,9 +30,4 @@ public byte[] serialize(final String topic, final Float data) {
             (byte) bits
         };
     }
-
-    @Override
-    public void close() {
-        // nothing to do
-    }
 }
\ No newline at end of file
diff --git a/clients/src/main/java/org/apache/kafka/common/serialization/IntegerDeserializer.java b/clients/src/main/java/org/apache/kafka/common/serialization/IntegerDeserializer.java
index 45f8cf18fd4ea..20ca63f022413 100644
--- a/clients/src/main/java/org/apache/kafka/common/serialization/IntegerDeserializer.java
+++ b/clients/src/main/java/org/apache/kafka/common/serialization/IntegerDeserializer.java
@@ -18,14 +18,7 @@
 
 import org.apache.kafka.common.errors.SerializationException;
 
-import java.util.Map;
-
 public class IntegerDeserializer implements Deserializer {
-
-    public void configure(Map configs, boolean isKey) {
-        // nothing to do
-    }
-
     public Integer deserialize(String topic, byte[] data) {
         if (data == null)
             return null;
@@ -40,8 +33,4 @@ public Integer deserialize(String topic, byte[] data) {
         }
         return value;
     }
-
-    public void close() {
-        // nothing to do
-    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/serialization/IntegerSerializer.java b/clients/src/main/java/org/apache/kafka/common/serialization/IntegerSerializer.java
index f2144ceee707c..8ab531046006a 100644
--- a/clients/src/main/java/org/apache/kafka/common/serialization/IntegerSerializer.java
+++ b/clients/src/main/java/org/apache/kafka/common/serialization/IntegerSerializer.java
@@ -16,14 +16,7 @@
  */
 package org.apache.kafka.common.serialization;
 
-import java.util.Map;
-
 public class IntegerSerializer implements Serializer {
-
-    public void configure(Map configs, boolean isKey) {
-        // nothing to do
-    }
-
     public byte[] serialize(String topic, Integer data) {
         if (data == null)
             return null;
@@ -35,8 +28,4 @@ public byte[] serialize(String topic, Integer data) {
             data.byteValue()
         };
     }
-
-    public void close() {
-        // nothing to do
-    }
 }
\ No newline at end of file
diff --git a/clients/src/main/java/org/apache/kafka/common/serialization/LongDeserializer.java b/clients/src/main/java/org/apache/kafka/common/serialization/LongDeserializer.java
index a58b1d38cf3c1..1e445d2452f02 100644
--- a/clients/src/main/java/org/apache/kafka/common/serialization/LongDeserializer.java
+++ b/clients/src/main/java/org/apache/kafka/common/serialization/LongDeserializer.java
@@ -18,14 +18,7 @@
 
 import org.apache.kafka.common.errors.SerializationException;
 
-import java.util.Map;
-
 public class LongDeserializer implements Deserializer {
-
-    public void configure(Map configs, boolean isKey) {
-        // nothing to do
-    }
-
     public Long deserialize(String topic, byte[] data) {
         if (data == null)
             return null;
@@ -40,8 +33,4 @@ public Long deserialize(String topic, byte[] data) {
         }
         return value;
     }
-
-    public void close() {
-        // nothing to do
-    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/serialization/LongSerializer.java b/clients/src/main/java/org/apache/kafka/common/serialization/LongSerializer.java
index d37842c391496..436f0e01095ca 100644
--- a/clients/src/main/java/org/apache/kafka/common/serialization/LongSerializer.java
+++ b/clients/src/main/java/org/apache/kafka/common/serialization/LongSerializer.java
@@ -16,14 +16,7 @@
  */
 package org.apache.kafka.common.serialization;
 
-import java.util.Map;
-
 public class LongSerializer implements Serializer {
-
-    public void configure(Map configs, boolean isKey) {
-        // nothing to do
-    }
-
     public byte[] serialize(String topic, Long data) {
         if (data == null)
             return null;
@@ -39,8 +32,4 @@ public byte[] serialize(String topic, Long data) {
             data.byteValue()
         };
     }
-
-    public void close() {
-        // nothing to do
-    }
 }
\ No newline at end of file
diff --git a/clients/src/main/java/org/apache/kafka/common/serialization/Serde.java b/clients/src/main/java/org/apache/kafka/common/serialization/Serde.java
index fbcc7c2e0fcae..5b052e69f6969 100644
--- a/clients/src/main/java/org/apache/kafka/common/serialization/Serde.java
+++ b/clients/src/main/java/org/apache/kafka/common/serialization/Serde.java
@@ -34,14 +34,19 @@ public interface Serde extends Closeable {
      * @param configs configs in key/value pairs
      * @param isKey whether is for key or value
      */
-    void configure(Map configs, boolean isKey);
+    default void configure(Map configs, boolean isKey) {
+        // intentionally left blank
+    }
 
     /**
      * Close this serde class, which will close the underlying serializer and deserializer.
+     * 
            
      * This method has to be idempotent because it might be called multiple times.
      */
     @Override
-    void close();
+    default void close() {
+        // intentionally left blank
+    }
 
     Serializer serializer();
 
diff --git a/clients/src/main/java/org/apache/kafka/common/serialization/Serdes.java b/clients/src/main/java/org/apache/kafka/common/serialization/Serdes.java
index d6b4d2da6110b..7b47dc6d29ba2 100644
--- a/clients/src/main/java/org/apache/kafka/common/serialization/Serdes.java
+++ b/clients/src/main/java/org/apache/kafka/common/serialization/Serdes.java
@@ -20,13 +20,14 @@
 
 import java.nio.ByteBuffer;
 import java.util.Map;
+import java.util.UUID;
 
 /**
  * Factory for creating serializers / deserializers.
  */
 public class Serdes {
 
-    static protected class WrapperSerde implements Serde {
+    static public class WrapperSerde implements Serde {
         final private Serializer serializer;
         final private Deserializer deserializer;
 
@@ -58,6 +59,12 @@ public Deserializer deserializer() {
         }
     }
 
+    static public final class VoidSerde extends WrapperSerde {
+        public VoidSerde() {
+            super(new VoidSerializer(), new VoidDeserializer());
+        }
+    }
+
     static public final class LongSerde extends WrapperSerde {
         public LongSerde() {
             super(new LongSerializer(), new LongDeserializer());
@@ -112,6 +119,12 @@ public ByteArraySerde() {
         }
     }
 
+    static public final class UUIDSerde extends WrapperSerde {
+        public UUIDSerde() {
+            super(new UUIDSerializer(), new UUIDDeserializer());
+        }
+    }
+
     @SuppressWarnings("unchecked")
     static public  Serde serdeFrom(Class type) {
         if (String.class.isAssignableFrom(type)) {
@@ -150,9 +163,13 @@ static public  Serde serdeFrom(Class type) {
             return (Serde) Bytes();
         }
 
+        if (UUID.class.isAssignableFrom(type)) {
+            return (Serde) UUID();
+        }
+
         // TODO: we can also serializes objects of type T using generic Java serialization by default
         throw new IllegalArgumentException("Unknown class for built-in serializer. Supported types are: " +
-            "String, Short, Integer, Long, Float, Double, ByteArray, ByteBuffer, Bytes");
+            "String, Short, Integer, Long, Float, Double, ByteArray, ByteBuffer, Bytes, UUID");
     }
 
     /**
@@ -228,10 +245,24 @@ static public Serde Bytes() {
         return new BytesSerde();
     }
 
+    /*
+     * A serde for nullable {@code UUID} type
+     */
+    static public Serde UUID() {
+        return new UUIDSerde();
+    }
+
     /*
      * A serde for nullable {@code byte[]} type.
      */
     static public Serde ByteArray() {
         return new ByteArraySerde();
     }
+
+    /*
+     * A serde for {@code Void} type.
+     */
+    static public Serde Void() {
+        return new VoidSerde();
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/serialization/Serializer.java b/clients/src/main/java/org/apache/kafka/common/serialization/Serializer.java
index 96fe86b4bd96e..144b5ab945ebf 100644
--- a/clients/src/main/java/org/apache/kafka/common/serialization/Serializer.java
+++ b/clients/src/main/java/org/apache/kafka/common/serialization/Serializer.java
@@ -16,6 +16,8 @@
  */
 package org.apache.kafka.common.serialization;
 
+import org.apache.kafka.common.header.Headers;
+
 import java.io.Closeable;
 import java.util.Map;
 
@@ -35,7 +37,9 @@ public interface Serializer extends Closeable {
      * @param configs configs in key/value pairs
      * @param isKey whether is for key or value
      */
-    void configure(Map configs, boolean isKey);
+    default void configure(Map configs, boolean isKey) {
+        // intentionally left blank
+    }
 
     /**
      * Convert {@code data} into a byte array.
@@ -47,10 +51,24 @@ public interface Serializer extends Closeable {
     byte[] serialize(String topic, T data);
 
     /**
-     * Close this serializer.
+     * Convert {@code data} into a byte array.
      *
+     * @param topic topic associated with data
+     * @param headers headers associated with the record
+     * @param data typed data
+     * @return serialized bytes
+     */
+    default byte[] serialize(String topic, Headers headers, T data) {
+        return serialize(topic, data);
+    }
+
+    /**
+     * Close this serializer.
+     * 
            
      * This method must be idempotent as it may be called multiple times.
      */
     @Override
-    void close();
+    default void close() {
+        // intentionally left blank
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/serialization/ShortDeserializer.java b/clients/src/main/java/org/apache/kafka/common/serialization/ShortDeserializer.java
index 45aa8ae7ae3b4..7814a7bd7122f 100644
--- a/clients/src/main/java/org/apache/kafka/common/serialization/ShortDeserializer.java
+++ b/clients/src/main/java/org/apache/kafka/common/serialization/ShortDeserializer.java
@@ -18,14 +18,8 @@
 
 import org.apache.kafka.common.errors.SerializationException;
 
-import java.util.Map;
-
 public class ShortDeserializer implements Deserializer {
 
-    public void configure(Map configs, boolean isKey) {
-        // nothing to do
-    }
-
     public Short deserialize(String topic, byte[] data) {
         if (data == null)
             return null;
@@ -40,8 +34,4 @@ public Short deserialize(String topic, byte[] data) {
         }
         return value;
     }
-
-    public void close() {
-        // nothing to do
-    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/serialization/ShortSerializer.java b/clients/src/main/java/org/apache/kafka/common/serialization/ShortSerializer.java
index a66aaa09685d5..e54354b4dea05 100644
--- a/clients/src/main/java/org/apache/kafka/common/serialization/ShortSerializer.java
+++ b/clients/src/main/java/org/apache/kafka/common/serialization/ShortSerializer.java
@@ -16,14 +16,7 @@
  */
 package org.apache.kafka.common.serialization;
 
-import java.util.Map;
-
 public class ShortSerializer implements Serializer {
-
-    public void configure(Map configs, boolean isKey) {
-        // nothing to do
-    }
-
     public byte[] serialize(String topic, Short data) {
         if (data == null)
             return null;
@@ -33,8 +26,4 @@ public byte[] serialize(String topic, Short data) {
             data.byteValue()
         };
     }
-
-    public void close() {
-        // nothing to do
-    }
 }
\ No newline at end of file
diff --git a/clients/src/main/java/org/apache/kafka/common/serialization/StringDeserializer.java b/clients/src/main/java/org/apache/kafka/common/serialization/StringDeserializer.java
index c647c9b73324d..68e6c409cbabe 100644
--- a/clients/src/main/java/org/apache/kafka/common/serialization/StringDeserializer.java
+++ b/clients/src/main/java/org/apache/kafka/common/serialization/StringDeserializer.java
@@ -34,7 +34,7 @@ public void configure(Map configs, boolean isKey) {
         Object encodingValue = configs.get(propertyName);
         if (encodingValue == null)
             encodingValue = configs.get("deserializer.encoding");
-        if (encodingValue != null && encodingValue instanceof String)
+        if (encodingValue instanceof String)
             encoding = (String) encodingValue;
     }
 
@@ -49,9 +49,4 @@ public String deserialize(String topic, byte[] data) {
             throw new SerializationException("Error when deserializing byte[] to string due to unsupported encoding " + encoding);
         }
     }
-
-    @Override
-    public void close() {
-        // nothing to do
-    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/serialization/StringSerializer.java b/clients/src/main/java/org/apache/kafka/common/serialization/StringSerializer.java
index 28e41741a37fd..e16e19ac75c5c 100644
--- a/clients/src/main/java/org/apache/kafka/common/serialization/StringSerializer.java
+++ b/clients/src/main/java/org/apache/kafka/common/serialization/StringSerializer.java
@@ -49,9 +49,4 @@ public byte[] serialize(String topic, String data) {
             throw new SerializationException("Error when serializing string to byte[] due to unsupported encoding " + encoding);
         }
     }
-
-    @Override
-    public void close() {
-        // nothing to do
-    }
 }
\ No newline at end of file
diff --git a/clients/src/main/java/org/apache/kafka/common/serialization/UUIDDeserializer.java b/clients/src/main/java/org/apache/kafka/common/serialization/UUIDDeserializer.java
new file mode 100644
index 0000000000000..e852fc95afd4d
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/serialization/UUIDDeserializer.java
@@ -0,0 +1,55 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.serialization;
+
+import org.apache.kafka.common.errors.SerializationException;
+
+import java.io.UnsupportedEncodingException;
+import java.util.Map;
+import java.util.UUID;
+
+/**
+ *  We are converting the byte array to String before deserializing to UUID. String encoding defaults to UTF8 and can be customized by setting
+ *  the property key.deserializer.encoding, value.deserializer.encoding or deserializer.encoding. The first two take precedence over the last.
+ */
+public class UUIDDeserializer implements Deserializer {
+    private String encoding = "UTF8";
+
+    @Override
+    public void configure(Map configs, boolean isKey) {
+        String propertyName = isKey ? "key.deserializer.encoding" : "value.deserializer.encoding";
+        Object encodingValue = configs.get(propertyName);
+        if (encodingValue == null)
+            encodingValue = configs.get("deserializer.encoding");
+        if (encodingValue != null && encodingValue instanceof String)
+            encoding = (String) encodingValue;
+    }
+
+    @Override
+    public UUID deserialize(String topic, byte[] data) {
+        try {
+            if (data == null)
+                return null;
+            else
+                return UUID.fromString(new String(data, encoding));
+        } catch (UnsupportedEncodingException e) {
+            throw new SerializationException("Error when deserializing byte[] to UUID due to unsupported encoding " + encoding, e);
+        } catch (IllegalArgumentException e) {
+            throw new SerializationException("Error parsing data into UUID", e);
+        }
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/serialization/UUIDSerializer.java b/clients/src/main/java/org/apache/kafka/common/serialization/UUIDSerializer.java
new file mode 100644
index 0000000000000..908c202c7483f
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/serialization/UUIDSerializer.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 org.apache.kafka.common.serialization;
+
+import org.apache.kafka.common.errors.SerializationException;
+
+import java.io.UnsupportedEncodingException;
+import java.util.Map;
+import java.util.UUID;
+
+/**
+ *  We are converting UUID to String before serializing. String encoding defaults to UTF8 and can be customized by setting
+ *  the property key.deserializer.encoding, value.deserializer.encoding or deserializer.encoding. The first two take precedence over the last.
+ */
+public class UUIDSerializer implements Serializer {
+    private String encoding = "UTF8";
+
+    @Override
+    public void configure(Map configs, boolean isKey) {
+        String propertyName = isKey ? "key.serializer.encoding" : "value.serializer.encoding";
+        Object encodingValue = configs.get(propertyName);
+        if (encodingValue == null)
+            encodingValue = configs.get("serializer.encoding");
+        if (encodingValue instanceof String)
+            encoding = (String) encodingValue;
+    }
+
+    @Override
+    public byte[] serialize(String topic, UUID data) {
+        try {
+            if (data == null)
+                return null;
+            else
+                return data.toString().getBytes(encoding);
+        } catch (UnsupportedEncodingException e) {
+            throw new SerializationException("Error when serializing UUID to byte[] due to unsupported encoding " + encoding);
+        }
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/serialization/VoidDeserializer.java b/clients/src/main/java/org/apache/kafka/common/serialization/VoidDeserializer.java
new file mode 100644
index 0000000000000..08ff57a7bb28e
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/serialization/VoidDeserializer.java
@@ -0,0 +1,27 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.serialization;
+
+public class VoidDeserializer implements Deserializer {
+    @Override
+    public Void deserialize(String topic, byte[] data) {
+        if (data != null)
+            throw new IllegalArgumentException("Data should be null for a VoidDeserializer.");
+
+        return null;
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/serialization/VoidSerializer.java b/clients/src/main/java/org/apache/kafka/common/serialization/VoidSerializer.java
new file mode 100644
index 0000000000000..f1f2c6071a832
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/serialization/VoidSerializer.java
@@ -0,0 +1,24 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.common.serialization;
+
+public class VoidSerializer implements Serializer {
+    @Override
+    public byte[] serialize(String topic, Void data) {
+        return null;
+    }
+}
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 0a17ecd26859d..19f98d1b652aa 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,15 +36,17 @@ 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")) {
             props.load(resourceStream);
         } catch (Exception e) {
-            log.warn("Error while loading kafka-version.properties :" + e.getMessage());
+            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,13 @@ 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();
+            AppInfo mBean = new AppInfo(nowMs);
             ManagementFactory.getPlatformMBeanServer().registerMBean(mBean, name);
 
-            registerMetrics(metrics); // prefix will be added later by JmxReporter
+            registerMetrics(metrics, mBean); // prefix will be added later by JmxReporter
         } catch (JMException e) {
             log.warn("Error registering AppInfo mbean", e);
         }
@@ -77,6 +79,8 @@ public static synchronized void unregisterAppInfo(String prefix, String id, Metr
             unregisterMetrics(metrics);
         } catch (JMException e) {
             log.warn("Error unregistering AppInfo mbean", e);
+        } finally {
+            log.info("App info {} for {} unregistered", prefix, id);
         }
     }
 
@@ -84,10 +88,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 +100,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 {
-        public String getVersion();
-        public String getCommitId();
+        String getVersion();
+        String getCommitId();
+        Long getStartTimeMs();
     }
 
     public static class AppInfo implements AppInfoMBean {
 
-        public AppInfo() {
-            log.info("Kafka version : " + AppInfoParser.getVersion());
-            log.info("Kafka commitId : " + AppInfoParser.getCommitId());
+        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 +131,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/Base64.java b/clients/src/main/java/org/apache/kafka/common/utils/Base64.java
deleted file mode 100644
index e06e1eed5a0d8..0000000000000
--- a/clients/src/main/java/org/apache/kafka/common/utils/Base64.java
+++ /dev/null
@@ -1,261 +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 java.lang.invoke.MethodHandle;
-import java.lang.invoke.MethodHandles;
-import java.lang.invoke.MethodType;
-
-/**
- * Temporary class in order to support Java 7 and Java 9. `DatatypeConverter` is not in the base module of Java 9
- * and `java.util.Base64` was only introduced in Java 8.
- */
-public final class Base64 {
-
-    private static final Factory FACTORY;
-
-    static {
-        if (Java.IS_JAVA8_COMPATIBLE)
-            FACTORY = new Java8Factory();
-        else
-            FACTORY = new Java7Factory();
-    }
-
-    private Base64() {}
-
-    public static Encoder encoder() {
-        return FACTORY.encoder();
-    }
-
-    public static Encoder urlEncoderNoPadding() {
-        return FACTORY.urlEncoderNoPadding();
-    }
-
-    public static Decoder decoder() {
-        return FACTORY.decoder();
-    }
-
-    /* Contains a subset of methods from java.util.Base64.Encoder (introduced in Java 8) */
-    public interface Encoder {
-        String encodeToString(byte[] bytes);
-    }
-
-    /* Contains a subset of methods from java.util.Base64.Decoder (introduced in Java 8) */
-    public interface Decoder {
-        byte[] decode(String string);
-    }
-
-    private interface Factory {
-        Encoder urlEncoderNoPadding();
-        Encoder encoder();
-        Decoder decoder();
-    }
-
-    private static class Java8Factory implements Factory {
-
-        // Static final MethodHandles are optimised better by HotSpot
-        private static final MethodHandle URL_ENCODE_NO_PADDING;
-        private static final MethodHandle ENCODE;
-        private static final MethodHandle DECODE;
-
-        private static final Encoder URL_ENCODER_NO_PADDING;
-        private static final Encoder ENCODER;
-        private static final Decoder DECODER;
-
-        static {
-            try {
-                Class base64Class = Class.forName("java.util.Base64");
-
-                MethodHandles.Lookup lookup = MethodHandles.publicLookup();
-
-                Class juEncoderClass = Class.forName("java.util.Base64$Encoder");
-
-                MethodHandle getEncoder = lookup.findStatic(base64Class, "getEncoder",
-                        MethodType.methodType(juEncoderClass));
-                Object juEncoder;
-                try {
-                    juEncoder = getEncoder.invoke();
-                } catch (Throwable throwable) {
-                    // Invoked method doesn't throw checked exceptions, so safe to cast
-                    throw (RuntimeException) throwable;
-                }
-                MethodHandle encode = lookup.findVirtual(juEncoderClass, "encodeToString",
-                        MethodType.methodType(String.class, byte[].class));
-                ENCODE = encode.bindTo(juEncoder);
-
-
-                MethodHandle getUrlEncoder = lookup.findStatic(base64Class, "getUrlEncoder",
-                        MethodType.methodType(juEncoderClass));
-                Object juUrlEncoderNoPassing;
-                try {
-                    juUrlEncoderNoPassing = lookup.findVirtual(juEncoderClass, "withoutPadding",
-                            MethodType.methodType(juEncoderClass)).invoke(getUrlEncoder.invoke());
-                } catch (Throwable throwable) {
-                    // Invoked method doesn't throw checked exceptions, so safe to cast
-                    throw (RuntimeException) throwable;
-                }
-                URL_ENCODE_NO_PADDING = encode.bindTo(juUrlEncoderNoPassing);
-
-                Class juDecoderClass = Class.forName("java.util.Base64$Decoder");
-                MethodHandle getDecoder = lookup.findStatic(base64Class, "getDecoder",
-                        MethodType.methodType(juDecoderClass));
-                MethodHandle decode = lookup.findVirtual(juDecoderClass, "decode",
-                        MethodType.methodType(byte[].class, String.class));
-                try {
-                    DECODE = decode.bindTo(getDecoder.invoke());
-                } catch (Throwable throwable) {
-                    // Invoked method doesn't throw checked exceptions, so safe to cast
-                    throw (RuntimeException) throwable;
-                }
-
-                URL_ENCODER_NO_PADDING = new Encoder() {
-                    @Override
-                    public String encodeToString(byte[] bytes) {
-                        try {
-                            return (String) URL_ENCODE_NO_PADDING.invokeExact(bytes);
-                        } catch (Throwable throwable) {
-                            // Invoked method doesn't throw checked exceptions, so safe to cast
-                            throw (RuntimeException) throwable;
-                        }
-                    }
-                };
-
-                ENCODER = new Encoder() {
-                    @Override
-                    public String encodeToString(byte[] bytes) {
-                        try {
-                            return (String) ENCODE.invokeExact(bytes);
-                        } catch (Throwable throwable) {
-                            // Invoked method doesn't throw checked exceptions, so safe to cast
-                            throw (RuntimeException) throwable;
-                        }
-                    }
-                };
-
-                DECODER = new Decoder() {
-                    @Override
-                    public byte[] decode(String string) {
-                        try {
-                            return (byte[]) DECODE.invokeExact(string);
-                        } catch (Throwable throwable) {
-                            // Invoked method doesn't throw checked exceptions, so safe to cast
-                            throw (RuntimeException) throwable;
-                        }
-                    }
-                };
-
-            } catch (ReflectiveOperationException e) {
-                // Should never happen
-                throw new RuntimeException(e);
-            }
-        }
-
-        @Override
-        public Encoder urlEncoderNoPadding() {
-            return URL_ENCODER_NO_PADDING;
-        }
-
-        @Override
-        public Encoder encoder() {
-            return ENCODER;
-        }
-
-        @Override
-        public Decoder decoder() {
-            return DECODER;
-        }
-    }
-
-    private static class Java7Factory implements Factory {
-
-        // Static final MethodHandles are optimised better by HotSpot
-        private static final MethodHandle PRINT;
-        private static final MethodHandle PARSE;
-
-        static {
-            try {
-                Class cls = Class.forName("javax.xml.bind.DatatypeConverter");
-                MethodHandles.Lookup lookup = MethodHandles.publicLookup();
-                PRINT = lookup.findStatic(cls, "printBase64Binary", MethodType.methodType(String.class,
-                        byte[].class));
-                PARSE = lookup.findStatic(cls, "parseBase64Binary", MethodType.methodType(byte[].class,
-                        String.class));
-            } catch (ReflectiveOperationException e) {
-                // Should never happen
-                throw new RuntimeException(e);
-            }
-        }
-
-        public static final Encoder URL_ENCODER_NO_PADDING = new Encoder() {
-
-            @Override
-            public String encodeToString(byte[] bytes) {
-                String base64EncodedUUID = Java7Factory.encodeToString(bytes);
-                //Convert to URL safe variant by replacing + and / with - and _ respectively.
-                String urlSafeBase64EncodedUUID = base64EncodedUUID.replace("+", "-")
-                        .replace("/", "_");
-                // Remove the "==" padding at the end.
-                return urlSafeBase64EncodedUUID.substring(0, urlSafeBase64EncodedUUID.length() - 2);
-            }
-
-        };
-
-        public static final Encoder ENCODER = new Encoder() {
-            @Override
-            public String encodeToString(byte[] bytes) {
-                return Java7Factory.encodeToString(bytes);
-            }
-        };
-
-        public static final Decoder DECODER = new Decoder() {
-            @Override
-            public byte[] decode(String string) {
-                try {
-                    return (byte[]) PARSE.invokeExact(string);
-                } catch (Throwable throwable) {
-                    // Invoked method doesn't throw checked exceptions, so safe to cast
-                    throw (RuntimeException) throwable;
-                }
-            }
-        };
-
-        private static String encodeToString(byte[] bytes) {
-            try {
-                return (String) PRINT.invokeExact(bytes);
-            } catch (Throwable throwable) {
-                // Invoked method doesn't throw checked exceptions, so safe to cast
-                throw (RuntimeException) throwable;
-            }
-        }
-
-        @Override
-        public Encoder urlEncoderNoPadding() {
-            return URL_ENCODER_NO_PADDING;
-        }
-
-        @Override
-        public Encoder encoder() {
-            return ENCODER;
-        }
-
-        @Override
-        public Decoder decoder() {
-            return DECODER;
-        }
-    }
-}
diff --git a/clients/src/main/java/org/apache/kafka/common/utils/ByteBufferUnmapper.java b/clients/src/main/java/org/apache/kafka/common/utils/ByteBufferUnmapper.java
new file mode 100644
index 0000000000000..4777f7bf96b49
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/utils/ByteBufferUnmapper.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.utils;
+
+import java.io.IOException;
+import java.lang.invoke.MethodHandle;
+import java.lang.invoke.MethodHandles;
+import java.lang.reflect.Field;
+import java.lang.reflect.Method;
+import java.nio.ByteBuffer;
+
+import static java.lang.invoke.MethodHandles.constant;
+import static java.lang.invoke.MethodHandles.dropArguments;
+import static java.lang.invoke.MethodHandles.filterReturnValue;
+import static java.lang.invoke.MethodHandles.guardWithTest;
+import static java.lang.invoke.MethodHandles.lookup;
+import static java.lang.invoke.MethodType.methodType;
+
+/**
+ * Provides a mechanism to unmap mapped and direct byte buffers.
+ *
+ * The implementation was inspired by the one in Lucene's MMapDirectory.
+ */
+public final class ByteBufferUnmapper {
+
+    // null if unmap is not supported
+    private static final MethodHandle UNMAP;
+
+    // null if unmap is supported
+    private static final RuntimeException UNMAP_NOT_SUPPORTED_EXCEPTION;
+
+    static {
+        Object unmap = null;
+        RuntimeException exception = null;
+        try {
+            unmap = lookupUnmapMethodHandle();
+        } catch (RuntimeException e) {
+            exception = e;
+        }
+        if (unmap != null) {
+            UNMAP = (MethodHandle) unmap;
+            UNMAP_NOT_SUPPORTED_EXCEPTION = null;
+        } else {
+            UNMAP = null;
+            UNMAP_NOT_SUPPORTED_EXCEPTION = exception;
+        }
+    }
+
+    private ByteBufferUnmapper() {}
+
+    /**
+     * Unmap the provided mapped or direct byte buffer.
+     *
+     * This buffer cannot be referenced after this call, so it's highly recommended that any fields referencing it
+     * should be set to null.
+     *
+     * @throws IllegalArgumentException if buffer is not mapped or direct.
+     */
+    public static void unmap(String resourceDescription, ByteBuffer buffer) throws IOException {
+        if (!buffer.isDirect())
+            throw new IllegalArgumentException("Unmapping only works with direct buffers");
+        if (UNMAP == null)
+            throw UNMAP_NOT_SUPPORTED_EXCEPTION;
+
+        try {
+            UNMAP.invokeExact(buffer);
+        } catch (Throwable throwable) {
+            throw new IOException("Unable to unmap the mapped buffer: " + resourceDescription, throwable);
+        }
+    }
+
+    private static MethodHandle lookupUnmapMethodHandle() {
+        final MethodHandles.Lookup lookup = lookup();
+        try {
+            if (Java.IS_JAVA9_COMPATIBLE)
+                return unmapJava9(lookup);
+            else
+                return unmapJava7Or8(lookup);
+        } catch (ReflectiveOperationException | RuntimeException e1) {
+            throw new UnsupportedOperationException("Unmapping is not supported on this platform, because internal " +
+                "Java APIs are not compatible with this Kafka version", e1);
+        }
+    }
+
+    private static MethodHandle unmapJava7Or8(MethodHandles.Lookup lookup) throws ReflectiveOperationException {
+        /* "Compile" a MethodHandle that is roughly equivalent to the following lambda:
+         *
+         * (ByteBuffer buffer) -> {
+         *   sun.misc.Cleaner cleaner = ((java.nio.DirectByteBuffer) byteBuffer).cleaner();
+         *   if (nonNull(cleaner))
+         *     cleaner.clean();
+         *   else
+         *     noop(cleaner); // the noop is needed because MethodHandles#guardWithTest always needs both if and else
+         * }
+         */
+        Class directBufferClass = Class.forName("java.nio.DirectByteBuffer");
+        Method m = directBufferClass.getMethod("cleaner");
+        m.setAccessible(true);
+        MethodHandle directBufferCleanerMethod = lookup.unreflect(m);
+        Class cleanerClass = directBufferCleanerMethod.type().returnType();
+        MethodHandle cleanMethod = lookup.findVirtual(cleanerClass, "clean", methodType(void.class));
+        MethodHandle nonNullTest = lookup.findStatic(ByteBufferUnmapper.class, "nonNull",
+                methodType(boolean.class, Object.class)).asType(methodType(boolean.class, cleanerClass));
+        MethodHandle noop = dropArguments(constant(Void.class, null).asType(methodType(void.class)), 0, cleanerClass);
+        MethodHandle unmapper = filterReturnValue(directBufferCleanerMethod, guardWithTest(nonNullTest, cleanMethod, noop))
+                .asType(methodType(void.class, ByteBuffer.class));
+        return unmapper;
+    }
+
+    private static MethodHandle unmapJava9(MethodHandles.Lookup lookup) throws ReflectiveOperationException {
+        Class unsafeClass = Class.forName("sun.misc.Unsafe");
+        MethodHandle unmapper = lookup.findVirtual(unsafeClass, "invokeCleaner",
+                methodType(void.class, ByteBuffer.class));
+        Field f = unsafeClass.getDeclaredField("theUnsafe");
+        f.setAccessible(true);
+        Object theUnsafe = f.get(null);
+        return unmapper.bindTo(theUnsafe);
+    }
+
+    private static boolean nonNull(Object o) {
+        return o != null;
+    }
+}
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..15868721da9ea 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
@@ -28,6 +28,8 @@
  */
 public final class ByteUtils {
 
+    public static final ByteBuffer EMPTY_BUF = ByteBuffer.wrap(new byte[0]);
+
     private ByteUtils() {}
 
     /**
@@ -129,7 +131,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 +139,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 +150,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 +163,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 +174,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 +253,60 @@ public static long readVarlong(ByteBuffer buffer)  {
         return (value >>> 1) ^ -(value & 1);
     }
 
+    /**
+     * Read a double-precision 64-bit format IEEE 754 value.
+     *
+     * @param in The input to read from
+     * @return The double value read
+     */
+    public static double readDouble(DataInput in) throws IOException {
+        return in.readDouble();
+    }
+
+    /**
+     * Read a double-precision 64-bit format IEEE 754 value.
+     *
+     * @param buffer The buffer to read from
+     * @return The long value read
+     */
+    public static double readDouble(ByteBuffer buffer) {
+        return buffer.getDouble();
+    }
+
+    /**
+     * 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 +316,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 +328,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 +367,48 @@ public static void writeVarlong(long value, ByteBuffer buffer) {
     }
 
     /**
-     * Number of bytes needed to encode an integer in variable-length format.
+     * Write the given double following the double-precision 64-bit format IEEE 754 value into the output.
+     *
+     * @param value The value to write
+     * @param out The output to write to
+     */
+    public static void writeDouble(double value, DataOutput out) throws IOException {
+        out.writeDouble(value);
+    }
+
+    /**
+     * Write the given double following the double-precision 64-bit format IEEE 754 value into the buffer.
+     *
+     * @param value The value to write
+     * @param buffer The buffer to write to
+     */
+    public static void writeDouble(double value, ByteBuffer buffer) {
+        buffer.putDouble(value);
+    }
+
+    /**
+     * 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 e531d1f1760f3..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
@@ -25,6 +25,8 @@
  */
 public class Bytes implements Comparable {
 
+    public static final byte[] EMPTY = new byte[0];
+
     private static final char[] HEX_CHARS_UPPER = {'0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'A', 'B', 'C', 'D', 'E', 'F'};
 
     private final byte[] bytes;
@@ -138,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/CollectionUtils.java b/clients/src/main/java/org/apache/kafka/common/utils/CollectionUtils.java
index d7ab4e0b39caf..f3a7dbd52bf08 100644
--- a/clients/src/main/java/org/apache/kafka/common/utils/CollectionUtils.java
+++ b/clients/src/main/java/org/apache/kafka/common/utils/CollectionUtils.java
@@ -22,45 +22,54 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Collection;
+import java.util.stream.Collectors;
+
+public final class CollectionUtils {
+
+    private CollectionUtils() {}
+
+    /**
+     * Given two maps (A, B), returns all the key-value pairs in A whose keys are not contained in B
+     */
+    public static  Map subtractMap(Map minuend, Map subtrahend) {
+        return minuend.entrySet().stream()
+                .filter(entry -> !subtrahend.containsKey(entry.getKey()))
+                .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+    }
 
-public class CollectionUtils {
     /**
      * group data by topic
+     *
      * @param data Data to be partitioned
      * @param  Partition data type
      * @return partitioned data
      */
-    public static  Map> groupDataByTopic(Map data) {
+    public static  Map> groupPartitionDataByTopic(Map data) {
         Map> dataByTopic = new HashMap<>();
-        for (Map.Entry entry: data.entrySet()) {
+        for (Map.Entry entry : data.entrySet()) {
             String topic = entry.getKey().topic();
             int partition = entry.getKey().partition();
-            Map topicData = dataByTopic.get(topic);
-            if (topicData == null) {
-                topicData = new HashMap<>();
-                dataByTopic.put(topic, topicData);
-            }
+            Map topicData = dataByTopic.computeIfAbsent(topic, t -> new HashMap<>());
             topicData.put(partition, entry.getValue());
         }
         return dataByTopic;
     }
 
     /**
-     * group partitions by topic
-     * @param partitions
+     * Group a list of partitions by the topic name.
+     *
+     * @param partitions The partitions to collect
      * @return partitions per topic
      */
-    public static Map> groupDataByTopic(List partitions) {
+    public static Map> groupPartitionsByTopic(Collection partitions) {
         Map> partitionsByTopic = new HashMap<>();
-        for (TopicPartition tp: partitions) {
+        for (TopicPartition tp : partitions) {
             String topic = tp.topic();
-            List topicData = partitionsByTopic.get(topic);
-            if (topicData == null) {
-                topicData = new ArrayList<>();
-                partitionsByTopic.put(topic, topicData);
-            }
+            List topicData = partitionsByTopic.computeIfAbsent(topic, t -> new ArrayList<>());
             topicData.add(tp.partition());
         }
-        return  partitionsByTopic;
+        return partitionsByTopic;
     }
+
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/utils/ConfigUtils.java b/clients/src/main/java/org/apache/kafka/common/utils/ConfigUtils.java
new file mode 100644
index 0000000000000..504a1f06431d0
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/utils/ConfigUtils.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.common.utils;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+public class ConfigUtils {
+
+    private static final Logger log = LoggerFactory.getLogger(ConfigUtils.class);
+
+    /**
+     * Translates deprecated configurations into their non-deprecated equivalents
+     *
+     * This is a convenience method for {@link ConfigUtils#translateDeprecatedConfigs(Map, Map)}
+     * until we can use Java 9+ {@code Map.of(..)} and {@code Set.of(...)}
+     *
+     * @param configs the input configuration
+     * @param aliasGroups An array of arrays of synonyms.  Each synonym array begins with the non-deprecated synonym
+     *                    For example, new String[][] { { a, b }, { c, d, e} }
+     *                    would declare b as a deprecated synonym for a,
+     *                    and d and e as deprecated synonyms for c.
+     *                    The ordering of synonyms determines the order of precedence
+     *                    (e.g. the first synonym takes precedence over the second one)
+     * @return a new configuration map with deprecated  keys translated to their non-deprecated equivalents
+     */
+    public static  Map translateDeprecatedConfigs(Map configs, String[][] aliasGroups) {
+        return translateDeprecatedConfigs(configs, Stream.of(aliasGroups)
+            .collect(Collectors.toMap(x -> x[0], x -> Stream.of(x).skip(1).collect(Collectors.toList()))));
+    }
+
+    /**
+     * Translates deprecated configurations into their non-deprecated equivalents
+     *
+     * @param configs the input configuration
+     * @param aliasGroups A map of config to synonyms.  Each key is the non-deprecated synonym
+     *                    For example, Map.of(a , Set.of(b), c, Set.of(d, e))
+     *                    would declare b as a deprecated synonym for a,
+     *                    and d and e as deprecated synonyms for c.
+     *                    The ordering of synonyms determines the order of precedence
+     *                    (e.g. the first synonym takes precedence over the second one)
+     * @return a new configuration map with deprecated  keys translated to their non-deprecated equivalents
+     */
+    public static  Map translateDeprecatedConfigs(Map configs,
+                                                                Map> aliasGroups) {
+        Set aliasSet = Stream.concat(
+            aliasGroups.keySet().stream(),
+            aliasGroups.values().stream().flatMap(Collection::stream))
+            .collect(Collectors.toSet());
+
+        // pass through all configurations without aliases
+        Map newConfigs = configs.entrySet().stream()
+            .filter(e -> !aliasSet.contains(e.getKey()))
+            // filter out null values
+            .filter(e -> Objects.nonNull(e.getValue()))
+            .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+
+        aliasGroups.forEach((target, aliases) -> {
+            List deprecated = aliases.stream()
+                .filter(configs::containsKey)
+                .collect(Collectors.toList());
+
+            if (deprecated.isEmpty()) {
+                // No deprecated key(s) found.
+                if (configs.containsKey(target)) {
+                    newConfigs.put(target, configs.get(target));
+                }
+                return;
+            }
+
+            String aliasString = String.join(", ", deprecated);
+
+            if (configs.containsKey(target)) {
+                // Ignore the deprecated key(s) because the actual key was set.
+                log.error(target + " was configured, as well as the deprecated alias(es) " +
+                          aliasString + ".  Using the value of " + target);
+                newConfigs.put(target, configs.get(target));
+            } else if (deprecated.size() > 1) {
+                log.error("The configuration keys " + aliasString + " are deprecated and may be " +
+                          "removed in the future.  Additionally, this configuration is ambigous because " +
+                          "these configuration keys are all aliases for " + target + ".  Please update " +
+                          "your configuration to have only " + target + " set.");
+                newConfigs.put(target, configs.get(deprecated.get(0)));
+            } else {
+                log.warn("Configuration key " + deprecated.get(0) + " is deprecated and may be removed " +
+                         "in the future.  Please update your configuration to use " + target + " instead.");
+                newConfigs.put(target, configs.get(deprecated.get(0)));
+            }
+        });
+
+        return newConfigs;
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/utils/Crc32.java b/clients/src/main/java/org/apache/kafka/common/utils/Crc32.java
index 9bfc57641b7f8..777ea2bdfd7fb 100644
--- a/clients/src/main/java/org/apache/kafka/common/utils/Crc32.java
+++ b/clients/src/main/java/org/apache/kafka/common/utils/Crc32.java
@@ -89,6 +89,7 @@ public void reset() {
         crc = 0xffffffff;
     }
 
+    @SuppressWarnings("fallthrough")
     @Override
     public void update(byte[] b, int off, int len) {
         if (off < 0 || len < 0 || off > b.length - len)
diff --git a/clients/src/main/java/org/apache/kafka/common/utils/Exit.java b/clients/src/main/java/org/apache/kafka/common/utils/Exit.java
index 7c2a663166e27..45f92e41e46c6 100644
--- a/clients/src/main/java/org/apache/kafka/common/utils/Exit.java
+++ b/clients/src/main/java/org/apache/kafka/common/utils/Exit.java
@@ -26,6 +26,10 @@ public interface Procedure {
         void execute(int statusCode, String message);
     }
 
+    public interface ShutdownHookAdder {
+        void addShutdownHook(String name, Runnable runnable);
+    }
+
     private static final Procedure DEFAULT_HALT_PROCEDURE = new Procedure() {
         @Override
         public void execute(int statusCode, String message) {
@@ -40,8 +44,19 @@ public void execute(int statusCode, String message) {
         }
     };
 
+    private static final ShutdownHookAdder DEFAULT_SHUTDOWN_HOOK_ADDER = new ShutdownHookAdder() {
+        @Override
+        public void addShutdownHook(String name, Runnable runnable) {
+            if (name != null)
+                Runtime.getRuntime().addShutdownHook(KafkaThread.nonDaemon(name, runnable));
+            else
+                Runtime.getRuntime().addShutdownHook(new Thread(runnable));
+        }
+    };
+
     private volatile static Procedure exitProcedure = DEFAULT_EXIT_PROCEDURE;
     private volatile static Procedure haltProcedure = DEFAULT_HALT_PROCEDURE;
+    private volatile static ShutdownHookAdder shutdownHookAdder = DEFAULT_SHUTDOWN_HOOK_ADDER;
 
     public static void exit(int statusCode) {
         exit(statusCode, null);
@@ -59,6 +74,10 @@ public static void halt(int statusCode, String message) {
         haltProcedure.execute(statusCode, message);
     }
 
+    public static void addShutdownHook(String name, Runnable runnable) {
+        shutdownHookAdder.addShutdownHook(name, runnable);
+    }
+
     public static void setExitProcedure(Procedure procedure) {
         exitProcedure = procedure;
     }
@@ -67,6 +86,10 @@ public static void setHaltProcedure(Procedure procedure) {
         haltProcedure = procedure;
     }
 
+    public static void setShutdownHookAdder(ShutdownHookAdder shutdownHookAdder) {
+        Exit.shutdownHookAdder = shutdownHookAdder;
+    }
+
     public static void resetExitProcedure() {
         exitProcedure = DEFAULT_EXIT_PROCEDURE;
     }
@@ -75,4 +98,7 @@ public static void resetHaltProcedure() {
         haltProcedure = DEFAULT_HALT_PROCEDURE;
     }
 
+    public static void resetShutdownHookAdder() {
+        shutdownHookAdder = DEFAULT_SHUTDOWN_HOOK_ADDER;
+    }
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/utils/ExponentialBackoff.java b/clients/src/main/java/org/apache/kafka/common/utils/ExponentialBackoff.java
new file mode 100644
index 0000000000000..73972c834ac8a
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/utils/ExponentialBackoff.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 org.apache.kafka.common.utils;
+
+import java.util.concurrent.ThreadLocalRandom;
+
+/**
+ * An utility class for keeping the parameters and providing the value of exponential
+ * retry backoff, exponential reconnect backoff, exponential timeout, etc.
+ * The formula is:
+ * Backoff(attempts) = random(1 - jitter, 1 + jitter) * initialInterval * multiplier ^ attempts
+ * If initialInterval is greater or equal than maxInterval, a constant backoff of will be provided
+ * This class is thread-safe
+ */
+public class ExponentialBackoff {
+    private final int multiplier;
+    private final double expMax;
+    private final long initialInterval;
+    private final double jitter;
+
+    public ExponentialBackoff(long initialInterval, int multiplier, long maxInterval, double jitter) {
+        this.initialInterval = initialInterval;
+        this.multiplier = multiplier;
+        this.jitter = jitter;
+        this.expMax = maxInterval > initialInterval ?
+                Math.log(maxInterval / (double) Math.max(initialInterval, 1)) / Math.log(multiplier) : 0;
+    }
+
+    public long backoff(long attempts) {
+        if (expMax == 0) {
+            return initialInterval;
+        }
+        double exp = Math.min(attempts, this.expMax);
+        double term = initialInterval * Math.pow(multiplier, exp);
+        double randomFactor = ThreadLocalRandom.current().nextDouble(1 - jitter, 1 + jitter);
+        return (long) (randomFactor * term);
+    }
+}
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..175282e7e09f6
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/utils/FixedOrderMap.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.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");
+    }
+
+    @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/ImplicitLinkedHashCollection.java b/clients/src/main/java/org/apache/kafka/common/utils/ImplicitLinkedHashCollection.java
new file mode 100644
index 0000000000000..3e1bccc476ed9
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/utils/ImplicitLinkedHashCollection.java
@@ -0,0 +1,679 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.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.
+ *
+ * Like java.util.LinkedHashSet, this collection maintains a linked list of elements.
+ * However, rather than using a separate linked list, this collection embeds the next
+ * and previous fields into the elements themselves.  This reduces memory consumption,
+ * because it means that we only have to store one Java object per element, rather
+ * than multiple.
+ *
+ * The next and previous fields are stored as array indices rather than pointers.
+ * This ensures that the fields only take 32 bits, even when pointers are 64 bits.
+ * It also makes the garbage collector's job easier, because it reduces the number of
+ * pointers that it must chase.
+ *
+ * This class uses linear probing.  Unlike HashMap (but like HashTable), we don't force
+ * the size to be a power of 2.  This saves memory.
+ *
+ * This set does not allow null elements.  It does not have internal synchronization.
+ */
+public class ImplicitLinkedHashCollection extends AbstractCollection {
+    /**
+     * The interface which elements of this collection must implement.  The prev,
+     * setPrev, next, and setNext functions handle manipulating the implicit linked
+     * list which these elements reside in inside the collection.
+     * elementKeysAreEqual() is the function which this collection uses to compare
+     * elements.
+     */
+    public interface Element {
+        int prev();
+        void setPrev(int prev);
+        int next();
+        void setNext(int next);
+        default boolean elementKeysAreEqual(Object other) {
+            return equals(other);
+        }
+    }
+
+    /**
+     * A special index value used to indicate that the next or previous field is
+     * the head.
+     */
+    private static final int HEAD_INDEX = -1;
+
+    /**
+     * A special index value used for next and previous indices which have not
+     * been initialized.
+     */
+    public static final int INVALID_INDEX = -2;
+
+    /**
+     * The minimum new capacity for a non-empty implicit hash set.
+     */
+    private static final int MIN_NONEMPTY_CAPACITY = 5;
+
+    /**
+     * A static empty array used to avoid object allocations when the capacity is zero.
+     */
+    private static final Element[] EMPTY_ELEMENTS = new Element[0];
+
+    private static class HeadElement implements Element {
+        static final HeadElement EMPTY = new HeadElement();
+
+        private int prev = HEAD_INDEX;
+        private int next = HEAD_INDEX;
+
+        @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;
+        }
+    }
+
+    private static Element indexToElement(Element head, Element[] elements, int index) {
+        if (index == HEAD_INDEX) {
+            return head;
+        }
+        return elements[index];
+    }
+
+    private static void addToListTail(Element head, Element[] elements, int elementIdx) {
+        int oldTailIdx = head.prev();
+        Element element = indexToElement(head, elements, elementIdx);
+        Element oldTail = indexToElement(head, elements, oldTailIdx);
+        head.setPrev(elementIdx);
+        oldTail.setNext(elementIdx);
+        element.setPrev(oldTailIdx);
+        element.setNext(HEAD_INDEX);
+    }
+
+    private static void removeFromList(Element head, Element[] elements, int elementIdx) {
+        Element element = indexToElement(head, elements, elementIdx);
+        elements[elementIdx] = null;
+        int prevIdx = element.prev();
+        int nextIdx = element.next();
+        Element prev = indexToElement(head, elements, prevIdx);
+        Element next = indexToElement(head, elements, nextIdx);
+        prev.setNext(nextIdx);
+        next.setPrev(prevIdx);
+        element.setNext(INVALID_INDEX);
+        element.setPrev(INVALID_INDEX);
+    }
+
+    private class ImplicitLinkedHashCollectionIterator implements ListIterator {
+        private int index = 0;
+        private Element cur;
+        private Element lastReturned;
+
+        ImplicitLinkedHashCollectionIterator(int index) {
+            this.cur = indexToElement(head, elements, head.next());
+            for (int i = 0; i < index; ++i) {
+                next();
+            }
+            this.lastReturned = null;
+        }
+
+        @Override
+        public boolean hasNext() {
+            return cur != head;
+        }
+
+        @Override
+        public boolean hasPrevious() {
+            return indexToElement(head, elements, cur.prev()) != head;
+        }
+
+        @Override
+        public E next() {
+            if (!hasNext()) {
+                throw new NoSuchElementException();
+            }
+            @SuppressWarnings("unchecked")
+            E returnValue = (E) cur;
+            lastReturned = cur;
+            cur = indexToElement(head, elements, cur.next());
+            ++index;
+            return returnValue;
+        }
+
+        @Override
+        public E previous() {
+            Element prev = indexToElement(head, elements, cur.prev());
+            if (prev == head) {
+                throw new NoSuchElementException();
+            }
+            cur = prev;
+            --index;
+            lastReturned = cur;
+            @SuppressWarnings("unchecked")
+            E returnValue = (E) cur;
+            return returnValue;
+        }
+
+        @Override
+        public int nextIndex() {
+            return index;
+        }
+
+        @Override
+        public int previousIndex() {
+            return index - 1;
+        }
+
+        @Override
+        public void remove() {
+            if (lastReturned == null) {
+                throw new IllegalStateException();
+            }
+            Element nextElement = indexToElement(head, elements, lastReturned.next());
+            ImplicitLinkedHashCollection.this.removeElementAtSlot(nextElement.prev());
+            if (lastReturned == cur) {
+                // If the element we are removing was cur, set cur to cur->next.
+                cur = nextElement;
+            } else {
+                // If the element we are removing comes before cur, decrement the index,
+                // since there are now fewer entries before cur.
+                --index;
+            }
+            lastReturned = null;
+        }
+
+        @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();
+        }
+    }
+
+    private Element head;
+
+    Element[] elements;
+
+    private int size;
+
+    /**
+     * Returns an iterator that will yield every element in the set.
+     * The elements will be returned in the order that they were inserted in.
+     *
+     * Do not modify the set while you are iterating over it (except by calling
+     * remove on the iterator itself, of course.)
+     */
+    @Override
+    final public Iterator iterator() {
+        return listIterator(0);
+    }
+
+    private ListIterator listIterator(int index) {
+        return new ImplicitLinkedHashCollectionIterator(index);
+    }
+
+    final int slot(Element[] curElements, Object e) {
+        return (e.hashCode() & 0x7fffffff) % curElements.length;
+    }
+
+    /**
+     * Find an element matching an example element.
+     *
+     * Using the element's hash code, we can look up the slot where it belongs.
+     * However, it may not have ended up in exactly this slot, due to a collision.
+     * Therefore, we must search forward in the array until we hit a null, before
+     * concluding that the element is not present.
+     *
+     * @param key               The element to match.
+     * @return                  The match index, or INVALID_INDEX if no match was found.
+     */
+    final private int findIndexOfEqualElement(Object key) {
+        if (key == null || size == 0) {
+            return INVALID_INDEX;
+        }
+        int slot = slot(elements, key);
+        for (int seen = 0; seen < elements.length; seen++) {
+            Element element = elements[slot];
+            if (element == null) {
+                return INVALID_INDEX;
+            }
+            if (element.elementKeysAreEqual(key)) {
+                return slot;
+            }
+            slot = (slot + 1) % elements.length;
+        }
+        return INVALID_INDEX;
+    }
+
+    /**
+     * An element e in the collection such that e.elementKeysAreEqual(key) and
+     * e.hashCode() == key.hashCode().
+     *
+     * @param key   The element to match.
+     * @return      The matching element, or null if there were none.
+     */
+    final public E find(E key) {
+        int index = findIndexOfEqualElement(key);
+        if (index == INVALID_INDEX) {
+            return null;
+        }
+        @SuppressWarnings("unchecked")
+        E result = (E) elements[index];
+        return result;
+    }
+
+    /**
+     * Returns the number of elements in the set.
+     */
+    @Override
+    final public int size() {
+        return size;
+    }
+
+    /**
+     * Returns true if there is at least one element e in the collection such
+     * that key.elementKeysAreEqual(e) and key.hashCode() == e.hashCode().
+     *
+     * @param key       The object to try to match.
+     */
+    @Override
+    final public boolean contains(Object key) {
+        return findIndexOfEqualElement(key) != INVALID_INDEX;
+    }
+
+    private static int calculateCapacity(int expectedNumElements) {
+        // Avoid using even-sized capacities, to get better key distribution.
+        int newCapacity = (2 * expectedNumElements) + 1;
+        // Don't use a capacity that is too small.
+        return Math.max(newCapacity, MIN_NONEMPTY_CAPACITY);
+    }
+
+    /**
+     * Add a new element to the collection.
+     *
+     * @param newElement    The new element.
+     *
+     * @return              True if the element was added to the collection;
+     *                      false if it was not, because there was an existing equal element.
+     */
+    @Override
+    final public boolean add(E newElement) {
+        if (newElement == null) {
+            return false;
+        }
+        if (newElement.prev() != INVALID_INDEX || newElement.next() != INVALID_INDEX) {
+            return false;
+        }
+        if ((size + 1) >= elements.length / 2) {
+            changeCapacity(calculateCapacity(elements.length));
+        }
+        int slot = addInternal(newElement, elements);
+        if (slot >= 0) {
+            addToListTail(head, elements, slot);
+            size++;
+            return true;
+        }
+        return false;
+    }
+
+    final public void mustAdd(E newElement) {
+        if (!add(newElement)) {
+            throw new RuntimeException("Unable to add " + newElement);
+        }
+    }
+
+    /**
+     * Adds a new element to the appropriate place in the elements array.
+     *
+     * @param newElement    The new element to add.
+     * @param addElements   The elements array.
+     * @return              The index at which the element was inserted, or INVALID_INDEX
+     *                      if the element could not be inserted.
+     */
+    int addInternal(Element newElement, Element[] addElements) {
+        int slot = slot(addElements, newElement);
+        for (int seen = 0; seen < addElements.length; seen++) {
+            Element element = addElements[slot];
+            if (element == null) {
+                addElements[slot] = newElement;
+                return slot;
+            }
+            if (element.elementKeysAreEqual(newElement)) {
+                return INVALID_INDEX;
+            }
+            slot = (slot + 1) % addElements.length;
+        }
+        throw new RuntimeException("Not enough hash table slots to add a new element.");
+    }
+
+    private void changeCapacity(int newCapacity) {
+        Element[] newElements = new Element[newCapacity];
+        HeadElement newHead = new HeadElement();
+        int oldSize = size;
+        for (Iterator iter = iterator(); iter.hasNext(); ) {
+            Element element = iter.next();
+            iter.remove();
+            int newSlot = addInternal(element, newElements);
+            addToListTail(newHead, newElements, newSlot);
+        }
+        this.elements = newElements;
+        this.head = newHead;
+        this.size = oldSize;
+    }
+
+    /**
+     * Remove the first element e such that key.elementKeysAreEqual(e)
+     * and key.hashCode == e.hashCode.
+     *
+     * @param key       The object to try to match.
+     * @return          True if an element was removed; false otherwise.
+     */
+    @Override
+    final public boolean remove(Object key) {
+        int slot = findElementToRemove(key);
+        if (slot == INVALID_INDEX) {
+            return false;
+        }
+        removeElementAtSlot(slot);
+        return true;
+    }
+
+    int findElementToRemove(Object key) {
+        return findIndexOfEqualElement(key);
+    }
+
+    /**
+     * Remove an element in a particular slot.
+     *
+     * @param slot      The slot of the element to remove.
+     *
+     * @return          True if an element was removed; false otherwise.
+     */
+    private boolean removeElementAtSlot(int slot) {
+        size--;
+        removeFromList(head, elements, slot);
+        slot = (slot + 1) % elements.length;
+
+        // Find the next empty slot
+        int endSlot = slot;
+        for (int seen = 0; seen < elements.length; seen++) {
+            Element element = elements[endSlot];
+            if (element == null) {
+                break;
+            }
+            endSlot = (endSlot + 1) % elements.length;
+        }
+
+        // We must preserve the denseness invariant.  The denseness invariant says that
+        // any element is either in the slot indicated by its hash code, or a slot which
+        // is not separated from that slot by any nulls.
+        // Reseat all elements in between the deleted element and the next empty slot.
+        while (slot != endSlot) {
+            reseat(slot);
+            slot = (slot + 1) % elements.length;
+        }
+        return true;
+    }
+
+    private void reseat(int prevSlot) {
+        Element element = elements[prevSlot];
+        int newSlot = slot(elements, element);
+        for (int seen = 0; seen < elements.length; seen++) {
+            Element e = elements[newSlot];
+            if ((e == null) || (e == element)) {
+                break;
+            }
+            newSlot = (newSlot + 1) % elements.length;
+        }
+        if (newSlot == prevSlot) {
+            return;
+        }
+        Element prev = indexToElement(head, elements, element.prev());
+        prev.setNext(newSlot);
+        Element next = indexToElement(head, elements, element.next());
+        next.setPrev(newSlot);
+        elements[prevSlot] = null;
+        elements[newSlot] = element;
+    }
+
+    /**
+     * Create a new ImplicitLinkedHashCollection.
+     */
+    public ImplicitLinkedHashCollection() {
+        this(0);
+    }
+
+    /**
+     * 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 ImplicitLinkedHashCollection(int expectedNumElements) {
+        clear(expectedNumElements);
+    }
+
+    /**
+     * Create a new ImplicitLinkedHashCollection.
+     *
+     * @param iter                  We will add all the elements accessible through this iterator
+     *                              to the set.
+     */
+    public ImplicitLinkedHashCollection(Iterator iter) {
+        clear(0);
+        while (iter.hasNext()) {
+            mustAdd(iter.next());
+        }
+    }
+
+    /**
+     * Removes all of the elements from this set.
+     */
+    @Override
+    final public void clear() {
+        clear(elements.length);
+    }
+
+    /**
+     * Moves an element which is already in the collection so that it comes last
+     * in iteration order.
+     */
+    final public void moveToEnd(E element) {
+        if (element.prev() == INVALID_INDEX || element.next() == INVALID_INDEX) {
+            throw new RuntimeException("Element " + element + " is not in the collection.");
+        }
+        Element prevElement = indexToElement(head, elements, element.prev());
+        Element nextElement = indexToElement(head, elements, element.next());
+        int slot = prevElement.next();
+        prevElement.setNext(element.next());
+        nextElement.setPrev(element.prev());
+        addToListTail(head, elements, slot);
+    }
+
+    /**
+     * Removes all of the elements from this set, and resets the set capacity
+     * based on the provided expected number of elements.
+     */
+    final public void clear(int expectedNumElements) {
+        if (expectedNumElements == 0) {
+            // Optimize away object allocations for empty sets.
+            this.head = HeadElement.EMPTY;
+            this.elements = EMPTY_ELEMENTS;
+            this.size = 0;
+        } else {
+            this.head = new HeadElement();
+            this.elements = new Element[calculateCapacity(expectedNumElements)];
+            this.size = 0;
+        }
+    }
+
+    /**
+     * 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/ImplicitLinkedHashMultiCollection.java b/clients/src/main/java/org/apache/kafka/common/utils/ImplicitLinkedHashMultiCollection.java
new file mode 100644
index 0000000000000..b95b7de21673b
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/utils/ImplicitLinkedHashMultiCollection.java
@@ -0,0 +1,142 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.ArrayList;
+import java.util.Collections;
+import java.util.Iterator;
+import java.util.List;
+
+/**
+ * A memory-efficient hash multiset which tracks the order of insertion of elements.
+ * 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
+ * have equivalent keys.
+ *
+ * We use reference equality when adding elements to the set.  A new element A can
+ * be added if there is no existing element B such that A == B.  If an element B
+ * exists such that A.elementKeysAreEqual(B), A will still be added.
+ *
+ * When deleting an element A from the set, we will try to delete the element B such
+ * that A == B.  If no such element can be found, we will try to delete an element B
+ * such that A.elementKeysAreEqual(B).
+ *
+ * contains() and find() are unchanged from the base class-- they will look for element
+ * based on object equality via elementKeysAreEqual, not reference equality.
+ *
+ * This multiset does not allow null elements.  It does not have internal synchronization.
+ */
+public class ImplicitLinkedHashMultiCollection
+        extends ImplicitLinkedHashCollection {
+    public ImplicitLinkedHashMultiCollection() {
+        super(0);
+    }
+
+    public ImplicitLinkedHashMultiCollection(int expectedNumElements) {
+        super(expectedNumElements);
+    }
+
+    public ImplicitLinkedHashMultiCollection(Iterator iter) {
+        super(iter);
+    }
+
+
+    /**
+     * Adds a new element to the appropriate place in the elements array.
+     *
+     * @param newElement    The new element to add.
+     * @param addElements   The elements array.
+     * @return              The index at which the element was inserted, or INVALID_INDEX
+     *                      if the element could not be inserted.
+     */
+    @Override
+    int addInternal(Element newElement, Element[] addElements) {
+        int slot = slot(addElements, newElement);
+        for (int seen = 0; seen < addElements.length; seen++) {
+            Element element = addElements[slot];
+            if (element == null) {
+                addElements[slot] = newElement;
+                return slot;
+            }
+            if (element == newElement) {
+                return INVALID_INDEX;
+            }
+            slot = (slot + 1) % addElements.length;
+        }
+        throw new RuntimeException("Not enough hash table slots to add a new element.");
+    }
+
+    /**
+     * Find an element matching an example element.
+     *
+     * @param key               The element to match.
+     *
+     * @return                  The match index, or INVALID_INDEX if no match was found.
+     */
+    @Override
+    int findElementToRemove(Object key) {
+        if (key == null || size() == 0) {
+            return INVALID_INDEX;
+        }
+        int slot = slot(elements, key);
+        int bestSlot = INVALID_INDEX;
+        for (int seen = 0; seen < elements.length; seen++) {
+            Element element = elements[slot];
+            if (element == null) {
+                return bestSlot;
+            }
+            if (key == element) {
+                return slot;
+            } else if (element.elementKeysAreEqual(key)) {
+                bestSlot = slot;
+            }
+            slot = (slot + 1) % elements.length;
+        }
+        return INVALID_INDEX;
+    }
+
+    /**
+     * Returns all of the elements e in the collection such that
+     * key.elementKeysAreEqual(e) and key.hashCode() == e.hashCode().
+     *
+     * @param key       The element to match.
+     *
+     * @return          All of the matching elements.
+     */
+    final public List findAll(E key) {
+        if (key == null || size() == 0) {
+            return Collections.emptyList();
+        }
+        ArrayList results = new ArrayList<>();
+        int slot = slot(elements, key);
+        for (int seen = 0; seen < elements.length; seen++) {
+            Element element = elements[slot];
+            if (element == null) {
+                break;
+            }
+            if (key.elementKeysAreEqual(element)) {
+                @SuppressWarnings("unchecked")
+                E result = (E) elements[slot];
+                results.add(result);
+            }
+            slot = (slot + 1) % elements.length;
+        }
+        return results;
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/utils/Java.java b/clients/src/main/java/org/apache/kafka/common/utils/Java.java
index fb8cafab4b779..c0c0a8970e809 100644
--- a/clients/src/main/java/org/apache/kafka/common/utils/Java.java
+++ b/clients/src/main/java/org/apache/kafka/common/utils/Java.java
@@ -38,7 +38,7 @@ static Version parseVersion(String versionString) {
 
     // Having these as static final provides the best opportunity for compilar optimization
     public static final boolean IS_JAVA9_COMPATIBLE = VERSION.isJava9Compatible();
-    public static final boolean IS_JAVA8_COMPATIBLE = VERSION.isJava8Compatible();
+    public static final boolean IS_JAVA11_COMPATIBLE = VERSION.isJava11Compatible();
 
     public static boolean isIbmJdk() {
         return System.getProperty("java.vendor").contains("IBM");
@@ -65,10 +65,10 @@ boolean isJava9Compatible() {
             return majorVersion >= 9;
         }
 
-        // Package private for testing
-        boolean isJava8Compatible() {
-            return majorVersion > 1 || (majorVersion == 1 && minorVersion >= 8);
+        boolean isJava11Compatible() {
+            return majorVersion >= 11;
         }
+
     }
 
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/utils/KafkaThread.java b/clients/src/main/java/org/apache/kafka/common/utils/KafkaThread.java
index dd634a50c2d26..13430e446a661 100644
--- a/clients/src/main/java/org/apache/kafka/common/utils/KafkaThread.java
+++ b/clients/src/main/java/org/apache/kafka/common/utils/KafkaThread.java
@@ -46,11 +46,7 @@ public KafkaThread(final String name, Runnable runnable, boolean daemon) {
 
     private void configureThread(final String name, boolean daemon) {
         setDaemon(daemon);
-        setUncaughtExceptionHandler(new UncaughtExceptionHandler() {
-            public void uncaughtException(Thread t, Throwable e) {
-                log.error("Uncaught exception in thread '{}':", name, e);
-            }
-        });
+        setUncaughtExceptionHandler((t, e) -> log.error("Uncaught exception in thread '{}':", name, e));
     }
 
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/utils/LogContext.java b/clients/src/main/java/org/apache/kafka/common/utils/LogContext.java
index 8d34ad058ef8e..774961fa418f3 100644
--- a/clients/src/main/java/org/apache/kafka/common/utils/LogContext.java
+++ b/clients/src/main/java/org/apache/kafka/common/utils/LogContext.java
@@ -19,6 +19,9 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.slf4j.Marker;
+import org.slf4j.helpers.FormattingTuple;
+import org.slf4j.helpers.MessageFormatter;
+import org.slf4j.spi.LocationAwareLogger;
 
 /**
  * This class provides a way to instrument loggers with a common context which can be used to
@@ -40,20 +43,404 @@ public LogContext() {
     }
 
     public Logger logger(Class clazz) {
-        return new KafkaLogger(clazz, logPrefix);
+        Logger logger = LoggerFactory.getLogger(clazz);
+        if (logger instanceof LocationAwareLogger) {
+            return new LocationAwareKafkaLogger(logPrefix, (LocationAwareLogger) logger);
+        } else {
+            return new LocationIgnorantKafkaLogger(logPrefix, logger);
+        }
     }
 
     public String logPrefix() {
         return logPrefix;
     }
 
-    private static class KafkaLogger implements Logger {
+    private static abstract class AbstractKafkaLogger implements Logger {
+        private final String prefix;
+
+        protected AbstractKafkaLogger(final String prefix) {
+            this.prefix = prefix;
+        }
+
+        protected String addPrefix(final String message) {
+            return prefix + message;
+        }
+    }
+
+    private static class LocationAwareKafkaLogger extends AbstractKafkaLogger {
+        private final LocationAwareLogger logger;
+        private final String fqcn;
+
+        LocationAwareKafkaLogger(String logPrefix, LocationAwareLogger logger) {
+            super(logPrefix);
+            this.logger = logger;
+            this.fqcn = LocationAwareKafkaLogger.class.getName();
+        }
+
+        @Override
+        public String getName() {
+            return logger.getName();
+        }
+
+        @Override
+        public boolean isTraceEnabled() {
+            return logger.isTraceEnabled();
+        }
+
+        @Override
+        public boolean isTraceEnabled(Marker marker) {
+            return logger.isTraceEnabled(marker);
+        }
+
+        @Override
+        public boolean isDebugEnabled() {
+            return logger.isDebugEnabled();
+        }
+
+        @Override
+        public boolean isDebugEnabled(Marker marker) {
+            return logger.isDebugEnabled(marker);
+        }
+
+        @Override
+        public boolean isInfoEnabled() {
+            return logger.isInfoEnabled();
+        }
+
+        @Override
+        public boolean isInfoEnabled(Marker marker) {
+            return logger.isInfoEnabled(marker);
+        }
+
+        @Override
+        public boolean isWarnEnabled() {
+            return logger.isWarnEnabled();
+        }
+
+        @Override
+        public boolean isWarnEnabled(Marker marker) {
+            return logger.isWarnEnabled(marker);
+        }
+
+        @Override
+        public boolean isErrorEnabled() {
+            return logger.isErrorEnabled();
+        }
+
+        @Override
+        public boolean isErrorEnabled(Marker marker) {
+            return logger.isErrorEnabled(marker);
+        }
+
+        @Override
+        public void trace(String message) {
+            if (logger.isTraceEnabled()) {
+                writeLog(null, LocationAwareLogger.TRACE_INT, message, null, null);
+            }
+        }
+
+        @Override
+        public void trace(String format, Object arg) {
+            if (logger.isTraceEnabled()) {
+                writeLog(null, LocationAwareLogger.TRACE_INT, format, new Object[]{arg}, null);
+            }
+        }
+
+        @Override
+        public void trace(String format, Object arg1, Object arg2) {
+            if (logger.isTraceEnabled()) {
+                writeLog(null, LocationAwareLogger.TRACE_INT, format, new Object[]{arg1, arg2}, null);
+            }
+        }
+
+        @Override
+        public void trace(String format, Object... args) {
+            if (logger.isTraceEnabled()) {
+                writeLog(null, LocationAwareLogger.TRACE_INT, format, args, null);
+            }
+        }
+
+        @Override
+        public void trace(String msg, Throwable t) {
+            if (logger.isTraceEnabled()) {
+                writeLog(null, LocationAwareLogger.TRACE_INT, msg, null, t);
+            }
+        }
+
+        @Override
+        public void trace(Marker marker, String msg) {
+            if (logger.isTraceEnabled()) {
+                writeLog(marker, LocationAwareLogger.TRACE_INT, msg, null, null);
+            }
+        }
+
+        @Override
+        public void trace(Marker marker, String format, Object arg) {
+            if (logger.isTraceEnabled()) {
+                writeLog(marker, LocationAwareLogger.TRACE_INT, format, new Object[]{arg}, null);
+            }
+        }
+
+        @Override
+        public void trace(Marker marker, String format, Object arg1, Object arg2) {
+            if (logger.isTraceEnabled()) {
+                writeLog(marker, LocationAwareLogger.TRACE_INT, format, new Object[]{arg1, arg2}, null);
+            }
+        }
+
+        @Override
+        public void trace(Marker marker, String format, Object... argArray) {
+            if (logger.isTraceEnabled()) {
+                writeLog(marker, LocationAwareLogger.TRACE_INT, format, argArray, null);
+            }
+        }
+
+        @Override
+        public void trace(Marker marker, String msg, Throwable t) {
+            if (logger.isTraceEnabled()) {
+                writeLog(marker, LocationAwareLogger.TRACE_INT, msg, null, t);
+            }
+        }
+
+        @Override
+        public void debug(String message) {
+            if (logger.isDebugEnabled()) {
+                writeLog(null, LocationAwareLogger.DEBUG_INT, message, null, null);
+            }
+        }
+
+        @Override
+        public void debug(String format, Object arg) {
+            if (logger.isDebugEnabled()) {
+                writeLog(null, LocationAwareLogger.DEBUG_INT, format, new Object[]{arg}, null);
+            }
+        }
+
+        @Override
+        public void debug(String format, Object arg1, Object arg2) {
+            if (logger.isDebugEnabled()) {
+                writeLog(null, LocationAwareLogger.DEBUG_INT, format, new Object[]{arg1, arg2}, null);
+            }
+        }
+
+        @Override
+        public void debug(String format, Object... args) {
+            if (logger.isDebugEnabled()) {
+                writeLog(null, LocationAwareLogger.DEBUG_INT, format, args, null);
+            }
+        }
+
+        @Override
+        public void debug(String msg, Throwable t) {
+            if (logger.isDebugEnabled()) {
+                writeLog(null, LocationAwareLogger.DEBUG_INT, msg, null, t);
+            }
+        }
+
+        @Override
+        public void debug(Marker marker, String msg) {
+            if (logger.isDebugEnabled()) {
+                writeLog(marker, LocationAwareLogger.DEBUG_INT, msg, null, null);
+            }
+        }
+
+        @Override
+        public void debug(Marker marker, String format, Object arg) {
+            if (logger.isDebugEnabled()) {
+                writeLog(marker, LocationAwareLogger.DEBUG_INT, format, new Object[]{arg}, null);
+            }
+        }
+
+        @Override
+        public void debug(Marker marker, String format, Object arg1, Object arg2) {
+            if (logger.isDebugEnabled()) {
+                writeLog(marker, LocationAwareLogger.DEBUG_INT, format, new Object[]{arg1, arg2}, null);
+            }
+        }
+
+        @Override
+        public void debug(Marker marker, String format, Object... arguments) {
+            if (logger.isDebugEnabled()) {
+                writeLog(marker, LocationAwareLogger.DEBUG_INT, format, arguments, null);
+            }
+        }
+
+        @Override
+        public void debug(Marker marker, String msg, Throwable t) {
+            if (logger.isDebugEnabled()) {
+                writeLog(marker, LocationAwareLogger.DEBUG_INT, msg, null, t);
+            }
+        }
+
+        @Override
+        public void warn(String message) {
+            writeLog(null, LocationAwareLogger.WARN_INT, message, null, null);
+        }
+
+        @Override
+        public void warn(String format, Object arg) {
+            writeLog(null, LocationAwareLogger.WARN_INT, format, new Object[]{arg}, null);
+        }
+
+        @Override
+        public void warn(String message, Object arg1, Object arg2) {
+            writeLog(null, LocationAwareLogger.WARN_INT, message, new Object[]{arg1, arg2}, null);
+        }
+
+        @Override
+        public void warn(String format, Object... args) {
+            writeLog(null, LocationAwareLogger.WARN_INT, format, args, null);
+        }
+
+        @Override
+        public void warn(String msg, Throwable t) {
+            writeLog(null, LocationAwareLogger.WARN_INT, msg, null, t);
+        }
+
+        @Override
+        public void warn(Marker marker, String msg) {
+            writeLog(marker, LocationAwareLogger.WARN_INT, msg, null, null);
+        }
+
+        @Override
+        public void warn(Marker marker, String format, Object arg) {
+            writeLog(marker, LocationAwareLogger.WARN_INT, format, new Object[]{arg}, null);
+        }
+
+        @Override
+        public void warn(Marker marker, String format, Object arg1, Object arg2) {
+            writeLog(marker, LocationAwareLogger.WARN_INT, format, new Object[]{arg1, arg2}, null);
+        }
+
+        @Override
+        public void warn(Marker marker, String format, Object... arguments) {
+            writeLog(marker, LocationAwareLogger.WARN_INT, format, arguments, null);
+        }
+
+        @Override
+        public void warn(Marker marker, String msg, Throwable t) {
+            writeLog(marker, LocationAwareLogger.WARN_INT, msg, null, t);
+        }
+
+        @Override
+        public void error(String message) {
+            writeLog(null, LocationAwareLogger.ERROR_INT, message, null, null);
+        }
+
+        @Override
+        public void error(String format, Object arg) {
+            writeLog(null, LocationAwareLogger.ERROR_INT, format, new Object[]{arg}, null);
+        }
+
+        @Override
+        public void error(String format, Object arg1, Object arg2) {
+            writeLog(null, LocationAwareLogger.ERROR_INT, format, new Object[]{arg1, arg2}, null);
+        }
+
+        @Override
+        public void error(String format, Object... args) {
+            writeLog(null, LocationAwareLogger.ERROR_INT, format, args, null);
+        }
+
+        @Override
+        public void error(String msg, Throwable t) {
+            writeLog(null, LocationAwareLogger.ERROR_INT, msg, null, t);
+        }
+
+        @Override
+        public void error(Marker marker, String msg) {
+            writeLog(marker, LocationAwareLogger.ERROR_INT, msg, null, null);
+        }
+
+        @Override
+        public void error(Marker marker, String format, Object arg) {
+            writeLog(marker, LocationAwareLogger.ERROR_INT, format, new Object[]{arg}, null);
+        }
+
+        @Override
+        public void error(Marker marker, String format, Object arg1, Object arg2) {
+            writeLog(marker, LocationAwareLogger.ERROR_INT, format, new Object[]{arg1, arg2}, null);
+        }
+
+        @Override
+        public void error(Marker marker, String format, Object... arguments) {
+            writeLog(marker, LocationAwareLogger.ERROR_INT, format, arguments, null);
+        }
+
+        @Override
+        public void error(Marker marker, String msg, Throwable t) {
+            writeLog(marker, LocationAwareLogger.ERROR_INT, msg, null, t);
+        }
+
+        @Override
+        public void info(String msg) {
+            writeLog(null, LocationAwareLogger.INFO_INT, msg, null, null);
+        }
+
+        @Override
+        public void info(String format, Object arg) {
+            writeLog(null, LocationAwareLogger.INFO_INT, format, new Object[]{arg}, null);
+        }
+
+        @Override
+        public void info(String format, Object arg1, Object arg2) {
+            writeLog(null, LocationAwareLogger.INFO_INT, format, new Object[]{arg1, arg2}, null);
+        }
+
+        @Override
+        public void info(String format, Object... args) {
+            writeLog(null, LocationAwareLogger.INFO_INT, format, args, null);
+        }
+
+        @Override
+        public void info(String msg, Throwable t) {
+            writeLog(null, LocationAwareLogger.INFO_INT, msg, null, t);
+        }
+
+        @Override
+        public void info(Marker marker, String msg) {
+            writeLog(marker, LocationAwareLogger.INFO_INT, msg, null, null);
+        }
+
+        @Override
+        public void info(Marker marker, String format, Object arg) {
+            writeLog(marker, LocationAwareLogger.INFO_INT, format, new Object[]{arg}, null);
+        }
+
+        @Override
+        public void info(Marker marker, String format, Object arg1, Object arg2) {
+            writeLog(marker, LocationAwareLogger.INFO_INT, format, new Object[]{arg1, arg2}, null);
+        }
+
+        @Override
+        public void info(Marker marker, String format, Object... arguments) {
+            writeLog(marker, LocationAwareLogger.INFO_INT, format, arguments, null);
+        }
+
+        @Override
+        public void info(Marker marker, String msg, Throwable t) {
+            writeLog(marker, LocationAwareLogger.INFO_INT, msg, null, t);
+        }
+
+        private void writeLog(Marker marker, int level, String format, Object[] args, Throwable exception) {
+            String message = format;
+            if (args != null && args.length > 0) {
+                FormattingTuple formatted = MessageFormatter.arrayFormat(format, args);
+                if (exception == null && formatted.getThrowable() != null) {
+                    exception = formatted.getThrowable();
+                }
+                message = formatted.getMessage();
+            }
+            logger.log(marker, fqcn, level, addPrefix(message), null, exception);
+        }
+    }
+
+    private static class LocationIgnorantKafkaLogger extends AbstractKafkaLogger {
         private final Logger logger;
-        private final String logPrefix;
 
-        public KafkaLogger(Class clazz, String logPrefix) {
-            this.logger = LoggerFactory.getLogger(clazz);
-            this.logPrefix = logPrefix;
+        LocationIgnorantKafkaLogger(String logPrefix, Logger logger) {
+            super(logPrefix);
+            this.logger = logger;
         }
 
         @Override
@@ -113,273 +500,294 @@ public boolean isErrorEnabled(Marker marker) {
 
         @Override
         public void trace(String message) {
-            if (logger.isTraceEnabled())
-                logger.trace(logPrefix + message);
+            if (logger.isTraceEnabled()) {
+                logger.trace(addPrefix(message));
+            }
         }
 
         @Override
         public void trace(String message, Object arg) {
-            if (logger.isTraceEnabled())
-                logger.trace(logPrefix + message, arg);
+            if (logger.isTraceEnabled()) {
+                logger.trace(addPrefix(message), arg);
+            }
         }
 
         @Override
         public void trace(String message, Object arg1, Object arg2) {
-            if (logger.isTraceEnabled())
-                logger.trace(logPrefix + message, arg1, arg2);
+            if (logger.isTraceEnabled()) {
+                logger.trace(addPrefix(message), arg1, arg2);
+            }
         }
 
         @Override
         public void trace(String message, Object... args) {
-            if (logger.isTraceEnabled())
-                logger.trace(logPrefix + message, args);
+            if (logger.isTraceEnabled()) {
+                logger.trace(addPrefix(message), args);
+            }
         }
 
         @Override
         public void trace(String msg, Throwable t) {
-            if (logger.isTraceEnabled())
-                logger.trace(logPrefix + msg, t);
+            if (logger.isTraceEnabled()) {
+                logger.trace(addPrefix(msg), t);
+            }
         }
 
         @Override
         public void trace(Marker marker, String msg) {
-            if (logger.isTraceEnabled())
-                logger.trace(marker, logPrefix + msg);
+            if (logger.isTraceEnabled()) {
+                logger.trace(marker, addPrefix(msg));
+            }
         }
 
         @Override
         public void trace(Marker marker, String format, Object arg) {
-            if (logger.isTraceEnabled())
-                logger.trace(marker, logPrefix + format, arg);
+            if (logger.isTraceEnabled()) {
+                logger.trace(marker, addPrefix(format), arg);
+            }
         }
 
         @Override
         public void trace(Marker marker, String format, Object arg1, Object arg2) {
-            if (logger.isTraceEnabled())
-                logger.trace(marker, logPrefix + format, arg1, arg2);
+            if (logger.isTraceEnabled()) {
+                logger.trace(marker, addPrefix(format), arg1, arg2);
+            }
         }
 
         @Override
         public void trace(Marker marker, String format, Object... argArray) {
-            if (logger.isTraceEnabled())
-                logger.trace(marker, logPrefix + format, argArray);
+            if (logger.isTraceEnabled()) {
+                logger.trace(marker, addPrefix(format), argArray);
+            }
         }
 
         @Override
         public void trace(Marker marker, String msg, Throwable t) {
-            if (logger.isTraceEnabled())
-                logger.trace(marker, logPrefix + msg, t);
+            if (logger.isTraceEnabled()) {
+                logger.trace(marker, addPrefix(msg), t);
+            }
         }
 
         @Override
         public void debug(String message) {
-            if (logger.isDebugEnabled())
-                logger.debug(logPrefix + message);
+            if (logger.isDebugEnabled()) {
+                logger.debug(addPrefix(message));
+            }
         }
 
         @Override
         public void debug(String message, Object arg) {
-            if (logger.isDebugEnabled())
-                logger.debug(logPrefix + message, arg);
+            if (logger.isDebugEnabled()) {
+                logger.debug(addPrefix(message), arg);
+            }
         }
 
         @Override
         public void debug(String message, Object arg1, Object arg2) {
-            if (logger.isDebugEnabled())
-                logger.debug(logPrefix + message, arg1, arg2);
+            if (logger.isDebugEnabled()) {
+                logger.debug(addPrefix(message), arg1, arg2);
+            }
         }
 
         @Override
         public void debug(String message, Object... args) {
-            if (logger.isDebugEnabled())
-                logger.debug(logPrefix + message, args);
+            if (logger.isDebugEnabled()) {
+                logger.debug(addPrefix(message), args);
+            }
         }
 
         @Override
         public void debug(String msg, Throwable t) {
-            if (logger.isDebugEnabled())
-                logger.debug(logPrefix + msg, t);
+            if (logger.isDebugEnabled()) {
+                logger.debug(addPrefix(msg), t);
+            }
         }
 
         @Override
         public void debug(Marker marker, String msg) {
-            if (logger.isDebugEnabled())
-                logger.debug(marker, logPrefix + msg);
+            if (logger.isDebugEnabled()) {
+                logger.debug(marker, addPrefix(msg));
+            }
         }
 
         @Override
         public void debug(Marker marker, String format, Object arg) {
-            if (logger.isDebugEnabled())
-                logger.debug(marker, logPrefix + format, arg);
+            if (logger.isDebugEnabled()) {
+                logger.debug(marker, addPrefix(format), arg);
+            }
         }
 
         @Override
         public void debug(Marker marker, String format, Object arg1, Object arg2) {
-            if (logger.isDebugEnabled())
-                logger.debug(marker, logPrefix + format, arg1, arg2);
+            if (logger.isDebugEnabled()) {
+                logger.debug(marker, addPrefix(format), arg1, arg2);
+            }
         }
 
         @Override
         public void debug(Marker marker, String format, Object... arguments) {
-            if (logger.isDebugEnabled())
-                logger.debug(marker, logPrefix + format, arguments);
+            if (logger.isDebugEnabled()) {
+                logger.debug(marker, addPrefix(format), arguments);
+            }
         }
 
         @Override
         public void debug(Marker marker, String msg, Throwable t) {
-            if (logger.isDebugEnabled())
-                logger.debug(marker, logPrefix + msg, t);
+            if (logger.isDebugEnabled()) {
+                logger.debug(marker, addPrefix(msg), t);
+            }
         }
 
         @Override
         public void warn(String message) {
-            logger.warn(logPrefix + message);
+            logger.warn(addPrefix(message));
         }
 
         @Override
         public void warn(String message, Object arg) {
-            logger.warn(logPrefix + message, arg);
+            logger.warn(addPrefix(message), arg);
         }
 
         @Override
         public void warn(String message, Object arg1, Object arg2) {
-            logger.warn(logPrefix + message, arg1, arg2);
+            logger.warn(addPrefix(message), arg1, arg2);
         }
 
         @Override
         public void warn(String message, Object... args) {
-            logger.warn(logPrefix + message, args);
+            logger.warn(addPrefix(message), args);
         }
 
         @Override
         public void warn(String msg, Throwable t) {
-            logger.warn(logPrefix + msg, t);
+            logger.warn(addPrefix(msg), t);
         }
 
         @Override
         public void warn(Marker marker, String msg) {
-            logger.warn(marker, logPrefix + msg);
+            logger.warn(marker, addPrefix(msg));
         }
 
         @Override
         public void warn(Marker marker, String format, Object arg) {
-            logger.warn(marker, logPrefix + format, arg);
+            logger.warn(marker, addPrefix(format), arg);
         }
 
         @Override
         public void warn(Marker marker, String format, Object arg1, Object arg2) {
-            logger.warn(marker, logPrefix + format, arg1, arg2);
+            logger.warn(marker, addPrefix(format), arg1, arg2);
         }
 
         @Override
         public void warn(Marker marker, String format, Object... arguments) {
-            logger.warn(marker, logPrefix + format, arguments);
+            logger.warn(marker, addPrefix(format), arguments);
         }
 
         @Override
         public void warn(Marker marker, String msg, Throwable t) {
-            logger.warn(marker, logPrefix + msg, t);
+            logger.warn(marker, addPrefix(msg), t);
         }
 
         @Override
         public void error(String message) {
-            logger.error(logPrefix + message);
+            logger.error(addPrefix(message));
         }
 
         @Override
         public void error(String message, Object arg) {
-            logger.error(logPrefix + message, arg);
+            logger.error(addPrefix(message), arg);
         }
 
         @Override
         public void error(String message, Object arg1, Object arg2) {
-            logger.error(logPrefix + message, arg1, arg2);
+            logger.error(addPrefix(message), arg1, arg2);
         }
 
         @Override
         public void error(String message, Object... args) {
-            logger.error(logPrefix + message, args);
+            logger.error(addPrefix(message), args);
         }
 
         @Override
         public void error(String msg, Throwable t) {
-            logger.error(logPrefix + msg, t);
+            logger.error(addPrefix(msg), t);
         }
 
         @Override
         public void error(Marker marker, String msg) {
-            logger.error(marker, logPrefix + msg);
+            logger.error(marker, addPrefix(msg));
         }
 
         @Override
         public void error(Marker marker, String format, Object arg) {
-            logger.error(marker, logPrefix + format, arg);
+            logger.error(marker, addPrefix(format), arg);
         }
 
         @Override
         public void error(Marker marker, String format, Object arg1, Object arg2) {
-            logger.error(marker, logPrefix + format, arg1, arg2);
+            logger.error(marker, addPrefix(format), arg1, arg2);
         }
 
         @Override
         public void error(Marker marker, String format, Object... arguments) {
-            logger.error(marker, logPrefix + format, arguments);
+            logger.error(marker, addPrefix(format), arguments);
         }
 
         @Override
         public void error(Marker marker, String msg, Throwable t) {
-            logger.error(marker, logPrefix + msg, t);
+            logger.error(marker, addPrefix(msg), t);
         }
 
         @Override
         public void info(String message) {
-            logger.info(logPrefix + message);
+            logger.info(addPrefix(message));
         }
 
         @Override
         public void info(String message, Object arg) {
-            logger.info(logPrefix + message, arg);
+            logger.info(addPrefix(message), arg);
         }
 
         @Override
         public void info(String message, Object arg1, Object arg2) {
-            logger.info(logPrefix + message, arg1, arg2);
+            logger.info(addPrefix(message), arg1, arg2);
         }
 
         @Override
         public void info(String message, Object... args) {
-            logger.info(logPrefix + message, args);
+            logger.info(addPrefix(message), args);
         }
 
         @Override
         public void info(String msg, Throwable t) {
-            logger.info(logPrefix + msg, t);
+            logger.info(addPrefix(msg), t);
         }
 
         @Override
         public void info(Marker marker, String msg) {
-            logger.info(marker, logPrefix + msg);
+            logger.info(marker, addPrefix(msg));
         }
 
         @Override
         public void info(Marker marker, String format, Object arg) {
-            logger.info(marker, logPrefix + format, arg);
+            logger.info(marker, addPrefix(format), arg);
         }
 
         @Override
         public void info(Marker marker, String format, Object arg1, Object arg2) {
-            logger.info(marker, logPrefix + format, arg1, arg2);
+            logger.info(marker, addPrefix(format), arg1, arg2);
         }
 
         @Override
         public void info(Marker marker, String format, Object... arguments) {
-            logger.info(marker, logPrefix + format, arguments);
+            logger.info(marker, addPrefix(format), arguments);
         }
 
         @Override
         public void info(Marker marker, String msg, Throwable t) {
-            logger.info(marker, logPrefix + msg, t);
+            logger.info(marker, addPrefix(msg), t);
         }
 
     }
+
 }
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
new file mode 100644
index 0000000000000..112d7fdb3edaa
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/utils/LoggingSignalHandler.java
@@ -0,0 +1,107 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.lang.reflect.Constructor;
+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;
+
+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;
+    private final Method signalHandleMethod;
+    private final Method signalGetNameMethod;
+    private final Method signalHandlerHandleMethod;
+
+    /**
+     * Create an instance of this class.
+     *
+     * @throws ReflectiveOperationException if the underlying API has changed in an incompatible manner.
+     */
+    public LoggingSignalHandler() throws ReflectiveOperationException {
+        signalClass = Class.forName("sun.misc.Signal");
+        signalConstructor = signalClass.getConstructor(String.class);
+        signalHandlerClass = Class.forName("sun.misc.SignalHandler");
+        signalHandlerHandleMethod = signalHandlerClass.getMethod("handle", signalClass);
+        signalHandleMethod = signalClass.getMethod("handle", signalClass, signalHandlerClass);
+        signalGetNameMethod = signalClass.getMethod("getName");
+    }
+
+    /**
+     * Register signal handler to log termination due to SIGTERM, SIGHUP and SIGINT (control-c). This method
+     * does not currently work on Windows.
+     *
+     * @implNote sun.misc.Signal and sun.misc.SignalHandler are described as "not encapsulated" in
+     * http://openjdk.java.net/jeps/260. However, they are not available in the compile classpath if the `--release`
+     * flag is used. As a workaround, we rely on reflection.
+     */
+    public void register() throws ReflectiveOperationException {
+        Map jvmSignalHandlers = new ConcurrentHashMap<>();
+
+        for (String signal : SIGNALS) {
+            register(signal, jvmSignalHandlers);
+        }
+        log.info("Registered signal handlers for " + String.join(", ", SIGNALS));
+    }
+
+    private Object createSignalHandler(final Map jvmSignalHandlers) {
+        InvocationHandler invocationHandler = new InvocationHandler() {
+
+            private String getName(Object signal) throws ReflectiveOperationException {
+                return (String) signalGetNameMethod.invoke(signal);
+            }
+
+            private void handle(Object signalHandler, Object signal) throws ReflectiveOperationException {
+                signalHandlerHandleMethod.invoke(signalHandler, signal);
+            }
+
+            @Override
+            public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
+                Object signal = args[0];
+                log.info("Terminating process due to signal {}", signal);
+                Object handler = jvmSignalHandlers.get(getName(signal));
+                if (handler != null)
+                    handle(handler, signal);
+                return null;
+            }
+        };
+        return Proxy.newProxyInstance(Utils.getContextOrKafkaClassLoader(), new Class[] {signalHandlerClass},
+                invocationHandler);
+    }
+
+    private void register(String signalName, final Map jvmSignalHandlers) throws ReflectiveOperationException {
+        Object signal = signalConstructor.newInstance(signalName);
+        Object signalHandler = createSignalHandler(jvmSignalHandlers);
+        Object oldHandler = signalHandleMethod.invoke(null, signal, signalHandler);
+        if (oldHandler != null)
+            jvmSignalHandlers.put(signalName, oldHandler);
+    }
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/utils/MappedByteBuffers.java b/clients/src/main/java/org/apache/kafka/common/utils/MappedByteBuffers.java
deleted file mode 100644
index 1faecb3eacd66..0000000000000
--- a/clients/src/main/java/org/apache/kafka/common/utils/MappedByteBuffers.java
+++ /dev/null
@@ -1,136 +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.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.io.IOException;
-import java.lang.invoke.MethodHandle;
-import java.lang.invoke.MethodHandles;
-import java.lang.reflect.Field;
-import java.lang.reflect.Method;
-import java.nio.ByteBuffer;
-import java.nio.MappedByteBuffer;
-
-import static java.lang.invoke.MethodHandles.constant;
-import static java.lang.invoke.MethodHandles.dropArguments;
-import static java.lang.invoke.MethodHandles.filterReturnValue;
-import static java.lang.invoke.MethodHandles.guardWithTest;
-import static java.lang.invoke.MethodHandles.lookup;
-import static java.lang.invoke.MethodType.methodType;
-
-/**
- * Utility methods for MappedByteBuffer implementations.
- *
- * The unmap implementation was inspired by the one in Lucene's MMapDirectory.
- */
-public final class MappedByteBuffers {
-
-    private static final Logger log = LoggerFactory.getLogger(MappedByteBuffers.class);
-
-    // null if unmap is not supported
-    private static final MethodHandle UNMAP;
-
-    // null if unmap is supported
-    private static final RuntimeException UNMAP_NOT_SUPPORTED_EXCEPTION;
-
-    static {
-        Object unmap = null;
-        RuntimeException exception = null;
-        try {
-            unmap = lookupUnmapMethodHandle();
-        } catch (RuntimeException e) {
-            exception = e;
-        }
-        if (unmap != null) {
-            UNMAP = (MethodHandle) unmap;
-            UNMAP_NOT_SUPPORTED_EXCEPTION = null;
-        } else {
-            UNMAP = null;
-            UNMAP_NOT_SUPPORTED_EXCEPTION = exception;
-        }
-    }
-
-    private MappedByteBuffers() {}
-
-    public static void unmap(String resourceDescription, MappedByteBuffer buffer) throws IOException {
-        if (!buffer.isDirect())
-            throw new IllegalArgumentException("Unmapping only works with direct buffers");
-        if (UNMAP == null)
-            throw UNMAP_NOT_SUPPORTED_EXCEPTION;
-
-        try {
-            UNMAP.invokeExact((ByteBuffer) buffer);
-        } catch (Throwable throwable) {
-            throw new IOException("Unable to unmap the mapped buffer: " + resourceDescription, throwable);
-        }
-    }
-
-    private static MethodHandle lookupUnmapMethodHandle() {
-        final MethodHandles.Lookup lookup = lookup();
-        try {
-            if (Java.IS_JAVA9_COMPATIBLE)
-                return unmapJava9(lookup);
-            else
-                return unmapJava7Or8(lookup);
-        } catch (ReflectiveOperationException | RuntimeException e1) {
-            throw new UnsupportedOperationException("Unmapping is not supported on this platform, because internal " +
-                "Java APIs are not compatible with this Kafka version", e1);
-        }
-    }
-
-    private static MethodHandle unmapJava7Or8(MethodHandles.Lookup lookup) throws ReflectiveOperationException {
-        /* "Compile" a MethodHandle that is roughly equivalent to the following lambda:
-         *
-         * (ByteBuffer buffer) -> {
-         *   sun.misc.Cleaner cleaner = ((java.nio.DirectByteBuffer) byteBuffer).cleaner();
-         *   if (nonNull(cleaner))
-         *     cleaner.clean();
-         *   else
-         *     noop(cleaner); // the noop is needed because MethodHandles#guardWithTest always needs both if and else
-         * }
-         */
-        Class directBufferClass = Class.forName("java.nio.DirectByteBuffer");
-        Method m = directBufferClass.getMethod("cleaner");
-        m.setAccessible(true);
-        MethodHandle directBufferCleanerMethod = lookup.unreflect(m);
-        Class cleanerClass = directBufferCleanerMethod.type().returnType();
-        MethodHandle cleanMethod = lookup.findVirtual(cleanerClass, "clean", methodType(void.class));
-        MethodHandle nonNullTest = lookup.findStatic(MappedByteBuffers.class, "nonNull",
-                methodType(boolean.class, Object.class)).asType(methodType(boolean.class, cleanerClass));
-        MethodHandle noop = dropArguments(constant(Void.class, null).asType(methodType(void.class)), 0, cleanerClass);
-        MethodHandle unmapper = filterReturnValue(directBufferCleanerMethod, guardWithTest(nonNullTest, cleanMethod, noop))
-                .asType(methodType(void.class, ByteBuffer.class));
-        return unmapper;
-    }
-
-    private static MethodHandle unmapJava9(MethodHandles.Lookup lookup) throws ReflectiveOperationException {
-        Class unsafeClass = Class.forName("sun.misc.Unsafe");
-        MethodHandle unmapper = lookup.findVirtual(unsafeClass, "invokeCleaner",
-                methodType(void.class, ByteBuffer.class));
-        Field f = unsafeClass.getDeclaredField("theUnsafe");
-        f.setAccessible(true);
-        Object theUnsafe = f.get(null);
-        return unmapper.bindTo(theUnsafe);
-    }
-
-    private static boolean nonNull(Object o) {
-        return o != null;
-    }
-}
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/OperatingSystem.java b/clients/src/main/java/org/apache/kafka/common/utils/OperatingSystem.java
index 9d80295b1237f..8dc8b86233f2e 100644
--- a/clients/src/main/java/org/apache/kafka/common/utils/OperatingSystem.java
+++ b/clients/src/main/java/org/apache/kafka/common/utils/OperatingSystem.java
@@ -27,8 +27,11 @@ private OperatingSystem() {
 
     public static final boolean IS_WINDOWS;
 
+    public static final boolean IS_ZOS;
+
     static {
         NAME = System.getProperty("os.name").toLowerCase(Locale.ROOT);
         IS_WINDOWS = NAME.startsWith("windows");
+        IS_ZOS = NAME.startsWith("z/os");
     }
 }
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/ProducerIdAndEpoch.java b/clients/src/main/java/org/apache/kafka/common/utils/ProducerIdAndEpoch.java
new file mode 100644
index 0000000000000..674b42352b787
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/utils/ProducerIdAndEpoch.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.common.utils;
+
+import org.apache.kafka.common.record.RecordBatch;
+
+public class ProducerIdAndEpoch {
+    public static final ProducerIdAndEpoch NONE = new ProducerIdAndEpoch(RecordBatch.NO_PRODUCER_ID, RecordBatch.NO_PRODUCER_EPOCH);
+
+    public final long producerId;
+    public final short epoch;
+
+    public ProducerIdAndEpoch(long producerId, short epoch) {
+        this.producerId = producerId;
+        this.epoch = epoch;
+    }
+
+    public boolean isValid() {
+        return RecordBatch.NO_PRODUCER_ID < producerId;
+    }
+
+    @Override
+    public String toString() {
+        return "(producerId=" + producerId + ", epoch=" + epoch + ")";
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+
+        ProducerIdAndEpoch that = (ProducerIdAndEpoch) o;
+
+        if (producerId != that.producerId) return false;
+        return epoch == that.epoch;
+    }
+
+    @Override
+    public int hashCode() {
+        int result = (int) (producerId ^ (producerId >>> 32));
+        result = 31 * result + (int) epoch;
+        return result;
+    }
+
+}
diff --git a/clients/src/main/java/org/apache/kafka/common/utils/PureJavaCrc32C.java b/clients/src/main/java/org/apache/kafka/common/utils/PureJavaCrc32C.java
index 3489dddec7102..8abc93dc3b486 100644
--- a/clients/src/main/java/org/apache/kafka/common/utils/PureJavaCrc32C.java
+++ b/clients/src/main/java/org/apache/kafka/common/utils/PureJavaCrc32C.java
@@ -53,6 +53,7 @@ public void reset() {
         crc = 0xffffffff;
     }
 
+    @SuppressWarnings("fallthrough")
     @Override
     public void update(byte[] b, int off, int len) {
         int localCrc = crc;
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..12defdc7173fa 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,50 @@
  */
 package org.apache.kafka.common.utils;
 
+import org.apache.kafka.common.acl.AclOperation;
+import org.apache.kafka.common.acl.AclPermissionType;
+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.Locale;
+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;
+    private static final Map NAME_TO_PERMISSION_TYPES;
+
+    static {
+        NAME_TO_RESOURCE_TYPES = new HashMap<>(ResourceType.values().length);
+        NAME_TO_OPERATIONS = new HashMap<>(AclOperation.values().length);
+        NAME_TO_PERMISSION_TYPES = new HashMap<>(AclPermissionType.values().length);
+
+        for (ResourceType resourceType : ResourceType.values()) {
+            String resourceTypeName = toPascalCase(resourceType.name());
+            NAME_TO_RESOURCE_TYPES.put(resourceTypeName, resourceType);
+            NAME_TO_RESOURCE_TYPES.put(resourceTypeName.toUpperCase(Locale.ROOT), resourceType);
+        }
+        for (AclOperation operation : AclOperation.values()) {
+            String operationName = toPascalCase(operation.name());
+            NAME_TO_OPERATIONS.put(operationName, operation);
+            NAME_TO_OPERATIONS.put(operationName.toUpperCase(Locale.ROOT), operation);
+        }
+        for (AclPermissionType permissionType : AclPermissionType.values()) {
+            String permissionName  = toPascalCase(permissionType.name());
+            NAME_TO_PERMISSION_TYPES.put(permissionName, permissionType);
+            NAME_TO_PERMISSION_TYPES.put(permissionName.toUpperCase(Locale.ROOT), permissionType);
+        }
+    }
+
     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 +74,76 @@ 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]).getConstructor().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 (ReflectiveOperationException e) {
+            LOGGER.error("Unexpected implementation of security provider creator class", e);
+        }
+    }
+
+    public static ResourceType resourceType(String name) {
+        return valueFromMap(NAME_TO_RESOURCE_TYPES, name, ResourceType.UNKNOWN);
+    }
+
+    public static AclOperation operation(String name) {
+        return valueFromMap(NAME_TO_OPERATIONS, name, AclOperation.UNKNOWN);
+    }
+
+    public static AclPermissionType permissionType(String name) {
+        return valueFromMap(NAME_TO_PERMISSION_TYPES, name, AclPermissionType.UNKNOWN);
+    }
+
+    // We use Pascal-case to store these values, so lookup using provided key first to avoid
+    // case conversion for the common case. For backward compatibility, also perform
+    // case-insensitive look up (without underscores) by converting the key to upper-case.
+    private static  T valueFromMap(Map map, String key, T unknown) {
+        T value = map.get(key);
+        if (value == null) {
+            value = map.get(key.toUpperCase(Locale.ROOT));
+        }
+        return value == null ? unknown : value;
+    }
+
+    public static String resourceTypeName(ResourceType resourceType) {
+        return toPascalCase(resourceType.name());
+    }
+
+    public static String operationName(AclOperation operation) {
+        return toPascalCase(operation.name());
+    }
+
+    public static String permissionTypeName(AclPermissionType permissionType) {
+        return toPascalCase(permissionType.name());
+    }
+
+    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/Shell.java b/clients/src/main/java/org/apache/kafka/common/utils/Shell.java
index ebfd0bacc5c42..a9b93ec8aa1fa 100644
--- a/clients/src/main/java/org/apache/kafka/common/utils/Shell.java
+++ b/clients/src/main/java/org/apache/kafka/common/utils/Shell.java
@@ -251,7 +251,7 @@ public String toString() {
      * @param cmd shell command to execute.
      * @return the output of the executed command.
      */
-    public static String execCommand(String ... cmd) throws IOException {
+    public static String execCommand(String... cmd) throws IOException {
         return execCommand(cmd, -1);
     }
 
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 60da06444d220..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
@@ -16,7 +16,9 @@
  */
 package org.apache.kafka.common.utils;
 
-import java.util.concurrent.TimeUnit;
+import org.apache.kafka.common.errors.TimeoutException;
+
+import java.util.function.Supplier;
 
 /**
  * A time implementation that uses the system clock and sleep call. Use `Time.SYSTEM` instead of creating an instance
@@ -29,11 +31,6 @@ public long milliseconds() {
         return System.currentTimeMillis();
     }
 
-    @Override
-    public long hiResClockMs() {
-        return TimeUnit.NANOSECONDS.toMillis(nanoseconds());
-    }
-
     @Override
     public long nanoseconds() {
         return System.nanoTime();
@@ -41,11 +38,22 @@ 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
+    public void waitObject(Object obj, Supplier condition, long deadlineMs) throws InterruptedException {
+        synchronized (obj) {
+            while (true) {
+                if (condition.get())
+                    return;
+
+                long currentTimeMs = milliseconds();
+                if (currentTimeMs >= deadlineMs)
+                    throw new TimeoutException("Condition not satisfied before deadline");
+
+                obj.wait(deadlineMs - currentTimeMs);
+            }
         }
     }
 
diff --git a/tools/src/main/java/org/apache/kafka/trogdor/common/ThreadUtils.java b/clients/src/main/java/org/apache/kafka/common/utils/ThreadUtils.java
similarity index 98%
rename from tools/src/main/java/org/apache/kafka/trogdor/common/ThreadUtils.java
rename to clients/src/main/java/org/apache/kafka/common/utils/ThreadUtils.java
index 7e0285648ad34..750c8d7a89f41 100644
--- a/tools/src/main/java/org/apache/kafka/trogdor/common/ThreadUtils.java
+++ b/clients/src/main/java/org/apache/kafka/common/utils/ThreadUtils.java
@@ -15,7 +15,7 @@
  * limitations under the License.
  */
 
-package org.apache.kafka.trogdor.common;
+package org.apache.kafka.common.utils;
 
 import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.atomic.AtomicLong;
diff --git a/clients/src/main/java/org/apache/kafka/common/utils/Time.java b/clients/src/main/java/org/apache/kafka/common/utils/Time.java
index c288bd33e54f1..9e0a475da5cf2 100644
--- a/clients/src/main/java/org/apache/kafka/common/utils/Time.java
+++ b/clients/src/main/java/org/apache/kafka/common/utils/Time.java
@@ -16,6 +16,10 @@
  */
 package org.apache.kafka.common.utils;
 
+import java.time.Duration;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Supplier;
+
 /**
  * An interface abstracting the clock to use in unit testing classes that make use of clock time.
  *
@@ -33,7 +37,9 @@ public interface Time {
     /**
      * Returns the value returned by `nanoseconds` converted into milliseconds.
      */
-    long hiResClockMs();
+    default long hiResClockMs() {
+        return TimeUnit.NANOSECONDS.toMillis(nanoseconds());
+    }
 
     /**
      * Returns the current value of the running JVM's high-resolution time source, in nanoseconds.
@@ -53,4 +59,31 @@ public interface Time {
      */
     void sleep(long ms);
 
+    /**
+     * Wait for a condition using the monitor of a given object. This avoids the implicit
+     * dependence on system time when calling {@link Object#wait()}.
+     *
+     * @param obj The object that will be waited with {@link Object#wait()}. Note that it is the responsibility
+     *      of the caller to call notify on this object when the condition is satisfied.
+     * @param condition The condition we are awaiting
+     * @param deadlineMs The deadline timestamp at which to raise a timeout error
+     *
+     * @throws org.apache.kafka.common.errors.TimeoutException if the timeout expires before the condition is satisfied
+     */
+    void waitObject(Object obj, Supplier condition, long deadlineMs) throws InterruptedException;
+
+    /**
+     * Get a timer which is bound to this time instance and expires after the given timeout
+     */
+    default Timer timer(long timeoutMs) {
+        return new Timer(this, timeoutMs);
+    }
+
+    /**
+     * Get a timer which is bound to this time instance and expires after the given timeout
+     */
+    default Timer timer(Duration timeout) {
+        return timer(timeout.toMillis());
+    }
+
 }
diff --git a/clients/src/main/java/org/apache/kafka/common/utils/Timer.java b/clients/src/main/java/org/apache/kafka/common/utils/Timer.java
new file mode 100644
index 0000000000000..98b09a38d3636
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/common/utils/Timer.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.common.utils;
+
+/**
+ * This is a helper class which makes blocking methods with a timeout easier to implement.
+ * In particular it enables use cases where a high-level blocking call with a timeout is
+ * composed of several lower level calls, each of which has their own respective timeouts. The idea
+ * is to create a single timer object for the high level timeout and carry it along to
+ * all of the lower level methods. This class also handles common problems such as integer overflow.
+ * This class also ensures monotonic updates to the timer even if the underlying clock is subject
+ * to non-monotonic behavior. For example, the remaining time returned by {@link #remainingMs()} is
+ * guaranteed to decrease monotonically until it hits zero.
+ *
+ * Note that it is up to the caller to ensure progress of the timer using one of the
+ * {@link #update()} methods or {@link #sleep(long)}. The timer will cache the current time and
+ * return it indefinitely until the timer has been updated. This allows the caller to limit
+ * unnecessary system calls and update the timer only when needed. For example, a timer which is
+ * waiting a request sent through the {@link org.apache.kafka.clients.NetworkClient} should call
+ * {@link #update()} following each blocking call to
+ * {@link org.apache.kafka.clients.NetworkClient#poll(long, long)}.
+ *
+ * A typical usage might look something like this:
+ *
+ * 
            + *     Time time = Time.SYSTEM;
+ *     Timer timer = time.timer(500);
+ *
+ *     while (!conditionSatisfied() && timer.notExpired) {
+ *         client.poll(timer.remainingMs(), timer.currentTimeMs());
+ *         timer.update();
+ *     }
+ * 
            
+ */
+public class Timer {
+    private final Time time;
+    private long startMs;
+    private long currentTimeMs;
+    private long deadlineMs;
+    private long timeoutMs;
+
+    Timer(Time time, long timeoutMs) {
+        this.time = time;
+        update();
+        reset(timeoutMs);
+    }
+
+    /**
+     * Check timer expiration. Like {@link #remainingMs()}, this depends on the current cached
+     * time in milliseconds, which is only updated through one of the {@link #update()} methods
+     * or with {@link #sleep(long)};
+     *
+     * @return true if the timer has expired, false otherwise
+     */
+    public boolean isExpired() {
+        return currentTimeMs >= deadlineMs;
+    }
+
+    /**
+     * Check whether the timer has not yet expired.
+     * @return true if there is still time remaining before expiration
+     */
+    public boolean notExpired() {
+        return !isExpired();
+    }
+
+    /**
+     * Reset the timer to the specific timeout. This will use the underlying {@link #Timer(Time, long)}
+     * implementation to update the current cached time in milliseconds and it will set a new timer
+     * deadline.
+     *
+     * @param timeoutMs The new timeout in milliseconds
+     */
+    public void updateAndReset(long timeoutMs) {
+        update();
+        reset(timeoutMs);
+    }
+
+    /**
+     * Reset the timer using a new timeout. Note that this does not update the cached current time
+     * in milliseconds, so it typically must be accompanied with a separate call to {@link #update()}.
+     * Typically, you can just use {@link #updateAndReset(long)}.
+     *
+     * @param timeoutMs The new timeout in milliseconds
+     */
+    public void reset(long timeoutMs) {
+        if (timeoutMs < 0)
+            throw new IllegalArgumentException("Invalid negative timeout " + timeoutMs);
+
+        this.timeoutMs = timeoutMs;
+        this.startMs = this.currentTimeMs;
+
+        if (currentTimeMs > Long.MAX_VALUE - timeoutMs)
+            this.deadlineMs = Long.MAX_VALUE;
+        else
+            this.deadlineMs = currentTimeMs + timeoutMs;
+    }
+
+    /**
+     * Reset the timer's deadline directly.
+     *
+     * @param deadlineMs The new deadline in milliseconds
+     */
+    public void resetDeadline(long deadlineMs) {
+        if (deadlineMs < 0)
+            throw new IllegalArgumentException("Invalid negative deadline " + deadlineMs);
+
+        this.timeoutMs = Math.max(0, deadlineMs - this.currentTimeMs);
+        this.startMs = this.currentTimeMs;
+        this.deadlineMs = deadlineMs;
+    }
+
+    /**
+     * Use the underlying {@link Time} implementation to update the current cached time. If
+     * the underlying time returns a value which is smaller than the current cached time,
+     * the update will be ignored.
+     */
+    public void update() {
+        update(time.milliseconds());
+    }
+
+    /**
+     * Update the cached current time to a specific value. In some contexts, the caller may already
+     * have an accurate time, so this avoids unnecessary calls to system time.
+     *
+     * Note that if the updated current time is smaller than the cached time, then the update
+     * is ignored.
+     *
+     * @param currentTimeMs The current time in milliseconds to cache
+     */
+    public void update(long currentTimeMs) {
+        this.currentTimeMs = Math.max(currentTimeMs, this.currentTimeMs);
+    }
+
+    /**
+     * Get the remaining time in milliseconds until the timer expires. Like {@link #currentTimeMs},
+     * this depends on the cached current time, so the returned value will not change until the timer
+     * has been updated using one of the {@link #update()} methods or {@link #sleep(long)}.
+     *
+     * @return The cached remaining time in milliseconds until timer expiration
+     */
+    public long remainingMs() {
+        return Math.max(0, deadlineMs - currentTimeMs);
+    }
+
+    /**
+     * Get the current time in milliseconds. This will return the same cached value until the timer
+     * has been updated using one of the {@link #update()} methods or {@link #sleep(long)} is used.
+     *
+     * Note that the value returned is guaranteed to increase monotonically even if the underlying
+     * {@link Time} implementation goes backwards. Effectively, the timer will just wait for the
+     * time to catch up.
+     *
+     * @return The current cached time in milliseconds
+     */
+    public long currentTimeMs() {
+        return currentTimeMs;
+    }
+
+    /**
+     * Get the amount of time that has elapsed since the timer began. If the timer was reset, this
+     * will be the amount of time since the last reset.
+     *
+     * @return The elapsed time since construction or the last reset
+     */
+    public long elapsedMs() {
+        return currentTimeMs - startMs;
+    }
+
+    /**
+     * Get the current timeout value specified through {@link #reset(long)} or {@link #resetDeadline(long)}.
+     * This value is constant until altered by one of these API calls.
+     *
+     * @return The timeout in milliseconds
+     */
+    public long timeoutMs() {
+        return timeoutMs;
+    }
+
+    /**
+     * Sleep for the requested duration and update the timer. Return when either the duration has
+     * elapsed or the timer has expired.
+     *
+     * @param durationMs The duration in milliseconds to sleep
+     */
+    public void sleep(long durationMs) {
+        long sleepDurationMs = Math.min(durationMs, remainingMs());
+        time.sleep(sleepDurationMs);
+        update();
+    }
+}
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 2843f3ef82c83..34aed9d98bf0d 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
@@ -16,7 +16,12 @@
  */
 package org.apache.kafka.common.utils;
 
+import java.nio.BufferUnderflowException;
+import java.util.EnumSet;
+import java.util.SortedSet;
+import java.util.TreeSet;
 import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.config.ConfigException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -24,8 +29,6 @@
 import java.io.DataOutput;
 import java.io.EOFException;
 import java.io.File;
-import java.io.FileInputStream;
-import java.io.FileNotFoundException;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.PrintWriter;
@@ -33,18 +36,18 @@
 import java.lang.reflect.Constructor;
 import java.lang.reflect.InvocationTargetException;
 import java.nio.ByteBuffer;
-import java.nio.MappedByteBuffer;
 import java.nio.channels.FileChannel;
-import java.nio.charset.Charset;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.FileVisitResult;
 import java.nio.file.Files;
 import java.nio.file.NoSuchFileException;
 import java.nio.file.Path;
+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;
@@ -52,21 +55,40 @@
 import java.util.HashMap;
 import java.util.HashSet;
 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.BiConsumer;
+import java.util.function.BinaryOperator;
+import java.util.function.Function;
+import java.util.function.Supplier;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
+import java.util.stream.Collector;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+import java.text.ParseException;
+import java.text.SimpleDateFormat;
+import java.util.Date;
 
-public class Utils {
+public final class Utils {
+
+    private Utils() {}
 
     // This matches URIs of formats: host:port and protocol:\\host:port
     // IPv6 is supported with [ip] pattern
     private static final Pattern HOST_PORT_PATTERN = Pattern.compile(".*?\\[?([0-9a-zA-Z\\-%._:]*)\\]?:([0-9]+)");
 
+    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"};
 
@@ -108,6 +130,17 @@ public static String utf8(ByteBuffer buffer, int length) {
         return utf8(buffer, 0, length);
     }
 
+    /**
+     * Read a UTF8 string from the current position till the end of a byte buffer. The position of the byte buffer is
+     * not affected by this method.
+     *
+     * @param buffer The buffer to read from
+     * @return The UTF8 string
+     */
+    public static String utf8(ByteBuffer buffer) {
+        return utf8(buffer, buffer.remaining());
+    }
+
     /**
      * Read a UTF8 string from a byte buffer at a given offset. Note that the position of the byte buffer
      * is not affected by this method.
@@ -148,7 +181,7 @@ public static int abs(int n) {
      * @param rest The remaining values to compare
      * @return The minimum of all passed values
      */
-    public static long min(long first, long ... rest) {
+    public static long min(long first, long... rest) {
         long min = first;
         for (long r : rest) {
             if (r < min)
@@ -163,7 +196,7 @@ public static long min(long first, long ... rest) {
      * @param rest The remaining values to compare
      * @return The maximum of all passed values
      */
-    public static long max(long first, long ... rest) {
+    public static long max(long first, long... rest) {
         long max = first;
         for (long r : rest) {
             if (r > max)
@@ -256,17 +289,43 @@ public static byte[] toArray(ByteBuffer buffer, int offset, int size) {
     }
 
     /**
-     * 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.
+     * Starting from the current position, read an integer indicating the size of the byte array to read,
+     * then read the array. Consumes the buffer: upon returning, the buffer's position is after the array
+     * that is returned.
+     * @param buffer The buffer to read a size-prefixed array from
+     * @return The array
      */
-    public static  T notNull(T t) {
-        if (t == null)
-            throw new NullPointerException();
-        else
-            return t;
+    public static byte[] getNullableSizePrefixedArray(final ByteBuffer buffer) {
+        final int size = buffer.getInt();
+        return getNullableArray(buffer, size);
+    }
+
+    /**
+     * Read a byte array of the given size. Consumes the buffer: upon returning, the buffer's position
+     * is after the array that is returned.
+     * @param buffer The buffer to read a size-prefixed array from
+     * @param size The number of bytes to read out of the buffer
+     * @return The array
+     */
+    public static byte[] getNullableArray(final ByteBuffer buffer, final int size) {
+        if (size > buffer.remaining()) {
+            // preemptively throw this when the read is doomed to fail, so we don't have to allocate the array.
+            throw new BufferUnderflowException();
+        }
+        final byte[] oldBytes = size == -1 ? null : new byte[size];
+        if (oldBytes != null) {
+            buffer.get(oldBytes);
+        }
+        return oldBytes;
+    }
+
+    /**
+     * Returns a copy of src byte array
+     * @param src The byte array to copy
+     * @return The copy
+     */
+    public static byte[] copyArray(byte[] src) {
+        return Arrays.copyOf(src, src.length);
     }
 
     /**
@@ -301,11 +360,34 @@ public static  T newInstance(Class c) {
      * Look up the class by name and instantiate it.
      * @param klass class name
      * @param base super class of the class to be instantiated
-     * @param 
+     * @param  the type of the base class
      * @return the new instance
      */
     public static  T newInstance(String klass, Class base) throws ClassNotFoundException {
-        return Utils.newInstance(Class.forName(klass, true, Utils.getContextOrKafkaClassLoader()).asSubclass(base));
+        return Utils.newInstance(loadClass(klass, base));
+    }
+
+    /**
+     * Look up a class by name.
+     * @param klass class name
+     * @param base super class of the class for verification
+     * @param  the type of the base class
+     * @return the new class
+     */
+    public static  Class loadClass(String klass, Class base) throws ClassNotFoundException {
+        return Class.forName(klass, true, Utils.getContextOrKafkaClassLoader()).asSubclass(base);
+    }
+
+    /**
+     * Cast {@code klass} to {@code base} and instantiate it.
+     * @param klass The class to instantiate
+     * @param base A know baseclass of klass.
+     * @param  the type of the base class
+     * @throws ClassCastException If {@code klass} is not a subclass of {@code base}.
+     * @return the new instance.
+     */
+    public static  T newInstance(Class klass, Class base) {
+        return Utils.newInstance(klass.asSubclass(base));
     }
 
     /**
@@ -350,6 +432,7 @@ public static  T newParameterizedInstance(String className, Object... params)
      * @param data byte array to hash
      * @return 32 bit hash of the given array
      */
+    @SuppressWarnings("fallthrough")
     public static int murmur2(final byte[] data) {
         int length = data.length;
         int seed = 0x9747b28c;
@@ -410,6 +493,15 @@ public static Integer getPort(String address) {
         return matcher.matches() ? Integer.parseInt(matcher.group(2)) : null;
     }
 
+    /**
+     * Basic validation of the supplied address. checks for valid characters
+     * @param address hostname string to validate
+     * @return true if address contains valid characters
+     */
+    public static boolean validHostPattern(String address) {
+        return VALID_HOST_CHARACTERS.matcher(address).matches();
+    }
+
     /**
      * Formats hostname and port number as a "host:port" address string,
      * surrounding IPv6 addresses with braces '[', ']'
@@ -430,7 +522,7 @@ public static String formatAddress(String host, Integer port) {
      */
     public static String formatBytes(long bytes) {
         if (bytes < 0) {
-            return "" + bytes;
+            return String.valueOf(bytes);
         }
         double asDouble = (double) bytes;
         int ordinal = (int) Math.floor(Math.log(asDouble) / Math.log(1024.0));
@@ -441,7 +533,7 @@ public static String formatBytes(long bytes) {
             return formatted + " " + BYTE_SCALE_SUFFIXES[ordinal];
         } catch (IndexOutOfBoundsException e) {
             //huge number?
-            return "" + asDouble;
+            return String.valueOf(asDouble);
         }
     }
 
@@ -462,6 +554,7 @@ public static  String join(T[] strs, String separator) {
      * @return The string representation.
      */
     public static  String join(Collection list, String separator) {
+        Objects.requireNonNull(list);
         StringBuilder sb = new StringBuilder();
         Iterator iter = list.iterator();
         while (iter.hasNext()) {
@@ -472,6 +565,12 @@ public static  String join(Collection list, String separator) {
         return sb.toString();
     }
 
+    /**
+     *  Converts a {@code Map} class into a string, concatenating keys and values
+     *  Example:
+     *      {@code mkString({ key: "hello", keyTwo: "hi" }, "|START|", "|END|", "=", ",")
+     *          => "|START|key=hello,keyTwo=hi|END|"}
+     */
     public static  String mkString(Map map, String begin, String end,
                                          String keyValueSeparator, String elementSeparator) {
         StringBuilder bld = new StringBuilder();
@@ -486,16 +585,61 @@ public static  String mkString(Map map, String begin, String end,
         return bld.toString();
     }
 
+    /**
+     *  Converts an extensions string into a {@code Map}.
+     *
+     *  Example:
+     *      {@code parseMap("key=hey,keyTwo=hi,keyThree=hello", "=", ",") => { key: "hey", keyTwo: "hi", keyThree: "hello" }}
+     *
+     */
+    public static Map parseMap(String mapStr, String keyValueSeparator, String elementSeparator) {
+        Map map = new HashMap<>();
+
+        if (!mapStr.isEmpty()) {
+            String[] attrvals = mapStr.split(elementSeparator);
+            for (String attrval : attrvals) {
+                String[] array = attrval.split(keyValueSeparator, 2);
+                map.put(array[0], array[1]);
+            }
+        }
+        return map;
+    }
+
+    /**
+     * Read a properties file from the given path
+     * @param filename The path of the file to read
+     * @return the loaded properties
+     */
+    public static Properties loadProps(String filename) throws IOException {
+        return loadProps(filename, null);
+    }
+
     /**
      * Read a properties file from the given path
      * @param filename The path of the file to read
+     * @param onlyIncludeKeys When non-null, only return values associated with these keys and ignore all others
+     * @return the loaded properties
      */
-    public static Properties loadProps(String filename) throws IOException, FileNotFoundException {
+    public static Properties loadProps(String filename, List onlyIncludeKeys) throws IOException {
         Properties props = new Properties();
-        try (InputStream propStream = new FileInputStream(filename)) {
-            props.load(propStream);
+
+        if (filename != null) {
+            try (InputStream propStream = Files.newInputStream(Paths.get(filename))) {
+                props.load(propStream);
+            }
+        } else {
+            System.out.println("Did not load any properties since the property file is not specified");
         }
-        return props;
+
+        if (onlyIncludeKeys == null || onlyIncludeKeys.isEmpty())
+            return props;
+        Properties requestedProps = new Properties();
+        onlyIncludeKeys.forEach(key -> {
+            String value = props.getProperty(key);
+            if (value != null)
+                requestedProps.setProperty(key, value);
+        });
+        return requestedProps;
     }
 
     /**
@@ -519,15 +663,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
      */
@@ -538,7 +673,7 @@ public static byte[] readBytes(ByteBuffer buffer, int offset, int length) {
         } else {
             buffer.mark();
             buffer.position(offset);
-            buffer.get(dest, 0, length);
+            buffer.get(dest);
             buffer.reset();
         }
         return dest;
@@ -552,22 +687,16 @@ public static byte[] readBytes(ByteBuffer buffer) {
     }
 
     /**
-     * Attempt to read a file as a string
-     * @throws IOException
+     * Read a file as string and return the content. The file is treated as a stream and no seek is performed.
+     * This allows the program to read from a regular file as well as from a pipe/fifo.
      */
-    public static String readFileAsString(String path, Charset charset) throws IOException {
-        if (charset == null) charset = Charset.defaultCharset();
-
-        try (FileInputStream stream = new FileInputStream(new File(path))) {
-            FileChannel fc = stream.getChannel();
-            MappedByteBuffer bb = fc.map(FileChannel.MapMode.READ_ONLY, 0, fc.size());
-            return charset.decode(bb).toString();
-        }
-
-    }
-
     public static String readFileAsString(String path) throws IOException {
-        return Utils.readFileAsString(path, Charset.defaultCharset());
+        try {
+            byte[] allBytes = Files.readAllBytes(Paths.get(path));
+            return new String(allBytes, StandardCharsets.UTF_8);
+        } catch (IOException ex) {
+            throw new IOException("Unable to read file " + path, ex);
+        }
     }
 
     /**
@@ -586,7 +715,7 @@ public static ByteBuffer ensureCapacity(ByteBuffer existingBuffer, int newLength
         return existingBuffer;
     }
 
-    /*
+    /**
      * Creates a set
      * @param elems the elements
      * @param  the type of element
@@ -594,46 +723,152 @@ public static ByteBuffer ensureCapacity(ByteBuffer existingBuffer, int newLength
      */
     @SafeVarargs
     public static  Set mkSet(T... elems) {
-        return new HashSet<>(Arrays.asList(elems));
+        Set result = new HashSet<>((int) (elems.length / 0.75) + 1);
+        for (T elem : elems)
+            result.add(elem);
+        return result;
     }
 
-    /*
-     * Creates a list
+    /**
+     * Creates a sorted set
      * @param elems the elements
-     * @param  the type of element
-     * @return List
+     * @param  the type of element, must be comparable
+     * @return SortedSet
      */
     @SafeVarargs
-    public static  List mkList(T... elems) {
-        return Arrays.asList(elems);
+    public static > SortedSet mkSortedSet(T... elems) {
+        SortedSet result = new TreeSet<>();
+        for (T elem : elems)
+            result.add(elem);
+        return result;
+    }
+
+    /**
+     * Creates a map entry (for use with {@link Utils#mkMap(java.util.Map.Entry[])})
+     *
+     * @param k   The key
+     * @param v   The value
+     * @param  The key type
+     * @param  The value type
+     * @return An entry
+     */
+    public static  Map.Entry mkEntry(final K k, final V v) {
+        return new Map.Entry() {
+            @Override
+            public K getKey() {
+                return k;
+            }
+
+            @Override
+            public V getValue() {
+                return v;
+            }
+
+            @Override
+            public V setValue(final V value) {
+                throw new UnsupportedOperationException();
+            }
+        };
+    }
+
+    /**
+     * Creates a map from a sequence of entries
+     *
+     * @param entries The entries to map
+     * @param      The key type
+     * @param      The value type
+     * @return A map
+     */
+    @SafeVarargs
+    public static  Map mkMap(final Map.Entry... entries) {
+        final LinkedHashMap result = new LinkedHashMap<>();
+        for (final Map.Entry entry : entries) {
+            result.put(entry.getKey(), entry.getValue());
+        }
+        return result;
+    }
+
+    /**
+     * Creates a {@link Properties} from a map
+     *
+     * @param properties A map of properties to add
+     * @return The properties object
+     */
+    public static Properties mkProperties(final Map properties) {
+        final Properties result = new Properties();
+        for (final Map.Entry entry : properties.entrySet()) {
+            result.setProperty(entry.getKey(), entry.getValue());
+        }
+        return result;
+    }
+
+    /**
+     * Creates a {@link Properties} from a map
+     *
+     * @param properties A map of properties to add
+     * @return The properties object
+     */
+    public static Properties mkObjectProperties(final Map properties) {
+        final Properties result = new Properties();
+        for (final Map.Entry entry : properties.entrySet()) {
+            result.put(entry.getKey(), entry.getValue());
+        }
+        return result;
     }
 
     /**
      * Recursively delete the given file/directory and any subfiles (if any exist)
      *
-     * @param file The root file at which to begin deleting
+     * @param rootFile The root file at which to begin deleting
      */
-    public static void delete(final File file) throws IOException {
-        if (file == null)
+    public static void delete(final File rootFile) throws IOException {
+        delete(rootFile, Collections.emptyList());
+    }
+
+    /**
+     * Recursively delete the subfiles (if any exist) of the passed in root file that are not included
+     * in the list to keep
+     *
+     * @param rootFile The root file at which to begin deleting
+     * @param filesToKeep The subfiles to keep (note that if a subfile is to be kept, so are all its parent
+     *                    files in its pat)h; if empty we would also delete the root file
+     */
+    public static void delete(final File rootFile, final List filesToKeep) throws IOException {
+        if (rootFile == null)
             return;
-        Files.walkFileTree(file.toPath(), new SimpleFileVisitor() {
+        Files.walkFileTree(rootFile.toPath(), new SimpleFileVisitor() {
             @Override
             public FileVisitResult visitFileFailed(Path path, IOException exc) throws IOException {
                 // If the root path did not exist, ignore the error; otherwise throw it.
-                if (exc instanceof NoSuchFileException && path.toFile().equals(file))
+                if (exc instanceof NoSuchFileException && path.toFile().equals(rootFile))
                     return FileVisitResult.TERMINATE;
                 throw exc;
             }
 
             @Override
             public FileVisitResult visitFile(Path path, BasicFileAttributes attrs) throws IOException {
-                Files.delete(path);
+                if (!filesToKeep.contains(path.toFile())) {
+                    Files.delete(path);
+                }
                 return FileVisitResult.CONTINUE;
             }
 
             @Override
             public FileVisitResult postVisitDirectory(Path path, IOException exc) throws IOException {
-                Files.delete(path);
+                // KAFKA-8999: if there's an exception thrown previously already, we should throw it
+                if (exc != null) {
+                    throw exc;
+                }
+
+                if (rootFile.toPath().equals(path)) {
+                    // only delete the parent directory if there's nothing to keep
+                    if (filesToKeep.isEmpty()) {
+                        Files.delete(path);
+                    }
+                } else if (!filesToKeep.contains(path.toFile())) {
+                    Files.delete(path);
+                }
+
                 return FileVisitResult.CONTINUE;
             }
         });
@@ -645,7 +880,7 @@ public FileVisitResult postVisitDirectory(Path path, IOException exc) throws IOE
      * @return
      */
     public static  List safe(List other) {
-        return other == null ? Collections.emptyList() : other;
+        return other == null ? Collections.emptyList() : other;
     }
 
    /**
@@ -712,6 +947,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.
      */
@@ -725,6 +972,27 @@ 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);
+            }
+        }
+    }
+
+    /**
+     * close all closable objects even if one of them throws exception.
+     * @param firstException keeps the first exception
+     * @param name message of closing those objects
+     * @param closeables closable objects
+     */
+    public static void closeAllQuietly(AtomicReference firstException, String name, AutoCloseable... closeables) {
+        for (AutoCloseable closeable : closeables) closeQuietly(closeable, name, firstException);
+    }
+
     /**
      * 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
@@ -742,10 +1010,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
@@ -866,6 +1130,10 @@ public static void writeTo(DataOutput out, ByteBuffer buffer, int length) throws
         }
     }
 
+    public static  List toList(Iterable iterable) {
+        return toList(iterable.iterator());
+    }
+
     public static  List toList(Iterator iterator) {
         List res = new ArrayList<>();
         while (iterator.hasNext())
@@ -873,4 +1141,187 @@ public static  List toList(Iterator iterator) {
         return res;
     }
 
+    public static  List concatListsUnmodifiable(List left, List right) {
+        return concatLists(left, right, Collections::unmodifiableList);
+    }
+
+    public static  List concatLists(List left, List right, Function, List> finisher) {
+        return Stream.concat(left.stream(), right.stream())
+                .collect(Collectors.collectingAndThen(Collectors.toList(), finisher));
+    }
+
+    public static int to32BitField(final Set bytes) {
+        int value = 0;
+        for (final byte b : bytes)
+            value |= 1 << checkRange(b);
+        return value;
+    }
+
+    private static byte checkRange(final byte i) {
+        if (i > 31)
+            throw new IllegalArgumentException("out of range: i>31, i = " + i);
+        if (i < 0)
+            throw new IllegalArgumentException("out of range: i<0, i = " + i);
+        return i;
+    }
+
+    public static Set from32BitField(final int intValue) {
+        Set result = new HashSet<>();
+        for (int itr = intValue, count = 0; itr != 0; itr >>>= 1) {
+            if ((itr & 1) != 0)
+                result.add((byte) count);
+            count++;
+        }
+        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())
+            )
+        );
+    }
+
+    /**
+     * A Collector that offers two kinds of convenience:
+     * 1. You can specify the concrete type of the returned Map
+     * 2. You can turn a stream of Entries directly into a Map without having to mess with a key function
+     *    and a value function. In particular, this is handy if all you need to do is apply a filter to a Map's entries.
+     *
+     *
+     * One thing to be wary of: These types are too "distant" for IDE type checkers to warn you if you
+     * try to do something like build a TreeMap of non-Comparable elements. You'd get a runtime exception for that.
+     *
+     * @param mapSupplier The constructor for your concrete map type.
+     * @param  The Map key type
+     * @param  The Map value type
+     * @param  The type of the Map itself.
+     * @return new Collector, M, M>
+     */
+    public static > Collector, M, M> entriesToMap(final Supplier mapSupplier) {
+        return new Collector, M, M>() {
+            @Override
+            public Supplier supplier() {
+                return mapSupplier;
+            }
+
+            @Override
+            public BiConsumer> accumulator() {
+                return (map, entry) -> map.put(entry.getKey(), entry.getValue());
+            }
+
+            @Override
+            public BinaryOperator combiner() {
+                return (map, map2) -> {
+                    map.putAll(map2);
+                    return map;
+                };
+            }
+
+            @Override
+            public Function finisher() {
+                return map -> map;
+            }
+
+            @Override
+            public Set characteristics() {
+                return EnumSet.of(Characteristics.UNORDERED, Characteristics.IDENTITY_FINISH);
+            }
+        };
+    }
+
+    @SafeVarargs
+    public static  Set union(final Supplier> constructor, final Set... set) {
+        final Set result = constructor.get();
+        for (final Set s : set) {
+            result.addAll(s);
+        }
+        return result;
+    }
+
+    @SafeVarargs
+    public static  Set intersection(final Supplier> constructor, final Set first, final Set... set) {
+        final Set result = constructor.get();
+        result.addAll(first);
+        for (final Set s : set) {
+            result.retainAll(s);
+        }
+        return result;
+    }
+
+    public static  Set diff(final Supplier> constructor, final Set left, final Set right) {
+        final Set result = constructor.get();
+        result.addAll(left);
+        result.removeAll(right);
+        return result;
+    }
+
+    /**
+     * Convert a properties to map. All keys in properties must be string type. Otherwise, a ConfigException is thrown.
+     * @param properties to be converted
+     * @return a map including all elements in properties
+     */
+    public static Map propsToMap(Properties properties) {
+        Map map = new HashMap<>(properties.size());
+        for (Map.Entry entry : properties.entrySet()) {
+            if (entry.getKey() instanceof String) {
+                String k = (String) entry.getKey();
+                map.put(k, properties.get(k));
+            } else {
+                throw new ConfigException(entry.getKey().toString(), entry.getValue(), "Key must be a string.");
+            }
+        }
+        return map;
+    }
+
+    /**
+     * Convert timestamp to an epoch value
+     * @param timestamp the timestamp to be converted, the accepted formats are:
+     *                 (1) yyyy-MM-dd'T'HH:mm:ss.SSS, ex: 2020-11-10T16:51:38.198
+     *                 (2) yyyy-MM-dd'T'HH:mm:ss.SSSZ, ex: 2020-11-10T16:51:38.198+0800
+     *                 (3) yyyy-MM-dd'T'HH:mm:ss.SSSX, ex: 2020-11-10T16:51:38.198+08
+     *                 (4) yyyy-MM-dd'T'HH:mm:ss.SSSXX, ex: 2020-11-10T16:51:38.198+0800
+     *                 (5) yyyy-MM-dd'T'HH:mm:ss.SSSXXX, ex: 2020-11-10T16:51:38.198+08:00
+     *
+     * @return epoch value of a given timestamp (i.e. the number of milliseconds since January 1, 1970, 00:00:00 GMT)
+     * @throws ParseException for timestamp that doesn't follow ISO8601 format or the format is not expected
+     */
+    public static long getDateTime(String timestamp) throws ParseException, IllegalArgumentException {
+        if (timestamp == null) {
+            throw new IllegalArgumentException("Error parsing timestamp with null value");
+        }
+
+        final String[] timestampParts = timestamp.split("T");
+        if (timestampParts.length < 2) {
+            throw new ParseException("Error parsing timestamp. It does not contain a 'T' according to ISO8601 format", timestamp.length());
+        }
+
+        final String secondPart = timestampParts[1];
+        if (!(secondPart.contains("+") || secondPart.contains("-") || secondPart.contains("Z"))) {
+            timestamp = timestamp + "Z";
+        }
+
+        SimpleDateFormat simpleDateFormat = new SimpleDateFormat();
+        // strictly parsing the date/time format
+        simpleDateFormat.setLenient(false);
+        try {
+            simpleDateFormat.applyPattern("yyyy-MM-dd'T'HH:mm:ss.SSSXXX");
+            final Date date = simpleDateFormat.parse(timestamp);
+            return date.getTime();
+        } catch (final ParseException e) {
+            simpleDateFormat.applyPattern("yyyy-MM-dd'T'HH:mm:ss.SSSX");
+            final Date date = simpleDateFormat.parse(timestamp);
+            return date.getTime();
+        }
+    }
+
+    @SuppressWarnings("unchecked")
+    public static  Iterator covariantCast(Iterator iterator) {
+        return (Iterator) iterator;
+    }
 }
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..1865e7e6ad4a4
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/server/authorizer/Authorizer.java
@@ -0,0 +1,142 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.
+     * 
            
+     * Refer to the authorizer implementation docs for details on concurrent update guarantees.
+     *
+     * @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/java/org/apache/kafka/server/policy/AlterConfigPolicy.java b/clients/src/main/java/org/apache/kafka/server/policy/AlterConfigPolicy.java
index ca47efa66458f..68174c14dc97c 100644
--- a/clients/src/main/java/org/apache/kafka/server/policy/AlterConfigPolicy.java
+++ b/clients/src/main/java/org/apache/kafka/server/policy/AlterConfigPolicy.java
@@ -23,12 +23,12 @@
 import java.util.Map;
 
 /**
- * An interface for enforcing a policy on alter configs requests.
+ * 
            An interface for enforcing a policy on alter configs requests.
  *
- * Common use cases are requiring that the replication factor, min.insync.replicas and/or retention settings for a
+ * 
            Common use cases are requiring that the replication factor, min.insync.replicas and/or retention settings for a
  * topic remain within an allowable range.
  *
- * If alter.config.policy.class.name is defined, Kafka will create an instance of the specified class
+ * 
            If alter.config.policy.class.name is defined, Kafka will create an instance of the specified class
  * using the default constructor and will then pass the broker configs to its configure() method. During
  * broker shutdown, the close() method will be invoked so that resources can be released (if necessary).
  */
diff --git a/clients/src/main/java/org/apache/kafka/server/policy/CreateTopicPolicy.java b/clients/src/main/java/org/apache/kafka/server/policy/CreateTopicPolicy.java
index e623d73f1ac16..e02fa5ff663e0 100644
--- a/clients/src/main/java/org/apache/kafka/server/policy/CreateTopicPolicy.java
+++ b/clients/src/main/java/org/apache/kafka/server/policy/CreateTopicPolicy.java
@@ -24,12 +24,12 @@
 import java.util.Map;
 
 /**
- * An interface for enforcing a policy on create topics requests.
+ * 
            An interface for enforcing a policy on create topics requests.
  *
- * Common use cases are requiring that the replication factor, min.insync.replicas and/or retention settings for a
+ * 
            Common use cases are requiring that the replication factor, min.insync.replicas and/or retention settings for a
  * topic are within an allowable range.
  *
- * If create.topic.policy.class.name is defined, Kafka will create an instance of the specified class
+ * 
            If create.topic.policy.class.name is defined, Kafka will create an instance of the specified class
  * using the default constructor and will then pass the broker configs to its configure() method. During
  * broker shutdown, the close() method will be invoked so that resources can be released (if necessary).
  */
diff --git a/clients/src/main/java/org/apache/kafka/server/quota/ClientQuotaCallback.java b/clients/src/main/java/org/apache/kafka/server/quota/ClientQuotaCallback.java
new file mode 100644
index 0000000000000..210e9f45840e2
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/server/quota/ClientQuotaCallback.java
@@ -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 org.apache.kafka.server.quota;
+
+import org.apache.kafka.common.Cluster;
+import org.apache.kafka.common.Configurable;
+import org.apache.kafka.common.security.auth.KafkaPrincipal;
+
+import java.util.Map;
+
+/**
+ * Quota callback interface for brokers that enables customization of client quota computation.
+ */
+public interface ClientQuotaCallback extends Configurable {
+
+    /**
+     * Quota callback invoked to determine the quota metric tags to be applied for a request.
+     * Quota limits are associated with quota metrics and all clients which use the same
+     * metric tags share the quota limit.
+     *
+     * @param quotaType Type of quota requested
+     * @param principal The user principal of the connection for which quota is requested
+     * @param clientId  The client id associated with the request
+     * @return quota metric tags that indicate which other clients share this quota
+     */
+    Map quotaMetricTags(ClientQuotaType quotaType, KafkaPrincipal principal, String clientId);
+
+    /**
+     * Returns the quota limit associated with the provided metric tags. These tags were returned from
+     * a previous call to {@link #quotaMetricTags(ClientQuotaType, KafkaPrincipal, String)}. This method is
+     * invoked by quota managers to obtain the current quota limit applied to a metric when the first request
+     * using these tags is processed. It is also invoked after a quota update or cluster metadata change.
+     * If the tags are no longer in use after the update, (e.g. this is a {user, client-id} quota metric
+     * and the quota now in use is a {user} quota), null is returned.
+     *
+     * @param quotaType  Type of quota requested
+     * @param metricTags Metric tags for a quota metric of type `quotaType`
+     * @return the quota limit for the provided metric tags or null if the metric tags are no longer in use
+     */
+    Double quotaLimit(ClientQuotaType quotaType, Map metricTags);
+
+    /**
+     * Quota configuration update callback that is invoked when quota configuration for an entity is
+     * updated in ZooKeeper. This is useful to track configured quotas if built-in quota configuration
+     * tools are used for quota management.
+     *
+     * @param quotaType   Type of quota being updated
+     * @param quotaEntity The quota entity for which quota is being updated
+     * @param newValue    The new quota value
+     */
+    void updateQuota(ClientQuotaType quotaType, ClientQuotaEntity quotaEntity, double newValue);
+
+    /**
+     * Quota configuration removal callback that is invoked when quota configuration for an entity is
+     * removed in ZooKeeper. This is useful to track configured quotas if built-in quota configuration
+     * tools are used for quota management.
+     *
+     * @param quotaType   Type of quota being updated
+     * @param quotaEntity The quota entity for which quota is being updated
+     */
+    void removeQuota(ClientQuotaType quotaType, ClientQuotaEntity quotaEntity);
+
+    /**
+     * Returns true if any of the existing quota configs may have been updated since the last call
+     * to this method for the provided quota type. Quota updates as a result of calls to
+     * {@link #updateClusterMetadata(Cluster)}, {@link #updateQuota(ClientQuotaType, ClientQuotaEntity, double)}
+     * and {@link #removeQuota(ClientQuotaType, ClientQuotaEntity)} are automatically processed.
+     * So callbacks that rely only on built-in quota configuration tools always return false. Quota callbacks
+     * with external quota configuration or custom reconfigurable quota configs that affect quota limits must
+     * return true if existing metric configs may need to be updated. This method is invoked on every request
+     * and hence is expected to be handled by callbacks as a simple flag that is updated when quotas change.
+     *
+     * @param quotaType Type of quota
+     */
+    boolean quotaResetRequired(ClientQuotaType quotaType);
+
+    /**
+     * Metadata update callback that is invoked whenever UpdateMetadata request is received from
+     * the controller. This is useful if quota computation takes partitions into account.
+     * Topics that are being deleted will not be included in `cluster`.
+     *
+     * @param cluster Cluster metadata including partitions and their leaders if known
+     * @return true if quotas have changed and metric configs may need to be updated
+     */
+    boolean updateClusterMetadata(Cluster cluster);
+
+    /**
+     * Closes this instance.
+     */
+    void close();
+}
+
diff --git a/clients/src/main/java/org/apache/kafka/server/quota/ClientQuotaEntity.java b/clients/src/main/java/org/apache/kafka/server/quota/ClientQuotaEntity.java
new file mode 100644
index 0000000000000..0f0eb6253c828
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/server/quota/ClientQuotaEntity.java
@@ -0,0 +1,62 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.quota;
+
+import java.util.List;
+
+/**
+ * The metadata for an entity for which quota is configured. Quotas may be defined at
+ * different levels and `configEntities` gives the list of config entities that define
+ * the level of this quota entity.
+ */
+public interface ClientQuotaEntity {
+
+    /**
+     * Entity type of a {@link ConfigEntity}
+     */
+    enum ConfigEntityType {
+        USER,
+        CLIENT_ID,
+        DEFAULT_USER,
+        DEFAULT_CLIENT_ID
+    }
+
+    /**
+     * Interface representing a quota configuration entity. Quota may be
+     * configured at levels that include one or more configuration entities.
+     * For example, {user, client-id} quota is represented using two
+     * instances of ConfigEntity with entity types USER and CLIENT_ID.
+     */
+    interface ConfigEntity {
+        /**
+         * Returns the name of this entity. For default quotas, an empty string is returned.
+         */
+        String name();
+
+        /**
+         * Returns the type of this entity.
+         */
+        ConfigEntityType entityType();
+    }
+
+    /**
+     * Returns the list of configuration entities that this quota entity is comprised of.
+     * For {user} or {clientId} quota, this is a single entity and for {user, clientId}
+     * quota, this is a list of two entities.
+     */
+    List configEntities();
+}
diff --git a/clients/src/main/java/org/apache/kafka/server/quota/ClientQuotaType.java b/clients/src/main/java/org/apache/kafka/server/quota/ClientQuotaType.java
new file mode 100644
index 0000000000000..5b0828a08211f
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/server/quota/ClientQuotaType.java
@@ -0,0 +1,27 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.quota;
+
+/**
+ * Types of quotas that may be configured on brokers for client requests.
+ */
+public enum ClientQuotaType {
+    PRODUCE,
+    FETCH,
+    REQUEST,
+    CONTROLLER_MUTATION
+}
diff --git a/clients/src/main/resources/META-INF/services/org.apache.kafka.common.config.provider.ConfigProvider b/clients/src/main/resources/META-INF/services/org.apache.kafka.common.config.provider.ConfigProvider
new file mode 100644
index 0000000000000..146de7abe4442
--- /dev/null
+++ b/clients/src/main/resources/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.
+
+org.apache.kafka.common.config.provider.FileConfigProvider
diff --git a/clients/src/main/resources/common/message/AddOffsetsToTxnRequest.json b/clients/src/main/resources/common/message/AddOffsetsToTxnRequest.json
new file mode 100644
index 0000000000000..b2bb9a78f6470
--- /dev/null
+++ b/clients/src/main/resources/common/message/AddOffsetsToTxnRequest.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": 25,
+  "type": "request",
+  "name": "AddOffsetsToTxnRequest",
+  // Version 1 is the same as version 0.
+  //
+  // Version 2 adds the support for new error code PRODUCER_FENCED.
+  //
+  // Version 3 enables flexible versions.
+  "validVersions": "0-3",
+  "flexibleVersions": "3+",
+  "fields": [
+    { "name": "TransactionalId", "type": "string", "versions": "0+", "entityType": "transactionalId",
+      "about": "The transactional id corresponding to the transaction."},
+    { "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+", "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
new file mode 100644
index 0000000000000..71fa655227081
--- /dev/null
+++ b/clients/src/main/resources/common/message/AddOffsetsToTxnResponse.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": 25,
+  "type": "response",
+  "name": "AddOffsetsToTxnResponse",
+  // Starting in version 1, on quota violation brokers send out responses before throttling.
+  //
+  // Version 2 adds the support for new error code PRODUCER_FENCED.
+  //
+  // Version 3 enables flexible versions.
+  "validVersions": "0-3",
+  "flexibleVersions": "3+",
+  "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": "ErrorCode", "type": "int16", "versions": "0+",
+      "about": "The response error code, or 0 if there was no error." }
+  ]
+}
diff --git a/clients/src/main/resources/common/message/AddPartitionsToTxnRequest.json b/clients/src/main/resources/common/message/AddPartitionsToTxnRequest.json
new file mode 100644
index 0000000000000..3e1d7207a51dc
--- /dev/null
+++ b/clients/src/main/resources/common/message/AddPartitionsToTxnRequest.json
@@ -0,0 +1,42 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements.  See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+{
+  "apiKey": 24,
+  "type": "request",
+  "name": "AddPartitionsToTxnRequest",
+  // Version 1 is the same as version 0.
+  //
+  // Version 2 adds the support for new error code PRODUCER_FENCED.
+  //
+  // Version 3 enables flexible versions.
+  "validVersions": "0-3",
+  "flexibleVersions": "3+",
+  "fields": [
+    { "name": "TransactionalId", "type": "string", "versions": "0+", "entityType": "transactionalId",
+      "about": "The transactional id corresponding to the transaction."},
+    { "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 transaction.", "fields": [
+      { "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
new file mode 100644
index 0000000000000..4241dc77b4a66
--- /dev/null
+++ b/clients/src/main/resources/common/message/AddPartitionsToTxnResponse.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": 24,
+  "type": "response",
+  "name": "AddPartitionsToTxnResponse",
+  // Starting in version 1, on quota violation brokers send out responses before throttling.
+  //
+  // Version 2 adds the support for new error code PRODUCER_FENCED.
+  //
+  // Version 3 enables flexible versions.
+  "validVersions": "0-3",
+  "flexibleVersions": "3+",
+  "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, "entityType": "topicName",
+        "about": "The topic name." },
+      { "name": "Results", "type": "[]AddPartitionsToTxnPartitionResult", "versions": "0+", 
+        "about": "The results for each partition", "fields": [
+        { "name": "PartitionIndex", "type": "int32", "versions": "0+", "mapKey": true,
+          "about": "The partition indexes." },
+        { "name": "ErrorCode", "type": "int16", "versions": "0+",
+          "about": "The response error code."}
+      ]}
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/AlterClientQuotasRequest.json b/clients/src/main/resources/common/message/AlterClientQuotasRequest.json
new file mode 100644
index 0000000000000..715a00644e9c0
--- /dev/null
+++ b/clients/src/main/resources/common/message/AlterClientQuotasRequest.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": 49,
+  "type": "request",
+  "name": "AlterClientQuotasRequest",
+  "validVersions": "0-1",
+  // Version 1 enables flexible versions.
+  "flexibleVersions": "1+",
+  "fields": [
+    { "name": "Entries", "type": "[]EntryData", "versions": "0+",
+      "about": "The quota configuration entries to alter.", "fields": [
+      { "name": "Entity", "type": "[]EntityData", "versions": "0+",
+        "about": "The quota entity to alter.", "fields": [
+        { "name": "EntityType", "type": "string", "versions": "0+",
+          "about": "The entity type." },
+        { "name": "EntityName", "type": "string", "versions": "0+", "nullableVersions": "0+",
+          "about": "The name of the entity, or null if the default." }
+      ]},
+      { "name": "Ops", "type": "[]OpData", "versions": "0+",
+        "about": "An individual quota configuration entry to alter.", "fields": [
+        { "name": "Key", "type": "string", "versions": "0+",
+          "about": "The quota configuration key." },
+        { "name": "Value", "type": "float64", "versions": "0+",
+          "about": "The value to set, otherwise ignored if the value is to be removed." },
+        { "name": "Remove", "type": "bool", "versions": "0+",
+          "about": "Whether the quota configuration value should be removed, otherwise set." }
+      ]}
+    ]},
+    { "name": "ValidateOnly", "type": "bool", "versions": "0+",
+      "about": "Whether the alteration should be validated, but not performed." }
+  ]
+}
diff --git a/clients/src/main/resources/common/message/AlterClientQuotasResponse.json b/clients/src/main/resources/common/message/AlterClientQuotasResponse.json
new file mode 100644
index 0000000000000..326b85d17769d
--- /dev/null
+++ b/clients/src/main/resources/common/message/AlterClientQuotasResponse.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": 49,
+  "type": "response",
+  "name": "AlterClientQuotasResponse",
+  // Version 1 enables flexible versions.
+  "validVersions": "0-1",
+  "flexibleVersions": "1+",
+  "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": "Entries", "type": "[]EntryData", "versions": "0+",
+      "about": "The quota configuration entries to alter.", "fields": [
+      { "name": "ErrorCode", "type": "int16", "versions": "0+",
+        "about": "The error code, or `0` if the quota alteration succeeded." },
+      { "name": "ErrorMessage", "type": "string", "versions": "0+", "nullableVersions": "0+",
+        "about": "The error message, or `null` if the quota alteration succeeded." },
+      { "name": "Entity", "type": "[]EntityData", "versions": "0+",
+        "about": "The quota entity to alter.", "fields": [
+        { "name": "EntityType", "type": "string", "versions": "0+",
+          "about": "The entity type." },
+        { "name": "EntityName", "type": "string", "versions": "0+", "nullableVersions": "0+",
+          "about": "The name of the entity, or null if the default." }
+      ]}
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/AlterConfigsRequest.json b/clients/src/main/resources/common/message/AlterConfigsRequest.json
new file mode 100644
index 0000000000000..a1d7d1d4ca4d4
--- /dev/null
+++ b/clients/src/main/resources/common/message/AlterConfigsRequest.json
@@ -0,0 +1,42 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements.  See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+{
+  "apiKey": 33,
+  "type": "request",
+  "name": "AlterConfigsRequest",
+  // Version 1 is the same as version 0.
+  // Version 2 enables flexible versions.
+  "validVersions": "0-2",
+  "flexibleVersions": "2+",
+  "fields": [
+    { "name": "Resources", "type": "[]AlterConfigsResource", "versions": "0+", 
+      "about": "The 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": "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/AlterConfigsResponse.json b/clients/src/main/resources/common/message/AlterConfigsResponse.json
new file mode 100644
index 0000000000000..fab102b6caf3b
--- /dev/null
+++ b/clients/src/main/resources/common/message/AlterConfigsResponse.json
@@ -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.
+
+{
+  "apiKey": 33,
+  "type": "response",
+  "name": "AlterConfigsResponse",
+  // Starting in version 1, on quota violation brokers send out responses before throttling.
+  // Version 2 enables flexible versions.
+  "validVersions": "0-2",
+  "flexibleVersions": "2+",
+  "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/AlterIsrRequest.json b/clients/src/main/resources/common/message/AlterIsrRequest.json
new file mode 100644
index 0000000000000..3d5084a4094ec
--- /dev/null
+++ b/clients/src/main/resources/common/message/AlterIsrRequest.json
@@ -0,0 +1,42 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements.  See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implie
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+{
+  "apiKey": 56,
+  "type": "request",
+  "name": "AlterIsrRequest",
+  "validVersions": "0",
+  "flexibleVersions": "0+",
+  "fields": [
+    { "name": "BrokerId", "type": "int32", "versions": "0+", "entityType": "brokerId",
+      "about": "The ID of the requesting broker" },
+    { "name": "BrokerEpoch", "type": "int64", "versions": "0+", "default": "-1",
+      "about": "The epoch of the requesting broker" },
+    { "name": "Topics", "type": "[]TopicData", "versions": "0+", "fields": [
+      { "name":  "Name", "type": "string", "versions": "0+", "entityType": "topicName",
+        "about": "The name of the topic to alter ISRs for" },
+      { "name": "Partitions", "type": "[]PartitionData", "versions": "0+", "fields": [
+        { "name": "PartitionIndex", "type": "int32", "versions": "0+",
+          "about": "The partition index" },
+        { "name": "LeaderEpoch", "type": "int32", "versions": "0+",
+          "about": "The leader epoch of this partition" },
+        { "name": "NewIsr", "type": "[]int32", "versions": "0+",
+          "about": "The ISR for this partition"},
+        { "name": "CurrentIsrVersion", "type": "int32", "versions": "0+",
+          "about": "The expected version of ISR which is being updated"}
+      ]}
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/AlterIsrResponse.json b/clients/src/main/resources/common/message/AlterIsrResponse.json
new file mode 100644
index 0000000000000..7a81339dda8c0
--- /dev/null
+++ b/clients/src/main/resources/common/message/AlterIsrResponse.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": 56,
+  "type": "response",
+  "name": "AlterIsrResponse",
+  "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 response error code" },
+    { "name": "Topics", "type": "[]TopicData", "versions": "0+", "fields": [
+      { "name":  "Name", "type": "string", "versions": "0+", "entityType": "topicName",
+        "about": "The name of the topic" },
+      { "name": "Partitions", "type": "[]PartitionData", "versions": "0+", "fields": [
+        { "name": "PartitionIndex", "type": "int32", "versions": "0+",
+          "about": "The partition index" },
+        { "name": "ErrorCode", "type": "int16", "versions": "0+",
+          "about": "The partition level error code" },
+        { "name": "LeaderId", "type": "int32", "versions": "0+", "entityType": "brokerId",
+          "about": "The broker ID of the leader." },
+        { "name": "LeaderEpoch", "type": "int32", "versions": "0+",
+          "about": "The leader epoch." },
+        { "name": "Isr", "type": "[]int32", "versions": "0+",
+          "about": "The in-sync replica IDs." },
+        { "name": "CurrentIsrVersion", "type": "int32", "versions": "0+",
+          "about": "The current ISR version." }
+      ]}
+    ]}
+  ]
+}
\ No newline at end of file
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
new file mode 100644
index 0000000000000..a6749077f4439
--- /dev/null
+++ b/clients/src/main/resources/common/message/AlterReplicaLogDirsRequest.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": 34,
+  "type": "request",
+  "name": "AlterReplicaLogDirsRequest",
+  // Version 1 is the same as version 0.
+  // Version 2 enables flexible versions.
+  "validVersions": "0-2",
+  "flexibleVersions": "2+",
+  "fields": [
+    { "name": "Dirs", "type": "[]AlterReplicaLogDir", "versions": "0+", 
+      "about": "The alterations to make for each directory.", "fields": [
+      { "name": "Path", "type": "string", "versions": "0+", "mapKey": true,
+        "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, "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
new file mode 100644
index 0000000000000..386e24e9d8600
--- /dev/null
+++ b/clients/src/main/resources/common/message/AlterReplicaLogDirsResponse.json
@@ -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.
+
+{
+  "apiKey": 34,
+  "type": "response",
+  "name": "AlterReplicaLogDirsResponse",
+  // Starting in version 1, on quota violation brokers send out responses before throttling.
+  // Version 2 enables flexible versions.
+  "validVersions": "0-2",
+  "flexibleVersions": "2+",
+  "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+", "entityType": "topicName",
+        "about": "The name of the topic." },
+      { "name": "Partitions", "type": "[]AlterReplicaLogDirPartitionResult", "versions": "0+",
+        "about": "The results for each partition.", "fields": [
+        { "name": "PartitionIndex", "type": "int32", "versions": "0+",
+          "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/AlterUserScramCredentialsRequest.json b/clients/src/main/resources/common/message/AlterUserScramCredentialsRequest.json
new file mode 100644
index 0000000000000..242bcb97dda79
--- /dev/null
+++ b/clients/src/main/resources/common/message/AlterUserScramCredentialsRequest.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": 51,
+  "type": "request",
+  "name": "AlterUserScramCredentialsRequest",
+  "validVersions": "0",
+  "flexibleVersions": "0+",
+  "fields": [
+    { "name": "Deletions", "type": "[]ScramCredentialDeletion", "versions": "0+",
+      "about": "The SCRAM credentials to remove.", "fields": [
+      { "name": "Name", "type": "string", "versions": "0+",
+        "about": "The user name." },
+      { "name": "Mechanism", "type": "int8", "versions": "0+",
+        "about": "The SCRAM mechanism." }
+    ]},
+    { "name": "Upsertions", "type": "[]ScramCredentialUpsertion", "versions": "0+",
+      "about": "The SCRAM credentials to update/insert.", "fields": [
+      { "name": "Name", "type": "string", "versions": "0+",
+        "about": "The user name." },
+      { "name": "Mechanism", "type": "int8", "versions": "0+",
+        "about": "The SCRAM mechanism." },
+      { "name": "Iterations", "type": "int32", "versions": "0+",
+        "about": "The number of iterations." },
+      { "name": "Salt", "type": "bytes", "versions": "0+",
+        "about": "A random salt generated by the client." },
+      { "name": "SaltedPassword", "type": "bytes", "versions": "0+",
+        "about": "The salted password." }
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/AlterUserScramCredentialsResponse.json b/clients/src/main/resources/common/message/AlterUserScramCredentialsResponse.json
new file mode 100644
index 0000000000000..92b62d52abdd5
--- /dev/null
+++ b/clients/src/main/resources/common/message/AlterUserScramCredentialsResponse.json
@@ -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.
+
+{
+  "apiKey": 51,
+  "type": "response",
+  "name": "AlterUserScramCredentialsResponse",
+  "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": "Results", "type": "[]AlterUserScramCredentialsResult", "versions": "0+",
+      "about": "The results for deletions and alterations, one per affected user.", "fields": [
+      { "name": "User", "type": "string", "versions": "0+",
+        "about": "The user name." },
+      { "name": "ErrorCode", "type": "int16", "versions": "0+",
+        "about": "The error code." },
+      { "name": "ErrorMessage", "type": "string", "versions": "0+", "nullableVersions": "0+",
+        "about": "The error message, if any." }
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/ApiVersionsRequest.json b/clients/src/main/resources/common/message/ApiVersionsRequest.json
new file mode 100644
index 0000000000000..66e4511a92e93
--- /dev/null
+++ b/clients/src/main/resources/common/message/ApiVersionsRequest.json
@@ -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.
+
+{
+  "apiKey": 18,
+  "type": "request",
+  "name": "ApiVersionsRequest",
+  // Versions 0 through 2 of ApiVersionsRequest are the same.
+  //
+  // Version 3 is the first flexible version and adds ClientSoftwareName and ClientSoftwareVersion.
+  "validVersions": "0-3",
+  "flexibleVersions": "3+",
+  "fields": [
+    { "name": "ClientSoftwareName", "type": "string", "versions": "3+",
+      "ignorable": true, "about": "The name of the client." },
+    { "name": "ClientSoftwareVersion", "type": "string", "versions": "3+",
+      "ignorable": true, "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
new file mode 100644
index 0000000000000..ba6f01cb9434a
--- /dev/null
+++ b/clients/src/main/resources/common/message/ApiVersionsResponse.json
@@ -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.
+
+{
+  "apiKey": 18,
+  "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.
+  //
+  // 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": "ApiKey", "type": "int16", "versions": "0+", "mapKey": true,
+        "about": "The API index." },
+      { "name": "MinVersion", "type": "int16", "versions": "0+",
+        "about": "The minimum supported version, inclusive." },
+      { "name": "MaxVersion", "type": "int16", "versions": "0+",
+        "about": "The maximum supported version, inclusive." }
+    ]},
+    { "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":  "SupportedFeatures", "type": "[]SupportedFeatureKey",
+      "versions":  "3+", "tag": 0, "taggedVersions": "3+",
+      "about": "Features supported by the broker.",
+      "fields":  [
+        { "name": "Name", "type": "string", "versions": "3+", "mapKey": true,
+          "about": "The name of the feature." },
+        { "name": "MinVersion", "type": "int16", "versions": "3+",
+          "about": "The minimum supported version for the feature." },
+        { "name": "MaxVersion", "type": "int16", "versions": "3+",
+          "about": "The maximum supported version for the feature." }
+      ]
+    },
+    { "name": "FinalizedFeaturesEpoch", "type": "int64", "versions": "3+",
+      "tag": 1, "taggedVersions": "3+", "default": "-1", "ignorable": true,
+      "about": "The monotonically increasing epoch for the finalized features information. Valid values are >= 0. A value of -1 is special and represents unknown epoch."},
+    { "name":  "FinalizedFeatures", "type": "[]FinalizedFeatureKey",
+      "versions":  "3+", "tag": 2, "taggedVersions": "3+",
+      "about": "List of cluster-wide finalized features. The information is valid only if FinalizedFeaturesEpoch >= 0.",
+      "fields":  [
+        {"name": "Name", "type": "string", "versions":  "3+", "mapKey": true,
+          "about": "The name of the feature."},
+        {"name":  "MaxVersionLevel", "type": "int16", "versions":  "3+",
+          "about": "The cluster-wide finalized max version level for the feature."},
+        {"name":  "MinVersionLevel", "type": "int16", "versions":  "3+",
+          "about": "The cluster-wide finalized min version level for the feature."}
+      ]
+    }
+  ]
+}
diff --git a/clients/src/main/resources/common/message/BeginQuorumEpochRequest.json b/clients/src/main/resources/common/message/BeginQuorumEpochRequest.json
new file mode 100644
index 0000000000000..fd986097085f6
--- /dev/null
+++ b/clients/src/main/resources/common/message/BeginQuorumEpochRequest.json
@@ -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.
+
+{
+  "apiKey": 53,
+  "type": "request",
+  "name": "BeginQuorumEpochRequest",
+  "validVersions": "0",
+  "fields": [
+    { "name": "ClusterId", "type": "string", "versions": "0+",
+      "nullableVersions": "0+", "default": "null"},
+    { "name": "Topics", "type": "[]TopicData",
+      "versions": "0+", "fields": [
+      { "name": "TopicName", "type": "string", "versions": "0+", "entityType": "topicName",
+        "about": "The topic name." },
+      { "name": "Partitions", "type": "[]PartitionData",
+        "versions": "0+", "fields": [
+        { "name": "PartitionIndex", "type": "int32", "versions": "0+",
+          "about": "The partition index." },
+        { "name": "LeaderId", "type": "int32", "versions": "0+",
+          "about": "The ID of the newly elected leader"},
+        { "name": "LeaderEpoch", "type": "int32", "versions": "0+",
+          "about": "The epoch of the newly elected leader"}
+      ]}
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/BeginQuorumEpochResponse.json b/clients/src/main/resources/common/message/BeginQuorumEpochResponse.json
new file mode 100644
index 0000000000000..41e2292b498dd
--- /dev/null
+++ b/clients/src/main/resources/common/message/BeginQuorumEpochResponse.json
@@ -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.
+
+{
+  "apiKey": 53,
+  "type": "response",
+  "name": "BeginQuorumEpochResponse",
+  "validVersions": "0",
+  "fields": [
+    { "name": "ErrorCode", "type": "int16", "versions": "0+",
+      "about": "The top level error code."},
+    { "name": "Topics", "type": "[]TopicData",
+      "versions": "0+", "fields": [
+      { "name": "TopicName", "type": "string", "versions": "0+", "entityType": "topicName",
+        "about": "The topic name." },
+      { "name": "Partitions", "type": "[]PartitionData",
+        "versions": "0+", "fields": [
+        { "name": "PartitionIndex", "type": "int32", "versions": "0+",
+          "about": "The partition index." },
+        { "name": "ErrorCode", "type": "int16", "versions": "0+"},
+        { "name": "LeaderId", "type": "int32", "versions": "0+",
+          "about": "The ID of the current leader or -1 if the leader is unknown."},
+        { "name": "LeaderEpoch", "type": "int32", "versions": "0+",
+          "about": "The latest known leader epoch"}
+      ]}
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/ConsumerProtocolAssignment.json b/clients/src/main/resources/common/message/ConsumerProtocolAssignment.json
new file mode 100644
index 0000000000000..2ad373a5659ec
--- /dev/null
+++ b/clients/src/main/resources/common/message/ConsumerProtocolAssignment.json
@@ -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.
+
+{
+  "type": "data",
+  "name": "ConsumerProtocolAssignment",
+  // Assignment part of the Consumer Protocol.
+  //
+  // 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.
+  "validVersions": "0-1",
+  "fields": [
+    { "name": "AssignedPartitions", "type": "[]TopicPartition", "versions": "0+",
+      "fields": [
+        { "name": "Topic", "type": "string", "versions": "0+" },
+        { "name": "Partitions", "type": "[]int32", "versions": "0+" }
+      ]
+    },
+    { "name": "UserData", "type": "bytes", "versions": "0+", "nullableVersions": "0+",
+      "default": "null", "zeroCopy": true }
+  ]
+}
diff --git a/clients/src/main/resources/common/message/ConsumerProtocolSubscription.json b/clients/src/main/resources/common/message/ConsumerProtocolSubscription.json
new file mode 100644
index 0000000000000..fa4c371d70de7
--- /dev/null
+++ b/clients/src/main/resources/common/message/ConsumerProtocolSubscription.json
@@ -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.
+
+{
+  "type": "data",
+  "name": "ConsumerProtocolSubscription",
+  // Subscription part of the Consumer Protocol.
+  //
+  // 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.
+  "validVersions": "0-1",
+  "fields": [
+    { "name": "Topics", "type": "[]string", "versions": "0+" },
+    { "name": "UserData", "type": "bytes", "versions": "0+", "nullableVersions": "0+",
+      "default": "null", "zeroCopy": true },
+    { "name": "OwnedPartitions", "type": "[]TopicPartition", "versions": "1+", "ignorable": true,
+      "fields": [
+        { "name": "Topic", "type": "string", "versions": "1+" },
+        { "name": "Partitions", "type": "[]int32", "versions": "1+"}
+      ]
+    }
+  ]
+}
diff --git a/clients/src/main/resources/common/message/ControlledShutdownRequest.json b/clients/src/main/resources/common/message/ControlledShutdownRequest.json
new file mode 100644
index 0000000000000..5756d1c16b5df
--- /dev/null
+++ b/clients/src/main/resources/common/message/ControlledShutdownRequest.json
@@ -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.
+
+{
+  "apiKey": 7,
+  "type": "request",
+  "name": "ControlledShutdownRequest",
+  // Version 0 of ControlledShutdownRequest has a non-standard request header
+  // which does not include clientId.  Version 1 and later use the standard
+  // request header.
+  //
+  // Version 1 is the same as version 0.
+  //
+  // Version 2 adds BrokerEpoch.
+  "validVersions": "0-3",
+  "flexibleVersions": "3+",
+  "fields": [
+    { "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
new file mode 100644
index 0000000000000..27feb1b69b087
--- /dev/null
+++ b/clients/src/main/resources/common/message/ControlledShutdownResponse.json
@@ -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.
+
+{
+  "apiKey": 7,
+  "type": "response",
+  "name": "ControlledShutdownResponse",
+  // Versions 1 and 2 are the same as version 0.
+  "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, "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
new file mode 100644
index 0000000000000..a9bd9c5f60a11
--- /dev/null
+++ b/clients/src/main/resources/common/message/CreateAclsRequest.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": 30,
+  "type": "request",
+  "name": "CreateAclsRequest",
+  // Version 1 adds resource pattern type.
+  // Version 2 enables flexible versions.
+  "validVersions": "0-2",
+  "flexibleVersions": "2+",
+  "fields": [
+    { "name": "Creations", "type": "[]AclCreation", "versions": "0+",
+      "about": "The ACLs that we want to create.", "fields": [
+      { "name": "ResourceType", "type": "int8", "versions": "0+",
+        "about": "The type of the resource." },
+      { "name": "ResourceName", "type": "string", "versions": "0+",
+        "about": "The resource name for the ACL." },
+      { "name": "ResourcePatternType", "type": "int8", "versions": "1+", "default": "3",
+        "about": "The pattern type for the ACL." },
+      { "name": "Principal", "type": "string", "versions": "0+",
+        "about": "The principal for the ACL." },
+      { "name": "Host", "type": "string", "versions": "0+",
+        "about": "The host for the ACL." },
+      { "name": "Operation", "type": "int8", "versions": "0+",
+        "about": "The operation type for the ACL (read, write, etc.)." },
+      { "name": "PermissionType", "type": "int8", "versions": "0+",
+        "about": "The permission type for the ACL (allow, deny, etc.)." }
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/CreateAclsResponse.json b/clients/src/main/resources/common/message/CreateAclsResponse.json
new file mode 100644
index 0000000000000..7b0de7e56b731
--- /dev/null
+++ b/clients/src/main/resources/common/message/CreateAclsResponse.json
@@ -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.
+
+{
+  "apiKey": 30,
+  "type": "response",
+  "name": "CreateAclsResponse",
+  // Starting in version 1, on quota violation, brokers send out responses before throttling.
+  // Version 2 enables flexible versions.
+  "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": "[]AclCreationResult", "versions": "0+",
+      "about": "The results for each ACL creation.", "fields": [
+      { "name": "ErrorCode", "type": "int16", "versions": "0+",
+        "about": "The result error, or zero if there was no error." },
+      { "name": "ErrorMessage", "type": "string", "nullableVersions": "0+", "versions": "0+",
+        "about": "The result message, or null if there was no error." }
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/CreateDelegationTokenRequest.json b/clients/src/main/resources/common/message/CreateDelegationTokenRequest.json
new file mode 100644
index 0000000000000..8d881355aa9d2
--- /dev/null
+++ b/clients/src/main/resources/common/message/CreateDelegationTokenRequest.json
@@ -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.
+
+{
+  "apiKey": 38,
+  "type": "request",
+  "name": "CreateDelegationTokenRequest",
+  // Version 1 is the same as version 0.
+  //
+  // 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": [
+      { "name": "PrincipalType", "type": "string", "versions": "0+",
+        "about": "The type of the Kafka principal." },
+      { "name": "PrincipalName", "type": "string", "versions": "0+",
+        "about": "The name of the Kafka principal." }
+    ]},
+    { "name": "MaxLifetimeMs", "type": "int64", "versions": "0+",
+      "about": "The maximum lifetime of the token in milliseconds, or -1 to use the server side default." }
+  ]
+}
diff --git a/clients/src/main/resources/common/message/CreateDelegationTokenResponse.json b/clients/src/main/resources/common/message/CreateDelegationTokenResponse.json
new file mode 100644
index 0000000000000..74ad905b94b26
--- /dev/null
+++ b/clients/src/main/resources/common/message/CreateDelegationTokenResponse.json
@@ -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.
+
+{
+  "apiKey": 38,
+  "type": "response",
+  "name": "CreateDelegationTokenResponse",
+  // Starting in version 1, on quota violation, brokers send out responses before throttling.
+  //
+  // 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."},
+    { "name": "PrincipalType", "type": "string", "versions": "0+",
+      "about": "The principal type of the token owner." },
+    { "name": "PrincipalName", "type": "string", "versions": "0+",
+      "about": "The name of the token owner." },
+    { "name": "IssueTimestampMs", "type": "int64", "versions": "0+",
+      "about": "When this token was generated." },
+    { "name": "ExpiryTimestampMs", "type": "int64", "versions": "0+",
+      "about": "When this token expires." },
+    { "name": "MaxTimestampMs", "type": "int64", "versions": "0+",
+      "about": "The maximum lifetime of this token." },
+    { "name": "TokenId", "type": "string", "versions": "0+",
+      "about": "The token UUID." },
+    { "name": "Hmac", "type": "bytes", "versions": "0+",
+      "about": "HMAC of the delegation token." },
+    { "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/CreatePartitionsRequest.json b/clients/src/main/resources/common/message/CreatePartitionsRequest.json
new file mode 100644
index 0000000000000..ba1138801c1ef
--- /dev/null
+++ b/clients/src/main/resources/common/message/CreatePartitionsRequest.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": 37,
+  "type": "request",
+  "name": "CreatePartitionsRequest",
+  // Version 1 is the same as version 0.
+  //
+  // Version 2 adds flexible version support
+  //
+  // Version 3 is identical to version 2 but may return a THROTTLING_QUOTA_EXCEEDED error
+  // in the response if the partitions creation is throttled (KIP-599).
+  "validVersions": "0-3",
+  "flexibleVersions": "2+",
+  "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+", "mapKey": true, "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+", "entityType": "brokerId",
+          "about": "The assigned broker IDs." }
+      ]}
+    ]},
+    { "name": "TimeoutMs", "type": "int32", "versions": "0+",
+      "about": "The time in ms to wait for the partitions to be created." },
+    { "name": "ValidateOnly", "type": "bool", "versions": "0+",
+      "about": "If true, then validate the request, but don't actually increase the number of partitions." }
+  ]
+}
diff --git a/clients/src/main/resources/common/message/CreatePartitionsResponse.json b/clients/src/main/resources/common/message/CreatePartitionsResponse.json
new file mode 100644
index 0000000000000..ef9f1f6f4bae2
--- /dev/null
+++ b/clients/src/main/resources/common/message/CreatePartitionsResponse.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": 37,
+  "type": "response",
+  "name": "CreatePartitionsResponse",
+  // Starting in version 1, on quota violation, brokers send out responses before throttling.
+  //
+  // Version 2 adds flexible version support
+  //
+  // Version 3 is identical to version 2 but may return a THROTTLING_QUOTA_EXCEEDED error
+  // in the response if the partitions creation is throttled (KIP-599).
+  "validVersions": "0-3",
+  "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": "[]CreatePartitionsTopicResult", "versions": "0+",
+      "about": "The partition creation results for each topic.", "fields": [
+      { "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."},
+      { "name": "ErrorMessage", "type": "string", "versions": "0+", "nullableVersions": "0+",
+        "default": "null", "about": "The result message, or null 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
new file mode 100644
index 0000000000000..c80add64a7b48
--- /dev/null
+++ b/clients/src/main/resources/common/message/CreateTopicsRequest.json
@@ -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.
+
+{
+  "apiKey": 19,
+  "type": "request",
+  "name": "CreateTopicsRequest",
+  // Version 1 adds validateOnly.
+  //
+  // 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).
+  //
+  // Version 6 is identical to version 5 but may return a THROTTLING_QUOTA_EXCEEDED error
+  // in the response if the topics creation is throttled (KIP-599).
+  "validVersions": "0-6",
+  "flexibleVersions": "5+",
+  "fields": [
+    { "name": "Topics", "type": "[]CreatableTopic", "versions": "0+",
+      "about": "The topics to create.", "fields": [
+      { "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 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 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+", "entityType": "brokerId",
+          "about": "The brokers to place the partition on." }
+      ]},
+      { "name": "Configs", "type": "[]CreateableTopicConfig", "versions": "0+",
+        "about": "The custom topic configurations to set.", "fields": [
+        { "name": "Name", "type": "string", "versions": "0+" , "mapKey": true,
+          "about": "The configuration name." },
+        { "name": "Value", "type": "string", "versions": "0+", "nullableVersions": "0+",
+          "about": "The configuration value." }
+      ]}
+    ]},
+    { "name": "timeoutMs", "type": "int32", "versions": "0+", "default": "60000",
+      "about": "How long to wait in milliseconds before timing out the request." },
+    { "name": "validateOnly", "type": "bool", "versions": "1+", "default": "false", "ignorable": false,
+      "about": "If true, check that the topics can be created as specified, but don't create anything." }
+  ]
+}
diff --git a/clients/src/main/resources/common/message/CreateTopicsResponse.json b/clients/src/main/resources/common/message/CreateTopicsResponse.json
new file mode 100644
index 0000000000000..9f2296124a783
--- /dev/null
+++ b/clients/src/main/resources/common/message/CreateTopicsResponse.json
@@ -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.
+
+{
+  "apiKey": 19,
+  "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.
+  //
+  // 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).
+  //
+  // Version 6 is identical to version 5 but may return a THROTTLING_QUOTA_EXCEEDED error
+  // in the response if the topics creation is throttled (KIP-599).
+  "validVersions": "0-6",
+  "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, "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." },
+      { "name": "TopicConfigErrorCode", "type": "int16", "versions": "5+", "tag": 0, "taggedVersions": "5+", "ignorable": true,
+        "about": "Optional topic config error returned if configs are not returned in the response." },
+      { "name": "NumPartitions", "type": "int32", "versions": "5+", "default": "-1", "ignorable": true,
+        "about": "Number of partitions of the topic." },
+      { "name": "ReplicationFactor", "type": "int16", "versions": "5+", "default": "-1", "ignorable": true,
+        "about": "Replication factor of the topic." },
+      { "name": "Configs", "type": "[]CreatableTopicConfigs", "versions": "5+", "nullableVersions": "5+", "ignorable": true,
+        "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/DefaultPrincipalData.json b/clients/src/main/resources/common/message/DefaultPrincipalData.json
new file mode 100644
index 0000000000000..e06295d1783f8
--- /dev/null
+++ b/clients/src/main/resources/common/message/DefaultPrincipalData.json
@@ -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.
+
+{
+  "type": "data",
+  "name": "DefaultPrincipalData",
+  // The encoding format for default Kafka principal in
+  // org.apache.kafka.common.security.authenticator.DefaultKafkaPrincipalBuilder.
+  "validVersions": "0",
+  "flexibleVersions": "0+",
+  "fields": [
+    {"name": "Type", "type": "string", "versions": "0+",
+      "about": "The principal type"},
+    {"name": "Name", "type": "string", "versions": "0+",
+      "about": "The principal name"},
+    {"name": "TokenAuthenticated", "type": "bool", "versions": "0+",
+      "about": "Whether the principal was authenticated by a delegation token on the forwarding broker."}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/DeleteAclsRequest.json b/clients/src/main/resources/common/message/DeleteAclsRequest.json
new file mode 100644
index 0000000000000..664737e5810f7
--- /dev/null
+++ b/clients/src/main/resources/common/message/DeleteAclsRequest.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": 31,
+  "type": "request",
+  "name": "DeleteAclsRequest",
+  // Version 1 adds the pattern type.
+  // Version 2 enables flexible versions.
+  "validVersions": "0-2",
+  "flexibleVersions": "2+",
+  "fields": [
+    { "name": "Filters", "type": "[]DeleteAclsFilter", "versions": "0+",
+      "about": "The filters to use when deleting ACLs.", "fields": [
+      { "name": "ResourceTypeFilter", "type": "int8", "versions": "0+",
+        "about": "The resource type." },
+      { "name": "ResourceNameFilter", "type": "string", "versions": "0+", "nullableVersions": "0+",
+        "about": "The resource name." },
+      { "name": "PatternTypeFilter", "type": "int8", "versions": "1+", "default": "3", "ignorable": false,
+        "about": "The pattern type." },
+      { "name": "PrincipalFilter", "type": "string", "versions": "0+", "nullableVersions": "0+",
+        "about": "The principal filter, or null to accept all principals." },
+      { "name": "HostFilter", "type": "string", "versions": "0+", "nullableVersions": "0+",
+        "about": "The host filter, or null to accept all hosts." },
+      { "name": "Operation", "type": "int8", "versions": "0+",
+        "about": "The ACL operation." },
+      { "name": "PermissionType", "type": "int8", "versions": "0+",
+        "about": "The permission type." }
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/DeleteAclsResponse.json b/clients/src/main/resources/common/message/DeleteAclsResponse.json
new file mode 100644
index 0000000000000..08f570283e153
--- /dev/null
+++ b/clients/src/main/resources/common/message/DeleteAclsResponse.json
@@ -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.
+
+{
+  "apiKey": 31,
+  "type": "response",
+  "name": "DeleteAclsResponse",
+  // Version 1 adds the resource pattern type.
+  // Starting in version 1, on quota violation, brokers send out responses before throttling.
+  // Version 2 enables flexible versions.
+  "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": "FilterResults", "type": "[]DeleteAclsFilterResult", "versions": "0+",
+      "about": "The results for each filter.", "fields": [
+      { "name": "ErrorCode", "type": "int16", "versions": "0+",
+        "about": "The error code, or 0 if the filter succeeded." },
+      { "name": "ErrorMessage", "type": "string", "versions": "0+", "nullableVersions": "0+",
+        "about": "The error message, or null if the filter succeeded." },
+      { "name": "MatchingAcls", "type": "[]DeleteAclsMatchingAcl", "versions": "0+",
+        "about": "The ACLs which matched this filter.", "fields": [
+        { "name": "ErrorCode", "type": "int16", "versions": "0+",
+          "about": "The deletion error code, or 0 if the deletion succeeded." },
+        { "name": "ErrorMessage", "type": "string", "versions": "0+", "nullableVersions": "0+",
+          "about": "The deletion error message, or null if the deletion succeeded." },
+        { "name": "ResourceType", "type": "int8", "versions": "0+",
+          "about": "The ACL resource type." },
+        { "name": "ResourceName", "type": "string", "versions": "0+",
+          "about": "The ACL resource name." },
+        { "name": "PatternType", "type": "int8", "versions": "1+", "default": "3", "ignorable": false,
+          "about": "The ACL resource pattern type." },
+        { "name": "Principal", "type": "string", "versions": "0+",
+          "about": "The ACL principal." },
+        { "name": "Host", "type": "string", "versions": "0+",
+          "about": "The ACL host." },
+        { "name": "Operation", "type": "int8", "versions": "0+",
+          "about": "The ACL operation." },
+        { "name": "PermissionType", "type": "int8", "versions": "0+",
+          "about": "The ACL permission type." }
+      ]}
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/DeleteGroupsRequest.json b/clients/src/main/resources/common/message/DeleteGroupsRequest.json
new file mode 100644
index 0000000000000..833ed7a4dc6d7
--- /dev/null
+++ b/clients/src/main/resources/common/message/DeleteGroupsRequest.json
@@ -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.
+
+{
+  "apiKey": 42,
+  "type": "request",
+  "name": "DeleteGroupsRequest",
+  // Version 1 is the same as version 0.
+  //
+  // Version 2 is the first flexible version.
+  "validVersions": "0-2",
+  "flexibleVersions": "2+",
+  "fields": [
+    { "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
new file mode 100644
index 0000000000000..37e06a55b9913
--- /dev/null
+++ b/clients/src/main/resources/common/message/DeleteGroupsResponse.json
@@ -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.
+
+{
+  "apiKey": 42,
+  "type": "response",
+  "name": "DeleteGroupsResponse",
+  // Starting in version 1, on quota violation, brokers send out responses before throttling.
+  //
+  // 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, "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
new file mode 100644
index 0000000000000..93cbd56d50a6c
--- /dev/null
+++ b/clients/src/main/resources/common/message/DeleteRecordsRequest.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": 21,
+  "type": "request",
+  "name": "DeleteRecordsRequest",
+  // Version 1 is the same as version 0.
+
+  // Version 2 is the first flexible version.
+  "validVersions": "0-2",
+  "flexibleVersions": "2+",
+  "fields": [
+    { "name": "Topics", "type": "[]DeleteRecordsTopic", "versions": "0+",
+      "about": "Each topic that we want to delete records from.", "fields": [
+      { "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": [
+        { "name": "PartitionIndex", "type": "int32", "versions": "0+",
+          "about": "The partition index." },
+        { "name": "Offset", "type": "int64", "versions": "0+",
+          "about": "The deletion offset." }
+      ]}
+    ]},
+    { "name": "TimeoutMs", "type": "int32", "versions": "0+",
+      "about": "How long to wait for the deletion to complete, in milliseconds." }
+  ]
+}
diff --git a/clients/src/main/resources/common/message/DeleteRecordsResponse.json b/clients/src/main/resources/common/message/DeleteRecordsResponse.json
new file mode 100644
index 0000000000000..bfc0a56390879
--- /dev/null
+++ b/clients/src/main/resources/common/message/DeleteRecordsResponse.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": 21,
+  "type": "response",
+  "name": "DeleteRecordsResponse",
+  // Starting in version 1, on quota violation, brokers send out responses before throttling.
+
+  // 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": "Topics", "type": "[]DeleteRecordsTopicResult", "versions": "0+",
+      "about": "Each topic that we wanted to delete records from.", "fields": [
+      { "name": "Name", "type": "string", "versions": "0+", "mapKey": true, "entityType": "topicName",
+        "about": "The topic name." },
+      { "name": "Partitions", "type": "[]DeleteRecordsPartitionResult", "versions": "0+",
+        "about": "Each partition that we wanted to delete records from.", "fields": [
+        { "name": "PartitionIndex", "type": "int32", "versions": "0+", "mapKey": true,
+          "about": "The partition index." },
+        { "name": "LowWatermark", "type": "int64", "versions": "0+",
+          "about": "The partition low water mark." },
+        { "name": "ErrorCode", "type": "int16", "versions": "0+",
+          "about": "The deletion error code, or 0 if the deletion succeeded." }
+      ]}
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/DeleteTopicsRequest.json b/clients/src/main/resources/common/message/DeleteTopicsRequest.json
new file mode 100644
index 0000000000000..9dfba16515633
--- /dev/null
+++ b/clients/src/main/resources/common/message/DeleteTopicsRequest.json
@@ -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.
+
+{
+  "apiKey": 20,
+  "type": "request",
+  "name": "DeleteTopicsRequest",
+  // Versions 0, 1, 2, and 3 are the same.
+  //
+  // Version 4 is the first flexible version.
+  //
+  // Version 5 adds ErrorMessage in the response and may return a THROTTLING_QUOTA_EXCEEDED error
+  // in the response if the topics deletion is throttled (KIP-599).
+  "validVersions": "0-5",
+  "flexibleVersions": "4+",
+  "fields": [
+    { "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
new file mode 100644
index 0000000000000..6986b0adb339a
--- /dev/null
+++ b/clients/src/main/resources/common/message/DeleteTopicsResponse.json
@@ -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.
+
+{
+  "apiKey": 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.
+  //
+  // Version 4 is the first flexible version.
+  //
+  // Version 5 adds ErrorMessage in the response and may return a THROTTLING_QUOTA_EXCEEDED error
+  // in the response if the topics deletion is throttled (KIP-599).
+  "validVersions": "0-5",
+  "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": "Responses", "type": "[]DeletableTopicResult", "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." },
+      { "name": "ErrorMessage", "type": "string", "versions": "5+", "nullableVersions": "5+", "ignorable": true, "default": "null",
+        "about": "The error message, or null if there was no error." }
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/DescribeAclsRequest.json b/clients/src/main/resources/common/message/DescribeAclsRequest.json
new file mode 100644
index 0000000000000..a9c36768c1447
--- /dev/null
+++ b/clients/src/main/resources/common/message/DescribeAclsRequest.json
@@ -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.
+
+{
+  "apiKey": 29,
+  "type": "request",
+  "name": "DescribeAclsRequest",
+  // Version 1 adds resource pattern type.
+  // Version 2 enables flexible versions.
+  "validVersions": "0-2",
+  "flexibleVersions": "2+",
+  "fields": [
+    { "name": "ResourceTypeFilter", "type": "int8", "versions": "0+",
+      "about": "The resource type." },
+    { "name": "ResourceNameFilter", "type": "string", "versions": "0+", "nullableVersions": "0+",
+      "about": "The resource name, or null to match any resource name." },
+    { "name": "PatternTypeFilter", "type": "int8", "versions": "1+", "default": "3", "ignorable": false,
+      "about": "The resource pattern to match." },
+    { "name": "PrincipalFilter", "type": "string", "versions": "0+", "nullableVersions": "0+",
+      "about": "The principal to match, or null to match any principal." },
+    { "name": "HostFilter", "type": "string", "versions": "0+", "nullableVersions": "0+",
+      "about": "The host to match, or null to match any host." },
+    { "name": "Operation", "type": "int8", "versions": "0+",
+      "about": "The operation to match." },
+    { "name": "PermissionType", "type": "int8", "versions": "0+",
+      "about": "The permission type to match." }
+  ]
+}
diff --git a/clients/src/main/resources/common/message/DescribeAclsResponse.json b/clients/src/main/resources/common/message/DescribeAclsResponse.json
new file mode 100644
index 0000000000000..0ae72d67c4641
--- /dev/null
+++ b/clients/src/main/resources/common/message/DescribeAclsResponse.json
@@ -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.
+
+{
+  "apiKey": 29,
+  "type": "response",
+  "name": "DescribeAclsResponse",
+  // Version 1 adds PatternType.
+  // Starting in version 1, on quota violation, brokers send out responses before throttling.
+  // Version 2 enables flexible versions.
+  "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": "0+",
+      "about": "The error code, or 0 if there was no error." },
+    { "name": "ErrorMessage", "type": "string", "versions": "0+", "nullableVersions": "0+",
+      "about": "The error message, or null if there was no error." },
+    { "name": "Resources", "type": "[]DescribeAclsResource", "versions": "0+",
+      "about": "Each Resource that is referenced in an ACL.", "fields": [
+      { "name": "ResourceType", "type": "int8", "versions": "0+",
+        "about": "The resource type." },
+      { "name": "ResourceName", "type": "string", "versions": "0+",
+        "about": "The resource name." },
+      { "name": "PatternType", "type": "int8", "versions": "1+", "default": "3", "ignorable": false,
+        "about": "The resource pattern type." },
+      { "name": "Acls", "type": "[]AclDescription", "versions": "0+",
+        "about": "The ACLs.", "fields": [
+        { "name": "Principal", "type": "string", "versions": "0+",
+          "about": "The ACL principal." },
+        { "name": "Host", "type": "string", "versions": "0+",
+          "about": "The ACL host." },
+        { "name": "Operation", "type": "int8", "versions": "0+",
+          "about": "The ACL operation." },
+        { "name": "PermissionType", "type": "int8", "versions": "0+",
+          "about": "The ACL permission type." }
+      ]}
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/DescribeClientQuotasRequest.json b/clients/src/main/resources/common/message/DescribeClientQuotasRequest.json
new file mode 100644
index 0000000000000..e35d3b0c91a01
--- /dev/null
+++ b/clients/src/main/resources/common/message/DescribeClientQuotasRequest.json
@@ -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.
+
+{
+  "apiKey": 48,
+  "type": "request",
+  "name": "DescribeClientQuotasRequest",
+  // Version 1 enables flexible versions.
+  "validVersions": "0-1",
+  "flexibleVersions": "1+",
+  "fields": [
+    { "name": "Components", "type": "[]ComponentData", "versions": "0+",
+      "about": "Filter components to apply to quota entities.", "fields": [
+      { "name": "EntityType", "type": "string", "versions": "0+",
+        "about": "The entity type that the filter component applies to." },
+      { "name": "MatchType", "type": "int8", "versions": "0+",
+        "about": "How to match the entity {0 = exact name, 1 = default name, 2 = any specified name}." },
+      { "name": "Match", "type": "string", "versions": "0+", "nullableVersions": "0+",
+        "about": "The string to match against, or null if unused for the match type." }
+    ]},
+    { "name": "Strict", "type": "bool", "versions": "0+",
+      "about": "Whether the match is strict, i.e. should exclude entities with unspecified entity types." }
+  ]
+}
diff --git a/clients/src/main/resources/common/message/DescribeClientQuotasResponse.json b/clients/src/main/resources/common/message/DescribeClientQuotasResponse.json
new file mode 100644
index 0000000000000..0dd0c9c7bf831
--- /dev/null
+++ b/clients/src/main/resources/common/message/DescribeClientQuotasResponse.json
@@ -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.
+{
+  "apiKey": 48,
+  "type": "response",
+  "name": "DescribeClientQuotasResponse",
+  // Version 1 enables flexible versions.
+  "validVersions": "0-1",
+  "flexibleVersions": "1+",
+  "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 error code, or `0` if the quota description succeeded." },
+    { "name": "ErrorMessage", "type": "string", "versions": "0+", "nullableVersions": "0+",
+      "about": "The error message, or `null` if the quota description succeeded." },
+    { "name": "Entries", "type": "[]EntryData", "versions": "0+", "nullableVersions": "0+",
+      "about": "A result entry.", "fields": [
+      { "name": "Entity", "type": "[]EntityData", "versions": "0+",
+        "about": "The quota entity description.", "fields": [
+        { "name": "EntityType", "type": "string", "versions": "0+",
+          "about": "The entity type." },
+        { "name": "EntityName", "type": "string", "versions": "0+", "nullableVersions": "0+",
+          "about": "The entity name, or null if the default." }
+      ]},
+      { "name": "Values", "type": "[]ValueData", "versions": "0+",
+	"about": "The quota values for the entity.", "fields": [
+        { "name": "Key", "type": "string", "versions": "0+",
+          "about": "The quota configuration key." },
+        { "name": "Value", "type": "float64", "versions": "0+",
+          "about": "The quota configuration value." }
+      ]}
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/DescribeConfigsRequest.json b/clients/src/main/resources/common/message/DescribeConfigsRequest.json
new file mode 100644
index 0000000000000..01a45eb2a5beb
--- /dev/null
+++ b/clients/src/main/resources/common/message/DescribeConfigsRequest.json
@@ -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.
+
+{
+  "apiKey": 32,
+  "type": "request",
+  "name": "DescribeConfigsRequest",
+  // Version 1 adds IncludeSynonyms.
+  // Version 2 is the same as version 1.
+  // Version 4 enables flexible versions.
+  "validVersions": "0-4",
+  "flexibleVersions": "4+",
+  "fields": [
+    { "name": "Resources", "type": "[]DescribeConfigsResource", "versions": "0+",
+      "about": "The resources whose configurations we want to describe.", "fields": [
+      { "name": "ResourceType", "type": "int8", "versions": "0+",
+        "about": "The resource type." },
+      { "name": "ResourceName", "type": "string", "versions": "0+",
+        "about": "The resource name." },
+      { "name": "ConfigurationKeys", "type": "[]string", "versions": "0+", "nullableVersions": "0+",
+        "about": "The configuration keys to list, or null to list all configuration keys." }
+    ]},
+    { "name": "IncludeSynonyms", "type": "bool", "versions": "1+", "default": "false", "ignorable": false,
+      "about": "True if we should include all synonyms." },
+    { "name": "IncludeDocumentation", "type": "bool", "versions": "3+", "default": "false", "ignorable": false,
+      "about": "True if we should include configuration documentation." }
+  ]
+}
diff --git a/clients/src/main/resources/common/message/DescribeConfigsResponse.json b/clients/src/main/resources/common/message/DescribeConfigsResponse.json
new file mode 100644
index 0000000000000..f2f57ad1a7367
--- /dev/null
+++ b/clients/src/main/resources/common/message/DescribeConfigsResponse.json
@@ -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.
+
+{
+  "apiKey": 32,
+  "type": "response",
+  "name": "DescribeConfigsResponse",
+  // Version 1 adds ConfigSource and the synonyms.
+  // Starting in version 2, on quota violation, brokers send out responses before throttling.
+  // Version 4 enables flexible versions.
+  "validVersions": "0-4",
+  "flexibleVersions": "4+",
+  "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": "[]DescribeConfigsResult", "versions": "0+",
+      "about": "The results for each resource.", "fields": [
+      { "name": "ErrorCode", "type": "int16", "versions": "0+",
+        "about": "The error code, or 0 if we were able to successfully describe the configurations." },
+      { "name": "ErrorMessage", "type": "string", "versions": "0+", "nullableVersions": "0+",
+        "about": "The error message, or null if we were able to successfully describe the configurations." },
+      { "name": "ResourceType", "type": "int8", "versions": "0+",
+        "about": "The resource type." },
+      { "name": "ResourceName", "type": "string", "versions": "0+",
+        "about": "The resource name." },
+      { "name": "Configs", "type": "[]DescribeConfigsResourceResult", "versions": "0+",
+        "about": "Each listed configuration.", "fields": [
+        { "name": "Name", "type": "string", "versions": "0+",
+          "about": "The configuration name." },
+        { "name": "Value", "type": "string", "versions": "0+", "nullableVersions": "0+",
+          "about": "The configuration value." },
+        { "name": "ReadOnly", "type": "bool", "versions": "0+",
+          "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 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,
+          "about": "The configuration source." },
+        { "name": "IsSensitive", "type": "bool", "versions": "0+",
+          "about": "True if this configuration is sensitive." },
+        { "name": "Synonyms", "type": "[]DescribeConfigsSynonym", "versions": "1+", "ignorable": true,
+          "about": "The synonyms for this configuration key.", "fields": [
+          { "name": "Name", "type": "string", "versions": "1+",
+            "about": "The synonym name." },
+          { "name": "Value", "type": "string", "versions": "1+", "nullableVersions": "0+",
+            "about": "The synonym value." },
+          { "name": "Source", "type": "int8", "versions": "1+",
+            "about": "The synonym source." }
+        ]},
+        { "name": "ConfigType", "type": "int8", "versions": "3+", "default": "0", "ignorable": true,
+          "about": "The configuration data type. Type can be one of the following values - BOOLEAN, STRING, INT, SHORT, LONG, DOUBLE, LIST, CLASS, PASSWORD" },
+        { "name": "Documentation", "type": "string", "versions": "3+", "nullableVersions": "0+", "ignorable": true,
+          "about": "The configuration documentation." }
+      ]}
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/DescribeDelegationTokenRequest.json b/clients/src/main/resources/common/message/DescribeDelegationTokenRequest.json
new file mode 100644
index 0000000000000..32a4f5c5150c1
--- /dev/null
+++ b/clients/src/main/resources/common/message/DescribeDelegationTokenRequest.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": 41,
+  "type": "request",
+  "name": "DescribeDelegationTokenRequest",
+  // Version 1 is the same as version 0.
+  // Version 2 adds flexible version support
+  "validVersions": "0-2",
+  "flexibleVersions": "2+",
+  "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": [
+      { "name": "PrincipalType", "type": "string", "versions": "0+",
+        "about": "The owner principal type." },
+      { "name": "PrincipalName", "type": "string", "versions": "0+",
+        "about": "The owner principal name." }
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/DescribeDelegationTokenResponse.json b/clients/src/main/resources/common/message/DescribeDelegationTokenResponse.json
new file mode 100644
index 0000000000000..09f69ce61c30b
--- /dev/null
+++ b/clients/src/main/resources/common/message/DescribeDelegationTokenResponse.json
@@ -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.
+
+{
+  "apiKey": 41,
+  "type": "response",
+  "name": "DescribeDelegationTokenResponse",
+  // Starting in version 1, on quota violation, brokers send out responses before throttling.
+  // Version 2 adds flexible version support
+  "validVersions": "0-2",
+  "flexibleVersions": "2+",
+  "fields": [
+    { "name": "ErrorCode", "type": "int16", "versions": "0+",
+      "about": "The error code, or 0 if there was no error." },
+    { "name": "Tokens", "type": "[]DescribedDelegationToken", "versions": "0+",
+      "about": "The tokens.", "fields": [
+      { "name": "PrincipalType", "type": "string", "versions": "0+",
+        "about": "The token principal type." },
+      { "name": "PrincipalName", "type": "string", "versions": "0+",
+        "about": "The token principal name." },
+      { "name": "IssueTimestamp", "type": "int64", "versions": "0+",
+        "about": "The token issue timestamp in milliseconds." },
+      { "name": "ExpiryTimestamp", "type": "int64", "versions": "0+",
+        "about": "The token expiry timestamp in milliseconds." },
+      { "name": "MaxTimestamp", "type": "int64", "versions": "0+",
+        "about": "The token maximum timestamp length in milliseconds." },
+      { "name": "TokenId", "type": "string", "versions": "0+",
+        "about": "The token ID." },
+      { "name": "Hmac", "type": "bytes", "versions": "0+",
+        "about": "The token HMAC." },
+      { "name": "Renewers", "type": "[]DescribedDelegationTokenRenewer", "versions": "0+",
+        "about": "Those who are able to renew this token before it expires.", "fields": [
+        { "name": "PrincipalType", "type": "string", "versions": "0+",
+          "about": "The renewer principal type" },
+        { "name": "PrincipalName", "type": "string", "versions": "0+",
+          "about": "The renewer principal name" }
+      ]}
+    ]},
+    { "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/DescribeGroupsRequest.json b/clients/src/main/resources/common/message/DescribeGroupsRequest.json
new file mode 100644
index 0000000000000..8a5887a5680a8
--- /dev/null
+++ b/clients/src/main/resources/common/message/DescribeGroupsRequest.json
@@ -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.
+
+{
+  "apiKey": 15,
+  "type": "request",
+  "name": "DescribeGroupsRequest",
+  // Versions 1 and 2 are the same as version 0.
+  //
+  // Starting in version 3, authorized operations can be requested.
+  //
+  // 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+", "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
new file mode 100644
index 0000000000000..f195843ce32b6
--- /dev/null
+++ b/clients/src/main/resources/common/message/DescribeGroupsResponse.json
@@ -0,0 +1,70 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements.  See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+{
+  "apiKey": 15,
+  "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.
+  //
+  // Starting in version 4, the response will optionally 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." },
+    { "name": "Groups", "type": "[]DescribedGroup", "versions": "0+",
+      "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+", "entityType": "groupId",
+        "about": "The group ID string." },
+      { "name": "GroupState", "type": "string", "versions": "0+",
+        "about": "The group state string, or the empty string." },
+      { "name": "ProtocolType", "type": "string", "versions": "0+",
+        "about": "The group protocol type, or the empty string." },
+      // ProtocolData is currently only filled in if the group state is in the Stable state.
+      { "name": "ProtocolData", "type": "string", "versions": "0+",
+        "about": "The group protocol data, or the empty string." },
+      // N.B. If the group is in the Dead state, the members array will always be empty.
+      { "name": "Members", "type": "[]DescribedGroupMember", "versions": "0+",
+        "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+", "ignorable": true,
+          "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+",
+          "about": "The client host." },
+        // This is currently only provided if the group is in the Stable state.
+        { "name": "MemberMetadata", "type": "bytes", "versions": "0+",
+          "about": "The metadata corresponding to the current group protocol in use." },
+        // This is currently only provided if the group is in the Stable state.
+        { "name": "MemberAssignment", "type": "bytes", "versions": "0+",
+          "about": "The current assignment provided by the group leader." }
+      ]},
+      { "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
new file mode 100644
index 0000000000000..577f2eb4cf6a8
--- /dev/null
+++ b/clients/src/main/resources/common/message/DescribeLogDirsRequest.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": 35,
+  "type": "request",
+  "name": "DescribeLogDirsRequest",
+  // Version 1 is the same as version 0.
+  "validVersions": "0-2",
+  // Version 2 is the first flexible version.
+  "flexibleVersions": "2+",
+  "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+", "entityType": "topicName", "mapKey": true,
+        "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
new file mode 100644
index 0000000000000..4322f1c52ee2f
--- /dev/null
+++ b/clients/src/main/resources/common/message/DescribeLogDirsResponse.json
@@ -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.
+
+{
+  "apiKey": 35,
+  "type": "response",
+  "name": "DescribeLogDirsResponse",
+  // Starting in version 1, on quota violation, brokers send out responses before throttling.
+  "validVersions": "0-2",
+  // Version 2 is the first flexible version.
+  "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": "[]DescribeLogDirsResult", "versions": "0+",
+      "about": "The log directories.", "fields": [
+      { "name": "ErrorCode", "type": "int16", "versions": "0+",
+        "about": "The error code, or 0 if there was no error." },
+      { "name": "LogDir", "type": "string", "versions": "0+",
+        "about": "The absolute log directory path." },
+      { "name": "Topics", "type": "[]DescribeLogDirsTopic", "versions": "0+",
+        "about": "Each topic.", "fields": [
+        { "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+",
+            "about": "The partition index." },
+          { "name": "PartitionSize", "type": "int64", "versions": "0+",
+            "about": "The size of the log segments in this partition in bytes." },
+          { "name": "OffsetLag", "type": "int64", "versions": "0+",
+            "about": "The lag of the log's LEO w.r.t. partition's HW (if it is the current log for the partition) or current replica's LEO (if it is the future log for the partition)" },
+          { "name": "IsFutureKey", "type": "bool", "versions": "0+",
+            "about": "True if this log is created by AlterReplicaLogDirsRequest and will replace the current log of the replica in the future." }
+        ]}
+      ]}
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/DescribeQuorumRequest.json b/clients/src/main/resources/common/message/DescribeQuorumRequest.json
new file mode 100644
index 0000000000000..f91d679db8b1b
--- /dev/null
+++ b/clients/src/main/resources/common/message/DescribeQuorumRequest.json
@@ -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.
+
+{
+  "apiKey": 55,
+  "type": "request",
+  "name": "DescribeQuorumRequest",
+  "validVersions": "0",
+  "flexibleVersions": "0+",
+  "fields": [
+    { "name": "Topics", "type": "[]TopicData",
+      "versions": "0+", "fields": [
+      { "name": "TopicName", "type": "string", "versions": "0+", "entityType": "topicName",
+        "about": "The topic name." },
+      { "name": "Partitions", "type": "[]PartitionData",
+        "versions": "0+", "fields": [
+        { "name": "PartitionIndex", "type": "int32", "versions": "0+",
+          "about": "The partition index." }
+      ]
+      }]
+    }
+  ]
+}
diff --git a/clients/src/main/resources/common/message/DescribeQuorumResponse.json b/clients/src/main/resources/common/message/DescribeQuorumResponse.json
new file mode 100644
index 0000000000000..98fbb43c391e5
--- /dev/null
+++ b/clients/src/main/resources/common/message/DescribeQuorumResponse.json
@@ -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.
+
+{
+  "apiKey": 55,
+  "type": "response",
+  "name": "DescribeQuorumResponse",
+  "validVersions": "0",
+  "flexibleVersions": "0+",
+  "fields": [
+    { "name": "ErrorCode", "type": "int16", "versions": "0+",
+      "about": "The top level error code."},
+    { "name": "Topics", "type": "[]TopicData",
+      "versions": "0+", "fields": [
+      { "name": "TopicName", "type": "string", "versions": "0+", "entityType": "topicName",
+        "about": "The topic name." },
+      { "name": "Partitions", "type": "[]PartitionData",
+        "versions": "0+", "fields": [
+        { "name": "PartitionIndex", "type": "int32", "versions": "0+",
+          "about": "The partition index." },
+        { "name": "ErrorCode", "type": "int16", "versions": "0+"},
+        { "name": "LeaderId", "type": "int32", "versions": "0+",
+          "about": "The ID of the current leader or -1 if the leader is unknown."},
+        { "name": "LeaderEpoch", "type": "int32", "versions": "0+",
+          "about": "The latest known leader epoch"},
+        { "name": "HighWatermark", "type": "int64", "versions": "0+"},
+        { "name": "CurrentVoters", "type": "[]ReplicaState", "versions": "0+" },
+        { "name": "Observers", "type": "[]ReplicaState", "versions": "0+" }
+      ]}
+    ]}],
+  "commonStructs": [
+    { "name": "ReplicaState", "versions": "0+", "fields": [
+      { "name": "ReplicaId", "type": "int32", "versions": "0+"},
+      { "name": "LogEndOffset", "type": "int64", "versions": "0+",
+        "about": "The last known log end offset of the follower or -1 if it is unknown"}
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/DescribeUserScramCredentialsRequest.json b/clients/src/main/resources/common/message/DescribeUserScramCredentialsRequest.json
new file mode 100644
index 0000000000000..f7f8c68991357
--- /dev/null
+++ b/clients/src/main/resources/common/message/DescribeUserScramCredentialsRequest.json
@@ -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.
+
+{
+  "apiKey": 50,
+  "type": "request",
+  "name": "DescribeUserScramCredentialsRequest",
+  "validVersions": "0",
+  "flexibleVersions": "0+",
+  "fields": [
+    { "name": "Users", "type": "[]UserName", "versions": "0+", "nullableVersions": "0+",
+      "about": "The users to describe, or null/empty to describe all users.", "fields": [
+      { "name": "Name", "type": "string", "versions": "0+",
+        "about": "The user name." }
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/DescribeUserScramCredentialsResponse.json b/clients/src/main/resources/common/message/DescribeUserScramCredentialsResponse.json
new file mode 100644
index 0000000000000..9e8b035281cd8
--- /dev/null
+++ b/clients/src/main/resources/common/message/DescribeUserScramCredentialsResponse.json
@@ -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.
+
+{
+  "apiKey": 50,
+  "type": "response",
+  "name": "DescribeUserScramCredentialsResponse",
+  "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 message-level error code, 0 except for user authorization or infrastructure issues." },
+    { "name": "ErrorMessage", "type": "string", "versions": "0+", "nullableVersions": "0+",
+      "about": "The message-level error message, if any." },
+    { "name": "Results", "type": "[]DescribeUserScramCredentialsResult", "versions": "0+",
+      "about": "The results for descriptions, one per user.", "fields": [
+      { "name": "User", "type": "string", "versions": "0+",
+        "about": "The user name." },
+      { "name": "ErrorCode", "type": "int16", "versions": "0+",
+        "about": "The user-level error code." },
+      { "name": "ErrorMessage", "type": "string", "versions": "0+", "nullableVersions": "0+",
+        "about": "The user-level error message, if any." },
+      { "name": "CredentialInfos", "type": "[]CredentialInfo", "versions": "0+",
+        "about": "The mechanism and related information associated with the user's SCRAM credentials.", "fields": [
+        { "name": "Mechanism", "type": "int8", "versions": "0+",
+          "about": "The SCRAM mechanism." },
+        { "name": "Iterations", "type": "int32", "versions": "0+",
+          "about": "The number of iterations used in the SCRAM credential." }]}
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/ElectLeadersRequest.json b/clients/src/main/resources/common/message/ElectLeadersRequest.json
new file mode 100644
index 0000000000000..5b86c96b04d60
--- /dev/null
+++ b/clients/src/main/resources/common/message/ElectLeadersRequest.json
@@ -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.
+
+{
+  "apiKey": 43,
+  "type": "request",
+  "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 leaders.",
+      "fields": [
+        { "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 leader should be elected." }
+      ]
+    },
+    { "name": "TimeoutMs", "type": "int32", "versions": "0+", "default": "60000",
+      "about": "The time in ms to wait for the election to complete." }
+  ]
+}
diff --git a/clients/src/main/resources/common/message/ElectLeadersResponse.json b/clients/src/main/resources/common/message/ElectLeadersResponse.json
new file mode 100644
index 0000000000000..15468c78d1f70
--- /dev/null
+++ b/clients/src/main/resources/common/message/ElectLeadersResponse.json
@@ -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.
+
+{
+  "apiKey": 43,
+  "type": "response",
+  "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 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": [
+        { "name": "PartitionId", "type": "int32", "versions": "0+",
+          "about": "The partition id" },
+        { "name": "ErrorCode", "type": "int16", "versions": "0+",
+          "about": "The result error, or zero if there was no error."},
+        { "name": "ErrorMessage", "type": "string", "versions": "0+", "nullableVersions": "0+",
+          "about": "The result message, or null if there was no error."}
+      ]}
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/EndQuorumEpochRequest.json b/clients/src/main/resources/common/message/EndQuorumEpochRequest.json
new file mode 100644
index 0000000000000..45dacde5fe99b
--- /dev/null
+++ b/clients/src/main/resources/common/message/EndQuorumEpochRequest.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": 54,
+  "type": "request",
+  "name": "EndQuorumEpochRequest",
+  "validVersions": "0",
+  "fields": [
+    { "name": "ClusterId", "type": "string", "versions": "0+",
+      "nullableVersions": "0+", "default": "null"},
+    { "name": "Topics", "type": "[]TopicData",
+      "versions": "0+", "fields": [
+      { "name": "TopicName", "type": "string", "versions": "0+", "entityType": "topicName",
+        "about": "The topic name." },
+      { "name": "Partitions", "type": "[]PartitionData",
+        "versions": "0+", "fields": [
+        { "name": "PartitionIndex", "type": "int32", "versions": "0+",
+          "about": "The partition index." },
+        { "name": "LeaderId", "type": "int32", "versions": "0+",
+          "about": "The current leader ID that is resigning"},
+        { "name": "LeaderEpoch", "type": "int32", "versions": "0+",
+          "about": "The current epoch"},
+        { "name": "PreferredSuccessors", "type": "[]int32", "versions": "0+",
+          "about": "A sorted list of preferred successors to start the election"}
+      ]}
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/EndQuorumEpochResponse.json b/clients/src/main/resources/common/message/EndQuorumEpochResponse.json
new file mode 100644
index 0000000000000..29d24d0b6f401
--- /dev/null
+++ b/clients/src/main/resources/common/message/EndQuorumEpochResponse.json
@@ -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.
+
+{
+  "apiKey": 54,
+  "type": "response",
+  "name": "EndQuorumEpochResponse",
+  "validVersions": "0",
+  "fields": [
+    { "name": "ErrorCode", "type": "int16", "versions": "0+",
+      "about": "The top level error code."},
+    { "name": "Topics", "type": "[]TopicData",
+      "versions": "0+", "fields": [
+      { "name": "TopicName", "type": "string", "versions": "0+", "entityType": "topicName",
+        "about": "The topic name." },
+      { "name": "Partitions", "type": "[]PartitionData",
+        "versions": "0+", "fields": [
+        { "name": "PartitionIndex", "type": "int32", "versions": "0+",
+          "about": "The partition index." },
+        { "name": "ErrorCode", "type": "int16", "versions": "0+"},
+        { "name": "LeaderId", "type": "int32", "versions": "0+",
+          "about": "The ID of the current leader or -1 if the leader is unknown."},
+        { "name": "LeaderEpoch", "type": "int32", "versions": "0+",
+          "about": "The latest known leader epoch"}
+      ]}
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/EndTxnRequest.json b/clients/src/main/resources/common/message/EndTxnRequest.json
new file mode 100644
index 0000000000000..de18b43b7d9d2
--- /dev/null
+++ b/clients/src/main/resources/common/message/EndTxnRequest.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": 26,
+  "type": "request",
+  "name": "EndTxnRequest",
+  // Version 1 is the same as version 0.
+  //
+  // Version 2 adds the support for new error code PRODUCER_FENCED.
+  //
+  // Version 3 enables flexible versions.
+  "validVersions": "0-3",
+  "flexibleVersions": "3+",
+  "fields": [
+    { "name": "TransactionalId", "type": "string", "versions": "0+", "entityType": "transactionalId",
+      "about": "The ID of the transaction to end." },
+    { "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." },
+    { "name": "Committed", "type": "bool", "versions": "0+",
+      "about": "True if the transaction was committed, false if it was aborted." }
+  ]
+}
diff --git a/clients/src/main/resources/common/message/EndTxnResponse.json b/clients/src/main/resources/common/message/EndTxnResponse.json
new file mode 100644
index 0000000000000..3071953185ffc
--- /dev/null
+++ b/clients/src/main/resources/common/message/EndTxnResponse.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": 26,
+  "type": "response",
+  "name": "EndTxnResponse",
+  // Starting in version 1, on quota violation, brokers send out responses before throttling.
+  //
+  // Version 2 adds the support for new error code PRODUCER_FENCED.
+  //
+  // Version 3 enables flexible versions.
+  "validVersions": "0-3",
+  "flexibleVersions": "3+",
+  "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 error code, or 0 if there was no error." }
+  ]
+}
diff --git a/clients/src/main/resources/common/message/EnvelopeRequest.json b/clients/src/main/resources/common/message/EnvelopeRequest.json
new file mode 100644
index 0000000000000..a1aa760a29b18
--- /dev/null
+++ b/clients/src/main/resources/common/message/EnvelopeRequest.json
@@ -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.
+
+{
+  "apiKey": 58,
+  "type": "request",
+  "name": "EnvelopeRequest",
+  // Request struct for forwarding.
+  "validVersions": "0",
+  "flexibleVersions": "0+",
+  "fields": [
+    { "name": "RequestData", "type": "bytes", "versions": "0+", "zeroCopy": true,
+      "about": "The embedded request header and data."},
+    { "name": "RequestPrincipal", "type": "bytes", "versions": "0+", "nullableVersions": "0+",
+      "about": "Value of the initial client principal when the request is redirected by a broker." },
+    { "name": "ClientHostAddress", "type": "bytes", "versions": "0+",
+      "about": "The original client's address in bytes." }
+  ]
+}
diff --git a/clients/src/main/resources/common/message/EnvelopeResponse.json b/clients/src/main/resources/common/message/EnvelopeResponse.json
new file mode 100644
index 0000000000000..900805284db12
--- /dev/null
+++ b/clients/src/main/resources/common/message/EnvelopeResponse.json
@@ -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.
+
+{
+  "apiKey": 58,
+  "type": "response",
+  "name": "EnvelopeResponse",
+  // Response struct for forwarding.
+  "validVersions": "0",
+  "flexibleVersions": "0+",
+  "fields": [
+    { "name": "ResponseData", "type": "bytes", "versions": "0+", "nullableVersions": "0+",
+      "zeroCopy": true, "default": "null",
+      "about": "The embedded response header and data."},
+    { "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/ExpireDelegationTokenRequest.json b/clients/src/main/resources/common/message/ExpireDelegationTokenRequest.json
new file mode 100644
index 0000000000000..a990862d2fc8a
--- /dev/null
+++ b/clients/src/main/resources/common/message/ExpireDelegationTokenRequest.json
@@ -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.
+
+{
+  "apiKey": 40,
+  "type": "request",
+  "name": "ExpireDelegationTokenRequest",
+  // Version 1 is the same as version 0.
+  // Version 2 adds flexible version support
+  "validVersions": "0-2",
+  "flexibleVersions": "2+",
+  "fields": [
+    { "name": "Hmac", "type": "bytes", "versions": "0+",
+      "about": "The HMAC of the delegation token to be expired." },
+    { "name": "ExpiryTimePeriodMs", "type": "int64", "versions": "0+",
+      "about": "The expiry time period in milliseconds." }
+  ]
+}
diff --git a/clients/src/main/resources/common/message/ExpireDelegationTokenResponse.json b/clients/src/main/resources/common/message/ExpireDelegationTokenResponse.json
new file mode 100644
index 0000000000000..f2d4bf48b6aeb
--- /dev/null
+++ b/clients/src/main/resources/common/message/ExpireDelegationTokenResponse.json
@@ -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.
+
+{
+  "apiKey": 40,
+  "type": "response",
+  "name": "ExpireDelegationTokenResponse",
+  // Starting in version 1, on quota violation, brokers send out responses before throttling.
+  // Version 2 adds flexible version support
+  "validVersions": "0-2",
+  "flexibleVersions": "2+",
+  "fields": [
+    { "name": "ErrorCode", "type": "int16", "versions": "0+",
+      "about": "The error code, or 0 if there was no error." },
+    { "name": "ExpiryTimestampMs", "type": "int64", "versions": "0+",
+      "about": "The timestamp in milliseconds at which this token expires." },
+    { "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/FetchRequest.json b/clients/src/main/resources/common/message/FetchRequest.json
new file mode 100644
index 0000000000000..0dcdd7af8bb00
--- /dev/null
+++ b/clients/src/main/resources/common/message/FetchRequest.json
@@ -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.
+
+{
+  "apiKey": 1,
+  "type": "request",
+  "name": "FetchRequest",
+  //
+  // Version 1 is the same as version 0.
+  //
+  // Starting in Version 2, the requestor must be able to handle Kafka Log
+  // Message format version 1.
+  //
+  // Version 3 adds MaxBytes.  Starting in version 3, the partition ordering in
+  // the request is now relevant.  Partitions will be processed in the order
+  // they appear in the request.
+  //
+  // Version 4 adds IsolationLevel.  Starting in version 4, the reqestor must be
+  // able to handle Kafka log message format version 2.
+  //
+  // Version 5 adds LogStartOffset to indicate the earliest available offset of
+  // partition data that can be consumed.
+  //
+  // Version 6 is the same as version 5.
+  //
+  // Version 7 adds incremental fetch request support.
+  //
+  // Version 8 is the same as version 7.
+  //
+  // Version 9 adds CurrentLeaderEpoch, as described in KIP-320.
+  //
+  // Version 10 indicates that we can use the ZStd compression algorithm, as
+  // described in KIP-110.
+  // Version 12 adds flexible versions support as well as epoch validation through
+  // the `LastFetchedEpoch` field
+  //
+  "validVersions": "0-12",
+  "flexibleVersions": "12+",
+  "fields": [
+    { "name": "ClusterId", "type": "string", "versions": "12+", "nullableVersions": "12+", "default": "null", "taggedVersions": "12+", "tag": 0,
+      "about": "The clusterId if known. This is used to validate metadata fetches prior to broker registration." },
+    { "name": "ReplicaId", "type": "int32", "versions": "0+",
+      "about": "The broker ID of the follower, of -1 if this request is from a consumer." },
+    { "name": "MaxWaitMs", "type": "int32", "versions": "0+",
+      "about": "The maximum time in milliseconds to wait for the response." },
+    { "name": "MinBytes", "type": "int32", "versions": "0+",
+      "about": "The minimum bytes to accumulate in the response." },
+    { "name": "MaxBytes", "type": "int32", "versions": "3+", "default": "0x7fffffff", "ignorable": true,
+      "about": "The maximum bytes to fetch.  See KIP-74 for cases where this limit may not be honored." },
+    { "name": "IsolationLevel", "type": "int8", "versions": "4+", "default": "0", "ignorable": true,
+      "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": "SessionId", "type": "int32", "versions": "7+", "default": "0", "ignorable": true,
+      "about": "The fetch session ID." },
+    { "name": "SessionEpoch", "type": "int32", "versions": "7+", "default": "-1", "ignorable": true,
+      "about": "The fetch session epoch, which is used for ordering requests in a session." },
+    { "name": "Topics", "type": "[]FetchTopic", "versions": "0+",
+      "about": "The topics to fetch.", "fields": [
+      { "name": "Topic", "type": "string", "versions": "0+", "entityType": "topicName",
+        "about": "The name of the topic to fetch." },
+      { "name": "Partitions", "type": "[]FetchPartition", "versions": "0+",
+        "about": "The partitions to fetch.", "fields": [
+        { "name": "Partition", "type": "int32", "versions": "0+",
+          "about": "The partition index." },
+        { "name": "CurrentLeaderEpoch", "type": "int32", "versions": "9+", "default": "-1", "ignorable": true,
+          "about": "The current leader epoch of the partition." },
+        { "name": "FetchOffset", "type": "int64", "versions": "0+",
+          "about": "The message offset." },
+        { "name": "LastFetchedEpoch", "type": "int32", "versions": "12+", "default": "-1", "ignorable": false,
+          "about": "The epoch of the last fetched record or -1 if there is none"},
+        { "name": "LogStartOffset", "type": "int64", "versions": "5+", "default": "-1", "ignorable": true,
+          "about": "The earliest available offset of the follower replica.  The field is only used when the request is sent by the follower."},
+        { "name": "PartitionMaxBytes", "type": "int32", "versions": "0+",
+          "about": "The maximum bytes to fetch from this partition.  See KIP-74 for cases where this limit may not be honored." }
+      ]}
+    ]},
+    { "name": "ForgottenTopicsData", "type": "[]ForgottenTopic", "versions": "7+", "ignorable": false,
+      "about": "In an incremental fetch request, the partitions to remove.", "fields": [
+      { "name": "Topic", "type": "string", "versions": "7+", "entityType": "topicName",
+        "about": "The partition name." },
+      { "name": "Partitions", "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
new file mode 100644
index 0000000000000..3db4359b3aaf2
--- /dev/null
+++ b/clients/src/main/resources/common/message/FetchResponse.json
@@ -0,0 +1,94 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements.  See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT 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": 1,
+  "type": "response",
+  "name": "FetchResponse",
+  //
+  // Version 1 adds throttle time.
+  //
+  // Version 2 and 3 are the same as version 1. 
+  //
+  // Version 4 adds features for transactional consumption.
+  //
+  // Version 5 adds LogStartOffset to indicate the earliest available offset of
+  // partition data that can be consumed.
+  //
+  // Starting in version 6, we may return KAFKA_STORAGE_ERROR as an error code.
+  //
+  // Version 7 adds incremental fetch request support.
+  //
+  // Starting in version 8, on quota violation, brokers send out responses before throttling.
+  //
+  // Version 9 is the same as version 8.
+  //
+  // Version 10 indicates that the response data can use the ZStd compression
+  // algorithm, as described in KIP-110.
+  // Version 12 adds support for flexible versions, epoch detection through the `TruncationOffset` field,
+  // and leader discovery through the `CurrentLeader` field
+  //
+  "validVersions": "0-12",
+  "flexibleVersions": "12+",
+  "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": "7+", "ignorable": true,
+      "about": "The top level response error code." },
+    { "name": "SessionId", "type": "int32", "versions": "7+", "default": "0", "ignorable": false,
+      "about": "The fetch session ID, or 0 if this is not part of a fetch session." },
+    { "name": "Responses", "type": "[]FetchableTopicResponse", "versions": "0+",
+      "about": "The response topics.", "fields": [
+      { "name": "Topic", "type": "string", "versions": "0+", "entityType": "topicName",
+        "about": "The topic name." },
+      { "name": "PartitionResponses", "type": "[]FetchablePartitionResponse", "versions": "0+",
+        "about": "The topic partitions.", "fields": [
+        { "name": "Partition", "type": "int32", "versions": "0+",
+          "about": "The partition index." },
+        { "name": "ErrorCode", "type": "int16", "versions": "0+",
+          "about": "The error code, or 0 if there was no fetch error." },
+        { "name": "HighWatermark", "type": "int64", "versions": "0+",
+          "about": "The current high water mark." },
+        { "name": "LastStableOffset", "type": "int64", "versions": "4+", "default": "-1", "ignorable": true,
+          "about": "The last stable offset (or LSO) of the partition. This is the last offset such that the state of all transactional records prior to this offset have been decided (ABORTED or COMMITTED)" },
+        { "name": "LogStartOffset", "type": "int64", "versions": "5+", "default": "-1", "ignorable": true,
+          "about": "The current log start offset." },
+        { "name": "DivergingEpoch", "type": "EpochEndOffset", "versions": "12+", "taggedVersions": "12+", "tag": 0,
+          "about": "In case divergence is detected based on the `LastFetchedEpoch` and `FetchOffset` in the request, this field indicates the largest epoch and its end offset such that subsequent records are known to diverge",
+          "fields": [
+            { "name": "Epoch", "type": "int32", "versions": "12+", "default": "-1" },
+            { "name": "EndOffset", "type": "int64", "versions": "12+", "default": "-1" }
+        ]},
+        { "name": "CurrentLeader", "type": "LeaderIdAndEpoch",
+          "versions": "12+", "taggedVersions": "12+", "tag": 1, "fields": [
+          { "name": "LeaderId", "type": "int32", "versions": "12+", "default": "-1",
+            "about": "The ID of the current leader or -1 if the leader is unknown."},
+          { "name": "LeaderEpoch", "type": "int32", "versions": "12+", "default": "-1",
+            "about": "The latest known leader epoch"}
+        ]},
+        { "name": "AbortedTransactions", "type": "[]AbortedTransaction", "versions": "4+", "nullableVersions": "4+", "ignorable": true,
+          "about": "The aborted transactions.",  "fields": [
+          { "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+", "default": "-1", "ignorable": false,
+          "about": "The preferred read replica for the consumer to use on its next fetch request"},
+        { "name": "RecordSet", "type": "records", "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
new file mode 100644
index 0000000000000..6a90887b1aa01
--- /dev/null
+++ b/clients/src/main/resources/common/message/FindCoordinatorRequest.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": 10,
+  "type": "request",
+  "name": "FindCoordinatorRequest",
+  // Version 1 adds KeyType.
+  //
+  // Version 2 is the same as version 1.
+  //
+  // Version 3 is the first flexible version.
+  "validVersions": "0-3",
+  "flexibleVersions": "3+",
+  "fields": [
+    { "name": "Key", "type": "string", "versions": "0+",
+      "about": "The coordinator key." },
+    { "name": "KeyType", "type": "int8", "versions": "1+", "default": "0", "ignorable": false,
+      "about": "The coordinator key type.  (Group, transaction, etc.)" }
+  ]
+}
diff --git a/clients/src/main/resources/common/message/FindCoordinatorResponse.json b/clients/src/main/resources/common/message/FindCoordinatorResponse.json
new file mode 100644
index 0000000000000..996e8fb5fba5c
--- /dev/null
+++ b/clients/src/main/resources/common/message/FindCoordinatorResponse.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": 10,
+  "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.
+  //
+  // 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." },
+    { "name": "ErrorCode", "type": "int16", "versions": "0+",
+      "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+", "entityType": "brokerId",
+      "about": "The node id." },
+    { "name": "Host", "type": "string", "versions": "0+",
+      "about": "The host name." },
+    { "name": "Port", "type": "int32", "versions": "0+",
+      "about": "The port." }
+  ]
+}
diff --git a/clients/src/main/resources/common/message/HeartbeatRequest.json b/clients/src/main/resources/common/message/HeartbeatRequest.json
new file mode 100644
index 0000000000000..4d799aa3c14c1
--- /dev/null
+++ b/clients/src/main/resources/common/message/HeartbeatRequest.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": 12,
+  "type": "request",
+  "name": "HeartbeatRequest",
+  // Version 1 and version 2 are the same as version 0.
+  //
+  // 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+", "entityType": "groupId",
+      "about": "The group id." },
+    { "name": "GenerationId", "type": "int32", "versions": "0+",
+      "about": "The generation of the group." },
+    { "name": "MemberId", "type": "string", "versions": "0+",
+      "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
new file mode 100644
index 0000000000000..280ba1103b437
--- /dev/null
+++ b/clients/src/main/resources/common/message/HeartbeatResponse.json
@@ -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.
+
+{
+  "apiKey": 12,
+  "type": "response",
+  "name": "HeartbeatResponse",
+  // Version 1 adds throttle time.
+  //
+  // Starting in version 2, on quota violation, brokers send out responses before throttling.
+  //
+  // 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." },
+    { "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/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
new file mode 100644
index 0000000000000..dc85063c29643
--- /dev/null
+++ b/clients/src/main/resources/common/message/InitProducerIdRequest.json
@@ -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.
+
+{
+  "apiKey": 22,
+  "type": "request",
+  "name": "InitProducerIdRequest",
+  // Version 1 is the same as version 0.
+  //
+  // Version 2 is the first flexible version.
+  //
+  // Version 3 adds ProducerId and ProducerEpoch, allowing producers to try to resume after an INVALID_PRODUCER_EPOCH error
+  //
+  // Version 4 adds the support for new error code PRODUCER_FENCED.
+  "validVersions": "0-4",
+  "flexibleVersions": "2+",
+  "fields": [
+    { "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 before aborting idle transactions sent by this producer. This is only relevant if a TransactionalId has been defined." },
+    { "name": "ProducerId", "type": "int64", "versions": "3+", "default": "-1",
+      "about": "The producer id. This is used to disambiguate requests if a transactional id is reused following its expiration." },
+    { "name": "ProducerEpoch", "type": "int16", "versions": "3+", "default": "-1",
+      "about": "The producer's current epoch. This will be checked against the producer epoch on the broker, and the request will return an error if they do not match." }
+  ]
+}
diff --git a/clients/src/main/resources/common/message/InitProducerIdResponse.json b/clients/src/main/resources/common/message/InitProducerIdResponse.json
new file mode 100644
index 0000000000000..f56c2fe67e3da
--- /dev/null
+++ b/clients/src/main/resources/common/message/InitProducerIdResponse.json
@@ -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.
+
+{
+  "apiKey": 22,
+  "type": "response",
+  "name": "InitProducerIdResponse",
+  // Starting in version 1, on quota violation, brokers send out responses before throttling.
+  //
+  // Version 2 is the first flexible version.
+  //
+  // Version 3 is the same as version 2.
+  //
+  // Version 4 adds the support for new error code PRODUCER_FENCED.
+  "validVersions": "0-4",
+  "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+", "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
new file mode 100644
index 0000000000000..2650d89e2b4a8
--- /dev/null
+++ b/clients/src/main/resources/common/message/JoinGroupRequest.json
@@ -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.
+
+{
+  "apiKey": 11,
+  "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.
+  //
+  // Version 6 is the first flexible version.
+  //
+  // Version 7 is the same as version 6.
+  "validVersions": "0-7",
+  "flexibleVersions": "6+",
+  "fields": [
+    { "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." },
+    // Note: if RebalanceTimeoutMs is not present, SessionTimeoutMs should be
+    // used instead.  The default of -1 here is just intended as a placeholder.
+    { "name": "RebalanceTimeoutMs", "type": "int32", "versions": "1+", "default": "-1", "ignorable": true,
+      "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+",
+      "about": "The list of protocols that the member supports.", "fields": [
+      { "name": "Name", "type": "string", "versions": "0+", "mapKey": true,
+        "about": "The protocol name." },
+      { "name": "Metadata", "type": "bytes", "versions": "0+",
+        "about": "The protocol metadata." }
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/JoinGroupResponse.json b/clients/src/main/resources/common/message/JoinGroupResponse.json
new file mode 100644
index 0000000000000..f95ec01bbb751
--- /dev/null
+++ b/clients/src/main/resources/common/message/JoinGroupResponse.json
@@ -0,0 +1,62 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements.  See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT 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": 11,
+  "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.
+  //
+  // Version 5 is bumped to apply group.instance.id to identify member across restarts.
+  //
+  // Version 6 is the first flexible version.
+  //
+  // Starting from version 7, the broker sends back the Protocol Type to the client (KIP-559).
+  "validVersions": "0-7",
+  "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+", "default": "-1",
+      "about": "The generation ID of the group." },
+    { "name": "ProtocolType", "type": "string", "versions": "7+",
+      "nullableVersions": "7+", "default": "null", "ignorable": true,
+      "about": "The group protocol name." },
+    { "name": "ProtocolName", "type": "string", "versions": "0+", "nullableVersions": "7+",
+      "about": "The group protocol selected by the coordinator." },
+    { "name": "Leader", "type": "string", "versions": "0+",
+      "about": "The leader of the group." },
+    { "name": "MemberId", "type": "string", "versions": "0+",
+      "about": "The member ID assigned by the group coordinator." },
+    { "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
new file mode 100644
index 0000000000000..852968801cee3
--- /dev/null
+++ b/clients/src/main/resources/common/message/LeaderAndIsrRequest.json
@@ -0,0 +1,82 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements.  See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+{
+  "apiKey": 4,
+  "type": "request",
+  "name": "LeaderAndIsrRequest",
+  // Version 1 adds IsNew.
+  //
+  // Version 2 adds broker epoch and reorganizes the partitions by topic.
+  //
+  // Version 3 adds AddingReplicas and RemovingReplicas
+  "validVersions": "0-4",
+  "flexibleVersions": "4+",
+  "fields": [
+    { "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",
+      "about": "The current broker epoch." },
+    { "name": "UngroupedPartitionStates", "type": "[]LeaderAndIsrPartitionState", "versions": "0-1",
+      "about": "The state of each partition, in a v0 or v1 message." },
+    // 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": "TopicName", "type": "string", "versions": "2+", "entityType": "topicName",
+        "about": "The topic name." },
+      { "name": "PartitionStates", "type": "[]LeaderAndIsrPartitionState", "versions": "2+",
+        "about": "The state of each partition" }
+    ]},
+    { "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+",
+        "about": "The controller epoch." },
+      { "name": "Leader", "type": "int32", "versions": "0+", "entityType": "brokerId",
+        "about": "The broker ID of the leader." },
+      { "name": "LeaderEpoch", "type": "int32", "versions": "0+",
+        "about": "The leader epoch." },
+      { "name": "Isr", "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": "AddingReplicas", "type": "[]int32", "versions": "3+", "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": "3+", "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." }
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/LeaderAndIsrResponse.json b/clients/src/main/resources/common/message/LeaderAndIsrResponse.json
new file mode 100644
index 0000000000000..10c3cd95e5fb2
--- /dev/null
+++ b/clients/src/main/resources/common/message/LeaderAndIsrResponse.json
@@ -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.
+
+{
+  "apiKey": 4,
+  "type": "response",
+  "name": "LeaderAndIsrResponse",
+  // Version 1 adds KAFKA_STORAGE_ERROR as a valid error code.
+  //
+  // Version 2 is the same as version 1.
+  //
+  // Version 3 is the same as version 2.
+  "validVersions": "0-4",
+  "flexibleVersions": "4+",
+  "fields": [
+    { "name": "ErrorCode", "type": "int16", "versions": "0+",
+      "about": "The error code, or 0 if there was no error." },
+    { "name": "PartitionErrors", "type": "[]LeaderAndIsrPartitionError", "versions": "0+",
+      "about": "Each partition.", "fields": [
+      { "name": "TopicName", "type": "string", "versions": "0+", "entityType": "topicName",
+        "about": "The topic name." },
+      { "name": "PartitionIndex", "type": "int32", "versions": "0+",
+        "about": "The partition index." },
+      { "name": "ErrorCode", "type": "int16", "versions": "0+",
+        "about": "The partition error code, or 0 if there was no error." }
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/LeaderChangeMessage.json b/clients/src/main/resources/common/message/LeaderChangeMessage.json
new file mode 100644
index 0000000000000..cb4d0fd36c907
--- /dev/null
+++ b/clients/src/main/resources/common/message/LeaderChangeMessage.json
@@ -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.
+
+{
+  "type": "data",
+  "name": "LeaderChangeMessage",
+  "validVersions": "0",
+  "flexibleVersions": "0+",
+  "fields": [
+    {"name": "LeaderId", "type": "int32", "versions": "0+",
+      "about": "The ID of the newly elected leader"},
+    {"name": "Voters", "type": "[]Voter", "versions": "0+",
+      "about": "The set of voters in the quorum for this epoch"},
+    {"name": "GrantingVoters", "type": "[]Voter", "versions": "0+",
+      "about": "The voters who voted for the leader at the time of election"}
+  ],
+  "commonStructs": [
+    { "name": "Voter", "versions": "0+", "fields": [
+      {"name": "VoterId", "type": "int32", "versions": "0+"}
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/LeaveGroupRequest.json b/clients/src/main/resources/common/message/LeaveGroupRequest.json
new file mode 100644
index 0000000000000..acc7938c387b1
--- /dev/null
+++ b/clients/src/main/resources/common/message/LeaveGroupRequest.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": 13,
+  "type": "request",
+  "name": "LeaveGroupRequest",
+  // Version 1 and 2 are the same as version 0.
+  //
+  // 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+", "entityType": "groupId",
+      "about": "The ID of the group to leave." },
+    { "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
new file mode 100644
index 0000000000000..0ddb4c6eb5323
--- /dev/null
+++ b/clients/src/main/resources/common/message/LeaveGroupResponse.json
@@ -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.
+
+{
+  "apiKey": 13,
+  "type": "response",
+  "name": "LeaveGroupResponse",
+  // Version 1 adds the throttle time.
+  //
+  // Starting in version 2, on quota violation, brokers send out responses before throttling.
+  //
+  // 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." },
+
+    { "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
new file mode 100644
index 0000000000000..dbe6d9b6f123a
--- /dev/null
+++ b/clients/src/main/resources/common/message/ListGroupsRequest.json
@@ -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.
+
+{
+  "apiKey": 16,
+  "type": "request",
+  "name": "ListGroupsRequest",
+  // Version 1 and 2 are the same as version 0.
+  //
+  // Version 3 is the first flexible version.
+  //
+  // Version 4 adds the StatesFilter field (KIP-518).
+  "validVersions": "0-4",
+  "flexibleVersions": "3+",
+  "fields": [
+    { "name": "StatesFilter", "type": "[]string", "versions": "4+",
+      "about": "The states of the groups we want to list. If empty all groups are returned with their state."
+    }
+  ]
+}
diff --git a/clients/src/main/resources/common/message/ListGroupsResponse.json b/clients/src/main/resources/common/message/ListGroupsResponse.json
new file mode 100644
index 0000000000000..87561c2ab964d
--- /dev/null
+++ b/clients/src/main/resources/common/message/ListGroupsResponse.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": 16,
+  "type": "response",
+  "name": "ListGroupsResponse",
+  // Version 1 adds the throttle time.
+  //
+  // Starting in version 2, on quota violation, brokers send out responses before throttling.
+  //
+  // Version 3 is the first flexible version.
+  //
+  // Version 4 adds the GroupState field (KIP-518).
+  "validVersions": "0-4",
+  "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." },
+    { "name": "ErrorCode", "type": "int16", "versions": "0+",
+      "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+", "entityType": "groupId",
+        "about": "The group ID." },
+      { "name": "ProtocolType", "type": "string", "versions": "0+",
+        "about": "The group protocol type." },
+      { "name": "GroupState", "type": "string", "versions": "4+", "ignorable": true,
+        "about": "The group state name." }
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/ListOffsetRequest.json b/clients/src/main/resources/common/message/ListOffsetRequest.json
new file mode 100644
index 0000000000000..8d1d6cadf8db2
--- /dev/null
+++ b/clients/src/main/resources/common/message/ListOffsetRequest.json
@@ -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.
+
+{
+  "apiKey": 2,
+  "type": "request",
+  "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 4.
+  //
+  // Version 6 enables flexible versions.
+  "validVersions": "0-6",
+  "flexibleVersions": "6+",
+  "fields": [
+    { "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+", "entityType": "topicName",
+        "about": "The topic name." },
+      { "name": "Partitions", "type": "[]ListOffsetPartition", "versions": "0+",
+        "about": "Each partition in the request.", "fields": [
+        { "name": "PartitionIndex", "type": "int32", "versions": "0+",
+          "about": "The partition index." },
+        { "name": "CurrentLeaderEpoch", "type": "int32", "versions": "4+", "default": "-1", "ignorable": true,
+          "about": "The current leader epoch." },
+        { "name": "Timestamp", "type": "int64", "versions": "0+",
+          "about": "The current timestamp." },
+        { "name": "MaxNumOffsets", "type": "int32", "versions": "0", "default": "1",
+          "about": "The maximum number of offsets to report." }
+      ]}
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/ListOffsetResponse.json b/clients/src/main/resources/common/message/ListOffsetResponse.json
new file mode 100644
index 0000000000000..ffa15621c7881
--- /dev/null
+++ b/clients/src/main/resources/common/message/ListOffsetResponse.json
@@ -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.
+
+{
+  "apiKey": 2,
+  "type": "response",
+  "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.
+  //
+  // Version 6 enables flexible versions.
+  "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": "Topics", "type": "[]ListOffsetTopicResponse", "versions": "0+",
+      "about": "Each topic in the response.", "fields": [
+      { "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": [
+        { "name": "PartitionIndex", "type": "int32", "versions": "0+",
+          "about": "The partition index." },
+        { "name": "ErrorCode", "type": "int16", "versions": "0+",
+          "about": "The partition error code, or 0 if there was no error." },
+        { "name": "OldStyleOffsets", "type": "[]int64", "versions": "0", "ignorable": false,
+          "about": "The result offsets." },
+        { "name": "Timestamp", "type": "int64", "versions": "1+", "default": "-1", "ignorable": false,
+          "about": "The timestamp associated with the returned offset." },
+        { "name": "Offset", "type": "int64", "versions": "1+", "default": "-1", "ignorable": false,
+          "about": "The returned offset." },
+        { "name": "LeaderEpoch", "type": "int32", "versions": "4+", "default": "-1" }
+      ]}
+    ]}
+  ]
+}
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
new file mode 100644
index 0000000000000..6129d427b7eef
--- /dev/null
+++ b/clients/src/main/resources/common/message/MetadataRequest.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": 3,
+  "type": "request",
+  "name": "MetadataRequest",
+  "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
+    // indiate "request metadata for all topics."
+    //
+    // 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+", "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." },
+    { "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
new file mode 100644
index 0000000000000..19faecff80d22
--- /dev/null
+++ b/clients/src/main/resources/common/message/MetadataResponse.json
@@ -0,0 +1,90 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements.  See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT 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": 3,
+  "type": "response",
+  "name": "MetadataResponse",
+  // Version 1 adds fields for the rack of each broker, the controller id, and
+  // whether or not the topic is internal.
+  //
+  // Version 2 adds the cluster ID field.
+  //
+  // Version 3 adds the throttle time.
+  //
+  // Version 4 is the same as version 3.
+  //
+  // Version 5 adds a per-partition offline_replicas field. This field specifies
+  // the list of replicas that are offline.
+  //
+  // Starting in version 6, on quota violation, brokers send out responses before throttling.
+  //
+  // Version 7 adds the leader epoch to the partition metadata.
+  //
+  // 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+", "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": "Brokers", "type": "[]MetadataResponseBroker", "versions": "0+",
+      "about": "Each broker in the response.", "fields": [
+      { "name": "NodeId", "type": "int32", "versions": "0+", "mapKey": true, "entityType": "brokerId",
+        "about": "The broker ID." },
+      { "name": "Host", "type": "string", "versions": "0+",
+        "about": "The broker hostname." },
+      { "name": "Port", "type": "int32", "versions": "0+",
+        "about": "The broker port." },
+      { "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, "default": "null",
+      "about": "The cluster ID that responding broker belongs to." },
+    { "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, "entityType": "topicName",
+        "about": "The topic name." },
+      { "name": "IsInternal", "type": "bool", "versions": "1+", "default": "false", "ignorable": true,
+        "about": "True if the topic is internal." },
+      { "name": "Partitions", "type": "[]MetadataResponsePartition", "versions": "0+",
+        "about": "Each partition in the topic.", "fields": [
+        { "name": "ErrorCode", "type": "int16", "versions": "0+",
+          "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+", "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+", "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
new file mode 100644
index 0000000000000..096b61917ae0f
--- /dev/null
+++ b/clients/src/main/resources/common/message/OffsetCommitRequest.json
@@ -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.
+
+{
+  "apiKey": 8,
+  "type": "request",
+  "name": "OffsetCommitRequest",
+  // Version 1 adds timestamp and group membership information, as well as the commit timestamp.
+  //
+  // Version 2 adds retention time.  It removes the commit timestamp added in version 1.
+  //
+  // Version 3 and 4 are the same as version 2. 
+  //
+  // Version 5 removes the retention time, which is now controlled only by a broker configuration.
+  //
+  // Version 6 adds the leader epoch for fencing.
+  //
+  // 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+", "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+", "entityType": "topicName",
+        "about": "The topic name." },
+      { "name": "Partitions", "type": "[]OffsetCommitRequestPartition", "versions": "0+",
+        "about": "Each partition to commit offsets for.", "fields": [
+        { "name": "PartitionIndex", "type": "int32", "versions": "0+",
+          "about": "The partition index." },
+        { "name": "CommittedOffset", "type": "int64", "versions": "0+",
+          "about": "The message offset to be committed." },
+        { "name": "CommittedLeaderEpoch", "type": "int32", "versions": "6+", "default": "-1", "ignorable": true,
+          "about": "The leader epoch of this partition." },
+        // CommitTimestamp has been removed from v2 and later.
+        { "name": "CommitTimestamp", "type": "int64", "versions": "1", "default": "-1",
+          "about": "The timestamp of the commit." },
+        { "name": "CommittedMetadata", "type": "string", "versions": "0+", "nullableVersions": "0+",
+          "about": "Any associated metadata the client wants to keep." }
+      ]}
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/OffsetCommitResponse.json b/clients/src/main/resources/common/message/OffsetCommitResponse.json
new file mode 100644
index 0000000000000..3d547794cfb42
--- /dev/null
+++ b/clients/src/main/resources/common/message/OffsetCommitResponse.json
@@ -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.
+
+{
+  "apiKey": 8,
+  "type": "response",
+  "name": "OffsetCommitResponse",
+  // Versions 1 and 2 are the same as version 0.
+  //
+  // Version 3 adds the throttle time to the response.
+  //
+  // Starting in version 4, on quota violation, brokers send out responses before throttling.
+  //
+  // Versions 5 and 6 are the same as version 4.
+  //
+  // 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+", "entityType": "topicName",
+        "about": "The topic name." },
+      { "name": "Partitions", "type": "[]OffsetCommitResponsePartition", "versions": "0+",
+        "about": "The responses for each partition in the topic.",  "fields": [
+        { "name": "PartitionIndex", "type": "int32", "versions": "0+",
+          "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/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
new file mode 100644
index 0000000000000..ddd53fc859600
--- /dev/null
+++ b/clients/src/main/resources/common/message/OffsetFetchRequest.json
@@ -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.
+
+{
+  "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. 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.
+  //
+  // Version 6 is the first flexible version.
+  //
+  // Version 7 is adding the require stable flag.
+  "validVersions": "0-7",
+  "flexibleVersions": "6+",
+  "fields": [
+    { "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+", "entityType": "topicName",
+        "about": "The topic name."},
+      { "name": "PartitionIndexes", "type": "[]int32", "versions": "0+",
+        "about": "The partition indexes we would like to fetch offsets for." }
+    ]},
+    {"name": "RequireStable", "type": "bool", "versions": "7+", "default": "false",
+     "about": "Whether broker should hold on returning unstable offsets but set a retriable error code for the partition."}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/OffsetFetchResponse.json b/clients/src/main/resources/common/message/OffsetFetchResponse.json
new file mode 100644
index 0000000000000..b9777017ed9ac
--- /dev/null
+++ b/clients/src/main/resources/common/message/OffsetFetchResponse.json
@@ -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.
+
+{
+  "apiKey": 9,
+  "type": "response",
+  "name": "OffsetFetchResponse",
+  // Version 1 is the same as version 0.
+  //
+  // Version 2 adds a top-level error code.
+  //
+  // Version 3 adds the throttle time.
+  //
+  // Starting in version 4, on quota violation, brokers send out responses before throttling.
+  //
+  // Version 5 adds the leader epoch to the committed offset.
+  //
+  // Version 6 is the first flexible version.
+  //
+  // Version 7 adds pending offset commit as new error response on partition level.
+  "validVersions": "0-7",
+  "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+", "entityType": "topicName",
+        "about": "The topic name." },
+      { "name": "Partitions", "type": "[]OffsetFetchResponsePartition", "versions": "0+",
+        "about": "The responses per partition", "fields": [
+        { "name": "PartitionIndex", "type": "int32", "versions": "0+",
+          "about": "The partition index." },
+        { "name": "CommittedOffset", "type": "int64", "versions": "0+",
+          "about": "The committed message offset." },
+        { "name": "CommittedLeaderEpoch", "type": "int32", "versions": "5+", "default": "-1",
+          "ignorable": true, "about": "The leader epoch." },
+        { "name": "Metadata", "type": "string", "versions": "0+", "nullableVersions": "0+",
+          "about": "The partition metadata." },
+        { "name": "ErrorCode", "type": "int16", "versions": "0+",
+          "about": "The error code, or 0 if there was no error." }
+      ]}
+    ]},
+    { "name": "ErrorCode", "type": "int16", "versions": "2+", "default": "0", "ignorable": true,
+      "about": "The top-level error code, or 0 if there was no error." }
+  ]
+}
diff --git a/clients/src/main/resources/common/message/OffsetForLeaderEpochRequest.json b/clients/src/main/resources/common/message/OffsetForLeaderEpochRequest.json
new file mode 100644
index 0000000000000..75b6c8dc3227c
--- /dev/null
+++ b/clients/src/main/resources/common/message/OffsetForLeaderEpochRequest.json
@@ -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.
+
+{
+  "apiKey": 23,
+  "type": "request",
+  "name": "OffsetForLeaderEpochRequest",
+  // Version 1 is the same as version 0.
+  //
+  // Version 2 adds the current leader epoch to support fencing.
+  //
+  // 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.
+  //
+  // Version 4 enables flexible versions.
+  "validVersions": "0-4",
+  "flexibleVersions": "4+",
+  "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": "Topic", "type": "string", "versions": "0+", "entityType": "topicName",
+        "mapKey": true, "about": "The topic name." },
+      { "name": "Partitions", "type": "[]OffsetForLeaderPartition", "versions": "0+",
+        "about": "Each partition to get offsets for.", "fields": [
+        { "name": "Partition", "type": "int32", "versions": "0+",
+          "about": "The partition index." },
+        { "name": "CurrentLeaderEpoch", "type": "int32", "versions": "2+", "default": "-1", "ignorable": true,
+          "about": "An epoch used to fence consumers/replicas with old metadata. If the epoch provided by the client is larger than the current epoch known to the broker, then the UNKNOWN_LEADER_EPOCH error code will be returned. If the provided epoch is smaller, then the FENCED_LEADER_EPOCH error code will be returned." },
+        { "name": "LeaderEpoch", "type": "int32", "versions": "0+",
+          "about": "The epoch to look up an offset for." }
+      ]}
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/OffsetForLeaderEpochResponse.json b/clients/src/main/resources/common/message/OffsetForLeaderEpochResponse.json
new file mode 100644
index 0000000000000..2b0810e1ec3ce
--- /dev/null
+++ b/clients/src/main/resources/common/message/OffsetForLeaderEpochResponse.json
@@ -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.
+
+{
+  "apiKey": 23,
+  "type": "response",
+  "name": "OffsetForLeaderEpochResponse",
+  // Version 1 added the leader epoch to the response.
+  //
+  // Version 2 added the throttle time.
+  //
+  // Version 3 is the same as version 2.
+  //
+  // Version 4 enables flexible versions.
+  "validVersions": "0-4",
+  "flexibleVersions": "4+",
+  "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": "Topic", "type": "string", "versions": "0+", "entityType": "topicName",
+        "mapKey": true, "about": "The topic name." },
+      { "name": "Partitions", "type": "[]EpochEndOffset", "versions": "0+",
+        "about": "Each partition in the topic we fetched offsets for.", "fields": [
+        { "name": "ErrorCode", "type": "int16", "versions": "0+",
+          "about": "The error code 0, or if there was no error." },
+        { "name": "Partition", "type": "int32", "versions": "0+",
+          "about": "The partition index." },
+        { "name": "LeaderEpoch", "type": "int32", "versions": "1+", "default": "-1", "ignorable": true,
+          "about": "The leader epoch of the partition." },
+        { "name": "EndOffset", "type": "int64", "versions": "0+", "default": "-1",
+          "about": "The end offset of the epoch." }
+      ]}
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/ProduceRequest.json b/clients/src/main/resources/common/message/ProduceRequest.json
new file mode 100644
index 0000000000000..73e72c352f4d4
--- /dev/null
+++ b/clients/src/main/resources/common/message/ProduceRequest.json
@@ -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.
+
+{
+  "apiKey": 0,
+  "type": "request",
+  "name": "ProduceRequest",
+  // Version 1 and 2 are the same as version 0.
+  //
+  // Version 3 adds the transactional ID, which is used for authorization when attempting to write
+  // transactional data.  Version 3 also adds support for Kafka Message Format v2.
+  //
+  // Version 4 is the same as version 3, but the requestor must be prepared to handle a
+  // KAFKA_STORAGE_ERROR. 
+  //
+  // Version 5 and 6 are the same as version 3.
+  //
+  // Starting in version 7, records can be produced using ZStandard compression.  See KIP-110.
+  //
+  // Starting in Version 8, response has RecordErrors and ErrorMEssage. See KIP-467.
+  //
+  // Version 9 enables flexible versions.
+  "validVersions": "0-9",
+  "flexibleVersions": "9+",
+  "fields": [
+    { "name": "TransactionalId", "type": "string", "versions": "3+", "nullableVersions": "3+", "default": "null", "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." },
+    { "name": "TimeoutMs", "type": "int32", "versions": "0+",
+      "about": "The timeout to await a response in miliseconds." },
+    { "name": "TopicData", "type": "[]TopicProduceData", "versions": "0+",
+      "about": "Each topic to produce to.", "fields": [
+      { "name": "Name", "type": "string", "versions": "0+", "entityType": "topicName", "mapKey": true,
+        "about": "The topic name." },
+      { "name": "PartitionData", "type": "[]PartitionProduceData", "versions": "0+",
+        "about": "Each partition to produce to.", "fields": [
+        { "name": "Index", "type": "int32", "versions": "0+",
+          "about": "The partition index." },
+        { "name": "Records", "type": "records", "versions": "0+", "nullableVersions": "0+",
+          "about": "The record data to be produced." }
+      ]}
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/ProduceResponse.json b/clients/src/main/resources/common/message/ProduceResponse.json
new file mode 100644
index 0000000000000..d6e5566eab3c2
--- /dev/null
+++ b/clients/src/main/resources/common/message/ProduceResponse.json
@@ -0,0 +1,68 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements.  See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT 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": 0,
+  "type": "response",
+  "name": "ProduceResponse",
+  // Version 1 added the throttle time.
+  //
+  // Version 2 added the log append time.
+  //
+  // Version 3 is the same as version 2.
+  //
+  // Version 4 added KAFKA_STORAGE_ERROR as a possible error code.
+  //
+  // Version 5 added LogStartOffset to filter out spurious
+  // OutOfOrderSequenceExceptions on the client.
+  //
+  // Version 8 added RecordErrors and ErrorMessage to include information about
+  // records that cause the whole batch to be dropped.  See KIP-467 for details.
+  //
+  // Version 9 enables flexible versions.
+  "validVersions": "0-9",
+  "flexibleVersions": "9+",
+  "fields": [
+    { "name": "Responses", "type": "[]TopicProduceResponse", "versions": "0+",
+      "about": "Each produce response", "fields": [
+      { "name": "Name", "type": "string", "versions": "0+", "entityType": "topicName", "mapKey": true,
+        "about": "The topic name" },
+      { "name": "PartitionResponses", "type": "[]PartitionProduceResponse", "versions": "0+",
+        "about": "Each partition that we produced to within the topic.", "fields": [
+        { "name": "Index", "type": "int32", "versions": "0+",
+          "about": "The partition index." },
+        { "name": "ErrorCode", "type": "int16", "versions": "0+",
+          "about": "The error code, or 0 if there was no error." },
+        { "name": "BaseOffset", "type": "int64", "versions": "0+",
+          "about": "The base offset." },
+        { "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." },
+        { "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, "default": "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/README.md b/clients/src/main/resources/common/message/README.md
new file mode 100644
index 0000000000000..774059399446a
--- /dev/null
+++ b/clients/src/main/resources/common/message/README.md
@@ -0,0 +1,261 @@
+Apache Kafka Message Definitions
+================================
+
+Introduction
+------------
+The JSON files in this directory define the Apache Kafka message protocol.
+This protocol describes what information clients and servers send to each
+other, and how it is serialized.  Note that this version of JSON supports
+comments.  Comments begin with a double forward slash.
+
+When Kafka is compiled, these specification files are translated into Java code
+to read and write messages.  Any change to these JSON files will trigger a
+recompilation of this generated code.
+
+These specification files replace an older system where hand-written
+serialization code was used.  Over time, we will migrate all messages to using
+automatically generated serialization and deserialization code.
+
+Requests and Responses
+----------------------
+The Kafka protocol features requests and responses.  Requests are sent to a
+server in order to get a response.  Each request is uniquely identified by a
+16-bit integer called the "api key".  The API key of the response will always
+match that of the request.
+
+Each message has a unique 16-bit version number.  The schema might be different
+for each version of the message.  Sometimes, the version is incremented even
+though the schema has not changed.  This may indicate that the server should
+behave differently in some way.  The version of a response must always match
+the version of the corresponding request.
+
+Each request or response has a top-level field named "validVersions."  This
+specifies the versions of the protocol that our code understands.  For example,
+specifying "0-2" indicates that we understand versions 0, 1, and 2.  You must
+always specify the highest message version which is supported.
+
+The only old message versions that are no longer supported are version 0 of
+MetadataRequest and MetadataResponse.  In general, since we adopted KIP-97,
+dropping support for old message versions is no longer allowed without a KIP.
+Therefore, please be careful not to increase the lower end of the version
+support interval for any message.
+
+MessageData Objects
+-------------------
+Using the JSON files in this directory, we generate Java code for MessageData
+objects.  These objects store request and response data for kafka.  MessageData
+objects do not contain a version number.  Instead, a single MessageData object
+represents every possible version of a Message.  This makes working with
+messages more convenient, because the same code path can be used for every
+version of a message.
+
+Fields
+------
+Each message contains an array of fields.  Fields specify the data that should
+be sent with the message.  In general, fields have a name, a type, and version
+information associated with them.
+
+The order that fields appear in a message is important.  Fields which come
+first in the message definition will be sent first over the network.  Changing
+the order of the fields in a message is an incompatible change.
+
+In each new message version, we may add or subtract fields.  For example, if we
+are creating a new version 3 of a message, we can add a new field with the
+version spec "3+".  This specifies that the field only appears in version 3 and
+later.  If a field is being removed, we should change its version from "0+" to
+"0-2" to indicate that it will not appear in version 3 and later.
+
+Field Types
+-----------
+There are several primitive field types available.
+
+* "boolean": either true or false.
+
+* "int8": an 8-bit integer.
+
+* "int16": a 16-bit integer.
+
+* "int32": a 32-bit integer.
+
+* "int64": a 64-bit integer.
+
+* "float64": is a double-precision floating point number (IEEE 754).
+
+* "string": a UTF-8 string.
+
+* "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, ints, and floats can never be null.  However, fields that are strings,
+bytes, or arrays may optionally be "nullable."  When a field is "nullable," that
+simply means that we are prepared to serialize and deserialize null entries for
+that field.
+
+If you want to declare a field as nullable, you set "nullableVersions" for that
+field.  Nullability is implemented as a version range in order to accomodate a
+very common pattern in Kafka where a field that was originally not nullable
+becomes nullable in a later version.
+
+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
+written out will depend on the version number that you supply to write().  When
+you write out a message using an older version, fields that are too old to be
+present in the schema will be omitted.
+
+When working with older message versions, please verify that the older message
+schema includes all the data that needs to be sent.  For example, it is probably
+OK to skip sending a timeout field.  However, a field which radically alters the
+meaning of the request, such as a "validateOnly" boolean, should not be ignored.
+
+It's often useful to know how much space a message will take up before writing
+it out to a buffer.  You can find this out by calling the Message#size method.
+
+You can also convert a message to a Struct by calling Message#toStruct.  This
+allows you to use the functions that serialize Structs to buffers.
+
+Deserializing Messages
+----------------------
+Message objects may be deserialized using the Message#read method.  This method
+overwrites all the data currently in the message object with new data.
+
+You can also deserialize a message from a Struct by calling Message#fromStruct.
+The Struct will not be modified.
+
+Any fields in the message object that are not present in the version that you
+are deserializing will be reset to default values.  Unless a custom default has
+been set:
+
+* Integer fields default to 0.
+
+* Floats default to 0.
+
+* Booleans default to false.
+
+* Strings default to the empty string.
+
+* Bytes fields default to the empty byte array.
+
+* Array fields default to empty.
+
+You can specify "null" as a default value for a string field by specifing the
+literal string "null".  Note that you can only specify null as a default if all
+versions of the field are nullable.
+
+Custom Default Values
+---------------------
+You may set a custom default for fields that are integers, booleans, floats, or
+strings.  Just add a "default" entry in the JSON object.  The custom default
+overrides the normal default for the type.  So for example, you could make a
+boolean field default to true rather than false, and so forth.
+
+Note that the default must be valid for the field type.  So the default for an
+int16 field must be an integer that fits in 16 bits, and so forth.  You may
+specify hex or octal values, as long as they are prefixed with 0x or 0.  It is
+currently not possible to specify a custom default for bytes or array fields.
+
+Custom defaults are useful when an older message version lacked some
+information.  For example, if an older request lacked a timeout field, you may
+want to specify that the server should assume that the timeout for such a
+request is 5000 ms (or some other arbitrary value).
+
+Ignorable Fields
+----------------
+When we write messages using an older or newer format, not all fields may be
+present.  The message receiver will fill in the default value for the field
+during deserialization.  Therefore, if the source field was set to a non-default
+value, that information will be lost.
+
+In some cases, this information loss is acceptable.  For example, if a timeout
+field does not get preserved, this is not a problem.  However, in other cases,
+the field is really quite important and should not be discarded.  One example is
+a "verify only" boolean which changes the whole meaning of the request.
+
+By default, we assume that information loss is not acceptable.  The message
+serialization code will throw an exception if the ignored field is not set to
+the default value.  If information loss for a field is OK, please set
+"ignorable" to true for the field to disable this behavior.  When ignorable is
+set to true, the field may be silently omitted during serialization.
+
+Hash Sets
+---------
+One very common pattern in Kafka is to load array elements from a message into
+a Map or Set for easier access.  The message protocol makes this easier with
+the "mapKey" concept.
+
+If some of the elements of an array are annotated with "mapKey": true, the
+entire array will be treated as a linked hash set rather than a list.  Elements
+in this set will be accessible in O(1) time with an automatically generated
+"find" function.  The order of elements in the set will still be preserved,
+however.  New entries that are added to the set always show up as last in the
+ordering.
+
+Incompatible Changes
+--------------------
+It's very important to avoid making incompatible changes to the message
+protocol.  Here are some examples of incompatible changes:
+
+#### Making changes to a protocol version which has already been released.
+Protocol versions that have been released must be regarded as done.  If there
+were mistakes, they should be corrected in a new version rather than changing
+the existing version.
+
+#### Re-ordering existing fields.
+It is OK to add new fields before or after existing fields.  However, existing
+fields should not be re-ordered with respect to each other.
+
+#### Changing the default of an existing field.
+You must never change the default of a field which already exists.  Otherwise,
+new clients and old servers will not agree on the default, and so forth.
+
+#### Changing the type of an existing field.
+One exception is that an array of primitives may be changed to an array of
+structures containing the same data, as long as the conversion is done
+correctly.  The Kafka protocol does not do any "boxing" of structures, so an
+array of structs that contain a single int32 is the same as an array of int32s.
diff --git a/clients/src/main/resources/common/message/RenewDelegationTokenRequest.json b/clients/src/main/resources/common/message/RenewDelegationTokenRequest.json
new file mode 100644
index 0000000000000..9fbb0fa07a5e8
--- /dev/null
+++ b/clients/src/main/resources/common/message/RenewDelegationTokenRequest.json
@@ -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.
+
+{
+  "apiKey": 39,
+  "type": "request",
+  "name": "RenewDelegationTokenRequest",
+  // Version 1 is the same as version 0.
+  // Version 2 adds flexible version support
+  "validVersions": "0-2",
+  "flexibleVersions": "2+",
+  "fields": [
+    { "name": "Hmac", "type": "bytes", "versions": "0+",
+      "about": "The HMAC of the delegation token to be renewed." },
+    { "name": "RenewPeriodMs", "type": "int64", "versions": "0+",
+      "about": "The renewal time period in milliseconds." }
+  ]
+}
diff --git a/clients/src/main/resources/common/message/RenewDelegationTokenResponse.json b/clients/src/main/resources/common/message/RenewDelegationTokenResponse.json
new file mode 100644
index 0000000000000..c429dadd0cf4f
--- /dev/null
+++ b/clients/src/main/resources/common/message/RenewDelegationTokenResponse.json
@@ -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.
+
+{
+  "apiKey": 39,
+  "type": "response",
+  "name": "RenewDelegationTokenResponse",
+  // Starting in version 1, on quota violation, brokers send out responses before throttling.
+  // Version 2 adds flexible version support
+  "validVersions": "0-2",
+  "flexibleVersions": "2+",
+  "fields": [
+    { "name": "ErrorCode", "type": "int16", "versions": "0+",
+      "about": "The error code, or 0 if there was no error." },
+    { "name": "ExpiryTimestampMs", "type": "int64", "versions": "0+",
+      "about": "The timestamp in milliseconds at which this token expires." },
+    { "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/RequestHeader.json b/clients/src/main/resources/common/message/RequestHeader.json
new file mode 100644
index 0000000000000..fbf4e2ca9f454
--- /dev/null
+++ b/clients/src/main/resources/common/message/RequestHeader.json
@@ -0,0 +1,42 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements.  See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+{
+  "type": "header",
+  "name": "RequestHeader",
+  // 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." },
+    { "name": "RequestApiVersion", "type": "int16", "versions": "0+",
+      "about": "The API version of this request." },
+    { "name": "CorrelationId", "type": "int32", "versions": "0+",
+      "about": "The correlation ID of this request." },
+
+    // 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
new file mode 100644
index 0000000000000..773673601a886
--- /dev/null
+++ b/clients/src/main/resources/common/message/ResponseHeader.json
@@ -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.
+
+{
+  "type": "header",
+  "name": "ResponseHeader",
+  // 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
new file mode 100644
index 0000000000000..122cef577076e
--- /dev/null
+++ b/clients/src/main/resources/common/message/SaslAuthenticateRequest.json
@@ -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.
+
+{
+  "apiKey": 36,
+  "type": "request",
+  "name": "SaslAuthenticateRequest",
+  // Version 1 is the same as version 0.
+  // Version 2 adds flexible version support
+  "validVersions": "0-2",
+  "flexibleVersions": "2+",
+  "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
new file mode 100644
index 0000000000000..0e26a51e8b2c9
--- /dev/null
+++ b/clients/src/main/resources/common/message/SaslAuthenticateResponse.json
@@ -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.
+
+{
+  "apiKey": 36,
+  "type": "response",
+  "name": "SaslAuthenticateResponse",
+  // Version 1 adds the session lifetime.
+  // Version 2 adds flexible version support
+  "validVersions": "0-2",
+  "flexibleVersions": "2+",
+  "fields": [
+    { "name": "ErrorCode", "type": "int16", "versions": "0+",
+      "about": "The error code, or 0 if there was no error." },
+    { "name": "ErrorMessage", "type": "string", "versions": "0+", "nullableVersions": "0+",
+      "about": "The error message, or null if there was no error." },
+    { "name": "AuthBytes", "type": "bytes", "versions": "0+",
+      "about": "The SASL authentication bytes from the server, as defined by the SASL mechanism." },
+    { "name": "SessionLifetimeMs", "type": "int64", "versions": "1+", "default": "0", "ignorable": true,
+      "about": "The SASL authentication bytes from the server, as defined by the SASL mechanism." }
+  ]
+}
diff --git a/clients/src/main/resources/common/message/SaslHandshakeRequest.json b/clients/src/main/resources/common/message/SaslHandshakeRequest.json
new file mode 100644
index 0000000000000..f384f414c5c04
--- /dev/null
+++ b/clients/src/main/resources/common/message/SaslHandshakeRequest.json
@@ -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.
+
+{
+  "apiKey": 17,
+  "type": "request",
+  "name": "SaslHandshakeRequest",
+  // Version 1 supports SASL_AUTHENTICATE.
+  // NOTE: Version cannot be easily bumped due to incorrect
+  // client negotiation for clients <= 2.4.
+  // See https://issues.apache.org/jira/browse/KAFKA-9577
+  "validVersions": "0-1",
+  "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
new file mode 100644
index 0000000000000..039841f4aff1b
--- /dev/null
+++ b/clients/src/main/resources/common/message/SaslHandshakeResponse.json
@@ -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.
+
+{
+  "apiKey": 17,
+  "type": "response",
+  "name": "SaslHandshakeResponse",
+  // Version 1 is the same as version 0.
+  // NOTE: Version cannot be easily bumped due to incorrect
+  // client negotiation for clients <= 2.4.
+  // See https://issues.apache.org/jira/browse/KAFKA-9577
+  "validVersions": "0-1",
+  "fields": [
+    { "name": "ErrorCode", "type": "int16", "versions": "0+",
+      "about": "The error code, or 0 if there was no error." },
+    { "name": "Mechanisms", "type": "[]string", "versions": "0+",
+      "about": "The mechanisms enabled in the server." }
+  ]
+}
diff --git a/clients/src/main/resources/common/message/StopReplicaRequest.json b/clients/src/main/resources/common/message/StopReplicaRequest.json
new file mode 100644
index 0000000000000..d43356cab89fe
--- /dev/null
+++ b/clients/src/main/resources/common/message/StopReplicaRequest.json
@@ -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.
+
+{
+  "apiKey": 5,
+  "type": "request",
+  "name": "StopReplicaRequest",
+  // Version 1 adds the broker epoch and reorganizes the partitions to be stored
+  // per topic.
+  //
+  // Version 2 is the first flexible version.
+  //
+  // Version 3 adds the leader epoch per partition (KIP-570).
+  "validVersions": "0-3",
+  "flexibleVersions": "2+",
+  "fields": [
+    { "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,
+      "about": "The broker epoch." },
+    { "name": "DeletePartitions", "type": "bool", "versions": "0-2",
+      "about": "Whether these partitions should be deleted." },
+    { "name": "UngroupedPartitions", "type": "[]StopReplicaPartitionV0", "versions": "0",
+      "about": "The partitions to stop.", "fields": [
+      { "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": "[]StopReplicaTopicV1", "versions": "1-2",
+      "about": "The topics to stop.", "fields": [
+      { "name": "Name", "type": "string", "versions": "1-2", "entityType": "topicName",
+        "about": "The topic name." },
+      { "name": "PartitionIndexes", "type": "[]int32", "versions": "1-2",
+        "about": "The partition indexes." }
+    ]},
+    { "name": "TopicStates", "type": "[]StopReplicaTopicState", "versions": "3+",
+      "about": "Each topic.", "fields": [
+      { "name": "TopicName", "type": "string", "versions": "3+", "entityType": "topicName",
+        "about": "The topic name." },
+      { "name": "PartitionStates", "type": "[]StopReplicaPartitionState", "versions": "3+",
+        "about": "The state of each partition", "fields": [
+        { "name": "PartitionIndex", "type": "int32", "versions": "3+",
+          "about": "The partition index." },
+        { "name": "LeaderEpoch", "type": "int32", "versions": "3+", "default": "-1",
+          "about": "The leader epoch." },
+        { "name": "DeletePartition", "type": "bool", "versions": "3+",
+          "about": "Whether this partition should be deleted." }
+      ]}
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/StopReplicaResponse.json b/clients/src/main/resources/common/message/StopReplicaResponse.json
new file mode 100644
index 0000000000000..64b355e040687
--- /dev/null
+++ b/clients/src/main/resources/common/message/StopReplicaResponse.json
@@ -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.
+
+{
+  "apiKey": 5,
+  "type": "response",
+  "name": "StopReplicaResponse",
+  // Version 1 is the same as version 0.
+  //
+  // Version 2 is the first flexible version.
+  //
+  // Version 3 returns FENCED_LEADER_EPOCH if the epoch of the leader is stale (KIP-570).
+  "validVersions": "0-3",
+  "flexibleVersions": "2+",
+  "fields": [
+    { "name": "ErrorCode", "type": "int16", "versions": "0+",
+      "about": "The top-level error code, or 0 if there was no top-level error." },
+    { "name": "PartitionErrors", "type": "[]StopReplicaPartitionError", "versions": "0+",
+      "about": "The responses for each partition.", "fields": [
+      { "name": "TopicName", "type": "string", "versions": "0+", "entityType": "topicName",
+        "about": "The topic name." },
+      { "name": "PartitionIndex", "type": "int32", "versions": "0+",
+        "about": "The partition index." },
+      { "name": "ErrorCode", "type": "int16", "versions": "0+",
+        "about": "The partition error code, or 0 if there was no partition error." }
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/SyncGroupRequest.json b/clients/src/main/resources/common/message/SyncGroupRequest.json
new file mode 100644
index 0000000000000..a0a599150ef53
--- /dev/null
+++ b/clients/src/main/resources/common/message/SyncGroupRequest.json
@@ -0,0 +1,55 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements.  See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT 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": 14,
+  "type": "request",
+  "name": "SyncGroupRequest",
+  // Versions 1 and 2 are the same as version 0.
+  //
+  // Starting from version 3, we add a new field called groupInstanceId to indicate member identity across restarts.
+  //
+  // Version 4 is the first flexible version.
+  //
+  // Starting from version 5, the client sends the Protocol Type and the Protocol Name
+  // to the broker (KIP-559). The broker will reject the request if they are inconsistent
+  // with the Type and Name known by the broker.
+  "validVersions": "0-5",
+  "flexibleVersions": "4+",
+  "fields": [
+    { "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": "ProtocolType", "type": "string", "versions": "5+",
+      "nullableVersions": "5+", "default": "null", "ignorable": true,
+      "about": "The group protocol type." },
+    { "name": "ProtocolName", "type": "string", "versions": "5+",
+      "nullableVersions": "5+", "default": "null", "ignorable": true,
+      "about": "The group protocol name." },
+    { "name": "Assignments", "type": "[]SyncGroupRequestAssignment", "versions": "0+",
+      "about": "Each assignment.", "fields": [
+      { "name": "MemberId", "type": "string", "versions": "0+",
+        "about": "The ID of the member to assign." },
+      { "name": "Assignment", "type": "bytes", "versions": "0+",
+        "about": "The member assignment." }
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/SyncGroupResponse.json b/clients/src/main/resources/common/message/SyncGroupResponse.json
new file mode 100644
index 0000000000000..4aa17e0d79ad9
--- /dev/null
+++ b/clients/src/main/resources/common/message/SyncGroupResponse.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": 14,
+  "type": "response",
+  "name": "SyncGroupResponse",
+  // Version 1 adds throttle time.
+  //
+  // Starting in version 2, on quota violation, brokers send out responses before throttling.
+  //
+  // Starting from version 3, syncGroupRequest supports a new field called groupInstanceId to indicate member identity across restarts.
+  //
+  // Version 4 is the first flexible version.
+  //
+  // Starting from version 5, the broker sends back the Protocol Type and the Protocol Name
+  // to the client (KIP-559).
+  "validVersions": "0-5",
+  "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." },
+    { "name": "ProtocolType", "type": "string", "versions": "5+",
+      "nullableVersions": "5+", "default": "null", "ignorable": true,
+      "about": "The group protocol type." },
+    { "name": "ProtocolName", "type": "string", "versions": "5+",
+      "nullableVersions": "5+", "default": "null", "ignorable": true,
+      "about": "The group protocol name." },
+    { "name": "Assignment", "type": "bytes", "versions": "0+",
+      "about": "The member assignment." }
+  ]
+}
diff --git a/clients/src/main/resources/common/message/TxnOffsetCommitRequest.json b/clients/src/main/resources/common/message/TxnOffsetCommitRequest.json
new file mode 100644
index 0000000000000..d4c4efba6702f
--- /dev/null
+++ b/clients/src/main/resources/common/message/TxnOffsetCommitRequest.json
@@ -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.
+
+{
+  "apiKey": 28,
+  "type": "request",
+  "name": "TxnOffsetCommitRequest",
+  // Version 1 is the same as version 0.
+  //
+  // Version 2 adds the committed leader epoch.
+  //
+  // Version 3 adds the member.id, group.instance.id and generation.id.
+  "validVersions": "0-3",
+  "flexibleVersions": "3+",
+  "fields": [
+    { "name": "TransactionalId", "type": "string", "versions": "0+",
+      "about": "The ID of the transaction." },
+    { "name": "GroupId", "type": "string", "versions": "0+", "entityType": "groupId",
+      "about": "The ID of the group." },
+    { "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": "GenerationId", "type": "int32", "versions": "3+", "default": "-1",
+      "about": "The generation of the consumer." },
+    { "name": "MemberId", "type": "string", "versions": "3+", "default": "",
+      "about": "The member ID assigned by the group coordinator." },
+    { "name": "GroupInstanceId", "type": "string", "versions": "3+",
+      "nullableVersions": "3+", "default": "null",
+      "about": "The unique identifier of the consumer instance provided by end user." },
+    { "name": "Topics", "type" : "[]TxnOffsetCommitRequestTopic", "versions": "0+",
+      "about": "Each topic that we want to commit offsets for.", "fields": [
+      { "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": [
+        { "name": "PartitionIndex", "type": "int32", "versions": "0+",
+          "about": "The index of the partition within the topic." },
+        { "name": "CommittedOffset", "type": "int64", "versions": "0+",
+          "about": "The message offset to be committed." },
+        { "name": "CommittedLeaderEpoch", "type": "int32", "versions": "2+", "default": "-1", "ignorable": true,
+          "about": "The leader epoch of the last consumed record." },
+        { "name": "CommittedMetadata", "type": "string", "versions": "0+", "nullableVersions": "0+",
+          "about": "Any associated metadata the client wants to keep." }
+      ]}
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/TxnOffsetCommitResponse.json b/clients/src/main/resources/common/message/TxnOffsetCommitResponse.json
new file mode 100644
index 0000000000000..96b03a078d877
--- /dev/null
+++ b/clients/src/main/resources/common/message/TxnOffsetCommitResponse.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": 28,
+  "type": "response",
+  "name": "TxnOffsetCommitResponse",
+  // Starting in version 1, on quota violation, brokers send out responses before throttling.
+  //
+  // Version 2 is the same as version 1.
+  //
+  // Version 3 adds illegal generation, fenced instance id, and unknown member id errors.
+  "validVersions": "0-3",
+  "flexibleVersions": "3+",
+  "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+", "entityType": "topicName",
+        "about": "The topic name." },
+      { "name": "Partitions", "type": "[]TxnOffsetCommitResponsePartition", "versions": "0+",
+        "about": "The responses for each partition in the topic.", "fields": [
+        { "name": "PartitionIndex", "type": "int32", "versions": "0+",
+          "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/UpdateFeaturesRequest.json b/clients/src/main/resources/common/message/UpdateFeaturesRequest.json
new file mode 100644
index 0000000000000..ab882dff1c754
--- /dev/null
+++ b/clients/src/main/resources/common/message/UpdateFeaturesRequest.json
@@ -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.
+
+{
+  "apiKey": 57,
+  "type": "request",
+  "name": "UpdateFeaturesRequest",
+  "validVersions": "0",
+  "flexibleVersions": "0+",
+  "fields": [
+    { "name": "timeoutMs", "type": "int32", "versions": "0+", "default": "60000",
+      "about": "How long to wait in milliseconds before timing out the request." },
+    { "name": "FeatureUpdates", "type": "[]FeatureUpdateKey", "versions": "0+",
+      "about": "The list of updates to finalized features.", "fields": [
+      {"name": "Feature", "type": "string", "versions": "0+", "mapKey": true,
+        "about": "The name of the finalized feature to be updated."},
+      {"name": "MaxVersionLevel", "type": "int16", "versions": "0+",
+        "about": "The new maximum version level for the finalized feature. A value >= 1 is valid. A value < 1, is special, and can be used to request the deletion of the finalized feature."},
+      {"name": "AllowDowngrade", "type": "bool", "versions": "0+",
+        "about": "When set to true, the finalized feature version level is allowed to be downgraded/deleted. The downgrade request will fail if the new maximum version level is a value that's not lower than the existing maximum finalized version level."}
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/UpdateFeaturesResponse.json b/clients/src/main/resources/common/message/UpdateFeaturesResponse.json
new file mode 100644
index 0000000000000..615f6177cfbee
--- /dev/null
+++ b/clients/src/main/resources/common/message/UpdateFeaturesResponse.json
@@ -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.
+
+{
+  "apiKey": 57,
+  "type": "response",
+  "name": "UpdateFeaturesResponse",
+  "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 top-level error." },
+    { "name": "ErrorMessage", "type": "string", "versions": "0+", "nullableVersions": "0+",
+      "about": "The top-level error message, or `null` if there was no top-level error." },
+    { "name": "Results", "type": "[]UpdatableFeatureResult", "versions": "0+",
+      "about": "Results for each feature update.", "fields": [
+      {"name": "Feature", "type": "string", "versions": "0+", "mapKey": true,
+        "about": "The name of the finalized feature."},
+      { "name": "ErrorCode", "type": "int16", "versions": "0+",
+        "about": "The feature update error code or `0` if the feature update succeeded." },
+      { "name": "ErrorMessage", "type": "string", "versions": "0+", "nullableVersions": "0+",
+        "about": "The feature update error, or `null` if the feature update succeeded." }
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/UpdateMetadataRequest.json b/clients/src/main/resources/common/message/UpdateMetadataRequest.json
new file mode 100644
index 0000000000000..3c45f833a7cab
--- /dev/null
+++ b/clients/src/main/resources/common/message/UpdateMetadataRequest.json
@@ -0,0 +1,93 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements.  See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT 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": 6,
+  "type": "request",
+  "name": "UpdateMetadataRequest",
+  // Version 1 allows specifying multiple endpoints for each broker.
+  //
+  // Version 2 adds the rack.
+  //
+  // Version 3 adds the listener name.
+  //
+  // Version 4 adds the offline replica list.
+  //
+  // Version 5 adds the broker epoch field and normalizes partitions by topic.
+  "validVersions": "0-6",
+  "flexibleVersions": "6+",
+  "fields": [
+    { "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",
+      "about": "The broker epoch." },
+    { "name": "UngroupedPartitionStates", "type": "[]UpdateMetadataPartitionState", "versions": "0-4",
+      "about": "In older versions of this RPC, each partition that we would like to update." },
+    { "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." }
+    ]},
+    { "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": "[]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+", "ignorable": true,
+            "about": "The listener name." },
+          { "name": "SecurityProtocol", "type": "int16", "versions": "1+",
+            "about": "The security protocol type." }
+        ]},
+        { "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
new file mode 100644
index 0000000000000..aaebed04ab108
--- /dev/null
+++ b/clients/src/main/resources/common/message/UpdateMetadataResponse.json
@@ -0,0 +1,27 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements.  See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT 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": 6,
+  "type": "response",
+  "name": "UpdateMetadataResponse",
+  // Versions 1, 2, 3, 4, and 5 are the same as version 0
+  "validVersions": "0-6",
+  "flexibleVersions": "6+",
+  "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/VoteRequest.json b/clients/src/main/resources/common/message/VoteRequest.json
new file mode 100644
index 0000000000000..3926ba7c7d6af
--- /dev/null
+++ b/clients/src/main/resources/common/message/VoteRequest.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": 52,
+  "type": "request",
+  "name": "VoteRequest",
+  "validVersions": "0",
+  "flexibleVersions": "0+",
+  "fields": [
+    { "name": "ClusterId", "type": "string", "versions": "0+",
+      "nullableVersions": "0+", "default": "null"},
+    { "name": "Topics", "type": "[]TopicData",
+      "versions": "0+", "fields": [
+      { "name": "TopicName", "type": "string", "versions": "0+", "entityType": "topicName",
+        "about": "The topic name." },
+      { "name": "Partitions", "type": "[]PartitionData",
+        "versions": "0+", "fields": [
+        { "name": "PartitionIndex", "type": "int32", "versions": "0+",
+          "about": "The partition index." },
+        { "name": "CandidateEpoch", "type": "int32", "versions": "0+",
+          "about": "The bumped epoch of the candidate sending the request"},
+        { "name": "CandidateId", "type": "int32", "versions": "0+",
+          "about": "The ID of the voter sending the request"},
+        { "name": "LastOffsetEpoch", "type": "int32", "versions": "0+",
+          "about": "The epoch of the last record written to the metadata log"},
+        { "name": "LastOffset", "type": "int64", "versions": "0+",
+          "about": "The offset of the last record written to the metadata log"}
+      ]
+      }
+    ]
+    }
+  ]
+}
diff --git a/clients/src/main/resources/common/message/VoteResponse.json b/clients/src/main/resources/common/message/VoteResponse.json
new file mode 100644
index 0000000000000..83d1a8dd4ee9c
--- /dev/null
+++ b/clients/src/main/resources/common/message/VoteResponse.json
@@ -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.
+
+{
+  "apiKey": 52,
+  "type": "response",
+  "name": "VoteResponse",
+  "validVersions": "0",
+  "flexibleVersions": "0+",
+  "fields": [
+    { "name": "ErrorCode", "type": "int16", "versions": "0+",
+      "about": "The top level error code."},
+    { "name": "Topics", "type": "[]TopicData",
+      "versions": "0+", "fields": [
+      { "name": "TopicName", "type": "string", "versions": "0+", "entityType": "topicName",
+        "about": "The topic name." },
+      { "name": "Partitions", "type": "[]PartitionData",
+        "versions": "0+", "fields": [
+        { "name": "PartitionIndex", "type": "int32", "versions": "0+",
+          "about": "The partition index." },
+        { "name": "ErrorCode", "type": "int16", "versions": "0+"},
+        { "name": "LeaderId", "type": "int32", "versions": "0+",
+          "about": "The ID of the current leader or -1 if the leader is unknown."},
+        { "name": "LeaderEpoch", "type": "int32", "versions": "0+",
+          "about": "The latest known leader epoch"},
+        { "name": "VoteGranted", "type": "bool", "versions": "0+",
+          "about": "True if the vote was granted and false otherwise"}
+      ]
+      }
+    ]
+    }
+  ]
+}
diff --git a/clients/src/main/resources/common/message/WriteTxnMarkersRequest.json b/clients/src/main/resources/common/message/WriteTxnMarkersRequest.json
new file mode 100644
index 0000000000000..3fdfa05bf9773
--- /dev/null
+++ b/clients/src/main/resources/common/message/WriteTxnMarkersRequest.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": 27,
+  "type": "request",
+  "name": "WriteTxnMarkersRequest",
+  // Version 1 enables flexible versions.
+  "validVersions": "0-1",
+  "flexibleVersions": "1+",
+  "fields": [
+    { "name": "Markers", "type": "[]WritableTxnMarker", "versions": "0+",
+      "about": "The transaction markers to be written.", "fields": [
+      { "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." },
+      { "name": "TransactionResult", "type": "bool", "versions": "0+",
+        "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+", "entityType": "topicName",
+          "about": "The topic name." },
+        { "name": "PartitionIndexes", "type": "[]int32", "versions": "0+",
+          "about": "The indexes of the partitions to write transaction markers for." }
+      ]},
+      { "name": "CoordinatorEpoch", "type": "int32", "versions": "0+",
+        "about": "Epoch associated with the transaction state partition hosted by this transaction coordinator" }
+    ]}
+  ]
+}
diff --git a/clients/src/main/resources/common/message/WriteTxnMarkersResponse.json b/clients/src/main/resources/common/message/WriteTxnMarkersResponse.json
new file mode 100644
index 0000000000000..59b8d66f2d47f
--- /dev/null
+++ b/clients/src/main/resources/common/message/WriteTxnMarkersResponse.json
@@ -0,0 +1,42 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements.  See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+{
+  "apiKey": 27,
+  "type": "response",
+  "name": "WriteTxnMarkersResponse",
+  "validVersions": "0-1",
+  // Version 1 enables flexible versions.
+  "flexibleVersions": "1+",
+  "fields": [
+    { "name": "Markers", "type": "[]WritableTxnMarkerResult", "versions": "0+",
+      "about": "The results for writing makers.", "fields": [
+      { "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+", "entityType": "topicName",
+          "about": "The topic name." },
+        { "name": "Partitions", "type": "[]WritableTxnMarkerPartitionResult", "versions": "0+",
+          "about": "The results by partition.", "fields": [
+          { "name": "PartitionIndex", "type": "int32", "versions": "0+",
+            "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/test/java/org/apache/kafka/clients/ApiVersionsTest.java b/clients/src/test/java/org/apache/kafka/clients/ApiVersionsTest.java
index 654a6064d5401..ab1a5cea7fe23 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,12 +32,10 @@ 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");
         assertEquals(RecordBatch.CURRENT_MAGIC_VALUE, apiVersions.maxUsableProduceMagic());
     }
-
 }
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 d19b0be4be49f..3a281f8e952b4 100644
--- a/clients/src/test/java/org/apache/kafka/clients/ClientUtilsTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/ClientUtilsTest.java
@@ -16,40 +16,114 @@
  */
 package org.apache.kafka.clients;
 
-import org.apache.kafka.common.config.ConfigException;
-import org.junit.Test;
-import static org.junit.Assert.assertEquals;
-
+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 {
 
+
     @Test
-    public void testParseAndValidateAddresses() {
-        check("127.0.0.1:8000");
-        check("mydomain.com:8080");
-        check("[::1]:8000");
-        check("[2001:db8:85a3:8d3:1319:8a2e:370:7348]:1234", "mydomain.com:10000");
-        List validatedAddresses = check("some.invalid.hostname.foo.bar.local:9999", "mydomain.com:10000");
+    public void testParseAndValidateAddresses() throws UnknownHostException {
+        checkWithoutLookup("127.0.0.1:8000");
+        checkWithoutLookup("localhost:8080");
+        checkWithoutLookup("[::1]:8000");
+        checkWithoutLookup("[2001:db8:85a3:8d3:1319:8a2e:370:7348]:1234", "localhost:10000");
+        List validatedAddresses = checkWithoutLookup("localhost:10000");
         assertEquals(1, validatedAddresses.size());
         InetSocketAddress onlyAddress = validatedAddresses.get(0);
-        assertEquals("mydomain.com", onlyAddress.getHostName());
+        assertEquals("localhost", onlyAddress.getHostName());
         assertEquals(10000, onlyAddress.getPort());
     }
 
+    @Test
+    public void testParseAndValidateAddressesWithReverseLookup() {
+        checkWithoutLookup("127.0.0.1:8000");
+        checkWithoutLookup("localhost:8080");
+        checkWithoutLookup("[::1]:8000");
+        checkWithoutLookup("[2001:db8:85a3:8d3:1319:8a2e:370:7348]:1234", "localhost:10000");
+
+        // With lookup of example.com, either one or two addresses are expected depending on
+        // whether ipv4 and ipv6 are enabled
+        List validatedAddresses = checkWithLookup(Arrays.asList("example.com:10000"));
+        assertTrue("Unexpected addresses " + validatedAddresses, validatedAddresses.size() >= 1);
+        List validatedHostNames = validatedAddresses.stream().map(InetSocketAddress::getHostName)
+                .collect(Collectors.toList());
+        List expectedHostNames = Arrays.asList("93.184.216.34", "2606:2800:220:1:248:1893:25c8:1946");
+        assertTrue("Unexpected addresses " + validatedHostNames, expectedHostNames.containsAll(validatedHostNames));
+        validatedAddresses.forEach(address -> assertEquals(10000, address.getPort()));
+    }
+
+    @Test(expected = IllegalArgumentException.class)
+    public void testInvalidConfig() {
+        ClientUtils.parseAndValidateAddresses(Arrays.asList("localhost:10000"), "random.value");
+    }
+
     @Test(expected = ConfigException.class)
     public void testNoPort() {
-        check("127.0.0.1");
+        checkWithoutLookup("127.0.0.1");
     }
 
     @Test(expected = ConfigException.class)
     public void testOnlyBadHostname() {
-        check("some.invalid.hostname.foo.bar.local:9999");
+        checkWithoutLookup("some.invalid.hostname.foo.bar.local:9999");
     }
 
-    private List check(String... url) {
-        return ClientUtils.parseAndValidateAddresses(Arrays.asList(url));
+    @Test
+    public void testFilterPreferredAddresses() throws UnknownHostException {
+        InetAddress ipv4 = InetAddress.getByName("192.0.0.1");
+        InetAddress ipv6 = InetAddress.getByName("::1");
+
+        InetAddress[] ipv4First = new InetAddress[]{ipv4, ipv6, ipv4};
+        List result = ClientUtils.filterPreferredAddresses(ipv4First);
+        assertTrue(result.contains(ipv4));
+        assertFalse(result.contains(ipv6));
+        assertEquals(2, result.size());
+
+        InetAddress[] ipv6First = new InetAddress[]{ipv6, ipv4, ipv4};
+        result = ClientUtils.filterPreferredAddresses(ipv6First);
+        assertTrue(result.contains(ipv6));
+        assertFalse(result.contains(ipv4));
+        assertEquals(1, result.size());
     }
+
+    @Test(expected = UnknownHostException.class)
+    public void testResolveUnknownHostException() throws UnknownHostException {
+        ClientUtils.resolve("some.invalid.hostname.foo.bar.local", ClientDnsLookup.USE_ALL_DNS_IPS);
+    }
+
+    @Test
+    public void testResolveDnsLookup() throws UnknownHostException {
+        // Note that kafka.apache.org resolves to at least 2 IP addresses
+        assertEquals(1, ClientUtils.resolve("kafka.apache.org", ClientDnsLookup.DEFAULT).size());
+    }
+
+    @Test
+    public void testResolveDnsLookupAllIps() throws UnknownHostException {
+        // Note that kafka.apache.org resolves to at least 2 IP addresses
+        assertTrue(ClientUtils.resolve("kafka.apache.org", ClientDnsLookup.USE_ALL_DNS_IPS).size() > 1);
+    }
+
+    @Test
+    public void testResolveDnsLookupResolveCanonicalBootstrapServers() throws UnknownHostException {
+        // Note that kafka.apache.org resolves to at least 2 IP addresses
+        assertTrue(ClientUtils.resolve("kafka.apache.org", ClientDnsLookup.RESOLVE_CANONICAL_BOOTSTRAP_SERVERS_ONLY).size() > 1);
+    }
+
+    private List checkWithoutLookup(String... url) {
+        return ClientUtils.parseAndValidateAddresses(Arrays.asList(url), ClientDnsLookup.USE_ALL_DNS_IPS);
+    }
+
+    private List checkWithLookup(List url) {
+        return ClientUtils.parseAndValidateAddresses(url, ClientDnsLookup.RESOLVE_CANONICAL_BOOTSTRAP_SERVERS_ONLY);
+    }
+
 }
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 3b246f4cad2c4..436c29fceb15d 100644
--- a/clients/src/test/java/org/apache/kafka/clients/ClusterConnectionStatesTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/ClusterConnectionStatesTest.java
@@ -19,9 +19,19 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotSame;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertTrue;
 
+import java.lang.reflect.Field;
+import java.lang.reflect.Method;
+import java.net.InetAddress;
+import java.net.UnknownHostException;
+
+import java.util.List;
 import org.apache.kafka.common.errors.AuthenticationException;
+import org.apache.kafka.common.utils.LogContext;
 import org.apache.kafka.common.utils.MockTime;
 import org.junit.Before;
 import org.junit.Test;
@@ -31,15 +41,24 @@ public class ClusterConnectionStatesTest {
     private final MockTime time = new MockTime();
     private final long reconnectBackoffMs = 10 * 1000;
     private final long reconnectBackoffMax = 60 * 1000;
-    private final double reconnectBackoffJitter = 0.2;
+    private final long connectionSetupTimeoutMs = 10 * 1000;
+    private final long connectionSetupTimeoutMaxMs = 127 * 1000;
+    private final int reconnectBackoffExpBase = ClusterConnectionStates.RECONNECT_BACKOFF_EXP_BASE;
+    private final double reconnectBackoffJitter = ClusterConnectionStates.RECONNECT_BACKOFF_JITTER;
+    private final int connectionSetupTimeoutExpBase = ClusterConnectionStates.CONNECTION_SETUP_TIMEOUT_EXP_BASE;
+    private final double connectionSetupTimeoutJitter = ClusterConnectionStates.CONNECTION_SETUP_TIMEOUT_JITTER;
     private final String nodeId1 = "1001";
     private final String nodeId2 = "2002";
+    private final String nodeId3 = "3003";
+    private final String hostTwoIps = "kafka.apache.org";
 
     private ClusterConnectionStates connectionStates;
 
     @Before
     public void setup() {
-        this.connectionStates = new ClusterConnectionStates(reconnectBackoffMs, reconnectBackoffMax);
+        this.connectionStates = new ClusterConnectionStates(
+                reconnectBackoffMs, reconnectBackoffMax,
+                connectionSetupTimeoutMs, connectionSetupTimeoutMaxMs, new LogContext());
     }
 
     @Test
@@ -47,20 +66,20 @@ public void testClusterConnectionStateChanges() {
         assertTrue(connectionStates.canConnect(nodeId1, time.milliseconds()));
 
         // Start connecting to Node and check state
-        connectionStates.connecting(nodeId1, time.milliseconds());
+        connectionStates.connecting(nodeId1, time.milliseconds(), "localhost", ClientDnsLookup.DEFAULT);
         assertEquals(connectionStates.connectionState(nodeId1), ConnectionState.CONNECTING);
         assertTrue(connectionStates.isConnecting(nodeId1));
-        assertFalse(connectionStates.isReady(nodeId1));
+        assertFalse(connectionStates.isReady(nodeId1, time.milliseconds()));
         assertFalse(connectionStates.isBlackedOut(nodeId1, time.milliseconds()));
-        assertFalse(connectionStates.hasReadyNodes());
+        assertFalse(connectionStates.hasReadyNodes(time.milliseconds()));
 
         time.sleep(100);
 
         // Successful connection
         connectionStates.ready(nodeId1);
         assertEquals(connectionStates.connectionState(nodeId1), ConnectionState.READY);
-        assertTrue(connectionStates.isReady(nodeId1));
-        assertTrue(connectionStates.hasReadyNodes());
+        assertTrue(connectionStates.isReady(nodeId1, time.milliseconds()));
+        assertTrue(connectionStates.hasReadyNodes(time.milliseconds()));
         assertFalse(connectionStates.isConnecting(nodeId1));
         assertFalse(connectionStates.isBlackedOut(nodeId1, time.milliseconds()));
         assertEquals(connectionStates.connectionDelay(nodeId1, time.milliseconds()), Long.MAX_VALUE);
@@ -73,7 +92,7 @@ public void testClusterConnectionStateChanges() {
         assertTrue(connectionStates.isDisconnected(nodeId1));
         assertTrue(connectionStates.isBlackedOut(nodeId1, time.milliseconds()));
         assertFalse(connectionStates.isConnecting(nodeId1));
-        assertFalse(connectionStates.hasReadyNodes());
+        assertFalse(connectionStates.hasReadyNodes(time.milliseconds()));
         assertFalse(connectionStates.canConnect(nodeId1, time.milliseconds()));
 
         // After disconnecting we expect a backoff value equal to the reconnect.backoff.ms setting (plus minus 20% jitter)
@@ -91,29 +110,29 @@ public void testMultipleNodeConnectionStates() {
         // Check initial state, allowed to connect to all nodes, but no nodes shown as ready
         assertTrue(connectionStates.canConnect(nodeId1, time.milliseconds()));
         assertTrue(connectionStates.canConnect(nodeId2, time.milliseconds()));
-        assertFalse(connectionStates.hasReadyNodes());
+        assertFalse(connectionStates.hasReadyNodes(time.milliseconds()));
 
         // Start connecting one node and check that the pool only shows ready nodes after
         // successful connect
-        connectionStates.connecting(nodeId2, time.milliseconds());
-        assertFalse(connectionStates.hasReadyNodes());
+        connectionStates.connecting(nodeId2, time.milliseconds(), "localhost", ClientDnsLookup.DEFAULT);
+        assertFalse(connectionStates.hasReadyNodes(time.milliseconds()));
         time.sleep(1000);
         connectionStates.ready(nodeId2);
-        assertTrue(connectionStates.hasReadyNodes());
+        assertTrue(connectionStates.hasReadyNodes(time.milliseconds()));
 
         // Connect second node and check that both are shown as ready, pool should immediately
         // show ready nodes, since node2 is already connected
-        connectionStates.connecting(nodeId1, time.milliseconds());
-        assertTrue(connectionStates.hasReadyNodes());
+        connectionStates.connecting(nodeId1, time.milliseconds(), "localhost", ClientDnsLookup.DEFAULT);
+        assertTrue(connectionStates.hasReadyNodes(time.milliseconds()));
         time.sleep(1000);
         connectionStates.ready(nodeId1);
-        assertTrue(connectionStates.hasReadyNodes());
+        assertTrue(connectionStates.hasReadyNodes(time.milliseconds()));
 
         time.sleep(12000);
 
         // disconnect nodes and check proper state of pool throughout
         connectionStates.disconnected(nodeId2, time.milliseconds());
-        assertTrue(connectionStates.hasReadyNodes());
+        assertTrue(connectionStates.hasReadyNodes(time.milliseconds()));
         assertTrue(connectionStates.isBlackedOut(nodeId2, time.milliseconds()));
         assertFalse(connectionStates.isBlackedOut(nodeId1, time.milliseconds()));
         time.sleep(connectionStates.connectionDelay(nodeId2, time.milliseconds()));
@@ -121,13 +140,13 @@ public void testMultipleNodeConnectionStates() {
         connectionStates.disconnected(nodeId1, time.milliseconds() + 1);
         assertTrue(connectionStates.isBlackedOut(nodeId1, time.milliseconds()));
         assertFalse(connectionStates.isBlackedOut(nodeId2, time.milliseconds()));
-        assertFalse(connectionStates.hasReadyNodes());
+        assertFalse(connectionStates.hasReadyNodes(time.milliseconds()));
     }
 
     @Test
     public void testAuthorizationFailed() {
         // Try connecting
-        connectionStates.connecting(nodeId1, time.milliseconds());
+        connectionStates.connecting(nodeId1, time.milliseconds(), "localhost", ClientDnsLookup.DEFAULT);
 
         time.sleep(100);
 
@@ -135,18 +154,19 @@ public void testAuthorizationFailed() {
         time.sleep(1000);
         assertEquals(connectionStates.connectionState(nodeId1), ConnectionState.AUTHENTICATION_FAILED);
         assertTrue(connectionStates.authenticationException(nodeId1) instanceof AuthenticationException);
-        assertFalse(connectionStates.hasReadyNodes());
+        assertFalse(connectionStates.hasReadyNodes(time.milliseconds()));
         assertFalse(connectionStates.canConnect(nodeId1, time.milliseconds()));
 
         time.sleep(connectionStates.connectionDelay(nodeId1, time.milliseconds()) + 1);
 
         assertTrue(connectionStates.canConnect(nodeId1, time.milliseconds()));
+        connectionStates.ready(nodeId1);
+        assertNull(connectionStates.authenticationException(nodeId1));
     }
 
-
     @Test
     public void testRemoveNode() {
-        connectionStates.connecting(nodeId1, time.milliseconds());
+        connectionStates.connecting(nodeId1, time.milliseconds(), "localhost", ClientDnsLookup.DEFAULT);
         time.sleep(1000);
         connectionStates.ready(nodeId1);
         time.sleep(10000);
@@ -162,7 +182,7 @@ public void testRemoveNode() {
     @Test
     public void testMaxReconnectBackoff() {
         long effectiveMaxReconnectBackoff = Math.round(reconnectBackoffMax * (1 + reconnectBackoffJitter));
-        connectionStates.connecting(nodeId1, time.milliseconds());
+        connectionStates.connecting(nodeId1, time.milliseconds(), "localhost", ClientDnsLookup.DEFAULT);
         time.sleep(1000);
         connectionStates.disconnected(nodeId1, time.milliseconds());
 
@@ -173,7 +193,7 @@ public void testMaxReconnectBackoff() {
             assertFalse(connectionStates.canConnect(nodeId1, time.milliseconds()));
             time.sleep(reconnectBackoff + 1);
             assertTrue(connectionStates.canConnect(nodeId1, time.milliseconds()));
-            connectionStates.connecting(nodeId1, time.milliseconds());
+            connectionStates.connecting(nodeId1, time.milliseconds(), "localhost", ClientDnsLookup.DEFAULT);
             time.sleep(10);
             connectionStates.disconnected(nodeId1, time.milliseconds());
         }
@@ -181,14 +201,12 @@ public void testMaxReconnectBackoff() {
 
     @Test
     public void testExponentialReconnectBackoff() {
-        // Calculate fixed components for backoff process
-        final int reconnectBackoffExpBase = 2;
         double reconnectBackoffMaxExp = Math.log(reconnectBackoffMax / (double) Math.max(reconnectBackoffMs, 1))
             / Math.log(reconnectBackoffExpBase);
 
         // Run through 10 disconnects and check that reconnect backoff value is within expected range for every attempt
         for (int i = 0; i < 10; i++) {
-            connectionStates.connecting(nodeId1, time.milliseconds());
+            connectionStates.connecting(nodeId1, time.milliseconds(), "localhost", ClientDnsLookup.DEFAULT);
             connectionStates.disconnected(nodeId1, time.milliseconds());
             // Calculate expected backoff value without jitter
             long expectedBackoff = Math.round(Math.pow(reconnectBackoffExpBase, Math.min(i, reconnectBackoffMaxExp))
@@ -198,4 +216,197 @@ public void testExponentialReconnectBackoff() {
             time.sleep(connectionStates.connectionDelay(nodeId1, time.milliseconds()) + 1);
         }
     }
+
+    @Test
+    public void testThrottled() {
+        connectionStates.connecting(nodeId1, time.milliseconds(), "localhost", ClientDnsLookup.DEFAULT);
+        time.sleep(1000);
+        connectionStates.ready(nodeId1);
+        time.sleep(10000);
+
+        // Initially not throttled.
+        assertEquals(0, connectionStates.throttleDelayMs(nodeId1, time.milliseconds()));
+
+        // Throttle for 100ms from now.
+        connectionStates.throttle(nodeId1, time.milliseconds() + 100);
+        assertEquals(100, connectionStates.throttleDelayMs(nodeId1, time.milliseconds()));
+
+        // Still throttled after 50ms. The remaining delay is 50ms. The poll delay should be same as throttling delay.
+        time.sleep(50);
+        assertEquals(50, connectionStates.throttleDelayMs(nodeId1, time.milliseconds()));
+        assertEquals(50, connectionStates.pollDelayMs(nodeId1, time.milliseconds()));
+
+        // Not throttled anymore when the deadline is reached. The poll delay should be same as connection delay.
+        time.sleep(50);
+        assertEquals(0, connectionStates.throttleDelayMs(nodeId1, time.milliseconds()));
+        assertEquals(connectionStates.connectionDelay(nodeId1, time.milliseconds()),
+            connectionStates.pollDelayMs(nodeId1, time.milliseconds()));
+    }
+
+    @Test
+    public void testSingleIPWithDefault() throws UnknownHostException {
+        connectionStates.connecting(nodeId1, time.milliseconds(), "localhost", ClientDnsLookup.DEFAULT);
+        InetAddress currAddress = connectionStates.currentAddress(nodeId1);
+        connectionStates.connecting(nodeId1, time.milliseconds(), "localhost", ClientDnsLookup.DEFAULT);
+        assertSame(currAddress, connectionStates.currentAddress(nodeId1));
+    }
+
+    @Test
+    public void testSingleIPWithUseAll() throws UnknownHostException {
+        assertEquals(1, ClientUtils.resolve("localhost", ClientDnsLookup.USE_ALL_DNS_IPS).size());
+
+        connectionStates.connecting(nodeId1, time.milliseconds(), "localhost", ClientDnsLookup.USE_ALL_DNS_IPS);
+        InetAddress currAddress = connectionStates.currentAddress(nodeId1);
+        connectionStates.connecting(nodeId1, time.milliseconds(), "localhost", ClientDnsLookup.USE_ALL_DNS_IPS);
+        assertSame(currAddress, connectionStates.currentAddress(nodeId1));
+    }
+
+    @Test
+    public void testMultipleIPsWithDefault() throws UnknownHostException {
+        assertTrue(ClientUtils.resolve(hostTwoIps, ClientDnsLookup.USE_ALL_DNS_IPS).size() > 1);
+
+        connectionStates.connecting(nodeId1, time.milliseconds(), hostTwoIps, ClientDnsLookup.DEFAULT);
+        InetAddress currAddress = connectionStates.currentAddress(nodeId1);
+        connectionStates.connecting(nodeId1, time.milliseconds(), hostTwoIps, ClientDnsLookup.DEFAULT);
+        assertSame(currAddress, connectionStates.currentAddress(nodeId1));
+    }
+
+    @Test
+    public void testMultipleIPsWithUseAll() throws UnknownHostException {
+        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);
+        connectionStates.connecting(nodeId1, time.milliseconds(), hostTwoIps, ClientDnsLookup.USE_ALL_DNS_IPS);
+        InetAddress addr2 = connectionStates.currentAddress(nodeId1);
+        assertNotSame(addr1, addr2);
+        connectionStates.connecting(nodeId1, time.milliseconds(), hostTwoIps, ClientDnsLookup.USE_ALL_DNS_IPS);
+        InetAddress addr3 = connectionStates.currentAddress(nodeId1);
+        assertNotSame(addr1, addr3);
+    }
+
+    @Test
+    public void testHostResolveChange() throws UnknownHostException, ReflectiveOperationException {
+        assertTrue(ClientUtils.resolve(hostTwoIps, ClientDnsLookup.USE_ALL_DNS_IPS).size() > 1);
+
+        connectionStates.connecting(nodeId1, time.milliseconds(), hostTwoIps, ClientDnsLookup.DEFAULT);
+        InetAddress addr1 = connectionStates.currentAddress(nodeId1);
+
+        // reflection to simulate host change in DNS lookup
+        Method nodeStateMethod = connectionStates.getClass().getDeclaredMethod("nodeState", String.class);
+        nodeStateMethod.setAccessible(true);
+        Object nodeState = nodeStateMethod.invoke(connectionStates, nodeId1);
+        Field hostField = nodeState.getClass().getDeclaredField("host");
+        hostField.setAccessible(true);
+        hostField.set(nodeState, "localhost");
+
+        connectionStates.connecting(nodeId1, time.milliseconds(), "localhost", ClientDnsLookup.DEFAULT);
+        InetAddress addr2 = connectionStates.currentAddress(nodeId1);
+
+        assertNotSame(addr1, addr2);
+    }
+
+    @Test
+    public void testNodeWithNewHostname() throws UnknownHostException {
+        connectionStates.connecting(nodeId1, time.milliseconds(), "localhost", ClientDnsLookup.DEFAULT);
+        InetAddress addr1 = connectionStates.currentAddress(nodeId1);
+
+        connectionStates.connecting(nodeId1, time.milliseconds(), hostTwoIps, ClientDnsLookup.DEFAULT);
+        InetAddress addr2 = connectionStates.currentAddress(nodeId1);
+
+        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));
+    }
+
+    @Test
+    public void testExponentialConnectionSetupTimeout() {
+        assertTrue(connectionStates.canConnect(nodeId1, time.milliseconds()));
+
+        // Check the exponential timeout growth
+        for (int n = 0; n <= Math.log((double) connectionSetupTimeoutMaxMs / connectionSetupTimeoutMs) / Math.log(connectionSetupTimeoutExpBase); n++) {
+            connectionStates.connecting(nodeId1, time.milliseconds(), "localhost", ClientDnsLookup.DEFAULT);
+            assertTrue(connectionStates.connectingNodes().contains(nodeId1));
+            assertEquals(connectionSetupTimeoutMs * Math.pow(connectionSetupTimeoutExpBase, n),
+                    connectionStates.connectionSetupTimeoutMs(nodeId1),
+                    connectionSetupTimeoutMs * Math.pow(connectionSetupTimeoutExpBase, n) * connectionSetupTimeoutJitter);
+            connectionStates.disconnected(nodeId1, time.milliseconds());
+            assertFalse(connectionStates.connectingNodes().contains(nodeId1));
+        }
+
+        // Check the timeout value upper bound
+        connectionStates.connecting(nodeId1, time.milliseconds(), "localhost", ClientDnsLookup.DEFAULT);
+        assertEquals(connectionSetupTimeoutMaxMs,
+                connectionStates.connectionSetupTimeoutMs(nodeId1),
+                connectionSetupTimeoutMaxMs * connectionSetupTimeoutJitter);
+        assertTrue(connectionStates.connectingNodes().contains(nodeId1));
+
+        // Should reset the timeout value to the init value
+        connectionStates.ready(nodeId1);
+        assertEquals(connectionSetupTimeoutMs,
+                connectionStates.connectionSetupTimeoutMs(nodeId1),
+                connectionSetupTimeoutMs * connectionSetupTimeoutJitter);
+        assertFalse(connectionStates.connectingNodes().contains(nodeId1));
+        connectionStates.disconnected(nodeId1, time.milliseconds());
+
+        // Check if the connection state transition from ready to disconnected
+        // won't increase the timeout value
+        connectionStates.connecting(nodeId1, time.milliseconds(), "localhost", ClientDnsLookup.DEFAULT);
+        assertEquals(connectionSetupTimeoutMs,
+                connectionStates.connectionSetupTimeoutMs(nodeId1),
+                connectionSetupTimeoutMs * connectionSetupTimeoutJitter);
+        assertTrue(connectionStates.connectingNodes().contains(nodeId1));
+    }
+
+    @Test
+    public void testTimedOutConnections() {
+        // Initiate two connections
+        connectionStates.connecting(nodeId1, time.milliseconds(), "localhost", ClientDnsLookup.DEFAULT);
+        connectionStates.connecting(nodeId2, time.milliseconds(), "localhost", ClientDnsLookup.DEFAULT);
+
+        // Expect no timed out connections
+        assertEquals(0, connectionStates.nodesWithConnectionSetupTimeout(time.milliseconds()).size());
+
+        // Advance time by half of the connection setup timeout
+        time.sleep(connectionSetupTimeoutMs / 2);
+
+        // Initiate a third connection
+        connectionStates.connecting(nodeId3, time.milliseconds(), "localhost", ClientDnsLookup.DEFAULT);
+
+        // Advance time beyond the connection setup timeout (+ max jitter) for the first two connections
+        time.sleep((long) (connectionSetupTimeoutMs / 2 + connectionSetupTimeoutMs * connectionSetupTimeoutJitter));
+
+        // Expect two timed out connections.
+        List timedOutConnections = connectionStates.nodesWithConnectionSetupTimeout(time.milliseconds());
+        assertEquals(2, timedOutConnections.size());
+        assertTrue(timedOutConnections.contains(nodeId1));
+        assertTrue(timedOutConnections.contains(nodeId2));
+
+        // Disconnect the first two connections
+        connectionStates.disconnected(nodeId1, time.milliseconds());
+        connectionStates.disconnected(nodeId2, time.milliseconds());
+
+        // Advance time beyond the connection setup timeout (+ max jitter) for the third connections
+        time.sleep((long) (connectionSetupTimeoutMs / 2 + connectionSetupTimeoutMs * connectionSetupTimeoutJitter));
+
+        // Expect one timed out connection
+        timedOutConnections = connectionStates.nodesWithConnectionSetupTimeout(time.milliseconds());
+        assertEquals(1, timedOutConnections.size());
+        assertTrue(timedOutConnections.contains(nodeId3));
+
+        // Disconnect the third connection
+        connectionStates.disconnected(nodeId3, time.milliseconds());
+
+        // Expect no timed out connections
+        assertEquals(0, connectionStates.nodesWithConnectionSetupTimeout(time.milliseconds()).size());
+    }
 }
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/FetchSessionHandlerTest.java b/clients/src/test/java/org/apache/kafka/clients/FetchSessionHandlerTest.java
new file mode 100644
index 0000000000000..810c4ecb6faeb
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/FetchSessionHandlerTest.java
@@ -0,0 +1,388 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.TopicPartition;
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.record.MemoryRecords;
+import org.apache.kafka.common.requests.FetchRequest;
+import org.apache.kafka.common.requests.FetchResponse;
+import org.apache.kafka.common.utils.LogContext;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.Timeout;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Comparator;
+import java.util.Iterator;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.TreeSet;
+
+import static org.apache.kafka.common.requests.FetchMetadata.INITIAL_EPOCH;
+import static org.apache.kafka.common.requests.FetchMetadata.INVALID_SESSION_ID;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+
+/**
+ * A unit test for FetchSessionHandler.
+ */
+public class FetchSessionHandlerTest {
+    @Rule
+    final public Timeout globalTimeout = Timeout.millis(120000);
+
+    private static final LogContext LOG_CONTEXT = new LogContext("[FetchSessionHandler]=");
+
+    /**
+     * Create a set of TopicPartitions.  We use a TreeSet, in order to get a deterministic
+     * ordering for test purposes.
+     */
+    private static Set toSet(TopicPartition... arr) {
+        TreeSet set = new TreeSet<>(new Comparator() {
+            @Override
+            public int compare(TopicPartition o1, TopicPartition o2) {
+                return o1.toString().compareTo(o2.toString());
+            }
+        });
+        set.addAll(Arrays.asList(arr));
+        return set;
+    }
+
+    @Test
+    public void testFindMissing() {
+        TopicPartition foo0 = new TopicPartition("foo", 0);
+        TopicPartition foo1 = new TopicPartition("foo", 1);
+        TopicPartition bar0 = new TopicPartition("bar", 0);
+        TopicPartition bar1 = new TopicPartition("bar", 1);
+        TopicPartition baz0 = new TopicPartition("baz", 0);
+        TopicPartition baz1 = new TopicPartition("baz", 1);
+        assertEquals(toSet(), FetchSessionHandler.findMissing(toSet(foo0), toSet(foo0)));
+        assertEquals(toSet(foo0), FetchSessionHandler.findMissing(toSet(foo0), toSet(foo1)));
+        assertEquals(toSet(foo0, foo1),
+            FetchSessionHandler.findMissing(toSet(foo0, foo1), toSet(baz0)));
+        assertEquals(toSet(bar1, foo0, foo1),
+            FetchSessionHandler.findMissing(toSet(foo0, foo1, bar0, bar1),
+                toSet(bar0, baz0, baz1)));
+        assertEquals(toSet(),
+            FetchSessionHandler.findMissing(toSet(foo0, foo1, bar0, bar1, baz1),
+                toSet(foo0, foo1, bar0, bar1, baz0, baz1)));
+    }
+
+    private static final class ReqEntry {
+        final TopicPartition part;
+        final FetchRequest.PartitionData data;
+
+        ReqEntry(String topic, int partition, long fetchOffset, long logStartOffset, int maxBytes) {
+            this.part = new TopicPartition(topic, partition);
+            this.data = new FetchRequest.PartitionData(fetchOffset, logStartOffset, maxBytes, Optional.empty());
+        }
+    }
+
+    private static LinkedHashMap reqMap(ReqEntry... entries) {
+        LinkedHashMap map = new LinkedHashMap<>();
+        for (ReqEntry entry : entries) {
+            map.put(entry.part, entry.data);
+        }
+        return map;
+    }
+
+    private static void assertMapEquals(Map expected,
+                                        Map actual) {
+        Iterator> expectedIter =
+            expected.entrySet().iterator();
+        Iterator> actualIter =
+            actual.entrySet().iterator();
+        int i = 1;
+        while (expectedIter.hasNext()) {
+            Map.Entry expectedEntry = expectedIter.next();
+            if (!actualIter.hasNext()) {
+                fail("Element " + i + " not found.");
+            }
+            Map.Entry actuaLEntry = actualIter.next();
+            assertEquals("Element " + i + " had a different TopicPartition than expected.",
+                expectedEntry.getKey(), actuaLEntry.getKey());
+            assertEquals("Element " + i + " had different PartitionData than expected.",
+                expectedEntry.getValue(), actuaLEntry.getValue());
+            i++;
+        }
+        if (expectedIter.hasNext()) {
+            fail("Unexpected element " + i + " found.");
+        }
+    }
+
+    @SafeVarargs
+    private static void assertMapsEqual(Map expected,
+                                        Map... actuals) {
+        for (Map actual : actuals) {
+            assertMapEquals(expected, actual);
+        }
+    }
+
+    private static void assertListEquals(List expected, List actual) {
+        for (TopicPartition expectedPart : expected) {
+            if (!actual.contains(expectedPart)) {
+                fail("Failed to find expected partition " + expectedPart);
+            }
+        }
+        for (TopicPartition actualPart : actual) {
+            if (!expected.contains(actualPart)) {
+                fail("Found unexpected partition " + actualPart);
+            }
+        }
+    }
+
+    private static final class RespEntry {
+        final TopicPartition part;
+        final FetchResponse.PartitionData data;
+
+        RespEntry(String topic, int partition, long highWatermark, long lastStableOffset) {
+            this.part = new TopicPartition(topic, partition);
+            this.data = new FetchResponse.PartitionData<>(
+                Errors.NONE,
+                highWatermark,
+                lastStableOffset,
+                0,
+                null,
+                null);
+        }
+    }
+
+    private static LinkedHashMap> respMap(RespEntry... entries) {
+        LinkedHashMap> map = new LinkedHashMap<>();
+        for (RespEntry entry : entries) {
+            map.put(entry.part, entry.data);
+        }
+        return map;
+    }
+
+    /**
+     * Test the handling of SESSIONLESS responses.
+     * Pre-KIP-227 brokers always supply this kind of response.
+     */
+    @Test
+    public void testSessionless() {
+        FetchSessionHandler handler = new FetchSessionHandler(LOG_CONTEXT, 1);
+        FetchSessionHandler.Builder builder = handler.newBuilder();
+        builder.add(new TopicPartition("foo", 0),
+            new FetchRequest.PartitionData(0, 100, 200, Optional.empty()));
+        builder.add(new TopicPartition("foo", 1),
+            new FetchRequest.PartitionData(10, 110, 210, Optional.empty()));
+        FetchSessionHandler.FetchRequestData data = builder.build();
+        assertMapsEqual(reqMap(new ReqEntry("foo", 0, 0, 100, 200),
+                               new ReqEntry("foo", 1, 10, 110, 210)),
+            data.toSend(), data.sessionPartitions());
+        assertEquals(INVALID_SESSION_ID, data.metadata().sessionId());
+        assertEquals(INITIAL_EPOCH, data.metadata().epoch());
+
+        FetchResponse resp = new FetchResponse<>(Errors.NONE,
+            respMap(new RespEntry("foo", 0, 0, 0),
+                    new RespEntry("foo", 1, 0, 0)),
+            0, INVALID_SESSION_ID);
+        handler.handleResponse(resp);
+
+        FetchSessionHandler.Builder builder2 = handler.newBuilder();
+        builder2.add(new TopicPartition("foo", 0),
+            new FetchRequest.PartitionData(0, 100, 200, Optional.empty()));
+        FetchSessionHandler.FetchRequestData data2 = builder2.build();
+        assertEquals(INVALID_SESSION_ID, data2.metadata().sessionId());
+        assertEquals(INITIAL_EPOCH, data2.metadata().epoch());
+        assertMapsEqual(reqMap(new ReqEntry("foo", 0, 0, 100, 200)),
+            data.toSend(), data.sessionPartitions());
+    }
+
+    /**
+     * Test handling an incremental fetch session.
+     */
+    @Test
+    public void testIncrementals() {
+        FetchSessionHandler handler = new FetchSessionHandler(LOG_CONTEXT, 1);
+        FetchSessionHandler.Builder builder = handler.newBuilder();
+        builder.add(new TopicPartition("foo", 0),
+            new FetchRequest.PartitionData(0, 100, 200, Optional.empty()));
+        builder.add(new TopicPartition("foo", 1),
+            new FetchRequest.PartitionData(10, 110, 210, Optional.empty()));
+        FetchSessionHandler.FetchRequestData data = builder.build();
+        assertMapsEqual(reqMap(new ReqEntry("foo", 0, 0, 100, 200),
+            new ReqEntry("foo", 1, 10, 110, 210)),
+            data.toSend(), data.sessionPartitions());
+        assertEquals(INVALID_SESSION_ID, data.metadata().sessionId());
+        assertEquals(INITIAL_EPOCH, data.metadata().epoch());
+
+        FetchResponse resp = new FetchResponse<>(Errors.NONE,
+            respMap(new RespEntry("foo", 0, 10, 20),
+                    new RespEntry("foo", 1, 10, 20)),
+            0, 123);
+        handler.handleResponse(resp);
+
+        // Test an incremental fetch request which adds one partition and modifies another.
+        FetchSessionHandler.Builder builder2 = handler.newBuilder();
+        builder2.add(new TopicPartition("foo", 0),
+            new FetchRequest.PartitionData(0, 100, 200, Optional.empty()));
+        builder2.add(new TopicPartition("foo", 1),
+            new FetchRequest.PartitionData(10, 120, 210, Optional.empty()));
+        builder2.add(new TopicPartition("bar", 0),
+            new FetchRequest.PartitionData(20, 200, 200, Optional.empty()));
+        FetchSessionHandler.FetchRequestData data2 = builder2.build();
+        assertFalse(data2.metadata().isFull());
+        assertMapEquals(reqMap(new ReqEntry("foo", 0, 0, 100, 200),
+                new ReqEntry("foo", 1, 10, 120, 210),
+                new ReqEntry("bar", 0, 20, 200, 200)),
+            data2.sessionPartitions());
+        assertMapEquals(reqMap(new ReqEntry("bar", 0, 20, 200, 200),
+                new ReqEntry("foo", 1, 10, 120, 210)),
+            data2.toSend());
+
+        FetchResponse resp2 = new FetchResponse<>(Errors.NONE,
+            respMap(new RespEntry("foo", 1, 20, 20)),
+            0, 123);
+        handler.handleResponse(resp2);
+
+        // Skip building a new request.  Test that handling an invalid fetch session epoch response results
+        // in a request which closes the session.
+        FetchResponse resp3 = new FetchResponse<>(Errors.INVALID_FETCH_SESSION_EPOCH, respMap(),
+            0, INVALID_SESSION_ID);
+        handler.handleResponse(resp3);
+
+        FetchSessionHandler.Builder builder4 = handler.newBuilder();
+        builder4.add(new TopicPartition("foo", 0),
+            new FetchRequest.PartitionData(0, 100, 200, Optional.empty()));
+        builder4.add(new TopicPartition("foo", 1),
+            new FetchRequest.PartitionData(10, 120, 210, Optional.empty()));
+        builder4.add(new TopicPartition("bar", 0),
+            new FetchRequest.PartitionData(20, 200, 200, Optional.empty()));
+        FetchSessionHandler.FetchRequestData data4 = builder4.build();
+        assertTrue(data4.metadata().isFull());
+        assertEquals(data2.metadata().sessionId(), data4.metadata().sessionId());
+        assertEquals(INITIAL_EPOCH, data4.metadata().epoch());
+        assertMapsEqual(reqMap(new ReqEntry("foo", 0, 0, 100, 200),
+            new ReqEntry("foo", 1, 10, 120, 210),
+            new ReqEntry("bar", 0, 20, 200, 200)),
+            data4.sessionPartitions(), data4.toSend());
+    }
+
+    /**
+     * Test that calling FetchSessionHandler#Builder#build twice fails.
+     */
+    @Test
+    public void testDoubleBuild() {
+        FetchSessionHandler handler = new FetchSessionHandler(LOG_CONTEXT, 1);
+        FetchSessionHandler.Builder builder = handler.newBuilder();
+        builder.add(new TopicPartition("foo", 0),
+            new FetchRequest.PartitionData(0, 100, 200, Optional.empty()));
+        builder.build();
+        try {
+            builder.build();
+            fail("Expected calling build twice to fail.");
+        } catch (Throwable t) {
+            // expected
+        }
+    }
+
+    @Test
+    public void testIncrementalPartitionRemoval() {
+        FetchSessionHandler handler = new FetchSessionHandler(LOG_CONTEXT, 1);
+        FetchSessionHandler.Builder builder = handler.newBuilder();
+        builder.add(new TopicPartition("foo", 0),
+            new FetchRequest.PartitionData(0, 100, 200, Optional.empty()));
+        builder.add(new TopicPartition("foo", 1),
+            new FetchRequest.PartitionData(10, 110, 210, Optional.empty()));
+        builder.add(new TopicPartition("bar", 0),
+            new FetchRequest.PartitionData(20, 120, 220, Optional.empty()));
+        FetchSessionHandler.FetchRequestData data = builder.build();
+        assertMapsEqual(reqMap(new ReqEntry("foo", 0, 0, 100, 200),
+            new ReqEntry("foo", 1, 10, 110, 210),
+            new ReqEntry("bar", 0, 20, 120, 220)),
+            data.toSend(), data.sessionPartitions());
+        assertTrue(data.metadata().isFull());
+
+        FetchResponse resp = new FetchResponse<>(Errors.NONE,
+            respMap(new RespEntry("foo", 0, 10, 20),
+                    new RespEntry("foo", 1, 10, 20),
+                    new RespEntry("bar", 0, 10, 20)),
+            0, 123);
+        handler.handleResponse(resp);
+
+        // Test an incremental fetch request which removes two partitions.
+        FetchSessionHandler.Builder builder2 = handler.newBuilder();
+        builder2.add(new TopicPartition("foo", 1),
+            new FetchRequest.PartitionData(10, 110, 210, Optional.empty()));
+        FetchSessionHandler.FetchRequestData data2 = builder2.build();
+        assertFalse(data2.metadata().isFull());
+        assertEquals(123, data2.metadata().sessionId());
+        assertEquals(1, data2.metadata().epoch());
+        assertMapEquals(reqMap(new ReqEntry("foo", 1, 10, 110, 210)),
+            data2.sessionPartitions());
+        assertMapEquals(reqMap(), data2.toSend());
+        ArrayList expectedToForget2 = new ArrayList<>();
+        expectedToForget2.add(new TopicPartition("foo", 0));
+        expectedToForget2.add(new TopicPartition("bar", 0));
+        assertListEquals(expectedToForget2, data2.toForget());
+
+        // A FETCH_SESSION_ID_NOT_FOUND response triggers us to close the session.
+        // The next request is a session establishing FULL request.
+        FetchResponse resp2 = new FetchResponse<>(Errors.FETCH_SESSION_ID_NOT_FOUND,
+            respMap(), 0, INVALID_SESSION_ID);
+        handler.handleResponse(resp2);
+        FetchSessionHandler.Builder builder3 = handler.newBuilder();
+        builder3.add(new TopicPartition("foo", 0),
+            new FetchRequest.PartitionData(0, 100, 200, Optional.empty()));
+        FetchSessionHandler.FetchRequestData data3 = builder3.build();
+        assertTrue(data3.metadata().isFull());
+        assertEquals(INVALID_SESSION_ID, data3.metadata().sessionId());
+        assertEquals(INITIAL_EPOCH, data3.metadata().epoch());
+        assertMapsEqual(reqMap(new ReqEntry("foo", 0, 0, 100, 200)),
+            data3.sessionPartitions(), data3.toSend());
+    }
+
+    @Test
+    public void testVerifyFullFetchResponsePartitions() throws Exception {
+        FetchSessionHandler handler = new FetchSessionHandler(LOG_CONTEXT, 1);
+        String issue = handler.verifyFullFetchResponsePartitions(new FetchResponse<>(Errors.NONE,
+            respMap(new RespEntry("foo", 0, 10, 20),
+                new RespEntry("foo", 1, 10, 20),
+                new RespEntry("bar", 0, 10, 20)),
+            0, INVALID_SESSION_ID));
+        assertTrue(issue.contains("extra"));
+        assertFalse(issue.contains("omitted"));
+        FetchSessionHandler.Builder builder = handler.newBuilder();
+        builder.add(new TopicPartition("foo", 0),
+            new FetchRequest.PartitionData(0, 100, 200, Optional.empty()));
+        builder.add(new TopicPartition("foo", 1),
+            new FetchRequest.PartitionData(10, 110, 210, Optional.empty()));
+        builder.add(new TopicPartition("bar", 0),
+            new FetchRequest.PartitionData(20, 120, 220, Optional.empty()));
+        builder.build();
+        String issue2 = handler.verifyFullFetchResponsePartitions(new FetchResponse<>(Errors.NONE,
+            respMap(new RespEntry("foo", 0, 10, 20),
+                new RespEntry("foo", 1, 10, 20),
+                new RespEntry("bar", 0, 10, 20)),
+            0, INVALID_SESSION_ID));
+        assertTrue(issue2 == null);
+        String issue3 = handler.verifyFullFetchResponsePartitions(new FetchResponse<>(Errors.NONE,
+            respMap(new RespEntry("foo", 0, 10, 20),
+                new RespEntry("foo", 1, 10, 20)),
+            0, INVALID_SESSION_ID));
+        assertFalse(issue3.contains("extra"));
+        assertTrue(issue3.contains("omitted"));
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/clients/InFlightRequestsTest.java b/clients/src/test/java/org/apache/kafka/clients/InFlightRequestsTest.java
new file mode 100644
index 0000000000000..c7b9eb903c8d4
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/InFlightRequestsTest.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;
+
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.requests.RequestHeader;
+import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.test.TestUtils;
+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;
+
+public class InFlightRequestsTest {
+
+    private InFlightRequests inFlightRequests;
+    private int correlationId;
+    private String dest = "dest";
+
+    @Before
+    public void setup() {
+        inFlightRequests = new InFlightRequests(12);
+        correlationId = 0;
+    }
+
+    @Test
+    public void testCompleteLastSent() {
+        int correlationId1 = addRequest(dest);
+        int correlationId2 = addRequest(dest);
+        assertEquals(2, inFlightRequests.count());
+
+        assertEquals(correlationId2, inFlightRequests.completeLastSent(dest).header.correlationId());
+        assertEquals(1, inFlightRequests.count());
+
+        assertEquals(correlationId1, inFlightRequests.completeLastSent(dest).header.correlationId());
+        assertEquals(0, inFlightRequests.count());
+    }
+
+    @Test
+    public void testClearAll() {
+        int correlationId1 = addRequest(dest);
+        int correlationId2 = addRequest(dest);
+
+        List clearedRequests = TestUtils.toList(this.inFlightRequests.clearAll(dest));
+        assertEquals(0, inFlightRequests.count());
+        assertEquals(2, clearedRequests.size());
+        assertEquals(correlationId1, clearedRequests.get(0).header.correlationId());
+        assertEquals(correlationId2, clearedRequests.get(1).header.correlationId());
+    }
+
+    @Test
+    public void testTimedOutNodes() {
+        Time time = new MockTime();
+
+        addRequest("A", time.milliseconds(), 50);
+        addRequest("B", time.milliseconds(), 200);
+        addRequest("B", time.milliseconds(), 100);
+
+        time.sleep(50);
+        assertEquals(Collections.emptyList(), inFlightRequests.nodesWithTimedOutRequests(time.milliseconds()));
+
+        time.sleep(25);
+        assertEquals(Collections.singletonList("A"), inFlightRequests.nodesWithTimedOutRequests(time.milliseconds()));
+
+        time.sleep(50);
+        assertEquals(Arrays.asList("A", "B"), inFlightRequests.nodesWithTimedOutRequests(time.milliseconds()));
+    }
+
+    @Test
+    public void testCompleteNext() {
+        int correlationId1 = addRequest(dest);
+        int correlationId2 = addRequest(dest);
+        assertEquals(2, inFlightRequests.count());
+
+        assertEquals(correlationId1, inFlightRequests.completeNext(dest).header.correlationId());
+        assertEquals(1, inFlightRequests.count());
+
+        assertEquals(correlationId2, inFlightRequests.completeNext(dest).header.correlationId());
+        assertEquals(0, inFlightRequests.count());
+    }
+
+    @Test(expected = IllegalStateException.class)
+    public void testCompleteNextThrowsIfNoInflights() {
+        inFlightRequests.completeNext(dest);
+    }
+
+    @Test(expected = IllegalStateException.class)
+    public void testCompleteLastSentThrowsIfNoInFlights() {
+        inFlightRequests.completeLastSent(dest);
+    }
+
+    private int addRequest(String destination) {
+        return addRequest(destination, 0, 10000);
+    }
+
+    private int addRequest(String destination, long sendTimeMs, int requestTimeoutMs) {
+        int correlationId = this.correlationId;
+        this.correlationId += 1;
+
+        RequestHeader requestHeader = new RequestHeader(ApiKeys.METADATA, (short) 0, "clientId", correlationId);
+        NetworkClient.InFlightRequest ifr = new NetworkClient.InFlightRequest(requestHeader, requestTimeoutMs, 0,
+                destination, null, false, false, null, null, sendTimeMs);
+        inFlightRequests.add(ifr);
+        return correlationId;
+    }
+
+}
diff --git a/clients/src/test/java/org/apache/kafka/clients/MetadataCacheTest.java b/clients/src/test/java/org/apache/kafka/clients/MetadataCacheTest.java
new file mode 100644
index 0000000000000..afdf89c6238c9
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/MetadataCacheTest.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 org.apache.kafka.clients;
+
+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.protocol.Errors;
+import org.apache.kafka.common.requests.MetadataResponse;
+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 java.util.function.Function;
+import java.util.stream.Collectors;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+
+public class MetadataCacheTest {
+
+    @Test
+    public void testMissingLeaderEndpoint() {
+        // Although the broker attempts to ensure leader information is available, the
+        // client metadata cache may retain partition metadata across multiple responses.
+        // For example, separate responses may contain conflicting leader epochs for
+        // separate partitions and the client will always retain the highest.
+
+        TopicPartition topicPartition = new TopicPartition("topic", 0);
+
+        MetadataResponse.PartitionMetadata partitionMetadata = new MetadataResponse.PartitionMetadata(
+                Errors.NONE,
+                topicPartition,
+                Optional.of(5),
+                Optional.of(10),
+                Arrays.asList(5, 6, 7),
+                Arrays.asList(5, 6, 7),
+                Collections.emptyList());
+
+        Map nodesById = new HashMap<>();
+        nodesById.put(6, new Node(6, "localhost", 2077));
+        nodesById.put(7, new Node(7, "localhost", 2078));
+        nodesById.put(8, new Node(8, "localhost", 2079));
+
+        MetadataCache cache = new MetadataCache("clusterId",
+                nodesById,
+                Collections.singleton(partitionMetadata),
+                Collections.emptySet(),
+                Collections.emptySet(),
+                Collections.emptySet(),
+                null);
+
+        Cluster cluster = cache.cluster();
+        assertNull(cluster.leaderFor(topicPartition));
+
+        PartitionInfo partitionInfo = cluster.partition(topicPartition);
+        Map replicas = Arrays.stream(partitionInfo.replicas())
+                .collect(Collectors.toMap(Node::id, Function.identity()));
+        assertNull(partitionInfo.leader());
+        assertEquals(3, replicas.size());
+        assertTrue(replicas.get(5).isEmpty());
+        assertEquals(nodesById.get(6), replicas.get(6));
+        assertEquals(nodesById.get(7), replicas.get(7));
+    }
+
+}
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 3f2a94c7ee952..713cb5afd7737 100644
--- a/clients/src/test/java/org/apache/kafka/clients/MetadataTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/MetadataTest.java
@@ -16,71 +16,70 @@
  */
 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.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.MessageUtil;
+import org.apache.kafka.common.requests.MetadataRequest;
+import org.apache.kafka.common.requests.MetadataResponse;
+import org.apache.kafka.common.requests.RequestTestUtils;
+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.MockClusterResourceListener;
+import org.junit.Test;
+
 import java.net.InetSocketAddress;
+import java.nio.ByteBuffer;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicReference;
 
-import org.apache.kafka.common.Cluster;
-import org.apache.kafka.common.internals.ClusterResourceListeners;
-import org.apache.kafka.common.Node;
-import org.apache.kafka.common.PartitionInfo;
-import org.apache.kafka.common.errors.TimeoutException;
-import org.apache.kafka.test.MockClusterResourceListener;
-import org.apache.kafka.test.TestUtils;
-import org.junit.After;
-import org.junit.Test;
-
-import static org.junit.Assert.assertArrayEquals;
+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;
 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 MetadataTest {
 
     private long refreshBackoffMs = 100;
     private long metadataExpireMs = 1000;
-    private Metadata metadata = new Metadata(refreshBackoffMs, metadataExpireMs, true);
-    private AtomicReference backgroundError = new AtomicReference<>();
+    private Metadata metadata = new Metadata(refreshBackoffMs, metadataExpireMs, new LogContext(),
+            new ClusterResourceListeners());
 
-    @After
-    public void tearDown() {
-        assertNull("Exception in background thread : " + backgroundError.get(), backgroundError.get());
+    private static MetadataResponse emptyMetadataResponse() {
+        return RequestTestUtils.metadataResponse(
+                Collections.emptyList(),
+                null,
+                -1,
+                Collections.emptyList());
     }
 
-    @Test
-    public void testMetadata() throws Exception {
-        long time = 0;
-        metadata.update(Cluster.empty(), Collections.emptySet(), time);
-        assertFalse("No update needed.", metadata.timeToNextUpdate(time) == 0);
-        metadata.requestUpdate();
-        assertFalse("Still no updated needed due to backoff", metadata.timeToNextUpdate(time) == 0);
-        time += refreshBackoffMs;
-        assertTrue("Update needed now that backoff time expired", metadata.timeToNextUpdate(time) == 0);
-        String topic = "my-topic";
-        Thread t1 = asyncFetch(topic, 500);
-        Thread t2 = asyncFetch(topic, 500);
-        assertTrue("Awaiting update", t1.isAlive());
-        assertTrue("Awaiting update", t2.isAlive());
-        // Perform metadata update when an update is requested on the async fetch thread
-        // This simulates the metadata update sequence in KafkaProducer
-        while (t1.isAlive() || t2.isAlive()) {
-            if (metadata.timeToNextUpdate(time) == 0) {
-                metadata.update(TestUtils.singletonCluster(topic, 1), Collections.emptySet(), time);
-                time += refreshBackoffMs;
-            }
-            Thread.sleep(1);
-        }
-        t1.join();
-        t2.join();
-        assertFalse("No update needed.", metadata.timeToNextUpdate(time) == 0);
-        time += metadataExpireMs;
-        assertTrue("Update needed due to stale metadata.", metadata.timeToNextUpdate(time) == 0);
+    @Test(expected = IllegalStateException.class)
+    public void testMetadataUpdateAfterClose() {
+        metadata.close();
+        metadata.updateWithCurrentRequestVersion(emptyMetadataResponse(), false, 1000);
     }
 
     private static void checkTimeToNextUpdate(long refreshBackoffMs, long metadataExpireMs) {
@@ -96,12 +95,13 @@ private static void checkTimeToNextUpdate(long refreshBackoffMs, long metadataEx
         }
 
         long largerOfBackoffAndExpire = Math.max(refreshBackoffMs, metadataExpireMs);
-        Metadata metadata = new Metadata(refreshBackoffMs, metadataExpireMs, true);
+        Metadata metadata = new Metadata(refreshBackoffMs, metadataExpireMs, new LogContext(),
+                new ClusterResourceListeners());
 
         assertEquals(0, metadata.timeToNextUpdate(now));
 
         // lastSuccessfulRefreshMs updated to now.
-        metadata.update(Cluster.empty(), Collections.emptySet(), now);
+        metadata.updateWithCurrentRequestVersion(emptyMetadataResponse(), false, now);
 
         // The last update was successful so the remaining time to expire the current metadata should be returned.
         assertEquals(largerOfBackoffAndExpire, metadata.timeToNextUpdate(now));
@@ -112,7 +112,7 @@ private static void checkTimeToNextUpdate(long refreshBackoffMs, long metadataEx
         assertEquals(refreshBackoffMs, metadata.timeToNextUpdate(now));
 
         // Reset needUpdate to false.
-        metadata.update(Cluster.empty(), Collections.emptySet(), now);
+        metadata.updateWithCurrentRequestVersion(emptyMetadataResponse(), false, now);
         assertEquals(largerOfBackoffAndExpire, metadata.timeToNextUpdate(now));
 
         // Both metadataExpireMs and refreshBackoffMs elapsed.
@@ -121,6 +121,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);
@@ -131,11 +143,11 @@ public void testTimeToNextUpdate() {
     }
 
     @Test
-    public void testTimeToNextUpdate_RetryBackoff() {
+    public void testTimeToNextUpdateRetryBackoff() {
         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));
@@ -151,127 +163,142 @@ public void testTimeToNextUpdate_RetryBackoff() {
         assertEquals(0, metadata.timeToNextUpdate(now + 1));
     }
 
-    @Test
-    public void testTimeToNextUpdate_OverwriteBackoff() {
-        long now = 10000;
-
-        // New topic added to fetch set and update requested. It should allow immediate update.
-        metadata.update(Cluster.empty(), Collections.emptySet(), now);
-        metadata.add("new-topic");
-        assertEquals(0, metadata.timeToNextUpdate(now));
-
-        // Even though setTopics called, immediate update isn't necessary if the new topic set isn't
-        // containing a new topic,
-        metadata.update(Cluster.empty(), Collections.emptySet(), now);
-        metadata.setTopics(metadata.topics());
-        assertEquals(metadataExpireMs, metadata.timeToNextUpdate(now));
-
-        // If the new set of topics containing a new topic then it should allow immediate update.
-        metadata.setTopics(Collections.singletonList("another-new-topic"));
-        assertEquals(0, metadata.timeToNextUpdate(now));
-
-        // If metadata requested for all topics it should allow immediate update.
-        metadata.update(Cluster.empty(), Collections.emptySet(), now);
-        metadata.needMetadataForAllTopics(true);
-        assertEquals(0, metadata.timeToNextUpdate(now));
-
-        // However if metadata is already capable to serve all topics it shouldn't override backoff.
-        metadata.update(Cluster.empty(), Collections.emptySet(), now);
-        metadata.needMetadataForAllTopics(true);
-        assertEquals(metadataExpireMs, metadata.timeToNextUpdate(now));
-    }
-
     /**
-     * Tests that {@link org.apache.kafka.clients.Metadata#awaitUpdate(int, long)} doesn't
-     * wait forever with a max timeout value of 0
-     *
-     * @throws Exception
-     * @see https://issues.apache.org/jira/browse/KAFKA-1836
+     * 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 testMetadataUpdateWaitTime() throws Exception {
-        long time = 0;
-        metadata.update(Cluster.empty(), Collections.emptySet(), time);
-        assertFalse("No update needed.", metadata.timeToNextUpdate(time) == 0);
-        // first try with a max wait time of 0 and ensure that this returns back without waiting forever
-        try {
-            metadata.awaitUpdate(metadata.requestUpdate(), 0);
-            fail("Wait on metadata update was expected to timeout, but it didn't");
-        } catch (TimeoutException te) {
-            // expected
+    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++) {
+            ByteBuffer buffer = MessageUtil.toByteBuffer(data, version);
+            MetadataResponse response = MetadataResponse.parse(buffer, version);
+            assertFalse(response.hasReliableLeaderEpochs());
+            metadata.updateWithCurrentRequestVersion(response, false, 100);
+            assertTrue(metadata.partitionMetadataIfCurrent(tp).isPresent());
+            MetadataResponse.PartitionMetadata responseMetadata = this.metadata.partitionMetadataIfCurrent(tp).get();
+            assertEquals(Optional.empty(), responseMetadata.leaderEpoch);
         }
-        // now try with a higher timeout value once
-        final long twoSecondWait = 2000;
-        try {
-            metadata.awaitUpdate(metadata.requestUpdate(), twoSecondWait);
-            fail("Wait on metadata update was expected to timeout, but it didn't");
-        } catch (TimeoutException te) {
-            // expected
+
+        for (short version = 9; version <= ApiKeys.METADATA.latestVersion(); version++) {
+            ByteBuffer buffer = MessageUtil.toByteBuffer(data, version);
+            MetadataResponse response = MetadataResponse.parse(buffer, version);
+            assertTrue(response.hasReliableLeaderEpochs());
+            metadata.updateWithCurrentRequestVersion(response, false, 100);
+            assertTrue(metadata.partitionMetadataIfCurrent(tp).isPresent());
+            MetadataResponse.PartitionMetadata responseMetadata = metadata.partitionMetadataIfCurrent(tp).get();
+            assertEquals(Optional.of(10), responseMetadata.leaderEpoch);
         }
     }
 
+    @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.updateWithCurrentRequestVersion(new MetadataResponse(data, ApiKeys.METADATA.latestVersion()), false, 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.updateWithCurrentRequestVersion(new MetadataResponse(data, ApiKeys.METADATA.latestVersion()), false, 101);
+        assertEquals(Optional.of(10), metadata.lastSeenLeaderEpoch(tp));
+
+        assertTrue(metadata.partitionMetadataIfCurrent(tp).isPresent());
+        MetadataResponse.PartitionMetadata responseMetadata = this.metadata.partitionMetadataIfCurrent(tp).get();
+
+        assertEquals(Arrays.asList(1, 2, 3), responseMetadata.inSyncReplicaIds);
+        assertEquals(Optional.of(10), responseMetadata.leaderEpoch);
+    }
+
     @Test
     public void testFailedUpdate() {
         long time = 100;
-        metadata.update(Cluster.empty(), Collections.emptySet(), time);
+        metadata.updateWithCurrentRequestVersion(emptyMetadataResponse(), false, time);
 
         assertEquals(100, metadata.timeToNextUpdate(1000));
-        metadata.failedUpdate(1100, null);
+        metadata.failedUpdate(1100);
 
         assertEquals(100, metadata.timeToNextUpdate(1100));
         assertEquals(100, metadata.lastSuccessfulUpdate());
 
-        metadata.needMetadataForAllTopics(true);
-        metadata.update(Cluster.empty(), Collections.emptySet(), time);
+        metadata.updateWithCurrentRequestVersion(emptyMetadataResponse(), false, time);
         assertEquals(100, metadata.timeToNextUpdate(1000));
     }
 
-    @Test
-    public void testUpdateWithNeedMetadataForAllTopics() {
-        long time = 0;
-        metadata.update(Cluster.empty(), Collections.emptySet(), time);
-        metadata.needMetadataForAllTopics(true);
-
-        final List expectedTopics = Collections.singletonList("topic");
-        metadata.setTopics(expectedTopics);
-        metadata.update(new Cluster(null,
-                Collections.singletonList(new Node(0, "host1", 1000)),
-                Arrays.asList(
-                    new PartitionInfo("topic", 0, null, null, null),
-                    new PartitionInfo("topic1", 0, null, null, null)),
-                Collections.emptySet(),
-                Collections.emptySet()),
-            Collections.emptySet(), 100);
-
-        assertArrayEquals("Metadata got updated with wrong set of topics.",
-            expectedTopics.toArray(), metadata.topics().toArray());
-
-        metadata.needMetadataForAllTopics(false);
-    }
-
     @Test
     public void testClusterListenerGetsNotifiedOfUpdate() {
-        long time = 0;
         MockClusterResourceListener mockClusterListener = new MockClusterResourceListener();
         ClusterResourceListeners listeners = new ClusterResourceListeners();
         listeners.maybeAdd(mockClusterListener);
-        metadata = new Metadata(refreshBackoffMs, metadataExpireMs, true, false, listeners);
+        metadata = new Metadata(refreshBackoffMs, metadataExpireMs, new LogContext(), listeners);
 
         String hostName = "www.example.com";
-        Cluster cluster = Cluster.bootstrap(Arrays.asList(new InetSocketAddress(hostName, 9002)));
-        metadata.update(cluster, Collections.emptySet(), 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());
 
-        metadata.update(new Cluster(
-                        "dummy",
-                        Arrays.asList(new Node(0, "host1", 1000)),
-                        Arrays.asList(
-                                new PartitionInfo("topic", 0, null, null, null),
-                                new PartitionInfo("topic1", 0, null, null, null)),
-                        Collections.emptySet(),
-                        Collections.emptySet()),
-                Collections.emptySet(), 100);
+        Map partitionCounts = new HashMap<>();
+        partitionCounts.put("topic", 1);
+        partitionCounts.put("topic1", 1);
+        MetadataResponse metadataResponse = RequestTestUtils.metadataUpdateWith("dummy", 1, partitionCounts);
+        metadata.updateWithCurrentRequestVersion(metadataResponse, false, 100);
 
         assertEquals("MockClusterResourceListener did not get cluster metadata correctly",
                 "dummy", mockClusterListener.clusterResource().clusterId());
@@ -279,149 +306,586 @@ public void testClusterListenerGetsNotifiedOfUpdate() {
                 MockClusterResourceListener.IS_ON_UPDATE_CALLED.get());
     }
 
+
     @Test
-    public void testListenerGetsNotifiedOfUpdate() {
-        long time = 0;
-        final Set topics = new HashSet<>();
-        metadata.update(Cluster.empty(), Collections.emptySet(), time);
-        metadata.addListener(new Metadata.Listener() {
-            @Override
-            public void onMetadataUpdate(Cluster cluster, Set unavailableTopics) {
-                topics.clear();
-                topics.addAll(cluster.topics());
+    public void testRequestUpdate() {
+        assertFalse(metadata.updateRequested());
+
+        int[] epochs =           {42,   42,    41,    41,    42,    43,   43,    42,    41,    44};
+        boolean[] updateResult = {true, false, false, false, false, true, false, false, false, true};
+        TopicPartition tp = new TopicPartition("topic", 0);
+
+        MetadataResponse metadataResponse = RequestTestUtils.metadataUpdateWith("dummy", 1,
+                Collections.emptyMap(), Collections.singletonMap("topic", 1), _tp -> 0);
+        metadata.updateWithCurrentRequestVersion(metadataResponse, false, 10L);
+
+        for (int i = 0; i < epochs.length; i++) {
+            metadata.updateLastSeenEpochIfNewer(tp, epochs[i]);
+            if (updateResult[i]) {
+                assertTrue("Expected metadata update to be requested [" + i + "]", metadata.updateRequested());
+            } else {
+                assertFalse("Did not expect metadata update to be requested [" + i + "]", metadata.updateRequested());
             }
-        });
-
-        metadata.update(new Cluster(
-                null,
-                Arrays.asList(new Node(0, "host1", 1000)),
-                Arrays.asList(
-                    new PartitionInfo("topic", 0, null, null, null),
-                    new PartitionInfo("topic1", 0, null, null, null)),
-                Collections.emptySet(),
-                Collections.emptySet()),
-            Collections.emptySet(), 100);
-
-        assertEquals("Listener did not update topics list correctly",
-            new HashSet<>(Arrays.asList("topic", "topic1")), topics);
+            metadata.updateWithCurrentRequestVersion(emptyMetadataResponse(), false, 0L);
+            assertFalse(metadata.updateRequested());
+        }
     }
 
     @Test
-    public void testListenerCanUnregister() {
-        long time = 0;
-        final Set topics = new HashSet<>();
-        metadata.update(Cluster.empty(), Collections.emptySet(), time);
-        final Metadata.Listener listener = new Metadata.Listener() {
-            @Override
-            public void onMetadataUpdate(Cluster cluster, Set unavailableTopics) {
-                topics.clear();
-                topics.addAll(cluster.topics());
-            }
-        };
-        metadata.addListener(listener);
-
-        metadata.update(new Cluster(
-                "cluster",
-                Collections.singletonList(new Node(0, "host1", 1000)),
-                Arrays.asList(
-                    new PartitionInfo("topic", 0, null, null, null),
-                    new PartitionInfo("topic1", 0, null, null, null)),
-                Collections.emptySet(),
-                Collections.emptySet()),
-            Collections.emptySet(), 100);
-
-        metadata.removeListener(listener);
-
-        metadata.update(new Cluster(
-                "cluster",
-                Arrays.asList(new Node(0, "host1", 1000)),
-                Arrays.asList(
-                    new PartitionInfo("topic2", 0, null, null, null),
-                    new PartitionInfo("topic3", 0, null, null, null)),
-                Collections.emptySet(),
-                Collections.emptySet()),
-            Collections.emptySet(), 100);
-
-        assertEquals("Listener did not update topics list correctly",
-            new HashSet<>(Arrays.asList("topic", "topic1")), topics);
+    public void testUpdateLastEpoch() {
+        TopicPartition tp = new TopicPartition("topic-1", 0);
+
+        MetadataResponse metadataResponse = emptyMetadataResponse();
+        metadata.updateWithCurrentRequestVersion(metadataResponse, false, 0L);
+
+        // if we have no leader epoch, this call shouldn't do anything
+        assertFalse(metadata.updateLastSeenEpochIfNewer(tp, 0));
+        assertFalse(metadata.updateLastSeenEpochIfNewer(tp, 1));
+        assertFalse(metadata.updateLastSeenEpochIfNewer(tp, 2));
+        assertFalse(metadata.lastSeenLeaderEpoch(tp).isPresent());
+
+        // Metadata with newer epoch is handled
+        metadataResponse = RequestTestUtils.metadataUpdateWith("dummy", 1, Collections.emptyMap(), Collections.singletonMap("topic-1", 1), _tp -> 10);
+        metadata.updateWithCurrentRequestVersion(metadataResponse, false, 1L);
+        assertOptional(metadata.lastSeenLeaderEpoch(tp), leaderAndEpoch -> assertEquals(leaderAndEpoch.intValue(), 10));
+
+        // Don't update to an older one
+        assertFalse(metadata.updateLastSeenEpochIfNewer(tp, 1));
+        assertOptional(metadata.lastSeenLeaderEpoch(tp), leaderAndEpoch -> assertEquals(leaderAndEpoch.intValue(), 10));
+
+        // Don't cause update if it's the same one
+        assertFalse(metadata.updateLastSeenEpochIfNewer(tp, 10));
+        assertOptional(metadata.lastSeenLeaderEpoch(tp), leaderAndEpoch -> assertEquals(leaderAndEpoch.intValue(), 10));
+
+        // Update if we see newer epoch
+        assertTrue(metadata.updateLastSeenEpochIfNewer(tp, 12));
+        assertOptional(metadata.lastSeenLeaderEpoch(tp), leaderAndEpoch -> assertEquals(leaderAndEpoch.intValue(), 12));
+
+        metadataResponse = RequestTestUtils.metadataUpdateWith("dummy", 1, Collections.emptyMap(), Collections.singletonMap("topic-1", 1), _tp -> 12);
+        metadata.updateWithCurrentRequestVersion(metadataResponse, false, 2L);
+        assertOptional(metadata.lastSeenLeaderEpoch(tp), leaderAndEpoch -> assertEquals(leaderAndEpoch.intValue(), 12));
+
+        // Don't overwrite metadata with older epoch
+        metadataResponse = RequestTestUtils.metadataUpdateWith("dummy", 1, Collections.emptyMap(), Collections.singletonMap("topic-1", 1), _tp -> 11);
+        metadata.updateWithCurrentRequestVersion(metadataResponse, false, 3L);
+        assertOptional(metadata.lastSeenLeaderEpoch(tp), leaderAndEpoch -> assertEquals(leaderAndEpoch.intValue(), 12));
     }
 
     @Test
-    public void testTopicExpiry() throws Exception {
-        metadata = new Metadata(refreshBackoffMs, metadataExpireMs, true, true, new ClusterResourceListeners());
-
-        // Test that topic is expired if not used within the expiry interval
-        long time = 0;
-        metadata.add("topic1");
-        metadata.update(Cluster.empty(), Collections.emptySet(), time);
-        time += Metadata.TOPIC_EXPIRY_MS;
-        metadata.update(Cluster.empty(), Collections.emptySet(), time);
-        assertFalse("Unused topic not expired", metadata.containsTopic("topic1"));
-
-        // Test that topic is not expired if used within the expiry interval
-        metadata.add("topic2");
-        metadata.update(Cluster.empty(), Collections.emptySet(), time);
-        for (int i = 0; i < 3; i++) {
-            time += Metadata.TOPIC_EXPIRY_MS / 2;
-            metadata.update(Cluster.empty(), Collections.emptySet(), time);
-            assertTrue("Topic expired even though in use", metadata.containsTopic("topic2"));
-            metadata.add("topic2");
+    public void testRejectOldMetadata() {
+        Map partitionCounts = new HashMap<>();
+        partitionCounts.put("topic-1", 1);
+        TopicPartition tp = new TopicPartition("topic-1", 0);
+
+        metadata.updateWithCurrentRequestVersion(emptyMetadataResponse(), false, 0L);
+
+        // First epoch seen, accept it
+        {
+            MetadataResponse metadataResponse = RequestTestUtils.metadataUpdateWith("dummy", 1, Collections.emptyMap(), partitionCounts, _tp -> 100);
+            metadata.updateWithCurrentRequestVersion(metadataResponse, false, 10L);
+            assertNotNull(metadata.fetch().partition(tp));
+            assertTrue(metadata.lastSeenLeaderEpoch(tp).isPresent());
+            assertEquals(metadata.lastSeenLeaderEpoch(tp).get().longValue(), 100);
         }
 
-        // Test that topics added using setTopics expire
-        HashSet topics = new HashSet<>();
-        topics.add("topic4");
-        metadata.setTopics(topics);
-        metadata.update(Cluster.empty(), Collections.emptySet(), time);
-        time += Metadata.TOPIC_EXPIRY_MS;
-        metadata.update(Cluster.empty(), Collections.emptySet(), time);
-        assertFalse("Unused topic not expired", metadata.containsTopic("topic4"));
+        // Fake an empty ISR, but with an older epoch, should reject it
+        {
+            MetadataResponse metadataResponse = RequestTestUtils.metadataUpdateWith("dummy", 1, Collections.emptyMap(), partitionCounts, _tp -> 99,
+                (error, partition, leader, leaderEpoch, replicas, isr, offlineReplicas) ->
+                        new MetadataResponse.PartitionMetadata(error, partition, leader,
+                            leaderEpoch, replicas, Collections.emptyList(), offlineReplicas), ApiKeys.METADATA.latestVersion());
+            metadata.updateWithCurrentRequestVersion(metadataResponse, false, 20L);
+            assertEquals(metadata.fetch().partition(tp).inSyncReplicas().length, 1);
+            assertEquals(metadata.lastSeenLeaderEpoch(tp).get().longValue(), 100);
+        }
+
+        // Fake an empty ISR, with same epoch, accept it
+        {
+            MetadataResponse metadataResponse = RequestTestUtils.metadataUpdateWith("dummy", 1, Collections.emptyMap(), partitionCounts, _tp -> 100,
+                (error, partition, leader, leaderEpoch, replicas, isr, offlineReplicas) ->
+                        new MetadataResponse.PartitionMetadata(error, partition, leader,
+                            leaderEpoch, replicas, Collections.emptyList(), offlineReplicas), ApiKeys.METADATA.latestVersion());
+            metadata.updateWithCurrentRequestVersion(metadataResponse, false, 20L);
+            assertEquals(metadata.fetch().partition(tp).inSyncReplicas().length, 0);
+            assertEquals(metadata.lastSeenLeaderEpoch(tp).get().longValue(), 100);
+        }
+
+        // Empty metadata response, should not keep old partition but should keep the last-seen epoch
+        {
+            MetadataResponse metadataResponse = RequestTestUtils.metadataUpdateWith("dummy", 1, Collections.emptyMap(), Collections.emptyMap());
+            metadata.updateWithCurrentRequestVersion(metadataResponse, false, 20L);
+            assertNull(metadata.fetch().partition(tp));
+            assertEquals(metadata.lastSeenLeaderEpoch(tp).get().longValue(), 100);
+        }
+
+        // Back in the metadata, with old epoch, should not get added
+        {
+            MetadataResponse metadataResponse = RequestTestUtils.metadataUpdateWith("dummy", 1, Collections.emptyMap(), partitionCounts, _tp -> 99);
+            metadata.updateWithCurrentRequestVersion(metadataResponse, false, 10L);
+            assertNull(metadata.fetch().partition(tp));
+            assertEquals(metadata.lastSeenLeaderEpoch(tp).get().longValue(), 100);
+        }
     }
 
     @Test
-    public void testNonExpiringMetadata() throws Exception {
-        metadata = new Metadata(refreshBackoffMs, metadataExpireMs, true, false, new ClusterResourceListeners());
-
-        // Test that topic is not expired if not used within the expiry interval
-        long time = 0;
-        metadata.add("topic1");
-        metadata.update(Cluster.empty(), Collections.emptySet(), time);
-        time += Metadata.TOPIC_EXPIRY_MS;
-        metadata.update(Cluster.empty(), Collections.emptySet(), time);
-        assertTrue("Unused topic expired when expiry disabled", metadata.containsTopic("topic1"));
-
-        // Test that topic is not expired if used within the expiry interval
-        metadata.add("topic2");
-        metadata.update(Cluster.empty(), Collections.emptySet(), time);
-        for (int i = 0; i < 3; i++) {
-            time += Metadata.TOPIC_EXPIRY_MS / 2;
-            metadata.update(Cluster.empty(), Collections.emptySet(), time);
-            assertTrue("Topic expired even though in use", metadata.containsTopic("topic2"));
-            metadata.add("topic2");
-        }
+    public void testOutOfBandEpochUpdate() {
+        Map partitionCounts = new HashMap<>();
+        partitionCounts.put("topic-1", 5);
+        TopicPartition tp = new TopicPartition("topic-1", 0);
+
+        metadata.updateWithCurrentRequestVersion(emptyMetadataResponse(), false, 0L);
+
+        assertFalse(metadata.updateLastSeenEpochIfNewer(tp, 99));
+
+        // Update epoch to 100
+        MetadataResponse metadataResponse = RequestTestUtils.metadataUpdateWith("dummy", 1, Collections.emptyMap(), partitionCounts, _tp -> 100);
+        metadata.updateWithCurrentRequestVersion(metadataResponse, false, 10L);
+        assertNotNull(metadata.fetch().partition(tp));
+        assertTrue(metadata.lastSeenLeaderEpoch(tp).isPresent());
+        assertEquals(metadata.lastSeenLeaderEpoch(tp).get().longValue(), 100);
+
+        // Simulate a leader epoch from another response, like a fetch response or list offsets
+        assertTrue(metadata.updateLastSeenEpochIfNewer(tp, 101));
+
+        // Cache of partition stays, but current partition info is not available since it's stale
+        assertNotNull(metadata.fetch().partition(tp));
+        assertEquals(Objects.requireNonNull(metadata.fetch().partitionCountForTopic("topic-1")).longValue(), 5);
+        assertFalse(metadata.partitionMetadataIfCurrent(tp).isPresent());
+        assertEquals(metadata.lastSeenLeaderEpoch(tp).get().longValue(), 101);
+
+        // Metadata with older epoch is rejected, metadata state is unchanged
+        metadata.updateWithCurrentRequestVersion(metadataResponse, false, 20L);
+        assertNotNull(metadata.fetch().partition(tp));
+        assertEquals(Objects.requireNonNull(metadata.fetch().partitionCountForTopic("topic-1")).longValue(), 5);
+        assertFalse(metadata.partitionMetadataIfCurrent(tp).isPresent());
+        assertEquals(metadata.lastSeenLeaderEpoch(tp).get().longValue(), 101);
+
+        // Metadata with equal or newer epoch is accepted
+        metadataResponse = RequestTestUtils.metadataUpdateWith("dummy", 1, Collections.emptyMap(), partitionCounts, _tp -> 101);
+        metadata.updateWithCurrentRequestVersion(metadataResponse, false, 30L);
+        assertNotNull(metadata.fetch().partition(tp));
+        assertEquals(Objects.requireNonNull(metadata.fetch().partitionCountForTopic("topic-1")).longValue(), 5);
+        assertTrue(metadata.partitionMetadataIfCurrent(tp).isPresent());
+        assertEquals(metadata.lastSeenLeaderEpoch(tp).get().longValue(), 101);
+    }
 
-        // Test that topics added using setTopics don't expire
-        HashSet topics = new HashSet<>();
-        topics.add("topic4");
-        metadata.setTopics(topics);
-        time += metadataExpireMs * 2;
-        metadata.update(Cluster.empty(), Collections.emptySet(), time);
-        assertTrue("Unused topic expired when expiry disabled", metadata.containsTopic("topic4"));
+    @Test
+    public void testNoEpoch() {
+        metadata.updateWithCurrentRequestVersion(emptyMetadataResponse(), false, 0L);
+        MetadataResponse metadataResponse = RequestTestUtils.metadataUpdateWith("dummy", 1, Collections.emptyMap(), Collections.singletonMap("topic-1", 1));
+        metadata.updateWithCurrentRequestVersion(metadataResponse, false, 10L);
+
+        TopicPartition tp = new TopicPartition("topic-1", 0);
+
+        // no epoch
+        assertFalse(metadata.lastSeenLeaderEpoch(tp).isPresent());
+
+        // still works
+        assertTrue(metadata.partitionMetadataIfCurrent(tp).isPresent());
+        assertEquals(0, metadata.partitionMetadataIfCurrent(tp).get().partition());
+        assertEquals(Optional.of(0), metadata.partitionMetadataIfCurrent(tp).get().leaderId);
+
+        // Since epoch was null, this shouldn't update it
+        metadata.updateLastSeenEpochIfNewer(tp, 10);
+        assertTrue(metadata.partitionMetadataIfCurrent(tp).isPresent());
+        assertFalse(metadata.partitionMetadataIfCurrent(tp).get().leaderEpoch.isPresent());
+    }
+
+    @Test
+    public void testClusterCopy() {
+        Map counts = new HashMap<>();
+        Map errors = new HashMap<>();
+        counts.put("topic1", 2);
+        counts.put("topic2", 3);
+        counts.put(Topic.GROUP_METADATA_TOPIC_NAME, 3);
+        errors.put("topic3", Errors.INVALID_TOPIC_EXCEPTION);
+        errors.put("topic4", Errors.TOPIC_AUTHORIZATION_FAILED);
+
+        MetadataResponse metadataResponse = RequestTestUtils.metadataUpdateWith("dummy", 4, errors, counts);
+        metadata.updateWithCurrentRequestVersion(metadataResponse, false, 0L);
+
+        Cluster cluster = metadata.fetch();
+        assertEquals(cluster.clusterResource().clusterId(), "dummy");
+        assertEquals(cluster.nodes().size(), 4);
+
+        // topic counts
+        assertEquals(cluster.invalidTopics(), Collections.singleton("topic3"));
+        assertEquals(cluster.unauthorizedTopics(), Collections.singleton("topic4"));
+        assertEquals(cluster.topics().size(), 3);
+        assertEquals(cluster.internalTopics(), Collections.singleton(Topic.GROUP_METADATA_TOPIC_NAME));
+
+        // partition counts
+        assertEquals(cluster.partitionsForTopic("topic1").size(), 2);
+        assertEquals(cluster.partitionsForTopic("topic2").size(), 3);
+
+        // Sentinel instances
+        InetSocketAddress address = InetSocketAddress.createUnresolved("localhost", 0);
+        Cluster fromMetadata = MetadataCache.bootstrap(Collections.singletonList(address)).cluster();
+        Cluster fromCluster = Cluster.bootstrap(Collections.singletonList(address));
+        assertEquals(fromMetadata, fromCluster);
+
+        Cluster fromMetadataEmpty = MetadataCache.empty().cluster();
+        Cluster fromClusterEmpty = Cluster.empty();
+        assertEquals(fromMetadataEmpty, fromClusterEmpty);
     }
 
-    private Thread asyncFetch(final String topic, final long maxWaitMs) {
-        Thread thread = new Thread() {
-            public void run() {
-                while (metadata.fetch().partitionsForTopic(topic).isEmpty()) {
-                    try {
-                        metadata.awaitUpdate(metadata.requestUpdate(), maxWaitMs);
-                    } catch (Exception e) {
-                        backgroundError.set(e.toString());
-                    }
+    @Test
+    public void testRequestVersion() {
+        Time time = new MockTime();
+
+        metadata.requestUpdate();
+        Metadata.MetadataRequestAndVersion versionAndBuilder = metadata.newMetadataRequestAndVersion(time.milliseconds());
+        metadata.update(versionAndBuilder.requestVersion,
+                RequestTestUtils.metadataUpdateWith(1, Collections.singletonMap("topic", 1)), false, time.milliseconds());
+        assertFalse(metadata.updateRequested());
+
+        // bump the request version for new topics added to the metadata
+        metadata.requestUpdateForNewTopics();
+
+        // simulating a bump while a metadata request is in flight
+        versionAndBuilder = metadata.newMetadataRequestAndVersion(time.milliseconds());
+        metadata.requestUpdateForNewTopics();
+        metadata.update(versionAndBuilder.requestVersion,
+                RequestTestUtils.metadataUpdateWith(1, Collections.singletonMap("topic", 1)), true, time.milliseconds());
+
+        // metadata update is still needed
+        assertTrue(metadata.updateRequested());
+
+        // the next update will resolve it
+        versionAndBuilder = metadata.newMetadataRequestAndVersion(time.milliseconds());
+        metadata.update(versionAndBuilder.requestVersion,
+                RequestTestUtils.metadataUpdateWith(1, Collections.singletonMap("topic", 1)), true, time.milliseconds());
+        assertFalse(metadata.updateRequested());
+    }
+
+    @Test
+    public void testPartialMetadataUpdate() {
+        Time time = new MockTime();
+
+        metadata = new Metadata(refreshBackoffMs, metadataExpireMs, new LogContext(), new ClusterResourceListeners()) {
+                @Override
+                protected MetadataRequest.Builder newMetadataRequestBuilderForNewTopics() {
+                    return newMetadataRequestBuilder();
                 }
-            }
-        };
-        thread.start();
-        return thread;
+            };
+
+        assertFalse(metadata.updateRequested());
+
+        // Request a metadata update. This must force a full metadata update request.
+        metadata.requestUpdate();
+        Metadata.MetadataRequestAndVersion versionAndBuilder = metadata.newMetadataRequestAndVersion(time.milliseconds());
+        assertFalse(versionAndBuilder.isPartialUpdate);
+        metadata.update(versionAndBuilder.requestVersion,
+                RequestTestUtils.metadataUpdateWith(1, Collections.singletonMap("topic", 1)), false, time.milliseconds());
+        assertFalse(metadata.updateRequested());
+
+        // Request a metadata update for a new topic. This should perform a partial metadata update.
+        metadata.requestUpdateForNewTopics();
+        versionAndBuilder = metadata.newMetadataRequestAndVersion(time.milliseconds());
+        assertTrue(versionAndBuilder.isPartialUpdate);
+        metadata.update(versionAndBuilder.requestVersion,
+                RequestTestUtils.metadataUpdateWith(1, Collections.singletonMap("topic", 1)), true, time.milliseconds());
+        assertFalse(metadata.updateRequested());
+
+        // Request both types of metadata updates. This should always perform a full update.
+        metadata.requestUpdate();
+        metadata.requestUpdateForNewTopics();
+        versionAndBuilder = metadata.newMetadataRequestAndVersion(time.milliseconds());
+        assertFalse(versionAndBuilder.isPartialUpdate);
+        metadata.update(versionAndBuilder.requestVersion,
+                RequestTestUtils.metadataUpdateWith(1, Collections.singletonMap("topic", 1)), false, time.milliseconds());
+        assertFalse(metadata.updateRequested());
+
+        // Request only a partial metadata update, but elapse enough time such that a full refresh is needed.
+        metadata.requestUpdateForNewTopics();
+        final long refreshTimeMs = time.milliseconds() + metadata.metadataExpireMs();
+        versionAndBuilder = metadata.newMetadataRequestAndVersion(refreshTimeMs);
+        assertFalse(versionAndBuilder.isPartialUpdate);
+        metadata.update(versionAndBuilder.requestVersion,
+                RequestTestUtils.metadataUpdateWith(1, Collections.singletonMap("topic", 1)), true, refreshTimeMs);
+        assertFalse(metadata.updateRequested());
+
+        // Request two partial metadata updates that are overlapping.
+        metadata.requestUpdateForNewTopics();
+        versionAndBuilder = metadata.newMetadataRequestAndVersion(time.milliseconds());
+        assertTrue(versionAndBuilder.isPartialUpdate);
+        metadata.requestUpdateForNewTopics();
+        Metadata.MetadataRequestAndVersion overlappingVersionAndBuilder = metadata.newMetadataRequestAndVersion(time.milliseconds());
+        assertTrue(overlappingVersionAndBuilder.isPartialUpdate);
+        assertTrue(metadata.updateRequested());
+        metadata.update(versionAndBuilder.requestVersion,
+                RequestTestUtils.metadataUpdateWith(1, Collections.singletonMap("topic-1", 1)), true, time.milliseconds());
+        assertTrue(metadata.updateRequested());
+        metadata.update(overlappingVersionAndBuilder.requestVersion,
+                RequestTestUtils.metadataUpdateWith(1, Collections.singletonMap("topic-2", 1)), true, time.milliseconds());
+        assertFalse(metadata.updateRequested());
+    }
+
+    @Test
+    public void testInvalidTopicError() {
+        Time time = new MockTime();
+
+        String invalidTopic = "topic dfsa";
+        MetadataResponse invalidTopicResponse = RequestTestUtils.metadataUpdateWith("clusterId", 1,
+                Collections.singletonMap(invalidTopic, Errors.INVALID_TOPIC_EXCEPTION), Collections.emptyMap());
+        metadata.updateWithCurrentRequestVersion(invalidTopicResponse, false, time.milliseconds());
+
+        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
+        metadata.maybeThrowAnyException();
+
+        // Reset the invalid topic error
+        metadata.updateWithCurrentRequestVersion(invalidTopicResponse, false, 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.updateWithCurrentRequestVersion(emptyMetadataResponse(), false, time.milliseconds());
+        metadata.maybeThrowAnyException();
+    }
+
+    @Test
+    public void testTopicAuthorizationError() {
+        Time time = new MockTime();
+
+        String invalidTopic = "foo";
+        MetadataResponse unauthorizedTopicResponse = RequestTestUtils.metadataUpdateWith("clusterId", 1,
+                Collections.singletonMap(invalidTopic, Errors.TOPIC_AUTHORIZATION_FAILED), Collections.emptyMap());
+        metadata.updateWithCurrentRequestVersion(unauthorizedTopicResponse, false, time.milliseconds());
+
+        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
+        metadata.maybeThrowAnyException();
+
+        // Reset the unauthorized topic error
+        metadata.updateWithCurrentRequestVersion(unauthorizedTopicResponse, false, 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.updateWithCurrentRequestVersion(emptyMetadataResponse(), false, time.milliseconds());
+        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 = RequestTestUtils.metadataUpdateWith("clusterId", 1, topicErrors, Collections.emptyMap());
+
+        metadata.updateWithCurrentRequestVersion(metadataResponse, false, 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.updateWithCurrentRequestVersion(metadataResponse, false, time.milliseconds());
+        TopicAuthorizationException e2 = assertThrows(TopicAuthorizationException.class,
+            () -> metadata.maybeThrowExceptionForTopic("sensitiveTopic2"));
+        assertEquals(Collections.singleton("sensitiveTopic2"), e2.unauthorizedTopics());
+        metadata.maybeThrowAnyException();
+
+        metadata.updateWithCurrentRequestVersion(metadataResponse, false, 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.updateWithCurrentRequestVersion(metadataResponse, false, 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 = RequestTestUtils.metadataUpdateWith("dummy", 2, Collections.emptyMap(), partitionCounts, _tp -> 99,
+            (error, partition, leader, leaderEpoch, replicas, isr, offlineReplicas) ->
+                new MetadataResponse.PartitionMetadata(error, partition, Optional.of(node0.id()), leaderEpoch,
+                    Collections.singletonList(node0.id()), Collections.emptyList(),
+                        Collections.singletonList(node1.id())), ApiKeys.METADATA.latestVersion());
+        metadata.updateWithCurrentRequestVersion(emptyMetadataResponse(), false, 0L);
+        metadata.updateWithCurrentRequestVersion(metadataResponse, false, 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.updateWithCurrentRequestVersion(new MetadataResponse(new MetadataResponseData()
+                        .setTopics(buildTopicCollection(tp.topic(), firstPartitionMetadata))
+                        .setBrokers(buildBrokerCollection(Arrays.asList(node0, node1, node2))),
+                        ApiKeys.METADATA.latestVersion()),
+                false, 10L);
+
+        metadata.updateWithCurrentRequestVersion(new MetadataResponse(new MetadataResponseData()
+                        .setTopics(buildTopicCollection(tp.topic(), secondPartitionMetadata))
+                        .setBrokers(buildBrokerCollection(Arrays.asList(node1, node2))),
+                        ApiKeys.METADATA.latestVersion()),
+                false, 20L);
+
+        assertNull(metadata.fetch().leaderFor(tp));
+        assertEquals(Optional.of(10), metadata.lastSeenLeaderEpoch(tp));
+        assertFalse(metadata.currentLeader(tp).leader.isPresent());
+    }
+
+    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) {
+            MetadataResponseData.MetadataResponseBroker broker = new MetadataResponseData.MetadataResponseBroker()
+                    .setNodeId(node.id())
+                    .setHost(node.host())
+                    .setPort(node.port())
+                    .setRack(node.rack());
+            brokers.add(broker);
+        }
+        return brokers;
+    }
+
+    @Test
+    public void testMetadataMerge() {
+        Time time = new MockTime();
+
+        final AtomicReference> retainTopics = new AtomicReference<>(new HashSet<>());
+        metadata = new Metadata(refreshBackoffMs, metadataExpireMs, new LogContext(), new ClusterResourceListeners()) {
+                @Override
+                protected boolean retainTopic(String topic, boolean isInternal, long nowMs) {
+                    return retainTopics.get().contains(topic);
+                }
+            };
+
+        // Initialize a metadata instance with two topic variants "old" and "keep". Both will be retained.
+        String oldClusterId = "oldClusterId";
+        int oldNodes = 2;
+        Map oldTopicErrors = new HashMap<>();
+        oldTopicErrors.put("oldInvalidTopic", Errors.INVALID_TOPIC_EXCEPTION);
+        oldTopicErrors.put("keepInvalidTopic", Errors.INVALID_TOPIC_EXCEPTION);
+        oldTopicErrors.put("oldUnauthorizedTopic", Errors.TOPIC_AUTHORIZATION_FAILED);
+        oldTopicErrors.put("keepUnauthorizedTopic", Errors.TOPIC_AUTHORIZATION_FAILED);
+        Map oldTopicPartitionCounts = new HashMap<>();
+        oldTopicPartitionCounts.put("oldValidTopic", 2);
+        oldTopicPartitionCounts.put("keepValidTopic", 3);
+
+        retainTopics.set(new HashSet<>(Arrays.asList(
+            "oldInvalidTopic",
+            "keepInvalidTopic",
+            "oldUnauthorizedTopic",
+            "keepUnauthorizedTopic",
+            "oldValidTopic",
+            "keepValidTopic")));
+
+        MetadataResponse metadataResponse =
+                RequestTestUtils.metadataUpdateWith(oldClusterId, oldNodes, oldTopicErrors, oldTopicPartitionCounts, _tp -> 100);
+        metadata.updateWithCurrentRequestVersion(metadataResponse, true, time.milliseconds());
+
+        // Update the metadata to add a new topic variant, "new", which will be retained with "keep". Note this
+        // means that all of the "old" topics should be dropped.
+        Cluster cluster = metadata.fetch();
+        assertEquals(cluster.clusterResource().clusterId(), oldClusterId);
+        assertEquals(cluster.nodes().size(), oldNodes);
+        assertEquals(cluster.invalidTopics(), new HashSet<>(Arrays.asList("oldInvalidTopic", "keepInvalidTopic")));
+        assertEquals(cluster.unauthorizedTopics(), new HashSet<>(Arrays.asList("oldUnauthorizedTopic", "keepUnauthorizedTopic")));
+        assertEquals(cluster.topics(), new HashSet<>(Arrays.asList("oldValidTopic", "keepValidTopic")));
+        assertEquals(cluster.partitionsForTopic("oldValidTopic").size(), 2);
+        assertEquals(cluster.partitionsForTopic("keepValidTopic").size(), 3);
+
+        String newClusterId = "newClusterId";
+        int newNodes = oldNodes + 1;
+        Map newTopicErrors = new HashMap<>();
+        newTopicErrors.put("newInvalidTopic", Errors.INVALID_TOPIC_EXCEPTION);
+        newTopicErrors.put("newUnauthorizedTopic", Errors.TOPIC_AUTHORIZATION_FAILED);
+        Map newTopicPartitionCounts = new HashMap<>();
+        newTopicPartitionCounts.put("keepValidTopic", 2);
+        newTopicPartitionCounts.put("newValidTopic", 4);
+
+        retainTopics.set(new HashSet<>(Arrays.asList(
+            "keepInvalidTopic",
+            "newInvalidTopic",
+            "keepUnauthorizedTopic",
+            "newUnauthorizedTopic",
+            "keepValidTopic",
+            "newValidTopic")));
+
+        metadataResponse = RequestTestUtils.metadataUpdateWith(newClusterId, newNodes, newTopicErrors, newTopicPartitionCounts, _tp -> 200);
+        metadata.updateWithCurrentRequestVersion(metadataResponse, true, time.milliseconds());
+
+        cluster = metadata.fetch();
+        assertEquals(cluster.clusterResource().clusterId(), newClusterId);
+        assertEquals(cluster.nodes().size(), newNodes);
+        assertEquals(cluster.invalidTopics(), new HashSet<>(Arrays.asList("keepInvalidTopic", "newInvalidTopic")));
+        assertEquals(cluster.unauthorizedTopics(), new HashSet<>(Arrays.asList("keepUnauthorizedTopic", "newUnauthorizedTopic")));
+        assertEquals(cluster.topics(), new HashSet<>(Arrays.asList("keepValidTopic", "newValidTopic")));
+        assertEquals(cluster.partitionsForTopic("keepValidTopic").size(), 2);
+        assertEquals(cluster.partitionsForTopic("newValidTopic").size(), 4);
+
+        // Perform another metadata update, but this time all topic metadata should be cleared.
+        retainTopics.set(Collections.emptySet());
+
+        metadataResponse = RequestTestUtils.metadataUpdateWith(newClusterId, newNodes, newTopicErrors, newTopicPartitionCounts, _tp -> 300);
+        metadata.updateWithCurrentRequestVersion(metadataResponse, true, time.milliseconds());
+
+        cluster = metadata.fetch();
+        assertEquals(cluster.clusterResource().clusterId(), newClusterId);
+        assertEquals(cluster.nodes().size(), newNodes);
+        assertEquals(cluster.invalidTopics(), Collections.emptySet());
+        assertEquals(cluster.unauthorizedTopics(), Collections.emptySet());
+        assertEquals(cluster.topics(), Collections.emptySet());
     }
 }
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 8b334729247dc..214ecc522270b 100644
--- a/clients/src/test/java/org/apache/kafka/clients/MockClient.java
+++ b/clients/src/test/java/org/apache/kafka/clients/MockClient.java
@@ -16,19 +16,20 @@
  */
 package org.apache.kafka.clients;
 
-import org.apache.kafka.common.Cluster;
 import org.apache.kafka.common.Node;
 import org.apache.kafka.common.errors.AuthenticationException;
+import org.apache.kafka.common.errors.InterruptException;
 import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.protocol.Errors;
 import org.apache.kafka.common.requests.AbstractRequest;
 import org.apache.kafka.common.requests.AbstractResponse;
+import org.apache.kafka.common.requests.MetadataRequest;
+import org.apache.kafka.common.requests.MetadataResponse;
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.test.TestCondition;
 import org.apache.kafka.test.TestUtils;
 
-import java.util.ArrayDeque;
 import java.util.ArrayList;
-import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Iterator;
@@ -37,17 +38,13 @@
 import java.util.Queue;
 import java.util.Set;
 import java.util.concurrent.ConcurrentLinkedDeque;
+import java.util.stream.Collectors;
 
 /**
  * A mock network client for use testing code
  */
 public class MockClient implements KafkaClient {
-    public static final RequestMatcher ALWAYS_TRUE = new RequestMatcher() {
-        @Override
-        public boolean matches(AbstractRequest body) {
-            return true;
-        }
-    };
+    public static final RequestMatcher ALWAYS_TRUE = body -> true;
 
     private static class FutureResponse {
         private final Node node;
@@ -70,74 +67,101 @@ public FutureResponse(Node node,
 
     }
 
+    private int correlation;
+    private Runnable wakeupHook;
     private final Time time;
-    private final Metadata metadata;
-    private Set unavailableTopics;
-    private Cluster cluster;
-    private Node node = null;
-    private final Set ready = new HashSet<>();
-    private final Map blackedOut = new HashMap<>();
+    private final MockMetadataUpdater metadataUpdater;
+    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
     private final Queue requests = new ConcurrentLinkedDeque<>();
     // Use concurrent queue for responses so that responses may be updated during poll() from a different thread.
     private final Queue responses = new ConcurrentLinkedDeque<>();
-    private final Queue futureResponses = new ArrayDeque<>();
-    private final Queue metadataUpdates = new ArrayDeque<>();
+    private final Queue futureResponses = new ConcurrentLinkedDeque<>();
+    private final Queue metadataUpdates = new ConcurrentLinkedDeque<>();
     private volatile NodeApiVersions nodeApiVersions = NodeApiVersions.create();
+    private volatile int numBlockingWakeups = 0;
+    private volatile boolean active = true;
 
-    public MockClient(Time time) {
-        this(time, null);
+    public MockClient(Time time, Metadata metadata) {
+        this(time, new DefaultMockMetadataUpdater(metadata));
     }
 
-    public MockClient(Time time, Metadata metadata) {
+    public MockClient(Time time, MockMetadataUpdater metadataUpdater) {
         this.time = time;
-        this.metadata = metadata;
-        this.unavailableTopics = Collections.emptySet();
+        this.metadataUpdater = metadataUpdater;
+    }
+
+    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 (isBlackedOut(node))
-            return false;
-        ready.add(node.idString());
-        return true;
+        return connectionState(node.idString()).ready(now);
     }
 
     @Override
     public long connectionDelay(Node node, long now) {
-        return 0;
+        return connectionState(node.idString()).connectionDelay(now);
     }
 
-    public void blackout(Node node, long duration) {
-        blackedOut.put(node, time.milliseconds() + duration);
+    @Override
+    public long pollDelayMs(Node node, long now) {
+        return connectionDelay(node, now);
     }
 
-    private boolean isBlackedOut(Node node) {
-        if (blackedOut.containsKey(node)) {
-            long expiration = blackedOut.get(node);
-            if (time.milliseconds() > expiration) {
-                blackedOut.remove(node);
-                return false;
-            } else {
-                return true;
-            }
-        }
-        return false;
+    public void backoff(Node node, long durationMs) {
+        connectionState(node.idString()).backoff(time.milliseconds() + durationMs);
+    }
+
+    public void setUnreachable(Node node, long durationMs) {
+        disconnect(node.idString());
+        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) {
+        connectionState(node.idString()).setReadyDelayed(time.milliseconds() + durationMs);
+    }
+
+    public void authenticationFailed(Node node, long backoffMs) {
+        pendingAuthenticationErrors.remove(node);
+        authenticationErrors.put(node, (AuthenticationException) Errors.SASL_AUTHENTICATION_FAILED.exception());
+        disconnect(node.idString());
+        backoff(node, backoffMs);
+    }
+
+    public void createPendingAuthenticationError(Node node, long backoffMs) {
+        pendingAuthenticationErrors.put(node, backoffMs);
     }
 
     @Override
     public boolean connectionFailed(Node node) {
-        return isBlackedOut(node);
+        return connectionState(node.idString()).isBackingOff(time.milliseconds());
     }
 
     @Override
     public AuthenticationException authenticationException(Node node) {
-        return null;
+        return authenticationErrors.get(node);
     }
 
     @Override
@@ -149,15 +173,39 @@ public void disconnect(String node) {
             if (request.destination().equals(node)) {
                 short version = request.requestBuilder().latestAllowedVersion();
                 responses.add(new ClientResponse(request.makeHeader(version), request.callback(), request.destination(),
-                        request.createdTimeMs(), now, true, null, null));
+                        request.createdTimeMs(), now, true, null, null, null));
                 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(); ) {
+            Map.Entry entry = authErrorIter.next();
+            Node node = entry.getKey();
+            long backoffMs = entry.getValue();
+            if (node.idString().equals(request.destination())) {
+                authErrorIter.remove();
+                // Set up a disconnected ClientResponse and create an authentication error
+                // for the affected node.
+                authenticationFailed(node, backoffMs);
+                AbstractRequest.Builder builder = request.requestBuilder();
+                short version = nodeApiVersions.latestUsableVersion(request.apiKey(), builder.oldestAllowedVersion(),
+                    builder.latestAllowedVersion());
+                ClientResponse resp = new ClientResponse(request.makeHeader(version), request.callback(), request.destination(),
+                    request.createdTimeMs(), time.milliseconds(), true, null,
+                        new AuthenticationException("Authentication failed"), null);
+                responses.add(resp);
+                return;
+            }
+        }
         Iterator iterator = futureResponses.iterator();
         while (iterator.hasNext()) {
             FutureResponse futureResp = iterator.next();
@@ -167,18 +215,20 @@ public void send(ClientRequest request, long now) {
             AbstractRequest.Builder builder = request.requestBuilder();
             short version = nodeApiVersions.latestUsableVersion(request.apiKey(), builder.oldestAllowedVersion(),
                     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);
 
             UnsupportedVersionException unsupportedVersionException = null;
-            if (futureResp.isUnsupportedRequest)
-                unsupportedVersionException = new UnsupportedVersionException("Api " +
-                        request.apiKey() + " with version " + version);
-
+            if (futureResp.isUnsupportedRequest) {
+                unsupportedVersionException = new UnsupportedVersionException(
+                        "Api " + request.apiKey() + " with version " + version);
+            } else {
+                AbstractRequest abstractRequest = request.requestBuilder().build(version);
+                if (!futureResp.requestMatcher.matches(abstractRequest))
+                    throw new IllegalStateException("Request matcher did not match next-in-line request "
+                            + abstractRequest + " with prepared response " + futureResp.responseBody);
+            }
             ClientResponse resp = new ClientResponse(request.makeHeader(version), request.callback(), request.destination(),
                     request.createdTimeMs(), time.milliseconds(), futureResp.disconnected,
-                    unsupportedVersionException, futureResp.responseBody);
+                    unsupportedVersionException, null, futureResp.responseBody);
             responses.add(resp);
             iterator.remove();
             return;
@@ -187,34 +237,89 @@ public void send(ClientRequest request, long now) {
         this.requests.add(request);
     }
 
+    /**
+     * Simulate a blocking poll in order to test wakeup behavior.
+     *
+     * @param numBlockingWakeups The number of polls which will block until woken up
+     */
+    public synchronized void enableBlockingUntilWakeup(int numBlockingWakeups) {
+        this.numBlockingWakeups = numBlockingWakeups;
+    }
+
+    @Override
+    public synchronized void wakeup() {
+        if (numBlockingWakeups > 0) {
+            numBlockingWakeups--;
+            notify();
+        }
+        if (wakeupHook != null) {
+            wakeupHook.run();
+        }
+    }
+
+    private synchronized void maybeAwaitWakeup() {
+        try {
+            int remainingBlockingWakeups = numBlockingWakeups;
+            if (remainingBlockingWakeups <= 0)
+                return;
+
+            while (numBlockingWakeups == remainingBlockingWakeups)
+                wait();
+        } catch (InterruptedException e) {
+            throw new InterruptException(e);
+        }
+    }
+
     @Override
     public List poll(long timeoutMs, long now) {
-        List copy = new ArrayList<>(this.responses);
+        maybeAwaitWakeup();
+        checkTimeoutOfPendingRequests(now);
 
-        if (metadata != null && metadata.updateRequested()) {
+        // We skip metadata updates if all nodes are currently blacked out
+        if (metadataUpdater.isUpdateNeeded() && leastLoadedNode(now) != null) {
             MetadataUpdate metadataUpdate = metadataUpdates.poll();
-            if (cluster != null)
-                metadata.update(cluster, this.unavailableTopics, time.milliseconds());
-            if (metadataUpdate == null)
-                metadata.update(metadata.fetch(), this.unavailableTopics, time.milliseconds());
-            else {
-                this.unavailableTopics = metadataUpdate.unavailableTopics;
-                metadata.update(metadataUpdate.cluster, metadataUpdate.unavailableTopics, time.milliseconds());
+            if (metadataUpdate != null) {
+                metadataUpdater.update(time, metadataUpdate);
+            } else {
+                metadataUpdater.updateWithCurrentMetadata(time);
             }
         }
 
+        List copy = new ArrayList<>();
         ClientResponse response;
         while ((response = this.responses.poll()) != null) {
             response.onComplete();
+            copy.add(response);
         }
 
         return copy;
     }
 
+    private long elapsedTimeMs(long currentTimeMs, long startTimeMs) {
+        return Math.max(0, currentTimeMs - startTimeMs);
+    }
+
+    private void checkTimeoutOfPendingRequests(long nowMs) {
+        ClientRequest request = requests.peek();
+        while (request != null && elapsedTimeMs(nowMs, request.createdTimeMs()) > request.requestTimeoutMs()) {
+            disconnect(request.destination());
+            requests.poll();
+            request = requests.peek();
+        }
+    }
+
     public Queue requests() {
         return this.requests;
     }
 
+    public Queue responses() {
+        return this.responses;
+    }
+
+    public Queue futureResponses() {
+        return this.futureResponses;
+    }
+
     public void respond(AbstractResponse response) {
         respond(response, false);
     }
@@ -233,19 +338,20 @@ public void respond(RequestMatcher matcher, AbstractResponse response) {
 
     // Utility method to enable out of order responses
     public void respondToRequest(ClientRequest clientRequest, AbstractResponse response) {
-        AbstractRequest request = clientRequest.requestBuilder().build();
         requests.remove(clientRequest);
         short version = clientRequest.requestBuilder().latestAllowedVersion();
         responses.add(new ClientResponse(clientRequest.makeHeader(version), clientRequest.callback(), clientRequest.destination(),
-                clientRequest.createdTimeMs(), time.milliseconds(), false, null, response));
+                clientRequest.createdTimeMs(), time.milliseconds(), false, null, null, response));
     }
 
 
     public void respond(AbstractResponse response, boolean disconnected) {
-        ClientRequest request = requests.remove();
+        if (requests.isEmpty())
+            throw new IllegalStateException("No requests pending for inbound response " + response);
+        ClientRequest request = requests.poll();
         short version = request.requestBuilder().latestAllowedVersion();
         responses.add(new ClientResponse(request.makeHeader(version), request.callback(), request.destination(),
-                request.createdTimeMs(), time.milliseconds(), disconnected, null, response));
+                request.createdTimeMs(), time.milliseconds(), disconnected, null, null, response));
     }
 
     public void respondFrom(AbstractResponse response, Node node) {
@@ -260,7 +366,7 @@ public void respondFrom(AbstractResponse response, Node node, boolean disconnect
                 iterator.remove();
                 short version = request.requestBuilder().latestAllowedVersion();
                 responses.add(new ClientResponse(request.makeHeader(version), request.callback(), request.destination(),
-                        request.createdTimeMs(), time.milliseconds(), disconnected, null, response));
+                        request.createdTimeMs(), time.milliseconds(), disconnected, null, null, response));
                 return;
             }
         }
@@ -335,24 +441,33 @@ public boolean conditionMet() {
     }
 
     public void reset() {
-        ready.clear();
-        blackedOut.clear();
+        connections.clear();
         requests.clear();
         responses.clear();
         futureResponses.clear();
         metadataUpdates.clear();
+        authenticationErrors.clear();
+    }
+
+    public boolean hasPendingMetadataUpdates() {
+        return !metadataUpdates.isEmpty();
+    }
+
+    public int numAwaitingResponses() {
+        return futureResponses.size();
     }
 
-    public void prepareMetadataUpdate(Cluster cluster, Set unavailableTopics) {
-        metadataUpdates.add(new MetadataUpdate(cluster, unavailableTopics));
+    public void prepareMetadataUpdate(MetadataResponse updateResponse) {
+        prepareMetadataUpdate(updateResponse, false);
     }
 
-    public void setNode(Node node) {
-        this.node = node;
+    public void prepareMetadataUpdate(MetadataResponse updateResponse,
+                                      boolean expectMatchMetadataTopics) {
+        metadataUpdates.add(new MetadataUpdate(updateResponse, expectMatchMetadataTopics));
     }
 
-    public void cluster(Cluster cluster) {
-        this.cluster = cluster;
+    public void updateMetadata(MetadataResponse updateResponse) {
+        metadataUpdater.update(time, new MetadataUpdate(updateResponse, false));
     }
 
     @Override
@@ -365,6 +480,10 @@ public boolean hasInFlightRequests() {
         return !requests.isEmpty();
     }
 
+    public boolean hasPendingResponses() {
+        return !responses.isEmpty() || !futureResponses.isEmpty();
+    }
+
     @Override
     public int inFlightRequestCount(String node) {
         int result = 0;
@@ -381,39 +500,60 @@ public boolean hasInFlightRequests(String node) {
     }
 
     @Override
-    public boolean hasReadyNodes() {
-        return !ready.isEmpty();
+    public boolean hasReadyNodes(long now) {
+        return connections.values().stream().anyMatch(cxn -> cxn.isReady(now));
     }
 
     @Override
     public ClientRequest newClientRequest(String nodeId, AbstractRequest.Builder requestBuilder, long createdTimeMs,
                                           boolean expectResponse) {
-        return newClientRequest(nodeId, requestBuilder, createdTimeMs, expectResponse, null);
+        return newClientRequest(nodeId, requestBuilder, createdTimeMs, expectResponse, 5000, null);
     }
 
     @Override
-    public ClientRequest newClientRequest(String nodeId, AbstractRequest.Builder requestBuilder, long createdTimeMs,
-                                          boolean expectResponse, RequestCompletionHandler callback) {
-        return new ClientRequest(nodeId, requestBuilder, 0, "mockClientId", createdTimeMs,
-                expectResponse, callback);
+    public ClientRequest newClientRequest(String nodeId,
+                                          AbstractRequest.Builder requestBuilder,
+                                          long createdTimeMs,
+                                          boolean expectResponse,
+                                          int requestTimeoutMs,
+                                          RequestCompletionHandler callback) {
+        return new ClientRequest(nodeId, requestBuilder, correlation++, "mockClientId", createdTimeMs,
+                expectResponse, requestTimeoutMs, callback);
+    }
+
+    @Override
+    public void initiateClose() {
+        close();
     }
 
     @Override
-    public void wakeup() {
+    public boolean active() {
+        return active;
     }
 
     @Override
     public void close() {
+        active = false;
+        metadataUpdater.close();
     }
 
     @Override
     public void close(String node) {
-        ready.remove(node);
+        connections.remove(node);
     }
 
     @Override
     public Node leastLoadedNode(long now) {
-        return this.node;
+        // Consistent with NetworkClient, we do not return nodes awaiting reconnect backoff
+        for (Node node : metadataUpdater.fetchNodes()) {
+            if (!connectionState(node.idString()).isBackingOff(now))
+                return node;
+        }
+        return null;
+    }
+
+    public void setWakeupHook(Runnable wakeupHook) {
+        this.wakeupHook = wakeupHook;
     }
 
     /**
@@ -422,6 +562,7 @@ public Node leastLoadedNode(long now) {
      * to inspect the request body for the type of the request or for specific fields that should be set,
      * and to fail the test if it doesn't match.
      */
+    @FunctionalInterface
     public interface RequestMatcher {
         boolean matches(AbstractRequest body);
     }
@@ -430,12 +571,181 @@ public void setNodeApiVersions(NodeApiVersions nodeApiVersions) {
         this.nodeApiVersions = nodeApiVersions;
     }
 
-    private static class MetadataUpdate {
-        final Cluster cluster;
-        final Set unavailableTopics;
-        MetadataUpdate(Cluster cluster, Set unavailableTopics) {
-            this.cluster = cluster;
-            this.unavailableTopics = unavailableTopics;
+    public static class MetadataUpdate {
+        final MetadataResponse updateResponse;
+        final boolean expectMatchRefreshTopics;
+
+        MetadataUpdate(MetadataResponse updateResponse, boolean expectMatchRefreshTopics) {
+            this.updateResponse = updateResponse;
+            this.expectMatchRefreshTopics = expectMatchRefreshTopics;
+        }
+
+        private Set topics() {
+            return updateResponse.topicMetadata().stream()
+                    .map(MetadataResponse.TopicMetadata::topic)
+                    .collect(Collectors.toSet());
         }
     }
+
+    /**
+     * 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)}
+     * and bookkeeping through {@link Metadata}. The extensibility allows AdminClient, which does
+     * not rely on {@link Metadata} to do its own thing.
+     */
+    public interface MockMetadataUpdater {
+        List fetchNodes();
+
+        boolean isUpdateNeeded();
+
+        void update(Time time, MetadataUpdate update);
+
+        default void updateWithCurrentMetadata(Time time) {}
+
+        default void close() {}
+    }
+
+    private static class DefaultMockMetadataUpdater implements MockMetadataUpdater {
+        private final Metadata metadata;
+        private MetadataUpdate lastUpdate;
+
+        public DefaultMockMetadataUpdater(Metadata metadata) {
+            this.metadata = metadata;
+        }
+
+        @Override
+        public List fetchNodes() {
+            return metadata.fetch().nodes();
+        }
+
+        @Override
+        public boolean isUpdateNeeded() {
+            return metadata.updateRequested();
+        }
+
+        @Override
+        public void updateWithCurrentMetadata(Time time) {
+            if (lastUpdate == null)
+                throw new IllegalStateException("No previous metadata update to use");
+            update(time, lastUpdate);
+        }
+
+        private void maybeCheckExpectedTopics(MetadataUpdate update, MetadataRequest.Builder builder) {
+            if (update.expectMatchRefreshTopics) {
+                if (builder.isAllTopics())
+                    throw new IllegalStateException("The metadata topics does not match expectation. "
+                            + "Expected topics: " + update.topics()
+                            + ", asked topics: ALL");
+
+                Set requestedTopics = new HashSet<>(builder.topics());
+                if (!requestedTopics.equals(update.topics())) {
+                    throw new IllegalStateException("The metadata topics does not match expectation. "
+                            + "Expected topics: " + update.topics()
+                            + ", asked topics: " + requestedTopics);
+                }
+            }
+        }
+
+        @Override
+        public void update(Time time, MetadataUpdate update) {
+            MetadataRequest.Builder builder = metadata.newMetadataRequestBuilder();
+            maybeCheckExpectedTopics(update, builder);
+            metadata.updateWithCurrentRequestVersion(update.updateResponse, false, time.milliseconds());
+            this.lastUpdate = update;
+        }
+
+        @Override
+        public void close() {
+            metadata.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 edbd72db88868..a6801ceb69f27 100644
--- a/clients/src/test/java/org/apache/kafka/clients/NetworkClientTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/NetworkClientTest.java
@@ -17,16 +17,28 @@
 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.message.ProduceRequestData;
+import org.apache.kafka.common.message.ProduceResponseData;
 import org.apache.kafka.common.network.NetworkReceive;
 import org.apache.kafka.common.protocol.ApiKeys;
-import org.apache.kafka.common.protocol.types.Struct;
-import org.apache.kafka.common.record.MemoryRecords;
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.requests.AbstractResponse;
 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.ResponseHeader;
+import org.apache.kafka.common.requests.ProduceResponse;
+import org.apache.kafka.common.requests.RequestHeader;
+import org.apache.kafka.common.requests.RequestTestUtils;
+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;
@@ -36,57 +48,80 @@
 import org.junit.Test;
 
 import java.nio.ByteBuffer;
-import java.util.Arrays;
+import java.util.ArrayList;
 import java.util.Collections;
+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.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;
 
 public class NetworkClientTest {
 
-    protected final int requestTimeoutMs = 1000;
+    protected final int defaultRequestTimeoutMs = 1000;
     protected final MockTime time = new MockTime();
     protected final MockSelector selector = new MockSelector(time);
-    protected final Metadata metadata = new Metadata(0, Long.MAX_VALUE, true);
-    protected final int nodeId = 1;
-    protected final Cluster cluster = TestUtils.singletonCluster("test", nodeId);
-    protected final Node node = cluster.nodes().get(0);
+    protected final Node node = TestUtils.singletonCluster().nodes().iterator().next();
     protected final long reconnectBackoffMsTest = 10 * 1000;
     protected final long reconnectBackoffMaxMsTest = 10 * 10000;
-
+    protected final long connectionSetupTimeoutMsTest = 5 * 1000;
+    protected final long connectionSetupTimeoutMaxMsTest = 127 * 1000;
+    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, metadata, "mock", Integer.MAX_VALUE,
+        return new NetworkClient(selector, metadataUpdater, "mock", Integer.MAX_VALUE,
                 reconnectBackoffMsTest, reconnectBackoffMaxMs, 64 * 1024, 64 * 1024,
-                requestTimeoutMs, time, true, new ApiVersions(), new LogContext());
+                defaultRequestTimeoutMs, connectionSetupTimeoutMsTest, connectionSetupTimeoutMaxMsTest, ClientDnsLookup.DEFAULT, time, true, new ApiVersions(), new LogContext());
+    }
+
+    private NetworkClient createNetworkClientWithMultipleNodes(long reconnectBackoffMaxMs, long connectionSetupTimeoutMsTest, int nodeNumber) {
+        List nodes = TestUtils.clusterWith(nodeNumber).nodes();
+        TestMetadataUpdater metadataUpdater = new TestMetadataUpdater(nodes);
+        return new NetworkClient(selector, metadataUpdater, "mock", Integer.MAX_VALUE,
+                reconnectBackoffMsTest, reconnectBackoffMaxMs, 64 * 1024, 64 * 1024,
+                defaultRequestTimeoutMs, connectionSetupTimeoutMsTest, connectionSetupTimeoutMaxMsTest, ClientDnsLookup.DEFAULT, time, true, new ApiVersions(), new LogContext());
     }
 
     private NetworkClient createNetworkClientWithStaticNodes() {
-        return new NetworkClient(selector, new ManualMetadataUpdater(Arrays.asList(node)),
-                "mock-static", Integer.MAX_VALUE, 0, 0, 64 * 1024, 64 * 1024, requestTimeoutMs,
-                time, true, new ApiVersions(), new LogContext());
+        return new NetworkClient(selector, metadataUpdater,
+                "mock-static", Integer.MAX_VALUE, 0, 0, 64 * 1024, 64 * 1024, defaultRequestTimeoutMs,
+                connectionSetupTimeoutMsTest, connectionSetupTimeoutMaxMsTest, ClientDnsLookup.DEFAULT, time, true, new ApiVersions(), new LogContext());
     }
 
-    private NetworkClient createNetworkClientWithNoVersionDiscovery() {
+    private NetworkClient createNetworkClientWithNoVersionDiscovery(Metadata metadata) {
         return new NetworkClient(selector, metadata, "mock", Integer.MAX_VALUE,
+                reconnectBackoffMsTest, 0, 64 * 1024, 64 * 1024,
+                defaultRequestTimeoutMs, connectionSetupTimeoutMsTest, connectionSetupTimeoutMaxMsTest, ClientDnsLookup.DEFAULT, time, false, new ApiVersions(), new LogContext());
+    }
+
+    private NetworkClient createNetworkClientWithNoVersionDiscovery() {
+        return new NetworkClient(selector, metadataUpdater, "mock", Integer.MAX_VALUE,
                 reconnectBackoffMsTest, reconnectBackoffMaxMsTest,
-                64 * 1024, 64 * 1024, requestTimeoutMs, time, false, new ApiVersions(), new LogContext());
+                64 * 1024, 64 * 1024, defaultRequestTimeoutMs,
+                connectionSetupTimeoutMsTest, connectionSetupTimeoutMaxMsTest, ClientDnsLookup.DEFAULT, time, false, new ApiVersions(), new LogContext());
     }
 
     @Before
     public void setup() {
-        metadata.update(cluster, Collections.emptySet(), time.milliseconds());
+        selector.reset();
     }
 
     @Test(expected = IllegalStateException.class)
     public void testSendToUnreadyNode() {
-        MetadataRequest.Builder builder = new MetadataRequest.Builder(Arrays.asList("test"), true);
+        MetadataRequest.Builder builder = new MetadataRequest.Builder(Collections.singletonList("test"), true);
         long now = time.milliseconds();
         ClientRequest request = client.newClientRequest("5", builder, now, false);
         client.send(request, now);
@@ -108,6 +143,12 @@ public void testSimpleRequestResponseWithNoBrokerDiscovery() {
         checkSimpleRequestResponse(clientWithNoVersionDiscovery);
     }
 
+    @Test
+    public void testDnsLookupFailure() {
+        /* Fail cleanly when the node has a bad hostname */
+        assertFalse(client.ready(new Node(1234, "badhost", 1234), time.milliseconds()));
+    }
+
     @Test
     public void testClose() {
         client.ready(node, time.milliseconds());
@@ -115,8 +156,10 @@ public void testClose() {
         client.poll(1, time.milliseconds());
         assertTrue("The client should be ready", client.isReady(node, time.milliseconds()));
 
-        ProduceRequest.Builder builder = ProduceRequest.Builder.forCurrentMagic((short) 1, 1000,
-                Collections.emptyMap());
+        ProduceRequest.Builder builder = ProduceRequest.forCurrentMagic(new ProduceRequestData()
+                .setTopicData(new ProduceRequestData.TopicProduceDataCollection())
+                .setAcks((short) 1)
+                .setTimeoutMs(1000));
         ClientRequest request = client.newClientRequest(node.idString(), builder, time.milliseconds(), true);
         client.send(request, time.milliseconds());
         assertEquals("There should be 1 in-flight request after send", 1,
@@ -131,25 +174,33 @@ public void testClose() {
         assertFalse("Connection should not be ready after close", client.isReady(node, 0));
     }
 
+    @Test
+    public void testUnsupportedVersionDuringInternalMetadataRequest() {
+        List topics = Collections.singletonList("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());
+        short requestVersion = PRODUCE.latestVersion();
+        ProduceRequest.Builder builder = new ProduceRequest.Builder(
+                requestVersion,
+                requestVersion,
+                new ProduceRequestData()
+                    .setAcks((short) 1)
+                    .setTimeoutMs(1000));
         TestCallbackHandler handler = new TestCallbackHandler();
-        ClientRequest request = networkClient.newClientRequest(
-                node.idString(), builder, time.milliseconds(), true, handler);
+        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()));
-        resp.set("responses", new Object[0]);
-        Struct responseHeaderStruct = respHeader.toStruct();
-        int size = responseHeaderStruct.sizeOf() + resp.sizeOf();
-        ByteBuffer buffer = ByteBuffer.allocate(size);
-        responseHeaderStruct.writeTo(buffer);
-        resp.writeTo(buffer);
-        buffer.flip();
+        ProduceResponse produceResponse = new ProduceResponse(new ProduceResponseData());
+        ByteBuffer buffer = RequestTestUtils.serializeResponseWithHeader(produceResponse, requestVersion, request.correlationId());
         selector.completeReceive(new NetworkReceive(node.idString(), buffer));
         List responses = networkClient.poll(1, time.milliseconds());
         assertEquals(1, responses.size());
@@ -159,61 +210,500 @@ private void checkSimpleRequestResponse(NetworkClient networkClient) {
                 request.correlationId(), handler.response.requestHeader().correlationId());
     }
 
-    private void maybeSetExpectedApiVersionsResponse() {
-        ApiVersionsResponse response = ApiVersionsResponse.defaultApiVersionsResponse();
-        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 = RequestTestUtils.serializeResponseWithHeader(response, 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()) {
-            maybeSetExpectedApiVersionsResponse();
+            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
-        ProduceRequest.Builder builder = ProduceRequest.Builder.forCurrentMagic((short) 1,
-                1000, Collections.emptyMap());
+        ProduceRequest.Builder builder = ProduceRequest.forCurrentMagic(new ProduceRequestData()
+                .setTopicData(new ProduceRequestData.TopicProduceDataCollection())
+                .setAcks((short) 1)
+                .setTimeoutMs(1000));
         TestCallbackHandler handler = new TestCallbackHandler();
-        long now = time.milliseconds();
-        ClientRequest request = client.newClientRequest(
-                node.idString(), builder, now, true, handler);
-        client.send(request, now);
-        // sleeping to make sure that the time since last send is greater than requestTimeOut
-        time.sleep(3000);
-        client.poll(3000, time.milliseconds());
-        assertEquals(1, selector.disconnected().size());
-        assertTrue("Node not found in disconnected map", selector.disconnected().containsKey(node.idString()));
+        int requestTimeoutMs = defaultRequestTimeoutMs + 5000;
+        ClientRequest request = client.newClientRequest(node.idString(), builder, time.milliseconds(), true,
+                requestTimeoutMs, handler);
+        assertEquals(requestTimeoutMs, request.requestTimeoutMs());
+        testRequestTimeout(request);
+    }
+
+    @Test
+    public void testDefaultRequestTimeout() {
+        awaitReady(client, 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.forCurrentMagic(new ProduceRequestData()
+                .setTopicData(new ProduceRequestData.TopicProduceDataCollection())
+                .setAcks((short) 1)
+                .setTimeoutMs(1000));
+        ClientRequest request = client.newClientRequest(node.idString(), builder, time.milliseconds(), true);
+        assertEquals(defaultRequestTimeoutMs, request.requestTimeoutMs());
+        testRequestTimeout(request);
+    }
+
+    private void testRequestTimeout(ClientRequest request) {
+        client.send(request, time.milliseconds());
+
+        time.sleep(request.requestTimeoutMs() + 1);
+        List responses = client.poll(0, time.milliseconds());
+
+        assertEquals(1, responses.size());
+        ClientResponse clientResponse = responses.get(0);
+        assertEquals(node.idString(), clientResponse.destination());
+        assertTrue("Expected response to fail due to disconnection", clientResponse.wasDisconnected());
+    }
+
+    @Test
+    public void testConnectionSetupTimeout() {
+        // Use two nodes to ensure that the logic iterate over a set of more than one
+        // element. ConcurrentModificationException is not triggered otherwise.
+        final Cluster cluster = TestUtils.clusterWith(2);
+        final Node node0 = cluster.nodeById(0);
+        final Node node1 = cluster.nodeById(1);
+
+        client.ready(node0, time.milliseconds());
+        selector.serverConnectionBlocked(node0.idString());
+
+        client.ready(node1, time.milliseconds());
+        selector.serverConnectionBlocked(node1.idString());
+
+        client.poll(0, time.milliseconds());
+        assertFalse(
+                "The connections should not fail before the socket connection setup timeout elapsed",
+                client.connectionFailed(node)
+        );
+
+        time.sleep((long) (connectionSetupTimeoutMsTest * 1.2) + 1);
+        client.poll(0, time.milliseconds());
+        assertTrue(
+                "Expected the connections to fail due to the socket connection setup timeout",
+                client.connectionFailed(node)
+        );
+    }
+
+    @Test
+    public void testConnectionThrottling() {
+        // Instrument the test to return a response with a 100ms throttle delay.
+        awaitReady(client, node);
+        short requestVersion = PRODUCE.latestVersion();
+        ProduceRequest.Builder builder = new ProduceRequest.Builder(
+            requestVersion,
+            requestVersion,
+            new ProduceRequestData()
+                .setAcks((short) 1)
+                .setTimeoutMs(1000));
+        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());
+        int throttleTime = 100;
+        ProduceResponse produceResponse = new ProduceResponse(new ProduceResponseData().setThrottleTimeMs(throttleTime));
+        ByteBuffer buffer = RequestTestUtils.serializeResponseWithHeader(produceResponse, requestVersion, request.correlationId());
+        selector.completeReceive(new NetworkReceive(node.idString(), buffer));
+        client.poll(1, time.milliseconds());
+
+        // The connection is not ready due to throttling.
+        assertFalse(client.ready(node, time.milliseconds()));
+        assertEquals(100, client.throttleDelayMs(node, time.milliseconds()));
+
+        // After 50ms, the connection is not ready yet.
+        time.sleep(50);
+        assertFalse(client.ready(node, time.milliseconds()));
+        assertEquals(50, client.throttleDelayMs(node, time.milliseconds()));
+
+        // After another 50ms, the throttling is done and the connection becomes ready again.
+        time.sleep(50);
+        assertTrue(client.ready(node, time.milliseconds()));
+        assertEquals(0, client.throttleDelayMs(node, time.milliseconds()));
+    }
+
+    // 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(ApiKeys key, short maxVersion) {
+        ApiVersionsResponseKeyCollection versionList = new ApiVersionsResponseKeyCollection();
+        for (ApiKeys apiKey : ApiKeys.values()) {
+            if (apiKey == key) {
+                versionList.add(new ApiVersionsResponseKey()
+                    .setApiKey(apiKey.id)
+                    .setMinVersion((short) 0)
+                    .setMaxVersion(maxVersion));
+            } else {
+                versionList.add(new ApiVersionsResponseKey()
+                    .setApiKey(apiKey.id)
+                    .setMinVersion(apiKey.oldestVersion())
+                    .setMaxVersion(apiKey.latestVersion()));
+            }
+        }
+        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(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, (short) 5);
+        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.forCurrentMagic(new ProduceRequestData()
+                .setTopicData(new ProduceRequestData.TopicProduceDataCollection())
+                .setAcks((short) 1)
+                .setTimeoutMs(1000));
+        TestCallbackHandler handler = new TestCallbackHandler();
+        ClientRequest request = client.newClientRequest(nodeId, builder, time.milliseconds(), true,
+                defaultRequestTimeoutMs, handler);
+        client.send(request, time.milliseconds());
+        return request.correlationId();
+    }
+
+    private void sendResponse(AbstractResponse response, short version, int correlationId) {
+        ByteBuffer buffer = RequestTestUtils.serializeResponseWithHeader(response, version, correlationId);
+        selector.completeReceive(new NetworkReceive(node.idString(), buffer));
+    }
+
+    private void sendThrottledProduceResponse(int correlationId, int throttleMs, short version) {
+        ProduceResponse response = new ProduceResponse(new ProduceResponseData().setThrottleTimeMs(throttleMs));
+        sendResponse(response, version, correlationId);
     }
 
     @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()));
-        
+
         // leastloadednode should be our single node
         Node leastNode = client.leastLoadedNode(time.milliseconds());
         assertEquals("There should be one leastloadednode", leastNode.id(), node.id());
-        
+
         // sleep for longer than reconnect backoff
         time.sleep(reconnectBackoffMsTest);
-        
-        // CLOSE node 
-        selector.close(node.idString());
-        
+
+        // CLOSE node
+        selector.serverDisconnect(node.idString());
+
         client.poll(1, time.milliseconds());
         assertFalse("After we forced the disconnection the client is no longer ready.", client.ready(node, time.milliseconds()));
         leastNode = client.leastLoadedNode(time.milliseconds());
-        assertEquals("There should be NO leastloadednode", leastNode, null);
-        
+        assertNull("There should be NO leastloadednode", leastNode);
+    }
+
+    @Test
+    public void testLeastLoadedNodeProvideDisconnectedNodesPrioritizedByLastConnectionTimestamp() {
+        int nodeNumber = 3;
+        NetworkClient client = createNetworkClientWithMultipleNodes(0, connectionSetupTimeoutMsTest, nodeNumber);
+
+        Set providedNodeIds = new HashSet<>();
+        for (int i = 0; i < nodeNumber * 10; i++) {
+            Node node = client.leastLoadedNode(time.milliseconds());
+            assertNotNull("Should provide a node", node);
+            providedNodeIds.add(node);
+            client.ready(node, time.milliseconds());
+            client.disconnect(node.idString());
+            time.sleep(connectionSetupTimeoutMsTest + 1);
+            client.poll(0, time.milliseconds());
+            // Define a round as nodeNumber of nodes have been provided
+            // In each round every node should be provided exactly once
+            if ((i + 1) % nodeNumber == 0) {
+                assertEquals("All the nodes should be provided", nodeNumber, providedNodeIds.size());
+                providedNodeIds.clear();
+            }
+        }
+    }
+
+    @Test
+    public void testAuthenticationFailureWithInFlightMetadataRequest() {
+        int refreshBackoffMs = 50;
+
+        MetadataResponse metadataResponse = RequestTestUtils.metadataUpdateWith(2, Collections.emptyMap());
+        Metadata metadata = new Metadata(refreshBackoffMs, 5000, new LogContext(), new ClusterResourceListeners());
+        metadata.updateWithCurrentRequestVersion(metadataResponse, false, 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 = RequestTestUtils.serializeResponseWithHeader(metadataResponse, 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, PRODUCE.latestVersion());
+        client.poll(1, time.milliseconds());
+
+        // leastloadednode should return null since the node is throttled
+        assertNull(client.leastLoadedNode(time.milliseconds()));
     }
 
     @Test
@@ -238,8 +728,8 @@ public void testConnectionDelayConnectedWithNoExponentialBackoff() {
     public void testConnectionDelayDisconnectedWithNoExponentialBackoff() {
         awaitReady(clientWithNoExponentialBackoff, node);
 
-        selector.close(node.idString());
-        clientWithNoExponentialBackoff.poll(requestTimeoutMs, time.milliseconds());
+        selector.serverDisconnect(node.idString());
+        clientWithNoExponentialBackoff.poll(defaultRequestTimeoutMs, time.milliseconds());
         long delay = clientWithNoExponentialBackoff.connectionDelay(node, time.milliseconds());
 
         assertEquals(reconnectBackoffMsTest, delay);
@@ -250,8 +740,8 @@ public void testConnectionDelayDisconnectedWithNoExponentialBackoff() {
 
         // Start connecting and disconnect before the connection is established
         client.ready(node, time.milliseconds());
-        selector.close(node.idString());
-        client.poll(requestTimeoutMs, time.milliseconds());
+        selector.serverDisconnect(node.idString());
+        client.poll(defaultRequestTimeoutMs, time.milliseconds());
 
         // Second attempt should have the same behaviour as exponential backoff is disabled
         assertEquals(reconnectBackoffMsTest, delay);
@@ -280,11 +770,11 @@ public void testConnectionDelayDisconnected() {
         awaitReady(client, node);
 
         // First disconnection
-        selector.close(node.idString());
-        client.poll(requestTimeoutMs, time.milliseconds());
+        selector.serverDisconnect(node.idString());
+        client.poll(defaultRequestTimeoutMs, time.milliseconds());
         long delay = client.connectionDelay(node, time.milliseconds());
         long expectedDelay = reconnectBackoffMsTest;
-        double jitter = 0.2;
+        double jitter = 0.3;
         assertEquals(expectedDelay, delay, expectedDelay * jitter);
 
         // Sleep until there is no connection delay
@@ -293,13 +783,13 @@ public void testConnectionDelayDisconnected() {
 
         // Start connecting and disconnect before the connection is established
         client.ready(node, time.milliseconds());
-        selector.close(node.idString());
-        client.poll(requestTimeoutMs, time.milliseconds());
+        selector.serverDisconnect(node.idString());
+        client.poll(defaultRequestTimeoutMs, time.milliseconds());
 
         // Second attempt should take twice as long with twice the jitter
         expectedDelay = Math.round(delay * 2);
         delay = client.connectionDelay(node, time.milliseconds());
-        jitter = 0.4;
+        jitter = 0.6;
         assertEquals(expectedDelay, delay, expectedDelay * jitter);
     }
 
@@ -309,21 +799,87 @@ public void testDisconnectDuringUserMetadataRequest() {
         // metadata request when the remote node disconnects with the request in-flight.
         awaitReady(client, node);
 
-        MetadataRequest.Builder builder = new MetadataRequest.Builder(Collections.emptyList(), true);
+        MetadataRequest.Builder builder = new MetadataRequest.Builder(Collections.emptyList(), true);
         long now = time.milliseconds();
         ClientRequest request = client.newClientRequest(node.idString(), builder, now, true);
         client.send(request, now);
-        client.poll(requestTimeoutMs, now);
+        client.poll(defaultRequestTimeoutMs, now);
         assertEquals(1, client.inFlightRequestCount(node.idString()));
         assertTrue(client.hasInFlightRequests(node.idString()));
         assertTrue(client.hasInFlightRequests());
 
         selector.close(node.idString());
-        List responses = client.poll(requestTimeoutMs, time.milliseconds());
+        List responses = client.poll(defaultRequestTimeoutMs, time.milliseconds());
         assertEquals(1, responses.size());
         assertTrue(responses.iterator().next().wasDisconnected());
     }
 
+    @Test
+    public void testServerDisconnectAfterInternalApiVersionRequest() throws Exception {
+        awaitInFlightApiVersionRequest();
+        selector.serverDisconnect(node.idString());
+
+        // The failed ApiVersion request should not be forwarded to upper layers
+        List responses = client.poll(0, time.milliseconds());
+        assertFalse(client.hasInFlightRequests(node.idString()));
+        assertTrue(responses.isEmpty());
+    }
+
+    @Test
+    public void testClientDisconnectAfterInternalApiVersionRequest() throws Exception {
+        awaitInFlightApiVersionRequest();
+        client.disconnect(node.idString());
+        assertFalse(client.hasInFlightRequests(node.idString()));
+
+        // The failed ApiVersion request should not be forwarded to upper layers
+        List responses = client.poll(0, time.milliseconds());
+        assertTrue(responses.isEmpty());
+    }
+
+    @Test
+    public void testDisconnectWithMultipleInFlights() {
+        NetworkClient client = this.clientWithNoVersionDiscovery;
+        awaitReady(client, node);
+        assertTrue("Expected NetworkClient to be ready to send to node " + node.idString(),
+                client.isReady(node, time.milliseconds()));
+
+        MetadataRequest.Builder builder = new MetadataRequest.Builder(Collections.emptyList(), true);
+        long now = time.milliseconds();
+
+        final List callbackResponses = new ArrayList<>();
+        RequestCompletionHandler callback = response -> callbackResponses.add(response);
+
+        ClientRequest request1 = client.newClientRequest(node.idString(), builder, now, true, defaultRequestTimeoutMs, callback);
+        client.send(request1, now);
+        client.poll(0, now);
+
+        ClientRequest request2 = client.newClientRequest(node.idString(), builder, now, true, defaultRequestTimeoutMs, callback);
+        client.send(request2, now);
+        client.poll(0, now);
+
+        assertNotEquals(request1.correlationId(), request2.correlationId());
+
+        assertEquals(2, client.inFlightRequestCount());
+        assertEquals(2, client.inFlightRequestCount(node.idString()));
+
+        client.disconnect(node.idString());
+
+        List responses = client.poll(0, time.milliseconds());
+        assertEquals(2, responses.size());
+        assertEquals(responses, callbackResponses);
+        assertEquals(0, client.inFlightRequestCount());
+        assertEquals(0, client.inFlightRequestCount(node.idString()));
+
+        // Ensure that the responses are returned in the order they were sent
+        ClientResponse response1 = responses.get(0);
+        assertTrue(response1.wasDisconnected());
+        assertEquals(request1.correlationId(), response1.requestHeader().correlationId());
+
+        ClientResponse response2 = responses.get(1);
+        assertTrue(response2.wasDisconnected());
+        assertEquals(request2.correlationId(), response2.requestHeader().correlationId());
+    }
+
     @Test
     public void testCallDisconnect() throws Exception {
         awaitReady(client, node);
@@ -338,13 +894,37 @@ public void testCallDisconnect() throws Exception {
             client.connectionFailed(node));
         assertFalse(client.canConnect(node, time.milliseconds()));
 
-        // ensure disconnect does not reset blackout period if already disconnected
+        // ensure disconnect does not reset backoff period if already disconnected
         time.sleep(reconnectBackoffMaxMsTest);
         assertTrue(client.canConnect(node, time.milliseconds()));
         client.disconnect(node.idString());
         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(() -> {
+            client.poll(0, time.milliseconds());
+            return client.hasInFlightRequests(node.idString());
+        }, 1000, "");
+        assertFalse(client.isReady(node, time.milliseconds()));
+    }
+
     private static class TestCallbackHandler implements RequestCompletionHandler {
         public boolean executed = false;
         public ClientResponse response;
@@ -354,4 +934,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..3004d053157d1 100644
--- a/clients/src/test/java/org/apache/kafka/clients/NodeApiVersionsTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/NodeApiVersionsTest.java
@@ -16,28 +16,28 @@
  */
 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.assertNull;
 import static org.junit.Assert.assertTrue;
 
 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()) {
+        for (ApiKeys apiKey : ApiKeys.enabledApis()) {
             bld.append(prefix).append(apiKey.name).
                     append("(").append(apiKey.id).append("): UNSUPPORTED");
             prefix = ", ";
@@ -48,8 +48,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 +91,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 +105,56 @@ 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));
+            if (apiKey.isEnabled) {
+                assertEquals(apiKey.latestVersion(), versions.latestUsableVersion(apiKey));
+            } else {
+                assertNull(versions.apiVersion(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/AdminClientTestUtils.java b/clients/src/test/java/org/apache/kafka/clients/admin/AdminClientTestUtils.java
new file mode 100644
index 0000000000000..b0f00552e851d
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/admin/AdminClientTestUtils.java
@@ -0,0 +1,92 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.Map;
+import org.apache.kafka.clients.admin.CreateTopicsResult.TopicMetadataAndConfig;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.internals.KafkaFutureImpl;
+
+public class AdminClientTestUtils {
+
+    /**
+     * Helper to create a ListPartitionReassignmentsResult instance for a given Throwable.
+     * ListPartitionReassignmentsResult's constructor is only accessible from within the
+     * admin package.
+     */
+    public static ListPartitionReassignmentsResult listPartitionReassignmentsResult(Throwable t) {
+        KafkaFutureImpl> future = new KafkaFutureImpl<>();
+        future.completeExceptionally(t);
+        return new ListPartitionReassignmentsResult(future);
+    }
+
+    /**
+     * Helper to create a CreateTopicsResult instance for a given Throwable.
+     * CreateTopicsResult's constructor is only accessible from within the
+     * admin package.
+     */
+    public static CreateTopicsResult createTopicsResult(String topic, Throwable t) {
+        KafkaFutureImpl future = new KafkaFutureImpl<>();
+        future.completeExceptionally(t);
+        return new CreateTopicsResult(Collections.singletonMap(topic, future));
+    }
+
+    /**
+     * Helper to create a DeleteTopicsResult instance for a given Throwable.
+     * DeleteTopicsResult's constructor is only accessible from within the
+     * admin package.
+     */
+    public static DeleteTopicsResult deleteTopicsResult(String topic, Throwable t) {
+        KafkaFutureImpl future = new KafkaFutureImpl<>();
+        future.completeExceptionally(t);
+        return new DeleteTopicsResult(Collections.singletonMap(topic, future));
+    }
+
+    /**
+     * Helper to create a ListTopicsResult instance for a given topic.
+     * ListTopicsResult's constructor is only accessible from within the
+     * admin package.
+     */
+    public static ListTopicsResult listTopicsResult(String topic) {
+        KafkaFutureImpl> future = new KafkaFutureImpl<>();
+        future.complete(Collections.singletonMap(topic, new TopicListing(topic, false)));
+        return new ListTopicsResult(future);
+    }
+
+    /**
+     * Helper to create a CreatePartitionsResult instance for a given Throwable.
+     * CreatePartitionsResult's constructor is only accessible from within the
+     * admin package.
+     */
+    public static CreatePartitionsResult createPartitionsResult(String topic, Throwable t) {
+        KafkaFutureImpl future = new KafkaFutureImpl<>();
+        future.completeExceptionally(t);
+        return new CreatePartitionsResult(Collections.singletonMap(topic, future));
+    }
+
+    /**
+     * Helper to create a DescribeTopicsResult instance for a given topic.
+     * DescribeTopicsResult's constructor is only accessible from within the
+     * admin package.
+     */
+    public static DescribeTopicsResult describeTopicsResult(String topic, TopicDescription description) {
+        KafkaFutureImpl future = new KafkaFutureImpl<>();
+        future.complete(description);
+        return new DescribeTopicsResult(Collections.singletonMap(topic, future));
+    }
+}
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
new file mode 100644
index 0000000000000..f6a808f74727e
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/admin/AdminClientUnitTestEnv.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.clients.admin;
+
+import org.apache.kafka.clients.MockClient;
+import org.apache.kafka.clients.admin.internals.AdminMetadataManager;
+import org.apache.kafka.common.Cluster;
+import org.apache.kafka.common.Node;
+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;
+
+/**
+ * Simple utility for setting up a mock {@link KafkaAdminClient} that uses a {@link MockClient} for a supplied
+ * {@link Cluster}. Create a {@link Cluster} manually or use {@link org.apache.kafka.test.TestUtils} methods to
+ * 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 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,
+ * users should be wary of calling its methods after the {@link #adminClient() AdminClient} is instantiated.
+ *
+ * 
            
+ * When finished, be sure to {@link #close() close} the environment object.
+ */
+public class AdminClientUnitTestEnv implements AutoCloseable {
+    private final Time time;
+    private final Cluster cluster;
+    private final MockClient mockClient;
+    private final KafkaAdminClient adminClient;
+
+    public AdminClientUnitTestEnv(Cluster cluster, String... vals) {
+        this(Time.SYSTEM, cluster, vals);
+    }
+
+    public AdminClientUnitTestEnv(Time time, Cluster cluster, String... vals) {
+        this(time, cluster, clientConfigs(vals));
+    }
+
+    public AdminClientUnitTestEnv(Time time, Cluster cluster) {
+        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);
+
+        AdminMetadataManager metadataManager = new AdminMetadataManager(new LogContext(),
+                adminClientConfig.getLong(AdminClientConfig.RETRY_BACKOFF_MS_CONFIG),
+                adminClientConfig.getLong(AdminClientConfig.METADATA_MAX_AGE_CONFIG));
+        this.mockClient = new MockClient(time, new MockClient.MockMetadataUpdater() {
+            @Override
+            public List fetchNodes() {
+                return cluster.nodes();
+            }
+
+            @Override
+            public boolean isUpdateNeeded() {
+                return false;
+            }
+
+            @Override
+            public void update(Time time, MockClient.MetadataUpdate update) {
+                throw new UnsupportedOperationException();
+            }
+        });
+
+        metadataManager.update(cluster, time.milliseconds());
+        unreachableNodes.forEach(mockClient::setUnreachable);
+        this.adminClient = KafkaAdminClient.createInternal(adminClientConfig, metadataManager, mockClient, time);
+    }
+
+    public Time time() {
+        return time;
+    }
+
+    public Cluster cluster() {
+        return cluster;
+    }
+
+    public Admin adminClient() {
+        return adminClient;
+    }
+
+    public MockClient kafkaClient() {
+        return mockClient;
+    }
+
+    @Override
+    public void close() {
+        this.adminClient.close();
+    }
+
+    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 (overrides.length % 2 != 0) {
+            throw new IllegalStateException();
+        }
+        for (int i = 0; i < overrides.length; i += 2) {
+            map.put(overrides[i], overrides[i + 1]);
+        }
+        return map;
+    }
+
+    public static String kafkaAdminClientNetworkThreadPrefix() {
+        return KafkaAdminClient.NETWORK_THREAD_PREFIX;
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/clients/admin/ConfigTest.java b/clients/src/test/java/org/apache/kafka/clients/admin/ConfigTest.java
new file mode 100644
index 0000000000000..ed4c0bde0b8c1
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/admin/ConfigTest.java
@@ -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 org.apache.kafka.clients.admin;
+
+import org.apache.kafka.clients.admin.ConfigEntry.ConfigType;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+
+import static org.hamcrest.CoreMatchers.containsString;
+import static org.hamcrest.CoreMatchers.equalTo;
+import static org.hamcrest.CoreMatchers.hasItems;
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.CoreMatchers.not;
+import static org.hamcrest.CoreMatchers.nullValue;
+import static org.hamcrest.MatcherAssert.assertThat;
+
+public class ConfigTest {
+    private static final ConfigEntry E1 = new ConfigEntry("a", "b");
+    private static final ConfigEntry E2 = new ConfigEntry("c", "d");
+    private Config config;
+
+    @Before
+    public void setUp() {
+        final Collection entries = new ArrayList<>();
+        entries.add(E1);
+        entries.add(E2);
+
+        config = new Config(entries);
+    }
+
+    @Test
+    public void shouldGetEntry() {
+        assertThat(config.get("a"), is(E1));
+        assertThat(config.get("c"), is(E2));
+    }
+
+    @Test
+    public void shouldReturnNullOnGetUnknownEntry() {
+        assertThat(config.get("unknown"), is(nullValue()));
+    }
+
+    @Test
+    public void shouldGetAllEntries() {
+        assertThat(config.entries().size(), is(2));
+        assertThat(config.entries(), hasItems(E1, E2));
+    }
+
+    @Test
+    public void shouldImplementEqualsProperly() {
+        final Collection entries = new ArrayList<>();
+        entries.add(E1);
+
+        assertThat(config, is(equalTo(config)));
+        assertThat(config, is(equalTo(new Config(config.entries()))));
+        assertThat(config, is(not(equalTo(new Config(entries)))));
+        assertThat(config, is(not(equalTo((Object) "this"))));
+    }
+
+    @Test
+    public void shouldImplementHashCodeProperly() {
+        final Collection entries = new ArrayList<>();
+        entries.add(E1);
+
+        assertThat(config.hashCode(), is(config.hashCode()));
+        assertThat(config.hashCode(), is(new Config(config.entries()).hashCode()));
+        assertThat(config.hashCode(), is(not(new Config(entries).hashCode())));
+    }
+
+    @Test
+    public void shouldImplementToStringProperly() {
+        assertThat(config.toString(), containsString(E1.toString()));
+        assertThat(config.toString(), containsString(E2.toString()));
+    }
+
+    public static ConfigEntry newConfigEntry(String name, String value, ConfigEntry.ConfigSource source, boolean isSensitive,
+                                             boolean isReadOnly, List synonyms) {
+        return new ConfigEntry(name, value, source, isSensitive, isReadOnly, synonyms, ConfigType.UNKNOWN, null);
+    }
+}
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/DescribeUserScramCredentialsResultTest.java b/clients/src/test/java/org/apache/kafka/clients/admin/DescribeUserScramCredentialsResultTest.java
new file mode 100644
index 0000000000000..9faf02ac90759
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/admin/DescribeUserScramCredentialsResultTest.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.internals.KafkaFutureImpl;
+import org.apache.kafka.common.message.DescribeUserScramCredentialsResponseData;
+import org.apache.kafka.common.protocol.Errors;
+import org.junit.Test;
+
+import java.util.Arrays;
+import java.util.Map;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.fail;
+
+public class DescribeUserScramCredentialsResultTest {
+    @Test
+    public void testTopLevelError() {
+        KafkaFutureImpl dataFuture = new KafkaFutureImpl<>();
+        dataFuture.completeExceptionally(new RuntimeException());
+        DescribeUserScramCredentialsResult results = new DescribeUserScramCredentialsResult(dataFuture);
+        try {
+            results.all().get();
+            fail("expected all() to fail when there is a top-level error");
+        } catch (Exception expected) {
+            // ignore, expected
+        }
+        try {
+            results.users().get();
+            fail("expected users() to fail when there is a top-level error");
+        } catch (Exception expected) {
+            // ignore, expected
+        }
+        try {
+            results.description("whatever").get();
+            fail("expected description() to fail when there is a top-level error");
+        } catch (Exception expected) {
+            // ignore, expected
+        }
+    }
+
+    @Test
+    public void testUserLevelErrors() throws Exception {
+        String goodUser = "goodUser";
+        String unknownUser = "unknownUser";
+        String failedUser = "failedUser";
+        KafkaFutureImpl dataFuture = new KafkaFutureImpl<>();
+        ScramMechanism scramSha256 = ScramMechanism.SCRAM_SHA_256;
+        int iterations = 4096;
+        dataFuture.complete(new DescribeUserScramCredentialsResponseData().setErrorCode(Errors.NONE.code()).setResults(Arrays.asList(
+                new DescribeUserScramCredentialsResponseData.DescribeUserScramCredentialsResult().setUser(goodUser).setCredentialInfos(
+                        Arrays.asList(new DescribeUserScramCredentialsResponseData.CredentialInfo().setMechanism(scramSha256.type()).setIterations(iterations))),
+                new DescribeUserScramCredentialsResponseData.DescribeUserScramCredentialsResult().setUser(unknownUser).setErrorCode(Errors.RESOURCE_NOT_FOUND.code()),
+                new DescribeUserScramCredentialsResponseData.DescribeUserScramCredentialsResult().setUser(failedUser).setErrorCode(Errors.DUPLICATE_RESOURCE.code()))));
+        DescribeUserScramCredentialsResult results = new DescribeUserScramCredentialsResult(dataFuture);
+        try {
+            results.all().get();
+            fail("expected all() to fail when there is a user-level error");
+        } catch (Exception expected) {
+            // ignore, expected
+        }
+        assertEquals("Expected 2 users with credentials", Arrays.asList(goodUser, failedUser), results.users().get());
+        UserScramCredentialsDescription goodUserDescription = results.description(goodUser).get();
+        assertEquals(new UserScramCredentialsDescription(goodUser, Arrays.asList(new ScramCredentialInfo(scramSha256, iterations))), goodUserDescription);
+        try {
+            results.description(failedUser).get();
+            fail("expected description(failedUser) to fail when there is a user-level error");
+        } catch (Exception expected) {
+            // ignore, expected
+        }
+        try {
+            results.description(unknownUser).get();
+            fail("expected description(unknownUser) to fail when there is no such user");
+        } catch (Exception expected) {
+            // ignore, expected
+        }
+    }
+
+    @Test
+    public void testSuccessfulDescription() throws Exception {
+        String goodUser = "goodUser";
+        String unknownUser = "unknownUser";
+        KafkaFutureImpl dataFuture = new KafkaFutureImpl<>();
+        ScramMechanism scramSha256 = ScramMechanism.SCRAM_SHA_256;
+        int iterations = 4096;
+        dataFuture.complete(new DescribeUserScramCredentialsResponseData().setErrorCode(Errors.NONE.code()).setResults(Arrays.asList(
+                new DescribeUserScramCredentialsResponseData.DescribeUserScramCredentialsResult().setUser(goodUser).setCredentialInfos(
+                        Arrays.asList(new DescribeUserScramCredentialsResponseData.CredentialInfo().setMechanism(scramSha256.type()).setIterations(iterations))))));
+        DescribeUserScramCredentialsResult results = new DescribeUserScramCredentialsResult(dataFuture);
+        assertEquals("Expected 1 user with credentials", Arrays.asList(goodUser), results.users().get());
+        Map allResults = results.all().get();
+        assertEquals(1, allResults.size());
+        UserScramCredentialsDescription goodUserDescriptionViaAll = allResults.get(goodUser);
+        assertEquals(new UserScramCredentialsDescription(goodUser, Arrays.asList(new ScramCredentialInfo(scramSha256, iterations))), goodUserDescriptionViaAll);
+        assertEquals("Expected same thing via all() and description()", goodUserDescriptionViaAll, results.description(goodUser).get());
+        try {
+            results.description(unknownUser).get();
+            fail("expected description(unknownUser) to fail when there is no such user even when all() succeeds");
+        } catch (Exception expected) {
+            // ignore, expected
+        }
+    }
+}
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 c0fe73c36ed20..8980289824e3f 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
@@ -16,73 +16,204 @@
  */
 package org.apache.kafka.clients.admin;
 
+import org.apache.kafka.clients.ApiVersion;
+import org.apache.kafka.clients.ClientDnsLookup;
+import org.apache.kafka.clients.ClientUtils;
+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.admin.ListOffsetsResult.ListOffsetsResultInfo;
+import org.apache.kafka.clients.admin.internals.ConsumerGroupOperationContext;
+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.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.Node;
 import org.apache.kafka.common.PartitionInfo;
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.TopicPartitionReplica;
 import org.apache.kafka.common.acl.AccessControlEntry;
 import org.apache.kafka.common.acl.AccessControlEntryFilter;
 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.acl.AclPermissionType;
+import org.apache.kafka.common.config.ConfigException;
 import org.apache.kafka.common.config.ConfigResource;
+import org.apache.kafka.common.errors.ApiException;
+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;
-import org.apache.kafka.common.errors.NotLeaderForPartitionException;
+import org.apache.kafka.common.errors.LogDirNotFoundException;
+import org.apache.kafka.common.errors.ThrottlingQuotaExceededException;
+import org.apache.kafka.common.errors.TimeoutException;
+import org.apache.kafka.common.errors.NotLeaderOrFollowerException;
 import org.apache.kafka.common.errors.OffsetOutOfRangeException;
+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.TopicExistsException;
+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.errors.UnsupportedVersionException;
+import org.apache.kafka.common.feature.Features;
+import org.apache.kafka.common.message.AlterPartitionReassignmentsResponseData;
+import org.apache.kafka.common.message.AlterReplicaLogDirsResponseData;
+import org.apache.kafka.common.message.AlterReplicaLogDirsResponseData.AlterReplicaLogDirPartitionResult;
+import org.apache.kafka.common.message.AlterReplicaLogDirsResponseData.AlterReplicaLogDirTopicResult;
+import org.apache.kafka.common.message.AlterUserScramCredentialsResponseData;
+import org.apache.kafka.common.message.ApiVersionsResponseData;
+import org.apache.kafka.common.message.CreatePartitionsResponseData;
+import org.apache.kafka.common.message.CreatePartitionsResponseData.CreatePartitionsTopicResult;
+import org.apache.kafka.common.message.CreateAclsResponseData;
+import org.apache.kafka.common.message.CreateTopicsResponseData;
+import org.apache.kafka.common.message.CreateTopicsResponseData.CreatableTopicResult;
+import org.apache.kafka.common.message.CreateTopicsResponseData.CreatableTopicResultCollection;
+import org.apache.kafka.common.message.DeleteAclsResponseData;
+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.DeleteRecordsResponseData;
+import org.apache.kafka.common.message.DeleteTopicsResponseData;
+import org.apache.kafka.common.message.DeleteTopicsResponseData.DeletableTopicResult;
+import org.apache.kafka.common.message.DeleteTopicsResponseData.DeletableTopicResultCollection;
+import org.apache.kafka.common.message.DescribeAclsResponseData;
+import org.apache.kafka.common.message.DescribeConfigsResponseData;
+import org.apache.kafka.common.message.DescribeGroupsResponseData;
+import org.apache.kafka.common.message.DescribeGroupsResponseData.DescribedGroupMember;
+import org.apache.kafka.common.message.DescribeLogDirsResponseData;
+import org.apache.kafka.common.message.DescribeLogDirsResponseData.DescribeLogDirsTopic;
+import org.apache.kafka.common.message.DescribeUserScramCredentialsResponseData;
+import org.apache.kafka.common.message.DescribeUserScramCredentialsResponseData.CredentialInfo;
+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.ListOffsetResponseData;
+import org.apache.kafka.common.message.ListOffsetResponseData.ListOffsetTopicResponse;
+import org.apache.kafka.common.message.ListPartitionReassignmentsResponseData;
+import org.apache.kafka.common.message.MetadataResponseData.MetadataResponseTopic;
+import org.apache.kafka.common.message.MetadataResponseData.MetadataResponsePartition;
+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.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.requests.CreatePartitionsResponse;
+import org.apache.kafka.common.protocol.Message;
+import org.apache.kafka.common.quota.ClientQuotaAlteration;
+import org.apache.kafka.common.quota.ClientQuotaEntity;
+import org.apache.kafka.common.quota.ClientQuotaFilter;
+import org.apache.kafka.common.quota.ClientQuotaFilterComponent;
+import org.apache.kafka.common.requests.AbstractResponse;
+import org.apache.kafka.common.requests.AlterClientQuotasResponse;
+import org.apache.kafka.common.requests.AlterPartitionReassignmentsResponse;
+import org.apache.kafka.common.requests.AlterReplicaLogDirsResponse;
+import org.apache.kafka.common.requests.AlterUserScramCredentialsResponse;
 import org.apache.kafka.common.requests.ApiError;
+import org.apache.kafka.common.requests.ApiVersionsRequest;
+import org.apache.kafka.common.requests.ApiVersionsResponse;
 import org.apache.kafka.common.requests.CreateAclsResponse;
-import org.apache.kafka.common.requests.CreateAclsResponse.AclCreationResponse;
+import org.apache.kafka.common.requests.CreatePartitionsRequest;
+import org.apache.kafka.common.requests.CreatePartitionsResponse;
+import org.apache.kafka.common.requests.CreateTopicsRequest;
 import org.apache.kafka.common.requests.CreateTopicsResponse;
 import org.apache.kafka.common.requests.DeleteAclsResponse;
-import org.apache.kafka.common.requests.DeleteAclsResponse.AclDeletionResult;
-import org.apache.kafka.common.requests.DeleteAclsResponse.AclFilterResponse;
+import org.apache.kafka.common.requests.DeleteGroupsResponse;
 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.DescribeAclsResponse;
+import org.apache.kafka.common.requests.DescribeClientQuotasResponse;
 import org.apache.kafka.common.requests.DescribeConfigsResponse;
+import org.apache.kafka.common.requests.DescribeGroupsResponse;
+import org.apache.kafka.common.requests.DescribeLogDirsResponse;
+import org.apache.kafka.common.requests.DescribeUserScramCredentialsResponse;
+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.JoinGroupRequest;
+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.ListOffsetResponse;
+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.resource.Resource;
-import org.apache.kafka.common.resource.ResourceFilter;
+import org.apache.kafka.common.requests.OffsetCommitResponse;
+import org.apache.kafka.common.requests.OffsetDeleteResponse;
+import org.apache.kafka.common.requests.OffsetFetchResponse;
+import org.apache.kafka.common.requests.RequestTestUtils;
+import org.apache.kafka.common.requests.UpdateFeaturesRequest;
+import org.apache.kafka.common.requests.UpdateFeaturesResponse;
+import org.apache.kafka.common.resource.PatternType;
+import org.apache.kafka.common.resource.ResourcePattern;
+import org.apache.kafka.common.resource.ResourcePatternFilter;
 import org.apache.kafka.common.resource.ResourceType;
 import org.apache.kafka.common.utils.MockTime;
-import org.apache.kafka.test.TestCondition;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.common.utils.Utils;
 import org.apache.kafka.test.TestUtils;
-import org.junit.Ignore;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.Timeout;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+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;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
+import java.util.OptionalInt;
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Future;
+import java.util.concurrent.Semaphore;
+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 org.apache.kafka.common.requests.ResourceType.TOPIC;
-import static org.apache.kafka.common.requests.ResourceType.BROKER;
+import static java.util.Collections.emptyList;
+import static java.util.Collections.emptySet;
+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.assertNotEquals;
 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;
 
-
 /**
  * A unit test for KafkaAdminClient.
  *
@@ -94,6 +225,14 @@ public class KafkaAdminClientTest {
     @Rule
     final public Timeout globalTimeout = Timeout.millis(120000);
 
+    @Test
+    public void testDefaultApiTimeoutAndRequestTimeoutConflicts() {
+        final AdminClientConfig config = newConfMap(AdminClientConfig.DEFAULT_API_TIMEOUT_MS_CONFIG, "500");
+        KafkaException exception = assertThrows(KafkaException.class,
+            () -> KafkaAdminClient.createInternal(config, null));
+        assertTrue(exception.getCause() instanceof ConfigException);
+    }
+
     @Test
     public void testGetOrCreateListValue() {
         Map> map = new HashMap<>();
@@ -155,34 +294,245 @@ public void testGenerateClientId() {
                 KafkaAdminClient.generateClientId(newConfMap(AdminClientConfig.CLIENT_ID_CONFIG, "myCustomId")));
     }
 
-    private static MockKafkaAdminClientEnv mockClientEnv(String... configVals) {
+    private static Cluster mockCluster(int numNodes, int controllerIndex) {
         HashMap nodes = new HashMap<>();
-        nodes.put(0, new Node(0, "localhost", 8121));
-        nodes.put(1, new Node(1, "localhost", 8122));
-        nodes.put(2, new Node(2, "localhost", 8123));
-        Cluster cluster = new Cluster("mockClusterId", nodes.values(),
-                Collections.emptySet(), Collections.emptySet(),
-                Collections.emptySet(), nodes.get(0));
-        return new MockKafkaAdminClientEnv(cluster, configVals);
+        for (int i = 0; i < numNodes; i++)
+            nodes.put(i, new Node(i, "localhost", 8121 + i));
+        return new Cluster("mockClusterId", nodes.values(),
+            Collections.emptySet(), Collections.emptySet(),
+            Collections.emptySet(), nodes.get(controllerIndex));
+    }
+
+    private static Cluster mockBootstrapCluster() {
+        return Cluster.bootstrap(ClientUtils.parseAndValidateAddresses(
+                singletonList("localhost:8121"), ClientDnsLookup.USE_ALL_DNS_IPS));
+    }
+
+    private static AdminClientUnitTestEnv mockClientEnv(String... configVals) {
+        return new AdminClientUnitTestEnv(mockCluster(3, 0), configVals);
+    }
+
+    private static AdminClientUnitTestEnv mockClientEnv(Time time, String... configVals) {
+        return new AdminClientUnitTestEnv(time, mockCluster(3, 0), configVals);
     }
 
     @Test
-    public void testCloseAdminClient() throws Exception {
-        try (MockKafkaAdminClientEnv env = mockClientEnv()) {
+    public void testCloseAdminClient() {
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
         }
     }
 
-    private static void assertFutureError(Future future, Class exceptionClass)
-            throws InterruptedException {
-        try {
-            future.get();
-            fail("Expected a " + exceptionClass.getSimpleName() + " exception, but got success.");
-        } catch (ExecutionException ee) {
-            Throwable cause = ee.getCause();
-            assertEquals("Expected a " + exceptionClass.getSimpleName() + " exception, but got " +
-                            cause.getClass().getSimpleName(),
-                    exceptionClass, cause.getClass());
+    /**
+     * Test if admin client can be closed in the callback invoked when
+     * an api call completes. If calling {@link Admin#close()} in callback, AdminClient thread hangs
+     */
+    @Test(timeout = 10_000)
+    public void testCloseAdminClientInCallback() throws InterruptedException {
+        MockTime time = new MockTime();
+        AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(time, mockCluster(3, 0));
+
+        final ListTopicsResult result = env.adminClient().listTopics(new ListTopicsOptions().timeoutMs(1000));
+        final KafkaFuture> kafkaFuture = result.listings();
+        final Semaphore callbackCalled = new Semaphore(0);
+        kafkaFuture.whenComplete((topicListings, throwable) -> {
+            env.close();
+            callbackCalled.release();
+        });
+
+        time.sleep(2000); // Advance time to timeout and complete listTopics request
+        callbackCalled.acquire();
+    }
+
+    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 OffsetCommitResponse prepareOffsetCommitResponse(TopicPartition tp, Errors error) {
+        Map responseData = new HashMap<>();
+        responseData.put(tp, error);
+        return new OffsetCommitResponse(0, responseData);
+    }
+
+    private static CreateTopicsResponse prepareCreateTopicsResponse(String topicName, Errors error) {
+        CreateTopicsResponseData data = new CreateTopicsResponseData();
+        data.topics().add(new CreatableTopicResult()
+            .setName(topicName)
+            .setErrorCode(error.code()));
+        return new CreateTopicsResponse(data);
+    }
+
+    public static CreateTopicsResponse prepareCreateTopicsResponse(int throttleTimeMs, CreatableTopicResult... topics) {
+        CreateTopicsResponseData data = new CreateTopicsResponseData()
+            .setThrottleTimeMs(throttleTimeMs)
+            .setTopics(new CreatableTopicResultCollection(Arrays.stream(topics).iterator()));
+        return new CreateTopicsResponse(data);
+    }
+
+    public static CreatableTopicResult creatableTopicResult(String name, Errors error) {
+        return new CreatableTopicResult()
+            .setName(name)
+            .setErrorCode(error.code());
+    }
+
+    public static DeleteTopicsResponse prepareDeleteTopicsResponse(int throttleTimeMs, DeletableTopicResult... topics) {
+        DeleteTopicsResponseData data = new DeleteTopicsResponseData()
+            .setThrottleTimeMs(throttleTimeMs)
+            .setResponses(new DeletableTopicResultCollection(Arrays.stream(topics).iterator()));
+        return new DeleteTopicsResponse(data);
+    }
+
+    public static DeletableTopicResult deletableTopicResult(String topicName, Errors error) {
+        return new DeletableTopicResult()
+            .setName(topicName)
+            .setErrorCode(error.code());
+    }
+
+    public static CreatePartitionsResponse prepareCreatePartitionsResponse(int throttleTimeMs, CreatePartitionsTopicResult... topics) {
+        CreatePartitionsResponseData data = new CreatePartitionsResponseData()
+            .setThrottleTimeMs(throttleTimeMs)
+            .setResults(Arrays.asList(topics));
+        return new CreatePartitionsResponse(data);
+    }
+
+    public static CreatePartitionsTopicResult createPartitionsTopicResult(String name, Errors error) {
+        return createPartitionsTopicResult(name, error, null);
+    }
+
+    public static CreatePartitionsTopicResult createPartitionsTopicResult(String name, Errors error, String errorMessage) {
+        return new CreatePartitionsTopicResult()
+            .setName(name)
+            .setErrorCode(error.code())
+            .setErrorMessage(errorMessage);
+    }
+
+    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);
+    }
+
+    private static MetadataResponse prepareMetadataResponse(Cluster cluster, Errors error) {
+        List metadata = new ArrayList<>();
+        for (String topic : cluster.topics()) {
+            List pms = new ArrayList<>();
+            for (PartitionInfo pInfo : cluster.availablePartitionsForTopic(topic)) {
+                MetadataResponsePartition pm  = new MetadataResponsePartition()
+                    .setErrorCode(error.code())
+                    .setPartitionIndex(pInfo.partition())
+                    .setLeaderId(pInfo.leader().id())
+                    .setLeaderEpoch(234)
+                    .setReplicaNodes(Arrays.stream(pInfo.replicas()).map(Node::id).collect(Collectors.toList()))
+                    .setIsrNodes(Arrays.stream(pInfo.inSyncReplicas()).map(Node::id).collect(Collectors.toList()))
+                    .setOfflineReplicas(Arrays.stream(pInfo.offlineReplicas()).map(Node::id).collect(Collectors.toList()));
+                pms.add(pm);
+            }
+            MetadataResponseTopic tm = new MetadataResponseTopic()
+                .setErrorCode(error.code())
+                .setName(topic)
+                .setIsInternal(false)
+                .setPartitions(pms);
+            metadata.add(tm);
+        }
+        return MetadataResponse.prepareResponse(true,
+            0,
+            cluster.nodes(),
+            cluster.clusterResource().clusterId(),
+            cluster.controller().id(),
+            metadata,
+            MetadataResponse.AUTHORIZED_OPERATIONS_OMITTED);
+    }
+
+    private static DescribeGroupsResponseData prepareDescribeGroupsResponseData(String groupId,
+                                                                                List groupInstances,
+                                                                                List topicPartitions) {
+        final ByteBuffer memberAssignment = ConsumerProtocol.serializeAssignment(new ConsumerPartitionAssignor.Assignment(topicPartitions));
+        List describedGroupMembers = groupInstances.stream().map(groupInstance -> DescribeGroupsResponse.groupMember(JoinGroupRequest.UNKNOWN_MEMBER_ID,
+                groupInstance, "clientId0", "clientHost", new byte[memberAssignment.remaining()], null)).collect(Collectors.toList());
+        DescribeGroupsResponseData data = new DescribeGroupsResponseData();
+        data.groups().add(DescribeGroupsResponse.groupMetadata(
+                groupId,
+                Errors.NONE,
+                "",
+                ConsumerProtocol.PROTOCOL_TYPE,
+                "",
+                describedGroupMembers,
+                Collections.emptySet()));
+        return data;
+    }
+
+    private static FeatureMetadata defaultFeatureMetadata() {
+        return new FeatureMetadata(
+            Utils.mkMap(Utils.mkEntry("test_feature_1", new FinalizedVersionRange((short) 2, (short) 3))),
+            Optional.of(1L),
+            Utils.mkMap(Utils.mkEntry("test_feature_1", new SupportedVersionRange((short) 1, (short) 5))));
+    }
+
+    private static Features convertSupportedFeaturesMap(Map features) {
+        final Map featuresMap = new HashMap<>();
+        for (final Map.Entry entry : features.entrySet()) {
+            final SupportedVersionRange versionRange = entry.getValue();
+            featuresMap.put(
+                entry.getKey(),
+                new org.apache.kafka.common.feature.SupportedVersionRange(versionRange.minVersion(),
+                                                                          versionRange.maxVersion()));
+        }
+
+        return Features.supportedFeatures(featuresMap);
+    }
+
+    private static Features convertFinalizedFeaturesMap(Map features) {
+        final Map featuresMap = new HashMap<>();
+        for (final Map.Entry entry : features.entrySet()) {
+            final FinalizedVersionRange versionRange = entry.getValue();
+            featuresMap.put(
+                entry.getKey(),
+                new org.apache.kafka.common.feature.FinalizedVersionRange(
+                    versionRange.minVersionLevel(), versionRange.maxVersionLevel()));
         }
+
+        return Features.finalizedFeatures(featuresMap);
+    }
+
+    private static ApiVersionsResponse prepareApiVersionsResponseForDescribeFeatures(Errors error) {
+        if (error == Errors.NONE) {
+            return new ApiVersionsResponse(ApiVersionsResponse.createApiVersionsResponseData(
+                ApiVersionsResponse.DEFAULT_API_VERSIONS_RESPONSE.throttleTimeMs(),
+                error,
+                ApiVersionsResponse.DEFAULT_API_VERSIONS_RESPONSE.data().apiKeys(),
+                convertSupportedFeaturesMap(defaultFeatureMetadata().supportedFeatures()),
+                convertFinalizedFeaturesMap(defaultFeatureMetadata().finalizedFeatures()),
+                defaultFeatureMetadata().finalizedFeaturesEpoch().get()));
+        }
+        return new ApiVersionsResponse(
+            new ApiVersionsResponseData()
+                .setThrottleTimeMs(ApiVersionsResponse.DEFAULT_API_VERSIONS_RESPONSE.throttleTimeMs())
+                .setErrorCode(error.code()));
     }
 
     /**
@@ -190,363 +540,4468 @@ private static void assertFutureError(Future future, Class future = env.adminClient().createTopics(
-                    Collections.singleton(new NewTopic("myTopic", Collections.singletonMap(Integer.valueOf(0), asList(new Integer[]{0, 1, 2})))),
+                    Collections.singleton(new NewTopic("myTopic", Collections.singletonMap(0, asList(0, 1, 2)))),
                     new CreateTopicsOptions().timeoutMs(1000)).all();
-            assertFutureError(future, TimeoutException.class);
+            TestUtils.assertFutureError(future, TimeoutException.class);
         }
     }
 
     @Test
-    public void testCreateTopics() throws Exception {
-        try (MockKafkaAdminClientEnv env = mockClientEnv()) {
+    public void testConnectionFailureOnMetadataUpdate() throws Exception {
+        // This tests the scenario in which we successfully connect to the bootstrap server, but
+        // the server disconnects before sending the full response
+
+        Cluster cluster = mockBootstrapCluster();
+        try (final AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(Time.SYSTEM, cluster)) {
+            Cluster discoveredCluster = mockCluster(3, 0);
             env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
-            env.kafkaClient().prepareMetadataUpdate(env.cluster(), Collections.emptySet());
-            env.kafkaClient().setNode(env.cluster().controller());
-            env.kafkaClient().prepareResponse(new CreateTopicsResponse(Collections.singletonMap("myTopic", new ApiError(Errors.NONE, ""))));
+            env.kafkaClient().prepareResponse(request -> request instanceof MetadataRequest, null, true);
+            env.kafkaClient().prepareResponse(request -> request instanceof MetadataRequest,
+                    RequestTestUtils.metadataResponse(discoveredCluster.nodes(), discoveredCluster.clusterResource().clusterId(),
+                            1, Collections.emptyList()));
+            env.kafkaClient().prepareResponse(body -> body instanceof CreateTopicsRequest,
+                    prepareCreateTopicsResponse("myTopic", Errors.NONE));
+
             KafkaFuture future = env.adminClient().createTopics(
-                    Collections.singleton(new NewTopic("myTopic", Collections.singletonMap(Integer.valueOf(0), asList(new Integer[]{0, 1, 2})))),
+                    Collections.singleton(new NewTopic("myTopic", Collections.singletonMap(0, asList(0, 1, 2)))),
                     new CreateTopicsOptions().timeoutMs(10000)).all();
+
             future.get();
         }
     }
 
-    private static final AclBinding ACL1 = new AclBinding(new Resource(ResourceType.TOPIC, "mytopic3"),
-        new AccessControlEntry("User:ANONYMOUS", "*", AclOperation.DESCRIBE, AclPermissionType.ALLOW));
-    private static final AclBinding ACL2 = new AclBinding(new Resource(ResourceType.TOPIC, "mytopic4"),
-        new AccessControlEntry("User:ANONYMOUS", "*", AclOperation.DESCRIBE, AclPermissionType.DENY));
-    private static final AclBindingFilter FILTER1 = new AclBindingFilter(new ResourceFilter(ResourceType.ANY, null),
-        new AccessControlEntryFilter("User:ANONYMOUS", null, AclOperation.ANY, AclPermissionType.ANY));
-    private static final AclBindingFilter FILTER2 = new AclBindingFilter(new ResourceFilter(ResourceType.ANY, null),
-        new AccessControlEntryFilter("User:bob", null, AclOperation.ANY, AclPermissionType.ANY));
-
     @Test
-    public void testDescribeAcls() throws Exception {
-        try (MockKafkaAdminClientEnv env = mockClientEnv()) {
-            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
-            env.kafkaClient().prepareMetadataUpdate(env.cluster(), Collections.emptySet());
-            env.kafkaClient().setNode(env.cluster().controller());
+    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
 
-            // Test a call where we get back ACL1 and ACL2.
-            env.kafkaClient().prepareResponse(new DescribeAclsResponse(0, ApiError.NONE,
-                    asList(ACL1, ACL2)));
-            assertCollectionIs(env.adminClient().describeAcls(FILTER1).values().get(), ACL1, ACL2);
+        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(3, 0);
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+            env.kafkaClient().prepareResponse(body -> body instanceof MetadataRequest,
+                    RequestTestUtils.metadataResponse(discoveredCluster.nodes(), discoveredCluster.clusterResource().clusterId(),
+                            1, Collections.emptyList()));
+            env.kafkaClient().prepareResponse(body -> body instanceof CreateTopicsRequest,
+                prepareCreateTopicsResponse("myTopic", Errors.NONE));
 
-            // Test a call where we get back no results.
-            env.kafkaClient().prepareResponse(new DescribeAclsResponse(0, ApiError.NONE,
-                Collections.emptySet()));
-            assertTrue(env.adminClient().describeAcls(FILTER2).values().get().isEmpty());
+            KafkaFuture future = env.adminClient().createTopics(
+                    Collections.singleton(new NewTopic("myTopic", Collections.singletonMap(0, asList(0, 1, 2)))),
+                    new CreateTopicsOptions().timeoutMs(10000)).all();
 
-            // Test a call where we get back an error.
-            env.kafkaClient().prepareResponse(new DescribeAclsResponse(0,
-                new ApiError(Errors.SECURITY_DISABLED, "Security is disabled"), Collections.emptySet()));
-            assertFutureError(env.adminClient().describeAcls(FILTER2).values(), SecurityDisabledException.class);
+            future.get();
         }
     }
 
+    /**
+     * Test that we propagate exceptions encountered when fetching metadata.
+     */
     @Test
-    public void testCreateAcls() throws Exception {
-        try (MockKafkaAdminClientEnv env = mockClientEnv()) {
+    public void testPropagatedMetadataFetchException() throws Exception {
+        try (final AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(Time.SYSTEM,
+                mockCluster(3, 0),
+                newStrMap(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:8121",
+                AdminClientConfig.REQUEST_TIMEOUT_MS_CONFIG, "10"))) {
             env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
-            env.kafkaClient().prepareMetadataUpdate(env.cluster(), Collections.emptySet());
-            env.kafkaClient().setNode(env.cluster().controller());
-
-            // Test a call where we successfully create two ACLs.
-            env.kafkaClient().prepareResponse(new CreateAclsResponse(0,
-                asList(new AclCreationResponse(ApiError.NONE), new AclCreationResponse(ApiError.NONE))));
-            CreateAclsResult results = env.adminClient().createAcls(asList(ACL1, ACL2));
-            assertCollectionIs(results.values().keySet(), ACL1, ACL2);
-            for (KafkaFuture future : results.values().values())
-                future.get();
-            results.all().get();
-
-            // Test a call where we fail to create one ACL.
-            env.kafkaClient().prepareResponse(new CreateAclsResponse(0, asList(
-                new AclCreationResponse(new ApiError(Errors.SECURITY_DISABLED, "Security is disabled")),
-                new AclCreationResponse(ApiError.NONE))
-            ));
-            results = env.adminClient().createAcls(asList(ACL1, ACL2));
-            assertCollectionIs(results.values().keySet(), ACL1, ACL2);
-            assertFutureError(results.values().get(ACL1), SecurityDisabledException.class);
-            results.values().get(ACL2).get();
-            assertFutureError(results.all(), SecurityDisabledException.class);
+            env.kafkaClient().createPendingAuthenticationError(env.cluster().nodeById(0),
+                    TimeUnit.DAYS.toMillis(1));
+            env.kafkaClient().prepareResponse(prepareCreateTopicsResponse("myTopic", Errors.NONE));
+            KafkaFuture future = env.adminClient().createTopics(
+                Collections.singleton(new NewTopic("myTopic", Collections.singletonMap(0, asList(0, 1, 2)))),
+                new CreateTopicsOptions().timeoutMs(1000)).all();
+            TestUtils.assertFutureError(future, SaslAuthenticationException.class);
         }
     }
 
     @Test
-    public void testDeleteAcls() throws Exception {
-        try (MockKafkaAdminClientEnv env = mockClientEnv()) {
+    public void testCreateTopics() throws Exception {
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
             env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
-            env.kafkaClient().prepareMetadataUpdate(env.cluster(), Collections.emptySet());
-            env.kafkaClient().setNode(env.cluster().controller());
-
-            // Test a call where one filter has an error.
-            env.kafkaClient().prepareResponse(new DeleteAclsResponse(0, asList(
-                    new AclFilterResponse(asList(new AclDeletionResult(ACL1), new AclDeletionResult(ACL2))),
-                    new AclFilterResponse(new ApiError(Errors.SECURITY_DISABLED, "No security"),
-                            Collections.emptySet()))));
-            DeleteAclsResult results = env.adminClient().deleteAcls(asList(FILTER1, FILTER2));
-            Map> filterResults = results.values();
-            FilterResults filter1Results = filterResults.get(FILTER1).get();
-            assertEquals(null, filter1Results.values().get(0).exception());
-            assertEquals(ACL1, filter1Results.values().get(0).binding());
-            assertEquals(null, filter1Results.values().get(1).exception());
-            assertEquals(ACL2, filter1Results.values().get(1).binding());
-            assertFutureError(filterResults.get(FILTER2), SecurityDisabledException.class);
-            assertFutureError(results.all(), SecurityDisabledException.class);
-
-            // Test a call where one deletion result has an error.
-            env.kafkaClient().prepareResponse(new DeleteAclsResponse(0, asList(
-                    new AclFilterResponse(asList(new AclDeletionResult(ACL1),
-                            new AclDeletionResult(new ApiError(Errors.SECURITY_DISABLED, "No security"), ACL2))),
-                    new AclFilterResponse(Collections.emptySet()))));
-            results = env.adminClient().deleteAcls(asList(FILTER1, FILTER2));
-            assertTrue(results.values().get(FILTER2).get().values().isEmpty());
-            assertFutureError(results.all(), SecurityDisabledException.class);
+            env.kafkaClient().prepareResponse(
+                expectCreateTopicsRequestWithTopics("myTopic"),
+                prepareCreateTopicsResponse("myTopic", Errors.NONE));
+            KafkaFuture future = env.adminClient().createTopics(
+                    Collections.singleton(new NewTopic("myTopic", Collections.singletonMap(0, asList(0, 1, 2)))),
+                    new CreateTopicsOptions().timeoutMs(10000)).all();
+            future.get();
+        }
+    }
 
-            // Test a call where there are no errors.
-            env.kafkaClient().prepareResponse(new DeleteAclsResponse(0, asList(
-                    new AclFilterResponse(asList(new AclDeletionResult(ACL1))),
-                    new AclFilterResponse(asList(new AclDeletionResult(ACL2))))));
-            results = env.adminClient().deleteAcls(asList(FILTER1, FILTER2));
-            Collection deleted = results.all().get();
-            assertCollectionIs(deleted, ACL1, ACL2);
+    @Test
+    public void testCreateTopicsPartialResponse() throws Exception {
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+            env.kafkaClient().prepareResponse(
+                expectCreateTopicsRequestWithTopics("myTopic", "myTopic2"),
+                prepareCreateTopicsResponse("myTopic", Errors.NONE));
+            CreateTopicsResult topicsResult = env.adminClient().createTopics(
+                    asList(new NewTopic("myTopic", Collections.singletonMap(0, asList(0, 1, 2))),
+                           new NewTopic("myTopic2", Collections.singletonMap(0, asList(0, 1, 2)))),
+                    new CreateTopicsOptions().timeoutMs(10000));
+            topicsResult.values().get("myTopic").get();
+            TestUtils.assertFutureThrows(topicsResult.values().get("myTopic2"), ApiException.class);
         }
     }
 
-    /**
-     * Test handling timeouts.
-     */
-    @Ignore // The test is flaky. Should be renabled when this JIRA is fixed: https://issues.apache.org/jira/browse/KAFKA-5792
     @Test
-    public void testHandleTimeout() throws Exception {
-        HashMap nodes = new HashMap<>();
+    public void testCreateTopicsRetryBackoff() throws Exception {
         MockTime time = new MockTime();
-        nodes.put(0, new Node(0, "localhost", 8121));
-        Cluster cluster = new Cluster("mockClusterId", nodes.values(),
-            Collections.emptySet(), Collections.emptySet(),
-            Collections.emptySet(), nodes.get(0));
-        try (MockKafkaAdminClientEnv env = new MockKafkaAdminClientEnv(time, cluster,
-            AdminClientConfig.RECONNECT_BACKOFF_MAX_MS_CONFIG, "1",
-                AdminClientConfig.RECONNECT_BACKOFF_MS_CONFIG, "1")) {
-            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
-            env.kafkaClient().prepareMetadataUpdate(env.cluster(), Collections.emptySet());
-            env.kafkaClient().setNode(nodes.get(0));
-            assertEquals(time, env.time());
-            assertEquals(env.time(), ((KafkaAdminClient) env.adminClient()).time());
-
-            // Make a request with an extremely short timeout.
-            // Then wait for it to fail by not supplying any response.
-            log.info("Starting AdminClient#listTopics...");
-            final ListTopicsResult result = env.adminClient().listTopics(new ListTopicsOptions().timeoutMs(1000));
-            TestUtils.waitForCondition(new TestCondition() {
-                @Override
-                public boolean conditionMet() {
-                    return env.kafkaClient().hasInFlightRequests();
-                }
-            }, "Timed out waiting for inFlightRequests");
-            time.sleep(5000);
-            TestUtils.waitForCondition(new TestCondition() {
-                @Override
-                public boolean conditionMet() {
-                    return result.listings().isDone();
-                }
-            }, "Timed out waiting for listTopics to complete");
-            assertFutureError(result.listings(), TimeoutException.class);
-            log.info("Verified the error result of AdminClient#listTopics");
+        int retryBackoff = 100;
+
+        try (final AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(time,
+                mockCluster(3, 0),
+                newStrMap(AdminClientConfig.RETRY_BACKOFF_MS_CONFIG, "" + retryBackoff))) {
+            MockClient mockClient = env.kafkaClient();
+
+            mockClient.setNodeApiVersions(NodeApiVersions.create());
+
+            AtomicLong firstAttemptTime = new AtomicLong(0);
+            AtomicLong secondAttemptTime = new AtomicLong(0);
+
+            mockClient.prepareResponse(body -> {
+                firstAttemptTime.set(time.milliseconds());
+                return body instanceof CreateTopicsRequest;
+            }, null, true);
+
+            mockClient.prepareResponse(body -> {
+                secondAttemptTime.set(time.milliseconds());
+                return body instanceof CreateTopicsRequest;
+            }, prepareCreateTopicsResponse("myTopic", Errors.NONE));
+
+            KafkaFuture future = env.adminClient().createTopics(
+                Collections.singleton(new NewTopic("myTopic", Collections.singletonMap(0, asList(0, 1, 2)))),
+                new CreateTopicsOptions().timeoutMs(10000)).all();
+
+            // 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);
 
-            // The next request should succeed.
-            time.sleep(5000);
-            env.kafkaClient().prepareResponse(new DescribeConfigsResponse(0,
-                Collections.singletonMap(new org.apache.kafka.common.requests.Resource(TOPIC, "foo"),
-                    new DescribeConfigsResponse.Config(ApiError.NONE,
-                        Collections.emptySet()))));
-            DescribeConfigsResult result2 = env.adminClient().describeConfigs(Collections.singleton(
-                new ConfigResource(ConfigResource.Type.TOPIC, "foo")));
-            time.sleep(5000);
-            result2.values().get(new ConfigResource(ConfigResource.Type.TOPIC, "foo")).get();
+            future.get();
+
+            long actualRetryBackoff = secondAttemptTime.get() - firstAttemptTime.get();
+            assertEquals("CreateTopics retry did not await expected backoff",
+                    retryBackoff, actualRetryBackoff);
         }
     }
 
     @Test
-    public void testDescribeConfigs() throws Exception {
-        try (MockKafkaAdminClientEnv env = mockClientEnv()) {
+    public void testCreateTopicsHandleNotControllerException() throws Exception {
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
             env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
-            env.kafkaClient().prepareMetadataUpdate(env.cluster(), Collections.emptySet());
-            env.kafkaClient().setNode(env.cluster().controller());
-            env.kafkaClient().prepareResponse(new DescribeConfigsResponse(0,
-                Collections.singletonMap(new org.apache.kafka.common.requests.Resource(BROKER, "0"),
-                    new DescribeConfigsResponse.Config(ApiError.NONE,
-                        Collections.emptySet()))));
-            DescribeConfigsResult result2 = env.adminClient().describeConfigs(Collections.singleton(
-                new ConfigResource(ConfigResource.Type.BROKER, "0")));
-            result2.all().get();
+            env.kafkaClient().prepareResponseFrom(
+                prepareCreateTopicsResponse("myTopic", Errors.NOT_CONTROLLER),
+                env.cluster().nodeById(0));
+            env.kafkaClient().prepareResponse(RequestTestUtils.metadataResponse(env.cluster().nodes(),
+                env.cluster().clusterResource().clusterId(),
+                1,
+                Collections.emptyList()));
+            env.kafkaClient().prepareResponseFrom(
+                prepareCreateTopicsResponse("myTopic", Errors.NONE),
+                env.cluster().nodeById(1));
+            KafkaFuture future = env.adminClient().createTopics(
+                Collections.singleton(new NewTopic("myTopic", Collections.singletonMap(0, asList(0, 1, 2)))),
+                new CreateTopicsOptions().timeoutMs(10000)).all();
+            future.get();
         }
     }
 
     @Test
-    public void testCreatePartitions() throws Exception {
-        try (MockKafkaAdminClientEnv env = mockClientEnv()) {
+    public void testCreateTopicsRetryThrottlingExceptionWhenEnabled() throws Exception {
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
             env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
-            env.kafkaClient().prepareMetadataUpdate(env.cluster(), Collections.emptySet());
-            env.kafkaClient().setNode(env.cluster().controller());
 
-            Map m = new HashMap<>();
-            m.put("my_topic", ApiError.NONE);
-            m.put("other_topic", ApiError.fromThrowable(new InvalidTopicException("some detailed reason")));
+            env.kafkaClient().prepareResponse(
+                expectCreateTopicsRequestWithTopics("topic1", "topic2", "topic3"),
+                prepareCreateTopicsResponse(1000,
+                    creatableTopicResult("topic1", Errors.NONE),
+                    creatableTopicResult("topic2", Errors.THROTTLING_QUOTA_EXCEEDED),
+                    creatableTopicResult("topic3", Errors.TOPIC_ALREADY_EXISTS)));
 
-            // Test a call where one filter has an error.
-            env.kafkaClient().prepareResponse(new CreatePartitionsResponse(0, m));
+            env.kafkaClient().prepareResponse(
+                expectCreateTopicsRequestWithTopics("topic2"),
+                prepareCreateTopicsResponse(1000,
+                    creatableTopicResult("topic2", Errors.THROTTLING_QUOTA_EXCEEDED)));
 
-            Map counts = new HashMap<>();
-            counts.put("my_topic", NewPartitions.increaseTo(3));
-            counts.put("other_topic", NewPartitions.increaseTo(3, asList(asList(2), asList(3))));
+            env.kafkaClient().prepareResponse(
+                expectCreateTopicsRequestWithTopics("topic2"),
+                prepareCreateTopicsResponse(0,
+                    creatableTopicResult("topic2", Errors.NONE)));
 
-            CreatePartitionsResult results = env.adminClient().createPartitions(counts);
-            Map> values = results.values();
-            KafkaFuture myTopicResult = values.get("my_topic");
-            myTopicResult.get();
-            KafkaFuture otherTopicResult = values.get("other_topic");
-            try {
-                otherTopicResult.get();
-                fail("get() should throw ExecutionException");
-            } catch (ExecutionException e0) {
-                assertTrue(e0.getCause() instanceof InvalidTopicException);
-                InvalidTopicException e = (InvalidTopicException) e0.getCause();
-                assertEquals("some detailed reason", e.getMessage());
-            }
+            CreateTopicsResult result = env.adminClient().createTopics(
+                asList(
+                    new NewTopic("topic1", 1, (short) 1),
+                    new NewTopic("topic2", 1, (short) 1),
+                    new NewTopic("topic3", 1, (short) 1)),
+                new CreateTopicsOptions().retryOnQuotaViolation(true));
+
+            assertNull(result.values().get("topic1").get());
+            assertNull(result.values().get("topic2").get());
+            TestUtils.assertFutureThrows(result.values().get("topic3"), TopicExistsException.class);
         }
     }
 
     @Test
-    public void testDeleteRecords() throws Exception {
-
-        HashMap nodes = new HashMap<>();
-        nodes.put(0, new Node(0, "localhost", 8121));
-        List partitionInfos = new ArrayList<>();
-        partitionInfos.add(new PartitionInfo("my_topic", 0, nodes.get(0), new Node[] {nodes.get(0)}, new Node[] {nodes.get(0)}));
-        partitionInfos.add(new PartitionInfo("my_topic", 1, nodes.get(0), new Node[] {nodes.get(0)}, new Node[] {nodes.get(0)}));
-        partitionInfos.add(new PartitionInfo("my_topic", 2, null, new Node[] {nodes.get(0)}, new Node[] {nodes.get(0)}));
-        partitionInfos.add(new PartitionInfo("my_topic", 3, nodes.get(0), new Node[] {nodes.get(0)}, new Node[] {nodes.get(0)}));
-        partitionInfos.add(new PartitionInfo("my_topic", 4, nodes.get(0), new Node[] {nodes.get(0)}, new Node[] {nodes.get(0)}));
-        Cluster cluster = new Cluster("mockClusterId", nodes.values(),
-                partitionInfos, Collections.emptySet(),
-                Collections.emptySet(), nodes.get(0));
+    public void testCreateTopicsRetryThrottlingExceptionWhenEnabledUntilRequestTimeOut() throws Exception {
+        long defaultApiTimeout = 60000;
+        MockTime time = new MockTime();
 
-        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);
-        TopicPartition myTopicPartition4 = new TopicPartition("my_topic", 4);
+        try (AdminClientUnitTestEnv env = mockClientEnv(time,
+            AdminClientConfig.DEFAULT_API_TIMEOUT_MS_CONFIG, String.valueOf(defaultApiTimeout))) {
 
-        try (MockKafkaAdminClientEnv env = new MockKafkaAdminClientEnv(cluster)) {
             env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
-            env.kafkaClient().prepareMetadataUpdate(env.cluster(), Collections.emptySet());
-            env.kafkaClient().setNode(env.cluster().nodes().get(0));
 
-            Map m = new HashMap<>();
-            m.put(myTopicPartition0,
-                    new DeleteRecordsResponse.PartitionResponse(3, Errors.NONE));
-            m.put(myTopicPartition1,
-                    new DeleteRecordsResponse.PartitionResponse(DeleteRecordsResponse.INVALID_LOW_WATERMARK, Errors.OFFSET_OUT_OF_RANGE));
-            m.put(myTopicPartition3,
-                    new DeleteRecordsResponse.PartitionResponse(DeleteRecordsResponse.INVALID_LOW_WATERMARK, Errors.NOT_LEADER_FOR_PARTITION));
-            m.put(myTopicPartition4,
-                    new DeleteRecordsResponse.PartitionResponse(DeleteRecordsResponse.INVALID_LOW_WATERMARK, Errors.UNKNOWN_TOPIC_OR_PARTITION));
+            env.kafkaClient().prepareResponse(
+                expectCreateTopicsRequestWithTopics("topic1", "topic2", "topic3"),
+                prepareCreateTopicsResponse(1000,
+                    creatableTopicResult("topic1", Errors.NONE),
+                    creatableTopicResult("topic2", Errors.THROTTLING_QUOTA_EXCEEDED),
+                    creatableTopicResult("topic3", Errors.TOPIC_ALREADY_EXISTS)));
 
-            List t = new ArrayList<>();
-            List p = new ArrayList<>();
-            p.add(new MetadataResponse.PartitionMetadata(Errors.NONE, 0, nodes.get(0),
-                    Collections.singletonList(nodes.get(0)), Collections.singletonList(nodes.get(0)), Collections.emptyList()));
-            p.add(new MetadataResponse.PartitionMetadata(Errors.NONE, 1, nodes.get(0),
-                    Collections.singletonList(nodes.get(0)), Collections.singletonList(nodes.get(0)), Collections.emptyList()));
-            p.add(new MetadataResponse.PartitionMetadata(Errors.LEADER_NOT_AVAILABLE, 2, null,
-                    Collections.singletonList(nodes.get(0)), Collections.singletonList(nodes.get(0)), Collections.emptyList()));
-            p.add(new MetadataResponse.PartitionMetadata(Errors.NONE, 3, nodes.get(0),
-                    Collections.singletonList(nodes.get(0)), Collections.singletonList(nodes.get(0)), Collections.emptyList()));
-            p.add(new MetadataResponse.PartitionMetadata(Errors.NONE, 4, nodes.get(0),
-                    Collections.singletonList(nodes.get(0)), Collections.singletonList(nodes.get(0)), Collections.emptyList()));
+            env.kafkaClient().prepareResponse(
+                expectCreateTopicsRequestWithTopics("topic2"),
+                prepareCreateTopicsResponse(1000,
+                    creatableTopicResult("topic2", Errors.THROTTLING_QUOTA_EXCEEDED)));
 
-            t.add(new MetadataResponse.TopicMetadata(Errors.NONE, "my_topic", false, p));
+            CreateTopicsResult result = env.adminClient().createTopics(
+                asList(
+                    new NewTopic("topic1", 1, (short) 1),
+                    new NewTopic("topic2", 1, (short) 1),
+                    new NewTopic("topic3", 1, (short) 1)),
+                new CreateTopicsOptions().retryOnQuotaViolation(true));
 
-            env.kafkaClient().prepareResponse(new MetadataResponse(cluster.nodes(), cluster.clusterResource().clusterId(), cluster.controller().id(), t));
-            env.kafkaClient().prepareResponse(new DeleteRecordsResponse(0, m));
+            // Wait until the prepared attempts have consumed
+            TestUtils.waitForCondition(() -> env.kafkaClient().numAwaitingResponses() == 0,
+                "Failed awaiting CreateTopics requests");
 
-            Map recordsToDelete = new HashMap<>();
-            recordsToDelete.put(myTopicPartition0, RecordsToDelete.beforeOffset(3L));
-            recordsToDelete.put(myTopicPartition1, RecordsToDelete.beforeOffset(10L));
-            recordsToDelete.put(myTopicPartition2, RecordsToDelete.beforeOffset(10L));
-            recordsToDelete.put(myTopicPartition3, RecordsToDelete.beforeOffset(10L));
-            recordsToDelete.put(myTopicPartition4, RecordsToDelete.beforeOffset(10L));
+            // Wait until the next request is sent out
+            TestUtils.waitForCondition(() -> env.kafkaClient().inFlightRequestCount() == 1,
+                "Failed awaiting next CreateTopics request");
 
-            DeleteRecordsResult results = env.adminClient().deleteRecords(recordsToDelete);
+            // Advance time past the default api timeout to time out the inflight request
+            time.sleep(defaultApiTimeout + 1);
 
-            // success on records deletion for partition 0
-            Map> values = results.lowWatermarks();
-            KafkaFuture myTopicPartition0Result = values.get(myTopicPartition0);
-            long lowWatermark = myTopicPartition0Result.get().lowWatermark();
-            assertEquals(lowWatermark, 3);
+            assertNull(result.values().get("topic1").get());
+            ThrottlingQuotaExceededException e = TestUtils.assertFutureThrows(result.values().get("topic2"),
+                ThrottlingQuotaExceededException.class);
+            assertEquals(0, e.throttleTimeMs());
+            TestUtils.assertFutureThrows(result.values().get("topic3"), TopicExistsException.class);
+        }
+    }
 
-            // "offset out of range" failure on records deletion for partition 1
-            KafkaFuture myTopicPartition1Result = values.get(myTopicPartition1);
-            try {
-                myTopicPartition1Result.get();
-                fail("get() should throw ExecutionException");
-            } catch (ExecutionException e0) {
-                assertTrue(e0.getCause() instanceof OffsetOutOfRangeException);
-            }
+    @Test
+    public void testCreateTopicsDontRetryThrottlingExceptionWhenDisabled() throws Exception {
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
 
-            // "leader not available" failure on metadata request for partition 2
-            KafkaFuture myTopicPartition2Result = values.get(myTopicPartition2);
-            try {
-                myTopicPartition2Result.get();
-                fail("get() should throw ExecutionException");
-            } catch (ExecutionException e1) {
-                assertTrue(e1.getCause() instanceof LeaderNotAvailableException);
-            }
+            env.kafkaClient().prepareResponse(
+                expectCreateTopicsRequestWithTopics("topic1", "topic2", "topic3"),
+                prepareCreateTopicsResponse(1000,
+                    creatableTopicResult("topic1", Errors.NONE),
+                    creatableTopicResult("topic2", Errors.THROTTLING_QUOTA_EXCEEDED),
+                    creatableTopicResult("topic3", Errors.TOPIC_ALREADY_EXISTS)));
 
-            // "not leader for partition" failure on records deletion for partition 3
-            KafkaFuture myTopicPartition3Result = values.get(myTopicPartition3);
-            try {
-                myTopicPartition3Result.get();
-                fail("get() should throw ExecutionException");
-            } catch (ExecutionException e1) {
-                assertTrue(e1.getCause() instanceof NotLeaderForPartitionException);
-            }
+            CreateTopicsResult result = env.adminClient().createTopics(
+                asList(
+                    new NewTopic("topic1", 1, (short) 1),
+                    new NewTopic("topic2", 1, (short) 1),
+                    new NewTopic("topic3", 1, (short) 1)),
+                new CreateTopicsOptions().retryOnQuotaViolation(false));
 
-            // "unknown topic or partition" failure on records deletion for partition 4
-            KafkaFuture myTopicPartition4Result = values.get(myTopicPartition4);
-            try {
-                myTopicPartition4Result.get();
-                fail("get() should throw ExecutionException");
-            } catch (ExecutionException e1) {
-                assertTrue(e1.getCause() instanceof UnknownTopicOrPartitionException);
-            }
+            assertNull(result.values().get("topic1").get());
+            ThrottlingQuotaExceededException e = TestUtils.assertFutureThrows(result.values().get("topic2"),
+                ThrottlingQuotaExceededException.class);
+            assertEquals(1000, e.throttleTimeMs());
+            TestUtils.assertFutureThrows(result.values().get("topic3"), TopicExistsException.class);
         }
     }
 
-    private static  void assertCollectionIs(Collection collection, T... elements) {
-        for (T element : elements) {
-            assertTrue("Did not find " + element, collection.contains(element));
-        }
-        assertEquals("There are unexpected extra elements in the collection.",
-            elements.length, collection.size());
+    private MockClient.RequestMatcher expectCreateTopicsRequestWithTopics(final String... topics) {
+        return body -> {
+            if (body instanceof CreateTopicsRequest) {
+                CreateTopicsRequest request = (CreateTopicsRequest) body;
+                for (String topic : topics) {
+                    if (request.data().topics().find(topic) == null)
+                        return false;
+                }
+                return topics.length == request.data().topics().size();
+            }
+            return false;
+        };
     }
 
-    public static KafkaAdminClient createInternal(AdminClientConfig config, KafkaAdminClient.TimeoutProcessorFactory timeoutProcessorFactory) {
-        return KafkaAdminClient.createInternal(config, timeoutProcessorFactory);
-    }
+    @Test
+    public void testDeleteTopics() throws Exception {
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
 
-    public static class FailureInjectingTimeoutProcessorFactory extends KafkaAdminClient.TimeoutProcessorFactory {
+            env.kafkaClient().prepareResponse(
+                expectDeleteTopicsRequestWithTopics("myTopic"),
+                prepareDeleteTopicsResponse("myTopic", Errors.NONE));
+            KafkaFuture future = env.adminClient().deleteTopics(singletonList("myTopic"),
+                new DeleteTopicsOptions()).all();
+            assertNull(future.get());
 
-        private int numTries = 0;
+            env.kafkaClient().prepareResponse(
+                expectDeleteTopicsRequestWithTopics("myTopic"),
+                prepareDeleteTopicsResponse("myTopic", Errors.TOPIC_DELETION_DISABLED));
+            future = env.adminClient().deleteTopics(singletonList("myTopic"),
+                new DeleteTopicsOptions()).all();
+            TestUtils.assertFutureError(future, TopicDeletionDisabledException.class);
+
+            env.kafkaClient().prepareResponse(
+                expectDeleteTopicsRequestWithTopics("myTopic"),
+                prepareDeleteTopicsResponse("myTopic", Errors.UNKNOWN_TOPIC_OR_PARTITION));
+            future = env.adminClient().deleteTopics(singletonList("myTopic"),
+                new DeleteTopicsOptions()).all();
+            TestUtils.assertFutureError(future, UnknownTopicOrPartitionException.class);
+        }
+    }
+
+    @Test
+    public void testDeleteTopicsPartialResponse() throws Exception {
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            env.kafkaClient().prepareResponse(
+                expectDeleteTopicsRequestWithTopics("myTopic", "myOtherTopic"),
+                prepareDeleteTopicsResponse(1000,
+                    deletableTopicResult("myTopic", Errors.NONE)));
+
+            DeleteTopicsResult result = env.adminClient().deleteTopics(
+                asList("myTopic", "myOtherTopic"), new DeleteTopicsOptions());
+
+            result.values().get("myTopic").get();
+            TestUtils.assertFutureThrows(result.values().get("myOtherTopic"), ApiException.class);
+        }
+    }
+
+    @Test
+    public void testDeleteTopicsRetryThrottlingExceptionWhenEnabled() throws Exception {
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            env.kafkaClient().prepareResponse(
+                expectDeleteTopicsRequestWithTopics("topic1", "topic2", "topic3"),
+                prepareDeleteTopicsResponse(1000,
+                    deletableTopicResult("topic1", Errors.NONE),
+                    deletableTopicResult("topic2", Errors.THROTTLING_QUOTA_EXCEEDED),
+                    deletableTopicResult("topic3", Errors.TOPIC_ALREADY_EXISTS)));
+
+            env.kafkaClient().prepareResponse(
+                expectDeleteTopicsRequestWithTopics("topic2"),
+                prepareDeleteTopicsResponse(1000,
+                    deletableTopicResult("topic2", Errors.THROTTLING_QUOTA_EXCEEDED)));
+
+            env.kafkaClient().prepareResponse(
+                expectDeleteTopicsRequestWithTopics("topic2"),
+                prepareDeleteTopicsResponse(0,
+                    deletableTopicResult("topic2", Errors.NONE)));
+
+            DeleteTopicsResult result = env.adminClient().deleteTopics(
+                asList("topic1", "topic2", "topic3"),
+                new DeleteTopicsOptions().retryOnQuotaViolation(true));
+
+            assertNull(result.values().get("topic1").get());
+            assertNull(result.values().get("topic2").get());
+            TestUtils.assertFutureThrows(result.values().get("topic3"), TopicExistsException.class);
+        }
+    }
+
+    @Test
+    public void testDeleteTopicsRetryThrottlingExceptionWhenEnabledUntilRequestTimeOut() throws Exception {
+        long defaultApiTimeout = 60000;
+        MockTime time = new MockTime();
+
+        try (AdminClientUnitTestEnv env = mockClientEnv(time,
+            AdminClientConfig.DEFAULT_API_TIMEOUT_MS_CONFIG, String.valueOf(defaultApiTimeout))) {
+
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            env.kafkaClient().prepareResponse(
+                expectDeleteTopicsRequestWithTopics("topic1", "topic2", "topic3"),
+                prepareDeleteTopicsResponse(1000,
+                    deletableTopicResult("topic1", Errors.NONE),
+                    deletableTopicResult("topic2", Errors.THROTTLING_QUOTA_EXCEEDED),
+                    deletableTopicResult("topic3", Errors.TOPIC_ALREADY_EXISTS)));
+
+            env.kafkaClient().prepareResponse(
+                expectDeleteTopicsRequestWithTopics("topic2"),
+                prepareDeleteTopicsResponse(1000,
+                    deletableTopicResult("topic2", Errors.THROTTLING_QUOTA_EXCEEDED)));
+
+            DeleteTopicsResult result = env.adminClient().deleteTopics(
+                asList("topic1", "topic2", "topic3"),
+                new DeleteTopicsOptions().retryOnQuotaViolation(true));
+
+            // Wait until the prepared attempts have consumed
+            TestUtils.waitForCondition(() -> env.kafkaClient().numAwaitingResponses() == 0,
+                "Failed awaiting DeleteTopics requests");
+
+            // Wait until the next request is sent out
+            TestUtils.waitForCondition(() -> env.kafkaClient().inFlightRequestCount() == 1,
+                "Failed awaiting next DeleteTopics request");
+
+            // Advance time past the default api timeout to time out the inflight request
+            time.sleep(defaultApiTimeout + 1);
+
+            assertNull(result.values().get("topic1").get());
+            ThrottlingQuotaExceededException e = TestUtils.assertFutureThrows(result.values().get("topic2"),
+                ThrottlingQuotaExceededException.class);
+            assertEquals(0, e.throttleTimeMs());
+            TestUtils.assertFutureThrows(result.values().get("topic3"), TopicExistsException.class);
+        }
+    }
+
+    @Test
+    public void testDeleteTopicsDontRetryThrottlingExceptionWhenDisabled() throws Exception {
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            env.kafkaClient().prepareResponse(
+                expectDeleteTopicsRequestWithTopics("topic1", "topic2", "topic3"),
+                prepareDeleteTopicsResponse(1000,
+                    deletableTopicResult("topic1", Errors.NONE),
+                    deletableTopicResult("topic2", Errors.THROTTLING_QUOTA_EXCEEDED),
+                    deletableTopicResult("topic3", Errors.TOPIC_ALREADY_EXISTS)));
+
+            DeleteTopicsResult result = env.adminClient().deleteTopics(
+                asList("topic1", "topic2", "topic3"),
+                new DeleteTopicsOptions().retryOnQuotaViolation(false));
+
+            assertNull(result.values().get("topic1").get());
+            ThrottlingQuotaExceededException e = TestUtils.assertFutureThrows(result.values().get("topic2"),
+                ThrottlingQuotaExceededException.class);
+            assertEquals(1000, e.throttleTimeMs());
+            TestUtils.assertFutureError(result.values().get("topic3"), TopicExistsException.class);
+        }
+    }
+
+    private MockClient.RequestMatcher expectDeleteTopicsRequestWithTopics(final String... topics) {
+        return body -> {
+            if (body instanceof DeleteTopicsRequest) {
+                DeleteTopicsRequest request = (DeleteTopicsRequest) body;
+                return request.data().topicNames().equals(Arrays.asList(topics));
+            }
+            return false;
+        };
+    }
+
+    @Test
+    public void testInvalidTopicNames() throws Exception {
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            List sillyTopicNames = asList("", null);
+            Map> deleteFutures = env.adminClient().deleteTopics(sillyTopicNames).values();
+            for (String sillyTopicName : sillyTopicNames) {
+                TestUtils.assertFutureError(deleteFutures.get(sillyTopicName), InvalidTopicException.class);
+            }
+            assertEquals(0, env.kafkaClient().inFlightRequestCount());
+
+            Map> describeFutures =
+                    env.adminClient().describeTopics(sillyTopicNames).values();
+            for (String sillyTopicName : sillyTopicNames) {
+                TestUtils.assertFutureError(describeFutures.get(sillyTopicName), InvalidTopicException.class);
+            }
+            assertEquals(0, env.kafkaClient().inFlightRequestCount());
+
+            List newTopics = new ArrayList<>();
+            for (String sillyTopicName : sillyTopicNames) {
+                newTopics.add(new NewTopic(sillyTopicName, 1, (short) 1));
+            }
+
+            Map> createFutures = env.adminClient().createTopics(newTopics).values();
+            for (String sillyTopicName : sillyTopicNames) {
+                TestUtils.assertFutureError(createFutures .get(sillyTopicName), InvalidTopicException.class);
+            }
+            assertEquals(0, env.kafkaClient().inFlightRequestCount());
+        }
+    }
+
+    @Test
+    public void testMetadataRetries() throws Exception {
+        // We should continue retrying on metadata update failures in spite of retry configuration
+
+        String topic = "topic";
+        Cluster bootstrapCluster = Cluster.bootstrap(singletonList(new InetSocketAddress("localhost", 9999)));
+        Cluster initializedCluster = mockCluster(3, 0);
+
+        try (final AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(Time.SYSTEM, bootstrapCluster,
+                newStrMap(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999",
+                        AdminClientConfig.DEFAULT_API_TIMEOUT_MS_CONFIG, "10000000",
+                        AdminClientConfig.RETRIES_CONFIG, "0"))) {
+
+            // The first request fails with a disconnect
+            env.kafkaClient().prepareResponse(null, true);
+
+            // The next one succeeds and gives us the controller id
+            env.kafkaClient().prepareResponse(RequestTestUtils.metadataResponse(initializedCluster.nodes(),
+                    initializedCluster.clusterResource().clusterId(),
+                    initializedCluster.controller().id(),
+                    Collections.emptyList()));
+
+            // Then we respond to the DescribeTopic request
+            Node leader = initializedCluster.nodes().get(0);
+            MetadataResponse.PartitionMetadata partitionMetadata = new MetadataResponse.PartitionMetadata(
+                    Errors.NONE, new TopicPartition(topic, 0), Optional.of(leader.id()), Optional.of(10),
+                    singletonList(leader.id()), singletonList(leader.id()), singletonList(leader.id()));
+            env.kafkaClient().prepareResponse(RequestTestUtils.metadataResponse(initializedCluster.nodes(),
+                    initializedCluster.clusterResource().clusterId(), 1,
+                    singletonList(new MetadataResponse.TopicMetadata(Errors.NONE, topic, false,
+                            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());
+        }
+    }
+
+    @Test
+    public void testAdminClientApisAuthenticationFailure() throws Exception {
+        Cluster cluster = mockBootstrapCluster();
+        try (final AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(Time.SYSTEM, cluster,
+                newStrMap(AdminClientConfig.REQUEST_TIMEOUT_MS_CONFIG, "1000"))) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+            env.kafkaClient().createPendingAuthenticationError(cluster.nodes().get(0),
+                    TimeUnit.DAYS.toMillis(1));
+            callAdminClientApisAndExpectAnAuthenticationError(env);
+            callClientQuotasApisAndExpectAnAuthenticationError(env);
+        }
+    }
+
+    private void callAdminClientApisAndExpectAnAuthenticationError(AdminClientUnitTestEnv env) throws InterruptedException {
+        try {
+            env.adminClient().createTopics(
+                    Collections.singleton(new NewTopic("myTopic", Collections.singletonMap(0, asList(0, 1, 2)))),
+                    new CreateTopicsOptions().timeoutMs(10000)).all().get();
+            fail("Expected an authentication error.");
+        } catch (ExecutionException e) {
+            assertTrue("Expected an authentication error, but got " + Utils.stackTrace(e),
+                e.getCause() instanceof AuthenticationException);
+        }
+
+        try {
+            Map counts = new HashMap<>();
+            counts.put("my_topic", NewPartitions.increaseTo(3));
+            counts.put("other_topic", NewPartitions.increaseTo(3, asList(asList(2), asList(3))));
+            env.adminClient().createPartitions(counts).all().get();
+            fail("Expected an authentication error.");
+        } catch (ExecutionException e) {
+            assertTrue("Expected an authentication error, but got " + Utils.stackTrace(e),
+                e.getCause() instanceof AuthenticationException);
+        }
+
+        try {
+            env.adminClient().createAcls(asList(ACL1, ACL2)).all().get();
+            fail("Expected an authentication error.");
+        } catch (ExecutionException e) {
+            assertTrue("Expected an authentication error, but got " + Utils.stackTrace(e),
+                e.getCause() instanceof AuthenticationException);
+        }
+
+        try {
+            env.adminClient().describeAcls(FILTER1).values().get();
+            fail("Expected an authentication error.");
+        } catch (ExecutionException e) {
+            assertTrue("Expected an authentication error, but got " + Utils.stackTrace(e),
+                e.getCause() instanceof AuthenticationException);
+        }
+
+        try {
+            env.adminClient().deleteAcls(asList(FILTER1, FILTER2)).all().get();
+            fail("Expected an authentication error.");
+        } catch (ExecutionException e) {
+            assertTrue("Expected an authentication error, but got " + Utils.stackTrace(e),
+                e.getCause() instanceof AuthenticationException);
+        }
+
+        try {
+            env.adminClient().describeConfigs(Collections.singleton(new ConfigResource(ConfigResource.Type.BROKER, "0"))).all().get();
+            fail("Expected an authentication error.");
+        } catch (ExecutionException e) {
+            assertTrue("Expected an authentication error, but got " + Utils.stackTrace(e),
+                e.getCause() instanceof AuthenticationException);
+        }
+    }
+
+    private void callClientQuotasApisAndExpectAnAuthenticationError(AdminClientUnitTestEnv env) throws InterruptedException {
+        try {
+            env.adminClient().describeClientQuotas(ClientQuotaFilter.all()).entities().get();
+            fail("Expected an authentication error.");
+        } catch (ExecutionException e) {
+            assertTrue("Expected an authentication error, but got " + Utils.stackTrace(e),
+                e.getCause() instanceof AuthenticationException);
+        }
+
+        try {
+            ClientQuotaEntity entity = new ClientQuotaEntity(Collections.singletonMap(ClientQuotaEntity.USER, "user"));
+            ClientQuotaAlteration alteration = new ClientQuotaAlteration(entity, asList(new ClientQuotaAlteration.Op("consumer_byte_rate", 1000.0)));
+            env.adminClient().alterClientQuotas(asList(alteration)).all().get();
+            fail("Expected an authentication error.");
+        } catch (ExecutionException e) {
+            assertTrue("Expected an authentication error, but got " + Utils.stackTrace(e),
+                e.getCause() instanceof AuthenticationException);
+        }
+    }
+
+    private static final AclBinding ACL1 = new AclBinding(new ResourcePattern(ResourceType.TOPIC, "mytopic3", PatternType.LITERAL),
+        new AccessControlEntry("User:ANONYMOUS", "*", AclOperation.DESCRIBE, AclPermissionType.ALLOW));
+    private static final AclBinding ACL2 = new AclBinding(new ResourcePattern(ResourceType.TOPIC, "mytopic4", PatternType.LITERAL),
+        new AccessControlEntry("User:ANONYMOUS", "*", AclOperation.DESCRIBE, AclPermissionType.DENY));
+    private static final AclBindingFilter FILTER1 = new AclBindingFilter(new ResourcePatternFilter(ResourceType.ANY, null, PatternType.LITERAL),
+        new AccessControlEntryFilter("User:ANONYMOUS", null, AclOperation.ANY, AclPermissionType.ANY));
+    private static final AclBindingFilter FILTER2 = new AclBindingFilter(new ResourcePatternFilter(ResourceType.ANY, null, PatternType.LITERAL),
+        new AccessControlEntryFilter("User:bob", null, AclOperation.ANY, AclPermissionType.ANY));
+    private static final AclBindingFilter UNKNOWN_FILTER = new AclBindingFilter(
+        new ResourcePatternFilter(ResourceType.UNKNOWN, null, PatternType.LITERAL),
+        new AccessControlEntryFilter("User:bob", null, AclOperation.ANY, AclPermissionType.ANY));
+
+    @Test
+    public void testDescribeAcls() throws Exception {
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            // Test a call where we get back ACL1 and ACL2.
+            env.kafkaClient().prepareResponse(new DescribeAclsResponse(new DescribeAclsResponseData()
+                .setResources(DescribeAclsResponse.aclsResources(asList(ACL1, ACL2))), ApiKeys.DESCRIBE_ACLS.latestVersion()));
+            assertCollectionIs(env.adminClient().describeAcls(FILTER1).values().get(), ACL1, ACL2);
+
+            // Test a call where we get back no results.
+            env.kafkaClient().prepareResponse(new DescribeAclsResponse(new DescribeAclsResponseData(),
+                    ApiKeys.DESCRIBE_ACLS.latestVersion()));
+            assertTrue(env.adminClient().describeAcls(FILTER2).values().get().isEmpty());
+
+            // Test a call where we get back an error.
+            env.kafkaClient().prepareResponse(new DescribeAclsResponse(new DescribeAclsResponseData()
+                .setErrorCode(Errors.SECURITY_DISABLED.code())
+                .setErrorMessage("Security is disabled"), ApiKeys.DESCRIBE_ACLS.latestVersion()));
+            TestUtils.assertFutureError(env.adminClient().describeAcls(FILTER2).values(), SecurityDisabledException.class);
+
+            // Test a call where we supply an invalid filter.
+            TestUtils.assertFutureError(env.adminClient().describeAcls(UNKNOWN_FILTER).values(),
+                InvalidRequestException.class);
+        }
+    }
+
+    @Test
+    public void testCreateAcls() throws Exception {
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            // Test a call where we successfully create two ACLs.
+            env.kafkaClient().prepareResponse(new CreateAclsResponse(new CreateAclsResponseData().setResults(asList(
+                new CreateAclsResponseData.AclCreationResult(),
+                new CreateAclsResponseData.AclCreationResult()))));
+            CreateAclsResult results = env.adminClient().createAcls(asList(ACL1, ACL2));
+            assertCollectionIs(results.values().keySet(), ACL1, ACL2);
+            for (KafkaFuture future : results.values().values())
+                future.get();
+            results.all().get();
+
+            // Test a call where we fail to create one ACL.
+            env.kafkaClient().prepareResponse(new CreateAclsResponse(new CreateAclsResponseData().setResults(asList(
+                new CreateAclsResponseData.AclCreationResult()
+                    .setErrorCode(Errors.SECURITY_DISABLED.code())
+                    .setErrorMessage("Security is disabled"),
+                new CreateAclsResponseData.AclCreationResult()))));
+            results = env.adminClient().createAcls(asList(ACL1, ACL2));
+            assertCollectionIs(results.values().keySet(), ACL1, ACL2);
+            TestUtils.assertFutureError(results.values().get(ACL1), SecurityDisabledException.class);
+            results.values().get(ACL2).get();
+            TestUtils.assertFutureError(results.all(), SecurityDisabledException.class);
+        }
+    }
+
+    @Test
+    public void testDeleteAcls() throws Exception {
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            // Test a call where one filter has an error.
+            env.kafkaClient().prepareResponse(new DeleteAclsResponse(new DeleteAclsResponseData()
+                .setThrottleTimeMs(0)
+                .setFilterResults(asList(
+                    new DeleteAclsResponseData.DeleteAclsFilterResult()
+                        .setMatchingAcls(asList(
+                            DeleteAclsResponse.matchingAcl(ACL1, ApiError.NONE),
+                            DeleteAclsResponse.matchingAcl(ACL2, ApiError.NONE))),
+                    new DeleteAclsResponseData.DeleteAclsFilterResult()
+                        .setErrorCode(Errors.SECURITY_DISABLED.code())
+                        .setErrorMessage("No security"))),
+                    ApiKeys.DELETE_ACLS.latestVersion()));
+            DeleteAclsResult results = env.adminClient().deleteAcls(asList(FILTER1, FILTER2));
+            Map> filterResults = results.values();
+            FilterResults filter1Results = filterResults.get(FILTER1).get();
+            assertEquals(null, filter1Results.values().get(0).exception());
+            assertEquals(ACL1, filter1Results.values().get(0).binding());
+            assertEquals(null, filter1Results.values().get(1).exception());
+            assertEquals(ACL2, filter1Results.values().get(1).binding());
+            TestUtils.assertFutureError(filterResults.get(FILTER2), SecurityDisabledException.class);
+            TestUtils.assertFutureError(results.all(), SecurityDisabledException.class);
+
+            // Test a call where one deletion result has an error.
+            env.kafkaClient().prepareResponse(new DeleteAclsResponse(new DeleteAclsResponseData()
+                .setThrottleTimeMs(0)
+                .setFilterResults(asList(
+                    new DeleteAclsResponseData.DeleteAclsFilterResult()
+                        .setMatchingAcls(asList(
+                            DeleteAclsResponse.matchingAcl(ACL1, ApiError.NONE),
+                            new DeleteAclsResponseData.DeleteAclsMatchingAcl()
+                                .setErrorCode(Errors.SECURITY_DISABLED.code())
+                                .setErrorMessage("No security")
+                                .setPermissionType(AclPermissionType.ALLOW.code())
+                                .setOperation(AclOperation.ALTER.code())
+                                .setResourceType(ResourceType.CLUSTER.code())
+                                .setPatternType(FILTER2.patternFilter().patternType().code()))),
+                    new DeleteAclsResponseData.DeleteAclsFilterResult())),
+                    ApiKeys.DELETE_ACLS.latestVersion()));
+            results = env.adminClient().deleteAcls(asList(FILTER1, FILTER2));
+            assertTrue(results.values().get(FILTER2).get().values().isEmpty());
+            TestUtils.assertFutureError(results.all(), SecurityDisabledException.class);
+
+            // Test a call where there are no errors.
+            env.kafkaClient().prepareResponse(new DeleteAclsResponse(new DeleteAclsResponseData()
+                .setThrottleTimeMs(0)
+                .setFilterResults(asList(
+                    new DeleteAclsResponseData.DeleteAclsFilterResult()
+                        .setMatchingAcls(asList(DeleteAclsResponse.matchingAcl(ACL1, ApiError.NONE))),
+                    new DeleteAclsResponseData.DeleteAclsFilterResult()
+                        .setMatchingAcls(asList(DeleteAclsResponse.matchingAcl(ACL2, ApiError.NONE))))),
+                    ApiKeys.DELETE_ACLS.latestVersion()));
+            results = env.adminClient().deleteAcls(asList(FILTER1, FILTER2));
+            Collection deleted = results.all().get();
+            assertCollectionIs(deleted, ACL1, ACL2);
+        }
+    }
+
+    @Test
+    public void testElectLeaders()  throws Exception {
+        TopicPartition topic1 = new TopicPartition("topic", 0);
+        TopicPartition topic2 = new TopicPartition("topic", 2);
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
+            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, ApiKeys.ELECT_LEADERS.latestVersion()));
+                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,
+                        ApiKeys.ELECT_LEADERS.latestVersion()));
+                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);
+            }
+        }
+    }
+
+    @Test
+    public void testDescribeBrokerConfigs() throws Exception {
+        ConfigResource broker0Resource = new ConfigResource(ConfigResource.Type.BROKER, "0");
+        ConfigResource broker1Resource = new ConfigResource(ConfigResource.Type.BROKER, "1");
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+            env.kafkaClient().prepareResponseFrom(new DescribeConfigsResponse(
+                    new DescribeConfigsResponseData().setResults(asList(new DescribeConfigsResponseData.DescribeConfigsResult()
+                            .setResourceName(broker0Resource.name()).setResourceType(broker0Resource.type().id()).setErrorCode(Errors.NONE.code())
+                            .setConfigs(emptyList())))), env.cluster().nodeById(0));
+            env.kafkaClient().prepareResponseFrom(new DescribeConfigsResponse(
+                    new DescribeConfigsResponseData().setResults(asList(new DescribeConfigsResponseData.DescribeConfigsResult()
+                            .setResourceName(broker1Resource.name()).setResourceType(broker1Resource.type().id()).setErrorCode(Errors.NONE.code())
+                            .setConfigs(emptyList())))), env.cluster().nodeById(1));
+            Map> result = env.adminClient().describeConfigs(asList(
+                    broker0Resource,
+                    broker1Resource)).values();
+            assertEquals(new HashSet<>(asList(broker0Resource, broker1Resource)), result.keySet());
+            result.get(broker0Resource).get();
+            result.get(broker1Resource).get();
+        }
+    }
+
+    @Test
+    public void testDescribeBrokerAndLogConfigs() throws Exception {
+        ConfigResource brokerResource = new ConfigResource(ConfigResource.Type.BROKER, "0");
+        ConfigResource brokerLoggerResource = new ConfigResource(ConfigResource.Type.BROKER_LOGGER, "0");
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+            env.kafkaClient().prepareResponseFrom(new DescribeConfigsResponse(
+                new DescribeConfigsResponseData().setResults(asList(new DescribeConfigsResponseData.DescribeConfigsResult()
+                    .setResourceName(brokerResource.name()).setResourceType(brokerResource.type().id()).setErrorCode(Errors.NONE.code())
+                    .setConfigs(emptyList()),
+                new DescribeConfigsResponseData.DescribeConfigsResult()
+                    .setResourceName(brokerLoggerResource.name()).setResourceType(brokerLoggerResource.type().id()).setErrorCode(Errors.NONE.code())
+                    .setConfigs(emptyList())))), env.cluster().nodeById(0));
+            Map> result = env.adminClient().describeConfigs(asList(
+                    brokerResource,
+                    brokerLoggerResource)).values();
+            assertEquals(new HashSet<>(asList(brokerResource, brokerLoggerResource)), result.keySet());
+            result.get(brokerResource).get();
+            result.get(brokerLoggerResource).get();
+        }
+    }
+
+    @Test
+    public void testDescribeConfigsPartialResponse() throws Exception {
+        ConfigResource topic = new ConfigResource(ConfigResource.Type.TOPIC, "topic");
+        ConfigResource topic2 = new ConfigResource(ConfigResource.Type.TOPIC, "topic2");
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+            env.kafkaClient().prepareResponse(new DescribeConfigsResponse(
+                    new DescribeConfigsResponseData().setResults(asList(new DescribeConfigsResponseData.DescribeConfigsResult()
+                                    .setResourceName(topic.name()).setResourceType(topic.type().id()).setErrorCode(Errors.NONE.code())
+                                    .setConfigs(emptyList())))));
+            Map> result = env.adminClient().describeConfigs(asList(
+                    topic,
+                    topic2)).values();
+            assertEquals(new HashSet<>(asList(topic, topic2)), result.keySet());
+            result.get(topic);
+            TestUtils.assertFutureThrows(result.get(topic2), ApiException.class);
+        }
+    }
+
+    @Test
+    public void testDescribeConfigsUnrequested() throws Exception {
+        ConfigResource topic = new ConfigResource(ConfigResource.Type.TOPIC, "topic");
+        ConfigResource unrequested = new ConfigResource(ConfigResource.Type.TOPIC, "unrequested");
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+            env.kafkaClient().prepareResponse(new DescribeConfigsResponse(
+                new DescribeConfigsResponseData().setResults(asList(new DescribeConfigsResponseData.DescribeConfigsResult()
+                        .setResourceName(topic.name()).setResourceType(topic.type().id()).setErrorCode(Errors.NONE.code())
+                        .setConfigs(emptyList()),
+                new DescribeConfigsResponseData.DescribeConfigsResult()
+                        .setResourceName(unrequested.name()).setResourceType(unrequested.type().id()).setErrorCode(Errors.NONE.code())
+                        .setConfigs(emptyList())))));
+            Map> result = env.adminClient().describeConfigs(asList(
+                    topic)).values();
+            assertEquals(new HashSet<>(asList(topic)), result.keySet());
+            assertNotNull(result.get(topic).get());
+            assertNull(result.get(unrequested));
+        }
+    }
+
+    private static DescribeLogDirsResponse prepareDescribeLogDirsResponse(Errors error, String logDir, TopicPartition tp, long partitionSize, long offsetLag) {
+        return prepareDescribeLogDirsResponse(error, logDir,
+                prepareDescribeLogDirsTopics(partitionSize, offsetLag, tp.topic(), tp.partition(), false));
+    }
+
+    private static List prepareDescribeLogDirsTopics(
+            long partitionSize, long offsetLag, String topic, int partition, boolean isFuture) {
+        return singletonList(new DescribeLogDirsTopic()
+                .setName(topic)
+                .setPartitions(singletonList(new DescribeLogDirsResponseData.DescribeLogDirsPartition()
+                        .setPartitionIndex(partition)
+                        .setPartitionSize(partitionSize)
+                        .setIsFutureKey(isFuture)
+                        .setOffsetLag(offsetLag))));
+    }
+
+    private static DescribeLogDirsResponse prepareDescribeLogDirsResponse(Errors error, String logDir,
+                                                                   List topics) {
+        return new DescribeLogDirsResponse(
+                new DescribeLogDirsResponseData().setResults(singletonList(new DescribeLogDirsResponseData.DescribeLogDirsResult()
+                        .setErrorCode(error.code())
+                        .setLogDir(logDir)
+                        .setTopics(topics)
+                )));
+    }
+
+    @Test
+    public void testDescribeLogDirs() throws ExecutionException, InterruptedException {
+        Set brokers = Collections.singleton(0);
+        String logDir = "/var/data/kafka";
+        TopicPartition tp = new TopicPartition("topic", 12);
+        long partitionSize = 1234567890;
+        long offsetLag = 24;
+
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+            env.kafkaClient().prepareResponseFrom(
+                    prepareDescribeLogDirsResponse(Errors.NONE, logDir, tp, partitionSize, offsetLag),
+                    env.cluster().nodeById(0));
+
+            DescribeLogDirsResult result = env.adminClient().describeLogDirs(brokers);
+
+            Map>> descriptions = result.descriptions();
+            assertEquals(brokers, descriptions.keySet());
+            assertNotNull(descriptions.get(0));
+            assertDescriptionContains(descriptions.get(0).get(), logDir, tp, partitionSize, offsetLag);
+
+            Map> allDescriptions = result.allDescriptions().get();
+            assertEquals(brokers, allDescriptions.keySet());
+            assertDescriptionContains(allDescriptions.get(0), logDir, tp, partitionSize, offsetLag);
+        }
+    }
+
+    private static void assertDescriptionContains(Map descriptionsMap, String logDir,
+                                           TopicPartition tp, long partitionSize, long offsetLag) {
+        assertNotNull(descriptionsMap);
+        assertEquals(Collections.singleton(logDir), descriptionsMap.keySet());
+        assertNull(descriptionsMap.get(logDir).error());
+        Map descriptionsReplicaInfos = descriptionsMap.get(logDir).replicaInfos();
+        assertEquals(Collections.singleton(tp), descriptionsReplicaInfos.keySet());
+        assertEquals(partitionSize, descriptionsReplicaInfos.get(tp).size());
+        assertEquals(offsetLag, descriptionsReplicaInfos.get(tp).offsetLag());
+        assertFalse(descriptionsReplicaInfos.get(tp).isFuture());
+    }
+
+    @SuppressWarnings("deprecation")
+    @Test
+    public void testDescribeLogDirsDeprecated() throws ExecutionException, InterruptedException {
+        Set brokers = Collections.singleton(0);
+        TopicPartition tp = new TopicPartition("topic", 12);
+        String logDir = "/var/data/kafka";
+        Errors error = Errors.NONE;
+        int offsetLag = 24;
+        long partitionSize = 1234567890;
+
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+            env.kafkaClient().prepareResponseFrom(
+                    prepareDescribeLogDirsResponse(error, logDir, tp, partitionSize, offsetLag),
+                    env.cluster().nodeById(0));
+
+            DescribeLogDirsResult result = env.adminClient().describeLogDirs(brokers);
+
+            Map>> deprecatedValues = result.values();
+            assertEquals(brokers, deprecatedValues.keySet());
+            assertNotNull(deprecatedValues.get(0));
+            assertDescriptionContains(deprecatedValues.get(0).get(), logDir, tp, error, offsetLag, partitionSize);
+
+            Map> deprecatedAll = result.all().get();
+            assertEquals(brokers, deprecatedAll.keySet());
+            assertDescriptionContains(deprecatedAll.get(0), logDir, tp, error, offsetLag, partitionSize);
+        }
+    }
+
+    @SuppressWarnings("deprecation")
+    private static  void assertDescriptionContains(Map descriptionsMap,
+                                           String logDir, TopicPartition tp, Errors error,
+                                           int offsetLag, long partitionSize) {
+        assertNotNull(descriptionsMap);
+        assertEquals(Collections.singleton(logDir), descriptionsMap.keySet());
+        assertEquals(error, descriptionsMap.get(logDir).error);
+        Map allReplicaInfos =
+                descriptionsMap.get(logDir).replicaInfos;
+        assertEquals(Collections.singleton(tp), allReplicaInfos.keySet());
+        assertEquals(partitionSize, allReplicaInfos.get(tp).size);
+        assertEquals(offsetLag, allReplicaInfos.get(tp).offsetLag);
+        assertFalse(allReplicaInfos.get(tp).isFuture);
+    }
+
+    @Test
+    public void testDescribeLogDirsOfflineDir() throws ExecutionException, InterruptedException {
+        Set brokers = Collections.singleton(0);
+        String logDir = "/var/data/kafka";
+        Errors error = Errors.KAFKA_STORAGE_ERROR;
+
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+            env.kafkaClient().prepareResponseFrom(
+                    prepareDescribeLogDirsResponse(error, logDir, emptyList()),
+                    env.cluster().nodeById(0));
+
+            DescribeLogDirsResult result = env.adminClient().describeLogDirs(brokers);
+
+            Map>> descriptions = result.descriptions();
+            assertEquals(brokers, descriptions.keySet());
+            assertNotNull(descriptions.get(0));
+            Map descriptionsMap = descriptions.get(0).get();
+            assertEquals(Collections.singleton(logDir), descriptionsMap.keySet());
+            assertEquals(error.exception().getClass(), descriptionsMap.get(logDir).error().getClass());
+            assertEquals(emptySet(), descriptionsMap.get(logDir).replicaInfos().keySet());
+
+            Map> allDescriptions = result.allDescriptions().get();
+            assertEquals(brokers, allDescriptions.keySet());
+            Map allMap = allDescriptions.get(0);
+            assertNotNull(allMap);
+            assertEquals(Collections.singleton(logDir), allMap.keySet());
+            assertEquals(error.exception().getClass(), allMap.get(logDir).error().getClass());
+            assertEquals(emptySet(), allMap.get(logDir).replicaInfos().keySet());
+        }
+    }
+
+    @SuppressWarnings("deprecation")
+    @Test
+    public void testDescribeLogDirsOfflineDirDeprecated() throws ExecutionException, InterruptedException {
+        Set brokers = Collections.singleton(0);
+        String logDir = "/var/data/kafka";
+        Errors error = Errors.KAFKA_STORAGE_ERROR;
+
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+            env.kafkaClient().prepareResponseFrom(
+                    prepareDescribeLogDirsResponse(error, logDir, emptyList()),
+                    env.cluster().nodeById(0));
+
+            DescribeLogDirsResult result = env.adminClient().describeLogDirs(brokers);
+
+            Map>> deprecatedValues = result.values();
+            assertEquals(brokers, deprecatedValues.keySet());
+            assertNotNull(deprecatedValues.get(0));
+            Map valuesMap = deprecatedValues.get(0).get();
+            assertEquals(Collections.singleton(logDir), valuesMap.keySet());
+            assertEquals(error, valuesMap.get(logDir).error);
+            assertEquals(emptySet(), valuesMap.get(logDir).replicaInfos.keySet());
+
+            Map> deprecatedAll = result.all().get();
+            assertEquals(brokers, deprecatedAll.keySet());
+            Map allMap = deprecatedAll.get(0);
+            assertNotNull(allMap);
+            assertEquals(Collections.singleton(logDir), allMap.keySet());
+            assertEquals(error, allMap.get(logDir).error);
+            assertEquals(emptySet(), allMap.get(logDir).replicaInfos.keySet());
+        }
+    }
+
+    @Test
+    public void testDescribeReplicaLogDirs() throws ExecutionException, InterruptedException {
+        TopicPartitionReplica tpr1 = new TopicPartitionReplica("topic", 12, 1);
+        TopicPartitionReplica tpr2 = new TopicPartitionReplica("topic", 12, 2);
+
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+            String broker1log0 = "/var/data/kafka0";
+            String broker1log1 = "/var/data/kafka1";
+            String broker2log0 = "/var/data/kafka2";
+            int broker1Log0OffsetLag = 24;
+            int broker1Log0PartitionSize = 987654321;
+            int broker1Log1PartitionSize = 123456789;
+            int broker1Log1OffsetLag = 4321;
+            env.kafkaClient().prepareResponseFrom(
+                    new DescribeLogDirsResponse(
+                            new DescribeLogDirsResponseData().setResults(asList(
+                                    prepareDescribeLogDirsResult(tpr1, broker1log0, broker1Log0PartitionSize, broker1Log0OffsetLag, false),
+                                    prepareDescribeLogDirsResult(tpr1, broker1log1, broker1Log1PartitionSize, broker1Log1OffsetLag, true)))),
+                    env.cluster().nodeById(tpr1.brokerId()));
+            env.kafkaClient().prepareResponseFrom(
+                    prepareDescribeLogDirsResponse(Errors.KAFKA_STORAGE_ERROR, broker2log0),
+                    env.cluster().nodeById(tpr2.brokerId()));
+
+            DescribeReplicaLogDirsResult result = env.adminClient().describeReplicaLogDirs(asList(tpr1, tpr2));
+
+            Map> values = result.values();
+            assertEquals(TestUtils.toSet(asList(tpr1, tpr2)), values.keySet());
+
+            assertNotNull(values.get(tpr1));
+            assertEquals(broker1log0, values.get(tpr1).get().getCurrentReplicaLogDir());
+            assertEquals(broker1Log0OffsetLag, values.get(tpr1).get().getCurrentReplicaOffsetLag());
+            assertEquals(broker1log1, values.get(tpr1).get().getFutureReplicaLogDir());
+            assertEquals(broker1Log1OffsetLag, values.get(tpr1).get().getFutureReplicaOffsetLag());
+
+            assertNotNull(values.get(tpr2));
+            assertNull(values.get(tpr2).get().getCurrentReplicaLogDir());
+            assertEquals(-1, values.get(tpr2).get().getCurrentReplicaOffsetLag());
+            assertNull(values.get(tpr2).get().getFutureReplicaLogDir());
+            assertEquals(-1, values.get(tpr2).get().getFutureReplicaOffsetLag());
+        }
+    }
+
+    private static DescribeLogDirsResponseData.DescribeLogDirsResult prepareDescribeLogDirsResult(TopicPartitionReplica tpr, String logDir, int partitionSize, int offsetLag, boolean isFuture) {
+        return new DescribeLogDirsResponseData.DescribeLogDirsResult()
+                .setErrorCode(Errors.NONE.code())
+                .setLogDir(logDir)
+                .setTopics(prepareDescribeLogDirsTopics(partitionSize, offsetLag, tpr.topic(), tpr.partition(), isFuture));
+    }
+
+    @Test
+    public void testDescribeReplicaLogDirsUnexpected() throws ExecutionException, InterruptedException {
+        TopicPartitionReplica expected = new TopicPartitionReplica("topic", 12, 1);
+        TopicPartitionReplica unexpected = new TopicPartitionReplica("topic", 12, 2);
+
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+            String broker1log0 = "/var/data/kafka0";
+            String broker1log1 = "/var/data/kafka1";
+            int broker1Log0PartitionSize = 987654321;
+            int broker1Log0OffsetLag = 24;
+            int broker1Log1PartitionSize = 123456789;
+            int broker1Log1OffsetLag = 4321;
+            env.kafkaClient().prepareResponseFrom(
+                    new DescribeLogDirsResponse(
+                            new DescribeLogDirsResponseData().setResults(asList(
+                                    prepareDescribeLogDirsResult(expected, broker1log0, broker1Log0PartitionSize, broker1Log0OffsetLag, false),
+                                    prepareDescribeLogDirsResult(unexpected, broker1log1, broker1Log1PartitionSize, broker1Log1OffsetLag, true)))),
+                    env.cluster().nodeById(expected.brokerId()));
+
+            DescribeReplicaLogDirsResult result = env.adminClient().describeReplicaLogDirs(asList(expected));
+
+            Map> values = result.values();
+            assertEquals(TestUtils.toSet(asList(expected)), values.keySet());
+
+            assertNotNull(values.get(expected));
+            assertEquals(broker1log0, values.get(expected).get().getCurrentReplicaLogDir());
+            assertEquals(broker1Log0OffsetLag, values.get(expected).get().getCurrentReplicaOffsetLag());
+            assertEquals(broker1log1, values.get(expected).get().getFutureReplicaLogDir());
+            assertEquals(broker1Log1OffsetLag, values.get(expected).get().getFutureReplicaOffsetLag());
+        }
+    }
+
+    @Test
+    public void testCreatePartitions() throws Exception {
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            // Test a call where one filter has an error.
+            env.kafkaClient().prepareResponse(
+                expectCreatePartitionsRequestWithTopics("my_topic", "other_topic"),
+                prepareCreatePartitionsResponse(1000,
+                    createPartitionsTopicResult("my_topic", Errors.NONE),
+                    createPartitionsTopicResult("other_topic", Errors.INVALID_TOPIC_EXCEPTION,
+                        "some detailed reason")));
+
+
+            Map counts = new HashMap<>();
+            counts.put("my_topic", NewPartitions.increaseTo(3));
+            counts.put("other_topic", NewPartitions.increaseTo(3, asList(asList(2), asList(3))));
+
+            CreatePartitionsResult results = env.adminClient().createPartitions(counts);
+            Map> values = results.values();
+            KafkaFuture myTopicResult = values.get("my_topic");
+            myTopicResult.get();
+            KafkaFuture otherTopicResult = values.get("other_topic");
+            try {
+                otherTopicResult.get();
+                fail("get() should throw ExecutionException");
+            } catch (ExecutionException e0) {
+                assertTrue(e0.getCause() instanceof InvalidTopicException);
+                InvalidTopicException e = (InvalidTopicException) e0.getCause();
+                assertEquals("some detailed reason", e.getMessage());
+            }
+        }
+    }
+
+    @Test
+    public void testCreatePartitionsRetryThrottlingExceptionWhenEnabled() throws Exception {
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            env.kafkaClient().prepareResponse(
+                expectCreatePartitionsRequestWithTopics("topic1", "topic2", "topic3"),
+                prepareCreatePartitionsResponse(1000,
+                    createPartitionsTopicResult("topic1", Errors.NONE),
+                    createPartitionsTopicResult("topic2", Errors.THROTTLING_QUOTA_EXCEEDED),
+                    createPartitionsTopicResult("topic3", Errors.TOPIC_ALREADY_EXISTS)));
+
+            env.kafkaClient().prepareResponse(
+                expectCreatePartitionsRequestWithTopics("topic2"),
+                prepareCreatePartitionsResponse(1000,
+                    createPartitionsTopicResult("topic2", Errors.THROTTLING_QUOTA_EXCEEDED)));
+
+            env.kafkaClient().prepareResponse(
+                expectCreatePartitionsRequestWithTopics("topic2"),
+                prepareCreatePartitionsResponse(0,
+                    createPartitionsTopicResult("topic2", Errors.NONE)));
+
+            Map counts = new HashMap<>();
+            counts.put("topic1", NewPartitions.increaseTo(1));
+            counts.put("topic2", NewPartitions.increaseTo(2));
+            counts.put("topic3", NewPartitions.increaseTo(3));
+
+            CreatePartitionsResult result = env.adminClient().createPartitions(
+                counts, new CreatePartitionsOptions().retryOnQuotaViolation(true));
+
+            assertNull(result.values().get("topic1").get());
+            assertNull(result.values().get("topic2").get());
+            TestUtils.assertFutureThrows(result.values().get("topic3"), TopicExistsException.class);
+        }
+    }
+
+    @Test
+    public void testCreatePartitionsRetryThrottlingExceptionWhenEnabledUntilRequestTimeOut() throws Exception {
+        long defaultApiTimeout = 60000;
+        MockTime time = new MockTime();
+
+        try (AdminClientUnitTestEnv env = mockClientEnv(time,
+            AdminClientConfig.DEFAULT_API_TIMEOUT_MS_CONFIG, String.valueOf(defaultApiTimeout))) {
+
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            env.kafkaClient().prepareResponse(
+                expectCreatePartitionsRequestWithTopics("topic1", "topic2", "topic3"),
+                prepareCreatePartitionsResponse(1000,
+                    createPartitionsTopicResult("topic1", Errors.NONE),
+                    createPartitionsTopicResult("topic2", Errors.THROTTLING_QUOTA_EXCEEDED),
+                    createPartitionsTopicResult("topic3", Errors.TOPIC_ALREADY_EXISTS)));
+
+            env.kafkaClient().prepareResponse(
+                expectCreatePartitionsRequestWithTopics("topic2"),
+                prepareCreatePartitionsResponse(1000,
+                    createPartitionsTopicResult("topic2", Errors.THROTTLING_QUOTA_EXCEEDED)));
+
+            Map counts = new HashMap<>();
+            counts.put("topic1", NewPartitions.increaseTo(1));
+            counts.put("topic2", NewPartitions.increaseTo(2));
+            counts.put("topic3", NewPartitions.increaseTo(3));
+
+            CreatePartitionsResult result = env.adminClient().createPartitions(
+                counts, new CreatePartitionsOptions().retryOnQuotaViolation(true));
+
+            // Wait until the prepared attempts have consumed
+            TestUtils.waitForCondition(() -> env.kafkaClient().numAwaitingResponses() == 0,
+                "Failed awaiting CreatePartitions requests");
+
+            // Wait until the next request is sent out
+            TestUtils.waitForCondition(() -> env.kafkaClient().inFlightRequestCount() == 1,
+                "Failed awaiting next CreatePartitions request");
+
+            // Advance time past the default api timeout to time out the inflight request
+            time.sleep(defaultApiTimeout + 1);
+
+            assertNull(result.values().get("topic1").get());
+            ThrottlingQuotaExceededException e = TestUtils.assertFutureThrows(result.values().get("topic2"),
+                ThrottlingQuotaExceededException.class);
+            assertEquals(0, e.throttleTimeMs());
+            TestUtils.assertFutureThrows(result.values().get("topic3"), TopicExistsException.class);
+        }
+    }
+
+    @Test
+    public void testCreatePartitionsDontRetryThrottlingExceptionWhenDisabled() throws Exception {
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            env.kafkaClient().prepareResponse(
+                expectCreatePartitionsRequestWithTopics("topic1", "topic2", "topic3"),
+                prepareCreatePartitionsResponse(1000,
+                    createPartitionsTopicResult("topic1", Errors.NONE),
+                    createPartitionsTopicResult("topic2", Errors.THROTTLING_QUOTA_EXCEEDED),
+                    createPartitionsTopicResult("topic3", Errors.TOPIC_ALREADY_EXISTS)));
+
+            Map counts = new HashMap<>();
+            counts.put("topic1", NewPartitions.increaseTo(1));
+            counts.put("topic2", NewPartitions.increaseTo(2));
+            counts.put("topic3", NewPartitions.increaseTo(3));
+
+            CreatePartitionsResult result = env.adminClient().createPartitions(
+                counts, new CreatePartitionsOptions().retryOnQuotaViolation(false));
+
+            assertNull(result.values().get("topic1").get());
+            ThrottlingQuotaExceededException e = TestUtils.assertFutureThrows(result.values().get("topic2"),
+                ThrottlingQuotaExceededException.class);
+            assertEquals(1000, e.throttleTimeMs());
+            TestUtils.assertFutureThrows(result.values().get("topic3"), TopicExistsException.class);
+        }
+    }
+
+    private MockClient.RequestMatcher expectCreatePartitionsRequestWithTopics(final String... topics) {
+        return body -> {
+            if (body instanceof CreatePartitionsRequest) {
+                CreatePartitionsRequest request = (CreatePartitionsRequest) body;
+                for (String topic : topics) {
+                    if (request.data().topics().find(topic) == null)
+                        return false;
+                }
+                return topics.length == request.data().topics().size();
+            }
+            return false;
+        };
+    }
+
+    @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(RequestTestUtils.metadataResponse(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);
+
+        MockTime time = new MockTime();
+
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(time, mockCluster(3, 0))) {
+            List nodes = env.cluster().nodes();
+
+            List partitionMetadata = new ArrayList<>();
+            partitionMetadata.add(new MetadataResponse.PartitionMetadata(Errors.NONE, tp0,
+                    Optional.of(nodes.get(0).id()), Optional.of(5), singletonList(nodes.get(0).id()),
+                    singletonList(nodes.get(0).id()), Collections.emptyList()));
+            partitionMetadata.add(new MetadataResponse.PartitionMetadata(Errors.NONE, tp1,
+                    Optional.of(nodes.get(1).id()), Optional.of(5), singletonList(nodes.get(1).id()),
+                    singletonList(nodes.get(1).id()), Collections.emptyList()));
+
+            List topicMetadata = new ArrayList<>();
+            topicMetadata.add(new MetadataResponse.TopicMetadata(Errors.NONE, topic, false, partitionMetadata));
+
+            env.kafkaClient().prepareResponse(RequestTestUtils.metadataResponse(env.cluster().nodes(),
+                    env.cluster().clusterResource().clusterId(), env.cluster().controller().id(), topicMetadata));
+
+            env.kafkaClient().prepareResponseFrom(new DeleteRecordsResponse(new DeleteRecordsResponseData().setTopics(
+                    new DeleteRecordsResponseData.DeleteRecordsTopicResultCollection(singletonList(new DeleteRecordsResponseData.DeleteRecordsTopicResult()
+                            .setName(tp0.topic())
+                            .setPartitions(new DeleteRecordsResponseData.DeleteRecordsPartitionResultCollection(singletonList(new DeleteRecordsResponseData.DeleteRecordsPartitionResult()
+                                    .setPartitionIndex(tp0.partition())
+                                    .setErrorCode(Errors.NONE.code())
+                                    .setLowWatermark(3)).iterator()))).iterator()))), 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 {
+        HashMap nodes = new HashMap<>();
+        nodes.put(0, new Node(0, "localhost", 8121));
+        List partitionInfos = new ArrayList<>();
+        partitionInfos.add(new PartitionInfo("my_topic", 0, nodes.get(0), new Node[] {nodes.get(0)}, new Node[] {nodes.get(0)}));
+        partitionInfos.add(new PartitionInfo("my_topic", 1, nodes.get(0), new Node[] {nodes.get(0)}, new Node[] {nodes.get(0)}));
+        partitionInfos.add(new PartitionInfo("my_topic", 2, null, new Node[] {nodes.get(0)}, new Node[] {nodes.get(0)}));
+        partitionInfos.add(new PartitionInfo("my_topic", 3, nodes.get(0), new Node[] {nodes.get(0)}, new Node[] {nodes.get(0)}));
+        partitionInfos.add(new PartitionInfo("my_topic", 4, nodes.get(0), new Node[] {nodes.get(0)}, new Node[] {nodes.get(0)}));
+        Cluster cluster = new Cluster("mockClusterId", nodes.values(),
+                partitionInfos, Collections.emptySet(),
+                Collections.emptySet(), nodes.get(0));
+
+        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);
+        TopicPartition myTopicPartition4 = new TopicPartition("my_topic", 4);
+
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(cluster)) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            DeleteRecordsResponseData m = new DeleteRecordsResponseData();
+            m.topics().add(new DeleteRecordsResponseData.DeleteRecordsTopicResult().setName(myTopicPartition0.topic())
+                    .setPartitions(new DeleteRecordsResponseData.DeleteRecordsPartitionResultCollection(asList(
+                        new DeleteRecordsResponseData.DeleteRecordsPartitionResult()
+                            .setPartitionIndex(myTopicPartition0.partition())
+                            .setLowWatermark(3)
+                            .setErrorCode(Errors.NONE.code()),
+                        new DeleteRecordsResponseData.DeleteRecordsPartitionResult()
+                            .setPartitionIndex(myTopicPartition1.partition())
+                            .setLowWatermark(DeleteRecordsResponse.INVALID_LOW_WATERMARK)
+                            .setErrorCode(Errors.OFFSET_OUT_OF_RANGE.code()),
+                        new DeleteRecordsResponseData.DeleteRecordsPartitionResult()
+                            .setPartitionIndex(myTopicPartition3.partition())
+                            .setLowWatermark(DeleteRecordsResponse.INVALID_LOW_WATERMARK)
+                            .setErrorCode(Errors.NOT_LEADER_OR_FOLLOWER.code()),
+                        new DeleteRecordsResponseData.DeleteRecordsPartitionResult()
+                            .setPartitionIndex(myTopicPartition4.partition())
+                            .setLowWatermark(DeleteRecordsResponse.INVALID_LOW_WATERMARK)
+                            .setErrorCode(Errors.UNKNOWN_TOPIC_OR_PARTITION.code())
+                    ).iterator())));
+
+            List t = new ArrayList<>();
+            List p = new ArrayList<>();
+            p.add(new MetadataResponse.PartitionMetadata(Errors.NONE, myTopicPartition0,
+                    Optional.of(nodes.get(0).id()), Optional.of(5), singletonList(nodes.get(0).id()),
+                    singletonList(nodes.get(0).id()), Collections.emptyList()));
+            p.add(new MetadataResponse.PartitionMetadata(Errors.NONE, myTopicPartition1,
+                    Optional.of(nodes.get(0).id()), Optional.of(5), singletonList(nodes.get(0).id()),
+                    singletonList(nodes.get(0).id()), Collections.emptyList()));
+            p.add(new MetadataResponse.PartitionMetadata(Errors.LEADER_NOT_AVAILABLE, myTopicPartition2,
+                    Optional.empty(), Optional.empty(), singletonList(nodes.get(0).id()),
+                    singletonList(nodes.get(0).id()), Collections.emptyList()));
+            p.add(new MetadataResponse.PartitionMetadata(Errors.NONE, myTopicPartition3,
+                    Optional.of(nodes.get(0).id()), Optional.of(5), singletonList(nodes.get(0).id()),
+                    singletonList(nodes.get(0).id()), Collections.emptyList()));
+            p.add(new MetadataResponse.PartitionMetadata(Errors.NONE, myTopicPartition4,
+                    Optional.of(nodes.get(0).id()), Optional.of(5), singletonList(nodes.get(0).id()),
+                    singletonList(nodes.get(0).id()), Collections.emptyList()));
+
+            t.add(new MetadataResponse.TopicMetadata(Errors.NONE, "my_topic", false, p));
+
+            env.kafkaClient().prepareResponse(RequestTestUtils.metadataResponse(cluster.nodes(), cluster.clusterResource().clusterId(), cluster.controller().id(), t));
+            env.kafkaClient().prepareResponse(new DeleteRecordsResponse(m));
+
+            Map recordsToDelete = new HashMap<>();
+            recordsToDelete.put(myTopicPartition0, RecordsToDelete.beforeOffset(3L));
+            recordsToDelete.put(myTopicPartition1, RecordsToDelete.beforeOffset(10L));
+            recordsToDelete.put(myTopicPartition2, RecordsToDelete.beforeOffset(10L));
+            recordsToDelete.put(myTopicPartition3, RecordsToDelete.beforeOffset(10L));
+            recordsToDelete.put(myTopicPartition4, RecordsToDelete.beforeOffset(10L));
+
+            DeleteRecordsResult results = env.adminClient().deleteRecords(recordsToDelete);
+
+            // success on records deletion for partition 0
+            Map> values = results.lowWatermarks();
+            KafkaFuture myTopicPartition0Result = values.get(myTopicPartition0);
+            long lowWatermark = myTopicPartition0Result.get().lowWatermark();
+            assertEquals(lowWatermark, 3);
+
+            // "offset out of range" failure on records deletion for partition 1
+            KafkaFuture myTopicPartition1Result = values.get(myTopicPartition1);
+            try {
+                myTopicPartition1Result.get();
+                fail("get() should throw ExecutionException");
+            } catch (ExecutionException e0) {
+                assertTrue(e0.getCause() instanceof OffsetOutOfRangeException);
+            }
+
+            // "leader not available" failure on metadata request for partition 2
+            KafkaFuture myTopicPartition2Result = values.get(myTopicPartition2);
+            try {
+                myTopicPartition2Result.get();
+                fail("get() should throw ExecutionException");
+            } catch (ExecutionException e1) {
+                assertTrue(e1.getCause() instanceof LeaderNotAvailableException);
+            }
+
+            // "not leader for partition" failure on records deletion for partition 3
+            KafkaFuture myTopicPartition3Result = values.get(myTopicPartition3);
+            try {
+                myTopicPartition3Result.get();
+                fail("get() should throw ExecutionException");
+            } catch (ExecutionException e1) {
+                assertTrue(e1.getCause() instanceof NotLeaderOrFollowerException);
+            }
+
+            // "unknown topic or partition" failure on records deletion for partition 4
+            KafkaFuture myTopicPartition4Result = values.get(myTopicPartition4);
+            try {
+                myTopicPartition4Result.get();
+                fail("get() should throw ExecutionException");
+            } catch (ExecutionException e1) {
+                assertTrue(e1.getCause() instanceof UnknownTopicOrPartitionException);
+            }
+        }
+    }
+
+    @Test
+    public void testDescribeCluster() throws Exception {
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(mockCluster(4, 0),
+                AdminClientConfig.RETRIES_CONFIG, "2")) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            // Prepare the metadata response used for the first describe cluster
+            MetadataResponse response = RequestTestUtils.metadataResponse(0,
+                    env.cluster().nodes(),
+                    env.cluster().clusterResource().clusterId(),
+                    2,
+                    Collections.emptyList(),
+                    MetadataResponse.AUTHORIZED_OPERATIONS_OMITTED,
+                    ApiKeys.METADATA.latestVersion());
+            env.kafkaClient().prepareResponse(response);
+
+            // Prepare the metadata response used for the second describe cluster
+            MetadataResponse response2 = RequestTestUtils.metadataResponse(0,
+                    env.cluster().nodes(),
+                    env.cluster().clusterResource().clusterId(),
+                    3,
+                    Collections.emptyList(),
+                    1 << AclOperation.DESCRIBE.code() | 1 << AclOperation.ALTER.code(),
+                    ApiKeys.METADATA.latestVersion());
+            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 {
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(mockCluster(4, 0),
+                AdminClientConfig.RETRIES_CONFIG, "2")) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            // Empty metadata response should be retried
+            env.kafkaClient().prepareResponse(
+                     RequestTestUtils.metadataResponse(
+                            Collections.emptyList(),
+                            env.cluster().clusterResource().clusterId(),
+                            -1,
+                            Collections.emptyList()));
+
+            env.kafkaClient().prepareResponse(
+                     RequestTestUtils.metadataResponse(
+                            env.cluster().nodes(),
+                            env.cluster().clusterResource().clusterId(),
+                            env.cluster().controller().id(),
+                            Collections.emptyList()));
+
+            env.kafkaClient().prepareResponseFrom(
+                    new ListGroupsResponse(
+                            new ListGroupsResponseData()
+                            .setErrorCode(Errors.NONE.code())
+                            .setGroups(Arrays.asList(
+                                    new ListGroupsResponseData.ListedGroup()
+                                            .setGroupId("group-1")
+                                            .setProtocolType(ConsumerProtocol.PROTOCOL_TYPE)
+                                            .setGroupState("Stable"),
+                                    new ListGroupsResponseData.ListedGroup()
+                                            .setGroupId("group-connect-1")
+                                            .setProtocolType("connector")
+                                            .setGroupState("Stable")
+                            ))),
+                    env.cluster().nodeById(0));
+
+            // handle retriable errors
+            env.kafkaClient().prepareResponseFrom(
+                    new ListGroupsResponse(
+                            new ListGroupsResponseData()
+                                    .setErrorCode(Errors.COORDINATOR_NOT_AVAILABLE.code())
+                                    .setGroups(Collections.emptyList())
+                    ),
+                    env.cluster().nodeById(1));
+            env.kafkaClient().prepareResponseFrom(
+                    new ListGroupsResponse(
+                            new ListGroupsResponseData()
+                                    .setErrorCode(Errors.COORDINATOR_LOAD_IN_PROGRESS.code())
+                                    .setGroups(Collections.emptyList())
+                    ),
+                    env.cluster().nodeById(1));
+            env.kafkaClient().prepareResponseFrom(
+                    new ListGroupsResponse(
+                            new ListGroupsResponseData()
+                                    .setErrorCode(Errors.NONE.code())
+                                    .setGroups(Arrays.asList(
+                                            new ListGroupsResponseData.ListedGroup()
+                                                    .setGroupId("group-2")
+                                                    .setProtocolType(ConsumerProtocol.PROTOCOL_TYPE)
+                                                    .setGroupState("Stable"),
+                                            new ListGroupsResponseData.ListedGroup()
+                                                    .setGroupId("group-connect-2")
+                                                    .setProtocolType("connector")
+                                                    .setGroupState("Stable")
+                            ))),
+                    env.cluster().nodeById(1));
+
+            env.kafkaClient().prepareResponseFrom(
+                    new ListGroupsResponse(
+                            new ListGroupsResponseData()
+                                    .setErrorCode(Errors.NONE.code())
+                                    .setGroups(Arrays.asList(
+                                            new ListGroupsResponseData.ListedGroup()
+                                                    .setGroupId("group-3")
+                                                    .setProtocolType(ConsumerProtocol.PROTOCOL_TYPE)
+                                                    .setGroupState("Stable"),
+                                            new ListGroupsResponseData.ListedGroup()
+                                                    .setGroupId("group-connect-3")
+                                                    .setProtocolType("connector")
+                                                    .setGroupState("Stable")
+                                    ))),
+                    env.cluster().nodeById(2));
+
+            // fatal error
+            env.kafkaClient().prepareResponseFrom(
+                    new ListGroupsResponse(
+                            new ListGroupsResponseData()
+                                    .setErrorCode(Errors.UNKNOWN_SERVER_ERROR.code())
+                                    .setGroups(Collections.emptyList())),
+                    env.cluster().nodeById(3));
+
+            final ListConsumerGroupsResult result = env.adminClient().listConsumerGroups();
+            TestUtils.assertFutureError(result.all(), UnknownServerException.class);
+
+            Collection listings = result.valid().get();
+            assertEquals(3, listings.size());
+
+            Set groupIds = new HashSet<>();
+            for (ConsumerGroupListing listing : listings) {
+                groupIds.add(listing.groupId());
+                assertTrue(listing.state().isPresent());
+            }
+
+            assertEquals(Utils.mkSet("group-1", "group-2", "group-3"), groupIds);
+            assertEquals(1, result.errors().get().size());
+        }
+    }
+
+    @Test
+    public void testListConsumerGroupsMetadataFailure() throws Exception {
+        final Cluster cluster = mockCluster(3, 0);
+        final Time time = new MockTime();
+
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(time, cluster,
+                AdminClientConfig.RETRIES_CONFIG, "0")) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            // Empty metadata causes the request to fail since we have no list of brokers
+            // to send the ListGroups requests to
+            env.kafkaClient().prepareResponse(
+                     RequestTestUtils.metadataResponse(
+                            Collections.emptyList(),
+                            env.cluster().clusterResource().clusterId(),
+                            -1,
+                            Collections.emptyList()));
+
+            final ListConsumerGroupsResult result = env.adminClient().listConsumerGroups();
+            TestUtils.assertFutureError(result.all(), KafkaException.class);
+        }
+    }
+
+    @Test
+    public void testListConsumerGroupsWithStates() throws Exception {
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(mockCluster(1, 0))) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            env.kafkaClient().prepareResponse(prepareMetadataResponse(env.cluster(), Errors.NONE));
+
+            env.kafkaClient().prepareResponseFrom(
+                new ListGroupsResponse(new ListGroupsResponseData()
+                    .setErrorCode(Errors.NONE.code())
+                    .setGroups(Arrays.asList(
+                            new ListGroupsResponseData.ListedGroup()
+                                .setGroupId("group-1")
+                                .setProtocolType(ConsumerProtocol.PROTOCOL_TYPE)
+                                .setGroupState("Stable"),
+                            new ListGroupsResponseData.ListedGroup()
+                                .setGroupId("group-2")
+                                .setGroupState("Empty")))),
+                env.cluster().nodeById(0));
+
+            final ListConsumerGroupsOptions options = new ListConsumerGroupsOptions();
+            final ListConsumerGroupsResult result = env.adminClient().listConsumerGroups(options);
+            Collection listings = result.valid().get();
+
+            assertEquals(2, listings.size());
+            List expected = new ArrayList<>();
+            expected.add(new ConsumerGroupListing("group-2", true, Optional.of(ConsumerGroupState.EMPTY)));
+            expected.add(new ConsumerGroupListing("group-1", false, Optional.of(ConsumerGroupState.STABLE)));
+            assertEquals(expected, listings);
+            assertEquals(0, result.errors().get().size());
+        }
+    }
+
+    @Test
+    public void testListConsumerGroupsWithStatesOlderBrokerVersion() throws Exception {
+        ApiVersion listGroupV3 = new ApiVersion(ApiKeys.LIST_GROUPS.id, (short) 0, (short) 3);
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(mockCluster(1, 0))) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create(Collections.singletonList(listGroupV3)));
+
+            env.kafkaClient().prepareResponse(prepareMetadataResponse(env.cluster(), Errors.NONE));
+
+            // Check we can list groups with older broker if we don't specify states
+            env.kafkaClient().prepareResponseFrom(
+                    new ListGroupsResponse(new ListGroupsResponseData()
+                        .setErrorCode(Errors.NONE.code())
+                        .setGroups(Collections.singletonList(
+                                new ListGroupsResponseData.ListedGroup()
+                                    .setGroupId("group-1")
+                                    .setProtocolType(ConsumerProtocol.PROTOCOL_TYPE)))),
+                    env.cluster().nodeById(0));
+            ListConsumerGroupsOptions options = new ListConsumerGroupsOptions();
+            ListConsumerGroupsResult result = env.adminClient().listConsumerGroups(options);
+            Collection listing = result.all().get();
+            assertEquals(1, listing.size());
+            List expected = Collections.singletonList(new ConsumerGroupListing("group-1", false, Optional.empty()));
+            assertEquals(expected, listing);
+
+            // But we cannot set a state filter with older broker
+            env.kafkaClient().prepareResponse(prepareMetadataResponse(env.cluster(), Errors.NONE));
+            env.kafkaClient().prepareUnsupportedVersionResponse(
+                body -> body instanceof ListGroupsRequest);
+
+            options = new ListConsumerGroupsOptions().inStates(Collections.singleton(ConsumerGroupState.STABLE));
+            result = env.adminClient().listConsumerGroups(options);
+            TestUtils.assertFutureThrows(result.all(), UnsupportedVersionException.class);
+        }
+    }
+
+    @Test
+    public void testOffsetCommitNumRetries() throws Exception {
+        final Cluster cluster = mockCluster(3, 0);
+        final Time time = new MockTime();
+
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(time, cluster,
+            AdminClientConfig.RETRIES_CONFIG, "0")) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            final String groupId = "group-0";
+            final TopicPartition tp1 = new TopicPartition("foo", 0);
+
+            env.kafkaClient().prepareResponse(prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+            env.kafkaClient().prepareResponse(prepareOffsetCommitResponse(tp1, Errors.NOT_COORDINATOR));
+            env.kafkaClient().prepareResponse(prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+
+            Map offsets = new HashMap<>();
+            offsets.put(tp1, new OffsetAndMetadata(123L));
+            final AlterConsumerGroupOffsetsResult result = env.adminClient().alterConsumerGroupOffsets(groupId, offsets);
+
+            TestUtils.assertFutureError(result.all(), TimeoutException.class);
+        }
+    }
+
+    @Test
+    public void testOffsetCommitRetryBackoff() throws Exception {
+        MockTime time = new MockTime();
+        int retryBackoff = 100;
+
+        try (final AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(time,
+            mockCluster(3, 0),
+            newStrMap(AdminClientConfig.RETRY_BACKOFF_MS_CONFIG, "" + retryBackoff))) {
+            MockClient mockClient = env.kafkaClient();
+
+            mockClient.setNodeApiVersions(NodeApiVersions.create());
+
+            AtomicLong firstAttemptTime = new AtomicLong(0);
+            AtomicLong secondAttemptTime = new AtomicLong(0);
+
+            final String groupId = "group-0";
+            final TopicPartition tp1 = new TopicPartition("foo", 0);
+
+            mockClient.prepareResponse(prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+            mockClient.prepareResponse(body -> {
+                firstAttemptTime.set(time.milliseconds());
+                return true;
+            }, prepareOffsetCommitResponse(tp1, Errors.NOT_COORDINATOR));
+
+
+            mockClient.prepareResponse(prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+            mockClient.prepareResponse(body -> {
+                secondAttemptTime.set(time.milliseconds());
+                return true;
+            }, prepareOffsetCommitResponse(tp1, Errors.NONE));
+
+
+            Map offsets = new HashMap<>();
+            offsets.put(tp1, new OffsetAndMetadata(123L));
+            final KafkaFuture future = env.adminClient().alterConsumerGroupOffsets(groupId, offsets).all();
+
+            TestUtils.waitForCondition(() -> mockClient.numAwaitingResponses() == 1, "Failed awaiting CommitOffsets first request failure");
+            TestUtils.waitForCondition(() -> ((KafkaAdminClient) env.adminClient()).numPendingCalls() == 1, "Failed to add retry CommitOffsets call on first failure");
+            time.sleep(retryBackoff);
+
+            future.get();
+
+            long actualRetryBackoff = secondAttemptTime.get() - firstAttemptTime.get();
+            assertEquals("CommitOffsets retry did not await expected backoff!", retryBackoff, actualRetryBackoff);
+        }
+    }
+
+    @Test
+    public void testDescribeConsumerGroupNumRetries() throws Exception {
+        final Cluster cluster = mockCluster(3, 0);
+        final Time time = new MockTime();
+
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(time, cluster,
+            AdminClientConfig.RETRIES_CONFIG, "0")) {
+            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.NOT_COORDINATOR,
+                "",
+                "",
+                "",
+                Collections.emptyList(),
+                Collections.emptySet()));
+            env.kafkaClient().prepareResponse(new DescribeGroupsResponse(data));
+            env.kafkaClient().prepareResponse(prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+
+            final DescribeConsumerGroupsResult result = env.adminClient().describeConsumerGroups(singletonList("group-0"));
+
+            TestUtils.assertFutureError(result.all(), TimeoutException.class);
+        }
+    }
+
+    @Test
+    public void testDescribeConsumerGroupRetryBackoff() throws Exception {
+        MockTime time = new MockTime();
+        int retryBackoff = 100;
+
+        try (final AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(time,
+            mockCluster(3, 0),
+            newStrMap(AdminClientConfig.RETRY_BACKOFF_MS_CONFIG, "" + retryBackoff))) {
+            MockClient mockClient = env.kafkaClient();
+
+            mockClient.setNodeApiVersions(NodeApiVersions.create());
+
+            AtomicLong firstAttemptTime = new AtomicLong(0);
+            AtomicLong secondAttemptTime = new AtomicLong(0);
+
+            mockClient.prepareResponse(prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+
+            DescribeGroupsResponseData data = new DescribeGroupsResponseData();
+            data.groups().add(DescribeGroupsResponse.groupMetadata(
+                "group-0",
+                Errors.NOT_COORDINATOR,
+                "",
+                "",
+                "",
+                Collections.emptyList(),
+                Collections.emptySet()));
+
+            mockClient.prepareResponse(body -> {
+                firstAttemptTime.set(time.milliseconds());
+                return true;
+            }, new DescribeGroupsResponse(data));
+
+            mockClient.prepareResponse(prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+
+            data = new DescribeGroupsResponseData();
+            data.groups().add(DescribeGroupsResponse.groupMetadata(
+                "group-0",
+                Errors.NONE,
+                "",
+                ConsumerProtocol.PROTOCOL_TYPE,
+                "",
+                Collections.emptyList(),
+                Collections.emptySet()));
+
+            mockClient.prepareResponse(body -> {
+                secondAttemptTime.set(time.milliseconds());
+                return true;
+            }, new DescribeGroupsResponse(data));
+
+            final KafkaFuture> future =
+                env.adminClient().describeConsumerGroups(singletonList("group-0")).all();
+
+            TestUtils.waitForCondition(() -> mockClient.numAwaitingResponses() == 1, "Failed awaiting DescribeConsumerGroup first request failure");
+            TestUtils.waitForCondition(() -> ((KafkaAdminClient) env.adminClient()).numPendingCalls() == 1, "Failed to add retry DescribeConsumerGroup call on first failure");
+            time.sleep(retryBackoff);
+
+            future.get();
+
+            long actualRetryBackoff = secondAttemptTime.get() - firstAttemptTime.get();
+            assertEquals("DescribeConsumerGroup retry did not await expected backoff!", retryBackoff, actualRetryBackoff);
+        }
+    }
+
+
+    @Test
+    public void testDescribeConsumerGroups() throws Exception {
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(mockCluster(1, 0))) {
+            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()));
+
+            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);
+
+            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);
+
+            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() {
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(mockCluster(1, 0))) {
+            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", null, "clientId0", "clientHost", memberAssignmentBytes, null),
+                            DescribeGroupsResponse.groupMember("1", null, "clientId1", "clientHost", memberAssignmentBytes, null)
+                    ),
+                    Collections.emptySet()));
+
+            DescribeGroupsResponseData groupConnectData = new DescribeGroupsResponseData();
+            group0Data.groups().add(DescribeGroupsResponse.groupMetadata(
+                    "group-connect-0",
+                    Errors.NONE,
+                    "",
+                    "connect",
+                    "",
+                    asList(
+                            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 {
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(mockCluster(1, 0))) {
+            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,
+                "",
+                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();
+
+            assertNull(groupDescription.authorizedOperations());
+        }
+    }
+
+    @Test
+    public void testDescribeNonConsumerGroups() throws Exception {
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(mockCluster(1, 0))) {
+            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);
+        }
+    }
+
+    @Test
+    public void testListConsumerGroupOffsetsNumRetries() throws Exception {
+        final Cluster cluster = mockCluster(3, 0);
+        final Time time = new MockTime();
+
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(time, cluster,
+            AdminClientConfig.RETRIES_CONFIG, "0")) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            env.kafkaClient().prepareResponse(prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+            env.kafkaClient().prepareResponse(new OffsetFetchResponse(Errors.NOT_COORDINATOR, Collections.emptyMap()));
+            env.kafkaClient().prepareResponse(prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+
+            final ListConsumerGroupOffsetsResult result = env.adminClient().listConsumerGroupOffsets("group-0");
+
+
+            TestUtils.assertFutureError(result.partitionsToOffsetAndMetadata(), TimeoutException.class);
+        }
+    }
+
+    @Test
+    public void testListConsumerGroupOffsetsRetryBackoff() throws Exception {
+        MockTime time = new MockTime();
+        int retryBackoff = 100;
+
+        try (final AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(time,
+            mockCluster(3, 0),
+            newStrMap(AdminClientConfig.RETRY_BACKOFF_MS_CONFIG, "" + retryBackoff))) {
+            MockClient mockClient = env.kafkaClient();
+
+            mockClient.setNodeApiVersions(NodeApiVersions.create());
+
+            AtomicLong firstAttemptTime = new AtomicLong(0);
+            AtomicLong secondAttemptTime = new AtomicLong(0);
+
+            mockClient.prepareResponse(prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+            mockClient.prepareResponse(body -> {
+                firstAttemptTime.set(time.milliseconds());
+                return true;
+            }, new OffsetFetchResponse(Errors.NOT_COORDINATOR, Collections.emptyMap()));
+
+            mockClient.prepareResponse(prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+
+            mockClient.prepareResponse(body -> {
+                secondAttemptTime.set(time.milliseconds());
+                return true;
+            }, new OffsetFetchResponse(Errors.NONE, Collections.emptyMap()));
+
+            final KafkaFuture> future = env.adminClient().listConsumerGroupOffsets("group-0").partitionsToOffsetAndMetadata();
+
+            TestUtils.waitForCondition(() -> mockClient.numAwaitingResponses() == 1, "Failed awaiting ListConsumerGroupOffsets first request failure");
+            TestUtils.waitForCondition(() -> ((KafkaAdminClient) env.adminClient()).numPendingCalls() == 1, "Failed to add retry ListConsumerGroupOffsets call on first failure");
+            time.sleep(retryBackoff);
+
+            future.get();
+
+            long actualRetryBackoff = secondAttemptTime.get() - firstAttemptTime.get();
+            assertEquals("ListConsumerGroupOffsets retry did not await expected backoff!", retryBackoff, actualRetryBackoff);
+        }
+    }
+
+    @Test
+    public void testListConsumerGroupOffsets() throws Exception {
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(mockCluster(1, 0))) {
+            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.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,
+                    Optional.empty(), "", Errors.NONE));
+            responseData.put(myTopicPartition1, new OffsetFetchResponse.PartitionData(0,
+                    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(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 testDeleteConsumerGroupsNumRetries() throws Exception {
+        final Cluster cluster = mockCluster(3, 0);
+        final Time time = new MockTime();
+        final List groupIds = singletonList("group-0");
+
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(time, cluster,
+            AdminClientConfig.RETRIES_CONFIG, "0")) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            env.kafkaClient().prepareResponse(prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+            final DeletableGroupResultCollection validResponse = new DeletableGroupResultCollection();
+            validResponse.add(new DeletableGroupResult()
+                .setGroupId("group-0")
+                .setErrorCode(Errors.NOT_COORDINATOR.code()));
+            env.kafkaClient().prepareResponse(new DeleteGroupsResponse(
+                new DeleteGroupsResponseData()
+                    .setResults(validResponse)
+            ));
+            env.kafkaClient().prepareResponse(prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+
+            final DeleteConsumerGroupsResult result = env.adminClient().deleteConsumerGroups(groupIds);
+
+            TestUtils.assertFutureError(result.all(), TimeoutException.class);
+        }
+    }
+
+    @Test
+    public void testDeleteConsumerGroupsRetryBackoff() throws Exception {
+        MockTime time = new MockTime();
+        int retryBackoff = 100;
+        final List groupIds = singletonList("group-0");
+
+        try (final AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(time,
+            mockCluster(3, 0),
+            newStrMap(AdminClientConfig.RETRY_BACKOFF_MS_CONFIG, "" + retryBackoff))) {
+            MockClient mockClient = env.kafkaClient();
+
+            mockClient.setNodeApiVersions(NodeApiVersions.create());
+
+            AtomicLong firstAttemptTime = new AtomicLong(0);
+            AtomicLong secondAttemptTime = new AtomicLong(0);
+
+            mockClient.prepareResponse(prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+
+            DeletableGroupResultCollection validResponse = new DeletableGroupResultCollection();
+            validResponse.add(new DeletableGroupResult()
+                .setGroupId("group-0")
+                .setErrorCode(Errors.NOT_COORDINATOR.code()));
+
+
+            mockClient.prepareResponse(body -> {
+                firstAttemptTime.set(time.milliseconds());
+                return true;
+            }, new DeleteGroupsResponse(new DeleteGroupsResponseData().setResults(validResponse)));
+
+            mockClient.prepareResponse(prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+
+            validResponse = new DeletableGroupResultCollection();
+            validResponse.add(new DeletableGroupResult()
+                .setGroupId("group-0")
+                .setErrorCode(Errors.NONE.code()));
+
+            mockClient.prepareResponse(body -> {
+                secondAttemptTime.set(time.milliseconds());
+                return true;
+            }, new DeleteGroupsResponse(new DeleteGroupsResponseData().setResults(validResponse)));
+
+            final KafkaFuture future = env.adminClient().deleteConsumerGroups(groupIds).all();
+
+            TestUtils.waitForCondition(() -> mockClient.numAwaitingResponses() == 1, "Failed awaiting DeleteConsumerGroups first request failure");
+            TestUtils.waitForCondition(() -> ((KafkaAdminClient) env.adminClient()).numPendingCalls() == 1, "Failed to add retry DeleteConsumerGroups call on first failure");
+            time.sleep(retryBackoff);
+
+            future.get();
+
+            long actualRetryBackoff = secondAttemptTime.get() - firstAttemptTime.get();
+            assertEquals("DeleteConsumerGroups retry did not await expected backoff!", retryBackoff, actualRetryBackoff);
+        }
+    }
+
+    @Test
+    public void testDeleteConsumerGroups() throws Exception {
+        final List groupIds = singletonList("group-0");
+
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(mockCluster(1, 0))) {
+            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()));
+
+            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);
+
+            final KafkaFuture results = result.deletedGroups().get("group-0");
+            assertNull(results.get());
+
+            //should throw error for non-retriable errors
+            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(
+                prepareFindCoordinatorResponse(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 testDeleteConsumerGroupOffsetsNumRetries() throws Exception {
+        final Cluster cluster = mockCluster(3, 0);
+        final Time time = new MockTime();
+
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(time, cluster,
+            AdminClientConfig.RETRIES_CONFIG, "0")) {
+            final TopicPartition tp1 = new TopicPartition("foo", 0);
+
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            env.kafkaClient().prepareResponse(prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+            env.kafkaClient().prepareResponse(prepareOffsetDeleteResponse("foo", 0, Errors.NOT_COORDINATOR));
+            env.kafkaClient().prepareResponse(prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+
+            final DeleteConsumerGroupOffsetsResult result = env.adminClient()
+                .deleteConsumerGroupOffsets("group-0", Stream.of(tp1).collect(Collectors.toSet()));
+
+            TestUtils.assertFutureError(result.all(), TimeoutException.class);
+        }
+    }
+
+    @Test
+    public void testDeleteConsumerGroupOffsetsRetryBackoff() throws Exception {
+        MockTime time = new MockTime();
+        int retryBackoff = 100;
+
+        try (final AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(time,
+            mockCluster(3, 0),
+            newStrMap(AdminClientConfig.RETRY_BACKOFF_MS_CONFIG, "" + retryBackoff))) {
+            MockClient mockClient = env.kafkaClient();
+
+            mockClient.setNodeApiVersions(NodeApiVersions.create());
+
+            AtomicLong firstAttemptTime = new AtomicLong(0);
+            AtomicLong secondAttemptTime = new AtomicLong(0);
+
+            final TopicPartition tp1 = new TopicPartition("foo", 0);
+
+            mockClient.prepareResponse(prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+
+            mockClient.prepareResponse(body -> {
+                firstAttemptTime.set(time.milliseconds());
+                return true;
+            }, prepareOffsetDeleteResponse("foo", 0, Errors.NOT_COORDINATOR));
+
+            mockClient.prepareResponse(prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+
+            mockClient.prepareResponse(body -> {
+                secondAttemptTime.set(time.milliseconds());
+                return true;
+            }, prepareOffsetDeleteResponse("foo", 0, Errors.NONE));
+
+            final KafkaFuture future = env.adminClient().deleteConsumerGroupOffsets("group-0", Stream.of(tp1).collect(Collectors.toSet())).all();
+
+            TestUtils.waitForCondition(() -> mockClient.numAwaitingResponses() == 1, "Failed awaiting DeleteConsumerGroupOffsets first request failure");
+            TestUtils.waitForCondition(() -> ((KafkaAdminClient) env.adminClient()).numPendingCalls() == 1, "Failed to add retry DeleteConsumerGroupOffsets call on first failure");
+            time.sleep(retryBackoff);
+
+            future.get();
+
+            long actualRetryBackoff = secondAttemptTime.get() - firstAttemptTime.get();
+            assertEquals("DeleteConsumerGroupOffsets retry did not await expected backoff!", retryBackoff, actualRetryBackoff);
+        }
+    }
+
+    @Test
+    public void testDeleteConsumerGroupOffsets() throws Exception {
+        // Happy path
+
+        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(mockCluster(1, 0))) {
+            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 String groupId = "group-0";
+        final TopicPartition tp1 = new TopicPartition("foo", 0);
+
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(mockCluster(1, 0))) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            env.kafkaClient().prepareResponse(
+                prepareFindCoordinatorResponse(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 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(mockCluster(1, 0))) {
+            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 String groupId = "group-0";
+        final TopicPartition tp1 = new TopicPartition("foo", 0);
+
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(mockCluster(1, 0))) {
+            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 String groupId = "group-0";
+        final TopicPartition tp1 = new TopicPartition("foo", 0);
+
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(mockCluster(1, 0))) {
+            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 testRemoveMembersFromGroupNumRetries() throws Exception {
+        final Cluster cluster = mockCluster(3, 0);
+        final Time time = new MockTime();
+
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(time, cluster,
+            AdminClientConfig.RETRIES_CONFIG, "0")) {
+
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            env.kafkaClient().prepareResponse(prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+            env.kafkaClient().prepareResponse(new LeaveGroupResponse(new LeaveGroupResponseData().setErrorCode(Errors.NOT_COORDINATOR.code())));
+            env.kafkaClient().prepareResponse(prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+
+            Collection membersToRemove = Arrays.asList(new MemberToRemove("instance-1"), new MemberToRemove("instance-2"));
+
+            final RemoveMembersFromConsumerGroupResult result = env.adminClient().removeMembersFromConsumerGroup(
+                "groupId", new RemoveMembersFromConsumerGroupOptions(membersToRemove));
+
+            TestUtils.assertFutureError(result.all(), TimeoutException.class);
+        }
+    }
+
+    @Test
+    public void testRemoveMembersFromGroupRetryBackoff() throws Exception {
+        MockTime time = new MockTime();
+        int retryBackoff = 100;
+
+        try (final AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(time,
+            mockCluster(3, 0),
+            newStrMap(AdminClientConfig.RETRY_BACKOFF_MS_CONFIG, "" + retryBackoff))) {
+            MockClient mockClient = env.kafkaClient();
+
+            mockClient.setNodeApiVersions(NodeApiVersions.create());
+
+            AtomicLong firstAttemptTime = new AtomicLong(0);
+            AtomicLong secondAttemptTime = new AtomicLong(0);
+
+            mockClient.prepareResponse(prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+
+            env.kafkaClient().prepareResponse(body -> {
+                firstAttemptTime.set(time.milliseconds());
+                return true;
+            }, new LeaveGroupResponse(new LeaveGroupResponseData().setErrorCode(Errors.NOT_COORDINATOR.code())));
+
+            mockClient.prepareResponse(prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+
+            MemberResponse responseOne = new MemberResponse()
+                .setGroupInstanceId("instance-1")
+                .setErrorCode(Errors.NONE.code());
+            env.kafkaClient().prepareResponse(body -> {
+                secondAttemptTime.set(time.milliseconds());
+                return true;
+            }, new LeaveGroupResponse(new LeaveGroupResponseData()
+                .setErrorCode(Errors.NONE.code())
+                .setMembers(Collections.singletonList(responseOne))));
+
+            Collection membersToRemove = singletonList(new MemberToRemove("instance-1"));
+
+            final KafkaFuture future = env.adminClient().removeMembersFromConsumerGroup(
+                "groupId", new RemoveMembersFromConsumerGroupOptions(membersToRemove)).all();
+
+
+            TestUtils.waitForCondition(() -> mockClient.numAwaitingResponses() == 1, "Failed awaiting RemoveMembersFromGroup first request failure");
+            TestUtils.waitForCondition(() -> ((KafkaAdminClient) env.adminClient()).numPendingCalls() == 1, "Failed to add retry RemoveMembersFromGroup call on first failure");
+            time.sleep(retryBackoff);
+
+            future.get();
+
+            long actualRetryBackoff = secondAttemptTime.get() - firstAttemptTime.get();
+            assertEquals("RemoveMembersFromGroup retry did not await expected backoff!", retryBackoff, actualRetryBackoff);
+        }
+    }
+
+    @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 the "removeAll" scenario
+            final List topicPartitions = Arrays.asList(1, 2, 3).stream().map(partition -> new TopicPartition("my_topic", partition))
+                    .collect(Collectors.toList());
+            // construct the DescribeGroupsResponse
+            DescribeGroupsResponseData data = prepareDescribeGroupsResponseData(groupId, Arrays.asList(instanceOne, instanceTwo), topicPartitions);
+
+            // Return with partial failure for "removeAll" scenario
+            // 1 prepare response for AdminClient.describeConsumerGroups
+            env.kafkaClient().prepareResponse(prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+            env.kafkaClient().prepareResponse(new DescribeGroupsResponse(data));
+
+            // 2 KafkaAdminClient encounter partial failure when trying to delete all members
+            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 partialFailureResults = env.adminClient().removeMembersFromConsumerGroup(
+                    groupId,
+                    new RemoveMembersFromConsumerGroupOptions()
+            );
+            ExecutionException exception = assertThrows(ExecutionException.class, () -> partialFailureResults.all().get());
+            assertTrue(exception.getCause() instanceof KafkaException);
+            assertTrue(exception.getCause().getCause() instanceof UnknownMemberIdException);
+
+            // Return with success for "removeAll" scenario
+            // 1 prepare response for AdminClient.describeConsumerGroups
+            env.kafkaClient().prepareResponse(prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+            env.kafkaClient().prepareResponse(new DescribeGroupsResponse(data));
+
+            // 2. KafkaAdminClient should delete all members correctly
+            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 successResult = env.adminClient().removeMembersFromConsumerGroup(
+                    groupId,
+                    new RemoveMembersFromConsumerGroupOptions()
+            );
+            assertNull(successResult.all().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 = RequestTestUtils.metadataResponse(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
+            String errorMessage = "this is custom error message";
+            AlterPartitionReassignmentsResponseData topLevelErrResponseData =
+                    new AlterPartitionReassignmentsResponseData()
+                            .setErrorCode(Errors.CLUSTER_AUTHORIZATION_FAILED.code())
+                            .setErrorMessage(errorMessage)
+                            .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);
+            assertEquals(errorMessage, TestUtils.assertFutureThrows(topLevelErrResult.all(), Errors.CLUSTER_AUTHORIZATION_FAILED.exception().getClass()).getMessage());
+            assertEquals(errorMessage, TestUtils.assertFutureThrows(topLevelErrResult.values().get(tp1), Errors.CLUSTER_AUTHORIZATION_FAILED.exception().getClass()).getMessage());
+            assertEquals(errorMessage, TestUtils.assertFutureThrows(topLevelErrResult.values().get(tp2), Errors.CLUSTER_AUTHORIZATION_FAILED.exception().getClass()).getMessage());
+
+            // 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 = RequestTestUtils.metadataResponse(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 testAlterConsumerGroupOffsets() throws Exception {
+        // Happy path
+
+        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(mockCluster(1, 0))) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            env.kafkaClient().prepareResponse(
+                prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+
+            Map responseData = new HashMap<>();
+            responseData.put(tp1, Errors.NONE);
+            responseData.put(tp2, Errors.NONE);
+            env.kafkaClient().prepareResponse(new OffsetCommitResponse(0, responseData));
+
+            Map offsets = new HashMap<>();
+            offsets.put(tp1, new OffsetAndMetadata(123L));
+            offsets.put(tp2, new OffsetAndMetadata(456L));
+            final AlterConsumerGroupOffsetsResult result = env.adminClient().alterConsumerGroupOffsets(
+                groupId, offsets);
+
+            assertNull(result.all().get());
+            assertNull(result.partitionResult(tp1).get());
+            assertNull(result.partitionResult(tp2).get());
+            TestUtils.assertFutureError(result.partitionResult(tp3), IllegalArgumentException.class);
+        }
+    }
+
+    @Test
+    public void testAlterConsumerGroupOffsetsRetriableErrors() throws Exception {
+        // Retriable errors should be retried
+
+        final String groupId = "group-0";
+        final TopicPartition tp1 = new TopicPartition("foo", 0);
+
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(mockCluster(1, 0))) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            env.kafkaClient().prepareResponse(
+                prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+
+            env.kafkaClient().prepareResponse(
+                prepareOffsetCommitResponse(tp1, Errors.COORDINATOR_NOT_AVAILABLE));
+
+            env.kafkaClient().prepareResponse(
+                prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+
+            env.kafkaClient().prepareResponse(
+                prepareOffsetCommitResponse(tp1, Errors.COORDINATOR_LOAD_IN_PROGRESS));
+
+            env.kafkaClient().prepareResponse(
+                prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+
+            env.kafkaClient().prepareResponse(
+                prepareOffsetCommitResponse(tp1, Errors.NOT_COORDINATOR));
+
+            env.kafkaClient().prepareResponse(
+                prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+
+            env.kafkaClient().prepareResponse(
+                prepareOffsetCommitResponse(tp1, Errors.NONE));
+
+            Map offsets = new HashMap<>();
+            offsets.put(tp1, new OffsetAndMetadata(123L));
+            final AlterConsumerGroupOffsetsResult result1 = env.adminClient()
+                .alterConsumerGroupOffsets(groupId, offsets);
+
+            assertNull(result1.all().get());
+            assertNull(result1.partitionResult(tp1).get());
+        }
+    }
+
+    @Test
+    public void testAlterConsumerGroupOffsetsNonRetriableErrors() throws Exception {
+        // Non-retriable errors throw an exception
+
+        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(mockCluster(1, 0))) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            for (Errors error : nonRetriableErrors) {
+                env.kafkaClient().prepareResponse(
+                    prepareFindCoordinatorResponse(Errors.NONE, env.cluster().controller()));
+
+                env.kafkaClient().prepareResponse(prepareOffsetCommitResponse(tp1, error));
+
+                Map offsets = new HashMap<>();
+                offsets.put(tp1,  new OffsetAndMetadata(123L));
+                AlterConsumerGroupOffsetsResult errorResult = env.adminClient()
+                    .alterConsumerGroupOffsets(groupId, offsets);
+
+                TestUtils.assertFutureError(errorResult.all(), error.exception().getClass());
+                TestUtils.assertFutureError(errorResult.partitionResult(tp1), error.exception().getClass());
+            }
+        }
+    }
+
+    @Test
+    public void testAlterConsumerGroupOffsetsFindCoordinatorRetriableErrors() throws Exception {
+        // Retriable FindCoordinatorResponse errors should be retried
+
+        final String groupId = "group-0";
+        final TopicPartition tp1 = new TopicPartition("foo", 0);
+
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(mockCluster(1, 0))) {
+            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(
+                prepareOffsetCommitResponse(tp1, Errors.NONE));
+
+            Map offsets = new HashMap<>();
+            offsets.put(tp1,  new OffsetAndMetadata(123L));
+            final AlterConsumerGroupOffsetsResult result = env.adminClient()
+                .alterConsumerGroupOffsets(groupId, offsets);
+
+            assertNull(result.all().get());
+            assertNull(result.partitionResult(tp1).get());
+        }
+    }
+
+    @Test
+    public void testAlterConsumerGroupOffsetsFindCoordinatorNonRetriableErrors() throws Exception {
+        // Non-retriable FindCoordinatorResponse errors throw an exception
+
+        final String groupId = "group-0";
+        final TopicPartition tp1 = new TopicPartition("foo", 0);
+
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(mockCluster(1, 0))) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            env.kafkaClient().prepareResponse(
+                prepareFindCoordinatorResponse(Errors.GROUP_AUTHORIZATION_FAILED,  Node.noNode()));
+
+            Map offsets = new HashMap<>();
+            offsets.put(tp1,  new OffsetAndMetadata(123L));
+            final AlterConsumerGroupOffsetsResult errorResult = env.adminClient()
+                .alterConsumerGroupOffsets(groupId, offsets);
+
+            TestUtils.assertFutureError(errorResult.all(), GroupAuthorizationException.class);
+            TestUtils.assertFutureError(errorResult.partitionResult(tp1), GroupAuthorizationException.class);
+        }
+    }
+
+    @Test
+    public void testListOffsets() throws Exception {
+        // Happy path
+
+        Node node0 = new Node(0, "localhost", 8120);
+        List pInfos = new ArrayList<>();
+        pInfos.add(new PartitionInfo("foo", 0, node0, new Node[]{node0}, new Node[]{node0}));
+        pInfos.add(new PartitionInfo("bar", 0, node0, new Node[]{node0}, new Node[]{node0}));
+        pInfos.add(new PartitionInfo("baz", 0, node0, new Node[]{node0}, new Node[]{node0}));
+        final Cluster cluster =
+            new Cluster(
+                "mockClusterId",
+                Arrays.asList(node0),
+                pInfos,
+                Collections.emptySet(),
+                Collections.emptySet(),
+                node0);
+
+        final TopicPartition tp0 = new TopicPartition("foo", 0);
+        final TopicPartition tp1 = new TopicPartition("bar", 0);
+        final TopicPartition tp2 = new TopicPartition("baz", 0);
+
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(cluster)) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            env.kafkaClient().prepareResponse(prepareMetadataResponse(cluster, Errors.NONE));
+
+            ListOffsetTopicResponse t0 = ListOffsetResponse.singletonListOffsetTopicResponse(tp0, Errors.NONE, -1L, 123L, 321);
+            ListOffsetTopicResponse t1 = ListOffsetResponse.singletonListOffsetTopicResponse(tp1, Errors.NONE, -1L, 234L, 432);
+            ListOffsetTopicResponse t2 = ListOffsetResponse.singletonListOffsetTopicResponse(tp2, Errors.NONE, 123456789L, 345L, 543);
+            ListOffsetResponseData responseData = new ListOffsetResponseData()
+                    .setThrottleTimeMs(0)
+                    .setTopics(Arrays.asList(t0, t1, t2));
+            env.kafkaClient().prepareResponse(new ListOffsetResponse(responseData));
+
+            Map partitions = new HashMap<>();
+            partitions.put(tp0, OffsetSpec.latest());
+            partitions.put(tp1, OffsetSpec.earliest());
+            partitions.put(tp2, OffsetSpec.forTimestamp(System.currentTimeMillis()));
+            ListOffsetsResult result = env.adminClient().listOffsets(partitions);
+
+            Map offsets = result.all().get();
+            assertFalse(offsets.isEmpty());
+            assertEquals(123L, offsets.get(tp0).offset());
+            assertEquals(321, offsets.get(tp0).leaderEpoch().get().intValue());
+            assertEquals(-1L, offsets.get(tp0).timestamp());
+            assertEquals(234L, offsets.get(tp1).offset());
+            assertEquals(432, offsets.get(tp1).leaderEpoch().get().intValue());
+            assertEquals(-1L, offsets.get(tp1).timestamp());
+            assertEquals(345L, offsets.get(tp2).offset());
+            assertEquals(543, offsets.get(tp2).leaderEpoch().get().intValue());
+            assertEquals(123456789L, offsets.get(tp2).timestamp());
+            assertEquals(offsets.get(tp0), result.partitionResult(tp0).get());
+            assertEquals(offsets.get(tp1), result.partitionResult(tp1).get());
+            assertEquals(offsets.get(tp2), result.partitionResult(tp2).get());
+            try {
+                result.partitionResult(new TopicPartition("unknown", 0)).get();
+                fail("should have thrown IllegalArgumentException");
+            } catch (IllegalArgumentException expected) { }
+        }
+    }
+
+    @Test
+    public void testListOffsetsRetriableErrors() throws Exception {
+
+        Node node0 = new Node(0, "localhost", 8120);
+        Node node1 = new Node(1, "localhost", 8121);
+        List nodes = Arrays.asList(node0, node1);
+        List pInfos = new ArrayList<>();
+        pInfos.add(new PartitionInfo("foo", 0, node0, new Node[]{node0, node1}, new Node[]{node0, node1}));
+        pInfos.add(new PartitionInfo("foo", 1, node0, new Node[]{node0, node1}, new Node[]{node0, node1}));
+        pInfos.add(new PartitionInfo("bar", 0, node1, new Node[]{node1, node0}, new Node[]{node1, node0}));
+        final Cluster cluster =
+            new Cluster(
+                "mockClusterId",
+                nodes,
+                pInfos,
+                Collections.emptySet(),
+                Collections.emptySet(),
+                node0);
+
+        final TopicPartition tp0 = new TopicPartition("foo", 0);
+        final TopicPartition tp1 = new TopicPartition("foo", 1);
+        final TopicPartition tp2 = new TopicPartition("bar", 0);
+
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(cluster)) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+            env.kafkaClient().prepareResponse(prepareMetadataResponse(cluster, Errors.NONE));
+            // listoffsets response from broker 0
+            ListOffsetTopicResponse t0 = ListOffsetResponse.singletonListOffsetTopicResponse(tp0, Errors.LEADER_NOT_AVAILABLE, -1L, 123L, 321);
+            ListOffsetTopicResponse t1 = ListOffsetResponse.singletonListOffsetTopicResponse(tp1, Errors.NONE, -1L, 987L, 789);
+            ListOffsetResponseData responseData = new ListOffsetResponseData()
+                    .setThrottleTimeMs(0)
+                    .setTopics(Arrays.asList(t0, t1));
+            env.kafkaClient().prepareResponseFrom(new ListOffsetResponse(responseData), node0);
+            // listoffsets response from broker 1
+            ListOffsetTopicResponse t2 = ListOffsetResponse.singletonListOffsetTopicResponse(tp2, Errors.NONE, -1L, 456L, 654);
+            responseData = new ListOffsetResponseData()
+                    .setThrottleTimeMs(0)
+                    .setTopics(Arrays.asList(t2));
+            env.kafkaClient().prepareResponseFrom(new ListOffsetResponse(responseData), node1);
+
+            // metadata refresh because of LEADER_NOT_AVAILABLE
+            env.kafkaClient().prepareResponse(prepareMetadataResponse(cluster, Errors.NONE));
+            // listoffsets response from broker 0
+            t0 = ListOffsetResponse.singletonListOffsetTopicResponse(tp0, Errors.NONE, -1L, 345L, 543);
+            responseData = new ListOffsetResponseData()
+                    .setThrottleTimeMs(0)
+                    .setTopics(Arrays.asList(t0));
+            env.kafkaClient().prepareResponseFrom(new ListOffsetResponse(responseData), node0);
+
+            Map partitions = new HashMap<>();
+            partitions.put(tp0, OffsetSpec.latest());
+            partitions.put(tp1, OffsetSpec.latest());
+            partitions.put(tp2, OffsetSpec.latest());
+            ListOffsetsResult result = env.adminClient().listOffsets(partitions);
+
+            Map offsets = result.all().get();
+            assertFalse(offsets.isEmpty());
+            assertEquals(345L, offsets.get(tp0).offset());
+            assertEquals(543, offsets.get(tp0).leaderEpoch().get().intValue());
+            assertEquals(-1L, offsets.get(tp0).timestamp());
+            assertEquals(987L, offsets.get(tp1).offset());
+            assertEquals(789, offsets.get(tp1).leaderEpoch().get().intValue());
+            assertEquals(-1L, offsets.get(tp1).timestamp());
+            assertEquals(456L, offsets.get(tp2).offset());
+            assertEquals(654, offsets.get(tp2).leaderEpoch().get().intValue());
+            assertEquals(-1L, offsets.get(tp2).timestamp());
+        }
+    }
+
+    @Test
+    public void testListOffsetsNonRetriableErrors() throws Exception {
+
+        Node node0 = new Node(0, "localhost", 8120);
+        Node node1 = new Node(1, "localhost", 8121);
+        List nodes = Arrays.asList(node0, node1);
+        List pInfos = new ArrayList<>();
+        pInfos.add(new PartitionInfo("foo", 0, node0, new Node[]{node0, node1}, new Node[]{node0, node1}));
+        final Cluster cluster =
+            new Cluster(
+                "mockClusterId",
+                nodes,
+                pInfos,
+                Collections.emptySet(),
+                Collections.emptySet(),
+                node0);
+
+        final TopicPartition tp0 = new TopicPartition("foo", 0);
+
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(cluster)) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            env.kafkaClient().prepareResponse(prepareMetadataResponse(cluster, Errors.NONE));
+
+            ListOffsetTopicResponse t0 = ListOffsetResponse.singletonListOffsetTopicResponse(tp0, Errors.TOPIC_AUTHORIZATION_FAILED, -1L, -1L, -1);
+            ListOffsetResponseData responseData = new ListOffsetResponseData()
+                    .setThrottleTimeMs(0)
+                    .setTopics(Arrays.asList(t0));
+            env.kafkaClient().prepareResponse(new ListOffsetResponse(responseData));
+
+            Map partitions = new HashMap<>();
+            partitions.put(tp0, OffsetSpec.latest());
+            ListOffsetsResult result = env.adminClient().listOffsets(partitions);
+
+            TestUtils.assertFutureError(result.all(), TopicAuthorizationException.class);
+        }
+    }
+
+    private Map makeTestFeatureUpdates() {
+        return Utils.mkMap(
+            Utils.mkEntry("test_feature_1", new FeatureUpdate((short) 2, false)),
+            Utils.mkEntry("test_feature_2", new FeatureUpdate((short) 3, true)));
+    }
+
+    private Map makeTestFeatureUpdateErrors(final Map updates, final Errors error) {
+        final Map errors = new HashMap<>();
+        for (Map.Entry entry : updates.entrySet()) {
+            errors.put(entry.getKey(), new ApiError(error));
+        }
+        return errors;
+    }
+
+    private void testUpdateFeatures(Map featureUpdates,
+                                    ApiError topLevelError,
+                                    Map featureUpdateErrors) throws Exception {
+        try (final AdminClientUnitTestEnv env = mockClientEnv()) {
+            env.kafkaClient().prepareResponse(
+                body -> body instanceof UpdateFeaturesRequest,
+                UpdateFeaturesResponse.createWithErrors(topLevelError, featureUpdateErrors, 0));
+            final Map> futures = env.adminClient().updateFeatures(
+                featureUpdates,
+                new UpdateFeaturesOptions().timeoutMs(10000)).values();
+            for (final Map.Entry> entry : futures.entrySet()) {
+                final KafkaFuture future = entry.getValue();
+                final ApiError error = featureUpdateErrors.get(entry.getKey());
+                if (topLevelError.error() == Errors.NONE) {
+                    assertNotNull(error);
+                    if (error.error() == Errors.NONE) {
+                        future.get();
+                    } else {
+                        final ExecutionException e = assertThrows(ExecutionException.class,
+                            () -> future.get());
+                        assertEquals(e.getCause().getClass(), error.exception().getClass());
+                    }
+                } else {
+                    final ExecutionException e = assertThrows(ExecutionException.class,
+                        () -> future.get());
+                    assertEquals(e.getCause().getClass(), topLevelError.exception().getClass());
+                }
+            }
+        }
+    }
+
+    @Test
+    public void testUpdateFeaturesDuringSuccess() throws Exception {
+        final Map updates = makeTestFeatureUpdates();
+        testUpdateFeatures(updates, ApiError.NONE, makeTestFeatureUpdateErrors(updates, Errors.NONE));
+    }
+
+    @Test
+    public void testUpdateFeaturesTopLevelError() throws Exception {
+        final Map updates = makeTestFeatureUpdates();
+        testUpdateFeatures(updates, new ApiError(Errors.INVALID_REQUEST), new HashMap<>());
+    }
+
+    @Test
+    public void testUpdateFeaturesInvalidRequestError() throws Exception {
+        final Map updates = makeTestFeatureUpdates();
+        testUpdateFeatures(updates, ApiError.NONE, makeTestFeatureUpdateErrors(updates, Errors.INVALID_REQUEST));
+    }
+
+    @Test
+    public void testUpdateFeaturesUpdateFailedError() throws Exception {
+        final Map updates = makeTestFeatureUpdates();
+        testUpdateFeatures(updates, ApiError.NONE, makeTestFeatureUpdateErrors(updates, Errors.FEATURE_UPDATE_FAILED));
+    }
+
+    @Test
+    public void testUpdateFeaturesPartialSuccess() throws Exception {
+        final Map errors = makeTestFeatureUpdateErrors(makeTestFeatureUpdates(), Errors.NONE);
+        errors.put("test_feature_2", new ApiError(Errors.INVALID_REQUEST));
+        testUpdateFeatures(makeTestFeatureUpdates(), ApiError.NONE, errors);
+    }
+
+    @Test
+    public void testUpdateFeaturesHandleNotControllerException() throws Exception {
+        try (final AdminClientUnitTestEnv env = mockClientEnv()) {
+            env.kafkaClient().prepareResponseFrom(
+                request -> request instanceof UpdateFeaturesRequest,
+                UpdateFeaturesResponse.createWithErrors(
+                    new ApiError(Errors.NOT_CONTROLLER),
+                    Utils.mkMap(),
+                    0),
+                env.cluster().nodeById(0));
+            final int controllerId = 1;
+            env.kafkaClient().prepareResponse(RequestTestUtils.metadataResponse(env.cluster().nodes(),
+                env.cluster().clusterResource().clusterId(),
+                controllerId,
+                Collections.emptyList()));
+            env.kafkaClient().prepareResponseFrom(
+                request -> request instanceof UpdateFeaturesRequest,
+                UpdateFeaturesResponse.createWithErrors(
+                    ApiError.NONE,
+                    Utils.mkMap(Utils.mkEntry("test_feature_1", ApiError.NONE),
+                                Utils.mkEntry("test_feature_2", ApiError.NONE)),
+                    0),
+                env.cluster().nodeById(controllerId));
+            final KafkaFuture future = env.adminClient().updateFeatures(
+                Utils.mkMap(
+                    Utils.mkEntry("test_feature_1", new FeatureUpdate((short) 2, false)),
+                    Utils.mkEntry("test_feature_2", new FeatureUpdate((short) 3, true))),
+                new UpdateFeaturesOptions().timeoutMs(10000)
+            ).all();
+            future.get();
+        }
+    }
+
+    @Test
+    public void testUpdateFeaturesShouldFailRequestForEmptyUpdates() {
+        try (final AdminClientUnitTestEnv env = mockClientEnv()) {
+            assertThrows(
+                IllegalArgumentException.class,
+                () -> env.adminClient().updateFeatures(
+                    new HashMap<>(), new UpdateFeaturesOptions()));
+        }
+    }
+
+    @Test
+    public void testUpdateFeaturesShouldFailRequestForInvalidFeatureName() {
+        try (final AdminClientUnitTestEnv env = mockClientEnv()) {
+            assertThrows(
+                IllegalArgumentException.class,
+                () -> env.adminClient().updateFeatures(
+                    Utils.mkMap(Utils.mkEntry("feature", new FeatureUpdate((short) 2, false)),
+                                Utils.mkEntry("", new FeatureUpdate((short) 2, false))),
+                    new UpdateFeaturesOptions()));
+        }
+    }
+
+    @Test
+    public void testUpdateFeaturesShouldFailRequestInClientWhenDowngradeFlagIsNotSetDuringDeletion() {
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> new FeatureUpdate((short) 0, false));
+    }
+
+    @Test
+    public void testDescribeFeaturesSuccess() throws Exception {
+        try (final AdminClientUnitTestEnv env = mockClientEnv()) {
+            env.kafkaClient().prepareResponse(
+                body -> body instanceof ApiVersionsRequest,
+                prepareApiVersionsResponseForDescribeFeatures(Errors.NONE));
+            final KafkaFuture future = env.adminClient().describeFeatures(
+                new DescribeFeaturesOptions().timeoutMs(10000)).featureMetadata();
+            final FeatureMetadata metadata = future.get();
+            assertEquals(defaultFeatureMetadata(), metadata);
+        }
+    }
+
+    @Test
+    public void testDescribeFeaturesFailure() {
+        try (final AdminClientUnitTestEnv env = mockClientEnv()) {
+            env.kafkaClient().prepareResponse(
+                body -> body instanceof ApiVersionsRequest,
+                prepareApiVersionsResponseForDescribeFeatures(Errors.INVALID_REQUEST));
+            final DescribeFeaturesOptions options = new DescribeFeaturesOptions();
+            options.timeoutMs(10000);
+            final KafkaFuture future = env.adminClient().describeFeatures(options).featureMetadata();
+            final ExecutionException e = assertThrows(ExecutionException.class, () -> future.get());
+            assertEquals(e.getCause().getClass(), Errors.INVALID_REQUEST.exception().getClass());
+        }
+    }
+
+    @Test
+    public void testListOffsetsMetadataRetriableErrors() throws Exception {
+
+        Node node0 = new Node(0, "localhost", 8120);
+        Node node1 = new Node(1, "localhost", 8121);
+        List nodes = Arrays.asList(node0, node1);
+        List pInfos = new ArrayList<>();
+        pInfos.add(new PartitionInfo("foo", 0, node0, new Node[]{node0}, new Node[]{node0}));
+        pInfos.add(new PartitionInfo("foo", 1, node1, new Node[]{node1}, new Node[]{node1}));
+        final Cluster cluster =
+            new Cluster(
+                "mockClusterId",
+                nodes,
+                pInfos,
+                Collections.emptySet(),
+                Collections.emptySet(),
+                node0);
+
+        final TopicPartition tp0 = new TopicPartition("foo", 0);
+        final TopicPartition tp1 = new TopicPartition("foo", 1);
+
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(cluster)) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            env.kafkaClient().prepareResponse(prepareMetadataResponse(cluster, Errors.LEADER_NOT_AVAILABLE));
+            env.kafkaClient().prepareResponse(prepareMetadataResponse(cluster, Errors.UNKNOWN_TOPIC_OR_PARTITION));
+            env.kafkaClient().prepareResponse(prepareMetadataResponse(cluster, Errors.NONE));
+
+            // listoffsets response from broker 0
+            ListOffsetTopicResponse t0 = ListOffsetResponse.singletonListOffsetTopicResponse(tp0, Errors.NONE, -1L, 345L, 543);
+            ListOffsetResponseData responseData = new ListOffsetResponseData()
+                    .setThrottleTimeMs(0)
+                    .setTopics(Arrays.asList(t0));
+            env.kafkaClient().prepareResponseFrom(new ListOffsetResponse(responseData), node0);
+            // listoffsets response from broker 1
+            ListOffsetTopicResponse t1 = ListOffsetResponse.singletonListOffsetTopicResponse(tp1, Errors.NONE, -1L, 789L, 987);
+            responseData = new ListOffsetResponseData()
+                    .setThrottleTimeMs(0)
+                    .setTopics(Arrays.asList(t1));
+            env.kafkaClient().prepareResponseFrom(new ListOffsetResponse(responseData), node1);
+
+            Map partitions = new HashMap<>();
+            partitions.put(tp0, OffsetSpec.latest());
+            partitions.put(tp1, OffsetSpec.latest());
+            ListOffsetsResult result = env.adminClient().listOffsets(partitions);
+
+            Map offsets = result.all().get();
+            assertFalse(offsets.isEmpty());
+            assertEquals(345L, offsets.get(tp0).offset());
+            assertEquals(543, offsets.get(tp0).leaderEpoch().get().intValue());
+            assertEquals(-1L, offsets.get(tp0).timestamp());
+            assertEquals(789L, offsets.get(tp1).offset());
+            assertEquals(987, offsets.get(tp1).leaderEpoch().get().intValue());
+            assertEquals(-1L, offsets.get(tp1).timestamp());
+        }
+    }
+
+    @Test
+    public void testListOffsetsWithMultiplePartitionsLeaderChange() throws Exception {
+        Node node0 = new Node(0, "localhost", 8120);
+        Node node1 = new Node(1, "localhost", 8121);
+        Node node2 = new Node(2, "localhost", 8122);
+        List nodes = Arrays.asList(node0, node1, node2);
+
+        final PartitionInfo oldPInfo1 = new PartitionInfo("foo", 0, node0,
+            new Node[]{node0, node1, node2}, new Node[]{node0, node1, node2});
+        final PartitionInfo oldPnfo2 = new PartitionInfo("foo", 1, node0,
+            new Node[]{node0, node1, node2}, new Node[]{node0, node1, node2});
+        List oldPInfos = Arrays.asList(oldPInfo1, oldPnfo2);
+
+        final Cluster oldCluster = new Cluster("mockClusterId", nodes, oldPInfos,
+            Collections.emptySet(), Collections.emptySet(), node0);
+        final TopicPartition tp0 = new TopicPartition("foo", 0);
+        final TopicPartition tp1 = new TopicPartition("foo", 1);
+
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(oldCluster)) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            env.kafkaClient().prepareResponse(prepareMetadataResponse(oldCluster, Errors.NONE));
+
+            ListOffsetTopicResponse t0 = ListOffsetResponse.singletonListOffsetTopicResponse(tp0, Errors.NOT_LEADER_OR_FOLLOWER, -1L, 345L, 543);
+            ListOffsetTopicResponse t1 = ListOffsetResponse.singletonListOffsetTopicResponse(tp1, Errors.LEADER_NOT_AVAILABLE, -2L, 123L, 456);
+            ListOffsetResponseData responseData = new ListOffsetResponseData()
+                    .setThrottleTimeMs(0)
+                    .setTopics(Arrays.asList(t0, t1));
+            env.kafkaClient().prepareResponseFrom(new ListOffsetResponse(responseData), node0);
+
+            final PartitionInfo newPInfo1 = new PartitionInfo("foo", 0, node1,
+                new Node[]{node0, node1, node2}, new Node[]{node0, node1, node2});
+            final PartitionInfo newPInfo2 = new PartitionInfo("foo", 1, node2,
+                new Node[]{node0, node1, node2}, new Node[]{node0, node1, node2});
+            List newPInfos = Arrays.asList(newPInfo1, newPInfo2);
+
+            final Cluster newCluster = new Cluster("mockClusterId", nodes, newPInfos,
+                Collections.emptySet(), Collections.emptySet(), node0);
+
+            env.kafkaClient().prepareResponse(prepareMetadataResponse(newCluster, Errors.NONE));
+
+            t0 = ListOffsetResponse.singletonListOffsetTopicResponse(tp0, Errors.NONE, -1L, 345L, 543);
+            responseData = new ListOffsetResponseData()
+                    .setThrottleTimeMs(0)
+                    .setTopics(Arrays.asList(t0));
+            env.kafkaClient().prepareResponseFrom(new ListOffsetResponse(responseData), node1);
+
+            t1 = ListOffsetResponse.singletonListOffsetTopicResponse(tp1, Errors.NONE, -2L, 123L, 456);
+            responseData = new ListOffsetResponseData()
+                    .setThrottleTimeMs(0)
+                    .setTopics(Arrays.asList(t1));
+            env.kafkaClient().prepareResponseFrom(new ListOffsetResponse(responseData), node2);
+
+            Map partitions = new HashMap<>();
+            partitions.put(tp0, OffsetSpec.latest());
+            partitions.put(tp1, OffsetSpec.latest());
+            ListOffsetsResult result = env.adminClient().listOffsets(partitions);
+            Map offsets = result.all().get();
+
+            assertFalse(offsets.isEmpty());
+            assertEquals(345L, offsets.get(tp0).offset());
+            assertEquals(543, offsets.get(tp0).leaderEpoch().get().intValue());
+            assertEquals(-1L, offsets.get(tp0).timestamp());
+            assertEquals(123L, offsets.get(tp1).offset());
+            assertEquals(456, offsets.get(tp1).leaderEpoch().get().intValue());
+            assertEquals(-2L, offsets.get(tp1).timestamp());
+        }
+    }
+
+    @Test
+    public void testListOffsetsWithLeaderChange() throws Exception {
+        Node node0 = new Node(0, "localhost", 8120);
+        Node node1 = new Node(1, "localhost", 8121);
+        Node node2 = new Node(2, "localhost", 8122);
+        List nodes = Arrays.asList(node0, node1, node2);
+
+        final PartitionInfo oldPartitionInfo = new PartitionInfo("foo", 0, node0,
+            new Node[]{node0, node1, node2}, new Node[]{node0, node1, node2});
+        final Cluster oldCluster = new Cluster("mockClusterId", nodes, singletonList(oldPartitionInfo),
+            Collections.emptySet(), Collections.emptySet(), node0);
+        final TopicPartition tp0 = new TopicPartition("foo", 0);
+
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(oldCluster)) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            env.kafkaClient().prepareResponse(prepareMetadataResponse(oldCluster, Errors.NONE));
+
+            ListOffsetTopicResponse t0 = ListOffsetResponse.singletonListOffsetTopicResponse(tp0, Errors.NOT_LEADER_OR_FOLLOWER, -1L, 345L, 543);
+            ListOffsetResponseData responseData = new ListOffsetResponseData()
+                    .setThrottleTimeMs(0)
+                    .setTopics(Arrays.asList(t0));
+            env.kafkaClient().prepareResponseFrom(new ListOffsetResponse(responseData), node0);
+
+            // updating leader from node0 to node1 and metadata refresh because of NOT_LEADER_OR_FOLLOWER
+            final PartitionInfo newPartitionInfo = new PartitionInfo("foo", 0, node1,
+                new Node[]{node0, node1, node2}, new Node[]{node0, node1, node2});
+            final Cluster newCluster = new Cluster("mockClusterId", nodes, singletonList(newPartitionInfo),
+                Collections.emptySet(), Collections.emptySet(), node0);
+
+            env.kafkaClient().prepareResponse(prepareMetadataResponse(newCluster, Errors.NONE));
+
+            t0 = ListOffsetResponse.singletonListOffsetTopicResponse(tp0, Errors.NONE, -2L, 123L, 456);
+            responseData = new ListOffsetResponseData()
+                    .setThrottleTimeMs(0)
+                    .setTopics(Arrays.asList(t0));
+            env.kafkaClient().prepareResponseFrom(new ListOffsetResponse(responseData), node1);
+
+            Map partitions = new HashMap<>();
+            partitions.put(tp0, OffsetSpec.latest());
+            ListOffsetsResult result = env.adminClient().listOffsets(partitions);
+            Map offsets = result.all().get();
+
+            assertFalse(offsets.isEmpty());
+            assertEquals(123L, offsets.get(tp0).offset());
+            assertEquals(456, offsets.get(tp0).leaderEpoch().get().intValue());
+            assertEquals(-2L, offsets.get(tp0).timestamp());
+        }
+    }
+
+    @Test
+    public void testListOffsetsMetadataNonRetriableErrors() throws Exception {
+
+        Node node0 = new Node(0, "localhost", 8120);
+        Node node1 = new Node(1, "localhost", 8121);
+        List nodes = Arrays.asList(node0, node1);
+        List pInfos = new ArrayList<>();
+        pInfos.add(new PartitionInfo("foo", 0, node0, new Node[]{node0, node1}, new Node[]{node0, node1}));
+        final Cluster cluster =
+            new Cluster(
+                "mockClusterId",
+                nodes,
+                pInfos,
+                Collections.emptySet(),
+                Collections.emptySet(),
+                node0);
+
+        final TopicPartition tp1 = new TopicPartition("foo", 0);
+
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(cluster)) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            env.kafkaClient().prepareResponse(prepareMetadataResponse(cluster, Errors.TOPIC_AUTHORIZATION_FAILED));
+
+            Map partitions = new HashMap<>();
+            partitions.put(tp1, OffsetSpec.latest());
+            ListOffsetsResult result = env.adminClient().listOffsets(partitions);
+
+            TestUtils.assertFutureError(result.all(), TopicAuthorizationException.class);
+        }
+    }
+
+    @Test
+    public void testListOffsetsPartialResponse() throws Exception {
+        Node node0 = new Node(0, "localhost", 8120);
+        Node node1 = new Node(1, "localhost", 8121);
+        List nodes = Arrays.asList(node0, node1);
+        List pInfos = new ArrayList<>();
+        pInfos.add(new PartitionInfo("foo", 0, node0, new Node[]{node0, node1}, new Node[]{node0, node1}));
+        pInfos.add(new PartitionInfo("foo", 1, node0, new Node[]{node0, node1}, new Node[]{node0, node1}));
+        final Cluster cluster =
+            new Cluster(
+                "mockClusterId",
+                nodes,
+                pInfos,
+                Collections.emptySet(),
+                Collections.emptySet(),
+                node0);
+
+        final TopicPartition tp0 = new TopicPartition("foo", 0);
+        final TopicPartition tp1 = new TopicPartition("foo", 1);
+
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(cluster)) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            env.kafkaClient().prepareResponse(prepareMetadataResponse(cluster, Errors.NONE));
+
+            ListOffsetTopicResponse t0 = ListOffsetResponse.singletonListOffsetTopicResponse(tp0, Errors.NONE, -2L, 123L, 456);
+            ListOffsetResponseData data = new ListOffsetResponseData()
+                    .setThrottleTimeMs(0)
+                    .setTopics(Arrays.asList(t0));
+            env.kafkaClient().prepareResponseFrom(new ListOffsetResponse(data), node0);
+
+            Map partitions = new HashMap<>();
+            partitions.put(tp0, OffsetSpec.latest());
+            partitions.put(tp1, OffsetSpec.latest());
+            ListOffsetsResult result = env.adminClient().listOffsets(partitions);
+            assertNotNull(result.partitionResult(tp0).get());
+            TestUtils.assertFutureThrows(result.partitionResult(tp1), ApiException.class);
+            TestUtils.assertFutureThrows(result.all(), ApiException.class);
+        }
+    }
+
+    @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());
+    }
+
+    @Test
+    public void testSuccessfulRetryAfterRequestTimeout() throws Exception {
+        HashMap nodes = new HashMap<>();
+        MockTime time = new MockTime();
+        Node node0 = new Node(0, "localhost", 8121);
+        nodes.put(0, node0);
+        Cluster cluster = new Cluster("mockClusterId", nodes.values(),
+                Arrays.asList(new PartitionInfo("foo", 0, node0, new Node[]{node0}, new Node[]{node0})),
+                Collections.emptySet(), Collections.emptySet(),
+                Collections.emptySet(), nodes.get(0));
+
+        final int requestTimeoutMs = 1000;
+        final int retryBackoffMs = 100;
+        final int apiTimeoutMs = 3000;
+
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(time, cluster,
+                AdminClientConfig.RETRY_BACKOFF_MS_CONFIG, String.valueOf(retryBackoffMs),
+                AdminClientConfig.REQUEST_TIMEOUT_MS_CONFIG, String.valueOf(requestTimeoutMs))) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            final ListTopicsResult result = env.adminClient()
+                    .listTopics(new ListTopicsOptions().timeoutMs(apiTimeoutMs));
+
+            // Wait until the first attempt has been sent, then advance the time
+            TestUtils.waitForCondition(() -> env.kafkaClient().hasInFlightRequests(),
+                    "Timed out waiting for Metadata request to be sent");
+            time.sleep(requestTimeoutMs + 1);
+
+            // Wait for the request to be timed out before backing off
+            TestUtils.waitForCondition(() -> !env.kafkaClient().hasInFlightRequests(),
+                    "Timed out waiting for inFlightRequests to be timed out");
+            time.sleep(retryBackoffMs);
+
+            // Since api timeout bound is not hit, AdminClient should retry
+            TestUtils.waitForCondition(() -> env.kafkaClient().hasInFlightRequests(),
+                    "Failed to retry Metadata request");
+            env.kafkaClient().respond(prepareMetadataResponse(cluster, Errors.NONE));
+
+            assertEquals(1, result.listings().get().size());
+            assertEquals("foo", result.listings().get().iterator().next().name());
+        }
+    }
+
+    @Test
+    public void testDefaultApiTimeout() throws Exception {
+        testApiTimeout(1500, 3000, OptionalInt.empty());
+    }
+
+    @Test
+    public void testDefaultApiTimeoutOverride() throws Exception {
+        testApiTimeout(1500, 10000, OptionalInt.of(3000));
+    }
+
+    private void testApiTimeout(int requestTimeoutMs,
+                                int defaultApiTimeoutMs,
+                                OptionalInt overrideApiTimeoutMs) throws Exception {
+        HashMap nodes = new HashMap<>();
+        MockTime time = new MockTime();
+        Node node0 = new Node(0, "localhost", 8121);
+        nodes.put(0, node0);
+        Cluster cluster = new Cluster("mockClusterId", nodes.values(),
+                Arrays.asList(new PartitionInfo("foo", 0, node0, new Node[]{node0}, new Node[]{node0})),
+                Collections.emptySet(), Collections.emptySet(),
+                Collections.emptySet(), nodes.get(0));
+
+        final int retryBackoffMs = 100;
+        final int effectiveTimeoutMs = overrideApiTimeoutMs.orElse(defaultApiTimeoutMs);
+        assertEquals("This test expects the effective timeout to be twice the request timeout",
+                2 * requestTimeoutMs, effectiveTimeoutMs);
+
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(time, cluster,
+                AdminClientConfig.RETRY_BACKOFF_MS_CONFIG, String.valueOf(retryBackoffMs),
+                AdminClientConfig.REQUEST_TIMEOUT_MS_CONFIG, String.valueOf(requestTimeoutMs),
+                AdminClientConfig.DEFAULT_API_TIMEOUT_MS_CONFIG, String.valueOf(defaultApiTimeoutMs))) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            ListTopicsOptions options = new ListTopicsOptions();
+            overrideApiTimeoutMs.ifPresent(options::timeoutMs);
+
+            final ListTopicsResult result = env.adminClient().listTopics(options);
+
+            // Wait until the first attempt has been sent, then advance the time
+            TestUtils.waitForCondition(() -> env.kafkaClient().hasInFlightRequests(),
+                    "Timed out waiting for Metadata request to be sent");
+            time.sleep(requestTimeoutMs + 1);
+
+            // Wait for the request to be timed out before backing off
+            TestUtils.waitForCondition(() -> !env.kafkaClient().hasInFlightRequests(),
+                    "Timed out waiting for inFlightRequests to be timed out");
+
+            // Since api timeout bound is not hit, AdminClient should retry
+            TestUtils.waitForCondition(() -> {
+                boolean hasInflightRequests = env.kafkaClient().hasInFlightRequests();
+                if (!hasInflightRequests)
+                    time.sleep(retryBackoffMs);
+                return hasInflightRequests;
+            }, "Timed out waiting for Metadata request to be sent");
+            time.sleep(requestTimeoutMs + 1);
+
+            TestUtils.assertFutureThrows(result.future, TimeoutException.class);
+        }
+    }
+
+    @Test
+    public void testRequestTimeoutExceedingDefaultApiTimeout() throws Exception {
+        HashMap nodes = new HashMap<>();
+        MockTime time = new MockTime();
+        Node node0 = new Node(0, "localhost", 8121);
+        nodes.put(0, node0);
+        Cluster cluster = new Cluster("mockClusterId", nodes.values(),
+                Arrays.asList(new PartitionInfo("foo", 0, node0, new Node[]{node0}, new Node[]{node0})),
+                Collections.emptySet(), Collections.emptySet(),
+                Collections.emptySet(), nodes.get(0));
+
+        // This test assumes the default api timeout value of 60000. When the request timeout
+        // is set to something larger, we should adjust the api timeout accordingly for compatibility.
+
+        final int retryBackoffMs = 100;
+        final int requestTimeoutMs = 120000;
+
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(time, cluster,
+                AdminClientConfig.RETRY_BACKOFF_MS_CONFIG, String.valueOf(retryBackoffMs),
+                AdminClientConfig.REQUEST_TIMEOUT_MS_CONFIG, String.valueOf(requestTimeoutMs))) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            ListTopicsOptions options = new ListTopicsOptions();
+
+            final ListTopicsResult result = env.adminClient().listTopics(options);
+
+            // Wait until the first attempt has been sent, then advance the time by the default api timeout
+            TestUtils.waitForCondition(() -> env.kafkaClient().hasInFlightRequests(),
+                    "Timed out waiting for Metadata request to be sent");
+            time.sleep(60001);
+
+            // The in-flight request should not be cancelled
+            assertTrue(env.kafkaClient().hasInFlightRequests());
+
+            // Now sleep the remaining time for the request timeout to expire
+            time.sleep(60000);
+            TestUtils.assertFutureThrows(result.future, TimeoutException.class);
+        }
+    }
+
+    private ClientQuotaEntity newClientQuotaEntity(String... args) {
+        assertTrue(args.length % 2 == 0);
+
+        Map entityMap = new HashMap<>(args.length / 2);
+        for (int index = 0; index < args.length; index += 2) {
+            entityMap.put(args[index], args[index + 1]);
+        }
+        return new ClientQuotaEntity(entityMap);
+    }
+
+    @Test
+    public void testDescribeClientQuotas() throws Exception {
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            final String value = "value";
+
+            Map> responseData = new HashMap<>();
+            ClientQuotaEntity entity1 = newClientQuotaEntity(ClientQuotaEntity.USER, "user-1", ClientQuotaEntity.CLIENT_ID, value);
+            ClientQuotaEntity entity2 = newClientQuotaEntity(ClientQuotaEntity.USER, "user-2", ClientQuotaEntity.CLIENT_ID, value);
+            responseData.put(entity1, Collections.singletonMap("consumer_byte_rate", 10000.0));
+            responseData.put(entity2, Collections.singletonMap("producer_byte_rate", 20000.0));
+
+            env.kafkaClient().prepareResponse(DescribeClientQuotasResponse.fromQuotaEntities(responseData, 0));
+
+            ClientQuotaFilter filter = ClientQuotaFilter.contains(asList(ClientQuotaFilterComponent.ofEntity(ClientQuotaEntity.USER, value)));
+
+            DescribeClientQuotasResult result = env.adminClient().describeClientQuotas(filter);
+            Map> resultData = result.entities().get();
+            assertEquals(resultData.size(), 2);
+            assertTrue(resultData.containsKey(entity1));
+            Map config1 = resultData.get(entity1);
+            assertEquals(config1.size(), 1);
+            assertEquals(config1.get("consumer_byte_rate"), 10000.0, 1e-6);
+            assertTrue(resultData.containsKey(entity2));
+            Map config2 = resultData.get(entity2);
+            assertEquals(config2.size(), 1);
+            assertEquals(config2.get("producer_byte_rate"), 20000.0, 1e-6);
+        }
+    }
+
+    @Test
+    public void testEqualsOfClientQuotaFilterComponent() {
+        assertEquals(ClientQuotaFilterComponent.ofDefaultEntity(ClientQuotaEntity.USER),
+            ClientQuotaFilterComponent.ofDefaultEntity(ClientQuotaEntity.USER));
+
+        assertEquals(ClientQuotaFilterComponent.ofEntityType(ClientQuotaEntity.USER),
+            ClientQuotaFilterComponent.ofEntityType(ClientQuotaEntity.USER));
+
+        // match = null is different from match = Empty
+        assertNotEquals(ClientQuotaFilterComponent.ofDefaultEntity(ClientQuotaEntity.USER),
+            ClientQuotaFilterComponent.ofEntityType(ClientQuotaEntity.USER));
+
+        assertEquals(ClientQuotaFilterComponent.ofEntity(ClientQuotaEntity.USER, "user"),
+            ClientQuotaFilterComponent.ofEntity(ClientQuotaEntity.USER, "user"));
+
+        assertNotEquals(ClientQuotaFilterComponent.ofEntity(ClientQuotaEntity.USER, "user"),
+            ClientQuotaFilterComponent.ofDefaultEntity(ClientQuotaEntity.USER));
+
+        assertNotEquals(ClientQuotaFilterComponent.ofEntity(ClientQuotaEntity.USER, "user"),
+            ClientQuotaFilterComponent.ofEntityType(ClientQuotaEntity.USER));
+    }
+
+    @Test
+    public void testAlterClientQuotas() throws Exception {
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            ClientQuotaEntity goodEntity = newClientQuotaEntity(ClientQuotaEntity.USER, "user-1");
+            ClientQuotaEntity unauthorizedEntity = newClientQuotaEntity(ClientQuotaEntity.USER, "user-0");
+            ClientQuotaEntity invalidEntity = newClientQuotaEntity("", "user-0");
+
+            Map responseData = new HashMap<>(2);
+            responseData.put(goodEntity, new ApiError(Errors.CLUSTER_AUTHORIZATION_FAILED, "Authorization failed"));
+            responseData.put(unauthorizedEntity, new ApiError(Errors.CLUSTER_AUTHORIZATION_FAILED, "Authorization failed"));
+            responseData.put(invalidEntity, new ApiError(Errors.INVALID_REQUEST, "Invalid quota entity"));
+
+            env.kafkaClient().prepareResponse(AlterClientQuotasResponse.fromQuotaEntities(responseData, 0));
+
+            List entries = new ArrayList<>(3);
+            entries.add(new ClientQuotaAlteration(goodEntity, Collections.singleton(new ClientQuotaAlteration.Op("consumer_byte_rate", 10000.0))));
+            entries.add(new ClientQuotaAlteration(unauthorizedEntity, Collections.singleton(new ClientQuotaAlteration.Op("producer_byte_rate", 10000.0))));
+            entries.add(new ClientQuotaAlteration(invalidEntity, Collections.singleton(new ClientQuotaAlteration.Op("producer_byte_rate", 100.0))));
+
+            AlterClientQuotasResult result = env.adminClient().alterClientQuotas(entries);
+            result.values().get(goodEntity);
+            TestUtils.assertFutureError(result.values().get(unauthorizedEntity), ClusterAuthorizationException.class);
+            TestUtils.assertFutureError(result.values().get(invalidEntity), InvalidRequestException.class);
+        }
+    }
+
+    @Test
+    public void testAlterReplicaLogDirsSuccess() throws Exception {
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
+            createAlterLogDirsResponse(env, env.cluster().nodeById(0), Errors.NONE, 0);
+            createAlterLogDirsResponse(env, env.cluster().nodeById(1), Errors.NONE, 0);
+
+            TopicPartitionReplica tpr0 = new TopicPartitionReplica("topic", 0, 0);
+            TopicPartitionReplica tpr1 = new TopicPartitionReplica("topic", 0, 1);
+
+            Map logDirs = new HashMap<>();
+            logDirs.put(tpr0, "/data0");
+            logDirs.put(tpr1, "/data1");
+            AlterReplicaLogDirsResult result = env.adminClient().alterReplicaLogDirs(logDirs);
+            assertNull(result.values().get(tpr0).get());
+            assertNull(result.values().get(tpr1).get());
+        }
+    }
+
+    @Test
+    public void testAlterReplicaLogDirsLogDirNotFound() throws Exception {
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
+            createAlterLogDirsResponse(env, env.cluster().nodeById(0), Errors.NONE, 0);
+            createAlterLogDirsResponse(env, env.cluster().nodeById(1), Errors.LOG_DIR_NOT_FOUND, 0);
+
+            TopicPartitionReplica tpr0 = new TopicPartitionReplica("topic", 0, 0);
+            TopicPartitionReplica tpr1 = new TopicPartitionReplica("topic", 0, 1);
+
+            Map logDirs = new HashMap<>();
+            logDirs.put(tpr0, "/data0");
+            logDirs.put(tpr1, "/data1");
+            AlterReplicaLogDirsResult result = env.adminClient().alterReplicaLogDirs(logDirs);
+            assertNull(result.values().get(tpr0).get());
+            TestUtils.assertFutureError(result.values().get(tpr1), LogDirNotFoundException.class);
+        }
+    }
+
+    @Test
+    public void testAlterReplicaLogDirsUnrequested() throws Exception {
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
+            createAlterLogDirsResponse(env, env.cluster().nodeById(0), Errors.NONE, 1, 2);
+
+            TopicPartitionReplica tpr1 = new TopicPartitionReplica("topic", 1, 0);
+
+            Map logDirs = new HashMap<>();
+            logDirs.put(tpr1, "/data1");
+            AlterReplicaLogDirsResult result = env.adminClient().alterReplicaLogDirs(logDirs);
+            assertNull(result.values().get(tpr1).get());
+        }
+    }
+
+    @Test
+    public void testAlterReplicaLogDirsPartialResponse() throws Exception {
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
+            createAlterLogDirsResponse(env, env.cluster().nodeById(0), Errors.NONE, 1);
+
+            TopicPartitionReplica tpr1 = new TopicPartitionReplica("topic", 1, 0);
+            TopicPartitionReplica tpr2 = new TopicPartitionReplica("topic", 2, 0);
+
+            Map logDirs = new HashMap<>();
+            logDirs.put(tpr1, "/data1");
+            logDirs.put(tpr2, "/data1");
+            AlterReplicaLogDirsResult result = env.adminClient().alterReplicaLogDirs(logDirs);
+            assertNull(result.values().get(tpr1).get());
+            TestUtils.assertFutureThrows(result.values().get(tpr2), ApiException.class);
+        }
+    }
+
+    @Test
+    public void testAlterReplicaLogDirsPartialFailure() throws Exception {
+        long defaultApiTimeout = 60000;
+        MockTime time = new MockTime();
+
+        try (AdminClientUnitTestEnv env = mockClientEnv(time, AdminClientConfig.RETRIES_CONFIG, "0")) {
+
+            // Provide only one prepared response from node 1
+            env.kafkaClient().prepareResponseFrom(
+                prepareAlterLogDirsResponse(Errors.NONE, "topic", 2),
+                env.cluster().nodeById(1));
+
+            TopicPartitionReplica tpr1 = new TopicPartitionReplica("topic", 1, 0);
+            TopicPartitionReplica tpr2 = new TopicPartitionReplica("topic", 2, 1);
+
+            Map logDirs = new HashMap<>();
+            logDirs.put(tpr1, "/data1");
+            logDirs.put(tpr2, "/data1");
+
+            AlterReplicaLogDirsResult result = env.adminClient().alterReplicaLogDirs(logDirs);
+
+            // Wait until the prepared attempt has been consumed
+            TestUtils.waitForCondition(() -> env.kafkaClient().numAwaitingResponses() == 0,
+                "Failed awaiting requests");
+
+            // Wait until the request is sent out
+            TestUtils.waitForCondition(() -> env.kafkaClient().inFlightRequestCount() == 1,
+                "Failed awaiting request");
+
+            // Advance time past the default api timeout to time out the inflight request
+            time.sleep(defaultApiTimeout + 1);
+
+            TestUtils.assertFutureThrows(result.values().get(tpr1), ApiException.class);
+            assertNull(result.values().get(tpr2).get());
+        }
+    }
+
+    @Test
+    public void testDescribeUserScramCredentials() throws Exception {
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            final String user0Name = "user0";
+            final ScramMechanism user0ScramMechanism0 = ScramMechanism.SCRAM_SHA_256;
+            final int user0Iterations0 = 4096;
+            final ScramMechanism user0ScramMechanism1 = ScramMechanism.SCRAM_SHA_512;
+            final int user0Iterations1 = 8192;
+
+            final CredentialInfo user0CredentialInfo0 = new CredentialInfo();
+            user0CredentialInfo0.setMechanism(user0ScramMechanism0.type());
+            user0CredentialInfo0.setIterations(user0Iterations0);
+            final CredentialInfo user0CredentialInfo1 = new CredentialInfo();
+            user0CredentialInfo1.setMechanism(user0ScramMechanism1.type());
+            user0CredentialInfo1.setIterations(user0Iterations1);
+
+            final String user1Name = "user1";
+            final ScramMechanism user1ScramMechanism = ScramMechanism.SCRAM_SHA_256;
+            final int user1Iterations = 4096;
+
+            final CredentialInfo user1CredentialInfo = new CredentialInfo();
+            user1CredentialInfo.setMechanism(user1ScramMechanism.type());
+            user1CredentialInfo.setIterations(user1Iterations);
+
+            final DescribeUserScramCredentialsResponseData responseData = new DescribeUserScramCredentialsResponseData();
+            responseData.setResults(Arrays.asList(
+                    new DescribeUserScramCredentialsResponseData.DescribeUserScramCredentialsResult()
+                            .setUser(user0Name)
+                            .setCredentialInfos(Arrays.asList(user0CredentialInfo0, user0CredentialInfo1)),
+                    new DescribeUserScramCredentialsResponseData.DescribeUserScramCredentialsResult()
+                            .setUser(user1Name)
+                            .setCredentialInfos(singletonList(user1CredentialInfo))));
+            final DescribeUserScramCredentialsResponse response = new DescribeUserScramCredentialsResponse(responseData);
+
+            final Set usersRequestedSet = new HashSet<>();
+            usersRequestedSet.add(user0Name);
+            usersRequestedSet.add(user1Name);
+
+            for (final List users : asList(null, new ArrayList(), asList(user0Name, null, user1Name))) {
+                env.kafkaClient().prepareResponse(response);
+
+                final DescribeUserScramCredentialsResult result = env.adminClient().describeUserScramCredentials(users);
+                final Map descriptionResults = result.all().get();
+                final KafkaFuture user0DescriptionFuture = result.description(user0Name);
+                final KafkaFuture user1DescriptionFuture = result.description(user1Name);
+
+                final Set usersDescribedFromUsersSet = new HashSet<>(result.users().get());
+                assertEquals(usersRequestedSet, usersDescribedFromUsersSet);
+
+                final Set usersDescribedFromMapKeySet = descriptionResults.keySet();
+                assertEquals(usersRequestedSet, usersDescribedFromMapKeySet);
+
+                final UserScramCredentialsDescription userScramCredentialsDescription0 = descriptionResults.get(user0Name);
+                assertEquals(user0Name, userScramCredentialsDescription0.name());
+                assertEquals(2, userScramCredentialsDescription0.credentialInfos().size());
+                assertEquals(user0ScramMechanism0, userScramCredentialsDescription0.credentialInfos().get(0).mechanism());
+                assertEquals(user0Iterations0, userScramCredentialsDescription0.credentialInfos().get(0).iterations());
+                assertEquals(user0ScramMechanism1, userScramCredentialsDescription0.credentialInfos().get(1).mechanism());
+                assertEquals(user0Iterations1, userScramCredentialsDescription0.credentialInfos().get(1).iterations());
+                assertEquals(userScramCredentialsDescription0, user0DescriptionFuture.get());
+
+                final UserScramCredentialsDescription userScramCredentialsDescription1 = descriptionResults.get(user1Name);
+                assertEquals(user1Name, userScramCredentialsDescription1.name());
+                assertEquals(1, userScramCredentialsDescription1.credentialInfos().size());
+                assertEquals(user1ScramMechanism, userScramCredentialsDescription1.credentialInfos().get(0).mechanism());
+                assertEquals(user1Iterations, userScramCredentialsDescription1.credentialInfos().get(0).iterations());
+                assertEquals(userScramCredentialsDescription1, user1DescriptionFuture.get());
+            }
+        }
+    }
+
+    @Test
+    public void testAlterUserScramCredentialsUnknownMechanism() throws Exception {
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            final String user0Name = "user0";
+            ScramMechanism user0ScramMechanism0 = ScramMechanism.UNKNOWN;
+
+            final String user1Name = "user1";
+            ScramMechanism user1ScramMechanism0 = ScramMechanism.UNKNOWN;
+
+            final String user2Name = "user2";
+            ScramMechanism user2ScramMechanism0 = ScramMechanism.SCRAM_SHA_256;
+
+            AlterUserScramCredentialsResponseData responseData = new AlterUserScramCredentialsResponseData();
+            responseData.setResults(Arrays.asList(
+                    new AlterUserScramCredentialsResponseData.AlterUserScramCredentialsResult().setUser(user2Name)));
+
+            env.kafkaClient().prepareResponse(new AlterUserScramCredentialsResponse(responseData));
+
+            AlterUserScramCredentialsResult result = env.adminClient().alterUserScramCredentials(Arrays.asList(
+                    new UserScramCredentialDeletion(user0Name, user0ScramMechanism0),
+                    new UserScramCredentialUpsertion(user1Name, new ScramCredentialInfo(user1ScramMechanism0, 8192), "password"),
+                    new UserScramCredentialUpsertion(user2Name, new ScramCredentialInfo(user2ScramMechanism0, 4096), "password")));
+            Map> resultData = result.values();
+            assertEquals(3, resultData.size());
+            Arrays.asList(user0Name, user1Name).stream().forEach(u -> {
+                assertTrue(resultData.containsKey(u));
+                try {
+                    resultData.get(u).get();
+                    fail("Expected request for user " + u + " to complete exceptionally, but it did not");
+                } catch (Exception expected) {
+                    // ignore
+                }
+            });
+            assertTrue(resultData.containsKey(user2Name));
+            try {
+                resultData.get(user2Name).get();
+            } catch (Exception e) {
+                fail("Expected request for user " + user2Name + " to NOT complete excdptionally, but it did");
+            }
+            try {
+                result.all().get();
+                fail("Expected 'result.all().get()' to throw an exception since at least one user failed, but it did not");
+            } catch (final Exception expected) {
+                // ignore, expected
+            }
+        }
+    }
+
+    @Test
+    public void testAlterUserScramCredentials() throws Exception {
+        try (AdminClientUnitTestEnv env = mockClientEnv()) {
+            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
+
+            final String user0Name = "user0";
+            ScramMechanism user0ScramMechanism0 = ScramMechanism.SCRAM_SHA_256;
+            ScramMechanism user0ScramMechanism1 = ScramMechanism.SCRAM_SHA_512;
+            final String user1Name = "user1";
+            ScramMechanism user1ScramMechanism0 = ScramMechanism.SCRAM_SHA_256;
+            final String user2Name = "user2";
+            ScramMechanism user2ScramMechanism0 = ScramMechanism.SCRAM_SHA_512;
+            AlterUserScramCredentialsResponseData responseData = new AlterUserScramCredentialsResponseData();
+            responseData.setResults(Arrays.asList(user0Name, user1Name, user2Name).stream().map(u ->
+                    new AlterUserScramCredentialsResponseData.AlterUserScramCredentialsResult()
+                    .setUser(u).setErrorCode(Errors.NONE.code())).collect(Collectors.toList()));
+
+            env.kafkaClient().prepareResponse(new AlterUserScramCredentialsResponse(responseData));
+
+            AlterUserScramCredentialsResult result = env.adminClient().alterUserScramCredentials(Arrays.asList(
+                    new UserScramCredentialDeletion(user0Name, user0ScramMechanism0),
+                    new UserScramCredentialUpsertion(user0Name, new ScramCredentialInfo(user0ScramMechanism1, 8192), "password"),
+                    new UserScramCredentialUpsertion(user1Name, new ScramCredentialInfo(user1ScramMechanism0, 8192), "password"),
+                    new UserScramCredentialDeletion(user2Name, user2ScramMechanism0)));
+            Map> resultData = result.values();
+            assertEquals(3, resultData.size());
+            Arrays.asList(user0Name, user1Name, user2Name).stream().forEach(u -> {
+                assertTrue(resultData.containsKey(u));
+                assertFalse(resultData.get(u).isCompletedExceptionally());
+            });
+        }
+    }
+
+    private void createAlterLogDirsResponse(AdminClientUnitTestEnv env, Node node, Errors error, int... partitions) {
+        env.kafkaClient().prepareResponseFrom(
+            prepareAlterLogDirsResponse(error, "topic", partitions), node);
+    }
+
+    private AlterReplicaLogDirsResponse prepareAlterLogDirsResponse(Errors error, String topic, int... partitions) {
+        return new AlterReplicaLogDirsResponse(
+            new AlterReplicaLogDirsResponseData().setResults(singletonList(
+                new AlterReplicaLogDirTopicResult()
+                    .setTopicName(topic)
+                    .setPartitions(Arrays.stream(partitions).boxed().map(partitionId ->
+                        new AlterReplicaLogDirPartitionResult()
+                            .setPartitionIndex(partitionId)
+                            .setErrorCode(error.code())).collect(Collectors.toList())))));
+    }
+
+    @Test
+    public void testDescribeLogDirsPartialFailure() throws Exception {
+        long defaultApiTimeout = 60000;
+        MockTime time = new MockTime();
+
+        try (AdminClientUnitTestEnv env = mockClientEnv(time, AdminClientConfig.RETRIES_CONFIG, "0")) {
+
+            env.kafkaClient().prepareResponseFrom(
+                prepareDescribeLogDirsResponse(Errors.NONE, "/data"),
+                env.cluster().nodeById(1));
+
+            DescribeLogDirsResult result = env.adminClient().describeLogDirs(Arrays.asList(0, 1));
+
+            // Wait until the prepared attempt has been consumed
+            TestUtils.waitForCondition(() -> env.kafkaClient().numAwaitingResponses() == 0,
+                "Failed awaiting requests");
+
+            // Wait until the request is sent out
+            TestUtils.waitForCondition(() -> env.kafkaClient().inFlightRequestCount() == 1,
+                "Failed awaiting request");
+
+            // Advance time past the default api timeout to time out the inflight request
+            time.sleep(defaultApiTimeout + 1);
+
+            TestUtils.assertFutureThrows(result.descriptions().get(0), ApiException.class);
+            assertNotNull(result.descriptions().get(1).get());
+        }
+    }
+
+    @Test
+    public void testHasCoordinatorMoved() {
+        Map errors = new HashMap<>();
+        AbstractResponse response = new AbstractResponse(ApiKeys.OFFSET_COMMIT) {
+            @Override
+            public Map errorCounts() {
+                return errors;
+            }
+
+            @Override
+            protected Message data() {
+                return null;
+            }
+
+            @Override
+            public int throttleTimeMs() {
+                return DEFAULT_THROTTLE_TIME;
+            }
+        };
+
+        assertFalse(ConsumerGroupOperationContext.hasCoordinatorMoved(response));
+
+        errors.put(Errors.NOT_COORDINATOR, 1);
+        assertTrue(ConsumerGroupOperationContext.hasCoordinatorMoved(response));
+    }
+
+    @Test
+    public void testShouldRefreshCoordinator() {
+        Map errorCounts = new HashMap<>();
+
+        assertFalse(ConsumerGroupOperationContext.shouldRefreshCoordinator(errorCounts));
+
+        errorCounts.put(Errors.COORDINATOR_LOAD_IN_PROGRESS, 1);
+        assertTrue(ConsumerGroupOperationContext.shouldRefreshCoordinator(errorCounts));
+
+        errorCounts.clear();
+        errorCounts.put(Errors.COORDINATOR_NOT_AVAILABLE, 1);
+        assertTrue(ConsumerGroupOperationContext.shouldRefreshCoordinator(errorCounts));
+    }
+
+    private DescribeLogDirsResponse prepareDescribeLogDirsResponse(Errors error, String logDir) {
+        return new DescribeLogDirsResponse(new DescribeLogDirsResponseData()
+            .setResults(Collections.singletonList(
+                new DescribeLogDirsResponseData.DescribeLogDirsResult()
+                    .setErrorCode(error.code())
+                    .setLogDir(logDir))));
+    }
+
+    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) {
+            assertTrue("Did not find " + element, collection.contains(element));
+        }
+        assertEquals("There are unexpected extra elements in the collection.",
+            elements.length, collection.size());
+    }
+
+    public static KafkaAdminClient createInternal(AdminClientConfig config, KafkaAdminClient.TimeoutProcessorFactory timeoutProcessorFactory) {
+        return KafkaAdminClient.createInternal(config, timeoutProcessorFactory);
+    }
+
+    public static class FailureInjectingTimeoutProcessorFactory extends KafkaAdminClient.TimeoutProcessorFactory {
+
+        private int numTries = 0;
 
         private int failuresInjected = 0;
-        
+
         @Override
         public KafkaAdminClient.TimeoutProcessor create(long now) {
             return new FailureInjectingTimeoutProcessor(now);
@@ -571,7 +5026,7 @@ public FailureInjectingTimeoutProcessor(long now) {
             }
 
             boolean callHasExpired(KafkaAdminClient.Call call) {
-                if (shouldInjectFailure()) {
+                if ((!call.isInternal()) && shouldInjectFailure()) {
                     log.debug("Injecting timeout for {}.", call);
                     return true;
                 } else {
@@ -581,7 +5036,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
new file mode 100644
index 0000000000000..3e9f605923bb7
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/admin/MockAdminClient.java
@@ -0,0 +1,906 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.clients.admin.DescribeReplicaLogDirsResult.ReplicaLogDirInfo;
+import org.apache.kafka.common.KafkaFuture;
+import org.apache.kafka.common.Metric;
+import org.apache.kafka.common.MetricName;
+import org.apache.kafka.common.Node;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.clients.consumer.OffsetAndMetadata;
+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.InvalidRequestException;
+import org.apache.kafka.common.errors.InvalidReplicationFactorException;
+import org.apache.kafka.common.errors.KafkaStorageException;
+import org.apache.kafka.common.errors.ReplicaNotAvailableException;
+import org.apache.kafka.common.errors.TimeoutException;
+import org.apache.kafka.common.errors.TopicExistsException;
+import org.apache.kafka.common.errors.UnknownTopicOrPartitionException;
+import org.apache.kafka.common.internals.KafkaFutureImpl;
+import org.apache.kafka.common.requests.DescribeLogDirsResponse;
+import org.apache.kafka.common.quota.ClientQuotaAlteration;
+import org.apache.kafka.common.quota.ClientQuotaFilter;
+
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+
+public class MockAdminClient extends AdminClient {
+    public static final String DEFAULT_CLUSTER_ID = "I4ZmrWqfT2e-upky_4fdPA";
+
+    public static final List DEFAULT_LOG_DIRS =
+        Collections.singletonList("/tmp/kafka-logs");
+
+    private final List brokers;
+    private final Map allTopics = new HashMap<>();
+    private final Map reassignments =
+        new HashMap<>();
+    private final Map replicaMoves =
+        new HashMap<>();
+    private final Map beginningOffsets;
+    private final Map endOffsets;
+    private final String clusterId;
+    private final List> brokerLogDirs;
+    private final List> brokerConfigs;
+
+    private Node controller;
+    private int timeoutNextRequests = 0;
+    private final int defaultPartitions;
+    private final int defaultReplicationFactor;
+
+    private Map mockMetrics = new HashMap<>();
+
+    public static Builder create() {
+        return new Builder();
+    }
+
+    public static class Builder {
+        private String clusterId = DEFAULT_CLUSTER_ID;
+        private List brokers = new ArrayList<>();
+        private Node controller = null;
+        private List> brokerLogDirs = new ArrayList<>();
+        private Short defaultPartitions;
+        private Integer defaultReplicationFactor;
+
+        public Builder() {
+            numBrokers(1);
+        }
+
+        public Builder clusterId(String clusterId) {
+            this.clusterId = clusterId;
+            return this;
+        }
+
+        public Builder brokers(List brokers) {
+            numBrokers(brokers.size());
+            this.brokers = brokers;
+            return this;
+        }
+
+        public Builder numBrokers(int numBrokers) {
+            if (brokers.size() >= numBrokers) {
+                brokers = brokers.subList(0, numBrokers);
+                brokerLogDirs = brokerLogDirs.subList(0, numBrokers);
+            } else {
+                for (int id = brokers.size(); id < numBrokers; id++) {
+                    brokers.add(new Node(id, "localhost", 1000 + id));
+                    brokerLogDirs.add(DEFAULT_LOG_DIRS);
+                }
+            }
+            return this;
+        }
+
+        public Builder controller(int index) {
+            this.controller = brokers.get(index);
+            return this;
+        }
+
+        public Builder brokerLogDirs(List> brokerLogDirs) {
+            this.brokerLogDirs = brokerLogDirs;
+            return this;
+        }
+
+        public Builder defaultReplicationFactor(int defaultReplicationFactor) {
+            this.defaultReplicationFactor = defaultReplicationFactor;
+            return this;
+        }
+
+        public Builder defaultPartitions(short numPartitions) {
+            this.defaultPartitions = numPartitions;
+            return this;
+        }
+
+        public MockAdminClient build() {
+            return new MockAdminClient(brokers,
+                controller == null ? brokers.get(0) : controller,
+                clusterId,
+                defaultPartitions != null ? defaultPartitions.shortValue() : 1,
+                defaultReplicationFactor != null ? defaultReplicationFactor.shortValue() : Math.min(brokers.size(), 3),
+                brokerLogDirs);
+        }
+    }
+
+    public MockAdminClient() {
+        this(Collections.singletonList(Node.noNode()), Node.noNode());
+    }
+
+    public MockAdminClient(List brokers, Node controller) {
+        this(brokers, controller, DEFAULT_CLUSTER_ID, 1, brokers.size(),
+            Collections.nCopies(brokers.size(), DEFAULT_LOG_DIRS));
+    }
+
+    private MockAdminClient(List brokers,
+                            Node controller,
+                            String clusterId,
+                            int defaultPartitions,
+                            int defaultReplicationFactor,
+                            List> brokerLogDirs) {
+        this.brokers = brokers;
+        controller(controller);
+        this.clusterId = clusterId;
+        this.defaultPartitions = defaultPartitions;
+        this.defaultReplicationFactor = defaultReplicationFactor;
+        this.brokerLogDirs = brokerLogDirs;
+        this.brokerConfigs = new ArrayList<>();
+        for (int i = 0; i < brokers.size(); i++) {
+            this.brokerConfigs.add(new HashMap<>());
+        }
+        this.beginningOffsets = new HashMap<>();
+        this.endOffsets = new HashMap<>();
+    }
+
+    synchronized public void controller(Node controller) {
+        if (!brokers.contains(controller))
+            throw new IllegalArgumentException("The controller node must be in the list of brokers");
+        this.controller = controller;
+    }
+
+    synchronized public void addTopic(boolean internal,
+                         String name,
+                         List partitions,
+                         Map configs) {
+        if (allTopics.containsKey(name)) {
+            throw new IllegalArgumentException(String.format("Topic %s was already added.", name));
+        }
+        for (TopicPartitionInfo partition : partitions) {
+            if (!brokers.contains(partition.leader())) {
+                throw new IllegalArgumentException("Leader broker unknown");
+            }
+            if (!brokers.containsAll(partition.replicas())) {
+                throw new IllegalArgumentException("Unknown brokers in replica list");
+            }
+            if (!brokers.containsAll(partition.isr())) {
+                throw new IllegalArgumentException("Unknown brokers in isr list");
+            }
+        }
+        ArrayList logDirs = new ArrayList<>();
+        for (TopicPartitionInfo partition : partitions) {
+            if (partition.leader() != null) {
+                logDirs.add(brokerLogDirs.get(partition.leader().id()).get(0));
+            }
+        }
+        allTopics.put(name, new TopicMetadata(internal, partitions, logDirs, configs));
+    }
+
+    synchronized public void markTopicForDeletion(final String name) {
+        if (!allTopics.containsKey(name)) {
+            throw new IllegalArgumentException(String.format("Topic %s did not exist.", name));
+        }
+
+        allTopics.get(name).markedForDeletion = true;
+    }
+
+    synchronized public void timeoutNextRequest(int numberOfRequest) {
+        timeoutNextRequests = numberOfRequest;
+    }
+
+    @Override
+    synchronized 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, authorizedOperationsFuture);
+    }
+
+    @Override
+    synchronized public CreateTopicsResult createTopics(Collection newTopics, CreateTopicsOptions options) {
+        Map> createTopicResult = new HashMap<>();
+
+        if (timeoutNextRequests > 0) {
+            for (final NewTopic newTopic : newTopics) {
+                String topicName = newTopic.name();
+
+                KafkaFutureImpl future = new KafkaFutureImpl<>();
+                future.completeExceptionally(new TimeoutException());
+                createTopicResult.put(topicName, future);
+            }
+
+            --timeoutNextRequests;
+            return new CreateTopicsResult(createTopicResult);
+        }
+
+        for (final NewTopic newTopic : newTopics) {
+            KafkaFutureImpl future = new KafkaFutureImpl<>();
+
+            String topicName = newTopic.name();
+            if (allTopics.containsKey(topicName)) {
+                future.completeExceptionally(new TopicExistsException(String.format("Topic %s exists already.", topicName)));
+                createTopicResult.put(topicName, future);
+                continue;
+            }
+            int replicationFactor = newTopic.replicationFactor();
+            if (replicationFactor == -1) {
+                replicationFactor = defaultReplicationFactor;
+            }
+            if (replicationFactor > brokers.size()) {
+                future.completeExceptionally(new InvalidReplicationFactorException(
+                        String.format("Replication factor: %d is larger than brokers: %d", newTopic.replicationFactor(), brokers.size())));
+                createTopicResult.put(topicName, future);
+                continue;
+            }
+
+            List replicas = new ArrayList<>(replicationFactor);
+            for (int i = 0; i < replicationFactor; ++i) {
+                replicas.add(brokers.get(i));
+            }
+
+            int numberOfPartitions = newTopic.numPartitions();
+            if (numberOfPartitions == -1) {
+                numberOfPartitions = defaultPartitions;
+            }
+            List partitions = new ArrayList<>(numberOfPartitions);
+            // Partitions start off on the first log directory of each broker, for now.
+            List logDirs = new ArrayList<>(numberOfPartitions);
+            for (int i = 0; i < numberOfPartitions; i++) {
+                partitions.add(new TopicPartitionInfo(i, brokers.get(0), replicas, Collections.emptyList()));
+                logDirs.add(brokerLogDirs.get(partitions.get(i).leader().id()).get(0));
+            }
+            allTopics.put(topicName, new TopicMetadata(false, partitions, logDirs, newTopic.configs()));
+            future.complete(null);
+            createTopicResult.put(topicName, future);
+        }
+
+        return new CreateTopicsResult(createTopicResult);
+    }
+
+    @Override
+    synchronized public ListTopicsResult listTopics(ListTopicsOptions options) {
+        Map topicListings = new HashMap<>();
+
+        if (timeoutNextRequests > 0) {
+            KafkaFutureImpl> future = new KafkaFutureImpl<>();
+            future.completeExceptionally(new TimeoutException());
+
+            --timeoutNextRequests;
+            return new ListTopicsResult(future);
+        }
+
+        for (Map.Entry topicDescription : allTopics.entrySet()) {
+            String topicName = topicDescription.getKey();
+            if (topicDescription.getValue().fetchesRemainingUntilVisible > 0) {
+                topicDescription.getValue().fetchesRemainingUntilVisible--;
+            } else {
+                topicListings.put(topicName, new TopicListing(topicName, topicDescription.getValue().isInternalTopic));
+            }
+        }
+
+        KafkaFutureImpl> future = new KafkaFutureImpl<>();
+        future.complete(topicListings);
+        return new ListTopicsResult(future);
+    }
+
+    @Override
+    synchronized public DescribeTopicsResult describeTopics(Collection topicNames, DescribeTopicsOptions options) {
+        Map> topicDescriptions = new HashMap<>();
+
+        if (timeoutNextRequests > 0) {
+            for (String requestedTopic : topicNames) {
+                KafkaFutureImpl future = new KafkaFutureImpl<>();
+                future.completeExceptionally(new TimeoutException());
+                topicDescriptions.put(requestedTopic, future);
+            }
+
+            --timeoutNextRequests;
+            return new DescribeTopicsResult(topicDescriptions);
+        }
+
+        for (String requestedTopic : topicNames) {
+            for (Map.Entry topicDescription : allTopics.entrySet()) {
+                String topicName = topicDescription.getKey();
+                if (topicName.equals(requestedTopic) && !topicDescription.getValue().markedForDeletion) {
+                    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)) {
+                KafkaFutureImpl future = new KafkaFutureImpl<>();
+                future.completeExceptionally(new UnknownTopicOrPartitionException("Topic " + requestedTopic + " not found."));
+                topicDescriptions.put(requestedTopic, future);
+            }
+        }
+
+        return new DescribeTopicsResult(topicDescriptions);
+    }
+
+    @Override
+    synchronized public DeleteTopicsResult deleteTopics(Collection topicsToDelete, DeleteTopicsOptions options) {
+        Map> deleteTopicsResult = new HashMap<>();
+
+        if (timeoutNextRequests > 0) {
+            for (final String topicName : topicsToDelete) {
+                KafkaFutureImpl future = new KafkaFutureImpl<>();
+                future.completeExceptionally(new TimeoutException());
+                deleteTopicsResult.put(topicName, future);
+            }
+
+            --timeoutNextRequests;
+            return new DeleteTopicsResult(deleteTopicsResult);
+        }
+
+        for (final String topicName : topicsToDelete) {
+            KafkaFutureImpl future = new KafkaFutureImpl<>();
+
+            if (allTopics.remove(topicName) == null) {
+                future.completeExceptionally(new UnknownTopicOrPartitionException(String.format("Topic %s does not exist.", topicName)));
+            } else {
+                future.complete(null);
+            }
+            deleteTopicsResult.put(topicName, future);
+        }
+
+        return new DeleteTopicsResult(deleteTopicsResult);
+    }
+
+    @Override
+    synchronized public CreatePartitionsResult createPartitions(Map newPartitions, CreatePartitionsOptions options) {
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
+    @Override
+    synchronized public DeleteRecordsResult deleteRecords(Map recordsToDelete, DeleteRecordsOptions options) {
+        Map> deletedRecordsResult = new HashMap<>();
+        if (recordsToDelete.isEmpty()) {
+            return new DeleteRecordsResult(deletedRecordsResult);
+        } else {
+            throw new UnsupportedOperationException("Not implemented yet");
+        }
+    }
+
+    @Override
+    synchronized public CreateDelegationTokenResult createDelegationToken(CreateDelegationTokenOptions options) {
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
+    @Override
+    synchronized public RenewDelegationTokenResult renewDelegationToken(byte[] hmac, RenewDelegationTokenOptions options) {
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
+    @Override
+    synchronized public ExpireDelegationTokenResult expireDelegationToken(byte[] hmac, ExpireDelegationTokenOptions options) {
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
+    @Override
+    synchronized public DescribeDelegationTokenResult describeDelegationToken(DescribeDelegationTokenOptions options) {
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
+    @Override
+    synchronized public DescribeConsumerGroupsResult describeConsumerGroups(Collection groupIds, DescribeConsumerGroupsOptions options) {
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
+    @Override
+    synchronized public ListConsumerGroupsResult listConsumerGroups(ListConsumerGroupsOptions options) {
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
+    @Override
+    synchronized public ListConsumerGroupOffsetsResult listConsumerGroupOffsets(String groupId, ListConsumerGroupOffsetsOptions options) {
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
+    @Override
+    synchronized public DeleteConsumerGroupsResult deleteConsumerGroups(Collection groupIds, DeleteConsumerGroupsOptions options) {
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
+    @Override
+    synchronized public DeleteConsumerGroupOffsetsResult deleteConsumerGroupOffsets(String groupId, Set partitions, DeleteConsumerGroupOffsetsOptions options) {
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
+    @Deprecated
+    @Override
+    synchronized public ElectPreferredLeadersResult electPreferredLeaders(Collection partitions, ElectPreferredLeadersOptions options) {
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
+    @Override
+    synchronized public ElectLeadersResult electLeaders(
+            ElectionType electionType,
+            Set partitions,
+            ElectLeadersOptions options) {
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
+    @Override
+    synchronized public RemoveMembersFromConsumerGroupResult removeMembersFromConsumerGroup(String groupId, RemoveMembersFromConsumerGroupOptions options) {
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
+    @Override
+    synchronized public CreateAclsResult createAcls(Collection acls, CreateAclsOptions options) {
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
+    @Override
+    synchronized public DescribeAclsResult describeAcls(AclBindingFilter filter, DescribeAclsOptions options) {
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
+    @Override
+    synchronized public DeleteAclsResult deleteAcls(Collection filters, DeleteAclsOptions options) {
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
+    @Override
+    synchronized public DescribeConfigsResult describeConfigs(Collection resources, DescribeConfigsOptions options) {
+
+        if (timeoutNextRequests > 0) {
+            Map> configs = new HashMap<>();
+            for (ConfigResource requestedResource : resources) {
+                KafkaFutureImpl future = new KafkaFutureImpl<>();
+                future.completeExceptionally(new TimeoutException());
+                configs.put(requestedResource, future);
+            }
+
+            --timeoutNextRequests;
+            return new DescribeConfigsResult(configs);
+        }
+
+        Map> results = new HashMap<>();
+        for (ConfigResource resource : resources) {
+            KafkaFutureImpl future = new KafkaFutureImpl<>();
+            results.put(resource, future);
+            try {
+                future.complete(getResourceDescription(resource));
+            } catch (Throwable e) {
+                future.completeExceptionally(e);
+            }
+        }
+        return new DescribeConfigsResult(results);
+    }
+
+    synchronized private Config getResourceDescription(ConfigResource resource) {
+        switch (resource.type()) {
+            case BROKER: {
+                int brokerId = Integer.parseInt(resource.name());
+                if (brokerId >= brokerConfigs.size()) {
+                    throw new InvalidRequestException("Broker " + resource.name() +
+                        " not found.");
+                }
+                return toConfigObject(brokerConfigs.get(brokerId));
+            }
+            case TOPIC: {
+                TopicMetadata topicMetadata = allTopics.get(resource.name());
+                if (topicMetadata != null && !topicMetadata.markedForDeletion) {
+                    if (topicMetadata.fetchesRemainingUntilVisible > 0)
+                        topicMetadata.fetchesRemainingUntilVisible = Math.max(0, topicMetadata.fetchesRemainingUntilVisible - 1);
+                    else return toConfigObject(topicMetadata.configs);
+
+                }
+                throw new UnknownTopicOrPartitionException("Resource " + resource + " not found.");
+            }
+            default:
+                throw new UnsupportedOperationException("Not implemented yet");
+        }
+    }
+
+    private static Config toConfigObject(Map map) {
+        List configEntries = new ArrayList<>();
+        for (Map.Entry entry : map.entrySet()) {
+            configEntries.add(new ConfigEntry(entry.getKey(), entry.getValue()));
+        }
+        return new Config(configEntries);
+    }
+
+    @Override
+    @Deprecated
+    synchronized public AlterConfigsResult alterConfigs(Map configs, AlterConfigsOptions options) {
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
+    @Override
+    synchronized public AlterConfigsResult incrementalAlterConfigs(
+            Map> configs,
+            AlterConfigsOptions options) {
+        Map> futures = new HashMap<>();
+        for (Map.Entry> entry :
+                configs.entrySet()) {
+            ConfigResource resource = entry.getKey();
+            KafkaFutureImpl future = new KafkaFutureImpl<>();
+            futures.put(resource, future);
+            Throwable throwable =
+                handleIncrementalResourceAlteration(resource, entry.getValue());
+            if (throwable == null) {
+                future.complete(null);
+            } else {
+                future.completeExceptionally(throwable);
+            }
+        }
+        return new AlterConfigsResult(futures);
+    }
+
+    synchronized private Throwable handleIncrementalResourceAlteration(
+            ConfigResource resource, Collection ops) {
+        switch (resource.type()) {
+            case BROKER: {
+                int brokerId;
+                try {
+                    brokerId = Integer.valueOf(resource.name());
+                } catch (NumberFormatException e) {
+                    return e;
+                }
+                if (brokerId >= brokerConfigs.size()) {
+                    return new InvalidRequestException("no such broker as " + brokerId);
+                }
+                HashMap newMap = new HashMap<>(brokerConfigs.get(brokerId));
+                for (AlterConfigOp op : ops) {
+                    switch (op.opType()) {
+                        case SET:
+                            newMap.put(op.configEntry().name(), op.configEntry().value());
+                            break;
+                        case DELETE:
+                            newMap.remove(op.configEntry().name());
+                            break;
+                        default:
+                            return new InvalidRequestException(
+                                "Unsupported op type " + op.opType());
+                    }
+                }
+                brokerConfigs.set(brokerId, newMap);
+                return null;
+            }
+            case TOPIC: {
+                TopicMetadata topicMetadata = allTopics.get(resource.name());
+                if (topicMetadata == null) {
+                    return new UnknownTopicOrPartitionException("No such topic as " +
+                        resource.name());
+                }
+                HashMap newMap = new HashMap<>(topicMetadata.configs);
+                for (AlterConfigOp op : ops) {
+                    switch (op.opType()) {
+                        case SET:
+                            newMap.put(op.configEntry().name(), op.configEntry().value());
+                            break;
+                        case DELETE:
+                            newMap.remove(op.configEntry().name());
+                            break;
+                        default:
+                            return new InvalidRequestException(
+                                "Unsupported op type " + op.opType());
+                    }
+                }
+                topicMetadata.configs = newMap;
+                return null;
+            }
+            default:
+                return new UnsupportedOperationException();
+        }
+    }
+
+    @Override
+    synchronized public AlterReplicaLogDirsResult alterReplicaLogDirs(
+            Map replicaAssignment,
+            AlterReplicaLogDirsOptions options) {
+        Map> results = new HashMap<>();
+        for (Map.Entry entry : replicaAssignment.entrySet()) {
+            TopicPartitionReplica replica = entry.getKey();
+            String newLogDir = entry.getValue();
+            KafkaFutureImpl future = new KafkaFutureImpl<>();
+            results.put(replica, future);
+            List dirs = brokerLogDirs.get(replica.brokerId());
+            if (dirs == null) {
+                future.completeExceptionally(
+                    new ReplicaNotAvailableException("Can't find " + replica));
+            } else if (!dirs.contains(newLogDir)) {
+                future.completeExceptionally(
+                    new KafkaStorageException("Log directory " + newLogDir + " is offline"));
+            } else {
+                TopicMetadata metadata = allTopics.get(replica.topic());
+                if (metadata == null || metadata.partitions.size() <= replica.partition()) {
+                    future.completeExceptionally(
+                        new ReplicaNotAvailableException("Can't find " + replica));
+                } else {
+                    String currentLogDir = metadata.partitionLogDirs.get(replica.partition());
+                    replicaMoves.put(replica,
+                        new ReplicaLogDirInfo(currentLogDir, 0, newLogDir, 0));
+                    future.complete(null);
+                }
+            }
+        }
+        return new AlterReplicaLogDirsResult(results);
+    }
+
+    @Override
+    synchronized public DescribeLogDirsResult describeLogDirs(Collection brokers,
+                                                              DescribeLogDirsOptions options) {
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
+    @Override
+    synchronized public DescribeReplicaLogDirsResult describeReplicaLogDirs(
+            Collection replicas, DescribeReplicaLogDirsOptions options) {
+        Map> results = new HashMap<>();
+        for (TopicPartitionReplica replica : replicas) {
+            TopicMetadata topicMetadata = allTopics.get(replica.topic());
+            if (topicMetadata != null) {
+                KafkaFutureImpl future = new KafkaFutureImpl<>();
+                results.put(replica, future);
+                String currentLogDir = currentLogDir(replica);
+                if (currentLogDir == null) {
+                    future.complete(new ReplicaLogDirInfo(null,
+                        DescribeLogDirsResponse.INVALID_OFFSET_LAG,
+                        null,
+                        DescribeLogDirsResponse.INVALID_OFFSET_LAG));
+                } else {
+                    ReplicaLogDirInfo info = replicaMoves.get(replica);
+                    if (info == null) {
+                        future.complete(new ReplicaLogDirInfo(currentLogDir, 0, null, 0));
+                    } else {
+                        future.complete(info);
+                    }
+                }
+            }
+        }
+        return new DescribeReplicaLogDirsResult(results);
+    }
+
+    private synchronized String currentLogDir(TopicPartitionReplica replica) {
+        TopicMetadata topicMetadata = allTopics.get(replica.topic());
+        if (topicMetadata == null) {
+            return null;
+        }
+        if (topicMetadata.partitionLogDirs.size() <= replica.partition()) {
+            return null;
+        }
+        return topicMetadata.partitionLogDirs.get(replica.partition());
+    }
+
+    @Override
+    synchronized public AlterPartitionReassignmentsResult alterPartitionReassignments(
+            Map> newReassignments,
+            AlterPartitionReassignmentsOptions options) {
+        Map> futures = new HashMap<>();
+        for (Map.Entry> entry :
+                newReassignments.entrySet()) {
+            TopicPartition partition = entry.getKey();
+            Optional newReassignment = entry.getValue();
+            KafkaFutureImpl future = new KafkaFutureImpl();
+            futures.put(partition, future);
+            TopicMetadata topicMetadata = allTopics.get(partition.topic());
+            if (partition.partition() < 0 ||
+                    topicMetadata == null ||
+                    topicMetadata.partitions.size() <= partition.partition()) {
+                future.completeExceptionally(new UnknownTopicOrPartitionException());
+            } else if (newReassignment.isPresent()) {
+                reassignments.put(partition, newReassignment.get());
+                future.complete(null);
+            } else {
+                reassignments.remove(partition);
+                future.complete(null);
+            }
+        }
+        return new AlterPartitionReassignmentsResult(futures);
+    }
+
+    @Override
+    synchronized public ListPartitionReassignmentsResult listPartitionReassignments(
+            Optional> partitions,
+            ListPartitionReassignmentsOptions options) {
+        Map map = new HashMap<>();
+        for (TopicPartition partition : partitions.isPresent() ?
+                partitions.get() : reassignments.keySet()) {
+            PartitionReassignment reassignment = findPartitionReassignment(partition);
+            if (reassignment != null) {
+                map.put(partition, reassignment);
+            }
+        }
+        return new ListPartitionReassignmentsResult(KafkaFutureImpl.completedFuture(map));
+    }
+
+    synchronized private PartitionReassignment findPartitionReassignment(TopicPartition partition) {
+        NewPartitionReassignment reassignment = reassignments.get(partition);
+        if (reassignment == null) {
+            return null;
+        }
+        TopicMetadata metadata = allTopics.get(partition.topic());
+        if (metadata == null) {
+            throw new RuntimeException("Internal MockAdminClient logic error: found " +
+                "reassignment for " + partition + ", but no TopicMetadata");
+        }
+        TopicPartitionInfo info = metadata.partitions.get(partition.partition());
+        if (info == null) {
+            throw new RuntimeException("Internal MockAdminClient logic error: found " +
+                "reassignment for " + partition + ", but no TopicPartitionInfo");
+        }
+        List replicas = new ArrayList<>();
+        List removingReplicas = new ArrayList<>();
+        List addingReplicas = new ArrayList<>(reassignment.targetReplicas());
+        for (Node node : info.replicas()) {
+            replicas.add(node.id());
+            if (!reassignment.targetReplicas().contains(node.id())) {
+                removingReplicas.add(node.id());
+            }
+            addingReplicas.remove(Integer.valueOf(node.id()));
+        }
+        return new PartitionReassignment(replicas, addingReplicas, removingReplicas);
+    }
+
+    @Override
+    synchronized public AlterConsumerGroupOffsetsResult alterConsumerGroupOffsets(String groupId, Map offsets, AlterConsumerGroupOffsetsOptions options) {
+        throw new UnsupportedOperationException("Not implement yet");
+    }
+
+    @Override
+    synchronized public ListOffsetsResult listOffsets(Map topicPartitionOffsets, ListOffsetsOptions options) {
+        Map> futures = new HashMap<>();
+
+        for (Map.Entry entry : topicPartitionOffsets.entrySet()) {
+            TopicPartition tp = entry.getKey();
+            OffsetSpec spec = entry.getValue();
+            KafkaFutureImpl future = new KafkaFutureImpl<>();
+
+            if (spec instanceof OffsetSpec.TimestampSpec)
+                throw new UnsupportedOperationException("Not implement yet");
+            else if (spec instanceof OffsetSpec.EarliestSpec)
+                future.complete(new ListOffsetsResult.ListOffsetsResultInfo(beginningOffsets.get(tp), -1, Optional.empty()));
+            else
+                future.complete(new ListOffsetsResult.ListOffsetsResultInfo(endOffsets.get(tp), -1, Optional.empty()));
+
+            futures.put(tp, future);
+        }
+
+        return new ListOffsetsResult(futures);
+    }
+
+    @Override
+    public DescribeClientQuotasResult describeClientQuotas(ClientQuotaFilter filter, DescribeClientQuotasOptions options) {
+        throw new UnsupportedOperationException("Not implement yet");
+    }
+
+    @Override
+    public AlterClientQuotasResult alterClientQuotas(Collection entries, AlterClientQuotasOptions options) {
+        throw new UnsupportedOperationException("Not implement yet");
+    }
+
+    @Override
+    public DescribeUserScramCredentialsResult describeUserScramCredentials(List users, DescribeUserScramCredentialsOptions options) {
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
+    @Override
+    public AlterUserScramCredentialsResult alterUserScramCredentials(List alterations, AlterUserScramCredentialsOptions options) {
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
+    @Override
+    public DescribeFeaturesResult describeFeatures(DescribeFeaturesOptions options) {
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
+    @Override
+    public UpdateFeaturesResult updateFeatures(Map featureUpdates, UpdateFeaturesOptions options) {
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
+    @Override
+    synchronized public void close(Duration timeout) {}
+
+    public synchronized void updateBeginningOffsets(Map newOffsets) {
+        beginningOffsets.putAll(newOffsets);
+    }
+
+    public synchronized void updateEndOffsets(final Map newOffsets) {
+        endOffsets.putAll(newOffsets);
+    }
+
+    private final static class TopicMetadata {
+        final boolean isInternalTopic;
+        final List partitions;
+        final List partitionLogDirs;
+        Map configs;
+        int fetchesRemainingUntilVisible;
+
+        public boolean markedForDeletion;
+
+        TopicMetadata(boolean isInternalTopic,
+                      List partitions,
+                      List partitionLogDirs,
+                      Map configs) {
+            this.isInternalTopic = isInternalTopic;
+            this.partitions = partitions;
+            this.partitionLogDirs = partitionLogDirs;
+            this.configs = configs != null ? configs : Collections.emptyMap();
+            this.markedForDeletion = false;
+            this.fetchesRemainingUntilVisible = 0;
+        }
+    }
+
+    synchronized public void setMockMetrics(MetricName name, Metric metric) {
+        mockMetrics.put(name, metric);
+    }
+
+    @Override
+    synchronized public Map metrics() {
+        return mockMetrics;
+    }
+
+    synchronized 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;
+    }
+
+    synchronized public List brokers() {
+        return new ArrayList<>(brokers);
+    }
+
+    synchronized public Node broker(int index) {
+        return brokers.get(index);
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/clients/admin/MockKafkaAdminClientEnv.java b/clients/src/test/java/org/apache/kafka/clients/admin/MockKafkaAdminClientEnv.java
deleted file mode 100644
index cca35ac22c977..0000000000000
--- a/clients/src/test/java/org/apache/kafka/clients/admin/MockKafkaAdminClientEnv.java
+++ /dev/null
@@ -1,95 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements. See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License. You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.kafka.clients.admin;
-
-import org.apache.kafka.clients.Metadata;
-import org.apache.kafka.clients.MockClient;
-import org.apache.kafka.common.Cluster;
-import org.apache.kafka.common.utils.Time;
-
-import java.util.HashMap;
-import java.util.Map;
-
-/**
- * Simple utility for setting up a mock {@link KafkaAdminClient} that uses a {@link MockClient} for a supplied
- * {@link Cluster}. Create a {@link Cluster} manually or use {@link org.apache.kafka.test.TestUtils} methods to
- * 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
- * and receive the responses you provided.
- * 
            
- * When finished, be sure to {@link #close() close} the environment object.
- */
-public class MockKafkaAdminClientEnv implements AutoCloseable {
-    private final Time time;
-    private final Cluster cluster;
-    private final MockClient mockClient;
-    private final KafkaAdminClient adminClient;
-
-    public MockKafkaAdminClientEnv(Cluster cluster, String...vals) {
-        this(Time.SYSTEM, cluster, vals);
-    }
-
-    public MockKafkaAdminClientEnv(Time time, Cluster cluster, String...vals) {
-        this(time, cluster, newStrMap(vals));
-    }
-
-    public MockKafkaAdminClientEnv(Time time, Cluster cluster, Map config) {
-        this.time = time;
-        this.cluster = cluster;
-        AdminClientConfig adminClientConfig = new AdminClientConfig(config);
-        Metadata metadata = new Metadata(adminClientConfig.getLong(AdminClientConfig.RETRY_BACKOFF_MS_CONFIG),
-                adminClientConfig.getLong(AdminClientConfig.METADATA_MAX_AGE_CONFIG), false);
-        this.mockClient = new MockClient(time, metadata);
-        this.adminClient = KafkaAdminClient.createInternal(adminClientConfig, mockClient, metadata, time);
-    }
-
-    public Time time() {
-        return time;
-    }
-
-    public Cluster cluster() {
-        return cluster;
-    }
-
-    public AdminClient adminClient() {
-        return adminClient;
-    }
-
-    public MockClient kafkaClient() {
-        return mockClient;
-    }
-
-    @Override
-    public void close() {
-        this.adminClient.close();
-    }
-
-    private static Map newStrMap(String... vals) {
-        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) {
-            throw new IllegalStateException();
-        }
-        for (int i = 0; i < vals.length; i += 2) {
-            map.put(vals[i], vals[i + 1]);
-        }
-        return map;
-    }
-}
diff --git a/clients/src/test/java/org/apache/kafka/clients/admin/RemoveMembersFromConsumerGroupOptionsTest.java b/clients/src/test/java/org/apache/kafka/clients/admin/RemoveMembersFromConsumerGroupOptionsTest.java
new file mode 100644
index 0000000000000..92b37e77df596
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/admin/RemoveMembersFromConsumerGroupOptionsTest.java
@@ -0,0 +1,39 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.clients.admin;
+
+import org.junit.Test;
+
+import java.util.Collections;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+
+public class RemoveMembersFromConsumerGroupOptionsTest {
+
+    @Test
+    public void testConstructor() {
+        RemoveMembersFromConsumerGroupOptions options = new RemoveMembersFromConsumerGroupOptions(
+            Collections.singleton(new MemberToRemove("instance-1")));
+
+        assertEquals(Collections.singleton(
+            new MemberToRemove("instance-1")), options.members());
+
+        // Construct will fail if illegal empty members provided
+        assertThrows(IllegalArgumentException.class, () -> new RemoveMembersFromConsumerGroupOptions(Collections.emptyList()));
+    }
+}
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/admin/ScramMechanismTest.java b/clients/src/test/java/org/apache/kafka/clients/admin/ScramMechanismTest.java
new file mode 100644
index 0000000000000..03ee0516b6521
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/admin/ScramMechanismTest.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.admin;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import org.junit.jupiter.api.Test;
+
+class ScramMechanismTest {
+
+    @Test
+    public void testFromMechanismName() {
+        assertEquals(ScramMechanism.UNKNOWN, ScramMechanism.fromMechanismName("UNKNOWN"));
+        assertEquals(ScramMechanism.SCRAM_SHA_256, ScramMechanism.fromMechanismName("SCRAM-SHA-256"));
+        assertEquals(ScramMechanism.SCRAM_SHA_512, ScramMechanism.fromMechanismName("SCRAM-SHA-512"));
+        assertEquals(ScramMechanism.UNKNOWN, ScramMechanism.fromMechanismName("some string"));
+        assertEquals(ScramMechanism.UNKNOWN, ScramMechanism.fromMechanismName("scram-sha-256"));
+    }
+
+}
diff --git a/clients/src/test/java/org/apache/kafka/clients/admin/internals/AdminMetadataManagerTest.java b/clients/src/test/java/org/apache/kafka/clients/admin/internals/AdminMetadataManagerTest.java
new file mode 100644
index 0000000000000..7254123024d2a
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/admin/internals/AdminMetadataManagerTest.java
@@ -0,0 +1,107 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.internals;
+
+import org.apache.kafka.common.Cluster;
+import org.apache.kafka.common.Node;
+import org.apache.kafka.common.PartitionInfo;
+import org.apache.kafka.common.errors.AuthenticationException;
+import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.common.utils.MockTime;
+import org.junit.Test;
+
+import java.net.InetSocketAddress;
+import java.util.Collections;
+import java.util.HashMap;
+
+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 AdminMetadataManagerTest {
+    private final MockTime time = new MockTime();
+    private final LogContext logContext = new LogContext();
+    private final long refreshBackoffMs = 100;
+    private final long metadataExpireMs = 60000;
+    private final AdminMetadataManager mgr = new AdminMetadataManager(
+            logContext, refreshBackoffMs, metadataExpireMs);
+
+    @Test
+    public void testMetadataReady() {
+        // Metadata is not ready on initialization
+        assertFalse(mgr.isReady());
+        assertEquals(0, mgr.metadataFetchDelayMs(time.milliseconds()));
+
+        // Metadata is not ready when bootstrap servers are set
+        mgr.update(Cluster.bootstrap(Collections.singletonList(new InetSocketAddress("localhost", 9999))),
+                time.milliseconds());
+        assertFalse(mgr.isReady());
+        assertEquals(0, mgr.metadataFetchDelayMs(time.milliseconds()));
+
+        mgr.update(mockCluster(), time.milliseconds());
+        assertTrue(mgr.isReady());
+        assertEquals(metadataExpireMs, mgr.metadataFetchDelayMs(time.milliseconds()));
+
+        time.sleep(metadataExpireMs);
+        assertEquals(0, mgr.metadataFetchDelayMs(time.milliseconds()));
+    }
+
+    @Test
+    public void testMetadataRefreshBackoff() {
+        mgr.transitionToUpdatePending(time.milliseconds());
+        assertEquals(Long.MAX_VALUE, mgr.metadataFetchDelayMs(time.milliseconds()));
+
+        mgr.updateFailed(new RuntimeException());
+        assertEquals(refreshBackoffMs, mgr.metadataFetchDelayMs(time.milliseconds()));
+
+        // Even if we explicitly request an update, the backoff should be respected
+        mgr.requestUpdate();
+        assertEquals(refreshBackoffMs, mgr.metadataFetchDelayMs(time.milliseconds()));
+
+        time.sleep(refreshBackoffMs);
+        assertEquals(0, mgr.metadataFetchDelayMs(time.milliseconds()));
+    }
+
+    @Test
+    public void testAuthenticationFailure() {
+        mgr.transitionToUpdatePending(time.milliseconds());
+        mgr.updateFailed(new AuthenticationException("Authentication failed"));
+        assertEquals(refreshBackoffMs, mgr.metadataFetchDelayMs(time.milliseconds()));
+        try {
+            mgr.isReady();
+            fail("Expected AuthenticationException to be thrown");
+        } catch (AuthenticationException e) {
+            // Expected
+        }
+
+        mgr.update(mockCluster(), time.milliseconds());
+        assertTrue(mgr.isReady());
+    }
+
+    private static Cluster mockCluster() {
+        HashMap nodes = new HashMap<>();
+        nodes.put(0, new Node(0, "localhost", 8121));
+        nodes.put(1, new Node(1, "localhost", 8122));
+        nodes.put(2, new Node(2, "localhost", 8123));
+        return new Cluster("mockClusterId", nodes.values(),
+                Collections.emptySet(), Collections.emptySet(),
+                Collections.emptySet(), nodes.get(0));
+    }
+
+}
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/ConsumerConfigTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/ConsumerConfigTest.java
index 48580c2657d12..ea7b170d58ef3 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/ConsumerConfigTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/ConsumerConfigTest.java
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.clients.consumer;
 
+import org.apache.kafka.common.errors.InvalidConfigurationException;
 import org.apache.kafka.common.serialization.ByteArrayDeserializer;
 import org.apache.kafka.common.serialization.Deserializer;
 import org.apache.kafka.common.serialization.StringDeserializer;
@@ -25,17 +26,49 @@
 import java.util.Map;
 import java.util.Properties;
 
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.fail;
 
 public class ConsumerConfigTest {
 
-    private final Deserializer keyDeserializer = new ByteArrayDeserializer();
-    private final Deserializer valueDeserializer = new StringDeserializer();
+    private final Deserializer keyDeserializer = new ByteArrayDeserializer();
+    private final Deserializer valueDeserializer = new StringDeserializer();
     private final String keyDeserializerClassName = keyDeserializer.getClass().getName();
     private final String valueDeserializerClassName = valueDeserializer.getClass().getName();
     private final Object keyDeserializerClass = keyDeserializer.getClass();
     private final Object valueDeserializerClass = valueDeserializer.getClass();
 
+    @Test
+    public void testOverrideClientId() {
+        Properties properties = new Properties();
+        properties.setProperty(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, keyDeserializerClassName);
+        properties.setProperty(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, valueDeserializerClassName);
+        properties.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "test-group");
+        ConsumerConfig config = new ConsumerConfig(properties);
+        assertFalse(config.getString(ConsumerConfig.CLIENT_ID_CONFIG).isEmpty());
+    }
+
+    @Test
+    public void testOverrideEnableAutoCommit() {
+        Properties properties = new Properties();
+        properties.setProperty(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, keyDeserializerClassName);
+        properties.setProperty(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, valueDeserializerClassName);
+        ConsumerConfig config = new ConsumerConfig(properties);
+        boolean overrideEnableAutoCommit = config.maybeOverrideEnableAutoCommit();
+        assertFalse(overrideEnableAutoCommit);
+
+        properties.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true");
+        config = new ConsumerConfig(properties);
+        try {
+            config.maybeOverrideEnableAutoCommit();
+            fail("Should have thrown an exception");
+        } catch (InvalidConfigurationException e) {
+            // expected
+        }
+    }
+
+    @SuppressWarnings("deprecation")
     @Test
     public void testDeserializerToPropertyConfig() {
         Properties properties = new Properties();
@@ -64,29 +97,37 @@ public void testDeserializerToPropertyConfig() {
     }
 
     @Test
-    public void testDeserializerToMapConfig() {
+    public void testAppendDeserializerToConfig() {
         Map configs = new HashMap<>();
         configs.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, keyDeserializerClass);
         configs.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, valueDeserializerClass);
-        Map newConfigs = ConsumerConfig.addDeserializerToConfig(configs, null, null);
+        Map newConfigs = ConsumerConfig.appendDeserializerToConfig(configs, null, null);
         assertEquals(newConfigs.get(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG), keyDeserializerClass);
         assertEquals(newConfigs.get(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG), valueDeserializerClass);
 
         configs.clear();
         configs.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, valueDeserializerClass);
-        newConfigs = ConsumerConfig.addDeserializerToConfig(configs, keyDeserializer, null);
+        newConfigs = ConsumerConfig.appendDeserializerToConfig(configs, keyDeserializer, null);
         assertEquals(newConfigs.get(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG), keyDeserializerClass);
         assertEquals(newConfigs.get(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG), valueDeserializerClass);
 
         configs.clear();
         configs.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, keyDeserializerClass);
-        newConfigs = ConsumerConfig.addDeserializerToConfig(configs, null, valueDeserializer);
+        newConfigs = ConsumerConfig.appendDeserializerToConfig(configs, null, valueDeserializer);
         assertEquals(newConfigs.get(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG), keyDeserializerClass);
         assertEquals(newConfigs.get(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG), valueDeserializerClass);
 
         configs.clear();
-        newConfigs = ConsumerConfig.addDeserializerToConfig(configs, keyDeserializer, valueDeserializer);
+        newConfigs = ConsumerConfig.appendDeserializerToConfig(configs, keyDeserializer, valueDeserializer);
         assertEquals(newConfigs.get(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG), keyDeserializerClass);
         assertEquals(newConfigs.get(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG), valueDeserializerClass);
     }
+
+    @Test
+    public void ensureDefaultThrowOnUnsupportedStableFlagToFalse() {
+        Properties properties = new Properties();
+        properties.setProperty(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, keyDeserializerClassName);
+        properties.setProperty(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, valueDeserializerClassName);
+        assertFalse(new ConsumerConfig(properties).getBoolean(ConsumerConfig.THROW_ON_FETCH_STABLE_OFFSET_UNSUPPORTED));
+    }
 }
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/ConsumerGroupMetadataTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/ConsumerGroupMetadataTest.java
new file mode 100644
index 0000000000000..0af6261447875
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/ConsumerGroupMetadataTest.java
@@ -0,0 +1,87 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.requests.JoinGroupRequest;
+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.assertThrows;
+import static org.junit.Assert.assertTrue;
+
+public class ConsumerGroupMetadataTest {
+
+    private String groupId = "group";
+
+    @Test
+    public void testAssignmentConstructor() {
+        String memberId = "member";
+        int generationId = 2;
+        String groupInstanceId = "instance";
+
+        ConsumerGroupMetadata groupMetadata = new ConsumerGroupMetadata(groupId,
+            generationId, memberId, Optional.of(groupInstanceId));
+
+        assertEquals(groupId, groupMetadata.groupId());
+        assertEquals(generationId, groupMetadata.generationId());
+        assertEquals(memberId, groupMetadata.memberId());
+        assertTrue(groupMetadata.groupInstanceId().isPresent());
+        assertEquals(groupInstanceId, groupMetadata.groupInstanceId().get());
+    }
+
+    @Test
+    public void testGroupIdConstructor() {
+        ConsumerGroupMetadata groupMetadata = new ConsumerGroupMetadata(groupId);
+
+        assertEquals(groupId, groupMetadata.groupId());
+        assertEquals(JoinGroupRequest.UNKNOWN_GENERATION_ID, groupMetadata.generationId());
+        assertEquals(JoinGroupRequest.UNKNOWN_MEMBER_ID, groupMetadata.memberId());
+        assertFalse(groupMetadata.groupInstanceId().isPresent());
+    }
+
+    @Test
+    public void testInvalidGroupId() {
+        String memberId = "member";
+        int generationId = 2;
+
+        assertThrows(NullPointerException.class, () -> new ConsumerGroupMetadata(
+            null, generationId, memberId, Optional.empty())
+        );
+    }
+
+    @Test
+    public void testInvalidMemberId() {
+        int generationId = 2;
+
+        assertThrows(NullPointerException.class, () -> new ConsumerGroupMetadata(
+            groupId, generationId, null, Optional.empty())
+        );
+    }
+
+    @Test
+    public void testInvalidInstanceId() {
+        String memberId = "member";
+        int generationId = 2;
+
+        assertThrows(NullPointerException.class, () -> new ConsumerGroupMetadata(
+            groupId, generationId, memberId, null)
+        );
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/ConsumerRecordTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/ConsumerRecordTest.java
index 327364520040c..aae269adf2bde 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/ConsumerRecordTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/ConsumerRecordTest.java
@@ -21,6 +21,8 @@
 import org.apache.kafka.common.record.TimestampType;
 import org.junit.Test;
 
+import java.util.Optional;
+
 import static org.junit.Assert.assertEquals;
 
 public class ConsumerRecordTest {
@@ -45,6 +47,7 @@ public void testOldConstructor() {
         assertEquals(ConsumerRecord.NULL_CHECKSUM, record.checksum());
         assertEquals(ConsumerRecord.NULL_SIZE, record.serializedKeySize());
         assertEquals(ConsumerRecord.NULL_SIZE, record.serializedValueSize());
+        assertEquals(Optional.empty(), record.leaderEpoch());
         assertEquals(new RecordHeaders(), record.headers());
     }
 
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 12254c9a4554b..0040452667431 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,24 +16,45 @@
  */
 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.Metadata;
 import org.apache.kafka.clients.MockClient;
+import org.apache.kafka.clients.NodeApiVersions;
 import org.apache.kafka.clients.consumer.internals.ConsumerCoordinator;
 import org.apache.kafka.clients.consumer.internals.ConsumerInterceptors;
+import org.apache.kafka.clients.consumer.internals.ConsumerMetadata;
 import org.apache.kafka.clients.consumer.internals.ConsumerMetrics;
 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.PartitionAssignor;
+import org.apache.kafka.clients.consumer.internals.MockRebalanceListener;
 import org.apache.kafka.clients.consumer.internals.SubscriptionState;
 import org.apache.kafka.common.Cluster;
+import org.apache.kafka.common.IsolationLevel;
 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.config.SslConfigs;
+import org.apache.kafka.common.errors.AuthenticationException;
 import org.apache.kafka.common.errors.InterruptException;
+import org.apache.kafka.common.errors.InvalidConfigurationException;
+import org.apache.kafka.common.errors.InvalidGroupIdException;
+import org.apache.kafka.common.errors.InvalidTopicException;
+import org.apache.kafka.common.errors.UnsupportedVersionException;
 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.message.ListOffsetResponseData;
+import org.apache.kafka.common.message.ListOffsetResponseData.ListOffsetTopicResponse;
+import org.apache.kafka.common.message.ListOffsetResponseData.ListOffsetPartitionResponse;
+import org.apache.kafka.common.internals.ClusterResourceListeners;
+import org.apache.kafka.common.message.SyncGroupResponseData;
+import org.apache.kafka.common.message.ListOffsetRequestData.ListOffsetPartition;
 import org.apache.kafka.common.metrics.Metrics;
 import org.apache.kafka.common.metrics.Sensor;
 import org.apache.kafka.common.network.Selectable;
@@ -47,17 +68,18 @@
 import org.apache.kafka.common.requests.AbstractResponse;
 import org.apache.kafka.common.requests.FetchRequest;
 import org.apache.kafka.common.requests.FetchResponse;
-import org.apache.kafka.common.requests.FetchResponse.PartitionData;
 import org.apache.kafka.common.requests.FindCoordinatorResponse;
 import org.apache.kafka.common.requests.HeartbeatResponse;
-import org.apache.kafka.common.requests.IsolationLevel;
 import org.apache.kafka.common.requests.JoinGroupRequest;
 import org.apache.kafka.common.requests.JoinGroupResponse;
 import org.apache.kafka.common.requests.LeaveGroupResponse;
+import org.apache.kafka.common.requests.ListOffsetRequest;
 import org.apache.kafka.common.requests.ListOffsetResponse;
+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.OffsetFetchResponse;
+import org.apache.kafka.common.requests.RequestTestUtils;
 import org.apache.kafka.common.requests.SyncGroupResponse;
 import org.apache.kafka.common.serialization.ByteArrayDeserializer;
 import org.apache.kafka.common.serialization.Deserializer;
@@ -69,21 +91,30 @@
 import org.apache.kafka.test.MockConsumerInterceptor;
 import org.apache.kafka.test.MockMetricsReporter;
 import org.apache.kafka.test.TestUtils;
+import org.apache.kafka.common.metrics.stats.Avg;
+
 import org.junit.Assert;
-import org.junit.Rule;
 import org.junit.Test;
-import org.junit.rules.ExpectedException;
 
+import javax.management.MBeanServer;
+import javax.management.ObjectName;
+import java.lang.management.ManagementFactory;
 import java.nio.ByteBuffer;
+import java.time.Duration;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.ConcurrentModificationException;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.Iterator;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Properties;
+import java.util.Queue;
 import java.util.Set;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -92,17 +123,30 @@
 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.AtomicReference;
 import java.util.regex.Pattern;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 import static java.util.Collections.singleton;
 import static java.util.Collections.singletonList;
 import static java.util.Collections.singletonMap;
+import static org.apache.kafka.common.requests.FetchMetadata.INVALID_SESSION_ID;
 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.assertNotSame;
 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.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.Mockito.doCallRealMethod;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
 
 public class KafkaConsumerTest {
     private final String topic = "test";
@@ -115,14 +159,43 @@ public class KafkaConsumerTest {
     private final String topic3 = "test3";
     private final TopicPartition t3p0 = new TopicPartition(topic3, 0);
 
-    @Rule
-    public ExpectedException expectedException = ExpectedException.none();
+    private final int sessionTimeoutMs = 10000;
+    private final int heartbeatIntervalMs = 1000;
+
+    // Set auto commit interval lower than heartbeat so we don't need to deal with
+    // a concurrent heartbeat request
+    private final int autoCommitIntervalMs = 500;
+
+    private final String groupId = "mock-group";
+    private final String memberId = "memberId";
+    private final String leaderId = "leaderId";
+    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() {
+        Properties props = new Properties();
+        props.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
+        props.setProperty(ConsumerConfig.METRIC_REPORTER_CLASSES_CONFIG, MockMetricsReporter.class.getName());
+        KafkaConsumer consumer = new KafkaConsumer<>(
+                props, new StringDeserializer(), new StringDeserializer());
+
+        MockMetricsReporter mockMetricsReporter = (MockMetricsReporter) consumer.metrics.reporters().get(0);
+
+        Assert.assertEquals(consumer.getClientId(), mockMetricsReporter.clientId);
+        consumer.close();
+    }
 
     @Test
-    public void testConstructorClose() throws Exception {
+    public void testConstructorClose() {
         Properties props = new Properties();
         props.setProperty(ConsumerConfig.CLIENT_ID_CONFIG, "testConstructorClose");
-        props.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "some.invalid.hostname.foo.bar.local:9999");
+        props.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "invalid-23-8409-adsfsdj");
         props.setProperty(ConsumerConfig.METRIC_REPORTER_CLASSES_CONFIG, MockMetricsReporter.class.getName());
 
         final int oldInitCount = MockMetricsReporter.INIT_COUNT.get();
@@ -138,7 +211,7 @@ public void testConstructorClose() throws Exception {
     }
 
     @Test
-    public void testOsDefaultSocketBufferSizes() throws Exception {
+    public void testOsDefaultSocketBufferSizes() {
         Map config = new HashMap<>();
         config.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
         config.put(ConsumerConfig.SEND_BUFFER_CONFIG, Selectable.USE_DEFAULT_BUFFER_SIZE);
@@ -149,7 +222,7 @@ public void testOsDefaultSocketBufferSizes() throws Exception {
     }
 
     @Test(expected = KafkaException.class)
-    public void testInvalidSocketSendBufferSize() throws Exception {
+    public void testInvalidSocketSendBufferSize() {
         Map config = new HashMap<>();
         config.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
         config.put(ConsumerConfig.SEND_BUFFER_CONFIG, -2);
@@ -157,22 +230,32 @@ public void testInvalidSocketSendBufferSize() throws Exception {
     }
 
     @Test(expected = KafkaException.class)
-    public void testInvalidSocketReceiveBufferSize() throws Exception {
+    public void testInvalidSocketReceiveBufferSize() {
         Map config = new HashMap<>();
         config.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
         config.put(ConsumerConfig.RECEIVE_BUFFER_CONFIG, -2);
         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();
+        KafkaConsumer consumer = newConsumer(groupId);
 
         consumer.subscribe(singletonList(topic));
         assertEquals(singleton(topic), consumer.subscription());
         assertTrue(consumer.assignment().isEmpty());
 
-        consumer.subscribe(Collections.emptyList());
+        consumer.subscribe(Collections.emptyList());
         assertTrue(consumer.subscription().isEmpty());
         assertTrue(consumer.assignment().isEmpty());
 
@@ -189,49 +272,37 @@ public void testSubscription() {
 
     @Test(expected = IllegalArgumentException.class)
     public void testSubscriptionOnNullTopicCollection() {
-        KafkaConsumer consumer = newConsumer();
-        List nullList = null;
-
-        try {
-            consumer.subscribe(nullList);
-        } finally {
-            consumer.close();
+        try (KafkaConsumer consumer = newConsumer(groupId)) {
+            consumer.subscribe((List) null);
         }
     }
 
     @Test(expected = IllegalArgumentException.class)
     public void testSubscriptionOnNullTopic() {
-        KafkaConsumer consumer = newConsumer();
-        String nullTopic = null;
-
-        try {
-            consumer.subscribe(singletonList(nullTopic));
-        } finally {
-            consumer.close();
+        try (KafkaConsumer consumer = newConsumer(groupId)) {
+            consumer.subscribe(singletonList(null));
         }
     }
 
     @Test(expected = IllegalArgumentException.class)
     public void testSubscriptionOnEmptyTopic() {
-        KafkaConsumer consumer = newConsumer();
-        String emptyTopic = "  ";
-
-        try {
+        try (KafkaConsumer consumer = newConsumer(groupId)) {
+            String emptyTopic = "  ";
             consumer.subscribe(singletonList(emptyTopic));
-        } finally {
-            consumer.close();
         }
     }
 
     @Test(expected = IllegalArgumentException.class)
     public void testSubscriptionOnNullPattern() {
-        KafkaConsumer consumer = newConsumer();
-        Pattern pattern = null;
+        try (KafkaConsumer consumer = newConsumer(groupId)) {
+            consumer.subscribe((Pattern) null);
+        }
+    }
 
-        try {
-            consumer.subscribe(pattern);
-        } finally {
-            consumer.close();
+    @Test(expected = IllegalArgumentException.class)
+    public void testSubscriptionOnEmptyPattern() {
+        try (KafkaConsumer consumer = newConsumer(groupId)) {
+            consumer.subscribe(Pattern.compile(""));
         }
     }
 
@@ -240,76 +311,59 @@ public void testSubscriptionWithEmptyPartitionAssignment() {
         Properties props = new Properties();
         props.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
         props.setProperty(ConsumerConfig.PARTITION_ASSIGNMENT_STRATEGY_CONFIG, "");
-
-        KafkaConsumer consumer = new KafkaConsumer<>(props, new ByteArrayDeserializer(), new ByteArrayDeserializer());
-        try {
+        props.setProperty(ConsumerConfig.GROUP_ID_CONFIG, groupId);
+        try (KafkaConsumer consumer = newConsumer(props)) {
             consumer.subscribe(singletonList(topic));
-        } finally {
-            consumer.close();
         }
     }
 
     @Test(expected = IllegalArgumentException.class)
     public void testSeekNegative() {
-        KafkaConsumer consumer = newConsumer();
-        try {
-            consumer.assign(Arrays.asList(new TopicPartition("nonExistTopic", 0)));
+        try (KafkaConsumer consumer = newConsumer((String) null)) {
+            consumer.assign(singleton(new TopicPartition("nonExistTopic", 0)));
             consumer.seek(new TopicPartition("nonExistTopic", 0), -1);
-        } finally {
-            consumer.close();
         }
     }
 
     @Test(expected = IllegalArgumentException.class)
     public void testAssignOnNullTopicPartition() {
-        KafkaConsumer consumer = newConsumer();
-        try {
+        try (KafkaConsumer consumer = newConsumer((String) null)) {
             consumer.assign(null);
-        } finally {
-            consumer.close();
         }
     }
 
     @Test
     public void testAssignOnEmptyTopicPartition() {
-        KafkaConsumer consumer = newConsumer();
-
-        consumer.assign(Collections.emptyList());
-        assertTrue(consumer.subscription().isEmpty());
-        assertTrue(consumer.assignment().isEmpty());
-
-        consumer.close();
+        try (KafkaConsumer consumer = newConsumer(groupId)) {
+            consumer.assign(Collections.emptyList());
+            assertTrue(consumer.subscription().isEmpty());
+            assertTrue(consumer.assignment().isEmpty());
+        }
     }
 
     @Test(expected = IllegalArgumentException.class)
     public void testAssignOnNullTopicInPartition() {
-        KafkaConsumer consumer = newConsumer();
-        try {
-            consumer.assign(Arrays.asList(new TopicPartition(null, 0)));
-        } finally {
-            consumer.close();
+        try (KafkaConsumer consumer = newConsumer((String) null)) {
+            consumer.assign(singleton(new TopicPartition(null, 0)));
         }
     }
 
     @Test(expected = IllegalArgumentException.class)
     public void testAssignOnEmptyTopicInPartition() {
-        KafkaConsumer consumer = newConsumer();
-        try {
-            consumer.assign(Arrays.asList(new TopicPartition("  ", 0)));
-        } finally {
-            consumer.close();
+        try (KafkaConsumer consumer = newConsumer((String) null)) {
+            consumer.assign(singleton(new TopicPartition("  ", 0)));
         }
     }
 
     @Test
-    public void testInterceptorConstructorClose() throws Exception {
+    public void testInterceptorConstructorClose() {
         try {
             Properties props = new Properties();
             // test with client ID assigned by KafkaConsumer
             props.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
             props.setProperty(ConsumerConfig.INTERCEPTOR_CLASSES_CONFIG, MockConsumerInterceptor.class.getName());
 
-            KafkaConsumer consumer = new KafkaConsumer(
+            KafkaConsumer consumer = new KafkaConsumer<>(
                     props, new StringDeserializer(), new StringDeserializer());
             assertEquals(1, MockConsumerInterceptor.INIT_COUNT.get());
             assertEquals(0, MockConsumerInterceptor.CLOSE_COUNT.get());
@@ -328,7 +382,7 @@ public void testInterceptorConstructorClose() throws Exception {
 
     @Test
     public void testPause() {
-        KafkaConsumer consumer = newConsumer();
+        KafkaConsumer consumer = newConsumer(groupId);
 
         consumer.assign(singletonList(tp0));
         assertEquals(singleton(tp0), consumer.assignment());
@@ -346,131 +400,334 @@ public void testPause() {
         consumer.close();
     }
 
-    private KafkaConsumer newConsumer() {
+    @Test
+    public void testConsumerJmxPrefix() throws  Exception {
+        Map config = new HashMap<>();
+        config.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
+        config.put(ConsumerConfig.SEND_BUFFER_CONFIG, Selectable.USE_DEFAULT_BUFFER_SIZE);
+        config.put(ConsumerConfig.RECEIVE_BUFFER_CONFIG, Selectable.USE_DEFAULT_BUFFER_SIZE);
+        config.put("client.id", "client-1");
+        KafkaConsumer consumer = new KafkaConsumer<>(
+                config, new ByteArrayDeserializer(), new ByteArrayDeserializer());
+        MBeanServer server = ManagementFactory.getPlatformMBeanServer();
+        MetricName testMetricName = consumer.metrics.metricName("test-metric",
+                "grp1", "test metric");
+        consumer.metrics.addMetric(testMetricName, new Avg());
+        Assert.assertNotNull(server.getObjectInstance(new ObjectName("kafka.consumer:type=grp1,client-id=client-1")));
+        consumer.close();
+    }
+
+    private KafkaConsumer newConsumer(String groupId) {
+        return newConsumer(groupId, Optional.empty());
+    }
+
+    private KafkaConsumer newConsumer(String groupId, Optional enableAutoCommit) {
         Properties props = new Properties();
         props.setProperty(ConsumerConfig.CLIENT_ID_CONFIG, "my.consumer");
         props.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
         props.setProperty(ConsumerConfig.METRIC_REPORTER_CLASSES_CONFIG, MockMetricsReporter.class.getName());
+        if (groupId != null)
+            props.setProperty(ConsumerConfig.GROUP_ID_CONFIG, groupId);
+        enableAutoCommit.ifPresent(
+            autoCommit -> props.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, autoCommit.toString()));
+        return newConsumer(props);
+    }
 
+    private KafkaConsumer newConsumer(Properties props) {
         return new KafkaConsumer<>(props, new ByteArrayDeserializer(), new ByteArrayDeserializer());
     }
 
     @Test
     public void verifyHeartbeatSent() throws Exception {
-        int rebalanceTimeoutMs = 60000;
-        int sessionTimeoutMs = 30000;
-        int heartbeatIntervalMs = 1000;
-        int autoCommitIntervalMs = 10000;
-
         Time time = new MockTime();
-        Cluster cluster = TestUtils.singletonCluster(topic, 1);
-        Node node = cluster.nodes().get(0);
+        SubscriptionState subscription = new SubscriptionState(new LogContext(), OffsetResetStrategy.EARLIEST);
+        ConsumerMetadata metadata = createMetadata(subscription);
+        MockClient client = new MockClient(time, metadata);
 
-        Metadata metadata = createMetadata();
-        metadata.update(cluster, Collections.emptySet(), time.milliseconds());
+        initMetadata(client, Collections.singletonMap(topic, 1));
+        Node node = metadata.fetch().nodes().get(0);
 
-        MockClient client = new MockClient(time, metadata);
-        client.setNode(node);
-        PartitionAssignor assignor = new RoundRobinAssignor();
+        ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
 
-        final KafkaConsumer consumer = newConsumer(time, client, metadata, assignor,
-                rebalanceTimeoutMs, sessionTimeoutMs, heartbeatIntervalMs, true, autoCommitIntervalMs);
+        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true, groupInstanceId);
 
-        consumer.subscribe(Arrays.asList(topic), getConsumerRebalanceListener(consumer));
-        Node coordinator = prepareRebalance(client, node, assignor, Arrays.asList(tp0), null);
+        consumer.subscribe(singleton(topic), getConsumerRebalanceListener(consumer));
+        Node coordinator = prepareRebalance(client, node, assignor, singletonList(tp0), null);
 
         // initial fetch
         client.prepareResponseFrom(fetchResponse(tp0, 0, 0), node);
+        consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));
 
-        consumer.poll(0);
-        assertEquals(Collections.singleton(tp0), consumer.assignment());
+        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);
         Thread.sleep(heartbeatIntervalMs);
 
-        consumer.poll(0);
+        consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));
 
         assertTrue(heartbeatReceived.get());
-        consumer.close(0, TimeUnit.MILLISECONDS);
+        consumer.close(Duration.ofMillis(0));
     }
 
     @Test
     public void verifyHeartbeatSentWhenFetchedDataReady() throws Exception {
-        int rebalanceTimeoutMs = 60000;
-        int sessionTimeoutMs = 30000;
-        int heartbeatIntervalMs = 1000;
-        int autoCommitIntervalMs = 10000;
-
         Time time = new MockTime();
-        Cluster cluster = TestUtils.singletonCluster(topic, 1);
-        Node node = cluster.nodes().get(0);
-
-        Metadata metadata = createMetadata();
-        metadata.update(cluster, Collections.emptySet(), time.milliseconds());
-
+        SubscriptionState subscription = new SubscriptionState(new LogContext(), OffsetResetStrategy.EARLIEST);
+        ConsumerMetadata metadata = createMetadata(subscription);
         MockClient client = new MockClient(time, metadata);
-        client.setNode(node);
-        PartitionAssignor assignor = new RoundRobinAssignor();
 
-        final KafkaConsumer consumer = newConsumer(time, client, metadata, assignor,
-                rebalanceTimeoutMs, sessionTimeoutMs, heartbeatIntervalMs, true, autoCommitIntervalMs);
+        initMetadata(client, Collections.singletonMap(topic, 1));
+        Node node = metadata.fetch().nodes().get(0);
+        ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
 
-        consumer.subscribe(Arrays.asList(topic), getConsumerRebalanceListener(consumer));
-        Node coordinator = prepareRebalance(client, node, assignor, Arrays.asList(tp0), null);
+        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);
 
-        consumer.poll(0);
+        consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));
+        consumer.poll(Duration.ZERO);
 
         // respond to the outstanding fetch so that we have data available on the next poll
         client.respondFrom(fetchResponse(tp0, 0, 5), node);
         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);
 
-        consumer.poll(0);
+        consumer.poll(Duration.ZERO);
 
         assertTrue(heartbeatReceived.get());
-        consumer.close(0, TimeUnit.MILLISECONDS);
+        consumer.close(Duration.ofMillis(0));
     }
 
     @Test
-    public void verifyNoCoordinatorLookupForManualAssignmentWithSeek() {
-        int rebalanceTimeoutMs = 60000;
-        int sessionTimeoutMs = 3000;
-        int heartbeatIntervalMs = 2000;
-        int autoCommitIntervalMs = 1000;
+    public void verifyPollTimesOutDuringMetadataUpdate() {
+        final Time time = new MockTime();
+        final SubscriptionState subscription = new SubscriptionState(new LogContext(), OffsetResetStrategy.EARLIEST);
+        final ConsumerMetadata metadata = createMetadata(subscription);
+        final MockClient client = new MockClient(time, metadata);
 
-        Time time = new MockTime();
-        Cluster cluster = TestUtils.singletonCluster(topic, 1);
-        Node node = cluster.nodes().get(0);
+        initMetadata(client, Collections.singletonMap(topic, 1));
+        Node node = metadata.fetch().nodes().get(0);
+
+        final ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
+
+        final KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true, groupInstanceId);
+        consumer.subscribe(singleton(topic), getConsumerRebalanceListener(consumer));
+        // Since we would enable the heartbeat thread after received join-response which could
+        // send the sync-group on behalf of the consumer if it is enqueued, we may still complete
+        // the rebalance and send out the fetch; in order to avoid it we do not prepare sync response here.
+        client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, node), node);
+        Node coordinator = new Node(Integer.MAX_VALUE - node.id(), node.host(), node.port());
+        client.prepareResponseFrom(joinGroupFollowerResponse(assignor, 1, memberId, leaderId, Errors.NONE), coordinator);
+
+        consumer.poll(Duration.ZERO);
+
+        final Queue requests = client.requests();
+        Assert.assertEquals(0, requests.stream().filter(request -> request.apiKey().equals(ApiKeys.FETCH)).count());
+    }
 
-        Metadata metadata = createMetadata();
-        metadata.update(cluster, Collections.emptySet(), time.milliseconds());
+    @SuppressWarnings("deprecation")
+    @Test
+    public void verifyDeprecatedPollDoesNotTimeOutDuringMetadataUpdate() {
+        final Time time = new MockTime();
+        final SubscriptionState subscription = new SubscriptionState(new LogContext(), OffsetResetStrategy.EARLIEST);
+        final ConsumerMetadata metadata = createMetadata(subscription);
+        final MockClient client = new MockClient(time, metadata);
+
+        initMetadata(client, Collections.singletonMap(topic, 1));
+        Node node = metadata.fetch().nodes().get(0);
+
+        final ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
+
+        final KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true, groupInstanceId);
+        consumer.subscribe(singleton(topic), getConsumerRebalanceListener(consumer));
+        prepareRebalance(client, node, assignor, singletonList(tp0), null);
+
+        consumer.poll(0L);
 
+        // The underlying client SHOULD get a fetch request
+        final Queue requests = client.requests();
+        Assert.assertEquals(1, requests.size());
+        final Class aClass = requests.peek().requestBuilder().getClass();
+        Assert.assertEquals(FetchRequest.Builder.class, aClass);
+    }
+
+    @Test
+    public void verifyNoCoordinatorLookupForManualAssignmentWithSeek() {
+        Time time = new MockTime();
+        SubscriptionState subscription = new SubscriptionState(new LogContext(), OffsetResetStrategy.EARLIEST);
+        ConsumerMetadata metadata = createMetadata(subscription);
         MockClient client = new MockClient(time, metadata);
-        client.setNode(node);
-        PartitionAssignor assignor = new RoundRobinAssignor();
 
-        final KafkaConsumer consumer = newConsumer(time, client, metadata, assignor,
-                rebalanceTimeoutMs, sessionTimeoutMs, heartbeatIntervalMs, true, autoCommitIntervalMs);
-        consumer.assign(Arrays.asList(tp0));
-        consumer.seekToBeginning(Arrays.asList(tp0));
+        initMetadata(client, Collections.singletonMap(topic, 1));
+        ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
+
+        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true, groupInstanceId);
+        consumer.assign(singleton(tp0));
+        consumer.seekToBeginning(singleton(tp0));
 
         // there shouldn't be any need to lookup the coordinator or fetch committed offsets.
         // we just lookup the starting position and send the record fetch.
-        client.prepareResponse(listOffsetsResponse(Collections.singletonMap(tp0, 50L), Errors.NONE));
+        client.prepareResponse(listOffsetsResponse(Collections.singletonMap(tp0, 50L)));
         client.prepareResponse(fetchResponse(tp0, 50L, 5));
 
-        ConsumerRecords records = consumer.poll(0);
+        ConsumerRecords records = consumer.poll(Duration.ofMillis(1));
         assertEquals(5, records.count());
         assertEquals(55L, consumer.position(tp0));
-        consumer.close(0, TimeUnit.MILLISECONDS);
+        consumer.close(Duration.ofMillis(0));
+    }
+
+    @Test
+    public void testFetchProgressWithMissingPartitionPosition() {
+        // Verifies that we can make progress on one partition while we are awaiting
+        // a reset on another partition.
+
+        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));
+
+        KafkaConsumer consumer = newConsumerNoAutoCommit(time, client, subscription, metadata);
+        consumer.assign(Arrays.asList(tp0, tp1));
+        consumer.seekToEnd(singleton(tp0));
+        consumer.seekToBeginning(singleton(tp1));
+
+        client.prepareResponse(body -> {
+            ListOffsetRequest request = (ListOffsetRequest) body;
+            List partitions = request.topics().stream().flatMap(t -> {
+                if (t.name().equals(topic))
+                    return Stream.of(t.partitions());
+                else
+                    return Stream.empty();
+            }).flatMap(List::stream).collect(Collectors.toList());
+            ListOffsetPartition expectedTp0 = new ListOffsetPartition()
+                    .setPartitionIndex(tp0.partition())
+                    .setTimestamp(ListOffsetRequest.LATEST_TIMESTAMP);
+            ListOffsetPartition expectedTp1 = new ListOffsetPartition()
+                    .setPartitionIndex(tp1.partition())
+                    .setTimestamp(ListOffsetRequest.EARLIEST_TIMESTAMP);
+            return partitions.contains(expectedTp0) && partitions.contains(expectedTp1);
+        }, listOffsetsResponse(Collections.singletonMap(tp0, 50L), Collections.singletonMap(tp1, Errors.NOT_LEADER_OR_FOLLOWER)));
+        client.prepareResponse(
+            body -> {
+                FetchRequest request = (FetchRequest) body;
+                return request.fetchData().keySet().equals(singleton(tp0)) &&
+                        request.fetchData().get(tp0).fetchOffset == 50L;
+
+            }, fetchResponse(tp0, 50L, 5));
+
+        ConsumerRecords records = consumer.poll(Duration.ofMillis(1));
+        assertEquals(5, records.count());
+        assertEquals(singleton(tp0), records.partitions());
+    }
+
+    private void initMetadata(MockClient mockClient, Map partitionCounts) {
+        MetadataResponse initialMetadata = RequestTestUtils.metadataUpdateWith(1, partitionCounts);
+
+        mockClient.updateMetadata(initialMetadata);
+    }
+
+    @Test(expected = NoOffsetForPartitionException.class)
+    public void testMissingOffsetNoResetPolicy() {
+        Time time = new MockTime();
+        SubscriptionState subscription = new SubscriptionState(new LogContext(), OffsetResetStrategy.NONE);
+        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, groupId, groupInstanceId, false);
+        consumer.assign(singletonList(tp0));
+
+        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
+        client.prepareResponseFrom(offsetResponse(Collections.singletonMap(tp0, -1L), Errors.NONE), coordinator);
+        consumer.poll(Duration.ZERO);
+    }
+
+    @Test
+    public void testResetToCommittedOffset() {
+        Time time = new MockTime();
+        SubscriptionState subscription = new SubscriptionState(new LogContext(), OffsetResetStrategy.NONE);
+        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, groupId, groupInstanceId, false);
+        consumer.assign(singletonList(tp0));
+
+        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);
+        consumer.poll(Duration.ZERO);
+
+        assertEquals(539L, consumer.position(tp0));
+    }
+
+    @Test
+    public void testResetUsingAutoResetPolicy() {
+        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));
+        Node node = metadata.fetch().nodes().get(0);
+
+        ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
+
+        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor,
+                true, groupId, groupInstanceId, false);
+        consumer.assign(singletonList(tp0));
+
+        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);
+        client.prepareResponse(listOffsetsResponse(Collections.singletonMap(tp0, 50L)));
+
+        consumer.poll(Duration.ZERO);
+
+        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(), false);
+        consumer.assign(singletonList(tp0));
+        consumer.seek(tp0, 20L);
+        consumer.poll(Duration.ZERO);
+        assertEquals(subscription.validPosition(tp0).offset, 20L);
     }
 
     @Test
@@ -478,36 +735,26 @@ public void testCommitsFetchedDuringAssign() {
         long offset1 = 10000;
         long offset2 = 20000;
 
-        int rebalanceTimeoutMs = 6000;
-        int sessionTimeoutMs = 3000;
-        int heartbeatIntervalMs = 2000;
-        int autoCommitIntervalMs = 1000;
-
         Time time = new MockTime();
-        Cluster cluster = TestUtils.singletonCluster(topic, 1);
-        Node node = cluster.nodes().get(0);
+        SubscriptionState subscription = new SubscriptionState(new LogContext(), OffsetResetStrategy.EARLIEST);
+        ConsumerMetadata metadata = createMetadata(subscription);
+        MockClient client = new MockClient(time, metadata);
 
-        Metadata metadata = createMetadata();
-        metadata.update(cluster, Collections.emptySet(), time.milliseconds());
+        initMetadata(client, Collections.singletonMap(topic, 2));
+        Node node = metadata.fetch().nodes().get(0);
 
-        MockClient client = new MockClient(time, metadata);
-        client.setNode(node);
-        PartitionAssignor assignor = new RoundRobinAssignor();
+        ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
 
-        final KafkaConsumer consumer = newConsumer(time, client, metadata, assignor,
-                rebalanceTimeoutMs, sessionTimeoutMs, heartbeatIntervalMs, true, autoCommitIntervalMs);
+        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());
+        client.prepareResponseFrom(offsetResponse(Collections.singletonMap(tp0, offset1), Errors.NONE), coordinator);
+        assertEquals(offset1, consumer.committed(Collections.singleton(tp0)).get(tp0).offset());
 
         consumer.assign(Arrays.asList(tp0, tp1));
 
@@ -515,43 +762,106 @@ 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());
-        consumer.close(0, TimeUnit.MILLISECONDS);
+        assertEquals(offset2, consumer.committed(Collections.singleton(tp1)).get(tp1).offset());
+        consumer.close(Duration.ofMillis(0));
     }
 
     @Test
-    public void testAutoCommitSentBeforePositionUpdate() {
-        int rebalanceTimeoutMs = 60000;
-        int sessionTimeoutMs = 30000;
-        int heartbeatIntervalMs = 3000;
+    public void testFetchStableOffsetThrowInCommitted() {
+        assertThrows(UnsupportedVersionException.class, () -> setupThrowableConsumer().committed(Collections.singleton(tp0)));
+    }
+
+    @Test
+    public void testFetchStableOffsetThrowInPoll() {
+        assertThrows(UnsupportedVersionException.class, () -> setupThrowableConsumer().poll(Duration.ZERO));
+    }
+
+    @Test
+    public void testFetchStableOffsetThrowInPosition() {
+        assertThrows(UnsupportedVersionException.class, () -> setupThrowableConsumer().position(tp0));
+    }
+
+    private KafkaConsumer setupThrowableConsumer() {
+        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));
+        client.setNodeApiVersions(NodeApiVersions.create(ApiKeys.OFFSET_FETCH.id, (short) 0, (short) 6));
+
+        Node node = metadata.fetch().nodes().get(0);
+
+        ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
 
-        // adjust auto commit interval lower than heartbeat so we don't need to deal with
-        // a concurrent heartbeat request
-        int autoCommitIntervalMs = 1000;
+        KafkaConsumer consumer = newConsumer(
+            time, client, subscription, metadata, assignor, true, groupId, groupInstanceId, true);
+        consumer.assign(singletonList(tp0));
+
+        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, offset1), Errors.NONE), coordinator);
+        return consumer;
+    }
+
+    @Test
+    public void testNoCommittedOffsets() {
+        long offset1 = 10000;
 
         Time time = new MockTime();
-        Cluster cluster = TestUtils.singletonCluster(topic, 1);
-        Node node = cluster.nodes().get(0);
+        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());
 
-        Metadata metadata = createMetadata();
-        metadata.update(cluster, Collections.emptySet(), time.milliseconds());
+        // 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));
+    }
 
+    @Test
+    public void testAutoCommitSentBeforePositionUpdate() {
+        Time time = new MockTime();
+        SubscriptionState subscription = new SubscriptionState(new LogContext(), OffsetResetStrategy.EARLIEST);
+        ConsumerMetadata metadata = createMetadata(subscription);
         MockClient client = new MockClient(time, metadata);
-        client.setNode(node);
-        PartitionAssignor assignor = new RoundRobinAssignor();
 
-        final KafkaConsumer consumer = newConsumer(time, client, metadata, assignor,
-                rebalanceTimeoutMs, sessionTimeoutMs, heartbeatIntervalMs, true, autoCommitIntervalMs);
+        initMetadata(client, Collections.singletonMap(topic, 1));
+        Node node = metadata.fetch().nodes().get(0);
 
-        consumer.subscribe(Arrays.asList(topic), getConsumerRebalanceListener(consumer));
-        Node coordinator = prepareRebalance(client, node, assignor, Arrays.asList(tp0), null);
+        ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
 
-        consumer.poll(0);
+        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);
+
+        consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));
+        consumer.poll(Duration.ZERO);
 
         // respond to the outstanding fetch so that we have data available on the next poll
         client.respondFrom(fetchResponse(tp0, 0, 5), node);
@@ -564,124 +874,98 @@ public void testAutoCommitSentBeforePositionUpdate() {
         // no data has been returned to the user yet, so the committed offset should be 0
         AtomicBoolean commitReceived = prepareOffsetCommitResponse(client, coordinator, tp0, 0);
 
-        consumer.poll(0);
+        consumer.poll(Duration.ZERO);
 
         assertTrue(commitReceived.get());
-        consumer.close(0, TimeUnit.MILLISECONDS);
+        consumer.close(Duration.ofMillis(0));
     }
 
     @Test
     public void testRegexSubscription() {
-        int rebalanceTimeoutMs = 60000;
-        int sessionTimeoutMs = 30000;
-        int heartbeatIntervalMs = 3000;
-        int autoCommitIntervalMs = 1000;
-
         String unmatchedTopic = "unmatched";
-
         Time time = new MockTime();
-
-        Map topicMetadata = new HashMap<>();
-        topicMetadata.put(topic, 1);
-        topicMetadata.put(unmatchedTopic, 1);
-
-        Cluster cluster = TestUtils.clusterWith(1, topicMetadata);
-        Metadata metadata = createMetadata();
-        Node node = cluster.nodes().get(0);
-
+        SubscriptionState subscription = new SubscriptionState(new LogContext(), OffsetResetStrategy.EARLIEST);
+        ConsumerMetadata metadata = createMetadata(subscription);
         MockClient client = new MockClient(time, metadata);
-        client.setNode(node);
-        PartitionAssignor assignor = new RoundRobinAssignor();
 
-        final KafkaConsumer consumer = newConsumer(time, client, metadata, assignor,
-                rebalanceTimeoutMs, sessionTimeoutMs, heartbeatIntervalMs, true, autoCommitIntervalMs);
+        Map partitionCounts = new HashMap<>();
+        partitionCounts.put(topic, 1);
+        partitionCounts.put(unmatchedTopic, 1);
+        initMetadata(client, partitionCounts);
+        Node node = metadata.fetch().nodes().get(0);
 
+        ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
 
+        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));
 
-        client.prepareMetadataUpdate(cluster, Collections.emptySet());
+        client.prepareMetadataUpdate(RequestTestUtils.metadataUpdateWith(1, partitionCounts));
+
+        consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));
 
-        consumer.poll(0);
         assertEquals(singleton(topic), consumer.subscription());
         assertEquals(singleton(tp0), consumer.assignment());
-        consumer.close(0, TimeUnit.MILLISECONDS);
+        consumer.close(Duration.ofMillis(0));
     }
 
     @Test
     public void testChangingRegexSubscription() {
-        int rebalanceTimeoutMs = 60000;
-        int sessionTimeoutMs = 30000;
-        int heartbeatIntervalMs = 3000;
-        int autoCommitIntervalMs = 1000;
-        PartitionAssignor assignor = new RoundRobinAssignor();
+        ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
 
         String otherTopic = "other";
         TopicPartition otherTopicPartition = new TopicPartition(otherTopic, 0);
 
         Time time = new MockTime();
-
-        Map topicMetadata = new HashMap<>();
-        topicMetadata.put(topic, 1);
-        topicMetadata.put(otherTopic, 1);
-
-        Cluster cluster = TestUtils.clusterWith(1, topicMetadata);
-        Metadata metadata = createMetadata();
-        Node node = cluster.nodes().get(0);
-
+        SubscriptionState subscription = new SubscriptionState(new LogContext(), OffsetResetStrategy.EARLIEST);
+        ConsumerMetadata metadata = createMetadata(subscription);
         MockClient client = new MockClient(time, metadata);
-        client.setNode(node);
-        client.cluster(cluster);
 
-        metadata.update(cluster, Collections.emptySet(), time.milliseconds());
+        Map partitionCounts = new HashMap<>();
+        partitionCounts.put(topic, 1);
+        partitionCounts.put(otherTopic, 1);
+        initMetadata(client, partitionCounts);
+        Node node = metadata.fetch().nodes().get(0);
 
-        final KafkaConsumer consumer = newConsumer(time, client, metadata, assignor,
-                rebalanceTimeoutMs, sessionTimeoutMs, heartbeatIntervalMs, false, autoCommitIntervalMs);
+        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));
 
-        consumer.poll(0);
+        consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));
+        consumer.poll(Duration.ZERO);
+
         assertEquals(singleton(topic), consumer.subscription());
 
         consumer.subscribe(Pattern.compile(otherTopic), getConsumerRebalanceListener(consumer));
 
+        client.prepareMetadataUpdate(RequestTestUtils.metadataUpdateWith(1, partitionCounts));
         prepareRebalance(client, node, singleton(otherTopic), assignor, singletonList(otherTopicPartition), coordinator);
-        consumer.poll(0);
+        consumer.poll(Duration.ZERO);
 
         assertEquals(singleton(otherTopic), consumer.subscription());
-        consumer.close(0, TimeUnit.MILLISECONDS);
+        consumer.close(Duration.ofMillis(0));
     }
 
     @Test
     public void testWakeupWithFetchDataAvailable() throws Exception {
-        int rebalanceTimeoutMs = 60000;
-        final int sessionTimeoutMs = 30000;
-        int heartbeatIntervalMs = 3000;
-
-        // adjust auto commit interval lower than heartbeat so we don't need to deal with
-        // a concurrent heartbeat request
-        int autoCommitIntervalMs = 1000;
-
         final Time time = new MockTime();
-        Cluster cluster = TestUtils.singletonCluster(topic, 1);
-        Node node = cluster.nodes().get(0);
-
-        Metadata metadata = createMetadata();
-        metadata.update(cluster, Collections.emptySet(), time.milliseconds());
-
+        SubscriptionState subscription = new SubscriptionState(new LogContext(), OffsetResetStrategy.EARLIEST);
+        ConsumerMetadata metadata = createMetadata(subscription);
         MockClient client = new MockClient(time, metadata);
-        client.setNode(node);
-        PartitionAssignor assignor = new RoundRobinAssignor();
 
-        final KafkaConsumer consumer = newConsumer(time, client, metadata, assignor,
-                rebalanceTimeoutMs, sessionTimeoutMs, heartbeatIntervalMs, true, autoCommitIntervalMs);
+        initMetadata(client, Collections.singletonMap(topic, 1));
+        Node node = metadata.fetch().nodes().get(0);
 
-        consumer.subscribe(Arrays.asList(topic), getConsumerRebalanceListener(consumer));
-        prepareRebalance(client, node, assignor, Arrays.asList(tp0), null);
+        ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
 
-        consumer.poll(0);
+        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true, groupInstanceId);
+        consumer.subscribe(singleton(topic), getConsumerRebalanceListener(consumer));
+        prepareRebalance(client, node, assignor, singletonList(tp0), null);
+
+        consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));
+        consumer.poll(Duration.ZERO);
 
         // respond to the outstanding fetch so that we have data available on the next poll
         client.respondFrom(fetchResponse(tp0, 0, 5), node);
@@ -689,92 +973,65 @@ public void testWakeupWithFetchDataAvailable() throws Exception {
 
         consumer.wakeup();
 
-        try {
-            consumer.poll(0);
-            fail();
-        } catch (WakeupException e) {
-        }
+        assertThrows(WakeupException.class, () -> consumer.poll(Duration.ZERO));
 
         // make sure the position hasn't been updated
         assertEquals(0, consumer.position(tp0));
 
         // the next poll should return the completed fetch
-        ConsumerRecords records = consumer.poll(0);
+        ConsumerRecords records = consumer.poll(Duration.ZERO);
         assertEquals(5, records.count());
         // Increment time asynchronously to clear timeouts in closing the consumer
         final ScheduledExecutorService exec = Executors.newSingleThreadScheduledExecutor();
-        exec.scheduleAtFixedRate(new Runnable() {
-            @Override
-            public void run() {
-                time.sleep(sessionTimeoutMs);
-            }
-        }, 0L, 10L, TimeUnit.MILLISECONDS);
+        exec.scheduleAtFixedRate(() -> time.sleep(sessionTimeoutMs), 0L, 10L, TimeUnit.MILLISECONDS);
         consumer.close();
         exec.shutdownNow();
         exec.awaitTermination(5L, TimeUnit.SECONDS);
     }
 
     @Test
-    public void testPollThrowsInterruptExceptionIfInterrupted() throws Exception {
-        int rebalanceTimeoutMs = 60000;
-        int sessionTimeoutMs = 30000;
-        int heartbeatIntervalMs = 3000;
-
+    public void testPollThrowsInterruptExceptionIfInterrupted() {
         final Time time = new MockTime();
-        Cluster cluster = TestUtils.singletonCluster(topic, 1);
-        final Node node = cluster.nodes().get(0);
-
-        Metadata metadata = createMetadata();
-        metadata.update(cluster, Collections.emptySet(), time.milliseconds());
-
+        final SubscriptionState subscription = new SubscriptionState(new LogContext(), OffsetResetStrategy.EARLIEST);
+        final ConsumerMetadata metadata = createMetadata(subscription);
         final MockClient client = new MockClient(time, metadata);
-        client.setNode(node);
-        final PartitionAssignor assignor = new RoundRobinAssignor();
 
-        final KafkaConsumer consumer = newConsumer(time, client, metadata, assignor,
-                rebalanceTimeoutMs, sessionTimeoutMs, heartbeatIntervalMs, false, 0);
+        initMetadata(client, Collections.singletonMap(topic, 1));
+        Node node = metadata.fetch().nodes().get(0);
 
-        consumer.subscribe(Arrays.asList(topic), getConsumerRebalanceListener(consumer));
-        prepareRebalance(client, node, assignor, Arrays.asList(tp0), null);
+        final ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
+
+        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, false, groupInstanceId);
+        consumer.subscribe(singleton(topic), getConsumerRebalanceListener(consumer));
+        prepareRebalance(client, node, assignor, singletonList(tp0), null);
 
-        consumer.poll(0);
+        consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));
+        consumer.poll(Duration.ZERO);
 
         // interrupt the thread and call poll
         try {
             Thread.currentThread().interrupt();
-            expectedException.expect(InterruptException.class);
-            consumer.poll(0);
+            assertThrows(InterruptException.class, () -> consumer.poll(Duration.ZERO));
         } finally {
             // clear interrupted state again since this thread may be reused by JUnit
             Thread.interrupted();
+            consumer.close(Duration.ofMillis(0));
         }
-        consumer.close(0, TimeUnit.MILLISECONDS);
     }
 
     @Test
     public void fetchResponseWithUnexpectedPartitionIsIgnored() {
-        int rebalanceTimeoutMs = 60000;
-        int sessionTimeoutMs = 30000;
-        int heartbeatIntervalMs = 3000;
-
-        // adjust auto commit interval lower than heartbeat so we don't need to deal with
-        // a concurrent heartbeat request
-        int autoCommitIntervalMs = 1000;
-
         Time time = new MockTime();
-        Cluster cluster = TestUtils.singletonCluster(singletonMap(topic, 1));
-        Node node = cluster.nodes().get(0);
-
-        Metadata metadata = createMetadata();
-        metadata.update(cluster, Collections.emptySet(), time.milliseconds());
-
+        SubscriptionState subscription = new SubscriptionState(new LogContext(), OffsetResetStrategy.EARLIEST);
+        ConsumerMetadata metadata = createMetadata(subscription);
         MockClient client = new MockClient(time, metadata);
-        client.setNode(node);
-        PartitionAssignor assignor = new RangeAssignor();
 
-        final KafkaConsumer consumer = newConsumer(time, client, metadata, assignor,
-                rebalanceTimeoutMs, sessionTimeoutMs, heartbeatIntervalMs, true, autoCommitIntervalMs);
+        initMetadata(client, Collections.singletonMap(topic, 1));
+        Node node = metadata.fetch().nodes().get(0);
+
+        ConsumerPartitionAssignor assignor = new RangeAssignor();
 
+        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true, groupInstanceId);
         consumer.subscribe(singletonList(topic), getConsumerRebalanceListener(consumer));
 
         prepareRebalance(client, node, assignor, singletonList(tp0), null);
@@ -784,9 +1041,11 @@ public void fetchResponseWithUnexpectedPartitionIsIgnored() {
         fetches1.put(t2p0, new FetchInfo(0, 10)); // not assigned and not fetched
         client.prepareResponseFrom(fetchResponse(fetches1), node);
 
-        ConsumerRecords records = consumer.poll(0);
+        consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));
+
+        ConsumerRecords records = consumer.poll(Duration.ZERO);
         assertEquals(0, records.count());
-        consumer.close(0, TimeUnit.MILLISECONDS);
+        consumer.close(Duration.ofMillis(0));
     }
 
     /**
@@ -798,49 +1057,40 @@ public void fetchResponseWithUnexpectedPartitionIsIgnored() {
      */
     @Test
     public void testSubscriptionChangesWithAutoCommitEnabled() {
-        int rebalanceTimeoutMs = 60000;
-        int sessionTimeoutMs = 30000;
-        int heartbeatIntervalMs = 3000;
-
-        // adjust auto commit interval lower than heartbeat so we don't need to deal with
-        // a concurrent heartbeat request
-        int autoCommitIntervalMs = 1000;
-
         Time time = new MockTime();
+        SubscriptionState subscription = new SubscriptionState(new LogContext(), OffsetResetStrategy.EARLIEST);
+        ConsumerMetadata metadata = createMetadata(subscription);
+        MockClient client = new MockClient(time, metadata);
+
         Map tpCounts = new HashMap<>();
         tpCounts.put(topic, 1);
         tpCounts.put(topic2, 1);
         tpCounts.put(topic3, 1);
-        Cluster cluster = TestUtils.singletonCluster(tpCounts);
-        Node node = cluster.nodes().get(0);
+        initMetadata(client, tpCounts);
+        Node node = metadata.fetch().nodes().get(0);
 
-        Metadata metadata = createMetadata();
-        metadata.update(cluster, Collections.emptySet(), time.milliseconds());
+        ConsumerPartitionAssignor assignor = new RangeAssignor();
 
-        MockClient client = new MockClient(time, metadata);
-        client.setNode(node);
-        PartitionAssignor assignor = new RangeAssignor();
-
-        final KafkaConsumer consumer = newConsumer(time, client, metadata, assignor,
-                rebalanceTimeoutMs, sessionTimeoutMs, heartbeatIntervalMs, true, autoCommitIntervalMs);
+        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true, groupInstanceId);
 
         // initial subscription
         consumer.subscribe(Arrays.asList(topic, topic2), getConsumerRebalanceListener(consumer));
 
         // verify that subscription has changed but assignment is still unchanged
-        assertTrue(consumer.subscription().size() == 2);
+        assertEquals(2, consumer.subscription().size());
         assertTrue(consumer.subscription().contains(topic) && consumer.subscription().contains(topic2));
         assertTrue(consumer.assignment().isEmpty());
 
         // mock rebalance responses
         Node coordinator = prepareRebalance(client, node, assignor, Arrays.asList(tp0, t2p0), null);
 
-        consumer.poll(0);
+        consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));
+        consumer.poll(Duration.ZERO);
 
         // verify that subscription is still the same, and now assignment has caught up
-        assertTrue(consumer.subscription().size() == 2);
+        assertEquals(2, consumer.subscription().size());
         assertTrue(consumer.subscription().contains(topic) && consumer.subscription().contains(topic2));
-        assertTrue(consumer.assignment().size() == 2);
+        assertEquals(2, consumer.assignment().size());
         assertTrue(consumer.assignment().contains(tp0) && consumer.assignment().contains(t2p0));
 
         // mock a response to the outstanding fetch so that we have data available on the next poll
@@ -850,7 +1100,7 @@ public void testSubscriptionChangesWithAutoCommitEnabled() {
         client.respondFrom(fetchResponse(fetches1), node);
         client.poll(0, time.milliseconds());
 
-        ConsumerRecords records = consumer.poll(0);
+        ConsumerRecords records = consumer.poll(Duration.ofMillis(1));
 
         // clear out the prefetch so it doesn't interfere with the rest of the test
         fetches1.put(tp0, new FetchInfo(1, 0));
@@ -867,9 +1117,9 @@ public void testSubscriptionChangesWithAutoCommitEnabled() {
         consumer.subscribe(Arrays.asList(topic, topic3), getConsumerRebalanceListener(consumer));
 
         // verify that subscription has changed but assignment is still unchanged
-        assertTrue(consumer.subscription().size() == 2);
+        assertEquals(2, consumer.subscription().size());
         assertTrue(consumer.subscription().contains(topic) && consumer.subscription().contains(topic3));
-        assertTrue(consumer.assignment().size() == 2);
+        assertEquals(2, consumer.assignment().size());
         assertTrue(consumer.assignment().contains(tp0) && consumer.assignment().contains(t2p0));
 
         // mock the offset commit response for to be revoked partitions
@@ -887,7 +1137,7 @@ public void testSubscriptionChangesWithAutoCommitEnabled() {
         fetches2.put(t3p0, new FetchInfo(0, 100));
         client.prepareResponse(fetchResponse(fetches2));
 
-        records = consumer.poll(0);
+        records = consumer.poll(Duration.ofMillis(1));
 
         // verify that the fetch occurred as expected
         assertEquals(101, records.count());
@@ -898,9 +1148,9 @@ public void testSubscriptionChangesWithAutoCommitEnabled() {
         assertTrue(commitReceived.get());
 
         // verify that subscription is still the same, and now assignment has caught up
-        assertTrue(consumer.subscription().size() == 2);
+        assertEquals(2, consumer.subscription().size());
         assertTrue(consumer.subscription().contains(topic) && consumer.subscription().contains(topic3));
-        assertTrue(consumer.assignment().size() == 2);
+        assertEquals(2, consumer.assignment().size());
         assertTrue(consumer.assignment().contains(tp0) && consumer.assignment().contains(t3p0));
 
         consumer.unsubscribe();
@@ -922,119 +1172,162 @@ public void testSubscriptionChangesWithAutoCommitEnabled() {
      */
     @Test
     public void testSubscriptionChangesWithAutoCommitDisabled() {
-        int rebalanceTimeoutMs = 60000;
-        int sessionTimeoutMs = 30000;
-        int heartbeatIntervalMs = 3000;
-        int autoCommitIntervalMs = 1000;
-
         Time time = new MockTime();
+        SubscriptionState subscription = new SubscriptionState(new LogContext(), OffsetResetStrategy.EARLIEST);
+        ConsumerMetadata metadata = createMetadata(subscription);
+        MockClient client = new MockClient(time, metadata);
+
         Map tpCounts = new HashMap<>();
         tpCounts.put(topic, 1);
         tpCounts.put(topic2, 1);
-        Cluster cluster = TestUtils.singletonCluster(tpCounts);
-        Node node = cluster.nodes().get(0);
+        initMetadata(client, tpCounts);
+        Node node = metadata.fetch().nodes().get(0);
 
-        Metadata metadata = createMetadata();
-        metadata.update(cluster, Collections.emptySet(), time.milliseconds());
+        ConsumerPartitionAssignor assignor = new RangeAssignor();
 
-        MockClient client = new MockClient(time, metadata);
-        client.setNode(node);
-        PartitionAssignor assignor = new RangeAssignor();
-
-        final KafkaConsumer consumer = newConsumer(time, client, metadata, assignor,
-                rebalanceTimeoutMs, sessionTimeoutMs, heartbeatIntervalMs, false, autoCommitIntervalMs);
-
-        // initial subscription
-        consumer.subscribe(Arrays.asList(topic), getConsumerRebalanceListener(consumer));
+        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, false, groupInstanceId);
 
-        // verify that subscription has changed but assignment is still unchanged
-        assertTrue(consumer.subscription().equals(Collections.singleton(topic)));
-        assertTrue(consumer.assignment().isEmpty());
+        initializeSubscriptionWithSingleTopic(consumer, getConsumerRebalanceListener(consumer));
 
         // mock rebalance responses
-        prepareRebalance(client, node, assignor, Arrays.asList(tp0), null);
+        prepareRebalance(client, node, assignor, singletonList(tp0), null);
 
-        consumer.poll(0);
+        consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));
+        consumer.poll(Duration.ZERO);
 
         // verify that subscription is still the same, and now assignment has caught up
-        assertTrue(consumer.subscription().equals(Collections.singleton(topic)));
-        assertTrue(consumer.assignment().equals(Collections.singleton(tp0)));
+        assertEquals(singleton(topic), consumer.subscription());
+        assertEquals(singleton(tp0), consumer.assignment());
 
-        consumer.poll(0);
+        consumer.poll(Duration.ZERO);
 
         // subscription change
-        consumer.subscribe(Arrays.asList(topic2), getConsumerRebalanceListener(consumer));
+        consumer.subscribe(singleton(topic2), getConsumerRebalanceListener(consumer));
 
         // verify that subscription has changed but assignment is still unchanged
-        assertTrue(consumer.subscription().equals(Collections.singleton(topic2)));
-        assertTrue(consumer.assignment().equals(Collections.singleton(tp0)));
+        assertEquals(singleton(topic2), consumer.subscription());
+        assertEquals(singleton(tp0), consumer.assignment());
 
         // the auto commit is disabled, so no offset commit request should be sent
         for (ClientRequest req: client.requests())
-            assertTrue(req.requestBuilder().apiKey() != ApiKeys.OFFSET_COMMIT);
+            assertNotSame(ApiKeys.OFFSET_COMMIT, req.requestBuilder().apiKey());
 
         // subscription change
         consumer.unsubscribe();
 
         // verify that subscription and assignment are both updated
-        assertTrue(consumer.subscription().isEmpty());
-        assertTrue(consumer.assignment().isEmpty());
+        assertEquals(Collections.emptySet(), consumer.subscription());
+        assertEquals(Collections.emptySet(), consumer.assignment());
 
         // the auto commit is disabled, so no offset commit request should be sent
         for (ClientRequest req: client.requests())
-            assertTrue(req.requestBuilder().apiKey() != ApiKeys.OFFSET_COMMIT);
+            assertNotSame(ApiKeys.OFFSET_COMMIT, req.requestBuilder().apiKey());
 
         client.requests().clear();
         consumer.close();
     }
 
     @Test
-    public void testManualAssignmentChangeWithAutoCommitEnabled() {
-        int rebalanceTimeoutMs = 60000;
-        int sessionTimeoutMs = 30000;
-        int heartbeatIntervalMs = 3000;
-        int autoCommitIntervalMs = 1000;
+    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.");
+
+        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();
+        SubscriptionState subscription = new SubscriptionState(new LogContext(), OffsetResetStrategy.EARLIEST);
+        ConsumerMetadata metadata = createMetadata(subscription);
+        MockClient client = new MockClient(time, metadata);
+
         Map tpCounts = new HashMap<>();
         tpCounts.put(topic, 1);
         tpCounts.put(topic2, 1);
-        Cluster cluster = TestUtils.singletonCluster(tpCounts);
-        Node node = cluster.nodes().get(0);
-
-        Metadata metadata = createMetadata();
-        metadata.update(cluster, Collections.emptySet(), time.milliseconds());
+        initMetadata(client, tpCounts);
+        Node node = metadata.fetch().nodes().get(0);
 
-        MockClient client = new MockClient(time, metadata);
-        client.setNode(node);
-        PartitionAssignor assignor = new RangeAssignor();
+        ConsumerPartitionAssignor assignor = new RangeAssignor();
 
-        final KafkaConsumer consumer = newConsumer(time, client, metadata, assignor,
-                rebalanceTimeoutMs, sessionTimeoutMs, heartbeatIntervalMs, true, autoCommitIntervalMs);
+        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
-        consumer.assign(Arrays.asList(tp0));
-        consumer.seekToBeginning(Arrays.asList(tp0));
+        consumer.assign(singleton(tp0));
+        consumer.seekToBeginning(singleton(tp0));
 
         // fetch offset for one topic
-        client.prepareResponseFrom(
-                offsetResponse(Collections.singletonMap(tp0, 0L), Errors.NONE),
-                coordinator);
-        assertEquals(0, consumer.committed(tp0).offset());
+        client.prepareResponseFrom(offsetResponse(Collections.singletonMap(tp0, 0L), Errors.NONE), coordinator);
+        assertEquals(0, consumer.committed(Collections.singleton(tp0)).get(tp0).offset());
 
         // verify that assignment immediately changes
-        assertTrue(consumer.assignment().equals(Collections.singleton(tp0)));
+        assertEquals(consumer.assignment(), singleton(tp0));
 
         // there shouldn't be any need to lookup the coordinator or fetch committed offsets.
         // we just lookup the starting position and send the record fetch.
-        client.prepareResponse(listOffsetsResponse(Collections.singletonMap(tp0, 10L), Errors.NONE));
+        client.prepareResponse(listOffsetsResponse(Collections.singletonMap(tp0, 10L)));
         client.prepareResponse(fetchResponse(tp0, 10L, 1));
 
-        ConsumerRecords records = consumer.poll(0);
+        ConsumerRecords records = consumer.poll(Duration.ofMillis(1));
+
         assertEquals(1, records.count());
         assertEquals(11L, consumer.position(tp0));
 
@@ -1042,10 +1335,10 @@ public void testManualAssignmentChangeWithAutoCommitEnabled() {
         AtomicBoolean commitReceived = prepareOffsetCommitResponse(client, coordinator, tp0, 11);
 
         // new manual assignment
-        consumer.assign(Arrays.asList(t2p0));
+        consumer.assign(singleton(t2p0));
 
         // verify that assignment immediately changes
-        assertTrue(consumer.assignment().equals(Collections.singleton(t2p0)));
+        assertEquals(consumer.assignment(), singleton(t2p0));
         // verify that the offset commits occurred as expected
         assertTrue(commitReceived.get());
 
@@ -1055,63 +1348,56 @@ public void testManualAssignmentChangeWithAutoCommitEnabled() {
 
     @Test
     public void testManualAssignmentChangeWithAutoCommitDisabled() {
-        int rebalanceTimeoutMs = 60000;
-        int sessionTimeoutMs = 30000;
-        int heartbeatIntervalMs = 3000;
-        int autoCommitIntervalMs = 1000;
-
         Time time = new MockTime();
+        SubscriptionState subscription = new SubscriptionState(new LogContext(), OffsetResetStrategy.EARLIEST);
+        ConsumerMetadata metadata = createMetadata(subscription);
+        MockClient client = new MockClient(time, metadata);
+
         Map tpCounts = new HashMap<>();
         tpCounts.put(topic, 1);
         tpCounts.put(topic2, 1);
-        Cluster cluster = TestUtils.singletonCluster(tpCounts);
-        Node node = cluster.nodes().get(0);
-
-        Metadata metadata = createMetadata();
-        metadata.update(cluster, Collections.emptySet(), time.milliseconds());
+        initMetadata(client, tpCounts);
+        Node node = metadata.fetch().nodes().get(0);
 
-        MockClient client = new MockClient(time, metadata);
-        client.setNode(node);
-        PartitionAssignor assignor = new RangeAssignor();
+        ConsumerPartitionAssignor assignor = new RangeAssignor();
 
-        final KafkaConsumer consumer = newConsumer(time, client, metadata, assignor,
-                rebalanceTimeoutMs, sessionTimeoutMs, heartbeatIntervalMs, false, autoCommitIntervalMs);
+        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
-        consumer.assign(Arrays.asList(tp0));
-        consumer.seekToBeginning(Arrays.asList(tp0));
+        consumer.assign(singleton(tp0));
+        consumer.seekToBeginning(singleton(tp0));
 
         // 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(Collections.singleton(tp0)));
+        assertEquals(consumer.assignment(), singleton(tp0));
 
         // there shouldn't be any need to lookup the coordinator or fetch committed offsets.
         // we just lookup the starting position and send the record fetch.
-        client.prepareResponse(listOffsetsResponse(Collections.singletonMap(tp0, 10L), Errors.NONE));
+        client.prepareResponse(listOffsetsResponse(Collections.singletonMap(tp0, 10L)));
         client.prepareResponse(fetchResponse(tp0, 10L, 1));
 
-        ConsumerRecords records = consumer.poll(0);
+        ConsumerRecords records = consumer.poll(Duration.ofMillis(1));
         assertEquals(1, records.count());
         assertEquals(11L, consumer.position(tp0));
 
         // new manual assignment
-        consumer.assign(Arrays.asList(t2p0));
+        consumer.assign(singleton(t2p0));
 
         // verify that assignment immediately changes
-        assertTrue(consumer.assignment().equals(Collections.singleton(t2p0)));
+        assertEquals(consumer.assignment(), singleton(t2p0));
 
         // the auto commit is disabled, so no offset commit request should be sent
         for (ClientRequest req : client.requests())
-            assertTrue(req.requestBuilder().apiKey() != ApiKeys.OFFSET_COMMIT);
+            assertNotSame(req.requestBuilder().apiKey(), ApiKeys.OFFSET_COMMIT);
 
         client.requests().clear();
         consumer.close();
@@ -1119,34 +1405,27 @@ public void testManualAssignmentChangeWithAutoCommitDisabled() {
 
     @Test
     public void testOffsetOfPausedPartitions() {
-        int rebalanceTimeoutMs = 60000;
-        int sessionTimeoutMs = 30000;
-        int heartbeatIntervalMs = 3000;
-        int autoCommitIntervalMs = 1000;
-
         Time time = new MockTime();
-        Cluster cluster = TestUtils.singletonCluster(topic, 2);
-        Node node = cluster.nodes().get(0);
+        SubscriptionState subscription = new SubscriptionState(new LogContext(), OffsetResetStrategy.EARLIEST);
+        ConsumerMetadata metadata = createMetadata(subscription);
+        MockClient client = new MockClient(time, metadata);
 
-        Metadata metadata = createMetadata();
-        metadata.update(cluster, Collections.emptySet(), time.milliseconds());
+        initMetadata(client, Collections.singletonMap(topic, 2));
+        Node node = metadata.fetch().nodes().get(0);
 
-        MockClient client = new MockClient(time, metadata);
-        client.setNode(node);
-        PartitionAssignor assignor = new RangeAssignor();
+        ConsumerPartitionAssignor assignor = new RangeAssignor();
 
-        final KafkaConsumer consumer = newConsumer(time, client, metadata, assignor,
-                rebalanceTimeoutMs, sessionTimeoutMs, heartbeatIntervalMs, true, autoCommitIntervalMs);
+        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
         Set partitions = Utils.mkSet(tp0, tp1);
         consumer.assign(partitions);
         // verify consumer's assignment
-        assertTrue(consumer.assignment().equals(partitions));
+        assertEquals(partitions, consumer.assignment());
 
         consumer.pause(partitions);
         consumer.seekToEnd(partitions);
@@ -1157,18 +1436,18 @@ 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<>();
         offsetResponse.put(tp0, 3L);
         offsetResponse.put(tp1, 3L);
-        client.prepareResponse(listOffsetsResponse(offsetResponse, Errors.NONE));
+        client.prepareResponse(listOffsetsResponse(offsetResponse));
         assertEquals(3L, consumer.position(tp0));
         assertEquals(3L, consumer.position(tp1));
 
@@ -1179,27 +1458,24 @@ public void testOffsetOfPausedPartitions() {
 
     @Test(expected = IllegalStateException.class)
     public void testPollWithNoSubscription() {
-        KafkaConsumer consumer = newConsumer();
-        try {
-            consumer.poll(0);
-        } finally {
-            consumer.close();
+        try (KafkaConsumer consumer = newConsumer((String) null)) {
+            consumer.poll(Duration.ZERO);
         }
     }
 
     @Test(expected = IllegalStateException.class)
     public void testPollWithEmptySubscription() {
-        try (KafkaConsumer consumer = newConsumer()) {
-            consumer.subscribe(Collections.emptyList());
-            consumer.poll(0);
+        try (KafkaConsumer consumer = newConsumer(groupId)) {
+            consumer.subscribe(Collections.emptyList());
+            consumer.poll(Duration.ZERO);
         }
     }
 
     @Test(expected = IllegalStateException.class)
     public void testPollWithEmptyUserAssignment() {
-        try (KafkaConsumer consumer = newConsumer()) {
-            consumer.assign(Collections.emptySet());
-            consumer.poll(0);
+        try (KafkaConsumer consumer = newConsumer(groupId)) {
+            consumer.assign(Collections.emptySet());
+            consumer.poll(Duration.ZERO);
         }
     }
 
@@ -1208,13 +1484,13 @@ public void testGracefulClose() throws Exception {
         Map response = new HashMap<>();
         response.put(tp0, Errors.NONE);
         OffsetCommitResponse commitResponse = offsetCommitResponse(response);
-        LeaveGroupResponse leaveGroupResponse = new LeaveGroupResponse(Errors.NONE);
+        LeaveGroupResponse leaveGroupResponse = new LeaveGroupResponse(new LeaveGroupResponseData().setErrorCode(Errors.NONE.code()));
         consumerCloseTest(5000, Arrays.asList(commitResponse, leaveGroupResponse), 0, false);
     }
 
     @Test
     public void testCloseTimeout() throws Exception {
-        consumerCloseTest(5000, Collections.emptyList(), 5000, false);
+        consumerCloseTest(5000, Collections.emptyList(), 5000, false);
     }
 
     @Test
@@ -1222,160 +1498,215 @@ public void testLeaveGroupTimeout() throws Exception {
         Map response = new HashMap<>();
         response.put(tp0, Errors.NONE);
         OffsetCommitResponse commitResponse = offsetCommitResponse(response);
-        consumerCloseTest(5000, Arrays.asList(commitResponse), 5000, false);
+        consumerCloseTest(5000, singletonList(commitResponse), 5000, false);
     }
 
     @Test
     public void testCloseNoWait() throws Exception {
-        consumerCloseTest(0, Collections.emptyList(), 0, false);
+        consumerCloseTest(0, Collections.emptyList(), 0, false);
     }
 
     @Test
     public void testCloseInterrupt() throws Exception {
-        consumerCloseTest(Long.MAX_VALUE, Collections.emptyList(), 0, true);
+        consumerCloseTest(Long.MAX_VALUE, Collections.emptyList(), 0, true);
     }
 
     @Test
-    public void closeShouldBeIdempotent() {
-        KafkaConsumer consumer = newConsumer();
+    public void testCloseShouldBeIdempotent() {
+        KafkaConsumer consumer = newConsumer((String) null);
         consumer.close();
         consumer.close();
         consumer.close();
     }
 
+    @Test
+    public void testOperationsBySubscribingConsumerWithDefaultGroupId() {
+        try {
+            newConsumer(null, Optional.of(Boolean.TRUE));
+            fail("Expected an InvalidConfigurationException");
+        } catch (KafkaException e) {
+            assertEquals(InvalidConfigurationException.class, e.getCause().getClass());
+        }
+
+        try {
+            newConsumer((String) null).subscribe(Collections.singleton(topic));
+            fail("Expected an InvalidGroupIdException");
+        } catch (InvalidGroupIdException e) {
+            // OK, expected
+        }
+
+        try {
+            newConsumer((String) null).committed(Collections.singleton(tp0)).get(tp0);
+            fail("Expected an InvalidGroupIdException");
+        } catch (InvalidGroupIdException e) {
+            // OK, expected
+        }
+
+        try {
+            newConsumer((String) null).commitAsync();
+            fail("Expected an InvalidGroupIdException");
+        } catch (InvalidGroupIdException e) {
+            // OK, expected
+        }
+
+        try {
+            newConsumer((String) null).commitSync();
+            fail("Expected an InvalidGroupIdException");
+        } catch (InvalidGroupIdException e) {
+            // OK, expected
+        }
+    }
+
+    @Test
+    public void testOperationsByAssigningConsumerWithDefaultGroupId() {
+        KafkaConsumer consumer = newConsumer((String) null);
+        consumer.assign(singleton(tp0));
+
+        try {
+            consumer.committed(Collections.singleton(tp0)).get(tp0);
+            fail("Expected an InvalidGroupIdException");
+        } catch (InvalidGroupIdException e) {
+            // OK, expected
+        }
+
+        try {
+            consumer.commitAsync();
+            fail("Expected an InvalidGroupIdException");
+        } catch (InvalidGroupIdException e) {
+            // OK, expected
+        }
+
+        try {
+            consumer.commitSync();
+            fail("Expected an InvalidGroupIdException");
+        } catch (InvalidGroupIdException e) {
+            // OK, expected
+        }
+    }
+
     @Test
     public void testMetricConfigRecordingLevel() {
         Properties props = new Properties();
         props.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9000");
-        try (KafkaConsumer consumer = new KafkaConsumer<>(props, new ByteArrayDeserializer(), new ByteArrayDeserializer())) {
+        try (KafkaConsumer consumer = new KafkaConsumer<>(props, new ByteArrayDeserializer(), new ByteArrayDeserializer())) {
             assertEquals(Sensor.RecordingLevel.INFO, consumer.metrics.config().recordLevel());
         }
 
         props.put(ConsumerConfig.METRICS_RECORDING_LEVEL_CONFIG, "DEBUG");
-        try (KafkaConsumer consumer = new KafkaConsumer<>(props, new ByteArrayDeserializer(), new ByteArrayDeserializer())) {
+        try (KafkaConsumer consumer = new KafkaConsumer<>(props, new ByteArrayDeserializer(), new ByteArrayDeserializer())) {
             assertEquals(Sensor.RecordingLevel.DEBUG, consumer.metrics.config().recordLevel());
         }
     }
 
     @Test
-    public void shouldAttemptToRejoinGroupAfterSyncGroupFailed() throws Exception {
-        int rebalanceTimeoutMs = 60000;
-        int sessionTimeoutMs = 30000;
-        int heartbeatIntervalMs = 500;
-
+    public void testShouldAttemptToRejoinGroupAfterSyncGroupFailed() throws Exception {
         Time time = new MockTime();
-        Cluster cluster = TestUtils.singletonCluster(topic, 1);
-        Node node = cluster.nodes().get(0);
-
-        Metadata metadata = new Metadata(0, Long.MAX_VALUE, false);
-        metadata.update(cluster, Collections.emptySet(), time.milliseconds());
-
+        SubscriptionState subscription = new SubscriptionState(new LogContext(), OffsetResetStrategy.EARLIEST);
+        ConsumerMetadata metadata = createMetadata(subscription);
         MockClient client = new MockClient(time, metadata);
-        client.setNode(node);
-        PartitionAssignor assignor = new RoundRobinAssignor();
 
-        final KafkaConsumer consumer = newConsumer(time, client, metadata, assignor,
-                                                                   rebalanceTimeoutMs, sessionTimeoutMs, heartbeatIntervalMs, false, 1000);
+        initMetadata(client, Collections.singletonMap(topic, 1));
+        Node node = metadata.fetch().nodes().get(0);
+
+        ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
 
-        consumer.subscribe(Collections.singleton(topic), getConsumerRebalanceListener(consumer));
-        client.prepareResponseFrom(new FindCoordinatorResponse(Errors.NONE, node), node);
+        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, false, groupInstanceId);
+        consumer.subscribe(singleton(topic), getConsumerRebalanceListener(consumer));
+        client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, node), node);
         Node coordinator = new Node(Integer.MAX_VALUE - node.id(), node.host(), node.port());
 
 
-        client.prepareResponseFrom(joinGroupFollowerResponse(assignor, 1, "memberId", "leaderId", Errors.NONE), coordinator);
-        client.prepareResponseFrom(syncGroupResponse(Collections.singletonList(tp0), Errors.NONE), coordinator);
+        client.prepareResponseFrom(joinGroupFollowerResponse(assignor, 1, memberId, leaderId, Errors.NONE), coordinator);
+        client.prepareResponseFrom(syncGroupResponse(singletonList(tp0), Errors.NONE), coordinator);
 
         client.prepareResponseFrom(fetchResponse(tp0, 0, 1), node);
         client.prepareResponseFrom(fetchResponse(tp0, 1, 0), node);
 
-        consumer.poll(0);
+        consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));
+        consumer.poll(Duration.ZERO);
 
         // heartbeat fails due to rebalance in progress
-        client.prepareResponseFrom(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                return true;
-            }
-        }, new HeartbeatResponse(Errors.REBALANCE_IN_PROGRESS), coordinator);
+        client.prepareResponseFrom(body -> true, new HeartbeatResponse(
+            new HeartbeatResponseData().setErrorCode(Errors.REBALANCE_IN_PROGRESS.code())), coordinator);
 
         // join group
-        final ByteBuffer byteBuffer = ConsumerProtocol.serializeSubscription(new PartitionAssignor.Subscription(Collections.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));
+        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(Collections.singletonList(tp0), Errors.NONE), coordinator, true);
+        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);
-        client.prepareResponseFrom(syncGroupResponse(Collections.singletonList(tp0), Errors.NONE), coordinator);
+        client.prepareResponseFrom(joinGroupFollowerResponse(assignor, 1, memberId, leaderId, Errors.NONE), coordinator);
+        client.prepareResponseFrom(syncGroupResponse(singletonList(tp0), Errors.NONE), coordinator);
 
-        client.prepareResponseFrom(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(final AbstractRequest body) {
-                return body instanceof FetchRequest && ((FetchRequest) body).fetchData().containsKey(tp0);
-            }
-        }, fetchResponse(tp0, 1, 1), node);
+        client.prepareResponseFrom(body -> body instanceof FetchRequest 
+            && ((FetchRequest) body).fetchData().containsKey(tp0), fetchResponse(tp0, 1, 1), node);
         time.sleep(heartbeatIntervalMs);
         Thread.sleep(heartbeatIntervalMs);
-        final ConsumerRecords records = consumer.poll(0);
+        consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));
+        final ConsumerRecords records = consumer.poll(Duration.ZERO);
         assertFalse(records.isEmpty());
-        consumer.close(0, TimeUnit.MILLISECONDS);
+        consumer.close(Duration.ofMillis(0));
     }
 
     private void consumerCloseTest(final long closeTimeoutMs,
-            List responses,
-            long waitMs,
-            boolean interrupt) throws Exception {
-        int rebalanceTimeoutMs = 60000;
-        int sessionTimeoutMs = 30000;
-        int heartbeatIntervalMs = 5000;
-
+                                   List responses,
+                                   long waitMs,
+                                   boolean interrupt) throws Exception {
         Time time = new MockTime();
-        Cluster cluster = TestUtils.singletonCluster(topic, 1);
-        Node node = cluster.nodes().get(0);
+        SubscriptionState subscription = new SubscriptionState(new LogContext(), OffsetResetStrategy.EARLIEST);
+        ConsumerMetadata metadata = createMetadata(subscription);
+        MockClient client = new MockClient(time, metadata);
 
-        Metadata metadata = createMetadata();
-        metadata.update(cluster, Collections.emptySet(), time.milliseconds());
+        initMetadata(client, Collections.singletonMap(topic, 1));
+        Node node = metadata.fetch().nodes().get(0);
 
-        MockClient client = new MockClient(time, metadata);
-        client.setNode(node);
-        PartitionAssignor assignor = new RoundRobinAssignor();
+        ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
 
-        final KafkaConsumer consumer = newConsumer(time, client, metadata, assignor,
-                rebalanceTimeoutMs, sessionTimeoutMs, heartbeatIntervalMs, false, 1000);
+        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);
 
-        consumer.subscribe(Arrays.asList(topic), getConsumerRebalanceListener(consumer));
-        Node coordinator = prepareRebalance(client, node, assignor, Arrays.asList(tp0), null);
+        client.prepareMetadataUpdate(RequestTestUtils.metadataUpdateWith(1, Collections.singletonMap(topic, 1)));
 
-        client.prepareMetadataUpdate(cluster, Collections.emptySet());
+        consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));
 
         // Poll with responses
         client.prepareResponseFrom(fetchResponse(tp0, 0, 1), node);
         client.prepareResponseFrom(fetchResponse(tp0, 1, 0), node);
-        consumer.poll(0);
+        consumer.poll(Duration.ZERO);
 
         // Initiate close() after a commit request on another thread.
         // Kafka consumer is single-threaded, but the implementation allows calls on a
         // different thread as long as the calls are not executed concurrently. So this is safe.
         ExecutorService executor = Executors.newSingleThreadExecutor();
-        final AtomicReference closeException = new AtomicReference();
+        final AtomicReference closeException = new AtomicReference<>();
         try {
-            Future future = executor.submit(new Runnable() {
-                @Override
-                public void run() {
-                    consumer.commitAsync();
-                    try {
-                        consumer.close(closeTimeoutMs, TimeUnit.MILLISECONDS);
-                    } catch (Exception e) {
-                        closeException.set(e);
-                    }
+            Future future = executor.submit(() -> {
+                consumer.commitAsync();
+                try {
+                    consumer.close(Duration.ofMillis(closeTimeoutMs));
+                } catch (Exception e) {
+                    closeException.set(e);
                 }
             });
 
@@ -1409,26 +1740,317 @@ public void run() {
 
             if (waitMs > 0)
                 time.sleep(waitMs);
-            if (interrupt)
+            if (interrupt) {
                 assertTrue("Close terminated prematurely", future.cancel(true));
 
-            // Make sure that close task completes and another task can be run on the single threaded executor
-            executor.submit(new Runnable() {
-                @Override
-                public void run() {
-                }
-            }).get(500, TimeUnit.MILLISECONDS);
+                TestUtils.waitForCondition(
+                    () -> closeException.get() != null, "InterruptException did not occur within timeout.");
 
-            if (!interrupt) {
+                assertTrue("Expected exception not thrown " + closeException, closeException.get() instanceof InterruptException);
+            } else {
                 future.get(500, TimeUnit.MILLISECONDS); // Should succeed without TimeoutException or ExecutionException
                 assertNull("Unexpected exception during close", closeException.get());
-            } else
-                assertTrue("Expected exception not thrown " + closeException, closeException.get() instanceof InterruptException);
+            }
         } finally {
             executor.shutdownNow();
         }
     }
 
+    @Test(expected = AuthenticationException.class)
+    public void testPartitionsForAuthenticationFailure() {
+        final KafkaConsumer consumer = consumerWithPendingAuthenticationError();
+        consumer.partitionsFor("some other topic");
+    }
+
+    @Test(expected = AuthenticationException.class)
+    public void testBeginningOffsetsAuthenticationFailure() {
+        final KafkaConsumer consumer = consumerWithPendingAuthenticationError();
+        consumer.beginningOffsets(Collections.singleton(tp0));
+    }
+
+    @Test(expected = AuthenticationException.class)
+    public void testEndOffsetsAuthenticationFailure() {
+        final KafkaConsumer consumer = consumerWithPendingAuthenticationError();
+        consumer.endOffsets(Collections.singleton(tp0));
+    }
+
+    @Test(expected = AuthenticationException.class)
+    public void testPollAuthenticationFailure() {
+        final KafkaConsumer consumer = consumerWithPendingAuthenticationError();
+        consumer.subscribe(singleton(topic));
+        consumer.poll(Duration.ZERO);
+    }
+
+    @Test(expected = AuthenticationException.class)
+    public void testOffsetsForTimesAuthenticationFailure() {
+        final KafkaConsumer consumer = consumerWithPendingAuthenticationError();
+        consumer.offsetsForTimes(singletonMap(tp0, 0L));
+    }
+
+    @Test(expected = AuthenticationException.class)
+    public void testCommitSyncAuthenticationFailure() {
+        final KafkaConsumer consumer = consumerWithPendingAuthenticationError();
+        Map offsets = new HashMap<>();
+        offsets.put(tp0, new OffsetAndMetadata(10L));
+        consumer.commitSync(offsets);
+    }
+
+    @Test(expected = AuthenticationException.class)
+    public void testCommittedAuthenticationFailure() {
+        final KafkaConsumer consumer = consumerWithPendingAuthenticationError();
+        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());
+    }
+
+    @Test
+    public void testReturnRecordsDuringRebalance() throws InterruptedException {
+        Time time = new MockTime(1L);
+        SubscriptionState subscription = new SubscriptionState(new LogContext(), OffsetResetStrategy.EARLIEST);
+        ConsumerMetadata metadata = createMetadata(subscription);
+        MockClient client = new MockClient(time, metadata);
+        ConsumerPartitionAssignor assignor = new CooperativeStickyAssignor();
+        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true, groupInstanceId);
+
+        initMetadata(client, Utils.mkMap(Utils.mkEntry(topic, 1), Utils.mkEntry(topic2, 1), Utils.mkEntry(topic3, 1)));
+
+        consumer.subscribe(Arrays.asList(topic, topic2), getConsumerRebalanceListener(consumer));
+
+        Node node = metadata.fetch().nodes().get(0);
+        Node coordinator = prepareRebalance(client, node, assignor, Arrays.asList(tp0, t2p0), null);
+
+        // a poll with non-zero milliseconds would complete three round-trips (discover, join, sync)
+        TestUtils.waitForCondition(() -> {
+            consumer.poll(Duration.ofMillis(100L));
+            return consumer.assignment().equals(Utils.mkSet(tp0, t2p0));
+        }, "Does not complete rebalance in time");
+
+        assertEquals(Utils.mkSet(topic, topic2), consumer.subscription());
+        assertEquals(Utils.mkSet(tp0, t2p0), consumer.assignment());
+
+        // prepare a response of the outstanding fetch so that we have data available on the next poll
+        Map fetches1 = new HashMap<>();
+        fetches1.put(tp0, new FetchInfo(0, 1));
+        fetches1.put(t2p0, new FetchInfo(0, 10));
+        client.respondFrom(fetchResponse(fetches1), node);
+
+        ConsumerRecords records = consumer.poll(Duration.ZERO);
+
+        // verify that the fetch occurred as expected
+        assertEquals(11, records.count());
+        assertEquals(1L, consumer.position(tp0));
+        assertEquals(10L, consumer.position(t2p0));
+
+        // prepare the next response of the prefetch
+        fetches1.clear();
+        fetches1.put(tp0, new FetchInfo(1, 1));
+        fetches1.put(t2p0, new FetchInfo(10, 20));
+        client.respondFrom(fetchResponse(fetches1), node);
+
+        // subscription change
+        consumer.subscribe(Arrays.asList(topic, topic3), getConsumerRebalanceListener(consumer));
+
+        // verify that subscription has changed but assignment is still unchanged
+        assertEquals(Utils.mkSet(topic, topic3), consumer.subscription());
+        assertEquals(Utils.mkSet(tp0, t2p0), consumer.assignment());
+
+        // mock the offset commit response for to be revoked partitions
+        Map partitionOffsets1 = new HashMap<>();
+        partitionOffsets1.put(t2p0, 10L);
+        AtomicBoolean commitReceived = prepareOffsetCommitResponse(client, coordinator, partitionOffsets1);
+
+        // poll once which would not complete the rebalance
+        records = consumer.poll(Duration.ZERO);
+
+        // clear out the prefetch so it doesn't interfere with the rest of the test
+        fetches1.clear();
+        fetches1.put(tp0, new FetchInfo(2, 1));
+        client.respondFrom(fetchResponse(fetches1), node);
+
+        // verify that the fetch still occurred as expected
+        assertEquals(Utils.mkSet(topic, topic3), consumer.subscription());
+        assertEquals(Collections.singleton(tp0), consumer.assignment());
+        assertEquals(1, records.count());
+        assertEquals(2L, consumer.position(tp0));
+
+        // verify that the offset commits occurred as expected
+        assertTrue(commitReceived.get());
+
+        // mock rebalance responses
+        client.respondFrom(joinGroupFollowerResponse(assignor, 2, "memberId", "leaderId", Errors.NONE), coordinator);
+
+        // we need to poll 1) for getting the join response, and then send the sync request;
+        //                 2) for getting the sync response
+        records = consumer.poll(Duration.ZERO);
+
+        // should not finish the response yet
+        assertEquals(Utils.mkSet(topic, topic3), consumer.subscription());
+        assertEquals(Collections.singleton(tp0), consumer.assignment());
+        assertEquals(1, records.count());
+        assertEquals(3L, consumer.position(tp0));
+
+        fetches1.clear();
+        fetches1.put(tp0, new FetchInfo(3, 1));
+        client.respondFrom(fetchResponse(fetches1), node);
+
+        // now complete the rebalance
+        client.respondFrom(syncGroupResponse(Arrays.asList(tp0, t3p0), Errors.NONE), coordinator);
+
+        AtomicInteger count = new AtomicInteger(0);
+        TestUtils.waitForCondition(() -> {
+            ConsumerRecords recs = consumer.poll(Duration.ofMillis(100L));
+            return consumer.assignment().equals(Utils.mkSet(tp0, t3p0)) && count.addAndGet(recs.count()) == 1;
+
+        }, "Does not complete rebalance in time");
+
+        // should have t3 but not sent yet the t3 records
+        assertEquals(Utils.mkSet(topic, topic3), consumer.subscription());
+        assertEquals(Utils.mkSet(tp0, t3p0), consumer.assignment());
+        assertEquals(4L, consumer.position(tp0));
+        assertEquals(0L, consumer.position(t3p0));
+
+        fetches1.clear();
+        fetches1.put(tp0, new FetchInfo(4, 1));
+        fetches1.put(t3p0, new FetchInfo(0, 100));
+        client.respondFrom(fetchResponse(fetches1), node);
+
+        count.set(0);
+        TestUtils.waitForCondition(() -> {
+            ConsumerRecords recs = consumer.poll(Duration.ofMillis(100L));
+            return count.addAndGet(recs.count()) == 101;
+
+        }, "Does not complete rebalance in time");
+
+        assertEquals(5L, consumer.position(tp0));
+        assertEquals(100L, consumer.position(t3p0));
+
+        client.requests().clear();
+        consumer.unsubscribe();
+        consumer.close();
+    }
+
+    @Test
+    public void testGetGroupMetadata() {
+        final Time time = new MockTime();
+        final SubscriptionState subscription = new SubscriptionState(new LogContext(), OffsetResetStrategy.EARLIEST);
+        final ConsumerMetadata metadata = createMetadata(subscription);
+        final MockClient client = new MockClient(time, metadata);
+
+        initMetadata(client, Collections.singletonMap(topic, 1));
+        final Node node = metadata.fetch().nodes().get(0);
+
+        final ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
+
+        final KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true, groupInstanceId);
+
+        final ConsumerGroupMetadata groupMetadataOnStart = consumer.groupMetadata();
+        assertEquals(groupId, groupMetadataOnStart.groupId());
+        assertEquals(JoinGroupRequest.UNKNOWN_MEMBER_ID, groupMetadataOnStart.memberId());
+        assertEquals(JoinGroupRequest.UNKNOWN_GENERATION_ID, groupMetadataOnStart.generationId());
+        assertEquals(groupInstanceId, groupMetadataOnStart.groupInstanceId());
+
+        consumer.subscribe(singleton(topic), getConsumerRebalanceListener(consumer));
+        prepareRebalance(client, node, assignor, singletonList(tp0), null);
+
+        // initial fetch
+        client.prepareResponseFrom(fetchResponse(tp0, 0, 0), node);
+        consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));
+
+        final ConsumerGroupMetadata groupMetadataAfterPoll = consumer.groupMetadata();
+        assertEquals(groupId, groupMetadataAfterPoll.groupId());
+        assertEquals(memberId, groupMetadataAfterPoll.memberId());
+        assertEquals(1, groupMetadataAfterPoll.generationId());
+        assertEquals(groupInstanceId, groupMetadataAfterPoll.groupInstanceId());
+    }
+
+    @Test
+    public void testInvalidGroupMetadata() throws InterruptedException {
+        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));
+        // concurrent access is illegal
+        client.enableBlockingUntilWakeup(1);
+        ExecutorService service = Executors.newSingleThreadExecutor();
+        service.execute(() -> consumer.poll(Duration.ofSeconds(5)));
+        try {
+            TimeUnit.SECONDS.sleep(1);
+            assertThrows(ConcurrentModificationException.class, consumer::groupMetadata);
+            client.wakeup();
+            consumer.wakeup();
+        } finally {
+            service.shutdown();
+            assertTrue(service.awaitTermination(10, TimeUnit.SECONDS));
+        }
+
+        // accessing closed consumer is illegal
+        consumer.close(Duration.ofSeconds(5));
+        assertThrows(IllegalStateException.class, consumer::groupMetadata);
+    }
+
+    private KafkaConsumer consumerWithPendingAuthenticationError() {
+        Time time = new MockTime();
+        SubscriptionState subscription = new SubscriptionState(new LogContext(), OffsetResetStrategy.EARLIEST);
+        ConsumerMetadata metadata = createMetadata(subscription);
+        MockClient client = new MockClient(time, metadata);
+
+        initMetadata(client, singletonMap(topic, 1));
+        Node node = metadata.fetch().nodes().get(0);
+
+        ConsumerPartitionAssignor assignor = new RangeAssignor();
+
+        client.createPendingAuthenticationError(node, 0);
+        return newConsumer(time, client, subscription, metadata, assignor, false, groupInstanceId);
+    }
+
     private ConsumerRebalanceListener getConsumerRebalanceListener(final KafkaConsumer consumer) {
         return new ConsumerRebalanceListener() {
             @Override
@@ -1444,26 +2066,48 @@ public void onPartitionsAssigned(Collection partitions) {
         };
     }
 
-    private Metadata createMetadata() {
-        return new Metadata(0, Long.MAX_VALUE, true);
+    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, 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());
         }
 
         // join group
-        client.prepareResponseFrom(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                JoinGroupRequest joinGroupRequest = (JoinGroupRequest) body;
-                PartitionAssignor.Subscription subscription = ConsumerProtocol.deserializeSubscription(joinGroupRequest.groupProtocols().get(0).metadata());
-                return subscribedTopics.equals(new HashSet<>(subscription.topics()));
-            }
-        }, joinGroupFollowerResponse(assignor, 1, "memberId", "leaderId", Errors.NONE), coordinator);
+        client.prepareResponseFrom(body -> {
+            JoinGroupRequest joinGroupRequest = (JoinGroupRequest) body;
+            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);
 
         // sync group
         client.prepareResponseFrom(syncGroupResponse(partitions, Errors.NONE), coordinator);
@@ -1471,15 +2115,15 @@ 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());
         }
 
         // join group
-        client.prepareResponseFrom(joinGroupFollowerResponse(assignor, 1, "memberId", "leaderId", Errors.NONE), coordinator);
+        client.prepareResponseFrom(joinGroupFollowerResponse(assignor, 1, memberId, leaderId, Errors.NONE), coordinator);
 
         // sync group
         client.prepareResponseFrom(syncGroupResponse(partitions, Errors.NONE), coordinator);
@@ -1487,15 +2131,12 @@ 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
-            public boolean matches(AbstractRequest body) {
-                heartbeatReceived.set(true);
-                return true;
-            }
-        }, new HeartbeatResponse(Errors.NONE), coordinator);
+        client.prepareResponseFrom(body -> {
+            heartbeatReceived.set(true);
+            return true;
+        }, new HeartbeatResponse(new HeartbeatResponseData().setErrorCode(error.code())), coordinator);
         return heartbeatReceived;
     }
 
@@ -1505,20 +2146,18 @@ private AtomicBoolean prepareOffsetCommitResponse(MockClient client, Node coordi
         for (TopicPartition partition : partitionOffsets.keySet())
             response.put(partition, Errors.NONE);
 
-        client.prepareResponseFrom(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                OffsetCommitRequest commitRequest = (OffsetCommitRequest) body;
-                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()) {
-                        commitReceived.set(false);
-                        return false;
-                    }
+        client.prepareResponseFrom(body -> {
+            OffsetCommitRequest commitRequest = (OffsetCommitRequest) body;
+            Map commitErrors = commitRequest.offsets();
+
+            for (Map.Entry partitionOffset : partitionOffsets.entrySet()) {
+                // verify that the expected offset has been committed
+                if (!commitErrors.get(partitionOffset.getKey()).equals(partitionOffset.getValue())) {
+                    commitReceived.set(false);
+                    return false;
                 }
-                return true;
             }
+            return true;
         }, offsetCommitResponse(response), coordinator);
         return commitReceived;
     }
@@ -1531,35 +2170,69 @@ 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) {
         Map partitionData = new HashMap<>();
         for (Map.Entry entry : offsets.entrySet()) {
-            partitionData.put(entry.getKey(), new OffsetFetchResponse.PartitionData(entry.getValue(), "", error));
+            partitionData.put(entry.getKey(), new OffsetFetchResponse.PartitionData(entry.getValue(),
+                    Optional.empty(), "", error));
         }
         return new OffsetFetchResponse(Errors.NONE, partitionData);
     }
 
-    private ListOffsetResponse listOffsetsResponse(Map offsets, Errors error) {
-        Map partitionData = new HashMap<>();
-        for (Map.Entry partitionOffset : offsets.entrySet()) {
-            partitionData.put(partitionOffset.getKey(), new ListOffsetResponse.PartitionData(error,
-                    1L, partitionOffset.getValue()));
+    private ListOffsetResponse listOffsetsResponse(Map offsets) {
+        return listOffsetsResponse(offsets, Collections.emptyMap());
+    }
+
+    private ListOffsetResponse listOffsetsResponse(Map partitionOffsets,
+                                                   Map partitionErrors) {
+        Map responses = new HashMap<>();
+        for (Map.Entry partitionOffset : partitionOffsets.entrySet()) {
+            TopicPartition tp = partitionOffset.getKey();
+            ListOffsetTopicResponse topic = responses.computeIfAbsent(tp.topic(), k -> new ListOffsetTopicResponse().setName(tp.topic()));
+            topic.partitions().add(new ListOffsetPartitionResponse()
+                    .setPartitionIndex(tp.partition())
+                    .setErrorCode(Errors.NONE.code())
+                    .setTimestamp(ListOffsetResponse.UNKNOWN_TIMESTAMP)
+                    .setOffset(partitionOffset.getValue()));
         }
-        return new ListOffsetResponse(partitionData);
+
+        for (Map.Entry partitionError : partitionErrors.entrySet()) {
+            TopicPartition tp = partitionError.getKey();
+            ListOffsetTopicResponse topic = responses.computeIfAbsent(tp.topic(), k -> new ListOffsetTopicResponse().setName(tp.topic()));
+            topic.partitions().add(new ListOffsetPartitionResponse()
+                    .setPartitionIndex(tp.partition())
+                    .setErrorCode(partitionError.getValue().code())
+                    .setTimestamp(ListOffsetResponse.UNKNOWN_TIMESTAMP)
+                    .setOffset(ListOffsetResponse.UNKNOWN_OFFSET));
+        }
+        ListOffsetResponseData data = new ListOffsetResponseData()
+                .setTopics(new ArrayList<>(responses.values()));
+        return new ListOffsetResponse(data);
     }
 
-    private FetchResponse fetchResponse(Map fetches) {
-        LinkedHashMap tpResponses = new LinkedHashMap<>();
+    private FetchResponse fetchResponse(Map fetches) {
+        LinkedHashMap> tpResponses = new LinkedHashMap<>();
         for (Map.Entry fetchEntry : fetches.entrySet()) {
             TopicPartition partition = fetchEntry.getKey();
             long fetchOffset = fetchEntry.getValue().offset;
@@ -1574,75 +2247,90 @@ private FetchResponse fetchResponse(Map fetches) {
                     builder.append(0L, ("key-" + i).getBytes(), ("value-" + i).getBytes());
                 records = builder.build();
             }
-            tpResponses.put(partition, new FetchResponse.PartitionData(Errors.NONE, 0, FetchResponse.INVALID_LAST_STABLE_OFFSET, 0L,
-                    null, records));
+            tpResponses.put(partition, new FetchResponse.PartitionData<>(
+                    Errors.NONE, 0, FetchResponse.INVALID_LAST_STABLE_OFFSET,
+                    0L, null, records));
         }
-        return new FetchResponse(tpResponses, 0);
+        return new FetchResponse<>(Errors.NONE, tpResponses, 0, INVALID_SESSION_ID);
     }
 
-    private FetchResponse fetchResponse(TopicPartition partition, long fetchOffset, int count) {
+    private FetchResponse fetchResponse(TopicPartition partition, long fetchOffset, int count) {
         FetchInfo fetchInfo = new FetchInfo(fetchOffset, count);
         return fetchResponse(Collections.singletonMap(partition, fetchInfo));
     }
 
     private KafkaConsumer newConsumer(Time time,
                                                       KafkaClient client,
-                                                      Metadata metadata,
-                                                      PartitionAssignor assignor,
-                                                      int rebalanceTimeoutMs,
-                                                      int sessionTimeoutMs,
-                                                      int heartbeatIntervalMs,
+                                                      SubscriptionState subscription,
+                                                      ConsumerMetadata metadata,
+                                                      ConsumerPartitionAssignor assignor,
                                                       boolean autoCommitEnabled,
-                                                      int autoCommitIntervalMs) {
-        // create a consumer with mocked time and mocked network
+                                                      Optional groupInstanceId) {
+        return newConsumer(time, client, subscription, metadata, assignor, autoCommitEnabled, groupId, groupInstanceId, false);
+    }
 
+    private KafkaConsumer newConsumerNoAutoCommit(Time time,
+                                                                  KafkaClient client,
+                                                                  SubscriptionState subscription,
+                                                                  ConsumerMetadata metadata) {
+        return newConsumer(time, client, subscription, metadata, new RangeAssignor(), false, groupId, groupInstanceId, false);
+    }
+
+    private KafkaConsumer newConsumer(Time time,
+                                                      KafkaClient client,
+                                                      SubscriptionState subscription,
+                                                      ConsumerMetadata metadata,
+                                                      ConsumerPartitionAssignor assignor,
+                                                      boolean autoCommitEnabled,
+                                                      String groupId,
+                                                      Optional groupInstanceId,
+                                                      boolean throwOnStableOffsetNotSupported) {
         String clientId = "mock-consumer";
-        String groupId = "mock-group";
         String metricGroupPrefix = "consumer";
         long retryBackoffMs = 100;
-        long requestTimeoutMs = 30000;
-        boolean excludeInternalTopics = true;
+        int requestTimeoutMs = 30000;
+        int defaultApiTimeoutMs = 30000;
         int minBytes = 1;
         int maxBytes = Integer.MAX_VALUE;
         int maxWaitMs = 500;
         int fetchSize = 1024 * 1024;
         int maxPollRecords = Integer.MAX_VALUE;
         boolean checkCrcs = true;
+        int rebalanceTimeoutMs = 60000;
 
         Deserializer keyDeserializer = new StringDeserializer();
         Deserializer valueDeserializer = new StringDeserializer();
 
-        OffsetResetStrategy autoResetStrategy = OffsetResetStrategy.EARLIEST;
-        List assignors = Arrays.asList(assignor);
-        ConsumerInterceptors interceptors = null;
+        List assignors = singletonList(assignor);
+        ConsumerInterceptors interceptors = new ConsumerInterceptors<>(Collections.emptyList());
 
-        Metrics metrics = new Metrics();
+        Metrics metrics = new Metrics(time);
         ConsumerMetrics metricsRegistry = new ConsumerMetrics(metricGroupPrefix);
 
-        SubscriptionState subscriptions = new SubscriptionState(autoResetStrategy);
         LogContext loggerFactory = new LogContext();
         ConsumerNetworkClient consumerClient = new ConsumerNetworkClient(loggerFactory, client, metadata, time,
                 retryBackoffMs, requestTimeoutMs, heartbeatIntervalMs);
-        ConsumerCoordinator consumerCoordinator = new ConsumerCoordinator(
-                loggerFactory,
-                consumerClient,
-                groupId,
-                rebalanceTimeoutMs,
-                sessionTimeoutMs,
-                heartbeatIntervalMs,
-                assignors,
-                metadata,
-                subscriptions,
-                metrics,
-                metricGroupPrefix,
-                time,
-                retryBackoffMs,
-                autoCommitEnabled,
-                autoCommitIntervalMs,
-                interceptors,
-                excludeInternalTopics,
-                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,
+                                                                          throwOnStableOffsetNotSupported);
         Fetcher fetcher = new Fetcher<>(
                 loggerFactory,
                 consumerClient,
@@ -1652,15 +2340,18 @@ private KafkaConsumer newConsumer(Time time,
                 fetchSize,
                 maxPollRecords,
                 checkCrcs,
+                "",
                 keyDeserializer,
                 valueDeserializer,
                 metadata,
-                subscriptions,
+                subscription,
                 metrics,
                 metricsRegistry.fetcherMetrics,
                 time,
                 retryBackoffMs,
-                IsolationLevel.READ_UNCOMMITTED);
+                requestTimeoutMs,
+                IsolationLevel.READ_UNCOMMITTED,
+                new ApiVersions());
 
         return new KafkaConsumer<>(
                 loggerFactory,
@@ -1673,11 +2364,13 @@ private KafkaConsumer newConsumer(Time time,
                 time,
                 consumerClient,
                 metrics,
-                subscriptions,
+                subscription,
                 metadata,
                 retryBackoffMs,
                 requestTimeoutMs,
-                assignors);
+                defaultApiTimeoutMs,
+                assignors,
+                groupId);
     }
 
     private static class FetchInfo {
@@ -1689,4 +2382,265 @@ private static class FetchInfo {
             this.count = count;
         }
     }
+
+    @Test
+    @SuppressWarnings("deprecation")
+    public void testCloseWithTimeUnit() {
+        KafkaConsumer consumer = mock(KafkaConsumer.class);
+        doCallRealMethod().when(consumer).close(anyLong(), any());
+        consumer.close(1, TimeUnit.SECONDS);
+        verify(consumer).close(Duration.ofSeconds(1));
+    }
+
+    @Test(expected = InvalidTopicException.class)
+    public void testSubscriptionOnInvalidTopic() {
+        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));
+        Cluster cluster = metadata.fetch();
+
+        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 = RequestTestUtils.metadataResponse(cluster.nodes(),
+                cluster.clusterResource().clusterId(),
+                cluster.controller().id(),
+                topicMetadata);
+        client.prepareMetadataUpdate(updateResponse);
+
+        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"));
+    }
+
+    @Test(expected = IllegalStateException.class)
+    public void testEnforceRebalanceWithManualAssignment() {
+        try (KafkaConsumer consumer = newConsumer((String) null)) {
+            consumer.assign(singleton(new TopicPartition("topic", 0)));
+            consumer.enforceRebalance();
+        }
+    }
+
+    @Test
+    public void testEnforceRebalanceTriggersRebalanceOnNextPoll() {
+        Time time = new MockTime(1L);
+        SubscriptionState subscription = new SubscriptionState(new LogContext(), OffsetResetStrategy.EARLIEST);
+        ConsumerMetadata metadata = createMetadata(subscription);
+        MockClient client = new MockClient(time, metadata);
+        ConsumerPartitionAssignor assignor = new RoundRobinAssignor();
+        KafkaConsumer consumer = newConsumer(time, client, subscription, metadata, assignor, true, groupInstanceId);
+        MockRebalanceListener countingRebalanceListener = new MockRebalanceListener();
+        initMetadata(client, Utils.mkMap(Utils.mkEntry(topic, 1), Utils.mkEntry(topic2, 1), Utils.mkEntry(topic3, 1)));
+
+        consumer.subscribe(Arrays.asList(topic, topic2), countingRebalanceListener);
+        Node node = metadata.fetch().nodes().get(0);
+        prepareRebalance(client, node, assignor, Arrays.asList(tp0, t2p0), null);
+
+        // a first rebalance to get the assignment, we need two poll calls since we need two round trips to finish join / sync-group
+        consumer.poll(Duration.ZERO);
+        consumer.poll(Duration.ZERO);
+
+        // onPartitionsRevoked is not invoked when first joining the group
+        assertEquals(countingRebalanceListener.revokedCount, 0);
+        assertEquals(countingRebalanceListener.assignedCount, 1);
+
+        consumer.enforceRebalance();
+
+        // the next poll should trigger a rebalance
+        consumer.poll(Duration.ZERO);
+
+        assertEquals(countingRebalanceListener.revokedCount, 1);
+    }
+
+    @Test
+    public void configurableObjectsShouldSeeGeneratedClientId() {
+        Properties props = new Properties();
+        props.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
+        props.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, DeserializerForClientId.class.getName());
+        props.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, DeserializerForClientId.class.getName());
+        props.put(ConsumerConfig.INTERCEPTOR_CLASSES_CONFIG, ConsumerInterceptorForClientId.class.getName());
+
+        KafkaConsumer consumer = new KafkaConsumer<>(props);
+        assertNotNull(consumer.getClientId());
+        assertNotEquals(0, consumer.getClientId().length());
+        assertEquals(3, CLIENT_IDS.size());
+        CLIENT_IDS.forEach(id -> assertEquals(id, consumer.getClientId()));
+        consumer.close();
+    }
+
+    @Test
+    public void testUnusedConfigs() {
+        Map props = new HashMap<>();
+        props.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
+        props.put(SslConfigs.SSL_PROTOCOL_CONFIG, "TLS");
+        ConsumerConfig config = new ConsumerConfig(ConsumerConfig.appendDeserializerToConfig(props, new StringDeserializer(), new StringDeserializer()));
+
+        assertTrue(config.unused().contains(SslConfigs.SSL_PROTOCOL_CONFIG));
+
+        try (KafkaConsumer consumer = new KafkaConsumer<>(config, null, null)) {
+            assertTrue(config.unused().contains(SslConfigs.SSL_PROTOCOL_CONFIG));
+        }
+    }
+
+    private static final List CLIENT_IDS = new ArrayList<>();
+    public static class DeserializerForClientId implements Deserializer {
+        @Override
+        public void configure(Map configs, boolean isKey) {
+            CLIENT_IDS.add(configs.get(ConsumerConfig.CLIENT_ID_CONFIG).toString());
+        }
+
+        @Override
+        public byte[] deserialize(String topic, byte[] data) {
+            return data;
+        }
+    }
+
+    public static class ConsumerInterceptorForClientId implements ConsumerInterceptor {
+
+        @Override
+        public ConsumerRecords onConsume(ConsumerRecords records) {
+            return records;
+        }
+
+        @Override
+        public void onCommit(Map offsets) {
+
+        }
+
+        @Override
+        public void close() {
+
+        }
+
+        @Override
+        public void configure(Map configs) {
+            CLIENT_IDS.add(configs.get(ConsumerConfig.CLIENT_ID_CONFIG).toString());
+        }
+    }
 }
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 6ed46f711cbaf..753d6d794c935 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
@@ -16,25 +16,59 @@
  */
 package org.apache.kafka.clients.consumer;
 
-import org.apache.kafka.clients.consumer.internals.NoOpConsumerRebalanceListener;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.record.TimestampType;
+import org.junit.Assert;
 import org.junit.Test;
 
+import java.time.Duration;
 import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.Iterator;
+import java.util.Optional;
 
+import static org.hamcrest.CoreMatchers.equalTo;
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 
 public class MockConsumerTest {
     
-    private MockConsumer consumer = new MockConsumer(OffsetResetStrategy.EARLIEST);
+    private MockConsumer consumer = new MockConsumer<>(OffsetResetStrategy.EARLIEST);
 
     @Test
     public void testSimpleMock() {
-        consumer.subscribe(Arrays.asList("test"), new NoOpConsumerRebalanceListener());
+        consumer.subscribe(Collections.singleton("test"));
+        assertEquals(0, consumer.poll(Duration.ZERO).count());
+        consumer.rebalance(Arrays.asList(new TopicPartition("test", 0), new TopicPartition("test", 1)));
+        // Mock consumers need to seek manually since they cannot automatically reset offsets
+        HashMap beginningOffsets = new HashMap<>();
+        beginningOffsets.put(new TopicPartition("test", 0), 0L);
+        beginningOffsets.put(new TopicPartition("test", 1), 0L);
+        consumer.updateBeginningOffsets(beginningOffsets);
+        consumer.seek(new TopicPartition("test", 0), 0);
+        ConsumerRecord rec1 = new ConsumerRecord<>("test", 0, 0, 0L, TimestampType.CREATE_TIME, 0L, 0, 0, "key1", "value1");
+        ConsumerRecord rec2 = new ConsumerRecord<>("test", 0, 1, 0L, TimestampType.CREATE_TIME, 0L, 0, 0, "key2", "value2");
+        consumer.addRecord(rec1);
+        consumer.addRecord(rec2);
+        ConsumerRecords recs = consumer.poll(Duration.ofMillis(1));
+        Iterator> iter = recs.iterator();
+        assertEquals(rec1, iter.next());
+        assertEquals(rec2, iter.next());
+        assertFalse(iter.hasNext());
+        final TopicPartition tp = new TopicPartition("test", 0);
+        assertEquals(2L, consumer.position(tp));
+        consumer.commitSync();
+        assertEquals(2L, consumer.committed(Collections.singleton(tp)).get(tp).offset());
+    }
+
+    @SuppressWarnings("deprecation")
+    @Test
+    public void testSimpleMockDeprecated() {
+        consumer.subscribe(Collections.singleton("test"));
         assertEquals(0, consumer.poll(1000).count());
         consumer.rebalance(Arrays.asList(new TopicPartition("test", 0), new TopicPartition("test", 1)));
         // Mock consumers need to seek manually since they cannot automatically reset offsets
@@ -43,8 +77,8 @@ public void testSimpleMock() {
         beginningOffsets.put(new TopicPartition("test", 1), 0L);
         consumer.updateBeginningOffsets(beginningOffsets);
         consumer.seek(new TopicPartition("test", 0), 0);
-        ConsumerRecord rec1 = new ConsumerRecord("test", 0, 0, 0L, TimestampType.CREATE_TIME, 0L, 0, 0, "key1", "value1");
-        ConsumerRecord rec2 = new ConsumerRecord("test", 0, 1, 0L, TimestampType.CREATE_TIME, 0L, 0, 0, "key2", "value2");
+        ConsumerRecord rec1 = new ConsumerRecord<>("test", 0, 0, 0L, TimestampType.CREATE_TIME, 0L, 0, 0, "key1", "value1");
+        ConsumerRecord rec2 = new ConsumerRecord<>("test", 0, 1, 0L, TimestampType.CREATE_TIME, 0L, 0, 0, "key2", "value2");
         consumer.addRecord(rec1);
         consumer.addRecord(rec2);
         ConsumerRecords recs = consumer.poll(1);
@@ -52,9 +86,54 @@ 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());
+        assertThat(consumer.groupMetadata(), equalTo(new ConsumerGroupMetadata("dummy.group.id", 1, "1", Optional.empty())));
+    }
+
+    @Test
+    public void testConsumerRecordsIsEmptyWhenReturningNoRecords() {
+        TopicPartition partition = new TopicPartition("test", 0);
+        consumer.assign(Collections.singleton(partition));
+        consumer.addRecord(new ConsumerRecord<>("test", 0, 0, null, null));
+        consumer.updateEndOffsets(Collections.singletonMap(partition, 1L));
+        consumer.seekToEnd(Collections.singleton(partition));
+        ConsumerRecords records = consumer.poll(Duration.ofMillis(1));
+        assertThat(records.count(), is(0));
+        assertThat(records.isEmpty(), is(true));
+    }
+
+    @Test
+    public void shouldNotClearRecordsForPausedPartitions() {
+        TopicPartition partition0 = new TopicPartition("test", 0);
+        Collection testPartitionList = Collections.singletonList(partition0);
+        consumer.assign(testPartitionList);
+        consumer.addRecord(new ConsumerRecord<>("test", 0, 0, null, null));
+        consumer.updateBeginningOffsets(Collections.singletonMap(partition0, 0L));
+        consumer.seekToBeginning(testPartitionList);
+
+        consumer.pause(testPartitionList);
+        consumer.poll(Duration.ofMillis(1));
+        consumer.resume(testPartitionList);
+        ConsumerRecords recordsSecondPoll = consumer.poll(Duration.ofMillis(1));
+        assertThat(recordsSecondPoll.count(), is(1));
+    }
+
+    @Test
+    public void endOffsetsShouldBeIdempotent() {
+        TopicPartition partition = new TopicPartition("test", 0);
+        consumer.updateEndOffsets(Collections.singletonMap(partition, 10L));
+        // consumer.endOffsets should NOT change the value of end offsets
+        Assert.assertEquals(10L, (long) consumer.endOffsets(Collections.singleton(partition)).get(partition));
+        Assert.assertEquals(10L, (long) consumer.endOffsets(Collections.singleton(partition)).get(partition));
+        Assert.assertEquals(10L, (long) consumer.endOffsets(Collections.singleton(partition)).get(partition));
+        consumer.updateEndOffsets(Collections.singletonMap(partition, 11L));
+        // consumer.endOffsets should NOT change the value of end offsets
+        Assert.assertEquals(11L, (long) consumer.endOffsets(Collections.singleton(partition)).get(partition));
+        Assert.assertEquals(11L, (long) consumer.endOffsets(Collections.singleton(partition)).get(partition));
+        Assert.assertEquals(11L, (long) consumer.endOffsets(Collections.singleton(partition)).get(partition));
     }
 
 }
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/OffsetAndMetadataTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/OffsetAndMetadataTest.java
new file mode 100644
index 0000000000000..5bdbf7c7a3279
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/OffsetAndMetadataTest.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.clients.consumer;
+
+import org.apache.kafka.common.utils.Serializer;
+import org.junit.Test;
+
+import java.io.IOException;
+import java.util.Optional;
+
+import static org.junit.Assert.assertEquals;
+
+/**
+ * This test case ensures OffsetAndMetadata class is serializable and is serialization compatible.
+ * Note: this ensures that the current code can deserialize data serialized with older versions of the code, but not the reverse.
+ * That is, older code won't necessarily be able to deserialize data serialized with newer code.
+ */
+public class OffsetAndMetadataTest {
+
+    @Test(expected = IllegalArgumentException.class)
+    public void testInvalidNegativeOffset() {
+        new OffsetAndMetadata(-239L, Optional.of(15), "");
+    }
+
+    @Test
+    public void testSerializationRoundtrip() throws IOException, ClassNotFoundException {
+        checkSerde(new OffsetAndMetadata(239L, Optional.of(15), "blah"));
+        checkSerde(new OffsetAndMetadata(239L, "blah"));
+        checkSerde(new OffsetAndMetadata(239L));
+    }
+
+    private void checkSerde(OffsetAndMetadata offsetAndMetadata) throws IOException, ClassNotFoundException {
+        byte[] bytes =  Serializer.serialize(offsetAndMetadata);
+        OffsetAndMetadata deserialized = (OffsetAndMetadata) Serializer.deserialize(bytes);
+        assertEquals(offsetAndMetadata, deserialized);
+    }
+
+    @Test
+    public void testDeserializationCompatibilityBeforeLeaderEpoch() throws IOException, ClassNotFoundException {
+        String fileName = "serializedData/offsetAndMetadataBeforeLeaderEpoch";
+        Object deserializedObject = Serializer.deserialize(fileName);
+        assertEquals(new OffsetAndMetadata(10, "test commit metadata"), deserializedObject);
+    }
+
+    @Test
+    public void testDeserializationCompatibilityWithLeaderEpoch() throws IOException, ClassNotFoundException {
+        String fileName = "serializedData/offsetAndMetadataWithLeaderEpoch";
+        Object deserializedObject = Serializer.deserialize(fileName);
+        assertEquals(new OffsetAndMetadata(10, Optional.of(235), "test commit metadata"), deserializedObject);
+    }
+
+}
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..74fa223a7bba0 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.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.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/SerializeCompatibilityOffsetAndMetadataTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/SerializeCompatibilityOffsetAndMetadataTest.java
deleted file mode 100644
index 324aeafd88646..0000000000000
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/SerializeCompatibilityOffsetAndMetadataTest.java
+++ /dev/null
@@ -1,64 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements. See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License. You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.kafka.clients.consumer;
-
-import org.apache.kafka.common.utils.Serializer;
-import org.junit.Test;
-
-import java.io.IOException;
-
-
-import static org.junit.Assert.assertTrue;
-import static org.junit.Assert.assertEquals;
-
-/**
- * This test case ensures OffsetAndMetadata class is serializable and is serialization compatible.
- * Note: this ensures that the current code can deserialize data serialized with older versions of the code, but not the reverse.
- * That is, older code won't necessarily be able to deserialize data serialized with newer code.
- */
-public class SerializeCompatibilityOffsetAndMetadataTest {
-    private String metadata = "test commit metadata";
-    private String fileName = "serializedData/offsetAndMetadataSerializedfile";
-    private long offset = 10;
-
-    private void checkValues(OffsetAndMetadata deSerOAM) {
-        //assert deserialized values are same as original
-        assertEquals("Offset should be " + offset + " but got " + deSerOAM.offset(), offset, deSerOAM.offset());
-        assertEquals("metadata should be " + metadata + " but got " + deSerOAM.metadata(), metadata, deSerOAM.metadata());
-    }
-
-    @Test
-    public void testSerializationRoundtrip() throws IOException, ClassNotFoundException {
-        //assert OffsetAndMetadata is serializable
-        OffsetAndMetadata origOAM = new OffsetAndMetadata(offset, metadata);
-        byte[] byteArray =  Serializer.serialize(origOAM);
-
-        //deserialize the byteArray and check if the values are same as original
-        Object deserializedObject = Serializer.deserialize(byteArray);
-        assertTrue(deserializedObject instanceof OffsetAndMetadata);
-        checkValues((OffsetAndMetadata) deserializedObject);
-    }
-
-    @Test
-    public void testOffsetMetadataSerializationCompatibility() throws IOException, ClassNotFoundException {
-        // assert serialized OffsetAndMetadata object in file (oamserializedfile under resources folder) is
-        // deserializable into OffsetAndMetadata and is compatible
-        Object deserializedObject = Serializer.deserialize(fileName);
-        assertTrue(deserializedObject instanceof OffsetAndMetadata);
-        checkValues((OffsetAndMetadata) deserializedObject);
-    }
-}
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 4a78919b7f677..2189700fe6230 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,720 +16,224 @@
  */
 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";
+public class StickyAssignorTest extends AbstractStickyAssignorTest {
 
-        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());
-
-        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));
+    @Override
+    public Subscription buildSubscription(List topics, List partitions) {
+        return new Subscription(topics,
+            serializeTopicPartitionAssignment(new MemberData(partitions, Optional.of(DEFAULT_GENERATION))));
     }
 
     @Test
-    public void testOneConsumerOneTopic() {
-        String topic = "topic";
-        String consumerId = "consumer";
+    public void testAssignmentWithMultipleGenerations1() {
+        String consumer1 = "consumer1";
+        String consumer2 = "consumer2";
+        String consumer3 = "consumer3";
 
         Map partitionsPerTopic = new HashMap<>();
-        partitionsPerTopic.put(topic, 3);
-        Map subscriptions = Collections.singletonMap(consumerId, new Subscription(topics(topic)));
+        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), tp(topic, 1), tp(topic, 2)), assignment.get(consumerId));
-
-        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 testOnlyAssignsPartitionsFromSubscribedTopics() {
-        String topic = "topic";
-        String otherTopic = "other";
-        String consumerId = "consumer";
+        subscriptions.put(consumer1, buildSubscription(topics(topic), r1partitions1));
+        subscriptions.put(consumer2, buildSubscription(topics(topic), r1partitions2));
+        subscriptions.remove(consumer3);
 
-        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);
+        assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        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));
-    }
+        assertFalse(Collections.disjoint(r2partitions2, r1partitions3));
 
-    @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));
+        subscriptions.remove(consumer1);
+        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));
     }
 
     @Test
-    public void testTwoConsumersOneTopicOnePartition() {
-        String topic = "topic";
+    public void testAssignmentWithMultipleGenerations2() {
         String consumer1 = "consumer1";
         String consumer2 = "consumer2";
+        String consumer3 = "consumer3";
 
         Map partitionsPerTopic = new HashMap<>();
-        partitionsPerTopic.put(topic, 1);
-
-        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(Collections.emptyList(), 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 testTwoConsumersOneTopicTwoPartitions() {
-        String topic = "topic";
-        String consumer1 = "consumer1";
-        String consumer2 = "consumer2";
 
-        Map partitionsPerTopic = new HashMap<>();
-        partitionsPerTopic.put(topic, 2);
-
-        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(partitions(tp(topic, 1)), assignment.get(consumer2));
+        subscriptions.remove(consumer1);
+        subscriptions.put(consumer2, buildSubscriptionWithGeneration(topics(topic), r1partitions2, 1));
+        subscriptions.remove(consumer3);
 
-        verifyValidityAndBalance(subscriptions, assignment);
+        assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        List r2partitions2 = assignment.get(consumer2);
+        assertEquals(6, r2partitions2.size());
+        assertTrue(r2partitions2.containsAll(r1partitions2));
+        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);
 
-        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)));
+        subscriptions.put(consumer1, buildSubscriptionWithGeneration(topics(topic), r1partitions1, 1));
+        subscriptions.put(consumer2, buildSubscriptionWithGeneration(topics(topic), r2partitions2, 2));
+        subscriptions.put(consumer3, buildSubscriptionWithGeneration(topics(topic), r1partitions3, 1));
 
-        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);
+        assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        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));
     }
 
     @Test
-    public void testTwoConsumersTwoTopicsSixPartitions() {
-        String topic1 = "topic1";
-        String topic2 = "topic2";
+    public void testAssignmentWithConflictingPreviousGenerations() {
         String consumer1 = "consumer1";
         String consumer2 = "consumer2";
+        String consumer3 = "consumer3";
 
         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);
-        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<>();
+        partitionsPerTopic.put(topic, 6);
         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);
-        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);
+        subscriptions.put(consumer3, new Subscription(topics(topic)));
+
+        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);
+
+        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);
+        List c1partitions = assignment.get(consumer1);
+        List c2partitions = assignment.get(consumer2);
+        List c3partitions = assignment.get(consumer3);
+
+        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());
-
-        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)));
-
-        verifyValidityAndBalance(subscriptions, assignment);
-        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 testSchemaBackwardCompatibility() {
+        String consumer1 = "consumer1";
         String consumer2 = "consumer2";
+        String consumer3 = "consumer3";
 
         Map partitionsPerTopic = new HashMap<>();
         partitionsPerTopic.put(topic, 3);
-        Map subscriptions = new HashMap<>();
         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
-        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")));
-
-        assignment = assignor.assign(partitionsPerTopic, subscriptions);
-        verifyValidityAndBalance(subscriptions, 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));
-
-        assignment = assignor.assign(partitionsPerTopic, subscriptions);
-        verifyValidityAndBalance(subscriptions, assignment);
-        assertTrue(assignor.isSticky());
-    }
-
-    @Test
-    public void testLargeAssignmentWithMultipleConsumersLeaving() {
-        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);
-
-        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);
-        }
-
-        assignment = assignor.assign(partitionsPerTopic, subscriptions);
-        verifyValidityAndBalance(subscriptions, assignment);
-        assertTrue(assignor.isSticky());
-    }
-
-    @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));
-        }
-
-        Map> assignment = assignor.assign(partitionsPerTopic, subscriptions);
-        verifyValidityAndBalance(subscriptions, assignment);
-
-        subscriptions.get(getConsumerName(0, 3)).topics().add(getTopicName(1, 5));
-
-        assignment = assignor.assign(partitionsPerTopic, subscriptions);
-        verifyValidityAndBalance(subscriptions, 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());
-        }
-    }
-
-    @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);
-    }
-
-    @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))));
-        }
-    }
-
-    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();
+        List c1partitions = assignment.get(consumer1);
+        List c2partitions = assignment.get(consumer2);
+        List c3partitions = assignment.get(consumer3);
 
-        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);
+        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));
     }
 
-    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 assignment: 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.groupDataByTopic(partitions);
-                Map> otherMap = CollectionUtils.groupDataByTopic(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 7eaca9826d915..c24fa8577cc1a 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,11 +16,24 @@
  */
 package org.apache.kafka.clients.consumer.internals;
 
-import org.apache.kafka.clients.Metadata;
+import org.apache.kafka.clients.GroupRebalanceConfig;
 import org.apache.kafka.clients.MockClient;
-import org.apache.kafka.common.Cluster;
+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.DisconnectException;
+import org.apache.kafka.common.errors.FencedInstanceIdException;
+import org.apache.kafka.common.errors.InconsistentGroupProtocolException;
+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;
@@ -29,102 +42,896 @@
 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.RequestTestUtils;
 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.Time;
-import org.apache.kafka.test.TestCondition;
+import org.apache.kafka.common.utils.Timer;
 import org.apache.kafka.test.TestUtils;
 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;
+import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 
+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;
 
 public class AbstractCoordinatorTest {
-
     private static final ByteBuffer EMPTY_DATA = ByteBuffer.wrap(new byte[0]);
     private static final int REBALANCE_TIMEOUT_MS = 60000;
     private static final int SESSION_TIMEOUT_MS = 10000;
     private static final int HEARTBEAT_INTERVAL_MS = 3000;
-    private static final long RETRY_BACKOFF_MS = 20;
-    private static final long LONG_RETRY_BACKOFF_MS = 10000;
-    private static final long REQUEST_TIMEOUT_MS = 40000;
+    private static final int RETRY_BACKOFF_MS = 100;
+    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 static final String PROTOCOL_TYPE = "dummy";
+    private static final String PROTOCOL_NAME = "dummy-subprotocol";
 
-    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 void setupCoordinator(long retryBackoffMs) {
-        this.mockTime = new MockTime();
-        this.mockClient = new MockClient(mockTime);
+    private final String memberId = "memberId";
+    private final String leaderId = "leaderId";
+    private final int defaultGeneration = -1;
 
-        Metadata metadata = new Metadata(100L, 60 * 60 * 1000L, true);
-        this.consumerClient = new ConsumerNetworkClient(new LogContext(), mockClient, metadata, mockTime,
-                retryBackoffMs, REQUEST_TIMEOUT_MS, HEARTBEAT_INTERVAL_MS);
-        Metrics metrics = new Metrics();
+    private void setupCoordinator() {
+        setupCoordinator(RETRY_BACKOFF_MS, REBALANCE_TIMEOUT_MS,
+            Optional.empty());
+    }
 
-        Cluster cluster = TestUtils.singletonCluster("topic", 1);
-        metadata.update(cluster, Collections.emptySet(), mockTime.milliseconds());
-        this.node = cluster.nodes().get(0);
-        mockClient.setNode(node);
+    private void setupCoordinator(int retryBackoffMs) {
+        setupCoordinator(retryBackoffMs, REBALANCE_TIMEOUT_MS,
+            Optional.empty());
+    }
 
+    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, 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 = new Metrics(mockTime);
+
+        mockClient.updateMetadata(RequestTestUtils.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);
+
+        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);
+    }
+
+    private void joinGroup() {
+        mockClient.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+
+        final int generation = 1;
+
+        mockClient.prepareResponse(joinGroupFollowerResponse(generation, memberId, JoinGroupRequest.UNKNOWN_MEMBER_ID, Errors.NONE));
+        mockClient.prepareResponse(syncGroupResponse(Errors.NONE));
+
+        coordinator.ensureActiveGroup();
+    }
+
+    @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
     public void testCoordinatorDiscoveryBackoff() {
-        setupCoordinator(RETRY_BACKOFF_MS);
+        setupCoordinator();
 
         mockClient.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
         mockClient.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
 
-        // blackout the coordinator for 10 milliseconds to simulate a disconnect.
+        // cut out the coordinator for 10 milliseconds to simulate a disconnect.
         // after backing off, we should be able to connect.
-        mockClient.blackout(coordinatorNode, 10L);
+        mockClient.backoff(coordinatorNode, 10L);
 
         long initialTime = mockTime.milliseconds();
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(mockTime.timer(Long.MAX_VALUE));
         long endTime = mockTime.milliseconds();
 
         assertTrue(endTime - initialTime >= RETRY_BACKOFF_MS);
     }
 
+    @Test
+    public void testTimeoutAndRetryJoinGroupIfNeeded() throws Exception {
+        setupCoordinator();
+        mockClient.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        coordinator.ensureCoordinatorReady(mockTime.timer(0));
+
+        ExecutorService executor = Executors.newFixedThreadPool(1);
+        try {
+            Timer firstAttemptTimer = mockTime.timer(REQUEST_TIMEOUT_MS);
+            Future firstAttempt = executor.submit(() -> coordinator.joinGroupIfNeeded(firstAttemptTimer));
+
+            mockTime.sleep(REQUEST_TIMEOUT_MS);
+            assertFalse(firstAttempt.get());
+            assertTrue(consumerClient.hasPendingRequests(coordinatorNode));
+
+            mockClient.respond(joinGroupFollowerResponse(1, memberId, leaderId, Errors.NONE));
+            mockClient.prepareResponse(syncGroupResponse(Errors.NONE));
+
+            Timer secondAttemptTimer = mockTime.timer(REQUEST_TIMEOUT_MS);
+            Future secondAttempt = executor.submit(() -> coordinator.joinGroupIfNeeded(secondAttemptTimer));
+
+            assertTrue(secondAttempt.get());
+        } finally {
+            executor.shutdownNow();
+            executor.awaitTermination(1000, TimeUnit.MILLISECONDS);
+        }
+    }
+
+    @Test
+    public void testGroupMaxSizeExceptionIsFatal() {
+        setupCoordinator();
+        mockClient.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        coordinator.ensureCoordinatorReady(mockTime.timer(0));
+
+        mockClient.prepareResponse(joinGroupFollowerResponse(defaultGeneration, memberId, JoinGroupRequest.UNKNOWN_MEMBER_ID, Errors.GROUP_MAX_SIZE_REACHED));
+
+        RequestFuture future = coordinator.sendJoinGroupRequest();
+        assertTrue(consumerClient.poll(future, mockTime.timer(REQUEST_TIMEOUT_MS)));
+        assertTrue(future.exception().getClass().isInstance(Errors.GROUP_MAX_SIZE_REACHED.exception()));
+        assertFalse(future.isRetriable());
+    }
+
+    @Test
+    public void testJoinGroupRequestTimeout() {
+        setupCoordinator(RETRY_BACKOFF_MS, REBALANCE_TIMEOUT_MS,
+            Optional.empty());
+        mockClient.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        coordinator.ensureCoordinatorReady(mockTime.timer(0));
+
+        RequestFuture future = coordinator.sendJoinGroupRequest();
+
+        mockTime.sleep(REQUEST_TIMEOUT_MS + 1);
+        assertFalse(consumerClient.poll(future, mockTime.timer(0)));
+
+        mockTime.sleep(REBALANCE_TIMEOUT_MS - REQUEST_TIMEOUT_MS + AbstractCoordinator.JOIN_GROUP_TIMEOUT_LAPSE);
+        assertTrue(consumerClient.poll(future, mockTime.timer(0)));
+        assertTrue(future.exception() instanceof DisconnectException);
+    }
+
+    @Test
+    public void testJoinGroupRequestTimeoutLowerBoundedByDefaultRequestTimeout() {
+        int rebalanceTimeoutMs = REQUEST_TIMEOUT_MS - 10000;
+        setupCoordinator(RETRY_BACKOFF_MS, rebalanceTimeoutMs, Optional.empty());
+        mockClient.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        coordinator.ensureCoordinatorReady(mockTime.timer(0));
+
+        RequestFuture future = coordinator.sendJoinGroupRequest();
+
+        long expectedRequestDeadline = mockTime.milliseconds() + REQUEST_TIMEOUT_MS;
+        mockTime.sleep(rebalanceTimeoutMs + AbstractCoordinator.JOIN_GROUP_TIMEOUT_LAPSE + 1);
+        assertFalse(consumerClient.poll(future, mockTime.timer(0)));
+
+        mockTime.sleep(expectedRequestDeadline - mockTime.milliseconds() + 1);
+        assertTrue(consumerClient.poll(future, mockTime.timer(0)));
+        assertTrue(future.exception() instanceof DisconnectException);
+    }
+
+    @Test
+    public void testJoinGroupRequestMaxTimeout() {
+        // Ensure we can handle the maximum allowed rebalance timeout
+
+        setupCoordinator(RETRY_BACKOFF_MS, Integer.MAX_VALUE,
+            Optional.empty());
+        mockClient.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        coordinator.ensureCoordinatorReady(mockTime.timer(0));
+
+        RequestFuture future = coordinator.sendJoinGroupRequest();
+        assertFalse(consumerClient.poll(future, mockTime.timer(0)));
+
+        mockTime.sleep(Integer.MAX_VALUE + 1L);
+        assertTrue(consumerClient.poll(future, mockTime.timer(0)));
+    }
+
+    @Test
+    public void testJoinGroupRequestWithMemberIdRequired() {
+        setupCoordinator();
+        mockClient.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        coordinator.ensureCoordinatorReady(mockTime.timer(0));
+
+        mockClient.prepareResponse(joinGroupFollowerResponse(defaultGeneration, memberId, JoinGroupRequest.UNKNOWN_MEMBER_ID, Errors.MEMBER_ID_REQUIRED));
+
+        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();
+        assertTrue(consumerClient.poll(future, mockTime.timer(REQUEST_TIMEOUT_MS)));
+        assertEquals(Errors.MEMBER_ID_REQUIRED.message(), future.exception().getMessage());
+        assertTrue(coordinator.rejoinNeededOrPending());
+        assertTrue(coordinator.hasValidMemberId());
+        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, JoinGroupRequest.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 testJoinGroupProtocolTypeAndName() {
+        final String wrongProtocolType = "wrong-type";
+        final String wrongProtocolName = "wrong-name";
+
+        // No Protocol Type in both JoinGroup and SyncGroup responses
+        assertTrue(joinGroupWithProtocolTypeAndName(null, null, null));
+
+        // Protocol Type in both JoinGroup and SyncGroup responses
+        assertTrue(joinGroupWithProtocolTypeAndName(PROTOCOL_TYPE, PROTOCOL_TYPE, PROTOCOL_NAME));
+
+        // Wrong protocol type in the JoinGroupResponse
+        assertThrows(InconsistentGroupProtocolException.class,
+            () -> joinGroupWithProtocolTypeAndName("wrong", null, null));
+
+        // Correct protocol type in the JoinGroupResponse
+        // Wrong protocol type in the SyncGroupResponse
+        // Correct protocol name in the SyncGroupResponse
+        assertThrows(InconsistentGroupProtocolException.class,
+            () -> joinGroupWithProtocolTypeAndName(PROTOCOL_TYPE, wrongProtocolType, PROTOCOL_NAME));
+
+        // Correct protocol type in the JoinGroupResponse
+        // Correct protocol type in the SyncGroupResponse
+        // Wrong protocol name in the SyncGroupResponse
+        assertThrows(InconsistentGroupProtocolException.class,
+            () -> joinGroupWithProtocolTypeAndName(PROTOCOL_TYPE, PROTOCOL_TYPE, wrongProtocolName));
+    }
+
+    @Test
+    public void testNoGenerationWillNotTriggerProtocolNameCheck() {
+        final String wrongProtocolName = "wrong-name";
+
+        setupCoordinator();
+        mockClient.reset();
+        mockClient.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        coordinator.ensureCoordinatorReady(mockTime.timer(0));
+
+        mockClient.prepareResponse(body -> {
+            if (!(body instanceof JoinGroupRequest)) {
+                return false;
+            }
+            JoinGroupRequest joinGroupRequest = (JoinGroupRequest) body;
+            return joinGroupRequest.data().protocolType().equals(PROTOCOL_TYPE);
+        }, joinGroupFollowerResponse(defaultGeneration, memberId,
+            "memberid", Errors.NONE, PROTOCOL_TYPE));
+
+        mockClient.prepareResponse(body -> {
+            if (!(body instanceof SyncGroupRequest)) {
+                return false;
+            }
+            coordinator.resetGenerationOnLeaveGroup();
+
+            SyncGroupRequest syncGroupRequest = (SyncGroupRequest) body;
+            return syncGroupRequest.data().protocolType().equals(PROTOCOL_TYPE)
+                       && syncGroupRequest.data().protocolName().equals(PROTOCOL_NAME);
+        }, syncGroupResponse(Errors.NONE, PROTOCOL_TYPE, wrongProtocolName));
+
+        // let the retry to complete successfully to break out of the while loop
+        mockClient.prepareResponse(body -> {
+            if (!(body instanceof JoinGroupRequest)) {
+                return false;
+            }
+            JoinGroupRequest joinGroupRequest = (JoinGroupRequest) body;
+            return joinGroupRequest.data().protocolType().equals(PROTOCOL_TYPE);
+        }, joinGroupFollowerResponse(1, memberId,
+                "memberid", Errors.NONE, PROTOCOL_TYPE));
+
+        mockClient.prepareResponse(body -> {
+            if (!(body instanceof SyncGroupRequest)) {
+                return false;
+            }
+
+            SyncGroupRequest syncGroupRequest = (SyncGroupRequest) body;
+            return syncGroupRequest.data().protocolType().equals(PROTOCOL_TYPE)
+                    && syncGroupRequest.data().protocolName().equals(PROTOCOL_NAME);
+        }, syncGroupResponse(Errors.NONE, PROTOCOL_TYPE, PROTOCOL_NAME));
+
+        // No exception shall be thrown as the generation is reset.
+        coordinator.joinGroupIfNeeded(mockTime.timer(100L));
+    }
+
+    private boolean joinGroupWithProtocolTypeAndName(String joinGroupResponseProtocolType,
+                                                     String syncGroupResponseProtocolType,
+                                                     String syncGroupResponseProtocolName) {
+        setupCoordinator();
+        mockClient.reset();
+        mockClient.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        coordinator.ensureCoordinatorReady(mockTime.timer(0));
+
+        mockClient.prepareResponse(body -> {
+            if (!(body instanceof JoinGroupRequest)) {
+                return false;
+            }
+            JoinGroupRequest joinGroupRequest = (JoinGroupRequest) body;
+            return joinGroupRequest.data().protocolType().equals(PROTOCOL_TYPE);
+        }, joinGroupFollowerResponse(defaultGeneration, memberId,
+            "memberid", Errors.NONE, joinGroupResponseProtocolType));
+
+        mockClient.prepareResponse(body -> {
+            if (!(body instanceof SyncGroupRequest)) {
+                return false;
+            }
+            SyncGroupRequest syncGroupRequest = (SyncGroupRequest) body;
+            return syncGroupRequest.data().protocolType().equals(PROTOCOL_TYPE)
+                && syncGroupRequest.data().protocolName().equals(PROTOCOL_NAME);
+        }, syncGroupResponse(Errors.NONE, syncGroupResponseProtocolType, syncGroupResponseProtocolName));
+
+        return coordinator.joinGroupIfNeeded(mockTime.timer(5000L));
+    }
+
+    @Test
+    public void testSyncGroupRequestWithFencedInstanceIdException() {
+        setupCoordinator();
+        mockClient.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+
+        final int generation = -1;
+
+        mockClient.prepareResponse(joinGroupFollowerResponse(generation, memberId, JoinGroupRequest.UNKNOWN_MEMBER_ID, Errors.NONE));
+        mockClient.prepareResponse(syncGroupResponse(Errors.FENCED_INSTANCE_ID));
+
+        assertThrows(FencedInstanceIdException.class, () -> coordinator.ensureActiveGroup());
+    }
+
+    @Test
+    public void testJoinGroupUnknownMemberResponseWithOldGeneration() throws InterruptedException {
+        setupCoordinator();
+        joinGroup();
+
+        final AbstractCoordinator.Generation currGen = coordinator.generation();
+
+        RequestFuture future = coordinator.sendJoinGroupRequest();
+
+        TestUtils.waitForCondition(() -> !mockClient.requests().isEmpty(), 2000,
+            "The join-group request was not sent");
+
+        // change the generation after the join-group request
+        final AbstractCoordinator.Generation newGen = new AbstractCoordinator.Generation(
+            currGen.generationId,
+            currGen.memberId + "-new",
+            currGen.protocolName);
+        coordinator.setNewGeneration(newGen);
+
+        mockClient.respond(joinGroupFollowerResponse(currGen.generationId + 1, memberId, JoinGroupRequest.UNKNOWN_MEMBER_ID, Errors.UNKNOWN_MEMBER_ID));
+
+        assertTrue(consumerClient.poll(future, mockTime.timer(REQUEST_TIMEOUT_MS)));
+        assertTrue(future.exception().getClass().isInstance(Errors.UNKNOWN_MEMBER_ID.exception()));
+
+        // the generation should not be reset
+        assertEquals(newGen, coordinator.generation());
+    }
+
+    @Test
+    public void testSyncGroupUnknownMemberResponseWithOldGeneration() throws InterruptedException {
+        setupCoordinator();
+        joinGroup();
+
+        final AbstractCoordinator.Generation currGen = coordinator.generation();
+
+        coordinator.setNewState(AbstractCoordinator.MemberState.PREPARING_REBALANCE);
+        RequestFuture future = coordinator.sendJoinGroupRequest();
+
+        TestUtils.waitForCondition(() -> {
+            consumerClient.poll(mockTime.timer(REQUEST_TIMEOUT_MS));
+            return !mockClient.requests().isEmpty();
+        }, 2000,
+            "The join-group request was not sent");
+
+        mockClient.respond(joinGroupFollowerResponse(currGen.generationId, memberId, JoinGroupRequest.UNKNOWN_MEMBER_ID, Errors.NONE));
+        assertTrue(mockClient.requests().isEmpty());
+
+        TestUtils.waitForCondition(() -> {
+            consumerClient.poll(mockTime.timer(REQUEST_TIMEOUT_MS));
+            return !mockClient.requests().isEmpty();
+        }, 2000,
+            "The sync-group request was not sent");
+
+        // change the generation after the sync-group request
+        final AbstractCoordinator.Generation newGen = new AbstractCoordinator.Generation(
+            currGen.generationId,
+            currGen.memberId + "-new",
+            currGen.protocolName);
+        coordinator.setNewGeneration(newGen);
+
+        mockClient.respond(syncGroupResponse(Errors.UNKNOWN_MEMBER_ID));
+        assertTrue(consumerClient.poll(future, mockTime.timer(REQUEST_TIMEOUT_MS)));
+        assertTrue(future.exception().getClass().isInstance(Errors.UNKNOWN_MEMBER_ID.exception()));
+
+        // the generation should not be reset
+        assertEquals(newGen, coordinator.generation());
+    }
+
+    @Test
+    public void testSyncGroupIllegalGenerationResponseWithOldGeneration() throws InterruptedException {
+        setupCoordinator();
+        joinGroup();
+
+        final AbstractCoordinator.Generation currGen = coordinator.generation();
+
+        coordinator.setNewState(AbstractCoordinator.MemberState.PREPARING_REBALANCE);
+        RequestFuture future = coordinator.sendJoinGroupRequest();
+
+        TestUtils.waitForCondition(() -> {
+            consumerClient.poll(mockTime.timer(REQUEST_TIMEOUT_MS));
+            return !mockClient.requests().isEmpty();
+        }, 2000,
+            "The join-group request was not sent");
+
+        mockClient.respond(joinGroupFollowerResponse(currGen.generationId, memberId, JoinGroupRequest.UNKNOWN_MEMBER_ID, Errors.NONE));
+        assertTrue(mockClient.requests().isEmpty());
+
+        TestUtils.waitForCondition(() -> {
+            consumerClient.poll(mockTime.timer(REQUEST_TIMEOUT_MS));
+            return !mockClient.requests().isEmpty();
+        }, 2000,
+            "The sync-group request was not sent");
+
+        // change the generation after the sync-group request
+        final AbstractCoordinator.Generation newGen = new AbstractCoordinator.Generation(
+            currGen.generationId,
+            currGen.memberId + "-new",
+            currGen.protocolName);
+        coordinator.setNewGeneration(newGen);
+
+        mockClient.respond(syncGroupResponse(Errors.ILLEGAL_GENERATION));
+        assertTrue(consumerClient.poll(future, mockTime.timer(REQUEST_TIMEOUT_MS)));
+        assertTrue(future.exception().getClass().isInstance(Errors.ILLEGAL_GENERATION.exception()));
+
+        // the generation should not be reset
+        assertEquals(newGen, coordinator.generation());
+    }
+
+    @Test
+    public void testHeartbeatSentWhenCompletingRebalance() throws Exception {
+        setupCoordinator();
+        joinGroup();
+
+        final AbstractCoordinator.Generation currGen = coordinator.generation();
+
+        coordinator.setNewState(AbstractCoordinator.MemberState.COMPLETING_REBALANCE);
+
+        // the heartbeat should be sent out during a rebalance
+        mockTime.sleep(HEARTBEAT_INTERVAL_MS);
+        TestUtils.waitForCondition(() -> !mockClient.requests().isEmpty(), 2000,
+                "The heartbeat request was not sent");
+        assertTrue(coordinator.heartbeat().hasInflight());
+
+        mockClient.respond(heartbeatResponse(Errors.REBALANCE_IN_PROGRESS));
+        assertEquals(currGen, coordinator.generation());
+    }
+
+    @Test
+    public void testHeartbeatIllegalGenerationResponseWithOldGeneration() throws InterruptedException {
+        setupCoordinator();
+        joinGroup();
+
+        final AbstractCoordinator.Generation currGen = coordinator.generation();
+
+        // let the heartbeat thread send out a request
+        mockTime.sleep(HEARTBEAT_INTERVAL_MS);
+
+        TestUtils.waitForCondition(() -> !mockClient.requests().isEmpty(), 2000,
+            "The heartbeat request was not sent");
+        assertTrue(coordinator.heartbeat().hasInflight());
+
+        // change the generation
+        final AbstractCoordinator.Generation newGen = new AbstractCoordinator.Generation(
+            currGen.generationId + 1,
+            currGen.memberId,
+            currGen.protocolName);
+        coordinator.setNewGeneration(newGen);
+
+        mockClient.respond(heartbeatResponse(Errors.ILLEGAL_GENERATION));
+
+        // the heartbeat error code should be ignored
+        TestUtils.waitForCondition(() -> {
+            coordinator.pollHeartbeat(mockTime.milliseconds());
+            return !coordinator.heartbeat().hasInflight();
+        }, 2000,
+            "The heartbeat response was not received");
+
+        // the generation should not be reset
+        assertEquals(newGen, coordinator.generation());
+    }
+
+    @Test
+    public void testHeartbeatUnknownMemberResponseWithOldGeneration() throws InterruptedException {
+        setupCoordinator();
+        joinGroup();
+
+        final AbstractCoordinator.Generation currGen = coordinator.generation();
+
+        // let the heartbeat request to send out a request
+        mockTime.sleep(HEARTBEAT_INTERVAL_MS);
+
+        TestUtils.waitForCondition(() -> !mockClient.requests().isEmpty(), 2000,
+            "The heartbeat request was not sent");
+        assertTrue(coordinator.heartbeat().hasInflight());
+
+        // change the generation
+        final AbstractCoordinator.Generation newGen = new AbstractCoordinator.Generation(
+            currGen.generationId,
+            currGen.memberId + "-new",
+            currGen.protocolName);
+        coordinator.setNewGeneration(newGen);
+
+        mockClient.respond(heartbeatResponse(Errors.UNKNOWN_MEMBER_ID));
+
+        // the heartbeat error code should be ignored
+        TestUtils.waitForCondition(() -> {
+            coordinator.pollHeartbeat(mockTime.milliseconds());
+            return !coordinator.heartbeat().hasInflight();
+        }, 2000,
+            "The heartbeat response was not received");
+
+        // the generation should not be reset
+        assertEquals(newGen, coordinator.generation());
+    }
+
+    @Test
+    public void testHeartbeatRebalanceInProgressResponseDuringRebalancing() throws InterruptedException {
+        setupCoordinator();
+        joinGroup();
+
+        final AbstractCoordinator.Generation currGen = coordinator.generation();
+
+        // let the heartbeat request to send out a request
+        mockTime.sleep(HEARTBEAT_INTERVAL_MS);
+
+        TestUtils.waitForCondition(() -> !mockClient.requests().isEmpty(), 2000,
+            "The heartbeat request was not sent");
+
+        assertTrue(coordinator.heartbeat().hasInflight());
+
+        mockClient.respond(heartbeatResponse(Errors.REBALANCE_IN_PROGRESS));
+
+        coordinator.requestRejoin();
+
+        TestUtils.waitForCondition(() -> {
+            coordinator.ensureActiveGroup(new MockTime(1L).timer(100L));
+            return !coordinator.heartbeat().hasInflight();
+        },
+            2000,
+            "The heartbeat response was not received");
+
+        // the generation would not be reset while the rebalance is in progress
+        assertEquals(currGen, coordinator.generation());
+
+        mockClient.respond(joinGroupFollowerResponse(currGen.generationId, memberId, JoinGroupRequest.UNKNOWN_MEMBER_ID, Errors.NONE));
+        mockClient.prepareResponse(syncGroupResponse(Errors.NONE));
+
+        coordinator.ensureActiveGroup();
+        assertEquals(currGen, coordinator.generation());
+    }
+
+    @Test
+    public void testHeartbeatInstanceFencedResponseWithOldGeneration() throws InterruptedException {
+        setupCoordinator();
+        joinGroup();
+
+        final AbstractCoordinator.Generation currGen = coordinator.generation();
+
+        // let the heartbeat request to send out a request
+        mockTime.sleep(HEARTBEAT_INTERVAL_MS);
+
+        TestUtils.waitForCondition(() -> !mockClient.requests().isEmpty(), 2000,
+            "The heartbeat request was not sent");
+        assertTrue(coordinator.heartbeat().hasInflight());
+
+        // change the generation
+        final AbstractCoordinator.Generation newGen = new AbstractCoordinator.Generation(
+            currGen.generationId,
+            currGen.memberId + "-new",
+            currGen.protocolName);
+        coordinator.setNewGeneration(newGen);
+
+        mockClient.respond(heartbeatResponse(Errors.FENCED_INSTANCE_ID));
+
+        // the heartbeat error code should be ignored
+        TestUtils.waitForCondition(() -> {
+            coordinator.pollHeartbeat(mockTime.milliseconds());
+            return !coordinator.heartbeat().hasInflight();
+        }, 2000,
+            "The heartbeat response was not received");
+
+        // the generation should not be reset
+        assertEquals(newGen, coordinator.generation());
+    }
+
+    @Test
+    public void testHeartbeatRequestWithFencedInstanceIdException() throws InterruptedException {
+        setupCoordinator();
+        mockClient.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+
+        final int generation = -1;
+
+        mockClient.prepareResponse(joinGroupFollowerResponse(generation, memberId, JoinGroupRequest.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, JoinGroupRequest.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(RETRY_BACKOFF_MS);
+        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 {
@@ -143,21 +950,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);
@@ -169,26 +974,26 @@ public boolean matches(AbstractRequest body) {
     }
 
     @Test
-    public void testLookupCoordinator() throws Exception {
-        setupCoordinator(RETRY_BACKOFF_MS);
+    public void testLookupCoordinator() {
+        setupCoordinator();
 
-        mockClient.setNode(null);
+        mockClient.backoff(node, 50);
         RequestFuture noBrokersAvailableFuture = coordinator.lookupCoordinator();
         assertTrue("Failed future expected", noBrokersAvailableFuture.failed());
+        mockTime.sleep(50);
 
-        mockClient.setNode(node);
         RequestFuture future = coordinator.lookupCoordinator();
         assertFalse("Request not sent", future.isDone());
-        assertTrue("New request sent while one is in progress", future == coordinator.lookupCoordinator());
+        assertSame("New request sent while one is in progress", future, coordinator.lookupCoordinator());
 
         mockClient.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
-        assertTrue("New request not sent after previous completed", future != coordinator.lookupCoordinator());
+        coordinator.ensureCoordinatorReady(mockTime.timer(Long.MAX_VALUE));
+        assertNotSame("New request not sent after previous completed", future, coordinator.lookupCoordinator());
     }
 
     @Test
     public void testWakeupAfterJoinGroupSent() throws Exception {
-        setupCoordinator(RETRY_BACKOFF_MS);
+        setupCoordinator();
 
         mockClient.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
         mockClient.prepareResponse(new MockClient.RequestMatcher() {
@@ -202,14 +1007,14 @@ 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();
 
         try {
             coordinator.ensureActiveGroup();
             fail("Should have woken up from ensureActiveGroup()");
-        } catch (WakeupException e) {
+        } catch (WakeupException ignored) {
         }
 
         assertEquals(1, coordinator.onJoinPrepareInvokes);
@@ -226,7 +1031,7 @@ public boolean matches(AbstractRequest body) {
 
     @Test
     public void testWakeupAfterJoinGroupSentExternalCompletion() throws Exception {
-        setupCoordinator(RETRY_BACKOFF_MS);
+        setupCoordinator();
 
         mockClient.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
         mockClient.prepareResponse(new MockClient.RequestMatcher() {
@@ -240,14 +1045,14 @@ 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();
 
         try {
             coordinator.ensureActiveGroup();
             fail("Should have woken up from ensureActiveGroup()");
-        } catch (WakeupException e) {
+        } catch (WakeupException ignored) {
         }
 
         assertEquals(1, coordinator.onJoinPrepareInvokes);
@@ -255,7 +1060,7 @@ public boolean matches(AbstractRequest body) {
         assertFalse(heartbeatReceived.get());
 
         // the join group completes in this poll()
-        consumerClient.poll(0);
+        consumerClient.poll(mockTime.timer(0));
         coordinator.ensureActiveGroup();
 
         assertEquals(1, coordinator.onJoinPrepareInvokes);
@@ -266,26 +1071,23 @@ public boolean matches(AbstractRequest body) {
 
     @Test
     public void testWakeupAfterJoinGroupReceived() throws Exception {
-        setupCoordinator(RETRY_BACKOFF_MS);
+        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();
 
         try {
             coordinator.ensureActiveGroup();
             fail("Should have woken up from ensureActiveGroup()");
-        } catch (WakeupException e) {
+        } catch (WakeupException ignored) {
         }
 
         assertEquals(1, coordinator.onJoinPrepareInvokes);
@@ -302,19 +1104,16 @@ public boolean matches(AbstractRequest body) {
 
     @Test
     public void testWakeupAfterJoinGroupReceivedExternalCompletion() throws Exception {
-        setupCoordinator(RETRY_BACKOFF_MS);
+        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();
 
@@ -329,45 +1128,7 @@ public boolean matches(AbstractRequest body) {
         assertFalse(heartbeatReceived.get());
 
         // the join group completes in this poll()
-        consumerClient.poll(0);
-        coordinator.ensureActiveGroup();
-
-        assertEquals(1, coordinator.onJoinPrepareInvokes);
-        assertEquals(1, coordinator.onJoinCompleteInvokes);
-
-        awaitFirstHeartbeat(heartbeatReceived);
-    }
-
-    @Test
-    public void testWakeupAfterSyncGroupSent() throws Exception {
-        setupCoordinator(RETRY_BACKOFF_MS);
-
-        mockClient.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        mockClient.prepareResponse(joinGroupFollowerResponse(1, "memberId", "leaderId", Errors.NONE));
-        mockClient.prepareResponse(new MockClient.RequestMatcher() {
-            private int invocations = 0;
-            @Override
-            public boolean matches(AbstractRequest body) {
-                invocations++;
-                boolean isSyncGroupRequest = body instanceof SyncGroupRequest;
-                if (isSyncGroupRequest && invocations == 1)
-                    // simulate wakeup after the request sent
-                    throw new WakeupException();
-                return isSyncGroupRequest;
-            }
-        }, syncGroupResponse(Errors.NONE));
-        AtomicBoolean heartbeatReceived = prepareFirstHeartbeat();
-
-        try {
-            coordinator.ensureActiveGroup();
-            fail("Should have woken up from ensureActiveGroup()");
-        } catch (WakeupException e) {
-        }
-
-        assertEquals(1, coordinator.onJoinPrepareInvokes);
-        assertEquals(0, coordinator.onJoinCompleteInvokes);
-        assertFalse(heartbeatReceived.get());
-
+        consumerClient.poll(mockTime.timer(0));
         coordinator.ensureActiveGroup();
 
         assertEquals(1, coordinator.onJoinPrepareInvokes);
@@ -378,10 +1139,10 @@ public boolean matches(AbstractRequest body) {
 
     @Test
     public void testWakeupAfterSyncGroupSentExternalCompletion() throws Exception {
-        setupCoordinator(RETRY_BACKOFF_MS);
+        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
@@ -389,8 +1150,8 @@ public boolean matches(AbstractRequest body) {
                 invocations++;
                 boolean isSyncGroupRequest = body instanceof SyncGroupRequest;
                 if (isSyncGroupRequest && invocations == 1)
-                    // simulate wakeup after the request sent
-                    throw new WakeupException();
+                    // wakeup after the request returns
+                    consumerClient.wakeup();
                 return isSyncGroupRequest;
             }
         }, syncGroupResponse(Errors.NONE));
@@ -407,7 +1168,7 @@ public boolean matches(AbstractRequest body) {
         assertFalse(heartbeatReceived.get());
 
         // the join group completes in this poll()
-        consumerClient.poll(0);
+        consumerClient.poll(mockTime.timer(0));
         coordinator.ensureActiveGroup();
 
         assertEquals(1, coordinator.onJoinPrepareInvokes);
@@ -418,26 +1179,23 @@ public boolean matches(AbstractRequest body) {
 
     @Test
     public void testWakeupAfterSyncGroupReceived() throws Exception {
-        setupCoordinator(RETRY_BACKOFF_MS);
+        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();
 
         try {
             coordinator.ensureActiveGroup();
             fail("Should have woken up from ensureActiveGroup()");
-        } catch (WakeupException e) {
+        } catch (WakeupException ignored) {
         }
 
         assertEquals(1, coordinator.onJoinPrepareInvokes);
@@ -454,19 +1212,16 @@ public boolean matches(AbstractRequest body) {
 
     @Test
     public void testWakeupAfterSyncGroupReceivedExternalCompletion() throws Exception {
-        setupCoordinator(RETRY_BACKOFF_MS);
+        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();
 
@@ -490,18 +1245,18 @@ public boolean matches(AbstractRequest body) {
 
     @Test
     public void testWakeupInOnJoinComplete() throws Exception {
-        setupCoordinator(RETRY_BACKOFF_MS);
+        setupCoordinator();
 
         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();
 
         try {
             coordinator.ensureActiveGroup();
             fail("Should have woken up from ensureActiveGroup()");
-        } catch (WakeupException e) {
+        } catch (WakeupException ignored) {
         }
 
         assertEquals(1, coordinator.onJoinPrepareInvokes);
@@ -510,7 +1265,7 @@ public void testWakeupInOnJoinComplete() throws Exception {
 
         // the join group completes in this poll()
         coordinator.wakeupOnJoinComplete = false;
-        consumerClient.poll(0);
+        consumerClient.poll(mockTime.timer(0));
         coordinator.ensureActiveGroup();
 
         assertEquals(1, coordinator.onJoinPrepareInvokes);
@@ -519,45 +1274,94 @@ public void testWakeupInOnJoinComplete() throws Exception {
         awaitFirstHeartbeat(heartbeatReceived);
     }
 
+    @Test
+    public void testAuthenticationErrorInEnsureCoordinatorReady() {
+        setupCoordinator();
+
+        mockClient.createPendingAuthenticationError(node, 300);
+
+        try {
+            coordinator.ensureCoordinatorReady(mockTime.timer(Long.MAX_VALUE));
+            fail("Expected an authentication error.");
+        } catch (AuthenticationException e) {
+            // OK
+        }
+    }
+
     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;
     }
 
     private void awaitFirstHeartbeat(final AtomicBoolean heartbeatReceived) throws Exception {
         mockTime.sleep(HEARTBEAT_INTERVAL_MS);
-        TestUtils.waitForCondition(new TestCondition() {
-            @Override
-            public boolean conditionMet() {
-                return heartbeatReceived.get();
-            }
-        }, 3000, "Should have received a heartbeat request after joining the group");
+        TestUtils.waitForCondition(heartbeatReceived::get,
+            3000, "Should have received a heartbeat request after joining the group");
     }
 
     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 joinGroupFollowerResponse(generationId, memberId, leaderId, error, null);
+    }
+
+    private JoinGroupResponse joinGroupFollowerResponse(int generationId,
+                                                        String memberId,
+                                                        String leaderId,
+                                                        Errors error,
+                                                        String protocolType) {
+        return new JoinGroupResponse(
+                new JoinGroupResponseData()
+                        .setErrorCode(error.code())
+                        .setGenerationId(generationId)
+                        .setProtocolType(protocolType)
+                        .setProtocolName(PROTOCOL_NAME)
+                        .setMemberId(memberId)
+                        .setLeader(leaderId)
+                        .setMembers(Collections.emptyList())
+        );
     }
 
-    private JoinGroupResponse joinGroupFollowerResponse(int generationId, String memberId, String leaderId, Errors error) {
-        return new JoinGroupResponse(error, generationId, "dummy-subprotocol", memberId, leaderId,
-                Collections.emptyMap());
+    private JoinGroupResponse joinGroupResponse(Errors error) {
+        return joinGroupFollowerResponse(JoinGroupRequest.UNKNOWN_GENERATION_ID,
+            JoinGroupRequest.UNKNOWN_MEMBER_ID, JoinGroupRequest.UNKNOWN_MEMBER_ID, error);
     }
 
     private SyncGroupResponse syncGroupResponse(Errors error) {
-        return new SyncGroupResponse(error, ByteBuffer.allocate(0));
+        return syncGroupResponse(error, null, null);
+    }
+
+    private SyncGroupResponse syncGroupResponse(Errors error,
+                                                String protocolType,
+                                                String protocolName) {
+        return new SyncGroupResponse(
+                new SyncGroupResponseData()
+                        .setErrorCode(error.code())
+                        .setProtocolType(protocolType)
+                        .setProtocolName(protocolName)
+                        .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 {
@@ -566,28 +1370,35 @@ public static class DummyCoordinator extends AbstractCoordinator {
         private int onJoinCompleteInvokes = 0;
         private boolean wakeupOnJoinComplete = false;
 
-        public DummyCoordinator(ConsumerNetworkClient client,
-                                Metrics metrics,
-                                Time time) {
-            super(new LogContext(), client, GROUP_ID, REBALANCE_TIMEOUT_MS, SESSION_TIMEOUT_MS,
-                    HEARTBEAT_INTERVAL_MS, metrics, METRIC_GROUP_PREFIX, time, RETRY_BACKOFF_MS, false);
+        DummyCoordinator(GroupRebalanceConfig rebalanceConfig,
+                         ConsumerNetworkClient client,
+                         Metrics metrics,
+                         Time time) {
+            super(rebalanceConfig, new LogContext(), client, metrics, METRIC_GROUP_PREFIX, time);
         }
 
         @Override
         protected String protocolType() {
-            return "dummy";
+            return PROTOCOL_TYPE;
         }
 
         @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(PROTOCOL_NAME)
+                            .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..7ce765ad11e59
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignorTest.java
@@ -0,0 +1,780 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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));
+
+        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));
+    }
+
+    /**
+     * 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));
+
+        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));
+    }
+
+
+    @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);
+    }
+
+    @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);
+    }
+
+    @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);
+
+        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);
+
+                int moreLoaded = len > otherLen ? i : j;
+                int lessLoaded = len > otherLen ? j : i;
+
+                // If there's any overlap in the subscribed topics, we should have been able to balance partitions
+                for (String topic: map.keySet()) {
+                    assertFalse("Error: Some partitions can be moved from c" + moreLoaded + " to c" + lessLoaded
+                            + " 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));
+                }
+            }
+        }
+    }
+
+}
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 90b839d5caacd..1c7cfb47bdbe3 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
@@ -16,47 +16,69 @@
  */
 package org.apache.kafka.clients.consumer.internals;
 
-import org.apache.kafka.clients.Metadata;
+import org.apache.kafka.clients.ClientResponse;
+import org.apache.kafka.clients.GroupRebalanceConfig;
 import org.apache.kafka.clients.MockClient;
+import org.apache.kafka.clients.NodeApiVersions;
 import org.apache.kafka.clients.consumer.CommitFailedException;
-import org.apache.kafka.clients.consumer.ConsumerConfig;
-import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
+import org.apache.kafka.clients.consumer.ConsumerGroupMetadata;
+import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor;
 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.Cluster;
 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.TopicPartition;
 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.RebalanceInProgressException;
+import org.apache.kafka.common.errors.TopicAuthorizationException;
+import org.apache.kafka.common.errors.UnsupportedVersionException;
 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.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.requests.AbstractRequest;
+import org.apache.kafka.common.record.RecordBatch;
 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;
+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.OffsetFetchResponse;
+import org.apache.kafka.common.requests.RequestTestUtils;
 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.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;
@@ -65,8 +87,10 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -74,75 +98,200 @@
 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.AtomicReference;
 import java.util.regex.Pattern;
 
 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.common.utils.Utils.mkEntry;
+import static org.apache.kafka.common.utils.Utils.mkMap;
+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;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertThrows;
 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;
+    private final long retryBackoffMs = 100;
+    private final int autoCommitIntervalMs = 2000;
+    private final int requestTimeoutMs = 30000;
+    private final MockTime time = new MockTime();
+    private GroupRebalanceConfig rebalanceConfig;
+
+    private final ConsumerPartitionAssignor.RebalanceProtocol protocol;
+    private final MockPartitionAssignor partitionAssignor;
+    private final ThrowOnAssignmentAssignor throwOnAssignmentAssignor;
+    private final ThrowOnAssignmentAssignor throwFatalErrorOnAssignmentAssignor;
+    private final List assignors;
+    private final Map assignorMap;
+    private final String consumerId = "consumer";
 
-    private String topic1 = "test1";
-    private String topic2 = "test2";
-    private String groupId = "test-group";
-    private TopicPartition t1p = new TopicPartition(topic1, 0);
-    private TopicPartition t2p = new TopicPartition(topic2, 0);
-    private int rebalanceTimeoutMs = 60000;
-    private int sessionTimeoutMs = 10000;
-    private int heartbeatIntervalMs = 5000;
-    private long retryBackoffMs = 100;
-    private boolean autoCommitEnabled = false;
-    private int autoCommitIntervalMs = 2000;
-    private MockPartitionAssignor partitionAssignor = new MockPartitionAssignor();
-    private List assignors = Collections.singletonList(partitionAssignor);
-    private MockTime time;
     private MockClient client;
-    private Cluster cluster = TestUtils.clusterWith(1, new HashMap() {
+    private MetadataResponse metadataResponse = RequestTestUtils.metadataUpdateWith(1, new HashMap() {
         {
             put(topic1, 1);
             put(topic2, 1);
         }
     });
-    private Node node = cluster.nodes().get(0);
+    private Node node = metadataResponse.brokers().iterator().next();
     private SubscriptionState subscriptions;
-    private Metadata metadata;
+    private ConsumerMetadata metadata;
     private Metrics metrics;
     private ConsumerNetworkClient consumerClient;
     private MockRebalanceListener rebalanceListener;
     private MockCommitCallback mockOffsetCommitCallback;
     private ConsumerCoordinator coordinator;
 
+    public ConsumerCoordinatorTest(final ConsumerPartitionAssignor.RebalanceProtocol protocol) {
+        this.protocol = protocol;
+
+        this.partitionAssignor = new MockPartitionAssignor(Collections.singletonList(protocol));
+        this.throwOnAssignmentAssignor = new ThrowOnAssignmentAssignor(Collections.singletonList(protocol),
+            new KafkaException("Kaboom for assignment!"),
+            "throw-on-assignment-assignor");
+        this.throwFatalErrorOnAssignmentAssignor = new ThrowOnAssignmentAssignor(Collections.singletonList(protocol),
+            new IllegalStateException("Illegal state for assignment!"),
+            "throw-fatal-error-on-assignment-assignor");
+        this.assignors = Arrays.asList(partitionAssignor, throwOnAssignmentAssignor, throwFatalErrorOnAssignmentAssignor);
+        this.assignorMap = mkMap(mkEntry(partitionAssignor.name(), partitionAssignor),
+            mkEntry(throwOnAssignmentAssignor.name(), throwOnAssignmentAssignor),
+            mkEntry(throwFatalErrorOnAssignmentAssignor.name(), throwFatalErrorOnAssignmentAssignor));
+    }
+
+    @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() {
-        this.time = new MockTime();
-        this.subscriptions = new SubscriptionState(OffsetResetStrategy.EARLIEST);
-        this.metadata = new Metadata(0, Long.MAX_VALUE, true);
-        this.metadata.update(cluster, Collections.emptySet(), time.milliseconds());
+        LogContext logContext = new LogContext();
+        this.subscriptions = new SubscriptionState(logContext, OffsetResetStrategy.EARLIEST);
+        this.metadata = new ConsumerMetadata(0, Long.MAX_VALUE, false,
+                false, subscriptions, logContext, new ClusterResourceListeners());
         this.client = new MockClient(time, metadata);
-        this.consumerClient = new ConsumerNetworkClient(new LogContext(), client, metadata, time, 100, 1000, Integer.MAX_VALUE);
+        this.client.updateMetadata(metadataResponse);
+        this.consumerClient = new ConsumerNetworkClient(logContext, client, metadata, time, 100,
+                requestTimeoutMs, Integer.MAX_VALUE);
         this.metrics = new Metrics(time);
         this.rebalanceListener = new MockRebalanceListener();
         this.mockOffsetCommitCallback = new MockCommitCallback();
         this.partitionAssignor.clear();
+        this.rebalanceConfig = buildRebalanceConfig(Optional.empty());
+        this.coordinator = buildCoordinator(rebalanceConfig,
+                                            metrics,
+                                            assignors,
+                                            false);
+    }
 
-        client.setNode(node);
-        this.coordinator = buildCoordinator(metrics, assignors, ConsumerConfig.DEFAULT_EXCLUDE_INTERNAL_TOPICS, autoCommitEnabled, true);
+    private GroupRebalanceConfig buildRebalanceConfig(Optional groupInstanceId) {
+        return new GroupRebalanceConfig(sessionTimeoutMs,
+                                        rebalanceTimeoutMs,
+                                        heartbeatIntervalMs,
+                                        groupId,
+                                        groupInstanceId,
+                                        retryBackoffMs,
+                                        !groupInstanceId.isPresent());
     }
 
     @After
     public void teardown() {
         this.metrics.close();
+        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, consumerId + groupId + "-coordinator-metrics"));
+    }
+
+    @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));
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         // normal heartbeat
         time.sleep(sessionTimeoutMs);
@@ -151,7 +300,7 @@ public void testNormalHeartbeat() {
         assertFalse(future.isDone());
 
         client.prepareResponse(heartbeatResponse(Errors.NONE));
-        consumerClient.poll(0);
+        consumerClient.poll(time.timer(0));
 
         assertTrue(future.isDone());
         assertTrue(future.succeeded());
@@ -160,7 +309,7 @@ public void testNormalHeartbeat() {
     @Test(expected = GroupAuthorizationException.class)
     public void testGroupDescribeUnauthorized() {
         client.prepareResponse(groupCoordinatorResponse(node, Errors.GROUP_AUTHORIZATION_FAILED));
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
     }
 
     @Test(expected = GroupAuthorizationException.class)
@@ -168,17 +317,17 @@ public void testGroupReadUnauthorized() {
         subscriptions.subscribe(singleton(topic1), rebalanceListener);
 
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        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.milliseconds(), Long.MAX_VALUE);
+        coordinator.poll(time.timer(Long.MAX_VALUE));
     }
 
     @Test
     public void testCoordinatorNotAvailable() {
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         // COORDINATOR_NOT_AVAILABLE will mark coordinator as unknown
         time.sleep(sessionTimeoutMs);
@@ -188,7 +337,7 @@ public void testCoordinatorNotAvailable() {
 
         client.prepareResponse(heartbeatResponse(Errors.COORDINATOR_NOT_AVAILABLE));
         time.sleep(sessionTimeoutMs);
-        consumerClient.poll(0);
+        consumerClient.poll(time.timer(0));
 
         assertTrue(future.isDone());
         assertTrue(future.failed());
@@ -196,10 +345,80 @@ public void testCoordinatorNotAvailable() {
         assertTrue(coordinator.coordinatorUnknown());
     }
 
+    @Test
+    public void testManyInFlightAsyncCommitsWithCoordinatorDisconnect() {
+        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
+
+        int numRequests = 1000;
+        TopicPartition tp = new TopicPartition("foo", 0);
+        final AtomicInteger responses = new AtomicInteger(0);
+
+        for (int i = 0; i < numRequests; i++) {
+            Map offsets = singletonMap(tp, new OffsetAndMetadata(i));
+            coordinator.commitOffsetsAsync(offsets, (offsets1, exception) -> {
+                responses.incrementAndGet();
+                Throwable cause = exception.getCause();
+                assertTrue("Unexpected exception cause type: " + (cause == null ? null : cause.getClass()),
+                        cause instanceof DisconnectException);
+            });
+        }
+
+        coordinator.markCoordinatorUnknown();
+        consumerClient.pollNoWakeup();
+        coordinator.invokeCompletedOffsetCommitCallbacks();
+        assertEquals(numRequests, responses.get());
+    }
+
+    @Test
+    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.
+
+        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
+
+        final AtomicBoolean asyncCallbackInvoked = new AtomicBoolean(false);
+
+        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) {}
+
+                    @Override
+                    public void onFailure(RuntimeException e, RequestFuture future) {
+                        assertTrue("Unexpected exception type: " + e.getClass(), e instanceof DisconnectException);
+                        assertTrue(coordinator.coordinatorUnknown());
+                        asyncCallbackInvoked.set(true);
+                    }
+                });
+
+        coordinator.markCoordinatorUnknown();
+        consumerClient.pollNoWakeup();
+        assertTrue(asyncCallbackInvoked.get());
+    }
+
     @Test
     public void testNotCoordinator() {
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         // not_coordinator will mark coordinator as unknown
         time.sleep(sessionTimeoutMs);
@@ -209,7 +428,7 @@ public void testNotCoordinator() {
 
         client.prepareResponse(heartbeatResponse(Errors.NOT_COORDINATOR));
         time.sleep(sessionTimeoutMs);
-        consumerClient.poll(0);
+        consumerClient.poll(time.timer(0));
 
         assertTrue(future.isDone());
         assertTrue(future.failed());
@@ -220,7 +439,7 @@ public void testNotCoordinator() {
     @Test
     public void testIllegalGeneration() {
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         // illegal_generation will cause re-partition
         subscriptions.subscribe(singleton(topic1), rebalanceListener);
@@ -233,18 +452,179 @@ public void testIllegalGeneration() {
 
         client.prepareResponse(heartbeatResponse(Errors.ILLEGAL_GENERATION));
         time.sleep(sessionTimeoutMs);
-        consumerClient.poll(0);
+        consumerClient.poll(time.timer(0));
 
         assertTrue(future.isDone());
         assertTrue(future.failed());
         assertEquals(Errors.ILLEGAL_GENERATION.exception(), future.exception());
-        assertTrue(coordinator.needRejoin());
+        assertTrue(coordinator.rejoinNeededOrPending());
+
+        coordinator.poll(time.timer(0));
+
+        assertEquals(1, rebalanceListener.lostCount);
+        assertEquals(Collections.singleton(t1p), rebalanceListener.lost);
+    }
+
+    @Test
+    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 testRevokeExceptionThrownFirstNonBlockingSubCallbacks() {
+        MockRebalanceListener throwOnRevokeListener = new MockRebalanceListener() {
+            @Override
+            public void onPartitionsRevoked(Collection partitions) {
+                super.onPartitionsRevoked(partitions);
+                throw new KafkaException("Kaboom on revoke!");
+            }
+        };
+
+        if (protocol == COOPERATIVE) {
+            verifyOnCallbackExceptions(throwOnRevokeListener,
+                throwOnAssignmentAssignor.name(), "Kaboom on revoke!", null);
+        } else {
+            // Eager protocol doesn't revoke partitions.
+            verifyOnCallbackExceptions(throwOnRevokeListener,
+                throwOnAssignmentAssignor.name(), "Kaboom for assignment!", null);
+        }
+    }
+
+    @Test
+    public void testOnAssignmentExceptionThrownFirstNonBlockingSubCallbacks() {
+        MockRebalanceListener throwOnAssignListener = new MockRebalanceListener() {
+            @Override
+            public void onPartitionsAssigned(Collection partitions) {
+                super.onPartitionsAssigned(partitions);
+                throw new KafkaException("Kaboom on partition assign!");
+            }
+        };
+
+        verifyOnCallbackExceptions(throwOnAssignListener,
+            throwOnAssignmentAssignor.name(), "Kaboom for assignment!", null);
+    }
+
+    @Test
+    public void testOnPartitionsAssignExceptionThrownWhenNoPreviousThrownCallbacks() {
+        MockRebalanceListener throwOnAssignListener = new MockRebalanceListener() {
+            @Override
+            public void onPartitionsAssigned(Collection partitions) {
+                super.onPartitionsAssigned(partitions);
+                throw new KafkaException("Kaboom on partition assign!");
+            }
+        };
+
+        verifyOnCallbackExceptions(throwOnAssignListener,
+            partitionAssignor.name(), "Kaboom on partition assign!", null);
     }
 
     @Test
-    public void testUnknownConsumerId() {
+    public void testOnRevokeExceptionShouldBeRenderedIfNotKafkaException() {
+        MockRebalanceListener throwOnRevokeListener = new MockRebalanceListener() {
+            @Override
+            public void onPartitionsRevoked(Collection partitions) {
+                super.onPartitionsRevoked(partitions);
+                throw new IllegalStateException("Illegal state on partition revoke!");
+            }
+        };
+
+        if (protocol == COOPERATIVE) {
+            verifyOnCallbackExceptions(throwOnRevokeListener,
+                throwOnAssignmentAssignor.name(),
+                "User rebalance callback throws an error", "Illegal state on partition revoke!");
+        } else {
+            // Eager protocol doesn't revoke partitions.
+            verifyOnCallbackExceptions(throwOnRevokeListener,
+                throwOnAssignmentAssignor.name(), "Kaboom for assignment!", null);
+        }
+    }
+
+    @Test
+    public void testOnAssignmentExceptionShouldBeRenderedIfNotKafkaException() {
+        MockRebalanceListener throwOnAssignListener = new MockRebalanceListener() {
+            @Override
+            public void onPartitionsAssigned(Collection partitions) {
+                super.onPartitionsAssigned(partitions);
+                throw new KafkaException("Kaboom on partition assign!");
+            }
+        };
+        verifyOnCallbackExceptions(throwOnAssignListener,
+            throwFatalErrorOnAssignmentAssignor.name(),
+            "User rebalance callback throws an error", "Illegal state for assignment!");
+    }
+
+    @Test
+    public void testOnPartitionsAssignExceptionShouldBeRenderedIfNotKafkaException() {
+        MockRebalanceListener throwOnAssignListener = new MockRebalanceListener() {
+            @Override
+            public void onPartitionsAssigned(Collection partitions) {
+                super.onPartitionsAssigned(partitions);
+                throw new IllegalStateException("Illegal state on partition assign!");
+            }
+        };
+
+        verifyOnCallbackExceptions(throwOnAssignListener,
+            partitionAssignor.name(), "User rebalance callback throws an error",
+            "Illegal state on partition assign!");
+    }
+
+    private void verifyOnCallbackExceptions(final MockRebalanceListener rebalanceListener,
+                                            final String assignorName,
+                                            final String exceptionMessage,
+                                            final String causeMessage) {
+        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])));
+        subscriptions.assignFromSubscribed(singleton(t2p));
+
+        if (exceptionMessage != null) {
+            final Exception exception = assertThrows(KafkaException.class,
+                () -> coordinator.onJoinComplete(1, "memberId", assignorName, buffer));
+            assertEquals(exceptionMessage, exception.getMessage());
+            if (causeMessage != null) {
+                assertEquals(causeMessage, exception.getCause().getMessage());
+            }
+        }
+
+        // Eager doesn't trigger on partition revoke.
+        assertEquals(protocol == COOPERATIVE ? 1 : 0, rebalanceListener.revokedCount);
+        assertEquals(0, rebalanceListener.lostCount);
+        assertEquals(1, rebalanceListener.assignedCount);
+        assertTrue("Unknown assignor name: " + assignorName,
+            assignorMap.containsKey(assignorName));
+        assertEquals(1, assignorMap.get(assignorName).numAssignment());
+    }
+
+    @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();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         // illegal_generation will cause re-partition
         subscriptions.subscribe(singleton(topic1), rebalanceListener);
@@ -257,18 +637,23 @@ public void testUnknownConsumerId() {
 
         client.prepareResponse(heartbeatResponse(Errors.UNKNOWN_MEMBER_ID));
         time.sleep(sessionTimeoutMs);
-        consumerClient.poll(0);
+        consumerClient.poll(time.timer(0));
 
         assertTrue(future.isDone());
         assertTrue(future.failed());
         assertEquals(Errors.UNKNOWN_MEMBER_ID.exception(), future.exception());
-        assertTrue(coordinator.needRejoin());
+        assertTrue(coordinator.rejoinNeededOrPending());
+
+        coordinator.poll(time.timer(0));
+
+        assertEquals(1, rebalanceListener.lostCount);
+        assertEquals(Collections.singleton(t1p), rebalanceListener.lost);
     }
 
     @Test
     public void testCoordinatorDisconnect() {
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         // coordinator disconnect will mark coordinator as unknown
         time.sleep(sessionTimeoutMs);
@@ -278,7 +663,7 @@ public void testCoordinatorDisconnect() {
 
         client.prepareResponse(heartbeatResponse(Errors.NONE), true); // return disconnected
         time.sleep(sessionTimeoutMs);
-        consumerClient.poll(0);
+        consumerClient.poll(time.timer(0));
 
         assertTrue(future.isDone());
         assertTrue(future.failed());
@@ -293,450 +678,711 @@ public void testJoinGroupInvalidGroupId() {
         subscriptions.subscribe(singleton(topic1), rebalanceListener);
 
         // ensure metadata is up-to-date for leader
-        metadata.setTopics(singletonList(topic1));
-        metadata.update(cluster, Collections.emptySet(), time.milliseconds());
+        client.updateMetadata(metadataResponse);
 
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        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.milliseconds(), Long.MAX_VALUE);
+        coordinator.poll(time.timer(Long.MAX_VALUE));
     }
 
     @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);
 
         // ensure metadata is up-to-date for leader
-        metadata.setTopics(singletonList(topic1));
-        metadata.update(cluster, Collections.emptySet(), time.milliseconds());
+        client.updateMetadata(metadataResponse);
 
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         // normal join group
-        Map> memberSubscriptions = Collections.singletonMap(consumerId, singletonList(topic1));
-        partitionAssignor.prepare(Collections.singletonMap(consumerId, singletonList(t1p)));
+        Map> memberSubscriptions = singletonMap(consumerId, singletonList(topic1));
+        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));
-        coordinator.poll(time.milliseconds(), Long.MAX_VALUE);
+        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(toSet(assigned), subscriptions.assignedPartitions());
+        assertEquals(subscription, subscriptions.metadataTopics());
+        assertEquals(0, rebalanceListener.revokedCount);
+        assertNull(rebalanceListener.revoked);
+        assertEquals(1, rebalanceListener.assignedCount);
+        assertEquals(getAdded(owned, assigned), rebalanceListener.assigned);
+    }
 
-        assertFalse(coordinator.needRejoin());
-        assertEquals(singleton(t1p), subscriptions.assignedPartitions());
-        assertEquals(singleton(topic1), subscriptions.groupSubscription());
-        assertEquals(1, rebalanceListener.revokedCount);
-        assertEquals(Collections.emptySet(), rebalanceListener.revoked);
+    @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(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, newAssignment));
+
+        // First incorrect assignment for subscription
+        client.prepareResponse(
+                joinGroupLeaderResponse(
+                    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, 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(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(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 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(toSet(oldSubscription), rebalanceListener);
+        assertEquals(toSet(oldSubscription), subscriptions.metadataTopics());
+
+        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
+
+        prepareJoinAndSyncResponse(consumerId, 1, oldSubscription, oldAssignment);
+
+        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);
 
         // partially update the metadata with one topic first,
         // let the leader to refresh metadata during assignment
-        metadata.setTopics(singletonList(topic1));
-        metadata.update(TestUtils.singletonCluster(topic1, 1), Collections.emptySet(), time.milliseconds());
+        client.updateMetadata(RequestTestUtils.metadataUpdateWith(1, singletonMap(topic1, 1)));
 
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         // normal join group
-        Map> memberSubscriptions = Collections.singletonMap(consumerId, singletonList(topic1));
-        partitionAssignor.prepare(Collections.singletonMap(consumerId, Arrays.asList(t1p, t2p)));
+        Map> memberSubscriptions = singletonMap(consumerId, singletonList(topic1));
+        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(cluster, Collections.emptySet());
+        client.prepareMetadataUpdate(metadataResponse);
 
-        coordinator.poll(time.milliseconds(), Long.MAX_VALUE);
+        coordinator.poll(time.timer(Long.MAX_VALUE));
 
-        assertFalse(coordinator.needRejoin());
-        assertEquals(2, subscriptions.assignedPartitions().size());
-        assertEquals(2, subscriptions.groupSubscription().size());
+        assertFalse(coordinator.rejoinNeededOrPending());
+        assertEquals(2, subscriptions.numAssignedPartitions());
+        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 = singletonList(t1p);
         subscriptions.subscribe(Pattern.compile(".*"), rebalanceListener);
-        metadata.needMetadataForAllTopics(true);
-        metadata.update(TestUtils.singletonCluster(topic1, 1), Collections.emptySet(), time.milliseconds());
+        client.updateMetadata(RequestTestUtils.metadataUpdateWith(1, singletonMap(topic1, 1)));
+        coordinator.maybeUpdateSubscriptionMetadata();
 
         assertEquals(singleton(topic1), subscriptions.subscription());
 
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        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(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                final Map updatedPartitions = new HashMap<>();
-                for (String topic : updatedSubscription)
-                    updatedPartitions.put(topic, 1);
-                metadata.update(TestUtils.clusterWith(1, updatedPartitions), Collections.emptySet(), time.milliseconds());
-                return true;
-            }
-        }, syncGroupResponse(singletonList(t1p), Errors.NONE));
+        client.prepareResponse(body -> {
+            final Map updatedPartitions = new HashMap<>();
+            for (String topic : updatedSubscription)
+                updatedPartitions.put(topic, 1);
+            client.updateMetadata(RequestTestUtils.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 newAssignment = Arrays.asList(t1p, t2p);
-        Set newAssignmentSet = new HashSet<>(newAssignment);
+        List newAssigned = Arrays.asList(t1p, t2p);
 
-        Map> updatedSubscriptions = singletonMap(consumerId, Arrays.asList(topic1, topic2));
-        partitionAssignor.prepare(singletonMap(consumerId, newAssignment));
+        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(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);
-            }
+        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(syncGroupResponse(newAssignment, Errors.NONE));
+        // update the metadata again back to topic1
+        client.prepareResponse(body -> {
+            client.updateMetadata(RequestTestUtils.metadataUpdateWith(1, singletonMap(topic1, 1)));
+            return true;
+        }, syncGroupResponse(newAssigned, Errors.NONE));
+
+        coordinator.poll(time.timer(Long.MAX_VALUE));
+
+        Collection revoked = getRevoked(oldAssigned, newAssigned);
+        int revokedCount = revoked.isEmpty() ? 0 : 1;
+
+        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));
+
+        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().contains(topic1);
+        }, 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(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
+    public void testForceMetadataRefreshForPatternSubscriptionDuringRebalance() {
+        // Set up a non-leader consumer with pattern subscription and a cluster containing one topic matching the
+        // pattern.
+        subscriptions.subscribe(Pattern.compile(".*"), rebalanceListener);
+        client.updateMetadata(RequestTestUtils.metadataUpdateWith(1, singletonMap(topic1, 1)));
+        coordinator.maybeUpdateSubscriptionMetadata();
+        assertEquals(singleton(topic1), subscriptions.subscription());
+
+        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
+
+        // Instrument the test so that metadata will contain two topics after next refresh.
+        client.prepareMetadataUpdate(metadataResponse);
+
+        client.prepareResponse(joinGroupFollowerResponse(1, consumerId, "leader", Errors.NONE));
+        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));
 
-        coordinator.poll(time.milliseconds(), Long.MAX_VALUE);
+        partitionAssignor.prepare(singletonMap(consumerId, singletonList(t1p)));
+
+        // This will trigger rebalance.
+        coordinator.poll(time.timer(Long.MAX_VALUE));
 
-        assertFalse(coordinator.needRejoin());
+        // Make sure that the metadata was refreshed during the rebalance and thus subscriptions now contain two topics.
+        final Set updatedSubscriptionSet = new HashSet<>(Arrays.asList(topic1, topic2));
         assertEquals(updatedSubscriptionSet, subscriptions.subscription());
-        assertEquals(newAssignmentSet, subscriptions.assignedPartitions());
-        assertEquals(2, rebalanceListener.revokedCount);
-        assertEquals(singleton(t1p), rebalanceListener.revoked);
+
+        // Refresh the metadata again. Since there have been no changes since the last refresh, it won't trigger
+        // rebalance again.
+        metadata.requestUpdate();
+        consumerClient.poll(time.timer(Long.MAX_VALUE));
+        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(RequestTestUtils.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(RequestTestUtils.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(RequestTestUtils.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(newAssignmentSet, rebalanceListener.assigned);
+        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 = singletonList(t1p);
 
         subscriptions.subscribe(singleton(topic1), rebalanceListener);
 
         // ensure metadata is up-to-date for leader
-        metadata.setTopics(singletonList(topic1));
-        metadata.update(cluster, Collections.emptySet(), time.milliseconds());
+        client.updateMetadata(metadataResponse);
 
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
-        Map> memberSubscriptions = Collections.singletonMap(consumerId, singletonList(topic1));
-        partitionAssignor.prepare(Collections.singletonMap(consumerId, singletonList(t1p)));
+        Map> memberSubscriptions = singletonMap(consumerId, singletonList(topic1));
+        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));
         consumerClient.wakeup();
 
         try {
-            coordinator.poll(time.milliseconds(), Long.MAX_VALUE);
+            coordinator.poll(time.timer(Long.MAX_VALUE));
         } catch (WakeupException e) {
             // ignore
         }
 
         // now complete the second half
-        client.prepareResponse(syncGroupResponse(singletonList(t1p), Errors.NONE));
-        coordinator.poll(time.milliseconds(), Long.MAX_VALUE);
+        client.prepareResponse(syncGroupResponse(assigned, Errors.NONE));
+        coordinator.poll(time.timer(Long.MAX_VALUE));
 
-        assertFalse(coordinator.needRejoin());
-        assertEquals(singleton(t1p), subscriptions.assignedPartitions());
-        assertEquals(1, rebalanceListener.revokedCount);
-        assertEquals(Collections.emptySet(), rebalanceListener.revoked);
+        assertFalse(coordinator.rejoinNeededOrPending());
+        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 = singletonList(t1p);
 
-        subscriptions.subscribe(singleton(topic1), rebalanceListener);
+        subscriptions.subscribe(subscription, rebalanceListener);
 
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        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(toSet(assigned), subscriptions.assignedPartitions());
+        assertEquals(subscription, subscriptions.metadataTopics());
+        assertEquals(0, rebalanceListener.revokedCount);
+        assertNull(rebalanceListener.revoked);
+        assertEquals(1, rebalanceListener.assignedCount);
+        assertEquals(getAdded(owned, assigned), rebalanceListener.assigned);
+    }
 
-        coordinator.joinGroupIfNeeded();
+    @Test
+    public void testUpdateLastHeartbeatPollWhenCoordinatorUnknown() throws Exception {
+        // If we are part of an active group and we cannot find the coordinator, we should nevertheless
+        // continue to update the last poll time so that we do not expire the consumer
+        subscriptions.subscribe(singleton(topic1), rebalanceListener);
 
-        assertFalse(coordinator.needRejoin());
-        assertEquals(singleton(t1p), subscriptions.assignedPartitions());
-        assertEquals(singleton(topic1), subscriptions.groupSubscription());
-        assertEquals(1, rebalanceListener.revokedCount);
-        assertEquals(Collections.emptySet(), rebalanceListener.revoked);
-        assertEquals(1, rebalanceListener.assignedCount);
-        assertEquals(singleton(t1p), rebalanceListener.assigned);
+        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
+
+        // Join the group, but signal a coordinator change after the first heartbeat
+        client.prepareResponse(joinGroupFollowerResponse(1, consumerId, "leader", Errors.NONE));
+        client.prepareResponse(syncGroupResponse(singletonList(t1p), Errors.NONE));
+        client.prepareResponse(heartbeatResponse(Errors.NOT_COORDINATOR));
+
+        coordinator.poll(time.timer(Long.MAX_VALUE));
+        time.sleep(heartbeatIntervalMs);
+
+        // Await the first heartbeat which forces us to find a new coordinator
+        TestUtils.waitForCondition(() -> !client.hasPendingResponses(),
+                "Failed to observe expected heartbeat from background thread");
+
+        assertTrue(coordinator.coordinatorUnknown());
+        assertFalse(coordinator.poll(time.timer(0)));
+        assertEquals(time.milliseconds(), coordinator.heartbeat().lastPollTime());
+
+        time.sleep(rebalanceTimeoutMs - 1);
+        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);
 
         // partially update the metadata with one topic first,
         // let the leader to refresh metadata during assignment
-        metadata.setTopics(singletonList(topic1));
-        metadata.update(TestUtils.singletonCluster(topic1, 1), Collections.emptySet(), time.milliseconds());
+        client.updateMetadata(RequestTestUtils.metadataUpdateWith(1, singletonMap(topic1, 1)));
 
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        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(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(cluster, Collections.emptySet());
+        client.prepareMetadataUpdate(metadataResponse);
 
-        coordinator.joinGroupIfNeeded();
+        coordinator.joinGroupIfNeeded(time.timer(Long.MAX_VALUE));
 
-        assertFalse(coordinator.needRejoin());
-        assertEquals(2, subscriptions.assignedPartitions().size());
-        assertEquals(2, subscriptions.subscription().size());
-        assertEquals(1, rebalanceListener.revokedCount);
+        assertFalse(coordinator.rejoinNeededOrPending());
+        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
     public void testLeaveGroupOnClose() {
-        final String consumerId = "consumer";
 
         subscriptions.subscribe(singleton(topic1), rebalanceListener);
-
-        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
-
-        client.prepareResponse(joinGroupFollowerResponse(1, consumerId, "leader", Errors.NONE));
-        client.prepareResponse(syncGroupResponse(singletonList(t1p), Errors.NONE));
-        coordinator.joinGroupIfNeeded();
+        joinAsFollowerAndReceiveAssignment(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.memberId().equals(consumerId) &&
-                        leaveRequest.groupId().equals(groupId);
-            }
-        }, new LeaveGroupResponse(Errors.NONE));
-        coordinator.close(0);
+        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());
     }
 
     @Test
     public void testMaybeLeaveGroup() {
-        final String consumerId = "consumer";
+        subscriptions.subscribe(singleton(topic1), rebalanceListener);
+        joinAsFollowerAndReceiveAssignment(coordinator, singletonList(t1p));
+
+        final AtomicBoolean received = new AtomicBoolean(false);
+        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("test maybe leave group");
+        assertTrue(received.get());
 
+        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})
+     * can still execute a leave group request. Such a situation may arise when a consumer has initiated a JoinGroup
+     * request without a memberId, but is shutdown or restarted before it has a chance to initiate and complete the
+     * second request.
+     */
+    @Test
+    public void testPendingMemberShouldLeaveGroup() {
+        final String consumerId = "consumer-id";
         subscriptions.subscribe(singleton(topic1), rebalanceListener);
 
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
-        client.prepareResponse(joinGroupFollowerResponse(1, consumerId, "leader", Errors.NONE));
-        client.prepareResponse(syncGroupResponse(singletonList(t1p), Errors.NONE));
-        coordinator.joinGroupIfNeeded();
+        // here we return a DEFAULT_GENERATION_ID, but valid member id and leader id.
+        client.prepareResponse(joinGroupFollowerResponse(-1, consumerId, "leader-id", Errors.MEMBER_ID_REQUIRED));
+
+        // execute join group
+        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.memberId().equals(consumerId) &&
-                        leaveRequest.groupId().equals(groupId);
-            }
-        }, new LeaveGroupResponse(Errors.NONE));
-        coordinator.maybeLeaveGroup();
-        assertTrue(received.get());
+        client.prepareResponse(body -> {
+            received.set(true);
+            LeaveGroupRequest leaveRequest = (LeaveGroupRequest) body;
+            return validateLeaveGroup(groupId, consumerId, leaveRequest);
+        }, new LeaveGroupResponse(new LeaveGroupResponseData().setErrorCode(Errors.NONE.code())));
 
-        AbstractCoordinator.Generation generation = coordinator.generation();
-        assertNull(generation);
+        coordinator.maybeLeaveGroup("pending member leaves");
+        assertTrue(received.get());
     }
 
     @Test(expected = KafkaException.class)
     public void testUnexpectedErrorOnSyncGroup() {
-        final String consumerId = "consumer";
-
         subscriptions.subscribe(singleton(topic1), rebalanceListener);
 
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         // join initially, but let coordinator rebalance on sync
         client.prepareResponse(joinGroupFollowerResponse(1, consumerId, "leader", Errors.NONE));
-        client.prepareResponse(syncGroupResponse(Collections.emptyList(), Errors.UNKNOWN_SERVER_ERROR));
-        coordinator.joinGroupIfNeeded();
+        client.prepareResponse(syncGroupResponse(Collections.emptyList(), Errors.UNKNOWN_SERVER_ERROR));
+        coordinator.joinGroupIfNeeded(time.timer(Long.MAX_VALUE));
     }
 
     @Test
     public void testUnknownMemberIdOnSyncGroup() {
-        final String consumerId = "consumer";
-
         subscriptions.subscribe(singleton(topic1), rebalanceListener);
 
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         // 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));
 
-        coordinator.joinGroupIfNeeded();
+        coordinator.joinGroupIfNeeded(time.timer(Long.MAX_VALUE));
 
-        assertFalse(coordinator.needRejoin());
+        assertFalse(coordinator.rejoinNeededOrPending());
         assertEquals(singleton(t1p), subscriptions.assignedPartitions());
     }
 
     @Test
     public void testRebalanceInProgressOnSyncGroup() {
-        final String consumerId = "consumer";
-
         subscriptions.subscribe(singleton(topic1), rebalanceListener);
 
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         // 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));
         client.prepareResponse(syncGroupResponse(singletonList(t1p), Errors.NONE));
 
-        coordinator.joinGroupIfNeeded();
+        coordinator.joinGroupIfNeeded(time.timer(Long.MAX_VALUE));
 
-        assertFalse(coordinator.needRejoin());
+        assertFalse(coordinator.rejoinNeededOrPending());
         assertEquals(singleton(t1p), subscriptions.assignedPartitions());
     }
 
     @Test
     public void testIllegalGenerationOnSyncGroup() {
-        final String consumerId = "consumer";
-
         subscriptions.subscribe(singleton(topic1), rebalanceListener);
 
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         // 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));
 
-        coordinator.joinGroupIfNeeded();
+        coordinator.joinGroupIfNeeded(time.timer(Long.MAX_VALUE));
 
-        assertFalse(coordinator.needRejoin());
+        assertFalse(coordinator.rejoinNeededOrPending());
         assertEquals(singleton(t1p), subscriptions.assignedPartitions());
     }
 
     @Test
     public void testMetadataChangeTriggersRebalance() {
-        final String consumerId = "consumer";
 
         // ensure metadata is up-to-date for leader
-        metadata.setTopics(singletonList(topic1));
-        metadata.update(cluster, Collections.emptySet(), time.milliseconds());
-
         subscriptions.subscribe(singleton(topic1), rebalanceListener);
+        client.updateMetadata(metadataResponse);
 
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
-        Map> memberSubscriptions = Collections.singletonMap(consumerId, singletonList(topic1));
-        partitionAssignor.prepare(Collections.singletonMap(consumerId, singletonList(t1p)));
+        Map> memberSubscriptions = singletonMap(consumerId, singletonList(topic1));
+        partitionAssignor.prepare(singletonMap(consumerId, singletonList(t1p)));
 
         // the leader is responsible for picking up metadata changes and forcing a group rebalance
         client.prepareResponse(joinGroupLeaderResponse(1, consumerId, memberSubscriptions, Errors.NONE));
         client.prepareResponse(syncGroupResponse(singletonList(t1p), Errors.NONE));
 
-        coordinator.poll(time.milliseconds(), Long.MAX_VALUE);
+        coordinator.poll(time.timer(Long.MAX_VALUE));
 
-        assertFalse(coordinator.needRejoin());
+        assertFalse(coordinator.rejoinNeededOrPending());
 
         // a new partition is added to the topic
-        metadata.update(TestUtils.singletonCluster(topic1, 2), Collections.emptySet(), time.milliseconds());
+        metadata.updateWithCurrentRequestVersion(RequestTestUtils.metadataUpdateWith(1, singletonMap(topic1, 2)), false, time.milliseconds());
+        coordinator.maybeUpdateSubscriptionMetadata();
 
         // we should detect the change and ask for reassignment
-        assertTrue(coordinator.needRejoin());
+        assertTrue(coordinator.rejoinNeededOrPending());
     }
 
     @Test
@@ -750,177 +1396,269 @@ public void testUpdateMetadataDuringRebalance() {
         List topics = Arrays.asList(topic1, topic2);
 
         subscriptions.subscribe(new HashSet<>(topics), rebalanceListener);
-        metadata.setTopics(topics);
 
         // we only have metadata for one topic initially
-        metadata.update(TestUtils.singletonCluster(topic1, 1), Collections.emptySet(), time.milliseconds());
+        client.updateMetadata(RequestTestUtils.metadataUpdateWith(1, singletonMap(topic1, 1)));
 
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         // prepare initial rebalance
-        Map> memberSubscriptions = Collections.singletonMap(consumerId, topics);
-        partitionAssignor.prepare(Collections.singletonMap(consumerId, Collections.singletonList(tp1)));
+        Map> memberSubscriptions = singletonMap(consumerId, topics);
+        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);
-                    metadata.update(TestUtils.singletonCluster(topicPartitionCounts), Collections.emptySet(), time.milliseconds());
-                    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(RequestTestUtils.metadataUpdateWith(1, topicPartitionCounts));
+                return true;
             }
+            return false;
         }, syncGroupResponse(Collections.singletonList(tp1), Errors.NONE));
+        coordinator.poll(time.timer(Long.MAX_VALUE));
 
         // the metadata update should trigger a second rebalance
         client.prepareResponse(joinGroupLeaderResponse(2, consumerId, memberSubscriptions, Errors.NONE));
         client.prepareResponse(syncGroupResponse(Arrays.asList(tp1, tp2), Errors.NONE));
 
-        coordinator.poll(time.milliseconds(), Long.MAX_VALUE);
+        coordinator.poll(time.timer(Long.MAX_VALUE));
 
-        assertFalse(coordinator.needRejoin());
+        assertFalse(coordinator.rejoinNeededOrPending());
         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 testRebalanceAfterTopicUnavailableWithSubscribe() {
-        unavailableTopicTest(false, false, Collections.emptySet());
-    }
-
-    @Test
-    public void testRebalanceAfterTopicUnavailableWithPatternSubscribe() {
-        unavailableTopicTest(true, false, Collections.emptySet());
-    }
-
-    @Test
-    public void testRebalanceAfterNotMatchingTopicUnavailableWithPatternSSubscribe() {
-        unavailableTopicTest(true, false, Collections.singleton("notmatching"));
-    }
+    public void testSubscriptionChangeWithAuthorizationFailure() {
+        // 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(RequestTestUtils.metadataUpdateWith("kafka-cluster", 1,
+                Collections.singletonMap(topic2, Errors.TOPIC_AUTHORIZATION_FAILED), singletonMap(topic1, 1)));
+        assertThrows(TopicAuthorizationException.class, () -> coordinator.poll(time.timer(Long.MAX_VALUE)));
 
-    @Test
-    public void testAssignWithTopicUnavailable() {
-        unavailableTopicTest(true, false, Collections.emptySet());
-    }
+        client.respond(groupCoordinatorResponse(node, Errors.NONE));
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
-    private void unavailableTopicTest(boolean patternSubscribe, boolean assign, Set unavailableTopicsInLastMetadata) {
-        final String consumerId = "consumer";
+        // 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)));
 
-        metadata.setTopics(singletonList(topic1));
-        client.prepareMetadataUpdate(Cluster.empty(), Collections.singleton("test1"));
+        // 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(RequestTestUtils.metadataUpdateWith("kafka-cluster", 1,
+                Collections.emptyMap(), singletonMap(topic1, 1)));
 
-        if (assign)
-            subscriptions.assignFromUser(singleton(t1p));
-        else if (patternSubscribe)
+        Map> memberSubscriptions = singletonMap(consumerId, singletonList(topic1));
+        partitionAssignor.prepare(singletonMap(consumerId, singletonList(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";
+        Set topics = Collections.singleton(topic);
+        MockRebalanceListener rebalanceListener = new MockRebalanceListener() {
+            @Override
+            public void onPartitionsAssigned(Collection partitions) {
+                boolean raiseWakeup = this.assignedCount == 0;
+                super.onPartitionsAssigned(partitions);
+
+                if (raiseWakeup)
+                    throw new WakeupException();
+            }
+        };
+
+        subscriptions.subscribe(topics, rebalanceListener);
+
+        // we only have metadata for one topic initially
+        client.updateMetadata(RequestTestUtils.metadataUpdateWith(1, singletonMap(topic1, 1)));
+
+        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
+
+        // prepare initial rebalance
+        partitionAssignor.prepare(singletonMap(consumerId, Collections.singletonList(partition)));
+
+        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));
+            fail("Expected exception thrown from assignment callback");
+        } catch (WakeupException e) {
+        }
+
+        // The second call should retry the assignment callback and succeed
+        coordinator.poll(time.timer(Long.MAX_VALUE));
+
+        assertFalse(coordinator.rejoinNeededOrPending());
+        assertEquals(0, rebalanceListener.revokedCount);
+        assertEquals(2, rebalanceListener.assignedCount);
+    }
+
+    @Test
+    public void testRebalanceAfterTopicUnavailableWithSubscribe() {
+        unavailableTopicTest(false, Collections.emptySet());
+    }
+
+    @Test
+    public void testRebalanceAfterTopicUnavailableWithPatternSubscribe() {
+        unavailableTopicTest(true, Collections.emptySet());
+    }
+
+    @Test
+    public void testRebalanceAfterNotMatchingTopicUnavailableWithPatternSubscribe() {
+        unavailableTopicTest(true, Collections.singleton("notmatching"));
+    }
+
+    private void unavailableTopicTest(boolean patternSubscribe, Set unavailableTopicsInLastMetadata) {
+        if (patternSubscribe)
             subscriptions.subscribe(Pattern.compile("test.*"), rebalanceListener);
         else
             subscriptions.subscribe(singleton(topic1), rebalanceListener);
 
+        client.prepareMetadataUpdate(RequestTestUtils.metadataUpdateWith("kafka-cluster", 1,
+                Collections.singletonMap(topic1, Errors.UNKNOWN_TOPIC_OR_PARTITION), Collections.emptyMap()));
+
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
-        Map> memberSubscriptions = Collections.singletonMap(consumerId, singletonList(topic1));
-        partitionAssignor.prepare(Collections.>emptyMap());
+        Map> memberSubscriptions = singletonMap(consumerId, singletonList(topic1));
+        partitionAssignor.prepare(Collections.emptyMap());
 
         client.prepareResponse(joinGroupLeaderResponse(1, consumerId, memberSubscriptions, Errors.NONE));
-        client.prepareResponse(syncGroupResponse(Collections.emptyList(), Errors.NONE));
-        coordinator.poll(time.milliseconds(), Long.MAX_VALUE);
-        if (!assign) {
-            assertFalse(coordinator.needRejoin());
-            assertEquals(Collections.emptySet(), rebalanceListener.assigned);
-        }
+        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());
 
-        client.prepareMetadataUpdate(cluster, unavailableTopicsInLastMetadata);
-        client.poll(0, time.milliseconds());
+        Map topicErrors = new HashMap<>();
+        for (String topic : unavailableTopicsInLastMetadata)
+            topicErrors.put(topic, Errors.UNKNOWN_TOPIC_OR_PARTITION);
+
+        client.prepareMetadataUpdate(RequestTestUtils.metadataUpdateWith("kafka-cluster", 1,
+                topicErrors, singletonMap(topic1, 1)));
+
+        consumerClient.poll(time.timer(0));
         client.prepareResponse(joinGroupLeaderResponse(2, consumerId, memberSubscriptions, Errors.NONE));
         client.prepareResponse(syncGroupResponse(singletonList(t1p), Errors.NONE));
-        coordinator.poll(time.milliseconds(), Long.MAX_VALUE);
+        coordinator.poll(time.timer(Long.MAX_VALUE));
 
         assertFalse("Metadata refresh requested unnecessarily", metadata.updateRequested());
-        if (!assign) {
-            assertFalse(coordinator.needRejoin());
-            assertEquals(singleton(t1p), rebalanceListener.assigned);
-        }
+        assertFalse(coordinator.rejoinNeededOrPending());
+        assertEquals(singleton(t1p), rebalanceListener.assigned);
     }
 
     @Test
     public void testExcludeInternalTopicsConfigOption() {
-        subscriptions.subscribe(Pattern.compile(".*"), rebalanceListener);
-
-        metadata.update(TestUtils.singletonCluster(TestUtils.GROUP_METADATA_TOPIC_NAME, 2), Collections.emptySet(), time.milliseconds());
-
-        assertFalse(subscriptions.subscription().contains(TestUtils.GROUP_METADATA_TOPIC_NAME));
+        testInternalTopicInclusion(false);
     }
 
     @Test
     public void testIncludeInternalTopicsConfigOption() {
-        coordinator = buildCoordinator(new Metrics(), assignors, false, false, true);
-        subscriptions.subscribe(Pattern.compile(".*"), rebalanceListener);
-
-        metadata.update(TestUtils.singletonCluster(TestUtils.GROUP_METADATA_TOPIC_NAME, 2), Collections.emptySet(), time.milliseconds());
-
-        assertTrue(subscriptions.subscription().contains(TestUtils.GROUP_METADATA_TOPIC_NAME));
+        testInternalTopicInclusion(true);
+    }
+
+    private void testInternalTopicInclusion(boolean includeInternalTopics) {
+        metadata = new ConsumerMetadata(0, Long.MAX_VALUE, includeInternalTopics,
+                false, subscriptions, new LogContext(), new ClusterResourceListeners());
+        client = new MockClient(time, metadata);
+        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, new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, 0),
+                        Optional.of(node.id()), Optional.empty(), singletonList(node.id()), singletonList(node.id()),
+                        singletonList(node.id()));
+            MetadataResponse.TopicMetadata topicMetadata = new MetadataResponse.TopicMetadata(Errors.NONE,
+                Topic.GROUP_METADATA_TOPIC_NAME, true, singletonList(partitionMetadata));
+
+            client.updateMetadata(RequestTestUtils.metadataResponse(singletonList(node), "clusterId", node.id(),
+                singletonList(topicMetadata)));
+            coordinator.maybeUpdateSubscriptionMetadata();
+
+            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);
 
-        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
-
         // join the group once
-        client.prepareResponse(joinGroupFollowerResponse(1, "consumer", "leader", Errors.NONE));
-        client.prepareResponse(syncGroupResponse(singletonList(t1p), Errors.NONE));
-        coordinator.joinGroupIfNeeded();
+        joinAsFollowerAndReceiveAssignment(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));
-        coordinator.joinGroupIfNeeded();
+        client.prepareResponse(joinGroupFollowerResponse(2, consumerId, "leader", 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();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         // disconnected from original coordinator will cause re-discover and join again
-        client.prepareResponse(joinGroupFollowerResponse(1, "consumer", "leader", Errors.NONE), true);
+        client.prepareResponse(joinGroupFollowerResponse(1, consumerId, "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));
-        coordinator.joinGroupIfNeeded();
-
-        assertFalse(coordinator.needRejoin());
-        assertEquals(singleton(t1p), subscriptions.assignedPartitions());
-        assertEquals(1, rebalanceListener.revokedCount);
+        client.prepareResponse(joinGroupFollowerResponse(1, consumerId, "leader", Errors.NONE));
+        client.prepareResponse(syncGroupResponse(assigned, Errors.NONE));
+        coordinator.joinGroupIfNeeded(time.timer(Long.MAX_VALUE));
+
+        assertFalse(coordinator.rejoinNeededOrPending());
+        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)
@@ -928,11 +1666,11 @@ public void testInvalidSessionTimeout() {
         subscriptions.subscribe(singleton(topic1), rebalanceListener);
 
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         // coordinator doesn't like the session timeout
-        client.prepareResponse(joinGroupFollowerResponse(0, "consumer", "", Errors.INVALID_SESSION_TIMEOUT));
-        coordinator.joinGroupIfNeeded();
+        client.prepareResponse(joinGroupFollowerResponse(0, consumerId, "", Errors.INVALID_SESSION_TIMEOUT));
+        coordinator.joinGroupIfNeeded(time.timer(Long.MAX_VALUE));
     }
 
     @Test
@@ -940,16 +1678,14 @@ public void testCommitOffsetOnly() {
         subscriptions.assignFromUser(singleton(t1p));
 
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
-        client.prepareResponse(offsetCommitResponse(Collections.singletonMap(t1p, Errors.NONE)));
+        prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.NONE);
 
         AtomicBoolean success = new AtomicBoolean(false);
-        coordinator.commitOffsetsAsync(Collections.singletonMap(t1p, new OffsetAndMetadata(100L)), callback(success));
+        coordinator.commitOffsetsAsync(singletonMap(t1p, new OffsetAndMetadata(100L)), callback(success));
         coordinator.invokeCompletedOffsetCommitCallbacks();
         assertTrue(success.get());
-
-        assertEquals(100L, subscriptions.committed(t1p).offset());
     }
 
     @Test
@@ -964,18 +1700,19 @@ public void testCoordinatorDisconnectAfterCoordinatorNotAvailableError() {
 
     private void testInFlightRequestsFailedAfterCoordinatorMarkedDead(Errors error) {
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         // Send two async commits and fail the first one with an error.
         // This should cause a coordinator disconnect which will cancel the second request.
 
         MockCommitCallback firstCommitCallback = new MockCommitCallback();
         MockCommitCallback secondCommitCallback = new MockCommitCallback();
-        coordinator.commitOffsetsAsync(Collections.singletonMap(t1p, new OffsetAndMetadata(100L)), firstCommitCallback);
-        coordinator.commitOffsetsAsync(Collections.singletonMap(t1p, new OffsetAndMetadata(100L)), secondCommitCallback);
+        coordinator.commitOffsetsAsync(singletonMap(t1p, new OffsetAndMetadata(100L)), firstCommitCallback);
+        coordinator.commitOffsetsAsync(singletonMap(t1p, new OffsetAndMetadata(100L)), secondCommitCallback);
 
-        client.respond(offsetCommitResponse(Collections.singletonMap(t1p, error)));
+        respondToOffsetCommitRequest(singletonMap(t1p, 100L), error);
         consumerClient.pollNoWakeup();
+        consumerClient.pollNoWakeup(); // second poll since coordinator disconnect is async
         coordinator.invokeCompletedOffsetCommitCallbacks();
 
         assertTrue(coordinator.coordinatorUnknown());
@@ -985,101 +1722,148 @@ private void testInFlightRequestsFailedAfterCoordinatorMarkedDead(Errors error)
 
     @Test
     public void testAutoCommitDynamicAssignment() {
-        final String consumerId = "consumer";
+        try (ConsumerCoordinator coordinator = buildCoordinator(rebalanceConfig, new Metrics(), assignors, true)
+        ) {
+            subscriptions.subscribe(singleton(topic1), rebalanceListener);
+            joinAsFollowerAndReceiveAssignment(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());
+        }
+    }
 
-        ConsumerCoordinator coordinator = buildCoordinator(new Metrics(), assignors,
-                                                           ConsumerConfig.DEFAULT_EXCLUDE_INTERNAL_TOPICS, true, true);
+    @Test
+    public void testAutoCommitRetryBackoff() {
+        try (ConsumerCoordinator coordinator = buildCoordinator(rebalanceConfig, new Metrics(), assignors, true)) {
+            subscriptions.subscribe(singleton(topic1), rebalanceListener);
+            joinAsFollowerAndReceiveAssignment(coordinator, singletonList(t1p));
 
-        subscriptions.subscribe(singleton(topic1), rebalanceListener);
+            subscriptions.seek(t1p, 100);
+            time.sleep(autoCommitIntervalMs);
 
-        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+            // 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());
 
-        client.prepareResponse(joinGroupFollowerResponse(1, consumerId, "leader", Errors.NONE));
-        client.prepareResponse(syncGroupResponse(singletonList(t1p), Errors.NONE));
-        coordinator.joinGroupIfNeeded();
+            // After the disconnect, we should rediscover the coordinator
+            client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+            coordinator.poll(time.timer(Long.MAX_VALUE));
 
-        subscriptions.seek(t1p, 100);
+            subscriptions.seek(t1p, 200);
 
-        client.prepareResponse(offsetCommitResponse(Collections.singletonMap(t1p, Errors.NONE)));
-        time.sleep(autoCommitIntervalMs);
-        coordinator.poll(time.milliseconds(), 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());
 
-        assertEquals(100L, subscriptions.committed(t1p).offset());
+            // 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 testAutoCommitDynamicAssignmentRebalance() {
-        final String consumerId = "consumer";
+    public void testAutoCommitAwaitsInterval() {
+        try (ConsumerCoordinator coordinator = buildCoordinator(rebalanceConfig, new Metrics(), assignors, true)) {
+            subscriptions.subscribe(singleton(topic1), rebalanceListener);
+            joinAsFollowerAndReceiveAssignment(coordinator, singletonList(t1p));
 
-        ConsumerCoordinator coordinator = buildCoordinator(new Metrics(), assignors,
-                                                           ConsumerConfig.DEFAULT_EXCLUDE_INTERNAL_TOPICS, true, true);
+            subscriptions.seek(t1p, 100);
+            time.sleep(autoCommitIntervalMs);
 
-        subscriptions.subscribe(singleton(topic1), rebalanceListener);
+            // Send the offset commit request, but do not respond
+            coordinator.poll(time.timer(Long.MAX_VALUE));
+            assertEquals(1, client.inFlightRequestCount());
 
-        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+            time.sleep(autoCommitIntervalMs / 2);
 
-        // haven't joined, so should not cause a commit
-        time.sleep(autoCommitIntervalMs);
-        consumerClient.poll(0);
+            // Ensure that no additional offset commit is sent
+            coordinator.poll(time.timer(Long.MAX_VALUE));
+            assertEquals(1, client.inFlightRequestCount());
 
-        client.prepareResponse(joinGroupFollowerResponse(1, consumerId, "leader", Errors.NONE));
-        client.prepareResponse(syncGroupResponse(singletonList(t1p), Errors.NONE));
-        coordinator.joinGroupIfNeeded();
+            respondToOffsetCommitRequest(singletonMap(t1p, 100L), Errors.NONE);
+            coordinator.poll(time.timer(Long.MAX_VALUE));
+            assertEquals(0, client.inFlightRequestCount());
 
-        subscriptions.seek(t1p, 100);
+            subscriptions.seek(t1p, 200);
 
-        client.prepareResponse(offsetCommitResponse(Collections.singletonMap(t1p, Errors.NONE)));
-        time.sleep(autoCommitIntervalMs);
-        coordinator.poll(time.milliseconds(), Long.MAX_VALUE);
+            // 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());
 
-        assertEquals(100L, subscriptions.committed(t1p).offset());
+            // 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 testAutoCommitManualAssignment() {
-        ConsumerCoordinator coordinator = buildCoordinator(new Metrics(), assignors,
-                                                           ConsumerConfig.DEFAULT_EXCLUDE_INTERNAL_TOPICS, true, true);
+    public void testAutoCommitDynamicAssignmentRebalance() {
+        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));
 
-        subscriptions.assignFromUser(singleton(t1p));
-        subscriptions.seek(t1p, 100);
+            // haven't joined, so should not cause a commit
+            time.sleep(autoCommitIntervalMs);
+            consumerClient.poll(time.timer(0));
 
-        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+            client.prepareResponse(joinGroupFollowerResponse(1, consumerId, "leader", Errors.NONE));
+            client.prepareResponse(syncGroupResponse(singletonList(t1p), Errors.NONE));
+            coordinator.joinGroupIfNeeded(time.timer(Long.MAX_VALUE));
 
-        client.prepareResponse(offsetCommitResponse(Collections.singletonMap(t1p, Errors.NONE)));
-        time.sleep(autoCommitIntervalMs);
-        coordinator.poll(time.milliseconds(), Long.MAX_VALUE);
+            subscriptions.seek(t1p, 100);
 
-        assertEquals(100L, subscriptions.committed(t1p).offset());
+            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,
-                                                           ConsumerConfig.DEFAULT_EXCLUDE_INTERNAL_TOPICS, true, true);
-
-        subscriptions.assignFromUser(singleton(t1p));
-        subscriptions.seek(t1p, 100);
-
-        // no commit initially since coordinator is unknown
-        consumerClient.poll(0);
-        time.sleep(autoCommitIntervalMs);
-        consumerClient.poll(0);
-
-        assertNull(subscriptions.committed(t1p));
-
-        // now find the coordinator
-        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
-
-        // sleep only for the retry backoff
-        time.sleep(retryBackoffMs);
-        client.prepareResponse(offsetCommitResponse(Collections.singletonMap(t1p, Errors.NONE)));
-        coordinator.poll(time.milliseconds(), Long.MAX_VALUE);
+    public void testAutoCommitManualAssignment() {
+        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());
+        }
+    }
 
-        assertEquals(100L, subscriptions.committed(t1p).offset());
+    @Test
+    public void testAutoCommitManualAssignmentCoordinatorUnknown() {
+        try (ConsumerCoordinator coordinator = buildCoordinator(rebalanceConfig, new Metrics(), assignors, true)) {
+            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));
+
+            // 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());
+        }
     }
 
     @Test
@@ -1087,26 +1871,25 @@ public void testCommitOffsetMetadata() {
         subscriptions.assignFromUser(singleton(t1p));
 
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
-        client.prepareResponse(offsetCommitResponse(Collections.singletonMap(t1p, Errors.NONE)));
+        prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.NONE);
 
         AtomicBoolean success = new AtomicBoolean(false);
-        coordinator.commitOffsetsAsync(Collections.singletonMap(t1p, new OffsetAndMetadata(100L, "hello")), callback(success));
+
+        Map offsets = singletonMap(t1p, new OffsetAndMetadata(100L, "hello"));
+        coordinator.commitOffsetsAsync(offsets, callback(offsets, success));
         coordinator.invokeCompletedOffsetCommitCallbacks();
         assertTrue(success.get());
-
-        assertEquals(100L, subscriptions.committed(t1p).offset());
-        assertEquals("hello", subscriptions.committed(t1p).metadata());
     }
 
     @Test
     public void testCommitOffsetAsyncWithDefaultCallback() {
         int invokedBeforeTest = mockOffsetCommitCallback.invoked;
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
-        client.prepareResponse(offsetCommitResponse(Collections.singletonMap(t1p, Errors.NONE)));
-        coordinator.commitOffsetsAsync(Collections.singletonMap(t1p, new OffsetAndMetadata(100L)), mockOffsetCommitCallback);
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
+        prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.NONE);
+        coordinator.commitOffsetsAsync(singletonMap(t1p, new OffsetAndMetadata(100L)), mockOffsetCommitCallback);
         coordinator.invokeCompletedOffsetCommitCallbacks();
         assertEquals(invokedBeforeTest + 1, mockOffsetCommitCallback.invoked);
         assertNull(mockOffsetCommitCallback.exception);
@@ -1117,34 +1900,24 @@ public void testCommitAfterLeaveGroup() {
         // enable auto-assignment
         subscriptions.subscribe(singleton(topic1), rebalanceListener);
 
-        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
-
-        client.prepareResponse(joinGroupFollowerResponse(1, "consumer", "leader", Errors.NONE));
-        client.prepareResponse(syncGroupResponse(singletonList(t1p), Errors.NONE));
-
-        client.prepareMetadataUpdate(cluster, Collections.emptySet());
-
-        coordinator.joinGroupIfNeeded();
+        joinAsFollowerAndReceiveAssignment(coordinator, singletonList(t1p));
 
         // now switch to manual assignment
-        client.prepareResponse(new LeaveGroupResponse(Errors.NONE));
+        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;
-            }
-        }, offsetCommitResponse(Collections.singletonMap(t1p, Errors.NONE)));
+        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);
-        coordinator.commitOffsetsAsync(Collections.singletonMap(t1p, new OffsetAndMetadata(100L)), callback(success));
+        coordinator.commitOffsetsAsync(singletonMap(t1p, new OffsetAndMetadata(100L)), callback(success));
         coordinator.invokeCompletedOffsetCommitCallbacks();
         assertTrue(success.get());
     }
@@ -1153,9 +1926,9 @@ public boolean matches(AbstractRequest body) {
     public void testCommitOffsetAsyncFailedWithDefaultCallback() {
         int invokedBeforeTest = mockOffsetCommitCallback.invoked;
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
-        client.prepareResponse(offsetCommitResponse(Collections.singletonMap(t1p, Errors.COORDINATOR_NOT_AVAILABLE)));
-        coordinator.commitOffsetsAsync(Collections.singletonMap(t1p, new OffsetAndMetadata(100L)), mockOffsetCommitCallback);
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
+        prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.COORDINATOR_NOT_AVAILABLE);
+        coordinator.commitOffsetsAsync(singletonMap(t1p, new OffsetAndMetadata(100L)), mockOffsetCommitCallback);
         coordinator.invokeCompletedOffsetCommitCallbacks();
         assertEquals(invokedBeforeTest + 1, mockOffsetCommitCallback.invoked);
         assertTrue(mockOffsetCommitCallback.exception instanceof RetriableCommitFailedException);
@@ -1164,12 +1937,12 @@ public void testCommitOffsetAsyncFailedWithDefaultCallback() {
     @Test
     public void testCommitOffsetAsyncCoordinatorNotAvailable() {
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         // async commit with coordinator not available
         MockCommitCallback cb = new MockCommitCallback();
-        client.prepareResponse(offsetCommitResponse(Collections.singletonMap(t1p, Errors.COORDINATOR_NOT_AVAILABLE)));
-        coordinator.commitOffsetsAsync(Collections.singletonMap(t1p, new OffsetAndMetadata(100L)), cb);
+        prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.COORDINATOR_NOT_AVAILABLE);
+        coordinator.commitOffsetsAsync(singletonMap(t1p, new OffsetAndMetadata(100L)), cb);
         coordinator.invokeCompletedOffsetCommitCallbacks();
 
         assertTrue(coordinator.coordinatorUnknown());
@@ -1180,12 +1953,12 @@ public void testCommitOffsetAsyncCoordinatorNotAvailable() {
     @Test
     public void testCommitOffsetAsyncNotCoordinator() {
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         // async commit with not coordinator
         MockCommitCallback cb = new MockCommitCallback();
-        client.prepareResponse(offsetCommitResponse(Collections.singletonMap(t1p, Errors.NOT_COORDINATOR)));
-        coordinator.commitOffsetsAsync(Collections.singletonMap(t1p, new OffsetAndMetadata(100L)), cb);
+        prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.COORDINATOR_NOT_AVAILABLE);
+        coordinator.commitOffsetsAsync(singletonMap(t1p, new OffsetAndMetadata(100L)), cb);
         coordinator.invokeCompletedOffsetCommitCallbacks();
 
         assertTrue(coordinator.coordinatorUnknown());
@@ -1196,12 +1969,12 @@ public void testCommitOffsetAsyncNotCoordinator() {
     @Test
     public void testCommitOffsetAsyncDisconnected() {
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         // async commit with coordinator disconnected
         MockCommitCallback cb = new MockCommitCallback();
-        client.prepareResponse(offsetCommitResponse(Collections.singletonMap(t1p, Errors.NONE)), true);
-        coordinator.commitOffsetsAsync(Collections.singletonMap(t1p, new OffsetAndMetadata(100L)), cb);
+        prepareOffsetCommitRequestDisconnect(singletonMap(t1p, 100L));
+        coordinator.commitOffsetsAsync(singletonMap(t1p, new OffsetAndMetadata(100L)), cb);
         coordinator.invokeCompletedOffsetCommitCallbacks();
 
         assertTrue(coordinator.coordinatorUnknown());
@@ -1212,49 +1985,49 @@ public void testCommitOffsetAsyncDisconnected() {
     @Test
     public void testCommitOffsetSyncNotCoordinator() {
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         // sync commit with coordinator disconnected (should connect, get metadata, and then submit the commit request)
-        client.prepareResponse(offsetCommitResponse(Collections.singletonMap(t1p, Errors.NOT_COORDINATOR)));
+        prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.NOT_COORDINATOR);
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        client.prepareResponse(offsetCommitResponse(Collections.singletonMap(t1p, Errors.NONE)));
-        coordinator.commitOffsetsSync(Collections.singletonMap(t1p, new OffsetAndMetadata(100L)), Long.MAX_VALUE);
+        prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.NONE);
+        coordinator.commitOffsetsSync(singletonMap(t1p, new OffsetAndMetadata(100L)), time.timer(Long.MAX_VALUE));
     }
 
     @Test
     public void testCommitOffsetSyncCoordinatorNotAvailable() {
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         // sync commit with coordinator disconnected (should connect, get metadata, and then submit the commit request)
-        client.prepareResponse(offsetCommitResponse(Collections.singletonMap(t1p, Errors.COORDINATOR_NOT_AVAILABLE)));
+        prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.COORDINATOR_NOT_AVAILABLE);
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        client.prepareResponse(offsetCommitResponse(Collections.singletonMap(t1p, Errors.NONE)));
-        coordinator.commitOffsetsSync(Collections.singletonMap(t1p, new OffsetAndMetadata(100L)), Long.MAX_VALUE);
+        prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.NONE);
+        coordinator.commitOffsetsSync(singletonMap(t1p, new OffsetAndMetadata(100L)), time.timer(Long.MAX_VALUE));
     }
 
     @Test
     public void testCommitOffsetSyncCoordinatorDisconnected() {
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         // sync commit with coordinator disconnected (should connect, get metadata, and then submit the commit request)
-        client.prepareResponse(offsetCommitResponse(Collections.singletonMap(t1p, Errors.NONE)), true);
+        prepareOffsetCommitRequestDisconnect(singletonMap(t1p, 100L));
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        client.prepareResponse(offsetCommitResponse(Collections.singletonMap(t1p, Errors.NONE)));
-        coordinator.commitOffsetsSync(Collections.singletonMap(t1p, new OffsetAndMetadata(100L)), Long.MAX_VALUE);
+        prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.NONE);
+        coordinator.commitOffsetsSync(singletonMap(t1p, new OffsetAndMetadata(100L)), time.timer(Long.MAX_VALUE));
     }
 
     @Test
     public void testAsyncCommitCallbacksInvokedPriorToSyncCommitCompletion() throws Exception {
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        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);
 
-        coordinator.commitOffsetsAsync(Collections.singletonMap(t1p, firstOffset), new OffsetCommitCallback() {
+        coordinator.commitOffsetsAsync(singletonMap(t1p, firstOffset), new OffsetCommitCallback() {
             @Override
             public void onComplete(Map offsets, Exception exception) {
                 committedOffsets.add(firstOffset);
@@ -1265,7 +2038,7 @@ public void onComplete(Map offsets, Exception
         Thread thread = new Thread() {
             @Override
             public void run() {
-                coordinator.commitOffsetsSync(Collections.singletonMap(t1p, secondOffset), 10000);
+                coordinator.commitOffsetsSync(singletonMap(t1p, secondOffset), time.timer(10000));
                 committedOffsets.add(secondOffset);
             }
         };
@@ -1273,318 +2046,879 @@ public void run() {
         thread.start();
 
         client.waitForRequests(2, 5000);
-        client.respond(offsetCommitResponse(Collections.singletonMap(t1p, Errors.NONE)));
-        client.respond(offsetCommitResponse(Collections.singletonMap(t1p, Errors.NONE)));
+        respondToOffsetCommitRequest(singletonMap(t1p, firstOffset.offset()), Errors.NONE);
+        respondToOffsetCommitRequest(singletonMap(t1p, secondOffset.offset()), Errors.NONE);
 
         thread.join();
 
         assertEquals(Arrays.asList(firstOffset, secondOffset), committedOffsets);
     }
 
-    @Test(expected = KafkaException.class)
-    public void testCommitUnknownTopicOrPartition() {
+    @Test
+    public void testRetryCommitUnknownTopicOrPartition() {
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
-        client.prepareResponse(offsetCommitResponse(Collections.singletonMap(t1p, Errors.UNKNOWN_TOPIC_OR_PARTITION)));
-        coordinator.commitOffsetsSync(Collections.singletonMap(t1p, new OffsetAndMetadata(100L, "metadata")), Long.MAX_VALUE);
+        client.prepareResponse(offsetCommitResponse(singletonMap(t1p, Errors.UNKNOWN_TOPIC_OR_PARTITION)));
+        client.prepareResponse(offsetCommitResponse(singletonMap(t1p, Errors.NONE)));
+
+        assertTrue(coordinator.commitOffsetsSync(singletonMap(t1p,
+                new OffsetAndMetadata(100L, "metadata")), time.timer(10000)));
     }
 
     @Test(expected = OffsetMetadataTooLarge.class)
     public void testCommitOffsetMetadataTooLarge() {
         // since offset metadata is provided by the user, we have to propagate the exception so they can handle it
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
-        client.prepareResponse(offsetCommitResponse(Collections.singletonMap(t1p, Errors.OFFSET_METADATA_TOO_LARGE)));
-        coordinator.commitOffsetsSync(Collections.singletonMap(t1p, new OffsetAndMetadata(100L, "metadata")), Long.MAX_VALUE);
+        prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.OFFSET_METADATA_TOO_LARGE);
+        coordinator.commitOffsetsSync(singletonMap(t1p,
+                new OffsetAndMetadata(100L, "metadata")), time.timer(Long.MAX_VALUE));
     }
 
     @Test(expected = CommitFailedException.class)
     public void testCommitOffsetIllegalGeneration() {
         // we cannot retry if a rebalance occurs before the commit completed
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
-        client.prepareResponse(offsetCommitResponse(Collections.singletonMap(t1p, Errors.ILLEGAL_GENERATION)));
-        coordinator.commitOffsetsSync(Collections.singletonMap(t1p, new OffsetAndMetadata(100L, "metadata")), Long.MAX_VALUE);
+        prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.ILLEGAL_GENERATION);
+        coordinator.commitOffsetsSync(singletonMap(t1p,
+                new OffsetAndMetadata(100L, "metadata")), time.timer(Long.MAX_VALUE));
     }
 
     @Test(expected = CommitFailedException.class)
     public void testCommitOffsetUnknownMemberId() {
         // we cannot retry if a rebalance occurs before the commit completed
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
-        client.prepareResponse(offsetCommitResponse(Collections.singletonMap(t1p, Errors.UNKNOWN_MEMBER_ID)));
-        coordinator.commitOffsetsSync(Collections.singletonMap(t1p, new OffsetAndMetadata(100L, "metadata")), Long.MAX_VALUE);
+        prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.UNKNOWN_MEMBER_ID);
+        coordinator.commitOffsetsSync(singletonMap(t1p,
+                new OffsetAndMetadata(100L, "metadata")), time.timer(Long.MAX_VALUE));
     }
 
-    @Test(expected = CommitFailedException.class)
+    @Test
+    public void testCommitOffsetIllegalGenerationWithNewGeneration() {
+        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
+
+        final AbstractCoordinator.Generation currGen = new AbstractCoordinator.Generation(
+            1,
+            "memberId",
+            null);
+        coordinator.setNewGeneration(currGen);
+
+        prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.ILLEGAL_GENERATION);
+        RequestFuture future = coordinator.sendOffsetCommitRequest(singletonMap(t1p,
+            new OffsetAndMetadata(100L, "metadata")));
+
+        // change the generation
+        final AbstractCoordinator.Generation newGen = new AbstractCoordinator.Generation(
+            2,
+            "memberId-new",
+            null);
+        coordinator.setNewGeneration(newGen);
+        coordinator.setNewState(AbstractCoordinator.MemberState.PREPARING_REBALANCE);
+
+        assertTrue(consumerClient.poll(future, time.timer(30000)));
+        assertTrue(future.exception().getClass().isInstance(Errors.REBALANCE_IN_PROGRESS.exception()));
+
+        // the generation should not be reset
+        assertEquals(newGen, coordinator.generation());
+    }
+
+    @Test
+    public void testCommitOffsetIllegalGenerationWithResetGenearion() {
+        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
+
+        final AbstractCoordinator.Generation currGen = new AbstractCoordinator.Generation(
+            1,
+            "memberId",
+            null);
+        coordinator.setNewGeneration(currGen);
+
+        prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.ILLEGAL_GENERATION);
+        RequestFuture future = coordinator.sendOffsetCommitRequest(singletonMap(t1p,
+            new OffsetAndMetadata(100L, "metadata")));
+
+        // reset the generation
+        coordinator.setNewGeneration(AbstractCoordinator.Generation.NO_GENERATION);
+
+        assertTrue(consumerClient.poll(future, time.timer(30000)));
+        assertTrue(future.exception().getClass().isInstance(new CommitFailedException()));
+
+        // the generation should not be reset
+        assertEquals(AbstractCoordinator.Generation.NO_GENERATION, coordinator.generation());
+    }
+
+    @Test
+    public void testCommitOffsetUnknownMemberWithNewGenearion() {
+        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
+
+        final AbstractCoordinator.Generation currGen = new AbstractCoordinator.Generation(
+            1,
+            "memberId",
+            null);
+        coordinator.setNewGeneration(currGen);
+
+        prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.UNKNOWN_MEMBER_ID);
+        RequestFuture future = coordinator.sendOffsetCommitRequest(singletonMap(t1p,
+            new OffsetAndMetadata(100L, "metadata")));
+
+        // change the generation
+        final AbstractCoordinator.Generation newGen = new AbstractCoordinator.Generation(
+            2,
+            "memberId-new",
+            null);
+        coordinator.setNewGeneration(newGen);
+        coordinator.setNewState(AbstractCoordinator.MemberState.PREPARING_REBALANCE);
+
+        assertTrue(consumerClient.poll(future, time.timer(30000)));
+        assertTrue(future.exception().getClass().isInstance(Errors.REBALANCE_IN_PROGRESS.exception()));
+
+        // the generation should not be reset
+        assertEquals(newGen, coordinator.generation());
+    }
+
+    @Test
+    public void testCommitOffsetUnknownMemberWithResetGenearion() {
+        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
+
+        final AbstractCoordinator.Generation currGen = new AbstractCoordinator.Generation(
+            1,
+            "memberId",
+            null);
+        coordinator.setNewGeneration(currGen);
+
+        prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.UNKNOWN_MEMBER_ID);
+        RequestFuture future = coordinator.sendOffsetCommitRequest(singletonMap(t1p,
+            new OffsetAndMetadata(100L, "metadata")));
+
+        // reset the generation
+        coordinator.setNewGeneration(AbstractCoordinator.Generation.NO_GENERATION);
+
+        assertTrue(consumerClient.poll(future, time.timer(30000)));
+        assertTrue(future.exception().getClass().isInstance(new CommitFailedException()));
+
+        // the generation should not be reset
+        assertEquals(AbstractCoordinator.Generation.NO_GENERATION, coordinator.generation());
+    }
+
+    @Test
+    public void testCommitOffsetFencedInstanceWithRebalancingGenearion() {
+        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
+
+        final AbstractCoordinator.Generation currGen = new AbstractCoordinator.Generation(
+            1,
+            "memberId",
+            null);
+        coordinator.setNewGeneration(currGen);
+        coordinator.setNewState(AbstractCoordinator.MemberState.PREPARING_REBALANCE);
+
+        prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.FENCED_INSTANCE_ID);
+        RequestFuture future = coordinator.sendOffsetCommitRequest(singletonMap(t1p,
+            new OffsetAndMetadata(100L, "metadata")));
+
+        // change the generation
+        final AbstractCoordinator.Generation newGen = new AbstractCoordinator.Generation(
+            2,
+            "memberId-new",
+            null);
+        coordinator.setNewGeneration(newGen);
+
+        assertTrue(consumerClient.poll(future, time.timer(30000)));
+        assertTrue(future.exception().getClass().isInstance(Errors.REBALANCE_IN_PROGRESS.exception()));
+
+        // the generation should not be reset
+        assertEquals(newGen, coordinator.generation());
+    }
+
+    @Test
+    public void testCommitOffsetFencedInstanceWithNewGenearion() {
+        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
+
+        final AbstractCoordinator.Generation currGen = new AbstractCoordinator.Generation(
+            1,
+            "memberId",
+            null);
+        coordinator.setNewGeneration(currGen);
+
+        prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.FENCED_INSTANCE_ID);
+        RequestFuture future = coordinator.sendOffsetCommitRequest(singletonMap(t1p,
+            new OffsetAndMetadata(100L, "metadata")));
+
+        // change the generation
+        final AbstractCoordinator.Generation newGen = new AbstractCoordinator.Generation(
+            2,
+            "memberId-new",
+            null);
+        coordinator.setNewGeneration(newGen);
+
+        assertTrue(consumerClient.poll(future, time.timer(30000)));
+        assertTrue(future.exception().getClass().isInstance(new CommitFailedException()));
+
+        // the generation should not be reset
+        assertEquals(newGen, coordinator.generation());
+    }
+
+    @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();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
+
+        // normal join group
+        Map> memberSubscriptions = singletonMap(consumerId, singletonList(topic1));
+        partitionAssignor.prepare(singletonMap(consumerId, singletonList(t1p)));
+
+        coordinator.ensureActiveGroup(time.timer(0L));
+
+        assertTrue(coordinator.rejoinNeededOrPending());
+        assertNull(coordinator.generationIfStable());
+
+        // when the state is REBALANCING, we would not even send out the request but fail immediately
+        assertThrows(RebalanceInProgressException.class, () -> coordinator.commitOffsetsSync(singletonMap(t1p,
+            new OffsetAndMetadata(100L, "metadata")), time.timer(Long.MAX_VALUE)));
+
+        final Node coordinatorNode = new Node(Integer.MAX_VALUE - node.id(), node.host(), node.port());
+        client.respondFrom(joinGroupLeaderResponse(1, consumerId, memberSubscriptions, Errors.NONE), coordinatorNode);
+
+        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());
 
-        client.prepareResponse(offsetCommitResponse(Collections.singletonMap(t1p, Errors.REBALANCE_IN_PROGRESS)));
-        coordinator.commitOffsetsSync(Collections.singletonMap(t1p, new OffsetAndMetadata(100L, "metadata")), Long.MAX_VALUE);
+        prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.REBALANCE_IN_PROGRESS);
+        assertThrows(RebalanceInProgressException.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)
     public void testCommitOffsetSyncCallbackWithNonRetriableException() {
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         // sync commit with invalid partitions should throw if we have no callback
-        client.prepareResponse(offsetCommitResponse(Collections.singletonMap(t1p, Errors.UNKNOWN_SERVER_ERROR)), false);
-        coordinator.commitOffsetsSync(Collections.singletonMap(t1p, new OffsetAndMetadata(100L)), Long.MAX_VALUE);
+        prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.UNKNOWN_SERVER_ERROR);
+        coordinator.commitOffsetsSync(singletonMap(t1p, new OffsetAndMetadata(100L)), time.timer(Long.MAX_VALUE));
     }
 
-    @Test(expected = IllegalArgumentException.class)
-    public void testCommitSyncNegativeOffset() {
+    @Test
+    public void testCommitOffsetSyncWithoutFutureGetsCompleted() {
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.commitOffsetsSync(Collections.singletonMap(t1p, new OffsetAndMetadata(-1L)), Long.MAX_VALUE);
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
+        assertFalse(coordinator.commitOffsetsSync(singletonMap(t1p, new OffsetAndMetadata(100L)), time.timer(0)));
     }
 
     @Test
-    public void testCommitAsyncNegativeOffset() {
-        int invokedBeforeTest = mockOffsetCommitCallback.invoked;
+    public void testRefreshOffset() {
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.commitOffsetsAsync(Collections.singletonMap(t1p, new OffsetAndMetadata(-1L)), mockOffsetCommitCallback);
-        coordinator.invokeCompletedOffsetCommitCallbacks();
-        assertEquals(invokedBeforeTest + 1, mockOffsetCommitCallback.invoked);
-        assertTrue(mockOffsetCommitCallback.exception instanceof IllegalArgumentException);
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
+
+        subscriptions.assignFromUser(singleton(t1p));
+        client.prepareResponse(offsetFetchResponse(t1p, Errors.NONE, "", 100L));
+        coordinator.refreshCommittedOffsetsIfNeeded(time.timer(Long.MAX_VALUE));
+
+        assertEquals(Collections.emptySet(), subscriptions.initializingPartitions());
+        assertTrue(subscriptions.hasAllFetchPositions());
+        assertEquals(100L, subscriptions.position(t1p).offset);
     }
 
     @Test
-    public void testCommitOffsetSyncWithoutFutureGetsCompleted() {
+    public void testRefreshOffsetWithValidation() {
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
-        assertFalse(coordinator.commitOffsetsSync(Collections.singletonMap(t1p, new OffsetAndMetadata(100L)), 0));
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
+
+        subscriptions.assignFromUser(singleton(t1p));
+
+        // Initial leader epoch of 4
+        MetadataResponse metadataResponse = RequestTestUtils.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, Optional.of(3)));
+        coordinator.refreshCommittedOffsetsIfNeeded(time.timer(Long.MAX_VALUE));
+
+        // Offset gets loaded, but requires validation
+        assertEquals(Collections.emptySet(), subscriptions.initializingPartitions());
+        assertFalse(subscriptions.hasAllFetchPositions());
+        assertTrue(subscriptions.awaitingValidation(t1p));
+        assertEquals(subscriptions.position(t1p).offset, 100L);
+        assertNull(subscriptions.validPosition(t1p));
     }
 
     @Test
-    public void testRefreshOffset() {
+    public void testFetchCommittedOffsets() {
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
-        subscriptions.assignFromUser(singleton(t1p));
-        subscriptions.needRefreshCommits();
-        client.prepareResponse(offsetFetchResponse(t1p, Errors.NONE, "", 100L));
-        coordinator.refreshCommittedOffsetsIfNeeded();
-        assertFalse(subscriptions.refreshCommitsNeeded());
-        assertEquals(100L, subscriptions.committed(t1p).offset());
+        long offset = 500L;
+        String metadata = "blahblah";
+        Optional leaderEpoch = Optional.of(15);
+        OffsetFetchResponse.PartitionData data = new OffsetFetchResponse.PartitionData(offset, leaderEpoch,
+                metadata, Errors.NONE);
+
+        client.prepareResponse(new OffsetFetchResponse(Errors.NONE, singletonMap(t1p, data)));
+        Map fetchedOffsets = coordinator.fetchCommittedOffsets(singleton(t1p),
+                time.timer(Long.MAX_VALUE));
+
+        assertNotNull(fetchedOffsets);
+        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));
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         subscriptions.assignFromUser(singleton(t1p));
-        subscriptions.needRefreshCommits();
         client.prepareResponse(offsetFetchResponse(Errors.COORDINATOR_LOAD_IN_PROGRESS));
         client.prepareResponse(offsetFetchResponse(t1p, Errors.NONE, "", 100L));
-        coordinator.refreshCommittedOffsetsIfNeeded();
-        assertFalse(subscriptions.refreshCommitsNeeded());
-        assertEquals(100L, subscriptions.committed(t1p).offset());
+        coordinator.refreshCommittedOffsetsIfNeeded(time.timer(Long.MAX_VALUE));
+
+        assertEquals(Collections.emptySet(), subscriptions.initializingPartitions());
+        assertTrue(subscriptions.hasAllFetchPositions());
+        assertEquals(100L, subscriptions.position(t1p).offset);
     }
 
     @Test
     public void testRefreshOffsetsGroupNotAuthorized() {
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         subscriptions.assignFromUser(singleton(t1p));
-        subscriptions.needRefreshCommits();
         client.prepareResponse(offsetFetchResponse(Errors.GROUP_AUTHORIZATION_FAILED));
         try {
-            coordinator.refreshCommittedOffsetsIfNeeded();
+            coordinator.refreshCommittedOffsetsIfNeeded(time.timer(Long.MAX_VALUE));
             fail("Expected group authorization error");
         } catch (GroupAuthorizationException e) {
             assertEquals(groupId, e.groupId());
         }
     }
 
+    @Test
+    public void testRefreshOffsetWithPendingTransactions() {
+        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
+
+        subscriptions.assignFromUser(singleton(t1p));
+        client.prepareResponse(offsetFetchResponse(t1p, Errors.UNSTABLE_OFFSET_COMMIT, "", -1L));
+        client.prepareResponse(offsetFetchResponse(t1p, Errors.NONE, "", 100L));
+        assertEquals(Collections.singleton(t1p), subscriptions.initializingPartitions());
+        coordinator.refreshCommittedOffsetsIfNeeded(time.timer(0L));
+        assertEquals(Collections.singleton(t1p), subscriptions.initializingPartitions());
+        coordinator.refreshCommittedOffsetsIfNeeded(time.timer(0L));
+
+        assertEquals(Collections.emptySet(), subscriptions.initializingPartitions());
+        assertTrue(subscriptions.hasAllFetchPositions());
+        assertEquals(100L, subscriptions.position(t1p).offset);
+    }
+
     @Test(expected = KafkaException.class)
     public void testRefreshOffsetUnknownTopicOrPartition() {
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         subscriptions.assignFromUser(singleton(t1p));
-        subscriptions.needRefreshCommits();
         client.prepareResponse(offsetFetchResponse(t1p, Errors.UNKNOWN_TOPIC_OR_PARTITION, "", 100L));
-        coordinator.refreshCommittedOffsetsIfNeeded();
+        coordinator.refreshCommittedOffsetsIfNeeded(time.timer(Long.MAX_VALUE));
     }
 
     @Test
     public void testRefreshOffsetNotCoordinatorForConsumer() {
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         subscriptions.assignFromUser(singleton(t1p));
-        subscriptions.needRefreshCommits();
         client.prepareResponse(offsetFetchResponse(Errors.NOT_COORDINATOR));
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
         client.prepareResponse(offsetFetchResponse(t1p, Errors.NONE, "", 100L));
-        coordinator.refreshCommittedOffsetsIfNeeded();
-        assertFalse(subscriptions.refreshCommitsNeeded());
-        assertEquals(100L, subscriptions.committed(t1p).offset());
+        coordinator.refreshCommittedOffsetsIfNeeded(time.timer(Long.MAX_VALUE));
+
+        assertEquals(Collections.emptySet(), subscriptions.initializingPartitions());
+        assertTrue(subscriptions.hasAllFetchPositions());
+        assertEquals(100L, subscriptions.position(t1p).offset);
     }
 
     @Test
     public void testRefreshOffsetWithNoFetchableOffsets() {
         client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         subscriptions.assignFromUser(singleton(t1p));
-        subscriptions.needRefreshCommits();
         client.prepareResponse(offsetFetchResponse(t1p, Errors.NONE, "", -1L));
-        coordinator.refreshCommittedOffsetsIfNeeded();
-        assertFalse(subscriptions.refreshCommitsNeeded());
-        assertEquals(null, subscriptions.committed(t1p));
+        coordinator.refreshCommittedOffsetsIfNeeded(time.timer(Long.MAX_VALUE));
+
+        assertEquals(Collections.singleton(t1p), subscriptions.initializingPartitions());
+        assertEquals(Collections.emptySet(), subscriptions.partitionsNeedingReset(time.milliseconds()));
+        assertFalse(subscriptions.hasAllFetchPositions());
+        assertNull(subscriptions.position(t1p));
     }
 
     @Test
-    public void testProtocolMetadataOrder() {
-        RoundRobinAssignor roundRobin = new RoundRobinAssignor();
-        RangeAssignor range = new RangeAssignor();
+    public void testNoCoordinatorDiscoveryIfPositionsKnown() {
+        assertTrue(coordinator.coordinatorUnknown());
 
-        try (Metrics metrics = new Metrics(time)) {
-            ConsumerCoordinator coordinator = buildCoordinator(metrics, Arrays.asList(roundRobin, range),
-                                                               ConsumerConfig.DEFAULT_EXCLUDE_INTERNAL_TOPICS, false, true);
-            List metadata = coordinator.metadata();
-            assertEquals(2, metadata.size());
-            assertEquals(roundRobin.name(), metadata.get(0).name());
-            assertEquals(range.name(), metadata.get(1).name());
-        }
+        subscriptions.assignFromUser(singleton(t1p));
+        subscriptions.seek(t1p, 500L);
+        coordinator.refreshCommittedOffsetsIfNeeded(time.timer(Long.MAX_VALUE));
+
+        assertEquals(Collections.emptySet(), subscriptions.initializingPartitions());
+        assertTrue(subscriptions.hasAllFetchPositions());
+        assertEquals(500L, subscriptions.position(t1p).offset);
+        assertTrue(coordinator.coordinatorUnknown());
+    }
+
+    @Test
+    public void testNoCoordinatorDiscoveryIfPartitionAwaitingReset() {
+        assertTrue(coordinator.coordinatorUnknown());
+
+        subscriptions.assignFromUser(singleton(t1p));
+        subscriptions.requestOffsetReset(t1p, OffsetResetStrategy.EARLIEST);
+        coordinator.refreshCommittedOffsetsIfNeeded(time.timer(Long.MAX_VALUE));
+
+        assertEquals(Collections.emptySet(), subscriptions.initializingPartitions());
+        assertFalse(subscriptions.hasAllFetchPositions());
+        assertEquals(Collections.singleton(t1p), subscriptions.partitionsNeedingReset(time.milliseconds()));
+        assertEquals(OffsetResetStrategy.EARLIEST, subscriptions.resetStrategy(t1p));
+        assertTrue(coordinator.coordinatorUnknown());
+    }
 
-        try (Metrics metrics = new Metrics(time)) {
-            ConsumerCoordinator coordinator = buildCoordinator(metrics, Arrays.asList(range, roundRobin),
-                                                               ConsumerConfig.DEFAULT_EXCLUDE_INTERNAL_TOPICS, 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 testAuthenticationFailureInEnsureActiveGroup() {
+        client.createPendingAuthenticationError(node, 300);
+
+        try {
+            coordinator.ensureActiveGroup();
+            fail("Expected an authentication error.");
+        } catch (AuthenticationException e) {
+            // OK
         }
     }
 
     @Test
-    public void testCloseDynamicAssignment() throws Exception {
-        ConsumerCoordinator coordinator = prepareCoordinatorForCloseTest(true, true, true);
-        gracefulCloseTest(coordinator, true);
+    public void testThreadSafeAssignedPartitionsMetric() throws Exception {
+        // Get the assigned-partitions metric
+        final Metric metric = metrics.metric(new MetricName("assigned-partitions", consumerId + groupId + "-coordinator-metrics",
+                "", Collections.emptyMap()));
+
+        // Start polling the metric in the background
+        final AtomicBoolean doStop = new AtomicBoolean();
+        final AtomicReference exceptionHolder = new AtomicReference<>();
+        final AtomicInteger observedSize = new AtomicInteger();
+
+        Thread poller = new Thread() {
+            @Override
+            public void run() {
+                // Poll as fast as possible to reproduce ConcurrentModificationException
+                while (!doStop.get()) {
+                    try {
+                        int size = ((Double) metric.metricValue()).intValue();
+                        observedSize.set(size);
+                    } catch (Exception e) {
+                        exceptionHolder.set(e);
+                        return;
+                    }
+                }
+            }
+        };
+        poller.start();
+
+        // Assign two partitions to trigger a metric change that can lead to ConcurrentModificationException
+        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
+
+        // Change the assignment several times to increase likelihood of concurrent updates
+        Set partitions = new HashSet<>();
+        int totalPartitions = 10;
+        for (int partition = 0; partition < totalPartitions; partition++) {
+            partitions.add(new TopicPartition(topic1, partition));
+            subscriptions.assignFromUser(partitions);
+        }
+
+        // Wait for the metric poller to observe the final assignment change or raise an error
+        TestUtils.waitForCondition(
+            () -> observedSize.get() == totalPartitions ||
+            exceptionHolder.get() != null, "Failed to observe expected assignment change");
+
+        doStop.set(true);
+        poller.join();
+
+        assertNull("Failed fetching the metric at least once", exceptionHolder.get());
     }
 
     @Test
-    public void testCloseManualAssignment() throws Exception {
-        ConsumerCoordinator coordinator = prepareCoordinatorForCloseTest(false, true, true);
-        gracefulCloseTest(coordinator, false);
+    public void testCloseDynamicAssignment() {
+        try (ConsumerCoordinator coordinator = prepareCoordinatorForCloseTest(true, true, Optional.empty())) {
+            gracefulCloseTest(coordinator, true);
+        }
     }
 
     @Test
-    public void shouldNotLeaveGroupWhenLeaveGroupFlagIsFalse() throws Exception {
-        final ConsumerCoordinator coordinator = prepareCoordinatorForCloseTest(true, true, false);
-        gracefulCloseTest(coordinator, false);
+    public void testCloseManualAssignment() {
+        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, 60000, 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, 60000, 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, 60000, 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, 60000, 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, 60000, 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, 60000, 60000, 60000);
+        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, 60000, 60000, 60000);
+        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, 60000, 60000, 60000);
+        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, 60000, 0, 0);
+        try (ConsumerCoordinator coordinator = prepareCoordinatorForCloseTest(true, true, groupInstanceId)) {
+            time.sleep(autoCommitIntervalMs);
+            closeVerifyTimeout(coordinator, 0, 0, 0);
+        }
     }
 
     @Test
     public void testHeartbeatThreadClose() throws Exception {
-        groupId = "testCloseTimeoutWithHeartbeatThread";
-        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, 60000, 60000, 60000);
-        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() {
+        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);
+        }
+    }
+
+    @Test(expected = FencedInstanceIdException.class)
+    public void testCommitOffsetRequestSyncWithFencedInstanceIdException() {
+        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
+
+        // 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 testGetGroupMetadata() {
+        final ConsumerGroupMetadata groupMetadata = coordinator.groupMetadata();
+        assertNotNull(groupMetadata);
+        assertEquals(groupId, groupMetadata.groupId());
+        assertEquals(JoinGroupRequest.UNKNOWN_GENERATION_ID, groupMetadata.generationId());
+        assertEquals(JoinGroupRequest.UNKNOWN_MEMBER_ID, groupMetadata.memberId());
+        assertFalse(groupMetadata.groupInstanceId().isPresent());
+
+        try (final ConsumerCoordinator coordinator = prepareCoordinatorForCloseTest(true, true, groupInstanceId)) {
+            coordinator.ensureActiveGroup();
+
+            final ConsumerGroupMetadata joinedGroupMetadata = coordinator.groupMetadata();
+            assertNotNull(joinedGroupMetadata);
+            assertEquals(groupId, joinedGroupMetadata.groupId());
+            assertEquals(1, joinedGroupMetadata.generationId());
+            assertEquals(consumerId, joinedGroupMetadata.memberId());
+            assertEquals(groupInstanceId, joinedGroupMetadata.groupInstanceId());
+        }
+    }
+
+    @Test
+    public void shouldUpdateConsumerGroupMetadataBeforeCallbacks() {
+        final MockRebalanceListener rebalanceListener = new MockRebalanceListener() {
+            @Override
+            public void onPartitionsRevoked(Collection partitions) {
+                assertEquals(2, coordinator.groupMetadata().generationId());
+            }
+        };
+
+        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);
+        }
+
+        ByteBuffer buffer = ConsumerProtocol.serializeAssignment(
+            new ConsumerPartitionAssignor.Assignment(Collections.emptyList(), ByteBuffer.wrap(new byte[0])));
+        coordinator.onJoinComplete(2, "memberId", partitionAssignor.name(), buffer);
+    }
+
+    @Test
+    public void testConsumerRejoinAfterRebalance() {
+        try (ConsumerCoordinator coordinator = prepareCoordinatorForCloseTest(true, false, Optional.of("group-id"))) {
+            coordinator.ensureActiveGroup();
+
+            prepareOffsetCommitRequest(singletonMap(t1p, 100L), Errors.REBALANCE_IN_PROGRESS);
+
+            assertThrows(RebalanceInProgressException.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);
+
+            // should have retried sending a join group request already
+            assertFalse(client.hasPendingResponses());
+            assertEquals(1, client.inFlightRequestCount());
+
+            System.out.println(client.requests());
+
+            // Retry join should then succeed
+            client.respond(joinGroupFollowerResponse(generationId, memberId, "leader", Errors.NONE));
+            client.prepareResponse(syncGroupResponse(singletonList(t1p), Errors.NONE));
+
+            res = coordinator.joinGroupIfNeeded(time.timer(3000));
+
+            assertTrue(res);
+            assertFalse(client.hasPendingResponses());
+            assertFalse(client.hasInFlightRequests());
+        }
+    }
+
+    @Test
+    public void testThrowOnUnsupportedStableFlag() {
+        supportStableFlag((short) 6, true);
+    }
+
+    @Test
+    public void testNoThrowWhenStableFlagIsSupported() {
+        supportStableFlag((short) 7, false);
+    }
+
+    private void supportStableFlag(final short upperVersion, final boolean expectThrows) {
+        ConsumerCoordinator coordinator = new ConsumerCoordinator(
+            rebalanceConfig,
+            new LogContext(),
+            consumerClient,
+            assignors,
+            metadata,
+            subscriptions,
+            new Metrics(time),
+            consumerId + groupId,
+            time,
+            false,
+            autoCommitIntervalMs,
+            null,
+            true);
+
+        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        client.setNodeApiVersions(NodeApiVersions.create(ApiKeys.OFFSET_FETCH.id, (short) 0, upperVersion));
+
+        long offset = 500L;
+        String metadata = "blahblah";
+        Optional leaderEpoch = Optional.of(15);
+        OffsetFetchResponse.PartitionData data = new OffsetFetchResponse.PartitionData(offset, leaderEpoch,
+            metadata, Errors.NONE);
+
+        client.prepareResponse(new OffsetFetchResponse(Errors.NONE, singletonMap(t1p, data)));
+        if (expectThrows) {
+            assertThrows(UnsupportedVersionException.class,
+                () -> coordinator.fetchCommittedOffsets(singleton(t1p), time.timer(Long.MAX_VALUE)));
+        } else {
+            Map fetchedOffsets = coordinator.fetchCommittedOffsets(singleton(t1p),
+                time.timer(Long.MAX_VALUE));
+
+            assertNotNull(fetchedOffsets);
+            assertEquals(new OffsetAndMetadata(offset, leaderEpoch, metadata), fetchedOffsets.get(t1p));
+        }
+    }
+
+    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 String consumerId = "consumer";
-        ConsumerCoordinator coordinator = buildCoordinator(new Metrics(), assignors,
-                ConsumerConfig.DEFAULT_EXCLUDE_INTERNAL_TOPICS, autoCommit, leaveGroup);
-        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+                                                               final Optional groupInstanceId) {
+        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) {
             subscriptions.subscribe(singleton(topic1), rebalanceListener);
             client.prepareResponse(joinGroupFollowerResponse(1, consumerId, "leader", Errors.NONE));
             client.prepareResponse(syncGroupResponse(singletonList(t1p), Errors.NONE));
-            coordinator.joinGroupIfNeeded();
-        } else
+            coordinator.joinGroupIfNeeded(time.timer(Long.MAX_VALUE));
+        } else {
             subscriptions.assignFromUser(singleton(t1p));
+        }
 
         subscriptions.seek(t1p, 100);
-        coordinator.poll(time.milliseconds(), Long.MAX_VALUE);
+        coordinator.poll(time.timer(Long.MAX_VALUE));
 
         return coordinator;
     }
@@ -1594,23 +2928,21 @@ private void makeCoordinatorUnknown(ConsumerCoordinator coordinator, Errors erro
         coordinator.sendHeartbeatRequest();
         client.prepareResponse(heartbeatResponse(error));
         time.sleep(sessionTimeoutMs);
-        consumerClient.poll(0);
+        consumerClient.poll(time.timer(0));
         assertTrue(coordinator.coordinatorUnknown());
     }
+
     private void closeVerifyTimeout(final ConsumerCoordinator coordinator,
-            final long closeTimeoutMs, final long requestTimeoutMs,
-            long expectedMinTimeMs, long expectedMaxTimeMs) throws Exception {
+                                    final long closeTimeoutMs,
+                                    final long expectedMinTimeMs,
+                                    final long expectedMaxTimeMs) throws Exception {
         ExecutorService executor = Executors.newSingleThreadExecutor();
         try {
             boolean coordinatorUnknown = coordinator.coordinatorUnknown();
             // Run close on a different thread. Coordinator is locked by this thread, so it is
             // not safe to use the coordinator from the main thread until the task completes.
-            Future future = executor.submit(new Runnable() {
-                @Override
-                public void run() {
-                    coordinator.close(Math.min(closeTimeoutMs, requestTimeoutMs));
-                }
-            });
+            Future future = executor.submit(
+                () -> coordinator.close(time.timer(Math.min(closeTimeoutMs, requestTimeoutMs))));
             // Wait for close to start. If coordinator is known, wait for close to queue
             // at least one request. Otherwise, sleep for a short time.
             if (!coordinatorUnknown)
@@ -1634,86 +2966,130 @@ public void run() {
         }
     }
 
-    private void gracefulCloseTest(ConsumerCoordinator coordinator, boolean shouldLeaveGroup) throws Exception {
+    private void gracefulCloseTest(ConsumerCoordinator coordinator, boolean shouldLeaveGroup) {
         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.groupId().equals(groupId);
-            }
-        }, new LeaveGroupResponse(Errors.NONE));
+        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 excludeInternalTopics,
-                                                 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,
-                heartbeatIntervalMs,
                 assignors,
                 metadata,
                 subscriptions,
                 metrics,
-                "consumer" + groupId,
+                consumerId + groupId,
                 time,
-                retryBackoffMs,
                 autoCommitEnabled,
                 autoCommitIntervalMs,
                 null,
-                excludeInternalTopics,
-                leaveGroup);
+                false);
+    }
+
+    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) {
@@ -1721,21 +3097,102 @@ private OffsetCommitResponse offsetCommitResponse(Map re
     }
 
     private OffsetFetchResponse offsetFetchResponse(Errors topLevelError) {
-        return new OffsetFetchResponse(topLevelError, Collections.emptyMap());
+        return new OffsetFetchResponse(topLevelError, Collections.emptyMap());
     }
 
     private OffsetFetchResponse offsetFetchResponse(TopicPartition tp, Errors partitionLevelError, String metadata, long offset) {
-        OffsetFetchResponse.PartitionData data = new OffsetFetchResponse.PartitionData(offset, metadata, partitionLevelError);
-        return new OffsetFetchResponse(Errors.NONE, Collections.singletonMap(tp, data));
+        return offsetFetchResponse(tp, partitionLevelError, metadata, offset, Optional.empty());
+    }
+
+    private OffsetFetchResponse offsetFetchResponse(TopicPartition tp, Errors partitionLevelError, String metadata, long offset, Optional epoch) {
+        OffsetFetchResponse.PartitionData data = new OffsetFetchResponse.PartitionData(offset,
+                epoch, metadata, partitionLevelError);
+        return new OffsetFetchResponse(Errors.NONE, singletonMap(tp, data));
     }
 
     private OffsetCommitCallback callback(final AtomicBoolean success) {
-        return new OffsetCommitCallback() {
-            @Override
-            public void onComplete(Map offsets, Exception exception) {
-                if (exception == null)
-                    success.set(true);
+        return (offsets, exception) -> {
+            if (exception == null)
+                success.set(true);
+        };
+    }
+
+    private void joinAsFollowerAndReceiveAssignment(ConsumerCoordinator coordinator,
+                                                    List assignment) {
+        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
+        client.prepareResponse(joinGroupFollowerResponse(1, consumerId, "leader", Errors.NONE));
+        client.prepareResponse(syncGroupResponse(assignment, Errors.NONE));
+        coordinator.joinGroupIfNeeded(time.timer(Long.MAX_VALUE));
+    }
+
+    private void prepareOffsetCommitRequest(Map expectedOffsets, Errors error) {
+        prepareOffsetCommitRequest(expectedOffsets, error, false);
+    }
+
+    private void prepareOffsetCommitRequestDisconnect(Map expectedOffsets) {
+        prepareOffsetCommitRequest(expectedOffsets, Errors.NONE, true);
+    }
+
+    private void prepareOffsetCommitRequest(final Map expectedOffsets,
+                                            Errors error,
+                                            boolean disconnected) {
+        Map errors = partitionErrors(expectedOffsets.keySet(), error);
+        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) {
+            errors.put(partition, error);
+        }
+        return errors;
+    }
+
+    private void respondToOffsetCommitRequest(final Map expectedOffsets, Errors error) {
+        Map errors = partitionErrors(expectedOffsets.keySet(), error);
+        client.respond(offsetCommitRequestMatcher(expectedOffsets), offsetCommitResponse(errors));
+    }
+
+    private MockClient.RequestMatcher offsetCommitRequestMatcher(final Map expectedOffsets) {
+        return body -> {
+            OffsetCommitRequest req = (OffsetCommitRequest) body;
+            Map offsets = req.offsets();
+            if (offsets.size() != expectedOffsets.size())
+                return false;
+
+            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;
+        };
+    }
+
+    private OffsetCommitCallback callback(final Map expectedOffsets,
+                                          final AtomicBoolean success) {
+        return (offsets, exception) -> {
+            if (expectedOffsets.equals(offsets) && exception == null)
+                success.set(true);
         };
     }
 
@@ -1749,25 +3206,4 @@ public void onComplete(Map offsets, Exception
             this.exception = exception;
         }
     }
-
-    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;
-            assignedCount++;
-        }
-
-        @Override
-        public void onPartitionsRevoked(Collection partitions) {
-            this.revoked = partitions;
-            revokedCount++;
-        }
-
-    }
 }
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
new file mode 100644
index 0000000000000..0853481d6c970
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerMetadataTest.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.clients.consumer.internals;
+
+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.ClusterResourceListeners;
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.requests.MetadataRequest;
+import org.apache.kafka.common.requests.MetadataResponse;
+import org.apache.kafka.common.requests.RequestTestUtils;
+import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.common.utils.Utils;
+import org.junit.Test;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.regex.Pattern;
+
+import static java.util.Collections.singleton;
+import static java.util.Collections.singletonList;
+import static java.util.Collections.singletonMap;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+public class ConsumerMetadataTest {
+
+    private final Node node = new Node(1, "localhost", 9092);
+    private final SubscriptionState subscription = new SubscriptionState(new LogContext(), OffsetResetStrategy.EARLIEST);
+    private final Time time = new MockTime();
+
+    @Test
+    public void testPatternSubscriptionNoInternalTopics() {
+        testPatternSubscription(false);
+    }
+
+    @Test
+    public void testPatternSubscriptionIncludeInternalTopics() {
+        testPatternSubscription(true);
+    }
+
+    private void testPatternSubscription(boolean includeInternalTopics) {
+        subscription.subscribe(Pattern.compile("__.*"), new NoOpConsumerRebalanceListener());
+        ConsumerMetadata metadata = newConsumerMetadata(includeInternalTopics);
+
+        MetadataRequest.Builder builder = metadata.newMetadataRequestBuilder();
+        assertTrue(builder.isAllTopics());
+
+        List topics = new ArrayList<>();
+        topics.add(topicMetadata("__consumer_offsets", true));
+        topics.add(topicMetadata("__matching_topic", false));
+        topics.add(topicMetadata("non_matching_topic", false));
+
+        MetadataResponse response = RequestTestUtils.metadataResponse(singletonList(node),
+            "clusterId", node.id(), topics);
+        metadata.updateWithCurrentRequestVersion(response, false, time.milliseconds());
+
+        if (includeInternalTopics)
+            assertEquals(Utils.mkSet("__matching_topic", "__consumer_offsets"), metadata.fetch().topics());
+        else
+            assertEquals(Collections.singleton("__matching_topic"), metadata.fetch().topics());
+    }
+
+    @Test
+    public void testUserAssignment() {
+        subscription.assignFromUser(Utils.mkSet(
+                new TopicPartition("foo", 0),
+                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", "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
+    public void testTransientTopics() {
+        subscription.subscribe(singleton("foo"), new NoOpConsumerRebalanceListener());
+        ConsumerMetadata metadata = newConsumerMetadata(false);
+        metadata.updateWithCurrentRequestVersion(RequestTestUtils.metadataUpdateWith(1, singletonMap("foo", 1)), false, time.milliseconds());
+        assertFalse(metadata.updateRequested());
+
+        metadata.addTransientTopics(singleton("foo"));
+        assertFalse(metadata.updateRequested());
+
+        metadata.addTransientTopics(singleton("bar"));
+        assertTrue(metadata.updateRequested());
+
+        Map topicPartitionCounts = new HashMap<>();
+        topicPartitionCounts.put("foo", 1);
+        topicPartitionCounts.put("bar", 1);
+        metadata.updateWithCurrentRequestVersion(RequestTestUtils.metadataUpdateWith(1, topicPartitionCounts), false, time.milliseconds());
+        assertFalse(metadata.updateRequested());
+
+        assertEquals(Utils.mkSet("foo", "bar"), new HashSet<>(metadata.fetch().topics()));
+
+        metadata.clearTransientTopics();
+        metadata.updateWithCurrentRequestVersion(RequestTestUtils.metadataUpdateWith(1, topicPartitionCounts), false, time.milliseconds());
+        assertEquals(singleton("foo"), new HashSet<>(metadata.fetch().topics()));
+    }
+
+    private void testBasicSubscription(Set expectedTopics, Set expectedInternalTopics) {
+        Set allTopics = new HashSet<>();
+        allTopics.addAll(expectedTopics);
+        allTopics.addAll(expectedInternalTopics);
+
+        ConsumerMetadata metadata = newConsumerMetadata(false);
+
+        MetadataRequest.Builder builder = metadata.newMetadataRequestBuilder();
+        assertEquals(allTopics, new HashSet<>(builder.topics()));
+
+        List topics = new ArrayList<>();
+        for (String expectedTopic : expectedTopics)
+            topics.add(topicMetadata(expectedTopic, false));
+        for (String expectedInternalTopic : expectedInternalTopics)
+            topics.add(topicMetadata(expectedInternalTopic, true));
+
+        MetadataResponse response = RequestTestUtils.metadataResponse(singletonList(node),
+            "clusterId", node.id(), topics);
+        metadata.updateWithCurrentRequestVersion(response, false, time.milliseconds());
+
+        assertEquals(allTopics, metadata.fetch().topics());
+    }
+
+    private MetadataResponse.TopicMetadata topicMetadata(String topic, boolean isInternal) {
+        MetadataResponse.PartitionMetadata partitionMetadata = new MetadataResponse.PartitionMetadata(Errors.NONE,
+                new TopicPartition(topic, 0), Optional.of(node.id()), Optional.of(5),
+                singletonList(node.id()), singletonList(node.id()), singletonList(node.id()));
+        return new MetadataResponse.TopicMetadata(Errors.NONE, topic, isInternal, singletonList(partitionMetadata));
+    }
+
+    private ConsumerMetadata newConsumerMetadata(boolean includeInternalTopics) {
+        long refreshBackoffMs = 50;
+        long expireMs = 50000;
+        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 93c6acd9c40b4..1c3b666347d52 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,34 +21,51 @@
 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;
+import org.apache.kafka.common.errors.InvalidTopicException;
+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.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;
+import org.apache.kafka.common.requests.MetadataResponse;
+import org.apache.kafka.common.requests.RequestTestUtils;
 import org.apache.kafka.common.utils.LogContext;
 import org.apache.kafka.common.utils.MockTime;
 import org.apache.kafka.test.TestUtils;
-import org.easymock.EasyMock;
 import org.junit.Test;
 
+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.assertTrue;
 import static org.junit.Assert.fail;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 public class ConsumerNetworkClientTest {
 
     private String topicName = "test";
     private MockTime time = new MockTime(1);
-    private MockClient client = new MockClient(time);
     private Cluster cluster = TestUtils.singletonCluster(topicName, 1);
     private Node node = cluster.nodes().get(0);
-    private Metadata metadata = new Metadata(0, Long.MAX_VALUE, true);
+    private Metadata metadata = new Metadata(100, 50000, new LogContext(),
+            new ClusterResourceListeners());
+    private MockClient client = new MockClient(time, metadata);
     private ConsumerNetworkClient consumerClient = new ConsumerNetworkClient(new LogContext(),
             client, metadata, time, 100, 1000, Integer.MAX_VALUE);
 
@@ -69,6 +86,24 @@ public void send() {
         assertEquals(Errors.NONE, response.error());
     }
 
+    @Test
+    public void sendWithinBackoffPeriodAfterAuthenticationFailure() {
+        client.authenticationFailed(node, 300);
+        client.prepareResponse(heartbeatResponse(Errors.NONE));
+        final RequestFuture future = consumerClient.send(node, heartbeat());
+        consumerClient.poll(future);
+        assertTrue(future.failed());
+        assertTrue("Expected only an authentication error.", future.exception() instanceof AuthenticationException);
+
+        time.sleep(30); // wait less than the backoff period
+        assertTrue(client.connectionFailed(node));
+
+        final RequestFuture future2 = consumerClient.send(node, heartbeat());
+        consumerClient.poll(future2);
+        assertTrue(future2.failed());
+        assertTrue("Expected only an authentication error.", future2.exception() instanceof AuthenticationException);
+    }
+
     @Test
     public void multiSend() {
         client.prepareResponse(heartbeatResponse(Errors.NONE));
@@ -78,7 +113,7 @@ public void multiSend() {
         assertEquals(2, consumerClient.pendingRequestCount());
         assertEquals(2, consumerClient.pendingRequestCount(node));
 
-        consumerClient.awaitPendingRequests(node, Long.MAX_VALUE);
+        consumerClient.awaitPendingRequests(node, time.timer(Long.MAX_VALUE));
         assertTrue(future1.succeeded());
         assertTrue(future2.succeeded());
     }
@@ -88,7 +123,8 @@ public void testDisconnectWithUnsentRequests() {
         RequestFuture future = consumerClient.send(node, heartbeat());
         assertTrue(consumerClient.hasPendingRequests(node));
         assertFalse(client.hasInFlightRequests(node.idString()));
-        consumerClient.disconnect(node);
+        consumerClient.disconnectAsync(node);
+        consumerClient.pollNoWakeup();
         assertTrue(future.failed());
         assertTrue(future.exception() instanceof DisconnectException);
     }
@@ -99,76 +135,69 @@ public void testDisconnectWithInFlightRequests() {
         consumerClient.pollNoWakeup();
         assertTrue(consumerClient.hasPendingRequests(node));
         assertTrue(client.hasInFlightRequests(node.idString()));
-        consumerClient.disconnect(node);
+        consumerClient.disconnectAsync(node);
+        consumerClient.pollNoWakeup();
         assertTrue(future.failed());
         assertTrue(future.exception() instanceof DisconnectException);
     }
 
+    @Test
+    public void testTimeoutUnsentRequest() {
+        // Delay connection to the node so that the request remains unsent
+        client.delayReady(node, 1000);
+
+        RequestFuture future = consumerClient.send(node, heartbeat(), 500);
+        consumerClient.pollNoWakeup();
+
+        // Ensure the request is pending, but hasn't been sent
+        assertTrue(consumerClient.hasPendingRequests());
+        assertFalse(client.hasInFlightRequests());
+
+        time.sleep(501);
+        consumerClient.pollNoWakeup();
+
+        assertFalse(consumerClient.hasPendingRequests());
+        assertTrue(future.failed());
+        assertTrue(future.exception() instanceof TimeoutException);
+    }
+
     @Test
     public void doNotBlockIfPollConditionIsSatisfied() {
-        NetworkClient mockNetworkClient = EasyMock.mock(NetworkClient.class);
+        NetworkClient mockNetworkClient = mock(NetworkClient.class);
         ConsumerNetworkClient consumerClient = new ConsumerNetworkClient(new LogContext(),
                 mockNetworkClient, metadata, time, 100, 1000, Integer.MAX_VALUE);
 
         // expect poll, but with no timeout
-        EasyMock.expect(mockNetworkClient.poll(EasyMock.eq(0L), EasyMock.anyLong())).andReturn(Collections.emptyList());
-
-        EasyMock.replay(mockNetworkClient);
-
-        consumerClient.poll(Long.MAX_VALUE, time.milliseconds(), new ConsumerNetworkClient.PollCondition() {
-            @Override
-            public boolean shouldBlock() {
-                return false;
-            }
-        });
-
-        EasyMock.verify(mockNetworkClient);
+        consumerClient.poll(time.timer(Long.MAX_VALUE), () -> false);
+        verify(mockNetworkClient).poll(eq(0L), anyLong());
     }
 
     @Test
     public void blockWhenPollConditionNotSatisfied() {
         long timeout = 4000L;
 
-        NetworkClient mockNetworkClient = EasyMock.mock(NetworkClient.class);
+        NetworkClient mockNetworkClient = mock(NetworkClient.class);
         ConsumerNetworkClient consumerClient = new ConsumerNetworkClient(new LogContext(),
                 mockNetworkClient, metadata, time, 100, 1000, Integer.MAX_VALUE);
 
-        EasyMock.expect(mockNetworkClient.inFlightRequestCount()).andReturn(1);
-        EasyMock.expect(mockNetworkClient.poll(EasyMock.eq(timeout), EasyMock.anyLong())).andReturn(Collections.emptyList());
-
-        EasyMock.replay(mockNetworkClient);
-
-        consumerClient.poll(timeout, time.milliseconds(), new ConsumerNetworkClient.PollCondition() {
-            @Override
-            public boolean shouldBlock() {
-                return true;
-            }
-        });
-
-        EasyMock.verify(mockNetworkClient);
+        when(mockNetworkClient.inFlightRequestCount()).thenReturn(1);
+        consumerClient.poll(time.timer(timeout), () -> true);
+        verify(mockNetworkClient).poll(eq(timeout), anyLong());
     }
 
     @Test
     public void blockOnlyForRetryBackoffIfNoInflightRequests() {
         long retryBackoffMs = 100L;
 
-        NetworkClient mockNetworkClient = EasyMock.mock(NetworkClient.class);
+        NetworkClient mockNetworkClient = mock(NetworkClient.class);
         ConsumerNetworkClient consumerClient = new ConsumerNetworkClient(new LogContext(),
-                mockNetworkClient, metadata, time, retryBackoffMs, 1000L, Integer.MAX_VALUE);
-
-        EasyMock.expect(mockNetworkClient.inFlightRequestCount()).andReturn(0);
-        EasyMock.expect(mockNetworkClient.poll(EasyMock.eq(retryBackoffMs), EasyMock.anyLong())).andReturn(Collections.emptyList());
+                mockNetworkClient, metadata, time, retryBackoffMs, 1000, Integer.MAX_VALUE);
 
-        EasyMock.replay(mockNetworkClient);
+        when(mockNetworkClient.inFlightRequestCount()).thenReturn(0);
 
-        consumerClient.poll(Long.MAX_VALUE, time.milliseconds(), new ConsumerNetworkClient.PollCondition() {
-            @Override
-            public boolean shouldBlock() {
-                return true;
-            }
-        });
+        consumerClient.poll(time.timer(Long.MAX_VALUE), () -> true);
 
-        EasyMock.verify(mockNetworkClient);
+        verify(mockNetworkClient).poll(eq(retryBackoffMs), anyLong());
     }
 
     @Test
@@ -176,7 +205,7 @@ public void wakeup() {
         RequestFuture future = consumerClient.send(node, heartbeat());
         consumerClient.wakeup();
         try {
-            consumerClient.poll(0);
+            consumerClient.poll(time.timer(0));
             fail();
         } catch (WakeupException e) {
         }
@@ -186,17 +215,117 @@ public void wakeup() {
         assertTrue(future.isDone());
     }
 
+    @Test
+    public void testDisconnectWakesUpPoll() throws Exception {
+        final RequestFuture future = consumerClient.send(node, heartbeat());
+
+        client.enableBlockingUntilWakeup(1);
+        Thread t = new Thread() {
+            @Override
+            public void run() {
+                consumerClient.poll(future);
+            }
+        };
+        t.start();
+
+        consumerClient.disconnectAsync(node);
+        t.join();
+        assertTrue(future.failed());
+        assertTrue(future.exception() instanceof DisconnectException);
+    }
+
+    @Test
+    public void testAuthenticationExceptionPropagatedFromMetadata() {
+        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
+            metadata.maybeThrowAnyException();
+        }
+    }
+
+    @Test(expected = InvalidTopicException.class)
+    public void testInvalidTopicExceptionPropagatedFromMetadata() {
+        MetadataResponse metadataResponse = RequestTestUtils.metadataUpdateWith("clusterId", 1,
+                Collections.singletonMap("topic", Errors.INVALID_TOPIC_EXCEPTION), Collections.emptyMap());
+        metadata.updateWithCurrentRequestVersion(metadataResponse, false, time.milliseconds());
+        consumerClient.poll(time.timer(Duration.ZERO));
+    }
+
+    @Test(expected = TopicAuthorizationException.class)
+    public void testTopicAuthorizationExceptionPropagatedFromMetadata() {
+        MetadataResponse metadataResponse = RequestTestUtils.metadataUpdateWith("clusterId", 1,
+                Collections.singletonMap("topic", Errors.TOPIC_AUTHORIZATION_FAILED), Collections.emptyMap());
+        metadata.updateWithCurrentRequestVersion(metadataResponse, false, time.milliseconds());
+        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
+        // is received and completed in another thread.
+
+        final RequestFuture future = consumerClient.send(node, heartbeat());
+        consumerClient.pollNoWakeup(); // dequeue and send the request
+
+        client.enableBlockingUntilWakeup(2);
+        Thread t1 = new Thread() {
+            @Override
+            public void run() {
+                consumerClient.pollNoWakeup();
+            }
+        };
+        t1.start();
+
+        // Sleep a little so that t1 is blocking in poll
+        Thread.sleep(50);
+
+        Thread t2 = new Thread() {
+            @Override
+            public void run() {
+                consumerClient.poll(future);
+            }
+        };
+        t2.start();
+
+        // Sleep a little so that t2 is awaiting the network client lock
+        Thread.sleep(50);
+
+        // Simulate a network response and return from the poll in t1
+        client.respond(heartbeatResponse(Errors.NONE));
+        client.wakeup();
+
+        // Both threads should complete since t1 should wakeup t2
+        t1.join();
+        t2.join();
+        assertTrue(future.succeeded());
+    }
+
     @Test
     public void testAwaitForMetadataUpdateWithTimeout() {
-        assertFalse(consumerClient.awaitMetadataUpdate(10L));
+        assertFalse(consumerClient.awaitMetadataUpdate(time.timer(10L)));
     }
 
     @Test
-    public void sendExpiry() throws InterruptedException {
-        long unsentExpiryMs = 10;
+    public void sendExpiry() {
+        int requestTimeoutMs = 10;
         final AtomicBoolean isReady = new AtomicBoolean();
         final AtomicBoolean disconnected = new AtomicBoolean();
-        client = new MockClient(time) {
+        client = new MockClient(time, metadata) {
             @Override
             public boolean ready(Node node, long now) {
                 if (isReady.get())
@@ -210,20 +339,20 @@ public boolean connectionFailed(Node node) {
             }
         };
         // Queue first send, sleep long enough for this to expire and then queue second send
-        consumerClient = new ConsumerNetworkClient(new LogContext(), client, metadata, time, 100, unsentExpiryMs, Integer.MAX_VALUE);
+        consumerClient = new ConsumerNetworkClient(new LogContext(), client, metadata, time, 100, requestTimeoutMs, Integer.MAX_VALUE);
         RequestFuture future1 = consumerClient.send(node, heartbeat());
         assertEquals(1, consumerClient.pendingRequestCount());
         assertEquals(1, consumerClient.pendingRequestCount(node));
         assertFalse(future1.isDone());
 
-        time.sleep(unsentExpiryMs + 1);
+        time.sleep(requestTimeoutMs + 1);
         RequestFuture future2 = consumerClient.send(node, heartbeat());
         assertEquals(2, consumerClient.pendingRequestCount());
         assertEquals(2, consumerClient.pendingRequestCount(node));
         assertFalse(future2.isDone());
 
         // First send should have expired and second send still pending
-        consumerClient.poll(0);
+        consumerClient.poll(time.timer(0));
         assertTrue(future1.isDone());
         assertFalse(future1.succeeded());
         assertEquals(1, consumerClient.pendingRequestCount());
@@ -244,19 +373,56 @@ public boolean connectionFailed(Node node) {
         assertEquals(1, consumerClient.pendingRequestCount());
         assertEquals(1, consumerClient.pendingRequestCount(node));
         disconnected.set(true);
-        consumerClient.poll(0);
+        consumerClient.poll(time.timer(0));
         assertTrue(future3.isDone());
         assertFalse(future3.succeeded());
         assertEquals(0, consumerClient.pendingRequestCount());
         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..5fc5f83826b59 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,8 +16,11 @@
  */
 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.message.ConsumerProtocolAssignment;
+import org.apache.kafka.common.message.ConsumerProtocolSubscription;
 import org.apache.kafka.common.protocol.types.ArrayOf;
 import org.apache.kafka.common.protocol.types.Field;
 import org.apache.kafka.common.protocol.types.Schema;
@@ -27,22 +30,66 @@
 
 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.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 serializeDeserializeSubscriptionAllVersions() {
+        List ownedPartitions = Arrays.asList(
+            new TopicPartition("foo", 0),
+            new TopicPartition("bar", 0));
+        Subscription subscription = new Subscription(Arrays.asList("foo", "bar"),
+            ByteBuffer.wrap("hello".getBytes()), ownedPartitions);
+
+        for (short version = ConsumerProtocolSubscription.LOWEST_SUPPORTED_VERSION; version <= ConsumerProtocolSubscription.HIGHEST_SUPPORTED_VERSION; version++) {
+            ByteBuffer buffer = ConsumerProtocol.serializeSubscription(subscription, version);
+            Subscription parsedSubscription = ConsumerProtocol.deserializeSubscription(buffer);
+
+            assertEquals(subscription.topics(), parsedSubscription.topics());
+            assertEquals(subscription.userData(), parsedSubscription.userData());
+            assertFalse(parsedSubscription.groupInstanceId().isPresent());
+
+            if (version >= 1) {
+                assertEquals(toSet(subscription.ownedPartitions()), toSet(parsedSubscription.ownedPartitions()));
+            } else {
+                assertEquals(Collections.emptyList(), parsedSubscription.ownedPartitions());
+            }
+        }
+    }
+
     @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 +98,55 @@ 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 deserializeOldSubscriptionVersion() {
+        Subscription subscription = new Subscription(Arrays.asList("foo", "bar"), null);
+        ByteBuffer buffer = ConsumerProtocol.serializeSubscription(subscription, (short) 0);
+        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
+        ConsumerProtocol.deserializeVersion(buffer);
+        Subscription parsedSubscription = ConsumerProtocol.deserializeSubscription(buffer, (short) 0);
+        assertEquals(subscription.topics(), parsedSubscription.topics());
+        assertNull(parsedSubscription.userData());
+        assertTrue(parsedSubscription.ownedPartitions().isEmpty());
+        assertFalse(parsedSubscription.groupInstanceId().isPresent());
     }
 
     @Test
-    public void deserializeNewSubscriptionVersion() {
+    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("foo", Type.STRING));
+            new Field("topics", new ArrayOf(Type.STRING)),
+            new Field("user_data", Type.NULLABLE_BYTES),
+            new Field("owned_partitions", new ArrayOf(
+                ConsumerProtocolSubscription.TopicPartition.SCHEMA_1)),
+            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", new Object[]{"topic"});
+        subscriptionV100.set("user_data", ByteBuffer.wrap(new byte[0]));
+        subscriptionV100.set("owned_partitions", new Object[]{new Struct(
+            ConsumerProtocolSubscription.TopicPartition.SCHEMA_1)
+            .set("topic", tp2.topic())
+            .set("partitions", 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(new Schema(new Field("version", Type.INT16)));
+        headerV100.set("version", version);
 
         ByteBuffer buffer = ByteBuffer.allocate(subscriptionV100.sizeOf() + headerV100.sizeOf());
         headerV100.writeTo(buffer);
@@ -79,46 +155,64 @@ 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 serializeDeserializeAssignmentAllVersions() {
+        List partitions = Arrays.asList(tp1, tp2);
+        Assignment assignment = new Assignment(partitions, ByteBuffer.wrap("hello".getBytes()));
+
+        for (short version = ConsumerProtocolAssignment.LOWEST_SUPPORTED_VERSION; version <= ConsumerProtocolAssignment.HIGHEST_SUPPORTED_VERSION; version++) {
+            ByteBuffer buffer = ConsumerProtocol.serializeAssignment(assignment, version);
+            Assignment parsedAssignment = ConsumerProtocol.deserializeAssignment(buffer);
+            assertEquals(toSet(partitions), toSet(parsedAssignment.partitions()));
+            assertEquals(assignment.userData(), parsedAssignment.userData());
+        }
     }
 
     @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("foo", Type.STRING));
+            new Field("assigned_partitions", new ArrayOf(
+                ConsumerProtocolAssignment.TopicPartition.SCHEMA_0)),
+            new Field("user_data", 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("assigned_partitions",
+            new Object[]{new Struct(ConsumerProtocolAssignment.TopicPartition.SCHEMA_0)
+                .set("topic", tp1.topic())
+                .set("partitions", new Object[]{tp1.partition()})});
+        assignmentV100.set("user_data", 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(new Schema(new Field("version", Type.INT16)));
+        headerV100.set("version", version);
 
         ByteBuffer buffer = ByteBuffer.allocate(assignmentV100.sizeOf() + headerV100.sizeOf());
         headerV100.writeTo(buffer);
@@ -126,11 +220,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 7c1d9ba420855..5fd65791927f6 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
@@ -17,19 +17,23 @@
 package org.apache.kafka.clients.consumer.internals;
 
 import org.apache.kafka.clients.ApiVersions;
+import org.apache.kafka.clients.ClientDnsLookup;
 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;
 import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
-import org.apache.kafka.clients.consumer.NoOffsetForPartitionException;
+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.IsolationLevel;
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.MetricName;
 import org.apache.kafka.common.MetricNameTemplate;
@@ -43,6 +47,17 @@
 import org.apache.kafka.common.errors.TopicAuthorizationException;
 import org.apache.kafka.common.header.Header;
 import org.apache.kafka.common.header.internals.RecordHeader;
+import org.apache.kafka.common.internals.ClusterResourceListeners;
+import org.apache.kafka.common.message.ListOffsetRequestData.ListOffsetTopic;
+import org.apache.kafka.common.message.ListOffsetRequestData.ListOffsetPartition;
+import org.apache.kafka.common.message.ListOffsetResponseData;
+import org.apache.kafka.common.message.ListOffsetResponseData.ListOffsetTopicResponse;
+import org.apache.kafka.common.message.ListOffsetResponseData.ListOffsetPartitionResponse;
+import org.apache.kafka.common.message.OffsetForLeaderEpochRequestData;
+import org.apache.kafka.common.message.OffsetForLeaderEpochRequestData.OffsetForLeaderPartition;
+import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData;
+import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData.EpochEndOffset;
+import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData.OffsetForLeaderTopicResult;
 import org.apache.kafka.common.metrics.KafkaMetric;
 import org.apache.kafka.common.metrics.MetricConfig;
 import org.apache.kafka.common.metrics.Metrics;
@@ -60,20 +75,21 @@
 import org.apache.kafka.common.record.MemoryRecordsBuilder;
 import org.apache.kafka.common.record.Record;
 import org.apache.kafka.common.record.RecordBatch;
+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.FetchRequest;
-import org.apache.kafka.common.requests.FetchRequest.PartitionData;
 import org.apache.kafka.common.requests.FetchResponse;
-import org.apache.kafka.common.requests.IsolationLevel;
 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.ResponseHeader;
+import org.apache.kafka.common.requests.OffsetsForLeaderEpochRequest;
+import org.apache.kafka.common.requests.OffsetsForLeaderEpochResponse;
+import org.apache.kafka.common.requests.RequestTestUtils;
 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;
@@ -84,12 +100,15 @@
 import org.apache.kafka.test.MockSelector;
 import org.apache.kafka.test.TestUtils;
 import org.junit.After;
+import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
 
 import java.io.DataOutputStream;
+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;
@@ -99,93 +118,128 @@
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
-
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.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.emptyList;
+import static java.util.Collections.emptyMap;
 import static java.util.Collections.singleton;
+import static java.util.Collections.singletonMap;
+import static org.apache.kafka.common.requests.OffsetsForLeaderEpochResponse.UNDEFINED_EPOCH;
+import static org.apache.kafka.common.requests.OffsetsForLeaderEpochResponse.UNDEFINED_EPOCH_OFFSET;
+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;
+import static org.junit.Assert.assertNotEquals;
+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;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.spy;
 
-@SuppressWarnings("deprecation")
 public class FetcherTest {
+    private static final double EPSILON = 0.0001;
+
     private ConsumerRebalanceListener listener = new NoOpConsumerRebalanceListener();
     private String topicName = "test";
     private String groupId = "test-group";
     private final String metricGroup = "consumer" + groupId + "-fetch-manager-metrics";
     private TopicPartition tp0 = new TopicPartition(topicName, 0);
     private TopicPartition tp1 = new TopicPartition(topicName, 1);
+    private TopicPartition tp2 = new TopicPartition(topicName, 2);
+    private TopicPartition tp3 = new TopicPartition(topicName, 3);
+    private int validLeaderEpoch = 0;
+    private MetadataResponse initialUpdateResponse =
+        RequestTestUtils.metadataUpdateWith(1, singletonMap(topicName, 4));
+
     private int minBytes = 1;
     private int maxBytes = Integer.MAX_VALUE;
     private int maxWaitMs = 0;
     private int fetchSize = 1000;
     private long retryBackoffMs = 100;
+    private long requestTimeoutMs = 30000;
     private MockTime time = new MockTime(1);
-    private Metadata metadata = new Metadata(0, Long.MAX_VALUE, true);
-    private MockClient client = new MockClient(time, metadata);
-    private Cluster cluster = TestUtils.singletonCluster(topicName, 2);
-    private Node node = cluster.nodes().get(0);
-    private Metrics metrics = new Metrics(time);
-    FetcherMetricsRegistry metricsRegistry = new FetcherMetricsRegistry("consumer" + groupId);
-
-    private SubscriptionState subscriptions = new SubscriptionState(OffsetResetStrategy.EARLIEST);
-    private SubscriptionState subscriptionsNoAutoReset = new SubscriptionState(OffsetResetStrategy.NONE);
-    private static final double EPSILON = 0.0001;
-    private ConsumerNetworkClient consumerClient = new ConsumerNetworkClient(new LogContext(),
-            client, metadata, time, 100, 1000, Integer.MAX_VALUE);
+    private SubscriptionState subscriptions;
+    private ConsumerMetadata metadata;
+    private FetcherMetricsRegistry metricsRegistry;
+    private MockClient client;
+    private Metrics metrics;
+    private ApiVersions apiVersions = new ApiVersions();
+    private ConsumerNetworkClient consumerClient;
+    private Fetcher fetcher;
 
     private MemoryRecords records;
     private MemoryRecords nextRecords;
-    private Fetcher fetcher = createFetcher(subscriptions, metrics);
-    private Metrics fetcherMetrics = new Metrics(time);
-    private Fetcher fetcherNoAutoReset = createFetcher(subscriptionsNoAutoReset, fetcherMetrics);
+    private MemoryRecords emptyRecords;
+    private MemoryRecords partialRecords;
+    private ExecutorService executorService;
 
     @Before
-    public void setup() throws Exception {
-        metadata.update(cluster, Collections.emptySet(), time.milliseconds());
-        client.setNode(node);
+    public void setup() {
+        records = buildRecords(1L, 3, 1);
+        nextRecords = buildRecords(4L, 2, 4);
+        emptyRecords = buildRecords(0L, 0, 0);
+        partialRecords = buildRecords(4L, 1, 0);
+        partialRecords.buffer().putInt(Records.SIZE_OFFSET, 10000);
+    }
 
-        MemoryRecordsBuilder builder = MemoryRecords.builder(ByteBuffer.allocate(1024), CompressionType.NONE, TimestampType.CREATE_TIME, 1L);
-        builder.append(0L, "key".getBytes(), "value-1".getBytes());
-        builder.append(0L, "key".getBytes(), "value-2".getBytes());
-        builder.append(0L, "key".getBytes(), "value-3".getBytes());
-        records = builder.build();
+    private void assignFromUser(Set partitions) {
+        subscriptions.assignFromUser(partitions);
+        client.updateMetadata(initialUpdateResponse);
 
-        builder = MemoryRecords.builder(ByteBuffer.allocate(1024), CompressionType.NONE, TimestampType.CREATE_TIME, 4L);
-        builder.append(0L, "key".getBytes(), "value-4".getBytes());
-        builder.append(0L, "key".getBytes(), "value-5".getBytes());
-        nextRecords = builder.build();
+        // A dummy metadata update to ensure valid leader epoch.
+        metadata.updateWithCurrentRequestVersion(RequestTestUtils.metadataUpdateWith("dummy", 1,
+            Collections.emptyMap(), singletonMap(topicName, 4),
+            tp -> validLeaderEpoch), false, 0L);
     }
 
     @After
-    public void teardown() {
-        this.metrics.close();
-        this.fetcherMetrics.close();
-        this.fetcher.close();
-        this.fetcherMetrics.close();
+    public void teardown() throws Exception {
+        if (metrics != null)
+            this.metrics.close();
+        if (fetcher != null)
+            this.fetcher.close();
+        if (executorService != null) {
+            executorService.shutdownNow();
+            assertTrue(executorService.awaitTermination(5, TimeUnit.SECONDS));
+        }
     }
 
     @Test
     public void testFetchNormal() {
-        subscriptions.assignFromUser(singleton(tp0));
+        buildFetcher();
+
+        assignFromUser(singleton(tp0));
         subscriptions.seek(tp0, 0);
 
         // normal fetch
         assertEquals(1, fetcher.sendFetches());
         assertFalse(fetcher.hasCompletedFetches());
 
-        client.prepareResponse(fetchResponse(tp0, this.records, Errors.NONE, 100L, 0));
-        consumerClient.poll(0);
+        client.prepareResponse(fullFetchResponse(tp0, this.records, Errors.NONE, 100L, 0));
+        consumerClient.poll(time.timer(0));
         assertTrue(fetcher.hasCompletedFetches());
 
-        Map>> partitionRecords = fetcher.fetchedRecords();
+        Map>> partitionRecords = fetchedRecords();
         assertTrue(partitionRecords.containsKey(tp0));
 
         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());
@@ -193,9 +247,130 @@ public void testFetchNormal() {
         }
     }
 
+    @Test
+    public void testMissingLeaderEpochInRecords() {
+        buildFetcher();
+
+        assignFromUser(singleton(tp0));
+        subscriptions.seek(tp0, 0);
+
+        ByteBuffer buffer = ByteBuffer.allocate(1024);
+        MemoryRecordsBuilder builder = MemoryRecords.builder(buffer, RecordBatch.MAGIC_VALUE_V0,
+                CompressionType.NONE, TimestampType.CREATE_TIME, 0L, System.currentTimeMillis(),
+                RecordBatch.NO_PARTITION_LEADER_EPOCH);
+        builder.append(0L, "key".getBytes(), "1".getBytes());
+        builder.append(0L, "key".getBytes(), "2".getBytes());
+        MemoryRecords records = builder.build();
+
+        assertEquals(1, fetcher.sendFetches());
+        assertFalse(fetcher.hasCompletedFetches());
+
+        client.prepareResponse(fullFetchResponse(tp0, records, Errors.NONE, 100L, 0));
+        consumerClient.poll(time.timer(0));
+        assertTrue(fetcher.hasCompletedFetches());
+
+        Map>> partitionRecords = fetchedRecords();
+        assertTrue(partitionRecords.containsKey(tp0));
+        assertEquals(2, partitionRecords.get(tp0).size());
+
+        for (ConsumerRecord record : partitionRecords.get(tp0)) {
+            assertEquals(Optional.empty(), record.leaderEpoch());
+        }
+    }
+
+    @Test
+    public void testLeaderEpochInConsumerRecord() {
+        buildFetcher();
+
+        assignFromUser(singleton(tp0));
+        subscriptions.seek(tp0, 0);
+
+        Integer partitionLeaderEpoch = 1;
+
+        ByteBuffer buffer = ByteBuffer.allocate(1024);
+        MemoryRecordsBuilder builder = MemoryRecords.builder(buffer, RecordBatch.CURRENT_MAGIC_VALUE,
+                CompressionType.NONE, TimestampType.CREATE_TIME, 0L, System.currentTimeMillis(),
+                partitionLeaderEpoch);
+        builder.append(0L, "key".getBytes(), partitionLeaderEpoch.toString().getBytes());
+        builder.append(0L, "key".getBytes(), partitionLeaderEpoch.toString().getBytes());
+        builder.close();
+
+        partitionLeaderEpoch += 7;
+
+        builder = MemoryRecords.builder(buffer, RecordBatch.CURRENT_MAGIC_VALUE, CompressionType.NONE,
+                TimestampType.CREATE_TIME, 2L, System.currentTimeMillis(), partitionLeaderEpoch);
+        builder.append(0L, "key".getBytes(), partitionLeaderEpoch.toString().getBytes());
+        builder.close();
+
+        partitionLeaderEpoch += 5;
+        builder = MemoryRecords.builder(buffer, RecordBatch.CURRENT_MAGIC_VALUE, CompressionType.NONE,
+                TimestampType.CREATE_TIME, 3L, System.currentTimeMillis(), partitionLeaderEpoch);
+        builder.append(0L, "key".getBytes(), partitionLeaderEpoch.toString().getBytes());
+        builder.append(0L, "key".getBytes(), partitionLeaderEpoch.toString().getBytes());
+        builder.append(0L, "key".getBytes(), partitionLeaderEpoch.toString().getBytes());
+        builder.close();
+
+        buffer.flip();
+        MemoryRecords records = MemoryRecords.readableRecords(buffer);
+
+        assertEquals(1, fetcher.sendFetches());
+        assertFalse(fetcher.hasCompletedFetches());
+
+        client.prepareResponse(fullFetchResponse(tp0, records, Errors.NONE, 100L, 0));
+        consumerClient.poll(time.timer(0));
+        assertTrue(fetcher.hasCompletedFetches());
+
+        Map>> partitionRecords = fetchedRecords();
+        assertTrue(partitionRecords.containsKey(tp0));
+        assertEquals(6, partitionRecords.get(tp0).size());
+
+        for (ConsumerRecord record : partitionRecords.get(tp0)) {
+            int expectedLeaderEpoch = Integer.parseInt(Utils.utf8(record.value()));
+            assertEquals(Optional.of(expectedLeaderEpoch), record.leaderEpoch());
+        }
+    }
+
+    @Test
+    public void testClearBufferedDataForTopicPartitions() {
+        buildFetcher();
+
+        assignFromUser(singleton(tp0));
+        subscriptions.seek(tp0, 0);
+
+        // normal fetch
+        assertEquals(1, fetcher.sendFetches());
+        assertFalse(fetcher.hasCompletedFetches());
+
+        client.prepareResponse(fullFetchResponse(tp0, this.records, Errors.NONE, 100L, 0));
+        consumerClient.poll(time.timer(0));
+        assertTrue(fetcher.hasCompletedFetches());
+        Set newAssignedTopicPartitions = new HashSet<>();
+        newAssignedTopicPartitions.add(tp1);
+
+        fetcher.clearBufferedDataForUnassignedPartitions(newAssignedTopicPartitions);
+        assertFalse(fetcher.hasCompletedFetches());
+    }
+
+    @Test
+    public void testFetchSkipsBlackedOutNodes() {
+        buildFetcher();
+
+        assignFromUser(singleton(tp0));
+        subscriptions.seek(tp0, 0);
+        Node node = initialUpdateResponse.brokers().iterator().next();
+
+        client.backoff(node, 500);
+        assertEquals(0, fetcher.sendFetches());
+
+        time.sleep(500);
+        assertEquals(1, fetcher.sendFetches());
+    }
+
     @Test
     public void testFetcherIgnoresControlRecords() {
-        subscriptions.assignFromUser(singleton(tp0));
+        buildFetcher();
+
+        assignFromUser(singleton(tp0));
         subscriptions.seek(tp0, 0);
 
         // normal fetch
@@ -218,16 +393,16 @@ public void testFetcherIgnoresControlRecords() {
 
         buffer.flip();
 
-        client.prepareResponse(fetchResponse(tp0, MemoryRecords.readableRecords(buffer), Errors.NONE, 100L, 0));
-        consumerClient.poll(0);
+        client.prepareResponse(fullFetchResponse(tp0, MemoryRecords.readableRecords(buffer), Errors.NONE, 100L, 0));
+        consumerClient.poll(time.timer(0));
         assertTrue(fetcher.hasCompletedFetches());
 
-        Map>> partitionRecords = fetcher.fetchedRecords();
+        Map>> partitionRecords = fetchedRecords();
         assertTrue(partitionRecords.containsKey(tp0));
 
         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());
@@ -235,28 +410,27 @@ public void testFetcherIgnoresControlRecords() {
 
     @Test
     public void testFetchError() {
-        subscriptions.assignFromUser(singleton(tp0));
+        buildFetcher();
+
+        assignFromUser(singleton(tp0));
         subscriptions.seek(tp0, 0);
 
         assertEquals(1, fetcher.sendFetches());
         assertFalse(fetcher.hasCompletedFetches());
 
-        client.prepareResponse(fetchResponse(tp0, this.records, Errors.NOT_LEADER_FOR_PARTITION, 100L, 0));
-        consumerClient.poll(0);
+        client.prepareResponse(fullFetchResponse(tp0, this.records, Errors.NOT_LEADER_OR_FOLLOWER, 100L, 0));
+        consumerClient.poll(time.timer(0));
         assertTrue(fetcher.hasCompletedFetches());
 
-        Map>> partitionRecords = fetcher.fetchedRecords();
+        Map>> partitionRecords = fetchedRecords();
         assertFalse(partitionRecords.containsKey(tp0));
     }
 
     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;
         };
     }
 
@@ -277,15 +451,15 @@ public byte[] deserialize(String topic, byte[] data) {
             }
         };
 
-        Fetcher fetcher = createFetcher(subscriptions, new Metrics(time), deserializer, deserializer);
+        buildFetcher(deserializer, deserializer);
 
-        subscriptions.assignFromUser(singleton(tp0));
+        assignFromUser(singleton(tp0));
         subscriptions.seek(tp0, 1);
 
-        client.prepareResponse(matchesOffset(tp0, 1), fetchResponse(tp0, this.records, Errors.NONE, 100L, 0));
+        client.prepareResponse(matchesOffset(tp0, 1), fullFetchResponse(tp0, this.records, Errors.NONE, 100L, 0));
 
         assertEquals(1, fetcher.sendFetches());
-        consumerClient.poll(0);
+        consumerClient.poll(time.timer(0));
         // The fetcher should block on Deserialization error
         for (int i = 0; i < 2; i++) {
             try {
@@ -293,13 +467,16 @@ 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);
             }
         }
     }
 
     @Test
     public void testParseCorruptedRecord() throws Exception {
+        buildFetcher();
+        assignFromUser(singleton(tp0));
+
         ByteBuffer buffer = ByteBuffer.allocate(1024);
         DataOutputStream out = new DataOutputStream(new ByteBufferOutputStream(buffer));
 
@@ -339,17 +516,16 @@ public void testParseCorruptedRecord() throws Exception {
 
         buffer.flip();
 
-        subscriptions.assignFromUser(singleton(tp0));
-        subscriptions.seek(tp0, 0);
+        subscriptions.seekUnvalidated(tp0, new SubscriptionState.FetchPosition(0, Optional.empty(), metadata.currentLeader(tp0)));
 
         // normal fetch
         assertEquals(1, fetcher.sendFetches());
-        client.prepareResponse(fetchResponse(tp0, MemoryRecords.readableRecords(buffer), Errors.NONE, 100L, 0));
-        consumerClient.poll(0);
+        client.prepareResponse(fullFetchResponse(tp0, MemoryRecords.readableRecords(buffer), Errors.NONE, 100L, 0));
+        consumerClient.poll(time.timer(0));
 
         // 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);
@@ -371,28 +547,31 @@ 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);
             }
         }
     }
 
     private void seekAndConsumeRecord(ByteBuffer responseBuffer, long toOffset) {
         // Seek to skip the bad record and fetch again.
-        subscriptions.seek(tp0, toOffset);
+        subscriptions.seekUnvalidated(tp0, new SubscriptionState.FetchPosition(toOffset, Optional.empty(), metadata.currentLeader(tp0)));
         // Should not throw exception after the seek.
         fetcher.fetchedRecords();
         assertEquals(1, fetcher.sendFetches());
-        client.prepareResponse(fetchResponse(tp0, MemoryRecords.readableRecords(responseBuffer), Errors.NONE, 100L, 0));
-        consumerClient.poll(0);
+        client.prepareResponse(fullFetchResponse(tp0, MemoryRecords.readableRecords(responseBuffer), Errors.NONE, 100L, 0));
+        consumerClient.poll(time.timer(0));
 
-        List> records = fetcher.fetchedRecords().get(tp0);
+        Map>> recordsByPartition = fetchedRecords();
+        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
     public void testInvalidDefaultRecordBatch() {
+        buildFetcher();
+
         ByteBuffer buffer = ByteBuffer.allocate(1024);
         ByteBufferOutputStream out = new ByteBufferOutputStream(buffer);
 
@@ -410,13 +589,13 @@ public void testInvalidDefaultRecordBatch() {
         buffer.put("beef".getBytes());
         buffer.position(0);
 
-        subscriptions.assignFromUser(singleton(tp0));
+        assignFromUser(singleton(tp0));
         subscriptions.seek(tp0, 0);
 
         // normal fetch
         assertEquals(1, fetcher.sendFetches());
-        client.prepareResponse(fetchResponse(tp0, MemoryRecords.readableRecords(buffer), Errors.NONE, 100L, 0));
-        consumerClient.poll(0);
+        client.prepareResponse(fullFetchResponse(tp0, MemoryRecords.readableRecords(buffer), Errors.NONE, 100L, 0));
+        consumerClient.poll(time.timer(0));
 
         // the fetchedRecords() should always throw exception due to the bad batch.
         for (int i = 0; i < 2; i++) {
@@ -424,13 +603,14 @@ 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);
             }
         }
     }
 
     @Test
-    public void testParseInvalidRecordBatch() throws Exception {
+    public void testParseInvalidRecordBatch() {
+        buildFetcher();
         MemoryRecords records = MemoryRecords.withRecords(RecordBatch.MAGIC_VALUE_V2, 0L,
                 CompressionType.NONE, TimestampType.CREATE_TIME,
                 new SimpleRecord(1L, "a".getBytes(), "1".getBytes()),
@@ -441,25 +621,25 @@ public void testParseInvalidRecordBatch() throws Exception {
         // flip some bits to fail the crc
         buffer.putInt(32, buffer.get(32) ^ 87238423);
 
-        subscriptions.assignFromUser(singleton(tp0));
+        assignFromUser(singleton(tp0));
         subscriptions.seek(tp0, 0);
 
         // normal fetch
         assertEquals(1, fetcher.sendFetches());
-        client.prepareResponse(fetchResponse(tp0, MemoryRecords.readableRecords(buffer), Errors.NONE, 100L, 0));
-        consumerClient.poll(0);
+        client.prepareResponse(fullFetchResponse(tp0, MemoryRecords.readableRecords(buffer), Errors.NONE, 100L, 0));
+        consumerClient.poll(time.timer(0));
         try {
             fetcher.fetchedRecords();
             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 testHeaders() {
-        Fetcher fetcher = createFetcher(subscriptions, new Metrics(time));
+        buildFetcher();
 
         MemoryRecordsBuilder builder = MemoryRecords.builder(ByteBuffer.allocate(1024), CompressionType.NONE, TimestampType.CREATE_TIME, 1L);
         builder.append(0L, "key".getBytes(), "value-1".getBytes());
@@ -476,14 +656,15 @@ public void testHeaders() {
         MemoryRecords memoryRecords = builder.build();
 
         List> records;
-        subscriptions.assignFromUser(singleton(tp0));
+        assignFromUser(singleton(tp0));
         subscriptions.seek(tp0, 1);
 
-        client.prepareResponse(matchesOffset(tp0, 1), fetchResponse(tp0, memoryRecords, Errors.NONE, 100L, 0));
+        client.prepareResponse(matchesOffset(tp0, 1), fullFetchResponse(tp0, memoryRecords, Errors.NONE, 100L, 0));
 
         assertEquals(1, fetcher.sendFetches());
-        consumerClient.poll(0);
-        records = fetcher.fetchedRecords().get(tp0);
+        consumerClient.poll(time.timer(0));
+        Map>> recordsByPartition = fetchedRecords();
+        records = recordsByPartition.get(tp0);
 
         assertEquals(3, records.size());
 
@@ -503,35 +684,38 @@ record = recordIterator.next();
 
     @Test
     public void testFetchMaxPollRecords() {
-        Fetcher fetcher = createFetcher(subscriptions, new Metrics(time), 2);
+        buildFetcher(2);
 
         List> records;
-        subscriptions.assignFromUser(singleton(tp0));
+        assignFromUser(singleton(tp0));
         subscriptions.seek(tp0, 1);
 
-        client.prepareResponse(matchesOffset(tp0, 1), fetchResponse(tp0, this.records, Errors.NONE, 100L, 0));
-        client.prepareResponse(matchesOffset(tp0, 4), fetchResponse(tp0, this.nextRecords, Errors.NONE, 100L, 0));
+        client.prepareResponse(matchesOffset(tp0, 1), fullFetchResponse(tp0, this.records, Errors.NONE, 100L, 0));
+        client.prepareResponse(matchesOffset(tp0, 4), fullFetchResponse(tp0, this.nextRecords, Errors.NONE, 100L, 0));
 
         assertEquals(1, fetcher.sendFetches());
-        consumerClient.poll(0);
-        records = fetcher.fetchedRecords().get(tp0);
+        consumerClient.poll(time.timer(0));
+        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());
 
         assertEquals(0, fetcher.sendFetches());
-        consumerClient.poll(0);
-        records = fetcher.fetchedRecords().get(tp0);
+        consumerClient.poll(time.timer(0));
+        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);
-        consumerClient.poll(0);
-        records = fetcher.fetchedRecords().get(tp0);
+        consumerClient.poll(time.timer(0));
+        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());
     }
@@ -543,34 +727,35 @@ public void testFetchMaxPollRecords() {
      */
     @Test
     public void testFetchAfterPartitionWithFetchedRecordsIsUnassigned() {
-        Fetcher fetcher = createFetcher(subscriptions, new Metrics(time), 2);
+        buildFetcher(2);
 
         List> records;
-        subscriptions.assignFromUser(singleton(tp0));
+        assignFromUser(singleton(tp0));
         subscriptions.seek(tp0, 1);
 
         // Returns 3 records while `max.poll.records` is configured to 2
-        client.prepareResponse(matchesOffset(tp0, 1), fetchResponse(tp0, this.records, Errors.NONE, 100L, 0));
+        client.prepareResponse(matchesOffset(tp0, 1), fullFetchResponse(tp0, this.records, Errors.NONE, 100L, 0));
 
         assertEquals(1, fetcher.sendFetches());
-        consumerClient.poll(0);
-        records = fetcher.fetchedRecords().get(tp0);
+        consumerClient.poll(time.timer(0));
+        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());
 
-        subscriptions.assignFromUser(singleton(tp1));
-        client.prepareResponse(matchesOffset(tp1, 4), fetchResponse(tp1, this.nextRecords, Errors.NONE, 100L, 0));
+        assignFromUser(singleton(tp1));
+        client.prepareResponse(matchesOffset(tp1, 4), fullFetchResponse(tp1, this.nextRecords, Errors.NONE, 100L, 0));
         subscriptions.seek(tp1, 4);
 
         assertEquals(1, fetcher.sendFetches());
-        consumerClient.poll(0);
-        Map>> fetchedRecords = fetcher.fetchedRecords();
+        consumerClient.poll(time.timer(0));
+        Map>> fetchedRecords = fetchedRecords();
         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());
     }
@@ -579,6 +764,7 @@ public void testFetchAfterPartitionWithFetchedRecordsIsUnassigned() {
     public void testFetchNonContinuousRecords() {
         // if we are fetching from a compacted topic, there may be gaps in the returned records
         // this test verifies the fetcher updates the current fetched/consumed positions correctly for this case
+        buildFetcher();
 
         MemoryRecordsBuilder builder = MemoryRecords.builder(ByteBuffer.allocate(1024), CompressionType.NONE,
                 TimestampType.CREATE_TIME, 0L);
@@ -588,16 +774,17 @@ public void testFetchNonContinuousRecords() {
         MemoryRecords records = builder.build();
 
         List> consumerRecords;
-        subscriptions.assignFromUser(singleton(tp0));
+        assignFromUser(singleton(tp0));
         subscriptions.seek(tp0, 0);
 
         // normal fetch
         assertEquals(1, fetcher.sendFetches());
-        client.prepareResponse(fetchResponse(tp0, records, Errors.NONE, 100L, 0));
-        consumerClient.poll(0);
-        consumerRecords = fetcher.fetchedRecords().get(tp0);
+        client.prepareResponse(fullFetchResponse(tp0, records, Errors.NONE, 100L, 0));
+        consumerClient.poll(time.timer(0));
+        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());
@@ -611,8 +798,9 @@ public void testFetchNonContinuousRecords() {
     @Test
     public void testFetchRequestWhenRecordTooLarge() {
         try {
-            client.setNodeApiVersions(NodeApiVersions.create(Collections.singletonList(
-                new ApiVersionsResponse.ApiVersion(ApiKeys.FETCH.id, (short) 2, (short) 2))));
+            buildFetcher();
+
+            client.setNodeApiVersions(NodeApiVersions.create(ApiKeys.FETCH.id, (short) 2, (short) 2));
             makeFetchRequestWithIncompleteRecord();
             try {
                 fetcher.fetchedRecords();
@@ -620,7 +808,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());
@@ -635,6 +823,7 @@ public void testFetchRequestWhenRecordTooLarge() {
      */
     @Test
     public void testFetchRequestInternalError() {
+        buildFetcher();
         makeFetchRequestWithIncompleteRecord();
         try {
             fetcher.fetchedRecords();
@@ -642,31 +831,33 @@ 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);
         }
     }
 
     private void makeFetchRequestWithIncompleteRecord() {
-        subscriptions.assignFromUser(singleton(tp0));
+        assignFromUser(singleton(tp0));
         subscriptions.seek(tp0, 0);
         assertEquals(1, fetcher.sendFetches());
         assertFalse(fetcher.hasCompletedFetches());
         MemoryRecords partialRecord = MemoryRecords.readableRecords(
             ByteBuffer.wrap(new byte[]{0, 0, 0, 0, 0, 0, 0, 0}));
-        client.prepareResponse(fetchResponse(tp0, partialRecord, Errors.NONE, 100L, 0));
-        consumerClient.poll(0);
+        client.prepareResponse(fullFetchResponse(tp0, partialRecord, Errors.NONE, 100L, 0));
+        consumerClient.poll(time.timer(0));
         assertTrue(fetcher.hasCompletedFetches());
     }
 
     @Test
     public void testUnauthorizedTopic() {
-        subscriptions.assignFromUser(singleton(tp0));
+        buildFetcher();
+
+        assignFromUser(singleton(tp0));
         subscriptions.seek(tp0, 0);
 
         // resize the limit of the buffer to pretend it is only fetch-size large
         assertEquals(1, fetcher.sendFetches());
-        client.prepareResponse(fetchResponse(tp0, this.records, Errors.TOPIC_AUTHORIZATION_FAILED, 100L, 0));
-        consumerClient.poll(0);
+        client.prepareResponse(fullFetchResponse(tp0, this.records, Errors.TOPIC_AUTHORIZATION_FAILED, 100L, 0));
+        consumerClient.poll(time.timer(0));
         try {
             fetcher.fetchedRecords();
             fail("fetchedRecords should have thrown");
@@ -676,38 +867,75 @@ public void testUnauthorizedTopic() {
     }
 
     @Test
-    public void testFetchDuringRebalance() {
+    public void testFetchDuringEagerRebalance() {
+        buildFetcher();
+
         subscriptions.subscribe(singleton(topicName), listener);
         subscriptions.assignFromSubscribed(singleton(tp0));
         subscriptions.seek(tp0, 0);
 
+        client.updateMetadata(RequestTestUtils.metadataUpdateWith(
+            1, singletonMap(topicName, 4), tp -> validLeaderEpoch));
+
         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(fetchResponse(tp0, this.records, Errors.NONE, 100L, 0));
-        consumerClient.poll(0);
+        client.prepareResponse(fullFetchResponse(tp0, this.records, Errors.NONE, 100L, 0));
+        consumerClient.poll(time.timer(0));
 
         // The active fetch should be ignored since its position is no longer valid
         assertTrue(fetcher.fetchedRecords().isEmpty());
     }
 
+    @Test
+    public void testFetchDuringCooperativeRebalance() {
+        buildFetcher();
+
+        subscriptions.subscribe(singleton(topicName), listener);
+        subscriptions.assignFromSubscribed(singleton(tp0));
+        subscriptions.seek(tp0, 0);
+
+        client.updateMetadata(RequestTestUtils.metadataUpdateWith(
+            1, singletonMap(topicName, 4), tp -> validLeaderEpoch));
+
+        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() {
-        subscriptions.assignFromUser(singleton(tp0));
+        buildFetcher();
+
+        assignFromUser(singleton(tp0));
         subscriptions.seek(tp0, 0);
 
         assertEquals(1, fetcher.sendFetches());
         subscriptions.pause(tp0);
 
-        client.prepareResponse(fetchResponse(tp0, this.records, Errors.NONE, 100L, 0));
-        consumerClient.poll(0);
+        client.prepareResponse(fullFetchResponse(tp0, this.records, Errors.NONE, 100L, 0));
+        consumerClient.poll(time.timer(0));
         assertNull(fetcher.fetchedRecords().get(tp0));
     }
 
     @Test
     public void testFetchOnPausedPartition() {
-        subscriptions.assignFromUser(singleton(tp0));
+        buildFetcher();
+
+        assignFromUser(singleton(tp0));
         subscriptions.seek(tp0, 0);
 
         subscriptions.pause(tp0);
@@ -716,89 +944,319 @@ public void testFetchOnPausedPartition() {
     }
 
     @Test
-    public void testFetchNotLeaderForPartition() {
-        subscriptions.assignFromUser(singleton(tp0));
+    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.seekUnvalidated(tp0, new SubscriptionState.FetchPosition(1, Optional.empty(), metadata.currentLeader(tp0)));
+        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.seekUnvalidated(tp1, new SubscriptionState.FetchPosition(1, Optional.empty(), metadata.currentLeader(tp1)));
+        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.seekUnvalidated(tp0, new SubscriptionState.FetchPosition(1, Optional.empty(), metadata.currentLeader(tp0)));
+        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.seekUnvalidated(tp1, new SubscriptionState.FetchPosition(1, Optional.empty(), metadata.currentLeader(tp1)));
+        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 testFetchNotLeaderOrFollower() {
+        buildFetcher();
+        assignFromUser(singleton(tp0));
         subscriptions.seek(tp0, 0);
 
         assertEquals(1, fetcher.sendFetches());
-        client.prepareResponse(fetchResponse(tp0, this.records, Errors.NOT_LEADER_FOR_PARTITION, 100L, 0));
-        consumerClient.poll(0);
+        client.prepareResponse(fullFetchResponse(tp0, this.records, Errors.NOT_LEADER_OR_FOLLOWER, 100L, 0));
+        consumerClient.poll(time.timer(0));
         assertEquals(0, fetcher.fetchedRecords().size());
         assertEquals(0L, metadata.timeToNextUpdate(time.milliseconds()));
     }
 
     @Test
     public void testFetchUnknownTopicOrPartition() {
-        subscriptions.assignFromUser(singleton(tp0));
+        buildFetcher();
+        assignFromUser(singleton(tp0));
         subscriptions.seek(tp0, 0);
 
         assertEquals(1, fetcher.sendFetches());
-        client.prepareResponse(fetchResponse(tp0, this.records, Errors.UNKNOWN_TOPIC_OR_PARTITION, 100L, 0));
-        consumerClient.poll(0);
+        client.prepareResponse(fullFetchResponse(tp0, this.records, Errors.UNKNOWN_TOPIC_OR_PARTITION, 100L, 0));
+        consumerClient.poll(time.timer(0));
         assertEquals(0, fetcher.fetchedRecords().size());
         assertEquals(0L, metadata.timeToNextUpdate(time.milliseconds()));
     }
 
     @Test
-    public void testFetchOffsetOutOfRange() {
+    public void testFetchFencedLeaderEpoch() {
+        buildFetcher();
+        assignFromUser(singleton(tp0));
+        subscriptions.seek(tp0, 0);
+
+        assertEquals(1, fetcher.sendFetches());
+        client.prepareResponse(fullFetchResponse(tp0, this.records, Errors.FENCED_LEADER_EPOCH, 100L, 0));
+        consumerClient.poll(time.timer(0));
+
+        assertEquals("Should not return any records", 0, fetcher.fetchedRecords().size());
+        assertEquals("Should have requested metadata update", 0L, metadata.timeToNextUpdate(time.milliseconds()));
+    }
+
+    @Test
+    public void testFetchUnknownLeaderEpoch() {
+        buildFetcher();
+        assignFromUser(singleton(tp0));
+        subscriptions.seek(tp0, 0);
+
+        assertEquals(1, fetcher.sendFetches());
+        client.prepareResponse(fullFetchResponse(tp0, this.records, Errors.UNKNOWN_LEADER_EPOCH, 100L, 0));
+        consumerClient.poll(time.timer(0));
+
+        assertEquals("Should not return any records", 0, fetcher.fetchedRecords().size());
+        assertNotEquals("Should not have requested metadata update", 0L, metadata.timeToNextUpdate(time.milliseconds()));
+    }
+
+    @Test
+    public void testEpochSetInFetchRequest() {
+        buildFetcher();
         subscriptions.assignFromUser(singleton(tp0));
+        MetadataResponse metadataResponse = RequestTestUtils.metadataUpdateWith("dummy", 1,
+                Collections.emptyMap(), Collections.singletonMap(topicName, 4), tp -> 99);
+        client.updateMetadata(metadataResponse);
+
+        subscriptions.seek(tp0, 10);
+        assertEquals(1, fetcher.sendFetches());
+
+        // Check for epoch in outgoing request
+        MockClient.RequestMatcher matcher = body -> {
+            if (body instanceof FetchRequest) {
+                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", 99, partitionData.currentLeaderEpoch.get().longValue());
+                });
+                return true;
+            } else {
+                fail("Should have seen FetchRequest");
+                return false;
+            }
+        };
+        client.prepareResponse(matcher, fullFetchResponse(tp0, this.records, Errors.NONE, 100L, 0));
+        consumerClient.pollNoWakeup();
+    }
+
+    @Test
+    public void testFetchOffsetOutOfRange() {
+        buildFetcher();
+        assignFromUser(singleton(tp0));
         subscriptions.seek(tp0, 0);
 
         assertEquals(1, fetcher.sendFetches());
-        client.prepareResponse(fetchResponse(tp0, this.records, Errors.OFFSET_OUT_OF_RANGE, 100L, 0));
-        consumerClient.poll(0);
+        client.prepareResponse(fullFetchResponse(tp0, this.records, Errors.OFFSET_OUT_OF_RANGE, 100L, 0));
+        consumerClient.poll(time.timer(0));
         assertEquals(0, fetcher.fetchedRecords().size());
         assertTrue(subscriptions.isOffsetResetNeeded(tp0));
-        assertEquals(null, subscriptions.position(tp0));
+        assertNull(subscriptions.validPosition(tp0));
+        assertNull(subscriptions.position(tp0));
     }
 
     @Test
     public void testStaleOutOfRangeError() {
         // verify that an out of range error which arrives after a seek
         // does not cause us to reset our position or throw an exception
-        subscriptions.assignFromUser(singleton(tp0));
+        buildFetcher();
+        assignFromUser(singleton(tp0));
         subscriptions.seek(tp0, 0);
 
         assertEquals(1, fetcher.sendFetches());
-        client.prepareResponse(fetchResponse(tp0, this.records, Errors.OFFSET_OUT_OF_RANGE, 100L, 0));
+        client.prepareResponse(fullFetchResponse(tp0, this.records, Errors.OFFSET_OUT_OF_RANGE, 100L, 0));
         subscriptions.seek(tp0, 1);
-        consumerClient.poll(0);
+        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
     public void testFetchedRecordsAfterSeek() {
-        subscriptionsNoAutoReset.assignFromUser(singleton(tp0));
-        subscriptionsNoAutoReset.seek(tp0, 0);
+        buildFetcher(OffsetResetStrategy.NONE, new ByteArrayDeserializer(),
+                new ByteArrayDeserializer(), 2, IsolationLevel.READ_UNCOMMITTED);
+
+        assignFromUser(singleton(tp0));
+        subscriptions.seek(tp0, 0);
 
-        assertTrue(fetcherNoAutoReset.sendFetches() > 0);
-        client.prepareResponse(fetchResponse(tp0, this.records, Errors.OFFSET_OUT_OF_RANGE, 100L, 0));
-        consumerClient.poll(0);
-        assertFalse(subscriptionsNoAutoReset.isOffsetResetNeeded(tp0));
-        subscriptionsNoAutoReset.seek(tp0, 2);
-        assertEquals(0, fetcherNoAutoReset.fetchedRecords().size());
+        assertTrue(fetcher.sendFetches() > 0);
+        client.prepareResponse(fullFetchResponse(tp0, this.records, Errors.OFFSET_OUT_OF_RANGE, 100L, 0));
+        consumerClient.poll(time.timer(0));
+        assertFalse(subscriptions.isOffsetResetNeeded(tp0));
+        subscriptions.seek(tp0, 2);
+        assertEquals(0, fetcher.fetchedRecords().size());
     }
 
     @Test
     public void testFetchOffsetOutOfRangeException() {
-        subscriptionsNoAutoReset.assignFromUser(singleton(tp0));
-        subscriptionsNoAutoReset.seek(tp0, 0);
+        buildFetcher(OffsetResetStrategy.NONE, new ByteArrayDeserializer(),
+                new ByteArrayDeserializer(), 2, IsolationLevel.READ_UNCOMMITTED);
 
-        fetcherNoAutoReset.sendFetches();
-        client.prepareResponse(fetchResponse(tp0, this.records, Errors.OFFSET_OUT_OF_RANGE, 100L, 0));
-        consumerClient.poll(0);
+        assignFromUser(singleton(tp0));
+        subscriptions.seek(tp0, 0);
+
+        fetcher.sendFetches();
+        client.prepareResponse(fullFetchResponse(tp0, this.records, Errors.OFFSET_OUT_OF_RANGE, 100L, 0));
+        consumerClient.poll(time.timer(0));
 
-        assertFalse(subscriptionsNoAutoReset.isOffsetResetNeeded(tp0));
+        assertFalse(subscriptions.isOffsetResetNeeded(tp0));
         for (int i = 0; i < 2; i++) {
-            try {
-                fetcherNoAutoReset.fetchedRecords();
-                fail("Should have thrown OffsetOutOfRangeException");
-            } catch (OffsetOutOfRangeException e) {
-                assertTrue(e.offsetOutOfRangePartitions().containsKey(tp0));
-                assertEquals(e.offsetOutOfRangePartitions().size(), 1);
-            }
+            OffsetOutOfRangeException e = assertThrows(OffsetOutOfRangeException.class, () ->
+                    fetcher.fetchedRecords());
+            assertEquals(singleton(tp0), e.offsetOutOfRangePartitions().keySet());
+            assertEquals(0L, e.offsetOutOfRangePartitions().get(tp0).longValue());
         }
     }
 
@@ -806,289 +1264,705 @@ public void testFetchOffsetOutOfRangeException() {
     public void testFetchPositionAfterException() {
         // verify the advancement in the next fetch offset equals to the number of fetched records when
         // some fetched partitions cause Exception. This ensures that consumer won't lose record upon exception
-        subscriptionsNoAutoReset.assignFromUser(Utils.mkSet(tp0, tp1));
-        subscriptionsNoAutoReset.seek(tp0, 1);
-        subscriptionsNoAutoReset.seek(tp1, 1);
+        buildFetcher(OffsetResetStrategy.NONE, new ByteArrayDeserializer(),
+                new ByteArrayDeserializer(), Integer.MAX_VALUE, IsolationLevel.READ_UNCOMMITTED);
+        assignFromUser(Utils.mkSet(tp0, tp1));
+        subscriptions.seek(tp0, 1);
+        subscriptions.seek(tp1, 1);
 
-        assertEquals(1, fetcherNoAutoReset.sendFetches());
+        assertEquals(1, fetcher.sendFetches());
 
-        Map partitions = new LinkedHashMap<>();
-        partitions.put(tp1, new FetchResponse.PartitionData(Errors.NONE, 100,
+        Map> partitions = new LinkedHashMap<>();
+        partitions.put(tp1, new FetchResponse.PartitionData<>(Errors.NONE, 100,
             FetchResponse.INVALID_LAST_STABLE_OFFSET, FetchResponse.INVALID_LOG_START_OFFSET, null, records));
-        partitions.put(tp0, new FetchResponse.PartitionData(Errors.OFFSET_OUT_OF_RANGE, 100,
-                FetchResponse.INVALID_LAST_STABLE_OFFSET, FetchResponse.INVALID_LOG_START_OFFSET, null, MemoryRecords.EMPTY));
-        client.prepareResponse(new FetchResponse(new LinkedHashMap<>(partitions), 0));
-        consumerClient.poll(0);
+        partitions.put(tp0, new FetchResponse.PartitionData<>(Errors.OFFSET_OUT_OF_RANGE, 100,
+            FetchResponse.INVALID_LAST_STABLE_OFFSET, FetchResponse.INVALID_LOG_START_OFFSET, null, MemoryRecords.EMPTY));
+        client.prepareResponse(new FetchResponse<>(Errors.NONE, new LinkedHashMap<>(partitions),
+            0, INVALID_SESSION_ID));
+        consumerClient.poll(time.timer(0));
 
-        List> fetchedRecords = new ArrayList<>();
-        List exceptions = new ArrayList<>();
+        List> allFetchedRecords = new ArrayList<>();
+        fetchRecordsInto(allFetchedRecords);
 
-        for (List> records: fetcherNoAutoReset.fetchedRecords().values())
-            fetchedRecords.addAll(records);
+        assertEquals(1, subscriptions.position(tp0).offset);
+        assertEquals(4, subscriptions.position(tp1).offset);
+        assertEquals(3, allFetchedRecords.size());
 
-        assertEquals(fetchedRecords.size(), subscriptionsNoAutoReset.position(tp1) - 1);
+        OffsetOutOfRangeException e = assertThrows(OffsetOutOfRangeException.class, () ->
+                fetchRecordsInto(allFetchedRecords));
 
-        try {
-            for (List> records: fetcherNoAutoReset.fetchedRecords().values())
-                fetchedRecords.addAll(records);
-        } catch (OffsetOutOfRangeException e) {
-            exceptions.add(e);
-        }
+        assertEquals(singleton(tp0), e.offsetOutOfRangePartitions().keySet());
+        assertEquals(1L, e.offsetOutOfRangePartitions().get(tp0).longValue());
 
-        assertEquals(4, subscriptionsNoAutoReset.position(tp1).longValue());
-        assertEquals(3, fetchedRecords.size());
+        assertEquals(1, subscriptions.position(tp0).offset);
+        assertEquals(4, subscriptions.position(tp1).offset);
+        assertEquals(3, allFetchedRecords.size());
+    }
 
-        // Should have received one OffsetOutOfRangeException for partition tp1
-        assertEquals(1, exceptions.size());
-        OffsetOutOfRangeException e = exceptions.get(0);
-        assertTrue(e.offsetOutOfRangePartitions().containsKey(tp0));
-        assertEquals(e.offsetOutOfRangePartitions().size(), 1);
+    private void fetchRecordsInto(List> allFetchedRecords) {
+        Map>> fetchedRecords = fetchedRecords();
+        fetchedRecords.values().forEach(allFetchedRecords::addAll);
     }
 
     @Test
-    public void testSeekBeforeException() {
-        Fetcher fetcher = createFetcher(subscriptionsNoAutoReset, new Metrics(time), 2);
-
-        subscriptionsNoAutoReset.assignFromUser(Utils.mkSet(tp0));
-        subscriptionsNoAutoReset.seek(tp0, 1);
-        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));
-        client.prepareResponse(fetchResponse(tp0, this.records, Errors.NONE, 100L, 0));
-        consumerClient.poll(0);
+    public void testCompletedFetchRemoval() {
+        // Ensure the removal of completed fetches that cause an Exception if and only if they contain empty records.
+        buildFetcher(OffsetResetStrategy.NONE, new ByteArrayDeserializer(),
+                new ByteArrayDeserializer(), Integer.MAX_VALUE, IsolationLevel.READ_UNCOMMITTED);
+        assignFromUser(Utils.mkSet(tp0, tp1, tp2, tp3));
 
-        assertEquals(2, fetcher.fetchedRecords().get(tp0).size());
+        subscriptions.seek(tp0, 1);
+        subscriptions.seek(tp1, 1);
+        subscriptions.seek(tp2, 1);
+        subscriptions.seek(tp3, 1);
 
-        subscriptionsNoAutoReset.assignFromUser(Utils.mkSet(tp0, tp1));
-        subscriptionsNoAutoReset.seek(tp1, 1);
         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));
-        client.prepareResponse(new FetchResponse(new LinkedHashMap<>(partitions), 0));
-        consumerClient.poll(0);
-        assertEquals(1, fetcher.fetchedRecords().get(tp0).size());
 
-        subscriptionsNoAutoReset.seek(tp1, 10);
+        Map> partitions = new LinkedHashMap<>();
+        partitions.put(tp1, new FetchResponse.PartitionData<>(Errors.NONE, 100, FetchResponse.INVALID_LAST_STABLE_OFFSET,
+                FetchResponse.INVALID_LOG_START_OFFSET, null, records));
+        partitions.put(tp0, new FetchResponse.PartitionData<>(Errors.OFFSET_OUT_OF_RANGE, 100,
+                FetchResponse.INVALID_LAST_STABLE_OFFSET, FetchResponse.INVALID_LOG_START_OFFSET, null, MemoryRecords.EMPTY));
+        partitions.put(tp2, new FetchResponse.PartitionData<>(Errors.NONE, 100L, 4,
+                0L, null, nextRecords));
+        partitions.put(tp3, new FetchResponse.PartitionData<>(Errors.NONE, 100L, 4,
+                0L, null, partialRecords));
+        client.prepareResponse(new FetchResponse<>(Errors.NONE, new LinkedHashMap<>(partitions),
+                0, INVALID_SESSION_ID));
+        consumerClient.poll(time.timer(0));
+
+        List> fetchedRecords = new ArrayList<>();
+        Map>> recordsByPartition = fetchedRecords();
+        for (List> records : recordsByPartition.values())
+            fetchedRecords.addAll(records);
+
+        assertEquals(fetchedRecords.size(), subscriptions.position(tp1).offset - 1);
+        assertEquals(4, subscriptions.position(tp1).offset);
+        assertEquals(3, fetchedRecords.size());
+
+        List oorExceptions = new ArrayList<>();
+        try {
+            recordsByPartition = fetchedRecords();
+            for (List> records : recordsByPartition.values())
+                fetchedRecords.addAll(records);
+        } catch (OffsetOutOfRangeException oor) {
+            oorExceptions.add(oor);
+        }
+
+        // Should have received one OffsetOutOfRangeException for partition tp1
+        assertEquals(1, oorExceptions.size());
+        OffsetOutOfRangeException oor = oorExceptions.get(0);
+        assertTrue(oor.offsetOutOfRangePartitions().containsKey(tp0));
+        assertEquals(oor.offsetOutOfRangePartitions().size(), 1);
+
+        recordsByPartition = fetchedRecords();
+        for (List> records : recordsByPartition.values())
+            fetchedRecords.addAll(records);
+
+        // Should not have received an Exception for tp2.
+        assertEquals(6, subscriptions.position(tp2).offset);
+        assertEquals(5, fetchedRecords.size());
+
+        int numExceptionsExpected = 3;
+        List kafkaExceptions = new ArrayList<>();
+        for (int i = 1; i <= numExceptionsExpected; i++) {
+            try {
+                recordsByPartition = fetchedRecords();
+                for (List> records : recordsByPartition.values())
+                    fetchedRecords.addAll(records);
+            } catch (KafkaException e) {
+                kafkaExceptions.add(e);
+            }
+        }
+        // Should have received as much as numExceptionsExpected Kafka exceptions for tp3.
+        assertEquals(numExceptionsExpected, kafkaExceptions.size());
+    }
+
+    @Test
+    public void testSeekBeforeException() {
+        buildFetcher(OffsetResetStrategy.NONE, new ByteArrayDeserializer(),
+                new ByteArrayDeserializer(), 2, IsolationLevel.READ_UNCOMMITTED);
+
+        assignFromUser(Utils.mkSet(tp0));
+        subscriptions.seek(tp0, 1);
+        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, Optional.empty(), null, records));
+        client.prepareResponse(fullFetchResponse(tp0, this.records, Errors.NONE, 100L, 0));
+        consumerClient.poll(time.timer(0));
+
+        assertEquals(2, fetcher.fetchedRecords().get(tp0).size());
+
+        subscriptions.assignFromUser(Utils.mkSet(tp0, tp1));
+        subscriptions.seekUnvalidated(tp1, new SubscriptionState.FetchPosition(1, Optional.empty(), metadata.currentLeader(tp1)));
+
+        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, 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());
+
+        subscriptions.seek(tp1, 10);
         // Should not throw OffsetOutOfRangeException after the seek
         assertEquals(0, fetcher.fetchedRecords().size());
     }
 
     @Test
     public void testFetchDisconnected() {
-        subscriptions.assignFromUser(singleton(tp0));
+        buildFetcher();
+
+        assignFromUser(singleton(tp0));
         subscriptions.seek(tp0, 0);
 
         assertEquals(1, fetcher.sendFetches());
-        client.prepareResponse(fetchResponse(tp0, this.records, Errors.NONE, 100L, 0), true);
-        consumerClient.poll(0);
+        client.prepareResponse(fullFetchResponse(tp0, this.records, Errors.NONE, 100L, 0), true);
+        consumerClient.poll(time.timer(0));
         assertEquals(0, fetcher.fetchedRecords().size());
 
         // 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
-    public void testUpdateFetchPositionsNoneCommittedNoResetStrategy() {
-        Set tps = new HashSet<>(Arrays.asList(tp0, tp1));
-        subscriptionsNoAutoReset.assignFromUser(tps);
-        try {
-            fetcherNoAutoReset.updateFetchPositions(tps);
-            fail("Should have thrown NoOffsetForPartitionException");
-        } catch (NoOffsetForPartitionException e) {
-            // we expect the exception to be thrown for both TPs at the same time
-            Set partitions = e.partitions();
-            assertEquals(tps, partitions);
-        }
+    public void testUpdateFetchPositionNoOpWithPositionSet() {
+        buildFetcher();
+        assignFromUser(singleton(tp0));
+        subscriptions.seek(tp0, 5L);
+
+        fetcher.resetOffsetsIfNeeded();
+        assertFalse(client.hasInFlightRequests());
+        assertTrue(subscriptions.isFetchable(tp0));
+        assertEquals(5, subscriptions.position(tp0).offset);
     }
 
     @Test
-    public void testUpdateFetchPositionToCommitted() {
-        // unless a specific reset is expected, the default behavior is to reset to the committed
-        // position if one is present
-        subscriptions.assignFromUser(singleton(tp0));
-        subscriptions.committed(tp0, new OffsetAndMetadata(5));
+    public void testUpdateFetchPositionResetToDefaultOffset() {
+        buildFetcher();
+        assignFromUser(singleton(tp0));
+        subscriptions.requestOffsetReset(tp0);
+
+        client.prepareResponse(listOffsetRequestMatcher(ListOffsetRequest.EARLIEST_TIMESTAMP,
+            validLeaderEpoch), listOffsetResponse(Errors.NONE, 1L, 5L));
+        fetcher.resetOffsetsIfNeeded();
+        consumerClient.pollNoWakeup();
+        assertFalse(subscriptions.isOffsetResetNeeded(tp0));
+        assertTrue(subscriptions.isFetchable(tp0));
+        assertEquals(5, subscriptions.position(tp0).offset);
+    }
 
-        fetcher.updateFetchPositions(singleton(tp0));
+    @Test
+    public void testUpdateFetchPositionResetToLatestOffset() {
+        buildFetcher();
+        assignFromUser(singleton(tp0));
+        subscriptions.requestOffsetReset(tp0, OffsetResetStrategy.LATEST);
+
+        client.updateMetadata(initialUpdateResponse);
+
+        client.prepareResponse(listOffsetRequestMatcher(ListOffsetRequest.LATEST_TIMESTAMP),
+                listOffsetResponse(Errors.NONE, 1L, 5L));
+        fetcher.resetOffsetsIfNeeded();
+        consumerClient.pollNoWakeup();
+        assertFalse(subscriptions.isOffsetResetNeeded(tp0));
         assertTrue(subscriptions.isFetchable(tp0));
-        assertEquals(5, subscriptions.position(tp0).longValue());
+        assertEquals(5, subscriptions.position(tp0).offset);
     }
 
+    /**
+     * Make sure the client behaves appropriately when receiving an exception for unavailable offsets
+     */
     @Test
-    public void testUpdateFetchPositionResetToDefaultOffset() {
-        subscriptions.assignFromUser(singleton(tp0));
-        // with no commit position, we should reset using the default strategy defined above (EARLIEST)
+    public void testFetchOffsetErrors() {
+        buildFetcher();
+        assignFromUser(singleton(tp0));
+        subscriptions.requestOffsetReset(tp0, OffsetResetStrategy.LATEST);
+
+        // Fail with OFFSET_NOT_AVAILABLE
+        client.prepareResponse(listOffsetRequestMatcher(ListOffsetRequest.LATEST_TIMESTAMP,
+            validLeaderEpoch), listOffsetResponse(Errors.OFFSET_NOT_AVAILABLE, 1L, 5L), false);
+        fetcher.resetOffsetsIfNeeded();
+        consumerClient.pollNoWakeup();
+        assertFalse(subscriptions.hasValidPosition(tp0));
+        assertTrue(subscriptions.isOffsetResetNeeded(tp0));
+        assertFalse(subscriptions.isFetchable(tp0));
+
+        // Fail with LEADER_NOT_AVAILABLE
+        time.sleep(retryBackoffMs);
+        client.prepareResponse(listOffsetRequestMatcher(ListOffsetRequest.LATEST_TIMESTAMP,
+            validLeaderEpoch), listOffsetResponse(Errors.LEADER_NOT_AVAILABLE, 1L, 5L), false);
+        fetcher.resetOffsetsIfNeeded();
+        consumerClient.pollNoWakeup();
+        assertFalse(subscriptions.hasValidPosition(tp0));
+        assertTrue(subscriptions.isOffsetResetNeeded(tp0));
+        assertFalse(subscriptions.isFetchable(tp0));
+
+        // Back to normal
+        time.sleep(retryBackoffMs);
+        client.prepareResponse(listOffsetRequestMatcher(ListOffsetRequest.LATEST_TIMESTAMP),
+                listOffsetResponse(Errors.NONE, 1L, 5L), false);
+        fetcher.resetOffsetsIfNeeded();
+        consumerClient.pollNoWakeup();
+        assertTrue(subscriptions.hasValidPosition(tp0));
+        assertFalse(subscriptions.isOffsetResetNeeded(tp0));
+        assertTrue(subscriptions.isFetchable(tp0));
+        assertEquals(subscriptions.position(tp0).offset, 5L);
+    }
+
+    @Test
+    public void testListOffsetSendsReadUncommitted() {
+        testListOffsetsSendsIsolationLevel(IsolationLevel.READ_UNCOMMITTED);
+    }
+
+    @Test
+    public void testListOffsetSendsReadCommitted() {
+        testListOffsetsSendsIsolationLevel(IsolationLevel.READ_COMMITTED);
+    }
+
+    private void testListOffsetsSendsIsolationLevel(IsolationLevel isolationLevel) {
+        buildFetcher(OffsetResetStrategy.EARLIEST, new ByteArrayDeserializer(), new ByteArrayDeserializer(),
+                Integer.MAX_VALUE, isolationLevel);
+
+        assignFromUser(singleton(tp0));
+        subscriptions.requestOffsetReset(tp0, OffsetResetStrategy.LATEST);
+
+        client.prepareResponse(body -> {
+            ListOffsetRequest request = (ListOffsetRequest) body;
+            return request.isolationLevel() == isolationLevel;
+        }, listOffsetResponse(Errors.NONE, 1L, 5L));
+        fetcher.resetOffsetsIfNeeded();
+        consumerClient.pollNoWakeup();
+
+        assertFalse(subscriptions.isOffsetResetNeeded(tp0));
+        assertTrue(subscriptions.isFetchable(tp0));
+        assertEquals(5, subscriptions.position(tp0).offset);
+    }
+
+    @Test
+    public void testResetOffsetsSkipsBlackedOutConnections() {
+        buildFetcher();
+        assignFromUser(singleton(tp0));
+        subscriptions.requestOffsetReset(tp0, OffsetResetStrategy.EARLIEST);
+
+        // Check that we skip sending the ListOffset request when the node is blacked out
+        client.updateMetadata(initialUpdateResponse);
+        Node node = initialUpdateResponse.brokers().iterator().next();
+        client.backoff(node, 500);
+        fetcher.resetOffsetsIfNeeded();
+        assertEquals(0, consumerClient.pendingRequestCount());
+        consumerClient.pollNoWakeup();
+        assertTrue(subscriptions.isOffsetResetNeeded(tp0));
+        assertEquals(OffsetResetStrategy.EARLIEST, subscriptions.resetStrategy(tp0));
 
+        time.sleep(500);
         client.prepareResponse(listOffsetRequestMatcher(ListOffsetRequest.EARLIEST_TIMESTAMP),
-                               listOffsetResponse(Errors.NONE, 1L, 5L));
-        fetcher.updateFetchPositions(singleton(tp0));
+                listOffsetResponse(Errors.NONE, 1L, 5L));
+        fetcher.resetOffsetsIfNeeded();
+        consumerClient.pollNoWakeup();
+
         assertFalse(subscriptions.isOffsetResetNeeded(tp0));
         assertTrue(subscriptions.isFetchable(tp0));
-        assertEquals(5, subscriptions.position(tp0).longValue());
+        assertEquals(5, subscriptions.position(tp0).offset);
     }
 
     @Test
-    public void testUpdateFetchPositionResetToLatestOffset() {
-        subscriptions.assignFromUser(singleton(tp0));
-        subscriptions.needOffsetReset(tp0, OffsetResetStrategy.LATEST);
+    public void testUpdateFetchPositionResetToEarliestOffset() {
+        buildFetcher();
+        assignFromUser(singleton(tp0));
+        subscriptions.requestOffsetReset(tp0, OffsetResetStrategy.EARLIEST);
+
+        client.prepareResponse(listOffsetRequestMatcher(ListOffsetRequest.EARLIEST_TIMESTAMP,
+            validLeaderEpoch), listOffsetResponse(Errors.NONE, 1L, 5L));
+        fetcher.resetOffsetsIfNeeded();
+        consumerClient.pollNoWakeup();
+
+        assertFalse(subscriptions.isOffsetResetNeeded(tp0));
+        assertTrue(subscriptions.isFetchable(tp0));
+        assertEquals(5, subscriptions.position(tp0).offset);
+    }
 
+    @Test
+    public void testResetOffsetsMetadataRefresh() {
+        buildFetcher();
+        assignFromUser(singleton(tp0));
+        subscriptions.requestOffsetReset(tp0, OffsetResetStrategy.LATEST);
+
+        // First fetch fails with stale metadata
+        client.prepareResponse(listOffsetRequestMatcher(ListOffsetRequest.LATEST_TIMESTAMP,
+            validLeaderEpoch), listOffsetResponse(Errors.NOT_LEADER_OR_FOLLOWER, 1L, 5L), false);
+        fetcher.resetOffsetsIfNeeded();
+        consumerClient.pollNoWakeup();
+        assertFalse(subscriptions.hasValidPosition(tp0));
+
+        // Expect a metadata refresh
+        client.prepareMetadataUpdate(initialUpdateResponse);
+        consumerClient.pollNoWakeup();
+        assertFalse(client.hasPendingMetadataUpdates());
+
+        // Next fetch succeeds
+        time.sleep(retryBackoffMs);
         client.prepareResponse(listOffsetRequestMatcher(ListOffsetRequest.LATEST_TIMESTAMP),
-                               listOffsetResponse(Errors.NONE, 1L, 5L));
-        fetcher.updateFetchPositions(singleton(tp0));
+                listOffsetResponse(Errors.NONE, 1L, 5L));
+        fetcher.resetOffsetsIfNeeded();
+        consumerClient.pollNoWakeup();
+
         assertFalse(subscriptions.isOffsetResetNeeded(tp0));
         assertTrue(subscriptions.isFetchable(tp0));
-        assertEquals(5, subscriptions.position(tp0).longValue());
+        assertEquals(5, subscriptions.position(tp0).offset);
     }
 
     @Test
-    public void testListOffsetsSendsIsolationLevel() {
-        for (final IsolationLevel isolationLevel : IsolationLevel.values()) {
-            Fetcher fetcher = createFetcher(subscriptions, new Metrics(), new ByteArrayDeserializer(),
-                    new ByteArrayDeserializer(), Integer.MAX_VALUE, isolationLevel);
+    public void testListOffsetNoUpdateMissingEpoch() {
+        buildFetcher();
 
-            subscriptions.assignFromUser(singleton(tp0));
-            subscriptions.needOffsetReset(tp0, OffsetResetStrategy.LATEST);
+        // Set up metadata with no leader epoch
+        subscriptions.assignFromUser(singleton(tp0));
+        MetadataResponse metadataWithNoLeaderEpochs = RequestTestUtils.metadataUpdateWith(
+                "kafka-cluster", 1, Collections.emptyMap(), singletonMap(topicName, 4), tp -> null);
+        client.updateMetadata(metadataWithNoLeaderEpochs);
 
-            client.prepareResponse(new MockClient.RequestMatcher() {
-                @Override
-                public boolean matches(AbstractRequest body) {
-                    ListOffsetRequest request = (ListOffsetRequest) body;
-                    return request.isolationLevel() == isolationLevel;
-                }
-            }, listOffsetResponse(Errors.NONE, 1L, 5L));
-            fetcher.updateFetchPositions(singleton(tp0));
-            assertFalse(subscriptions.isOffsetResetNeeded(tp0));
-            assertTrue(subscriptions.isFetchable(tp0));
-            assertEquals(5, subscriptions.position(tp0).longValue());
-        }
+        // Return a ListOffsets response with leaderEpoch=1, we should ignore it
+        subscriptions.requestOffsetReset(tp0, OffsetResetStrategy.LATEST);
+        client.prepareResponse(listOffsetRequestMatcher(ListOffsetRequest.LATEST_TIMESTAMP),
+                listOffsetResponse(tp0, Errors.NONE, 1L, 5L, 1));
+        fetcher.resetOffsetsIfNeeded();
+        consumerClient.pollNoWakeup();
+
+        // Reset should be satisfied and no metadata update requested
+        assertFalse(subscriptions.isOffsetResetNeeded(tp0));
+        assertFalse(metadata.updateRequested());
+        assertFalse(metadata.lastSeenLeaderEpoch(tp0).isPresent());
     }
 
     @Test
-    public void testUpdateFetchPositionResetToEarliestOffset() {
+    public void testListOffsetUpdateEpoch() {
+        buildFetcher();
+
+        // Set up metadata with leaderEpoch=1
         subscriptions.assignFromUser(singleton(tp0));
-        subscriptions.needOffsetReset(tp0, OffsetResetStrategy.EARLIEST);
+        MetadataResponse metadataWithLeaderEpochs = RequestTestUtils.metadataUpdateWith(
+                "kafka-cluster", 1, Collections.emptyMap(), singletonMap(topicName, 4), tp -> 1);
+        client.updateMetadata(metadataWithLeaderEpochs);
+
+        // Reset offsets to trigger ListOffsets call
+        subscriptions.requestOffsetReset(tp0, OffsetResetStrategy.LATEST);
+
+        // Now we see a ListOffsets with leaderEpoch=2 epoch, we trigger a metadata update
+        client.prepareResponse(listOffsetRequestMatcher(ListOffsetRequest.LATEST_TIMESTAMP, 1),
+                listOffsetResponse(tp0, Errors.NONE, 1L, 5L, 2));
+        fetcher.resetOffsetsIfNeeded();
+        consumerClient.pollNoWakeup();
 
-        client.prepareResponse(listOffsetRequestMatcher(ListOffsetRequest.EARLIEST_TIMESTAMP),
-                               listOffsetResponse(Errors.NONE, 1L, 5L));
-        fetcher.updateFetchPositions(singleton(tp0));
         assertFalse(subscriptions.isOffsetResetNeeded(tp0));
-        assertTrue(subscriptions.isFetchable(tp0));
-        assertEquals(5, subscriptions.position(tp0).longValue());
+        assertTrue(metadata.updateRequested());
+        assertOptional(metadata.lastSeenLeaderEpoch(tp0), epoch -> assertEquals((long) epoch, 2));
     }
 
     @Test
     public void testUpdateFetchPositionDisconnect() {
-        subscriptions.assignFromUser(singleton(tp0));
-        subscriptions.needOffsetReset(tp0, OffsetResetStrategy.LATEST);
+        buildFetcher();
+        assignFromUser(singleton(tp0));
+        subscriptions.requestOffsetReset(tp0, OffsetResetStrategy.LATEST);
 
         // First request gets a disconnect
-        client.prepareResponse(listOffsetRequestMatcher(ListOffsetRequest.LATEST_TIMESTAMP),
-                               listOffsetResponse(Errors.NONE, 1L, 5L), true);
+        client.prepareResponse(listOffsetRequestMatcher(ListOffsetRequest.LATEST_TIMESTAMP,
+            validLeaderEpoch), listOffsetResponse(Errors.NONE, 1L, 5L), true);
+        fetcher.resetOffsetsIfNeeded();
+        consumerClient.pollNoWakeup();
+        assertFalse(subscriptions.hasValidPosition(tp0));
+
+        // Expect a metadata refresh
+        client.prepareMetadataUpdate(initialUpdateResponse);
+        consumerClient.pollNoWakeup();
+        assertFalse(client.hasPendingMetadataUpdates());
+
+        // No retry until the backoff passes
+        fetcher.resetOffsetsIfNeeded();
+        consumerClient.pollNoWakeup();
+        assertFalse(client.hasInFlightRequests());
+        assertFalse(subscriptions.hasValidPosition(tp0));
 
         // Next one succeeds
+        time.sleep(retryBackoffMs);
         client.prepareResponse(listOffsetRequestMatcher(ListOffsetRequest.LATEST_TIMESTAMP),
-                               listOffsetResponse(Errors.NONE, 1L, 5L));
-        fetcher.updateFetchPositions(singleton(tp0));
+                listOffsetResponse(Errors.NONE, 1L, 5L));
+        fetcher.resetOffsetsIfNeeded();
+        consumerClient.pollNoWakeup();
+
         assertFalse(subscriptions.isOffsetResetNeeded(tp0));
         assertTrue(subscriptions.isFetchable(tp0));
-        assertEquals(5, subscriptions.position(tp0).longValue());
+        assertEquals(5, subscriptions.position(tp0).offset);
     }
 
     @Test
-    public void testUpdateFetchPositionOfPausedPartitionsRequiringOffsetReset() {
-        subscriptions.assignFromUser(singleton(tp0));
-        subscriptions.committed(tp0, new OffsetAndMetadata(0));
-        subscriptions.pause(tp0); // paused partition does not have a valid position
-        subscriptions.needOffsetReset(tp0, OffsetResetStrategy.LATEST);
+    public void testAssignmentChangeWithInFlightReset() {
+        buildFetcher();
+        assignFromUser(singleton(tp0));
+        subscriptions.requestOffsetReset(tp0, OffsetResetStrategy.LATEST);
+
+        // Send the ListOffsets request to reset the position
+        fetcher.resetOffsetsIfNeeded();
+        consumerClient.pollNoWakeup();
+        assertFalse(subscriptions.hasValidPosition(tp0));
+        assertTrue(client.hasInFlightRequests());
+
+        // Now we have an assignment change
+        assignFromUser(singleton(tp1));
+
+        // The response returns and is discarded
+        client.respond(listOffsetResponse(Errors.NONE, 1L, 5L));
+        consumerClient.pollNoWakeup();
+
+        assertFalse(client.hasPendingResponses());
+        assertFalse(client.hasInFlightRequests());
+        assertFalse(subscriptions.isAssigned(tp0));
+    }
+
+    @Test
+    public void testSeekWithInFlightReset() {
+        buildFetcher();
+        assignFromUser(singleton(tp0));
+        subscriptions.requestOffsetReset(tp0, OffsetResetStrategy.LATEST);
+
+        // Send the ListOffsets request to reset the position
+        fetcher.resetOffsetsIfNeeded();
+        consumerClient.pollNoWakeup();
+        assertFalse(subscriptions.hasValidPosition(tp0));
+        assertTrue(client.hasInFlightRequests());
+
+        // Now we get a seek from the user
+        subscriptions.seek(tp0, 237);
+
+        // The response returns and is discarded
+        client.respond(listOffsetResponse(Errors.NONE, 1L, 5L));
+        consumerClient.pollNoWakeup();
+
+        assertFalse(client.hasPendingResponses());
+        assertFalse(client.hasInFlightRequests());
+        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, new SubscriptionState.FetchPosition(0L,
+                Optional.empty(), metadata.currentLeader(tp0)), 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
+    public void testChangeResetWithInFlightReset() {
+        buildFetcher();
+        assignFromUser(singleton(tp0));
+        subscriptions.requestOffsetReset(tp0, OffsetResetStrategy.LATEST);
+
+        // Send the ListOffsets request to reset the position
+        fetcher.resetOffsetsIfNeeded();
+        consumerClient.pollNoWakeup();
+        assertFalse(subscriptions.hasValidPosition(tp0));
+        assertTrue(client.hasInFlightRequests());
+
+        // Now we get a seek from the user
+        subscriptions.requestOffsetReset(tp0, OffsetResetStrategy.EARLIEST);
+
+        // The response returns and is discarded
+        client.respond(listOffsetResponse(Errors.NONE, 1L, 5L));
+        consumerClient.pollNoWakeup();
+
+        assertFalse(client.hasPendingResponses());
+        assertFalse(client.hasInFlightRequests());
+        assertTrue(subscriptions.isOffsetResetNeeded(tp0));
+        assertEquals(OffsetResetStrategy.EARLIEST, subscriptions.resetStrategy(tp0));
+    }
+
+    @Test
+    public void testIdempotentResetWithInFlightReset() {
+        buildFetcher();
+        assignFromUser(singleton(tp0));
+        subscriptions.requestOffsetReset(tp0, OffsetResetStrategy.LATEST);
+
+        // Send the ListOffsets request to reset the position
+        fetcher.resetOffsetsIfNeeded();
+        consumerClient.pollNoWakeup();
+        assertFalse(subscriptions.hasValidPosition(tp0));
+        assertTrue(client.hasInFlightRequests());
 
+        // Now we get a seek from the user
+        subscriptions.requestOffsetReset(tp0, OffsetResetStrategy.LATEST);
+
+        client.respond(listOffsetResponse(Errors.NONE, 1L, 5L));
+        consumerClient.pollNoWakeup();
+
+        assertFalse(client.hasInFlightRequests());
+        assertFalse(subscriptions.isOffsetResetNeeded(tp0));
+        assertEquals(5L, subscriptions.position(tp0).offset);
+    }
+
+    @Test
+    public void testRestOffsetsAuthorizationFailure() {
+        buildFetcher();
+        assignFromUser(singleton(tp0));
+        subscriptions.requestOffsetReset(tp0, OffsetResetStrategy.LATEST);
+
+        // First request gets a disconnect
+        client.prepareResponse(listOffsetRequestMatcher(ListOffsetRequest.LATEST_TIMESTAMP,
+            validLeaderEpoch), listOffsetResponse(Errors.TOPIC_AUTHORIZATION_FAILED, -1, -1), false);
+        fetcher.resetOffsetsIfNeeded();
+        consumerClient.pollNoWakeup();
+        assertFalse(subscriptions.hasValidPosition(tp0));
+
+        try {
+            fetcher.resetOffsetsIfNeeded();
+            fail("Expected authorization error to be raised");
+        } catch (TopicAuthorizationException e) {
+            assertEquals(singleton(tp0.topic()), e.unauthorizedTopics());
+        }
+
+        // The exception should clear after being raised, but no retry until the backoff
+        fetcher.resetOffsetsIfNeeded();
+        consumerClient.pollNoWakeup();
+        assertFalse(client.hasInFlightRequests());
+        assertFalse(subscriptions.hasValidPosition(tp0));
+
+        // Next one succeeds
+        time.sleep(retryBackoffMs);
         client.prepareResponse(listOffsetRequestMatcher(ListOffsetRequest.LATEST_TIMESTAMP),
-                               listOffsetResponse(Errors.NONE, 1L, 10L));
-        fetcher.updateFetchPositions(singleton(tp0));
+                listOffsetResponse(Errors.NONE, 1L, 5L));
+        fetcher.resetOffsetsIfNeeded();
+        consumerClient.pollNoWakeup();
 
         assertFalse(subscriptions.isOffsetResetNeeded(tp0));
-        assertFalse(subscriptions.isFetchable(tp0)); // because tp is paused
-        assertTrue(subscriptions.hasValidPosition(tp0));
-        assertEquals(10, subscriptions.position(tp0).longValue());
+        assertTrue(subscriptions.isFetchable(tp0));
+        assertEquals(5, subscriptions.position(tp0).offset);
     }
 
     @Test
-    public void testUpdateFetchPositionOfPausedPartitionsWithoutACommittedOffset() {
-        subscriptions.assignFromUser(singleton(tp0));
+    public void testUpdateFetchPositionOfPausedPartitionsRequiringOffsetReset() {
+        buildFetcher();
+        assignFromUser(singleton(tp0));
         subscriptions.pause(tp0); // paused partition does not have a valid position
+        subscriptions.requestOffsetReset(tp0, OffsetResetStrategy.LATEST);
 
-        client.prepareResponse(listOffsetRequestMatcher(ListOffsetRequest.EARLIEST_TIMESTAMP),
-                               listOffsetResponse(Errors.NONE, 1L, 0L));
-        fetcher.updateFetchPositions(singleton(tp0));
+        client.prepareResponse(listOffsetRequestMatcher(ListOffsetRequest.LATEST_TIMESTAMP,
+            validLeaderEpoch), listOffsetResponse(Errors.NONE, 1L, 10L));
+        fetcher.resetOffsetsIfNeeded();
+        consumerClient.pollNoWakeup();
 
         assertFalse(subscriptions.isOffsetResetNeeded(tp0));
         assertFalse(subscriptions.isFetchable(tp0)); // because tp is paused
         assertTrue(subscriptions.hasValidPosition(tp0));
-        assertEquals(0, subscriptions.position(tp0).longValue());
+        assertEquals(10, subscriptions.position(tp0).offset);
     }
 
     @Test
     public void testUpdateFetchPositionOfPausedPartitionsWithoutAValidPosition() {
-        subscriptions.assignFromUser(singleton(tp0));
-        subscriptions.committed(tp0, new OffsetAndMetadata(0));
+        buildFetcher();
+        assignFromUser(singleton(tp0));
+        subscriptions.requestOffsetReset(tp0);
         subscriptions.pause(tp0); // paused partition does not have a valid position
-        subscriptions.seek(tp0, 10);
 
-        fetcher.updateFetchPositions(singleton(tp0));
+        fetcher.resetOffsetsIfNeeded();
+        consumerClient.pollNoWakeup();
 
-        assertFalse(subscriptions.isOffsetResetNeeded(tp0));
+        assertTrue(subscriptions.isOffsetResetNeeded(tp0));
         assertFalse(subscriptions.isFetchable(tp0)); // because tp is paused
-        assertTrue(subscriptions.hasValidPosition(tp0));
-        assertEquals(10, subscriptions.position(tp0).longValue());
+        assertFalse(subscriptions.hasValidPosition(tp0));
     }
 
     @Test
     public void testUpdateFetchPositionOfPausedPartitionsWithAValidPosition() {
-        subscriptions.assignFromUser(singleton(tp0));
-        subscriptions.committed(tp0, new OffsetAndMetadata(0));
+        buildFetcher();
+        assignFromUser(singleton(tp0));
         subscriptions.seek(tp0, 10);
         subscriptions.pause(tp0); // paused partition already has a valid position
 
-        fetcher.updateFetchPositions(singleton(tp0));
+        fetcher.resetOffsetsIfNeeded();
 
         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
     public void testGetAllTopics() {
         // sending response before request, as getTopicMetadata is a blocking call
+        buildFetcher();
+        assignFromUser(singleton(tp0));
         client.prepareResponse(newMetadataResponse(topicName, Errors.NONE));
 
-        Map> allTopics = fetcher.getAllTopicMetadata(5000L);
+        Map> allTopics = fetcher.getAllTopicMetadata(time.timer(5000L));
 
-        assertEquals(cluster.topics().size(), allTopics.size());
+        assertEquals(initialUpdateResponse.topicMetadata().size(), allTopics.size());
     }
 
     @Test
     public void testGetAllTopicsDisconnect() {
         // first try gets a disconnect, next succeeds
+        buildFetcher();
+        assignFromUser(singleton(tp0));
         client.prepareResponse(null, true);
         client.prepareResponse(newMetadataResponse(topicName, Errors.NONE));
-        Map> allTopics = fetcher.getAllTopicMetadata(5000L);
-        assertEquals(cluster.topics().size(), allTopics.size());
+        Map> allTopics = fetcher.getAllTopicMetadata(time.timer(5000L));
+        assertEquals(initialUpdateResponse.topicMetadata().size(), allTopics.size());
     }
 
     @Test(expected = TimeoutException.class)
     public void testGetAllTopicsTimeout() {
         // since no response is prepared, the request should timeout
-        fetcher.getAllTopicMetadata(50L);
+        buildFetcher();
+        assignFromUser(singleton(tp0));
+        fetcher.getAllTopicMetadata(time.timer(50L));
     }
 
     @Test
     public void testGetAllTopicsUnauthorized() {
+        buildFetcher();
+        assignFromUser(singleton(tp0));
         client.prepareResponse(newMetadataResponse(topicName, Errors.TOPIC_AUTHORIZATION_FAILED));
         try {
-            fetcher.getAllTopicMetadata(10L);
+            fetcher.getAllTopicMetadata(time.timer(10L));
             fail();
         } catch (TopicAuthorizationException e) {
             assertEquals(singleton(topicName), e.unauthorizedTopics());
@@ -1097,68 +1971,132 @@ public void testGetAllTopicsUnauthorized() {
 
     @Test(expected = InvalidTopicException.class)
     public void testGetTopicMetadataInvalidTopic() {
+        buildFetcher();
+        assignFromUser(singleton(tp0));
         client.prepareResponse(newMetadataResponse(topicName, Errors.INVALID_TOPIC_EXCEPTION));
         fetcher.getTopicMetadata(
-                new MetadataRequest.Builder(Collections.singletonList(topicName), true), 5000L);
+                new MetadataRequest.Builder(Collections.singletonList(topicName), true), time.timer(5000L));
     }
 
     @Test
     public void testGetTopicMetadataUnknownTopic() {
+        buildFetcher();
+        assignFromUser(singleton(tp0));
         client.prepareResponse(newMetadataResponse(topicName, Errors.UNKNOWN_TOPIC_OR_PARTITION));
 
         Map> topicMetadata = fetcher.getTopicMetadata(
-                new MetadataRequest.Builder(Collections.singletonList(topicName), true), 5000L);
+                new MetadataRequest.Builder(Collections.singletonList(topicName), true), time.timer(5000L));
         assertNull(topicMetadata.get(topicName));
     }
 
     @Test
     public void testGetTopicMetadataLeaderNotAvailable() {
+        buildFetcher();
+        assignFromUser(singleton(tp0));
         client.prepareResponse(newMetadataResponse(topicName, Errors.LEADER_NOT_AVAILABLE));
         client.prepareResponse(newMetadataResponse(topicName, Errors.NONE));
 
         Map> topicMetadata = fetcher.getTopicMetadata(
-                new MetadataRequest.Builder(Collections.singletonList(topicName), true), 5000L);
+                new MetadataRequest.Builder(Collections.singletonList(topicName), true), time.timer(5000L));
         assertTrue(topicMetadata.containsKey(topicName));
     }
 
+    @Test
+    public void testGetTopicMetadataOfflinePartitions() {
+        buildFetcher();
+        assignFromUser(singleton(tp0));
+        MetadataResponse originalResponse = newMetadataResponse(topicName, Errors.NONE); //baseline ok response
+
+        //create a response based on the above one with all partitions being leaderless
+        List altTopics = new ArrayList<>();
+        for (MetadataResponse.TopicMetadata item : originalResponse.topicMetadata()) {
+            List partitions = item.partitionMetadata();
+            List altPartitions = new ArrayList<>();
+            for (MetadataResponse.PartitionMetadata p : partitions) {
+                altPartitions.add(new MetadataResponse.PartitionMetadata(
+                    p.error,
+                    p.topicPartition,
+                    Optional.empty(), //no leader
+                    Optional.empty(),
+                    p.replicaIds,
+                    p.inSyncReplicaIds,
+                    p.offlineReplicaIds
+                ));
+            }
+            MetadataResponse.TopicMetadata alteredTopic = new MetadataResponse.TopicMetadata(
+                item.error(),
+                item.topic(),
+                item.isInternal(),
+                altPartitions
+            );
+            altTopics.add(alteredTopic);
+        }
+        Node controller = originalResponse.controller();
+        MetadataResponse altered = RequestTestUtils.metadataResponse(
+            originalResponse.brokers(),
+            originalResponse.clusterId(),
+            controller != null ? controller.id() : MetadataResponse.NO_CONTROLLER_ID,
+            altTopics);
+
+        client.prepareResponse(altered);
+
+        Map> topicMetadata =
+            fetcher.getTopicMetadata(new MetadataRequest.Builder(Collections.singletonList(topicName), false),
+                    time.timer(5000L));
+
+        Assert.assertNotNull(topicMetadata);
+        Assert.assertNotNull(topicMetadata.get(topicName));
+        //noinspection ConstantConditions
+        Assert.assertEquals(metadata.fetch().partitionCountForTopic(topicName).longValue(), topicMetadata.get(topicName).size());
+    }
+
     /*
      * Send multiple requests. Verify that the client side quota metrics have the right values
      */
     @Test
-    public void testQuotaMetrics() throws Exception {
+    public void testQuotaMetrics() {
+        buildFetcher();
+
         MockSelector selector = new MockSelector(time);
         Sensor throttleTimeSensor = Fetcher.throttleTimeSensor(metrics, metricsRegistry);
         Cluster cluster = TestUtils.singletonCluster("test", 1);
         Node node = cluster.nodes().get(0);
         NetworkClient client = new NetworkClient(selector, metadata, "mock", Integer.MAX_VALUE,
-                1000, 1000, 64 * 1024, 64 * 1024, 1000,
+                1000, 1000, 64 * 1024, 64 * 1024, 1000, 10 * 1000, 127 * 1000, ClientDnsLookup.USE_ALL_DNS_IPS,
                 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 = RequestTestUtils.serializeResponseWithHeader(ApiVersionsResponse.createApiVersionsResponse(
+            400, RecordBatch.CURRENT_MAGIC_VALUE), ApiKeys.API_VERSIONS.latestVersion(), 0);
+
         selector.delayedReceive(new DelayedReceive(node.idString(), new NetworkReceive(node.idString(), buffer)));
-        while (!client.ready(node, time.milliseconds()))
+        while (!client.ready(node, time.milliseconds())) {
             client.poll(1, time.milliseconds());
+            // If a throttled response is received, advance the time to ensure progress.
+            time.sleep(client.throttleDelayMs(node, time.milliseconds()));
+        }
         selector.clear();
 
         for (int i = 1; i <= 3; i++) {
             int throttleTimeMs = 100 * i;
-            FetchRequest.Builder builder = FetchRequest.Builder.forConsumer(100, 100, new LinkedHashMap());
-            ClientRequest request = client.newClientRequest(node.idString(), builder, time.milliseconds(), true, null);
+            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 = fetchResponse(tp0, nextRecords, Errors.NONE, i, throttleTimeMs);
-            buffer = response.serialize(ApiKeys.FETCH.latestVersion(), new ResponseHeader(request.correlationId()));
+            FetchResponse response = fullFetchResponse(tp0, nextRecords, Errors.NONE, i, throttleTimeMs);
+            buffer = RequestTestUtils.serializeResponseWithHeader(response, 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.
+            time.sleep(client.throttleDelayMs(node, time.milliseconds()));
             selector.clear();
         }
         Map allMetrics = metrics.metrics();
         KafkaMetric avgMetric = allMetrics.get(metrics.metricInstance(metricsRegistry.fetchThrottleTimeAvg));
         KafkaMetric maxMetric = allMetrics.get(metrics.metricInstance(metricsRegistry.fetchThrottleTimeMax));
         // Throttle times are ApiVersions=400, Fetch=(100, 200, 300)
-        assertEquals(250, avgMetric.value(), EPSILON);
-        assertEquals(400, maxMetric.value(), EPSILON);
+        assertEquals(250, (Double) avgMetric.metricValue(), EPSILON);
+        assertEquals(400, (Double) maxMetric.metricValue(), EPSILON);
         client.close();
     }
 
@@ -1167,24 +2105,28 @@ public void testQuotaMetrics() throws Exception {
      */
     @Test
     public void testFetcherMetrics() {
-        subscriptions.assignFromUser(singleton(tp0));
+        buildFetcher();
+        assignFromUser(singleton(tp0));
         subscriptions.seek(tp0, 0);
 
         MetricName maxLagMetric = metrics.metricInstance(metricsRegistry.recordsLagMax);
-        MetricName partitionLagMetric = metrics.metricName(tp0 + ".records-lag", metricGroup);
+        Map tags = new HashMap<>();
+        tags.put("topic", tp0.topic());
+        tags.put("partition", String.valueOf(tp0.partition()));
+        MetricName partitionLagMetric = metrics.metricName("records-lag", metricGroup, tags);
 
         Map allMetrics = metrics.metrics();
         KafkaMetric recordsFetchLagMax = allMetrics.get(maxLagMetric);
 
-        // recordsFetchLagMax should be initialized to negative infinity
-        assertEquals(Double.NEGATIVE_INFINITY, recordsFetchLagMax.value(), EPSILON);
+        // recordsFetchLagMax should be initialized to NaN
+        assertEquals(Double.NaN, (Double) recordsFetchLagMax.metricValue(), EPSILON);
 
         // recordsFetchLagMax should be hw - fetchOffset after receiving an empty FetchResponse
         fetchRecords(tp0, MemoryRecords.EMPTY, Errors.NONE, 100L, 0);
-        assertEquals(100, recordsFetchLagMax.value(), EPSILON);
+        assertEquals(100, (Double) recordsFetchLagMax.metricValue(), EPSILON);
 
         KafkaMetric partitionLag = allMetrics.get(partitionLagMetric);
-        assertEquals(100, partitionLag.value(), EPSILON);
+        assertEquals(100, (Double) partitionLag.metricValue(), EPSILON);
 
         // recordsFetchLagMax should be hw - offset of the last message after receiving a non-empty FetchResponse
         MemoryRecordsBuilder builder = MemoryRecords.builder(ByteBuffer.allocate(1024), CompressionType.NONE,
@@ -1192,38 +2134,84 @@ public void testFetcherMetrics() {
         for (int v = 0; v < 3; v++)
             builder.appendWithOffset(v, RecordBatch.NO_TIMESTAMP, "key".getBytes(), ("value-" + v).getBytes());
         fetchRecords(tp0, builder.build(), Errors.NONE, 200L, 0);
-        assertEquals(197, recordsFetchLagMax.value(), EPSILON);
-        assertEquals(197, partitionLag.value(), EPSILON);
+        assertEquals(197, (Double) recordsFetchLagMax.metricValue(), EPSILON);
+        assertEquals(197, (Double) partitionLag.metricValue(), EPSILON);
 
         // verify de-registration of partition lag
         subscriptions.unsubscribe();
+        fetcher.sendFetches();
         assertFalse(allMetrics.containsKey(partitionLagMetric));
     }
 
+    @Test
+    public void testFetcherLeadMetric() {
+        buildFetcher();
+
+        assignFromUser(singleton(tp0));
+        subscriptions.seek(tp0, 0);
+
+        MetricName minLeadMetric = metrics.metricInstance(metricsRegistry.recordsLeadMin);
+        Map tags = new HashMap<>(2);
+        tags.put("topic", tp0.topic());
+        tags.put("partition", String.valueOf(tp0.partition()));
+        MetricName partitionLeadMetric = metrics.metricName("records-lead", metricGroup, "", tags);
+
+        Map allMetrics = metrics.metrics();
+        KafkaMetric recordsFetchLeadMin = allMetrics.get(minLeadMetric);
+
+        // recordsFetchLeadMin should be initialized to NaN
+        assertEquals(Double.NaN, (Double) recordsFetchLeadMin.metricValue(), EPSILON);
+
+        // recordsFetchLeadMin should be position - logStartOffset after receiving an empty FetchResponse
+        fetchRecords(tp0, MemoryRecords.EMPTY, Errors.NONE, 100L, -1L, 0L, 0);
+        assertEquals(0L, (Double) recordsFetchLeadMin.metricValue(), EPSILON);
+
+        KafkaMetric partitionLead = allMetrics.get(partitionLeadMetric);
+        assertEquals(0L, (Double) partitionLead.metricValue(), EPSILON);
+
+        // recordsFetchLeadMin should be position - logStartOffset after receiving a non-empty FetchResponse
+        MemoryRecordsBuilder builder = MemoryRecords.builder(ByteBuffer.allocate(1024), CompressionType.NONE,
+                TimestampType.CREATE_TIME, 0L);
+        for (int v = 0; v < 3; v++) {
+            builder.appendWithOffset(v, RecordBatch.NO_TIMESTAMP, "key".getBytes(), ("value-" + v).getBytes());
+        }
+        fetchRecords(tp0, builder.build(), Errors.NONE, 200L, -1L, 0L, 0);
+        assertEquals(0L, (Double) recordsFetchLeadMin.metricValue(), EPSILON);
+        assertEquals(3L, (Double) partitionLead.metricValue(), EPSILON);
+
+        // verify de-registration of partition lag
+        subscriptions.unsubscribe();
+        fetcher.sendFetches();
+        assertFalse(allMetrics.containsKey(partitionLeadMetric));
+    }
+
     @Test
     public void testReadCommittedLagMetric() {
-        Metrics metrics = new Metrics();
-        fetcher = createFetcher(subscriptions, metrics, new ByteArrayDeserializer(),
+        buildFetcher(OffsetResetStrategy.EARLIEST, new ByteArrayDeserializer(),
                 new ByteArrayDeserializer(), Integer.MAX_VALUE, IsolationLevel.READ_COMMITTED);
 
-        subscriptions.assignFromUser(singleton(tp0));
+        assignFromUser(singleton(tp0));
         subscriptions.seek(tp0, 0);
 
         MetricName maxLagMetric = metrics.metricInstance(metricsRegistry.recordsLagMax);
-        MetricName partitionLagMetric = metrics.metricName(tp0 + ".records-lag", metricGroup);
+
+        Map tags = new HashMap<>();
+        tags.put("topic", tp0.topic());
+        tags.put("partition", String.valueOf(tp0.partition()));
+        MetricName partitionLagMetric = metrics.metricName("records-lag", metricGroup, tags);
 
         Map allMetrics = metrics.metrics();
         KafkaMetric recordsFetchLagMax = allMetrics.get(maxLagMetric);
 
-        // recordsFetchLagMax should be initialized to negative infinity
-        assertEquals(Double.NEGATIVE_INFINITY, recordsFetchLagMax.value(), EPSILON);
+        // recordsFetchLagMax should be initialized to NaN
+        assertEquals(Double.NaN, (Double) recordsFetchLagMax.metricValue(), EPSILON);
 
         // recordsFetchLagMax should be lso - fetchOffset after receiving an empty FetchResponse
         fetchRecords(tp0, MemoryRecords.EMPTY, Errors.NONE, 100L, 50L, 0);
-        assertEquals(50, recordsFetchLagMax.value(), EPSILON);
+        assertEquals(50, (Double) recordsFetchLagMax.metricValue(), EPSILON);
 
         KafkaMetric partitionLag = allMetrics.get(partitionLagMetric);
-        assertEquals(50, partitionLag.value(), EPSILON);
+        assertEquals(50, (Double) partitionLag.metricValue(), EPSILON);
 
         // recordsFetchLagMax should be lso - offset of the last message after receiving a non-empty FetchResponse
         MemoryRecordsBuilder builder = MemoryRecords.builder(ByteBuffer.allocate(1024), CompressionType.NONE,
@@ -1231,41 +2219,69 @@ public void testReadCommittedLagMetric() {
         for (int v = 0; v < 3; v++)
             builder.appendWithOffset(v, RecordBatch.NO_TIMESTAMP, "key".getBytes(), ("value-" + v).getBytes());
         fetchRecords(tp0, builder.build(), Errors.NONE, 200L, 150L, 0);
-        assertEquals(147, recordsFetchLagMax.value(), EPSILON);
-        assertEquals(147, partitionLag.value(), EPSILON);
+        assertEquals(147, (Double) recordsFetchLagMax.metricValue(), EPSILON);
+        assertEquals(147, (Double) partitionLag.metricValue(), EPSILON);
 
         // verify de-registration of partition lag
         subscriptions.unsubscribe();
+        fetcher.sendFetches();
         assertFalse(allMetrics.containsKey(partitionLagMetric));
     }
 
     @Test
     public void testFetchResponseMetrics() {
-        subscriptions.assignFromUser(singleton(tp0));
-        subscriptions.seek(tp0, 0);
+        buildFetcher();
 
-        Map allMetrics = metrics.metrics();
-        KafkaMetric fetchSizeAverage = allMetrics.get(metrics.metricInstance(metricsRegistry.fetchSizeAvg));
-        KafkaMetric recordsCountAverage = allMetrics.get(metrics.metricInstance(metricsRegistry.recordsPerRequestAvg));
+        String topic1 = "foo";
+        String topic2 = "bar";
+        TopicPartition tp1 = new TopicPartition(topic1, 0);
+        TopicPartition tp2 = new TopicPartition(topic2, 0);
 
-        MemoryRecordsBuilder builder = MemoryRecords.builder(ByteBuffer.allocate(1024), CompressionType.NONE,
-                TimestampType.CREATE_TIME, 0L);
-        for (int v = 0; v < 3; v++)
-            builder.appendWithOffset(v, RecordBatch.NO_TIMESTAMP, "key".getBytes(), ("value-" + v).getBytes());
-        MemoryRecords records = builder.build();
+        subscriptions.assignFromUser(Utils.mkSet(tp1, tp2));
+
+        Map partitionCounts = new HashMap<>();
+        partitionCounts.put(topic1, 1);
+        partitionCounts.put(topic2, 1);
+        client.updateMetadata(RequestTestUtils.metadataUpdateWith(1, partitionCounts, tp -> validLeaderEpoch));
 
         int expectedBytes = 0;
-        for (Record record : records.records())
-            expectedBytes += record.sizeInBytes();
+        LinkedHashMap> fetchPartitionData = new LinkedHashMap<>();
 
-        fetchRecords(tp0, records, Errors.NONE, 100L, 0);
-        assertEquals(expectedBytes, fetchSizeAverage.value(), EPSILON);
-        assertEquals(3, recordsCountAverage.value(), EPSILON);
+        for (TopicPartition tp : Utils.mkSet(tp1, tp2)) {
+            subscriptions.seek(tp, 0);
+
+            MemoryRecordsBuilder builder = MemoryRecords.builder(ByteBuffer.allocate(1024), CompressionType.NONE,
+                    TimestampType.CREATE_TIME, 0L);
+            for (int v = 0; v < 3; v++)
+                builder.appendWithOffset(v, RecordBatch.NO_TIMESTAMP, "key".getBytes(), ("value-" + v).getBytes());
+            MemoryRecords records = builder.build();
+            for (Record record : records.records())
+                expectedBytes += record.sizeInBytes();
+
+            fetchPartitionData.put(tp, new FetchResponse.PartitionData<>(Errors.NONE, 15L,
+                    FetchResponse.INVALID_LAST_STABLE_OFFSET, 0L, null, records));
+        }
+
+        assertEquals(1, fetcher.sendFetches());
+        client.prepareResponse(new FetchResponse<>(Errors.NONE, fetchPartitionData, 0, INVALID_SESSION_ID));
+        consumerClient.poll(time.timer(0));
+
+        Map>> fetchedRecords = fetchedRecords();
+        assertEquals(3, fetchedRecords.get(tp1).size());
+        assertEquals(3, fetchedRecords.get(tp2).size());
+
+        Map allMetrics = metrics.metrics();
+        KafkaMetric fetchSizeAverage = allMetrics.get(metrics.metricInstance(metricsRegistry.fetchSizeAvg));
+        KafkaMetric recordsCountAverage = allMetrics.get(metrics.metricInstance(metricsRegistry.recordsPerRequestAvg));
+        assertEquals(expectedBytes, (Double) fetchSizeAverage.metricValue(), EPSILON);
+        assertEquals(6, (Double) recordsCountAverage.metricValue(), EPSILON);
     }
 
     @Test
     public void testFetchResponseMetricsPartialResponse() {
-        subscriptions.assignFromUser(singleton(tp0));
+        buildFetcher();
+
+        assignFromUser(singleton(tp0));
         subscriptions.seek(tp0, 1);
 
         Map allMetrics = metrics.metrics();
@@ -1285,13 +2301,14 @@ public void testFetchResponseMetricsPartialResponse() {
         }
 
         fetchRecords(tp0, records, Errors.NONE, 100L, 0);
-        assertEquals(expectedBytes, fetchSizeAverage.value(), EPSILON);
-        assertEquals(2, recordsCountAverage.value(), EPSILON);
+        assertEquals(expectedBytes, (Double) fetchSizeAverage.metricValue(), EPSILON);
+        assertEquals(2, (Double) recordsCountAverage.metricValue(), EPSILON);
     }
 
     @Test
     public void testFetchResponseMetricsWithOnePartitionError() {
-        subscriptions.assignFromUser(Utils.mkSet(tp0, tp1));
+        buildFetcher();
+        assignFromUser(Utils.mkSet(tp0, tp1));
         subscriptions.seek(tp0, 0);
         subscriptions.seek(tp1, 0);
 
@@ -1305,28 +2322,31 @@ public void testFetchResponseMetricsWithOnePartitionError() {
             builder.appendWithOffset(v, RecordBatch.NO_TIMESTAMP, "key".getBytes(), ("value-" + v).getBytes());
         MemoryRecords records = builder.build();
 
-        Map partitions = new HashMap<>();
-        partitions.put(tp0, new FetchResponse.PartitionData(Errors.NONE, 100,
+        Map> partitions = new HashMap<>();
+        partitions.put(tp0, new FetchResponse.PartitionData<>(Errors.NONE, 100,
                 FetchResponse.INVALID_LAST_STABLE_OFFSET, 0L, null, records));
-        partitions.put(tp1, new FetchResponse.PartitionData(Errors.OFFSET_OUT_OF_RANGE, 100,
+        partitions.put(tp1, new FetchResponse.PartitionData<>(Errors.OFFSET_OUT_OF_RANGE, 100,
                 FetchResponse.INVALID_LAST_STABLE_OFFSET, 0L, null, MemoryRecords.EMPTY));
 
         assertEquals(1, fetcher.sendFetches());
-        client.prepareResponse(new FetchResponse(new LinkedHashMap<>(partitions), 0));
-        consumerClient.poll(0);
+        client.prepareResponse(new FetchResponse<>(Errors.NONE, new LinkedHashMap<>(partitions),
+                0, INVALID_SESSION_ID));
+        consumerClient.poll(time.timer(0));
         fetcher.fetchedRecords();
 
         int expectedBytes = 0;
         for (Record record : records.records())
             expectedBytes += record.sizeInBytes();
 
-        assertEquals(expectedBytes, fetchSizeAverage.value(), EPSILON);
-        assertEquals(3, recordsCountAverage.value(), EPSILON);
+        assertEquals(expectedBytes, (Double) fetchSizeAverage.metricValue(), EPSILON);
+        assertEquals(3, (Double) recordsCountAverage.metricValue(), EPSILON);
     }
 
     @Test
     public void testFetchResponseMetricsWithOnePartitionAtTheWrongOffset() {
-        subscriptions.assignFromUser(Utils.mkSet(tp0, tp1));
+        buildFetcher();
+
+        assignFromUser(Utils.mkSet(tp0, tp1));
         subscriptions.seek(tp0, 0);
         subscriptions.seek(tp1, 0);
 
@@ -1344,15 +2364,16 @@ public void testFetchResponseMetricsWithOnePartitionAtTheWrongOffset() {
             builder.appendWithOffset(v, RecordBatch.NO_TIMESTAMP, "key".getBytes(), ("value-" + v).getBytes());
         MemoryRecords records = builder.build();
 
-        Map partitions = new HashMap<>();
-        partitions.put(tp0, new FetchResponse.PartitionData(Errors.NONE, 100,
+        Map> partitions = new HashMap<>();
+        partitions.put(tp0, new FetchResponse.PartitionData<>(Errors.NONE, 100,
                 FetchResponse.INVALID_LAST_STABLE_OFFSET, 0L, null, records));
-        partitions.put(tp1, new FetchResponse.PartitionData(Errors.NONE, 100,
+        partitions.put(tp1, new FetchResponse.PartitionData<>(Errors.NONE, 100,
                 FetchResponse.INVALID_LAST_STABLE_OFFSET, 0L, null,
                 MemoryRecords.withRecords(CompressionType.NONE, new SimpleRecord("val".getBytes()))));
 
-        client.prepareResponse(new FetchResponse(new LinkedHashMap<>(partitions), 0));
-        consumerClient.poll(0);
+        client.prepareResponse(new FetchResponse<>(Errors.NONE, new LinkedHashMap<>(partitions),
+                0, INVALID_SESSION_ID));
+        consumerClient.poll(time.timer(0));
         fetcher.fetchedRecords();
 
         // we should have ignored the record at the wrong offset
@@ -1360,27 +2381,24 @@ public void testFetchResponseMetricsWithOnePartitionAtTheWrongOffset() {
         for (Record record : records.records())
             expectedBytes += record.sizeInBytes();
 
-        assertEquals(expectedBytes, fetchSizeAverage.value(), EPSILON);
-        assertEquals(3, recordsCountAverage.value(), EPSILON);
+        assertEquals(expectedBytes, (Double) fetchSizeAverage.metricValue(), EPSILON);
+        assertEquals(3, (Double) recordsCountAverage.metricValue(), EPSILON);
     }
 
     @Test
-    public void testFetcherMetricsTemplates() throws Exception {
-        metrics.close();
+    public void testFetcherMetricsTemplates() {
         Map clientTags = Collections.singletonMap("client-id", "clientA");
-        metrics = new Metrics(new MetricConfig().tags(clientTags));
-        metricsRegistry = new FetcherMetricsRegistry(clientTags.keySet(), "consumer" + groupId);
-        fetcher.close();
-        fetcher = createFetcher(subscriptions, metrics);
+        buildFetcher(new MetricConfig().tags(clientTags), OffsetResetStrategy.EARLIEST, new ByteArrayDeserializer(),
+                new ByteArrayDeserializer(), Integer.MAX_VALUE, IsolationLevel.READ_UNCOMMITTED);
 
         // Fetch from topic to generate topic metrics
-        subscriptions.assignFromUser(singleton(tp0));
+        assignFromUser(singleton(tp0));
         subscriptions.seek(tp0, 0);
         assertEquals(1, fetcher.sendFetches());
-        client.prepareResponse(fetchResponse(tp0, this.records, Errors.NONE, 100L, 0));
-        consumerClient.poll(0);
+        client.prepareResponse(fullFetchResponse(tp0, this.records, Errors.NONE, 100L, 0));
+        consumerClient.poll(time.timer(0));
         assertTrue(fetcher.hasCompletedFetches());
-        Map>> partitionRecords = fetcher.fetchedRecords();
+        Map>> partitionRecords = fetchedRecords();
         assertTrue(partitionRecords.containsKey(tp0));
 
         // Create throttle metrics
@@ -1404,25 +2422,32 @@ private Map>> fetchRecords(
     private Map>> fetchRecords(
             TopicPartition tp, MemoryRecords records, Errors error, long hw, long lastStableOffset, int throttleTime) {
         assertEquals(1, fetcher.sendFetches());
-        client.prepareResponse(fetchResponse(tp, records, error, hw, lastStableOffset, throttleTime));
-        consumerClient.poll(0);
-        return fetcher.fetchedRecords();
+        client.prepareResponse(fullFetchResponse(tp, records, error, hw, lastStableOffset, throttleTime));
+        consumerClient.poll(time.timer(0));
+        return fetchedRecords();
+    }
+
+    private Map>> fetchRecords(
+            TopicPartition tp, MemoryRecords records, Errors error, long hw, long lastStableOffset, long logStartOffset, int throttleTime) {
+        assertEquals(1, fetcher.sendFetches());
+        client.prepareResponse(fetchResponse(tp, records, error, hw, lastStableOffset, logStartOffset, throttleTime));
+        consumerClient.poll(time.timer(0));
+        return fetchedRecords();
     }
 
     @Test
     public void testGetOffsetsForTimesTimeout() {
-        try {
-            fetcher.getOffsetsByTimes(Collections.singletonMap(new TopicPartition(topicName, 2), 1000L), 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
     public void testGetOffsetsForTimes() {
+        buildFetcher();
+
         // Empty map
-        assertTrue(fetcher.getOffsetsByTimes(new HashMap(), 100L).isEmpty());
+        assertTrue(fetcher.offsetsForTimes(new HashMap<>(), time.timer(100L)).isEmpty());
         // Unknown Offset
         testGetOffsetsForTimesWithUnknownOffset();
         // Error code none with unknown offset
@@ -1430,69 +2455,344 @@ public void testGetOffsetsForTimes() {
         // Error code none with known offset
         testGetOffsetsForTimesWithError(Errors.NONE, Errors.NONE, 10L, 100L, 10L, 100L);
         // Test both of partition has error.
-        testGetOffsetsForTimesWithError(Errors.NOT_LEADER_FOR_PARTITION, Errors.INVALID_REQUEST, 10L, 100L, 10L, 100L);
+        testGetOffsetsForTimesWithError(Errors.NOT_LEADER_OR_FOLLOWER, Errors.INVALID_REQUEST, 10L, 100L, 10L, 100L);
         // Test the second partition has error.
-        testGetOffsetsForTimesWithError(Errors.NONE, Errors.NOT_LEADER_FOR_PARTITION, 10L, 100L, 10L, 100L);
+        testGetOffsetsForTimesWithError(Errors.NONE, Errors.NOT_LEADER_OR_FOLLOWER, 10L, 100L, 10L, 100L);
         // Test different errors.
-        testGetOffsetsForTimesWithError(Errors.NOT_LEADER_FOR_PARTITION, Errors.NONE, 10L, 100L, 10L, 100L);
+        testGetOffsetsForTimesWithError(Errors.NOT_LEADER_OR_FOLLOWER, Errors.NONE, 10L, 100L, 10L, 100L);
         testGetOffsetsForTimesWithError(Errors.UNKNOWN_TOPIC_OR_PARTITION, Errors.NONE, 10L, 100L, 10L, 100L);
         testGetOffsetsForTimesWithError(Errors.UNSUPPORTED_FOR_MESSAGE_FORMAT, Errors.NONE, 10L, 100L, null, 100L);
         testGetOffsetsForTimesWithError(Errors.BROKER_NOT_AVAILABLE, Errors.NONE, 10L, 100L, 10L, 100L);
     }
 
-    @Test(expected = TimeoutException.class)
-    public void testBatchedListOffsetsMetadataErrors() {
-        Map partitionData = new HashMap<>();
-        partitionData.put(tp0, new ListOffsetResponse.PartitionData(Errors.NOT_LEADER_FOR_PARTITION,
-                ListOffsetResponse.UNKNOWN_TIMESTAMP, ListOffsetResponse.UNKNOWN_OFFSET));
-        partitionData.put(tp1, new ListOffsetResponse.PartitionData(Errors.UNKNOWN_TOPIC_OR_PARTITION,
-                ListOffsetResponse.UNKNOWN_TIMESTAMP, ListOffsetResponse.UNKNOWN_OFFSET));
-        client.prepareResponse(new ListOffsetResponse(0, partitionData));
+    @Test
+    public void testGetOffsetsFencedLeaderEpoch() {
+        buildFetcher();
+        subscriptions.assignFromUser(singleton(tp0));
+        client.updateMetadata(initialUpdateResponse);
 
-        Map offsetsToSearch = new HashMap<>();
-        offsetsToSearch.put(tp0, ListOffsetRequest.EARLIEST_TIMESTAMP);
-        offsetsToSearch.put(tp1, ListOffsetRequest.EARLIEST_TIMESTAMP);
+        subscriptions.requestOffsetReset(tp0, OffsetResetStrategy.LATEST);
+
+        client.prepareResponse(listOffsetResponse(Errors.FENCED_LEADER_EPOCH, 1L, 5L));
+        fetcher.resetOffsetsIfNeeded();
+        consumerClient.pollNoWakeup();
 
-        fetcher.getOffsetsByTimes(offsetsToSearch, 0);
+        assertTrue(subscriptions.isOffsetResetNeeded(tp0));
+        assertFalse(subscriptions.isFetchable(tp0));
+        assertFalse(subscriptions.hasValidPosition(tp0));
+        assertEquals(0L, metadata.timeToNextUpdate(time.milliseconds()));
     }
 
     @Test
-    public void testSkippingAbortedTransactions() {
-        Fetcher fetcher = createFetcher(subscriptions, new Metrics(), new ByteArrayDeserializer(),
-                new ByteArrayDeserializer(), Integer.MAX_VALUE, IsolationLevel.READ_COMMITTED);
-        ByteBuffer buffer = ByteBuffer.allocate(1024);
-        int currentOffset = 0;
+    public void testGetOffsetByTimeWithPartitionsRetryCouldTriggerMetadataUpdate() {
+        List retriableErrors = Arrays.asList(Errors.NOT_LEADER_OR_FOLLOWER,
+            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 = RequestTestUtils.metadataUpdateWith("dummy", 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;
+            ListOffsetPartitionResponse tp0NoError = new ListOffsetPartitionResponse()
+                .setPartitionIndex(tp0.partition())
+                .setErrorCode(Errors.NONE.code())
+                .setTimestamp(fetchTimestamp)
+                .setOffset(4L);
+            List topics = Collections.singletonList(
+                    new ListOffsetTopicResponse()
+                        .setName(tp0.topic())
+                        .setPartitions(Arrays.asList(
+                                tp0NoError,
+                                new ListOffsetPartitionResponse()
+                                    .setPartitionIndex(tp1.partition())
+                                    .setErrorCode(retriableError.code())
+                                    .setTimestamp(ListOffsetRequest.LATEST_TIMESTAMP)
+                                    .setOffset(-1L))));
+            ListOffsetResponseData data = new ListOffsetResponseData()
+                    .setThrottleTimeMs(0)
+                    .setTopics(topics);
+
+            client.prepareResponseFrom(body -> {
+                boolean isListOffsetRequest = body instanceof ListOffsetRequest;
+                if (isListOffsetRequest) {
+                    ListOffsetRequest request = (ListOffsetRequest) body;
+                    List expectedTopics = Collections.singletonList(
+                            new ListOffsetTopic()
+                                .setName(tp0.topic())
+                                .setPartitions(Arrays.asList(
+                                        new ListOffsetPartition()
+                                            .setPartitionIndex(tp1.partition())
+                                            .setTimestamp(fetchTimestamp)
+                                            .setCurrentLeaderEpoch(ListOffsetResponse.UNKNOWN_EPOCH),
+                                        new ListOffsetPartition()
+                                            .setPartitionIndex(tp0.partition())
+                                            .setTimestamp(fetchTimestamp)
+                                            .setCurrentLeaderEpoch(ListOffsetResponse.UNKNOWN_EPOCH))));
+                    return request.topics().equals(expectedTopics);
+                } else {
+                    return false;
+                }
+            }, new ListOffsetResponse(data), 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.
+            List topicsWithFatalError = Collections.singletonList(
+                    new ListOffsetTopicResponse()
+                        .setName(tp0.topic())
+                        .setPartitions(Arrays.asList(
+                                tp0NoError,
+                                new ListOffsetPartitionResponse()
+                                    .setPartitionIndex(tp1.partition())
+                                    .setErrorCode(Errors.NOT_LEADER_OR_FOLLOWER.code())
+                                    .setTimestamp(ListOffsetRequest.LATEST_TIMESTAMP)
+                                    .setOffset(-1L))));
+            ListOffsetResponseData dataWithFatalError = new ListOffsetResponseData()
+                    .setThrottleTimeMs(0)
+                    .setTopics(topicsWithFatalError);
+            client.prepareResponseFrom(new ListOffsetResponse(dataWithFatalError), 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;
 
-        currentOffset += appendTransactionalRecords(buffer, 1L, currentOffset,
-                new SimpleRecord(time.milliseconds(), "key".getBytes(), "value".getBytes()),
-                new SimpleRecord(time.milliseconds(), "key".getBytes(), "value".getBytes()));
+                    ListOffsetTopic requestTopic = request.topics().get(0);
+                    ListOffsetPartition expectedPartition = new ListOffsetPartition()
+                            .setPartitionIndex(tp1.partition())
+                            .setTimestamp(fetchTimestamp)
+                            .setCurrentLeaderEpoch(newLeaderEpoch);
+                    return expectedPartition.equals(requestTopic.partitions().get(0));
+                } else {
+                    return false;
+                }
+            }, listOffsetResponse(tp1, Errors.NONE, fetchTimestamp, 5L), newLeader);
 
-        abortTransaction(buffer, 1L, currentOffset);
+            Map offsetAndTimestampMap =
+                fetcher.offsetsForTimes(
+                    Utils.mkMap(Utils.mkEntry(tp0, fetchTimestamp),
+                    Utils.mkEntry(tp1, fetchTimestamp)), time.timer(Integer.MAX_VALUE));
 
-        buffer.flip();
+            assertEquals(Utils.mkMap(
+                Utils.mkEntry(tp0, new OffsetAndTimestamp(4L, fetchTimestamp)),
+                Utils.mkEntry(tp1, new OffsetAndTimestamp(5L, fetchTimestamp))), offsetAndTimestampMap);
 
-        List abortedTransactions = new ArrayList<>();
-        abortedTransactions.add(new FetchResponse.AbortedTransaction(1, 0));
-        MemoryRecords records = MemoryRecords.readableRecords(buffer);
+            // 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();
         subscriptions.assignFromUser(singleton(tp0));
+        subscriptions.requestOffsetReset(tp0, OffsetResetStrategy.LATEST);
 
-        subscriptions.seek(tp0, 0);
+        client.prepareResponse(listOffsetResponse(Errors.UNKNOWN_LEADER_EPOCH, 1L, 5L));
+        fetcher.resetOffsetsIfNeeded();
+        consumerClient.pollNoWakeup();
+
+        assertTrue(subscriptions.isOffsetResetNeeded(tp0));
+        assertFalse(subscriptions.isFetchable(tp0));
+        assertFalse(subscriptions.hasValidPosition(tp0));
+        assertEquals(0L, metadata.timeToNextUpdate(time.milliseconds()));
+    }
+
+    @Test
+    public void testGetOffsetsIncludesLeaderEpoch() {
+        buildFetcher();
+        subscriptions.assignFromUser(singleton(tp0));
+
+        client.updateMetadata(initialUpdateResponse);
+
+        // Metadata update with leader epochs
+        MetadataResponse metadataResponse = RequestTestUtils.metadataUpdateWith("dummy", 1,
+                Collections.emptyMap(), Collections.singletonMap(topicName, 4), tp -> 99);
+        client.updateMetadata(metadataResponse);
+
+        // Request latest offset
+        subscriptions.requestOffsetReset(tp0);
+        fetcher.resetOffsetsIfNeeded();
+
+        // Check for epoch in outgoing request
+        MockClient.RequestMatcher matcher = body -> {
+            if (body instanceof ListOffsetRequest) {
+                ListOffsetRequest offsetRequest = (ListOffsetRequest) body;
+                int epoch = offsetRequest.topics().get(0).partitions().get(0).currentLeaderEpoch();
+                assertTrue("Expected Fetcher to set leader epoch in request", epoch != ListOffsetResponse.UNKNOWN_EPOCH);
+                assertEquals("Expected leader epoch to match epoch from metadata update", epoch, 99);
+                return true;
+            } else {
+                fail("Should have seen ListOffsetRequest");
+                return false;
+            }
+        };
+
+        client.prepareResponse(matcher, listOffsetResponse(Errors.NONE, 1L, 5L));
+        consumerClient.pollNoWakeup();
+    }
+
+    @Test
+    public void testGetOffsetsForTimesWhenSomeTopicPartitionLeadersNotKnownInitially() {
+        buildFetcher();
+
+        subscriptions.assignFromUser(Utils.mkSet(tp0, tp1));
+        final String anotherTopic = "another-topic";
+        final TopicPartition t2p0 = new TopicPartition(anotherTopic, 0);
+
+        client.reset();
+
+        // Metadata initially has one topic
+        MetadataResponse initialMetadata = RequestTestUtils.metadataUpdateWith(3, singletonMap(topicName, 2));
+        client.updateMetadata(initialMetadata);
+
+        // The first metadata refresh should contain one topic
+        client.prepareMetadataUpdate(initialMetadata);
+        client.prepareResponseFrom(listOffsetResponse(tp0, Errors.NONE, 1000L, 11L),
+                metadata.fetch().leaderFor(tp0));
+        client.prepareResponseFrom(listOffsetResponse(tp1, Errors.NONE, 1000L, 32L),
+                metadata.fetch().leaderFor(tp1));
+
+        // Second metadata refresh should contain two topics
+        Map partitionNumByTopic = new HashMap<>();
+        partitionNumByTopic.put(topicName, 2);
+        partitionNumByTopic.put(anotherTopic, 1);
+        MetadataResponse updatedMetadata = RequestTestUtils.metadataUpdateWith(3, partitionNumByTopic);
+        client.prepareMetadataUpdate(updatedMetadata);
+        client.prepareResponseFrom(listOffsetResponse(t2p0, Errors.NONE, 1000L, 54L),
+                metadata.fetch().leaderFor(t2p0));
+
+        Map timestampToSearch = new HashMap<>();
+        timestampToSearch.put(tp0, ListOffsetRequest.LATEST_TIMESTAMP);
+        timestampToSearch.put(tp1, ListOffsetRequest.LATEST_TIMESTAMP);
+        timestampToSearch.put(t2p0, ListOffsetRequest.LATEST_TIMESTAMP);
+        Map offsetAndTimestampMap =
+            fetcher.offsetsForTimes(timestampToSearch, time.timer(Long.MAX_VALUE));
+
+        assertNotNull("Expect Fetcher.offsetsForTimes() to return non-null result for " + tp0,
+                      offsetAndTimestampMap.get(tp0));
+        assertNotNull("Expect Fetcher.offsetsForTimes() to return non-null result for " + tp1,
+                      offsetAndTimestampMap.get(tp1));
+        assertNotNull("Expect Fetcher.offsetsForTimes() to return non-null result for " + t2p0,
+                      offsetAndTimestampMap.get(t2p0));
+        assertEquals(11L, offsetAndTimestampMap.get(tp0).offset());
+        assertEquals(32L, offsetAndTimestampMap.get(tp1).offset());
+        assertEquals(54L, offsetAndTimestampMap.get(t2p0).offset());
+    }
+
+    @Test
+    public void testGetOffsetsForTimesWhenSomeTopicPartitionLeadersDisconnectException() {
+        buildFetcher();
+        final String anotherTopic = "another-topic";
+        final TopicPartition t2p0 = new TopicPartition(anotherTopic, 0);
+        subscriptions.assignFromUser(Utils.mkSet(tp0, t2p0));
+
+        client.reset();
+
+        MetadataResponse initialMetadata = RequestTestUtils.metadataUpdateWith(1, singletonMap(topicName, 1));
+        client.updateMetadata(initialMetadata);
+
+        Map partitionNumByTopic = new HashMap<>();
+        partitionNumByTopic.put(topicName, 1);
+        partitionNumByTopic.put(anotherTopic, 1);
+        MetadataResponse updatedMetadata = RequestTestUtils.metadataUpdateWith(1, partitionNumByTopic);
+        client.prepareMetadataUpdate(updatedMetadata);
+
+        client.prepareResponse(listOffsetRequestMatcher(ListOffsetRequest.LATEST_TIMESTAMP),
+                listOffsetResponse(tp0, Errors.NONE, 1000L, 11L), true);
+        client.prepareResponseFrom(listOffsetResponse(tp0, Errors.NONE, 1000L, 11L), metadata.fetch().leaderFor(tp0));
+
+        Map timestampToSearch = new HashMap<>();
+        timestampToSearch.put(tp0, ListOffsetRequest.LATEST_TIMESTAMP);
+        Map offsetAndTimestampMap = fetcher.offsetsForTimes(timestampToSearch, time.timer(Long.MAX_VALUE));
+
+        assertNotNull("Expect Fetcher.offsetsForTimes() to return non-null result for " + tp0,
+                     offsetAndTimestampMap.get(tp0));
+        assertEquals(11L, offsetAndTimestampMap.get(tp0).offset());
+        Assert.assertNotNull(metadata.fetch().partitionCountForTopic(anotherTopic));
+    }
+
+    @Test(expected = TimeoutException.class)
+    public void testBatchedListOffsetsMetadataErrors() {
+        buildFetcher();
+
+        ListOffsetResponseData data = new ListOffsetResponseData()
+                .setThrottleTimeMs(0)
+                .setTopics(Collections.singletonList(new ListOffsetTopicResponse()
+                        .setName(tp0.topic())
+                        .setPartitions(Arrays.asList(
+                                new ListOffsetPartitionResponse()
+                                    .setPartitionIndex(tp0.partition())
+                                    .setErrorCode(Errors.NOT_LEADER_OR_FOLLOWER.code())
+                                    .setTimestamp(ListOffsetResponse.UNKNOWN_TIMESTAMP)
+                                    .setOffset(ListOffsetResponse.UNKNOWN_OFFSET),
+                                new ListOffsetPartitionResponse()
+                                    .setPartitionIndex(tp1.partition())
+                                    .setErrorCode(Errors.UNKNOWN_TOPIC_OR_PARTITION.code())
+                                    .setTimestamp(ListOffsetResponse.UNKNOWN_TIMESTAMP)
+                                    .setOffset(ListOffsetResponse.UNKNOWN_OFFSET)))));
+        client.prepareResponse(new ListOffsetResponse(data));
+
+        Map offsetsToSearch = new HashMap<>();
+        offsetsToSearch.put(tp0, ListOffsetRequest.EARLIEST_TIMESTAMP);
+        offsetsToSearch.put(tp1, ListOffsetRequest.EARLIEST_TIMESTAMP);
+
+        fetcher.offsetsForTimes(offsetsToSearch, time.timer(0));
+    }
+
+    @Test
+    public void testSkippingAbortedTransactions() {
+        buildFetcher(OffsetResetStrategy.EARLIEST, new ByteArrayDeserializer(),
+                new ByteArrayDeserializer(), Integer.MAX_VALUE, IsolationLevel.READ_COMMITTED);
+        ByteBuffer buffer = ByteBuffer.allocate(1024);
+        int currentOffset = 0;
+
+        currentOffset += appendTransactionalRecords(buffer, 1L, currentOffset,
+                new SimpleRecord(time.milliseconds(), "key".getBytes(), "value".getBytes()),
+                new SimpleRecord(time.milliseconds(), "key".getBytes(), "value".getBytes()));
+
+        abortTransaction(buffer, 1L, currentOffset);
+
+        buffer.flip();
+
+        List abortedTransactions = new ArrayList<>();
+        abortedTransactions.add(new FetchResponse.AbortedTransaction(1, 0));
+        MemoryRecords records = MemoryRecords.readableRecords(buffer);
+        assignFromUser(singleton(tp0));
+
+        subscriptions.seek(tp0, 0);
 
         // normal fetch
         assertEquals(1, fetcher.sendFetches());
         assertFalse(fetcher.hasCompletedFetches());
 
-        client.prepareResponse(fetchResponseWithAbortedTransactions(records, abortedTransactions, Errors.NONE, 100L, 100L, 0));
-        consumerClient.poll(0);
+        client.prepareResponse(fullFetchResponseWithAbortedTransactions(records, abortedTransactions, Errors.NONE, 100L, 100L, 0));
+        consumerClient.poll(time.timer(0));
         assertTrue(fetcher.hasCompletedFetches());
 
-        Map>> fetchedRecords = fetcher.fetchedRecords();
+        Map>> fetchedRecords = fetchedRecords();
         assertFalse(fetchedRecords.containsKey(tp0));
     }
 
     @Test
     public void testReturnCommittedTransactions() {
-        Fetcher fetcher = createFetcher(subscriptions, new Metrics(), new ByteArrayDeserializer(),
+        buildFetcher(OffsetResetStrategy.EARLIEST, new ByteArrayDeserializer(),
                 new ByteArrayDeserializer(), Integer.MAX_VALUE, IsolationLevel.READ_COMMITTED);
         ByteBuffer buffer = ByteBuffer.allocate(1024);
         int currentOffset = 0;
@@ -1501,38 +2801,35 @@ 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<>();
         MemoryRecords records = MemoryRecords.readableRecords(buffer);
-        subscriptions.assignFromUser(singleton(tp0));
+        assignFromUser(singleton(tp0));
 
         subscriptions.seek(tp0, 0);
 
         // 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;
-            }
-        }, fetchResponseWithAbortedTransactions(records, abortedTransactions, Errors.NONE, 100L, 100L, 0));
+        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(0);
+        consumerClient.poll(time.timer(0));
         assertTrue(fetcher.hasCompletedFetches());
 
-        Map>> fetchedRecords = fetcher.fetchedRecords();
+        Map>> fetchedRecords = fetchedRecords();
         assertTrue(fetchedRecords.containsKey(tp0));
         assertEquals(fetchedRecords.get(tp0).size(), 2);
     }
 
     @Test
     public void testReadCommittedWithCommittedAndAbortedTransactions() {
-        Fetcher fetcher = createFetcher(subscriptions, new Metrics(), new ByteArrayDeserializer(),
+        buildFetcher(OffsetResetStrategy.EARLIEST, new ByteArrayDeserializer(),
                 new ByteArrayDeserializer(), Integer.MAX_VALUE, IsolationLevel.READ_COMMITTED);
         ByteBuffer buffer = ByteBuffer.allocate(1024);
 
@@ -1583,7 +2880,7 @@ public void testReadCommittedWithCommittedAndAbortedTransactions() {
         buffer.flip();
 
         MemoryRecords records = MemoryRecords.readableRecords(buffer);
-        subscriptions.assignFromUser(singleton(tp0));
+        assignFromUser(singleton(tp0));
 
         subscriptions.seek(tp0, 0);
 
@@ -1591,11 +2888,11 @@ public void testReadCommittedWithCommittedAndAbortedTransactions() {
         assertEquals(1, fetcher.sendFetches());
         assertFalse(fetcher.hasCompletedFetches());
 
-        client.prepareResponse(fetchResponseWithAbortedTransactions(records, abortedTransactions, Errors.NONE, 100L, 100L, 0));
-        consumerClient.poll(0);
+        client.prepareResponse(fullFetchResponseWithAbortedTransactions(records, abortedTransactions, Errors.NONE, 100L, 100L, 0));
+        consumerClient.poll(time.timer(0));
         assertTrue(fetcher.hasCompletedFetches());
 
-        Map>> fetchedRecords = fetcher.fetchedRecords();
+        Map>> fetchedRecords = fetchedRecords();
         assertTrue(fetchedRecords.containsKey(tp0));
         // There are only 3 committed records
         List> fetchedConsumerRecords = fetchedRecords.get(tp0);
@@ -1608,7 +2905,7 @@ public void testReadCommittedWithCommittedAndAbortedTransactions() {
 
     @Test
     public void testMultipleAbortMarkers() {
-        Fetcher fetcher = createFetcher(subscriptions, new Metrics(), new ByteArrayDeserializer(),
+        buildFetcher(OffsetResetStrategy.EARLIEST, new ByteArrayDeserializer(),
                 new ByteArrayDeserializer(), Integer.MAX_VALUE, IsolationLevel.READ_COMMITTED);
         ByteBuffer buffer = ByteBuffer.allocate(1024);
         int currentOffset = 0;
@@ -1630,7 +2927,7 @@ public void testMultipleAbortMarkers() {
         List abortedTransactions = new ArrayList<>();
         abortedTransactions.add(new FetchResponse.AbortedTransaction(1, 0));
         MemoryRecords records = MemoryRecords.readableRecords(buffer);
-        subscriptions.assignFromUser(singleton(tp0));
+        assignFromUser(singleton(tp0));
 
         subscriptions.seek(tp0, 0);
 
@@ -1638,11 +2935,11 @@ public void testMultipleAbortMarkers() {
         assertEquals(1, fetcher.sendFetches());
         assertFalse(fetcher.hasCompletedFetches());
 
-        client.prepareResponse(fetchResponseWithAbortedTransactions(records, abortedTransactions, Errors.NONE, 100L, 100L, 0));
-        consumerClient.poll(0);
+        client.prepareResponse(fullFetchResponseWithAbortedTransactions(records, abortedTransactions, Errors.NONE, 100L, 100L, 0));
+        consumerClient.poll(time.timer(0));
         assertTrue(fetcher.hasCompletedFetches());
 
-        Map>> fetchedRecords = fetcher.fetchedRecords();
+        Map>> fetchedRecords = fetchedRecords();
         assertTrue(fetchedRecords.containsKey(tp0));
         assertEquals(fetchedRecords.get(tp0).size(), 2);
         List> fetchedConsumerRecords = fetchedRecords.get(tp0);
@@ -1651,12 +2948,12 @@ public void testMultipleAbortMarkers() {
         for (ConsumerRecord consumerRecord : fetchedConsumerRecords) {
             actuallyCommittedKeys.add(new String(consumerRecord.key(), StandardCharsets.UTF_8));
         }
-        assertTrue(actuallyCommittedKeys.equals(committedKeys));
+        assertEquals(actuallyCommittedKeys, committedKeys);
     }
 
     @Test
     public void testReadCommittedAbortMarkerWithNoData() {
-        Fetcher fetcher = createFetcher(subscriptions, new Metrics(), new StringDeserializer(),
+        buildFetcher(OffsetResetStrategy.EARLIEST, new StringDeserializer(),
                 new StringDeserializer(), Integer.MAX_VALUE, IsolationLevel.READ_COMMITTED);
         ByteBuffer buffer = ByteBuffer.allocate(1024);
 
@@ -1674,7 +2971,7 @@ public void testReadCommittedAbortMarkerWithNoData() {
         buffer.flip();
 
         // send the fetch
-        subscriptions.assignFromUser(singleton(tp0));
+        assignFromUser(singleton(tp0));
         subscriptions.seek(tp0, 0);
         assertEquals(1, fetcher.sendFetches());
 
@@ -1682,12 +2979,12 @@ public void testReadCommittedAbortMarkerWithNoData() {
         List abortedTransactions = new ArrayList<>();
         abortedTransactions.add(new FetchResponse.AbortedTransaction(producerId, 0L));
 
-        client.prepareResponse(fetchResponseWithAbortedTransactions(MemoryRecords.readableRecords(buffer),
+        client.prepareResponse(fullFetchResponseWithAbortedTransactions(MemoryRecords.readableRecords(buffer),
                 abortedTransactions, Errors.NONE, 100L, 100L, 0));
-        consumerClient.poll(0);
+        consumerClient.poll(time.timer(0));
         assertTrue(fetcher.hasCompletedFetches());
 
-        Map>> allFetchedRecords = fetcher.fetchedRecords();
+        Map>> allFetchedRecords = fetchedRecords();
         assertTrue(allFetchedRecords.containsKey(tp0));
         List> fetchedRecords = allFetchedRecords.get(tp0);
         assertEquals(3, fetchedRecords.size());
@@ -1696,6 +2993,8 @@ public void testReadCommittedAbortMarkerWithNoData() {
 
     @Test
     public void testUpdatePositionWithLastRecordMissingFromBatch() {
+        buildFetcher();
+
         MemoryRecords records = MemoryRecords.withRecords(CompressionType.NONE,
                 new SimpleRecord("0".getBytes(), "v".getBytes()),
                 new SimpleRecord("1".getBytes(), "v".getBytes()),
@@ -1714,17 +3013,17 @@ protected boolean shouldRetainRecord(RecordBatch recordBatch, Record record) {
                 return record.key() != null;
             }
         }, ByteBuffer.allocate(1024), Integer.MAX_VALUE, BufferSupplier.NO_CACHING);
-        result.output.flip();
-        MemoryRecords compactedRecords = MemoryRecords.readableRecords(result.output);
+        result.outputBuffer().flip();
+        MemoryRecords compactedRecords = MemoryRecords.readableRecords(result.outputBuffer());
 
-        subscriptions.assignFromUser(singleton(tp0));
+        assignFromUser(singleton(tp0));
         subscriptions.seek(tp0, 0);
         assertEquals(1, fetcher.sendFetches());
-        client.prepareResponse(fetchResponse(tp0, compactedRecords, Errors.NONE, 100L, 0));
-        consumerClient.poll(0);
+        client.prepareResponse(fullFetchResponse(tp0, compactedRecords, Errors.NONE, 100L, 0));
+        consumerClient.poll(time.timer(0));
         assertTrue(fetcher.hasCompletedFetches());
 
-        Map>> allFetchedRecords = fetcher.fetchedRecords();
+        Map>> allFetchedRecords = fetchedRecords();
         assertTrue(allFetchedRecords.containsKey(tp0));
         List> fetchedRecords = allFetchedRecords.get(tp0);
         assertEquals(3, fetchedRecords.size());
@@ -1734,11 +3033,13 @@ 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
     public void testUpdatePositionOnEmptyBatch() {
+        buildFetcher();
+
         long producerId = 1;
         short producerEpoch = 0;
         int sequence = 1;
@@ -1752,23 +3053,23 @@ public void testUpdatePositionOnEmptyBatch() {
         buffer.flip();
         MemoryRecords recordsWithEmptyBatch = MemoryRecords.readableRecords(buffer);
 
-        subscriptions.assignFromUser(singleton(tp0));
+        assignFromUser(singleton(tp0));
         subscriptions.seek(tp0, 0);
         assertEquals(1, fetcher.sendFetches());
-        client.prepareResponse(fetchResponse(tp0, recordsWithEmptyBatch, Errors.NONE, 100L, 0));
-        consumerClient.poll(0);
+        client.prepareResponse(fullFetchResponse(tp0, recordsWithEmptyBatch, Errors.NONE, 100L, 0));
+        consumerClient.poll(time.timer(0));
         assertTrue(fetcher.hasCompletedFetches());
 
-        Map>> allFetchedRecords = fetcher.fetchedRecords();
+        Map>> allFetchedRecords = fetchedRecords();
         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
     public void testReadCommittedWithCompactedTopic() {
-        Fetcher fetcher = createFetcher(subscriptions, new Metrics(), new StringDeserializer(),
+        buildFetcher(OffsetResetStrategy.EARLIEST, new StringDeserializer(),
                 new StringDeserializer(), Integer.MAX_VALUE, IsolationLevel.READ_COMMITTED);
         ByteBuffer buffer = ByteBuffer.allocate(1024);
 
@@ -1807,7 +3108,7 @@ public void testReadCommittedWithCompactedTopic() {
         buffer.flip();
 
         // send the fetch
-        subscriptions.assignFromUser(singleton(tp0));
+        assignFromUser(singleton(tp0));
         subscriptions.seek(tp0, 0);
         assertEquals(1, fetcher.sendFetches());
 
@@ -1816,12 +3117,12 @@ public void testReadCommittedWithCompactedTopic() {
         abortedTransactions.add(new FetchResponse.AbortedTransaction(pid2, 6L));
         abortedTransactions.add(new FetchResponse.AbortedTransaction(pid1, 0L));
 
-        client.prepareResponse(fetchResponseWithAbortedTransactions(MemoryRecords.readableRecords(buffer),
+        client.prepareResponse(fullFetchResponseWithAbortedTransactions(MemoryRecords.readableRecords(buffer),
                 abortedTransactions, Errors.NONE, 100L, 100L, 0));
-        consumerClient.poll(0);
+        consumerClient.poll(time.timer(0));
         assertTrue(fetcher.hasCompletedFetches());
 
-        Map>> allFetchedRecords = fetcher.fetchedRecords();
+        Map>> allFetchedRecords = fetchedRecords();
         assertTrue(allFetchedRecords.containsKey(tp0));
         List> fetchedRecords = allFetchedRecords.get(tp0);
         assertEquals(5, fetchedRecords.size());
@@ -1830,7 +3131,7 @@ public void testReadCommittedWithCompactedTopic() {
 
     @Test
     public void testReturnAbortedTransactionsinUncommittedMode() {
-        Fetcher fetcher = createFetcher(subscriptions, new Metrics(), new ByteArrayDeserializer(),
+        buildFetcher(OffsetResetStrategy.EARLIEST, new ByteArrayDeserializer(),
                 new ByteArrayDeserializer(), Integer.MAX_VALUE, IsolationLevel.READ_UNCOMMITTED);
         ByteBuffer buffer = ByteBuffer.allocate(1024);
         int currentOffset = 0;
@@ -1846,7 +3147,7 @@ public void testReturnAbortedTransactionsinUncommittedMode() {
         List abortedTransactions = new ArrayList<>();
         abortedTransactions.add(new FetchResponse.AbortedTransaction(1, 0));
         MemoryRecords records = MemoryRecords.readableRecords(buffer);
-        subscriptions.assignFromUser(singleton(tp0));
+        assignFromUser(singleton(tp0));
 
         subscriptions.seek(tp0, 0);
 
@@ -1854,17 +3155,17 @@ public void testReturnAbortedTransactionsinUncommittedMode() {
         assertEquals(1, fetcher.sendFetches());
         assertFalse(fetcher.hasCompletedFetches());
 
-        client.prepareResponse(fetchResponseWithAbortedTransactions(records, abortedTransactions, Errors.NONE, 100L, 100L, 0));
-        consumerClient.poll(0);
+        client.prepareResponse(fullFetchResponseWithAbortedTransactions(records, abortedTransactions, Errors.NONE, 100L, 100L, 0));
+        consumerClient.poll(time.timer(0));
         assertTrue(fetcher.hasCompletedFetches());
 
-        Map>> fetchedRecords = fetcher.fetchedRecords();
+        Map>> fetchedRecords = fetchedRecords();
         assertTrue(fetchedRecords.containsKey(tp0));
     }
 
     @Test
     public void testConsumerPositionUpdatedWhenSkippingAbortedTransactions() {
-        Fetcher fetcher = createFetcher(subscriptions, new Metrics(), new ByteArrayDeserializer(),
+        buildFetcher(OffsetResetStrategy.EARLIEST, new ByteArrayDeserializer(),
                 new ByteArrayDeserializer(), Integer.MAX_VALUE, IsolationLevel.READ_COMMITTED);
         ByteBuffer buffer = ByteBuffer.allocate(1024);
         long currentOffset = 0;
@@ -1879,7 +3180,7 @@ public void testConsumerPositionUpdatedWhenSkippingAbortedTransactions() {
         List abortedTransactions = new ArrayList<>();
         abortedTransactions.add(new FetchResponse.AbortedTransaction(1, 0));
         MemoryRecords records = MemoryRecords.readableRecords(buffer);
-        subscriptions.assignFromUser(singleton(tp0));
+        assignFromUser(singleton(tp0));
 
         subscriptions.seek(tp0, 0);
 
@@ -1887,15 +3188,321 @@ public void testConsumerPositionUpdatedWhenSkippingAbortedTransactions() {
         assertEquals(1, fetcher.sendFetches());
         assertFalse(fetcher.hasCompletedFetches());
 
-        client.prepareResponse(fetchResponseWithAbortedTransactions(records, abortedTransactions, Errors.NONE, 100L, 100L, 0));
-        consumerClient.poll(0);
+        client.prepareResponse(fullFetchResponseWithAbortedTransactions(records, abortedTransactions, Errors.NONE, 100L, 100L, 0));
+        consumerClient.poll(time.timer(0));
         assertTrue(fetcher.hasCompletedFetches());
 
-        Map>> fetchedRecords = fetcher.fetchedRecords();
+        Map>> fetchedRecords = fetchedRecords();
 
         // 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
+    public void testConsumingViaIncrementalFetchRequests() {
+        buildFetcher(2);
+
+        List> records;
+        assignFromUser(new HashSet<>(Arrays.asList(tp0, tp1)));
+        subscriptions.seekValidated(tp0, new SubscriptionState.FetchPosition(0, Optional.empty(), metadata.currentLeader(tp0)));
+        subscriptions.seekValidated(tp1, new SubscriptionState.FetchPosition(1, Optional.empty(), metadata.currentLeader(tp1)));
+
+        // Fetch some records and establish an incremental fetch session.
+        LinkedHashMap> partitions1 = new LinkedHashMap<>();
+        partitions1.put(tp0, new FetchResponse.PartitionData<>(Errors.NONE, 2L,
+                2, 0L, null, this.records));
+        partitions1.put(tp1, new FetchResponse.PartitionData<>(Errors.NONE, 100L,
+                FetchResponse.INVALID_LAST_STABLE_OFFSET, 0L, null, emptyRecords));
+        FetchResponse resp1 = new FetchResponse<>(Errors.NONE, partitions1, 0, 123);
+        client.prepareResponse(resp1);
+        assertEquals(1, fetcher.sendFetches());
+        assertFalse(fetcher.hasCompletedFetches());
+        consumerClient.poll(time.timer(0));
+        assertTrue(fetcher.hasCompletedFetches());
+        Map>> fetchedRecords = fetchedRecords();
+        assertFalse(fetchedRecords.containsKey(tp1));
+        records = fetchedRecords.get(tp0);
+        assertEquals(2, records.size());
+        assertEquals(3L, subscriptions.position(tp0).offset);
+        assertEquals(1L, subscriptions.position(tp1).offset);
+        assertEquals(1, records.get(0).offset());
+        assertEquals(2, records.get(1).offset());
+
+        // There is still a buffered record.
+        assertEquals(0, fetcher.sendFetches());
+        fetchedRecords = fetchedRecords();
+        assertFalse(fetchedRecords.containsKey(tp1));
+        records = fetchedRecords.get(tp0);
+        assertEquals(1, records.size());
+        assertEquals(3, records.get(0).offset());
+        assertEquals(4L, subscriptions.position(tp0).offset);
+
+        // The second response contains no new records.
+        LinkedHashMap> partitions2 = new LinkedHashMap<>();
+        FetchResponse resp2 = new FetchResponse<>(Errors.NONE, partitions2, 0, 123);
+        client.prepareResponse(resp2);
+        assertEquals(1, fetcher.sendFetches());
+        consumerClient.poll(time.timer(0));
+        fetchedRecords = fetchedRecords();
+        assertTrue(fetchedRecords.isEmpty());
+        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<>();
+        partitions3.put(tp0, new FetchResponse.PartitionData<>(Errors.NONE, 100L,
+                4, 0L, null, this.nextRecords));
+        FetchResponse resp3 = new FetchResponse<>(Errors.NONE, partitions3, 0, 123);
+        client.prepareResponse(resp3);
+        assertEquals(1, fetcher.sendFetches());
+        consumerClient.poll(time.timer(0));
+        fetchedRecords = fetchedRecords();
+        assertFalse(fetchedRecords.containsKey(tp1));
+        records = fetchedRecords.get(tp0);
+        assertEquals(2, records.size());
+        assertEquals(6L, subscriptions.position(tp0).offset);
+        assertEquals(1L, subscriptions.position(tp1).offset);
+        assertEquals(4, records.get(0).offset());
+        assertEquals(5, records.get(1).offset());
+    }
+
+    @Test
+    public void testFetcherConcurrency() throws Exception {
+        int numPartitions = 20;
+        Set topicPartitions = new HashSet<>();
+        for (int i = 0; i < numPartitions; i++)
+            topicPartitions.add(new TopicPartition(topicName, i));
+
+        LogContext logContext = new LogContext();
+        buildDependencies(new MetricConfig(), Long.MAX_VALUE, new SubscriptionState(logContext, OffsetResetStrategy.EARLIEST), logContext);
+
+        fetcher = new Fetcher(
+                new LogContext(),
+                consumerClient,
+                minBytes,
+                maxBytes,
+                maxWaitMs,
+                fetchSize,
+                2 * numPartitions,
+                true,
+                "",
+                new ByteArrayDeserializer(),
+                new ByteArrayDeserializer(),
+                metadata,
+                subscriptions,
+                metrics,
+                metricsRegistry,
+                time,
+                retryBackoffMs,
+                requestTimeoutMs,
+                IsolationLevel.READ_UNCOMMITTED,
+                apiVersions) {
+            @Override
+            protected FetchSessionHandler sessionHandler(int id) {
+                final FetchSessionHandler handler = super.sessionHandler(id);
+                if (handler == null)
+                    return null;
+                else {
+                    return new FetchSessionHandler(new LogContext(), id) {
+                        @Override
+                        public Builder newBuilder() {
+                            verifySessionPartitions();
+                            return handler.newBuilder();
+                        }
+
+                        @Override
+                        public boolean handleResponse(FetchResponse response) {
+                            verifySessionPartitions();
+                            return handler.handleResponse(response);
+                        }
+
+                        @Override
+                        public void handleError(Throwable t) {
+                            verifySessionPartitions();
+                            handler.handleError(t);
+                        }
+
+                        // Verify that session partitions can be traversed safely.
+                        private void verifySessionPartitions() {
+                            try {
+                                Field field = FetchSessionHandler.class.getDeclaredField("sessionPartitions");
+                                field.setAccessible(true);
+                                LinkedHashMap sessionPartitions =
+                                        (LinkedHashMap) field.get(handler);
+                                for (Map.Entry entry : sessionPartitions.entrySet()) {
+                                    // If `sessionPartitions` are modified on another thread, Thread.yield will increase the
+                                    // possibility of ConcurrentModificationException if appropriate synchronization is not used.
+                                    Thread.yield();
+                                }
+                            } catch (Exception e) {
+                                throw new RuntimeException(e);
+                            }
+                        }
+                    };
+                }
+            }
+        };
+
+        MetadataResponse initialMetadataResponse = RequestTestUtils.metadataUpdateWith(1,
+                singletonMap(topicName, numPartitions), tp -> validLeaderEpoch);
+        client.updateMetadata(initialMetadataResponse);
+        fetchSize = 10000;
+
+        assignFromUser(topicPartitions);
+        topicPartitions.forEach(tp -> subscriptions.seek(tp, 0L));
+
+        AtomicInteger fetchesRemaining = new AtomicInteger(1000);
+        executorService = Executors.newSingleThreadExecutor();
+        Future future = executorService.submit(() -> {
+            while (fetchesRemaining.get() > 0) {
+                synchronized (consumerClient) {
+                    if (!client.requests().isEmpty()) {
+                        ClientRequest request = client.requests().peek();
+                        FetchRequest fetchRequest = (FetchRequest) request.requestBuilder().build();
+                        LinkedHashMap> responseMap = new LinkedHashMap<>();
+                        for (Map.Entry entry : fetchRequest.fetchData().entrySet()) {
+                            TopicPartition tp = entry.getKey();
+                            long offset = entry.getValue().fetchOffset;
+                            responseMap.put(tp, new FetchResponse.PartitionData<>(Errors.NONE, offset + 2L, offset + 2,
+                                    0L, null, buildRecords(offset, 2, offset)));
+                        }
+                        client.respondToRequest(request, new FetchResponse<>(Errors.NONE, responseMap, 0, 123));
+                        consumerClient.poll(time.timer(0));
+                    }
+                }
+            }
+            return fetchesRemaining.get();
+        });
+        Map nextFetchOffsets = topicPartitions.stream()
+                .collect(Collectors.toMap(Function.identity(), t -> 0L));
+        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();
+                    fetchedRecords.forEach((tp, records) -> {
+                        assertEquals(2, records.size());
+                        long nextOffset = nextFetchOffsets.get(tp);
+                        assertEquals(nextOffset, records.get(0).offset());
+                        assertEquals(nextOffset + 1, records.get(1).offset());
+                        nextFetchOffsets.put(tp, nextOffset + 2);
+                    });
+                }
+            }
+        }
+        assertEquals(0, future.get());
+    }
+
+    @Test
+    public void testFetcherSessionEpochUpdate() throws Exception {
+        buildFetcher(2);
+
+        MetadataResponse initialMetadataResponse = RequestTestUtils.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(),
+                new ByteArrayDeserializer(), Integer.MAX_VALUE, IsolationLevel.READ_COMMITTED);
+        ByteBuffer buffer = ByteBuffer.allocate(1024);
+        int currentOffset = 1;
+
+        // Empty control batch should not cause an exception
+        DefaultRecordBatch.writeEmptyHeader(buffer, RecordBatch.MAGIC_VALUE_V2, 1L,
+                (short) 0, -1, 0, 0,
+                RecordBatch.NO_PARTITION_LEADER_EPOCH, TimestampType.CREATE_TIME, time.milliseconds(),
+                true, true);
+
+        currentOffset += appendTransactionalRecords(buffer, 1L, currentOffset,
+                new SimpleRecord(time.milliseconds(), "key".getBytes(), "value".getBytes()),
+                new SimpleRecord(time.milliseconds(), "key".getBytes(), "value".getBytes()));
+
+        commitTransaction(buffer, 1L, currentOffset);
+        buffer.flip();
+
+        List abortedTransactions = new ArrayList<>();
+        MemoryRecords records = MemoryRecords.readableRecords(buffer);
+        assignFromUser(singleton(tp0));
+
+        subscriptions.seek(tp0, 0);
+
+        // normal fetch
+        assertEquals(1, fetcher.sendFetches());
+        assertFalse(fetcher.hasCompletedFetches());
+        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));
+        assertTrue(fetcher.hasCompletedFetches());
+
+        Map>> fetchedRecords = fetchedRecords();
+        assertTrue(fetchedRecords.containsKey(tp0));
+        assertEquals(fetchedRecords.get(tp0).size(), 2);
+    }
+
+    private MemoryRecords buildRecords(long baseOffset, int count, long firstMessageId) {
+        MemoryRecordsBuilder builder = MemoryRecords.builder(ByteBuffer.allocate(1024), CompressionType.NONE, TimestampType.CREATE_TIME, baseOffset);
+        for (int i = 0; i < count; i++)
+            builder.append(0L, "key".getBytes(), ("value-" + (firstMessageId + i)).getBytes());
+        return builder.build();
     }
 
     private int appendTransactionalRecords(ByteBuffer buffer, long pid, long baseOffset, int baseSequence, SimpleRecord... records) {
@@ -1914,12 +3521,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) {
@@ -1930,72 +3536,886 @@ private int abortTransaction(ByteBuffer buffer, long producerId, long baseOffset
         return 1;
     }
 
-    private void testGetOffsetsForTimesWithError(Errors errorForTp0,
-                                                 Errors errorForTp1,
-                                                 long offsetForTp0,
-                                                 long offsetForTp1,
-                                                 Long expectedOffsetForTp0,
-                                                 Long expectedOffsetForTp1) {
+    private void testGetOffsetsForTimesWithError(Errors errorForP0,
+                                                 Errors errorForP1,
+                                                 long offsetForP0,
+                                                 long offsetForP1,
+                                                 Long expectedOffsetForP0,
+                                                 Long expectedOffsetForP1) {
         client.reset();
-        // Ensure metadata has both partition.
-        Cluster cluster = TestUtils.clusterWith(2, topicName, 2);
-        metadata.update(cluster, Collections.emptySet(), time.milliseconds());
+        String topicName2 = "topic2";
+        TopicPartition t2p0 = new TopicPartition(topicName2, 0);
+        // Expect a metadata refresh.
+        metadata.bootstrap(ClientUtils.parseAndValidateAddresses(Collections.singletonList("1.1.1.1:1111"),
+                ClientDnsLookup.USE_ALL_DNS_IPS));
+
+        Map partitionNumByTopic = new HashMap<>();
+        partitionNumByTopic.put(topicName, 2);
+        partitionNumByTopic.put(topicName2, 1);
+        MetadataResponse updateMetadataResponse = RequestTestUtils.metadataUpdateWith(2, partitionNumByTopic);
+        Cluster updatedCluster = updateMetadataResponse.cluster();
+
+        // The metadata refresh should contain all the topics.
+        client.prepareMetadataUpdate(updateMetadataResponse, true);
 
         // First try should fail due to metadata error.
-        client.prepareResponseFrom(listOffsetResponse(tp0, errorForTp0, offsetForTp0, offsetForTp0), cluster.leaderFor(tp0));
-        client.prepareResponseFrom(listOffsetResponse(tp1, errorForTp1, offsetForTp1, offsetForTp1), cluster.leaderFor(tp1));
+        client.prepareResponseFrom(listOffsetResponse(t2p0, errorForP0, offsetForP0, offsetForP0),
+                updatedCluster.leaderFor(t2p0));
+        client.prepareResponseFrom(listOffsetResponse(tp1, errorForP1, offsetForP1, offsetForP1),
+                updatedCluster.leaderFor(tp1));
         // Second try should succeed.
-        client.prepareResponseFrom(listOffsetResponse(tp0, Errors.NONE, offsetForTp0, offsetForTp0), cluster.leaderFor(tp0));
-        client.prepareResponseFrom(listOffsetResponse(tp1, Errors.NONE, offsetForTp1, offsetForTp1), cluster.leaderFor(tp1));
+        client.prepareResponseFrom(listOffsetResponse(t2p0, Errors.NONE, offsetForP0, offsetForP0),
+                updatedCluster.leaderFor(t2p0));
+        client.prepareResponseFrom(listOffsetResponse(tp1, Errors.NONE, offsetForP1, offsetForP1),
+                updatedCluster.leaderFor(tp1));
 
         Map timestampToSearch = new HashMap<>();
-        timestampToSearch.put(tp0, 0L);
+        timestampToSearch.put(t2p0, 0L);
         timestampToSearch.put(tp1, 0L);
-        Map offsetAndTimestampMap = fetcher.getOffsetsByTimes(timestampToSearch, Long.MAX_VALUE);
+        Map offsetAndTimestampMap =
+                fetcher.offsetsForTimes(timestampToSearch, time.timer(Long.MAX_VALUE));
 
-        if (expectedOffsetForTp0 == null)
-            assertNull(offsetAndTimestampMap.get(tp0));
+        if (expectedOffsetForP0 == null)
+            assertNull(offsetAndTimestampMap.get(t2p0));
         else {
-            assertEquals(expectedOffsetForTp0.longValue(), offsetAndTimestampMap.get(tp0).timestamp());
-            assertEquals(expectedOffsetForTp0.longValue(), offsetAndTimestampMap.get(tp0).offset());
+            assertEquals(expectedOffsetForP0.longValue(), offsetAndTimestampMap.get(t2p0).timestamp());
+            assertEquals(expectedOffsetForP0.longValue(), offsetAndTimestampMap.get(t2p0).offset());
         }
 
-        if (expectedOffsetForTp1 == null)
+        if (expectedOffsetForP1 == null)
             assertNull(offsetAndTimestampMap.get(tp1));
         else {
-            assertEquals(expectedOffsetForTp1.longValue(), offsetAndTimestampMap.get(tp1).timestamp());
-            assertEquals(expectedOffsetForTp1.longValue(), offsetAndTimestampMap.get(tp1).offset());
+            assertEquals(expectedOffsetForP1.longValue(), offsetAndTimestampMap.get(tp1).timestamp());
+            assertEquals(expectedOffsetForP1.longValue(), offsetAndTimestampMap.get(tp1).offset());
         }
     }
 
     private void testGetOffsetsForTimesWithUnknownOffset() {
         client.reset();
-        // Ensure metadata has both partition.
-        Cluster cluster = TestUtils.clusterWith(1, topicName, 1);
-        metadata.update(cluster, Collections.emptySet(), time.milliseconds());
+        // Ensure metadata has both partitions.
+        MetadataResponse initialMetadataUpdate = RequestTestUtils.metadataUpdateWith(1, singletonMap(topicName, 1));
+        client.updateMetadata(initialMetadataUpdate);
+
+        ListOffsetResponseData data = new ListOffsetResponseData()
+                .setThrottleTimeMs(0)
+                .setTopics(Collections.singletonList(new ListOffsetTopicResponse()
+                        .setName(tp0.topic())
+                        .setPartitions(Collections.singletonList(new ListOffsetPartitionResponse()
+                                .setPartitionIndex(tp0.partition())
+                                .setErrorCode(Errors.NONE.code())
+                                .setTimestamp(ListOffsetResponse.UNKNOWN_TIMESTAMP)
+                                .setOffset(ListOffsetResponse.UNKNOWN_OFFSET)))));
+
+        client.prepareResponseFrom(new ListOffsetResponse(data),
+                metadata.fetch().leaderFor(tp0));
+
+        Map timestampToSearch = new HashMap<>();
+        timestampToSearch.put(tp0, 0L);
+        Map offsetAndTimestampMap =
+                fetcher.offsetsForTimes(timestampToSearch, time.timer(Long.MAX_VALUE));
 
-        Map partitionData = new HashMap<>();
-        partitionData.put(tp0, new ListOffsetResponse.PartitionData(Errors.NONE,
-                ListOffsetResponse.UNKNOWN_TIMESTAMP, ListOffsetResponse.UNKNOWN_OFFSET));
+        assertTrue(offsetAndTimestampMap.containsKey(tp0));
+        assertNull(offsetAndTimestampMap.get(tp0));
+    }
 
-        client.prepareResponseFrom(new ListOffsetResponse(0, partitionData), cluster.leaderFor(tp0));
+    @Test
+    public void testGetOffsetsForTimesWithUnknownOffsetV0() {
+        buildFetcher();
+        // Empty map
+        assertTrue(fetcher.offsetsForTimes(new HashMap<>(), time.timer(100L)).isEmpty());
+        // Unknown Offset
+        client.reset();
+        // Ensure metadata has both partition.
+        MetadataResponse initialMetadataUpdate = RequestTestUtils.metadataUpdateWith(1, singletonMap(topicName, 1));
+        client.updateMetadata(initialMetadataUpdate);
+        // Force LIST_OFFSETS version 0
+        Node node = metadata.fetch().nodes().get(0);
+        apiVersions.update(node.idString(), NodeApiVersions.create(
+            ApiKeys.LIST_OFFSETS.id, (short) 0, (short) 0));
+
+        ListOffsetResponseData data = new ListOffsetResponseData()
+                .setThrottleTimeMs(0)
+                .setTopics(Collections.singletonList(new ListOffsetTopicResponse()
+                        .setName(tp0.topic())
+                        .setPartitions(Collections.singletonList(new ListOffsetPartitionResponse()
+                                .setPartitionIndex(tp0.partition())
+                                .setErrorCode(Errors.NONE.code())
+                                .setTimestamp(ListOffsetResponse.UNKNOWN_TIMESTAMP)
+                                .setOldStyleOffsets(Collections.emptyList())))));
+
+        client.prepareResponseFrom(new ListOffsetResponse(data),
+                metadata.fetch().leaderFor(tp0));
 
         Map timestampToSearch = new HashMap<>();
         timestampToSearch.put(tp0, 0L);
-        Map offsetAndTimestampMap = fetcher.getOffsetsByTimes(timestampToSearch, Long.MAX_VALUE);
+        Map offsetAndTimestampMap =
+                fetcher.offsetsForTimes(timestampToSearch, time.timer(Long.MAX_VALUE));
 
         assertTrue(offsetAndTimestampMap.containsKey(tp0));
         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 = RequestTestUtils.metadataUpdateWith("dummy", 1, Collections.emptyMap(), partitionCounts, tp -> 1);
+        metadata.updateWithCurrentRequestVersion(metadataResponse, false, 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 testOffsetValidationRequestGrouping() {
+        buildFetcher();
+        assignFromUser(Utils.mkSet(tp0, tp1, tp2, tp3));
+
+        metadata.updateWithCurrentRequestVersion(RequestTestUtils.metadataUpdateWith("dummy", 3,
+            Collections.emptyMap(), singletonMap(topicName, 4),
+            tp -> 5), false, 0L);
+
+        for (TopicPartition tp : subscriptions.assignedPartitions()) {
+            Metadata.LeaderAndEpoch leaderAndEpoch = new Metadata.LeaderAndEpoch(
+                metadata.currentLeader(tp).leader, Optional.of(4));
+            subscriptions.seekUnvalidated(tp,
+                new SubscriptionState.FetchPosition(0, Optional.of(4), leaderAndEpoch));
+        }
+
+        Set allRequestedPartitions = new HashSet<>();
+
+        for (Node node : metadata.fetch().nodes()) {
+            apiVersions.update(node.idString(), NodeApiVersions.create());
+
+            Set expectedPartitions = subscriptions.assignedPartitions().stream()
+                .filter(tp ->
+                    metadata.currentLeader(tp).leader.equals(Optional.of(node)))
+                .collect(Collectors.toSet());
+
+            assertTrue(expectedPartitions.stream().noneMatch(allRequestedPartitions::contains));
+            assertTrue(expectedPartitions.size() > 0);
+            allRequestedPartitions.addAll(expectedPartitions);
+
+            OffsetForLeaderEpochResponseData data = new OffsetForLeaderEpochResponseData();
+            expectedPartitions.forEach(tp -> {
+                OffsetForLeaderTopicResult topic = data.topics().find(tp.topic());
+                if (topic == null) {
+                    topic = new OffsetForLeaderTopicResult().setTopic(tp.topic());
+                    data.topics().add(topic);
+                }
+                topic.partitions().add(new EpochEndOffset()
+                    .setPartition(tp.partition())
+                    .setErrorCode(Errors.NONE.code())
+                    .setLeaderEpoch(4)
+                    .setEndOffset(0));
+            });
+
+            OffsetsForLeaderEpochResponse response = new OffsetsForLeaderEpochResponse(data);
+            client.prepareResponseFrom(body -> {
+                OffsetsForLeaderEpochRequest request = (OffsetsForLeaderEpochRequest) body;
+                return expectedPartitions.equals(offsetForLeaderPartitionMap(request.data()).keySet());
+            }, response, node);
+        }
+
+        assertEquals(subscriptions.assignedPartitions(), allRequestedPartitions);
+
+        fetcher.validateOffsetsIfNeeded();
+        consumerClient.pollNoWakeup();
+
+        assertTrue(subscriptions.assignedPartitions()
+            .stream().noneMatch(subscriptions::awaitingValidation));
+    }
+
+    @Test
+    public void testOffsetValidationAwaitsNodeApiVersion() {
+        buildFetcher();
+        assignFromUser(singleton(tp0));
+
+        Map partitionCounts = new HashMap<>();
+        partitionCounts.put(tp0.topic(), 4);
+
+        final int epochOne = 1;
+
+        metadata.updateWithCurrentRequestVersion(RequestTestUtils.metadataUpdateWith("dummy", 1,
+                Collections.emptyMap(), partitionCounts, tp -> epochOne), false, 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.currentLeader(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
+        client.prepareResponseFrom(
+            prepareOffsetsForLeaderEpochResponse(tp0, Errors.NONE, epochOne, 30L),
+            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.updateWithCurrentRequestVersion(RequestTestUtils.metadataUpdateWith("dummy", 1,
+                Collections.emptyMap(), partitionCounts, tp -> epochOne), false, 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.currentLeader(tp0).leader, Optional.of(epochOne));
+            subscriptions.seekUnvalidated(tp0, new SubscriptionState.FetchPosition(0, Optional.of(epochOne), leaderAndEpoch));
+
+            // Update metadata to epoch=2, enter validation
+            metadata.updateWithCurrentRequestVersion(RequestTestUtils.metadataUpdateWith("dummy", 1,
+                    Collections.emptyMap(), partitionCounts, tp -> epochTwo), false, 0L);
+            fetcher.validateOffsetsIfNeeded();
+
+            // Offset validation is skipped
+            assertFalse(subscriptions.awaitingValidation(tp0));
+        }
+
+        {
+            // Seek with a position and leader+epoch
+            Metadata.LeaderAndEpoch leaderAndEpoch = new Metadata.LeaderAndEpoch(
+                    metadata.currentLeader(tp0).leader, Optional.of(epochOne));
+            subscriptions.seekUnvalidated(tp0, new SubscriptionState.FetchPosition(0, Optional.of(epochOne), leaderAndEpoch));
+
+            // Update metadata to epoch=2, enter validation
+            metadata.updateWithCurrentRequestVersion(RequestTestUtils.metadataUpdateWith("dummy", 1,
+                    Collections.emptyMap(), partitionCounts, tp -> epochTwo), false, 0L);
+
+            // Subscription should not stay in AWAITING_VALIDATION in prepareFetchRequest
+            assertEquals(1, fetcher.sendFetches());
+            assertFalse(subscriptions.awaitingValidation(tp0));
+        }
+    }
+
+    @Test
+    public void testOffsetValidationSkippedForOldResponse() {
+        // Old responses may provide unreliable leader epoch,
+        // 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;
+
+        metadata.updateWithCurrentRequestVersion(RequestTestUtils.metadataUpdateWith("dummy", 1,
+            Collections.emptyMap(), partitionCounts, tp -> epochOne), false, 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.currentLeader(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));
+
+        // Inject an older version of the metadata response
+        final short responseVersion = 8;
+        metadata.updateWithCurrentRequestVersion(RequestTestUtils.metadataUpdateWith("dummy", 1,
+            Collections.emptyMap(), partitionCounts, responseVersion), false, 0L);
+        fetcher.validateOffsetsIfNeeded();
+        // Offset validation is skipped
+        assertFalse(subscriptions.awaitingValidation(tp0));
+    }
+
+    @Test
+    public void testOffsetValidationResetOffsetForUndefinedEpochWithDefinedResetPolicy() {
+        testOffsetValidationWithGivenEpochOffset(
+            UNDEFINED_EPOCH, 0L, OffsetResetStrategy.EARLIEST);
+    }
+
+    @Test
+    public void testOffsetValidationResetOffsetForUndefinedOffsetWithDefinedResetPolicy() {
+        testOffsetValidationWithGivenEpochOffset(
+            2, UNDEFINED_EPOCH_OFFSET, OffsetResetStrategy.EARLIEST);
+    }
+
+    @Test
+    public void testOffsetValidationResetOffsetForUndefinedEpochWithUndefinedResetPolicy() {
+        testOffsetValidationWithGivenEpochOffset(
+            UNDEFINED_EPOCH, 0L, OffsetResetStrategy.NONE);
+    }
+
+    @Test
+    public void testOffsetValidationResetOffsetForUndefinedOffsetWithUndefinedResetPolicy() {
+        testOffsetValidationWithGivenEpochOffset(
+            2, UNDEFINED_EPOCH_OFFSET, OffsetResetStrategy.NONE);
+    }
+
+    @Test
+    public void testOffsetValidationTriggerLogTruncationForBadOffsetWithUndefinedResetPolicy() {
+        testOffsetValidationWithGivenEpochOffset(
+            1, 1L, OffsetResetStrategy.NONE);
+    }
+
+    private void testOffsetValidationWithGivenEpochOffset(int leaderEpoch,
+                                                          long endOffset,
+                                                          OffsetResetStrategy offsetResetStrategy) {
+        buildFetcher(offsetResetStrategy);
+        assignFromUser(singleton(tp0));
+
+        Map partitionCounts = new HashMap<>();
+        partitionCounts.put(tp0.topic(), 4);
+
+        final int epochOne = 1;
+        final long initialOffset = 5;
+
+        metadata.updateWithCurrentRequestVersion(RequestTestUtils.metadataUpdateWith("dummy", 1,
+            Collections.emptyMap(), partitionCounts, tp -> epochOne), false, 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.currentLeader(tp0).leader, Optional.of(epochOne));
+        subscriptions.seekUnvalidated(tp0, new SubscriptionState.FetchPosition(initialOffset, Optional.of(epochOne), leaderAndEpoch));
+
+        fetcher.validateOffsetsIfNeeded();
+
+        consumerClient.poll(time.timer(Duration.ZERO));
+        assertTrue(subscriptions.awaitingValidation(tp0));
+        assertTrue(client.hasInFlightRequests());
+
+        client.respond(
+            offsetsForLeaderEpochRequestMatcher(tp0, epochOne, epochOne),
+            prepareOffsetsForLeaderEpochResponse(tp0, Errors.NONE, leaderEpoch, endOffset));
+        consumerClient.poll(time.timer(Duration.ZERO));
+
+        if (offsetResetStrategy == OffsetResetStrategy.NONE) {
+            LogTruncationException thrown =
+                assertThrows(LogTruncationException.class, () -> fetcher.validateOffsetsIfNeeded());
+            assertEquals(singletonMap(tp0, initialOffset), thrown.offsetOutOfRangePartitions());
+
+            if (endOffset == UNDEFINED_EPOCH_OFFSET || leaderEpoch == UNDEFINED_EPOCH) {
+                assertEquals(Collections.emptyMap(), thrown.divergentOffsets());
+            } else {
+                OffsetAndMetadata expectedDivergentOffset = new OffsetAndMetadata(
+                    endOffset, Optional.of(leaderEpoch), "");
+                assertEquals(singletonMap(tp0, expectedDivergentOffset), thrown.divergentOffsets());
+            }
+            assertTrue(subscriptions.awaitingValidation(tp0));
+        } else {
+            fetcher.validateOffsetsIfNeeded();
+            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.updateWithCurrentRequestVersion(RequestTestUtils.metadataUpdateWith("dummy", 1,
+                Collections.emptyMap(), partitionCounts, tp -> epochOne), false, 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.currentLeader(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(
+            offsetsForLeaderEpochRequestMatcher(tp0, epochOne, epochOne),
+            prepareOffsetsForLeaderEpochResponse(tp0, Errors.NONE, 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.updateWithCurrentRequestVersion(RequestTestUtils.metadataUpdateWith("dummy", 1,
+                Collections.emptyMap(), partitionCounts, tp -> epochOne), false, 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.currentLeader(tp0).leader, Optional.of(epochOne));
+        subscriptions.seekValidated(tp0, new SubscriptionState.FetchPosition(0, Optional.of(epochOne), leaderAndEpoch));
+
+        // Update metadata to epoch=2, enter validation
+        metadata.updateWithCurrentRequestVersion(RequestTestUtils.metadataUpdateWith("dummy", 1,
+                Collections.emptyMap(), partitionCounts, tp -> epochTwo), false, 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(apiVersions, tp0, new Metadata.LeaderAndEpoch(leaderAndEpoch.leader, Optional.of(epochThree)));
+
+        // Prepare offset list response from async validation with epoch=2
+        client.prepareResponse(prepareOffsetsForLeaderEpochResponse(tp0, Errors.NONE, epochTwo, 10L));
+        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();
+        client.prepareResponse(prepareOffsetsForLeaderEpochResponse(tp0, Errors.NONE, epochThree, 10L));
+        consumerClient.pollNoWakeup();
+        assertFalse("Expected validation to succeed with latest epoch", subscriptions.awaitingValidation(tp0));
+    }
+
+    @Test
+    public void testSkipValidationForOlderApiVersion() {
+        buildFetcher();
+        assignFromUser(singleton(tp0));
+
+        Map partitionCounts = new HashMap<>();
+        partitionCounts.put(tp0.topic(), 4);
+
+        apiVersions.update("0", NodeApiVersions.create(ApiKeys.OFFSET_FOR_LEADER_EPOCH.id, (short) 0, (short) 2));
+
+        // Start with metadata, epoch=1
+        metadata.updateWithCurrentRequestVersion(RequestTestUtils.metadataUpdateWith("dummy", 1,
+                Collections.emptyMap(), partitionCounts, tp -> 1), false, 0L);
+
+        // Request offset reset
+        subscriptions.requestOffsetReset(tp0, OffsetResetStrategy.LATEST);
+
+        // Since we have no position due to reset, no fetch is sent
+        assertEquals(0, fetcher.sendFetches());
+
+        // Still no position, ensure offset validation logic did not transition us to FETCHING state
+        assertEquals(0, fetcher.sendFetches());
+
+        // Complete reset and now we can fetch
+        fetcher.resetOffsetIfNeeded(tp0, OffsetResetStrategy.LATEST,
+                new Fetcher.ListOffsetData(100, 1L, Optional.empty()));
+        assertEquals(1, fetcher.sendFetches());
+    }
+
+    @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 = RequestTestUtils.metadataUpdateWith("dummy", 1, Collections.emptyMap(), partitionCounts, tp -> 2);
+        metadata.updateWithCurrentRequestVersion(metadataResponse, false, 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.currentLeader(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.
+        client.prepareResponse(prepareOffsetsForLeaderEpochResponse(tp0, Errors.NONE, 1, 10L));
+        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(RequestTestUtils.metadataUpdateWith(2, singletonMap(topicName, 4), tp -> validLeaderEpoch));
+        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(RequestTestUtils.metadataUpdateWith(2, singletonMap(topicName, 4), tp -> validLeaderEpoch));
+
+        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();
+
+        Metadata.LeaderAndEpoch leaderAndEpoch = subscriptions.position(tp0).currentLeader;
+        assertTrue(leaderAndEpoch.leader.isPresent());
+        Node readReplica = fetcher.selectReadReplica(tp0, leaderAndEpoch.leader.get(), time.milliseconds());
+
+        AtomicBoolean wokenUp = new AtomicBoolean(false);
+        client.setWakeupHook(() -> {
+            if (!wokenUp.getAndSet(true)) {
+                consumerClient.disconnectAsync(readReplica);
+                consumerClient.poll(time.timer(0));
+            }
+        });
+
+        assertEquals(1, fetcher.sendFetches());
+
+        consumerClient.disconnectAsync(readReplica);
+        consumerClient.poll(time.timer(0));
+
+        assertEquals(1, fetcher.sendFetches());
+    }
+
+    @Test
+    public void testCorruptMessageError() {
+        buildFetcher();
+        assignFromUser(singleton(tp0));
+        subscriptions.seek(tp0, 0);
+
+        assertEquals(1, fetcher.sendFetches());
+        assertFalse(fetcher.hasCompletedFetches());
+
+        // Prepare a response with the CORRUPT_MESSAGE error.
+        client.prepareResponse(fullFetchResponse(
+                tp0,
+                buildRecords(1L, 1, 1),
+                Errors.CORRUPT_MESSAGE,
+                100L, 0));
+        consumerClient.poll(time.timer(0));
+        assertTrue(fetcher.hasCompletedFetches());
+
+        // Trigger the exception.
+        assertThrows(KafkaException.class, this::fetchedRecords);
+    }
+
+    @Test
+    public void testBeginningOffsets() {
+        buildFetcher();
+        assignFromUser(singleton(tp0));
+        client.prepareResponse(listOffsetResponse(tp0, Errors.NONE, ListOffsetRequest.EARLIEST_TIMESTAMP, 2L));
+        assertEquals(singletonMap(tp0, 2L), fetcher.beginningOffsets(singleton(tp0), time.timer(5000L)));
+    }
+
+    @Test
+    public void testBeginningOffsetsDuplicateTopicPartition() {
+        buildFetcher();
+        assignFromUser(singleton(tp0));
+        client.prepareResponse(listOffsetResponse(tp0, Errors.NONE, ListOffsetRequest.EARLIEST_TIMESTAMP, 2L));
+        assertEquals(singletonMap(tp0, 2L), fetcher.beginningOffsets(asList(tp0, tp0), time.timer(5000L)));
+    }
+
+    @Test
+    public void testBeginningOffsetsMultipleTopicPartitions() {
+        buildFetcher();
+        Map expectedOffsets = new HashMap<>();
+        expectedOffsets.put(tp0, 2L);
+        expectedOffsets.put(tp1, 4L);
+        expectedOffsets.put(tp2, 6L);
+        assignFromUser(expectedOffsets.keySet());
+        client.prepareResponse(listOffsetResponse(expectedOffsets, Errors.NONE, ListOffsetRequest.EARLIEST_TIMESTAMP, ListOffsetResponse.UNKNOWN_EPOCH));
+        assertEquals(expectedOffsets, fetcher.beginningOffsets(asList(tp0, tp1, tp2), time.timer(5000L)));
+    }
+
+    @Test
+    public void testBeginningOffsetsEmpty() {
+        buildFetcher();
+        assertEquals(emptyMap(), fetcher.beginningOffsets(emptyList(), time.timer(5000L)));
+    }
+
+    @Test
+    public void testEndOffsets() {
+        buildFetcher();
+        assignFromUser(singleton(tp0));
+        client.prepareResponse(listOffsetResponse(tp0, Errors.NONE, ListOffsetRequest.LATEST_TIMESTAMP, 5L));
+        assertEquals(singletonMap(tp0, 5L), fetcher.endOffsets(singleton(tp0), time.timer(5000L)));
+    }
+
+    @Test
+    public void testEndOffsetsDuplicateTopicPartition() {
+        buildFetcher();
+        assignFromUser(singleton(tp0));
+        client.prepareResponse(listOffsetResponse(tp0, Errors.NONE, ListOffsetRequest.LATEST_TIMESTAMP, 5L));
+        assertEquals(singletonMap(tp0, 5L), fetcher.endOffsets(asList(tp0, tp0), time.timer(5000L)));
+    }
+
+    @Test
+    public void testEndOffsetsMultipleTopicPartitions() {
+        buildFetcher();
+        Map expectedOffsets = new HashMap<>();
+        expectedOffsets.put(tp0, 5L);
+        expectedOffsets.put(tp1, 7L);
+        expectedOffsets.put(tp2, 9L);
+        assignFromUser(expectedOffsets.keySet());
+        client.prepareResponse(listOffsetResponse(expectedOffsets, Errors.NONE, ListOffsetRequest.LATEST_TIMESTAMP, ListOffsetResponse.UNKNOWN_EPOCH));
+        assertEquals(expectedOffsets, fetcher.endOffsets(asList(tp0, tp1, tp2), time.timer(5000L)));
+    }
+
+    @Test
+    public void testEndOffsetsEmpty() {
+        buildFetcher();
+        assertEquals(emptyMap(), fetcher.endOffsets(emptyList(), time.timer(5000L)));
+    }
+
+    private MockClient.RequestMatcher offsetsForLeaderEpochRequestMatcher(
+        TopicPartition topicPartition,
+        int currentLeaderEpoch,
+        int leaderEpoch
+    ) {
+        return request -> {
+            OffsetsForLeaderEpochRequest epochRequest = (OffsetsForLeaderEpochRequest) request;
+            OffsetForLeaderPartition partition = offsetForLeaderPartitionMap(epochRequest.data())
+                .get(topicPartition);
+            return partition != null
+                && partition.currentLeaderEpoch() == currentLeaderEpoch
+                && partition.leaderEpoch() == leaderEpoch;
+        };
+    }
+
+    private OffsetsForLeaderEpochResponse prepareOffsetsForLeaderEpochResponse(
+        TopicPartition topicPartition,
+        Errors error,
+        int leaderEpoch,
+        long endOffset
+    ) {
+        OffsetForLeaderEpochResponseData data = new OffsetForLeaderEpochResponseData();
+        data.topics().add(new OffsetForLeaderTopicResult()
+            .setTopic(topicPartition.topic())
+            .setPartitions(Collections.singletonList(new EpochEndOffset()
+                .setPartition(topicPartition.partition())
+                .setErrorCode(error.code())
+                .setLeaderEpoch(leaderEpoch)
+                .setEndOffset(endOffset))));
+        return new OffsetsForLeaderEpochResponse(data);
+    }
+
+    private Map offsetForLeaderPartitionMap(
+        OffsetForLeaderEpochRequestData data
+    ) {
+        Map result = new HashMap<>();
+        data.topics().forEach(topic ->
+            topic.partitions().forEach(partition ->
+                result.put(new TopicPartition(topic.topic(), partition.partition()), partition)));
+        return result;
+    }
+
     private MockClient.RequestMatcher listOffsetRequestMatcher(final long timestamp) {
+        return listOffsetRequestMatcher(timestamp, ListOffsetResponse.UNKNOWN_EPOCH);
+    }
+
+    private MockClient.RequestMatcher listOffsetRequestMatcher(final long timestamp, final int leaderEpoch) {
         // 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);
-            }
+        return body -> {
+            ListOffsetRequest req = (ListOffsetRequest) body;
+            ListOffsetTopic topic = req.topics().get(0);
+            ListOffsetPartition partition = topic.partitions().get(0);
+            return tp0.topic().equals(topic.name())
+                    && tp0.partition() == partition.partitionIndex()
+                    && timestamp == partition.timestamp()
+                    && leaderEpoch == partition.currentLeaderEpoch();
         };
     }
 
@@ -2004,78 +4424,165 @@ private ListOffsetResponse listOffsetResponse(Errors error, long timestamp, long
     }
 
     private ListOffsetResponse listOffsetResponse(TopicPartition tp, Errors error, long timestamp, long offset) {
-        ListOffsetResponse.PartitionData partitionData = new ListOffsetResponse.PartitionData(error, timestamp, offset);
-        Map allPartitionData = new HashMap<>();
-        allPartitionData.put(tp, partitionData);
-        return new ListOffsetResponse(allPartitionData);
+        return listOffsetResponse(tp, error, timestamp, offset, ListOffsetResponse.UNKNOWN_EPOCH);
+    }
+
+    private ListOffsetResponse listOffsetResponse(TopicPartition tp, Errors error, long timestamp, long offset, int leaderEpoch) {
+        Map offsets = new HashMap<>();
+        offsets.put(tp, offset);
+        return listOffsetResponse(offsets, error, timestamp, leaderEpoch);
+    }
+
+    private ListOffsetResponse listOffsetResponse(Map offsets, Errors error, long timestamp, int leaderEpoch) {
+        Map> responses = new HashMap<>();
+        for (Map.Entry entry : offsets.entrySet()) {
+            TopicPartition tp = entry.getKey();
+            responses.putIfAbsent(tp.topic(), new ArrayList<>());
+            responses.get(tp.topic()).add(new ListOffsetPartitionResponse()
+                                .setPartitionIndex(tp.partition())
+                                .setErrorCode(error.code())
+                                .setOffset(entry.getValue())
+                                .setTimestamp(timestamp)
+                                .setLeaderEpoch(leaderEpoch));
+        }
+        List topics = new ArrayList<>();
+        for (Map.Entry> response : responses.entrySet()) {
+            topics.add(new ListOffsetTopicResponse()
+                    .setName(response.getKey())
+                    .setPartitions(response.getValue()));
+        }
+        ListOffsetResponseData data = new ListOffsetResponseData().setTopics(topics);
+        return new ListOffsetResponse(data);
+    }
+
+    private FetchResponse fullFetchResponseWithAbortedTransactions(MemoryRecords records,
+                                                                                  List abortedTransactions,
+                                                                                  Errors error,
+                                                                                  long lastStableOffset,
+                                                                                  long hw,
+                                                                                  int throttleTime) {
+        Map> partitions = Collections.singletonMap(tp0,
+                new FetchResponse.PartitionData<>(error, hw, lastStableOffset, 0L, abortedTransactions, records));
+        return new FetchResponse<>(Errors.NONE, new LinkedHashMap<>(partitions), throttleTime, INVALID_SESSION_ID);
     }
 
-    private FetchResponse fetchResponseWithAbortedTransactions(MemoryRecords records,
-                                                               List abortedTransactions,
-                                                               Errors error,
-                                                               long lastStableOffset,
-                                                               long hw,
-                                                               int throttleTime) {
-        Map partitions = Collections.singletonMap(tp0,
-                new FetchResponse.PartitionData(error, hw, lastStableOffset, 0L, abortedTransactions, records));
-        return new FetchResponse(new LinkedHashMap<>(partitions), throttleTime);
+    private FetchResponse fullFetchResponse(TopicPartition tp, MemoryRecords records, Errors error, long hw, int throttleTime) {
+        return fullFetchResponse(tp, records, error, hw, FetchResponse.INVALID_LAST_STABLE_OFFSET, throttleTime);
     }
 
-    private FetchResponse fetchResponse(TopicPartition tp, MemoryRecords records, Errors error, long hw, int throttleTime) {
-        return fetchResponse(tp, records, error, hw, FetchResponse.INVALID_LAST_STABLE_OFFSET, throttleTime);
+    private FetchResponse fullFetchResponse(TopicPartition tp, MemoryRecords records, Errors error, long hw,
+                                            long lastStableOffset, int throttleTime) {
+        Map> partitions = Collections.singletonMap(tp,
+                new FetchResponse.PartitionData<>(error, hw, lastStableOffset, 0L, 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, int throttleTime) {
-        Map partitions = Collections.singletonMap(tp,
-                new FetchResponse.PartitionData(error, hw, lastStableOffset, 0L, null, records));
-        return new FetchResponse(new LinkedHashMap<>(partitions), throttleTime);
+    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,
+                new FetchResponse.PartitionData<>(error, hw, lastStableOffset, logStartOffset, null, records));
+        return new FetchResponse<>(Errors.NONE, new LinkedHashMap<>(partitions), throttleTime, INVALID_SESSION_ID);
     }
 
     private MetadataResponse newMetadataResponse(String topic, Errors error) {
         List partitionsMetadata = new ArrayList<>();
         if (error == Errors.NONE) {
-            for (PartitionInfo partitionInfo : cluster.partitionsForTopic(topic)) {
-                partitionsMetadata.add(new MetadataResponse.PartitionMetadata(
-                        Errors.NONE,
-                        partitionInfo.partition(),
-                        partitionInfo.leader(),
-                        Arrays.asList(partitionInfo.replicas()),
-                        Arrays.asList(partitionInfo.inSyncReplicas()),
-                        Arrays.asList(partitionInfo.offlineReplicas())));
-            }
+            Optional foundMetadata = initialUpdateResponse.topicMetadata()
+                    .stream()
+                    .filter(topicMetadata -> topicMetadata.topic().equals(topic))
+                    .findFirst();
+            foundMetadata.ifPresent(topicMetadata -> {
+                partitionsMetadata.addAll(topicMetadata.partitionMetadata());
+            });
         }
 
-        MetadataResponse.TopicMetadata topicMetadata = new MetadataResponse.TopicMetadata(error, topic, false, partitionsMetadata);
-        return new MetadataResponse(cluster.nodes(), null, MetadataResponse.NO_CONTROLLER_ID, Arrays.asList(topicMetadata));
+        MetadataResponse.TopicMetadata topicMetadata = new MetadataResponse.TopicMetadata(error, topic, false,
+                partitionsMetadata);
+        List brokers = new ArrayList<>(initialUpdateResponse.brokers());
+        return RequestTestUtils.metadataResponse(brokers, initialUpdateResponse.clusterId(),
+                initialUpdateResponse.controller().id(), Collections.singletonList(topicMetadata));
     }
 
-    private Fetcher createFetcher(SubscriptionState subscriptions,
-                                                  Metrics metrics,
-                                                  int maxPollRecords) {
-        return createFetcher(subscriptions, metrics, new ByteArrayDeserializer(), new ByteArrayDeserializer(),
+    @SuppressWarnings("unchecked")
+    private  Map>> fetchedRecords() {
+        return (Map) fetcher.fetchedRecords();
+    }
+
+    private void buildFetcher(int maxPollRecords) {
+        buildFetcher(OffsetResetStrategy.EARLIEST, new ByteArrayDeserializer(), new ByteArrayDeserializer(),
                 maxPollRecords, IsolationLevel.READ_UNCOMMITTED);
     }
 
-    private Fetcher createFetcher(SubscriptionState subscriptions, Metrics metrics) {
-        return createFetcher(subscriptions, metrics, Integer.MAX_VALUE);
+    private void buildFetcher() {
+        buildFetcher(Integer.MAX_VALUE);
+    }
+
+    private void buildFetcher(Deserializer keyDeserializer,
+                              Deserializer valueDeserializer) {
+        buildFetcher(OffsetResetStrategy.EARLIEST, keyDeserializer, valueDeserializer,
+                Integer.MAX_VALUE, IsolationLevel.READ_UNCOMMITTED);
+    }
+
+    private void buildFetcher(OffsetResetStrategy offsetResetStrategy) {
+        buildFetcher(new MetricConfig(), offsetResetStrategy,
+            new ByteArrayDeserializer(), new ByteArrayDeserializer(),
+            Integer.MAX_VALUE, IsolationLevel.READ_UNCOMMITTED);
+    }
+
+    private  void buildFetcher(OffsetResetStrategy offsetResetStrategy,
+                                     Deserializer keyDeserializer,
+                                     Deserializer valueDeserializer,
+                                     int maxPollRecords,
+                                     IsolationLevel isolationLevel) {
+        buildFetcher(new MetricConfig(), offsetResetStrategy, keyDeserializer, valueDeserializer,
+                maxPollRecords, isolationLevel);
+    }
+
+    private  void buildFetcher(MetricConfig metricConfig,
+                                     OffsetResetStrategy offsetResetStrategy,
+                                     Deserializer keyDeserializer,
+                                     Deserializer valueDeserializer,
+                                     int maxPollRecords,
+                                     IsolationLevel isolationLevel) {
+        buildFetcher(metricConfig, offsetResetStrategy, keyDeserializer, valueDeserializer, maxPollRecords, isolationLevel, Long.MAX_VALUE);
     }
 
-    private  Fetcher createFetcher(SubscriptionState subscriptions,
-                                               Metrics metrics,
-                                               Deserializer keyDeserializer,
-                                               Deserializer valueDeserializer) {
-        return createFetcher(subscriptions, metrics, keyDeserializer, valueDeserializer, Integer.MAX_VALUE,
-                IsolationLevel.READ_UNCOMMITTED);
+    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  Fetcher createFetcher(SubscriptionState subscriptions,
-                                               Metrics metrics,
-                                               Deserializer keyDeserializer,
-                                               Deserializer valueDeserializer,
-                                               int maxPollRecords,
-                                               IsolationLevel isolationLevel) {
-        return new Fetcher<>(
+    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,
                 minBytes,
@@ -2084,6 +4591,7 @@ private  Fetcher createFetcher(SubscriptionState subscriptions,
                 fetchSize,
                 maxPollRecords,
                 true, // check crc
+                "",
                 keyDeserializer,
                 valueDeserializer,
                 metadata,
@@ -2092,13 +4600,27 @@ private  Fetcher createFetcher(SubscriptionState subscriptions,
                 metricsRegistry,
                 time,
                 retryBackoffMs,
-                isolationLevel);
+                requestTimeoutMs,
+                isolationLevel,
+                apiVersions);
+    }
+
+    private void buildDependencies(MetricConfig metricConfig,
+                                   long metadataExpireMs,
+                                   SubscriptionState subscriptionState,
+                                   LogContext logContext) {
+        time = new MockTime(1);
+        subscriptions = subscriptionState;
+        metadata = new ConsumerMetadata(0, metadataExpireMs, false, false,
+                subscriptions, logContext, new ClusterResourceListeners());
+        client = new MockClient(time, metadata);
+        metrics = new Metrics(metricConfig, time);
+        consumerClient = new ConsumerNetworkClient(logContext, client, metadata, time,
+                100, 1000, Integer.MAX_VALUE);
+        metricsRegistry = new FetcherMetricsRegistry(metricConfig.tags().keySet(), "consumer" + groupId);
     }
 
     private  List collectRecordOffsets(List> records) {
-        List res = new ArrayList<>(records.size());
-        for (ConsumerRecord record : records)
-            res.add(record.offset());
-        return res;
+        return records.stream().map(ConsumerRecord::offset).collect(Collectors.toList());
     }
 }
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 06cdae7ad6942..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,57 +16,107 @@
  */
 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 long timeout = 300L;
-    private long interval = 100L;
-    private long maxPollInterval = 900L;
-    private long retryBackoff = 10L;
+    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(timeout, interval, maxPollInterval, retryBackoff);
+
+    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() {
         heartbeat.sentHeartbeat(time.milliseconds());
-        time.sleep((long) ((float) interval * 1.1));
+        time.sleep((long) ((float) heartbeatIntervalMs * 1.1));
         assertTrue(heartbeat.shouldHeartbeat(time.milliseconds()));
     }
 
     @Test
     public void testShouldNotHeartbeat() {
         heartbeat.sentHeartbeat(time.milliseconds());
-        time.sleep(interval / 2);
+        time.sleep(heartbeatIntervalMs / 2);
         assertFalse(heartbeat.shouldHeartbeat(time.milliseconds()));
     }
 
     @Test
     public void testTimeToNextHeartbeat() {
-        heartbeat.sentHeartbeat(0);
-        assertEquals(100, heartbeat.timeToNextHeartbeat(0));
-        assertEquals(0, heartbeat.timeToNextHeartbeat(100));
-        assertEquals(0, heartbeat.timeToNextHeartbeat(200));
+        heartbeat.sentHeartbeat(time.milliseconds());
+        assertEquals(heartbeatIntervalMs, heartbeat.timeToNextHeartbeat(time.milliseconds()));
+
+        time.sleep(heartbeatIntervalMs);
+        assertEquals(0, heartbeat.timeToNextHeartbeat(time.milliseconds()));
+
+        time.sleep(heartbeatIntervalMs);
+        assertEquals(0, heartbeat.timeToNextHeartbeat(time.milliseconds()));
     }
 
     @Test
     public void testSessionTimeoutExpired() {
         heartbeat.sentHeartbeat(time.milliseconds());
-        time.sleep(305);
+        time.sleep(sessionTimeoutMs + 5);
         assertTrue(heartbeat.sessionTimeoutExpired(time.milliseconds()));
     }
 
     @Test
     public void testResetSession() {
         heartbeat.sentHeartbeat(time.milliseconds());
-        time.sleep(305);
-        heartbeat.resetTimeouts(time.milliseconds());
+        time.sleep(sessionTimeoutMs + 5);
+        heartbeat.resetSessionTimeout();
         assertFalse(heartbeat.sessionTimeoutExpired(time.milliseconds()));
+
+        // Resetting the session timeout should not reset the poll timeout
+        time.sleep(maxPollIntervalMs + 1);
+        heartbeat.resetSessionTimeout();
+        assertTrue(heartbeat.pollTimeoutExpired(time.milliseconds()));
     }
+
+    @Test
+    public void testResetTimeouts() {
+        time.sleep(maxPollIntervalMs);
+        assertTrue(heartbeat.sessionTimeoutExpired(time.milliseconds()));
+        assertEquals(0, heartbeat.timeToNextHeartbeat(time.milliseconds()));
+        assertTrue(heartbeat.pollTimeoutExpired(time.milliseconds()));
+
+        heartbeat.resetTimeouts();
+        assertFalse(heartbeat.sessionTimeoutExpired(time.milliseconds()));
+        assertEquals(heartbeatIntervalMs, heartbeat.timeToNextHeartbeat(time.milliseconds()));
+        assertFalse(heartbeat.pollTimeoutExpired(time.milliseconds()));
+    }
+
+    @Test
+    public void testPollTimeout() {
+        assertFalse(heartbeat.pollTimeoutExpired(time.milliseconds()));
+        time.sleep(maxPollIntervalMs / 2);
+
+        assertFalse(heartbeat.pollTimeoutExpired(time.milliseconds()));
+        time.sleep(maxPollIntervalMs / 2 + 1);
+
+        assertTrue(heartbeat.pollTimeoutExpired(time.milliseconds()));
+    }
+
 }
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..ef95f2ffb5594 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
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.clients.consumer.internals;
 
+import org.apache.kafka.clients.consumer.ConsumerGroupMetadata;
 import org.apache.kafka.common.TopicPartition;
 
 import java.util.List;
@@ -23,8 +24,17 @@
 
 public class MockPartitionAssignor extends AbstractPartitionAssignor {
 
+    private final List supportedProtocols;
+
+    private int numAssignment;
+
     private Map> result = null;
 
+    MockPartitionAssignor(final List supportedProtocols) {
+        this.supportedProtocols = supportedProtocols;
+        numAssignment = 0;
+    }
+
     @Override
     public Map> assign(Map partitionsPerTopic,
                                                     Map subscriptions) {
@@ -38,6 +48,11 @@ public String name() {
         return "consumer-mock-assignor";
     }
 
+    @Override
+    public List supportedProtocols() {
+        return supportedProtocols;
+    }
+
     public void clear() {
         this.result = null;
     }
@@ -46,4 +61,12 @@ public void prepare(Map> result) {
         this.result = result;
     }
 
+    @Override
+    public void onAssignment(Assignment assignment, ConsumerGroupMetadata metadata) {
+        numAssignment += 1;
+    }
+
+    int numAssignment() {
+        return numAssignment;
+    }
 }
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MockRebalanceListener.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MockRebalanceListener.java
new file mode 100644
index 0000000000000..be802542e26a0
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/MockRebalanceListener.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.clients.consumer.internals;
+
+import java.util.Collection;
+import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
+import org.apache.kafka.common.TopicPartition;
+
+public 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;
+        assignedCount++;
+    }
+
+    @Override
+    public void onPartitionsRevoked(Collection partitions) {
+        this.revoked = 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/OffsetForLeaderEpochClientTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/OffsetForLeaderEpochClientTest.java
new file mode 100644
index 0000000000000..8774a5fba59a4
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/OffsetForLeaderEpochClientTest.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.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.message.OffsetForLeaderEpochResponseData;
+import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData.EpochEndOffset;
+import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData.OffsetForLeaderTopicResult;
+import org.apache.kafka.common.protocol.Errors;
+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(
+            new OffsetForLeaderEpochResponseData());
+        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(Optional.empty(), Optional.of(1))));
+
+        OffsetsForLeaderEpochClient offsetClient = newOffsetClient();
+        RequestFuture future =
+                offsetClient.sendAsyncRequest(Node.noNode(), positionMap);
+
+        OffsetsForLeaderEpochResponse resp = new OffsetsForLeaderEpochResponse(
+            new OffsetForLeaderEpochResponseData());
+        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(Optional.empty(), Optional.of(1))));
+
+        OffsetsForLeaderEpochClient offsetClient = newOffsetClient();
+        RequestFuture future =
+                offsetClient.sendAsyncRequest(Node.noNode(), positionMap);
+
+        client.prepareResponse(prepareOffsetForLeaderEpochResponse(
+            tp0, Errors.NONE, 1, 10L));
+        consumerClient.pollNoWakeup();
+
+        OffsetsForLeaderEpochClient.OffsetForEpochResult result = future.value();
+        assertTrue(result.partitionsToRetry().isEmpty());
+        assertTrue(result.endOffsets().containsKey(tp0));
+        assertEquals(result.endOffsets().get(tp0).errorCode(), Errors.NONE.code());
+        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(Optional.empty(), Optional.of(1))));
+
+        OffsetsForLeaderEpochClient offsetClient = newOffsetClient();
+        RequestFuture future =
+                offsetClient.sendAsyncRequest(Node.noNode(), positionMap);
+
+        client.prepareResponse(prepareOffsetForLeaderEpochResponse(
+            tp0, Errors.TOPIC_AUTHORIZATION_FAILED, -1, -1));
+        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(Optional.empty(), Optional.of(1))));
+
+        OffsetsForLeaderEpochClient offsetClient = newOffsetClient();
+        RequestFuture future =
+                offsetClient.sendAsyncRequest(Node.noNode(), positionMap);
+
+        client.prepareResponse(prepareOffsetForLeaderEpochResponse(
+            tp0, Errors.LEADER_NOT_AVAILABLE, -1, -1));
+        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);
+    }
+
+    private static OffsetsForLeaderEpochResponse prepareOffsetForLeaderEpochResponse(
+            TopicPartition tp, Errors error, int leaderEpoch, long endOffset) {
+        OffsetForLeaderEpochResponseData data = new OffsetForLeaderEpochResponseData();
+        OffsetForLeaderTopicResult topic = new OffsetForLeaderTopicResult()
+            .setTopic(tp.topic());
+        data.topics().add(topic);
+        topic.partitions().add(new EpochEndOffset()
+            .setPartition(tp.partition())
+            .setErrorCode(error.code())
+            .setLeaderEpoch(leaderEpoch)
+            .setEndOffset(endOffset));
+        return new OffsetsForLeaderEpochResponse(data);
+    }
+}
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..ba16a52a7d79e
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/PartitionAssignorAdapterTest.java
@@ -0,0 +1,171 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.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
+    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(""));
+
+        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 c0a2df96d2876..cbcdb822802e1 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,48 +16,61 @@
  */
 package org.apache.kafka.clients.consumer.internals;
 
+import org.apache.kafka.clients.ApiVersions;
+import org.apache.kafka.clients.Metadata;
+import org.apache.kafka.clients.NodeApiVersions;
 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.clients.consumer.internals.SubscriptionState.LogTruncation;
+import org.apache.kafka.common.Node;
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData.EpochEndOffset;
+import org.apache.kafka.common.protocol.ApiKeys;
+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.apache.kafka.common.requests.OffsetsForLeaderEpochResponse.UNDEFINED_EPOCH;
+import static org.apache.kafka.common.requests.OffsetsForLeaderEpochResponse.UNDEFINED_EPOCH_OFFSET;
 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 SubscriptionStateTest {
 
-    private final SubscriptionState state = new SubscriptionState(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 = Metadata.LeaderAndEpoch.noLeaderOrEpoch();
 
     @Test
     public void partitionAssignment() {
         state.assignFromUser(singleton(tp0));
         assertEquals(singleton(tp0), state.assignedPartitions());
+        assertEquals(1, state.numAssignedPartitions());
         assertFalse(state.hasAllFetchPositions());
-        assertTrue(state.refreshCommitsNeeded());
-        state.committed(tp0, new OffsetAndMetadata(1));
         state.seek(tp0, 1);
         assertTrue(state.isFetchable(tp0));
-        assertAllPositions(tp0, 1L);
-        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));
         assertFalse(state.isFetchable(tp0));
     }
@@ -67,28 +80,57 @@ public void partitionAssignmentChangeOnTopicSubscription() {
         state.assignFromUser(new HashSet<>(Arrays.asList(tp0, tp1)));
         // assigned partitions should immediately change
         assertEquals(2, state.assignedPartitions().size());
+        assertEquals(2, state.numAssignedPartitions());
         assertTrue(state.assignedPartitions().contains(tp0));
         assertTrue(state.assignedPartitions().contains(tp1));
 
         state.unsubscribe();
         // assigned partitions should immediately change
         assertTrue(state.assignedPartitions().isEmpty());
+        assertEquals(0, state.numAssignedPartitions());
 
         state.subscribe(singleton(topic1), rebalanceListener);
         // assigned partitions should remain unchanged
         assertTrue(state.assignedPartitions().isEmpty());
+        assertEquals(0, state.numAssignedPartitions());
 
+        assertTrue(state.checkAssignmentMatchedSubscription(singleton(t1p0)));
         state.assignFromSubscribed(singleton(t1p0));
         // assigned partitions should immediately change
         assertEquals(singleton(t1p0), state.assignedPartitions());
+        assertEquals(1, state.numAssignedPartitions());
 
         state.subscribe(singleton(topic), rebalanceListener);
         // assigned partitions should remain unchanged
         assertEquals(singleton(t1p0), state.assignedPartitions());
+        assertEquals(1, state.numAssignedPartitions());
 
         state.unsubscribe();
         // assigned partitions should immediately change
         assertTrue(state.assignedPartitions().isEmpty());
+        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
@@ -96,70 +138,82 @@ public void partitionAssignmentChangeOnPatternSubscription() {
         state.subscribe(Pattern.compile(".*"), rebalanceListener);
         // assigned partitions should remain unchanged
         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.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());
 
-        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());
         assertEquals(singleton(topic), state.subscription());
 
         state.subscribe(Pattern.compile(".*t"), rebalanceListener);
         // assigned partitions should remain unchanged
         assertEquals(singleton(t1p0), state.assignedPartitions());
+        assertEquals(1, state.numAssignedPartitions());
 
         state.subscribeFromPattern(singleton(topic));
         // assigned partitions should remain unchanged
         assertEquals(singleton(t1p0), state.assignedPartitions());
+        assertEquals(1, state.numAssignedPartitions());
+
+        assertTrue(state.checkAssignmentMatchedSubscription(singleton(tp0)));
+        state.assignFromSubscribed(singleton(tp0));
 
-        state.assignFromSubscribed(Collections.singletonList(tp0));
         // assigned partitions should immediately change
         assertEquals(singleton(tp0), state.assignedPartitions());
+        assertEquals(1, state.numAssignedPartitions());
         assertEquals(singleton(topic), state.subscription());
 
         state.unsubscribe();
         // assigned partitions should immediately change
         assertTrue(state.assignedPartitions().isEmpty());
+        assertEquals(0, state.numAssignedPartitions());
     }
 
     @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.checkAssignmentMatchedSubscription(autoAssignment));
         state.assignFromSubscribed(autoAssignment);
-        assertEquals(autoAssignment, assignmentRef.get());
+        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));
-        state.needOffsetReset(tp0);
+        assertEquals(5L, state.position(tp0).offset);
+        state.requestOffsetReset(tp0);
         assertFalse(state.isFetchable(tp0));
         assertTrue(state.isOffsetResetNeeded(tp0));
-        assertEquals(null, state.position(tp0));
+        assertNull(state.position(tp0));
 
         // seek should clear the reset and make the partition fetchable
         state.seek(tp0, 0);
@@ -172,16 +226,21 @@ public void topicSubscription() {
         state.subscribe(singleton(topic), rebalanceListener);
         assertEquals(1, state.subscription().size());
         assertTrue(state.assignedPartitions().isEmpty());
-        assertTrue(state.partitionsAutoAssigned());
+        assertEquals(0, state.numAssignedPartitions());
+        assertTrue(state.hasAutoAssignedPartitions());
+        assertTrue(state.checkAssignmentMatchedSubscription(singleton(tp0)));
         state.assignFromSubscribed(singleton(tp0));
+
         state.seek(tp0, 1);
-        state.committed(tp0, new OffsetAndMetadata(1));
-        assertAllPositions(tp0, 1L);
+        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));
         assertEquals(singleton(tp1), state.assignedPartitions());
+        assertEquals(1, state.numAssignedPartitions());
     }
 
     @Test
@@ -195,38 +254,31 @@ public void partitionPause() {
         assertTrue(state.isFetchable(tp0));
     }
 
-    @Test
-    public void commitOffsetMetadata() {
-        state.assignFromUser(singleton(tp0));
-        state.committed(tp0, new OffsetAndMetadata(5, "hi"));
-
-        assertEquals(5, state.committed(tp0).offset());
-        assertEquals("hi", state.committed(tp0).metadata());
-    }
-
     @Test(expected = IllegalStateException.class)
     public void invalidPositionUpdate() {
         state.subscribe(singleton(topic), rebalanceListener);
+        assertTrue(state.checkAssignmentMatchedSubscription(singleton(tp0)));
         state.assignFromSubscribed(singleton(tp0));
-        state.position(tp0, 0);
+
+        state.position(tp0, new SubscriptionState.FetchPosition(0, Optional.empty(), leaderAndEpoch));
     }
 
-    @Test(expected = IllegalArgumentException.class)
+    @Test
     public void cantAssignPartitionForUnsubscribedTopics() {
         state.subscribe(singleton(topic), rebalanceListener);
-        state.assignFromSubscribed(Collections.singletonList(t1p0));
+        assertFalse(state.checkAssignmentMatchedSubscription(Collections.singletonList(t1p0)));
     }
 
-    @Test(expected = IllegalArgumentException.class)
+    @Test
     public void cantAssignPartitionForUnmatchedPattern() {
         state.subscribe(Pattern.compile(".*t"), rebalanceListener);
-        state.subscribeFromPattern(new HashSet<>(Collections.singletonList(topic)));
-        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)
@@ -274,38 +326,407 @@ public void unsubscribeUserSubscribe() {
         state.unsubscribe();
         state.assignFromUser(singleton(tp0));
         assertEquals(singleton(tp0), state.assignedPartitions());
+        assertEquals(1, state.numAssignedPartitions());
     }
 
     @Test
     public void unsubscription() {
         state.subscribe(Pattern.compile(".*"), rebalanceListener);
         state.subscribeFromPattern(new HashSet<>(Arrays.asList(topic, topic1)));
+        assertTrue(state.checkAssignmentMatchedSubscription(singleton(tp1)));
         state.assignFromSubscribed(singleton(tp1));
+
         assertEquals(singleton(tp1), state.assignedPartitions());
+        assertEquals(1, state.numAssignedPartitions());
 
         state.unsubscribe();
         assertEquals(0, state.subscription().size());
         assertTrue(state.assignedPartitions().isEmpty());
+        assertEquals(0, state.numAssignedPartitions());
 
         state.assignFromUser(singleton(tp0));
         assertEquals(singleton(tp0), state.assignedPartitions());
+        assertEquals(1, state.numAssignedPartitions());
 
         state.unsubscribe();
         assertEquals(0, state.subscription().size());
         assertTrue(state.assignedPartitions().isEmpty());
+        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(Optional.of(broker1), Optional.of(5))));
+        assertTrue(state.hasValidPosition(tp0));
+        assertFalse(state.awaitingValidation(tp0));
+        ApiVersions apiVersions = new ApiVersions();
+        apiVersions.update(broker1.idString(), NodeApiVersions.create());
+
+        assertFalse(state.maybeValidatePositionForCurrentLeader(apiVersions, tp0, new Metadata.LeaderAndEpoch(
+                Optional.of(broker1), Optional.empty())));
+        assertTrue(state.hasValidPosition(tp0));
+        assertFalse(state.awaitingValidation(tp0));
+
+        assertFalse(state.maybeValidatePositionForCurrentLeader(apiVersions, tp0, new Metadata.LeaderAndEpoch(
+                Optional.of(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(Optional.of(broker1), Optional.of(5))));
+        assertFalse(state.hasValidPosition(tp0));
+        assertTrue(state.awaitingValidation(tp0));
+
+        state.seekUnvalidated(tp0, new SubscriptionState.FetchPosition(0L, Optional.empty(),
+                new Metadata.LeaderAndEpoch(Optional.of(broker1), Optional.of(5))));
+        assertTrue(state.hasValidPosition(tp0));
+        assertFalse(state.awaitingValidation(tp0));
+    }
+
+    @Test
+    public void testSeekUnvalidatedWithOffsetEpoch() {
+        Node broker1 = new Node(1, "localhost", 9092);
+        ApiVersions apiVersions = new ApiVersions();
+        apiVersions.update(broker1.idString(), NodeApiVersions.create());
+
+        state.assignFromUser(Collections.singleton(tp0));
+
+        state.seekUnvalidated(tp0, new SubscriptionState.FetchPosition(0L, Optional.of(2),
+                new Metadata.LeaderAndEpoch(Optional.of(broker1), Optional.of(5))));
+        assertFalse(state.hasValidPosition(tp0));
+        assertTrue(state.awaitingValidation(tp0));
+
+        // Update using the current leader and epoch
+        assertTrue(state.maybeValidatePositionForCurrentLeader(apiVersions, tp0, new Metadata.LeaderAndEpoch(
+                Optional.of(broker1), Optional.of(5))));
+        assertFalse(state.hasValidPosition(tp0));
+        assertTrue(state.awaitingValidation(tp0));
+
+        // Update with a newer leader and epoch
+        assertTrue(state.maybeValidatePositionForCurrentLeader(apiVersions, tp0, new Metadata.LeaderAndEpoch(
+                Optional.of(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(apiVersions, tp0, new Metadata.LeaderAndEpoch(
+                Optional.of(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(Optional.of(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(Optional.of(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(Optional.of(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);
     }
 
-    private void assertAllPositions(TopicPartition tp, Long offset) {
-        assertEquals(offset.longValue(), state.committed(tp).offset());
-        assertEquals(offset, state.position(tp));
+    @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(Optional.of(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(Optional.of(broker1), Optional.of(currentEpoch)));
+        state.seekUnvalidated(tp0, initialPosition);
+        assertTrue(state.awaitingValidation(tp0));
+
+        Optional truncationOpt = state.maybeCompleteValidation(tp0, initialPosition,
+                new EpochEndOffset()
+                    .setLeaderEpoch(initialOffsetEpoch)
+                    .setEndOffset(initialOffset + 5));
+        assertEquals(Optional.empty(), truncationOpt);
+        assertFalse(state.awaitingValidation(tp0));
+        assertEquals(initialPosition, state.position(tp0));
+    }
+
+    @Test
+    public void testMaybeValidatePositionForCurrentLeader() {
+        NodeApiVersions oldApis = NodeApiVersions.create(ApiKeys.OFFSET_FOR_LEADER_EPOCH.id, (short) 0, (short) 2);
+        ApiVersions apiVersions = new ApiVersions();
+        apiVersions.update("1", oldApis);
+
+        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(Optional.of(broker1), Optional.of(10))));
+
+        // if API is too old to be usable, we just skip validation
+        assertFalse(state.maybeValidatePositionForCurrentLeader(apiVersions, tp0, new Metadata.LeaderAndEpoch(
+                Optional.of(broker1), Optional.of(10))));
+        assertTrue(state.hasValidPosition(tp0));
+
+        // New API
+        apiVersions.update("1", NodeApiVersions.create());
+        state.seekUnvalidated(tp0, new SubscriptionState.FetchPosition(10L, Optional.of(5),
+                new Metadata.LeaderAndEpoch(Optional.of(broker1), Optional.of(10))));
+
+        // API is too old to be usable, we just skip validation
+        assertTrue(state.maybeValidatePositionForCurrentLeader(apiVersions, tp0, new Metadata.LeaderAndEpoch(
+                Optional.of(broker1), Optional.of(10))));
+        assertFalse(state.hasValidPosition(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(Optional.of(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(Optional.of(broker1), Optional.of(currentEpoch)));
+        state.seekUnvalidated(tp0, updatePosition);
+
+        Optional truncationOpt = state.maybeCompleteValidation(tp0, initialPosition,
+                new EpochEndOffset()
+                    .setLeaderEpoch(initialOffsetEpoch)
+                    .setEndOffset(initialOffset + 5));
+        assertEquals(Optional.empty(), truncationOpt);
+        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(Optional.of(broker1), Optional.of(currentEpoch)));
+        state.seekUnvalidated(tp0, initialPosition);
+        assertTrue(state.awaitingValidation(tp0));
+
+        state.requestOffsetReset(tp0);
+
+        Optional truncationOpt = state.maybeCompleteValidation(tp0, initialPosition,
+                new EpochEndOffset()
+                    .setLeaderEpoch(initialOffsetEpoch)
+                    .setEndOffset(initialOffset + 5));
+        assertEquals(Optional.empty(), truncationOpt);
+        assertFalse(state.awaitingValidation(tp0));
+        assertTrue(state.isOffsetResetNeeded(tp0));
+        assertNull(state.position(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(Optional.of(broker1), Optional.of(currentEpoch)));
+        state.seekUnvalidated(tp0, initialPosition);
+        assertTrue(state.awaitingValidation(tp0));
+
+        Optional truncationOpt = state.maybeCompleteValidation(tp0, initialPosition,
+                new EpochEndOffset()
+                    .setLeaderEpoch(divergentOffsetEpoch)
+                    .setEndOffset(divergentOffset));
+        assertEquals(Optional.empty(), truncationOpt);
+        assertFalse(state.awaitingValidation(tp0));
+
+        SubscriptionState.FetchPosition updatedPosition = new SubscriptionState.FetchPosition(divergentOffset,
+                Optional.of(divergentOffsetEpoch), new Metadata.LeaderAndEpoch(Optional.of(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(Optional.of(broker1), Optional.of(currentEpoch)));
+        state.seekUnvalidated(tp0, initialPosition);
+        assertTrue(state.awaitingValidation(tp0));
+
+        Optional truncationOpt = state.maybeCompleteValidation(tp0, initialPosition,
+                new EpochEndOffset()
+                    .setLeaderEpoch(divergentOffsetEpoch)
+                    .setEndOffset(divergentOffset));
+        assertTrue(truncationOpt.isPresent());
+        LogTruncation truncation = truncationOpt.get();
+
+        assertEquals(Optional.of(new OffsetAndMetadata(divergentOffset, Optional.of(divergentOffsetEpoch), "")),
+                truncation.divergentOffsetOpt);
+        assertEquals(initialPosition, truncation.fetchPosition);
+        assertTrue(state.awaitingValidation(tp0));
+    }
+
+    @Test
+    public void testTruncationDetectionUnknownDivergentOffsetWithResetPolicy() {
+        Node broker1 = new Node(1, "localhost", 9092);
+        state = new SubscriptionState(new LogContext(), OffsetResetStrategy.EARLIEST);
+        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(Optional.of(broker1), Optional.of(currentEpoch)));
+        state.seekUnvalidated(tp0, initialPosition);
+        assertTrue(state.awaitingValidation(tp0));
+
+        Optional truncationOpt = state.maybeCompleteValidation(tp0, initialPosition,
+                new EpochEndOffset()
+                    .setLeaderEpoch(UNDEFINED_EPOCH)
+                    .setEndOffset(UNDEFINED_EPOCH_OFFSET));
+        assertEquals(Optional.empty(), truncationOpt);
+        assertFalse(state.awaitingValidation(tp0));
+        assertTrue(state.isOffsetResetNeeded(tp0));
+        assertEquals(OffsetResetStrategy.EARLIEST, state.resetStrategy(tp0));
+    }
+
+    @Test
+    public void testTruncationDetectionUnknownDivergentOffsetWithoutResetPolicy() {
+        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;
+
+        SubscriptionState.FetchPosition initialPosition = new SubscriptionState.FetchPosition(initialOffset,
+            Optional.of(initialOffsetEpoch), new Metadata.LeaderAndEpoch(Optional.of(broker1), Optional.of(currentEpoch)));
+        state.seekUnvalidated(tp0, initialPosition);
+        assertTrue(state.awaitingValidation(tp0));
+
+        Optional truncationOpt = state.maybeCompleteValidation(tp0, initialPosition,
+                new EpochEndOffset()
+                    .setLeaderEpoch(UNDEFINED_EPOCH)
+                    .setEndOffset(UNDEFINED_EPOCH_OFFSET));
+        assertTrue(truncationOpt.isPresent());
+        LogTruncation truncation = truncationOpt.get();
+
+        assertEquals(Optional.empty(), truncation.divergentOffsetOpt);
+        assertEquals(initialPosition, truncation.fetchPosition);
+        assertTrue(state.awaitingValidation(tp0));
     }
 
     private static class MockRebalanceListener implements ConsumerRebalanceListener {
-        public Collection revoked;
+        Collection revoked;
         public Collection assigned;
-        public int revokedCount = 0;
-        public int assignedCount = 0;
-
+        int revokedCount = 0;
+        int assignedCount = 0;
 
         @Override
         public void onPartitionsAssigned(Collection partitions) {
@@ -321,4 +742,53 @@ public void onPartitionsRevoked(Collection partitions) {
 
     }
 
+    @Test
+    public void resetOffsetNoValidation() {
+        // Check that offset reset works when we can't validate offsets (older brokers)
+
+        Node broker1 = new Node(1, "localhost", 9092);
+        state.assignFromUser(Collections.singleton(tp0));
+
+        // Reset offsets
+        state.requestOffsetReset(tp0, OffsetResetStrategy.EARLIEST);
+
+        // Attempt to validate with older API version, should do nothing
+        ApiVersions oldApis = new ApiVersions();
+        oldApis.update("1", NodeApiVersions.create(ApiKeys.OFFSET_FOR_LEADER_EPOCH.id, (short) 0, (short) 2));
+        assertFalse(state.maybeValidatePositionForCurrentLeader(oldApis, tp0, new Metadata.LeaderAndEpoch(
+                Optional.of(broker1), Optional.empty())));
+        assertFalse(state.hasValidPosition(tp0));
+        assertFalse(state.awaitingValidation(tp0));
+        assertTrue(state.isOffsetResetNeeded(tp0));
+
+        // Complete the reset via unvalidated seek
+        state.seekUnvalidated(tp0, new SubscriptionState.FetchPosition(10L));
+        assertTrue(state.hasValidPosition(tp0));
+        assertFalse(state.awaitingValidation(tp0));
+        assertFalse(state.isOffsetResetNeeded(tp0));
+
+        // Next call to validate offsets does nothing
+        assertFalse(state.maybeValidatePositionForCurrentLeader(oldApis, tp0, new Metadata.LeaderAndEpoch(
+                Optional.of(broker1), Optional.empty())));
+        assertTrue(state.hasValidPosition(tp0));
+        assertFalse(state.awaitingValidation(tp0));
+        assertFalse(state.isOffsetResetNeeded(tp0));
+
+        // Reset again, and complete it with a seek that would normally require validation
+        state.requestOffsetReset(tp0, OffsetResetStrategy.EARLIEST);
+        state.seekUnvalidated(tp0, new SubscriptionState.FetchPosition(10L, Optional.of(10), new Metadata.LeaderAndEpoch(
+                Optional.of(broker1), Optional.of(2))));
+        // We are now in AWAIT_VALIDATION
+        assertFalse(state.hasValidPosition(tp0));
+        assertTrue(state.awaitingValidation(tp0));
+        assertFalse(state.isOffsetResetNeeded(tp0));
+
+        // Now ensure next call to validate clears the validation state
+        assertFalse(state.maybeValidatePositionForCurrentLeader(oldApis, tp0, new Metadata.LeaderAndEpoch(
+                Optional.of(broker1), Optional.of(2))));
+        assertTrue(state.hasValidPosition(tp0));
+        assertFalse(state.awaitingValidation(tp0));
+        assertFalse(state.isOffsetResetNeeded(tp0));
+    }
+
 }
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ThrowOnAssignmentAssignor.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ThrowOnAssignmentAssignor.java
new file mode 100644
index 0000000000000..782ca26006985
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ThrowOnAssignmentAssignor.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.clients.consumer.internals;
+
+import org.apache.kafka.clients.consumer.ConsumerGroupMetadata;
+
+import java.util.List;
+
+/**
+ * A mock assignor which throws for {@link org.apache.kafka.clients.consumer.ConsumerPartitionAssignor#onAssignment}.
+ */
+public class ThrowOnAssignmentAssignor extends MockPartitionAssignor {
+
+    private final RuntimeException bookeepedException;
+    private final String name;
+
+    ThrowOnAssignmentAssignor(final List supportedProtocols,
+                              final RuntimeException bookeepedException,
+                              final String name) {
+        super(supportedProtocols);
+        this.bookeepedException = bookeepedException;
+        this.name = name;
+    }
+
+    @Override
+    public void onAssignment(Assignment assignment, ConsumerGroupMetadata metadata) {
+        super.onAssignment(assignment, metadata);
+        throw bookeepedException;
+    }
+
+    @Override
+    public String name() {
+        return name;
+    }
+}
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 26f758855e792..57c8474996dce 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
@@ -17,52 +17,173 @@
 package org.apache.kafka.clients.producer;
 
 import org.apache.kafka.clients.CommonClientConfigs;
-import org.apache.kafka.clients.Metadata;
+import org.apache.kafka.clients.KafkaClient;
+import org.apache.kafka.clients.MockClient;
+import org.apache.kafka.clients.NodeApiVersions;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerGroupMetadata;
 import org.apache.kafka.clients.producer.internals.ProducerInterceptors;
+import org.apache.kafka.clients.producer.internals.ProducerMetadata;
+import org.apache.kafka.clients.producer.internals.Sender;
 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.config.ConfigException;
+import org.apache.kafka.common.config.SslConfigs;
+import org.apache.kafka.common.errors.InterruptException;
+import org.apache.kafka.common.errors.InvalidTopicException;
 import org.apache.kafka.common.errors.TimeoutException;
-import org.apache.kafka.common.internals.ClusterResourceListeners;
 import org.apache.kafka.common.header.internals.RecordHeader;
+import org.apache.kafka.common.internals.ClusterResourceListeners;
+import org.apache.kafka.common.message.AddOffsetsToTxnResponseData;
+import org.apache.kafka.common.message.EndTxnResponseData;
+import org.apache.kafka.common.message.InitProducerIdResponseData;
+import org.apache.kafka.common.message.TxnOffsetCommitRequestData;
 import org.apache.kafka.common.metrics.Sensor;
+import org.apache.kafka.common.metrics.stats.Avg;
 import org.apache.kafka.common.network.Selectable;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.requests.AddOffsetsToTxnResponse;
+import org.apache.kafka.common.requests.EndTxnResponse;
+import org.apache.kafka.common.requests.FindCoordinatorRequest;
+import org.apache.kafka.common.requests.FindCoordinatorResponse;
+import org.apache.kafka.common.requests.InitProducerIdResponse;
+import org.apache.kafka.common.requests.JoinGroupRequest;
+import org.apache.kafka.common.requests.MetadataResponse;
+import org.apache.kafka.common.requests.RequestTestUtils;
+import org.apache.kafka.common.requests.TxnOffsetCommitRequest;
+import org.apache.kafka.common.requests.TxnOffsetCommitResponse;
 import org.apache.kafka.common.serialization.ByteArraySerializer;
-import org.apache.kafka.common.serialization.ExtendedSerializer;
+import org.apache.kafka.common.serialization.Serializer;
 import org.apache.kafka.common.serialization.StringSerializer;
+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.MockMetricsReporter;
+import org.apache.kafka.test.MockPartitioner;
 import org.apache.kafka.test.MockProducerInterceptor;
 import org.apache.kafka.test.MockSerializer;
-import org.apache.kafka.test.MockPartitioner;
-import org.easymock.EasyMock;
+import org.apache.kafka.test.TestUtils;
 import org.junit.Assert;
 import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.powermock.api.easymock.PowerMock;
-import org.powermock.api.support.membermodification.MemberModifier;
-import org.powermock.core.classloader.annotations.PowerMockIgnore;
-import org.powermock.core.classloader.annotations.PrepareOnlyThisForTest;
-import org.powermock.modules.junit4.PowerMockRunner;
 
+import javax.management.MBeanServer;
+import javax.management.ObjectName;
+import java.lang.management.ManagementFactory;
+import java.time.Duration;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Properties;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.Exchanger;
+import java.util.concurrent.ExecutionException;
+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.concurrent.atomic.AtomicReference;
+import java.util.stream.Stream;
 
+import static java.util.Collections.emptyMap;
+import static java.util.Collections.singletonMap;
+import static org.junit.Assert.assertArrayEquals;
 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.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.ArgumentMatchers.notNull;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
-@RunWith(PowerMockRunner.class)
-@PowerMockIgnore("javax.management.*")
 public class KafkaProducerTest {
+    private final String topic = "topic";
+    private final Node host1 = new Node(0, "host1", 1000);
+    private final 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",
+            nodes,
+            Collections.singletonList(new PartitionInfo(topic, 0, null, null, null)),
+            Collections.emptySet(),
+            Collections.emptySet());
+    private final Cluster threePartitionCluster = new Cluster(
+            "dummy",
+            nodes,
+            Arrays.asList(
+                    new PartitionInfo(topic, 0, null, null, null),
+                    new PartitionInfo(topic, 1, null, null, null),
+                    new PartitionInfo(topic, 2, null, null, null)),
+            Collections.emptySet(),
+            Collections.emptySet());
+    private static final int DEFAULT_METADATA_IDLE_MS = 5 * 60 * 1000;
+
+
+    private static  KafkaProducer kafkaProducer(Map configs,
+                  Serializer keySerializer,
+                  Serializer valueSerializer,
+                  ProducerMetadata metadata,
+                  KafkaClient kafkaClient,
+                  ProducerInterceptors interceptors,
+                  Time time) {
+        return new KafkaProducer<>(new ProducerConfig(ProducerConfig.appendSerializerToConfig(configs, keySerializer, valueSerializer)),
+                keySerializer, valueSerializer, metadata, kafkaClient, interceptors, time);
+    }
+
+    @Test
+    public void testOverwriteAcksAndRetriesForIdempotentProducers() {
+        Properties props = new Properties();
+        props.setProperty(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
+        props.setProperty(ProducerConfig.TRANSACTIONAL_ID_CONFIG, "transactionalId");
+        props.setProperty(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, StringSerializer.class.getName());
+        props.setProperty(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, StringSerializer.class.getName());
+
+        ProducerConfig config = new ProducerConfig(props);
+        assertTrue(config.getBoolean(ProducerConfig.ENABLE_IDEMPOTENCE_CONFIG));
+        assertTrue(Stream.of("-1", "all").anyMatch(each -> each.equalsIgnoreCase(config.getString(ProducerConfig.ACKS_CONFIG))));
+        assertEquals((int) config.getInt(ProducerConfig.RETRIES_CONFIG), Integer.MAX_VALUE);
+        assertTrue(config.getString(ProducerConfig.CLIENT_ID_CONFIG).equalsIgnoreCase("producer-" +
+                config.getString(ProducerConfig.TRANSACTIONAL_ID_CONFIG)));
+    }
+
+    @Test
+    public void testMetricsReporterAutoGeneratedClientId() {
+        Properties props = new Properties();
+        props.setProperty(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
+        props.setProperty(ProducerConfig.METRIC_REPORTER_CLASSES_CONFIG, MockMetricsReporter.class.getName());
+        KafkaProducer producer = new KafkaProducer<>(
+                props, new StringSerializer(), new StringSerializer());
+
+        MockMetricsReporter mockMetricsReporter = (MockMetricsReporter) producer.metrics.reporters().get(0);
+
+        Assert.assertEquals(producer.getClientId(), mockMetricsReporter.clientId);
+        producer.close();
+    }
 
     @Test
     public void testConstructorWithSerializers() {
@@ -87,7 +208,7 @@ public void testConstructorFailureCloseResource() {
 
         final int oldInitCount = MockMetricsReporter.INIT_COUNT.get();
         final int oldCloseCount = MockMetricsReporter.CLOSE_COUNT.get();
-        try (KafkaProducer producer = new KafkaProducer<>(props, new ByteArraySerializer(), new ByteArraySerializer())) {
+        try (KafkaProducer ignored = new KafkaProducer<>(props, new ByteArraySerializer(), new ByteArraySerializer())) {
             fail("should have caught an exception and returned");
         } catch (KafkaException e) {
             assertEquals(oldInitCount + 1, MockMetricsReporter.INIT_COUNT.get());
@@ -97,7 +218,19 @@ public void testConstructorFailureCloseResource() {
     }
 
     @Test
-    public void testSerializerClose() throws Exception {
+    public void testConstructorWithNotStringKey() {
+        Properties props = new Properties();
+        props.setProperty(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
+        props.put(1, "not string key");
+        try (KafkaProducer ff = new KafkaProducer<>(props, new StringSerializer(), new StringSerializer())) {
+            fail("Constructor should throw exception");
+        } catch (ConfigException e) {
+            assertTrue("Unexpected exception message: " + e.getMessage(), e.getMessage().contains("not string key"));
+        }
+    }
+
+    @Test
+    public void testSerializerClose() {
         Map configs = new HashMap<>();
         configs.put(ProducerConfig.CLIENT_ID_CONFIG, "testConstructorClose");
         configs.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
@@ -106,7 +239,7 @@ public void testSerializerClose() throws Exception {
         final int oldInitCount = MockSerializer.INIT_COUNT.get();
         final int oldCloseCount = MockSerializer.CLOSE_COUNT.get();
 
-        KafkaProducer producer = new KafkaProducer(
+        KafkaProducer producer = new KafkaProducer<>(
                 configs, new MockSerializer(), new MockSerializer());
         assertEquals(oldInitCount + 2, MockSerializer.INIT_COUNT.get());
         assertEquals(oldCloseCount, MockSerializer.CLOSE_COUNT.get());
@@ -117,7 +250,7 @@ public void testSerializerClose() throws Exception {
     }
 
     @Test
-    public void testInterceptorConstructClose() throws Exception {
+    public void testInterceptorConstructClose() {
         try {
             Properties props = new Properties();
             // test with client ID assigned by KafkaProducer
@@ -125,7 +258,7 @@ public void testInterceptorConstructClose() throws Exception {
             props.setProperty(ProducerConfig.INTERCEPTOR_CLASSES_CONFIG, MockProducerInterceptor.class.getName());
             props.setProperty(MockProducerInterceptor.APPEND_STRING_PROP, "something");
 
-            KafkaProducer producer = new KafkaProducer(
+            KafkaProducer producer = new KafkaProducer<>(
                     props, new StringSerializer(), new StringSerializer());
             assertEquals(1, MockProducerInterceptor.INIT_COUNT.get());
             assertEquals(0, MockProducerInterceptor.CLOSE_COUNT.get());
@@ -143,13 +276,14 @@ public void testInterceptorConstructClose() throws Exception {
     }
 
     @Test
-    public void testPartitionerClose() throws Exception {
+    public void testPartitionerClose() {
         try {
             Properties props = new Properties();
             props.setProperty(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
+            MockPartitioner.resetCounters();
             props.setProperty(ProducerConfig.PARTITIONER_CLASS_CONFIG, MockPartitioner.class.getName());
 
-            KafkaProducer producer = new KafkaProducer(
+            KafkaProducer producer = new KafkaProducer<>(
                     props, new StringSerializer(), new StringSerializer());
             assertEquals(1, MockPartitioner.INIT_COUNT.get());
             assertEquals(0, MockPartitioner.CLOSE_COUNT.get());
@@ -164,7 +298,58 @@ public void testPartitionerClose() throws Exception {
     }
 
     @Test
-    public void testOsDefaultSocketBufferSizes() throws Exception {
+    public void shouldCloseProperlyAndThrowIfInterrupted() throws Exception {
+        Map configs = new HashMap<>();
+        configs.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
+        configs.put(ProducerConfig.PARTITIONER_CLASS_CONFIG, MockPartitioner.class.getName());
+        configs.put(ProducerConfig.BATCH_SIZE_CONFIG, "1");
+
+        Time time = new MockTime();
+        MetadataResponse initialUpdateResponse = RequestTestUtils.metadataUpdateWith(1, singletonMap("topic", 1));
+        ProducerMetadata metadata = newMetadata(0, Long.MAX_VALUE);
+        MockClient client = new MockClient(time, metadata);
+        client.updateMetadata(initialUpdateResponse);
+
+        final Producer producer = kafkaProducer(configs, new StringSerializer(),
+                new StringSerializer(), metadata, client, null, time);
+
+        ExecutorService executor = Executors.newSingleThreadExecutor();
+        final AtomicReference closeException = new AtomicReference<>();
+        try {
+            Future future = executor.submit(() -> {
+                producer.send(new ProducerRecord<>("topic", "key", "value"));
+                try {
+                    producer.close();
+                    fail("Close should block and throw.");
+                } catch (Exception e) {
+                    closeException.set(e);
+                }
+            });
+
+            // Close producer should not complete until send succeeds
+            try {
+                future.get(100, TimeUnit.MILLISECONDS);
+                fail("Close completed without waiting for send");
+            } catch (java.util.concurrent.TimeoutException expected) { /* ignore */ }
+
+            // Ensure send has started
+            client.waitForRequests(1, 1000);
+
+            assertTrue("Close terminated prematurely", future.cancel(true));
+
+            TestUtils.waitForCondition(() -> closeException.get() != null,
+                    "InterruptException did not occur within timeout.");
+
+            assertTrue("Expected exception not thrown " + closeException,
+                    closeException.get() instanceof InterruptException);
+        } finally {
+            executor.shutdownNow();
+        }
+
+    }
+
+    @Test
+    public void testOsDefaultSocketBufferSizes() {
         Map config = new HashMap<>();
         config.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
         config.put(ProducerConfig.SEND_BUFFER_CONFIG, Selectable.USE_DEFAULT_BUFFER_SIZE);
@@ -173,7 +358,7 @@ public void testOsDefaultSocketBufferSizes() throws Exception {
     }
 
     @Test(expected = KafkaException.class)
-    public void testInvalidSocketSendBufferSize() throws Exception {
+    public void testInvalidSocketSendBufferSize() {
         Map config = new HashMap<>();
         config.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
         config.put(ProducerConfig.SEND_BUFFER_CONFIG, -2);
@@ -181,288 +366,418 @@ public void testInvalidSocketSendBufferSize() throws Exception {
     }
 
     @Test(expected = KafkaException.class)
-    public void testInvalidSocketReceiveBufferSize() throws Exception {
+    public void testInvalidSocketReceiveBufferSize() {
         Map config = new HashMap<>();
         config.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
         config.put(ProducerConfig.RECEIVE_BUFFER_CONFIG, -2);
         new KafkaProducer<>(config, new ByteArraySerializer(), new ByteArraySerializer());
     }
 
-    @PrepareOnlyThisForTest(Metadata.class)
+    private static KafkaProducer producerWithOverrideNewSender(Map configs,
+                                                                               ProducerMetadata metadata) {
+        return producerWithOverrideNewSender(configs, metadata, Time.SYSTEM);
+    }
+
+    private static KafkaProducer producerWithOverrideNewSender(Map configs,
+                                                                               ProducerMetadata metadata,
+                                                                               Time timer) {
+        return new KafkaProducer(
+                new ProducerConfig(ProducerConfig.appendSerializerToConfig(configs, new StringSerializer(), new StringSerializer())),
+                new StringSerializer(), new StringSerializer(), metadata, new MockClient(Time.SYSTEM, metadata), null, timer) {
+            @Override
+            Sender newSender(LogContext logContext, KafkaClient kafkaClient, ProducerMetadata metadata) {
+                // give Sender its own Metadata instance so that we can isolate Metadata calls from KafkaProducer
+                return super.newSender(logContext, kafkaClient, newMetadata(0, 100_000));
+            }
+        };
+    }
+
     @Test
-    public void testMetadataFetch() throws Exception {
-        Properties props = new Properties();
-        props.setProperty(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
-        KafkaProducer producer = new KafkaProducer<>(props, new StringSerializer(), new StringSerializer());
-        Metadata metadata = PowerMock.createNiceMock(Metadata.class);
-        MemberModifier.field(KafkaProducer.class, "metadata").set(producer, metadata);
+    public void testMetadataFetch() throws InterruptedException {
+        Map configs = new HashMap<>();
+        configs.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
+        ProducerMetadata metadata = mock(ProducerMetadata.class);
 
-        String topic = "topic";
+        // Return empty cluster 4 times and cluster from then on
+        when(metadata.fetch()).thenReturn(emptyCluster, emptyCluster, emptyCluster, emptyCluster, onePartitionCluster);
+
+        KafkaProducer producer = producerWithOverrideNewSender(configs, metadata);
         ProducerRecord record = new ProducerRecord<>(topic, "value");
-        Collection nodes = Collections.singletonList(new Node(0, "host1", 1000));
-        final Cluster emptyCluster = new Cluster(null, nodes,
-                Collections.emptySet(),
-                Collections.emptySet(),
-                Collections.emptySet());
-        final Cluster cluster = new Cluster(
-                "dummy",
-                Collections.singletonList(new Node(0, "host1", 1000)),
-                Arrays.asList(new PartitionInfo(topic, 0, null, null, null)),
-                Collections.emptySet(),
-                Collections.emptySet());
-
-        // Expect exactly one fetch for each attempt to refresh while topic metadata is not available
-        final int refreshAttempts = 5;
-        EasyMock.expect(metadata.fetch()).andReturn(emptyCluster).times(refreshAttempts - 1);
-        EasyMock.expect(metadata.fetch()).andReturn(cluster).once();
-        EasyMock.expect(metadata.fetch()).andThrow(new IllegalStateException("Unexpected call to metadata.fetch()")).anyTimes();
-        PowerMock.replay(metadata);
         producer.send(record);
-        PowerMock.verify(metadata);
 
-        // Expect exactly one fetch if topic metadata is available
-        PowerMock.reset(metadata);
-        EasyMock.expect(metadata.fetch()).andReturn(cluster).once();
-        EasyMock.expect(metadata.fetch()).andThrow(new IllegalStateException("Unexpected call to metadata.fetch()")).anyTimes();
-        PowerMock.replay(metadata);
+        // One request update for each empty cluster returned
+        verify(metadata, times(4)).requestUpdateForTopic(topic);
+        verify(metadata, times(4)).awaitUpdate(anyInt(), anyLong());
+        verify(metadata, times(5)).fetch();
+
+        // Should not request update for subsequent `send`
         producer.send(record, null);
-        PowerMock.verify(metadata);
+        verify(metadata, times(4)).requestUpdateForTopic(topic);
+        verify(metadata, times(4)).awaitUpdate(anyInt(), anyLong());
+        verify(metadata, times(6)).fetch();
 
-        // Expect exactly one fetch if topic metadata is available
-        PowerMock.reset(metadata);
-        EasyMock.expect(metadata.fetch()).andReturn(cluster).once();
-        EasyMock.expect(metadata.fetch()).andThrow(new IllegalStateException("Unexpected call to metadata.fetch()")).anyTimes();
-        PowerMock.replay(metadata);
+        // Should not request update for subsequent `partitionsFor`
         producer.partitionsFor(topic);
-        PowerMock.verify(metadata);
+        verify(metadata, times(4)).requestUpdateForTopic(topic);
+        verify(metadata, times(4)).awaitUpdate(anyInt(), anyLong());
+        verify(metadata, times(7)).fetch();
+
+        producer.close(Duration.ofMillis(0));
     }
 
-    @PrepareOnlyThisForTest(Metadata.class)
     @Test
-    public void testMetadataFetchOnStaleMetadata() throws Exception {
-        Properties props = new Properties();
-        props.setProperty(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
-        KafkaProducer producer = new KafkaProducer<>(props, new StringSerializer(), new StringSerializer());
-        Metadata metadata = PowerMock.createNiceMock(Metadata.class);
-        MemberModifier.field(KafkaProducer.class, "metadata").set(producer, metadata);
+    public void testMetadataExpiry() throws InterruptedException {
+        Map configs = new HashMap<>();
+        configs.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
+        ProducerMetadata metadata = mock(ProducerMetadata.class);
+
+        Cluster emptyCluster = new Cluster(
+            "dummy",
+            Collections.singletonList(host1),
+            Collections.emptySet(),
+            Collections.emptySet(),
+            Collections.emptySet());
+        when(metadata.fetch()).thenReturn(onePartitionCluster, emptyCluster, onePartitionCluster);
+
+        KafkaProducer producer = producerWithOverrideNewSender(configs, metadata);
+        ProducerRecord record = new ProducerRecord<>(topic, "value");
+        producer.send(record);
+
+        // Verify the topic's metadata isn't requested since it's already present.
+        verify(metadata, times(0)).requestUpdateForTopic(topic);
+        verify(metadata, times(0)).awaitUpdate(anyInt(), anyLong());
+        verify(metadata, times(1)).fetch();
+
+        // The metadata has been expired. Verify the producer requests the topic's metadata.
+        producer.send(record, null);
+        verify(metadata, times(1)).requestUpdateForTopic(topic);
+        verify(metadata, times(1)).awaitUpdate(anyInt(), anyLong());
+        verify(metadata, times(3)).fetch();
+
+        producer.close(Duration.ofMillis(0));
+    }
+
+    @Test
+    public void testMetadataTimeoutWithMissingTopic() throws Exception {
+        Map configs = new HashMap<>();
+        configs.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
+        configs.put(ProducerConfig.MAX_BLOCK_MS_CONFIG, 60000);
 
-        String topic = "topic";
-        ProducerRecord initialRecord = new ProducerRecord<>(topic, "value");
         // Create a record with a partition higher than the initial (outdated) partition range
-        ProducerRecord extendedRecord = new ProducerRecord<>(topic, 2, null, "value");
-        Collection nodes = Collections.singletonList(new Node(0, "host1", 1000));
-        final Cluster emptyCluster = new Cluster(null, nodes,
-                Collections.emptySet(),
-                Collections.emptySet(),
-                Collections.emptySet());
-        final Cluster initialCluster = new Cluster(
-                "dummy",
-                Collections.singletonList(new Node(0, "host1", 1000)),
-                Arrays.asList(new PartitionInfo(topic, 0, null, null, null)),
-                Collections.emptySet(),
-                Collections.emptySet());
-        final Cluster extendedCluster = new Cluster(
-                "dummy",
-                Collections.singletonList(new Node(0, "host1", 1000)),
-                Arrays.asList(
-                        new PartitionInfo(topic, 0, null, null, null),
-                        new PartitionInfo(topic, 1, null, null, null),
-                        new PartitionInfo(topic, 2, null, null, null)),
-                Collections.emptySet(),
-                Collections.emptySet());
-
-        // Expect exactly one fetch for each attempt to refresh while topic metadata is not available
-        final int refreshAttempts = 5;
-        EasyMock.expect(metadata.fetch()).andReturn(emptyCluster).times(refreshAttempts - 1);
-        EasyMock.expect(metadata.fetch()).andReturn(initialCluster).once();
-        EasyMock.expect(metadata.fetch()).andThrow(new IllegalStateException("Unexpected call to metadata.fetch()")).anyTimes();
-        PowerMock.replay(metadata);
-        producer.send(initialRecord);
-        PowerMock.verify(metadata);
-
-        // Expect exactly one fetch if topic metadata is available and records are still within range
-        PowerMock.reset(metadata);
-        EasyMock.expect(metadata.fetch()).andReturn(initialCluster).once();
-        EasyMock.expect(metadata.fetch()).andThrow(new IllegalStateException("Unexpected call to metadata.fetch()")).anyTimes();
-        PowerMock.replay(metadata);
-        producer.send(initialRecord, null);
-        PowerMock.verify(metadata);
-
-        // Expect exactly two fetches if topic metadata is available but metadata response still returns
-        // the same partition size (either because metadata are still stale at the broker too or because
-        // there weren't any partitions added in the first place).
-        PowerMock.reset(metadata);
-        EasyMock.expect(metadata.fetch()).andReturn(initialCluster).once();
-        EasyMock.expect(metadata.fetch()).andReturn(initialCluster).once();
-        EasyMock.expect(metadata.fetch()).andThrow(new IllegalStateException("Unexpected call to metadata.fetch()")).anyTimes();
-        PowerMock.replay(metadata);
+        ProducerRecord record = new ProducerRecord<>(topic, 2, null, "value");
+        ProducerMetadata metadata = mock(ProducerMetadata.class);
+
+        MockTime mockTime = new MockTime();
+        AtomicInteger invocationCount = new AtomicInteger(0);
+        when(metadata.fetch()).then(invocation -> {
+            invocationCount.incrementAndGet();
+            if (invocationCount.get() == 5) {
+                mockTime.setCurrentTimeMs(mockTime.milliseconds() + 70000);
+            }
+
+            return emptyCluster;
+        });
+
+        KafkaProducer producer = producerWithOverrideNewSender(configs, metadata, mockTime);
+
+        // Four request updates where the topic isn't present, at which point the timeout expires and a
+        // TimeoutException is thrown
+        Future future = producer.send(record);
+        verify(metadata, times(4)).requestUpdateForTopic(topic);
+        verify(metadata, times(4)).awaitUpdate(anyInt(), anyLong());
+        verify(metadata, times(5)).fetch();
         try {
-            producer.send(extendedRecord, null);
-            fail("Expected KafkaException to be raised");
-        } catch (KafkaException e) {
-            // expected
+            future.get();
+        } catch (ExecutionException e) {
+            assertTrue(e.getCause() instanceof TimeoutException);
+        } finally {
+            producer.close(Duration.ofMillis(0));
         }
-        PowerMock.verify(metadata);
+    }
 
-        // Expect exactly two fetches if topic metadata is available but outdated for the given record
-        PowerMock.reset(metadata);
-        EasyMock.expect(metadata.fetch()).andReturn(initialCluster).once();
-        EasyMock.expect(metadata.fetch()).andReturn(extendedCluster).once();
-        EasyMock.expect(metadata.fetch()).andThrow(new IllegalStateException("Unexpected call to metadata.fetch()")).anyTimes();
-        PowerMock.replay(metadata);
-        producer.send(extendedRecord, null);
-        PowerMock.verify(metadata);
+    @Test
+    public void testMetadataWithPartitionOutOfRange() throws Exception {
+        Map configs = new HashMap<>();
+        configs.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
+        configs.put(ProducerConfig.MAX_BLOCK_MS_CONFIG, 60000);
+
+        // Create a record with a partition higher than the initial (outdated) partition range
+        ProducerRecord record = new ProducerRecord<>(topic, 2, null, "value");
+        ProducerMetadata metadata = mock(ProducerMetadata.class);
+
+        MockTime mockTime = new MockTime();
+
+        when(metadata.fetch()).thenReturn(onePartitionCluster, onePartitionCluster, threePartitionCluster);
+
+        KafkaProducer producer = producerWithOverrideNewSender(configs, metadata, mockTime);
+        // One request update if metadata is available but outdated for the given record
+        producer.send(record);
+        verify(metadata, times(2)).requestUpdateForTopic(topic);
+        verify(metadata, times(2)).awaitUpdate(anyInt(), anyLong());
+        verify(metadata, times(3)).fetch();
+
+        producer.close(Duration.ofMillis(0));
     }
 
     @Test
-    public void testTopicRefreshInMetadata() throws Exception {
-        Properties props = new Properties();
-        props.setProperty(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
-        props.setProperty(ProducerConfig.MAX_BLOCK_MS_CONFIG, "600000");
-        KafkaProducer producer = new KafkaProducer<>(props, new StringSerializer(), new StringSerializer());
+    public void testMetadataTimeoutWithPartitionOutOfRange() throws Exception {
+        Map configs = new HashMap<>();
+        configs.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
+        configs.put(ProducerConfig.MAX_BLOCK_MS_CONFIG, 60000);
+
+        // Create a record with a partition higher than the initial (outdated) partition range
+        ProducerRecord record = new ProducerRecord<>(topic, 2, null, "value");
+        ProducerMetadata metadata = mock(ProducerMetadata.class);
+
+        MockTime mockTime = new MockTime();
+        AtomicInteger invocationCount = new AtomicInteger(0);
+        when(metadata.fetch()).then(invocation -> {
+            invocationCount.incrementAndGet();
+            if (invocationCount.get() == 5) {
+                mockTime.setCurrentTimeMs(mockTime.milliseconds() + 70000);
+            }
+
+            return onePartitionCluster;
+        });
+
+        KafkaProducer producer = producerWithOverrideNewSender(configs, metadata, mockTime);
+
+        // Four request updates where the requested partition is out of range, at which point the timeout expires
+        // and a TimeoutException is thrown
+        Future future = producer.send(record);
+        verify(metadata, times(4)).requestUpdateForTopic(topic);
+        verify(metadata, times(4)).awaitUpdate(anyInt(), anyLong());
+        verify(metadata, times(5)).fetch();
+        try {
+            future.get();
+        } catch (ExecutionException e) {
+            assertTrue(e.getCause() instanceof TimeoutException);
+        } finally {
+            producer.close(Duration.ofMillis(0));
+        }
+    }
+
+    @Test
+    public void testTopicRefreshInMetadata() throws InterruptedException {
+        Map configs = new HashMap<>();
+        configs.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
+        configs.put(ProducerConfig.MAX_BLOCK_MS_CONFIG, "600000");
         long refreshBackoffMs = 500L;
         long metadataExpireMs = 60000L;
-        final Metadata metadata = new Metadata(refreshBackoffMs, metadataExpireMs, true,
-                true, new ClusterResourceListeners());
+        long metadataIdleMs = 60000L;
         final Time time = new MockTime();
-        MemberModifier.field(KafkaProducer.class, "metadata").set(producer, metadata);
-        MemberModifier.field(KafkaProducer.class, "time").set(producer, time);
+        final ProducerMetadata metadata = new ProducerMetadata(refreshBackoffMs, metadataExpireMs, metadataIdleMs,
+                new LogContext(), new ClusterResourceListeners(), time);
         final String topic = "topic";
+        try (KafkaProducer producer = kafkaProducer(configs,
+                new StringSerializer(), new StringSerializer(), metadata, new MockClient(time, metadata), null, time)) {
 
-        Thread t = new Thread() {
-            @Override
-            public void run() {
+            AtomicBoolean running = new AtomicBoolean(true);
+            Thread t = new Thread(() -> {
                 long startTimeMs = System.currentTimeMillis();
-                for (int i = 0; i < 10; i++) {
-                    while (!metadata.updateRequested() && System.currentTimeMillis() - startTimeMs < 1000)
-                        yield();
-                    metadata.update(Cluster.empty(), Collections.singleton(topic), time.milliseconds());
+                while (running.get()) {
+                    while (!metadata.updateRequested() && System.currentTimeMillis() - startTimeMs < 100)
+                        Thread.yield();
+                    MetadataResponse updateResponse = RequestTestUtils.metadataUpdateWith("kafka-cluster", 1,
+                            singletonMap(topic, Errors.UNKNOWN_TOPIC_OR_PARTITION), emptyMap());
+                    metadata.updateWithCurrentRequestVersion(updateResponse, false, time.milliseconds());
                     time.sleep(60 * 1000L);
                 }
-            }
-        };
-        t.start();
-        try {
-            producer.partitionsFor(topic);
-            fail("Expect TimeoutException");
-        } catch (TimeoutException e) {
-            // skip
+            });
+            t.start();
+            assertThrows(TimeoutException.class, () -> producer.partitionsFor(topic));
+            running.set(false);
+            t.join();
         }
-        Assert.assertTrue("Topic should still exist in metadata", metadata.containsTopic(topic));
     }
 
-    @PrepareOnlyThisForTest(Metadata.class)
     @Test
-    public void testHeaders() throws Exception {
-        Properties props = new Properties();
-        props.setProperty(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
-        ExtendedSerializer keySerializer = PowerMock.createNiceMock(ExtendedSerializer.class);
-        ExtendedSerializer valueSerializer = PowerMock.createNiceMock(ExtendedSerializer.class);
+    public void testTopicExpiryInMetadata() throws InterruptedException {
+        Map configs = new HashMap<>();
+        configs.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
+        configs.put(ProducerConfig.MAX_BLOCK_MS_CONFIG, "30000");
+        long refreshBackoffMs = 500L;
+        long metadataExpireMs = 60000L;
+        long metadataIdleMs = 60000L;
+        final Time time = new MockTime();
+        final ProducerMetadata metadata = new ProducerMetadata(refreshBackoffMs, metadataExpireMs, metadataIdleMs,
+                new LogContext(), new ClusterResourceListeners(), time);
+        final String topic = "topic";
+        try (KafkaProducer producer = kafkaProducer(configs, new StringSerializer(),
+                new StringSerializer(), metadata, new MockClient(time, metadata), null, time)) {
 
-        KafkaProducer producer = new KafkaProducer<>(props, keySerializer, valueSerializer);
-        Metadata metadata = PowerMock.createNiceMock(Metadata.class);
-        MemberModifier.field(KafkaProducer.class, "metadata").set(producer, metadata);
+            Exchanger exchanger = new Exchanger<>();
 
-        String topic = "topic";
-        final Cluster cluster = new Cluster(
-                "dummy",
-                Collections.singletonList(new Node(0, "host1", 1000)),
-                Arrays.asList(new PartitionInfo(topic, 0, null, null, null)),
-                Collections.emptySet(),
-                Collections.emptySet());
+            Thread t = new Thread(() -> {
+                try {
+                    exchanger.exchange(null);  // 1
+                    while (!metadata.updateRequested())
+                        Thread.sleep(100);
+                    MetadataResponse updateResponse = RequestTestUtils.metadataUpdateWith(1, singletonMap(topic, 1));
+                    metadata.updateWithCurrentRequestVersion(updateResponse, false, time.milliseconds());
+                    exchanger.exchange(null);  // 2
+                    time.sleep(120 * 1000L);
 
+                    // Update the metadata again, but it should be expired at this point.
+                    updateResponse = RequestTestUtils.metadataUpdateWith(1, singletonMap(topic, 1));
+                    metadata.updateWithCurrentRequestVersion(updateResponse, false, time.milliseconds());
+                    exchanger.exchange(null);  // 3
+                    while (!metadata.updateRequested())
+                        Thread.sleep(100);
+                    time.sleep(30 * 1000L);
+                } catch (Exception e) {
+                    throw new RuntimeException(e);
+                }
+            });
+            t.start();
+            exchanger.exchange(null);  // 1
+            assertNotNull(producer.partitionsFor(topic));
+            exchanger.exchange(null);  // 2
+            exchanger.exchange(null);  // 3
+            assertThrows(TimeoutException.class, () -> producer.partitionsFor(topic));
+            t.join();
+        }
+    }
 
-        EasyMock.expect(metadata.fetch()).andReturn(cluster).anyTimes();
+    @SuppressWarnings("unchecked")
+    @Test
+    @Deprecated
+    public void testHeadersWithExtendedClasses() {
+        doTestHeaders(org.apache.kafka.common.serialization.ExtendedSerializer.class);
+    }
 
-        PowerMock.replay(metadata);
+    @SuppressWarnings("unchecked")
+    @Test
+    public void testHeaders() {
+        doTestHeaders(Serializer.class);
+    }
 
-        String value = "value";
+    private > void doTestHeaders(Class serializerClassToMock) {
+        Map configs = new HashMap<>();
+        configs.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
+        Serializer keySerializer = mock(serializerClassToMock);
+        Serializer valueSerializer = mock(serializerClassToMock);
+
+        long nowMs = Time.SYSTEM.milliseconds();
+        String topic = "topic";
+        ProducerMetadata metadata = newMetadata(0, 90000);
+        metadata.add(topic, nowMs);
+
+        MetadataResponse initialUpdateResponse = RequestTestUtils.metadataUpdateWith(1, singletonMap(topic, 1));
+        metadata.updateWithCurrentRequestVersion(initialUpdateResponse, false, nowMs);
 
-        ProducerRecord record = new ProducerRecord<>(topic, value);
-        EasyMock.expect(keySerializer.serialize(topic, record.headers(), null)).andReturn(null).once();
-        EasyMock.expect(valueSerializer.serialize(topic, record.headers(), value)).andReturn(value.getBytes()).once();
+        KafkaProducer producer = kafkaProducer(configs, keySerializer, valueSerializer, metadata,
+                null, null, Time.SYSTEM);
 
-        PowerMock.replay(keySerializer);
-        PowerMock.replay(valueSerializer);
+        when(keySerializer.serialize(any(), any(), any())).then(invocation ->
+                invocation.getArgument(2).getBytes());
+        when(valueSerializer.serialize(any(), any(), any())).then(invocation ->
+                invocation.getArgument(2).getBytes());
 
+        String value = "value";
+        String key = "key";
+        ProducerRecord record = new ProducerRecord<>(topic, key, value);
 
         //ensure headers can be mutated pre send.
         record.headers().add(new RecordHeader("test", "header2".getBytes()));
-
         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
-        assertTrue(Arrays.equals(record.headers().lastHeader("test").value(), "header2".getBytes()));
+        assertArrayEquals(record.headers().lastHeader("test").value(), "header2".getBytes());
 
-        PowerMock.verify(valueSerializer);
-        PowerMock.verify(keySerializer);
+        verify(valueSerializer).serialize(topic, record.headers(), value);
+        verify(keySerializer).serialize(topic, record.headers(), key);
 
+        producer.close(Duration.ofMillis(0));
     }
 
     @Test
     public void closeShouldBeIdempotent() {
         Properties producerProps = new Properties();
         producerProps.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9000");
-        Producer producer = new KafkaProducer<>(producerProps, new ByteArraySerializer(), new ByteArraySerializer());
+        Producer producer = new KafkaProducer<>(producerProps, new ByteArraySerializer(), new ByteArraySerializer());
         producer.close();
         producer.close();
     }
 
+    @Test
+    public void closeWithNegativeTimestampShouldThrow() {
+        Properties producerProps = new Properties();
+        producerProps.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9000");
+        try (Producer producer = new KafkaProducer<>(producerProps, new ByteArraySerializer(), new ByteArraySerializer())) {
+            assertThrows(IllegalArgumentException.class, () -> producer.close(Duration.ofMillis(-100)));
+        }
+    }
+
+    @Test
+    public void testFlushCompleteSendOfInflightBatches() {
+        Map configs = new HashMap<>();
+        configs.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9000");
+
+        Time time = new MockTime(1);
+        MetadataResponse initialUpdateResponse = RequestTestUtils.metadataUpdateWith(1, singletonMap("topic", 1));
+        ProducerMetadata metadata = newMetadata(0, Long.MAX_VALUE);
+
+        MockClient client = new MockClient(time, metadata);
+        client.updateMetadata(initialUpdateResponse);
+
+        try (Producer producer = kafkaProducer(configs, new StringSerializer(),
+                new StringSerializer(), metadata, client, null, time)) {
+            ArrayList> futureResponses = new ArrayList<>();
+            for (int i = 0; i < 50; i++) {
+                Future response = producer.send(new ProducerRecord<>("topic", "value" + i));
+                futureResponses.add(response);
+            }
+            futureResponses.forEach(res -> assertFalse(res.isDone()));
+            producer.flush();
+            futureResponses.forEach(res -> assertTrue(res.isDone()));
+        }
+    }
+
     @Test
     public void testMetricConfigRecordingLevel() {
         Properties props = new Properties();
         props.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9000");
-        try (KafkaProducer producer = new KafkaProducer<>(props, new ByteArraySerializer(), new ByteArraySerializer())) {
+        try (KafkaProducer producer = new KafkaProducer<>(props, new ByteArraySerializer(), new ByteArraySerializer())) {
             assertEquals(Sensor.RecordingLevel.INFO, producer.metrics.config().recordLevel());
         }
 
         props.put(ProducerConfig.METRICS_RECORDING_LEVEL_CONFIG, "DEBUG");
-        try (KafkaProducer producer = new KafkaProducer<>(props, new ByteArraySerializer(), new ByteArraySerializer())) {
+        try (KafkaProducer producer = new KafkaProducer<>(props, new ByteArraySerializer(), new ByteArraySerializer())) {
             assertEquals(Sensor.RecordingLevel.DEBUG, producer.metrics.config().recordLevel());
         }
     }
 
-    @PrepareOnlyThisForTest(Metadata.class)
     @Test
-    public void testInterceptorPartitionSetOnTooLargeRecord() throws Exception {
-        Properties props = new Properties();
-        props.setProperty(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
-        props.setProperty(ProducerConfig.MAX_REQUEST_SIZE_CONFIG, "1");
+    public void testInterceptorPartitionSetOnTooLargeRecord() {
+        Map configs = new HashMap<>();
+        configs.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
+        configs.put(ProducerConfig.MAX_REQUEST_SIZE_CONFIG, "1");
         String topic = "topic";
         ProducerRecord record = new ProducerRecord<>(topic, "value");
 
-        KafkaProducer producer = new KafkaProducer<>(props, new StringSerializer(),
-                new StringSerializer());
-        Metadata metadata = PowerMock.createNiceMock(Metadata.class);
-        MemberModifier.field(KafkaProducer.class, "metadata").set(producer, metadata);
-        final Cluster cluster = new Cluster(
-            "dummy",
-            Collections.singletonList(new Node(0, "host1", 1000)),
-            Arrays.asList(new PartitionInfo(topic, 0, null, null, null)),
-            Collections.emptySet(),
-            Collections.emptySet());
-        EasyMock.expect(metadata.fetch()).andReturn(cluster).once();
-
-        // Mock interceptors field
-        ProducerInterceptors interceptors = PowerMock.createMock(ProducerInterceptors.class);
-        EasyMock.expect(interceptors.onSend(record)).andReturn(record);
-        interceptors.onSendError(EasyMock.eq(record), EasyMock.notNull(), EasyMock.notNull());
-        EasyMock.expectLastCall();
-        MemberModifier.field(KafkaProducer.class, "interceptors").set(producer, interceptors);
-
-        PowerMock.replay(metadata);
-        EasyMock.replay(interceptors);
+        long nowMs = Time.SYSTEM.milliseconds();
+        ProducerMetadata metadata = newMetadata(0, 90000);
+        metadata.add(topic, nowMs);
+        MetadataResponse initialUpdateResponse = RequestTestUtils.metadataUpdateWith(1, singletonMap(topic, 1));
+        metadata.updateWithCurrentRequestVersion(initialUpdateResponse, false, nowMs);
+
+        @SuppressWarnings("unchecked") // it is safe to suppress, since this is a mock class
+                ProducerInterceptors interceptors = mock(ProducerInterceptors.class);
+        KafkaProducer producer = kafkaProducer(configs, new StringSerializer(),
+                new StringSerializer(), metadata, null, interceptors, Time.SYSTEM);
+
+        when(interceptors.onSend(any())).then(invocation -> invocation.getArgument(0));
+
         producer.send(record);
 
-        EasyMock.verify(interceptors);
+        verify(interceptors).onSend(record);
+        verify(interceptors).onSendError(eq(record), notNull(), notNull());
+
+        producer.close(Duration.ofMillis(0));
     }
 
     @Test
@@ -470,10 +785,580 @@ 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
+    public void testInitTransactionTimeout() {
+        Map configs = new HashMap<>();
+        configs.put(ProducerConfig.TRANSACTIONAL_ID_CONFIG, "bad-transaction");
+        configs.put(ProducerConfig.MAX_BLOCK_MS_CONFIG, 500);
+        configs.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9000");
+
+        Time time = new MockTime(1);
+        MetadataResponse initialUpdateResponse = RequestTestUtils.metadataUpdateWith(1, singletonMap("topic", 1));
+        ProducerMetadata metadata = newMetadata(0, Long.MAX_VALUE);
+        metadata.updateWithCurrentRequestVersion(initialUpdateResponse, false, time.milliseconds());
+
+        MockClient client = new MockClient(time, metadata);
+
+        try (Producer producer = kafkaProducer(configs, new StringSerializer(),
+                new StringSerializer(), metadata, client, null, time)) {
+            client.prepareResponse(
+                request -> request instanceof FindCoordinatorRequest &&
+                    ((FindCoordinatorRequest) request).data().keyType() == FindCoordinatorRequest.CoordinatorType.TRANSACTION.id(),
+                FindCoordinatorResponse.prepareResponse(Errors.NONE, host1));
+
+            assertThrows(TimeoutException.class, producer::initTransactions);
+
+            client.prepareResponse(
+                request -> request instanceof FindCoordinatorRequest &&
+                               ((FindCoordinatorRequest) request).data().keyType() == FindCoordinatorRequest.CoordinatorType.TRANSACTION.id(),
+                FindCoordinatorResponse.prepareResponse(Errors.NONE, host1));
+
+            client.prepareResponse(initProducerIdResponse(1L, (short) 5, Errors.NONE));
+
+            // retry initialization should work
+            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 = RequestTestUtils.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 = kafkaProducer(configs, new StringSerializer(),
+                new StringSerializer(), metadata, client, null, time)) {
+            producer.initTransactions();
+        }
+    }
+
+    @Test
+    public void testAbortTransaction() {
+        Map configs = new HashMap<>();
+        configs.put(ProducerConfig.TRANSACTIONAL_ID_CONFIG, "some.id");
+        configs.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9000");
+
+        Time time = new MockTime(1);
+        MetadataResponse initialUpdateResponse = RequestTestUtils.metadataUpdateWith(1, singletonMap("topic", 1));
+        ProducerMetadata metadata = newMetadata(0, Long.MAX_VALUE);
+
+        MockClient client = new MockClient(time, metadata);
+        client.updateMetadata(initialUpdateResponse);
+
+        client.prepareResponse(FindCoordinatorResponse.prepareResponse(Errors.NONE, host1));
+        client.prepareResponse(initProducerIdResponse(1L, (short) 5, Errors.NONE));
+        client.prepareResponse(endTxnResponse(Errors.NONE));
+
+        try (Producer producer = kafkaProducer(configs, new StringSerializer(),
+                new StringSerializer(), metadata, client, null, time)) {
+            producer.initTransactions();
+            producer.beginTransaction();
+            producer.abortTransaction();
+        }
+    }
+
+    @Test
+    public void testSendTxnOffsetsWithGroupId() {
+        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 = RequestTestUtils.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));
+        client.prepareResponse(addOffsetsToTxnResponse(Errors.NONE));
+        client.prepareResponse(FindCoordinatorResponse.prepareResponse(Errors.NONE, host1));
+        String groupId = "group";
+        client.prepareResponse(request ->
+            ((TxnOffsetCommitRequest) request).data().groupId().equals(groupId),
+            txnOffsetsCommitResponse(Collections.singletonMap(
+                new TopicPartition("topic", 0), Errors.NONE)));
+        client.prepareResponse(endTxnResponse(Errors.NONE));
+
+        try (Producer producer = kafkaProducer(configs, new StringSerializer(),
+            new StringSerializer(), metadata, client, null, time)) {
+            producer.initTransactions();
+            producer.beginTransaction();
+            producer.sendOffsetsToTransaction(Collections.emptyMap(), groupId);
+            producer.commitTransaction();
+        }
+    }
+
+    @Test
+    public void testSendTxnOffsetsWithGroupMetadata() {
+        sendOffsetsWithGroupMetadata((short) 3);
+    }
+
+    @Test
+    public void testSendTxnOffsetsWithGroupMetadataDowngrade() {
+        sendOffsetsWithGroupMetadata((short) 2);
+    }
+
+    private void sendOffsetsWithGroupMetadata(final short maxVersion) {
+        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");
+        configs.put(ProducerConfig.AUTO_DOWNGRADE_TXN_COMMIT, true);
+
+        Time time = new MockTime(1);
+        MetadataResponse initialUpdateResponse = RequestTestUtils.metadataUpdateWith(1, singletonMap("topic", 1));
+        ProducerMetadata metadata = newMetadata(0, Long.MAX_VALUE);
+
+        MockClient client = new MockClient(time, metadata);
+        client.updateMetadata(initialUpdateResponse);
+        client.setNodeApiVersions(NodeApiVersions.create(ApiKeys.TXN_OFFSET_COMMIT.id, (short) 0, maxVersion));
+
+        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));
+        client.prepareResponse(addOffsetsToTxnResponse(Errors.NONE));
+        client.prepareResponse(FindCoordinatorResponse.prepareResponse(Errors.NONE, host1));
+        String groupId = "group";
+        String memberId = "member";
+        int generationId = 5;
+        String groupInstanceId = "instance";
+        client.prepareResponse(request -> {
+            TxnOffsetCommitRequestData data = ((TxnOffsetCommitRequest) request).data();
+            if (maxVersion < 3) {
+                return data.groupId().equals(groupId) &&
+                           data.memberId().equals(JoinGroupRequest.UNKNOWN_MEMBER_ID) &&
+                           data.generationId() == JoinGroupRequest.UNKNOWN_GENERATION_ID &&
+                           data.groupInstanceId() == null;
+            } else {
+                return data.groupId().equals(groupId) &&
+                           data.memberId().equals(memberId) &&
+                           data.generationId() == generationId &&
+                           data.groupInstanceId().equals(groupInstanceId);
+            }
+        }, txnOffsetsCommitResponse(Collections.singletonMap(
+            new TopicPartition("topic", 0), Errors.NONE)));
+        client.prepareResponse(endTxnResponse(Errors.NONE));
+
+        try (Producer producer = kafkaProducer(configs, new StringSerializer(),
+            new StringSerializer(), metadata, client, null, time)) {
+            producer.initTransactions();
+            producer.beginTransaction();
+            ConsumerGroupMetadata groupMetadata = new ConsumerGroupMetadata(groupId,
+                generationId, memberId, Optional.of(groupInstanceId));
+
+            producer.sendOffsetsToTransaction(Collections.emptyMap(), groupMetadata);
+            producer.commitTransaction();
+        }
+    }
+
+    @Test
+    public void testNullGroupMetadataInSendOffsets() {
+        verifyInvalidGroupMetadata(null);
+    }
+
+    @Test
+    public void testInvalidGenerationIdAndMemberIdCombinedInSendOffsets() {
+        verifyInvalidGroupMetadata(new ConsumerGroupMetadata("group", 2, JoinGroupRequest.UNKNOWN_MEMBER_ID, Optional.empty()));
+    }
+
+    private void verifyInvalidGroupMetadata(ConsumerGroupMetadata groupMetadata) {
+        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 = RequestTestUtils.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 = kafkaProducer(configs, new StringSerializer(),
+            new StringSerializer(), metadata, client, null, time)) {
+            producer.initTransactions();
+            producer.beginTransaction();
+            assertThrows(IllegalArgumentException.class,
+                () -> producer.sendOffsetsToTransaction(Collections.emptyMap(), groupMetadata));
+        }
+    }
+
+    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 AddOffsetsToTxnResponse addOffsetsToTxnResponse(Errors error) {
+        return new AddOffsetsToTxnResponse(new AddOffsetsToTxnResponseData()
+                                               .setErrorCode(error.code())
+                                               .setThrottleTimeMs(10));
+    }
+
+    private TxnOffsetCommitResponse txnOffsetsCommitResponse(Map errorMap) {
+        return new TxnOffsetCommitResponse(10, errorMap);
+    }
+
+    private EndTxnResponse endTxnResponse(Errors error) {
+        return new EndTxnResponse(new EndTxnResponseData()
+                                      .setErrorCode(error.code())
+                                      .setThrottleTimeMs(0));
+    }
+
+    @Test
+    public void testOnlyCanExecuteCloseAfterInitTransactionsTimeout() {
+        Map configs = new HashMap<>();
+        configs.put(ProducerConfig.TRANSACTIONAL_ID_CONFIG, "bad-transaction");
+        configs.put(ProducerConfig.MAX_BLOCK_MS_CONFIG, 5);
+        configs.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9000");
+
+        Time time = new MockTime();
+        MetadataResponse initialUpdateResponse = RequestTestUtils.metadataUpdateWith(1, singletonMap("topic", 1));
+        ProducerMetadata metadata = newMetadata(0, Long.MAX_VALUE);
+        metadata.updateWithCurrentRequestVersion(initialUpdateResponse, false, time.milliseconds());
+
+        MockClient client = new MockClient(time, metadata);
+
+        Producer producer = kafkaProducer(configs, new StringSerializer(), new StringSerializer(),
+                metadata, client, null, time);
+        assertThrows(TimeoutException.class, producer::initTransactions);
+        // other transactional operations should not be allowed if we catch the error after initTransactions failed
+        try {
+            assertThrows(KafkaException.class, producer::beginTransaction);
+        } finally {
+            producer.close(Duration.ofMillis(0));
+        }
+    }
+
+    @Test
+    public void testSendToInvalidTopic() throws Exception {
+        Map configs = new HashMap<>();
+        configs.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9000");
+        configs.put(ProducerConfig.MAX_BLOCK_MS_CONFIG, "15000");
+
+        Time time = new MockTime();
+        MetadataResponse initialUpdateResponse = RequestTestUtils.metadataUpdateWith(1, emptyMap());
+        ProducerMetadata metadata = newMetadata(0, Long.MAX_VALUE);
+        metadata.updateWithCurrentRequestVersion(initialUpdateResponse, false, time.milliseconds());
+
+        MockClient client = new MockClient(time, metadata);
+
+        Producer producer = kafkaProducer(configs, new StringSerializer(), new StringSerializer(),
+                metadata, client, null, time);
+
+        String invalidTopicName = "topic abc"; // Invalid topic name due to space
+        ProducerRecord record = new ProducerRecord<>(invalidTopicName, "HelloKafka");
+
+        List topicMetadata = new ArrayList<>();
+        topicMetadata.add(new MetadataResponse.TopicMetadata(Errors.INVALID_TOPIC_EXCEPTION,
+                invalidTopicName, false, Collections.emptyList()));
+        MetadataResponse updateResponse =  RequestTestUtils.metadataResponse(
+                new ArrayList<>(initialUpdateResponse.brokers()),
+                initialUpdateResponse.clusterId(),
+                initialUpdateResponse.controller().id(),
+                topicMetadata);
+        client.prepareMetadataUpdate(updateResponse);
+
+        Future future = producer.send(record);
+
+        assertEquals("Cluster has incorrect invalid topic list.", Collections.singleton(invalidTopicName),
+                metadata.fetch().invalidTopics());
+        TestUtils.assertFutureError(future, InvalidTopicException.class);
+
+        producer.close(Duration.ofMillis(0));
+    }
+
+    @Test
+    public void testCloseWhenWaitingForMetadataUpdate() throws InterruptedException {
+        Map configs = new HashMap<>();
+        configs.put(ProducerConfig.MAX_BLOCK_MS_CONFIG, Long.MAX_VALUE);
+        configs.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9000");
+
+        // Simulate a case where metadata for a particular topic is not available. This will cause KafkaProducer#send to
+        // block in Metadata#awaitUpdate for the configured max.block.ms. When close() is invoked, KafkaProducer#send should
+        // return with a KafkaException.
+        String topicName = "test";
+        Time time = Time.SYSTEM;
+        MetadataResponse initialUpdateResponse = RequestTestUtils.metadataUpdateWith(1, emptyMap());
+        ProducerMetadata metadata = new ProducerMetadata(0, Long.MAX_VALUE, Long.MAX_VALUE,
+                new LogContext(), new ClusterResourceListeners(), time);
+        metadata.updateWithCurrentRequestVersion(initialUpdateResponse, false, time.milliseconds());
+        MockClient client = new MockClient(time, metadata);
+
+        Producer producer = kafkaProducer(configs, new StringSerializer(), new StringSerializer(),
+                metadata, client, null, time);
+
+        ExecutorService executor = Executors.newSingleThreadExecutor();
+        final AtomicReference sendException = new AtomicReference<>();
+
+        try {
+            executor.submit(() -> {
+                try {
+                    // Metadata for topic "test" will not be available which will cause us to block indefinitely until
+                    // KafkaProducer#close is invoked.
+                    producer.send(new ProducerRecord<>(topicName, "key", "value"));
+                    fail();
+                } catch (Exception e) {
+                    sendException.set(e);
+                }
+            });
+
+            // Wait until metadata update for the topic has been requested
+            TestUtils.waitForCondition(() -> metadata.containsTopic(topicName),
+                    "Timeout when waiting for topic to be added to metadata");
+            producer.close(Duration.ofMillis(0));
+            TestUtils.waitForCondition(() -> sendException.get() != null, "No producer exception within timeout");
+            assertEquals(KafkaException.class, sendException.get().getClass());
+        } finally {
+            executor.shutdownNow();
+        }
+    }
+
+    @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 = RequestTestUtils.metadataUpdateWith(1, emptyMap());
+        ProducerMetadata metadata = newMetadata(0, Long.MAX_VALUE);
+        metadata.updateWithCurrentRequestVersion(initialUpdateResponse, false, time.milliseconds());
+
+        MockClient client = new MockClient(time, metadata);
+
+        Producer producer = 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 = RequestTestUtils.metadataUpdateWith(1, singletonMap("testTopic", 1));
+        ProducerMetadata metadata = newMetadata(0, Long.MAX_VALUE);
+        metadata.updateWithCurrentRequestVersion(initialUpdateResponse, false, time.milliseconds());
+
+        MockClient client = new MockClient(time, metadata);
+
+        Producer producer = 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 = RequestTestUtils.metadataUpdateWith(1, singletonMap("testTopic", 1));
+        ProducerMetadata metadata = newMetadata(0, Long.MAX_VALUE);
+        metadata.updateWithCurrentRequestVersion(initialUpdateResponse, false, time.milliseconds());
+
+        MockClient client = new MockClient(time, metadata);
+
+        Producer producer = 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 = RequestTestUtils.metadataUpdateWith(1, singletonMap("testTopic", 1));
+        ProducerMetadata metadata = newMetadata(0, Long.MAX_VALUE);
+        metadata.updateWithCurrentRequestVersion(initialUpdateResponse, false, time.milliseconds());
+
+        MockClient client = new MockClient(time, metadata);
+
+        Producer producer = 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 testProducerJmxPrefix() throws  Exception {
+        Map props = new HashMap<>();
+        props.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
+        props.put("client.id", "client-1");
+
+        KafkaProducer producer = new KafkaProducer<>(
+                props, new StringSerializer(), new StringSerializer());
+
+        MBeanServer server = ManagementFactory.getPlatformMBeanServer();
+        MetricName testMetricName = producer.metrics.metricName("test-metric",
+                "grp1", "test metric");
+        producer.metrics.addMetric(testMetricName, new Avg());
+        Assert.assertNotNull(server.getObjectInstance(new ObjectName("kafka.producer:type=grp1,client-id=client-1")));
+        producer.close();
+    }
+
+    private static ProducerMetadata newMetadata(long refreshBackoffMs, long expirationMs) {
+        return new ProducerMetadata(refreshBackoffMs, expirationMs, DEFAULT_METADATA_IDLE_MS,
+                new LogContext(), new ClusterResourceListeners(), Time.SYSTEM);
+    }
+
+    @Test
+    public void configurableObjectsShouldSeeGeneratedClientId() {
+        Properties props = new Properties();
+        props.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
+        props.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, SerializerForClientId.class.getName());
+        props.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, SerializerForClientId.class.getName());
+        props.put(ProducerConfig.PARTITIONER_CLASS_CONFIG, PartitionerForClientId.class.getName());
+        props.put(ProducerConfig.INTERCEPTOR_CLASSES_CONFIG, ProducerInterceptorForClientId.class.getName());
+
+        KafkaProducer producer = new KafkaProducer<>(props);
+        assertNotNull(producer.getClientId());
+        assertNotEquals(0, producer.getClientId().length());
+        assertEquals(4, CLIENT_IDS.size());
+        CLIENT_IDS.forEach(id -> assertEquals(id, producer.getClientId()));
+        producer.close();
+    }
+
+    @Test
+    public void testUnusedConfigs() {
+        Map props = new HashMap<>();
+        props.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
+        props.put(SslConfigs.SSL_PROTOCOL_CONFIG, "TLS");
+        ProducerConfig config = new ProducerConfig(ProducerConfig.appendSerializerToConfig(props,
+                new StringSerializer(), new StringSerializer()));
+
+        assertTrue(config.unused().contains(SslConfigs.SSL_PROTOCOL_CONFIG));
+
+        try (KafkaProducer producer = new KafkaProducer<>(config, null, null,
+                null, null, null, Time.SYSTEM)) {
+            assertTrue(config.unused().contains(SslConfigs.SSL_PROTOCOL_CONFIG));
+        }
+    }
+
+    private static final List CLIENT_IDS = new ArrayList<>();
+
+    public static class SerializerForClientId implements Serializer {
+        @Override
+        public void configure(Map configs, boolean isKey) {
+            CLIENT_IDS.add(configs.get(ProducerConfig.CLIENT_ID_CONFIG).toString());
+        }
+
+        @Override
+        public byte[] serialize(String topic, byte[] data) {
+            return data;
+        }
+    }
+
+    public static class PartitionerForClientId implements Partitioner {
+
+        @Override
+        public int partition(String topic, Object key, byte[] keyBytes, Object value, byte[] valueBytes, Cluster cluster) {
+            return 0;
+        }
+
+        @Override
+        public void close() {
+
+        }
+
+        @Override
+        public void configure(Map configs) {
+            CLIENT_IDS.add(configs.get(ProducerConfig.CLIENT_ID_CONFIG).toString());
+        }
+    }
+
+    public static class ProducerInterceptorForClientId implements ProducerInterceptor {
+
+        @Override
+        public ProducerRecord onSend(ProducerRecord record) {
+            return record;
+        }
+
+        @Override
+        public void onAcknowledgement(RecordMetadata metadata, Exception exception) {
+        }
+
+        @Override
+        public void close() {
+        }
+
+        @Override
+        public void configure(Map configs) {
+            CLIENT_IDS.add(configs.get(ProducerConfig.CLIENT_ID_CONFIG).toString());
         }
     }
 }
diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/MockProducerTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/MockProducerTest.java
index 27fac280afcbb..d1239c91eb3c1 100644
--- a/clients/src/test/java/org/apache/kafka/clients/producer/MockProducerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/producer/MockProducerTest.java
@@ -16,13 +16,15 @@
  */
 package org.apache.kafka.clients.producer;
 
+import org.apache.kafka.clients.consumer.ConsumerGroupMetadata;
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.clients.producer.internals.DefaultPartitioner;
 import org.apache.kafka.common.Cluster;
-import org.apache.kafka.common.Node;
+import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.PartitionInfo;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.errors.ProducerFencedException;
+import org.apache.kafka.common.serialization.IntegerSerializer;
 import org.apache.kafka.common.serialization.StringSerializer;
 import org.apache.kafka.test.MockSerializer;
 import org.junit.After;
@@ -42,6 +44,7 @@
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
@@ -51,6 +54,7 @@ public class MockProducerTest {
     private MockProducer producer;
     private final ProducerRecord record1 = new ProducerRecord<>(topic, "key1".getBytes(), "value1".getBytes());
     private final ProducerRecord record2 = new ProducerRecord<>(topic, "key2".getBytes(), "value2".getBytes());
+    private final String groupId = "group";
 
     private void buildMockProducer(boolean autoComplete) {
         this.producer = new MockProducer<>(autoComplete, new MockSerializer(), new MockSerializer());
@@ -79,8 +83,8 @@ public void testAutoCompleteMock() throws Exception {
     public void testPartitioner() throws Exception {
         PartitionInfo partitionInfo0 = new PartitionInfo(topic, 0, null, null, null);
         PartitionInfo partitionInfo1 = new PartitionInfo(topic, 1, null, null, null);
-        Cluster cluster = new Cluster(null, new ArrayList(0), asList(partitionInfo0, partitionInfo1),
-                Collections.emptySet(), Collections.emptySet());
+        Cluster cluster = new Cluster(null, new ArrayList<>(0), asList(partitionInfo0, partitionInfo1),
+                Collections.emptySet(), Collections.emptySet());
         MockProducer producer = new MockProducer<>(cluster, true, new DefaultPartitioner(), new StringSerializer(), new StringSerializer());
         ProducerRecord record = new ProducerRecord<>(topic, "key", "value");
         Future metadata = producer.send(record);
@@ -148,10 +152,18 @@ public void shouldBeginTransactions() {
         assertTrue(producer.transactionInFlight());
     }
 
+    @Test(expected = IllegalStateException.class)
+    public void shouldThrowOnBeginTransactionsIfTransactionInflight() {
+        buildMockProducer(true);
+        producer.initTransactions();
+        producer.beginTransaction();
+        producer.beginTransaction();
+    }
+
     @Test(expected = IllegalStateException.class)
     public void shouldThrowOnSendOffsetsToTransactionIfTransactionsNotInitialized() {
         buildMockProducer(true);
-        producer.sendOffsetsToTransaction(null, null);
+        producer.sendOffsetsToTransaction(null, groupId);
     }
 
     @Test
@@ -159,7 +171,7 @@ public void shouldThrowOnSendOffsetsToTransactionTransactionIfNoTransactionGotSt
         buildMockProducer(true);
         producer.initTransactions();
         try {
-            producer.sendOffsetsToTransaction(null, null);
+            producer.sendOffsetsToTransaction(null, groupId);
             fail("Should have thrown as producer has no open transaction");
         } catch (IllegalStateException e) { }
     }
@@ -267,27 +279,29 @@ public void shouldThrowOnSendIfProducerGotFenced() {
         try {
             producer.send(null);
             fail("Should have thrown as producer is fenced off");
-        } catch (ProducerFencedException e) { }
+        } catch (KafkaException e) {
+            assertTrue("The root cause of the exception should be ProducerFenced", e.getCause() instanceof ProducerFencedException);
+        }
     }
 
     @Test
-    public void shouldThrowOnFlushIfProducerGotFenced() {
+    public void shouldThrowOnSendOffsetsToTransactionByGroupIdIfProducerGotFenced() {
         buildMockProducer(true);
         producer.initTransactions();
         producer.fenceProducer();
         try {
-            producer.flush();
+            producer.sendOffsetsToTransaction(null, groupId);
             fail("Should have thrown as producer is fenced off");
         } catch (ProducerFencedException e) { }
     }
 
     @Test
-    public void shouldThrowOnSendOffsetsToTransactionIfProducerGotFenced() {
+    public void shouldThrowOnSendOffsetsToTransactionByGroupMetadataIfProducerGotFenced() {
         buildMockProducer(true);
         producer.initTransactions();
         producer.fenceProducer();
         try {
-            producer.sendOffsetsToTransaction(null, null);
+            producer.sendOffsetsToTransaction(null, new ConsumerGroupMetadata(groupId));
             fail("Should have thrown as producer is fenced off");
         } catch (ProducerFencedException e) { }
     }
@@ -441,13 +455,26 @@ public void shouldThrowOnNullConsumerGroupIdWhenSendOffsetsToTransaction() {
         producer.beginTransaction();
 
         try {
-            producer.sendOffsetsToTransaction(Collections.emptyMap(), null);
+            String consumerGroupId = null;
+            producer.sendOffsetsToTransaction(Collections.emptyMap(), consumerGroupId);
+            fail("Should have thrown NullPointerException");
+        } catch (NullPointerException e) { }
+    }
+
+    @Test
+    public void shouldThrowOnNullConsumerGroupMetadataWhenSendOffsetsToTransaction() {
+        buildMockProducer(true);
+        producer.initTransactions();
+        producer.beginTransaction();
+
+        try {
+            producer.sendOffsetsToTransaction(Collections.emptyMap(), new ConsumerGroupMetadata(null));
             fail("Should have thrown NullPointerException");
         } catch (NullPointerException e) { }
     }
 
     @Test
-    public void shouldIgnoreEmptyOffsetsWhenSendOffsetsToTransaction() {
+    public void shouldIgnoreEmptyOffsetsWhenSendOffsetsToTransactionByGroupId() {
         buildMockProducer(true);
         producer.initTransactions();
         producer.beginTransaction();
@@ -456,7 +483,16 @@ public void shouldIgnoreEmptyOffsetsWhenSendOffsetsToTransaction() {
     }
 
     @Test
-    public void shouldAddOffsetsWhenSendOffsetsToTransaction() {
+    public void shouldIgnoreEmptyOffsetsWhenSendOffsetsToTransactionByGroupMetadata() {
+        buildMockProducer(true);
+        producer.initTransactions();
+        producer.beginTransaction();
+        producer.sendOffsetsToTransaction(Collections.emptyMap(), new ConsumerGroupMetadata("groupId"));
+        assertFalse(producer.sentOffsets());
+    }
+
+    @Test
+    public void shouldAddOffsetsWhenSendOffsetsToTransactionByGroupId() {
         buildMockProducer(true);
         producer.initTransactions();
         producer.beginTransaction();
@@ -472,6 +508,23 @@ public void shouldAddOffsetsWhenSendOffsetsToTransaction() {
         assertTrue(producer.sentOffsets());
     }
 
+    @Test
+    public void shouldAddOffsetsWhenSendOffsetsToTransactionByGroupMetadata() {
+        buildMockProducer(true);
+        producer.initTransactions();
+        producer.beginTransaction();
+
+        assertFalse(producer.sentOffsets());
+
+        Map groupCommit = new HashMap() {
+            {
+                put(new TopicPartition(topic, 0), new OffsetAndMetadata(42L, null));
+            }
+        };
+        producer.sendOffsetsToTransaction(groupCommit, new ConsumerGroupMetadata("groupId"));
+        assertTrue(producer.sentOffsets());
+    }
+
     @Test
     public void shouldResetSentOffsetsFlagOnlyWhenBeginningNewTransaction() {
         buildMockProducer(true);
@@ -491,6 +544,13 @@ public void shouldResetSentOffsetsFlagOnlyWhenBeginningNewTransaction() {
 
         producer.beginTransaction();
         assertFalse(producer.sentOffsets());
+
+        producer.sendOffsetsToTransaction(groupCommit, new ConsumerGroupMetadata("groupId"));
+        producer.commitTransaction(); // commit should not reset "sentOffsets" flag
+        assertTrue(producer.sentOffsets());
+
+        producer.beginTransaction();
+        assertFalse(producer.sentOffsets());
     }
 
     @Test
@@ -513,7 +573,7 @@ public void shouldPublishLatestAndCumulativeConsumerGroupOffsetsOnlyAfterCommitI
             }
         };
         producer.sendOffsetsToTransaction(groupCommit1, group);
-        producer.sendOffsetsToTransaction(groupCommit2, group);
+        producer.sendOffsetsToTransaction(groupCommit2, new ConsumerGroupMetadata(group));
 
         assertTrue(producer.consumerGroupOffsetsHistory().isEmpty());
 
@@ -549,10 +609,18 @@ public void shouldDropConsumerGroupOffsetsOnAbortIfTransactionsAreEnabled() {
         producer.beginTransaction();
         producer.commitTransaction();
         assertTrue(producer.consumerGroupOffsetsHistory().isEmpty());
+
+        producer.beginTransaction();
+        producer.sendOffsetsToTransaction(groupCommit, new ConsumerGroupMetadata(group));
+        producer.abortTransaction();
+
+        producer.beginTransaction();
+        producer.commitTransaction();
+        assertTrue(producer.consumerGroupOffsetsHistory().isEmpty());
     }
 
     @Test
-    public void shouldPreserveCommittedConsumerGroupsOffsetsOnAbortIfTransactionsAreEnabled() {
+    public void shouldPreserveOffsetsFromCommitByGroupIdOnAbortIfTransactionsAreEnabled() {
         buildMockProducer(true);
         producer.initTransactions();
         producer.beginTransaction();
@@ -576,6 +644,40 @@ public void shouldPreserveCommittedConsumerGroupsOffsetsOnAbortIfTransactionsAre
         assertThat(producer.consumerGroupOffsetsHistory(), equalTo(Collections.singletonList(expectedResult)));
     }
 
+    @Test
+    public void shouldPreserveOffsetsFromCommitByGroupMetadataOnAbortIfTransactionsAreEnabled() {
+        buildMockProducer(true);
+        producer.initTransactions();
+        producer.beginTransaction();
+
+        String group = "g";
+        Map groupCommit = new HashMap() {
+            {
+                put(new TopicPartition(topic, 0), new OffsetAndMetadata(42L, null));
+                put(new TopicPartition(topic, 1), new OffsetAndMetadata(73L, null));
+            }
+        };
+        producer.sendOffsetsToTransaction(groupCommit, new ConsumerGroupMetadata(group));
+        producer.commitTransaction();
+
+        producer.beginTransaction();
+
+        String group2 = "g2";
+        Map groupCommit2 = new HashMap() {
+            {
+                put(new TopicPartition(topic, 2), new OffsetAndMetadata(53L, null));
+                put(new TopicPartition(topic, 3), new OffsetAndMetadata(84L, null));
+            }
+        };
+        producer.sendOffsetsToTransaction(groupCommit, new ConsumerGroupMetadata(group2));
+        producer.abortTransaction();
+
+        Map> expectedResult = new HashMap<>();
+        expectedResult.put(group, groupCommit);
+
+        assertThat(producer.consumerGroupOffsetsHistory(), equalTo(Collections.singletonList(expectedResult)));
+    }
+
     @Test
     public void shouldThrowOnInitTransactionIfProducerIsClosed() {
         buildMockProducer(true);
@@ -607,41 +709,41 @@ public void shouldThrowOnBeginTransactionIfProducerIsClosed() {
     }
 
     @Test
-    public void shouldThrowSendOffsetsToTransactionIfProducerIsClosed() {
+    public void shouldThrowSendOffsetsToTransactionByGroupIdIfProducerIsClosed() {
         buildMockProducer(true);
         producer.close();
         try {
-            producer.sendOffsetsToTransaction(null, null);
+            producer.sendOffsetsToTransaction(null, groupId);
             fail("Should have thrown as producer is already closed");
         } catch (IllegalStateException e) { }
     }
 
     @Test
-    public void shouldThrowOnCommitTransactionIfProducerIsClosed() {
+    public void shouldThrowSendOffsetsToTransactionByGroupMetadataIfProducerIsClosed() {
         buildMockProducer(true);
         producer.close();
         try {
-            producer.commitTransaction();
+            producer.sendOffsetsToTransaction(null, new ConsumerGroupMetadata(groupId));
             fail("Should have thrown as producer is already closed");
         } catch (IllegalStateException e) { }
     }
 
     @Test
-    public void shouldThrowOnAbortTransactionIfProducerIsClosed() {
+    public void shouldThrowOnCommitTransactionIfProducerIsClosed() {
         buildMockProducer(true);
         producer.close();
         try {
-            producer.abortTransaction();
+            producer.commitTransaction();
             fail("Should have thrown as producer is already closed");
         } catch (IllegalStateException e) { }
     }
 
     @Test
-    public void shouldThrowOnCloseIfProducerIsClosed() {
+    public void shouldThrowOnAbortTransactionIfProducerIsClosed() {
         buildMockProducer(true);
         producer.close();
         try {
-            producer.close();
+            producer.abortTransaction();
             fail("Should have thrown as producer is already closed");
         } catch (IllegalStateException e) { }
     }
@@ -665,6 +767,14 @@ public void shouldThrowOnFlushProducerIfProducerIsClosed() {
             fail("Should have thrown as producer is already closed");
         } catch (IllegalStateException e) { }
     }
+    
+    @Test
+    @SuppressWarnings("unchecked")
+    public void shouldThrowClassCastException() {
+        try (MockProducer customProducer = new MockProducer<>(true, new IntegerSerializer(), new StringSerializer());) {
+            assertThrows(ClassCastException.class, () -> customProducer.send(new ProducerRecord(topic, "key1", "value1")));
+        }
+    }
 
     @Test
     public void shouldBeFlushedIfNoBufferedRecords() {
diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/ProducerConfigTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/ProducerConfigTest.java
new file mode 100644
index 0000000000000..d538a64641360
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/producer/ProducerConfigTest.java
@@ -0,0 +1,93 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.serialization.ByteArraySerializer;
+import org.apache.kafka.common.serialization.Serializer;
+import org.apache.kafka.common.serialization.StringSerializer;
+import org.junit.Test;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Properties;
+
+import static org.junit.Assert.assertEquals;
+
+public class ProducerConfigTest {
+
+    private final Serializer keySerializer = new ByteArraySerializer();
+    private final Serializer valueSerializer = new StringSerializer();
+    private final String keySerializerClassName = keySerializer.getClass().getName();
+    private final String valueSerializerClassName = valueSerializer.getClass().getName();
+    private final Object keySerializerClass = keySerializer.getClass();
+    private final Object valueSerializerClass = valueSerializer.getClass();
+
+    @SuppressWarnings("deprecation")
+    @Test
+    public void testSerializerToPropertyConfig() {
+        Properties properties = new Properties();
+        properties.setProperty(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, keySerializerClassName);
+        properties.setProperty(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, valueSerializerClassName);
+        Properties newProperties = ProducerConfig.addSerializerToConfig(properties, null, null);
+        assertEquals(newProperties.get(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG), keySerializerClassName);
+        assertEquals(newProperties.get(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG), valueSerializerClassName);
+
+        properties.clear();
+        properties.setProperty(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, valueSerializerClassName);
+        newProperties = ProducerConfig.addSerializerToConfig(properties, keySerializer, null);
+        assertEquals(newProperties.get(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG), keySerializerClassName);
+        assertEquals(newProperties.get(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG), valueSerializerClassName);
+
+        properties.clear();
+        properties.setProperty(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, keySerializerClassName);
+        newProperties = ProducerConfig.addSerializerToConfig(properties, null, valueSerializer);
+        assertEquals(newProperties.get(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG), keySerializerClassName);
+        assertEquals(newProperties.get(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG), valueSerializerClassName);
+
+        properties.clear();
+        newProperties = ProducerConfig.addSerializerToConfig(properties, keySerializer, valueSerializer);
+        assertEquals(newProperties.get(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG), keySerializerClassName);
+        assertEquals(newProperties.get(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG), valueSerializerClassName);
+    }
+    
+    @Test
+    public void testAppendSerializerToConfig() {
+        Map configs = new HashMap<>();
+        configs.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, keySerializerClass);
+        configs.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, valueSerializerClass);
+        Map newConfigs = ProducerConfig.appendSerializerToConfig(configs, null, null);
+        assertEquals(newConfigs.get(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG), keySerializerClass);
+        assertEquals(newConfigs.get(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG), valueSerializerClass);
+
+        configs.clear();
+        configs.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, valueSerializerClass);
+        newConfigs = ProducerConfig.appendSerializerToConfig(configs, keySerializer, null);
+        assertEquals(newConfigs.get(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG), keySerializerClass);
+        assertEquals(newConfigs.get(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG), valueSerializerClass);
+
+        configs.clear();
+        configs.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, keySerializerClass);
+        newConfigs = ProducerConfig.appendSerializerToConfig(configs, null, valueSerializer);
+        assertEquals(newConfigs.get(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG), keySerializerClass);
+        assertEquals(newConfigs.get(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG), valueSerializerClass);
+
+        configs.clear();
+        newConfigs = ProducerConfig.appendSerializerToConfig(configs, keySerializer, valueSerializer);
+        assertEquals(newConfigs.get(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG), keySerializerClass);
+        assertEquals(newConfigs.get(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG), valueSerializerClass);
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/clients/producer/RecordSendTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/RecordSendTest.java
index c083db3f29118..45be1171c8f07 100644
--- a/clients/src/test/java/org/apache/kafka/clients/producer/RecordSendTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/producer/RecordSendTest.java
@@ -31,6 +31,7 @@
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.errors.CorruptRecordException;
 import org.apache.kafka.common.record.RecordBatch;
+import org.apache.kafka.common.utils.Time;
 import org.junit.Test;
 
 public class RecordSendTest {
@@ -46,7 +47,7 @@ public class RecordSendTest {
     public void testTimeout() throws Exception {
         ProduceRequestResult request = new ProduceRequestResult(topicPartition);
         FutureRecordMetadata future = new FutureRecordMetadata(request, relOffset,
-                RecordBatch.NO_TIMESTAMP, 0L, 0, 0);
+                RecordBatch.NO_TIMESTAMP, 0L, 0, 0, Time.SYSTEM);
         assertFalse("Request is not completed", future.isDone());
         try {
             future.get(5, TimeUnit.MILLISECONDS);
@@ -66,7 +67,7 @@ public void testTimeout() throws Exception {
     @Test(expected = ExecutionException.class)
     public void testError() throws Exception {
         FutureRecordMetadata future = new FutureRecordMetadata(asyncRequest(baseOffset, new CorruptRecordException(), 50L),
-                relOffset, RecordBatch.NO_TIMESTAMP, 0L, 0, 0);
+                relOffset, RecordBatch.NO_TIMESTAMP, 0L, 0, 0, Time.SYSTEM);
         future.get();
     }
 
@@ -76,7 +77,7 @@ public void testError() throws Exception {
     @Test
     public void testBlocking() throws Exception {
         FutureRecordMetadata future = new FutureRecordMetadata(asyncRequest(baseOffset, null, 50L),
-                relOffset, RecordBatch.NO_TIMESTAMP, 0L, 0, 0);
+                relOffset, RecordBatch.NO_TIMESTAMP, 0L, 0, 0, Time.SYSTEM);
         assertEquals(baseOffset + relOffset, future.get().offset());
     }
 
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 8bcc775df560f..5a3bce8c453ef 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,11 +16,9 @@
  */
 package org.apache.kafka.clients.producer.internals;
 
-import org.apache.kafka.common.MetricName;
-import org.apache.kafka.common.errors.TimeoutException;
+import org.apache.kafka.clients.producer.BufferExhaustedException;
+import org.apache.kafka.common.KafkaException;
 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.MockTime;
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.test.TestUtils;
@@ -31,34 +29,29 @@
 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;
 import java.util.concurrent.locks.Condition;
-import org.junit.runner.RunWith;
-import org.powermock.core.classloader.annotations.PrepareForTest;
-import org.powermock.modules.junit4.PowerMockRunner;
-
-import static org.easymock.EasyMock.eq;
-import static org.easymock.EasyMock.createNiceMock;
-import static org.easymock.EasyMock.replay;
-import static org.easymock.EasyMock.anyLong;
-import static org.easymock.EasyMock.anyDouble;
-import static org.easymock.EasyMock.expectLastCall;
-import static org.easymock.EasyMock.expect;
-import static org.easymock.EasyMock.anyString;
+
 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;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.verify;
 
-
-@RunWith(PowerMockRunner.class)
 public class BufferPoolTest {
     private final MockTime time = new MockTime();
     private final Metrics metrics = new Metrics(time);
-    private final long maxBlockTimeMs = 2000;
+    private final long maxBlockTimeMs = 10;
     private final String metricGroup = "TestMetrics";
 
     @After
@@ -123,50 +116,54 @@ public void testDelayedAllocation() throws Exception {
 
     private CountDownLatch asyncDeallocate(final BufferPool pool, final ByteBuffer buffer) {
         final CountDownLatch latch = new CountDownLatch(1);
-        Thread thread = new Thread() {
-            public void run() {
-                try {
-                    latch.await();
-                } catch (InterruptedException e) {
-                    e.printStackTrace();
-                }
-                pool.deallocate(buffer);
+        Thread thread = new Thread(() -> {
+            try {
+                latch.await();
+            } catch (InterruptedException e) {
+                e.printStackTrace();
             }
-        };
+            pool.deallocate(buffer);
+        });
         thread.start();
         return latch;
     }
 
     private void delayedDeallocate(final BufferPool pool, final ByteBuffer buffer, final long delayMs) {
-        Thread thread = new Thread() {
-            public void run() {
-                Time.SYSTEM.sleep(delayMs);
-                pool.deallocate(buffer);
-            }
-        };
+        Thread thread = new Thread(() -> {
+            Time.SYSTEM.sleep(delayMs);
+            pool.deallocate(buffer);
+        });
         thread.start();
     }
 
     private CountDownLatch asyncAllocate(final BufferPool pool, final int size) {
         final CountDownLatch completed = new CountDownLatch(1);
-        Thread thread = new Thread() {
-            public void run() {
-                try {
-                    pool.allocate(size, maxBlockTimeMs);
-                } catch (InterruptedException e) {
-                    e.printStackTrace();
-                } finally {
-                    completed.countDown();
-                }
+        Thread thread = new Thread(() -> {
+            try {
+                pool.allocate(size, maxBlockTimeMs);
+            } catch (InterruptedException e) {
+                e.printStackTrace();
+            } finally {
+                completed.countDown();
             }
-        };
+        });
         thread.start();
         return completed;
     }
 
     /**
-     * Test if Timeout exception is thrown when there is not enough memory to allocate and the elapsed time is greater than the max specified block time.
-     * And verify that the allocation should finish soon after the maxBlockTimeMs.
+     * Test if BufferExhausted exception is thrown when there is not enough memory to allocate and the elapsed
+     * time is greater than the max specified block time.
+     */
+    @Test(expected = BufferExhaustedException.class)
+    public void testBufferExhaustedExceptionIsThrown() throws Exception {
+        BufferPool pool = new BufferPool(2, 1, metrics, time, metricGroup);
+        pool.allocate(1, maxBlockTimeMs);
+        pool.allocate(2, maxBlockTimeMs);
+    }
+
+    /**
+     * Verify that a failed allocation attempt due to not enough memory finishes soon after the maxBlockTimeMs.
      */
     @Test
     public void testBlockTimeout() throws Exception {
@@ -174,22 +171,24 @@ public void testBlockTimeout() throws Exception {
         ByteBuffer buffer1 = pool.allocate(1, maxBlockTimeMs);
         ByteBuffer buffer2 = pool.allocate(1, maxBlockTimeMs);
         ByteBuffer buffer3 = pool.allocate(1, maxBlockTimeMs);
-        // First two buffers will be de-allocated within maxBlockTimeMs since the most recent de-allocation
+        // The first two buffers will be de-allocated within maxBlockTimeMs since the most recent allocation
         delayedDeallocate(pool, buffer1, maxBlockTimeMs / 2);
         delayedDeallocate(pool, buffer2, maxBlockTimeMs);
-        // The third buffer will be de-allocated after maxBlockTimeMs since the most recent de-allocation
+        // The third buffer will be de-allocated after maxBlockTimeMs since the most recent allocation
         delayedDeallocate(pool, buffer3, maxBlockTimeMs / 2 * 5);
 
         long beginTimeMs = Time.SYSTEM.milliseconds();
         try {
             pool.allocate(10, maxBlockTimeMs);
             fail("The buffer allocated more memory than its maximum value 10");
-        } catch (TimeoutException e) {
+        } catch (BufferExhaustedException e) {
             // this is good
         }
-        assertTrue("available memory" + pool.availableMemory(), pool.availableMemory() >= 9 && pool.availableMemory() <= 10);
-        long endTimeMs = Time.SYSTEM.milliseconds();
-        assertTrue("Allocation should finish not much later than maxBlockTimeMs", endTimeMs - beginTimeMs < maxBlockTimeMs + 1000);
+        // Thread scheduling sometimes means that deallocation varies by this point
+        assertTrue("available memory " + pool.availableMemory(), pool.availableMemory() >= 7 && pool.availableMemory() <= 10);
+        long durationMs = Time.SYSTEM.milliseconds() - beginTimeMs;
+        assertTrue("BufferExhaustedException should not throw before maxBlockTimeMs", durationMs >= maxBlockTimeMs);
+        assertTrue("BufferExhaustedException should throw soon after maxBlockTimeMs", durationMs < maxBlockTimeMs + 1000);
     }
 
     /**
@@ -202,10 +201,11 @@ public void testCleanupMemoryAvailabilityWaiterOnBlockTimeout() throws Exception
         try {
             pool.allocate(2, maxBlockTimeMs);
             fail("The buffer allocated more memory than its maximum value 2");
-        } catch (TimeoutException e) {
+        } catch (BufferExhaustedException e) {
             // this is good
         }
-        assertTrue(pool.queued() == 0);
+        assertEquals(0, pool.queued());
+        assertEquals(1, pool.availableMemory());
     }
 
     /**
@@ -242,35 +242,24 @@ public void testCleanupMemoryAvailabilityWaiterOnInterruption() throws Exception
         assertEquals(pool.queued(), 0);
     }
 
-    @PrepareForTest({Sensor.class, MetricName.class})
     @Test
     public void testCleanupMemoryAvailabilityOnMetricsException() throws Exception {
-        Metrics mockedMetrics = createNiceMock(Metrics.class);
-        Sensor mockedSensor = createNiceMock(Sensor.class);
-        MetricName metricName = createNiceMock(MetricName.class);
-        MetricName rateMetricName = createNiceMock(MetricName.class);
-        MetricName totalMetricName = createNiceMock(MetricName.class);
-
-        expect(mockedMetrics.sensor(BufferPool.WAIT_TIME_SENSOR_NAME)).andReturn(mockedSensor);
-
-        mockedSensor.record(anyDouble(), anyLong());
-        expectLastCall().andThrow(new OutOfMemoryError());
-        expect(mockedMetrics.metricName(anyString(), eq(metricGroup), anyString())).andReturn(metricName);
-        mockedSensor.add(new Meter(TimeUnit.NANOSECONDS, rateMetricName, totalMetricName));
-
-        replay(mockedMetrics, mockedSensor, metricName);
+        BufferPool bufferPool = spy(new BufferPool(2, 1, new Metrics(), time, metricGroup));
+        doThrow(new OutOfMemoryError()).when(bufferPool).recordWaitTime(anyLong());
 
-        BufferPool bufferPool = new BufferPool(2, 1, mockedMetrics, time,  metricGroup);
         bufferPool.allocate(1, 0);
         try {
             bufferPool.allocate(2, 1000);
-            assertTrue("Expected oom.", false);
+            fail("Expected oom.");
         } catch (OutOfMemoryError expected) {
         }
         assertEquals(1, bufferPool.availableMemory());
         assertEquals(0, bufferPool.queued());
+        assertEquals(1, bufferPool.unallocatedMemory());
         //This shouldn't timeout
         bufferPool.allocate(1, 0);
+
+        verify(bufferPool).recordWaitTime(anyLong());
     }
 
     private static class BufferPoolAllocator implements Runnable {
@@ -287,7 +276,7 @@ public void run() {
             try {
                 pool.allocate(2, maxBlockTimeMs);
                 fail("The buffer allocated more memory than its maximum value 2");
-            } catch (TimeoutException e) {
+            } catch (BufferExhaustedException e) {
                 // this is good
             } catch (InterruptedException e) {
                 // this can be neglected
@@ -377,6 +366,7 @@ public StressTestThread(BufferPool pool, int iterations) {
             this.pool = pool;
         }
 
+        @Override
         public void run() {
             try {
                 for (int i = 0; i < iterations; i++) {
@@ -397,4 +387,46 @@ 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);
+
+        ExecutorService executor = Executors.newFixedThreadPool(numWorkers);
+        Callable work = new Callable() {
+                public Void call() throws Exception {
+                    assertThrows(KafkaException.class, () -> pool.allocate(1, Long.MAX_VALUE));
+                    return null;
+                }
+            };
+        for (int i = 0; i < numWorkers; ++i) {
+            executor.submit(work);
+        }
+
+        TestUtils.waitForCondition(() -> pool.queued() == numWorkers, "Awaiting " + numWorkers + " workers to be blocked on allocation");
+
+        // Close the buffer pool. This should notify all waiters.
+        pool.close();
+
+        TestUtils.waitForCondition(() -> pool.queued() == 0, "Awaiting " + numWorkers + " workers to be interrupted from allocation");
+
+        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/FutureRecordMetadataTest.java b/clients/src/test/java/org/apache/kafka/clients/producer/internals/FutureRecordMetadataTest.java
new file mode 100644
index 0000000000000..aee24e8fc6459
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/FutureRecordMetadataTest.java
@@ -0,0 +1,82 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.clients.producer.internals;
+
+import org.apache.kafka.common.record.RecordBatch;
+import org.apache.kafka.common.utils.MockTime;
+import org.junit.Test;
+
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+public class FutureRecordMetadataTest {
+
+    private final MockTime time = new MockTime();
+
+    @Test
+    public void testFutureGetWithSeconds() throws ExecutionException, InterruptedException, TimeoutException {
+        ProduceRequestResult produceRequestResult = mockProduceRequestResult();
+        FutureRecordMetadata future = futureRecordMetadata(produceRequestResult);
+
+        ProduceRequestResult chainedProduceRequestResult = mockProduceRequestResult();
+        future.chain(futureRecordMetadata(chainedProduceRequestResult));
+
+        future.get(1L, TimeUnit.SECONDS);
+
+        verify(produceRequestResult).await(1L, TimeUnit.SECONDS);
+        verify(chainedProduceRequestResult).await(1000L, TimeUnit.MILLISECONDS);
+    }
+
+    @Test
+    public void testFutureGetWithMilliSeconds() throws ExecutionException, InterruptedException, TimeoutException {
+        ProduceRequestResult produceRequestResult = mockProduceRequestResult();
+        FutureRecordMetadata future = futureRecordMetadata(produceRequestResult);
+
+        ProduceRequestResult chainedProduceRequestResult = mockProduceRequestResult();
+        future.chain(futureRecordMetadata(chainedProduceRequestResult));
+
+        future.get(1000L, TimeUnit.MILLISECONDS);
+
+        verify(produceRequestResult).await(1000L, TimeUnit.MILLISECONDS);
+        verify(chainedProduceRequestResult).await(1000L, TimeUnit.MILLISECONDS);
+    }
+
+    private FutureRecordMetadata futureRecordMetadata(ProduceRequestResult produceRequestResult) {
+        return new FutureRecordMetadata(
+                produceRequestResult,
+                0,
+                RecordBatch.NO_TIMESTAMP,
+                0L,
+                0,
+                0,
+                time
+        );
+    }
+
+    private ProduceRequestResult mockProduceRequestResult() throws InterruptedException {
+        ProduceRequestResult mockProduceRequestResult = mock(ProduceRequestResult.class);
+        when(mockProduceRequestResult.await(anyLong(), any(TimeUnit.class))).thenReturn(true);
+        return mockProduceRequestResult;
+    }
+}
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 2f89d7949e702..f9887f9033e81 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
@@ -197,6 +197,9 @@ public void testSplitPreservesMagicAndCompressionType() {
                 if (compressionType == CompressionType.NONE && magic < MAGIC_VALUE_V2)
                     continue;
 
+                if (compressionType == CompressionType.ZSTD && magic < MAGIC_VALUE_V2)
+                    continue;
+
                 MemoryRecordsBuilder builder = MemoryRecords.builder(ByteBuffer.allocate(1024), magic,
                         compressionType, TimestampType.CREATE_TIME, 0L);
 
@@ -226,40 +229,30 @@ public void testSplitPreservesMagicAndCompressionType() {
     }
 
     /**
-     * A {@link ProducerBatch} configured using a very large linger value and a timestamp preceding its create
-     * time is interpreted correctly as not expired when the linger time is larger than the difference
-     * between now and create time by {@link ProducerBatch#maybeExpire(int, long, long, long, boolean)}.
+     * A {@link ProducerBatch} configured using a timestamp preceding its create time is interpreted correctly
+     * as not expired by {@link ProducerBatch#hasReachedDeliveryTimeout(long, long)}.
      */
     @Test
-    public void testLargeLingerOldNowExpire() {
+    public void testBatchExpiration() {
+        long deliveryTimeoutMs = 10240;
         ProducerBatch batch = new ProducerBatch(new TopicPartition("topic", 1), memoryRecordsBuilder, now);
         // Set `now` to 2ms before the create time.
-        assertFalse(batch.maybeExpire(10240, 100L, now - 2L, Long.MAX_VALUE, false));
+        assertFalse(batch.hasReachedDeliveryTimeout(deliveryTimeoutMs, now - 2));
+        // Set `now` to deliveryTimeoutMs.
+        assertTrue(batch.hasReachedDeliveryTimeout(deliveryTimeoutMs, now + deliveryTimeoutMs));
     }
 
     /**
-     * A {@link ProducerBatch} configured using a very large retryBackoff value with retry = true and a timestamp
-     * preceding its create time is interpreted correctly as not expired when the retryBackoff time is larger than the
-     * difference between now and create time by {@link ProducerBatch#maybeExpire(int, long, long, long, boolean)}.
+     * A {@link ProducerBatch} configured using a timestamp preceding its create time is interpreted correctly
+     * * as not expired by {@link ProducerBatch#hasReachedDeliveryTimeout(long, long)}.
      */
     @Test
-    public void testLargeRetryBackoffOldNowExpire() {
+    public void testBatchExpirationAfterReenqueue() {
         ProducerBatch batch = new ProducerBatch(new TopicPartition("topic", 1), memoryRecordsBuilder, now);
         // Set batch.retry = true
         batch.reenqueued(now);
         // Set `now` to 2ms before the create time.
-        assertFalse(batch.maybeExpire(10240, Long.MAX_VALUE, now - 2L, 10240L, false));
-    }
-
-    /**
-     * A {@link ProducerBatch#maybeExpire(int, long, long, long, boolean)} call with a now value before the create
-     * time of the ProducerBatch is correctly recognized as not expired when invoked with parameter isFull = true.
-     */
-    @Test
-    public void testLargeFullOldNowExpire() {
-        ProducerBatch batch = new ProducerBatch(new TopicPartition("topic", 1), memoryRecordsBuilder, now);
-        // Set `now` to 2ms before the create time.
-        assertFalse(batch.maybeExpire(10240, 10240L, now - 2L, 10240L, true));
+        assertFalse(batch.hasReachedDeliveryTimeout(10240, now - 2L));
     }
 
     @Test
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
new file mode 100644
index 0000000000000..98d09da2034d6
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/producer/internals/ProducerMetadataTest.java
@@ -0,0 +1,310 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.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;
+import org.apache.kafka.common.requests.RequestTestUtils;
+import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.common.utils.Time;
+import org.junit.After;
+import org.junit.Test;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicReference;
+
+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;
+
+public class ProducerMetadataTest {
+    private static final long METADATA_IDLE_MS = 60 * 1000;
+    private long refreshBackoffMs = 100;
+    private long metadataExpireMs = 1000;
+    private ProducerMetadata metadata = new ProducerMetadata(refreshBackoffMs, metadataExpireMs, METADATA_IDLE_MS,
+            new LogContext(), new ClusterResourceListeners(), Time.SYSTEM);
+    private AtomicReference backgroundError = new AtomicReference<>();
+
+    @After
+    public void tearDown() {
+        assertNull("Exception in background thread : " + backgroundError.get(), backgroundError.get());
+    }
+
+    @Test
+    public void testMetadata() throws Exception {
+        long time = Time.SYSTEM.milliseconds();
+        String topic = "my-topic";
+        metadata.add(topic, time);
+
+        metadata.updateWithCurrentRequestVersion(responseWithTopics(Collections.emptySet()), false, time);
+        assertTrue("No update needed.", metadata.timeToNextUpdate(time) > 0);
+        metadata.requestUpdate();
+        assertTrue("Still no updated needed due to backoff", metadata.timeToNextUpdate(time) > 0);
+        time += refreshBackoffMs;
+        assertEquals("Update needed now that backoff time expired", 0, metadata.timeToNextUpdate(time));
+        Thread t1 = asyncFetch(topic, 500);
+        Thread t2 = asyncFetch(topic, 500);
+        assertTrue("Awaiting update", t1.isAlive());
+        assertTrue("Awaiting update", t2.isAlive());
+        // Perform metadata update when an update is requested on the async fetch thread
+        // This simulates the metadata update sequence in KafkaProducer
+        while (t1.isAlive() || t2.isAlive()) {
+            if (metadata.timeToNextUpdate(time) == 0) {
+                metadata.updateWithCurrentRequestVersion(responseWithCurrentTopics(), false, time);
+                time += refreshBackoffMs;
+            }
+            Thread.sleep(1);
+        }
+        t1.join();
+        t2.join();
+        assertTrue("No update needed.", metadata.timeToNextUpdate(time) > 0);
+        time += metadataExpireMs;
+        assertEquals("Update needed due to stale metadata.", 0, metadata.timeToNextUpdate(time));
+    }
+
+    @Test
+    public void testMetadataAwaitAfterClose() throws InterruptedException {
+        long time = 0;
+        metadata.updateWithCurrentRequestVersion(responseWithCurrentTopics(), false, time);
+        assertTrue("No update needed.", metadata.timeToNextUpdate(time) > 0);
+        metadata.requestUpdate();
+        assertTrue("Still no updated needed due to backoff", metadata.timeToNextUpdate(time) > 0);
+        time += refreshBackoffMs;
+        assertEquals("Update needed now that backoff time expired", 0, metadata.timeToNextUpdate(time));
+        String topic = "my-topic";
+        metadata.close();
+        Thread t1 = asyncFetch(topic, 500);
+        t1.join();
+        assertEquals(KafkaException.class, backgroundError.get().getClass());
+        assertTrue(backgroundError.get().toString().contains("Requested metadata update after close"));
+        clearBackgroundError();
+    }
+
+    /**
+     * Tests that {@link org.apache.kafka.clients.producer.internals.ProducerMetadata#awaitUpdate(int, long)} doesn't
+     * wait forever with a max timeout value of 0
+     *
+     * @throws Exception
+     * @see KAFKA-1836
+     */
+    @Test
+    public void testMetadataUpdateWaitTime() throws Exception {
+        long time = 0;
+        metadata.updateWithCurrentRequestVersion(responseWithCurrentTopics(), false, time);
+        assertTrue("No update needed.", metadata.timeToNextUpdate(time) > 0);
+        // first try with a max wait time of 0 and ensure that this returns back without waiting forever
+        try {
+            metadata.awaitUpdate(metadata.requestUpdate(), 0);
+            fail("Wait on metadata update was expected to timeout, but it didn't");
+        } catch (TimeoutException te) {
+            // expected
+        }
+        // now try with a higher timeout value once
+        final long twoSecondWait = 2000;
+        try {
+            metadata.awaitUpdate(metadata.requestUpdate(), twoSecondWait);
+            fail("Wait on metadata update was expected to timeout, but it didn't");
+        } catch (TimeoutException te) {
+            // expected
+        }
+    }
+
+    @Test
+    public void testTimeToNextUpdateOverwriteBackoff() {
+        long now = 10000;
+
+        // New topic added to fetch set and update requested. It should allow immediate update.
+        metadata.updateWithCurrentRequestVersion(responseWithCurrentTopics(), false, now);
+        metadata.add("new-topic", now);
+        assertEquals(0, metadata.timeToNextUpdate(now));
+
+        // Even though add is called, immediate update isn't necessary if the new topic set isn't
+        // containing a new topic,
+        metadata.updateWithCurrentRequestVersion(responseWithCurrentTopics(), false, now);
+        metadata.add("new-topic", now);
+        assertEquals(metadataExpireMs, metadata.timeToNextUpdate(now));
+
+        // If the new set of topics containing a new topic then it should allow immediate update.
+        metadata.add("another-new-topic", now);
+        assertEquals(0, metadata.timeToNextUpdate(now));
+    }
+
+    @Test
+    public void testTopicExpiry() {
+        // Test that topic is expired if not used within the expiry interval
+        long time = 0;
+        final String topic1 = "topic1";
+        metadata.add(topic1, time);
+        metadata.updateWithCurrentRequestVersion(responseWithCurrentTopics(), false, time);
+        assertTrue(metadata.containsTopic(topic1));
+
+        time += METADATA_IDLE_MS;
+        metadata.updateWithCurrentRequestVersion(responseWithCurrentTopics(), false, time);
+        assertFalse("Unused topic not expired", metadata.containsTopic(topic1));
+
+        // Test that topic is not expired if used within the expiry interval
+        final String topic2 = "topic2";
+        metadata.add(topic2, time);
+        metadata.updateWithCurrentRequestVersion(responseWithCurrentTopics(), false, time);
+        for (int i = 0; i < 3; i++) {
+            time += METADATA_IDLE_MS / 2;
+            metadata.updateWithCurrentRequestVersion(responseWithCurrentTopics(), false, time);
+            assertTrue("Topic expired even though in use", metadata.containsTopic(topic2));
+            metadata.add(topic2, time);
+        }
+
+        // Add a new topic, but update its metadata after the expiry would have occurred.
+        // The topic should still be retained.
+        final String topic3 = "topic3";
+        metadata.add(topic3, time);
+        time += METADATA_IDLE_MS * 2;
+        metadata.updateWithCurrentRequestVersion(responseWithCurrentTopics(), false, time);
+        assertTrue("Topic expired while awaiting metadata", metadata.containsTopic(topic3));
+    }
+
+    @Test
+    public void testMetadataWaitAbortedOnFatalException() {
+        metadata.fatalError(new AuthenticationException("Fatal exception from test"));
+        assertThrows(AuthenticationException.class, () -> metadata.awaitUpdate(0, 1000));
+    }
+
+    @Test
+    public void testMetadataPartialUpdate() {
+        long now = 10000;
+
+        // Add a new topic and fetch its metadata in a partial update.
+        final String topic1 = "topic-one";
+        metadata.add(topic1, now);
+        assertTrue(metadata.updateRequested());
+        assertEquals(0, metadata.timeToNextUpdate(now));
+        assertEquals(metadata.topics(), Collections.singleton(topic1));
+        assertEquals(metadata.newTopics(), Collections.singleton(topic1));
+
+        // Perform the partial update. Verify the topic is no longer considered "new".
+        now += 1000;
+        metadata.updateWithCurrentRequestVersion(responseWithTopics(Collections.singleton(topic1)), true, now);
+        assertFalse(metadata.updateRequested());
+        assertEquals(metadata.topics(), Collections.singleton(topic1));
+        assertEquals(metadata.newTopics(), Collections.emptySet());
+
+        // Add the topic again. It should not be considered "new".
+        metadata.add(topic1, now);
+        assertFalse(metadata.updateRequested());
+        assertTrue(metadata.timeToNextUpdate(now) > 0);
+        assertEquals(metadata.topics(), Collections.singleton(topic1));
+        assertEquals(metadata.newTopics(), Collections.emptySet());
+
+        // Add two new topics. However, we'll only apply a partial update for one of them.
+        now += 1000;
+        final String topic2 = "topic-two";
+        metadata.add(topic2, now);
+
+        now += 1000;
+        final String topic3 = "topic-three";
+        metadata.add(topic3, now);
+
+        assertTrue(metadata.updateRequested());
+        assertEquals(0, metadata.timeToNextUpdate(now));
+        assertEquals(metadata.topics(), new HashSet<>(Arrays.asList(topic1, topic2, topic3)));
+        assertEquals(metadata.newTopics(), new HashSet<>(Arrays.asList(topic2, topic3)));
+
+        // Perform the partial update for a subset of the new topics.
+        now += 1000;
+        assertTrue(metadata.updateRequested());
+        metadata.updateWithCurrentRequestVersion(responseWithTopics(Collections.singleton(topic2)), true, now);
+        assertEquals(metadata.topics(), new HashSet<>(Arrays.asList(topic1, topic2, topic3)));
+        assertEquals(metadata.newTopics(), Collections.singleton(topic3));
+    }
+
+    @Test
+    public void testRequestUpdateForTopic() {
+        long now = 10000;
+
+        final String topic1 = "topic-1";
+        final String topic2 = "topic-2";
+
+        // Add the topics to the metadata.
+        metadata.add(topic1, now);
+        metadata.add(topic2, now);
+        assertTrue(metadata.updateRequested());
+
+        // Request an update for topic1. Since the topic is considered new, it should not trigger
+        // the metadata to require a full update.
+        metadata.requestUpdateForTopic(topic1);
+        assertTrue(metadata.updateRequested());
+
+        // Perform the partial update. Verify no additional (full) updates are requested.
+        now += 1000;
+        metadata.updateWithCurrentRequestVersion(responseWithTopics(Collections.singleton(topic1)), true, now);
+        assertFalse(metadata.updateRequested());
+
+        // Request an update for topic1 again. Such a request may occur when the leader
+        // changes, which may affect many topics, and should therefore request a full update.
+        metadata.requestUpdateForTopic(topic1);
+        assertTrue(metadata.updateRequested());
+
+        // Perform a partial update for the topic. This should not clear the full update.
+        now += 1000;
+        metadata.updateWithCurrentRequestVersion(responseWithTopics(Collections.singleton(topic1)), true, now);
+        assertTrue(metadata.updateRequested());
+
+        // Perform the full update. This should clear the update request.
+        now += 1000;
+        metadata.updateWithCurrentRequestVersion(responseWithTopics(new HashSet<>(Arrays.asList(topic1, topic2))), false, now);
+        assertFalse(metadata.updateRequested());
+    }
+
+    private MetadataResponse responseWithCurrentTopics() {
+        return responseWithTopics(metadata.topics());
+    }
+
+    private MetadataResponse responseWithTopics(Set topics) {
+        Map partitionCounts = new HashMap<>();
+        for (String topic : topics)
+            partitionCounts.put(topic, 1);
+        return RequestTestUtils.metadataUpdateWith(1, partitionCounts);
+    }
+
+    private void clearBackgroundError() {
+        backgroundError.set(null);
+    }
+
+    private Thread asyncFetch(final String topic, final long maxWaitMs) {
+        Thread thread = new Thread(() -> {
+            try {
+                while (metadata.fetch().partitionsForTopic(topic).isEmpty())
+                    metadata.awaitUpdate(metadata.requestUpdate(), maxWaitMs);
+            } catch (Exception e) {
+                backgroundError.set(e);
+            }
+        });
+        thread.start();
+        return thread;
+    }
+
+}
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 d486c10a94e16..6b4271cdb0845 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
@@ -19,6 +19,7 @@
 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;
@@ -37,7 +38,6 @@
 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;
@@ -83,10 +83,9 @@ public class RecordAccumulatorTest {
     private MockTime time = new MockTime();
     private byte[] key = "key".getBytes();
     private byte[] value = "value".getBytes();
-    private int msgSize = DefaultRecord.sizeInBytes(0, 0, key.length, value.length,
-            Record.EMPTY_HEADERS);
+    private int msgSize = DefaultRecord.sizeInBytes(0, 0, key.length, value.length, Record.EMPTY_HEADERS);
     private Cluster cluster = new Cluster(null, Arrays.asList(node1, node2), Arrays.asList(part1, part2, part3),
-            Collections.emptySet(), Collections.emptySet());
+        Collections.emptySet(), Collections.emptySet());
     private Metrics metrics = new Metrics(time);
     private final long maxBlockTimeMs = 1000;
     private final LogContext logContext = new LogContext();
@@ -104,11 +103,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, time.milliseconds());
             Deque partitionBatches = accum.batches().get(tp1);
             assertEquals(1, partitionBatches.size());
 
@@ -119,7 +118,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, time.milliseconds());
         Deque partitionBatches = accum.batches().get(tp1);
         assertEquals(2, partitionBatches.size());
         Iterator partitionBatchesIterator = partitionBatches.iterator();
@@ -153,8 +152,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, time.milliseconds());
         assertEquals("Our partition's leader should be ready", Collections.singleton(node1), accum.ready(cluster, time.milliseconds()).readyNodes);
 
         Deque batches = accum.batches().get(tp1);
@@ -188,12 +187,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, time.milliseconds());
         assertEquals("Our partition's leader should be ready", Collections.singleton(node1), accum.ready(cluster, time.milliseconds()).readyNodes);
 
         Deque batches = accum.batches().get(tp1);
@@ -214,10 +212,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, time.milliseconds());
         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);
@@ -235,12 +233,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, time.milliseconds());
         }
         assertEquals("Partition's leader should be ready", Collections.singleton(node1), accum.ready(cluster, time.milliseconds()).readyNodes);
 
@@ -255,14 +253,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, time.milliseconds());
                         } catch (Exception e) {
                             e.printStackTrace();
                         }
@@ -294,19 +292,19 @@ 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;
 
-        RecordAccumulator accum = createTestRecordAccumulator(
-                batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10 * batchSize, CompressionType.NONE, lingerMs);
+        RecordAccumulator accum = createTestRecordAccumulator(batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD,
+                10 * batchSize, CompressionType.NONE, lingerMs);
         // Just short of going over the limit so we trigger linger time
         int appends = expectedNumAppends(batchSize);
 
         // 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, time.milliseconds());
         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);
@@ -315,14 +313,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, time.milliseconds());
         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, time.milliseconds());
         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
@@ -332,13 +330,19 @@ public void testNextReadyCheckDelay() throws Exception {
 
     @Test
     public void testRetryBackoff() throws Exception {
-        long lingerMs = Long.MAX_VALUE / 4;
-        long retryBackoffMs = Long.MAX_VALUE / 2;
-        final RecordAccumulator accum = new RecordAccumulator(logContext, 1024 + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10 * 1024,
-                CompressionType.NONE, lingerMs, retryBackoffMs, metrics, time, new ApiVersions(), null);
+        int lingerMs = Integer.MAX_VALUE / 16;
+        long retryBackoffMs = Integer.MAX_VALUE / 8;
+        int deliveryTimeoutMs = Integer.MAX_VALUE;
+        long totalSize = 10 * 1024;
+        int batchSize = 1024 + DefaultRecordBatch.RECORD_BATCH_OVERHEAD;
+        String metricGrpName = "producer-metrics";
+
+        final RecordAccumulator accum = new RecordAccumulator(logContext, batchSize,
+            CompressionType.NONE, lingerMs, retryBackoffMs, deliveryTimeoutMs, metrics, metricGrpName, time, new ApiVersions(), null,
+            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, time.milliseconds());
         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);
@@ -350,7 +354,7 @@ public void testRetryBackoff() throws Exception {
         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, time.milliseconds());
         result = accum.ready(cluster, now + lingerMs + 1);
         assertEquals("Node1 should be ready", Collections.singleton(node1), result.readyNodes);
 
@@ -371,12 +375,12 @@ public void testRetryBackoff() throws Exception {
 
     @Test
     public void testFlush() throws Exception {
-        long lingerMs = Long.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, time.milliseconds());
             assertTrue(accum.hasIncomplete());
         }
         RecordAccumulator.ReadyCheckResult result = accum.ready(cluster, time.milliseconds());
@@ -413,8 +417,8 @@ public void run() {
     @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, time.milliseconds());
 
         accum.beginFlush();
         assertTrue(accum.flushInProgress());
@@ -429,12 +433,12 @@ public void testAwaitFlushComplete() throws Exception {
 
     @Test
     public void testAbortIncompleteBatches() throws Exception {
-        long lingerMs = Long.MAX_VALUE;
+        int lingerMs = Integer.MAX_VALUE;
         int numRecords = 100;
 
         final AtomicInteger numExceptionReceivedInCallback = new AtomicInteger(0);
         final RecordAccumulator accum = createTestRecordAccumulator(
-                128 + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 64 * 1024, CompressionType.NONE, lingerMs);
+            128 + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 64 * 1024, CompressionType.NONE, lingerMs);
         class TestCallback implements Callback {
             @Override
             public void onCompletion(RecordMetadata metadata, Exception exception) {
@@ -443,7 +447,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, time.milliseconds());
         RecordAccumulator.ReadyCheckResult result = accum.ready(cluster, time.milliseconds());
         assertFalse(result.readyNodes.isEmpty());
         Map> drained = accum.drain(cluster, result.readyNodes, Integer.MAX_VALUE, time.milliseconds());
@@ -468,7 +472,7 @@ public void onCompletion(RecordMetadata metadata, Exception exception) {
 
     @Test
     public void testAbortUnsentBatches() throws Exception {
-        long lingerMs = Long.MAX_VALUE;
+        int lingerMs = Integer.MAX_VALUE;
         int numRecords = 100;
 
         final AtomicInteger numExceptionReceivedInCallback = new AtomicInteger(0);
@@ -484,7 +488,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, time.milliseconds());
         RecordAccumulator.ReadyCheckResult result = accum.ready(cluster, time.milliseconds());
         assertFalse(result.readyNodes.isEmpty());
         Map> drained = accum.drain(cluster, result.readyNodes, Integer.MAX_VALUE,
@@ -509,37 +513,85 @@ public void onCompletion(RecordMetadata metadata, Exception exception) {
         assertTrue(accum.hasIncomplete());
     }
 
+    private void doExpireBatchSingle(int deliveryTimeoutMs) throws InterruptedException {
+        int lingerMs = 300;
+        List muteStates = Arrays.asList(false, true);
+        Set readyNodes = null;
+        List expiredBatches = new ArrayList<>();
+        // test case assumes that the records do not fill the batch completely
+        int batchSize = 1025;
+        RecordAccumulator accum = createTestRecordAccumulator(deliveryTimeoutMs,
+            batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10 * batchSize, CompressionType.NONE, lingerMs);
+
+        // Make the batches ready due to linger. These batches are not in retry
+        for (Boolean mute: muteStates) {
+            if (time.milliseconds() < System.currentTimeMillis())
+                time.setCurrentTimeMs(System.currentTimeMillis());
+            accum.append(tp1, 0L, key, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false, time.milliseconds());
+            assertEquals("No partition should be ready.", 0, accum.ready(cluster, time.milliseconds()).readyNodes.size());
+
+            time.sleep(lingerMs);
+            readyNodes = accum.ready(cluster, time.milliseconds()).readyNodes;
+            assertEquals("Our partition's leader should be ready", Collections.singleton(node1), readyNodes);
+
+            expiredBatches = accum.expiredBatches(time.milliseconds());
+            assertEquals("The batch should not expire when just linger has passed", 0, expiredBatches.size());
+
+            if (mute)
+                accum.mutePartition(tp1);
+            else
+                accum.unmutePartition(tp1);
+
+            // Advance the clock to expire the batch.
+            time.sleep(deliveryTimeoutMs - lingerMs);
+            expiredBatches = accum.expiredBatches(time.milliseconds());
+            assertEquals("The batch may expire when the partition is muted", 1, expiredBatches.size());
+            assertEquals("No partitions should be ready.", 0, accum.ready(cluster, time.milliseconds()).readyNodes.size());
+        }
+    }
+
+    @Test
+    public void testExpiredBatchSingle() throws InterruptedException {
+        doExpireBatchSingle(3200);
+    }
+
+    @Test
+    public void testExpiredBatchSingleMaxValue() throws InterruptedException {
+        doExpireBatchSingle(Integer.MAX_VALUE);
+    }
+
     @Test
     public void testExpiredBatches() throws InterruptedException {
         long retryBackoffMs = 100L;
-        long lingerMs = 3000L;
+        int lingerMs = 30;
         int requestTimeout = 60;
+        int deliveryTimeoutMs = 3200;
 
         // test case assumes that the records do not fill the batch completely
         int batchSize = 1025;
 
         RecordAccumulator accum = createTestRecordAccumulator(
-                batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10 * batchSize, CompressionType.NONE, lingerMs);
+            deliveryTimeoutMs, batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10 * batchSize, CompressionType.NONE, lingerMs);
         int appends = expectedNumAppends(batchSize);
 
         // 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, time.milliseconds());
             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, time.milliseconds());
         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.
-        time.sleep(requestTimeout + 1);
+        time.sleep(deliveryTimeoutMs + 1);
         accum.mutePartition(tp1);
-        List expiredBatches = accum.expiredBatches(requestTimeout, time.milliseconds());
-        assertEquals("The batch should not be expired when the partition is muted", 0, expiredBatches.size());
+        List expiredBatches = accum.expiredBatches(time.milliseconds());
+        assertEquals("The batches will be muted no matter if the partition is muted or not", 2, expiredBatches.size());
 
         accum.unmutePartition(tp1);
-        expiredBatches = accum.expiredBatches(requestTimeout, time.milliseconds());
-        assertEquals("The batch should be expired", 1, expiredBatches.size());
+        expiredBatches = accum.expiredBatches(time.milliseconds());
+        assertEquals("All batches should have been expired earlier", 0, expiredBatches.size());
         assertEquals("No partitions should be ready.", 0, accum.ready(cluster, time.milliseconds()).readyNodes.size());
 
         // Advance the clock to make the next batch ready due to linger.ms
@@ -548,17 +600,17 @@ public void testExpiredBatches() throws InterruptedException {
         time.sleep(requestTimeout + 1);
 
         accum.mutePartition(tp1);
-        expiredBatches = accum.expiredBatches(requestTimeout, time.milliseconds());
+        expiredBatches = accum.expiredBatches(time.milliseconds());
         assertEquals("The batch should not be expired when metadata is still available and partition is muted", 0, expiredBatches.size());
 
         accum.unmutePartition(tp1);
-        expiredBatches = accum.expiredBatches(requestTimeout, time.milliseconds());
-        assertEquals("The batch should be expired when the partition is not muted", 1, expiredBatches.size());
+        expiredBatches = accum.expiredBatches(time.milliseconds());
+        assertEquals("All batches should have been expired", 0, expiredBatches.size());
         assertEquals("No partitions should be ready.", 0, accum.ready(cluster, time.milliseconds()).readyNodes.size());
 
         // 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.milliseconds());
         time.sleep(lingerMs);
         readyNodes = accum.ready(cluster, time.milliseconds()).readyNodes;
         assertEquals("Our partition's leader should be ready", Collections.singleton(node1), readyNodes);
@@ -569,17 +621,40 @@ public void testExpiredBatches() throws InterruptedException {
 
         // test expiration.
         time.sleep(requestTimeout + retryBackoffMs);
-        expiredBatches = accum.expiredBatches(requestTimeout, time.milliseconds());
+        expiredBatches = accum.expiredBatches(time.milliseconds());
         assertEquals("The batch should not be expired.", 0, expiredBatches.size());
         time.sleep(1L);
 
         accum.mutePartition(tp1);
-        expiredBatches = accum.expiredBatches(requestTimeout, time.milliseconds());
+        expiredBatches = accum.expiredBatches(time.milliseconds());
+        assertEquals("The batch should not be expired when the partition is muted", 0, expiredBatches.size());
+
+        accum.unmutePartition(tp1);
+        expiredBatches = accum.expiredBatches(time.milliseconds());
+        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, false, time.milliseconds());
+        time.sleep(lingerMs);
+        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.
+        time.sleep(requestTimeout + 1);
+        accum.mutePartition(tp1);
+        expiredBatches = accum.expiredBatches(time.milliseconds());
         assertEquals("The batch should not be expired when the partition is muted", 0, expiredBatches.size());
 
+        long throttleTimeMs = 100L;
         accum.unmutePartition(tp1);
-        expiredBatches = accum.expiredBatches(requestTimeout, time.milliseconds());
-        assertEquals("The batch should be expired when the partition is not muted.", 1, expiredBatches.size());
+        // The batch shouldn't be expired yet.
+        expiredBatches = accum.expiredBatches(time.milliseconds());
+        assertEquals("The batch should not be expired when the partition is muted", 0, expiredBatches.size());
+
+        // Once the throttle time is over, the batch can be expired.
+        time.sleep(throttleTimeMs);
+        expiredBatches = accum.expiredBatches(time.milliseconds());
+        assertEquals("All batches should have been expired earlier", 0, expiredBatches.size());
+        assertEquals("No partitions should be ready.", 1, accum.ready(cluster, time.milliseconds()).readyNodes.size());
     }
 
     @Test
@@ -592,7 +667,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, time.milliseconds());
             assertEquals("No partitions should be ready.", 0, accum.ready(cluster, now).readyNodes.size());
         }
         time.sleep(2000);
@@ -623,11 +698,18 @@ 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;
-        apiVersions.update("foobar", NodeApiVersions.create(Arrays.asList(new ApiVersionsResponse.ApiVersion(ApiKeys.PRODUCE.id,
-                (short) 0, (short) 2))));
-        RecordAccumulator accum = new RecordAccumulator(logContext, batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10 * batchSize,
-                CompressionType.NONE, 10, 100L, metrics, time, apiVersions, new TransactionManager());
-        accum.append(tp1, 0L, key, value, Record.EMPTY_HEADERS, null, 0);
+        int deliveryTimeoutMs = 3200;
+        int lingerMs = 10;
+        long retryBackoffMs = 100L;
+        long totalSize = 10 * batchSize;
+        String metricGrpName = "producer-metrics";
+
+        apiVersions.update("foobar", NodeApiVersions.create(ApiKeys.PRODUCE.id, (short) 0, (short) 2));
+        TransactionManager transactionManager = new TransactionManager(new LogContext(), null, 0, 100L, new ApiVersions(), false);
+        RecordAccumulator accum = new RecordAccumulator(logContext, batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD,
+            CompressionType.NONE, lingerMs, retryBackoffMs, deliveryTimeoutMs, metrics, metricGrpName, time, apiVersions, transactionManager,
+            new BufferPool(totalSize, batchSize, metrics, time, metricGrpName));
+        accum.append(tp1, 0L, key, value, Record.EMPTY_HEADERS, null, 0, false, time.milliseconds());
     }
 
     @Test
@@ -692,7 +774,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.
@@ -704,9 +786,9 @@ public void testSplitBatchOffAccumulator() throws InterruptedException {
             assertFalse(drained.get(node1.id()).isEmpty());
         }
         assertTrue("All the batches should have been drained.",
-                   accum.ready(cluster, time.milliseconds()).readyNodes.isEmpty());
+                accum.ready(cluster, time.milliseconds()).readyNodes.isEmpty());
         assertEquals("The split batches should be allocated off the accumulator",
-                     bufferCapacity, accum.bufferPoolAvailableMemory());
+                bufferCapacity, accum.bufferPoolAvailableMemory());
     }
 
     @Test
@@ -727,7 +809,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, time.milliseconds());
                 BatchDrainedResult result = completeOrSplitBatches(accum, batchSize);
                 numSplit += result.numSplit;
                 numBatches += result.numBatches;
@@ -737,11 +819,155 @@ public void testSplitFrequency() throws InterruptedException {
             numSplit += result.numSplit;
             numBatches += result.numBatches;
             assertTrue(String.format("Total num batches = %d, split batches = %d, more than 10%% of the batch splits. "
-                                         + "Random seed is " + seed,
-                                     numBatches, numSplit), (double) numSplit / numBatches < 0.1f);
+                    + "Random seed is " + seed,
+                numBatches, numSplit), (double) numSplit / numBatches < 0.1f);
         }
     }
 
+    @Test
+    public void testSoonToExpireBatchesArePickedUpForExpiry() throws InterruptedException {
+        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, false, time.milliseconds());
+        Set readyNodes = accum.ready(cluster, time.milliseconds()).readyNodes;
+        Map> drained = accum.drain(cluster, readyNodes, Integer.MAX_VALUE, time.milliseconds());
+        assertTrue(drained.isEmpty());
+        //assertTrue(accum.soonToExpireInFlightBatches().isEmpty());
+
+        // advanced clock and send one batch out but it should not be included in soon to expire inflight
+        // batches because batch's expiry is quite far.
+        time.sleep(lingerMs + 1);
+        readyNodes = accum.ready(cluster, time.milliseconds()).readyNodes;
+        drained = accum.drain(cluster, readyNodes, Integer.MAX_VALUE, time.milliseconds());
+        assertEquals("A batch did not drain after linger", 1, drained.size());
+        //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, false, time.milliseconds());
+        time.sleep(lingerMs * 4);
+
+        // Now drain and check that accumulator picked up the drained batch because its expiry is soon.
+        readyNodes = accum.ready(cluster, time.milliseconds()).readyNodes;
+        drained = accum.drain(cluster, readyNodes, Integer.MAX_VALUE, time.milliseconds());
+        assertEquals("A batch did not drain after linger", 1, drained.size());
+    }
+
+    @Test
+    public void testExpiredBatchesRetry() throws InterruptedException {
+        int lingerMs = 3000;
+        int rtt = 1000;
+        int deliveryTimeoutMs = 3200;
+        Set readyNodes;
+        List expiredBatches;
+        List muteStates = Arrays.asList(false, true);
+
+        // test case assumes that the records do not fill the batch completely
+        int batchSize = 1025;
+        RecordAccumulator accum = createTestRecordAccumulator(
+            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, false, time.milliseconds());
+            time.sleep(lingerMs);
+            readyNodes = accum.ready(cluster, time.milliseconds()).readyNodes;
+            assertEquals("Our partition's leader should be ready", Collections.singleton(node1), readyNodes);
+            Map> drained = accum.drain(cluster, readyNodes, Integer.MAX_VALUE, time.milliseconds());
+            assertEquals("There should be only one batch.", 1, drained.get(node1.id()).size());
+            time.sleep(rtt);
+            accum.reenqueue(drained.get(node1.id()).get(0), time.milliseconds());
+
+            if (mute)
+                accum.mutePartition(tp1);
+            else
+                accum.unmutePartition(tp1);
+
+            // test expiration
+            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());
+        }
+    }
+
+    @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, time.milliseconds());
+        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, time.milliseconds());
+            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, time.milliseconds());
+        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, time.milliseconds());
+            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)
         throws InterruptedException {
         Random random = new Random();
@@ -751,7 +977,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, time.milliseconds());
         }
 
         RecordAccumulator.ReadyCheckResult result = accum.ready(cluster, time.milliseconds());
@@ -829,7 +1055,23 @@ private int expectedNumAppends(int batchSize) {
         int offsetDelta = 0;
         while (true) {
             int recordSize = DefaultRecord.sizeInBytes(offsetDelta, 0, key.length, value.length,
-                    Record.EMPTY_HEADERS);
+                Record.EMPTY_HEADERS);
+            if (size + recordSize > batchSize)
+                return offsetDelta;
+            offsetDelta += 1;
+            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;
@@ -837,20 +1079,31 @@ private int expectedNumAppends(int batchSize) {
         }
     }
 
+
+    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(int batchSize, long totalSize, CompressionType type, long lingerMs) {
+    private RecordAccumulator createTestRecordAccumulator(int deliveryTimeoutMs, int batchSize, long totalSize, CompressionType type, int lingerMs) {
+        long retryBackoffMs = 100L;
+        String metricGrpName = "producer-metrics";
+
         return new RecordAccumulator(
-                logContext,
-                batchSize,
-                totalSize,
-                type,
-                lingerMs,
-                100L,
-                metrics,
-                time,
-                new ApiVersions(),
-                null);
+            logContext,
+            batchSize,
+            type,
+            lingerMs,
+            retryBackoffMs,
+            deliveryTimeoutMs,
+            metrics,
+            metricGrpName,
+            time,
+            new ApiVersions(),
+            null,
+            new BufferPool(totalSize, batchSize, metrics, time, metricGrpName));
     }
 }
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 1ce8e5a4bbdb4..a054f6f1811d7 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
@@ -17,27 +17,38 @@
 package org.apache.kafka.clients.producer.internals;
 
 import org.apache.kafka.clients.ApiVersions;
+import org.apache.kafka.clients.ClientDnsLookup;
 import org.apache.kafka.clients.ClientRequest;
-import org.apache.kafka.clients.Metadata;
+import org.apache.kafka.clients.ClientResponse;
 import org.apache.kafka.clients.MockClient;
 import org.apache.kafka.clients.NetworkClient;
 import org.apache.kafka.clients.NodeApiVersions;
 import org.apache.kafka.clients.producer.Callback;
+import org.apache.kafka.common.errors.InvalidRequestException;
+import org.apache.kafka.common.errors.TransactionAbortedException;
+import org.apache.kafka.common.message.ProduceRequestData;
+import org.apache.kafka.common.requests.FindCoordinatorRequest.CoordinatorType;
+import org.apache.kafka.common.requests.MetadataRequest;
+import org.apache.kafka.common.requests.RequestTestUtils;
+import org.apache.kafka.common.requests.RequestUtils;
+import org.apache.kafka.common.utils.ProducerIdAndEpoch;
 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;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.errors.ClusterAuthorizationException;
 import org.apache.kafka.common.errors.NetworkException;
-import org.apache.kafka.common.errors.OutOfOrderSequenceException;
 import org.apache.kafka.common.errors.RecordTooLargeException;
 import org.apache.kafka.common.errors.TimeoutException;
 import org.apache.kafka.common.errors.TopicAuthorizationException;
 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.EndTxnResponseData;
+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,59 +65,84 @@
 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;
 import org.apache.kafka.test.DelayedReceive;
 import org.apache.kafka.test.MockSelector;
 import org.apache.kafka.test.TestUtils;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
+import org.mockito.InOrder;
 
 import java.nio.ByteBuffer;
 import java.util.Collections;
 import java.util.Deque;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.IdentityHashMap;
 import java.util.Iterator;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.OptionalInt;
+import java.util.OptionalLong;
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 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.assertNotNull;
 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;
+import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.eq;
+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 static final long TOPIC_IDLE_MS = 60 * 1000;
 
     private TopicPartition tp0 = new TopicPartition("test", 0);
     private TopicPartition tp1 = new TopicPartition("test", 1);
     private MockTime time = new MockTime();
-    private MockClient client = new MockClient(time);
     private int batchSize = 16 * 1024;
-    private Metadata metadata = new Metadata(0, Long.MAX_VALUE, true, true, new ClusterResourceListeners());
+    private ProducerMetadata metadata = new ProducerMetadata(0, Long.MAX_VALUE, TOPIC_IDLE_MS,
+            new LogContext(), new ClusterResourceListeners(), time);
+    private MockClient client = new MockClient(time, metadata);
     private ApiVersions apiVersions = new ApiVersions();
-    private Cluster cluster = TestUtils.singletonCluster("test", 2);
     private Metrics metrics = null;
     private RecordAccumulator accumulator = null;
     private Sender sender = null;
@@ -115,7 +151,6 @@ public class SenderTest {
 
     @Before
     public void setup() {
-        client.setNode(cluster.nodes().get(0));
         setupWithTransactionState(null);
     }
 
@@ -124,19 +159,30 @@ public void tearDown() {
         this.metrics.close();
     }
 
+    private static Map partitionRecords(ProduceRequest request) {
+        Map partitionRecords = new HashMap<>();
+        request.data().topicData().forEach(tpData -> tpData.partitionData().forEach(p -> {
+            TopicPartition tp = new TopicPartition(tpData.name(), p.index());
+            partitionRecords.put(tp, (MemoryRecords) p.records());
+        }));
+        return Collections.unmodifiableMap(partitionRecords);
+    }
+ 
     @Test
     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
+        Future future = appendToAccumulator(tp0, 0L, "key", "value");
+        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());
     }
@@ -151,34 +197,30 @@ public void testMessageFormatDownConversion() throws Exception {
         // start off support produce request v3
         apiVersions.update("0", NodeApiVersions.create());
 
-        Future future = accumulator.append(tp0, 0L, "key".getBytes(), "value".getBytes(),
-                null, null, MAX_BLOCK_TIMEOUT).future;
+        Future future = appendToAccumulator(tp0, 0L, "key", "value");
 
         // 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
-            public boolean matches(AbstractRequest body) {
-                ProduceRequest request = (ProduceRequest) body;
-                if (request.version() != 2)
-                    return false;
+        client.prepareResponse(body -> {
+            ProduceRequest request = (ProduceRequest) body;
+            if (request.version() != 2)
+                return false;
 
-                MemoryRecords records = request.partitionRecordsOrFail().get(tp0);
-                return records != null &&
-                        records.sizeInBytes() > 0 &&
-                        records.hasMatchingMagic(RecordBatch.MAGIC_VALUE_V1);
-            }
+            MemoryRecords records = partitionRecords(request).get(tp0);
+            return records != null &&
+                    records.sizeInBytes() > 0 &&
+                    records.hasMatchingMagic(RecordBatch.MAGIC_VALUE_V1);
         }, 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());
     }
 
+    @SuppressWarnings("deprecation")
     @Test
     public void testDownConversionForMismatchedMagicValues() throws Exception {
         // it can happen that we construct a record set with mismatching magic values (perhaps
@@ -191,15 +233,12 @@ public void testDownConversionForMismatchedMagicValues() throws Exception {
         // start off support produce request v3
         apiVersions.update("0", NodeApiVersions.create());
 
-        Future future1 = accumulator.append(tp0, 0L, "key".getBytes(), "value".getBytes(),
-                null, null, MAX_BLOCK_TIMEOUT).future;
+        Future future1 = appendToAccumulator(tp0, 0L, "key", "value");
 
         // 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;
+        Future future2 = appendToAccumulator(tp1, 0L, "key", "value");
 
         // start off support produce request v3
         apiVersions.update("0", NodeApiVersions.create());
@@ -210,27 +249,24 @@ public void testDownConversionForMismatchedMagicValues() throws Exception {
         partResp.put(tp1, resp);
         ProduceResponse produceResponse = new ProduceResponse(partResp, 0);
 
-        client.prepareResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                ProduceRequest request = (ProduceRequest) body;
-                if (request.version() != 2)
-                    return false;
+        client.prepareResponse(body -> {
+            ProduceRequest request = (ProduceRequest) body;
+            if (request.version() != 2)
+                return false;
 
-                Map recordsMap = request.partitionRecordsOrFail();
-                if (recordsMap.size() != 2)
-                    return false;
+            Map recordsMap = partitionRecords(request);
+            if (recordsMap.size() != 2)
+                return false;
 
-                for (MemoryRecords records : recordsMap.values()) {
-                    if (records == null || records.sizeInBytes() == 0 || !records.hasMatchingMagic(RecordBatch.MAGIC_VALUE_V1))
-                        return false;
-                }
-                return true;
+            for (MemoryRecords records : recordsMap.values()) {
+                if (records == null || records.sizeInBytes() == 0 || !records.hasMatchingMagic(RecordBatch.MAGIC_VALUE_V1))
+                    return false;
             }
+            return true;
         }, 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());
@@ -239,43 +275,52 @@ 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 {
+    @Test
+    public void testQuotaMetrics() {
         MockSelector selector = new MockSelector(time);
         Sensor throttleTimeSensor = Sender.throttleTimeSensor(this.senderMetricsRegistry);
         Cluster cluster = TestUtils.singletonCluster("test", 1);
         Node node = cluster.nodes().get(0);
         NetworkClient client = new NetworkClient(selector, metadata, "mock", Integer.MAX_VALUE,
-                1000, 1000, 64 * 1024, 64 * 1024, 1000,
+                1000, 1000, 64 * 1024, 64 * 1024, 1000, 10 * 1000, 127 * 1000, ClientDnsLookup.USE_ALL_DNS_IPS,
                 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 = RequestTestUtils.serializeResponseWithHeader(
+            ApiVersionsResponse.createApiVersionsResponse(400, RecordBatch.CURRENT_MAGIC_VALUE),
+            ApiKeys.API_VERSIONS.latestVersion(), 0);
+
         selector.delayedReceive(new DelayedReceive(node.idString(), new NetworkReceive(node.idString(), buffer)));
-        while (!client.ready(node, time.milliseconds()))
+        while (!client.ready(node, time.milliseconds())) {
             client.poll(1, time.milliseconds());
+            // If a throttled response is received, advance the time to ensure progress.
+            time.sleep(client.throttleDelayMs(node, time.milliseconds()));
+        }
         selector.clear();
 
         for (int i = 1; i <= 3; i++) {
             int throttleTimeMs = 100 * i;
-            ProduceRequest.Builder builder = ProduceRequest.Builder.forCurrentMagic((short) 1, 1000,
-                            Collections.emptyMap());
-            ClientRequest request = client.newClientRequest(node.idString(), builder, time.milliseconds(), true, null);
+            ProduceRequest.Builder builder = ProduceRequest.forCurrentMagic(new ProduceRequestData()
+                    .setTopicData(new ProduceRequestData.TopicProduceDataCollection())
+                    .setAcks((short) 1)
+                    .setTimeoutMs(1000));
+            ClientRequest request = client.newClientRequest(node.idString(), builder, time.milliseconds(), true);
             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 = RequestTestUtils.serializeResponseWithHeader(response, 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.
+            time.sleep(client.throttleDelayMs(node, time.milliseconds()));
             selector.clear();
         }
         Map allMetrics = metrics.metrics();
         KafkaMetric avgMetric = allMetrics.get(this.senderMetricsRegistry.produceThrottleTimeAvg);
         KafkaMetric maxMetric = allMetrics.get(this.senderMetricsRegistry.produceThrottleTimeMax);
         // Throttle times are ApiVersions=400, Produce=(100, 200, 300)
-        assertEquals(250, avgMetric.value(), EPS);
-        assertEquals(400, maxMetric.value(), EPS);
+        assertEquals(250, (Double) avgMetric.metricValue(), EPS);
+        assertEquals(400, (Double) maxMetric.metricValue(), EPS);
         client.close();
     }
 
@@ -286,14 +331,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
+        appendToAccumulator(tp0, 0L, "key", "value");
+        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);
 
@@ -314,42 +359,51 @@ 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 = appendToAccumulator(tp0, 0L, "key", "value");
+            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());
-            assertTrue("Client ready status should be true", client.isReady(node, 0L));
+            assertEquals(1, sender.inFlightBatches(tp0).size());
+            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));
-            sender.run(time.milliseconds()); // receive error
-            sender.run(time.milliseconds()); // reconnect
-            sender.run(time.milliseconds()); // resend
+            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.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 = appendToAccumulator(tp0, 0L, "key", "value");
+            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.run(time.milliseconds()); // reconnect
-                sender.run(time.milliseconds()); // resend
+                sender.runOnce(); // receive error
+                assertEquals(0, sender.inFlightBatches(tp0).size());
+                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 {
             m.close();
         }
@@ -363,34 +417,37 @@ 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
-            Cluster cluster1 = TestUtils.clusterWith(2, "test", 2);
-            metadata.update(cluster1, Collections.emptySet(), time.milliseconds());
+            MetadataResponse metadataUpdate1 = RequestTestUtils.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
+            appendToAccumulator(tp2, 0L, "key1", "value1");
+            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);
+            appendToAccumulator(tp2, 0L, "key2", "value2");
 
             // Update metadata before sender receives response from broker 0. Now partition 2 moves to broker 0
-            Cluster cluster2 = TestUtils.singletonCluster("test", 2);
-            metadata.update(cluster2, Collections.emptySet(), time.milliseconds());
+            MetadataResponse metadataUpdate2 = RequestTestUtils.metadataUpdateWith(1, Collections.singletonMap("test", 2));
+            client.prepareMetadataUpdate(metadataUpdate2);
             // Sender should not send the second message to node 0.
-            sender.run(time.milliseconds());
+            assertEquals(1, sender.inFlightBatches(tp2).size());
+            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());
         } finally {
             m.close();
         }
@@ -404,34 +461,37 @@ public void testAppendInExpiryCallback() throws InterruptedException {
         final byte[] key = "key".getBytes();
         final byte[] value = "value".getBytes();
         final long maxBlockTimeMs = 1000;
-        Callback callback = new Callback() {
-            @Override
-            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);
-                    } catch (InterruptedException e) {
-                        throw new RuntimeException("Unexpected interruption", e);
-                    }
-                } else if (exception != null)
-                    unexpectedException.compareAndSet(null, exception);
-            }
+        Callback callback = (metadata, exception) -> {
+            if (exception instanceof TimeoutException) {
+                expiryCallbackCount.incrementAndGet();
+                try {
+                    accumulator.append(tp1, 0L, key, value,
+                        Record.EMPTY_HEADERS, null, maxBlockTimeMs, false, time.milliseconds());
+                } catch (InterruptedException e) {
+                    throw new RuntimeException("Unexpected interruption", e);
+                }
+            } else if (exception != null)
+                unexpectedException.compareAndSet(null, exception);
         };
 
+        final long nowMs = time.milliseconds();
         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, nowMs);
 
         // Advance the clock to expire the first batch.
         time.sleep(10000);
+
+        Node clusterNode = metadata.fetch().nodes().get(0);
+        Map> drainedBatches =
+            accumulator.drain(metadata.fetch(), Collections.singleton(clusterNode), Integer.MAX_VALUE, time.milliseconds());
+        sender.addToInflightBatches(drainedBatches);
+
         // Disconnect the target node for the pending produce request. This will ensure that sender will try to
         // expire the batch.
-        Node clusterNode = this.cluster.nodes().get(0);
         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.
+        client.backoff(clusterNode, 100);
 
+        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.
@@ -447,55 +507,139 @@ public void onCompletion(RecordMetadata metadata, Exception exception) {
     @Test
     public void testMetadataTopicExpiry() throws Exception {
         long offset = 0;
-        metadata.update(Cluster.empty(), Collections.emptySet(), time.milliseconds());
+        client.updateMetadata(RequestTestUtils.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 = appendToAccumulator(tp0);
+        sender.runOnce();
         assertTrue("Topic not added to metadata", metadata.containsTopic(tp0.topic()));
-        metadata.update(cluster, Collections.emptySet(), time.milliseconds());
-        sender.run(time.milliseconds());  // send produce request
-        client.respond(produceResponse(tp0, offset++, Errors.NONE, 0));
-        sender.run(time.milliseconds());
+        client.updateMetadata(RequestTestUtils.metadataUpdateWith(1, Collections.singletonMap("test", 2)));
+        sender.runOnce();  // send produce request
+        client.respond(produceResponse(tp0, offset, Errors.NONE, 0));
+        sender.runOnce();
         assertEquals("Request completed.", 0, client.inFlightRequestCount());
         assertFalse(client.hasInFlightRequests());
-        sender.run(time.milliseconds());
+        assertEquals(0, sender.inFlightBatches(tp0).size());
+        sender.runOnce();
         assertTrue("Request should be completed", future.isDone());
 
         assertTrue("Topic not retained in metadata list", metadata.containsTopic(tp0.topic()));
-        time.sleep(Metadata.TOPIC_EXPIRY_MS);
-        metadata.update(Cluster.empty(), Collections.emptySet(), time.milliseconds());
+        time.sleep(TOPIC_IDLE_MS);
+        client.updateMetadata(RequestTestUtils.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 = appendToAccumulator(tp0);
+        sender.runOnce();
         assertTrue("Topic not added to metadata", metadata.containsTopic(tp0.topic()));
-        metadata.update(cluster, Collections.emptySet(), time.milliseconds());
-        sender.run(time.milliseconds());  // send produce request
-        client.respond(produceResponse(tp0, offset++, Errors.NONE, 0));
-        sender.run(time.milliseconds());
+        client.updateMetadata(RequestTestUtils.metadataUpdateWith(1, Collections.singletonMap("test", 2)));
+        sender.runOnce();  // send produce request
+        client.respond(produceResponse(tp0, offset + 1, Errors.NONE, 0));
+        sender.runOnce();
         assertEquals("Request completed.", 0, client.inFlightRequestCount());
         assertFalse(client.hasInFlightRequests());
-        sender.run(time.milliseconds());
+        assertEquals(0, sender.inFlightBatches(tp0).size());
+        sender.runOnce();
         assertTrue("Request should be completed", future.isDone());
     }
 
     @Test
-    public void testInitProducerIdRequest() throws Exception {
+    public void testInitProducerIdRequest() {
         final long producerId = 343434L;
-        TransactionManager transactionManager = new TransactionManager();
+        TransactionManager transactionManager = createTransactionManager();
         setupWithTransactionState(transactionManager);
-        client.setNode(new Node(1, "localhost", 33343));
         prepareAndReceiveInitProducerId(producerId, Errors.NONE);
         assertTrue(transactionManager.hasProducerId());
         assertEquals(producerId, transactionManager.producerIdAndEpoch().producerId);
         assertEquals((short) 0, transactionManager.producerIdAndEpoch().epoch);
     }
 
+    /**
+     * Verifies that InitProducerId of transactional producer succeeds even if metadata requests
+     * are pending with only one bootstrap node available and maxInFlight=1, where multiple
+     * polls are necessary to send requests.
+     */
+    @Test
+    public void testInitProducerIdWithMaxInFlightOne() throws Exception {
+        final long producerId = 123456L;
+        createMockClientWithMaxFlightOneMetadataPending();
+
+        // Initialize transaction manager. InitProducerId will be queued up until metadata response
+        // is processed and FindCoordinator can be sent to `leastLoadedNode`.
+        TransactionManager transactionManager = new TransactionManager(new LogContext(), "testInitProducerIdWithPendingMetadataRequest",
+                60000, 100L, new ApiVersions(), false);
+        setupWithTransactionState(transactionManager, false, null, false);
+        ProducerIdAndEpoch producerIdAndEpoch = new ProducerIdAndEpoch(producerId, (short) 0);
+        transactionManager.initializeTransactions();
+        sender.runOnce();
+
+        // Process metadata response, prepare FindCoordinator and InitProducerId responses.
+        // Verify producerId after the sender is run to process responses.
+        MetadataResponse metadataUpdate = RequestTestUtils.metadataUpdateWith(1, Collections.emptyMap());
+        client.respond(metadataUpdate);
+        prepareFindCoordinatorResponse(Errors.NONE);
+        prepareInitProducerResponse(Errors.NONE, producerIdAndEpoch.producerId, producerIdAndEpoch.epoch);
+        waitForProducerId(transactionManager, producerIdAndEpoch);
+    }
+
+    /**
+     * Verifies that InitProducerId of idempotent producer succeeds even if metadata requests
+     * are pending with only one bootstrap node available and maxInFlight=1, where multiple
+     * polls are necessary to send requests.
+     */
+    @Test
+    public void testIdempotentInitProducerIdWithMaxInFlightOne() throws Exception {
+        final long producerId = 123456L;
+        createMockClientWithMaxFlightOneMetadataPending();
+
+        // Initialize transaction manager. InitProducerId will be queued up until metadata response
+        // is processed.
+        TransactionManager transactionManager = createTransactionManager();
+        setupWithTransactionState(transactionManager, false, null, false);
+        ProducerIdAndEpoch producerIdAndEpoch = new ProducerIdAndEpoch(producerId, (short) 0);
+
+        // Process metadata and InitProducerId responses.
+        // Verify producerId after the sender is run to process responses.
+        MetadataResponse metadataUpdate = RequestTestUtils.metadataUpdateWith(1, Collections.emptyMap());
+        client.respond(metadataUpdate);
+        sender.runOnce();
+        sender.runOnce();
+        client.respond(initProducerIdResponse(producerIdAndEpoch.producerId, producerIdAndEpoch.epoch, Errors.NONE));
+        waitForProducerId(transactionManager, producerIdAndEpoch);
+    }
+
+    /**
+     * Tests the code path where the target node to send FindCoordinator or InitProducerId
+     * is not ready.
+     */
+    @Test
+    public void testNodeNotReady() throws Exception {
+        final long producerId = 123456L;
+        time = new MockTime(10);
+        client = new MockClient(time, metadata);
+
+        TransactionManager transactionManager = new TransactionManager(new LogContext(), "testNodeNotReady",
+                60000, 100L, new ApiVersions(), false);
+        setupWithTransactionState(transactionManager, false, null, true);
+        ProducerIdAndEpoch producerIdAndEpoch = new ProducerIdAndEpoch(producerId, (short) 0);
+        transactionManager.initializeTransactions();
+        sender.runOnce();
+
+        Node node = metadata.fetch().nodes().get(0);
+        client.delayReady(node, REQUEST_TIMEOUT + 20);
+        prepareFindCoordinatorResponse(Errors.NONE);
+        sender.runOnce();
+        sender.runOnce();
+        assertNotNull("Coordinator not found", transactionManager.coordinator(CoordinatorType.TRANSACTION));
+
+        client.throttle(node, REQUEST_TIMEOUT + 20);
+        prepareFindCoordinatorResponse(Errors.NONE);
+        prepareInitProducerResponse(Errors.NONE, producerIdAndEpoch.producerId, producerIdAndEpoch.epoch);
+        waitForProducerId(transactionManager, producerIdAndEpoch);
+    }
+
     @Test
     public void testClusterAuthorizationExceptionInInitProducerIdRequest() throws Exception {
         final long producerId = 343434L;
-        TransactionManager transactionManager = new TransactionManager();
+        TransactionManager transactionManager = createTransactionManager();
         setupWithTransactionState(transactionManager);
-        client.setNode(new Node(1, "localhost", 33343));
         prepareAndReceiveInitProducerId(producerId, Errors.CLUSTER_AUTHORIZATION_FAILED);
         assertFalse(transactionManager.hasProducerId());
         assertTrue(transactionManager.hasError());
@@ -508,25 +652,23 @@ 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 = appendToAccumulator(tp0, 0L, "key", "value");
+        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());
-        assertTrue("Client ready status should be true", client.isReady(node, 0L));
+        assertEquals(1, sender.inFlightBatches(tp0).size());
+        assertTrue("Client ready status should be true", client.isReady(node, time.milliseconds()));
         assertFalse(future.isDone());
 
-        client.respond(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                ProduceRequest request = (ProduceRequest) body;
-                assertFalse(request.isIdempotent());
-                return true;
-            }
+        client.respond(body -> {
+            ProduceRequest request = (ProduceRequest) body;
+            assertFalse(RequestUtils.hasIdempotentRecords(request));
+            return true;
         }, produceResponse(tp0, -1L, Errors.TOPIC_AUTHORIZATION_FAILED, 0));
-        sender.run(time.milliseconds());
+        sender.runOnce();
         assertTrue(future.isDone());
         try {
             future.get();
@@ -538,7 +680,7 @@ public boolean matches(AbstractRequest body) {
     @Test
     public void testIdempotenceWithMultipleInflights() throws Exception {
         final long producerId = 343434L;
-        TransactionManager transactionManager = new TransactionManager();
+        TransactionManager transactionManager = createTransactionManager();
         setupWithTransactionState(transactionManager);
         prepareAndReceiveInitProducerId(producerId, Errors.NONE);
         assertTrue(transactionManager.hasProducerId());
@@ -546,38 +688,39 @@ 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 = appendToAccumulator(tp0);
+        sender.runOnce();
         String nodeId = client.requests().peek().destination();
         Node node = new Node(Integer.valueOf(nodeId), "localhost", 0);
         assertEquals(1, client.inFlightRequestCount());
         assertEquals(1, transactionManager.sequenceNumber(tp0).longValue());
-        assertEquals(-1, transactionManager.lastAckedSequence(tp0));
+        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 = appendToAccumulator(tp0);
+        sender.runOnce();
         assertEquals(2, client.inFlightRequestCount());
         assertEquals(2, transactionManager.sequenceNumber(tp0).longValue());
-        assertEquals(-1, transactionManager.lastAckedSequence(tp0));
+        assertEquals(OptionalInt.empty(), transactionManager.lastAckedSequence(tp0));
         assertFalse(request1.isDone());
         assertFalse(request2.isDone());
         assertTrue(client.isReady(node, time.milliseconds()));
 
         sendIdempotentProducerResponse(0, tp0, Errors.NONE, 0L);
 
-        sender.run(time.milliseconds()); // receive response 0
+        sender.runOnce(); // receive response 0
 
         assertEquals(1, client.inFlightRequestCount());
-        assertEquals(0, transactionManager.lastAckedSequence(tp0));
+        assertEquals(OptionalInt.of(0), transactionManager.lastAckedSequence(tp0));
         assertTrue(request1.isDone());
         assertEquals(0, request1.get().offset());
         assertFalse(request2.isDone());
 
         sendIdempotentProducerResponse(1, tp0, Errors.NONE, 1L);
-        sender.run(time.milliseconds()); // receive response 1
-        assertEquals(1, transactionManager.lastAckedSequence(tp0));
+        sender.runOnce(); // receive response 1
+        assertEquals(OptionalInt.of(1), transactionManager.lastAckedSequence(tp0));
         assertFalse(client.hasInFlightRequests());
+        assertEquals(0, sender.inFlightBatches(tp0).size());
         assertTrue(request2.isDone());
         assertEquals(1, request2.get().offset());
     }
@@ -587,7 +730,7 @@ public void testIdempotenceWithMultipleInflights() throws Exception {
     public void testIdempotenceWithMultipleInflightsRetriedInOrder() throws Exception {
         // Send multiple in flight requests, retry them all one at a time, in the correct order.
         final long producerId = 343434L;
-        TransactionManager transactionManager = new TransactionManager();
+        TransactionManager transactionManager = createTransactionManager();
         setupWithTransactionState(transactionManager);
         prepareAndReceiveInitProducerId(producerId, Errors.NONE);
         assertTrue(transactionManager.hasProducerId());
@@ -595,88 +738,91 @@ 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 = appendToAccumulator(tp0);
+        sender.runOnce();
         String nodeId = client.requests().peek().destination();
         Node node = new Node(Integer.valueOf(nodeId), "localhost", 0);
         assertEquals(1, client.inFlightRequestCount());
         assertEquals(1, transactionManager.sequenceNumber(tp0).longValue());
-        assertEquals(-1, transactionManager.lastAckedSequence(tp0));
+        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 = appendToAccumulator(tp0);
+        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 = appendToAccumulator(tp0);
+        sender.runOnce();
 
         assertEquals(3, client.inFlightRequestCount());
         assertEquals(3, transactionManager.sequenceNumber(tp0).longValue());
-        assertEquals(-1, transactionManager.lastAckedSequence(tp0));
+        assertEquals(OptionalInt.empty(), transactionManager.lastAckedSequence(tp0));
         assertFalse(request1.isDone());
         assertFalse(request2.isDone());
         assertFalse(request3.isDone());
         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 = appendToAccumulator(tp0);
 
         assertEquals(2, client.inFlightRequestCount());
-        assertEquals(-1, transactionManager.lastAckedSequence(tp0));
+        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(-1, transactionManager.lastAckedSequence(tp0));
+        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());
 
-        assertEquals(-1, transactionManager.lastAckedSequence(tp0));
+        assertEquals(OptionalInt.empty(), transactionManager.lastAckedSequence(tp0));
 
         sendIdempotentProducerResponse(0, tp0, Errors.NONE, 0L);
-        sender.run(time.milliseconds());  // receive response 1
-        assertEquals(0, transactionManager.lastAckedSequence(tp0));
+        sender.runOnce();  // receive response 1
+        assertEquals(OptionalInt.of(0), transactionManager.lastAckedSequence(tp0));
         assertTrue(request1.isDone());
         assertEquals(0, request1.get().offset());
-
-
         assertFalse(client.hasInFlightRequests());
-        sender.run(time.milliseconds()); // send request 2;
+        assertEquals(0, sender.inFlightBatches(tp0).size());
+
+        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
-        assertEquals(1, transactionManager.lastAckedSequence(tp0));
+        sender.runOnce();  // receive response 2
+        assertEquals(OptionalInt.of(1), transactionManager.lastAckedSequence(tp0));
         assertTrue(request2.isDone());
         assertEquals(1, request2.get().offset());
 
         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.
-        assertEquals(2, transactionManager.lastAckedSequence(tp0));
+        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());
-
         assertEquals(1, client.inFlightRequestCount());
+        assertEquals(1, sender.inFlightBatches(tp0).size());
 
         sendIdempotentProducerResponse(3, tp0, Errors.NONE, 3L);
-        sender.run(time.milliseconds());  // receive response 4
-        assertEquals(3, transactionManager.lastAckedSequence(tp0));
+        sender.runOnce();  // receive response 4
+        assertEquals(OptionalInt.of(3), transactionManager.lastAckedSequence(tp0));
         assertTrue(request4.isDone());
         assertEquals(3, request4.get().offset());
     }
@@ -684,7 +830,7 @@ public void testIdempotenceWithMultipleInflightsRetriedInOrder() throws Exceptio
     @Test
     public void testIdempotenceWithMultipleInflightsWhereFirstFailsFatallyAndSequenceOfFutureBatchesIsAdjusted() throws Exception {
         final long producerId = 343434L;
-        TransactionManager transactionManager = new TransactionManager();
+        TransactionManager transactionManager = createTransactionManager();
         setupWithTransactionState(transactionManager);
         prepareAndReceiveInitProducerId(producerId, Errors.NONE);
         assertTrue(transactionManager.hasProducerId());
@@ -692,48 +838,48 @@ 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 = appendToAccumulator(tp0);
+        sender.runOnce();
         String nodeId = client.requests().peek().destination();
         Node node = new Node(Integer.valueOf(nodeId), "localhost", 0);
         assertEquals(1, client.inFlightRequestCount());
         assertEquals(1, transactionManager.sequenceNumber(tp0).longValue());
-        assertEquals(-1, transactionManager.lastAckedSequence(tp0));
+        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 = appendToAccumulator(tp0);
+        sender.runOnce();
         assertEquals(2, client.inFlightRequestCount());
         assertEquals(2, transactionManager.sequenceNumber(tp0).longValue());
-        assertEquals(-1, transactionManager.lastAckedSequence(tp0));
+        assertEquals(OptionalInt.empty(), transactionManager.lastAckedSequence(tp0));
         assertFalse(request1.isDone());
         assertFalse(request2.isDone());
         assertTrue(client.isReady(node, time.milliseconds()));
 
         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());
-        assertEquals(-1, transactionManager.lastAckedSequence(tp0));
+        assertEquals(OptionalInt.empty(), transactionManager.lastAckedSequence(tp0));
 
         sendIdempotentProducerResponse(1, tp0, Errors.OUT_OF_ORDER_SEQUENCE_NUMBER, -1L);
 
-        sender.run(time.milliseconds()); // receive response 1
+        sender.runOnce(); // receive response 1
 
-        assertEquals(-1, transactionManager.lastAckedSequence(tp0));
+        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(-1, transactionManager.lastAckedSequence(tp0));
+        assertEquals(OptionalInt.empty(), transactionManager.lastAckedSequence(tp0));
 
         sendIdempotentProducerResponse(0, tp0, Errors.NONE, 0L);
-        sender.run(time.milliseconds());  // receive response 1
-        assertEquals(0, transactionManager.lastAckedSequence(tp0));
+        sender.runOnce();  // receive response 1
+        assertEquals(OptionalInt.of(0), transactionManager.lastAckedSequence(tp0));
         assertEquals(0, client.inFlightRequestCount());
 
         assertTrue(request1.isDone());
@@ -741,9 +887,9 @@ public void testIdempotenceWithMultipleInflightsWhereFirstFailsFatallyAndSequenc
     }
 
     @Test
-    public void testMustNotRetryOutOfOrderSequenceForNextBatch() throws Exception {
+    public void testEpochBumpOnOutOfOrderSequenceForNextBatch() throws Exception {
         final long producerId = 343434L;
-        TransactionManager transactionManager = new TransactionManager();
+        TransactionManager transactionManager = createTransactionManager();
         setupWithTransactionState(transactionManager);
         prepareAndReceiveInitProducerId(producerId, Errors.NONE);
         assertTrue(transactionManager.hasProducerId());
@@ -751,63 +897,67 @@ 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 = appendToAccumulator(tp0);
+        appendToAccumulator(tp0);
+        sender.runOnce();
         String nodeId = client.requests().peek().destination();
         Node node = new Node(Integer.valueOf(nodeId), "localhost", 0);
         assertEquals(1, client.inFlightRequestCount());
 
         // make sure the next sequence number accounts for multi-message batches.
         assertEquals(2, transactionManager.sequenceNumber(tp0).longValue());
-        assertEquals(-1, transactionManager.lastAckedSequence(tp0));
+        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 = appendToAccumulator(tp0);
+        sender.runOnce();
         assertEquals(1, client.inFlightRequestCount());
         assertEquals(3, transactionManager.sequenceNumber(tp0).longValue());
-        assertEquals(1, transactionManager.lastAckedSequence(tp0));
+        assertEquals(OptionalInt.of(1), transactionManager.lastAckedSequence(tp0));
         assertTrue(request1.isDone());
         assertEquals(0, request1.get().offset());
         assertFalse(request2.isDone());
         assertTrue(client.isReady(node, time.milliseconds()));
 
-        // This OutOfOrderSequence is fatal since it is returned for the batch succeeding the last acknowledged batch.
+        // This OutOfOrderSequence triggers an epoch bump 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());
-        assertFutureFailure(request2, OutOfOrderSequenceException.class);
+        sender.runOnce();
+        sender.runOnce();
+
+        // epoch should be bumped and sequence numbers reset
+        assertEquals(1, transactionManager.producerIdAndEpoch().epoch);
+        assertEquals(1, transactionManager.sequenceNumber(tp0).intValue());
+        assertEquals(0, transactionManager.firstInFlightSequence(tp0));
     }
 
     @Test
     public void testCorrectHandlingOfOutOfOrderResponses() throws Exception {
         final long producerId = 343434L;
-        TransactionManager transactionManager = new TransactionManager();
+        TransactionManager transactionManager = createTransactionManager();
         setupWithTransactionState(transactionManager);
         prepareAndReceiveInitProducerId(producerId, Errors.NONE);
         assertTrue(transactionManager.hasProducerId());
-
         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 = appendToAccumulator(tp0);
+        sender.runOnce();
         String nodeId = client.requests().peek().destination();
         Node node = new Node(Integer.valueOf(nodeId), "localhost", 0);
         assertEquals(1, client.inFlightRequestCount());
         assertEquals(1, transactionManager.sequenceNumber(tp0).longValue());
-        assertEquals(-1, transactionManager.lastAckedSequence(tp0));
+        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 = appendToAccumulator(tp0);
+        sender.runOnce();
         assertEquals(2, client.inFlightRequestCount());
         assertEquals(2, transactionManager.sequenceNumber(tp0).longValue());
-        assertEquals(-1, transactionManager.lastAckedSequence(tp0));
+        assertEquals(OptionalInt.empty(), transactionManager.lastAckedSequence(tp0));
         assertFalse(request1.isDone());
         assertFalse(request2.isDone());
         assertTrue(client.isReady(node, time.milliseconds()));
@@ -817,50 +967,50 @@ 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.
         assertEquals(1, queuedBatches.size());
         assertEquals(1, queuedBatches.peekFirst().baseSequence());
         assertEquals(1, client.inFlightRequestCount());
-        assertEquals(-1, transactionManager.lastAckedSequence(tp0));
+        assertEquals(OptionalInt.empty(), transactionManager.lastAckedSequence(tp0));
 
-        client.respondToRequest(firstClientRequest, produceResponse(tp0, -1, Errors.NOT_LEADER_FOR_PARTITION, -1));
+        client.respondToRequest(firstClientRequest, produceResponse(tp0, -1, Errors.NOT_LEADER_OR_FOLLOWER, -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());
         assertEquals(0, queuedBatches.peekFirst().baseSequence());
         assertEquals(1, queuedBatches.peekLast().baseSequence());
-        assertEquals(-1, transactionManager.lastAckedSequence(tp0));
+        assertEquals(OptionalInt.empty(), transactionManager.lastAckedSequence(tp0));
         assertEquals(0, client.inFlightRequestCount());
         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(-1, transactionManager.lastAckedSequence(tp0));
+        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
-        assertEquals(0, transactionManager.lastAckedSequence(tp0));
+        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(1, transactionManager.lastAckedSequence(tp0));
+        assertEquals(OptionalInt.of(1), transactionManager.lastAckedSequence(tp0));
         assertTrue(request2.isDone());
         assertEquals(1, request2.get().offset());
     }
@@ -868,7 +1018,7 @@ public void testCorrectHandlingOfOutOfOrderResponses() throws Exception {
     @Test
     public void testCorrectHandlingOfOutOfOrderResponsesWhenSecondSucceeds() throws Exception {
         final long producerId = 343434L;
-        TransactionManager transactionManager = new TransactionManager();
+        TransactionManager transactionManager = createTransactionManager();
         setupWithTransactionState(transactionManager);
         prepareAndReceiveInitProducerId(producerId, Errors.NONE);
         assertTrue(transactionManager.hasProducerId());
@@ -876,15 +1026,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 = appendToAccumulator(tp0);
+        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 = appendToAccumulator(tp0);
+        sender.runOnce();
         assertEquals(2, client.inFlightRequestCount());
         assertFalse(request1.isDone());
         assertFalse(request2.isDone());
@@ -895,7 +1045,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());
@@ -903,29 +1053,29 @@ public void testCorrectHandlingOfOutOfOrderResponsesWhenSecondSucceeds() throws
 
         assertEquals(0, queuedBatches.size());
         assertEquals(1, client.inFlightRequestCount());
-        assertEquals(1, transactionManager.lastAckedSequence(tp0));
+        assertEquals(OptionalInt.of(1), transactionManager.lastAckedSequence(tp0));
 
         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());
         assertEquals(0, queuedBatches.peekFirst().baseSequence());
-        assertEquals(1, transactionManager.lastAckedSequence(tp0));
+        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());
-        assertEquals(1, transactionManager.lastAckedSequence(tp0));
+        assertEquals(OptionalInt.of(1), transactionManager.lastAckedSequence(tp0));
 
         // 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(1, transactionManager.lastAckedSequence(tp0));
+        assertEquals(OptionalInt.of(1), transactionManager.lastAckedSequence(tp0));
         assertEquals(0, client.inFlightRequestCount());
 
         assertFalse(client.hasInFlightRequests());
@@ -936,7 +1086,7 @@ public void testCorrectHandlingOfOutOfOrderResponsesWhenSecondSucceeds() throws
     @Test
     public void testExpiryOfUnsentBatchesShouldNotCauseUnresolvedSequences() throws Exception {
         final long producerId = 343434L;
-        TransactionManager transactionManager = new TransactionManager();
+        TransactionManager transactionManager = createTransactionManager();
         setupWithTransactionState(transactionManager);
         prepareAndReceiveInitProducerId(producerId, Errors.NONE);
         assertTrue(transactionManager.hasProducerId());
@@ -944,13 +1094,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;
-        Node node = this.cluster.nodes().get(0);
+        Future request1 = appendToAccumulator(tp0, 0L, "key", "value");
+        Node node = metadata.fetch().nodes().get(0);
         time.sleep(10000L);
         client.disconnect(node.idString());
-        client.blackout(node, 10);
+        client.backoff(node, 10);
 
-        sender.run(time.milliseconds());
+        sender.runOnce();
 
         assertFutureFailure(request1, TimeoutException.class);
         assertFalse(transactionManager.hasUnresolvedSequence(tp0));
@@ -959,119 +1109,166 @@ public void testExpiryOfUnsentBatchesShouldNotCauseUnresolvedSequences() throws
     @Test
     public void testExpiryOfFirstBatchShouldNotCauseUnresolvedSequencesIfFutureBatchesSucceed() throws Exception {
         final long producerId = 343434L;
-        TransactionManager transactionManager = new TransactionManager();
-        setupWithTransactionState(transactionManager);
+        TransactionManager transactionManager = createTransactionManager();
+        setupWithTransactionState(transactionManager, false, null);
         prepareAndReceiveInitProducerId(producerId, Errors.NONE);
         assertTrue(transactionManager.hasProducerId());
-
         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 request2 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());  // send request
-
+        Future request1 = appendToAccumulator(tp0);
+        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 = appendToAccumulator(tp0);
+        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 = this.cluster.nodes().get(0);
-        time.sleep(10000L);
+        Node node = metadata.fetch().nodes().get(0);
+        // We add 600 millis to expire the first batch but not the second.
+        // Note deliveryTimeoutMs is 1500.
+        time.sleep(600L);
         client.disconnect(node.idString());
-        client.blackout(node, 10);
+        client.backoff(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;
+        assertEquals(0, sender.inFlightBatches(tp0).size());
 
+        // let's enqueue another batch, which should not be dequeued until the unresolved state is clear.
+        Future request3 = appendToAccumulator(tp0);
         time.sleep(20);
-
         assertFalse(request2.isDone());
 
-        sender.run(time.milliseconds());  // send second request
+        sender.runOnce();  // send second request
         sendIdempotentProducerResponse(1, tp0, Errors.NONE, 1);
-        sender.run(time.milliseconds()); // receive second response, the third request shouldn't be sent since we are in an unresolved state.
+        assertEquals(1, sender.inFlightBatches(tp0).size());
+
+        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());
-        Deque batches = accumulator.batches().get(tp0);
+        assertEquals(0, sender.inFlightBatches(tp0).size());
 
+        Deque batches = accumulator.batches().get(tp0);
         assertEquals(1, batches.size());
         assertFalse(batches.peekFirst().hasSequence());
         assertFalse(client.hasInFlightRequests());
         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());
         assertEquals(1, client.inFlightRequestCount());
         assertFalse(request3.isDone());
+        assertEquals(1, sender.inFlightBatches(tp0).size());
     }
 
     @Test
-    public void testExpiryOfFirstBatchShouldCauseResetIfFutureBatchesFail() throws Exception {
+    public void testExpiryOfFirstBatchShouldCauseEpochBumpIfFutureBatchesFail() throws Exception {
         final long producerId = 343434L;
-        TransactionManager transactionManager = new TransactionManager();
+        TransactionManager transactionManager = createTransactionManager();
         setupWithTransactionState(transactionManager);
         prepareAndReceiveInitProducerId(producerId, Errors.NONE);
         assertTrue(transactionManager.hasProducerId());
-
         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 = appendToAccumulator(tp0);
+        sender.runOnce();  // send request
 
-        Future request2 = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());  // send request
+        time.sleep(1000L);
+        Future request2 = appendToAccumulator(tp0);
+        sender.runOnce();  // send request
 
         assertEquals(2, client.inFlightRequestCount());
 
-        sendIdempotentProducerResponse(0, tp0, Errors.NOT_LEADER_FOR_PARTITION, -1);
-        sender.run(time.milliseconds());  // receive first response
+        sendIdempotentProducerResponse(0, tp0, Errors.NOT_LEADER_OR_FOLLOWER, -1);
+        sender.runOnce();  // receive first response
 
-        Node node = this.cluster.nodes().get(0);
-        time.sleep(10000L);
+        Node node = metadata.fetch().nodes().get(0);
+        time.sleep(1000L);
         client.disconnect(node.idString());
-        client.blackout(node, 10);
+        client.backoff(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;
+        appendToAccumulator(tp0);
 
         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.
-        assertFutureFailure(request2, OutOfOrderSequenceException.class);
+        sender.runOnce(); // receive second response, the third request shouldn't be sent since we are in an unresolved state.
 
         Deque batches = accumulator.batches().get(tp0);
 
-        // The second request should not be requeued.
-        assertEquals(1, batches.size());
-        assertFalse(batches.peekFirst().hasSequence());
-        assertFalse(client.hasInFlightRequests());
+        // The epoch should be bumped and the second request should be requeued
+        assertEquals(2, batches.size());
 
-        // The producer state should be reset.
-        assertFalse(transactionManager.hasProducerId());
+        sender.runOnce();
+        assertEquals((short) 1, transactionManager.producerIdAndEpoch().epoch);
+        assertEquals(1, transactionManager.sequenceNumber(tp0).longValue());
         assertFalse(transactionManager.hasUnresolvedSequence(tp0));
     }
 
+    @Test
+    public void testUnresolvedSequencesAreNotFatal() throws Exception {
+        ProducerIdAndEpoch producerIdAndEpoch = new ProducerIdAndEpoch(123456L, (short) 0);
+        apiVersions.update("0", NodeApiVersions.create(ApiKeys.INIT_PRODUCER_ID.id, (short) 0, (short) 3));
+        TransactionManager txnManager = new TransactionManager(logContext, "testUnresolvedSeq", 60000, 100, apiVersions, false);
+
+        setupWithTransactionState(txnManager);
+        doInitTransactions(txnManager, producerIdAndEpoch);
+
+        txnManager.beginTransaction();
+        txnManager.failIfNotReadyForSend();
+        txnManager.maybeAddPartitionToTransaction(tp0);
+        client.prepareResponse(new AddPartitionsToTxnResponse(0, Collections.singletonMap(tp0, Errors.NONE)));
+        sender.runOnce();
+
+        // Send first ProduceRequest
+        Future request1 = appendToAccumulator(tp0);
+        sender.runOnce();  // send request
+
+        time.sleep(1000L);
+        appendToAccumulator(tp0);
+        sender.runOnce();  // send request
+
+        assertEquals(2, client.inFlightRequestCount());
+
+        sendIdempotentProducerResponse(0, tp0, Errors.NOT_LEADER_OR_FOLLOWER, -1);
+        sender.runOnce();  // receive first response
+
+        Node node = metadata.fetch().nodes().get(0);
+        time.sleep(1000L);
+        client.disconnect(node.idString());
+        client.backoff(node, 10);
+
+        sender.runOnce(); // now expire the first batch.
+        assertFutureFailure(request1, TimeoutException.class);
+        assertTrue(txnManager.hasUnresolvedSequence(tp0));
+
+        // Loop once and confirm that the transaction manager does not enter a fatal error state
+        sender.runOnce();
+        assertTrue(txnManager.hasAbortableError());
+    }
+
     @Test
     public void testExpiryOfAllSentBatchesShouldCauseUnresolvedSequences() throws Exception {
         final long producerId = 343434L;
-        TransactionManager transactionManager = new TransactionManager();
+        TransactionManager transactionManager = createTransactionManager();
         setupWithTransactionState(transactionManager);
         prepareAndReceiveInitProducerId(producerId, Errors.NONE);
         assertTrue(transactionManager.hasProducerId());
@@ -1079,69 +1276,67 @@ 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
-        sendIdempotentProducerResponse(0, tp0, Errors.NOT_LEADER_FOR_PARTITION, -1);
-        sender.run(time.milliseconds());  // receive response
+        Future request1 = appendToAccumulator(tp0, 0L, "key", "value");
+        sender.runOnce();  // send request
+        sendIdempotentProducerResponse(0, tp0, Errors.NOT_LEADER_OR_FOLLOWER, -1);
 
+        sender.runOnce();  // receive response
         assertEquals(1L, transactionManager.sequenceNumber(tp0).longValue());
 
-        Node node = this.cluster.nodes().get(0);
-        time.sleep(10000L);
+        Node node = metadata.fetch().nodes().get(0);
+        time.sleep(15000L);
         client.disconnect(node.idString());
-        client.blackout(node, 10);
+        client.backoff(node, 10);
 
-        sender.run(time.milliseconds()); // now expire the batch.
+        sender.runOnce(); // now expire the batch.
 
         assertFutureFailure(request1, TimeoutException.class);
         assertTrue(transactionManager.hasUnresolvedSequence(tp0));
         assertFalse(client.hasInFlightRequests());
         Deque batches = accumulator.batches().get(tp0);
         assertEquals(0, batches.size());
-        assertTrue(transactionManager.hasProducerId(producerId));
-        // We should now clear the old producerId and get a new one in a single run loop.
-        prepareAndReceiveInitProducerId(producerId + 1, Errors.NONE);
-        assertTrue(transactionManager.hasProducerId(producerId + 1));
+        assertEquals(producerId, transactionManager.producerIdAndEpoch().producerId);
+
+        // In the next run loop, we bump the epoch and clear the unresolved sequences
+        sender.runOnce();
+        assertEquals(1, transactionManager.producerIdAndEpoch().epoch);
+        assertFalse(transactionManager.hasUnresolvedSequence(tp0));
     }
 
     @Test
     public void testResetOfProducerStateShouldAllowQueuedBatchesToDrain() throws Exception {
         final long producerId = 343434L;
-        TransactionManager transactionManager = new TransactionManager();
-        transactionManager.setProducerIdAndEpoch(new ProducerIdAndEpoch(producerId, (short) 0));
+        TransactionManager transactionManager = createTransactionManager();
         setupWithTransactionState(transactionManager);
-        client.setNode(new Node(1, "localhost", 33343));
+        prepareAndReceiveInitProducerId(producerId, Short.MAX_VALUE, Errors.NONE);
+        assertTrue(transactionManager.hasProducerId());
 
         int maxRetries = 10;
         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 failedResponse = accumulator.append(tp0, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).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.
+        appendToAccumulator(tp0); // failed response
+        Future successfulResponse = appendToAccumulator(tp1);
+        sender.runOnce();  // connect and send.
 
         assertEquals(1, client.inFlightRequestCount());
 
         Map responses = new LinkedHashMap<>();
-        responses.put(tp1, new OffsetAndError(-1, Errors.NOT_LEADER_FOR_PARTITION));
+        responses.put(tp1, new OffsetAndError(-1, Errors.NOT_LEADER_OR_FOLLOWER));
         responses.put(tp0, new OffsetAndError(-1, Errors.OUT_OF_ORDER_SEQUENCE_NUMBER));
         client.respond(produceResponse(responses));
-        sender.run(time.milliseconds());
-        assertTrue(failedResponse.isDone());
-        assertFalse("Expected transaction state to be reset upon receiving an OutOfOrderSequenceException", transactionManager.hasProducerId());
-        prepareAndReceiveInitProducerId(producerId + 1, Errors.NONE);
+
+        sender.runOnce(); // trigger epoch bump
+        prepareAndReceiveInitProducerId(producerId + 1, Errors.NONE); // also send request to tp1
+        sender.runOnce(); // reset producer ID because epoch is maxed out
         assertEquals(producerId + 1, transactionManager.producerIdAndEpoch().producerId);
-        sender.run(time.milliseconds());  // send request to tp1
 
         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());
@@ -1151,60 +1346,120 @@ public void testResetOfProducerStateShouldAllowQueuedBatchesToDrain() throws Exc
     }
 
     @Test
-    public void testBatchesDrainedWithOldProducerIdShouldFailWithOutOfOrderSequenceOnSubsequentRetry() throws Exception {
+    public void testCloseWithProducerIdReset() throws Exception {
         final long producerId = 343434L;
-        TransactionManager transactionManager = new TransactionManager();
-        transactionManager.setProducerIdAndEpoch(new ProducerIdAndEpoch(producerId, (short) 0));
+        TransactionManager transactionManager = createTransactionManager();
         setupWithTransactionState(transactionManager);
-        client.setNode(new Node(1, "localhost", 33343));
+        prepareAndReceiveInitProducerId(producerId, Short.MAX_VALUE, Errors.NONE);
+        assertTrue(transactionManager.hasProducerId());
+
+        Metrics m = new Metrics();
+        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, RETRY_BACKOFF_MS, transactionManager, apiVersions);
+
+        appendToAccumulator(tp0); // failed response
+        appendToAccumulator(tp1); // success response
+        sender.runOnce();  // connect and send.
+
+        assertEquals(1, client.inFlightRequestCount());
+
+        Map responses = new LinkedHashMap<>();
+        responses.put(tp1, new OffsetAndError(-1, Errors.NOT_LEADER_OR_FOLLOWER));
+        responses.put(tp0, new OffsetAndError(-1, Errors.OUT_OF_ORDER_SEQUENCE_NUMBER));
+        client.respond(produceResponse(responses));
+        sender.initiateClose(); // initiate close
+        sender.runOnce(); // out of order sequence error triggers producer ID reset because epoch is maxed out
+
+        TestUtils.waitForCondition(() -> {
+            prepareInitProducerResponse(Errors.NONE, producerId + 1, (short) 1);
+            sender.runOnce();
+            return !accumulator.hasUndrained();
+        }, 5000, "Failed to drain batches");
+    }
+
+    @Test
+    public void testForceCloseWithProducerIdReset() throws Exception {
+        TransactionManager transactionManager = createTransactionManager();
+        setupWithTransactionState(transactionManager);
+        prepareAndReceiveInitProducerId(1L, Short.MAX_VALUE, Errors.NONE);
+        assertTrue(transactionManager.hasProducerId());
+
+        Metrics m = new Metrics();
+        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, RETRY_BACKOFF_MS, transactionManager, apiVersions);
+
+        Future failedResponse = appendToAccumulator(tp0);
+        Future successfulResponse = appendToAccumulator(tp1);
+        sender.runOnce();  // connect and send.
+
+        assertEquals(1, client.inFlightRequestCount());
+
+        Map responses = new LinkedHashMap<>();
+        responses.put(tp1, new OffsetAndError(-1, Errors.NOT_LEADER_OR_FOLLOWER));
+        responses.put(tp0, new OffsetAndError(-1, Errors.OUT_OF_ORDER_SEQUENCE_NUMBER));
+        client.respond(produceResponse(responses));
+        sender.runOnce(); // out of order sequence error triggers producer ID reset because epoch is maxed out
+        sender.forceClose(); // initiate force close
+        sender.runOnce(); // this should not block
+        sender.run(); // run main loop to test forceClose flag
+        assertFalse("Pending batches are not aborted.", accumulator.hasUndrained());
+        assertTrue(successfulResponse.isDone());
+    }
+
+    @Test
+    public void testBatchesDrainedWithOldProducerIdShouldSucceedOnSubsequentRetry() throws Exception {
+        final long producerId = 343434L;
+        TransactionManager transactionManager = createTransactionManager();
+        setupWithTransactionState(transactionManager);
+        prepareAndReceiveInitProducerId(producerId, Errors.NONE);
+        assertTrue(transactionManager.hasProducerId());
 
         int maxRetries = 10;
         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 failedResponse = accumulator.append(tp0, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), null, null, MAX_BLOCK_TIMEOUT).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.
+        Future outOfOrderResponse = appendToAccumulator(tp0);
+        Future successfulResponse = appendToAccumulator(tp1);
+        sender.runOnce();  // connect.
+        sender.runOnce();  // send.
 
         assertEquals(1, client.inFlightRequestCount());
 
         Map responses = new LinkedHashMap<>();
-        responses.put(tp1, new OffsetAndError(-1, Errors.NOT_LEADER_FOR_PARTITION));
+        responses.put(tp1, new OffsetAndError(-1, Errors.NOT_LEADER_OR_FOLLOWER));
         responses.put(tp0, new OffsetAndError(-1, Errors.OUT_OF_ORDER_SEQUENCE_NUMBER));
         client.respond(produceResponse(responses));
-        sender.run(time.milliseconds());
-        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();
+        assertFalse(outOfOrderResponse.isDone());
+
+        sender.runOnce();  // bump epoch send request to tp1 with the old producerId
+        assertEquals(1, transactionManager.producerIdAndEpoch().epoch);
 
         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());
+        client.respond(produceResponse(tp1, 0, Errors.NOT_LEADER_OR_FOLLOWER, -1));
+        sender.runOnce();
 
+        // The response
+        assertFalse(successfulResponse.isDone());
+        sender.runOnce(); // retry one more time
+        client.respond(produceResponse(tp1, 0, Errors.NONE, -1));
+        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
-        // exception.
-        try {
-            successfulResponse.get();
-            fail("Should have raised an OutOfOrderSequenceException");
-        } catch (Exception e) {
-            assertTrue(e.getCause() instanceof OutOfOrderSequenceException);
-        }
+        assertEquals(0, transactionManager.sequenceNumber(tp1).intValue());
     }
 
     @Test
     public void testCorrectHandlingOfDuplicateSequenceError() throws Exception {
         final long producerId = 343434L;
-        TransactionManager transactionManager = new TransactionManager();
+        TransactionManager transactionManager = createTransactionManager();
         setupWithTransactionState(transactionManager);
         prepareAndReceiveInitProducerId(producerId, Errors.NONE);
         assertTrue(transactionManager.hasProducerId());
@@ -1212,20 +1467,20 @@ 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 = appendToAccumulator(tp0);
+        sender.runOnce();
         String nodeId = client.requests().peek().destination();
         Node node = new Node(Integer.valueOf(nodeId), "localhost", 0);
         assertEquals(1, client.inFlightRequestCount());
         assertEquals(1, transactionManager.sequenceNumber(tp0).longValue());
-        assertEquals(-1, transactionManager.lastAckedSequence(tp0));
+        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 = appendToAccumulator(tp0);
+        sender.runOnce();
         assertEquals(2, client.inFlightRequestCount());
         assertEquals(2, transactionManager.sequenceNumber(tp0).longValue());
-        assertEquals(-1, transactionManager.lastAckedSequence(tp0));
+        assertEquals(OptionalInt.empty(), transactionManager.lastAckedSequence(tp0));
         assertFalse(request1.isDone());
         assertFalse(request2.isDone());
         assertTrue(client.isReady(node, time.milliseconds()));
@@ -1235,18 +1490,18 @@ 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(1000, transactionManager.lastAckedOffset(tp0));
-        assertEquals(1, transactionManager.lastAckedSequence(tp0));
+        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(1, transactionManager.lastAckedSequence(tp0));
-        assertEquals(1000, transactionManager.lastAckedOffset(tp0));
+        assertEquals(OptionalInt.of(1), transactionManager.lastAckedSequence(tp0));
+        assertEquals(OptionalLong.of(1000), transactionManager.lastAckedOffset(tp0));
         assertFalse(client.hasInFlightRequests());
 
         RecordMetadata unknownMetadata = request1.get();
@@ -1255,67 +1510,130 @@ public void testCorrectHandlingOfDuplicateSequenceError() throws Exception {
     }
 
     @Test
-    public void testUnknownProducerHandlingWhenRetentionLimitReached() throws Exception {
+    public void testTransactionalUnknownProducerHandlingWhenRetentionLimitReached() throws Exception {
         final long producerId = 343434L;
-        TransactionManager transactionManager = new TransactionManager();
+        TransactionManager transactionManager = new TransactionManager(logContext, "testUnresolvedSeq", 60000, 100, apiVersions, false);
+
         setupWithTransactionState(transactionManager);
-        prepareAndReceiveInitProducerId(producerId, Errors.NONE);
+        doInitTransactions(transactionManager, new ProducerIdAndEpoch(producerId, (short) 0));
         assertTrue(transactionManager.hasProducerId());
 
+        transactionManager.maybeAddPartitionToTransaction(tp0);
+        client.prepareResponse(new AddPartitionsToTxnResponse(0, Collections.singletonMap(tp0, Errors.NONE)));
+        sender.runOnce(); // Receive AddPartitions response
+
         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 = appendToAccumulator(tp0);
+        sender.runOnce();
 
         assertEquals(1, client.inFlightRequestCount());
         assertEquals(1, transactionManager.sequenceNumber(tp0).longValue());
-        assertEquals(-1, transactionManager.lastAckedSequence(tp0));
+        assertEquals(OptionalInt.empty(), transactionManager.lastAckedSequence(tp0));
 
         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());
-        assertEquals(0L, transactionManager.lastAckedSequence(tp0));
-        assertEquals(1000L, transactionManager.lastAckedOffset(tp0));
+        assertEquals(OptionalInt.of(0), transactionManager.lastAckedSequence(tp0));
+        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());
+        appendToAccumulator(tp0);
+        Future request2 = appendToAccumulator(tp0);
+        sender.runOnce();
         assertEquals(3, transactionManager.sequenceNumber(tp0).longValue());
-        assertEquals(0, transactionManager.lastAckedSequence(tp0));
+        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(-1, transactionManager.lastAckedSequence(tp0));
+        assertEquals(OptionalInt.empty(), transactionManager.lastAckedSequence(tp0));
         assertEquals(2, transactionManager.sequenceNumber(tp0).longValue());
         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
-        assertEquals(1, transactionManager.lastAckedSequence(tp0));
+        sender.runOnce(); // receive response 1
+        assertEquals(OptionalInt.of(1), transactionManager.lastAckedSequence(tp0));
         assertEquals(2, transactionManager.sequenceNumber(tp0).longValue());
         assertFalse(client.hasInFlightRequests());
         assertTrue(request2.isDone());
         assertEquals(1012L, request2.get().offset());
-        assertEquals(1012L, transactionManager.lastAckedOffset(tp0));
+        assertEquals(OptionalLong.of(1012L), transactionManager.lastAckedOffset(tp0));
+    }
+
+    @Test
+    public void testIdempotentUnknownProducerHandlingWhenRetentionLimitReached() throws Exception {
+        final long producerId = 343434L;
+        TransactionManager transactionManager = createTransactionManager();
+        setupWithTransactionState(transactionManager);
+        prepareAndReceiveInitProducerId(producerId, Errors.NONE);
+        assertTrue(transactionManager.hasProducerId());
+
+        assertEquals(0, transactionManager.sequenceNumber(tp0).longValue());
+
+        // Send first ProduceRequest
+        Future request1 = appendToAccumulator(tp0);
+        sender.runOnce();
+
+        assertEquals(1, client.inFlightRequestCount());
+        assertEquals(1, transactionManager.sequenceNumber(tp0).longValue());
+        assertEquals(OptionalInt.empty(), transactionManager.lastAckedSequence(tp0));
+
+        sendIdempotentProducerResponse(0, tp0, Errors.NONE, 1000L, 10L);
+
+        sender.runOnce();  // receive the response.
+
+        assertTrue(request1.isDone());
+        assertEquals(1000L, request1.get().offset());
+        assertEquals(OptionalInt.of(0), transactionManager.lastAckedSequence(tp0));
+        assertEquals(OptionalLong.of(1000L), transactionManager.lastAckedOffset(tp0));
+
+        // Send second ProduceRequest, a single batch with 2 records.
+        appendToAccumulator(tp0);
+        Future request2 = appendToAccumulator(tp0);
+        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.runOnce(); // receive response 0, should be retried since the logStartOffset > lastAckedOffset.
+        sender.runOnce(); // bump epoch and retry request
+
+        // 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));
+        assertEquals(2, transactionManager.sequenceNumber(tp0).longValue());
+        assertFalse(request2.isDone());
+        assertTrue(client.hasInFlightRequests());
+        assertEquals((short) 1, transactionManager.producerIdAndEpoch().epoch);
+
+        // 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.runOnce(); // receive response 1
+        assertEquals(OptionalInt.of(1), transactionManager.lastAckedSequence(tp0));
+        assertEquals(2, transactionManager.sequenceNumber(tp0).longValue());
+        assertFalse(client.hasInFlightRequests());
+        assertTrue(request2.isDone());
+        assertEquals(1012L, request2.get().offset());
+        assertEquals(OptionalLong.of(1012L), transactionManager.lastAckedOffset(tp0));
     }
 
     @Test
     public void testUnknownProducerErrorShouldBeRetriedWhenLogStartOffsetIsUnknown() throws Exception {
         final long producerId = 343434L;
-        TransactionManager transactionManager = new TransactionManager();
+        TransactionManager transactionManager = createTransactionManager();
         setupWithTransactionState(transactionManager);
         prepareAndReceiveInitProducerId(producerId, Errors.NONE);
         assertTrue(transactionManager.hasProducerId());
@@ -1323,57 +1641,57 @@ 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 = appendToAccumulator(tp0);
+        sender.runOnce();
 
         assertEquals(1, client.inFlightRequestCount());
         assertEquals(1, transactionManager.sequenceNumber(tp0).longValue());
-        assertEquals(-1, transactionManager.lastAckedSequence(tp0));
+        assertEquals(OptionalInt.empty(), transactionManager.lastAckedSequence(tp0));
 
         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());
-        assertEquals(0L, transactionManager.lastAckedSequence(tp0));
-        assertEquals(1000L, transactionManager.lastAckedOffset(tp0));
+        assertEquals(OptionalInt.of(0), transactionManager.lastAckedSequence(tp0));
+        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 = appendToAccumulator(tp0);
+        sender.runOnce();
         assertEquals(2, transactionManager.sequenceNumber(tp0).longValue());
-        assertEquals(0, transactionManager.lastAckedSequence(tp0));
+        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(0, transactionManager.lastAckedSequence(tp0));
+        assertEquals(OptionalInt.of(0), transactionManager.lastAckedSequence(tp0));
         assertEquals(2, transactionManager.sequenceNumber(tp0).longValue());
         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
-        assertEquals(1, transactionManager.lastAckedSequence(tp0));
+        sender.runOnce(); // receive response 1
+        assertEquals(OptionalInt.of(1), transactionManager.lastAckedSequence(tp0));
         assertEquals(2, transactionManager.sequenceNumber(tp0).longValue());
         assertFalse(client.hasInFlightRequests());
         assertTrue(request2.isDone());
         assertEquals(1011L, request2.get().offset());
-        assertEquals(1011L, transactionManager.lastAckedOffset(tp0));
+        assertEquals(OptionalLong.of(1011L), transactionManager.lastAckedOffset(tp0));
     }
 
     @Test
     public void testUnknownProducerErrorShouldBeRetriedForFutureBatchesWhenFirstFails() throws Exception {
         final long producerId = 343434L;
-        TransactionManager transactionManager = new TransactionManager();
+        TransactionManager transactionManager = createTransactionManager();
         setupWithTransactionState(transactionManager);
         prepareAndReceiveInitProducerId(producerId, Errors.NONE);
         assertTrue(transactionManager.hasProducerId());
@@ -1381,88 +1699,85 @@ 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 = appendToAccumulator(tp0);
+        sender.runOnce();
 
         assertEquals(1, client.inFlightRequestCount());
         assertEquals(1, transactionManager.sequenceNumber(tp0).longValue());
-        assertEquals(-1, transactionManager.lastAckedSequence(tp0));
+        assertEquals(OptionalInt.empty(), transactionManager.lastAckedSequence(tp0));
 
         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());
-        assertEquals(0L, transactionManager.lastAckedSequence(tp0));
-        assertEquals(1000L, transactionManager.lastAckedOffset(tp0));
+        assertEquals(OptionalInt.of(0), transactionManager.lastAckedSequence(tp0));
+        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 = appendToAccumulator(tp0);
+        sender.runOnce();
         assertEquals(2, transactionManager.sequenceNumber(tp0).longValue());
-        assertEquals(0, transactionManager.lastAckedSequence(tp0));
+        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 = appendToAccumulator(tp0);
+        sender.runOnce();
         assertEquals(3, transactionManager.sequenceNumber(tp0).longValue());
-        assertEquals(0, transactionManager.lastAckedSequence(tp0));
+        assertEquals(OptionalInt.of(0), transactionManager.lastAckedSequence(tp0));
 
         assertFalse(request2.isDone());
         assertFalse(request3.isDone());
         assertEquals(2, client.inFlightRequestCount());
 
-
         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.
+        sender.runOnce(); // bump epoch and retry request 2
 
         // We should have reset the sequence number state of the partition because the state was lost on the broker.
-        assertEquals(-1, transactionManager.lastAckedSequence(tp0));
+        assertEquals(OptionalInt.empty(), transactionManager.lastAckedSequence(tp0));
         assertEquals(2, transactionManager.sequenceNumber(tp0).longValue());
         assertFalse(request2.isDone());
         assertFalse(request3.isDone());
-        assertEquals(1, client.inFlightRequestCount());
-
-        sender.run(time.milliseconds()); // resend request 2.
-
         assertEquals(2, client.inFlightRequestCount());
+        assertEquals((short) 1, transactionManager.producerIdAndEpoch().epoch);
 
         // 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(-1, transactionManager.lastAckedSequence(tp0));
+        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());
         assertFalse(client.hasInFlightRequests());
-        assertEquals(0, transactionManager.lastAckedSequence(tp0));
+        assertEquals(OptionalInt.of(0), transactionManager.lastAckedSequence(tp0));
         assertEquals(1011L, request2.get().offset());
-        assertEquals(1011L, transactionManager.lastAckedOffset(tp0));
+        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());
         assertEquals(1012L, request3.get().offset());
-        assertEquals(1012L, transactionManager.lastAckedOffset(tp0));
+        assertEquals(OptionalLong.of(1012L), transactionManager.lastAckedOffset(tp0));
     }
 
     @Test
     public void testShouldRaiseOutOfOrderSequenceExceptionToUserIfLogWasNotTruncated() throws Exception {
         final long producerId = 343434L;
-        TransactionManager transactionManager = new TransactionManager();
+        TransactionManager transactionManager = createTransactionManager();
         setupWithTransactionState(transactionManager);
         prepareAndReceiveInitProducerId(producerId, Errors.NONE);
         assertTrue(transactionManager.hasProducerId());
@@ -1470,78 +1785,71 @@ 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 = appendToAccumulator(tp0);
+        sender.runOnce();
 
         assertEquals(1, client.inFlightRequestCount());
         assertEquals(1, transactionManager.sequenceNumber(tp0).longValue());
-        assertEquals(-1, transactionManager.lastAckedSequence(tp0));
+        assertEquals(OptionalInt.empty(), transactionManager.lastAckedSequence(tp0));
 
         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());
-        assertEquals(0L, transactionManager.lastAckedSequence(tp0));
-        assertEquals(1000L, transactionManager.lastAckedOffset(tp0));
+        assertEquals(OptionalInt.of(0), transactionManager.lastAckedSequence(tp0));
+        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 = appendToAccumulator(tp0);
+        sender.runOnce();
         assertEquals(2, transactionManager.sequenceNumber(tp0).longValue());
-        assertEquals(0, transactionManager.lastAckedSequence(tp0));
+        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
-        assertFutureFailure(request2, OutOfOrderSequenceException.class);
-
+        sender.runOnce(); // receive response 0, should request an epoch bump
+        sender.runOnce(); // bump epoch
+        assertEquals(1, transactionManager.producerIdAndEpoch().epoch);
+        assertEquals(OptionalInt.empty(), transactionManager.lastAckedSequence(tp0));
+        assertFalse(request2.isDone());
     }
     void sendIdempotentProducerResponse(int expectedSequence, TopicPartition tp, Errors responseError, long responseOffset) {
         sendIdempotentProducerResponse(expectedSequence, tp, responseError, responseOffset, -1L);
     }
 
     void sendIdempotentProducerResponse(final int expectedSequence, TopicPartition tp, Errors responseError, long responseOffset, long logStartOffset) {
-        client.respond(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                ProduceRequest produceRequest = (ProduceRequest) body;
-                assertTrue(produceRequest.isIdempotent());
-
-                MemoryRecords records = produceRequest.partitionRecordsOrFail().get(tp0);
-                Iterator batchIterator = records.batches().iterator();
-                RecordBatch firstBatch = batchIterator.next();
-                assertFalse(batchIterator.hasNext());
-                assertEquals(expectedSequence, firstBatch.baseSequence());
-
-                return true;
-            }
-        }, produceResponse(tp, responseOffset, responseError, 0, logStartOffset));
+        client.respond(body -> {
+            ProduceRequest produceRequest = (ProduceRequest) body;
+            assertTrue(RequestUtils.hasIdempotentRecords(produceRequest));
+
+            MemoryRecords records = partitionRecords(produceRequest).get(tp0);
+            Iterator batchIterator = records.batches().iterator();
+            RecordBatch firstBatch = batchIterator.next();
+            assertFalse(batchIterator.hasNext());
+            assertEquals(expectedSequence, firstBatch.baseSequence());
+            return true;
+        }, produceResponse(tp, responseOffset, responseError, 0, logStartOffset, null));
     }
 
     @Test
     public void testClusterAuthorizationExceptionInProduceRequest() throws Exception {
         final long producerId = 343434L;
-        TransactionManager transactionManager = new TransactionManager();
+        TransactionManager transactionManager = createTransactionManager();
         setupWithTransactionState(transactionManager);
 
-        client.setNode(new Node(1, "localhost", 33343));
         prepareAndReceiveInitProducerId(producerId, Errors.NONE);
         assertTrue(transactionManager.hasProducerId());
 
         // 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;
-        client.prepareResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                return body instanceof ProduceRequest && ((ProduceRequest) body).isIdempotent();
-            }
-        }, produceResponse(tp0, -1, Errors.CLUSTER_AUTHORIZATION_FAILED, 0));
+        Future future = appendToAccumulator(tp0);
+        client.prepareResponse(
+            body -> body instanceof ProduceRequest && RequestUtils.hasIdempotentRecords((ProduceRequest) 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
@@ -1552,66 +1860,52 @@ public boolean matches(AbstractRequest body) {
     @Test
     public void testCancelInFlightRequestAfterFatalError() throws Exception {
         final long producerId = 343434L;
-        TransactionManager transactionManager = new TransactionManager();
+        TransactionManager transactionManager = createTransactionManager();
         setupWithTransactionState(transactionManager);
 
-        client.setNode(new Node(1, "localhost", 33343));
         prepareAndReceiveInitProducerId(producerId, Errors.NONE);
         assertTrue(transactionManager.hasProducerId());
 
         // 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());
+        Future future1 = appendToAccumulator(tp0);
+        sender.runOnce();
 
-        Future future2 = accumulator.append(tp1, time.milliseconds(), "key".getBytes(), "value".getBytes(),
-                null, null, MAX_BLOCK_TIMEOUT).future;
-        sender.run(time.milliseconds());
+        Future future2 = appendToAccumulator(tp1);
+        sender.runOnce();
 
-        client.respond(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                return body instanceof ProduceRequest && ((ProduceRequest) body).isIdempotent();
-            }
-        }, produceResponse(tp0, -1, Errors.CLUSTER_AUTHORIZATION_FAILED, 0));
+        client.respond(
+            body -> body instanceof ProduceRequest && RequestUtils.hasIdempotentRecords((ProduceRequest) 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
-        client.respond(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                return body instanceof ProduceRequest && ((ProduceRequest) body).isIdempotent();
-            }
-        }, produceResponse(tp1, 0, Errors.NONE, 0));
-        sender.run(time.milliseconds());
+        client.respond(
+            body -> body instanceof ProduceRequest && RequestUtils.hasIdempotentRecords((ProduceRequest) body),
+            produceResponse(tp1, 0, Errors.NONE, 0));
+        sender.runOnce();
     }
 
     @Test
     public void testUnsupportedForMessageFormatInProduceRequest() throws Exception {
         final long producerId = 343434L;
-        TransactionManager transactionManager = new TransactionManager();
+        TransactionManager transactionManager = createTransactionManager();
         setupWithTransactionState(transactionManager);
 
-        client.setNode(new Node(1, "localhost", 33343));
         prepareAndReceiveInitProducerId(producerId, Errors.NONE);
         assertTrue(transactionManager.hasProducerId());
 
-        Future future = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(),
-                null, null, MAX_BLOCK_TIMEOUT).future;
-        client.prepareResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                return body instanceof ProduceRequest && ((ProduceRequest) body).isIdempotent();
-            }
-        }, produceResponse(tp0, -1, Errors.UNSUPPORTED_FOR_MESSAGE_FORMAT, 0));
+        Future future = appendToAccumulator(tp0);
+        client.prepareResponse(
+            body -> body instanceof ProduceRequest && RequestUtils.hasIdempotentRecords((ProduceRequest) 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
@@ -1621,23 +1915,17 @@ public boolean matches(AbstractRequest body) {
     @Test
     public void testUnsupportedVersionInProduceRequest() throws Exception {
         final long producerId = 343434L;
-        TransactionManager transactionManager = new TransactionManager();
+        TransactionManager transactionManager = createTransactionManager();
         setupWithTransactionState(transactionManager);
 
-        client.setNode(new Node(1, "localhost", 33343));
         prepareAndReceiveInitProducerId(producerId, Errors.NONE);
         assertTrue(transactionManager.hasProducerId());
 
-        Future future = accumulator.append(tp0, time.milliseconds(), "key".getBytes(), "value".getBytes(),
-                null, null, MAX_BLOCK_TIMEOUT).future;
-        client.prepareUnsupportedVersionResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                return body instanceof ProduceRequest && ((ProduceRequest) body).isIdempotent();
-            }
-        });
+        Future future = appendToAccumulator(tp0);
+        client.prepareUnsupportedVersionResponse(
+            body -> body instanceof ProduceRequest && RequestUtils.hasIdempotentRecords((ProduceRequest) 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
@@ -1648,121 +1936,118 @@ public boolean matches(AbstractRequest body) {
     @Test
     public void testSequenceNumberIncrement() throws InterruptedException {
         final long producerId = 343434L;
-        TransactionManager transactionManager = new TransactionManager();
-        transactionManager.setProducerIdAndEpoch(new ProducerIdAndEpoch(producerId, (short) 0));
+        TransactionManager transactionManager = createTransactionManager();
         setupWithTransactionState(transactionManager);
-        client.setNode(new Node(1, "localhost", 33343));
+        prepareAndReceiveInitProducerId(producerId, Errors.NONE);
+        assertTrue(transactionManager.hasProducerId());
 
         int maxRetries = 10;
         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;
-        client.prepareResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                if (body instanceof ProduceRequest) {
-                    ProduceRequest request = (ProduceRequest) body;
-                    MemoryRecords records = request.partitionRecordsOrFail().get(tp0);
-                    Iterator batchIterator = records.batches().iterator();
-                    assertTrue(batchIterator.hasNext());
-                    RecordBatch batch = batchIterator.next();
-                    assertFalse(batchIterator.hasNext());
-                    assertEquals(0, batch.baseSequence());
-                    assertEquals(producerId, batch.producerId());
-                    assertEquals(0, batch.producerEpoch());
-                    return true;
-                }
-                return false;
+        Future responseFuture = appendToAccumulator(tp0);
+        client.prepareResponse(body -> {
+            if (body instanceof ProduceRequest) {
+                ProduceRequest request = (ProduceRequest) body;
+                MemoryRecords records = partitionRecords(request).get(tp0);
+                Iterator batchIterator = records.batches().iterator();
+                assertTrue(batchIterator.hasNext());
+                RecordBatch batch = batchIterator.next();
+                assertFalse(batchIterator.hasNext());
+                assertEquals(0, batch.baseSequence());
+                assertEquals(producerId, batch.producerId());
+                assertEquals(0, batch.producerEpoch());
+                return true;
             }
+            return false;
         }, 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(0L, (long) transactionManager.lastAckedSequence(tp0));
+        assertEquals(OptionalInt.of(0), transactionManager.lastAckedSequence(tp0));
         assertEquals(1L, (long) transactionManager.sequenceNumber(tp0));
     }
 
     @Test
-    @SuppressWarnings("deprecation")
-    public void testAbortRetryWhenProducerIdChanges() throws InterruptedException {
+    public void testRetryWhenProducerIdChanges() throws InterruptedException {
         final long producerId = 343434L;
-        TransactionManager transactionManager = new TransactionManager();
-        transactionManager.setProducerIdAndEpoch(new ProducerIdAndEpoch(producerId, (short) 0));
+        TransactionManager transactionManager = createTransactionManager();
         setupWithTransactionState(transactionManager);
-        client.setNode(new Node(1, "localhost", 33343));
+        prepareAndReceiveInitProducerId(producerId, Short.MAX_VALUE, Errors.NONE);
+        assertTrue(transactionManager.hasProducerId());
 
         int maxRetries = 10;
         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 = appendToAccumulator(tp0);
+        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));
-
-        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.
-        assertEquals("Expected requests to be aborted after pid change", 0, client.inFlightRequestCount());
+        assertFalse("Client ready status should be false", client.isReady(node, time.milliseconds()));
+        sender.runOnce(); // receive error
+        sender.runOnce(); // reset producer ID because epoch is maxed out
 
-        KafkaMetric recordErrors = m.metrics().get(senderMetrics.recordErrorRate);
-        assertTrue("Expected non-zero value for record send errors", recordErrors.value() > 0);
+        prepareAndReceiveInitProducerId(producerId + 1, Errors.NONE);
+        sender.runOnce(); // nothing to do, since the pid has changed. We should check the metrics for errors.
+        assertEquals("Expected requests to be retried after pid change", 1, client.inFlightRequestCount());
 
-        assertTrue(responseFuture.isDone());
-        assertEquals(0, (long) transactionManager.sequenceNumber(tp0));
+        assertFalse(responseFuture.isDone());
+        assertEquals(1, (long) transactionManager.sequenceNumber(tp0));
     }
 
     @Test
-    public void testResetWhenOutOfOrderSequenceReceived() throws InterruptedException {
+    public void testBumpEpochWhenOutOfOrderSequenceReceived() throws InterruptedException {
         final long producerId = 343434L;
-        TransactionManager transactionManager = new TransactionManager();
-        transactionManager.setProducerIdAndEpoch(new ProducerIdAndEpoch(producerId, (short) 0));
+        TransactionManager transactionManager = createTransactionManager();
         setupWithTransactionState(transactionManager);
-        client.setNode(new Node(1, "localhost", 33343));
+        prepareAndReceiveInitProducerId(producerId, Errors.NONE);
+        assertTrue(transactionManager.hasProducerId());
 
         int maxRetries = 10;
         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 = appendToAccumulator(tp0);
+        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());
-        assertTrue(responseFuture.isDone());
-        assertFalse("Expected transaction state to be reset upon receiving an OutOfOrderSequenceException", transactionManager.hasProducerId());
+        sender.runOnce(); // receive the out of order sequence error
+        sender.runOnce(); // bump the epoch
+        assertFalse(responseFuture.isDone());
+        assertEquals(1, sender.inFlightBatches(tp0).size());
+        assertEquals(1, transactionManager.producerIdAndEpoch().epoch);
     }
 
     @Test
     public void testIdempotentSplitBatchAndSend() throws Exception {
         TopicPartition tp = new TopicPartition("testSplitBatchAndSend", 1);
-        TransactionManager txnManager = new TransactionManager();
+        TransactionManager txnManager = createTransactionManager();
         ProducerIdAndEpoch producerIdAndEpoch = new ProducerIdAndEpoch(123456L, (short) 0);
-        txnManager.setProducerIdAndEpoch(producerIdAndEpoch);
+        setupWithTransactionState(txnManager);
+        prepareAndReceiveInitProducerId(123456L, Errors.NONE);
+        assertTrue(txnManager.hasProducerId());
         testSplitBatchAndSend(txnManager, producerIdAndEpoch, tp);
     }
 
@@ -1770,15 +2055,16 @@ public void testIdempotentSplitBatchAndSend() throws Exception {
     public void testTransactionalSplitBatchAndSend() throws Exception {
         ProducerIdAndEpoch producerIdAndEpoch = new ProducerIdAndEpoch(123456L, (short) 0);
         TopicPartition tp = new TopicPartition("testSplitBatchAndSend", 1);
-        TransactionManager txnManager = new TransactionManager(logContext, "testSplitBatchAndSend", 60000, 100);
+        TransactionManager txnManager = new TransactionManager(logContext, "testSplitBatchAndSend", 60000, 100, apiVersions, false);
 
         setupWithTransactionState(txnManager);
         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);
     }
@@ -1789,41 +2075,46 @@ private void testSplitBatchAndSend(TransactionManager txnManager,
                                        TopicPartition tp) throws Exception {
         int maxRetries = 1;
         String topic = tp.topic();
+        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, 1024 * 1024, CompressionType.GZIP, 0L, 0L, m, time,
-                    new ApiVersions(), txnManager);
+            accumulator = new RecordAccumulator(logContext, batchSize, CompressionType.GZIP,
+                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,
                     senderMetrics, time, REQUEST_TIMEOUT, 1000L, txnManager, new ApiVersions());
             // Create a two broker cluster, with partition 0 on broker 0 and partition 1 on broker 1
-            Cluster cluster1 = TestUtils.clusterWith(2, topic, 2);
-            metadata.update(cluster1, Collections.emptySet(), time.milliseconds());
+            MetadataResponse metadataUpdate1 = RequestTestUtils.metadataUpdateWith(2, Collections.singletonMap(topic, 2));
+            client.prepareMetadataUpdate(metadataUpdate1);
             // Send the first message.
+            long nowMs = time.milliseconds();
             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, nowMs).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, nowMs).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());
@@ -1831,38 +2122,488 @@ 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", 0, txnManager.lastAckedSequence(tp));
+            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", 1, txnManager.lastAckedSequence(tp));
+            assertEquals("The last ack'd sequence number should be 1", OptionalInt.of(1), txnManager.lastAckedSequence(tp));
             assertEquals("Offset of the first message should be 1", 1L, f2.get().offset());
             assertTrue("There should be no batch in the accumulator", accumulator.batches().get(tp).isEmpty());
+            assertTrue("There should be a split", (Double) (m.metrics().get(senderMetrics.batchSplitRate).metricValue()) > 0);
+        }
+    }
+
+    @Test
+    public void testNoDoubleDeallocation() throws Exception {
+        long deliverTimeoutMs = 1500L;
+        long totalSize = 1024 * 1024;
+        String metricGrpName = "producer-custom-metrics";
+        MatchingBufferPool pool = new MatchingBufferPool(totalSize, batchSize, metrics, time, metricGrpName);
+        setupWithTransactionState(null, false, pool);
+
+        // Send first ProduceRequest
+        Future request1 = appendToAccumulator(tp0);
+        sender.runOnce();  // send request
+        assertEquals(1, client.inFlightRequestCount());
+        assertEquals(1, sender.inFlightBatches(tp0).size());
+
+        time.sleep(deliverTimeoutMs);
+        assertFalse(pool.allMatch());
+
+        sender.runOnce();  // expire the batch
+        assertTrue(request1.isDone());
+        assertTrue("The batch should have been de-allocated", pool.allMatch());
+        assertTrue(pool.allMatch());
+
+        sender.runOnce();
+        assertTrue("The batch should have been de-allocated", pool.allMatch());
+        assertEquals(0, client.inFlightRequestCount());
+        assertEquals(0, sender.inFlightBatches(tp0).size());
+    }
+
+    @SuppressWarnings("deprecation")
+    @Test
+    public void testInflightBatchesExpireOnDeliveryTimeout() throws InterruptedException {
+        long deliveryTimeoutMs = 1500L;
+        setupWithTransactionState(null, true, null);
+
+        // Send first ProduceRequest
+        Future request = appendToAccumulator(tp0);
+        sender.runOnce();  // send request
+        assertEquals(1, client.inFlightRequestCount());
+        assertEquals("Expect one in-flight batch in accumulator", 1, sender.inFlightBatches(tp0).size());
+
+        Map responseMap = new HashMap<>();
+        responseMap.put(tp0, new ProduceResponse.PartitionResponse(Errors.NONE, 0L, 0L, 0L));
+        client.respond(new ProduceResponse(responseMap));
+
+        time.sleep(deliveryTimeoutMs);
+        sender.runOnce();  // receive first response
+        assertEquals("Expect zero in-flight batch in accumulator", 0, sender.inFlightBatches(tp0).size());
+        try {
+            request.get();
+            fail("The expired batch should throw a TimeoutException");
+        } catch (ExecutionException e) {
+            assertTrue(e.getCause() instanceof TimeoutException);
+        }
+    }
+
+    @Test
+    public void testWhenFirstBatchExpireNoSendSecondBatchIfGuaranteeOrder() throws InterruptedException {
+        long deliveryTimeoutMs = 1500L;
+        setupWithTransactionState(null, true, null);
+
+        // Send first ProduceRequest
+        appendToAccumulator(tp0);
+        sender.runOnce();  // send request
+        assertEquals(1, client.inFlightRequestCount());
+        assertEquals(1, sender.inFlightBatches(tp0).size());
+
+        time.sleep(deliveryTimeoutMs / 2);
+
+        // Send second ProduceRequest
+        appendToAccumulator(tp0);
+        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, null));
+        sender.runOnce();  // receive response (offset=0)
+        assertEquals(0, client.inFlightRequestCount());
+        assertEquals(0, sender.inFlightBatches(tp0).size());
+
+        sender.runOnce();  // Drain the second request only this time
+        assertEquals(1, client.inFlightRequestCount());
+        assertEquals(1, sender.inFlightBatches(tp0).size());
+    }
+
+    @Test
+    public void testExpiredBatchDoesNotRetry() throws Exception {
+        long deliverTimeoutMs = 1500L;
+        setupWithTransactionState(null, false, null);
+
+        // Send first ProduceRequest
+        Future request1 = appendToAccumulator(tp0);
+        sender.runOnce();  // send request
+        assertEquals(1, client.inFlightRequestCount());
+        time.sleep(deliverTimeoutMs);
+
+        client.respond(produceResponse(tp0, -1, Errors.NOT_LEADER_OR_FOLLOWER, -1)); // return a retriable error
+
+        sender.runOnce();  // expire the batch
+        assertTrue(request1.isDone());
+        assertEquals(0, client.inFlightRequestCount());
+        assertEquals(0, sender.inFlightBatches(tp0).size());
 
-            assertTrue("There should be a split",
-                    m.metrics().get(senderMetrics.batchSplitRate).value() > 0);
+        sender.runOnce(); // receive first response and do not reenqueue.
+        assertEquals(0, client.inFlightRequestCount());
+        assertEquals(0, sender.inFlightBatches(tp0).size());
+
+        sender.runOnce(); // run again and must not send anything.
+        assertEquals(0, client.inFlightRequestCount());
+        assertEquals(0, sender.inFlightBatches(tp0).size());
+    }
+
+    @Test
+    public void testExpiredBatchDoesNotSplitOnMessageTooLargeError() throws Exception {
+        long deliverTimeoutMs = 1500L;
+        // create a producer batch with more than one record so it is eligible for splitting
+        Future request1 = appendToAccumulator(tp0);
+        Future request2 = appendToAccumulator(tp0);
+
+        // 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.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.runOnce();
+        assertEquals(0, client.inFlightRequestCount());
+        assertEquals(0, sender.inFlightBatches(tp0).size());
+    }
+
+    @Test
+    public void testResetNextBatchExpiry() throws Exception {
+        client = spy(new MockClient(time, metadata));
+
+        setupWithTransactionState(null);
+
+        appendToAccumulator(tp0, 0L, "key", "value");
+
+        sender.runOnce();
+        sender.runOnce();
+        time.setCurrentTimeMs(time.milliseconds() + accumulator.getDeliveryTimeoutMs() + 1);
+        sender.runOnce();
+
+        InOrder inOrder = inOrder(client);
+        inOrder.verify(client, atLeastOnce()).ready(any(), anyLong());
+        inOrder.verify(client, atLeastOnce()).newClientRequest(anyString(), any(), anyLong(), anyBoolean(), anyInt(), any());
+        inOrder.verify(client, atLeastOnce()).send(any(), anyLong());
+        inOrder.verify(client).poll(eq(0L), anyLong());
+        inOrder.verify(client).poll(eq(accumulator.getDeliveryTimeoutMs()), anyLong());
+        inOrder.verify(client).poll(geq(1L), anyLong());
+
+    }
+
+    @SuppressWarnings("deprecation")
+    @Test
+    public void testExpiredBatchesInMultiplePartitions() throws Exception {
+        long deliveryTimeoutMs = 1500L;
+        setupWithTransactionState(null, true, null);
+
+        // Send multiple ProduceRequest across multiple partitions.
+        Future request1 = appendToAccumulator(tp0, time.milliseconds(), "k1", "v1");
+        Future request2 = appendToAccumulator(tp1, time.milliseconds(), "k2", "v2");
+
+        // Send request.
+        sender.runOnce();
+        assertEquals(1, client.inFlightRequestCount());
+        assertEquals("Expect one in-flight batch in accumulator", 1, sender.inFlightBatches(tp0).size());
+
+        Map responseMap = new HashMap<>();
+        responseMap.put(tp0, new ProduceResponse.PartitionResponse(Errors.NONE, 0L, 0L, 0L));
+        client.respond(new ProduceResponse(responseMap));
+
+        // Successfully expire both batches.
+        time.sleep(deliveryTimeoutMs);
+        sender.runOnce();
+        assertEquals("Expect zero in-flight batch in accumulator", 0, sender.inFlightBatches(tp0).size());
+
+        try {
+            request1.get();
+            fail("The expired batch should throw a TimeoutException");
+        } catch (ExecutionException e) {
+            assertTrue(e.getCause() instanceof TimeoutException);
+        }
+
+        try {
+            request2.get();
+            fail("The expired batch should throw a TimeoutException");
+        } catch (ExecutionException e) {
+            assertTrue(e.getCause() instanceof TimeoutException);
+        }
+    }
+
+    @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, apiVersions, false);
+            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(new EndTxnResponseData()
+                                                                         .setErrorCode(Errors.NONE.code())
+                                                                         .setThrottleTimeMs(0)));
+            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, apiVersions, false);
+            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(new EndTxnResponseData()
+                                                                         .setErrorCode(Errors.NONE.code())
+                                                                         .setThrottleTimeMs(0)));
+            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, apiVersions, false);
+            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 testTransactionAbortedExceptionOnAbortWithoutError() throws InterruptedException, ExecutionException {
+        ProducerIdAndEpoch producerIdAndEpoch = new ProducerIdAndEpoch(123456L, (short) 0);
+        TransactionManager txnManager = new TransactionManager(logContext, "testTransactionAbortedExceptionOnAbortWithoutError", 60000, 100, apiVersions, false);
+
+        setupWithTransactionState(txnManager, false, null);
+        doInitTransactions(txnManager, producerIdAndEpoch);
+        // Begin the transaction
+        txnManager.beginTransaction();
+        txnManager.maybeAddPartitionToTransaction(tp0);
+        client.prepareResponse(new AddPartitionsToTxnResponse(0, Collections.singletonMap(tp0, Errors.NONE)));
+        // Run it once so that the partition is added to the transaction.
+        sender.runOnce();
+        // Append a record to the accumulator.
+        FutureRecordMetadata metadata = appendToAccumulator(tp0, time.milliseconds(), "key", "value");
+        // Now abort the transaction manually.
+        txnManager.beginAbort();
+        // Try to send.
+        // This should abort the existing transaction and
+        // drain all the unsent batches with a TransactionAbortedException.
+        sender.runOnce();
+        // Now attempt to fetch the result for the record.
+        TestUtils.assertFutureThrows(metadata, TransactionAbortedException.class);
+    }
+
+    @Test
+    public void testDoNotPollWhenNoRequestSent() {
+        client = spy(new MockClient(time, metadata));
+
+        TransactionManager txnManager = new TransactionManager(logContext, "testDoNotPollWhenNoRequestSent", 6000, 100, apiVersions, false);
+        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, apiVersions, false);
+
+        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
+        appendToAccumulator(tp0, time.milliseconds(), "key1", "value1");
+        appendToAccumulator(tp0, time.milliseconds(), "key2", "value2");
+
+        // 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();
+    }
+
+    @Test
+    public void testDefaultErrorMessage() throws Exception {
+        verifyErrorMessage(produceResponse(tp0, 0L, Errors.INVALID_REQUEST, 0), Errors.INVALID_REQUEST.message());
+    }
+
+    @Test
+    public void testCustomErrorMessage() throws Exception {
+        String errorMessage = "testCustomErrorMessage";
+        verifyErrorMessage(produceResponse(tp0, 0L, Errors.INVALID_REQUEST, 0, -1, errorMessage), errorMessage);
+    }
+
+    private void verifyErrorMessage(ProduceResponse response, String expectedMessage) throws Exception {
+        Future future = appendToAccumulator(tp0, 0L, "key", "value");
+        sender.runOnce(); // connect
+        sender.runOnce(); // send produce request
+        client.respond(response);
+        sender.runOnce();
+        sender.runOnce();
+        ExecutionException e1 = assertThrows(ExecutionException.class, () -> future.get(5, TimeUnit.SECONDS));
+        assertEquals(InvalidRequestException.class, e1.getCause().getClass());
+        assertEquals(expectedMessage, e1.getCause().getMessage());
+    }
+
+    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).result());
+                matched = true;
+                return true;
+            } else {
+                return false;
+            }
+        }
+    }
+
+    private class MatchingBufferPool extends BufferPool {
+        IdentityHashMap allocatedBuffers;
+
+        MatchingBufferPool(long totalSize, int batchSize, Metrics metrics, Time time, String metricGrpName) {
+            super(totalSize, batchSize, metrics, time, metricGrpName);
+            allocatedBuffers = new IdentityHashMap<>();
+        }
+
+        @Override
+        public ByteBuffer allocate(int size, long maxTimeToBlockMs) throws InterruptedException {
+            ByteBuffer buffer = super.allocate(size, maxTimeToBlockMs);
+            allocatedBuffers.put(buffer, Boolean.TRUE);
+            return buffer;
+        }
+
+        @Override
+        public void deallocate(ByteBuffer buffer, int size) {
+            if (!allocatedBuffers.containsKey(buffer)) {
+                throw new IllegalStateException("Deallocating a buffer that is not allocated");
+            }
+            allocatedBuffers.remove(buffer);
+            super.deallocate(buffer, size);
+        }
+
+        public boolean allMatch() {
+            return allocatedBuffers.isEmpty();
         }
     }
 
@@ -1870,29 +2611,26 @@ private MockClient.RequestMatcher produceRequestMatcher(final TopicPartition tp,
                                                             final ProducerIdAndEpoch producerIdAndEpoch,
                                                             final int sequence,
                                                             final boolean isTransactional) {
-        return new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                if (!(body instanceof ProduceRequest))
-                    return false;
+        return body -> {
+            if (!(body instanceof ProduceRequest))
+                return false;
 
-                ProduceRequest request = (ProduceRequest) body;
-                Map recordsMap = request.partitionRecordsOrFail();
-                MemoryRecords records = recordsMap.get(tp);
-                if (records == null)
-                    return false;
+            ProduceRequest request = (ProduceRequest) body;
+            Map recordsMap = partitionRecords(request);
+            MemoryRecords records = recordsMap.get(tp);
+            if (records == null)
+                return false;
 
-                List batches = TestUtils.toList(records.batches());
-                if (batches.isEmpty() || batches.size() > 1)
-                    return false;
+            List batches = TestUtils.toList(records.batches());
+            if (batches.size() != 1)
+                return false;
 
-                MutableRecordBatch batch = batches.get(0);
-                return batch.baseOffset() == 0L &&
-                        batch.baseSequence() == sequence &&
-                        batch.producerId() == producerIdAndEpoch.producerId &&
-                        batch.producerEpoch() == producerIdAndEpoch.epoch &&
-                        batch.isTransactional() == isTransactional;
-            }
+            MutableRecordBatch batch = batches.get(0);
+            return batch.baseOffset() == 0L &&
+                    batch.baseSequence() == sequence &&
+                    batch.producerId() == producerIdAndEpoch.producerId &&
+                    batch.producerEpoch() == producerIdAndEpoch.epoch &&
+                    batch.isTransactional() == isTransactional;
         };
     }
 
@@ -1905,12 +2643,24 @@ class OffsetAndError {
         }
     }
 
-    private ProduceResponse produceResponse(TopicPartition tp, long offset, Errors error, int throttleTimeMs, long logStartOffset) {
-        ProduceResponse.PartitionResponse resp = new ProduceResponse.PartitionResponse(error, offset, RecordBatch.NO_TIMESTAMP, logStartOffset);
+    private FutureRecordMetadata appendToAccumulator(TopicPartition tp) throws InterruptedException {
+        return appendToAccumulator(tp, time.milliseconds(), "key", "value");
+    }
+
+    private FutureRecordMetadata appendToAccumulator(TopicPartition tp, long timestamp, String key, String value) throws InterruptedException {
+        return accumulator.append(tp, timestamp, key.getBytes(), value.getBytes(), Record.EMPTY_HEADERS,
+                null, MAX_BLOCK_TIMEOUT, false, time.milliseconds()).future;
+    }
+
+    @SuppressWarnings("deprecation")
+    private ProduceResponse produceResponse(TopicPartition tp, long offset, Errors error, int throttleTimeMs, long logStartOffset, String errorMessage) {
+        ProduceResponse.PartitionResponse resp = new ProduceResponse.PartitionResponse(error, offset,
+                RecordBatch.NO_TIMESTAMP, logStartOffset, Collections.emptyList(), errorMessage);
         Map partResp = Collections.singletonMap(tp, resp);
         return new ProduceResponse(partResp, throttleTimeMs);
     }
 
+    @SuppressWarnings("deprecation")
     private ProduceResponse produceResponse(Map responses) {
         Map partResponses = new LinkedHashMap<>();
         for (Map.Entry entry : responses.entrySet()) {
@@ -1922,27 +2672,43 @@ private ProduceResponse produceResponse(Map resp
 
     }
     private ProduceResponse produceResponse(TopicPartition tp, long offset, Errors error, int throttleTimeMs) {
-        return produceResponse(tp, offset, error, throttleTimeMs, -1L);
+        return produceResponse(tp, offset, error, throttleTimeMs, -1L, null);
     }
 
+    private TransactionManager createTransactionManager() {
+        return new TransactionManager(new LogContext(), null, 0, 100L, new ApiVersions(), false);
+    }
+    
     private void setupWithTransactionState(TransactionManager transactionManager) {
-        Map metricTags = new LinkedHashMap<>();
-        metricTags.put("client-id", CLIENT_ID);
-        MetricConfig metricConfig = new MetricConfig().tags(metricTags);
+        setupWithTransactionState(transactionManager, false, null, true);
+    }
+
+    private void setupWithTransactionState(TransactionManager transactionManager, boolean guaranteeOrder, BufferPool customPool) {
+        setupWithTransactionState(transactionManager, guaranteeOrder, customPool, true);
+    }
+
+    private void setupWithTransactionState(TransactionManager transactionManager, boolean guaranteeOrder, BufferPool customPool, boolean updateMetadata) {
+        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);
-        this.accumulator = new RecordAccumulator(logContext, batchSize, 1024 * 1024, CompressionType.NONE, 0L, 0L, metrics, time,
-                apiVersions, transactionManager);
+        BufferPool pool = (customPool == null) ? new BufferPool(totalSize, batchSize, metrics, time, metricGrpName) : customPool;
+
+        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, RETRY_BACKOFF_MS, transactionManager, apiVersions);
 
-        this.sender = new Sender(logContext, this.client, this.metadata, this.accumulator, false, MAX_REQUEST_SIZE, ACKS_ALL,
-                Integer.MAX_VALUE, this.senderMetricsRegistry, this.time, REQUEST_TIMEOUT, 50, transactionManager, apiVersions);
-        this.metadata.update(this.cluster, Collections.emptySet(), time.milliseconds());
+        metadata.add("test", time.milliseconds());
+        if (updateMetadata)
+            this.client.updateMetadata(RequestTestUtils.metadataUpdateWith(1, Collections.singletonMap("test", 2)));
     }
 
     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());
+        Future future = appendToAccumulator(tp0);
+        sender.runOnce();
         assertTrue(future.isDone());
         try {
             future.get();
@@ -1953,36 +2719,47 @@ private void assertSendFailure(Class expectedError)
     }
 
     private void prepareAndReceiveInitProducerId(long producerId, Errors error) {
-        short producerEpoch = 0;
+        prepareAndReceiveInitProducerId(producerId, (short) 0, error);
+    }
+
+    private void prepareAndReceiveInitProducerId(long producerId, short producerEpoch, 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 -> 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, cluster.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)
@@ -1997,4 +2774,43 @@ private void assertFutureFailure(Future future, Class ex
         }
     }
 
+    private void createMockClientWithMaxFlightOneMetadataPending() {
+        client = new MockClient(time, metadata) {
+            volatile boolean canSendMore = true;
+            @Override
+            public Node leastLoadedNode(long now) {
+                for (Node node : metadata.fetch().nodes()) {
+                    if (isReady(node, now) && canSendMore)
+                        return node;
+                }
+                return null;
+            }
+
+            @Override
+            public List poll(long timeoutMs, long now) {
+                canSendMore = inFlightRequestCount() < 1;
+                return super.poll(timeoutMs, now);
+            }
+        };
+
+        // Send metadata request and wait until request is sent. `leastLoadedNode` will be null once
+        // request is in progress since no more requests can be sent to the node. Node will be ready
+        // on the next poll() after response is processed later on in tests which use this method.
+        MetadataRequest.Builder builder = new MetadataRequest.Builder(Collections.emptyList(), false);
+        Node node = metadata.fetch().nodes().get(0);
+        ClientRequest request = client.newClientRequest(node.idString(), builder, time.milliseconds(), true);
+        while (!client.ready(node, time.milliseconds()))
+            client.poll(0, time.milliseconds());
+        client.send(request, time.milliseconds());
+        while (client.leastLoadedNode(time.milliseconds()) != null)
+            client.poll(0, time.milliseconds());
+    }
+
+    private void waitForProducerId(TransactionManager transactionManager, ProducerIdAndEpoch producerIdAndEpoch) {
+        for (int i = 0; i < 5 && !transactionManager.hasProducerId(); i++)
+            sender.runOnce();
+
+        assertTrue(transactionManager.hasProducerId());
+        assertEquals(producerIdAndEpoch, transactionManager.producerIdAndEpoch());
+    }
 }
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 fab139ab816c5..c2977ea9f66d1 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
@@ -16,10 +16,17 @@
  */
 package org.apache.kafka.clients.producer.internals;
 
+import org.apache.kafka.clients.ApiVersion;
 import org.apache.kafka.clients.ApiVersions;
-import org.apache.kafka.clients.Metadata;
 import org.apache.kafka.clients.MockClient;
+import org.apache.kafka.clients.consumer.CommitFailedException;
+import org.apache.kafka.clients.consumer.ConsumerGroupMetadata;
+import org.apache.kafka.clients.NodeApiVersions;
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
+import org.apache.kafka.common.errors.FencedInstanceIdException;
+import org.apache.kafka.common.requests.JoinGroupRequest;
+import org.apache.kafka.common.requests.RequestTestUtils;
+import org.apache.kafka.common.utils.ProducerIdAndEpoch;
 import org.apache.kafka.clients.producer.RecordMetadata;
 import org.apache.kafka.common.Cluster;
 import org.apache.kafka.common.KafkaException;
@@ -28,21 +35,27 @@
 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.metrics.MetricConfig;
+import org.apache.kafka.common.message.AddOffsetsToTxnResponseData;
+import org.apache.kafka.common.message.EndTxnResponseData;
+import org.apache.kafka.common.message.InitProducerIdResponseData;
 import org.apache.kafka.common.metrics.Metrics;
+import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.protocol.Errors;
 import org.apache.kafka.common.record.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.requests.AbstractRequest;
+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;
@@ -65,25 +78,35 @@
 import org.junit.Before;
 import org.junit.Test;
 
+import java.nio.ByteBuffer;
 import java.util.Arrays;
 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;
+import java.util.Optional;
+import java.util.OptionalInt;
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Supplier;
 
 import static java.util.Collections.singleton;
 import static java.util.Collections.singletonList;
 import static java.util.Collections.singletonMap;
+import static org.hamcrest.CoreMatchers.instanceOf;
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.MatcherAssert.assertThat;
 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.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
@@ -91,62 +114,103 @@ public class TransactionManagerTest {
     private static final int MAX_REQUEST_SIZE = 1024 * 1024;
     private static final short ACKS_ALL = -1;
     private static final int MAX_RETRIES = Integer.MAX_VALUE;
-    private static final String CLIENT_ID = "clientId";
     private static final int MAX_BLOCK_TIMEOUT = 1000;
     private static final int REQUEST_TIMEOUT = 1000;
     private static final long DEFAULT_RETRY_BACKOFF_MS = 100L;
+
     private final String transactionalId = "foobar";
     private final int transactionTimeoutMs = 1121;
 
     private final String topic = "test";
-    private TopicPartition tp0 = new TopicPartition(topic, 0);
-    private TopicPartition tp1 = new TopicPartition(topic, 1);
-    private MockTime time = new MockTime();
-    private MockClient client = new MockClient(time);
-
-    private Metadata metadata = new Metadata(0, Long.MAX_VALUE, true, true, new ClusterResourceListeners());
-    private ApiVersions apiVersions = new ApiVersions();
-    private Cluster cluster = TestUtils.singletonCluster("test", 2);
+    private final TopicPartition tp0 = new TopicPartition(topic, 0);
+    private final TopicPartition tp1 = new TopicPartition(topic, 1);
+    private final long producerId = 13131L;
+    private final short epoch = 1;
+    private final String consumerGroupId = "myConsumerGroup";
+    private final String memberId = "member";
+    private final int generationId = 5;
+    private final String groupInstanceId = "instance";
+
+    private final LogContext logContext = new LogContext();
+    private final MockTime time = new MockTime();
+    private final ProducerMetadata metadata = new ProducerMetadata(0, Long.MAX_VALUE, Long.MAX_VALUE,
+            logContext, new ClusterResourceListeners(), time);
+    private final MockClient client = new MockClient(time, metadata);
+    private final ApiVersions apiVersions = new ApiVersions();
+
     private RecordAccumulator accumulator = null;
     private Sender sender = null;
     private TransactionManager transactionManager = null;
     private Node brokerNode = null;
-    private final LogContext logContext = new LogContext();
 
     @Before
     public void setup() {
-        Map metricTags = new LinkedHashMap<>();
-        metricTags.put("client-id", CLIENT_ID);
+        this.metadata.add("test", time.milliseconds());
+        this.client.updateMetadata(RequestTestUtils.metadataUpdateWith(1, singletonMap("test", 2)));
+        this.brokerNode = new Node(0, "localhost", 2211);
+
+        initializeTransactionManager(Optional.of(transactionalId), false);
+    }
+
+    private void initializeTransactionManager(Optional transactionalId, boolean autoDowngrade) {
+        Metrics metrics = new Metrics(time);
+
+        apiVersions.update("0", new NodeApiVersions(Arrays.asList(
+                new ApiVersion(ApiKeys.INIT_PRODUCER_ID.id, (short) 0, (short) 3),
+                new ApiVersion(ApiKeys.PRODUCE.id, (short) 0, (short) 7))));
+        this.transactionManager = new TransactionManager(logContext, transactionalId.orElse(null),
+                transactionTimeoutMs, DEFAULT_RETRY_BACKOFF_MS, apiVersions, autoDowngrade);
+
         int batchSize = 16 * 1024;
-        MetricConfig metricConfig = new MetricConfig().tags(metricTags);
+        int deliveryTimeoutMs = 3000;
+        long totalSize = 1024 * 1024;
+        String metricGrpName = "producer-metrics";
+
         this.brokerNode = new Node(0, "localhost", 2211);
-        this.transactionManager = new TransactionManager(logContext, transactionalId, transactionTimeoutMs,
-                DEFAULT_RETRY_BACKOFF_MS);
-        Metrics metrics = new Metrics(metricConfig, time);
-        SenderMetricsRegistry senderMetrics = new SenderMetricsRegistry(metrics);
+        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, MAX_RETRIES, new SenderMetricsRegistry(metrics), this.time, REQUEST_TIMEOUT,
+                50, transactionManager, apiVersions);
+    }
+
+    @Test
+    public void testSenderShutdownWithPendingTransactions() throws Exception {
+        doInitTransactions();
+        transactionManager.beginTransaction();
+
+        transactionManager.failIfNotReadyForSend();
+        transactionManager.maybeAddPartitionToTransaction(tp0);
+        FutureRecordMetadata sendFuture = appendToAccumulator(tp0);
 
-        this.accumulator = new RecordAccumulator(logContext, batchSize, 1024 * 1024, CompressionType.NONE, 0L, 0L, metrics, time, apiVersions, transactionManager);
-        this.sender = new Sender(logContext, this.client, this.metadata, this.accumulator, true, MAX_REQUEST_SIZE, ACKS_ALL,
-                MAX_RETRIES, senderMetrics, this.time, REQUEST_TIMEOUT, 50, transactionManager, apiVersions);
-        this.metadata.update(this.cluster, Collections.emptySet(), time.milliseconds());
-        client.setNode(brokerNode);
+        prepareAddPartitionsToTxn(tp0, Errors.NONE);
+        prepareProduceResponse(Errors.NONE, producerId, epoch);
+        runUntil(() -> !client.hasPendingResponses());
+
+        sender.initiateClose();
+        sender.runOnce();
+
+        TransactionalRequestResult result = transactionManager.beginCommit();
+        prepareEndTxnResponse(Errors.NONE, TransactionResult.COMMIT, producerId, epoch);
+        runUntil(result::isCompleted);
+        runUntil(sendFuture::isDone);
     }
 
     @Test
     public void testEndTxnNotSentIfIncompleteBatches() {
-        long pid = 13131L;
-        short epoch = 1;
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
         transactionManager.beginTransaction();
 
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
         prepareAddPartitionsToTxn(tp0, Errors.NONE);
-        sender.run(time.milliseconds());
-        assertTrue(transactionManager.isPartitionAdded(tp0));
+        runUntil(() -> transactionManager.isPartitionAdded(tp0));
 
         transactionManager.beginCommit();
-        assertNull(transactionManager.nextRequestHandler(true));
-        assertTrue(transactionManager.nextRequestHandler(false).isEndTxn());
+        assertNull(transactionManager.nextRequest(true));
+        assertTrue(transactionManager.nextRequest(false).isEndTxn());
     }
 
     @Test(expected = IllegalStateException.class)
@@ -156,30 +220,26 @@ public void testFailIfNotReadyForSendNoProducerId() {
 
     @Test
     public void testFailIfNotReadyForSendIdempotentProducer() {
-        TransactionManager idempotentTransactionManager = new TransactionManager();
-        idempotentTransactionManager.failIfNotReadyForSend();
+        initializeTransactionManager(Optional.empty(), false);
+        transactionManager.failIfNotReadyForSend();
     }
 
     @Test(expected = KafkaException.class)
     public void testFailIfNotReadyForSendIdempotentProducerFatalError() {
-        TransactionManager idempotentTransactionManager = new TransactionManager();
-        idempotentTransactionManager.transitionToFatalError(new KafkaException());
-        idempotentTransactionManager.failIfNotReadyForSend();
+        initializeTransactionManager(Optional.empty(), false);
+        transactionManager.transitionToFatalError(new KafkaException());
+        transactionManager.failIfNotReadyForSend();
     }
 
     @Test(expected = IllegalStateException.class)
     public void testFailIfNotReadyForSendNoOngoingTransaction() {
-        long pid = 13131L;
-        short epoch = 1;
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
         transactionManager.failIfNotReadyForSend();
     }
 
     @Test(expected = KafkaException.class)
     public void testFailIfNotReadyForSendAfterAbortableError() {
-        long pid = 13131L;
-        short epoch = 1;
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
         transactionManager.beginTransaction();
         transactionManager.transitionToAbortableError(new KafkaException());
         transactionManager.failIfNotReadyForSend();
@@ -187,85 +247,78 @@ public void testFailIfNotReadyForSendAfterAbortableError() {
 
     @Test(expected = KafkaException.class)
     public void testFailIfNotReadyForSendAfterFatalError() {
-        long pid = 13131L;
-        short epoch = 1;
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
         transactionManager.transitionToFatalError(new KafkaException());
         transactionManager.failIfNotReadyForSend();
     }
 
     @Test
     public void testHasOngoingTransactionSuccessfulAbort() {
-        long pid = 13131L;
-        short epoch = 1;
         TopicPartition partition = new TopicPartition("foo", 0);
 
         assertFalse(transactionManager.hasOngoingTransaction());
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
         assertFalse(transactionManager.hasOngoingTransaction());
 
         transactionManager.beginTransaction();
         assertTrue(transactionManager.hasOngoingTransaction());
 
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(partition);
-        assertTrue(transactionManager.hasOngoingTransaction());
+        runUntil(transactionManager::hasOngoingTransaction);
 
         prepareAddPartitionsToTxn(partition, Errors.NONE);
-        sender.run(time.milliseconds());
+        runUntil(() -> transactionManager.isPartitionAdded(partition));
 
         transactionManager.beginAbort();
         assertTrue(transactionManager.hasOngoingTransaction());
 
-        prepareEndTxnResponse(Errors.NONE, TransactionResult.ABORT, pid, epoch);
-        sender.run(time.milliseconds());
-        assertFalse(transactionManager.hasOngoingTransaction());
+        prepareEndTxnResponse(Errors.NONE, TransactionResult.ABORT, producerId, epoch);
+        runUntil(() -> !transactionManager.hasOngoingTransaction());
     }
 
     @Test
     public void testHasOngoingTransactionSuccessfulCommit() {
-        long pid = 13131L;
-        short epoch = 1;
         TopicPartition partition = new TopicPartition("foo", 0);
 
         assertFalse(transactionManager.hasOngoingTransaction());
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
         assertFalse(transactionManager.hasOngoingTransaction());
 
         transactionManager.beginTransaction();
         assertTrue(transactionManager.hasOngoingTransaction());
 
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(partition);
         assertTrue(transactionManager.hasOngoingTransaction());
 
         prepareAddPartitionsToTxn(partition, Errors.NONE);
-        sender.run(time.milliseconds());
+        runUntil(() -> transactionManager.isPartitionAdded(partition));
 
         transactionManager.beginCommit();
         assertTrue(transactionManager.hasOngoingTransaction());
 
-        prepareEndTxnResponse(Errors.NONE, TransactionResult.COMMIT, pid, epoch);
-        sender.run(time.milliseconds());
-        assertFalse(transactionManager.hasOngoingTransaction());
+        prepareEndTxnResponse(Errors.NONE, TransactionResult.COMMIT, producerId, epoch);
+        runUntil(() -> !transactionManager.hasOngoingTransaction());
     }
 
     @Test
     public void testHasOngoingTransactionAbortableError() {
-        long pid = 13131L;
-        short epoch = 1;
         TopicPartition partition = new TopicPartition("foo", 0);
 
         assertFalse(transactionManager.hasOngoingTransaction());
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
         assertFalse(transactionManager.hasOngoingTransaction());
 
         transactionManager.beginTransaction();
         assertTrue(transactionManager.hasOngoingTransaction());
 
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(partition);
         assertTrue(transactionManager.hasOngoingTransaction());
 
         prepareAddPartitionsToTxn(partition, Errors.NONE);
-        sender.run(time.milliseconds());
+        runUntil(() -> transactionManager.isPartitionAdded(partition));
 
         transactionManager.transitionToAbortableError(new KafkaException());
         assertTrue(transactionManager.hasOngoingTransaction());
@@ -273,29 +326,27 @@ public void testHasOngoingTransactionAbortableError() {
         transactionManager.beginAbort();
         assertTrue(transactionManager.hasOngoingTransaction());
 
-        prepareEndTxnResponse(Errors.NONE, TransactionResult.ABORT, pid, epoch);
-        sender.run(time.milliseconds());
-        assertFalse(transactionManager.hasOngoingTransaction());
+        prepareEndTxnResponse(Errors.NONE, TransactionResult.ABORT, producerId, epoch);
+        runUntil(() -> !transactionManager.hasOngoingTransaction());
     }
 
     @Test
     public void testHasOngoingTransactionFatalError() {
-        long pid = 13131L;
-        short epoch = 1;
         TopicPartition partition = new TopicPartition("foo", 0);
 
         assertFalse(transactionManager.hasOngoingTransaction());
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
         assertFalse(transactionManager.hasOngoingTransaction());
 
         transactionManager.beginTransaction();
         assertTrue(transactionManager.hasOngoingTransaction());
 
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(partition);
         assertTrue(transactionManager.hasOngoingTransaction());
 
         prepareAddPartitionsToTxn(partition, Errors.NONE);
-        sender.run(time.milliseconds());
+        runUntil(() -> transactionManager.isPartitionAdded(partition));
 
         transactionManager.transitionToFatalError(new KafkaException());
         assertFalse(transactionManager.hasOngoingTransaction());
@@ -303,25 +354,25 @@ public void testHasOngoingTransactionFatalError() {
 
     @Test
     public void testMaybeAddPartitionToTransaction() {
-        long pid = 13131L;
-        short epoch = 1;
         TopicPartition partition = new TopicPartition("foo", 0);
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
         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());
+        assertTrue(transactionManager.hasPartitionsToAdd());
 
+        runUntil(() -> transactionManager.isPartitionAdded(partition));
         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));
@@ -330,112 +381,105 @@ public void testMaybeAddPartitionToTransaction() {
 
     @Test
     public void testAddPartitionToTransactionOverridesRetryBackoffForConcurrentTransactions() {
-        long pid = 13131L;
-        short epoch = 1;
         TopicPartition partition = new TopicPartition("foo", 0);
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
         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());
+        runUntil(() -> !client.hasPendingResponses());
 
-        TransactionManager.TxnRequestHandler handler = transactionManager.nextRequestHandler(false);
+        TransactionManager.TxnRequestHandler handler = transactionManager.nextRequest(false);
         assertNotNull(handler);
         assertEquals(20, handler.retryBackoffMs());
     }
 
     @Test
     public void testAddPartitionToTransactionRetainsRetryBackoffForRegularRetriableError() {
-        long pid = 13131L;
-        short epoch = 1;
         TopicPartition partition = new TopicPartition("foo", 0);
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
         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());
+        runUntil(() -> !client.hasPendingResponses());
 
-        TransactionManager.TxnRequestHandler handler = transactionManager.nextRequestHandler(false);
+        TransactionManager.TxnRequestHandler handler = transactionManager.nextRequest(false);
         assertNotNull(handler);
         assertEquals(DEFAULT_RETRY_BACKOFF_MS, handler.retryBackoffMs());
     }
 
     @Test
     public void testAddPartitionToTransactionRetainsRetryBackoffWhenPartitionsAlreadyAdded() {
-        long pid = 13131L;
-        short epoch = 1;
         TopicPartition partition = new TopicPartition("foo", 0);
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
         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());
-        assertTrue(transactionManager.isPartitionAdded(partition));
+        runUntil(() -> transactionManager.isPartitionAdded(partition));
 
         TopicPartition otherPartition = new TopicPartition("foo", 1);
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(otherPartition);
 
         prepareAddPartitionsToTxn(otherPartition, Errors.CONCURRENT_TRANSACTIONS);
-        TransactionManager.TxnRequestHandler handler = transactionManager.nextRequestHandler(false);
+        TransactionManager.TxnRequestHandler handler = transactionManager.nextRequest(false);
         assertNotNull(handler);
         assertEquals(DEFAULT_RETRY_BACKOFF_MS, handler.retryBackoffMs());
     }
 
     @Test(expected = IllegalStateException.class)
     public void testMaybeAddPartitionToTransactionBeforeInitTransactions() {
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(new TopicPartition("foo", 0));
     }
 
     @Test(expected = IllegalStateException.class)
     public void testMaybeAddPartitionToTransactionBeforeBeginTransaction() {
-        long pid = 13131L;
-        short epoch = 1;
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(new TopicPartition("foo", 0));
     }
 
     @Test(expected = KafkaException.class)
     public void testMaybeAddPartitionToTransactionAfterAbortableError() {
-        long pid = 13131L;
-        short epoch = 1;
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
         transactionManager.beginTransaction();
         transactionManager.transitionToAbortableError(new KafkaException());
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(new TopicPartition("foo", 0));
     }
 
     @Test(expected = KafkaException.class)
     public void testMaybeAddPartitionToTransactionAfterFatalError() {
-        long pid = 13131L;
-        short epoch = 1;
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
         transactionManager.transitionToFatalError(new KafkaException());
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(new TopicPartition("foo", 0));
     }
 
     @Test
     public void testIsSendToPartitionAllowedWithPendingPartitionAfterAbortableError() {
-        final long pid = 13131L;
-        final short epoch = 1;
-
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
         transactionManager.transitionToAbortableError(new KafkaException());
 
@@ -445,16 +489,14 @@ public void testIsSendToPartitionAllowedWithPendingPartitionAfterAbortableError(
 
     @Test
     public void testIsSendToPartitionAllowedWithInFlightPartitionAddAfterAbortableError() {
-        final long pid = 13131L;
-        final short epoch = 1;
-
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
 
         // Send the AddPartitionsToTxn request and leave it in-flight
-        sender.run(time.milliseconds());
+        runUntil(transactionManager::hasInFlightRequest);
         transactionManager.transitionToAbortableError(new KafkaException());
 
         assertFalse(transactionManager.isSendToPartitionAllowed(tp0));
@@ -463,12 +505,10 @@ public void testIsSendToPartitionAllowedWithInFlightPartitionAddAfterAbortableEr
 
     @Test
     public void testIsSendToPartitionAllowedWithPendingPartitionAfterFatalError() {
-        final long pid = 13131L;
-        final short epoch = 1;
-
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
         transactionManager.transitionToFatalError(new KafkaException());
 
@@ -478,16 +518,14 @@ public void testIsSendToPartitionAllowedWithPendingPartitionAfterFatalError() {
 
     @Test
     public void testIsSendToPartitionAllowedWithInFlightPartitionAddAfterFatalError() {
-        final long pid = 13131L;
-        final short epoch = 1;
-
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
 
         // Send the AddPartitionsToTxn request and leave it in-flight
-        sender.run(time.milliseconds());
+        runUntil(transactionManager::hasInFlightRequest);
         transactionManager.transitionToFatalError(new KafkaException());
 
         assertFalse(transactionManager.isSendToPartitionAllowed(tp0));
@@ -496,17 +534,15 @@ public void testIsSendToPartitionAllowedWithInFlightPartitionAddAfterFatalError(
 
     @Test
     public void testIsSendToPartitionAllowedWithAddedPartitionAfterAbortableError() {
-        final long pid = 13131L;
-        final short epoch = 1;
-
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
 
         transactionManager.beginTransaction();
 
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
-        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, pid);
-        sender.run(time.milliseconds());
-        assertFalse(transactionManager.hasPartitionsToAdd());
+        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, producerId);
+
+        runUntil(() -> !transactionManager.hasPartitionsToAdd());
         transactionManager.transitionToAbortableError(new KafkaException());
 
         assertTrue(transactionManager.isSendToPartitionAllowed(tp0));
@@ -515,16 +551,14 @@ public void testIsSendToPartitionAllowedWithAddedPartitionAfterAbortableError()
 
     @Test
     public void testIsSendToPartitionAllowedWithAddedPartitionAfterFatalError() {
-        final long pid = 13131L;
-        final short epoch = 1;
-
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
-        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, pid);
-        sender.run(time.milliseconds());
-        assertFalse(transactionManager.hasPartitionsToAdd());
+        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, producerId);
+
+        runUntil(() -> !transactionManager.hasPartitionsToAdd());
         transactionManager.transitionToFatalError(new KafkaException());
 
         assertFalse(transactionManager.isSendToPartitionAllowed(tp0));
@@ -533,100 +567,230 @@ public void testIsSendToPartitionAllowedWithAddedPartitionAfterFatalError() {
 
     @Test
     public void testIsSendToPartitionAllowedWithPartitionNotAdded() {
-        final long pid = 13131L;
-        final short epoch = 1;
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
         transactionManager.beginTransaction();
         assertFalse(transactionManager.isSendToPartitionAllowed(tp0));
     }
 
-    @Test(expected = IllegalStateException.class)
-    public void testInvalidSequenceIncrement() {
-        TransactionManager transactionManager = new TransactionManager();
-        transactionManager.incrementSequenceNumber(tp0, 3333);
-    }
-
     @Test
     public void testDefaultSequenceNumber() {
-        TransactionManager transactionManager = new TransactionManager();
+        initializeTransactionManager(Optional.empty(), false);
         assertEquals((int) transactionManager.sequenceNumber(tp0), 0);
         transactionManager.incrementSequenceNumber(tp0, 3);
         assertEquals((int) transactionManager.sequenceNumber(tp0), 3);
     }
 
+    @Test
+    public void testBumpEpochAndResetSequenceNumbersAfterUnknownProducerId() {
+        initializeTransactionManager(Optional.empty(), false);
+        initializeIdempotentProducerId(producerId, epoch);
+
+        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);
+
+        // We get an UNKNOWN_PRODUCER_ID, so bump the epoch and set sequence numbers back to 0
+        ProduceResponse.PartitionResponse b2Response = new ProduceResponse.PartitionResponse(
+                Errors.UNKNOWN_PRODUCER_ID, -1, -1, 500L);
+        assertTrue(transactionManager.canRetry(b2Response, b2));
+
+        // Run sender loop to trigger epoch bump
+        runUntil(() -> transactionManager.producerIdAndEpoch().epoch == 2);
+        assertEquals(2, b2.producerEpoch());
+        assertEquals(0, b2.baseSequence());
+        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 partition(s) that triggered the reset will have their sequence number reset, while any others will not
+        final short epoch = Short.MAX_VALUE;
+
+        initializeTransactionManager(Optional.empty(), false);
+        initializeIdempotentProducerId(producerId, epoch);
+
+        ProducerBatch tp0b1 = writeIdempotentBatchWithValue(transactionManager, tp0, "1");
+        ProducerBatch tp1b1 = writeIdempotentBatchWithValue(transactionManager, tp1, "1");
+
+        ProduceResponse.PartitionResponse tp0b1Response = new ProduceResponse.PartitionResponse(
+                Errors.NONE, -1, -1, 400L);
+        transactionManager.handleCompletedBatch(tp0b1, tp0b1Response);
+
+        ProduceResponse.PartitionResponse tp1b1Response = new ProduceResponse.PartitionResponse(
+                Errors.NONE, -1, -1, 400L);
+        transactionManager.handleCompletedBatch(tp1b1, tp1b1Response);
+
+        ProducerBatch tp0b2 = writeIdempotentBatchWithValue(transactionManager, tp0, "2");
+        ProducerBatch tp1b2 = writeIdempotentBatchWithValue(transactionManager, tp1, "2");
+        assertEquals(2, transactionManager.sequenceNumber(tp0).intValue());
+        assertEquals(2, transactionManager.sequenceNumber(tp1).intValue());
+
+        ProduceResponse.PartitionResponse b1Response = new ProduceResponse.PartitionResponse(
+                Errors.UNKNOWN_PRODUCER_ID, -1, -1, 400L);
+        assertTrue(transactionManager.canRetry(b1Response, tp0b1));
+
+        ProduceResponse.PartitionResponse b2Response = new ProduceResponse.PartitionResponse(
+                Errors.NONE, -1, -1, 400L);
+        transactionManager.handleCompletedBatch(tp1b1, b2Response);
+
+        transactionManager.bumpIdempotentEpochAndResetIdIfNeeded();
+
+        assertEquals(1, transactionManager.sequenceNumber(tp0).intValue());
+        assertEquals(tp0b2, transactionManager.nextBatchBySequence(tp0));
+        assertEquals(2, transactionManager.sequenceNumber(tp1).intValue());
+        assertEquals(tp1b2, transactionManager.nextBatchBySequence(tp1));
+    }
+
+    @Test
+    public void testBatchCompletedAfterProducerReset() {
+        final short epoch = Short.MAX_VALUE;
+
+        initializeTransactionManager(Optional.empty(), false);
+        initializeIdempotentProducerId(producerId, epoch);
+
+        ProducerBatch b1 = writeIdempotentBatchWithValue(transactionManager, tp0, "1");
+        writeIdempotentBatchWithValue(transactionManager, tp1, "1");
+
+        ProducerBatch b2 = writeIdempotentBatchWithValue(transactionManager, tp0, "2");
+        assertEquals(2, transactionManager.sequenceNumber(tp0).intValue());
+
+        // The producerId might be reset due to a failure on another partition
+        transactionManager.requestEpochBumpForPartition(tp1);
+        transactionManager.bumpIdempotentEpochAndResetIdIfNeeded();
+        initializeIdempotentProducerId(producerId + 1, (short) 0);
+
+        // We continue to track the state of tp0 until in-flight requests complete
+        ProduceResponse.PartitionResponse b1Response = new ProduceResponse.PartitionResponse(
+                Errors.NONE, 500L, time.milliseconds(), 0L);
+        transactionManager.handleCompletedBatch(b1, b1Response);
+
+        assertEquals(2, transactionManager.sequenceNumber(tp0).intValue());
+        assertEquals(0, transactionManager.lastAckedSequence(tp0).getAsInt());
+        assertEquals(b2, transactionManager.nextBatchBySequence(tp0));
+        assertEquals(epoch, transactionManager.nextBatchBySequence(tp0).producerEpoch());
+
+        ProduceResponse.PartitionResponse b2Response = new ProduceResponse.PartitionResponse(
+                Errors.NONE, 500L, time.milliseconds(), 0L);
+        transactionManager.handleCompletedBatch(b2, b2Response);
+
+        assertEquals(0, transactionManager.sequenceNumber(tp0).intValue());
+        assertFalse(transactionManager.lastAckedSequence(tp0).isPresent());
+        assertNull(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() {
+        initializeTransactionManager(Optional.empty(), false);
+        assertEquals((int) transactionManager.sequenceNumber(tp0), 0);
+        transactionManager.incrementSequenceNumber(tp0, Integer.MAX_VALUE);
+        assertEquals((int) transactionManager.sequenceNumber(tp0), Integer.MAX_VALUE);
+        transactionManager.incrementSequenceNumber(tp0, 100);
+        assertEquals((int) transactionManager.sequenceNumber(tp0), 99);
+        transactionManager.incrementSequenceNumber(tp0, Integer.MAX_VALUE);
+        assertEquals((int) transactionManager.sequenceNumber(tp0), 98);
+    }
+
     @Test
     public void testProducerIdReset() {
-        TransactionManager transactionManager = new TransactionManager();
+        initializeTransactionManager(Optional.empty(), false);
+        initializeIdempotentProducerId(15L, Short.MAX_VALUE);
         assertEquals((int) transactionManager.sequenceNumber(tp0), 0);
+        assertEquals((int) transactionManager.sequenceNumber(tp1), 0);
         transactionManager.incrementSequenceNumber(tp0, 3);
         assertEquals((int) transactionManager.sequenceNumber(tp0), 3);
-        transactionManager.resetProducerId();
+        transactionManager.incrementSequenceNumber(tp1, 3);
+        assertEquals((int) transactionManager.sequenceNumber(tp1), 3);
+
+        transactionManager.requestEpochBumpForPartition(tp0);
+        transactionManager.bumpIdempotentEpochAndResetIdIfNeeded();
         assertEquals((int) transactionManager.sequenceNumber(tp0), 0);
+        assertEquals((int) transactionManager.sequenceNumber(tp1), 3);
     }
 
     @Test
     public void testBasicTransaction() throws InterruptedException {
-        final long pid = 13131L;
-        final short epoch = 1;
-
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
 
         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;
+        Future responseFuture = appendToAccumulator(tp0);
 
         assertFalse(responseFuture.isDone());
-        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, pid);
+        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, producerId);
 
-        prepareProduceResponse(Errors.NONE, pid, epoch);
+        prepareProduceResponse(Errors.NONE, producerId, epoch);
         assertFalse(transactionManager.transactionContainsPartition(tp0));
         assertFalse(transactionManager.isSendToPartitionAllowed(tp0));
-        sender.run(time.milliseconds());  // send addPartitions.
-        // Check that only addPartitions was sent.
-        assertTrue(transactionManager.transactionContainsPartition(tp0));
+        runUntil(() -> transactionManager.transactionContainsPartition(tp0));
         assertTrue(transactionManager.isSendToPartitionAllowed(tp0));
         assertFalse(responseFuture.isDone());
-
-        sender.run(time.milliseconds());  // send produce request.
-        assertTrue(responseFuture.isDone());
+        runUntil(responseFuture::isDone);
 
         Map offsets = new HashMap<>();
         offsets.put(tp1, new OffsetAndMetadata(1));
-        final String consumerGroupId = "myconsumergroup";
-        TransactionalRequestResult addOffsetsResult = transactionManager.sendOffsetsToTransaction(offsets, consumerGroupId);
+
+        TransactionalRequestResult addOffsetsResult = transactionManager.sendOffsetsToTransaction(
+            offsets, new ConsumerGroupMetadata(consumerGroupId));
 
         assertFalse(transactionManager.hasPendingOffsetCommits());
 
-        prepareAddOffsetsToTxnResponse(Errors.NONE, consumerGroupId, pid, epoch);
+        prepareAddOffsetsToTxnResponse(Errors.NONE, consumerGroupId, producerId, epoch);
 
-        sender.run(time.milliseconds());  // Send AddOffsetsRequest
-        assertTrue(transactionManager.hasPendingOffsetCommits());  // We should now have created and queued the offset commit request.
+        runUntil(transactionManager::hasPendingOffsetCommits);
         assertFalse(addOffsetsResult.isCompleted()); // the result doesn't complete until TxnOffsetCommit returns
 
         Map txnOffsetCommitResponse = new HashMap<>();
         txnOffsetCommitResponse.put(tp1, Errors.NONE);
 
         prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.GROUP, consumerGroupId);
-        prepareTxnOffsetCommitResponse(consumerGroupId, pid, epoch, txnOffsetCommitResponse);
+        prepareTxnOffsetCommitResponse(consumerGroupId, producerId, epoch, txnOffsetCommitResponse);
 
-        assertEquals(null, 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
-        assertNotNull(transactionManager.coordinator(CoordinatorType.GROUP));
+        assertNull(transactionManager.coordinator(CoordinatorType.GROUP));
+        runUntil(() -> transactionManager.coordinator(CoordinatorType.GROUP) != null);
         assertTrue(transactionManager.hasPendingOffsetCommits());
 
-        sender.run(time.milliseconds());  // send TxnOffsetCommitRequest commit.
-
-        assertFalse(transactionManager.hasPendingOffsetCommits());
+        runUntil(() -> !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.
-
-        assertFalse(transactionManager.hasOngoingTransaction());
+        prepareEndTxnResponse(Errors.NONE, TransactionResult.COMMIT, producerId, epoch);
+        runUntil(() -> !transactionManager.hasOngoingTransaction());
         assertFalse(transactionManager.isCompleting());
         assertFalse(transactionManager.transactionContainsPartition(tp0));
     }
@@ -637,29 +801,31 @@ 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.
+        runUntil(() -> transactionManager.coordinator(CoordinatorType.TRANSACTION) == null);
 
         prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.TRANSACTION, transactionalId);
-        sender.run(time.milliseconds());  // find coordinator
-        sender.run(time.milliseconds());
+        runUntil(() -> transactionManager.coordinator(CoordinatorType.TRANSACTION) != null);
         assertEquals(brokerNode, transactionManager.coordinator(CoordinatorType.TRANSACTION));
     }
 
+    @Test
+    public void testInitializeTransactionsTwiceRaisesError() {
+        doInitTransactions(producerId, epoch);
+        assertTrue(transactionManager.hasProducerId());
+        assertThrows(KafkaException.class, () -> transactionManager.initializeTransactions());
+    }
+
     @Test
     public void testUnsupportedFindCoordinator() {
         transactionManager.initializeTransactions();
-        client.prepareUnsupportedVersionResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                FindCoordinatorRequest findCoordinatorRequest = (FindCoordinatorRequest) body;
-                assertEquals(findCoordinatorRequest.coordinatorType(), CoordinatorType.TRANSACTION);
-                assertEquals(findCoordinatorRequest.coordinatorKey(), transactionalId);
-                return true;
-            }
+        client.prepareUnsupportedVersionResponse(body -> {
+            FindCoordinatorRequest findCoordinatorRequest = (FindCoordinatorRequest) body;
+            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
+        runUntil(transactionManager::hasFatalError);
         assertTrue(transactionManager.hasFatalError());
         assertTrue(transactionManager.lastError() instanceof UnsupportedVersionException);
     }
@@ -668,51 +834,36 @@ public boolean matches(AbstractRequest body) {
     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
-
+        runUntil(() -> transactionManager.coordinator(CoordinatorType.TRANSACTION) != null);
         assertFalse(transactionManager.hasError());
-        assertNotNull(transactionManager.coordinator(CoordinatorType.TRANSACTION));
-
-        client.prepareUnsupportedVersionResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                InitProducerIdRequest initProducerIdRequest = (InitProducerIdRequest) body;
-                assertEquals(initProducerIdRequest.transactionalId(), transactionalId);
-                assertEquals(initProducerIdRequest.transactionTimeoutMs(), transactionTimeoutMs);
-                return true;
-            }
+
+        client.prepareUnsupportedVersionResponse(body -> {
+            InitProducerIdRequest initProducerIdRequest = (InitProducerIdRequest) body;
+            assertEquals(initProducerIdRequest.data().transactionalId(), transactionalId);
+            assertEquals(initProducerIdRequest.data().transactionTimeoutMs(), transactionTimeoutMs);
+            return true;
         });
 
-        sender.run(time.milliseconds()); // InitProducerRequest is dequeued
+        runUntil(transactionManager::hasFatalError);
         assertTrue(transactionManager.hasFatalError());
         assertTrue(transactionManager.lastError() instanceof UnsupportedVersionException);
     }
 
     @Test
     public void testUnsupportedForMessageFormatInTxnOffsetCommit() {
-        final String consumerGroupId = "consumer";
-        final long pid = 13131L;
-        final short epoch = 1;
         final TopicPartition tp = new TopicPartition("foo", 0);
 
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
 
         transactionManager.beginTransaction();
         TransactionalRequestResult sendOffsetsResult = transactionManager.sendOffsetsToTransaction(
-                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
+                singletonMap(tp, new OffsetAndMetadata(39L)), new ConsumerGroupMetadata(consumerGroupId));
 
+        prepareAddOffsetsToTxnResponse(Errors.NONE, consumerGroupId, producerId, epoch);
         prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.GROUP, consumerGroupId);
-        sender.run(time.milliseconds());  // FindCoordinator Returned
-
-        prepareTxnOffsetCommitResponse(consumerGroupId, pid, epoch, singletonMap(tp, Errors.UNSUPPORTED_FOR_MESSAGE_FORMAT));
-        sender.run(time.milliseconds());  // TxnOffsetCommit Handled
+        prepareTxnOffsetCommitResponse(consumerGroupId, producerId, epoch, singletonMap(tp, Errors.UNSUPPORTED_FOR_MESSAGE_FORMAT));
+        runUntil(transactionManager::hasError);
 
-        assertTrue(transactionManager.hasError());
         assertTrue(transactionManager.lastError() instanceof UnsupportedForMessageFormatException);
         assertTrue(sendOffsetsResult.isCompleted());
         assertFalse(sendOffsetsResult.isSuccessful());
@@ -720,36 +871,137 @@ public void testUnsupportedForMessageFormatInTxnOffsetCommit() {
         assertFatalError(UnsupportedForMessageFormatException.class);
     }
 
+    @Test
+    public void testFencedInstanceIdInTxnOffsetCommitByGroupMetadata() {
+        final TopicPartition tp = new TopicPartition("foo", 0);
+        final String fencedMemberId = "fenced_member";
+
+        doInitTransactions();
+
+        transactionManager.beginTransaction();
+
+        TransactionalRequestResult sendOffsetsResult = transactionManager.sendOffsetsToTransaction(
+            singletonMap(tp, new OffsetAndMetadata(39L)),
+            new ConsumerGroupMetadata(consumerGroupId, 5, fencedMemberId, Optional.of(groupInstanceId)));
+
+        prepareAddOffsetsToTxnResponse(Errors.NONE, consumerGroupId, producerId, epoch);
+        prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.GROUP, consumerGroupId);
+        runUntil(() -> transactionManager.coordinator(CoordinatorType.GROUP) != null);
+
+        client.prepareResponse(request -> {
+            TxnOffsetCommitRequest txnOffsetCommitRequest = (TxnOffsetCommitRequest) request;
+            assertEquals(consumerGroupId, txnOffsetCommitRequest.data().groupId());
+            assertEquals(producerId, txnOffsetCommitRequest.data().producerId());
+            assertEquals(epoch, txnOffsetCommitRequest.data().producerEpoch());
+            return txnOffsetCommitRequest.data().groupInstanceId().equals(groupInstanceId)
+                && !txnOffsetCommitRequest.data().memberId().equals(memberId);
+        }, new TxnOffsetCommitResponse(0, singletonMap(tp, Errors.FENCED_INSTANCE_ID)));
+
+        runUntil(transactionManager::hasError);
+        assertTrue(transactionManager.lastError() instanceof FencedInstanceIdException);
+        assertTrue(sendOffsetsResult.isCompleted());
+        assertFalse(sendOffsetsResult.isSuccessful());
+        assertTrue(sendOffsetsResult.error() instanceof FencedInstanceIdException);
+        assertAbortableError(FencedInstanceIdException.class);
+    }
+
+    @Test
+    public void testUnknownMemberIdInTxnOffsetCommitByGroupMetadata() {
+        final TopicPartition tp = new TopicPartition("foo", 0);
+        final String unknownMemberId = "unknownMember";
+
+        doInitTransactions();
+
+        transactionManager.beginTransaction();
+
+        TransactionalRequestResult sendOffsetsResult = transactionManager.sendOffsetsToTransaction(
+            singletonMap(tp, new OffsetAndMetadata(39L)),
+            new ConsumerGroupMetadata(consumerGroupId, 5, unknownMemberId, Optional.empty()));
+
+        prepareAddOffsetsToTxnResponse(Errors.NONE, consumerGroupId, producerId, epoch);
+        prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.GROUP, consumerGroupId);
+        runUntil(() -> transactionManager.coordinator(CoordinatorType.GROUP) != null);
+
+        client.prepareResponse(request -> {
+            TxnOffsetCommitRequest txnOffsetCommitRequest = (TxnOffsetCommitRequest) request;
+            assertEquals(consumerGroupId, txnOffsetCommitRequest.data().groupId());
+            assertEquals(producerId, txnOffsetCommitRequest.data().producerId());
+            assertEquals(epoch, txnOffsetCommitRequest.data().producerEpoch());
+            return !txnOffsetCommitRequest.data().memberId().equals(memberId);
+        }, new TxnOffsetCommitResponse(0, singletonMap(tp, Errors.UNKNOWN_MEMBER_ID)));
+
+        runUntil(transactionManager::hasError);
+        assertTrue(transactionManager.lastError() instanceof CommitFailedException);
+        assertTrue(sendOffsetsResult.isCompleted());
+        assertFalse(sendOffsetsResult.isSuccessful());
+        assertTrue(sendOffsetsResult.error() instanceof CommitFailedException);
+        assertAbortableError(CommitFailedException.class);
+    }
+
+    @Test
+    public void testIllegalGenerationInTxnOffsetCommitByGroupMetadata() {
+        final TopicPartition tp = new TopicPartition("foo", 0);
+        final int illegalGenerationId = 1;
+
+        doInitTransactions();
+
+        transactionManager.beginTransaction();
+
+        TransactionalRequestResult sendOffsetsResult = transactionManager.sendOffsetsToTransaction(
+            singletonMap(tp, new OffsetAndMetadata(39L)),
+            new ConsumerGroupMetadata(consumerGroupId, illegalGenerationId, JoinGroupRequest.UNKNOWN_MEMBER_ID,
+                Optional.empty()));
+
+        prepareAddOffsetsToTxnResponse(Errors.NONE, consumerGroupId, producerId, epoch);
+        prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.GROUP, consumerGroupId);
+        runUntil(() -> transactionManager.coordinator(CoordinatorType.GROUP) != null);
+
+        prepareTxnOffsetCommitResponse(consumerGroupId, producerId, epoch, singletonMap(tp, Errors.ILLEGAL_GENERATION));
+        client.prepareResponse(request -> {
+            TxnOffsetCommitRequest txnOffsetCommitRequest = (TxnOffsetCommitRequest) request;
+            assertEquals(consumerGroupId, txnOffsetCommitRequest.data().groupId());
+            assertEquals(producerId, txnOffsetCommitRequest.data().producerId());
+            assertEquals(epoch, txnOffsetCommitRequest.data().producerEpoch());
+            return txnOffsetCommitRequest.data().generationId() != generationId;
+        }, new TxnOffsetCommitResponse(0, singletonMap(tp, Errors.ILLEGAL_GENERATION)));
+
+        runUntil(transactionManager::hasError);
+        assertTrue(transactionManager.lastError() instanceof CommitFailedException);
+        assertTrue(sendOffsetsResult.isCompleted());
+        assertFalse(sendOffsetsResult.isSuccessful());
+        assertTrue(sendOffsetsResult.error() instanceof CommitFailedException);
+        assertAbortableError(CommitFailedException.class);
+    }
+
     @Test
     public void testLookupCoordinatorOnDisconnectAfterSend() {
         // This is called from the initTransactions method in the producer as the first order of business.
         // It finds the coordinator and then gets a PID.
-        final long pid = 13131L;
-        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());
+        runUntil(() -> transactionManager.coordinator(CoordinatorType.TRANSACTION) != null);
         assertEquals(brokerNode, transactionManager.coordinator(CoordinatorType.TRANSACTION));
-        prepareInitPidResponse(Errors.NONE, true, pid, epoch);
+
+        prepareInitPidResponse(Errors.NONE, true, producerId, epoch);
         // send pid to coordinator, should get disconnected before receiving the response, and resend the
         // FindCoordinator and InitPid requests.
-        sender.run(time.milliseconds());
+        runUntil(() -> transactionManager.coordinator(CoordinatorType.TRANSACTION) == null);
 
-        assertEquals(null, transactionManager.coordinator(CoordinatorType.TRANSACTION));
+        assertNull(transactionManager.coordinator(CoordinatorType.TRANSACTION));
         assertFalse(initPidResult.isCompleted());
         assertFalse(transactionManager.hasProducerId());
 
         prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.TRANSACTION, transactionalId);
-        sender.run(time.milliseconds());
+        runUntil(() -> transactionManager.coordinator(CoordinatorType.TRANSACTION) != null);
+
         assertEquals(brokerNode, transactionManager.coordinator(CoordinatorType.TRANSACTION));
         assertFalse(initPidResult.isCompleted());
-        prepareInitPidResponse(Errors.NONE, false, pid, epoch);
-        sender.run(time.milliseconds());  // get pid and epoch
+        prepareInitPidResponse(Errors.NONE, false, producerId, epoch);
+        runUntil(initPidResult::isCompleted);
 
         assertTrue(initPidResult.isCompleted()); // The future should only return after the second round of retries succeed.
         assertTrue(transactionManager.hasProducerId());
-        assertEquals(pid, transactionManager.producerIdAndEpoch().producerId);
+        assertEquals(producerId, transactionManager.producerIdAndEpoch().producerId);
         assertEquals(epoch, transactionManager.producerIdAndEpoch().epoch);
     }
 
@@ -757,35 +1009,30 @@ public void testLookupCoordinatorOnDisconnectAfterSend() {
     public void testLookupCoordinatorOnDisconnectBeforeSend() {
         // This is called from the initTransactions method in the producer as the first order of business.
         // It finds the coordinator and then gets a PID.
-        final long pid = 13131L;
-        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.
+        runUntil(() -> transactionManager.coordinator(CoordinatorType.TRANSACTION) != null);
         assertEquals(brokerNode, transactionManager.coordinator(CoordinatorType.TRANSACTION));
 
         client.disconnect(brokerNode.idString());
-        client.blackout(brokerNode, 100);
+        client.backoff(brokerNode, 100);
         // send pid to coordinator. Should get disconnected before the send and resend the FindCoordinator
         // and InitPid requests.
-        sender.run(time.milliseconds());
-        time.sleep(110);  // waiting for the blackout period for the node to expire.
+        runUntil(() -> transactionManager.coordinator(CoordinatorType.TRANSACTION) == null);
+        time.sleep(110);  // waiting for the backoff period for the node to expire.
 
-        assertEquals(null, transactionManager.coordinator(CoordinatorType.TRANSACTION));
         assertFalse(initPidResult.isCompleted());
         assertFalse(transactionManager.hasProducerId());
 
         prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.TRANSACTION, transactionalId);
-        sender.run(time.milliseconds());
+        runUntil(() -> transactionManager.coordinator(CoordinatorType.TRANSACTION) != null);
         assertEquals(brokerNode, transactionManager.coordinator(CoordinatorType.TRANSACTION));
         assertFalse(initPidResult.isCompleted());
-        prepareInitPidResponse(Errors.NONE, false, pid, epoch);
-        sender.run(time.milliseconds());  // get pid and epoch
+        prepareInitPidResponse(Errors.NONE, false, producerId, epoch);
 
-        assertTrue(initPidResult.isCompleted()); // The future should only return after the second round of retries succeed.
+        runUntil(initPidResult::isCompleted);
         assertTrue(transactionManager.hasProducerId());
-        assertEquals(pid, transactionManager.producerIdAndEpoch().producerId);
+        assertEquals(producerId, transactionManager.producerIdAndEpoch().producerId);
         assertEquals(epoch, transactionManager.producerIdAndEpoch().epoch);
     }
 
@@ -793,31 +1040,26 @@ public void testLookupCoordinatorOnDisconnectBeforeSend() {
     public void testLookupCoordinatorOnNotCoordinatorError() {
         // This is called from the initTransactions method in the producer as the first order of business.
         // It finds the coordinator and then gets a PID.
-        final long pid = 13131L;
-        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());
+        runUntil(() -> transactionManager.coordinator(CoordinatorType.TRANSACTION) != null);
         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
+        prepareInitPidResponse(Errors.NOT_COORDINATOR, false, producerId, epoch);
+        runUntil(() -> transactionManager.coordinator(CoordinatorType.TRANSACTION) == null);
 
-        assertEquals(null, transactionManager.coordinator(CoordinatorType.TRANSACTION));
         assertFalse(initPidResult.isCompleted());
         assertFalse(transactionManager.hasProducerId());
 
         prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.TRANSACTION, transactionalId);
-        sender.run(time.milliseconds());
+        runUntil(() -> transactionManager.coordinator(CoordinatorType.TRANSACTION) != null);
         assertEquals(brokerNode, transactionManager.coordinator(CoordinatorType.TRANSACTION));
         assertFalse(initPidResult.isCompleted());
-        prepareInitPidResponse(Errors.NONE, false, pid, epoch);
-        sender.run(time.milliseconds());  // get pid and epoch
+        prepareInitPidResponse(Errors.NONE, false, producerId, epoch);
 
-        assertTrue(initPidResult.isCompleted()); // The future should only return after the second round of retries succeed.
+        runUntil(initPidResult::isCompleted);
         assertTrue(transactionManager.hasProducerId());
-        assertEquals(pid, transactionManager.producerIdAndEpoch().producerId);
+        assertEquals(producerId, transactionManager.producerIdAndEpoch().producerId);
         assertEquals(epoch, transactionManager.producerIdAndEpoch().epoch);
     }
 
@@ -826,33 +1068,25 @@ 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());
 
-        assertTrue(transactionManager.hasError());
-        assertTrue(transactionManager.lastError() instanceof TransactionalIdAuthorizationException);
+        runUntil(transactionManager::hasError);
 
-        sender.run(time.milliseconds()); // one more run to fail the InitProducerId future
-        assertTrue(initPidResult.isCompleted());
+        assertTrue(transactionManager.hasFatalError());
+        assertTrue(transactionManager.lastError() instanceof TransactionalIdAuthorizationException);
         assertFalse(initPidResult.isSuccessful());
         assertTrue(initPidResult.error() instanceof TransactionalIdAuthorizationException);
-
         assertFatalError(TransactionalIdAuthorizationException.class);
     }
 
     @Test
     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());
+        runUntil(() -> transactionManager.coordinator(CoordinatorType.TRANSACTION) != null);
         assertEquals(brokerNode, transactionManager.coordinator(CoordinatorType.TRANSACTION));
 
-        prepareInitPidResponse(Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED, false, pid, RecordBatch.NO_PRODUCER_EPOCH);
-        sender.run(time.milliseconds());
-
-        assertTrue(transactionManager.hasError());
+        prepareInitPidResponse(Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED, false, producerId, RecordBatch.NO_PRODUCER_EPOCH);
+        runUntil(transactionManager::hasError);
         assertTrue(initPidResult.isCompleted());
         assertFalse(initPidResult.isSuccessful());
         assertTrue(initPidResult.error() instanceof TransactionalIdAuthorizationException);
@@ -862,26 +1096,20 @@ public void testTransactionalIdAuthorizationFailureInInitProducerId() {
 
     @Test
     public void testGroupAuthorizationFailureInFindCoordinator() {
-        final String consumerGroupId = "consumer";
-        final long pid = 13131L;
-        final short epoch = 1;
-
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
 
         transactionManager.beginTransaction();
         TransactionalRequestResult sendOffsetsResult = transactionManager.sendOffsetsToTransaction(
-                singletonMap(new TopicPartition("foo", 0), new OffsetAndMetadata(39L)), consumerGroupId);
+                singletonMap(new TopicPartition("foo", 0), new OffsetAndMetadata(39L)), new ConsumerGroupMetadata(consumerGroupId));
 
-        prepareAddOffsetsToTxnResponse(Errors.NONE, consumerGroupId, pid, epoch);
-        sender.run(time.milliseconds());  // AddOffsetsToTxn Handled, TxnOffsetCommit Enqueued
-        sender.run(time.milliseconds());  // FindCoordinator Enqueued
+        prepareAddOffsetsToTxnResponse(Errors.NONE, consumerGroupId, producerId, epoch);
+        runUntil(() -> !transactionManager.hasPartitionsToAdd());
 
         prepareFindCoordinatorResponse(Errors.GROUP_AUTHORIZATION_FAILED, false, CoordinatorType.GROUP, consumerGroupId);
-        sender.run(time.milliseconds());  // FindCoordinator Failed
-        sender.run(time.milliseconds());  // TxnOffsetCommit Aborted
-        assertTrue(transactionManager.hasError());
+        runUntil(transactionManager::hasError);
         assertTrue(transactionManager.lastError() instanceof GroupAuthorizationException);
-        assertTrue(sendOffsetsResult.isCompleted());
+
+        runUntil(sendOffsetsResult::isCompleted);
         assertFalse(sendOffsetsResult.isSuccessful());
         assertTrue(sendOffsetsResult.error() instanceof GroupAuthorizationException);
 
@@ -893,32 +1121,26 @@ public void testGroupAuthorizationFailureInFindCoordinator() {
 
     @Test
     public void testGroupAuthorizationFailureInTxnOffsetCommit() {
-        final String consumerGroupId = "consumer";
-        final long pid = 13131L;
-        final short epoch = 1;
-        final TopicPartition tp = new TopicPartition("foo", 0);
+        final TopicPartition tp1 = new TopicPartition("foo", 0);
 
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
 
         transactionManager.beginTransaction();
         TransactionalRequestResult sendOffsetsResult = transactionManager.sendOffsetsToTransaction(
-                singletonMap(tp, new OffsetAndMetadata(39L)), consumerGroupId);
+                singletonMap(tp1, new OffsetAndMetadata(39L)), new ConsumerGroupMetadata(consumerGroupId));
 
-        prepareAddOffsetsToTxnResponse(Errors.NONE, consumerGroupId, pid, epoch);
-        sender.run(time.milliseconds());  // AddOffsetsToTxn Handled, TxnOffsetCommit Enqueued
-        sender.run(time.milliseconds());  // FindCoordinator Enqueued
+        prepareAddOffsetsToTxnResponse(Errors.NONE, consumerGroupId, producerId, epoch);
+        runUntil(() -> !transactionManager.hasPartitionsToAdd());
 
         prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.GROUP, consumerGroupId);
-        sender.run(time.milliseconds());  // FindCoordinator Returned
+        prepareTxnOffsetCommitResponse(consumerGroupId, producerId, epoch, singletonMap(tp1, Errors.GROUP_AUTHORIZATION_FAILED));
 
-        prepareTxnOffsetCommitResponse(consumerGroupId, pid, epoch, singletonMap(tp, Errors.GROUP_AUTHORIZATION_FAILED));
-        sender.run(time.milliseconds());  // TxnOffsetCommit Handled
-
-        assertTrue(transactionManager.hasError());
+        runUntil(transactionManager::hasError);
         assertTrue(transactionManager.lastError() instanceof GroupAuthorizationException);
         assertTrue(sendOffsetsResult.isCompleted());
         assertFalse(sendOffsetsResult.isSuccessful());
         assertTrue(sendOffsetsResult.error() instanceof GroupAuthorizationException);
+        assertFalse(transactionManager.hasPendingOffsetCommits());
 
         GroupAuthorizationException exception = (GroupAuthorizationException) sendOffsetsResult.error();
         assertEquals(consumerGroupId, exception.groupId());
@@ -928,21 +1150,16 @@ public void testGroupAuthorizationFailureInTxnOffsetCommit() {
 
     @Test
     public void testTransactionalIdAuthorizationFailureInAddOffsetsToTxn() {
-        final String consumerGroupId = "consumer";
-        final long pid = 13131L;
-        final short epoch = 1;
         final TopicPartition tp = new TopicPartition("foo", 0);
 
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
 
         transactionManager.beginTransaction();
         TransactionalRequestResult sendOffsetsResult = transactionManager.sendOffsetsToTransaction(
-                singletonMap(tp, new OffsetAndMetadata(39L)), consumerGroupId);
-
-        prepareAddOffsetsToTxnResponse(Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED, consumerGroupId, pid, epoch);
-        sender.run(time.milliseconds());  // AddOffsetsToTxn Handled
+                singletonMap(tp, new OffsetAndMetadata(39L)), new ConsumerGroupMetadata(consumerGroupId));
 
-        assertTrue(transactionManager.hasError());
+        prepareAddOffsetsToTxnResponse(Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED, consumerGroupId, producerId, epoch);
+        runUntil(transactionManager::hasError);
         assertTrue(transactionManager.lastError() instanceof TransactionalIdAuthorizationException);
         assertTrue(sendOffsetsResult.isCompleted());
         assertFalse(sendOffsetsResult.isSuccessful());
@@ -953,28 +1170,21 @@ public void testTransactionalIdAuthorizationFailureInAddOffsetsToTxn() {
 
     @Test
     public void testTransactionalIdAuthorizationFailureInTxnOffsetCommit() {
-        final String consumerGroupId = "consumer";
-        final long pid = 13131L;
-        final short epoch = 1;
         final TopicPartition tp = new TopicPartition("foo", 0);
 
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
 
         transactionManager.beginTransaction();
         TransactionalRequestResult sendOffsetsResult = transactionManager.sendOffsetsToTransaction(
-                singletonMap(tp, new OffsetAndMetadata(39L)), consumerGroupId);
+                singletonMap(tp, new OffsetAndMetadata(39L)), new ConsumerGroupMetadata(consumerGroupId));
 
-        prepareAddOffsetsToTxnResponse(Errors.NONE, consumerGroupId, pid, epoch);
-        sender.run(time.milliseconds());  // AddOffsetsToTxn Handled, TxnOffsetCommit Enqueued
-        sender.run(time.milliseconds());  // FindCoordinator Enqueued
+        prepareAddOffsetsToTxnResponse(Errors.NONE, consumerGroupId, producerId, epoch);
+        runUntil(() -> !transactionManager.hasPartitionsToAdd());
 
         prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.GROUP, consumerGroupId);
-        sender.run(time.milliseconds());  // FindCoordinator Returned
-
-        prepareTxnOffsetCommitResponse(consumerGroupId, pid, epoch, singletonMap(tp, Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED));
-        sender.run(time.milliseconds());  // TxnOffsetCommit Handled
+        prepareTxnOffsetCommitResponse(consumerGroupId, producerId, epoch, singletonMap(tp, Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED));
+        runUntil(transactionManager::hasError);
 
-        assertTrue(transactionManager.hasError());
         assertTrue(transactionManager.lastError() instanceof TransactionalIdAuthorizationException);
         assertTrue(sendOffsetsResult.isCompleted());
         assertFalse(sendOffsetsResult.isSuccessful());
@@ -984,25 +1194,28 @@ public void testTransactionalIdAuthorizationFailureInTxnOffsetCommit() {
     }
 
     @Test
-    public void testTopicAuthorizationFailureInAddPartitions() {
-        final long pid = 13131L;
-        final short epoch = 1;
+    public void testTopicAuthorizationFailureInAddPartitions() throws InterruptedException {
         final TopicPartition tp0 = new TopicPartition("foo", 0);
         final TopicPartition tp1 = new TopicPartition("bar", 0);
 
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp1);
+
+        FutureRecordMetadata firstPartitionAppend = appendToAccumulator(tp0);
+        FutureRecordMetadata secondPartitionAppend = appendToAccumulator(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());
+        runUntil(transactionManager::hasError);
 
-        assertTrue(transactionManager.hasError());
         assertTrue(transactionManager.lastError() instanceof TopicAuthorizationException);
         assertFalse(transactionManager.isPartitionPendingAdd(tp0));
         assertFalse(transactionManager.isPartitionPendingAdd(tp1));
@@ -1012,172 +1225,199 @@ public void testTopicAuthorizationFailureInAddPartitions() {
 
         TopicAuthorizationException exception = (TopicAuthorizationException) transactionManager.lastError();
         assertEquals(singleton(tp0.topic()), exception.unauthorizedTopics());
-
         assertAbortableError(TopicAuthorizationException.class);
+        sender.runOnce();
+
+        TestUtils.assertFutureThrows(firstPartitionAppend, KafkaException.class);
+        TestUtils.assertFutureThrows(secondPartitionAppend, KafkaException.class);
+    }
+
+    @Test
+    public void testCommitWithTopicAuthorizationFailureInAddPartitionsInFlight() throws InterruptedException {
+        final TopicPartition tp0 = new TopicPartition("foo", 0);
+        final TopicPartition tp1 = new TopicPartition("bar", 0);
+
+        doInitTransactions();
+
+        // Begin a transaction, send two records, and begin commit
+        transactionManager.beginTransaction();
+        transactionManager.maybeAddPartitionToTransaction(tp0);
+        transactionManager.maybeAddPartitionToTransaction(tp1);
+        FutureRecordMetadata firstPartitionAppend = appendToAccumulator(tp0);
+        FutureRecordMetadata secondPartitionAppend = appendToAccumulator(tp1);
+        TransactionalRequestResult commitResult = transactionManager.beginCommit();
+
+        // We send the AddPartitionsToTxn request in the first sender call
+        sender.runOnce();
+        assertFalse(transactionManager.hasError());
+        assertFalse(commitResult.isCompleted());
+        assertFalse(firstPartitionAppend.isDone());
+
+        // The AddPartitionsToTxn response returns in the next call with the error
+        Map errors = new HashMap<>();
+        errors.put(tp0, Errors.TOPIC_AUTHORIZATION_FAILED);
+        errors.put(tp1, Errors.OPERATION_NOT_ATTEMPTED);
+        client.respond(body -> {
+            AddPartitionsToTxnRequest request = (AddPartitionsToTxnRequest) body;
+            assertEquals(new HashSet<>(request.partitions()), new HashSet<>(errors.keySet()));
+            return true;
+        }, new AddPartitionsToTxnResponse(0, errors));
+
+        sender.runOnce();
+        assertTrue(transactionManager.hasError());
+        assertFalse(commitResult.isCompleted());
+        assertFalse(firstPartitionAppend.isDone());
+        assertFalse(secondPartitionAppend.isDone());
+
+        // The next call aborts the records, which have not yet been sent. It should
+        // not block because there are no requests pending and we still need to cancel
+        // the pending transaction commit.
+        sender.runOnce();
+        assertTrue(commitResult.isCompleted());
+        TestUtils.assertFutureThrows(firstPartitionAppend, KafkaException.class);
+        TestUtils.assertFutureThrows(secondPartitionAppend, KafkaException.class);
+        assertTrue(commitResult.error() instanceof TopicAuthorizationException);
     }
 
     @Test
     public void testRecoveryFromAbortableErrorTransactionNotStarted() throws Exception {
-        final long pid = 13131L;
-        final short epoch = 1;
         final TopicPartition unauthorizedPartition = new TopicPartition("foo", 0);
 
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
 
         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;
+        Future responseFuture = appendToAccumulator(unauthorizedPartition);
 
         prepareAddPartitionsToTxn(singletonMap(unauthorizedPartition, Errors.TOPIC_AUTHORIZATION_FAILED));
-        sender.run(time.milliseconds());
+        runUntil(() -> !client.hasPendingResponses());
 
         assertTrue(transactionManager.hasAbortableError());
         transactionManager.beginAbort();
-        sender.run(time.milliseconds());
-        assertTrue(responseFuture.isDone());
-        assertFutureFailed(responseFuture);
+        runUntil(responseFuture::isDone);
+        assertProduceFutureFailed(responseFuture);
 
         // No partitions added, so no need to prepare EndTxn response
-        sender.run(time.milliseconds());
-        assertTrue(transactionManager.isReady());
+        runUntil(transactionManager::isReady);
         assertFalse(transactionManager.hasPartitionsToAdd());
         assertFalse(accumulator.hasIncomplete());
 
         // 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;
+        responseFuture = appendToAccumulator(tp0);
 
         prepareAddPartitionsToTxn(singletonMap(tp0, Errors.NONE));
-        sender.run(time.milliseconds());
-        assertTrue(transactionManager.isPartitionAdded(tp0));
+        runUntil(() -> transactionManager.isPartitionAdded(tp0));
         assertFalse(transactionManager.hasPartitionsToAdd());
 
         transactionManager.beginCommit();
-        prepareProduceResponse(Errors.NONE, pid, epoch);
-        sender.run(time.milliseconds());
-
-        assertTrue(responseFuture.isDone());
+        prepareProduceResponse(Errors.NONE, producerId, epoch);
+        runUntil(responseFuture::isDone);
         assertNotNull(responseFuture.get());
 
-        prepareEndTxnResponse(Errors.NONE, TransactionResult.COMMIT, pid, epoch);
-        sender.run(time.milliseconds());
-
-        assertTrue(transactionManager.isReady());
+        prepareEndTxnResponse(Errors.NONE, TransactionResult.COMMIT, producerId, epoch);
+        runUntil(transactionManager::isReady);
     }
 
     @Test
     public void testRecoveryFromAbortableErrorTransactionStarted() throws Exception {
-        final long pid = 13131L;
-        final short epoch = 1;
         final TopicPartition unauthorizedPartition = new TopicPartition("foo", 0);
 
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
 
         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());
-        assertTrue(transactionManager.isPartitionAdded(tp0));
+        Future authorizedTopicProduceFuture = appendToAccumulator(unauthorizedPartition);
+        runUntil(() -> 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;
+        Future unauthorizedTopicProduceFuture = appendToAccumulator(unauthorizedPartition);
         prepareAddPartitionsToTxn(singletonMap(unauthorizedPartition, Errors.TOPIC_AUTHORIZATION_FAILED));
-        sender.run(time.milliseconds());
-        assertTrue(transactionManager.hasAbortableError());
+        runUntil(transactionManager::hasAbortableError);
         assertTrue(transactionManager.isPartitionAdded(tp0));
         assertFalse(transactionManager.isPartitionAdded(unauthorizedPartition));
         assertFalse(authorizedTopicProduceFuture.isDone());
         assertFalse(unauthorizedTopicProduceFuture.isDone());
 
-        prepareEndTxnResponse(Errors.NONE, TransactionResult.ABORT, pid, epoch);
+        prepareEndTxnResponse(Errors.NONE, TransactionResult.ABORT, producerId, epoch);
         transactionManager.beginAbort();
-        sender.run(time.milliseconds());
+        runUntil(transactionManager::isReady);
         // neither produce request has been sent, so they should both be failed immediately
-        assertFutureFailed(authorizedTopicProduceFuture);
-        assertFutureFailed(unauthorizedTopicProduceFuture);
-        assertTrue(transactionManager.isReady());
+        assertProduceFutureFailed(authorizedTopicProduceFuture);
+        assertProduceFutureFailed(unauthorizedTopicProduceFuture);
         assertFalse(transactionManager.hasPartitionsToAdd());
         assertFalse(accumulator.hasIncomplete());
 
         // 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;
+        FutureRecordMetadata nextTransactionFuture = appendToAccumulator(tp0);
 
         prepareAddPartitionsToTxn(singletonMap(tp0, Errors.NONE));
-        sender.run(time.milliseconds());
-        assertTrue(transactionManager.isPartitionAdded(tp0));
+        runUntil(() -> transactionManager.isPartitionAdded(tp0));
         assertFalse(transactionManager.hasPartitionsToAdd());
 
         transactionManager.beginCommit();
-        prepareProduceResponse(Errors.NONE, pid, epoch);
-        sender.run(time.milliseconds());
-
-        assertTrue(nextTransactionFuture.isDone());
+        prepareProduceResponse(Errors.NONE, producerId, epoch);
+        runUntil(nextTransactionFuture::isDone);
         assertNotNull(nextTransactionFuture.get());
 
-        prepareEndTxnResponse(Errors.NONE, TransactionResult.COMMIT, pid, epoch);
-        sender.run(time.milliseconds());
-
-        assertTrue(transactionManager.isReady());
+        prepareEndTxnResponse(Errors.NONE, TransactionResult.COMMIT, producerId, epoch);
+        runUntil(transactionManager::isReady);
     }
 
     @Test
     public void testRecoveryFromAbortableErrorProduceRequestInRetry() throws Exception {
-        final long pid = 13131L;
-        final short epoch = 1;
         final TopicPartition unauthorizedPartition = new TopicPartition("foo", 0);
 
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
 
         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());
-        assertTrue(transactionManager.isPartitionAdded(tp0));
+        Future authorizedTopicProduceFuture = appendToAccumulator(tp0);
+        runUntil(() -> transactionManager.isPartitionAdded(tp0));
 
         accumulator.beginFlush();
-        prepareProduceResponse(Errors.REQUEST_TIMED_OUT, pid, epoch);
-        sender.run(time.milliseconds());
+        prepareProduceResponse(Errors.REQUEST_TIMED_OUT, producerId, epoch);
+        runUntil(() -> !client.hasPendingResponses());
         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;
+        Future unauthorizedTopicProduceFuture = appendToAccumulator(unauthorizedPartition);
         prepareAddPartitionsToTxn(singletonMap(unauthorizedPartition, Errors.TOPIC_AUTHORIZATION_FAILED));
-        sender.run(time.milliseconds());
-        assertTrue(transactionManager.hasAbortableError());
+        runUntil(transactionManager::hasAbortableError);
         assertTrue(transactionManager.isPartitionAdded(tp0));
         assertFalse(transactionManager.isPartitionAdded(unauthorizedPartition));
         assertFalse(authorizedTopicProduceFuture.isDone());
 
-        prepareProduceResponse(Errors.NONE, pid, epoch);
-        sender.run(time.milliseconds());
-        assertFutureFailed(unauthorizedTopicProduceFuture);
-        assertTrue(authorizedTopicProduceFuture.isDone());
+        prepareProduceResponse(Errors.NONE, producerId, epoch);
+        runUntil(authorizedTopicProduceFuture::isDone);
+
+        assertProduceFutureFailed(unauthorizedTopicProduceFuture);
         assertNotNull(authorizedTopicProduceFuture.get());
         assertTrue(authorizedTopicProduceFuture.isDone());
 
-        prepareEndTxnResponse(Errors.NONE, TransactionResult.ABORT, pid, epoch);
+        prepareEndTxnResponse(Errors.NONE, TransactionResult.ABORT, producerId, epoch);
         transactionManager.beginAbort();
-        sender.run(time.milliseconds());
+        runUntil(transactionManager::isReady);
         // neither produce request has been sent, so they should both be failed immediately
         assertTrue(transactionManager.isReady());
         assertFalse(transactionManager.hasPartitionsToAdd());
@@ -1186,44 +1426,36 @@ 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;
+        FutureRecordMetadata nextTransactionFuture = appendToAccumulator(tp0);
 
         prepareAddPartitionsToTxn(singletonMap(tp0, Errors.NONE));
-        sender.run(time.milliseconds());
-        assertTrue(transactionManager.isPartitionAdded(tp0));
+        runUntil(() -> transactionManager.isPartitionAdded(tp0));
         assertFalse(transactionManager.hasPartitionsToAdd());
 
         transactionManager.beginCommit();
-        prepareProduceResponse(Errors.NONE, pid, epoch);
-        sender.run(time.milliseconds());
-
-        assertTrue(nextTransactionFuture.isDone());
+        prepareProduceResponse(Errors.NONE, producerId, epoch);
+        runUntil(nextTransactionFuture::isDone);
         assertNotNull(nextTransactionFuture.get());
 
-        prepareEndTxnResponse(Errors.NONE, TransactionResult.COMMIT, pid, epoch);
-        sender.run(time.milliseconds());
-
-        assertTrue(transactionManager.isReady());
+        prepareEndTxnResponse(Errors.NONE, TransactionResult.COMMIT, producerId, epoch);
+        runUntil(transactionManager::isReady);
     }
 
     @Test
     public void testTransactionalIdAuthorizationFailureInAddPartitions() {
-        final long pid = 13131L;
-        final short epoch = 1;
         final TopicPartition tp = new TopicPartition("foo", 0);
 
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp);
 
         prepareAddPartitionsToTxn(tp, Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED);
-        sender.run(time.milliseconds());
-
-        assertTrue(transactionManager.hasError());
+        runUntil(transactionManager::hasError);
         assertTrue(transactionManager.lastError() instanceof TransactionalIdAuthorizationException);
 
         assertFatalError(TransactionalIdAuthorizationException.class);
@@ -1231,16 +1463,13 @@ public void testTransactionalIdAuthorizationFailureInAddPartitions() {
 
     @Test
     public void testFlushPendingPartitionsOnCommit() throws InterruptedException {
-        final long pid = 13131L;
-        final short epoch = 1;
-
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
 
         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;
+        Future responseFuture = appendToAccumulator(tp0);
 
         assertFalse(responseFuture.isDone());
 
@@ -1252,58 +1481,50 @@ public void testFlushPendingPartitionsOnCommit() throws InterruptedException {
         //  2. Produce
         //  3. EndTxn.
         assertFalse(transactionManager.transactionContainsPartition(tp0));
-        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, pid);
+        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, producerId);
 
-        sender.run(time.milliseconds());  // AddPartitions.
-        assertTrue(transactionManager.transactionContainsPartition(tp0));
+        runUntil(() -> transactionManager.transactionContainsPartition(tp0));
         assertFalse(responseFuture.isDone());
         assertFalse(commitResult.isCompleted());
 
-        prepareProduceResponse(Errors.NONE, pid, epoch);
-        sender.run(time.milliseconds());  // Produce.
-        assertTrue(responseFuture.isDone());
+        prepareProduceResponse(Errors.NONE, producerId, epoch);
+        runUntil(responseFuture::isDone);
 
-        prepareEndTxnResponse(Errors.NONE, TransactionResult.COMMIT, pid, epoch);
+        prepareEndTxnResponse(Errors.NONE, TransactionResult.COMMIT, producerId, epoch);
         assertFalse(commitResult.isCompleted());
         assertTrue(transactionManager.hasOngoingTransaction());
         assertTrue(transactionManager.isCompleting());
 
-        sender.run(time.milliseconds());
-        assertTrue(commitResult.isCompleted());
+        runUntil(commitResult::isCompleted);
         assertFalse(transactionManager.hasOngoingTransaction());
     }
 
     @Test
     public void testMultipleAddPartitionsPerForOneProduce() throws InterruptedException {
-        final long pid = 13131L;
-        final short epoch = 1;
-
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
 
         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;
+        Future responseFuture = appendToAccumulator(tp0);
 
         assertFalse(responseFuture.isDone());
-        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, pid);
+        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, producerId);
 
         assertFalse(transactionManager.transactionContainsPartition(tp0));
 
         // Sender flushes one add partitions. The produce goes next.
-        sender.run(time.milliseconds());  // send addPartitions.
-        // Check that only addPartitions was sent.
-        assertTrue(transactionManager.transactionContainsPartition(tp0));
+        runUntil(() -> 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;
+        Future secondResponseFuture = appendToAccumulator(tp0);
 
-        prepareAddPartitionsToTxnResponse(Errors.NONE, tp1, epoch, pid);
-        prepareProduceResponse(Errors.NONE, pid, epoch);
+        prepareAddPartitionsToTxnResponse(Errors.NONE, tp1, epoch, producerId);
+        prepareProduceResponse(Errors.NONE, producerId, epoch);
 
         assertFalse(transactionManager.transactionContainsPartition(tp1));
 
@@ -1311,73 +1532,199 @@ 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
-        assertTrue(transactionManager.transactionContainsPartition(tp1));
+        runUntil(() -> transactionManager.transactionContainsPartition(tp1));
 
         assertFalse(responseFuture.isDone());
         assertFalse(secondResponseFuture.isDone());
 
         // Finally we get to the produce.
-        sender.run(time.milliseconds());  // send produce request
-
-        assertTrue(responseFuture.isDone());
+        runUntil(responseFuture::isDone);
         assertTrue(secondResponseFuture.isDone());
     }
 
-    @Test(expected = ExecutionException.class)
-    public void testProducerFencedException() throws InterruptedException, ExecutionException {
-        final long pid = 13131L;
-        final short epoch = 1;
+    @Test
+    public void testProducerFencedExceptionInInitProducerId() {
+        verifyProducerFencedForInitProducerId(Errors.PRODUCER_FENCED);
+    }
+
+    @Test
+    public void testInvalidProducerEpochConvertToProducerFencedInInitProducerId() {
+        verifyProducerFencedForInitProducerId(Errors.INVALID_PRODUCER_EPOCH);
+    }
+
+    private void verifyProducerFencedForInitProducerId(Errors error) {
+        TransactionalRequestResult result = transactionManager.initializeTransactions();
+        prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.TRANSACTION, transactionalId);
+        runUntil(() -> transactionManager.coordinator(CoordinatorType.TRANSACTION) != null);
+        assertEquals(brokerNode, transactionManager.coordinator(CoordinatorType.TRANSACTION));
+
+        prepareInitPidResponse(error, false, producerId, epoch);
+
+        runUntil(transactionManager::hasError);
 
-        doInitTransactions(pid, epoch);
+        assertEquals(ProducerFencedException.class, result.error().getClass());
+
+        assertThrows(ProducerFencedException.class, () -> transactionManager.beginTransaction());
+        assertThrows(ProducerFencedException.class, () -> transactionManager.beginCommit());
+        assertThrows(ProducerFencedException.class, () -> transactionManager.beginAbort());
+        assertThrows(ProducerFencedException.class, () -> transactionManager.sendOffsetsToTransaction(
+            Collections.emptyMap(), new ConsumerGroupMetadata("dummyId")));
+    }
+
+    @Test
+    public void testProducerFencedInAddPartitionToTxn() throws InterruptedException {
+        verifyProducerFencedForAddPartitionsToTxn(Errors.PRODUCER_FENCED);
+    }
+
+    @Test
+    public void testInvalidProducerEpochConvertToProducerFencedInAddPartitionToTxn() throws InterruptedException {
+        verifyProducerFencedForAddPartitionsToTxn(Errors.INVALID_PRODUCER_EPOCH);
+    }
+
+    private void verifyProducerFencedForAddPartitionsToTxn(Errors error) throws InterruptedException {
+        doInitTransactions();
 
         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;
+        Future responseFuture = appendToAccumulator(tp0);
 
         assertFalse(responseFuture.isDone());
-        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, pid);
-        prepareProduceResponse(Errors.INVALID_PRODUCER_EPOCH, pid, epoch);
-        sender.run(time.milliseconds()); // Add partitions.
+        prepareAddPartitionsToTxnResponse(error, tp0, epoch, producerId);
 
-        sender.run(time.milliseconds());  // send produce.
+        verifyProducerFenced(responseFuture);
+    }
+
+    @Test
+    public void testProducerFencedInAddOffSetsToTxn() throws InterruptedException {
+        verifyProducerFencedForAddOffsetsToTxn(Errors.INVALID_PRODUCER_EPOCH);
+    }
+
+    @Test
+    public void testInvalidProducerEpochConvertToProducerFencedInAddOffSetsToTxn() throws InterruptedException {
+        verifyProducerFencedForAddOffsetsToTxn(Errors.INVALID_PRODUCER_EPOCH);
+    }
 
-        assertTrue(responseFuture.isDone());
+    private void verifyProducerFencedForAddOffsetsToTxn(Errors error) throws InterruptedException {
+        doInitTransactions();
+
+        transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
+        transactionManager.sendOffsetsToTransaction(Collections.emptyMap(), new ConsumerGroupMetadata(consumerGroupId));
+
+        Future responseFuture = appendToAccumulator(tp0);
+
+        assertFalse(responseFuture.isDone());
+        prepareAddOffsetsToTxnResponse(error, consumerGroupId, producerId, epoch);
+
+        verifyProducerFenced(responseFuture);
+    }
+
+    private void verifyProducerFenced(Future responseFuture) throws InterruptedException {
+        runUntil(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(), new ConsumerGroupMetadata("dummyId")));
     }
 
     @Test
-    public void testDisallowCommitOnProduceFailure() throws InterruptedException {
-        final long pid = 13131L;
-        final short epoch = 1;
+    public void testInvalidProducerEpochConvertToProducerFencedInEndTxn() throws InterruptedException {
+        doInitTransactions();
+
+        transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
+        transactionManager.maybeAddPartitionToTransaction(tp0);
+        TransactionalRequestResult commitResult = transactionManager.beginCommit();
+
+        Future responseFuture = appendToAccumulator(tp0);
+
+        assertFalse(responseFuture.isDone());
+        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, producerId);
+        prepareProduceResponse(Errors.NONE, producerId, epoch);
+        prepareEndTxnResponse(Errors.INVALID_PRODUCER_EPOCH, TransactionResult.COMMIT, producerId, epoch);
+
+        runUntil(commitResult::isCompleted);
+        runUntil(responseFuture::isDone);
+
+        // make sure the exception was thrown directly from the follow-up calls.
+        assertThrows(KafkaException.class, () -> transactionManager.beginTransaction());
+        assertThrows(KafkaException.class, () -> transactionManager.beginCommit());
+        assertThrows(KafkaException.class, () -> transactionManager.beginAbort());
+        assertThrows(KafkaException.class, () -> transactionManager.sendOffsetsToTransaction(
+            Collections.emptyMap(), new ConsumerGroupMetadata("dummyId")));
+    }
+
+    @Test
+    public void testInvalidProducerEpochFromProduce() throws InterruptedException {
+        doInitTransactions();
 
-        doInitTransactions(pid, epoch);
+        transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
+        transactionManager.maybeAddPartitionToTransaction(tp0);
+
+        Future responseFuture = appendToAccumulator(tp0);
+
+        assertFalse(responseFuture.isDone());
+        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, producerId);
+        prepareProduceResponse(Errors.INVALID_PRODUCER_EPOCH, producerId, epoch);
+        prepareProduceResponse(Errors.NONE, producerId, epoch);
+
+        sender.runOnce();
+
+        runUntil(responseFuture::isDone);
+        assertTrue(transactionManager.hasError());
+
+        transactionManager.beginAbort();
+
+        TransactionManager.TxnRequestHandler handler = transactionManager.nextRequest(false);
+
+        // First we will get an EndTxn for abort.
+        assertNotNull(handler);
+        assertTrue(handler.requestBuilder() instanceof EndTxnRequest.Builder);
+
+        handler = transactionManager.nextRequest(false);
+
+        // Second we will see an InitPid for handling InvalidProducerEpoch.
+        assertNotNull(handler);
+        assertTrue(handler.requestBuilder() instanceof InitProducerIdRequest.Builder);
+    }
+
+    @Test
+    public void testDisallowCommitOnProduceFailure() throws InterruptedException {
+        doInitTransactions();
 
         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;
+        Future responseFuture = appendToAccumulator(tp0);
 
         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
-        assertFalse(commitResult.isCompleted());
-        sender.run(time.milliseconds());  // Send Produce Request, returns OutOfOrderSequenceException.
+        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, producerId);
+        prepareProduceResponse(Errors.OUT_OF_ORDER_SEQUENCE_NUMBER, producerId, epoch);
 
-        sender.run(time.milliseconds());  // try to commit.
-        assertTrue(commitResult.isCompleted());  // commit should be cancelled with exception without being sent.
+        runUntil(commitResult::isCompleted);  // commit should be cancelled with exception without being sent.
 
         try {
             commitResult.await();
             fail();  // the get() must throw an exception.
         } catch (KafkaException e) {
+            // Expected
         }
 
         try {
@@ -1389,231 +1736,194 @@ 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
-
-        assertTrue(abortResult.isCompleted());
+        prepareEndTxnResponse(Errors.NONE, TransactionResult.ABORT, producerId, epoch);
+        prepareInitPidResponse(Errors.NONE, false, producerId, (short) (epoch + 1));
+        runUntil(abortResult::isCompleted);
         assertTrue(abortResult.isSuccessful());
         assertTrue(transactionManager.isReady());  // make sure we are ready for a transaction now.
     }
 
     @Test
     public void testAllowAbortOnProduceFailure() throws InterruptedException {
-        final long pid = 13131L;
-        final short epoch = 1;
-
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
 
         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;
+        Future responseFuture = appendToAccumulator(tp0);
 
         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.
+        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, producerId);
+        prepareProduceResponse(Errors.OUT_OF_ORDER_SEQUENCE_NUMBER, producerId, epoch);
 
+        // Because this is a failure that triggers an epoch bump, the abort will trigger an InitProducerId call
+        runUntil(transactionManager::hasAbortableError);
         TransactionalRequestResult abortResult = transactionManager.beginAbort();
-        sender.run(time.milliseconds());  // try to abort
-        assertTrue(abortResult.isCompleted());
+        prepareEndTxnResponse(Errors.NONE, TransactionResult.ABORT, producerId, epoch);
+        prepareInitPidResponse(Errors.NONE, false, producerId, (short) (epoch + 1));
+        runUntil(abortResult::isCompleted);
         assertTrue(abortResult.isSuccessful());
         assertTrue(transactionManager.isReady());  // make sure we are ready for a transaction now.
     }
 
     @Test
     public void testAbortableErrorWhileAbortInProgress() throws InterruptedException {
-        final long pid = 13131L;
-        final short epoch = 1;
-
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
 
         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;
+        Future responseFuture = appendToAccumulator(tp0);
 
         assertFalse(responseFuture.isDone());
-        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, pid);
-
-        sender.run(time.milliseconds());  // Send AddPartitionsRequest
-        sender.run(time.milliseconds());  // Send Produce Request
+        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, producerId);
+        runUntil(() -> !accumulator.hasUndrained());
 
         TransactionalRequestResult abortResult = transactionManager.beginAbort();
         assertTrue(transactionManager.isAborting());
         assertFalse(transactionManager.hasError());
 
-        sendProduceResponse(Errors.OUT_OF_ORDER_SEQUENCE_NUMBER, pid, epoch);
-        prepareEndTxnResponse(Errors.NONE, TransactionResult.ABORT, pid, epoch);
-        sender.run(time.milliseconds());  // receive the produce response
+        sendProduceResponse(Errors.OUT_OF_ORDER_SEQUENCE_NUMBER, producerId, epoch);
+        prepareEndTxnResponse(Errors.NONE, TransactionResult.ABORT, producerId, epoch);
+        runUntil(responseFuture::isDone);
 
         // 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
-        assertTrue(abortResult.isCompleted());
+        runUntil(abortResult::isCompleted);
         assertTrue(abortResult.isSuccessful());
         assertTrue(transactionManager.isReady());  // make sure we are ready for a transaction now.
     }
 
     @Test
     public void testCommitTransactionWithUnsentProduceRequest() throws Exception {
-        final long pid = 13131L;
-        final short epoch = 1;
-
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
 
         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;
+        Future responseFuture = appendToAccumulator(tp0);
 
         prepareAddPartitionsToTxn(tp0, Errors.NONE);
-        sender.run(time.milliseconds());
+        runUntil(() -> !client.hasPendingResponses());
         assertTrue(accumulator.hasUndrained());
 
         // committing the transaction should cause the unsent batch to be flushed
         transactionManager.beginCommit();
-        sender.run(time.milliseconds());
-        assertFalse(accumulator.hasUndrained());
+        runUntil(() -> !accumulator.hasUndrained());
         assertTrue(accumulator.hasIncomplete());
-        assertFalse(transactionManager.hasInFlightTransactionalRequest());
+        assertFalse(transactionManager.hasInFlightRequest());
         assertFalse(responseFuture.isDone());
 
         // until the produce future returns, we will not send EndTxn
-        sender.run(time.milliseconds());
+        AtomicInteger numRuns = new AtomicInteger(0);
+        runUntil(() -> numRuns.incrementAndGet() >= 4);
         assertFalse(accumulator.hasUndrained());
         assertTrue(accumulator.hasIncomplete());
-        assertFalse(transactionManager.hasInFlightTransactionalRequest());
+        assertFalse(transactionManager.hasInFlightRequest());
         assertFalse(responseFuture.isDone());
 
         // now the produce response returns
-        sendProduceResponse(Errors.NONE, pid, epoch);
-        sender.run(time.milliseconds());
-        assertTrue(responseFuture.isDone());
+        sendProduceResponse(Errors.NONE, producerId, epoch);
+        runUntil(responseFuture::isDone);
         assertFalse(accumulator.hasUndrained());
         assertFalse(accumulator.hasIncomplete());
-        assertFalse(transactionManager.hasInFlightTransactionalRequest());
+        assertFalse(transactionManager.hasInFlightRequest());
 
         // now we send EndTxn
-        sender.run(time.milliseconds());
-        assertTrue(transactionManager.hasInFlightTransactionalRequest());
-        sendEndTxnResponse(Errors.NONE, TransactionResult.COMMIT, pid, epoch);
-        sender.run(time.milliseconds());
-        assertFalse(transactionManager.hasInFlightTransactionalRequest());
-        assertTrue(transactionManager.isReady());
+        runUntil(transactionManager::hasInFlightRequest);
+        sendEndTxnResponse(Errors.NONE, TransactionResult.COMMIT, producerId, epoch);
+
+        runUntil(transactionManager::isReady);
+        assertFalse(transactionManager.hasInFlightRequest());
     }
 
     @Test
     public void testCommitTransactionWithInFlightProduceRequest() throws Exception {
-        final long pid = 13131L;
-        final short epoch = 1;
-
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
 
         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;
+        Future responseFuture = appendToAccumulator(tp0);
 
         prepareAddPartitionsToTxn(tp0, Errors.NONE);
-        sender.run(time.milliseconds());
+        runUntil(() -> !transactionManager.hasPartitionsToAdd());
         assertTrue(accumulator.hasUndrained());
 
         accumulator.beginFlush();
-        sender.run(time.milliseconds());
+        runUntil(() -> !accumulator.hasUndrained());
         assertFalse(accumulator.hasUndrained());
         assertTrue(accumulator.hasIncomplete());
-        assertFalse(transactionManager.hasInFlightTransactionalRequest());
+        assertFalse(transactionManager.hasInFlightRequest());
 
         // now we begin the commit with the produce request still pending
         transactionManager.beginCommit();
-        sender.run(time.milliseconds());
-        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());
+        AtomicInteger numRuns = new AtomicInteger(0);
+        runUntil(() -> numRuns.incrementAndGet() >= 4);
         assertFalse(accumulator.hasUndrained());
         assertTrue(accumulator.hasIncomplete());
-        assertFalse(transactionManager.hasInFlightTransactionalRequest());
+        assertFalse(transactionManager.hasInFlightRequest());
         assertFalse(responseFuture.isDone());
 
         // now the produce response returns
-        sendProduceResponse(Errors.NONE, pid, epoch);
-        sender.run(time.milliseconds());
-        assertTrue(responseFuture.isDone());
+        sendProduceResponse(Errors.NONE, producerId, epoch);
+        runUntil(responseFuture::isDone);
         assertFalse(accumulator.hasUndrained());
         assertFalse(accumulator.hasIncomplete());
-        assertFalse(transactionManager.hasInFlightTransactionalRequest());
+        assertFalse(transactionManager.hasInFlightRequest());
 
         // now we send EndTxn
-        sender.run(time.milliseconds());
-        assertTrue(transactionManager.hasInFlightTransactionalRequest());
-        sendEndTxnResponse(Errors.NONE, TransactionResult.COMMIT, pid, epoch);
-        sender.run(time.milliseconds());
-        assertFalse(transactionManager.hasInFlightTransactionalRequest());
-        assertTrue(transactionManager.isReady());
+        runUntil(transactionManager::hasInFlightRequest);
+        sendEndTxnResponse(Errors.NONE, TransactionResult.COMMIT, producerId, epoch);
+        runUntil(transactionManager::isReady);
+        assertFalse(transactionManager.hasInFlightRequest());
     }
 
     @Test
     public void testFindCoordinatorAllowedInAbortableErrorState() throws InterruptedException {
-        final long pid = 13131L;
-        final short epoch = 1;
-
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
 
         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;
+        Future responseFuture = appendToAccumulator(tp0);
 
         assertFalse(responseFuture.isDone());
-        sender.run(time.milliseconds());  // Send AddPartitionsRequest
+        runUntil(transactionManager::hasInFlightRequest);
 
         transactionManager.transitionToAbortableError(new KafkaException());
-        sendAddPartitionsToTxnResponse(Errors.NOT_COORDINATOR, tp0, epoch, pid);
-        sender.run(time.milliseconds()); // AddPartitions returns
-        assertTrue(transactionManager.hasAbortableError());
+        sendAddPartitionsToTxnResponse(Errors.NOT_COORDINATOR, tp0, epoch, producerId);
+        runUntil(() -> transactionManager.coordinator(CoordinatorType.TRANSACTION) == null);
 
-        assertNull(transactionManager.coordinator(CoordinatorType.TRANSACTION));
         prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.TRANSACTION, transactionalId);
-        sender.run(time.milliseconds()); // FindCoordinator handled
+        runUntil(() -> transactionManager.coordinator(CoordinatorType.TRANSACTION) != null);
         assertEquals(brokerNode, transactionManager.coordinator(CoordinatorType.TRANSACTION));
         assertTrue(transactionManager.hasAbortableError());
     }
 
     @Test
     public void testCancelUnsentAddPartitionsAndProduceOnAbort() throws InterruptedException {
-        final long pid = 13131L;
-        final short epoch = 1;
-
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
 
         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;
+        Future responseFuture = appendToAccumulator(tp0);
 
         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
-        assertTrue(abortResult.isCompleted());
+        runUntil(abortResult::isCompleted);
         assertTrue(abortResult.isSuccessful());
         assertTrue(transactionManager.isReady());  // make sure we are ready for a transaction now.
 
@@ -1627,31 +1937,25 @@ public void testCancelUnsentAddPartitionsAndProduceOnAbort() throws InterruptedE
 
     @Test
     public void testAbortResendsAddPartitionErrorIfRetried() throws InterruptedException {
-        final long producerId = 13131L;
-        final short producerEpoch = 1;
-
-        doInitTransactions(producerId, producerEpoch);
+        doInitTransactions(producerId, epoch);
 
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
-        prepareAddPartitionsToTxnResponse(Errors.UNKNOWN_TOPIC_OR_PARTITION, tp0, producerEpoch, producerId);
+        prepareAddPartitionsToTxnResponse(Errors.UNKNOWN_TOPIC_OR_PARTITION, tp0, epoch, producerId);
 
-        Future responseFuture = accumulator.append(tp0, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT).future;
+        Future responseFuture = appendToAccumulator(tp0);
 
-        sender.run(time.milliseconds());  // Send AddPartitions and let it fail
+        runUntil(() -> !client.hasPendingResponses());
         assertFalse(responseFuture.isDone());
 
         TransactionalRequestResult abortResult = transactionManager.beginAbort();
 
         // we should resend the AddPartitions
-        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, producerEpoch, producerId);
-        prepareEndTxnResponse(Errors.NONE, TransactionResult.ABORT, producerId, producerEpoch);
+        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, producerId);
+        prepareEndTxnResponse(Errors.NONE, TransactionResult.ABORT, producerId, epoch);
 
-        sender.run(time.milliseconds());  // Resend AddPartitions
-        sender.run(time.milliseconds());  // Send EndTxn
-
-        assertTrue(abortResult.isCompleted());
+        runUntil(abortResult::isCompleted);
         assertTrue(abortResult.isSuccessful());
         assertTrue(transactionManager.isReady());  // make sure we are ready for a transaction now.
 
@@ -1665,34 +1969,26 @@ public void testAbortResendsAddPartitionErrorIfRetried() throws InterruptedExcep
 
     @Test
     public void testAbortResendsProduceRequestIfRetried() throws Exception {
-        final long producerId = 13131L;
-        final short producerEpoch = 1;
-
-        doInitTransactions(producerId, producerEpoch);
+        doInitTransactions(producerId, epoch);
 
         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;
+        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, producerId);
+        prepareProduceResponse(Errors.REQUEST_TIMED_OUT, producerId, epoch);
 
-        sender.run(time.milliseconds());  // Send AddPartitions
-        sender.run(time.milliseconds());  // Send ProduceRequest and let it fail
+        Future responseFuture = appendToAccumulator(tp0);
 
+        runUntil(() -> !client.hasPendingResponses());
         assertFalse(responseFuture.isDone());
 
         TransactionalRequestResult abortResult = transactionManager.beginAbort();
 
         // we should resend the ProduceRequest before aborting
-        prepareProduceResponse(Errors.NONE, producerId, producerEpoch);
-        prepareEndTxnResponse(Errors.NONE, TransactionResult.ABORT, producerId, producerEpoch);
+        prepareProduceResponse(Errors.NONE, producerId, epoch);
+        prepareEndTxnResponse(Errors.NONE, TransactionResult.ABORT, producerId, epoch);
 
-        sender.run(time.milliseconds());  // Resend ProduceRequest
-        sender.run(time.milliseconds());  // Send EndTxn
-
-        assertTrue(abortResult.isCompleted());
+        runUntil(abortResult::isCompleted);
         assertTrue(abortResult.isSuccessful());
         assertTrue(transactionManager.isReady());  // make sure we are ready for a transaction now.
 
@@ -1702,167 +1998,290 @@ public void testAbortResendsProduceRequestIfRetried() throws Exception {
 
     @Test
     public void testHandlingOfUnknownTopicPartitionErrorOnAddPartitions() throws InterruptedException {
-        final long pid = 13131L;
-        final short epoch = 1;
-
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
 
         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;
+        Future responseFuture = appendToAccumulator(tp0);
 
         assertFalse(responseFuture.isDone());
-        prepareAddPartitionsToTxnResponse(Errors.UNKNOWN_TOPIC_OR_PARTITION, tp0, epoch, pid);
+        prepareAddPartitionsToTxnResponse(Errors.UNKNOWN_TOPIC_OR_PARTITION, tp0, epoch, producerId);
 
-        sender.run(time.milliseconds());  // Send AddPartitionsRequest
+        runUntil(() -> !client.hasPendingResponses());
         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.
-        assertTrue(transactionManager.transactionContainsPartition(tp0));
+        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, producerId);
+        prepareProduceResponse(Errors.NONE, producerId, epoch);
+        runUntil(() -> transactionManager.transactionContainsPartition(tp0));
+        runUntil(responseFuture::isDone);
+    }
 
-        sender.run(time.milliseconds());  // Send ProduceRequest.
-        assertTrue(responseFuture.isDone());
+    @Test
+    public void testHandlingOfUnknownTopicPartitionErrorOnTxnOffsetCommit() {
+        testRetriableErrorInTxnOffsetCommit(Errors.UNKNOWN_TOPIC_OR_PARTITION);
     }
 
     @Test
-    public void testHandlingOfUnknownTopicPartitionErrorOnTxnOffsetCommit() throws InterruptedException {
-        final long pid = 13131L;
-        final short epoch = 1;
+    public void testHandlingOfCoordinatorLoadingErrorOnTxnOffsetCommit() {
+        testRetriableErrorInTxnOffsetCommit(Errors.COORDINATOR_LOAD_IN_PROGRESS);
+    }
 
-        doInitTransactions(pid, epoch);
+    private void testRetriableErrorInTxnOffsetCommit(Errors error) {
+        doInitTransactions();
 
         transactionManager.beginTransaction();
 
         Map offsets = new HashMap<>();
+        offsets.put(tp0, new OffsetAndMetadata(1));
         offsets.put(tp1, new OffsetAndMetadata(1));
-        final String consumerGroupId = "myconsumergroup";
-
-        TransactionalRequestResult addOffsetsResult = transactionManager.sendOffsetsToTransaction(offsets, consumerGroupId);
-        prepareAddOffsetsToTxnResponse(Errors.NONE, consumerGroupId, pid, epoch);
-
-        sender.run(time.milliseconds());  // send AddOffsetsToTxnResult
 
+        TransactionalRequestResult addOffsetsResult = transactionManager.sendOffsetsToTransaction(
+            offsets, new ConsumerGroupMetadata(consumerGroupId));
+        prepareAddOffsetsToTxnResponse(Errors.NONE, consumerGroupId, producerId, epoch);
+        runUntil(() -> !client.hasPendingResponses());
         assertFalse(addOffsetsResult.isCompleted());  // The request should complete only after the TxnOffsetCommit completes.
 
         Map txnOffsetCommitResponse = new HashMap<>();
-        txnOffsetCommitResponse.put(tp1, Errors.UNKNOWN_TOPIC_OR_PARTITION);
+        txnOffsetCommitResponse.put(tp0, Errors.NONE);
+        txnOffsetCommitResponse.put(tp1, error);
 
         prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.GROUP, consumerGroupId);
-        prepareTxnOffsetCommitResponse(consumerGroupId, pid, epoch, txnOffsetCommitResponse);
+        prepareTxnOffsetCommitResponse(consumerGroupId, producerId, epoch, txnOffsetCommitResponse);
 
-        assertEquals(null, 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
-        assertNotNull(transactionManager.coordinator(CoordinatorType.GROUP));
+        assertNull(transactionManager.coordinator(CoordinatorType.GROUP));
+        runUntil(() -> transactionManager.coordinator(CoordinatorType.GROUP) != null);
         assertTrue(transactionManager.hasPendingOffsetCommits());
 
-        sender.run(time.milliseconds());  // send TxnOffsetCommitRequest request.
-
-        assertTrue(transactionManager.hasPendingOffsetCommits());  // The TxnOffsetCommit failed.
+        runUntil(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.
-
-        assertTrue(addOffsetsResult.isCompleted());
+        prepareTxnOffsetCommitResponse(consumerGroupId, producerId, epoch, txnOffsetCommitResponse);
+        runUntil(addOffsetsResult::isCompleted);
         assertTrue(addOffsetsResult.isSuccessful());
     }
 
     @Test
-    public void shouldNotAddPartitionsToTransactionWhenTopicAuthorizationFailed() throws Exception {
-        verifyAddPartitionsFailsWithPartitionLevelError(Errors.TOPIC_AUTHORIZATION_FAILED);
+    public void testHandlingOfProducerFencedErrorOnTxnOffsetCommit() {
+        testFatalErrorInTxnOffsetCommit(Errors.PRODUCER_FENCED);
     }
 
     @Test
-    public void shouldNotSendAbortTxnRequestWhenOnlyAddPartitionsRequestFailed() throws Exception {
-        final long pid = 13131L;
-        final short epoch = 1;
+    public void testHandlingOfTransactionalIdAuthorizationFailedErrorOnTxnOffsetCommit() {
+        testFatalErrorInTxnOffsetCommit(Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED);
+    }
+
+    @Test
+    public void testHandlingOfInvalidProducerEpochErrorOnTxnOffsetCommit() {
+        testFatalErrorInTxnOffsetCommit(Errors.INVALID_PRODUCER_EPOCH);
+    }
 
-        doInitTransactions(pid, epoch);
+    @Test
+    public void testHandlingOfUnsupportedForMessageFormatErrorOnTxnOffsetCommit() {
+        testFatalErrorInTxnOffsetCommit(Errors.UNSUPPORTED_FOR_MESSAGE_FORMAT);
+    }
+
+    private void testFatalErrorInTxnOffsetCommit(final Errors error) {
+        doInitTransactions();
 
         transactionManager.beginTransaction();
+
+        Map offsets = new HashMap<>();
+        offsets.put(tp0, new OffsetAndMetadata(1));
+        offsets.put(tp1, new OffsetAndMetadata(1));
+
+        TransactionalRequestResult addOffsetsResult = transactionManager.sendOffsetsToTransaction(
+            offsets, new ConsumerGroupMetadata(consumerGroupId));
+        prepareAddOffsetsToTxnResponse(Errors.NONE, consumerGroupId, producerId, epoch);
+        runUntil(() -> !client.hasPendingResponses());
+        assertThat(addOffsetsResult.isCompleted(), is(false));  // The request should complete only after the TxnOffsetCommit completes.
+
+        Map txnOffsetCommitResponse = new HashMap<>();
+        txnOffsetCommitResponse.put(tp0, Errors.NONE);
+        txnOffsetCommitResponse.put(tp1, error);
+
+        prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.GROUP, consumerGroupId);
+        prepareTxnOffsetCommitResponse(consumerGroupId, producerId, epoch, txnOffsetCommitResponse);
+
+        runUntil(addOffsetsResult::isCompleted);
+        assertThat(addOffsetsResult.isSuccessful(), is(false));
+        assertThat(addOffsetsResult.error(), instanceOf(error.exception().getClass()));
+    }
+
+    @Test
+    public void shouldNotAddPartitionsToTransactionWhenTopicAuthorizationFailed() throws Exception {
+        doInitTransactions();
+
+        transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
+        transactionManager.maybeAddPartitionToTransaction(tp0);
+
+        Future responseFuture = appendToAccumulator(tp0);
+        assertFalse(responseFuture.isDone());
+        prepareAddPartitionsToTxn(tp0, Errors.TOPIC_AUTHORIZATION_FAILED);
+        runUntil(transactionManager::hasError);
+        assertFalse(transactionManager.transactionContainsPartition(tp0));
+    }
+
+    @Test
+    public void shouldNotSendAbortTxnRequestWhenOnlyAddPartitionsRequestFailed() {
+        doInitTransactions();
+
+        transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
 
-        prepareAddPartitionsToTxnResponse(Errors.TOPIC_AUTHORIZATION_FAILED, tp0, epoch, pid);
-        sender.run(time.milliseconds());  // Send AddPartitionsRequest
+        prepareAddPartitionsToTxnResponse(Errors.TOPIC_AUTHORIZATION_FAILED, tp0, epoch, producerId);
+        runUntil(() -> !client.hasPendingResponses());
 
         TransactionalRequestResult abortResult = transactionManager.beginAbort();
         assertFalse(abortResult.isCompleted());
 
-        sender.run(time.milliseconds());
-        assertTrue(abortResult.isCompleted());
+        runUntil(abortResult::isCompleted);
         assertTrue(abortResult.isSuccessful());
     }
 
     @Test
-    public void shouldNotSendAbortTxnRequestWhenOnlyAddOffsetsRequestFailed() throws Exception {
-        final long pid = 13131L;
-        final short epoch = 1;
-
-        doInitTransactions(pid, epoch);
+    public void shouldNotSendAbortTxnRequestWhenOnlyAddOffsetsRequestFailed() {
+        doInitTransactions();
 
         transactionManager.beginTransaction();
         Map offsets = new HashMap<>();
         offsets.put(tp1, new OffsetAndMetadata(1));
-        final String consumerGroupId = "myconsumergroup";
 
-        transactionManager.sendOffsetsToTransaction(offsets, consumerGroupId);
+        transactionManager.sendOffsetsToTransaction(offsets, new ConsumerGroupMetadata(consumerGroupId));
 
         TransactionalRequestResult abortResult = transactionManager.beginAbort();
 
-        prepareAddOffsetsToTxnResponse(Errors.GROUP_AUTHORIZATION_FAILED, consumerGroupId, pid, epoch);
-        sender.run(time.milliseconds());  // Send AddOffsetsToTxnRequest
-        assertFalse(abortResult.isCompleted());
-
-        sender.run(time.milliseconds());
+        prepareAddOffsetsToTxnResponse(Errors.GROUP_AUTHORIZATION_FAILED, consumerGroupId, producerId, epoch);
+        runUntil(abortResult::isCompleted);
         assertTrue(transactionManager.isReady());
         assertTrue(abortResult.isCompleted());
         assertTrue(abortResult.isSuccessful());
     }
 
     @Test
-    public void shouldFailAbortIfAddOffsetsFailsWithFatalError() throws Exception {
-        final long pid = 13131L;
-        final short epoch = 1;
-
-        doInitTransactions(pid, epoch);
+    public void shouldFailAbortIfAddOffsetsFailsWithFatalError() {
+        doInitTransactions();
 
         transactionManager.beginTransaction();
         Map offsets = new HashMap<>();
         offsets.put(tp1, new OffsetAndMetadata(1));
-        final String consumerGroupId = "myconsumergroup";
 
-        transactionManager.sendOffsetsToTransaction(offsets, consumerGroupId);
+        transactionManager.sendOffsetsToTransaction(offsets, new ConsumerGroupMetadata(consumerGroupId));
 
         TransactionalRequestResult abortResult = transactionManager.beginAbort();
 
-        prepareAddOffsetsToTxnResponse(Errors.UNKNOWN_SERVER_ERROR, consumerGroupId, pid, epoch);
-        sender.run(time.milliseconds());  // Send AddOffsetsToTxnRequest
-        assertFalse(abortResult.isCompleted());
+        prepareAddOffsetsToTxnResponse(Errors.UNKNOWN_SERVER_ERROR, consumerGroupId, producerId, epoch);
 
-        sender.run(time.milliseconds());
-        assertTrue(abortResult.isCompleted());
+        runUntil(abortResult::isCompleted);
         assertFalse(abortResult.isSuccessful());
         assertTrue(transactionManager.hasFatalError());
     }
 
+    @Test
+    public void testSendOffsetsWithGroupMetadata() {
+        Map txnOffsetCommitResponse = new HashMap<>();
+        txnOffsetCommitResponse.put(tp0, Errors.NONE);
+        txnOffsetCommitResponse.put(tp1, Errors.COORDINATOR_LOAD_IN_PROGRESS);
+
+        TransactionalRequestResult addOffsetsResult = prepareGroupMetadataCommit(
+            () -> prepareTxnOffsetCommitResponse(consumerGroupId, producerId,
+                epoch, groupInstanceId, memberId, generationId, txnOffsetCommitResponse));
+
+        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, producerId, epoch, groupInstanceId, memberId, generationId, txnOffsetCommitResponse);
+        sender.runOnce();  // Send TxnOffsetCommitRequest again.
+
+        assertTrue(addOffsetsResult.isCompleted());
+        assertTrue(addOffsetsResult.isSuccessful());
+    }
+
+    @Test
+    public void testSendOffsetWithGroupMetadataFailAsAutoDowngradeTxnCommitNotEnabled() {
+        client.setNodeApiVersions(NodeApiVersions.create(ApiKeys.TXN_OFFSET_COMMIT.id, (short) 0, (short) 2));
+
+        Map txnOffsetCommitResponse = new HashMap<>();
+        txnOffsetCommitResponse.put(tp0, Errors.NONE);
+        txnOffsetCommitResponse.put(tp1, Errors.COORDINATOR_LOAD_IN_PROGRESS);
+
+        prepareGroupMetadataCommit(
+            () -> prepareTxnOffsetCommitResponse(consumerGroupId, producerId, epoch, txnOffsetCommitResponse));
+
+        assertThrows(UnsupportedVersionException.class, () -> sender.runOnce());
+    }
+
+    @Test
+    public void testSendOffsetWithGroupMetadataSuccessAsAutoDowngradeTxnCommitEnabled() {
+        initializeTransactionManager(Optional.of(transactionalId), true);
+
+        client.setNodeApiVersions(NodeApiVersions.create(ApiKeys.TXN_OFFSET_COMMIT.id, (short) 0, (short) 2));
+
+        Map txnOffsetCommitResponse = new HashMap<>();
+        txnOffsetCommitResponse.put(tp0, Errors.NONE);
+        txnOffsetCommitResponse.put(tp1, Errors.COORDINATOR_LOAD_IN_PROGRESS);
+
+        TransactionalRequestResult addOffsetsResult = prepareGroupMetadataCommit(
+            () -> prepareTxnOffsetCommitResponse(consumerGroupId, producerId, epoch, txnOffsetCommitResponse));
+
+        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, producerId, epoch, txnOffsetCommitResponse);
+        sender.runOnce();  // Send TxnOffsetCommitRequest again.
+
+        assertTrue(addOffsetsResult.isCompleted());
+        assertTrue(addOffsetsResult.isSuccessful());
+    }
+
+    private TransactionalRequestResult prepareGroupMetadataCommit(Runnable prepareTxnCommitResponse) {
+        doInitTransactions();
+
+        transactionManager.beginTransaction();
+        Map offsets = new HashMap<>();
+        offsets.put(tp0, new OffsetAndMetadata(1));
+        offsets.put(tp1, new OffsetAndMetadata(1));
+
+        TransactionalRequestResult addOffsetsResult = transactionManager.sendOffsetsToTransaction(
+            offsets, new ConsumerGroupMetadata(consumerGroupId, generationId, memberId, Optional.of(groupInstanceId)));
+        prepareAddOffsetsToTxnResponse(Errors.NONE, consumerGroupId, producerId, epoch);
+
+        sender.runOnce();  // send AddOffsetsToTxnResult
+
+        assertFalse(addOffsetsResult.isCompleted());  // The request should complete only after the TxnOffsetCommit completes
+
+        prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.GROUP, consumerGroupId);
+//        prepareTxnOffsetCommitResponse(consumerGroupId, producerId, epoch, groupInstanceId, memberId, generationId, txnOffsetCommitResponse);
+        prepareTxnCommitResponse.run();
+
+        assertNull(transactionManager.coordinator(CoordinatorType.GROUP));
+        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());
+        return addOffsetsResult;
+    }
+
     @Test
     public void testNoDrainWhenPartitionsPending() throws InterruptedException {
-        final long pid = 13131L;
-        final short epoch = 1;
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp0);
-        accumulator.append(tp0, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT);
+        appendToAccumulator(tp0);
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp1);
-        accumulator.append(tp1, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT);
+        appendToAccumulator(tp1);
 
         assertFalse(transactionManager.isSendToPartitionAllowed(tp0));
         assertFalse(transactionManager.isSendToPartitionAllowed(tp1));
@@ -1873,7 +2292,7 @@ public void testNoDrainWhenPartitionsPending() throws InterruptedException {
         PartitionInfo part2 = new PartitionInfo(topic, 1, node2, null, null);
 
         Cluster cluster = new Cluster(null, Arrays.asList(node1, node2), Arrays.asList(part1, part2),
-                Collections.emptySet(), Collections.emptySet());
+                Collections.emptySet(), Collections.emptySet());
         Set nodes = new HashSet<>();
         nodes.add(node1);
         nodes.add(node2);
@@ -1890,30 +2309,25 @@ public void testNoDrainWhenPartitionsPending() throws InterruptedException {
 
     @Test
     public void testAllowDrainInAbortableErrorState() throws InterruptedException {
-        final long pid = 13131L;
-        final short epoch = 1;
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
         transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
         transactionManager.maybeAddPartitionToTransaction(tp1);
         prepareAddPartitionsToTxn(tp1, Errors.NONE);
-        sender.run(time.milliseconds());  // Send AddPartitions, tp1 should be in the transaction now.
-
-        assertTrue(transactionManager.transactionContainsPartition(tp1));
+        runUntil(() -> 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.
-
-        assertTrue(transactionManager.hasAbortableError());
+        runUntil(transactionManager::hasAbortableError);
         assertTrue(transactionManager.isSendToPartitionAllowed(tp1));
 
         // Try to drain a message destined for tp1, it should get drained.
         Node node1 = new Node(1, "localhost", 1112);
         PartitionInfo part1 = new PartitionInfo(topic, 1, node1, null, null);
-        Cluster cluster = new Cluster(null, Arrays.asList(node1), Arrays.asList(part1),
-                Collections.emptySet(), Collections.emptySet());
-        accumulator.append(tp1, time.milliseconds(), "key".getBytes(),
-                "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT);
+        Cluster cluster = new Cluster(null, Collections.singletonList(node1), Collections.singletonList(part1),
+                Collections.emptySet(), Collections.emptySet());
+        appendToAccumulator(tp1);
         Map> drainedBatches = accumulator.drain(cluster, Collections.singleton(node1),
                 Integer.MAX_VALUE,
                 time.milliseconds());
@@ -1927,18 +2341,15 @@ public void testAllowDrainInAbortableErrorState() throws InterruptedException {
 
     @Test
     public void testRaiseErrorWhenNoPartitionsPendingOnDrain() throws InterruptedException {
-        final long pid = 13131L;
-        final short epoch = 1;
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
         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);
+        appendToAccumulator(tp0);
         Node node1 = new Node(0, "localhost", 1111);
         PartitionInfo part1 = new PartitionInfo(topic, 0, node1, null, null);
 
-        Cluster cluster = new Cluster(null, Arrays.asList(node1), Arrays.asList(part1),
-                Collections.emptySet(), Collections.emptySet());
+        Cluster cluster = new Cluster(null, Collections.singletonList(node1), Collections.singletonList(part1),
+                Collections.emptySet(), Collections.emptySet());
         Set nodes = new HashSet<>();
         nodes.add(node1);
         Map> drainedBatches = accumulator.drain(cluster, nodes, Integer.MAX_VALUE,
@@ -1951,53 +2362,44 @@ public void testRaiseErrorWhenNoPartitionsPendingOnDrain() throws InterruptedExc
 
     @Test
     public void resendFailedProduceRequestAfterAbortableError() throws Exception {
-        final long pid = 13131L;
-        final short epoch = 1;
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
         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;
+        Future responseFuture = appendToAccumulator(tp0);
 
-        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
+        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, producerId);
+        prepareProduceResponse(Errors.NOT_LEADER_OR_FOLLOWER, producerId, epoch);
+        runUntil(() -> !client.hasPendingResponses());
 
         assertFalse(responseFuture.isDone());
 
         transactionManager.transitionToAbortableError(new KafkaException());
-        prepareProduceResponse(Errors.NONE, pid, epoch);
-        sender.run(time.milliseconds());
-
-        assertTrue(responseFuture.isDone());
+        prepareProduceResponse(Errors.NONE, producerId, epoch);
+        runUntil(responseFuture::isDone);
         assertNotNull(responseFuture.get()); // should throw the exception which caused the transaction to be aborted.
     }
 
     @Test
-    public void testTransitionToAbortableErrorOnBatchExpiry() throws InterruptedException, ExecutionException {
-        final long pid = 13131L;
-        final short epoch = 1;
-
-        doInitTransactions(pid, epoch);
+    public void testTransitionToAbortableErrorOnBatchExpiry() throws InterruptedException {
+        doInitTransactions();
 
         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;
+        Future responseFuture = appendToAccumulator(tp0);
 
         assertFalse(responseFuture.isDone());
 
-        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, pid);
+        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, producerId);
 
         assertFalse(transactionManager.transactionContainsPartition(tp0));
         assertFalse(transactionManager.isSendToPartitionAllowed(tp0));
-        sender.run(time.milliseconds());  // send addPartitions.
         // Check that only addPartitions was sent.
-        assertTrue(transactionManager.transactionContainsPartition(tp0));
+        runUntil(() -> transactionManager.transactionContainsPartition(tp0));
         assertTrue(transactionManager.isSendToPartitionAllowed(tp0));
         assertFalse(responseFuture.isDone());
 
@@ -2005,12 +2407,11 @@ public void testTransitionToAbortableErrorOnBatchExpiry() throws InterruptedExce
         time.sleep(10000);
         // Disconnect the target node for the pending produce request. This will ensure that sender will try to
         // expire the batch.
-        Node clusterNode = this.cluster.nodes().get(0);
+        Node clusterNode = metadata.fetch().nodes().get(0);
         client.disconnect(clusterNode.idString());
-        client.blackout(clusterNode, 100);
+        client.backoff(clusterNode, 100);
 
-        sender.run(time.milliseconds());  // We should try to flush the produce, but expire it instead without sending anything.
-        assertTrue(responseFuture.isDone());
+        runUntil(responseFuture::isDone);
 
         try {
             // make sure the produce was expired.
@@ -2023,20 +2424,17 @@ public void testTransitionToAbortableErrorOnBatchExpiry() throws InterruptedExce
     }
 
     @Test
-    public void testTransitionToAbortableErrorOnMultipleBatchExpiry() throws InterruptedException, ExecutionException {
-        final long pid = 13131L;
-        final short epoch = 1;
-
-        doInitTransactions(pid, epoch);
+    public void testTransitionToAbortableErrorOnMultipleBatchExpiry() throws InterruptedException {
+        doInitTransactions();
 
         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;
-        Future secondBatchResponse = accumulator.append(tp1, time.milliseconds(), "key".getBytes(),
-               "value".getBytes(), Record.EMPTY_HEADERS, null, MAX_BLOCK_TIMEOUT).future;
+        Future firstBatchResponse = appendToAccumulator(tp0);
+        Future secondBatchResponse = appendToAccumulator(tp1);
 
         assertFalse(firstBatchResponse.isDone());
         assertFalse(secondBatchResponse.isDone());
@@ -2048,9 +2446,8 @@ public void testTransitionToAbortableErrorOnMultipleBatchExpiry() throws Interru
 
         assertFalse(transactionManager.transactionContainsPartition(tp0));
         assertFalse(transactionManager.isSendToPartitionAllowed(tp0));
-        sender.run(time.milliseconds());  // send addPartitions.
         // Check that only addPartitions was sent.
-        assertTrue(transactionManager.transactionContainsPartition(tp0));
+        runUntil(() -> transactionManager.transactionContainsPartition(tp0));
         assertTrue(transactionManager.transactionContainsPartition(tp1));
         assertTrue(transactionManager.isSendToPartitionAllowed(tp1));
         assertTrue(transactionManager.isSendToPartitionAllowed(tp1));
@@ -2061,13 +2458,12 @@ public void testTransitionToAbortableErrorOnMultipleBatchExpiry() throws Interru
         time.sleep(10000);
         // Disconnect the target node for the pending produce request. This will ensure that sender will try to
         // expire the batch.
-        Node clusterNode = this.cluster.nodes().get(0);
+        Node clusterNode = metadata.fetch().nodes().get(0);
         client.disconnect(clusterNode.idString());
-        client.blackout(clusterNode, 100);
+        client.backoff(clusterNode, 100);
 
-        sender.run(time.milliseconds());  // We should try to flush the produce, but expire it instead without sending anything.
-        assertTrue(firstBatchResponse.isDone());
-        assertTrue(secondBatchResponse.isDone());
+        runUntil(firstBatchResponse::isDone);
+        runUntil(secondBatchResponse::isDone);
 
         try {
             // make sure the produce was expired.
@@ -2088,27 +2484,23 @@ public void testTransitionToAbortableErrorOnMultipleBatchExpiry() throws Interru
     }
 
     @Test
-    public void testDropCommitOnBatchExpiry() throws InterruptedException, ExecutionException {
-        final long pid = 13131L;
-        final short epoch = 1;
-
-        doInitTransactions(pid, epoch);
+    public void testDropCommitOnBatchExpiry() throws InterruptedException {
+        doInitTransactions();
 
         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;
+        Future responseFuture = appendToAccumulator(tp0);
 
         assertFalse(responseFuture.isDone());
 
-        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, pid);
+        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, producerId);
 
         assertFalse(transactionManager.transactionContainsPartition(tp0));
         assertFalse(transactionManager.isSendToPartitionAllowed(tp0));
-        sender.run(time.milliseconds());  // send addPartitions.
         // Check that only addPartitions was sent.
-        assertTrue(transactionManager.transactionContainsPartition(tp0));
+        runUntil(() -> transactionManager.transactionContainsPartition(tp0));
         assertTrue(transactionManager.isSendToPartitionAllowed(tp0));
         assertFalse(responseFuture.isDone());
 
@@ -2118,12 +2510,10 @@ public void testDropCommitOnBatchExpiry() throws InterruptedException, Execution
         time.sleep(10000);
         // Disconnect the target node for the pending produce request. This will ensure that sender will try to
         // expire the batch.
-        Node clusterNode = this.cluster.nodes().get(0);
+        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.
-        assertTrue(responseFuture.isDone());
+        runUntil(responseFuture::isDone);  // We should try to flush the produce, but expire it instead without sending anything.
 
         try {
             // make sure the produce was expired.
@@ -2132,9 +2522,7 @@ public void testDropCommitOnBatchExpiry() throws InterruptedException, Execution
         } 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.
-
-        assertTrue(commitResult.isCompleted());
+        runUntil(commitResult::isCompleted);  // the commit shouldn't be completed without being sent since the produce request failed.
         assertFalse(commitResult.isSuccessful());  // the commit shouldn't succeed since the produce request failed.
 
         assertTrue(transactionManager.hasAbortableError());
@@ -2144,43 +2532,40 @@ public void testDropCommitOnBatchExpiry() throws InterruptedException, Execution
 
         TransactionalRequestResult abortResult = transactionManager.beginAbort();
 
-        prepareEndTxnResponse(Errors.NONE, TransactionResult.ABORT, pid, epoch);
-
-        sender.run(time.milliseconds());  // send the abort.
-
-        assertTrue(abortResult.isCompleted());
+        prepareEndTxnResponse(Errors.NONE, TransactionResult.ABORT, producerId, epoch);
+        prepareInitPidResponse(Errors.NONE, false, producerId, (short) (epoch + 1));
+        runUntil(abortResult::isCompleted);
         assertTrue(abortResult.isSuccessful());
         assertFalse(transactionManager.hasOngoingTransaction());
         assertFalse(transactionManager.transactionContainsPartition(tp0));
     }
 
     @Test
-    public void testTransitionToFatalErrorWhenRetriedBatchIsExpired() throws InterruptedException, ExecutionException {
-        final long pid = 13131L;
-        final short epoch = 1;
+    public void testTransitionToFatalErrorWhenRetriedBatchIsExpired() throws InterruptedException {
+        apiVersions.update("0", new NodeApiVersions(Arrays.asList(
+                new ApiVersion(ApiKeys.INIT_PRODUCER_ID.id, (short) 0, (short) 1),
+                new ApiVersion(ApiKeys.PRODUCE.id, (short) 0, (short) 7))));
 
-        doInitTransactions(pid, epoch);
+        doInitTransactions();
 
         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;
+        Future responseFuture = appendToAccumulator(tp0);
 
         assertFalse(responseFuture.isDone());
 
-        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, pid);
+        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, producerId);
 
         assertFalse(transactionManager.transactionContainsPartition(tp0));
         assertFalse(transactionManager.isSendToPartitionAllowed(tp0));
-        sender.run(time.milliseconds());  // send addPartitions.
         // Check that only addPartitions was sent.
-        assertTrue(transactionManager.transactionContainsPartition(tp0));
+        runUntil(() -> transactionManager.transactionContainsPartition(tp0));
         assertTrue(transactionManager.isSendToPartitionAllowed(tp0));
 
-        prepareProduceResponse(Errors.NOT_LEADER_FOR_PARTITION, pid, epoch);
-        sender.run(time.milliseconds());  // send the produce request.
-
+        prepareProduceResponse(Errors.NOT_LEADER_OR_FOLLOWER, producerId, epoch);
+        runUntil(() -> !client.hasPendingResponses());
         assertFalse(responseFuture.isDone());
 
         TransactionalRequestResult commitResult = transactionManager.beginCommit();
@@ -2189,12 +2574,11 @@ public void testTransitionToFatalErrorWhenRetriedBatchIsExpired() throws Interru
         time.sleep(10000);
         // Disconnect the target node for the pending produce request. This will ensure that sender will try to
         // expire the batch.
-        Node clusterNode = this.cluster.nodes().get(0);
+        Node clusterNode = metadata.fetch().nodes().get(0);
         client.disconnect(clusterNode.idString());
-        client.blackout(clusterNode, 100);
+        client.backoff(clusterNode, 100);
 
-        sender.run(time.milliseconds());  // We should try to flush the produce, but expire it instead without sending anything.
-        assertTrue(responseFuture.isDone());
+        runUntil(responseFuture::isDone);  // We should try to flush the produce, but expire it instead without sending anything.
 
         try {
             // make sure the produce was expired.
@@ -2203,42 +2587,725 @@ 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
-
-        assertTrue(commitResult.isCompleted());
+        runUntil(commitResult::isCompleted);
         assertFalse(commitResult.isSuccessful());  // the commit should have been dropped.
 
         assertTrue(transactionManager.hasFatalError());
         assertFalse(transactionManager.hasOngoingTransaction());
     }
 
-    private void verifyAddPartitionsFailsWithPartitionLevelError(final Errors error) throws InterruptedException {
-        final long pid = 1L;
-        final short epoch = 1;
+    @Test
+    public void testBumpEpochAfterTimeoutWithoutPendingInflightRequests() {
+        initializeTransactionManager(Optional.empty(), false);
+        long producerId = 15L;
+        short epoch = 5;
+        ProducerIdAndEpoch producerIdAndEpoch = new ProducerIdAndEpoch(producerId, epoch);
+        initializeIdempotentProducerId(producerId, epoch);
+
+        // Nothing to resolve, so no reset is needed
+        transactionManager.bumpIdempotentEpochAndResetIdIfNeeded();
+        assertEquals(producerIdAndEpoch, transactionManager.producerIdAndEpoch());
+
+        TopicPartition tp0 = new TopicPartition("foo", 0);
+        assertEquals(Integer.valueOf(0), transactionManager.sequenceNumber(tp0));
+
+        ProducerBatch b1 = writeIdempotentBatchWithValue(transactionManager, tp0, "1");
+        assertEquals(Integer.valueOf(1), transactionManager.sequenceNumber(tp0));
+        transactionManager.handleCompletedBatch(b1, new ProduceResponse.PartitionResponse(
+                Errors.NONE, 500L, time.milliseconds(), 0L));
+        assertEquals(OptionalInt.of(0), transactionManager.lastAckedSequence(tp0));
+
+        // Marking sequence numbers unresolved without inflight requests is basically a no-op.
+        transactionManager.markSequenceUnresolved(b1);
+        transactionManager.maybeResolveSequences();
+        assertEquals(producerIdAndEpoch, transactionManager.producerIdAndEpoch());
+        assertFalse(transactionManager.hasUnresolvedSequences());
+
+        // We have a new batch which fails with a timeout
+        ProducerBatch b2 = writeIdempotentBatchWithValue(transactionManager, tp0, "2");
+        assertEquals(Integer.valueOf(2), transactionManager.sequenceNumber(tp0));
+        transactionManager.markSequenceUnresolved(b2);
+        transactionManager.handleFailedBatch(b2, new TimeoutException(), false);
+        assertTrue(transactionManager.hasUnresolvedSequences());
+
+        // We only had one inflight batch, so we should be able to clear the unresolved status
+        // and bump the epoch
+        transactionManager.maybeResolveSequences();
+        assertFalse(transactionManager.hasUnresolvedSequences());
+
+        // Run sender loop to trigger epoch bump
+        runUntil(() -> transactionManager.producerIdAndEpoch().epoch == 6);
+    }
+
+    @Test
+    public void testNoProducerIdResetAfterLastInFlightBatchSucceeds() {
+        initializeTransactionManager(Optional.empty(), false);
+        long producerId = 15L;
+        short epoch = 5;
+        ProducerIdAndEpoch producerIdAndEpoch = new ProducerIdAndEpoch(producerId, epoch);
+        initializeIdempotentProducerId(producerId, epoch);
+
+        TopicPartition tp0 = new TopicPartition("foo", 0);
+        ProducerBatch b1 = writeIdempotentBatchWithValue(transactionManager, tp0, "1");
+        ProducerBatch b2 = writeIdempotentBatchWithValue(transactionManager, tp0, "2");
+        ProducerBatch b3 = writeIdempotentBatchWithValue(transactionManager, tp0, "3");
+        assertEquals(3, transactionManager.sequenceNumber(tp0).intValue());
+
+        // The first batch fails with a timeout
+        transactionManager.markSequenceUnresolved(b1);
+        transactionManager.handleFailedBatch(b1, new TimeoutException(), false);
+        assertTrue(transactionManager.hasUnresolvedSequences());
+
+        // The reset should not occur until sequence numbers have been resolved
+        transactionManager.bumpIdempotentEpochAndResetIdIfNeeded();
+        assertEquals(producerIdAndEpoch, transactionManager.producerIdAndEpoch());
+        assertTrue(transactionManager.hasUnresolvedSequences());
+
+        // The second batch fails as well with a timeout
+        transactionManager.handleFailedBatch(b2, new TimeoutException(), false);
+        transactionManager.bumpIdempotentEpochAndResetIdIfNeeded();
+        assertEquals(producerIdAndEpoch, transactionManager.producerIdAndEpoch());
+        assertTrue(transactionManager.hasUnresolvedSequences());
+
+        // The third batch succeeds, which should resolve the sequence number without
+        // requiring a producerId reset.
+        transactionManager.handleCompletedBatch(b3, new ProduceResponse.PartitionResponse(
+                Errors.NONE, 500L, time.milliseconds(), 0L));
+        transactionManager.maybeResolveSequences();
+        assertEquals(producerIdAndEpoch, transactionManager.producerIdAndEpoch());
+        assertFalse(transactionManager.hasUnresolvedSequences());
+        assertEquals(3, transactionManager.sequenceNumber(tp0).intValue());
+    }
+
+    @Test
+    public void testEpochBumpAfterLastInflightBatchFails() {
+        initializeTransactionManager(Optional.empty(), false);
+        ProducerIdAndEpoch producerIdAndEpoch = new ProducerIdAndEpoch(producerId, epoch);
+        initializeIdempotentProducerId(producerId, epoch);
+
+        TopicPartition tp0 = new TopicPartition("foo", 0);
+        ProducerBatch b1 = writeIdempotentBatchWithValue(transactionManager, tp0, "1");
+        ProducerBatch b2 = writeIdempotentBatchWithValue(transactionManager, tp0, "2");
+        ProducerBatch b3 = writeIdempotentBatchWithValue(transactionManager, tp0, "3");
+        assertEquals(Integer.valueOf(3), transactionManager.sequenceNumber(tp0));
+
+        // The first batch fails with a timeout
+        transactionManager.markSequenceUnresolved(b1);
+        transactionManager.handleFailedBatch(b1, new TimeoutException(), false);
+        assertTrue(transactionManager.hasUnresolvedSequences());
+
+        // The second batch succeeds, but sequence numbers are still not resolved
+        transactionManager.handleCompletedBatch(b2, new ProduceResponse.PartitionResponse(
+                Errors.NONE, 500L, time.milliseconds(), 0L));
+        transactionManager.bumpIdempotentEpochAndResetIdIfNeeded();
+        assertEquals(producerIdAndEpoch, transactionManager.producerIdAndEpoch());
+        assertTrue(transactionManager.hasUnresolvedSequences());
+
+        // When the last inflight batch fails, we have to bump the epoch
+        transactionManager.handleFailedBatch(b3, new TimeoutException(), false);
+
+        // Run sender loop to trigger epoch bump
+        runUntil(() -> transactionManager.producerIdAndEpoch().epoch == 2);
+        assertFalse(transactionManager.hasUnresolvedSequences());
+        assertEquals(0, transactionManager.sequenceNumber(tp0).intValue());
+    }
+
+    @Test
+    public void testNoFailedBatchHandlingWhenTxnManagerIsInFatalError() {
+        initializeTransactionManager(Optional.empty(), false);
+        long producerId = 15L;
+        short epoch = 5;
+        initializeIdempotentProducerId(producerId, epoch);
+
+        TopicPartition tp0 = new TopicPartition("foo", 0);
+        ProducerBatch b1 = writeIdempotentBatchWithValue(transactionManager, tp0, "1");
+        // Handling b1 should bump the epoch after OutOfOrderSequenceException
+        transactionManager.handleFailedBatch(b1, new OutOfOrderSequenceException("out of sequence"), false);
+        transactionManager.bumpIdempotentEpochAndResetIdIfNeeded();
+        ProducerIdAndEpoch idAndEpochAfterFirstBatch = new ProducerIdAndEpoch(producerId, (short) (epoch + 1));
+        assertEquals(idAndEpochAfterFirstBatch, transactionManager.producerIdAndEpoch());
+
+        transactionManager.transitionToFatalError(new KafkaException());
+
+        // The second batch should not bump the epoch as txn manager is already in fatal error state
+        ProducerBatch b2 = writeIdempotentBatchWithValue(transactionManager, tp0, "2");
+        transactionManager.handleFailedBatch(b2, new TimeoutException(), true);
+        transactionManager.bumpIdempotentEpochAndResetIdIfNeeded();
+        assertEquals(idAndEpochAfterFirstBatch, transactionManager.producerIdAndEpoch());
+    }
+
+    @Test
+    public void testAbortTransactionAndReuseSequenceNumberOnError() throws InterruptedException {
+        apiVersions.update("0", new NodeApiVersions(Arrays.asList(
+                new ApiVersion(ApiKeys.INIT_PRODUCER_ID.id, (short) 0, (short) 1),
+                new ApiVersion(ApiKeys.PRODUCE.id, (short) 0, (short) 7))));
+
+        doInitTransactions();
+
+        transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
+        transactionManager.maybeAddPartitionToTransaction(tp0);
+
+        Future responseFuture0 = appendToAccumulator(tp0);
+        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, producerId);
+        prepareProduceResponse(Errors.NONE, producerId, epoch);
+        runUntil(() -> transactionManager.isPartitionAdded(tp0));  // Send AddPartitionsRequest
+        runUntil(responseFuture0::isDone);
+
+        Future responseFuture1 = appendToAccumulator(tp0);
+        prepareProduceResponse(Errors.NONE, producerId, epoch);
+        runUntil(responseFuture1::isDone);
+
+        Future responseFuture2 = appendToAccumulator(tp0);
+        prepareProduceResponse(Errors.TOPIC_AUTHORIZATION_FAILED, producerId, epoch);
+        runUntil(responseFuture2::isDone); // Receive abortable error
+
+        assertTrue(transactionManager.hasAbortableError());
+
+        TransactionalRequestResult abortResult = transactionManager.beginAbort();
+        prepareEndTxnResponse(Errors.NONE, TransactionResult.ABORT, producerId, epoch);
+        runUntil(abortResult::isCompleted);
+        assertTrue(abortResult.isSuccessful());
+        assertTrue(transactionManager.isReady());  // make sure we are ready for a transaction now.
+
+        transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
+        transactionManager.maybeAddPartitionToTransaction(tp0);
+
+        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, producerId);
+        runUntil(() -> transactionManager.isPartitionAdded(tp0));  // Send AddPartitionsRequest
+
+        assertEquals(2, transactionManager.sequenceNumber(tp0).intValue());
+    }
+
+    @Test
+    public void testAbortTransactionAndResetSequenceNumberOnUnknownProducerId() throws InterruptedException {
+        // Set the InitProducerId version such that bumping the epoch number is not supported. This will test the case
+        // where the sequence number is reset on an UnknownProducerId error, allowing subsequent transactions to
+        // append to the log successfully
+        apiVersions.update("0", new NodeApiVersions(Arrays.asList(
+                new ApiVersion(ApiKeys.INIT_PRODUCER_ID.id, (short) 0, (short) 1),
+                new ApiVersion(ApiKeys.PRODUCE.id, (short) 0, (short) 7))));
+
+        doInitTransactions();
+
+        transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
+
+        transactionManager.maybeAddPartitionToTransaction(tp1);
+        Future successPartitionResponseFuture = appendToAccumulator(tp1);
+        prepareAddPartitionsToTxnResponse(Errors.NONE, tp1, epoch, producerId);
+        prepareProduceResponse(Errors.NONE, producerId, epoch, tp1);
+        runUntil(successPartitionResponseFuture::isDone);
+        assertTrue(transactionManager.isPartitionAdded(tp1));
+
+        transactionManager.maybeAddPartitionToTransaction(tp0);
+        Future responseFuture0 = appendToAccumulator(tp0);
+        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, producerId);
+        prepareProduceResponse(Errors.NONE, producerId, epoch);
+        runUntil(responseFuture0::isDone);
+        assertTrue(transactionManager.isPartitionAdded(tp0));
+
+        Future responseFuture1 = appendToAccumulator(tp0);
+        prepareProduceResponse(Errors.NONE, producerId, epoch);
+        runUntil(responseFuture1::isDone);
+
+        Future responseFuture2 = appendToAccumulator(tp0);
+        client.prepareResponse(produceRequestMatcher(producerId, epoch, tp0),
+                produceResponse(tp0, 0, Errors.UNKNOWN_PRODUCER_ID, 0, 0));
+        runUntil(responseFuture2::isDone);
+
+        assertTrue(transactionManager.hasAbortableError());
+
+        TransactionalRequestResult abortResult = transactionManager.beginAbort();
+        prepareEndTxnResponse(Errors.NONE, TransactionResult.ABORT, producerId, epoch);
+        runUntil(abortResult::isCompleted);
+        assertTrue(abortResult.isSuccessful());
+        assertTrue(transactionManager.isReady());  // make sure we are ready for a transaction now.
+
+        transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
+        transactionManager.maybeAddPartitionToTransaction(tp0);
+
+        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, producerId);
+        runUntil(() -> transactionManager.isPartitionAdded(tp0));
+
+        assertEquals(0, transactionManager.sequenceNumber(tp0).intValue());
+        assertEquals(1, transactionManager.sequenceNumber(tp1).intValue());
+    }
+
+    @Test
+    public void testBumpTransactionalEpochOnAbortableError() throws InterruptedException {
+        final short initialEpoch = 1;
+        final short bumpedEpoch = initialEpoch + 1;
 
-        doInitTransactions(pid, epoch);
+        doInitTransactions(producerId, initialEpoch);
 
         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;
+        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, initialEpoch, producerId);
+        runUntil(() -> transactionManager.isPartitionAdded(tp0));
+
+        Future responseFuture0 = appendToAccumulator(tp0);
+        prepareProduceResponse(Errors.NONE, producerId, initialEpoch);
+        runUntil(responseFuture0::isDone);
+
+        Future responseFuture1 = appendToAccumulator(tp0);
+        prepareProduceResponse(Errors.NONE, producerId, initialEpoch);
+        runUntil(responseFuture1::isDone);
+
+        Future responseFuture2 = appendToAccumulator(tp0);
+        prepareProduceResponse(Errors.TOPIC_AUTHORIZATION_FAILED, producerId, initialEpoch);
+        runUntil(responseFuture2::isDone);
+
+        assertTrue(transactionManager.hasAbortableError());
+        TransactionalRequestResult abortResult = transactionManager.beginAbort();
+
+        prepareEndTxnResponse(Errors.NONE, TransactionResult.ABORT, producerId, initialEpoch);
+        prepareInitPidResponse(Errors.NONE, false, producerId, bumpedEpoch);
+        runUntil(() -> transactionManager.producerIdAndEpoch().epoch == bumpedEpoch);
+
+        assertTrue(abortResult.isCompleted());
+        assertTrue(abortResult.isSuccessful());
+        assertTrue(transactionManager.isReady());  // make sure we are ready for a transaction now.
+
+        transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
+        transactionManager.maybeAddPartitionToTransaction(tp0);
+
+        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, bumpedEpoch, producerId);
+        runUntil(() -> transactionManager.isPartitionAdded(tp0));
+
+        assertEquals(0, transactionManager.sequenceNumber(tp0).intValue());
+    }
+
+    @Test
+    public void testBumpTransactionalEpochOnUnknownProducerIdError() throws InterruptedException {
+        final short initialEpoch = 1;
+        final short bumpedEpoch = 2;
+
+        doInitTransactions(producerId, initialEpoch);
+
+        transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
+        transactionManager.maybeAddPartitionToTransaction(tp0);
+
+        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, initialEpoch, producerId);
+        runUntil(() -> transactionManager.isPartitionAdded(tp0));
+
+        Future responseFuture0 = appendToAccumulator(tp0);
+        prepareProduceResponse(Errors.NONE, producerId, initialEpoch);
+        runUntil(responseFuture0::isDone);
+
+        Future responseFuture1 = appendToAccumulator(tp0);
+        prepareProduceResponse(Errors.NONE, producerId, initialEpoch);
+        runUntil(responseFuture1::isDone);
+
+        Future responseFuture2 = appendToAccumulator(tp0);
+        client.prepareResponse(produceRequestMatcher(producerId, initialEpoch, tp0),
+                produceResponse(tp0, 0, Errors.UNKNOWN_PRODUCER_ID, 0, 0));
+        runUntil(responseFuture2::isDone);
+
+        assertTrue(transactionManager.hasAbortableError());
+        TransactionalRequestResult abortResult = transactionManager.beginAbort();
+
+        prepareEndTxnResponse(Errors.NONE, TransactionResult.ABORT, producerId, initialEpoch);
+        prepareInitPidResponse(Errors.NONE, false, producerId, bumpedEpoch);
+        runUntil(() -> transactionManager.producerIdAndEpoch().epoch == bumpedEpoch);
+
+        assertTrue(abortResult.isCompleted());
+        assertTrue(abortResult.isSuccessful());
+        assertTrue(transactionManager.isReady());  // make sure we are ready for a transaction now.
+
+        transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
+        transactionManager.maybeAddPartitionToTransaction(tp0);
+
+        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, bumpedEpoch, producerId);
+        runUntil(() -> transactionManager.isPartitionAdded(tp0));
+
+        assertEquals(0, transactionManager.sequenceNumber(tp0).intValue());
+    }
+
+    @Test
+    public void testBumpTransactionalEpochOnTimeout() throws InterruptedException {
+        final short initialEpoch = 1;
+        final short bumpedEpoch = 2;
+
+        doInitTransactions(producerId, initialEpoch);
+
+        transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
+        transactionManager.maybeAddPartitionToTransaction(tp0);
+
+        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, initialEpoch, producerId);
+        runUntil(() -> transactionManager.isPartitionAdded(tp0));
+
+        Future responseFuture0 = appendToAccumulator(tp0);
+        prepareProduceResponse(Errors.NONE, producerId, initialEpoch);
+        runUntil(responseFuture0::isDone);
+
+        Future responseFuture1 = appendToAccumulator(tp0);
+        prepareProduceResponse(Errors.NONE, producerId, initialEpoch);
+        runUntil(responseFuture1::isDone);
+
+        Future responseFuture2 = appendToAccumulator(tp0);
+        runUntil(client::hasInFlightRequests); // Send Produce Request
+
+        // Sleep 10 seconds to make sure that the batches in the queue would be expired if they can't be drained.
+        time.sleep(10000);
+        // Disconnect the target node for the pending produce request. This will ensure that sender will try to
+        // expire the batch.
+        Node clusterNode = metadata.fetch().nodes().get(0);
+        client.disconnect(clusterNode.idString());
+        client.backoff(clusterNode, 100);
+
+        runUntil(responseFuture2::isDone); // We should try to flush the produce, but expire it instead without sending anything.
+
+        assertTrue(transactionManager.hasAbortableError());
+        TransactionalRequestResult abortResult = transactionManager.beginAbort();
+
+        sender.runOnce();  // handle the abort
+        time.sleep(110);  // Sleep to make sure the node backoff period has passed
+
+        prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.TRANSACTION, transactionalId);
+        prepareEndTxnResponse(Errors.NONE, TransactionResult.ABORT, producerId, initialEpoch);
+        prepareInitPidResponse(Errors.NONE, false, producerId, bumpedEpoch);
+        runUntil(() -> transactionManager.producerIdAndEpoch().epoch == bumpedEpoch);
+
+        assertTrue(abortResult.isCompleted());
+        assertTrue(abortResult.isSuccessful());
+        assertTrue(transactionManager.isReady());  // make sure we are ready for a transaction now.
+
+        transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
+        transactionManager.maybeAddPartitionToTransaction(tp0);
+
+        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, bumpedEpoch, producerId);
+        runUntil(() -> transactionManager.isPartitionAdded(tp0));
+
+        assertEquals(0, transactionManager.sequenceNumber(tp0).intValue());
+    }
+
+    @Test
+    public void testBumpTransactionalEpochOnRecoverableAddPartitionRequestError() {
+        final short initialEpoch = 1;
+        final short bumpedEpoch = 2;
+
+        doInitTransactions(producerId, initialEpoch);
+
+        transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
+        transactionManager.maybeAddPartitionToTransaction(tp0);
+        prepareAddPartitionsToTxnResponse(Errors.INVALID_PRODUCER_ID_MAPPING, tp0, initialEpoch, producerId);
+        runUntil(transactionManager::hasAbortableError);
+        TransactionalRequestResult abortResult = transactionManager.beginAbort();
+
+        prepareInitPidResponse(Errors.NONE, false, producerId, bumpedEpoch);
+        runUntil(abortResult::isCompleted);
+        assertEquals(bumpedEpoch, transactionManager.producerIdAndEpoch().epoch);
+        assertTrue(abortResult.isSuccessful());
+        assertTrue(transactionManager.isReady());  // make sure we are ready for a transaction now.
+    }
+
+    @Test
+    public void testBumpTransactionalEpochOnRecoverableAddOffsetsRequestError() throws InterruptedException {
+        final short initialEpoch = 1;
+        final short bumpedEpoch = 2;
+
+        doInitTransactions(producerId, initialEpoch);
+
+        transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
+        transactionManager.maybeAddPartitionToTransaction(tp0);
+
+        Future responseFuture = appendToAccumulator(tp0);
+
         assertFalse(responseFuture.isDone());
-        prepareAddPartitionsToTxn(tp0, error);
-        sender.run(time.milliseconds());  // attempt send addPartitions.
-        assertTrue(transactionManager.hasError());
-        assertFalse(transactionManager.transactionContainsPartition(tp0));
+        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, initialEpoch, producerId);
+        prepareProduceResponse(Errors.NONE, producerId, initialEpoch);
+        runUntil(responseFuture::isDone);
+
+        Map offsets = new HashMap<>();
+        offsets.put(tp0, new OffsetAndMetadata(1));
+        transactionManager.sendOffsetsToTransaction(offsets, new ConsumerGroupMetadata(consumerGroupId));
+        assertFalse(transactionManager.hasPendingOffsetCommits());
+        prepareAddOffsetsToTxnResponse(Errors.INVALID_PRODUCER_ID_MAPPING, consumerGroupId, producerId, initialEpoch);
+        runUntil(transactionManager::hasAbortableError);  // Send AddOffsetsRequest
+        TransactionalRequestResult abortResult = transactionManager.beginAbort();
+
+        prepareEndTxnResponse(Errors.NONE, TransactionResult.ABORT, producerId, initialEpoch);
+        prepareInitPidResponse(Errors.NONE, false, producerId, bumpedEpoch);
+        runUntil(abortResult::isCompleted);
+        assertEquals(bumpedEpoch, transactionManager.producerIdAndEpoch().epoch);
+        assertTrue(abortResult.isSuccessful());
+        assertTrue(transactionManager.isReady());  // make sure we are ready for a transaction now.
+    }
+
+    @Test
+    public void testHealthyPartitionRetriesDuringEpochBump() throws InterruptedException {
+        // Use a custom Sender to allow multiple inflight requests
+        initializeTransactionManager(Optional.empty(), false);
+        Sender sender = new Sender(logContext, this.client, this.metadata, this.accumulator, false,
+                MAX_REQUEST_SIZE, ACKS_ALL, MAX_RETRIES, new SenderMetricsRegistry(new Metrics(time)), this.time,
+                REQUEST_TIMEOUT, 50, transactionManager, apiVersions);
+        initializeIdempotentProducerId(producerId, epoch);
+
+        ProducerBatch tp0b1 = writeIdempotentBatchWithValue(transactionManager, tp0, "1");
+        ProducerBatch tp0b2 = writeIdempotentBatchWithValue(transactionManager, tp0, "2");
+        writeIdempotentBatchWithValue(transactionManager, tp0, "3");
+        ProducerBatch tp1b1 = writeIdempotentBatchWithValue(transactionManager, tp1, "4");
+        ProducerBatch tp1b2 = writeIdempotentBatchWithValue(transactionManager, tp1, "5");
+        assertEquals(3, transactionManager.sequenceNumber(tp0).intValue());
+        assertEquals(2, transactionManager.sequenceNumber(tp1).intValue());
+
+        // First batch of each partition succeeds
+        long b1AppendTime = time.milliseconds();
+        ProduceResponse.PartitionResponse t0b1Response = new ProduceResponse.PartitionResponse(
+                Errors.NONE, 500L, b1AppendTime, 0L);
+        tp0b1.done(500L, b1AppendTime, null);
+        transactionManager.handleCompletedBatch(tp0b1, t0b1Response);
+
+        ProduceResponse.PartitionResponse t1b1Response = new ProduceResponse.PartitionResponse(
+                Errors.NONE, 500L, b1AppendTime, 0L);
+        tp1b1.done(500L, b1AppendTime, null);
+        transactionManager.handleCompletedBatch(tp1b1, t1b1Response);
+
+        // We bump the epoch and set sequence numbers back to 0
+        ProduceResponse.PartitionResponse t0b2Response = new ProduceResponse.PartitionResponse(
+                Errors.UNKNOWN_PRODUCER_ID, -1, -1, 500L);
+        assertTrue(transactionManager.canRetry(t0b2Response, tp0b2));
+
+        // Run sender loop to trigger epoch bump
+        runUntil(() -> transactionManager.producerIdAndEpoch().epoch == 2);
+
+        // tp0 batches should have had sequence and epoch rewritten, but tp1 batches should not
+        assertEquals(tp0b2, transactionManager.nextBatchBySequence(tp0));
+        assertEquals(0, transactionManager.firstInFlightSequence(tp0));
+        assertEquals(0, tp0b2.baseSequence());
+        assertTrue(tp0b2.sequenceHasBeenReset());
+        assertEquals(2, tp0b2.producerEpoch());
+
+        assertEquals(tp1b2, transactionManager.nextBatchBySequence(tp1));
+        assertEquals(1, transactionManager.firstInFlightSequence(tp1));
+        assertEquals(1, tp1b2.baseSequence());
+        assertFalse(tp1b2.sequenceHasBeenReset());
+        assertEquals(1, tp1b2.producerEpoch());
+
+        // New tp1 batches should not be drained from the accumulator while tp1 has in-flight requests using the old epoch
+        appendToAccumulator(tp1);
+        sender.runOnce();
+        assertEquals(1, accumulator.batches().get(tp1).size());
+
+        // Partition failover occurs and tp1 returns a NOT_LEADER_OR_FOLLOWER error
+        // Despite having the old epoch, the batch should retry
+        ProduceResponse.PartitionResponse t1b2Response = new ProduceResponse.PartitionResponse(
+                Errors.NOT_LEADER_OR_FOLLOWER, -1, -1, 600L);
+        assertTrue(transactionManager.canRetry(t1b2Response, tp1b2));
+        accumulator.reenqueue(tp1b2, time.milliseconds());
+
+        // The batch with the old epoch should be successfully drained, leaving the new one in the queue
+        sender.runOnce();
+        assertEquals(1, accumulator.batches().get(tp1).size());
+        assertNotEquals(tp1b2, accumulator.batches().get(tp1).peek());
+        assertEquals(epoch, tp1b2.producerEpoch());
+
+        // After successfully retrying, there should be no in-flight batches for tp1 and the sequence should be 0
+        t1b2Response = new ProduceResponse.PartitionResponse(
+                Errors.NONE, 500L, b1AppendTime, 0L);
+        tp1b2.done(500L, b1AppendTime, null);
+        transactionManager.handleCompletedBatch(tp1b2, t1b2Response);
+
+        assertFalse(transactionManager.hasInflightBatches(tp1));
+        assertEquals(0, transactionManager.sequenceNumber(tp1).intValue());
+
+        // The last batch should now be drained and sent
+        runUntil(() -> transactionManager.hasInflightBatches(tp1));
+        assertTrue(accumulator.batches().get(tp1).isEmpty());
+        ProducerBatch tp1b3 = transactionManager.nextBatchBySequence(tp1);
+        assertEquals(epoch + 1, tp1b3.producerEpoch());
+
+        ProduceResponse.PartitionResponse t1b3Response = new ProduceResponse.PartitionResponse(
+                Errors.NONE, 500L, b1AppendTime, 0L);
+        tp1b3.done(500L, b1AppendTime, null);
+        transactionManager.handleCompletedBatch(tp1b3, t1b3Response);
+
+        assertFalse(transactionManager.hasInflightBatches(tp1));
+        assertEquals(1, transactionManager.sequenceNumber(tp1).intValue());
+    }
+
+    @Test
+    public void testRetryAbortTransaction() throws InterruptedException {
+        verifyCommitOrAbortTransactionRetriable(TransactionResult.ABORT, TransactionResult.ABORT);
+    }
+
+    @Test
+    public void testRetryCommitTransaction() throws InterruptedException {
+        verifyCommitOrAbortTransactionRetriable(TransactionResult.COMMIT, TransactionResult.COMMIT);
+    }
+
+    @Test(expected = KafkaException.class)
+    public void testRetryAbortTransactionAfterCommitTimeout() throws InterruptedException {
+        verifyCommitOrAbortTransactionRetriable(TransactionResult.COMMIT, TransactionResult.ABORT);
+    }
+
+    @Test(expected = KafkaException.class)
+    public void testRetryCommitTransactionAfterAbortTimeout() throws InterruptedException {
+        verifyCommitOrAbortTransactionRetriable(TransactionResult.ABORT, TransactionResult.COMMIT);
+    }
+
+    @Test
+    public void testCanBumpEpochDuringCoordinatorDisconnect() {
+        doInitTransactions(0, (short) 0);
+        runUntil(() -> transactionManager.coordinator(CoordinatorType.TRANSACTION) != null);
+        assertTrue(transactionManager.canBumpEpoch());
+
+        apiVersions.remove(transactionManager.coordinator(CoordinatorType.TRANSACTION).idString());
+        assertTrue(transactionManager.canBumpEpoch());
+    }
+
+    @Test
+    public void testFailedInflightBatchAfterEpochBump() throws InterruptedException {
+        // Use a custom Sender to allow multiple inflight requests
+        initializeTransactionManager(Optional.empty(), false);
+        Sender sender = new Sender(logContext, this.client, this.metadata, this.accumulator, false,
+                MAX_REQUEST_SIZE, ACKS_ALL, MAX_RETRIES, new SenderMetricsRegistry(new Metrics(time)), this.time,
+                REQUEST_TIMEOUT, 50, transactionManager, apiVersions);
+        initializeIdempotentProducerId(producerId, epoch);
+
+        ProducerBatch tp0b1 = writeIdempotentBatchWithValue(transactionManager, tp0, "1");
+        ProducerBatch tp0b2 = writeIdempotentBatchWithValue(transactionManager, tp0, "2");
+        writeIdempotentBatchWithValue(transactionManager, tp0, "3");
+        ProducerBatch tp1b1 = writeIdempotentBatchWithValue(transactionManager, tp1, "4");
+        ProducerBatch tp1b2 = writeIdempotentBatchWithValue(transactionManager, tp1, "5");
+        assertEquals(3, transactionManager.sequenceNumber(tp0).intValue());
+        assertEquals(2, transactionManager.sequenceNumber(tp1).intValue());
+
+        // First batch of each partition succeeds
+        long b1AppendTime = time.milliseconds();
+        ProduceResponse.PartitionResponse t0b1Response = new ProduceResponse.PartitionResponse(
+                Errors.NONE, 500L, b1AppendTime, 0L);
+        tp0b1.done(500L, b1AppendTime, null);
+        transactionManager.handleCompletedBatch(tp0b1, t0b1Response);
+
+        ProduceResponse.PartitionResponse t1b1Response = new ProduceResponse.PartitionResponse(
+                Errors.NONE, 500L, b1AppendTime, 0L);
+        tp1b1.done(500L, b1AppendTime, null);
+        transactionManager.handleCompletedBatch(tp1b1, t1b1Response);
+
+        // We bump the epoch and set sequence numbers back to 0
+        ProduceResponse.PartitionResponse t0b2Response = new ProduceResponse.PartitionResponse(
+                Errors.UNKNOWN_PRODUCER_ID, -1, -1, 500L);
+        assertTrue(transactionManager.canRetry(t0b2Response, tp0b2));
+
+        // Run sender loop to trigger epoch bump
+        runUntil(() -> transactionManager.producerIdAndEpoch().epoch == 2);
+
+        // tp0 batches should have had sequence and epoch rewritten, but tp1 batches should not
+        assertEquals(tp0b2, transactionManager.nextBatchBySequence(tp0));
+        assertEquals(0, transactionManager.firstInFlightSequence(tp0));
+        assertEquals(0, tp0b2.baseSequence());
+        assertTrue(tp0b2.sequenceHasBeenReset());
+        assertEquals(2, tp0b2.producerEpoch());
+
+        assertEquals(tp1b2, transactionManager.nextBatchBySequence(tp1));
+        assertEquals(1, transactionManager.firstInFlightSequence(tp1));
+        assertEquals(1, tp1b2.baseSequence());
+        assertFalse(tp1b2.sequenceHasBeenReset());
+        assertEquals(1, tp1b2.producerEpoch());
+
+        // New tp1 batches should not be drained from the accumulator while tp1 has in-flight requests using the old epoch
+        appendToAccumulator(tp1);
+        sender.runOnce();
+        assertEquals(1, accumulator.batches().get(tp1).size());
+
+        // Partition failover occurs and tp1 returns a NOT_LEADER_OR_FOLLOWER error
+        // Despite having the old epoch, the batch should retry
+        ProduceResponse.PartitionResponse t1b2Response = new ProduceResponse.PartitionResponse(
+                Errors.NOT_LEADER_OR_FOLLOWER, -1, -1, 600L);
+        assertTrue(transactionManager.canRetry(t1b2Response, tp1b2));
+        accumulator.reenqueue(tp1b2, time.milliseconds());
+
+        // The batch with the old epoch should be successfully drained, leaving the new one in the queue
+        sender.runOnce();
+        assertEquals(1, accumulator.batches().get(tp1).size());
+        assertNotEquals(tp1b2, accumulator.batches().get(tp1).peek());
+        assertEquals(epoch, tp1b2.producerEpoch());
+
+        // After successfully retrying, there should be no in-flight batches for tp1 and the sequence should be 0
+        t1b2Response = new ProduceResponse.PartitionResponse(
+                Errors.NONE, 500L, b1AppendTime, 0L);
+        tp1b2.done(500L, b1AppendTime, null);
+        transactionManager.handleCompletedBatch(tp1b2, t1b2Response);
+
+        assertFalse(transactionManager.hasInflightBatches(tp1));
+        assertEquals(0, transactionManager.sequenceNumber(tp1).intValue());
+
+        // The last batch should now be drained and sent
+        runUntil(() -> transactionManager.hasInflightBatches(tp1));
+        assertTrue(accumulator.batches().get(tp1).isEmpty());
+        ProducerBatch tp1b3 = transactionManager.nextBatchBySequence(tp1);
+        assertEquals(epoch + 1, tp1b3.producerEpoch());
+
+        ProduceResponse.PartitionResponse t1b3Response = new ProduceResponse.PartitionResponse(
+                Errors.NONE, 500L, b1AppendTime, 0L);
+        tp1b3.done(500L, b1AppendTime, null);
+        transactionManager.handleCompletedBatch(tp1b3, t1b3Response);
+
+        assertFalse(transactionManager.hasInflightBatches(tp1));
+        assertEquals(1, transactionManager.sequenceNumber(tp1).intValue());
+    }
+
+    private FutureRecordMetadata appendToAccumulator(TopicPartition tp) throws InterruptedException {
+        final long nowMs = time.milliseconds();
+        return accumulator.append(tp, nowMs, "key".getBytes(), "value".getBytes(), Record.EMPTY_HEADERS,
+                null, MAX_BLOCK_TIMEOUT, false, nowMs).future;
+    }
+
+    private void verifyCommitOrAbortTransactionRetriable(TransactionResult firstTransactionResult,
+                                                         TransactionResult retryTransactionResult) throws InterruptedException {
+        doInitTransactions();
+
+        transactionManager.beginTransaction();
+        transactionManager.failIfNotReadyForSend();
+        transactionManager.maybeAddPartitionToTransaction(tp0);
+
+        appendToAccumulator(tp0);
+
+        prepareAddPartitionsToTxnResponse(Errors.NONE, tp0, epoch, producerId);
+        prepareProduceResponse(Errors.NONE, producerId, epoch);
+        runUntil(() -> !client.hasPendingResponses());
+
+        TransactionalRequestResult result = firstTransactionResult == TransactionResult.COMMIT ?
+                transactionManager.beginCommit() : transactionManager.beginAbort();
+        prepareEndTxnResponse(Errors.NONE, firstTransactionResult, producerId, epoch, true);
+        runUntil(() -> !client.hasPendingResponses());
+        assertFalse(result.isCompleted());
+
+        try {
+            result.await(MAX_BLOCK_TIMEOUT, TimeUnit.MILLISECONDS);
+            fail("Should have raised TimeoutException");
+        } catch (TimeoutException ignored) {
+        }
+
+        prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.TRANSACTION, transactionalId);
+        runUntil(() -> !client.hasPendingResponses());
+        TransactionalRequestResult retryResult = retryTransactionResult == TransactionResult.COMMIT ?
+                transactionManager.beginCommit() : transactionManager.beginAbort();
+        assertEquals(retryResult, result); // check if cached result is reused.
+
+        prepareEndTxnResponse(Errors.NONE, retryTransactionResult, producerId, epoch, false);
+        runUntil(retryResult::isCompleted);
+        assertFalse(transactionManager.hasOngoingTransaction());
     }
 
     private void prepareAddPartitionsToTxn(final Map errors) {
-        client.prepareResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                AddPartitionsToTxnRequest request = (AddPartitionsToTxnRequest) body;
-                assertEquals(new HashSet<>(request.partitions()), new HashSet<>(errors.keySet()));
-                return true;
-            }
+        client.prepareResponse(body -> {
+            AddPartitionsToTxnRequest request = (AddPartitionsToTxnRequest) body;
+            assertEquals(new HashSet<>(request.partitions()), new HashSet<>(errors.keySet()));
+            return true;
         }, new AddPartitionsToTxnResponse(0, errors));
     }
 
@@ -2249,156 +3316,220 @@ private void prepareAddPartitionsToTxn(final TopicPartition tp, final Errors err
     private void prepareFindCoordinatorResponse(Errors error, boolean shouldDisconnect,
                                                 final CoordinatorType coordinatorType,
                                                 final String coordinatorKey) {
-        client.prepareResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                FindCoordinatorRequest findCoordinatorRequest = (FindCoordinatorRequest) body;
-                assertEquals(findCoordinatorRequest.coordinatorType(), coordinatorType);
-                assertEquals(findCoordinatorRequest.coordinatorKey(), coordinatorKey);
-                return true;
-            }
-        }, new FindCoordinatorResponse(error, brokerNode), shouldDisconnect);
-    }
-
-    private void prepareInitPidResponse(Errors error, boolean shouldDisconnect, long pid, short epoch) {
-        client.prepareResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                InitProducerIdRequest initProducerIdRequest = (InitProducerIdRequest) body;
-                assertEquals(initProducerIdRequest.transactionalId(), transactionalId);
-                assertEquals(initProducerIdRequest.transactionTimeoutMs(), transactionTimeoutMs);
-                return true;
-            }
-        }, new InitProducerIdResponse(0, error, pid, epoch), shouldDisconnect);
-    }
-
-    private void sendProduceResponse(Errors error, final long pid, final short epoch) {
-        client.respond(produceRequestMatcher(pid, epoch), 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 MockClient.RequestMatcher produceRequestMatcher(final long pid, final short epoch) {
-        return new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                ProduceRequest produceRequest = (ProduceRequest) body;
-                MemoryRecords records = produceRequest.partitionRecordsOrFail().get(tp0);
-                assertNotNull(records);
-                Iterator batchIterator = records.batches().iterator();
-                assertTrue(batchIterator.hasNext());
-                MutableRecordBatch batch = batchIterator.next();
-                assertFalse(batchIterator.hasNext());
-                assertTrue(batch.isTransactional());
-                assertEquals(pid, batch.producerId());
-                assertEquals(epoch, batch.producerEpoch());
-                assertEquals(transactionalId, produceRequest.transactionalId());
-                return true;
-            }
+        client.prepareResponse(body -> {
+            FindCoordinatorRequest findCoordinatorRequest = (FindCoordinatorRequest) body;
+            assertEquals(CoordinatorType.forId(findCoordinatorRequest.data().keyType()), coordinatorType);
+            assertEquals(findCoordinatorRequest.data().key(), coordinatorKey);
+            return true;
+        }, FindCoordinatorResponse.prepareResponse(error, brokerNode), shouldDisconnect);
+    }
+
+    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(transactionalId, initProducerIdRequest.data().transactionalId());
+            assertEquals(transactionTimeoutMs, initProducerIdRequest.data().transactionTimeoutMs());
+            return true;
+        }, new InitProducerIdResponse(responseData), shouldDisconnect);
+    }
+
+    private void sendProduceResponse(Errors error, final long producerId, final short producerEpoch) {
+        sendProduceResponse(error, producerId, producerEpoch, tp0);
+    }
+
+    private void sendProduceResponse(Errors error, final long producerId, final short producerEpoch, TopicPartition tp) {
+        client.respond(produceRequestMatcher(producerId, producerEpoch, tp), produceResponse(tp, 0, error, 0));
+    }
+
+    private void prepareProduceResponse(Errors error, final long producerId, final short producerEpoch) {
+        prepareProduceResponse(error, producerId, producerEpoch, tp0);
+    }
+
+    private void prepareProduceResponse(Errors error, final long producerId, final short producerEpoch, TopicPartition tp) {
+        client.prepareResponse(produceRequestMatcher(producerId, producerEpoch, tp), produceResponse(tp, 0, error, 0));
+    }
+
+    private MockClient.RequestMatcher produceRequestMatcher(final long producerId, final short epoch, TopicPartition tp) {
+        return body -> {
+            ProduceRequest produceRequest = (ProduceRequest) body;
+            MemoryRecords records = produceRequest.data().topicData()
+                    .stream()
+                    .filter(t -> t.name().equals(tp.topic()))
+                    .findAny()
+                    .get()
+                    .partitionData()
+                    .stream()
+                    .filter(p -> p.index() == tp.partition())
+                    .map(p -> (MemoryRecords) p.records())
+                    .findAny().get();
+            assertNotNull(records);
+            Iterator batchIterator = records.batches().iterator();
+            assertTrue(batchIterator.hasNext());
+            MutableRecordBatch batch = batchIterator.next();
+            assertFalse(batchIterator.hasNext());
+            assertTrue(batch.isTransactional());
+            assertEquals(producerId, batch.producerId());
+            assertEquals(epoch, batch.producerEpoch());
+            assertEquals(transactionalId, produceRequest.transactionalId());
+            return true;
         };
     }
 
     private void prepareAddPartitionsToTxnResponse(Errors error, final TopicPartition topicPartition,
-                                                   final short epoch, final long pid) {
-        client.prepareResponse(addPartitionsRequestMatcher(topicPartition, epoch, pid),
+                                                   final short epoch, final long producerId) {
+        client.prepareResponse(addPartitionsRequestMatcher(topicPartition, epoch, producerId),
                 new AddPartitionsToTxnResponse(0, singletonMap(topicPartition, error)));
     }
 
     private void sendAddPartitionsToTxnResponse(Errors error, final TopicPartition topicPartition,
-                                                final short epoch, final long pid) {
-        client.respond(addPartitionsRequestMatcher(topicPartition, epoch, pid),
+                                                final short epoch, final long producerId) {
+        client.respond(addPartitionsRequestMatcher(topicPartition, epoch, producerId),
                 new AddPartitionsToTxnResponse(0, singletonMap(topicPartition, error)));
     }
 
     private MockClient.RequestMatcher addPartitionsRequestMatcher(final TopicPartition topicPartition,
-                                                                  final short epoch, final long pid) {
-        return new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                AddPartitionsToTxnRequest addPartitionsToTxnRequest = (AddPartitionsToTxnRequest) body;
-                assertEquals(pid, addPartitionsToTxnRequest.producerId());
-                assertEquals(epoch, addPartitionsToTxnRequest.producerEpoch());
-                assertEquals(singletonList(topicPartition), addPartitionsToTxnRequest.partitions());
-                assertEquals(transactionalId, addPartitionsToTxnRequest.transactionalId());
-                return true;
-            }
+                                                                  final short epoch, final long producerId) {
+        return body -> {
+            AddPartitionsToTxnRequest addPartitionsToTxnRequest = (AddPartitionsToTxnRequest) body;
+            assertEquals(producerId, addPartitionsToTxnRequest.data().producerId());
+            assertEquals(epoch, addPartitionsToTxnRequest.data().producerEpoch());
+            assertEquals(singletonList(topicPartition), addPartitionsToTxnRequest.partitions());
+            assertEquals(transactionalId, addPartitionsToTxnRequest.data().transactionalId());
+            return true;
         };
     }
 
-    private void prepareEndTxnResponse(Errors error, final TransactionResult result, final long pid, final short epoch) {
-        client.prepareResponse(endTxnMatcher(result, pid, epoch), new EndTxnResponse(0, error));
+    private void prepareEndTxnResponse(Errors error, final TransactionResult result, final long producerId, final short epoch) {
+        this.prepareEndTxnResponse(error, result, producerId, epoch, false);
     }
 
-    private void sendEndTxnResponse(Errors error, final TransactionResult result, final long pid, final short epoch) {
-        client.respond(endTxnMatcher(result, pid, epoch), new EndTxnResponse(0, error));
+    private void prepareEndTxnResponse(Errors error,
+                                       final TransactionResult result,
+                                       final long producerId,
+                                       final short epoch,
+                                       final boolean shouldDisconnect) {
+        client.prepareResponse(endTxnMatcher(result, producerId, epoch),
+                               new EndTxnResponse(new EndTxnResponseData()
+                                                      .setErrorCode(error.code())
+                                                      .setThrottleTimeMs(0)), shouldDisconnect);
     }
 
-    private MockClient.RequestMatcher endTxnMatcher(final TransactionResult result, final long pid, final short epoch) {
-        return new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                EndTxnRequest endTxnRequest = (EndTxnRequest) body;
-                assertEquals(transactionalId, endTxnRequest.transactionalId());
-                assertEquals(pid, endTxnRequest.producerId());
-                assertEquals(epoch, endTxnRequest.producerEpoch());
-                assertEquals(result, endTxnRequest.command());
-                return true;
-            }
+    private void sendEndTxnResponse(Errors error, final TransactionResult result, final long producerId, final short epoch) {
+        client.respond(endTxnMatcher(result, producerId, epoch), new EndTxnResponse(
+            new EndTxnResponseData()
+                .setErrorCode(error.code())
+                .setThrottleTimeMs(0)
+        ));
+    }
+
+    private MockClient.RequestMatcher endTxnMatcher(final TransactionResult result, final long producerId, final short epoch) {
+        return body -> {
+            EndTxnRequest endTxnRequest = (EndTxnRequest) body;
+            assertEquals(transactionalId, endTxnRequest.data().transactionalId());
+            assertEquals(producerId, endTxnRequest.data().producerId());
+            assertEquals(epoch, endTxnRequest.data().producerEpoch());
+            assertEquals(result, endTxnRequest.result());
+            return true;
         };
     }
 
-    private void prepareAddOffsetsToTxnResponse(Errors error, final String consumerGroupId, final long producerId,
+    private void prepareAddOffsetsToTxnResponse(final Errors error,
+                                                final String consumerGroupId,
+                                                final long producerId,
                                                 final short producerEpoch) {
-        client.prepareResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                AddOffsetsToTxnRequest addOffsetsToTxnRequest = (AddOffsetsToTxnRequest) body;
-                assertEquals(consumerGroupId, addOffsetsToTxnRequest.consumerGroupId());
-                assertEquals(transactionalId, addOffsetsToTxnRequest.transactionalId());
-                assertEquals(producerId, addOffsetsToTxnRequest.producerId());
-                assertEquals(producerEpoch, addOffsetsToTxnRequest.producerEpoch());
-                return true;
-            }
-        }, new AddOffsetsToTxnResponse(0, error));
-    }
-
-    private void prepareTxnOffsetCommitResponse(final String consumerGroupId, final long producerId,
-                                                final short producerEpoch, Map txnOffsetCommitResponse) {
-        client.prepareResponse(new MockClient.RequestMatcher() {
-            @Override
-            public boolean matches(AbstractRequest body) {
-                TxnOffsetCommitRequest txnOffsetCommitRequest = (TxnOffsetCommitRequest) body;
-                assertEquals(consumerGroupId, txnOffsetCommitRequest.consumerGroupId());
-                assertEquals(producerId, txnOffsetCommitRequest.producerId());
-                assertEquals(producerEpoch, txnOffsetCommitRequest.producerEpoch());
-                return true;
-            }
+        client.prepareResponse(body -> {
+            AddOffsetsToTxnRequest addOffsetsToTxnRequest = (AddOffsetsToTxnRequest) body;
+            assertEquals(consumerGroupId, addOffsetsToTxnRequest.data().groupId());
+            assertEquals(transactionalId, addOffsetsToTxnRequest.data().transactionalId());
+            assertEquals(producerId, addOffsetsToTxnRequest.data().producerId());
+            assertEquals(producerEpoch, addOffsetsToTxnRequest.data().producerEpoch());
+            return true;
+        }, new AddOffsetsToTxnResponse(
+            new AddOffsetsToTxnResponseData()
+                .setErrorCode(error.code()))
+        );
+    }
+
+    private void prepareTxnOffsetCommitResponse(final String consumerGroupId,
+                                                final long producerId,
+                                                final short producerEpoch,
+                                                Map txnOffsetCommitResponse) {
+        client.prepareResponse(request -> {
+            TxnOffsetCommitRequest txnOffsetCommitRequest = (TxnOffsetCommitRequest) request;
+            assertEquals(consumerGroupId, txnOffsetCommitRequest.data().groupId());
+            assertEquals(producerId, txnOffsetCommitRequest.data().producerId());
+            assertEquals(producerEpoch, txnOffsetCommitRequest.data().producerEpoch());
+            return true;
         }, new TxnOffsetCommitResponse(0, txnOffsetCommitResponse));
+    }
 
+    private void prepareTxnOffsetCommitResponse(final String consumerGroupId,
+                                                final long producerId,
+                                                final short producerEpoch,
+                                                final String groupInstanceId,
+                                                final String memberId,
+                                                final int generationId,
+                                                Map txnOffsetCommitResponse) {
+        client.prepareResponse(request -> {
+            TxnOffsetCommitRequest txnOffsetCommitRequest = (TxnOffsetCommitRequest) request;
+            assertEquals(consumerGroupId, txnOffsetCommitRequest.data().groupId());
+            assertEquals(producerId, txnOffsetCommitRequest.data().producerId());
+            assertEquals(producerEpoch, txnOffsetCommitRequest.data().producerEpoch());
+            assertEquals(groupInstanceId, txnOffsetCommitRequest.data().groupInstanceId());
+            assertEquals(memberId, txnOffsetCommitRequest.data().memberId());
+            assertEquals(generationId, txnOffsetCommitRequest.data().generationId());
+            return true;
+        }, new TxnOffsetCommitResponse(0, txnOffsetCommitResponse));
     }
 
     private ProduceResponse produceResponse(TopicPartition tp, long offset, Errors error, int throttleTimeMs) {
-        ProduceResponse.PartitionResponse resp = new ProduceResponse.PartitionResponse(error, offset, RecordBatch.NO_TIMESTAMP, 10);
+        return produceResponse(tp, offset, error, throttleTimeMs, 10);
+    }
+
+    @SuppressWarnings("deprecation")
+    private ProduceResponse produceResponse(TopicPartition tp, long offset, Errors error, int throttleTimeMs, int logStartOffset) {
+        ProduceResponse.PartitionResponse resp = new ProduceResponse.PartitionResponse(error, offset, RecordBatch.NO_TIMESTAMP, logStartOffset);
         Map partResp = singletonMap(tp, resp);
         return new ProduceResponse(partResp, throttleTimeMs);
     }
 
-    private void doInitTransactions(long pid, short epoch) {
+    private void initializeIdempotentProducerId(long producerId, short epoch) {
+        InitProducerIdResponseData responseData = new InitProducerIdResponseData()
+                .setErrorCode(Errors.NONE.code())
+                .setProducerEpoch(epoch)
+                .setProducerId(producerId)
+                .setThrottleTimeMs(0);
+        client.prepareResponse(body -> {
+            InitProducerIdRequest initProducerIdRequest = (InitProducerIdRequest) body;
+            assertNull(initProducerIdRequest.data().transactionalId());
+            return true;
+        }, new InitProducerIdResponse(responseData), false);
+
+        runUntil(transactionManager::hasProducerId);
+    }
+
+    private void doInitTransactions() {
+        doInitTransactions(producerId, epoch);
+    }
+
+    private void doInitTransactions(long producerId, short epoch) {
         transactionManager.initializeTransactions();
         prepareFindCoordinatorResponse(Errors.NONE, false, CoordinatorType.TRANSACTION, transactionalId);
-        sender.run(time.milliseconds());  // find coordinator
-        sender.run(time.milliseconds());
+        runUntil(() -> transactionManager.coordinator(CoordinatorType.TRANSACTION) != null);
         assertEquals(brokerNode, transactionManager.coordinator(CoordinatorType.TRANSACTION));
 
-        prepareInitPidResponse(Errors.NONE, false, pid, epoch);
-        sender.run(time.milliseconds());  // get pid.
-        assertTrue(transactionManager.hasProducerId());
+        prepareInitPidResponse(Errors.NONE, false, producerId, epoch);
+        runUntil(transactionManager::hasProducerId);
     }
 
     private void assertAbortableError(Class cause) {
         try {
-            transactionManager.beginTransaction();
+            transactionManager.beginCommit();
             fail("Should have raised " + cause.getSimpleName());
         } catch (KafkaException e) {
             assertTrue(cause.isAssignableFrom(e.getCause().getClass()));
@@ -2431,7 +3562,7 @@ private void assertFatalError(Class cause) {
         }
     }
 
-    private void assertFutureFailed(Future future) throws InterruptedException {
+    private void assertProduceFutureFailed(Future future) throws InterruptedException {
         assertTrue(future.isDone());
 
         try {
@@ -2442,4 +3573,14 @@ private void assertFutureFailed(Future future) throws Interrupte
         }
     }
 
+    private void runUntil(Supplier condition) {
+        for (int i = 0; i < 5; i++) {
+            if (condition.get())
+                break;
+            sender.runOnce();
+        }
+        if (!condition.get())
+            throw new AssertionError("Condition was not satisfied after multiple runs");
+    }
+
 }
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/KafkaFutureTest.java b/clients/src/test/java/org/apache/kafka/common/KafkaFutureTest.java
index 71f3c3c984ddd..37cfd1eefa18f 100644
--- a/clients/src/test/java/org/apache/kafka/common/KafkaFutureTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/KafkaFutureTest.java
@@ -67,7 +67,7 @@ public void testCompleteFutures() throws Exception {
     @Test
     public void testCompletingFutures() throws Exception {
         final KafkaFutureImpl future = new KafkaFutureImpl<>();
-        CompleterThread myThread = new CompleterThread(future, "You must construct additional pylons.");
+        CompleterThread myThread = new CompleterThread<>(future, "You must construct additional pylons.");
         assertFalse(future.isDone());
         assertFalse(future.isCompletedExceptionally());
         assertFalse(future.isCancelled());
@@ -83,6 +83,27 @@ public void testCompletingFutures() throws Exception {
         assertEquals(null, myThread.testException);
     }
 
+    @Test
+    public void testThenApply() throws Exception {
+        KafkaFutureImpl future = new KafkaFutureImpl<>();
+        KafkaFuture doubledFuture = future.thenApply(integer -> 2 * integer);
+        assertFalse(doubledFuture.isDone());
+        KafkaFuture tripledFuture = future.thenApply(integer -> 3 * integer);
+        assertFalse(tripledFuture.isDone());
+        future.complete(21);
+        assertEquals(Integer.valueOf(21), future.getNow(-1));
+        assertEquals(Integer.valueOf(42), doubledFuture.getNow(-1));
+        assertEquals(Integer.valueOf(63), tripledFuture.getNow(-1));
+        KafkaFuture quadrupledFuture = future.thenApply(integer -> 4 * integer);
+        assertEquals(Integer.valueOf(84), quadrupledFuture.getNow(-1));
+
+        KafkaFutureImpl futureFail = new KafkaFutureImpl<>();
+        KafkaFuture futureAppliedFail = futureFail.thenApply(integer -> 2 * integer);
+        futureFail.completeExceptionally(new RuntimeException());
+        assertTrue(futureFail.isCompletedExceptionally());
+        assertTrue(futureAppliedFail.isCompletedExceptionally());
+    }
+
     private static class CompleterThread extends Thread {
 
         private final KafkaFutureImpl future;
@@ -135,7 +156,7 @@ public void testAllOfFutures() throws Exception {
         final int numThreads = 5;
         final List> futures = new ArrayList<>();
         for (int i = 0; i < numThreads; i++) {
-            futures.add(new KafkaFutureImpl());
+            futures.add(new KafkaFutureImpl<>());
         }
         KafkaFuture allFuture = KafkaFuture.allOf(futures.toArray(new KafkaFuture[0]));
         final List completerThreads = new ArrayList<>();
diff --git a/clients/src/test/java/org/apache/kafka/common/SerializeCompatibilityTopicPartitionTest.java b/clients/src/test/java/org/apache/kafka/common/SerializeCompatibilityTopicPartitionTest.java
deleted file mode 100644
index 8b4df5ff01f6a..0000000000000
--- a/clients/src/test/java/org/apache/kafka/common/SerializeCompatibilityTopicPartitionTest.java
+++ /dev/null
@@ -1,64 +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;
-
-import org.apache.kafka.common.utils.Serializer;
-import org.junit.Test;
-
-import java.io.IOException;
-
-import static org.junit.Assert.assertTrue;
-import static org.junit.Assert.assertEquals;
-
-/**
- * This test ensures TopicPartition class is serializable and is serialization compatible.
- * Note: this ensures that the current code can deserialize data serialized with older versions of the code, but not the reverse.
- * That is, older code won't necessarily be able to deserialize data serialized with newer code.
- */
-public class SerializeCompatibilityTopicPartitionTest {
-
-    private String topicName = "mytopic";
-    private String fileName = "serializedData/topicPartitionSerializedfile";
-    private int partNum = 5;
-
-    private void checkValues(TopicPartition deSerTP) {
-        //assert deserialized values are same as original
-        assertEquals("partition number should be " + partNum + " but got " + deSerTP.partition(), partNum, deSerTP.partition());
-        assertEquals("topic should be " + topicName + " but got " + deSerTP.topic(), topicName, deSerTP.topic());
-    }
-
-    @Test
-    public void testSerializationRoundtrip() throws IOException, ClassNotFoundException {
-        //assert TopicPartition is serializable and deserialization renders the clone of original properly
-        TopicPartition origTp = new TopicPartition(topicName, partNum);
-        byte[] byteArray = Serializer.serialize(origTp);
-
-        //deserialize the byteArray and check if the values are same as original
-        Object deserializedObject = Serializer.deserialize(byteArray);
-        assertTrue(deserializedObject instanceof TopicPartition);
-        checkValues((TopicPartition) deserializedObject);
-    }
-
-    @Test
-    public void testTopiPartitionSerializationCompatibility() throws IOException, ClassNotFoundException {
-        // assert serialized TopicPartition object in file (serializedData/topicPartitionSerializedfile) is
-        // deserializable into TopicPartition and is compatible
-        Object deserializedObject = Serializer.deserialize(fileName);
-        assertTrue(deserializedObject instanceof TopicPartition);
-        checkValues((TopicPartition) deserializedObject);
-    }
-}
diff --git a/clients/src/test/java/org/apache/kafka/common/TopicPartitionTest.java b/clients/src/test/java/org/apache/kafka/common/TopicPartitionTest.java
new file mode 100644
index 0000000000000..2a90338c64e6c
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/TopicPartitionTest.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;
+
+import org.apache.kafka.common.utils.Serializer;
+import org.junit.Test;
+
+import java.io.IOException;
+
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.assertEquals;
+
+/**
+ * This test ensures TopicPartition class is serializable and is serialization compatible.
+ * Note: this ensures that the current code can deserialize data serialized with older versions of the code, but not the reverse.
+ * That is, older code won't necessarily be able to deserialize data serialized with newer code.
+ */
+public class TopicPartitionTest {
+    private String topicName = "mytopic";
+    private String fileName = "serializedData/topicPartitionSerializedfile";
+    private int partNum = 5;
+
+    private void checkValues(TopicPartition deSerTP) {
+        //assert deserialized values are same as original
+        assertEquals("partition number should be " + partNum + " but got " + deSerTP.partition(), partNum, deSerTP.partition());
+        assertEquals("topic should be " + topicName + " but got " + deSerTP.topic(), topicName, deSerTP.topic());
+    }
+
+    @Test
+    public void testSerializationRoundtrip() throws IOException, ClassNotFoundException {
+        //assert TopicPartition is serializable and deserialization renders the clone of original properly
+        TopicPartition origTp = new TopicPartition(topicName, partNum);
+        byte[] byteArray = Serializer.serialize(origTp);
+
+        //deserialize the byteArray and check if the values are same as original
+        Object deserializedObject = Serializer.deserialize(byteArray);
+        assertTrue(deserializedObject instanceof TopicPartition);
+        checkValues((TopicPartition) deserializedObject);
+    }
+
+    @Test
+    public void testTopiPartitionSerializationCompatibility() throws IOException, ClassNotFoundException {
+        // assert serialized TopicPartition object in file (serializedData/topicPartitionSerializedfile) is
+        // deserializable into TopicPartition and is compatible
+        Object deserializedObject = Serializer.deserialize(fileName);
+        assertTrue(deserializedObject instanceof TopicPartition);
+        checkValues((TopicPartition) deserializedObject);
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/UuidTest.java b/clients/src/test/java/org/apache/kafka/common/UuidTest.java
new file mode 100644
index 0000000000000..99e85bfe687ef
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/UuidTest.java
@@ -0,0 +1,81 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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 org.junit.Test;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotEquals;
+
+public class UuidTest {
+
+    @Test
+    public void testSignificantBits() {
+        Uuid id = new Uuid(34L, 98L);
+
+        assertEquals(id.getMostSignificantBits(), 34L);
+        assertEquals(id.getLeastSignificantBits(), 98L);
+    }
+
+    @Test
+    public void testUuidEquality() {
+        Uuid id1 = new Uuid(12L, 13L);
+        Uuid id2 = new Uuid(12L, 13L);
+        Uuid id3 = new Uuid(24L, 38L);
+
+        assertEquals(Uuid.ZERO_UUID, Uuid.ZERO_UUID);
+        assertEquals(id1, id2);
+        assertNotEquals(id1, id3);
+
+        assertEquals(Uuid.ZERO_UUID.hashCode(), Uuid.ZERO_UUID.hashCode());
+        assertEquals(id1.hashCode(), id2.hashCode());
+        assertNotEquals(id1.hashCode(), id3.hashCode());
+    }
+    
+    @Test
+    public void testHashCode() {
+        Uuid id1 = new Uuid(16L, 7L);
+        Uuid id2 = new Uuid(1043L, 20075L);
+        Uuid id3 = new Uuid(104312423523523L, 200732425676585L);
+        
+        assertEquals(23, id1.hashCode());
+        assertEquals(19064, id2.hashCode());
+        assertEquals(-2011255899, id3.hashCode());
+    }
+
+    @Test
+    public void testStringConversion() {
+        Uuid id = Uuid.randomUuid();
+        String idString = id.toString();
+
+        assertEquals(Uuid.fromString(idString), id);
+
+        String zeroIdString = Uuid.ZERO_UUID.toString();
+
+        assertEquals(Uuid.fromString(zeroIdString), Uuid.ZERO_UUID);
+    }
+
+    @Test
+    public void testRandomUuid() {
+        Uuid randomID = Uuid.randomUuid();
+        // reservedSentinel is based on the value of SENTINEL_ID_INTERNAL in Uuid.
+        Uuid reservedSentinel = new Uuid(0L, 1L);
+
+        assertNotEquals(randomID, Uuid.ZERO_UUID);
+        assertNotEquals(randomID, reservedSentinel);
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/acl/AclBindingTest.java b/clients/src/test/java/org/apache/kafka/common/acl/AclBindingTest.java
index 0ebcdfedb4fa0..461661fb21ac1 100644
--- a/clients/src/test/java/org/apache/kafka/common/acl/AclBindingTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/acl/AclBindingTest.java
@@ -16,68 +16,71 @@
  */
 package org.apache.kafka.common.acl;
 
-import org.apache.kafka.common.resource.Resource;
-import org.apache.kafka.common.resource.ResourceFilter;
+import org.apache.kafka.common.resource.PatternType;
+import org.apache.kafka.common.resource.ResourcePattern;
+import org.apache.kafka.common.resource.ResourcePatternFilter;
 import org.apache.kafka.common.resource.ResourceType;
 import org.junit.Test;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 
 public class AclBindingTest {
     private static final AclBinding ACL1 = new AclBinding(
-        new Resource(ResourceType.TOPIC, "mytopic"),
+        new ResourcePattern(ResourceType.TOPIC, "mytopic", PatternType.LITERAL),
         new AccessControlEntry("User:ANONYMOUS", "", AclOperation.ALL, AclPermissionType.ALLOW));
 
     private static final AclBinding ACL2 = new AclBinding(
-        new Resource(ResourceType.TOPIC, "mytopic"),
+        new ResourcePattern(ResourceType.TOPIC, "mytopic", PatternType.LITERAL),
         new AccessControlEntry("User:*", "", AclOperation.READ, AclPermissionType.ALLOW));
 
     private static final AclBinding ACL3 = new AclBinding(
-        new Resource(ResourceType.TOPIC, "mytopic2"),
+        new ResourcePattern(ResourceType.TOPIC, "mytopic2", PatternType.LITERAL),
         new AccessControlEntry("User:ANONYMOUS", "127.0.0.1", AclOperation.READ, AclPermissionType.DENY));
 
     private static final AclBinding UNKNOWN_ACL = new AclBinding(
-        new Resource(ResourceType.TOPIC, "mytopic2"),
+        new ResourcePattern(ResourceType.TOPIC, "mytopic2", PatternType.LITERAL),
         new AccessControlEntry("User:ANONYMOUS", "127.0.0.1", AclOperation.UNKNOWN, AclPermissionType.DENY));
 
     private static final AclBindingFilter ANY_ANONYMOUS = new AclBindingFilter(
-        new ResourceFilter(ResourceType.ANY, null),
+        ResourcePatternFilter.ANY,
         new AccessControlEntryFilter("User:ANONYMOUS", null, AclOperation.ANY, AclPermissionType.ANY));
 
     private static final AclBindingFilter ANY_DENY = new AclBindingFilter(
-        new ResourceFilter(ResourceType.ANY, null),
+        ResourcePatternFilter.ANY,
         new AccessControlEntryFilter(null, null, AclOperation.ANY, AclPermissionType.DENY));
 
     private static final AclBindingFilter ANY_MYTOPIC = new AclBindingFilter(
-        new ResourceFilter(ResourceType.TOPIC, "mytopic"),
+        new ResourcePatternFilter(ResourceType.TOPIC, "mytopic", PatternType.LITERAL),
         new AccessControlEntryFilter(null, null, AclOperation.ANY, AclPermissionType.ANY));
 
     @Test
-    public void testMatching() throws Exception {
-        assertTrue(ACL1.equals(ACL1));
+    public void testMatching() {
+        assertEquals(ACL1, ACL1);
         final AclBinding acl1Copy = new AclBinding(
-            new Resource(ResourceType.TOPIC, "mytopic"),
+            new ResourcePattern(ResourceType.TOPIC, "mytopic", PatternType.LITERAL),
             new AccessControlEntry("User:ANONYMOUS", "", AclOperation.ALL, AclPermissionType.ALLOW));
-        assertTrue(ACL1.equals(acl1Copy));
-        assertTrue(acl1Copy.equals(ACL1));
-        assertTrue(ACL2.equals(ACL2));
-        assertFalse(ACL1.equals(ACL2));
-        assertFalse(ACL2.equals(ACL1));
+        assertEquals(ACL1, acl1Copy);
+        assertEquals(acl1Copy, ACL1);
+        assertEquals(ACL2, ACL2);
+        assertNotEquals(ACL1, ACL2);
+        assertNotEquals(ACL2, ACL1);
         assertTrue(AclBindingFilter.ANY.matches(ACL1));
-        assertFalse(AclBindingFilter.ANY.equals(ACL1));
+        assertNotEquals(AclBindingFilter.ANY, ACL1);
         assertTrue(AclBindingFilter.ANY.matches(ACL2));
-        assertFalse(AclBindingFilter.ANY.equals(ACL2));
+        assertNotEquals(AclBindingFilter.ANY, ACL2);
         assertTrue(AclBindingFilter.ANY.matches(ACL3));
-        assertFalse(AclBindingFilter.ANY.equals(ACL3));
-        assertTrue(AclBindingFilter.ANY.equals(AclBindingFilter.ANY));
+        assertNotEquals(AclBindingFilter.ANY, ACL3);
+        assertEquals(AclBindingFilter.ANY, AclBindingFilter.ANY);
         assertTrue(ANY_ANONYMOUS.matches(ACL1));
-        assertFalse(ANY_ANONYMOUS.equals(ACL1));
+        assertNotEquals(ANY_ANONYMOUS, ACL1);
         assertFalse(ANY_ANONYMOUS.matches(ACL2));
-        assertFalse(ANY_ANONYMOUS.equals(ACL2));
+        assertNotEquals(ANY_ANONYMOUS, ACL2);
         assertTrue(ANY_ANONYMOUS.matches(ACL3));
-        assertFalse(ANY_ANONYMOUS.equals(ACL3));
+        assertNotEquals(ANY_ANONYMOUS, ACL3);
         assertFalse(ANY_DENY.matches(ACL1));
         assertFalse(ANY_DENY.matches(ACL2));
         assertTrue(ANY_DENY.matches(ACL3));
@@ -86,12 +89,12 @@ public void testMatching() throws Exception {
         assertFalse(ANY_MYTOPIC.matches(ACL3));
         assertTrue(ANY_ANONYMOUS.matches(UNKNOWN_ACL));
         assertTrue(ANY_DENY.matches(UNKNOWN_ACL));
-        assertTrue(UNKNOWN_ACL.equals(UNKNOWN_ACL));
+        assertEquals(UNKNOWN_ACL, UNKNOWN_ACL);
         assertFalse(ANY_MYTOPIC.matches(UNKNOWN_ACL));
     }
 
     @Test
-    public void testUnknowns() throws Exception {
+    public void testUnknowns() {
         assertFalse(ACL1.isUnknown());
         assertFalse(ACL2.isUnknown());
         assertFalse(ACL3.isUnknown());
@@ -102,12 +105,37 @@ public void testUnknowns() throws Exception {
     }
 
     @Test
-    public void testMatchesAtMostOne() throws Exception {
-        assertEquals(null, ACL1.toFilter().findIndefiniteField());
-        assertEquals(null, ACL2.toFilter().findIndefiniteField());
-        assertEquals(null, ACL3.toFilter().findIndefiniteField());
+    public void testMatchesAtMostOne() {
+        assertNull(ACL1.toFilter().findIndefiniteField());
+        assertNull(ACL2.toFilter().findIndefiniteField());
+        assertNull(ACL3.toFilter().findIndefiniteField());
         assertFalse(ANY_ANONYMOUS.matchesAtMostOne());
         assertFalse(ANY_DENY.matchesAtMostOne());
         assertFalse(ANY_MYTOPIC.matchesAtMostOne());
     }
+
+    @Test
+    public void shouldNotThrowOnUnknownPatternType() {
+        new AclBinding(new ResourcePattern(ResourceType.TOPIC, "foo", PatternType.UNKNOWN), ACL1.entry());
+    }
+
+    @Test
+    public void shouldNotThrowOnUnknownResourceType() {
+        new AclBinding(new ResourcePattern(ResourceType.UNKNOWN, "foo", PatternType.LITERAL), ACL1.entry());
+    }
+
+    @Test(expected = IllegalArgumentException.class)
+    public void shouldThrowOnMatchPatternType() {
+        new AclBinding(new ResourcePattern(ResourceType.TOPIC, "foo", PatternType.MATCH), ACL1.entry());
+    }
+
+    @Test(expected = IllegalArgumentException.class)
+    public void shouldThrowOnAnyPatternType() {
+        new AclBinding(new ResourcePattern(ResourceType.TOPIC, "foo", PatternType.ANY), ACL1.entry());
+    }
+
+    @Test(expected = IllegalArgumentException.class)
+    public void shouldThrowOnAnyResourceType() {
+        new AclBinding(new ResourcePattern(ResourceType.ANY, "foo", PatternType.LITERAL), ACL1.entry());
+    }
 }
diff --git a/clients/src/test/java/org/apache/kafka/common/acl/ResourcePatternFilterTest.java b/clients/src/test/java/org/apache/kafka/common/acl/ResourcePatternFilterTest.java
new file mode 100644
index 0000000000000..08d5a63450841
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/acl/ResourcePatternFilterTest.java
@@ -0,0 +1,176 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.acl;
+
+import org.apache.kafka.common.resource.PatternType;
+import org.apache.kafka.common.resource.ResourcePattern;
+import org.apache.kafka.common.resource.ResourcePatternFilter;
+import org.junit.Test;
+
+import static org.apache.kafka.common.resource.PatternType.LITERAL;
+import static org.apache.kafka.common.resource.PatternType.PREFIXED;
+import static org.apache.kafka.common.resource.ResourceType.ANY;
+import static org.apache.kafka.common.resource.ResourceType.GROUP;
+import static org.apache.kafka.common.resource.ResourceType.TOPIC;
+import static org.apache.kafka.common.resource.ResourceType.UNKNOWN;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+public class ResourcePatternFilterTest {
+    @Test
+    public void shouldBeUnknownIfResourceTypeUnknown() {
+        assertTrue(new ResourcePatternFilter(UNKNOWN, null, PatternType.LITERAL).isUnknown());
+    }
+
+    @Test
+    public void shouldBeUnknownIfPatternTypeUnknown() {
+        assertTrue(new ResourcePatternFilter(GROUP, null, PatternType.UNKNOWN).isUnknown());
+    }
+
+    @Test
+    public void shouldNotMatchIfDifferentResourceType() {
+        assertFalse(new ResourcePatternFilter(TOPIC, "Name", LITERAL)
+            .matches(new ResourcePattern(GROUP, "Name", LITERAL)));
+    }
+
+    @Test
+    public void shouldNotMatchIfDifferentName() {
+        assertFalse(new ResourcePatternFilter(TOPIC, "Different", PREFIXED)
+            .matches(new ResourcePattern(TOPIC, "Name", PREFIXED)));
+    }
+
+    @Test
+    public void shouldNotMatchIfDifferentNameCase() {
+        assertFalse(new ResourcePatternFilter(TOPIC, "NAME", LITERAL)
+            .matches(new ResourcePattern(TOPIC, "Name", LITERAL)));
+    }
+
+    @Test
+    public void shouldNotMatchIfDifferentPatternType() {
+        assertFalse(new ResourcePatternFilter(TOPIC, "Name", LITERAL)
+            .matches(new ResourcePattern(TOPIC, "Name", PREFIXED)));
+    }
+
+    @Test
+    public void shouldMatchWhereResourceTypeIsAny() {
+        assertTrue(new ResourcePatternFilter(ANY, "Name", PREFIXED)
+            .matches(new ResourcePattern(TOPIC, "Name", PREFIXED)));
+    }
+
+    @Test
+    public void shouldMatchWhereResourceNameIsAny() {
+        assertTrue(new ResourcePatternFilter(TOPIC, null, PREFIXED)
+            .matches(new ResourcePattern(TOPIC, "Name", PREFIXED)));
+    }
+
+    @Test
+    public void shouldMatchWherePatternTypeIsAny() {
+        assertTrue(new ResourcePatternFilter(TOPIC, null, PatternType.ANY)
+            .matches(new ResourcePattern(TOPIC, "Name", PREFIXED)));
+    }
+
+    @Test
+    public void shouldMatchWherePatternTypeIsMatch() {
+        assertTrue(new ResourcePatternFilter(TOPIC, null, PatternType.MATCH)
+            .matches(new ResourcePattern(TOPIC, "Name", PREFIXED)));
+    }
+
+    @Test
+    public void shouldMatchLiteralIfExactMatch() {
+        assertTrue(new ResourcePatternFilter(TOPIC, "Name", LITERAL)
+            .matches(new ResourcePattern(TOPIC, "Name", LITERAL)));
+    }
+
+    @Test
+    public void shouldMatchLiteralIfNameMatchesAndFilterIsOnPatternTypeAny() {
+        assertTrue(new ResourcePatternFilter(TOPIC, "Name", PatternType.ANY)
+            .matches(new ResourcePattern(TOPIC, "Name", LITERAL)));
+    }
+
+    @Test
+    public void shouldMatchLiteralIfNameMatchesAndFilterIsOnPatternTypeMatch() {
+        assertTrue(new ResourcePatternFilter(TOPIC, "Name", PatternType.MATCH)
+            .matches(new ResourcePattern(TOPIC, "Name", LITERAL)));
+    }
+
+    @Test
+    public void shouldNotMatchLiteralIfNamePrefixed() {
+        assertFalse(new ResourcePatternFilter(TOPIC, "Name-something", PatternType.MATCH)
+            .matches(new ResourcePattern(TOPIC, "Name", LITERAL)));
+    }
+
+    @Test
+    public void shouldMatchLiteralWildcardIfExactMatch() {
+        assertTrue(new ResourcePatternFilter(TOPIC, "*", LITERAL)
+            .matches(new ResourcePattern(TOPIC, "*", LITERAL)));
+    }
+
+    @Test
+    public void shouldNotMatchLiteralWildcardAgainstOtherName() {
+        assertFalse(new ResourcePatternFilter(TOPIC, "Name", LITERAL)
+            .matches(new ResourcePattern(TOPIC, "*", LITERAL)));
+    }
+
+    @Test
+    public void shouldNotMatchLiteralWildcardTheWayAround() {
+        assertFalse(new ResourcePatternFilter(TOPIC, "*", LITERAL)
+            .matches(new ResourcePattern(TOPIC, "Name", LITERAL)));
+    }
+
+    @Test
+    public void shouldNotMatchLiteralWildcardIfFilterHasPatternTypeOfAny() {
+        assertFalse(new ResourcePatternFilter(TOPIC, "Name", PatternType.ANY)
+            .matches(new ResourcePattern(TOPIC, "*", LITERAL)));
+    }
+
+    @Test
+    public void shouldMatchLiteralWildcardIfFilterHasPatternTypeOfMatch() {
+        assertTrue(new ResourcePatternFilter(TOPIC, "Name", PatternType.MATCH)
+            .matches(new ResourcePattern(TOPIC, "*", LITERAL)));
+    }
+
+    @Test
+    public void shouldMatchPrefixedIfExactMatch() {
+        assertTrue(new ResourcePatternFilter(TOPIC, "Name", PREFIXED)
+            .matches(new ResourcePattern(TOPIC, "Name", PREFIXED)));
+    }
+
+    @Test
+    public void shouldNotMatchIfBothPrefixedAndFilterIsPrefixOfResource() {
+        assertFalse(new ResourcePatternFilter(TOPIC, "Name", PREFIXED)
+            .matches(new ResourcePattern(TOPIC, "Name-something", PREFIXED)));
+    }
+
+    @Test
+    public void shouldNotMatchIfBothPrefixedAndResourceIsPrefixOfFilter() {
+        assertFalse(new ResourcePatternFilter(TOPIC, "Name-something", PREFIXED)
+            .matches(new ResourcePattern(TOPIC, "Name", PREFIXED)));
+    }
+
+    @Test
+    public void shouldNotMatchPrefixedIfNamePrefixedAnyFilterTypeIsAny() {
+        assertFalse(new ResourcePatternFilter(TOPIC, "Name-something", PatternType.ANY)
+            .matches(new ResourcePattern(TOPIC, "Name", PREFIXED)));
+    }
+
+    @Test
+    public void shouldMatchPrefixedIfNamePrefixedAnyFilterTypeIsMatch() {
+        assertTrue(new ResourcePatternFilter(TOPIC, "Name-something", PatternType.MATCH)
+            .matches(new ResourcePattern(TOPIC, "Name", PREFIXED)));
+    }
+}
\ No newline at end of file
diff --git a/clients/src/test/java/org/apache/kafka/common/acl/ResourcePatternTest.java b/clients/src/test/java/org/apache/kafka/common/acl/ResourcePatternTest.java
new file mode 100644
index 0000000000000..d3538e012e334
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/acl/ResourcePatternTest.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.acl;
+
+import org.apache.kafka.common.resource.PatternType;
+import org.apache.kafka.common.resource.ResourcePattern;
+import org.apache.kafka.common.resource.ResourceType;
+import org.junit.Test;
+
+public class ResourcePatternTest {
+    @Test(expected = IllegalArgumentException.class)
+    public void shouldThrowIfResourceTypeIsAny() {
+        new ResourcePattern(ResourceType.ANY, "name", PatternType.LITERAL);
+    }
+
+    @Test(expected = IllegalArgumentException.class)
+    public void shouldThrowIfPatternTypeIsMatch() {
+        new ResourcePattern(ResourceType.TOPIC, "name", PatternType.MATCH);
+    }
+
+    @Test(expected = IllegalArgumentException.class)
+    public void shouldThrowIfPatternTypeIsAny() {
+        new ResourcePattern(ResourceType.TOPIC, "name", PatternType.ANY);
+    }
+
+    @Test(expected = NullPointerException.class)
+    public void shouldThrowIfResourceNameIsNull() {
+        new ResourcePattern(ResourceType.TOPIC, null, PatternType.ANY);
+    }
+}
\ No newline at end of file
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 9e2117975a957..a4fbf04f89579 100644
--- a/clients/src/test/java/org/apache/kafka/common/config/AbstractConfigTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/config/AbstractConfigTest.java
@@ -19,17 +19,22 @@
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.config.ConfigDef.Importance;
 import org.apache.kafka.common.config.ConfigDef.Type;
+import org.apache.kafka.common.config.types.Password;
 import org.apache.kafka.common.metrics.FakeMetricsReporter;
 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;
+import java.util.Collections;
 import java.util.HashMap;
 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;
@@ -50,6 +55,21 @@ public void testConfiguredInstances() {
         testInvalidInputs("org.apache.kafka.common.metrics.FakeMetricsReporter,");
     }
 
+    @Test
+    public void testEmptyList() {
+        AbstractConfig conf;
+        ConfigDef configDef = new ConfigDef().define("a", Type.LIST, "", new ConfigDef.NonNullValidator(), Importance.HIGH, "doc");
+
+        conf = new AbstractConfig(configDef, Collections.emptyMap());
+        assertEquals(Collections.emptyList(), conf.getList("a"));
+
+        conf = new AbstractConfig(configDef, Collections.singletonMap("a", ""));
+        assertEquals(Collections.emptyList(), conf.getList("a"));
+
+        conf = new AbstractConfig(configDef, Collections.singletonMap("a", "b,c,d"));
+        assertEquals(Arrays.asList("b", "c", "d"), conf.getList("a"));
+    }
+
     @Test
     public void testOriginalsWithPrefix() {
         Properties props = new Properties();
@@ -124,18 +144,95 @@ public void testValuesWithPrefixOverride() {
     }
 
     @Test
-    public void testUnused() {
+    public void testValuesWithSecondaryPrefix() {
+        String prefix = "listener.name.listener1.";
+        Password saslJaasConfig1 =  new Password("test.myLoginModule1 required;");
+        Password saslJaasConfig2 =  new Password("test.myLoginModule2 required;");
+        Password saslJaasConfig3 =  new Password("test.myLoginModule3 required;");
+        Properties props = new Properties();
+        props.put("listener.name.listener1.test-mechanism.sasl.jaas.config", saslJaasConfig1.value());
+        props.put("test-mechanism.sasl.jaas.config", saslJaasConfig2.value());
+        props.put("sasl.jaas.config", saslJaasConfig3.value());
+        props.put("listener.name.listener1.gssapi.sasl.kerberos.kinit.cmd", "/usr/bin/kinit2");
+        props.put("listener.name.listener1.gssapi.sasl.kerberos.service.name", "testkafka");
+        props.put("listener.name.listener1.gssapi.sasl.kerberos.min.time.before.relogin", "60000");
+        props.put("ssl.provider", "TEST");
+        TestSecurityConfig config = new TestSecurityConfig(props);
+        Map valuesWithPrefixOverride = config.valuesWithPrefixOverride(prefix);
+
+        // prefix with mechanism overrides global
+        assertTrue(config.unused().contains("listener.name.listener1.test-mechanism.sasl.jaas.config"));
+        assertTrue(config.unused().contains("test-mechanism.sasl.jaas.config"));
+        assertEquals(saslJaasConfig1, valuesWithPrefixOverride.get("test-mechanism.sasl.jaas.config"));
+        assertEquals(saslJaasConfig3, valuesWithPrefixOverride.get("sasl.jaas.config"));
+        assertFalse(config.unused().contains("listener.name.listener1.test-mechanism.sasl.jaas.config"));
+        assertFalse(config.unused().contains("test-mechanism.sasl.jaas.config"));
+        assertFalse(config.unused().contains("sasl.jaas.config"));
+
+        // prefix with mechanism overrides default
+        assertFalse(config.unused().contains("sasl.kerberos.kinit.cmd"));
+        assertTrue(config.unused().contains("listener.name.listener1.gssapi.sasl.kerberos.kinit.cmd"));
+        assertFalse(config.unused().contains("gssapi.sasl.kerberos.kinit.cmd"));
+        assertFalse(config.unused().contains("sasl.kerberos.kinit.cmd"));
+        assertEquals("/usr/bin/kinit2", valuesWithPrefixOverride.get("gssapi.sasl.kerberos.kinit.cmd"));
+        assertFalse(config.unused().contains("listener.name.listener1.sasl.kerberos.kinit.cmd"));
+
+        // prefix override for mechanism with no default
+        assertFalse(config.unused().contains("sasl.kerberos.service.name"));
+        assertTrue(config.unused().contains("listener.name.listener1.gssapi.sasl.kerberos.service.name"));
+        assertFalse(config.unused().contains("gssapi.sasl.kerberos.service.name"));
+        assertFalse(config.unused().contains("sasl.kerberos.service.name"));
+        assertEquals("testkafka", valuesWithPrefixOverride.get("gssapi.sasl.kerberos.service.name"));
+        assertFalse(config.unused().contains("listener.name.listener1.gssapi.sasl.kerberos.service.name"));
+
+        // unset with no default
+        assertTrue(config.unused().contains("ssl.provider"));
+        assertNull(valuesWithPrefixOverride.get("gssapi.ssl.provider"));
+        assertTrue(config.unused().contains("ssl.provider"));
+    }
+
+    @Test
+    public void testValuesWithPrefixAllOrNothing() {
+        String prefix1 = "prefix1.";
+        String prefix2 = "prefix2.";
+        Properties props = new Properties();
+        props.put("sasl.mechanism", "PLAIN");
+        props.put("prefix1.sasl.mechanism", "GSSAPI");
+        props.put("prefix1.sasl.kerberos.kinit.cmd", "/usr/bin/kinit2");
+        props.put("prefix1.ssl.truststore.location", "my location");
+        props.put("sasl.kerberos.service.name", "service name");
+        props.put("ssl.keymanager.algorithm", "algorithm");
+        TestSecurityConfig config = new TestSecurityConfig(props);
+        Map valuesWithPrefixAllOrNothing1 = config.valuesWithPrefixAllOrNothing(prefix1);
+
+        // All prefixed values are there
+        assertEquals("GSSAPI", valuesWithPrefixAllOrNothing1.get("sasl.mechanism"));
+        assertEquals("/usr/bin/kinit2", valuesWithPrefixAllOrNothing1.get("sasl.kerberos.kinit.cmd"));
+        assertEquals("my location", valuesWithPrefixAllOrNothing1.get("ssl.truststore.location"));
+
+        // Non-prefixed values are missing
+        assertFalse(valuesWithPrefixAllOrNothing1.containsKey("sasl.kerberos.service.name"));
+        assertFalse(valuesWithPrefixAllOrNothing1.containsKey("ssl.keymanager.algorithm"));
+
+        Map valuesWithPrefixAllOrNothing2 = config.valuesWithPrefixAllOrNothing(prefix2);
+        assertTrue(valuesWithPrefixAllOrNothing2.containsKey("sasl.kerberos.service.name"));
+        assertTrue(valuesWithPrefixAllOrNothing2.containsKey("ssl.keymanager.algorithm"));
+    }
+
+    @Test
+    public void testUnusedConfigs() {
         Properties props = new Properties();
         String configValue = "org.apache.kafka.common.config.AbstractConfigTest$ConfiguredFakeMetricsReporter";
         props.put(TestConfig.METRIC_REPORTER_CLASSES_CONFIG, configValue);
-        props.put(FakeMetricsReporterConfig.EXTRA_CONFIG, "my_value");
+        props.put(ConfiguredFakeMetricsReporter.EXTRA_CONFIG, "my_value");
         TestConfig config = new TestConfig(props);
 
-        assertTrue("metric.extra_config should be marked unused before getConfiguredInstances is called",
-            config.unused().contains(FakeMetricsReporterConfig.EXTRA_CONFIG));
+        assertTrue(ConfiguredFakeMetricsReporter.EXTRA_CONFIG + " should be marked unused before getConfiguredInstances is called",
+            config.unused().contains(ConfiguredFakeMetricsReporter.EXTRA_CONFIG));
 
         config.getConfiguredInstances(TestConfig.METRIC_REPORTER_CLASSES_CONFIG, MetricsReporter.class);
-        assertTrue("All defined configurations should be marked as used", config.unused().isEmpty());
+        assertFalse(ConfiguredFakeMetricsReporter.EXTRA_CONFIG + " should be marked as used",
+            config.unused().contains(ConfiguredFakeMetricsReporter.EXTRA_CONFIG));
     }
 
     private void testValidInputs(String configValue) {
@@ -170,7 +267,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);
             }
@@ -201,7 +298,7 @@ protected Class findClass(String name) throws ClassNotFoundException {
         // Test class overrides where some classes are not visible to thread context classloader
         Thread.currentThread().setContextClassLoader(restrictedClassLoader);
         // Properties specified as classes should succeed
-        testConfig = new ClassTestConfig(ClassTestConfig.RESTRICTED_CLASS, Arrays.asList(ClassTestConfig.RESTRICTED_CLASS));
+        testConfig = new ClassTestConfig(ClassTestConfig.RESTRICTED_CLASS, Collections.singletonList(ClassTestConfig.RESTRICTED_CLASS));
         testConfig.checkInstances(ClassTestConfig.RESTRICTED_CLASS, ClassTestConfig.RESTRICTED_CLASS);
         testConfig = new ClassTestConfig(ClassTestConfig.RESTRICTED_CLASS, Arrays.asList(ClassTestConfig.VISIBLE_CLASS, ClassTestConfig.RESTRICTED_CLASS));
         testConfig.checkInstances(ClassTestConfig.RESTRICTED_CLASS, ClassTestConfig.VISIBLE_CLASS, ClassTestConfig.RESTRICTED_CLASS);
@@ -228,6 +325,216 @@ 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 testOriginalWithOverrides() {
+        Properties props = new Properties();
+        props.put("config.providers", "file");
+        TestIndirectConfigResolution config = new TestIndirectConfigResolution(props);
+        assertEquals(config.originals().get("config.providers"), "file");
+        assertEquals(config.originals(Collections.singletonMap("config.providers", "file2")).get("config.providers"), "file2");
+    }
+
+    @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("testKey", config.originals().get("sasl.kerberos.key"));
+        assertEquals("randomPassword", config.originals().get("sasl.kerberos.password"));
+        assertEquals(5, config.originals().get("prefix.ssl.truststore.location.number"));
+        assertEquals("service name", config.originals().get("sasl.kerberos.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("testKey", config.originals().get("sasl.kerberos.key"));
+        assertEquals("randomPassword", config.originals().get("sasl.kerberos.password"));
+        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("testKey", config.originals().get("sasl.kerberos.key"));
+        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("testKey", config.originals().get("sasl.kerberos.key"));
+        assertEquals("randomPassword", config.originals().get("sasl.kerberos.password"));
+        assertEquals("testTruststoreKey", config.originals().get("sasl.truststore.key"));
+        assertEquals("randomtruststorePassword", config.originals().get("sasl.truststore.password"));
+        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 {
+            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("${test:/foo/bar/testpath:testKey}", config.originals().get("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("${test:/foo/bar/testpath:random}", config.originals().get("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("${file:/usr/kerberos:key}", config.originals().get("sasl.kerberos.key"));
+    }
+
+    @Test
+    public void testConfigProviderConfigurationWithConfigParams() {
+        // Test Case: Valid Test Case With Multiple ConfigProviders as a separate variable
+        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.key", "${vault:/usr/truststore:truststoreKey}");
+        props.put("sasl.truststore.password", "${vault:/usr/truststore:truststorePassword}");
+        props.put("sasl.truststore.location", "${vault:/usr/truststore:truststoreLocation}");
+        TestIndirectConfigResolution config = new TestIndirectConfigResolution(props, convertPropertiesToMap(providers));
+        assertEquals("/usr/vault", config.originals().get("sasl.truststore.location"));
+    }
+
+    @Test
+    public void testDocumentationOf() {
+        Properties props = new Properties();
+        TestIndirectConfigResolution config = new TestIndirectConfigResolution(props);
+
+        assertEquals(
+                TestIndirectConfigResolution.INDIRECT_CONFIGS_DOC,
+                    config.documentationOf(TestIndirectConfigResolution.INDIRECT_CONFIGS)
+        );
+    }
+
+    @Test
+    public void testDocumentationOfExpectNull() {
+        Properties props = new Properties();
+        TestIndirectConfigResolution config = new TestIndirectConfigResolution(props);
+
+        assertNull(config.documentationOf("xyz"));
+    }
+
+    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;
@@ -236,7 +543,7 @@ private static class ClassTestConfig extends AbstractConfig {
         private static final ConfigDef CONFIG;
         static {
             CONFIG = new ConfigDef().define("class.prop", Type.CLASS, DEFAULT_CLASS, Importance.HIGH, "docs")
-                                    .define("list.prop", Type.LIST, Arrays.asList(DEFAULT_CLASS), Importance.HIGH, "docs");
+                                    .define("list.prop", Type.LIST, Collections.singletonList(DEFAULT_CLASS), Importance.HIGH, "docs");
         }
 
         public ClassTestConfig() {
@@ -296,26 +603,12 @@ public TestConfig(Map props) {
     }
 
     public static class ConfiguredFakeMetricsReporter extends FakeMetricsReporter {
+        public static final String EXTRA_CONFIG = "metric.extra_config";
         @Override
         public void configure(Map configs) {
-            FakeMetricsReporterConfig config = new FakeMetricsReporterConfig(configs);
-
-            // Calling getString() should have the side effect of marking that config as used.
-            config.getString(FakeMetricsReporterConfig.EXTRA_CONFIG);
-        }
-    }
-
-    public static class FakeMetricsReporterConfig extends AbstractConfig {
-
-        public static final String EXTRA_CONFIG = "metric.extra_config";
-        private static final String EXTRA_CONFIG_DOC = "An extraneous configuration string.";
-        private static final ConfigDef CONFIG = new ConfigDef().define(
-                EXTRA_CONFIG, ConfigDef.Type.STRING, "",
-                ConfigDef.Importance.LOW, EXTRA_CONFIG_DOC);
-
-
-        public FakeMetricsReporterConfig(Map props) {
-            super(CONFIG, props);
+            // Calling get() should have the side effect of marking that config as used.
+            // this is required by testUnusedConfigs
+            configs.get(EXTRA_CONFIG);
         }
     }
 }
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 ed4997dea4ba9..389ec7c3bc2c0 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;
@@ -25,6 +26,7 @@
 import org.apache.kafka.common.config.types.Password;
 import org.junit.Test;
 
+import java.time.Duration;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
@@ -39,6 +41,7 @@
 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 ConfigDefTest {
@@ -156,8 +159,15 @@ 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"});
+        testValidators(Type.STRING, new ConfigDef.NonEmptyStringWithoutControlChars(), "defaultname",
+                new Object[]{"test", "name", "test/test", "test\u1234", "\u1324name\\", "/+%>&):??<&()?-", "+1", "\uD83D\uDE01", "\uF3B1", "     test   \n\r", "\n  hello \t"},
+                new Object[]{"nontrailing\nnotallowed", "as\u0001cii control char", "tes\rt", "test\btest", "1\t2", ""});
     }
 
     @Test
@@ -362,6 +372,32 @@ public void testInternalConfigDoesntShowUpInDocs() throws Exception {
         assertFalse(configDef.toRst().contains("my.config"));
     }
 
+    @Test
+    public void testDynamicUpdateModeInDocs() throws Exception {
+        final ConfigDef configDef = new ConfigDef()
+                .define("my.broker.config", Type.LONG, Importance.HIGH, "docs")
+                .define("my.cluster.config", Type.LONG, Importance.HIGH, "docs")
+                .define("my.readonly.config", Type.LONG, Importance.HIGH, "docs");
+        final Map updateModes = new HashMap<>();
+        updateModes.put("my.broker.config", "per-broker");
+        updateModes.put("my.cluster.config", "cluster-wide");
+        final String html = configDef.toHtmlTable(updateModes);
+        Set configsInHtml = new HashSet<>();
+        for (String line : html.split("\n")) {
+            if (line.contains("my.broker.config")) {
+                assertTrue(line.contains("per-broker"));
+                configsInHtml.add("my.broker.config");
+            } else if (line.contains("my.cluster.config")) {
+                assertTrue(line.contains("cluster-wide"));
+                configsInHtml.add("my.cluster.config");
+            } else if (line.contains("my.readonly.config")) {
+                assertTrue(line.contains("read-only"));
+                configsInHtml.add("my.readonly.config");
+            }
+        }
+        assertEquals(configDef.names(), configsInHtml);
+    }
+
     @Test
     public void testNames() {
         final ConfigDef configDef = new ConfigDef()
@@ -607,6 +643,68 @@ 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 {
     }
+
+    @Test
+    public void testNiceMemoryUnits() {
+        assertEquals("", ConfigDef.niceMemoryUnits(0L));
+        assertEquals("", ConfigDef.niceMemoryUnits(1023));
+        assertEquals(" (1 kibibyte)", ConfigDef.niceMemoryUnits(1024));
+        assertEquals("", ConfigDef.niceMemoryUnits(1025));
+        assertEquals(" (2 kibibytes)", ConfigDef.niceMemoryUnits(2 * 1024));
+        assertEquals(" (1 mebibyte)", ConfigDef.niceMemoryUnits(1024 * 1024));
+        assertEquals(" (2 mebibytes)", ConfigDef.niceMemoryUnits(2 * 1024 * 1024));
+        assertEquals(" (1 gibibyte)", ConfigDef.niceMemoryUnits(1024 * 1024 * 1024));
+        assertEquals(" (2 gibibytes)", ConfigDef.niceMemoryUnits(2L * 1024 * 1024 * 1024));
+        assertEquals(" (1 tebibyte)", ConfigDef.niceMemoryUnits(1024L * 1024 * 1024 * 1024));
+        assertEquals(" (2 tebibytes)", ConfigDef.niceMemoryUnits(2L * 1024 * 1024 * 1024 * 1024));
+        assertEquals(" (1024 tebibytes)", ConfigDef.niceMemoryUnits(1024L * 1024 * 1024 * 1024 * 1024));
+        assertEquals(" (2048 tebibytes)", ConfigDef.niceMemoryUnits(2L * 1024 * 1024 * 1024 * 1024 * 1024));
+    }
+
+    @Test
+    public void testNiceTimeUnits() {
+        assertEquals("", ConfigDef.niceTimeUnits(0));
+        assertEquals("", ConfigDef.niceTimeUnits(Duration.ofSeconds(1).toMillis() - 1));
+        assertEquals(" (1 second)", ConfigDef.niceTimeUnits(Duration.ofSeconds(1).toMillis()));
+        assertEquals("", ConfigDef.niceTimeUnits(Duration.ofSeconds(1).toMillis() + 1));
+        assertEquals(" (2 seconds)", ConfigDef.niceTimeUnits(Duration.ofSeconds(2).toMillis()));
+
+        assertEquals(" (1 minute)", ConfigDef.niceTimeUnits(Duration.ofMinutes(1).toMillis()));
+        assertEquals(" (2 minutes)", ConfigDef.niceTimeUnits(Duration.ofMinutes(2).toMillis()));
+
+        assertEquals(" (1 hour)", ConfigDef.niceTimeUnits(Duration.ofHours(1).toMillis()));
+        assertEquals(" (2 hours)", ConfigDef.niceTimeUnits(Duration.ofHours(2).toMillis()));
+
+        assertEquals(" (1 day)", ConfigDef.niceTimeUnits(Duration.ofDays(1).toMillis()));
+        assertEquals(" (2 days)", ConfigDef.niceTimeUnits(Duration.ofDays(2).toMillis()));
+
+        assertEquals(" (7 days)", ConfigDef.niceTimeUnits(Duration.ofDays(7).toMillis()));
+        assertEquals(" (365 days)", ConfigDef.niceTimeUnits(Duration.ofDays(365).toMillis()));
+    }
+
 }
diff --git a/clients/src/test/java/org/apache/kafka/common/config/ConfigResourceTest.java b/clients/src/test/java/org/apache/kafka/common/config/ConfigResourceTest.java
new file mode 100644
index 0000000000000..73effeead48d3
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/config/ConfigResourceTest.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;
+
+import org.junit.Test;
+
+import java.util.Arrays;
+
+import static org.junit.Assert.assertEquals;
+
+public class ConfigResourceTest {
+    @Test
+    public void shouldGetTypeFromId() {
+        assertEquals(ConfigResource.Type.TOPIC, ConfigResource.Type.forId((byte) 2));
+        assertEquals(ConfigResource.Type.BROKER, ConfigResource.Type.forId((byte) 4));
+    }
+
+    @Test
+    public void shouldReturnUnknownForUnknownCode() {
+        assertEquals(ConfigResource.Type.UNKNOWN, ConfigResource.Type.forId((byte) -1));
+        assertEquals(ConfigResource.Type.UNKNOWN, ConfigResource.Type.forId((byte) 0));
+        assertEquals(ConfigResource.Type.UNKNOWN, ConfigResource.Type.forId((byte) 1));
+    }
+
+    @Test
+    public void shouldRoundTripEveryType() {
+        Arrays.stream(ConfigResource.Type.values()).forEach(type ->
+            assertEquals(type.toString(), type, ConfigResource.Type.forId(type.id())));
+    }
+}
\ No newline at end of file
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
new file mode 100644
index 0000000000000..12c6b1f4a2785
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/config/ConfigTransformerTest.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.common.config;
+
+import org.apache.kafka.common.config.provider.ConfigProvider;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Set;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+
+public class ConfigTransformerTest {
+
+    public static final String MY_KEY = "myKey";
+    public static final String TEST_INDIRECTION = "testIndirection";
+    public static final String TEST_KEY = "testKey";
+    public static final String TEST_KEY_WITH_TTL = "testKeyWithTTL";
+    public static final String TEST_PATH = "testPath";
+    public static final String TEST_RESULT = "testResult";
+    public static final String TEST_RESULT_WITH_TTL = "testResultWithTTL";
+    public static final String TEST_RESULT_NO_PATH = "testResultNoPath";
+
+    private ConfigTransformer configTransformer;
+
+    @Before
+    public void setup() {
+        configTransformer = new ConfigTransformer(Collections.singletonMap("test", new TestConfigProvider()));
+    }
+
+    @Test
+    public void testReplaceVariable() throws Exception {
+        ConfigTransformerResult result = configTransformer.transform(Collections.singletonMap(MY_KEY, "${test:testPath:testKey}"));
+        Map data = result.data();
+        Map ttls = result.ttls();
+        assertEquals(TEST_RESULT, data.get(MY_KEY));
+        assertTrue(ttls.isEmpty());
+    }
+
+    @Test
+    public void testReplaceVariableWithTTL() throws Exception {
+        ConfigTransformerResult result = configTransformer.transform(Collections.singletonMap(MY_KEY, "${test:testPath:testKeyWithTTL}"));
+        Map data = result.data();
+        Map ttls = result.ttls();
+        assertEquals(TEST_RESULT_WITH_TTL, data.get(MY_KEY));
+        assertEquals(1L, ttls.get(TEST_PATH).longValue());
+    }
+
+    @Test
+    public void testReplaceMultipleVariablesInValue() throws Exception {
+        ConfigTransformerResult result = configTransformer.transform(Collections.singletonMap(MY_KEY, "hello, ${test:testPath:testKey}; goodbye, ${test:testPath:testKeyWithTTL}!!!"));
+        Map data = result.data();
+        assertEquals("hello, testResult; goodbye, testResultWithTTL!!!", data.get(MY_KEY));
+    }
+
+    @Test
+    public void testNoReplacement() throws Exception {
+        ConfigTransformerResult result = configTransformer.transform(Collections.singletonMap(MY_KEY, "${test:testPath:missingKey}"));
+        Map data = result.data();
+        assertEquals("${test:testPath:missingKey}", data.get(MY_KEY));
+    }
+
+    @Test
+    public void testSingleLevelOfIndirection() throws Exception {
+        ConfigTransformerResult result = configTransformer.transform(Collections.singletonMap(MY_KEY, "${test:testPath:testIndirection}"));
+        Map data = result.data();
+        assertEquals("${test:testPath:testResult}", data.get(MY_KEY));
+    }
+
+    @Test
+    public void testReplaceVariableNoPath() throws Exception {
+        ConfigTransformerResult result = configTransformer.transform(Collections.singletonMap(MY_KEY, "${test:testKey}"));
+        Map data = result.data();
+        Map ttls = result.ttls();
+        assertEquals(TEST_RESULT_NO_PATH, data.get(MY_KEY));
+        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));
+        Map data = result.data();
+        Map ttls = result.ttls();
+        assertNull(data.get(MY_KEY));
+        assertTrue(ttls.isEmpty());
+    }
+
+    public static class TestConfigProvider implements ConfigProvider {
+
+        public void configure(Map configs) {
+        }
+
+        public ConfigData get(String path) {
+            return null;
+        }
+
+        public ConfigData get(String path, Set keys) {
+            Map data = new HashMap<>();
+            Long ttl = null;
+            if (TEST_PATH.equals(path)) {
+                if (keys.contains(TEST_KEY)) {
+                    data.put(TEST_KEY, TEST_RESULT);
+                }
+                if (keys.contains(TEST_KEY_WITH_TTL)) {
+                    data.put(TEST_KEY_WITH_TTL, TEST_RESULT_WITH_TTL);
+                    ttl = 1L;
+                }
+                if (keys.contains(TEST_INDIRECTION)) {
+                    data.put(TEST_INDIRECTION, "${test:testPath:testResult}");
+                }
+            } else {
+                if (keys.contains(TEST_KEY)) {
+                    data.put(TEST_KEY, TEST_RESULT_NO_PATH);
+                }
+            }
+            return new ConfigData(data, ttl);
+        }
+
+        public void close() {
+        }
+    }
+
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/config/SaslConfigsTest.java b/clients/src/test/java/org/apache/kafka/common/config/SaslConfigsTest.java
new file mode 100644
index 0000000000000..760fcd23380b6
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/config/SaslConfigsTest.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.config;
+
+import static org.junit.Assert.assertEquals;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+import org.junit.Test;
+
+public class SaslConfigsTest {
+    @Test
+    public void testSaslLoginRefreshDefaults() {
+        Map vals = new ConfigDef().withClientSaslSupport().parse(Collections.emptyMap());
+        assertEquals(SaslConfigs.DEFAULT_LOGIN_REFRESH_WINDOW_FACTOR,
+                vals.get(SaslConfigs.SASL_LOGIN_REFRESH_WINDOW_FACTOR));
+        assertEquals(SaslConfigs.DEFAULT_LOGIN_REFRESH_WINDOW_JITTER,
+                vals.get(SaslConfigs.SASL_LOGIN_REFRESH_WINDOW_JITTER));
+        assertEquals(SaslConfigs.DEFAULT_LOGIN_REFRESH_MIN_PERIOD_SECONDS,
+                vals.get(SaslConfigs.SASL_LOGIN_REFRESH_MIN_PERIOD_SECONDS));
+        assertEquals(SaslConfigs.DEFAULT_LOGIN_REFRESH_BUFFER_SECONDS,
+                vals.get(SaslConfigs.SASL_LOGIN_REFRESH_BUFFER_SECONDS));
+    }
+
+    @Test
+    public void testSaslLoginRefreshMinValuesAreValid() {
+        Map props = new HashMap<>();
+        props.put(SaslConfigs.SASL_LOGIN_REFRESH_WINDOW_FACTOR, "0.5");
+        props.put(SaslConfigs.SASL_LOGIN_REFRESH_WINDOW_JITTER, "0.0");
+        props.put(SaslConfigs.SASL_LOGIN_REFRESH_MIN_PERIOD_SECONDS, "0");
+        props.put(SaslConfigs.SASL_LOGIN_REFRESH_BUFFER_SECONDS, "0");
+        Map vals = new ConfigDef().withClientSaslSupport().parse(props);
+        assertEquals(Double.valueOf("0.5"), vals.get(SaslConfigs.SASL_LOGIN_REFRESH_WINDOW_FACTOR));
+        assertEquals(Double.valueOf("0.0"), vals.get(SaslConfigs.SASL_LOGIN_REFRESH_WINDOW_JITTER));
+        assertEquals(Short.valueOf("0"), vals.get(SaslConfigs.SASL_LOGIN_REFRESH_MIN_PERIOD_SECONDS));
+        assertEquals(Short.valueOf("0"), vals.get(SaslConfigs.SASL_LOGIN_REFRESH_BUFFER_SECONDS));
+    }
+
+    @Test
+    public void testSaslLoginRefreshMaxValuesAreValid() {
+        Map props = new HashMap<>();
+        props.put(SaslConfigs.SASL_LOGIN_REFRESH_WINDOW_FACTOR, "1.0");
+        props.put(SaslConfigs.SASL_LOGIN_REFRESH_WINDOW_JITTER, "0.25");
+        props.put(SaslConfigs.SASL_LOGIN_REFRESH_MIN_PERIOD_SECONDS, "900");
+        props.put(SaslConfigs.SASL_LOGIN_REFRESH_BUFFER_SECONDS, "3600");
+        Map vals = new ConfigDef().withClientSaslSupport().parse(props);
+        assertEquals(Double.valueOf("1.0"), vals.get(SaslConfigs.SASL_LOGIN_REFRESH_WINDOW_FACTOR));
+        assertEquals(Double.valueOf("0.25"), vals.get(SaslConfigs.SASL_LOGIN_REFRESH_WINDOW_JITTER));
+        assertEquals(Short.valueOf("900"), vals.get(SaslConfigs.SASL_LOGIN_REFRESH_MIN_PERIOD_SECONDS));
+        assertEquals(Short.valueOf("3600"), vals.get(SaslConfigs.SASL_LOGIN_REFRESH_BUFFER_SECONDS));
+    }
+
+    @Test(expected = ConfigException.class)
+    public void testSaslLoginRefreshWindowFactorMinValueIsReallyMinimum() {
+        Map props = new HashMap<>();
+        props.put(SaslConfigs.SASL_LOGIN_REFRESH_WINDOW_FACTOR, "0.499999");
+        new ConfigDef().withClientSaslSupport().parse(props);
+    }
+
+    @Test(expected = ConfigException.class)
+    public void testSaslLoginRefreshWindowFactorMaxValueIsReallyMaximum() {
+        Map props = new HashMap<>();
+        props.put(SaslConfigs.SASL_LOGIN_REFRESH_WINDOW_FACTOR, "1.0001");
+        new ConfigDef().withClientSaslSupport().parse(props);
+    }
+
+    @Test(expected = ConfigException.class)
+    public void testSaslLoginRefreshWindowJitterMinValueIsReallyMinimum() {
+        Map props = new HashMap<>();
+        props.put(SaslConfigs.SASL_LOGIN_REFRESH_WINDOW_JITTER, "-0.000001");
+        new ConfigDef().withClientSaslSupport().parse(props);
+    }
+
+    @Test(expected = ConfigException.class)
+    public void testSaslLoginRefreshWindowJitterMaxValueIsReallyMaximum() {
+        Map props = new HashMap<>();
+        props.put(SaslConfigs.SASL_LOGIN_REFRESH_WINDOW_JITTER, "0.251");
+        new ConfigDef().withClientSaslSupport().parse(props);
+    }
+
+    @Test(expected = ConfigException.class)
+    public void testSaslLoginRefreshMinPeriodSecondsMinValueIsReallyMinimum() {
+        Map props = new HashMap<>();
+        props.put(SaslConfigs.SASL_LOGIN_REFRESH_MIN_PERIOD_SECONDS, "-1");
+        new ConfigDef().withClientSaslSupport().parse(props);
+    }
+
+    @Test(expected = ConfigException.class)
+    public void testSaslLoginRefreshMinPeriodSecondsMaxValueIsReallyMaximum() {
+        Map props = new HashMap<>();
+        props.put(SaslConfigs.SASL_LOGIN_REFRESH_MIN_PERIOD_SECONDS, "901");
+        new ConfigDef().withClientSaslSupport().parse(props);
+    }
+
+    @Test(expected = ConfigException.class)
+    public void testSaslLoginRefreshBufferSecondsMinValueIsReallyMinimum() {
+        Map props = new HashMap<>();
+        props.put(SaslConfigs.SASL_LOGIN_REFRESH_BUFFER_SECONDS, "-1");
+        new ConfigDef().withClientSaslSupport().parse(props);
+    }
+
+    @Test(expected = ConfigException.class)
+    public void testSaslLoginRefreshBufferSecondsMaxValueIsReallyMaximum() {
+        Map props = new HashMap<>();
+        props.put(SaslConfigs.SASL_LOGIN_REFRESH_BUFFER_SECONDS, "3601");
+        new ConfigDef().withClientSaslSupport().parse(props);
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/config/provider/DirectoryConfigProviderTest.java b/clients/src/test/java/org/apache/kafka/common/config/provider/DirectoryConfigProviderTest.java
new file mode 100644
index 0000000000000..9d7139bca6f4d
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/config/provider/DirectoryConfigProviderTest.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.common.config.provider;
+
+import org.apache.kafka.common.config.ConfigData;
+import org.apache.kafka.common.utils.Utils;
+import org.apache.kafka.test.TestUtils;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.io.File;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.util.Collections;
+import java.util.Locale;
+import java.util.Set;
+
+import static java.util.Arrays.asList;
+import static org.apache.kafka.test.TestUtils.toSet;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+
+public class DirectoryConfigProviderTest {
+
+    private DirectoryConfigProvider provider;
+    private File parent;
+    private File dir;
+    private File bar;
+    private File foo;
+    private File subdir;
+    private File subdirFile;
+    private File siblingDir;
+    private File siblingDirFile;
+    private File siblingFile;
+
+    private static File writeFile(File file) throws IOException {
+        Files.write(file.toPath(), file.getName().toUpperCase(Locale.ENGLISH).getBytes(StandardCharsets.UTF_8));
+        return file;
+    }
+
+    @Before
+    public void setup() throws IOException {
+        provider = new DirectoryConfigProvider();
+        provider.configure(Collections.emptyMap());
+        parent = TestUtils.tempDirectory();
+        dir = new File(parent, "dir");
+        dir.mkdir();
+        foo = writeFile(new File(dir, "foo"));
+        bar = writeFile(new File(dir, "bar"));
+        subdir = new File(dir, "subdir");
+        subdir.mkdir();
+        subdirFile = writeFile(new File(subdir, "subdirFile"));
+        siblingDir = new File(parent, "siblingdir");
+        siblingDir.mkdir();
+        siblingDirFile = writeFile(new File(siblingDir, "siblingdirFile"));
+        siblingFile = writeFile(new File(parent, "siblingFile"));
+    }
+
+    @After
+    public void close() throws IOException {
+        provider.close();
+        Utils.delete(parent);
+    }
+
+    @Test
+    public void testGetAllKeysAtPath() throws IOException {
+        ConfigData configData = provider.get(dir.getAbsolutePath());
+        assertEquals(toSet(asList(foo.getName(), bar.getName())), configData.data().keySet());
+        assertEquals("FOO", configData.data().get(foo.getName()));
+        assertEquals("BAR", configData.data().get(bar.getName()));
+        assertNull(configData.ttl());
+    }
+
+    @Test
+    public void testGetSetOfKeysAtPath() {
+        Set keys = toSet(asList(foo.getName(), "baz"));
+        ConfigData configData = provider.get(dir.getAbsolutePath(), keys);
+        assertEquals(Collections.singleton(foo.getName()), configData.data().keySet());
+        assertEquals("FOO", configData.data().get(foo.getName()));
+        assertNull(configData.ttl());
+    }
+
+    @Test
+    public void testNoSubdirs() {
+        // Only regular files directly in the path directory are allowed, not in subdirs
+        Set keys = toSet(asList(subdir.getName(), String.join(File.separator, subdir.getName(), subdirFile.getName())));
+        ConfigData configData = provider.get(dir.getAbsolutePath(), keys);
+        assertTrue(configData.data().isEmpty());
+        assertNull(configData.ttl());
+    }
+
+    @Test
+    public void testNoTraversal() {
+        // Check we can't escape outside the path directory
+        Set keys = toSet(asList(
+                String.join(File.separator, "..", siblingFile.getName()),
+                String.join(File.separator, "..", siblingDir.getName()),
+                String.join(File.separator, "..", siblingDir.getName(), siblingDirFile.getName())));
+        ConfigData configData = provider.get(dir.getAbsolutePath(), keys);
+        assertTrue(configData.data().isEmpty());
+        assertNull(configData.ttl());
+    }
+
+    @Test
+    public void testEmptyPath() {
+        ConfigData configData = provider.get("");
+        assertTrue(configData.data().isEmpty());
+        assertNull(configData.ttl());
+    }
+
+    @Test
+    public void testEmptyPathWithKey() {
+        ConfigData configData = provider.get("", Collections.singleton("foo"));
+        assertTrue(configData.data().isEmpty());
+        assertEquals(null, configData.ttl());
+    }
+
+    @Test
+    public void testNullPath() {
+        ConfigData configData = provider.get(null);
+        assertTrue(configData.data().isEmpty());
+        assertEquals(null, configData.ttl());
+    }
+
+    @Test
+    public void testNullPathWithKey() {
+        ConfigData configData = provider.get(null, Collections.singleton("foo"));
+        assertTrue(configData.data().isEmpty());
+        assertEquals(null, configData.ttl());
+    }
+}
+
diff --git a/clients/src/test/java/org/apache/kafka/common/config/provider/FileConfigProviderTest.java b/clients/src/test/java/org/apache/kafka/common/config/provider/FileConfigProviderTest.java
new file mode 100644
index 0000000000000..b2c791afebeaf
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/config/provider/FileConfigProviderTest.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.common.config.provider;
+
+import org.apache.kafka.common.config.ConfigData;
+import org.junit.Before;
+import org.junit.Test;
+
+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.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+public class FileConfigProviderTest {
+
+    private FileConfigProvider configProvider;
+
+    @Before
+    public void setup() {
+        configProvider = new TestFileConfigProvider();
+    }
+
+    @Test
+    public void testGetAllKeysAtPath() throws Exception {
+        ConfigData configData = configProvider.get("dummy");
+        Map result = new HashMap<>();
+        result.put("testKey", "testResult");
+        result.put("testKey2", "testResult2");
+        assertEquals(result, configData.data());
+        assertEquals(null, configData.ttl());
+    }
+
+    @Test
+    public void testGetOneKeyAtPath() throws Exception {
+        ConfigData configData = configProvider.get("dummy", Collections.singleton("testKey"));
+        Map result = new HashMap<>();
+        result.put("testKey", "testResult");
+        assertEquals(result, configData.data());
+        assertEquals(null, configData.ttl());
+    }
+
+    @Test
+    public void testEmptyPath() throws Exception {
+        ConfigData configData = configProvider.get("", Collections.singleton("testKey"));
+        assertTrue(configData.data().isEmpty());
+        assertEquals(null, configData.ttl());
+    }
+
+    @Test
+    public void testEmptyPathWithKey() throws Exception {
+        ConfigData configData = configProvider.get("");
+        assertTrue(configData.data().isEmpty());
+        assertEquals(null, configData.ttl());
+    }
+
+    @Test
+    public void testNullPath() throws Exception {
+        ConfigData configData = configProvider.get(null);
+        assertTrue(configData.data().isEmpty());
+        assertEquals(null, configData.ttl());
+    }
+
+    @Test
+    public void testNullPathWithKey() throws Exception {
+        ConfigData configData = configProvider.get(null, Collections.singleton("testKey"));
+        assertTrue(configData.data().isEmpty());
+        assertEquals(null, configData.ttl());
+    }
+
+    public static class TestFileConfigProvider extends FileConfigProvider {
+
+        @Override
+        protected Reader reader(String path) throws IOException {
+            return new StringReader("testKey=testResult\ntestKey2=testResult2");
+        }
+    }
+}
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/feature/FeaturesTest.java b/clients/src/test/java/org/apache/kafka/common/feature/FeaturesTest.java
new file mode 100644
index 0000000000000..896196dc550a4
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/feature/FeaturesTest.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.common.feature;
+
+import java.util.HashMap;
+import java.util.Map;
+
+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;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertThrows;
+
+public class FeaturesTest {
+
+    @Test
+    public void testEmptyFeatures() {
+        Map> emptyMap = new HashMap<>();
+
+        Features emptyFinalizedFeatures = Features.emptyFinalizedFeatures();
+        assertTrue(emptyFinalizedFeatures.features().isEmpty());
+        assertTrue(emptyFinalizedFeatures.toMap().isEmpty());
+        assertEquals(emptyFinalizedFeatures, Features.fromFinalizedFeaturesMap(emptyMap));
+
+        Features emptySupportedFeatures = Features.emptySupportedFeatures();
+        assertTrue(emptySupportedFeatures.features().isEmpty());
+        assertTrue(emptySupportedFeatures.toMap().isEmpty());
+        assertEquals(emptySupportedFeatures, Features.fromSupportedFeaturesMap(emptyMap));
+    }
+
+    @Test
+    public void testNullFeatures() {
+        assertThrows(
+            NullPointerException.class,
+            () -> Features.finalizedFeatures(null));
+        assertThrows(
+            NullPointerException.class,
+            () -> Features.supportedFeatures(null));
+    }
+
+    @Test
+    public void testGetAllFeaturesAPI() {
+        SupportedVersionRange v1 = new SupportedVersionRange((short) 1, (short) 2);
+        SupportedVersionRange v2 = new SupportedVersionRange((short) 3, (short) 4);
+        Map allFeatures =
+            mkMap(mkEntry("feature_1", v1), mkEntry("feature_2", v2));
+        Features features = Features.supportedFeatures(allFeatures);
+        assertEquals(allFeatures, features.features());
+    }
+
+    @Test
+    public void testGetAPI() {
+        SupportedVersionRange v1 = new SupportedVersionRange((short) 1, (short) 2);
+        SupportedVersionRange v2 = new SupportedVersionRange((short) 3, (short) 4);
+        Map allFeatures = mkMap(mkEntry("feature_1", v1), mkEntry("feature_2", v2));
+        Features features = Features.supportedFeatures(allFeatures);
+        assertEquals(v1, features.get("feature_1"));
+        assertEquals(v2, features.get("feature_2"));
+        assertNull(features.get("nonexistent_feature"));
+    }
+
+    @Test
+    public void testFromFeaturesMapToFeaturesMap() {
+        SupportedVersionRange v1 = new SupportedVersionRange((short) 1, (short) 2);
+        SupportedVersionRange v2 = new SupportedVersionRange((short) 3, (short) 4);
+        Map allFeatures = mkMap(mkEntry("feature_1", v1), mkEntry("feature_2", v2));
+
+        Features features = Features.supportedFeatures(allFeatures);
+
+        Map> expected = mkMap(
+            mkEntry("feature_1", mkMap(mkEntry("min_version", (short) 1), mkEntry("max_version", (short) 2))),
+            mkEntry("feature_2", mkMap(mkEntry("min_version", (short) 3), mkEntry("max_version", (short) 4))));
+        assertEquals(expected, features.toMap());
+        assertEquals(features, Features.fromSupportedFeaturesMap(expected));
+    }
+
+    @Test
+    public void testFromToFinalizedFeaturesMap() {
+        FinalizedVersionRange v1 = new FinalizedVersionRange((short) 1, (short) 2);
+        FinalizedVersionRange v2 = new FinalizedVersionRange((short) 3, (short) 4);
+        Map allFeatures = mkMap(mkEntry("feature_1", v1), mkEntry("feature_2", v2));
+
+        Features features = Features.finalizedFeatures(allFeatures);
+
+        Map> expected = mkMap(
+            mkEntry("feature_1", mkMap(mkEntry("min_version_level", (short) 1), mkEntry("max_version_level", (short) 2))),
+            mkEntry("feature_2", mkMap(mkEntry("min_version_level", (short) 3), mkEntry("max_version_level", (short) 4))));
+        assertEquals(expected, features.toMap());
+        assertEquals(features, Features.fromFinalizedFeaturesMap(expected));
+    }
+
+    @Test
+    public void testToStringFinalizedFeatures() {
+        FinalizedVersionRange v1 = new FinalizedVersionRange((short) 1, (short) 2);
+        FinalizedVersionRange v2 = new FinalizedVersionRange((short) 3, (short) 4);
+        Map allFeatures = mkMap(mkEntry("feature_1", v1), mkEntry("feature_2", v2));
+
+        Features features = Features.finalizedFeatures(allFeatures);
+
+        assertEquals(
+            "Features{(feature_1 -> FinalizedVersionRange[min_version_level:1, max_version_level:2]), (feature_2 -> FinalizedVersionRange[min_version_level:3, max_version_level:4])}",
+            features.toString());
+    }
+
+    @Test
+    public void testToStringSupportedFeatures() {
+        SupportedVersionRange v1 = new SupportedVersionRange((short) 1, (short) 2);
+        SupportedVersionRange v2 = new SupportedVersionRange((short) 3, (short) 4);
+        Map allFeatures
+            = mkMap(mkEntry("feature_1", v1), mkEntry("feature_2", v2));
+
+        Features features = Features.supportedFeatures(allFeatures);
+
+        assertEquals(
+            "Features{(feature_1 -> SupportedVersionRange[min_version:1, max_version:2]), (feature_2 -> SupportedVersionRange[min_version:3, max_version:4])}",
+            features.toString());
+    }
+
+    @Test
+    public void testSuppportedFeaturesFromMapFailureWithInvalidMissingMaxVersion() {
+        // This is invalid because 'max_version' key is missing.
+        Map> invalidFeatures = mkMap(
+            mkEntry("feature_1", mkMap(mkEntry("min_version", (short) 1))));
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> Features.fromSupportedFeaturesMap(invalidFeatures));
+    }
+
+    @Test
+    public void testFinalizedFeaturesFromMapFailureWithInvalidMissingMaxVersionLevel() {
+        // This is invalid because 'max_version_level' key is missing.
+        Map> invalidFeatures = mkMap(
+            mkEntry("feature_1", mkMap(mkEntry("min_version_level", (short) 1))));
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> Features.fromFinalizedFeaturesMap(invalidFeatures));
+    }
+
+    @Test
+    public void testEquals() {
+        SupportedVersionRange v1 = new SupportedVersionRange((short) 1, (short) 2);
+        Map allFeatures = mkMap(mkEntry("feature_1", v1));
+        Features features = Features.supportedFeatures(allFeatures);
+        Features featuresClone = Features.supportedFeatures(allFeatures);
+        assertTrue(features.equals(featuresClone));
+
+        SupportedVersionRange v2 = new SupportedVersionRange((short) 1, (short) 3);
+        Map allFeaturesDifferent = mkMap(mkEntry("feature_1", v2));
+        Features featuresDifferent = Features.supportedFeatures(allFeaturesDifferent);
+        assertFalse(features.equals(featuresDifferent));
+
+        assertFalse(features.equals(null));
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/feature/FinalizedVersionRangeTest.java b/clients/src/test/java/org/apache/kafka/common/feature/FinalizedVersionRangeTest.java
new file mode 100644
index 0000000000000..3d62a8fdbad8a
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/feature/FinalizedVersionRangeTest.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.feature;
+
+import java.util.Map;
+
+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;
+
+/**
+ * Unit tests for the FinalizedVersionRange class.
+ *
+ * Most of the unit tests required for BaseVersionRange are part of the SupportedVersionRangeTest
+ * suite. This suite only tests behavior very specific to FinalizedVersionRange.
+ */
+public class FinalizedVersionRangeTest {
+
+    @Test
+    public void testFromToMap() {
+        FinalizedVersionRange versionRange = new FinalizedVersionRange((short) 1, (short) 2);
+        assertEquals(1, versionRange.min());
+        assertEquals(2, versionRange.max());
+
+        Map versionRangeMap = versionRange.toMap();
+        assertEquals(
+            mkMap(
+                mkEntry("min_version_level", versionRange.min()),
+                mkEntry("max_version_level", versionRange.max())),
+            versionRangeMap);
+
+        FinalizedVersionRange newVersionRange = FinalizedVersionRange.fromMap(versionRangeMap);
+        assertEquals(1, newVersionRange.min());
+        assertEquals(2, newVersionRange.max());
+        assertEquals(versionRange, newVersionRange);
+    }
+
+    @Test
+    public void testToString() {
+        assertEquals("FinalizedVersionRange[min_version_level:1, max_version_level:1]", new FinalizedVersionRange((short) 1, (short) 1).toString());
+        assertEquals("FinalizedVersionRange[min_version_level:1, max_version_level:2]", new FinalizedVersionRange((short) 1, (short) 2).toString());
+    }
+
+    @Test
+    public void testIsCompatibleWith() {
+        assertFalse(new FinalizedVersionRange((short) 1, (short) 1).isIncompatibleWith(new SupportedVersionRange((short) 1, (short) 1)));
+        assertFalse(new FinalizedVersionRange((short) 2, (short) 3).isIncompatibleWith(new SupportedVersionRange((short) 1, (short) 4)));
+        assertFalse(new FinalizedVersionRange((short) 1, (short) 4).isIncompatibleWith(new SupportedVersionRange((short) 1, (short) 4)));
+
+        assertTrue(new FinalizedVersionRange((short) 1, (short) 4).isIncompatibleWith(new SupportedVersionRange((short) 2, (short) 3)));
+        assertTrue(new FinalizedVersionRange((short) 1, (short) 4).isIncompatibleWith(new SupportedVersionRange((short) 2, (short) 4)));
+        assertTrue(new FinalizedVersionRange((short) 2, (short) 4).isIncompatibleWith(new SupportedVersionRange((short) 2, (short) 3)));
+    }
+
+    @Test
+    public void testMinMax() {
+        FinalizedVersionRange versionRange = new FinalizedVersionRange((short) 1, (short) 2);
+        assertEquals(1, versionRange.min());
+        assertEquals(2, versionRange.max());
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/feature/SupportedVersionRangeTest.java b/clients/src/test/java/org/apache/kafka/common/feature/SupportedVersionRangeTest.java
new file mode 100644
index 0000000000000..4c7b5591dad0f
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/feature/SupportedVersionRangeTest.java
@@ -0,0 +1,142 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.feature;
+
+import java.util.Map;
+
+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.assertThrows;
+import static org.junit.Assert.assertTrue;
+
+/**
+ * Unit tests for the SupportedVersionRange class.
+ * Along the way, this suite also includes extensive tests for the base class BaseVersionRange.
+ */
+public class SupportedVersionRangeTest {
+    @Test
+    public void testFailDueToInvalidParams() {
+        // min and max can't be < 1.
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> new SupportedVersionRange((short) 0, (short) 0));
+        // min can't be < 1.
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> new SupportedVersionRange((short) 0, (short) 1));
+        // max can't be < 1.
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> new SupportedVersionRange((short) 1, (short) 0));
+        // min can't be > max.
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> new SupportedVersionRange((short) 2, (short) 1));
+    }
+
+    @Test
+    public void testFromToMap() {
+        SupportedVersionRange versionRange = new SupportedVersionRange((short) 1, (short) 2);
+        assertEquals(1, versionRange.min());
+        assertEquals(2, versionRange.max());
+
+        Map versionRangeMap = versionRange.toMap();
+        assertEquals(
+            mkMap(mkEntry("min_version", versionRange.min()), mkEntry("max_version", versionRange.max())),
+            versionRangeMap);
+
+        SupportedVersionRange newVersionRange = SupportedVersionRange.fromMap(versionRangeMap);
+        assertEquals(1, newVersionRange.min());
+        assertEquals(2, newVersionRange.max());
+        assertEquals(versionRange, newVersionRange);
+    }
+
+    @Test
+    public void testFromMapFailure() {
+        // min_version can't be < 1.
+        Map invalidWithBadMinVersion =
+            mkMap(mkEntry("min_version", (short) 0), mkEntry("max_version", (short) 1));
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> SupportedVersionRange.fromMap(invalidWithBadMinVersion));
+
+        // max_version can't be < 1.
+        Map invalidWithBadMaxVersion =
+            mkMap(mkEntry("min_version", (short) 1), mkEntry("max_version", (short) 0));
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> SupportedVersionRange.fromMap(invalidWithBadMaxVersion));
+
+        // min_version and max_version can't be < 1.
+        Map invalidWithBadMinMaxVersion =
+            mkMap(mkEntry("min_version", (short) 0), mkEntry("max_version", (short) 0));
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> SupportedVersionRange.fromMap(invalidWithBadMinMaxVersion));
+
+        // min_version can't be > max_version.
+        Map invalidWithLowerMaxVersion =
+            mkMap(mkEntry("min_version", (short) 2), mkEntry("max_version", (short) 1));
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> SupportedVersionRange.fromMap(invalidWithLowerMaxVersion));
+
+        // min_version key missing.
+        Map invalidWithMinKeyMissing =
+            mkMap(mkEntry("max_version", (short) 1));
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> SupportedVersionRange.fromMap(invalidWithMinKeyMissing));
+
+        // max_version key missing.
+        Map invalidWithMaxKeyMissing =
+            mkMap(mkEntry("min_version", (short) 1));
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> SupportedVersionRange.fromMap(invalidWithMaxKeyMissing));
+    }
+
+    @Test
+    public void testToString() {
+        assertEquals(
+            "SupportedVersionRange[min_version:1, max_version:1]",
+            new SupportedVersionRange((short) 1, (short) 1).toString());
+        assertEquals(
+            "SupportedVersionRange[min_version:1, max_version:2]",
+            new SupportedVersionRange((short) 1, (short) 2).toString());
+    }
+
+    @Test
+    public void testEquals() {
+        SupportedVersionRange tested = new SupportedVersionRange((short) 1, (short) 1);
+        assertTrue(tested.equals(tested));
+        assertFalse(tested.equals(new SupportedVersionRange((short) 1, (short) 2)));
+        assertFalse(tested.equals(null));
+    }
+
+    @Test
+    public void testMinMax() {
+        SupportedVersionRange versionRange = new SupportedVersionRange((short) 1, (short) 2);
+        assertEquals(1, versionRange.min());
+        assertEquals(2, versionRange.max());
+    }
+}
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..8a6379992051f 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,20 @@
  */
 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.assertThrows;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
 
 public class RecordHeadersTest {
 
@@ -206,6 +207,17 @@ public void testNew() throws IOException {
         assertEquals(2, getCount(newHeaders));
     }
 
+    @Test
+    public void shouldThrowNpeWhenAddingNullHeader() {
+        final RecordHeaders recordHeaders = new RecordHeaders();
+        assertThrows(NullPointerException.class, () -> recordHeaders.add(null));
+    }
+
+    @Test
+    public void shouldThrowNpeWhenAddingCollectionWithNullHeader() {
+        assertThrows(NullPointerException.class, () -> new RecordHeaders(new Header[1]));
+    }
+
     private int getCount(Headers headers) {
         int count = 0;
         Iterator
             headerIterator = headers.iterator();
diff --git a/clients/src/test/java/org/apache/kafka/common/internals/PartitionStatesTest.java b/clients/src/test/java/org/apache/kafka/common/internals/PartitionStatesTest.java
index 65a812a77b172..f9cacb1aab1ba 100644
--- a/clients/src/test/java/org/apache/kafka/common/internals/PartitionStatesTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/internals/PartitionStatesTest.java
@@ -22,10 +22,8 @@
 import java.util.ArrayList;
 import java.util.LinkedHashMap;
 import java.util.List;
-import java.util.Map;
 
 import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertTrue;
 
 public class PartitionStatesTest {
 
@@ -43,8 +41,8 @@ public void testSet() {
         expected.put(new TopicPartition("baz", 3), "baz 3");
         checkState(states, expected);
 
-        states.set(new LinkedHashMap());
-        checkState(states, new LinkedHashMap());
+        states.set(new LinkedHashMap<>());
+        checkState(states, new LinkedHashMap<>());
     }
 
     private LinkedHashMap createMap() {
@@ -61,12 +59,7 @@ private LinkedHashMap createMap() {
     private void checkState(PartitionStates states, LinkedHashMap expected) {
         assertEquals(expected.keySet(), states.partitionSet());
         assertEquals(expected.size(), states.size());
-        List> statesList = new ArrayList<>();
-        for (Map.Entry entry : expected.entrySet()) {
-            statesList.add(new PartitionStates.PartitionState<>(entry.getKey(), entry.getValue()));
-            assertTrue(states.contains(entry.getKey()));
-        }
-        assertEquals(statesList, states.partitionStates());
+        assertEquals(expected, states.partitionStateMap());
     }
 
     @Test
diff --git a/clients/src/test/java/org/apache/kafka/common/memory/GarbageCollectedMemoryPoolTest.java b/clients/src/test/java/org/apache/kafka/common/memory/GarbageCollectedMemoryPoolTest.java
index 788d447551125..02b38564e30a3 100644
--- a/clients/src/test/java/org/apache/kafka/common/memory/GarbageCollectedMemoryPoolTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/memory/GarbageCollectedMemoryPoolTest.java
@@ -80,6 +80,7 @@ public void testReleaseForeignBuffer() throws Exception {
         GarbageCollectedMemoryPool pool = new GarbageCollectedMemoryPool(1000, 10, true, null);
         ByteBuffer fellOffATruck = ByteBuffer.allocate(1);
         pool.release(fellOffATruck);
+        pool.close();
     }
 
     @Test
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
new file mode 100644
index 0000000000000..52cf41584f70d
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/message/MessageTest.java
@@ -0,0 +1,1122 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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 com.fasterxml.jackson.databind.JsonNode;
+
+import org.apache.kafka.common.IsolationLevel;
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.message.AddPartitionsToTxnRequestData.AddPartitionsToTxnTopic;
+import org.apache.kafka.common.message.AddPartitionsToTxnRequestData.AddPartitionsToTxnTopicCollection;
+import org.apache.kafka.common.message.DescribeGroupsResponseData.DescribedGroup;
+import org.apache.kafka.common.message.DescribeGroupsResponseData.DescribedGroupMember;
+import org.apache.kafka.common.message.JoinGroupResponseData.JoinGroupResponseMember;
+import org.apache.kafka.common.message.LeaveGroupResponseData.MemberResponse;
+import org.apache.kafka.common.message.ListOffsetRequestData.ListOffsetPartition;
+import org.apache.kafka.common.message.ListOffsetRequestData.ListOffsetTopic;
+import org.apache.kafka.common.message.ListOffsetResponseData.ListOffsetPartitionResponse;
+import org.apache.kafka.common.message.ListOffsetResponseData.ListOffsetTopicResponse;
+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.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 org.apache.kafka.common.protocol.types.Type;
+import org.apache.kafka.common.utils.Utils;
+import org.junit.Assert;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.Timeout;
+
+import java.lang.reflect.Method;
+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.assertNotEquals;
+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().
+                setResourceTypeFilter((byte) 42).
+                setResourceNameFilter(null).
+                setPatternTypeFilter((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 testListOffsetsRequestVersions() throws Exception {
+        List v = Collections.singletonList(new ListOffsetTopic()
+                .setName("topic")
+                .setPartitions(Collections.singletonList(new ListOffsetPartition()
+                        .setPartitionIndex(0)
+                        .setTimestamp(123L))));
+        Supplier newRequest = () -> new ListOffsetRequestData()
+                .setTopics(v)
+                .setReplicaId(0);
+        testAllMessageRoundTrips(newRequest.get());
+        testAllMessageRoundTripsFromVersion((short) 2, newRequest.get().setIsolationLevel(IsolationLevel.READ_COMMITTED.id()));
+    }
+
+    @Test
+    public void testListOffsetsResponseVersions() throws Exception {
+        ListOffsetPartitionResponse partition = new ListOffsetPartitionResponse()
+                .setErrorCode(Errors.NONE.code())
+                .setPartitionIndex(0)
+                .setOldStyleOffsets(Collections.singletonList(321L));
+        List topics = Collections.singletonList(new ListOffsetTopicResponse()
+                .setName("topic")
+                .setPartitions(Collections.singletonList(partition)));
+        Supplier response = () -> new ListOffsetResponseData()
+                .setTopics(topics);
+        for (short version = 0; version <= ApiKeys.LIST_OFFSETS.latestVersion(); version++) {
+            ListOffsetResponseData responseData = response.get();
+            if (version > 0) {
+                responseData.topics().get(0).partitions().get(0)
+                    .setOldStyleOffsets(Collections.emptyList())
+                    .setOffset(456L)
+                    .setTimestamp(123L);
+            }
+            if (version > 1) {
+                responseData.setThrottleTimeMs(1000);
+            }
+            if (version > 3) {
+                partition.setLeaderEpoch(1);
+            }
+            testEquivalentMessageRoundTrip(version, responseData);
+        }
+    }
+
+    @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 testDescribeGroupsRequestVersions() throws Exception {
+        testAllMessageRoundTrips(new DescribeGroupsRequestData()
+                .setGroups(Collections.singletonList("group"))
+                .setIncludeAuthorizedOperations(false));
+    }
+
+    @Test
+    public void testDescribeGroupsResponseVersions() throws Exception {
+        DescribedGroupMember baseMember = new DescribedGroupMember()
+            .setMemberId(memberId);
+
+        DescribedGroup baseGroup = new DescribedGroup()
+                                       .setGroupId("group")
+                                       .setGroupState("Stable").setErrorCode(Errors.NONE.code())
+                                       .setMembers(Collections.singletonList(baseMember))
+                                       .setProtocolType("consumer");
+        DescribeGroupsResponseData baseResponse = new DescribeGroupsResponseData()
+                                                      .setGroups(Collections.singletonList(baseGroup));
+        testAllMessageRoundTrips(baseResponse);
+
+        testAllMessageRoundTripsFromVersion((short) 1, baseResponse.setThrottleTimeMs(10));
+
+        baseGroup.setAuthorizedOperations(1);
+        testAllMessageRoundTripsFromVersion((short) 3, baseResponse);
+
+        baseMember.setGroupInstanceId(instanceId);
+        testAllMessageRoundTripsFromVersion((short) 4, baseResponse);
+    }
+
+    @Test
+    public void testGroupInstanceIdIgnorableInDescribeGroupsResponse() throws Exception {
+        DescribeGroupsResponseData responseWithGroupInstanceId =
+            new DescribeGroupsResponseData()
+                .setGroups(Collections.singletonList(
+                    new DescribedGroup()
+                        .setGroupId("group")
+                        .setGroupState("Stable")
+                        .setErrorCode(Errors.NONE.code())
+                        .setMembers(Collections.singletonList(
+                            new DescribedGroupMember()
+                                .setMemberId(memberId)
+                                .setGroupInstanceId(instanceId)))
+                        .setProtocolType("consumer")
+                ));
+
+        DescribeGroupsResponseData expectedResponse = responseWithGroupInstanceId.duplicate();
+        // Unset GroupInstanceId
+        expectedResponse.groups().get(0).members().get(0).setGroupInstanceId(null);
+
+        testAllMessageRoundTripsBeforeVersion((short) 4, responseWithGroupInstanceId, expectedResponse);
+    }
+
+    @Test
+    public void testThrottleTimeIgnorableInDescribeGroupsResponse() throws Exception {
+        DescribeGroupsResponseData responseWithGroupInstanceId =
+            new DescribeGroupsResponseData()
+                .setGroups(Collections.singletonList(
+                    new DescribedGroup()
+                        .setGroupId("group")
+                        .setGroupState("Stable")
+                        .setErrorCode(Errors.NONE.code())
+                        .setMembers(Collections.singletonList(
+                            new DescribedGroupMember()
+                                .setMemberId(memberId)))
+                        .setProtocolType("consumer")
+                ))
+                .setThrottleTimeMs(10);
+
+        DescribeGroupsResponseData expectedResponse = responseWithGroupInstanceId.duplicate();
+        // Unset throttle time
+        expectedResponse.setThrottleTimeMs(0);
+
+        testAllMessageRoundTripsBeforeVersion((short) 1, responseWithGroupInstanceId, expectedResponse);
+    }
+
+    @Test
+    public void testOffsetForLeaderEpochVersions() throws Exception {
+        // Version 2 adds optional current leader epoch
+        OffsetForLeaderEpochRequestData.OffsetForLeaderPartition partitionDataNoCurrentEpoch =
+                new OffsetForLeaderEpochRequestData.OffsetForLeaderPartition()
+                        .setPartition(0)
+                        .setLeaderEpoch(3);
+        OffsetForLeaderEpochRequestData.OffsetForLeaderPartition partitionDataWithCurrentEpoch =
+                new OffsetForLeaderEpochRequestData.OffsetForLeaderPartition()
+                        .setPartition(0)
+                        .setLeaderEpoch(3)
+                        .setCurrentLeaderEpoch(5);
+        OffsetForLeaderEpochRequestData data = new OffsetForLeaderEpochRequestData();
+        data.topics().add(new OffsetForLeaderEpochRequestData.OffsetForLeaderTopic()
+                .setTopic("foo")
+                .setPartitions(singletonList(partitionDataNoCurrentEpoch)));
+
+        testAllMessageRoundTrips(data);
+        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) 3, 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;
+        String instanceId = "instance";
+        String memberId = "member";
+        int generationId = 1;
+
+        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)
+                      .setGroupInstanceId(instanceId)
+                      .setMemberId(memberId)
+                      .setGenerationId(generationId)
+                      .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 < 2) {
+                requestData.topics().get(0).partitions().get(0).setCommittedLeaderEpoch(-1);
+            }
+
+            if (version < 3) {
+                final short finalVersion = version;
+                assertThrows(UnsupportedVersionException.class, () -> testEquivalentMessageRoundTrip(finalVersion, requestData));
+                requestData.setGroupInstanceId(null);
+                assertThrows(UnsupportedVersionException.class, () -> testEquivalentMessageRoundTrip(finalVersion, requestData));
+                requestData.setMemberId("");
+                assertThrows(UnsupportedVersionException.class, () -> testEquivalentMessageRoundTrip(finalVersion, requestData));
+                requestData.setGenerationId(-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";
+
+        List topics = Collections.singletonList(
+            new OffsetFetchRequestTopic()
+                .setName(topicName)
+                .setPartitionIndexes(Collections.singletonList(5)));
+        testAllMessageRoundTrips(new OffsetFetchRequestData()
+                                     .setTopics(new ArrayList<>())
+                                     .setGroupId(groupId));
+
+        testAllMessageRoundTrips(new OffsetFetchRequestData()
+                                     .setGroupId(groupId)
+                                     .setTopics(topics));
+
+        OffsetFetchRequestData allPartitionData = new OffsetFetchRequestData()
+                                                      .setGroupId(groupId)
+                                                      .setTopics(null);
+
+        OffsetFetchRequestData requireStableData = new OffsetFetchRequestData()
+                                                       .setGroupId(groupId)
+                                                       .setTopics(topics)
+                                                       .setRequireStable(true);
+
+        for (short version = 0; version <= ApiKeys.OFFSET_FETCH.latestVersion(); version++) {
+            final short finalVersion = version;
+            if (version < 2) {
+                assertThrows(SchemaException.class, () -> testAllMessageRoundTripsFromVersion(finalVersion, allPartitionData));
+            } else {
+                testAllMessageRoundTripsFromVersion(version, allPartitionData);
+            }
+
+            if (version < 7) {
+                assertThrows(UnsupportedVersionException.class, () -> testAllMessageRoundTripsFromVersion(finalVersion, requireStableData));
+            } else {
+                testAllMessageRoundTripsFromVersion(finalVersion, requireStableData);
+            }
+        }
+
+        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(new ProduceResponseData.TopicProduceResponseCollection(singletonList(
+                new ProduceResponseData.TopicProduceResponse()
+                    .setName(topicName)
+                    .setPartitionResponses(singletonList(
+                        new ProduceResponseData.PartitionProduceResponse()
+                            .setIndex(partitionIndex)
+                            .setErrorCode(errorCode)
+                            .setBaseOffset(baseOffset)))).iterator())));
+
+        Supplier response = () -> new ProduceResponseData()
+                .setResponses(new ProduceResponseData.TopicProduceResponseCollection(singletonList(
+                    new ProduceResponseData.TopicProduceResponse()
+                        .setName(topicName)
+                        .setPartitionResponses(singletonList(
+                             new ProduceResponseData.PartitionProduceResponse()
+                                 .setIndex(partitionIndex)
+                                 .setErrorCode(errorCode)
+                                 .setBaseOffset(baseOffset)
+                                 .setLogAppendTimeMs(logAppendTimeMs)
+                                 .setLogStartOffset(logStartOffset)
+                                 .setRecordErrors(singletonList(
+                                     new ProduceResponseData.BatchIndexAndErrorMessage()
+                                         .setBatchIndex(batchIndex)
+                                         .setBatchIndexErrorMessage(batchIndexErrorMessage)))
+                                 .setErrorMessage(errorMessage)))).iterator()))
+                .setThrottleTimeMs(throttleTimeMs);
+
+        for (short version = 0; version <= ApiKeys.PRODUCE.latestVersion(); version++) {
+            ProduceResponseData responseData = response.get();
+
+            if (version < 8) {
+                responseData.responses().iterator().next().partitionResponses().get(0).setRecordErrors(Collections.emptyList());
+                responseData.responses().iterator().next().partitionResponses().get(0).setErrorMessage(null);
+            }
+
+            if (version < 5) {
+                responseData.responses().iterator().next().partitionResponses().get(0).setLogStartOffset(-1);
+            }
+
+            if (version < 2) {
+                responseData.responses().iterator().next().partitionResponses().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);
+            }
+        }
+    }
+
+    @Test
+    public void testSimpleMessage() throws Exception {
+        final SimpleExampleMessageData message = new SimpleExampleMessageData();
+        message.setMyStruct(new SimpleExampleMessageData.MyStruct().setStructId(25).setArrayInStruct(
+            Collections.singletonList(new SimpleExampleMessageData.StructArray().setArrayFieldId(20))
+        ));
+        message.setMyTaggedStruct(new SimpleExampleMessageData.TaggedStruct().setStructId("abc"));
+
+        message.setProcessId(Uuid.randomUuid());
+        message.setMyNullableString("notNull");
+        message.setMyInt16((short) 3);
+        message.setMyString("test string");
+        SimpleExampleMessageData duplicate = message.duplicate();
+        assertEquals(duplicate, message);
+        assertEquals(message, duplicate);
+        duplicate.setMyTaggedIntArray(Collections.singletonList(123));
+        assertNotEquals(duplicate, message);
+        assertNotEquals(message, duplicate);
+
+        testAllMessageRoundTripsFromVersion((short) 2, message);
+    }
+
+    private void testAllMessageRoundTrips(Message message) throws Exception {
+        testDuplication(message);
+        testAllMessageRoundTripsFromVersion(message.lowestSupportedVersion(), message);
+    }
+
+    private void testDuplication(Message message) {
+        Message duplicate = message.duplicate();
+        assertEquals(duplicate, message);
+        assertEquals(message, duplicate);
+        assertEquals(duplicate.hashCode(), message.hashCode());
+        assertEquals(message.hashCode(), duplicate.hashCode());
+    }
+
+    private void testAllMessageRoundTripsBeforeVersion(short beforeVersion, Message message, Message expected) throws Exception {
+        testAllMessageRoundTripsBetweenVersions((short) 0, beforeVersion, message, expected);
+    }
+
+    /**
+     * @param startVersion - the version we want to start at, inclusive
+     * @param endVersion - the version we want to end at, exclusive
+     */
+    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);
+        testJsonRoundTrip(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, 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().getConstructor().newInstance();
+        buf.flip();
+        message2.read(byteBufferAccessor, version);
+        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(short version, Message message, Message expected) throws Exception {
+        Struct struct = message.toStruct(version);
+        Message message2 = message.getClass().getConstructor().newInstance();
+        message2.fromStruct(struct, version);
+        assertEquals(expected, message2);
+        assertEquals(expected.hashCode(), message2.hashCode());
+        assertEquals(expected.toString(), message2.toString());
+    }
+
+    @SuppressWarnings("unchecked")
+    private void testJsonRoundTrip(short version, Message message, Message expected) throws Exception {
+        String jsonConverter = jsonConverterTypeName(message.getClass().getTypeName());
+        Class converter = Class.forName(jsonConverter);
+        Method writeMethod = converter.getMethod("write", message.getClass(), short.class);
+        JsonNode jsonNode = (JsonNode) writeMethod.invoke(null, message, version);
+        Method readMethod = converter.getMethod("read", JsonNode.class, short.class);
+        Message message2 = (Message) readMethod.invoke(null, jsonNode, version);
+        assertEquals(expected, message2);
+        assertEquals(expected.hashCode(), message2.hashCode());
+        assertEquals(expected.toString(), message2.toString());
+    }
+
+    private static String jsonConverterTypeName(String source) {
+        int outerClassIndex = source.lastIndexOf('$');
+        if (outerClassIndex == -1) {
+            return  source + "JsonConverter";
+        } else {
+            return source.substring(0, outerClassIndex) + "JsonConverter$" +
+                source.substring(outerClassIndex + 1) + "JsonConverter";
+        }
+    }
+
+    /**
+     * Verify that the JSON files support the same message versions as the
+     * schemas accessible through the ApiKey class.
+     */
+    @Test
+    public void testMessageVersions() {
+        for (ApiKeys apiKey : ApiKeys.values()) {
+            Message message = null;
+            try {
+                message = ApiMessageType.fromApiKey(apiKey.id).newRequest();
+            } catch (UnsupportedVersionException e) {
+                fail("No request message spec found for API " + apiKey);
+            }
+            assertTrue("Request message spec for " + apiKey + " only " +
+                    "supports versions up to " + message.highestSupportedVersion(),
+                apiKey.latestVersion() <= message.highestSupportedVersion());
+            try {
+                message = ApiMessageType.fromApiKey(apiKey.id).newResponse();
+            } catch (UnsupportedVersionException e) {
+                fail("No response message spec found for API " + apiKey);
+            }
+            assertTrue("Response message spec for " + apiKey + " only " +
+                    "supports versions up to " + message.highestSupportedVersion(),
+                apiKey.latestVersion() <= message.highestSupportedVersion());
+        }
+    }
+
+    /**
+     * Test that the JSON request files match the schemas accessible through the ApiKey class.
+     */
+    @Test
+    public void testRequestSchemas() {
+        for (ApiKeys apiKey : ApiKeys.values()) {
+            Schema[] manualSchemas = apiKey.requestSchemas;
+            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++) {
+                try {
+                    if (generatedSchemas[v] != null) {
+                        compareTypes(manualSchemas[v], generatedSchemas[v]);
+                    }
+                } catch (Exception e) {
+                    throw new RuntimeException("Failed to compare request schemas " +
+                        "for version " + v + " of " + apiKey, e);
+                }
+            }
+        }
+    }
+
+    /**
+     * Test that the JSON response files match the schemas accessible through the ApiKey class.
+     */
+    @Test
+    public void testResponseSchemas() {
+        for (ApiKeys apiKey : ApiKeys.values()) {
+            Schema[] manualSchemas = apiKey.responseSchemas;
+            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++) {
+                try {
+                    if (generatedSchemas[v] != null) {
+                        compareTypes(manualSchemas[v], generatedSchemas[v]);
+                    }
+                } catch (Exception e) {
+                    throw new RuntimeException("Failed to compare response schemas " +
+                        "for version " + v + " of " + apiKey, e);
+                }
+            }
+        }
+    }
+
+    private static class NamedType {
+        final String name;
+        final Type type;
+
+        NamedType(String name, Type type) {
+            this.name = name;
+            this.type = type;
+        }
+
+        boolean hasSimilarType(NamedType other) {
+            if (type.getClass().equals(other.type.getClass())) {
+                return true;
+            }
+            if (type.getClass().equals(Type.RECORDS.getClass())) {
+                return other.type.getClass().equals(Type.NULLABLE_BYTES.getClass());
+            } else if (type.getClass().equals(Type.NULLABLE_BYTES.getClass())) {
+                return other.type.getClass().equals(Type.RECORDS.getClass());
+            }
+            return false;
+        }
+
+        @Override
+        public String toString() {
+            return name + "[" + type + "]";
+        }
+    }
+
+    private static void compareTypes(Schema schemaA, Schema schemaB) {
+        compareTypes(new NamedType("schemaA", schemaA),
+                     new NamedType("schemaB", schemaB));
+    }
+
+    private static void compareTypes(NamedType typeA, NamedType typeB) {
+        List listA = flatten(typeA);
+        List listB = flatten(typeB);
+        if (listA.size() != listB.size()) {
+            throw new RuntimeException("Can't match up structures: typeA has " +
+                Utils.join(listA, ", ") + ", but typeB has " +
+                Utils.join(listB, ", "));
+        }
+        for (int i = 0; i < listA.size(); i++) {
+            NamedType entryA = listA.get(i);
+            NamedType entryB = listB.get(i);
+            if (!entryA.hasSimilarType(entryB)) {
+                throw new RuntimeException("Type " + entryA + " in schema A " +
+                    "does not match type " + entryB + " in schema B.");
+            }
+            if (entryA.type.isNullable() != entryB.type.isNullable()) {
+                throw new RuntimeException(String.format(
+                    "Type %s in Schema A is %s, but type %s in " +
+                        "Schema B is %s",
+                    entryA, entryA.type.isNullable() ? "nullable" : "non-nullable",
+                    entryB, entryB.type.isNullable() ? "nullable" : "non-nullable"));
+            }
+            if (entryA.type.isArray()) {
+                compareTypes(new NamedType(entryA.name, entryA.type.arrayElementType().get()),
+                             new NamedType(entryB.name, entryB.type.arrayElementType().get()));
+            }
+        }
+    }
+
+    /**
+     * We want to remove Schema nodes from the hierarchy before doing
+     * our comparison.  The reason is because Schema nodes don't change what
+     * is written to the wire.  Schema(STRING, Schema(INT, SHORT)) is equivalent to
+     * Schema(STRING, INT, SHORT).  This function translates schema nodes into their
+     * component types.
+     */
+    private static List flatten(NamedType type) {
+        if (!(type.type instanceof Schema)) {
+            return singletonList(type);
+        }
+        Schema schema = (Schema) type.type;
+        ArrayList results = new ArrayList<>();
+        for (BoundField field : schema.fields()) {
+            results.addAll(flatten(new NamedType(field.def.name, field.def.type)));
+        }
+        return results;
+    }
+
+    @Test
+    public void testDefaultValues() {
+        verifyWriteRaisesUve((short) 0, "validateOnly",
+            new CreateTopicsRequestData().setValidateOnly(true));
+        verifyWriteSucceeds((short) 0,
+            new CreateTopicsRequestData().setValidateOnly(false));
+        verifyWriteSucceeds((short) 0,
+            new OffsetCommitRequestData().setRetentionTimeMs(123));
+        verifyWriteRaisesUve((short) 5, "forgotten",
+            new FetchRequestData().setForgottenTopicsData(singletonList(
+                new FetchRequestData.ForgottenTopic().setTopic("foo"))));
+    }
+
+    @Test
+    public void testNonIgnorableFieldWithDefaultNull() {
+        // 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() {
+        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);
+    }
+
+    @Test
+    public void testUnknownTaggedFields() {
+        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) {
+        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) {
+        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) {
+        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());
+    }
+
+    @Test
+    public void testCompareWithUnknownTaggedFields() {
+        CreateTopicsRequestData createTopics = new CreateTopicsRequestData();
+        createTopics.setTimeoutMs(123);
+        CreateTopicsRequestData createTopics2 = new CreateTopicsRequestData();
+        createTopics2.setTimeoutMs(123);
+        assertEquals(createTopics, createTopics2);
+        assertEquals(createTopics2, createTopics);
+        // Call the accessor, which will create a new empty list.
+        createTopics.unknownTaggedFields();
+        // Verify that the equalities still hold after the new empty list has been created.
+        assertEquals(createTopics, createTopics2);
+        assertEquals(createTopics2, createTopics);
+        createTopics.unknownTaggedFields().add(new RawTaggedField(0, new byte[] {0}));
+        assertNotEquals(createTopics, createTopics2);
+        assertNotEquals(createTopics2, createTopics);
+        createTopics2.unknownTaggedFields().add(new RawTaggedField(0, new byte[] {0}));
+        assertEquals(createTopics, createTopics2);
+        assertEquals(createTopics2, createTopics);
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/message/RecordsSerdeTest.java b/clients/src/test/java/org/apache/kafka/common/message/RecordsSerdeTest.java
new file mode 100644
index 0000000000000..f3bc05d23da97
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/message/RecordsSerdeTest.java
@@ -0,0 +1,117 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.ObjectSerializationCache;
+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.MemoryRecords;
+import org.apache.kafka.common.record.SimpleRecord;
+import org.junit.Test;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+
+public class RecordsSerdeTest {
+
+    @Test
+    public void testSerdeRecords() throws Exception {
+        MemoryRecords records = MemoryRecords.withRecords(CompressionType.NONE,
+            new SimpleRecord("foo".getBytes()),
+            new SimpleRecord("bar".getBytes()));
+
+        SimpleRecordsMessageData message = new SimpleRecordsMessageData()
+            .setTopic("foo")
+            .setRecordSet(records);
+
+        testAllRoundTrips(message);
+    }
+
+    @Test
+    public void testSerdeNullRecords() throws Exception {
+        SimpleRecordsMessageData message = new SimpleRecordsMessageData()
+            .setTopic("foo");
+        assertNull(message.recordSet());
+
+        testAllRoundTrips(message);
+    }
+
+    @Test
+    public void testSerdeEmptyRecords() throws Exception {
+        SimpleRecordsMessageData message = new SimpleRecordsMessageData()
+            .setTopic("foo")
+            .setRecordSet(MemoryRecords.EMPTY);
+        testAllRoundTrips(message);
+    }
+
+    private void testAllRoundTrips(SimpleRecordsMessageData message) throws Exception {
+        for (short version = SimpleRecordsMessageData.LOWEST_SUPPORTED_VERSION;
+             version <= SimpleRecordsMessageData.HIGHEST_SUPPORTED_VERSION;
+             version++) {
+            testRoundTrip(message, version);
+        }
+    }
+
+    private void testRoundTrip(SimpleRecordsMessageData message, short version) throws IOException {
+        ByteBuffer buf = serialize(message, version);
+
+        SimpleRecordsMessageData message2 = deserialize(buf.duplicate(), version);
+        assertEquals(message, message2);
+        assertEquals(message.hashCode(), message2.hashCode());
+
+        // Check struct serialization as well
+        assertEquals(buf, serializeThroughStruct(message, version));
+        SimpleRecordsMessageData messageFromStruct = deserializeThroughStruct(buf.duplicate(), version);
+        assertEquals(message, messageFromStruct);
+        assertEquals(message.hashCode(), messageFromStruct.hashCode());
+    }
+
+    private SimpleRecordsMessageData deserializeThroughStruct(ByteBuffer buffer, short version) {
+        Schema schema = SimpleRecordsMessageData.SCHEMAS[version];
+        Struct struct = schema.read(buffer);
+        return new SimpleRecordsMessageData(struct, version);
+    }
+
+    private SimpleRecordsMessageData deserialize(ByteBuffer buffer, short version) {
+        ByteBufferAccessor readable = new ByteBufferAccessor(buffer);
+        return new SimpleRecordsMessageData(readable, version);
+    }
+
+    private ByteBuffer serializeThroughStruct(SimpleRecordsMessageData message, short version) {
+        Struct struct = message.toStruct(version);
+        ByteBuffer buffer = ByteBuffer.allocate(struct.sizeOf());
+        struct.writeTo(buffer);
+        buffer.flip();
+        return buffer;
+    }
+
+    private ByteBuffer serialize(SimpleRecordsMessageData message, short version) {
+        ObjectSerializationCache cache = new ObjectSerializationCache();
+        int totalMessageSize = message.size(cache, version);
+        ByteBuffer buffer = ByteBuffer.allocate(totalMessageSize);
+        ByteBufferAccessor writer = new ByteBufferAccessor(buffer);
+        message.write(writer, cache, version);
+        buffer.flip();
+        return buffer;
+    }
+
+}
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..b3e45e5a64e09
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/message/SimpleExampleMessageTest.java
@@ -0,0 +1,404 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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 com.fasterxml.jackson.databind.JsonNode;
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.ObjectSerializationCache;
+import org.apache.kafka.common.protocol.types.Schema;
+import org.apache.kafka.common.protocol.types.Struct;
+import org.apache.kafka.common.utils.ByteUtils;
+import org.junit.Test;
+
+import java.nio.ByteBuffer;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.function.Consumer;
+
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+
+public class SimpleExampleMessageTest {
+
+    @Test
+    public void shouldStoreField() {
+        final Uuid uuid = Uuid.randomUuid();
+        final ByteBuffer buf = ByteBuffer.wrap(new byte[] {1, 2, 3});
+
+        final SimpleExampleMessageData out = new SimpleExampleMessageData();
+        out.setProcessId(uuid);
+        out.setZeroCopyByteBuffer(buf);
+
+        assertEquals(uuid, out.processId());
+        assertEquals(buf, out.zeroCopyByteBuffer());
+
+        out.setNullableZeroCopyByteBuffer(null);
+        assertNull(out.nullableZeroCopyByteBuffer());
+        out.setNullableZeroCopyByteBuffer(buf);
+        assertEquals(buf, out.nullableZeroCopyByteBuffer());
+    }
+
+    @Test
+    public void shouldThrowIfCannotWriteNonIgnorableField() {
+        // processId is not supported in v0 and is not marked as ignorable
+
+        final SimpleExampleMessageData out = new SimpleExampleMessageData().setProcessId(Uuid.randomUuid());
+        assertThrows(UnsupportedVersionException.class, () ->
+                out.write(new ByteBufferAccessor(ByteBuffer.allocate(64)), new ObjectSerializationCache(), (short) 0));
+        assertThrows(UnsupportedVersionException.class, () -> out.toStruct((short) 0));
+    }
+
+    @Test
+    public void shouldDefaultField() {
+        final SimpleExampleMessageData out = new SimpleExampleMessageData();
+        assertEquals(Uuid.fromString("AAAAAAAAAAAAAAAAAAAAAA"), out.processId());
+        assertEquals(ByteUtils.EMPTY_BUF, out.zeroCopyByteBuffer());
+        assertEquals(ByteUtils.EMPTY_BUF, out.nullableZeroCopyByteBuffer());
+    }
+
+    @Test
+    public void shouldRoundTripFieldThroughStruct() {
+        final Uuid uuid = Uuid.randomUuid();
+        final ByteBuffer buf = ByteBuffer.wrap(new byte[] {1, 2, 3});
+        final SimpleExampleMessageData out = new SimpleExampleMessageData();
+        out.setProcessId(uuid);
+        out.setZeroCopyByteBuffer(buf);
+
+        final Struct struct = out.toStruct((short) 1);
+        final SimpleExampleMessageData in = new SimpleExampleMessageData();
+        in.fromStruct(struct, (short) 1);
+
+        buf.rewind();
+
+        assertEquals(uuid, in.processId());
+        assertEquals(buf, in.zeroCopyByteBuffer());
+        assertEquals(ByteUtils.EMPTY_BUF, in.nullableZeroCopyByteBuffer());
+    }
+
+    @Test
+    public void shouldRoundTripFieldThroughStructWithNullable() {
+        final Uuid uuid = Uuid.randomUuid();
+        final ByteBuffer buf1 = ByteBuffer.wrap(new byte[] {1, 2, 3});
+        final ByteBuffer buf2 = ByteBuffer.wrap(new byte[] {4, 5, 6});
+        final SimpleExampleMessageData out = new SimpleExampleMessageData();
+        out.setProcessId(uuid);
+        out.setZeroCopyByteBuffer(buf1);
+        out.setNullableZeroCopyByteBuffer(buf2);
+
+        final Struct struct = out.toStruct((short) 1);
+        final SimpleExampleMessageData in = new SimpleExampleMessageData();
+        in.fromStruct(struct, (short) 1);
+
+        buf1.rewind();
+        buf2.rewind();
+
+        assertEquals(uuid, in.processId());
+        assertEquals(buf1, in.zeroCopyByteBuffer());
+        assertEquals(buf2, in.nullableZeroCopyByteBuffer());
+    }
+
+    @Test
+    public void shouldRoundTripFieldThroughBuffer() {
+        final Uuid uuid = Uuid.randomUuid();
+        final ByteBuffer buf = ByteBuffer.wrap(new byte[] {1, 2, 3});
+        final SimpleExampleMessageData out = new SimpleExampleMessageData();
+        out.setProcessId(uuid);
+        out.setZeroCopyByteBuffer(buf);
+
+        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);
+
+        buf.rewind();
+
+        assertEquals(uuid, in.processId());
+        assertEquals(buf, in.zeroCopyByteBuffer());
+        assertEquals(ByteUtils.EMPTY_BUF, in.nullableZeroCopyByteBuffer());
+    }
+
+    @Test
+    public void shouldRoundTripFieldThroughBufferWithNullable() {
+        final Uuid uuid = Uuid.randomUuid();
+        final ByteBuffer buf1 = ByteBuffer.wrap(new byte[] {1, 2, 3});
+        final ByteBuffer buf2 = ByteBuffer.wrap(new byte[] {4, 5, 6});
+        final SimpleExampleMessageData out = new SimpleExampleMessageData();
+        out.setProcessId(uuid);
+        out.setZeroCopyByteBuffer(buf1);
+        out.setNullableZeroCopyByteBuffer(buf2);
+
+        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);
+
+        buf1.rewind();
+        buf2.rewind();
+
+        assertEquals(uuid, in.processId());
+        assertEquals(buf1, in.zeroCopyByteBuffer());
+        assertEquals(buf2, in.nullableZeroCopyByteBuffer());
+    }
+
+    @Test
+    public void shouldImplementEqualsAndHashCode() {
+        final Uuid uuid = Uuid.randomUuid();
+        final ByteBuffer buf = ByteBuffer.wrap(new byte[] {1, 2, 3});
+        final SimpleExampleMessageData a = new SimpleExampleMessageData();
+        a.setProcessId(uuid);
+        a.setZeroCopyByteBuffer(buf);
+
+        final SimpleExampleMessageData b = new SimpleExampleMessageData();
+        b.setProcessId(uuid);
+        b.setZeroCopyByteBuffer(buf);
+
+        assertEquals(a, b);
+        assertEquals(a.hashCode(), b.hashCode());
+        // just tagging this on here
+        assertEquals(a.toString(), b.toString());
+
+        a.setNullableZeroCopyByteBuffer(buf);
+        b.setNullableZeroCopyByteBuffer(buf);
+
+        assertEquals(a, b);
+        assertEquals(a.hashCode(), b.hashCode());
+        assertEquals(a.toString(), b.toString());
+
+        a.setNullableZeroCopyByteBuffer(null);
+        b.setNullableZeroCopyByteBuffer(null);
+
+        assertEquals(a, b);
+        assertEquals(a.hashCode(), b.hashCode());
+        assertEquals(a.toString(), b.toString());
+    }
+
+    @Test
+    public void testMyTaggedIntArray() {
+        // Verify that the tagged int array reads as empty when not set.
+        testRoundTrip(new SimpleExampleMessageData(),
+            message -> assertEquals(Collections.emptyList(), message.myTaggedIntArray()));
+
+        // Verify that we can set a tagged array of ints.
+        testRoundTrip(new SimpleExampleMessageData().
+                setMyTaggedIntArray(Arrays.asList(1, 2, 3)),
+            message -> assertEquals(Arrays.asList(1, 2, 3), message.myTaggedIntArray()));
+    }
+
+    @Test
+    public void testMyNullableString() {
+        // Verify that the tagged field reads as null when not set.
+        testRoundTrip(new SimpleExampleMessageData(),
+            message -> assertTrue(message.myNullableString() == null));
+
+        // Verify that we can set and retrieve a string for the tagged field.
+        testRoundTrip(new SimpleExampleMessageData().setMyNullableString("foobar"),
+            message -> assertEquals("foobar", message.myNullableString()));
+    }
+
+    @Test
+    public void testMyInt16() {
+        // Verify that the tagged field reads as 123 when not set.
+        testRoundTrip(new SimpleExampleMessageData(),
+            message -> assertEquals((short) 123, message.myInt16()));
+
+        testRoundTrip(new SimpleExampleMessageData().setMyInt16((short) 456),
+            message -> assertEquals((short) 456, message.myInt16()));
+    }
+
+    @Test
+    public void testMyString() {
+        // Verify that the tagged field reads as empty when not set.
+        testRoundTrip(new SimpleExampleMessageData(),
+            message -> assertEquals("", message.myString()));
+
+        testRoundTrip(new SimpleExampleMessageData().setMyString("abc"),
+            message -> assertEquals("abc", message.myString()));
+    }
+
+    @Test
+    public void testMyBytes() {
+        // Verify that the tagged field reads as empty when not set.
+        testRoundTrip(new SimpleExampleMessageData(),
+            message -> assertArrayEquals(new byte[0], message.myBytes()));
+
+        testRoundTrip(new SimpleExampleMessageData().
+                setMyBytes(new byte[] {0x43, 0x66}),
+            message -> assertArrayEquals(new byte[] {0x43, 0x66},
+                message.myBytes()));
+
+        testRoundTrip(new SimpleExampleMessageData().setMyBytes(null),
+            message -> assertTrue(message.myBytes() == null));
+    }
+
+    @Test
+    public void testTaggedUuid() {
+        testRoundTrip(new SimpleExampleMessageData(),
+            message -> assertEquals(
+                Uuid.fromString("212d54944a8b4fdf94b388b470beb367"),
+                message.taggedUuid()));
+
+        testRoundTrip(new SimpleExampleMessageData().
+                setTaggedUuid(Uuid.fromString("0123456789abcdef0123456789abcdef")),
+            message -> assertEquals(
+                Uuid.fromString("0123456789abcdef0123456789abcdef"),
+                message.taggedUuid()));
+    }
+
+    @Test
+    public void testTaggedLong() {
+        testRoundTrip(new SimpleExampleMessageData(),
+            message -> assertEquals(0xcafcacafcacafcaL,
+                message.taggedLong()));
+
+        testRoundTrip(new SimpleExampleMessageData().
+                setMyString("blah").
+                setMyTaggedIntArray(Arrays.asList(4)).
+                setTaggedLong(0x123443211234432L),
+            message -> assertEquals(0x123443211234432L,
+                message.taggedLong()));
+    }
+
+    @Test
+    public void testMyStruct() {
+        // Verify that we can set and retrieve a nullable struct object.
+        SimpleExampleMessageData.MyStruct myStruct =
+            new SimpleExampleMessageData.MyStruct().setStructId(10).setArrayInStruct(
+                Collections.singletonList(new SimpleExampleMessageData.StructArray().setArrayFieldId(20))
+            );
+        testRoundTrip(new SimpleExampleMessageData().setMyStruct(myStruct),
+            message -> assertEquals(myStruct, message.myStruct()), (short) 2);
+    }
+
+    @Test(expected = UnsupportedVersionException.class)
+    public void testMyStructUnsupportedVersion() {
+        SimpleExampleMessageData.MyStruct myStruct =
+                new SimpleExampleMessageData.MyStruct().setStructId(10);
+        // Check serialization throws exception for unsupported version
+        testRoundTrip(new SimpleExampleMessageData().setMyStruct(myStruct), (short) 1);
+    }
+
+    /**
+     * Check following cases:
+     * 1. Tagged struct can be serialized/deserialized for version it is supported
+     * 2. Tagged struct doesn't matter for versions it is not declared.
+     */
+    @Test
+    public void testMyTaggedStruct() {
+        // Verify that we can set and retrieve a nullable struct object.
+        SimpleExampleMessageData.TaggedStruct myStruct =
+            new SimpleExampleMessageData.TaggedStruct().setStructId("abc");
+        testRoundTrip(new SimpleExampleMessageData().setMyTaggedStruct(myStruct),
+            message -> assertEquals(myStruct, message.myTaggedStruct()), (short) 2);
+
+        // Not setting field works for both version 1 and version 2 protocol
+        testRoundTrip(new SimpleExampleMessageData().setMyString("abc"),
+            message -> assertEquals("abc", message.myString()), (short) 1);
+        testRoundTrip(new SimpleExampleMessageData().setMyString("abc"),
+            message -> assertEquals("abc", message.myString()), (short) 2);
+    }
+
+    @Test
+    public void testCommonStruct() {
+        SimpleExampleMessageData message = new SimpleExampleMessageData();
+        message.setMyCommonStruct(new SimpleExampleMessageData.TestCommonStruct()
+            .setFoo(1)
+            .setBar(2));
+        message.setMyOtherCommonStruct(new SimpleExampleMessageData.TestCommonStruct()
+            .setFoo(3)
+            .setBar(4));
+        testRoundTrip(message, (short) 2);
+    }
+
+    private ByteBuffer serialize(SimpleExampleMessageData message, short version) {
+        ObjectSerializationCache cache = new ObjectSerializationCache();
+        int size = message.size(cache, version);
+        ByteBuffer buf = ByteBuffer.allocate(size);
+        message.write(new ByteBufferAccessor(buf), cache, version);
+        buf.flip();
+        assertEquals(size, buf.remaining());
+        return buf;
+    }
+
+    private SimpleExampleMessageData deserialize(ByteBuffer buf, short version) {
+        SimpleExampleMessageData message = new SimpleExampleMessageData();
+        message.read(new ByteBufferAccessor(buf.duplicate()), version);
+        return message;
+    }
+
+    private ByteBuffer serializeThroughStruct(SimpleExampleMessageData message, short version) {
+        Struct struct = message.toStruct(version);
+        int size = struct.sizeOf();
+        ByteBuffer buf = ByteBuffer.allocate(size);
+        struct.writeTo(buf);
+        buf.flip();
+        assertEquals(size, buf.remaining());
+        return buf;
+    }
+
+    private SimpleExampleMessageData deserializeThroughStruct(ByteBuffer buf, short version) {
+        Schema schema = SimpleExampleMessageData.SCHEMAS[version];
+        Struct struct = schema.read(buf);
+        return new SimpleExampleMessageData(struct, version);
+    }
+
+    private void testRoundTrip(SimpleExampleMessageData message, short version) {
+        testRoundTrip(message, m -> { }, version);
+    }
+
+    private void testRoundTrip(SimpleExampleMessageData message,
+                               Consumer validator) {
+        testRoundTrip(message, validator, (short) 1);
+    }
+
+    private void testRoundTrip(SimpleExampleMessageData message,
+                               Consumer validator,
+                               short version) {
+        validator.accept(message);
+        ByteBuffer buf = serialize(message, version);
+
+        SimpleExampleMessageData message2 = deserialize(buf.duplicate(), version);
+        validator.accept(message2);
+        assertEquals(message, message2);
+        assertEquals(message.hashCode(), message2.hashCode());
+
+        // Check struct serialization as well
+        assertEquals(buf, serializeThroughStruct(message, version));
+        SimpleExampleMessageData messageFromStruct = deserializeThroughStruct(buf.duplicate(), version);
+        validator.accept(messageFromStruct);
+        assertEquals(message, messageFromStruct);
+        assertEquals(message.hashCode(), messageFromStruct.hashCode());
+
+        // Check JSON serialization
+        JsonNode serializedJson = SimpleExampleMessageDataJsonConverter.write(message, version);
+        SimpleExampleMessageData messageFromJson = SimpleExampleMessageDataJsonConverter.read(serializedJson, version);
+        validator.accept(messageFromJson);
+        assertEquals(message, messageFromJson);
+        assertEquals(message.hashCode(), messageFromJson.hashCode());
+    }
+
+}
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 98e49f3abf5f4..15eebdcebf853 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
@@ -16,13 +16,19 @@
  */
 package org.apache.kafka.common.metrics;
 
+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.apache.kafka.common.utils.Time;
 import org.junit.Test;
 
 import javax.management.MBeanServer;
 import javax.management.ObjectName;
 import java.lang.management.ManagementFactory;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.Map;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
@@ -35,26 +41,33 @@ public void testJmxRegistration() throws Exception {
         Metrics metrics = new Metrics();
         MBeanServer server = ManagementFactory.getPlatformMBeanServer();
         try {
-            metrics.addReporter(new JmxReporter());
+            JmxReporter reporter = new JmxReporter();
+            metrics.addReporter(reporter);
 
             assertFalse(server.isRegistered(new ObjectName(":type=grp1")));
 
             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(0.0, server.getAttribute(new ObjectName(":type=grp1"), "pack.bean1.avg"));
+            assertEquals(Double.NaN, server.getAttribute(new ObjectName(":type=grp1"), "pack.bean1.avg"));
             assertTrue(server.isRegistered(new ObjectName(":type=grp2")));
             assertEquals(0.0, server.getAttribute(new ObjectName(":type=grp2"), "pack.bean2.total"));
 
-            metrics.removeMetric(metrics.metricName("pack.bean1.avg", "grp1"));
+            MetricName metricName = metrics.metricName("pack.bean1.avg", "grp1");
+            String mBeanName = JmxReporter.getMBeanName("", metricName);
+            assertTrue(reporter.containsMbean(mBeanName));
+            metrics.removeMetric(metricName);
+            assertFalse(reporter.containsMbean(mBeanName));
 
             assertFalse(server.isRegistered(new ObjectName(":type=grp1")));
             assertTrue(server.isRegistered(new ObjectName(":type=grp2")));
             assertEquals(0.0, server.getAttribute(new ObjectName(":type=grp2"), "pack.bean2.total"));
 
-            metrics.removeMetric(metrics.metricName("pack.bean2.total", "grp2"));
+            metricName = metrics.metricName("pack.bean2.total", "grp2");
+            metrics.removeMetric(metricName);
+            assertFalse(reporter.containsMbean(mBeanName));
 
             assertFalse(server.isRegistered(new ObjectName(":type=grp1")));
             assertFalse(server.isRegistered(new ObjectName(":type=grp2")));
@@ -71,11 +84,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"));
@@ -103,4 +116,82 @@ public void testJmxRegistrationSanitization() throws Exception {
             metrics.close();
         }
     }
+
+    @Test
+    public void testPredicateAndDynamicReload() throws Exception {
+        Metrics metrics = new Metrics();
+        MBeanServer server = ManagementFactory.getPlatformMBeanServer();
+
+        Map configs = new HashMap<>();
+
+        configs.put(JmxReporter.EXCLUDE_CONFIG,
+                    JmxReporter.getMBeanName("", metrics.metricName("pack.bean2.total", "grp2")));
+
+        try {
+            JmxReporter reporter = new JmxReporter();
+            reporter.configure(configs);
+            metrics.addReporter(reporter);
+
+            Sensor sensor = metrics.sensor("kafka.requests");
+            sensor.add(metrics.metricName("pack.bean2.avg", "grp1"), new Avg());
+            sensor.add(metrics.metricName("pack.bean2.total", "grp2"), new CumulativeSum());
+            sensor.record();
+
+            assertTrue(server.isRegistered(new ObjectName(":type=grp1")));
+            assertEquals(1.0, server.getAttribute(new ObjectName(":type=grp1"), "pack.bean2.avg"));
+            assertFalse(server.isRegistered(new ObjectName(":type=grp2")));
+
+            sensor.record();
+
+            configs.put(JmxReporter.EXCLUDE_CONFIG,
+                        JmxReporter.getMBeanName("", metrics.metricName("pack.bean2.avg", "grp1")));
+
+            reporter.reconfigure(configs);
+
+            assertFalse(server.isRegistered(new ObjectName(":type=grp1")));
+            assertTrue(server.isRegistered(new ObjectName(":type=grp2")));
+            assertEquals(2.0, server.getAttribute(new ObjectName(":type=grp2"), "pack.bean2.total"));
+
+            metrics.removeMetric(metrics.metricName("pack.bean2.total", "grp2"));
+            assertFalse(server.isRegistered(new ObjectName(":type=grp2")));
+        } finally {
+            metrics.close();
+        }
+    }
+
+    @Test
+    public void testJmxPrefix() throws Exception {
+        JmxReporter reporter = new JmxReporter();
+        MetricsContext metricsContext = new KafkaMetricsContext("kafka.server");
+        MetricConfig metricConfig = new MetricConfig();
+        Metrics metrics = new Metrics(metricConfig, new ArrayList<>(Arrays.asList(reporter)), Time.SYSTEM, metricsContext);
+
+        MBeanServer server = ManagementFactory.getPlatformMBeanServer();
+        try {
+            Sensor sensor = metrics.sensor("kafka.requests");
+            sensor.add(metrics.metricName("pack.bean1.avg", "grp1"), new Avg());
+            assertEquals("kafka.server", server.getObjectInstance(new ObjectName("kafka.server:type=grp1")).getObjectName().getDomain());
+        } finally {
+            metrics.close();
+        }
+    }
+
+    @Test
+    public void testDeprecatedJmxPrefixWithDefaultMetrics() throws Exception {
+        @SuppressWarnings("deprecation")
+        JmxReporter reporter = new JmxReporter("my-prefix");
+
+        // for backwards compatibility, ensure prefix does not get overridden by the default empty namespace in metricscontext
+        MetricConfig metricConfig = new MetricConfig();
+        Metrics metrics = new Metrics(metricConfig, new ArrayList<>(Arrays.asList(reporter)), Time.SYSTEM);
+
+        MBeanServer server = ManagementFactory.getPlatformMBeanServer();
+        try {
+            Sensor sensor = metrics.sensor("my-sensor");
+            sensor.add(metrics.metricName("pack.bean1.avg", "grp1"), new Avg());
+            assertEquals("my-prefix", server.getObjectInstance(new ObjectName("my-prefix:type=grp1")).getObjectName().getDomain());
+        } finally {
+            metrics.close();
+        }
+    }
 }
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 963a66abb2914..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;
@@ -33,8 +33,8 @@
 import java.util.List;
 
 import static org.hamcrest.core.IsInstanceOf.instanceOf;
+import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertThat;
 import static org.junit.Assert.fail;
 
 public class KafkaMbeanTest {
@@ -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/KafkaMetricsContextTest.java b/clients/src/test/java/org/apache/kafka/common/metrics/KafkaMetricsContextTest.java
new file mode 100644
index 0000000000000..55b24f89b0744
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/metrics/KafkaMetricsContextTest.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.common.metrics;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import org.junit.Before;
+import org.junit.Test;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertThrows;
+
+public class KafkaMetricsContextTest {
+
+    private static final String SAMPLE_NAMESPACE = "sample-ns";
+
+    private static final String LABEL_A_KEY = "label-a";
+    private static final String LABEL_A_VALUE = "label-a-value";
+
+    private String namespace;
+    private Map labels;
+    private KafkaMetricsContext context;
+
+    @Before
+    public void beforeEach() {
+        namespace = SAMPLE_NAMESPACE;
+        labels = new HashMap<>();
+        labels.put(LABEL_A_KEY, LABEL_A_VALUE);
+    }
+
+    @Test
+    public void testCreationWithValidNamespaceAndNoLabels() {
+        labels.clear();
+        context = new KafkaMetricsContext(namespace, labels);
+
+        assertEquals(1, context.contextLabels().size());
+        assertEquals(namespace, context.contextLabels().get(MetricsContext.NAMESPACE));
+    }
+
+    @Test
+    public void testCreationWithValidNamespaceAndLabels() {
+        context = new KafkaMetricsContext(namespace, labels);
+
+        assertEquals(2, context.contextLabels().size());
+        assertEquals(namespace, context.contextLabels().get(MetricsContext.NAMESPACE));
+        assertEquals(LABEL_A_VALUE, context.contextLabels().get(LABEL_A_KEY));
+    }
+
+    @Test
+    public void testCreationWithValidNamespaceAndNullLabelValues() {
+        labels.put(LABEL_A_KEY, null);
+        context = new KafkaMetricsContext(namespace, labels);
+
+        assertEquals(2, context.contextLabels().size());
+        assertEquals(namespace, context.contextLabels().get(MetricsContext.NAMESPACE));
+        assertNull(context.contextLabels().get(LABEL_A_KEY));
+    }
+
+    @Test
+    public void testCreationWithNullNamespaceAndLabels() {
+        context = new KafkaMetricsContext(null, labels);
+
+        assertEquals(2, context.contextLabels().size());
+        assertNull(context.contextLabels().get(MetricsContext.NAMESPACE));
+        assertEquals(LABEL_A_VALUE, context.contextLabels().get(LABEL_A_KEY));
+    }
+
+    @Test
+    public void testKafkaMetricsContextLabelsAreImmutable() {
+        context = new KafkaMetricsContext(namespace, labels);
+        assertThrows(UnsupportedOperationException.class, () -> context.contextLabels().clear());
+    }
+}
\ No newline at end of file
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 55f8e2349bf65..daa0ae4d6ad6a 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
@@ -16,6 +16,8 @@
  */
 package org.apache.kafka.common.metrics;
 
+import static java.util.Collections.emptyList;
+import static java.util.Collections.singletonList;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
@@ -23,15 +25,27 @@
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
+import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
+import java.util.Deque;
+import java.util.List;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Random;
+import java.util.concurrent.ConcurrentLinkedDeque;
+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.function.Function;
+
 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;
@@ -39,28 +53,37 @@
 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.Total;
+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.Value;
 import org.apache.kafka.common.utils.MockTime;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
-@SuppressWarnings("deprecation")
 public class MetricsTest {
+    private static final Logger log = LoggerFactory.getLogger(MetricsTest.class);
 
     private static final double EPS = 0.000001;
     private MockTime time = new MockTime();
     private MetricConfig config = new MetricConfig();
     private Metrics metrics;
+    private ExecutorService executorService;
 
     @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
-    public void tearDown() {
+    public void tearDown() throws Exception {
+        if (executorService != null) {
+            executorService.shutdownNow();
+            executorService.awaitTermination(5, TimeUnit.SECONDS);
+        }
         this.metrics.close();
     }
 
@@ -83,6 +106,10 @@ public void testMetricName() {
 
     @Test
     public void testSimpleStats() throws Exception {
+        verifyStats(m -> (double) m.metricValue());
+    }
+
+    private void verifyStats(Function metricValueFunc) {
         ConstantMeasurable measurable = new ConstantMeasurable();
 
         metrics.addMetric(metrics.metricName("direct.measurable", "grp1", "The fraction of time an appender waits for space allocation."), measurable);
@@ -92,15 +119,15 @@ public void testSimpleStats() throws Exception {
         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;
@@ -112,38 +139,38 @@ public void testSimpleStats() throws Exception {
         // prior to any time passing
         double elapsedSecs = (config.timeWindowMs() * (config.samples() - 1)) / 1000.0;
         assertEquals(String.format("Occurrences(0...%d) = %f", count, count / elapsedSecs), count / elapsedSecs,
-                     metrics.metrics().get(metrics.metricName("test.occurences", "grp1")).value(), EPS);
+                     metricValueFunc.apply(metrics.metrics().get(metrics.metricName("test.occurences", "grp1"))), EPS);
 
         // pretend 2 seconds passed...
         long sleepTimeMs = 2;
         time.sleep(sleepTimeMs * 1000);
         elapsedSecs += sleepTimeMs;
 
-        assertEquals("s2 reflects the constant value", 5.0, metrics.metrics().get(metrics.metricName("s2.total", "grp1")).value(), EPS);
-        assertEquals("Avg(0...9) = 4.5", 4.5, metrics.metrics().get(metrics.metricName("test.avg", "grp1")).value(), EPS);
-        assertEquals("Max(0...9) = 9", count - 1, metrics.metrics().get(metrics.metricName("test.max", "grp1")).value(), EPS);
-        assertEquals("Min(0...9) = 0", 0.0, metrics.metrics().get(metrics.metricName("test.min", "grp1")).value(), EPS);
+        assertEquals("s2 reflects the constant value", 5.0, metricValueFunc.apply(metrics.metrics().get(metrics.metricName("s2.total", "grp1"))), EPS);
+        assertEquals("Avg(0...9) = 4.5", 4.5, metricValueFunc.apply(metrics.metrics().get(metrics.metricName("test.avg", "grp1"))), EPS);
+        assertEquals("Max(0...9) = 9", count - 1,  metricValueFunc.apply(metrics.metrics().get(metrics.metricName("test.max", "grp1"))), EPS);
+        assertEquals("Min(0...9) = 0", 0.0, metricValueFunc.apply(metrics.metrics().get(metrics.metricName("test.min", "grp1"))), EPS);
         assertEquals("Rate(0...9) = 1.40625",
-                     sum / elapsedSecs, metrics.metrics().get(metrics.metricName("test.rate", "grp1")).value(), EPS);
+                     sum / elapsedSecs, metricValueFunc.apply(metrics.metrics().get(metrics.metricName("test.rate", "grp1"))), EPS);
         assertEquals(String.format("Occurrences(0...%d) = %f", count, count / elapsedSecs),
                      count / elapsedSecs,
-                     metrics.metrics().get(metrics.metricName("test.occurences", "grp1")).value(), EPS);
+                     metricValueFunc.apply(metrics.metrics().get(metrics.metricName("test.occurences", "grp1"))), EPS);
         assertEquals("Count(0...9) = 10",
-                     (double) count, metrics.metrics().get(metrics.metricName("test.count", "grp1")).value(), EPS);
+                     (double) count, metricValueFunc.apply(metrics.metrics().get(metrics.metricName("test.count", "grp1"))), EPS);
     }
 
     @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();
@@ -152,11 +179,11 @@ public void testHierarchicalSensors() {
         child2.record();
         grandchild.record();
 
-        double p1 = parent1.metrics().get(0).value();
-        double p2 = parent2.metrics().get(0).value();
-        double c1 = child1.metrics().get(0).value();
-        double c2 = child2.metrics().get(0).value();
-        double gc = grandchild.metrics().get(0).value();
+        double p1 = (double) parent1.metrics().get(0).metricValue();
+        double p2 = (double) parent2.metrics().get(0).metricValue();
+        double c1 = (double) child1.metrics().get(0).metricValue();
+        double c2 = (double) child2.metrics().get(0).metricValue();
+        double gc = (double) grandchild.metrics().get(0).metricValue();
 
         /* each metric should have a count equal to one + its children's count */
         assertEquals(1.0, gc, EPS);
@@ -177,19 +204,33 @@ public void testBadSensorHierarchy() {
         metrics.sensor("gc", c1, c2); // should fail
     }
 
+    @Test
+    public void testRemoveChildSensor() {
+        final Metrics metrics = new Metrics();
+
+        final Sensor parent = metrics.sensor("parent");
+        final Sensor child = metrics.sensor("child", parent);
+
+        assertEquals(singletonList(child), metrics.childrenSensors().get(parent));
+
+        metrics.removeSensor("child");
+
+        assertEquals(emptyList(), metrics.childrenSensors().get(parent));
+    }
+
     @Test
     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);
@@ -227,10 +268,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();
@@ -268,7 +309,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")));
@@ -277,8 +318,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")));
@@ -292,7 +333,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());
@@ -303,7 +344,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);
@@ -322,37 +363,51 @@ public void testOldDataHasNoEffect() {
         MetricConfig config = new MetricConfig().timeWindow(windowMs, TimeUnit.MILLISECONDS).samples(samples);
         max.record(config, 50, time.milliseconds());
         time.sleep(samples * windowMs);
-        assertEquals(Double.NEGATIVE_INFINITY, max.measure(config, time.milliseconds()), EPS);
+        assertEquals(Double.NaN, max.measure(config, time.milliseconds()), EPS);
     }
 
+    /**
+     * Some implementations of SampledStat make sense to return NaN
+     * when there are no values set rather than the initial value
+     */
     @Test
-    public void testSampledStatInitialValue() {
-        // initialValue from each SampledStat is set as the initialValue on its Sample.
-        // The only way to test the initialValue is to infer it by having a SampledStat
-        // with expired Stats, because their values are reset to the initial values.
-        // Most implementations of combine on SampledStat end up returning the default
-        // value, so we can use this. This doesn't work for Percentiles though.
-        // This test looks a lot like testOldDataHasNoEffect because it's the same
-        // flow that leads to this state.
+    public void testSampledStatReturnsNaNWhenNoValuesExist() {
+        // This is tested by having a SampledStat with expired Stats,
+        // because their values get reset to the initial values.
         Max max = new Max();
         Min min = new Min();
         Avg avg = new Avg();
-        Count count = new Count();
-        Rate.SampledTotal sampledTotal = new Rate.SampledTotal();
-
         long windowMs = 100;
         int samples = 2;
         MetricConfig config = new MetricConfig().timeWindow(windowMs, TimeUnit.MILLISECONDS).samples(samples);
         max.record(config, 50, time.milliseconds());
         min.record(config, 50, time.milliseconds());
         avg.record(config, 50, time.milliseconds());
+
+        time.sleep(samples * windowMs);
+
+        assertEquals(Double.NaN, max.measure(config, time.milliseconds()), EPS);
+        assertEquals(Double.NaN, min.measure(config, time.milliseconds()), EPS);
+        assertEquals(Double.NaN, avg.measure(config, time.milliseconds()), EPS);
+    }
+
+    /**
+     * Some implementations of SampledStat make sense to return the initial value
+     * when there are no values set
+     */
+    @Test
+    public void testSampledStatReturnsInitialValueWhenNoValuesExist() {
+        WindowedCount count = new WindowedCount();
+        WindowedSum sampledTotal = new WindowedSum();
+        long windowMs = 100;
+        int samples = 2;
+        MetricConfig config = new MetricConfig().timeWindow(windowMs, TimeUnit.MILLISECONDS).samples(samples);
+
         count.record(config, 50, time.milliseconds());
         sampledTotal.record(config, 50, time.milliseconds());
+
         time.sleep(samples * windowMs);
 
-        assertEquals(Double.NEGATIVE_INFINITY, max.measure(config, time.milliseconds()), EPS);
-        assertEquals(Double.MAX_VALUE, min.measure(config, time.milliseconds()), EPS);
-        assertEquals(0.0, avg.measure(config, time.milliseconds()), EPS);
         assertEquals(0, count.measure(config, time.milliseconds()), EPS);
         assertEquals(0.0, sampledTotal.measure(config, time.milliseconds()), EPS);
     }
@@ -360,14 +415,14 @@ public void testSampledStatInitialValue() {
     @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 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);
@@ -375,7 +430,7 @@ public void testQuotas() {
         } catch (QuotaViolationException e) {
             // this is good
         }
-        assertEquals(6.0, metrics.metrics().get(metrics.metricName("test1.total", "grp1")).value(), EPS);
+        assertEquals(6.0, (Double) metrics.metrics().get(metrics.metricName("test1.total", "grp1")).metricValue(), EPS);
         sensor.record(-6.0);
         try {
             sensor.record(-1.0);
@@ -418,24 +473,106 @@ public void testPercentiles() {
         for (int i = 0; i < buckets; i++)
             sensor.record(i);
 
-        assertEquals(25, p25.value(), 1.0);
-        assertEquals(50, p50.value(), 1.0);
-        assertEquals(75, p75.value(), 1.0);
+        assertEquals(25, (Double) p25.metricValue(), 1.0);
+        assertEquals(50, (Double) p50.metricValue(), 1.0);
+        assertEquals(75, (Double) p75.metricValue(), 1.0);
 
         for (int i = 0; i < buckets; i++)
             sensor.record(0.0);
 
-        assertEquals(0.0, p25.value(), 1.0);
-        assertEquals(0.0, p50.value(), 1.0);
-        assertEquals(0.0, p75.value(), 1.0);
+        assertEquals(0.0, (Double) p25.metricValue(), 1.0);
+        assertEquals(0.0, (Double) p50.metricValue(), 1.0);
+        assertEquals(0.0, (Double) p75.metricValue(), 1.0);
 
         // record two more windows worth of sequential values
         for (int i = 0; i < buckets; i++)
             sensor.record(i);
 
-        assertEquals(25, p25.value(), 1.0);
-        assertEquals(50, p50.value(), 1.0);
-        assertEquals(75, p75.value(), 1.0);
+        assertEquals(25, (Double) p25.metricValue(), 1.0);
+        assertEquals(50, (Double) p50.metricValue(), 1.0);
+        assertEquals(75, (Double) p75.metricValue(), 1.0);
+    }
+
+    @Test
+    public void shouldPinSmallerValuesToMin() {
+        final double min = 0.0d;
+        final double max = 100d;
+        Percentiles percs = new Percentiles(1000,
+                                            min,
+                                            max,
+                                            BucketSizing.LINEAR,
+                                            new Percentile(metrics.metricName("test.p50", "grp1"), 50));
+        MetricConfig config = new MetricConfig().eventWindow(50).samples(2);
+        Sensor sensor = metrics.sensor("test", config);
+        sensor.add(percs);
+        Metric p50 = this.metrics.metrics().get(metrics.metricName("test.p50", "grp1"));
+
+        sensor.record(min - 100);
+        sensor.record(min - 100);
+        assertEquals(min, (double) p50.metricValue(), 0d);
+    }
+
+    @Test
+    public void shouldPinLargerValuesToMax() {
+        final double min = 0.0d;
+        final double max = 100d;
+        Percentiles percs = new Percentiles(1000,
+                                            min,
+                                            max,
+                                            BucketSizing.LINEAR,
+                                            new Percentile(metrics.metricName("test.p50", "grp1"), 50));
+        MetricConfig config = new MetricConfig().eventWindow(50).samples(2);
+        Sensor sensor = metrics.sensor("test", config);
+        sensor.add(percs);
+        Metric p50 = this.metrics.metrics().get(metrics.metricName("test.p50", "grp1"));
+
+        sensor.record(max + 100);
+        sensor.record(max + 100);
+        assertEquals(max, (double) p50.metricValue(), 0d);
+    }
+
+    @Test
+    public void testPercentilesWithRandomNumbersAndLinearBucketing() {
+        long seed = new Random().nextLong();
+        int sizeInBytes = 100 * 1000;   // 100kB
+        long maximumValue = 1000 * 24 * 60 * 60 * 1000L; // if values are ms, max is 1000 days
+
+        try {
+            Random prng = new Random(seed);
+            int numberOfValues = 5000 + prng.nextInt(10_000);  // range is [5000, 15000]
+
+            Percentiles percs = new Percentiles(sizeInBytes,
+                                                maximumValue,
+                                                BucketSizing.LINEAR,
+                                                new Percentile(metrics.metricName("test.p90", "grp1"), 90),
+                                                new Percentile(metrics.metricName("test.p99", "grp1"), 99));
+            MetricConfig config = new MetricConfig().eventWindow(50).samples(2);
+            Sensor sensor = metrics.sensor("test", config);
+            sensor.add(percs);
+            Metric p90 = this.metrics.metrics().get(metrics.metricName("test.p90", "grp1"));
+            Metric p99 = this.metrics.metrics().get(metrics.metricName("test.p99", "grp1"));
+
+            final List values = new ArrayList<>(numberOfValues);
+            // record two windows worth of sequential values
+            for (int i = 0; i < numberOfValues; ++i) {
+                long value = (Math.abs(prng.nextLong()) - 1) % maximumValue;
+                values.add(value);
+                sensor.record(value);
+            }
+
+            Collections.sort(values);
+
+            int p90Index = (int) Math.ceil(((double) (90 * numberOfValues)) / 100);
+            int p99Index = (int) Math.ceil(((double) (99 * numberOfValues)) / 100);
+
+            double expectedP90 = values.get(p90Index - 1);
+            double expectedP99 = values.get(p99Index - 1);
+
+            assertEquals(expectedP90, (Double) p90.metricValue(), expectedP90 / 5);
+            assertEquals(expectedP99, (Double) p99.metricValue(), expectedP99 / 5);
+        } catch (AssertionError e) {
+            throw new AssertionError("Assertion failed in randomized test. Reproduce with seed = " + seed + " .", e);
+        }
     }
 
     @Test
@@ -445,8 +582,12 @@ public void testRateWindowing() throws Exception {
         Sensor s = metrics.sensor("test.sensor", cfg);
         MetricName rateMetricName = metrics.metricName("test.rate", "grp1");
         MetricName totalMetricName = metrics.metricName("test.total", "grp1");
+        MetricName countRateMetricName = metrics.metricName("test.count.rate", "grp1");
+        MetricName countTotalMetricName = metrics.metricName("test.count.total", "grp1");
         s.add(new Meter(TimeUnit.SECONDS, rateMetricName, totalMetricName));
-        KafkaMetric totalMetric = metrics.metrics().get(metrics.metricName("test.total", "grp1"));
+        s.add(new Meter(TimeUnit.SECONDS, new WindowedCount(), countRateMetricName, countTotalMetricName));
+        KafkaMetric totalMetric = metrics.metrics().get(totalMetricName);
+        KafkaMetric countTotalMetric = metrics.metrics().get(countTotalMetricName);
 
         int sum = 0;
         int count = cfg.samples() - 1;
@@ -455,7 +596,7 @@ public void testRateWindowing() throws Exception {
             s.record(100);
             sum += 100;
             time.sleep(cfg.timeWindowMs());
-            assertEquals(sum, totalMetric.value(), EPS);
+            assertEquals(sum, (Double) totalMetric.metricValue(), EPS);
         }
 
         // Sleep for half the window.
@@ -464,11 +605,21 @@ public void testRateWindowing() throws Exception {
         // prior to any time passing
         double elapsedSecs = (cfg.timeWindowMs() * (cfg.samples() - 1) + cfg.timeWindowMs() / 2) / 1000.0;
 
-        KafkaMetric rateMetric = metrics.metrics().get(metrics.metricName("test.rate", "grp1"));
-        assertEquals("Rate(0...2) = 2.666", sum / elapsedSecs, rateMetric.value(), EPS);
+        KafkaMetric rateMetric = metrics.metrics().get(rateMetricName);
+        KafkaMetric countRateMetric = metrics.metrics().get(countRateMetricName);
+        assertEquals("Rate(0...2) = 2.666", sum / elapsedSecs, (Double) rateMetric.metricValue(), EPS);
+        assertEquals("Count rate(0...2) = 0.02666", count / elapsedSecs, (Double) countRateMetric.metricValue(), EPS);
         assertEquals("Elapsed Time = 75 seconds", elapsedSecs,
                 ((Rate) rateMetric.measurable()).windowSize(cfg, time.milliseconds()) / 1000, EPS);
-        assertEquals(sum, totalMetric.value(), EPS);
+        assertEquals(sum, (Double) totalMetric.metricValue(), EPS);
+        assertEquals(count, (Double) countTotalMetric.metricValue(), EPS);
+
+        // Verify that rates are expired, but total is cumulative
+        time.sleep(cfg.timeWindowMs() * cfg.samples());
+        assertEquals(0, (Double) rateMetric.metricValue(), EPS);
+        assertEquals(0, (Double) countRateMetric.metricValue(), EPS);
+        assertEquals(sum, (Double) totalMetric.metricValue(), EPS);
+        assertEquals(count, (Double) countTotalMetric.metricValue(), EPS);
     }
 
     public static class ConstantMeasurable implements Measurable {
@@ -562,7 +713,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();
@@ -588,9 +739,222 @@ public void testMetricInstances() {
                 // this is expected
             }
         }
+    }
 
+    /**
+     * Verifies that concurrent sensor add, remove, updates and read don't result
+     * in errors or deadlock.
+     */
+    @Test
+    public void testConcurrentReadUpdate() throws Exception {
+        final Random random = new Random();
+        final Deque sensors = new ConcurrentLinkedDeque<>();
+        metrics = new Metrics(new MockTime(10));
+        SensorCreator sensorCreator = new SensorCreator(metrics);
+
+        final AtomicBoolean alive = new AtomicBoolean(true);
+        executorService = Executors.newSingleThreadExecutor();
+        executorService.submit(new ConcurrentMetricOperation(alive, "record",
+            () -> sensors.forEach(sensor -> sensor.record(random.nextInt(10000)))));
+
+        for (int i = 0; i < 10000; i++) {
+            if (sensors.size() > 5) {
+                Sensor sensor = random.nextBoolean() ? sensors.removeFirst() : sensors.removeLast();
+                metrics.removeSensor(sensor.name());
+            }
+            StatType statType = StatType.forId(random.nextInt(StatType.values().length));
+            sensors.add(sensorCreator.createSensor(statType, i));
+            for (Sensor sensor : sensors) {
+                for (KafkaMetric metric : sensor.metrics()) {
+                    assertNotNull("Invalid metric value", metric.metricValue());
+                }
+            }
+        }
+        alive.set(false);
     }
 
-    
+    /**
+     * Verifies that concurrent sensor add, remove, updates and read with a metrics reporter
+     * that synchronizes on every reporter method doesn't result in errors or deadlock.
+     */
+    @Test
+    public void testConcurrentReadUpdateReport() throws Exception {
+
+        class LockingReporter implements MetricsReporter {
+            Map activeMetrics = new HashMap<>();
+            @Override
+            public synchronized void init(List metrics) {
+            }
+
+            @Override
+            public synchronized void metricChange(KafkaMetric metric) {
+                activeMetrics.put(metric.metricName(), metric);
+            }
+
+            @Override
+            public synchronized void metricRemoval(KafkaMetric metric) {
+                activeMetrics.remove(metric.metricName(), metric);
+            }
 
+            @Override
+            public synchronized void close() {
+            }
+
+            @Override
+            public void configure(Map configs) {
+            }
+
+            synchronized void processMetrics() {
+                for (KafkaMetric metric : activeMetrics.values()) {
+                    assertNotNull("Invalid metric value", metric.metricValue());
+                }
+            }
+        }
+
+        final LockingReporter reporter = new LockingReporter();
+        this.metrics.close();
+        this.metrics = new Metrics(config, Arrays.asList(reporter), new MockTime(10), true);
+        final Deque sensors = new ConcurrentLinkedDeque<>();
+        SensorCreator sensorCreator = new SensorCreator(metrics);
+
+        final Random random = new Random();
+        final AtomicBoolean alive = new AtomicBoolean(true);
+        executorService = Executors.newFixedThreadPool(3);
+
+        Future writeFuture = executorService.submit(new ConcurrentMetricOperation(alive, "record",
+            () -> sensors.forEach(sensor -> sensor.record(random.nextInt(10000)))));
+        Future readFuture = executorService.submit(new ConcurrentMetricOperation(alive, "read",
+            () -> sensors.forEach(sensor -> sensor.metrics().forEach(metric ->
+                assertNotNull("Invalid metric value", metric.metricValue())))));
+        Future reportFuture = executorService.submit(new ConcurrentMetricOperation(alive, "report",
+            () -> reporter.processMetrics()));
+
+        for (int i = 0; i < 10000; i++) {
+            if (sensors.size() > 10) {
+                Sensor sensor = random.nextBoolean() ? sensors.removeFirst() : sensors.removeLast();
+                metrics.removeSensor(sensor.name());
+            }
+            StatType statType = StatType.forId(random.nextInt(StatType.values().length));
+            sensors.add(sensorCreator.createSensor(statType, i));
+        }
+        assertFalse("Read failed", readFuture.isDone());
+        assertFalse("Write failed", writeFuture.isDone());
+        assertFalse("Report failed", reportFuture.isDone());
+
+        alive.set(false);
+    }
+
+    private class ConcurrentMetricOperation implements Runnable {
+        private final AtomicBoolean alive;
+        private final String opName;
+        private final Runnable op;
+        ConcurrentMetricOperation(AtomicBoolean alive, String opName, Runnable op) {
+            this.alive = alive;
+            this.opName = opName;
+            this.op = op;
+        }
+        @Override
+        public void run() {
+            try {
+                while (alive.get()) {
+                    op.run();
+                }
+            } catch (Throwable t) {
+                log.error("Metric {} failed with exception", opName, t);
+            }
+        }
+    }
+
+    enum StatType {
+        AVG(0),
+        TOTAL(1),
+        COUNT(2),
+        MAX(3),
+        MIN(4),
+        RATE(5),
+        SIMPLE_RATE(6),
+        SUM(7),
+        VALUE(8),
+        PERCENTILES(9),
+        METER(10);
+
+        int id;
+        StatType(int id) {
+            this.id = id;
+        }
+
+        static StatType forId(int id) {
+            for (StatType statType : StatType.values()) {
+                if (statType.id == id)
+                    return statType;
+            }
+            return null;
+        }
+    }
+
+    private static class SensorCreator {
+
+        private final Metrics metrics;
+
+        SensorCreator(Metrics metrics) {
+            this.metrics = metrics;
+        }
+
+        private Sensor createSensor(StatType statType, int index) {
+            Sensor sensor = metrics.sensor("kafka.requests." + index);
+            Map tags = Collections.singletonMap("tag", "tag" + index);
+            switch (statType) {
+                case AVG:
+                    sensor.add(metrics.metricName("test.metric.avg", "avg", tags), new Avg());
+                    break;
+                case 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 WindowedCount());
+                    break;
+                case MAX:
+                    sensor.add(metrics.metricName("test.metric.max", "max", tags), new Max());
+                    break;
+                case MIN:
+                    sensor.add(metrics.metricName("test.metric.min", "min", tags), new Min());
+                    break;
+                case RATE:
+                    sensor.add(metrics.metricName("test.metric.rate", "rate", tags), new Rate());
+                    break;
+                case SIMPLE_RATE:
+                    sensor.add(metrics.metricName("test.metric.simpleRate", "simpleRate", tags), new SimpleRate());
+                    break;
+                case 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());
+                    break;
+                case PERCENTILES:
+                    sensor.add(metrics.metricName("test.metric.percentiles", "percentiles", tags),
+                               new Percentiles(100, -100, 100, Percentiles.BucketSizing.CONSTANT,
+                                               new Percentile(metrics.metricName("test.median", "percentiles"), 50.0),
+                                               new Percentile(metrics.metricName("test.perc99_9", "percentiles"), 99.9)));
+                    break;
+                case METER:
+                    sensor.add(new Meter(metrics.metricName("test.metric.meter.rate", "meter", tags),
+                               metrics.metricName("test.metric.meter.total", "meter", tags)));
+                    break;
+                default:
+                    throw new IllegalStateException("Invalid stat type " + statType);
+            }
+            return sensor;
+        }
+    }
+
+    /**
+     * This test is to verify the deprecated {@link Metric#value()} method.
+     * @deprecated This will be removed in a future major release.
+     */
+    @Deprecated
+    @Test
+    public void testDeprecatedMetricValueMethod() {
+        verifyStats(KafkaMetric::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 d22111e128e25..de83ce120ba54 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
@@ -16,47 +16,360 @@
  */
 package org.apache.kafka.common.metrics;
 
+import org.apache.kafka.common.MetricName;
+import org.apache.kafka.common.metrics.stats.Avg;
+import org.apache.kafka.common.metrics.stats.CumulativeCount;
+import org.apache.kafka.common.metrics.stats.Meter;
+import org.apache.kafka.common.metrics.stats.Rate;
+import org.apache.kafka.common.metrics.stats.TokenBucket;
+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;
+import org.junit.Test;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.Callable;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import org.mockito.Mockito;
+
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.MatcherAssert.assertThat;
 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 org.apache.kafka.common.utils.SystemTime;
-import org.junit.Test;
+import static org.junit.Assert.fail;
 
 public class SensorTest {
+
+    private static final MetricConfig INFO_CONFIG = new MetricConfig().recordLevel(Sensor.RecordingLevel.INFO);
+    private static final MetricConfig DEBUG_CONFIG = new MetricConfig().recordLevel(Sensor.RecordingLevel.DEBUG);
+    private static final MetricConfig TRACE_CONFIG = new MetricConfig().recordLevel(Sensor.RecordingLevel.TRACE);
+
     @Test
     public void testRecordLevelEnum() {
         Sensor.RecordingLevel configLevel = Sensor.RecordingLevel.INFO;
         assertTrue(Sensor.RecordingLevel.INFO.shouldRecord(configLevel.id));
         assertFalse(Sensor.RecordingLevel.DEBUG.shouldRecord(configLevel.id));
+        assertFalse(Sensor.RecordingLevel.TRACE.shouldRecord(configLevel.id));
 
         configLevel = Sensor.RecordingLevel.DEBUG;
         assertTrue(Sensor.RecordingLevel.INFO.shouldRecord(configLevel.id));
         assertTrue(Sensor.RecordingLevel.DEBUG.shouldRecord(configLevel.id));
+        assertFalse(Sensor.RecordingLevel.TRACE.shouldRecord(configLevel.id));
+
+        configLevel = Sensor.RecordingLevel.TRACE;
+        assertTrue(Sensor.RecordingLevel.INFO.shouldRecord(configLevel.id));
+        assertTrue(Sensor.RecordingLevel.DEBUG.shouldRecord(configLevel.id));
+        assertTrue(Sensor.RecordingLevel.TRACE.shouldRecord(configLevel.id));
 
         assertEquals(Sensor.RecordingLevel.valueOf(Sensor.RecordingLevel.DEBUG.toString()),
             Sensor.RecordingLevel.DEBUG);
         assertEquals(Sensor.RecordingLevel.valueOf(Sensor.RecordingLevel.INFO.toString()),
             Sensor.RecordingLevel.INFO);
+        assertEquals(Sensor.RecordingLevel.valueOf(Sensor.RecordingLevel.TRACE.toString()),
+            Sensor.RecordingLevel.TRACE);
     }
 
     @Test
-    public void testShouldRecord() {
-        MetricConfig debugConfig = new MetricConfig().recordLevel(Sensor.RecordingLevel.DEBUG);
-        MetricConfig infoConfig = new MetricConfig().recordLevel(Sensor.RecordingLevel.INFO);
-
-        Sensor infoSensor = new Sensor(null, "infoSensor", null, debugConfig, new SystemTime(),
+    public void testShouldRecordForInfoLevelSensor() {
+        Sensor infoSensor = new Sensor(null, "infoSensor", null, INFO_CONFIG, new SystemTime(),
             0, Sensor.RecordingLevel.INFO);
         assertTrue(infoSensor.shouldRecord());
-        infoSensor = new Sensor(null, "infoSensor", null, debugConfig, new SystemTime(),
-            0, Sensor.RecordingLevel.DEBUG);
+
+        infoSensor = new Sensor(null, "infoSensor", null, DEBUG_CONFIG, new SystemTime(),
+            0, Sensor.RecordingLevel.INFO);
         assertTrue(infoSensor.shouldRecord());
 
-        Sensor debugSensor = new Sensor(null, "debugSensor", null, infoConfig, new SystemTime(),
+        infoSensor = new Sensor(null, "infoSensor", null, TRACE_CONFIG, new SystemTime(),
             0, Sensor.RecordingLevel.INFO);
-        assertTrue(debugSensor.shouldRecord());
-        debugSensor = new Sensor(null, "debugSensor", null, infoConfig, new SystemTime(),
+        assertTrue(infoSensor.shouldRecord());
+    }
+
+    @Test
+    public void testShouldRecordForDebugLevelSensor() {
+        Sensor debugSensor = new Sensor(null, "debugSensor", null, INFO_CONFIG, new SystemTime(),
             0, Sensor.RecordingLevel.DEBUG);
         assertFalse(debugSensor.shouldRecord());
+
+        debugSensor = new Sensor(null, "debugSensor", null, DEBUG_CONFIG, new SystemTime(),
+             0, Sensor.RecordingLevel.DEBUG);
+        assertTrue(debugSensor.shouldRecord());
+
+        debugSensor = new Sensor(null, "debugSensor", null, TRACE_CONFIG, new SystemTime(),
+             0, Sensor.RecordingLevel.DEBUG);
+        assertTrue(debugSensor.shouldRecord());
+    }
+
+    @Test
+    public void testShouldRecordForTraceLevelSensor() {
+        Sensor traceSensor = new Sensor(null, "traceSensor", null, INFO_CONFIG, new SystemTime(),
+             0, Sensor.RecordingLevel.TRACE);
+        assertFalse(traceSensor.shouldRecord());
+
+        traceSensor = new Sensor(null, "traceSensor", null, DEBUG_CONFIG, new SystemTime(),
+             0, Sensor.RecordingLevel.TRACE);
+        assertFalse(traceSensor.shouldRecord());
+
+        traceSensor = new Sensor(null, "traceSensor", null, TRACE_CONFIG, new SystemTime(),
+             0, Sensor.RecordingLevel.TRACE);
+        assertTrue(traceSensor.shouldRecord());
+    }
+
+    @Test
+    public void testExpiredSensor() {
+        MetricConfig config = new MetricConfig();
+        Time mockTime = new MockTime();
+        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
+    public void testIdempotentAdd() {
+        final Metrics metrics = new Metrics();
+        final Sensor sensor = metrics.sensor("sensor");
+
+        assertTrue(sensor.add(metrics.metricName("test-metric", "test-group"), new Avg()));
+
+        // adding the same metric to the same sensor is a no-op
+        assertTrue(sensor.add(metrics.metricName("test-metric", "test-group"), new Avg()));
+
+
+        // but adding the same metric to a DIFFERENT sensor is an error
+        final Sensor anotherSensor = metrics.sensor("another-sensor");
+        try {
+            anotherSensor.add(metrics.metricName("test-metric", "test-group"), new Avg());
+            fail("should have thrown");
+        } catch (final IllegalArgumentException ignored) {
+            // pass
+        }
+
+        // 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 WindowedSum()));
+
+        // so after all this, we still just have the original metric registered
+        assertEquals(1, sensor.metrics().size());
+        assertEquals(org.apache.kafka.common.metrics.stats.Avg.class, sensor.metrics().get(0).measurable().getClass());
+    }
+
+    /**
+     * The Sensor#checkQuotas should be thread-safe since the method may be used by many ReplicaFetcherThreads.
+     */
+    @Test
+    public void testCheckQuotasInMultiThreads() throws InterruptedException, ExecutionException {
+        final Metrics metrics = new Metrics(new MetricConfig().quota(Quota.upperBound(Double.MAX_VALUE))
+            // decreasing the value of time window make SampledStat always record the given value
+            .timeWindow(1, TimeUnit.MILLISECONDS)
+            // increasing the value of samples make SampledStat store more samples
+            .samples(100));
+        final Sensor sensor = metrics.sensor("sensor");
+
+        assertTrue(sensor.add(metrics.metricName("test-metric", "test-group"), new Rate()));
+        final int threadCount = 10;
+        final CountDownLatch latch = new CountDownLatch(1);
+        ExecutorService service = Executors.newFixedThreadPool(threadCount);
+        List> workers = new ArrayList<>(threadCount);
+        boolean needShutdown = true;
+        try {
+            for (int i = 0; i != threadCount; ++i) {
+                final int index = i;
+                workers.add(service.submit(new Callable() {
+                    @Override
+                    public Throwable call() {
+                        try {
+                            assertTrue(latch.await(5, TimeUnit.SECONDS));
+                            for (int j = 0; j != 20; ++j) {
+                                sensor.record(j * index, System.currentTimeMillis() + j, false);
+                                sensor.checkQuotas();
+                            }
+                            return null;
+                        } catch (Throwable e) {
+                            return e;
+                        }
+                    }
+                }));
+            }
+            latch.countDown();
+            service.shutdown();
+            assertTrue(service.awaitTermination(10, TimeUnit.SECONDS));
+            needShutdown = false;
+            for (Future callable : workers) {
+                assertTrue("If this failure happen frequently, we can try to increase the wait time", callable.isDone());
+                assertNull("Sensor#checkQuotas SHOULD be thread-safe!", callable.get());
+            }
+        } finally {
+            if (needShutdown) {
+                service.shutdownNow();
+            }
+        }
+    }
+
+    @Test
+    public void shouldReturnPresenceOfMetrics() {
+        final Metrics metrics = new Metrics();
+        final Sensor sensor = metrics.sensor("sensor");
+
+        assertThat(sensor.hasMetrics(), is(false));
+
+        sensor.add(
+            new MetricName("name1", "group1", "description1", Collections.emptyMap()),
+            new WindowedSum()
+        );
+
+        assertThat(sensor.hasMetrics(), is(true));
+
+        sensor.add(
+            new MetricName("name2", "group2", "description2", Collections.emptyMap()),
+            new CumulativeCount()
+        );
+
+        assertThat(sensor.hasMetrics(), is(true));
+    }
+
+    @Test
+    public void testStrictQuotaEnforcementWithRate() {
+        final Time time = new MockTime(0, System.currentTimeMillis(), 0);
+        final Metrics metrics = new Metrics(time);
+        final Sensor sensor = metrics.sensor("sensor", new MetricConfig()
+            .quota(Quota.upperBound(2))
+            .timeWindow(1, TimeUnit.SECONDS)
+            .samples(11));
+        final MetricName metricName = metrics.metricName("rate", "test-group");
+        assertTrue(sensor.add(metricName, new Rate()));
+        final KafkaMetric rateMetric = metrics.metric(metricName);
+
+        // Recording a first value at T+0 to bring the avg rate to 3 which is already
+        // above the quota.
+        strictRecord(sensor, 30, time.milliseconds());
+        assertEquals(3, rateMetric.measurableValue(time.milliseconds()), 0.1);
+
+        // Theoretically, we should wait 5s to bring back the avg rate to the define quota:
+        // ((30 / 10) - 2) / 2 * 10 = 5s
+        time.sleep(5000);
+
+        // But, recording a second value is rejected because the avg rate is still equal
+        // to 3 after 5s.
+        assertEquals(3, rateMetric.measurableValue(time.milliseconds()), 0.1);
+        assertThrows(QuotaViolationException.class, () -> strictRecord(sensor, 30, time.milliseconds()));
+
+        metrics.close();
+    }
+
+    @Test
+    public void testStrictQuotaEnforcementWithTokenBucket() {
+        final Time time = new MockTime(0, System.currentTimeMillis(), 0);
+        final Metrics metrics = new Metrics(time);
+        final Sensor sensor = metrics.sensor("sensor", new MetricConfig()
+            .quota(Quota.upperBound(2))
+            .timeWindow(1, TimeUnit.SECONDS)
+            .samples(10));
+        final MetricName metricName = metrics.metricName("credits", "test-group");
+        assertTrue(sensor.add(metricName, new TokenBucket()));
+        final KafkaMetric tkMetric = metrics.metric(metricName);
+
+        // Recording a first value at T+0 to bring the remaining credits below zero
+        strictRecord(sensor, 30, time.milliseconds());
+        assertEquals(-10, tkMetric.measurableValue(time.milliseconds()), 0.1);
+
+        // Theoretically, we should wait 5s to bring back the avg rate to the define quota:
+        // 10 / 2 = 5s
+        time.sleep(5000);
+
+        // Unlike the default rate based on a windowed sum, it works as expected.
+        assertEquals(0, tkMetric.measurableValue(time.milliseconds()), 0.1);
+        strictRecord(sensor, 30, time.milliseconds());
+        assertEquals(-30, tkMetric.measurableValue(time.milliseconds()), 0.1);
+
+        metrics.close();
+    }
+
+    private void strictRecord(Sensor sensor, double value, long timeMs) {
+        synchronized (sensor) {
+            sensor.checkQuotas(timeMs);
+            sensor.record(value, timeMs, false);
+        }
+    }
+
+    @Test
+    public void testRecordAndCheckQuotaUseMetricConfigOfEachStat() {
+        final Time time = new MockTime(0, System.currentTimeMillis(), 0);
+        final Metrics metrics = new Metrics(time);
+        final Sensor sensor = metrics.sensor("sensor");
+
+        final MeasurableStat stat1 = Mockito.mock(MeasurableStat.class);
+        final MetricName stat1Name = metrics.metricName("stat1", "test-group");
+        final MetricConfig stat1Config = new MetricConfig().quota(Quota.upperBound(5));
+        sensor.add(stat1Name, stat1, stat1Config);
+
+        final MeasurableStat stat2 = Mockito.mock(MeasurableStat.class);
+        final MetricName stat2Name = metrics.metricName("stat2", "test-group");
+        final MetricConfig stat2Config = new MetricConfig().quota(Quota.upperBound(10));
+        sensor.add(stat2Name, stat2, stat2Config);
+
+        sensor.record(10, 1);
+        Mockito.verify(stat1).record(stat1Config, 10, 1);
+        Mockito.verify(stat2).record(stat2Config, 10, 1);
+
+        sensor.checkQuotas(2);
+        Mockito.verify(stat1).measure(stat1Config, 2);
+        Mockito.verify(stat2).measure(stat2Config, 2);
+
+        metrics.close();
+    }
+
+    @Test
+    public void testUpdatingMetricConfigIsReflectedInTheSensor() {
+        final Time time = new MockTime(0, System.currentTimeMillis(), 0);
+        final Metrics metrics = new Metrics(time);
+        final Sensor sensor = metrics.sensor("sensor");
+
+        final MeasurableStat stat = Mockito.mock(MeasurableStat.class);
+        final MetricName statName = metrics.metricName("stat", "test-group");
+        final MetricConfig statConfig = new MetricConfig().quota(Quota.upperBound(5));
+        sensor.add(statName, stat, statConfig);
+
+        sensor.record(10, 1);
+        Mockito.verify(stat).record(statConfig, 10, 1);
+
+        sensor.checkQuotas(2);
+        Mockito.verify(stat).measure(statConfig, 2);
+
+        // Update the config of the KafkaMetric
+        final MetricConfig newConfig = new MetricConfig().quota(Quota.upperBound(10));
+        metrics.metric(statName).config(newConfig);
+
+        sensor.record(10, 3);
+        Mockito.verify(stat).record(newConfig, 10, 3);
+
+        sensor.checkQuotas(4);
+        Mockito.verify(stat).measure(newConfig, 4);
+
+        metrics.close();
     }
-}
+}
\ No newline at end of file
diff --git a/clients/src/test/java/org/apache/kafka/common/metrics/TokenBucketTest.java b/clients/src/test/java/org/apache/kafka/common/metrics/TokenBucketTest.java
new file mode 100644
index 0000000000000..db69af9656bd0
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/metrics/TokenBucketTest.java
@@ -0,0 +1,93 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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;
+
+import static org.junit.Assert.assertEquals;
+
+import java.util.concurrent.TimeUnit;
+import org.apache.kafka.common.metrics.stats.TokenBucket;
+import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.common.utils.Time;
+import org.junit.Before;
+import org.junit.Test;
+
+public class TokenBucketTest {
+    Time time;
+
+    @Before
+    public void setup() {
+        time = new MockTime(0, System.currentTimeMillis(), System.nanoTime());
+    }
+
+    @Test
+    public void testRecord() {
+        // Rate  = 5 unit / sec
+        // Burst = 2 * 10 = 20 units
+        MetricConfig config = new MetricConfig()
+            .quota(Quota.upperBound(5))
+            .timeWindow(2, TimeUnit.SECONDS)
+            .samples(10);
+
+        TokenBucket tk = new TokenBucket();
+
+        // Expect 100 credits at T
+        assertEquals(100, tk.measure(config, time.milliseconds()), 0.1);
+
+        // Record 60 at T, expect 13 credits
+        tk.record(config, 60, time.milliseconds());
+        assertEquals(40, tk.measure(config, time.milliseconds()), 0.1);
+
+        // Advance by 2s, record 5, expect 45 credits
+        time.sleep(2000);
+        tk.record(config, 5, time.milliseconds());
+        assertEquals(45, tk.measure(config, time.milliseconds()), 0.1);
+
+        // Advance by 2s, record 60, expect -5 credits
+        time.sleep(2000);
+        tk.record(config, 60, time.milliseconds());
+        assertEquals(-5, tk.measure(config, time.milliseconds()), 0.1);
+    }
+
+    @Test
+    public void testUnrecord() {
+        // Rate  = 5 unit / sec
+        // Burst = 2 * 10 = 20 units
+        MetricConfig config = new MetricConfig()
+            .quota(Quota.upperBound(5))
+            .timeWindow(2, TimeUnit.SECONDS)
+            .samples(10);
+
+        TokenBucket tk = new TokenBucket();
+
+        // Expect 100 credits at T
+        assertEquals(100, tk.measure(config, time.milliseconds()), 0.1);
+
+        // Record -60 at T, expect 100 credits
+        tk.record(config, -60, time.milliseconds());
+        assertEquals(100, tk.measure(config, time.milliseconds()), 0.1);
+
+        // Advance by 2s, record 60, expect 40 credits
+        time.sleep(2000);
+        tk.record(config, 60, time.milliseconds());
+        assertEquals(40, tk.measure(config, time.milliseconds()), 0.1);
+
+        // Advance by 2s, record -60, expect 100 credits
+        time.sleep(2000);
+        tk.record(config, -60, time.milliseconds());
+        assertEquals(100, tk.measure(config, time.milliseconds()), 0.1);
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/metrics/internals/IntGaugeSuiteTest.java b/clients/src/test/java/org/apache/kafka/common/metrics/internals/IntGaugeSuiteTest.java
new file mode 100644
index 0000000000000..461e3a97b9cc6
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/metrics/internals/IntGaugeSuiteTest.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.metrics.internals;
+
+import org.apache.kafka.common.MetricName;
+import org.apache.kafka.common.metrics.MetricConfig;
+import org.apache.kafka.common.metrics.Metrics;
+import org.junit.Test;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.Collections;
+import java.util.Map;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+
+public class IntGaugeSuiteTest {
+    private static final Logger log = LoggerFactory.getLogger(IntGaugeSuiteTest.class);
+
+    private static IntGaugeSuite createIntGaugeSuite() {
+        MetricConfig config = new MetricConfig();
+        Metrics metrics = new Metrics(config);
+        IntGaugeSuite suite = new IntGaugeSuite<>(log,
+            "mySuite",
+            metrics,
+            name -> new MetricName(name, "group", "myMetric", Collections.emptyMap()),
+            3);
+        return suite;
+    }
+
+    @Test
+    public void testCreateAndClose() {
+        IntGaugeSuite suite = createIntGaugeSuite();
+        assertEquals(3, suite.maxEntries());
+        suite.close();
+        suite.close();
+        suite.metrics().close();
+    }
+
+    @Test
+    public void testCreateMetrics() {
+        IntGaugeSuite suite = createIntGaugeSuite();
+        suite.increment("foo");
+        Map values = suite.values();
+        assertEquals(Integer.valueOf(1), values.get("foo"));
+        assertEquals(1, values.size());
+        suite.increment("foo");
+        suite.increment("bar");
+        suite.increment("baz");
+        suite.increment("quux");
+        values = suite.values();
+        assertEquals(Integer.valueOf(2), values.get("foo"));
+        assertEquals(Integer.valueOf(1), values.get("bar"));
+        assertEquals(Integer.valueOf(1), values.get("baz"));
+        assertEquals(3, values.size());
+        assertFalse(values.containsKey("quux"));
+        suite.close();
+        suite.metrics().close();
+    }
+
+    @Test
+    public void testCreateAndRemoveMetrics() {
+        IntGaugeSuite suite = createIntGaugeSuite();
+        suite.increment("foo");
+        suite.decrement("foo");
+        suite.increment("foo");
+        suite.increment("foo");
+        suite.increment("bar");
+        suite.decrement("bar");
+        suite.increment("baz");
+        suite.increment("quux");
+        Map values = suite.values();
+        assertEquals(Integer.valueOf(2), values.get("foo"));
+        assertFalse(values.containsKey("bar"));
+        assertEquals(Integer.valueOf(1), values.get("baz"));
+        assertEquals(Integer.valueOf(1), values.get("quux"));
+        assertEquals(3, values.size());
+        suite.close();
+        suite.metrics().close();
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/metrics/internals/MetricsUtilsTest.java b/clients/src/test/java/org/apache/kafka/common/metrics/internals/MetricsUtilsTest.java
new file mode 100644
index 0000000000000..0cc2cad6f30e6
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/metrics/internals/MetricsUtilsTest.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.common.metrics.internals;
+
+import org.junit.Test;
+
+import java.util.Map;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+
+public class MetricsUtilsTest {
+
+    @Test
+    public void testCreatingTags() {
+        Map tags = MetricsUtils.getTags("k1", "v1", "k2", "v2");
+        assertEquals("v1", tags.get("k1"));
+        assertEquals("v2", tags.get("k2"));
+        assertEquals(2, tags.size());
+    }
+
+    @Test
+    public void testCreatingTagsWithOddNumberOfTags() {
+        assertThrows(IllegalArgumentException.class, () -> MetricsUtils.getTags("k1", "v1", "k2", "v2", "extra"));
+    }
+}
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 9b6f686bd5e7b..061f60c18fda3 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
@@ -102,7 +101,6 @@ public void testBooleanFrequencies() {
     }
 
     @Test
-    @SuppressWarnings("deprecation")
     public void testUseWithMetrics() {
         MetricName name1 = name("1");
         MetricName name2 = name("2");
@@ -124,29 +122,29 @@ public void testUseWithMetrics() {
         for (int i = 0; i != 100; ++i) {
             frequencies.record(config, i % 4 + 1, time.milliseconds());
         }
-        assertEquals(0.25, metric1.value(), DELTA);
-        assertEquals(0.25, metric2.value(), DELTA);
-        assertEquals(0.25, metric3.value(), DELTA);
-        assertEquals(0.25, metric4.value(), DELTA);
+        assertEquals(0.25, (Double) metric1.metricValue(), DELTA);
+        assertEquals(0.25, (Double) metric2.metricValue(), DELTA);
+        assertEquals(0.25, (Double) metric3.metricValue(), DELTA);
+        assertEquals(0.25, (Double) metric4.metricValue(), DELTA);
 
         // Record 2 windows worth of values
         for (int i = 0; i != 100; ++i) {
             frequencies.record(config, i % 2 + 1, time.milliseconds());
         }
-        assertEquals(0.50, metric1.value(), DELTA);
-        assertEquals(0.50, metric2.value(), DELTA);
-        assertEquals(0.00, metric3.value(), DELTA);
-        assertEquals(0.00, metric4.value(), DELTA);
+        assertEquals(0.50, (Double) metric1.metricValue(), DELTA);
+        assertEquals(0.50, (Double) metric2.metricValue(), DELTA);
+        assertEquals(0.00, (Double) metric3.metricValue(), DELTA);
+        assertEquals(0.00, (Double) metric4.metricValue(), DELTA);
 
         // Record 1 window worth of values to overlap with the last window
         // that is half 1.0 and half 2.0
         for (int i = 0; i != 50; ++i) {
             frequencies.record(config, 4.0, time.milliseconds());
         }
-        assertEquals(0.25, metric1.value(), DELTA);
-        assertEquals(0.25, metric2.value(), DELTA);
-        assertEquals(0.00, metric3.value(), DELTA);
-        assertEquals(0.50, metric4.value(), DELTA);
+        assertEquals(0.25, (Double) metric1.metricValue(), DELTA);
+        assertEquals(0.25, (Double) metric2.metricValue(), DELTA);
+        assertEquals(0.00, (Double) metric3.metricValue(), DELTA);
+        assertEquals(0.50, (Double) metric4.metricValue(), DELTA);
     }
 
     protected MetricName name(String metricName) {
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 916e61989784e..3230da2c8340a 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
@@ -16,16 +16,35 @@
  */
 package org.apache.kafka.common.network;
 
+import java.util.ArrayList;
+import java.util.List;
 import org.apache.kafka.common.config.SslConfigs;
+import org.apache.kafka.common.utils.Utils;
 import org.apache.kafka.test.TestSslUtils;
 
 import java.io.File;
 import java.net.InetAddress;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Set;
+import org.apache.kafka.test.TestSslUtils.SslConfigsBuilder;
 
 public class CertStores {
 
+    public static final Set KEYSTORE_PROPS = Utils.mkSet(
+            SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG,
+            SslConfigs.SSL_KEYSTORE_TYPE_CONFIG,
+            SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG,
+            SslConfigs.SSL_KEY_PASSWORD_CONFIG,
+            SslConfigs.SSL_KEYSTORE_KEY_CONFIG,
+            SslConfigs.SSL_KEYSTORE_CERTIFICATE_CHAIN_CONFIG);
+
+    public static final Set TRUSTSTORE_PROPS = Utils.mkSet(
+            SslConfigs.SSL_TRUSTSTORE_LOCATION_CONFIG,
+            SslConfigs.SSL_TRUSTSTORE_TYPE_CONFIG,
+            SslConfigs.SSL_TRUSTSTORE_PASSWORD_CONFIG,
+            SslConfigs.SSL_TRUSTSTORE_CERTIFICATES_CONFIG);
+
     private final Map sslConfig;
 
     public CertStores(boolean server, String hostName) throws Exception {
@@ -33,7 +52,7 @@ public CertStores(boolean server, String hostName) throws Exception {
     }
 
     public CertStores(boolean server, String commonName, String sanHostName) throws Exception {
-        this(server, commonName, new TestSslUtils.CertificateBuilder().sanDnsName(sanHostName));
+        this(server, commonName, new TestSslUtils.CertificateBuilder().sanDnsNames(sanHostName));
     }
 
     public CertStores(boolean server, String commonName, InetAddress hostAddress) throws Exception {
@@ -41,21 +60,98 @@ public CertStores(boolean server, String commonName, InetAddress hostAddress) th
     }
 
     private CertStores(boolean server, String commonName, TestSslUtils.CertificateBuilder certBuilder) throws Exception {
+        this(server, commonName, "RSA", certBuilder, false);
+    }
+
+    private CertStores(boolean server, String commonName, String keyAlgorithm, TestSslUtils.CertificateBuilder certBuilder, boolean usePem) 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);
+        File truststoreFile = usePem ? null : File.createTempFile(name + "TS", ".jks");
+        sslConfig = new SslConfigsBuilder(mode)
+                .useClientCert(!server)
+                .certAlias(name)
+                .cn(commonName)
+                .createNewTrustStore(truststoreFile)
+                .certBuilder(certBuilder)
+                .algorithm(keyAlgorithm)
+                .usePem(usePem)
+                .build();
     }
 
+
     public Map getTrustingConfig(CertStores truststoreConfig) {
         Map config = new HashMap<>(sslConfig);
-        config.put(SslConfigs.SSL_TRUSTSTORE_LOCATION_CONFIG, truststoreConfig.sslConfig.get(SslConfigs.SSL_TRUSTSTORE_LOCATION_CONFIG));
-        config.put(SslConfigs.SSL_TRUSTSTORE_PASSWORD_CONFIG, truststoreConfig.sslConfig.get(SslConfigs.SSL_TRUSTSTORE_PASSWORD_CONFIG));
-        config.put(SslConfigs.SSL_TRUSTSTORE_TYPE_CONFIG, truststoreConfig.sslConfig.get(SslConfigs.SSL_TRUSTSTORE_TYPE_CONFIG));
+        for (String propName : TRUSTSTORE_PROPS) {
+            config.put(propName, truststoreConfig.sslConfig.get(propName));
+        }
         return config;
     }
 
     public Map getUntrustingConfig() {
         return sslConfig;
     }
+
+    public Map keyStoreProps() {
+        Map props = new HashMap<>();
+        for (String propName : KEYSTORE_PROPS) {
+            props.put(propName, sslConfig.get(propName));
+        }
+        return props;
+    }
+
+    public Map trustStoreProps() {
+        Map props = new HashMap<>();
+        for (String propName : TRUSTSTORE_PROPS) {
+            props.put(propName, sslConfig.get(propName));
+        }
+        return props;
+    }
+
+    public static class Builder {
+        private final boolean isServer;
+        private String cn;
+        private List sanDns;
+        private InetAddress sanIp;
+        private String keyAlgorithm;
+        private boolean usePem;
+
+        public Builder(boolean isServer) {
+            this.isServer = isServer;
+            this.sanDns = new ArrayList<>();
+            this.keyAlgorithm = "RSA";
+        }
+
+        public Builder cn(String cn) {
+            this.cn = cn;
+            return this;
+        }
+
+        public Builder addHostName(String hostname) {
+            this.sanDns.add(hostname);
+            return this;
+        }
+
+        public Builder hostAddress(InetAddress hostAddress) {
+            this.sanIp = hostAddress;
+            return this;
+        }
+
+        public Builder keyAlgorithm(String keyAlgorithm) {
+            this.keyAlgorithm = keyAlgorithm;
+            return this;
+        }
+
+        public Builder usePem(boolean usePem) {
+            this.usePem = usePem;
+            return this;
+        }
+
+        public CertStores build() throws Exception {
+            TestSslUtils.CertificateBuilder certBuilder = new TestSslUtils.CertificateBuilder()
+                .sanDnsNames(sanDns.toArray(new String[0]));
+            if (sanIp != null)
+                certBuilder = certBuilder.sanIpAddress(sanIp);
+            return new CertStores(isServer, cn, keyAlgorithm, certBuilder, usePem);
+        }
+    }
 }
\ No newline at end of file
diff --git a/clients/src/test/java/org/apache/kafka/common/network/ChannelBuildersTest.java b/clients/src/test/java/org/apache/kafka/common/network/ChannelBuildersTest.java
index de210e7496cb1..9c3fdc90c54c1 100644
--- a/clients/src/test/java/org/apache/kafka/common/network/ChannelBuildersTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/network/ChannelBuildersTest.java
@@ -19,39 +19,42 @@
 import org.apache.kafka.common.Configurable;
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.config.internals.BrokerSecurityConfigs;
+import org.apache.kafka.common.security.TestSecurityConfig;
 import org.apache.kafka.common.security.auth.AuthenticationContext;
 import org.apache.kafka.common.security.auth.KafkaPrincipal;
 import org.apache.kafka.common.security.auth.KafkaPrincipalBuilder;
 import org.apache.kafka.common.security.auth.PlaintextAuthenticationContext;
-import org.apache.kafka.common.security.auth.PrincipalBuilder;
-import org.easymock.EasyMock;
+import org.apache.kafka.common.security.auth.SecurityProtocol;
 import org.junit.Test;
 
 import java.net.InetAddress;
 import java.security.Principal;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Properties;
 
 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.mockito.Mockito.mock;
 
 public class ChannelBuildersTest {
 
     @Test
-    @SuppressWarnings("deprecation")
     public void testCreateOldPrincipalBuilder() throws Exception {
-        TransportLayer transportLayer = EasyMock.mock(TransportLayer.class);
-        Authenticator authenticator = EasyMock.mock(Authenticator.class);
+        TransportLayer transportLayer = mock(TransportLayer.class);
+        Authenticator authenticator = mock(Authenticator.class);
 
         Map configs = new HashMap<>();
         configs.put(BrokerSecurityConfigs.PRINCIPAL_BUILDER_CLASS_CONFIG, OldPrincipalBuilder.class);
-        KafkaPrincipalBuilder builder = ChannelBuilders.createPrincipalBuilder(configs, transportLayer, authenticator, null);
+        KafkaPrincipalBuilder builder = ChannelBuilders.createPrincipalBuilder(configs, transportLayer, authenticator, null, null);
 
         // test old principal builder is properly configured and delegated to
         assertTrue(OldPrincipalBuilder.configured);
 
         // test delegation
-        KafkaPrincipal principal = builder.build(new PlaintextAuthenticationContext(InetAddress.getLocalHost()));
+        KafkaPrincipal principal = builder.build(new PlaintextAuthenticationContext(InetAddress.getLocalHost(), SecurityProtocol.PLAINTEXT.name()));
         assertEquals(OldPrincipalBuilder.PRINCIPAL_NAME, principal.getName());
         assertEquals(KafkaPrincipal.USER_TYPE, principal.getPrincipalType());
     }
@@ -60,13 +63,73 @@ public void testCreateOldPrincipalBuilder() throws Exception {
     public void testCreateConfigurableKafkaPrincipalBuilder() {
         Map configs = new HashMap<>();
         configs.put(BrokerSecurityConfigs.PRINCIPAL_BUILDER_CLASS_CONFIG, ConfigurableKafkaPrincipalBuilder.class);
-        KafkaPrincipalBuilder builder = ChannelBuilders.createPrincipalBuilder(configs, null, null, null);
+        KafkaPrincipalBuilder builder = ChannelBuilders.createPrincipalBuilder(configs, null, null, null, null);
         assertTrue(builder instanceof ConfigurableKafkaPrincipalBuilder);
         assertTrue(((ConfigurableKafkaPrincipalBuilder) builder).configured);
     }
 
+    @Test
+    public void testChannelBuilderConfigs() {
+        Properties props = new Properties();
+        props.put("listener.name.listener1.gssapi.sasl.kerberos.service.name", "testkafka");
+        props.put("listener.name.listener1.sasl.kerberos.service.name", "testkafkaglobal");
+        props.put("plain.sasl.server.callback.handler.class", "callback");
+        props.put("listener.name.listener1.gssapi.config1.key", "custom.config1");
+        props.put("custom.config2.key", "custom.config2");
+        TestSecurityConfig securityConfig = new TestSecurityConfig(props);
+
+        // test configs with listener prefix
+        Map configs = ChannelBuilders.channelBuilderConfigs(securityConfig, new ListenerName("listener1"));
+
+        assertNull(configs.get("listener.name.listener1.gssapi.sasl.kerberos.service.name"));
+        assertFalse(securityConfig.unused().contains("listener.name.listener1.gssapi.sasl.kerberos.service.name"));
+
+        assertEquals(configs.get("gssapi.sasl.kerberos.service.name"), "testkafka");
+        assertFalse(securityConfig.unused().contains("gssapi.sasl.kerberos.service.name"));
+
+        assertEquals(configs.get("sasl.kerberos.service.name"), "testkafkaglobal");
+        assertFalse(securityConfig.unused().contains("sasl.kerberos.service.name"));
+
+        assertNull(configs.get("listener.name.listener1.sasl.kerberos.service.name"));
+        assertFalse(securityConfig.unused().contains("listener.name.listener1.sasl.kerberos.service.name"));
+
+        assertNull(configs.get("plain.sasl.server.callback.handler.class"));
+        assertFalse(securityConfig.unused().contains("plain.sasl.server.callback.handler.class"));
+
+        assertEquals(configs.get("listener.name.listener1.gssapi.config1.key"), "custom.config1");
+        assertFalse(securityConfig.unused().contains("listener.name.listener1.gssapi.config1.key"));
+
+        assertEquals(configs.get("custom.config2.key"), "custom.config2");
+        assertFalse(securityConfig.unused().contains("custom.config2.key"));
+
+        // test configs without listener prefix
+        securityConfig = new TestSecurityConfig(props);
+        configs = ChannelBuilders.channelBuilderConfigs(securityConfig, null);
+
+        assertEquals(configs.get("listener.name.listener1.gssapi.sasl.kerberos.service.name"), "testkafka");
+        assertFalse(securityConfig.unused().contains("listener.name.listener1.gssapi.sasl.kerberos.service.name"));
+
+        assertNull(configs.get("gssapi.sasl.kerberos.service.name"));
+        assertFalse(securityConfig.unused().contains("gssapi.sasl.kerberos.service.name"));
+
+        assertEquals(configs.get("listener.name.listener1.sasl.kerberos.service.name"), "testkafkaglobal");
+        assertFalse(securityConfig.unused().contains("listener.name.listener1.sasl.kerberos.service.name"));
+
+        assertNull(configs.get("sasl.kerberos.service.name"));
+        assertFalse(securityConfig.unused().contains("sasl.kerberos.service.name"));
+
+        assertEquals(configs.get("plain.sasl.server.callback.handler.class"), "callback");
+        assertFalse(securityConfig.unused().contains("plain.sasl.server.callback.handler.class"));
+
+        assertEquals(configs.get("listener.name.listener1.gssapi.config1.key"), "custom.config1");
+        assertFalse(securityConfig.unused().contains("listener.name.listener1.gssapi.config1.key"));
+
+        assertEquals(configs.get("custom.config2.key"), "custom.config2");
+        assertFalse(securityConfig.unused().contains("custom.config2.key"));
+    }
+
     @SuppressWarnings("deprecation")
-    public static class OldPrincipalBuilder implements PrincipalBuilder {
+    public static class OldPrincipalBuilder implements org.apache.kafka.common.security.auth.PrincipalBuilder {
         private static boolean configured = false;
         private static final String PRINCIPAL_NAME = "bob";
 
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 aa7a15ece7cf7..d0cc05942e326 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
@@ -17,6 +17,7 @@
 package org.apache.kafka.common.network;
 
 import org.apache.kafka.common.security.auth.SecurityProtocol;
+import org.apache.kafka.common.security.ssl.DefaultSslEngineFactory;
 import org.apache.kafka.common.security.ssl.SslFactory;
 
 import javax.net.ssl.SSLContext;
@@ -41,6 +42,7 @@ class EchoServer extends Thread {
     private final ServerSocket serverSocket;
     private final List threads;
     private final List sockets;
+    private volatile boolean closing = false;
     private final SslFactory sslFactory;
     private final AtomicBoolean renegotiate = new AtomicBoolean();
 
@@ -49,7 +51,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 = ((DefaultSslEngineFactory) this.sslFactory.sslEngineFactory()).sslContext();
                 this.serverSocket = sslContext.getServerSocketFactory().createServerSocket(0);
                 break;
             case PLAINTEXT:
@@ -71,40 +73,45 @@ public void renegotiate() {
     @Override
     public void run() {
         try {
-            while (true) {
+            while (!closing) {
                 final Socket socket = serverSocket.accept();
-                sockets.add(socket);
-                Thread thread = new Thread() {
-                    @Override
-                    public void run() {
-                        try {
-                            DataInputStream input = new DataInputStream(socket.getInputStream());
-                            DataOutputStream output = new DataOutputStream(socket.getOutputStream());
-                            while (socket.isConnected() && !socket.isClosed()) {
-                                int size = input.readInt();
-                                if (renegotiate.get()) {
-                                    renegotiate.set(false);
-                                    ((SSLSocket) socket).startHandshake();
-                                }
-                                byte[] bytes = new byte[size];
-                                input.readFully(bytes);
-                                output.writeInt(size);
-                                output.write(bytes);
-                                output.flush();
-                            }
-                        } catch (IOException e) {
-                            // ignore
-                        } finally {
+                synchronized (sockets) {
+                    if (closing) {
+                        break;
+                    }
+                    sockets.add(socket);
+                    Thread thread = new Thread() {
+                        @Override
+                        public void run() {
                             try {
-                                socket.close();
+                                DataInputStream input = new DataInputStream(socket.getInputStream());
+                                DataOutputStream output = new DataOutputStream(socket.getOutputStream());
+                                while (socket.isConnected() && !socket.isClosed()) {
+                                    int size = input.readInt();
+                                    if (renegotiate.get()) {
+                                        renegotiate.set(false);
+                                        ((SSLSocket) socket).startHandshake();
+                                    }
+                                    byte[] bytes = new byte[size];
+                                    input.readFully(bytes);
+                                    output.writeInt(size);
+                                    output.write(bytes);
+                                    output.flush();
+                                }
                             } catch (IOException e) {
                                 // ignore
+                            } finally {
+                                try {
+                                    socket.close();
+                                } catch (IOException e) {
+                                    // ignore
+                                }
                             }
                         }
-                    }
-                };
-                thread.start();
-                threads.add(thread);
+                    };
+                    thread.start();
+                    threads.add(thread);
+                }
             }
         } catch (IOException e) {
             // ignore
@@ -112,11 +119,14 @@ public void run() {
     }
 
     public void closeConnections() throws IOException {
-        for (Socket socket : sockets)
-            socket.close();
+        synchronized (sockets) {
+            for (Socket socket : sockets)
+                socket.close();
+        }
     }
 
     public void close() throws IOException, InterruptedException {
+        closing = true;
         this.serverSocket.close();
         closeConnections();
         for (Thread t : threads)
diff --git a/clients/src/test/java/org/apache/kafka/common/network/KafkaChannelTest.java b/clients/src/test/java/org/apache/kafka/common/network/KafkaChannelTest.java
new file mode 100644
index 0000000000000..51bd788ebd2a0
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/network/KafkaChannelTest.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.network;
+
+import org.apache.kafka.common.memory.MemoryPool;
+import org.apache.kafka.test.TestUtils;
+import org.junit.Test;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mockito;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+
+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;
+
+public class KafkaChannelTest {
+
+    @Test
+    public void testSending() throws IOException {
+        Authenticator authenticator = Mockito.mock(Authenticator.class);
+        TransportLayer transport = Mockito.mock(TransportLayer.class);
+        MemoryPool pool = Mockito.mock(MemoryPool.class);
+        ChannelMetadataRegistry metadataRegistry = Mockito.mock(ChannelMetadataRegistry.class);
+
+        KafkaChannel channel = new KafkaChannel("0", transport, () -> authenticator,
+            1024, pool, metadataRegistry);
+        ByteBufferSend send = ByteBufferSend.sizePrefixed(ByteBuffer.wrap(TestUtils.randomBytes(128)));
+        NetworkSend networkSend = new NetworkSend("0", send);
+
+        channel.setSend(networkSend);
+        assertTrue(channel.hasSend());
+        assertThrows(IllegalStateException.class, () -> channel.setSend(networkSend));
+
+        Mockito.when(transport.write(Mockito.any(ByteBuffer[].class))).thenReturn(4L);
+        assertEquals(4L, channel.write());
+        assertEquals(128, send.remaining());
+        assertNull(channel.maybeCompleteSend());
+
+        Mockito.when(transport.write(Mockito.any(ByteBuffer[].class))).thenReturn(64L);
+        assertEquals(64, channel.write());
+        assertEquals(64, send.remaining());
+        assertNull(channel.maybeCompleteSend());
+
+        Mockito.when(transport.write(Mockito.any(ByteBuffer[].class))).thenReturn(64L);
+        assertEquals(64, channel.write());
+        assertEquals(0, send.remaining());
+        assertEquals(networkSend, channel.maybeCompleteSend());
+    }
+
+    @Test
+    public void testReceiving() throws IOException {
+        Authenticator authenticator = Mockito.mock(Authenticator.class);
+        TransportLayer transport = Mockito.mock(TransportLayer.class);
+        MemoryPool pool = Mockito.mock(MemoryPool.class);
+        ChannelMetadataRegistry metadataRegistry = Mockito.mock(ChannelMetadataRegistry.class);
+
+        ArgumentCaptor sizeCaptor = ArgumentCaptor.forClass(Integer.class);
+        Mockito.when(pool.tryAllocate(sizeCaptor.capture())).thenAnswer(invocation -> {
+            return ByteBuffer.allocate(sizeCaptor.getValue());
+        });
+
+        KafkaChannel channel = new KafkaChannel("0", transport, () -> authenticator,
+            1024, pool, metadataRegistry);
+
+        ArgumentCaptor bufferCaptor = ArgumentCaptor.forClass(ByteBuffer.class);
+        Mockito.when(transport.read(bufferCaptor.capture())).thenAnswer(invocation -> {
+            bufferCaptor.getValue().putInt(128);
+            return 4;
+        }).thenReturn(0);
+        assertEquals(4, channel.read());
+        assertEquals(4, channel.currentReceive().bytesRead());
+        assertNull(channel.maybeCompleteReceive());
+
+        Mockito.reset(transport);
+        Mockito.when(transport.read(bufferCaptor.capture())).thenAnswer(invocation -> {
+            bufferCaptor.getValue().put(TestUtils.randomBytes(64));
+            return 64;
+        });
+        assertEquals(64, channel.read());
+        assertEquals(68, channel.currentReceive().bytesRead());
+        assertNull(channel.maybeCompleteReceive());
+
+        Mockito.reset(transport);
+        Mockito.when(transport.read(bufferCaptor.capture())).thenAnswer(invocation -> {
+            bufferCaptor.getValue().put(TestUtils.randomBytes(64));
+            return 64;
+        });
+        assertEquals(64, channel.read());
+        assertEquals(132, channel.currentReceive().bytesRead());
+        assertNotNull(channel.maybeCompleteReceive());
+        assertNull(channel.currentReceive());
+    }
+
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/network/NetworkReceiveTest.java b/clients/src/test/java/org/apache/kafka/common/network/NetworkReceiveTest.java
new file mode 100644
index 0000000000000..f03aea4faab68
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/network/NetworkReceiveTest.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.network;
+
+import org.apache.kafka.test.TestUtils;
+import org.junit.Test;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mockito;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.channels.ScatteringByteChannel;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+public class NetworkReceiveTest {
+
+    @Test
+    public void testBytesRead() throws IOException {
+        NetworkReceive receive = new NetworkReceive(128, "0");
+        assertEquals(0, receive.bytesRead());
+
+        ScatteringByteChannel channel = Mockito.mock(ScatteringByteChannel.class);
+
+        ArgumentCaptor bufferCaptor = ArgumentCaptor.forClass(ByteBuffer.class);
+        Mockito.when(channel.read(bufferCaptor.capture())).thenAnswer(invocation -> {
+            bufferCaptor.getValue().putInt(128);
+            return 4;
+        }).thenReturn(0);
+
+        assertEquals(4, receive.readFrom(channel));
+        assertEquals(4, receive.bytesRead());
+        assertFalse(receive.complete());
+
+        Mockito.reset(channel);
+        Mockito.when(channel.read(bufferCaptor.capture())).thenAnswer(invocation -> {
+            bufferCaptor.getValue().put(TestUtils.randomBytes(64));
+            return 64;
+        });
+
+        assertEquals(64, receive.readFrom(channel));
+        assertEquals(68, receive.bytesRead());
+        assertFalse(receive.complete());
+
+        Mockito.reset(channel);
+        Mockito.when(channel.read(bufferCaptor.capture())).thenAnswer(invocation -> {
+            bufferCaptor.getValue().put(TestUtils.randomBytes(64));
+            return 64;
+        });
+
+        assertEquals(64, receive.readFrom(channel));
+        assertEquals(132, receive.bytesRead());
+        assertTrue(receive.complete());
+    }
+
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/network/NetworkTestUtils.java b/clients/src/test/java/org/apache/kafka/common/network/NetworkTestUtils.java
index 59980490e6807..6ffa0afb01759 100644
--- a/clients/src/test/java/org/apache/kafka/common/network/NetworkTestUtils.java
+++ b/clients/src/test/java/org/apache/kafka/common/network/NetworkTestUtils.java
@@ -18,6 +18,8 @@
 
 import java.io.IOException;
 import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
+import java.util.Map;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
@@ -27,7 +29,8 @@
 import org.apache.kafka.common.security.auth.SecurityProtocol;
 import org.apache.kafka.common.utils.LogContext;
 import org.apache.kafka.common.security.authenticator.CredentialCache;
-import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.common.security.token.delegation.internals.DelegationTokenCache;
+import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.common.utils.Utils;
 import org.apache.kafka.test.TestUtils;
 
@@ -35,37 +38,51 @@
  * Common utility functions used by transport layer and authenticator tests.
  */
 public class NetworkTestUtils {
+    public static NioEchoServer createEchoServer(ListenerName listenerName, SecurityProtocol securityProtocol,
+                                                 AbstractConfig serverConfig, CredentialCache credentialCache, Time time) throws Exception {
+        return createEchoServer(listenerName, securityProtocol, serverConfig, credentialCache, 100, time);
+    }
 
     public static NioEchoServer createEchoServer(ListenerName listenerName, SecurityProtocol securityProtocol,
-            AbstractConfig serverConfig, CredentialCache credentialCache) throws Exception {
-        NioEchoServer server = new NioEchoServer(listenerName, securityProtocol, serverConfig, "localhost", null, credentialCache);
+                                                 AbstractConfig serverConfig, CredentialCache credentialCache,
+                                                 int failedAuthenticationDelayMs, Time time) throws Exception {
+        NioEchoServer server = new NioEchoServer(listenerName, securityProtocol, serverConfig, "localhost",
+                null, credentialCache, failedAuthenticationDelayMs, time);
         server.start();
         return server;
     }
 
-    public static Selector createSelector(ChannelBuilder channelBuilder) {
-        return new Selector(5000, new Metrics(), new MockTime(), "MetricGroup", channelBuilder, new LogContext());
+    public static NioEchoServer createEchoServer(ListenerName listenerName, SecurityProtocol securityProtocol,
+            AbstractConfig serverConfig, CredentialCache credentialCache,
+            int failedAuthenticationDelayMs, Time time, DelegationTokenCache tokenCache) throws Exception {
+        NioEchoServer server = new NioEchoServer(listenerName, securityProtocol, serverConfig, "localhost",
+                null, credentialCache, failedAuthenticationDelayMs, time, tokenCache);
+        server.start();
+        return server;
     }
 
-    public static void checkClientConnection(Selector selector, String node, int minMessageSize, int messageCount) throws Exception {
+    public static Selector createSelector(ChannelBuilder channelBuilder, Time time) {
+        return new Selector(5000, new Metrics(), time, "MetricGroup", channelBuilder, new LogContext());
+    }
 
+    public static void checkClientConnection(Selector selector, String node, int minMessageSize, int messageCount) throws Exception {
         waitForChannelReady(selector, node);
         String prefix = TestUtils.randomString(minMessageSize);
         int requests = 0;
         int responses = 0;
-        selector.send(new NetworkSend(node, ByteBuffer.wrap((prefix + "-0").getBytes())));
+        selector.send(new NetworkSend(node, ByteBufferSend.sizePrefixed(ByteBuffer.wrap((prefix + "-0").getBytes(StandardCharsets.UTF_8)))));
         requests++;
         while (responses < messageCount) {
             selector.poll(0L);
-            assertEquals("No disconnects should have occurred.", 0, selector.disconnected().size());
+            assertEquals("No disconnects should have occurred ." + selector.disconnected(),  0, selector.disconnected().size());
 
             for (NetworkReceive receive : selector.completedReceives()) {
-                assertEquals(prefix + "-" + responses, new String(Utils.toArray(receive.payload())));
+                assertEquals(prefix + "-" + responses, new String(Utils.toArray(receive.payload()), StandardCharsets.UTF_8));
                 responses++;
             }
 
             for (int i = 0; i < selector.completedSends().size() && requests < messageCount && selector.isChannelReady(node); i++, requests++) {
-                selector.send(new NetworkSend(node, ByteBuffer.wrap((prefix + "-" + requests).getBytes())));
+                selector.send(new NetworkSend(node, ByteBufferSend.sizePrefixed(ByteBuffer.wrap((prefix + "-" + requests).getBytes()))));
             }
         }
     }
@@ -82,8 +99,8 @@ public static void waitForChannelReady(Selector selector, String node) throws IO
     public static ChannelState waitForChannelClose(Selector selector, String node, ChannelState.State channelState)
             throws IOException {
         boolean closed = false;
-        for (int i = 0; i < 30; i++) {
-            selector.poll(1000L);
+        for (int i = 0; i < 300; i++) {
+            selector.poll(100L);
             if (selector.channel(node) == null && selector.closingChannel(node) == null) {
                 closed = true;
                 break;
@@ -94,4 +111,12 @@ public static ChannelState waitForChannelClose(Selector selector, String node, C
         assertEquals(channelState, finalState.state());
         return finalState;
     }
+
+    public static void completeDelayedChannelClose(Selector selector, long currentTimeNanos) {
+        selector.completeDelayedChannelClose(currentTimeNanos);
+    }
+
+    public static Map delayedClosingChannels(Selector selector) {
+        return selector.delayedClosingChannels();
+    }
 }
diff --git a/clients/src/test/java/org/apache/kafka/common/network/NioEchoServer.java b/clients/src/test/java/org/apache/kafka/common/network/NioEchoServer.java
index ad587b915ff19..078ebb84804a0 100644
--- a/clients/src/test/java/org/apache/kafka/common/network/NioEchoServer.java
+++ b/clients/src/test/java/org/apache/kafka/common/network/NioEchoServer.java
@@ -22,27 +22,35 @@
 import org.apache.kafka.common.config.AbstractConfig;
 import org.apache.kafka.common.metrics.KafkaMetric;
 import org.apache.kafka.common.metrics.Metrics;
+import org.apache.kafka.common.protocol.ApiKeys;
 import org.apache.kafka.common.security.auth.SecurityProtocol;
 import org.apache.kafka.common.security.authenticator.CredentialCache;
-import org.apache.kafka.common.security.scram.ScramCredentialUtils;
-import org.apache.kafka.common.security.scram.ScramMechanism;
+import org.apache.kafka.common.security.scram.ScramCredential;
+import org.apache.kafka.common.security.scram.internals.ScramMechanism;
 import org.apache.kafka.common.utils.LogContext;
 import org.apache.kafka.common.utils.MockTime;
-import org.apache.kafka.test.TestCondition;
+import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.test.TestUtils;
 
 import java.io.IOException;
 import java.net.InetSocketAddress;
 import java.nio.ByteBuffer;
+import java.nio.channels.GatheringByteChannel;
 import java.nio.channels.SelectionKey;
 import java.nio.channels.ServerSocketChannel;
 import java.nio.channels.SocketChannel;
 import java.nio.channels.WritableByteChannel;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
+import java.util.EnumSet;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Locale;
 import java.util.Map;
+import java.util.Set;
+
+import org.apache.kafka.common.security.token.delegation.internals.DelegationTokenCache;
 
 /**
  * Non-blocking EchoServer implementation that uses ChannelBuilder to create channels
@@ -50,6 +58,15 @@
  *
  */
 public class NioEchoServer extends Thread {
+    public enum MetricType {
+        TOTAL, RATE, AVG, MAX;
+
+        private final String metricNameSuffix = "-" + name().toLowerCase(Locale.ROOT);
+
+        public String metricNameSuffix() {
+            return metricNameSuffix;
+        }
+    }
 
     private static final double EPS = 0.0001;
 
@@ -59,12 +76,29 @@ public class NioEchoServer extends Thread {
     private final List socketChannels;
     private final AcceptorThread acceptorThread;
     private final Selector selector;
-    private volatile WritableByteChannel outputChannel;
+    private volatile GatheringByteChannel outputChannel;
     private final CredentialCache credentialCache;
     private final Metrics metrics;
+    private volatile int numSent = 0;
+    private volatile boolean closeKafkaChannels;
+    private final DelegationTokenCache tokenCache;
+    private final Time time;
+    
+    public NioEchoServer(ListenerName listenerName, SecurityProtocol securityProtocol, AbstractConfig config,
+                         String serverHost, ChannelBuilder channelBuilder, CredentialCache credentialCache, Time time) throws Exception {
+        this(listenerName, securityProtocol, config, serverHost, channelBuilder, credentialCache, 100, time);
+    }
+
+    public NioEchoServer(ListenerName listenerName, SecurityProtocol securityProtocol, AbstractConfig config,
+                         String serverHost, ChannelBuilder channelBuilder, CredentialCache credentialCache,
+                         int failedAuthenticationDelayMs, Time time) throws Exception {
+        this(listenerName, securityProtocol, config, serverHost, channelBuilder, credentialCache, failedAuthenticationDelayMs, time,
+                new DelegationTokenCache(ScramMechanism.mechanismNames()));
+    }
 
     public NioEchoServer(ListenerName listenerName, SecurityProtocol securityProtocol, AbstractConfig config,
-            String serverHost, ChannelBuilder channelBuilder, CredentialCache credentialCache) throws Exception {
+            String serverHost, ChannelBuilder channelBuilder, CredentialCache credentialCache,
+            int failedAuthenticationDelayMs, Time time, DelegationTokenCache tokenCache) throws Exception {
         super("echoserver");
         setDaemon(true);
         serverSocketChannel = ServerSocketChannel.open();
@@ -74,13 +108,22 @@ public NioEchoServer(ListenerName listenerName, SecurityProtocol securityProtoco
         this.socketChannels = Collections.synchronizedList(new ArrayList());
         this.newChannels = Collections.synchronizedList(new ArrayList());
         this.credentialCache = credentialCache;
-        if (securityProtocol == SecurityProtocol.SASL_PLAINTEXT || securityProtocol == SecurityProtocol.SASL_SSL)
-            ScramCredentialUtils.createCache(credentialCache, ScramMechanism.mechanismNames());
+        this.tokenCache = tokenCache;
+        if (securityProtocol == SecurityProtocol.SASL_PLAINTEXT || securityProtocol == SecurityProtocol.SASL_SSL) {
+            for (String mechanism : ScramMechanism.mechanismNames()) {
+                if (credentialCache.cache(mechanism, ScramCredential.class) == null)
+                    credentialCache.createCache(mechanism, ScramCredential.class);
+            }
+        }
+        LogContext logContext = new LogContext();
         if (channelBuilder == null)
-            channelBuilder = ChannelBuilders.serverChannelBuilder(listenerName, securityProtocol, config, credentialCache);
+            channelBuilder = ChannelBuilders.serverChannelBuilder(listenerName, false,
+                    securityProtocol, config, credentialCache, tokenCache, time, logContext);
         this.metrics = new Metrics();
-        this.selector = new Selector(5000, metrics, new MockTime(), "MetricGroup", channelBuilder, new LogContext());
+        this.selector = new Selector(10000, failedAuthenticationDelayMs, metrics, time,
+                "MetricGroup", channelBuilder, logContext);
         acceptorThread = new AcceptorThread();
+        this.time = time;
     }
 
     public int port() {
@@ -91,40 +134,70 @@ public CredentialCache credentialCache() {
         return credentialCache;
     }
 
-    @SuppressWarnings("deprecation")
+    public DelegationTokenCache tokenCache() {
+        return tokenCache;
+    }
+
     public double metricValue(String name) {
         for (Map.Entry entry : metrics.metrics().entrySet()) {
             if (entry.getKey().name().equals(name))
-                return entry.getValue().value();
+                return (double) entry.getValue().metricValue();
         }
         throw new IllegalStateException("Metric not found, " + name + ", found=" + metrics.metrics().keySet());
     }
 
     public void verifyAuthenticationMetrics(int successfulAuthentications, final int failedAuthentications)
             throws InterruptedException {
-        waitForMetric("successful-authentication", successfulAuthentications);
-        waitForMetric("failed-authentication", failedAuthentications);
-    }
-
-    private void waitForMetric(String name, final double expectedValue) throws InterruptedException {
-        final String totalName = name + "-total";
-        final String rateName = name + "-rate";
-        if (expectedValue == 0.0) {
-            assertEquals(expectedValue, metricValue(totalName), EPS);
-            assertEquals(expectedValue, metricValue(rateName), EPS);
-        } else {
-            TestUtils.waitForCondition(new TestCondition() {
-                @Override
-                public boolean conditionMet() {
-                    return Math.abs(metricValue(totalName) - expectedValue) <= EPS;
-                }
-            }, "Metric not updated " + totalName);
-            TestUtils.waitForCondition(new TestCondition() {
-                @Override
-                public boolean conditionMet() {
-                    return metricValue(rateName) > 0.0;
-                }
-            }, "Metric not updated " + rateName);
+        waitForMetrics("successful-authentication", successfulAuthentications,
+                EnumSet.of(MetricType.TOTAL, MetricType.RATE));
+        waitForMetrics("failed-authentication", failedAuthentications, EnumSet.of(MetricType.TOTAL, MetricType.RATE));
+    }
+
+    public void verifyReauthenticationMetrics(int successfulReauthentications, final int failedReauthentications)
+            throws InterruptedException {
+        waitForMetrics("successful-reauthentication", successfulReauthentications,
+                EnumSet.of(MetricType.TOTAL, MetricType.RATE));
+        waitForMetrics("failed-reauthentication", failedReauthentications,
+                EnumSet.of(MetricType.TOTAL, MetricType.RATE));
+        waitForMetrics("successful-authentication-no-reauth", 0, EnumSet.of(MetricType.TOTAL));
+        if (!(time instanceof MockTime)) {
+            waitForMetrics("reauthentication-latency", Math.signum(successfulReauthentications),
+                    EnumSet.of(MetricType.MAX, MetricType.AVG));
+        }
+    }
+
+    public void verifyAuthenticationNoReauthMetric(int successfulAuthenticationNoReauths) throws InterruptedException {
+        waitForMetrics("successful-authentication-no-reauth", successfulAuthenticationNoReauths,
+                EnumSet.of(MetricType.TOTAL));
+    }
+
+    public void waitForMetric(String name, final double expectedValue) throws InterruptedException {
+        waitForMetrics(name, expectedValue, EnumSet.of(MetricType.TOTAL, MetricType.RATE));
+    }
+
+    public void waitForMetrics(String namePrefix, final double expectedValue, Set metricTypes)
+            throws InterruptedException {
+        long maxAggregateWaitMs = 15000;
+        long startMs = time.milliseconds();
+        for (MetricType metricType : metricTypes) {
+            long currentElapsedMs = time.milliseconds() - startMs;
+            long thisMaxWaitMs = maxAggregateWaitMs - currentElapsedMs;
+            String metricName = namePrefix + metricType.metricNameSuffix();
+            if (expectedValue == 0.0) {
+                Double expected = expectedValue;
+                if (metricType == MetricType.MAX || metricType == MetricType.AVG)
+                    expected = Double.NaN;
+
+                assertEquals("Metric not updated " + metricName + " expected:<" + expectedValue + "> but was:<"
+                    + metricValue(metricName) + ">", expected, metricValue(metricName), EPS);
+            } else if (metricType == MetricType.TOTAL)
+                TestUtils.waitForCondition(() -> Math.abs(metricValue(metricName) - expectedValue) <= EPS,
+                        thisMaxWaitMs, () -> "Metric not updated " + metricName + " expected:<" + expectedValue
+                                + "> but was:<" + metricValue(metricName) + ">");
+            else
+                TestUtils.waitForCondition(() -> metricValue(metricName) > 0.0, thisMaxWaitMs,
+                    () -> "Metric not updated " + metricName + " expected: but was:<"
+                                + metricValue(metricName) + ">");
         }
     }
 
@@ -133,7 +206,7 @@ public void run() {
         try {
             acceptorThread.start();
             while (serverSocketChannel.isOpen()) {
-                selector.poll(1000);
+                selector.poll(100);
                 synchronized (newChannels) {
                     for (SocketChannel socketChannel : newChannels) {
                         String id = id(socketChannel);
@@ -142,30 +215,51 @@ public void run() {
                     }
                     newChannels.clear();
                 }
+                if (closeKafkaChannels) {
+                    for (KafkaChannel channel : selector.channels())
+                        selector.close(channel.id());
+                }
 
-                List completedReceives = selector.completedReceives();
+                Collection completedReceives = selector.completedReceives();
                 for (NetworkReceive rcv : completedReceives) {
                     KafkaChannel channel = channel(rcv.source());
-                    String channelId = channel.id();
-                    selector.mute(channelId);
-                    NetworkSend send = new NetworkSend(rcv.source(), rcv.payload());
-                    if (outputChannel == null)
-                        selector.send(send);
-                    else {
-                        for (ByteBuffer buffer : send.buffers)
-                            outputChannel.write(buffer);
-                        selector.unmute(channelId);
+                    if (!maybeBeginServerReauthentication(channel, rcv, time)) {
+                        String channelId = channel.id();
+                        selector.mute(channelId);
+                        NetworkSend send = new NetworkSend(rcv.source(), ByteBufferSend.sizePrefixed(rcv.payload()));
+                        if (outputChannel == null)
+                            selector.send(send);
+                        else {
+                            send.writeTo(outputChannel);
+                            selector.unmute(channelId);
+                        }
                     }
                 }
-                for (Send send : selector.completedSends())
-                    selector.unmute(send.destination());
-
+                for (NetworkSend send : selector.completedSends()) {
+                    selector.unmute(send.destinationId());
+                    numSent += 1;
+                }
             }
         } catch (IOException e) {
             // ignore
         }
     }
 
+    public int numSent() {
+        return numSent;
+    }
+
+    private static boolean maybeBeginServerReauthentication(KafkaChannel channel, NetworkReceive networkReceive, Time time) {
+        try {
+            if (TestUtils.apiKeyFrom(networkReceive) == ApiKeys.SASL_HANDSHAKE) {
+                return channel.maybeBeginServerReauthentication(networkReceive, () -> time.nanoseconds());
+            }
+        } catch (Exception e) {
+            // ignore
+        }
+        return false;
+    }
+
     private String id(SocketChannel channel) {
         return channel.socket().getLocalAddress().getHostAddress() + ":" + channel.socket().getLocalPort() + "-" +
                 channel.socket().getInetAddress().getHostAddress() + ":" + channel.socket().getPort();
@@ -182,22 +276,67 @@ private KafkaChannel channel(String id) {
      * the responses (eg. testing graceful close).
      */
     public void outputChannel(WritableByteChannel channel) {
-        this.outputChannel = channel;
+        if (channel instanceof GatheringByteChannel)
+            this.outputChannel = (GatheringByteChannel) channel;
+        else {
+            this.outputChannel = new GatheringByteChannel() {
+                @Override
+                public boolean isOpen() {
+                    return channel.isOpen();
+                }
+
+                @Override
+                public void close() throws IOException {
+                    channel.close();
+                }
+
+                @Override
+                public int write(ByteBuffer src) throws IOException {
+                    return channel.write(src);
+                }
+
+                @Override
+                public long write(ByteBuffer[] srcs, int offset, int length) throws IOException {
+                    long result = 0;
+                    for (int i = offset; i < offset + length; ++i)
+                        result += write(srcs[i]);
+                    return result;
+                }
+
+                @Override
+                public long write(ByteBuffer[] srcs) throws IOException {
+                    return write(srcs, 0, srcs.length);
+                }
+            };
+        }
     }
 
     public Selector selector() {
         return selector;
     }
 
-    public void closeConnections() throws IOException {
-        for (SocketChannel channel : socketChannels)
+    public void closeKafkaChannels() {
+        closeKafkaChannels = true;
+        selector.wakeup();
+        try {
+            TestUtils.waitForCondition(() -> selector.channels().isEmpty(), "Channels not closed");
+        } catch (InterruptedException e) {
+            throw new RuntimeException(e);
+        } finally {
+            closeKafkaChannels = false;
+        }
+    }
+
+    public void closeSocketChannels() throws IOException {
+        for (SocketChannel channel : socketChannels) {
             channel.close();
+        }
         socketChannels.clear();
     }
 
     public void close() throws IOException, InterruptedException {
         this.serverSocketChannel.close();
-        closeConnections();
+        closeSocketChannels();
         acceptorThread.interrupt();
         acceptorThread.join();
         interrupt();
@@ -205,9 +344,10 @@ public void close() throws IOException, InterruptedException {
     }
 
     private class AcceptorThread extends Thread {
-        public AcceptorThread() throws IOException {
+        public AcceptorThread() {
             setName("acceptor");
         }
+        @Override
         public void run() {
             try {
                 java.nio.channels.Selector acceptSelector = java.nio.channels.Selector.open();
diff --git a/clients/src/test/java/org/apache/kafka/common/network/SaslChannelBuilderTest.java b/clients/src/test/java/org/apache/kafka/common/network/SaslChannelBuilderTest.java
index 750fd01b87b90..fd633ac2d76ba 100644
--- a/clients/src/test/java/org/apache/kafka/common/network/SaslChannelBuilderTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/network/SaslChannelBuilderTest.java
@@ -17,62 +17,217 @@
 package org.apache.kafka.common.network;
 
 import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.common.config.SaslConfigs;
 import org.apache.kafka.common.config.SslConfigs;
+import org.apache.kafka.common.config.internals.BrokerSecurityConfigs;
+import org.apache.kafka.common.security.TestSecurityConfig;
+import org.apache.kafka.common.security.auth.KafkaPrincipal;
 import org.apache.kafka.common.security.auth.SecurityProtocol;
 import org.apache.kafka.common.security.JaasContext;
 import org.apache.kafka.common.security.authenticator.TestJaasConfig;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerLoginModule;
 import org.apache.kafka.common.security.plain.PlainLoginModule;
+import org.apache.kafka.common.security.scram.ScramLoginModule;
+import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.common.utils.Time;
+import org.ietf.jgss.GSSContext;
+import org.ietf.jgss.GSSCredential;
+import org.ietf.jgss.GSSManager;
+import org.ietf.jgss.GSSName;
+import org.ietf.jgss.Oid;
+import org.junit.After;
 import org.junit.Test;
+import org.mockito.Mockito;
 
+import javax.security.auth.Subject;
+import javax.security.auth.callback.CallbackHandler;
+import javax.security.auth.login.LoginException;
+import javax.security.auth.spi.LoginModule;
+import java.lang.reflect.Field;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.Map;
 
+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;
 import static org.junit.Assert.fail;
 
+
 public class SaslChannelBuilderTest {
 
+    @After
+    public void tearDown() {
+        System.clearProperty(SaslChannelBuilder.GSS_NATIVE_PROP);
+    }
+
     @Test
     public void testCloseBeforeConfigureIsIdempotent() {
-        SaslChannelBuilder builder = createChannelBuilder(SecurityProtocol.SASL_PLAINTEXT);
+        SaslChannelBuilder builder = createChannelBuilder(SecurityProtocol.SASL_PLAINTEXT, "PLAIN");
         builder.close();
-        assertNull(builder.loginManager());
+        assertTrue(builder.loginManagers().isEmpty());
         builder.close();
-        assertNull(builder.loginManager());
+        assertTrue(builder.loginManagers().isEmpty());
     }
 
     @Test
     public void testCloseAfterConfigIsIdempotent() {
-        SaslChannelBuilder builder = createChannelBuilder(SecurityProtocol.SASL_PLAINTEXT);
-        builder.configure(new HashMap());
-        assertNotNull(builder.loginManager());
+        SaslChannelBuilder builder = createChannelBuilder(SecurityProtocol.SASL_PLAINTEXT, "PLAIN");
+        builder.configure(new HashMap<>());
+        assertNotNull(builder.loginManagers().get("PLAIN"));
         builder.close();
-        assertNull(builder.loginManager());
+        assertTrue(builder.loginManagers().isEmpty());
         builder.close();
-        assertNull(builder.loginManager());
+        assertTrue(builder.loginManagers().isEmpty());
     }
 
     @Test
     public void testLoginManagerReleasedIfConfigureThrowsException() {
-        SaslChannelBuilder builder = createChannelBuilder(SecurityProtocol.SASL_SSL);
+        SaslChannelBuilder builder = createChannelBuilder(SecurityProtocol.SASL_SSL, "PLAIN");
         try {
             // Use invalid config so that an exception is thrown
             builder.configure(Collections.singletonMap(SslConfigs.SSL_ENABLED_PROTOCOLS_CONFIG, "1"));
             fail("Exception should have been thrown");
         } catch (KafkaException e) {
-            assertNull(builder.loginManager());
+            assertTrue(builder.loginManagers().isEmpty());
         }
         builder.close();
-        assertNull(builder.loginManager());
+        assertTrue(builder.loginManagers().isEmpty());
     }
 
-    private SaslChannelBuilder createChannelBuilder(SecurityProtocol securityProtocol) {
+    @Test
+    public void testNativeGssapiCredentials() throws Exception {
+        System.setProperty(SaslChannelBuilder.GSS_NATIVE_PROP, "true");
+
         TestJaasConfig jaasConfig = new TestJaasConfig();
-        jaasConfig.addEntry("jaasContext", PlainLoginModule.class.getName(), new HashMap());
-        JaasContext jaasContext = new JaasContext("jaasContext", JaasContext.Type.SERVER, jaasConfig);
-        return new SaslChannelBuilder(Mode.CLIENT, jaasContext, securityProtocol, new ListenerName("PLAIN"),
-                "PLAIN", true, null);
+        jaasConfig.addEntry("jaasContext", TestGssapiLoginModule.class.getName(), new HashMap<>());
+        JaasContext jaasContext = new JaasContext("jaasContext", JaasContext.Type.SERVER, jaasConfig, null);
+        Map jaasContexts = Collections.singletonMap("GSSAPI", jaasContext);
+        GSSManager gssManager = Mockito.mock(GSSManager.class);
+        GSSName gssName = Mockito.mock(GSSName.class);
+        Mockito.when(gssManager.createName(Mockito.anyString(), Mockito.any()))
+                .thenAnswer(unused -> gssName);
+        Oid oid = new Oid("1.2.840.113554.1.2.2");
+        Mockito.when(gssManager.createCredential(gssName, GSSContext.INDEFINITE_LIFETIME, oid, GSSCredential.ACCEPT_ONLY))
+                .thenAnswer(unused -> Mockito.mock(GSSCredential.class));
+
+        SaslChannelBuilder channelBuilder1 = createGssapiChannelBuilder(jaasContexts, gssManager);
+        assertEquals(1, channelBuilder1.subject("GSSAPI").getPrincipals().size());
+        assertEquals(1, channelBuilder1.subject("GSSAPI").getPrivateCredentials().size());
+
+        SaslChannelBuilder channelBuilder2 = createGssapiChannelBuilder(jaasContexts, gssManager);
+        assertEquals(1, channelBuilder2.subject("GSSAPI").getPrincipals().size());
+        assertEquals(1, channelBuilder2.subject("GSSAPI").getPrivateCredentials().size());
+        assertSame(channelBuilder1.subject("GSSAPI"), channelBuilder2.subject("GSSAPI"));
+
+        Mockito.verify(gssManager, Mockito.times(1))
+                .createCredential(gssName, GSSContext.INDEFINITE_LIFETIME, oid, GSSCredential.ACCEPT_ONLY);
+    }
+
+    /**
+     * Verify that unparsed broker configs don't break clients. This is to ensure that clients
+     * created by brokers are not broken if broker configs are passed to clients.
+     */
+    @Test
+    public void testClientChannelBuilderWithBrokerConfigs() throws Exception {
+        Map configs = new HashMap<>();
+        CertStores certStores = new CertStores(false, "client", "localhost");
+        configs.putAll(certStores.getTrustingConfig(certStores));
+        configs.put(SaslConfigs.SASL_KERBEROS_SERVICE_NAME, "kafka");
+        configs.putAll(new ConfigDef().withClientSaslSupport().parse(configs));
+        for (Field field : BrokerSecurityConfigs.class.getFields()) {
+            if (field.getName().endsWith("_CONFIG"))
+                configs.put(field.get(BrokerSecurityConfigs.class).toString(), "somevalue");
+        }
+
+        SaslChannelBuilder plainBuilder = createChannelBuilder(SecurityProtocol.SASL_PLAINTEXT, "PLAIN");
+        plainBuilder.configure(configs);
+
+        SaslChannelBuilder gssapiBuilder = createChannelBuilder(SecurityProtocol.SASL_PLAINTEXT, "GSSAPI");
+        gssapiBuilder.configure(configs);
+
+        SaslChannelBuilder oauthBearerBuilder = createChannelBuilder(SecurityProtocol.SASL_PLAINTEXT, "OAUTHBEARER");
+        oauthBearerBuilder.configure(configs);
+
+        SaslChannelBuilder scramBuilder = createChannelBuilder(SecurityProtocol.SASL_PLAINTEXT, "SCRAM-SHA-256");
+        scramBuilder.configure(configs);
+
+        SaslChannelBuilder saslSslBuilder = createChannelBuilder(SecurityProtocol.SASL_SSL, "PLAIN");
+        saslSslBuilder.configure(configs);
     }
 
+    private SaslChannelBuilder createGssapiChannelBuilder(Map jaasContexts, GSSManager gssManager) {
+        SaslChannelBuilder channelBuilder = new SaslChannelBuilder(Mode.SERVER, jaasContexts,
+                SecurityProtocol.SASL_PLAINTEXT,
+                new ListenerName("GSSAPI"), false, "GSSAPI",
+                true, null, null, Time.SYSTEM, new LogContext()) {
+
+            @Override
+            protected GSSManager gssManager() {
+                return gssManager;
+            }
+        };
+        Map props = Collections.singletonMap(SaslConfigs.SASL_KERBEROS_SERVICE_NAME, "kafka");
+        channelBuilder.configure(new TestSecurityConfig(props).values());
+        return channelBuilder;
+    }
+
+    private SaslChannelBuilder createChannelBuilder(SecurityProtocol securityProtocol, String saslMechanism) {
+        Class loginModule = null;
+        switch (saslMechanism) {
+            case "PLAIN":
+                loginModule = PlainLoginModule.class;
+                break;
+            case "SCRAM-SHA-256":
+                loginModule = ScramLoginModule.class;
+                break;
+            case "OAUTHBEARER":
+                loginModule = OAuthBearerLoginModule.class;
+                break;
+            case "GSSAPI":
+                loginModule = TestGssapiLoginModule.class;
+                break;
+            default:
+                throw new IllegalArgumentException("Unsupported SASL mechanism " + saslMechanism);
+        }
+        TestJaasConfig jaasConfig = new TestJaasConfig();
+        jaasConfig.addEntry("jaasContext", loginModule.getName(), new HashMap<>());
+        JaasContext jaasContext = new JaasContext("jaasContext", JaasContext.Type.SERVER, jaasConfig, null);
+        Map jaasContexts = Collections.singletonMap(saslMechanism, jaasContext);
+        return new SaslChannelBuilder(Mode.CLIENT, jaasContexts, securityProtocol, new ListenerName(saslMechanism),
+                false, saslMechanism, true, null,
+                null, Time.SYSTEM, new LogContext());
+    }
+
+    public static final class TestGssapiLoginModule implements LoginModule {
+        private Subject subject;
+
+        @Override
+        public void initialize(Subject subject, CallbackHandler callbackHandler, Map sharedState, Map options) {
+            this.subject = subject;
+        }
+
+        @Override
+        public boolean login() throws LoginException {
+            subject.getPrincipals().add(new KafkaPrincipal("User", "kafka@kafka1.example.com"));
+            return true;
+        }
+
+        @Override
+        public boolean commit() throws LoginException {
+            return true;
+        }
+
+        @Override
+        public boolean abort() throws LoginException {
+            return true;
+        }
+
+        @Override
+        public boolean logout() throws LoginException {
+            return true;
+        }
+    }
 }
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 36b847b4fee03..14be0390ee90d 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
@@ -17,19 +17,23 @@
 package org.apache.kafka.common.network;
 
 import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.MetricName;
 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.Metrics;
 import org.apache.kafka.common.security.auth.SecurityProtocol;
 import org.apache.kafka.common.utils.LogContext;
 import org.apache.kafka.common.utils.MockTime;
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.common.utils.Utils;
+import org.apache.kafka.test.TestCondition;
 import org.apache.kafka.test.TestUtils;
-import org.easymock.IMocksControl;
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.Timeout;
 
 import java.io.ByteArrayOutputStream;
 import java.io.DataOutputStream;
@@ -41,30 +45,44 @@
 import java.nio.channels.SelectionKey;
 import java.nio.channels.ServerSocketChannel;
 import java.nio.channels.SocketChannel;
+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.Optional;
 import java.util.Random;
 import java.util.Set;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Supplier;
+import java.util.stream.Collectors;
 
-import static org.easymock.EasyMock.createControl;
-import static org.easymock.EasyMock.expect;
-import static org.easymock.EasyMock.expectLastCall;
+import static java.util.Arrays.asList;
 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;
+import static org.mockito.Mockito.atLeastOnce;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 
 /**
  * A set of tests for the selector. These use a test harness that runs a simple socket server that echos back responses.
  */
 public class SelectorTest {
+    @Rule
+    final public Timeout globalTimeout = Timeout.millis(240000);
 
     protected static final int BUFFER_SIZE = 4 * 1024;
+    private static final String METRIC_GROUP = "MetricGroup";
 
     protected EchoServer server;
     protected Time time;
@@ -78,10 +96,10 @@ public void setUp() throws Exception {
         this.server = new EchoServer(SecurityProtocol.PLAINTEXT, configs);
         this.server.start();
         this.time = new MockTime();
-        this.channelBuilder = new PlaintextChannelBuilder();
-        this.channelBuilder.configure(configs);
+        this.channelBuilder = new PlaintextChannelBuilder(ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT));
+        this.channelBuilder.configure(clientConfigs());
         this.metrics = new Metrics();
-        this.selector = new Selector(5000, this.metrics, time, "MetricGroup", channelBuilder, new LogContext());
+        this.selector = new Selector(5000, this.metrics, time, METRIC_GROUP, channelBuilder, new LogContext());
     }
 
     @After
@@ -99,21 +117,38 @@ public SecurityProtocol securityProtocol() {
         return SecurityProtocol.PLAINTEXT;
     }
 
+    protected Map clientConfigs() {
+        return new HashMap<>();
+    }
+
     /**
      * Validate that when the server disconnects, a client send ends up with that node in the disconnected list.
      */
     @Test
     public void testServerDisconnect() throws Exception {
-        String node = "0";
+        final String node = "0";
 
         // connect and do a simple request
         blockingConnect(node);
         assertEquals("hello", blockingRequest(node, "hello"));
 
+        KafkaChannel channel = selector.channel(node);
+
         // disconnect
         this.server.closeConnections();
-        while (!selector.disconnected().containsKey(node))
-            selector.poll(1000L);
+        TestUtils.waitForCondition(new TestCondition() {
+            @Override
+            public boolean conditionMet() {
+                try {
+                    selector.poll(1000L);
+                    return selector.disconnected().containsKey(node);
+                } catch (IOException e) {
+                    throw new RuntimeException(e);
+                }
+            }
+        }, 5000, "Failed to observe disconnected node in disconnected set");
+
+        assertNull(channel.selectionKey().attachment());
 
         // reconnect and do another request
         blockingConnect(node);
@@ -186,8 +221,8 @@ public void testNormalOperation() throws Exception {
         for (int i = 0; i < conns; i++)
             connect(Integer.toString(i), addr);
         // send echo requests and receive responses
-        Map requests = new HashMap();
-        Map responses = new HashMap();
+        Map requests = new HashMap<>();
+        Map responses = new HashMap<>();
         int responseCount = 0;
         for (int i = 0; i < conns; i++) {
             String node = Integer.toString(i);
@@ -218,8 +253,8 @@ public void testNormalOperation() throws Exception {
             }
 
             // prepare new sends for the next round
-            for (Send send : selector.completedSends()) {
-                String dest = send.destination();
+            for (NetworkSend send : selector.completedSends()) {
+                String dest = send.destinationId();
                 if (requests.containsKey(dest))
                     requests.put(dest, requests.get(dest) + 1);
                 else
@@ -228,6 +263,21 @@ public void testNormalOperation() throws Exception {
                     selector.send(createSend(dest, dest + "-" + requests.get(dest)));
             }
         }
+        if (channelBuilder instanceof PlaintextChannelBuilder) {
+            assertEquals(0, cipherMetrics(metrics).size());
+        } else {
+            TestUtils.waitForCondition(() -> cipherMetrics(metrics).size() == 1,
+                "Waiting for cipher metrics to be created.");
+            assertEquals(Integer.valueOf(5), cipherMetrics(metrics).get(0).metricValue());
+        }
+    }
+
+    static List cipherMetrics(Metrics metrics) {
+        return metrics.metrics().entrySet().stream().
+            filter(e -> e.getKey().description().
+                contains("The number of connections with this SSL cipher and protocol.")).
+            map(e -> e.getValue()).
+            collect(Collectors.toList());
     }
 
     /**
@@ -241,6 +291,42 @@ public void testSendLargeRequest() throws Exception {
         assertEquals(big, blockingRequest(node, big));
     }
 
+    @Test
+    public void testPartialSendAndReceiveReflectedInMetrics() throws Exception {
+        // We use a large payload to attempt to trigger the partial send and receive logic.
+        int payloadSize = 20 * BUFFER_SIZE;
+        String payload = TestUtils.randomString(payloadSize);
+        String nodeId = "0";
+        blockingConnect(nodeId);
+        ByteBufferSend send = ByteBufferSend.sizePrefixed(ByteBuffer.wrap(payload.getBytes()));
+        NetworkSend networkSend = new NetworkSend(nodeId, send);
+
+        selector.send(networkSend);
+        KafkaChannel channel = selector.channel(nodeId);
+
+        KafkaMetric outgoingByteTotal = findUntaggedMetricByName("outgoing-byte-total");
+        KafkaMetric incomingByteTotal = findUntaggedMetricByName("incoming-byte-total");
+
+        TestUtils.waitForCondition(() -> {
+            long bytesSent = send.size() - send.remaining();
+            assertEquals(bytesSent, ((Double) outgoingByteTotal.metricValue()).longValue());
+
+            NetworkReceive currentReceive = channel.currentReceive();
+            if (currentReceive != null) {
+                assertEquals(currentReceive.bytesRead(), ((Double) incomingByteTotal.metricValue()).intValue());
+            }
+
+            selector.poll(50);
+            return !selector.completedReceives().isEmpty();
+        }, "Failed to receive expected response");
+
+        KafkaMetric requestTotal = findUntaggedMetricByName("request-total");
+        assertEquals(1, ((Double) requestTotal.metricValue()).intValue());
+
+        KafkaMetric responseTotal = findUntaggedMetricByName("response-total");
+        assertEquals(1, ((Double) responseTotal.metricValue()).intValue());
+    }
+
     @Test
     public void testLargeMessageSequence() throws Exception {
         int bufferSize = 512 * 1024;
@@ -252,11 +338,6 @@ public void testLargeMessageSequence() throws Exception {
         sendAndReceive(node, requestPrefix, 0, reqs);
     }
 
-
-
-    /**
-     * Test sending an empty string
-     */
     @Test
     public void testEmptyRequest() throws Exception {
         String node = "0";
@@ -264,6 +345,36 @@ public void testEmptyRequest() throws Exception {
         assertEquals("", blockingRequest(node, ""));
     }
 
+    @Test
+    public void testClearCompletedSendsAndReceives() throws Exception {
+        int bufferSize = 1024;
+        String node = "0";
+        InetSocketAddress addr = new InetSocketAddress("localhost", server.port);
+        connect(node, addr);
+        String request = TestUtils.randomString(bufferSize);
+        selector.send(createSend(node, request));
+        boolean sent = false;
+        boolean received = false;
+        while (!sent || !received) {
+            selector.poll(1000L);
+            assertEquals("No disconnects should have occurred.", 0, selector.disconnected().size());
+            if (!selector.completedSends().isEmpty()) {
+                assertEquals(1, selector.completedSends().size());
+                selector.clearCompletedSends();
+                assertEquals(0, selector.completedSends().size());
+                sent = true;
+            }
+
+            if (!selector.completedReceives().isEmpty()) {
+                assertEquals(1, selector.completedReceives().size());
+                assertEquals(request, asString(selector.completedReceives().iterator().next()));
+                selector.clearCompletedReceives();
+                assertEquals(0, selector.completedReceives().size());
+                received = true;
+            }
+        }
+    }
+
     @Test(expected = IllegalStateException.class)
     public void testExistingConnectionId() throws IOException {
         blockingConnect("0");
@@ -283,22 +394,49 @@ public void testMute() throws Exception {
         while (selector.completedReceives().isEmpty())
             selector.poll(5);
         assertEquals("We should have only one response", 1, selector.completedReceives().size());
-        assertEquals("The response should not be from the muted node", "0", selector.completedReceives().get(0).source());
+        assertEquals("The response should not be from the muted node", "0", selector.completedReceives().iterator().next().source());
 
         selector.unmute("1");
         do {
             selector.poll(5);
         } while (selector.completedReceives().isEmpty());
         assertEquals("We should have only one response", 1, selector.completedReceives().size());
-        assertEquals("The response should be from the previously muted node", "1", selector.completedReceives().get(0).source());
+        assertEquals("The response should be from the previously muted node", "1", selector.completedReceives().iterator().next().source());
+    }
+
+    @Test
+    public void testCloseAllChannels() throws Exception {
+        AtomicInteger closedChannelsCount = new AtomicInteger(0);
+        ChannelBuilder channelBuilder = new PlaintextChannelBuilder(null) {
+            private int channelIndex = 0;
+            @Override
+            KafkaChannel buildChannel(String id, TransportLayer transportLayer, Supplier authenticatorCreator,
+                                      int maxReceiveSize, MemoryPool memoryPool, ChannelMetadataRegistry metadataRegistry) {
+                return new KafkaChannel(id, transportLayer, authenticatorCreator, maxReceiveSize, memoryPool, metadataRegistry) {
+                    private final int index = channelIndex++;
+                    @Override
+                    public void close() throws IOException {
+                        closedChannelsCount.getAndIncrement();
+                        if (index == 0) throw new RuntimeException("you should fail");
+                        else super.close();
+                    }
+                };
+            }
+        };
+        channelBuilder.configure(clientConfigs());
+        Selector selector = new Selector(5000, new Metrics(), new MockTime(), "MetricGroup", channelBuilder, new LogContext());
+        selector.connect("0", new InetSocketAddress("localhost", server.port), BUFFER_SIZE, BUFFER_SIZE);
+        selector.connect("1", new InetSocketAddress("localhost", server.port), BUFFER_SIZE, BUFFER_SIZE);
+        assertThrows(RuntimeException.class, selector::close);
+        assertEquals(2, closedChannelsCount.get());
     }
 
     @Test
     public void registerFailure() throws Exception {
-        ChannelBuilder channelBuilder = new PlaintextChannelBuilder() {
+        ChannelBuilder channelBuilder = new PlaintextChannelBuilder(null) {
             @Override
             public KafkaChannel buildChannel(String id, SelectionKey key, int maxReceiveSize,
-                    MemoryPool memoryPool) throws KafkaException {
+                    MemoryPool memoryPool, ChannelMetadataRegistry metadataRegistry) throws KafkaException {
                 throw new RuntimeException("Test exception");
             }
             @Override
@@ -318,25 +456,6 @@ public void close() {
         selector.close();
     }
 
-    @Test
-    public void testCloseConnectionInClosingState() throws Exception {
-        KafkaChannel channel = createConnectionWithStagedReceives(5);
-        String id = channel.id();
-        selector.mute(id); // Mute to allow channel to be expired even if more data is available for read
-        time.sleep(6000);  // The max idle time is 5000ms
-        selector.poll(0);
-        assertNull("Channel not expired", selector.channel(id));
-        assertEquals(channel, selector.closingChannel(id));
-        assertEquals(ChannelState.EXPIRED, channel.state());
-        selector.close(id);
-        assertNull("Channel not removed from channels", selector.channel(id));
-        assertNull("Channel not removed from closingChannels", selector.closingChannel(id));
-        assertTrue("Unexpected disconnect notification", selector.disconnected().isEmpty());
-        assertEquals(ChannelState.EXPIRED, channel.state());
-        selector.poll(0);
-        assertTrue("Unexpected disconnect notification", selector.disconnected().isEmpty());
-    }
-
     @Test
     public void testCloseOldestConnection() throws Exception {
         String id = "0";
@@ -350,67 +469,226 @@ public void testCloseOldestConnection() throws Exception {
     }
 
     @Test
-    public void testCloseOldestConnectionWithOneStagedReceive() throws Exception {
-        verifyCloseOldestConnectionWithStagedReceives(1);
+    public void testIdleExpiryWithoutReadyKeys() throws IOException {
+        String id = "0";
+        selector.connect(id, new InetSocketAddress("localhost", server.port), BUFFER_SIZE, BUFFER_SIZE);
+        KafkaChannel channel = selector.channel(id);
+        channel.selectionKey().interestOps(0);
+
+        time.sleep(6000); // The max idle time is 5000ms
+        selector.poll(0);
+        assertTrue("The idle connection should have been closed", selector.disconnected().containsKey(id));
+        assertEquals(ChannelState.EXPIRED, selector.disconnected().get(id));
     }
 
     @Test
-    public void testCloseOldestConnectionWithMultipleStagedReceives() throws Exception {
-        verifyCloseOldestConnectionWithStagedReceives(5);
+    public void testImmediatelyConnectedCleaned() throws Exception {
+        Metrics metrics = new Metrics(); // new metrics object to avoid metric registration conflicts
+        Selector selector = new ImmediatelyConnectingSelector(5000, metrics, time, "MetricGroup", channelBuilder, new LogContext());
+
+        try {
+            testImmediatelyConnectedCleaned(selector, true);
+            testImmediatelyConnectedCleaned(selector, false);
+        } finally {
+            selector.close();
+            metrics.close();
+        }
+    }
+
+    private static class ImmediatelyConnectingSelector extends Selector {
+        public ImmediatelyConnectingSelector(long connectionMaxIdleMS,
+                                             Metrics metrics,
+                                             Time time,
+                                             String metricGrpPrefix,
+                                             ChannelBuilder channelBuilder,
+                                             LogContext logContext) {
+            super(connectionMaxIdleMS, metrics, time, metricGrpPrefix, channelBuilder, logContext);
+        }
+
+        @Override
+        protected boolean doConnect(SocketChannel channel, InetSocketAddress address) throws IOException {
+            // Use a blocking connect to trigger the immediately connected path
+            channel.configureBlocking(true);
+            boolean connected = super.doConnect(channel, address);
+            channel.configureBlocking(false);
+            return connected;
+        }
     }
 
-    private KafkaChannel createConnectionWithStagedReceives(int maxStagedReceives) throws Exception {
+    private void testImmediatelyConnectedCleaned(Selector selector, boolean closeAfterFirstPoll) throws Exception {
         String id = "0";
-        blockingConnect(id);
-        KafkaChannel channel = selector.channel(id);
-        int retries = 100;
+        selector.connect(id, new InetSocketAddress("localhost", server.port), BUFFER_SIZE, BUFFER_SIZE);
+        verifyNonEmptyImmediatelyConnectedKeys(selector);
+        if (closeAfterFirstPoll) {
+            selector.poll(0);
+            verifyEmptyImmediatelyConnectedKeys(selector);
+        }
+        selector.close(id);
+        verifySelectorEmpty(selector);
+    }
 
-        do {
-            selector.mute(id);
-            for (int i = 0; i <= maxStagedReceives; i++) {
-                selector.send(createSend(id, String.valueOf(i)));
-                selector.poll(1000);
+    /**
+     * Verify that if Selector#connect fails and throws an Exception, all related objects
+     * are cleared immediately before the exception is propagated.
+     */
+    @Test
+    public void testConnectException() throws Exception {
+        Metrics metrics = new Metrics();
+        AtomicBoolean throwIOException = new AtomicBoolean();
+        Selector selector = new ImmediatelyConnectingSelector(5000, metrics, time, "MetricGroup", channelBuilder, new LogContext()) {
+            @Override
+            protected SelectionKey registerChannel(String id, SocketChannel socketChannel, int interestedOps) throws IOException {
+                SelectionKey key = super.registerChannel(id, socketChannel, interestedOps);
+                key.cancel();
+                if (throwIOException.get())
+                    throw new IOException("Test exception");
+                return key;
             }
+        };
 
-            selector.unmute(id);
-            do {
-                selector.poll(1000);
-            } while (selector.completedReceives().isEmpty());
-        } while (selector.numStagedReceives(channel) == 0 && --retries > 0);
-        assertTrue("No staged receives after 100 attempts", selector.numStagedReceives(channel) > 0);
+        try {
+            verifyImmediatelyConnectedException(selector, "0");
+            throwIOException.set(true);
+            verifyImmediatelyConnectedException(selector, "1");
+        } finally {
+            selector.close();
+            metrics.close();
+        }
+    }
 
-        return channel;
+    private void verifyImmediatelyConnectedException(Selector selector, String id) throws Exception {
+        try {
+            selector.connect(id, new InetSocketAddress("localhost", server.port), BUFFER_SIZE, BUFFER_SIZE);
+            fail("Expected exception not thrown");
+        } catch (Exception e) {
+            verifyEmptyImmediatelyConnectedKeys(selector);
+            assertNull("Channel not removed", selector.channel(id));
+            ensureEmptySelectorFields(selector);
+        }
+    }
+
+    /*
+     * Verifies that a muted connection is expired on idle timeout even if there are pending
+     * receives on the socket.
+     */
+    @Test
+    public void testExpireConnectionWithPendingReceives() throws Exception {
+        KafkaChannel channel = createConnectionWithPendingReceives(5);
+        verifyChannelExpiry(channel);
+    }
+
+    /**
+     * Verifies that a muted connection closed by peer is expired on idle timeout even if there are pending
+     * receives on the socket.
+     */
+    @Test
+    public void testExpireClosedConnectionWithPendingReceives() throws Exception {
+        KafkaChannel channel = createConnectionWithPendingReceives(5);
+        server.closeConnections();
+        verifyChannelExpiry(channel);
+    }
+
+    private void verifyChannelExpiry(KafkaChannel channel) throws Exception {
+        String id = channel.id();
+        selector.mute(id); // Mute to allow channel to be expired even if more data is available for read
+        time.sleep(6000);  // The max idle time is 5000ms
+        selector.poll(0);
+        assertNull("Channel not expired", selector.channel(id));
+        assertNull("Channel not removed from closingChannels", selector.closingChannel(id));
+        assertEquals(ChannelState.EXPIRED, channel.state());
+        assertNull(channel.selectionKey().attachment());
+        assertTrue("Disconnect not notified", selector.disconnected().containsKey(id));
+        assertEquals(ChannelState.EXPIRED, selector.disconnected().get(id));
+        verifySelectorEmpty();
     }
 
-    private void verifyCloseOldestConnectionWithStagedReceives(int maxStagedReceives) throws Exception {
-        KafkaChannel channel = createConnectionWithStagedReceives(maxStagedReceives);
+    /**
+     * Verifies that sockets with incoming data available are not expired.
+     * For PLAINTEXT, pending receives are always read from socket without any buffering, so this
+     * test is only verifying that channels are not expired while there is data to read from socket.
+     * For SSL, pending receives may also be in SSL netReadBuffer or appReadBuffer. So the test verifies
+     * that connection is not expired when data is available from buffers or network.
+     */
+    @Test
+    public void testCloseOldestConnectionWithMultiplePendingReceives() throws Exception {
+        int expectedReceives = 5;
+        KafkaChannel channel = createConnectionWithPendingReceives(expectedReceives);
         String id = channel.id();
-        int stagedReceives = selector.numStagedReceives(channel);
         int completedReceives = 0;
         while (selector.disconnected().isEmpty()) {
             time.sleep(6000); // The max idle time is 5000ms
-            selector.poll(0);
+            selector.poll(completedReceives == expectedReceives ? 0 : 1000);
             completedReceives += selector.completedReceives().size();
-            // With SSL, more receives may be staged from buffered data
-            int newStaged = selector.numStagedReceives(channel) - (stagedReceives - completedReceives);
-            if (newStaged > 0) {
-                stagedReceives += newStaged;
-                assertNotNull("Channel should not have been expired", selector.channel(id));
-                assertFalse("Channel should not have been disconnected", selector.disconnected().containsKey(id));
-            } else if (!selector.completedReceives().isEmpty()) {
+            if (!selector.completedReceives().isEmpty()) {
                 assertEquals(1, selector.completedReceives().size());
+                assertNotNull("Channel should not have been expired", selector.channel(id));
                 assertTrue("Channel not found", selector.closingChannel(id) != null || selector.channel(id) != null);
                 assertFalse("Disconnect notified too early", selector.disconnected().containsKey(id));
             }
         }
-        assertEquals(stagedReceives, completedReceives);
+        assertEquals(expectedReceives, completedReceives);
         assertNull("Channel not removed", selector.channel(id));
         assertNull("Channel not removed", selector.closingChannel(id));
         assertTrue("Disconnect not notified", selector.disconnected().containsKey(id));
         assertTrue("Unexpected receive", selector.completedReceives().isEmpty());
     }
 
+    /**
+     * Tests that graceful close of channel processes remaining data from socket read buffers.
+     * Since we cannot determine how much data is available in the buffers, this test verifies that
+     * multiple receives are completed after server shuts down connections, with retries to tolerate
+     * cases where data may not be available in the socket buffer.
+     */
+    @Test
+    public void testGracefulClose() throws Exception {
+        int maxReceiveCountAfterClose = 0;
+        for (int i = 6; i <= 100 && maxReceiveCountAfterClose < 5; i++) {
+            int receiveCount = 0;
+            KafkaChannel channel = createConnectionWithPendingReceives(i);
+            // Poll until one or more receives complete and then close the server-side connection
+            TestUtils.waitForCondition(() -> {
+                selector.poll(1000);
+                return selector.completedReceives().size() > 0;
+            }, 5000, "Receive not completed");
+            server.closeConnections();
+            while (selector.disconnected().isEmpty()) {
+                selector.poll(1);
+                receiveCount += selector.completedReceives().size();
+                assertTrue("Too many completed receives in one poll", selector.completedReceives().size() <= 1);
+            }
+            assertEquals(channel.id(), selector.disconnected().keySet().iterator().next());
+            maxReceiveCountAfterClose = Math.max(maxReceiveCountAfterClose, receiveCount);
+        }
+        assertTrue("Too few receives after close: " + maxReceiveCountAfterClose, maxReceiveCountAfterClose >= 5);
+    }
 
+    /**
+     * Tests that graceful close is not delayed if only part of an incoming receive is
+     * available in the socket buffer.
+     */
+    @Test
+    public void testPartialReceiveGracefulClose() throws Exception {
+        String id = "0";
+        blockingConnect(id);
+        KafkaChannel channel = selector.channel(id);
+        // Inject a NetworkReceive into Kafka channel with a large size
+        injectNetworkReceive(channel, 100000);
+        sendNoReceive(channel, 2); // Send some data that gets received as part of injected receive
+        selector.poll(1000); // Wait until some data arrives, but not a completed receive
+        assertEquals(0, selector.completedReceives().size());
+        server.closeConnections();
+        TestUtils.waitForCondition(() -> {
+            try {
+                selector.poll(100);
+                return !selector.disconnected().isEmpty();
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+        }, 10000, "Channel not disconnected");
+        assertEquals(1, selector.disconnected().size());
+        assertEquals(channel.id(), selector.disconnected().keySet().iterator().next());
+        assertEquals(0, selector.completedReceives().size());
+    }
 
     @Test
     public void testMuteOnOOM() throws Exception {
@@ -442,14 +720,14 @@ public void testMuteOnOOM() throws Exception {
             selector.register("clientX", channelX);
             selector.register("clientY", channelY);
 
-            List completed = Collections.emptyList();
+            Collection completed = Collections.emptyList();
             long deadline = System.currentTimeMillis() + 5000;
             while (System.currentTimeMillis() < deadline && completed.isEmpty()) {
                 selector.poll(1000);
                 completed = selector.completedReceives();
             }
             assertEquals("could not read a single request within timeout", 1, completed.size());
-            NetworkReceive firstReceive = completed.get(0);
+            NetworkReceive firstReceive = completed.iterator().next();
             assertEquals(0, pool.availableMemory());
             assertTrue(selector.isOutOfMemory());
 
@@ -499,28 +777,20 @@ protected byte[] randomPayload(int sizeBytes) throws Exception {
      */
     @Test
     public void testConnectDisconnectDuringInSinglePoll() throws Exception {
-        IMocksControl control = createControl();
-
         // channel is connected, not ready and it throws an exception during prepare
-        KafkaChannel kafkaChannel = control.createMock(KafkaChannel.class);
-        expect(kafkaChannel.id()).andStubReturn("1");
-        expect(kafkaChannel.socketDescription()).andStubReturn("");
-        expect(kafkaChannel.state()).andStubReturn(ChannelState.NOT_CONNECTED);
-        expect(kafkaChannel.finishConnect()).andReturn(true);
-        expect(kafkaChannel.isConnected()).andStubReturn(true);
-        // record void method invocations
-        kafkaChannel.disconnect();
-        kafkaChannel.close();
-        expect(kafkaChannel.ready()).andReturn(false).anyTimes();
-        // prepare throws an exception
-        kafkaChannel.prepare();
-        expectLastCall().andThrow(new IOException());
-
-        SelectionKey selectionKey = control.createMock(SelectionKey.class);
-        expect(selectionKey.channel()).andReturn(SocketChannel.open());
-        expect(selectionKey.readyOps()).andStubReturn(SelectionKey.OP_CONNECT);
-
-        control.replay();
+        KafkaChannel kafkaChannel = mock(KafkaChannel.class);
+        when(kafkaChannel.id()).thenReturn("1");
+        when(kafkaChannel.socketDescription()).thenReturn("");
+        when(kafkaChannel.state()).thenReturn(ChannelState.NOT_CONNECTED);
+        when(kafkaChannel.finishConnect()).thenReturn(true);
+        when(kafkaChannel.isConnected()).thenReturn(true);
+        when(kafkaChannel.ready()).thenReturn(false);
+        doThrow(new IOException()).when(kafkaChannel).prepare();
+
+        SelectionKey selectionKey = mock(SelectionKey.class);
+        when(kafkaChannel.selectionKey()).thenReturn(selectionKey);
+        when(selectionKey.channel()).thenReturn(SocketChannel.open());
+        when(selectionKey.readyOps()).thenReturn(SelectionKey.OP_CONNECT);
 
         selectionKey.attach(kafkaChannel);
         Set selectionKeys = Utils.mkSet(selectionKey);
@@ -528,10 +798,227 @@ public void testConnectDisconnectDuringInSinglePoll() throws Exception {
 
         assertFalse(selector.connected().contains(kafkaChannel.id()));
         assertTrue(selector.disconnected().containsKey(kafkaChannel.id()));
+        assertNull(selectionKey.attachment());
+
+        verify(kafkaChannel, atLeastOnce()).ready();
+        verify(kafkaChannel).disconnect();
+        verify(kafkaChannel).close();
+        verify(selectionKey).cancel();
+    }
+
+    @Test
+    public void testOutboundConnectionsCountInConnectionCreationMetric() throws Exception {
+        // create connections
+        int expectedConnections = 5;
+        InetSocketAddress addr = new InetSocketAddress("localhost", server.port);
+        for (int i = 0; i < expectedConnections; i++)
+            connect(Integer.toString(i), addr);
+
+        // Poll continuously, as we cannot guarantee that the first call will see all connections
+        int seenConnections = 0;
+        for (int i = 0; i < 10; i++) {
+            selector.poll(100L);
+            seenConnections += selector.connected().size();
+            if (seenConnections == expectedConnections)
+                break;
+        }
+
+        assertEquals((double) expectedConnections, getMetric("connection-creation-total").metricValue());
+        assertEquals((double) expectedConnections, getMetric("connection-count").metricValue());
+    }
+
+    @Test
+    public void testInboundConnectionsCountInConnectionCreationMetric() throws Exception {
+        int conns = 5;
+
+        try (ServerSocketChannel ss = ServerSocketChannel.open()) {
+            ss.bind(new InetSocketAddress(0));
+            InetSocketAddress serverAddress = (InetSocketAddress) ss.getLocalAddress();
 
-        control.verify();
+            for (int i = 0; i < conns; i++) {
+                Thread sender = createSender(serverAddress, randomPayload(1));
+                sender.start();
+                SocketChannel channel = ss.accept();
+                channel.configureBlocking(false);
+
+                selector.register(Integer.toString(i), channel);
+            }
+        }
+
+        assertEquals((double) conns, getMetric("connection-creation-total").metricValue());
+        assertEquals((double) conns, getMetric("connection-count").metricValue());
+    }
+
+    @Test
+    public void testConnectionsByClientMetric() throws Exception {
+        String node = "0";
+        Map unknownNameAndVersion = softwareNameAndVersionTags(
+            ClientInformation.UNKNOWN_NAME_OR_VERSION, ClientInformation.UNKNOWN_NAME_OR_VERSION);
+        Map knownNameAndVersion = softwareNameAndVersionTags("A", "B");
+
+        try (ServerSocketChannel ss = ServerSocketChannel.open()) {
+            ss.bind(new InetSocketAddress(0));
+            InetSocketAddress serverAddress = (InetSocketAddress) ss.getLocalAddress();
+
+            Thread sender = createSender(serverAddress, randomPayload(1));
+            sender.start();
+            SocketChannel channel = ss.accept();
+            channel.configureBlocking(false);
+
+            // Metric with unknown / unknown should be there
+            selector.register(node, channel);
+            assertEquals(1,
+                getMetric("connections", unknownNameAndVersion).metricValue());
+            assertEquals(ClientInformation.EMPTY,
+                selector.channel(node).channelMetadataRegistry().clientInformation());
+
+            // Metric with unknown / unknown should not be there, metric with A / B should be there
+            ClientInformation clientInformation = new ClientInformation("A", "B");
+            selector.channel(node).channelMetadataRegistry()
+                .registerClientInformation(clientInformation);
+            assertEquals(clientInformation,
+                selector.channel(node).channelMetadataRegistry().clientInformation());
+            assertEquals(0, getMetric("connections", unknownNameAndVersion).metricValue());
+            assertEquals(1, getMetric("connections", knownNameAndVersion).metricValue());
+
+            // Metric with A / B should not be there,
+            selector.close(node);
+            assertEquals(0, getMetric("connections", knownNameAndVersion).metricValue());
+        }
+    }
+
+    private Map softwareNameAndVersionTags(String clientSoftwareName, String clientSoftwareVersion) {
+        Map tags = new HashMap<>(2);
+        tags.put("clientSoftwareName", clientSoftwareName);
+        tags.put("clientSoftwareVersion", clientSoftwareVersion);
+        return tags;
+    }
+
+    private KafkaMetric getMetric(String name, Map tags) throws Exception {
+        Optional> metric = metrics.metrics().entrySet().stream()
+            .filter(entry ->
+                entry.getKey().name().equals(name) && entry.getKey().tags().equals(tags))
+            .findFirst();
+        if (!metric.isPresent())
+            throw new Exception(String.format("Could not find metric called %s with tags %s", name, tags.toString()));
+
+        return metric.get().getValue();
+    }
+
+    @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());
+    }
+
+    @Test
+    public void testWriteCompletesSendWithNoBytesWritten() throws IOException {
+        KafkaChannel channel = mock(KafkaChannel.class);
+        when(channel.id()).thenReturn("1");
+        when(channel.write()).thenReturn(0L);
+        NetworkSend send = new NetworkSend("destination", new ByteBufferSend(ByteBuffer.allocate(0)));
+        when(channel.maybeCompleteSend()).thenReturn(send);
+        selector.write(channel);
+        assertEquals(asList(send), selector.completedSends());
+    }
+
+    /**
+     * Ensure that no errors are thrown if channels are closed while processing multiple completed receives
+     */
+    @Test
+    public void testChannelCloseWhileProcessingReceives() throws Exception {
+        int numChannels = 4;
+        Map channels = TestUtils.fieldValue(selector, Selector.class, "channels");
+        Set selectionKeys = new HashSet<>();
+        for (int i = 0; i < numChannels; i++) {
+            String id = String.valueOf(i);
+            KafkaChannel channel = mock(KafkaChannel.class);
+            channels.put(id, channel);
+            when(channel.id()).thenReturn(id);
+            when(channel.state()).thenReturn(ChannelState.READY);
+            when(channel.isConnected()).thenReturn(true);
+            when(channel.ready()).thenReturn(true);
+            when(channel.read()).thenReturn(1L);
+
+            SelectionKey selectionKey = mock(SelectionKey.class);
+            when(channel.selectionKey()).thenReturn(selectionKey);
+            when(selectionKey.isValid()).thenReturn(true);
+            when(selectionKey.readyOps()).thenReturn(SelectionKey.OP_READ);
+            selectionKey.attach(channel);
+            selectionKeys.add(selectionKey);
+
+            NetworkReceive receive = mock(NetworkReceive.class);
+            when(receive.source()).thenReturn(id);
+            when(receive.size()).thenReturn(10);
+            when(receive.bytesRead()).thenReturn(1);
+            when(receive.payload()).thenReturn(ByteBuffer.allocate(10));
+            when(channel.maybeCompleteReceive()).thenReturn(receive);
+        }
+
+        selector.pollSelectionKeys(selectionKeys, false, System.nanoTime());
+        assertEquals(numChannels, selector.completedReceives().size());
+        Set closed = new HashSet<>();
+        Set notClosed = new HashSet<>();
+        for (NetworkReceive receive : selector.completedReceives()) {
+            KafkaChannel channel = selector.channel(receive.source());
+            assertNotNull(channel);
+            if (closed.size() < 2) {
+                selector.close(channel.id());
+                closed.add(channel);
+            } else
+                notClosed.add(channel);
+        }
+        assertEquals(notClosed, new HashSet<>(selector.channels()));
+        closed.forEach(channel -> assertNull(selector.channel(channel.id())));
+
+        selector.poll(0);
+        assertEquals(0, selector.completedReceives().size());
+    }
+
+
     private String blockingRequest(String node, String s) throws IOException {
         selector.send(createSend(node, s));
         selector.poll(1000L);
@@ -551,6 +1038,7 @@ protected void connect(String node, InetSocketAddress serverAddr) throws IOExcep
     private void blockingConnect(String node) throws IOException {
         blockingConnect(node, new InetSocketAddress("localhost", server.port));
     }
+
     protected void blockingConnect(String node, InetSocketAddress serverAddr) throws IOException {
         selector.connect(node, serverAddr, BUFFER_SIZE, BUFFER_SIZE);
         while (!selector.connected().contains(node))
@@ -559,8 +1047,8 @@ protected void blockingConnect(String node, InetSocketAddress serverAddr) throws
             selector.poll(10000L);
     }
 
-    protected NetworkSend createSend(String node, String s) {
-        return new NetworkSend(node, ByteBuffer.wrap(s.getBytes()));
+    protected final NetworkSend createSend(String node, String payload) {
+        return new NetworkSend(node, ByteBufferSend.sizePrefixed(ByteBuffer.wrap(payload.getBytes())));
     }
 
     protected String asString(NetworkReceive receive) {
@@ -576,7 +1064,6 @@ private void sendAndReceive(String node, String requestPrefix, int startIndex, i
             // do the i/o
             selector.poll(0L);
             assertEquals("No disconnects should have occurred.", 0, selector.disconnected().size());
-
             // handle requests and responses of the fast node
             for (NetworkReceive receive : selector.completedReceives()) {
                 assertEquals(requestPrefix + "-" + responses, asString(receive));
@@ -589,21 +1076,104 @@ private void sendAndReceive(String node, String requestPrefix, int startIndex, i
         }
     }
 
-    private void verifySelectorEmpty() throws Exception {
-        for (KafkaChannel channel : selector.channels())
+    private void verifyNonEmptyImmediatelyConnectedKeys(Selector selector) throws Exception {
+        Field field = Selector.class.getDeclaredField("immediatelyConnectedKeys");
+        field.setAccessible(true);
+        Collection immediatelyConnectedKeys = (Collection) field.get(selector);
+        assertFalse(immediatelyConnectedKeys.isEmpty());
+    }
+
+    private void verifyEmptyImmediatelyConnectedKeys(Selector selector) throws Exception {
+        Field field = Selector.class.getDeclaredField("immediatelyConnectedKeys");
+        ensureEmptySelectorField(selector, field);
+    }
+
+    protected void verifySelectorEmpty() throws Exception {
+        verifySelectorEmpty(this.selector);
+    }
+
+    public void verifySelectorEmpty(Selector selector) throws Exception {
+        for (KafkaChannel channel : selector.channels()) {
             selector.close(channel.id());
+            assertNull(channel.selectionKey().attachment());
+        }
         selector.poll(0);
         selector.poll(0); // Poll a second time to clear everything
-        for (Field field : selector.getClass().getDeclaredFields()) {
-            field.setAccessible(true);
-            Object obj = field.get(selector);
-            if (obj instanceof Set)
-                assertTrue("Field not empty: " + field + " " + obj, ((Set) obj).isEmpty());
-            else if (obj instanceof Map)
-                assertTrue("Field not empty: " + field + " " + obj, ((Map) obj).isEmpty());
-            else if (obj instanceof List)
-                assertTrue("Field not empty: " + field + " " + obj, ((List) obj).isEmpty());
+        ensureEmptySelectorFields(selector);
+    }
+
+    private void ensureEmptySelectorFields(Selector selector) throws Exception {
+        for (Field field : Selector.class.getDeclaredFields()) {
+            ensureEmptySelectorField(selector, field);
+        }
+    }
+
+    private void ensureEmptySelectorField(Selector selector, Field field) throws Exception {
+        field.setAccessible(true);
+        Object obj = field.get(selector);
+        if (obj instanceof Collection)
+            assertTrue("Field not empty: " + field + " " + obj, ((Collection) obj).isEmpty());
+        else if (obj instanceof Map)
+            assertTrue("Field not empty: " + field + " " + obj, ((Map) obj).isEmpty());
+    }
+
+    private KafkaMetric getMetric(String name) throws Exception {
+        Optional> metric = metrics.metrics().entrySet().stream()
+                .filter(entry -> entry.getKey().name().equals(name))
+                .findFirst();
+        if (!metric.isPresent())
+            throw new Exception(String.format("Could not find metric called %s", name));
+
+        return metric.get().getValue();
+    }
+
+    private KafkaMetric findUntaggedMetricByName(String name) {
+        MetricName metricName = new MetricName(name, METRIC_GROUP + "-metrics", "", new HashMap<>());
+        KafkaMetric metric = metrics.metrics().get(metricName);
+        assertNotNull(metric);
+        return metric;
+    }
+
+    /**
+     * Creates a connection, sends the specified number of requests and returns without reading
+     * any incoming data. Some of the incoming data may be in the socket buffers when this method
+     * returns, but there is no guarantee that all the data from the server will be available
+     * immediately.
+     */
+    private KafkaChannel createConnectionWithPendingReceives(int pendingReceives) throws Exception {
+        String id = "0";
+        blockingConnect(id);
+        KafkaChannel channel = selector.channel(id);
+        sendNoReceive(channel, pendingReceives);
+        return channel;
+    }
+
+    /**
+     * Sends the specified number of requests and waits for the requests to be sent. The channel
+     * is muted during polling to ensure that incoming data is not received.
+     */
+    private KafkaChannel sendNoReceive(KafkaChannel channel, int numRequests) throws Exception {
+        channel.mute();
+        for (int i = 0; i < numRequests; i++) {
+            selector.send(createSend(channel.id(), String.valueOf(i)));
+            do {
+                selector.poll(10);
+            } while (selector.completedSends().isEmpty());
         }
+        channel.maybeUnmute();
+
+        return channel;
     }
 
+    /**
+     * Injects a NetworkReceive for channel with size buffer filled in with the provided size
+     * and a payload buffer allocated with that size, but no data in the payload buffer.
+     */
+    private void injectNetworkReceive(KafkaChannel channel, int size) throws Exception {
+        NetworkReceive receive = new NetworkReceive();
+        TestUtils.setFieldValue(channel, "receive", receive);
+        ByteBuffer sizeBuffer = TestUtils.fieldValue(receive, NetworkReceive.class, "size");
+        sizeBuffer.putInt(size);
+        TestUtils.setFieldValue(receive, "buffer", ByteBuffer.allocate(size));
+    }
 }
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 bf2e77c6ed921..8a36fbfa594c5 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
@@ -16,13 +16,23 @@
  */
 package org.apache.kafka.common.network;
 
+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.Metrics;
 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.TestSslUtils;
+import org.apache.kafka.test.TestUtils;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
@@ -32,11 +42,16 @@
 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.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Consumer;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
@@ -59,10 +74,11 @@ public void setUp() throws Exception {
         this.server.start();
         this.time = new MockTime();
         sslClientConfigs = TestSslUtils.createSslConfig(false, false, Mode.CLIENT, trustStoreFile, "client");
-        this.channelBuilder = new SslChannelBuilder(Mode.CLIENT);
+        LogContext logContext = new LogContext();
+        this.channelBuilder = new SslChannelBuilder(Mode.CLIENT, null, false, logContext);
         this.channelBuilder.configure(sslClientConfigs);
         this.metrics = new Metrics();
-        this.selector = new Selector(5000, metrics, time, "MetricGroup", channelBuilder, new LogContext());
+        this.selector = new Selector(5000, metrics, time, "MetricGroup", channelBuilder, logContext);
     }
 
     @After
@@ -77,6 +93,162 @@ public SecurityProtocol securityProtocol() {
         return SecurityProtocol.PLAINTEXT;
     }
 
+    @Override
+    protected Map clientConfigs() {
+        return sslClientConfigs;
+    }
+
+    @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,
+                TestSslUtils.DEFAULT_TLS_PROTOCOL_FOR_TESTS
+        );
+        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);
+
+        TestUtils.waitForCondition(() -> cipherMetrics(metrics).size() == 1,
+            "Waiting for cipher metrics to be created.");
+        assertEquals(Integer.valueOf(1), cipherMetrics(metrics).get(0).metricValue());
+        assertNotNull(selector.channel(node).channelMetadataRegistry().cipherInformation());
+
+        selector.close(node);
+        super.verifySelectorEmpty(selector);
+
+        assertEquals(1, cipherMetrics(metrics).size());
+        assertEquals(Integer.valueOf(0), cipherMetrics(metrics).get(0).metricValue());
+
+        Security.removeProvider(testProviderCreator.getProvider().getName());
+        selector.close();
+        server.close();
+        metrics.close();
+    }
+
+    @Test
+    public void testDisconnectWithIntermediateBufferedBytes() throws Exception {
+        int requestSize = 100 * 1024;
+        final String node = "0";
+        String request = TestUtils.randomString(requestSize);
+
+        this.selector.close();
+
+        this.channelBuilder = new TestSslChannelBuilder(Mode.CLIENT);
+        this.channelBuilder.configure(sslClientConfigs);
+        this.selector = new Selector(5000, metrics, time, "MetricGroup", channelBuilder, new LogContext());
+        connect(node, new InetSocketAddress("localhost", server.port));
+        selector.send(createSend(node, request));
+
+        waitForBytesBuffered(selector, node);
+
+        selector.close(node);
+        verifySelectorEmpty();
+    }
+
+    private void waitForBytesBuffered(Selector selector, String node) throws Exception {
+        TestUtils.waitForCondition(() -> {
+            try {
+                selector.poll(0L);
+                return selector.channel(node).hasBytesBuffered();
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+        }, 2000L, "Failed to reach socket state with bytes buffered");
+    }
+
+    @Test
+    public void testBytesBufferedChannelWithNoIncomingBytes() throws Exception {
+        verifyNoUnnecessaryPollWithBytesBuffered(key ->
+            key.interestOps(key.interestOps() & ~SelectionKey.OP_READ));
+    }
+
+    @Test
+    public void testBytesBufferedChannelAfterMute() throws Exception {
+        verifyNoUnnecessaryPollWithBytesBuffered(key -> ((KafkaChannel) key.attachment()).mute());
+    }
+
+    private void verifyNoUnnecessaryPollWithBytesBuffered(Consumer disableRead)
+            throws Exception {
+        this.selector.close();
+
+        String node1 = "1";
+        String node2 = "2";
+        final AtomicInteger node1Polls = new AtomicInteger();
+
+        this.channelBuilder = new TestSslChannelBuilder(Mode.CLIENT);
+        this.channelBuilder.configure(sslClientConfigs);
+        this.selector = new Selector(5000, metrics, time, "MetricGroup", channelBuilder, new LogContext()) {
+            @Override
+            void pollSelectionKeys(Set selectionKeys, boolean isImmediatelyConnected, long currentTimeNanos) {
+                for (SelectionKey key : selectionKeys) {
+                    KafkaChannel channel = (KafkaChannel) key.attachment();
+                    if (channel != null && channel.id().equals(node1))
+                        node1Polls.incrementAndGet();
+                }
+                super.pollSelectionKeys(selectionKeys, isImmediatelyConnected, currentTimeNanos);
+            }
+        };
+
+        // Get node1 into bytes buffered state and then disable read on the socket.
+        // Truncate the read buffers to ensure that there is buffered data, but not enough to make progress.
+        int largeRequestSize = 100 * 1024;
+        connect(node1, new InetSocketAddress("localhost", server.port));
+        selector.send(createSend(node1,  TestUtils.randomString(largeRequestSize)));
+        waitForBytesBuffered(selector, node1);
+        TestSslChannelBuilder.TestSslTransportLayer.transportLayers.get(node1).truncateReadBuffer();
+        disableRead.accept(selector.channel(node1).selectionKey());
+
+        // Clear poll count and count the polls from now on
+        node1Polls.set(0);
+
+        // Process sends and receives on node2. Test verifies that we don't process node1
+        // unnecessarily on each of these polls.
+        connect(node2, new InetSocketAddress("localhost", server.port));
+        int received = 0;
+        String request = TestUtils.randomString(10);
+        selector.send(createSend(node2, request));
+        while (received < 100) {
+            received += selector.completedReceives().size();
+            if (!selector.completedSends().isEmpty()) {
+                selector.send(createSend(node2, request));
+            }
+            selector.poll(5);
+        }
+
+        // Verify that pollSelectionKeys was invoked once to process buffered data
+        // but not again since there isn't sufficient data to process.
+        assertEquals(1, node1Polls.get());
+        selector.close(node1);
+        selector.close(node2);
+        verifySelectorEmpty();
+    }
+
     /**
      * Renegotiation is not supported since it is potentially unsafe and it has been removed in TLS 1.3
      */
@@ -112,11 +284,15 @@ public void testMuteOnOOM() throws Exception {
         selector.close();
         MemoryPool pool = new SimpleMemoryPool(900, 900, false, null);
         //the initial channel builder is for clients, we need a server one
+        String tlsProtocol = "TLSv1.2";
         File trustStoreFile = File.createTempFile("truststore", ".jks");
-        Map sslServerConfigs = TestSslUtils.createSslConfig(false, true, Mode.SERVER, trustStoreFile, "server");
-        channelBuilder = new SslChannelBuilder(Mode.SERVER);
+        Map sslServerConfigs = new TestSslUtils.SslConfigsBuilder(Mode.SERVER)
+                .tlsProtocol(tlsProtocol)
+                .createNewTrustStore(trustStoreFile)
+                .build();
+        channelBuilder = new SslChannelBuilder(Mode.SERVER, null, false, new LogContext());
         channelBuilder.configure(sslServerConfigs);
-        selector = new Selector(NetworkReceive.UNLIMITED, 5000, metrics, time, "MetricGroup", 
+        selector = new Selector(NetworkReceive.UNLIMITED, 5000, metrics, time, "MetricGroup",
                 new HashMap(), true, false, channelBuilder, pool, new LogContext());
 
         try (ServerSocketChannel ss = ServerSocketChannel.open()) {
@@ -124,8 +300,8 @@ public void testMuteOnOOM() throws Exception {
 
             InetSocketAddress serverAddress = (InetSocketAddress) ss.getLocalAddress();
 
-            SslSender sender1 = createSender(serverAddress, randomPayload(900));
-            SslSender sender2 = createSender(serverAddress, randomPayload(900));
+            SslSender sender1 = createSender(tlsProtocol, serverAddress, randomPayload(900));
+            SslSender sender2 = createSender(tlsProtocol, serverAddress, randomPayload(900));
             sender1.start();
             sender2.start();
 
@@ -145,11 +321,11 @@ public void testMuteOnOOM() throws Exception {
             while (System.currentTimeMillis() < deadline) {
                 selector.poll(10);
 
-                List completed = selector.completedReceives();
+                Collection completed = selector.completedReceives();
                 if (firstReceive == null) {
                     if (!completed.isEmpty()) {
                         assertEquals("expecting a single request", 1, completed.size());
-                        firstReceive = completed.get(0);
+                        firstReceive = completed.iterator().next();
                         assertTrue(selector.isMadeReadProgressLastPoll());
                         assertEquals(0, pool.availableMemory());
                     }
@@ -173,7 +349,7 @@ public void testMuteOnOOM() throws Exception {
             firstReceive.close();
             assertEquals(900, pool.availableMemory()); //memory has been released back to pool
 
-            List completed = Collections.emptyList();
+            Collection completed = Collections.emptyList();
             deadline = System.currentTimeMillis() + 5000;
             while (System.currentTimeMillis() < deadline && completed.isEmpty()) {
                 selector.poll(1000);
@@ -194,7 +370,57 @@ protected void connect(String node, InetSocketAddress serverAddr) throws IOExcep
         blockingConnect(node, serverAddr);
     }
 
-    private SslSender createSender(InetSocketAddress serverAddress, byte[] payload) {
-        return new SslSender(serverAddress, payload);
+    private SslSender createSender(String tlsProtocol, InetSocketAddress serverAddress, byte[] payload) {
+        return new SslSender(tlsProtocol, serverAddress, payload);
+    }
+
+    private static class TestSslChannelBuilder extends SslChannelBuilder {
+
+        public TestSslChannelBuilder(Mode mode) {
+            super(mode, null, false, new LogContext());
+        }
+
+        @Override
+        protected SslTransportLayer buildTransportLayer(SslFactory sslFactory, String id, SelectionKey key,
+                                                        String host, ChannelMetadataRegistry metadataRegistry) throws IOException {
+            SocketChannel socketChannel = (SocketChannel) key.channel();
+            SSLEngine sslEngine = sslFactory.createSslEngine(host, socketChannel.socket().getPort());
+            TestSslTransportLayer transportLayer = new TestSslTransportLayer(id, key, sslEngine, metadataRegistry);
+            return transportLayer;
+        }
+
+        /*
+         * TestSslTransportLayer will read from socket once every two tries. This increases
+         * the chance that there will be bytes buffered in the transport layer after read().
+         */
+        static class TestSslTransportLayer extends SslTransportLayer {
+            static Map transportLayers = new HashMap<>();
+            boolean muteSocket = false;
+
+            public TestSslTransportLayer(String channelId, SelectionKey key, SSLEngine sslEngine,
+                                         ChannelMetadataRegistry metadataRegistry) throws IOException {
+                super(channelId, key, sslEngine, metadataRegistry);
+                transportLayers.put(channelId, this);
+            }
+
+            @Override
+            protected int readFromSocketChannel() throws IOException {
+                if (muteSocket) {
+                    if ((selectionKey().interestOps() & SelectionKey.OP_READ) != 0)
+                        muteSocket = false;
+                    return 0;
+                }
+                muteSocket = true;
+                return super.readFromSocketChannel();
+            }
+
+            // Leave one byte in network read buffer so that some buffered bytes are present,
+            // but not enough to make progress on a read.
+            void truncateReadBuffer() throws Exception {
+                netReadBuffer().position(1);
+                appReadBuffer().position(0);
+                muteSocket = true;
+            }
+        }
     }
 }
diff --git a/clients/src/test/java/org/apache/kafka/common/network/SslSender.java b/clients/src/test/java/org/apache/kafka/common/network/SslSender.java
index cae69cb675d3c..22196dde13a54 100644
--- a/clients/src/test/java/org/apache/kafka/common/network/SslSender.java
+++ b/clients/src/test/java/org/apache/kafka/common/network/SslSender.java
@@ -29,11 +29,13 @@
 
 public class SslSender extends Thread {
 
+    private final String tlsProtocol;
     private final InetSocketAddress serverAddress;
     private final byte[] payload;
     private final CountDownLatch handshaked = new CountDownLatch(1);
 
-    public SslSender(InetSocketAddress serverAddress, byte[] payload) {
+    public SslSender(String tlsProtocol, InetSocketAddress serverAddress, byte[] payload) {
+        this.tlsProtocol = tlsProtocol;
         this.serverAddress = serverAddress;
         this.payload = payload;
         setDaemon(true);
@@ -43,7 +45,7 @@ public SslSender(InetSocketAddress serverAddress, byte[] payload) {
     @Override
     public void run() {
         try {
-            SSLContext sc = SSLContext.getInstance("TLSv1.2");
+            SSLContext sc = SSLContext.getInstance(tlsProtocol);
             sc.init(null, new TrustManager[]{new NaiveTrustManager()}, new java.security.SecureRandom());
             try (SSLSocket connection = (SSLSocket) sc.getSocketFactory().createSocket(serverAddress.getAddress(), serverAddress.getPort())) {
                 OutputStream os = connection.getOutputStream();
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 7cc2808134ff6..bb00edfc0fd87 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
@@ -20,24 +20,27 @@
 import org.apache.kafka.common.config.SslConfigs;
 import org.apache.kafka.common.config.internals.BrokerSecurityConfigs;
 import org.apache.kafka.common.config.types.Password;
+import org.apache.kafka.common.errors.InvalidConfigurationException;
 import org.apache.kafka.common.memory.MemoryPool;
 import org.apache.kafka.common.metrics.Metrics;
-import org.apache.kafka.common.security.auth.SecurityProtocol;
 import org.apache.kafka.common.security.TestSecurityConfig;
+import org.apache.kafka.common.security.auth.SecurityProtocol;
+import org.apache.kafka.common.security.ssl.DefaultSslEngineFactory;
 import org.apache.kafka.common.security.ssl.SslFactory;
+import org.apache.kafka.common.utils.Java;
 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.common.utils.Utils;
+import org.apache.kafka.test.TestSslUtils;
 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 javax.net.ssl.SSLContext;
-import javax.net.ssl.SSLEngine;
-import javax.net.ssl.SSLParameters;
 import java.io.ByteArrayOutputStream;
+import java.io.File;
 import java.io.IOException;
 import java.net.InetAddress;
 import java.net.InetSocketAddress;
@@ -45,42 +48,84 @@
 import java.nio.channels.Channels;
 import java.nio.channels.SelectionKey;
 import java.nio.channels.SocketChannel;
+import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
 
+import javax.net.ssl.SSLContext;
+import javax.net.ssl.SSLEngine;
+import javax.net.ssl.SSLParameters;
+
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
+import static org.junit.Assume.assumeTrue;
 
 /**
  * Tests for the SSL transport layer. These use a test harness that runs a simple socket server that echos back responses.
  */
+@RunWith(value = Parameterized.class)
 public class SslTransportLayerTest {
 
     private static final int BUFFER_SIZE = 4 * 1024;
+    private static Time time = Time.SYSTEM;
 
+    private final String tlsProtocol;
+    private final boolean useInlinePem;
     private NioEchoServer server;
     private Selector selector;
-    private ChannelBuilder channelBuilder;
     private CertStores serverCertStores;
     private CertStores clientCertStores;
     private Map sslClientConfigs;
     private Map sslServerConfigs;
+    private Map sslConfigOverrides;
+
+    @Parameterized.Parameters(name = "tlsProtocol={0}, useInlinePem={1}")
+    public static Collection data() {
+        List values = new ArrayList<>();
+        values.add(new Object[] {"TLSv1.2", false});
+        values.add(new Object[] {"TLSv1.2", true});
+        if (Java.IS_JAVA11_COMPATIBLE) {
+            values.add(new Object[] {"TLSv1.3", false});
+        }
+        return values;
+    }
+
+    public SslTransportLayerTest(String tlsProtocol, boolean useInlinePem) {
+        this.tlsProtocol = tlsProtocol;
+        this.useInlinePem = useInlinePem;
+        sslConfigOverrides = new HashMap<>();
+        sslConfigOverrides.put(SslConfigs.SSL_PROTOCOL_CONFIG, tlsProtocol);
+        sslConfigOverrides.put(SslConfigs.SSL_ENABLED_PROTOCOLS_CONFIG, Collections.singletonList(tlsProtocol));
+    }
+    private CertStores.Builder certBuilder(boolean isServer, String cn) {
+        return new CertStores.Builder(isServer)
+            .cn(cn)
+            .usePem(useInlinePem);
+    }
 
     @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");
-        sslServerConfigs = serverCertStores.getTrustingConfig(clientCertStores);
-        sslClientConfigs = clientCertStores.getTrustingConfig(serverCertStores);
-        this.channelBuilder = new SslChannelBuilder(Mode.CLIENT);
-        this.channelBuilder.configure(sslClientConfigs);
-        this.selector = new Selector(5000, new Metrics(), new MockTime(), "MetricGroup", channelBuilder, new LogContext());
+        serverCertStores = certBuilder(true, "server").addHostName("localhost").build();
+        clientCertStores = certBuilder(false, "client").addHostName("localhost").build();
+        sslServerConfigs = getTrustingConfig(serverCertStores, clientCertStores);
+        sslClientConfigs = getTrustingConfig(clientCertStores, serverCertStores);
+        sslServerConfigs.put(SslConfigs.SSL_ENGINE_FACTORY_CLASS_CONFIG, DefaultSslEngineFactory.class);
+        sslClientConfigs.put(SslConfigs.SSL_ENGINE_FACTORY_CLASS_CONFIG, DefaultSslEngineFactory.class);
+
+        LogContext logContext = new LogContext();
+        ChannelBuilder channelBuilder = new SslChannelBuilder(Mode.CLIENT, null, false, logContext);
+        channelBuilder.configure(sslClientConfigs);
+        this.selector = new Selector(5000, new Metrics(), time, "MetricGroup", channelBuilder, logContext);
     }
 
     @After
@@ -115,10 +160,10 @@ 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"));
-        sslServerConfigs = serverCertStores.getTrustingConfig(clientCertStores);
-        sslClientConfigs = clientCertStores.getTrustingConfig(serverCertStores);
+        serverCertStores = certBuilder(true, "server").hostAddress(InetAddress.getByName("127.0.0.1")).build();
+        clientCertStores = certBuilder(false, "client").hostAddress(InetAddress.getByName("127.0.0.1")).build();
+        sslServerConfigs = getTrustingConfig(serverCertStores, clientCertStores);
+        sslClientConfigs = getTrustingConfig(clientCertStores, serverCertStores);
         server = createEchoServer(SecurityProtocol.SSL);
         sslClientConfigs.put(SslConfigs.SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG, "HTTPS");
         createSelector(sslClientConfigs);
@@ -134,18 +179,12 @@ public void testValidEndpointIdentificationSanIp() throws Exception {
      */
     @Test
     public void testValidEndpointIdentificationCN() throws Exception {
-        String node = "0";
-        serverCertStores = new CertStores(true, "localhost");
-        clientCertStores = new CertStores(false, "localhost");
-        sslServerConfigs = serverCertStores.getTrustingConfig(clientCertStores);
-        sslClientConfigs = clientCertStores.getTrustingConfig(serverCertStores);
-        server = createEchoServer(SecurityProtocol.SSL);
+        serverCertStores = certBuilder(true, "localhost").build();
+        clientCertStores = certBuilder(false, "localhost").build();
+        sslServerConfigs = getTrustingConfig(serverCertStores, clientCertStores);
+        sslClientConfigs = getTrustingConfig(clientCertStores, serverCertStores);
         sslClientConfigs.put(SslConfigs.SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG, "HTTPS");
-        createSelector(sslClientConfigs);
-        InetSocketAddress addr = new InetSocketAddress("localhost", server.port());
-        selector.connect(node, addr, BUFFER_SIZE, BUFFER_SIZE);
-
-        NetworkTestUtils.checkClientConnection(selector, node, 100, 10);
+        verifySslConfigs();
     }
 
     /**
@@ -185,10 +224,10 @@ 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");
-        sslServerConfigs = serverCertStores.getTrustingConfig(clientCertStores);
-        sslClientConfigs = clientCertStores.getTrustingConfig(serverCertStores);
+        clientCertStores = certBuilder(false, "non-existent.com").build();
+        serverCertStores = certBuilder(true, "localhost").build();
+        sslServerConfigs = getTrustingConfig(serverCertStores, clientCertStores);
+        sslClientConfigs = getTrustingConfig(clientCertStores, serverCertStores);
 
         // Create a server with endpoint validation enabled on the server SSL engine
         SslChannelBuilder serverChannelBuilder = new TestSslChannelBuilder(Mode.SERVER) {
@@ -202,7 +241,7 @@ protected TestSslTransportLayer newTransportLayer(String id, SelectionKey key, S
         };
         serverChannelBuilder.configure(sslServerConfigs);
         server = new NioEchoServer(ListenerName.forSecurityProtocol(SecurityProtocol.SSL), SecurityProtocol.SSL,
-                new TestSecurityConfig(sslServerConfigs), "localhost", serverChannelBuilder, null);
+                new TestSecurityConfig(sslServerConfigs), "localhost", serverChannelBuilder, null, time);
         server.start();
 
         createSelector(sslClientConfigs);
@@ -219,39 +258,49 @@ 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");
-        sslServerConfigs = serverCertStores.getTrustingConfig(clientCertStores);
-        sslClientConfigs = clientCertStores.getTrustingConfig(serverCertStores);
+        serverCertStores = certBuilder(true, "server").addHostName("notahost").build();
+        clientCertStores = certBuilder(false, "client").addHostName("localhost").build();
+        sslServerConfigs = getTrustingConfig(serverCertStores, clientCertStores);
+        sslClientConfigs = getTrustingConfig(clientCertStores, serverCertStores);
         sslClientConfigs.put(SslConfigs.SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG, "HTTPS");
-        server = createEchoServer(SecurityProtocol.SSL);
-        createSelector(sslClientConfigs);
-        InetSocketAddress addr = new InetSocketAddress("localhost", server.port());
-        selector.connect(node, addr, BUFFER_SIZE, BUFFER_SIZE);
-
-        NetworkTestUtils.waitForChannelClose(selector, node, ChannelState.State.AUTHENTICATION_FAILED);
-        server.verifyAuthenticationMetrics(0, 1);
+        verifySslConfigsWithHandshakeFailure();
     }
 
     /**
-     * Tests that server certificate with invalid IP address is accepted by
+     * Tests that server certificate with invalid host name is accepted by
      * a client that has disabled endpoint validation
      */
     @Test
     public void testEndpointIdentificationDisabled() throws Exception {
-        String node = "0";
-        String serverHost = InetAddress.getLocalHost().getHostAddress();
-        SecurityProtocol securityProtocol = SecurityProtocol.SSL;
-        server = new NioEchoServer(ListenerName.forSecurityProtocol(securityProtocol), securityProtocol,
-                new TestSecurityConfig(sslServerConfigs), serverHost, null, null);
-        server.start();
-        sslClientConfigs.remove(SslConfigs.SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG);
+        serverCertStores = certBuilder(true, "server").addHostName("notahost").build();
+        clientCertStores = certBuilder(false, "client").addHostName("localhost").build();
+        sslServerConfigs = getTrustingConfig(serverCertStores, clientCertStores);
+        sslClientConfigs = getTrustingConfig(clientCertStores, serverCertStores);
+
+        server = createEchoServer(SecurityProtocol.SSL);
+        InetSocketAddress addr = new InetSocketAddress("localhost", server.port());
+
+        // Disable endpoint validation, connection should succeed
+        String node = "1";
+        sslClientConfigs.put(SslConfigs.SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG, "");
         createSelector(sslClientConfigs);
-        InetSocketAddress addr = new InetSocketAddress(serverHost, server.port());
         selector.connect(node, addr, BUFFER_SIZE, BUFFER_SIZE);
-
         NetworkTestUtils.checkClientConnection(selector, node, 100, 10);
+
+        // Disable endpoint validation using null value, connection should succeed
+        String node2 = "2";
+        sslClientConfigs.put(SslConfigs.SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG, null);
+        createSelector(sslClientConfigs);
+        selector.connect(node2, addr, BUFFER_SIZE, BUFFER_SIZE);
+        NetworkTestUtils.checkClientConnection(selector, node2, 100, 10);
+
+        // Connection should fail with endpoint validation enabled
+        String node3 = "3";
+        sslClientConfigs.put(SslConfigs.SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG, "HTTPS");
+        createSelector(sslClientConfigs);
+        selector.connect(node3, addr, BUFFER_SIZE, BUFFER_SIZE);
+        NetworkTestUtils.waitForChannelClose(selector, node3, ChannelState.State.AUTHENTICATION_FAILED);
+        selector.close();
     }
 
     /**
@@ -260,14 +309,8 @@ public void testEndpointIdentificationDisabled() throws Exception {
      */
     @Test
     public void testClientAuthenticationRequiredValidProvided() throws Exception {
-        String node = "0";
         sslServerConfigs.put(BrokerSecurityConfigs.SSL_CLIENT_AUTH_CONFIG, "required");
-        server = createEchoServer(SecurityProtocol.SSL);
-        createSelector(sslClientConfigs);
-        InetSocketAddress addr = new InetSocketAddress("localhost", server.port());
-        selector.connect(node, addr, BUFFER_SIZE, BUFFER_SIZE);
-
-        NetworkTestUtils.checkClientConnection(selector, node, 100, 10);
+        verifySslConfigs();
     }
 
     /**
@@ -291,9 +334,7 @@ public void testListenerConfigOverride() throws Exception {
         selector.close();
 
         // Remove client auth, so connection should fail
-        sslClientConfigs.remove(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG);
-        sslClientConfigs.remove(SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG);
-        sslClientConfigs.remove(SslConfigs.SSL_KEY_PASSWORD_CONFIG);
+        CertStores.KEYSTORE_PROPS.forEach(sslClientConfigs::remove);
         createSelector(sslClientConfigs);
         selector.connect(node, addr, BUFFER_SIZE, BUFFER_SIZE);
         NetworkTestUtils.waitForChannelClose(selector, node, ChannelState.State.AUTHENTICATION_FAILED);
@@ -316,16 +357,10 @@ public void testListenerConfigOverride() throws Exception {
      */
     @Test
     public void testClientAuthenticationRequiredUntrustedProvided() throws Exception {
-        String node = "0";
         sslServerConfigs = serverCertStores.getUntrustingConfig();
+        sslServerConfigs.putAll(sslConfigOverrides);
         sslServerConfigs.put(BrokerSecurityConfigs.SSL_CLIENT_AUTH_CONFIG, "required");
-        server = createEchoServer(SecurityProtocol.SSL);
-        createSelector(sslClientConfigs);
-        InetSocketAddress addr = new InetSocketAddress("localhost", server.port());
-        selector.connect(node, addr, BUFFER_SIZE, BUFFER_SIZE);
-
-        NetworkTestUtils.waitForChannelClose(selector, node, ChannelState.State.AUTHENTICATION_FAILED);
-        server.verifyAuthenticationMetrics(0, 1);
+        verifySslConfigsWithHandshakeFailure();
     }
 
     /**
@@ -334,19 +369,9 @@ public void testClientAuthenticationRequiredUntrustedProvided() throws Exception
      */
     @Test
     public void testClientAuthenticationRequiredNotProvided() throws Exception {
-        String node = "0";
         sslServerConfigs.put(BrokerSecurityConfigs.SSL_CLIENT_AUTH_CONFIG, "required");
-        server = createEchoServer(SecurityProtocol.SSL);
-
-        sslClientConfigs.remove(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG);
-        sslClientConfigs.remove(SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG);
-        sslClientConfigs.remove(SslConfigs.SSL_KEY_PASSWORD_CONFIG);
-        createSelector(sslClientConfigs);
-        InetSocketAddress addr = new InetSocketAddress("localhost", server.port());
-        selector.connect(node, addr, BUFFER_SIZE, BUFFER_SIZE);
-
-        NetworkTestUtils.waitForChannelClose(selector, node, ChannelState.State.AUTHENTICATION_FAILED);
-        server.verifyAuthenticationMetrics(0, 1);
+        CertStores.KEYSTORE_PROPS.forEach(sslClientConfigs::remove);
+        verifySslConfigsWithHandshakeFailure();
     }
 
     /**
@@ -355,15 +380,10 @@ public void testClientAuthenticationRequiredNotProvided() throws Exception {
      */
     @Test
     public void testClientAuthenticationDisabledUntrustedProvided() throws Exception {
-        String node = "0";
         sslServerConfigs = serverCertStores.getUntrustingConfig();
+        sslServerConfigs.putAll(sslConfigOverrides);
         sslServerConfigs.put(BrokerSecurityConfigs.SSL_CLIENT_AUTH_CONFIG, "none");
-        server = createEchoServer(SecurityProtocol.SSL);
-        createSelector(sslClientConfigs);
-        InetSocketAddress addr = new InetSocketAddress("localhost", server.port());
-        selector.connect(node, addr, BUFFER_SIZE, BUFFER_SIZE);
-
-        NetworkTestUtils.checkClientConnection(selector, node, 100, 10);
+        verifySslConfigs();
     }
 
     /**
@@ -372,18 +392,10 @@ public void testClientAuthenticationDisabledUntrustedProvided() throws Exception
      */
     @Test
     public void testClientAuthenticationDisabledNotProvided() throws Exception {
-        String node = "0";
         sslServerConfigs.put(BrokerSecurityConfigs.SSL_CLIENT_AUTH_CONFIG, "none");
-        server = createEchoServer(SecurityProtocol.SSL);
 
-        sslClientConfigs.remove(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG);
-        sslClientConfigs.remove(SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG);
-        sslClientConfigs.remove(SslConfigs.SSL_KEY_PASSWORD_CONFIG);
-        createSelector(sslClientConfigs);
-        InetSocketAddress addr = new InetSocketAddress("localhost", server.port());
-        selector.connect(node, addr, BUFFER_SIZE, BUFFER_SIZE);
-
-        NetworkTestUtils.checkClientConnection(selector, node, 100, 10);
+        CertStores.KEYSTORE_PROPS.forEach(sslClientConfigs::remove);
+        verifySslConfigs();
     }
 
     /**
@@ -392,14 +404,8 @@ public void testClientAuthenticationDisabledNotProvided() throws Exception {
      */
     @Test
     public void testClientAuthenticationRequestedValidProvided() throws Exception {
-        String node = "0";
         sslServerConfigs.put(BrokerSecurityConfigs.SSL_CLIENT_AUTH_CONFIG, "requested");
-        server = createEchoServer(SecurityProtocol.SSL);
-        createSelector(sslClientConfigs);
-        InetSocketAddress addr = new InetSocketAddress("localhost", server.port());
-        selector.connect(node, addr, BUFFER_SIZE, BUFFER_SIZE);
-
-        NetworkTestUtils.checkClientConnection(selector, node, 100, 10);
+        verifySslConfigs();
     }
 
     /**
@@ -408,27 +414,90 @@ public void testClientAuthenticationRequestedValidProvided() throws Exception {
      */
     @Test
     public void testClientAuthenticationRequestedNotProvided() throws Exception {
-        String node = "0";
         sslServerConfigs.put(BrokerSecurityConfigs.SSL_CLIENT_AUTH_CONFIG, "requested");
-        server = createEchoServer(SecurityProtocol.SSL);
 
-        sslClientConfigs.remove(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG);
-        sslClientConfigs.remove(SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG);
-        sslClientConfigs.remove(SslConfigs.SSL_KEY_PASSWORD_CONFIG);
-        createSelector(sslClientConfigs);
-        InetSocketAddress addr = new InetSocketAddress("localhost", server.port());
-        selector.connect(node, addr, BUFFER_SIZE, BUFFER_SIZE);
+        CertStores.KEYSTORE_PROPS.forEach(sslClientConfigs::remove);
+        verifySslConfigs();
+    }
 
-        NetworkTestUtils.checkClientConnection(selector, node, 100, 10);
+    /**
+     * Tests key-pair created using DSA.
+     */
+    @Test
+    public void testDsaKeyPair() throws Exception {
+        // DSA algorithms are not supported for TLSv1.3.
+        assumeTrue(tlsProtocol.equals("TLSv1.2"));
+        serverCertStores = certBuilder(true, "server").keyAlgorithm("DSA").build();
+        clientCertStores = certBuilder(false, "client").keyAlgorithm("DSA").build();
+        sslServerConfigs = getTrustingConfig(serverCertStores, clientCertStores);
+        sslClientConfigs = getTrustingConfig(clientCertStores, serverCertStores);
+        sslServerConfigs.put(BrokerSecurityConfigs.SSL_CLIENT_AUTH_CONFIG, "required");
+        verifySslConfigs();
+    }
+
+    /**
+     * Tests key-pair created using EC.
+     */
+    @Test
+    public void testECKeyPair() throws Exception {
+        serverCertStores = certBuilder(true, "server").keyAlgorithm("EC").build();
+        clientCertStores = certBuilder(false, "client").keyAlgorithm("EC").build();
+        sslServerConfigs = getTrustingConfig(serverCertStores, clientCertStores);
+        sslClientConfigs = getTrustingConfig(clientCertStores, serverCertStores);
+        sslServerConfigs.put(BrokerSecurityConfigs.SSL_CLIENT_AUTH_CONFIG, "required");
+        verifySslConfigs();
+    }
+
+    /**
+     * Tests PEM key store and trust store files which don't have store passwords.
+     */
+    @Test
+    public void testPemFiles() throws Exception {
+        TestSslUtils.convertToPem(sslServerConfigs, true, true);
+        TestSslUtils.convertToPem(sslClientConfigs, true, true);
+        sslServerConfigs.put(BrokerSecurityConfigs.SSL_CLIENT_AUTH_CONFIG, "required");
+        verifySslConfigs();
+    }
+
+    /**
+     * Test with PEM key store files without key password for client key store. We don't allow this
+     * with PEM files since unprotected private key on disk is not safe. We do allow with inline
+     * PEM config since key config can be encrypted or externalized similar to other password configs.
+     */
+    @Test
+    public void testPemFilesWithoutClientKeyPassword() throws Exception {
+        TestSslUtils.convertToPem(sslServerConfigs, !useInlinePem, true);
+        TestSslUtils.convertToPem(sslClientConfigs, !useInlinePem, false);
+        sslServerConfigs.put(BrokerSecurityConfigs.SSL_CLIENT_AUTH_CONFIG, "required");
+        server = createEchoServer(SecurityProtocol.SSL);
+        if (useInlinePem)
+            verifySslConfigs();
+        else
+            assertThrows(KafkaException.class, () -> createSelector(sslClientConfigs));
+    }
+
+    /**
+     * Test with PEM key store files without key password for server key store.We don't allow this
+     * with PEM files since unprotected private key on disk is not safe. We do allow with inline
+     * PEM config since key config can be encrypted or externalized similar to other password configs.
+     */
+    @Test
+    public void testPemFilesWithoutServerKeyPassword() throws Exception {
+        TestSslUtils.convertToPem(sslServerConfigs, !useInlinePem, false);
+        TestSslUtils.convertToPem(sslClientConfigs, !useInlinePem, true);
+
+        if (useInlinePem)
+            verifySslConfigs();
+        else
+            assertThrows(KafkaException.class, () -> createEchoServer(SecurityProtocol.SSL));
     }
 
     /**
      * Tests that an invalid SecureRandom implementation cannot be configured
      */
     @Test
-    public void testInvalidSecureRandomImplementation() throws Exception {
-        SslChannelBuilder channelBuilder = new SslChannelBuilder(Mode.CLIENT);
-        try {
+    public void testInvalidSecureRandomImplementation() {
+        try (SslChannelBuilder channelBuilder = newClientChannelBuilder()) {
             sslClientConfigs.put(SslConfigs.SSL_SECURE_RANDOM_IMPLEMENTATION_CONFIG, "invalid");
             channelBuilder.configure(sslClientConfigs);
             fail("SSL channel configured with invalid SecureRandom implementation");
@@ -441,9 +510,8 @@ public void testInvalidSecureRandomImplementation() throws Exception {
      * Tests that channels cannot be created if truststore cannot be loaded
      */
     @Test
-    public void testInvalidTruststorePassword() throws Exception {
-        SslChannelBuilder channelBuilder = new SslChannelBuilder(Mode.CLIENT);
-        try {
+    public void testInvalidTruststorePassword() {
+        try (SslChannelBuilder channelBuilder = newClientChannelBuilder()) {
             sslClientConfigs.put(SslConfigs.SSL_TRUSTSTORE_PASSWORD_CONFIG, "invalid");
             channelBuilder.configure(sslClientConfigs);
             fail("SSL channel configured with invalid truststore password");
@@ -456,9 +524,8 @@ public void testInvalidTruststorePassword() throws Exception {
      * Tests that channels cannot be created if keystore cannot be loaded
      */
     @Test
-    public void testInvalidKeystorePassword() throws Exception {
-        SslChannelBuilder channelBuilder = new SslChannelBuilder(Mode.CLIENT);
-        try {
+    public void testInvalidKeystorePassword() {
+        try (SslChannelBuilder channelBuilder = newClientChannelBuilder()) {
             sslClientConfigs.put(SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG, "invalid");
             channelBuilder.configure(sslClientConfigs);
             fail("SSL channel configured with invalid keystore password");
@@ -473,16 +540,10 @@ public void testInvalidKeystorePassword() throws Exception {
      */
     @Test
     public void testNullTruststorePassword() throws Exception {
-        String node = "0";
         sslClientConfigs.remove(SslConfigs.SSL_TRUSTSTORE_PASSWORD_CONFIG);
         sslServerConfigs.remove(SslConfigs.SSL_TRUSTSTORE_PASSWORD_CONFIG);
 
-        server = createEchoServer(SecurityProtocol.SSL);
-        createSelector(sslClientConfigs);
-        InetSocketAddress addr = new InetSocketAddress("localhost", server.port());
-        selector.connect(node, addr, BUFFER_SIZE, BUFFER_SIZE);
-
-        NetworkTestUtils.checkClientConnection(selector, node, 100, 10);
+        verifySslConfigs();
     }
 
     /**
@@ -491,15 +552,53 @@ public void testNullTruststorePassword() throws Exception {
      */
     @Test
     public void testInvalidKeyPassword() throws Exception {
-        String node = "0";
         sslServerConfigs.put(SslConfigs.SSL_KEY_PASSWORD_CONFIG, new Password("invalid"));
+        if (useInlinePem) {
+            // We fail fast for PEM
+            assertThrows(InvalidConfigurationException.class, () -> createEchoServer(SecurityProtocol.SSL));
+            return;
+        }
+        verifySslConfigsWithHandshakeFailure();
+    }
+
+    /**
+     * Tests that connection success with the default TLS version.
+     */
+    @Test
+    public void testTlsDefaults() throws Exception {
+        sslServerConfigs = serverCertStores.getTrustingConfig(clientCertStores);
+        sslClientConfigs = clientCertStores.getTrustingConfig(serverCertStores);
+
+        assertEquals(SslConfigs.DEFAULT_SSL_PROTOCOL, sslServerConfigs.get(SslConfigs.SSL_PROTOCOL_CONFIG));
+        assertEquals(SslConfigs.DEFAULT_SSL_PROTOCOL, sslClientConfigs.get(SslConfigs.SSL_PROTOCOL_CONFIG));
+
         server = createEchoServer(SecurityProtocol.SSL);
         createSelector(sslClientConfigs);
+
+        InetSocketAddress addr = new InetSocketAddress("localhost", server.port());
+        selector.connect("0", addr, BUFFER_SIZE, BUFFER_SIZE);
+
+        NetworkTestUtils.checkClientConnection(selector, "0", 10, 100);
+        server.verifyAuthenticationMetrics(1, 0);
+        selector.close();
+
+        checkAuthenticationFailed("1", "TLSv1.1");
+        server.verifyAuthenticationMetrics(1, 1);
+
+        checkAuthenticationFailed("2", "TLSv1");
+        server.verifyAuthenticationMetrics(1, 2);
+    }
+
+    /** Checks connection failed using the specified {@code tlsVersion}. */
+    private void checkAuthenticationFailed(String node, String tlsVersion) throws IOException {
+        sslClientConfigs.put(SslConfigs.SSL_ENABLED_PROTOCOLS_CONFIG, Arrays.asList(tlsVersion));
+        createSelector(sslClientConfigs);
         InetSocketAddress addr = new InetSocketAddress("localhost", server.port());
         selector.connect(node, addr, BUFFER_SIZE, BUFFER_SIZE);
 
         NetworkTestUtils.waitForChannelClose(selector, node, ChannelState.State.AUTHENTICATION_FAILED);
-        server.verifyAuthenticationMetrics(0, 1);
+
+        selector.close();
     }
 
     /**
@@ -507,16 +606,10 @@ public void testInvalidKeyPassword() throws Exception {
      */
     @Test
     public void testUnsupportedTLSVersion() throws Exception {
-        String node = "0";
         sslServerConfigs.put(SslConfigs.SSL_ENABLED_PROTOCOLS_CONFIG, Arrays.asList("TLSv1.2"));
         server = createEchoServer(SecurityProtocol.SSL);
 
-        sslClientConfigs.put(SslConfigs.SSL_ENABLED_PROTOCOLS_CONFIG, Arrays.asList("TLSv1.1"));
-        createSelector(sslClientConfigs);
-        InetSocketAddress addr = new InetSocketAddress("localhost", server.port());
-        selector.connect(node, addr, BUFFER_SIZE, BUFFER_SIZE);
-
-        NetworkTestUtils.waitForChannelClose(selector, node, ChannelState.State.AUTHENTICATION_FAILED);
+        checkAuthenticationFailed("0", "TLSv1.1");
         server.verifyAuthenticationMetrics(0, 1);
     }
 
@@ -525,20 +618,79 @@ public void testUnsupportedTLSVersion() throws Exception {
      */
     @Test
     public void testUnsupportedCiphers() throws Exception {
-        String node = "0";
-        String[] cipherSuites = SSLContext.getDefault().getDefaultSSLParameters().getCipherSuites();
+        SSLContext context = SSLContext.getInstance(tlsProtocol);
+        context.init(null, null, null);
+        String[] cipherSuites = context.getDefaultSSLParameters().getCipherSuites();
         sslServerConfigs.put(SslConfigs.SSL_CIPHER_SUITES_CONFIG, Arrays.asList(cipherSuites[0]));
         server = createEchoServer(SecurityProtocol.SSL);
 
         sslClientConfigs.put(SslConfigs.SSL_CIPHER_SUITES_CONFIG, Arrays.asList(cipherSuites[1]));
         createSelector(sslClientConfigs);
-        InetSocketAddress addr = new InetSocketAddress("localhost", server.port());
-        selector.connect(node, addr, BUFFER_SIZE, BUFFER_SIZE);
 
-        NetworkTestUtils.waitForChannelClose(selector, node, ChannelState.State.AUTHENTICATION_FAILED);
+        checkAuthenticationFailed("1", tlsProtocol);
         server.verifyAuthenticationMetrics(0, 1);
     }
 
+    @Test
+    public void testServerRequestMetrics() throws Exception {
+        String node = "0";
+        server = createEchoServer(SecurityProtocol.SSL);
+        createSelector(sslClientConfigs, 16384, 16384, 16384);
+        InetSocketAddress addr = new InetSocketAddress("localhost", server.port());
+        selector.connect(node, addr, 102400, 102400);
+        NetworkTestUtils.waitForChannelReady(selector, node);
+        int messageSize = 1024 * 1024;
+        String message = TestUtils.randomString(messageSize);
+        selector.send(new NetworkSend(node, ByteBufferSend.sizePrefixed(ByteBuffer.wrap(message.getBytes()))));
+        while (selector.completedReceives().isEmpty()) {
+            selector.poll(100L);
+        }
+        int totalBytes = messageSize + 4; // including 4-byte size
+        server.waitForMetric("incoming-byte", totalBytes);
+        server.waitForMetric("outgoing-byte", totalBytes);
+        server.waitForMetric("request", 1);
+        server.waitForMetric("response", 1);
+    }
+
+    /**
+     * selector.poll() should be able to fetch more data than netReadBuffer from the socket.
+     */
+    @Test
+    public void testSelectorPollReadSize() throws Exception {
+        String node = "0";
+        server = createEchoServer(SecurityProtocol.SSL);
+        createSelector(sslClientConfigs, 16384, 16384, 16384);
+        InetSocketAddress addr = new InetSocketAddress("localhost", server.port());
+        selector.connect(node, addr, 102400, 102400);
+        NetworkTestUtils.checkClientConnection(selector, node, 81920, 1);
+
+        // Send a message of 80K. This is 5X as large as the socket buffer. It should take at least three selector.poll()
+        // to read this message from socket if the SslTransportLayer.read() does not read all data from socket buffer.
+        String message = TestUtils.randomString(81920);
+        selector.send(new NetworkSend(node, ByteBufferSend.sizePrefixed(ByteBuffer.wrap(message.getBytes()))));
+
+        // Send the message to echo server
+        TestUtils.waitForCondition(() -> {
+            try {
+                selector.poll(100L);
+            } catch (IOException e) {
+                return false;
+            }
+            return selector.completedSends().size() > 0;
+        }, "Timed out waiting for message to be sent");
+
+        // Wait for echo server to send the message back
+        TestUtils.waitForCondition(() ->
+            server.numSent() >= 2, "Timed out waiting for echo server to send message");
+
+        // Read the message from socket with only one poll()
+        selector.poll(1000L);
+
+        Collection receiveList = selector.completedReceives();
+        assertEquals(1, receiveList.size());
+        assertEquals(message, new String(Utils.toArray(receiveList.iterator().next().payload())));
+    }
+
     /**
      * Tests handling of BUFFER_UNDERFLOW during unwrap when network read buffer is smaller than SSL session packet buffer size.
      */
@@ -586,39 +738,48 @@ public void testApplicationBufferResize() throws Exception {
      */
     @Test
     public void testNetworkThreadTimeRecorded() throws Exception {
-        selector.close();
-        this.selector = new Selector(NetworkReceive.UNLIMITED, 5000, new Metrics(), Time.SYSTEM,
-                "MetricGroup", new HashMap(), false, true, channelBuilder, MemoryPool.NONE, new LogContext());
+        LogContext logContext = new LogContext();
+        ChannelBuilder channelBuilder = new SslChannelBuilder(Mode.CLIENT, null, false, logContext);
+        channelBuilder.configure(sslClientConfigs);
+        try (Selector selector = new Selector(NetworkReceive.UNLIMITED, Selector.NO_IDLE_TIMEOUT_MS, new Metrics(), Time.SYSTEM,
+                "MetricGroup", new HashMap<>(), false, true, channelBuilder, MemoryPool.NONE, logContext)) {
 
-        String node = "0";
-        server = createEchoServer(SecurityProtocol.SSL);
-        InetSocketAddress addr = new InetSocketAddress("localhost", server.port());
-        selector.connect(node, addr, BUFFER_SIZE, BUFFER_SIZE);
+            String node = "0";
+            server = createEchoServer(SecurityProtocol.SSL);
+            InetSocketAddress addr = new InetSocketAddress("localhost", server.port());
+            selector.connect(node, addr, BUFFER_SIZE, BUFFER_SIZE);
 
-        String message = TestUtils.randomString(1024 * 1024);
-        NetworkTestUtils.waitForChannelReady(selector, node);
-        KafkaChannel channel = selector.channel(node);
-        assertTrue("SSL handshake time not recorded", channel.getAndResetNetworkThreadTimeNanos() > 0);
-        assertEquals("Time not reset", 0, channel.getAndResetNetworkThreadTimeNanos());
+            String message = TestUtils.randomString(1024 * 1024);
+            NetworkTestUtils.waitForChannelReady(selector, node);
+            final KafkaChannel channel = selector.channel(node);
+            assertTrue("SSL handshake time not recorded", channel.getAndResetNetworkThreadTimeNanos() > 0);
+            assertEquals("Time not reset", 0, channel.getAndResetNetworkThreadTimeNanos());
 
-        selector.mute(node);
-        selector.send(new NetworkSend(node, ByteBuffer.wrap(message.getBytes())));
-        while (selector.completedSends().isEmpty()) {
-            selector.poll(100L);
-        }
-        long sendTimeNanos = channel.getAndResetNetworkThreadTimeNanos();
-        assertTrue("Send time not recorded: " + sendTimeNanos, sendTimeNanos > 0);
-        assertEquals("Time not reset", 0, channel.getAndResetNetworkThreadTimeNanos());
-        assertFalse("Unexpected bytes buffered", channel.hasBytesBuffered());
-        assertEquals(0, selector.completedReceives().size());
+            selector.mute(node);
+            selector.send(new NetworkSend(node, ByteBufferSend.sizePrefixed(ByteBuffer.wrap(message.getBytes()))));
+            while (selector.completedSends().isEmpty()) {
+                selector.poll(100L);
+            }
+            long sendTimeNanos = channel.getAndResetNetworkThreadTimeNanos();
+            assertTrue("Send time not recorded: " + sendTimeNanos, sendTimeNanos > 0);
+            assertEquals("Time not reset", 0, channel.getAndResetNetworkThreadTimeNanos());
+            assertFalse("Unexpected bytes buffered", channel.hasBytesBuffered());
+            assertEquals(0, selector.completedReceives().size());
+
+            selector.unmute(node);
+            // Wait for echo server to send the message back
+            TestUtils.waitForCondition(() -> {
+                try {
+                    selector.poll(100L);
+                } catch (IOException e) {
+                    return false;
+                }
+                return !selector.completedReceives().isEmpty();
+            }, "Timed out waiting for a message to receive from echo server");
 
-        selector.unmute(node);
-        while (selector.completedReceives().isEmpty()) {
-            selector.poll(100L);
-            assertEquals(0, selector.numStagedReceives(channel));
+            long receiveTimeNanos = channel.getAndResetNetworkThreadTimeNanos();
+            assertTrue("Receive time not recorded: " + receiveTimeNanos, receiveTimeNanos > 0);
         }
-        long receiveTimeNanos = channel.getAndResetNetworkThreadTimeNanos();
-        assertTrue("Receive time not recorded: " + receiveTimeNanos, receiveTimeNanos > 0);
     }
 
     /**
@@ -626,7 +787,8 @@ public void testNetworkThreadTimeRecorded() throws Exception {
      */
     @Test
     public void testIOExceptionsDuringHandshakeRead() throws Exception {
-        testIOExceptionsDuringHandshake(true, false);
+        server = createEchoServer(SecurityProtocol.SSL);
+        testIOExceptionsDuringHandshake(FailureAction.THROW_IO_EXCEPTION, FailureAction.NO_OP);
     }
 
     /**
@@ -634,22 +796,62 @@ public void testIOExceptionsDuringHandshakeRead() throws Exception {
      */
     @Test
     public void testIOExceptionsDuringHandshakeWrite() throws Exception {
-        testIOExceptionsDuringHandshake(false, true);
+        server = createEchoServer(SecurityProtocol.SSL);
+        testIOExceptionsDuringHandshake(FailureAction.NO_OP, FailureAction.THROW_IO_EXCEPTION);
     }
 
-    private void testIOExceptionsDuringHandshake(boolean failRead, boolean failWrite) throws Exception {
+    /**
+     * Tests that if the remote end closes connection ungracefully  during SSL handshake while reading data,
+     * the disconnection is not treated as an authentication failure.
+     */
+    @Test
+    public void testUngracefulRemoteCloseDuringHandshakeRead() throws Exception {
+        server = createEchoServer(SecurityProtocol.SSL);
+        testIOExceptionsDuringHandshake(server::closeSocketChannels, FailureAction.NO_OP);
+    }
+
+    /**
+     * Tests that if the remote end closes connection ungracefully during SSL handshake while writing data,
+     * the disconnection is not treated as an authentication failure.
+     */
+    @Test
+    public void testUngracefulRemoteCloseDuringHandshakeWrite() throws Exception {
+        server = createEchoServer(SecurityProtocol.SSL);
+        testIOExceptionsDuringHandshake(FailureAction.NO_OP, server::closeSocketChannels);
+    }
+
+    /**
+     * Tests that if the remote end closes the connection during SSL handshake while reading data,
+     * the disconnection is not treated as an authentication failure.
+     */
+    @Test
+    public void testGracefulRemoteCloseDuringHandshakeRead() throws Exception {
         server = createEchoServer(SecurityProtocol.SSL);
+        testIOExceptionsDuringHandshake(FailureAction.NO_OP, server::closeKafkaChannels);
+    }
+
+    /**
+     * Tests that if the remote end closes the connection during SSL handshake while writing data,
+     * the disconnection is not treated as an authentication failure.
+     */
+    @Test
+    public void testGracefulRemoteCloseDuringHandshakeWrite() throws Exception {
+        server = createEchoServer(SecurityProtocol.SSL);
+        testIOExceptionsDuringHandshake(server::closeKafkaChannels, FailureAction.NO_OP);
+    }
+
+    private void testIOExceptionsDuringHandshake(FailureAction readFailureAction,
+                                                 FailureAction flushFailureAction) throws Exception {
         TestSslChannelBuilder channelBuilder = new TestSslChannelBuilder(Mode.CLIENT);
         boolean done = false;
         for (int i = 1; i <= 100; i++) {
-            int readFailureIndex = failRead ? i : Integer.MAX_VALUE;
-            int flushFailureIndex = failWrite ? i : Integer.MAX_VALUE;
             String node = String.valueOf(i);
 
-            channelBuilder.readFailureIndex = readFailureIndex;
-            channelBuilder.flushFailureIndex = flushFailureIndex;
+            channelBuilder.readFailureAction = readFailureAction;
+            channelBuilder.flushFailureAction = flushFailureAction;
+            channelBuilder.failureIndex = i;
             channelBuilder.configure(sslClientConfigs);
-            this.selector = new Selector(5000, new Metrics(), new MockTime(), "MetricGroup", channelBuilder, new LogContext());
+            this.selector = new Selector(5000, new Metrics(), time, "MetricGroup", channelBuilder, new LogContext());
 
             InetSocketAddress addr = new InetSocketAddress("localhost", server.port());
             selector.connect(node, addr, BUFFER_SIZE, BUFFER_SIZE);
@@ -661,13 +863,16 @@ private void testIOExceptionsDuringHandshake(boolean failRead, boolean failWrite
                     break;
                 }
                 if (selector.disconnected().containsKey(node)) {
-                    assertEquals(ChannelState.State.AUTHENTICATE, selector.disconnected().get(node).state());
+                    ChannelState.State state = selector.disconnected().get(node).state();
+                    assertTrue("Unexpected channel state " + state,
+                            state == ChannelState.State.AUTHENTICATE || state == ChannelState.State.READY);
                     break;
                 }
             }
             KafkaChannel channel = selector.channel(node);
             if (channel != null)
                 assertTrue("Channel not ready or disconnected:" + channel.state().state(), channel.ready());
+            selector.close();
         }
         assertTrue("Too many invocations of read/write during SslTransportLayer.handshake()", done);
     }
@@ -680,6 +885,7 @@ private void testIOExceptionsDuringHandshake(boolean failRead, boolean failWrite
     @Test
     public void testPeerNotifiedOfHandshakeFailure() throws Exception {
         sslServerConfigs = serverCertStores.getUntrustingConfig();
+        sslServerConfigs.putAll(sslConfigOverrides);
         sslServerConfigs.put(BrokerSecurityConfigs.SSL_CLIENT_AUTH_CONFIG, "required");
 
         // Test without delay and a couple of delay counts to ensure delay applies to handshake failure
@@ -690,7 +896,7 @@ public void testPeerNotifiedOfHandshakeFailure() throws Exception {
             serverChannelBuilder.flushDelayCount = i;
             server = new NioEchoServer(ListenerName.forSecurityProtocol(SecurityProtocol.SSL),
                     SecurityProtocol.SSL, new TestSecurityConfig(sslServerConfigs),
-                    "localhost", serverChannelBuilder, null);
+                    "localhost", serverChannelBuilder, null, time);
             server.start();
             createSelector(sslClientConfigs);
             InetSocketAddress addr = new InetSocketAddress("localhost", server.port());
@@ -699,28 +905,38 @@ SecurityProtocol.SSL, new TestSecurityConfig(sslServerConfigs),
             NetworkTestUtils.waitForChannelClose(selector, node, ChannelState.State.AUTHENTICATION_FAILED);
             server.close();
             selector.close();
+            serverChannelBuilder.close();
         }
     }
 
     @Test
     public void testCloseSsl() throws Exception {
-        testClose(SecurityProtocol.SSL, new SslChannelBuilder(Mode.CLIENT));
+        testClose(SecurityProtocol.SSL, newClientChannelBuilder());
     }
 
     @Test
     public void testClosePlaintext() throws Exception {
-        testClose(SecurityProtocol.PLAINTEXT, new PlaintextChannelBuilder());
+        testClose(SecurityProtocol.PLAINTEXT, new PlaintextChannelBuilder(null));
+    }
+
+    private SslChannelBuilder newClientChannelBuilder() {
+        return new SslChannelBuilder(Mode.CLIENT, null, false, new LogContext());
     }
 
     private void testClose(SecurityProtocol securityProtocol, ChannelBuilder clientChannelBuilder) throws Exception {
         String node = "0";
         server = createEchoServer(securityProtocol);
         clientChannelBuilder.configure(sslClientConfigs);
-        this.selector = new Selector(5000, new Metrics(), new MockTime(), "MetricGroup", clientChannelBuilder, new LogContext());
+        this.selector = new Selector(5000, new Metrics(), time, "MetricGroup", clientChannelBuilder, new LogContext());
         InetSocketAddress addr = new InetSocketAddress("localhost", server.port());
         selector.connect(node, addr, BUFFER_SIZE, BUFFER_SIZE);
 
         NetworkTestUtils.waitForChannelReady(selector, node);
+        // `waitForChannelReady` waits for client-side channel to be ready. This is sufficient for other tests
+        // operating on the client-side channel. But here, we are muting the server-side channel below, so we
+        // need to wait for the server-side channel to be ready as well.
+        TestUtils.waitForCondition(() -> server.selector().channels().stream().allMatch(KafkaChannel::ready),
+                "Channel not ready");
 
         final ByteArrayOutputStream bytesOut = new ByteArrayOutputStream();
         server.outputChannel(Channels.newChannel(bytesOut));
@@ -729,53 +945,355 @@ private void testClose(SecurityProtocol securityProtocol, ChannelBuilder clientC
         int count = 20;
         final int totalSendSize = count * (message.length + 4);
         for (int i = 0; i < count; i++) {
-            selector.send(new NetworkSend(node, ByteBuffer.wrap(message)));
+            selector.send(new NetworkSend(node, ByteBufferSend.sizePrefixed(ByteBuffer.wrap(message))));
             do {
                 selector.poll(0L);
             } while (selector.completedSends().isEmpty());
         }
         server.selector().unmuteAll();
         selector.close(node);
-        TestUtils.waitForCondition(new TestCondition() {
-            @Override
-            public boolean conditionMet() {
-                return bytesOut.toByteArray().length == totalSendSize;
-            }
-        }, 5000, "All requests sent were not processed");
+        TestUtils.waitForCondition(() ->
+            bytesOut.toByteArray().length == totalSendSize, 5000, "All requests sent were not processed");
     }
 
-    private void createSelector(Map sslClientConfigs) {
-        createSelector(sslClientConfigs, null, null, null);
+    /**
+     * 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, new LogContext());
+        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() {
+        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, new LogContext());
     }
 
-    private void createSelector(Map sslClientConfigs, final Integer netReadBufSize,
+    /**
+     * Tests reconfiguration of server keystore. Verifies that existing connections continue
+     * to work with old keystore and new connections work with new keystore.
+     */
+    @Test
+    public void testServerKeystoreDynamicUpdate() throws Exception {
+        SecurityProtocol securityProtocol = SecurityProtocol.SSL;
+        TestSecurityConfig config = new TestSecurityConfig(sslServerConfigs);
+        ListenerName listenerName = ListenerName.forSecurityProtocol(securityProtocol);
+        ChannelBuilder serverChannelBuilder = ChannelBuilders.serverChannelBuilder(listenerName,
+                false, securityProtocol, config, null, null, time, new LogContext());
+        server = new NioEchoServer(listenerName, securityProtocol, config,
+                "localhost", serverChannelBuilder, null, time);
+        server.start();
+        InetSocketAddress addr = new InetSocketAddress("localhost", server.port());
+
+        // Verify that client with matching truststore can authenticate, send and receive
+        String oldNode = "0";
+        Selector oldClientSelector = createSelector(sslClientConfigs);
+        oldClientSelector.connect(oldNode, addr, BUFFER_SIZE, BUFFER_SIZE);
+        NetworkTestUtils.checkClientConnection(selector, oldNode, 100, 10);
+
+        CertStores newServerCertStores = certBuilder(true, "server").addHostName("localhost").build();
+        Map newKeystoreConfigs = newServerCertStores.keyStoreProps();
+        assertTrue("SslChannelBuilder not reconfigurable", serverChannelBuilder instanceof ListenerReconfigurable);
+        ListenerReconfigurable reconfigurableBuilder = (ListenerReconfigurable) serverChannelBuilder;
+        assertEquals(listenerName, reconfigurableBuilder.listenerName());
+        reconfigurableBuilder.validateReconfiguration(newKeystoreConfigs);
+        reconfigurableBuilder.reconfigure(newKeystoreConfigs);
+
+        // Verify that new client with old truststore fails
+        oldClientSelector.connect("1", addr, BUFFER_SIZE, BUFFER_SIZE);
+        NetworkTestUtils.waitForChannelClose(oldClientSelector, "1", ChannelState.State.AUTHENTICATION_FAILED);
+
+        // Verify that new client with new truststore can authenticate, send and receive
+        sslClientConfigs = getTrustingConfig(clientCertStores, newServerCertStores);
+        Selector newClientSelector = createSelector(sslClientConfigs);
+        newClientSelector.connect("2", addr, BUFFER_SIZE, BUFFER_SIZE);
+        NetworkTestUtils.checkClientConnection(newClientSelector, "2", 100, 10);
+
+        // Verify that old client continues to work
+        NetworkTestUtils.checkClientConnection(oldClientSelector, oldNode, 100, 10);
+
+        CertStores invalidCertStores = certBuilder(true, "server").addHostName("127.0.0.1").build();
+        Map  invalidConfigs = getTrustingConfig(invalidCertStores, clientCertStores);
+        verifyInvalidReconfigure(reconfigurableBuilder, invalidConfigs, "keystore with different SubjectAltName");
+
+        Map  missingStoreConfigs = new HashMap<>();
+        missingStoreConfigs.put(SslConfigs.SSL_KEYSTORE_TYPE_CONFIG, "PKCS12");
+        missingStoreConfigs.put(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG, "some.keystore.path");
+        missingStoreConfigs.put(SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG, new Password("some.keystore.password"));
+        missingStoreConfigs.put(SslConfigs.SSL_KEY_PASSWORD_CONFIG, new Password("some.key.password"));
+        verifyInvalidReconfigure(reconfigurableBuilder, missingStoreConfigs, "keystore not found");
+
+        // Verify that new connections continue to work with the server with previously configured keystore after failed reconfiguration
+        newClientSelector.connect("3", addr, BUFFER_SIZE, BUFFER_SIZE);
+        NetworkTestUtils.checkClientConnection(newClientSelector, "3", 100, 10);
+    }
+
+    @Test
+    public void testServerKeystoreDynamicUpdateWithNewSubjectAltName() throws Exception {
+        SecurityProtocol securityProtocol = SecurityProtocol.SSL;
+        TestSecurityConfig config = new TestSecurityConfig(sslServerConfigs);
+        ListenerName listenerName = ListenerName.forSecurityProtocol(securityProtocol);
+        ChannelBuilder serverChannelBuilder = ChannelBuilders.serverChannelBuilder(listenerName,
+            false, securityProtocol, config, null, null, time, new LogContext());
+        server = new NioEchoServer(listenerName, securityProtocol, config,
+            "localhost", serverChannelBuilder, null, time);
+        server.start();
+        InetSocketAddress addr = new InetSocketAddress("localhost", server.port());
+
+        Selector selector = createSelector(sslClientConfigs);
+        String node1 = "1";
+        selector.connect(node1, addr, BUFFER_SIZE, BUFFER_SIZE);
+        NetworkTestUtils.checkClientConnection(selector, node1, 100, 10);
+        selector.close();
+
+        TestSslUtils.CertificateBuilder certBuilder = new TestSslUtils.CertificateBuilder().sanDnsNames("localhost", "*.example.com");
+        String truststorePath = (String) sslClientConfigs.get(SslConfigs.SSL_TRUSTSTORE_LOCATION_CONFIG);
+        File truststoreFile = truststorePath != null ? new File(truststorePath) : null;
+        TestSslUtils.SslConfigsBuilder builder = new TestSslUtils.SslConfigsBuilder(Mode.SERVER)
+                .useClientCert(false)
+                .certAlias("server")
+                .cn("server")
+                .certBuilder(certBuilder)
+                .createNewTrustStore(truststoreFile)
+                .usePem(useInlinePem);
+        Map newConfigs = builder.build();
+        Map newKeystoreConfigs = new HashMap<>();
+        for (String propName : CertStores.KEYSTORE_PROPS) {
+            newKeystoreConfigs.put(propName, newConfigs.get(propName));
+        }
+        ListenerReconfigurable reconfigurableBuilder = (ListenerReconfigurable) serverChannelBuilder;
+        reconfigurableBuilder.validateReconfiguration(newKeystoreConfigs);
+        reconfigurableBuilder.reconfigure(newKeystoreConfigs);
+
+        for (String propName : CertStores.TRUSTSTORE_PROPS) {
+            sslClientConfigs.put(propName, newConfigs.get(propName));
+        }
+        selector = createSelector(sslClientConfigs);
+        String node2 = "2";
+        selector.connect(node2, addr, BUFFER_SIZE, BUFFER_SIZE);
+        NetworkTestUtils.checkClientConnection(selector, node2, 100, 10);
+
+        TestSslUtils.CertificateBuilder invalidBuilder = new TestSslUtils.CertificateBuilder().sanDnsNames("localhost");
+        if (!useInlinePem)
+            builder.useExistingTrustStore(truststoreFile);
+        Map invalidConfig = builder.certBuilder(invalidBuilder).build();
+        Map invalidKeystoreConfigs = new HashMap<>();
+        for (String propName : CertStores.KEYSTORE_PROPS) {
+            invalidKeystoreConfigs.put(propName, invalidConfig.get(propName));
+        }
+        verifyInvalidReconfigure(reconfigurableBuilder, invalidKeystoreConfigs, "keystore without existing SubjectAltName");
+        String node3 = "3";
+        selector.connect(node3, addr, BUFFER_SIZE, BUFFER_SIZE);
+        NetworkTestUtils.checkClientConnection(selector, node3, 100, 10);
+    }
+
+    /**
+     * Tests reconfiguration of server truststore. Verifies that existing connections continue
+     * to work with old truststore and new connections work with new truststore.
+     */
+    @Test
+    public void testServerTruststoreDynamicUpdate() throws Exception {
+        SecurityProtocol securityProtocol = SecurityProtocol.SSL;
+        sslServerConfigs.put(BrokerSecurityConfigs.SSL_CLIENT_AUTH_CONFIG, "required");
+        TestSecurityConfig config = new TestSecurityConfig(sslServerConfigs);
+        ListenerName listenerName = ListenerName.forSecurityProtocol(securityProtocol);
+        ChannelBuilder serverChannelBuilder = ChannelBuilders.serverChannelBuilder(listenerName,
+                false, securityProtocol, config, null, null, time, new LogContext());
+        server = new NioEchoServer(listenerName, securityProtocol, config,
+                "localhost", serverChannelBuilder, null, time);
+        server.start();
+        InetSocketAddress addr = new InetSocketAddress("localhost", server.port());
+
+        // Verify that client with matching keystore can authenticate, send and receive
+        String oldNode = "0";
+        Selector oldClientSelector = createSelector(sslClientConfigs);
+        oldClientSelector.connect(oldNode, addr, BUFFER_SIZE, BUFFER_SIZE);
+        NetworkTestUtils.checkClientConnection(selector, oldNode, 100, 10);
+
+        CertStores newClientCertStores = certBuilder(true, "client").addHostName("localhost").build();
+        sslClientConfigs = getTrustingConfig(newClientCertStores, serverCertStores);
+        Map newTruststoreConfigs = newClientCertStores.trustStoreProps();
+        assertTrue("SslChannelBuilder not reconfigurable", serverChannelBuilder instanceof ListenerReconfigurable);
+        ListenerReconfigurable reconfigurableBuilder = (ListenerReconfigurable) serverChannelBuilder;
+        assertEquals(listenerName, reconfigurableBuilder.listenerName());
+        reconfigurableBuilder.validateReconfiguration(newTruststoreConfigs);
+        reconfigurableBuilder.reconfigure(newTruststoreConfigs);
+
+        // Verify that new client with old truststore fails
+        oldClientSelector.connect("1", addr, BUFFER_SIZE, BUFFER_SIZE);
+        NetworkTestUtils.waitForChannelClose(oldClientSelector, "1", ChannelState.State.AUTHENTICATION_FAILED);
+
+        // Verify that new client with new truststore can authenticate, send and receive
+        Selector newClientSelector = createSelector(sslClientConfigs);
+        newClientSelector.connect("2", addr, BUFFER_SIZE, BUFFER_SIZE);
+        NetworkTestUtils.checkClientConnection(newClientSelector, "2", 100, 10);
+
+        // Verify that old client continues to work
+        NetworkTestUtils.checkClientConnection(oldClientSelector, oldNode, 100, 10);
+
+        Map  invalidConfigs = new HashMap<>(newTruststoreConfigs);
+        invalidConfigs.put(SslConfigs.SSL_TRUSTSTORE_TYPE_CONFIG, "INVALID_TYPE");
+        verifyInvalidReconfigure(reconfigurableBuilder, invalidConfigs, "invalid truststore type");
+
+        Map  missingStoreConfigs = new HashMap<>();
+        missingStoreConfigs.put(SslConfigs.SSL_TRUSTSTORE_TYPE_CONFIG, "PKCS12");
+        missingStoreConfigs.put(SslConfigs.SSL_TRUSTSTORE_LOCATION_CONFIG, "some.truststore.path");
+        missingStoreConfigs.put(SslConfigs.SSL_TRUSTSTORE_PASSWORD_CONFIG, new Password("some.truststore.password"));
+        verifyInvalidReconfigure(reconfigurableBuilder, missingStoreConfigs, "truststore not found");
+
+        // Verify that new connections continue to work with the server with previously configured keystore after failed reconfiguration
+        newClientSelector.connect("3", addr, BUFFER_SIZE, BUFFER_SIZE);
+        NetworkTestUtils.checkClientConnection(newClientSelector, "3", 100, 10);
+    }
+
+    /**
+     * Tests if client can plugin customize ssl.engine.factory
+     */
+    @Test
+    public void testCustomClientSslEngineFactory() throws Exception {
+        sslClientConfigs.put(SslConfigs.SSL_ENGINE_FACTORY_CLASS_CONFIG, TestSslUtils.TestSslEngineFactory.class);
+        verifySslConfigs();
+    }
+
+    /**
+     * Tests if server can plugin customize ssl.engine.factory
+     */
+    @Test
+    public void testCustomServerSslEngineFactory() throws Exception {
+        sslServerConfigs.put(SslConfigs.SSL_ENGINE_FACTORY_CLASS_CONFIG, TestSslUtils.TestSslEngineFactory.class);
+        verifySslConfigs();
+    }
+
+    /**
+     * Tests if client and server both can plugin customize ssl.engine.factory and talk to each other!
+     */
+    @Test
+    public void testCustomClientAndServerSslEngineFactory() throws Exception {
+        sslClientConfigs.put(SslConfigs.SSL_ENGINE_FACTORY_CLASS_CONFIG, TestSslUtils.TestSslEngineFactory.class);
+        sslServerConfigs.put(SslConfigs.SSL_ENGINE_FACTORY_CLASS_CONFIG, TestSslUtils.TestSslEngineFactory.class);
+        verifySslConfigs();
+    }
+
+    /**
+     * Tests invalid ssl.engine.factory plugin class
+     */
+    @Test(expected = KafkaException.class)
+    public void testInvalidSslEngineFactory() {
+        sslClientConfigs.put(SslConfigs.SSL_ENGINE_FACTORY_CLASS_CONFIG, String.class);
+        createSelector(sslClientConfigs);
+    }
+
+    private void verifyInvalidReconfigure(ListenerReconfigurable reconfigurable,
+                                          Map  invalidConfigs, String errorMessage) {
+        try {
+            reconfigurable.validateReconfiguration(invalidConfigs);
+            fail("Should have failed validation with an exception: " + errorMessage);
+        } catch (KafkaException e) {
+            // expected exception
+        }
+        try {
+            reconfigurable.reconfigure(invalidConfigs);
+            fail("Should have failed to reconfigure: " + errorMessage);
+        } catch (KafkaException e) {
+            // expected exception
+        }
+    }
+
+    private Selector createSelector(Map sslClientConfigs) {
+        return createSelector(sslClientConfigs, null, null, null);
+    }
+
+    private Selector createSelector(Map sslClientConfigs, final Integer netReadBufSize,
                                 final Integer netWriteBufSize, final Integer appBufSize) {
         TestSslChannelBuilder channelBuilder = new TestSslChannelBuilder(Mode.CLIENT);
         channelBuilder.configureBufferSizes(netReadBufSize, netWriteBufSize, appBufSize);
-        this.channelBuilder = channelBuilder;
-        this.channelBuilder.configure(sslClientConfigs);
-        this.selector = new Selector(5000, new Metrics(), new MockTime(), "MetricGroup", channelBuilder, new LogContext());
+        channelBuilder.configure(sslClientConfigs);
+        this.selector = new Selector(100 * 5000, new Metrics(), time, "MetricGroup", channelBuilder, new LogContext());
+        return selector;
     }
 
     private NioEchoServer createEchoServer(ListenerName listenerName, SecurityProtocol securityProtocol) throws Exception {
-        return NetworkTestUtils.createEchoServer(listenerName, securityProtocol, new TestSecurityConfig(sslServerConfigs), null);
+        return NetworkTestUtils.createEchoServer(listenerName, securityProtocol, new TestSecurityConfig(sslServerConfigs), null, time);
     }
 
     private NioEchoServer createEchoServer(SecurityProtocol securityProtocol) throws Exception {
         return createEchoServer(ListenerName.forSecurityProtocol(securityProtocol), securityProtocol);
     }
 
-    private static class TestSslChannelBuilder extends SslChannelBuilder {
+    private Map getTrustingConfig(CertStores certStores, CertStores peerCertStores) {
+        Map configs = certStores.getTrustingConfig(peerCertStores);
+        configs.putAll(sslConfigOverrides);
+        return configs;
+    }
+
+    private void verifySslConfigs() throws Exception {
+        server = createEchoServer(SecurityProtocol.SSL);
+        createSelector(sslClientConfigs);
+        InetSocketAddress addr = new InetSocketAddress("localhost", server.port());
+        String node = "0";
+        selector.connect(node, addr, BUFFER_SIZE, BUFFER_SIZE);
+        NetworkTestUtils.checkClientConnection(selector, node, 100, 10);
+    }
+
+    public void verifySslConfigsWithHandshakeFailure() throws Exception {
+        server = createEchoServer(SecurityProtocol.SSL);
+        createSelector(sslClientConfigs);
+        InetSocketAddress addr = new InetSocketAddress("localhost", server.port());
+        String node = "0";
+        selector.connect(node, addr, BUFFER_SIZE, BUFFER_SIZE);
+        NetworkTestUtils.waitForChannelClose(selector, node, ChannelState.State.AUTHENTICATION_FAILED);
+        server.verifyAuthenticationMetrics(0, 1);
+    }
+
+    @FunctionalInterface
+    private interface FailureAction {
+        FailureAction NO_OP = () -> { };
+        FailureAction THROW_IO_EXCEPTION = () -> {
+            throw new IOException("Test IO exception");
+        };
+        void run() throws IOException;
+    }
+
+    static class TestSslChannelBuilder extends SslChannelBuilder {
 
         private Integer netReadBufSizeOverride;
         private Integer netWriteBufSizeOverride;
         private Integer appBufSizeOverride;
-        long readFailureIndex = Long.MAX_VALUE;
-        long flushFailureIndex = Long.MAX_VALUE;
+        private long failureIndex = Long.MAX_VALUE;
+        FailureAction readFailureAction = FailureAction.NO_OP;
+        FailureAction flushFailureAction = FailureAction.NO_OP;
         int flushDelayCount = 0;
 
         public TestSslChannelBuilder(Mode mode) {
-            super(mode);
+            super(mode, null, false, new LogContext());
         }
 
         public void configureBufferSizes(Integer netReadBufSize, Integer netWriteBufSize, Integer appBufSize) {
@@ -785,12 +1303,11 @@ public void configureBufferSizes(Integer netReadBufSize, Integer netWriteBufSize
         }
 
         @Override
-        protected SslTransportLayer buildTransportLayer(SslFactory sslFactory, String id, SelectionKey key, String host) throws IOException {
+        protected SslTransportLayer buildTransportLayer(SslFactory sslFactory, String id, SelectionKey key,
+                                                        String host, ChannelMetadataRegistry metadataRegistry) throws IOException {
             SocketChannel socketChannel = (SocketChannel) key.channel();
             SSLEngine sslEngine = sslFactory.createSslEngine(host, socketChannel.socket().getPort());
-            TestSslTransportLayer transportLayer = newTransportLayer(id, key, sslEngine);
-            transportLayer.startHandshake();
-            return transportLayer;
+            return newTransportLayer(id, key, sslEngine);
         }
 
         protected TestSslTransportLayer newTransportLayer(String id, SelectionKey key, SSLEngine sslEngine) throws IOException {
@@ -816,13 +1333,13 @@ class TestSslTransportLayer extends SslTransportLayer {
             private final AtomicLong numFlushesRemaining;
             private final AtomicInteger numDelayedFlushesRemaining;
 
-            public TestSslTransportLayer(String channelId, SelectionKey key, SSLEngine sslEngine) throws IOException {
-                super(channelId, key, sslEngine);
+            public TestSslTransportLayer(String channelId, SelectionKey key, SSLEngine sslEngine) {
+                super(channelId, key, sslEngine, new DefaultChannelMetadataRegistry());
                 this.netReadBufSize = new ResizeableBufferSize(netReadBufSizeOverride);
                 this.netWriteBufSize = new ResizeableBufferSize(netWriteBufSizeOverride);
                 this.appBufSize = new ResizeableBufferSize(appBufSizeOverride);
-                numReadsRemaining = new AtomicLong(readFailureIndex);
-                numFlushesRemaining = new AtomicLong(flushFailureIndex);
+                numReadsRemaining = new AtomicLong(failureIndex);
+                numFlushesRemaining = new AtomicLong(failureIndex);
                 numDelayedFlushesRemaining = new AtomicInteger(flushDelayCount);
             }
 
@@ -850,26 +1367,32 @@ protected int applicationBufferSize() {
             @Override
             protected int readFromSocketChannel() throws IOException {
                 if (numReadsRemaining.decrementAndGet() == 0 && !ready())
-                    throw new IOException("Test exception during read");
+                    readFailureAction.run();
                 return super.readFromSocketChannel();
             }
 
             @Override
             protected boolean flush(ByteBuffer buf) throws IOException {
                 if (numFlushesRemaining.decrementAndGet() == 0 && !ready())
-                    throw new IOException("Test exception during write");
+                    flushFailureAction.run();
                 else if (numDelayedFlushesRemaining.getAndDecrement() != 0)
                     return false;
                 resetDelayedFlush();
                 return super.flush(buf);
             }
 
+            @Override
+            protected void startHandshake() throws IOException {
+                assertTrue("SSL handshake initialized too early", socketChannel().isConnected());
+                super.startHandshake();
+            }
+
             private void resetDelayedFlush() {
                 numDelayedFlushesRemaining.set(flushDelayCount);
             }
         }
 
-        private static class ResizeableBufferSize {
+        static class ResizeableBufferSize {
             private Integer bufSizeOverride;
             ResizeableBufferSize(Integer bufSizeOverride) {
                 this.bufSizeOverride = bufSizeOverride;
diff --git a/clients/src/test/java/org/apache/kafka/common/network/SslTransportTls12Tls13Test.java b/clients/src/test/java/org/apache/kafka/common/network/SslTransportTls12Tls13Test.java
new file mode 100644
index 0000000000000..81b86d4e6a1b7
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/network/SslTransportTls12Tls13Test.java
@@ -0,0 +1,169 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.network;
+
+import java.io.IOException;
+import java.net.InetSocketAddress;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.Map;
+import org.apache.kafka.common.config.SslConfigs;
+import org.apache.kafka.common.metrics.Metrics;
+import org.apache.kafka.common.security.TestSecurityConfig;
+import org.apache.kafka.common.security.auth.SecurityProtocol;
+import org.apache.kafka.common.utils.Java;
+import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.common.utils.Time;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+import static org.junit.Assume.assumeTrue;
+
+public class SslTransportTls12Tls13Test {
+    private static final int BUFFER_SIZE = 4 * 1024;
+    private static final Time TIME = Time.SYSTEM;
+
+    private NioEchoServer server;
+    private Selector selector;
+    private Map sslClientConfigs;
+    private Map sslServerConfigs;
+
+    @Before
+    public void setup() throws Exception {
+        // Create certificates for use by client and server. Add server cert to client truststore and vice versa.
+        CertStores serverCertStores = new CertStores(true, "server", "localhost");
+        CertStores clientCertStores = new CertStores(false, "client", "localhost");
+        sslServerConfigs = serverCertStores.getTrustingConfig(clientCertStores);
+        sslClientConfigs = clientCertStores.getTrustingConfig(serverCertStores);
+
+        LogContext logContext = new LogContext();
+        ChannelBuilder channelBuilder = new SslChannelBuilder(Mode.CLIENT, null, false, logContext);
+        channelBuilder.configure(sslClientConfigs);
+        this.selector = new Selector(5000, new Metrics(), TIME, "MetricGroup", channelBuilder, logContext);
+    }
+
+    @After
+    public void teardown() throws Exception {
+        if (selector != null)
+            this.selector.close();
+        if (server != null)
+            this.server.close();
+    }
+
+    /**
+     * Tests that connections fails if TLSv1.3 enabled but cipher suite suitable only for TLSv1.2 used.
+     */
+    @Test
+    public void testCiphersSuiteForTls12FailsForTls13() throws Exception {
+        assumeTrue(Java.IS_JAVA11_COMPATIBLE);
+
+        String cipherSuite = "TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384";
+
+        sslServerConfigs.put(SslConfigs.SSL_ENABLED_PROTOCOLS_CONFIG, Collections.singletonList("TLSv1.3"));
+        sslServerConfigs.put(SslConfigs.SSL_CIPHER_SUITES_CONFIG, Collections.singletonList(cipherSuite));
+        server = NetworkTestUtils.createEchoServer(ListenerName.forSecurityProtocol(SecurityProtocol.SSL),
+            SecurityProtocol.SSL, new TestSecurityConfig(sslServerConfigs), null, TIME);
+
+        sslClientConfigs.put(SslConfigs.SSL_ENABLED_PROTOCOLS_CONFIG, Collections.singletonList("TLSv1.3"));
+        sslClientConfigs.put(SslConfigs.SSL_CIPHER_SUITES_CONFIG, Collections.singletonList(cipherSuite));
+
+        checkAuthentiationFailed();
+    }
+
+    /**
+     * Tests that connections can't be made if server uses TLSv1.2 with custom cipher suite and client uses TLSv1.3.
+     */
+    @Test
+    public void testCiphersSuiteFailForServerTls12ClientTls13() throws Exception {
+        assumeTrue(Java.IS_JAVA11_COMPATIBLE);
+
+        String tls12CipherSuite = "TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384";
+        String tls13CipherSuite = "TLS_AES_128_GCM_SHA256";
+
+        sslServerConfigs.put(SslConfigs.SSL_PROTOCOL_CONFIG, "TLSv1.2");
+        sslServerConfigs.put(SslConfigs.SSL_ENABLED_PROTOCOLS_CONFIG, Collections.singletonList("TLSv1.2"));
+        sslServerConfigs.put(SslConfigs.SSL_CIPHER_SUITES_CONFIG, Collections.singletonList(tls12CipherSuite));
+        server = NetworkTestUtils.createEchoServer(ListenerName.forSecurityProtocol(SecurityProtocol.SSL),
+            SecurityProtocol.SSL, new TestSecurityConfig(sslServerConfigs), null, TIME);
+
+        sslClientConfigs.put(SslConfigs.SSL_PROTOCOL_CONFIG, "TLSv1.3");
+        sslClientConfigs.put(SslConfigs.SSL_CIPHER_SUITES_CONFIG, Collections.singletonList(tls13CipherSuite));
+
+        checkAuthentiationFailed();
+    }
+
+    /**
+     * Tests that connections can be made with TLSv1.3 cipher suite.
+     */
+    @Test
+    public void testCiphersSuiteForTls13() throws Exception {
+        assumeTrue(Java.IS_JAVA11_COMPATIBLE);
+
+        String cipherSuite = "TLS_AES_128_GCM_SHA256";
+
+        sslServerConfigs.put(SslConfigs.SSL_CIPHER_SUITES_CONFIG, Collections.singletonList(cipherSuite));
+        server = NetworkTestUtils.createEchoServer(ListenerName.forSecurityProtocol(SecurityProtocol.SSL),
+            SecurityProtocol.SSL, new TestSecurityConfig(sslServerConfigs), null, TIME);
+
+        sslClientConfigs.put(SslConfigs.SSL_CIPHER_SUITES_CONFIG, Collections.singletonList(cipherSuite));
+        checkAuthenticationSucceed();
+    }
+
+    /**
+     * Tests that connections can be made with TLSv1.2 cipher suite.
+     */
+    @Test
+    public void testCiphersSuiteForTls12() throws Exception {
+        String cipherSuite = "TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384";
+
+        sslServerConfigs.put(SslConfigs.SSL_ENABLED_PROTOCOLS_CONFIG, Arrays.asList(SslConfigs.DEFAULT_SSL_ENABLED_PROTOCOLS.split(",")));
+        sslServerConfigs.put(SslConfigs.SSL_CIPHER_SUITES_CONFIG, Collections.singletonList(cipherSuite));
+        server = NetworkTestUtils.createEchoServer(ListenerName.forSecurityProtocol(SecurityProtocol.SSL),
+            SecurityProtocol.SSL, new TestSecurityConfig(sslServerConfigs), null, TIME);
+
+        sslClientConfigs.put(SslConfigs.SSL_ENABLED_PROTOCOLS_CONFIG, Arrays.asList(SslConfigs.DEFAULT_SSL_ENABLED_PROTOCOLS.split(",")));
+        sslClientConfigs.put(SslConfigs.SSL_CIPHER_SUITES_CONFIG, Collections.singletonList(cipherSuite));
+        checkAuthenticationSucceed();
+    }
+
+    /** Checks connection failed using the specified {@code tlsVersion}. */
+    private void checkAuthentiationFailed() throws IOException, InterruptedException {
+        sslClientConfigs.put(SslConfigs.SSL_ENABLED_PROTOCOLS_CONFIG, Arrays.asList("TLSv1.3"));
+        createSelector(sslClientConfigs);
+        InetSocketAddress addr = new InetSocketAddress("localhost", server.port());
+        selector.connect("0", addr, BUFFER_SIZE, BUFFER_SIZE);
+
+        NetworkTestUtils.waitForChannelClose(selector, "0", ChannelState.State.AUTHENTICATION_FAILED);
+        server.verifyAuthenticationMetrics(0, 1);
+    }
+
+    private void checkAuthenticationSucceed() throws IOException, InterruptedException {
+        createSelector(sslClientConfigs);
+        InetSocketAddress addr = new InetSocketAddress("localhost", server.port());
+        selector.connect("0", addr, BUFFER_SIZE, BUFFER_SIZE);
+        NetworkTestUtils.waitForChannelReady(selector, "0");
+        server.verifyAuthenticationMetrics(1, 0);
+    }
+
+    private void createSelector(Map sslClientConfigs) {
+        SslTransportLayerTest.TestSslChannelBuilder channelBuilder = new SslTransportLayerTest.TestSslChannelBuilder(Mode.CLIENT);
+        channelBuilder.configureBufferSizes(null, null, null);
+        channelBuilder.configure(sslClientConfigs);
+        this.selector = new Selector(100 * 5000, new Metrics(), TIME, "MetricGroup", channelBuilder, new LogContext());
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/network/SslVersionsTransportLayerTest.java b/clients/src/test/java/org/apache/kafka/common/network/SslVersionsTransportLayerTest.java
new file mode 100644
index 0000000000000..e972aa52cee45
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/network/SslVersionsTransportLayerTest.java
@@ -0,0 +1,183 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.network;
+
+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;
+import java.util.List;
+import java.util.Map;
+import org.apache.kafka.common.config.SslConfigs;
+import org.apache.kafka.common.metrics.Metrics;
+import org.apache.kafka.common.security.TestSecurityConfig;
+import org.apache.kafka.common.security.auth.SecurityProtocol;
+import org.apache.kafka.common.utils.Java;
+import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.test.TestUtils;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+
+/**
+ * Tests for the SSL transport layer.
+ * Checks different versions of the protocol usage on the server and client.
+ */
+@RunWith(value = Parameterized.class)
+public class SslVersionsTransportLayerTest {
+    private static final int BUFFER_SIZE = 4 * 1024;
+    private static final Time TIME = Time.SYSTEM;
+
+    private final List serverProtocols;
+    private final List clientProtocols;
+
+    @Parameterized.Parameters(name = "tlsServerProtocol={0},tlsClientProtocol={1}")
+    public static Collection data() {
+        List values = new ArrayList<>();
+
+        values.add(new Object[] {Collections.singletonList("TLSv1.2"), Collections.singletonList("TLSv1.2")});
+
+        if (Java.IS_JAVA11_COMPATIBLE) {
+            values.add(new Object[] {Collections.singletonList("TLSv1.2"), Collections.singletonList("TLSv1.3")});
+            values.add(new Object[] {Collections.singletonList("TLSv1.3"), Collections.singletonList("TLSv1.2")});
+            values.add(new Object[] {Collections.singletonList("TLSv1.3"), Collections.singletonList("TLSv1.3")});
+            values.add(new Object[] {Collections.singletonList("TLSv1.2"), Arrays.asList("TLSv1.2", "TLSv1.3")});
+            values.add(new Object[] {Collections.singletonList("TLSv1.2"), Arrays.asList("TLSv1.3", "TLSv1.2")});
+            values.add(new Object[] {Collections.singletonList("TLSv1.3"), Arrays.asList("TLSv1.2", "TLSv1.3")});
+            values.add(new Object[] {Collections.singletonList("TLSv1.3"), Arrays.asList("TLSv1.3", "TLSv1.2")});
+            values.add(new Object[] {Arrays.asList("TLSv1.3", "TLSv1.2"), Collections.singletonList("TLSv1.3")});
+            values.add(new Object[] {Arrays.asList("TLSv1.3", "TLSv1.2"), Collections.singletonList("TLSv1.2")});
+            values.add(new Object[] {Arrays.asList("TLSv1.3", "TLSv1.2"), Arrays.asList("TLSv1.2", "TLSv1.3")});
+            values.add(new Object[] {Arrays.asList("TLSv1.3", "TLSv1.2"), Arrays.asList("TLSv1.3", "TLSv1.2")});
+            values.add(new Object[] {Arrays.asList("TLSv1.2", "TLSv1.3"), Collections.singletonList("TLSv1.3")});
+            values.add(new Object[] {Arrays.asList("TLSv1.2", "TLSv1.3"), Collections.singletonList("TLSv1.2")});
+            values.add(new Object[] {Arrays.asList("TLSv1.2", "TLSv1.3"), Arrays.asList("TLSv1.2", "TLSv1.3")});
+            values.add(new Object[] {Arrays.asList("TLSv1.2", "TLSv1.3"), Arrays.asList("TLSv1.3", "TLSv1.2")});
+        }
+        return values;
+    }
+
+    /**
+     * Be aware that you can turn on debug mode for a javax.net.ssl library with the line {@code System.setProperty("javax.net.debug", "ssl:handshake");}
+     * @param serverProtocols Server protocols.
+     * @param clientProtocols Client protocols.
+     */
+    public SslVersionsTransportLayerTest(List serverProtocols, List clientProtocols) {
+        this.serverProtocols = serverProtocols;
+        this.clientProtocols = clientProtocols;
+    }
+
+    /**
+     * Tests that connection success with the default TLS version.
+     */
+    @Test
+    public void testTlsDefaults() throws Exception {
+        // Create certificates for use by client and server. Add server cert to client truststore and vice versa.
+        CertStores serverCertStores = new CertStores(true, "server",  "localhost");
+        CertStores clientCertStores = new CertStores(false, "client", "localhost");
+
+        Map sslClientConfigs = getTrustingConfig(clientCertStores, serverCertStores, clientProtocols);
+        Map sslServerConfigs = getTrustingConfig(serverCertStores, clientCertStores, serverProtocols);
+
+        NioEchoServer server = NetworkTestUtils.createEchoServer(ListenerName.forSecurityProtocol(SecurityProtocol.SSL),
+            SecurityProtocol.SSL,
+            new TestSecurityConfig(sslServerConfigs),
+            null,
+            TIME);
+        Selector selector = createClientSelector(sslClientConfigs);
+
+        String node = "0";
+        selector.connect(node, new InetSocketAddress("localhost", server.port()), BUFFER_SIZE, BUFFER_SIZE);
+
+        if (isCompatible(serverProtocols, clientProtocols)) {
+            NetworkTestUtils.waitForChannelReady(selector, node);
+
+            int msgSz = 1024 * 1024;
+            String message = TestUtils.randomString(msgSz);
+            selector.send(new NetworkSend(node, ByteBufferSend.sizePrefixed(ByteBuffer.wrap(message.getBytes()))));
+            while (selector.completedReceives().isEmpty()) {
+                selector.poll(100L);
+            }
+            int totalBytes = msgSz + 4; // including 4-byte size
+            server.waitForMetric("incoming-byte", totalBytes);
+            server.waitForMetric("outgoing-byte", totalBytes);
+            server.waitForMetric("request", 1);
+            server.waitForMetric("response", 1);
+        } else {
+            NetworkTestUtils.waitForChannelClose(selector, node, ChannelState.State.AUTHENTICATION_FAILED);
+            server.verifyAuthenticationMetrics(0, 1);
+        }
+    }
+
+    /**
+     * 
            
+     * The explanation of this check in the structure of the ClientHello SSL message.
+     * Please, take a look at the Guide,
+     * "Send ClientHello Message" section.
+     * 
            
+     * > Client version: For TLS 1.3, this has a fixed value, TLSv1.2; TLS 1.3 uses the extension supported_versions and not this field to negotiate protocol version
+     * ...
+     * > supported_versions: Lists which versions of TLS the client supports. In particular, if the client
+     * > requests TLS 1.3, then the client version field has the value TLSv1.2 and this extension
+     * > contains the value TLSv1.3; if the client requests TLS 1.2, then the client version field has the
+     * > value TLSv1.2 and this extension either doesn’t exist or contains the value TLSv1.2 but not the value TLSv1.3.
+     * 
            
+     *
+     * This mean that TLSv1.3 client can fallback to TLSv1.2 but TLSv1.2 client can't change protocol to TLSv1.3.
+     *
+     * @param serverProtocols Server protocols. Expected to be non empty.
+     * @param clientProtocols Client protocols. Expected to be non empty.
+     * @return {@code true} if client should be able to connect to the server.
+     */
+    private boolean isCompatible(List serverProtocols, List clientProtocols) {
+        assertNotNull(serverProtocols);
+        assertFalse(serverProtocols.isEmpty());
+        assertNotNull(clientProtocols);
+        assertFalse(clientProtocols.isEmpty());
+
+        return serverProtocols.contains(clientProtocols.get(0)) ||
+            (clientProtocols.get(0).equals("TLSv1.3") && !Collections.disjoint(serverProtocols, clientProtocols));
+    }
+
+    private static Map getTrustingConfig(CertStores certStores, CertStores peerCertStores, List tlsProtocols) {
+        Map configs = certStores.getTrustingConfig(peerCertStores);
+        configs.putAll(sslConfig(tlsProtocols));
+        return configs;
+    }
+
+    private static Map sslConfig(List tlsProtocols) {
+        Map sslConfig = new HashMap<>();
+        sslConfig.put(SslConfigs.SSL_PROTOCOL_CONFIG, tlsProtocols.get(0));
+        sslConfig.put(SslConfigs.SSL_ENABLED_PROTOCOLS_CONFIG, tlsProtocols);
+        return sslConfig;
+    }
+
+    private Selector createClientSelector(Map sslClientConfigs) {
+        SslTransportLayerTest.TestSslChannelBuilder channelBuilder =
+            new SslTransportLayerTest.TestSslChannelBuilder(Mode.CLIENT);
+        channelBuilder.configureBufferSizes(null, null, null);
+        channelBuilder.configure(sslClientConfigs);
+        return new Selector(100 * 5000, new Metrics(), TIME, "MetricGroup", channelBuilder, new LogContext());
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/protocol/ApiKeysTest.java b/clients/src/test/java/org/apache/kafka/common/protocol/ApiKeysTest.java
index a1dd7750159f3..5634fd202f732 100644
--- a/clients/src/test/java/org/apache/kafka/common/protocol/ApiKeysTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/protocol/ApiKeysTest.java
@@ -20,11 +20,12 @@
 import org.apache.kafka.common.protocol.types.Schema;
 import org.junit.Test;
 
-import java.util.Arrays;
-import java.util.List;
+import java.util.EnumSet;
+import java.util.Set;
 
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
 public class ApiKeysTest {
 
@@ -43,23 +44,31 @@ public void schemaVersionOutOfRange() {
         ApiKeys.PRODUCE.requestSchema((short) ApiKeys.PRODUCE.requestSchemas.length);
     }
 
+    @Test
+    public void testAlterIsrIsClusterAction() {
+        assertTrue(ApiKeys.ALTER_ISR.clusterAction);
+    }
+
     /**
      * All valid client responses which may be throttled should have a field named
      * 'throttle_time_ms' to return the throttle time to the client. Exclusions are
      * 
              
-     *   
            • Cluster actions used only for inter-broker are throttled only if unauthorized
+     *   
            •  Cluster actions used only for inter-broker are throttled only if unauthorized
      *   
            •  SASL_HANDSHAKE and SASL_AUTHENTICATE are not throttled when used for authentication
-     *        when a connection is established. At any other time, this request returns an error
-     *        response that may be throttled.
+     *        when a connection is established or for re-authentication thereafter; these requests
+     *        return an error response that may be throttled if they are sent otherwise.
      * 
            
      */
     @Test
     public void testResponseThrottleTime() {
-        List authenticationKeys = Arrays.asList(ApiKeys.SASL_HANDSHAKE, ApiKeys.SASL_AUTHENTICATE);
+        Set authenticationKeys = EnumSet.of(ApiKeys.SASL_HANDSHAKE, ApiKeys.SASL_AUTHENTICATE);
+        // Newer protocol apis include throttle time ms even for cluster actions
+        Set clusterActionsWithThrottleTimeMs = EnumSet.of(ApiKeys.ALTER_ISR);
         for (ApiKeys apiKey: ApiKeys.values()) {
             Schema responseSchema = apiKey.responseSchema(apiKey.latestVersion());
             BoundField throttleTimeField = responseSchema.get(CommonFields.THROTTLE_TIME_MS.name);
-            if (apiKey.clusterAction || authenticationKeys.contains(apiKey))
+            if ((apiKey.clusterAction && !clusterActionsWithThrottleTimeMs.contains(apiKey))
+                || authenticationKeys.contains(apiKey))
                 assertNull("Unexpected throttle time field: " + apiKey, throttleTimeField);
             else
                 assertNotNull("Throttle time field missing: " + apiKey, throttleTimeField);
diff --git a/clients/src/test/java/org/apache/kafka/common/protocol/ErrorsTest.java b/clients/src/test/java/org/apache/kafka/common/protocol/ErrorsTest.java
index e4243842b8bd7..2c68b86fb8ae4 100644
--- a/clients/src/test/java/org/apache/kafka/common/protocol/ErrorsTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/protocol/ErrorsTest.java
@@ -40,7 +40,7 @@ public void testUniqueErrorCodes() {
 
     @Test
     public void testUniqueExceptions() {
-        Set exceptionSet = new HashSet<>();
+        Set> exceptionSet = new HashSet<>();
         for (Errors error : Errors.values()) {
             if (error != Errors.NONE)
                 exceptionSet.add(error.exception().getClass());
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
new file mode 100755
index 0000000000000..178641ea5f7bf
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/protocol/MessageUtilTest.java
@@ -0,0 +1,81 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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 org.apache.kafka.common.protocol.types.RawTaggedField;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.Timeout;
+
+import java.nio.ByteBuffer;
+import java.util.Arrays;
+import java.util.Collections;
+
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+public final class MessageUtilTest {
+    @Rule
+    final public Timeout globalTimeout = Timeout.millis(120000);
+
+    @Test
+    public void testDeepToString() {
+        assertEquals("[1, 2, 3]",
+            MessageUtil.deepToString(Arrays.asList(1, 2, 3).iterator()));
+        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[]{})));
+    }
+
+    @Test
+    public void testDuplicate() {
+        assertEquals(null, MessageUtil.duplicate(null));
+        assertArrayEquals(new byte[] {},
+            MessageUtil.duplicate(new byte[] {}));
+        assertArrayEquals(new byte[] {1, 2, 3},
+            MessageUtil.duplicate(new byte[] {1, 2, 3}));
+    }
+
+    @Test
+    public void testCompareRawTaggedFields() {
+        assertTrue(MessageUtil.compareRawTaggedFields(null, null));
+        assertTrue(MessageUtil.compareRawTaggedFields(null, Collections.emptyList()));
+        assertTrue(MessageUtil.compareRawTaggedFields(Collections.emptyList(), null));
+        assertFalse(MessageUtil.compareRawTaggedFields(Collections.emptyList(),
+            Collections.singletonList(new RawTaggedField(1, new byte[] {1}))));
+        assertFalse(MessageUtil.compareRawTaggedFields(null,
+            Collections.singletonList(new RawTaggedField(1, new byte[] {1}))));
+        assertFalse(MessageUtil.compareRawTaggedFields(
+            Collections.singletonList(new RawTaggedField(1, new byte[] {1})),
+            Collections.emptyList()));
+        assertTrue(MessageUtil.compareRawTaggedFields(
+            Arrays.asList(new RawTaggedField(1, new byte[] {1}),
+                new RawTaggedField(2, new byte[] {})),
+            Arrays.asList(new RawTaggedField(1, new byte[] {1}),
+                new RawTaggedField(2, new byte[] {}))));
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/protocol/ProtoUtilsTest.java b/clients/src/test/java/org/apache/kafka/common/protocol/ProtoUtilsTest.java
index c4b3fc406af0f..613f07973857b 100644
--- a/clients/src/test/java/org/apache/kafka/common/protocol/ProtoUtilsTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/protocol/ProtoUtilsTest.java
@@ -23,13 +23,29 @@
 
 public class ProtoUtilsTest {
     @Test
-    public void testDelayedAllocationSchemaDetection() throws Exception {
+    public void testDelayedAllocationSchemaDetection() {
         //verifies that schemas known to retain a reference to the underlying byte buffer are correctly detected.
         for (ApiKeys key : ApiKeys.values()) {
-            if (key == ApiKeys.PRODUCE || key == ApiKeys.JOIN_GROUP || key == ApiKeys.SYNC_GROUP || key == ApiKeys.SASL_AUTHENTICATE) {
-                assertTrue(key + " should require delayed allocation", key.requiresDelayedAllocation);
-            } else {
-                assertFalse(key + " should not require delayed allocation", key.requiresDelayedAllocation);
+            switch (key) {
+                case PRODUCE:
+                case JOIN_GROUP:
+                case SYNC_GROUP:
+                case SASL_AUTHENTICATE:
+                case EXPIRE_DELEGATION_TOKEN:
+                case RENEW_DELEGATION_TOKEN:
+                case ALTER_USER_SCRAM_CREDENTIALS:
+                case ENVELOPE:
+                    assertTrue(key + " should require delayed allocation", key.requiresDelayedAllocation);
+                    break;
+                default:
+                    if (key.forwardable) {
+                        assertTrue(key + " should require delayed allocation since it is forwardable",
+                            key.requiresDelayedAllocation);
+                    } else {
+                        assertFalse(key + " should not require delayed allocation",
+                            key.requiresDelayedAllocation);
+                    }
+                    break;
             }
         }
     }
diff --git a/clients/src/test/java/org/apache/kafka/common/protocol/SendBuilderTest.java b/clients/src/test/java/org/apache/kafka/common/protocol/SendBuilderTest.java
new file mode 100644
index 0000000000000..39aa0945a6a72
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/protocol/SendBuilderTest.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.protocol;
+
+import org.apache.kafka.common.network.Send;
+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.SimpleRecord;
+import org.apache.kafka.common.record.TimestampType;
+import org.apache.kafka.common.utils.Utils;
+import org.apache.kafka.test.TestUtils;
+import org.junit.Test;
+
+import java.nio.ByteBuffer;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+public class SendBuilderTest {
+
+    @Test
+    public void testZeroCopyByteBuffer() {
+        byte[] data = Utils.utf8("foo");
+        ByteBuffer zeroCopyBuffer = ByteBuffer.wrap(data);
+        SendBuilder builder = new SendBuilder(8);
+
+        builder.writeInt(5);
+        builder.writeByteBuffer(zeroCopyBuffer);
+        builder.writeInt(15);
+        Send send = builder.build();
+
+        // Overwrite the original buffer in order to prove the data was not copied
+        byte[] overwrittenData = Utils.utf8("bar");
+        assertEquals(data.length, overwrittenData.length);
+        zeroCopyBuffer.rewind();
+        zeroCopyBuffer.put(overwrittenData);
+        zeroCopyBuffer.rewind();
+
+        ByteBuffer buffer = TestUtils.toBuffer(send);
+        assertEquals(8 + data.length, buffer.remaining());
+        assertEquals(5, buffer.getInt());
+        assertEquals("bar", getString(buffer, data.length));
+        assertEquals(15, buffer.getInt());
+    }
+
+    @Test
+    public void testWriteByteBufferRespectsPosition() {
+        byte[] data = Utils.utf8("yolo");
+        assertEquals(4, data.length);
+
+        ByteBuffer buffer = ByteBuffer.wrap(data);
+        SendBuilder builder = new SendBuilder(0);
+
+        buffer.limit(2);
+        builder.writeByteBuffer(buffer);
+        assertEquals(0, buffer.position());
+
+        buffer.position(2);
+        buffer.limit(4);
+        builder.writeByteBuffer(buffer);
+        assertEquals(2, buffer.position());
+
+        Send send = builder.build();
+        ByteBuffer readBuffer = TestUtils.toBuffer(send);
+        assertEquals("yolo", getString(readBuffer, 4));
+    }
+
+    @Test
+    public void testZeroCopyRecords() {
+        ByteBuffer buffer = ByteBuffer.allocate(128);
+        MemoryRecords records = createRecords(buffer, "foo");
+
+        SendBuilder builder = new SendBuilder(8);
+        builder.writeInt(5);
+        builder.writeRecords(records);
+        builder.writeInt(15);
+        Send send = builder.build();
+
+        // Overwrite the original buffer in order to prove the data was not copied
+        buffer.rewind();
+        MemoryRecords overwrittenRecords = createRecords(buffer, "bar");
+
+        ByteBuffer readBuffer = TestUtils.toBuffer(send);
+        assertEquals(5, readBuffer.getInt());
+        assertEquals(overwrittenRecords, getRecords(readBuffer, records.sizeInBytes()));
+        assertEquals(15, readBuffer.getInt());
+    }
+
+    private String getString(ByteBuffer buffer, int size) {
+        byte[] readData = new byte[size];
+        buffer.get(readData);
+        return Utils.utf8(readData);
+    }
+
+    private MemoryRecords getRecords(ByteBuffer buffer, int size) {
+        int initialPosition = buffer.position();
+        int initialLimit = buffer.limit();
+        int recordsLimit = initialPosition + size;
+
+        buffer.limit(recordsLimit);
+        MemoryRecords records = MemoryRecords.readableRecords(buffer.slice());
+
+        buffer.position(recordsLimit);
+        buffer.limit(initialLimit);
+        return records;
+    }
+
+    private MemoryRecords createRecords(ByteBuffer buffer, String value) {
+        MemoryRecordsBuilder recordsBuilder = MemoryRecords.builder(
+            buffer,
+            CompressionType.NONE,
+            TimestampType.CREATE_TIME,
+            0L
+        );
+        recordsBuilder.append(new SimpleRecord(Utils.utf8(value)));
+        return recordsBuilder.build();
+    }
+
+}
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..dfce6aad53604 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 {
@@ -42,12 +44,19 @@ public void setup() {
                                  new Field("int64", Type.INT64),
                                  new Field("varint", Type.VARINT),
                                  new Field("varlong", Type.VARLONG),
+                                 new Field("float64", Type.FLOAT64),
                                  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)
@@ -56,42 +65,81 @@ public void setup() {
                                              .set("int64", 1L)
                                              .set("varint", 300)
                                              .set("varlong", 500L)
+                                             .set("float64", 0.5D)
                                              .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.FLOAT64, 2.5, "FLOAT64");
+        check(Type.FLOAT64, -0.5, "FLOAT64");
+        check(Type.FLOAT64, 1e300, "FLOAT64");
+        check(Type.FLOAT64, 0.0, "FLOAT64");
+        check(Type.FLOAT64, -0.0, "FLOAT64");
+        check(Type.FLOAT64, Double.MAX_VALUE, "FLOAT64");
+        check(Type.FLOAT64, Double.MIN_VALUE, "FLOAT64");
+        check(Type.FLOAT64, Double.NaN, "FLOAT64");
+        check(Type.FLOAT64, Double.NEGATIVE_INFINITY, "FLOAT64");
+        check(Type.FLOAT64, Double.POSITIVE_INFINITY, "FLOAT64");
+        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 +152,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 +170,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 +200,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 +235,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 +344,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 +365,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..12e561c4f53c0
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/protocol/types/StructTest.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.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.Float64("float64", ""),
+        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("float64", 0.5)
+                .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("float64", 0.5)
+                .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("float64", 0.5)
+                .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/protocol/types/TypeTest.java b/clients/src/test/java/org/apache/kafka/common/protocol/types/TypeTest.java
new file mode 100644
index 0000000000000..08a912f8fecaf
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/protocol/types/TypeTest.java
@@ -0,0 +1,93 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.record.CompressionType;
+import org.apache.kafka.common.record.MemoryRecords;
+import org.apache.kafka.common.record.SimpleRecord;
+import org.junit.Test;
+
+import java.nio.ByteBuffer;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+
+public class TypeTest {
+
+    @Test
+    public void testEmptyRecordsSerde() {
+        ByteBuffer buffer = ByteBuffer.allocate(4);
+        Type.RECORDS.write(buffer, MemoryRecords.EMPTY);
+        buffer.flip();
+        assertEquals(4, Type.RECORDS.sizeOf(MemoryRecords.EMPTY));
+        assertEquals(4, buffer.limit());
+        assertEquals(MemoryRecords.EMPTY, Type.RECORDS.read(buffer));
+    }
+
+    @Test
+    public void testNullRecordsSerde() {
+        ByteBuffer buffer = ByteBuffer.allocate(4);
+        Type.RECORDS.write(buffer, null);
+        buffer.flip();
+        assertEquals(4, Type.RECORDS.sizeOf(MemoryRecords.EMPTY));
+        assertEquals(4, buffer.limit());
+        assertNull(Type.RECORDS.read(buffer));
+    }
+
+    @Test
+    public void testRecordsSerde() {
+        MemoryRecords records = MemoryRecords.withRecords(CompressionType.NONE,
+            new SimpleRecord("foo".getBytes()),
+            new SimpleRecord("bar".getBytes()));
+        ByteBuffer buffer = ByteBuffer.allocate(Type.RECORDS.sizeOf(records));
+        Type.RECORDS.write(buffer, records);
+        buffer.flip();
+        assertEquals(records, Type.RECORDS.read(buffer));
+    }
+
+    @Test
+    public void testEmptyCompactRecordsSerde() {
+        ByteBuffer buffer = ByteBuffer.allocate(4);
+        Type.COMPACT_RECORDS.write(buffer, MemoryRecords.EMPTY);
+        buffer.flip();
+        assertEquals(1, Type.COMPACT_RECORDS.sizeOf(MemoryRecords.EMPTY));
+        assertEquals(1, buffer.limit());
+        assertEquals(MemoryRecords.EMPTY, Type.COMPACT_RECORDS.read(buffer));
+    }
+
+    @Test
+    public void testNullCompactRecordsSerde() {
+        ByteBuffer buffer = ByteBuffer.allocate(4);
+        Type.COMPACT_RECORDS.write(buffer, null);
+        buffer.flip();
+        assertEquals(1, Type.COMPACT_RECORDS.sizeOf(MemoryRecords.EMPTY));
+        assertEquals(1, buffer.limit());
+        assertNull(Type.COMPACT_RECORDS.read(buffer));
+    }
+
+    @Test
+    public void testCompactRecordsSerde() {
+        MemoryRecords records = MemoryRecords.withRecords(CompressionType.NONE,
+            new SimpleRecord("foo".getBytes()),
+            new SimpleRecord("bar".getBytes()));
+        ByteBuffer buffer = ByteBuffer.allocate(Type.COMPACT_RECORDS.sizeOf(records));
+        Type.COMPACT_RECORDS.write(buffer, records);
+        buffer.flip();
+        assertEquals(records, Type.COMPACT_RECORDS.read(buffer));
+    }
+
+}
\ No newline at end of file
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 83ada730ce24f..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;
@@ -25,6 +26,7 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
 
 public class AbstractLegacyRecordBatchTest {
 
@@ -208,4 +210,41 @@ public void testSetCreateTimeV1() {
             assertEquals(expectedTimestamp++, record.timestamp());
     }
 
+    @Test
+    public void testZStdCompressionTypeWithV0OrV1() {
+        SimpleRecord[] simpleRecords = new SimpleRecord[] {
+            new SimpleRecord(1L, "a".getBytes(), "1".getBytes()),
+            new SimpleRecord(2L, "b".getBytes(), "2".getBytes()),
+            new SimpleRecord(3L, "c".getBytes(), "3".getBytes())
+        };
+
+        // Check V0
+        try {
+            MemoryRecords records = MemoryRecords.withRecords(RecordBatch.MAGIC_VALUE_V0, 0L,
+                CompressionType.ZSTD, TimestampType.CREATE_TIME, simpleRecords);
+
+            ByteBufferLegacyRecordBatch batch = new ByteBufferLegacyRecordBatch(records.buffer());
+            batch.setLastOffset(1L);
+
+            batch.iterator();
+            fail("Can't reach here");
+        } catch (IllegalArgumentException e) {
+            assertEquals("ZStandard compression is not supported for magic 0", e.getMessage());
+        }
+
+        // Check V1
+        try {
+            MemoryRecords records = MemoryRecords.withRecords(RecordBatch.MAGIC_VALUE_V1, 0L,
+                CompressionType.ZSTD, TimestampType.CREATE_TIME, simpleRecords);
+
+            ByteBufferLegacyRecordBatch batch = new ByteBufferLegacyRecordBatch(records.buffer());
+            batch.setLastOffset(1L);
+
+            batch.iterator();
+            fail("Can't reach here");
+        } catch (IllegalArgumentException e) {
+            assertEquals("ZStandard compression is not supported for magic 1", e.getMessage());
+        }
+    }
+
 }
diff --git a/clients/src/test/java/org/apache/kafka/common/record/BufferSupplierTest.java b/clients/src/test/java/org/apache/kafka/common/record/BufferSupplierTest.java
new file mode 100644
index 0000000000000..dea0c9854133e
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/record/BufferSupplierTest.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.common.record;
+
+import org.junit.Test;
+
+import java.nio.ByteBuffer;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertSame;
+
+public class BufferSupplierTest {
+
+    @Test
+    public void testGrowableBuffer() {
+        BufferSupplier.GrowableBufferSupplier supplier = new BufferSupplier.GrowableBufferSupplier();
+        ByteBuffer buffer = supplier.get(1024);
+        assertEquals(0, buffer.position());
+        assertEquals(1024, buffer.capacity());
+        supplier.release(buffer);
+
+        ByteBuffer cached = supplier.get(512);
+        assertEquals(0, cached.position());
+        assertSame(buffer, cached);
+
+        ByteBuffer increased = supplier.get(2048);
+        assertEquals(2048, increased.capacity());
+        assertEquals(0, increased.position());
+    }
+
+}
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/CompressionRatioEstimatorTest.java b/clients/src/test/java/org/apache/kafka/common/record/CompressionRatioEstimatorTest.java
new file mode 100644
index 0000000000000..4a2317c864847
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/record/CompressionRatioEstimatorTest.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 org.apache.kafka.common.record;
+
+import org.junit.Test;
+import static org.junit.Assert.assertTrue;
+import java.util.Arrays;
+import java.util.List;
+
+public class CompressionRatioEstimatorTest {
+
+    @Test
+    public void testUpdateEstimation() {
+        class EstimationsObservedRatios {
+            float currentEstimation;
+            float observedRatio;
+            EstimationsObservedRatios(float currentEstimation, float observedRatio) {
+                this.currentEstimation = currentEstimation;
+                this.observedRatio = observedRatio;
+            }
+        }
+
+        // If currentEstimation is smaller than observedRatio, the updatedCompressionRatio is currentEstimation plus
+        // COMPRESSION_RATIO_DETERIORATE_STEP 0.05, otherwise currentEstimation minus COMPRESSION_RATIO_IMPROVING_STEP
+        // 0.005. There are four cases,and updatedCompressionRatio shouldn't smaller than observedRatio in all of cases.
+        // Refer to non test code for more details.
+        List estimationsObservedRatios = Arrays.asList(
+            new EstimationsObservedRatios(0.8f, 0.84f),
+            new EstimationsObservedRatios(0.6f, 0.7f),
+            new EstimationsObservedRatios(0.6f, 0.4f),
+            new EstimationsObservedRatios(0.004f, 0.001f));
+        for (EstimationsObservedRatios estimationsObservedRatio : estimationsObservedRatios) {
+            String topic = "tp";
+            CompressionRatioEstimator.setEstimation(topic, CompressionType.ZSTD, estimationsObservedRatio.currentEstimation);
+            float updatedCompressionRatio = CompressionRatioEstimator.updateEstimation(topic, CompressionType.ZSTD, estimationsObservedRatio.observedRatio);
+            assertTrue(updatedCompressionRatio >= estimationsObservedRatio.observedRatio);
+        }
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/record/ControlRecordUtilsTest.java b/clients/src/test/java/org/apache/kafka/common/record/ControlRecordUtilsTest.java
new file mode 100644
index 0000000000000..f89366c7a3c6d
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/record/ControlRecordUtilsTest.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.record;
+
+import org.apache.kafka.common.message.LeaderChangeMessage;
+import org.apache.kafka.common.message.LeaderChangeMessage.Voter;
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.nio.ByteBuffer;
+import java.util.Collections;
+
+public class ControlRecordUtilsTest {
+
+    @Test
+    public void testInvalidControlRecordType() {
+        IllegalArgumentException thrown = Assert.assertThrows(
+            IllegalArgumentException.class, () -> testDeserializeRecord(ControlRecordType.COMMIT));
+        Assert.assertEquals("Expected LEADER_CHANGE control record type(3), but found COMMIT", thrown.getMessage());
+    }
+
+    @Test
+    public void testDeserializeByteData() {
+        testDeserializeRecord(ControlRecordType.LEADER_CHANGE);
+    }
+
+    private void testDeserializeRecord(ControlRecordType controlRecordType) {
+        final int leaderId = 1;
+        final int voterId = 2;
+        LeaderChangeMessage data = new LeaderChangeMessage()
+                                           .setLeaderId(leaderId)
+                                           .setVoters(Collections.singletonList(
+                                               new Voter().setVoterId(voterId)));
+
+        ByteBuffer valueBuffer = ByteBuffer.allocate(256);
+        data.toStruct(data.highestSupportedVersion()).writeTo(valueBuffer);
+        valueBuffer.flip();
+
+        byte[] keyData = new byte[]{0, 0, 0, (byte) controlRecordType.type};
+
+        DefaultRecord record = new DefaultRecord(
+            256, (byte) 0, 0, 0L, 0, ByteBuffer.wrap(keyData),  valueBuffer, null
+        );
+
+        LeaderChangeMessage deserializedData = ControlRecordUtils.deserializeLeaderChangeMessage(record);
+
+        Assert.assertEquals(leaderId, deserializedData.leaderId());
+        Assert.assertEquals(Collections.singletonList(
+            new Voter().setVoterId(voterId)), deserializedData.voters());
+    }
+}
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..11854dc4e0773 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 = InvalidRecordException.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 d5de4bd9b5208..783a5b531ef27 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
@@ -56,12 +56,14 @@ public FileLogInputStreamTest(byte magic, CompressionType compression) {
 
     @Test
     public void testWriteTo() throws IOException {
+        if (compression == CompressionType.ZSTD && magic < MAGIC_VALUE_V2)
+            return;
+
         try (FileRecords fileRecords = FileRecords.open(tempFile())) {
             fileRecords.append(MemoryRecords.withRecords(magic, compression, new SimpleRecord("foo".getBytes())));
             fileRecords.flush();
 
-            FileLogInputStream logInputStream = new FileLogInputStream(fileRecords.channel(), 0,
-                    fileRecords.sizeInBytes());
+            FileLogInputStream logInputStream = new FileLogInputStream(fileRecords, 0, fileRecords.sizeInBytes());
 
             FileChannelRecordBatch batch = logInputStream.nextBatch();
             assertNotNull(batch);
@@ -82,6 +84,9 @@ public void testWriteTo() throws IOException {
 
     @Test
     public void testSimpleBatchIteration() throws IOException {
+        if (compression == CompressionType.ZSTD && magic < MAGIC_VALUE_V2)
+            return;
+
         try (FileRecords fileRecords = FileRecords.open(tempFile())) {
             SimpleRecord firstBatchRecord = new SimpleRecord(3241324L, "a".getBytes(), "foo".getBytes());
             SimpleRecord secondBatchRecord = new SimpleRecord(234280L, "b".getBytes(), "bar".getBytes());
@@ -90,8 +95,7 @@ public void testSimpleBatchIteration() throws IOException {
             fileRecords.append(MemoryRecords.withRecords(magic, 1L, compression, CREATE_TIME, secondBatchRecord));
             fileRecords.flush();
 
-            FileLogInputStream logInputStream = new FileLogInputStream(fileRecords.channel(), 0,
-                    fileRecords.sizeInBytes());
+            FileLogInputStream logInputStream = new FileLogInputStream(fileRecords, 0, fileRecords.sizeInBytes());
 
             FileChannelRecordBatch firstBatch = logInputStream.nextBatch();
             assertGenericRecordBatchData(firstBatch, 0L, 3241324L, firstBatchRecord);
@@ -110,12 +114,15 @@ public void testBatchIterationWithMultipleRecordsPerBatch() throws IOException {
         if (magic < MAGIC_VALUE_V2 && compression == CompressionType.NONE)
             return;
 
+        if (compression == CompressionType.ZSTD && magic < MAGIC_VALUE_V2)
+            return;
+
         try (FileRecords fileRecords = FileRecords.open(tempFile())) {
             SimpleRecord[] firstBatchRecords = new SimpleRecord[]{
                 new SimpleRecord(3241324L, "a".getBytes(), "1".getBytes()),
                 new SimpleRecord(234280L, "b".getBytes(), "2".getBytes())
-
             };
+
             SimpleRecord[] secondBatchRecords = new SimpleRecord[]{
                 new SimpleRecord(238423489L, "c".getBytes(), "3".getBytes()),
                 new SimpleRecord(897839L, null, "4".getBytes()),
@@ -126,8 +133,7 @@ public void testBatchIterationWithMultipleRecordsPerBatch() throws IOException {
             fileRecords.append(MemoryRecords.withRecords(magic, 1L, compression, CREATE_TIME, secondBatchRecords));
             fileRecords.flush();
 
-            FileLogInputStream logInputStream = new FileLogInputStream(fileRecords.channel(), 0,
-                    fileRecords.sizeInBytes());
+            FileLogInputStream logInputStream = new FileLogInputStream(fileRecords, 0, fileRecords.sizeInBytes());
 
             FileChannelRecordBatch firstBatch = logInputStream.nextBatch();
             assertNoProducerData(firstBatch);
@@ -155,8 +161,8 @@ public void testBatchIterationV2() throws IOException {
             SimpleRecord[] firstBatchRecords = new SimpleRecord[]{
                 new SimpleRecord(3241324L, "a".getBytes(), "1".getBytes()),
                 new SimpleRecord(234280L, "b".getBytes(), "2".getBytes())
-
             };
+
             SimpleRecord[] secondBatchRecords = new SimpleRecord[]{
                 new SimpleRecord(238423489L, "c".getBytes(), "3".getBytes()),
                 new SimpleRecord(897839L, null, "4".getBytes()),
@@ -169,8 +175,7 @@ public void testBatchIterationV2() throws IOException {
                     producerEpoch, baseSequence + firstBatchRecords.length, partitionLeaderEpoch, secondBatchRecords));
             fileRecords.flush();
 
-            FileLogInputStream logInputStream = new FileLogInputStream(fileRecords.channel(), 0,
-                    fileRecords.sizeInBytes());
+            FileLogInputStream logInputStream = new FileLogInputStream(fileRecords, 0, fileRecords.sizeInBytes());
 
             FileChannelRecordBatch firstBatch = logInputStream.nextBatch();
             assertProducerData(firstBatch, producerId, producerEpoch, baseSequence, false, firstBatchRecords);
@@ -189,6 +194,9 @@ public void testBatchIterationV2() throws IOException {
 
     @Test
     public void testBatchIterationIncompleteBatch() throws IOException {
+        if (compression == CompressionType.ZSTD && magic < MAGIC_VALUE_V2)
+            return;
+
         try (FileRecords fileRecords = FileRecords.open(tempFile())) {
             SimpleRecord firstBatchRecord = new SimpleRecord(100L, "foo".getBytes());
             SimpleRecord secondBatchRecord = new SimpleRecord(200L, "bar".getBytes());
@@ -198,8 +206,7 @@ public void testBatchIterationIncompleteBatch() throws IOException {
             fileRecords.flush();
             fileRecords.truncateTo(fileRecords.sizeInBytes() - 13);
 
-            FileLogInputStream logInputStream = new FileLogInputStream(fileRecords.channel(), 0,
-                    fileRecords.sizeInBytes());
+            FileLogInputStream logInputStream = new FileLogInputStream(fileRecords, 0, fileRecords.sizeInBytes());
 
             FileChannelRecordBatch firstBatch = logInputStream.nextBatch();
             assertNoProducerData(firstBatch);
@@ -209,8 +216,24 @@ public void testBatchIterationIncompleteBatch() throws IOException {
         }
     }
 
+    @Test
+    public void testNextBatchSelectionWithMaxedParams() throws IOException {
+        try (FileRecords fileRecords = FileRecords.open(tempFile())) {
+            FileLogInputStream logInputStream = new FileLogInputStream(fileRecords, Integer.MAX_VALUE, Integer.MAX_VALUE);
+            assertNull(logInputStream.nextBatch());
+        }
+    }
+
+    @Test
+    public void testNextBatchSelectionWithZeroedParams() throws IOException {
+        try (FileRecords fileRecords = FileRecords.open(tempFile())) {
+            FileLogInputStream logInputStream = new FileLogInputStream(fileRecords, 0, 0);
+            assertNull(logInputStream.nextBatch());
+        }
+    }
+
     private void assertProducerData(RecordBatch batch, long producerId, short producerEpoch, int baseSequence,
-                                    boolean isTransactional, SimpleRecord ... records) {
+                                    boolean isTransactional, SimpleRecord... records) {
         assertEquals(producerId, batch.producerId());
         assertEquals(producerEpoch, batch.producerEpoch());
         assertEquals(baseSequence, batch.baseSequence());
@@ -226,7 +249,7 @@ private void assertNoProducerData(RecordBatch batch) {
         assertFalse(batch.isTransactional());
     }
 
-    private void assertGenericRecordBatchData(RecordBatch batch, long baseOffset, long maxTimestamp, SimpleRecord ... records) {
+    private void assertGenericRecordBatchData(RecordBatch batch, long baseOffset, long maxTimestamp, SimpleRecord... records) {
         assertEquals(magic, batch.magic());
         assertEquals(compression, batch.compressionType());
 
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 53ac200358626..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
@@ -17,11 +17,12 @@
 package org.apache.kafka.common.record;
 
 import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.header.Header;
+import org.apache.kafka.common.header.internals.RecordHeader;
 import org.apache.kafka.common.utils.MockTime;
 import org.apache.kafka.common.utils.Time;
-import org.apache.kafka.common.utils.Utils;
 import org.apache.kafka.test.TestUtils;
-import org.easymock.EasyMock;
 import org.junit.Before;
 import org.junit.Test;
 
@@ -33,13 +34,28 @@
 import java.util.Collections;
 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;
 import static org.apache.kafka.test.TestUtils.tempFile;
+import static org.junit.Assert.assertArrayEquals;
 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.junit.Assert.fail;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.Mockito.atLeastOnce;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 public class FileRecordsTest {
 
@@ -57,6 +73,30 @@ public void setup() throws IOException {
         this.time = new MockTime();
     }
 
+    @Test(expected = IllegalArgumentException.class)
+    public void testAppendProtectsFromOverflow() throws Exception {
+        File fileMock = mock(File.class);
+        FileChannel fileChannelMock = mock(FileChannel.class);
+        when(fileChannelMock.size()).thenReturn((long) Integer.MAX_VALUE);
+
+        FileRecords records = new FileRecords(fileMock, fileChannelMock, 0, Integer.MAX_VALUE, false);
+        append(records, values);
+    }
+
+    @Test(expected = KafkaException.class)
+    public void testOpenOversizeFile() throws Exception {
+        File fileMock = mock(File.class);
+        FileChannel fileChannelMock = mock(FileChannel.class);
+        when(fileChannelMock.size()).thenReturn(Integer.MAX_VALUE + 5L);
+
+        new FileRecords(fileMock, fileChannelMock, 0, Integer.MAX_VALUE, false);
+    }
+
+    @Test(expected = IllegalArgumentException.class)
+    public void testOutOfRangeSlice() throws Exception {
+        this.fileRecords.slice(fileRecords.sizeInBytes() + 1, 15).sizeInBytes();
+    }
+
     /**
      * Test that the cached size variable matches the actual file size as we append messages
      */
@@ -81,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++)
@@ -117,7 +187,7 @@ public void testIterationDoesntChangePosition() throws IOException {
      */
     @Test
     public void testRead() throws IOException {
-        FileRecords read = fileRecords.read(0, fileRecords.sizeInBytes());
+        FileRecords read = fileRecords.slice(0, fileRecords.sizeInBytes());
         assertEquals(fileRecords.sizeInBytes(), read.sizeInBytes());
         TestUtils.checkEquals(fileRecords.batches(), read.batches());
 
@@ -125,35 +195,35 @@ public void testRead() throws IOException {
         RecordBatch first = items.get(0);
 
         // read from second message until the end
-        read = fileRecords.read(first.sizeInBytes(), fileRecords.sizeInBytes() - first.sizeInBytes());
+        read = fileRecords.slice(first.sizeInBytes(), fileRecords.sizeInBytes() - first.sizeInBytes());
         assertEquals(fileRecords.sizeInBytes() - first.sizeInBytes(), read.sizeInBytes());
         assertEquals("Read starting from the second message", items.subList(1, items.size()), batches(read));
 
         // read from second message and size is past the end of the file
-        read = fileRecords.read(first.sizeInBytes(), fileRecords.sizeInBytes());
+        read = fileRecords.slice(first.sizeInBytes(), fileRecords.sizeInBytes());
         assertEquals(fileRecords.sizeInBytes() - first.sizeInBytes(), read.sizeInBytes());
         assertEquals("Read starting from the second message", items.subList(1, items.size()), batches(read));
 
         // read from second message and position + size overflows
-        read = fileRecords.read(first.sizeInBytes(), Integer.MAX_VALUE);
+        read = fileRecords.slice(first.sizeInBytes(), Integer.MAX_VALUE);
         assertEquals(fileRecords.sizeInBytes() - first.sizeInBytes(), read.sizeInBytes());
         assertEquals("Read starting from the second message", items.subList(1, items.size()), batches(read));
 
         // read from second message and size is past the end of the file on a view/slice
-        read = fileRecords.read(1, fileRecords.sizeInBytes() - 1)
-                .read(first.sizeInBytes() - 1, fileRecords.sizeInBytes());
+        read = fileRecords.slice(1, fileRecords.sizeInBytes() - 1)
+                .slice(first.sizeInBytes() - 1, fileRecords.sizeInBytes());
         assertEquals(fileRecords.sizeInBytes() - first.sizeInBytes(), read.sizeInBytes());
         assertEquals("Read starting from the second message", items.subList(1, items.size()), batches(read));
 
         // read from second message and position + size overflows on a view/slice
-        read = fileRecords.read(1, fileRecords.sizeInBytes() - 1)
-                .read(first.sizeInBytes() - 1, Integer.MAX_VALUE);
+        read = fileRecords.slice(1, fileRecords.sizeInBytes() - 1)
+                .slice(first.sizeInBytes() - 1, Integer.MAX_VALUE);
         assertEquals(fileRecords.sizeInBytes() - first.sizeInBytes(), read.sizeInBytes());
         assertEquals("Read starting from the second message", items.subList(1, items.size()), batches(read));
 
         // read a single message starting from second message
         RecordBatch second = items.get(1);
-        read = fileRecords.read(first.sizeInBytes(), second.sizeInBytes());
+        read = fileRecords.slice(first.sizeInBytes(), second.sizeInBytes());
         assertEquals(second.sizeInBytes(), read.sizeInBytes());
         assertEquals("Read a single message starting from the second message",
                 Collections.singletonList(second), batches(read));
@@ -187,7 +257,7 @@ public void testSearch() throws IOException {
         position += message2Size + batches.get(2).sizeInBytes();
 
         int message4Size = batches.get(3).sizeInBytes();
-        assertEquals("Should be able to find fourth message from a non-existant offset",
+        assertEquals("Should be able to find fourth message from a non-existent offset",
                 new FileRecords.LogOffsetPosition(50L, position, message4Size),
                 fileRecords.searchForOffsetWithSize(3, position));
         assertEquals("Should be able to find fourth message by correct offset",
@@ -203,9 +273,9 @@ public void testIteratorWithLimits() throws IOException {
         RecordBatch batch = batches(fileRecords).get(1);
         int start = fileRecords.searchForOffsetWithSize(1, 0).position;
         int size = batch.sizeInBytes();
-        FileRecords slice = fileRecords.read(start, size);
+        FileRecords slice = fileRecords.slice(start, size);
         assertEquals(Collections.singletonList(batch), batches(slice));
-        FileRecords slice2 = fileRecords.read(start, size - 1);
+        FileRecords slice2 = fileRecords.slice(start, size - 1);
         assertEquals(Collections.emptyList(), batches(slice2));
     }
 
@@ -228,16 +298,16 @@ public void testTruncate() throws IOException {
      */
     @Test
     public void testTruncateNotCalledIfSizeIsSameAsTargetSize() throws IOException {
-        FileChannel channelMock = EasyMock.createMock(FileChannel.class);
+        FileChannel channelMock = mock(FileChannel.class);
 
-        EasyMock.expect(channelMock.size()).andReturn(42L).atLeastOnce();
-        EasyMock.expect(channelMock.position(42L)).andReturn(null);
-        EasyMock.replay(channelMock);
+        when(channelMock.size()).thenReturn(42L);
+        when(channelMock.position(42L)).thenReturn(null);
 
         FileRecords fileRecords = new FileRecords(tempFile(), channelMock, 0, Integer.MAX_VALUE, false);
         fileRecords.truncateTo(42);
 
-        EasyMock.verify(channelMock);
+        verify(channelMock, atLeastOnce()).size();
+        verify(channelMock, times(0)).truncate(anyLong());
     }
 
     /**
@@ -246,11 +316,9 @@ public void testTruncateNotCalledIfSizeIsSameAsTargetSize() throws IOException {
      */
     @Test
     public void testTruncateNotCalledIfSizeIsBiggerThanTargetSize() throws IOException {
-        FileChannel channelMock = EasyMock.createMock(FileChannel.class);
+        FileChannel channelMock = mock(FileChannel.class);
 
-        EasyMock.expect(channelMock.size()).andReturn(42L).atLeastOnce();
-        EasyMock.expect(channelMock.position(42L)).andReturn(null);
-        EasyMock.replay(channelMock);
+        when(channelMock.size()).thenReturn(42L);
 
         FileRecords fileRecords = new FileRecords(tempFile(), channelMock, 0, Integer.MAX_VALUE, false);
 
@@ -261,7 +329,7 @@ public void testTruncateNotCalledIfSizeIsBiggerThanTargetSize() throws IOExcepti
             // expected
         }
 
-        EasyMock.verify(channelMock);
+        verify(channelMock, atLeastOnce()).size();
     }
 
     /**
@@ -269,17 +337,16 @@ public void testTruncateNotCalledIfSizeIsBiggerThanTargetSize() throws IOExcepti
      */
     @Test
     public void testTruncateIfSizeIsDifferentToTargetSize() throws IOException {
-        FileChannel channelMock = EasyMock.createMock(FileChannel.class);
+        FileChannel channelMock = mock(FileChannel.class);
 
-        EasyMock.expect(channelMock.size()).andReturn(42L).atLeastOnce();
-        EasyMock.expect(channelMock.position(42L)).andReturn(null).once();
-        EasyMock.expect(channelMock.truncate(23L)).andReturn(null).once();
-        EasyMock.replay(channelMock);
+        when(channelMock.size()).thenReturn(42L);
+        when(channelMock.truncate(anyLong())).thenReturn(channelMock);
 
         FileRecords fileRecords = new FileRecords(tempFile(), channelMock, 0, Integer.MAX_VALUE, false);
         fileRecords.truncateTo(23);
 
-        EasyMock.verify(channelMock);
+        verify(channelMock, atLeastOnce()).size();
+        verify(channelMock).truncate(23);
     }
 
     /**
@@ -288,12 +355,12 @@ public void testTruncateIfSizeIsDifferentToTargetSize() throws IOException {
     @Test
     public void testPreallocateTrue() throws IOException {
         File temp = tempFile();
-        FileRecords fileRecords = FileRecords.open(temp, false, 512 * 1024 * 1024, true);
+        FileRecords fileRecords = FileRecords.open(temp, false, 1024 * 1024, true);
         long position = fileRecords.channel().position();
         int size = fileRecords.sizeInBytes();
         assertEquals(0, position);
         assertEquals(0, size);
-        assertEquals(512 * 1024 * 1024, temp.length());
+        assertEquals(1024 * 1024, temp.length());
     }
 
     /**
@@ -302,7 +369,7 @@ public void testPreallocateTrue() throws IOException {
     @Test
     public void testPreallocateFalse() throws IOException {
         File temp = tempFile();
-        FileRecords set = FileRecords.open(temp, false, 512 * 1024 * 1024, false);
+        FileRecords set = FileRecords.open(temp, false, 1024 * 1024, false);
         long position = set.channel().position();
         int size = set.sizeInBytes();
         assertEquals(0, position);
@@ -316,7 +383,7 @@ public void testPreallocateFalse() throws IOException {
     @Test
     public void testPreallocateClearShutdown() throws IOException {
         File temp = tempFile();
-        FileRecords fileRecords = FileRecords.open(temp, false, 512 * 1024 * 1024, true);
+        FileRecords fileRecords = FileRecords.open(temp, false, 1024 * 1024, true);
         append(fileRecords, values);
 
         int oldPosition = (int) fileRecords.channel().position();
@@ -326,7 +393,7 @@ public void testPreallocateClearShutdown() throws IOException {
         fileRecords.close();
 
         File tempReopen = new File(temp.getAbsolutePath());
-        FileRecords setReopen = FileRecords.open(tempReopen, true, 512 * 1024 * 1024, true);
+        FileRecords setReopen = FileRecords.open(tempReopen, true, 1024 * 1024, true);
         int position = (int) setReopen.channel().position();
         int size = setReopen.sizeInBytes();
 
@@ -340,10 +407,98 @@ public void testFormatConversionWithPartialMessage() throws IOException {
         RecordBatch batch = batches(fileRecords).get(1);
         int start = fileRecords.searchForOffsetWithSize(1, 0).position;
         int size = batch.sizeInBytes();
-        FileRecords slice = fileRecords.read(start, size - 1);
+        FileRecords slice = fileRecords.slice(start, size - 1);
         Records messageV0 = slice.downConvert(RecordBatch.MAGIC_VALUE_V0, 0, time).records();
         assertTrue("No message should be there", batches(messageV0).isEmpty());
         assertEquals("There should be " + (size - 1) + " bytes", size - 1, messageV0.sizeInBytes());
+
+        // Lazy down-conversion will not return any messages for a partial input batch
+        TopicPartition tp = new TopicPartition("topic-1", 0);
+        LazyDownConversionRecords lazyRecords = new LazyDownConversionRecords(tp, slice, RecordBatch.MAGIC_VALUE_V0, 0, Time.SYSTEM);
+        Iterator> it = lazyRecords.iterator(16 * 1024L);
+        assertTrue("No messages should be returned", !it.hasNext());
+    }
+
+    @Test
+    public void testSearchForTimestamp() throws IOException {
+        for (RecordVersion version : RecordVersion.values()) {
+            testSearchForTimestamp(version);
+        }
+    }
+
+    private void testSearchForTimestamp(RecordVersion version) throws IOException {
+        File temp = tempFile();
+        FileRecords fileRecords = FileRecords.open(temp, false, 1024 * 1024, true);
+        appendWithOffsetAndTimestamp(fileRecords, version, 10L, 5, 0);
+        appendWithOffsetAndTimestamp(fileRecords, version, 11L, 6, 1);
+
+        assertFoundTimestamp(new FileRecords.TimestampAndOffset(10L, 5, Optional.of(0)),
+                fileRecords.searchForTimestamp(9L, 0, 0L), version);
+        assertFoundTimestamp(new FileRecords.TimestampAndOffset(10L, 5, Optional.of(0)),
+                fileRecords.searchForTimestamp(10L, 0, 0L), version);
+        assertFoundTimestamp(new FileRecords.TimestampAndOffset(11L, 6, Optional.of(1)),
+                fileRecords.searchForTimestamp(11L, 0, 0L), version);
+        assertNull(fileRecords.searchForTimestamp(12L, 0, 0L));
+    }
+
+    private void assertFoundTimestamp(FileRecords.TimestampAndOffset expected,
+                                      FileRecords.TimestampAndOffset actual,
+                                      RecordVersion version) {
+        if (version == RecordVersion.V0) {
+            assertNull("Expected no match for message format v0", actual);
+        } else {
+            assertNotNull("Expected to find timestamp for message format " + version, actual);
+            assertEquals("Expected matching timestamps for message format" + version, expected.timestamp, actual.timestamp);
+            assertEquals("Expected matching offsets for message format " + version, expected.offset, actual.offset);
+            Optional expectedLeaderEpoch = version.value >= RecordVersion.V2.value ?
+                    expected.leaderEpoch : Optional.empty();
+            assertEquals("Non-matching leader epoch for version " + version, expectedLeaderEpoch, actual.leaderEpoch);
+        }
+    }
+
+    private void appendWithOffsetAndTimestamp(FileRecords fileRecords,
+                                              RecordVersion recordVersion,
+                                              long timestamp,
+                                              long offset,
+                                              int leaderEpoch) throws IOException {
+        ByteBuffer buffer = ByteBuffer.allocate(128);
+        MemoryRecordsBuilder builder = MemoryRecords.builder(buffer, recordVersion.value,
+                CompressionType.NONE, TimestampType.CREATE_TIME, offset, timestamp, leaderEpoch);
+        builder.append(new SimpleRecord(timestamp, new byte[0], new byte[0]));
+        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
@@ -358,6 +513,11 @@ public void testConversion() throws IOException {
 
     private void doTestConversion(CompressionType compressionType, byte toMagic) throws IOException {
         List offsets = asList(0L, 2L, 3L, 9L, 11L, 15L, 16L, 17L, 22L, 24L);
+
+        Header[] headers = {new RecordHeader("headerKey1", "headerValue1".getBytes()),
+                            new RecordHeader("headerKey2", "headerValue2".getBytes()),
+                            new RecordHeader("headerKey3", "headerValue3".getBytes())};
+
         List records = asList(
                 new SimpleRecord(1L, "k1".getBytes(), "hello".getBytes()),
                 new SimpleRecord(2L, "k2".getBytes(), "goodbye".getBytes()),
@@ -366,9 +526,10 @@ private void doTestConversion(CompressionType compressionType, byte toMagic) thr
                 new SimpleRecord(5L, "k5".getBytes(), "hello again".getBytes()),
                 new SimpleRecord(6L, "k6".getBytes(), "I sense indecision".getBytes()),
                 new SimpleRecord(7L, "k7".getBytes(), "what now".getBytes()),
-                new SimpleRecord(8L, "k8".getBytes(), "running out".getBytes()),
+                new SimpleRecord(8L, "k8".getBytes(), "running out".getBytes(), headers),
                 new SimpleRecord(9L, "k9".getBytes(), "ok, almost done".getBytes()),
-                new SimpleRecord(10L, "k10".getBytes(), "finally".getBytes()));
+                new SimpleRecord(10L, "k10".getBytes(), "finally".getBytes(), headers));
+        assertEquals("incorrect test setup", offsets.size(), records.size());
 
         ByteBuffer buffer = ByteBuffer.allocate(1024);
         MemoryRecordsBuilder builder = MemoryRecords.builder(buffer, RecordBatch.MAGIC_VALUE_V0, compressionType,
@@ -393,8 +554,7 @@ private void doTestConversion(CompressionType compressionType, byte toMagic) thr
         try (FileRecords fileRecords = FileRecords.open(tempFile())) {
             fileRecords.append(MemoryRecords.readableRecords(buffer));
             fileRecords.flush();
-            Records convertedRecords = fileRecords.downConvert(toMagic, 0L, time).records();
-            verifyConvertedRecords(records, offsets, convertedRecords, compressionType, toMagic);
+            downConvertAndVerifyRecords(records, offsets, fileRecords, compressionType, toMagic, 0L, time);
 
             if (toMagic <= RecordBatch.MAGIC_VALUE_V1 && compressionType == CompressionType.NONE) {
                 long firstOffset;
@@ -402,58 +562,94 @@ private void doTestConversion(CompressionType compressionType, byte toMagic) thr
                     firstOffset = 11L; // v1 record
                 else
                     firstOffset = 17; // v2 record
-                Records convertedRecords2 = fileRecords.downConvert(toMagic, firstOffset, time).records();
                 List filteredOffsets = new ArrayList<>(offsets);
                 List filteredRecords = new ArrayList<>(records);
                 int index = filteredOffsets.indexOf(firstOffset) - 1;
                 filteredRecords.remove(index);
                 filteredOffsets.remove(index);
-                verifyConvertedRecords(filteredRecords, filteredOffsets, convertedRecords2, compressionType, toMagic);
+                downConvertAndVerifyRecords(filteredRecords, filteredOffsets, fileRecords, compressionType, toMagic, firstOffset, time);
             } else {
                 // firstOffset doesn't have any effect in this case
-                Records convertedRecords2 = fileRecords.downConvert(toMagic, 10L, time).records();
-                verifyConvertedRecords(records, offsets, convertedRecords2, compressionType, toMagic);
+                downConvertAndVerifyRecords(records, offsets, fileRecords, compressionType, toMagic, 10L, time);
             }
         }
     }
 
-    private String utf8(ByteBuffer buffer) {
-        return Utils.utf8(buffer, buffer.remaining());
+    private void downConvertAndVerifyRecords(List initialRecords,
+                                             List initialOffsets,
+                                             FileRecords fileRecords,
+                                             CompressionType compressionType,
+                                             byte toMagic,
+                                             long firstOffset,
+                                             Time time) {
+        long minBatchSize = Long.MAX_VALUE;
+        long maxBatchSize = Long.MIN_VALUE;
+        for (RecordBatch batch : fileRecords.batches()) {
+            minBatchSize = Math.min(minBatchSize, batch.sizeInBytes());
+            maxBatchSize = Math.max(maxBatchSize, batch.sizeInBytes());
+        }
+
+        // Test the normal down-conversion path
+        List convertedRecords = new ArrayList<>();
+        convertedRecords.add(fileRecords.downConvert(toMagic, firstOffset, time).records());
+        verifyConvertedRecords(initialRecords, initialOffsets, convertedRecords, compressionType, toMagic);
+        convertedRecords.clear();
+
+        // Test the lazy down-conversion path
+        List maximumReadSize = asList(16L * 1024L,
+                (long) fileRecords.sizeInBytes(),
+                (long) fileRecords.sizeInBytes() - 1,
+                (long) fileRecords.sizeInBytes() / 4,
+                maxBatchSize + 1,
+                1L);
+        for (long readSize : maximumReadSize) {
+            TopicPartition tp = new TopicPartition("topic-1", 0);
+            LazyDownConversionRecords lazyRecords = new LazyDownConversionRecords(tp, fileRecords, toMagic, firstOffset, Time.SYSTEM);
+            Iterator> it = lazyRecords.iterator(readSize);
+            while (it.hasNext())
+                convertedRecords.add(it.next().records());
+            verifyConvertedRecords(initialRecords, initialOffsets, convertedRecords, compressionType, toMagic);
+            convertedRecords.clear();
+        }
     }
 
     private void verifyConvertedRecords(List initialRecords,
                                         List initialOffsets,
-                                        Records convertedRecords,
+                                        List convertedRecordsList,
                                         CompressionType compressionType,
                                         byte magicByte) {
         int i = 0;
-        for (RecordBatch batch : convertedRecords.batches()) {
-            assertTrue("Magic byte should be lower than or equal to " + magicByte, batch.magic() <= magicByte);
-            if (batch.magic() == RecordBatch.MAGIC_VALUE_V0)
-                assertEquals(TimestampType.NO_TIMESTAMP_TYPE, batch.timestampType());
-            else
-                assertEquals(TimestampType.CREATE_TIME, batch.timestampType());
-            assertEquals("Compression type should not be affected by conversion", compressionType, batch.compressionType());
-            for (Record record : batch) {
-                assertTrue("Inner record should have magic " + magicByte, record.hasMagic(batch.magic()));
-                assertEquals("Offset should not change", initialOffsets.get(i).longValue(), record.offset());
-                assertEquals("Key should not change", utf8(initialRecords.get(i).key()), utf8(record.key()));
-                assertEquals("Value should not change", utf8(initialRecords.get(i).value()), utf8(record.value()));
-                assertFalse(record.hasTimestampType(TimestampType.LOG_APPEND_TIME));
-                if (batch.magic() == RecordBatch.MAGIC_VALUE_V0) {
-                    assertEquals(RecordBatch.NO_TIMESTAMP, record.timestamp());
-                    assertFalse(record.hasTimestampType(TimestampType.CREATE_TIME));
-                    assertTrue(record.hasTimestampType(TimestampType.NO_TIMESTAMP_TYPE));
-                } else if (batch.magic() == RecordBatch.MAGIC_VALUE_V1) {
-                    assertEquals("Timestamp should not change", initialRecords.get(i).timestamp(), record.timestamp());
-                    assertTrue(record.hasTimestampType(TimestampType.CREATE_TIME));
-                    assertFalse(record.hasTimestampType(TimestampType.NO_TIMESTAMP_TYPE));
-                } else {
-                    assertEquals("Timestamp should not change", initialRecords.get(i).timestamp(), record.timestamp());
-                    assertFalse(record.hasTimestampType(TimestampType.CREATE_TIME));
-                    assertFalse(record.hasTimestampType(TimestampType.NO_TIMESTAMP_TYPE));
+
+        for (Records convertedRecords : convertedRecordsList) {
+            for (RecordBatch batch : convertedRecords.batches()) {
+                assertTrue("Magic byte should be lower than or equal to " + magicByte, batch.magic() <= magicByte);
+                if (batch.magic() == RecordBatch.MAGIC_VALUE_V0)
+                    assertEquals(TimestampType.NO_TIMESTAMP_TYPE, batch.timestampType());
+                else
+                    assertEquals(TimestampType.CREATE_TIME, batch.timestampType());
+                assertEquals("Compression type should not be affected by conversion", compressionType, batch.compressionType());
+                for (Record record : batch) {
+                    assertTrue("Inner record should have magic " + magicByte, record.hasMagic(batch.magic()));
+                    assertEquals("Offset should not change", initialOffsets.get(i).longValue(), record.offset());
+                    assertEquals("Key should not change", utf8(initialRecords.get(i).key()), utf8(record.key()));
+                    assertEquals("Value should not change", utf8(initialRecords.get(i).value()), utf8(record.value()));
+                    assertFalse(record.hasTimestampType(TimestampType.LOG_APPEND_TIME));
+                    if (batch.magic() == RecordBatch.MAGIC_VALUE_V0) {
+                        assertEquals(RecordBatch.NO_TIMESTAMP, record.timestamp());
+                        assertFalse(record.hasTimestampType(TimestampType.CREATE_TIME));
+                        assertTrue(record.hasTimestampType(TimestampType.NO_TIMESTAMP_TYPE));
+                    } else if (batch.magic() == RecordBatch.MAGIC_VALUE_V1) {
+                        assertEquals("Timestamp should not change", initialRecords.get(i).timestamp(), record.timestamp());
+                        assertTrue(record.hasTimestampType(TimestampType.CREATE_TIME));
+                        assertFalse(record.hasTimestampType(TimestampType.NO_TIMESTAMP_TYPE));
+                    } else {
+                        assertEquals("Timestamp should not change", initialRecords.get(i).timestamp(), record.timestamp());
+                        assertFalse(record.hasTimestampType(TimestampType.CREATE_TIME));
+                        assertFalse(record.hasTimestampType(TimestampType.NO_TIMESTAMP_TYPE));
+                        assertArrayEquals("Headers should not change", initialRecords.get(i).headers(), record.headers());
+                    }
+                    i += 1;
                 }
-                i += 1;
             }
         }
         assertEquals(initialOffsets.size(), i);
@@ -480,5 +676,4 @@ private void append(FileRecords fileRecords, byte[][] values) throws IOException
         }
         fileRecords.flush();
     }
-
 }
diff --git a/clients/src/test/java/org/apache/kafka/common/record/KafkaLZ4Test.java b/clients/src/test/java/org/apache/kafka/common/record/KafkaLZ4Test.java
index 222599b8b7317..5e44c49825416 100644
--- a/clients/src/test/java/org/apache/kafka/common/record/KafkaLZ4Test.java
+++ b/clients/src/test/java/org/apache/kafka/common/record/KafkaLZ4Test.java
@@ -19,9 +19,7 @@
 import net.jpountz.xxhash.XXHashFactory;
 
 import org.hamcrest.CoreMatchers;
-import org.junit.Rule;
 import org.junit.Test;
-import org.junit.rules.ExpectedException;
 import org.junit.runner.RunWith;
 import org.junit.runners.Parameterized;
 import org.junit.runners.Parameterized.Parameters;
@@ -37,9 +35,11 @@
 import java.util.Random;
 
 import static org.apache.kafka.common.record.KafkaLZ4BlockOutputStream.LZ4_FRAME_INCOMPRESSIBLE_MASK;
+import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 
 @RunWith(value = Parameterized.class)
@@ -71,9 +71,6 @@ public String toString() {
         }
     }
 
-    @Rule
-    public ExpectedException thrown = ExpectedException.none();
-
     @Parameters(name = "{index} useBrokenFlagDescriptorChecksum={0}, ignoreFlagDescriptorChecksum={1}, blockChecksum={2}, close={3}, payload={4}")
     public static Collection data() {
         List payloads = new ArrayList<>();
@@ -111,12 +108,10 @@ public KafkaLZ4Test(boolean useBrokenFlagDescriptorChecksum, boolean ignoreFlagD
     }
 
     @Test
-    public void testHeaderPrematureEnd() throws Exception {
-        thrown.expect(IOException.class);
-        thrown.expectMessage(KafkaLZ4BlockInputStream.PREMATURE_EOS);
-
-        final ByteBuffer buffer = ByteBuffer.allocate(2);
-        makeInputStream(buffer);
+    public void testHeaderPrematureEnd() {
+        ByteBuffer buffer = ByteBuffer.allocate(2);
+        IOException e = assertThrows(IOException.class, () -> makeInputStream(buffer));
+        assertEquals(KafkaLZ4BlockInputStream.PREMATURE_EOS, e.getMessage());
     }
 
     private KafkaLZ4BlockInputStream makeInputStream(ByteBuffer buffer) throws IOException {
@@ -125,43 +120,41 @@ private KafkaLZ4BlockInputStream makeInputStream(ByteBuffer buffer) throws IOExc
 
     @Test
     public void testNotSupported() throws Exception {
-        thrown.expect(IOException.class);
-        thrown.expectMessage(KafkaLZ4BlockInputStream.NOT_SUPPORTED);
-
         byte[] compressed = compressedBytes();
         compressed[0] = 0x00;
-
-        makeInputStream(ByteBuffer.wrap(compressed));
+        ByteBuffer buffer = ByteBuffer.wrap(compressed);
+        IOException e = assertThrows(IOException.class, () -> makeInputStream(buffer));
+        assertEquals(KafkaLZ4BlockInputStream.NOT_SUPPORTED, e.getMessage());
     }
 
     @Test
     public void testBadFrameChecksum() throws Exception {
-        if (!ignoreFlagDescriptorChecksum) {
-            thrown.expect(IOException.class);
-            thrown.expectMessage(KafkaLZ4BlockInputStream.DESCRIPTOR_HASH_MISMATCH);
-        }
-
         byte[] compressed = compressedBytes();
         compressed[6] = (byte) 0xFF;
+        ByteBuffer buffer = ByteBuffer.wrap(compressed);
 
-        makeInputStream(ByteBuffer.wrap(compressed));
+        if (ignoreFlagDescriptorChecksum) {
+            makeInputStream(buffer);
+        } else {
+            IOException e = assertThrows(IOException.class, () -> makeInputStream(buffer));
+            assertEquals(KafkaLZ4BlockInputStream.DESCRIPTOR_HASH_MISMATCH, e.getMessage());
+        }
     }
 
     @Test
     public void testBadBlockSize() throws Exception {
-        if (!close || (useBrokenFlagDescriptorChecksum && !ignoreFlagDescriptorChecksum)) return;
-
-        thrown.expect(IOException.class);
-        thrown.expectMessage(CoreMatchers.containsString("exceeded max"));
+        if (!close || (useBrokenFlagDescriptorChecksum && !ignoreFlagDescriptorChecksum))
+            return;
 
         byte[] compressed = compressedBytes();
-        final ByteBuffer buffer = ByteBuffer.wrap(compressed).order(ByteOrder.LITTLE_ENDIAN);
+        ByteBuffer buffer = ByteBuffer.wrap(compressed).order(ByteOrder.LITTLE_ENDIAN);
 
         int blockSize = buffer.getInt(7);
         blockSize = (blockSize & LZ4_FRAME_INCOMPRESSIBLE_MASK) | (1 << 24 & ~LZ4_FRAME_INCOMPRESSIBLE_MASK);
         buffer.putInt(7, blockSize);
 
-        testDecompression(buffer);
+        IOException e = assertThrows(IOException.class, () -> testDecompression(buffer));
+        assertThat(e.getMessage(), CoreMatchers.containsString("exceeded max"));
     }
 
 
diff --git a/clients/src/test/java/org/apache/kafka/common/record/LazyDownConversionRecordsTest.java b/clients/src/test/java/org/apache/kafka/common/record/LazyDownConversionRecordsTest.java
new file mode 100644
index 0000000000000..a25d279ba4816
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/record/LazyDownConversionRecordsTest.java
@@ -0,0 +1,219 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.TopicPartition;
+import org.apache.kafka.common.header.Header;
+import org.apache.kafka.common.header.internals.RecordHeader;
+import org.apache.kafka.common.utils.Time;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+
+import java.io.File;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.channels.FileChannel;
+import java.nio.file.StandardOpenOption;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+
+import static java.util.Arrays.asList;
+import static org.apache.kafka.common.utils.Utils.utf8;
+import static org.apache.kafka.test.TestUtils.tempFile;
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+public class LazyDownConversionRecordsTest {
+    /**
+     * Test the lazy down-conversion path in the presence of commit markers. When converting to V0 or V1, these batches
+     * are dropped. If there happen to be no more batches left to convert, we must get an overflow message batch after
+     * conversion.
+     */
+    @Test
+    public void testConversionOfCommitMarker() throws IOException {
+        MemoryRecords recordsToConvert = MemoryRecords.withEndTransactionMarker(0, Time.SYSTEM.milliseconds(), RecordBatch.NO_PARTITION_LEADER_EPOCH,
+                1, (short) 1, new EndTransactionMarker(ControlRecordType.COMMIT, 0));
+        MemoryRecords convertedRecords = convertRecords(recordsToConvert, (byte) 1, recordsToConvert.sizeInBytes());
+        ByteBuffer buffer = convertedRecords.buffer();
+
+        // read the offset and the batch length
+        buffer.getLong();
+        int sizeOfConvertedRecords = buffer.getInt();
+
+        // assert we got an overflow message batch
+        assertTrue(sizeOfConvertedRecords > buffer.limit());
+        assertFalse(convertedRecords.batchIterator().hasNext());
+    }
+
+    @RunWith(value = Parameterized.class)
+    public static class ParameterizedConversionTest {
+        private final CompressionType compressionType;
+        private final byte toMagic;
+
+        public ParameterizedConversionTest(CompressionType compressionType, byte toMagic) {
+            this.compressionType = compressionType;
+            this.toMagic = toMagic;
+        }
+
+        @Parameterized.Parameters(name = "compressionType={0}, toMagic={1}")
+        public static Collection data() {
+            List values = new ArrayList<>();
+            for (byte toMagic = RecordBatch.MAGIC_VALUE_V0; toMagic <= RecordBatch.CURRENT_MAGIC_VALUE; toMagic++) {
+                values.add(new Object[]{CompressionType.NONE, toMagic});
+                values.add(new Object[]{CompressionType.GZIP, toMagic});
+            }
+            return values;
+        }
+
+        /**
+         * Test the lazy down-conversion path.
+         */
+        @Test
+        public void testConversion() throws IOException {
+            doTestConversion(false);
+        }
+
+        /**
+         * Test the lazy down-conversion path where the number of bytes we want to convert is much larger than the
+         * number of bytes we get after conversion. This causes overflow message batch(es) to be appended towards the
+         * end of the converted output.
+         */
+        @Test
+        public void testConversionWithOverflow() throws IOException {
+            doTestConversion(true);
+        }
+
+        private void doTestConversion(boolean testConversionOverflow) throws IOException {
+            List offsets = asList(0L, 2L, 3L, 9L, 11L, 15L, 16L, 17L, 22L, 24L);
+
+            Header[] headers = {new RecordHeader("headerKey1", "headerValue1".getBytes()),
+                                new RecordHeader("headerKey2", "headerValue2".getBytes()),
+                                new RecordHeader("headerKey3", "headerValue3".getBytes())};
+
+            List records = asList(
+                    new SimpleRecord(1L, "k1".getBytes(), "hello".getBytes()),
+                    new SimpleRecord(2L, "k2".getBytes(), "goodbye".getBytes()),
+                    new SimpleRecord(3L, "k3".getBytes(), "hello again".getBytes()),
+                    new SimpleRecord(4L, "k4".getBytes(), "goodbye for now".getBytes()),
+                    new SimpleRecord(5L, "k5".getBytes(), "hello again".getBytes()),
+                    new SimpleRecord(6L, "k6".getBytes(), "I sense indecision".getBytes()),
+                    new SimpleRecord(7L, "k7".getBytes(), "what now".getBytes()),
+                    new SimpleRecord(8L, "k8".getBytes(), "running out".getBytes(), headers),
+                    new SimpleRecord(9L, "k9".getBytes(), "ok, almost done".getBytes()),
+                    new SimpleRecord(10L, "k10".getBytes(), "finally".getBytes(), headers));
+            assertEquals("incorrect test setup", offsets.size(), records.size());
+
+            ByteBuffer buffer = ByteBuffer.allocate(1024);
+            MemoryRecordsBuilder builder = MemoryRecords.builder(buffer, RecordBatch.CURRENT_MAGIC_VALUE, compressionType,
+                    TimestampType.CREATE_TIME, 0L);
+            for (int i = 0; i < 3; i++)
+                builder.appendWithOffset(offsets.get(i), records.get(i));
+            builder.close();
+
+            builder = MemoryRecords.builder(buffer, RecordBatch.CURRENT_MAGIC_VALUE, compressionType, TimestampType.CREATE_TIME,
+                    0L);
+            for (int i = 3; i < 6; i++)
+                builder.appendWithOffset(offsets.get(i), records.get(i));
+            builder.close();
+
+            builder = MemoryRecords.builder(buffer, RecordBatch.CURRENT_MAGIC_VALUE, compressionType, TimestampType.CREATE_TIME,
+                    0L);
+            for (int i = 6; i < 10; i++)
+                builder.appendWithOffset(offsets.get(i), records.get(i));
+            builder.close();
+            buffer.flip();
+
+            MemoryRecords recordsToConvert = MemoryRecords.readableRecords(buffer);
+            int numBytesToConvert = recordsToConvert.sizeInBytes();
+            if (testConversionOverflow)
+                numBytesToConvert *= 2;
+
+            MemoryRecords convertedRecords = convertRecords(recordsToConvert, toMagic, numBytesToConvert);
+            verifyDownConvertedRecords(records, offsets, convertedRecords, compressionType, toMagic);
+        }
+    }
+
+    private static MemoryRecords convertRecords(MemoryRecords recordsToConvert, byte toMagic, int bytesToConvert) throws IOException {
+        try (FileRecords inputRecords = FileRecords.open(tempFile())) {
+            inputRecords.append(recordsToConvert);
+            inputRecords.flush();
+
+            LazyDownConversionRecords lazyRecords = new LazyDownConversionRecords(new TopicPartition("test", 1),
+                    inputRecords, toMagic, 0L, Time.SYSTEM);
+            LazyDownConversionRecordsSend lazySend = lazyRecords.toSend();
+            File outputFile = tempFile();
+            FileChannel channel = FileChannel.open(outputFile.toPath(), StandardOpenOption.READ, StandardOpenOption.WRITE);
+
+            int written = 0;
+            while (written < bytesToConvert)
+                written += lazySend.writeTo(channel, written, bytesToConvert - written);
+
+            FileRecords convertedRecords = FileRecords.open(outputFile, true, (int) channel.size(), false);
+            ByteBuffer convertedRecordsBuffer = ByteBuffer.allocate(convertedRecords.sizeInBytes());
+            convertedRecords.readInto(convertedRecordsBuffer, 0);
+
+            // cleanup
+            convertedRecords.close();
+            channel.close();
+
+            return MemoryRecords.readableRecords(convertedRecordsBuffer);
+        }
+    }
+
+    private static void verifyDownConvertedRecords(List initialRecords,
+                                                   List initialOffsets,
+                                                   MemoryRecords downConvertedRecords,
+                                                   CompressionType compressionType,
+                                                   byte toMagic) {
+        int i = 0;
+        for (RecordBatch batch : downConvertedRecords.batches()) {
+            assertTrue("Magic byte should be lower than or equal to " + toMagic, batch.magic() <= toMagic);
+            if (batch.magic() == RecordBatch.MAGIC_VALUE_V0)
+                assertEquals(TimestampType.NO_TIMESTAMP_TYPE, batch.timestampType());
+            else
+                assertEquals(TimestampType.CREATE_TIME, batch.timestampType());
+            assertEquals("Compression type should not be affected by conversion", compressionType, batch.compressionType());
+            for (Record record : batch) {
+                assertTrue("Inner record should have magic " + toMagic, record.hasMagic(batch.magic()));
+                assertEquals("Offset should not change", initialOffsets.get(i).longValue(), record.offset());
+                assertEquals("Key should not change", utf8(initialRecords.get(i).key()), utf8(record.key()));
+                assertEquals("Value should not change", utf8(initialRecords.get(i).value()), utf8(record.value()));
+                assertFalse(record.hasTimestampType(TimestampType.LOG_APPEND_TIME));
+                if (batch.magic() == RecordBatch.MAGIC_VALUE_V0) {
+                    assertEquals(RecordBatch.NO_TIMESTAMP, record.timestamp());
+                    assertFalse(record.hasTimestampType(TimestampType.CREATE_TIME));
+                    assertTrue(record.hasTimestampType(TimestampType.NO_TIMESTAMP_TYPE));
+                } else if (batch.magic() == RecordBatch.MAGIC_VALUE_V1) {
+                    assertEquals("Timestamp should not change", initialRecords.get(i).timestamp(), record.timestamp());
+                    assertTrue(record.hasTimestampType(TimestampType.CREATE_TIME));
+                    assertFalse(record.hasTimestampType(TimestampType.NO_TIMESTAMP_TYPE));
+                } else {
+                    assertEquals("Timestamp should not change", initialRecords.get(i).timestamp(), record.timestamp());
+                    assertFalse(record.hasTimestampType(TimestampType.CREATE_TIME));
+                    assertFalse(record.hasTimestampType(TimestampType.NO_TIMESTAMP_TYPE));
+                    assertArrayEquals("Headers should not change", initialRecords.get(i).headers(), record.headers());
+                }
+                i += 1;
+            }
+        }
+        assertEquals(initialOffsets.size(), i);
+    }
+}
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 c713d17f61545..885b20affe78d 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
@@ -16,6 +16,9 @@
  */
 package org.apache.kafka.common.record;
 
+import org.apache.kafka.common.errors.UnsupportedCompressionTypeException;
+import org.apache.kafka.common.message.LeaderChangeMessage;
+import org.apache.kafka.common.message.LeaderChangeMessage.Voter;
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.common.utils.Utils;
 import org.apache.kafka.test.TestUtils;
@@ -27,17 +30,24 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.List;
+import java.util.Random;
+import java.util.function.Supplier;
+import java.util.stream.Collectors;
 
+import static org.apache.kafka.common.record.RecordBatch.MAGIC_VALUE_V2;
+import static org.apache.kafka.common.utils.Utils.utf8;
 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;
+import static org.junit.Assume.assumeTrue;
 
 @RunWith(value = Parameterized.class)
 public class MemoryRecordsBuilderTest {
-
     private final CompressionType compressionType;
     private final int bufferOffset;
     private final Time time;
@@ -50,15 +60,25 @@ public MemoryRecordsBuilderTest(int bufferOffset, CompressionType compressionTyp
 
     @Test
     public void testWriteEmptyRecordSet() {
+        byte magic = RecordBatch.MAGIC_VALUE_V0;
+        assumeAtLeastV2OrNotZstd(magic);
+
         ByteBuffer buffer = ByteBuffer.allocate(128);
         buffer.position(bufferOffset);
 
-        MemoryRecordsBuilder builder = new MemoryRecordsBuilder(buffer, RecordBatch.MAGIC_VALUE_V0, compressionType,
-                TimestampType.CREATE_TIME, 0L, 0L, RecordBatch.NO_PRODUCER_ID, RecordBatch.NO_PRODUCER_EPOCH, RecordBatch.NO_SEQUENCE,
-                false, false, RecordBatch.NO_PARTITION_LEADER_EPOCH, buffer.capacity());
-        MemoryRecords records = builder.build();
-        assertEquals(0, records.sizeInBytes());
-        assertEquals(bufferOffset, buffer.position());
+        Supplier builderSupplier = () -> new MemoryRecordsBuilder(buffer, magic,
+            compressionType, TimestampType.CREATE_TIME, 0L, 0L,
+            RecordBatch.NO_PRODUCER_ID, RecordBatch.NO_PRODUCER_EPOCH, RecordBatch.NO_SEQUENCE,
+            false, false, RecordBatch.NO_PARTITION_LEADER_EPOCH, buffer.capacity());
+
+        if (compressionType != CompressionType.ZSTD) {
+            MemoryRecords records = builderSupplier.get().build();
+            assertEquals(0, records.sizeInBytes());
+            assertEquals(bufferOffset, buffer.position());
+        } else {
+            Exception e = assertThrows(IllegalArgumentException.class, () -> builderSupplier.get().build());
+            assertEquals(e.getMessage(), "ZStandard compression is not supported for magic " + magic);
+        }
     }
 
     @Test
@@ -203,18 +223,65 @@ public void testWriteEndTxnMarkerNonControlBatch() {
         builder.appendEndTxnMarker(RecordBatch.NO_TIMESTAMP, new EndTransactionMarker(ControlRecordType.ABORT, 0));
     }
 
+    @Test(expected = IllegalArgumentException.class)
+    public void testWriteLeaderChangeControlBatchWithoutLeaderEpoch() {
+        ByteBuffer buffer = ByteBuffer.allocate(128);
+        buffer.position(bufferOffset);
+
+        final int leaderId = 1;
+        MemoryRecordsBuilder builder = new MemoryRecordsBuilder(buffer, RecordBatch.CURRENT_MAGIC_VALUE, compressionType, TimestampType.CREATE_TIME,
+            0L, 0L,
+            RecordBatch.NO_PRODUCER_ID, RecordBatch.NO_PRODUCER_EPOCH, RecordBatch.NO_SEQUENCE,
+            false, true, RecordBatch.NO_PARTITION_LEADER_EPOCH, buffer.capacity());
+        builder.appendLeaderChangeMessage(RecordBatch.NO_TIMESTAMP,
+            new LeaderChangeMessage()
+                .setLeaderId(leaderId)
+                .setVoters(Collections.emptyList()));
+    }
+
+    @Test
+    public void testWriteLeaderChangeControlBatch() {
+        ByteBuffer buffer = ByteBuffer.allocate(128);
+        buffer.position(bufferOffset);
+
+        final int leaderId = 1;
+        final int leaderEpoch = 5;
+        final List voters = Arrays.asList(2, 3);
+
+        MemoryRecordsBuilder builder = new MemoryRecordsBuilder(buffer, RecordBatch.CURRENT_MAGIC_VALUE, compressionType, TimestampType.CREATE_TIME,
+            0L, 0L,
+            RecordBatch.NO_PRODUCER_ID, RecordBatch.NO_PRODUCER_EPOCH, RecordBatch.NO_SEQUENCE,
+            false, true, leaderEpoch, buffer.capacity());
+        builder.appendLeaderChangeMessage(RecordBatch.NO_TIMESTAMP,
+            new LeaderChangeMessage()
+                .setLeaderId(leaderId)
+                .setVoters(voters.stream().map(
+                    voterId -> new Voter().setVoterId(voterId)).collect(Collectors.toList())));
+
+        MemoryRecords built = builder.build();
+        List records = TestUtils.toList(built.records());
+        assertEquals(1, records.size());
+        LeaderChangeMessage leaderChangeMessage = ControlRecordUtils.deserializeLeaderChangeMessage(records.get(0));
+
+        assertEquals(leaderId, leaderChangeMessage.leaderId());
+        assertEquals(voters, leaderChangeMessage.voters().stream().map(Voter::voterId).collect(Collectors.toList()));
+    }
+
     @Test
     public void testCompressionRateV0() {
+        byte magic = RecordBatch.MAGIC_VALUE_V0;
+        assumeAtLeastV2OrNotZstd(magic);
+
         ByteBuffer buffer = ByteBuffer.allocate(1024);
         buffer.position(bufferOffset);
 
         LegacyRecord[] records = new LegacyRecord[] {
-                LegacyRecord.create(RecordBatch.MAGIC_VALUE_V0, 0L, "a".getBytes(), "1".getBytes()),
-                LegacyRecord.create(RecordBatch.MAGIC_VALUE_V0, 1L, "b".getBytes(), "2".getBytes()),
-                LegacyRecord.create(RecordBatch.MAGIC_VALUE_V0, 2L, "c".getBytes(), "3".getBytes()),
+                LegacyRecord.create(magic, 0L, "a".getBytes(), "1".getBytes()),
+                LegacyRecord.create(magic, 1L, "b".getBytes(), "2".getBytes()),
+                LegacyRecord.create(magic, 2L, "c".getBytes(), "3".getBytes()),
         };
 
-        MemoryRecordsBuilder builder = new MemoryRecordsBuilder(buffer, RecordBatch.MAGIC_VALUE_V0, compressionType,
+        MemoryRecordsBuilder builder = new MemoryRecordsBuilder(buffer, magic, compressionType,
                 TimestampType.CREATE_TIME, 0L, 0L, RecordBatch.NO_PRODUCER_ID, RecordBatch.NO_PRODUCER_EPOCH, RecordBatch.NO_SEQUENCE,
                 false, false, RecordBatch.NO_PARTITION_LEADER_EPOCH, buffer.capacity());
 
@@ -260,16 +327,19 @@ public void testEstimatedSizeInBytes() {
 
     @Test
     public void testCompressionRateV1() {
+        byte magic = RecordBatch.MAGIC_VALUE_V1;
+        assumeAtLeastV2OrNotZstd(magic);
+
         ByteBuffer buffer = ByteBuffer.allocate(1024);
         buffer.position(bufferOffset);
 
         LegacyRecord[] records = new LegacyRecord[] {
-                LegacyRecord.create(RecordBatch.MAGIC_VALUE_V1, 0L, "a".getBytes(), "1".getBytes()),
-                LegacyRecord.create(RecordBatch.MAGIC_VALUE_V1, 1L, "b".getBytes(), "2".getBytes()),
-                LegacyRecord.create(RecordBatch.MAGIC_VALUE_V1, 2L, "c".getBytes(), "3".getBytes()),
+                LegacyRecord.create(magic, 0L, "a".getBytes(), "1".getBytes()),
+                LegacyRecord.create(magic, 1L, "b".getBytes(), "2".getBytes()),
+                LegacyRecord.create(magic, 2L, "c".getBytes(), "3".getBytes()),
         };
 
-        MemoryRecordsBuilder builder = new MemoryRecordsBuilder(buffer, RecordBatch.MAGIC_VALUE_V1, compressionType,
+        MemoryRecordsBuilder builder = new MemoryRecordsBuilder(buffer, magic, compressionType,
                 TimestampType.CREATE_TIME, 0L, 0L, RecordBatch.NO_PRODUCER_ID, RecordBatch.NO_PRODUCER_EPOCH, RecordBatch.NO_SEQUENCE,
                 false, false, RecordBatch.NO_PARTITION_LEADER_EPOCH, buffer.capacity());
 
@@ -291,11 +361,14 @@ public void testCompressionRateV1() {
 
     @Test
     public void buildUsingLogAppendTime() {
+        byte magic = RecordBatch.MAGIC_VALUE_V1;
+        assumeAtLeastV2OrNotZstd(magic);
+
         ByteBuffer buffer = ByteBuffer.allocate(1024);
         buffer.position(bufferOffset);
 
         long logAppendTime = System.currentTimeMillis();
-        MemoryRecordsBuilder builder = new MemoryRecordsBuilder(buffer, RecordBatch.MAGIC_VALUE_V1, compressionType,
+        MemoryRecordsBuilder builder = new MemoryRecordsBuilder(buffer, magic, compressionType,
                 TimestampType.LOG_APPEND_TIME, 0L, logAppendTime, RecordBatch.NO_PRODUCER_ID, RecordBatch.NO_PRODUCER_EPOCH,
                 RecordBatch.NO_SEQUENCE, false, false, RecordBatch.NO_PARTITION_LEADER_EPOCH, buffer.capacity());
         builder.append(0L, "a".getBytes(), "1".getBytes());
@@ -320,11 +393,14 @@ public void buildUsingLogAppendTime() {
 
     @Test
     public void buildUsingCreateTime() {
+        byte magic = RecordBatch.MAGIC_VALUE_V1;
+        assumeAtLeastV2OrNotZstd(magic);
+
         ByteBuffer buffer = ByteBuffer.allocate(1024);
         buffer.position(bufferOffset);
 
         long logAppendTime = System.currentTimeMillis();
-        MemoryRecordsBuilder builder = new MemoryRecordsBuilder(buffer, RecordBatch.MAGIC_VALUE_V1, compressionType,
+        MemoryRecordsBuilder builder = new MemoryRecordsBuilder(buffer, magic, compressionType,
                 TimestampType.CREATE_TIME, 0L, logAppendTime, RecordBatch.NO_PRODUCER_ID, RecordBatch.NO_PRODUCER_EPOCH, RecordBatch.NO_SEQUENCE,
                 false, false, RecordBatch.NO_PARTITION_LEADER_EPOCH, buffer.capacity());
         builder.append(0L, "a".getBytes(), "1".getBytes());
@@ -351,6 +427,9 @@ public void buildUsingCreateTime() {
 
     @Test
     public void testAppendedChecksumConsistency() {
+        assumeAtLeastV2OrNotZstd(RecordBatch.MAGIC_VALUE_V0);
+        assumeAtLeastV2OrNotZstd(RecordBatch.MAGIC_VALUE_V1);
+
         ByteBuffer buffer = ByteBuffer.allocate(512);
         for (byte magic : Arrays.asList(RecordBatch.MAGIC_VALUE_V0, RecordBatch.MAGIC_VALUE_V1, RecordBatch.MAGIC_VALUE_V2)) {
             MemoryRecordsBuilder builder = new MemoryRecordsBuilder(buffer, magic, compressionType,
@@ -395,11 +474,14 @@ public void testSmallWriteLimit() {
 
     @Test
     public void writePastLimit() {
+        byte magic = RecordBatch.MAGIC_VALUE_V1;
+        assumeAtLeastV2OrNotZstd(magic);
+
         ByteBuffer buffer = ByteBuffer.allocate(64);
         buffer.position(bufferOffset);
 
         long logAppendTime = System.currentTimeMillis();
-        MemoryRecordsBuilder builder = new MemoryRecordsBuilder(buffer, RecordBatch.MAGIC_VALUE_V1, compressionType,
+        MemoryRecordsBuilder builder = new MemoryRecordsBuilder(buffer, magic, compressionType,
                 TimestampType.CREATE_TIME, 0L, logAppendTime, RecordBatch.NO_PRODUCER_ID, RecordBatch.NO_PRODUCER_EPOCH, RecordBatch.NO_SEQUENCE,
                 false, false, RecordBatch.NO_PARTITION_LEADER_EPOCH, buffer.capacity());
         builder.setEstimatedCompressionRatio(0.5f);
@@ -466,35 +548,45 @@ public void convertV2ToV1UsingMixedCreateAndLogAppendTime() {
 
         buffer.flip();
 
-        ConvertedRecords convertedRecords = MemoryRecords.readableRecords(buffer)
-                .downConvert(RecordBatch.MAGIC_VALUE_V1, 0, time);
-        MemoryRecords records = convertedRecords.records();
+        Supplier> convertedRecordsSupplier = () ->
+            MemoryRecords.readableRecords(buffer).downConvert(RecordBatch.MAGIC_VALUE_V1, 0, time);
+
+        if (compressionType != CompressionType.ZSTD) {
+            ConvertedRecords convertedRecords = convertedRecordsSupplier.get();
+            MemoryRecords records = convertedRecords.records();
 
-        // Transactional markers are skipped when down converting to V1, so exclude them from size
-        verifyRecordsProcessingStats(convertedRecords.recordsProcessingStats(),
+            // Transactional markers are skipped when down converting to V1, so exclude them from size
+            verifyRecordsProcessingStats(convertedRecords.recordConversionStats(),
                 3, 3, records.sizeInBytes(), sizeExcludingTxnMarkers);
 
-        List batches = Utils.toList(records.batches().iterator());
-        if (compressionType != CompressionType.NONE) {
-            assertEquals(2, batches.size());
-            assertEquals(TimestampType.LOG_APPEND_TIME, batches.get(0).timestampType());
-            assertEquals(TimestampType.CREATE_TIME, batches.get(1).timestampType());
+            List batches = Utils.toList(records.batches().iterator());
+            if (compressionType != CompressionType.NONE) {
+                assertEquals(2, batches.size());
+                assertEquals(TimestampType.LOG_APPEND_TIME, batches.get(0).timestampType());
+                assertEquals(TimestampType.CREATE_TIME, batches.get(1).timestampType());
+            } else {
+                assertEquals(3, batches.size());
+                assertEquals(TimestampType.LOG_APPEND_TIME, batches.get(0).timestampType());
+                assertEquals(TimestampType.CREATE_TIME, batches.get(1).timestampType());
+                assertEquals(TimestampType.CREATE_TIME, batches.get(2).timestampType());
+            }
+
+            List logRecords = Utils.toList(records.records().iterator());
+            assertEquals(3, logRecords.size());
+            assertEquals(ByteBuffer.wrap("1".getBytes()), logRecords.get(0).key());
+            assertEquals(ByteBuffer.wrap("2".getBytes()), logRecords.get(1).key());
+            assertEquals(ByteBuffer.wrap("3".getBytes()), logRecords.get(2).key());
         } else {
-            assertEquals(3, batches.size());
-            assertEquals(TimestampType.LOG_APPEND_TIME, batches.get(0).timestampType());
-            assertEquals(TimestampType.CREATE_TIME, batches.get(1).timestampType());
-            assertEquals(TimestampType.CREATE_TIME, batches.get(2).timestampType());
+            Exception e = assertThrows(UnsupportedCompressionTypeException.class, convertedRecordsSupplier::get);
+            assertEquals("Down-conversion of zstandard-compressed batches is not supported", e.getMessage());
         }
-
-        List logRecords = Utils.toList(records.records().iterator());
-        assertEquals(3, logRecords.size());
-        assertEquals(ByteBuffer.wrap("1".getBytes()), logRecords.get(0).key());
-        assertEquals(ByteBuffer.wrap("2".getBytes()), logRecords.get(1).key());
-        assertEquals(ByteBuffer.wrap("3".getBytes()), logRecords.get(2).key());
     }
 
     @Test
     public void convertToV1WithMixedV0AndV2Data() {
+        assumeAtLeastV2OrNotZstd(RecordBatch.MAGIC_VALUE_V0);
+        assumeAtLeastV2OrNotZstd(RecordBatch.MAGIC_VALUE_V1);
+
         ByteBuffer buffer = ByteBuffer.allocate(512);
         MemoryRecordsBuilder builder = MemoryRecords.builder(buffer, RecordBatch.MAGIC_VALUE_V0,
                 compressionType, TimestampType.NO_TIMESTAMP_TYPE, 0L);
@@ -512,7 +604,7 @@ public void convertToV1WithMixedV0AndV2Data() {
         ConvertedRecords convertedRecords = MemoryRecords.readableRecords(buffer)
                 .downConvert(RecordBatch.MAGIC_VALUE_V1, 0, time);
         MemoryRecords records = convertedRecords.records();
-        verifyRecordsProcessingStats(convertedRecords.recordsProcessingStats(), 3, 2,
+        verifyRecordsProcessingStats(convertedRecords.recordConversionStats(), 3, 2,
                 records.sizeInBytes(), buffer.limit());
 
         List batches = Utils.toList(records.batches().iterator());
@@ -552,7 +644,7 @@ public void convertToV1WithMixedV0AndV2Data() {
             assertEquals("1", utf8(logRecords.get(0).key()));
             assertEquals("2", utf8(logRecords.get(1).key()));
             assertEquals("3", utf8(logRecords.get(2).key()));
-            verifyRecordsProcessingStats(convertedRecords.recordsProcessingStats(), 3, 2,
+            verifyRecordsProcessingStats(convertedRecords.recordConversionStats(), 3, 2,
                     records.sizeInBytes(), buffer.limit());
         } else {
             assertEquals(2, batches.size());
@@ -562,38 +654,35 @@ public void convertToV1WithMixedV0AndV2Data() {
             assertEquals(2, batches.get(1).baseOffset());
             assertEquals("1", utf8(logRecords.get(0).key()));
             assertEquals("3", utf8(logRecords.get(1).key()));
-            verifyRecordsProcessingStats(convertedRecords.recordsProcessingStats(), 3, 1,
+            verifyRecordsProcessingStats(convertedRecords.recordConversionStats(), 3, 1,
                     records.sizeInBytes(), buffer.limit());
         }
     }
 
-    private String utf8(ByteBuffer buffer) {
-        return Utils.utf8(buffer, buffer.remaining());
-    }
-
     @Test
-    public void shouldThrowIllegalStateExceptionOnBuildWhenAborted() throws Exception {
+    public void shouldThrowIllegalStateExceptionOnBuildWhenAborted() {
+        byte magic = RecordBatch.MAGIC_VALUE_V0;
+        assumeAtLeastV2OrNotZstd(magic);
+
         ByteBuffer buffer = ByteBuffer.allocate(128);
         buffer.position(bufferOffset);
 
-        MemoryRecordsBuilder builder = new MemoryRecordsBuilder(buffer, RecordBatch.MAGIC_VALUE_V0, compressionType,
+        MemoryRecordsBuilder builder = new MemoryRecordsBuilder(buffer, magic, compressionType,
                 TimestampType.CREATE_TIME, 0L, 0L, RecordBatch.NO_PRODUCER_ID, RecordBatch.NO_PRODUCER_EPOCH,
                 RecordBatch.NO_SEQUENCE, false, false, RecordBatch.NO_PARTITION_LEADER_EPOCH, buffer.capacity());
         builder.abort();
-        try {
-            builder.build();
-            fail("Should have thrown KafkaException");
-        } catch (IllegalStateException e) {
-            // ok
-        }
+        assertThrows(IllegalStateException.class, builder::build);
     }
 
     @Test
-    public void shouldResetBufferToInitialPositionOnAbort() throws Exception {
+    public void shouldResetBufferToInitialPositionOnAbort() {
+        byte magic = RecordBatch.MAGIC_VALUE_V0;
+        assumeAtLeastV2OrNotZstd(magic);
+
         ByteBuffer buffer = ByteBuffer.allocate(128);
         buffer.position(bufferOffset);
 
-        MemoryRecordsBuilder builder = new MemoryRecordsBuilder(buffer, RecordBatch.MAGIC_VALUE_V0, compressionType,
+        MemoryRecordsBuilder builder = new MemoryRecordsBuilder(buffer, magic, compressionType,
                                                                 TimestampType.CREATE_TIME, 0L, 0L, RecordBatch.NO_PRODUCER_ID, RecordBatch.NO_PRODUCER_EPOCH, RecordBatch.NO_SEQUENCE,
                                                                 false, false, RecordBatch.NO_PARTITION_LEADER_EPOCH, buffer.capacity());
         builder.append(0L, "a".getBytes(), "1".getBytes());
@@ -602,11 +691,14 @@ public void shouldResetBufferToInitialPositionOnAbort() throws Exception {
     }
 
     @Test
-    public void shouldThrowIllegalStateExceptionOnCloseWhenAborted() throws Exception {
+    public void shouldThrowIllegalStateExceptionOnCloseWhenAborted() {
+        byte magic = RecordBatch.MAGIC_VALUE_V0;
+        assumeAtLeastV2OrNotZstd(magic);
+
         ByteBuffer buffer = ByteBuffer.allocate(128);
         buffer.position(bufferOffset);
 
-        MemoryRecordsBuilder builder = new MemoryRecordsBuilder(buffer, RecordBatch.MAGIC_VALUE_V0, compressionType,
+        MemoryRecordsBuilder builder = new MemoryRecordsBuilder(buffer, magic, compressionType,
                                                                 TimestampType.CREATE_TIME, 0L, 0L, RecordBatch.NO_PRODUCER_ID, RecordBatch.NO_PRODUCER_EPOCH, RecordBatch.NO_SEQUENCE,
                                                                 false, false, RecordBatch.NO_PARTITION_LEADER_EPOCH, buffer.capacity());
         builder.abort();
@@ -619,11 +711,14 @@ public void shouldThrowIllegalStateExceptionOnCloseWhenAborted() throws Exceptio
     }
 
     @Test
-    public void shouldThrowIllegalStateExceptionOnAppendWhenAborted() throws Exception {
+    public void shouldThrowIllegalStateExceptionOnAppendWhenAborted() {
+        byte magic = RecordBatch.MAGIC_VALUE_V0;
+        assumeAtLeastV2OrNotZstd(magic);
+
         ByteBuffer buffer = ByteBuffer.allocate(128);
         buffer.position(bufferOffset);
 
-        MemoryRecordsBuilder builder = new MemoryRecordsBuilder(buffer, RecordBatch.MAGIC_VALUE_V0, compressionType,
+        MemoryRecordsBuilder builder = new MemoryRecordsBuilder(buffer, magic, compressionType,
                                                                 TimestampType.CREATE_TIME, 0L, 0L, RecordBatch.NO_PRODUCER_ID, RecordBatch.NO_PRODUCER_EPOCH, RecordBatch.NO_SEQUENCE,
                                                                 false, false, RecordBatch.NO_PARTITION_LEADER_EPOCH, buffer.capacity());
         builder.abort();
@@ -644,8 +739,41 @@ public static Collection data() {
         return values;
     }
 
-    private void verifyRecordsProcessingStats(RecordsProcessingStats processingStats, int numRecords,
-            int numRecordsConverted, long finalBytes, long preConvertedBytes) {
+    @Test
+    public void testBuffersDereferencedOnClose() {
+        Runtime runtime = Runtime.getRuntime();
+        int payloadLen = 1024 * 1024;
+        ByteBuffer buffer = ByteBuffer.allocate(payloadLen * 2);
+        byte[] key = new byte[0];
+        byte[] value = new byte[payloadLen];
+        new Random().nextBytes(value); // Use random payload so that compressed buffer is large
+        List builders = new ArrayList<>(100);
+        long startMem = 0;
+        long memUsed = 0;
+        int iterations =  0;
+        while (iterations++ < 100) {
+            buffer.rewind();
+            MemoryRecordsBuilder builder = new MemoryRecordsBuilder(buffer, RecordBatch.MAGIC_VALUE_V2, compressionType,
+                    TimestampType.CREATE_TIME, 0L, 0L, RecordBatch.NO_PRODUCER_ID,
+                    RecordBatch.NO_PRODUCER_EPOCH, RecordBatch.NO_SEQUENCE, false, false,
+                    RecordBatch.NO_PARTITION_LEADER_EPOCH, 0);
+            builder.append(1L, new byte[0], value);
+            builder.build();
+            builders.add(builder);
+
+            System.gc();
+            memUsed = runtime.totalMemory() - runtime.freeMemory() - startMem;
+            // Ignore memory usage during initialization
+            if (iterations == 2)
+                startMem = memUsed;
+            else if (iterations > 2 && memUsed < (iterations - 2) * 1024)
+                break;
+        }
+        assertTrue("Memory usage too high: " + memUsed, iterations < 100);
+    }
+
+    private void verifyRecordsProcessingStats(RecordConversionStats processingStats, int numRecords,
+                                              int numRecordsConverted, long finalBytes, long preConvertedBytes) {
         assertNotNull("Records processing info is null", processingStats);
         assertEquals(numRecordsConverted, processingStats.numRecordsConverted());
         // Since nanoTime accuracy on build machines may not be sufficient to measure small conversion times,
@@ -668,4 +796,7 @@ else if (numRecordsConverted == numRecords)
         }
     }
 
+    private void assumeAtLeastV2OrNotZstd(byte magic) {
+        assumeTrue(compressionType != CompressionType.ZSTD || magic >= MAGIC_VALUE_V2);
+    }
 }
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 de00378880208..e432a7af904bb 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
@@ -17,7 +17,10 @@
 package org.apache.kafka.common.record;
 
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.errors.CorruptRecordException;
 import org.apache.kafka.common.header.internals.RecordHeaders;
+import org.apache.kafka.common.message.LeaderChangeMessage;
+import org.apache.kafka.common.message.LeaderChangeMessage.Voter;
 import org.apache.kafka.common.record.MemoryRecords.RecordFilter.BatchRetention;
 import org.apache.kafka.common.utils.Utils;
 import org.apache.kafka.test.TestUtils;
@@ -29,18 +32,22 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
+import java.util.Collections;
 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;
 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.Assume.assumeTrue;
 
 @RunWith(value = Parameterized.class)
 public class MemoryRecordsTest {
-
     private CompressionType compression;
     private byte magic;
     private long firstOffset;
@@ -67,6 +74,8 @@ public MemoryRecordsTest(byte magic, long firstOffset, CompressionType compressi
 
     @Test
     public void testIterator() {
+        assumeAtLeastV2OrNotZstd();
+
         ByteBuffer buffer = ByteBuffer.allocate(1024);
 
         MemoryRecordsBuilder builder = new MemoryRecordsBuilder(buffer, magic, compression,
@@ -150,6 +159,7 @@ public void testIterator() {
 
     @Test
     public void testHasRoomForMethod() {
+        assumeAtLeastV2OrNotZstd();
         MemoryRecordsBuilder builder = MemoryRecords.builder(ByteBuffer.allocate(1024), magic, compression,
                 TimestampType.CREATE_TIME, 0L);
         builder.append(0L, "a".getBytes(), "1".getBytes());
@@ -250,6 +260,7 @@ public void testFilterToEmptyBatchRetention() {
                 long baseOffset = 3L;
                 int baseSequence = 10;
                 int partitionLeaderEpoch = 293;
+                int numRecords = 2;
 
                 MemoryRecordsBuilder builder = MemoryRecords.builder(buffer, magic, compression, TimestampType.CREATE_TIME,
                         baseOffset, RecordBatch.NO_TIMESTAMP, producerId, producerEpoch, baseSequence, isTransactional,
@@ -257,22 +268,34 @@ public void testFilterToEmptyBatchRetention() {
                 builder.append(11L, "2".getBytes(), "b".getBytes());
                 builder.append(12L, "3".getBytes(), "c".getBytes());
                 builder.close();
+                MemoryRecords records = builder.build();
 
                 ByteBuffer filtered = ByteBuffer.allocate(2048);
-                builder.build().filterTo(new TopicPartition("foo", 0), new MemoryRecords.RecordFilter() {
-                    @Override
-                    protected BatchRetention checkBatchRetention(RecordBatch batch) {
-                        // retain all batches
-                        return BatchRetention.RETAIN_EMPTY;
-                    }
-
-                    @Override
-                    protected boolean shouldRetainRecord(RecordBatch recordBatch, Record record) {
-                        // delete the records
-                        return false;
-                    }
-                }, filtered, Integer.MAX_VALUE, BufferSupplier.NO_CACHING);
-
+                MemoryRecords.FilterResult filterResult = records.filterTo(new TopicPartition("foo", 0),
+                        new MemoryRecords.RecordFilter() {
+                            @Override
+                            protected BatchRetention checkBatchRetention(RecordBatch batch) {
+                                // retain all batches
+                                return BatchRetention.RETAIN_EMPTY;
+                            }
+
+                            @Override
+                            protected boolean shouldRetainRecord(RecordBatch recordBatch, Record record) {
+                                // delete the records
+                                return false;
+                            }
+                        }, filtered, Integer.MAX_VALUE, BufferSupplier.NO_CACHING);
+
+                // Verify filter result
+                assertEquals(numRecords, filterResult.messagesRead());
+                assertEquals(records.sizeInBytes(), filterResult.bytesRead());
+                assertEquals(baseOffset + 1, filterResult.maxOffset());
+                assertEquals(0, filterResult.messagesRetained());
+                assertEquals(DefaultRecordBatch.RECORD_BATCH_OVERHEAD, filterResult.bytesRetained());
+                assertEquals(12, filterResult.maxTimestamp());
+                assertEquals(baseOffset + 1, filterResult.shallowOffsetOfMaxTimestamp());
+
+                // Verify filtered records
                 filtered.flip();
                 MemoryRecords filteredRecords = MemoryRecords.readableRecords(filtered);
 
@@ -292,6 +315,55 @@ protected boolean shouldRetainRecord(RecordBatch recordBatch, Record record) {
         }
     }
 
+    @Test
+    public void testEmptyBatchRetention() {
+        if (magic >= RecordBatch.MAGIC_VALUE_V2) {
+            ByteBuffer buffer = ByteBuffer.allocate(DefaultRecordBatch.RECORD_BATCH_OVERHEAD);
+            long producerId = 23L;
+            short producerEpoch = 5;
+            long baseOffset = 3L;
+            int baseSequence = 10;
+            int partitionLeaderEpoch = 293;
+            long timestamp = System.currentTimeMillis();
+
+            DefaultRecordBatch.writeEmptyHeader(buffer, RecordBatch.MAGIC_VALUE_V2, producerId, producerEpoch,
+                    baseSequence, baseOffset, baseOffset, partitionLeaderEpoch, TimestampType.CREATE_TIME,
+                    timestamp, false, false);
+            buffer.flip();
+
+            ByteBuffer filtered = ByteBuffer.allocate(2048);
+            MemoryRecords records = MemoryRecords.readableRecords(buffer);
+            MemoryRecords.FilterResult filterResult = records.filterTo(new TopicPartition("foo", 0),
+                    new MemoryRecords.RecordFilter() {
+                        @Override
+                        protected BatchRetention checkBatchRetention(RecordBatch batch) {
+                            // retain all batches
+                            return BatchRetention.RETAIN_EMPTY;
+                        }
+
+                        @Override
+                        protected boolean shouldRetainRecord(RecordBatch recordBatch, Record record) {
+                            return false;
+                        }
+                    }, filtered, Integer.MAX_VALUE, BufferSupplier.NO_CACHING);
+
+            // Verify filter result
+            assertEquals(0, filterResult.messagesRead());
+            assertEquals(records.sizeInBytes(), filterResult.bytesRead());
+            assertEquals(baseOffset, filterResult.maxOffset());
+            assertEquals(0, filterResult.messagesRetained());
+            assertEquals(DefaultRecordBatch.RECORD_BATCH_OVERHEAD, filterResult.bytesRetained());
+            assertEquals(timestamp, filterResult.maxTimestamp());
+            assertEquals(baseOffset, filterResult.shallowOffsetOfMaxTimestamp());
+            assertTrue(filterResult.outputBuffer().position() > 0);
+
+            // Verify filtered records
+            filtered.flip();
+            MemoryRecords filteredRecords = MemoryRecords.readableRecords(filtered);
+            assertEquals(DefaultRecordBatch.RECORD_BATCH_OVERHEAD, filteredRecords.sizeInBytes());
+        }
+    }
+
     @Test
     public void testEmptyBatchDeletion() {
         if (magic >= RecordBatch.MAGIC_VALUE_V2) {
@@ -302,25 +374,32 @@ public void testEmptyBatchDeletion() {
                 long baseOffset = 3L;
                 int baseSequence = 10;
                 int partitionLeaderEpoch = 293;
+                long timestamp = System.currentTimeMillis();
 
                 DefaultRecordBatch.writeEmptyHeader(buffer, RecordBatch.MAGIC_VALUE_V2, producerId, producerEpoch,
                         baseSequence, baseOffset, baseOffset, partitionLeaderEpoch, TimestampType.CREATE_TIME,
-                        System.currentTimeMillis(), false, false);
+                        timestamp, false, false);
                 buffer.flip();
 
                 ByteBuffer filtered = ByteBuffer.allocate(2048);
-                MemoryRecords.readableRecords(buffer).filterTo(new TopicPartition("foo", 0), new MemoryRecords.RecordFilter() {
-                    @Override
-                    protected BatchRetention checkBatchRetention(RecordBatch batch) {
-                        return deleteRetention;
-                    }
-
-                    @Override
-                    protected boolean shouldRetainRecord(RecordBatch recordBatch, Record record) {
-                        return false;
-                    }
-                }, filtered, Integer.MAX_VALUE, BufferSupplier.NO_CACHING);
-
+                MemoryRecords records = MemoryRecords.readableRecords(buffer);
+                MemoryRecords.FilterResult filterResult = records.filterTo(new TopicPartition("foo", 0),
+                        new MemoryRecords.RecordFilter() {
+                            @Override
+                            protected BatchRetention checkBatchRetention(RecordBatch batch) {
+                                return deleteRetention;
+                            }
+
+                            @Override
+                            protected boolean shouldRetainRecord(RecordBatch recordBatch, Record record) {
+                                return false;
+                            }
+                        }, filtered, Integer.MAX_VALUE, BufferSupplier.NO_CACHING);
+
+                // Verify filter result
+                assertEquals(0, filterResult.outputBuffer().position());
+
+                // Verify filtered records
                 filtered.flip();
                 MemoryRecords filteredRecords = MemoryRecords.readableRecords(filtered);
                 assertEquals(0, filteredRecords.sizeInBytes());
@@ -365,59 +444,100 @@ public void testBuildEndTxnMarker() {
         }
     }
 
+    @Test
+    public void testBuildLeaderChangeMessage() {
+        if (magic >= RecordBatch.MAGIC_VALUE_V2) {
+
+            final int leaderId = 5;
+            final int leaderEpoch = 20;
+            final int voterId = 6;
+
+            LeaderChangeMessage leaderChangeMessage = new LeaderChangeMessage()
+                .setLeaderId(leaderId)
+                .setVoters(Collections.singletonList(
+                    new Voter().setVoterId(voterId)));
+            MemoryRecords records = MemoryRecords.withLeaderChangeMessage(System.currentTimeMillis(),
+                leaderEpoch, leaderChangeMessage);
+
+            List batches = TestUtils.toList(records.batches());
+            assertEquals(1, batches.size());
+
+            RecordBatch batch = batches.get(0);
+            assertTrue(batch.isControlBatch());
+            assertEquals(0, batch.baseOffset());
+            assertEquals(leaderEpoch, batch.partitionLeaderEpoch());
+            assertTrue(batch.isValid());
+
+            List createdRecords = TestUtils.toList(batch);
+            assertEquals(1, createdRecords.size());
+
+            Record record = createdRecords.get(0);
+            assertTrue(record.isValid());
+            assertEquals(ControlRecordType.LEADER_CHANGE, ControlRecordType.parse(record.key()));
+
+            LeaderChangeMessage deserializedMessage = ControlRecordUtils.deserializeLeaderChangeMessage(record);
+            assertEquals(leaderId, deserializedMessage.leaderId());
+            assertEquals(1, deserializedMessage.voters().size());
+            assertEquals(voterId, deserializedMessage.voters().get(0).voterId());
+        }
+    }
+
     @Test
     public void testFilterToBatchDiscard() {
-        if (compression != CompressionType.NONE || magic >= RecordBatch.MAGIC_VALUE_V2) {
-            ByteBuffer buffer = ByteBuffer.allocate(2048);
-            MemoryRecordsBuilder builder = MemoryRecords.builder(buffer, magic, compression, TimestampType.CREATE_TIME, 0L);
-            builder.append(10L, "1".getBytes(), "a".getBytes());
-            builder.close();
+        assumeAtLeastV2OrNotZstd();
+        assumeTrue(compression != CompressionType.NONE || magic >= MAGIC_VALUE_V2);
 
-            builder = MemoryRecords.builder(buffer, magic, compression, TimestampType.CREATE_TIME, 1L);
-            builder.append(11L, "2".getBytes(), "b".getBytes());
-            builder.append(12L, "3".getBytes(), "c".getBytes());
-            builder.close();
+        ByteBuffer buffer = ByteBuffer.allocate(2048);
+        MemoryRecordsBuilder builder = MemoryRecords.builder(buffer, magic, compression, TimestampType.CREATE_TIME, 0L);
+        builder.append(10L, "1".getBytes(), "a".getBytes());
+        builder.close();
 
-            builder = MemoryRecords.builder(buffer, magic, compression, TimestampType.CREATE_TIME, 3L);
-            builder.append(13L, "4".getBytes(), "d".getBytes());
-            builder.append(20L, "5".getBytes(), "e".getBytes());
-            builder.append(15L, "6".getBytes(), "f".getBytes());
-            builder.close();
+        builder = MemoryRecords.builder(buffer, magic, compression, TimestampType.CREATE_TIME, 1L);
+        builder.append(11L, "2".getBytes(), "b".getBytes());
+        builder.append(12L, "3".getBytes(), "c".getBytes());
+        builder.close();
 
-            builder = MemoryRecords.builder(buffer, magic, compression, TimestampType.CREATE_TIME, 6L);
-            builder.append(16L, "7".getBytes(), "g".getBytes());
-            builder.close();
+        builder = MemoryRecords.builder(buffer, magic, compression, TimestampType.CREATE_TIME, 3L);
+        builder.append(13L, "4".getBytes(), "d".getBytes());
+        builder.append(20L, "5".getBytes(), "e".getBytes());
+        builder.append(15L, "6".getBytes(), "f".getBytes());
+        builder.close();
 
-            buffer.flip();
+        builder = MemoryRecords.builder(buffer, magic, compression, TimestampType.CREATE_TIME, 6L);
+        builder.append(16L, "7".getBytes(), "g".getBytes());
+        builder.close();
 
-            ByteBuffer filtered = ByteBuffer.allocate(2048);
-            MemoryRecords.readableRecords(buffer).filterTo(new TopicPartition("foo", 0), new MemoryRecords.RecordFilter() {
-                @Override
-                protected BatchRetention checkBatchRetention(RecordBatch batch) {
-                    // discard the second and fourth batches
-                    if (batch.lastOffset() == 2L || batch.lastOffset() == 6L)
-                        return BatchRetention.DELETE;
-                    return BatchRetention.DELETE_EMPTY;
-                }
+        buffer.flip();
 
-                @Override
-                protected boolean shouldRetainRecord(RecordBatch recordBatch, Record record) {
-                    return true;
-                }
-            }, filtered, Integer.MAX_VALUE, BufferSupplier.NO_CACHING);
+        ByteBuffer filtered = ByteBuffer.allocate(2048);
+        MemoryRecords.readableRecords(buffer).filterTo(new TopicPartition("foo", 0), new MemoryRecords.RecordFilter() {
+            @Override
+            protected BatchRetention checkBatchRetention(RecordBatch batch) {
+                // discard the second and fourth batches
+                if (batch.lastOffset() == 2L || batch.lastOffset() == 6L)
+                    return BatchRetention.DELETE;
+                return BatchRetention.DELETE_EMPTY;
+            }
 
-            filtered.flip();
-            MemoryRecords filteredRecords = MemoryRecords.readableRecords(filtered);
+            @Override
+            protected boolean shouldRetainRecord(RecordBatch recordBatch, Record record) {
+                return true;
+            }
+        }, filtered, Integer.MAX_VALUE, BufferSupplier.NO_CACHING);
 
-            List batches = TestUtils.toList(filteredRecords.batches());
-            assertEquals(2, batches.size());
-            assertEquals(0L, batches.get(0).lastOffset());
-            assertEquals(5L, batches.get(1).lastOffset());
-        }
+        filtered.flip();
+        MemoryRecords filteredRecords = MemoryRecords.readableRecords(filtered);
+
+        List batches = TestUtils.toList(filteredRecords.batches());
+        assertEquals(2, batches.size());
+        assertEquals(0L, batches.get(0).lastOffset());
+        assertEquals(5L, batches.get(1).lastOffset());
     }
 
     @Test
     public void testFilterToAlreadyCompactedLog() {
+        assumeAtLeastV2OrNotZstd();
+
         ByteBuffer buffer = ByteBuffer.allocate(2048);
 
         // create a batch with some offset gaps to simulate a compacted batch
@@ -558,6 +678,8 @@ public void testFilterToPreservesProducerInfo() {
 
     @Test
     public void testFilterToWithUndersizedBuffer() {
+        assumeAtLeastV2OrNotZstd();
+
         ByteBuffer buffer = ByteBuffer.allocate(1024);
         MemoryRecordsBuilder builder = MemoryRecords.builder(buffer, magic, compression, TimestampType.CREATE_TIME, 0L);
         builder.append(10L, null, "a".getBytes());
@@ -589,15 +711,15 @@ public void testFilterToWithUndersizedBuffer() {
 
             MemoryRecords.FilterResult result = MemoryRecords.readableRecords(buffer)
                     .filterTo(new TopicPartition("foo", 0), new RetainNonNullKeysFilter(), output, Integer.MAX_VALUE,
-                              BufferSupplier.NO_CACHING);
+                            BufferSupplier.NO_CACHING);
 
-            buffer.position(buffer.position() + result.bytesRead);
-            result.output.flip();
+            buffer.position(buffer.position() + result.bytesRead());
+            result.outputBuffer().flip();
 
-            if (output != result.output)
+            if (output != result.outputBuffer())
                 assertEquals(0, output.position());
 
-            MemoryRecords filtered = MemoryRecords.readableRecords(result.output);
+            MemoryRecords filtered = MemoryRecords.readableRecords(result.outputBuffer());
             records.addAll(TestUtils.toList(filtered.records()));
         }
 
@@ -608,6 +730,8 @@ public void testFilterToWithUndersizedBuffer() {
 
     @Test
     public void testFilterTo() {
+        assumeAtLeastV2OrNotZstd();
+
         ByteBuffer buffer = ByteBuffer.allocate(2048);
         MemoryRecordsBuilder builder = MemoryRecords.builder(buffer, magic, compression, TimestampType.CREATE_TIME, 0L);
         builder.append(10L, null, "a".getBytes());
@@ -637,16 +761,16 @@ public void testFilterTo() {
 
         filtered.flip();
 
-        assertEquals(7, result.messagesRead);
-        assertEquals(4, result.messagesRetained);
-        assertEquals(buffer.limit(), result.bytesRead);
-        assertEquals(filtered.limit(), result.bytesRetained);
+        assertEquals(7, result.messagesRead());
+        assertEquals(4, result.messagesRetained());
+        assertEquals(buffer.limit(), result.bytesRead());
+        assertEquals(filtered.limit(), result.bytesRetained());
         if (magic > RecordBatch.MAGIC_VALUE_V0) {
-            assertEquals(20L, result.maxTimestamp);
+            assertEquals(20L, result.maxTimestamp());
             if (compression == CompressionType.NONE && magic < RecordBatch.MAGIC_VALUE_V2)
-                assertEquals(4L, result.shallowOffsetOfMaxTimestamp);
+                assertEquals(4L, result.shallowOffsetOfMaxTimestamp());
             else
-                assertEquals(5L, result.shallowOffsetOfMaxTimestamp);
+                assertEquals(5L, result.shallowOffsetOfMaxTimestamp());
         }
 
         MemoryRecords filteredRecords = MemoryRecords.readableRecords(filtered);
@@ -721,6 +845,8 @@ public void testFilterTo() {
 
     @Test
     public void testFilterToPreservesLogAppendTime() {
+        assumeAtLeastV2OrNotZstd();
+
         long logAppendTime = System.currentTimeMillis();
 
         ByteBuffer buffer = ByteBuffer.allocate(2048);
@@ -763,6 +889,56 @@ public void testFilterToPreservesLogAppendTime() {
         }
     }
 
+    @Test
+    public void testNextBatchSize() {
+        assumeAtLeastV2OrNotZstd();
+
+        ByteBuffer buffer = ByteBuffer.allocate(2048);
+        MemoryRecordsBuilder builder = MemoryRecords.builder(buffer, magic, compression,
+                TimestampType.LOG_APPEND_TIME, 0L, logAppendTime, pid, epoch, firstSequence);
+        builder.append(10L, null, "abc".getBytes());
+        builder.close();
+
+        buffer.flip();
+        int size = buffer.remaining();
+        MemoryRecords records = MemoryRecords.readableRecords(buffer);
+        assertEquals(size, records.firstBatchSize().intValue());
+        assertEquals(0, buffer.position());
+
+        buffer.limit(1); // size not in buffer
+        assertNull(records.firstBatchSize());
+        buffer.limit(Records.LOG_OVERHEAD); // magic not in buffer
+        assertNull(records.firstBatchSize());
+        buffer.limit(Records.HEADER_SIZE_UP_TO_MAGIC); // payload not in buffer
+        assertEquals(size, records.firstBatchSize().intValue());
+
+        buffer.limit(size);
+        byte magic = buffer.get(Records.MAGIC_OFFSET);
+        buffer.put(Records.MAGIC_OFFSET, (byte) 10);
+        assertThrows(CorruptRecordException.class, records::firstBatchSize);
+        buffer.put(Records.MAGIC_OFFSET, magic);
+
+        buffer.put(Records.SIZE_OFFSET + 3, (byte) 0);
+        assertThrows(CorruptRecordException.class, records::firstBatchSize);
+    }
+
+    @Test
+    public void testWithRecords() {
+        Supplier recordsSupplier = () -> MemoryRecords.withRecords(magic, compression,
+            new SimpleRecord(10L, "key1".getBytes(), "value1".getBytes()));
+        if (compression != CompressionType.ZSTD || magic >= MAGIC_VALUE_V2) {
+            MemoryRecords memoryRecords = recordsSupplier.get();
+            String key = Utils.utf8(memoryRecords.batches().iterator().next().iterator().next().key());
+            assertEquals("key1", key);
+        } else {
+            assertThrows(IllegalArgumentException.class, recordsSupplier::get);
+        }
+    }
+
+    private void assumeAtLeastV2OrNotZstd() {
+        assumeTrue(compression != CompressionType.ZSTD || magic >= MAGIC_VALUE_V2);
+    }
+
     @Parameterized.Parameters(name = "{index} magic={0}, firstOffset={1}, compressionType={2}")
     public static Collection data() {
         List values = new ArrayList<>();
diff --git a/clients/src/test/java/org/apache/kafka/common/record/MultiRecordsSendTest.java b/clients/src/test/java/org/apache/kafka/common/record/MultiRecordsSendTest.java
new file mode 100644
index 0000000000000..348bbf9a06562
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/record/MultiRecordsSendTest.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.record;
+
+import org.apache.kafka.common.network.ByteBufferSend;
+import org.apache.kafka.common.network.Send;
+import org.apache.kafka.test.TestUtils;
+import org.junit.Test;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.LinkedList;
+import java.util.Queue;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+public class MultiRecordsSendTest {
+
+    @Test
+    public void testSendsFreedAfterWriting() throws IOException {
+        int numChunks = 4;
+        int chunkSize = 32;
+        int totalSize = numChunks * chunkSize;
+
+        Queue sends = new LinkedList<>();
+        ByteBuffer[] chunks = new ByteBuffer[numChunks];
+
+        for (int i = 0; i < numChunks; i++) {
+            ByteBuffer buffer = ByteBuffer.wrap(TestUtils.randomBytes(chunkSize));
+            chunks[i] = buffer;
+            sends.add(new ByteBufferSend(buffer));
+        }
+
+        MultiRecordsSend send = new MultiRecordsSend(sends);
+        assertEquals(totalSize, send.size());
+
+        for (int i = 0; i < numChunks; i++) {
+            assertEquals(numChunks - i, send.numResidentSends());
+            NonOverflowingByteBufferChannel out = new NonOverflowingByteBufferChannel(chunkSize);
+            send.writeTo(out);
+            out.close();
+            assertEquals(chunks[i], out.buffer());
+        }
+
+        assertEquals(0, send.numResidentSends());
+        assertTrue(send.completed());
+    }
+
+    private static class NonOverflowingByteBufferChannel extends org.apache.kafka.common.requests.ByteBufferChannel {
+
+        private NonOverflowingByteBufferChannel(long size) {
+            super(size);
+        }
+
+        @Override
+        public long write(ByteBuffer[] srcs) {
+            // Instead of overflowing, this channel refuses additional writes once the buffer is full,
+            // which allows us to test the MultiRecordsSend behavior on a per-send basis.
+            if (!buffer().hasRemaining())
+                return 0;
+            return super.write(srcs);
+        }
+    }
+
+}
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/AddPartitionsToTxnRequestTest.java b/clients/src/test/java/org/apache/kafka/common/requests/AddPartitionsToTxnRequestTest.java
new file mode 100644
index 0000000000000..1758a978f0b2b
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/AddPartitionsToTxnRequestTest.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.TopicPartition;
+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.List;
+
+import static org.junit.Assert.assertEquals;
+
+public class AddPartitionsToTxnRequestTest {
+
+    private static String transactionalId = "transactionalId";
+    private static int producerId = 10;
+    private static short producerEpoch = 1;
+    private static int throttleTimeMs = 10;
+
+    @Test
+    public void testConstructor() {
+        List partitions = new ArrayList<>();
+        partitions.add(new TopicPartition("topic", 0));
+        partitions.add(new TopicPartition("topic", 1));
+
+        AddPartitionsToTxnRequest.Builder builder = new AddPartitionsToTxnRequest.Builder(transactionalId, producerId, producerEpoch, partitions);
+
+        for (short version = 0; version <= ApiKeys.ADD_PARTITIONS_TO_TXN.latestVersion(); version++) {
+            AddPartitionsToTxnRequest request = builder.build(version);
+
+            assertEquals(transactionalId, request.data().transactionalId());
+            assertEquals(producerId, request.data().producerId());
+            assertEquals(producerEpoch, request.data().producerEpoch());
+            assertEquals(partitions, request.partitions());
+
+            AddPartitionsToTxnResponse response = request.getErrorResponse(throttleTimeMs, Errors.UNKNOWN_TOPIC_OR_PARTITION.exception());
+
+            assertEquals(Collections.singletonMap(Errors.UNKNOWN_TOPIC_OR_PARTITION, 2), response.errorCounts());
+            assertEquals(throttleTimeMs, response.throttleTimeMs());
+        }
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/AddPartitionsToTxnResponseTest.java b/clients/src/test/java/org/apache/kafka/common/requests/AddPartitionsToTxnResponseTest.java
new file mode 100644
index 0000000000000..cd9434f5a76fd
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/AddPartitionsToTxnResponseTest.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.AddPartitionsToTxnResponseData;
+import org.apache.kafka.common.message.AddPartitionsToTxnResponseData.AddPartitionsToTxnPartitionResult;
+import org.apache.kafka.common.message.AddPartitionsToTxnResponseData.AddPartitionsToTxnTopicResult;
+import org.apache.kafka.common.message.AddPartitionsToTxnResponseData.AddPartitionsToTxnTopicResultCollection;
+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.HashMap;
+import java.util.Map;
+
+import static org.junit.Assert.assertEquals;
+
+public class AddPartitionsToTxnResponseTest {
+
+    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() {
+        AddPartitionsToTxnResponse response = new AddPartitionsToTxnResponse(throttleTimeMs, errorsMap);
+
+        assertEquals(expectedErrorCounts, response.errorCounts());
+        assertEquals(throttleTimeMs, response.throttleTimeMs());
+    }
+
+    @Test
+    public void testParse() {
+
+        AddPartitionsToTxnTopicResultCollection topicCollection = new AddPartitionsToTxnTopicResultCollection();
+
+        AddPartitionsToTxnTopicResult topicResult = new AddPartitionsToTxnTopicResult();
+        topicResult.setName(topicOne);
+
+        topicResult.results().add(new AddPartitionsToTxnPartitionResult()
+                                      .setErrorCode(errorOne.code())
+                                      .setPartitionIndex(partitionOne));
+
+        topicResult.results().add(new AddPartitionsToTxnPartitionResult()
+                                      .setErrorCode(errorTwo.code())
+                                      .setPartitionIndex(partitionTwo));
+
+        topicCollection.add(topicResult);
+
+        AddPartitionsToTxnResponseData data = new AddPartitionsToTxnResponseData()
+                                                  .setResults(topicCollection)
+                                                  .setThrottleTimeMs(throttleTimeMs);
+        AddPartitionsToTxnResponse response = new AddPartitionsToTxnResponse(data);
+
+        for (short version = 0; version <= ApiKeys.ADD_PARTITIONS_TO_TXN.latestVersion(); version++) {
+            AddPartitionsToTxnResponse parsedResponse = AddPartitionsToTxnResponse.parse(response.serialize(version), version);
+            assertEquals(expectedErrorCounts, parsedResponse.errorCounts());
+            assertEquals(throttleTimeMs, parsedResponse.throttleTimeMs());
+            assertEquals(version >= 1, parsedResponse.shouldClientThrottle(version));
+        }
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/AlterReplicaLogDirsRequestTest.java b/clients/src/test/java/org/apache/kafka/common/requests/AlterReplicaLogDirsRequestTest.java
new file mode 100644
index 0000000000000..5ed919938a313
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/AlterReplicaLogDirsRequestTest.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.requests;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.errors.LogDirNotFoundException;
+import org.apache.kafka.common.message.AlterReplicaLogDirsRequestData;
+import org.apache.kafka.common.message.AlterReplicaLogDirsRequestData.AlterReplicaLogDir;
+import org.apache.kafka.common.message.AlterReplicaLogDirsRequestData.AlterReplicaLogDirCollection;
+import org.apache.kafka.common.message.AlterReplicaLogDirsRequestData.AlterReplicaLogDirTopic;
+import org.apache.kafka.common.message.AlterReplicaLogDirsRequestData.AlterReplicaLogDirTopicCollection;
+import org.apache.kafka.common.message.AlterReplicaLogDirsResponseData.AlterReplicaLogDirTopicResult;
+import org.apache.kafka.common.protocol.Errors;
+import org.junit.Test;
+
+import static java.util.Arrays.asList;
+import static java.util.Collections.singletonList;
+import static org.junit.Assert.assertEquals;
+
+public class AlterReplicaLogDirsRequestTest {
+
+    @Test
+    public void testErrorResponse() {
+        AlterReplicaLogDirsRequestData data = new AlterReplicaLogDirsRequestData()
+                .setDirs(new AlterReplicaLogDirCollection(
+                        singletonList(new AlterReplicaLogDir()
+                                .setPath("/data0")
+                                .setTopics(new AlterReplicaLogDirTopicCollection(
+                                        singletonList(new AlterReplicaLogDirTopic()
+                                                .setName("topic")
+                                                .setPartitions(asList(0, 1, 2))).iterator()))).iterator()));
+        AlterReplicaLogDirsResponse errorResponse = new AlterReplicaLogDirsRequest.Builder(data).build()
+                .getErrorResponse(123, new LogDirNotFoundException("/data0"));
+        assertEquals(1, errorResponse.data().results().size());
+        AlterReplicaLogDirTopicResult topicResponse = errorResponse.data().results().get(0);
+        assertEquals("topic", topicResponse.topicName());
+        assertEquals(3, topicResponse.partitions().size());
+        for (int i = 0; i < 3; i++) {
+            assertEquals(i, topicResponse.partitions().get(i).partitionIndex());
+            assertEquals(Errors.LOG_DIR_NOT_FOUND.code(), topicResponse.partitions().get(i).errorCode());
+        }
+    }
+
+    @Test
+    public void testPartitionDir() {
+        AlterReplicaLogDirsRequestData data = new AlterReplicaLogDirsRequestData()
+                .setDirs(new AlterReplicaLogDirCollection(
+                        asList(new AlterReplicaLogDir()
+                                .setPath("/data0")
+                                .setTopics(new AlterReplicaLogDirTopicCollection(
+                                        asList(new AlterReplicaLogDirTopic()
+                                                .setName("topic")
+                                                .setPartitions(asList(0, 1)),
+                                               new AlterReplicaLogDirTopic()
+                                                .setName("topic2")
+                                                .setPartitions(asList(7))).iterator())),
+                                new AlterReplicaLogDir()
+                                        .setPath("/data1")
+                                        .setTopics(new AlterReplicaLogDirTopicCollection(
+                                                asList(new AlterReplicaLogDirTopic()
+                                                        .setName("topic3")
+                                                        .setPartitions(asList(12))).iterator()))).iterator()));
+        AlterReplicaLogDirsRequest request = new AlterReplicaLogDirsRequest.Builder(data).build();
+        Map expect = new HashMap<>();
+        expect.put(new TopicPartition("topic", 0), "/data0");
+        expect.put(new TopicPartition("topic", 1), "/data0");
+        expect.put(new TopicPartition("topic2", 7), "/data0");
+        expect.put(new TopicPartition("topic3", 12), "/data1");
+        assertEquals(expect, request.partitionDirs());
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/AlterReplicaLogDirsResponseTest.java b/clients/src/test/java/org/apache/kafka/common/requests/AlterReplicaLogDirsResponseTest.java
new file mode 100644
index 0000000000000..01826c69bc2ce
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/AlterReplicaLogDirsResponseTest.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.requests;
+
+import java.util.Map;
+
+import org.apache.kafka.common.message.AlterReplicaLogDirsResponseData;
+import org.apache.kafka.common.message.AlterReplicaLogDirsResponseData.AlterReplicaLogDirPartitionResult;
+import org.apache.kafka.common.message.AlterReplicaLogDirsResponseData.AlterReplicaLogDirTopicResult;
+import org.apache.kafka.common.protocol.Errors;
+import org.junit.Test;
+
+import static java.util.Arrays.asList;
+import static org.junit.Assert.assertEquals;
+
+public class AlterReplicaLogDirsResponseTest {
+
+    @Test
+    public void testErrorCounts() {
+        AlterReplicaLogDirsResponseData data = new AlterReplicaLogDirsResponseData()
+                .setResults(asList(
+                        new AlterReplicaLogDirTopicResult()
+                                .setTopicName("t0")
+                                .setPartitions(asList(
+                                        new AlterReplicaLogDirPartitionResult()
+                                                .setPartitionIndex(0)
+                                                .setErrorCode(Errors.LOG_DIR_NOT_FOUND.code()),
+                                        new AlterReplicaLogDirPartitionResult()
+                                                .setPartitionIndex(1)
+                                                .setErrorCode(Errors.NONE.code()))),
+                        new AlterReplicaLogDirTopicResult()
+                                .setTopicName("t1")
+                                .setPartitions(asList(
+                                        new AlterReplicaLogDirPartitionResult()
+                                                .setPartitionIndex(0)
+                                                .setErrorCode(Errors.LOG_DIR_NOT_FOUND.code())))));
+        Map counts = new AlterReplicaLogDirsResponse(data).errorCounts();
+        assertEquals(2, counts.size());
+        assertEquals(Integer.valueOf(2), counts.get(Errors.LOG_DIR_NOT_FOUND));
+        assertEquals(Integer.valueOf(1), counts.get(Errors.NONE));
+
+    }
+}
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..b578e2c39db69 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,9 +17,8 @@
 
 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;
 import org.junit.Test;
 
 import java.util.Collection;
@@ -31,66 +30,50 @@
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 
-
 public class ApiVersionsResponseTest {
 
-    @Test
-    public void shouldCreateApiResponseOnlyWithKeysSupportedByMagicValue() {
-        final ApiVersionsResponse response = ApiVersionsResponse.apiVersionsResponse(10, RecordBatch.MAGIC_VALUE_V1);
-        verifyApiKeysForMagic(response, RecordBatch.MAGIC_VALUE_V1);
-        assertEquals(10, response.throttleTimeMs());
-    }
-
     @Test
     public void shouldCreateApiResponseThatHasAllApiKeysSupportedByBroker() {
-        assertEquals(apiKeysInResponse(ApiVersionsResponse.defaultApiVersionsResponse()), Utils.mkSet(ApiKeys.values()));
-    }
-
-    @Test
-    public void shouldReturnAllKeysWhenMagicIsCurrentValueAndThrottleMsIsDefaultThrottle() {
-        ApiVersionsResponse response = ApiVersionsResponse.apiVersionsResponse(AbstractResponse.DEFAULT_THROTTLE_TIME, RecordBatch.CURRENT_MAGIC_VALUE);
-        assertEquals(Utils.mkSet(ApiKeys.values()), apiKeysInResponse(response));
-        assertEquals(AbstractResponse.DEFAULT_THROTTLE_TIME, response.throttleTimeMs());
+        assertEquals(apiKeysInResponse(ApiVersionsResponse.DEFAULT_API_VERSIONS_RESPONSE), new HashSet<>(ApiKeys.enabledApis()));
+        assertTrue(ApiVersionsResponse.DEFAULT_API_VERSIONS_RESPONSE.data().supportedFeatures().isEmpty());
+        assertTrue(ApiVersionsResponse.DEFAULT_API_VERSIONS_RESPONSE.data().finalizedFeatures().isEmpty());
+        assertEquals(ApiVersionsResponse.UNKNOWN_FINALIZED_FEATURES_EPOCH, ApiVersionsResponse.DEFAULT_API_VERSIONS_RESPONSE.data().finalizedFeaturesEpoch());
     }
 
     @Test
     public void shouldHaveCorrectDefaultApiVersionsResponse() {
-        Collection apiVersions = ApiVersionsResponse.defaultApiVersionsResponse().apiVersions();
-        assertEquals("API versions for all API keys must be maintained.", apiVersions.size(), ApiKeys.values().length);
+        Collection apiVersions = ApiVersionsResponse.DEFAULT_API_VERSIONS_RESPONSE.data().apiKeys();
+        assertEquals("API versions for all API keys must be maintained.", apiVersions.size(), ApiKeys.enabledApis().size());
 
-        for (ApiKeys key : ApiKeys.values()) {
-            ApiVersionsResponse.ApiVersion version = ApiVersionsResponse.defaultApiVersionsResponse().apiVersion(key.id);
+        for (ApiKeys key : ApiKeys.enabledApis()) {
+            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);
-        }
+
+        assertTrue(ApiVersionsResponse.DEFAULT_API_VERSIONS_RESPONSE.data().supportedFeatures().isEmpty());
+        assertTrue(ApiVersionsResponse.DEFAULT_API_VERSIONS_RESPONSE.data().finalizedFeatures().isEmpty());
+        assertEquals(ApiVersionsResponse.UNKNOWN_FINALIZED_FEATURES_EPOCH, ApiVersionsResponse.DEFAULT_API_VERSIONS_RESPONSE.data().finalizedFeaturesEpoch());
     }
 
     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/ByteBufferChannel.java b/clients/src/test/java/org/apache/kafka/common/requests/ByteBufferChannel.java
index 0446df97c322a..44bb358159eb0 100644
--- a/clients/src/test/java/org/apache/kafka/common/requests/ByteBufferChannel.java
+++ b/clients/src/test/java/org/apache/kafka/common/requests/ByteBufferChannel.java
@@ -16,7 +16,6 @@
  */
 package org.apache.kafka.common.requests;
 
-import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.nio.channels.GatheringByteChannel;
 
@@ -31,24 +30,25 @@ public ByteBufferChannel(long size) {
     }
 
     @Override
-    public long write(ByteBuffer[] srcs, int offset, int length) throws IOException {
+    public long write(ByteBuffer[] srcs, int offset, int length) {
         int position = buf.position();
         for (int i = 0; i < length; i++) {
             ByteBuffer src = srcs[i].duplicate();
-            if (i == 0)
-                src.position(offset);
+            if (i == 0) {
+                src.position(src.position() + offset);
+            }
             buf.put(src);
         }
         return buf.position() - position;
     }
 
     @Override
-    public long write(ByteBuffer[] srcs) throws IOException {
+    public long write(ByteBuffer[] srcs) {
         return write(srcs, 0, srcs.length);
     }
 
     @Override
-    public int write(ByteBuffer src) throws IOException {
+    public int write(ByteBuffer src) {
         int position = buf.position();
         buf.put(src);
         return buf.position() - position;
@@ -60,7 +60,7 @@ public boolean isOpen() {
     }
 
     @Override
-    public void close() throws IOException {
+    public void close() {
         buf.flip();
         closed = true;
     }
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/ByteBufferChannelTest.java b/clients/src/test/java/org/apache/kafka/common/requests/ByteBufferChannelTest.java
new file mode 100644
index 0000000000000..f8798b8a51cf8
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/ByteBufferChannelTest.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.requests;
+
+import org.apache.kafka.common.utils.Utils;
+import org.junit.Test;
+
+import java.nio.ByteBuffer;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+public class ByteBufferChannelTest {
+
+    @Test
+    public void testWriteBufferArrayWithNonZeroPosition() {
+        byte[] data = Utils.utf8("hello");
+        ByteBuffer buffer = ByteBuffer.allocate(32);
+        buffer.position(10);
+        buffer.put(data);
+
+        int limit = buffer.position();
+        buffer.position(10);
+        buffer.limit(limit);
+
+        ByteBufferChannel channel = new ByteBufferChannel(buffer.remaining());
+        ByteBuffer[] buffers = new ByteBuffer[] {buffer};
+        channel.write(buffers);
+        channel.close();
+        ByteBuffer channelBuffer = channel.buffer();
+        assertEquals(data.length, channelBuffer.remaining());
+        assertEquals("hello", Utils.utf8(channelBuffer));
+    }
+
+}
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/CreateAclsRequestTest.java b/clients/src/test/java/org/apache/kafka/common/requests/CreateAclsRequestTest.java
new file mode 100644
index 0000000000000..fa25f812ef88d
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/CreateAclsRequestTest.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.common.requests;
+
+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.errors.UnsupportedVersionException;
+import org.apache.kafka.common.message.CreateAclsRequestData;
+import org.apache.kafka.common.resource.PatternType;
+import org.apache.kafka.common.resource.ResourcePattern;
+import org.apache.kafka.common.resource.ResourceType;
+import org.junit.Test;
+
+import java.nio.ByteBuffer;
+import java.util.Arrays;
+import java.util.List;
+import java.util.stream.Collectors;
+
+import static org.junit.Assert.assertEquals;
+
+public class CreateAclsRequestTest {
+    private static final short V0 = 0;
+    private static final short V1 = 1;
+
+    private static final AclBinding LITERAL_ACL1 = new AclBinding(new ResourcePattern(ResourceType.TOPIC, "foo", PatternType.LITERAL),
+        new AccessControlEntry("User:ANONYMOUS", "127.0.0.1", AclOperation.READ, AclPermissionType.DENY));
+
+    private static final AclBinding LITERAL_ACL2 = new AclBinding(new ResourcePattern(ResourceType.GROUP, "group", PatternType.LITERAL),
+        new AccessControlEntry("User:*", "127.0.0.1", AclOperation.WRITE, AclPermissionType.ALLOW));
+
+    private static final AclBinding PREFIXED_ACL1 = new AclBinding(new ResourcePattern(ResourceType.GROUP, "prefix", PatternType.PREFIXED),
+        new AccessControlEntry("User:*", "127.0.0.1", AclOperation.CREATE, AclPermissionType.ALLOW));
+
+    private static final AclBinding UNKNOWN_ACL1 = new AclBinding(new ResourcePattern(ResourceType.UNKNOWN, "unknown", PatternType.LITERAL),
+        new AccessControlEntry("User:*", "127.0.0.1", AclOperation.CREATE, AclPermissionType.ALLOW));
+
+    @Test(expected = UnsupportedVersionException.class)
+    public void shouldThrowOnV0IfNotLiteral() {
+        new CreateAclsRequest(data(PREFIXED_ACL1), V0);
+    }
+
+    @Test(expected = IllegalArgumentException.class)
+    public void shouldThrowOnIfUnknown() {
+        new CreateAclsRequest(data(UNKNOWN_ACL1), V0);
+    }
+
+    @Test
+    public void shouldRoundTripV0() {
+        final CreateAclsRequest original = new CreateAclsRequest(data(LITERAL_ACL1, LITERAL_ACL2), V0);
+        final ByteBuffer buffer = original.serialize();
+
+        final CreateAclsRequest result = CreateAclsRequest.parse(buffer, V0);
+
+        assertRequestEquals(original, result);
+    }
+
+    @Test
+    public void shouldRoundTripV1() {
+        final CreateAclsRequest original = new CreateAclsRequest(data(LITERAL_ACL1, PREFIXED_ACL1), V1);
+        final ByteBuffer buffer = original.serialize();
+
+        final CreateAclsRequest result = CreateAclsRequest.parse(buffer, V1);
+
+        assertRequestEquals(original, result);
+    }
+
+    private static void assertRequestEquals(final CreateAclsRequest original, final CreateAclsRequest actual) {
+        assertEquals("Number of Acls wrong", original.aclCreations().size(), actual.aclCreations().size());
+
+        for (int idx = 0; idx != original.aclCreations().size(); ++idx) {
+            final AclBinding originalBinding = CreateAclsRequest.aclBinding(original.aclCreations().get(idx));
+            final AclBinding actualBinding = CreateAclsRequest.aclBinding(actual.aclCreations().get(idx));
+            assertEquals(originalBinding, actualBinding);
+        }
+    }
+
+    private static CreateAclsRequestData data(final AclBinding... acls) {
+        List aclCreations = Arrays.stream(acls)
+            .map(CreateAclsRequest::aclCreation)
+            .collect(Collectors.toList());
+        return new CreateAclsRequestData().setCreations(aclCreations);
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/DeleteAclsRequestTest.java b/clients/src/test/java/org/apache/kafka/common/requests/DeleteAclsRequestTest.java
new file mode 100644
index 0000000000000..20b2d68a0b4a0
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/DeleteAclsRequestTest.java
@@ -0,0 +1,117 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.acl.AccessControlEntryFilter;
+import org.apache.kafka.common.acl.AclBindingFilter;
+import org.apache.kafka.common.acl.AclOperation;
+import org.apache.kafka.common.acl.AclPermissionType;
+import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.message.DeleteAclsRequestData;
+import org.apache.kafka.common.resource.PatternType;
+import org.apache.kafka.common.resource.ResourcePatternFilter;
+import org.apache.kafka.common.resource.ResourceType;
+import org.junit.Test;
+
+import java.nio.ByteBuffer;
+import java.util.stream.Collectors;
+
+import static java.util.Arrays.asList;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+
+public class DeleteAclsRequestTest {
+    private static final short V0 = 0;
+    private static final short V1 = 1;
+
+    private static final AclBindingFilter LITERAL_FILTER = new AclBindingFilter(new ResourcePatternFilter(ResourceType.TOPIC, "foo", PatternType.LITERAL),
+        new AccessControlEntryFilter("User:ANONYMOUS", "127.0.0.1", AclOperation.READ, AclPermissionType.DENY));
+
+    private static final AclBindingFilter PREFIXED_FILTER = new AclBindingFilter(new ResourcePatternFilter(ResourceType.GROUP, "prefix", PatternType.PREFIXED),
+        new AccessControlEntryFilter("User:*", "127.0.0.1", AclOperation.CREATE, AclPermissionType.ALLOW));
+
+    private static final AclBindingFilter ANY_FILTER = new AclBindingFilter(new ResourcePatternFilter(ResourceType.GROUP, "bar", PatternType.ANY),
+        new AccessControlEntryFilter("User:*", "127.0.0.1", AclOperation.CREATE, AclPermissionType.ALLOW));
+
+    private static final AclBindingFilter UNKNOWN_FILTER = new AclBindingFilter(new ResourcePatternFilter(ResourceType.UNKNOWN, "prefix", PatternType.PREFIXED),
+        new AccessControlEntryFilter("User:*", "127.0.0.1", AclOperation.CREATE, AclPermissionType.ALLOW));
+
+    @Test
+    public void shouldThrowOnV0IfPrefixed() {
+        assertThrows(UnsupportedVersionException.class, () -> new DeleteAclsRequest.Builder(requestData(PREFIXED_FILTER)).build(V0));
+    }
+
+    @Test
+    public void shouldThrowOnUnknownElements() {
+        assertThrows(IllegalArgumentException.class, () -> new DeleteAclsRequest.Builder(requestData(UNKNOWN_FILTER)).build(V1));
+    }
+
+    @Test
+    public void shouldRoundTripLiteralV0() {
+        final DeleteAclsRequest original = new DeleteAclsRequest.Builder(requestData(LITERAL_FILTER)).build(V0);
+        final ByteBuffer buffer = original.serialize();
+
+        final DeleteAclsRequest result = DeleteAclsRequest.parse(buffer, V0);
+
+        assertRequestEquals(original, result);
+    }
+
+    @Test
+    public void shouldRoundTripAnyV0AsLiteral() {
+        final DeleteAclsRequest original = new DeleteAclsRequest.Builder(requestData(ANY_FILTER)).build(V0);
+        final DeleteAclsRequest expected = new DeleteAclsRequest.Builder(requestData(
+            new AclBindingFilter(new ResourcePatternFilter(
+                ANY_FILTER.patternFilter().resourceType(),
+                ANY_FILTER.patternFilter().name(),
+                PatternType.LITERAL),
+                ANY_FILTER.entryFilter()))
+        ).build(V0);
+
+        final DeleteAclsRequest result = DeleteAclsRequest.parse(original.serialize(), V0);
+
+        assertRequestEquals(expected, result);
+    }
+
+    @Test
+    public void shouldRoundTripV1() {
+        final DeleteAclsRequest original = new DeleteAclsRequest.Builder(
+                requestData(LITERAL_FILTER, PREFIXED_FILTER, ANY_FILTER)
+        ).build(V1);
+        final ByteBuffer buffer = original.serialize();
+
+        final DeleteAclsRequest result = DeleteAclsRequest.parse(buffer, V1);
+
+        assertRequestEquals(original, result);
+    }
+
+    private static void assertRequestEquals(final DeleteAclsRequest original, final DeleteAclsRequest actual) {
+        assertEquals("Number of filters wrong", original.filters().size(), actual.filters().size());
+
+        for (int idx = 0; idx != original.filters().size(); ++idx) {
+            final AclBindingFilter originalFilter = original.filters().get(idx);
+            final AclBindingFilter actualFilter = actual.filters().get(idx);
+            assertEquals(originalFilter, actualFilter);
+        }
+    }
+
+    private static DeleteAclsRequestData requestData(AclBindingFilter... acls) {
+        return new DeleteAclsRequestData().setFilters(asList(acls).stream()
+            .map(DeleteAclsRequest::deleteAclsFilter)
+            .collect(Collectors.toList()));
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/DeleteAclsResponseTest.java b/clients/src/test/java/org/apache/kafka/common/requests/DeleteAclsResponseTest.java
new file mode 100644
index 0000000000000..a4a06eca17628
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/DeleteAclsResponseTest.java
@@ -0,0 +1,130 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.kafka.common.requests;
+
+import org.apache.kafka.common.acl.AclOperation;
+import org.apache.kafka.common.acl.AclPermissionType;
+import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.message.DeleteAclsResponseData;
+import org.apache.kafka.common.message.DeleteAclsResponseData.DeleteAclsFilterResult;
+import org.apache.kafka.common.message.DeleteAclsResponseData.DeleteAclsMatchingAcl;
+import org.apache.kafka.common.resource.PatternType;
+import org.apache.kafka.common.resource.ResourceType;
+import org.junit.Test;
+
+import java.nio.ByteBuffer;
+
+import static java.util.Arrays.asList;
+import static java.util.Collections.singletonList;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+
+public class DeleteAclsResponseTest {
+    private static final short V0 = 0;
+    private static final short V1 = 1;
+
+    private static final DeleteAclsMatchingAcl LITERAL_ACL1 = new DeleteAclsMatchingAcl()
+        .setResourceType(ResourceType.TOPIC.code())
+        .setResourceName("foo")
+        .setPatternType(PatternType.LITERAL.code())
+        .setPrincipal("User:ANONYMOUS")
+        .setHost("127.0.0.1")
+        .setOperation(AclOperation.READ.code())
+        .setPermissionType(AclPermissionType.DENY.code());
+
+    private static final DeleteAclsMatchingAcl LITERAL_ACL2 = new DeleteAclsMatchingAcl()
+            .setResourceType(ResourceType.GROUP.code())
+            .setResourceName("group")
+            .setPatternType(PatternType.LITERAL.code())
+            .setPrincipal("User:*")
+            .setHost("127.0.0.1")
+            .setOperation(AclOperation.WRITE.code())
+            .setPermissionType(AclPermissionType.ALLOW.code());
+
+    private static final DeleteAclsMatchingAcl PREFIXED_ACL1 = new DeleteAclsMatchingAcl()
+            .setResourceType(ResourceType.GROUP.code())
+            .setResourceName("prefix")
+            .setPatternType(PatternType.PREFIXED.code())
+            .setPrincipal("User:*")
+            .setHost("127.0.0.1")
+            .setOperation(AclOperation.CREATE.code())
+            .setPermissionType(AclPermissionType.ALLOW.code());
+
+    private static final DeleteAclsMatchingAcl UNKNOWN_ACL = new DeleteAclsMatchingAcl()
+            .setResourceType(ResourceType.UNKNOWN.code())
+            .setResourceName("group")
+            .setPatternType(PatternType.LITERAL.code())
+            .setPrincipal("User:*")
+            .setHost("127.0.0.1")
+            .setOperation(AclOperation.WRITE.code())
+            .setPermissionType(AclPermissionType.ALLOW.code());
+
+    private static final DeleteAclsFilterResult LITERAL_RESPONSE = new DeleteAclsFilterResult().setMatchingAcls(asList(
+        LITERAL_ACL1, LITERAL_ACL2));
+
+    private static final DeleteAclsFilterResult PREFIXED_RESPONSE = new DeleteAclsFilterResult().setMatchingAcls(asList(
+        LITERAL_ACL1, PREFIXED_ACL1));
+
+    private static final DeleteAclsFilterResult UNKNOWN_RESPONSE = new DeleteAclsFilterResult().setMatchingAcls(asList(
+            UNKNOWN_ACL));
+
+    @Test
+    public void shouldThrowOnV0IfNotLiteral() {
+        assertThrows(UnsupportedVersionException.class, () -> new DeleteAclsResponse(
+            new DeleteAclsResponseData()
+                .setThrottleTimeMs(10)
+                .setFilterResults(singletonList(PREFIXED_RESPONSE)),
+            V0));
+    }
+
+    @Test
+    public void shouldThrowOnIfUnknown() {
+        assertThrows(IllegalArgumentException.class, () -> new DeleteAclsResponse(
+            new DeleteAclsResponseData()
+                .setThrottleTimeMs(10)
+                .setFilterResults(singletonList(UNKNOWN_RESPONSE)),
+            V1));
+    }
+
+    @Test
+    public void shouldRoundTripV0() {
+        final DeleteAclsResponse original = new DeleteAclsResponse(
+            new DeleteAclsResponseData()
+                .setThrottleTimeMs(10)
+                .setFilterResults(singletonList(LITERAL_RESPONSE)),
+            V0);
+        final ByteBuffer buffer = original.serialize(V0);
+
+        final DeleteAclsResponse result = DeleteAclsResponse.parse(buffer, V0);
+        assertEquals(original.filterResults(), result.filterResults());
+    }
+
+    @Test
+    public void shouldRoundTripV1() {
+        final DeleteAclsResponse original = new DeleteAclsResponse(
+            new DeleteAclsResponseData()
+                .setThrottleTimeMs(10)
+                .setFilterResults(asList(LITERAL_RESPONSE, PREFIXED_RESPONSE)),
+            V1);
+        final ByteBuffer buffer = original.serialize(V1);
+
+        final DeleteAclsResponse result = DeleteAclsResponse.parse(buffer, V1);
+        assertEquals(original.filterResults(), result.filterResults());
+    }
+
+}
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/test/java/org/apache/kafka/common/requests/DescribeAclsRequestTest.java b/clients/src/test/java/org/apache/kafka/common/requests/DescribeAclsRequestTest.java
new file mode 100644
index 0000000000000..c35a19b7e2b99
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/DescribeAclsRequestTest.java
@@ -0,0 +1,107 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.acl.AccessControlEntryFilter;
+import org.apache.kafka.common.acl.AclBindingFilter;
+import org.apache.kafka.common.acl.AclOperation;
+import org.apache.kafka.common.acl.AclPermissionType;
+import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.resource.PatternType;
+import org.apache.kafka.common.resource.ResourcePatternFilter;
+import org.apache.kafka.common.resource.ResourceType;
+import org.junit.Test;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+
+public class DescribeAclsRequestTest {
+    private static final short V0 = 0;
+    private static final short V1 = 1;
+
+    private static final AclBindingFilter LITERAL_FILTER = new AclBindingFilter(new ResourcePatternFilter(ResourceType.TOPIC, "foo", PatternType.LITERAL),
+        new AccessControlEntryFilter("User:ANONYMOUS", "127.0.0.1", AclOperation.READ, AclPermissionType.DENY));
+
+    private static final AclBindingFilter PREFIXED_FILTER = new AclBindingFilter(new ResourcePatternFilter(ResourceType.GROUP, "prefix", PatternType.PREFIXED),
+        new AccessControlEntryFilter("User:*", "127.0.0.1", AclOperation.CREATE, AclPermissionType.ALLOW));
+
+    private static final AclBindingFilter ANY_FILTER = new AclBindingFilter(new ResourcePatternFilter(ResourceType.GROUP, "bar", PatternType.ANY),
+        new AccessControlEntryFilter("User:*", "127.0.0.1", AclOperation.CREATE, AclPermissionType.ALLOW));
+
+    private static final AclBindingFilter UNKNOWN_FILTER = new AclBindingFilter(new ResourcePatternFilter(ResourceType.UNKNOWN, "foo", PatternType.LITERAL),
+        new AccessControlEntryFilter("User:ANONYMOUS", "127.0.0.1", AclOperation.READ, AclPermissionType.DENY));
+
+    @Test
+    public void shouldThrowOnV0IfPrefixed() {
+        assertThrows(UnsupportedVersionException.class, () -> new DescribeAclsRequest.Builder(PREFIXED_FILTER).build(V0));
+    }
+
+    @Test
+    public void shouldThrowIfUnknown() {
+        assertThrows(IllegalArgumentException.class, () -> new DescribeAclsRequest.Builder(UNKNOWN_FILTER).build(V0));
+    }
+
+    @Test
+    public void shouldRoundTripLiteralV0() {
+        final DescribeAclsRequest original = new DescribeAclsRequest.Builder(LITERAL_FILTER).build(V0);
+        final DescribeAclsRequest result = DescribeAclsRequest.parse(original.serialize(), V0);
+
+        assertRequestEquals(original, result);
+    }
+
+    @Test
+    public void shouldRoundTripAnyV0AsLiteral() {
+        final DescribeAclsRequest original = new DescribeAclsRequest.Builder(ANY_FILTER).build(V0);
+        final DescribeAclsRequest expected = new DescribeAclsRequest.Builder(
+            new AclBindingFilter(new ResourcePatternFilter(
+                ANY_FILTER.patternFilter().resourceType(),
+                ANY_FILTER.patternFilter().name(),
+                PatternType.LITERAL),
+                ANY_FILTER.entryFilter())).build(V0);
+
+        final DescribeAclsRequest result = DescribeAclsRequest.parse(original.serialize(), V0);
+        assertRequestEquals(expected, result);
+    }
+
+    @Test
+    public void shouldRoundTripLiteralV1() {
+        final DescribeAclsRequest original = new DescribeAclsRequest.Builder(LITERAL_FILTER).build(V1);
+        final DescribeAclsRequest result = DescribeAclsRequest.parse(original.serialize(), V1);
+        assertRequestEquals(original, result);
+    }
+
+    @Test
+    public void shouldRoundTripPrefixedV1() {
+        final DescribeAclsRequest original = new DescribeAclsRequest.Builder(PREFIXED_FILTER).build(V1);
+        final DescribeAclsRequest result = DescribeAclsRequest.parse(original.serialize(), V1);
+        assertRequestEquals(original, result);
+    }
+
+    @Test
+    public void shouldRoundTripAnyV1() {
+        final DescribeAclsRequest original = new DescribeAclsRequest.Builder(ANY_FILTER).build(V1);
+        final DescribeAclsRequest result = DescribeAclsRequest.parse(original.serialize(), V1);
+        assertRequestEquals(original, result);
+    }
+
+    private static void assertRequestEquals(final DescribeAclsRequest original, final DescribeAclsRequest actual) {
+        final AclBindingFilter originalFilter = original.filter();
+        final AclBindingFilter acttualFilter = actual.filter();
+        assertEquals(originalFilter, acttualFilter);
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/DescribeAclsResponseTest.java b/clients/src/test/java/org/apache/kafka/common/requests/DescribeAclsResponseTest.java
new file mode 100644
index 0000000000000..baa64ee696573
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/DescribeAclsResponseTest.java
@@ -0,0 +1,161 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.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.errors.UnsupportedVersionException;
+import org.apache.kafka.common.message.DescribeAclsResponseData;
+import org.apache.kafka.common.message.DescribeAclsResponseData.AclDescription;
+import org.apache.kafka.common.message.DescribeAclsResponseData.DescribeAclsResource;
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.resource.PatternType;
+import org.apache.kafka.common.resource.ResourcePattern;
+import org.apache.kafka.common.resource.ResourceType;
+import org.junit.Test;
+
+import java.nio.ByteBuffer;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+
+import static org.junit.Assert.assertEquals;
+
+public class DescribeAclsResponseTest {
+    private static final short V0 = 0;
+    private static final short V1 = 1;
+
+    private static final AclDescription ALLOW_CREATE_ACL = buildAclDescription(
+            "127.0.0.1",
+            "User:ANONYMOUS",
+            AclOperation.CREATE,
+            AclPermissionType.ALLOW);
+
+    private static final AclDescription DENY_READ_ACL = buildAclDescription(
+            "127.0.0.1",
+            "User:ANONYMOUS",
+            AclOperation.READ,
+            AclPermissionType.DENY);
+
+    private static final DescribeAclsResource UNKNOWN_ACL = buildResource(
+            "foo",
+            ResourceType.UNKNOWN,
+            PatternType.LITERAL,
+            Collections.singletonList(DENY_READ_ACL));
+
+    private static final DescribeAclsResource PREFIXED_ACL1 = buildResource(
+            "prefix",
+            ResourceType.GROUP,
+            PatternType.PREFIXED,
+            Collections.singletonList(ALLOW_CREATE_ACL));
+
+    private static final DescribeAclsResource LITERAL_ACL1 = buildResource(
+            "foo",
+            ResourceType.TOPIC,
+            PatternType.LITERAL,
+            Collections.singletonList(ALLOW_CREATE_ACL));
+
+    private static final DescribeAclsResource LITERAL_ACL2 = buildResource(
+            "group",
+            ResourceType.GROUP,
+            PatternType.LITERAL,
+            Collections.singletonList(DENY_READ_ACL));
+
+    @Test(expected = UnsupportedVersionException.class)
+    public void shouldThrowOnV0IfNotLiteral() {
+        buildResponse(10, Errors.NONE, Collections.singletonList(PREFIXED_ACL1)).serialize(V0);
+    }
+
+    @Test(expected = IllegalArgumentException.class)
+    public void shouldThrowIfUnknown() {
+        buildResponse(10, Errors.NONE, Collections.singletonList(UNKNOWN_ACL)).serialize(V0);
+    }
+
+    @Test
+    public void shouldRoundTripV0() {
+        List resources = Arrays.asList(LITERAL_ACL1, LITERAL_ACL2);
+        final DescribeAclsResponse original = buildResponse(10, Errors.NONE, resources);
+        final ByteBuffer buffer = original.serialize(V0);
+
+        final DescribeAclsResponse result = DescribeAclsResponse.parse(buffer, V0);
+        assertResponseEquals(original, result);
+
+        final DescribeAclsResponse result2 = buildResponse(10, Errors.NONE, DescribeAclsResponse.aclsResources(
+            DescribeAclsResponse.aclBindings(resources)));
+        assertResponseEquals(original, result2);
+    }
+
+    @Test
+    public void shouldRoundTripV1() {
+        List resources = Arrays.asList(LITERAL_ACL1, PREFIXED_ACL1);
+        final DescribeAclsResponse original = buildResponse(100, Errors.NONE, resources);
+        final ByteBuffer buffer = original.serialize(V1);
+
+        final DescribeAclsResponse result = DescribeAclsResponse.parse(buffer, V1);
+        assertResponseEquals(original, result);
+
+        final DescribeAclsResponse result2 = buildResponse(100, Errors.NONE, DescribeAclsResponse.aclsResources(
+            DescribeAclsResponse.aclBindings(resources)));
+        assertResponseEquals(original, result2);
+    }
+
+    @Test
+    public void testAclBindings() {
+        final AclBinding original = new AclBinding(new ResourcePattern(ResourceType.TOPIC, "foo", PatternType.LITERAL),
+                new AccessControlEntry("User:ANONYMOUS", "127.0.0.1", AclOperation.CREATE, AclPermissionType.ALLOW));
+
+        final List result = DescribeAclsResponse.aclBindings(Collections.singletonList(LITERAL_ACL1));
+        assertEquals(1, result.size());
+        assertEquals(original, result.get(0));
+    }
+
+    private static void assertResponseEquals(final DescribeAclsResponse original, final DescribeAclsResponse actual) {
+        final Set originalBindings = new HashSet<>(original.acls());
+        final Set actualBindings = new HashSet<>(actual.acls());
+
+        assertEquals(originalBindings, actualBindings);
+    }
+
+    private static DescribeAclsResponse buildResponse(int throttleTimeMs, Errors error, List resources) {
+        return new DescribeAclsResponse(new DescribeAclsResponseData()
+            .setThrottleTimeMs(throttleTimeMs)
+            .setErrorCode(error.code())
+            .setErrorMessage(error.message())
+            .setResources(resources));
+    }
+
+    private static DescribeAclsResource buildResource(String name, ResourceType type, PatternType patternType, List acls) {
+        return new DescribeAclsResource()
+            .setResourceName(name)
+            .setResourceType(type.code())
+            .setPatternType(patternType.code())
+            .setAcls(acls);
+    }
+
+    private static AclDescription buildAclDescription(String host, String principal, AclOperation operation, AclPermissionType permission) {
+        return new AclDescription()
+            .setHost(host)
+            .setPrincipal(principal)
+            .setOperation(operation.code())
+            .setPermissionType(permission.code());
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/EndTxnRequestTest.java b/clients/src/test/java/org/apache/kafka/common/requests/EndTxnRequestTest.java
new file mode 100644
index 0000000000000..3a6c82eb17bbc
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/EndTxnRequestTest.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.requests;
+
+import org.apache.kafka.common.message.EndTxnRequestData;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.Errors;
+import org.junit.Test;
+
+import java.util.Collections;
+
+import static org.junit.Assert.assertEquals;
+
+public class EndTxnRequestTest {
+
+    @Test
+    public void testConstructor() {
+        short producerEpoch = 0;
+        int producerId = 1;
+        String transactionId = "txn_id";
+        int throttleTimeMs = 10;
+
+        EndTxnRequest.Builder builder = new EndTxnRequest.Builder(
+            new EndTxnRequestData()
+                .setCommitted(true)
+                .setProducerEpoch(producerEpoch)
+                .setProducerId(producerId)
+                .setTransactionalId(transactionId));
+
+        for (short version = 0; version <= ApiKeys.END_TXN.latestVersion(); version++) {
+            EndTxnRequest request = builder.build(version);
+
+            EndTxnResponse response = request.getErrorResponse(throttleTimeMs, Errors.NOT_COORDINATOR.exception());
+
+            assertEquals(Collections.singletonMap(Errors.NOT_COORDINATOR, 1), response.errorCounts());
+
+            assertEquals(TransactionResult.COMMIT, request.result());
+
+            assertEquals(throttleTimeMs, response.throttleTimeMs());
+        }
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/EndTxnResponseTest.java b/clients/src/test/java/org/apache/kafka/common/requests/EndTxnResponseTest.java
new file mode 100644
index 0000000000000..ab9b9f99ad605
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/EndTxnResponseTest.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 org.apache.kafka.common.requests;
+
+import org.apache.kafka.common.message.EndTxnResponseData;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.Errors;
+import org.junit.Test;
+
+import java.util.Collections;
+import java.util.Map;
+
+import static org.junit.Assert.assertEquals;
+
+public class EndTxnResponseTest {
+
+    @Test
+    public void testConstructor() {
+        int throttleTimeMs = 10;
+
+        EndTxnResponseData data = new EndTxnResponseData()
+            .setErrorCode(Errors.NOT_COORDINATOR.code())
+            .setThrottleTimeMs(throttleTimeMs);
+
+        Map expectedErrorCounts = Collections.singletonMap(Errors.NOT_COORDINATOR, 1);
+
+        for (short version = 0; version <= ApiKeys.END_TXN.latestVersion(); version++) {
+            EndTxnResponse response = new EndTxnResponse(data);
+            assertEquals(expectedErrorCounts, response.errorCounts());
+            assertEquals(throttleTimeMs, response.throttleTimeMs());
+            assertEquals(version >= 1, response.shouldClientThrottle(version));
+
+            response = EndTxnResponse.parse(response.serialize(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/EnvelopeRequestTest.java b/clients/src/test/java/org/apache/kafka/common/requests/EnvelopeRequestTest.java
new file mode 100644
index 0000000000000..ef92eb16215d6
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/EnvelopeRequestTest.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.common.requests;
+
+import org.apache.kafka.common.message.EnvelopeRequestData;
+import org.apache.kafka.common.network.Send;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.security.auth.KafkaPrincipal;
+import org.apache.kafka.common.security.authenticator.DefaultKafkaPrincipalBuilder;
+import org.apache.kafka.test.TestUtils;
+import org.junit.Test;
+
+import java.io.IOException;
+import java.net.InetAddress;
+import java.nio.ByteBuffer;
+
+import static org.junit.Assert.assertEquals;
+
+public class EnvelopeRequestTest {
+
+    @Test
+    public void testGetPrincipal() {
+        KafkaPrincipal kafkaPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "principal", true);
+        DefaultKafkaPrincipalBuilder kafkaPrincipalBuilder = new DefaultKafkaPrincipalBuilder(null, null);
+
+        EnvelopeRequest.Builder requestBuilder = new EnvelopeRequest.Builder(ByteBuffer.allocate(0),
+            kafkaPrincipalBuilder.serialize(kafkaPrincipal), "client-address".getBytes());
+        EnvelopeRequest request = requestBuilder.build(EnvelopeRequestData.HIGHEST_SUPPORTED_VERSION);
+        assertEquals(kafkaPrincipal, kafkaPrincipalBuilder.deserialize(request.requestPrincipal()));
+    }
+
+    @Test
+    public void testToSend() throws IOException {
+        for (short version = ApiKeys.ENVELOPE.oldestVersion(); version <= ApiKeys.ENVELOPE.latestVersion(); version++) {
+            ByteBuffer requestData = ByteBuffer.wrap("foobar".getBytes());
+            RequestHeader header = new RequestHeader(ApiKeys.ENVELOPE, version, "clientId", 15);
+            EnvelopeRequest request = new EnvelopeRequest.Builder(
+                requestData,
+                "principal".getBytes(),
+                InetAddress.getLocalHost().getAddress()
+            ).build(version);
+
+            Send send = request.toSend(header);
+            ByteBuffer buffer = TestUtils.toBuffer(send);
+            assertEquals(send.size() - 4, buffer.getInt());
+            assertEquals(header, RequestHeader.parse(buffer));
+
+            EnvelopeRequestData parsedRequestData = new EnvelopeRequestData();
+            parsedRequestData.read(new ByteBufferAccessor(buffer), version);
+            assertEquals(request.data(), parsedRequestData);
+        }
+    }
+
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/EnvelopeResponseTest.java b/clients/src/test/java/org/apache/kafka/common/requests/EnvelopeResponseTest.java
new file mode 100644
index 0000000000000..0a384cdd1661a
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/EnvelopeResponseTest.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.common.requests;
+
+import org.apache.kafka.common.message.EnvelopeResponseData;
+import org.apache.kafka.common.network.Send;
+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.test.TestUtils;
+import org.junit.jupiter.api.Test;
+
+import java.nio.ByteBuffer;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+class EnvelopeResponseTest {
+
+    @Test
+    public void testToSend() {
+        for (short version = ApiKeys.ENVELOPE.oldestVersion(); version <= ApiKeys.ENVELOPE.latestVersion(); version++) {
+            ByteBuffer responseData = ByteBuffer.wrap("foobar".getBytes());
+            EnvelopeResponse response = new EnvelopeResponse(responseData, Errors.NONE);
+            short headerVersion = ApiKeys.ENVELOPE.responseHeaderVersion(version);
+            ResponseHeader header = new ResponseHeader(15, headerVersion);
+
+            Send send = response.toSend(header, version);
+            ByteBuffer buffer = TestUtils.toBuffer(send);
+            assertEquals(send.size() - 4, buffer.getInt());
+            assertEquals(header, ResponseHeader.parse(buffer, headerVersion));
+
+            EnvelopeResponseData parsedResponseData = new EnvelopeResponseData();
+            parsedResponseData.read(new ByteBufferAccessor(buffer), version);
+            assertEquals(response.data(), parsedResponseData);
+        }
+    }
+
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/HeartbeatRequestTest.java b/clients/src/test/java/org/apache/kafka/common/requests/HeartbeatRequestTest.java
new file mode 100644
index 0000000000000..2532213041d93
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/HeartbeatRequestTest.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.HeartbeatRequestData;
+import org.junit.Test;
+
+public class HeartbeatRequestTest {
+
+    @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..125a328839d25
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/JoinGroupRequestTest.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.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")
+                .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..939514ed7339b
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/LeaderAndIsrRequestTest.java
@@ -0,0 +1,171 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.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, 0,
+                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,
+                    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, 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 = request.serialize();
+            LeaderAndIsrRequest deserializedRequest = new LeaderAndIsrRequest(new LeaderAndIsrRequestData(
+                new ByteBufferAccessor(byteBuffer), version), version);
+
+            // Adding/removing replicas is only supported from version 3, so the deserialized request won't have
+            // them for earlier versions.
+            if (version < 3) {
+                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,
+            partitionStates, Collections.emptySet());
+
+        LeaderAndIsrRequest v2 = builder.build((short) 2);
+        LeaderAndIsrRequest v1 = builder.build((short) 1);
+        assertTrue("Expected v2 < v1: v2=" + v2.sizeInBytes() + ", v1=" + v1.sizeInBytes(), v2.sizeInBytes() < v1.sizeInBytes());
+    }
+
+    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 3eda77827437c..fbd7d48556b03 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,52 +16,99 @@
  */
 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;
 
 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, partitionStates, Collections.emptySet()).build();
+                15, 20, 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());
+        assertEquals(Collections.singletonMap(Errors.CLUSTER_AUTHORIZATION_FAILED, 3), 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);
-        assertEquals(Collections.singletonMap(Errors.UNKNOWN_SERVER_ERROR, 2), response.errorCounts());
+        List partitions = createPartitions("foo",
+            asList(Errors.NONE, Errors.NOT_LEADER_OR_FOLLOWER));
+        LeaderAndIsrResponse response = new LeaderAndIsrResponse(new LeaderAndIsrResponseData()
+            .setErrorCode(Errors.UNKNOWN_SERVER_ERROR.code())
+            .setPartitionErrors(partitions));
+        assertEquals(Collections.singletonMap(Errors.UNKNOWN_SERVER_ERROR, 3), 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());
+        assertEquals(2, errorCounts.get(Errors.NONE).intValue());
         assertEquals(1, errorCounts.get(Errors.CLUSTER_AUTHORIZATION_FAILED).intValue());
     }
 
+    @Test
+    public void testToString() {
+        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(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..976f17fe11e29
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/LeaveGroupResponseTest.java
@@ -0,0 +1,168 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.apache.kafka.common.protocol.MessageUtil;
+import org.junit.Before;
+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 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 testConstructorWithMemberResponses() {
+        Map expectedErrorCounts = new HashMap<>();
+        expectedErrorCounts.put(Errors.NONE, 1); // top level
+        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 testEqualityWithSerialization() {
+        LeaveGroupResponseData responseData = new LeaveGroupResponseData()
+                .setErrorCode(Errors.NONE.code())
+                .setThrottleTimeMs(throttleTimeMs);
+        for (short version = 0; version <= ApiKeys.LEAVE_GROUP.latestVersion(); version++) {
+            LeaveGroupResponse primaryResponse = LeaveGroupResponse.parse(
+                MessageUtil.toByteBuffer(responseData, version), version);
+            LeaveGroupResponse secondaryResponse = LeaveGroupResponse.parse(
+                MessageUtil.toByteBuffer(responseData, version), version);
+
+            assertEquals(primaryResponse, primaryResponse);
+            assertEquals(primaryResponse, secondaryResponse);
+            assertEquals(primaryResponse.hashCode(), secondaryResponse.hashCode());
+        }
+    }
+
+    @Test
+    public void testParse() {
+        Map expectedErrorCounts = Collections.singletonMap(Errors.NOT_COORDINATOR, 1);
+
+        LeaveGroupResponseData data = new LeaveGroupResponseData()
+            .setErrorCode(Errors.NOT_COORDINATOR.code())
+            .setThrottleTimeMs(throttleTimeMs);
+
+        for (short version = 0; version <= ApiKeys.LEAVE_GROUP.latestVersion(); version++) {
+            ByteBuffer buffer = MessageUtil.toByteBuffer(data, version);
+            LeaveGroupResponse leaveGroupResponse = LeaveGroupResponse.parse(buffer, 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 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/ListOffsetRequestTest.java b/clients/src/test/java/org/apache/kafka/common/requests/ListOffsetRequestTest.java
new file mode 100644
index 0000000000000..80af40e653332
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/ListOffsetRequestTest.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.common.requests;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import org.apache.kafka.common.IsolationLevel;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.message.ListOffsetRequestData;
+import org.apache.kafka.common.message.ListOffsetRequestData.ListOffsetPartition;
+import org.apache.kafka.common.message.ListOffsetRequestData.ListOffsetTopic;
+import org.apache.kafka.common.message.ListOffsetResponseData;
+import org.apache.kafka.common.message.ListOffsetResponseData.ListOffsetPartitionResponse;
+import org.apache.kafka.common.message.ListOffsetResponseData.ListOffsetTopicResponse;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.protocol.MessageUtil;
+import org.junit.Test;
+
+public class ListOffsetRequestTest {
+
+    @Test
+    public void testDuplicatePartitions() {
+        List topics = Collections.singletonList(
+                new ListOffsetTopic()
+                    .setName("topic")
+                    .setPartitions(Arrays.asList(
+                            new ListOffsetPartition()
+                                .setPartitionIndex(0),
+                            new ListOffsetPartition()
+                                .setPartitionIndex(0))));
+        ListOffsetRequestData data = new ListOffsetRequestData()
+                .setTopics(topics)
+                .setReplicaId(-1);
+        ListOffsetRequest request = ListOffsetRequest.parse(MessageUtil.toByteBuffer(data, (short) 0), (short) 0);
+        assertEquals(Collections.singleton(new TopicPartition("topic", 0)), request.duplicatePartitions());
+    }
+
+    @Test
+    public void testGetErrorResponse() {
+        for (short version = 1; version <= ApiKeys.LIST_OFFSETS.latestVersion(); version++) {
+            List topics = Arrays.asList(
+                    new ListOffsetTopic()
+                        .setName("topic")
+                        .setPartitions(Collections.singletonList(
+                                new ListOffsetPartition()
+                                    .setPartitionIndex(0))));
+            ListOffsetRequest request = ListOffsetRequest.Builder
+                    .forConsumer(true, IsolationLevel.READ_COMMITTED)
+                    .setTargetTimes(topics)
+                    .build(version);
+            ListOffsetResponse response = (ListOffsetResponse) request.getErrorResponse(0, Errors.NOT_LEADER_OR_FOLLOWER.exception());
+    
+            List v = Collections.singletonList(
+                    new ListOffsetTopicResponse()
+                        .setName("topic")
+                        .setPartitions(Collections.singletonList(
+                                new ListOffsetPartitionResponse()
+                                    .setErrorCode(Errors.NOT_LEADER_OR_FOLLOWER.code())
+                                    .setLeaderEpoch(ListOffsetResponse.UNKNOWN_EPOCH)
+                                    .setOffset(ListOffsetResponse.UNKNOWN_OFFSET)
+                                    .setPartitionIndex(0)
+                                    .setTimestamp(ListOffsetResponse.UNKNOWN_TIMESTAMP))));
+            ListOffsetResponseData data = new ListOffsetResponseData()
+                    .setThrottleTimeMs(0)
+                    .setTopics(v);
+            ListOffsetResponse expectedResponse = new ListOffsetResponse(data);
+            assertEquals(expectedResponse.data().topics(), response.data().topics());
+            assertEquals(expectedResponse.throttleTimeMs(), response.throttleTimeMs());
+        }
+    }
+
+    @Test
+    public void testGetErrorResponseV0() {
+        List topics = Arrays.asList(
+                new ListOffsetTopic()
+                    .setName("topic")
+                    .setPartitions(Collections.singletonList(
+                            new ListOffsetPartition()
+                                .setPartitionIndex(0))));
+        ListOffsetRequest request = ListOffsetRequest.Builder
+                .forConsumer(true, IsolationLevel.READ_UNCOMMITTED)
+                .setTargetTimes(topics)
+                .build((short) 0);
+        ListOffsetResponse response = (ListOffsetResponse) request.getErrorResponse(0, Errors.NOT_LEADER_OR_FOLLOWER.exception());
+
+        List v = Collections.singletonList(
+                new ListOffsetTopicResponse()
+                    .setName("topic")
+                    .setPartitions(Collections.singletonList(
+                            new ListOffsetPartitionResponse()
+                                .setErrorCode(Errors.NOT_LEADER_OR_FOLLOWER.code())
+                                .setOldStyleOffsets(Collections.emptyList())
+                                .setPartitionIndex(0))));
+        ListOffsetResponseData data = new ListOffsetResponseData()
+                .setThrottleTimeMs(0)
+                .setTopics(v);
+        ListOffsetResponse expectedResponse = new ListOffsetResponse(data);
+        assertEquals(expectedResponse.data().topics(), response.data().topics());
+        assertEquals(expectedResponse.throttleTimeMs(), response.throttleTimeMs());
+    }
+
+    @Test
+    public void testToListOffsetTopics() {
+        ListOffsetPartition lop0 = new ListOffsetPartition()
+                .setPartitionIndex(0)
+                .setCurrentLeaderEpoch(1)
+                .setMaxNumOffsets(2)
+                .setTimestamp(123L);
+        ListOffsetPartition lop1 = new ListOffsetPartition()
+                .setPartitionIndex(1)
+                .setCurrentLeaderEpoch(3)
+                .setMaxNumOffsets(4)
+                .setTimestamp(567L);
+        Map timestampsToSearch = new HashMap<>();
+        timestampsToSearch.put(new TopicPartition("topic", 0), lop0);
+        timestampsToSearch.put(new TopicPartition("topic", 1), lop1);
+        List listOffsetTopics = ListOffsetRequest.toListOffsetTopics(timestampsToSearch);
+        assertEquals(1, listOffsetTopics.size());
+        ListOffsetTopic topic = listOffsetTopics.get(0);
+        assertEquals("topic", topic.name());
+        assertEquals(2, topic.partitions().size());
+        assertTrue(topic.partitions().contains(lop0));
+        assertTrue(topic.partitions().contains(lop1));
+    }
+
+}
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
new file mode 100644
index 0000000000000..31e22625ce8d1
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/MetadataRequestTest.java
@@ -0,0 +1,68 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.MetadataRequestData;
+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.assertFalse;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+
+public class MetadataRequestTest {
+
+    @Test
+    public void testEmptyMeansAllTopicsV0() {
+        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 < 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..b379f2157262f
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/OffsetCommitRequestTest.java
@@ -0,0 +1,142 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.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 memberId = "consumerId";
+    protected static String groupInstanceId = "groupInstanceId";
+    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());
+            assertEquals(throttleTimeMs, 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(memberId)
+                .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..17a53458a0676
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/OffsetCommitResponseTest.java
@@ -0,0 +1,103 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.apache.kafka.common.protocol.MessageUtil;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.nio.ByteBuffer;
+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 testParse() {
+        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++) {
+            ByteBuffer buffer = MessageUtil.toByteBuffer(data, version);
+            OffsetCommitResponse response = OffsetCommitResponse.parse(buffer, 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..79fd9aec05f3d
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/OffsetFetchRequestTest.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.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,
+            false,
+            partitions,
+            false);
+    }
+
+    @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("Incorrect error count for version " + version,
+                    Collections.singletonMap(Errors.NONE, version <= (short) 1 ? 3 : 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 testConstructorFailForUnsupportedRequireStable() {
+        for (short version = 0; version <= ApiKeys.OFFSET_FETCH.latestVersion(); version++) {
+            // The builder needs to be initialized every cycle as the internal data `requireStable` flag is flipped.
+            builder = new OffsetFetchRequest.Builder(groupId, true, null, false);
+            final short finalVersion = version;
+            if (version < 2) {
+                assertThrows(UnsupportedVersionException.class, () -> builder.build(finalVersion));
+            } else {
+                OffsetFetchRequest request = builder.build(finalVersion);
+                assertEquals(groupId, request.groupId());
+                assertNull(request.partitions());
+                assertTrue(request.isAllPartitions());
+                if (version < 7) {
+                    assertFalse(request.requireStable());
+                } else {
+                    assertTrue(request.requireStable());
+                }
+            }
+        }
+    }
+
+    @Test
+    public void testBuildThrowForUnsupportedRequireStable() {
+        for (short version = 0; version <= ApiKeys.OFFSET_FETCH.latestVersion(); version++) {
+            builder = new OffsetFetchRequest.Builder(groupId, true, null, true);
+            if (version < 7) {
+                final short finalVersion = version;
+                assertThrows(UnsupportedVersionException.class, () -> builder.build(finalVersion));
+            } else {
+                OffsetFetchRequest request = builder.build(version);
+                assertTrue(request.requireStable());
+            }
+        }
+    }
+}
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..06e611f49a6da
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/OffsetFetchResponseTest.java
@@ -0,0 +1,231 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.apache.kafka.common.utils.Utils;
+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(3, response.errorCounts().size());
+        assertEquals(Utils.mkMap(Utils.mkEntry(Errors.NOT_COORDINATOR, 1),
+                Utils.mkEntry(Errors.TOPIC_AUTHORIZATION_FAILED, 1),
+                Utils.mkEntry(Errors.UNKNOWN_TOPIC_OR_PARTITION, 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 = OffsetFetchResponse.parse(latestResponse.serialize(version), version);
+
+            if (version <= 1) {
+                assertFalse(struct.hasField(ERROR_CODE));
+
+                // Partition level error populated in older versions.
+                assertEquals(Errors.GROUP_AUTHORIZATION_FAILED, oldResponse.error());
+                assertEquals(Utils.mkMap(Utils.mkEntry(Errors.GROUP_AUTHORIZATION_FAILED, 2),
+                        Utils.mkEntry(Errors.TOPIC_AUTHORIZATION_FAILED, 1)), oldResponse.errorCounts());
+
+            } else {
+                assertTrue(struct.hasField(ERROR_CODE));
+
+                assertEquals(Errors.NONE, oldResponse.error());
+                assertEquals(Utils.mkMap(
+                        Utils.mkEntry(Errors.NONE, 1),
+                        Utils.mkEntry(Errors.GROUP_AUTHORIZATION_FAILED, 1),
+                        Utils.mkEntry(Errors.TOPIC_AUTHORIZATION_FAILED, 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() {
+        PartitionData pd = new PartitionData(
+            offset,
+            leaderEpochOne,
+            null,
+            Errors.UNKNOWN_TOPIC_OR_PARTITION);
+        // test PartitionData.equals with null metadata
+        assertEquals(pd, pd);
+        partitionDataMap.clear();
+        partitionDataMap.put(new TopicPartition(topicOne, partitionOne), pd);
+
+        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..9754161eddde8
--- /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.message.OffsetForLeaderEpochRequestData.OffsetForLeaderTopicCollection;
+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(new OffsetForLeaderTopicCollection());
+        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.parse(request.serialize(), version);
+            if (version < 3)
+                assertEquals(OffsetsForLeaderEpochRequest.DEBUGGING_REPLICA_ID, parsed.replicaId());
+            else
+                assertEquals(replicaId, parsed.replicaId());
+        }
+    }
+
+}
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 ef17c96c5ad01..e2b4cee3c8435 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,21 +17,21 @@
 
 package org.apache.kafka.common.requests;
 
-import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.InvalidRecordException;
+import org.apache.kafka.common.message.ProduceRequestData;
 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;
+import org.apache.kafka.common.record.RecordVersion;
 import org.apache.kafka.common.record.SimpleRecord;
 import org.apache.kafka.common.record.TimestampType;
 import org.junit.Test;
 
 import java.nio.ByteBuffer;
+import java.util.Arrays;
 import java.util.Collections;
-import java.util.HashMap;
-import java.util.Map;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
@@ -48,30 +48,47 @@ public class ProduceRequestTest {
     public void shouldBeFlaggedAsTransactionalWhenTransactionalRecords() throws Exception {
         final MemoryRecords memoryRecords = MemoryRecords.withTransactionalRecords(0, CompressionType.NONE, 1L,
                 (short) 1, 1, 1, simpleRecord);
-        final ProduceRequest request = ProduceRequest.Builder.forCurrentMagic((short) -1,
-                10, Collections.singletonMap(new TopicPartition("topic", 1), memoryRecords)).build();
-        assertTrue(request.isTransactional());
+
+        final ProduceRequest request = ProduceRequest.forCurrentMagic(new ProduceRequestData()
+                .setTopicData(new ProduceRequestData.TopicProduceDataCollection(Collections.singletonList(
+                    new ProduceRequestData.TopicProduceData()
+                        .setName("topic")
+                        .setPartitionData(Collections.singletonList(
+                            new ProduceRequestData.PartitionProduceData()
+                                .setIndex(1)
+                                .setRecords(memoryRecords)))).iterator()))
+                .setAcks((short) -1)
+                .setTimeoutMs(10)).build();
+        assertTrue(RequestUtils.hasTransactionalRecords(request));
     }
 
     @Test
     public void shouldNotBeFlaggedAsTransactionalWhenNoRecords() throws Exception {
         final ProduceRequest request = createNonIdempotentNonTransactionalRecords();
-        assertFalse(request.isTransactional());
+        assertFalse(RequestUtils.hasTransactionalRecords(request));
     }
 
     @Test
     public void shouldNotBeFlaggedAsIdempotentWhenRecordsNotIdempotent() throws Exception {
         final ProduceRequest request = createNonIdempotentNonTransactionalRecords();
-        assertFalse(request.isTransactional());
+        assertFalse(RequestUtils.hasTransactionalRecords(request));
     }
 
     @Test
     public void shouldBeFlaggedAsIdempotentWhenIdempotentRecords() throws Exception {
         final MemoryRecords memoryRecords = MemoryRecords.withIdempotentRecords(1, CompressionType.NONE, 1L,
                 (short) 1, 1, 1, simpleRecord);
-        final ProduceRequest request = ProduceRequest.Builder.forCurrentMagic((short) -1, 10,
-                Collections.singletonMap(new TopicPartition("topic", 1), memoryRecords)).build();
-        assertTrue(request.isIdempotent());
+        final ProduceRequest request = ProduceRequest.forCurrentMagic(new ProduceRequestData()
+                .setTopicData(new ProduceRequestData.TopicProduceDataCollection(Collections.singletonList(
+                    new ProduceRequestData.TopicProduceData()
+                        .setName("topic")
+                        .setPartitionData(Collections.singletonList(
+                            new ProduceRequestData.PartitionProduceData()
+                                .setIndex(1)
+                                .setRecords(memoryRecords)))).iterator()))
+                .setAcks((short) -1)
+                .setTimeoutMs(10)).build();
+        assertTrue(RequestUtils.hasIdempotentRecords(request));
     }
 
     @Test
@@ -80,11 +97,14 @@ public void testBuildWithOldMessageFormat() {
         MemoryRecordsBuilder builder = MemoryRecords.builder(buffer, RecordBatch.MAGIC_VALUE_V1, CompressionType.NONE,
                 TimestampType.CREATE_TIME, 0L);
         builder.append(10L, null, "a".getBytes());
-        Map produceData = new HashMap<>();
-        produceData.put(new TopicPartition("test", 0), builder.build());
-
-        ProduceRequest.Builder requestBuilder = ProduceRequest.Builder.forMagic(RecordBatch.MAGIC_VALUE_V1, (short) 1,
-                5000, produceData, null);
+        ProduceRequest.Builder requestBuilder = ProduceRequest.forMagic(RecordBatch.MAGIC_VALUE_V1,
+            new ProduceRequestData()
+                .setTopicData(new ProduceRequestData.TopicProduceDataCollection(Collections.singletonList(
+                    new ProduceRequestData.TopicProduceData().setName("test").setPartitionData(Collections.singletonList(
+                        new ProduceRequestData.PartitionProduceData().setIndex(9).setRecords(builder.build()))))
+                    .iterator()))
+                .setAcks((short) 1)
+                .setTimeoutMs(5000));
         assertEquals(2, requestBuilder.oldestAllowedVersion());
         assertEquals(2, requestBuilder.latestAllowedVersion());
     }
@@ -95,11 +115,14 @@ public void testBuildWithCurrentMessageFormat() {
         MemoryRecordsBuilder builder = MemoryRecords.builder(buffer, RecordBatch.CURRENT_MAGIC_VALUE,
                 CompressionType.NONE, TimestampType.CREATE_TIME, 0L);
         builder.append(10L, null, "a".getBytes());
-        Map produceData = new HashMap<>();
-        produceData.put(new TopicPartition("test", 0), builder.build());
-
-        ProduceRequest.Builder requestBuilder = ProduceRequest.Builder.forMagic(RecordBatch.CURRENT_MAGIC_VALUE,
-                (short) 1, 5000, produceData, null);
+        ProduceRequest.Builder requestBuilder = ProduceRequest.forMagic(RecordBatch.CURRENT_MAGIC_VALUE,
+                new ProduceRequestData()
+                        .setTopicData(new ProduceRequestData.TopicProduceDataCollection(Collections.singletonList(
+                                new ProduceRequestData.TopicProduceData().setName("test").setPartitionData(Collections.singletonList(
+                                        new ProduceRequestData.PartitionProduceData().setIndex(9).setRecords(builder.build()))))
+                                .iterator()))
+                        .setAcks((short) 1)
+                        .setTimeoutMs(5000));
         assertEquals(3, requestBuilder.oldestAllowedVersion());
         assertEquals(ApiKeys.PRODUCE.latestVersion(), requestBuilder.latestAllowedVersion());
     }
@@ -118,17 +141,31 @@ public void testV3AndAboveShouldContainOnlyOneRecordBatch() {
 
         buffer.flip();
 
-        Map produceData = new HashMap<>();
-        produceData.put(new TopicPartition("test", 0), MemoryRecords.readableRecords(buffer));
-        ProduceRequest.Builder requestBuilder = ProduceRequest.Builder.forCurrentMagic((short) 1, 5000, produceData);
+        ProduceRequest.Builder requestBuilder = ProduceRequest.forCurrentMagic(new ProduceRequestData()
+                .setTopicData(new ProduceRequestData.TopicProduceDataCollection(Collections.singletonList(
+                    new ProduceRequestData.TopicProduceData()
+                        .setName("test")
+                        .setPartitionData(Collections.singletonList(
+                            new ProduceRequestData.PartitionProduceData()
+                                .setIndex(0)
+                                    .setRecords(MemoryRecords.readableRecords(buffer))))).iterator()))
+                .setAcks((short) 1)
+                .setTimeoutMs(5000));
         assertThrowsInvalidRecordExceptionForAllVersions(requestBuilder);
     }
 
     @Test
     public void testV3AndAboveCannotHaveNoRecordBatches() {
-        Map produceData = new HashMap<>();
-        produceData.put(new TopicPartition("test", 0), MemoryRecords.EMPTY);
-        ProduceRequest.Builder requestBuilder = ProduceRequest.Builder.forCurrentMagic((short) 1, 5000, produceData);
+        ProduceRequest.Builder requestBuilder = ProduceRequest.forCurrentMagic(new ProduceRequestData()
+                .setTopicData(new ProduceRequestData.TopicProduceDataCollection(Collections.singletonList(
+                        new ProduceRequestData.TopicProduceData()
+                                .setName("test")
+                                .setPartitionData(Collections.singletonList(
+                                        new ProduceRequestData.PartitionProduceData()
+                                                .setIndex(0)
+                                                .setRecords(MemoryRecords.EMPTY)))).iterator()))
+                .setAcks((short) 1)
+                .setTimeoutMs(5000));
         assertThrowsInvalidRecordExceptionForAllVersions(requestBuilder);
     }
 
@@ -139,9 +176,16 @@ public void testV3AndAboveCannotUseMagicV0() {
                 TimestampType.NO_TIMESTAMP_TYPE, 0L);
         builder.append(10L, null, "a".getBytes());
 
-        Map produceData = new HashMap<>();
-        produceData.put(new TopicPartition("test", 0), builder.build());
-        ProduceRequest.Builder requestBuilder = ProduceRequest.Builder.forCurrentMagic((short) 1, 5000, produceData);
+        ProduceRequest.Builder requestBuilder = ProduceRequest.forCurrentMagic(new ProduceRequestData()
+                .setTopicData(new ProduceRequestData.TopicProduceDataCollection(Collections.singletonList(
+                        new ProduceRequestData.TopicProduceData()
+                                .setName("test")
+                                .setPartitionData(Collections.singletonList(
+                                        new ProduceRequestData.PartitionProduceData()
+                                                .setIndex(0)
+                                                .setRecords(builder.build())))).iterator()))
+                .setAcks((short) 1)
+                .setTimeoutMs(5000));
         assertThrowsInvalidRecordExceptionForAllVersions(requestBuilder);
     }
 
@@ -152,12 +196,101 @@ public void testV3AndAboveCannotUseMagicV1() {
                 TimestampType.CREATE_TIME, 0L);
         builder.append(10L, null, "a".getBytes());
 
-        Map produceData = new HashMap<>();
-        produceData.put(new TopicPartition("test", 0), builder.build());
-        ProduceRequest.Builder requestBuilder = ProduceRequest.Builder.forCurrentMagic((short) 1, 5000, produceData);
+        ProduceRequest.Builder requestBuilder = ProduceRequest.forCurrentMagic(new ProduceRequestData()
+                .setTopicData(new ProduceRequestData.TopicProduceDataCollection(Collections.singletonList(
+                        new ProduceRequestData.TopicProduceData()
+                                .setName("test")
+                                .setPartitionData(Collections.singletonList(new ProduceRequestData.PartitionProduceData()
+                                        .setIndex(0)
+                                        .setRecords(builder.build()))))
+                        .iterator()))
+                .setAcks((short) 1)
+                .setTimeoutMs(5000));
         assertThrowsInvalidRecordExceptionForAllVersions(requestBuilder);
     }
 
+    @Test
+    public void testV6AndBelowCannotUseZStdCompression() {
+        ByteBuffer buffer = ByteBuffer.allocate(256);
+        MemoryRecordsBuilder builder = MemoryRecords.builder(buffer, RecordBatch.MAGIC_VALUE_V2, CompressionType.ZSTD,
+            TimestampType.CREATE_TIME, 0L);
+        builder.append(10L, null, "a".getBytes());
+
+        ProduceRequestData produceData = new ProduceRequestData()
+                .setTopicData(new ProduceRequestData.TopicProduceDataCollection(Collections.singletonList(
+                    new ProduceRequestData.TopicProduceData()
+                        .setName("test")
+                        .setPartitionData(Collections.singletonList(new ProduceRequestData.PartitionProduceData()
+                            .setIndex(0)
+                            .setRecords(builder.build()))))
+                    .iterator()))
+                .setAcks((short) 1)
+                .setTimeoutMs(1000);
+        // Can't create ProduceRequest instance with version within [3, 7)
+        for (short version = 3; version < 7; version++) {
+
+            ProduceRequest.Builder requestBuilder = new ProduceRequest.Builder(version, version, produceData);
+            assertThrowsInvalidRecordExceptionForAllVersions(requestBuilder);
+        }
+
+        // Works fine with current version (>= 7)
+        ProduceRequest.forCurrentMagic(produceData);
+    }
+
+    @Test
+    public void testMixedTransactionalData() {
+        final long producerId = 15L;
+        final short producerEpoch = 5;
+        final int sequence = 10;
+        final String transactionalId = "txnlId";
+
+        final MemoryRecords nonTxnRecords = MemoryRecords.withRecords(CompressionType.NONE,
+                new SimpleRecord("foo".getBytes()));
+        final MemoryRecords txnRecords = MemoryRecords.withTransactionalRecords(CompressionType.NONE, producerId,
+                producerEpoch, sequence, new SimpleRecord("bar".getBytes()));
+
+        ProduceRequest.Builder builder = ProduceRequest.forMagic(RecordBatch.CURRENT_MAGIC_VALUE,
+                new ProduceRequestData()
+                        .setTopicData(new ProduceRequestData.TopicProduceDataCollection(Arrays.asList(
+                                new ProduceRequestData.TopicProduceData().setName("foo").setPartitionData(Collections.singletonList(
+                                        new ProduceRequestData.PartitionProduceData().setIndex(0).setRecords(txnRecords))),
+                                new ProduceRequestData.TopicProduceData().setName("foo").setPartitionData(Collections.singletonList(
+                                        new ProduceRequestData.PartitionProduceData().setIndex(1).setRecords(nonTxnRecords))))
+                                .iterator()))
+                        .setAcks((short) -1)
+                        .setTimeoutMs(5000));
+        final ProduceRequest request = builder.build();
+        assertTrue(RequestUtils.hasTransactionalRecords(request));
+        assertTrue(RequestUtils.hasIdempotentRecords(request));
+    }
+
+    @Test
+    public void testMixedIdempotentData() {
+        final long producerId = 15L;
+        final short producerEpoch = 5;
+        final int sequence = 10;
+
+        final MemoryRecords nonTxnRecords = MemoryRecords.withRecords(CompressionType.NONE,
+                new SimpleRecord("foo".getBytes()));
+        final MemoryRecords txnRecords = MemoryRecords.withIdempotentRecords(CompressionType.NONE, producerId,
+                producerEpoch, sequence, new SimpleRecord("bar".getBytes()));
+
+        ProduceRequest.Builder builder = ProduceRequest.forMagic(RecordVersion.current().value,
+                new ProduceRequestData()
+                        .setTopicData(new ProduceRequestData.TopicProduceDataCollection(Arrays.asList(
+                                new ProduceRequestData.TopicProduceData().setName("foo").setPartitionData(Collections.singletonList(
+                                        new ProduceRequestData.PartitionProduceData().setIndex(0).setRecords(txnRecords))),
+                                new ProduceRequestData.TopicProduceData().setName("foo").setPartitionData(Collections.singletonList(
+                                        new ProduceRequestData.PartitionProduceData().setIndex(1).setRecords(nonTxnRecords))))
+                                .iterator()))
+                        .setAcks((short) -1)
+                        .setTimeoutMs(5000));
+
+        final ProduceRequest request = builder.build();
+        assertFalse(RequestUtils.hasTransactionalRecords(request));
+        assertTrue(RequestUtils.hasIdempotentRecords(request));
+    }
+
     private void assertThrowsInvalidRecordExceptionForAllVersions(ProduceRequest.Builder builder) {
         for (short version = builder.oldestAllowedVersion(); version < builder.latestAllowedVersion(); version++) {
             assertThrowsInvalidRecordException(builder, version);
@@ -166,7 +299,7 @@ private void assertThrowsInvalidRecordExceptionForAllVersions(ProduceRequest.Bui
 
     private void assertThrowsInvalidRecordException(ProduceRequest.Builder builder, short version) {
         try {
-            builder.build(version).toStruct();
+            builder.build(version).serialize();
             fail("Builder did not raise " + InvalidRecordException.class.getName() + " as expected");
         } catch (RuntimeException e) {
             assertTrue("Unexpected exception type " + e.getClass().getName(),
@@ -175,8 +308,15 @@ private void assertThrowsInvalidRecordException(ProduceRequest.Builder builder,
     }
 
     private ProduceRequest createNonIdempotentNonTransactionalRecords() {
-        final MemoryRecords memoryRecords = MemoryRecords.withRecords(CompressionType.NONE, simpleRecord);
-        return ProduceRequest.Builder.forCurrentMagic((short) -1, 10,
-                Collections.singletonMap(new TopicPartition("topic", 1), memoryRecords)).build();
+        return ProduceRequest.forCurrentMagic(new ProduceRequestData()
+                .setTopicData(new ProduceRequestData.TopicProduceDataCollection(Collections.singletonList(
+                        new ProduceRequestData.TopicProduceData()
+                                .setName("topic")
+                                .setPartitionData(Collections.singletonList(new ProduceRequestData.PartitionProduceData()
+                                        .setIndex(1)
+                                        .setRecords(MemoryRecords.withRecords(CompressionType.NONE, simpleRecord)))))
+                        .iterator()))
+                .setAcks((short) -1)
+                .setTimeoutMs(10)).build();
     }
 }
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..6af14142efd82
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/ProduceResponseTest.java
@@ -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 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.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.assertNull;
+import static org.junit.Assert.assertTrue;
+
+public class ProduceResponseTest {
+
+    @SuppressWarnings("deprecation")
+    @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 = RequestTestUtils.serializeResponseWithHeader(v5Response, version, 0);
+
+        ResponseHeader.parse(buffer, ApiKeys.PRODUCE.responseHeaderVersion(version)); // throw away.
+        ProduceResponse v5FromBytes = (ProduceResponse) AbstractResponse.parseResponse(ApiKeys.PRODUCE,
+                buffer, 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());
+    }
+
+    @SuppressWarnings("deprecation")
+    @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("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());
+    }
+
+    @SuppressWarnings("deprecation")
+    @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);
+            ProduceResponse.PartitionResponse deserialized = ProduceResponse.parse(response.serialize(ver), ver).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());
+                assertNull(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 baf0faf01c916..5bef2e42a530c 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,18 +16,20 @@
  */
 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.message.CreateTopicsResponseData;
+import org.apache.kafka.common.network.ClientInformation;
 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.security.auth.SecurityProtocol;
-import org.apache.kafka.common.protocol.types.Struct;
 import org.apache.kafka.common.security.auth.KafkaPrincipal;
+import org.apache.kafka.common.security.auth.SecurityProtocol;
 import org.junit.Test;
 
 import java.net.InetAddress;
 import java.nio.ByteBuffer;
-import java.util.Collections;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
@@ -40,7 +42,8 @@ public void testSerdeUnsupportedApiVersionRequest() throws Exception {
 
         RequestHeader header = new RequestHeader(ApiKeys.API_VERSIONS, Short.MAX_VALUE, "", correlationId);
         RequestContext context = new RequestContext(header, "0", InetAddress.getLocalHost(), KafkaPrincipal.ANONYMOUS,
-                new ListenerName("ssl"), SecurityProtocol.SASL_SSL);
+                new ListenerName("ssl"), SecurityProtocol.SASL_SSL, ClientInformation.EMPTY, false);
+        assertEquals(0, context.apiVersion());
 
         // Write some garbage to the request buffer. This should be ignored since we will treat
         // the unknown version type as v0 which has an empty request body.
@@ -54,8 +57,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.buildResponseSend(new ApiVersionsResponse(new ApiVersionsResponseData()
+            .setThrottleTimeMs(0)
+            .setErrorCode(Errors.UNSUPPORTED_VERSION.code())
+            .setApiKeys(new ApiVersionsResponseKeyCollection())));
         ByteBufferChannel channel = new ByteBufferChannel(256);
         send.writeTo(channel);
 
@@ -63,13 +68,40 @@ 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);
-        assertEquals(Errors.UNSUPPORTED_VERSION, response.error());
-        assertTrue(response.apiVersions().isEmpty());
+        ApiVersionsResponse response = (ApiVersionsResponse) AbstractResponse.parseResponse(ApiKeys.API_VERSIONS,
+            responseBuffer, (short) 0);
+        assertEquals(Errors.UNSUPPORTED_VERSION.code(), response.data().errorCode());
+        assertTrue(response.data().apiKeys().isEmpty());
+    }
+
+    @Test
+    public void testEnvelopeResponseSerde() throws Exception {
+        CreateTopicsResponseData.CreatableTopicResultCollection collection =
+            new CreateTopicsResponseData.CreatableTopicResultCollection();
+        collection.add(new CreateTopicsResponseData.CreatableTopicResult()
+            .setTopicConfigErrorCode(Errors.CLUSTER_AUTHORIZATION_FAILED.code())
+            .setNumPartitions(5));
+        CreateTopicsResponseData expectedResponse = new CreateTopicsResponseData()
+            .setThrottleTimeMs(10)
+            .setTopics(collection);
+
+        int correlationId = 15;
+        String clientId = "clientId";
+        RequestHeader header = new RequestHeader(ApiKeys.CREATE_TOPICS, ApiKeys.CREATE_TOPICS.latestVersion(),
+            clientId, correlationId);
+
+        RequestContext context = new RequestContext(header, "0", InetAddress.getLocalHost(),
+            KafkaPrincipal.ANONYMOUS, new ListenerName("ssl"), SecurityProtocol.SASL_SSL,
+            ClientInformation.EMPTY, true);
+
+        ByteBuffer buffer = context.buildResponseEnvelopePayload(new CreateTopicsResponse(expectedResponse));
+        assertEquals("Buffer limit and capacity should be the same", buffer.capacity(), buffer.limit());
+        CreateTopicsResponse parsedResponse = (CreateTopicsResponse) AbstractResponse.parseResponse(buffer, header);
+        assertEquals(expectedResponse, parsedResponse.data());
     }
 
 }
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..e9b72777f0a63 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
@@ -17,12 +17,11 @@
 package org.apache.kafka.common.requests;
 
 import org.apache.kafka.common.protocol.ApiKeys;
-import org.apache.kafka.common.protocol.types.Struct;
+import org.apache.kafka.common.protocol.ObjectSerializationCache;
 import org.junit.Test;
 
 import java.nio.ByteBuffer;
 
-import static org.apache.kafka.test.TestUtils.toBuffer;
 import static org.junit.Assert.assertEquals;
 
 public class RequestHeaderTest {
@@ -30,11 +29,11 @@ public class RequestHeaderTest {
     @Test
     public void testSerdeControlledShutdownV0() {
         // Verify that version 0 of controlled shutdown does not include the clientId field
-
+        short apiVersion = 0;
         int correlationId = 2342;
         ByteBuffer rawBuffer = ByteBuffer.allocate(32);
         rawBuffer.putShort(ApiKeys.CONTROLLED_SHUTDOWN.id);
-        rawBuffer.putShort((short) 0);
+        rawBuffer.putShort(apiVersion);
         rawBuffer.putInt(correlationId);
         rawBuffer.flip();
 
@@ -43,9 +42,9 @@ 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);
+        ByteBuffer serializedBuffer = RequestTestUtils.serializeRequestHeader(deserialized);
 
         assertEquals(ApiKeys.CONTROLLED_SHUTDOWN.id, serializedBuffer.getShort(0));
         assertEquals(0, serializedBuffer.getShort(2));
@@ -54,23 +53,44 @@ public void testSerdeControlledShutdownV0() {
     }
 
     @Test
-    public void testRequestHeader() {
-        RequestHeader header = new RequestHeader(ApiKeys.FIND_COORDINATOR, (short) 1, "", 10);
-        ByteBuffer buffer = toBuffer(header.toStruct());
+    public void testRequestHeaderV1() {
+        short apiVersion = 1;
+        RequestHeader header = new RequestHeader(ApiKeys.FIND_COORDINATOR, apiVersion, "", 10);
+        assertEquals(1, header.headerVersion());
+
+        ByteBuffer buffer = RequestTestUtils.serializeRequestHeader(header);
+        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() {
+        short apiVersion = 2;
+        RequestHeader header = new RequestHeader(ApiKeys.CREATE_DELEGATION_TOKEN, apiVersion, "", 10);
+        assertEquals(2, header.headerVersion());
+
+        ByteBuffer buffer = RequestTestUtils.serializeRequestHeader(header);
+        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);
     }
 
+    @Test
+    public void parseHeaderFromBufferWithNonZeroPosition() {
+        ByteBuffer buffer = ByteBuffer.allocate(64);
+        buffer.position(10);
+
+        RequestHeader header = new RequestHeader(ApiKeys.FIND_COORDINATOR, (short) 1, "", 10);
+        ObjectSerializationCache serializationCache = new ObjectSerializationCache();
+        // size must be called before write to avoid an NPE with the current implementation
+        header.size(serializationCache);
+        header.write(buffer, serializationCache);
+        int limit = buffer.position();
+        buffer.position(10);
+        buffer.limit(limit);
+
+        RequestHeader parsed = RequestHeader.parse(buffer);
+        assertEquals(header, parsed);
+    }
 }
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 edd1314c48bde..2ebc3c4b9cbe3 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
@@ -16,7 +16,9 @@
  */
 package org.apache.kafka.common.requests;
 
-import org.apache.kafka.clients.admin.NewPartitions;
+import org.apache.kafka.common.ConsumerGroupState;
+import org.apache.kafka.common.ElectionType;
+import org.apache.kafka.common.IsolationLevel;
 import org.apache.kafka.common.Node;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.acl.AccessControlEntry;
@@ -25,281 +27,634 @@
 import org.apache.kafka.common.acl.AclBindingFilter;
 import org.apache.kafka.common.acl.AclOperation;
 import org.apache.kafka.common.acl.AclPermissionType;
-import org.apache.kafka.common.errors.InvalidReplicaAssignmentException;
-import org.apache.kafka.common.errors.InvalidTopicException;
+import org.apache.kafka.common.config.ConfigResource;
 import org.apache.kafka.common.errors.NotCoordinatorException;
 import org.apache.kafka.common.errors.NotEnoughReplicasException;
 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.AddOffsetsToTxnRequestData;
+import org.apache.kafka.common.message.AddOffsetsToTxnResponseData;
+import org.apache.kafka.common.message.AlterClientQuotasResponseData;
+import org.apache.kafka.common.message.AlterConfigsResponseData;
+import org.apache.kafka.common.message.AlterPartitionReassignmentsRequestData;
+import org.apache.kafka.common.message.AlterPartitionReassignmentsResponseData;
+import org.apache.kafka.common.message.AlterReplicaLogDirsRequestData;
+import org.apache.kafka.common.message.AlterReplicaLogDirsRequestData.AlterReplicaLogDirTopic;
+import org.apache.kafka.common.message.AlterReplicaLogDirsRequestData.AlterReplicaLogDirTopicCollection;
+import org.apache.kafka.common.message.AlterReplicaLogDirsResponseData;
+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;
+import org.apache.kafka.common.message.ControlledShutdownResponseData.RemainingPartition;
+import org.apache.kafka.common.message.ControlledShutdownResponseData.RemainingPartitionCollection;
+import org.apache.kafka.common.message.CreateAclsRequestData;
+import org.apache.kafka.common.message.CreateAclsResponseData;
+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.CreatePartitionsRequestData;
+import org.apache.kafka.common.message.CreatePartitionsRequestData.CreatePartitionsAssignment;
+import org.apache.kafka.common.message.CreatePartitionsRequestData.CreatePartitionsTopic;
+import org.apache.kafka.common.message.CreatePartitionsRequestData.CreatePartitionsTopicCollection;
+import org.apache.kafka.common.message.CreatePartitionsResponseData;
+import org.apache.kafka.common.message.CreatePartitionsResponseData.CreatePartitionsTopicResult;
+import org.apache.kafka.common.message.CreateTopicsRequestData;
+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.CreateTopicsResponseData.CreatableTopicConfigs;
+import org.apache.kafka.common.message.CreateTopicsResponseData.CreatableTopicResult;
+import org.apache.kafka.common.message.DeleteAclsRequestData;
+import org.apache.kafka.common.message.DeleteAclsResponseData;
+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;
+import org.apache.kafka.common.message.DeleteTopicsResponseData.DeletableTopicResult;
+import org.apache.kafka.common.message.DescribeAclsResponseData;
+import org.apache.kafka.common.message.DescribeAclsResponseData.AclDescription;
+import org.apache.kafka.common.message.DescribeAclsResponseData.DescribeAclsResource;
+import org.apache.kafka.common.message.DescribeClientQuotasResponseData;
+import org.apache.kafka.common.message.DescribeConfigsRequestData;
+import org.apache.kafka.common.message.DescribeConfigsResponseData;
+import org.apache.kafka.common.message.DescribeConfigsResponseData.DescribeConfigsResourceResult;
+import org.apache.kafka.common.message.DescribeConfigsResponseData.DescribeConfigsResult;
+import org.apache.kafka.common.message.DescribeGroupsRequestData;
+import org.apache.kafka.common.message.DescribeGroupsResponseData;
+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.EndTxnRequestData;
+import org.apache.kafka.common.message.EndTxnResponseData;
+import org.apache.kafka.common.message.ExpireDelegationTokenRequestData;
+import org.apache.kafka.common.message.ExpireDelegationTokenResponseData;
+import org.apache.kafka.common.message.FetchRequestData;
+import org.apache.kafka.common.message.FetchResponseData;
+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;
+import org.apache.kafka.common.message.IncrementalAlterConfigsRequestData.AlterConfigsResource;
+import org.apache.kafka.common.message.IncrementalAlterConfigsRequestData.AlterableConfig;
+import org.apache.kafka.common.message.IncrementalAlterConfigsResponseData;
+import org.apache.kafka.common.message.IncrementalAlterConfigsResponseData.AlterConfigsResourceResponse;
+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.JoinGroupResponseData.JoinGroupResponseMember;
+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.ListOffsetRequestData.ListOffsetPartition;
+import org.apache.kafka.common.message.ListOffsetRequestData.ListOffsetTopic;
+import org.apache.kafka.common.message.ListOffsetResponseData;
+import org.apache.kafka.common.message.ListOffsetResponseData.ListOffsetPartitionResponse;
+import org.apache.kafka.common.message.ListOffsetResponseData.ListOffsetTopicResponse;
+import org.apache.kafka.common.message.ListPartitionReassignmentsRequestData;
+import org.apache.kafka.common.message.ListPartitionReassignmentsResponseData;
+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.OffsetForLeaderEpochResponseData.EpochEndOffset;
+import org.apache.kafka.common.message.OffsetForLeaderEpochRequestData.OffsetForLeaderPartition;
+import org.apache.kafka.common.message.OffsetForLeaderEpochRequestData.OffsetForLeaderTopic;
+import org.apache.kafka.common.message.OffsetForLeaderEpochRequestData.OffsetForLeaderTopicCollection;
+import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData;
+import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData.OffsetForLeaderTopicResult;
+import org.apache.kafka.common.message.ProduceRequestData;
+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.StopReplicaPartitionState;
+import org.apache.kafka.common.message.StopReplicaRequestData.StopReplicaTopicState;
+import org.apache.kafka.common.message.StopReplicaResponseData;
+import org.apache.kafka.common.message.SyncGroupRequestData;
+import org.apache.kafka.common.message.SyncGroupRequestData.SyncGroupRequestAssignment;
+import org.apache.kafka.common.message.SyncGroupResponseData;
+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.ByteBufferAccessor;
 import org.apache.kafka.common.protocol.Errors;
-import org.apache.kafka.common.security.auth.SecurityProtocol;
-import org.apache.kafka.common.protocol.types.Struct;
+import org.apache.kafka.common.protocol.ObjectSerializationCache;
+import org.apache.kafka.common.quota.ClientQuotaAlteration;
+import org.apache.kafka.common.quota.ClientQuotaEntity;
+import org.apache.kafka.common.quota.ClientQuotaFilter;
 import org.apache.kafka.common.record.CompressionType;
 import org.apache.kafka.common.record.MemoryRecords;
 import org.apache.kafka.common.record.RecordBatch;
 import org.apache.kafka.common.record.SimpleRecord;
-import org.apache.kafka.common.requests.CreateAclsRequest.AclCreation;
-import org.apache.kafka.common.requests.CreateAclsResponse.AclCreationResponse;
-import org.apache.kafka.common.requests.DeleteAclsResponse.AclDeletionResult;
-import org.apache.kafka.common.requests.DeleteAclsResponse.AclFilterResponse;
-import org.apache.kafka.common.resource.Resource;
-import org.apache.kafka.common.resource.ResourceFilter;
+import org.apache.kafka.common.requests.CreateTopicsRequest.Builder;
+import org.apache.kafka.common.requests.DescribeConfigsResponse.ConfigType;
+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;
 import org.apache.kafka.common.resource.ResourceType;
+import org.apache.kafka.common.security.auth.KafkaPrincipal;
+import org.apache.kafka.common.security.auth.SecurityProtocol;
+import org.apache.kafka.common.security.token.delegation.DelegationToken;
+import org.apache.kafka.common.security.token.delegation.TokenInformation;
+import org.apache.kafka.common.utils.SecurityUtils;
 import org.apache.kafka.common.utils.Utils;
 import org.junit.Test;
 
-import java.lang.reflect.InvocationTargetException;
-import java.lang.reflect.Method;
+import java.nio.BufferUnderflowException;
 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.HashSet;
 import java.util.LinkedHashMap;
+import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 
 import static java.util.Arrays.asList;
+import static java.util.Collections.emptyList;
 import static java.util.Collections.singletonList;
-import static org.apache.kafka.test.TestUtils.toBuffer;
+import static org.apache.kafka.common.protocol.ApiKeys.DESCRIBE_CONFIGS;
+import static org.apache.kafka.common.protocol.ApiKeys.FETCH;
+import static org.apache.kafka.common.protocol.ApiKeys.JOIN_GROUP;
+import static org.apache.kafka.common.protocol.ApiKeys.LIST_GROUPS;
+import static org.apache.kafka.common.protocol.ApiKeys.LIST_OFFSETS;
+import static org.apache.kafka.common.protocol.ApiKeys.SYNC_GROUP;
+import static org.apache.kafka.common.requests.FetchMetadata.INVALID_SESSION_ID;
 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;
 
 public class RequestResponseTest {
 
+    // Exception includes a message that we verify is not included in error responses
+    private final UnknownServerException unknownServerException = new UnknownServerException("secret");
+
     @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);
-        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(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(createOffsetCommitRequest(2));
-        checkErrorResponse(createOffsetCommitRequest(2), new UnknownServerException());
-        checkResponse(createOffsetCommitResponse(), 0);
-        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(true));
-        checkRequest(createStopReplicaRequest(false));
-        checkErrorResponse(createStopReplicaRequest(true), new UnknownServerException());
-        checkResponse(createStopReplicaResponse(), 0);
-        checkRequest(createLeaderAndIsrRequest());
-        checkErrorResponse(createLeaderAndIsrRequest(), new UnknownServerException());
-        checkResponse(createLeaderAndIsrResponse(), 0);
-        checkRequest(createSaslHandshakeRequest());
-        checkErrorResponse(createSaslHandshakeRequest(), new UnknownServerException());
-        checkResponse(createSaslHandshakeResponse(), 0);
-        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(createFindCoordinatorRequest(0), true);
+        checkRequest(createFindCoordinatorRequest(1), true);
+        checkErrorResponse(createFindCoordinatorRequest(0), unknownServerException, true);
+        checkErrorResponse(createFindCoordinatorRequest(1), unknownServerException, true);
+        checkResponse(createFindCoordinatorResponse(), 0, true);
+        checkResponse(createFindCoordinatorResponse(), 1, true);
+        checkRequest(createControlledShutdownRequest(), true);
+        checkResponse(createControlledShutdownResponse(), 1, true);
+        checkErrorResponse(createControlledShutdownRequest(), unknownServerException, true);
+        checkErrorResponse(createControlledShutdownRequest(0), unknownServerException, true);
+        checkRequest(createFetchRequest(4), true);
+        checkResponse(createFetchResponse(true), 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), true);
+        checkResponse(createFetchResponse(123), 7, true);
+        checkResponse(createFetchResponse(Errors.FETCH_SESSION_ID_NOT_FOUND, 123), 7, true);
+        checkErrorResponse(createFetchRequest(7), unknownServerException, true);
+        checkRequest(createHeartBeatRequest(), true);
+        checkErrorResponse(createHeartBeatRequest(), unknownServerException, true);
+        checkResponse(createHeartBeatResponse(), 0, true);
+
+        for (int v = ApiKeys.JOIN_GROUP.oldestVersion(); v <= ApiKeys.JOIN_GROUP.latestVersion(); v++) {
+            checkRequest(createJoinGroupRequest(v), true);
+            checkErrorResponse(createJoinGroupRequest(v), unknownServerException, true);
+            checkResponse(createJoinGroupResponse(v), v, true);
+        }
+
+        for (int v = ApiKeys.SYNC_GROUP.oldestVersion(); v <= ApiKeys.SYNC_GROUP.latestVersion(); v++) {
+            checkRequest(createSyncGroupRequest(v), true);
+            checkErrorResponse(createSyncGroupRequest(v), unknownServerException, true);
+            checkResponse(createSyncGroupResponse(v), v, true);
+        }
+
+        checkRequest(createLeaveGroupRequest(), true);
+        checkErrorResponse(createLeaveGroupRequest(), unknownServerException, true);
+        checkResponse(createLeaveGroupResponse(), 0, true);
+
+        for (short v = ApiKeys.LIST_GROUPS.oldestVersion(); v <= ApiKeys.LIST_GROUPS.latestVersion(); v++) {
+            checkRequest(createListGroupsRequest(v), false);
+            checkErrorResponse(createListGroupsRequest(v), unknownServerException, true);
+            checkResponse(createListGroupsResponse(v), v, true);
+        }
+
+        checkRequest(createDescribeGroupRequest(), true);
+        checkErrorResponse(createDescribeGroupRequest(), unknownServerException, true);
+        checkResponse(createDescribeGroupResponse(), 0, true);
+        checkRequest(createDeleteGroupsRequest(), true);
+        checkErrorResponse(createDeleteGroupsRequest(), unknownServerException, true);
+        checkResponse(createDeleteGroupsResponse(), 0, true);
+        for (int i = 0; i < ApiKeys.LIST_OFFSETS.latestVersion(); i++) {
+            checkRequest(createListOffsetRequest(i), true);
+            checkErrorResponse(createListOffsetRequest(i), 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")), unknownServerException, true);
+        checkResponse(createMetadataResponse(), 2, true);
+        checkErrorResponse(createMetadataRequest(2, Collections.singletonList("topic1")), unknownServerException, true);
+        checkResponse(createMetadataResponse(), 3, true);
+        checkErrorResponse(createMetadataRequest(3, Collections.singletonList("topic1")), unknownServerException, true);
+        checkResponse(createMetadataResponse(), 4, true);
+        checkErrorResponse(createMetadataRequest(4, Collections.singletonList("topic1")), unknownServerException, true);
+        checkRequest(createOffsetFetchRequestForAllPartition("group1", false), true);
+        checkRequest(createOffsetFetchRequestForAllPartition("group1", true), true);
+        checkErrorResponse(createOffsetFetchRequestForAllPartition("group1", false), new NotCoordinatorException("Not Coordinator"), true);
+        checkErrorResponse(createOffsetFetchRequestForAllPartition("group1", true), new NotCoordinatorException("Not Coordinator"), true);
+        checkRequest(createOffsetFetchRequest(0, false), true);
+        checkRequest(createOffsetFetchRequest(1, false), true);
+        checkRequest(createOffsetFetchRequest(2, false), true);
+        checkRequest(createOffsetFetchRequest(7, true), true);
+        checkRequest(createOffsetFetchRequestForAllPartition("group1", false), true);
+        checkRequest(createOffsetFetchRequestForAllPartition("group1", true), true);
+        checkErrorResponse(createOffsetFetchRequest(0, false), unknownServerException, true);
+        checkErrorResponse(createOffsetFetchRequest(1, false), unknownServerException, true);
+        checkErrorResponse(createOffsetFetchRequest(2, false), unknownServerException, true);
+        checkErrorResponse(createOffsetFetchRequest(7, true), unknownServerException, true);
+        checkResponse(createOffsetFetchResponse(), 0, true);
+        checkRequest(createProduceRequest(2), true);
+        checkErrorResponse(createProduceRequest(2), unknownServerException, true);
+        checkRequest(createProduceRequest(3), true);
+        checkErrorResponse(createProduceRequest(3), unknownServerException, true);
+        checkResponse(createProduceResponse(), 2, true);
+        checkResponse(createProduceResponseWithErrorMessage(), 8, true);
+
+        for (int v = ApiKeys.STOP_REPLICA.oldestVersion(); v <= ApiKeys.STOP_REPLICA.latestVersion(); v++) {
+            checkRequest(createStopReplicaRequest(v, true), true);
+            checkRequest(createStopReplicaRequest(v, false), true);
+            checkErrorResponse(createStopReplicaRequest(v, true), unknownServerException, true);
+            checkErrorResponse(createStopReplicaRequest(v, false), unknownServerException, true);
+            checkResponse(createStopReplicaResponse(), v, true);
+        }
+
+        checkRequest(createLeaderAndIsrRequest(0), true);
+        checkErrorResponse(createLeaderAndIsrRequest(0), unknownServerException, false);
+        checkRequest(createLeaderAndIsrRequest(1), true);
+        checkErrorResponse(createLeaderAndIsrRequest(1), unknownServerException, false);
+        checkRequest(createLeaderAndIsrRequest(2), true);
+        checkErrorResponse(createLeaderAndIsrRequest(2), unknownServerException, false);
+        checkResponse(createLeaderAndIsrResponse(), 0, true);
+        checkRequest(createSaslHandshakeRequest(), true);
+        checkErrorResponse(createSaslHandshakeRequest(), unknownServerException, true);
+        checkResponse(createSaslHandshakeResponse(), 0, true);
+        checkRequest(createSaslAuthenticateRequest(), true);
+        checkErrorResponse(createSaslAuthenticateRequest(), unknownServerException, true);
+        checkResponse(createSaslAuthenticateResponse(), 0, true);
+        checkResponse(createSaslAuthenticateResponse(), 1, true);
+        checkRequest(createApiVersionRequest(), true);
+        checkErrorResponse(createApiVersionRequest(), 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);
+
+        for (int v = ApiKeys.CREATE_TOPICS.oldestVersion(); v <= ApiKeys.CREATE_TOPICS.latestVersion(); v++) {
+            checkRequest(createCreateTopicRequest(v), true);
+            checkErrorResponse(createCreateTopicRequest(v), unknownServerException, true);
+            checkResponse(createCreateTopicResponse(), v, true);
+        }
+
+        for (int v = ApiKeys.DELETE_TOPICS.oldestVersion(); v <= ApiKeys.DELETE_TOPICS.latestVersion(); v++) {
+            checkRequest(createDeleteTopicsRequest(v), true);
+            checkErrorResponse(createDeleteTopicsRequest(v), unknownServerException, true);
+            checkResponse(createDeleteTopicsResponse(), v, true);
+        }
+
+        for (int v = ApiKeys.CREATE_PARTITIONS.oldestVersion(); v <= ApiKeys.CREATE_PARTITIONS.latestVersion(); v++) {
+            checkRequest(createCreatePartitionsRequest(v), true);
+            checkRequest(createCreatePartitionsRequestWithAssignments(v), false);
+            checkErrorResponse(createCreatePartitionsRequest(v), unknownServerException, true);
+            checkResponse(createCreatePartitionsResponse(), v, true);
+        }
+
+        checkRequest(createInitPidRequest(), true);
+        checkErrorResponse(createInitPidRequest(), unknownServerException, true);
+        checkResponse(createInitPidResponse(), 0, true);
+
+        checkRequest(createAddPartitionsToTxnRequest(), true);
+        checkResponse(createAddPartitionsToTxnResponse(), 0, true);
+        checkErrorResponse(createAddPartitionsToTxnRequest(), unknownServerException, true);
+        checkRequest(createAddOffsetsToTxnRequest(), true);
+        checkResponse(createAddOffsetsToTxnResponse(), 0, true);
+        checkErrorResponse(createAddOffsetsToTxnRequest(), unknownServerException, true);
+        checkRequest(createEndTxnRequest(), true);
+        checkResponse(createEndTxnResponse(), 0, true);
+        checkErrorResponse(createEndTxnRequest(), unknownServerException, true);
+        checkRequest(createWriteTxnMarkersRequest(), true);
+        checkResponse(createWriteTxnMarkersResponse(), 0, true);
+        checkErrorResponse(createWriteTxnMarkersRequest(), 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(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());
-        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());
-        checkRequest(createDescribeConfigsRequestWithConfigEntries());
-        checkErrorResponse(createDescribeConfigsRequest(), new UnknownServerException());
-        checkResponse(createDescribeConfigsResponse(), 0);
-        checkRequest(createCreatePartitionsRequest());
-        checkRequest(createCreatePartitionsRequestWithAssignments());
-        checkErrorResponse(createCreatePartitionsRequest(), new InvalidTopicException());
-        checkResponse(createCreatePartitionsResponse(), 0);
+        checkResponse(createMetadataResponse(), 0, true);
+        checkResponse(createMetadataResponse(), 1, true);
+        checkErrorResponse(createMetadataRequest(1, Collections.singletonList("topic1")), unknownServerException, true);
+        checkRequest(createOffsetCommitRequest(0), true);
+        checkErrorResponse(createOffsetCommitRequest(0), unknownServerException, true);
+        checkRequest(createOffsetCommitRequest(1), true);
+        checkErrorResponse(createOffsetCommitRequest(1), unknownServerException, true);
+        checkRequest(createOffsetCommitRequest(2), true);
+        checkErrorResponse(createOffsetCommitRequest(2), unknownServerException, true);
+        checkRequest(createOffsetCommitRequest(3), true);
+        checkErrorResponse(createOffsetCommitRequest(3), unknownServerException, true);
+        checkRequest(createOffsetCommitRequest(4), true);
+        checkErrorResponse(createOffsetCommitRequest(4), unknownServerException, true);
+        checkResponse(createOffsetCommitResponse(), 4, true);
+        checkRequest(createOffsetCommitRequest(5), true);
+        checkErrorResponse(createOffsetCommitRequest(5), unknownServerException, true);
+        checkResponse(createOffsetCommitResponse(), 5, true);
+        checkRequest(createJoinGroupRequest(0), true);
+        checkRequest(createUpdateMetadataRequest(0, null), false);
+        checkErrorResponse(createUpdateMetadataRequest(0, null), unknownServerException, true);
+        checkRequest(createUpdateMetadataRequest(1, null), false);
+        checkRequest(createUpdateMetadataRequest(1, "rack1"), false);
+        checkErrorResponse(createUpdateMetadataRequest(1, null), unknownServerException, true);
+        checkRequest(createUpdateMetadataRequest(2, "rack1"), false);
+        checkRequest(createUpdateMetadataRequest(2, null), false);
+        checkErrorResponse(createUpdateMetadataRequest(2, "rack1"), unknownServerException, true);
+        checkRequest(createUpdateMetadataRequest(3, "rack1"), false);
+        checkRequest(createUpdateMetadataRequest(3, null), false);
+        checkErrorResponse(createUpdateMetadataRequest(3, "rack1"), unknownServerException, true);
+        checkRequest(createUpdateMetadataRequest(4, "rack1"), false);
+        checkRequest(createUpdateMetadataRequest(4, null), false);
+        checkErrorResponse(createUpdateMetadataRequest(4, "rack1"), unknownServerException, true);
+        checkRequest(createUpdateMetadataRequest(5, "rack1"), false);
+        checkRequest(createUpdateMetadataRequest(5, null), false);
+        checkErrorResponse(createUpdateMetadataRequest(5, "rack1"), unknownServerException, true);
+        checkResponse(createUpdateMetadataResponse(), 0, true);
+        checkRequest(createListOffsetRequest(0), true);
+        checkErrorResponse(createListOffsetRequest(0), unknownServerException, true);
+        checkResponse(createListOffsetResponse(0), 0, true);
+        checkRequest(createLeaderEpochRequestForReplica(0, 1), true);
+        checkRequest(createLeaderEpochRequestForConsumer(), true);
+        checkResponse(createLeaderEpochResponse(), 0, true);
+        checkErrorResponse(createLeaderEpochRequestForConsumer(), unknownServerException, true);
+        checkRequest(createAddPartitionsToTxnRequest(), true);
+        checkErrorResponse(createAddPartitionsToTxnRequest(), unknownServerException, true);
+        checkResponse(createAddPartitionsToTxnResponse(), 0, true);
+        checkRequest(createAddOffsetsToTxnRequest(), true);
+        checkErrorResponse(createAddOffsetsToTxnRequest(), unknownServerException, true);
+        checkResponse(createAddOffsetsToTxnResponse(), 0, true);
+        checkRequest(createEndTxnRequest(), true);
+        checkErrorResponse(createEndTxnRequest(), unknownServerException, true);
+        checkResponse(createEndTxnResponse(), 0, true);
+        checkRequest(createWriteTxnMarkersRequest(), true);
+        checkErrorResponse(createWriteTxnMarkersRequest(), unknownServerException, true);
+        checkResponse(createWriteTxnMarkersResponse(), 0, true);
+        checkRequest(createTxnOffsetCommitRequest(0), true);
+        checkRequest(createTxnOffsetCommitRequest(3), true);
+        checkRequest(createTxnOffsetCommitRequestWithAutoDowngrade(2), true);
+        checkErrorResponse(createTxnOffsetCommitRequest(0), unknownServerException, true);
+        checkErrorResponse(createTxnOffsetCommitRequest(3), unknownServerException, true);
+        checkErrorResponse(createTxnOffsetCommitRequestWithAutoDowngrade(2), unknownServerException, true);
+        checkResponse(createTxnOffsetCommitResponse(), 0, true);
+        checkRequest(createDescribeAclsRequest(), true);
+        checkErrorResponse(createDescribeAclsRequest(), 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()), ApiKeys.DELETE_ACLS.latestVersion(), true);
+        checkRequest(createAlterConfigsRequest(), false);
+        checkErrorResponse(createAlterConfigsRequest(), unknownServerException, true);
+        checkResponse(createAlterConfigsResponse(), 0, false);
+        checkRequest(createDescribeConfigsRequest(0), true);
+        checkRequest(createDescribeConfigsRequestWithConfigEntries(0), false);
+        checkErrorResponse(createDescribeConfigsRequest(0), unknownServerException, true);
+        checkResponse(createDescribeConfigsResponse((short) 0), 0, false);
+        checkRequest(createDescribeConfigsRequest(1), true);
+        checkRequest(createDescribeConfigsRequestWithConfigEntries(1), false);
+        checkRequest(createDescribeConfigsRequestWithDocumentation(1), false);
+        checkRequest(createDescribeConfigsRequestWithDocumentation(2), false);
+        checkRequest(createDescribeConfigsRequestWithDocumentation(3), false);
+        checkErrorResponse(createDescribeConfigsRequest(1), unknownServerException, true);
+        checkResponse(createDescribeConfigsResponse((short) 1), 1, false);
+        checkDescribeConfigsResponseVersions();
+        checkRequest(createCreateTokenRequest(), true);
+        checkErrorResponse(createCreateTokenRequest(), unknownServerException, true);
+        checkResponse(createCreateTokenResponse(), 0, true);
+        checkRequest(createDescribeTokenRequest(), true);
+        checkErrorResponse(createDescribeTokenRequest(), unknownServerException, true);
+        checkResponse(createDescribeTokenResponse(), 0, true);
+        checkRequest(createExpireTokenRequest(), true);
+        checkErrorResponse(createExpireTokenRequest(), unknownServerException, true);
+        checkResponse(createExpireTokenResponse(), 0, true);
+        checkRequest(createRenewTokenRequest(), true);
+        checkErrorResponse(createRenewTokenRequest(), unknownServerException, true);
+        checkResponse(createRenewTokenResponse(), 0, true);
+        checkRequest(createElectLeadersRequest(), true);
+        checkRequest(createElectLeadersRequestNullPartitions(), true);
+        checkErrorResponse(createElectLeadersRequest(), unknownServerException, true);
+        checkResponse(createElectLeadersResponse(), 1, true);
+        checkRequest(createIncrementalAlterConfigsRequest(), true);
+        checkErrorResponse(createIncrementalAlterConfigsRequest(), unknownServerException, true);
+        checkResponse(createIncrementalAlterConfigsResponse(), 0, true);
+        checkRequest(createAlterPartitionReassignmentsRequest(), true);
+        checkErrorResponse(createAlterPartitionReassignmentsRequest(), unknownServerException, true);
+        checkResponse(createAlterPartitionReassignmentsResponse(), 0, true);
+        checkRequest(createListPartitionReassignmentsRequest(), true);
+        checkErrorResponse(createListPartitionReassignmentsRequest(), unknownServerException, true);
+        checkResponse(createListPartitionReassignmentsResponse(), 0, true);
+        checkRequest(createOffsetDeleteRequest(), true);
+        checkErrorResponse(createOffsetDeleteRequest(), unknownServerException, true);
+        checkResponse(createOffsetDeleteResponse(), 0, true);
+        checkRequest(createAlterReplicaLogDirsRequest(), true);
+        checkErrorResponse(createAlterReplicaLogDirsRequest(), unknownServerException, true);
+        checkResponse(createAlterReplicaLogDirsResponse(), 0, true);
+
+        checkRequest(createDescribeClientQuotasRequest(), true);
+        checkErrorResponse(createDescribeClientQuotasRequest(), unknownServerException, true);
+        checkResponse(createDescribeClientQuotasResponse(), 0, true);
+        checkRequest(createAlterClientQuotasRequest(), true);
+        checkErrorResponse(createAlterClientQuotasRequest(), unknownServerException, true);
+        checkResponse(createAlterClientQuotasResponse(), 0, true);
     }
 
     @Test
     public void testResponseHeader() {
-        ResponseHeader header = createResponseHeader();
-        ByteBuffer buffer = toBuffer(header.toStruct());
-        ResponseHeader deserialized = ResponseHeader.parse(buffer);
+        ResponseHeader header = createResponseHeader((short) 1);
+        ObjectSerializationCache serializationCache = new ObjectSerializationCache();
+        ByteBuffer buffer = ByteBuffer.allocate(header.size(serializationCache));
+        header.write(buffer, serializationCache);
+        buffer.flip();
+        ResponseHeader deserialized = ResponseHeader.parse(buffer, header.headerVersion());
         assertEquals(header.correlationId(), deserialized.correlationId());
     }
 
-    private void checkOlderFetchVersions() throws Exception {
-        int latestVersion = ApiKeys.FETCH.latestVersion();
+    private void checkOlderFetchVersions() {
+        int latestVersion = FETCH.latestVersion();
         for (int i = 0; i < latestVersion; ++i) {
-            checkErrorResponse(createFetchRequest(i), new UnknownServerException());
-            checkRequest(createFetchRequest(i));
-            checkResponse(createFetchResponse(), i);
+            if (i > 7) {
+                checkErrorResponse(createFetchRequest(i), unknownServerException, true);
+            }
+            checkRequest(createFetchRequest(i), true);
+            checkResponse(createFetchResponse(i >= 4), i, true);
         }
     }
 
-    private void checkErrorResponse(AbstractRequest req, Throwable e) throws Exception {
-        checkResponse(req.getErrorResponse(e), req.version());
+    private void verifyDescribeConfigsResponse(DescribeConfigsResponse expected, DescribeConfigsResponse actual,
+                                               int version) {
+        for (Map.Entry resource : expected.resultMap().entrySet()) {
+            List actualEntries = actual.resultMap().get(resource.getKey()).configs();
+            List expectedEntries = expected.resultMap().get(resource.getKey()).configs();
+            assertEquals(expectedEntries.size(), actualEntries.size());
+            for (int i = 0; i < actualEntries.size(); ++i) {
+                DescribeConfigsResourceResult actualEntry = actualEntries.get(i);
+                DescribeConfigsResourceResult expectedEntry = expectedEntries.get(i);
+                assertEquals(expectedEntry.name(), actualEntry.name());
+                assertEquals("Non-matching values for " + actualEntry.name() + " in version " + version,
+                        expectedEntry.value(), actualEntry.value());
+                assertEquals("Non-matching readonly for " + actualEntry.name() + " in version " + version,
+                        expectedEntry.readOnly(), actualEntry.readOnly());
+                assertEquals("Non-matching isSensitive for " + actualEntry.name() + " in version " + version,
+                        expectedEntry.isSensitive(), actualEntry.isSensitive());
+                if (version < 3) {
+                    assertEquals("Non-matching configType for " + actualEntry.name() + " in version " + version,
+                            ConfigType.UNKNOWN.id(), actualEntry.configType());
+                } else {
+                    assertEquals("Non-matching configType for " + actualEntry.name() + " in version " + version,
+                            expectedEntry.configType(), actualEntry.configType());
+                }
+                if (version == 0) {
+                    assertEquals("Non matching configSource for " + actualEntry.name() + " in version " + version,
+                            DescribeConfigsResponse.ConfigSource.STATIC_BROKER_CONFIG.id(), actualEntry.configSource());
+                } else {
+                    assertEquals("Non-matching configSource for " + actualEntry.name() + " in version " + version,
+                            expectedEntry.configSource(), actualEntry.configSource());
+                }
+            }
+        }
     }
 
-    private void checkRequest(AbstractRequest req) throws Exception {
-        // 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
-        Struct struct = req.toStruct();
-        AbstractRequest deserialized = (AbstractRequest) deserialize(req, struct, req.version());
-        deserialized.toStruct();
+    private void checkDescribeConfigsResponseVersions() {
+        for (int version = ApiKeys.DESCRIBE_CONFIGS.oldestVersion(); version < ApiKeys.DESCRIBE_CONFIGS.latestVersion(); ++version) {
+            short apiVersion = (short) version;
+            DescribeConfigsResponse response = createDescribeConfigsResponse(apiVersion);
+            DescribeConfigsResponse deserialized0 = (DescribeConfigsResponse) AbstractResponse.parseResponse(ApiKeys.DESCRIBE_CONFIGS,
+                    response.serialize(apiVersion), apiVersion);
+            verifyDescribeConfigsResponse(response, deserialized0, apiVersion);
+        }
+    }
+
+    private void checkErrorResponse(AbstractRequest req, Throwable e, boolean checkEqualityAndHashCode) {
+        AbstractResponse response = req.getErrorResponse(e);
+        checkResponse(response, req.version(), checkEqualityAndHashCode);
+        if (e instanceof UnknownServerException) {
+            String responseStr = response.toString();
+            assertFalse(String.format("Unknown message included in response for %s: %s ", req.apiKey(), responseStr),
+                    responseStr.contains(e.getMessage()));
+        }
     }
 
-    private void checkResponse(AbstractResponse response, int version) throws Exception {
+    private void checkRequest(AbstractRequest req, boolean checkEquality) {
         // 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);
-        deserialized.toStruct((short) version);
+        // Check for equality of the ByteBuffer 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 {
+            ByteBuffer serializedBytes = req.serialize();
+            AbstractRequest deserialized = AbstractRequest.parseRequest(req.apiKey(), req.version(), serializedBytes).request;
+            ByteBuffer serializedBytes2 = deserialized.serialize();
+            serializedBytes.rewind();
+            if (checkEquality)
+                assertEquals("Request " + req + "failed equality test", serializedBytes, serializedBytes2);
+        } catch (Exception e) {
+            throw new RuntimeException("Failed to deserialize request " + req + " with type " + req.getClass(), e);
+        }
     }
 
-    private AbstractRequestResponse deserialize(AbstractRequestResponse req, Struct struct, short version) throws NoSuchMethodException, IllegalAccessException, InvocationTargetException {
-        ByteBuffer buffer = toBuffer(struct);
-        Method deserializer = req.getClass().getDeclaredMethod("parse", ByteBuffer.class, Short.TYPE);
-        return (AbstractRequestResponse) deserializer.invoke(null, buffer, version);
+    private void checkResponse(AbstractResponse response, int version, boolean checkEquality) {
+        // Check that we can serialize, deserialize and serialize again
+        // 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 {
+            ByteBuffer serializedBytes = response.serialize((short) version);
+            AbstractResponse deserialized = AbstractResponse.parseResponse(response.apiKey(), serializedBytes, (short) version);
+            ByteBuffer serializedBytes2 = deserialized.serialize((short) version);
+            serializedBytes.rewind();
+            if (checkEquality)
+                assertEquals("Response " + response + "failed equality test", serializedBytes, serializedBytes2);
+        } catch (Exception e) {
+            throw new RuntimeException("Failed to deserialize response " + response + " with type " + response.getClass(), e);
+        }
     }
 
     @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);
     }
 
+    @Test
+    public void testPartitionSize() {
+        TopicPartition tp0 = new TopicPartition("test", 0);
+        TopicPartition tp1 = new TopicPartition("test", 1);
+        MemoryRecords records0 = MemoryRecords.withRecords(RecordBatch.MAGIC_VALUE_V2,
+            CompressionType.NONE, new SimpleRecord("woot".getBytes()));
+        MemoryRecords records1 = MemoryRecords.withRecords(RecordBatch.MAGIC_VALUE_V2,
+            CompressionType.NONE, new SimpleRecord("woot".getBytes()), new SimpleRecord("woot".getBytes()));
+        ProduceRequest request = ProduceRequest.forMagic(RecordBatch.MAGIC_VALUE_V2,
+                new ProduceRequestData()
+                        .setTopicData(new ProduceRequestData.TopicProduceDataCollection(Arrays.asList(
+                                new ProduceRequestData.TopicProduceData().setName(tp0.topic()).setPartitionData(
+                                        Collections.singletonList(new ProduceRequestData.PartitionProduceData().setIndex(tp0.partition()).setRecords(records0))),
+                                new ProduceRequestData.TopicProduceData().setName(tp1.topic()).setPartitionData(
+                                        Collections.singletonList(new ProduceRequestData.PartitionProduceData().setIndex(tp1.partition()).setRecords(records1))))
+                                .iterator()))
+                        .setAcks((short) 1)
+                        .setTimeoutMs(5000)
+                        .setTransactionalId("transactionalId"))
+            .build((short) 3);
+        assertEquals(2, request.partitionSizes().size());
+        assertEquals(records0.sizeInBytes(), (int) request.partitionSizes().get(tp0));
+        assertEquals(records1.sizeInBytes(), (int) request.partitionSizes().get(tp1));
+    }
+
     @Test
     public void produceRequestToStringTest() {
         ProduceRequest request = createProduceRequest(ApiKeys.PRODUCE.latestVersion());
-        assertEquals(1, request.partitionRecordsOrFail().size());
+        assertEquals(1, request.data().topicData().size());
         assertFalse(request.toString(false).contains("partitionSizes"));
         assertTrue(request.toString(false).contains("numPartitions=1"));
         assertTrue(request.toString(true).contains("partitionSizes"));
@@ -307,8 +662,8 @@ public void produceRequestToStringTest() {
 
         request.clearPartitionRecords();
         try {
-            request.partitionRecordsOrFail();
-            fail("partitionRecordsOrFail should fail after clearPartitionRecords()");
+            request.data();
+            fail("dataOrException should fail after clearPartitionRecords()");
         } catch (IllegalStateException e) {
             // OK
         }
@@ -320,10 +675,11 @@ public void produceRequestToStringTest() {
         assertFalse(request.toString(true).contains("numPartitions"));
     }
 
+    @SuppressWarnings("deprecation")
     @Test
     public void produceRequestGetErrorResponseTest() {
         ProduceRequest request = createProduceRequest(ApiKeys.PRODUCE.latestVersion());
-        Set partitions = new HashSet<>(request.partitionRecordsOrFail().keySet());
+        Set partitions = new HashSet<>(request.partitionSizes().keySet());
 
         ProduceResponse errorResponse = (ProduceResponse) request.getErrorResponse(new NotEnoughReplicasException());
         assertEquals(partitions, errorResponse.responses().keySet());
@@ -343,106 +699,59 @@ 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);
-
-        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.getThrottleTime());
-        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.getThrottleTime());
-        assertEquals("Throttle time must be 10", 10, v1Response.getThrottleTime());
-        assertEquals("Throttle time must be 10", 10, v2Response.getThrottleTime());
-        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<>();
+        LinkedHashMap> responseData = new LinkedHashMap<>();
 
         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));
+        responseData.put(new TopicPartition("test", 0), new FetchResponse.PartitionData<>(
+                Errors.NONE, 1000000, FetchResponse.INVALID_LAST_STABLE_OFFSET,
+                0L, Optional.empty(), Collections.emptyList(), records));
 
-        FetchResponse v0Response = new FetchResponse(responseData, 0);
-        FetchResponse v1Response = new FetchResponse(responseData, 10);
+        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),
-                v0Response.toStruct((short) 0).schema());
-        assertEquals("Should use schema version 1", ApiKeys.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());
     }
 
     @Test
     public void testFetchResponseV4() {
-        LinkedHashMap responseData = new LinkedHashMap<>();
-
+        LinkedHashMap> responseData = new LinkedHashMap<>();
         MemoryRecords records = MemoryRecords.readableRecords(ByteBuffer.allocate(10));
 
         List abortedTransactions = asList(
                 new FetchResponse.AbortedTransaction(10, 100),
                 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));
-        responseData.put(new TopicPartition("bar", 1), new FetchResponse.PartitionData(Errors.NONE, 900000,
-                5, FetchResponse.INVALID_LOG_START_OFFSET, null, records));
-        responseData.put(new TopicPartition("foo", 0), new FetchResponse.PartitionData(Errors.NONE, 70000,
-                6, FetchResponse.INVALID_LOG_START_OFFSET, Collections.emptyList(), records));
-
-        FetchResponse response = new FetchResponse(responseData, 10);
-        FetchResponse deserialized = FetchResponse.parse(toBuffer(response.toStruct((short) 4)), (short) 4);
+        responseData.put(new TopicPartition("bar", 0), new FetchResponse.PartitionData<>(Errors.NONE, 100000,
+                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, Optional.empty(), null, records));
+        responseData.put(new TopicPartition("foo", 0), new FetchResponse.PartitionData<>(Errors.NONE, 70000,
+                6, FetchResponse.INVALID_LOG_START_OFFSET, Optional.empty(), emptyList(), records));
+
+        FetchResponse response = new FetchResponse<>(Errors.NONE, responseData, 10, INVALID_SESSION_ID);
+        FetchResponse deserialized = FetchResponse.parse(response.serialize((short) 4), (short) 4);
         assertEquals(responseData, deserialized.responseData());
     }
 
     @Test
-    public void verifyFetchResponseFullWrite() throws Exception {
-        FetchResponse fetchResponse = createFetchResponse();
-        short apiVersion = ApiKeys.FETCH.latestVersion();
+    public void verifyFetchResponseFullWrites() throws Exception {
+        verifyFetchResponseFullWrite(FETCH.latestVersion(), createFetchResponse(123));
+        verifyFetchResponseFullWrite(FETCH.latestVersion(),
+            createFetchResponse(Errors.FETCH_SESSION_ID_NOT_FOUND, 123));
+        for (short version = 0; version <= FETCH.latestVersion(); version++) {
+            verifyFetchResponseFullWrite(version, createFetchResponse(version >= 4));
+        }
+    }
+
+    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(new ResponseHeader(correlationId, responseHeaderVersion), apiVersion);
         ByteBufferChannel channel = new ByteBufferChannel(send.size());
         send.writeTo(channel);
         channel.close();
@@ -454,25 +763,23 @@ public void verifyFetchResponseFullWrite() throws Exception {
         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);
-        assertEquals(fetchResponse.toStruct(apiVersion), responseBody);
-
-        assertEquals(size, responseHeader.sizeOf() + responseBody.sizeOf());
+        assertEquals(fetchResponse.serialize(apiVersion), buf);
+        FetchResponseData deserialized = new FetchResponseData(new ByteBufferAccessor(buf), apiVersion);
+        ObjectSerializationCache serializationCache = new ObjectSerializationCache();
+        assertEquals(size, responseHeader.size(serializationCache) + deserialized.size(serializationCache, apiVersion));
     }
 
     @Test
     public void testControlledShutdownResponse() {
         ControlledShutdownResponse response = createControlledShutdownResponse();
         short version = ApiKeys.CONTROLLED_SHUTDOWN.latestVersion();
-        Struct struct = response.toStruct(version);
-        ByteBuffer buffer = toBuffer(struct);
+        ByteBuffer buffer = response.serialize(version);
         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)
@@ -481,186 +788,539 @@ public void testCreateTopicRequestV0FailsIfValidateOnly() {
     }
 
     @Test
-    public void testFetchRequestMaxBytesOldVersions() throws Exception {
+    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() {
         final short version = 1;
         FetchRequest fr = createFetchRequest(version);
-        FetchRequest fr2 = new FetchRequest(fr.toStruct(), version);
+        FetchRequest fr2 = FetchRequest.parse(fr.serialize(), version);
         assertEquals(fr2.maxBytes(), fr.maxBytes());
     }
 
     @Test
     public void testFetchRequestIsolationLevel() throws Exception {
         FetchRequest request = createFetchRequest(4, IsolationLevel.READ_COMMITTED);
-        Struct struct = request.toStruct();
-        FetchRequest deserialized = (FetchRequest) deserialize(request, struct, request.version());
+        FetchRequest deserialized = (FetchRequest) AbstractRequest.parseRequest(request.apiKey(), request.version(),
+                request.serialize()).request;
+        assertEquals(request.isolationLevel(), deserialized.isolationLevel());
+
+        request = createFetchRequest(4, IsolationLevel.READ_UNCOMMITTED);
+        deserialized = (FetchRequest) AbstractRequest.parseRequest(request.apiKey(), request.version(),
+                request.serialize()).request;
+        assertEquals(request.isolationLevel(), deserialized.isolationLevel());
+    }
+
+    @Test
+    public void testFetchRequestWithMetadata() throws Exception {
+        FetchRequest request = createFetchRequest(4, IsolationLevel.READ_COMMITTED);
+        FetchRequest deserialized = (FetchRequest) AbstractRequest.parseRequest(ApiKeys.FETCH, request.version(),
+                request.serialize()).request;
         assertEquals(request.isolationLevel(), deserialized.isolationLevel());
 
         request = createFetchRequest(4, IsolationLevel.READ_UNCOMMITTED);
-        struct = request.toStruct();
-        deserialized = (FetchRequest) deserialize(request, struct, request.version());
+        deserialized = (FetchRequest) AbstractRequest.parseRequest(ApiKeys.FETCH, request.version(),
+                request.serialize()).request;
         assertEquals(request.isolationLevel(), deserialized.isolationLevel());
     }
 
     @Test
-    public void testJoinGroupRequestVersion0RebalanceTimeout() throws Exception {
+    public void testFetchRequestCompat() {
+        Map fetchData = new HashMap<>();
+        fetchData.put(new TopicPartition("test", 0), new FetchRequest.PartitionData(100, 2, 100, Optional.of(42)));
+        FetchRequest req = FetchRequest.Builder
+                .forConsumer(100, 100, fetchData)
+                .metadata(new FetchMetadata(10, 20))
+                .isolationLevel(IsolationLevel.READ_COMMITTED)
+                .build((short) 2);
+
+        FetchRequestData data = req.data();
+        ObjectSerializationCache cache = new ObjectSerializationCache();
+        int size = data.size(cache, (short) 2);
+
+        ByteBufferAccessor writer = new ByteBufferAccessor(ByteBuffer.allocate(size));
+        data.write(writer, cache, (short) 2);
+    }
+
+    @Test
+    public void testJoinGroupRequestVersion0RebalanceTimeout() {
         final short version = 0;
         JoinGroupRequest jgr = createJoinGroupRequest(version);
-        JoinGroupRequest jgr2 = new JoinGroupRequest(jgr.toStruct(), version);
-        assertEquals(jgr2.rebalanceTimeout(), jgr.rebalanceTimeout());
+        JoinGroupRequest jgr2 = JoinGroupRequest.parse(jgr.serialize(), version);
+        assertEquals(jgr2.data().rebalanceTimeoutMs(), jgr.data().rebalanceTimeoutMs());
     }
 
     @Test
     public void testOffsetFetchRequestBuilderToString() {
-        String allTopicPartitionsString = OffsetFetchRequest.Builder.allTopicPartitions("someGroup").toString();
-        assertTrue(allTopicPartitionsString.contains(""));
-        String string = new OffsetFetchRequest.Builder("group1",
-                singletonList(new TopicPartition("test11", 1))).toString();
-        assertTrue(string.contains("test11"));
-        assertTrue(string.contains("group1"));
+        List stableFlags = Arrays.asList(true, false);
+        for (Boolean requireStable : stableFlags) {
+            String allTopicPartitionsString = new OffsetFetchRequest.Builder("someGroup", requireStable, null, false).toString();
+
+            assertTrue(allTopicPartitionsString.contains("groupId='someGroup', topics=null, requireStable="
+                                                             + requireStable.toString()));
+            String string = new OffsetFetchRequest.Builder("group1",
+                requireStable, Collections.singletonList(new TopicPartition("test11", 1)), false).toString();
+            assertTrue(string.contains("test11"));
+            assertTrue(string.contains("group1"));
+            assertTrue(string.contains("requireStable=" + requireStable.toString()));
+        }
+    }
+
+    @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(expected = UnsupportedVersionException.class)
+    public void testListGroupRequestV3FailsWithStates() {
+        ListGroupsRequestData data = new ListGroupsRequestData()
+                .setStatesFilter(asList(ConsumerGroupState.STABLE.name()));
+        new ListGroupsRequest.Builder(data).build((short) 3);
+    }
+
+    @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());
     }
 
-    private ResponseHeader createResponseHeader() {
-        return new ResponseHeader(10);
+    @Test
+    public void testApiVersionResponseParsingFallback() {
+        ByteBuffer buffer = ApiVersionsResponse.DEFAULT_API_VERSIONS_RESPONSE.serialize((short) 0);
+        ApiVersionsResponse response = ApiVersionsResponse.parse(buffer, ApiKeys.API_VERSIONS.latestVersion());
+
+        assertEquals(Errors.NONE.code(), response.data().errorCode());
+    }
+
+    @Test
+    public void testApiVersionResponseParsingFallbackException() {
+        short version = 0;
+        assertThrows(BufferUnderflowException.class, () -> ApiVersionsResponse.parse(ByteBuffer.allocate(0), version));
+    }
+
+    @Test
+    public void testApiVersionResponseParsing() {
+        ByteBuffer buffer = ApiVersionsResponse.DEFAULT_API_VERSIONS_RESPONSE.serialize(ApiKeys.API_VERSIONS.latestVersion());
+        ApiVersionsResponse response = ApiVersionsResponse.parse(buffer, ApiKeys.API_VERSIONS.latestVersion());
+
+        assertEquals(Errors.NONE.code(), response.data().errorCode());
+    }
+
+    @Test
+    public void testInitProducerIdRequestVersions() {
+        InitProducerIdRequest.Builder bld = new InitProducerIdRequest.Builder(
+            new InitProducerIdRequestData().setTransactionTimeoutMs(1000).
+                setTransactionalId("abracadabra").
+                setProducerId(123));
+        final UnsupportedVersionException exception = assertThrows(
+            UnsupportedVersionException.class, () -> bld.build((short) 2).serialize());
+        assertTrue(exception.getMessage().contains("Attempted to write a non-default producerId at version 2"));
+        bld.build((short) 3);
+    }
+
+    @Test
+    public void testDeletableTopicResultErrorMessageIsNullByDefault() {
+        DeletableTopicResult result = new DeletableTopicResult()
+            .setName("topic")
+            .setErrorCode(Errors.THROTTLING_QUOTA_EXCEEDED.code());
+
+        assertEquals("topic", result.name());
+        assertEquals(Errors.THROTTLING_QUOTA_EXCEEDED.code(), result.errorCode());
+        assertNull(result.errorMessage());
+    }
+
+    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) {
+        LinkedHashMap fetchData = new LinkedHashMap<>();
+        fetchData.put(new TopicPartition("test1", 0), new FetchRequest.PartitionData(100, -1L,
+                1000000, Optional.empty()));
+        fetchData.put(new TopicPartition("test2", 0), new FetchRequest.PartitionData(200, -1L,
+                1000000, Optional.empty()));
+        return FetchRequest.Builder.forConsumer(100, 100000, fetchData).
+            metadata(metadata).setMaxBytes(1000).toForget(toForget).build((short) version);
     }
 
     private FetchRequest createFetchRequest(int version, IsolationLevel isolationLevel) {
         LinkedHashMap fetchData = new LinkedHashMap<>();
-        fetchData.put(new TopicPartition("test1", 0), new FetchRequest.PartitionData(100, 0L, 1000000));
-        fetchData.put(new TopicPartition("test2", 0), new FetchRequest.PartitionData(200, 0L, 1000000));
-        return FetchRequest.Builder.forConsumer(100, 100000, fetchData, isolationLevel).setMaxBytes(1000).build((short) version);
+        fetchData.put(new TopicPartition("test1", 0), new FetchRequest.PartitionData(100, -1L,
+                1000000, Optional.empty()));
+        fetchData.put(new TopicPartition("test2", 0), new FetchRequest.PartitionData(200, -1L,
+                1000000, Optional.empty()));
+        return FetchRequest.Builder.forConsumer(100, 100000, fetchData).
+            isolationLevel(isolationLevel).setMaxBytes(1000).build((short) version);
     }
 
     private FetchRequest createFetchRequest(int version) {
         LinkedHashMap fetchData = new LinkedHashMap<>();
-        fetchData.put(new TopicPartition("test1", 0), new FetchRequest.PartitionData(100, 0L, 1000000));
-        fetchData.put(new TopicPartition("test2", 0), new FetchRequest.PartitionData(200, 0L, 1000000));
+        fetchData.put(new TopicPartition("test1", 0), new FetchRequest.PartitionData(100, -1L,
+                1000000, Optional.empty()));
+        fetchData.put(new TopicPartition("test2", 0), new FetchRequest.PartitionData(200, -1L,
+                1000000, Optional.empty()));
         return FetchRequest.Builder.forConsumer(100, 100000, fetchData).setMaxBytes(1000).build((short) version);
     }
 
-    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));
+    private FetchResponse createFetchResponse(Errors error, int sessionId) {
+        return new FetchResponse<>(error, new LinkedHashMap<>(), 25, sessionId);
+    }
 
+    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, Optional.empty(), Collections.emptyList(), 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));
+            new FetchResponse.AbortedTransaction(234L, 999L));
+        responseData.put(new TopicPartition("test", 1), new FetchResponse.PartitionData<>(Errors.NONE,
+            1000000, FetchResponse.INVALID_LAST_STABLE_OFFSET, 0L, Optional.empty(), abortedTransactions, MemoryRecords.EMPTY));
+        return new FetchResponse<>(Errors.NONE, responseData, 25, sessionId);
+    }
+
+    private FetchResponse createFetchResponse(boolean includeAborted) {
+        LinkedHashMap> responseData = new LinkedHashMap<>();
+        MemoryRecords records = MemoryRecords.withRecords(CompressionType.NONE, new SimpleRecord("blah".getBytes()));
 
-        return new FetchResponse(responseData, 25);
+        responseData.put(new TopicPartition("test", 0), new FetchResponse.PartitionData<>(Errors.NONE,
+                1000000, FetchResponse.INVALID_LAST_STABLE_OFFSET, 0L, Optional.empty(), Collections.emptyList(), records));
+
+        List abortedTransactions = Collections.emptyList();
+        if (includeAborted) {
+            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, 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);
-        } else {
-            return new JoinGroupRequest.Builder("group1", 10000, "consumer1", "consumer", protocols).
-                    setRebalanceTimeout(60000).build();
+        JoinGroupRequestData.JoinGroupRequestProtocolCollection protocols =
+            new JoinGroupRequestData.JoinGroupRequestProtocolCollection(
+                Collections.singleton(
+                new JoinGroupRequestData.JoinGroupRequestProtocol()
+                        .setName("consumer-range")
+                        .setMetadata(new byte[0])).iterator()
+        );
+
+        JoinGroupRequestData data = new JoinGroupRequestData()
+            .setGroupId("group1")
+            .setSessionTimeoutMs(30000)
+            .setMemberId("consumer1")
+            .setProtocolType("consumer")
+            .setProtocols(protocols);
+
+        // v1 and above contains rebalance timeout
+        if (version >= 1)
+            data.setRebalanceTimeoutMs(60000);
+
+        // v5 and above could set group instance id
+        if (version >= 5)
+            data.setGroupInstanceId("groupInstanceId");
+
+        return new JoinGroupRequest.Builder(data).build((short) version);
+    }
+
+    private JoinGroupResponse createJoinGroupResponse(int version) {
+        List members = new ArrayList<>();
+
+        for (int i = 0; i < 2; i++) {
+            JoinGroupResponseMember member = new JoinGroupResponseData.JoinGroupResponseMember()
+                .setMemberId("consumer" + i)
+                .setMetadata(new byte[0]);
+
+            if (version >= 5)
+                member.setGroupInstanceId("instance" + i);
+
+            members.add(member);
         }
+
+        JoinGroupResponseData data = new JoinGroupResponseData()
+            .setErrorCode(Errors.NONE.code())
+            .setGenerationId(1)
+            .setProtocolType("consumer") // Added in v7 but ignorable
+            .setProtocolName("range")
+            .setLeader("leader")
+            .setMemberId("consumer1")
+            .setMembers(members);
+
+        // v1 and above could set throttle time
+        if (version >= 1)
+            data.setThrottleTimeMs(1000);
+
+        return new JoinGroupResponse(data);
     }
 
-    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);
+    private SyncGroupRequest createSyncGroupRequest(int version) {
+        List assignments = Collections.singletonList(
+            new SyncGroupRequestAssignment()
+                .setMemberId("member")
+                .setAssignment(new byte[0])
+        );
+
+        SyncGroupRequestData data = new SyncGroupRequestData()
+            .setGroupId("group1")
+            .setGenerationId(1)
+            .setMemberId("member")
+            .setProtocolType("consumer") // Added in v5 but ignorable
+            .setProtocolName("range")    // Added in v5 but ignorable
+            .setAssignments(assignments);
+
+        // v3 and above could set group instance id
+        if (version >= 3)
+            data.setGroupInstanceId("groupInstanceId");
+
+        return new SyncGroupRequest.Builder(data).build((short) version);
+    }
+
+    private SyncGroupResponse createSyncGroupResponse(int version) {
+        SyncGroupResponseData data = new SyncGroupResponseData()
+            .setErrorCode(Errors.NONE.code())
+            .setProtocolType("consumer") // Added in v5 but ignorable
+            .setProtocolName("range")    // Added in v5 but ignorable
+            .setAssignment(new byte[0]);
+
+        // v1 and above could set throttle time
+        if (version >= 1)
+            data.setThrottleTimeMs(1000);
+
+        return new SyncGroupResponse(data);
     }
 
-    private ListGroupsRequest createListGroupsRequest() {
-        return new ListGroupsRequest.Builder().build();
+    private ListGroupsRequest createListGroupsRequest(short version) {
+        ListGroupsRequestData data = new ListGroupsRequestData();
+        if (version >= 4)
+            data.setStatesFilter(Arrays.asList("Stable"));
+        return new ListGroupsRequest.Builder(data).build(version);
     }
 
-    private ListGroupsResponse createListGroupsResponse() {
-        List groups = Collections.singletonList(new ListGroupsResponse.Group("test-group", "consumer"));
-        return new ListGroupsResponse(Errors.NONE, groups);
+    private ListGroupsResponse createListGroupsResponse(int version) {
+        ListGroupsResponseData.ListedGroup group = new ListGroupsResponseData.ListedGroup()
+                .setGroupId("test-group")
+                .setProtocolType("consumer");
+        if (version >= 4)
+            group.setGroupState("Stable");
+        ListGroupsResponseData data = new ListGroupsResponseData()
+                .setErrorCode(Errors.NONE.code())
+                .setGroups(Collections.singletonList(group));
+        return new ListGroupsResponse(data);
     }
 
     private DescribeGroupsRequest createDescribeGroupRequest() {
-        return new DescribeGroupsRequest.Builder(singletonList("test-group")).build();
+        return new DescribeGroupsRequest.Builder(
+            new DescribeGroupsRequestData().
+                setGroups(Collections.singletonList("test-group"))).build();
     }
 
     private DescribeGroupsResponse createDescribeGroupResponse() {
         String clientId = "consumer-1";
         String clientHost = "localhost";
-        ByteBuffer empty = ByteBuffer.allocate(0);
-        DescribeGroupsResponse.GroupMember member = new DescribeGroupsResponse.GroupMember("memberId",
-                clientId, clientHost, empty, empty);
-        DescribeGroupsResponse.GroupMetadata metadata = new DescribeGroupsResponse.GroupMetadata(Errors.NONE,
-                "STABLE", "consumer", "roundrobin", asList(member));
-        return new DescribeGroupsResponse(Collections.singletonMap("test-group", metadata));
+        DescribeGroupsResponseData describeGroupsResponseData = new DescribeGroupsResponseData();
+        DescribeGroupsResponseData.DescribedGroupMember member = DescribeGroupsResponse.groupMember("memberId", null,
+                clientId, clientHost, new byte[0], new byte[0]);
+        DescribedGroup metadata = DescribeGroupsResponse.groupMetadata("test-group",
+                Errors.NONE,
+                "STABLE",
+                "consumer",
+                "roundrobin",
+                Collections.singletonList(member),
+                DescribeGroupsResponse.AUTHORIZED_OPERATIONS_OMITTED);
+        describeGroupsResponseData.groups().add(metadata);
+        return new DescribeGroupsResponse(describeGroupsResponseData);
     }
 
     private LeaveGroupRequest createLeaveGroupRequest() {
-        return new LeaveGroupRequest.Builder("group1", "consumer1").build();
+        return new LeaveGroupRequest.Builder(
+            "group1", Collections.singletonList(new MemberIdentity()
+                                                    .setMemberId("consumer1"))
+            ).build();
     }
 
     private LeaveGroupResponse createLeaveGroupResponse() {
-        return new LeaveGroupResponse(Errors.NONE);
+        return new LeaveGroupResponse(new LeaveGroupResponseData().setErrorCode(Errors.NONE.code()));
+    }
+
+    private DeleteGroupsRequest createDeleteGroupsRequest() {
+        return new DeleteGroupsRequest.Builder(
+            new DeleteGroupsRequestData()
+                .setGroupsNames(Collections.singletonList("test-group"))
+        ).build();
+    }
+
+    private DeleteGroupsResponse createDeleteGroupsResponse() {
+        DeletableGroupResultCollection result = new DeletableGroupResultCollection();
+        result.add(new DeletableGroupResult()
+                       .setGroupId("test-group")
+                       .setErrorCode(Errors.NONE.code()));
+        return new DeleteGroupsResponse(
+            new DeleteGroupsResponseData()
+                .setResults(result)
+        );
     }
 
-    @SuppressWarnings("deprecation")
     private ListOffsetRequest createListOffsetRequest(int version) {
         if (version == 0) {
-            Map offsetData = Collections.singletonMap(
-                    new TopicPartition("test", 0),
-                    new ListOffsetRequest.PartitionData(1000000L, 10));
+            ListOffsetTopic topic = new ListOffsetTopic()
+                    .setName("test")
+                    .setPartitions(Arrays.asList(new ListOffsetPartition()
+                            .setPartitionIndex(0)
+                            .setTimestamp(1000000L)
+                            .setMaxNumOffsets(10)
+                            .setCurrentLeaderEpoch(5)));
             return ListOffsetRequest.Builder
                     .forConsumer(false, IsolationLevel.READ_UNCOMMITTED)
-                    .setOffsetData(offsetData)
+                    .setTargetTimes(Collections.singletonList(topic))
                     .build((short) version);
         } else if (version == 1) {
-            Map offsetData = Collections.singletonMap(
-                    new TopicPartition("test", 0), 1000000L);
+            ListOffsetTopic topic = new ListOffsetTopic()
+                    .setName("test")
+                    .setPartitions(Arrays.asList(new ListOffsetPartition()
+                            .setPartitionIndex(0)
+                            .setTimestamp(1000000L)
+                            .setCurrentLeaderEpoch(5)));
             return ListOffsetRequest.Builder
                     .forConsumer(true, IsolationLevel.READ_UNCOMMITTED)
-                    .setTargetTimes(offsetData)
+                    .setTargetTimes(Collections.singletonList(topic))
                     .build((short) version);
-        } else if (version == 2) {
-            Map offsetData = Collections.singletonMap(
-                    new TopicPartition("test", 0), 1000000L);
+        } else if (version >= 2 && version <= LIST_OFFSETS.latestVersion()) {
+            ListOffsetPartition partition = new ListOffsetPartition()
+                    .setPartitionIndex(0)
+                    .setTimestamp(1000000L)
+                    .setCurrentLeaderEpoch(5);
+
+            ListOffsetTopic topic = new ListOffsetTopic()
+                    .setName("test")
+                    .setPartitions(Arrays.asList(partition));
             return ListOffsetRequest.Builder
                     .forConsumer(true, IsolationLevel.READ_COMMITTED)
-                    .setTargetTimes(offsetData)
+                    .setTargetTimes(Collections.singletonList(topic))
                     .build((short) version);
         } else {
             throw new IllegalArgumentException("Illegal ListOffsetRequest version " + version);
         }
     }
 
-    @SuppressWarnings("deprecation")
     private ListOffsetResponse createListOffsetResponse(int version) {
         if (version == 0) {
-            Map responseData = new HashMap<>();
-            responseData.put(new TopicPartition("test", 0),
-                    new ListOffsetResponse.PartitionData(Errors.NONE, asList(100L)));
-            return new ListOffsetResponse(responseData);
-        } else if (version == 1 || version == 2) {
-            Map responseData = new HashMap<>();
-            responseData.put(new TopicPartition("test", 0),
-                    new ListOffsetResponse.PartitionData(Errors.NONE, 10000L, 100L));
-            return new ListOffsetResponse(responseData);
+            ListOffsetResponseData data = new ListOffsetResponseData()
+                    .setTopics(Collections.singletonList(new ListOffsetTopicResponse()
+                            .setName("test")
+                            .setPartitions(Collections.singletonList(new ListOffsetPartitionResponse()
+                                    .setPartitionIndex(0)
+                                    .setErrorCode(Errors.NONE.code())
+                                    .setOldStyleOffsets(asList(100L))))));
+            return new ListOffsetResponse(data);
+        } else if (version >= 1 && version <= LIST_OFFSETS.latestVersion()) {
+            ListOffsetPartitionResponse partition = new ListOffsetPartitionResponse()
+                    .setPartitionIndex(0)
+                    .setErrorCode(Errors.NONE.code())
+                    .setTimestamp(10000L)
+                    .setOffset(100L);
+            if (version >= 4) {
+                partition.setLeaderEpoch(27);
+            }
+            ListOffsetResponseData data = new ListOffsetResponseData()
+                    .setTopics(Collections.singletonList(new ListOffsetTopicResponse()
+                            .setName("test")
+                            .setPartitions(Collections.singletonList(partition))));
+            return new ListOffsetResponse(data);
         } else {
             throw new IllegalArgumentException("Illegal ListOffsetResponse version " + version);
         }
@@ -672,59 +1332,103 @@ private MetadataRequest createMetadataRequest(int version, List topics)
 
     private MetadataResponse createMetadataResponse() {
         Node node = new Node(1, "host1", 1001);
-        List replicas = asList(node);
-        List isr = asList(node);
-        List offlineReplicas = asList();
+        List replicas = singletonList(node.id());
+        List isr = singletonList(node.id());
+        List offlineReplicas = emptyList();
 
         List allTopicMetadata = new ArrayList<>();
         allTopicMetadata.add(new MetadataResponse.TopicMetadata(Errors.NONE, "__consumer_offsets", true,
-                asList(new MetadataResponse.PartitionMetadata(Errors.NONE, 1, node, replicas, isr, offlineReplicas))));
+                asList(new MetadataResponse.PartitionMetadata(Errors.NONE,
+                        new TopicPartition("__consumer_offsets", 1),
+                        Optional.of(node.id()), Optional.of(5), replicas, isr, offlineReplicas))));
         allTopicMetadata.add(new MetadataResponse.TopicMetadata(Errors.LEADER_NOT_AVAILABLE, "topic2", false,
-                Collections.emptyList()));
+                emptyList()));
+        allTopicMetadata.add(new MetadataResponse.TopicMetadata(Errors.NONE, "topic3", false,
+            asList(new MetadataResponse.PartitionMetadata(Errors.LEADER_NOT_AVAILABLE,
+                    new TopicPartition("topic3", 0), Optional.empty(),
+                    Optional.empty(), replicas, isr, offlineReplicas))));
 
-        return new MetadataResponse(asList(node), null, MetadataResponse.NO_CONTROLLER_ID, allTopicMetadata);
+        return RequestTestUtils.metadataResponse(asList(node), null, MetadataResponse.NO_CONTROLLER_ID, allTopicMetadata);
     }
 
     private OffsetCommitRequest createOffsetCommitRequest(int version) {
-        Map commitData = new HashMap<>();
-        commitData.put(new TopicPartition("test", 0), new OffsetCommitRequest.PartitionData(100, ""));
-        commitData.put(new TopicPartition("test", 1), new OffsetCommitRequest.PartitionData(200, null));
-        return new OffsetCommitRequest.Builder("group1", commitData)
-                .setGenerationId(100)
+        return new OffsetCommitRequest.Builder(new OffsetCommitRequestData()
+                .setGroupId("group1")
                 .setMemberId("consumer1")
-                .setRetentionTime(1000000)
-                .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)))
+    private OffsetFetchRequest createOffsetFetchRequest(int version, boolean requireStable) {
+        return new OffsetFetchRequest.Builder("group1", requireStable, Collections.singletonList(new TopicPartition("test11", 1)), false)
                 .build((short) version);
     }
 
+    private OffsetFetchRequest createOffsetFetchRequestForAllPartition(String groupId, boolean requireStable) {
+        return new OffsetFetchRequest.Builder(groupId, requireStable, null, false).build();
+    }
+
     private OffsetFetchResponse createOffsetFetchResponse() {
         Map responseData = new HashMap<>();
-        responseData.put(new TopicPartition("test", 0), new OffsetFetchResponse.PartitionData(100L, "", Errors.NONE));
-        responseData.put(new TopicPartition("test", 1), new OffsetFetchResponse.PartitionData(100L, null, Errors.NONE));
+        responseData.put(new TopicPartition("test", 0), new OffsetFetchResponse.PartitionData(
+                100L, Optional.empty(), "", Errors.NONE));
+        responseData.put(new TopicPartition("test", 1), new OffsetFetchResponse.PartitionData(
+                100L, Optional.of(10), null, Errors.NONE));
         return new OffsetFetchResponse(Errors.NONE, responseData);
     }
 
+    @SuppressWarnings("deprecation")
     private ProduceRequest createProduceRequest(int version) {
         if (version < 2)
             throw new IllegalArgumentException("Produce request version 2 is not supported");
-
         byte magic = version == 2 ? RecordBatch.MAGIC_VALUE_V1 : RecordBatch.MAGIC_VALUE_V2;
         MemoryRecords records = MemoryRecords.withRecords(magic, CompressionType.NONE, new SimpleRecord("woot".getBytes()));
-        Map produceData = Collections.singletonMap(new TopicPartition("test", 0), records);
-        return ProduceRequest.Builder.forMagic(magic, (short) 1, 5000, produceData, "transactionalId")
+        return ProduceRequest.forMagic(magic,
+                new ProduceRequestData()
+                        .setTopicData(new ProduceRequestData.TopicProduceDataCollection(Collections.singletonList(
+                                new ProduceRequestData.TopicProduceData()
+                                        .setName("test")
+                                        .setPartitionData(Collections.singletonList(new ProduceRequestData.PartitionProduceData()
+                                                .setIndex(0)
+                                                .setRecords(records)))).iterator()))
+                        .setAcks((short) 1)
+                        .setTimeoutMs(5000)
+                        .setTransactionalId(version >= 3 ? "transactionalId" : null))
                 .build((short) version);
     }
 
+    @SuppressWarnings("deprecation")
     private ProduceResponse createProduceResponse() {
         Map responseData = new HashMap<>();
         responseData.put(new TopicPartition("test", 0), new ProduceResponse.PartitionResponse(Errors.NONE,
@@ -732,103 +1436,239 @@ private ProduceResponse createProduceResponse() {
         return new ProduceResponse(responseData, 0);
     }
 
-    private StopReplicaRequest createStopReplicaRequest(boolean deletePartitions) {
-        Set partitions = Utils.mkSet(new TopicPartition("test", 0));
-        return new StopReplicaRequest.Builder(0, 1, deletePartitions, partitions).build();
+    @SuppressWarnings("deprecation")
+    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) {
+        List topicStates = new ArrayList<>();
+        StopReplicaTopicState topic1 = new StopReplicaTopicState()
+            .setTopicName("topic1")
+            .setPartitionStates(Collections.singletonList(new StopReplicaPartitionState()
+                .setPartitionIndex(0)
+                .setLeaderEpoch(1)
+                .setDeletePartition(deletePartitions)));
+        topicStates.add(topic1);
+        StopReplicaTopicState topic2 = new StopReplicaTopicState()
+            .setTopicName("topic2")
+            .setPartitionStates(Collections.singletonList(new StopReplicaPartitionState()
+                .setPartitionIndex(1)
+                .setLeaderEpoch(2)
+                .setDeletePartition(deletePartitions)));
+        topicStates.add(topic2);
+
+        return new StopReplicaRequest.Builder((short) version, 0, 1, 0,
+            deletePartitions, topicStates).build((short) version);
     }
 
     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, 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, 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);
-    }
-
-    private LeaderAndIsrRequest createLeaderAndIsrRequest() {
-        Map partitionStates = new HashMap<>();
+        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) {
+        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(ApiKeys.LEADER_AND_ISR.latestVersion(), 1, 10, partitionStates, leaders).build();
+        return new LeaderAndIsrRequest.Builder((short) version, 1, 10, 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, partitionStates,
-                liveBrokers).build();
+        return new UpdateMetadataRequest.Builder((short) version, 1, 10, 0, partitionStates,
+            liveBrokers).build();
     }
 
     private UpdateMetadataResponse createUpdateMetadataResponse() {
-        return new UpdateMetadataResponse(Errors.NONE);
+        return new UpdateMetadataResponse(new UpdateMetadataResponseData().setErrorCode(Errors.NONE.code()));
     }
 
     private SaslHandshakeRequest createSaslHandshakeRequest() {
-        return new SaslHandshakeRequest("PLAIN");
+        return new SaslHandshakeRequest.Builder(
+                new SaslHandshakeRequestData().setMechanism("PLAIN")).build();
     }
 
     private SaslHandshakeResponse createSaslHandshakeResponse() {
-        return new SaslHandshakeResponse(Errors.NONE, singletonList("GSSAPI"));
+        return new SaslHandshakeResponse(
+                new SaslHandshakeResponseData()
+                .setErrorCode(Errors.NONE.code()).setMechanisms(Collections.singletonList("GSSAPI")));
+    }
+
+    private SaslAuthenticateRequest createSaslAuthenticateRequest() {
+        SaslAuthenticateRequestData data = new SaslAuthenticateRequestData().setAuthBytes(new byte[0]);
+        return new SaslAuthenticateRequest(data, ApiKeys.SASL_AUTHENTICATE.latestVersion());
+    }
+
+    private SaslAuthenticateResponse createSaslAuthenticateResponse() {
+        SaslAuthenticateResponseData data = new SaslAuthenticateResponseData()
+                .setErrorCode(Errors.NONE.code())
+                .setAuthBytes(new byte[0])
+                .setSessionLifetimeMs(Long.MAX_VALUE);
+        return new SaslAuthenticateResponse(data);
     }
 
     private ApiVersionsRequest createApiVersionRequest() {
@@ -836,8 +1676,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) {
@@ -845,68 +1693,159 @@ private CreateTopicsRequest createCreateTopicRequest(int version) {
     }
 
     private CreateTopicsRequest createCreateTopicRequest(int version, boolean validateOnly) {
-        CreateTopicsRequest.TopicDetails request1 = new CreateTopicsRequest.TopicDetails(3, (short) 5);
-
-        Map> replicaAssignments = new HashMap<>();
-        replicaAssignments.put(1, asList(1, 2, 3));
-        replicaAssignments.put(2, asList(2, 3, 4));
-
-        Map configs = new HashMap<>();
-        configs.put("config1", "value1");
-
-        CreateTopicsRequest.TopicDetails request2 = new CreateTopicsRequest.TopicDetails(replicaAssignments, configs);
-
-        Map request = new HashMap<>();
-        request.put("my_t1", request1);
-        request.put("my_t2", request2);
-        return new CreateTopicsRequest.Builder(request, 0, validateOnly).build((short) version);
+        CreateTopicsRequestData data = new CreateTopicsRequestData()
+            .setTimeoutMs(123)
+            .setValidateOnly(validateOnly);
+        data.topics().add(new CreatableTopic()
+            .setNumPartitions(3)
+            .setReplicationFactor((short) 5));
+
+        CreatableTopic topic2 = new CreatableTopic();
+        data.topics().add(topic2);
+        topic2.assignments().add(new CreatableReplicaAssignment()
+            .setPartitionIndex(0)
+            .setBrokerIds(Arrays.asList(1, 2, 3)));
+        topic2.assignments().add(new CreatableReplicaAssignment()
+            .setPartitionIndex(1)
+            .setBrokerIds(Arrays.asList(2, 3, 4)));
+        topic2.configs().add(new CreateableTopicConfig()
+            .setName("config1").setValue("value1"));
+
+        return new CreateTopicsRequest.Builder(data).build((short) version);
     }
 
     private CreateTopicsResponse createCreateTopicResponse() {
-        Map errors = new HashMap<>();
-        errors.put("t1", new ApiError(Errors.INVALID_TOPIC_EXCEPTION, null));
-        errors.put("t2", new ApiError(Errors.LEADER_NOT_AVAILABLE, "Leader with id 5 is not available."));
-        return new CreateTopicsResponse(errors);
-    }
-
-    private DeleteTopicsRequest createDeleteTopicsRequest() {
-        return new DeleteTopicsRequest.Builder(Utils.mkSet("my_t1", "my_t2"), 10000).build();
+        CreateTopicsResponseData data = new CreateTopicsResponseData();
+        data.topics().add(new CreatableTopicResult()
+            .setName("t1")
+            .setErrorCode(Errors.INVALID_TOPIC_EXCEPTION.code())
+            .setErrorMessage(null));
+        data.topics().add(new CreatableTopicResult()
+            .setName("t2")
+            .setErrorCode(Errors.LEADER_NOT_AVAILABLE.code())
+            .setErrorMessage("Leader with id 5 is not available."));
+        data.topics().add(new CreatableTopicResult()
+            .setName("t3")
+            .setErrorCode(Errors.NONE.code())
+            .setNumPartitions(1)
+            .setReplicationFactor((short) 2)
+            .setConfigs(Collections.singletonList(new CreatableTopicConfigs()
+                .setName("min.insync.replicas")
+                .setValue("2"))));
+        return new CreateTopicsResponse(data);
+    }
+
+    private DeleteTopicsRequest createDeleteTopicsRequest(int version) {
+        return new DeleteTopicsRequest.Builder(new DeleteTopicsRequestData()
+            .setTopicNames(Arrays.asList("my_t1", "my_t2"))
+            .setTimeoutMs(1000)
+        ).build((short) version);
     }
 
     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())
+            .setErrorMessage("Error Message"));
+        data.responses().add(new DeletableTopicResult()
+            .setName("t2")
+            .setErrorCode(Errors.TOPIC_AUTHORIZATION_FAILED.code())
+            .setErrorMessage("Error Message"));
+        data.responses().add(new DeletableTopicResult()
+            .setName("t3")
+            .setErrorCode(Errors.NOT_CONTROLLER.code()));
+        data.responses().add(new DeletableTopicResult()
+                .setName("t4")
+                .setErrorCode(Errors.NONE.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);
-    }
-
-
-    private OffsetsForLeaderEpochRequest createLeaderEpochRequest() {
-        Map epochs = new HashMap<>();
-
-        epochs.put(new TopicPartition("topic1", 0), 1);
-        epochs.put(new TopicPartition("topic1", 1), 1);
-        epochs.put(new TopicPartition("topic2", 2), 3);
-
-        return new OffsetsForLeaderEpochRequest.Builder(epochs).build();
+        InitProducerIdResponseData responseData = new InitProducerIdResponseData()
+                .setErrorCode(Errors.NONE.code())
+                .setProducerEpoch((short) 3)
+                .setProducerId(3332)
+                .setThrottleTimeMs(0);
+        return new InitProducerIdResponse(responseData);
+    }
+
+    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 OffsetForLeaderTopicCollection createOffsetForLeaderTopicCollection() {
+        OffsetForLeaderTopicCollection topics = new OffsetForLeaderTopicCollection();
+        topics.add(new OffsetForLeaderTopic()
+            .setTopic("topic1")
+            .setPartitions(Arrays.asList(
+                new OffsetForLeaderPartition()
+                    .setPartition(0)
+                    .setLeaderEpoch(1)
+                    .setCurrentLeaderEpoch(0),
+                new OffsetForLeaderPartition()
+                    .setPartition(1)
+                    .setLeaderEpoch(1)
+                    .setCurrentLeaderEpoch(0))));
+        topics.add(new OffsetForLeaderTopic()
+            .setTopic("topic2")
+            .setPartitions(Arrays.asList(
+                new OffsetForLeaderPartition()
+                    .setPartition(2)
+                    .setLeaderEpoch(3)
+                    .setCurrentLeaderEpoch(RecordBatch.NO_PARTITION_LEADER_EPOCH))));
+        return topics;
+    }
+
+    private OffsetsForLeaderEpochRequest createLeaderEpochRequestForConsumer() {
+        OffsetForLeaderTopicCollection epochs = createOffsetForLeaderTopicCollection();
+        return OffsetsForLeaderEpochRequest.Builder.forConsumer(epochs).build();
+    }
+
+    private OffsetsForLeaderEpochRequest createLeaderEpochRequestForReplica(int version, int replicaId) {
+        Map epochs = createOffsetForLeaderEpochPartitionData();
+        return OffsetsForLeaderEpochRequest.Builder.forFollower((short) version, epochs, replicaId).build();
     }
 
     private OffsetsForLeaderEpochResponse createLeaderEpochResponse() {
-        Map epochs = new HashMap<>();
-
-        epochs.put(new TopicPartition("topic1", 0), new EpochEndOffset(Errors.NONE, 0));
-        epochs.put(new TopicPartition("topic1", 1), new EpochEndOffset(Errors.NONE, 1));
-        epochs.put(new TopicPartition("topic2", 2), new EpochEndOffset(Errors.NONE, 2));
-
-        return new OffsetsForLeaderEpochResponse(epochs);
+        OffsetForLeaderEpochResponseData data = new OffsetForLeaderEpochResponseData();
+        data.topics().add(new OffsetForLeaderTopicResult()
+            .setTopic("topic1")
+            .setPartitions(Arrays.asList(
+                new EpochEndOffset()
+                    .setPartition(0)
+                    .setErrorCode(Errors.NONE.code())
+                    .setLeaderEpoch(1)
+                    .setEndOffset(0),
+                new EpochEndOffset()
+                    .setPartition(1)
+                    .setErrorCode(Errors.NONE.code())
+                    .setLeaderEpoch(1)
+                    .setEndOffset(1))));
+        data.topics().add(new OffsetForLeaderTopicResult()
+            .setTopic("topic2")
+            .setPartitions(Arrays.asList(
+                new EpochEndOffset()
+                    .setPartition(2)
+                    .setErrorCode(Errors.NONE.code())
+                    .setLeaderEpoch(1)
+                    .setEndOffset(1))));
+
+        return new OffsetsForLeaderEpochResponse(data);
     }
 
     private AddPartitionsToTxnRequest createAddPartitionsToTxnRequest() {
@@ -919,25 +1858,43 @@ private AddPartitionsToTxnResponse createAddPartitionsToTxnResponse() {
     }
 
     private AddOffsetsToTxnRequest createAddOffsetsToTxnRequest() {
-        return new AddOffsetsToTxnRequest.Builder("tid", 21L, (short) 42, "gid").build();
+        return new AddOffsetsToTxnRequest.Builder(
+            new AddOffsetsToTxnRequestData()
+                .setTransactionalId("tid")
+                .setProducerId(21L)
+                .setProducerEpoch((short) 42)
+                .setGroupId("gid")
+        ).build();
     }
 
     private AddOffsetsToTxnResponse createAddOffsetsToTxnResponse() {
-        return new AddOffsetsToTxnResponse(0, Errors.NONE);
+        return new AddOffsetsToTxnResponse(new AddOffsetsToTxnResponseData()
+                                               .setErrorCode(Errors.NONE.code())
+                                               .setThrottleTimeMs(0));
     }
 
     private EndTxnRequest createEndTxnRequest() {
-        return new EndTxnRequest.Builder("tid", 21L, (short) 42, TransactionResult.COMMIT).build();
+        return new EndTxnRequest.Builder(
+            new EndTxnRequestData()
+                .setTransactionalId("tid")
+                .setProducerId(21L)
+                .setProducerEpoch((short) 42)
+                .setCommitted(TransactionResult.COMMIT.id)
+            ).build();
     }
 
     private EndTxnResponse createEndTxnResponse() {
-        return new EndTxnResponse(0, Errors.NONE);
+        return new EndTxnResponse(
+            new EndTxnResponseData()
+                .setErrorCode(Errors.NONE.code())
+                .setThrottleTimeMs(0)
+        );
     }
 
     private WriteTxnMarkersRequest createWriteTxnMarkersRequest() {
-        return new WriteTxnMarkersRequest.Builder(
-            Collections.singletonList(new WriteTxnMarkersRequest.TxnMarkerEntry(21L, (short) 42, 73, TransactionResult.ABORT,
-                                                                                Collections.singletonList(new TopicPartition("topic", 73))))).build();
+        List partitions = Collections.singletonList(new TopicPartition("topic", 73));
+        WriteTxnMarkersRequest.TxnMarkerEntry txnMarkerEntry = new WriteTxnMarkersRequest.TxnMarkerEntry(21L, (short) 42, 73, TransactionResult.ABORT, partitions);
+        return new WriteTxnMarkersRequest.Builder(ApiKeys.WRITE_TXN_MARKERS.latestVersion(), Collections.singletonList(txnMarkerEntry)).build();
     }
 
     private WriteTxnMarkersResponse createWriteTxnMarkersResponse() {
@@ -948,11 +1905,49 @@ private WriteTxnMarkersResponse createWriteTxnMarkersResponse() {
         return new WriteTxnMarkersResponse(response);
     }
 
-    private TxnOffsetCommitRequest createTxnOffsetCommitRequest() {
+    private TxnOffsetCommitRequest createTxnOffsetCommitRequest(int version) {
         final Map offsets = new HashMap<>();
         offsets.put(new TopicPartition("topic", 73),
-                    new TxnOffsetCommitRequest.CommittedOffset(100, null));
-        return new TxnOffsetCommitRequest.Builder("transactionalId", "groupId", 21L, (short) 42, offsets).build();
+                    new TxnOffsetCommitRequest.CommittedOffset(100, null, Optional.empty()));
+        offsets.put(new TopicPartition("topic", 74),
+                new TxnOffsetCommitRequest.CommittedOffset(100, "blah", Optional.of(27)));
+
+        if (version < 3) {
+            return new TxnOffsetCommitRequest.Builder("transactionalId",
+                "groupId",
+                21L,
+                (short) 42,
+                offsets,
+                false).build();
+        } else {
+            return new TxnOffsetCommitRequest.Builder("transactionalId",
+                "groupId",
+                21L,
+                (short) 42,
+                offsets,
+                "member",
+                2,
+                Optional.of("instance"),
+                false).build();
+        }
+    }
+
+    private TxnOffsetCommitRequest createTxnOffsetCommitRequestWithAutoDowngrade(int version) {
+        final Map offsets = new HashMap<>();
+        offsets.put(new TopicPartition("topic", 73),
+            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,
+            "member",
+            2,
+            Optional.of("instance"),
+            true).build();
     }
 
     private TxnOffsetCommitResponse createTxnOffsetCommitResponse() {
@@ -961,125 +1956,620 @@ private TxnOffsetCommitResponse createTxnOffsetCommitResponse() {
         return new TxnOffsetCommitResponse(0, errorPerPartitions);
     }
 
-    private DescribeAclsRequest createListAclsRequest() {
+    private DescribeAclsRequest createDescribeAclsRequest() {
         return new DescribeAclsRequest.Builder(new AclBindingFilter(
-                new ResourceFilter(ResourceType.TOPIC, "mytopic"),
+                new ResourcePatternFilter(ResourceType.TOPIC, "mytopic", PatternType.LITERAL),
                 new AccessControlEntryFilter(null, null, AclOperation.ANY, AclPermissionType.ANY))).build();
     }
 
     private DescribeAclsResponse createDescribeAclsResponse() {
-        return new DescribeAclsResponse(0, ApiError.NONE, Collections.singleton(new AclBinding(
-            new Resource(ResourceType.TOPIC, "mytopic"),
-            new AccessControlEntry("User:ANONYMOUS", "*", AclOperation.WRITE, AclPermissionType.ALLOW))));
+        DescribeAclsResponseData data = new DescribeAclsResponseData()
+                .setErrorCode(Errors.NONE.code())
+                .setErrorMessage(Errors.NONE.message())
+                .setThrottleTimeMs(0)
+                .setResources(Collections.singletonList(new DescribeAclsResource()
+                        .setResourceType(ResourceType.TOPIC.code())
+                        .setResourceName("mytopic")
+                        .setPatternType(PatternType.LITERAL.code())
+                        .setAcls(Collections.singletonList(new AclDescription()
+                                .setHost("*")
+                                .setOperation(AclOperation.WRITE.code())
+                                .setPermissionType(AclPermissionType.ALLOW.code())
+                                .setPrincipal("User:ANONYMOUS")))));
+        return new DescribeAclsResponse(data);
     }
 
     private CreateAclsRequest createCreateAclsRequest() {
-        List creations = new ArrayList<>();
-        creations.add(new AclCreation(new AclBinding(
-            new Resource(ResourceType.TOPIC, "mytopic"),
+        List creations = new ArrayList<>();
+        creations.add(CreateAclsRequest.aclCreation(new AclBinding(
+            new ResourcePattern(ResourceType.TOPIC, "mytopic", PatternType.LITERAL),
             new AccessControlEntry("User:ANONYMOUS", "127.0.0.1", AclOperation.READ, AclPermissionType.ALLOW))));
-        creations.add(new AclCreation(new AclBinding(
-            new Resource(ResourceType.GROUP, "mygroup"),
+        creations.add(CreateAclsRequest.aclCreation(new AclBinding(
+            new ResourcePattern(ResourceType.GROUP, "mygroup", PatternType.LITERAL),
             new AccessControlEntry("User:ANONYMOUS", "*", AclOperation.WRITE, AclPermissionType.DENY))));
-        return new CreateAclsRequest.Builder(creations).build();
+        CreateAclsRequestData data = new CreateAclsRequestData().setCreations(creations);
+        return new CreateAclsRequest.Builder(data).build();
     }
 
     private CreateAclsResponse createCreateAclsResponse() {
-        return new CreateAclsResponse(0, Arrays.asList(new AclCreationResponse(ApiError.NONE),
-            new AclCreationResponse(new ApiError(Errors.INVALID_REQUEST, "Foo bar"))));
+        return new CreateAclsResponse(new CreateAclsResponseData().setResults(asList(
+            new CreateAclsResponseData.AclCreationResult(),
+            new CreateAclsResponseData.AclCreationResult()
+                .setErrorCode(Errors.NONE.code())
+                .setErrorMessage("Foo bar"))));
     }
 
     private DeleteAclsRequest createDeleteAclsRequest() {
-        List filters = new ArrayList<>();
-        filters.add(new AclBindingFilter(
-            new ResourceFilter(ResourceType.ANY, null),
-            new AccessControlEntryFilter("User:ANONYMOUS", null, AclOperation.ANY, AclPermissionType.ANY)));
-        filters.add(new AclBindingFilter(
-            new ResourceFilter(ResourceType.ANY, null),
-            new AccessControlEntryFilter("User:bob", null, AclOperation.ANY, AclPermissionType.ANY)));
-        return new DeleteAclsRequest.Builder(filters).build();
-    }
-
-    private DeleteAclsResponse createDeleteAclsResponse() {
-        List responses = new ArrayList<>();
-        responses.add(new AclFilterResponse(Utils.mkSet(
-                new AclDeletionResult(new AclBinding(
-                        new Resource(ResourceType.TOPIC, "mytopic3"),
-                        new AccessControlEntry("User:ANONYMOUS", "*", AclOperation.DESCRIBE, AclPermissionType.ALLOW))),
-                new AclDeletionResult(new AclBinding(
-                        new Resource(ResourceType.TOPIC, "mytopic4"),
-                        new AccessControlEntry("User:ANONYMOUS", "*", AclOperation.DESCRIBE, AclPermissionType.DENY))))));
-        responses.add(new AclFilterResponse(new ApiError(Errors.SECURITY_DISABLED, "No security"),
-            Collections.emptySet()));
-        return new DeleteAclsResponse(0, responses);
-    }
-
-    private DescribeConfigsRequest createDescribeConfigsRequest() {
-        return new DescribeConfigsRequest.Builder(asList(
-                new org.apache.kafka.common.requests.Resource(org.apache.kafka.common.requests.ResourceType.BROKER, "0"),
-                new org.apache.kafka.common.requests.Resource(org.apache.kafka.common.requests.ResourceType.TOPIC, "topic"))).build((short) 0);
-    }
-
-    private DescribeConfigsRequest createDescribeConfigsRequestWithConfigEntries() {
-        Map> resources = new HashMap<>();
-        resources.put(new org.apache.kafka.common.requests.Resource(org.apache.kafka.common.requests.ResourceType.BROKER, "0"), asList("foo", "bar"));
-        resources.put(new org.apache.kafka.common.requests.Resource(org.apache.kafka.common.requests.ResourceType.TOPIC, "topic"), null);
-        resources.put(new org.apache.kafka.common.requests.Resource(org.apache.kafka.common.requests.ResourceType.TOPIC, "topic a"), Collections.emptyList());
-        return new DescribeConfigsRequest.Builder(resources).build((short) 0);
-    }
-
-    private DescribeConfigsResponse createDescribeConfigsResponse() {
-        Map configs = new HashMap<>();
-        List configEntries = asList(
-                new DescribeConfigsResponse.ConfigEntry("config_name", "config_value", false, true, false),
-                new DescribeConfigsResponse.ConfigEntry("another_name", "another value", true, false, true)
-        );
-        configs.put(new org.apache.kafka.common.requests.Resource(org.apache.kafka.common.requests.ResourceType.BROKER, "0"), new DescribeConfigsResponse.Config(
-                ApiError.NONE, configEntries));
-        configs.put(new org.apache.kafka.common.requests.Resource(org.apache.kafka.common.requests.ResourceType.TOPIC, "topic"), new DescribeConfigsResponse.Config(
-                ApiError.NONE, Collections.emptyList()));
-        return new DescribeConfigsResponse(200, configs);
+        DeleteAclsRequestData data = new DeleteAclsRequestData().setFilters(asList(
+            new DeleteAclsRequestData.DeleteAclsFilter()
+                .setResourceTypeFilter(ResourceType.ANY.code())
+                .setResourceNameFilter(null)
+                .setPatternTypeFilter(PatternType.LITERAL.code())
+                .setPrincipalFilter("User:ANONYMOUS")
+                .setHostFilter(null)
+                .setOperation(AclOperation.ANY.code())
+                .setPermissionType(AclPermissionType.ANY.code()),
+            new DeleteAclsRequestData.DeleteAclsFilter()
+                .setResourceTypeFilter(ResourceType.ANY.code())
+                .setResourceNameFilter(null)
+                .setPatternTypeFilter(PatternType.LITERAL.code())
+                .setPrincipalFilter("User:bob")
+                .setHostFilter(null)
+                .setOperation(AclOperation.ANY.code())
+                .setPermissionType(AclPermissionType.ANY.code())
+        ));
+        return new DeleteAclsRequest.Builder(data).build();
+    }
+
+    private DeleteAclsResponse createDeleteAclsResponse(int version) {
+        List filterResults = new ArrayList<>();
+        filterResults.add(new DeleteAclsResponseData.DeleteAclsFilterResult().setMatchingAcls(asList(
+                new DeleteAclsResponseData.DeleteAclsMatchingAcl()
+                    .setResourceType(ResourceType.TOPIC.code())
+                    .setResourceName("mytopic3")
+                    .setPatternType(PatternType.LITERAL.code())
+                    .setPrincipal("User:ANONYMOUS")
+                    .setHost("*")
+                    .setOperation(AclOperation.DESCRIBE.code())
+                    .setPermissionType(AclPermissionType.ALLOW.code()),
+                new DeleteAclsResponseData.DeleteAclsMatchingAcl()
+                    .setResourceType(ResourceType.TOPIC.code())
+                    .setResourceName("mytopic4")
+                    .setPatternType(PatternType.LITERAL.code())
+                    .setPrincipal("User:ANONYMOUS")
+                    .setHost("*")
+                    .setOperation(AclOperation.DESCRIBE.code())
+                    .setPermissionType(AclPermissionType.DENY.code()))));
+        filterResults.add(new DeleteAclsResponseData.DeleteAclsFilterResult()
+            .setErrorCode(Errors.SECURITY_DISABLED.code())
+            .setErrorMessage("No security"));
+        return new DeleteAclsResponse(new DeleteAclsResponseData()
+            .setThrottleTimeMs(0)
+            .setFilterResults(filterResults), (short) version);
+    }
+
+    private DescribeConfigsRequest createDescribeConfigsRequest(int version) {
+        return new DescribeConfigsRequest.Builder(new DescribeConfigsRequestData()
+                .setResources(asList(
+                        new DescribeConfigsRequestData.DescribeConfigsResource()
+                                .setResourceType(ConfigResource.Type.BROKER.id())
+                                .setResourceName("0"),
+                        new DescribeConfigsRequestData.DescribeConfigsResource()
+                                .setResourceType(ConfigResource.Type.TOPIC.id())
+                                .setResourceName("topic"))))
+                .build((short) version);
+    }
+
+    private DescribeConfigsRequest createDescribeConfigsRequestWithConfigEntries(int version) {
+        return new DescribeConfigsRequest.Builder(new DescribeConfigsRequestData()
+            .setResources(asList(
+                new DescribeConfigsRequestData.DescribeConfigsResource()
+                    .setResourceType(ConfigResource.Type.BROKER.id())
+                    .setResourceName("0")
+                    .setConfigurationKeys(asList("foo", "bar")),
+                new DescribeConfigsRequestData.DescribeConfigsResource()
+                    .setResourceType(ConfigResource.Type.TOPIC.id())
+                    .setResourceName("topic")
+                    .setConfigurationKeys(null),
+                new DescribeConfigsRequestData.DescribeConfigsResource()
+                    .setResourceType(ConfigResource.Type.TOPIC.id())
+                    .setResourceName("topic a")
+                    .setConfigurationKeys(emptyList())))).build((short) version);
+    }
+
+    private DescribeConfigsRequest createDescribeConfigsRequestWithDocumentation(int version) {
+        DescribeConfigsRequestData data = new DescribeConfigsRequestData()
+                .setResources(asList(
+                        new DescribeConfigsRequestData.DescribeConfigsResource()
+                                .setResourceType(ConfigResource.Type.BROKER.id())
+                                .setResourceName("0")
+                                .setConfigurationKeys(asList("foo", "bar"))));
+        if (version == 3) {
+            data.setIncludeDocumentation(true);
+        }
+        return new DescribeConfigsRequest.Builder(data).build((short) version);
+    }
+
+    private DescribeConfigsResponse createDescribeConfigsResponse(short version) {
+        return new DescribeConfigsResponse(new DescribeConfigsResponseData().setResults(asList(
+                new DescribeConfigsResult()
+                        .setErrorCode(Errors.NONE.code())
+                        .setResourceType(ConfigResource.Type.BROKER.id())
+                        .setResourceName("0")
+                        .setConfigs(asList(
+                                new DescribeConfigsResourceResult()
+                                        .setName("config_name")
+                                        .setValue("config_value")
+                                        // 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.
+                                        .setConfigSource(version == 0 ? DescribeConfigsResponse.ConfigSource.STATIC_BROKER_CONFIG.id() : DescribeConfigsResponse.ConfigSource.DYNAMIC_BROKER_CONFIG.id)
+                                        .setIsSensitive(true).setReadOnly(false)
+                                        .setSynonyms(emptyList()),
+                                new DescribeConfigsResourceResult()
+                                        .setName("yet_another_name")
+                                        .setValue("yet another value")
+                                        .setConfigSource(version == 0 ? DescribeConfigsResponse.ConfigSource.STATIC_BROKER_CONFIG.id() : DescribeConfigsResponse.ConfigSource.DEFAULT_CONFIG.id)
+                                        .setIsSensitive(false).setReadOnly(true)
+                                        .setSynonyms(emptyList())
+                                        .setConfigType(ConfigType.BOOLEAN.id())
+                                        .setDocumentation("some description"),
+                                new DescribeConfigsResourceResult()
+                                        .setName("another_name")
+                                        .setValue("another value")
+                                        .setConfigSource(version == 0 ? DescribeConfigsResponse.ConfigSource.STATIC_BROKER_CONFIG.id() : DescribeConfigsResponse.ConfigSource.DEFAULT_CONFIG.id)
+                                        .setIsSensitive(false).setReadOnly(true)
+                                        .setSynonyms(emptyList())
+                        )),
+                new DescribeConfigsResult()
+                        .setErrorCode(Errors.NONE.code())
+                        .setResourceType(ConfigResource.Type.TOPIC.id())
+                        .setResourceName("topic")
+                        .setConfigs(emptyList())
+        )));
+
     }
 
     private AlterConfigsRequest createAlterConfigsRequest() {
-        Map configs = new HashMap<>();
+        Map configs = new HashMap<>();
         List configEntries = asList(
                 new AlterConfigsRequest.ConfigEntry("config_name", "config_value"),
                 new AlterConfigsRequest.ConfigEntry("another_name", "another value")
         );
-        configs.put(new org.apache.kafka.common.requests.Resource(org.apache.kafka.common.requests.ResourceType.BROKER, "0"), new AlterConfigsRequest.Config(configEntries));
-        configs.put(new org.apache.kafka.common.requests.Resource(org.apache.kafka.common.requests.ResourceType.TOPIC, "topic"),
+        configs.put(new ConfigResource(ConfigResource.Type.BROKER, "0"), new AlterConfigsRequest.Config(configEntries));
+        configs.put(new ConfigResource(ConfigResource.Type.TOPIC, "topic"),
                 new AlterConfigsRequest.Config(Collections.emptyList()));
-        return new AlterConfigsRequest((short) 0, configs, false);
+        return new AlterConfigsRequest.Builder(configs, false).build((short) 0);
     }
 
     private AlterConfigsResponse createAlterConfigsResponse() {
-        Map errors = new HashMap<>();
-        errors.put(new org.apache.kafka.common.requests.Resource(org.apache.kafka.common.requests.ResourceType.BROKER, "0"), ApiError.NONE);
-        errors.put(new org.apache.kafka.common.requests.Resource(org.apache.kafka.common.requests.ResourceType.TOPIC, "topic"), new ApiError(Errors.INVALID_REQUEST, "This request is invalid"));
-        return new AlterConfigsResponse(20, errors);
-    }
+        AlterConfigsResponseData data = new AlterConfigsResponseData()
+                .setThrottleTimeMs(20);
+        data.responses().add(new AlterConfigsResponseData.AlterConfigsResourceResponse()
+                .setErrorCode(Errors.NONE.code())
+                .setErrorMessage(null)
+                .setResourceName("0")
+                .setResourceType(ConfigResource.Type.BROKER.id()));
+        data.responses().add(new AlterConfigsResponseData.AlterConfigsResourceResponse()
+                .setErrorCode(Errors.INVALID_REQUEST.code())
+                .setErrorMessage("This request is invalid")
+                .setResourceName("topic")
+                .setResourceType(ConfigResource.Type.TOPIC.id()));
+        return new AlterConfigsResponse(data);
+    }
+
+    private CreatePartitionsRequest createCreatePartitionsRequest(int version) {
+        CreatePartitionsTopicCollection topics = new CreatePartitionsTopicCollection();
+        topics.add(new CreatePartitionsTopic()
+                .setName("my_topic")
+                .setCount(3)
+        );
+        topics.add(new CreatePartitionsTopic()
+                .setName("my_other_topic")
+                .setCount(3)
+        );
+
+        CreatePartitionsRequestData data = new CreatePartitionsRequestData()
+                .setTimeoutMs(0)
+                .setValidateOnly(false)
+                .setTopics(topics);
 
-    private CreatePartitionsRequest createCreatePartitionsRequest() {
-        Map assignments = new HashMap<>();
-        assignments.put("my_topic", NewPartitions.increaseTo(3));
-        assignments.put("my_other_topic", NewPartitions.increaseTo(3));
-        return new CreatePartitionsRequest(assignments, 0, false, (short) 0);
+        return new CreatePartitionsRequest(data, (short) version);
     }
 
-    private CreatePartitionsRequest createCreatePartitionsRequestWithAssignments() {
-        Map assignments = new HashMap<>();
-        assignments.put("my_topic", NewPartitions.increaseTo(3, asList(asList(2))));
-        assignments.put("my_other_topic", NewPartitions.increaseTo(3, asList(asList(2, 3), asList(3, 1))));
-        return new CreatePartitionsRequest(assignments, 0, false, (short) 0);
+    private CreatePartitionsRequest createCreatePartitionsRequestWithAssignments(int version) {
+        CreatePartitionsTopicCollection topics = new CreatePartitionsTopicCollection();
+        CreatePartitionsAssignment myTopicAssignment = new CreatePartitionsAssignment()
+                .setBrokerIds(Collections.singletonList(2));
+        topics.add(new CreatePartitionsTopic()
+                .setName("my_topic")
+                .setCount(3)
+                .setAssignments(Collections.singletonList(myTopicAssignment))
+        );
+
+        topics.add(new CreatePartitionsTopic()
+                .setName("my_other_topic")
+                .setCount(3)
+                .setAssignments(asList(
+                    new CreatePartitionsAssignment().setBrokerIds(asList(2, 3)),
+                    new CreatePartitionsAssignment().setBrokerIds(asList(3, 1))
+                ))
+        );
+
+        CreatePartitionsRequestData data = new CreatePartitionsRequestData()
+                .setTimeoutMs(0)
+                .setValidateOnly(false)
+                .setTopics(topics);
+
+        return new CreatePartitionsRequest(data, (short) version);
     }
 
     private CreatePartitionsResponse createCreatePartitionsResponse() {
-        Map results = new HashMap<>();
-        results.put("my_topic", ApiError.fromThrowable(
-                new InvalidReplicaAssignmentException("The assigned brokers included an unknown broker")));
-        results.put("my_topic", ApiError.NONE);
-        return new CreatePartitionsResponse(42, results);
+        List results = new LinkedList<>();
+        results.add(new CreatePartitionsTopicResult()
+                .setName("my_topic")
+                .setErrorCode(Errors.INVALID_REPLICA_ASSIGNMENT.code()));
+        results.add(new CreatePartitionsTopicResult()
+                .setName("my_topic")
+                .setErrorCode(Errors.NONE.code()));
+        CreatePartitionsResponseData data = new CreatePartitionsResponseData()
+                .setThrottleTimeMs(42)
+                .setResults(results);
+        return new CreatePartitionsResponse(data);
+    }
+
+    private CreateDelegationTokenRequest createCreateTokenRequest() {
+        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() {
+        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() {
+        RenewDelegationTokenRequestData data = new RenewDelegationTokenRequestData()
+                .setHmac("test".getBytes())
+                .setRenewPeriodMs(System.currentTimeMillis());
+        return new RenewDelegationTokenRequest.Builder(data).build();
+    }
+
+    private RenewDelegationTokenResponse createRenewTokenResponse() {
+        RenewDelegationTokenResponseData data = new RenewDelegationTokenResponseData()
+                .setThrottleTimeMs(20)
+                .setErrorCode(Errors.NONE.code())
+                .setExpiryTimestampMs(System.currentTimeMillis());
+        return new RenewDelegationTokenResponse(data);
+    }
+
+    private ExpireDelegationTokenRequest createExpireTokenRequest() {
+        ExpireDelegationTokenRequestData data = new ExpireDelegationTokenRequestData()
+                .setHmac("test".getBytes())
+                .setExpiryTimePeriodMs(System.currentTimeMillis());
+        return new ExpireDelegationTokenRequest.Builder(data).build();
+    }
+
+    private ExpireDelegationTokenResponse createExpireTokenResponse() {
+        ExpireDelegationTokenResponseData data = new ExpireDelegationTokenResponseData()
+                .setThrottleTimeMs(20)
+                .setErrorCode(Errors.NONE.code())
+                .setExpiryTimestampMs(System.currentTimeMillis());
+        return new ExpireDelegationTokenResponse(data);
+    }
+
+    private DescribeDelegationTokenRequest createDescribeTokenRequest() {
+        List owners = new ArrayList<>();
+        owners.add(SecurityUtils.parseKafkaPrincipal("User:user1"));
+        owners.add(SecurityUtils.parseKafkaPrincipal("User:user2"));
+        return new DescribeDelegationTokenRequest.Builder(owners).build();
+    }
+
+    private DescribeDelegationTokenResponse createDescribeTokenResponse() {
+        List renewers = new ArrayList<>();
+        renewers.add(SecurityUtils.parseKafkaPrincipal("User:user1"));
+        renewers.add(SecurityUtils.parseKafkaPrincipal("User:user2"));
+
+        List tokenList = new LinkedList<>();
+
+        TokenInformation tokenInfo1 = new TokenInformation("1", SecurityUtils.parseKafkaPrincipal("User:owner"), renewers,
+            System.currentTimeMillis(), System.currentTimeMillis(), System.currentTimeMillis());
+
+        TokenInformation tokenInfo2 = new TokenInformation("2", SecurityUtils.parseKafkaPrincipal("User:owner1"), renewers,
+            System.currentTimeMillis(), System.currentTimeMillis(), System.currentTimeMillis());
+
+        tokenList.add(new DelegationToken(tokenInfo1, "test".getBytes()));
+        tokenList.add(new DelegationToken(tokenInfo2, "test".getBytes()));
+
+        return new DescribeDelegationTokenResponse(20, Errors.NONE, tokenList);
+    }
+
+    private ElectLeadersRequest createElectLeadersRequestNullPartitions() {
+        return new ElectLeadersRequest.Builder(ElectionType.PREFERRED, null, 100).build((short) 1);
+    }
+
+    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 ElectLeadersResponse createElectLeadersResponse() {
+        String topic = "myTopic";
+        List electionResults = new ArrayList<>();
+        ReplicaElectionResult electionResult = new ReplicaElectionResult();
+        electionResults.add(electionResult);
+        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, ApiKeys.ELECT_LEADERS.latestVersion());
+    }
+
+    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.NONE.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.NONE.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);
+    }
+
+    private AlterReplicaLogDirsRequest createAlterReplicaLogDirsRequest() {
+        AlterReplicaLogDirsRequestData data = new AlterReplicaLogDirsRequestData();
+        data.dirs().add(
+                new AlterReplicaLogDirsRequestData.AlterReplicaLogDir()
+                        .setPath("/data0")
+                        .setTopics(new AlterReplicaLogDirTopicCollection(Collections.singletonList(
+                                new AlterReplicaLogDirTopic()
+                                        .setPartitions(singletonList(0))
+                                        .setName("topic")
+                        ).iterator())
+                )
+        );
+        return new AlterReplicaLogDirsRequest.Builder(data).build((short) 0);
+    }
+
+    private AlterReplicaLogDirsResponse createAlterReplicaLogDirsResponse() {
+        AlterReplicaLogDirsResponseData data = new AlterReplicaLogDirsResponseData();
+        data.results().add(
+                new AlterReplicaLogDirsResponseData.AlterReplicaLogDirTopicResult()
+                        .setTopicName("topic")
+                        .setPartitions(Collections.singletonList(
+                                new AlterReplicaLogDirsResponseData.AlterReplicaLogDirPartitionResult()
+                                        .setPartitionIndex(0)
+                                        .setErrorCode(Errors.NONE.code())
+                                )
+                        )
+        );
+        return new AlterReplicaLogDirsResponse(data);
+    }
+
+    private DescribeClientQuotasRequest createDescribeClientQuotasRequest() {
+        ClientQuotaFilter filter = ClientQuotaFilter.all();
+        return new DescribeClientQuotasRequest.Builder(filter).build((short) 0);
+    }
+
+    private DescribeClientQuotasResponse createDescribeClientQuotasResponse() {
+        DescribeClientQuotasResponseData data = new DescribeClientQuotasResponseData().setEntries(asList(
+                new DescribeClientQuotasResponseData.EntryData()
+                        .setEntity(asList(new DescribeClientQuotasResponseData.EntityData()
+                            .setEntityType(ClientQuotaEntity.USER)
+                            .setEntityName("user")))
+                        .setValues(asList(new DescribeClientQuotasResponseData.ValueData()
+                            .setKey("request_percentage")
+                            .setValue(1.0)))));
+        return new DescribeClientQuotasResponse(data);
+    }
+
+    private AlterClientQuotasRequest createAlterClientQuotasRequest() {
+        ClientQuotaEntity entity = new ClientQuotaEntity(Collections.singletonMap(ClientQuotaEntity.USER, "user"));
+        ClientQuotaAlteration.Op op = new ClientQuotaAlteration.Op("request_percentage", 2.0);
+        ClientQuotaAlteration alteration = new ClientQuotaAlteration(entity, Collections.singleton(op));
+        return new AlterClientQuotasRequest.Builder(Collections.singleton(alteration), false).build((short) 0);
+    }
+
+    private AlterClientQuotasResponse createAlterClientQuotasResponse() {
+        AlterClientQuotasResponseData data = new AlterClientQuotasResponseData()
+            .setEntries(asList(new AlterClientQuotasResponseData.EntryData()
+                .setEntity(asList(new AlterClientQuotasResponseData.EntityData()
+                    .setEntityType(ClientQuotaEntity.USER)
+                    .setEntityName("user")))));
+        return new AlterClientQuotasResponse(data);
+    }
+
+    /**
+     * Check that all error codes in the response get included in {@link AbstractResponse#errorCounts()}.
+     */
+    @Test
+    public void testErrorCountsIncludesNone() {
+        assertEquals(Integer.valueOf(1), createAddOffsetsToTxnResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(1), createAddPartitionsToTxnResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(1), createAlterClientQuotasResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(1), createAlterConfigsResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(2), createAlterPartitionReassignmentsResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(1), createAlterReplicaLogDirsResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(1), createApiVersionResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(1), createControlledShutdownResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(2), createCreateAclsResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(1), createCreatePartitionsResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(1), createCreateTokenResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(1), createCreateTopicResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(1), createDeleteAclsResponse(ApiKeys.DELETE_ACLS.latestVersion()).errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(1), createDeleteGroupsResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(1), createDeleteTopicsResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(1), createDescribeAclsResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(1), createDescribeClientQuotasResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(2), createDescribeConfigsResponse(DESCRIBE_CONFIGS.latestVersion()).errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(1), createDescribeGroupResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(1), createDescribeTokenResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(2), createElectLeadersResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(1), createEndTxnResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(1), createExpireTokenResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(3), createFetchResponse(123).errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(1), createFindCoordinatorResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(1), createHeartBeatResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(1), createIncrementalAlterConfigsResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(1), createJoinGroupResponse(JOIN_GROUP.latestVersion()).errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(2), createLeaderAndIsrResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(3), createLeaderEpochResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(1), createLeaveGroupResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(1), createListGroupsResponse(LIST_GROUPS.latestVersion()).errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(1), createListOffsetResponse(LIST_OFFSETS.latestVersion()).errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(1), createListPartitionReassignmentsResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(3), createMetadataResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(1), createOffsetCommitResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(2), createOffsetDeleteResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(3), createOffsetFetchResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(1), createProduceResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(1), createRenewTokenResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(1), createSaslAuthenticateResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(1), createSaslHandshakeResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(2), createStopReplicaResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(1), createSyncGroupResponse(SYNC_GROUP.latestVersion()).errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(1), createTxnOffsetCommitResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(1), createUpdateMetadataResponse().errorCounts().get(Errors.NONE));
+        assertEquals(Integer.valueOf(1), createWriteTxnMarkersResponse().errorCounts().get(Errors.NONE));
+    }
 }
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/RequestTestUtils.java b/clients/src/test/java/org/apache/kafka/common/requests/RequestTestUtils.java
new file mode 100644
index 0000000000000..69a74a29a1997
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/RequestTestUtils.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.common.requests;
+
+import org.apache.kafka.common.Node;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.internals.Topic;
+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.ObjectSerializationCache;
+import org.apache.kafka.common.record.RecordBatch;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.function.Function;
+
+public class RequestTestUtils {
+
+    public static ByteBuffer serializeRequestHeader(RequestHeader header) {
+        ObjectSerializationCache serializationCache = new ObjectSerializationCache();
+        ByteBuffer buffer = ByteBuffer.allocate(header.size(serializationCache));
+        header.write(buffer, serializationCache);
+        buffer.flip();
+        return buffer;
+    }
+
+    public static ByteBuffer serializeRequestWithHeader(RequestHeader header, AbstractRequest request) {
+        return RequestUtils.serialize(header.data(), header.headerVersion(), request.data(), request.version());
+    }
+
+    public static ByteBuffer serializeResponseWithHeader(AbstractResponse response, short version, int correlationId) {
+        return response.serializeWithHeader(new ResponseHeader(correlationId,
+            response.apiKey().responseHeaderVersion(version)), version);
+    }
+
+    public static MetadataResponse metadataResponse(Collection brokers,
+                                                    String clusterId, int controllerId,
+                                                    List topicMetadataList) {
+        return metadataResponse(brokers, clusterId, controllerId, topicMetadataList, ApiKeys.METADATA.latestVersion());
+    }
+
+    public static MetadataResponse metadataResponse(Collection brokers,
+                                                    String clusterId, int controllerId,
+                                                    List topicMetadataList,
+                                                    short responseVersion) {
+        return metadataResponse(MetadataResponse.DEFAULT_THROTTLE_TIME, brokers, clusterId, controllerId,
+                topicMetadataList, MetadataResponse.AUTHORIZED_OPERATIONS_OMITTED, responseVersion);
+    }
+
+    public static MetadataResponse metadataResponse(int throttleTimeMs, Collection brokers,
+                                                    String clusterId, int controllerId,
+                                                    List topicMetadatas,
+                                                    int clusterAuthorizedOperations,
+                                                    short responseVersion) {
+        List topics = new ArrayList<>();
+        topicMetadatas.forEach(topicMetadata -> {
+            MetadataResponseData.MetadataResponseTopic metadataResponseTopic = new MetadataResponseData.MetadataResponseTopic();
+            metadataResponseTopic
+                    .setErrorCode(topicMetadata.error().code())
+                    .setName(topicMetadata.topic())
+                    .setIsInternal(topicMetadata.isInternal())
+                    .setTopicAuthorizedOperations(topicMetadata.authorizedOperations());
+
+            for (MetadataResponse.PartitionMetadata partitionMetadata : topicMetadata.partitionMetadata()) {
+                metadataResponseTopic.partitions().add(new MetadataResponseData.MetadataResponsePartition()
+                        .setErrorCode(partitionMetadata.error.code())
+                        .setPartitionIndex(partitionMetadata.partition())
+                        .setLeaderId(partitionMetadata.leaderId.orElse(MetadataResponse.NO_LEADER_ID))
+                        .setLeaderEpoch(partitionMetadata.leaderEpoch.orElse(RecordBatch.NO_PARTITION_LEADER_EPOCH))
+                        .setReplicaNodes(partitionMetadata.replicaIds)
+                        .setIsrNodes(partitionMetadata.inSyncReplicaIds)
+                        .setOfflineReplicas(partitionMetadata.offlineReplicaIds));
+            }
+            topics.add(metadataResponseTopic);
+        });
+        return MetadataResponse.prepareResponse(responseVersion, throttleTimeMs, brokers, clusterId, controllerId,
+                topics, clusterAuthorizedOperations); }
+
+    public static MetadataResponse metadataUpdateWith(final int numNodes,
+                                                      final Map topicPartitionCounts) {
+        return metadataUpdateWith("kafka-cluster", numNodes, topicPartitionCounts);
+    }
+
+    public static MetadataResponse metadataUpdateWith(final int numNodes,
+                                                      final Map topicPartitionCounts,
+                                                      final Function epochSupplier) {
+        return metadataUpdateWith("kafka-cluster", numNodes, Collections.emptyMap(),
+                topicPartitionCounts, epochSupplier, MetadataResponse.PartitionMetadata::new, ApiKeys.METADATA.latestVersion());
+    }
+
+    public static MetadataResponse metadataUpdateWith(final String clusterId,
+                                                      final int numNodes,
+                                                      final Map topicPartitionCounts) {
+        return metadataUpdateWith(clusterId, numNodes, Collections.emptyMap(),
+                topicPartitionCounts, tp -> null, MetadataResponse.PartitionMetadata::new, ApiKeys.METADATA.latestVersion());
+    }
+
+    public static MetadataResponse metadataUpdateWith(final String clusterId,
+                                                      final int numNodes,
+                                                      final Map topicErrors,
+                                                      final Map topicPartitionCounts) {
+        return metadataUpdateWith(clusterId, numNodes, topicErrors,
+                topicPartitionCounts, tp -> null, MetadataResponse.PartitionMetadata::new, ApiKeys.METADATA.latestVersion());
+    }
+
+    public static MetadataResponse metadataUpdateWith(final String clusterId,
+                                                      final int numNodes,
+                                                      final Map topicErrors,
+                                                      final Map topicPartitionCounts,
+                                                      final short responseVersion) {
+        return metadataUpdateWith(clusterId, numNodes, topicErrors,
+                topicPartitionCounts, tp -> null, MetadataResponse.PartitionMetadata::new, responseVersion);
+    }
+
+    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, ApiKeys.METADATA.latestVersion());
+    }
+
+    public static MetadataResponse metadataUpdateWith(final String clusterId,
+                                                      final int numNodes,
+                                                      final Map topicErrors,
+                                                      final Map topicPartitionCounts,
+                                                      final Function epochSupplier,
+                                                      final PartitionMetadataSupplier partitionSupplier,
+                                                      final short responseVersion) {
+        final List nodes = new ArrayList<>(numNodes);
+        for (int i = 0; i < numNodes; i++)
+            nodes.add(new Node(i, "localhost", 1969 + i));
+
+        List topicMetadata = new ArrayList<>();
+        for (Map.Entry topicPartitionCountEntry : topicPartitionCounts.entrySet()) {
+            String topic = topicPartitionCountEntry.getKey();
+            int numPartitions = topicPartitionCountEntry.getValue();
+
+            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 replicaIds = Collections.singletonList(leader.id());
+                partitionMetadata.add(partitionSupplier.supply(
+                        Errors.NONE, tp, Optional.of(leader.id()), Optional.ofNullable(epochSupplier.apply(tp)),
+                        replicaIds, replicaIds, replicaIds));
+            }
+
+            topicMetadata.add(new MetadataResponse.TopicMetadata(Errors.NONE, topic,
+                    Topic.isInternal(topic), partitionMetadata));
+        }
+
+        for (Map.Entry topicErrorEntry : topicErrors.entrySet()) {
+            String topic = topicErrorEntry.getKey();
+            topicMetadata.add(new MetadataResponse.TopicMetadata(topicErrorEntry.getValue(), topic,
+                    Topic.isInternal(topic), Collections.emptyList()));
+        }
+
+        return metadataResponse(nodes, clusterId, 0, topicMetadata, responseVersion);
+    }
+
+    @FunctionalInterface
+    public interface PartitionMetadataSupplier {
+        MetadataResponse.PartitionMetadata supply(Errors error,
+                                                  TopicPartition partition,
+                                                  Optional leaderId,
+                                                  Optional leaderEpoch,
+                                                  List replicas,
+                                                  List isr,
+                                                  List offlineReplicas);
+    }
+}
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..bdaa0dc0a9011
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/StopReplicaRequestTest.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.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.StopReplicaRequestData;
+import org.apache.kafka.common.message.StopReplicaRequestData.StopReplicaPartitionState;
+import org.apache.kafka.common.message.StopReplicaRequestData.StopReplicaPartitionV0;
+import org.apache.kafka.common.message.StopReplicaRequestData.StopReplicaTopicV1;
+import org.apache.kafka.common.message.StopReplicaRequestData.StopReplicaTopicState;
+import org.apache.kafka.common.message.StopReplicaResponseData.StopReplicaPartitionError;
+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.HashSet;
+import java.util.List;
+import java.util.Map;
+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.assertFalse;
+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, false, Collections.emptyList());
+        assertThrows(UnsupportedVersionException.class, builder::build);
+    }
+
+    @Test
+    public void testGetErrorResponse() {
+        List topicStates = topicStates(true);
+
+        Set expectedPartitions = new HashSet<>();
+        for (StopReplicaTopicState topicState : topicStates) {
+            for (StopReplicaPartitionState partitionState: topicState.partitionStates()) {
+                expectedPartitions.add(new StopReplicaPartitionError()
+                    .setTopicName(topicState.topicName())
+                    .setPartitionIndex(partitionState.partitionIndex())
+                    .setErrorCode(Errors.CLUSTER_AUTHORIZATION_FAILED.code()));
+            }
+        }
+
+        for (short version = STOP_REPLICA.oldestVersion(); version < STOP_REPLICA.latestVersion(); version++) {
+            StopReplicaRequest.Builder builder = new StopReplicaRequest.Builder(version,
+                    0, 0, 0L, false, topicStates);
+            StopReplicaRequest request = builder.build();
+            StopReplicaResponse response = request.getErrorResponse(0,
+                    new ClusterAuthorizationException("Not authorized"));
+            assertEquals(Errors.CLUSTER_AUTHORIZATION_FAILED, response.error());
+            assertEquals(expectedPartitions, new HashSet<>(response.partitionErrors()));
+        }
+    }
+
+    @Test
+    public void testBuilderNormalizationWithAllDeletePartitionEqualToTrue() {
+        testBuilderNormalization(true);
+    }
+
+    @Test
+    public void testBuilderNormalizationWithAllDeletePartitionEqualToFalse() {
+        testBuilderNormalization(false);
+    }
+
+    private void testBuilderNormalization(boolean deletePartitions) {
+        List topicStates = topicStates(deletePartitions);
+
+        Map expectedPartitionStates =
+            StopReplicaRequestTest.partitionStates(topicStates);
+
+        for (short version = STOP_REPLICA.oldestVersion(); version < STOP_REPLICA.latestVersion(); version++) {
+            StopReplicaRequest request = new StopReplicaRequest.Builder(version, 0, 1, 0,
+                deletePartitions, topicStates).build(version);
+            StopReplicaRequestData data = request.data();
+
+            if (version < 1) {
+                Set partitions = new HashSet<>();
+                for (StopReplicaPartitionV0 partition : data.ungroupedPartitions()) {
+                    partitions.add(new TopicPartition(partition.topicName(), partition.partitionIndex()));
+                }
+                assertEquals(expectedPartitionStates.keySet(), partitions);
+                assertEquals(deletePartitions, data.deletePartitions());
+            } else if (version < 3) {
+                Set partitions = new HashSet<>();
+                for (StopReplicaTopicV1 topic : data.topics()) {
+                    for (Integer partition : topic.partitionIndexes()) {
+                        partitions.add(new TopicPartition(topic.name(), partition));
+                    }
+                }
+                assertEquals(expectedPartitionStates.keySet(), partitions);
+                assertEquals(deletePartitions, data.deletePartitions());
+            } else {
+                Map partitionStates =
+                    StopReplicaRequestTest.partitionStates(data.topicStates());
+                assertEquals(expectedPartitionStates, partitionStates);
+                // Always false from V3 on
+                assertFalse(data.deletePartitions());
+            }
+        }
+    }
+
+    @Test
+    public void testTopicStatesNormalization() {
+        List topicStates = topicStates(true);
+
+        for (short version = STOP_REPLICA.oldestVersion(); version < STOP_REPLICA.latestVersion(); version++) {
+            // Create a request for version to get its serialized form
+            StopReplicaRequest baseRequest = new StopReplicaRequest.Builder(version, 0, 1, 0,
+                true, topicStates).build(version);
+
+            // Construct the request from the buffer
+            StopReplicaRequest request = StopReplicaRequest.parse(baseRequest.serialize(), version);
+
+            Map partitionStates =
+                StopReplicaRequestTest.partitionStates(request.topicStates());
+            assertEquals(6, partitionStates.size());
+
+            for (StopReplicaTopicState expectedTopicState : topicStates) {
+                for (StopReplicaPartitionState expectedPartitionState: expectedTopicState.partitionStates()) {
+                    TopicPartition tp = new TopicPartition(expectedTopicState.topicName(),
+                        expectedPartitionState.partitionIndex());
+                    StopReplicaPartitionState partitionState = partitionStates.get(tp);
+
+                    assertEquals(expectedPartitionState.partitionIndex(), partitionState.partitionIndex());
+                    assertTrue(partitionState.deletePartition());
+
+                    if (version >= 3) {
+                        assertEquals(expectedPartitionState.leaderEpoch(), partitionState.leaderEpoch());
+                    } else {
+                        assertEquals(-1, partitionState.leaderEpoch());
+                    }
+                }
+            }
+        }
+    }
+
+    @Test
+    public void testPartitionStatesNormalization() {
+        List topicStates = topicStates(true);
+
+        for (short version = STOP_REPLICA.oldestVersion(); version < STOP_REPLICA.latestVersion(); version++) {
+            // Create a request for version to get its serialized form
+            StopReplicaRequest baseRequest = new StopReplicaRequest.Builder(version, 0, 1, 0,
+                true, topicStates).build(version);
+
+            // Construct the request from the buffer
+            StopReplicaRequest request = StopReplicaRequest.parse(baseRequest.serialize(), version);
+
+            Map partitionStates = request.partitionStates();
+            assertEquals(6, partitionStates.size());
+
+            for (StopReplicaTopicState expectedTopicState : topicStates) {
+                for (StopReplicaPartitionState expectedPartitionState: expectedTopicState.partitionStates()) {
+                    TopicPartition tp = new TopicPartition(expectedTopicState.topicName(),
+                        expectedPartitionState.partitionIndex());
+                    StopReplicaPartitionState partitionState = partitionStates.get(tp);
+
+                    assertEquals(expectedPartitionState.partitionIndex(), partitionState.partitionIndex());
+                    assertTrue(partitionState.deletePartition());
+
+                    if (version >= 3) {
+                        assertEquals(expectedPartitionState.leaderEpoch(), partitionState.leaderEpoch());
+                    } else {
+                        assertEquals(-1, partitionState.leaderEpoch());
+                    }
+                }
+            }
+        }
+    }
+
+    private List topicStates(boolean deletePartition) {
+        List topicStates = new ArrayList<>();
+        StopReplicaTopicState topic0 = new StopReplicaTopicState()
+            .setTopicName("topic0");
+        topic0.partitionStates().add(new StopReplicaPartitionState()
+            .setPartitionIndex(0)
+            .setLeaderEpoch(0)
+            .setDeletePartition(deletePartition));
+        topic0.partitionStates().add(new StopReplicaPartitionState()
+            .setPartitionIndex(1)
+            .setLeaderEpoch(1)
+            .setDeletePartition(deletePartition));
+        topicStates.add(topic0);
+        StopReplicaTopicState topic1 = new StopReplicaTopicState()
+            .setTopicName("topic1");
+        topic1.partitionStates().add(new StopReplicaPartitionState()
+            .setPartitionIndex(2)
+            .setLeaderEpoch(2)
+            .setDeletePartition(deletePartition));
+        topic1.partitionStates().add(new StopReplicaPartitionState()
+            .setPartitionIndex(3)
+            .setLeaderEpoch(3)
+            .setDeletePartition(deletePartition));
+        topicStates.add(topic1);
+        StopReplicaTopicState topic3 = new StopReplicaTopicState()
+            .setTopicName("topic1");
+        topic3.partitionStates().add(new StopReplicaPartitionState()
+            .setPartitionIndex(4)
+            .setLeaderEpoch(-2)
+            .setDeletePartition(deletePartition));
+        topic3.partitionStates().add(new StopReplicaPartitionState()
+            .setPartitionIndex(5)
+            .setLeaderEpoch(-2)
+            .setDeletePartition(deletePartition));
+        topicStates.add(topic3);
+        return topicStates;
+    }
+
+    public static Map partitionStates(
+        Iterable topicStates) {
+        Map partitionStates = new HashMap<>();
+        for (StopReplicaTopicState topicState : topicStates) {
+            for (StopReplicaPartitionState partitionState: topicState.partitionStates()) {
+                partitionStates.put(
+                    new TopicPartition(topicState.topicName(), partitionState.partitionIndex()),
+                    partitionState);
+            }
+        }
+        return partitionStates;
+    }
+}
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 95cb3aa39e2a8..cdeff46678422 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
@@ -16,46 +16,82 @@
  */
 package org.apache.kafka.common.requests;
 
-import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.message.StopReplicaRequestData.StopReplicaPartitionState;
+import org.apache.kafka.common.message.StopReplicaRequestData.StopReplicaTopicState;
+import org.apache.kafka.common.message.StopReplicaResponseData;
+import org.apache.kafka.common.message.StopReplicaResponseData.StopReplicaPartitionError;
 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.Arrays;
 import java.util.Collections;
-import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 
+import static org.apache.kafka.common.protocol.ApiKeys.STOP_REPLICA;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
 
 public class StopReplicaResponseTest {
 
     @Test
     public void testErrorCountsFromGetErrorResponse() {
-        StopReplicaRequest request = new StopReplicaRequest.Builder(15, 20, 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());
+        List topicStates = new ArrayList<>();
+        topicStates.add(new StopReplicaTopicState()
+            .setTopicName("foo")
+            .setPartitionStates(Arrays.asList(
+                new StopReplicaPartitionState().setPartitionIndex(0),
+                new StopReplicaPartitionState().setPartitionIndex(1))));
+
+        for (short version = STOP_REPLICA.oldestVersion(); version < STOP_REPLICA.latestVersion(); version++) {
+            StopReplicaRequest request = new StopReplicaRequest.Builder(version,
+                15, 20, 0, false, topicStates).build(version);
+            StopReplicaResponse response = request
+                .getErrorResponse(0, Errors.CLUSTER_AUTHORIZATION_FAILED.exception());
+            assertEquals(Collections.singletonMap(Errors.CLUSTER_AUTHORIZATION_FAILED, 3),
+                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);
-        StopReplicaResponse response = new StopReplicaResponse(Errors.UNKNOWN_SERVER_ERROR, errors);
-        assertEquals(Collections.singletonMap(Errors.UNKNOWN_SERVER_ERROR, 2), response.errorCounts());
+        List errors = new ArrayList<>();
+        errors.add(new StopReplicaPartitionError().setTopicName("foo").setPartitionIndex(0));
+        errors.add(new StopReplicaPartitionError().setTopicName("foo").setPartitionIndex(1)
+            .setErrorCode(Errors.NOT_LEADER_OR_FOLLOWER.code()));
+        StopReplicaResponse response = new StopReplicaResponse(new StopReplicaResponseData()
+            .setErrorCode(Errors.UNKNOWN_SERVER_ERROR.code())
+            .setPartitionErrors(errors));
+        assertEquals(Collections.singletonMap(Errors.UNKNOWN_SERVER_ERROR, 3), 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());
+        assertEquals(2, errorCounts.get(Errors.NONE).intValue());
         assertEquals(1, errorCounts.get(Errors.CLUSTER_AUTHORIZATION_FAILED).intValue());
     }
 
+    @Test
+    public void testToString() {
+        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("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..98da6bcd9f44f
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/TxnOffsetCommitRequestTest.java
@@ -0,0 +1,158 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.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.List;
+import java.util.Map;
+import java.util.Optional;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertThrows;
+
+public class TxnOffsetCommitRequestTest extends OffsetCommitRequestTest {
+
+    private static String transactionalId = "transactionalId";
+    private static int producerId = 10;
+    private static short producerEpoch = 1;
+    private static int generationId = 5;
+    private static Map offsets = new HashMap<>();
+    private static TxnOffsetCommitRequest.Builder builder;
+    private static TxnOffsetCommitRequest.Builder builderWithGroupMetadata;
+
+    @Before
+    @Override
+    public void setUp() {
+        super.setUp();
+        offsets.clear();
+        offsets.put(new TopicPartition(topicOne, partitionOne),
+            new CommittedOffset(
+                offset,
+                metadata,
+                Optional.of((int) leaderEpoch)));
+        offsets.put(new TopicPartition(topicTwo, partitionTwo),
+            new CommittedOffset(
+                offset,
+                metadata,
+                Optional.of((int) leaderEpoch)));
+
+        builder = new TxnOffsetCommitRequest.Builder(
+            transactionalId,
+            groupId,
+            producerId,
+            producerEpoch,
+            offsets,
+            false);
+
+        initializeBuilderWithGroupMetadata(false);
+    }
+
+    private void initializeBuilderWithGroupMetadata(final boolean autoDowngrade) {
+        builderWithGroupMetadata = new TxnOffsetCommitRequest.Builder(
+            transactionalId,
+            groupId,
+            producerId,
+            producerEpoch,
+            offsets,
+            memberId,
+            generationId,
+            Optional.of(groupInstanceId),
+            autoDowngrade);
+    }
+
+    @Test
+    @Override
+    public void testConstructor() {
+
+        Map errorsMap = new HashMap<>();
+        errorsMap.put(new TopicPartition(topicOne, partitionOne), Errors.NOT_COORDINATOR);
+        errorsMap.put(new TopicPartition(topicTwo, partitionTwo), Errors.NOT_COORDINATOR);
+
+        List expectedTopics = 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)
+                ))
+        );
+
+        for (short version = 0; version <= ApiKeys.TXN_OFFSET_COMMIT.latestVersion(); version++) {
+            final TxnOffsetCommitRequest request;
+            if (version < 3) {
+                request = builder.build(version);
+            } else {
+                request = builderWithGroupMetadata.build(version);
+            }
+            assertEquals(offsets, request.offsets());
+            assertEquals(expectedTopics, 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());
+        }
+    }
+
+    @Test
+    public void testEnableGroupMetadataAutoDowngrade() {
+        for (short version = 0; version <= 2; version++) {
+            initializeBuilderWithGroupMetadata(true);
+            final TxnOffsetCommitRequest request = builderWithGroupMetadata.build(version);
+
+            assertEquals(JoinGroupRequest.UNKNOWN_MEMBER_ID, request.data().memberId());
+            assertEquals(JoinGroupRequest.UNKNOWN_GENERATION_ID, request.data().generationId());
+            assertNull(request.data().groupInstanceId());
+        }
+    }
+
+    @Test
+    public void testDisableGroupMetadataAutoDowngrade() {
+        for (short version = 0; version <= 2; version++) {
+            initializeBuilderWithGroupMetadata(false);
+            final short finalVersion = version;
+            assertThrows(UnsupportedVersionException.class, () -> builderWithGroupMetadata.build(finalVersion));
+        }
+    }
+}
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..397928201e313
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/TxnOffsetCommitResponseTest.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.requests;
+
+import org.apache.kafka.common.message.TxnOffsetCommitResponseData;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.MessageUtil;
+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 testParse() {
+        TxnOffsetCommitResponseData data = new TxnOffsetCommitResponseData()
+            .setThrottleTimeMs(throttleTimeMs)
+            .setTopics(Arrays.asList(
+                new TxnOffsetCommitResponseData.TxnOffsetCommitResponseTopic().setPartitions(
+                    Collections.singletonList(new TxnOffsetCommitResponseData.TxnOffsetCommitResponsePartition()
+                        .setPartitionIndex(partitionOne)
+                        .setErrorCode(errorOne.code()))),
+                    new TxnOffsetCommitResponseData.TxnOffsetCommitResponseTopic().setPartitions(
+                        Collections.singletonList(new TxnOffsetCommitResponseData.TxnOffsetCommitResponsePartition()
+                            .setPartitionIndex(partitionTwo)
+                            .setErrorCode(errorTwo.code())))
+                ));
+
+        for (short version = 0; version <= ApiKeys.TXN_OFFSET_COMMIT.latestVersion(); version++) {
+            TxnOffsetCommitResponse response = TxnOffsetCommitResponse.parse(
+                MessageUtil.toByteBuffer(data, 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..787ad0619aa90
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/UpdateMetadataRequestTest.java
@@ -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 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.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,
+                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, 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,
+                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 = request.serialize();
+            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, deserializedRequest.controllerId());
+            assertEquals(2, deserializedRequest.controllerEpoch());
+            // Broker epoch is only supported from version 5
+            if (version >= 5)
+                assertEquals(3, deserializedRequest.brokerEpoch());
+            else
+                assertEquals(-1, deserializedRequest.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,
+                partitionStates, Collections.emptyList());
+
+        assertTrue(builder.build((short) 5).sizeInBytes() <  builder.build((short) 4).sizeInBytes());
+    }
+
+    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/WriteTxnMarkersRequestTest.java b/clients/src/test/java/org/apache/kafka/common/requests/WriteTxnMarkersRequestTest.java
new file mode 100644
index 0000000000000..12fa71b92e015
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/WriteTxnMarkersRequestTest.java
@@ -0,0 +1,81 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.junit.Before;
+import org.junit.Test;
+
+import java.util.Collections;
+import java.util.List;
+
+import static org.junit.Assert.assertEquals;
+
+public class WriteTxnMarkersRequestTest {
+
+    private static long producerId = 10L;
+    private static short producerEpoch = 2;
+    private static int coordinatorEpoch = 1;
+    private static TransactionResult result = TransactionResult.COMMIT;
+    private static TopicPartition topicPartition = new TopicPartition("topic", 73);
+
+    protected static int throttleTimeMs = 10;
+
+    private static List markers;
+
+    @Before
+    public void setUp() {
+        markers = Collections.singletonList(
+             new WriteTxnMarkersRequest.TxnMarkerEntry(
+                 producerId, producerEpoch, coordinatorEpoch,
+                 result, Collections.singletonList(topicPartition))
+        );
+    }
+
+    @Test
+    public void testConstructor() {
+        WriteTxnMarkersRequest.Builder builder = new WriteTxnMarkersRequest.Builder(ApiKeys.WRITE_TXN_MARKERS.latestVersion(), markers);
+        for (short version = 0; version <= ApiKeys.WRITE_TXN_MARKERS.latestVersion(); version++) {
+            WriteTxnMarkersRequest request = builder.build(version);
+            assertEquals(1, request.markers().size());
+            WriteTxnMarkersRequest.TxnMarkerEntry marker = request.markers().get(0);
+            assertEquals(producerId, marker.producerId());
+            assertEquals(producerEpoch, marker.producerEpoch());
+            assertEquals(coordinatorEpoch, marker.coordinatorEpoch());
+            assertEquals(result, marker.transactionResult());
+            assertEquals(Collections.singletonList(topicPartition), marker.partitions());
+        }
+    }
+
+    @Test
+    public void testGetErrorResponse() {
+        WriteTxnMarkersRequest.Builder builder = new WriteTxnMarkersRequest.Builder(ApiKeys.WRITE_TXN_MARKERS.latestVersion(), markers);
+        for (short version = 0; version <= ApiKeys.WRITE_TXN_MARKERS.latestVersion(); version++) {
+            WriteTxnMarkersRequest request = builder.build(version);
+            WriteTxnMarkersResponse errorResponse =
+                request.getErrorResponse(throttleTimeMs, Errors.UNKNOWN_PRODUCER_ID.exception());
+
+            assertEquals(Collections.singletonMap(
+                topicPartition, Errors.UNKNOWN_PRODUCER_ID), errorResponse.errorsByProducerId().get(producerId));
+            assertEquals(Collections.singletonMap(Errors.UNKNOWN_PRODUCER_ID, 1), errorResponse.errorCounts());
+            // Write txn marker has no throttle time defined in response.
+            assertEquals(0, errorResponse.throttleTimeMs());
+        }
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/WriteTxnMarkersResponseTest.java b/clients/src/test/java/org/apache/kafka/common/requests/WriteTxnMarkersResponseTest.java
new file mode 100644
index 0000000000000..3860b186f072e
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/requests/WriteTxnMarkersResponseTest.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.TopicPartition;
+import org.apache.kafka.common.protocol.Errors;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.junit.Assert.assertEquals;
+
+public class WriteTxnMarkersResponseTest {
+
+    private static long producerIdOne = 1L;
+    private static long producerIdTwo = 2L;
+
+    private static TopicPartition tp1 = new TopicPartition("topic", 1);
+    private static TopicPartition tp2 = new TopicPartition("topic", 2);
+
+    private static Errors pidOneError = Errors.UNKNOWN_PRODUCER_ID;
+    private static Errors pidTwoError = Errors.INVALID_PRODUCER_EPOCH;
+
+    private static Map> errorMap;
+
+    @Before
+    public void setUp() {
+        errorMap = new HashMap<>();
+        errorMap.put(producerIdOne, Collections.singletonMap(tp1, pidOneError));
+        errorMap.put(producerIdTwo, Collections.singletonMap(tp2, pidTwoError));
+    }
+
+    @Test
+    public void testConstructor() {
+        Map expectedErrorCounts = new HashMap<>();
+        expectedErrorCounts.put(Errors.UNKNOWN_PRODUCER_ID, 1);
+        expectedErrorCounts.put(Errors.INVALID_PRODUCER_EPOCH, 1);
+        WriteTxnMarkersResponse response = new WriteTxnMarkersResponse(errorMap);
+        assertEquals(expectedErrorCounts, response.errorCounts());
+        assertEquals(Collections.singletonMap(tp1, pidOneError), response.errorsByProducerId().get(producerIdOne));
+        assertEquals(Collections.singletonMap(tp2, pidTwoError), response.errorsByProducerId().get(producerIdTwo));
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/resource/ResourceFilterTest.java b/clients/src/test/java/org/apache/kafka/common/resource/ResourceFilterTest.java
new file mode 100644
index 0000000000000..4399744d91d64
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/resource/ResourceFilterTest.java
@@ -0,0 +1,82 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.kafka.common.resource;
+
+import org.junit.Test;
+
+import static org.apache.kafka.common.resource.ResourceType.ANY;
+import static org.apache.kafka.common.resource.ResourceType.GROUP;
+import static org.apache.kafka.common.resource.ResourceType.TOPIC;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+public class ResourceFilterTest {
+    @Test
+    public void shouldNotMatchIfDifferentResourceType() {
+        assertFalse(new ResourceFilter(TOPIC, "Name")
+            .matches(new Resource(GROUP, "Name")));
+    }
+
+    @Test
+    public void shouldNotMatchIfDifferentName() {
+        assertFalse(new ResourceFilter(TOPIC, "Different")
+            .matches(new Resource(TOPIC, "Name")));
+    }
+
+    @Test
+    public void shouldNotMatchIfDifferentNameCase() {
+        assertFalse(new ResourceFilter(TOPIC, "NAME")
+            .matches(new Resource(TOPIC, "Name")));
+    }
+
+    @Test
+    public void shouldMatchWhereResourceTypeIsAny() {
+        assertTrue(new ResourceFilter(ANY, "Name")
+            .matches(new Resource(TOPIC, "Name")));
+    }
+
+    @Test
+    public void shouldMatchWhereResourceNameIsAny() {
+        assertTrue(new ResourceFilter(TOPIC, null)
+            .matches(new Resource(TOPIC, "Name")));
+    }
+
+    @Test
+    public void shouldMatchIfExactMatch() {
+        assertTrue(new ResourceFilter(TOPIC, "Name")
+            .matches(new Resource(TOPIC, "Name")));
+    }
+
+    @Test
+    public void shouldMatchWildcardIfExactMatch() {
+        assertTrue(new ResourceFilter(TOPIC, "*")
+            .matches(new Resource(TOPIC, "*")));
+    }
+
+    @Test
+    public void shouldNotMatchWildcardAgainstOtherName() {
+        assertFalse(new ResourceFilter(TOPIC, "Name")
+            .matches(new Resource(TOPIC, "*")));
+    }
+
+    @Test
+    public void shouldNotMatchLiteralWildcardTheWayAround() {
+        assertFalse(new ResourceFilter(TOPIC, "*")
+            .matches(new Resource(TOPIC, "Name")));
+    }
+}
\ No newline at end of file
diff --git a/clients/src/test/java/org/apache/kafka/common/resource/ResourceTypeTest.java b/clients/src/test/java/org/apache/kafka/common/resource/ResourceTypeTest.java
index d5f13bcff3d14..8eafba728db3f 100644
--- a/clients/src/test/java/org/apache/kafka/common/resource/ResourceTypeTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/resource/ResourceTypeTest.java
@@ -41,11 +41,12 @@ private static class AclResourceTypeTestInfo {
         new AclResourceTypeTestInfo(ResourceType.TOPIC, 2, "topic", false),
         new AclResourceTypeTestInfo(ResourceType.GROUP, 3, "group", false),
         new AclResourceTypeTestInfo(ResourceType.CLUSTER, 4, "cluster", false),
-        new AclResourceTypeTestInfo(ResourceType.TRANSACTIONAL_ID, 5, "transactional_id", false)
+        new AclResourceTypeTestInfo(ResourceType.TRANSACTIONAL_ID, 5, "transactional_id", false),
+        new AclResourceTypeTestInfo(ResourceType.DELEGATION_TOKEN, 6, "delegation_token", false)
     };
 
     @Test
-    public void testIsUnknown() throws Exception {
+    public void testIsUnknown() {
         for (AclResourceTypeTestInfo info : INFOS) {
             assertEquals(info.resourceType + " was supposed to have unknown == " + info.unknown,
                 info.unknown, info.resourceType.isUnknown());
@@ -53,7 +54,7 @@ public void testIsUnknown() throws Exception {
     }
 
     @Test
-    public void testCode() throws Exception {
+    public void testCode() {
         assertEquals(ResourceType.values().length, INFOS.length);
         for (AclResourceTypeTestInfo info : INFOS) {
             assertEquals(info.resourceType + " was supposed to have code == " + info.code,
@@ -65,7 +66,7 @@ public void testCode() throws Exception {
     }
 
     @Test
-    public void testName() throws Exception {
+    public void testName() {
         for (AclResourceTypeTestInfo info : INFOS) {
             assertEquals("ResourceType.fromString(" + info.name + ") was supposed to be " +
                 info.resourceType, info.resourceType, ResourceType.fromString(info.name));
@@ -74,7 +75,7 @@ public void testName() throws Exception {
     }
 
     @Test
-    public void testExhaustive() throws Exception {
+    public void testExhaustive() {
         assertEquals(INFOS.length, ResourceType.values().length);
         for (int i = 0; i < INFOS.length; i++) {
             assertEquals(INFOS[i].resourceType, ResourceType.values()[i]);
diff --git a/clients/src/test/java/org/apache/kafka/common/security/JaasContextTest.java b/clients/src/test/java/org/apache/kafka/common/security/JaasContextTest.java
index 49d5a86ee2ad1..8b8c251f359ba 100644
--- a/clients/src/test/java/org/apache/kafka/common/security/JaasContextTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/security/JaasContextTest.java
@@ -174,6 +174,11 @@ public void testNumericOptionWithoutQuotes() throws Exception {
         checkInvalidConfiguration("test.testNumericOptionWithoutQuotes required option1=3;");
     }
 
+    @Test
+    public void testInvalidControlFlag() throws Exception {
+        checkInvalidConfiguration("test.testInvalidControlFlag { option1=3;");
+    }
+
     @Test
     public void testNumericOptionWithQuotes() throws Exception {
         Map options = new HashMap<>();
@@ -185,54 +190,43 @@ public void testNumericOptionWithQuotes() throws Exception {
     @Test
     public void testLoadForServerWithListenerNameOverride() throws IOException {
         writeConfiguration(Arrays.asList(
-                "KafkaServer { test.LoginModuleDefault required; };",
-                "plaintext.KafkaServer { test.LoginModuleOverride requisite; };"
+            "KafkaServer { test.LoginModuleDefault required; };",
+            "plaintext.KafkaServer { test.LoginModuleOverride requisite; };"
         ));
-        JaasContext context = JaasContext.load(JaasContext.Type.SERVER, new ListenerName("plaintext"),
-                Collections.emptyMap());
+        JaasContext context = JaasContext.loadServerContext(new ListenerName("plaintext"),
+            "SOME-MECHANISM", Collections.emptyMap());
         assertEquals("plaintext.KafkaServer", context.name());
         assertEquals(JaasContext.Type.SERVER, context.type());
         assertEquals(1, context.configurationEntries().size());
         checkEntry(context.configurationEntries().get(0), "test.LoginModuleOverride",
-                LoginModuleControlFlag.REQUISITE, Collections.emptyMap());
+            LoginModuleControlFlag.REQUISITE, Collections.emptyMap());
     }
 
     @Test
     public void testLoadForServerWithListenerNameAndFallback() throws IOException {
         writeConfiguration(Arrays.asList(
-                "KafkaServer { test.LoginModule required; };",
-                "other.KafkaServer { test.LoginModuleOther requisite; };"
+            "KafkaServer { test.LoginModule required; };",
+            "other.KafkaServer { test.LoginModuleOther requisite; };"
         ));
-        JaasContext context = JaasContext.load(JaasContext.Type.SERVER, new ListenerName("plaintext"),
-                Collections.emptyMap());
+        JaasContext context = JaasContext.loadServerContext(new ListenerName("plaintext"),
+            "SOME-MECHANISM", Collections.emptyMap());
         assertEquals("KafkaServer", context.name());
         assertEquals(JaasContext.Type.SERVER, context.type());
         assertEquals(1, context.configurationEntries().size());
         checkEntry(context.configurationEntries().get(0), "test.LoginModule", LoginModuleControlFlag.REQUIRED,
-                Collections.emptyMap());
+            Collections.emptyMap());
     }
 
     @Test(expected = IllegalArgumentException.class)
     public void testLoadForServerWithWrongListenerName() throws IOException {
         writeConfiguration("Server", "test.LoginModule required;");
-        JaasContext.load(JaasContext.Type.SERVER, new ListenerName("plaintext"),
-                Collections.emptyMap());
-    }
-
-    /**
-     * ListenerName can only be used with Type.SERVER.
-     */
-    @Test(expected = IllegalArgumentException.class)
-    public void testLoadForClientWithListenerName() {
-        JaasContext.load(JaasContext.Type.CLIENT, new ListenerName("foo"),
-                Collections.emptyMap());
+        JaasContext.loadServerContext(new ListenerName("plaintext"), "SOME-MECHANISM",
+            Collections.emptyMap());
     }
 
     private AppConfigurationEntry configurationEntry(JaasContext.Type contextType, String jaasConfigProp) {
-        Map configs = new HashMap<>();
-        if (jaasConfigProp != null)
-            configs.put(SaslConfigs.SASL_JAAS_CONFIG, new Password(jaasConfigProp));
-        JaasContext context = JaasContext.load(contextType, null, contextType.name(), configs);
+        Password saslJaasConfig = jaasConfigProp == null ? null : new Password(jaasConfigProp);
+        JaasContext context = JaasContext.load(contextType, null, contextType.name(), saslJaasConfig);
         List entries = context.configurationEntries();
         assertEquals(1, entries.size());
         return entries.get(0);
diff --git a/clients/src/test/java/org/apache/kafka/common/security/SaslExtensionsTest.java b/clients/src/test/java/org/apache/kafka/common/security/SaslExtensionsTest.java
new file mode 100644
index 0000000000000..77a45235ea5d9
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/security/SaslExtensionsTest.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.common.security;
+
+import org.apache.kafka.common.security.auth.SaslExtensions;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.junit.Assert.assertNull;
+
+public class SaslExtensionsTest {
+    Map map;
+
+    @Before
+    public void setUp() {
+        this.map = new HashMap<>();
+        this.map.put("what", "42");
+        this.map.put("who", "me");
+    }
+
+    @Test(expected = UnsupportedOperationException.class)
+    public void testReturnedMapIsImmutable() {
+        SaslExtensions extensions = new SaslExtensions(this.map);
+        extensions.map().put("hello", "test");
+    }
+
+    @Test
+    public void testCannotAddValueToMapReferenceAndGetFromExtensions() {
+        SaslExtensions extensions = new SaslExtensions(this.map);
+
+        assertNull(extensions.map().get("hello"));
+        this.map.put("hello", "42");
+        assertNull(extensions.map().get("hello"));
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/security/TestSecurityConfig.java b/clients/src/test/java/org/apache/kafka/common/security/TestSecurityConfig.java
index 05294cfc5c268..07cbb7856dded 100644
--- a/clients/src/test/java/org/apache/kafka/common/security/TestSecurityConfig.java
+++ b/clients/src/test/java/org/apache/kafka/common/security/TestSecurityConfig.java
@@ -31,8 +31,13 @@ public class TestSecurityConfig extends AbstractConfig {
             .define(BrokerSecurityConfigs.SASL_ENABLED_MECHANISMS_CONFIG, Type.LIST,
                     BrokerSecurityConfigs.DEFAULT_SASL_ENABLED_MECHANISMS,
                     Importance.MEDIUM, BrokerSecurityConfigs.SASL_ENABLED_MECHANISMS_DOC)
+            .define(BrokerSecurityConfigs.SASL_SERVER_CALLBACK_HANDLER_CLASS, Type.CLASS,
+                    null,
+                    Importance.MEDIUM, BrokerSecurityConfigs.SASL_SERVER_CALLBACK_HANDLER_CLASS_DOC)
             .define(BrokerSecurityConfigs.PRINCIPAL_BUILDER_CLASS_CONFIG, Type.CLASS,
                     null, Importance.MEDIUM, BrokerSecurityConfigs.PRINCIPAL_BUILDER_CLASS_DOC)
+            .define(BrokerSecurityConfigs.CONNECTIONS_MAX_REAUTH_MS, Type.LONG, 0L, Importance.MEDIUM,
+                    BrokerSecurityConfigs.CONNECTIONS_MAX_REAUTH_MS_DOC)
             .withClientSslSupport()
             .withClientSaslSupport();
 
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 a30c09ff3167e..1fec9a1819c2f 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
@@ -16,15 +16,14 @@
  */
 package org.apache.kafka.common.security.auth;
 
+import javax.security.auth.x500.X500Principal;
 import org.apache.kafka.common.config.SaslConfigs;
 import org.apache.kafka.common.network.Authenticator;
 import org.apache.kafka.common.network.TransportLayer;
 import org.apache.kafka.common.security.authenticator.DefaultKafkaPrincipalBuilder;
-import org.apache.kafka.common.security.kerberos.KerberosName;
 import org.apache.kafka.common.security.kerberos.KerberosShortNamer;
-import org.apache.kafka.common.security.scram.ScramMechanism;
-import org.easymock.EasyMock;
-import org.easymock.EasyMockSupport;
+import org.apache.kafka.common.security.scram.internals.ScramMechanism;
+import org.apache.kafka.common.security.ssl.SslPrincipalMapper;
 import org.junit.Test;
 
 import javax.net.ssl.SSLSession;
@@ -33,8 +32,14 @@
 import java.security.Principal;
 
 import static org.junit.Assert.assertEquals;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.atLeastOnce;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
-public class DefaultKafkaPrincipalBuilderTest extends EasyMockSupport {
+public class DefaultKafkaPrincipalBuilderTest {
 
     @Test
     @SuppressWarnings("deprecation")
@@ -43,29 +48,28 @@ public void testUseOldPrincipalBuilderForPlaintextIfProvided() throws Exception
         Authenticator authenticator = mock(Authenticator.class);
         PrincipalBuilder oldPrincipalBuilder = mock(PrincipalBuilder.class);
 
-        EasyMock.expect(oldPrincipalBuilder.buildPrincipal(transportLayer, authenticator))
-                .andReturn(new DummyPrincipal("foo"));
-        oldPrincipalBuilder.close();
-        EasyMock.expectLastCall();
-
-        replayAll();
+        when(oldPrincipalBuilder.buildPrincipal(any(), any())).thenReturn(new DummyPrincipal("foo"));
 
         DefaultKafkaPrincipalBuilder builder = DefaultKafkaPrincipalBuilder.fromOldPrincipalBuilder(authenticator,
                 transportLayer, oldPrincipalBuilder, null);
 
-        KafkaPrincipal principal = builder.build(new PlaintextAuthenticationContext(InetAddress.getLocalHost()));
+        KafkaPrincipal principal = builder.build(new PlaintextAuthenticationContext(
+                InetAddress.getLocalHost(), SecurityProtocol.PLAINTEXT.name()));
         assertEquals(KafkaPrincipal.USER_TYPE, principal.getPrincipalType());
         assertEquals("foo", principal.getName());
 
         builder.close();
 
-        verifyAll();
+        verify(oldPrincipalBuilder).buildPrincipal(transportLayer, authenticator);
+        verify(oldPrincipalBuilder).close();
     }
 
     @Test
     public void testReturnAnonymousPrincipalForPlaintext() throws Exception {
-        DefaultKafkaPrincipalBuilder builder = new DefaultKafkaPrincipalBuilder(null);
-        assertEquals(KafkaPrincipal.ANONYMOUS, builder.build(new PlaintextAuthenticationContext(InetAddress.getLocalHost())));
+        try (DefaultKafkaPrincipalBuilder builder = new DefaultKafkaPrincipalBuilder(null, null)) {
+            assertEquals(KafkaPrincipal.ANONYMOUS, builder.build(
+                    new PlaintextAuthenticationContext(InetAddress.getLocalHost(), SecurityProtocol.PLAINTEXT.name())));
+        }
     }
 
     @Test
@@ -76,59 +80,113 @@ public void testUseOldPrincipalBuilderForSslIfProvided() throws Exception {
         PrincipalBuilder oldPrincipalBuilder = mock(PrincipalBuilder.class);
         SSLSession session = mock(SSLSession.class);
 
-        EasyMock.expect(oldPrincipalBuilder.buildPrincipal(transportLayer, authenticator))
-                .andReturn(new DummyPrincipal("foo"));
-        oldPrincipalBuilder.close();
-        EasyMock.expectLastCall();
-
-        replayAll();
+        when(oldPrincipalBuilder.buildPrincipal(any(), any()))
+                .thenReturn(new DummyPrincipal("foo"));
 
         DefaultKafkaPrincipalBuilder builder = DefaultKafkaPrincipalBuilder.fromOldPrincipalBuilder(authenticator,
                 transportLayer, oldPrincipalBuilder, null);
 
-        KafkaPrincipal principal = builder.build(new SslAuthenticationContext(session, InetAddress.getLocalHost()));
+        KafkaPrincipal principal = builder.build(
+                new SslAuthenticationContext(session, InetAddress.getLocalHost(), SecurityProtocol.PLAINTEXT.name()));
         assertEquals(KafkaPrincipal.USER_TYPE, principal.getPrincipalType());
         assertEquals("foo", principal.getName());
 
         builder.close();
 
-        verifyAll();
+        verify(oldPrincipalBuilder).buildPrincipal(transportLayer, authenticator);
+        verify(oldPrincipalBuilder).close();
     }
 
     @Test
     public void testUseSessionPeerPrincipalForSsl() throws Exception {
         SSLSession session = mock(SSLSession.class);
 
-        EasyMock.expect(session.getPeerPrincipal()).andReturn(new DummyPrincipal("foo"));
-
-        replayAll();
+        when(session.getPeerPrincipal()).thenReturn(new DummyPrincipal("foo"));
 
-        DefaultKafkaPrincipalBuilder builder = new DefaultKafkaPrincipalBuilder(null);
+        DefaultKafkaPrincipalBuilder builder = new DefaultKafkaPrincipalBuilder(null, null);
 
-        KafkaPrincipal principal = builder.build(new SslAuthenticationContext(session, InetAddress.getLocalHost()));
+        KafkaPrincipal principal = builder.build(
+                new SslAuthenticationContext(session, InetAddress.getLocalHost(), SecurityProtocol.PLAINTEXT.name()));
         assertEquals(KafkaPrincipal.USER_TYPE, principal.getPrincipalType());
         assertEquals("foo", principal.getName());
 
-        verifyAll();
+        builder.close();
+
+        verify(session, atLeastOnce()).getPeerPrincipal();
+    }
+
+    @Test
+    public void testPrincipalIfSSLPeerIsNotAuthenticated() throws Exception {
+        SSLSession session = mock(SSLSession.class);
+
+        when(session.getPeerPrincipal()).thenReturn(KafkaPrincipal.ANONYMOUS);
+
+        DefaultKafkaPrincipalBuilder builder = new DefaultKafkaPrincipalBuilder(null, null);
+
+        KafkaPrincipal principal = builder.build(
+                new SslAuthenticationContext(session, InetAddress.getLocalHost(), SecurityProtocol.PLAINTEXT.name()));
+        assertEquals(KafkaPrincipal.ANONYMOUS, principal);
+
+        builder.close();
+        verify(session, atLeastOnce()).getPeerPrincipal();
+    }
+
+
+    @Test
+    public void testPrincipalWithSslPrincipalMapper() throws Exception {
+        SSLSession session = mock(SSLSession.class);
+
+        when(session.getPeerPrincipal()).thenReturn(new X500Principal("CN=Duke, OU=ServiceUsers, O=Org, C=US"))
+                                        .thenReturn(new X500Principal("CN=Duke, OU=SME, O=mycp, L=Fulton, ST=MD, C=US"))
+                                        .thenReturn(new X500Principal("CN=duke, OU=JavaSoft, O=Sun Microsystems"))
+                                        .thenReturn(new X500Principal("OU=JavaSoft, O=Sun Microsystems, C=US"));
+
+        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",
+            "DEFAULT"
+        );
+
+        SslPrincipalMapper mapper = SslPrincipalMapper.fromRules(rules);
+        DefaultKafkaPrincipalBuilder builder = new DefaultKafkaPrincipalBuilder(null, mapper);
+
+        SslAuthenticationContext sslContext = new SslAuthenticationContext(session, InetAddress.getLocalHost(), SecurityProtocol.PLAINTEXT.name());
+
+        KafkaPrincipal principal = builder.build(sslContext);
+        assertEquals("duke", principal.getName());
+
+        principal = builder.build(sslContext);
+        assertEquals("duke@sme", principal.getName());
+
+        principal = builder.build(sslContext);
+        assertEquals("DUKE", principal.getName());
+
+        principal = builder.build(sslContext);
+        assertEquals("OU=JavaSoft,O=Sun Microsystems,C=US", principal.getName());
+
+        builder.close();
+        verify(session, times(4)).getPeerPrincipal();
     }
 
     @Test
     public void testPrincipalBuilderScram() throws Exception {
         SaslServer server = mock(SaslServer.class);
 
-        EasyMock.expect(server.getMechanismName()).andReturn(ScramMechanism.SCRAM_SHA_256.mechanismName());
-        EasyMock.expect(server.getAuthorizationID()).andReturn("foo");
-
-        replayAll();
+        when(server.getMechanismName()).thenReturn(ScramMechanism.SCRAM_SHA_256.mechanismName());
+        when(server.getAuthorizationID()).thenReturn("foo");
 
-        DefaultKafkaPrincipalBuilder builder = new DefaultKafkaPrincipalBuilder(null);
+        DefaultKafkaPrincipalBuilder builder = new DefaultKafkaPrincipalBuilder(null, null);
 
         KafkaPrincipal principal = builder.build(new SaslAuthenticationContext(server,
-                SecurityProtocol.SASL_PLAINTEXT, InetAddress.getLocalHost()));
+                SecurityProtocol.SASL_PLAINTEXT, InetAddress.getLocalHost(), SecurityProtocol.SASL_PLAINTEXT.name()));
         assertEquals(KafkaPrincipal.USER_TYPE, principal.getPrincipalType());
         assertEquals("foo", principal.getName());
 
-        verifyAll();
+        builder.close();
+
+        verify(server, atLeastOnce()).getMechanismName();
+        verify(server, atLeastOnce()).getAuthorizationID();
     }
 
     @Test
@@ -136,21 +194,49 @@ public void testPrincipalBuilderGssapi() throws Exception {
         SaslServer server = mock(SaslServer.class);
         KerberosShortNamer kerberosShortNamer = mock(KerberosShortNamer.class);
 
-        EasyMock.expect(server.getMechanismName()).andReturn(SaslConfigs.GSSAPI_MECHANISM);
-        EasyMock.expect(server.getAuthorizationID()).andReturn("foo/host@REALM.COM");
-        EasyMock.expect(kerberosShortNamer.shortName(EasyMock.anyObject(KerberosName.class)))
-                .andReturn("foo");
+        when(server.getMechanismName()).thenReturn(SaslConfigs.GSSAPI_MECHANISM);
+        when(server.getAuthorizationID()).thenReturn("foo/host@REALM.COM");
+        when(kerberosShortNamer.shortName(any())).thenReturn("foo");
+
+        DefaultKafkaPrincipalBuilder builder = new DefaultKafkaPrincipalBuilder(kerberosShortNamer, null);
+
+        KafkaPrincipal principal = builder.build(new SaslAuthenticationContext(server,
+                SecurityProtocol.SASL_PLAINTEXT, InetAddress.getLocalHost(), SecurityProtocol.SASL_PLAINTEXT.name()));
+        assertEquals(KafkaPrincipal.USER_TYPE, principal.getPrincipalType());
+        assertEquals("foo", principal.getName());
+
+        builder.close();
+
+        verify(server, atLeastOnce()).getMechanismName();
+        verify(server, atLeastOnce()).getAuthorizationID();
+        verify(kerberosShortNamer, atLeastOnce()).shortName(any());
+    }
+
+    @Test
+    public void testPrincipalBuilderSerde() throws Exception {
+        SaslServer server = mock(SaslServer.class);
+        KerberosShortNamer kerberosShortNamer = mock(KerberosShortNamer.class);
 
-        replayAll();
+        when(server.getMechanismName()).thenReturn(SaslConfigs.GSSAPI_MECHANISM);
+        when(server.getAuthorizationID()).thenReturn("foo/host@REALM.COM");
+        when(kerberosShortNamer.shortName(any())).thenReturn("foo");
 
-        DefaultKafkaPrincipalBuilder builder = new DefaultKafkaPrincipalBuilder(kerberosShortNamer);
+        DefaultKafkaPrincipalBuilder builder = new DefaultKafkaPrincipalBuilder(kerberosShortNamer, null);
 
         KafkaPrincipal principal = builder.build(new SaslAuthenticationContext(server,
-                SecurityProtocol.SASL_PLAINTEXT, InetAddress.getLocalHost()));
+            SecurityProtocol.SASL_PLAINTEXT, InetAddress.getLocalHost(), SecurityProtocol.SASL_PLAINTEXT.name()));
         assertEquals(KafkaPrincipal.USER_TYPE, principal.getPrincipalType());
         assertEquals("foo", principal.getName());
 
-        verifyAll();
+        byte[] serializedPrincipal = builder.serialize(principal);
+        KafkaPrincipal deserializedPrincipal = builder.deserialize(serializedPrincipal);
+        assertEquals(principal, deserializedPrincipal);
+
+        builder.close();
+
+        verify(server, atLeastOnce()).getMechanismName();
+        verify(server, atLeastOnce()).getAuthorizationID();
+        verify(kerberosShortNamer, atLeastOnce()).shortName(any());
     }
 
     private static class DummyPrincipal implements Principal {
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 7c028c4db6171..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;
@@ -34,19 +36,23 @@
 import org.apache.kafka.common.security.auth.SecurityProtocol;
 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;
 
+import java.time.Duration;
 import java.util.Arrays;
 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);
 
     private NioEchoServer server;
     private Map saslServerConfigs;
@@ -56,6 +62,7 @@ public class ClientAuthenticationFailureTest {
 
     @Before
     public void setup() throws Exception {
+        LoginManager.closeAll();
         SecurityProtocol securityProtocol = SecurityProtocol.SASL_PLAINTEXT;
 
         saslServerConfigs = new HashMap<>();
@@ -80,16 +87,14 @@ public void teardown() throws Exception {
     public void testConsumerWithInvalidCredentials() {
         Map props = new HashMap<>(saslClientConfigs);
         props.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:" + server.port());
+        props.put(ConsumerConfig.GROUP_ID_CONFIG, "");
         StringDeserializer deserializer = new StringDeserializer();
 
         try (KafkaConsumer consumer = new KafkaConsumer<>(props, deserializer, deserializer)) {
-            consumer.subscribe(Arrays.asList(topic));
-            consumer.poll(100);
-            fail("Expected an authentication error!");
-        } catch (SaslAuthenticationException e) {
-            // OK
-        } catch (Exception e) {
-            fail("Expected only an authentication error, but another error occurred: " + e.getMessage());
+            assertThrows(SaslAuthenticationException.class, () -> {
+                consumer.subscribe(Collections.singleton(topic));
+                consumer.poll(Duration.ofSeconds(10));
+            });
         }
     }
 
@@ -101,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);
         }
     }
 
@@ -113,18 +115,14 @@ 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);
         }
     }
 
     @Test
-    public void testTransactionalProducerWithInvalidCredentials() throws Exception {
+    public void testTransactionalProducerWithInvalidCredentials() {
         Map props = new HashMap<>(saslClientConfigs);
         props.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:" + server.port());
         props.put(ProducerConfig.TRANSACTIONAL_ID_CONFIG, "txclient-1");
@@ -132,10 +130,7 @@ public void testTransactionalProducerWithInvalidCredentials() throws Exception {
         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);
         }
     }
 
@@ -145,6 +140,6 @@ private NioEchoServer createEchoServer(SecurityProtocol securityProtocol) throws
 
     private NioEchoServer createEchoServer(ListenerName listenerName, SecurityProtocol securityProtocol) throws Exception {
         return NetworkTestUtils.createEchoServer(listenerName, securityProtocol,
-                new TestSecurityConfig(saslServerConfigs), new CredentialCache());
+                new TestSecurityConfig(saslServerConfigs), new CredentialCache(), time);
     }
 }
diff --git a/clients/src/test/java/org/apache/kafka/common/security/authenticator/LoginManagerTest.java b/clients/src/test/java/org/apache/kafka/common/security/authenticator/LoginManagerTest.java
new file mode 100644
index 0000000000000..34ff8817cba42
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/security/authenticator/LoginManagerTest.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.security.authenticator;
+
+import org.apache.kafka.common.config.types.Password;
+import org.apache.kafka.common.network.ListenerName;
+import org.apache.kafka.common.security.JaasContext;
+import org.apache.kafka.common.security.plain.PlainLoginModule;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.junit.Assert.assertNotSame;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertEquals;
+
+public class LoginManagerTest {
+
+    private Password dynamicPlainContext;
+    private Password dynamicDigestContext;
+
+    @Before
+    public void setUp() {
+        dynamicPlainContext = new Password(PlainLoginModule.class.getName() +
+                " required user=\"plainuser\" password=\"plain-secret\";");
+        dynamicDigestContext = new Password(TestDigestLoginModule.class.getName() +
+                " required user=\"digestuser\" password=\"digest-secret\";");
+        TestJaasConfig.createConfiguration("SCRAM-SHA-256",
+                Collections.singletonList("SCRAM-SHA-256"));
+    }
+
+    @After
+    public void tearDown() {
+        LoginManager.closeAll();
+    }
+
+    @Test
+    public void testClientLoginManager() throws Exception {
+        Map configs = Collections.singletonMap("sasl.jaas.config", dynamicPlainContext);
+        JaasContext dynamicContext = JaasContext.loadClientContext(configs);
+        JaasContext staticContext = JaasContext.loadClientContext(Collections.emptyMap());
+
+        LoginManager dynamicLogin = LoginManager.acquireLoginManager(dynamicContext, "PLAIN",
+                DefaultLogin.class, configs);
+        assertEquals(dynamicPlainContext, dynamicLogin.cacheKey());
+        LoginManager staticLogin = LoginManager.acquireLoginManager(staticContext, "SCRAM-SHA-256",
+                DefaultLogin.class, configs);
+        assertNotSame(dynamicLogin, staticLogin);
+        assertEquals("KafkaClient", staticLogin.cacheKey());
+
+        assertSame(dynamicLogin, LoginManager.acquireLoginManager(dynamicContext, "PLAIN",
+                DefaultLogin.class, configs));
+        assertSame(staticLogin, LoginManager.acquireLoginManager(staticContext, "SCRAM-SHA-256",
+                DefaultLogin.class, configs));
+
+        verifyLoginManagerRelease(dynamicLogin, 2, dynamicContext, configs);
+        verifyLoginManagerRelease(staticLogin, 2, staticContext, configs);
+    }
+
+    @Test
+    public void testServerLoginManager() throws Exception {
+        Map configs = new HashMap<>();
+        configs.put("plain.sasl.jaas.config", dynamicPlainContext);
+        configs.put("digest-md5.sasl.jaas.config", dynamicDigestContext);
+        ListenerName listenerName = new ListenerName("listener1");
+        JaasContext plainJaasContext = JaasContext.loadServerContext(listenerName, "PLAIN", configs);
+        JaasContext digestJaasContext = JaasContext.loadServerContext(listenerName, "DIGEST-MD5", configs);
+        JaasContext scramJaasContext = JaasContext.loadServerContext(listenerName, "SCRAM-SHA-256", configs);
+
+        LoginManager dynamicPlainLogin = LoginManager.acquireLoginManager(plainJaasContext, "PLAIN",
+                DefaultLogin.class, configs);
+        assertEquals(dynamicPlainContext, dynamicPlainLogin.cacheKey());
+        LoginManager dynamicDigestLogin = LoginManager.acquireLoginManager(digestJaasContext, "DIGEST-MD5",
+                DefaultLogin.class, configs);
+        assertNotSame(dynamicPlainLogin, dynamicDigestLogin);
+        assertEquals(dynamicDigestContext, dynamicDigestLogin.cacheKey());
+        LoginManager staticScramLogin = LoginManager.acquireLoginManager(scramJaasContext, "SCRAM-SHA-256",
+                DefaultLogin.class, configs);
+        assertNotSame(dynamicPlainLogin, staticScramLogin);
+        assertEquals("KafkaServer", staticScramLogin.cacheKey());
+
+        assertSame(dynamicPlainLogin, LoginManager.acquireLoginManager(plainJaasContext, "PLAIN",
+                DefaultLogin.class, configs));
+        assertSame(dynamicDigestLogin, LoginManager.acquireLoginManager(digestJaasContext, "DIGEST-MD5",
+                DefaultLogin.class, configs));
+        assertSame(staticScramLogin, LoginManager.acquireLoginManager(scramJaasContext, "SCRAM-SHA-256",
+                DefaultLogin.class, configs));
+
+        verifyLoginManagerRelease(dynamicPlainLogin, 2, plainJaasContext, configs);
+        verifyLoginManagerRelease(dynamicDigestLogin, 2, digestJaasContext, configs);
+        verifyLoginManagerRelease(staticScramLogin, 2, scramJaasContext, configs);
+    }
+
+    private void verifyLoginManagerRelease(LoginManager loginManager, int acquireCount, JaasContext jaasContext,
+                                           Map configs) throws Exception {
+
+        // Release all except one reference and verify that the loginManager is still cached
+        for (int i = 0; i < acquireCount - 1; i++)
+            loginManager.release();
+        assertSame(loginManager, LoginManager.acquireLoginManager(jaasContext, "PLAIN",
+                DefaultLogin.class, configs));
+
+        // Release all references and verify that new LoginManager is created on next acquire
+        for (int i = 0; i < 2; i++) // release all references
+            loginManager.release();
+        LoginManager newLoginManager = LoginManager.acquireLoginManager(jaasContext, "PLAIN",
+                DefaultLogin.class, configs);
+        assertNotSame(loginManager, newLoginManager);
+        newLoginManager.release();
+    }
+}
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
new file mode 100644
index 0000000000000..19003ed56da90
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/security/authenticator/SaslAuthenticatorFailureDelayTest.java
@@ -0,0 +1,264 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.authenticator;
+
+import org.apache.kafka.common.config.SaslConfigs;
+import org.apache.kafka.common.config.internals.BrokerSecurityConfigs;
+import org.apache.kafka.common.errors.SaslAuthenticationException;
+import org.apache.kafka.common.network.CertStores;
+import org.apache.kafka.common.network.ChannelBuilder;
+import org.apache.kafka.common.network.ChannelBuilders;
+import org.apache.kafka.common.network.ChannelState;
+import org.apache.kafka.common.network.ListenerName;
+import org.apache.kafka.common.network.NetworkTestUtils;
+import org.apache.kafka.common.network.NioEchoServer;
+import org.apache.kafka.common.network.Selector;
+import org.apache.kafka.common.security.JaasContext;
+import org.apache.kafka.common.security.TestSecurityConfig;
+import org.apache.kafka.common.security.auth.SecurityProtocol;
+import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.test.TestUtils;
+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 java.io.IOException;
+import java.net.InetSocketAddress;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+@RunWith(value = Parameterized.class)
+public class SaslAuthenticatorFailureDelayTest {
+    private static final int BUFFER_SIZE = 4 * 1024;
+
+    private final MockTime time = new MockTime(1);
+    private NioEchoServer server;
+    private Selector selector;
+    private ChannelBuilder channelBuilder;
+    private CertStores serverCertStores;
+    private CertStores clientCertStores;
+    private Map saslClientConfigs;
+    private Map saslServerConfigs;
+    private CredentialCache credentialCache;
+    private long startTimeMs;
+    private final int failedAuthenticationDelayMs;
+
+    public SaslAuthenticatorFailureDelayTest(int failedAuthenticationDelayMs) {
+        this.failedAuthenticationDelayMs = failedAuthenticationDelayMs;
+    }
+
+    @Parameterized.Parameters(name = "failedAuthenticationDelayMs={0}")
+    public static Collection data() {
+        List values = new ArrayList<>();
+        values.add(new Object[]{0});
+        values.add(new Object[]{200});
+        return values;
+    }
+
+    @Before
+    public void setup() throws Exception {
+        LoginManager.closeAll();
+        serverCertStores = new CertStores(true, "localhost");
+        clientCertStores = new CertStores(false, "localhost");
+        saslServerConfigs = serverCertStores.getTrustingConfig(clientCertStores);
+        saslClientConfigs = clientCertStores.getTrustingConfig(serverCertStores);
+        credentialCache = new CredentialCache();
+        SaslAuthenticatorTest.TestLogin.loginCount.set(0);
+        startTimeMs = time.milliseconds();
+    }
+
+    @After
+    public void teardown() throws Exception {
+        long now = time.milliseconds();
+        if (server != null)
+            this.server.close();
+        if (selector != null)
+            this.selector.close();
+        if (failedAuthenticationDelayMs != -1)
+            assertTrue("timeSpent: " + (now - startTimeMs), now - startTimeMs >= failedAuthenticationDelayMs);
+    }
+
+    /**
+     * Tests that SASL/PLAIN clients with invalid password fail authentication.
+     */
+    @Test
+    public void testInvalidPasswordSaslPlain() throws Exception {
+        String node = "0";
+        SecurityProtocol securityProtocol = SecurityProtocol.SASL_SSL;
+        TestJaasConfig jaasConfig = configureMechanisms("PLAIN", Arrays.asList("PLAIN"));
+        jaasConfig.setClientOptions("PLAIN", TestJaasConfig.USERNAME, "invalidpassword");
+
+        server = createEchoServer(securityProtocol);
+        createAndCheckClientAuthenticationFailure(securityProtocol, node, "PLAIN",
+                "Authentication failed: Invalid username or password");
+        server.verifyAuthenticationMetrics(0, 1);
+    }
+
+    /**
+     * Tests that SASL/SCRAM clients with invalid password fail authentication with
+     * connection close delay if configured.
+     */
+    @Test
+    public void testInvalidPasswordSaslScram() throws Exception {
+        String node = "0";
+        SecurityProtocol securityProtocol = SecurityProtocol.SASL_SSL;
+        TestJaasConfig jaasConfig = configureMechanisms("SCRAM-SHA-256", Collections.singletonList("SCRAM-SHA-256"));
+        jaasConfig.setClientOptions("SCRAM-SHA-256", TestJaasConfig.USERNAME, "invalidpassword");
+
+        server = createEchoServer(securityProtocol);
+        createAndCheckClientAuthenticationFailure(securityProtocol, node, "SCRAM-SHA-256", null);
+        server.verifyAuthenticationMetrics(0, 1);
+    }
+
+    /**
+     * Tests that clients with disabled SASL mechanism fail authentication with
+     * connection close delay if configured.
+     */
+    @Test
+    public void testDisabledSaslMechanism() throws Exception {
+        String node = "0";
+        SecurityProtocol securityProtocol = SecurityProtocol.SASL_SSL;
+        TestJaasConfig jaasConfig = configureMechanisms("SCRAM-SHA-256", Collections.singletonList("SCRAM-SHA-256"));
+        jaasConfig.setClientOptions("PLAIN", TestJaasConfig.USERNAME, "invalidpassword");
+
+        server = createEchoServer(securityProtocol);
+        createAndCheckClientAuthenticationFailure(securityProtocol, node, "SCRAM-SHA-256", null);
+        server.verifyAuthenticationMetrics(0, 1);
+    }
+
+    /**
+     * Tests client connection close before response for authentication failure is sent.
+     */
+    @Test
+    public void testClientConnectionClose() throws Exception {
+        String node = "0";
+        SecurityProtocol securityProtocol = SecurityProtocol.SASL_SSL;
+        TestJaasConfig jaasConfig = configureMechanisms("PLAIN", Arrays.asList("PLAIN"));
+        jaasConfig.setClientOptions("PLAIN", TestJaasConfig.USERNAME, "invalidpassword");
+
+        server = createEchoServer(securityProtocol);
+        createClientConnection(securityProtocol, node);
+
+        Map delayedClosingChannels = NetworkTestUtils.delayedClosingChannels(server.selector());
+
+        // Wait until server has established connection with client and has processed the auth failure
+        TestUtils.waitForCondition(() -> {
+            poll(selector);
+            return !server.selector().channels().isEmpty();
+        }, "Timeout waiting for connection");
+        TestUtils.waitForCondition(() -> {
+            poll(selector);
+            return failedAuthenticationDelayMs == 0 || !delayedClosingChannels.isEmpty();
+        }, "Timeout waiting for auth failure");
+
+        selector.close();
+        selector = null;
+
+        // Now that client connection is closed, wait until server notices the disconnection and removes it from the
+        // list of connected channels and from delayed response for auth failure
+        TestUtils.waitForCondition(() -> failedAuthenticationDelayMs == 0 || delayedClosingChannels.isEmpty(),
+                "Timeout waiting for delayed response remove");
+        TestUtils.waitForCondition(() -> server.selector().channels().isEmpty(),
+                "Timeout waiting for connection close");
+
+        // Try forcing completion of delayed channel close
+        TestUtils.waitForCondition(() -> time.milliseconds() > startTimeMs + failedAuthenticationDelayMs + 1,
+                "Timeout when waiting for auth failure response timeout to elapse");
+        NetworkTestUtils.completeDelayedChannelClose(server.selector(), time.nanoseconds());
+    }
+
+    private void poll(Selector selector) {
+        try {
+            selector.poll(50);
+        } catch (IOException e) {
+            Assert.fail("Caught unexpected exception " + e);
+        }
+    }
+
+    private TestJaasConfig configureMechanisms(String clientMechanism, List serverMechanisms) {
+        saslClientConfigs.put(SaslConfigs.SASL_MECHANISM, clientMechanism);
+        saslServerConfigs.put(BrokerSecurityConfigs.SASL_ENABLED_MECHANISMS_CONFIG, serverMechanisms);
+        if (serverMechanisms.contains("DIGEST-MD5")) {
+            saslServerConfigs.put("digest-md5." + BrokerSecurityConfigs.SASL_SERVER_CALLBACK_HANDLER_CLASS,
+                    TestDigestLoginModule.DigestServerCallbackHandler.class.getName());
+        }
+        return TestJaasConfig.createConfiguration(clientMechanism, serverMechanisms);
+    }
+
+    private void createSelector(SecurityProtocol securityProtocol, Map clientConfigs) {
+        if (selector != null) {
+            selector.close();
+            selector = null;
+        }
+
+        String saslMechanism = (String) saslClientConfigs.get(SaslConfigs.SASL_MECHANISM);
+        this.channelBuilder = ChannelBuilders.clientChannelBuilder(securityProtocol, JaasContext.Type.CLIENT,
+                new TestSecurityConfig(clientConfigs), null, saslMechanism, time, true,
+                new LogContext());
+        this.selector = NetworkTestUtils.createSelector(channelBuilder, time);
+    }
+
+    private NioEchoServer createEchoServer(SecurityProtocol securityProtocol) throws Exception {
+        return createEchoServer(ListenerName.forSecurityProtocol(securityProtocol), securityProtocol);
+    }
+
+    private NioEchoServer createEchoServer(ListenerName listenerName, SecurityProtocol securityProtocol) throws Exception {
+        if (failedAuthenticationDelayMs != -1)
+            return NetworkTestUtils.createEchoServer(listenerName, securityProtocol,
+                    new TestSecurityConfig(saslServerConfigs), credentialCache, failedAuthenticationDelayMs, time);
+        else
+            return NetworkTestUtils.createEchoServer(listenerName, securityProtocol,
+                    new TestSecurityConfig(saslServerConfigs), credentialCache, time);
+    }
+
+    private void createClientConnection(SecurityProtocol securityProtocol, String node) throws Exception {
+        createSelector(securityProtocol, saslClientConfigs);
+        InetSocketAddress addr = new InetSocketAddress("127.0.0.1", server.port());
+        selector.connect(node, addr, BUFFER_SIZE, BUFFER_SIZE);
+    }
+
+    private void createAndCheckClientAuthenticationFailure(SecurityProtocol securityProtocol, String node,
+                                                           String mechanism, String expectedErrorMessage) throws Exception {
+        ChannelState finalState = createAndCheckClientConnectionFailure(securityProtocol, node);
+        Exception exception = finalState.exception();
+        assertTrue("Invalid exception class " + exception.getClass(), exception instanceof SaslAuthenticationException);
+        if (expectedErrorMessage == null)
+            expectedErrorMessage = "Authentication failed during authentication due to invalid credentials with SASL mechanism " + mechanism;
+        assertEquals(expectedErrorMessage, exception.getMessage());
+    }
+
+    private ChannelState createAndCheckClientConnectionFailure(SecurityProtocol securityProtocol, String node)
+            throws Exception {
+        createClientConnection(securityProtocol, node);
+        ChannelState finalState = NetworkTestUtils.waitForChannelClose(selector, node,
+                ChannelState.State.AUTHENTICATION_FAILED);
+        selector.close();
+        selector = null;
+        return finalState;
+    }
+}
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 c4426727ba628..34ee542a5c6cd 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
@@ -16,15 +16,61 @@
  */
 package org.apache.kafka.common.security.authenticator;
 
+import java.io.IOException;
+import java.net.InetSocketAddress;
+import java.nio.ByteBuffer;
+import java.nio.channels.SelectionKey;
+import java.nio.charset.StandardCharsets;
+import java.security.NoSuchAlgorithmException;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Base64;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+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;
+import javax.security.auth.callback.CallbackHandler;
+import javax.security.auth.callback.NameCallback;
+import javax.security.auth.callback.PasswordCallback;
+import javax.security.auth.callback.UnsupportedCallbackException;
+import javax.security.auth.login.Configuration;
+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.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.ListOffsetResponseData;
+import org.apache.kafka.common.message.ListOffsetResponseData.ListOffsetTopicResponse;
+import org.apache.kafka.common.message.ListOffsetResponseData.ListOffsetPartitionResponse;
+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.ByteBufferSend;
 import org.apache.kafka.common.network.CertStores;
 import org.apache.kafka.common.network.ChannelBuilder;
 import org.apache.kafka.common.network.ChannelBuilders;
+import org.apache.kafka.common.network.ChannelMetadataRegistry;
 import org.apache.kafka.common.network.ChannelState;
 import org.apache.kafka.common.network.ListenerName;
 import org.apache.kafka.common.network.Mode;
@@ -33,16 +79,18 @@
 import org.apache.kafka.common.network.NioEchoServer;
 import org.apache.kafka.common.network.SaslChannelBuilder;
 import org.apache.kafka.common.network.Selector;
-import org.apache.kafka.common.network.Send;
 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.requests.RequestTestUtils;
+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;
@@ -51,34 +99,41 @@
 import org.apache.kafka.common.requests.SaslHandshakeResponse;
 import org.apache.kafka.common.security.JaasContext;
 import org.apache.kafka.common.security.TestSecurityConfig;
+import org.apache.kafka.common.security.auth.KafkaPrincipal;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerLoginModule;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerToken;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerTokenCallback;
+import org.apache.kafka.common.security.oauthbearer.internals.unsecured.OAuthBearerConfigException;
+import org.apache.kafka.common.security.oauthbearer.internals.unsecured.OAuthBearerIllegalTokenException;
+import org.apache.kafka.common.security.oauthbearer.internals.unsecured.OAuthBearerUnsecuredJws;
+import org.apache.kafka.common.security.oauthbearer.internals.unsecured.OAuthBearerUnsecuredLoginCallbackHandler;
 import org.apache.kafka.common.security.plain.PlainLoginModule;
 import org.apache.kafka.common.security.scram.ScramCredential;
-import org.apache.kafka.common.security.scram.ScramCredentialUtils;
-import org.apache.kafka.common.security.scram.ScramFormatter;
+import org.apache.kafka.common.security.scram.internals.ScramCredentialUtils;
+import org.apache.kafka.common.security.scram.internals.ScramFormatter;
 import org.apache.kafka.common.security.scram.ScramLoginModule;
-import org.apache.kafka.common.security.scram.ScramMechanism;
+import org.apache.kafka.common.security.scram.internals.ScramMechanism;
+import org.apache.kafka.common.security.token.delegation.TokenInformation;
+import org.apache.kafka.common.security.token.delegation.internals.DelegationTokenCache;
+import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.common.utils.SecurityUtils;
+import org.apache.kafka.common.security.auth.AuthenticateCallbackHandler;
+import org.apache.kafka.common.security.authenticator.TestDigestLoginModule.DigestServerCallbackHandler;
+import org.apache.kafka.common.security.plain.internals.PlainServerCallbackHandler;
+
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.common.utils.Utils;
 import org.junit.After;
+import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
 
-import javax.security.auth.Subject;
-import javax.security.auth.login.Configuration;
-import java.io.IOException;
-import java.net.InetSocketAddress;
-import java.nio.ByteBuffer;
-import java.nio.channels.SelectionKey;
-import java.security.NoSuchAlgorithmException;
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.Iterator;
-import java.util.List;
-import java.util.Map;
-import java.util.Random;
-
+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.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
@@ -87,7 +142,9 @@
  */
 public class SaslAuthenticatorTest {
 
+    private static final long CONNECTIONS_MAX_REAUTH_MS_VALUE = 100L;
     private static final int BUFFER_SIZE = 4 * 1024;
+    private static Time time = Time.SYSTEM;
 
     private NioEchoServer server;
     private Selector selector;
@@ -101,11 +158,14 @@ public class SaslAuthenticatorTest {
 
     @Before
     public void setup() throws Exception {
+        LoginManager.closeAll();
+        time = Time.SYSTEM;
         serverCertStores = new CertStores(true, "localhost");
         clientCertStores = new CertStores(false, "localhost");
         saslServerConfigs = serverCertStores.getTrustingConfig(clientCertStores);
         saslClientConfigs = clientCertStores.getTrustingConfig(serverCertStores);
         credentialCache = new CredentialCache();
+        TestLogin.loginCount.set(0);
     }
 
     @After
@@ -118,6 +178,7 @@ public void teardown() throws Exception {
 
     /**
      * Tests good path SASL/PLAIN client and server channels using SSL transport layer.
+     * Also tests successful re-authentication.
      */
     @Test
     public void testValidSaslPlainOverSsl() throws Exception {
@@ -126,12 +187,12 @@ public void testValidSaslPlainOverSsl() throws Exception {
         configureMechanisms("PLAIN", Arrays.asList("PLAIN"));
 
         server = createEchoServer(securityProtocol);
-        createAndCheckClientConnection(securityProtocol, node);
-        server.verifyAuthenticationMetrics(1, 0);
+        checkAuthenticationAndReauthentication(securityProtocol, node);
     }
 
     /**
      * Tests good path SASL/PLAIN client and server channels using PLAINTEXT transport layer.
+     * Also tests successful re-authentication.
      */
     @Test
     public void testValidSaslPlainOverPlaintext() throws Exception {
@@ -140,8 +201,7 @@ public void testValidSaslPlainOverPlaintext() throws Exception {
         configureMechanisms("PLAIN", Arrays.asList("PLAIN"));
 
         server = createEchoServer(securityProtocol);
-        createAndCheckClientConnection(securityProtocol, node);
-        server.verifyAuthenticationMetrics(1, 0);
+        checkAuthenticationAndReauthentication(securityProtocol, node);
     }
 
     /**
@@ -158,6 +218,7 @@ public void testInvalidPasswordSaslPlain() throws Exception {
         createAndCheckClientAuthenticationFailure(securityProtocol, node, "PLAIN",
                 "Authentication failed: Invalid username or password");
         server.verifyAuthenticationMetrics(0, 1);
+        server.verifyReauthenticationMetrics(0, 0);
     }
 
     /**
@@ -174,6 +235,7 @@ public void testInvalidUsernameSaslPlain() throws Exception {
         createAndCheckClientAuthenticationFailure(securityProtocol, node, "PLAIN",
                 "Authentication failed: Invalid username or password");
         server.verifyAuthenticationMetrics(0, 1);
+        server.verifyReauthenticationMetrics(0, 0);
     }
 
     /**
@@ -221,6 +283,74 @@ public void testMissingPasswordSaslPlain() throws Exception {
         }
     }
 
+    /**
+     * Verify that messages from SaslExceptions thrown in the server during authentication are not
+     * propagated to the client since these may contain sensitive data.
+     */
+    @Test
+    public void testClientExceptionDoesNotContainSensitiveData() throws Exception {
+        InvalidScramServerCallbackHandler.reset();
+
+        SecurityProtocol securityProtocol = SecurityProtocol.SASL_PLAINTEXT;
+        TestJaasConfig jaasConfig = configureMechanisms("SCRAM-SHA-256", Collections.singletonList("SCRAM-SHA-256"));
+        jaasConfig.createOrUpdateEntry(TestJaasConfig.LOGIN_CONTEXT_SERVER, PlainLoginModule.class.getName(), new HashMap<>());
+        String callbackPrefix = ListenerName.forSecurityProtocol(securityProtocol).saslMechanismConfigPrefix("SCRAM-SHA-256");
+        saslServerConfigs.put(callbackPrefix + BrokerSecurityConfigs.SASL_SERVER_CALLBACK_HANDLER_CLASS,
+                InvalidScramServerCallbackHandler.class.getName());
+        server = createEchoServer(securityProtocol);
+
+        try {
+            InvalidScramServerCallbackHandler.sensitiveException =
+                    new IOException("Could not connect to password database locahost:8000");
+            createAndCheckClientAuthenticationFailure(securityProtocol, "1", "SCRAM-SHA-256", null);
+
+            InvalidScramServerCallbackHandler.sensitiveException =
+                    new SaslException("Password for existing user " + TestServerCallbackHandler.USERNAME + " is invalid");
+            createAndCheckClientAuthenticationFailure(securityProtocol, "1", "SCRAM-SHA-256", null);
+
+            InvalidScramServerCallbackHandler.reset();
+            InvalidScramServerCallbackHandler.clientFriendlyException =
+                    new SaslAuthenticationException("Credential verification failed");
+            createAndCheckClientAuthenticationFailure(securityProtocol, "1", "SCRAM-SHA-256",
+                    InvalidScramServerCallbackHandler.clientFriendlyException.getMessage());
+        } finally {
+            InvalidScramServerCallbackHandler.reset();
+        }
+    }
+
+    public static class InvalidScramServerCallbackHandler implements AuthenticateCallbackHandler {
+        // We want to test three types of exceptions:
+        //   1) IOException since we can throw this from callback handlers. This may be sensitive.
+        //   2) SaslException (also an IOException) which may contain data from external (or JRE) servers and callbacks and may be sensitive
+        //   3) SaslAuthenticationException which is from our own code and is used only for client-friendly exceptions
+        // We use two different exceptions here since the only checked exception CallbackHandler can throw is IOException,
+        // covering case 1) and 2). For case 3), SaslAuthenticationException is a RuntimeExceptiom.
+        static volatile IOException sensitiveException;
+        static volatile SaslAuthenticationException clientFriendlyException;
+
+        @Override
+        public void configure(Map configs, String saslMechanism, List jaasConfigEntries) {
+        }
+
+        @Override
+        public void handle(Callback[] callbacks) throws IOException {
+            if (sensitiveException != null)
+                throw sensitiveException;
+            if (clientFriendlyException != null)
+                throw clientFriendlyException;
+        }
+
+        @Override
+        public void close() {
+            reset();
+        }
+
+        static void reset() {
+            sensitiveException = null;
+            clientFriendlyException = null;
+        }
+    }
+
     /**
      * Tests that mechanisms that are not supported in Kafka can be plugged in without modifying
      * Kafka code if Sasl client and server providers are available.
@@ -230,6 +360,7 @@ public void testMechanismPluggability() throws Exception {
         String node = "0";
         SecurityProtocol securityProtocol = SecurityProtocol.SASL_SSL;
         configureMechanisms("DIGEST-MD5", Arrays.asList("DIGEST-MD5"));
+        configureDigestMd5ServerCallback(securityProtocol);
 
         server = createEchoServer(securityProtocol);
         createAndCheckClientConnection(securityProtocol, node);
@@ -238,34 +369,66 @@ public void testMechanismPluggability() throws Exception {
     /**
      * Tests that servers supporting multiple SASL mechanisms work with clients using
      * any of the enabled mechanisms.
+     * Also tests successful re-authentication over multiple mechanisms.
      */
     @Test
     public void testMultipleServerMechanisms() throws Exception {
         SecurityProtocol securityProtocol = SecurityProtocol.SASL_SSL;
         configureMechanisms("DIGEST-MD5", Arrays.asList("DIGEST-MD5", "PLAIN", "SCRAM-SHA-256"));
+        configureDigestMd5ServerCallback(securityProtocol);
         server = createEchoServer(securityProtocol);
         updateScramCredentialCache(TestJaasConfig.USERNAME, TestJaasConfig.PASSWORD);
 
         String node1 = "1";
         saslClientConfigs.put(SaslConfigs.SASL_MECHANISM, "PLAIN");
         createAndCheckClientConnection(securityProtocol, node1);
+        server.verifyAuthenticationMetrics(1, 0);
 
-        String node2 = "2";
-        saslClientConfigs.put(SaslConfigs.SASL_MECHANISM, "DIGEST-MD5");
-        createSelector(securityProtocol, saslClientConfigs);
-        InetSocketAddress addr = new InetSocketAddress("127.0.0.1", server.port());
-        selector.connect(node2, addr, BUFFER_SIZE, BUFFER_SIZE);
-        NetworkTestUtils.checkClientConnection(selector, node2, 100, 10);
-
-        String node3 = "3";
-        saslClientConfigs.put(SaslConfigs.SASL_MECHANISM, "SCRAM-SHA-256");
-        createSelector(securityProtocol, saslClientConfigs);
-        selector.connect(node3, new InetSocketAddress("127.0.0.1", server.port()), BUFFER_SIZE, BUFFER_SIZE);
-        NetworkTestUtils.checkClientConnection(selector, node3, 100, 10);
+        Selector selector2 = null;
+        Selector selector3 = null;
+        try {
+            String node2 = "2";
+            saslClientConfigs.put(SaslConfigs.SASL_MECHANISM, "DIGEST-MD5");
+            createSelector(securityProtocol, saslClientConfigs);
+            selector2 = selector;
+            InetSocketAddress addr = new InetSocketAddress("127.0.0.1", server.port());
+            selector.connect(node2, addr, BUFFER_SIZE, BUFFER_SIZE);
+            NetworkTestUtils.checkClientConnection(selector, node2, 100, 10);
+            selector = null; // keeps it from being closed when next one is created
+            server.verifyAuthenticationMetrics(2, 0);
+
+            String node3 = "3";
+            saslClientConfigs.put(SaslConfigs.SASL_MECHANISM, "SCRAM-SHA-256");
+            createSelector(securityProtocol, saslClientConfigs);
+            selector3 = selector;
+            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.
+             */
+            delay((long) (CONNECTIONS_MAX_REAUTH_MS_VALUE * 1.1));
+            server.verifyReauthenticationMetrics(0, 0);
+
+            NetworkTestUtils.checkClientConnection(selector2, node2, 100, 10);
+            server.verifyReauthenticationMetrics(1, 0);
+
+            NetworkTestUtils.checkClientConnection(selector3, node3, 100, 10);
+            server.verifyReauthenticationMetrics(2, 0);
+            
+        } finally {
+            if (selector2 != null)
+                selector2.close();
+            if (selector3 != null)
+                selector3.close();
+        }
     }
 
     /**
      * Tests good path SASL/SCRAM-SHA-256 client and server channels.
+     * Also tests successful re-authentication.
      */
     @Test
     public void testValidSaslScramSha256() throws Exception {
@@ -274,8 +437,7 @@ public void testValidSaslScramSha256() throws Exception {
 
         server = createEchoServer(securityProtocol);
         updateScramCredentialCache(TestJaasConfig.USERNAME, TestJaasConfig.PASSWORD);
-        createAndCheckClientConnection(securityProtocol, "0");
-        server.verifyAuthenticationMetrics(1, 0);
+        checkAuthenticationAndReauthentication(securityProtocol, "0");
     }
 
     /**
@@ -312,6 +474,7 @@ public void testInvalidPasswordSaslScram() throws Exception {
         updateScramCredentialCache(TestJaasConfig.USERNAME, TestJaasConfig.PASSWORD);
         createAndCheckClientAuthenticationFailure(securityProtocol, node, "SCRAM-SHA-256", null);
         server.verifyAuthenticationMetrics(0, 1);
+        server.verifyReauthenticationMetrics(0, 0);
     }
 
     /**
@@ -331,6 +494,7 @@ public void testUnknownUserSaslScram() throws Exception {
         updateScramCredentialCache(TestJaasConfig.USERNAME, TestJaasConfig.PASSWORD);
         createAndCheckClientAuthenticationFailure(securityProtocol, node, "SCRAM-SHA-256", null);
         server.verifyAuthenticationMetrics(0, 1);
+        server.verifyReauthenticationMetrics(0, 0);
     }
 
     /**
@@ -353,6 +517,7 @@ public void testUserCredentialsUnavailableForScramMechanism() throws Exception {
         saslClientConfigs.put(SaslConfigs.SASL_MECHANISM, "SCRAM-SHA-512");
         createAndCheckClientConnection(securityProtocol, "2");
         server.verifyAuthenticationMetrics(1, 1);
+        server.verifyReauthenticationMetrics(0, 0);
     }
 
     /**
@@ -375,6 +540,104 @@ public void testScramUsernameWithSpecialCharacters() throws Exception {
         createAndCheckClientConnection(securityProtocol, "0");
     }
 
+
+    @Test
+    public void testTokenAuthenticationOverSaslScram() throws Exception {
+        SecurityProtocol securityProtocol = SecurityProtocol.SASL_SSL;
+        TestJaasConfig jaasConfig = configureMechanisms("SCRAM-SHA-256", Arrays.asList("SCRAM-SHA-256"));
+
+        //create jaas config for token auth
+        Map options = new HashMap<>();
+        String tokenId = "token1";
+        String tokenHmac = "abcdefghijkl";
+        options.put("username", tokenId); //tokenId
+        options.put("password", tokenHmac); //token hmac
+        options.put(ScramLoginModule.TOKEN_AUTH_CONFIG, "true"); //enable token authentication
+        jaasConfig.createOrUpdateEntry(TestJaasConfig.LOGIN_CONTEXT_CLIENT, ScramLoginModule.class.getName(), options);
+
+        server = createEchoServer(securityProtocol);
+
+        //Check invalid tokenId/tokenInfo in tokenCache
+        createAndCheckClientConnectionFailure(securityProtocol, "0");
+        server.verifyAuthenticationMetrics(0, 1);
+
+        //Check valid token Info and invalid credentials
+        KafkaPrincipal owner = SecurityUtils.parseKafkaPrincipal("User:Owner");
+        KafkaPrincipal renewer = SecurityUtils.parseKafkaPrincipal("User:Renewer1");
+        TokenInformation tokenInfo = new TokenInformation(tokenId, owner, Collections.singleton(renewer),
+            System.currentTimeMillis(), System.currentTimeMillis(), System.currentTimeMillis());
+        server.tokenCache().addToken(tokenId, tokenInfo);
+        createAndCheckClientConnectionFailure(securityProtocol, "0");
+        server.verifyAuthenticationMetrics(0, 2);
+
+        //Check with valid token Info and credentials
+        updateTokenCredentialCache(tokenId, tokenHmac);
+        createAndCheckClientConnection(securityProtocol, "0");
+        server.verifyAuthenticationMetrics(1, 2);
+        server.verifyReauthenticationMetrics(0, 0);
+    }
+
+    @Test
+    public void testTokenReauthenticationOverSaslScram() throws Exception {
+        SecurityProtocol securityProtocol = SecurityProtocol.SASL_SSL;
+        TestJaasConfig jaasConfig = configureMechanisms("SCRAM-SHA-256", Arrays.asList("SCRAM-SHA-256"));
+
+        // create jaas config for token auth
+        Map options = new HashMap<>();
+        String tokenId = "token1";
+        String tokenHmac = "abcdefghijkl";
+        options.put("username", tokenId); // tokenId
+        options.put("password", tokenHmac); // token hmac
+        options.put(ScramLoginModule.TOKEN_AUTH_CONFIG, "true"); // enable token authentication
+        jaasConfig.createOrUpdateEntry(TestJaasConfig.LOGIN_CONTEXT_CLIENT, ScramLoginModule.class.getName(), options);
+
+        // ensure re-authentication based on token expiry rather than a default value
+        saslServerConfigs.put(BrokerSecurityConfigs.CONNECTIONS_MAX_REAUTH_MS, Long.MAX_VALUE);
+        /*
+         * create a token cache that adjusts the token expiration dynamically so that
+         * the first time the expiry is read during authentication we use it to define a
+         * session expiration time that we can then sleep through; then the second time
+         * the value is read (during re-authentication) it will be in the future.
+         */
+        Function tokenLifetime = callNum -> 10 * callNum * CONNECTIONS_MAX_REAUTH_MS_VALUE;
+        DelegationTokenCache tokenCache = new DelegationTokenCache(ScramMechanism.mechanismNames()) {
+            int callNum = 0;
+
+            @Override
+            public TokenInformation token(String tokenId) {
+                TokenInformation baseTokenInfo = super.token(tokenId);
+                long thisLifetimeMs = System.currentTimeMillis() + tokenLifetime.apply(++callNum).longValue();
+                TokenInformation retvalTokenInfo = new TokenInformation(baseTokenInfo.tokenId(), baseTokenInfo.owner(),
+                        baseTokenInfo.renewers(), baseTokenInfo.issueTimestamp(), thisLifetimeMs, thisLifetimeMs);
+                return retvalTokenInfo;
+            }
+        };
+        server = createEchoServer(ListenerName.forSecurityProtocol(securityProtocol), securityProtocol, tokenCache);
+
+        KafkaPrincipal owner = SecurityUtils.parseKafkaPrincipal("User:Owner");
+        KafkaPrincipal renewer = SecurityUtils.parseKafkaPrincipal("User:Renewer1");
+        TokenInformation tokenInfo = new TokenInformation(tokenId, owner, Collections.singleton(renewer),
+                System.currentTimeMillis(), System.currentTimeMillis(), System.currentTimeMillis());
+        server.tokenCache().addToken(tokenId, tokenInfo);
+        updateTokenCredentialCache(tokenId, tokenHmac);
+        // initial authentication must succeed
+        createClientConnection(securityProtocol, "0");
+        checkClientConnection("0");
+        // ensure metrics are as expected before trying to re-authenticate
+        server.verifyAuthenticationMetrics(1, 0);
+        server.verifyReauthenticationMetrics(0, 0);
+        /*
+         * Now re-authenticate and ensure it succeeds. We have to sleep long enough so
+         * that the current delegation token will be expired when the next write occurs;
+         * this will trigger a re-authentication. Then the second time the delegation
+         * token is read and transmitted to the server it will again have an expiration
+         * date in the future.
+         */
+        delay(tokenLifetime.apply(1));
+        checkClientConnection("0");
+        server.verifyReauthenticationMetrics(1, 0);
+    }
+
     /**
      * Tests that Kafka ApiVersionsRequests are handled by the SASL server authenticator
      * prior to SASL handshake flow and that subsequent authentication succeeds
@@ -416,7 +679,7 @@ public void testUnauthenticatedApiVersionsRequestOverSslHandshakeVersion0() thro
      */
     @Test
     public void testUnauthenticatedApiVersionsRequestOverSslHandshakeVersion1() throws Exception {
-        testUnauthenticatedApiVersionsRequest(SecurityProtocol.SASL_PLAINTEXT, (short) 1);
+        testUnauthenticatedApiVersionsRequest(SecurityProtocol.SASL_SSL, (short) 1);
     }
 
     /**
@@ -427,7 +690,7 @@ public void testUnauthenticatedApiVersionsRequestOverSslHandshakeVersion1() thro
      * {@link SaslServerAuthenticator} of the server prior to client authentication.
      */
     @Test
-    public void testApiVersionsRequestWithUnsupportedVersion() throws Exception {
+    public void testApiVersionsRequestWithServerUnsupportedVersion() throws Exception {
         short handshakeVersion = ApiKeys.SASL_HANDSHAKE.latestVersion();
         SecurityProtocol securityProtocol = SecurityProtocol.SASL_PLAINTEXT;
         configureMechanisms("PLAIN", Arrays.asList("PLAIN"));
@@ -436,13 +699,78 @@ 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));
+        selector.send(new NetworkSend(node, request.toSend(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);
+
+        // Test that client can authenticate successfully
+        sendHandshakeRequestReceiveResponse(node, handshakeVersion);
+        authenticateUsingSaslPlainAndCheckConnection(node, handshakeVersion > 0);
+    }
+
+    /**
+     * Tests correct negotiation of handshake and authenticate api versions by having the server
+     * return a higher version than supported on the client.
+     * Note, that due to KAFKA-9577 this will require a workaround to effectively bump
+     * SASL_HANDSHAKE in the future.
+     */
+    @Test
+    public void testSaslUnsupportedClientVersions() throws Exception {
+        configureMechanisms("SCRAM-SHA-512", Arrays.asList("SCRAM-SHA-512"));
+
+        server = startServerApiVersionsUnsupportedByClient(SecurityProtocol.SASL_SSL, "SCRAM-SHA-512");
+        updateScramCredentialCache(TestJaasConfig.USERNAME, TestJaasConfig.PASSWORD);
+
+        String node = "0";
+
+        createClientConnection(SecurityProtocol.SASL_SSL, "SCRAM-SHA-512", node, true);
+        NetworkTestUtils.checkClientConnection(selector, "0", 100, 10);
+    }
+
+    /**
+     * 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(new NetworkSend(node, request.toSend(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);
@@ -452,6 +780,44 @@ public void testApiVersionsRequestWithUnsupportedVersion() throws Exception {
         authenticateUsingSaslPlainAndCheckConnection(node, handshakeVersion > 0);
     }
 
+
+    @Test
+    public void testForBrokenSaslHandshakeVersionBump() {
+        assertEquals("It is not possible to easily bump SASL_HANDSHAKE schema"
+                        + " due to improper version negotiation in clients < 2.5."
+                        + " Please see https://issues.apache.org/jira/browse/KAFKA-9577",
+                ApiKeys.SASL_HANDSHAKE.latestVersion(),
+                1);
+    }
+
+    /**
+     * 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(new NetworkSend(node, request.toSend(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
@@ -468,9 +834,10 @@ public void testSaslHandshakeRequestWithUnsupportedVersion() throws Exception {
         // Send SaslHandshakeRequest and validate that connection is closed by server.
         String node1 = "invalid1";
         createClientConnection(SecurityProtocol.PLAINTEXT, node1);
-        SaslHandshakeRequest request = new SaslHandshakeRequest("PLAIN");
+        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));
+        
+        selector.send(new NetworkSend(node1, request.toSend(header)));
         // This test uses a non-SASL PLAINTEXT client in order to do manual handshake.
         // So the channel is in READY state.
         NetworkTestUtils.waitForChannelClose(selector, node1, ChannelState.READY.state());
@@ -498,7 +865,7 @@ public void testInvalidSaslPacket() throws Exception {
         Random random = new Random();
         byte[] bytes = new byte[1024];
         random.nextBytes(bytes);
-        selector.send(new NetworkSend(node1, ByteBuffer.wrap(bytes)));
+        selector.send(new NetworkSend(node1, ByteBufferSend.sizePrefixed(ByteBuffer.wrap(bytes))));
         NetworkTestUtils.waitForChannelClose(selector, node1, ChannelState.READY.state());
         selector.close();
 
@@ -509,7 +876,7 @@ public void testInvalidSaslPacket() throws Exception {
         String node2 = "invalid2";
         createClientConnection(SecurityProtocol.PLAINTEXT, node2);
         random.nextBytes(bytes);
-        selector.send(new NetworkSend(node2, ByteBuffer.wrap(bytes)));
+        selector.send(new NetworkSend(node2, ByteBufferSend.sizePrefixed(ByteBuffer.wrap(bytes))));
         NetworkTestUtils.waitForChannelClose(selector, node2, ChannelState.READY.state());
         selector.close();
 
@@ -537,7 +904,7 @@ public void testInvalidApiVersionsRequestSequence() throws Exception {
 
         ApiVersionsRequest request = createApiVersionsRequestV0();
         RequestHeader versionsHeader = new RequestHeader(ApiKeys.API_VERSIONS, request.version(), "someclient", 2);
-        selector.send(request.toSend(node1, versionsHeader));
+        selector.send(new NetworkSend(node1, request.toSend(versionsHeader)));
         NetworkTestUtils.waitForChannelClose(selector, node1, ChannelState.READY.state());
         selector.close();
 
@@ -564,7 +931,7 @@ public void testPacketSizeTooBig() throws Exception {
         buffer.putInt(Integer.MAX_VALUE);
         buffer.put(new byte[buffer.capacity() - 4]);
         buffer.rewind();
-        selector.send(new NetworkSend(node1, buffer));
+        selector.send(new NetworkSend(node1, ByteBufferSend.sizePrefixed(buffer)));
         NetworkTestUtils.waitForChannelClose(selector, node1, ChannelState.READY.state());
         selector.close();
 
@@ -578,7 +945,7 @@ public void testPacketSizeTooBig() throws Exception {
         buffer.putInt(Integer.MAX_VALUE);
         buffer.put(new byte[buffer.capacity() - 4]);
         buffer.rewind();
-        selector.send(new NetworkSend(node2, buffer));
+        selector.send(new NetworkSend(node2, ByteBufferSend.sizePrefixed(buffer)));
         NetworkTestUtils.waitForChannelClose(selector, node2, ChannelState.READY.state());
         selector.close();
 
@@ -603,7 +970,7 @@ public void testDisallowedKafkaRequestsBeforeAuthentication() throws Exception {
                 true).build();
         RequestHeader metadataRequestHeader1 = new RequestHeader(ApiKeys.METADATA, metadataRequest1.version(),
                 "someclient", 1);
-        selector.send(metadataRequest1.toSend(node1, metadataRequestHeader1));
+        selector.send(new NetworkSend(node1, metadataRequest1.toSend(metadataRequestHeader1)));
         NetworkTestUtils.waitForChannelClose(selector, node1, ChannelState.READY.state());
         selector.close();
 
@@ -617,7 +984,7 @@ public void testDisallowedKafkaRequestsBeforeAuthentication() throws Exception {
         MetadataRequest metadataRequest2 = new MetadataRequest.Builder(Collections.singletonList("sometopic"), true).build();
         RequestHeader metadataRequestHeader2 = new RequestHeader(ApiKeys.METADATA,
                 metadataRequest2.version(), "someclient", 2);
-        selector.send(metadataRequest2.toSend(node2, metadataRequestHeader2));
+        selector.send(new NetworkSend(node2, metadataRequest2.toSend(metadataRequestHeader2)));
         NetworkTestUtils.waitForChannelClose(selector, node2, ChannelState.READY.state());
         selector.close();
 
@@ -643,6 +1010,207 @@ public void testInvalidLoginModule() throws Exception {
         }
     }
 
+    /**
+     * Tests SASL client authentication callback handler override.
+     */
+    @Test
+    public void testClientAuthenticateCallbackHandler() throws Exception {
+        SecurityProtocol securityProtocol = SecurityProtocol.SASL_PLAINTEXT;
+        TestJaasConfig jaasConfig = configureMechanisms("PLAIN", Collections.singletonList("PLAIN"));
+        saslClientConfigs.put(SaslConfigs.SASL_CLIENT_CALLBACK_HANDLER_CLASS, TestClientCallbackHandler.class.getName());
+        jaasConfig.setClientOptions("PLAIN", "", ""); // remove username, password in login context
+
+        Map options = new HashMap<>();
+        options.put("user_" + TestClientCallbackHandler.USERNAME, TestClientCallbackHandler.PASSWORD);
+        jaasConfig.createOrUpdateEntry(TestJaasConfig.LOGIN_CONTEXT_SERVER, PlainLoginModule.class.getName(), options);
+        server = createEchoServer(securityProtocol);
+        createAndCheckClientConnection(securityProtocol, "good");
+
+        options.clear();
+        options.put("user_" + TestClientCallbackHandler.USERNAME, "invalid-password");
+        jaasConfig.createOrUpdateEntry(TestJaasConfig.LOGIN_CONTEXT_SERVER, PlainLoginModule.class.getName(), options);
+        createAndCheckClientConnectionFailure(securityProtocol, "invalid");
+    }
+
+    /**
+     * Tests SASL server authentication callback handler override.
+     */
+    @Test
+    public void testServerAuthenticateCallbackHandler() throws Exception {
+        SecurityProtocol securityProtocol = SecurityProtocol.SASL_PLAINTEXT;
+        TestJaasConfig jaasConfig = configureMechanisms("PLAIN", Collections.singletonList("PLAIN"));
+        jaasConfig.createOrUpdateEntry(TestJaasConfig.LOGIN_CONTEXT_SERVER, PlainLoginModule.class.getName(), new HashMap());
+        String callbackPrefix = ListenerName.forSecurityProtocol(securityProtocol).saslMechanismConfigPrefix("PLAIN");
+        saslServerConfigs.put(callbackPrefix + BrokerSecurityConfigs.SASL_SERVER_CALLBACK_HANDLER_CLASS,
+                TestServerCallbackHandler.class.getName());
+        server = createEchoServer(securityProtocol);
+
+        // Set client username/password to the values used by `TestServerCallbackHandler`
+        jaasConfig.setClientOptions("PLAIN", TestServerCallbackHandler.USERNAME, TestServerCallbackHandler.PASSWORD);
+        createAndCheckClientConnection(securityProtocol, "good");
+
+        // Set client username/password to the invalid values
+        jaasConfig.setClientOptions("PLAIN", TestJaasConfig.USERNAME, "invalid-password");
+        createAndCheckClientConnectionFailure(securityProtocol, "invalid");
+    }
+
+    /**
+     * Test that callback handlers are only applied to connections for the mechanisms
+     * configured for the handler. Test enables two mechanisms 'PLAIN` and `DIGEST-MD5`
+     * on the servers with different callback handlers for the two mechanisms. Verifies
+     * that clients using both mechanisms authenticate successfully.
+     */
+    @Test
+    public void testAuthenticateCallbackHandlerMechanisms() throws Exception {
+        SecurityProtocol securityProtocol = SecurityProtocol.SASL_PLAINTEXT;
+        TestJaasConfig jaasConfig = configureMechanisms("DIGEST-MD5", Arrays.asList("DIGEST-MD5", "PLAIN"));
+
+        // Connections should fail using the digest callback handler if listener.mechanism prefix not specified
+        saslServerConfigs.put("plain." + BrokerSecurityConfigs.SASL_SERVER_CALLBACK_HANDLER_CLASS,
+                TestServerCallbackHandler.class);
+        saslServerConfigs.put("digest-md5." + BrokerSecurityConfigs.SASL_SERVER_CALLBACK_HANDLER_CLASS,
+                DigestServerCallbackHandler.class);
+        server = createEchoServer(securityProtocol);
+        createAndCheckClientConnectionFailure(securityProtocol, "invalid");
+
+        // Connections should succeed using the server callback handler associated with the listener
+        ListenerName listener = ListenerName.forSecurityProtocol(securityProtocol);
+        saslServerConfigs.remove("plain." + BrokerSecurityConfigs.SASL_SERVER_CALLBACK_HANDLER_CLASS);
+        saslServerConfigs.remove("digest-md5." + BrokerSecurityConfigs.SASL_SERVER_CALLBACK_HANDLER_CLASS);
+        saslServerConfigs.put(listener.saslMechanismConfigPrefix("plain") + BrokerSecurityConfigs.SASL_SERVER_CALLBACK_HANDLER_CLASS,
+                TestServerCallbackHandler.class);
+        saslServerConfigs.put(listener.saslMechanismConfigPrefix("digest-md5") + BrokerSecurityConfigs.SASL_SERVER_CALLBACK_HANDLER_CLASS,
+                DigestServerCallbackHandler.class);
+        server = createEchoServer(securityProtocol);
+
+        // Verify that DIGEST-MD5 (currently configured for client) works with `DigestServerCallbackHandler`
+        createAndCheckClientConnection(securityProtocol, "good-digest-md5");
+
+        // Verify that PLAIN works with `TestServerCallbackHandler`
+        jaasConfig.setClientOptions("PLAIN", TestServerCallbackHandler.USERNAME, TestServerCallbackHandler.PASSWORD);
+        saslClientConfigs.put(SaslConfigs.SASL_MECHANISM, "PLAIN");
+        createAndCheckClientConnection(securityProtocol, "good-plain");
+    }
+
+    /**
+     * Tests SASL login class override.
+     */
+    @Test
+    public void testClientLoginOverride() throws Exception {
+        SecurityProtocol securityProtocol = SecurityProtocol.SASL_PLAINTEXT;
+        TestJaasConfig jaasConfig = configureMechanisms("PLAIN", Collections.singletonList("PLAIN"));
+        jaasConfig.setClientOptions("PLAIN", "invaliduser", "invalidpassword");
+        server = createEchoServer(securityProtocol);
+
+        // Connection should succeed using login override that sets correct username/password in Subject
+        saslClientConfigs.put(SaslConfigs.SASL_LOGIN_CLASS, TestLogin.class.getName());
+        createAndCheckClientConnection(securityProtocol, "1");
+        assertEquals(1, TestLogin.loginCount.get());
+
+        // Connection should fail without login override since username/password in jaas config is invalid
+        saslClientConfigs.remove(SaslConfigs.SASL_LOGIN_CLASS);
+        createAndCheckClientConnectionFailure(securityProtocol, "invalid");
+        assertEquals(1, TestLogin.loginCount.get());
+    }
+
+    /**
+     * Tests SASL server login class override.
+     */
+    @Test
+    public void testServerLoginOverride() throws Exception {
+        SecurityProtocol securityProtocol = SecurityProtocol.SASL_PLAINTEXT;
+        configureMechanisms("PLAIN", Collections.singletonList("PLAIN"));
+        String prefix = ListenerName.forSecurityProtocol(securityProtocol).saslMechanismConfigPrefix("PLAIN");
+        saslServerConfigs.put(prefix + SaslConfigs.SASL_LOGIN_CLASS, TestLogin.class.getName());
+        server = createEchoServer(securityProtocol);
+
+        // Login is performed when server channel builder is created (before any connections are made on the server)
+        assertEquals(1, TestLogin.loginCount.get());
+
+        createAndCheckClientConnection(securityProtocol, "1");
+        assertEquals(1, TestLogin.loginCount.get());
+    }
+
+    /**
+     * Tests SASL login callback class override.
+     */
+    @Test
+    public void testClientLoginCallbackOverride() throws Exception {
+        SecurityProtocol securityProtocol = SecurityProtocol.SASL_PLAINTEXT;
+        TestJaasConfig jaasConfig = configureMechanisms("PLAIN", Collections.singletonList("PLAIN"));
+        jaasConfig.createOrUpdateEntry(TestJaasConfig.LOGIN_CONTEXT_CLIENT, TestPlainLoginModule.class.getName(),
+                Collections.emptyMap());
+        server = createEchoServer(securityProtocol);
+
+        // Connection should succeed using login callback override that sets correct username/password
+        saslClientConfigs.put(SaslConfigs.SASL_LOGIN_CALLBACK_HANDLER_CLASS, TestLoginCallbackHandler.class.getName());
+        createAndCheckClientConnection(securityProtocol, "1");
+
+        // Connection should fail without login callback override since username/password in jaas config is invalid
+        saslClientConfigs.remove(SaslConfigs.SASL_LOGIN_CALLBACK_HANDLER_CLASS);
+        try {
+            createClientConnection(securityProtocol, "invalid");
+        } catch (Exception e) {
+            assertTrue("Unexpected exception " + e.getCause(), e.getCause() instanceof LoginException);
+        }
+    }
+
+    /**
+     * Tests SASL server login callback class override.
+     */
+    @Test
+    public void testServerLoginCallbackOverride() throws Exception {
+        SecurityProtocol securityProtocol = SecurityProtocol.SASL_PLAINTEXT;
+        TestJaasConfig jaasConfig = configureMechanisms("PLAIN", Collections.singletonList("PLAIN"));
+        jaasConfig.createOrUpdateEntry(TestJaasConfig.LOGIN_CONTEXT_SERVER, TestPlainLoginModule.class.getName(),
+                Collections.emptyMap());
+        jaasConfig.setClientOptions("PLAIN", TestServerCallbackHandler.USERNAME, TestServerCallbackHandler.PASSWORD);
+        ListenerName listenerName = ListenerName.forSecurityProtocol(securityProtocol);
+        String prefix = listenerName.saslMechanismConfigPrefix("PLAIN");
+        saslServerConfigs.put(prefix + BrokerSecurityConfigs.SASL_SERVER_CALLBACK_HANDLER_CLASS,
+                TestServerCallbackHandler.class);
+        Class loginCallback = TestLoginCallbackHandler.class;
+
+        try {
+            createEchoServer(securityProtocol);
+            fail("Should have failed to create server with default login handler");
+        } catch (KafkaException e) {
+            // Expected exception
+        }
+
+        try {
+            saslServerConfigs.put(SaslConfigs.SASL_LOGIN_CALLBACK_HANDLER_CLASS, loginCallback);
+            createEchoServer(securityProtocol);
+            fail("Should have failed to create server with login handler config without listener+mechanism prefix");
+        } catch (KafkaException e) {
+            // Expected exception
+            saslServerConfigs.remove(SaslConfigs.SASL_LOGIN_CALLBACK_HANDLER_CLASS);
+        }
+
+        try {
+            saslServerConfigs.put("plain." + SaslConfigs.SASL_LOGIN_CALLBACK_HANDLER_CLASS, loginCallback);
+            createEchoServer(securityProtocol);
+            fail("Should have failed to create server with login handler config without listener prefix");
+        } catch (KafkaException e) {
+            // Expected exception
+            saslServerConfigs.remove("plain." + SaslConfigs.SASL_LOGIN_CALLBACK_HANDLER_CLASS);
+        }
+
+        try {
+            saslServerConfigs.put(listenerName.configPrefix() + SaslConfigs.SASL_LOGIN_CALLBACK_HANDLER_CLASS, loginCallback);
+            createEchoServer(securityProtocol);
+            fail("Should have failed to create server with login handler config without mechanism prefix");
+        } catch (KafkaException e) {
+            // Expected exception
+            saslServerConfigs.remove("plain." + SaslConfigs.SASL_LOGIN_CALLBACK_HANDLER_CLASS);
+        }
+
+        // Connection should succeed using login callback override for mechanism
+        saslServerConfigs.put(prefix + SaslConfigs.SASL_LOGIN_CALLBACK_HANDLER_CLASS, loginCallback);
+        server = createEchoServer(securityProtocol);
+        createAndCheckClientConnection(securityProtocol, "1");
+    }
+
     /**
      * Tests that mechanisms with default implementation in Kafka may be disabled in
      * the Kafka server by removing from the enabled mechanism list.
@@ -656,6 +1224,7 @@ public void testDisabledMechanism() throws Exception {
         server = createEchoServer(securityProtocol);
         createAndCheckClientConnectionFailure(securityProtocol, node);
         server.verifyAuthenticationMetrics(0, 1);
+        server.verifyReauthenticationMetrics(0, 0);
     }
 
     /**
@@ -669,8 +1238,18 @@ public void testInvalidMechanism() throws Exception {
         saslClientConfigs.put(SaslConfigs.SASL_MECHANISM, "INVALID");
 
         server = createEchoServer(securityProtocol);
-        createAndCheckClientConnectionFailure(securityProtocol, node);
-        server.verifyAuthenticationMetrics(0, 1);
+        try {
+            createAndCheckClientConnectionFailure(securityProtocol, node);
+            fail("Did not generate exception prior to creating channel");
+        } catch (IOException expected) {
+            server.verifyAuthenticationMetrics(0, 0);
+            server.verifyReauthenticationMetrics(0, 0);
+            Throwable underlyingCause = expected.getCause().getCause().getCause();
+            assertEquals(SaslAuthenticationException.class, underlyingCause.getClass());
+            assertEquals("Failed to create SaslClient with mechanism INVALID", underlyingCause.getMessage());
+        } finally {
+            closeClientConnectionIfNecessary();
+        }
     }
 
     /**
@@ -679,7 +1258,7 @@ public void testInvalidMechanism() throws Exception {
      * property override is used during authentication.
      */
     @Test
-    public void testDynamicJaasConfiguration() throws Exception {
+    public void testClientDynamicJaasConfiguration() throws Exception {
         SecurityProtocol securityProtocol = SecurityProtocol.SASL_SSL;
         saslClientConfigs.put(SaslConfigs.SASL_MECHANISM, "PLAIN");
         saslServerConfigs.put(BrokerSecurityConfigs.SASL_ENABLED_MECHANISMS_CONFIG, Arrays.asList("PLAIN"));
@@ -720,6 +1299,37 @@ public void testDynamicJaasConfiguration() throws Exception {
         }
     }
 
+    /**
+     * Tests dynamic JAAS configuration property for SASL server. Invalid server credentials
+     * are set in the static JVM-wide configuration instance to ensure that the dynamic
+     * property override is used during authentication.
+     */
+    @Test
+    public void testServerDynamicJaasConfiguration() throws Exception {
+        SecurityProtocol securityProtocol = SecurityProtocol.SASL_SSL;
+        saslClientConfigs.put(SaslConfigs.SASL_MECHANISM, "PLAIN");
+        saslServerConfigs.put(BrokerSecurityConfigs.SASL_ENABLED_MECHANISMS_CONFIG, Arrays.asList("PLAIN"));
+        Map serverOptions = new HashMap<>();
+        serverOptions.put("user_user1", "user1-secret");
+        serverOptions.put("user_user2", "user2-secret");
+        saslServerConfigs.put("listener.name.sasl_ssl.plain." + SaslConfigs.SASL_JAAS_CONFIG,
+                TestJaasConfig.jaasConfigProperty("PLAIN", serverOptions));
+        TestJaasConfig staticJaasConfig = new TestJaasConfig();
+        staticJaasConfig.createOrUpdateEntry(TestJaasConfig.LOGIN_CONTEXT_SERVER, PlainLoginModule.class.getName(),
+                Collections.emptyMap());
+        staticJaasConfig.setClientOptions("PLAIN", "user1", "user1-secret");
+        Configuration.setConfiguration(staticJaasConfig);
+        server = createEchoServer(securityProtocol);
+
+        // Check that 'user1' can connect with static Jaas config
+        createAndCheckClientConnection(securityProtocol, "1");
+
+        // Check that user 'user2' can also connect with a Jaas config override
+        saslClientConfigs.put(SaslConfigs.SASL_JAAS_CONFIG,
+                TestJaasConfig.jaasConfigProperty("PLAIN", "user2", "user2-secret"));
+        createAndCheckClientConnection(securityProtocol, "2");
+    }
+
     @Test
     public void testJaasConfigurationForListener() throws Exception {
         SecurityProtocol securityProtocol = SecurityProtocol.SASL_PLAINTEXT;
@@ -905,6 +1515,281 @@ public void oldSaslScramSslClientWithoutSaslAuthenticateHeaderFailure() throws E
         verifySaslAuthenticateHeaderInteropWithFailure(true, false, SecurityProtocol.SASL_SSL, "SCRAM-SHA-512");
     }
 
+    /**
+     * Tests OAUTHBEARER client and server channels.
+     */
+    @Test
+    public void testValidSaslOauthBearerMechanism() throws Exception {
+        String node = "0";
+        SecurityProtocol securityProtocol = SecurityProtocol.SASL_SSL;
+        configureMechanisms("OAUTHBEARER", Arrays.asList("OAUTHBEARER"));
+        server = createEchoServer(securityProtocol);
+        createAndCheckClientConnection(securityProtocol, node);
+    }
+    
+    /**
+     * Re-authentication must fail if principal changes
+     */
+    @Test
+    public void testCannotReauthenticateWithDifferentPrincipal() throws Exception {
+        String node = "0";
+        SecurityProtocol securityProtocol = SecurityProtocol.SASL_SSL;
+        saslClientConfigs.put(SaslConfigs.SASL_LOGIN_CALLBACK_HANDLER_CLASS,
+                AlternateLoginCallbackHandler.class.getName());
+        configureMechanisms(OAuthBearerLoginModule.OAUTHBEARER_MECHANISM,
+                Arrays.asList(OAuthBearerLoginModule.OAUTHBEARER_MECHANISM));
+        server = createEchoServer(securityProtocol);
+        // initial authentication must succeed
+        createClientConnection(securityProtocol, node);
+        checkClientConnection(node);
+        // ensure metrics are as expected before trying to re-authenticate
+        server.verifyAuthenticationMetrics(1, 0);
+        server.verifyReauthenticationMetrics(0, 0);
+        /*
+         * Now re-authenticate with a different principal and ensure it fails. We first
+         * have to sleep long enough for the background refresh thread to replace the
+         * original token with a new one.
+         */
+        delay(1000L);
+        try {
+            checkClientConnection(node);
+            fail("Re-authentication with a different principal should have failed but did not");
+        } catch (AssertionError e) {
+            // ignore, expected
+            server.verifyReauthenticationMetrics(0, 1);
+        }
+    }
+
+    @Test
+    public void testCorrelationId() {
+        SaslClientAuthenticator authenticator = new SaslClientAuthenticator(
+              Collections.emptyMap(),
+              null,
+              "node",
+              null,
+              null,
+              null,
+              "plain",
+              false,
+              null,
+              null,
+            new LogContext()
+        ) {
+            @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() {
+        ListOffsetResponseData data = new ListOffsetResponseData()
+                .setThrottleTimeMs(0)
+                .setTopics(Collections.singletonList(new ListOffsetTopicResponse()
+                        .setName("topic")
+                        .setPartitions(Collections.singletonList(new ListOffsetPartitionResponse()
+                                .setErrorCode(Errors.NONE.code())
+                                .setLeaderEpoch(ListOffsetResponse.UNKNOWN_EPOCH)
+                                .setPartitionIndex(0)
+                                .setOffset(0)
+                                .setTimestamp(0)))));
+        ListOffsetResponse response = new ListOffsetResponse(data);
+        ByteBuffer buffer = RequestTestUtils.serializeResponseWithHeader(response, 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
+     */
+    @Test
+    public void testCannotReauthenticateWithDifferentMechanism() throws Exception {
+        String node = "0";
+        SecurityProtocol securityProtocol = SecurityProtocol.SASL_SSL;
+        configureMechanisms("DIGEST-MD5", Arrays.asList("DIGEST-MD5", "PLAIN"));
+        configureDigestMd5ServerCallback(securityProtocol);
+        server = createEchoServer(securityProtocol);
+
+        String saslMechanism = (String) saslClientConfigs.get(SaslConfigs.SASL_MECHANISM);
+        Map configs = new TestSecurityConfig(saslClientConfigs).values();
+        this.channelBuilder = new AlternateSaslChannelBuilder(Mode.CLIENT,
+                Collections.singletonMap(saslMechanism, JaasContext.loadClientContext(configs)), securityProtocol, null,
+                false, saslMechanism, true, credentialCache, null, time);
+        this.channelBuilder.configure(configs);
+        // initial authentication must succeed
+        this.selector = NetworkTestUtils.createSelector(channelBuilder, time);
+        InetSocketAddress addr = new InetSocketAddress("localhost", server.port());
+        selector.connect(node, addr, BUFFER_SIZE, BUFFER_SIZE);
+        checkClientConnection(node);
+        // ensure metrics are as expected before trying to re-authenticate
+        server.verifyAuthenticationMetrics(1, 0);
+        server.verifyReauthenticationMetrics(0, 0);
+        /*
+         * Now re-authenticate with a different mechanism and ensure it fails. We have
+         * to sleep long enough so that the next write will trigger a re-authentication.
+         */
+        delay((long) (CONNECTIONS_MAX_REAUTH_MS_VALUE * 1.1));
+        try {
+            checkClientConnection(node);
+            fail("Re-authentication with a different mechanism should have failed but did not");
+        } catch (AssertionError e) {
+            // ignore, expected
+            server.verifyAuthenticationMetrics(1, 0);
+            server.verifyReauthenticationMetrics(0, 1);
+        }
+    }
+
+    /**
+     * Second re-authentication must fail if it is sooner than one second after the first
+     */
+    @Test
+    public void testCannotReauthenticateAgainFasterThanOneSecond() throws Exception {
+        String node = "0";
+        time = new MockTime();
+        SecurityProtocol securityProtocol = SecurityProtocol.SASL_SSL;
+        configureMechanisms(OAuthBearerLoginModule.OAUTHBEARER_MECHANISM,
+                Arrays.asList(OAuthBearerLoginModule.OAUTHBEARER_MECHANISM));
+        server = createEchoServer(securityProtocol);
+        try {
+            createClientConnection(securityProtocol, node);
+            checkClientConnection(node);
+            server.verifyAuthenticationMetrics(1, 0);
+            server.verifyReauthenticationMetrics(0, 0);
+            /*
+             * Now sleep long enough so that the next write will cause re-authentication,
+             * which we expect to succeed.
+             */
+            time.sleep((long) (CONNECTIONS_MAX_REAUTH_MS_VALUE * 1.1));
+            checkClientConnection(node);
+            server.verifyAuthenticationMetrics(1, 0);
+            server.verifyReauthenticationMetrics(1, 0);
+            /*
+             * Now sleep long enough so that the next write will cause re-authentication,
+             * but this time we expect re-authentication to not occur since it has been too
+             * soon. The checkClientConnection() call should return an error saying it
+             * expected the one byte-plus-node response but got the SaslHandshakeRequest
+             * instead
+             */
+            time.sleep((long) (CONNECTIONS_MAX_REAUTH_MS_VALUE * 1.1));
+            AssertionError exception = assertThrows(AssertionError.class,
+                () -> NetworkTestUtils.checkClientConnection(selector, node, 1, 1));
+            String expectedResponseTextRegex = "\\w-" + node;
+            String receivedResponseTextRegex = ".*" + OAuthBearerLoginModule.OAUTHBEARER_MECHANISM;
+            assertTrue(
+                    "Should have received the SaslHandshakeRequest bytes back since we re-authenticated too quickly, " +
+                    "but instead we got our generated message echoed back, implying re-auth succeeded when it " +
+                    "should not have: " + exception,
+                    exception.getMessage().matches(
+                            ".*<\\[" + expectedResponseTextRegex + "]>.*<\\[" + receivedResponseTextRegex + ".*?]>"));
+            server.verifyReauthenticationMetrics(1, 0); // unchanged
+        } finally {
+            selector.close();
+            selector = null;
+        }
+    }
+
+    /**
+     * Tests good path SASL/PLAIN client and server channels using SSL transport layer.
+     * Repeatedly tests successful re-authentication over several seconds.
+     */
+    @Test
+    public void testRepeatedValidSaslPlainOverSsl() throws Exception {
+        String node = "0";
+        SecurityProtocol securityProtocol = SecurityProtocol.SASL_SSL;
+        configureMechanisms("PLAIN", Arrays.asList("PLAIN"));
+        /*
+         * Make sure 85% of this value is at least 1 second otherwise it is possible for
+         * the client to start re-authenticating but the server does not start due to
+         * the 1-second minimum. If this happens the SASL HANDSHAKE request that was
+         * injected to start re-authentication will be echoed back to the client instead
+         * of the data that the client explicitly sent, and then the client will not
+         * recognize that data and will throw an assertion error.
+         */
+        saslServerConfigs.put(BrokerSecurityConfigs.CONNECTIONS_MAX_REAUTH_MS,
+                Double.valueOf(1.1 * 1000L / 0.85).longValue());
+
+        server = createEchoServer(securityProtocol);
+        createClientConnection(securityProtocol, node);
+        checkClientConnection(node);
+        server.verifyAuthenticationMetrics(1, 0);
+        server.verifyReauthenticationMetrics(0, 0);
+        double successfulReauthentications = 0;
+        int desiredNumReauthentications = 5;
+        long startMs = Time.SYSTEM.milliseconds();
+        long timeoutMs = startMs + 1000 * 15; // stop after 15 seconds
+        while (successfulReauthentications < desiredNumReauthentications
+                && Time.SYSTEM.milliseconds() < timeoutMs) {
+            checkClientConnection(node);
+            successfulReauthentications = server.metricValue("successful-reauthentication-total");
+        }
+        server.verifyReauthenticationMetrics(desiredNumReauthentications, 0);
+    }
+    
+    /**
+     * Tests OAUTHBEARER client channels without tokens for the server.
+     */
+    @Test
+    public void testValidSaslOauthBearerMechanismWithoutServerTokens() throws Exception {
+        String node = "0";
+        SecurityProtocol securityProtocol = SecurityProtocol.SASL_SSL;
+        saslClientConfigs.put(SaslConfigs.SASL_MECHANISM, "OAUTHBEARER");
+        saslServerConfigs.put(BrokerSecurityConfigs.SASL_ENABLED_MECHANISMS_CONFIG, Arrays.asList("OAUTHBEARER"));
+        saslClientConfigs.put(SaslConfigs.SASL_JAAS_CONFIG,
+                TestJaasConfig.jaasConfigProperty("OAUTHBEARER", Collections.singletonMap("unsecuredLoginStringClaim_sub", TestJaasConfig.USERNAME)));
+        saslServerConfigs.put("listener.name.sasl_ssl.oauthbearer." + SaslConfigs.SASL_JAAS_CONFIG,
+                TestJaasConfig.jaasConfigProperty("OAUTHBEARER", Collections.emptyMap()));
+
+        // Server without a token should start up successfully and authenticate clients.
+        server = createEchoServer(securityProtocol);
+        createAndCheckClientConnection(securityProtocol, node);
+
+        // Client without a token should fail to connect
+        saslClientConfigs.put(SaslConfigs.SASL_JAAS_CONFIG,
+                TestJaasConfig.jaasConfigProperty("OAUTHBEARER", Collections.emptyMap()));
+        createAndCheckClientConnectionFailure(securityProtocol, node);
+
+        // Server with extensions, but without a token should fail to start up since it could indicate a configuration error
+        saslServerConfigs.put("listener.name.sasl_ssl.oauthbearer." + SaslConfigs.SASL_JAAS_CONFIG,
+                TestJaasConfig.jaasConfigProperty("OAUTHBEARER", Collections.singletonMap("unsecuredLoginExtension_test", "something")));
+        try {
+            createEchoServer(securityProtocol);
+            fail("Server created with invalid login config containing extensions without a token");
+        } catch (Throwable e) {
+            assertTrue("Unexpected exception " + Utils.stackTrace(e), e.getCause() instanceof LoginException);
+        }
+    }
+
+    /**
+     * Tests OAUTHBEARER fails the connection when the client presents a token with
+     * insufficient scope .
+     */
+    @Test
+    public void testInsufficientScopeSaslOauthBearerMechanism() throws Exception {
+        SecurityProtocol securityProtocol = SecurityProtocol.SASL_SSL;
+        TestJaasConfig jaasConfig = configureMechanisms("OAUTHBEARER", Arrays.asList("OAUTHBEARER"));
+        // now update the server side to require a scope the client does not provide
+        Map serverJaasConfigOptionsMap = TestJaasConfig.defaultServerOptions("OAUTHBEARER");
+        serverJaasConfigOptionsMap.put("unsecuredValidatorRequiredScope", "LOGIN_TO_KAFKA"); // causes the failure
+        jaasConfig.createOrUpdateEntry("KafkaServer",
+                "org.apache.kafka.common.security.oauthbearer.OAuthBearerLoginModule", serverJaasConfigOptionsMap);
+        server = createEchoServer(securityProtocol);
+        createAndCheckClientAuthenticationFailure(securityProtocol,
+                "node-" + OAuthBearerLoginModule.OAUTHBEARER_MECHANISM, OAuthBearerLoginModule.OAUTHBEARER_MECHANISM,
+                "{\"status\":\"insufficient_scope\", \"scope\":\"[LOGIN_TO_KAFKA]\"}");
+    }
+
     private void verifySaslAuthenticateHeaderInterop(boolean enableHeaderOnServer, boolean enableHeaderOnClient,
             SecurityProtocol securityProtocol, String saslMechanism) throws Exception {
         configureMechanisms(saslMechanism, Arrays.asList(saslMechanism));
@@ -926,7 +1811,7 @@ private void verifySaslAuthenticateHeaderInteropWithFailure(boolean enableHeader
         // Without SASL_AUTHENTICATE headers, disconnect state is ChannelState.AUTHENTICATE which is
         // a hint that channel was closed during authentication, unlike ChannelState.AUTHENTICATE_FAILED
         // which is an actual authentication failure reported by the broker.
-        NetworkTestUtils.waitForChannelClose(selector, node, ChannelState.AUTHENTICATE.state());
+        NetworkTestUtils.waitForChannelClose(selector, node, ChannelState.State.AUTHENTICATE);
     }
 
     private void createServer(SecurityProtocol securityProtocol, String saslMechanism,
@@ -946,35 +1831,93 @@ private void createClientConnection(SecurityProtocol securityProtocol, String sa
             createClientConnectionWithoutSaslAuthenticateHeader(securityProtocol, saslMechanism, node);
     }
 
+    private NioEchoServer startServerApiVersionsUnsupportedByClient(final SecurityProtocol securityProtocol, String saslMechanism) throws Exception {
+        final ListenerName listenerName = ListenerName.forSecurityProtocol(securityProtocol);
+        final Map configs = Collections.emptyMap();
+        final JaasContext jaasContext = JaasContext.loadServerContext(listenerName, saslMechanism, configs);
+        final Map jaasContexts = Collections.singletonMap(saslMechanism, jaasContext);
+
+        boolean isScram = ScramMechanism.isScram(saslMechanism);
+        if (isScram)
+            ScramCredentialUtils.createCache(credentialCache, Arrays.asList(saslMechanism));
+        SaslChannelBuilder serverChannelBuilder = new SaslChannelBuilder(Mode.SERVER, jaasContexts,
+                securityProtocol, listenerName, false, saslMechanism, true,
+                credentialCache, null, time, new LogContext()) {
+
+            @Override
+            protected SaslServerAuthenticator buildServerAuthenticator(Map configs,
+                                                                       Map callbackHandlers,
+                                                                       String id,
+                                                                       TransportLayer transportLayer,
+                                                                       Map subjects,
+                                                                       Map connectionsMaxReauthMsByMechanism,
+                                                                       ChannelMetadataRegistry metadataRegistry) {
+                return new SaslServerAuthenticator(configs, callbackHandlers, id, subjects, null, listenerName,
+                        securityProtocol, transportLayer, connectionsMaxReauthMsByMechanism, metadataRegistry, time) {
+
+                    @Override
+                    protected ApiVersionsResponse apiVersionsResponse() {
+                        ApiVersionsResponseKeyCollection versionCollection = new ApiVersionsResponseKeyCollection(2);
+                        versionCollection.add(new ApiVersionsResponseKey().setApiKey(ApiKeys.SASL_HANDSHAKE.id).setMinVersion((short) 0).setMaxVersion((short) 100));
+                        versionCollection.add(new ApiVersionsResponseKey().setApiKey(ApiKeys.SASL_AUTHENTICATE.id).setMinVersion((short) 0).setMaxVersion((short) 100));
+                        return new ApiVersionsResponse(new ApiVersionsResponseData().setApiKeys(versionCollection));
+                    }
+                };
+            }
+        };
+        serverChannelBuilder.configure(saslServerConfigs);
+        server = new NioEchoServer(listenerName, securityProtocol, new TestSecurityConfig(saslServerConfigs),
+                "localhost", serverChannelBuilder, credentialCache, time);
+        server.start();
+        return server;
+    }
+
     private NioEchoServer startServerWithoutSaslAuthenticateHeader(final SecurityProtocol securityProtocol, String saslMechanism)
             throws Exception {
         final ListenerName listenerName = ListenerName.forSecurityProtocol(securityProtocol);
         final Map configs = Collections.emptyMap();
-        final JaasContext jaasContext = JaasContext.load(JaasContext.Type.SERVER, listenerName, configs);
+        final JaasContext jaasContext = JaasContext.loadServerContext(listenerName, saslMechanism, configs);
+        final Map jaasContexts = Collections.singletonMap(saslMechanism, jaasContext);
 
         boolean isScram = ScramMechanism.isScram(saslMechanism);
         if (isScram)
             ScramCredentialUtils.createCache(credentialCache, Arrays.asList(saslMechanism));
-        SaslChannelBuilder serverChannelBuilder = new SaslChannelBuilder(Mode.SERVER, jaasContext,
-                securityProtocol, listenerName, saslMechanism, true, credentialCache) {
+        SaslChannelBuilder serverChannelBuilder = new SaslChannelBuilder(Mode.SERVER, jaasContexts,
+                securityProtocol, listenerName, false, saslMechanism, true,
+                credentialCache, null, time, new LogContext()) {
 
             @Override
-            protected SaslServerAuthenticator buildServerAuthenticator(Map configs, String id,
-                            TransportLayer transportLayer, Subject subject) throws IOException {
-                return new SaslServerAuthenticator(configs, id, jaasContext, subject, null,
-                                credentialCache, listenerName, securityProtocol, transportLayer) {
+            protected SaslServerAuthenticator buildServerAuthenticator(Map configs,
+                                                                       Map callbackHandlers,
+                                                                       String id,
+                                                                       TransportLayer transportLayer,
+                                                                       Map subjects,
+                                                                       Map connectionsMaxReauthMsByMechanism,
+                                                                       ChannelMetadataRegistry metadataRegistry) {
+                return new SaslServerAuthenticator(configs, callbackHandlers, id, subjects, null, listenerName,
+                    securityProtocol, transportLayer, connectionsMaxReauthMsByMechanism, metadataRegistry, time) {
 
                     @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
@@ -986,7 +1929,7 @@ protected void enableKafkaSaslAuthenticateHeaders(boolean flag) {
         };
         serverChannelBuilder.configure(saslServerConfigs);
         server = new NioEchoServer(listenerName, securityProtocol, new TestSecurityConfig(saslServerConfigs),
-                "localhost", serverChannelBuilder, credentialCache);
+                "localhost", serverChannelBuilder, credentialCache, time);
         server.start();
         return server;
     }
@@ -996,30 +1939,38 @@ private void createClientConnectionWithoutSaslAuthenticateHeader(final SecurityP
 
         final ListenerName listenerName = ListenerName.forSecurityProtocol(securityProtocol);
         final Map configs = Collections.emptyMap();
-        final JaasContext jaasContext = JaasContext.load(JaasContext.Type.CLIENT, null, configs);
-        SaslChannelBuilder clientChannelBuilder = new SaslChannelBuilder(Mode.CLIENT, jaasContext,
-                securityProtocol, listenerName, saslMechanism, true, null) {
+        final JaasContext jaasContext = JaasContext.loadClientContext(configs);
+        final Map jaasContexts = Collections.singletonMap(saslMechanism, jaasContext);
 
-            @Override
-            protected SaslClientAuthenticator buildClientAuthenticator(Map configs, String id,
-                    String serverHost, String servicePrincipal,
-                    TransportLayer transportLayer, Subject subject) throws IOException {
+        SaslChannelBuilder clientChannelBuilder = new SaslChannelBuilder(Mode.CLIENT, jaasContexts,
+                securityProtocol, listenerName, false, saslMechanism, true,
+                null, null, time, new LogContext()) {
 
-                return new SaslClientAuthenticator(configs, id, subject,
-                        servicePrincipal, serverHost, saslMechanism, true, transportLayer) {
+            @Override
+            protected SaslClientAuthenticator buildClientAuthenticator(Map configs,
+                                                                       AuthenticateCallbackHandler callbackHandler,
+                                                                       String id,
+                                                                       String serverHost,
+                                                                       String servicePrincipal,
+                                                                       TransportLayer transportLayer,
+                                                                       Subject subject) {
+
+                return new SaslClientAuthenticator(configs, callbackHandler, id, subject,
+                        servicePrincipal, serverHost, saslMechanism, true,
+                        transportLayer, time, new LogContext()) {
                     @Override
                     protected SaslHandshakeRequest createSaslHandshakeRequest(short version) {
-                        return new SaslHandshakeRequest.Builder(saslMechanism).build((short) 0);
+                        return buildSaslHandshakeRequest(saslMechanism, (short) 0);
                     }
                     @Override
-                    protected void saslAuthenticateVersion(short version) {
+                    protected void setSaslAuthenticateAndHandshakeVersions(ApiVersionsResponse apiVersionsResponse) {
                         // Don't set version so that headers are disabled
                     }
                 };
             }
         };
         clientChannelBuilder.configure(saslClientConfigs);
-        this.selector = NetworkTestUtils.createSelector(clientChannelBuilder);
+        this.selector = NetworkTestUtils.createSelector(clientChannelBuilder, time);
         InetSocketAddress addr = new InetSocketAddress("127.0.0.1", server.port());
         selector.connect(node, addr, BUFFER_SIZE, BUFFER_SIZE);
     }
@@ -1070,10 +2021,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);
@@ -1088,10 +2039,11 @@ private void authenticateUsingSaslPlainAndCheckConnection(String node, boolean e
         String authString = "\u0000" + TestJaasConfig.USERNAME + "\u0000" + TestJaasConfig.PASSWORD;
         ByteBuffer authBuf = ByteBuffer.wrap(authString.getBytes("UTF-8"));
         if (enableSaslAuthenticateHeader) {
-            SaslAuthenticateRequest request = new SaslAuthenticateRequest.Builder(authBuf).build();
+            SaslAuthenticateRequestData data = new SaslAuthenticateRequestData().setAuthBytes(authBuf.array());
+            SaslAuthenticateRequest request = new SaslAuthenticateRequest.Builder(data).build();
             sendKafkaRequestReceiveResponse(node, ApiKeys.SASL_AUTHENTICATE, request);
         } else {
-            selector.send(new NetworkSend(node, authBuf));
+            selector.send(new NetworkSend(node, ByteBufferSend.sizePrefixed(authBuf)));
             waitForResponse();
         }
 
@@ -1102,9 +2054,20 @@ private void authenticateUsingSaslPlainAndCheckConnection(String node, boolean e
     private TestJaasConfig configureMechanisms(String clientMechanism, List serverMechanisms) {
         saslClientConfigs.put(SaslConfigs.SASL_MECHANISM, clientMechanism);
         saslServerConfigs.put(BrokerSecurityConfigs.SASL_ENABLED_MECHANISMS_CONFIG, serverMechanisms);
+        saslServerConfigs.put(BrokerSecurityConfigs.CONNECTIONS_MAX_REAUTH_MS, CONNECTIONS_MAX_REAUTH_MS_VALUE);
+        if (serverMechanisms.contains("DIGEST-MD5")) {
+            saslServerConfigs.put("digest-md5." + BrokerSecurityConfigs.SASL_SERVER_CALLBACK_HANDLER_CLASS,
+                    TestDigestLoginModule.DigestServerCallbackHandler.class.getName());
+        }
         return TestJaasConfig.createConfiguration(clientMechanism, serverMechanisms);
     }
 
+    private void configureDigestMd5ServerCallback(SecurityProtocol securityProtocol) {
+        String callbackPrefix = ListenerName.forSecurityProtocol(securityProtocol).saslMechanismConfigPrefix("DIGEST-MD5");
+        saslServerConfigs.put(callbackPrefix + BrokerSecurityConfigs.SASL_SERVER_CALLBACK_HANDLER_CLASS,
+                TestDigestLoginModule.DigestServerCallbackHandler.class);
+    }
+
     private void createSelector(SecurityProtocol securityProtocol, Map clientConfigs) {
         if (selector != null) {
             selector.close();
@@ -1113,8 +2076,9 @@ private void createSelector(SecurityProtocol securityProtocol, Map 0);
         assertEquals(1, selector.completedReceives().size());
-        return selector.completedReceives().get(0).payload();
+        return selector.completedReceives().iterator().next().payload();
+    }
+
+    public static class TestServerCallbackHandler extends PlainServerCallbackHandler {
+
+        static final String USERNAME = "TestServerCallbackHandler-user";
+        static final String PASSWORD = "TestServerCallbackHandler-password";
+        private volatile boolean configured;
+
+        @Override
+        public void configure(Map configs, String mechanism, List jaasConfigEntries) {
+            if (configured)
+                throw new IllegalStateException("Server callback handler configured twice");
+            configured = true;
+            super.configure(configs, mechanism, jaasConfigEntries);
+        }
+
+        @Override
+        protected boolean authenticate(String username, char[] password) {
+            if (!configured)
+                throw new IllegalStateException("Server callback handler not configured");
+            return USERNAME.equals(username) && new String(password).equals(PASSWORD);
+        }
+    }
+
+    private SaslHandshakeRequest buildSaslHandshakeRequest(String mechanism, short version) {
+        return new SaslHandshakeRequest.Builder(
+                new SaslHandshakeRequestData().setMechanism(mechanism)).build(version);
     }
 
     @SuppressWarnings("unchecked")
@@ -1207,4 +2240,250 @@ private void updateScramCredentialCache(String username, String password) throws
     private ApiVersionsRequest createApiVersionsRequestV0() {
         return new ApiVersionsRequest.Builder((short) 0).build();
     }
+
+    @SuppressWarnings("unchecked")
+    private void updateTokenCredentialCache(String username, String password) throws NoSuchAlgorithmException {
+        for (String mechanism : (List) saslServerConfigs.get(BrokerSecurityConfigs.SASL_ENABLED_MECHANISMS_CONFIG)) {
+            ScramMechanism scramMechanism = ScramMechanism.forMechanismName(mechanism);
+            if (scramMechanism != null) {
+                ScramFormatter formatter = new ScramFormatter(scramMechanism);
+                ScramCredential credential = formatter.generateCredential(password, 4096);
+                server.tokenCache().credentialCache(scramMechanism.mechanismName()).put(username, credential);
+            }
+        }
+    }
+
+    private static void delay(long delayMillis) throws InterruptedException {
+        final long startTime = System.currentTimeMillis();
+        while ((System.currentTimeMillis() - startTime) < delayMillis)
+            Thread.sleep(CONNECTIONS_MAX_REAUTH_MS_VALUE / 5);
+    }
+
+    public static class TestClientCallbackHandler implements AuthenticateCallbackHandler {
+
+        static final String USERNAME = "TestClientCallbackHandler-user";
+        static final String PASSWORD = "TestClientCallbackHandler-password";
+        private volatile boolean configured;
+
+        @Override
+        public void configure(Map configs, String mechanism, List jaasConfigEntries) {
+            if (configured)
+                throw new IllegalStateException("Client callback handler configured twice");
+            configured = true;
+        }
+
+        @Override
+        public void handle(Callback[] callbacks) throws UnsupportedCallbackException {
+            if (!configured)
+                throw new IllegalStateException("Client callback handler not configured");
+            for (Callback callback : callbacks) {
+                if (callback instanceof NameCallback)
+                    ((NameCallback) callback).setName(USERNAME);
+                else if (callback instanceof PasswordCallback)
+                    ((PasswordCallback) callback).setPassword(PASSWORD.toCharArray());
+                else
+                    throw new UnsupportedCallbackException(callback);
+            }
+        }
+
+        @Override
+        public void close() {
+        }
+    }
+
+    public static class TestLogin implements Login {
+
+        static AtomicInteger loginCount = new AtomicInteger();
+
+        private String contextName;
+        private Configuration configuration;
+        private Subject subject;
+        @Override
+        public void configure(Map configs, String contextName, Configuration configuration,
+                              AuthenticateCallbackHandler callbackHandler) {
+            assertEquals(1, configuration.getAppConfigurationEntry(contextName).length);
+            this.contextName = contextName;
+            this.configuration = configuration;
+        }
+
+        @Override
+        public LoginContext login() throws LoginException {
+            LoginContext context = new LoginContext(contextName, null, new AbstractLogin.DefaultLoginCallbackHandler(), configuration);
+            context.login();
+            subject = context.getSubject();
+            subject.getPublicCredentials().clear();
+            subject.getPrivateCredentials().clear();
+            subject.getPublicCredentials().add(TestJaasConfig.USERNAME);
+            subject.getPrivateCredentials().add(TestJaasConfig.PASSWORD);
+            loginCount.incrementAndGet();
+            return context;
+        }
+
+        @Override
+        public Subject subject() {
+            return subject;
+        }
+
+        @Override
+        public String serviceName() {
+            return "kafka";
+        }
+
+        @Override
+        public void close() {
+        }
+    }
+
+    public static class TestLoginCallbackHandler implements AuthenticateCallbackHandler {
+        private volatile boolean configured = false;
+        @Override
+        public void configure(Map configs, String saslMechanism, List jaasConfigEntries) {
+            if (configured)
+                throw new IllegalStateException("Login callback handler configured twice");
+            configured = true;
+        }
+
+        @Override
+        public void handle(Callback[] callbacks) {
+            if (!configured)
+                throw new IllegalStateException("Login callback handler not configured");
+
+            for (Callback callback : callbacks) {
+                if (callback instanceof NameCallback)
+                    ((NameCallback) callback).setName(TestJaasConfig.USERNAME);
+                else if (callback instanceof PasswordCallback)
+                    ((PasswordCallback) callback).setPassword(TestJaasConfig.PASSWORD.toCharArray());
+            }
+        }
+
+        @Override
+        public void close() {
+        }
+    }
+
+    public static final class TestPlainLoginModule extends PlainLoginModule {
+        @Override
+        public void initialize(Subject subject, CallbackHandler callbackHandler, Map sharedState, Map options) {
+            try {
+                NameCallback nameCallback = new NameCallback("name:");
+                PasswordCallback passwordCallback = new PasswordCallback("password:", false);
+                callbackHandler.handle(new Callback[]{nameCallback, passwordCallback});
+                subject.getPublicCredentials().add(nameCallback.getName());
+                subject.getPrivateCredentials().add(new String(passwordCallback.getPassword()));
+            } catch (Exception e) {
+                throw new SaslAuthenticationException("Login initialization failed", e);
+            }
+        }
+    }
+
+    /*
+     * Create an alternate login callback handler that continually returns a
+     * different principal
+     */
+    public static class AlternateLoginCallbackHandler implements AuthenticateCallbackHandler {
+        private static final OAuthBearerUnsecuredLoginCallbackHandler DELEGATE = new OAuthBearerUnsecuredLoginCallbackHandler();
+        private static final String QUOTE = "\"";
+        private static int numInvocations = 0;
+
+        @Override
+        public void handle(Callback[] callbacks) throws IOException, UnsupportedCallbackException {
+            DELEGATE.handle(callbacks);
+            // now change any returned token to have a different principal name
+            if (callbacks.length > 0)
+                for (Callback callback : callbacks) {
+                    if (callback instanceof OAuthBearerTokenCallback) {
+                        OAuthBearerTokenCallback oauthBearerTokenCallback = (OAuthBearerTokenCallback) callback;
+                        OAuthBearerToken token = oauthBearerTokenCallback.token();
+                        if (token != null) {
+                            String changedPrincipalNameToUse = token.principalName()
+                                    + String.valueOf(++numInvocations);
+                            String headerJson = "{" + claimOrHeaderJsonText("alg", "none") + "}";
+                            /*
+                             * Use a short lifetime so the background refresh thread replaces it before we
+                             * re-authenticate
+                             */
+                            String lifetimeSecondsValueToUse = "1";
+                            String claimsJson;
+                            try {
+                                claimsJson = String.format("{%s,%s,%s}",
+                                        expClaimText(Long.parseLong(lifetimeSecondsValueToUse)),
+                                        claimOrHeaderJsonText("iat", time.milliseconds() / 1000.0),
+                                        claimOrHeaderJsonText("sub", changedPrincipalNameToUse));
+                            } catch (NumberFormatException e) {
+                                throw new OAuthBearerConfigException(e.getMessage());
+                            }
+                            try {
+                                Encoder urlEncoderNoPadding = Base64.getUrlEncoder().withoutPadding();
+                                OAuthBearerUnsecuredJws jws = new OAuthBearerUnsecuredJws(String.format("%s.%s.",
+                                        urlEncoderNoPadding.encodeToString(headerJson.getBytes(StandardCharsets.UTF_8)),
+                                        urlEncoderNoPadding
+                                                .encodeToString(claimsJson.getBytes(StandardCharsets.UTF_8))),
+                                        "sub", "scope");
+                                oauthBearerTokenCallback.token(jws);
+                            } catch (OAuthBearerIllegalTokenException e) {
+                                // occurs if the principal claim doesn't exist or has an empty value
+                                throw new OAuthBearerConfigException(e.getMessage(), e);
+                            }
+                        }
+                    }
+                }
+        }
+
+        private static String claimOrHeaderJsonText(String claimName, String claimValue) {
+            return QUOTE + claimName + QUOTE + ":" + QUOTE + claimValue + QUOTE;
+        }
+
+        private static String claimOrHeaderJsonText(String claimName, Number claimValue) {
+            return QUOTE + claimName + QUOTE + ":" + claimValue;
+        }
+
+        private static String expClaimText(long lifetimeSeconds) {
+            return claimOrHeaderJsonText("exp", time.milliseconds() / 1000.0 + lifetimeSeconds);
+        }
+
+        @Override
+        public void configure(Map configs, String saslMechanism,
+                List jaasConfigEntries) {
+            DELEGATE.configure(configs, saslMechanism, jaasConfigEntries);
+        }
+
+        @Override
+        public void close() {
+            DELEGATE.close();
+        }
+    }
+
+    /*
+     * Define a channel builder that starts with the DIGEST-MD5 mechanism and then
+     * switches to the PLAIN mechanism
+     */
+    private static class AlternateSaslChannelBuilder extends SaslChannelBuilder {
+        private int numInvocations = 0;
+
+        public AlternateSaslChannelBuilder(Mode mode, Map jaasContexts,
+                SecurityProtocol securityProtocol, ListenerName listenerName, boolean isInterBrokerListener,
+                String clientSaslMechanism, boolean handshakeRequestEnable, CredentialCache credentialCache,
+                DelegationTokenCache tokenCache, Time time) {
+            super(mode, jaasContexts, securityProtocol, listenerName, isInterBrokerListener, clientSaslMechanism,
+                    handshakeRequestEnable, credentialCache, tokenCache, time, new LogContext());
+        }
+
+        @Override
+        protected SaslClientAuthenticator buildClientAuthenticator(Map configs,
+                AuthenticateCallbackHandler callbackHandler, String id, String serverHost, String servicePrincipal,
+                TransportLayer transportLayer, Subject subject) {
+            if (++numInvocations == 1)
+                return new SaslClientAuthenticator(configs, callbackHandler, id, subject, servicePrincipal, serverHost,
+                        "DIGEST-MD5", true, transportLayer, time, new LogContext());
+            else
+                return new SaslClientAuthenticator(configs, callbackHandler, id, subject, servicePrincipal, serverHost,
+                        "PLAIN", true, transportLayer, time, new LogContext()) {
+                    @Override
+                    protected SaslHandshakeRequest createSaslHandshakeRequest(short version) {
+                        return new SaslHandshakeRequest.Builder(
+                                new SaslHandshakeRequestData().setMechanism("PLAIN")).build(version);
+                    }
+                };
+        }
+    }
 }
diff --git a/clients/src/test/java/org/apache/kafka/common/security/authenticator/SaslServerAuthenticatorTest.java b/clients/src/test/java/org/apache/kafka/common/security/authenticator/SaslServerAuthenticatorTest.java
index 022a099bef5e2..20853b6e63c8c 100644
--- a/clients/src/test/java/org/apache/kafka/common/security/authenticator/SaslServerAuthenticatorTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/security/authenticator/SaslServerAuthenticatorTest.java
@@ -16,20 +16,24 @@
  */
 package org.apache.kafka.common.security.authenticator;
 
+import java.net.InetAddress;
 import org.apache.kafka.common.config.internals.BrokerSecurityConfigs;
 import org.apache.kafka.common.errors.IllegalSaslStateException;
+import org.apache.kafka.common.network.ChannelMetadataRegistry;
+import org.apache.kafka.common.network.ClientInformation;
+import org.apache.kafka.common.network.DefaultChannelMetadataRegistry;
 import org.apache.kafka.common.network.InvalidReceiveException;
 import org.apache.kafka.common.network.ListenerName;
 import org.apache.kafka.common.network.TransportLayer;
 import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.requests.ApiVersionsRequest;
+import org.apache.kafka.common.requests.RequestTestUtils;
+import org.apache.kafka.common.security.auth.AuthenticateCallbackHandler;
 import org.apache.kafka.common.security.auth.SecurityProtocol;
-import org.apache.kafka.common.protocol.types.Struct;
 import org.apache.kafka.common.requests.RequestHeader;
-import org.apache.kafka.common.security.JaasContext;
 import org.apache.kafka.common.security.plain.PlainLoginModule;
-import org.easymock.Capture;
-import org.easymock.EasyMock;
-import org.easymock.IAnswer;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.apache.kafka.common.utils.Time;
 import org.junit.Test;
 
 import javax.security.auth.Subject;
@@ -38,79 +42,123 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
+import org.mockito.Answers;
 
-import static org.apache.kafka.common.security.scram.ScramMechanism.SCRAM_SHA_256;
+import static org.apache.kafka.common.security.scram.internals.ScramMechanism.SCRAM_SHA_256;
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.fail;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 public class SaslServerAuthenticatorTest {
 
     @Test(expected = InvalidReceiveException.class)
     public void testOversizeRequest() throws IOException {
-        TransportLayer transportLayer = EasyMock.mock(TransportLayer.class);
+        TransportLayer transportLayer = mock(TransportLayer.class);
         Map configs = Collections.singletonMap(BrokerSecurityConfigs.SASL_ENABLED_MECHANISMS_CONFIG,
                 Collections.singletonList(SCRAM_SHA_256.mechanismName()));
-        SaslServerAuthenticator authenticator = setupAuthenticator(configs, transportLayer);
-
-        final Capture size = EasyMock.newCapture();
-        EasyMock.expect(transportLayer.read(EasyMock.capture(size))).andAnswer(new IAnswer() {
-            @Override
-            public Integer answer() throws Throwable {
-                size.getValue().putInt(SaslServerAuthenticator.MAX_RECEIVE_SIZE + 1);
-                return 4;
-            }
-        });
-
-        EasyMock.replay(transportLayer);
+        SaslServerAuthenticator authenticator = setupAuthenticator(configs, transportLayer,
+            SCRAM_SHA_256.mechanismName(), new DefaultChannelMetadataRegistry());
 
+        when(transportLayer.read(any(ByteBuffer.class))).then(invocation -> {
+            invocation.getArgument(0).putInt(SaslServerAuthenticator.MAX_RECEIVE_SIZE + 1);
+            return 4;
+        });
         authenticator.authenticate();
+        verify(transportLayer).read(any(ByteBuffer.class));
     }
 
     @Test
     public void testUnexpectedRequestType() throws IOException {
-        TransportLayer transportLayer = EasyMock.mock(TransportLayer.class);
+        TransportLayer transportLayer = mock(TransportLayer.class);
         Map configs = Collections.singletonMap(BrokerSecurityConfigs.SASL_ENABLED_MECHANISMS_CONFIG,
                 Collections.singletonList(SCRAM_SHA_256.mechanismName()));
-        SaslServerAuthenticator authenticator = setupAuthenticator(configs, transportLayer);
-
-        final RequestHeader header = new RequestHeader(ApiKeys.METADATA, (short) 0, "clientId", 13243);
-        final Struct headerStruct = header.toStruct();
-
-        final Capture size = EasyMock.newCapture();
-        EasyMock.expect(transportLayer.read(EasyMock.capture(size))).andAnswer(new IAnswer() {
-            @Override
-            public Integer answer() throws Throwable {
-                size.getValue().putInt(headerStruct.sizeOf());
-                return 4;
-            }
+        SaslServerAuthenticator authenticator = setupAuthenticator(configs, transportLayer,
+            SCRAM_SHA_256.mechanismName(), new DefaultChannelMetadataRegistry());
+
+        RequestHeader header = new RequestHeader(ApiKeys.METADATA, (short) 0, "clientId", 13243);
+        ByteBuffer headerBuffer = RequestTestUtils.serializeRequestHeader(header);
+
+        when(transportLayer.read(any(ByteBuffer.class))).then(invocation -> {
+            invocation.getArgument(0).putInt(headerBuffer.remaining());
+            return 4;
+        }).then(invocation -> {
+            // serialize only the request header. the authenticator should not parse beyond this
+            invocation.getArgument(0).put(headerBuffer.duplicate());
+            return headerBuffer.remaining();
         });
 
-        final Capture payload = EasyMock.newCapture();
-        EasyMock.expect(transportLayer.read(EasyMock.capture(payload))).andAnswer(new IAnswer() {
-            @Override
-            public Integer answer() throws Throwable {
-                // serialize only the request header. the authenticator should not parse beyond this
-                headerStruct.writeTo(payload.getValue());
-                return headerStruct.sizeOf();
-            }
-        });
-
-        EasyMock.replay(transportLayer);
-
         try {
             authenticator.authenticate();
             fail("Expected authenticate() to raise an exception");
         } catch (IllegalSaslStateException e) {
             // expected exception
         }
+
+        verify(transportLayer, times(2)).read(any(ByteBuffer.class));
+    }
+
+    @Test
+    public void testOldestApiVersionsRequest() throws IOException {
+        testApiVersionsRequest(ApiKeys.API_VERSIONS.oldestVersion(),
+            ClientInformation.UNKNOWN_NAME_OR_VERSION, ClientInformation.UNKNOWN_NAME_OR_VERSION);
+    }
+
+    @Test
+    public void testLatestApiVersionsRequest() throws IOException {
+        testApiVersionsRequest(ApiKeys.API_VERSIONS.latestVersion(),
+            "apache-kafka-java", AppInfoParser.getVersion());
+    }
+
+    private void testApiVersionsRequest(short version, String expectedSoftwareName,
+                                       String expectedSoftwareVersion) throws IOException {
+        TransportLayer transportLayer = mock(TransportLayer.class, Answers.RETURNS_DEEP_STUBS);
+        Map configs = Collections.singletonMap(BrokerSecurityConfigs.SASL_ENABLED_MECHANISMS_CONFIG,
+            Collections.singletonList(SCRAM_SHA_256.mechanismName()));
+        ChannelMetadataRegistry metadataRegistry = new DefaultChannelMetadataRegistry();
+        SaslServerAuthenticator authenticator = setupAuthenticator(configs, transportLayer,
+            SCRAM_SHA_256.mechanismName(), metadataRegistry);
+
+        RequestHeader header = new RequestHeader(ApiKeys.API_VERSIONS, version, "clientId", 0);
+        ByteBuffer headerBuffer = RequestTestUtils.serializeRequestHeader(header);
+
+        ApiVersionsRequest request = new ApiVersionsRequest.Builder().build(version);
+        ByteBuffer requestBuffer = request.serialize();
+        requestBuffer.rewind();
+
+        when(transportLayer.socketChannel().socket().getInetAddress()).thenReturn(InetAddress.getLoopbackAddress());
+
+        when(transportLayer.read(any(ByteBuffer.class))).then(invocation -> {
+            invocation.getArgument(0).putInt(headerBuffer.remaining() + requestBuffer.remaining());
+            return 4;
+        }).then(invocation -> {
+            invocation.getArgument(0)
+                .put(headerBuffer.duplicate())
+                .put(requestBuffer.duplicate());
+            return headerBuffer.remaining() + requestBuffer.remaining();
+        });
+
+        authenticator.authenticate();
+
+        assertEquals(expectedSoftwareName, metadataRegistry.clientInformation().softwareName());
+        assertEquals(expectedSoftwareVersion, metadataRegistry.clientInformation().softwareVersion());
+
+        verify(transportLayer, times(2)).read(any(ByteBuffer.class));
     }
 
-    private SaslServerAuthenticator setupAuthenticator(Map configs, TransportLayer transportLayer) throws IOException {
+    private SaslServerAuthenticator setupAuthenticator(Map configs, TransportLayer transportLayer,
+                                                       String mechanism, ChannelMetadataRegistry metadataRegistry) throws IOException {
         TestJaasConfig jaasConfig = new TestJaasConfig();
         jaasConfig.addEntry("jaasContext", PlainLoginModule.class.getName(), new HashMap());
-        JaasContext jaasContext = new JaasContext("jaasContext", JaasContext.Type.SERVER, jaasConfig);
-        Subject subject = new Subject();
-        return new SaslServerAuthenticator(configs, "node", jaasContext, subject, null, new CredentialCache(),
-                new ListenerName("ssl"), SecurityProtocol.SASL_SSL, transportLayer);
+        Map subjects = Collections.singletonMap(mechanism, new Subject());
+        Map callbackHandlers = Collections.singletonMap(
+                mechanism, new SaslServerCallbackHandler());
+        return new SaslServerAuthenticator(configs, callbackHandlers, "node", subjects, null,
+                new ListenerName("ssl"), SecurityProtocol.SASL_SSL, transportLayer, Collections.emptyMap(),
+                metadataRegistry, Time.SYSTEM);
     }
 
 }
diff --git a/clients/src/test/java/org/apache/kafka/common/security/authenticator/TestDigestLoginModule.java b/clients/src/test/java/org/apache/kafka/common/security/authenticator/TestDigestLoginModule.java
index f1ef740e5dbe3..c27e853d3915c 100644
--- a/clients/src/test/java/org/apache/kafka/common/security/authenticator/TestDigestLoginModule.java
+++ b/clients/src/test/java/org/apache/kafka/common/security/authenticator/TestDigestLoginModule.java
@@ -16,63 +16,46 @@
  */
 package org.apache.kafka.common.security.authenticator;
 
-import java.io.IOException;
-import java.security.Provider;
-import java.security.Security;
-import java.util.Arrays;
-import java.util.Enumeration;
-import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 
 import javax.security.auth.callback.Callback;
-import javax.security.auth.callback.CallbackHandler;
 import javax.security.auth.callback.NameCallback;
 import javax.security.auth.callback.PasswordCallback;
-import javax.security.auth.callback.UnsupportedCallbackException;
+import javax.security.auth.login.AppConfigurationEntry;
 import javax.security.sasl.AuthorizeCallback;
 import javax.security.sasl.RealmCallback;
-import javax.security.sasl.Sasl;
-import javax.security.sasl.SaslException;
-import javax.security.sasl.SaslServer;
-import javax.security.sasl.SaslServerFactory;
 
+import org.apache.kafka.common.security.auth.AuthenticateCallbackHandler;
 import org.apache.kafka.common.security.plain.PlainLoginModule;
 
 /**
- * Digest-MD5 login module for multi-mechanism tests. Since callback handlers are not configurable in Kafka
- * yet, this replaces the standard Digest-MD5 SASL server provider with one that invokes the test callback handler.
+ * Digest-MD5 login module for multi-mechanism tests.
  * This login module uses the same format as PlainLoginModule and hence simply reuses the same methods.
  *
  */
 public class TestDigestLoginModule extends PlainLoginModule {
 
-    private static final SaslServerFactory STANDARD_DIGEST_SASL_SERVER_FACTORY;
-    static {
-        SaslServerFactory digestSaslServerFactory = null;
-        Enumeration factories = Sasl.getSaslServerFactories();
-        Map emptyProps = new HashMap<>();
-        while (factories.hasMoreElements()) {
-            SaslServerFactory factory = factories.nextElement();
-            if (Arrays.asList(factory.getMechanismNames(emptyProps)).contains("DIGEST-MD5")) {
-                digestSaslServerFactory = factory;
-                break;
-            }
-        }
-        STANDARD_DIGEST_SASL_SERVER_FACTORY = digestSaslServerFactory;
-        Security.insertProviderAt(new DigestSaslServerProvider(), 1);
-    }
+    public static class DigestServerCallbackHandler implements AuthenticateCallbackHandler {
 
-    public static class DigestServerCallbackHandler implements CallbackHandler {
+        @Override
+        public void configure(Map configs, String saslMechanism, List jaasConfigEntries) {
+        }
 
         @Override
-        public void handle(Callback[] callbacks) throws IOException, UnsupportedCallbackException {
+        public void handle(Callback[] callbacks) {
+            String username = null;
             for (Callback callback : callbacks) {
                 if (callback instanceof NameCallback) {
                     NameCallback nameCallback = (NameCallback) callback;
-                    nameCallback.setName(nameCallback.getDefaultName());
+                    if (TestJaasConfig.USERNAME.equals(nameCallback.getDefaultName())) {
+                        nameCallback.setName(nameCallback.getDefaultName());
+                        username = TestJaasConfig.USERNAME;
+                    }
                 } else if (callback instanceof PasswordCallback) {
                     PasswordCallback passwordCallback = (PasswordCallback) callback;
-                    passwordCallback.setPassword(TestJaasConfig.PASSWORD.toCharArray());
+                    if (TestJaasConfig.USERNAME.equals(username))
+                        passwordCallback.setPassword(TestJaasConfig.PASSWORD.toCharArray());
                 } else if (callback instanceof RealmCallback) {
                     RealmCallback realmCallback = (RealmCallback) callback;
                     realmCallback.setText(realmCallback.getDefaultText());
@@ -85,30 +68,9 @@ public void handle(Callback[] callbacks) throws IOException, UnsupportedCallback
                 }
             }
         }
-    }
-
-    public static class DigestSaslServerFactory implements SaslServerFactory {
-
-        @Override
-        public SaslServer createSaslServer(String mechanism, String protocol, String serverName, Map props, CallbackHandler cbh)
-                throws SaslException {
-            return STANDARD_DIGEST_SASL_SERVER_FACTORY.createSaslServer(mechanism, protocol, serverName, props, new DigestServerCallbackHandler());
-        }
 
         @Override
-        public String[] getMechanismNames(Map props) {
-            return new String[] {"DIGEST-MD5"};
-        }
-    }
-
-    public static class DigestSaslServerProvider extends Provider {
-
-        private static final long serialVersionUID = 1L;
-
-        @SuppressWarnings("deprecation")
-        protected DigestSaslServerProvider() {
-            super("Test SASL/Digest-MD5 Server Provider", 1.0, "Test SASL/Digest-MD5 Server Provider for Kafka");
-            put("SaslServerFactory.DIGEST-MD5", TestDigestLoginModule.DigestSaslServerFactory.class.getName());
+        public void close() {
         }
     }
 }
diff --git a/clients/src/test/java/org/apache/kafka/common/security/authenticator/TestJaasConfig.java b/clients/src/test/java/org/apache/kafka/common/security/authenticator/TestJaasConfig.java
index 5336fd79dbbfd..f7ad140cb075e 100644
--- a/clients/src/test/java/org/apache/kafka/common/security/authenticator/TestJaasConfig.java
+++ b/clients/src/test/java/org/apache/kafka/common/security/authenticator/TestJaasConfig.java
@@ -26,9 +26,10 @@
 import javax.security.auth.login.AppConfigurationEntry.LoginModuleControlFlag;
 
 import org.apache.kafka.common.config.types.Password;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerLoginModule;
 import org.apache.kafka.common.security.plain.PlainLoginModule;
 import org.apache.kafka.common.security.scram.ScramLoginModule;
-import org.apache.kafka.common.security.scram.ScramMechanism;
+import org.apache.kafka.common.security.scram.internals.ScramMechanism;
 
 public class TestJaasConfig extends Configuration {
 
@@ -42,7 +43,7 @@ public class TestJaasConfig extends Configuration {
 
     public static TestJaasConfig createConfiguration(String clientMechanism, List serverMechanisms) {
         TestJaasConfig config = new TestJaasConfig();
-        config.createOrUpdateEntry(LOGIN_CONTEXT_CLIENT, loginModule(clientMechanism), defaultClientOptions());
+        config.createOrUpdateEntry(LOGIN_CONTEXT_CLIENT, loginModule(clientMechanism), defaultClientOptions(clientMechanism));
         for (String mechanism : serverMechanisms) {
             config.addEntry(LOGIN_CONTEXT_SERVER, loginModule(mechanism), defaultServerOptions(mechanism));
         }
@@ -54,6 +55,20 @@ public static Password jaasConfigProperty(String mechanism, String username, Str
         return new Password(loginModule(mechanism) + " required username=" + username + " password=" + password + ";");
     }
 
+    public static Password jaasConfigProperty(String mechanism, Map options) {
+        StringBuilder builder = new StringBuilder();
+        builder.append(loginModule(mechanism));
+        builder.append(" required");
+        for (Map.Entry option : options.entrySet()) {
+            builder.append(' ');
+            builder.append(option.getKey());
+            builder.append('=');
+            builder.append(option.getValue());
+        }
+        builder.append(';');
+        return new Password(builder.toString());
+    }
+
     public void setClientOptions(String saslMechanism, String clientUsername, String clientPassword) {
         Map options = new HashMap<>();
         if (clientUsername != null)
@@ -91,6 +106,9 @@ private static String loginModule(String mechanism) {
             case "DIGEST-MD5":
                 loginModule = TestDigestLoginModule.class.getName();
                 break;
+            case "OAUTHBEARER":
+                loginModule = OAuthBearerLoginModule.class.getName();
+                break;
             default:
                 if (ScramMechanism.isScram(mechanism))
                     loginModule = ScramLoginModule.class.getName();
@@ -100,6 +118,17 @@ private static String loginModule(String mechanism) {
         return loginModule;
     }
 
+    public static Map defaultClientOptions(String mechanism) {
+        switch (mechanism) {
+            case "OAUTHBEARER":
+                Map options = new HashMap<>();
+                options.put("unsecuredLoginStringClaim_sub", USERNAME);
+                return options;
+            default:
+                return defaultClientOptions();
+        }
+    }
+    
     public static Map defaultClientOptions() {
         Map options = new HashMap<>();
         options.put("username", USERNAME);
@@ -114,6 +143,9 @@ public static Map defaultServerOptions(String mechanism) {
             case "DIGEST-MD5":
                 options.put("user_" + USERNAME, PASSWORD);
                 break;
+            case "OAUTHBEARER":
+                options.put("unsecuredLoginStringClaim_sub", USERNAME);
+                break;
             default:
                 if (!ScramMechanism.isScram(mechanism))
                     throw new IllegalArgumentException("Unsupported mechanism " + mechanism);
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 3fbd310351e87..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
@@ -19,22 +19,23 @@
 import org.junit.Test;
 
 import java.io.IOException;
-import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.fail;
 
 public class KerberosNameTest {
 
     @Test
     public void testParse() throws IOException {
-        List rules = new ArrayList<>(Arrays.asList(
+        List rules = Arrays.asList(
             "RULE:[1:$1](App\\..*)s/App\\.(.*)/$1/g",
             "RULE:[2:$1](App\\..*)s/App\\.(.*)/$1/g",
             "DEFAULT"
-        ));
+        );
+
         KerberosShortNamer shortNamer = KerberosShortNamer.fromUnparsedRules("REALM.COM", rules);
 
         KerberosName name = KerberosName.parse("App.service-name/example.com@REALM.COM");
@@ -55,4 +56,86 @@ public void testParse() throws IOException {
         assertEquals("REALM.COM", name.realm());
         assertEquals("user", shortNamer.shortName(name));
     }
+
+    @Test
+    public void testToLowerCase() throws Exception {
+        List rules = Arrays.asList(
+            "RULE:[1:$1]/L",
+            "RULE:[2:$1](Test.*)s/ABC///L",
+            "RULE:[2:$1](ABC.*)s/ABC/XYZ/g/L",
+            "RULE:[2:$1](App\\..*)s/App\\.(.*)/$1/g/L",
+            "RULE:[2:$1]/L",
+            "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 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"));
+        testInvalidRule(Arrays.asList("DEFAUL"));
+        testInvalidRule(Arrays.asList("DEFAULT/L"));
+        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"));
+    }
+
+    private void testInvalidRule(List rules) {
+        try {
+            KerberosShortNamer.fromUnparsedRules("REALM.COM", rules);
+            fail("should have thrown IllegalArgumentException");
+        } catch (IllegalArgumentException e) {
+        }
+    }
 }
diff --git a/clients/src/test/java/org/apache/kafka/common/security/kerberos/KerberosRuleTest.java b/clients/src/test/java/org/apache/kafka/common/security/kerberos/KerberosRuleTest.java
new file mode 100644
index 0000000000000..f79c47af05fdb
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/security/kerberos/KerberosRuleTest.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.security.kerberos;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.fail;
+
+import org.junit.Test;
+
+public class KerberosRuleTest {
+
+    @Test
+    public void testReplaceParameters() throws BadFormatString {
+        // positive test cases
+        assertEquals(KerberosRule.replaceParameters("", new String[0]), "");
+        assertEquals(KerberosRule.replaceParameters("hello", new String[0]), "hello");
+        assertEquals(KerberosRule.replaceParameters("", new String[]{"too", "many", "parameters", "are", "ok"}), "");
+        assertEquals(KerberosRule.replaceParameters("hello", new String[]{"too", "many", "parameters", "are", "ok"}), "hello");
+        assertEquals(KerberosRule.replaceParameters("hello $0", new String[]{"too", "many", "parameters", "are", "ok"}), "hello too");
+        assertEquals(KerberosRule.replaceParameters("hello $0", new String[]{"no recursion $1"}), "hello no recursion $1");
+
+        // negative test cases
+        try {
+            KerberosRule.replaceParameters("$0", new String[]{});
+            fail("An out-of-bounds parameter number should trigger an exception!");
+        } catch (BadFormatString bfs) {
+        }
+        try {
+            KerberosRule.replaceParameters("hello $a", new String[]{"does not matter"});
+            fail("A malformed parameter name should trigger an exception!");
+        } catch (BadFormatString bfs) {
+        }
+    }
+
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerExtensionsValidatorCallbackTest.java b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerExtensionsValidatorCallbackTest.java
new file mode 100644
index 0000000000000..f65031ff38ae6
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerExtensionsValidatorCallbackTest.java
@@ -0,0 +1,90 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.oauthbearer;
+
+import org.apache.kafka.common.security.auth.SaslExtensions;
+import org.junit.Test;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+public class OAuthBearerExtensionsValidatorCallbackTest {
+    private static final OAuthBearerToken TOKEN = new OAuthBearerTokenMock();
+
+    @Test
+    public void testValidatedExtensionsAreReturned() {
+        Map extensions = new HashMap<>();
+        extensions.put("hello", "bye");
+
+        OAuthBearerExtensionsValidatorCallback callback = new OAuthBearerExtensionsValidatorCallback(TOKEN, new SaslExtensions(extensions));
+
+        assertTrue(callback.validatedExtensions().isEmpty());
+        assertTrue(callback.invalidExtensions().isEmpty());
+        callback.valid("hello");
+        assertFalse(callback.validatedExtensions().isEmpty());
+        assertEquals("bye", callback.validatedExtensions().get("hello"));
+        assertTrue(callback.invalidExtensions().isEmpty());
+    }
+
+    @Test
+    public void testInvalidExtensionsAndErrorMessagesAreReturned() {
+        Map extensions = new HashMap<>();
+        extensions.put("hello", "bye");
+
+        OAuthBearerExtensionsValidatorCallback callback = new OAuthBearerExtensionsValidatorCallback(TOKEN, new SaslExtensions(extensions));
+
+        assertTrue(callback.validatedExtensions().isEmpty());
+        assertTrue(callback.invalidExtensions().isEmpty());
+        callback.error("hello", "error");
+        assertFalse(callback.invalidExtensions().isEmpty());
+        assertEquals("error", callback.invalidExtensions().get("hello"));
+        assertTrue(callback.validatedExtensions().isEmpty());
+    }
+
+    /**
+     * Extensions that are neither validated or invalidated must not be present in either maps
+     */
+    @Test
+    public void testUnvalidatedExtensionsAreIgnored() {
+        Map extensions = new HashMap<>();
+        extensions.put("valid", "valid");
+        extensions.put("error", "error");
+        extensions.put("nothing", "nothing");
+
+        OAuthBearerExtensionsValidatorCallback callback = new OAuthBearerExtensionsValidatorCallback(TOKEN, new SaslExtensions(extensions));
+        callback.error("error", "error");
+        callback.valid("valid");
+
+        assertFalse(callback.validatedExtensions().containsKey("nothing"));
+        assertFalse(callback.invalidExtensions().containsKey("nothing"));
+        assertEquals("nothing", callback.ignoredExtensions().get("nothing"));
+    }
+
+    @Test(expected = IllegalArgumentException.class)
+    public void testCannotValidateExtensionWhichWasNotGiven() {
+        Map extensions = new HashMap<>();
+        extensions.put("hello", "bye");
+
+        OAuthBearerExtensionsValidatorCallback callback = new OAuthBearerExtensionsValidatorCallback(TOKEN, new SaslExtensions(extensions));
+
+        callback.valid("???");
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerLoginModuleTest.java b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerLoginModuleTest.java
new file mode 100644
index 0000000000000..7c6d998bd3e41
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerLoginModuleTest.java
@@ -0,0 +1,439 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.oauthbearer;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotSame;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.assertNotNull;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verifyNoInteractions;
+
+import java.io.IOException;
+import java.util.Collections;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+
+import javax.security.auth.Subject;
+import javax.security.auth.callback.Callback;
+import javax.security.auth.callback.UnsupportedCallbackException;
+import javax.security.auth.login.AppConfigurationEntry;
+import javax.security.auth.login.LoginException;
+
+import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.security.auth.AuthenticateCallbackHandler;
+import org.apache.kafka.common.security.auth.SaslExtensionsCallback;
+import org.apache.kafka.common.security.auth.SaslExtensions;
+import org.junit.Test;
+
+public class OAuthBearerLoginModuleTest {
+
+    public static final SaslExtensions RAISE_UNSUPPORTED_CB_EXCEPTION_FLAG = null;
+
+    private static class TestCallbackHandler implements AuthenticateCallbackHandler {
+        private final OAuthBearerToken[] tokens;
+        private int index = 0;
+        private int extensionsIndex = 0;
+        private final SaslExtensions[] extensions;
+
+        public TestCallbackHandler(OAuthBearerToken[] tokens, SaslExtensions[] extensions) {
+            this.tokens = Objects.requireNonNull(tokens);
+            this.extensions = extensions;
+        }
+
+        @Override
+        public void handle(Callback[] callbacks) throws IOException, UnsupportedCallbackException {
+            for (Callback callback : callbacks) {
+                if (callback instanceof OAuthBearerTokenCallback)
+                    try {
+                        handleCallback((OAuthBearerTokenCallback) callback);
+                    } catch (KafkaException e) {
+                        throw new IOException(e.getMessage(), e);
+                    }
+                else if (callback instanceof SaslExtensionsCallback) {
+                    try {
+                        handleExtensionsCallback((SaslExtensionsCallback) callback);
+                    } catch (KafkaException e) {
+                        throw new IOException(e.getMessage(), e);
+                    }
+                } else
+                    throw new UnsupportedCallbackException(callback);
+            }
+        }
+
+        @Override
+        public void configure(Map configs, String saslMechanism,
+                List jaasConfigEntries) {
+            // empty
+        }
+
+        @Override
+        public void close() {
+            // empty
+        }
+
+        private void handleCallback(OAuthBearerTokenCallback callback) throws IOException {
+            if (callback.token() != null)
+                throw new IllegalArgumentException("Callback had a token already");
+            if (tokens.length > index)
+                callback.token(tokens[index++]);
+            else
+                throw new IOException("no more tokens");
+        }
+
+        private void handleExtensionsCallback(SaslExtensionsCallback callback) throws IOException, UnsupportedCallbackException {
+            if (extensions.length > extensionsIndex) {
+                SaslExtensions extension = extensions[extensionsIndex++];
+
+                if (extension == RAISE_UNSUPPORTED_CB_EXCEPTION_FLAG) {
+                    throw new UnsupportedCallbackException(callback);
+                }
+
+                callback.extensions(extension);
+            } else
+                throw new IOException("no more extensions");
+        }
+    }
+
+    @Test
+    public void login1Commit1Login2Commit2Logout1Login3Commit3Logout2() throws LoginException {
+        /*
+         * Invoke login()/commit() on loginModule1; invoke login/commit() on
+         * loginModule2; invoke logout() on loginModule1; invoke login()/commit() on
+         * loginModule3; invoke logout() on loginModule2
+         */
+        Subject subject = new Subject();
+        Set privateCredentials = subject.getPrivateCredentials();
+        Set publicCredentials = subject.getPublicCredentials();
+
+        // Create callback handler
+        OAuthBearerToken[] tokens = new OAuthBearerToken[] {mock(OAuthBearerToken.class),
+            mock(OAuthBearerToken.class), mock(OAuthBearerToken.class)};
+        SaslExtensions[] extensions = new SaslExtensions[] {mock(SaslExtensions.class),
+            mock(SaslExtensions.class), mock(SaslExtensions.class)};
+        TestCallbackHandler testTokenCallbackHandler = new TestCallbackHandler(tokens, extensions);
+
+        // Create login modules
+        OAuthBearerLoginModule loginModule1 = new OAuthBearerLoginModule();
+        loginModule1.initialize(subject, testTokenCallbackHandler, Collections.emptyMap(),
+                Collections.emptyMap());
+        OAuthBearerLoginModule loginModule2 = new OAuthBearerLoginModule();
+        loginModule2.initialize(subject, testTokenCallbackHandler, Collections.emptyMap(),
+                Collections.emptyMap());
+        OAuthBearerLoginModule loginModule3 = new OAuthBearerLoginModule();
+        loginModule3.initialize(subject, testTokenCallbackHandler, Collections.emptyMap(),
+                Collections.emptyMap());
+
+        // Should start with nothing
+        assertEquals(0, privateCredentials.size());
+        assertEquals(0, publicCredentials.size());
+        loginModule1.login();
+        // Should still have nothing until commit() is called
+        assertEquals(0, privateCredentials.size());
+        assertEquals(0, publicCredentials.size());
+        loginModule1.commit();
+        // Now we should have the first token and extensions
+        assertEquals(1, privateCredentials.size());
+        assertEquals(1, publicCredentials.size());
+        assertSame(tokens[0], privateCredentials.iterator().next());
+        assertSame(extensions[0], publicCredentials.iterator().next());
+
+        // Now login on loginModule2 to get the second token
+        // loginModule2 does not support the extensions callback and will raise UnsupportedCallbackException
+        loginModule2.login();
+        // Should still have just the first token and extensions
+        assertEquals(1, privateCredentials.size());
+        assertEquals(1, publicCredentials.size());
+        assertSame(tokens[0], privateCredentials.iterator().next());
+        assertSame(extensions[0], publicCredentials.iterator().next());
+        loginModule2.commit();
+        // Should have the first and second tokens at this point
+        assertEquals(2, privateCredentials.size());
+        assertEquals(2, publicCredentials.size());
+        Iterator iterator = privateCredentials.iterator();
+        Iterator publicIterator = publicCredentials.iterator();
+        assertNotSame(tokens[2], iterator.next());
+        assertNotSame(tokens[2], iterator.next());
+        assertNotSame(extensions[2], publicIterator.next());
+        assertNotSame(extensions[2], publicIterator.next());
+        // finally logout() on loginModule1
+        loginModule1.logout();
+        // Now we should have just the second token and extension
+        assertEquals(1, privateCredentials.size());
+        assertEquals(1, publicCredentials.size());
+        assertSame(tokens[1], privateCredentials.iterator().next());
+        assertSame(extensions[1], publicCredentials.iterator().next());
+
+        // Now login on loginModule3 to get the third token
+        loginModule3.login();
+        // Should still have just the second token and extensions
+        assertEquals(1, privateCredentials.size());
+        assertEquals(1, publicCredentials.size());
+        assertSame(tokens[1], privateCredentials.iterator().next());
+        assertSame(extensions[1], publicCredentials.iterator().next());
+        loginModule3.commit();
+        // Should have the second and third tokens at this point
+        assertEquals(2, privateCredentials.size());
+        assertEquals(2, publicCredentials.size());
+        iterator = privateCredentials.iterator();
+        publicIterator = publicCredentials.iterator();
+        assertNotSame(tokens[0], iterator.next());
+        assertNotSame(tokens[0], iterator.next());
+        assertNotSame(extensions[0], publicIterator.next());
+        assertNotSame(extensions[0], publicIterator.next());
+        // finally logout() on loginModule2
+        loginModule2.logout();
+        // Now we should have just the third token
+        assertEquals(1, privateCredentials.size());
+        assertEquals(1, publicCredentials.size());
+        assertSame(tokens[2], privateCredentials.iterator().next());
+        assertSame(extensions[2], publicCredentials.iterator().next());
+
+        verifyNoInteractions((Object[]) tokens);
+        verifyNoInteractions((Object[]) extensions);
+    }
+
+    @Test
+    public void login1Commit1Logout1Login2Commit2Logout2() throws LoginException {
+        /*
+         * Invoke login()/commit() on loginModule1; invoke logout() on loginModule1;
+         * invoke login()/commit() on loginModule2; invoke logout() on loginModule2
+         */
+        Subject subject = new Subject();
+        Set privateCredentials = subject.getPrivateCredentials();
+        Set publicCredentials = subject.getPublicCredentials();
+
+        // Create callback handler
+        OAuthBearerToken[] tokens = new OAuthBearerToken[] {mock(OAuthBearerToken.class),
+            mock(OAuthBearerToken.class)};
+        SaslExtensions[] extensions = new SaslExtensions[] {mock(SaslExtensions.class),
+            mock(SaslExtensions.class)};
+        TestCallbackHandler testTokenCallbackHandler = new TestCallbackHandler(tokens, extensions);
+
+        // Create login modules
+        OAuthBearerLoginModule loginModule1 = new OAuthBearerLoginModule();
+        loginModule1.initialize(subject, testTokenCallbackHandler, Collections.emptyMap(),
+                Collections.emptyMap());
+        OAuthBearerLoginModule loginModule2 = new OAuthBearerLoginModule();
+        loginModule2.initialize(subject, testTokenCallbackHandler, Collections.emptyMap(),
+                Collections.emptyMap());
+
+        // Should start with nothing
+        assertEquals(0, privateCredentials.size());
+        assertEquals(0, publicCredentials.size());
+        loginModule1.login();
+        // Should still have nothing until commit() is called
+        assertEquals(0, privateCredentials.size());
+        assertEquals(0, publicCredentials.size());
+        loginModule1.commit();
+        // Now we should have the first token
+        assertEquals(1, privateCredentials.size());
+        assertEquals(1, publicCredentials.size());
+        assertSame(tokens[0], privateCredentials.iterator().next());
+        assertSame(extensions[0], publicCredentials.iterator().next());
+        loginModule1.logout();
+        // Should have nothing again
+        assertEquals(0, privateCredentials.size());
+        assertEquals(0, publicCredentials.size());
+
+        loginModule2.login();
+        // Should still have nothing until commit() is called
+        assertEquals(0, privateCredentials.size());
+        assertEquals(0, publicCredentials.size());
+        loginModule2.commit();
+        // Now we should have the second token
+        assertEquals(1, privateCredentials.size());
+        assertEquals(1, publicCredentials.size());
+        assertSame(tokens[1], privateCredentials.iterator().next());
+        assertSame(extensions[1], publicCredentials.iterator().next());
+        loginModule2.logout();
+        // Should have nothing again
+        assertEquals(0, privateCredentials.size());
+        assertEquals(0, publicCredentials.size());
+
+        verifyNoInteractions((Object[]) tokens);
+        verifyNoInteractions((Object[]) extensions);
+    }
+
+    @Test
+    public void loginAbortLoginCommitLogout() throws LoginException {
+        /*
+         * Invoke login(); invoke abort(); invoke login(); logout()
+         */
+        Subject subject = new Subject();
+        Set privateCredentials = subject.getPrivateCredentials();
+        Set publicCredentials = subject.getPublicCredentials();
+
+        // Create callback handler
+        OAuthBearerToken[] tokens = new OAuthBearerToken[] {mock(OAuthBearerToken.class),
+            mock(OAuthBearerToken.class)};
+        SaslExtensions[] extensions = new SaslExtensions[] {mock(SaslExtensions.class),
+            mock(SaslExtensions.class)};
+        TestCallbackHandler testTokenCallbackHandler = new TestCallbackHandler(tokens, extensions);
+
+        // Create login module
+        OAuthBearerLoginModule loginModule = new OAuthBearerLoginModule();
+        loginModule.initialize(subject, testTokenCallbackHandler, Collections.emptyMap(),
+                Collections.emptyMap());
+
+        // Should start with nothing
+        assertEquals(0, privateCredentials.size());
+        assertEquals(0, publicCredentials.size());
+        loginModule.login();
+        // Should still have nothing until commit() is called
+        assertEquals(0, privateCredentials.size());
+        assertEquals(0, publicCredentials.size());
+        loginModule.abort();
+        // Should still have nothing since we aborted
+        assertEquals(0, privateCredentials.size());
+        assertEquals(0, publicCredentials.size());
+
+        loginModule.login();
+        // Should still have nothing until commit() is called
+        assertEquals(0, privateCredentials.size());
+        assertEquals(0, publicCredentials.size());
+        loginModule.commit();
+        // Now we should have the second token
+        assertEquals(1, privateCredentials.size());
+        assertEquals(1, publicCredentials.size());
+        assertSame(tokens[1], privateCredentials.iterator().next());
+        assertSame(extensions[1], publicCredentials.iterator().next());
+        loginModule.logout();
+        // Should have nothing again
+        assertEquals(0, privateCredentials.size());
+        assertEquals(0, publicCredentials.size());
+
+        verifyNoInteractions((Object[]) tokens);
+        verifyNoInteractions((Object[]) extensions);
+    }
+
+    @Test
+    public void login1Commit1Login2Abort2Login3Commit3Logout3() throws LoginException {
+        /*
+         * Invoke login()/commit() on loginModule1; invoke login()/abort() on
+         * loginModule2; invoke login()/commit()/logout() on loginModule3
+         */
+        Subject subject = new Subject();
+        Set privateCredentials = subject.getPrivateCredentials();
+        Set publicCredentials = subject.getPublicCredentials();
+
+        // Create callback handler
+        OAuthBearerToken[] tokens = new OAuthBearerToken[] {mock(OAuthBearerToken.class),
+            mock(OAuthBearerToken.class), mock(OAuthBearerToken.class)};
+        SaslExtensions[] extensions = new SaslExtensions[] {mock(SaslExtensions.class),
+            mock(SaslExtensions.class), mock(SaslExtensions.class)};
+        TestCallbackHandler testTokenCallbackHandler = new TestCallbackHandler(tokens, extensions);
+
+        // Create login modules
+        OAuthBearerLoginModule loginModule1 = new OAuthBearerLoginModule();
+        loginModule1.initialize(subject, testTokenCallbackHandler, Collections.emptyMap(),
+                Collections.emptyMap());
+        OAuthBearerLoginModule loginModule2 = new OAuthBearerLoginModule();
+        loginModule2.initialize(subject, testTokenCallbackHandler, Collections.emptyMap(),
+                Collections.emptyMap());
+        OAuthBearerLoginModule loginModule3 = new OAuthBearerLoginModule();
+        loginModule3.initialize(subject, testTokenCallbackHandler, Collections.emptyMap(),
+                Collections.emptyMap());
+
+        // Should start with nothing
+        assertEquals(0, privateCredentials.size());
+        assertEquals(0, publicCredentials.size());
+        loginModule1.login();
+        // Should still have nothing until commit() is called
+        assertEquals(0, privateCredentials.size());
+        assertEquals(0, publicCredentials.size());
+        loginModule1.commit();
+        // Now we should have the first token
+        assertEquals(1, privateCredentials.size());
+        assertEquals(1, publicCredentials.size());
+        assertSame(tokens[0], privateCredentials.iterator().next());
+        assertSame(extensions[0], publicCredentials.iterator().next());
+
+        // Now go get the second token
+        loginModule2.login();
+        // Should still have first token
+        assertEquals(1, privateCredentials.size());
+        assertEquals(1, publicCredentials.size());
+        assertSame(tokens[0], privateCredentials.iterator().next());
+        assertSame(extensions[0], publicCredentials.iterator().next());
+        loginModule2.abort();
+        // Should still have just the first token because we aborted
+        assertEquals(1, privateCredentials.size());
+        assertSame(tokens[0], privateCredentials.iterator().next());
+        assertEquals(1, publicCredentials.size());
+        assertSame(extensions[0], publicCredentials.iterator().next());
+
+        // Now go get the third token
+        loginModule2.login();
+        // Should still have first token
+        assertEquals(1, privateCredentials.size());
+        assertSame(tokens[0], privateCredentials.iterator().next());
+        assertEquals(1, publicCredentials.size());
+        assertSame(extensions[0], publicCredentials.iterator().next());
+        loginModule2.commit();
+        // Should have first and third tokens at this point
+        assertEquals(2, privateCredentials.size());
+        Iterator iterator = privateCredentials.iterator();
+        assertNotSame(tokens[1], iterator.next());
+        assertNotSame(tokens[1], iterator.next());
+        assertEquals(2, publicCredentials.size());
+        Iterator publicIterator = publicCredentials.iterator();
+        assertNotSame(extensions[1], publicIterator.next());
+        assertNotSame(extensions[1], publicIterator.next());
+        loginModule1.logout();
+        // Now we should have just the third token
+        assertEquals(1, privateCredentials.size());
+        assertSame(tokens[2], privateCredentials.iterator().next());
+        assertEquals(1, publicCredentials.size());
+        assertSame(extensions[2], publicCredentials.iterator().next());
+
+        verifyNoInteractions((Object[]) tokens);
+        verifyNoInteractions((Object[]) extensions);
+    }
+
+    /**
+     * 2.1.0 added customizable SASL extensions and a new callback type.
+     * Ensure that old, custom-written callbackHandlers that do not handle the callback work
+     */
+    @Test
+    public void commitDoesNotThrowOnUnsupportedExtensionsCallback() throws LoginException {
+        Subject subject = new Subject();
+
+        // Create callback handler
+        OAuthBearerToken[] tokens = new OAuthBearerToken[] {mock(OAuthBearerToken.class),
+                mock(OAuthBearerToken.class), mock(OAuthBearerToken.class)};
+        TestCallbackHandler testTokenCallbackHandler = new TestCallbackHandler(tokens, new SaslExtensions[] {RAISE_UNSUPPORTED_CB_EXCEPTION_FLAG});
+
+        // Create login modules
+        OAuthBearerLoginModule loginModule1 = new OAuthBearerLoginModule();
+        loginModule1.initialize(subject, testTokenCallbackHandler, Collections.emptyMap(),
+                Collections.emptyMap());
+
+        loginModule1.login();
+        // Should populate public credentials with SaslExtensions and not throw an exception
+        loginModule1.commit();
+        SaslExtensions extensions = subject.getPublicCredentials(SaslExtensions.class).iterator().next();
+        assertNotNull(extensions);
+        assertTrue(extensions.map().isEmpty());
+
+        verifyNoInteractions((Object[]) tokens);
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerSaslClienCallbackHandlerTest.java b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerSaslClienCallbackHandlerTest.java
new file mode 100644
index 0000000000000..4115c5227c27b
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerSaslClienCallbackHandlerTest.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.security.oauthbearer;
+
+import static org.junit.Assert.assertEquals;
+
+import java.io.IOException;
+import java.security.AccessController;
+import java.security.PrivilegedActionException;
+import java.security.PrivilegedExceptionAction;
+import java.util.Collections;
+import java.util.Set;
+
+import javax.security.auth.Subject;
+import javax.security.auth.callback.Callback;
+
+import org.apache.kafka.common.security.oauthbearer.internals.OAuthBearerSaslClientCallbackHandler;
+import org.junit.Test;
+
+public class OAuthBearerSaslClienCallbackHandlerTest {
+    private static OAuthBearerToken createTokenWithLifetimeMillis(final long lifetimeMillis) {
+        return new OAuthBearerToken() {
+            @Override
+            public String value() {
+                return null;
+            }
+
+            @Override
+            public Long startTimeMs() {
+                return null;
+            }
+
+            @Override
+            public Set scope() {
+                return null;
+            }
+
+            @Override
+            public String principalName() {
+                return null;
+            }
+
+            @Override
+            public long lifetimeMs() {
+                return lifetimeMillis;
+            }
+        };
+    }
+
+    @Test(expected = IOException.class)
+    public void testWithZeroTokens() throws Throwable {
+        OAuthBearerSaslClientCallbackHandler handler = createCallbackHandler();
+        try {
+            Subject.doAs(new Subject(), (PrivilegedExceptionAction) () -> {
+                OAuthBearerTokenCallback callback = new OAuthBearerTokenCallback();
+                handler.handle(new Callback[] {callback});
+                return null;
+            });
+        } catch (PrivilegedActionException e) {
+            throw e.getCause();
+        }
+    }
+
+    @Test()
+    public void testWithPotentiallyMultipleTokens() throws Exception {
+        OAuthBearerSaslClientCallbackHandler handler = createCallbackHandler();
+        Subject.doAs(new Subject(), (PrivilegedExceptionAction) () -> {
+            final int maxTokens = 4;
+            final Set privateCredentials = Subject.getSubject(AccessController.getContext())
+                    .getPrivateCredentials();
+            privateCredentials.clear();
+            for (int num = 1; num <= maxTokens; ++num) {
+                privateCredentials.add(createTokenWithLifetimeMillis(num));
+                OAuthBearerTokenCallback callback = new OAuthBearerTokenCallback();
+                handler.handle(new Callback[] {callback});
+                assertEquals(num, callback.token().lifetimeMs());
+            }
+            return null;
+        });
+    }
+
+    private static OAuthBearerSaslClientCallbackHandler createCallbackHandler() {
+        OAuthBearerSaslClientCallbackHandler handler = new OAuthBearerSaslClientCallbackHandler();
+        handler.configure(Collections.emptyMap(), OAuthBearerLoginModule.OAUTHBEARER_MECHANISM,
+                Collections.emptyList());
+        return handler;
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerTokenCallbackTest.java b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerTokenCallbackTest.java
new file mode 100644
index 0000000000000..be97ea2af4f95
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerTokenCallbackTest.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.security.oauthbearer;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertSame;
+
+import java.util.Collections;
+import java.util.Set;
+
+import org.junit.Test;
+
+public class OAuthBearerTokenCallbackTest {
+    private static final OAuthBearerToken TOKEN = new OAuthBearerToken() {
+        @Override
+        public String value() {
+            return "value";
+        }
+
+        @Override
+        public Long startTimeMs() {
+            return null;
+        }
+
+        @Override
+        public Set scope() {
+            return Collections.emptySet();
+        }
+
+        @Override
+        public String principalName() {
+            return "principalName";
+        }
+
+        @Override
+        public long lifetimeMs() {
+            return 0;
+        }
+    };
+
+    @Test
+    public void testError() {
+        String errorCode = "errorCode";
+        String errorDescription = "errorDescription";
+        String errorUri = "errorUri";
+        OAuthBearerTokenCallback callback = new OAuthBearerTokenCallback();
+        callback.error(errorCode, errorDescription, errorUri);
+        assertEquals(errorCode, callback.errorCode());
+        assertEquals(errorDescription, callback.errorDescription());
+        assertEquals(errorUri, callback.errorUri());
+        assertNull(callback.token());
+    }
+
+    @Test
+    public void testToken() {
+        OAuthBearerTokenCallback callback = new OAuthBearerTokenCallback();
+        callback.token(TOKEN);
+        assertSame(TOKEN, callback.token());
+        assertNull(callback.errorCode());
+        assertNull(callback.errorDescription());
+        assertNull(callback.errorUri());
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerTokenMock.java b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerTokenMock.java
new file mode 100644
index 0000000000000..994c923a4c17e
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerTokenMock.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.common.security.oauthbearer;
+
+import java.util.Set;
+
+public class OAuthBearerTokenMock implements OAuthBearerToken {
+    @Override
+    public String value() {
+        return null;
+    }
+
+    @Override
+    public Set scope() {
+        return null;
+    }
+
+    @Override
+    public long lifetimeMs() {
+        return 0;
+    }
+
+    @Override
+    public String principalName() {
+        return null;
+    }
+
+    @Override
+    public Long startTimeMs() {
+        return null;
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerValidatorCallbackTest.java b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerValidatorCallbackTest.java
new file mode 100644
index 0000000000000..ed827661bcda9
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/OAuthBearerValidatorCallbackTest.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.security.oauthbearer;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertSame;
+
+import java.util.Collections;
+import java.util.Set;
+
+import org.junit.Test;
+
+public class OAuthBearerValidatorCallbackTest {
+    private static final OAuthBearerToken TOKEN = new OAuthBearerToken() {
+        @Override
+        public String value() {
+            return "value";
+        }
+
+        @Override
+        public Long startTimeMs() {
+            return null;
+        }
+
+        @Override
+        public Set scope() {
+            return Collections.emptySet();
+        }
+
+        @Override
+        public String principalName() {
+            return "principalName";
+        }
+
+        @Override
+        public long lifetimeMs() {
+            return 0;
+        }
+    };
+
+    @Test
+    public void testError() {
+        String errorStatus = "errorStatus";
+        String errorScope = "errorScope";
+        String errorOpenIDConfiguration = "errorOpenIDConfiguration";
+        OAuthBearerValidatorCallback callback = new OAuthBearerValidatorCallback(TOKEN.value());
+        callback.error(errorStatus, errorScope, errorOpenIDConfiguration);
+        assertEquals(errorStatus, callback.errorStatus());
+        assertEquals(errorScope, callback.errorScope());
+        assertEquals(errorOpenIDConfiguration, callback.errorOpenIDConfiguration());
+        assertNull(callback.token());
+    }
+
+    @Test
+    public void testToken() {
+        OAuthBearerValidatorCallback callback = new OAuthBearerValidatorCallback(TOKEN.value());
+        callback.token(TOKEN);
+        assertSame(TOKEN, callback.token());
+        assertNull(callback.errorStatus());
+        assertNull(callback.errorScope());
+        assertNull(callback.errorOpenIDConfiguration());
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/OAuthBearerClientInitialResponseTest.java b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/OAuthBearerClientInitialResponseTest.java
new file mode 100644
index 0000000000000..0ba956561dfd5
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/OAuthBearerClientInitialResponseTest.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.security.oauthbearer.internals;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+import org.apache.kafka.common.security.auth.SaslExtensions;
+import org.junit.Test;
+
+import javax.security.sasl.SaslException;
+import java.nio.charset.StandardCharsets;
+import java.util.HashMap;
+import java.util.Map;
+
+public class OAuthBearerClientInitialResponseTest {
+
+    /*
+        Test how a client would build a response
+     */
+    @Test
+    public void testBuildClientResponseToBytes() throws Exception {
+        String expectedMesssage = "n,,\u0001auth=Bearer 123.345.567\u0001nineteen=42\u0001\u0001";
+
+        Map extensions = new HashMap<>();
+        extensions.put("nineteen", "42");
+        OAuthBearerClientInitialResponse response = new OAuthBearerClientInitialResponse("123.345.567", new SaslExtensions(extensions));
+
+        String message = new String(response.toBytes(), StandardCharsets.UTF_8);
+
+        assertEquals(expectedMesssage, message);
+    }
+
+    @Test
+    public void testBuildServerResponseToBytes() throws Exception {
+        String serverMessage = "n,,\u0001auth=Bearer 123.345.567\u0001nineteen=42\u0001\u0001";
+        OAuthBearerClientInitialResponse response = new OAuthBearerClientInitialResponse(serverMessage.getBytes(StandardCharsets.UTF_8));
+
+        String message = new String(response.toBytes(), StandardCharsets.UTF_8);
+
+        assertEquals(serverMessage, message);
+    }
+
+    @Test(expected = SaslException.class)
+    public void testThrowsSaslExceptionOnInvalidExtensionKey() throws Exception {
+        Map extensions = new HashMap<>();
+        extensions.put("19", "42"); // keys can only be a-z
+        new OAuthBearerClientInitialResponse("123.345.567", new SaslExtensions(extensions));
+    }
+
+    @Test
+    public void testToken() throws Exception {
+        String message = "n,,\u0001auth=Bearer 123.345.567\u0001\u0001";
+        OAuthBearerClientInitialResponse response = new OAuthBearerClientInitialResponse(message.getBytes(StandardCharsets.UTF_8));
+        assertEquals("123.345.567", response.tokenValue());
+        assertEquals("", response.authorizationId());
+    }
+
+    @Test
+    public void testAuthorizationId() throws Exception {
+        String message = "n,a=myuser,\u0001auth=Bearer 345\u0001\u0001";
+        OAuthBearerClientInitialResponse response = new OAuthBearerClientInitialResponse(message.getBytes(StandardCharsets.UTF_8));
+        assertEquals("345", response.tokenValue());
+        assertEquals("myuser", response.authorizationId());
+    }
+
+    @Test
+    public void testExtensions() throws Exception {
+        String message = "n,,\u0001propA=valueA1, valueA2\u0001auth=Bearer 567\u0001propB=valueB\u0001\u0001";
+        OAuthBearerClientInitialResponse response = new OAuthBearerClientInitialResponse(message.getBytes(StandardCharsets.UTF_8));
+        assertEquals("567", response.tokenValue());
+        assertEquals("", response.authorizationId());
+        assertEquals("valueA1, valueA2", response.extensions().map().get("propA"));
+        assertEquals("valueB", response.extensions().map().get("propB"));
+    }
+
+    // The example in the RFC uses `vF9dft4qmTc2Nvb3RlckBhbHRhdmlzdGEuY29tCg==` as the token
+    // But since we use Base64Url encoding, padding is omitted. Hence this test verifies without '='.
+    @Test
+    public void testRfc7688Example() throws Exception {
+        String message = "n,a=user@example.com,\u0001host=server.example.com\u0001port=143\u0001" +
+                "auth=Bearer vF9dft4qmTc2Nvb3RlckBhbHRhdmlzdGEuY29tCg\u0001\u0001";
+        OAuthBearerClientInitialResponse response = new OAuthBearerClientInitialResponse(message.getBytes(StandardCharsets.UTF_8));
+        assertEquals("vF9dft4qmTc2Nvb3RlckBhbHRhdmlzdGEuY29tCg", response.tokenValue());
+        assertEquals("user@example.com", response.authorizationId());
+        assertEquals("server.example.com", response.extensions().map().get("host"));
+        assertEquals("143", response.extensions().map().get("port"));
+    }
+
+    @Test
+    public void testNoExtensionsFromByteArray() throws Exception {
+        String message = "n,a=user@example.com,\u0001" +
+                "auth=Bearer vF9dft4qmTc2Nvb3RlckBhbHRhdmlzdGEuY29tCg\u0001\u0001";
+        OAuthBearerClientInitialResponse response = new OAuthBearerClientInitialResponse(message.getBytes(StandardCharsets.UTF_8));
+        assertEquals("vF9dft4qmTc2Nvb3RlckBhbHRhdmlzdGEuY29tCg", response.tokenValue());
+        assertEquals("user@example.com", response.authorizationId());
+        assertTrue(response.extensions().map().isEmpty());
+    }
+
+    @Test
+    public void testNoExtensionsFromTokenAndNullExtensions() throws Exception {
+        OAuthBearerClientInitialResponse response = new OAuthBearerClientInitialResponse("token", null);
+        assertTrue(response.extensions().map().isEmpty());
+    }
+
+    @Test
+    public void testValidateNullExtensions() throws Exception {
+        OAuthBearerClientInitialResponse.validateExtensions(null);
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/OAuthBearerSaslClientTest.java b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/OAuthBearerSaslClientTest.java
new file mode 100644
index 0000000000000..6d23f62059990
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/OAuthBearerSaslClientTest.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.common.security.oauthbearer.internals;
+
+import org.apache.kafka.common.config.ConfigException;
+import org.apache.kafka.common.security.auth.SaslExtensionsCallback;
+import org.apache.kafka.common.security.auth.AuthenticateCallbackHandler;
+import org.apache.kafka.common.security.auth.SaslExtensions;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerToken;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerTokenCallback;
+import org.junit.Test;
+
+import javax.security.auth.callback.Callback;
+import javax.security.auth.callback.UnsupportedCallbackException;
+import javax.security.auth.login.AppConfigurationEntry;
+import javax.security.sasl.SaslException;
+import java.nio.charset.StandardCharsets;
+import java.util.Collections;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.fail;
+
+public class OAuthBearerSaslClientTest {
+
+    private static final Map TEST_PROPERTIES = new LinkedHashMap() {
+        {
+            put("One", "1");
+            put("Two", "2");
+            put("Three", "3");
+        }
+    };
+    private SaslExtensions testExtensions = new SaslExtensions(TEST_PROPERTIES);
+    private final String errorMessage = "Error as expected!";
+
+    public class ExtensionsCallbackHandler implements AuthenticateCallbackHandler {
+        private boolean configured = false;
+        private boolean toThrow;
+
+        ExtensionsCallbackHandler(boolean toThrow) {
+            this.toThrow = toThrow;
+        }
+
+        public boolean configured() {
+            return configured;
+        }
+
+        @Override
+        public void configure(Map configs, String saslMechanism, List jaasConfigEntries) {
+            configured = true;
+        }
+
+        @Override
+        public void handle(Callback[] callbacks) throws UnsupportedCallbackException {
+            for (Callback callback : callbacks) {
+                if (callback instanceof OAuthBearerTokenCallback)
+                    ((OAuthBearerTokenCallback) callback).token(new OAuthBearerToken() {
+                        @Override
+                        public String value() {
+                            return "";
+                        }
+
+                        @Override
+                        public Set scope() {
+                            return Collections.emptySet();
+                        }
+
+                        @Override
+                        public long lifetimeMs() {
+                            return 100;
+                        }
+
+                        @Override
+                        public String principalName() {
+                            return "principalName";
+                        }
+
+                        @Override
+                        public Long startTimeMs() {
+                            return null;
+                        }
+                    });
+                else if (callback instanceof SaslExtensionsCallback) {
+                    if (toThrow)
+                        throw new ConfigException(errorMessage);
+                    else
+                        ((SaslExtensionsCallback) callback).extensions(testExtensions);
+                } else
+                    throw new UnsupportedCallbackException(callback);
+            }
+        }
+
+        @Override
+        public void close() {
+        }
+    }
+
+    @Test
+    public void testAttachesExtensionsToFirstClientMessage() throws Exception {
+        String expectedToken = new String(new OAuthBearerClientInitialResponse("", testExtensions).toBytes(), StandardCharsets.UTF_8);
+
+        OAuthBearerSaslClient client = new OAuthBearerSaslClient(new ExtensionsCallbackHandler(false));
+
+        String message = new String(client.evaluateChallenge("".getBytes()), StandardCharsets.UTF_8);
+
+        assertEquals(expectedToken, message);
+    }
+
+    @Test
+    public void testNoExtensionsDoesNotAttachAnythingToFirstClientMessage() throws Exception {
+        TEST_PROPERTIES.clear();
+        testExtensions = new SaslExtensions(TEST_PROPERTIES);
+        String expectedToken = new String(new OAuthBearerClientInitialResponse("", new SaslExtensions(TEST_PROPERTIES)).toBytes(), StandardCharsets.UTF_8);
+        OAuthBearerSaslClient client = new OAuthBearerSaslClient(new ExtensionsCallbackHandler(false));
+
+        String message = new String(client.evaluateChallenge("".getBytes()), StandardCharsets.UTF_8);
+
+        assertEquals(expectedToken, message);
+    }
+
+    @Test
+    public void testWrapsExtensionsCallbackHandlingErrorInSaslExceptionInFirstClientMessage() {
+        OAuthBearerSaslClient client = new OAuthBearerSaslClient(new ExtensionsCallbackHandler(true));
+        try {
+            client.evaluateChallenge("".getBytes());
+            fail("Should have failed with " + SaslException.class.getName());
+        } catch (SaslException e) {
+            // assert it has caught our expected exception
+            assertEquals(ConfigException.class, e.getCause().getClass());
+            assertEquals(errorMessage, e.getCause().getMessage());
+        }
+
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/OAuthBearerSaslServerTest.java b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/OAuthBearerSaslServerTest.java
new file mode 100644
index 0000000000000..5d8b84c17b02c
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/OAuthBearerSaslServerTest.java
@@ -0,0 +1,222 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.oauthbearer.internals;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.assertNull;
+
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+import javax.security.auth.callback.Callback;
+import javax.security.auth.callback.UnsupportedCallbackException;
+import javax.security.auth.login.LoginException;
+
+import org.apache.kafka.common.config.SaslConfigs;
+import org.apache.kafka.common.config.types.Password;
+import org.apache.kafka.common.errors.SaslAuthenticationException;
+import org.apache.kafka.common.security.JaasContext;
+import org.apache.kafka.common.security.auth.AuthenticateCallbackHandler;
+import org.apache.kafka.common.security.auth.SaslExtensions;
+import org.apache.kafka.common.security.authenticator.SaslInternalConfigs;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerExtensionsValidatorCallback;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerLoginModule;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerToken;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerTokenCallback;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerTokenMock;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerValidatorCallback;
+import org.apache.kafka.common.security.oauthbearer.internals.unsecured.OAuthBearerConfigException;
+import org.apache.kafka.common.security.oauthbearer.internals.unsecured.OAuthBearerUnsecuredLoginCallbackHandler;
+import org.apache.kafka.common.security.oauthbearer.internals.unsecured.OAuthBearerUnsecuredValidatorCallbackHandler;
+import org.junit.Before;
+import org.junit.Test;
+
+public class OAuthBearerSaslServerTest {
+    private static final String USER = "user";
+    private static final Map CONFIGS;
+    static {
+        String jaasConfigText = "org.apache.kafka.common.security.oauthbearer.OAuthBearerLoginModule Required"
+                + " unsecuredLoginStringClaim_sub=\"" + USER + "\";";
+        Map tmp = new HashMap<>();
+        tmp.put(SaslConfigs.SASL_JAAS_CONFIG, new Password(jaasConfigText));
+        CONFIGS = Collections.unmodifiableMap(tmp);
+    }
+    private static final AuthenticateCallbackHandler LOGIN_CALLBACK_HANDLER;
+    static {
+        LOGIN_CALLBACK_HANDLER = new OAuthBearerUnsecuredLoginCallbackHandler();
+        LOGIN_CALLBACK_HANDLER.configure(CONFIGS, OAuthBearerLoginModule.OAUTHBEARER_MECHANISM,
+                JaasContext.loadClientContext(CONFIGS).configurationEntries());
+    }
+    private static final AuthenticateCallbackHandler VALIDATOR_CALLBACK_HANDLER;
+    private static final AuthenticateCallbackHandler EXTENSIONS_VALIDATOR_CALLBACK_HANDLER;
+    static {
+        VALIDATOR_CALLBACK_HANDLER = new OAuthBearerUnsecuredValidatorCallbackHandler();
+        VALIDATOR_CALLBACK_HANDLER.configure(CONFIGS, OAuthBearerLoginModule.OAUTHBEARER_MECHANISM,
+                JaasContext.loadClientContext(CONFIGS).configurationEntries());
+        // only validate extensions "firstKey" and "secondKey"
+        EXTENSIONS_VALIDATOR_CALLBACK_HANDLER = new OAuthBearerUnsecuredValidatorCallbackHandler() {
+            @Override
+            public void handle(Callback[] callbacks) throws UnsupportedCallbackException {
+                for (Callback callback : callbacks) {
+                    if (callback instanceof OAuthBearerValidatorCallback) {
+                        OAuthBearerValidatorCallback validationCallback = (OAuthBearerValidatorCallback) callback;
+                        validationCallback.token(new OAuthBearerTokenMock());
+                    } else if (callback instanceof OAuthBearerExtensionsValidatorCallback) {
+                        OAuthBearerExtensionsValidatorCallback extensionsCallback = (OAuthBearerExtensionsValidatorCallback) callback;
+                        extensionsCallback.valid("firstKey");
+                        extensionsCallback.valid("secondKey");
+                    } else
+                        throw new UnsupportedCallbackException(callback);
+                }
+            }
+        };
+    }
+    private OAuthBearerSaslServer saslServer;
+
+    @Before
+    public void setUp() {
+        saslServer = new OAuthBearerSaslServer(VALIDATOR_CALLBACK_HANDLER);
+    }
+
+    @Test
+    public void noAuthorizationIdSpecified() throws Exception {
+        byte[] nextChallenge = saslServer
+                .evaluateResponse(clientInitialResponse(null));
+        // also asserts that no authentication error is thrown if OAuthBearerExtensionsValidatorCallback is not supported
+        assertTrue("Next challenge is not empty", nextChallenge.length == 0);
+    }
+
+    @Test
+    public void negotiatedProperty() throws Exception {
+        saslServer.evaluateResponse(clientInitialResponse(USER));
+        OAuthBearerToken token = (OAuthBearerToken) saslServer.getNegotiatedProperty("OAUTHBEARER.token");
+        assertNotNull(token);
+        assertEquals(token.lifetimeMs(),
+                saslServer.getNegotiatedProperty(SaslInternalConfigs.CREDENTIAL_LIFETIME_MS_SASL_NEGOTIATED_PROPERTY_KEY));
+    }
+
+    /**
+     * SASL Extensions that are validated by the callback handler should be accessible through the {@code #getNegotiatedProperty()} method
+     */
+    @Test
+    public void savesCustomExtensionAsNegotiatedProperty() throws Exception {
+        Map customExtensions = new HashMap<>();
+        customExtensions.put("firstKey", "value1");
+        customExtensions.put("secondKey", "value2");
+
+        byte[] nextChallenge = saslServer
+                .evaluateResponse(clientInitialResponse(null, false, customExtensions));
+
+        assertTrue("Next challenge is not empty", nextChallenge.length == 0);
+        assertEquals("value1", saslServer.getNegotiatedProperty("firstKey"));
+        assertEquals("value2", saslServer.getNegotiatedProperty("secondKey"));
+    }
+
+    /**
+     * SASL Extensions that were not recognized (neither validated nor invalidated)
+     * by the callback handler must not be accessible through the {@code #getNegotiatedProperty()} method
+     */
+    @Test
+    public void unrecognizedExtensionsAreNotSaved() throws Exception {
+        saslServer = new OAuthBearerSaslServer(EXTENSIONS_VALIDATOR_CALLBACK_HANDLER);
+        Map customExtensions = new HashMap<>();
+        customExtensions.put("firstKey", "value1");
+        customExtensions.put("secondKey", "value1");
+        customExtensions.put("thirdKey", "value1");
+
+        byte[] nextChallenge = saslServer
+                .evaluateResponse(clientInitialResponse(null, false, customExtensions));
+
+        assertTrue("Next challenge is not empty", nextChallenge.length == 0);
+        assertNull("Extensions not recognized by the server must be ignored", saslServer.getNegotiatedProperty("thirdKey"));
+    }
+
+    /**
+     * If the callback handler handles the `OAuthBearerExtensionsValidatorCallback`
+     *  and finds an invalid extension, SaslServer should throw an authentication exception
+     */
+    @Test(expected = SaslAuthenticationException.class)
+    public void throwsAuthenticationExceptionOnInvalidExtensions() throws Exception {
+        OAuthBearerUnsecuredValidatorCallbackHandler invalidHandler = new OAuthBearerUnsecuredValidatorCallbackHandler() {
+            @Override
+            public void handle(Callback[] callbacks) throws UnsupportedCallbackException {
+                for (Callback callback : callbacks) {
+                    if (callback instanceof OAuthBearerValidatorCallback) {
+                        OAuthBearerValidatorCallback validationCallback = (OAuthBearerValidatorCallback) callback;
+                        validationCallback.token(new OAuthBearerTokenMock());
+                    } else if (callback instanceof OAuthBearerExtensionsValidatorCallback) {
+                        OAuthBearerExtensionsValidatorCallback extensionsCallback = (OAuthBearerExtensionsValidatorCallback) callback;
+                        extensionsCallback.error("firstKey", "is not valid");
+                        extensionsCallback.error("secondKey", "is not valid either");
+                    } else
+                        throw new UnsupportedCallbackException(callback);
+                }
+            }
+        };
+        saslServer = new OAuthBearerSaslServer(invalidHandler);
+        Map customExtensions = new HashMap<>();
+        customExtensions.put("firstKey", "value");
+        customExtensions.put("secondKey", "value");
+
+        saslServer.evaluateResponse(clientInitialResponse(null, false, customExtensions));
+    }
+
+    @Test
+    public void authorizatonIdEqualsAuthenticationId() throws Exception {
+        byte[] nextChallenge = saslServer
+                .evaluateResponse(clientInitialResponse(USER));
+        assertTrue("Next challenge is not empty", nextChallenge.length == 0);
+    }
+
+    @Test(expected = SaslAuthenticationException.class)
+    public void authorizatonIdNotEqualsAuthenticationId() throws Exception {
+        saslServer.evaluateResponse(clientInitialResponse(USER + "x"));
+    }
+
+    @Test
+    public void illegalToken() throws Exception {
+        byte[] bytes = saslServer.evaluateResponse(clientInitialResponse(null, true, Collections.emptyMap()));
+        String challenge = new String(bytes, StandardCharsets.UTF_8);
+        assertEquals("{\"status\":\"invalid_token\"}", challenge);
+    }
+
+    private byte[] clientInitialResponse(String authorizationId)
+            throws OAuthBearerConfigException, IOException, UnsupportedCallbackException, LoginException {
+        return clientInitialResponse(authorizationId, false);
+    }
+
+    private byte[] clientInitialResponse(String authorizationId, boolean illegalToken)
+            throws OAuthBearerConfigException, IOException, UnsupportedCallbackException {
+        return clientInitialResponse(authorizationId, false, Collections.emptyMap());
+    }
+
+    private byte[] clientInitialResponse(String authorizationId, boolean illegalToken, Map customExtensions)
+            throws OAuthBearerConfigException, IOException, UnsupportedCallbackException {
+        OAuthBearerTokenCallback callback = new OAuthBearerTokenCallback();
+        LOGIN_CALLBACK_HANDLER.handle(new Callback[] {callback});
+        OAuthBearerToken token = callback.token();
+        String compactSerialization = token.value();
+
+        String tokenValue = compactSerialization + (illegalToken ? "AB" : "");
+        return new OAuthBearerClientInitialResponse(tokenValue, authorizationId, new SaslExtensions(customExtensions)).toBytes();
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/expiring/ExpiringCredentialRefreshConfigTest.java b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/expiring/ExpiringCredentialRefreshConfigTest.java
new file mode 100644
index 0000000000000..f188898aad26c
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/expiring/ExpiringCredentialRefreshConfigTest.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.oauthbearer.internals.expiring;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+import java.util.Collections;
+
+import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.common.config.SaslConfigs;
+import org.junit.Test;
+
+public class ExpiringCredentialRefreshConfigTest {
+    @Test
+    public void fromGoodConfig() {
+        ExpiringCredentialRefreshConfig expiringCredentialRefreshConfig = new ExpiringCredentialRefreshConfig(
+                new ConfigDef().withClientSaslSupport().parse(Collections.emptyMap()), true);
+        assertEquals(Double.valueOf(SaslConfigs.DEFAULT_LOGIN_REFRESH_WINDOW_FACTOR),
+                Double.valueOf(expiringCredentialRefreshConfig.loginRefreshWindowFactor()));
+        assertEquals(Double.valueOf(SaslConfigs.DEFAULT_LOGIN_REFRESH_WINDOW_JITTER),
+                Double.valueOf(expiringCredentialRefreshConfig.loginRefreshWindowJitter()));
+        assertEquals(Short.valueOf(SaslConfigs.DEFAULT_LOGIN_REFRESH_MIN_PERIOD_SECONDS),
+                Short.valueOf(expiringCredentialRefreshConfig.loginRefreshMinPeriodSeconds()));
+        assertEquals(Short.valueOf(SaslConfigs.DEFAULT_LOGIN_REFRESH_BUFFER_SECONDS),
+                Short.valueOf(expiringCredentialRefreshConfig.loginRefreshBufferSeconds()));
+        assertTrue(expiringCredentialRefreshConfig.loginRefreshReloginAllowedBeforeLogout());
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/expiring/ExpiringCredentialRefreshingLoginTest.java b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/expiring/ExpiringCredentialRefreshingLoginTest.java
new file mode 100644
index 0000000000000..f87405a1f8a77
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/expiring/ExpiringCredentialRefreshingLoginTest.java
@@ -0,0 +1,775 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.oauthbearer.internals.expiring;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.inOrder;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import java.util.ArrayList;
+import java.util.Date;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+
+import javax.security.auth.Subject;
+import javax.security.auth.login.AppConfigurationEntry;
+import javax.security.auth.login.Configuration;
+import javax.security.auth.login.LoginContext;
+import javax.security.auth.login.LoginException;
+
+import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.common.config.SaslConfigs;
+import org.apache.kafka.common.internals.KafkaFutureImpl;
+import org.apache.kafka.common.security.oauthbearer.internals.expiring.ExpiringCredentialRefreshingLogin.LoginContextFactory;
+import org.apache.kafka.common.utils.MockScheduler;
+import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.common.utils.Time;
+import org.junit.Test;
+import org.mockito.InOrder;
+import org.mockito.Mockito;
+
+public class ExpiringCredentialRefreshingLoginTest {
+    private static final Configuration EMPTY_WILDCARD_CONFIGURATION;
+    static {
+        EMPTY_WILDCARD_CONFIGURATION = new Configuration() {
+            @Override
+            public AppConfigurationEntry[] getAppConfigurationEntry(String name) {
+                return new AppConfigurationEntry[0]; // match any name
+            }
+        };
+    }
+
+    /*
+     * An ExpiringCredentialRefreshingLogin that we can tell explicitly to
+     * create/remove an expiring credential with specific
+     * create/expire/absoluteLastRefresh times
+     */
+    private static class TestExpiringCredentialRefreshingLogin extends ExpiringCredentialRefreshingLogin {
+        private ExpiringCredential expiringCredential;
+        private ExpiringCredential tmpExpiringCredential;
+        private final Time time;
+        private final long lifetimeMillis;
+        private final long absoluteLastRefreshTimeMs;
+        private final boolean clientReloginAllowedBeforeLogout;
+
+        public TestExpiringCredentialRefreshingLogin(ExpiringCredentialRefreshConfig refreshConfig,
+                LoginContextFactory loginContextFactory, Time time, final long lifetimeMillis,
+                final long absoluteLastRefreshMs, boolean clientReloginAllowedBeforeLogout) {
+            super("contextName", EMPTY_WILDCARD_CONFIGURATION, refreshConfig, null,
+                    TestExpiringCredentialRefreshingLogin.class, loginContextFactory, Objects.requireNonNull(time));
+            this.time = time;
+            this.lifetimeMillis = lifetimeMillis;
+            this.absoluteLastRefreshTimeMs = absoluteLastRefreshMs;
+            this.clientReloginAllowedBeforeLogout = clientReloginAllowedBeforeLogout;
+        }
+
+        public long getCreateMs() {
+            return time.milliseconds();
+        }
+
+        public long getExpireTimeMs() {
+            return time.milliseconds() + lifetimeMillis;
+        }
+
+        /*
+         * Invoke at login time
+         */
+        public void createNewExpiringCredential() {
+            if (!clientReloginAllowedBeforeLogout)
+                /*
+                 * Was preceded by logout
+                 */
+                expiringCredential = internalNewExpiringCredential();
+            else {
+                boolean initialLogin = expiringCredential == null;
+                if (initialLogin)
+                    // no logout immediately after the initial login
+                    this.expiringCredential = internalNewExpiringCredential();
+                else
+                    /*
+                     * This is at least the second invocation of login; we will move the credential
+                     * over upon logout, which should be invoked next
+                     */
+                    this.tmpExpiringCredential = internalNewExpiringCredential();
+            }
+        }
+
+        /*
+         * Invoke at logout time
+         */
+        public void clearExpiringCredential() {
+            if (!clientReloginAllowedBeforeLogout)
+                /*
+                 * Have not yet invoked login
+                 */
+                expiringCredential = null;
+            else
+                /*
+                 * login has already been invoked
+                 */
+                expiringCredential = tmpExpiringCredential;
+        }
+
+        @Override
+        public ExpiringCredential expiringCredential() {
+            return expiringCredential;
+        }
+
+        private ExpiringCredential internalNewExpiringCredential() {
+            return new ExpiringCredential() {
+                private final long createMs = getCreateMs();
+                private final long expireTimeMs = getExpireTimeMs();
+
+                @Override
+                public String principalName() {
+                    return "Created at " + new Date(createMs);
+                }
+
+                @Override
+                public Long startTimeMs() {
+                    return createMs;
+                }
+
+                @Override
+                public long expireTimeMs() {
+                    return expireTimeMs;
+                }
+
+                @Override
+                public Long absoluteLastRefreshTimeMs() {
+                    return absoluteLastRefreshTimeMs;
+                }
+
+                // useful in debugger
+                @Override
+                public String toString() {
+                    return String.format("startTimeMs=%d, expireTimeMs=%d, absoluteLastRefreshTimeMs=%s", startTimeMs(),
+                            expireTimeMs(), absoluteLastRefreshTimeMs());
+                }
+
+            };
+        }
+    }
+
+    /*
+     * A class that will forward all login/logout/getSubject() calls to a mock while
+     * also telling an instance of TestExpiringCredentialRefreshingLogin to
+     * create/remove an expiring credential upon login/logout(). Basically we are
+     * getting the functionality of a mock while simultaneously in the same method
+     * call performing creation/removal of expiring credentials.
+     */
+    private static class TestLoginContext extends LoginContext {
+        private final TestExpiringCredentialRefreshingLogin testExpiringCredentialRefreshingLogin;
+        private final LoginContext mockLoginContext;
+
+        public TestLoginContext(TestExpiringCredentialRefreshingLogin testExpiringCredentialRefreshingLogin,
+                LoginContext mockLoginContext) throws LoginException {
+            super("contextName", null, null, EMPTY_WILDCARD_CONFIGURATION);
+            this.testExpiringCredentialRefreshingLogin = Objects.requireNonNull(testExpiringCredentialRefreshingLogin);
+            // sanity check to make sure it is likely a mock
+            if (Objects.requireNonNull(mockLoginContext).getClass().equals(LoginContext.class)
+                    || mockLoginContext.getClass().equals(getClass()))
+                throw new IllegalArgumentException();
+            this.mockLoginContext = mockLoginContext;
+        }
+
+        @Override
+        public void login() throws LoginException {
+            /*
+             * Here is where we get the functionality of a mock while simultaneously
+             * performing the creation of an expiring credential
+             */
+            mockLoginContext.login();
+            testExpiringCredentialRefreshingLogin.createNewExpiringCredential();
+        }
+
+        @Override
+        public void logout() throws LoginException {
+            /*
+             * Here is where we get the functionality of a mock while simultaneously
+             * performing the removal of an expiring credential
+             */
+            mockLoginContext.logout();
+            testExpiringCredentialRefreshingLogin.clearExpiringCredential();
+        }
+
+        @Override
+        public Subject getSubject() {
+            // here we just need the functionality of a mock
+            return mockLoginContext.getSubject();
+        }
+    }
+
+    /*
+     * An implementation of LoginContextFactory that returns an instance of
+     * TestLoginContext
+     */
+    private static class TestLoginContextFactory extends LoginContextFactory {
+        private final KafkaFutureImpl refresherThreadStartedFuture = new KafkaFutureImpl<>();
+        private final KafkaFutureImpl refresherThreadDoneFuture = new KafkaFutureImpl<>();
+        private TestLoginContext testLoginContext;
+
+        public void configure(LoginContext mockLoginContext,
+                TestExpiringCredentialRefreshingLogin testExpiringCredentialRefreshingLogin) throws LoginException {
+            // sanity check to make sure it is likely a mock
+            if (Objects.requireNonNull(mockLoginContext).getClass().equals(LoginContext.class)
+                    || mockLoginContext.getClass().equals(TestLoginContext.class))
+                throw new IllegalArgumentException();
+            this.testLoginContext = new TestLoginContext(Objects.requireNonNull(testExpiringCredentialRefreshingLogin),
+                    mockLoginContext);
+        }
+
+        @Override
+        public LoginContext createLoginContext(ExpiringCredentialRefreshingLogin expiringCredentialRefreshingLogin) throws LoginException {
+            return new LoginContext("", null, null, EMPTY_WILDCARD_CONFIGURATION) {
+                private boolean loginSuccess = false;
+                @Override
+                public void login() throws LoginException {
+                    testLoginContext.login();
+                    loginSuccess = true;
+                }
+        
+                @Override
+                public void logout() throws LoginException {
+                    if (!loginSuccess)
+                        // will cause the refresher thread to exit
+                        throw new IllegalStateException("logout called without a successful login");
+                    testLoginContext.logout();
+                }
+        
+                @Override
+                public Subject getSubject() {
+                    return testLoginContext.getSubject();
+                }
+            };
+        }
+
+        @Override
+        public void refresherThreadStarted() {
+            refresherThreadStartedFuture.complete(null);
+        }
+
+        @Override
+        public void refresherThreadDone() {
+            refresherThreadDoneFuture.complete(null);
+        }
+
+        public Future refresherThreadStartedFuture() {
+            return refresherThreadStartedFuture;
+        }
+
+        public Future refresherThreadDoneFuture() {
+            return refresherThreadDoneFuture;
+        }
+    }
+
+    @Test
+    public void testRefresh() throws Exception {
+        for (int numExpectedRefreshes : new int[] {0, 1, 2}) {
+            for (boolean clientReloginAllowedBeforeLogout : new boolean[] {true, false}) {
+                Subject subject = new Subject();
+                final LoginContext mockLoginContext = mock(LoginContext.class);
+                when(mockLoginContext.getSubject()).thenReturn(subject);
+
+                MockTime mockTime = new MockTime();
+                long startMs = mockTime.milliseconds();
+                /*
+                 * Identify the lifetime of each expiring credential
+                 */
+                long lifetimeMinutes = 100L;
+                /*
+                 * Identify the point at which refresh will occur in that lifetime
+                 */
+                long refreshEveryMinutes = 80L;
+                /*
+                 * Set an absolute last refresh time that will cause the login thread to exit
+                 * after a certain number of re-logins (by adding an extra half of a refresh
+                 * interval).
+                 */
+                long absoluteLastRefreshMs = startMs + (1 + numExpectedRefreshes) * 1000 * 60 * refreshEveryMinutes
+                        - 1000 * 60 * refreshEveryMinutes / 2;
+                /*
+                 * Identify buffer time on either side for the refresh algorithm
+                 */
+                short minPeriodSeconds = (short) 0;
+                short bufferSeconds = minPeriodSeconds;
+
+                /*
+                 * Define some listeners so we can keep track of who gets done and when. All
+                 * added listeners should end up done except the last, extra one, which should
+                 * not.
+                 */
+                MockScheduler mockScheduler = new MockScheduler(mockTime);
+                List> waiters = addWaiters(mockScheduler, 1000 * 60 * refreshEveryMinutes,
+                        numExpectedRefreshes + 1);
+
+                // Create the ExpiringCredentialRefreshingLogin instance under test
+                TestLoginContextFactory testLoginContextFactory = new TestLoginContextFactory();
+                TestExpiringCredentialRefreshingLogin testExpiringCredentialRefreshingLogin = new TestExpiringCredentialRefreshingLogin(
+                        refreshConfigThatPerformsReloginEveryGivenPercentageOfLifetime(
+                                1.0 * refreshEveryMinutes / lifetimeMinutes, minPeriodSeconds, bufferSeconds,
+                                clientReloginAllowedBeforeLogout),
+                        testLoginContextFactory, mockTime, 1000 * 60 * lifetimeMinutes, absoluteLastRefreshMs,
+                        clientReloginAllowedBeforeLogout);
+                testLoginContextFactory.configure(mockLoginContext, testExpiringCredentialRefreshingLogin);
+
+                /*
+                 * Perform the login, wait up to a certain amount of time for the refresher
+                 * thread to exit, and make sure the correct calls happened at the correct times
+                 */
+                long expectedFinalMs = startMs + numExpectedRefreshes * 1000 * 60 * refreshEveryMinutes;
+                assertFalse(testLoginContextFactory.refresherThreadStartedFuture().isDone());
+                assertFalse(testLoginContextFactory.refresherThreadDoneFuture().isDone());
+                testExpiringCredentialRefreshingLogin.login();
+                assertTrue(testLoginContextFactory.refresherThreadStartedFuture().isDone());
+                testLoginContextFactory.refresherThreadDoneFuture().get(1L, TimeUnit.SECONDS);
+                assertEquals(expectedFinalMs, mockTime.milliseconds());
+                for (int i = 0; i < numExpectedRefreshes; ++i) {
+                    KafkaFutureImpl waiter = waiters.get(i);
+                    assertTrue(waiter.isDone());
+                    assertEquals((i + 1) * 1000 * 60 * refreshEveryMinutes, waiter.get().longValue() - startMs);
+                }
+                assertFalse(waiters.get(numExpectedRefreshes).isDone());
+
+                /*
+                 * We expect login() to be invoked followed by getSubject() and then ultimately followed by
+                 * numExpectedRefreshes pairs of either login()/logout() or logout()/login() calls
+                 */
+                InOrder inOrder = inOrder(mockLoginContext);
+                inOrder.verify(mockLoginContext).login();
+                inOrder.verify(mockLoginContext).getSubject();
+                for (int i = 0; i < numExpectedRefreshes; ++i) {
+                    if (clientReloginAllowedBeforeLogout) {
+                        inOrder.verify(mockLoginContext).login();
+                        inOrder.verify(mockLoginContext).logout();
+                    } else {
+                        inOrder.verify(mockLoginContext).logout();
+                        inOrder.verify(mockLoginContext).login();
+                    }
+                }
+                testExpiringCredentialRefreshingLogin.close();
+            }
+        }
+    }
+
+    @Test
+    public void testRefreshWithExpirationSmallerThanConfiguredBuffers() throws Exception {
+        int numExpectedRefreshes = 1;
+        boolean clientReloginAllowedBeforeLogout = true;
+        final LoginContext mockLoginContext = mock(LoginContext.class);
+        Subject subject = new Subject();
+        when(mockLoginContext.getSubject()).thenReturn(subject);
+
+        MockTime mockTime = new MockTime();
+        long startMs = mockTime.milliseconds();
+        /*
+         * Identify the lifetime of each expiring credential
+         */
+        long lifetimeMinutes = 10L;
+        /*
+         * Identify the point at which refresh will occur in that lifetime
+         */
+        long refreshEveryMinutes = 8L;
+        /*
+         * Set an absolute last refresh time that will cause the login thread to exit
+         * after a certain number of re-logins (by adding an extra half of a refresh
+         * interval).
+         */
+        long absoluteLastRefreshMs = startMs + (1 + numExpectedRefreshes) * 1000 * 60 * refreshEveryMinutes
+                - 1000 * 60 * refreshEveryMinutes / 2;
+        /*
+         * Identify buffer time on either side for the refresh algorithm that will cause
+         * the entire lifetime to be taken up. In other words, make sure there is no way
+         * to honor the buffers.
+         */
+        short minPeriodSeconds = (short) (1 + lifetimeMinutes * 60 / 2);
+        short bufferSeconds = minPeriodSeconds;
+
+        /*
+         * Define some listeners so we can keep track of who gets done and when. All
+         * added listeners should end up done except the last, extra one, which should
+         * not.
+         */
+        MockScheduler mockScheduler = new MockScheduler(mockTime);
+        List> waiters = addWaiters(mockScheduler, 1000 * 60 * refreshEveryMinutes,
+                numExpectedRefreshes + 1);
+
+        // Create the ExpiringCredentialRefreshingLogin instance under test
+        TestLoginContextFactory testLoginContextFactory = new TestLoginContextFactory();
+        TestExpiringCredentialRefreshingLogin testExpiringCredentialRefreshingLogin = new TestExpiringCredentialRefreshingLogin(
+                refreshConfigThatPerformsReloginEveryGivenPercentageOfLifetime(
+                        1.0 * refreshEveryMinutes / lifetimeMinutes, minPeriodSeconds, bufferSeconds,
+                        clientReloginAllowedBeforeLogout),
+                testLoginContextFactory, mockTime, 1000 * 60 * lifetimeMinutes, absoluteLastRefreshMs,
+                clientReloginAllowedBeforeLogout);
+        testLoginContextFactory.configure(mockLoginContext, testExpiringCredentialRefreshingLogin);
+
+        /*
+         * Perform the login, wait up to a certain amount of time for the refresher
+         * thread to exit, and make sure the correct calls happened at the correct times
+         */
+        long expectedFinalMs = startMs + numExpectedRefreshes * 1000 * 60 * refreshEveryMinutes;
+        assertFalse(testLoginContextFactory.refresherThreadStartedFuture().isDone());
+        assertFalse(testLoginContextFactory.refresherThreadDoneFuture().isDone());
+        testExpiringCredentialRefreshingLogin.login();
+        assertTrue(testLoginContextFactory.refresherThreadStartedFuture().isDone());
+        testLoginContextFactory.refresherThreadDoneFuture().get(1L, TimeUnit.SECONDS);
+        assertEquals(expectedFinalMs, mockTime.milliseconds());
+        for (int i = 0; i < numExpectedRefreshes; ++i) {
+            KafkaFutureImpl waiter = waiters.get(i);
+            assertTrue(waiter.isDone());
+            assertEquals((i + 1) * 1000 * 60 * refreshEveryMinutes, waiter.get().longValue() - startMs);
+        }
+        assertFalse(waiters.get(numExpectedRefreshes).isDone());
+
+        InOrder inOrder = inOrder(mockLoginContext);
+        inOrder.verify(mockLoginContext).login();
+        for (int i = 0; i < numExpectedRefreshes; ++i) {
+            inOrder.verify(mockLoginContext).login();
+            inOrder.verify(mockLoginContext).logout();
+        }
+    }
+
+    @Test
+    public void testRefreshWithExpirationSmallerThanConfiguredBuffersAndOlderCreateTime() throws Exception {
+        int numExpectedRefreshes = 1;
+        boolean clientReloginAllowedBeforeLogout = true;
+        final LoginContext mockLoginContext = mock(LoginContext.class);
+        Subject subject = new Subject();
+        when(mockLoginContext.getSubject()).thenReturn(subject);
+
+        MockTime mockTime = new MockTime();
+        long startMs = mockTime.milliseconds();
+        /*
+         * Identify the lifetime of each expiring credential
+         */
+        long lifetimeMinutes = 10L;
+        /*
+         * Identify the point at which refresh will occur in that lifetime
+         */
+        long refreshEveryMinutes = 8L;
+        /*
+         * Set an absolute last refresh time that will cause the login thread to exit
+         * after a certain number of re-logins (by adding an extra half of a refresh
+         * interval).
+         */
+        long absoluteLastRefreshMs = startMs + (1 + numExpectedRefreshes) * 1000 * 60 * refreshEveryMinutes
+                - 1000 * 60 * refreshEveryMinutes / 2;
+        /*
+         * Identify buffer time on either side for the refresh algorithm that will cause
+         * the entire lifetime to be taken up. In other words, make sure there is no way
+         * to honor the buffers.
+         */
+        short minPeriodSeconds = (short) (1 + lifetimeMinutes * 60 / 2);
+        short bufferSeconds = minPeriodSeconds;
+
+        /*
+         * Define some listeners so we can keep track of who gets done and when. All
+         * added listeners should end up done except the last, extra one, which should
+         * not.
+         */
+        MockScheduler mockScheduler = new MockScheduler(mockTime);
+        List> waiters = addWaiters(mockScheduler, 1000 * 60 * refreshEveryMinutes,
+                numExpectedRefreshes + 1);
+
+        // Create the ExpiringCredentialRefreshingLogin instance under test
+        TestLoginContextFactory testLoginContextFactory = new TestLoginContextFactory();
+        TestExpiringCredentialRefreshingLogin testExpiringCredentialRefreshingLogin = new TestExpiringCredentialRefreshingLogin(
+                refreshConfigThatPerformsReloginEveryGivenPercentageOfLifetime(
+                        1.0 * refreshEveryMinutes / lifetimeMinutes, minPeriodSeconds, bufferSeconds,
+                        clientReloginAllowedBeforeLogout),
+                testLoginContextFactory, mockTime, 1000 * 60 * lifetimeMinutes, absoluteLastRefreshMs,
+                clientReloginAllowedBeforeLogout) {
+
+            @Override
+            public long getCreateMs() {
+                return super.getCreateMs() - 1000 * 60 * 60; // distant past
+            }
+        };
+        testLoginContextFactory.configure(mockLoginContext, testExpiringCredentialRefreshingLogin);
+
+        /*
+         * Perform the login, wait up to a certain amount of time for the refresher
+         * thread to exit, and make sure the correct calls happened at the correct times
+         */
+        long expectedFinalMs = startMs + numExpectedRefreshes * 1000 * 60 * refreshEveryMinutes;
+        assertFalse(testLoginContextFactory.refresherThreadStartedFuture().isDone());
+        assertFalse(testLoginContextFactory.refresherThreadDoneFuture().isDone());
+        testExpiringCredentialRefreshingLogin.login();
+        assertTrue(testLoginContextFactory.refresherThreadStartedFuture().isDone());
+        testLoginContextFactory.refresherThreadDoneFuture().get(1L, TimeUnit.SECONDS);
+        assertEquals(expectedFinalMs, mockTime.milliseconds());
+        for (int i = 0; i < numExpectedRefreshes; ++i) {
+            KafkaFutureImpl waiter = waiters.get(i);
+            assertTrue(waiter.isDone());
+            assertEquals((i + 1) * 1000 * 60 * refreshEveryMinutes, waiter.get().longValue() - startMs);
+        }
+        assertFalse(waiters.get(numExpectedRefreshes).isDone());
+
+        InOrder inOrder = inOrder(mockLoginContext);
+        inOrder.verify(mockLoginContext).login();
+        for (int i = 0; i < numExpectedRefreshes; ++i) {
+            inOrder.verify(mockLoginContext).login();
+            inOrder.verify(mockLoginContext).logout();
+        }
+    }
+
+    @Test
+    public void testRefreshWithMinPeriodIntrusion() throws Exception {
+        int numExpectedRefreshes = 1;
+        boolean clientReloginAllowedBeforeLogout = true;
+        Subject subject = new Subject();
+        final LoginContext mockLoginContext = mock(LoginContext.class);
+        when(mockLoginContext.getSubject()).thenReturn(subject);
+
+        MockTime mockTime = new MockTime();
+        long startMs = mockTime.milliseconds();
+        /*
+         * Identify the lifetime of each expiring credential
+         */
+        long lifetimeMinutes = 10L;
+        /*
+         * Identify the point at which refresh will occur in that lifetime
+         */
+        long refreshEveryMinutes = 8L;
+        /*
+         * Set an absolute last refresh time that will cause the login thread to exit
+         * after a certain number of re-logins (by adding an extra half of a refresh
+         * interval).
+         */
+        long absoluteLastRefreshMs = startMs + (1 + numExpectedRefreshes) * 1000 * 60 * refreshEveryMinutes
+                - 1000 * 60 * refreshEveryMinutes / 2;
+
+        /*
+         * Identify a minimum period that will cause the refresh time to be delayed a
+         * bit.
+         */
+        int bufferIntrusionSeconds = 1;
+        short minPeriodSeconds = (short) (refreshEveryMinutes * 60 + bufferIntrusionSeconds);
+        short bufferSeconds = (short) 0;
+
+        /*
+         * Define some listeners so we can keep track of who gets done and when. All
+         * added listeners should end up done except the last, extra one, which should
+         * not.
+         */
+        MockScheduler mockScheduler = new MockScheduler(mockTime);
+        List> waiters = addWaiters(mockScheduler,
+                1000 * (60 * refreshEveryMinutes + bufferIntrusionSeconds), numExpectedRefreshes + 1);
+
+        // Create the ExpiringCredentialRefreshingLogin instance under test
+        TestLoginContextFactory testLoginContextFactory = new TestLoginContextFactory();
+        TestExpiringCredentialRefreshingLogin testExpiringCredentialRefreshingLogin = new TestExpiringCredentialRefreshingLogin(
+                refreshConfigThatPerformsReloginEveryGivenPercentageOfLifetime(
+                        1.0 * refreshEveryMinutes / lifetimeMinutes, minPeriodSeconds, bufferSeconds,
+                        clientReloginAllowedBeforeLogout),
+                testLoginContextFactory, mockTime, 1000 * 60 * lifetimeMinutes, absoluteLastRefreshMs,
+                clientReloginAllowedBeforeLogout);
+        testLoginContextFactory.configure(mockLoginContext, testExpiringCredentialRefreshingLogin);
+
+        /*
+         * Perform the login, wait up to a certain amount of time for the refresher
+         * thread to exit, and make sure the correct calls happened at the correct times
+         */
+        long expectedFinalMs = startMs
+                + numExpectedRefreshes * 1000 * (60 * refreshEveryMinutes + bufferIntrusionSeconds);
+        assertFalse(testLoginContextFactory.refresherThreadStartedFuture().isDone());
+        assertFalse(testLoginContextFactory.refresherThreadDoneFuture().isDone());
+        testExpiringCredentialRefreshingLogin.login();
+        assertTrue(testLoginContextFactory.refresherThreadStartedFuture().isDone());
+        testLoginContextFactory.refresherThreadDoneFuture().get(1L, TimeUnit.SECONDS);
+        assertEquals(expectedFinalMs, mockTime.milliseconds());
+        for (int i = 0; i < numExpectedRefreshes; ++i) {
+            KafkaFutureImpl waiter = waiters.get(i);
+            assertTrue(waiter.isDone());
+            assertEquals((i + 1) * 1000 * (60 * refreshEveryMinutes + bufferIntrusionSeconds),
+                    waiter.get().longValue() - startMs);
+        }
+        assertFalse(waiters.get(numExpectedRefreshes).isDone());
+
+        InOrder inOrder = inOrder(mockLoginContext);
+        inOrder.verify(mockLoginContext).login();
+        for (int i = 0; i < numExpectedRefreshes; ++i) {
+            inOrder.verify(mockLoginContext).login();
+            inOrder.verify(mockLoginContext).logout();
+        }
+    }
+
+    @Test
+    public void testRefreshWithPreExpirationBufferIntrusion() throws Exception {
+        int numExpectedRefreshes = 1;
+        boolean clientReloginAllowedBeforeLogout = true;
+        Subject subject = new Subject();
+        final LoginContext mockLoginContext = mock(LoginContext.class);
+        when(mockLoginContext.getSubject()).thenReturn(subject);
+
+        MockTime mockTime = new MockTime();
+        long startMs = mockTime.milliseconds();
+        /*
+         * Identify the lifetime of each expiring credential
+         */
+        long lifetimeMinutes = 10L;
+        /*
+         * Identify the point at which refresh will occur in that lifetime
+         */
+        long refreshEveryMinutes = 8L;
+        /*
+         * Set an absolute last refresh time that will cause the login thread to exit
+         * after a certain number of re-logins (by adding an extra half of a refresh
+         * interval).
+         */
+        long absoluteLastRefreshMs = startMs + (1 + numExpectedRefreshes) * 1000 * 60 * refreshEveryMinutes
+                - 1000 * 60 * refreshEveryMinutes / 2;
+        /*
+         * Identify a minimum period that will cause the refresh time to be delayed a
+         * bit.
+         */
+        int bufferIntrusionSeconds = 1;
+        short bufferSeconds = (short) ((lifetimeMinutes - refreshEveryMinutes) * 60 + bufferIntrusionSeconds);
+        short minPeriodSeconds = (short) 0;
+
+        /*
+         * Define some listeners so we can keep track of who gets done and when. All
+         * added listeners should end up done except the last, extra one, which should
+         * not.
+         */
+        MockScheduler mockScheduler = new MockScheduler(mockTime);
+        List> waiters = addWaiters(mockScheduler,
+                1000 * (60 * refreshEveryMinutes - bufferIntrusionSeconds), numExpectedRefreshes + 1);
+
+        // Create the ExpiringCredentialRefreshingLogin instance under test
+        TestLoginContextFactory testLoginContextFactory = new TestLoginContextFactory();
+        TestExpiringCredentialRefreshingLogin testExpiringCredentialRefreshingLogin = new TestExpiringCredentialRefreshingLogin(
+                refreshConfigThatPerformsReloginEveryGivenPercentageOfLifetime(
+                        1.0 * refreshEveryMinutes / lifetimeMinutes, minPeriodSeconds, bufferSeconds,
+                        clientReloginAllowedBeforeLogout),
+                testLoginContextFactory, mockTime, 1000 * 60 * lifetimeMinutes, absoluteLastRefreshMs,
+                clientReloginAllowedBeforeLogout);
+        testLoginContextFactory.configure(mockLoginContext, testExpiringCredentialRefreshingLogin);
+
+        /*
+         * Perform the login, wait up to a certain amount of time for the refresher
+         * thread to exit, and make sure the correct calls happened at the correct times
+         */
+        long expectedFinalMs = startMs
+                + numExpectedRefreshes * 1000 * (60 * refreshEveryMinutes - bufferIntrusionSeconds);
+        assertFalse(testLoginContextFactory.refresherThreadStartedFuture().isDone());
+        assertFalse(testLoginContextFactory.refresherThreadDoneFuture().isDone());
+        testExpiringCredentialRefreshingLogin.login();
+        assertTrue(testLoginContextFactory.refresherThreadStartedFuture().isDone());
+        testLoginContextFactory.refresherThreadDoneFuture().get(1L, TimeUnit.SECONDS);
+        assertEquals(expectedFinalMs, mockTime.milliseconds());
+        for (int i = 0; i < numExpectedRefreshes; ++i) {
+            KafkaFutureImpl waiter = waiters.get(i);
+            assertTrue(waiter.isDone());
+            assertEquals((i + 1) * 1000 * (60 * refreshEveryMinutes - bufferIntrusionSeconds),
+                    waiter.get().longValue() - startMs);
+        }
+        assertFalse(waiters.get(numExpectedRefreshes).isDone());
+
+        InOrder inOrder = inOrder(mockLoginContext);
+        inOrder.verify(mockLoginContext).login();
+        for (int i = 0; i < numExpectedRefreshes; ++i) {
+            inOrder.verify(mockLoginContext).login();
+            inOrder.verify(mockLoginContext).logout();
+        }
+    }
+
+    @Test
+    public void testLoginExceptionCausesCorrectLogout() throws Exception {
+        int numExpectedRefreshes = 3;
+        boolean clientReloginAllowedBeforeLogout = true;
+        Subject subject = new Subject();
+        final LoginContext mockLoginContext = mock(LoginContext.class);
+        when(mockLoginContext.getSubject()).thenReturn(subject);
+        Mockito.doNothing().doThrow(new LoginException()).doNothing().when(mockLoginContext).login();
+
+        MockTime mockTime = new MockTime();
+        long startMs = mockTime.milliseconds();
+        /*
+         * Identify the lifetime of each expiring credential
+         */
+        long lifetimeMinutes = 100L;
+        /*
+         * Identify the point at which refresh will occur in that lifetime
+         */
+        long refreshEveryMinutes = 80L;
+        /*
+         * Set an absolute last refresh time that will cause the login thread to exit
+         * after a certain number of re-logins (by adding an extra half of a refresh
+         * interval).
+         */
+        long absoluteLastRefreshMs = startMs + (1 + numExpectedRefreshes) * 1000 * 60 * refreshEveryMinutes
+                - 1000 * 60 * refreshEveryMinutes / 2;
+        /*
+         * Identify buffer time on either side for the refresh algorithm
+         */
+        short minPeriodSeconds = (short) 0;
+        short bufferSeconds = minPeriodSeconds;
+
+        // Create the ExpiringCredentialRefreshingLogin instance under test
+        TestLoginContextFactory testLoginContextFactory = new TestLoginContextFactory();
+        TestExpiringCredentialRefreshingLogin testExpiringCredentialRefreshingLogin = new TestExpiringCredentialRefreshingLogin(
+                refreshConfigThatPerformsReloginEveryGivenPercentageOfLifetime(
+                        1.0 * refreshEveryMinutes / lifetimeMinutes, minPeriodSeconds, bufferSeconds,
+                        clientReloginAllowedBeforeLogout),
+                testLoginContextFactory, mockTime, 1000 * 60 * lifetimeMinutes, absoluteLastRefreshMs,
+                clientReloginAllowedBeforeLogout);
+        testLoginContextFactory.configure(mockLoginContext, testExpiringCredentialRefreshingLogin);
+
+        /*
+         * Perform the login and wait up to a certain amount of time for the refresher
+         * thread to exit.  A timeout indicates the thread died due to logout()
+         * being invoked on an instance where the login() invocation had failed.
+         */
+        assertFalse(testLoginContextFactory.refresherThreadStartedFuture().isDone());
+        assertFalse(testLoginContextFactory.refresherThreadDoneFuture().isDone());
+        testExpiringCredentialRefreshingLogin.login();
+        assertTrue(testLoginContextFactory.refresherThreadStartedFuture().isDone());
+        testLoginContextFactory.refresherThreadDoneFuture().get(1L, TimeUnit.SECONDS);
+    }
+
+    private static List> addWaiters(MockScheduler mockScheduler, long refreshEveryMillis,
+            int numWaiters) {
+        List> retvalWaiters = new ArrayList<>(numWaiters);
+        for (int i = 1; i <= numWaiters; ++i) {
+            KafkaFutureImpl waiter = new KafkaFutureImpl();
+            mockScheduler.addWaiter(i * refreshEveryMillis, waiter);
+            retvalWaiters.add(waiter);
+        }
+        return retvalWaiters;
+    }
+
+    private static ExpiringCredentialRefreshConfig refreshConfigThatPerformsReloginEveryGivenPercentageOfLifetime(
+            double refreshWindowFactor, short minPeriodSeconds, short bufferSeconds,
+            boolean clientReloginAllowedBeforeLogout) {
+        Map configs = new HashMap<>();
+        configs.put(SaslConfigs.SASL_LOGIN_REFRESH_WINDOW_FACTOR, refreshWindowFactor);
+        configs.put(SaslConfigs.SASL_LOGIN_REFRESH_WINDOW_JITTER, 0);
+        configs.put(SaslConfigs.SASL_LOGIN_REFRESH_MIN_PERIOD_SECONDS, minPeriodSeconds);
+        configs.put(SaslConfigs.SASL_LOGIN_REFRESH_BUFFER_SECONDS, bufferSeconds);
+        return new ExpiringCredentialRefreshConfig(new ConfigDef().withClientSaslSupport().parse(configs),
+                clientReloginAllowedBeforeLogout);
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerScopeUtilsTest.java b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerScopeUtilsTest.java
new file mode 100644
index 0000000000000..4fb301167b022
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerScopeUtilsTest.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 org.apache.kafka.common.security.oauthbearer.internals.unsecured;
+
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+
+import java.util.List;
+
+import org.junit.Test;
+
+public class OAuthBearerScopeUtilsTest {
+    @Test
+    public void validScope() {
+        for (String validScope : new String[] {"", "   ", "scope1", " scope1 ", "scope1 Scope2", "scope1   Scope2"}) {
+            List parsedScope = OAuthBearerScopeUtils.parseScope(validScope);
+            if (validScope.trim().isEmpty()) {
+                assertTrue(parsedScope.isEmpty());
+            } else if (validScope.contains("Scope2")) {
+                assertTrue(parsedScope.size() == 2 && parsedScope.get(0).equals("scope1")
+                        && parsedScope.get(1).equals("Scope2"));
+            } else {
+                assertTrue(parsedScope.size() == 1 && parsedScope.get(0).equals("scope1"));
+            }
+        }
+    }
+
+    @Test
+    public void invalidScope() {
+        for (String invalidScope : new String[] {"\"foo", "\\foo"}) {
+            try {
+                OAuthBearerScopeUtils.parseScope(invalidScope);
+                fail("did not detect invalid scope: " + invalidScope);
+            } catch (OAuthBearerConfigException expected) {
+                // empty
+            }
+        }
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerUnsecuredJwsTest.java b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerUnsecuredJwsTest.java
new file mode 100644
index 0000000000000..a99e75bfb8ec8
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerUnsecuredJwsTest.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.common.security.oauthbearer.internals.unsecured;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
+import java.util.Base64;
+import java.util.Base64.Encoder;
+import java.util.HashSet;
+import java.util.List;
+
+import org.junit.Test;
+
+public class OAuthBearerUnsecuredJwsTest {
+    private static final String QUOTE = "\"";
+    private static final String HEADER_COMPACT_SERIALIZATION = Base64.getUrlEncoder().withoutPadding()
+            .encodeToString("{\"alg\":\"none\"}".getBytes(StandardCharsets.UTF_8)) + ".";
+
+    @Test
+    public void validClaims() throws OAuthBearerIllegalTokenException {
+        double issuedAtSeconds = 100.1;
+        double expirationTimeSeconds = 300.3;
+        StringBuilder sb = new StringBuilder("{");
+        appendJsonText(sb, "sub", "SUBJECT");
+        appendCommaJsonText(sb, "iat", issuedAtSeconds);
+        appendCommaJsonText(sb, "exp", expirationTimeSeconds);
+        sb.append("}");
+        String compactSerialization = HEADER_COMPACT_SERIALIZATION
+                + Base64.getUrlEncoder().withoutPadding().encodeToString(sb.toString().getBytes(StandardCharsets.UTF_8))
+                + ".";
+        OAuthBearerUnsecuredJws testJwt = new OAuthBearerUnsecuredJws(compactSerialization, "sub", "scope");
+        assertEquals(compactSerialization, testJwt.value());
+        assertEquals("sub", testJwt.principalClaimName());
+        assertEquals(1, testJwt.header().size());
+        assertEquals("none", testJwt.header().get("alg"));
+        assertEquals("scope", testJwt.scopeClaimName());
+        assertEquals(expirationTimeSeconds, testJwt.expirationTime());
+        assertTrue(testJwt.isClaimType("exp", Number.class));
+        assertEquals(issuedAtSeconds, testJwt.issuedAt());
+        assertEquals("SUBJECT", testJwt.subject());
+    }
+
+    @Test
+    public void validCompactSerialization() {
+        String subject = "foo";
+        long issuedAt = 100;
+        long expirationTime = issuedAt + 60 * 60;
+        List scope = Arrays.asList("scopeValue1", "scopeValue2");
+        String validCompactSerialization = compactSerialization(subject, issuedAt, expirationTime, scope);
+        OAuthBearerUnsecuredJws jws = new OAuthBearerUnsecuredJws(validCompactSerialization, "sub", "scope");
+        assertEquals(1, jws.header().size());
+        assertEquals("none", jws.header().get("alg"));
+        assertEquals(4, jws.claims().size());
+        assertEquals(subject, jws.claims().get("sub"));
+        assertEquals(subject, jws.principalName());
+        assertEquals(issuedAt, Number.class.cast(jws.claims().get("iat")).longValue());
+        assertEquals(expirationTime, Number.class.cast(jws.claims().get("exp")).longValue());
+        assertEquals(expirationTime * 1000, jws.lifetimeMs());
+        assertEquals(scope, jws.claims().get("scope"));
+        assertEquals(new HashSet<>(scope), jws.scope());
+        assertEquals(3, jws.splits().size());
+        assertEquals(validCompactSerialization.split("\\.")[0], jws.splits().get(0));
+        assertEquals(validCompactSerialization.split("\\.")[1], jws.splits().get(1));
+        assertEquals("", jws.splits().get(2));
+    }
+
+    @Test(expected = OAuthBearerIllegalTokenException.class)
+    public void missingPrincipal() {
+        String subject = null;
+        long issuedAt = 100;
+        Long expirationTime = null;
+        List scope = Arrays.asList("scopeValue1", "scopeValue2");
+        String validCompactSerialization = compactSerialization(subject, issuedAt, expirationTime, scope);
+        new OAuthBearerUnsecuredJws(validCompactSerialization, "sub", "scope");
+    }
+
+    @Test(expected = OAuthBearerIllegalTokenException.class)
+    public void blankPrincipalName() {
+        String subject = "   ";
+        long issuedAt = 100;
+        long expirationTime = issuedAt + 60 * 60;
+        List scope = Arrays.asList("scopeValue1", "scopeValue2");
+        String validCompactSerialization = compactSerialization(subject, issuedAt, expirationTime, scope);
+        new OAuthBearerUnsecuredJws(validCompactSerialization, "sub", "scope");
+    }
+
+    private static String compactSerialization(String subject, Long issuedAt, Long expirationTime, List scope) {
+        Encoder encoder = Base64.getUrlEncoder().withoutPadding();
+        String algorithm = "none";
+        String headerJson = "{\"alg\":\"" + algorithm + "\"}";
+        String encodedHeader = encoder.encodeToString(headerJson.getBytes(StandardCharsets.UTF_8));
+        String subjectJson = subject != null ? "\"sub\":\"" + subject + "\"" : null;
+        String issuedAtJson = issuedAt != null ? "\"iat\":" + issuedAt.longValue() : null;
+        String expirationTimeJson = expirationTime != null ? "\"exp\":" + expirationTime.longValue() : null;
+        String scopeJson = scope != null ? scopeJson(scope) : null;
+        String claimsJson = claimsJson(subjectJson, issuedAtJson, expirationTimeJson, scopeJson);
+        String encodedClaims = encoder.encodeToString(claimsJson.getBytes(StandardCharsets.UTF_8));
+        return encodedHeader + "." + encodedClaims + ".";
+    }
+
+    private static String claimsJson(String... jsonValues) {
+        StringBuilder claimsJsonBuilder = new StringBuilder("{");
+        int initialLength = claimsJsonBuilder.length();
+        for (String jsonValue : jsonValues) {
+            if (jsonValue != null) {
+                if (claimsJsonBuilder.length() > initialLength)
+                    claimsJsonBuilder.append(',');
+                claimsJsonBuilder.append(jsonValue);
+            }
+        }
+        claimsJsonBuilder.append('}');
+        return claimsJsonBuilder.toString();
+    }
+
+    private static String scopeJson(List scope) {
+        StringBuilder scopeJsonBuilder = new StringBuilder("\"scope\":[");
+        int initialLength = scopeJsonBuilder.length();
+        for (String scopeValue : scope) {
+            if (scopeJsonBuilder.length() > initialLength)
+                scopeJsonBuilder.append(',');
+            scopeJsonBuilder.append('"').append(scopeValue).append('"');
+        }
+        scopeJsonBuilder.append(']');
+        return scopeJsonBuilder.toString();
+    }
+
+    private static void appendCommaJsonText(StringBuilder sb, String claimName, Number claimValue) {
+        sb.append(',').append(QUOTE).append(escape(claimName)).append(QUOTE).append(":").append(claimValue);
+    }
+
+    private static void appendJsonText(StringBuilder sb, String claimName, String claimValue) {
+        sb.append(QUOTE).append(escape(claimName)).append(QUOTE).append(":").append(QUOTE).append(escape(claimValue))
+                .append(QUOTE);
+    }
+
+    private static String escape(String jsonStringValue) {
+        return jsonStringValue.replace("\"", "\\\"").replace("\\", "\\\\");
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerUnsecuredLoginCallbackHandlerTest.java b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerUnsecuredLoginCallbackHandlerTest.java
new file mode 100644
index 0000000000000..34b8209abe59c
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerUnsecuredLoginCallbackHandlerTest.java
@@ -0,0 +1,160 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.oauthbearer.internals.unsecured;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+
+import java.io.IOException;
+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 javax.security.auth.callback.Callback;
+import javax.security.auth.callback.UnsupportedCallbackException;
+
+import org.apache.kafka.common.security.auth.SaslExtensionsCallback;
+import org.apache.kafka.common.security.authenticator.TestJaasConfig;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerLoginModule;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerTokenCallback;
+import org.apache.kafka.common.utils.MockTime;
+import org.junit.Test;
+
+public class OAuthBearerUnsecuredLoginCallbackHandlerTest {
+
+    @Test
+    public void addsExtensions() throws IOException, UnsupportedCallbackException {
+        Map options = new HashMap<>();
+        options.put("unsecuredLoginExtension_testId", "1");
+        OAuthBearerUnsecuredLoginCallbackHandler callbackHandler = createCallbackHandler(options, new MockTime());
+        SaslExtensionsCallback callback = new SaslExtensionsCallback();
+
+        callbackHandler.handle(new Callback[] {callback});
+
+        assertEquals("1", callback.extensions().map().get("testId"));
+    }
+
+    @Test(expected = IOException.class)
+    public void throwsErrorOnInvalidExtensionName() throws IOException, UnsupportedCallbackException {
+        Map options = new HashMap<>();
+        options.put("unsecuredLoginExtension_test.Id", "1");
+        OAuthBearerUnsecuredLoginCallbackHandler callbackHandler = createCallbackHandler(options, new MockTime());
+        SaslExtensionsCallback callback = new SaslExtensionsCallback();
+
+        callbackHandler.handle(new Callback[] {callback});
+    }
+
+    @Test(expected = IOException.class)
+    public void throwsErrorOnInvalidExtensionValue() throws IOException, UnsupportedCallbackException {
+        Map options = new HashMap<>();
+        options.put("unsecuredLoginExtension_testId", "Çalifornia");
+        OAuthBearerUnsecuredLoginCallbackHandler callbackHandler = createCallbackHandler(options, new MockTime());
+        SaslExtensionsCallback callback = new SaslExtensionsCallback();
+
+        callbackHandler.handle(new Callback[] {callback});
+    }
+
+    @Test
+    public void minimalToken() throws IOException, UnsupportedCallbackException {
+        Map options = new HashMap<>();
+        String user = "user";
+        options.put("unsecuredLoginStringClaim_sub", user);
+        MockTime mockTime = new MockTime();
+        OAuthBearerUnsecuredLoginCallbackHandler callbackHandler = createCallbackHandler(options, mockTime);
+        OAuthBearerTokenCallback callback = new OAuthBearerTokenCallback();
+        callbackHandler.handle(new Callback[] {callback});
+        OAuthBearerUnsecuredJws jws = (OAuthBearerUnsecuredJws) callback.token();
+        assertNotNull("create token failed", jws);
+        long startMs = mockTime.milliseconds();
+        confirmCorrectValues(jws, user, startMs, 1000 * 60 * 60);
+        assertEquals(new HashSet<>(Arrays.asList("sub", "iat", "exp")), jws.claims().keySet());
+    }
+
+    @SuppressWarnings("unchecked")
+    @Test
+    public void validOptionsWithExplicitOptionValues()
+            throws IOException, UnsupportedCallbackException {
+        String explicitScope1 = "scope1";
+        String explicitScope2 = "scope2";
+        String explicitScopeClaimName = "putScopeInHere";
+        String principalClaimName = "principal";
+        final String[] scopeClaimNameOptionValues = {null, explicitScopeClaimName};
+        for (String scopeClaimNameOptionValue : scopeClaimNameOptionValues) {
+            Map options = new HashMap<>();
+            String user = "user";
+            options.put("unsecuredLoginStringClaim_" + principalClaimName, user);
+            options.put("unsecuredLoginListClaim_" + "list", ",1,2,");
+            options.put("unsecuredLoginListClaim_" + "emptyList1", "");
+            options.put("unsecuredLoginListClaim_" + "emptyList2", ",");
+            options.put("unsecuredLoginNumberClaim_" + "number", "1");
+            long lifetmeSeconds = 10000;
+            options.put("unsecuredLoginLifetimeSeconds", String.valueOf(lifetmeSeconds));
+            options.put("unsecuredLoginPrincipalClaimName", principalClaimName);
+            if (scopeClaimNameOptionValue != null)
+                options.put("unsecuredLoginScopeClaimName", scopeClaimNameOptionValue);
+            String actualScopeClaimName = scopeClaimNameOptionValue == null ? "scope" : explicitScopeClaimName;
+            options.put("unsecuredLoginListClaim_" + actualScopeClaimName,
+                    String.format("|%s|%s", explicitScope1, explicitScope2));
+            MockTime mockTime = new MockTime();
+            OAuthBearerUnsecuredLoginCallbackHandler callbackHandler = createCallbackHandler(options, mockTime);
+            OAuthBearerTokenCallback callback = new OAuthBearerTokenCallback();
+            callbackHandler.handle(new Callback[] {callback});
+            OAuthBearerUnsecuredJws jws = (OAuthBearerUnsecuredJws) callback.token();
+            assertNotNull("create token failed", jws);
+            long startMs = mockTime.milliseconds();
+            confirmCorrectValues(jws, user, startMs, lifetmeSeconds * 1000);
+            Map claims = jws.claims();
+            assertEquals(new HashSet<>(Arrays.asList(actualScopeClaimName, principalClaimName, "iat", "exp", "number",
+                    "list", "emptyList1", "emptyList2")), claims.keySet());
+            assertEquals(new HashSet<>(Arrays.asList(explicitScope1, explicitScope2)),
+                    new HashSet<>((List) claims.get(actualScopeClaimName)));
+            assertEquals(new HashSet<>(Arrays.asList(explicitScope1, explicitScope2)), jws.scope());
+            assertEquals(1.0, jws.claim("number", Number.class));
+            assertEquals(Arrays.asList("1", "2", ""), jws.claim("list", List.class));
+            assertEquals(Collections.emptyList(), jws.claim("emptyList1", List.class));
+            assertEquals(Collections.emptyList(), jws.claim("emptyList2", List.class));
+        }
+    }
+
+    @SuppressWarnings({"unchecked", "rawtypes"})
+    private static OAuthBearerUnsecuredLoginCallbackHandler createCallbackHandler(Map options,
+            MockTime mockTime) {
+        TestJaasConfig config = new TestJaasConfig();
+        config.createOrUpdateEntry("KafkaClient", "org.apache.kafka.common.security.oauthbearer.OAuthBearerLoginModule",
+                (Map) options);
+        OAuthBearerUnsecuredLoginCallbackHandler callbackHandler = new OAuthBearerUnsecuredLoginCallbackHandler();
+        callbackHandler.time(mockTime);
+        callbackHandler.configure(Collections.emptyMap(), OAuthBearerLoginModule.OAUTHBEARER_MECHANISM,
+                Arrays.asList(config.getAppConfigurationEntry("KafkaClient")[0]));
+        return callbackHandler;
+    }
+
+    private static void confirmCorrectValues(OAuthBearerUnsecuredJws jws, String user, long startMs,
+            long lifetimeSeconds) throws OAuthBearerIllegalTokenException {
+        Map header = jws.header();
+        assertEquals(header.size(), 1);
+        assertEquals("none", header.get("alg"));
+        assertEquals(user != null ? user : "", jws.principalName());
+        assertEquals(Long.valueOf(startMs), jws.startTimeMs());
+        assertEquals(startMs, Math.round(jws.issuedAt().doubleValue() * 1000));
+        assertEquals(startMs + lifetimeSeconds, jws.lifetimeMs());
+        assertEquals(jws.lifetimeMs(), Math.round(jws.expirationTime().doubleValue() * 1000));
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerUnsecuredValidatorCallbackHandlerTest.java b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerUnsecuredValidatorCallbackHandlerTest.java
new file mode 100644
index 0000000000000..2449a83ecf25c
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerUnsecuredValidatorCallbackHandlerTest.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.common.security.oauthbearer.internals.unsecured;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
+import java.util.Base64;
+import java.util.Base64.Encoder;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+import javax.security.auth.callback.Callback;
+import javax.security.auth.callback.UnsupportedCallbackException;
+
+import org.apache.kafka.common.security.authenticator.TestJaasConfig;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerLoginModule;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerValidatorCallback;
+import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.common.utils.Time;
+import org.junit.Test;
+
+public class OAuthBearerUnsecuredValidatorCallbackHandlerTest {
+    private static final String UNSECURED_JWT_HEADER_JSON = "{" + claimOrHeaderText("alg", "none") + "}";
+    private static final Time MOCK_TIME = new MockTime();
+    private static final String QUOTE = "\"";
+    private static final String PRINCIPAL_CLAIM_VALUE = "username";
+    private static final String PRINCIPAL_CLAIM_TEXT = claimOrHeaderText("principal", PRINCIPAL_CLAIM_VALUE);
+    private static final String SUB_CLAIM_TEXT = claimOrHeaderText("sub", PRINCIPAL_CLAIM_VALUE);
+    private static final String BAD_PRINCIPAL_CLAIM_TEXT = claimOrHeaderText("principal", 1);
+    private static final long LIFETIME_SECONDS_TO_USE = 1000 * 60 * 60;
+    private static final String EXPIRATION_TIME_CLAIM_TEXT = expClaimText(LIFETIME_SECONDS_TO_USE);
+    private static final String TOO_EARLY_EXPIRATION_TIME_CLAIM_TEXT = expClaimText(0);
+    private static final String ISSUED_AT_CLAIM_TEXT = claimOrHeaderText("iat", MOCK_TIME.milliseconds() / 1000.0);
+    private static final String SCOPE_CLAIM_TEXT = claimOrHeaderText("scope", "scope1");
+    private static final Map MODULE_OPTIONS_MAP_NO_SCOPE_REQUIRED;
+    static {
+        Map tmp = new HashMap<>();
+        tmp.put("unsecuredValidatorPrincipalClaimName", "principal");
+        tmp.put("unsecuredValidatorAllowableClockSkewMs", "1");
+        MODULE_OPTIONS_MAP_NO_SCOPE_REQUIRED = Collections.unmodifiableMap(tmp);
+    }
+    private static final Map MODULE_OPTIONS_MAP_REQUIRE_EXISTING_SCOPE;
+    static {
+        Map tmp = new HashMap<>();
+        tmp.put("unsecuredValidatorRequiredScope", "scope1");
+        MODULE_OPTIONS_MAP_REQUIRE_EXISTING_SCOPE = Collections.unmodifiableMap(tmp);
+    }
+    private static final Map MODULE_OPTIONS_MAP_REQUIRE_ADDITIONAL_SCOPE;
+    static {
+        Map tmp = new HashMap<>();
+        tmp.put("unsecuredValidatorRequiredScope", "scope1 scope2");
+        MODULE_OPTIONS_MAP_REQUIRE_ADDITIONAL_SCOPE = Collections.unmodifiableMap(tmp);
+    }
+
+    @Test
+    public void validToken() {
+        for (final boolean includeOptionalIssuedAtClaim : new boolean[] {true, false}) {
+            String claimsJson = "{" + PRINCIPAL_CLAIM_TEXT + comma(EXPIRATION_TIME_CLAIM_TEXT)
+                    + (includeOptionalIssuedAtClaim ? comma(ISSUED_AT_CLAIM_TEXT) : "") + "}";
+            Object validationResult = validationResult(UNSECURED_JWT_HEADER_JSON, claimsJson,
+                    MODULE_OPTIONS_MAP_NO_SCOPE_REQUIRED);
+            assertTrue(validationResult instanceof OAuthBearerValidatorCallback);
+            assertTrue(((OAuthBearerValidatorCallback) validationResult).token() instanceof OAuthBearerUnsecuredJws);
+        }
+    }
+
+    @Test
+    public void badOrMissingPrincipal() throws IOException, UnsupportedCallbackException {
+        for (boolean exists : new boolean[] {true, false}) {
+            String claimsJson = "{" + EXPIRATION_TIME_CLAIM_TEXT + (exists ? comma(BAD_PRINCIPAL_CLAIM_TEXT) : "")
+                    + "}";
+            confirmFailsValidation(UNSECURED_JWT_HEADER_JSON, claimsJson, MODULE_OPTIONS_MAP_NO_SCOPE_REQUIRED);
+        }
+    }
+
+    @Test
+    public void tooEarlyExpirationTime() throws IOException, UnsupportedCallbackException {
+        String claimsJson = "{" + PRINCIPAL_CLAIM_TEXT + comma(ISSUED_AT_CLAIM_TEXT)
+                + comma(TOO_EARLY_EXPIRATION_TIME_CLAIM_TEXT) + "}";
+        confirmFailsValidation(UNSECURED_JWT_HEADER_JSON, claimsJson, MODULE_OPTIONS_MAP_NO_SCOPE_REQUIRED);
+    }
+
+    @Test
+    public void includesRequiredScope() {
+        String claimsJson = "{" + SUB_CLAIM_TEXT + comma(EXPIRATION_TIME_CLAIM_TEXT) + comma(SCOPE_CLAIM_TEXT) + "}";
+        Object validationResult = validationResult(UNSECURED_JWT_HEADER_JSON, claimsJson,
+                MODULE_OPTIONS_MAP_REQUIRE_EXISTING_SCOPE);
+        assertTrue(validationResult instanceof OAuthBearerValidatorCallback);
+        assertTrue(((OAuthBearerValidatorCallback) validationResult).token() instanceof OAuthBearerUnsecuredJws);
+    }
+
+    @Test
+    public void missingRequiredScope() throws IOException, UnsupportedCallbackException {
+        String claimsJson = "{" + SUB_CLAIM_TEXT + comma(EXPIRATION_TIME_CLAIM_TEXT) + comma(SCOPE_CLAIM_TEXT) + "}";
+        confirmFailsValidation(UNSECURED_JWT_HEADER_JSON, claimsJson, MODULE_OPTIONS_MAP_REQUIRE_ADDITIONAL_SCOPE,
+                "[scope1, scope2]");
+    }
+
+    private static void confirmFailsValidation(String headerJson, String claimsJson,
+            Map moduleOptionsMap) throws OAuthBearerConfigException, OAuthBearerIllegalTokenException,
+            IOException, UnsupportedCallbackException {
+        confirmFailsValidation(headerJson, claimsJson, moduleOptionsMap, null);
+    }
+
+    private static void confirmFailsValidation(String headerJson, String claimsJson,
+            Map moduleOptionsMap, String optionalFailureScope) throws OAuthBearerConfigException,
+            OAuthBearerIllegalTokenException {
+        Object validationResultObj = validationResult(headerJson, claimsJson, moduleOptionsMap);
+        assertTrue(validationResultObj instanceof OAuthBearerValidatorCallback);
+        OAuthBearerValidatorCallback callback = (OAuthBearerValidatorCallback) validationResultObj;
+        assertNull(callback.token());
+        assertNull(callback.errorOpenIDConfiguration());
+        if (optionalFailureScope == null) {
+            assertEquals("invalid_token", callback.errorStatus());
+            assertNull(callback.errorScope());
+        } else {
+            assertEquals("insufficient_scope", callback.errorStatus());
+            assertEquals(optionalFailureScope, callback.errorScope());
+        }
+    }
+
+    private static Object validationResult(String headerJson, String claimsJson, Map moduleOptionsMap) {
+        Encoder urlEncoderNoPadding = Base64.getUrlEncoder().withoutPadding();
+        try {
+            String tokenValue = String.format("%s.%s.",
+                    urlEncoderNoPadding.encodeToString(headerJson.getBytes(StandardCharsets.UTF_8)),
+                    urlEncoderNoPadding.encodeToString(claimsJson.getBytes(StandardCharsets.UTF_8)));
+            OAuthBearerValidatorCallback callback = new OAuthBearerValidatorCallback(tokenValue);
+            createCallbackHandler(moduleOptionsMap).handle(new Callback[] {callback});
+            return callback;
+        } catch (Exception e) {
+            return e;
+        }
+    }
+
+    @SuppressWarnings({"unchecked", "rawtypes"})
+    private static OAuthBearerUnsecuredValidatorCallbackHandler createCallbackHandler(Map options) {
+        TestJaasConfig config = new TestJaasConfig();
+        config.createOrUpdateEntry("KafkaClient", "org.apache.kafka.common.security.oauthbearer.OAuthBearerLoginModule",
+                (Map) options);
+        OAuthBearerUnsecuredValidatorCallbackHandler callbackHandler = new OAuthBearerUnsecuredValidatorCallbackHandler();
+        callbackHandler.configure(Collections.emptyMap(), OAuthBearerLoginModule.OAUTHBEARER_MECHANISM,
+                Arrays.asList(config.getAppConfigurationEntry("KafkaClient")[0]));
+        return callbackHandler;
+    }
+
+    private static String comma(String value) {
+        return "," + value;
+    }
+
+    private static String claimOrHeaderText(String claimName, Number claimValue) {
+        return QUOTE + claimName + QUOTE + ":" + claimValue;
+    }
+
+    private static String claimOrHeaderText(String claimName, String claimValue) {
+        return QUOTE + claimName + QUOTE + ":" + QUOTE + claimValue + QUOTE;
+    }
+
+    private static String expClaimText(long lifetimeSeconds) {
+        return claimOrHeaderText("exp", MOCK_TIME.milliseconds() / 1000.0 + lifetimeSeconds);
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerValidationUtilsTest.java b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerValidationUtilsTest.java
new file mode 100644
index 0000000000000..7477dbc89d9df
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/security/oauthbearer/internals/unsecured/OAuthBearerValidationUtilsTest.java
@@ -0,0 +1,247 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.oauthbearer.internals.unsecured;
+
+import static org.junit.Assert.assertTrue;
+
+import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
+import java.util.Base64;
+import java.util.Collections;
+import java.util.List;
+
+import org.apache.kafka.common.utils.Time;
+import org.junit.Test;
+
+public class OAuthBearerValidationUtilsTest {
+    private static final String QUOTE = "\"";
+    private static final String HEADER_COMPACT_SERIALIZATION = Base64.getUrlEncoder().withoutPadding()
+            .encodeToString("{\"alg\":\"none\"}".getBytes(StandardCharsets.UTF_8)) + ".";
+    private static final Time TIME = Time.SYSTEM;
+
+    @Test
+    public void validateClaimForExistenceAndType() throws OAuthBearerIllegalTokenException {
+        String claimName = "foo";
+        for (Boolean exists : new Boolean[] {null, Boolean.TRUE, Boolean.FALSE}) {
+            boolean useErrorValue = exists == null;
+            for (Boolean required : new boolean[] {true, false}) {
+                StringBuilder sb = new StringBuilder("{");
+                appendJsonText(sb, "exp", 100);
+                appendCommaJsonText(sb, "sub", "principalName");
+                if (useErrorValue)
+                    appendCommaJsonText(sb, claimName, 1);
+                else if (exists != null && exists.booleanValue())
+                    appendCommaJsonText(sb, claimName, claimName);
+                sb.append("}");
+                String compactSerialization = HEADER_COMPACT_SERIALIZATION + Base64.getUrlEncoder().withoutPadding()
+                        .encodeToString(sb.toString().getBytes(StandardCharsets.UTF_8)) + ".";
+                OAuthBearerUnsecuredJws testJwt = new OAuthBearerUnsecuredJws(compactSerialization, "sub", "scope");
+                OAuthBearerValidationResult result = OAuthBearerValidationUtils
+                        .validateClaimForExistenceAndType(testJwt, required, claimName, String.class);
+                if (useErrorValue || required && !exists.booleanValue())
+                    assertTrue(isFailureWithMessageAndNoFailureScope(result));
+                else
+                    assertTrue(isSuccess(result));
+            }
+        }
+    }
+
+    @Test
+    public void validateIssuedAt() {
+        long nowMs = TIME.milliseconds();
+        double nowClaimValue = ((double) nowMs) / 1000;
+        for (boolean exists : new boolean[] {true, false}) {
+            StringBuilder sb = new StringBuilder("{");
+            appendJsonText(sb, "exp", nowClaimValue);
+            appendCommaJsonText(sb, "sub", "principalName");
+            if (exists)
+                appendCommaJsonText(sb, "iat", nowClaimValue);
+            sb.append("}");
+            String compactSerialization = HEADER_COMPACT_SERIALIZATION + Base64.getUrlEncoder().withoutPadding()
+                    .encodeToString(sb.toString().getBytes(StandardCharsets.UTF_8)) + ".";
+            OAuthBearerUnsecuredJws testJwt = new OAuthBearerUnsecuredJws(compactSerialization, "sub", "scope");
+            for (boolean required : new boolean[] {true, false}) {
+                for (int allowableClockSkewMs : new int[] {0, 5, 10, 20}) {
+                    for (long whenCheckOffsetMs : new long[] {-10, 0, 10}) {
+                        long whenCheckMs = nowMs + whenCheckOffsetMs;
+                        OAuthBearerValidationResult result = OAuthBearerValidationUtils.validateIssuedAt(testJwt,
+                                required, whenCheckMs, allowableClockSkewMs);
+                        if (required && !exists)
+                            assertTrue("useErrorValue || required && !exists",
+                                    isFailureWithMessageAndNoFailureScope(result));
+                        else if (!required && !exists)
+                            assertTrue("!required && !exists", isSuccess(result));
+                        else if (nowClaimValue * 1000 > whenCheckMs + allowableClockSkewMs) // issued in future
+                            assertTrue(assertionFailureMessage(nowClaimValue, allowableClockSkewMs, whenCheckMs),
+                                    isFailureWithMessageAndNoFailureScope(result));
+                        else
+                            assertTrue(assertionFailureMessage(nowClaimValue, allowableClockSkewMs, whenCheckMs),
+                                    isSuccess(result));
+                    }
+                }
+            }
+        }
+    }
+
+    @Test
+    public void validateExpirationTime() {
+        long nowMs = TIME.milliseconds();
+        double nowClaimValue = ((double) nowMs) / 1000;
+        StringBuilder sb = new StringBuilder("{");
+        appendJsonText(sb, "exp", nowClaimValue);
+        appendCommaJsonText(sb, "sub", "principalName");
+        sb.append("}");
+        String compactSerialization = HEADER_COMPACT_SERIALIZATION
+                + Base64.getUrlEncoder().withoutPadding().encodeToString(sb.toString().getBytes(StandardCharsets.UTF_8))
+                + ".";
+        OAuthBearerUnsecuredJws testJwt = new OAuthBearerUnsecuredJws(compactSerialization, "sub", "scope");
+        for (int allowableClockSkewMs : new int[] {0, 5, 10, 20}) {
+            for (long whenCheckOffsetMs : new long[] {-10, 0, 10}) {
+                long whenCheckMs = nowMs + whenCheckOffsetMs;
+                OAuthBearerValidationResult result = OAuthBearerValidationUtils.validateExpirationTime(testJwt,
+                        whenCheckMs, allowableClockSkewMs);
+                if (whenCheckMs - allowableClockSkewMs >= nowClaimValue * 1000) // expired
+                    assertTrue(assertionFailureMessage(nowClaimValue, allowableClockSkewMs, whenCheckMs),
+                            isFailureWithMessageAndNoFailureScope(result));
+                else
+                    assertTrue(assertionFailureMessage(nowClaimValue, allowableClockSkewMs, whenCheckMs),
+                            isSuccess(result));
+            }
+        }
+    }
+
+    @Test
+    public void validateExpirationTimeAndIssuedAtConsistency() throws OAuthBearerIllegalTokenException {
+        long nowMs = TIME.milliseconds();
+        double nowClaimValue = ((double) nowMs) / 1000;
+        for (boolean issuedAtExists : new boolean[] {true, false}) {
+            if (!issuedAtExists) {
+                StringBuilder sb = new StringBuilder("{");
+                appendJsonText(sb, "exp", nowClaimValue);
+                appendCommaJsonText(sb, "sub", "principalName");
+                sb.append("}");
+                String compactSerialization = HEADER_COMPACT_SERIALIZATION + Base64.getUrlEncoder().withoutPadding()
+                        .encodeToString(sb.toString().getBytes(StandardCharsets.UTF_8)) + ".";
+                OAuthBearerUnsecuredJws testJwt = new OAuthBearerUnsecuredJws(compactSerialization, "sub", "scope");
+                assertTrue(isSuccess(OAuthBearerValidationUtils.validateTimeConsistency(testJwt)));
+            } else
+                for (int expirationTimeOffset = -1; expirationTimeOffset <= 1; ++expirationTimeOffset) {
+                    StringBuilder sb = new StringBuilder("{");
+                    appendJsonText(sb, "iat", nowClaimValue);
+                    appendCommaJsonText(sb, "exp", nowClaimValue + expirationTimeOffset);
+                    appendCommaJsonText(sb, "sub", "principalName");
+                    sb.append("}");
+                    String compactSerialization = HEADER_COMPACT_SERIALIZATION + Base64.getUrlEncoder().withoutPadding()
+                            .encodeToString(sb.toString().getBytes(StandardCharsets.UTF_8)) + ".";
+                    OAuthBearerUnsecuredJws testJwt = new OAuthBearerUnsecuredJws(compactSerialization, "sub", "scope");
+                    OAuthBearerValidationResult result = OAuthBearerValidationUtils.validateTimeConsistency(testJwt);
+                    if (expirationTimeOffset <= 0)
+                        assertTrue(isFailureWithMessageAndNoFailureScope(result));
+                    else
+                        assertTrue(isSuccess(result));
+                }
+        }
+    }
+
+    @SuppressWarnings({"unchecked", "rawtypes"})
+    @Test
+    public void validateScope() {
+        long nowMs = TIME.milliseconds();
+        double nowClaimValue = ((double) nowMs) / 1000;
+        final List noScope = Collections.emptyList();
+        final List scope1 = Arrays.asList("scope1");
+        final List scope1And2 = Arrays.asList("scope1", "scope2");
+        for (boolean actualScopeExists : new boolean[] {true, false}) {
+            List scopes = !actualScopeExists ? Arrays.asList((List) null)
+                    : Arrays.asList(noScope, scope1, scope1And2);
+            for (List actualScope : scopes) {
+                for (boolean requiredScopeExists : new boolean[] {true, false}) {
+                    List requiredScopes = !requiredScopeExists ? Arrays.asList((List) null)
+                            : Arrays.asList(noScope, scope1, scope1And2);
+                    for (List requiredScope : requiredScopes) {
+                        StringBuilder sb = new StringBuilder("{");
+                        appendJsonText(sb, "exp", nowClaimValue);
+                        appendCommaJsonText(sb, "sub", "principalName");
+                        if (actualScope != null)
+                            sb.append(',').append(scopeJson(actualScope));
+                        sb.append("}");
+                        String compactSerialization = HEADER_COMPACT_SERIALIZATION + Base64.getUrlEncoder()
+                                .withoutPadding().encodeToString(sb.toString().getBytes(StandardCharsets.UTF_8)) + ".";
+                        OAuthBearerUnsecuredJws testJwt = new OAuthBearerUnsecuredJws(compactSerialization, "sub",
+                                "scope");
+                        OAuthBearerValidationResult result = OAuthBearerValidationUtils.validateScope(testJwt,
+                                requiredScope);
+                        if (!requiredScopeExists || requiredScope.isEmpty())
+                            assertTrue(isSuccess(result));
+                        else if (!actualScopeExists || actualScope.size() < requiredScope.size())
+                            assertTrue(isFailureWithMessageAndFailureScope(result));
+                        else
+                            assertTrue(isSuccess(result));
+                    }
+                }
+            }
+        }
+    }
+
+    private static String assertionFailureMessage(double claimValue, int allowableClockSkewMs, long whenCheckMs) {
+        return String.format("time=%f seconds, whenCheck = %d ms, allowableClockSkew=%d ms", claimValue, whenCheckMs,
+                allowableClockSkewMs);
+    }
+
+    private static boolean isSuccess(OAuthBearerValidationResult result) {
+        return result.success();
+    }
+
+    private static boolean isFailureWithMessageAndNoFailureScope(OAuthBearerValidationResult result) {
+        return !result.success() && !result.failureDescription().isEmpty() && result.failureScope() == null
+                && result.failureOpenIdConfig() == null;
+    }
+
+    private static boolean isFailureWithMessageAndFailureScope(OAuthBearerValidationResult result) {
+        return !result.success() && !result.failureDescription().isEmpty() && !result.failureScope().isEmpty()
+                && result.failureOpenIdConfig() == null;
+    }
+
+    private static void appendCommaJsonText(StringBuilder sb, String claimName, Number claimValue) {
+        sb.append(',').append(QUOTE).append(escape(claimName)).append(QUOTE).append(":").append(claimValue);
+    }
+
+    private static void appendCommaJsonText(StringBuilder sb, String claimName, String claimValue) {
+        sb.append(',').append(QUOTE).append(escape(claimName)).append(QUOTE).append(":").append(QUOTE)
+                .append(escape(claimValue)).append(QUOTE);
+    }
+
+    private static void appendJsonText(StringBuilder sb, String claimName, Number claimValue) {
+        sb.append(QUOTE).append(escape(claimName)).append(QUOTE).append(":").append(claimValue);
+    }
+
+    private static String escape(String jsonStringValue) {
+        return jsonStringValue.replace("\"", "\\\"").replace("\\", "\\\\");
+    }
+
+    private static String scopeJson(List scope) {
+        StringBuilder scopeJsonBuilder = new StringBuilder("\"scope\":[");
+        int initialLength = scopeJsonBuilder.length();
+        for (String scopeValue : scope) {
+            if (scopeJsonBuilder.length() > initialLength)
+                scopeJsonBuilder.append(',');
+            scopeJsonBuilder.append('"').append(scopeValue).append('"');
+        }
+        scopeJsonBuilder.append(']');
+        return scopeJsonBuilder.toString();
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/security/plain/PlainSaslServerTest.java b/clients/src/test/java/org/apache/kafka/common/security/plain/PlainSaslServerTest.java
deleted file mode 100644
index 4196db6ac0f53..0000000000000
--- a/clients/src/test/java/org/apache/kafka/common/security/plain/PlainSaslServerTest.java
+++ /dev/null
@@ -1,74 +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.security.plain;
-
-import org.junit.Before;
-import org.junit.Test;
-
-import java.nio.charset.StandardCharsets;
-import java.util.HashMap;
-import java.util.Map;
-
-import static org.junit.Assert.assertEquals;
-
-import org.apache.kafka.common.errors.SaslAuthenticationException;
-import org.apache.kafka.common.security.JaasContext;
-import org.apache.kafka.common.security.authenticator.TestJaasConfig;
-
-public class PlainSaslServerTest {
-
-    private static final String USER_A = "userA";
-    private static final String PASSWORD_A = "passwordA";
-    private static final String USER_B = "userB";
-    private static final String PASSWORD_B = "passwordB";
-
-    private PlainSaslServer saslServer;
-
-    @Before
-    public void setUp() throws Exception {
-        TestJaasConfig jaasConfig = new TestJaasConfig();
-        Map options = new HashMap<>();
-        options.put("user_" + USER_A, PASSWORD_A);
-        options.put("user_" + USER_B, PASSWORD_B);
-        jaasConfig.addEntry("jaasContext", PlainLoginModule.class.getName(), options);
-        JaasContext jaasContext = new JaasContext("jaasContext", JaasContext.Type.SERVER, jaasConfig);
-        saslServer = new PlainSaslServer(jaasContext);
-    }
-
-    @Test
-    public void noAuthorizationIdSpecified() throws Exception {
-        byte[] nextChallenge = saslServer.evaluateResponse(saslMessage("", USER_A, PASSWORD_A));
-        assertEquals(0, nextChallenge.length);
-    }
-
-    @Test
-    public void authorizatonIdEqualsAuthenticationId() throws Exception {
-        byte[] nextChallenge = saslServer.evaluateResponse(saslMessage(USER_A, USER_A, PASSWORD_A));
-        assertEquals(0, nextChallenge.length);
-    }
-
-    @Test(expected = SaslAuthenticationException.class)
-    public void authorizatonIdNotEqualsAuthenticationId() throws Exception {
-        saslServer.evaluateResponse(saslMessage(USER_B, USER_A, PASSWORD_A));
-    }
-
-    private byte[] saslMessage(String authorizationId, String userName, String password) {
-        String nul = "\u0000";
-        String message = String.format("%s%s%s%s%s", authorizationId, nul, userName, nul, password);
-        return message.getBytes(StandardCharsets.UTF_8);
-    }
-}
diff --git a/clients/src/test/java/org/apache/kafka/common/security/plain/internals/PlainSaslServerTest.java b/clients/src/test/java/org/apache/kafka/common/security/plain/internals/PlainSaslServerTest.java
new file mode 100644
index 0000000000000..5fb2042bbbd36
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/security/plain/internals/PlainSaslServerTest.java
@@ -0,0 +1,117 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.plain.internals;
+
+import org.apache.kafka.common.security.plain.PlainLoginModule;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.nio.charset.StandardCharsets;
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+
+import org.apache.kafka.common.errors.SaslAuthenticationException;
+import org.apache.kafka.common.security.JaasContext;
+import org.apache.kafka.common.security.authenticator.TestJaasConfig;
+
+public class PlainSaslServerTest {
+
+    private static final String USER_A = "userA";
+    private static final String PASSWORD_A = "passwordA";
+    private static final String USER_B = "userB";
+    private static final String PASSWORD_B = "passwordB";
+
+    private PlainSaslServer saslServer;
+
+    @Before
+    public void setUp() {
+        TestJaasConfig jaasConfig = new TestJaasConfig();
+        Map options = new HashMap<>();
+        options.put("user_" + USER_A, PASSWORD_A);
+        options.put("user_" + USER_B, PASSWORD_B);
+        jaasConfig.addEntry("jaasContext", PlainLoginModule.class.getName(), options);
+        JaasContext jaasContext = new JaasContext("jaasContext", JaasContext.Type.SERVER, jaasConfig, null);
+        PlainServerCallbackHandler callbackHandler = new PlainServerCallbackHandler();
+        callbackHandler.configure(null, "PLAIN", jaasContext.configurationEntries());
+        saslServer = new PlainSaslServer(callbackHandler);
+    }
+
+    @Test
+    public void noAuthorizationIdSpecified() throws Exception {
+        byte[] nextChallenge = saslServer.evaluateResponse(saslMessage("", USER_A, PASSWORD_A));
+        assertEquals(0, nextChallenge.length);
+    }
+
+    @Test
+    public void authorizatonIdEqualsAuthenticationId() throws Exception {
+        byte[] nextChallenge = saslServer.evaluateResponse(saslMessage(USER_A, USER_A, PASSWORD_A));
+        assertEquals(0, nextChallenge.length);
+    }
+
+    @Test(expected = SaslAuthenticationException.class)
+    public void authorizatonIdNotEqualsAuthenticationId() throws Exception {
+        saslServer.evaluateResponse(saslMessage(USER_B, USER_A, PASSWORD_A));
+    }
+
+    @Test
+    public void emptyTokens() {
+        Exception e = assertThrows(SaslAuthenticationException.class, () ->
+            saslServer.evaluateResponse(saslMessage("", "", "")));
+        assertEquals("Authentication failed: username not specified", e.getMessage());
+
+        e = assertThrows(SaslAuthenticationException.class, () ->
+            saslServer.evaluateResponse(saslMessage("", "", "p")));
+        assertEquals("Authentication failed: username not specified", e.getMessage());
+
+        e = assertThrows(SaslAuthenticationException.class, () ->
+            saslServer.evaluateResponse(saslMessage("", "u", "")));
+        assertEquals("Authentication failed: password not specified", e.getMessage());
+
+        e = assertThrows(SaslAuthenticationException.class, () ->
+            saslServer.evaluateResponse(saslMessage("a", "", "")));
+        assertEquals("Authentication failed: username not specified", e.getMessage());
+
+        e = assertThrows(SaslAuthenticationException.class, () ->
+            saslServer.evaluateResponse(saslMessage("a", "", "p")));
+        assertEquals("Authentication failed: username not specified", e.getMessage());
+
+        e = assertThrows(SaslAuthenticationException.class, () ->
+            saslServer.evaluateResponse(saslMessage("a", "u", "")));
+        assertEquals("Authentication failed: password not specified", e.getMessage());
+
+        String nul = "\u0000";
+
+        e = assertThrows(SaslAuthenticationException.class, () ->
+            saslServer.evaluateResponse(
+                String.format("%s%s%s%s%s%s", "a", nul, "u", nul, "p", nul).getBytes(StandardCharsets.UTF_8)));
+        assertEquals("Invalid SASL/PLAIN response: expected 3 tokens, got 4", e.getMessage());
+
+        e = assertThrows(SaslAuthenticationException.class, () ->
+            saslServer.evaluateResponse(
+                String.format("%s%s%s", "", nul, "u").getBytes(StandardCharsets.UTF_8)));
+        assertEquals("Invalid SASL/PLAIN response: expected 3 tokens, got 2", e.getMessage());
+    }
+
+    private byte[] saslMessage(String authorizationId, String userName, String password) {
+        String nul = "\u0000";
+        String message = String.format("%s%s%s%s%s", authorizationId, nul, userName, nul, password);
+        return message.getBytes(StandardCharsets.UTF_8);
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/security/scram/ScramCredentialUtilsTest.java b/clients/src/test/java/org/apache/kafka/common/security/scram/internals/ScramCredentialUtilsTest.java
similarity index 97%
rename from clients/src/test/java/org/apache/kafka/common/security/scram/ScramCredentialUtilsTest.java
rename to clients/src/test/java/org/apache/kafka/common/security/scram/internals/ScramCredentialUtilsTest.java
index e9dd285f54e62..bc9ef069c9837 100644
--- a/clients/src/test/java/org/apache/kafka/common/security/scram/ScramCredentialUtilsTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/security/scram/internals/ScramCredentialUtilsTest.java
@@ -14,23 +14,23 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.kafka.common.security.scram;
+package org.apache.kafka.common.security.scram.internals;
 
-import org.junit.Test;
+import java.security.NoSuchAlgorithmException;
+import java.util.Arrays;
 
+import org.apache.kafka.common.security.authenticator.CredentialCache;
+import org.apache.kafka.common.security.scram.ScramCredential;
+
+import org.junit.Before;
+import org.junit.Test;
 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 java.security.NoSuchAlgorithmException;
-import java.util.Arrays;
-
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertTrue;
 
-import org.apache.kafka.common.security.authenticator.CredentialCache;
-import org.junit.Before;
 
 public class ScramCredentialUtilsTest {
 
diff --git a/clients/src/test/java/org/apache/kafka/common/security/scram/ScramFormatterTest.java b/clients/src/test/java/org/apache/kafka/common/security/scram/internals/ScramFormatterTest.java
similarity index 80%
rename from clients/src/test/java/org/apache/kafka/common/security/scram/ScramFormatterTest.java
rename to clients/src/test/java/org/apache/kafka/common/security/scram/internals/ScramFormatterTest.java
index a86e0ddb80079..251f4ca25638c 100644
--- a/clients/src/test/java/org/apache/kafka/common/security/scram/ScramFormatterTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/security/scram/internals/ScramFormatterTest.java
@@ -14,19 +14,20 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.kafka.common.security.scram;
+package org.apache.kafka.common.security.scram.internals;
+
+import org.apache.kafka.common.security.scram.internals.ScramMessages.ClientFinalMessage;
+import org.apache.kafka.common.security.scram.internals.ScramMessages.ClientFirstMessage;
+import org.apache.kafka.common.security.scram.internals.ScramMessages.ServerFinalMessage;
+import org.apache.kafka.common.security.scram.internals.ScramMessages.ServerFirstMessage;
 
-import org.apache.kafka.common.utils.Base64;
 import org.junit.Test;
 
+import java.util.Base64;
+
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 
-import org.apache.kafka.common.security.scram.ScramMessages.ClientFinalMessage;
-import org.apache.kafka.common.security.scram.ScramMessages.ClientFirstMessage;
-import org.apache.kafka.common.security.scram.ScramMessages.ServerFinalMessage;
-import org.apache.kafka.common.security.scram.ScramMessages.ServerFirstMessage;
-
 public class ScramFormatterTest {
 
     /**
@@ -42,10 +43,10 @@ public void rfc7677Example() throws Exception {
         String s1 = "r=rOprNGfwEbeRWgbNEkqO%hvYDpWUa2RaTCAfuxFIlj)hNlF$k0,s=W22ZaJ0SNY7soEsUEjb6gQ==,i=4096";
         String c2 = "c=biws,r=rOprNGfwEbeRWgbNEkqO%hvYDpWUa2RaTCAfuxFIlj)hNlF$k0,p=dHzbZapWIk4jUhN+Ute9ytag9zjfMHgsqmmiz7AndVQ=";
         String s2 = "v=6rriTRBi23WpRR/wtup+mMhUZUn/dB5nLTJRsjl95G4=";
-        ClientFirstMessage clientFirst = new ClientFirstMessage(formatter.toBytes(c1));
-        ServerFirstMessage serverFirst = new ServerFirstMessage(formatter.toBytes(s1));
-        ClientFinalMessage clientFinal = new ClientFinalMessage(formatter.toBytes(c2));
-        ServerFinalMessage serverFinal = new ServerFinalMessage(formatter.toBytes(s2));
+        ClientFirstMessage clientFirst = new ClientFirstMessage(ScramFormatter.toBytes(c1));
+        ServerFirstMessage serverFirst = new ServerFirstMessage(ScramFormatter.toBytes(s1));
+        ClientFinalMessage clientFinal = new ClientFinalMessage(ScramFormatter.toBytes(c2));
+        ServerFinalMessage serverFinal = new ServerFinalMessage(ScramFormatter.toBytes(s2));
 
         String username = clientFirst.saslName();
         assertEquals("user", username);
@@ -54,13 +55,13 @@ public void rfc7677Example() throws Exception {
         String serverNonce = serverFirst.nonce().substring(clientNonce.length());
         assertEquals("%hvYDpWUa2RaTCAfuxFIlj)hNlF$k0", serverNonce);
         byte[] salt = serverFirst.salt();
-        assertArrayEquals(Base64.decoder().decode("W22ZaJ0SNY7soEsUEjb6gQ=="), salt);
+        assertArrayEquals(Base64.getDecoder().decode("W22ZaJ0SNY7soEsUEjb6gQ=="), salt);
         int iterations = serverFirst.iterations();
         assertEquals(4096, iterations);
         byte[] channelBinding = clientFinal.channelBinding();
-        assertArrayEquals(Base64.decoder().decode("biws"), channelBinding);
+        assertArrayEquals(Base64.getDecoder().decode("biws"), channelBinding);
         byte[] serverSignature = serverFinal.serverSignature();
-        assertArrayEquals(Base64.decoder().decode("6rriTRBi23WpRR/wtup+mMhUZUn/dB5nLTJRsjl95G4="), serverSignature);
+        assertArrayEquals(Base64.getDecoder().decode("6rriTRBi23WpRR/wtup+mMhUZUn/dB5nLTJRsjl95G4="), serverSignature);
 
         byte[] saltedPassword = formatter.saltedPassword(password, salt, iterations);
         byte[] serverKey = formatter.serverKey(saltedPassword);
@@ -81,12 +82,12 @@ public void saslName() throws Exception {
         String[] usernames = {"user1", "123", "1,2", "user=A", "user==B", "user,1", "user 1", ",", "=", ",=", "=="};
         ScramFormatter formatter = new ScramFormatter(ScramMechanism.SCRAM_SHA_256);
         for (String username : usernames) {
-            String saslName = formatter.saslName(username);
+            String saslName = ScramFormatter.saslName(username);
             // There should be no commas in saslName (comma is used as field separator in SASL messages)
             assertEquals(-1, saslName.indexOf(','));
             // There should be no "=" in the saslName apart from those used in encoding (comma is =2C and equals is =3D)
             assertEquals(-1, saslName.replace("=2C", "").replace("=3D", "").indexOf('='));
-            assertEquals(username, formatter.username(saslName));
+            assertEquals(username, ScramFormatter.username(saslName));
         }
     }
 }
diff --git a/clients/src/test/java/org/apache/kafka/common/security/scram/ScramMessagesTest.java b/clients/src/test/java/org/apache/kafka/common/security/scram/internals/ScramMessagesTest.java
similarity index 88%
rename from clients/src/test/java/org/apache/kafka/common/security/scram/ScramMessagesTest.java
rename to clients/src/test/java/org/apache/kafka/common/security/scram/internals/ScramMessagesTest.java
index 89e626095e8bf..f1e9f33814fac 100644
--- a/clients/src/test/java/org/apache/kafka/common/security/scram/ScramMessagesTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/security/scram/internals/ScramMessagesTest.java
@@ -14,27 +14,28 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.kafka.common.security.scram;
-
-import org.apache.kafka.common.utils.Base64;
-import org.junit.Before;
-import org.junit.Test;
+package org.apache.kafka.common.security.scram.internals;
 
 import java.nio.charset.StandardCharsets;
+import java.util.Base64;
+import java.util.Collections;
 
 import javax.security.sasl.SaslException;
 
+import org.apache.kafka.common.security.scram.internals.ScramMessages.AbstractScramMessage;
+import org.apache.kafka.common.security.scram.internals.ScramMessages.ClientFinalMessage;
+import org.apache.kafka.common.security.scram.internals.ScramMessages.ClientFirstMessage;
+import org.apache.kafka.common.security.scram.internals.ScramMessages.ServerFinalMessage;
+import org.apache.kafka.common.security.scram.internals.ScramMessages.ServerFirstMessage;
+
+import org.junit.Before;
+import org.junit.Test;
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
-import org.apache.kafka.common.security.scram.ScramMessages.AbstractScramMessage;
-import org.apache.kafka.common.security.scram.ScramMessages.ClientFinalMessage;
-import org.apache.kafka.common.security.scram.ScramMessages.ClientFirstMessage;
-import org.apache.kafka.common.security.scram.ScramMessages.ServerFinalMessage;
-import org.apache.kafka.common.security.scram.ScramMessages.ServerFirstMessage;
-
 public class ScramMessagesTest {
 
     private static final String[] VALID_EXTENSIONS = {
@@ -70,7 +71,7 @@ public void setUp() throws Exception {
     @Test
     public void validClientFirstMessage() throws SaslException {
         String nonce = formatter.secureRandomString();
-        ClientFirstMessage m = new ClientFirstMessage("someuser", nonce);
+        ClientFirstMessage m = new ClientFirstMessage("someuser", nonce, Collections.emptyMap());
         checkClientFirstMessage(m, "someuser", nonce, "");
 
         // Default format used by Kafka client: only user and nonce are specified
@@ -84,13 +85,13 @@ public void validClientFirstMessage() throws SaslException {
         str = String.format("n,,n=test=2Cuser,r=%s", nonce);
         m = createScramMessage(ClientFirstMessage.class, str);
         checkClientFirstMessage(m, "test=2Cuser", nonce, "");
-        assertEquals("test,user", formatter.username(m.saslName()));
+        assertEquals("test,user", ScramFormatter.username(m.saslName()));
 
         // Username containing equals, encoded as =3D
         str = String.format("n,,n=test=3Duser,r=%s", nonce);
         m = createScramMessage(ClientFirstMessage.class, str);
         checkClientFirstMessage(m, "test=3Duser", nonce, "");
-        assertEquals("test=user", formatter.username(m.saslName()));
+        assertEquals("test=user", ScramFormatter.username(m.saslName()));
 
         // Optional authorization id specified
         str = String.format("n,a=testauthzid,n=testuser,r=%s", nonce);
@@ -107,10 +108,15 @@ public void validClientFirstMessage() throws SaslException {
             str = String.format("n,,n=testuser,r=%s,%s", nonce, extension);
             checkClientFirstMessage(createScramMessage(ClientFirstMessage.class, str), "testuser", nonce, "");
         }
+
+        //optional tokenauth specified as extensions
+        str = String.format("n,,n=testuser,r=%s,%s", nonce, "tokenauth=true");
+        m = createScramMessage(ClientFirstMessage.class, str);
+        assertTrue("Token authentication not set from extensions", m.extensions().tokenAuthenticated());
     }
 
     @Test
-    public void invalidClientFirstMessage() throws SaslException {
+    public void invalidClientFirstMessage() {
         String nonce = formatter.secureRandomString();
         // Invalid entry in gs2-header
         String invalid = String.format("n,x=something,n=testuser,r=%s", nonce);
@@ -160,7 +166,7 @@ public void validServerFirstMessage() throws SaslException {
     }
 
     @Test
-    public void invalidServerFirstMessage() throws SaslException {
+    public void invalidServerFirstMessage() {
         String nonce = formatter.secureRandomString();
         String salt = randomBytesAsString();
 
@@ -215,7 +221,7 @@ public void validClientFinalMessage() throws SaslException {
     }
 
     @Test
-    public void invalidClientFinalMessage() throws SaslException {
+    public void invalidClientFinalMessage() {
         String nonce = formatter.secureRandomString();
         String channelBinding = randomBytesAsString();
         String proof = randomBytesAsString();
@@ -266,7 +272,7 @@ public void validServerFinalMessage() throws SaslException {
     }
 
     @Test
-    public void invalidServerFinalMessage() throws SaslException {
+    public void invalidServerFinalMessage() {
         String serverSignature = randomBytesAsString();
 
         // Invalid error
@@ -288,12 +294,12 @@ public void invalidServerFinalMessage() throws SaslException {
     }
 
     private String randomBytesAsString() {
-        return Base64.encoder().encodeToString(formatter.secureRandomBytes());
+        return Base64.getEncoder().encodeToString(formatter.secureRandomBytes());
     }
 
     private byte[] toBytes(String base64Str) {
-        return Base64.decoder().decode(base64Str);
-    };
+        return Base64.getDecoder().decode(base64Str);
+    }
 
     private void checkClientFirstMessage(ClientFirstMessage message, String saslName, String nonce, String authzid) {
         assertEquals(saslName, message.saslName());
@@ -303,14 +309,14 @@ private void checkClientFirstMessage(ClientFirstMessage message, String saslName
 
     private void checkServerFirstMessage(ServerFirstMessage message, String nonce, String salt, int iterations) {
         assertEquals(nonce, message.nonce());
-        assertArrayEquals(Base64.decoder().decode(salt), message.salt());
+        assertArrayEquals(Base64.getDecoder().decode(salt), message.salt());
         assertEquals(iterations, message.iterations());
     }
 
     private void checkClientFinalMessage(ClientFinalMessage message, String channelBinding, String nonce, String proof) {
-        assertArrayEquals(Base64.decoder().decode(channelBinding), message.channelBinding());
+        assertArrayEquals(Base64.getDecoder().decode(channelBinding), message.channelBinding());
         assertEquals(nonce, message.nonce());
-        assertArrayEquals(Base64.decoder().decode(proof), message.proof());
+        assertArrayEquals(Base64.getDecoder().decode(proof), message.proof());
     }
 
     private void checkServerFinalMessage(ServerFinalMessage message, String error, String serverSignature) {
@@ -318,7 +324,7 @@ private void checkServerFinalMessage(ServerFinalMessage message, String error, S
         if (serverSignature == null)
             assertNull("Unexpected server signature", message.serverSignature());
         else
-            assertArrayEquals(Base64.decoder().decode(serverSignature), message.serverSignature());
+            assertArrayEquals(Base64.getDecoder().decode(serverSignature), message.serverSignature());
     }
 
     @SuppressWarnings("unchecked")
diff --git a/clients/src/test/java/org/apache/kafka/common/security/scram/ScramSaslServerTest.java b/clients/src/test/java/org/apache/kafka/common/security/scram/internals/ScramSaslServerTest.java
similarity index 91%
rename from clients/src/test/java/org/apache/kafka/common/security/scram/ScramSaslServerTest.java
rename to clients/src/test/java/org/apache/kafka/common/security/scram/internals/ScramSaslServerTest.java
index fd7f98838b784..270d972f99cf7 100644
--- a/clients/src/test/java/org/apache/kafka/common/security/scram/ScramSaslServerTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/security/scram/internals/ScramSaslServerTest.java
@@ -14,18 +14,20 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.kafka.common.security.scram;
+package org.apache.kafka.common.security.scram.internals;
 
-import org.junit.Before;
-import org.junit.Test;
 
 import java.nio.charset.StandardCharsets;
 import java.util.HashMap;
 
-import static org.junit.Assert.assertTrue;
-
 import org.apache.kafka.common.errors.SaslAuthenticationException;
 import org.apache.kafka.common.security.authenticator.CredentialCache;
+import org.apache.kafka.common.security.scram.ScramCredential;
+import org.apache.kafka.common.security.token.delegation.internals.DelegationTokenCache;
+
+import org.junit.Before;
+import org.junit.Test;
+import static org.junit.Assert.assertTrue;
 
 public class ScramSaslServerTest {
 
@@ -43,7 +45,7 @@ public void setUp() throws Exception {
         CredentialCache.Cache credentialCache = new CredentialCache().createCache(mechanism.mechanismName(), ScramCredential.class);
         credentialCache.put(USER_A, formatter.generateCredential("passwordA", 4096));
         credentialCache.put(USER_B, formatter.generateCredential("passwordB", 4096));
-        ScramServerCallbackHandler callbackHandler = new ScramServerCallbackHandler(credentialCache);
+        ScramServerCallbackHandler callbackHandler = new ScramServerCallbackHandler(credentialCache, new DelegationTokenCache(ScramMechanism.mechanismNames()));
         saslServer = new ScramSaslServer(mechanism, new HashMap(), callbackHandler);
     }
 
diff --git a/clients/src/test/java/org/apache/kafka/common/security/ssl/DefaultSslEngineFactoryTest.java b/clients/src/test/java/org/apache/kafka/common/security/ssl/DefaultSslEngineFactoryTest.java
new file mode 100644
index 0000000000000..c8c0566015de2
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/security/ssl/DefaultSslEngineFactoryTest.java
@@ -0,0 +1,324 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.config.SslConfigs;
+import org.apache.kafka.common.config.types.Password;
+import org.apache.kafka.common.errors.InvalidConfigurationException;
+import org.apache.kafka.test.TestUtils;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.security.KeyStore;
+import java.util.Arrays;
+import java.util.Collections;
+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.assertNull;
+import static org.junit.Assert.assertThrows;
+
+public class DefaultSslEngineFactoryTest {
+
+    /*
+     * Key and certificates were extracted using openssl from a key store file created with 100 years validity using:
+     *
+     * openssl pkcs12 -in server.keystore.p12 -nodes -nocerts -out test.key.pem -passin pass:key-password
+     * openssl pkcs12 -in server.keystore.p12 -nodes -nokeys -out test.certchain.pem  -passin pass:key-password
+     * openssl pkcs12 -in server.keystore.p12 -nodes  -out test.keystore.pem -passin pass:key-password
+     * openssl pkcs8 -topk8 -v1 pbeWithSHA1And3-KeyTripleDES-CBC -in test.key.pem -out test.key.encrypted.pem -passout pass:key-password
+     */
+
+    private static final String CA1 = "-----BEGIN CERTIFICATE-----\n"
+            + "MIIC0zCCAbugAwIBAgIEStdXHTANBgkqhkiG9w0BAQsFADASMRAwDgYDVQQDEwdU\n"
+            + "ZXN0Q0ExMCAXDTIwMDkyODA5MDI0MFoYDzIxMjAwOTA0MDkwMjQwWjASMRAwDgYD\n"
+            + "VQQDEwdUZXN0Q0ExMIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAo3Gr\n"
+            + "WJAkjnvgcuIfjArDhNdtAlRTt094WMUXhYDibgGtd+CLcWqA+c4PEoK4oybnKZqU\n"
+            + "6MlDfPgesIK2YiNBuSVWMtZ2doageOBnd80Iwbg8DqGtQpUsvw8X5fOmuza+4inv\n"
+            + "/8IpiTizq8YjSMT4nYDmIjyyRCSNY4atjgMnskutJ0v6i69+ZAA520Y6nn2n4RD5\n"
+            + "8Yc+y7yCkbZXnYS5xBOFEExmtc0Xa7S9nM157xqKws9Z+rTKZYLrryaHI9JNcXgG\n"
+            + "kzQEH9fBePASeWfi9AGRvAyS2GMSIBOsihIDIha/mqQcJOGCEqTMtefIj2FaErO2\n"
+            + "bL9yU7OpW53iIC8y0QIDAQABoy8wLTAMBgNVHRMEBTADAQH/MB0GA1UdDgQWBBRf\n"
+            + "svKcoQ9ZBvjwyUSV2uMFzlkOWDANBgkqhkiG9w0BAQsFAAOCAQEAEE1ZG2MGE248\n"
+            + "glO83ROrHbxmnVWSQHt/JZANR1i362sY1ekL83wlhkriuvGVBlHQYWezIfo/4l9y\n"
+            + "JTHNX3Mrs9eWUkaDXADkHWj3AyLXN3nfeU307x1wA7OvI4YKpwvfb4aYS8RTPz9d\n"
+            + "JtrfR0r8aGTgsXvCe4SgwDBKv7bckctOwD3S7D/b6y3w7X0s7JCU5+8ZjgoYfcLE\n"
+            + "gNqQEaOwdT2LHCvxHmGn/2VGs/yatPQIYYuufe5i8yX7pp4Xbd2eD6LULYkHFs3x\n"
+            + "uJzMRI7BukmIIWuBbAkYI0atxLQIysnVFXdL9pBgvgso2nA3FgP/XeORhkyHVvtL\n"
+            + "REH2YTlftQ==\n"
+            + "-----END CERTIFICATE-----";
+
+    private static final String CA2 = "-----BEGIN CERTIFICATE-----\n"
+            + "MIIC0zCCAbugAwIBAgIEfk9e9DANBgkqhkiG9w0BAQsFADASMRAwDgYDVQQDEwdU\n"
+            + "ZXN0Q0EyMCAXDTIwMDkyODA5MDI0MVoYDzIxMjAwOTA0MDkwMjQxWjASMRAwDgYD\n"
+            + "VQQDEwdUZXN0Q0EyMIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAvCh0\n"
+            + "UO5op9eHfz7mvZ7IySK7AOCTC56QYFJcU+hD6yk1wKg2qot7naI5ozAc8n7c4pMt\n"
+            + "LjI3D0VtC/oHC29R2HNMSWyHcxIXw8z127XeCLRkCqYWuVAl3nBuWfWVPObjKetH\n"
+            + "TWlQANYWAfk1VbS6wfzgp9cMaK7wQ+VoGEo4x3pjlrdlyg4k4O2yubcpWmJ2TjxS\n"
+            + "gg7TfKGizUVAvF9wUG9Q4AlCg4uuww5RN9w6vnzDKGhWJhkQ6pf/m1xB+WueFOeU\n"
+            + "aASGhGqCTqiz3p3M3M4OZzG3KptjQ/yb67x4T5U5RxqoiN4L57E7ZJLREpa6ZZNs\n"
+            + "ps/gQ8dR9Uo/PRyAkQIDAQABoy8wLTAMBgNVHRMEBTADAQH/MB0GA1UdDgQWBBRg\n"
+            + "IAOVH5LeE6nZmdScEE3JO/AhvTANBgkqhkiG9w0BAQsFAAOCAQEAHkk1iybwy/Lf\n"
+            + "iEQMVRy7XfuC008O7jfCUBMgUvE+oO2RadH5MmsXHG3YerdsDM90dui4JqQNZOUh\n"
+            + "kF8dIWPQHE0xDsR9jiUsemZFpVMN7DcvVZ3eFhbvJA8Q50rxcNGA+tn9xT/xdQ6z\n"
+            + "1eRq9IPoYcRexQ7s9mincM4T4lLm8GGcd7ZPHy8kw0Bp3E/enRHWaF5b8KbXezXD\n"
+            + "I3SEYUyRL2K3px4FImT4X9XQm2EX6EONlu4GRcJpD6RPc0zC7c9dwEnSo+0NnewR\n"
+            + "gjgO34CLzShB/kASLS9VQXcUC6bsggAVK2rWQMmy35SOEUufSuvg8kUFoyuTzfhn\n"
+            + "hL+PVwIu7g==\n"
+            + "-----END CERTIFICATE-----";
+
+    private static final String CERTCHAIN = "Bag Attributes\n"
+            + "    friendlyName: server\n"
+            + "    localKeyID: 54 69 6D 65 20 31 36 30 31 32 38 33 37 36 35 34 32 33 \n"
+            + "subject=/CN=TestBroker\n"
+            + "issuer=/CN=TestCA1\n"
+            + "-----BEGIN CERTIFICATE-----\n"
+            + "MIIC/zCCAeegAwIBAgIEatBnEzANBgkqhkiG9w0BAQsFADASMRAwDgYDVQQDEwdU\n"
+            + "ZXN0Q0ExMCAXDTIwMDkyODA5MDI0NFoYDzIxMjAwOTA0MDkwMjQ0WjAVMRMwEQYD\n"
+            + "VQQDEwpUZXN0QnJva2VyMIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEA\n"
+            + "pkw1AS71ej/iOMvzVgVL1dkQOYzI842NcPmx0yFFsue2umL8WVd3085NgWRb3SS1\n"
+            + "4X676t7zxjPGzYi7jwmA8stCrDt0NAPWd/Ko6ErsCs87CUs4u1Cinf+b3o9NF5u0\n"
+            + "UPYBQLF4Ir8T1jQ+tKiqsChGDt6urRAg1Cro5i7r10jN1uofY2tBs+r8mALhJ17c\n"
+            + "T5LKawXeYwNOQ86c5djClbcP0RrfcPyRyj1/Cp1axo28iO0fXFyO2Zf3a4vtt+Ih\n"
+            + "PW+A2tL+t3JTBd8g7Fl3ozzpcotAi7MDcZaYA9GiTP4DOiKUeDt6yMYQQr3VEqGa\n"
+            + "pXp4fKY+t9slqnAmcBZ4kQIDAQABo1gwVjAfBgNVHSMEGDAWgBRfsvKcoQ9ZBvjw\n"
+            + "yUSV2uMFzlkOWDAUBgNVHREEDTALgglsb2NhbGhvc3QwHQYDVR0OBBYEFGWt+27P\n"
+            + "INk/S5X+PRV/jW3WOhtaMA0GCSqGSIb3DQEBCwUAA4IBAQCLHCjFFvqa+0GcG9eq\n"
+            + "v1QWaXDohY5t5CCwD8Z+lT9wcSruTxDPwL7LrR36h++D6xJYfiw4iaRighoA40xP\n"
+            + "W6+0zGK/UtWV4t+ODTDzyAWgls5w+0R5ki6447qGqu5tXlW5DCHkkxWiozMnhNU2\n"
+            + "G3P/Drh7DhmADDBjtVLsu5M1sagF/xwTP/qCLMdChlJNdeqyLnAUa9SYG1eNZS/i\n"
+            + "wrCC8m9RUQb4+OlQuFtr0KhaaCkBXfmhigQAmh44zSyO+oa3qQDEavVFo/Mcui9o\n"
+            + "WBYetcgVbXPNoti+hQEMqmJYBHlLbhxMnkooGn2fa70f453Bdu/Xh6Yphi5NeCHn\n"
+            + "1I+y\n"
+            + "-----END CERTIFICATE-----\n"
+            + "Bag Attributes\n"
+            + "    friendlyName: CN=TestCA1\n"
+            + "subject=/CN=TestCA1\n"
+            + "issuer=/CN=TestCA1\n"
+            + "-----BEGIN CERTIFICATE-----\n"
+            + "MIIC0zCCAbugAwIBAgIEStdXHTANBgkqhkiG9w0BAQsFADASMRAwDgYDVQQDEwdU\n"
+            + "ZXN0Q0ExMCAXDTIwMDkyODA5MDI0MFoYDzIxMjAwOTA0MDkwMjQwWjASMRAwDgYD\n"
+            + "VQQDEwdUZXN0Q0ExMIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAo3Gr\n"
+            + "WJAkjnvgcuIfjArDhNdtAlRTt094WMUXhYDibgGtd+CLcWqA+c4PEoK4oybnKZqU\n"
+            + "6MlDfPgesIK2YiNBuSVWMtZ2doageOBnd80Iwbg8DqGtQpUsvw8X5fOmuza+4inv\n"
+            + "/8IpiTizq8YjSMT4nYDmIjyyRCSNY4atjgMnskutJ0v6i69+ZAA520Y6nn2n4RD5\n"
+            + "8Yc+y7yCkbZXnYS5xBOFEExmtc0Xa7S9nM157xqKws9Z+rTKZYLrryaHI9JNcXgG\n"
+            + "kzQEH9fBePASeWfi9AGRvAyS2GMSIBOsihIDIha/mqQcJOGCEqTMtefIj2FaErO2\n"
+            + "bL9yU7OpW53iIC8y0QIDAQABoy8wLTAMBgNVHRMEBTADAQH/MB0GA1UdDgQWBBRf\n"
+            + "svKcoQ9ZBvjwyUSV2uMFzlkOWDANBgkqhkiG9w0BAQsFAAOCAQEAEE1ZG2MGE248\n"
+            + "glO83ROrHbxmnVWSQHt/JZANR1i362sY1ekL83wlhkriuvGVBlHQYWezIfo/4l9y\n"
+            + "JTHNX3Mrs9eWUkaDXADkHWj3AyLXN3nfeU307x1wA7OvI4YKpwvfb4aYS8RTPz9d\n"
+            + "JtrfR0r8aGTgsXvCe4SgwDBKv7bckctOwD3S7D/b6y3w7X0s7JCU5+8ZjgoYfcLE\n"
+            + "gNqQEaOwdT2LHCvxHmGn/2VGs/yatPQIYYuufe5i8yX7pp4Xbd2eD6LULYkHFs3x\n"
+            + "uJzMRI7BukmIIWuBbAkYI0atxLQIysnVFXdL9pBgvgso2nA3FgP/XeORhkyHVvtL\n"
+            + "REH2YTlftQ==\n"
+            + "-----END CERTIFICATE-----";
+
+    private static final String KEY =  "Bag Attributes\n"
+            + "    friendlyName: server\n"
+            + "    localKeyID: 54 69 6D 65 20 31 36 30 31 32 38 33 37 36 35 34 32 33\n"
+            + "Key Attributes: \n"
+            + "-----BEGIN PRIVATE KEY-----\n"
+            + "MIIEvAIBADANBgkqhkiG9w0BAQEFAASCBKYwggSiAgEAAoIBAQCmTDUBLvV6P+I4\n"
+            + "y/NWBUvV2RA5jMjzjY1w+bHTIUWy57a6YvxZV3fTzk2BZFvdJLXhfrvq3vPGM8bN\n"
+            + "iLuPCYDyy0KsO3Q0A9Z38qjoSuwKzzsJSzi7UKKd/5vej00Xm7RQ9gFAsXgivxPW\n"
+            + "ND60qKqwKEYO3q6tECDUKujmLuvXSM3W6h9ja0Gz6vyYAuEnXtxPksprBd5jA05D\n"
+            + "zpzl2MKVtw/RGt9w/JHKPX8KnVrGjbyI7R9cXI7Zl/dri+234iE9b4Da0v63clMF\n"
+            + "3yDsWXejPOlyi0CLswNxlpgD0aJM/gM6IpR4O3rIxhBCvdUSoZqlenh8pj632yWq\n"
+            + "cCZwFniRAgMBAAECggEAOfC/XwQvf0KW3VciF0yNGZshbgvBUCp3p284J+ml0Smu\n"
+            + "ns4yQiaZl3B/zJ9c6nYJ8OEpNDIuGVac46vKPZIAHZf4SO4GFMFpji078IN6LmH5\n"
+            + "nclZoNn9brNKaYbgQ2N6teKgmRu8Uc7laHKXjnZd0jaWAkRP8/h0l7fDob+jaERj\n"
+            + "oJBx4ux2Z62TTCP6W4VY3KZgSL1p6dQswqlukPVytMeI2XEwWnO+w8ED0BxCxM4F\n"
+            + "K//dw7nUMGS9GUNkgyDcH1akYSCDzdBeymQBp2latBotVfGNK1hq9nC1iaxmRkJL\n"
+            + "sYjwVc24n37u+txOovy3daq2ySj9trF7ySAPVYkh4QKBgQDWeN/MR6cy1TLF2j3g\n"
+            + "eMMeM32LxXArIPsar+mft+uisKWk5LDpsKpph93sl0JjFi4x0t1mqw23h23I+B2c\n"
+            + "JWiPAHUG3FGvvkPPcfMUvd7pODyE2XaXi+36UZAH7qc94VZGJEb+sPITckSruREE\n"
+            + "QErWZyrbBRgvQXsmVme5B2/kRQKBgQDGf2HQH0KHl54O2r9vrhiQxWIIMSWlizJC\n"
+            + "hjboY6DkIsAMwnXp3wn3Bk4tSgeLk8DEVlmEaE3gvGpiIp0vQnSOlME2TXfEthdM\n"
+            + "uS3+BFXN4Vxxx/qjKL2WfZloyzdaaaF7s+LIwmXgLsFFCUSq+uLtBqfpH2Qv+paX\n"
+            + "Xqm7LN3V3QKBgH5ssj/Q3RZx5oQKqf7wMNRUteT2dbB2uI56s9SariQwzPPuevrG\n"
+            + "US30ETWt1ExkfsaP7kLfAi71fhnBaHLq+j+RnWp15REbrw1RtmC7q/L+W25UYjvj\n"
+            + "GF0+RxDl9V/cvOaL6+2mkIw2B5TSet1uqK7KEdEZp6/zgYyP0oSXhbWhAoGAdnlZ\n"
+            + "HCtMPjnUcPFHCZVTvDTTSihrW9805FfPNe0g/olvLy5xymEBRZtR1d41mq1ZhNY1\n"
+            + "H75RnS1YIbKfNrHnd6J5n7ulHJfCWFy+grp7rCIyVwcRJYkPf17/zXhdVW1uoLLB\n"
+            + "TSoaPDAr0tSxU4vjHa23UoEV/z0F3Nr3W2xwC1ECgYBHKjv6ekLhx7HbP797+Ai+\n"
+            + "wkHvS2L/MqEBxuHzcQ9G6Mj3ANAeyDB8YSC8qGtDQoEyukv2dO73lpodNgbR8P+Q\n"
+            + "PDBb6eyntAo2sSeo0jZkiXvDOfRaGuGVrxjuTfaqcVB33jC6BYfi61/3Sr5oG9Nd\n"
+            + "tDGh1HlOIRm1jD9KQNVZ/Q==\n"
+            + "-----END PRIVATE KEY-----";
+
+    private static final String ENCRYPTED_KEY =  "-----BEGIN ENCRYPTED PRIVATE KEY-----\n"
+            + "MIIE6jAcBgoqhkiG9w0BDAEDMA4ECGyAEWAXlaXzAgIIAASCBMgt7QD1Bbz7MAHI\n"
+            + "Ni0eTrwNiuAPluHirLXzsV57d1O9i4EXVp5nzRy6753cjXbGXARbBeaJD+/+jbZp\n"
+            + "CBZTHMG8rTCfbsg5kMqxT6XuuqWlKLKc4gaq+QNgHHleKqnpwZQmOQ+awKWEK/Ow\n"
+            + "Z0KxXqkp+b4/qJK3MqKZDsJtVdyUhO0tLVxd+BHDg9B93oExc87F16h3R0+T4rxE\n"
+            + "Tvz2c2upBqva49AbLDxpWXLCJC8CRkxM+KHrPkYjpNx3jCjtwiiXfzJCWjuCkVrL\n"
+            + "2F4bqvpYPIseoPtMvWaplNtoPwhpzBB/hoJ+R+URr4XHX3Y+bz6k6iQnhoCOIviy\n"
+            + "oEEUvWtKnaEEKSauR+Wyj3MoeB64g9NWMEHv7+SQeA4WqlgV2s4txwRxFGKyKLPq\n"
+            + "caMSpfxvYujtSh0DOv9GI3cVHPM8WsebCz9cNrbKSR8/8JufcoonTitwF/4vm1Et\n"
+            + "AdmCuH9JIYVvmFKFVxY9SvRAvo43OQaPmJQHMUa4yDfMtpTSgmB/7HFgxtksYs++\n"
+            + "Gbrq6F/hon+0bLx+bMz2FK635UU+iVno+qaScKWN3BFqDl+KnZprBhLSXTT3aHmp\n"
+            + "fisQit/HWp71a0Vzq85WwI4ucMKNc8LemlwNBxWLLiJDp7sNPLb5dIl8yIwSEIgd\n"
+            + "vC5px9KWEdt3GxTUEqtIeBmagbBhahcv+c9Dq924DLI+Slv6TJKZpIcMqUECgzvi\n"
+            + "hb8gegyEscBEcDSzl0ojlFVz4Va5eZS/linTjNJhnkx8BKLn/QFco7FpEE6uOmQ3\n"
+            + "0kF64M2Rv67cJbYVrhD46TgIzH3Y/FOMSi1zFHQ14nVXWMu0yAlBX+QGk7Xl+/aF\n"
+            + "BIq+i9WcBqbttR3CwyeTnIFXkdC66iTZYhDl9HT6yMcazql2Or2TjIIWr6tfNWH/\n"
+            + "5dWSEHYM5m8F2/wF0ANWJyR1oPr4ckcUsfl5TfOWVj5wz4QVF6EGV7FxEnQHrdx0\n"
+            + "6rXThRKFjqxUubsNt1yUEwdlTNz2UFhobGF9MmFeB97BZ6T4v8G825de/Caq9FzO\n"
+            + "yMFFCRcGC7gIzMXRPEjHIvBdTThm9rbNzKPXHqw0LHG478yIqzxvraCYTRw/4eWN\n"
+            + "Q+hyOL/5T5QNXHpR8Udp/7sptw7HfRnecQ/Vz9hOKShQq3h4Sz6eQMQm7P9qGo/N\n"
+            + "bltEAIECRVcNYLN8LuEORfeecNcV3BX+4BBniFtdD2bIRsWC0ZUsGf14Yhr4P1OA\n"
+            + "PtMJzy99mrcq3h+o+hEW6bhIj1gA88JSMJ4iRuwTLRKE81w7EyziScDsotYKvDPu\n"
+            + "w4+PFbQO3fr/Zga3LgYis8/DMqZoWjVCjAeVoypuOZreieZYC/BgBS8qSUAmDPKq\n"
+            + "jK+T5pwMMchfXbkV80LTu1kqLfKWdE0AmZfGy8COE/NNZ/FeiWZPdwu2Ix6u/RoY\n"
+            + "LTjNy4YLIBdVELFXaFJF2GfzLpnwrW5tyNPVVrGmUoiyOzgx8gMyCLGavGtduyoY\n"
+            + "tBiUTmd05Ugscn4Rz9X30S4NbnjL/h+bWl1m6/M+9FHEe85FPxmt/GRmJPbFPMR5\n"
+            + "q5EgQGkt4ifiaP6qvyFulwvVwx+m0bf1q6Vb/k3clIyLMcVZWFE1TqNH2Ife46AE\n"
+            + "2I39ZnGTt0mbWskpHBA=\n"
+            + "-----END ENCRYPTED PRIVATE KEY-----";
+
+    private static final Password KEY_PASSWORD = new Password("key-password");
+
+    private DefaultSslEngineFactory factory = new DefaultSslEngineFactory();
+    Map configs = new HashMap<>();
+
+    @Before
+    public void setUp() {
+        factory = new DefaultSslEngineFactory();
+        configs.put(SslConfigs.SSL_PROTOCOL_CONFIG, "TLSv1.2");
+    }
+
+    @Test
+    public void testPemTrustStoreConfigWithOneCert() throws Exception {
+        configs.put(SslConfigs.SSL_TRUSTSTORE_CERTIFICATES_CONFIG, pemAsConfigValue(CA1));
+        configs.put(SslConfigs.SSL_TRUSTSTORE_TYPE_CONFIG, DefaultSslEngineFactory.PEM_TYPE);
+        factory.configure(configs);
+
+        KeyStore trustStore = factory.truststore();
+        List aliases = Collections.list(trustStore.aliases());
+        assertEquals(Collections.singletonList("kafka0"), aliases);
+        assertNotNull("Certificate not loaded", trustStore.getCertificate("kafka0"));
+        assertNull("Unexpected private key", trustStore.getKey("kafka0", null));
+    }
+
+    @Test
+    public void testPemTrustStoreConfigWithMultipleCerts() throws Exception {
+        configs.put(SslConfigs.SSL_TRUSTSTORE_CERTIFICATES_CONFIG, pemAsConfigValue(CA1, CA2));
+        configs.put(SslConfigs.SSL_TRUSTSTORE_TYPE_CONFIG, DefaultSslEngineFactory.PEM_TYPE);
+        factory.configure(configs);
+
+        KeyStore trustStore = factory.truststore();
+        List aliases = Collections.list(trustStore.aliases());
+        assertEquals(Arrays.asList("kafka0", "kafka1"), aliases);
+        assertNotNull("Certificate not loaded", trustStore.getCertificate("kafka0"));
+        assertNull("Unexpected private key", trustStore.getKey("kafka0", null));
+        assertNotNull("Certificate not loaded", trustStore.getCertificate("kafka1"));
+        assertNull("Unexpected private key", trustStore.getKey("kafka1", null));
+    }
+
+    @Test
+    public void testPemKeyStoreConfigNoPassword() throws Exception {
+        verifyPemKeyStoreConfig(KEY, null);
+    }
+
+    @Test
+    public void testPemKeyStoreConfigWithKeyPassword() throws Exception {
+        verifyPemKeyStoreConfig(ENCRYPTED_KEY, KEY_PASSWORD);
+    }
+
+    @Test
+    public void testTrailingNewLines() throws Exception {
+        verifyPemKeyStoreConfig(ENCRYPTED_KEY + "\n\n", KEY_PASSWORD);
+    }
+
+    @Test
+    public void testLeadingNewLines() throws Exception {
+        verifyPemKeyStoreConfig("\n\n" + ENCRYPTED_KEY, KEY_PASSWORD);
+    }
+
+    @Test
+    public void testCarriageReturnLineFeed() throws Exception {
+        verifyPemKeyStoreConfig(ENCRYPTED_KEY.replaceAll("\n", "\r\n"), KEY_PASSWORD);
+    }
+
+    private void verifyPemKeyStoreConfig(String keyFileName, Password keyPassword) throws Exception {
+        configs.put(SslConfigs.SSL_KEYSTORE_KEY_CONFIG, pemAsConfigValue(keyFileName));
+        configs.put(SslConfigs.SSL_KEYSTORE_CERTIFICATE_CHAIN_CONFIG, pemAsConfigValue(CERTCHAIN));
+        configs.put(SslConfigs.SSL_KEY_PASSWORD_CONFIG, keyPassword);
+        configs.put(SslConfigs.SSL_KEYSTORE_TYPE_CONFIG, DefaultSslEngineFactory.PEM_TYPE);
+        factory.configure(configs);
+
+        KeyStore keyStore = factory.keystore();
+        List aliases = Collections.list(keyStore.aliases());
+        assertEquals(Collections.singletonList("kafka"), aliases);
+        assertNotNull("Certificate not loaded", keyStore.getCertificate("kafka"));
+        assertNotNull("Private key not loaded",
+                keyStore.getKey("kafka", keyPassword == null ? null : keyPassword.value().toCharArray()));
+    }
+
+    @Test
+    public void testPemTrustStoreFile() throws Exception {
+        configs.put(SslConfigs.SSL_TRUSTSTORE_LOCATION_CONFIG, pemFilePath(CA1));
+        configs.put(SslConfigs.SSL_TRUSTSTORE_TYPE_CONFIG, DefaultSslEngineFactory.PEM_TYPE);
+        factory.configure(configs);
+
+        KeyStore trustStore = factory.truststore();
+        List aliases = Collections.list(trustStore.aliases());
+        assertEquals(Collections.singletonList("kafka0"), aliases);
+        assertNotNull("Certificate not found", trustStore.getCertificate("kafka0"));
+        assertNull("Unexpected private key", trustStore.getKey("kafka0", null));
+    }
+
+    @Test
+    public void testPemKeyStoreFileNoKeyPassword() throws Exception {
+        configs.put(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG,
+                pemFilePath(pemAsConfigValue(KEY, CERTCHAIN).value()));
+        configs.put(SslConfigs.SSL_KEYSTORE_TYPE_CONFIG, DefaultSslEngineFactory.PEM_TYPE);
+        assertThrows(InvalidConfigurationException.class, () -> factory.configure(configs));
+    }
+
+    @Test
+    public void testPemKeyStoreFileWithKeyPassword() throws Exception {
+        configs.put(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG,
+                pemFilePath(pemAsConfigValue(ENCRYPTED_KEY, CERTCHAIN).value()));
+        configs.put(SslConfigs.SSL_KEY_PASSWORD_CONFIG, KEY_PASSWORD);
+        configs.put(SslConfigs.SSL_KEYSTORE_TYPE_CONFIG, DefaultSslEngineFactory.PEM_TYPE);
+        factory.configure(configs);
+
+        KeyStore keyStore = factory.keystore();
+        List aliases = Collections.list(keyStore.aliases());
+        assertEquals(Collections.singletonList("kafka"), aliases);
+        assertNotNull("Certificate not found", keyStore.getCertificate("kafka"));
+        assertNotNull("Private key not found", keyStore.getKey("kafka", KEY_PASSWORD.value().toCharArray()));
+    }
+
+    private String pemFilePath(String pem) throws Exception {
+        return TestUtils.tempFile(pem).getAbsolutePath();
+    }
+
+    private Password pemAsConfigValue(String... pemValues)  throws Exception {
+        StringBuilder builder = new StringBuilder();
+        for (String pem : pemValues) {
+            builder.append(pem);
+            builder.append("\n");
+        }
+        return new Password(builder.toString().trim());
+    }
+}
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 5546a55123f79..d5ffdf9c5d424 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
@@ -17,44 +17,132 @@
 package org.apache.kafka.common.security.ssl;
 
 import java.io.File;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.security.GeneralSecurityException;
+import java.security.KeyStore;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.List;
 import java.util.Map;
 
+import javax.net.ssl.SSLContext;
 import javax.net.ssl.SSLEngine;
 
+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.TestSecurityConfig;
+import org.apache.kafka.common.security.auth.SslEngineFactory;
+import org.apache.kafka.common.security.ssl.DefaultSslEngineFactory.FileBasedStore;
+import org.apache.kafka.common.security.ssl.DefaultSslEngineFactory.PemStore;
+import org.apache.kafka.common.security.ssl.DefaultSslEngineFactory.SecurityStore;
+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.Java;
+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 org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
 
-import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNotSame;
+import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
+import java.security.Security;
+import java.util.Properties;
 
-/**
- * A set of tests for the selector over ssl. These use a test harness that runs a simple socket server that echos back responses.
- */
+@RunWith(value = Parameterized.class)
 public class SslFactoryTest {
 
+    private final String tlsProtocol;
+
+    @Parameterized.Parameters(name = "tlsProtocol={0}")
+    public static Collection data() {
+        List values = new ArrayList<>();
+        values.add(new Object[] {"TLSv1.2"});
+        if (Java.IS_JAVA11_COMPATIBLE) {
+            values.add(new Object[] {"TLSv1.3"});
+        }
+        return values;
+    }
+
+    public SslFactoryTest(String tlsProtocol) {
+        this.tlsProtocol = tlsProtocol;
+    }
+
     @Test
     public void testSslFactoryConfiguration() throws Exception {
         File trustStoreFile = File.createTempFile("truststore", ".jks");
-        Map serverSslConfig = TestSslUtils.createSslConfig(false, true, Mode.SERVER, trustStoreFile, "server");
+        Map serverSslConfig = sslConfigsBuilder(Mode.SERVER)
+                .createNewTrustStore(trustStoreFile)
+                .build();
         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(tlsProtocol), Utils.mkSet(engine.getEnabledProtocols()));
         assertEquals(false, engine.getUseClientMode());
     }
 
+    @Test
+    public void testSslFactoryWithCustomKeyManagerConfiguration() {
+        TestProviderCreator testProviderCreator = new TestProviderCreator();
+        Map serverSslConfig = TestSslUtils.createSslConfig(
+                TestKeyManagerFactory.ALGORITHM,
+                TestTrustManagerFactory.ALGORITHM,
+                tlsProtocol
+        );
+        serverSslConfig.put(SecurityConfig.SECURITY_PROVIDERS_CONFIG, testProviderCreator.getClass().getName());
+        SslFactory sslFactory = new SslFactory(Mode.SERVER);
+        sslFactory.configure(serverSslConfig);
+        assertNotNull("SslEngineFactory not created", sslFactory.sslEngineFactory());
+        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,
+                tlsProtocol
+        );
+        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,
+                tlsProtocol
+        );
+        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");
-        Map serverSslConfig = TestSslUtils.createSslConfig(false, true, Mode.SERVER, trustStoreFile, "server");
+        Map serverSslConfig = sslConfigsBuilder(Mode.SERVER)
+                .createNewTrustStore(trustStoreFile)
+                .build();
         // unset the password
         serverSslConfig.remove(SslConfigs.SSL_TRUSTSTORE_PASSWORD_CONFIG);
         SslFactory sslFactory = new SslFactory(Mode.SERVER);
@@ -68,7 +156,10 @@ 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 = sslConfigsBuilder(Mode.CLIENT)
+                .createNewTrustStore(trustStoreFile)
+                .useClientCert(false)
+                .build();
         SslFactory sslFactory = new SslFactory(Mode.CLIENT);
         sslFactory.configure(clientSslConfig);
         //host and port are hints
@@ -76,4 +167,403 @@ public void testClientMode() throws Exception {
         assertTrue(engine.getUseClientMode());
     }
 
+    @Test
+    public void staleSslEngineFactoryShouldBeClosed() throws IOException, GeneralSecurityException {
+        File trustStoreFile = File.createTempFile("truststore", ".jks");
+        Map clientSslConfig = sslConfigsBuilder(Mode.SERVER)
+                .createNewTrustStore(trustStoreFile)
+                .useClientCert(false)
+                .build();
+        clientSslConfig.put(SslConfigs.SSL_ENGINE_FACTORY_CLASS_CONFIG, TestSslUtils.TestSslEngineFactory.class);
+        SslFactory sslFactory = new SslFactory(Mode.SERVER);
+        sslFactory.configure(clientSslConfig);
+        TestSslUtils.TestSslEngineFactory sslEngineFactory = (TestSslUtils.TestSslEngineFactory) sslFactory.sslEngineFactory();
+        assertNotNull(sslEngineFactory);
+        assertFalse(sslEngineFactory.closed);
+
+        trustStoreFile = File.createTempFile("truststore", ".jks");
+        clientSslConfig = sslConfigsBuilder(Mode.SERVER)
+                .createNewTrustStore(trustStoreFile)
+                .build();
+        clientSslConfig.put(SslConfigs.SSL_ENGINE_FACTORY_CLASS_CONFIG, TestSslUtils.TestSslEngineFactory.class);
+        sslFactory.reconfigure(clientSslConfig);
+        TestSslUtils.TestSslEngineFactory newSslEngineFactory = (TestSslUtils.TestSslEngineFactory) sslFactory.sslEngineFactory();
+        assertNotEquals(sslEngineFactory, newSslEngineFactory);
+        // the older one should be closed
+        assertTrue(sslEngineFactory.closed);
+    }
+
+    @Test
+    public void testReconfiguration() throws Exception {
+        File trustStoreFile = File.createTempFile("truststore", ".jks");
+        Map sslConfig = sslConfigsBuilder(Mode.SERVER)
+                .createNewTrustStore(trustStoreFile)
+                .build();
+        SslFactory sslFactory = new SslFactory(Mode.SERVER);
+        sslFactory.configure(sslConfig);
+        SslEngineFactory sslEngineFactory = sslFactory.sslEngineFactory();
+        assertNotNull("SslEngineFactory not created", sslEngineFactory);
+
+        // Verify that SslEngineFactory is not recreated on reconfigure() if config and
+        // file are not changed
+        sslFactory.reconfigure(sslConfig);
+        assertSame("SslEngineFactory recreated unnecessarily",
+                sslEngineFactory, sslFactory.sslEngineFactory());
+
+        // Verify that the SslEngineFactory is recreated on reconfigure() if config is changed
+        trustStoreFile = File.createTempFile("truststore", ".jks");
+        sslConfig = sslConfigsBuilder(Mode.SERVER)
+                .createNewTrustStore(trustStoreFile)
+                .build();
+        sslFactory.reconfigure(sslConfig);
+        assertNotSame("SslEngineFactory not recreated",
+                sslEngineFactory, sslFactory.sslEngineFactory());
+        sslEngineFactory = sslFactory.sslEngineFactory();
+
+        // 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("SslEngineFactory not recreated",
+                sslEngineFactory, sslFactory.sslEngineFactory());
+        sslEngineFactory = sslFactory.sslEngineFactory();
+
+        // 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("SslEngineFactory not recreated",
+                sslEngineFactory, sslFactory.sslEngineFactory());
+        sslEngineFactory = sslFactory.sslEngineFactory();
+
+        // 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("SslEngineFactory not recreated",
+                sslEngineFactory, sslFactory.sslEngineFactory());
+        sslEngineFactory = sslFactory.sslEngineFactory();
+
+        // 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("SslEngineFactory recreated unnecessarily",
+                sslEngineFactory, sslFactory.sslEngineFactory());
+    }
+
+    @Test
+    public void testReconfigurationWithoutTruststore() throws Exception {
+        File trustStoreFile = File.createTempFile("truststore", ".jks");
+        Map sslConfig = sslConfigsBuilder(Mode.SERVER)
+                .createNewTrustStore(trustStoreFile)
+                .build();
+        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(sslConfig);
+        SSLContext sslContext = ((DefaultSslEngineFactory) sslFactory.sslEngineFactory()).sslContext();
+        assertNotNull("SSL context not created", sslContext);
+        assertSame("SSL context recreated unnecessarily", sslContext,
+                ((DefaultSslEngineFactory) sslFactory.sslEngineFactory()).sslContext());
+        assertFalse(sslFactory.createSslEngine("localhost", 0).getUseClientMode());
+
+        Map sslConfig2 = sslConfigsBuilder(Mode.SERVER)
+                .createNewTrustStore(trustStoreFile)
+                .build();
+        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 = sslConfigsBuilder(Mode.SERVER)
+                .createNewTrustStore(trustStoreFile)
+                .build();
+        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 = ((DefaultSslEngineFactory) sslFactory.sslEngineFactory()).sslContext();
+        assertNotNull("SSL context not created", sslContext);
+        assertSame("SSL context recreated unnecessarily", sslContext,
+                ((DefaultSslEngineFactory) sslFactory.sslEngineFactory()).sslContext());
+        assertFalse(sslFactory.createSslEngine("localhost", 0).getUseClientMode());
+
+        File newTrustStoreFile = File.createTempFile("truststore", ".jks");
+        sslConfig = sslConfigsBuilder(Mode.SERVER)
+                .createNewTrustStore(newTrustStoreFile)
+                .build();
+        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,
+                ((DefaultSslEngineFactory) sslFactory.sslEngineFactory()).sslContext());
+
+        sslConfig = sslConfigsBuilder(Mode.SERVER)
+                .createNewTrustStore(newTrustStoreFile)
+                .build();
+        try {
+            sslFactory.validateReconfiguration(sslConfig);
+            fail("Keystore configured dynamically for listener without previous keystore");
+        } catch (ConfigException e) {
+            // Expected exception
+        }
+    }
+
+    @Test
+    public void testPemReconfiguration() throws Exception {
+        Properties props = new Properties();
+        props.putAll(sslConfigsBuilder(Mode.SERVER)
+                .createNewTrustStore(null)
+                .usePem(true)
+                .build());
+        TestSecurityConfig sslConfig = new TestSecurityConfig(props);
+
+        SslFactory sslFactory = new SslFactory(Mode.SERVER);
+        sslFactory.configure(sslConfig.values());
+        SslEngineFactory sslEngineFactory = sslFactory.sslEngineFactory();
+        assertNotNull("SslEngineFactory not created", sslEngineFactory);
+
+        props.put("some.config", "some.value");
+        sslConfig = new TestSecurityConfig(props);
+        sslFactory.reconfigure(sslConfig.values());
+        assertSame("SslEngineFactory recreated unnecessarily", sslEngineFactory, sslFactory.sslEngineFactory());
+
+        props.put(SslConfigs.SSL_KEYSTORE_KEY_CONFIG,
+                new Password(((Password) props.get(SslConfigs.SSL_KEYSTORE_KEY_CONFIG)).value() + " "));
+        sslConfig = new TestSecurityConfig(props);
+        sslFactory.reconfigure(sslConfig.values());
+        assertNotSame("SslEngineFactory not recreated", sslEngineFactory, sslFactory.sslEngineFactory());
+        sslEngineFactory = sslFactory.sslEngineFactory();
+
+        props.put(SslConfigs.SSL_KEYSTORE_CERTIFICATE_CHAIN_CONFIG,
+                new Password(((Password) props.get(SslConfigs.SSL_KEYSTORE_CERTIFICATE_CHAIN_CONFIG)).value() + " "));
+        sslConfig = new TestSecurityConfig(props);
+        sslFactory.reconfigure(sslConfig.values());
+        assertNotSame("SslEngineFactory not recreated", sslEngineFactory, sslFactory.sslEngineFactory());
+        sslEngineFactory = sslFactory.sslEngineFactory();
+
+        props.put(SslConfigs.SSL_TRUSTSTORE_CERTIFICATES_CONFIG,
+                new Password(((Password) props.get(SslConfigs.SSL_TRUSTSTORE_CERTIFICATES_CONFIG)).value() + " "));
+        sslConfig = new TestSecurityConfig(props);
+        sslFactory.reconfigure(sslConfig.values());
+        assertNotSame("SslEngineFactory not recreated", sslEngineFactory, sslFactory.sslEngineFactory());
+        sslEngineFactory = sslFactory.sslEngineFactory();
+    }
+
+    @Test
+    public void testKeyStoreTrustStoreValidation() throws Exception {
+        File trustStoreFile = File.createTempFile("truststore", ".jks");
+        Map serverSslConfig = sslConfigsBuilder(Mode.SERVER)
+                .createNewTrustStore(trustStoreFile)
+                .build();
+        SslFactory sslFactory = new SslFactory(Mode.SERVER);
+        sslFactory.configure(serverSslConfig);
+        assertNotNull("SslEngineFactory not created", sslFactory.sslEngineFactory());
+    }
+
+    @Test
+    public void testUntrustedKeyStoreValidationFails() throws Exception {
+        File trustStoreFile1 = File.createTempFile("truststore1", ".jks");
+        File trustStoreFile2 = File.createTempFile("truststore2", ".jks");
+        Map sslConfig1 = sslConfigsBuilder(Mode.SERVER)
+                .createNewTrustStore(trustStoreFile1)
+                .build();
+        Map sslConfig2 = sslConfigsBuilder(Mode.SERVER)
+                .createNewTrustStore(trustStoreFile2)
+                .build();
+        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.configure(sslConfig1);
+            fail("Validation did not fail with untrusted truststore");
+        } catch (ConfigException e) {
+            // Expected exception
+        }
+    }
+
+    @Test
+    public void testKeystoreVerifiableUsingTruststore() throws Exception {
+        verifyKeystoreVerifiableUsingTruststore(false);
+    }
+
+    @Test
+    public void testPemKeystoreVerifiableUsingTruststore() throws Exception {
+        verifyKeystoreVerifiableUsingTruststore(true);
+    }
+
+    private void verifyKeystoreVerifiableUsingTruststore(boolean usePem) throws Exception {
+        File trustStoreFile1 = usePem ? null : File.createTempFile("truststore1", ".jks");
+        Map sslConfig1 = sslConfigsBuilder(Mode.SERVER)
+                .createNewTrustStore(trustStoreFile1)
+                .usePem(usePem)
+                .build();
+        SslFactory sslFactory = new SslFactory(Mode.SERVER, null, true);
+        sslFactory.configure(sslConfig1);
+
+        File trustStoreFile2 = usePem ? null : File.createTempFile("truststore2", ".jks");
+        Map sslConfig2 = sslConfigsBuilder(Mode.SERVER)
+                .createNewTrustStore(trustStoreFile2)
+                .usePem(usePem)
+                .build();
+        // 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.validateReconfiguration(sslConfig2);
+            fail("ValidateReconfiguration did not fail as expected");
+        } catch (ConfigException e) {
+            // Expected exception
+        }
+    }
+
+    @Test
+    public void testCertificateEntriesValidation() throws Exception {
+        verifyCertificateEntriesValidation(false);
+    }
+
+    @Test
+    public void testPemCertificateEntriesValidation() throws Exception {
+        verifyCertificateEntriesValidation(true);
+    }
+
+    private void verifyCertificateEntriesValidation(boolean usePem) throws Exception {
+        File trustStoreFile = usePem ? null : File.createTempFile("truststore", ".jks");
+        Map serverSslConfig = sslConfigsBuilder(Mode.SERVER)
+                .createNewTrustStore(trustStoreFile)
+                .usePem(usePem)
+                .build();
+        File newTrustStoreFile = usePem ? null : File.createTempFile("truststore", ".jks");
+        Map newCnConfig = sslConfigsBuilder(Mode.SERVER)
+                .createNewTrustStore(newTrustStoreFile)
+                .cn("Another CN")
+                .usePem(usePem)
+                .build();
+        KeyStore ks1 = sslKeyStore(serverSslConfig);
+        KeyStore ks2 = sslKeyStore(serverSslConfig);
+        assertEquals(SslFactory.CertificateEntries.create(ks1), SslFactory.CertificateEntries.create(ks2));
+
+        // Use different alias name, validation should succeed
+        ks2.setCertificateEntry("another", ks1.getCertificate("localhost"));
+        assertEquals(SslFactory.CertificateEntries.create(ks1), SslFactory.CertificateEntries.create(ks2));
+
+        KeyStore ks3 = sslKeyStore(newCnConfig);
+        assertNotEquals(SslFactory.CertificateEntries.create(ks1), SslFactory.CertificateEntries.create(ks3));
+    }
+
+    /**
+     * Tests client side ssl.engine.factory configuration is used when specified
+     */
+    @Test
+    public void testClientSpecifiedSslEngineFactoryUsed() throws Exception {
+        File trustStoreFile = File.createTempFile("truststore", ".jks");
+        Map clientSslConfig = sslConfigsBuilder(Mode.CLIENT)
+                .createNewTrustStore(trustStoreFile)
+                .useClientCert(false)
+                .build();
+        clientSslConfig.put(SslConfigs.SSL_ENGINE_FACTORY_CLASS_CONFIG, TestSslUtils.TestSslEngineFactory.class);
+        SslFactory sslFactory = new SslFactory(Mode.CLIENT);
+        sslFactory.configure(clientSslConfig);
+        assertTrue("SslEngineFactory must be of expected type",
+                sslFactory.sslEngineFactory() instanceof TestSslUtils.TestSslEngineFactory);
+    }
+
+    @Test
+    public void testEngineFactoryClosed() throws Exception {
+        File trustStoreFile = File.createTempFile("truststore", ".jks");
+        Map clientSslConfig = sslConfigsBuilder(Mode.CLIENT)
+                .createNewTrustStore(trustStoreFile)
+                .useClientCert(false)
+                .build();
+        clientSslConfig.put(SslConfigs.SSL_ENGINE_FACTORY_CLASS_CONFIG, TestSslUtils.TestSslEngineFactory.class);
+        SslFactory sslFactory = new SslFactory(Mode.CLIENT);
+        sslFactory.configure(clientSslConfig);
+        TestSslUtils.TestSslEngineFactory engine = (TestSslUtils.TestSslEngineFactory) sslFactory.sslEngineFactory();
+        assertFalse(engine.closed);
+        sslFactory.close();
+        assertTrue(engine.closed);
+    }
+
+    /**
+     * Tests server side ssl.engine.factory configuration is used when specified
+     */
+    @Test
+    public void testServerSpecifiedSslEngineFactoryUsed() throws Exception {
+        File trustStoreFile = File.createTempFile("truststore", ".jks");
+        Map serverSslConfig = sslConfigsBuilder(Mode.SERVER)
+                .createNewTrustStore(trustStoreFile)
+                .useClientCert(false)
+                .build();
+        serverSslConfig.put(SslConfigs.SSL_ENGINE_FACTORY_CLASS_CONFIG, TestSslUtils.TestSslEngineFactory.class);
+        SslFactory sslFactory = new SslFactory(Mode.SERVER);
+        sslFactory.configure(serverSslConfig);
+        assertTrue("SslEngineFactory must be of expected type",
+                sslFactory.sslEngineFactory() instanceof TestSslUtils.TestSslEngineFactory);
+    }
+
+    /**
+     * Tests invalid ssl.engine.factory configuration
+     */
+    @Test(expected = ClassCastException.class)
+    public void testInvalidSslEngineFactory() throws Exception {
+        File trustStoreFile = File.createTempFile("truststore", ".jks");
+        Map clientSslConfig = sslConfigsBuilder(Mode.CLIENT)
+                .createNewTrustStore(trustStoreFile)
+                .useClientCert(false)
+                .build();
+        clientSslConfig.put(SslConfigs.SSL_ENGINE_FACTORY_CLASS_CONFIG, String.class);
+        SslFactory sslFactory = new SslFactory(Mode.CLIENT);
+        sslFactory.configure(clientSslConfig);
+    }
+
+    @Test
+    public void testUsedConfigs() throws IOException, GeneralSecurityException {
+        Map serverSslConfig = sslConfigsBuilder(Mode.SERVER)
+                .createNewTrustStore(File.createTempFile("truststore", ".jks"))
+                .useClientCert(false)
+                .build();
+        serverSslConfig.put(SslConfigs.SSL_ENGINE_FACTORY_CLASS_CONFIG, TestSslUtils.TestSslEngineFactory.class);
+        TestSecurityConfig securityConfig = new TestSecurityConfig(serverSslConfig);
+        SslFactory sslFactory = new SslFactory(Mode.SERVER);
+        sslFactory.configure(securityConfig.values());
+        assertFalse(securityConfig.unused().contains(SslConfigs.SSL_ENGINE_FACTORY_CLASS_CONFIG));
+    }
+
+    private KeyStore sslKeyStore(Map sslConfig) {
+        SecurityStore store;
+        if (sslConfig.get(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG) != null) {
+            store = new FileBasedStore(
+                    (String) sslConfig.get(SslConfigs.SSL_KEYSTORE_TYPE_CONFIG),
+                    (String) sslConfig.get(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG),
+                    (Password) sslConfig.get(SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG),
+                    (Password) sslConfig.get(SslConfigs.SSL_KEY_PASSWORD_CONFIG),
+                    true
+            );
+        } else {
+            store = new PemStore(
+                    (Password) sslConfig.get(SslConfigs.SSL_KEYSTORE_CERTIFICATE_CHAIN_CONFIG),
+                    (Password) sslConfig.get(SslConfigs.SSL_KEYSTORE_KEY_CONFIG),
+                    (Password) sslConfig.get(SslConfigs.SSL_KEY_PASSWORD_CONFIG)
+            );
+        }
+        return store.get();
+    }
+
+    private TestSslUtils.SslConfigsBuilder sslConfigsBuilder(Mode mode) {
+        return new TestSslUtils.SslConfigsBuilder(mode).tlsProtocol(tlsProtocol);
+    }
 }
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
new file mode 100644
index 0000000000000..7b7b67b1dbc3a
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/security/ssl/SslPrincipalMapperTest.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.security.ssl;
+
+import org.junit.Test;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.fail;
+
+public class SslPrincipalMapperTest {
+
+    @Test
+    public void testValidRules() {
+        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/");
+    }
+
+    private void testValidRule(String rules) {
+        SslPrincipalMapper.fromRules(rules);
+    }
+
+    @Test
+    public void testInvalidRules() {
+        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");
+    }
+
+    private void testInvalidRule(String rules) {
+        try {
+            System.out.println(SslPrincipalMapper.fromRules(rules));
+            fail("should have thrown IllegalArgumentException");
+        } catch (IllegalArgumentException e) {
+        }
+    }
+
+    @Test
+    public void testSslPrincipalMapper() throws Exception {
+        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",
+            "RULE:^.*[Cc][Nn]=([a-zA-Z0-9.]*).*$/$1/U",
+            "DEFAULT"
+        );
+
+        SslPrincipalMapper mapper = SslPrincipalMapper.fromRules(rules);
+
+        assertEquals("duke", mapper.getName("CN=Duke,OU=ServiceUsers,O=Org,C=US"));
+        assertEquals("duke@sme", mapper.getName("CN=Duke,OU=SME,O=mycp,L=Fulton,ST=MD,C=US"));
+        assertEquals("DUKE@SME", mapper.getName("cn=duke,ou=sme,dc=mycp,dc=com"));
+        assertEquals("DUKE", mapper.getName("cN=duke,OU=JavaSoft,O=Sun Microsystems"));
+        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/serialization/SerializationTest.java b/clients/src/test/java/org/apache/kafka/common/serialization/SerializationTest.java
index 134882f002b3d..8f231b271b292 100644
--- a/clients/src/test/java/org/apache/kafka/common/serialization/SerializationTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/serialization/SerializationTest.java
@@ -25,6 +25,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.UUID;
 
 import static org.hamcrest.CoreMatchers.equalTo;
 import static org.hamcrest.MatcherAssert.assertThat;
@@ -34,27 +35,29 @@
 public class SerializationTest {
 
     final private String topic = "testTopic";
-    final private Map, List> testData = new HashMap() {
+    final private Map, List> testData = new HashMap, List>() {
         {
             put(String.class, Arrays.asList("my string"));
             put(Short.class, Arrays.asList((short) 32767, (short) -32768));
-            put(Integer.class, Arrays.asList((int) 423412424, (int) -41243432));
+            put(Integer.class, Arrays.asList(423412424, -41243432));
             put(Long.class, Arrays.asList(922337203685477580L, -922337203685477581L));
             put(Float.class, Arrays.asList(5678567.12312f, -5678567.12341f));
             put(Double.class, Arrays.asList(5678567.12312d, -5678567.12341d));
             put(byte[].class, Arrays.asList("my string".getBytes()));
             put(ByteBuffer.class, Arrays.asList(ByteBuffer.allocate(10).put("my string".getBytes())));
             put(Bytes.class, Arrays.asList(new Bytes("my string".getBytes())));
+            put(UUID.class, Arrays.asList(UUID.randomUUID()));
         }
     };
 
     private class DummyClass {
     }
 
+    @SuppressWarnings("unchecked")
     @Test
     public void allSerdesShouldRoundtripInput() {
-        for (Map.Entry, List> test : testData.entrySet()) {
-            try (Serde serde = Serdes.serdeFrom(test.getKey())) {
+        for (Map.Entry, List> test : testData.entrySet()) {
+            try (Serde serde = Serdes.serdeFrom((Class) test.getKey())) {
                 for (Object value : test.getValue()) {
                     assertEquals("Should get the original " + test.getKey().getSimpleName() +
                                     " after serialization and deserialization", value,
@@ -144,6 +147,27 @@ public void floatSerdeShouldPreserveNaNValues() {
         }
     }
 
+    @Test
+    public void testSerializeVoid() {
+        try (Serde serde = Serdes.Void()) {
+            serde.serializer().serialize(topic, null);
+        }
+    }
+
+    @Test
+    public void testDeserializeVoid() {
+        try (Serde serde = Serdes.Void()) {
+            serde.deserializer().deserialize(topic, null);
+        }
+    }
+
+    @Test(expected = IllegalArgumentException.class)
+    public void voidDeserializerShouldThrowOnNotNullValues() {
+        try (Serde serde = Serdes.Void()) {
+            serde.deserializer().deserialize(topic, new byte[5]);
+        }
+    }
+
     private Serde getStringSerde(String encoder) {
         Map serializerConfigs = new HashMap();
         serializerConfigs.put("key.serializer.encoding", encoder);
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/ByteBufferUnmapperTest.java b/clients/src/test/java/org/apache/kafka/common/utils/ByteBufferUnmapperTest.java
new file mode 100644
index 0000000000000..ca9fcd907b504
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/utils/ByteBufferUnmapperTest.java
@@ -0,0 +1,41 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.kafka.common.utils;
+
+import org.apache.kafka.test.TestUtils;
+import org.junit.Test;
+
+import java.io.File;
+import java.nio.MappedByteBuffer;
+import java.nio.channels.FileChannel;
+
+public class ByteBufferUnmapperTest {
+
+    /**
+     * Checks that unmap doesn't throw exceptions.
+     */
+    @Test
+    public void testUnmap() throws Exception {
+        File file = TestUtils.tempFile();
+        try (FileChannel channel = FileChannel.open(file.toPath())) {
+            MappedByteBuffer map = channel.map(FileChannel.MapMode.READ_ONLY, 0, 0);
+            ByteBufferUnmapper.unmap(file.getAbsolutePath(), map);
+        }
+    }
+
+}
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..7e599f9614027 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,39 @@ public void testInvalidVarlong() {
         ByteUtils.readVarlong(buf);
     }
 
+    @Test
+    public void testDouble() throws IOException {
+        assertDoubleSerde(0.0, 0x0L);
+        assertDoubleSerde(-0.0, 0x8000000000000000L);
+        assertDoubleSerde(1.0, 0x3FF0000000000000L);
+        assertDoubleSerde(-1.0, 0xBFF0000000000000L);
+        assertDoubleSerde(123e45, 0x49B58B82C0E0BB00L);
+        assertDoubleSerde(-123e45, 0xC9B58B82C0E0BB00L);
+        assertDoubleSerde(Double.MIN_VALUE, 0x1L);
+        assertDoubleSerde(-Double.MIN_VALUE, 0x8000000000000001L);
+        assertDoubleSerde(Double.MAX_VALUE, 0x7FEFFFFFFFFFFFFFL);
+        assertDoubleSerde(-Double.MAX_VALUE, 0xFFEFFFFFFFFFFFFFL);
+        assertDoubleSerde(Double.NaN, 0x7FF8000000000000L);
+        assertDoubleSerde(Double.POSITIVE_INFINITY, 0x7FF0000000000000L);
+        assertDoubleSerde(Double.NEGATIVE_INFINITY, 0xFFF0000000000000L);
+    }
+
+    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);
@@ -229,4 +286,25 @@ private void assertVarlongSerde(long value, byte[] expectedEncoding) throws IOEx
         assertEquals(value, ByteUtils.readVarlong(in));
     }
 
+    private void assertDoubleSerde(double value, long expectedLongValue) throws IOException {
+        byte[] expectedEncoding = new byte[8];
+        for (int i = 0; i < 8; i++) {
+            expectedEncoding[7 - i] = (byte) (expectedLongValue & 0xFF);
+            expectedLongValue >>= 8;
+        }
+
+        ByteBuffer buf = ByteBuffer.allocate(8);
+        ByteUtils.writeDouble(value, buf);
+        buf.flip();
+        assertEquals(value, ByteUtils.readDouble(buf.duplicate()), 0.0);
+        assertArrayEquals(expectedEncoding, Utils.toArray(buf));
+
+        buf.rewind();
+        DataOutputStream out = new DataOutputStream(new ByteBufferOutputStream(buf));
+        ByteUtils.writeDouble(value, out);
+        buf.flip();
+        assertArrayEquals(expectedEncoding, Utils.toArray(buf));
+        DataInputStream in = new DataInputStream(new ByteBufferInputStream(buf));
+        assertEquals(value, ByteUtils.readDouble(in), 0.0);
+    }
 }
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/CollectionUtilsTest.java b/clients/src/test/java/org/apache/kafka/common/utils/CollectionUtilsTest.java
new file mode 100644
index 0000000000000..7abf08a4e76a1
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/utils/CollectionUtilsTest.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.utils;
+
+import org.junit.Test;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.apache.kafka.common.utils.CollectionUtils.subtractMap;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.assertNotSame;
+
+public class CollectionUtilsTest {
+
+    @Test
+    public void testSubtractMapRemovesSecondMapsKeys() {
+        Map mainMap = new HashMap<>();
+        mainMap.put("one", "1");
+        mainMap.put("two", "2");
+        mainMap.put("three", "3");
+        Map secondaryMap = new HashMap<>();
+        secondaryMap.put("one", "4");
+        secondaryMap.put("two", "5");
+
+        Map newMap = subtractMap(mainMap, secondaryMap);
+
+        assertEquals(3, mainMap.size());  // original map should not be modified
+        assertEquals(1, newMap.size());
+        assertTrue(newMap.containsKey("three"));
+        assertEquals("3", newMap.get("three"));
+    }
+
+    @Test
+    public void testSubtractMapDoesntRemoveAnythingWhenEmptyMap() {
+        Map mainMap = new HashMap<>();
+        mainMap.put("one", "1");
+        mainMap.put("two", "2");
+        mainMap.put("three", "3");
+        Map secondaryMap = new HashMap<>();
+
+        Map newMap = subtractMap(mainMap, secondaryMap);
+
+        assertEquals(3, newMap.size());
+        assertEquals("1", newMap.get("one"));
+        assertEquals("2", newMap.get("two"));
+        assertEquals("3", newMap.get("three"));
+        assertNotSame(newMap, mainMap);
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/utils/ConfigUtilsTest.java b/clients/src/test/java/org/apache/kafka/common/utils/ConfigUtilsTest.java
new file mode 100644
index 0000000000000..c77cb0e987427
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/utils/ConfigUtilsTest.java
@@ -0,0 +1,143 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.HashMap;
+import java.util.Map;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+
+public class ConfigUtilsTest {
+
+    @Test
+    public void testTranslateDeprecated() {
+        Map config = new HashMap<>();
+        config.put("foo.bar", "baz");
+        config.put("foo.bar.deprecated", "quux");
+        config.put("chicken", "1");
+        config.put("rooster", "2");
+        config.put("hen", "3");
+        config.put("heifer", "moo");
+        config.put("blah", "blah");
+        config.put("unexpected.non.string.object", 42);
+        Map newConfig = ConfigUtils.translateDeprecatedConfigs(config, new String[][]{
+            {"foo.bar", "foo.bar.deprecated"},
+            {"chicken", "rooster", "hen"},
+            {"cow", "beef", "heifer", "steer"}
+        });
+        assertEquals("baz", newConfig.get("foo.bar"));
+        assertEquals(null, newConfig.get("foobar.deprecated"));
+        assertEquals("1", newConfig.get("chicken"));
+        assertEquals(null, newConfig.get("rooster"));
+        assertEquals(null, newConfig.get("hen"));
+        assertEquals("moo", newConfig.get("cow"));
+        assertEquals(null, newConfig.get("beef"));
+        assertEquals(null, newConfig.get("heifer"));
+        assertEquals(null, newConfig.get("steer"));
+        assertEquals(null, config.get("cow"));
+        assertEquals("blah", config.get("blah"));
+        assertEquals("blah", newConfig.get("blah"));
+        assertEquals(42, newConfig.get("unexpected.non.string.object"));
+        assertEquals(42, config.get("unexpected.non.string.object"));
+
+    }
+
+    @Test
+    public void testAllowsNewKey() {
+        Map config = new HashMap<>();
+        config.put("foo.bar", "baz");
+        Map newConfig = ConfigUtils.translateDeprecatedConfigs(config, new String[][]{
+            {"foo.bar", "foo.bar.deprecated"},
+            {"chicken", "rooster", "hen"},
+            {"cow", "beef", "heifer", "steer"}
+        });
+        assertNotNull(newConfig);
+        assertEquals("baz", newConfig.get("foo.bar"));
+        assertNull(newConfig.get("foo.bar.deprecated"));
+    }
+
+    @Test
+    public void testAllowDeprecatedNulls() {
+        Map config = new HashMap<>();
+        config.put("foo.bar.deprecated", null);
+        config.put("foo.bar", "baz");
+        Map newConfig = ConfigUtils.translateDeprecatedConfigs(config, new String[][]{
+            {"foo.bar", "foo.bar.deprecated"}
+        });
+        assertNotNull(newConfig);
+        assertEquals("baz", newConfig.get("foo.bar"));
+        assertNull(newConfig.get("foo.bar.deprecated"));
+    }
+
+    @Test
+    public void testAllowNullOverride() {
+        Map config = new HashMap<>();
+        config.put("foo.bar.deprecated", "baz");
+        config.put("foo.bar", null);
+        Map newConfig = ConfigUtils.translateDeprecatedConfigs(config, new String[][]{
+            {"foo.bar", "foo.bar.deprecated"}
+        });
+        assertNotNull(newConfig);
+        assertNull(newConfig.get("foo.bar"));
+        assertNull(newConfig.get("foo.bar.deprecated"));
+    }
+
+    @Test
+    public void testNullMapEntriesWithoutAliasesDoNotThrowNPE() {
+        Map config = new HashMap<>();
+        config.put("other", null);
+        Map newConfig = ConfigUtils.translateDeprecatedConfigs(config, new String[][]{
+            {"foo.bar", "foo.bar.deprecated"}
+        });
+        assertNotNull(newConfig);
+        assertNull(newConfig.get("other"));
+    }
+
+    @Test
+    public void testDuplicateSynonyms() {
+        Map config = new HashMap<>();
+        config.put("foo.bar", "baz");
+        config.put("foo.bar.deprecated", "derp");
+        Map newConfig = ConfigUtils.translateDeprecatedConfigs(config, new String[][]{
+            {"foo.bar", "foo.bar.deprecated"},
+            {"chicken", "foo.bar.deprecated"}
+        });
+        assertNotNull(newConfig);
+        assertEquals("baz", newConfig.get("foo.bar"));
+        assertEquals("derp", newConfig.get("chicken"));
+        assertNull(newConfig.get("foo.bar.deprecated"));
+    }
+
+    @Test
+    public void testMultipleDeprecations() {
+        Map config = new HashMap<>();
+        config.put("foo.bar.deprecated", "derp");
+        config.put("foo.bar.even.more.deprecated", "very old configuration");
+        Map newConfig = ConfigUtils.translateDeprecatedConfigs(config, new String[][]{
+            {"foo.bar", "foo.bar.deprecated", "foo.bar.even.more.deprecated"}
+        });
+        assertNotNull(newConfig);
+        assertEquals("derp", newConfig.get("foo.bar"));
+        assertNull(newConfig.get("foo.bar.deprecated"));
+        assertNull(newConfig.get("foo.bar.even.more.deprecated"));
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/utils/ExitTest.java b/clients/src/test/java/org/apache/kafka/common/utils/ExitTest.java
new file mode 100644
index 0000000000000..1377c99b54178
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/utils/ExitTest.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.junit.Test;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+
+import static org.junit.Assert.assertEquals;
+
+public class ExitTest {
+    @Test
+    public void shouldHaltImmediately() {
+        List list = new ArrayList<>();
+        Exit.setHaltProcedure((statusCode, message) -> {
+            list.add(statusCode);
+            list.add(message);
+        });
+        try {
+            int statusCode = 0;
+            String message = "mesaage";
+            Exit.halt(statusCode);
+            Exit.halt(statusCode, message);
+            assertEquals(Arrays.asList(statusCode, null, statusCode, message), list);
+        } finally {
+            Exit.resetHaltProcedure();
+        }
+    }
+
+    @Test
+    public void shouldExitImmediately() {
+        List list = new ArrayList<>();
+        Exit.setExitProcedure((statusCode, message) -> {
+            list.add(statusCode);
+            list.add(message);
+        });
+        try {
+            int statusCode = 0;
+            String message = "mesaage";
+            Exit.exit(statusCode);
+            Exit.exit(statusCode, message);
+            assertEquals(Arrays.asList(statusCode, null, statusCode, message), list);
+        } finally {
+            Exit.resetExitProcedure();
+        }
+    }
+
+    @Test
+    public void shouldAddShutdownHookImmediately() {
+        List list = new ArrayList<>();
+        Exit.setShutdownHookAdder((name, runnable) -> {
+            list.add(name);
+            list.add(runnable);
+        });
+        try {
+            Runnable runnable = () -> { };
+            String name = "name";
+            Exit.addShutdownHook(name, runnable);
+            assertEquals(Arrays.asList(name, runnable), list);
+        } finally {
+            Exit.resetShutdownHookAdder();
+        }
+    }
+
+    @Test
+    public void shouldNotInvokeShutdownHookImmediately() {
+        List list = new ArrayList<>();
+        Runnable runnable = () -> list.add(this);
+        Exit.addShutdownHook("message", runnable);
+        assertEquals(0, list.size());
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/utils/ExponentialBackoffTest.java b/clients/src/test/java/org/apache/kafka/common/utils/ExponentialBackoffTest.java
new file mode 100644
index 0000000000000..7ac6680ebb146
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/utils/ExponentialBackoffTest.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.utils;
+
+import org.junit.Test;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+public class ExponentialBackoffTest {
+    @Test
+    public void testExponentialBackoff() {
+        long scaleFactor = 100;
+        int ratio = 2;
+        long backoffMax = 2000;
+        double jitter = 0.2;
+        ExponentialBackoff exponentialBackoff = new ExponentialBackoff(
+                scaleFactor, ratio, backoffMax, jitter
+        );
+
+        for (int i = 0; i <= 100; i++) {
+            for (int attempts = 0; attempts <= 10; attempts++) {
+                if (attempts <= 4) {
+                    assertEquals(scaleFactor * Math.pow(ratio, attempts),
+                            exponentialBackoff.backoff(attempts),
+                            scaleFactor * Math.pow(ratio, attempts) * jitter);
+                } else {
+                    assertTrue(exponentialBackoff.backoff(attempts) <= backoffMax * (1 + jitter));
+                }
+            }
+        }
+    }
+}
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..d4b41519d79d2
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/utils/FixedOrderMapTest.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.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));
+    }
+}
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..152d5be7f407b
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/utils/ImplicitLinkedHashCollectionTest.java
@@ -0,0 +1,621 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.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 key;
+        private final int val;
+
+        TestElement(int key) {
+            this.key = key;
+            this.val = 0;
+        }
+
+        TestElement(int key, int val) {
+            this.key = key;
+            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 elementKeysAreEqual(Object o) {
+            if (this == o) return true;
+            if ((o == null) || (o.getClass() != TestElement.class)) return false;
+            TestElement that = (TestElement) o;
+            return key == that.key;
+        }
+
+        @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 key == that.key && val == that.val;
+        }
+
+        @Override
+        public String toString() {
+            return "TestElement(key=" + key + ", val=" + val + ")";
+        }
+
+        @Override
+        public int hashCode() {
+            long hashCode = 2654435761L * key;
+            return (int) (hashCode >> 32);
+        }
+    }
+
+    @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.key);
+            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.key);
+            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 key = 3;
+        for (TestElement e : coll) {
+            assertEquals(key, e.key);
+            ++key;
+        }
+    }
+
+    @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).key);
+        assertEquals(2, list.get(1).key);
+        assertEquals(3, list.get(2).key);
+        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).key);
+        assertEquals(1, list.size());
+
+        // Addition to collection is reflected in list
+        coll.add(new TestElement(4));
+        assertEquals(3, list.get(0).key);
+        assertEquals(4, list.get(1).key);
+        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().key);
+        assertTrue(iter.hasNext());
+        assertTrue(iter.hasPrevious());
+        assertEquals(1, iter.nextIndex());
+        assertEquals(0, iter.previousIndex());
+
+        assertEquals(2, iter.next().key);
+        assertTrue(iter.hasNext());
+        assertTrue(iter.hasPrevious());
+        assertEquals(2, iter.nextIndex());
+        assertEquals(1, iter.previousIndex());
+
+        assertEquals(3, iter.next().key);
+        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().key);
+        assertTrue(iter.hasNext());
+        assertTrue(iter.hasPrevious());
+        assertEquals(2, iter.nextIndex());
+        assertEquals(1, iter.previousIndex());
+
+        assertEquals(2, iter.previous().key);
+        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().key);
+        assertTrue(iter.hasNext());
+        assertTrue(iter.hasPrevious());
+        assertEquals(2, iter.nextIndex());
+        assertEquals(1, iter.previousIndex());
+
+        assertEquals(2, iter.previous().key);
+        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().key);
+        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().key);
+        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().key);
+        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().key);
+        assertEquals(5, iter.next().key);
+        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().key);
+        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);
+            expectTraversal(coll.iterator(), existing.iterator());
+        }
+    }
+
+    @Test
+    public void testInsertingTheSameObjectMultipleTimes() {
+        ImplicitLinkedHashCollection coll = new ImplicitLinkedHashCollection<>();
+        TestElement element = new TestElement(123);
+        assertTrue(coll.add(element));
+        assertFalse(coll.add(element));
+        assertFalse(coll.add(element));
+        assertTrue(coll.remove(element));
+        assertFalse(coll.remove(element));
+        assertTrue(coll.add(element));
+        assertFalse(coll.add(element));
+    }
+
+    @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) {
+        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));
+    }
+
+    @Test
+    public void testSameKeysDifferentValues() {
+        ImplicitLinkedHashCollection coll = new ImplicitLinkedHashCollection<>();
+        assertTrue(coll.add(new TestElement(1, 1)));
+        assertFalse(coll.add(new TestElement(1, 2)));
+        TestElement element2 = new TestElement(1, 2);
+        TestElement element1 = coll.find(element2);
+        assertFalse(element2.equals(element1));
+        assertTrue(element2.elementKeysAreEqual(element1));
+    }
+
+    @Test
+    public void testMoveToEnd() {
+        ImplicitLinkedHashCollection coll = new ImplicitLinkedHashCollection<>();
+        TestElement e1 = new TestElement(1, 1);
+        TestElement e2 = new TestElement(2, 2);
+        TestElement e3 = new TestElement(3, 3);
+        assertTrue(coll.add(e1));
+        assertTrue(coll.add(e2));
+        assertTrue(coll.add(e3));
+        coll.moveToEnd(e1);
+        expectTraversal(coll.iterator(), 2, 3, 1);
+        Assert.assertThrows(RuntimeException.class, () ->
+            coll.moveToEnd(new TestElement(4, 4)));
+    }
+
+    @Test
+    public void testRemovals() {
+        ImplicitLinkedHashCollection coll = new ImplicitLinkedHashCollection<>();
+        List elements  = new ArrayList<>();
+        for (int i = 0; i < 100; i++) {
+            TestElement element  = new TestElement(i, i);
+            elements.add(element);
+            coll.add(element);
+        }
+        assertEquals(100, coll.size());
+        Iterator iter = coll.iterator();
+        for (int i = 0; i < 50; i++) {
+            iter.next();
+            iter.remove();
+        }
+        assertEquals(50, coll.size());
+        for (int i = 50; i < 100; i++) {
+            assertEquals(new TestElement(i, i), coll.find(elements.get(i)));
+        }
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/utils/ImplicitLinkedHashMultiCollectionTest.java b/clients/src/test/java/org/apache/kafka/common/utils/ImplicitLinkedHashMultiCollectionTest.java
new file mode 100644
index 0000000000000..ad87b55c493c4
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/utils/ImplicitLinkedHashMultiCollectionTest.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.common.utils;
+
+import org.apache.kafka.common.utils.ImplicitLinkedHashCollectionTest.TestElement;
+import org.junit.Assert;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.Timeout;
+
+import java.util.Iterator;
+import java.util.LinkedList;
+import java.util.Random;
+
+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 ImplicitLinkedHashMultiCollection.
+ */
+public class ImplicitLinkedHashMultiCollectionTest {
+    @Rule
+    final public Timeout globalTimeout = Timeout.millis(120000);
+
+    @Test
+    public void testNullForbidden() {
+        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() {
+        ImplicitLinkedHashMultiCollection multiSet = new ImplicitLinkedHashMultiCollection<>(100);
+        TestElement e1 = new TestElement(1);
+        TestElement e2 = new TestElement(1);
+        TestElement e3 = new TestElement(2);
+        multiSet.mustAdd(e1);
+        multiSet.mustAdd(e2);
+        multiSet.mustAdd(e3);
+        assertFalse(multiSet.add(e3));
+        assertEquals(3, multiSet.size());
+        expectExactTraversal(multiSet.findAll(e1).iterator(), e1, e2);
+        expectExactTraversal(multiSet.findAll(e3).iterator(), e3);
+        multiSet.remove(e2);
+        expectExactTraversal(multiSet.findAll(e1).iterator(), e1);
+        assertTrue(multiSet.contains(e2));
+    }
+
+    @Test
+    public void testTraversal() {
+        ImplicitLinkedHashMultiCollection multiSet = new ImplicitLinkedHashMultiCollection<>();
+        expectExactTraversal(multiSet.iterator());
+        TestElement e1 = new TestElement(1);
+        TestElement e2 = new TestElement(1);
+        TestElement e3 = new TestElement(2);
+        assertTrue(multiSet.add(e1));
+        assertTrue(multiSet.add(e2));
+        assertTrue(multiSet.add(e3));
+        expectExactTraversal(multiSet.iterator(), e1, e2, e3);
+        assertTrue(multiSet.remove(e2));
+        expectExactTraversal(multiSet.iterator(), e1, e3);
+        assertTrue(multiSet.remove(e1));
+        expectExactTraversal(multiSet.iterator(), e3);
+    }
+
+    static void expectExactTraversal(Iterator iterator, TestElement... sequence) {
+        int i = 0;
+        while (iterator.hasNext()) {
+            TestElement element = iterator.next();
+            assertTrue("Iterator yieled " + (i + 1) + " elements, but only " +
+                sequence.length + " were expected.", i < sequence.length);
+            if (sequence[i] != element) {
+                fail("Iterator value number " + (i + 1) + " was incorrect.");
+            }
+            i = i + 1;
+        }
+        assertTrue("Iterator yieled " + (i + 1) + " elements, but " +
+            sequence.length + " were expected.", i == sequence.length);
+    }
+
+    @Test
+    public void testEnlargement() {
+        ImplicitLinkedHashMultiCollection multiSet = new ImplicitLinkedHashMultiCollection<>(5);
+        assertEquals(11, multiSet.numSlots());
+        TestElement[] testElements = {
+            new TestElement(100),
+            new TestElement(101),
+            new TestElement(102),
+            new TestElement(100),
+            new TestElement(101),
+            new TestElement(105)
+        };
+        for (int i = 0; i < testElements.length; i++) {
+            assertTrue(multiSet.add(testElements[i]));
+        }
+        for (int i = 0; i < testElements.length; i++) {
+            assertFalse(multiSet.add(testElements[i]));
+        }
+        assertEquals(23, multiSet.numSlots());
+        assertEquals(testElements.length, multiSet.size());
+        expectExactTraversal(multiSet.iterator(), testElements);
+        multiSet.remove(testElements[1]);
+        assertEquals(23, multiSet.numSlots());
+        assertEquals(5, multiSet.size());
+        expectExactTraversal(multiSet.iterator(),
+            testElements[0], testElements[2], testElements[3], testElements[4], testElements[5]);
+    }
+
+    @Test
+    public void testManyInsertsAndDeletes() {
+        Random random = new Random(123);
+        LinkedList existing = new LinkedList<>();
+        ImplicitLinkedHashMultiCollection multiSet = new ImplicitLinkedHashMultiCollection<>();
+        for (int i = 0; i < 100; i++) {
+            for (int j = 0; j < 4; j++) {
+                TestElement testElement = new TestElement(random.nextInt());
+                multiSet.mustAdd(testElement);
+                existing.add(testElement);
+            }
+            int elementToRemove = random.nextInt(multiSet.size());
+            Iterator iter1 = multiSet.iterator();
+            Iterator iter2 = existing.iterator();
+            for (int j = 0; j <= elementToRemove; j++) {
+                iter1.next();
+                iter2.next();
+            }
+            iter1.remove();
+            iter2.remove();
+            expectTraversal(multiSet.iterator(), existing.iterator());
+        }
+    }
+
+    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());
+            TestElement expected = expectedIter.next();
+            assertTrue("Iterator value number " + (i + 1) + " was incorrect.",
+                expected == element);
+            i = i + 1;
+        }
+        Assert.assertFalse("Iterator yieled " + i + " elements, but at least " +
+            (i + 1) + " were expected.", expectedIter.hasNext());
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/utils/JavaTest.java b/clients/src/test/java/org/apache/kafka/common/utils/JavaTest.java
index 4810f927c5aa5..0ffe06e6a78e2 100644
--- a/clients/src/test/java/org/apache/kafka/common/utils/JavaTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/utils/JavaTest.java
@@ -60,37 +60,30 @@ public void testJavaVersion() {
         assertEquals(9, v.majorVersion);
         assertEquals(0, v.minorVersion);
         assertTrue(v.isJava9Compatible());
-        assertTrue(v.isJava8Compatible());
 
         v = Java.parseVersion("9.0.1");
         assertEquals(9, v.majorVersion);
         assertEquals(0, v.minorVersion);
         assertTrue(v.isJava9Compatible());
-        assertTrue(v.isJava8Compatible());
 
         v = Java.parseVersion("9.0.0.15"); // Azul Zulu
         assertEquals(9, v.majorVersion);
         assertEquals(0, v.minorVersion);
         assertTrue(v.isJava9Compatible());
-        assertTrue(v.isJava8Compatible());
 
         v = Java.parseVersion("9.1");
         assertEquals(9, v.majorVersion);
         assertEquals(1, v.minorVersion);
         assertTrue(v.isJava9Compatible());
-        assertTrue(v.isJava8Compatible());
 
         v = Java.parseVersion("1.8.0_152");
         assertEquals(1, v.majorVersion);
         assertEquals(8, v.minorVersion);
         assertFalse(v.isJava9Compatible());
-        assertTrue(v.isJava8Compatible());
 
         v = Java.parseVersion("1.7.0_80");
         assertEquals(1, v.majorVersion);
         assertEquals(7, v.minorVersion);
         assertFalse(v.isJava9Compatible());
-        assertFalse(v.isJava8Compatible());
-
     }
 }
diff --git a/clients/src/test/java/org/apache/kafka/common/utils/LoggingSignalHandlerTest.java b/clients/src/test/java/org/apache/kafka/common/utils/LoggingSignalHandlerTest.java
new file mode 100644
index 0000000000000..468e10e84f898
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/utils/LoggingSignalHandlerTest.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.utils;
+
+import org.junit.Test;
+
+public class LoggingSignalHandlerTest {
+
+    @Test
+    public void testRegister() throws ReflectiveOperationException {
+        new LoggingSignalHandler().register();
+    }
+
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/utils/MappedByteBuffersTest.java b/clients/src/test/java/org/apache/kafka/common/utils/MappedByteBuffersTest.java
deleted file mode 100644
index 38fe9ddd4f70c..0000000000000
--- a/clients/src/test/java/org/apache/kafka/common/utils/MappedByteBuffersTest.java
+++ /dev/null
@@ -1,41 +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.apache.kafka.test.TestUtils;
-import org.junit.Test;
-
-import java.io.File;
-import java.nio.MappedByteBuffer;
-import java.nio.channels.FileChannel;
-
-public class MappedByteBuffersTest {
-
-    /**
-     * Checks that unmap doesn't throw exceptions.
-     */
-    @Test
-    public void testUnmap() throws Exception {
-        File file = TestUtils.tempFile();
-        try (FileChannel channel = FileChannel.open(file.toPath())) {
-            MappedByteBuffer map = channel.map(FileChannel.MapMode.READ_ONLY, 0, 0);
-            MappedByteBuffers.unmap(file.getAbsolutePath(), map);
-        }
-    }
-
-}
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/MockScheduler.java b/clients/src/test/java/org/apache/kafka/common/utils/MockScheduler.java
index ba5e1ed0c59df..78a9060c4046c 100644
--- a/clients/src/test/java/org/apache/kafka/common/utils/MockScheduler.java
+++ b/clients/src/test/java/org/apache/kafka/common/utils/MockScheduler.java
@@ -29,7 +29,7 @@
 import java.util.concurrent.Future;
 import java.util.concurrent.ScheduledExecutorService;
 
-public class MockScheduler implements Scheduler, MockTime.MockTimeListener {
+public class MockScheduler implements Scheduler, MockTime.Listener {
     private static final Logger log = LoggerFactory.getLogger(MockScheduler.class);
 
     /**
@@ -53,7 +53,7 @@ public Time time() {
     }
 
     @Override
-    public synchronized void tick() {
+    public synchronized void onTimeUpdated() {
         long timeMs = time.milliseconds();
         while (true) {
             Map.Entry>> entry = waiters.firstEntry();
@@ -87,12 +87,12 @@ public  Future schedule(final ScheduledExecutorService executor,
                                   final Callable callable, long delayMs) {
         final KafkaFutureImpl future = new KafkaFutureImpl<>();
         KafkaFutureImpl waiter = new KafkaFutureImpl<>();
-        waiter.thenApply(new KafkaFuture.Function() {
+        waiter.thenApply(new KafkaFuture.BaseFunction() {
             @Override
             public Void apply(final Long now) {
                 executor.submit(new Callable() {
                     @Override
-                    public Void call() throws Exception {
+                    public Void call() {
                         // Note: it is possible that we'll execute Callable#call right after
                         // the future is cancelled.  This is a valid sequence of events
                         // that the author of the Callable needs to be able to handle.
diff --git a/clients/src/test/java/org/apache/kafka/common/utils/MockTime.java b/clients/src/test/java/org/apache/kafka/common/utils/MockTime.java
index 011eba2111b64..ccf3eec196e61 100644
--- a/clients/src/test/java/org/apache/kafka/common/utils/MockTime.java
+++ b/clients/src/test/java/org/apache/kafka/common/utils/MockTime.java
@@ -16,23 +16,26 @@
  */
 package org.apache.kafka.common.utils;
 
+import org.apache.kafka.common.errors.TimeoutException;
+
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Supplier;
 
 /**
  * A clock that you can manually advance by calling sleep
  */
 public class MockTime implements Time {
 
-    interface MockTimeListener {
-        void tick();
+    public interface Listener {
+        void onTimeUpdated();
     }
 
     /**
      * Listeners which are waiting for time changes.
      */
-    private final CopyOnWriteArrayList listeners = new CopyOnWriteArrayList<>();
+    private final CopyOnWriteArrayList listeners = new CopyOnWriteArrayList<>();
 
     private final long autoTickMs;
 
@@ -55,7 +58,7 @@ public MockTime(long autoTickMs, long currentTimeMs, long currentHighResTimeNs)
         this.autoTickMs = autoTickMs;
     }
 
-    public void addListener(MockTimeListener listener) {
+    public void addListener(Listener listener) {
         listeners.add(listener);
     }
 
@@ -71,11 +74,6 @@ public long nanoseconds() {
         return highResTimeNs.get();
     }
 
-    @Override
-    public long hiResClockMs() {
-        return TimeUnit.NANOSECONDS.toMillis(nanoseconds());
-    }
-
     private void maybeSleep(long ms) {
         if (ms != 0)
             sleep(ms);
@@ -88,6 +86,27 @@ public void sleep(long ms) {
         tick();
     }
 
+    @Override
+    public void waitObject(Object obj, Supplier condition, long deadlineMs) throws InterruptedException {
+        Listener listener = () -> {
+            synchronized (obj) {
+                obj.notify();
+            }
+        };
+        listeners.add(listener);
+        try {
+            synchronized (obj) {
+                while (milliseconds() < deadlineMs && !condition.get()) {
+                    obj.wait();
+                }
+                if (!condition.get())
+                    throw new TimeoutException("Condition not satisfied before deadline");
+            }
+        } finally {
+            listeners.remove(listener);
+        }
+    }
+
     public void setCurrentTimeMs(long newMs) {
         long oldMs = timeMs.getAndSet(newMs);
 
@@ -100,8 +119,8 @@ public void setCurrentTimeMs(long newMs) {
     }
 
     private void tick() {
-        for (MockTimeListener listener : listeners) {
-            listener.tick();
+        for (Listener listener : listeners) {
+            listener.onTimeUpdated();
         }
     }
 }
diff --git a/clients/src/test/java/org/apache/kafka/common/utils/MockTimeTest.java b/clients/src/test/java/org/apache/kafka/common/utils/MockTimeTest.java
index 58bcb19361ed6..d8101ac8fc185 100644
--- a/clients/src/test/java/org/apache/kafka/common/utils/MockTimeTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/utils/MockTimeTest.java
@@ -16,35 +16,38 @@
  */
 package org.apache.kafka.common.utils;
 
-import org.junit.Assert;
 import org.junit.Rule;
-import org.junit.rules.Timeout;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 import org.junit.Test;
+import org.junit.rules.Timeout;
+
+import static org.junit.Assert.assertEquals;
 
-public class MockTimeTest {
-    private static final Logger log = LoggerFactory.getLogger(MockTimeTest.class);
+public class MockTimeTest extends TimeTest {
 
     @Rule
     final public Timeout globalTimeout = Timeout.millis(120000);
 
     @Test
-    public void testAdvanceClock() throws Exception {
+    public void testAdvanceClock() {
         MockTime time = new MockTime(0, 100, 200);
-        Assert.assertEquals(100, time.milliseconds());
-        Assert.assertEquals(200, time.nanoseconds());
+        assertEquals(100, time.milliseconds());
+        assertEquals(200, time.nanoseconds());
         time.sleep(1);
-        Assert.assertEquals(101, time.milliseconds());
-        Assert.assertEquals(1000200, time.nanoseconds());
+        assertEquals(101, time.milliseconds());
+        assertEquals(1000200, time.nanoseconds());
     }
 
     @Test
-    public void testAutoTickMs() throws Exception {
+    public void testAutoTickMs() {
         MockTime time = new MockTime(1, 100, 200);
-        Assert.assertEquals(101, time.milliseconds());
-        Assert.assertEquals(2000200, time.nanoseconds());
-        Assert.assertEquals(103, time.milliseconds());
-        Assert.assertEquals(104, time.milliseconds());
+        assertEquals(101, time.milliseconds());
+        assertEquals(2000200, time.nanoseconds());
+        assertEquals(103, time.milliseconds());
+        assertEquals(104, time.milliseconds());
+    }
+
+    @Override
+    protected Time createTime() {
+        return new MockTime();
     }
 }
diff --git a/clients/src/test/java/org/apache/kafka/common/utils/SanitizerTest.java b/clients/src/test/java/org/apache/kafka/common/utils/SanitizerTest.java
index 59ac6c0c0d4f8..bba19a87576d4 100644
--- a/clients/src/test/java/org/apache/kafka/common/utils/SanitizerTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/utils/SanitizerTest.java
@@ -20,7 +20,6 @@
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
-import java.io.UnsupportedEncodingException;
 import java.lang.management.ManagementFactory;
 
 import javax.management.MBeanException;
@@ -34,7 +33,7 @@
 public class SanitizerTest {
 
     @Test
-    public void testSanitize() throws UnsupportedEncodingException {
+    public void testSanitize() {
         String principal = "CN=Some characters !@#$%&*()_-+=';:,/~";
         String sanitizedPrincipal = Sanitizer.sanitize(principal);
         assertTrue(sanitizedPrincipal.replace('%', '_').matches("[a-zA-Z0-9\\._\\-]+"));
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/ShellTest.java b/clients/src/test/java/org/apache/kafka/common/utils/ShellTest.java
index d77f4d0284bff..7bd96f00a4bd7 100644
--- a/clients/src/test/java/org/apache/kafka/common/utils/ShellTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/utils/ShellTest.java
@@ -49,7 +49,7 @@ public void testHeadDevZero() throws Exception {
     private final static String NONEXISTENT_PATH = "/dev/a/path/that/does/not/exist/in/the/filesystem";
 
     @Test
-    public void testAttemptToRunNonExistentProgram() throws Exception {
+    public void testAttemptToRunNonExistentProgram() {
         assumeTrue(!OperatingSystem.IS_WINDOWS);
         try {
             Shell.execCommand(NONEXISTENT_PATH);
diff --git a/clients/src/test/java/org/apache/kafka/common/utils/SystemTimeTest.java b/clients/src/test/java/org/apache/kafka/common/utils/SystemTimeTest.java
new file mode 100644
index 0000000000000..edc53d2293e15
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/utils/SystemTimeTest.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.common.utils;
+
+public class SystemTimeTest extends TimeTest {
+
+    @Override
+    protected Time createTime() {
+        return Time.SYSTEM;
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/utils/ThreadUtilsTest.java b/clients/src/test/java/org/apache/kafka/common/utils/ThreadUtilsTest.java
new file mode 100644
index 0000000000000..c49447c17d20d
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/utils/ThreadUtilsTest.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.common.utils;
+
+import org.junit.Test;
+
+import java.util.concurrent.ThreadFactory;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+public class ThreadUtilsTest {
+
+    private static final Runnable EMPTY_RUNNABLE = () -> {
+    };
+    private static final String THREAD_NAME = "ThreadName";
+    private static final String THREAD_NAME_WITH_NUMBER = THREAD_NAME + "%d";
+
+
+    @Test
+    public void testThreadNameWithoutNumberNoDemon() {
+        assertEquals(THREAD_NAME, ThreadUtils.createThreadFactory(THREAD_NAME, false).
+                newThread(EMPTY_RUNNABLE).getName());
+    }
+
+    @Test
+    public void testThreadNameWithoutNumberDemon() {
+        Thread daemonThread = ThreadUtils.createThreadFactory(THREAD_NAME, true).newThread(EMPTY_RUNNABLE);
+        try {
+            assertEquals(THREAD_NAME, daemonThread.getName());
+            assertTrue(daemonThread.isDaemon());
+        } finally {
+            try {
+                daemonThread.join();
+            } catch (InterruptedException e) {
+                // can be ignored
+            }
+        }
+    }
+
+    @Test
+    public void testThreadNameWithNumberNoDemon() {
+        ThreadFactory localThreadFactory = ThreadUtils.createThreadFactory(THREAD_NAME_WITH_NUMBER, false);
+        assertEquals(THREAD_NAME + "1", localThreadFactory.newThread(EMPTY_RUNNABLE).getName());
+        assertEquals(THREAD_NAME + "2", localThreadFactory.newThread(EMPTY_RUNNABLE).getName());
+    }
+
+    @Test
+    public void testThreadNameWithNumberDemon() {
+        ThreadFactory localThreadFactory = ThreadUtils.createThreadFactory(THREAD_NAME_WITH_NUMBER, true);
+        Thread daemonThread1 = localThreadFactory.newThread(EMPTY_RUNNABLE);
+        Thread daemonThread2 = localThreadFactory.newThread(EMPTY_RUNNABLE);
+
+        try {
+            assertEquals(THREAD_NAME + "1", daemonThread1.getName());
+            assertTrue(daemonThread1.isDaemon());
+        } finally {
+            try {
+                daemonThread1.join();
+            } catch (InterruptedException e) {
+                // can be ignored
+            }
+        }
+        try {
+            assertEquals(THREAD_NAME + "2", daemonThread2.getName());
+            assertTrue(daemonThread2.isDaemon());
+        } finally {
+            try {
+                daemonThread2.join();
+            } catch (InterruptedException e) {
+                // can be ignored
+            }
+        }
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/utils/TimeTest.java b/clients/src/test/java/org/apache/kafka/common/utils/TimeTest.java
new file mode 100644
index 0000000000000..98878e2617bfb
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/utils/TimeTest.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.utils;
+
+import org.apache.kafka.common.errors.TimeoutException;
+import org.junit.Test;
+
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+
+public abstract class TimeTest {
+
+    protected abstract Time createTime();
+
+    @Test
+    public void testWaitObjectTimeout() throws InterruptedException {
+        Object obj = new Object();
+        Time time = createTime();
+        long timeoutMs = 100;
+        long deadlineMs = time.milliseconds() + timeoutMs;
+        AtomicReference caughtException = new AtomicReference<>();
+        Thread t = new Thread(() -> {
+            try {
+                time.waitObject(obj, () -> false, deadlineMs);
+            } catch (Exception e) {
+                caughtException.set(e);
+            }
+        });
+
+        t.start();
+        time.sleep(timeoutMs);
+        t.join();
+
+        assertEquals(TimeoutException.class, caughtException.get().getClass());
+    }
+
+    @Test
+    public void testWaitObjectConditionSatisfied() throws InterruptedException {
+        Object obj = new Object();
+        Time time = createTime();
+        long timeoutMs = 1000000000;
+        long deadlineMs = time.milliseconds() + timeoutMs;
+        AtomicBoolean condition = new AtomicBoolean(false);
+        AtomicReference caughtException = new AtomicReference<>();
+        Thread t = new Thread(() -> {
+            try {
+                time.waitObject(obj, condition::get, deadlineMs);
+            } catch (Exception e) {
+                caughtException.set(e);
+            }
+        });
+
+        t.start();
+
+        synchronized (obj) {
+            condition.set(true);
+            obj.notify();
+        }
+
+        t.join();
+
+        assertTrue(time.milliseconds() < deadlineMs);
+        assertNull(caughtException.get());
+    }
+}
diff --git a/clients/src/test/java/org/apache/kafka/common/utils/TimerTest.java b/clients/src/test/java/org/apache/kafka/common/utils/TimerTest.java
new file mode 100644
index 0000000000000..52a1ac6367e2d
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/common/utils/TimerTest.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.common.utils;
+
+import org.junit.Test;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+public class TimerTest {
+
+    private final MockTime time = new MockTime();
+
+    @Test
+    public void testTimerUpdate() {
+        Timer timer = time.timer(500);
+        assertEquals(500, timer.timeoutMs());
+        assertEquals(500, timer.remainingMs());
+        assertEquals(0, timer.elapsedMs());
+
+        time.sleep(100);
+        timer.update();
+
+        assertEquals(500, timer.timeoutMs());
+        assertEquals(400, timer.remainingMs());
+        assertEquals(100, timer.elapsedMs());
+
+        time.sleep(400);
+        timer.update(time.milliseconds());
+
+        assertEquals(500, timer.timeoutMs());
+        assertEquals(0, timer.remainingMs());
+        assertEquals(500, timer.elapsedMs());
+        assertTrue(timer.isExpired());
+
+        // Going over the expiration is fine and the elapsed time can exceed
+        // the initial timeout. However, remaining time should be stuck at 0.
+        time.sleep(200);
+        timer.update(time.milliseconds());
+        assertTrue(timer.isExpired());
+        assertEquals(500, timer.timeoutMs());
+        assertEquals(0, timer.remainingMs());
+        assertEquals(700, timer.elapsedMs());
+    }
+
+    @Test
+    public void testTimerUpdateAndReset() {
+        Timer timer = time.timer(500);
+        timer.sleep(200);
+        assertEquals(500, timer.timeoutMs());
+        assertEquals(300, timer.remainingMs());
+        assertEquals(200, timer.elapsedMs());
+
+        timer.updateAndReset(400);
+        assertEquals(400, timer.timeoutMs());
+        assertEquals(400, timer.remainingMs());
+        assertEquals(0, timer.elapsedMs());
+
+        timer.sleep(400);
+        assertTrue(timer.isExpired());
+
+        timer.updateAndReset(200);
+        assertEquals(200, timer.timeoutMs());
+        assertEquals(200, timer.remainingMs());
+        assertEquals(0, timer.elapsedMs());
+        assertFalse(timer.isExpired());
+    }
+
+    @Test
+    public void testTimerResetUsesCurrentTime() {
+        Timer timer = time.timer(500);
+        timer.sleep(200);
+        assertEquals(300, timer.remainingMs());
+        assertEquals(200, timer.elapsedMs());
+
+        time.sleep(300);
+        timer.reset(500);
+        assertEquals(500, timer.remainingMs());
+
+        timer.update();
+        assertEquals(200, timer.remainingMs());
+    }
+
+    @Test
+    public void testTimerResetDeadlineUsesCurrentTime() {
+        Timer timer = time.timer(500);
+        timer.sleep(200);
+        assertEquals(300, timer.remainingMs());
+        assertEquals(200, timer.elapsedMs());
+
+        timer.sleep(100);
+        timer.resetDeadline(time.milliseconds() + 200);
+        assertEquals(200, timer.timeoutMs());
+        assertEquals(200, timer.remainingMs());
+
+        timer.sleep(100);
+        assertEquals(200, timer.timeoutMs());
+        assertEquals(100, timer.remainingMs());
+    }
+
+    @Test
+    public void testTimeoutOverflow() {
+        Timer timer = time.timer(Long.MAX_VALUE);
+        assertEquals(Long.MAX_VALUE - timer.currentTimeMs(), timer.remainingMs());
+        assertEquals(0, timer.elapsedMs());
+    }
+
+    @Test
+    public void testNonMonotonicUpdate() {
+        Timer timer = time.timer(100);
+        long currentTimeMs = timer.currentTimeMs();
+
+        timer.update(currentTimeMs - 1);
+        assertEquals(currentTimeMs, timer.currentTimeMs());
+
+        assertEquals(100, timer.remainingMs());
+        assertEquals(0, timer.elapsedMs());
+    }
+
+    @Test
+    public void testTimerSleep() {
+        Timer timer = time.timer(500);
+        long currentTimeMs = timer.currentTimeMs();
+
+        timer.sleep(200);
+        assertEquals(time.milliseconds(), timer.currentTimeMs());
+        assertEquals(currentTimeMs + 200, timer.currentTimeMs());
+
+        timer.sleep(1000);
+        assertEquals(time.milliseconds(), timer.currentTimeMs());
+        assertEquals(currentTimeMs + 500, timer.currentTimeMs());
+        assertTrue(timer.isExpired());
+    }
+
+}
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 3feeff25e791b..320c10648eba4 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
@@ -16,37 +16,92 @@
  */
 package org.apache.kafka.common.utils;
 
+import org.apache.kafka.common.config.ConfigException;
 import org.apache.kafka.test.TestUtils;
-import org.easymock.EasyMock;
-import org.easymock.IAnswer;
 import org.junit.Test;
+import org.junit.function.ThrowingRunnable;
+import org.mockito.stubbing.OngoingStubbing;
 
 import java.io.Closeable;
 import java.io.DataOutputStream;
 import java.io.EOFException;
 import java.io.File;
 import java.io.IOException;
+import java.nio.BufferUnderflowException;
 import java.nio.ByteBuffer;
 import java.nio.channels.FileChannel;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.nio.file.StandardOpenOption;
+import java.time.LocalDateTime;
+import java.time.format.DateTimeFormatter;
+import java.time.format.DateTimeFormatterBuilder;
+import java.time.temporal.ChronoField;
+import java.time.temporal.ChronoUnit;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Properties;
 import java.util.Random;
-
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+import java.util.stream.Stream;
+
+import static java.util.Arrays.asList;
+import static java.util.Collections.emptySet;
+import static org.apache.kafka.common.utils.Utils.diff;
 import static org.apache.kafka.common.utils.Utils.formatAddress;
 import static org.apache.kafka.common.utils.Utils.formatBytes;
 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.intersection;
+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.union;
+import static org.apache.kafka.common.utils.Utils.validHostPattern;
+import static org.hamcrest.CoreMatchers.containsString;
+import static org.hamcrest.CoreMatchers.equalTo;
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.Assert.assertArrayEquals;
 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;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.Mockito.atLeastOnce;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+import java.text.ParseException;
+import java.text.SimpleDateFormat;
+import java.util.Date;
 
 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(c.getValue().intValue(), murmur2(c.getKey()));
+        }
+    }
+
     @Test
     public void testGetHost() {
         assertEquals("127.0.0.1", getHost("127.0.0.1:8000"));
@@ -59,6 +114,16 @@ public void testGetHost() {
         assertEquals("fe80::b1da:69ca:57f7:63d8%3", getHost("PLAINTEXT://[fe80::b1da:69ca:57f7:63d8%3]:5678"));
     }
 
+    @Test
+    public void testHostPattern() {
+        assertTrue(validHostPattern("127.0.0.1"));
+        assertTrue(validHostPattern("mydomain.com"));
+        assertTrue(validHostPattern("MyDomain.com"));
+        assertTrue(validHostPattern("My_Domain.com"));
+        assertTrue(validHostPattern("::1"));
+        assertTrue(validHostPattern("2001:db8:85a3:8d3:1319:8a2e:370"));
+    }
+
     @Test
     public void testGetPort() {
         assertEquals(8000, getPort("127.0.0.1:8000").intValue());
@@ -92,10 +157,10 @@ public void testFormatBytes() {
     @Test
     public void testJoin() {
         assertEquals("", Utils.join(Collections.emptyList(), ","));
-        assertEquals("1", Utils.join(Arrays.asList("1"), ","));
-        assertEquals("1,2,3", Utils.join(Arrays.asList(1, 2, 3), ","));
+        assertEquals("1", Utils.join(asList("1"), ","));
+        assertEquals("1,2,3", Utils.join(asList(1, 2, 3), ","));
     }
-    
+
     @Test
     public void testAbs() {
         assertEquals(0, Utils.abs(Integer.MIN_VALUE));
@@ -163,6 +228,65 @@ public void toArrayDirectByteBuffer() {
         assertEquals(2, buffer.position());
     }
 
+    @Test
+    public void getNullableSizePrefixedArrayExact() {
+        byte[] input = {0, 0, 0, 2, 1, 0};
+        final ByteBuffer buffer = ByteBuffer.wrap(input);
+        final byte[] array = Utils.getNullableSizePrefixedArray(buffer);
+        assertArrayEquals(new byte[] {1, 0}, array);
+        assertEquals(6, buffer.position());
+        assertFalse(buffer.hasRemaining());
+    }
+
+    @Test
+    public void getNullableSizePrefixedArrayExactEmpty() {
+        byte[] input = {0, 0, 0, 0};
+        final ByteBuffer buffer = ByteBuffer.wrap(input);
+        final byte[] array = Utils.getNullableSizePrefixedArray(buffer);
+        assertArrayEquals(new byte[] {}, array);
+        assertEquals(4, buffer.position());
+        assertFalse(buffer.hasRemaining());
+    }
+
+    @Test
+    public void getNullableSizePrefixedArrayRemainder() {
+        byte[] input = {0, 0, 0, 2, 1, 0, 9};
+        final ByteBuffer buffer = ByteBuffer.wrap(input);
+        final byte[] array = Utils.getNullableSizePrefixedArray(buffer);
+        assertArrayEquals(new byte[] {1, 0}, array);
+        assertEquals(6, buffer.position());
+        assertTrue(buffer.hasRemaining());
+    }
+
+    @Test
+    public void getNullableSizePrefixedArrayNull() {
+        // -1
+        byte[] input = {-1, -1, -1, -1};
+        final ByteBuffer buffer = ByteBuffer.wrap(input);
+        final byte[] array = Utils.getNullableSizePrefixedArray(buffer);
+        assertNull(array);
+        assertEquals(4, buffer.position());
+        assertFalse(buffer.hasRemaining());
+    }
+
+    @Test
+    public void getNullableSizePrefixedArrayInvalid() {
+        // -2
+        byte[] input = {-1, -1, -1, -2};
+        final ByteBuffer buffer = ByteBuffer.wrap(input);
+        assertThrows(NegativeArraySizeException.class, () -> Utils.getNullableSizePrefixedArray(buffer));
+    }
+
+    @Test
+    public void getNullableSizePrefixedArrayUnderflow() {
+        // Integer.MAX_VALUE
+        byte[] input = {127, -1, -1, -1};
+        final ByteBuffer buffer = ByteBuffer.wrap(input);
+        // note, we get a buffer underflow exception instead of an OOME, even though the encoded size
+        // would be 2,147,483,647 aka 2.1 GB, probably larger than the available heap
+        assertThrows(BufferUnderflowException.class, () -> Utils.getNullableSizePrefixedArray(buffer));
+    }
+
     @Test
     public void utf8ByteArraySerde() {
         String utf8String = "A\u00ea\u00f1\u00fcC";
@@ -226,6 +350,61 @@ public void testReadBytes() {
         this.subTest(buffer);
     }
 
+    @Test
+    public void testFileAsStringSimpleFile() throws IOException {
+        File tempFile = TestUtils.tempFile();
+        try {
+            String testContent = "Test Content";
+            Files.write(tempFile.toPath(), testContent.getBytes());
+            assertEquals(testContent, Utils.readFileAsString(tempFile.getPath()));
+        } finally {
+            Files.deleteIfExists(tempFile.toPath());
+        }
+    }
+
+    /**
+     * Test to read content of named pipe as string. As reading/writing to a pipe can block,
+     * timeout test after a minute (test finishes within 100 ms normally).
+     */
+    @Test(timeout = 60 * 1000)
+    public void testFileAsStringNamedPipe() throws Exception {
+
+        // Create a temporary name for named pipe
+        Random random = new Random();
+        long n = random.nextLong();
+        n = n == Long.MIN_VALUE ? 0 : Math.abs(n);
+
+        // Use the name to create a FIFO in tmp directory
+        String tmpDir = System.getProperty("java.io.tmpdir");
+        String fifoName = "fifo-" + n + ".tmp";
+        File fifo = new File(tmpDir, fifoName);
+        Thread producerThread = null;
+        try {
+            Process mkFifoCommand = new ProcessBuilder("mkfifo", fifo.getCanonicalPath()).start();
+            mkFifoCommand.waitFor();
+
+            // Send some data to fifo and then read it back, but as FIFO blocks if the consumer isn't present,
+            // we need to send data in a separate thread.
+            final String testFileContent = "This is test";
+            producerThread = new Thread(() -> {
+                try {
+                    Files.write(fifo.toPath(), testFileContent.getBytes());
+                } catch (IOException e) {
+                    fail("Error when producing to fifo : " + e.getMessage());
+                }
+            }, "FIFO-Producer");
+            producerThread.start();
+
+            assertEquals(testFileContent, Utils.readFileAsString(fifo.getCanonicalPath()));
+        } finally {
+            Files.deleteIfExists(fifo.toPath());
+            if (producerThread != null) {
+                producerThread.join(30 * 1000); // Wait for thread to terminate
+                assertFalse(producerThread.isAlive());
+            }
+        }
+    }
+
     @Test
     public void testMin() {
         assertEquals(1, Utils.min(1));
@@ -313,17 +492,15 @@ public void testReadFullyOrFailWithRealFile() throws IOException {
      */
     @Test
     public void testReadFullyOrFailWithPartialFileChannelReads() throws IOException {
-        FileChannel channelMock = EasyMock.createMock(FileChannel.class);
+        FileChannel channelMock = mock(FileChannel.class);
         final int bufferSize = 100;
         ByteBuffer buffer = ByteBuffer.allocate(bufferSize);
-        StringBuilder expectedBufferContent = new StringBuilder();
-        fileChannelMockExpectReadWithRandomBytes(channelMock, expectedBufferContent, bufferSize);
-        EasyMock.replay(channelMock);
+        String expectedBufferContent = fileChannelMockExpectReadWithRandomBytes(channelMock, bufferSize);
         Utils.readFullyOrFail(channelMock, buffer, 0L, "test");
-        assertEquals("The buffer should be populated correctly", expectedBufferContent.toString(),
-                new String(buffer.array()));
+        assertEquals("The buffer should be populated correctly", expectedBufferContent,
+                     new String(buffer.array()));
         assertFalse("The buffer should be filled", buffer.hasRemaining());
-        EasyMock.verify(channelMock);
+        verify(channelMock, atLeastOnce()).read(any(), anyLong());
     }
 
     /**
@@ -332,73 +509,83 @@ public void testReadFullyOrFailWithPartialFileChannelReads() throws IOException
      */
     @Test
     public void testReadFullyWithPartialFileChannelReads() throws IOException {
-        FileChannel channelMock = EasyMock.createMock(FileChannel.class);
+        FileChannel channelMock = mock(FileChannel.class);
         final int bufferSize = 100;
-        StringBuilder expectedBufferContent = new StringBuilder();
-        fileChannelMockExpectReadWithRandomBytes(channelMock, expectedBufferContent, bufferSize);
-        EasyMock.replay(channelMock);
+        String expectedBufferContent = fileChannelMockExpectReadWithRandomBytes(channelMock, bufferSize);
         ByteBuffer buffer = ByteBuffer.allocate(bufferSize);
         Utils.readFully(channelMock, buffer, 0L);
-        assertEquals("The buffer should be populated correctly.", expectedBufferContent.toString(),
-                new String(buffer.array()));
+        assertEquals("The buffer should be populated correctly.", expectedBufferContent,
+                     new String(buffer.array()));
         assertFalse("The buffer should be filled", buffer.hasRemaining());
-        EasyMock.verify(channelMock);
+        verify(channelMock, atLeastOnce()).read(any(), anyLong());
     }
 
     @Test
     public void testReadFullyIfEofIsReached() throws IOException {
-        final FileChannel channelMock = EasyMock.createMock(FileChannel.class);
+        final FileChannel channelMock = mock(FileChannel.class);
         final int bufferSize = 100;
         final String fileChannelContent = "abcdefghkl";
         ByteBuffer buffer = ByteBuffer.allocate(bufferSize);
-        EasyMock.expect(channelMock.size()).andReturn((long) fileChannelContent.length());
-        EasyMock.expect(channelMock.read(EasyMock.anyObject(ByteBuffer.class), EasyMock.anyInt())).andAnswer(new IAnswer() {
-            @Override
-            public Integer answer() throws Throwable {
-                ByteBuffer buffer = (ByteBuffer) EasyMock.getCurrentArguments()[0];
-                buffer.put(fileChannelContent.getBytes());
-                return -1;
-            }
+        when(channelMock.read(any(), anyLong())).then(invocation -> {
+            ByteBuffer bufferArg = invocation.getArgument(0);
+            bufferArg.put(fileChannelContent.getBytes());
+            return -1;
         });
-        EasyMock.replay(channelMock);
         Utils.readFully(channelMock, buffer, 0L);
         assertEquals("abcdefghkl", new String(buffer.array(), 0, buffer.position()));
-        assertEquals(buffer.position(), channelMock.size());
+        assertEquals(fileChannelContent.length(), buffer.position());
         assertTrue(buffer.hasRemaining());
-        EasyMock.verify(channelMock);
+        verify(channelMock, atLeastOnce()).read(any(), anyLong());
+    }
+
+    @Test
+    public void testLoadProps() throws IOException {
+        File tempFile = TestUtils.tempFile();
+        try {
+            String testContent = "a=1\nb=2\n#a comment\n\nc=3\nd=";
+            Files.write(tempFile.toPath(), testContent.getBytes());
+            Properties props = Utils.loadProps(tempFile.getPath());
+            assertEquals(4, props.size());
+            assertEquals("1", props.get("a"));
+            assertEquals("2", props.get("b"));
+            assertEquals("3", props.get("c"));
+            assertEquals("", props.get("d"));
+            Properties restrictedProps = Utils.loadProps(tempFile.getPath(), Arrays.asList("b", "d", "e"));
+            assertEquals(2, restrictedProps.size());
+            assertEquals("2", restrictedProps.get("b"));
+            assertEquals("", restrictedProps.get("d"));
+        } finally {
+            Files.deleteIfExists(tempFile.toPath());
+        }
     }
 
     /**
      * Expectation setter for multiple reads where each one reads random bytes to the buffer.
      *
      * @param channelMock           The mocked FileChannel object
-     * @param expectedBufferContent buffer that will be updated to contain the expected buffer content after each
-     *                              `FileChannel.read` invocation
      * @param bufferSize            The buffer size
+     * @return Expected buffer string
      * @throws IOException          If an I/O error occurs
      */
-    private void fileChannelMockExpectReadWithRandomBytes(final FileChannel channelMock,
-                                                          final StringBuilder expectedBufferContent,
-                                                          final int bufferSize) throws IOException {
+    private String fileChannelMockExpectReadWithRandomBytes(final FileChannel channelMock,
+                                                            final int bufferSize) throws IOException {
         final int step = 20;
         final Random random = new Random();
         int remainingBytes = bufferSize;
+        OngoingStubbing when = when(channelMock.read(any(), anyLong()));
+        StringBuilder expectedBufferContent = new StringBuilder();
         while (remainingBytes > 0) {
-            final int mockedBytesRead = remainingBytes < step ? remainingBytes : random.nextInt(step);
-            final StringBuilder sb = new StringBuilder();
-            EasyMock.expect(channelMock.read(EasyMock.anyObject(ByteBuffer.class), EasyMock.anyInt())).andAnswer(new IAnswer() {
-                @Override
-                public Integer answer() throws Throwable {
-                    ByteBuffer buffer = (ByteBuffer) EasyMock.getCurrentArguments()[0];
-                    for (int i = 0; i < mockedBytesRead; i++)
-                        sb.append("a");
-                    buffer.put(sb.toString().getBytes());
-                    expectedBufferContent.append(sb);
-                    return mockedBytesRead;
-                }
+            final int bytesRead = remainingBytes < step ? remainingBytes : random.nextInt(step);
+            final String stringRead = IntStream.range(0, bytesRead).mapToObj(i -> "a").collect(Collectors.joining());
+            expectedBufferContent.append(stringRead);
+            when = when.then(invocation -> {
+                ByteBuffer buffer = invocation.getArgument(0);
+                buffer.put(stringRead.getBytes());
+                return bytesRead;
             });
-            remainingBytes -= mockedBytesRead;
+            remainingBytes -= bytesRead;
         }
+        return expectedBufferContent.toString();
     }
 
     private static class TestCloseable implements Closeable {
@@ -414,8 +601,9 @@ private static class TestCloseable implements Closeable {
         @Override
         public void close() throws IOException {
             closed = true;
-            if (closeException != null)
+            if (closeException != null) {
                 throw closeException;
+            }
         }
 
         static TestCloseable[] createCloseables(boolean... exceptionOnClose) {
@@ -461,4 +649,214 @@ public void testRecursiveDelete() throws IOException {
         Utils.delete(tempDir);
         assertFalse(Files.exists(tempDir.toPath()));
     }
+
+    @Test
+    public void testConvertTo32BitField() {
+        Set bytes = mkSet((byte) 0, (byte) 1, (byte) 5, (byte) 10, (byte) 31);
+        int bitField = Utils.to32BitField(bytes);
+        assertEquals(bytes, Utils.from32BitField(bitField));
+
+        bytes = new HashSet<>();
+        bitField = Utils.to32BitField(bytes);
+        assertEquals(bytes, Utils.from32BitField(bitField));
+
+        bytes = mkSet((byte) 0, (byte) 11, (byte) 32);
+        try {
+            Utils.to32BitField(bytes);
+            fail("Expected exception not thrown");
+        } catch (IllegalArgumentException e) {
+        }
+    }
+
+    @Test
+    public void testUnion() {
+        final Set oneSet = mkSet("a", "b", "c");
+        final Set anotherSet = mkSet("c", "d", "e");
+        final Set union = union(TreeSet::new, oneSet, anotherSet);
+
+        assertThat(union, is(mkSet("a", "b", "c", "d", "e")));
+        assertThat(union.getClass(), equalTo(TreeSet.class));
+    }
+
+    @Test
+    public void testUnionOfOne() {
+        final Set oneSet = mkSet("a", "b", "c");
+        final Set union = union(TreeSet::new, oneSet);
+
+        assertThat(union, is(mkSet("a", "b", "c")));
+        assertThat(union.getClass(), equalTo(TreeSet.class));
+    }
+
+    @Test
+    public void testUnionOfMany() {
+        final Set oneSet = mkSet("a", "b", "c");
+        final Set twoSet = mkSet("c", "d", "e");
+        final Set threeSet = mkSet("b", "c", "d");
+        final Set fourSet = mkSet("x", "y", "z");
+        final Set union = union(TreeSet::new, oneSet, twoSet, threeSet, fourSet);
+
+        assertThat(union, is(mkSet("a", "b", "c", "d", "e", "x", "y", "z")));
+        assertThat(union.getClass(), equalTo(TreeSet.class));
+    }
+
+    @Test
+    public void testUnionOfNone() {
+        final Set union = union(TreeSet::new);
+
+        assertThat(union, is(emptySet()));
+        assertThat(union.getClass(), equalTo(TreeSet.class));
+    }
+
+    @Test
+    public void testIntersection() {
+        final Set oneSet = mkSet("a", "b", "c");
+        final Set anotherSet = mkSet("c", "d", "e");
+        final Set intersection = intersection(TreeSet::new, oneSet, anotherSet);
+
+        assertThat(intersection, is(mkSet("c")));
+        assertThat(intersection.getClass(), equalTo(TreeSet.class));
+    }
+
+    @Test
+    public void testIntersectionOfOne() {
+        final Set oneSet = mkSet("a", "b", "c");
+        final Set intersection = intersection(TreeSet::new, oneSet);
+
+        assertThat(intersection, is(mkSet("a", "b", "c")));
+        assertThat(intersection.getClass(), equalTo(TreeSet.class));
+    }
+
+    @Test
+    public void testIntersectionOfMany() {
+        final Set oneSet = mkSet("a", "b", "c");
+        final Set twoSet = mkSet("c", "d", "e");
+        final Set threeSet = mkSet("b", "c", "d");
+        final Set union = intersection(TreeSet::new, oneSet, twoSet, threeSet);
+
+        assertThat(union, is(mkSet("c")));
+        assertThat(union.getClass(), equalTo(TreeSet.class));
+    }
+
+    @Test
+    public void testDisjointIntersectionOfMany() {
+        final Set oneSet = mkSet("a", "b", "c");
+        final Set twoSet = mkSet("c", "d", "e");
+        final Set threeSet = mkSet("b", "c", "d");
+        final Set fourSet = mkSet("x", "y", "z");
+        final Set union = intersection(TreeSet::new, oneSet, twoSet, threeSet, fourSet);
+
+        assertThat(union, is(emptySet()));
+        assertThat(union.getClass(), equalTo(TreeSet.class));
+    }
+
+    @Test
+    public void testDiff() {
+        final Set oneSet = mkSet("a", "b", "c");
+        final Set anotherSet = mkSet("c", "d", "e");
+        final Set diff = diff(TreeSet::new, oneSet, anotherSet);
+
+        assertThat(diff, is(mkSet("a", "b")));
+        assertThat(diff.getClass(), equalTo(TreeSet.class));
+    }
+
+    @Test
+    public void testPropsToMap() {
+        assertThrows(ConfigException.class, () -> {
+            Properties props = new Properties();
+            props.put(1, 2);
+            Utils.propsToMap(props);
+        });
+        assertValue(false);
+        assertValue(1);
+        assertValue("string");
+        assertValue(1.1);
+        assertValue(Collections.emptySet());
+        assertValue(Collections.emptyList());
+        assertValue(Collections.emptyMap());
+    }
+
+    private static void assertValue(Object value) {
+        Properties props = new Properties();
+        props.put("key", value);
+        assertEquals(Utils.propsToMap(props).get("key"), value);
+    }
+
+    @Test
+    public void testCloseAllQuietly() {
+        AtomicReference exception = new AtomicReference<>();
+        String msg = "you should fail";
+        AtomicInteger count = new AtomicInteger(0);
+        AutoCloseable c0 = () -> {
+            throw new RuntimeException(msg);
+        };
+        AutoCloseable c1 = count::incrementAndGet;
+        Utils.closeAllQuietly(exception, "test", Stream.of(c0, c1).toArray(AutoCloseable[]::new));
+        assertEquals(msg, exception.get().getMessage());
+        assertEquals(1, count.get());
+    }
+
+    @Test
+    public void shouldAcceptValidDateFormats() throws ParseException {
+        //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"));
+        invokeGetDateTimeMethod(new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss.SSSX"));
+        invokeGetDateTimeMethod(new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss.SSSXX"));
+        invokeGetDateTimeMethod(new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss.SSSXXX"));
+    }
+
+    @Test
+    public void shouldThrowOnInvalidDateFormatOrNullTimestamp() {
+        //check some invalid formats
+        // test null timestamp
+        assertThat(assertThrows(IllegalArgumentException.class, () -> {
+            Utils.getDateTime(null);
+        }).getMessage(), containsString("Error parsing timestamp with null value"));
+
+        // test pattern: yyyy-MM-dd'T'HH:mm:ss.X
+        checkExceptionForGetDateTimeMethod(() -> {
+            invokeGetDateTimeMethod(new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss.X"));
+        });
+
+        // test pattern: yyyy-MM-dd HH:mm:ss
+        assertThat(assertThrows(ParseException.class, () -> {
+            invokeGetDateTimeMethod(new SimpleDateFormat("yyyy-MM-dd HH:mm:ss"));
+        }).getMessage(), containsString("It does not contain a 'T' according to ISO8601 format"));
+
+        // KAFKA-10685: use DateTimeFormatter generate micro/nano second timestamp
+        final DateTimeFormatter formatter = new DateTimeFormatterBuilder()
+            .appendPattern("yyyy-MM-dd'T'HH:mm:ss")
+            .appendFraction(ChronoField.NANO_OF_SECOND, 0, 9, true)
+            .toFormatter();
+        final LocalDateTime timestampWithNanoSeconds = LocalDateTime.of(2020, 11, 9, 12, 34, 56, 123456789);
+        final LocalDateTime timestampWithMicroSeconds = timestampWithNanoSeconds.truncatedTo(ChronoUnit.MICROS);
+        final LocalDateTime timestampWithSeconds = timestampWithNanoSeconds.truncatedTo(ChronoUnit.SECONDS);
+
+        // test pattern: yyyy-MM-dd'T'HH:mm:ss.SSSSSSSSS
+        checkExceptionForGetDateTimeMethod(() -> {
+            Utils.getDateTime(formatter.format(timestampWithNanoSeconds));
+        });
+
+        // test pattern: yyyy-MM-dd'T'HH:mm:ss.SSSSSS
+        checkExceptionForGetDateTimeMethod(() -> {
+            Utils.getDateTime(formatter.format(timestampWithMicroSeconds));
+        });
+
+        // test pattern: yyyy-MM-dd'T'HH:mm:ss
+        checkExceptionForGetDateTimeMethod(() -> {
+            Utils.getDateTime(formatter.format(timestampWithSeconds));
+        });
+    }
+
+    private void checkExceptionForGetDateTimeMethod(ThrowingRunnable runnable) {
+        assertThat(assertThrows(ParseException.class, runnable)
+            .getMessage(), containsString("Unparseable date"));
+    }
+
+    private void invokeGetDateTimeMethod(final SimpleDateFormat format) throws ParseException {
+        final Date checkpoint = new Date();
+        final String formattedCheckpoint = format.format(checkpoint);
+        Utils.getDateTime(formattedCheckpoint);
+    }
+
 }
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/MockDeserializer.java b/clients/src/test/java/org/apache/kafka/test/MockDeserializer.java
index 99551f718a9c2..ac2865e9bb8b6 100644
--- a/clients/src/test/java/org/apache/kafka/test/MockDeserializer.java
+++ b/clients/src/test/java/org/apache/kafka/test/MockDeserializer.java
@@ -31,6 +31,9 @@ public class MockDeserializer implements ClusterResourceListener, Deserializer clusterIdBeforeDeserialize = new AtomicReference<>(noClusterId);
 
+    public boolean isKey;
+    public Map configs;
+
     public static void resetStaticVariables() {
         initCount = new AtomicInteger(0);
         closeCount = new AtomicInteger(0);
@@ -44,6 +47,8 @@ public MockDeserializer() {
 
     @Override
     public void configure(Map configs, boolean isKey) {
+        this.configs = configs;
+        this.isKey = isKey;
     }
 
     @Override
diff --git a/clients/src/test/java/org/apache/kafka/test/MockMetricsReporter.java b/clients/src/test/java/org/apache/kafka/test/MockMetricsReporter.java
index b5f3855034fe2..40521f583dd26 100644
--- a/clients/src/test/java/org/apache/kafka/test/MockMetricsReporter.java
+++ b/clients/src/test/java/org/apache/kafka/test/MockMetricsReporter.java
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.test;
 
+import org.apache.kafka.clients.CommonClientConfigs;
 import org.apache.kafka.common.metrics.KafkaMetric;
 import org.apache.kafka.common.metrics.MetricsReporter;
 
@@ -26,6 +27,7 @@
 public class MockMetricsReporter implements MetricsReporter {
     public static final AtomicInteger INIT_COUNT = new AtomicInteger(0);
     public static final AtomicInteger CLOSE_COUNT = new AtomicInteger(0);
+    public String clientId;
 
     public MockMetricsReporter() {
     }
@@ -48,5 +50,6 @@ public void close() {
 
     @Override
     public void configure(Map configs) {
+        clientId = (String) configs.get(CommonClientConfigs.CLIENT_ID_CONFIG);
     }
 }
\ No newline at end of file
diff --git a/clients/src/test/java/org/apache/kafka/test/MockProducerInterceptor.java b/clients/src/test/java/org/apache/kafka/test/MockProducerInterceptor.java
index 2a0751efc8ee8..133ff567d47fa 100644
--- a/clients/src/test/java/org/apache/kafka/test/MockProducerInterceptor.java
+++ b/clients/src/test/java/org/apache/kafka/test/MockProducerInterceptor.java
@@ -51,7 +51,7 @@ public void configure(Map configs) {
         Object o = configs.get(APPEND_STRING_PROP);
         if (o == null)
             throw new ConfigException("Mock producer interceptor expects configuration " + APPEND_STRING_PROP);
-        if (o != null && o instanceof String)
+        if (o instanceof String)
             appendStr = (String) o;
 
         // clientId also must be in configs
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 225aba4755c78..f98c80e73a880 100644
--- a/clients/src/test/java/org/apache/kafka/test/MockSelector.java
+++ b/clients/src/test/java/org/apache/kafka/test/MockSelector.java
@@ -16,32 +16,35 @@
  */
 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.requests.ByteBufferChannel;
+import org.apache.kafka.common.utils.Time;
+
 import java.io.IOException;
 import java.net.InetSocketAddress;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 
-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.utils.Time;
-
 /**
  * A fake selector to use for testing
  */
 public class MockSelector implements Selectable {
 
     private final Time time;
-    private final List initiatedSends = new ArrayList();
-    private final List completedSends = new ArrayList();
-    private final List completedReceives = new ArrayList();
+    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();
+    private final List connected = new ArrayList<>();
     private final List delayedReceives = new ArrayList<>();
 
     public MockSelector(Time time) {
@@ -63,7 +66,11 @@ public void close() {
 
     @Override
     public void close(String id) {
-        this.disconnected.put(id, ChannelState.LOCAL_CLOSE);
+        // Note that there are no notifications for client-side disconnects
+
+        removeSendsForNode(id, completedSends);
+        removeSendsForNode(id, initiatedSends);
+
         for (int i = 0; i < this.connected.size(); i++) {
             if (this.connected.get(i).equals(id)) {
                 this.connected.remove(i);
@@ -72,42 +79,96 @@ public void close(String id) {
         }
     }
 
+    /**
+     * Since MockSelector.connect will always succeed and add the
+     * connection id to the Set connected, we can only simulate
+     * that the connection is still pending by remove the connection
+     * id from the Set connected
+     *
+     * @param id connection id
+     */
+    public void serverConnectionBlocked(String id) {
+        this.connected.remove(id);
+    }
+
+    /**
+     * Simulate a server disconnect. This id will be present in {@link #disconnected()} on
+     * the next {@link #poll(long)}.
+     */
+    public void serverDisconnect(String id) {
+        this.disconnected.put(id, ChannelState.READY);
+        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) {
+        sends.removeIf(send -> id.equals(send.destinationId()));
+    }
+
     public void clear() {
         this.completedSends.clear();
         this.completedReceives.clear();
+        this.completedSendBuffers.clear();
         this.disconnected.clear();
         this.connected.clear();
     }
 
     @Override
-    public void send(Send send) {
+    public void send(NetworkSend send) {
         this.initiatedSends.add(send);
     }
 
     @Override
     public void poll(long timeout) throws IOException {
-        this.completedSends.addAll(this.initiatedSends);
+        completeInitiatedSends();
+        completeDelayedReceives();
+        time.sleep(timeout);
+    }
+
+    private void completeInitiatedSends() throws IOException {
+        for (NetworkSend send : initiatedSends) {
+            completeSend(send);
+        }
         this.initiatedSends.clear();
-        for (Send completedSend : completedSends) {
+    }
+
+    private void completeSend(NetworkSend send) throws IOException {
+        // Consume the send so that we will be able to send more requests to the destination
+        try (ByteBufferChannel discardChannel = new ByteBufferChannel(send.size())) {
+            while (!send.completed()) {
+                send.writeTo(discardChannel);
+            }
+            completedSends.add(send);
+            completedSendBuffers.add(discardChannel);
+        }
+    }
+
+    private void completeDelayedReceives() {
+        for (NetworkSend completedSend : completedSends) {
             Iterator delayedReceiveIterator = delayedReceives.iterator();
             while (delayedReceiveIterator.hasNext()) {
                 DelayedReceive delayedReceive = delayedReceiveIterator.next();
-                if (delayedReceive.source().equals(completedSend.destination())) {
+                if (delayedReceive.source().equals(completedSend.destinationId())) {
                     completedReceives.add(delayedReceive.receive());
                     delayedReceiveIterator.remove();
                 }
             }
         }
-        time.sleep(timeout);
     }
 
     @Override
-    public List completedSends() {
+    public List completedSends() {
         return completedSends;
     }
 
-    public void completeSend(NetworkSend send) {
-        this.completedSends.add(send);
+    public List completedSendBuffers() {
+        return completedSendBuffers;
     }
 
     @Override
@@ -155,4 +216,10 @@ public void unmuteAll() {
     public boolean isChannelReady(String id) {
         return true;
     }
+
+    public void reset() {
+        clear();
+        initiatedSends.clear();
+        delayedReceives.clear();
+    }
 }
diff --git a/clients/src/test/java/org/apache/kafka/test/MockSerializer.java b/clients/src/test/java/org/apache/kafka/test/MockSerializer.java
index 0c597c8b72aee..1c1444550801a 100644
--- a/clients/src/test/java/org/apache/kafka/test/MockSerializer.java
+++ b/clients/src/test/java/org/apache/kafka/test/MockSerializer.java
@@ -20,7 +20,6 @@
 import org.apache.kafka.common.ClusterResource;
 import org.apache.kafka.common.serialization.Serializer;
 
-import java.util.Map;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 
@@ -35,10 +34,6 @@ public MockSerializer() {
         INIT_COUNT.incrementAndGet();
     }
 
-    @Override
-    public void configure(Map configs, boolean isKey) {
-    }
-
     @Override
     public byte[] serialize(String topic, byte[] data) {
         // This will ensure that we get the cluster metadata when serialize is called for the first time
diff --git a/clients/src/test/java/org/apache/kafka/test/NoRetryException.java b/clients/src/test/java/org/apache/kafka/test/NoRetryException.java
new file mode 100644
index 0000000000000..a6f7db95afdeb
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/test/NoRetryException.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.test;
+
+/**
+ * This class can be used in the callback given to {@link TestUtils#retryOnExceptionWithTimeout(long, long, ValuelessCallable)}
+ * to indicate that a particular exception should not be retried. Instead the retry operation will
+ * be aborted immediately and the exception will be rethrown.
+ */
+public class NoRetryException extends RuntimeException {
+    private final Throwable cause;
+
+    public NoRetryException(Throwable cause) {
+        this.cause = cause;
+    }
+
+    @Override
+    public Throwable getCause() {
+        return this.cause;
+    }
+}
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 9087ad41a9b2e..a30ee68d0295a 100644
--- a/clients/src/test/java/org/apache/kafka/test/TestCondition.java
+++ b/clients/src/test/java/org/apache/kafka/test/TestCondition.java
@@ -20,7 +20,8 @@
  * Interface to wrap actions that are required to wait until a condition is met
  * for testing purposes.  Note that this is not intended to do any assertions.
  */
+@FunctionalInterface
 public interface TestCondition {
 
-    boolean conditionMet();
+    boolean conditionMet() throws Exception;
 }
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 6c057d0ef31d5..fc72d3da68030 100644
--- a/clients/src/test/java/org/apache/kafka/test/TestSslUtils.java
+++ b/clients/src/test/java/org/apache/kafka/test/TestSslUtils.java
@@ -17,32 +17,12 @@
 package org.apache.kafka.test;
 
 import org.apache.kafka.common.config.SslConfigs;
-import org.apache.kafka.common.network.Mode;
-
-import java.io.File;
-import java.io.FileOutputStream;
-import java.io.FileInputStream;
-import java.io.IOException;
-import java.io.EOFException;
-import java.math.BigInteger;
-import java.net.InetAddress;
-
-import javax.net.ssl.TrustManagerFactory;
-
-import java.security.GeneralSecurityException;
-import java.security.Key;
-import java.security.KeyPair;
-import java.security.KeyPairGenerator;
-import java.security.KeyStore;
-import java.security.NoSuchAlgorithmException;
-import java.security.SecureRandom;
-import java.security.Security;
-import java.security.cert.Certificate;
-import java.security.cert.CertificateException;
-import java.security.cert.X509Certificate;
-
 import org.apache.kafka.common.config.types.Password;
+import org.apache.kafka.common.network.Mode;
+import org.apache.kafka.common.security.auth.SslEngineFactory;
+import org.apache.kafka.common.security.ssl.DefaultSslEngineFactory;
 import org.bouncycastle.asn1.DEROctetString;
+import org.bouncycastle.asn1.DERSequence;
 import org.bouncycastle.asn1.x500.X500Name;
 import org.bouncycastle.asn1.x509.AlgorithmIdentifier;
 import org.bouncycastle.asn1.x509.Extension;
@@ -55,19 +35,66 @@
 import org.bouncycastle.crypto.params.AsymmetricKeyParameter;
 import org.bouncycastle.crypto.util.PrivateKeyFactory;
 import org.bouncycastle.jce.provider.BouncyCastleProvider;
+import org.bouncycastle.openssl.PKCS8Generator;
+import org.bouncycastle.openssl.jcajce.JcaMiscPEMGenerator;
+import org.bouncycastle.openssl.jcajce.JcaPKCS8Generator;
+import org.bouncycastle.openssl.jcajce.JceOpenSSLPKCS8EncryptorBuilder;
 import org.bouncycastle.operator.ContentSigner;
 import org.bouncycastle.operator.DefaultDigestAlgorithmIdentifierFinder;
 import org.bouncycastle.operator.DefaultSignatureAlgorithmIdentifierFinder;
+import org.bouncycastle.operator.bc.BcContentSignerBuilder;
+import org.bouncycastle.operator.bc.BcDSAContentSignerBuilder;
+import org.bouncycastle.operator.bc.BcECContentSignerBuilder;
 import org.bouncycastle.operator.bc.BcRSAContentSignerBuilder;
+import org.bouncycastle.util.io.pem.PemWriter;
 
+import java.io.ByteArrayOutputStream;
+import java.io.EOFException;
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.io.OutputStreamWriter;
+import java.math.BigInteger;
+import java.net.InetAddress;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Paths;
+import java.security.GeneralSecurityException;
+import java.security.Key;
+import java.security.KeyPair;
+import java.security.KeyPairGenerator;
+import java.security.KeyStore;
+import java.security.NoSuchAlgorithmException;
+import java.security.PrivateKey;
+import java.security.SecureRandom;
+import java.security.Security;
+import java.security.cert.Certificate;
+import java.security.cert.CertificateException;
+import java.security.cert.X509Certificate;
+import java.util.ArrayList;
+import java.util.Collections;
 import java.util.Date;
+import java.util.Enumeration;
 import java.util.HashMap;
-import java.util.Map;
 import java.util.List;
-import java.util.ArrayList;
+import java.util.Locale;
+import java.util.Map;
+import java.util.Properties;
+import java.util.Set;
+
+import javax.net.ssl.SSLEngine;
+import javax.net.ssl.TrustManagerFactory;
+
+import static org.apache.kafka.common.security.ssl.DefaultSslEngineFactory.PEM_TYPE;
 
 public class TestSslUtils {
 
+    public static final String TRUST_STORE_PASSWORD = "TrustStorePassword";
+    public static final String DEFAULT_TLS_PROTOCOL_FOR_TESTS = SslConfigs.DEFAULT_SSL_PROTOCOL;
+
     /**
      * 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.
@@ -87,7 +114,7 @@ public static X509Certificate generateCertificate(String dn, KeyPair pair,
 
     public static KeyPair generateKeyPair(String algorithm) throws NoSuchAlgorithmException {
         KeyPairGenerator keyGen = KeyPairGenerator.getInstance(algorithm);
-        keyGen.initialize(1024);
+        keyGen.initialize(algorithm.equals("EC") ? 256 : 2048);
         return keyGen.genKeyPair();
     }
 
@@ -99,20 +126,11 @@ private static KeyStore createEmptyKeyStore() throws GeneralSecurityException, I
 
     private static void saveKeyStore(KeyStore ks, String filename,
                                      Password password) throws GeneralSecurityException, IOException {
-        try (FileOutputStream out = new FileOutputStream(filename)) {
+        try (OutputStream out = Files.newOutputStream(Paths.get(filename))) {
             ks.store(out, password.value().toCharArray());
         }
     }
 
-    public static void createKeyStore(String filename,
-                                      Password password, String alias,
-                                      Key privateKey, Certificate cert) throws GeneralSecurityException, IOException {
-        KeyStore ks = createEmptyKeyStore();
-        ks.setKeyEntry(alias, privateKey, password.value().toCharArray(),
-                new Certificate[]{cert});
-        saveKeyStore(ks, filename, password);
-    }
-
     /**
      * Creates a keystore with a single key and saves it to a file.
      *
@@ -137,7 +155,7 @@ public static void createKeyStore(String filename,
     public static  void createTrustStore(
             String filename, Password password, Map certs) throws GeneralSecurityException, IOException {
         KeyStore ks = KeyStore.getInstance("JKS");
-        try (FileInputStream in = new FileInputStream(filename)) {
+        try (InputStream in = Files.newInputStream(Paths.get(filename))) {
             ks.load(in, password.value().toCharArray());
         } catch (EOFException e) {
             ks = createEmptyKeyStore();
@@ -148,26 +166,15 @@ public static  void createTrustStore(
         saveKeyStore(ks, filename, password);
     }
 
-    private static Map createSslConfig(Mode mode, File keyStoreFile, Password password, Password keyPassword,
-                                                       File trustStoreFile, Password trustStorePassword) {
+    public static Map createSslConfig(String keyManagerAlgorithm, String trustManagerAlgorithm, String tlsProtocol) {
         Map sslConfigs = new HashMap<>();
-        sslConfigs.put(SslConfigs.SSL_PROTOCOL_CONFIG, "TLSv1.2"); // protocol to create SSLContext
+        sslConfigs.put(SslConfigs.SSL_PROTOCOL_CONFIG, tlsProtocol); // protocol to create SSLContext
 
-        if (mode == Mode.SERVER || (mode == Mode.CLIENT && keyStoreFile != null)) {
-            sslConfigs.put(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG, keyStoreFile.getPath());
-            sslConfigs.put(SslConfigs.SSL_KEYSTORE_TYPE_CONFIG, "JKS");
-            sslConfigs.put(SslConfigs.SSL_KEYMANAGER_ALGORITHM_CONFIG, TrustManagerFactory.getDefaultAlgorithm());
-            sslConfigs.put(SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG, password);
-            sslConfigs.put(SslConfigs.SSL_KEY_PASSWORD_CONFIG, keyPassword);
-        }
-
-        sslConfigs.put(SslConfigs.SSL_TRUSTSTORE_LOCATION_CONFIG, trustStoreFile.getPath());
-        sslConfigs.put(SslConfigs.SSL_TRUSTSTORE_PASSWORD_CONFIG, trustStorePassword);
-        sslConfigs.put(SslConfigs.SSL_TRUSTSTORE_TYPE_CONFIG, "JKS");
-        sslConfigs.put(SslConfigs.SSL_TRUSTMANAGER_ALGORITHM_CONFIG, TrustManagerFactory.getDefaultAlgorithm());
+        sslConfigs.put(SslConfigs.SSL_KEYMANAGER_ALGORITHM_CONFIG, keyManagerAlgorithm);
+        sslConfigs.put(SslConfigs.SSL_TRUSTMANAGER_ALGORITHM_CONFIG, trustManagerAlgorithm);
 
         List enabledProtocols  = new ArrayList<>();
-        enabledProtocols.add("TLSv1.2");
+        enabledProtocols.add(tlsProtocol);
         sslConfigs.put(SslConfigs.SSL_ENABLED_PROTOCOLS_CONFIG, enabledProtocols);
 
         return sslConfigs;
@@ -179,42 +186,172 @@ public static  Map createSslConfig(boolean useClientCert, boolea
     }
 
     public static  Map createSslConfig(boolean useClientCert, boolean trustStore,
-            Mode mode, File trustStoreFile, String certAlias, String hostName)
+            Mode mode, File trustStoreFile, String certAlias, String cn)
         throws IOException, GeneralSecurityException {
-        return createSslConfig(useClientCert, trustStore, mode, trustStoreFile, certAlias, hostName, new CertificateBuilder());
+        return createSslConfig(useClientCert, trustStore, mode, trustStoreFile, certAlias, cn, new CertificateBuilder());
     }
 
-    public static  Map createSslConfig(boolean useClientCert, boolean trustStore,
+    public static  Map createSslConfig(boolean useClientCert, boolean createTrustStore,
             Mode mode, File trustStoreFile, String certAlias, String cn, CertificateBuilder certBuilder)
             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");
-
-        if (mode == Mode.CLIENT && useClientCert) {
-            keyStoreFile = File.createTempFile("clientKS", ".jks");
-            KeyPair cKP = generateKeyPair("RSA");
-            X509Certificate cCert = certBuilder.generate("CN=" + cn + ", O=A client", cKP);
-            createKeyStore(keyStoreFile.getPath(), password, "client", cKP.getPrivate(), cCert);
-            certs.put(certAlias, cCert);
-            keyStoreFile.deleteOnExit();
-        } else if (mode == Mode.SERVER) {
-            keyStoreFile = File.createTempFile("serverKS", ".jks");
-            KeyPair sKP = generateKeyPair("RSA");
-            X509Certificate sCert = certBuilder.generate("CN=" + cn + ", O=A server", sKP);
-            createKeyStore(keyStoreFile.getPath(), password, password, "server", sKP.getPrivate(), sCert);
-            certs.put(certAlias, sCert);
-            keyStoreFile.deleteOnExit();
-        }
-
-        if (trustStore) {
-            createTrustStore(trustStoreFile.getPath(), trustStorePassword, certs);
-            trustStoreFile.deleteOnExit();
-        }
-
-        return createSslConfig(mode, keyStoreFile, password, password, trustStoreFile, trustStorePassword);
+        SslConfigsBuilder builder = new SslConfigsBuilder(mode)
+                .useClientCert(useClientCert)
+                .certAlias(certAlias)
+                .cn(cn)
+                .certBuilder(certBuilder);
+        if (createTrustStore)
+            builder = builder.createNewTrustStore(trustStoreFile);
+        else
+            builder = builder.useExistingTrustStore(trustStoreFile);
+        return builder.build();
+    }
+
+    public static void convertToPem(Map sslProps, boolean writeToFile, boolean encryptPrivateKey) throws Exception {
+        String tsPath = (String) sslProps.get(SslConfigs.SSL_TRUSTSTORE_LOCATION_CONFIG);
+        String tsType = (String) sslProps.get(SslConfigs.SSL_TRUSTSTORE_TYPE_CONFIG);
+        Password tsPassword = (Password) sslProps.remove(SslConfigs.SSL_TRUSTSTORE_PASSWORD_CONFIG);
+        Password trustCerts = (Password) sslProps.remove(SslConfigs.SSL_TRUSTSTORE_CERTIFICATES_CONFIG);
+        if (trustCerts == null && tsPath != null) {
+            trustCerts = exportCertificates(tsPath, tsPassword, tsType);
+        }
+        if (trustCerts != null) {
+            if (tsPath == null) {
+                tsPath = File.createTempFile("truststore", ".pem").getPath();
+                sslProps.put(SslConfigs.SSL_TRUSTSTORE_LOCATION_CONFIG, tsPath);
+            }
+            sslProps.put(SslConfigs.SSL_TRUSTSTORE_TYPE_CONFIG, PEM_TYPE);
+            if (writeToFile)
+                writeToFile(tsPath, trustCerts);
+            else {
+                sslProps.put(SslConfigs.SSL_TRUSTSTORE_CERTIFICATES_CONFIG, trustCerts);
+                sslProps.remove(SslConfigs.SSL_TRUSTSTORE_LOCATION_CONFIG);
+            }
+        }
+
+        String ksPath = (String) sslProps.get(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG);
+        Password certChain = (Password) sslProps.remove(SslConfigs.SSL_KEYSTORE_CERTIFICATE_CHAIN_CONFIG);
+        Password key = (Password) sslProps.remove(SslConfigs.SSL_KEYSTORE_KEY_CONFIG);
+        if (certChain == null && ksPath != null) {
+            String ksType = (String) sslProps.get(SslConfigs.SSL_KEYSTORE_TYPE_CONFIG);
+            Password ksPassword = (Password) sslProps.remove(SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG);
+            Password keyPassword = (Password) sslProps.get(SslConfigs.SSL_KEY_PASSWORD_CONFIG);
+            certChain = exportCertificates(ksPath, ksPassword, ksType);
+            Password pemKeyPassword = encryptPrivateKey ? keyPassword : null;
+            key = exportPrivateKey(ksPath, ksPassword, keyPassword, ksType, pemKeyPassword);
+            if (!encryptPrivateKey)
+                sslProps.remove(SslConfigs.SSL_KEY_PASSWORD_CONFIG);
+        }
+
+        if (certChain != null) {
+            if (ksPath == null) {
+                ksPath = File.createTempFile("keystore", ".pem").getPath();
+                sslProps.put(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG, ksPath);
+            }
+            sslProps.put(SslConfigs.SSL_KEYSTORE_TYPE_CONFIG, PEM_TYPE);
+            if (writeToFile)
+                writeToFile(ksPath, key, certChain);
+            else {
+                sslProps.put(SslConfigs.SSL_KEYSTORE_KEY_CONFIG, key);
+                sslProps.put(SslConfigs.SSL_KEYSTORE_CERTIFICATE_CHAIN_CONFIG, certChain);
+                sslProps.remove(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG);
+            }
+        }
+    }
+
+    private static void writeToFile(String path, Password... entries) throws IOException {
+        try (FileOutputStream out = new FileOutputStream(path)) {
+            for (Password entry: entries) {
+                out.write(entry.value().getBytes(StandardCharsets.UTF_8));
+            }
+        }
+    }
+
+    public static void convertToPemWithoutFiles(Properties sslProps) throws Exception {
+        String tsPath = sslProps.getProperty(SslConfigs.SSL_TRUSTSTORE_LOCATION_CONFIG);
+        if (tsPath != null) {
+            Password trustCerts = exportCertificates(tsPath,
+                    (Password) sslProps.get(SslConfigs.SSL_TRUSTSTORE_PASSWORD_CONFIG),
+                    sslProps.getProperty(SslConfigs.SSL_TRUSTSTORE_TYPE_CONFIG));
+            sslProps.remove(SslConfigs.SSL_TRUSTSTORE_LOCATION_CONFIG);
+            sslProps.remove(SslConfigs.SSL_TRUSTSTORE_PASSWORD_CONFIG);
+            sslProps.setProperty(SslConfigs.SSL_TRUSTSTORE_TYPE_CONFIG, PEM_TYPE);
+            sslProps.put(SslConfigs.SSL_TRUSTSTORE_CERTIFICATES_CONFIG, trustCerts);
+        }
+        String ksPath = sslProps.getProperty(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG);
+        if (ksPath != null) {
+            String ksType = sslProps.getProperty(SslConfigs.SSL_KEYSTORE_TYPE_CONFIG);
+            Password ksPassword = (Password) sslProps.get(SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG);
+            Password keyPassword = (Password) sslProps.get(SslConfigs.SSL_KEY_PASSWORD_CONFIG);
+            Password certChain = exportCertificates(ksPath, ksPassword, ksType);
+            Password key = exportPrivateKey(ksPath, ksPassword, keyPassword, ksType, keyPassword);
+            sslProps.remove(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG);
+            sslProps.remove(SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG);
+            sslProps.setProperty(SslConfigs.SSL_KEYSTORE_TYPE_CONFIG, PEM_TYPE);
+            sslProps.put(SslConfigs.SSL_KEYSTORE_CERTIFICATE_CHAIN_CONFIG, certChain);
+            sslProps.put(SslConfigs.SSL_KEYSTORE_KEY_CONFIG, key);
+        }
+    }
+
+    public static Password exportCertificates(String storePath, Password storePassword, String storeType) throws Exception {
+        StringBuilder builder = new StringBuilder();
+        try (FileInputStream in = new FileInputStream(storePath)) {
+            KeyStore ks = KeyStore.getInstance(storeType);
+            ks.load(in, storePassword.value().toCharArray());
+            Enumeration aliases = ks.aliases();
+            if (!aliases.hasMoreElements())
+                throw new IllegalArgumentException("No certificates found in file " + storePath);
+            while (aliases.hasMoreElements()) {
+                String alias = aliases.nextElement();
+                Certificate[] certs = ks.getCertificateChain(alias);
+                if (certs != null) {
+                    for (Certificate cert : certs) {
+                        builder.append(pem(cert));
+                    }
+                } else {
+                    builder.append(pem(ks.getCertificate(alias)));
+                }
+            }
+        }
+        return new Password(builder.toString());
+    }
+
+    public static Password exportPrivateKey(String storePath,
+                                            Password storePassword,
+                                            Password keyPassword,
+                                            String storeType,
+                                            Password pemKeyPassword) throws Exception {
+        try (FileInputStream in = new FileInputStream(storePath)) {
+            KeyStore ks = KeyStore.getInstance(storeType);
+            ks.load(in, storePassword.value().toCharArray());
+            String alias = ks.aliases().nextElement();
+            return new Password(pem((PrivateKey) ks.getKey(alias, keyPassword.value().toCharArray()), pemKeyPassword));
+        }
+    }
+
+    static String pem(Certificate cert) throws IOException {
+        ByteArrayOutputStream out = new ByteArrayOutputStream();
+        try (PemWriter pemWriter = new PemWriter(new OutputStreamWriter(out, StandardCharsets.UTF_8.name()))) {
+            pemWriter.writeObject(new JcaMiscPEMGenerator(cert));
+        }
+        return new String(out.toByteArray(), StandardCharsets.UTF_8);
+    }
+
+    static String pem(PrivateKey privateKey, Password password) throws IOException {
+        ByteArrayOutputStream out = new ByteArrayOutputStream();
+        try (PemWriter pemWriter = new PemWriter(new OutputStreamWriter(out, StandardCharsets.UTF_8.name()))) {
+            if (password == null) {
+                pemWriter.writeObject(new JcaPKCS8Generator(privateKey, null));
+            } else {
+                JceOpenSSLPKCS8EncryptorBuilder encryptorBuilder = new JceOpenSSLPKCS8EncryptorBuilder(PKCS8Generator.PBE_SHA1_3DES);
+                encryptorBuilder.setPasssword(password.value().toCharArray());
+                try {
+                    pemWriter.writeObject(new JcaPKCS8Generator(privateKey, encryptorBuilder.build()));
+                } catch (Exception e) {
+                    throw new RuntimeException(e);
+                }
+            }
+        }
+        return new String(out.toByteArray(), StandardCharsets.UTF_8);
     }
 
     public static class CertificateBuilder {
@@ -231,8 +368,11 @@ public CertificateBuilder(int days, String algorithm) {
             this.algorithm = algorithm;
         }
 
-        public CertificateBuilder sanDnsName(String hostName) throws IOException {
-            subjectAltName = new GeneralNames(new GeneralName(GeneralName.dNSName, hostName)).getEncoded();
+        public CertificateBuilder sanDnsNames(String... hostNames) throws IOException {
+            GeneralName[] altNames = new GeneralName[hostNames.length];
+            for (int i = 0; i < hostNames.length; i++)
+                altNames[i] = new GeneralName(GeneralName.dNSName, hostNames[i]);
+            subjectAltName = GeneralNames.getInstance(new DERSequence(altNames)).getEncoded();
             return this;
         }
 
@@ -248,7 +388,17 @@ public X509Certificate generate(String dn, KeyPair keyPair) throws CertificateEx
                 AlgorithmIdentifier digAlgId = new DefaultDigestAlgorithmIdentifierFinder().find(sigAlgId);
                 AsymmetricKeyParameter privateKeyAsymKeyParam = PrivateKeyFactory.createKey(keyPair.getPrivate().getEncoded());
                 SubjectPublicKeyInfo subPubKeyInfo = SubjectPublicKeyInfo.getInstance(keyPair.getPublic().getEncoded());
-                ContentSigner sigGen = new BcRSAContentSignerBuilder(sigAlgId, digAlgId).build(privateKeyAsymKeyParam);
+                BcContentSignerBuilder signerBuilder;
+                String keyAlgorithm = keyPair.getPublic().getAlgorithm();
+                if (keyAlgorithm.equals("RSA"))
+                    signerBuilder = new BcRSAContentSignerBuilder(sigAlgId, digAlgId);
+                else if (keyAlgorithm.equals("DSA"))
+                    signerBuilder = new BcDSAContentSignerBuilder(sigAlgId, digAlgId);
+                else if (keyAlgorithm.equals("EC"))
+                    signerBuilder = new BcECContentSignerBuilder(sigAlgId, digAlgId);
+                else
+                    throw new IllegalArgumentException("Unsupported algorithm " + keyAlgorithm);
+                ContentSigner sigGen = signerBuilder.build(privateKeyAsymKeyParam);
                 X500Name name = new X500Name(dn);
                 Date from = new Date();
                 Date to = new Date(from.getTime() + days * 86400000L);
@@ -266,4 +416,208 @@ public X509Certificate generate(String dn, KeyPair keyPair) throws CertificateEx
             }
         }
     }
+
+    public static class SslConfigsBuilder {
+        final Mode mode;
+        String tlsProtocol;
+        boolean useClientCert;
+        boolean createTrustStore;
+        File trustStoreFile;
+        Password trustStorePassword;
+        Password keyStorePassword;
+        Password keyPassword;
+        String certAlias;
+        String cn;
+        String algorithm;
+        CertificateBuilder certBuilder;
+        boolean usePem;
+
+        public SslConfigsBuilder(Mode mode) {
+            this.mode = mode;
+            this.tlsProtocol = DEFAULT_TLS_PROTOCOL_FOR_TESTS;
+            trustStorePassword = new Password(TRUST_STORE_PASSWORD);
+            keyStorePassword = mode == Mode.SERVER ? new Password("ServerPassword") : new Password("ClientPassword");
+            keyPassword = keyStorePassword;
+            this.certBuilder = new CertificateBuilder();
+            this.cn = "localhost";
+            this.certAlias = mode.name().toLowerCase(Locale.ROOT);
+            this.algorithm = "RSA";
+            this.createTrustStore = true;
+        }
+
+        public SslConfigsBuilder tlsProtocol(String tlsProtocol) {
+            this.tlsProtocol = tlsProtocol;
+            return this;
+        }
+
+        public SslConfigsBuilder createNewTrustStore(File trustStoreFile) {
+            this.trustStoreFile = trustStoreFile;
+            this.createTrustStore = true;
+            return this;
+        }
+
+        public SslConfigsBuilder useExistingTrustStore(File trustStoreFile) {
+            this.trustStoreFile = trustStoreFile;
+            this.createTrustStore = false;
+            return this;
+        }
+
+        public SslConfigsBuilder useClientCert(boolean useClientCert) {
+            this.useClientCert = useClientCert;
+            return this;
+        }
+
+        public SslConfigsBuilder certAlias(String certAlias) {
+            this.certAlias = certAlias;
+            return this;
+        }
+
+        public SslConfigsBuilder cn(String cn) {
+            this.cn = cn;
+            return this;
+        }
+
+        public SslConfigsBuilder algorithm(String algorithm) {
+            this.algorithm = algorithm;
+            return this;
+        }
+
+        public SslConfigsBuilder certBuilder(CertificateBuilder certBuilder) {
+            this.certBuilder = certBuilder;
+            return this;
+        }
+
+        public SslConfigsBuilder usePem(boolean usePem) {
+            this.usePem = usePem;
+            return this;
+        }
+
+        public  Map build() throws IOException, GeneralSecurityException {
+            if (usePem) {
+                return buildPem();
+            } else
+                return buildJks();
+        }
+
+        private Map buildJks() throws IOException, GeneralSecurityException {
+            Map certs = new HashMap<>();
+            File keyStoreFile = null;
+
+            if (mode == Mode.CLIENT && useClientCert) {
+                keyStoreFile = File.createTempFile("clientKS", ".jks");
+                KeyPair cKP = generateKeyPair(algorithm);
+                X509Certificate cCert = certBuilder.generate("CN=" + cn + ", O=A client", cKP);
+                createKeyStore(keyStoreFile.getPath(), keyStorePassword, keyPassword, "client", cKP.getPrivate(), cCert);
+                certs.put(certAlias, cCert);
+            } else if (mode == Mode.SERVER) {
+                keyStoreFile = File.createTempFile("serverKS", ".jks");
+                KeyPair sKP = generateKeyPair(algorithm);
+                X509Certificate sCert = certBuilder.generate("CN=" + cn + ", O=A server", sKP);
+                createKeyStore(keyStoreFile.getPath(), keyStorePassword, keyPassword, "server", sKP.getPrivate(), sCert);
+                certs.put(certAlias, sCert);
+                keyStoreFile.deleteOnExit();
+            }
+
+            if (createTrustStore) {
+                createTrustStore(trustStoreFile.getPath(), trustStorePassword, certs);
+                trustStoreFile.deleteOnExit();
+            }
+
+            Map sslConfigs = new HashMap<>();
+
+            sslConfigs.put(SslConfigs.SSL_PROTOCOL_CONFIG, tlsProtocol); // protocol to create SSLContext
+
+            if (mode == Mode.SERVER || (mode == Mode.CLIENT && keyStoreFile != null)) {
+                sslConfigs.put(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG, keyStoreFile.getPath());
+                sslConfigs.put(SslConfigs.SSL_KEYSTORE_TYPE_CONFIG, "JKS");
+                sslConfigs.put(SslConfigs.SSL_KEYMANAGER_ALGORITHM_CONFIG, TrustManagerFactory.getDefaultAlgorithm());
+                sslConfigs.put(SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG, keyStorePassword);
+                sslConfigs.put(SslConfigs.SSL_KEY_PASSWORD_CONFIG, keyPassword);
+            }
+
+            sslConfigs.put(SslConfigs.SSL_TRUSTSTORE_LOCATION_CONFIG, trustStoreFile.getPath());
+            sslConfigs.put(SslConfigs.SSL_TRUSTSTORE_PASSWORD_CONFIG, trustStorePassword);
+            sslConfigs.put(SslConfigs.SSL_TRUSTSTORE_TYPE_CONFIG, "JKS");
+            sslConfigs.put(SslConfigs.SSL_TRUSTMANAGER_ALGORITHM_CONFIG, TrustManagerFactory.getDefaultAlgorithm());
+
+            List enabledProtocols  = new ArrayList<>();
+            enabledProtocols.add(tlsProtocol);
+            sslConfigs.put(SslConfigs.SSL_ENABLED_PROTOCOLS_CONFIG, enabledProtocols);
+
+            return sslConfigs;
+        }
+
+        private Map buildPem() throws IOException, GeneralSecurityException {
+            if (!createTrustStore) {
+                throw new IllegalArgumentException("PEM configs cannot be created with existing trust stores");
+            }
+
+            Map sslConfigs = new HashMap<>();
+            sslConfigs.put(SslConfigs.SSL_PROTOCOL_CONFIG, tlsProtocol);
+            sslConfigs.put(SslConfigs.SSL_ENABLED_PROTOCOLS_CONFIG, Collections.singletonList(tlsProtocol));
+
+            if (mode != Mode.CLIENT || useClientCert) {
+                KeyPair keyPair = generateKeyPair(algorithm);
+                X509Certificate cert = certBuilder.generate("CN=" + cn + ", O=A " + mode.name().toLowerCase(Locale.ROOT), keyPair);
+
+                Password privateKeyPem = new Password(pem(keyPair.getPrivate(), keyPassword));
+                Password certPem = new Password(pem(cert));
+                sslConfigs.put(SslConfigs.SSL_KEYSTORE_TYPE_CONFIG, PEM_TYPE);
+                sslConfigs.put(SslConfigs.SSL_TRUSTSTORE_TYPE_CONFIG, PEM_TYPE);
+                sslConfigs.put(SslConfigs.SSL_KEYSTORE_KEY_CONFIG, privateKeyPem);
+                sslConfigs.put(SslConfigs.SSL_KEYSTORE_CERTIFICATE_CHAIN_CONFIG, certPem);
+                sslConfigs.put(SslConfigs.SSL_KEY_PASSWORD_CONFIG, keyPassword);
+                sslConfigs.put(SslConfigs.SSL_TRUSTSTORE_CERTIFICATES_CONFIG, certPem);
+            }
+            return sslConfigs;
+        }
+    }
+
+    public static final class TestSslEngineFactory implements SslEngineFactory {
+
+        public boolean closed = false;
+
+        DefaultSslEngineFactory defaultSslEngineFactory = new DefaultSslEngineFactory();
+
+        @Override
+        public SSLEngine createClientSslEngine(String peerHost, int peerPort, String endpointIdentification) {
+            return defaultSslEngineFactory.createClientSslEngine(peerHost, peerPort, endpointIdentification);
+        }
+
+        @Override
+        public SSLEngine createServerSslEngine(String peerHost, int peerPort) {
+            return defaultSslEngineFactory.createServerSslEngine(peerHost, peerPort);
+        }
+
+        @Override
+        public boolean shouldBeRebuilt(Map nextConfigs) {
+            return defaultSslEngineFactory.shouldBeRebuilt(nextConfigs);
+        }
+
+        @Override
+        public Set reconfigurableConfigs() {
+            return defaultSslEngineFactory.reconfigurableConfigs();
+        }
+
+        @Override
+        public KeyStore keystore() {
+            return defaultSslEngineFactory.keystore();
+        }
+
+        @Override
+        public KeyStore truststore() {
+            return defaultSslEngineFactory.truststore();
+        }
+
+        @Override
+        public void close() throws IOException {
+            defaultSslEngineFactory.close();
+            closed = true;
+        }
+
+        @Override
+        public void configure(Map configs) {
+            defaultSslEngineFactory.configure(configs);
+        }
+    }
 }
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 958ab2cc636a8..db6f17951c73a 100644
--- a/clients/src/test/java/org/apache/kafka/test/TestUtils.java
+++ b/clients/src/test/java/org/apache/kafka/test/TestUtils.java
@@ -16,41 +16,56 @@
  */
 package org.apache.kafka.test;
 
+import java.io.FileWriter;
 import org.apache.kafka.clients.consumer.ConsumerConfig;
 import org.apache.kafka.clients.producer.ProducerConfig;
 import org.apache.kafka.common.Cluster;
 import org.apache.kafka.common.Node;
 import org.apache.kafka.common.PartitionInfo;
-import org.apache.kafka.common.protocol.types.Struct;
-import org.apache.kafka.common.utils.Base64;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.network.NetworkReceive;
+import org.apache.kafka.common.network.Send;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.requests.ByteBufferChannel;
+import org.apache.kafka.common.requests.RequestHeader;
+import org.apache.kafka.common.utils.Exit;
 import org.apache.kafka.common.utils.Utils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 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;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Properties;
 import java.util.Random;
 import java.util.Set;
 import java.util.UUID;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Future;
+import java.util.function.Consumer;
+import java.util.function.Supplier;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
 import static java.util.Arrays.asList;
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertNotNull;
-import static org.junit.Assert.assertTrue;
-import static org.junit.Assert.fail;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.fail;
 
 /**
  * Helper functions for writing unit tests
@@ -64,28 +79,22 @@ public class TestUtils {
     public static final String DIGITS = "0123456789";
     public static final String LETTERS_AND_DIGITS = LETTERS + DIGITS;
 
-    public static final String GROUP_METADATA_TOPIC_NAME = "__consumer_offsets";
-    public static final Set INTERNAL_TOPICS = Collections.singleton(GROUP_METADATA_TOPIC_NAME);
-
     /* 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() {
         return clusterWith(1);
     }
 
-    public static Cluster singletonCluster(final Map topicPartitionCounts) {
-        return clusterWith(1, topicPartitionCounts);
-    }
-
     public static Cluster singletonCluster(final String topic, final int partitions) {
         return clusterWith(1, topic, partitions);
     }
 
     public static Cluster clusterWith(int nodes) {
-        return clusterWith(nodes, new HashMap());
+        return clusterWith(nodes, new HashMap<>());
     }
 
     public static Cluster clusterWith(final int nodes, final Map topicPartitionCounts) {
@@ -99,7 +108,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(), INTERNAL_TOPICS);
+        return new Cluster("kafka-cluster", asList(ns), parts, Collections.emptySet(), Collections.emptySet());
     }
 
     public static Cluster clusterWith(final int nodes, final String topic, final int partitions) {
@@ -141,6 +150,19 @@ public static File tempFile() throws IOException {
         return file;
     }
 
+    /**
+     * Create a file with the given contents in the default temporary-file directory,
+     * using `kafka` as the prefix and `tmp` as the suffix to generate its name.
+     */
+    public static File tempFile(final String contents) throws IOException {
+        final File file = tempFile();
+        final FileWriter writer = new FileWriter(file);
+        writer.write(contents);
+        writer.close();
+
+        return file;
+    }
+
     /**
      * Create a temporary relative directory in the default temporary-file directory with the given prefix.
      *
@@ -177,14 +199,11 @@ public static File tempDirectory(final Path parent, String prefix) {
         }
         file.deleteOnExit();
 
-        Runtime.getRuntime().addShutdownHook(new Thread() {
-            @Override
-            public void run() {
-                try {
-                    Utils.delete(file);
-                } catch (IOException e) {
-                    log.error("Error deleting {}", file.getAbsolutePath(), e);
-                }
+        Exit.addShutdownHook("delete-temp-file-shutdown-hook", () -> {
+            try {
+                Utils.delete(file);
+            } catch (IOException e) {
+                log.error("Error deleting {}", file.getAbsolutePath(), e);
             }
         });
 
@@ -192,27 +211,26 @@ public void run() {
     }
 
     public static Properties producerConfig(final String bootstrapServers,
-                                            final Class keySerializer,
-                                            final Class valueSerializer,
+                                            final Class keySerializer,
+                                            final Class valueSerializer,
                                             final Properties additional) {
         final Properties properties = new Properties();
         properties.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers);
         properties.put(ProducerConfig.ACKS_CONFIG, "all");
-        properties.put(ProducerConfig.RETRIES_CONFIG, 0);
         properties.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, keySerializer);
         properties.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, valueSerializer);
         properties.putAll(additional);
         return properties;
     }
 
-    public static Properties producerConfig(final String bootstrapServers, final Class keySerializer, final Class valueSerializer) {
+    public static Properties producerConfig(final String bootstrapServers, final Class keySerializer, final Class valueSerializer) {
         return producerConfig(bootstrapServers, keySerializer, valueSerializer, new Properties());
     }
 
     public static Properties consumerConfig(final String bootstrapServers,
                                             final String groupId,
-                                            final Class keyDeserializer,
-                                            final Class valueDeserializer,
+                                            final Class keyDeserializer,
+                                            final Class valueDeserializer,
                                             final Properties additional) {
 
         final Properties consumerConfig = new Properties();
@@ -227,8 +245,8 @@ public static Properties consumerConfig(final String bootstrapServers,
 
     public static Properties consumerConfig(final String bootstrapServers,
                                             final String groupId,
-                                            final Class keyDeserializer,
-                                            final Class valueDeserializer) {
+                                            final Class keyDeserializer,
+                                            final Class valueDeserializer) {
         return consumerConfig(bootstrapServers,
             groupId,
             keyDeserializer,
@@ -239,7 +257,7 @@ public static Properties consumerConfig(final String bootstrapServers,
     /**
      * returns consumer config with random UUID for the Group ID
      */
-    public static Properties consumerConfig(final String bootstrapServers, final Class keyDeserializer, final Class valueDeserializer) {
+    public static Properties consumerConfig(final String bootstrapServers, final Class keyDeserializer, final Class valueDeserializer) {
         return consumerConfig(bootstrapServers,
             UUID.randomUUID().toString(),
             keyDeserializer,
@@ -251,7 +269,14 @@ public static Properties consumerConfig(final String bootstrapServers, final Cla
      * uses default value of 15 seconds for timeout
      */
     public static void waitForCondition(final TestCondition testCondition, final String conditionDetails) throws InterruptedException {
-        waitForCondition(testCondition, DEFAULT_MAX_WAIT_MS, conditionDetails);
+        waitForCondition(testCondition, DEFAULT_MAX_WAIT_MS, () -> conditionDetails);
+    }
+
+    /**
+     * uses default value of 15 seconds for timeout
+     */
+    public static void waitForCondition(final TestCondition testCondition, final Supplier conditionDetailsSupplier) throws InterruptedException {
+        waitForCondition(testCondition, DEFAULT_MAX_WAIT_MS, conditionDetailsSupplier);
     }
 
     /**
@@ -261,19 +286,81 @@ public static void waitForCondition(final TestCondition testCondition, final Str
      * avoid transient failures due to slow or overloaded machines.
      */
     public static void waitForCondition(final TestCondition testCondition, final long maxWaitMs, String conditionDetails) throws InterruptedException {
-        final long startTime = System.currentTimeMillis();
+        waitForCondition(testCondition, maxWaitMs, () -> conditionDetails);
+    }
 
-        boolean testConditionMet;
-        while (!(testConditionMet = testCondition.conditionMet()) && ((System.currentTimeMillis() - startTime) < maxWaitMs)) {
-            Thread.sleep(Math.min(maxWaitMs, 100L));
-        }
+    /**
+     * 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
+     * without unnecessarily increasing test time (as the condition is checked frequently). The longer timeout is needed to
+     * avoid transient failures due to slow or overloaded machines.
+     */
+    public static void waitForCondition(final TestCondition testCondition, final long maxWaitMs, Supplier conditionDetailsSupplier) throws InterruptedException {
+        retryOnExceptionWithTimeout(maxWaitMs, () -> {
+            String conditionDetailsSupplied = conditionDetailsSupplier != null ? conditionDetailsSupplier.get() : null;
+            String conditionDetails = conditionDetailsSupplied != null ? conditionDetailsSupplied : "";
+            assertTrue(testCondition.conditionMet(),
+                "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 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);
+    }
+
+    /**
+     * 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);
+    }
 
-        // 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) {
-            conditionDetails = conditionDetails != null ? conditionDetails : "";
-            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 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 NoRetryException e) {
+                throw e;
+            } 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));
         }
     }
 
@@ -293,7 +380,7 @@ public static void isValidClusterId(String clusterId) {
 
         // Convert into normal variant and add padding at the end.
         String originalClusterId = String.format("%s==", clusterId.replace("_", "/").replace("-", "+"));
-        byte[] decodedUuid = Base64.decoder().decode(originalClusterId);
+        byte[] decodedUuid = Base64.getDecoder().decode(originalClusterId);
 
         // We expect 16 bytes, same as the input UUID.
         assertEquals(decodedUuid.length, 16);
@@ -335,10 +422,87 @@ public static  List toList(Iterable iterable) {
         return list;
     }
 
-    public static ByteBuffer toBuffer(Struct struct) {
-        ByteBuffer buffer = ByteBuffer.allocate(struct.sizeOf());
-        struct.writeTo(buffer);
-        buffer.rewind();
-        return buffer;
+    public static  Set toSet(Collection collection) {
+        return new HashSet<>(collection);
+    }
+
+    public static ByteBuffer toBuffer(Send send) {
+        ByteBufferChannel channel = new ByteBufferChannel(send.size());
+        try {
+            assertEquals(send.size(), send.writeTo(channel));
+            channel.close();
+            return channel.buffer();
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    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(exceptionCauseClass.isInstance(exception.getCause()),
+            "Unexpected exception cause " + exception.getCause());
+        return exceptionCauseClass.cast(exception.getCause());
+    }
+
+    public static void assertFutureError(Future future, Class exceptionClass)
+        throws InterruptedException {
+        try {
+            future.get();
+            fail("Expected a " + exceptionClass.getSimpleName() + " exception, but got success.");
+        } catch (ExecutionException ee) {
+            Throwable cause = ee.getCause();
+            assertEquals(exceptionClass, cause.getClass(),
+                "Expected a " + exceptionClass.getSimpleName() + " exception, but got " +
+                    cause.getClass().getSimpleName());
+        }
+    }
+
+    public static ApiKeys apiKeyFrom(NetworkReceive networkReceive) {
+        return RequestHeader.parse(networkReceive.payload().duplicate()).apiKey();
+    }
+
+    public static  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);
+        }
+    }
+
+    public static void setFieldValue(Object obj, String fieldName, Object value) throws Exception {
+        Field field = obj.getClass().getDeclaredField(fieldName);
+        field.setAccessible(true);
+        field.set(obj, value);
     }
 }
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..0a0fb1f6c6987
--- /dev/null
+++ b/clients/src/test/resources/common/message/SimpleExampleMessage.json
@@ -0,0 +1,62 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements.  See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT 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-2",
+  "flexibleVersions": "1+",
+  "fields": [
+    { "name": "processId", "versions": "1+", "type": "uuid" },
+    { "name": "myTaggedIntArray", "type": "[]int32",
+        "taggedVersions": "1+", "tag": 0 },
+    { "name": "myNullableString", "type": "string", "default": "null",
+        "nullableVersions": "1+", "taggedVersions": "1+", "tag": 1 },
+    { "name": "myInt16", "type": "int16", "default": "123",
+        "taggedVersions": "1+", "tag": 2 },
+    { "name": "myFloat64", "type": "float64", "default": "12.34",
+        "taggedVersions": "1+", "tag": 3 },
+    { "name": "myString", "type": "string", "taggedVersions": "1+", "tag": 4 },
+    { "name": "myBytes", "type": "bytes",
+        "nullableVersions": "1+", "taggedVersions": "1+", "tag": 5 },
+    { "name": "taggedUuid", "type": "uuid", "default": "212d54944a8b4fdf94b388b470beb367",
+        "taggedVersions": "1+", "tag": 6 },
+    { "name": "taggedLong", "type": "int64", "default": "0xcafcacafcacafca",
+        "taggedVersions": "1+", "tag": 7 },
+    { "name": "zeroCopyByteBuffer", "versions": "1", "type": "bytes", "zeroCopy": true },
+    { "name": "nullableZeroCopyByteBuffer", "versions": "1", "nullableVersions": "0+",
+      "type": "bytes", "zeroCopy": true },
+    { "name": "myStruct", "type": "MyStruct", "versions": "2+", "about": "Test Struct field",
+      "fields": [
+        { "name": "structId", "type": "int32", "versions": "2+", "about": "Int field in struct"},
+        { "name": "arrayInStruct", "type": "[]StructArray", "versions": "2+",
+          "fields": [
+            { "name": "arrayFieldId", "type": "int32", "versions": "2+"}
+          ]}
+    ]},
+    { "name": "myTaggedStruct", "type": "TaggedStruct", "versions": "2+", "about": "Test Tagged Struct field",
+      "taggedVersions": "2+", "tag": 8,
+      "fields": [
+        { "name": "structId", "type": "string", "versions": "2+", "about": "String field in struct"}
+    ]},
+    { "name":  "myCommonStruct", "type": "TestCommonStruct", "versions": "0+"},
+    { "name":  "myOtherCommonStruct", "type": "TestCommonStruct", "versions": "0+"}
+  ],
+  "commonStructs": [
+    { "name": "TestCommonStruct", "versions": "0+", "fields": [
+      { "name": "foo", "type": "int32", "default": "123", "versions": "0+" },
+      { "name": "bar", "type": "int32", "default": "123", "versions": "0+" }
+    ]}
+  ]
+}
diff --git a/clients/src/test/resources/common/message/SimpleRecordsMessage.json b/clients/src/test/resources/common/message/SimpleRecordsMessage.json
new file mode 100644
index 0000000000000..ce81b29b88e06
--- /dev/null
+++ b/clients/src/test/resources/common/message/SimpleRecordsMessage.json
@@ -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.
+{
+  "name": "SimpleRecordsMessage",
+  "type": "header",
+  "validVersions": "0-1",
+  "flexibleVersions": "1+",
+  "fields": [
+      { "name": "Topic", "type": "string", "versions": "0+", "entityType": "topicName",
+        "about": "The topic name." },
+      { "name": "RecordSet", "type": "records", "versions": "0+",
+        "nullableVersions": "0+", "about": "The record data." }
+  ]
+}
diff --git a/clients/src/test/resources/serializedData/offsetAndMetadataSerializedfile b/clients/src/test/resources/serializedData/offsetAndMetadataBeforeLeaderEpoch
similarity index 100%
rename from clients/src/test/resources/serializedData/offsetAndMetadataSerializedfile
rename to clients/src/test/resources/serializedData/offsetAndMetadataBeforeLeaderEpoch
diff --git a/clients/src/test/resources/serializedData/offsetAndMetadataWithLeaderEpoch b/clients/src/test/resources/serializedData/offsetAndMetadataWithLeaderEpoch
new file mode 100644
index 0000000000000..ddf3956a0cee8
Binary files /dev/null and b/clients/src/test/resources/serializedData/offsetAndMetadataWithLeaderEpoch differ
diff --git a/config/connect-distributed.properties b/config/connect-distributed.properties
index 5f3f35802cf55..72db145f3f892 100644
--- a/config/connect-distributed.properties
+++ b/config/connect-distributed.properties
@@ -34,13 +34,6 @@ value.converter=org.apache.kafka.connect.json.JsonConverter
 key.converter.schemas.enable=true
 value.converter.schemas.enable=true
 
-# The internal converter used for offsets, config, and status data is configurable and must be specified, but most users will
-# always want to use the built-in default. Offset, config, and status data is never visible outside of Kafka Connect in this format.
-internal.key.converter=org.apache.kafka.connect.json.JsonConverter
-internal.value.converter=org.apache.kafka.connect.json.JsonConverter
-internal.key.converter.schemas.enable=false
-internal.value.converter.schemas.enable=false
-
 # Topic to use for storing offsets. This topic should have many partitions and be replicated and compacted.
 # Kafka Connect will attempt to create the topic automatically when needed, but you can always manually create
 # the topic before starting Kafka Connect if a specific topic configuration is needed.
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..40afda5e4ad68
--- /dev/null
+++ b/config/connect-mirror-maker.properties
@@ -0,0 +1,59 @@
+# 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
+
+# connection information for each cluster
+# This is a comma separated host:port pairs for each cluster
+# for e.g. "A_host1:9092, A_host2:9092, A_host3:9092"
+A.bootstrap.servers = A_host1:9092, A_host2:9092, A_host3:9092
+B.bootstrap.servers = B_host1:9092, B_host2:9092, B_host3:9092
+
+# enable and configure individual replication flows
+A->B.enabled = true
+
+# regex which defines which topics gets replicated. For eg "foo-.*"
+A->B.topics = .*
+
+B->A.enabled = true
+B->A.topics = .*
+
+# Setting replication factor of newly created remote topics
+replication.factor=1
+
+############################# Internal Topic Settings  #############################
+# The replication factor for mm2 internal topics "heartbeats", "B.checkpoints.internal" and
+# "mm2-offset-syncs.B.internal"
+# For anything other than development testing, a value greater than 1 is recommended to ensure availability such as 3.
+checkpoints.topic.replication.factor=1
+heartbeats.topic.replication.factor=1
+offset-syncs.topic.replication.factor=1
+
+# The replication factor for connect internal topics "mm2-configs.B.internal", "mm2-offsets.B.internal" and
+# "mm2-status.B.internal"
+# For anything other than development testing, a value greater than 1 is recommended to ensure availability such as 3.
+offset.storage.replication.factor=1
+status.storage.replication.factor=1
+config.storage.replication.factor=1
+
+# customize as needed
+# replication.policy.separator = _
+# sync.topic.acls.enabled = false
+# emit.heartbeats.interval.seconds = 5
diff --git a/config/connect-standalone.properties b/config/connect-standalone.properties
index 0039796ddccae..a340a3bf315cc 100644
--- a/config/connect-standalone.properties
+++ b/config/connect-standalone.properties
@@ -25,13 +25,6 @@ value.converter=org.apache.kafka.connect.json.JsonConverter
 key.converter.schemas.enable=true
 value.converter.schemas.enable=true
 
-# The internal converter used for offsets and config data is configurable and must be specified, but most users will
-# always want to use the built-in default. Offset and config data is never visible outside of Kafka Connect in this format.
-internal.key.converter=org.apache.kafka.connect.json.JsonConverter
-internal.value.converter=org.apache.kafka.connect.json.JsonConverter
-internal.key.converter.schemas.enable=false
-internal.value.converter.schemas.enable=false
-
 offset.storage.file.filename=/tmp/connect.offsets
 # Flush much faster than normal, which is useful for testing/debugging
 offset.flush.interval.ms=10000
diff --git a/config/log4j.properties b/config/log4j.properties
index 3ff3f9e4f6c11..4cbce9d104291 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)
@@ -83,7 +82,7 @@ log4j.additivity.kafka.controller=false
 log4j.logger.kafka.log.LogCleaner=INFO, cleanerAppender
 log4j.additivity.kafka.log.LogCleaner=false
 
-log4j.logger.state.change.logger=TRACE, stateChangeAppender
+log4j.logger.state.change.logger=INFO, stateChangeAppender
 log4j.additivity.state.change.logger=false
 
 # Access denials are logged at INFO level, change to DEBUG to also log allowed accesses
diff --git a/config/producer.properties b/config/producer.properties
index 750b95ee0aeb7..4786b988a29b8 100644
--- a/config/producer.properties
+++ b/config/producer.properties
@@ -20,7 +20,7 @@
 # format: host1:port1,host2:port2 ...
 bootstrap.servers=localhost:9092
 
-# specify the compression codec for all data generated: none, gzip, snappy, lz4
+# specify the compression codec for all data generated: none, gzip, snappy, lz4, zstd
 compression.type=none
 
 # name of the partitioner class for partitioning events; default partition spreads data randomly
diff --git a/config/server.properties b/config/server.properties
index dc224049cfb78..b1cf5c454169c 100644
--- a/config/server.properties
+++ b/config/server.properties
@@ -56,7 +56,7 @@ socket.request.max.bytes=104857600
 
 ############################# Log Basics #############################
 
-# A comma seperated list of directories under which to store log files
+# A comma separated list of directories under which to store log files
 log.dirs=/tmp/kafka-logs
 
 # The default number of log partitions per topic. More partitions allow greater
@@ -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
@@ -82,7 +82,7 @@ transaction.state.log.min.isr=1
 # There are a few important trade-offs here:
 #    1. Durability: Unflushed data may be lost if you are not using replication.
 #    2. Latency: Very large flush intervals may lead to latency spikes when the flush does occur as there will be a lot of data to flush.
-#    3. Throughput: The flush is generally the most expensive operation, and a small flush interval may lead to exceessive seeks.
+#    3. Throughput: The flush is generally the most expensive operation, and a small flush interval may lead to excessive seeks.
 # The settings below allow one to configure the flush policy to flush data after a period of time or
 # every N messages (or both). This can be done globally and overridden on a per-topic basis.
 
@@ -123,7 +123,7 @@ log.retention.check.interval.ms=300000
 zookeeper.connect=localhost:2181
 
 # Timeout in ms for connecting to zookeeper
-zookeeper.connection.timeout.ms=6000
+zookeeper.connection.timeout.ms=18000
 
 
 ############################# Group Coordinator Settings #############################
@@ -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/trogdor.conf b/config/trogdor.conf
new file mode 100644
index 0000000000000..320cbe7560cd0
--- /dev/null
+++ b/config/trogdor.conf
@@ -0,0 +1,25 @@
+{
+    "_comment": [
+        "Licensed to the Apache Software Foundation (ASF) under one or more",
+        "contributor license agreements.  See the NOTICE file distributed with",
+        "this work for additional information regarding copyright ownership.",
+        "The ASF licenses this file to You under the Apache License, Version 2.0",
+        "(the \"License\"); you may not use this file except in compliance with",
+        "the License.  You may obtain a copy of the License at",
+        "",
+        "http://www.apache.org/licenses/LICENSE-2.0",
+        "",
+        "Unless required by applicable law or agreed to in writing, software",
+        "distributed under the License is distributed on an \"AS IS\" BASIS,",
+        "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.",
+        "See the License for the specific language governing permissions and",
+        "limitations under the License."
+    ],
+    "platform": "org.apache.kafka.trogdor.basic.BasicPlatform", "nodes": {
+        "node0": {
+            "hostname": "localhost",
+            "trogdor.agent.port": 8888,
+            "trogdor.coordinator.port": 8889
+        }
+    }
+}
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/components/Versioned.java b/connect/api/src/main/java/org/apache/kafka/connect/components/Versioned.java
new file mode 100644
index 0000000000000..adabe8fbce154
--- /dev/null
+++ b/connect/api/src/main/java/org/apache/kafka/connect/components/Versioned.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.connect.components;
+
+/**
+ * Connect requires some components implement this interface to define a version string.
+ */
+public interface Versioned {
+    /**
+     * Get the version of this component.
+     *
+     * @return the version, formatted as a String. The version may not be (@code null} or empty.
+     */
+    String version();
+}
diff --git a/connect/api/src/main/java/org/apache/kafka/connect/connector/ConnectRecord.java b/connect/api/src/main/java/org/apache/kafka/connect/connector/ConnectRecord.java
index 344e365fd3f08..b181209163926 100644
--- a/connect/api/src/main/java/org/apache/kafka/connect/connector/ConnectRecord.java
+++ b/connect/api/src/main/java/org/apache/kafka/connect/connector/ConnectRecord.java
@@ -17,11 +17,16 @@
 package org.apache.kafka.connect.connector;
 
 import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.header.ConnectHeaders;
+import org.apache.kafka.connect.header.Header;
+import org.apache.kafka.connect.header.Headers;
+
+import java.util.Objects;
 
 /**
  * 
            
  * Base class for records containing data to be copied to/from Kafka. This corresponds closely to
- * Kafka's ProducerRecord and ConsumerRecord classes, and holds the data that may be used by both
+ * Kafka's {@link org.apache.kafka.clients.producer.ProducerRecord ProducerRecord} and {@link org.apache.kafka.clients.consumer.ConsumerRecord ConsumerRecord} classes, and holds the data that may be used by both
  * sources and sinks (topic, kafkaPartition, key, value). Although both implementations include a
  * notion of offset, it is not included here because they differ in type.
  * 
            
@@ -34,11 +39,19 @@ public abstract class ConnectRecord> {
     private final Schema valueSchema;
     private final Object value;
     private final Long timestamp;
+    private final Headers headers;
 
     public ConnectRecord(String topic, Integer kafkaPartition,
                          Schema keySchema, Object key,
                          Schema valueSchema, Object value,
                          Long timestamp) {
+        this(topic, kafkaPartition, keySchema, key, valueSchema, value, timestamp, new ConnectHeaders());
+    }
+
+    public ConnectRecord(String topic, Integer kafkaPartition,
+                         Schema keySchema, Object key,
+                         Schema valueSchema, Object value,
+                         Long timestamp, Iterable
             headers) {
         this.topic = topic;
         this.kafkaPartition = kafkaPartition;
         this.keySchema = keySchema;
@@ -46,6 +59,11 @@ public ConnectRecord(String topic, Integer kafkaPartition,
         this.valueSchema = valueSchema;
         this.value = value;
         this.timestamp = timestamp;
+        if (headers instanceof ConnectHeaders) {
+            this.headers = (ConnectHeaders) headers;
+        } else {
+            this.headers = new ConnectHeaders(headers);
+        }
     }
 
     public String topic() {
@@ -76,17 +94,57 @@ public Long timestamp() {
         return timestamp;
     }
 
-    /** Generate a new record of the same type as itself, with the specified parameter values. **/
+    /**
+     * Get the headers for this record.
+     *
+     * @return the headers; never null
+     */
+    public Headers headers() {
+        return headers;
+    }
+
+    /**
+     * Create a new record of the same type as itself, with the specified parameter values. All other fields in this record will be copied
+     * over to the new record. Since the headers are mutable, the resulting record will have a copy of this record's headers.
+     *
+     * @param topic the name of the topic; may be null
+     * @param kafkaPartition the partition number for the Kafka topic; may be null
+     * @param keySchema the schema for the key; may be null
+     * @param key the key; may be null
+     * @param valueSchema the schema for the value; may be null
+     * @param value the value; may be null
+     * @param timestamp the timestamp; may be null
+     * @return the new record
+     */
     public abstract R newRecord(String topic, Integer kafkaPartition, Schema keySchema, Object key, Schema valueSchema, Object value, Long timestamp);
 
+    /**
+     * Create a new record of the same type as itself, with the specified parameter values. All other fields in this record will be copied
+     * over to the new record.
+     *
+     * @param topic the name of the topic; may be null
+     * @param kafkaPartition the partition number for the Kafka topic; may be null
+     * @param keySchema the schema for the key; may be null
+     * @param key the key; may be null
+     * @param valueSchema the schema for the value; may be null
+     * @param value the value; may be null
+     * @param timestamp the timestamp; may be null
+     * @param headers the headers; may be null or empty
+     * @return the new record
+     */
+    public abstract R newRecord(String topic, Integer kafkaPartition, Schema keySchema, Object key, Schema valueSchema, Object value, Long timestamp, Iterable
             headers);
+
     @Override
     public String toString() {
         return "ConnectRecord{" +
                 "topic='" + topic + '\'' +
                 ", kafkaPartition=" + kafkaPartition +
                 ", key=" + key +
+                ", keySchema=" + keySchema +
                 ", value=" + value +
+                ", valueSchema=" + valueSchema +
                 ", timestamp=" + timestamp +
+                ", headers=" + headers +
                 '}';
     }
 
@@ -99,22 +157,14 @@ public boolean equals(Object o) {
 
         ConnectRecord that = (ConnectRecord) o;
 
-        if (kafkaPartition != null ? !kafkaPartition.equals(that.kafkaPartition) : that.kafkaPartition != null)
-            return false;
-        if (topic != null ? !topic.equals(that.topic) : that.topic != null)
-            return false;
-        if (keySchema != null ? !keySchema.equals(that.keySchema) : that.keySchema != null)
-            return false;
-        if (key != null ? !key.equals(that.key) : that.key != null)
-            return false;
-        if (valueSchema != null ? !valueSchema.equals(that.valueSchema) : that.valueSchema != null)
-            return false;
-        if (value != null ? !value.equals(that.value) : that.value != null)
-            return false;
-        if (timestamp != null ? !timestamp.equals(that.timestamp) : that.timestamp != null)
-            return false;
-
-        return true;
+        return Objects.equals(kafkaPartition, that.kafkaPartition)
+               && Objects.equals(topic, that.topic)
+               && Objects.equals(keySchema, that.keySchema)
+               && Objects.equals(key, that.key)
+               && Objects.equals(valueSchema, that.valueSchema)
+               && Objects.equals(value, that.value)
+               && Objects.equals(timestamp, that.timestamp)
+               && Objects.equals(headers, that.headers);
     }
 
     @Override
@@ -126,6 +176,7 @@ public int hashCode() {
         result = 31 * result + (valueSchema != null ? valueSchema.hashCode() : 0);
         result = 31 * result + (value != null ? value.hashCode() : 0);
         result = 31 * result + (timestamp != null ? timestamp.hashCode() : 0);
+        result = 31 * result + headers.hashCode();
         return result;
     }
 }
diff --git a/connect/api/src/main/java/org/apache/kafka/connect/connector/Connector.java b/connect/api/src/main/java/org/apache/kafka/connect/connector/Connector.java
index a8a5dabfc839c..6d54aabf9f587 100644
--- a/connect/api/src/main/java/org/apache/kafka/connect/connector/Connector.java
+++ b/connect/api/src/main/java/org/apache/kafka/connect/connector/Connector.java
@@ -19,6 +19,8 @@
 import org.apache.kafka.common.config.Config;
 import org.apache.kafka.common.config.ConfigDef;
 import org.apache.kafka.common.config.ConfigValue;
+import org.apache.kafka.connect.errors.ConnectException;
+import org.apache.kafka.connect.components.Versioned;
 
 import java.util.List;
 import java.util.Map;
@@ -27,29 +29,24 @@
  * 
            
  * Connectors manage integration of Kafka Connect with another system, either as an input that ingests
  * data into Kafka or an output that passes data to an external system. Implementations should
- * not use this class directly; they should inherit from SourceConnector or SinkConnector.
+ * not use this class directly; they should inherit from {@link org.apache.kafka.connect.source.SourceConnector SourceConnector}
+ * or {@link org.apache.kafka.connect.sink.SinkConnector SinkConnector}.
  * 
            
  * 
            
  * Connectors have two primary tasks. First, given some configuration, they are responsible for
  * creating configurations for a set of {@link Task}s that split up the data processing. For
  * example, a database Connector might create Tasks by dividing the set of tables evenly among
  * tasks. Second, they are responsible for monitoring inputs for changes that require
- * reconfiguration and notifying the Kafka Connect runtime via the ConnectorContext. Continuing the
+ * reconfiguration and notifying the Kafka Connect runtime via the {@link ConnectorContext}. Continuing the
  * previous example, the connector might periodically check for new tables and notify Kafka Connect of
  * additions and deletions. Kafka Connect will then request new configurations and update the running
  * Tasks.
  * 
            
  */
-public abstract class Connector {
+public abstract class Connector implements Versioned {
 
     protected ConnectorContext context;
 
-    /**
-     * Get the version of this connector.
-     *
-     * @return the version, formatted as a String
-     */
-    public abstract String version();
 
     /**
      * Initialize this connector, using the provided ConnectorContext to notify the runtime of
@@ -82,6 +79,15 @@ public void initialize(ConnectorContext ctx, List> taskConfi
         // are very different, but reduces the difficulty of implementing a Connector
     }
 
+    /**
+     * Returns the context object used to interact with the Kafka Connect runtime.
+     *
+     * @return the context for this Connector.
+     */
+    protected ConnectorContext context() {
+        return context;
+    }
+
     /**
      * Start this Connector. This method will only be called on a clean Connector, i.e. it has
      * either just been instantiated and initialized or {@link #stop()} has been invoked.
@@ -130,13 +136,18 @@ public void reconfigure(Map props) {
      */
     public Config validate(Map connectorConfigs) {
         ConfigDef configDef = config();
+        if (null == configDef) {
+            throw new ConnectException(
+                String.format("%s.config() must return a ConfigDef that is not null.", this.getClass().getName())
+            );
+        }
         List configValues = configDef.validate(connectorConfigs);
         return new Config(configValues);
     }
 
     /**
      * Define the configuration for the connector.
-     * @return The ConfigDef for this connector.
+     * @return The ConfigDef for this connector; may not be null.
      */
     public abstract ConfigDef config();
 }
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/ConnectSchema.java b/connect/api/src/main/java/org/apache/kafka/connect/data/ConnectSchema.java
index 5b0579ee30c76..a465b12eeb573 100644
--- a/connect/api/src/main/java/org/apache/kafka/connect/data/ConnectSchema.java
+++ b/connect/api/src/main/java/org/apache/kafka/connect/data/ConnectSchema.java
@@ -22,6 +22,7 @@
 import java.nio.ByteBuffer;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.EnumMap;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -31,7 +32,7 @@ public class ConnectSchema implements Schema {
     /**
      * Maps Schema.Types to a list of Java classes that can be used to represent them.
      */
-    private static final Map> SCHEMA_TYPE_CLASSES = new HashMap<>();
+    private static final Map> SCHEMA_TYPE_CLASSES = new EnumMap<>(Type.class);
     /**
      * Maps known logical types to a list of Java classes that can be used to represent them.
      */
@@ -43,31 +44,31 @@ public class ConnectSchema implements Schema {
     private static final Map, Type> JAVA_CLASS_SCHEMA_TYPES = new HashMap<>();
 
     static {
-        SCHEMA_TYPE_CLASSES.put(Type.INT8, Arrays.asList((Class) Byte.class));
-        SCHEMA_TYPE_CLASSES.put(Type.INT16, Arrays.asList((Class) Short.class));
-        SCHEMA_TYPE_CLASSES.put(Type.INT32, Arrays.asList((Class) Integer.class));
-        SCHEMA_TYPE_CLASSES.put(Type.INT64, Arrays.asList((Class) Long.class));
-        SCHEMA_TYPE_CLASSES.put(Type.FLOAT32, Arrays.asList((Class) Float.class));
-        SCHEMA_TYPE_CLASSES.put(Type.FLOAT64, Arrays.asList((Class) Double.class));
-        SCHEMA_TYPE_CLASSES.put(Type.BOOLEAN, Arrays.asList((Class) Boolean.class));
-        SCHEMA_TYPE_CLASSES.put(Type.STRING, Arrays.asList((Class) String.class));
+        SCHEMA_TYPE_CLASSES.put(Type.INT8, Collections.singletonList((Class) Byte.class));
+        SCHEMA_TYPE_CLASSES.put(Type.INT16, Collections.singletonList((Class) Short.class));
+        SCHEMA_TYPE_CLASSES.put(Type.INT32, Collections.singletonList((Class) Integer.class));
+        SCHEMA_TYPE_CLASSES.put(Type.INT64, Collections.singletonList((Class) Long.class));
+        SCHEMA_TYPE_CLASSES.put(Type.FLOAT32, Collections.singletonList((Class) Float.class));
+        SCHEMA_TYPE_CLASSES.put(Type.FLOAT64, Collections.singletonList((Class) Double.class));
+        SCHEMA_TYPE_CLASSES.put(Type.BOOLEAN, Collections.singletonList((Class) Boolean.class));
+        SCHEMA_TYPE_CLASSES.put(Type.STRING, Collections.singletonList((Class) String.class));
         // Bytes are special and have 2 representations. byte[] causes problems because it doesn't handle equals() and
         // hashCode() like we want objects to, so we support both byte[] and ByteBuffer. Using plain byte[] can cause
         // those methods to fail, so ByteBuffers are recommended
         SCHEMA_TYPE_CLASSES.put(Type.BYTES, Arrays.asList((Class) byte[].class, (Class) ByteBuffer.class));
-        SCHEMA_TYPE_CLASSES.put(Type.ARRAY, Arrays.asList((Class) List.class));
-        SCHEMA_TYPE_CLASSES.put(Type.MAP, Arrays.asList((Class) Map.class));
-        SCHEMA_TYPE_CLASSES.put(Type.STRUCT, Arrays.asList((Class) Struct.class));
+        SCHEMA_TYPE_CLASSES.put(Type.ARRAY, Collections.singletonList((Class) List.class));
+        SCHEMA_TYPE_CLASSES.put(Type.MAP, Collections.singletonList((Class) Map.class));
+        SCHEMA_TYPE_CLASSES.put(Type.STRUCT, Collections.singletonList((Class) Struct.class));
 
         for (Map.Entry> schemaClasses : SCHEMA_TYPE_CLASSES.entrySet()) {
             for (Class schemaClass : schemaClasses.getValue())
                 JAVA_CLASS_SCHEMA_TYPES.put(schemaClass, schemaClasses.getKey());
         }
 
-        LOGICAL_TYPE_CLASSES.put(Decimal.LOGICAL_NAME, Arrays.asList((Class) BigDecimal.class));
-        LOGICAL_TYPE_CLASSES.put(Date.LOGICAL_NAME, Arrays.asList((Class) java.util.Date.class));
-        LOGICAL_TYPE_CLASSES.put(Time.LOGICAL_NAME, Arrays.asList((Class) java.util.Date.class));
-        LOGICAL_TYPE_CLASSES.put(Timestamp.LOGICAL_NAME, Arrays.asList((Class) java.util.Date.class));
+        LOGICAL_TYPE_CLASSES.put(Decimal.LOGICAL_NAME, Collections.singletonList((Class) BigDecimal.class));
+        LOGICAL_TYPE_CLASSES.put(Date.LOGICAL_NAME, Collections.singletonList((Class) java.util.Date.class));
+        LOGICAL_TYPE_CLASSES.put(Time.LOGICAL_NAME, Collections.singletonList((Class) java.util.Date.class));
+        LOGICAL_TYPE_CLASSES.put(Timestamp.LOGICAL_NAME, Collections.singletonList((Class) java.util.Date.class));
         // We don't need to put these into JAVA_CLASS_SCHEMA_TYPES since that's only used to determine schemas for
         // schemaless data and logical types will have ambiguous schemas (e.g. many of them use the same Java class) so
         // they should not be used without schemas.
@@ -179,6 +180,7 @@ public List fields() {
         return fields;
     }
 
+    @Override
     public Field field(String fieldName) {
         if (type != Type.STRUCT)
             throw new DataException("Cannot look up fields on non-struct type");
@@ -216,14 +218,10 @@ public static void validateValue(String name, Schema schema, Object value) {
             if (!schema.isOptional())
                 throw new DataException("Invalid value: null used for required field: \"" + name
                         + "\", schema type: " + schema.type());
-            else
-                return;
+            return;
         }
 
-        List expectedClasses = LOGICAL_TYPE_CLASSES.get(schema.name());
-
-        if (expectedClasses == null)
-                expectedClasses = SCHEMA_TYPE_CLASSES.get(schema.type());
+        List expectedClasses = expectedClassesFor(schema);
 
         if (expectedClasses == null)
             throw new DataException("Invalid Java object for schema type " + schema.type()
@@ -231,12 +229,17 @@ public static void validateValue(String name, Schema schema, Object value) {
                     + " for field: \"" + name + "\"");
 
         boolean foundMatch = false;
-        for (Class expectedClass : expectedClasses) {
-            if (expectedClass.isInstance(value)) {
-                foundMatch = true;
-                break;
+        if (expectedClasses.size() == 1) {
+            foundMatch = expectedClasses.get(0).isInstance(value);
+        } else {
+            for (Class expectedClass : expectedClasses) {
+                if (expectedClass.isInstance(value)) {
+                    foundMatch = true;
+                    break;
+                }
             }
         }
+
         if (!foundMatch)
             throw new DataException("Invalid Java object for schema type " + schema.type()
                     + ": " + value.getClass()
@@ -264,6 +267,13 @@ public static void validateValue(String name, Schema schema, Object value) {
         }
     }
 
+    private static List expectedClassesFor(Schema schema) {
+        List expectedClasses = LOGICAL_TYPE_CLASSES.get(schema.name());
+        if (expectedClasses == null)
+            expectedClasses = SCHEMA_TYPE_CLASSES.get(schema.type());
+        return expectedClasses;
+    }
+
     /**
      * Validate that the value can be used for this schema, i.e. that its type matches the schema type and optional
      * requirements. Throws a DataException if the value is invalid.
@@ -289,7 +299,7 @@ public boolean equals(Object o) {
                 Objects.equals(name, schema.name) &&
                 Objects.equals(doc, schema.doc) &&
                 Objects.equals(type, schema.type) &&
-                Objects.equals(defaultValue, schema.defaultValue) &&
+                Objects.deepEquals(defaultValue, schema.defaultValue) &&
                 Objects.equals(fields, schema.fields) &&
                 Objects.equals(keySchema, schema.keySchema) &&
                 Objects.equals(valueSchema, schema.valueSchema) &&
@@ -318,7 +328,7 @@ public String toString() {
      * Get the {@link Schema.Type} associated with the given class.
      *
      * @param klass the Class to
-     * @return the corresponding type, nor null if there is no matching type
+     * @return the corresponding type, or null if there is no matching type
      */
     public static Type schemaType(Class klass) {
         synchronized (JAVA_CLASS_SCHEMA_TYPES) {
diff --git a/connect/api/src/main/java/org/apache/kafka/connect/data/Field.java b/connect/api/src/main/java/org/apache/kafka/connect/data/Field.java
index 3a8df17a91017..b5d3f027968cb 100644
--- a/connect/api/src/main/java/org/apache/kafka/connect/data/Field.java
+++ b/connect/api/src/main/java/org/apache/kafka/connect/data/Field.java
@@ -73,4 +73,13 @@ public boolean equals(Object o) {
     public int hashCode() {
         return Objects.hash(name, index, schema);
     }
+
+    @Override
+    public String toString() {
+        return "Field{" +
+                "name=" + name +
+                ", index=" + index +
+                ", schema=" + schema +
+                "}";
+    }
 }
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 a9064b841fabf..9b4be4f801d7e 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
@@ -336,12 +336,14 @@ public SchemaBuilder field(String fieldName, Schema fieldSchema) {
      * Get the list of fields for this Schema. Throws a DataException if this schema is not a struct.
      * @return the list of fields for this Schema
      */
+    @Override
     public List fields() {
         if (type != Type.STRUCT)
             throw new DataException("Cannot list fields on non-struct type");
         return new ArrayList<>(fields.values());
     }
 
+    @Override
     public Field field(String fieldName) {
         if (type != Type.STRUCT)
             throw new DataException("Cannot look up fields on non-struct type");
@@ -380,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;
@@ -398,7 +420,7 @@ public Schema valueSchema() {
     public Schema build() {
         return new ConnectSchema(type, isOptional(), defaultValue, name, version, doc,
                 parameters == null ? null : Collections.unmodifiableMap(parameters),
-                fields == null ? null : Collections.unmodifiableList(new ArrayList(fields.values())), keySchema, valueSchema);
+                fields == null ? null : Collections.unmodifiableList(new ArrayList<>(fields.values())), keySchema, valueSchema);
     }
 
     /**
diff --git a/connect/api/src/main/java/org/apache/kafka/connect/data/SchemaProjector.java b/connect/api/src/main/java/org/apache/kafka/connect/data/SchemaProjector.java
index ea31752f60e58..5400705cd105c 100644
--- a/connect/api/src/main/java/org/apache/kafka/connect/data/SchemaProjector.java
+++ b/connect/api/src/main/java/org/apache/kafka/connect/data/SchemaProjector.java
@@ -160,7 +160,7 @@ private static Object projectPrimitive(Schema source, Object record, Schema targ
         assert source.type().isPrimitive();
         assert target.type().isPrimitive();
         Object result;
-        if (isPromotable(source.type(), target.type())) {
+        if (isPromotable(source.type(), target.type()) && record instanceof Number) {
             Number numberRecord = (Number) record;
             switch (target.type()) {
                 case INT8:
diff --git a/connect/api/src/main/java/org/apache/kafka/connect/data/Struct.java b/connect/api/src/main/java/org/apache/kafka/connect/data/Struct.java
index 6e7b5d23bbd78..1f542e5cf88c3 100644
--- a/connect/api/src/main/java/org/apache/kafka/connect/data/Struct.java
+++ b/connect/api/src/main/java/org/apache/kafka/connect/data/Struct.java
@@ -211,7 +211,9 @@ public Struct put(String fieldName, Object value) {
      * @return the Struct, to allow chaining of {@link #put(String, Object)} calls
      */
     public Struct put(Field field, Object value) {
-        ConnectSchema.validateValue(field.schema(), value);
+        if (null == field)
+            throw new DataException("field cannot be null.");
+        ConnectSchema.validateValue(field.name(), field.schema(), value);
         values[field.index()] = value;
         return this;
     }
@@ -238,12 +240,12 @@ public boolean equals(Object o) {
         if (o == null || getClass() != o.getClass()) return false;
         Struct struct = (Struct) o;
         return Objects.equals(schema, struct.schema) &&
-                Arrays.equals(values, struct.values);
+                Arrays.deepEquals(values, struct.values);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(schema, Arrays.hashCode(values));
+        return Objects.hash(schema, Arrays.deepHashCode(values));
     }
 
     private Field lookupField(String fieldName) {
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
new file mode 100644
index 0000000000000..067c91bbe3025
--- /dev/null
+++ b/connect/api/src/main/java/org/apache/kafka/connect/data/Values.java
@@ -0,0 +1,1265 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.data;
+
+import org.apache.kafka.common.utils.Utils;
+import org.apache.kafka.connect.data.Schema.Type;
+import org.apache.kafka.connect.errors.DataException;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.math.BigDecimal;
+import java.nio.ByteBuffer;
+import java.text.CharacterIterator;
+import java.text.DateFormat;
+import java.text.ParseException;
+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;
+
+/**
+ * Utility for converting from one Connect value to a different form. This is useful when the caller expects a value of a particular type
+ * but is uncertain whether the actual value is one that isn't directly that type but can be converted into that type.
+ *
+ * 
            For example, a caller might expects a particular {@link org.apache.kafka.connect.header.Header} to contain an {@link Type#INT64}
+ * value, when in fact that header contains a string representation of a 32-bit integer. Here, the caller can use the methods in this
+ * class to convert the value to the desired type:
+ * 
            + *     Header header = ...
+ *     long value = Values.convertToLong(header.schema(), header.value());
+ * 
            
+ *
+ * 
            This class is able to convert any value to a string representation as well as parse those string representations back into most of
+ * the types. The only exception is {@link Struct} values that require a schema and thus cannot be parsed from a simple string.
+ */
+public class Values {
+
+    private static final Logger LOG = LoggerFactory.getLogger(Values.class);
+
+    private static final TimeZone UTC = TimeZone.getTimeZone("UTC");
+    private static final SchemaAndValue NULL_SCHEMA_AND_VALUE = new SchemaAndValue(null, null);
+    private static final SchemaAndValue TRUE_SCHEMA_AND_VALUE = new SchemaAndValue(Schema.BOOLEAN_SCHEMA, Boolean.TRUE);
+    private static final SchemaAndValue FALSE_SCHEMA_AND_VALUE = new SchemaAndValue(Schema.BOOLEAN_SCHEMA, Boolean.FALSE);
+    private static final Schema ARRAY_SELECTOR_SCHEMA = SchemaBuilder.array(Schema.STRING_SCHEMA).build();
+    private static final Schema MAP_SELECTOR_SCHEMA = SchemaBuilder.map(Schema.STRING_SCHEMA, Schema.STRING_SCHEMA).build();
+    private static final Schema STRUCT_SELECTOR_SCHEMA = SchemaBuilder.struct().build();
+    private static final String TRUE_LITERAL = Boolean.TRUE.toString();
+    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";
+    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 = ",";
+    private static final String ENTRY_DELIMITER = ":";
+    private static final String ARRAY_BEGIN_DELIMITER = "[";
+    private static final String ARRAY_END_DELIMITER = "]";
+    private static final String MAP_BEGIN_DELIMITER = "{";
+    private static final String MAP_END_DELIMITER = "}";
+    private static final int ISO_8601_DATE_LENGTH = ISO_8601_DATE_FORMAT_PATTERN.length();
+    private static final int ISO_8601_TIME_LENGTH = ISO_8601_TIME_FORMAT_PATTERN.length() - 2; // subtract single quotes
+    private static final int ISO_8601_TIMESTAMP_LENGTH = ISO_8601_TIMESTAMP_FORMAT_PATTERN.length() - 4; // subtract single quotes
+
+    private static final Pattern TWO_BACKSLASHES = Pattern.compile("\\\\");
+
+    private static final Pattern DOUBLEQOUTE = Pattern.compile("\"");
+
+    /**
+     * Convert the specified value to an {@link Type#BOOLEAN} value. The supplied schema is required if the value is a logical
+     * type when the schema contains critical information that might be necessary for converting to a boolean.
+     *
+     * @param schema the schema for the value; may be null
+     * @param value  the value to be converted; may be null
+     * @return the representation as a boolean, or null if the supplied value was null
+     * @throws DataException if the value could not be converted to a boolean
+     */
+    public static Boolean convertToBoolean(Schema schema, Object value) throws DataException {
+        return (Boolean) convertTo(Schema.OPTIONAL_BOOLEAN_SCHEMA, schema, value);
+    }
+
+    /**
+     * Convert the specified value to an {@link Type#INT8} byte value. The supplied schema is required if the value is a logical
+     * type when the schema contains critical information that might be necessary for converting to a byte.
+     *
+     * @param schema the schema for the value; may be null
+     * @param value  the value to be converted; may be null
+     * @return the representation as a byte, or null if the supplied value was null
+     * @throws DataException if the value could not be converted to a byte
+     */
+    public static Byte convertToByte(Schema schema, Object value) throws DataException {
+        return (Byte) convertTo(Schema.OPTIONAL_INT8_SCHEMA, schema, value);
+    }
+
+    /**
+     * Convert the specified value to an {@link Type#INT16} short value. The supplied schema is required if the value is a logical
+     * type when the schema contains critical information that might be necessary for converting to a short.
+     *
+     * @param schema the schema for the value; may be null
+     * @param value  the value to be converted; may be null
+     * @return the representation as a short, or null if the supplied value was null
+     * @throws DataException if the value could not be converted to a short
+     */
+    public static Short convertToShort(Schema schema, Object value) throws DataException {
+        return (Short) convertTo(Schema.OPTIONAL_INT16_SCHEMA, schema, value);
+    }
+
+    /**
+     * Convert the specified value to an {@link Type#INT32} int value. The supplied schema is required if the value is a logical
+     * type when the schema contains critical information that might be necessary for converting to an integer.
+     *
+     * @param schema the schema for the value; may be null
+     * @param value  the value to be converted; may be null
+     * @return the representation as an integer, or null if the supplied value was null
+     * @throws DataException if the value could not be converted to an integer
+     */
+    public static Integer convertToInteger(Schema schema, Object value) throws DataException {
+        return (Integer) convertTo(Schema.OPTIONAL_INT32_SCHEMA, schema, value);
+    }
+
+    /**
+     * Convert the specified value to an {@link Type#INT64} long value. The supplied schema is required if the value is a logical
+     * type when the schema contains critical information that might be necessary for converting to a long.
+     *
+     * @param schema the schema for the value; may be null
+     * @param value  the value to be converted; may be null
+     * @return the representation as a long, or null if the supplied value was null
+     * @throws DataException if the value could not be converted to a long
+     */
+    public static Long convertToLong(Schema schema, Object value) throws DataException {
+        return (Long) convertTo(Schema.OPTIONAL_INT64_SCHEMA, schema, value);
+    }
+
+    /**
+     * Convert the specified value to an {@link Type#FLOAT32} float value. The supplied schema is required if the value is a logical
+     * type when the schema contains critical information that might be necessary for converting to a floating point number.
+     *
+     * @param schema the schema for the value; may be null
+     * @param value  the value to be converted; may be null
+     * @return the representation as a float, or null if the supplied value was null
+     * @throws DataException if the value could not be converted to a float
+     */
+    public static Float convertToFloat(Schema schema, Object value) throws DataException {
+        return (Float) convertTo(Schema.OPTIONAL_FLOAT32_SCHEMA, schema, value);
+    }
+
+    /**
+     * Convert the specified value to an {@link Type#FLOAT64} double value. The supplied schema is required if the value is a logical
+     * type when the schema contains critical information that might be necessary for converting to a floating point number.
+     *
+     * @param schema the schema for the value; may be null
+     * @param value  the value to be converted; may be null
+     * @return the representation as a double, or null if the supplied value was null
+     * @throws DataException if the value could not be converted to a double
+     */
+    public static Double convertToDouble(Schema schema, Object value) throws DataException {
+        return (Double) convertTo(Schema.OPTIONAL_FLOAT64_SCHEMA, schema, value);
+    }
+
+    /**
+     * Convert the specified value to an {@link Type#STRING} value.
+     * Not supplying a schema may limit the ability to convert to the desired type.
+     *
+     * @param schema the schema for the value; may be null
+     * @param value  the value to be converted; may be null
+     * @return the representation as a string, or null if the supplied value was null
+     */
+    public static String convertToString(Schema schema, Object value) {
+        return (String) convertTo(Schema.OPTIONAL_STRING_SCHEMA, schema, value);
+    }
+
+    /**
+     * Convert the specified value to an {@link Type#ARRAY} value. If the value is a string representation of an array, this method
+     * will parse the string and its elements to infer the schemas for those elements. Thus, this method supports
+     * arrays of other primitives and structured types. If the value is already an array (or list), this method simply casts and
+     * returns it.
+     *
+     * 
            This method currently does not use the schema, though it may be used in the future.
            
+     *
+     * @param schema the schema for the value; may be null
+     * @param value  the value to be converted; may be null
+     * @return the representation as a list, or null if the supplied value was null
+     * @throws DataException if the value cannot be converted to a list value
+     */
+    public static List convertToList(Schema schema, Object value) {
+        return (List) convertTo(ARRAY_SELECTOR_SCHEMA, schema, value);
+    }
+
+    /**
+     * Convert the specified value to an {@link Type#MAP} value. If the value is a string representation of a map, this method
+     * will parse the string and its entries to infer the schemas for those entries. Thus, this method supports
+     * maps with primitives and structured keys and values. If the value is already a map, this method simply casts and returns it.
+     *
+     * 
            This method currently does not use the schema, though it may be used in the future.
            
+     *
+     * @param schema the schema for the value; may be null
+     * @param value  the value to be converted; may be null
+     * @return the representation as a map, or null if the supplied value was null
+     * @throws DataException if the value cannot be converted to a map value
+     */
+    public static Map convertToMap(Schema schema, Object value) {
+        return (Map) convertTo(MAP_SELECTOR_SCHEMA, schema, value);
+    }
+
+    /**
+     * Convert the specified value to an {@link Type#STRUCT} value. Structs cannot be converted from other types, so this method returns
+     * a struct only if the supplied value is a struct. If not a struct, this method throws an exception.
+     *
+     * 
            This method currently does not use the schema, though it may be used in the future.
            
+     *
+     * @param schema the schema for the value; may be null
+     * @param value  the value to be converted; may be null
+     * @return the representation as a struct, or null if the supplied value was null
+     * @throws DataException if the value is not a struct
+     */
+    public static Struct convertToStruct(Schema schema, Object value) {
+        return (Struct) convertTo(STRUCT_SELECTOR_SCHEMA, schema, value);
+    }
+
+    /**
+     * Convert the specified value to an {@link Time#SCHEMA time} value.
+     * Not supplying a schema may limit the ability to convert to the desired type.
+     *
+     * @param schema the schema for the value; may be null
+     * @param value  the value to be converted; may be null
+     * @return the representation as a time, or null if the supplied value was null
+     * @throws DataException if the value cannot be converted to a time value
+     */
+    public static java.util.Date convertToTime(Schema schema, Object value) {
+        return (java.util.Date) convertTo(Time.SCHEMA, schema, value);
+    }
+
+    /**
+     * Convert the specified value to an {@link Date#SCHEMA date} value.
+     * Not supplying a schema may limit the ability to convert to the desired type.
+     *
+     * @param schema the schema for the value; may be null
+     * @param value  the value to be converted; may be null
+     * @return the representation as a date, or null if the supplied value was null
+     * @throws DataException if the value cannot be converted to a date value
+     */
+    public static java.util.Date convertToDate(Schema schema, Object value) {
+        return (java.util.Date) convertTo(Date.SCHEMA, schema, value);
+    }
+
+    /**
+     * Convert the specified value to an {@link Timestamp#SCHEMA timestamp} value.
+     * Not supplying a schema may limit the ability to convert to the desired type.
+     *
+     * @param schema the schema for the value; may be null
+     * @param value  the value to be converted; may be null
+     * @return the representation as a timestamp, or null if the supplied value was null
+     * @throws DataException if the value cannot be converted to a timestamp value
+     */
+    public static java.util.Date convertToTimestamp(Schema schema, Object value) {
+        return (java.util.Date) convertTo(Timestamp.SCHEMA, schema, value);
+    }
+
+    /**
+     * Convert the specified value to an {@link Decimal decimal} value.
+     * Not supplying a schema may limit the ability to convert to the desired type.
+     *
+     * @param schema the schema for the value; may be null
+     * @param value  the value to be converted; may be null
+     * @return the representation as a decimal, or null if the supplied value was null
+     * @throws DataException if the value cannot be converted to a decimal value
+     */
+    public static BigDecimal convertToDecimal(Schema schema, Object value, int scale) {
+        return (BigDecimal) convertTo(Decimal.schema(scale), schema, value);
+    }
+
+    /**
+     * If possible infer a schema for the given value.
+     *
+     * @param value the value whose schema is to be inferred; may be null
+     * @return the inferred schema, or null if the value is null or no schema could be inferred
+     */
+    public static Schema inferSchema(Object value) {
+        if (value instanceof String) {
+            return Schema.STRING_SCHEMA;
+        }
+        if (value instanceof Boolean) {
+            return Schema.BOOLEAN_SCHEMA;
+        }
+        if (value instanceof Byte) {
+            return Schema.INT8_SCHEMA;
+        }
+        if (value instanceof Short) {
+            return Schema.INT16_SCHEMA;
+        }
+        if (value instanceof Integer) {
+            return Schema.INT32_SCHEMA;
+        }
+        if (value instanceof Long) {
+            return Schema.INT64_SCHEMA;
+        }
+        if (value instanceof Float) {
+            return Schema.FLOAT32_SCHEMA;
+        }
+        if (value instanceof Double) {
+            return Schema.FLOAT64_SCHEMA;
+        }
+        if (value instanceof byte[] || value instanceof ByteBuffer) {
+            return Schema.BYTES_SCHEMA;
+        }
+        if (value instanceof List) {
+            List list = (List) value;
+            if (list.isEmpty()) {
+                return null;
+            }
+            SchemaDetector detector = new SchemaDetector();
+            for (Object element : list) {
+                if (!detector.canDetect(element)) {
+                    return null;
+                }
+            }
+            return SchemaBuilder.array(detector.schema()).build();
+        }
+        if (value instanceof Map) {
+            Map map = (Map) value;
+            if (map.isEmpty()) {
+                return null;
+            }
+            SchemaDetector keyDetector = new SchemaDetector();
+            SchemaDetector valueDetector = new SchemaDetector();
+            for (Map.Entry entry : map.entrySet()) {
+                if (!keyDetector.canDetect(entry.getKey()) || !valueDetector.canDetect(entry.getValue())) {
+                    return null;
+                }
+            }
+            return SchemaBuilder.map(keyDetector.schema(), valueDetector.schema()).build();
+        }
+        if (value instanceof Struct) {
+            return ((Struct) value).schema();
+        }
+        return null;
+    }
+
+
+    /**
+     * Parse the specified string representation of a value into its schema and value.
+     *
+     * @param value the string form of the value
+     * @return the schema and value; never null, but whose schema and value may be null
+     * @see #convertToString
+     */
+    public static SchemaAndValue parseString(String value) {
+        if (value == null) {
+            return NULL_SCHEMA_AND_VALUE;
+        }
+        if (value.isEmpty()) {
+            return new SchemaAndValue(Schema.STRING_SCHEMA, value);
+        }
+        Parser parser = new Parser(value);
+        return parse(parser, false);
+    }
+
+    /**
+     * Convert the value to the desired type.
+     *
+     * @param toSchema   the schema for the desired type; may not be null
+     * @param fromSchema the schema for the supplied value; may be null if not known
+     * @return the converted value; never null
+     * @throws DataException if the value could not be converted to the desired type
+     */
+    protected static Object convertTo(Schema toSchema, Schema fromSchema, Object value) throws DataException {
+        if (value == null) {
+            if (toSchema.isOptional()) {
+                return null;
+            }
+            throw new DataException("Unable to convert a null value to a schema that requires a value");
+        }
+        switch (toSchema.type()) {
+            case BYTES:
+                if (Decimal.LOGICAL_NAME.equals(toSchema.name())) {
+                    if (value instanceof ByteBuffer) {
+                        value = Utils.toArray((ByteBuffer) value);
+                    }
+                    if (value instanceof byte[]) {
+                        return Decimal.toLogical(toSchema, (byte[]) value);
+                    }
+                    if (value instanceof BigDecimal) {
+                        return value;
+                    }
+                    if (value instanceof Number) {
+                        // Not already a decimal, so treat it as a double ...
+                        double converted = ((Number) value).doubleValue();
+                        return BigDecimal.valueOf(converted);
+                    }
+                    if (value instanceof String) {
+                        return new BigDecimal(value.toString()).doubleValue();
+                    }
+                }
+                if (value instanceof ByteBuffer) {
+                    return Utils.toArray((ByteBuffer) value);
+                }
+                if (value instanceof byte[]) {
+                    return value;
+                }
+                if (value instanceof BigDecimal) {
+                    return Decimal.fromLogical(toSchema, (BigDecimal) value);
+                }
+                break;
+            case STRING:
+                StringBuilder sb = new StringBuilder();
+                append(sb, value, false);
+                return sb.toString();
+            case BOOLEAN:
+                if (value instanceof Boolean) {
+                    return value;
+                }
+                if (value instanceof String) {
+                    SchemaAndValue parsed = parseString(value.toString());
+                    if (parsed.value() instanceof Boolean) {
+                        return parsed.value();
+                    }
+                }
+                return asLong(value, fromSchema, null) == 0L ? Boolean.FALSE : Boolean.TRUE;
+            case INT8:
+                if (value instanceof Byte) {
+                    return value;
+                }
+                return (byte) asLong(value, fromSchema, null);
+            case INT16:
+                if (value instanceof Short) {
+                    return value;
+                }
+                return (short) asLong(value, fromSchema, null);
+            case INT32:
+                if (Date.LOGICAL_NAME.equals(toSchema.name())) {
+                    if (value instanceof String) {
+                        SchemaAndValue parsed = parseString(value.toString());
+                        value = parsed.value();
+                    }
+                    if (value instanceof java.util.Date) {
+                        if (fromSchema != null) {
+                            String fromSchemaName = fromSchema.name();
+                            if (Date.LOGICAL_NAME.equals(fromSchemaName)) {
+                                return value;
+                            }
+                            if (Timestamp.LOGICAL_NAME.equals(fromSchemaName)) {
+                                // Just get the number of days from this timestamp
+                                long millis = ((java.util.Date) value).getTime();
+                                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);
+                    return Date.toLogical(toSchema, (int) numeric);
+                }
+                if (Time.LOGICAL_NAME.equals(toSchema.name())) {
+                    if (value instanceof String) {
+                        SchemaAndValue parsed = parseString(value.toString());
+                        value = parsed.value();
+                    }
+                    if (value instanceof java.util.Date) {
+                        if (fromSchema != null) {
+                            String fromSchemaName = fromSchema.name();
+                            if (Time.LOGICAL_NAME.equals(fromSchemaName)) {
+                                return value;
+                            }
+                            if (Timestamp.LOGICAL_NAME.equals(fromSchemaName)) {
+                                // Just get the time portion of this timestamp
+                                Calendar calendar = Calendar.getInstance(UTC);
+                                calendar.setTime((java.util.Date) value);
+                                calendar.set(Calendar.YEAR, 1970);
+                                calendar.set(Calendar.MONTH, 0); // Months are zero-based
+                                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);
+                    return Time.toLogical(toSchema, (int) numeric);
+                }
+                if (value instanceof Integer) {
+                    return value;
+                }
+                return (int) asLong(value, fromSchema, null);
+            case INT64:
+                if (Timestamp.LOGICAL_NAME.equals(toSchema.name())) {
+                    if (value instanceof String) {
+                        SchemaAndValue parsed = parseString(value.toString());
+                        value = parsed.value();
+                    }
+                    if (value instanceof java.util.Date) {
+                        java.util.Date date = (java.util.Date) value;
+                        if (fromSchema != null) {
+                            String fromSchemaName = fromSchema.name();
+                            if (Date.LOGICAL_NAME.equals(fromSchemaName)) {
+                                int days = Date.fromLogical(fromSchema, date);
+                                long millis = days * MILLIS_PER_DAY;
+                                return Timestamp.toLogical(toSchema, millis);
+                            }
+                            if (Time.LOGICAL_NAME.equals(fromSchemaName)) {
+                                long millis = Time.fromLogical(fromSchema, date);
+                                return Timestamp.toLogical(toSchema, millis);
+                            }
+                            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);
+                    return Timestamp.toLogical(toSchema, numeric);
+                }
+                if (value instanceof Long) {
+                    return value;
+                }
+                return asLong(value, fromSchema, null);
+            case FLOAT32:
+                if (value instanceof Float) {
+                    return value;
+                }
+                return (float) asDouble(value, fromSchema, null);
+            case FLOAT64:
+                if (value instanceof Double) {
+                    return value;
+                }
+                return asDouble(value, fromSchema, null);
+            case ARRAY:
+                if (value instanceof String) {
+                    SchemaAndValue schemaAndValue = parseString(value.toString());
+                    value = schemaAndValue.value();
+                }
+                if (value instanceof List) {
+                    return value;
+                }
+                break;
+            case MAP:
+                if (value instanceof String) {
+                    SchemaAndValue schemaAndValue = parseString(value.toString());
+                    value = schemaAndValue.value();
+                }
+                if (value instanceof Map) {
+                    return value;
+                }
+                break;
+            case STRUCT:
+                if (value instanceof Struct) {
+                    Struct struct = (Struct) value;
+                    return struct;
+                }
+        }
+        throw new DataException("Unable to convert " + value + " (" + value.getClass() + ") to " + toSchema);
+    }
+
+    /**
+     * Convert the specified value to the desired scalar value type.
+     *
+     * @param value      the value to be converted; may not be null
+     * @param fromSchema the schema for the current value type; may not be null
+     * @param error      any previous error that should be included in an exception message; may be null
+     * @return the long value after conversion; never null
+     * @throws DataException if the value could not be converted to a long
+     */
+    protected static long asLong(Object value, Schema fromSchema, Throwable error) {
+        try {
+            if (value instanceof Number) {
+                Number number = (Number) value;
+                return number.longValue();
+            }
+            if (value instanceof String) {
+                return new BigDecimal(value.toString()).longValue();
+            }
+        } catch (NumberFormatException e) {
+            error = e;
+            // fall through
+        }
+        if (fromSchema != null) {
+            String schemaName = fromSchema.name();
+            if (value instanceof java.util.Date) {
+                if (Date.LOGICAL_NAME.equals(schemaName)) {
+                    return Date.fromLogical(fromSchema, (java.util.Date) value);
+                }
+                if (Time.LOGICAL_NAME.equals(schemaName)) {
+                    return Time.fromLogical(fromSchema, (java.util.Date) value);
+                }
+                if (Timestamp.LOGICAL_NAME.equals(schemaName)) {
+                    return Timestamp.fromLogical(fromSchema, (java.util.Date) value);
+                }
+            }
+            throw new DataException("Unable to convert " + value + " (" + value.getClass() + ") to " + fromSchema, error);
+        }
+        throw new DataException("Unable to convert " + value + " (" + value.getClass() + ") to a number", error);
+    }
+
+    /**
+     * Convert the specified value with the desired floating point type.
+     *
+     * @param value  the value to be converted; may not be null
+     * @param schema the schema for the current value type; may not be null
+     * @param error  any previous error that should be included in an exception message; may be null
+     * @return the double value after conversion; never null
+     * @throws DataException if the value could not be converted to a double
+     */
+    protected static double asDouble(Object value, Schema schema, Throwable error) {
+        try {
+            if (value instanceof Number) {
+                Number number = (Number) value;
+                return number.doubleValue();
+            }
+            if (value instanceof String) {
+                return new BigDecimal(value.toString()).doubleValue();
+            }
+        } catch (NumberFormatException e) {
+            error = e;
+            // fall through
+        }
+        return asLong(value, schema, error);
+    }
+
+    protected static void append(StringBuilder sb, Object value, boolean embedded) {
+        if (value == null) {
+            sb.append(NULL_VALUE);
+        } else if (value instanceof Number) {
+            sb.append(value);
+        } else if (value instanceof Boolean) {
+            sb.append(value);
+        } else if (value instanceof String) {
+            if (embedded) {
+                String escaped = escape((String) value);
+                sb.append('"').append(escaped).append('"');
+            } else {
+                sb.append(value);
+            }
+        } else if (value instanceof byte[]) {
+            value = Base64.getEncoder().encodeToString((byte[]) value);
+            if (embedded) {
+                sb.append('"').append(value).append('"');
+            } else {
+                sb.append(value);
+            }
+        } else if (value instanceof ByteBuffer) {
+            byte[] bytes = Utils.readBytes((ByteBuffer) value);
+            append(sb, bytes, embedded);
+        } else if (value instanceof List) {
+            List list = (List) value;
+            sb.append('[');
+            appendIterable(sb, list.iterator());
+            sb.append(']');
+        } else if (value instanceof Map) {
+            Map map = (Map) value;
+            sb.append('{');
+            appendIterable(sb, map.entrySet().iterator());
+            sb.append('}');
+        } else if (value instanceof Struct) {
+            Struct struct = (Struct) value;
+            Schema schema = struct.schema();
+            boolean first = true;
+            sb.append('{');
+            for (Field field : schema.fields()) {
+                if (first) {
+                    first = false;
+                } else {
+                    sb.append(',');
+                }
+                append(sb, field.name(), true);
+                sb.append(':');
+                append(sb, struct.get(field), true);
+            }
+            sb.append('}');
+        } else if (value instanceof Map.Entry) {
+            Map.Entry entry = (Map.Entry) value;
+            append(sb, entry.getKey(), true);
+            sb.append(':');
+            append(sb, entry.getValue(), true);
+        } else if (value instanceof java.util.Date) {
+            java.util.Date dateValue = (java.util.Date) value;
+            String formatted = dateFormatFor(dateValue).format(dateValue);
+            sb.append(formatted);
+        } else {
+            throw new DataException("Failed to serialize unexpected value type " + value.getClass().getName() + ": " + value);
+        }
+    }
+
+    protected static void appendIterable(StringBuilder sb, Iterator iter) {
+        if (iter.hasNext()) {
+            append(sb, iter.next(), true);
+            while (iter.hasNext()) {
+                sb.append(',');
+                append(sb, iter.next(), true);
+            }
+        }
+    }
+
+    protected static String escape(String value) {
+        String replace1 = TWO_BACKSLASHES.matcher(value).replaceAll("\\\\\\\\");
+        return DOUBLEQOUTE.matcher(replace1).replaceAll("\\\\\"");
+    }
+
+    public static DateFormat dateFormatFor(java.util.Date value) {
+        if (value.getTime() < MILLIS_PER_DAY) {
+            return new SimpleDateFormat(ISO_8601_TIME_FORMAT_PATTERN);
+        }
+        if (value.getTime() % MILLIS_PER_DAY == 0) {
+            return new SimpleDateFormat(ISO_8601_DATE_FORMAT_PATTERN);
+        }
+        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(QUOTE_DELIMITER)) {
+                StringBuilder sb = new StringBuilder();
+                while (parser.hasNext()) {
+                    if (parser.canConsume(QUOTE_DELIMITER)) {
+                        break;
+                    }
+                    sb.append(parser.next());
+                }
+                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 (canParseSingleTokenLiteral(parser, embedded, NULL_VALUE)) {
+            return null;
+        }
+        if (canParseSingleTokenLiteral(parser, embedded, TRUE_LITERAL)) {
+            return TRUE_SCHEMA_AND_VALUE;
+        }
+        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;
+                        if (elementSchema != null) {
+                            listSchema = SchemaBuilder.array(elementSchema).schema();
+                            result = alignListEntriesWithSchema(listSchema, result);
+                        } else {
+                            // Every value is null
+                            listSchema = SchemaBuilder.arrayOfNull().build();
+                        }
+                        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 != 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());
+                }
+                throw new DataException("Array is missing terminating ']': " + parser.original());
+            }
+
+            if (parser.canConsume(MAP_BEGIN_DELIMITER)) {
+                Map result = new LinkedHashMap<>();
+                Schema keySchema = null;
+                Schema valueSchema = null;
+                while (parser.hasNext()) {
+                    if (parser.canConsume(MAP_END_DELIMITER)) {
+                        Schema mapSchema;
+                        if (keySchema != null && valueSchema != null) {
+                            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();
+                        }
+                        return new SchemaAndValue(mapSchema, result);
+                    }
+
+                    if (parser.canConsume(COMMA_DELIMITER)) {
+                        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 '" + 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);
+                }
+                // Missing either a comma or an end delimiter
+                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());
+            }
+        } catch (DataException 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();
+        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 ...
+                BigDecimal decimal = new BigDecimal(token);
+                try {
+                    return new SchemaAndValue(Schema.INT8_SCHEMA, decimal.byteValueExact());
+                } catch (ArithmeticException e) {
+                    // continue
+                }
+                try {
+                    return new SchemaAndValue(Schema.INT16_SCHEMA, decimal.shortValueExact());
+                } catch (ArithmeticException e) {
+                    // continue
+                }
+                try {
+                    return new SchemaAndValue(Schema.INT32_SCHEMA, decimal.intValueExact());
+                } catch (ArithmeticException e) {
+                    // continue
+                }
+                try {
+                    return new SchemaAndValue(Schema.INT64_SCHEMA, decimal.longValueExact());
+                } catch (ArithmeticException e) {
+                    // continue
+                }
+                float fValue = decimal.floatValue();
+                if (fValue != Float.NEGATIVE_INFINITY && fValue != Float.POSITIVE_INFINITY
+                    && decimal.scale() != 0) {
+                    return new SchemaAndValue(Schema.FLOAT32_SCHEMA, fValue);
+                }
+                double dValue = decimal.doubleValue();
+                if (dValue != Double.NEGATIVE_INFINITY && dValue != Double.POSITIVE_INFINITY
+                    && decimal.scale() != 0) {
+                    return new SchemaAndValue(Schema.FLOAT64_SCHEMA, dValue);
+                }
+                Schema schema = Decimal.schema(decimal.scale());
+                return new SchemaAndValue(schema, decimal);
+            } catch (NumberFormatException e) {
+                // can't parse as a number
+            }
+        }
+        if (embedded) {
+            throw new DataException("Failed to parse embedded value");
+        }
+        // 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;
+        }
+        if (previous == null) {
+            return latest.schema();
+        }
+        Schema newSchema = latest.schema();
+        Type previousType = previous.type();
+        Type newType = newSchema.type();
+        if (previousType != newType) {
+            switch (previous.type()) {
+                case INT8:
+                    if (newType == Type.INT16 || newType == Type.INT32 || newType == Type.INT64 || newType == Type.FLOAT32 || newType ==
+                                                                                                                              Type.FLOAT64) {
+                        return newSchema;
+                    }
+                    break;
+                case INT16:
+                    if (newType == Type.INT8) {
+                        return previous;
+                    }
+                    if (newType == Type.INT32 || newType == Type.INT64 || newType == Type.FLOAT32 || newType == Type.FLOAT64) {
+                        return newSchema;
+                    }
+                    break;
+                case INT32:
+                    if (newType == Type.INT8 || newType == Type.INT16) {
+                        return previous;
+                    }
+                    if (newType == Type.INT64 || newType == Type.FLOAT32 || newType == Type.FLOAT64) {
+                        return newSchema;
+                    }
+                    break;
+                case INT64:
+                    if (newType == Type.INT8 || newType == Type.INT16 || newType == Type.INT32) {
+                        return previous;
+                    }
+                    if (newType == Type.FLOAT32 || newType == Type.FLOAT64) {
+                        return newSchema;
+                    }
+                    break;
+                case FLOAT32:
+                    if (newType == Type.INT8 || newType == Type.INT16 || newType == Type.INT32 || newType == Type.INT64) {
+                        return previous;
+                    }
+                    if (newType == Type.FLOAT64) {
+                        return newSchema;
+                    }
+                    break;
+                case FLOAT64:
+                    if (newType == Type.INT8 || newType == Type.INT16 || newType == Type.INT32 || newType == Type.INT64 || newType ==
+                                                                                                                           Type.FLOAT32) {
+                        return previous;
+                    }
+                    break;
+            }
+            return null;
+        }
+        if (previous.isOptional() == newSchema.isOptional()) {
+            // Use the optional one
+            return previous.isOptional() ? previous : newSchema;
+        }
+        if (!previous.equals(newSchema)) {
+            return null;
+        }
+        return previous;
+    }
+
+    protected static List alignListEntriesWithSchema(Schema schema, List input) {
+        Schema valueSchema = schema.valueSchema();
+        List result = new ArrayList<>();
+        for (Object value : input) {
+            Object newValue = convertTo(valueSchema, null, value);
+            result.add(newValue);
+        }
+        return result;
+    }
+
+    protected static Map alignMapKeysAndValuesWithSchema(Schema mapSchema, Map input) {
+        Schema keySchema = mapSchema.keySchema();
+        Schema valueSchema = mapSchema.valueSchema();
+        Map result = new LinkedHashMap<>();
+        for (Map.Entry entry : input.entrySet()) {
+            Object newKey = convertTo(keySchema, null, entry.getKey());
+            Object newValue = convertTo(valueSchema, null, entry.getValue());
+            result.put(newKey, newValue);
+        }
+        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;
+
+        public SchemaDetector() {
+        }
+
+        public boolean canDetect(Object value) {
+            if (value == null) {
+                optional = true;
+                return true;
+            }
+            Schema schema = inferSchema(value);
+            if (schema == null) {
+                return false;
+            }
+            if (knownType == null) {
+                knownType = schema.type();
+            } else if (knownType != schema.type()) {
+                return false;
+            }
+            return true;
+        }
+
+        public Schema schema() {
+            SchemaBuilder builder = SchemaBuilder.type(knownType);
+            if (optional) {
+                builder.optional();
+            }
+            return builder.schema();
+        }
+    }
+
+    protected static class Parser {
+        private final String original;
+        private final CharacterIterator iter;
+        private String nextToken = null;
+        private String previousToken = null;
+
+        public Parser(String original) {
+            this.original = original;
+            this.iter = new StringCharacterIterator(this.original);
+        }
+
+        public int position() {
+            return iter.getIndex();
+        }
+
+        public int mark() {
+            return iter.getIndex() - (nextToken != null ? nextToken.length() : 0);
+        }
+
+        public void rewindTo(int position) {
+            iter.setIndex(position);
+            nextToken = null;
+            previousToken = null;
+        }
+
+        public String original() {
+            return original;
+        }
+
+        public boolean hasNext() {
+            return nextToken != null || canConsumeNextToken();
+        }
+
+        protected boolean canConsumeNextToken() {
+            return iter.getEndIndex() > iter.getIndex();
+        }
+
+        public String next() {
+            if (nextToken != null) {
+                previousToken = nextToken;
+                nextToken = null;
+            } else {
+                previousToken = consumeNextToken();
+            }
+            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();
+            char c = iter.current();
+            while (canConsumeNextToken()) {
+                switch (c) {
+                    case '\\':
+                        escaped = !escaped;
+                        break;
+                    case ':':
+                    case ',':
+                    case '{':
+                    case '}':
+                    case '[':
+                    case ']':
+                    case '\"':
+                        if (!escaped) {
+                            if (start < iter.getIndex()) {
+                                // Return the previous token
+                                return original.substring(start, iter.getIndex());
+                            }
+                            // Consume and return this delimiter as a token
+                            iter.next();
+                            return original.substring(start, start + 1);
+                        }
+                        // escaped, so continue
+                        escaped = false;
+                        break;
+                    default:
+                        // If escaped, then we don't care what was escaped
+                        escaped = false;
+                        break;
+                }
+                c = iter.next();
+            }
+            return original.substring(start, iter.getIndex());
+        }
+
+        public String previous() {
+            return previousToken;
+        }
+
+        public boolean canConsume(String expected) {
+            return canConsume(expected, true);
+        }
+
+        public boolean canConsume(String expected, boolean ignoreLeadingAndTrailingWhitespace) {
+            if (isNext(expected, ignoreLeadingAndTrailingWhitespace)) {
+                // consume this token ...
+                nextToken = null;
+                return true;
+            }
+            return false;
+        }
+
+        protected boolean isNext(String expected, boolean ignoreLeadingAndTrailingWhitespace) {
+            if (nextToken == null) {
+                if (!hasNext()) {
+                    return false;
+                }
+                // There's another token, so consume it
+                nextToken = consumeNextToken();
+            }
+            if (ignoreLeadingAndTrailingWhitespace) {
+                while (nextToken.trim().isEmpty() && canConsumeNextToken()) {
+                    nextToken = consumeNextToken();
+                }
+            }
+            return ignoreLeadingAndTrailingWhitespace
+                ? nextToken.trim().equals(expected)
+                : nextToken.equals(expected);
+        }
+    }
+}
diff --git a/connect/api/src/main/java/org/apache/kafka/connect/header/ConnectHeader.java b/connect/api/src/main/java/org/apache/kafka/connect/header/ConnectHeader.java
new file mode 100644
index 0000000000000..3b9f3470f6c14
--- /dev/null
+++ b/connect/api/src/main/java/org/apache/kafka/connect/header/ConnectHeader.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.connect.header;
+
+import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.data.SchemaAndValue;
+import org.apache.kafka.connect.data.Struct;
+
+import java.util.Objects;
+
+/**
+ * A {@link Header} implementation.
+ */
+class ConnectHeader implements Header {
+
+    private static final SchemaAndValue NULL_SCHEMA_AND_VALUE = new SchemaAndValue(null, null);
+
+    private final String key;
+    private final SchemaAndValue schemaAndValue;
+
+    protected ConnectHeader(String key, SchemaAndValue schemaAndValue) {
+        Objects.requireNonNull(key, "Null header keys are not permitted");
+        this.key = key;
+        this.schemaAndValue = schemaAndValue != null ? schemaAndValue : NULL_SCHEMA_AND_VALUE;
+    }
+
+    @Override
+    public String key() {
+        return key;
+    }
+
+    @Override
+    public Object value() {
+        return schemaAndValue.value();
+    }
+
+    @Override
+    public Schema schema() {
+        Schema schema = schemaAndValue.schema();
+        if (schema == null && value() instanceof Struct) {
+            schema = ((Struct) value()).schema();
+        }
+        return schema;
+    }
+
+    @Override
+    public Header rename(String key) {
+        Objects.requireNonNull(key, "Null header keys are not permitted");
+        if (this.key.equals(key)) {
+            return this;
+        }
+        return new ConnectHeader(key, schemaAndValue);
+    }
+
+    @Override
+    public Header with(Schema schema, Object value) {
+        return new ConnectHeader(key, new SchemaAndValue(schema, value));
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(key, schemaAndValue);
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        if (obj == this) {
+            return true;
+        }
+        if (obj instanceof Header) {
+            Header that = (Header) obj;
+            return Objects.equals(this.key, that.key()) && Objects.equals(this.schema(), that.schema()) && Objects.equals(this.value(),
+                                                                                                                          that.value());
+        }
+        return false;
+    }
+
+    @Override
+    public String toString() {
+        return "ConnectHeader(key=" + key + ", value=" + value() + ", schema=" + schema() + ")";
+    }
+}
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
new file mode 100644
index 0000000000000..5c37ddc5e58b4
--- /dev/null
+++ b/connect/api/src/main/java/org/apache/kafka/connect/header/ConnectHeaders.java
@@ -0,0 +1,497 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.header;
+
+import org.apache.kafka.common.utils.AbstractIterator;
+import org.apache.kafka.connect.data.Date;
+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.SchemaAndValue;
+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.errors.DataException;
+
+import java.math.BigDecimal;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.ListIterator;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+
+/**
+ * A basic {@link Headers} implementation.
+ */
+public class ConnectHeaders implements Headers {
+
+    private static final int EMPTY_HASH = Objects.hash(new LinkedList<>());
+
+    private LinkedList
             headers;
+
+    public ConnectHeaders() {
+    }
+
+    public ConnectHeaders(Iterable
             original) {
+        if (original == null) {
+            return;
+        }
+        if (original instanceof ConnectHeaders) {
+            ConnectHeaders originalHeaders = (ConnectHeaders) original;
+            if (!originalHeaders.isEmpty()) {
+                headers = new LinkedList<>(originalHeaders.headers);
+            }
+        } else {
+            headers = new LinkedList<>();
+            for (Header header : original) {
+                Objects.requireNonNull(header, "Unable to add a null header.");
+                headers.add(header);
+            }
+        }
+    }
+
+    @Override
+    public int size() {
+        return headers == null ? 0 : headers.size();
+    }
+
+    @Override
+    public boolean isEmpty() {
+        return headers == null || headers.isEmpty();
+    }
+
+    @Override
+    public Headers clear() {
+        if (headers != null) {
+            headers.clear();
+        }
+        return this;
+    }
+
+    @Override
+    public Headers add(Header header) {
+        Objects.requireNonNull(header, "Unable to add a null header.");
+        if (headers == null) {
+            headers = new LinkedList<>();
+        }
+        headers.add(header);
+        return this;
+    }
+
+    protected Headers addWithoutValidating(String key, Object value, Schema schema) {
+        return add(new ConnectHeader(key, new SchemaAndValue(schema, value)));
+    }
+
+    @Override
+    public Headers add(String key, SchemaAndValue schemaAndValue) {
+        checkSchemaMatches(schemaAndValue);
+        return add(new ConnectHeader(key, schemaAndValue != null ? schemaAndValue : SchemaAndValue.NULL));
+    }
+
+    @Override
+    public Headers add(String key, Object value, Schema schema) {
+        return add(key, value != null || schema != null ? new SchemaAndValue(schema, value) : SchemaAndValue.NULL);
+    }
+
+    @Override
+    public Headers addString(String key, String value) {
+        return addWithoutValidating(key, value, value != null ? Schema.STRING_SCHEMA : Schema.OPTIONAL_STRING_SCHEMA);
+    }
+
+    @Override
+    public Headers addBytes(String key, byte[] value) {
+        return addWithoutValidating(key, value, value != null ? Schema.BYTES_SCHEMA : Schema.OPTIONAL_BYTES_SCHEMA);
+    }
+
+    @Override
+    public Headers addBoolean(String key, boolean value) {
+        return addWithoutValidating(key, value, Schema.BOOLEAN_SCHEMA);
+    }
+
+    @Override
+    public Headers addByte(String key, byte value) {
+        return addWithoutValidating(key, value, Schema.INT8_SCHEMA);
+    }
+
+    @Override
+    public Headers addShort(String key, short value) {
+        return addWithoutValidating(key, value, Schema.INT16_SCHEMA);
+    }
+
+    @Override
+    public Headers addInt(String key, int value) {
+        return addWithoutValidating(key, value, Schema.INT32_SCHEMA);
+    }
+
+    @Override
+    public Headers addLong(String key, long value) {
+        return addWithoutValidating(key, value, Schema.INT64_SCHEMA);
+    }
+
+    @Override
+    public Headers addFloat(String key, float value) {
+        return addWithoutValidating(key, value, Schema.FLOAT32_SCHEMA);
+    }
+
+    @Override
+    public Headers addDouble(String key, double value) {
+        return addWithoutValidating(key, value, Schema.FLOAT64_SCHEMA);
+    }
+
+    @Override
+    public Headers addList(String key, List value, Schema schema) {
+        if (value == null) {
+            return add(key, null, null);
+        }
+        checkSchemaType(schema, Type.ARRAY);
+        return addWithoutValidating(key, value, schema);
+    }
+
+    @Override
+    public Headers addMap(String key, Map value, Schema schema) {
+        if (value == null) {
+            return add(key, null, null);
+        }
+        checkSchemaType(schema, Type.MAP);
+        return addWithoutValidating(key, value, schema);
+    }
+
+    @Override
+    public Headers addStruct(String key, Struct value) {
+        if (value == null) {
+            return add(key, null, null);
+        }
+        checkSchemaType(value.schema(), Type.STRUCT);
+        return addWithoutValidating(key, value, value.schema());
+    }
+
+    @Override
+    public Headers addDecimal(String key, BigDecimal value) {
+        if (value == null) {
+            return add(key, null, null);
+        }
+        // Check that this is a decimal ...
+        Schema schema = Decimal.schema(value.scale());
+        Decimal.fromLogical(schema, value);
+        return addWithoutValidating(key, value, schema);
+    }
+
+    @Override
+    public Headers addDate(String key, java.util.Date value) {
+        if (value != null) {
+            // Check that this is a date ...
+            Date.fromLogical(Date.SCHEMA, value);
+        }
+        return addWithoutValidating(key, value, Date.SCHEMA);
+    }
+
+    @Override
+    public Headers addTime(String key, java.util.Date value) {
+        if (value != null) {
+            // Check that this is a time ...
+            Time.fromLogical(Time.SCHEMA, value);
+        }
+        return addWithoutValidating(key, value, Time.SCHEMA);
+    }
+
+    @Override
+    public Headers addTimestamp(String key, java.util.Date value) {
+        if (value != null) {
+            // Check that this is a timestamp ...
+            Timestamp.fromLogical(Timestamp.SCHEMA, value);
+        }
+        return addWithoutValidating(key, value, Timestamp.SCHEMA);
+    }
+
+    @Override
+    public Header lastWithName(String key) {
+        checkKey(key);
+        if (headers != null) {
+            ListIterator
             iter = headers.listIterator(headers.size());
+            while (iter.hasPrevious()) {
+                Header header = iter.previous();
+                if (key.equals(header.key())) {
+                    return header;
+                }
+            }
+        }
+        return null;
+    }
+
+    @Override
+    public Iterator
             allWithName(String key) {
+        return new FilterByKeyIterator(iterator(), key);
+    }
+
+    @Override
+    public Iterator
             iterator() {
+        return headers == null ? Collections.emptyIterator() :
+            headers.iterator();
+    }
+
+    @Override
+    public Headers remove(String key) {
+        checkKey(key);
+        if (!isEmpty()) {
+            Iterator
             iterator = iterator();
+            while (iterator.hasNext()) {
+                if (iterator.next().key().equals(key)) {
+                    iterator.remove();
+                }
+            }
+        }
+        return this;
+    }
+
+    @Override
+    public Headers retainLatest() {
+        if (!isEmpty()) {
+            Set keys = new HashSet<>();
+            ListIterator
             iter = headers.listIterator(headers.size());
+            while (iter.hasPrevious()) {
+                Header header = iter.previous();
+                String key = header.key();
+                if (!keys.add(key)) {
+                    iter.remove();
+                }
+            }
+        }
+        return this;
+    }
+
+    @Override
+    public Headers retainLatest(String key) {
+        checkKey(key);
+        if (!isEmpty()) {
+            boolean found = false;
+            ListIterator
             iter = headers.listIterator(headers.size());
+            while (iter.hasPrevious()) {
+                String headerKey = iter.previous().key();
+                if (key.equals(headerKey)) {
+                    if (found)
+                        iter.remove();
+                    found = true;
+                }
+            }
+        }
+        return this;
+    }
+
+    @Override
+    public Headers apply(String key, HeaderTransform transform) {
+        checkKey(key);
+        if (!isEmpty()) {
+            ListIterator
             iter = headers.listIterator();
+            while (iter.hasNext()) {
+                Header orig = iter.next();
+                if (orig.key().equals(key)) {
+                    Header updated = transform.apply(orig);
+                    if (updated != null) {
+                        iter.set(updated);
+                    } else {
+                        iter.remove();
+                    }
+                }
+            }
+        }
+        return this;
+    }
+
+    @Override
+    public Headers apply(HeaderTransform transform) {
+        if (!isEmpty()) {
+            ListIterator
             iter = headers.listIterator();
+            while (iter.hasNext()) {
+                Header orig = iter.next();
+                Header updated = transform.apply(orig);
+                if (updated != null) {
+                    iter.set(updated);
+                } else {
+                    iter.remove();
+                }
+            }
+        }
+        return this;
+    }
+
+    @Override
+    public int hashCode() {
+        return isEmpty() ? EMPTY_HASH : Objects.hash(headers);
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        if (obj == this) {
+            return true;
+        }
+        if (obj instanceof Headers) {
+            Headers that = (Headers) obj;
+            Iterator
             thisIter = this.iterator();
+            Iterator
             thatIter = that.iterator();
+            while (thisIter.hasNext() && thatIter.hasNext()) {
+                if (!Objects.equals(thisIter.next(), thatIter.next()))
+                    return false;
+            }
+            return !thisIter.hasNext() && !thatIter.hasNext();
+        }
+        return false;
+    }
+
+    @Override
+    public String toString() {
+        return "ConnectHeaders(headers=" + (headers != null ? headers : "") + ")";
+    }
+
+    @Override
+    public ConnectHeaders duplicate() {
+        return new ConnectHeaders(this);
+    }
+
+    /**
+     * Check that the key is not null
+     *
+     * @param key the key; may not be null
+     * @throws NullPointerException if the supplied key is null
+     */
+    private void checkKey(String key) {
+        Objects.requireNonNull(key, "Header key cannot be null");
+    }
+
+    /**
+     * Check the {@link Schema#type() schema's type} matches the specified type.
+     *
+     * @param schema the schema; never null
+     * @param type   the expected type
+     * @throws DataException if the schema's type does not match the expected type
+     */
+    private void checkSchemaType(Schema schema, Type type) {
+        if (schema.type() != type) {
+            throw new DataException("Expecting " + type + " but instead found " + schema.type());
+        }
+    }
+
+    /**
+     * Check that the value and its schema are compatible.
+     *
+     * @param schemaAndValue the schema and value pair
+     * @throws DataException if the schema is not compatible with the value
+     */
+    // visible for testing
+    void checkSchemaMatches(SchemaAndValue schemaAndValue) {
+        if (schemaAndValue != null) {
+            Schema schema = schemaAndValue.schema();
+            if (schema == null)
+                return;
+            schema = schema.schema(); // in case a SchemaBuilder is used
+            Object value = schemaAndValue.value();
+            if (value == null && !schema.isOptional()) {
+                throw new DataException("A null value requires an optional schema but was " + schema);
+            }
+            if (value != null) {
+                switch (schema.type()) {
+                    case BYTES:
+                        if (value instanceof ByteBuffer)
+                            return;
+                        if (value instanceof byte[])
+                            return;
+                        if (value instanceof BigDecimal && Decimal.LOGICAL_NAME.equals(schema.name()))
+                            return;
+                        break;
+                    case STRING:
+                        if (value instanceof String)
+                            return;
+                        break;
+                    case BOOLEAN:
+                        if (value instanceof Boolean)
+                            return;
+                        break;
+                    case INT8:
+                        if (value instanceof Byte)
+                            return;
+                        break;
+                    case INT16:
+                        if (value instanceof Short)
+                            return;
+                        break;
+                    case INT32:
+                        if (value instanceof Integer)
+                            return;
+                        if (value instanceof java.util.Date && Date.LOGICAL_NAME.equals(schema.name()))
+                            return;
+                        if (value instanceof java.util.Date && Time.LOGICAL_NAME.equals(schema.name()))
+                            return;
+                        break;
+                    case INT64:
+                        if (value instanceof Long)
+                            return;
+                        if (value instanceof java.util.Date && Timestamp.LOGICAL_NAME.equals(schema.name()))
+                            return;
+                        break;
+                    case FLOAT32:
+                        if (value instanceof Float)
+                            return;
+                        break;
+                    case FLOAT64:
+                        if (value instanceof Double)
+                            return;
+                        break;
+                    case ARRAY:
+                        if (value instanceof List)
+                            return;
+                        break;
+                    case MAP:
+                        if (value instanceof Map)
+                            return;
+                        break;
+                    case STRUCT:
+                        if (value instanceof Struct)
+                            return;
+                        break;
+                }
+                throw new DataException("The value " + value + " is not compatible with the schema " + schema);
+            }
+        }
+    }
+
+    private static final class FilterByKeyIterator extends AbstractIterator
             {
+
+        private final Iterator
             original;
+        private final String key;
+
+        private FilterByKeyIterator(Iterator
             original, String key) {
+            this.original = original;
+            this.key = key;
+        }
+
+        @Override
+        protected Header makeNext() {
+            while (original.hasNext()) {
+                Header header = original.next();
+                if (!header.key().equals(key)) {
+                    continue;
+                }
+                return header;
+            }
+            return this.allDone();
+        }
+    }
+}
diff --git a/connect/api/src/main/java/org/apache/kafka/connect/header/Header.java b/connect/api/src/main/java/org/apache/kafka/connect/header/Header.java
new file mode 100644
index 0000000000000..a70d1dc77ea6d
--- /dev/null
+++ b/connect/api/src/main/java/org/apache/kafka/connect/header/Header.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.connect.header;
+
+import org.apache.kafka.connect.data.Schema;
+
+/**
+ * A {@link Header} is a key-value pair, and multiple headers can be included with the key, value, and timestamp in each Kafka message.
+ * If the value contains schema information, then the header will have a non-null {@link #schema() schema}.
+ * 
            
+ * This is an immutable interface.
+ */
+public interface Header {
+
+    /**
+     * The header's key, which is not necessarily unique within the set of headers on a Kafka message.
+     *
+     * @return the header's key; never null
+     */
+    String key();
+
+    /**
+     * Return the {@link Schema} associated with this header, if there is one. Not all headers will have schemas.
+     *
+     * @return the header's schema, or null if no schema is associated with this header
+     */
+    Schema schema();
+
+    /**
+     * Get the header's value as deserialized by Connect's header converter.
+     *
+     * @return the deserialized object representation of the header's value; may be null
+     */
+    Object value();
+
+    /**
+     * Return a new {@link Header} object that has the same key but with the supplied value.
+     *
+     * @param schema the schema for the new value; may be null
+     * @param value  the new value
+     * @return the new {@link Header}; never null
+     */
+    Header with(Schema schema, Object value);
+
+    /**
+     * Return a new {@link Header} object that has the same schema and value but with the supplied key.
+     *
+     * @param key the key for the new header; may not be null
+     * @return the new {@link Header}; never null
+     */
+    Header rename(String key);
+}
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
new file mode 100644
index 0000000000000..d7bd77952cf24
--- /dev/null
+++ b/connect/api/src/main/java/org/apache/kafka/connect/header/Headers.java
@@ -0,0 +1,308 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.header;
+
+import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.data.SchemaAndValue;
+import org.apache.kafka.connect.data.Struct;
+import org.apache.kafka.connect.errors.DataException;
+
+import java.math.BigDecimal;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * A mutable ordered collection of {@link Header} objects. Note that multiple headers may have the same {@link Header#key() key}.
+ */
+public interface Headers extends Iterable
             {
+
+    /**
+     * Get the number of headers in this object.
+     *
+     * @return the number of headers; never negative
+     */
+    int size();
+
+    /**
+     * Determine whether this object has no headers.
+     *
+     * @return true if there are no headers, or false if there is at least one header
+     */
+    boolean isEmpty();
+
+    /**
+     * Get the collection of {@link Header} objects whose {@link Header#key() keys} all match the specified key.
+     *
+     * @param key the key; may not be null
+     * @return the iterator over headers with the specified key; may be null if there are no headers with the specified key
+     */
+    Iterator
             allWithName(String key);
+
+    /**
+     * Return the last {@link Header} with the specified key.
+     *
+     * @param key the key for the header; may not be null
+     * @return the last Header, or null if there are no headers with the specified key
+     */
+    Header lastWithName(String key);
+
+    /**
+     * Add the given {@link Header} to this collection.
+     *
+     * @param header the header; may not be null
+     * @return this object to facilitate chaining multiple methods; never null
+     */
+    Headers add(Header header);
+
+    /**
+     * Add to this collection a {@link Header} with the given key and value.
+     *
+     * @param key            the header's key; may not be null
+     * @param schemaAndValue the {@link SchemaAndValue} for the header; may be null
+     * @return this object to facilitate chaining multiple methods; never null
+     */
+    Headers add(String key, SchemaAndValue schemaAndValue);
+
+    /**
+     * Add to this collection a {@link Header} with the given key and value.
+     *
+     * @param key    the header's key; may not be null
+     * @param value  the header's value; may be null
+     * @param schema the schema for the header's value; may not be null if the value is not null
+     * @return this object to facilitate chaining multiple methods; never null
+     */
+    Headers add(String key, Object value, Schema schema);
+
+    /**
+     * Add to this collection a {@link Header} with the given key and value.
+     *
+     * @param key   the header's key; may not be null
+     * @param value the header's value; may be null
+     * @return this object to facilitate chaining multiple methods; never null
+     */
+    Headers addString(String key, String value);
+
+    /**
+     * Add to this collection a {@link Header} with the given key and value.
+     *
+     * @param key   the header's key; may not be null
+     * @param value the header's value; may be null
+     * @return this object to facilitate chaining multiple methods; never null
+     */
+    Headers addBoolean(String key, boolean value);
+
+    /**
+     * Add to this collection a {@link Header} with the given key and value.
+     *
+     * @param key   the header's key; may not be null
+     * @param value the header's value; may be null
+     * @return this object to facilitate chaining multiple methods; never null
+     */
+    Headers addByte(String key, byte value);
+
+    /**
+     * Add to this collection a {@link Header} with the given key and value.
+     *
+     * @param key   the header's key; may not be null
+     * @param value the header's value; may be null
+     * @return this object to facilitate chaining multiple methods; never null
+     */
+    Headers addShort(String key, short value);
+
+    /**
+     * Add to this collection a {@link Header} with the given key and value.
+     *
+     * @param key   the header's key; may not be null
+     * @param value the header's value; may be null
+     * @return this object to facilitate chaining multiple methods; never null
+     */
+    Headers addInt(String key, int value);
+
+    /**
+     * Add to this collection a {@link Header} with the given key and value.
+     *
+     * @param key   the header's key; may not be null
+     * @param value the header's value; may be null
+     * @return this object to facilitate chaining multiple methods; never null
+     */
+    Headers addLong(String key, long value);
+
+    /**
+     * Add to this collection a {@link Header} with the given key and value.
+     *
+     * @param key   the header's key; may not be null
+     * @param value the header's value; may be null
+     * @return this object to facilitate chaining multiple methods; never null
+     */
+    Headers addFloat(String key, float value);
+
+    /**
+     * Add to this collection a {@link Header} with the given key and value.
+     *
+     * @param key   the header's key; may not be null
+     * @param value the header's value; may be null
+     * @return this object to facilitate chaining multiple methods; never null
+     */
+    Headers addDouble(String key, double value);
+
+    /**
+     * Add to this collection a {@link Header} with the given key and value.
+     *
+     * @param key   the header's key; may not be null
+     * @param value the header's value; may be null
+     * @return this object to facilitate chaining multiple methods; never null
+     */
+    Headers addBytes(String key, byte[] value);
+
+    /**
+     * Add to this collection a {@link Header} with the given key and value.
+     *
+     * @param key    the header's key; may not be null
+     * @param value  the header's value; may be null
+     * @param schema the schema describing the list value; may not be null
+     * @return this object to facilitate chaining multiple methods; never null
+     * @throws DataException if the header's value is invalid
+     */
+    Headers addList(String key, List value, Schema schema);
+
+    /**
+     * Add to this collection a {@link Header} with the given key and value.
+     *
+     * @param key    the header's key; may not be null
+     * @param value  the header's value; may be null
+     * @param schema the schema describing the map value; may not be null
+     * @return this object to facilitate chaining multiple methods; never null
+     * @throws DataException if the header's value is invalid
+     */
+    Headers addMap(String key, Map value, Schema schema);
+
+    /**
+     * Add to this collection a {@link Header} with the given key and value.
+     *
+     * @param key   the header's key; may not be null
+     * @param value the header's value; may be null
+     * @return this object to facilitate chaining multiple methods; never null
+     * @throws DataException if the header's value is invalid
+     */
+    Headers addStruct(String key, Struct value);
+
+    /**
+     * Add to this collection a {@link Header} with the given key and {@link org.apache.kafka.connect.data.Decimal} value.
+     *
+     * @param key   the header's key; may not be null
+     * @param value the header's {@link org.apache.kafka.connect.data.Decimal} value; may be null
+     * @return this object to facilitate chaining multiple methods; never null
+     */
+    Headers addDecimal(String key, BigDecimal value);
+
+    /**
+     * Add to this collection a {@link Header} with the given key and {@link org.apache.kafka.connect.data.Date} value.
+     *
+     * @param key   the header's key; may not be null
+     * @param value the header's {@link org.apache.kafka.connect.data.Date} value; may be null
+     * @return this object to facilitate chaining multiple methods; never null
+     */
+    Headers addDate(String key, java.util.Date value);
+
+    /**
+     * Add to this collection a {@link Header} with the given key and {@link org.apache.kafka.connect.data.Time} value.
+     *
+     * @param key   the header's key; may not be null
+     * @param value the header's {@link org.apache.kafka.connect.data.Time} value; may be null
+     * @return this object to facilitate chaining multiple methods; never null
+     */
+    Headers addTime(String key, java.util.Date value);
+
+    /**
+     * Add to this collection a {@link Header} with the given key and {@link org.apache.kafka.connect.data.Timestamp} value.
+     *
+     * @param key   the header's key; may not be null
+     * @param value the header's {@link org.apache.kafka.connect.data.Timestamp} value; may be null
+     * @return this object to facilitate chaining multiple methods; never null
+     */
+    Headers addTimestamp(String key, java.util.Date value);
+
+    /**
+     * Removes all {@link Header} objects whose {@link Header#key() key} matches the specified key.
+     *
+     * @param key the key; may not be null
+     * @return this object to facilitate chaining multiple methods; never null
+     */
+    Headers remove(String key);
+
+    /**
+     * Removes all but the latest {@link Header} objects whose {@link Header#key() key} matches the specified key.
+     *
+     * @param key the key; may not be null
+     * @return this object to facilitate chaining multiple methods; never null
+     */
+    Headers retainLatest(String key);
+
+    /**
+     * Removes all but the last {@link Header} object with each key.
+     *
+     * @return this object to facilitate chaining multiple methods; never null
+     */
+    Headers retainLatest();
+
+    /**
+     * Removes all headers from this object.
+     *
+     * @return this object to facilitate chaining multiple methods; never null
+     */
+    Headers clear();
+
+    /**
+     * Create a copy of this {@link Headers} object. The new copy will contain all of the same {@link Header} objects as this object.
+     * @return the copy; never null
+     */
+    Headers duplicate();
+
+    /**
+     * Get all {@link Header}s, apply the transform to each and store the result in place of the original.
+     *
+     * @param transform the transform to apply; may not be null
+     * @return this object to facilitate chaining multiple methods; never null
+     * @throws DataException if the header's value is invalid
+     */
+    Headers apply(HeaderTransform transform);
+
+    /**
+     * Get all {@link Header}s with the given key, apply the transform to each and store the result in place of the original.
+     *
+     * @param key       the header's key; may not be null
+     * @param transform the transform to apply; may not be null
+     * @return this object to facilitate chaining multiple methods; never null
+     * @throws DataException if the header's value is invalid
+     */
+    Headers apply(String key, HeaderTransform transform);
+
+    /**
+     * A function to transform the supplied {@link Header}. Implementations will likely need to use {@link Header#with(Schema, Object)}
+     * to create the new instance.
+     */
+    interface HeaderTransform {
+        /**
+         * Transform the given {@link Header} and return the updated {@link Header}.
+         *
+         * @param header the input header; never null
+         * @return the new header, or null if the supplied {@link Header} is to be removed
+         */
+        Header apply(Header header);
+    }
+}
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
new file mode 100644
index 0000000000000..f707b3c3026ed
--- /dev/null
+++ b/connect/api/src/main/java/org/apache/kafka/connect/health/AbstractState.java
@@ -0,0 +1,93 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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;
+
+import java.util.Objects;
+
+/**
+ * Provides the current status along with identifier for Connect worker and tasks.
+ */
+public abstract class AbstractState {
+
+    private final String state;
+    private final String traceMessage;
+    private final String workerId;
+
+    /**
+     * Construct a state for connector or task.
+     *
+     * @param state  the status of connector or task; may not be null or empty
+     * @param workerId  the workerId associated with the connector or the task; may not be null or empty
+     * @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()) {
+            throw new IllegalArgumentException("State must not be null or empty");
+        }
+        if (workerId == null || workerId.trim().isEmpty()) {
+            throw new IllegalArgumentException("Worker ID must not be null or empty");
+        }
+        this.state = state;
+        this.workerId = workerId;
+        this.traceMessage = traceMessage;
+    }
+
+    /**
+     * Provides the current state of the connector or task.
+     *
+     * @return state, never {@code null} or empty
+     */
+    public String state() {
+        return state;
+    }
+
+    /**
+     * The identifier of the worker associated with the connector or the task.
+     *
+     * @return workerId, never {@code null} or empty.
+     */
+    public String workerId() {
+        return workerId;
+    }
+
+    /**
+     * The error message associated with the connector or task.
+     *
+     * @return traceMessage, can be {@code null} or empty.
+     */
+    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
new file mode 100644
index 0000000000000..753ee1a7613cd
--- /dev/null
+++ b/connect/api/src/main/java/org/apache/kafka/connect/health/ConnectClusterState.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.connect.health;
+
+import java.util.Collection;
+import java.util.Map;
+
+/**
+ * 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 {
+
+    /**
+     * Get the names of the connectors currently deployed in this cluster. This is a full list of connectors in the cluster gathered from
+     * the current configuration, which may change over time.
+     *
+     * @return collection of connector names, never {@code null}
+     */
+    Collection connectors();
+
+    /**
+     * Lookup the current health of a connector and its tasks. This provides the current snapshot of health 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 health of the connector for the connector name
+     * @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
new file mode 100644
index 0000000000000..12fa6b76aff1e
--- /dev/null
+++ b/connect/api/src/main/java/org/apache/kafka/connect/health/ConnectorHealth.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.connect.health;
+
+
+import java.util.Map;
+import java.util.Objects;
+
+/**
+ * Provides basic health information about the connector and its tasks.
+ */
+public class ConnectorHealth {
+
+    private final String name;
+    private final ConnectorState connectorState;
+    private final Map tasks;
+    private final ConnectorType type;
+
+
+    public ConnectorHealth(String name,
+                           ConnectorState connectorState,
+                           Map tasks,
+                           ConnectorType type) {
+        if (name == null || name.trim().isEmpty()) {
+            throw new IllegalArgumentException("Connector name is required");
+        }
+        Objects.requireNonNull(connectorState, "connectorState can't be null");
+        Objects.requireNonNull(tasks, "tasks can't be null");
+        Objects.requireNonNull(type, "type can't be null");
+        this.name = name;
+        this.connectorState = connectorState;
+        this.tasks = tasks;
+        this.type = type;
+    }
+
+    /**
+     * Provides the name of the connector.
+     *
+     * @return name, never {@code null} or empty
+     */
+    public String name() {
+        return name;
+    }
+
+    /**
+     * Provides the current state of the connector.
+     *
+     * @return the connector state, never {@code null}
+     */
+    public ConnectorState connectorState() {
+        return connectorState;
+    }
+
+    /**
+     * Provides the current state of the connector tasks.
+     *
+     * @return the state for each task ID; never {@code null}
+     */
+    public Map tasksState() {
+        return tasks;
+    }
+
+    /**
+     * Provides the type of the connector.
+     *
+     * @return type, never {@code null}
+     */
+    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
new file mode 100644
index 0000000000000..63044265bb999
--- /dev/null
+++ b/connect/api/src/main/java/org/apache/kafka/connect/health/ConnectorState.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.connect.health;
+
+/**
+ * Describes the status, worker ID, and any errors associated with a connector.
+ */
+public class ConnectorState extends AbstractState {
+
+    /**
+     * Provides an instance of the ConnectorState.
+     *
+     * @param state - the status of connector, may not be {@code null} or empty
+     * @param workerId - the workerId associated with the connector, may not be {@code null} or empty
+     * @param traceMessage - any error message associated with the connector, may be {@code null} or empty
+     */
+    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/ConnectorType.java b/connect/api/src/main/java/org/apache/kafka/connect/health/ConnectorType.java
new file mode 100644
index 0000000000000..fa9db6f6b60b6
--- /dev/null
+++ b/connect/api/src/main/java/org/apache/kafka/connect/health/ConnectorType.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.health;
+
+import java.util.Locale;
+
+/**
+ * Enum definition that identifies the type of the connector.
+ */
+public enum ConnectorType {
+    /**
+     * Identifies a source connector
+     */
+    SOURCE,
+    /**
+     * Identifies a sink connector
+     */
+    SINK,
+    /**
+     * Identifies a connector whose type could not be inferred
+     */
+    UNKNOWN;
+
+    @Override
+    public String toString() {
+        return super.toString().toLowerCase(Locale.ROOT);
+    }
+
+}
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
new file mode 100644
index 0000000000000..ae78a5f3af990
--- /dev/null
+++ b/connect/api/src/main/java/org/apache/kafka/connect/health/TaskState.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.health;
+
+import java.util.Objects;
+
+/**
+ * Describes the state, IDs, and any errors of a connector task.
+ */
+public class TaskState extends AbstractState {
+
+    private final int taskId;
+
+    /**
+     * Provides an instance of {@link TaskState}.
+     *
+     * @param taskId   the id associated with the connector task
+     * @param state    the status of the task, may not be {@code null} or empty
+     * @param workerId id of the worker the task is associated with, may not be {@code null} or empty
+     * @param trace    error message if that task had failed or errored out, may be {@code null} or empty
+     */
+    public TaskState(int taskId, String state, String workerId, String trace) {
+        super(state, workerId, trace);
+        this.taskId = taskId;
+    }
+
+    /**
+     * Provides the ID of the task.
+     *
+     * @return the task ID
+     */
+    public int taskId() {
+        return taskId;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o)
+            return true;
+        if (o == null || getClass() != o.getClass())
+            return false;
+        if (!super.equals(o))
+            return false;
+        TaskState taskState = (TaskState) o;
+        return taskId == taskState.taskId;
+    }
+
+    @Override
+    public int hashCode() {
+        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/rest/ConnectRestExtension.java b/connect/api/src/main/java/org/apache/kafka/connect/rest/ConnectRestExtension.java
new file mode 100644
index 0000000000000..aa479a3d449ac
--- /dev/null
+++ b/connect/api/src/main/java/org/apache/kafka/connect/rest/ConnectRestExtension.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.connect.rest;
+
+import org.apache.kafka.common.Configurable;
+import org.apache.kafka.connect.components.Versioned;
+import org.apache.kafka.connect.health.ConnectClusterState;
+
+import java.io.Closeable;
+import java.util.Map;
+
+/**
+ * A plugin interface to allow registration of new JAX-RS resources like Filters, REST endpoints, providers, etc. The implementations will
+ * be discovered using the standard Java {@link java.util.ServiceLoader} mechanism by  Connect's plugin class loading mechanism.
+ *
+ * 
            The extension class(es) must be packaged as a plugin, with one JAR containing the implementation classes and a {@code
+ * META-INF/services/org.apache.kafka.connect.rest.extension.ConnectRestExtension} file that contains the fully qualified name of the
+ * class(es) that implement the ConnectRestExtension interface. The plugin should also include the JARs of all dependencies except those
+ * already provided by the Connect framework.
+ *
+ * 
            To install into a Connect installation, add a directory named for the plugin and containing the plugin's JARs into a directory that is
+ * on Connect's {@code plugin.path}, and (re)start the Connect worker.
+ *
+ * 
            When the Connect worker process starts up, it will read its configuration and instantiate all of the REST extension implementation
+ * classes that are specified in the `rest.extension.classes` configuration property. Connect will then pass its configuration to each
+ * extension via the {@link Configurable#configure(Map)} method, and will then call {@link #register} with a provided context.
+ *
+ * 
            When the Connect worker shuts down, it will call the extension's {@link #close} method to allow the implementation to release all of
+ * its resources.
+ */
+public interface ConnectRestExtension extends Configurable, Versioned, Closeable {
+
+    /**
+     * ConnectRestExtension implementations can register custom JAX-RS resources via the {@link #register(ConnectRestExtensionContext)}
+     * method. The Connect framework will invoke this method after registering the default Connect resources. If the implementations attempt
+     * to re-register any of the Connect resources, it will be be ignored and will be logged.
+     *
+     * @param restPluginContext The context provides access to JAX-RS {@link javax.ws.rs.core.Configurable} and {@link
+     *                          ConnectClusterState}.The custom JAX-RS resources can be registered via the {@link
+     *                          ConnectRestExtensionContext#configurable()}
+     */
+    void register(ConnectRestExtensionContext restPluginContext);
+}
diff --git a/connect/api/src/main/java/org/apache/kafka/connect/rest/ConnectRestExtensionContext.java b/connect/api/src/main/java/org/apache/kafka/connect/rest/ConnectRestExtensionContext.java
new file mode 100644
index 0000000000000..76085971c379b
--- /dev/null
+++ b/connect/api/src/main/java/org/apache/kafka/connect/rest/ConnectRestExtensionContext.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.connect.rest;
+
+import org.apache.kafka.connect.health.ConnectClusterState;
+
+import javax.ws.rs.core.Configurable;
+
+/**
+ * The interface provides the ability for {@link ConnectRestExtension} implementations to access the JAX-RS
+ * {@link javax.ws.rs.core.Configurable} and cluster state {@link ConnectClusterState}. The implementation for the interface is provided
+ * by the Connect framework.
+ */
+public interface ConnectRestExtensionContext {
+
+    /**
+     * Provides an implementation of {@link javax.ws.rs.core.Configurable} that be used to register JAX-RS resources.
+     *
+     * @return @return the JAX-RS {@link javax.ws.rs.core.Configurable}; never {@code null}
+     */
+    Configurable configurable();
+
+    /**
+     * Provides the cluster state and health information about the connectors and tasks.
+     *
+     * @return the cluster state information; never {@code null}
+     */
+    ConnectClusterState clusterState();
+}
diff --git a/connect/api/src/main/java/org/apache/kafka/connect/sink/ErrantRecordReporter.java b/connect/api/src/main/java/org/apache/kafka/connect/sink/ErrantRecordReporter.java
new file mode 100644
index 0000000000000..a20e1e32f200d
--- /dev/null
+++ b/connect/api/src/main/java/org/apache/kafka/connect/sink/ErrantRecordReporter.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 org.apache.kafka.connect.sink;
+
+import java.util.concurrent.Future;
+import org.apache.kafka.connect.errors.ConnectException;
+
+/**
+ * Component that the sink task can use as it {@link SinkTask#put(java.util.Collection)}.
+ * Reporter of problematic records and the corresponding problems.
+ *
+ * @since 2.6
+ */
+public interface ErrantRecordReporter {
+
+    /**
+     * Report a problematic record and the corresponding error to be written to the sink
+     * connector's dead letter queue (DLQ).
+     *
+     * 
            This call is asynchronous and returns a {@link java.util.concurrent.Future Future}.
+     * Invoking {@link java.util.concurrent.Future#get() get()} on this future will block until the
+     * record has been written or throw any exception that occurred while sending the record.
+     * If you want to simulate a simple blocking call you can call the get() method
+     * immediately.
+     *
+     * Connect guarantees that sink records reported through this reporter will be written to the error topic
+     * before the framework calls the {@link SinkTask#preCommit(java.util.Map)} method and therefore before
+     * committing the consumer offsets. SinkTask implementations can use the Future when stronger guarantees
+     * are required.
+     *
+     * @param record the problematic record; may not be null
+     * @param error  the error capturing the problem with the record; may not be null
+     * @return a future that can be used to block until the record and error are reported
+     *         to the DLQ
+     * @throws ConnectException if the error reporter and DLQ fails to write a reported record
+     * @since 2.6
+     */
+    Future report(SinkRecord record, Throwable error);
+}
diff --git a/connect/api/src/main/java/org/apache/kafka/connect/sink/SinkConnector.java b/connect/api/src/main/java/org/apache/kafka/connect/sink/SinkConnector.java
index 5c5886163e23d..9627571482bc4 100644
--- a/connect/api/src/main/java/org/apache/kafka/connect/sink/SinkConnector.java
+++ b/connect/api/src/main/java/org/apache/kafka/connect/sink/SinkConnector.java
@@ -34,4 +34,9 @@ public abstract class SinkConnector extends Connector {
      */
     public static final String TOPICS_CONFIG = "topics";
 
+    @Override
+    protected SinkConnectorContext context() {
+        return (SinkConnectorContext) context;
+    }
+
 }
diff --git a/connect/api/src/main/java/org/apache/kafka/connect/sink/SinkConnectorContext.java b/connect/api/src/main/java/org/apache/kafka/connect/sink/SinkConnectorContext.java
new file mode 100644
index 0000000000000..5e2b07a9fbb99
--- /dev/null
+++ b/connect/api/src/main/java/org/apache/kafka/connect/sink/SinkConnectorContext.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.connect.sink;
+
+import org.apache.kafka.connect.connector.ConnectorContext;
+
+/**
+ * A context to allow a {@link SinkConnector} to interact with the Kafka Connect runtime.
+ */
+public interface SinkConnectorContext extends ConnectorContext {
+}
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 e03a1f1396f07..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
@@ -19,6 +19,7 @@
 import org.apache.kafka.common.record.TimestampType;
 import org.apache.kafka.connect.connector.ConnectRecord;
 import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.header.Header;
 
 /**
  * SinkRecord is a {@link ConnectRecord} that has been read from Kafka and includes the kafkaOffset of
@@ -38,7 +39,12 @@ public SinkRecord(String topic, int partition, Schema keySchema, Object key, Sch
 
     public SinkRecord(String topic, int partition, Schema keySchema, Object key, Schema valueSchema, Object value, long kafkaOffset,
                       Long timestamp, TimestampType timestampType) {
-        super(topic, partition, keySchema, key, valueSchema, value, timestamp);
+        this(topic, partition, keySchema, key, valueSchema, value, kafkaOffset, timestamp, timestampType, null);
+    }
+
+    public SinkRecord(String topic, int partition, Schema keySchema, Object key, Schema valueSchema, Object value, long kafkaOffset,
+                      Long timestamp, TimestampType timestampType, Iterable
             headers) {
+        super(topic, partition, keySchema, key, valueSchema, value, timestamp, headers);
         this.kafkaOffset = kafkaOffset;
         this.timestampType = timestampType;
     }
@@ -53,7 +59,13 @@ public TimestampType timestampType() {
 
     @Override
     public SinkRecord newRecord(String topic, Integer kafkaPartition, Schema keySchema, Object key, Schema valueSchema, Object value, Long timestamp) {
-        return new SinkRecord(topic, kafkaPartition, keySchema, key, valueSchema, value, kafkaOffset(), timestamp, timestampType);
+        return newRecord(topic, kafkaPartition, keySchema, key, valueSchema, value, timestamp, headers().duplicate());
+    }
+
+    @Override
+    public SinkRecord newRecord(String topic, Integer kafkaPartition, Schema keySchema, Object key, Schema valueSchema, Object value,
+                                Long timestamp, Iterable
             headers) {
+        return new SinkRecord(topic, kafkaPartition, keySchema, key, valueSchema, value, kafkaOffset(), timestamp, timestampType, headers);
     }
 
     @Override
@@ -76,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/sink/SinkTask.java b/connect/api/src/main/java/org/apache/kafka/connect/sink/SinkTask.java
index 1406b30ca8d52..5d308c439625a 100644
--- a/connect/api/src/main/java/org/apache/kafka/connect/sink/SinkTask.java
+++ b/connect/api/src/main/java/org/apache/kafka/connect/sink/SinkTask.java
@@ -169,5 +169,6 @@ public void onPartitionsRevoked(Collection partitions) {
      * commit has completed. Implementations of this method should only need to perform final cleanup operations, such
      * as closing network connections to the sink system.
      */
+    @Override
     public abstract void stop();
 }
diff --git a/connect/api/src/main/java/org/apache/kafka/connect/sink/SinkTaskContext.java b/connect/api/src/main/java/org/apache/kafka/connect/sink/SinkTaskContext.java
index 1e214be4b8c05..c4522c796286e 100644
--- a/connect/api/src/main/java/org/apache/kafka/connect/sink/SinkTaskContext.java
+++ b/connect/api/src/main/java/org/apache/kafka/connect/sink/SinkTaskContext.java
@@ -25,6 +25,16 @@
  * Context passed to SinkTasks, allowing them to access utilities in the Kafka Connect runtime.
  */
 public interface SinkTaskContext {
+
+    /**
+     * Get the Task configuration.  This is the latest configuration and may differ from that passed on startup.
+     *
+     * For example, this method can be used to obtain the latest configuration if an external secret has changed,
+     * and the configuration is using variable references such as those compatible with
+     * {@link org.apache.kafka.common.config.ConfigTransformer}.
+     */
+    Map configs();
+
     /**
      * Reset the consumer offsets for the given topic partitions. SinkTasks should use this if they manage offsets
      * in the sink data store rather than using Kafka consumer offsets. For example, an HDFS connector might record
@@ -85,4 +95,32 @@ public interface SinkTaskContext {
      */
     void requestCommit();
 
+    /**
+     * Get the reporter to which the sink task can report problematic or failed {@link SinkRecord records}
+     * passed to the {@link SinkTask#put(java.util.Collection)} method. When reporting a failed record,
+     * the sink task will receive a {@link java.util.concurrent.Future} that the task can optionally use to wait until
+     * the failed record and exception have been written to Kafka. Note that the result of
+     * this method may be null if this connector has not been configured to use a reporter.
+     *
+     * 
            This method was added in Apache Kafka 2.6. Sink tasks that use this method but want to
+     * maintain backward compatibility so they can also be deployed to older Connect runtimes
+     * should guard the call to this method with a try-catch block, since calling this method will result in a
+     * {@link NoSuchMethodException} or {@link NoClassDefFoundError} when the sink connector is deployed to
+     * Connect runtimes older than Kafka 2.6. For example:
+     * 
            +     *     ErrantRecordReporter reporter;
+     *     try {
+     *         reporter = context.errantRecordReporter();
+     *     } catch (NoSuchMethodError | NoClassDefFoundError e) {
+     *         reporter = null;
+     *     }
+     * 
            
+     *
+     * @return the reporter; null if no error reporter has been configured for the connector
+     * @since 2.6
+     */
+    default ErrantRecordReporter errantRecordReporter() {
+        return null;
+    }
+
 }
diff --git a/connect/api/src/main/java/org/apache/kafka/connect/source/SourceConnector.java b/connect/api/src/main/java/org/apache/kafka/connect/source/SourceConnector.java
index 0ca3b335fc289..6e9694024d334 100644
--- a/connect/api/src/main/java/org/apache/kafka/connect/source/SourceConnector.java
+++ b/connect/api/src/main/java/org/apache/kafka/connect/source/SourceConnector.java
@@ -24,4 +24,8 @@
  */
 public abstract class SourceConnector extends Connector {
 
+    @Override
+    protected SourceConnectorContext context() {
+        return (SourceConnectorContext) context;
+    }
 }
diff --git a/connect/api/src/main/java/org/apache/kafka/connect/source/SourceConnectorContext.java b/connect/api/src/main/java/org/apache/kafka/connect/source/SourceConnectorContext.java
new file mode 100644
index 0000000000000..417fbddd08d46
--- /dev/null
+++ b/connect/api/src/main/java/org/apache/kafka/connect/source/SourceConnectorContext.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.source;
+
+import org.apache.kafka.connect.connector.ConnectorContext;
+import org.apache.kafka.connect.storage.OffsetStorageReader;
+
+/**
+ * A context to allow a {@link SourceConnector} to interact with the Kafka Connect runtime.
+ */
+public interface SourceConnectorContext extends ConnectorContext {
+
+    /**
+     * Returns the {@link OffsetStorageReader} for this SourceConnectorContext.
+     * @return the OffsetStorageReader for this connector.
+     */
+    OffsetStorageReader offsetStorageReader();
+}
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 2f3e5e471c934..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
@@ -18,8 +18,10 @@
 
 import org.apache.kafka.connect.connector.ConnectRecord;
 import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.header.Header;
 
 import java.util.Map;
+import java.util.Objects;
 
 /**
  * 
            
@@ -69,7 +71,15 @@ public SourceRecord(Map sourcePartition, Map sourceOffset,
                         Schema keySchema, Object key,
                         Schema valueSchema, Object value,
                         Long timestamp) {
-        super(topic, partition, keySchema, key, valueSchema, value, timestamp);
+        this(sourcePartition, sourceOffset, topic, partition, keySchema, key, valueSchema, value, timestamp, null);
+    }
+
+    public SourceRecord(Map sourcePartition, Map sourceOffset,
+                        String topic, Integer partition,
+                        Schema keySchema, Object key,
+                        Schema valueSchema, Object value,
+                        Long timestamp, Iterable
             headers) {
+        super(topic, partition, keySchema, key, valueSchema, value, timestamp, headers);
         this.sourcePartition = sourcePartition;
         this.sourceOffset = sourceOffset;
     }
@@ -84,7 +94,13 @@ public SourceRecord(Map sourcePartition, Map sourceOffset,
 
     @Override
     public SourceRecord newRecord(String topic, Integer kafkaPartition, Schema keySchema, Object key, Schema valueSchema, Object value, Long timestamp) {
-        return new SourceRecord(sourcePartition, sourceOffset, topic, kafkaPartition, keySchema, key, valueSchema, value, timestamp);
+        return newRecord(topic, kafkaPartition, keySchema, key, valueSchema, value, timestamp, headers().duplicate());
+    }
+
+    @Override
+    public SourceRecord newRecord(String topic, Integer kafkaPartition, Schema keySchema, Object key, Schema valueSchema, Object value,
+                                  Long timestamp, Iterable
             headers) {
+        return new SourceRecord(sourcePartition, sourceOffset, topic, kafkaPartition, keySchema, key, valueSchema, value, timestamp, headers);
     }
 
     @Override
@@ -98,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 f3e636bb87af5..225b080ee18d0 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;
@@ -43,8 +44,16 @@ public void initialize(SourceTaskContext context) {
     public abstract void start(Map props);
 
     /**
-     * Poll this SourceTask for new records. This method should block if no data is currently
-     * available.
+     * 
            
+     * Poll this source task for new records. If no data is currently available, this method
+     * should block but return control to the caller regularly (by returning {@code null}) in
+     * order for the task to transition to the {@code PAUSED} state if requested to do so.
+     * 
            
+     * 
            
+     * The task will be {@link #stop() stopped} on a separate thread, and when that happens
+     * this method is expected to unblock, quickly finish up any remaining processing, and
+     * return.
+     * 
            
      *
      * @return a list of source records
      */
@@ -75,11 +84,18 @@ public void commit() throws InterruptedException {
      * could set a flag that will force {@link #poll()} to exit immediately and invoke
      * {@link java.nio.channels.Selector#wakeup() wakeup()} to interrupt any ongoing requests.
      */
+    @Override
     public abstract void stop();
 
     /**
      * 
            
-     * 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
@@ -87,10 +103,38 @@ 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
+     * @deprecated Use {@link #commitRecord(SourceRecord, RecordMetadata)} instead.
      */
+    @Deprecated
     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 compatibility
+        commitRecord(record);
+    }
 }
diff --git a/connect/api/src/main/java/org/apache/kafka/connect/source/SourceTaskContext.java b/connect/api/src/main/java/org/apache/kafka/connect/source/SourceTaskContext.java
index 8eec1dfb1386f..ddb0a78718351 100644
--- a/connect/api/src/main/java/org/apache/kafka/connect/source/SourceTaskContext.java
+++ b/connect/api/src/main/java/org/apache/kafka/connect/source/SourceTaskContext.java
@@ -18,11 +18,22 @@
 
 import org.apache.kafka.connect.storage.OffsetStorageReader;
 
+import java.util.Map;
+
 /**
  * SourceTaskContext is provided to SourceTasks to allow them to interact with the underlying
  * runtime.
  */
 public interface SourceTaskContext {
+    /**
+     * Get the Task configuration.  This is the latest configuration and may differ from that passed on startup.
+     *
+     * For example, this method can be used to obtain the latest configuration if an external secret has changed,
+     * and the configuration is using variable references such as those compatible with
+     * {@link org.apache.kafka.common.config.ConfigTransformer}.
+     */
+    Map configs();
+
     /**
      * Get the OffsetStorageReader for this SourceTask.
      */
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..2d2ef4a00cecc 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,24 @@ 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; any changes done to the headers
+     *        are applied to the message sent to the broker
+     * @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 +70,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/main/java/org/apache/kafka/connect/storage/ConverterConfig.java b/connect/api/src/main/java/org/apache/kafka/connect/storage/ConverterConfig.java
new file mode 100644
index 0000000000000..cea69959fc89c
--- /dev/null
+++ b/connect/api/src/main/java/org/apache/kafka/connect/storage/ConverterConfig.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.connect.storage;
+
+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 java.util.Map;
+
+import static org.apache.kafka.common.config.ConfigDef.ValidString.in;
+
+/**
+ * Abstract class that defines the configuration options for {@link Converter} and {@link HeaderConverter} instances.
+ */
+public abstract class ConverterConfig extends AbstractConfig {
+
+    public static final String TYPE_CONFIG = "converter.type";
+    private static final String TYPE_DOC = "How this converter will be used.";
+
+    /**
+     * Create a new {@link ConfigDef} instance containing the configurations defined by ConverterConfig. This can be called by subclasses.
+     *
+     * @return the ConfigDef; never null
+     */
+    public static ConfigDef newConfigDef() {
+        return new ConfigDef().define(TYPE_CONFIG, Type.STRING, ConfigDef.NO_DEFAULT_VALUE,
+                                      in(ConverterType.KEY.getName(), ConverterType.VALUE.getName(), ConverterType.HEADER.getName()),
+                                      Importance.LOW, TYPE_DOC);
+    }
+
+    protected ConverterConfig(ConfigDef configDef, Map props) {
+        super(configDef, props, true);
+    }
+
+    /**
+     * Get the type of converter as defined by the {@link #TYPE_CONFIG} configuration.
+     * @return the converter type; never null
+     */
+    public ConverterType type() {
+        return ConverterType.withName(getString(TYPE_CONFIG));
+    }
+}
diff --git a/connect/api/src/main/java/org/apache/kafka/connect/storage/ConverterType.java b/connect/api/src/main/java/org/apache/kafka/connect/storage/ConverterType.java
new file mode 100644
index 0000000000000..446ff8b962cde
--- /dev/null
+++ b/connect/api/src/main/java/org/apache/kafka/connect/storage/ConverterType.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.connect.storage;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Locale;
+import java.util.Map;
+
+/**
+ * The type of {@link Converter} and {@link HeaderConverter}.
+ */
+public enum ConverterType {
+    KEY,
+    VALUE,
+    HEADER;
+
+    private static final Map NAME_TO_TYPE;
+
+    static {
+        ConverterType[] types = ConverterType.values();
+        Map nameToType = new HashMap<>(types.length);
+        for (ConverterType type : types) {
+            nameToType.put(type.name, type);
+        }
+        NAME_TO_TYPE = Collections.unmodifiableMap(nameToType);
+    }
+
+    /**
+     * Find the ConverterType with the given name, using a case-insensitive match.
+     * @param name the name of the converter type; may be null
+     * @return the matching converter type, or null if the supplied name is null or does not match the name of the known types
+     */
+    public static ConverterType withName(String name) {
+        if (name == null) {
+            return null;
+        }
+        return NAME_TO_TYPE.get(name.toLowerCase(Locale.getDefault()));
+    }
+
+    private String name;
+
+    ConverterType() {
+        this.name = this.name().toLowerCase(Locale.ROOT);
+    }
+
+    public String getName() {
+        return name;
+    }
+}
diff --git a/connect/api/src/main/java/org/apache/kafka/connect/storage/HeaderConverter.java b/connect/api/src/main/java/org/apache/kafka/connect/storage/HeaderConverter.java
new file mode 100644
index 0000000000000..3f9d504205d66
--- /dev/null
+++ b/connect/api/src/main/java/org/apache/kafka/connect/storage/HeaderConverter.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 org.apache.kafka.connect.storage;
+
+import org.apache.kafka.common.Configurable;
+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.header.Header;
+
+import java.io.Closeable;
+
+public interface HeaderConverter extends Configurable, Closeable {
+
+    /**
+     * Convert the header name and byte array value into a {@link Header} object.
+     * @param topic the name of the topic for the record containing the header
+     * @param headerKey the header's key; may not be null
+     * @param value the header's raw value; may be null
+     * @return the {@link SchemaAndValue}; may not be null
+     */
+    SchemaAndValue toConnectHeader(String topic, String headerKey, byte[] value);
+
+    /**
+     * Convert the {@link Header}'s {@link Header#value() value} into its byte array representation.
+     * @param topic the name of the topic for the record containing the header
+     * @param headerKey the header's key; may not be null
+     * @param schema the schema for the header's value; may be null
+     * @param value the header's value to convert; may be null
+     * @return the byte array form of the Header's value; may be null if the value is null
+     */
+    byte[] fromConnectHeader(String topic, String headerKey, Schema schema, Object value);
+
+    /**
+     * Configuration specification for this set of header converters.
+     * @return the configuration specification; may not be null
+     */
+    ConfigDef config();
+}
diff --git a/connect/api/src/main/java/org/apache/kafka/connect/storage/SimpleHeaderConverter.java b/connect/api/src/main/java/org/apache/kafka/connect/storage/SimpleHeaderConverter.java
new file mode 100644
index 0000000000000..69c4b86189878
--- /dev/null
+++ b/connect/api/src/main/java/org/apache/kafka/connect/storage/SimpleHeaderConverter.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 org.apache.kafka.connect.storage;
+
+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.data.Values;
+import org.apache.kafka.connect.errors.DataException;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.nio.charset.Charset;
+import java.nio.charset.StandardCharsets;
+import java.util.Map;
+import java.util.NoSuchElementException;
+
+/**
+ * A {@link HeaderConverter} that serializes header values as strings and that deserializes header values to the most appropriate
+ * numeric, boolean, array, or map representation. Schemas are not serialized, but are inferred upon deserialization when possible.
+ */
+public class SimpleHeaderConverter implements HeaderConverter {
+
+    private static final Logger LOG = LoggerFactory.getLogger(SimpleHeaderConverter.class);
+    private static final ConfigDef CONFIG_DEF = new ConfigDef();
+    private static final SchemaAndValue NULL_SCHEMA_AND_VALUE = new SchemaAndValue(null, null);
+    private static final Charset UTF_8 = StandardCharsets.UTF_8;
+
+    @Override
+    public ConfigDef config() {
+        return CONFIG_DEF;
+    }
+
+    @Override
+    public void configure(Map configs) {
+        // do nothing
+    }
+
+    @Override
+    public SchemaAndValue toConnectHeader(String topic, String headerKey, byte[] value) {
+        if (value == null) {
+            return NULL_SCHEMA_AND_VALUE;
+        }
+        try {
+            String str = new String(value, UTF_8);
+            if (str.isEmpty()) {
+                return new SchemaAndValue(Schema.STRING_SCHEMA, str);
+            }
+            return Values.parseString(str);
+        } catch (NoSuchElementException e) {
+            throw new DataException("Failed to deserialize value for header '" + headerKey + "' on topic '" + topic + "'", e);
+        } catch (Throwable t) {
+            LOG.warn("Failed to deserialize value for header '{}' on topic '{}', so using byte array", headerKey, topic, t);
+            return new SchemaAndValue(Schema.BYTES_SCHEMA, value);
+        }
+    }
+
+    @Override
+    public byte[] fromConnectHeader(String topic, String headerKey, Schema schema, Object value) {
+        if (value == null) {
+            return null;
+        }
+        return Values.convertToString(schema, value).getBytes(UTF_8);
+    }
+
+    @Override
+    public void close() throws IOException {
+        // do nothing
+    }
+}
diff --git a/connect/api/src/main/java/org/apache/kafka/connect/storage/StringConverter.java b/connect/api/src/main/java/org/apache/kafka/connect/storage/StringConverter.java
index 85fef841b0a3d..534cdddfa1331 100644
--- a/connect/api/src/main/java/org/apache/kafka/connect/storage/StringConverter.java
+++ b/connect/api/src/main/java/org/apache/kafka/connect/storage/StringConverter.java
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.connect.storage;
 
+import org.apache.kafka.common.config.ConfigDef;
 import org.apache.kafka.common.errors.SerializationException;
 import org.apache.kafka.common.serialization.StringDeserializer;
 import org.apache.kafka.common.serialization.StringSerializer;
@@ -27,16 +28,19 @@
 import java.util.Map;
 
 /**
- * {@link Converter} implementation that only supports serializing to strings. When converting Kafka Connect data to bytes,
- * the schema will be ignored and {@link Object#toString()} will always be invoked to convert the data to a String.
+ * {@link Converter} and {@link HeaderConverter} implementation that only supports serializing to strings. When converting Kafka Connect
+ * data to bytes, the schema will be ignored and {@link Object#toString()} will always be invoked to convert the data to a String.
  * When converting from bytes to Kafka Connect format, the converter will only ever return an optional string schema and
  * a string or null.
  *
  * Encoding configuration is identical to {@link StringSerializer} and {@link StringDeserializer}, but for convenience
- * this class can also be configured to use the same encoding for both encoding and decoding with the converter.encoding
- * setting.
+ * this class can also be configured to use the same encoding for both encoding and decoding with the
+ * {@link StringConverterConfig#ENCODING_CONFIG converter.encoding} setting.
+ *
+ * This implementation currently does nothing with the topic names or header names.
  */
-public class StringConverter implements Converter {
+public class StringConverter implements Converter, HeaderConverter {
+
     private final StringSerializer serializer = new StringSerializer();
     private final StringDeserializer deserializer = new StringDeserializer();
 
@@ -44,22 +48,32 @@ public StringConverter() {
     }
 
     @Override
-    public void configure(Map configs, boolean isKey) {
-        Map serializerConfigs = new HashMap<>();
-        serializerConfigs.putAll(configs);
-        Map deserializerConfigs = new HashMap<>();
-        deserializerConfigs.putAll(configs);
+    public ConfigDef config() {
+        return StringConverterConfig.configDef();
+    }
 
-        Object encodingValue = configs.get("converter.encoding");
-        if (encodingValue != null) {
-            serializerConfigs.put("serializer.encoding", encodingValue);
-            deserializerConfigs.put("deserializer.encoding", encodingValue);
-        }
+    @Override
+    public void configure(Map configs) {
+        StringConverterConfig conf = new StringConverterConfig(configs);
+        String encoding = conf.encoding();
+
+        Map serializerConfigs = new HashMap<>(configs);
+        Map deserializerConfigs = new HashMap<>(configs);
+        serializerConfigs.put("serializer.encoding", encoding);
+        deserializerConfigs.put("deserializer.encoding", encoding);
 
+        boolean isKey = conf.type() == ConverterType.KEY;
         serializer.configure(serializerConfigs, isKey);
         deserializer.configure(deserializerConfigs, isKey);
     }
 
+    @Override
+    public void configure(Map configs, boolean isKey) {
+        Map conf = new HashMap<>(configs);
+        conf.put(StringConverterConfig.TYPE_CONFIG, isKey ? ConverterType.KEY.getName() : ConverterType.VALUE.getName());
+        configure(conf);
+    }
+
     @Override
     public byte[] fromConnectData(String topic, Schema schema, Object value) {
         try {
@@ -77,4 +91,19 @@ public SchemaAndValue toConnectData(String topic, byte[] value) {
             throw new DataException("Failed to deserialize string: ", e);
         }
     }
+
+    @Override
+    public byte[] fromConnectHeader(String topic, String headerKey, Schema schema, Object value) {
+        return fromConnectData(topic, schema, value);
+    }
+
+    @Override
+    public SchemaAndValue toConnectHeader(String topic, String headerKey, byte[] value) {
+        return toConnectData(topic, value);
+    }
+
+    @Override
+    public void close() {
+        // do nothing
+    }
 }
diff --git a/connect/api/src/main/java/org/apache/kafka/connect/storage/StringConverterConfig.java b/connect/api/src/main/java/org/apache/kafka/connect/storage/StringConverterConfig.java
new file mode 100644
index 0000000000000..26b01332c1972
--- /dev/null
+++ b/connect/api/src/main/java/org/apache/kafka/connect/storage/StringConverterConfig.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.storage;
+
+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.ConfigDef.Width;
+
+import java.util.Map;
+
+/**
+ * Configuration options for {@link StringConverter} instances.
+ */
+public class StringConverterConfig extends ConverterConfig {
+
+    public static final String ENCODING_CONFIG = "converter.encoding";
+    public static final String ENCODING_DEFAULT = "UTF8";
+    private static final String ENCODING_DOC = "The name of the Java character set to use for encoding strings as byte arrays.";
+    private static final String ENCODING_DISPLAY = "Encoding";
+
+    private final static ConfigDef CONFIG;
+
+    static {
+        CONFIG = ConverterConfig.newConfigDef();
+        CONFIG.define(ENCODING_CONFIG, Type.STRING, ENCODING_DEFAULT, Importance.HIGH, ENCODING_DOC, null, -1, Width.MEDIUM,
+                      ENCODING_DISPLAY);
+    }
+
+    public static ConfigDef configDef() {
+        return CONFIG;
+    }
+
+    public StringConverterConfig(Map props) {
+        super(CONFIG, props);
+    }
+
+    /**
+     * Get the string encoding.
+     *
+     * @return the encoding; never null
+     */
+    public String encoding() {
+        return getString(ENCODING_CONFIG);
+    }
+}
diff --git a/connect/api/src/main/java/org/apache/kafka/connect/transforms/Transformation.java b/connect/api/src/main/java/org/apache/kafka/connect/transforms/Transformation.java
index 0803fa880dc57..238a6429d17a9 100644
--- a/connect/api/src/main/java/org/apache/kafka/connect/transforms/Transformation.java
+++ b/connect/api/src/main/java/org/apache/kafka/connect/transforms/Transformation.java
@@ -33,6 +33,11 @@ public interface Transformation> extends Configurable
      * Apply transformation to the {@code record} and return another record object (which may be {@code record} itself) or {@code null},
      * corresponding to a map or filter operation respectively.
      *
+     * A transformation must not mutate objects reachable from the given {@code record}
+     * (including, but not limited to, {@link org.apache.kafka.connect.header.Headers Headers},
+     * {@link org.apache.kafka.connect.data.Struct Structs}, {@code Lists}, and {@code Maps}).
+     * If such objects need to be changed, a new ConnectRecord should be created and returned.
+     *
      * The implementation must be thread-safe.
      */
     R apply(R record);
diff --git a/connect/api/src/main/java/org/apache/kafka/connect/transforms/predicates/Predicate.java b/connect/api/src/main/java/org/apache/kafka/connect/transforms/predicates/Predicate.java
new file mode 100644
index 0000000000000..cc38fc049b5aa
--- /dev/null
+++ b/connect/api/src/main/java/org/apache/kafka/connect/transforms/predicates/Predicate.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 org.apache.kafka.connect.transforms.predicates;
+
+import org.apache.kafka.common.Configurable;
+import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.connect.connector.ConnectRecord;
+
+/**
+ * 
            A predicate on records.
+ * Predicates can be used to conditionally apply a {@link org.apache.kafka.connect.transforms.Transformation}
+ * by configuring the transformation's {@code predicate} (and {@code negate}) configuration parameters.
+ * In particular, the {@code Filter} transformation can be conditionally applied in order to filter
+ * certain records from further processing.
+ *
+ * 
            Implementations of this interface must be public and have a public constructor with no parameters.
+ *
+ * @param  The type of record.
+ */
+public interface Predicate> extends Configurable, AutoCloseable {
+
+    /**
+     * Configuration specification for this predicate.
+     *
+     * @return the configuration definition for this predicate; never null
+     */
+    ConfigDef config();
+
+    /**
+     * Returns whether the given record satisfies this predicate.
+     *
+     * @param record the record to evaluate; may not be null
+     * @return true if the predicate matches, or false otherwise
+     */
+    boolean test(R record);
+
+    @Override
+    void close();
+}
diff --git a/connect/api/src/test/java/org/apache/kafka/connect/connector/ConnectorReconfigurationTest.java b/connect/api/src/test/java/org/apache/kafka/connect/connector/ConnectorReconfigurationTest.java
index 85a6cfae3f1de..6f768f9f2dc29 100644
--- a/connect/api/src/test/java/org/apache/kafka/connect/connector/ConnectorReconfigurationTest.java
+++ b/connect/api/src/test/java/org/apache/kafka/connect/connector/ConnectorReconfigurationTest.java
@@ -29,7 +29,7 @@
 public class ConnectorReconfigurationTest {
 
     @Test
-    public void testDefaultReconfigure() throws Exception {
+    public void testDefaultReconfigure() {
         TestConnector conn = new TestConnector(false);
         conn.reconfigure(Collections.emptyMap());
         assertEquals(conn.stopOrder, 0);
@@ -37,7 +37,7 @@ public void testDefaultReconfigure() throws Exception {
     }
 
     @Test(expected = ConnectException.class)
-    public void testReconfigureStopException() throws Exception {
+    public void testReconfigureStopException() {
         TestConnector conn = new TestConnector(true);
         conn.reconfigure(Collections.emptyMap());
     }
diff --git a/connect/api/src/test/java/org/apache/kafka/connect/connector/ConnectorTest.java b/connect/api/src/test/java/org/apache/kafka/connect/connector/ConnectorTest.java
new file mode 100644
index 0000000000000..7addf8f0f048e
--- /dev/null
+++ b/connect/api/src/test/java/org/apache/kafka/connect/connector/ConnectorTest.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.connector;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import org.junit.Before;
+import org.junit.Test;
+
+public abstract class ConnectorTest {
+
+    protected ConnectorContext context;
+    protected Connector connector;
+    protected AssertableConnector assertableConnector;
+
+    @Before
+    public void beforeEach() {
+        connector = createConnector();
+        context = createContext();
+        assertableConnector = (AssertableConnector) connector;
+    }
+
+    @Test
+    public void shouldInitializeContext() {
+        connector.initialize(context);
+        assertableConnector.assertInitialized();
+        assertableConnector.assertContext(context);
+        assertableConnector.assertTaskConfigs(null);
+    }
+
+    @Test
+    public void shouldInitializeContextWithTaskConfigs() {
+        List> taskConfigs = new ArrayList<>();
+        connector.initialize(context, taskConfigs);
+        assertableConnector.assertInitialized();
+        assertableConnector.assertContext(context);
+        assertableConnector.assertTaskConfigs(taskConfigs);
+    }
+
+    @Test
+    public void shouldStopAndStartWhenReconfigure() {
+        Map props = new HashMap<>();
+        connector.initialize(context);
+        assertableConnector.assertContext(context);
+        assertableConnector.assertStarted(false);
+        assertableConnector.assertStopped(false);
+        connector.reconfigure(props);
+        assertableConnector.assertStarted(true);
+        assertableConnector.assertStopped(true);
+        assertableConnector.assertProperties(props);
+    }
+
+    protected abstract ConnectorContext createContext();
+
+    protected abstract Connector createConnector();
+
+    public interface AssertableConnector {
+
+        void assertContext(ConnectorContext expected);
+
+        void assertInitialized();
+
+        void assertTaskConfigs(List> expectedTaskConfigs);
+
+        void assertStarted(boolean expected);
+
+        void assertStopped(boolean expected);
+
+        void assertProperties(Map expected);
+    }
+}
\ No newline at end of file
diff --git a/connect/api/src/test/java/org/apache/kafka/connect/data/ConnectSchemaTest.java b/connect/api/src/test/java/org/apache/kafka/connect/data/ConnectSchemaTest.java
index 339ef23ca54e3..048784e3335d6 100644
--- a/connect/api/src/test/java/org/apache/kafka/connect/data/ConnectSchemaTest.java
+++ b/connect/api/src/test/java/org/apache/kafka/connect/data/ConnectSchemaTest.java
@@ -268,6 +268,16 @@ public void testArrayEquality() {
         assertNotEquals(s1, differentValueSchema);
     }
 
+    @Test
+    public void testArrayDefaultValueEquality() {
+        ConnectSchema s1 = new ConnectSchema(Schema.Type.ARRAY, false, new String[] {"a", "b"}, null, null, null, null, null, null, SchemaBuilder.int8().build());
+        ConnectSchema s2 = new ConnectSchema(Schema.Type.ARRAY, false, new String[] {"a", "b"}, null, null, null, null, null, null, SchemaBuilder.int8().build());
+        ConnectSchema differentValueSchema = new ConnectSchema(Schema.Type.ARRAY, false, new String[] {"b", "c"}, null, null, null, null, null, null, SchemaBuilder.int8().build());
+
+        assertEquals(s1, s2);
+        assertNotEquals(s1, differentValueSchema);
+    }
+
     @Test
     public void testMapEquality() {
         // Same as testArrayEquality, but for both key and value schemas
diff --git a/connect/api/src/test/java/org/apache/kafka/connect/data/SchemaProjectorTest.java b/connect/api/src/test/java/org/apache/kafka/connect/data/SchemaProjectorTest.java
index 151114e2d2937..fa68c0194a278 100644
--- a/connect/api/src/test/java/org/apache/kafka/connect/data/SchemaProjectorTest.java
+++ b/connect/api/src/test/java/org/apache/kafka/connect/data/SchemaProjectorTest.java
@@ -35,7 +35,7 @@
 public class SchemaProjectorTest {
 
     @Test
-    public void testPrimitiveTypeProjection() throws Exception {
+    public void testPrimitiveTypeProjection() {
         Object projected;
         projected = SchemaProjector.project(Schema.BOOLEAN_SCHEMA, false, Schema.BOOLEAN_SCHEMA);
         assertEquals(false, projected);
@@ -79,7 +79,7 @@ public void testPrimitiveTypeProjection() throws Exception {
     }
 
     @Test
-    public void testNumericTypeProjection() throws Exception {
+    public void testNumericTypeProjection() {
         Schema[] promotableSchemas = {Schema.INT8_SCHEMA, Schema.INT16_SCHEMA, Schema.INT32_SCHEMA, Schema.INT64_SCHEMA, Schema.FLOAT32_SCHEMA, Schema.FLOAT64_SCHEMA};
         Schema[] promotableOptionalSchemas = {Schema.OPTIONAL_INT8_SCHEMA, Schema.OPTIONAL_INT16_SCHEMA, Schema.OPTIONAL_INT32_SCHEMA, Schema.OPTIONAL_INT64_SCHEMA,
                                               Schema.OPTIONAL_FLOAT32_SCHEMA, Schema.OPTIONAL_FLOAT64_SCHEMA};
@@ -146,7 +146,7 @@ public void testNumericTypeProjection() throws Exception {
     }
 
     @Test
-    public void testPrimitiveOptionalProjection() throws Exception {
+    public void testPrimitiveOptionalProjection() {
         verifyOptionalProjection(Schema.OPTIONAL_BOOLEAN_SCHEMA, Type.BOOLEAN, false, true, false, true);
         verifyOptionalProjection(Schema.OPTIONAL_BOOLEAN_SCHEMA, Type.BOOLEAN, false, true, false, false);
 
@@ -208,7 +208,7 @@ public void testPrimitiveOptionalProjection() throws Exception {
     }
 
     @Test
-    public void testStructAddField() throws Exception {
+    public void testStructAddField() {
         Schema source = SchemaBuilder.struct()
                 .field("field", Schema.INT32_SCHEMA)
                 .build();
@@ -240,7 +240,7 @@ public void testStructAddField() throws Exception {
     }
 
     @Test
-    public void testStructRemoveField() throws Exception {
+    public void testStructRemoveField() {
         Schema source = SchemaBuilder.struct()
                 .field("field", Schema.INT32_SCHEMA)
                 .field("field2", Schema.INT32_SCHEMA)
@@ -264,7 +264,7 @@ public void testStructRemoveField() throws Exception {
     }
 
     @Test
-    public void testStructDefaultValue() throws Exception {
+    public void testStructDefaultValue() {
         Schema source = SchemaBuilder.struct().optional()
                 .field("field", Schema.INT32_SCHEMA)
                 .field("field2", Schema.INT32_SCHEMA)
@@ -289,7 +289,7 @@ public void testStructDefaultValue() throws Exception {
     }
 
     @Test
-    public void testNestedSchemaProjection() throws Exception {
+    public void testNestedSchemaProjection() {
         Schema sourceFlatSchema = SchemaBuilder.struct()
                 .field("field", Schema.INT32_SCHEMA)
                 .build();
@@ -326,8 +326,8 @@ public void testNestedSchemaProjection() throws Exception {
                                                                      targetNestedSchema);
         assertEquals(1, targetNestedStruct.get("first"));
         assertEquals("abc", targetNestedStruct.get("second"));
-        assertEquals(Arrays.asList(1, 2), (List) targetNestedStruct.get("array"));
-        assertEquals(Collections.singletonMap(5, "def"), (Map) targetNestedStruct.get("map"));
+        assertEquals(Arrays.asList(1, 2), targetNestedStruct.get("array"));
+        assertEquals(Collections.singletonMap(5, "def"), targetNestedStruct.get("map"));
 
         Struct projectedStruct = (Struct) targetNestedStruct.get("nested");
         assertEquals(113, projectedStruct.get("field"));
@@ -335,7 +335,7 @@ public void testNestedSchemaProjection() throws Exception {
     }
 
     @Test
-    public void testLogicalTypeProjection() throws Exception {
+    public void testLogicalTypeProjection() {
         Schema[] logicalTypeSchemas = {Decimal.schema(2), Date.SCHEMA, Time.SCHEMA, Timestamp.SCHEMA};
         Object projected;
 
@@ -352,6 +352,17 @@ public void testLogicalTypeProjection() throws Exception {
         projected = SchemaProjector.project(Timestamp.SCHEMA, 34567L, Timestamp.SCHEMA);
         assertEquals(34567L, projected);
 
+        java.util.Date date = new java.util.Date();
+
+        projected = SchemaProjector.project(Date.SCHEMA, date, Date.SCHEMA);
+        assertEquals(date, projected);
+
+        projected = SchemaProjector.project(Time.SCHEMA, date, Time.SCHEMA);
+        assertEquals(date, projected);
+
+        projected = SchemaProjector.project(Timestamp.SCHEMA, date, Timestamp.SCHEMA);
+        assertEquals(date, projected);
+
         Schema namedSchema = SchemaBuilder.int32().name("invalidLogicalTypeName").build();
         for (Schema logicalTypeSchema: logicalTypeSchemas) {
             try {
@@ -378,25 +389,25 @@ public void testLogicalTypeProjection() throws Exception {
     }
 
     @Test
-    public void testArrayProjection() throws Exception {
+    public void testArrayProjection() {
         Schema source = SchemaBuilder.array(Schema.INT32_SCHEMA).build();
 
         Object projected = SchemaProjector.project(source, Arrays.asList(1, 2, 3), source);
-        assertEquals(Arrays.asList(1, 2, 3), (List) projected);
+        assertEquals(Arrays.asList(1, 2, 3), projected);
 
         Schema optionalSource = SchemaBuilder.array(Schema.INT32_SCHEMA).optional().build();
         Schema target = SchemaBuilder.array(Schema.INT32_SCHEMA).defaultValue(Arrays.asList(1, 2, 3)).build();
         projected = SchemaProjector.project(optionalSource, Arrays.asList(4, 5), target);
-        assertEquals(Arrays.asList(4, 5), (List) projected);
+        assertEquals(Arrays.asList(4, 5), projected);
         projected = SchemaProjector.project(optionalSource, null, target);
-        assertEquals(Arrays.asList(1, 2, 3), (List) projected);
+        assertEquals(Arrays.asList(1, 2, 3), projected);
 
         Schema promotedTarget = SchemaBuilder.array(Schema.INT64_SCHEMA).defaultValue(Arrays.asList(1L, 2L, 3L)).build();
         projected = SchemaProjector.project(optionalSource, Arrays.asList(4, 5), promotedTarget);
         List expectedProjected = Arrays.asList(4L, 5L);
-        assertEquals(expectedProjected, (List) projected);
+        assertEquals(expectedProjected, projected);
         projected = SchemaProjector.project(optionalSource, null, promotedTarget);
-        assertEquals(Arrays.asList(1L, 2L, 3L), (List) projected);
+        assertEquals(Arrays.asList(1L, 2L, 3L), projected);
 
         Schema noDefaultValueTarget = SchemaBuilder.array(Schema.INT32_SCHEMA).build();
         try {
@@ -416,21 +427,21 @@ public void testArrayProjection() throws Exception {
     }
 
     @Test
-    public void testMapProjection() throws Exception {
+    public void testMapProjection() {
         Schema source = SchemaBuilder.map(Schema.INT32_SCHEMA, Schema.INT32_SCHEMA).optional().build();
 
         Schema target = SchemaBuilder.map(Schema.INT32_SCHEMA, Schema.INT32_SCHEMA).defaultValue(Collections.singletonMap(1, 2)).build();
         Object projected = SchemaProjector.project(source, Collections.singletonMap(3, 4), target);
-        assertEquals(Collections.singletonMap(3, 4), (Map) projected);
+        assertEquals(Collections.singletonMap(3, 4), projected);
         projected = SchemaProjector.project(source, null, target);
-        assertEquals(Collections.singletonMap(1, 2), (Map) projected);
+        assertEquals(Collections.singletonMap(1, 2), projected);
 
         Schema promotedTarget = SchemaBuilder.map(Schema.INT64_SCHEMA, Schema.FLOAT32_SCHEMA).defaultValue(
                 Collections.singletonMap(3L, 4.5F)).build();
         projected = SchemaProjector.project(source, Collections.singletonMap(3, 4), promotedTarget);
-        assertEquals(Collections.singletonMap(3L, 4.F), (Map) projected);
+        assertEquals(Collections.singletonMap(3L, 4.F), projected);
         projected = SchemaProjector.project(source, null, promotedTarget);
-        assertEquals(Collections.singletonMap(3L, 4.5F), (Map) projected);
+        assertEquals(Collections.singletonMap(3L, 4.5F), projected);
 
         Schema noDefaultValueTarget = SchemaBuilder.map(Schema.INT32_SCHEMA, Schema.INT32_SCHEMA).build();
         try {
@@ -450,7 +461,7 @@ public void testMapProjection() throws Exception {
     }
 
     @Test
-    public void testMaybeCompatible() throws Exception {
+    public void testMaybeCompatible() {
         Schema source = SchemaBuilder.int32().name("source").build();
         Schema target = SchemaBuilder.int32().name("target").build();
 
diff --git a/connect/api/src/test/java/org/apache/kafka/connect/data/StructTest.java b/connect/api/src/test/java/org/apache/kafka/connect/data/StructTest.java
index 42345b1986cdf..415b2951816b8 100644
--- a/connect/api/src/test/java/org/apache/kafka/connect/data/StructTest.java
+++ b/connect/api/src/test/java/org/apache/kafka/connect/data/StructTest.java
@@ -17,9 +17,7 @@
 package org.apache.kafka.connect.data;
 
 import org.apache.kafka.connect.errors.DataException;
-import org.junit.Rule;
 import org.junit.Test;
-import org.junit.rules.ExpectedException;
 
 import java.nio.ByteBuffer;
 import java.util.Arrays;
@@ -30,6 +28,8 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertThrows;
+
 
 public class StructTest {
 
@@ -236,8 +236,53 @@ public void testEquals() {
         assertNotEquals(struct1, struct3);
     }
 
-    @Rule
-    public ExpectedException thrown = ExpectedException.none();
+    @Test
+    public void testEqualsAndHashCodeWithByteArrayValue() {
+        Struct struct1 = new Struct(FLAT_STRUCT_SCHEMA)
+                .put("int8", (byte) 12)
+                .put("int16", (short) 12)
+                .put("int32", 12)
+                .put("int64", (long) 12)
+                .put("float32", 12.f)
+                .put("float64", 12.)
+                .put("boolean", true)
+                .put("string", "foobar")
+                .put("bytes", "foobar".getBytes());
+
+        Struct struct2 = new Struct(FLAT_STRUCT_SCHEMA)
+                .put("int8", (byte) 12)
+                .put("int16", (short) 12)
+                .put("int32", 12)
+                .put("int64", (long) 12)
+                .put("float32", 12.f)
+                .put("float64", 12.)
+                .put("boolean", true)
+                .put("string", "foobar")
+                .put("bytes", "foobar".getBytes());
+
+        Struct struct3 = new Struct(FLAT_STRUCT_SCHEMA)
+                .put("int8", (byte) 12)
+                .put("int16", (short) 12)
+                .put("int32", 12)
+                .put("int64", (long) 12)
+                .put("float32", 12.f)
+                .put("float64", 12.)
+                .put("boolean", true)
+                .put("string", "foobar")
+                .put("bytes", "mismatching_string".getBytes());
+
+        // Verify contract for equals: method must be reflexive and transitive
+        assertEquals(struct1, struct2);
+        assertEquals(struct2, struct1);
+        assertNotEquals(struct1, struct3);
+        assertNotEquals(struct2, struct3);
+        // Testing hashCode against a hardcoded value here would be incorrect: hashCode values need not be equal for any
+        // two distinct executions. However, based on the general contract for hashCode, if two objects are equal, their
+        // hashCodes must be equal. If they are not equal, their hashCodes should not be equal for performance reasons.
+        assertEquals(struct1.hashCode(), struct2.hashCode());
+        assertNotEquals(struct1.hashCode(), struct3.hashCode());
+        assertNotEquals(struct2.hashCode(), struct3.hashCode());
+    }
 
     @Test
     public void testValidateStructWithNullValue() {
@@ -248,9 +293,9 @@ public void testValidateStructWithNullValue() {
                 .build();
 
         Struct struct = new Struct(schema);
-        thrown.expect(DataException.class);
-        thrown.expectMessage("Invalid value: null used for required field: \"one\", schema type: STRING");
-        struct.validate();
+        Exception e = assertThrows(DataException.class, struct::validate);
+        assertEquals("Invalid value: null used for required field: \"one\", schema type: STRING",
+            e.getMessage());
     }
 
     @Test
@@ -258,12 +303,36 @@ public void testValidateFieldWithInvalidValueType() {
         String fieldName = "field";
         FakeSchema fakeSchema = new FakeSchema();
 
-        thrown.expect(DataException.class);
-        thrown.expectMessage("Invalid Java object for schema type null: class java.lang.Object for field: \"field\"");
-        ConnectSchema.validateValue(fieldName, fakeSchema, new Object());
+        Exception e = assertThrows(DataException.class, () -> ConnectSchema.validateValue(fieldName,
+            fakeSchema, new Object()));
+        assertEquals("Invalid Java object for schema type null: class java.lang.Object for field: \"field\"",
+            e.getMessage());
+
+        e = assertThrows(DataException.class, () -> ConnectSchema.validateValue(fieldName,
+            Schema.INT8_SCHEMA, new Object()));
+        assertEquals("Invalid Java object for schema type INT8: class java.lang.Object for field: \"field\"",
+            e.getMessage());
+    }
 
-        thrown.expect(DataException.class);
-        thrown.expectMessage("Invalid Java object for schema type INT8: class java.lang.Object for field: \"field\"");
-        ConnectSchema.validateValue(fieldName, Schema.INT8_SCHEMA, new Object());
+    @Test
+    public void testPutNullField() {
+        final String fieldName = "fieldName";
+        Schema testSchema = SchemaBuilder.struct()
+            .field(fieldName, Schema.STRING_SCHEMA);
+        Struct struct = new Struct(testSchema);
+
+        assertThrows(DataException.class, () -> struct.put((Field) null, "valid"));
+    }
+
+    @Test
+    public void testInvalidPutIncludesFieldName() {
+        final String fieldName = "fieldName";
+        Schema testSchema = SchemaBuilder.struct()
+            .field(fieldName, Schema.STRING_SCHEMA);
+        Struct struct = new Struct(testSchema);
+
+        Exception e = assertThrows(DataException.class, () -> struct.put(fieldName, null));
+        assertEquals("Invalid value: null used for required field: \"fieldName\", schema type: STRING",
+            e.getMessage());
     }
 }
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
new file mode 100644
index 0000000000000..cb6f71a60dfa3
--- /dev/null
+++ b/connect/api/src/test/java/org/apache/kafka/connect/data/ValuesTest.java
@@ -0,0 +1,996 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.data;
+
+import org.apache.kafka.connect.data.Schema.Type;
+import org.apache.kafka.connect.data.Values.Parser;
+import org.apache.kafka.connect.errors.DataException;
+import org.junit.Test;
+
+import java.math.BigDecimal;
+import java.math.BigInteger;
+import java.nio.charset.StandardCharsets;
+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;
+
+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.junit.Assert.fail;
+
+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<>();
+    private static final Schema STRING_MAP_SCHEMA = SchemaBuilder.map(Schema.STRING_SCHEMA, Schema.STRING_SCHEMA).schema();
+
+    private static final Map STRING_SHORT_MAP = new LinkedHashMap<>();
+    private static final Schema STRING_SHORT_MAP_SCHEMA = SchemaBuilder.map(Schema.STRING_SCHEMA, Schema.INT16_SCHEMA).schema();
+
+    private static final Map STRING_INT_MAP = new LinkedHashMap<>();
+    private static final Schema STRING_INT_MAP_SCHEMA = SchemaBuilder.map(Schema.STRING_SCHEMA, Schema.INT32_SCHEMA).schema();
+
+    private static final List INT_LIST = new ArrayList<>();
+    private static final Schema INT_LIST_SCHEMA = SchemaBuilder.array(Schema.INT32_SCHEMA).schema();
+
+    private static final List STRING_LIST = new ArrayList<>();
+    private static final Schema STRING_LIST_SCHEMA = SchemaBuilder.array(Schema.STRING_SCHEMA).schema();
+
+    static {
+        STRING_MAP.put("foo", "123");
+        STRING_MAP.put("bar", "baz");
+        STRING_SHORT_MAP.put("foo", (short) 12345);
+        STRING_SHORT_MAP.put("bar", (short) 0);
+        STRING_SHORT_MAP.put("baz", (short) -4321);
+        STRING_INT_MAP.put("foo", 1234567890);
+        STRING_INT_MAP.put("bar", 0);
+        STRING_INT_MAP.put("baz", -987654321);
+        STRING_LIST.add("foo");
+        STRING_LIST.add("bar");
+        INT_LIST.add(1234567890);
+        INT_LIST.add(-987654321);
+    }
+
+    @Test(timeout = 5000)
+    public void shouldNotEncounterInfiniteLoop() {
+        // This byte sequence gets parsed as CharacterIterator.DONE and can cause issues if
+        // comparisons to that character are done to check if the end of a string has been reached.
+        // For more information, see https://issues.apache.org/jira/browse/KAFKA-10574
+        byte[] bytes = new byte[] {-17, -65,  -65};
+        String str = new String(bytes, StandardCharsets.UTF_8);
+        SchemaAndValue schemaAndValue = Values.parseString(str);
+        assertEquals(Type.STRING, schemaAndValue.schema().type());
+        assertEquals(str, schemaAndValue.value());
+    }
+
+    @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";
+        String expected = "three\\\"blind\\\\\\\"mice";
+        assertEquals(expected, Values.escape(original));
+    }
+
+    @Test
+    public void shouldConvertNullValue() {
+        assertRoundTrip(Schema.STRING_SCHEMA, Schema.STRING_SCHEMA, null);
+        assertRoundTrip(Schema.OPTIONAL_STRING_SCHEMA, Schema.STRING_SCHEMA, null);
+    }
+
+    @Test
+    public void shouldConvertBooleanValues() {
+        assertRoundTrip(Schema.BOOLEAN_SCHEMA, Schema.BOOLEAN_SCHEMA, Boolean.FALSE);
+        SchemaAndValue resultFalse = roundTrip(Schema.BOOLEAN_SCHEMA, "false");
+        assertEquals(Schema.BOOLEAN_SCHEMA, resultFalse.schema());
+        assertEquals(Boolean.FALSE, resultFalse.value());
+
+        assertRoundTrip(Schema.BOOLEAN_SCHEMA, Schema.BOOLEAN_SCHEMA, Boolean.TRUE);
+        SchemaAndValue resultTrue = roundTrip(Schema.BOOLEAN_SCHEMA, "true");
+        assertEquals(Schema.BOOLEAN_SCHEMA, resultTrue.schema());
+        assertEquals(Boolean.TRUE, resultTrue.value());
+    }
+
+    @Test(expected = DataException.class)
+    public void shouldFailToParseInvalidBooleanValueString() {
+        Values.convertToBoolean(Schema.STRING_SCHEMA, "\"green\"");
+    }
+
+    @Test
+    public void shouldConvertSimpleString() {
+        assertRoundTrip(Schema.STRING_SCHEMA,  "simple");
+    }
+
+    @Test
+    public void shouldConvertEmptyString() {
+        assertRoundTrip(Schema.STRING_SCHEMA, "");
+    }
+
+    @Test
+    public void shouldConvertStringWithQuotesAndOtherDelimiterCharacters() {
+        assertRoundTrip(Schema.STRING_SCHEMA, Schema.STRING_SCHEMA, "three\"blind\\\"mice");
+        assertRoundTrip(Schema.STRING_SCHEMA, Schema.STRING_SCHEMA, "string with delimiters: <>?,./\\=+-!@#$%^&*(){}[]|;':");
+    }
+
+    @Test
+    public void shouldConvertMapWithStringKeys() {
+        assertRoundTrip(STRING_MAP_SCHEMA, STRING_MAP_SCHEMA, STRING_MAP);
+    }
+
+    @Test
+    public void shouldParseStringOfMapWithStringValuesWithoutWhitespaceAsMap() {
+        SchemaAndValue result = roundTrip(STRING_MAP_SCHEMA, "{\"foo\":\"123\",\"bar\":\"baz\"}");
+        assertEquals(STRING_MAP_SCHEMA, result.schema());
+        assertEquals(STRING_MAP, result.value());
+    }
+
+    @Test
+    public void shouldParseStringOfMapWithStringValuesWithWhitespaceAsMap() {
+        SchemaAndValue result = roundTrip(STRING_MAP_SCHEMA, "{ \"foo\" : \"123\", \n\"bar\" : \"baz\" } ");
+        assertEquals(STRING_MAP_SCHEMA, result.schema());
+        assertEquals(STRING_MAP, result.value());
+    }
+
+    @Test
+    public void shouldConvertMapWithStringKeysAndShortValues() {
+        assertRoundTrip(STRING_SHORT_MAP_SCHEMA, STRING_SHORT_MAP_SCHEMA, STRING_SHORT_MAP);
+    }
+
+    @Test
+    public void shouldParseStringOfMapWithShortValuesWithoutWhitespaceAsMap() {
+        SchemaAndValue result = roundTrip(STRING_SHORT_MAP_SCHEMA, "{\"foo\":12345,\"bar\":0,\"baz\":-4321}");
+        assertEquals(STRING_SHORT_MAP_SCHEMA, result.schema());
+        assertEquals(STRING_SHORT_MAP, result.value());
+    }
+
+    @Test
+    public void shouldParseStringOfMapWithShortValuesWithWhitespaceAsMap() {
+        SchemaAndValue result = roundTrip(STRING_SHORT_MAP_SCHEMA, " { \"foo\" :  12345 , \"bar\" : 0,  \"baz\" : -4321 }  ");
+        assertEquals(STRING_SHORT_MAP_SCHEMA, result.schema());
+        assertEquals(STRING_SHORT_MAP, result.value());
+    }
+
+    @Test
+    public void shouldConvertMapWithStringKeysAndIntegerValues() {
+        assertRoundTrip(STRING_INT_MAP_SCHEMA, STRING_INT_MAP_SCHEMA, STRING_INT_MAP);
+    }
+
+    @Test
+    public void shouldParseStringOfMapWithIntValuesWithoutWhitespaceAsMap() {
+        SchemaAndValue result = roundTrip(STRING_INT_MAP_SCHEMA, "{\"foo\":1234567890,\"bar\":0,\"baz\":-987654321}");
+        assertEquals(STRING_INT_MAP_SCHEMA, result.schema());
+        assertEquals(STRING_INT_MAP, result.value());
+    }
+
+    @Test
+    public void shouldParseStringOfMapWithIntValuesWithWhitespaceAsMap() {
+        SchemaAndValue result = roundTrip(STRING_INT_MAP_SCHEMA, " { \"foo\" :  1234567890 , \"bar\" : 0,  \"baz\" : -987654321 }  ");
+        assertEquals(STRING_INT_MAP_SCHEMA, result.schema());
+        assertEquals(STRING_INT_MAP, result.value());
+    }
+
+    @Test
+    public void shouldConvertListWithStringValues() {
+        assertRoundTrip(STRING_LIST_SCHEMA, STRING_LIST_SCHEMA, STRING_LIST);
+    }
+
+    @Test
+    public void shouldConvertListWithIntegerValues() {
+        assertRoundTrip(INT_LIST_SCHEMA, INT_LIST_SCHEMA, INT_LIST);
+    }
+
+    /**
+     * The parsed array has byte values and one int value, so we should return list with single unified type of integers.
+     */
+    @Test
+    public void shouldConvertStringOfListWithOnlyNumericElementTypesIntoListOfLargestNumericType() {
+        int thirdValue = Short.MAX_VALUE + 1;
+        List list = Values.convertToList(Schema.STRING_SCHEMA, "[1, 2, " + thirdValue + "]");
+        assertEquals(3, list.size());
+        assertEquals(1, ((Number) list.get(0)).intValue());
+        assertEquals(2, ((Number) list.get(1)).intValue());
+        assertEquals(thirdValue, ((Number) list.get(2)).intValue());
+    }
+
+    /**
+     * The parsed array has byte values and one int value, so we should return list with single unified type of integers.
+     */
+    @Test
+    public void shouldConvertStringOfListWithMixedElementTypesIntoListWithDifferentElementTypes() {
+        String str = "[1, 2, \"three\"]";
+        List list = Values.convertToList(Schema.STRING_SCHEMA, str);
+        assertEquals(3, list.size());
+        assertEquals(1, ((Number) list.get(0)).intValue());
+        assertEquals(2, ((Number) list.get(1)).intValue());
+        assertEquals("three", list.get(2));
+    }
+
+    /**
+     * We parse into different element types, but cannot infer a common element schema.
+     */
+    @Test
+    public void shouldParseStringListWithMultipleElementTypesAndReturnListWithNoSchema() {
+        String str = "[1, 2, 3, \"four\"]";
+        SchemaAndValue result = Values.parseString(str);
+        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());
+        assertEquals(2, ((Number) list.get(1)).intValue());
+        assertEquals(3, ((Number) list.get(2)).intValue());
+        assertEquals("four", list.get(3));
+    }
+
+    /**
+     * We can't infer or successfully parse into a different type, so this returns the same string.
+     */
+    @Test
+    public void shouldParseStringListWithExtraDelimitersAndReturnString() {
+        String str = "[1, 2, 3,,,]";
+        SchemaAndValue result = Values.parseString(str);
+        assertEquals(Type.STRING, result.schema().type());
+        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.
+     */
+    @Test(expected = DataException.class)
+    public void shouldFailToConvertToListFromStringWithExtraDelimiters() {
+        Values.convertToList(Schema.STRING_SCHEMA, "[1, 2, 3,,,]");
+    }
+
+    /**
+     * Schema of type ARRAY requires a schema for the values, but Connect has no union or "any" schema type.
+     * Therefore, we can't represent this.
+     */
+    @Test(expected = DataException.class)
+    public void shouldFailToConvertToListFromStringWithNonCommonElementTypeAndBlankElement() {
+        Values.convertToList(Schema.STRING_SCHEMA, "[1, 2, 3, \"four\",,,]");
+    }
+
+    /**
+     * This is technically invalid JSON, and we don't want to simply ignore the blank entry.
+     */
+    @Test(expected = DataException.class)
+    public void shouldFailToParseStringOfMapWithIntValuesWithBlankEntry() {
+        Values.convertToMap(Schema.STRING_SCHEMA, " { \"foo\" :  1234567890 ,, \"bar\" : 0,  \"baz\" : -987654321 }  ");
+    }
+
+    /**
+     * This is technically invalid JSON, and we don't want to simply ignore the malformed entry.
+     */
+    @Test(expected = DataException.class)
+    public void shouldFailToParseStringOfMalformedMap() {
+        Values.convertToMap(Schema.STRING_SCHEMA, " { \"foo\" :  1234567890 , \"a\", \"bar\" : 0,  \"baz\" : -987654321 }  ");
+    }
+
+    /**
+     * This is technically invalid JSON, and we don't want to simply ignore the blank entries.
+     */
+    @Test(expected = DataException.class)
+    public void shouldFailToParseStringOfMapWithIntValuesWithOnlyBlankEntries() {
+        Values.convertToMap(Schema.STRING_SCHEMA, " { ,,  , , }  ");
+    }
+
+    /**
+     * This is technically invalid JSON, and we don't want to simply ignore the blank entry.
+     */
+    @Test(expected = DataException.class)
+    public void shouldFailToParseStringOfMapWithIntValuesWithBlankEntries() {
+        Values.convertToMap(Schema.STRING_SCHEMA, " { \"foo\" :  \"1234567890\" ,, \"bar\" : \"0\",  \"baz\" : \"boz\" }  ");
+    }
+
+    @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
+    public void shouldParseStringsWithoutDelimiters() {
+        //assertParsed("");
+        assertParsed("  ");
+        assertParsed("simple");
+        assertParsed("simple string");
+        assertParsed("simple \n\t\bstring");
+        assertParsed("'simple' string");
+        assertParsed("si\\mple");
+        assertParsed("si\\\\mple");
+    }
+
+    @Test
+    public void shouldParseStringsWithEscapedDelimiters() {
+        assertParsed("si\\\"mple");
+        assertParsed("si\\{mple");
+        assertParsed("si\\}mple");
+        assertParsed("si\\]mple");
+        assertParsed("si\\[mple");
+        assertParsed("si\\:mple");
+        assertParsed("si\\,mple");
+    }
+
+    @Test
+    public void shouldParseStringsWithSingleDelimiter() {
+        assertParsed("a{b", "a", "{", "b");
+        assertParsed("a}b", "a", "}", "b");
+        assertParsed("a[b", "a", "[", "b");
+        assertParsed("a]b", "a", "]", "b");
+        assertParsed("a:b", "a", ":", "b");
+        assertParsed("a,b", "a", ",", "b");
+        assertParsed("a\"b", "a", "\"", "b");
+        assertParsed("{b", "{", "b");
+        assertParsed("}b", "}", "b");
+        assertParsed("[b", "[", "b");
+        assertParsed("]b", "]", "b");
+        assertParsed(":b", ":", "b");
+        assertParsed(",b", ",", "b");
+        assertParsed("\"b", "\"", "b");
+        assertParsed("{", "{");
+        assertParsed("}", "}");
+        assertParsed("[", "[");
+        assertParsed("]", "]");
+        assertParsed(":", ":");
+        assertParsed(",", ",");
+        assertParsed("\"", "\"");
+    }
+
+    @Test
+    public void shouldParseStringsWithMultipleDelimiters() {
+        assertParsed("\"simple\" string", "\"", "simple", "\"", " string");
+        assertParsed("a{bc}d", "a", "{", "bc", "}", "d");
+        assertParsed("a { b c } d", "a ", "{", " b c ", "}", " d");
+        assertParsed("a { b c } d", "a ", "{", " b c ", "}", " d");
+    }
+
+    @Test
+    public void shouldConvertTimeValues() {
+        java.util.Date current = new java.util.Date();
+        long currentMillis = current.getTime() % MILLIS_PER_DAY;
+
+        // java.util.Date - just copy
+        java.util.Date t1 = Values.convertToTime(Time.SCHEMA, current);
+        assertEquals(current, t1);
+
+        // java.util.Date as a Timestamp - discard the date and keep just day's milliseconds
+        t1 = Values.convertToTime(Timestamp.SCHEMA, current);
+        assertEquals(new java.util.Date(currentMillis), t1);
+
+        // ISO8601 strings - currently broken because tokenization breaks at colon
+
+        // Millis as string
+        java.util.Date t3 = Values.convertToTime(Time.SCHEMA, Long.toString(currentMillis));
+        assertEquals(currentMillis, t3.getTime());
+
+        // Millis as long
+        java.util.Date t4 = Values.convertToTime(Time.SCHEMA, currentMillis);
+        assertEquals(currentMillis, t4.getTime());
+    }
+
+    @Test
+    public void shouldConvertDateValues() {
+        java.util.Date current = new java.util.Date();
+        long currentMillis = current.getTime() % MILLIS_PER_DAY;
+        long days = current.getTime() / MILLIS_PER_DAY;
+
+        // java.util.Date - just copy
+        java.util.Date d1 = Values.convertToDate(Date.SCHEMA, current);
+        assertEquals(current, d1);
+
+        // java.util.Date as a Timestamp - discard the day's milliseconds and keep the date
+        java.util.Date currentDate = new java.util.Date(current.getTime() - currentMillis);
+        d1 = Values.convertToDate(Timestamp.SCHEMA, currentDate);
+        assertEquals(currentDate, d1);
+
+        // ISO8601 strings - currently broken because tokenization breaks at colon
+
+        // Days as string
+        java.util.Date d3 = Values.convertToDate(Date.SCHEMA, Long.toString(days));
+        assertEquals(currentDate, d3);
+
+        // Days as long
+        java.util.Date d4 = Values.convertToDate(Date.SCHEMA, days);
+        assertEquals(currentDate, d4);
+    }
+
+    @Test
+    public void shouldConvertTimestampValues() {
+        java.util.Date current = new java.util.Date();
+        long currentMillis = current.getTime() % MILLIS_PER_DAY;
+
+        // java.util.Date - just copy
+        java.util.Date ts1 = Values.convertToTimestamp(Timestamp.SCHEMA, current);
+        assertEquals(current, ts1);
+
+        // java.util.Date as a Timestamp - discard the day's milliseconds and keep the date
+        java.util.Date currentDate = new java.util.Date(current.getTime() - currentMillis);
+        ts1 = Values.convertToTimestamp(Date.SCHEMA, currentDate);
+        assertEquals(currentDate, ts1);
+
+        // java.util.Date as a Time - discard the date and keep the day's milliseconds
+        ts1 = Values.convertToTimestamp(Time.SCHEMA, currentMillis);
+        assertEquals(new java.util.Date(currentMillis), ts1);
+
+        // ISO8601 strings - currently broken because tokenization breaks at colon
+
+        // Millis as string
+        java.util.Date ts3 = Values.convertToTimestamp(Timestamp.SCHEMA, Long.toString(current.getTime()));
+        assertEquals(current, ts3);
+
+        // Millis as long
+        java.util.Date ts4 = Values.convertToTimestamp(Timestamp.SCHEMA, current.getTime());
+        assertEquals(current, ts4);
+    }
+
+    @Test
+    public void canConsume() {
+    }
+
+    @Test
+    public void shouldParseBigIntegerAsDecimalWithZeroScale() {
+        BigInteger value = BigInteger.valueOf(Long.MAX_VALUE).add(new BigInteger("1"));
+        SchemaAndValue schemaAndValue = Values.parseString(
+            String.valueOf(value)
+        );
+        assertEquals(Decimal.schema(0), schemaAndValue.schema());
+        assertTrue(schemaAndValue.value() instanceof BigDecimal);
+        assertEquals(value, ((BigDecimal) schemaAndValue.value()).unscaledValue());
+        value = BigInteger.valueOf(Long.MIN_VALUE).subtract(new BigInteger("1"));
+        schemaAndValue = Values.parseString(
+            String.valueOf(value)
+        );
+        assertEquals(Decimal.schema(0), schemaAndValue.schema());
+        assertTrue(schemaAndValue.value() instanceof BigDecimal);
+        assertEquals(value, ((BigDecimal) schemaAndValue.value()).unscaledValue());
+    }
+
+    @Test
+    public void shouldParseByteAsInt8() {
+        Byte value = Byte.MAX_VALUE;
+        SchemaAndValue schemaAndValue = Values.parseString(
+            String.valueOf(value)
+        );
+        assertEquals(Schema.INT8_SCHEMA, schemaAndValue.schema());
+        assertTrue(schemaAndValue.value() instanceof Byte);
+        assertEquals(value.byteValue(), ((Byte) schemaAndValue.value()).byteValue());
+        value = Byte.MIN_VALUE;
+        schemaAndValue = Values.parseString(
+            String.valueOf(value)
+        );
+        assertEquals(Schema.INT8_SCHEMA, schemaAndValue.schema());
+        assertTrue(schemaAndValue.value() instanceof Byte);
+        assertEquals(value.byteValue(), ((Byte) schemaAndValue.value()).byteValue());
+    }
+
+    @Test
+    public void shouldParseShortAsInt16() {
+        Short value = Short.MAX_VALUE;
+        SchemaAndValue schemaAndValue = Values.parseString(
+            String.valueOf(value)
+        );
+        assertEquals(Schema.INT16_SCHEMA, schemaAndValue.schema());
+        assertTrue(schemaAndValue.value() instanceof Short);
+        assertEquals(value.shortValue(), ((Short) schemaAndValue.value()).shortValue());
+        value = Short.MIN_VALUE;
+        schemaAndValue = Values.parseString(
+            String.valueOf(value)
+        );
+        assertEquals(Schema.INT16_SCHEMA, schemaAndValue.schema());
+        assertTrue(schemaAndValue.value() instanceof Short);
+        assertEquals(value.shortValue(), ((Short) schemaAndValue.value()).shortValue());
+    }
+
+    @Test
+    public void shouldParseIntegerAsInt32() {
+        Integer value = Integer.MAX_VALUE;
+        SchemaAndValue schemaAndValue = Values.parseString(
+            String.valueOf(value)
+        );
+        assertEquals(Schema.INT32_SCHEMA, schemaAndValue.schema());
+        assertTrue(schemaAndValue.value() instanceof Integer);
+        assertEquals(value.intValue(), ((Integer) schemaAndValue.value()).intValue());
+        value = Integer.MIN_VALUE;
+        schemaAndValue = Values.parseString(
+            String.valueOf(value)
+        );
+        assertEquals(Schema.INT32_SCHEMA, schemaAndValue.schema());
+        assertTrue(schemaAndValue.value() instanceof Integer);
+        assertEquals(value.intValue(), ((Integer) schemaAndValue.value()).intValue());
+    }
+
+    @Test
+    public void shouldParseLongAsInt64() {
+        Long value = Long.MAX_VALUE;
+        SchemaAndValue schemaAndValue = Values.parseString(
+            String.valueOf(value)
+        );
+        assertEquals(Schema.INT64_SCHEMA, schemaAndValue.schema());
+        assertTrue(schemaAndValue.value() instanceof Long);
+        assertEquals(value.longValue(), ((Long) schemaAndValue.value()).longValue());
+        value = Long.MIN_VALUE;
+        schemaAndValue = Values.parseString(
+            String.valueOf(value)
+        );
+        assertEquals(Schema.INT64_SCHEMA, schemaAndValue.schema());
+        assertTrue(schemaAndValue.value() instanceof Long);
+        assertEquals(value.longValue(), ((Long) schemaAndValue.value()).longValue());
+    }
+
+    @Test
+    public void shouldParseFloatAsFloat32() {
+        Float value = Float.MAX_VALUE;
+        SchemaAndValue schemaAndValue = Values.parseString(
+            String.valueOf(value)
+        );
+        assertEquals(Schema.FLOAT32_SCHEMA, schemaAndValue.schema());
+        assertTrue(schemaAndValue.value() instanceof Float);
+        assertEquals(value.floatValue(), ((Float) schemaAndValue.value()).floatValue(), 0);
+        value = -Float.MAX_VALUE;
+        schemaAndValue = Values.parseString(
+            String.valueOf(value)
+        );
+        assertEquals(Schema.FLOAT32_SCHEMA, schemaAndValue.schema());
+        assertTrue(schemaAndValue.value() instanceof Float);
+        assertEquals(value.floatValue(), ((Float) schemaAndValue.value()).floatValue(), 0);
+    }
+
+    @Test
+    public void shouldParseDoubleAsFloat64() {
+        Double value = Double.MAX_VALUE;
+        SchemaAndValue schemaAndValue = Values.parseString(
+            String.valueOf(value)
+        );
+        assertEquals(Schema.FLOAT64_SCHEMA, schemaAndValue.schema());
+        assertTrue(schemaAndValue.value() instanceof Double);
+        assertEquals(value.doubleValue(), ((Double) schemaAndValue.value()).doubleValue(), 0);
+        value = -Double.MAX_VALUE;
+        schemaAndValue = Values.parseString(
+            String.valueOf(value)
+        );
+        assertEquals(Schema.FLOAT64_SCHEMA, schemaAndValue.schema());
+        assertTrue(schemaAndValue.value() instanceof Double);
+        assertEquals(value.doubleValue(), ((Double) schemaAndValue.value()).doubleValue(), 0);
+    }
+
+    protected void assertParsed(String input) {
+        assertParsed(input, input);
+    }
+
+    protected void assertParsed(String input, String... expectedTokens) {
+        Parser parser = new Parser(input);
+        if (!parser.hasNext()) {
+            assertEquals(1, expectedTokens.length);
+            assertTrue(expectedTokens[0].isEmpty());
+            return;
+        }
+
+        for (String expectedToken : expectedTokens) {
+            assertTrue(parser.hasNext());
+            int position = parser.mark();
+            assertEquals(expectedToken, parser.next());
+            assertEquals(position + expectedToken.length(), parser.position());
+            assertEquals(expectedToken, parser.previous());
+            parser.rewindTo(position);
+            assertEquals(position, parser.position());
+            assertEquals(expectedToken, parser.next());
+            int newPosition = parser.mark();
+            assertEquals(position + expectedToken.length(), newPosition);
+            assertEquals(expectedToken, parser.previous());
+        }
+        assertFalse(parser.hasNext());
+
+        // Rewind and try consuming expected tokens ...
+        parser.rewindTo(0);
+        assertConsumable(parser, expectedTokens);
+
+        // Parse again and try consuming expected tokens ...
+        parser = new Parser(input);
+        assertConsumable(parser, expectedTokens);
+    }
+
+    protected void assertConsumable(Parser parser, String... expectedTokens) {
+        for (String expectedToken : expectedTokens) {
+            if (!expectedToken.trim().isEmpty()) {
+                int position = parser.mark();
+                assertTrue(parser.canConsume(expectedToken.trim()));
+                parser.rewindTo(position);
+                assertTrue(parser.canConsume(expectedToken.trim(), true));
+                parser.rewindTo(position);
+                assertTrue(parser.canConsume(expectedToken, false));
+            }
+        }
+    }
+
+    protected SchemaAndValue roundTrip(Schema desiredSchema, String currentValue) {
+        return roundTrip(desiredSchema, new SchemaAndValue(Schema.STRING_SCHEMA, currentValue));
+    }
+
+    protected SchemaAndValue roundTrip(Schema desiredSchema, SchemaAndValue input) {
+        String serialized = Values.convertToString(input.schema(), input.value());
+        if (input != null && input.value() != null) {
+            assertNotNull(serialized);
+        }
+        if (desiredSchema == null) {
+            desiredSchema = Values.inferSchema(input);
+            assertNotNull(desiredSchema);
+        }
+        Object newValue = null;
+        Schema newSchema = null;
+        switch (desiredSchema.type()) {
+            case STRING:
+                newValue = Values.convertToString(Schema.STRING_SCHEMA, serialized);
+                break;
+            case INT8:
+                newValue = Values.convertToByte(Schema.STRING_SCHEMA, serialized);
+                break;
+            case INT16:
+                newValue = Values.convertToShort(Schema.STRING_SCHEMA, serialized);
+                break;
+            case INT32:
+                newValue = Values.convertToInteger(Schema.STRING_SCHEMA, serialized);
+                break;
+            case INT64:
+                newValue = Values.convertToLong(Schema.STRING_SCHEMA, serialized);
+                break;
+            case FLOAT32:
+                newValue = Values.convertToFloat(Schema.STRING_SCHEMA, serialized);
+                break;
+            case FLOAT64:
+                newValue = Values.convertToDouble(Schema.STRING_SCHEMA, serialized);
+                break;
+            case BOOLEAN:
+                newValue = Values.convertToBoolean(Schema.STRING_SCHEMA, serialized);
+                break;
+            case ARRAY:
+                newValue = Values.convertToList(Schema.STRING_SCHEMA, serialized);
+                break;
+            case MAP:
+                newValue = Values.convertToMap(Schema.STRING_SCHEMA, serialized);
+                break;
+            case STRUCT:
+                newValue = Values.convertToStruct(Schema.STRING_SCHEMA, serialized);
+                break;
+            case BYTES:
+                fail("unexpected schema type");
+                break;
+        }
+        newSchema = Values.inferSchema(newValue);
+        return new SchemaAndValue(newSchema, newValue);
+    }
+
+    protected void assertRoundTrip(Schema schema, String value) {
+        assertRoundTrip(schema, Schema.STRING_SCHEMA, value);
+    }
+
+    protected void assertRoundTrip(Schema schema, Schema currentSchema, Object value) {
+        SchemaAndValue result = roundTrip(schema, new SchemaAndValue(currentSchema, value));
+
+        if (value == null) {
+            assertNull(result.schema());
+            assertNull(result.value());
+        } else {
+            assertEquals(value, result.value());
+            assertEquals(schema, result.schema());
+
+            SchemaAndValue result2 = roundTrip(result.schema(), result);
+            assertEquals(schema, result2.schema());
+            assertEquals(value, result2.value());
+            assertEquals(result, result2);
+        }
+    }
+}
diff --git a/connect/api/src/test/java/org/apache/kafka/connect/header/ConnectHeaderTest.java b/connect/api/src/test/java/org/apache/kafka/connect/header/ConnectHeaderTest.java
new file mode 100644
index 0000000000000..e30109f3ffe72
--- /dev/null
+++ b/connect/api/src/test/java/org/apache/kafka/connect/header/ConnectHeaderTest.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.header;
+
+import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.data.SchemaAndValue;
+import org.apache.kafka.connect.data.SchemaBuilder;
+import org.apache.kafka.connect.data.Struct;
+import org.junit.Before;
+import org.junit.Test;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertSame;
+
+public class ConnectHeaderTest {
+
+    private String key;
+    private ConnectHeader header;
+
+    @Before
+    public void beforeEach() {
+        key = "key";
+        withString("value");
+    }
+
+    protected Header withValue(Schema schema, Object value) {
+        header = new ConnectHeader(key, new SchemaAndValue(schema, value));
+        return header;
+    }
+
+    protected Header withString(String value) {
+        return withValue(Schema.STRING_SCHEMA, value);
+    }
+
+    @Test
+    public void shouldAllowNullValues() {
+        withValue(Schema.OPTIONAL_STRING_SCHEMA, null);
+    }
+
+    @Test
+    public void shouldAllowNullSchema() {
+        withValue(null, null);
+        assertNull(header.schema());
+        assertNull(header.value());
+
+        String value = "non-null value";
+        withValue(null, value);
+        assertNull(header.schema());
+        assertSame(value, header.value());
+    }
+
+    @Test
+    public void shouldAllowNonNullValue() {
+        String value = "non-null value";
+        withValue(Schema.STRING_SCHEMA, value);
+        assertSame(Schema.STRING_SCHEMA, header.schema());
+        assertEquals(value, header.value());
+
+        withValue(Schema.BOOLEAN_SCHEMA, true);
+        assertSame(Schema.BOOLEAN_SCHEMA, header.schema());
+        assertEquals(true, header.value());
+    }
+
+    @Test
+    public void shouldGetSchemaFromStruct() {
+        Schema schema = SchemaBuilder.struct()
+                                     .field("foo", Schema.STRING_SCHEMA)
+                                     .field("bar", Schema.INT32_SCHEMA)
+                                     .build();
+        Struct value = new Struct(schema);
+        value.put("foo", "value");
+        value.put("bar", 100);
+        withValue(null, value);
+        assertSame(schema, header.schema());
+        assertSame(value, header.value());
+    }
+
+    @Test
+    public void shouldSatisfyEquals() {
+        String value = "non-null value";
+        Header h1 = withValue(Schema.STRING_SCHEMA, value);
+        assertSame(Schema.STRING_SCHEMA, header.schema());
+        assertEquals(value, header.value());
+
+        Header h2 = withValue(Schema.STRING_SCHEMA, value);
+        assertEquals(h1, h2);
+        assertEquals(h1.hashCode(), h2.hashCode());
+
+        Header h3 = withValue(Schema.INT8_SCHEMA, 100);
+        assertNotEquals(h3, h2);
+    }
+}
\ No newline at end of file
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
new file mode 100644
index 0000000000000..f4f11d004d6e8
--- /dev/null
+++ b/connect/api/src/test/java/org/apache/kafka/connect/header/ConnectHeadersTest.java
@@ -0,0 +1,582 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.header;
+
+import org.apache.kafka.connect.data.Date;
+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.SchemaAndValue;
+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.header.Headers.HeaderTransform;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.math.BigDecimal;
+import java.math.RoundingMode;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Calendar;
+import java.util.Collections;
+import java.util.GregorianCalendar;
+import java.util.HashMap;
+import java.util.Iterator;
+import java.util.TimeZone;
+
+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.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;
+
+public class ConnectHeadersTest {
+
+    private static final GregorianCalendar EPOCH_PLUS_TEN_THOUSAND_DAYS;
+    private static final GregorianCalendar EPOCH_PLUS_TEN_THOUSAND_MILLIS;
+
+    static {
+        EPOCH_PLUS_TEN_THOUSAND_DAYS = new GregorianCalendar(1970, Calendar.JANUARY, 1, 0, 0, 0);
+        EPOCH_PLUS_TEN_THOUSAND_DAYS.setTimeZone(TimeZone.getTimeZone("UTC"));
+        EPOCH_PLUS_TEN_THOUSAND_DAYS.add(Calendar.DATE, 10000);
+
+        EPOCH_PLUS_TEN_THOUSAND_MILLIS = new GregorianCalendar(1970, Calendar.JANUARY, 1, 0, 0, 0);
+        EPOCH_PLUS_TEN_THOUSAND_MILLIS.setTimeZone(TimeZone.getTimeZone("UTC"));
+        EPOCH_PLUS_TEN_THOUSAND_MILLIS.add(Calendar.MILLISECOND, 10000);
+    }
+
+    private ConnectHeaders headers;
+    private Iterator
             iter;
+    private String key;
+    private String other;
+
+    @Before
+    public void beforeEach() {
+        headers = new ConnectHeaders();
+        key = "k1";
+        other = "other key";
+    }
+
+    @Test(expected = NullPointerException.class)
+    public void shouldNotAllowNullKey() {
+        headers.add(null, "value", Schema.STRING_SCHEMA);
+    }
+
+    protected void populate(Headers headers) {
+        headers.addBoolean(key, true);
+        headers.addInt(key, 0);
+        headers.addString(other, "other value");
+        headers.addString(key, null);
+        headers.addString(key, "third");
+    }
+
+    @Test
+    public void shouldBeEquals() {
+        Headers other = new ConnectHeaders();
+        assertEquals(headers, other);
+        assertEquals(headers.hashCode(), other.hashCode());
+
+        populate(headers);
+        assertNotEquals(headers, other);
+        assertNotEquals(headers.hashCode(), other.hashCode());
+
+        populate(other);
+        assertEquals(headers, other);
+        assertEquals(headers.hashCode(), other.hashCode());
+
+        headers.addString("wow", "some value");
+        assertNotEquals(headers, other);
+    }
+
+    @Test
+    public void shouldHaveToString() {
+        // empty
+        assertNotNull(headers.toString());
+
+        // not empty
+        populate(headers);
+        assertNotNull(headers.toString());
+    }
+
+    @Test
+    public void shouldRetainLatestWhenEmpty() {
+        headers.retainLatest(other);
+        headers.retainLatest(key);
+        headers.retainLatest();
+        assertTrue(headers.isEmpty());
+    }
+
+    @Test
+    public void shouldAddMultipleHeadersWithSameKeyAndRetainLatest() {
+        populate(headers);
+
+        Header header = headers.lastWithName(key);
+        assertHeader(header, key, Schema.STRING_SCHEMA, "third");
+
+        iter = headers.allWithName(key);
+        assertNextHeader(iter, key, Schema.BOOLEAN_SCHEMA, true);
+        assertNextHeader(iter, key, Schema.INT32_SCHEMA, 0);
+        assertNextHeader(iter, key, Schema.OPTIONAL_STRING_SCHEMA, null);
+        assertNextHeader(iter, key, Schema.STRING_SCHEMA, "third");
+        assertNoNextHeader(iter);
+
+        iter = headers.allWithName(other);
+        assertOnlyNextHeader(iter, other, Schema.STRING_SCHEMA, "other value");
+
+        headers.retainLatest(other);
+        assertOnlySingleHeader(other, Schema.STRING_SCHEMA, "other value");
+
+        headers.retainLatest(key);
+        assertOnlySingleHeader(key, Schema.STRING_SCHEMA, "third");
+
+        headers.retainLatest();
+        assertOnlySingleHeader(other, Schema.STRING_SCHEMA, "other value");
+        assertOnlySingleHeader(key, Schema.STRING_SCHEMA, "third");
+    }
+
+    @Test
+    public void shouldAddHeadersWithPrimitiveValues() {
+        String key = "k1";
+        headers.addBoolean(key, true);
+        headers.addByte(key, (byte) 0);
+        headers.addShort(key, (short) 0);
+        headers.addInt(key, 0);
+        headers.addLong(key, 0);
+        headers.addFloat(key, 1.0f);
+        headers.addDouble(key, 1.0d);
+        headers.addString(key, null);
+        headers.addString(key, "third");
+    }
+
+    @Test
+    public void shouldAddHeadersWithNullObjectValuesWithOptionalSchema() {
+        addHeader("k1", Schema.BOOLEAN_SCHEMA, true);
+        addHeader("k2", Schema.STRING_SCHEMA, "hello");
+        addHeader("k3", Schema.OPTIONAL_STRING_SCHEMA, null);
+    }
+
+    @Test
+    public void shouldNotAddHeadersWithNullObjectValuesWithNonOptionalSchema() {
+        attemptAndFailToAddHeader("k1", Schema.BOOLEAN_SCHEMA, null);
+        attemptAndFailToAddHeader("k2", Schema.STRING_SCHEMA, null);
+    }
+
+    @Test
+    public void shouldNotAddHeadersWithObjectValuesAndMismatchedSchema() {
+        attemptAndFailToAddHeader("k1", Schema.BOOLEAN_SCHEMA, "wrong");
+        attemptAndFailToAddHeader("k2", Schema.OPTIONAL_STRING_SCHEMA, 0L);
+    }
+
+    @Test
+    public void shouldRemoveAllHeadersWithSameKeyWhenEmpty() {
+        headers.remove(key);
+        assertNoHeaderWithKey(key);
+    }
+
+    @Test
+    public void shouldRemoveAllHeadersWithSameKey() {
+        populate(headers);
+
+        iter = headers.allWithName(key);
+        assertContainsHeader(key, Schema.BOOLEAN_SCHEMA, true);
+        assertContainsHeader(key, Schema.INT32_SCHEMA, 0);
+        assertContainsHeader(key, Schema.STRING_SCHEMA, "third");
+        assertOnlySingleHeader(other, Schema.STRING_SCHEMA, "other value");
+
+        headers.remove(key);
+        assertNoHeaderWithKey(key);
+        assertOnlySingleHeader(other, Schema.STRING_SCHEMA, "other value");
+    }
+
+    @Test
+    public void shouldRemoveAllHeaders() {
+        populate(headers);
+
+        iter = headers.allWithName(key);
+        assertContainsHeader(key, Schema.BOOLEAN_SCHEMA, true);
+        assertContainsHeader(key, Schema.INT32_SCHEMA, 0);
+        assertContainsHeader(key, Schema.STRING_SCHEMA, "third");
+        assertOnlySingleHeader(other, Schema.STRING_SCHEMA, "other value");
+
+        headers.clear();
+        assertNoHeaderWithKey(key);
+        assertNoHeaderWithKey(other);
+        assertEquals(0, headers.size());
+        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);
+
+        iter = headers.allWithName(key);
+        assertNextHeader(iter, key, Schema.BOOLEAN_SCHEMA, true);
+        assertNextHeader(iter, key, Schema.INT32_SCHEMA, 0);
+        assertNextHeader(iter, key, Schema.OPTIONAL_STRING_SCHEMA, null);
+        assertNextHeader(iter, key, Schema.STRING_SCHEMA, "third");
+        assertNoNextHeader(iter);
+
+        iter = headers.allWithName(other);
+        assertOnlyNextHeader(iter, other, Schema.STRING_SCHEMA, "other value");
+
+        // Transform the headers
+        assertEquals(5, headers.size());
+        headers.apply(appendToKey("-suffix"));
+        assertEquals(5, headers.size());
+
+        assertNoHeaderWithKey(key);
+        assertNoHeaderWithKey(other);
+
+        String altKey = key + "-suffix";
+        iter = headers.allWithName(altKey);
+        assertNextHeader(iter, altKey, Schema.BOOLEAN_SCHEMA, true);
+        assertNextHeader(iter, altKey, Schema.INT32_SCHEMA, 0);
+        assertNextHeader(iter, altKey, Schema.OPTIONAL_STRING_SCHEMA, null);
+        assertNextHeader(iter, altKey, Schema.STRING_SCHEMA, "third");
+        assertNoNextHeader(iter);
+
+        iter = headers.allWithName(other + "-suffix");
+        assertOnlyNextHeader(iter, other + "-suffix", Schema.STRING_SCHEMA, "other value");
+    }
+
+    @Test
+    public void shouldTransformHeadersWithKey() {
+        populate(headers);
+
+        iter = headers.allWithName(key);
+        assertNextHeader(iter, key, Schema.BOOLEAN_SCHEMA, true);
+        assertNextHeader(iter, key, Schema.INT32_SCHEMA, 0);
+        assertNextHeader(iter, key, Schema.OPTIONAL_STRING_SCHEMA, null);
+        assertNextHeader(iter, key, Schema.STRING_SCHEMA, "third");
+        assertNoNextHeader(iter);
+
+        iter = headers.allWithName(other);
+        assertOnlyNextHeader(iter, other, Schema.STRING_SCHEMA, "other value");
+
+        // Transform the headers
+        assertEquals(5, headers.size());
+        headers.apply(key, appendToKey("-suffix"));
+        assertEquals(5, headers.size());
+
+        assertNoHeaderWithKey(key);
+
+        String altKey = key + "-suffix";
+        iter = headers.allWithName(altKey);
+        assertNextHeader(iter, altKey, Schema.BOOLEAN_SCHEMA, true);
+        assertNextHeader(iter, altKey, Schema.INT32_SCHEMA, 0);
+        assertNextHeader(iter, altKey, Schema.OPTIONAL_STRING_SCHEMA, null);
+        assertNextHeader(iter, altKey, Schema.STRING_SCHEMA, "third");
+        assertNoNextHeader(iter);
+
+        iter = headers.allWithName(other);
+        assertOnlyNextHeader(iter, other, Schema.STRING_SCHEMA, "other value");
+    }
+
+    @Test
+    public void shouldTransformAndRemoveHeaders() {
+        populate(headers);
+
+        iter = headers.allWithName(key);
+        assertNextHeader(iter, key, Schema.BOOLEAN_SCHEMA, true);
+        assertNextHeader(iter, key, Schema.INT32_SCHEMA, 0);
+        assertNextHeader(iter, key, Schema.OPTIONAL_STRING_SCHEMA, null);
+        assertNextHeader(iter, key, Schema.STRING_SCHEMA, "third");
+        assertNoNextHeader(iter);
+
+        iter = headers.allWithName(other);
+        assertOnlyNextHeader(iter, other, Schema.STRING_SCHEMA, "other value");
+
+        // Transform the headers
+        assertEquals(5, headers.size());
+        headers.apply(key, removeHeadersOfType(Type.STRING));
+        assertEquals(3, headers.size());
+
+        iter = headers.allWithName(key);
+        assertNextHeader(iter, key, Schema.BOOLEAN_SCHEMA, true);
+        assertNextHeader(iter, key, Schema.INT32_SCHEMA, 0);
+        assertNoNextHeader(iter);
+
+        assertHeader(headers.lastWithName(key), key, Schema.INT32_SCHEMA, 0);
+
+        iter = headers.allWithName(other);
+        assertOnlyNextHeader(iter, other, Schema.STRING_SCHEMA, "other value");
+
+        // Transform the headers
+        assertEquals(3, headers.size());
+        headers.apply(removeHeadersOfType(Type.STRING));
+        assertEquals(2, headers.size());
+
+        assertNoHeaderWithKey(other);
+
+        iter = headers.allWithName(key);
+        assertNextHeader(iter, key, Schema.BOOLEAN_SCHEMA, true);
+        assertNextHeader(iter, key, Schema.INT32_SCHEMA, 0);
+        assertNoNextHeader(iter);
+    }
+
+    protected HeaderTransform appendToKey(final String suffix) {
+        return new HeaderTransform() {
+            @Override
+            public Header apply(Header header) {
+                return header.rename(header.key() + suffix);
+            }
+        };
+    }
+
+    protected HeaderTransform removeHeadersOfType(final Type type) {
+        return new HeaderTransform() {
+            @Override
+            public Header apply(Header header) {
+                Schema schema = header.schema();
+                if (schema != null && schema.type() == type) {
+                    return null;
+                }
+                return header;
+            }
+        };
+    }
+
+    @Test
+    public void shouldValidateBuildInTypes() {
+        assertSchemaMatches(Schema.OPTIONAL_BOOLEAN_SCHEMA, null);
+        assertSchemaMatches(Schema.OPTIONAL_BYTES_SCHEMA, null);
+        assertSchemaMatches(Schema.OPTIONAL_INT8_SCHEMA, null);
+        assertSchemaMatches(Schema.OPTIONAL_INT16_SCHEMA, null);
+        assertSchemaMatches(Schema.OPTIONAL_INT32_SCHEMA, null);
+        assertSchemaMatches(Schema.OPTIONAL_INT64_SCHEMA, null);
+        assertSchemaMatches(Schema.OPTIONAL_FLOAT32_SCHEMA, null);
+        assertSchemaMatches(Schema.OPTIONAL_FLOAT64_SCHEMA, null);
+        assertSchemaMatches(Schema.OPTIONAL_STRING_SCHEMA, null);
+        assertSchemaMatches(Schema.BOOLEAN_SCHEMA, true);
+        assertSchemaMatches(Schema.BYTES_SCHEMA, new byte[]{});
+        assertSchemaMatches(Schema.INT8_SCHEMA, (byte) 0);
+        assertSchemaMatches(Schema.INT16_SCHEMA, (short) 0);
+        assertSchemaMatches(Schema.INT32_SCHEMA, 0);
+        assertSchemaMatches(Schema.INT64_SCHEMA, 0L);
+        assertSchemaMatches(Schema.FLOAT32_SCHEMA, 1.0f);
+        assertSchemaMatches(Schema.FLOAT64_SCHEMA, 1.0d);
+        assertSchemaMatches(Schema.STRING_SCHEMA, "value");
+        assertSchemaMatches(SchemaBuilder.array(Schema.STRING_SCHEMA), new ArrayList());
+        assertSchemaMatches(SchemaBuilder.array(Schema.STRING_SCHEMA), Collections.singletonList("value"));
+        assertSchemaMatches(SchemaBuilder.map(Schema.STRING_SCHEMA, Schema.INT32_SCHEMA), new HashMap());
+        assertSchemaMatches(SchemaBuilder.map(Schema.STRING_SCHEMA, Schema.INT32_SCHEMA), Collections.singletonMap("a", 0));
+        Schema emptyStructSchema = SchemaBuilder.struct();
+        assertSchemaMatches(emptyStructSchema, new Struct(emptyStructSchema));
+        Schema structSchema = SchemaBuilder.struct().field("foo", Schema.OPTIONAL_BOOLEAN_SCHEMA).field("bar", Schema.STRING_SCHEMA)
+                                           .schema();
+        assertSchemaMatches(structSchema, new Struct(structSchema).put("foo", true).put("bar", "v"));
+    }
+
+    @Test
+    public void shouldValidateLogicalTypes() {
+        assertSchemaMatches(Decimal.schema(3), new BigDecimal(100.00));
+        assertSchemaMatches(Time.SCHEMA, new java.util.Date());
+        assertSchemaMatches(Date.SCHEMA, new java.util.Date());
+        assertSchemaMatches(Timestamp.SCHEMA, new java.util.Date());
+    }
+
+    @Test
+    public void shouldNotValidateNullValuesWithBuiltInTypes() {
+        assertSchemaDoesNotMatch(Schema.BOOLEAN_SCHEMA, null);
+        assertSchemaDoesNotMatch(Schema.BYTES_SCHEMA, null);
+        assertSchemaDoesNotMatch(Schema.INT8_SCHEMA, null);
+        assertSchemaDoesNotMatch(Schema.INT16_SCHEMA, null);
+        assertSchemaDoesNotMatch(Schema.INT32_SCHEMA, null);
+        assertSchemaDoesNotMatch(Schema.INT64_SCHEMA, null);
+        assertSchemaDoesNotMatch(Schema.FLOAT32_SCHEMA, null);
+        assertSchemaDoesNotMatch(Schema.FLOAT64_SCHEMA, null);
+        assertSchemaDoesNotMatch(Schema.STRING_SCHEMA, null);
+        assertSchemaDoesNotMatch(SchemaBuilder.array(Schema.STRING_SCHEMA), null);
+        assertSchemaDoesNotMatch(SchemaBuilder.map(Schema.STRING_SCHEMA, Schema.INT32_SCHEMA), null);
+        assertSchemaDoesNotMatch(SchemaBuilder.struct(), null);
+    }
+
+    @Test
+    public void shouldNotValidateMismatchedValuesWithBuiltInTypes() {
+        assertSchemaDoesNotMatch(Schema.BOOLEAN_SCHEMA, 0L);
+        assertSchemaDoesNotMatch(Schema.BYTES_SCHEMA, "oops");
+        assertSchemaDoesNotMatch(Schema.INT8_SCHEMA, 1.0f);
+        assertSchemaDoesNotMatch(Schema.INT16_SCHEMA, 1.0f);
+        assertSchemaDoesNotMatch(Schema.INT32_SCHEMA, 0L);
+        assertSchemaDoesNotMatch(Schema.INT64_SCHEMA, 1.0f);
+        assertSchemaDoesNotMatch(Schema.FLOAT32_SCHEMA, 1L);
+        assertSchemaDoesNotMatch(Schema.FLOAT64_SCHEMA, 1L);
+        assertSchemaDoesNotMatch(Schema.STRING_SCHEMA, true);
+        assertSchemaDoesNotMatch(SchemaBuilder.array(Schema.STRING_SCHEMA), "value");
+        assertSchemaDoesNotMatch(SchemaBuilder.map(Schema.STRING_SCHEMA, Schema.INT32_SCHEMA), "value");
+        assertSchemaDoesNotMatch(SchemaBuilder.struct(), new ArrayList());
+    }
+
+    @Test
+    public void shouldAddDate() {
+        java.util.Date dateObj = EPOCH_PLUS_TEN_THOUSAND_DAYS.getTime();
+        int days = Date.fromLogical(Date.SCHEMA, dateObj);
+        headers.addDate(key, dateObj);
+        Header header = headers.lastWithName(key);
+        assertEquals(days, (int) Values.convertToInteger(header.schema(), header.value()));
+        assertSame(dateObj, Values.convertToDate(header.schema(), header.value()));
+
+        headers.addInt(other, days);
+        header = headers.lastWithName(other);
+        assertEquals(days, (int) Values.convertToInteger(header.schema(), header.value()));
+        assertEquals(dateObj, Values.convertToDate(header.schema(), header.value()));
+    }
+
+    @Test
+    public void shouldAddTime() {
+        java.util.Date dateObj = EPOCH_PLUS_TEN_THOUSAND_MILLIS.getTime();
+        long millis = Time.fromLogical(Time.SCHEMA, dateObj);
+        headers.addTime(key, dateObj);
+        Header header = headers.lastWithName(key);
+        assertEquals(millis, (long) Values.convertToLong(header.schema(), header.value()));
+        assertSame(dateObj, Values.convertToTime(header.schema(), header.value()));
+
+        headers.addLong(other, millis);
+        header = headers.lastWithName(other);
+        assertEquals(millis, (long) Values.convertToLong(header.schema(), header.value()));
+        assertEquals(dateObj, Values.convertToTime(header.schema(), header.value()));
+    }
+
+    @Test
+    public void shouldAddTimestamp() {
+        java.util.Date dateObj = EPOCH_PLUS_TEN_THOUSAND_MILLIS.getTime();
+        long millis = Timestamp.fromLogical(Timestamp.SCHEMA, dateObj);
+        headers.addTimestamp(key, dateObj);
+        Header header = headers.lastWithName(key);
+        assertEquals(millis, (long) Values.convertToLong(header.schema(), header.value()));
+        assertSame(dateObj, Values.convertToTimestamp(header.schema(), header.value()));
+
+        headers.addLong(other, millis);
+        header = headers.lastWithName(other);
+        assertEquals(millis, (long) Values.convertToLong(header.schema(), header.value()));
+        assertEquals(dateObj, Values.convertToTimestamp(header.schema(), header.value()));
+    }
+
+    @Test
+    public void shouldAddDecimal() {
+        BigDecimal value = new BigDecimal("3.038573478e+3");
+        headers.addDecimal(key, value);
+        Header header = headers.lastWithName(key);
+        assertEquals(value.doubleValue(), Values.convertToDouble(header.schema(), header.value()), 0.00001d);
+        assertEquals(value, Values.convertToDecimal(header.schema(), header.value(), value.scale()));
+
+        value = value.setScale(3, RoundingMode.DOWN);
+        BigDecimal decimal = Values.convertToDecimal(header.schema(), header.value(), value.scale());
+        assertEquals(value, decimal.setScale(value.scale(), RoundingMode.DOWN));
+    }
+
+    @Test
+    public void shouldDuplicateAndAlwaysReturnEquivalentButDifferentObject() {
+        assertEquals(headers, headers.duplicate());
+        assertNotSame(headers, headers.duplicate());
+    }
+
+    @Test
+    public void shouldNotAllowToAddNullHeader() {
+        final ConnectHeaders headers = new ConnectHeaders();
+        assertThrows(NullPointerException.class, () -> headers.add(null));
+    }
+
+    @Test
+    public void shouldThrowNpeWhenAddingCollectionWithNullHeader() {
+        final Iterable
             header = Arrays.asList(new ConnectHeader[1]);
+        assertThrows(NullPointerException.class, () -> new ConnectHeaders(header));
+    }
+
+    protected void assertSchemaMatches(Schema schema, Object value) {
+        headers.checkSchemaMatches(new SchemaAndValue(schema.schema(), value));
+    }
+
+    protected void assertSchemaDoesNotMatch(Schema schema, Object value) {
+        try {
+            assertSchemaMatches(schema, value);
+            fail("Should have failed to validate value '" + value + "' and schema: " + schema);
+        } catch (DataException e) {
+            // expected
+        }
+    }
+
+    protected void attemptAndFailToAddHeader(String key, Schema schema, Object value) {
+        try {
+            headers.add(key, value, schema);
+            fail("Should have failed to add header with key '" + key + "', value '" + value + "', and schema: " + schema);
+        } catch (DataException e) {
+            // expected
+        }
+    }
+
+    protected void addHeader(String key, Schema schema, Object value) {
+        headers.add(key, value, schema);
+        Header header = headers.lastWithName(key);
+        assertNotNull(header);
+        assertHeader(header, key, schema, value);
+    }
+
+    protected void assertNoHeaderWithKey(String key) {
+        assertNoNextHeader(headers.allWithName(key));
+    }
+
+    protected void assertContainsHeader(String key, Schema schema, Object value) {
+        Header expected = new ConnectHeader(key, new SchemaAndValue(schema, value));
+        Iterator
             iter = headers.allWithName(key);
+        while (iter.hasNext()) {
+            Header header = iter.next();
+            if (header.equals(expected))
+                return;
+        }
+        fail("Should have found header " + expected);
+    }
+
+    protected void assertOnlySingleHeader(String key, Schema schema, Object value) {
+        assertOnlyNextHeader(headers.allWithName(key), key, schema, value);
+    }
+
+    protected void assertOnlyNextHeader(Iterator
             iter, String key, Schema schema, Object value) {
+        assertNextHeader(iter, key, schema, value);
+        assertNoNextHeader(iter);
+    }
+
+    protected void assertNextHeader(Iterator
             iter, String key, Schema schema, Object value) {
+        Header header = iter.next();
+        assertHeader(header, key, schema, value);
+    }
+
+    protected void assertNoNextHeader(Iterator
             iter) {
+        assertFalse(iter.hasNext());
+    }
+
+    protected void assertHeader(Header header, String key, Schema schema, Object value) {
+        assertNotNull(header);
+        assertSame(schema, header.schema());
+        assertSame(value, header.value());
+    }
+}
diff --git a/connect/api/src/test/java/org/apache/kafka/connect/sink/SinkConnectorTest.java b/connect/api/src/test/java/org/apache/kafka/connect/sink/SinkConnectorTest.java
new file mode 100644
index 0000000000000..a4924ae3bbff6
--- /dev/null
+++ b/connect/api/src/test/java/org/apache/kafka/connect/sink/SinkConnectorTest.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.connect.sink;
+
+import java.util.List;
+import java.util.Map;
+
+import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.connect.connector.ConnectorContext;
+import org.apache.kafka.connect.connector.ConnectorTest;
+import org.apache.kafka.connect.connector.Task;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
+
+public class SinkConnectorTest extends ConnectorTest {
+
+    @Override
+    protected TestSinkConnectorContext createContext() {
+        return new TestSinkConnectorContext();
+    }
+
+    @Override
+    protected TestSinkConnector createConnector() {
+        return new TestSinkConnector();
+    }
+
+    private static class TestSinkConnectorContext implements SinkConnectorContext {
+
+        @Override
+        public void requestTaskReconfiguration() {
+            // Unexpected in these tests
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public void raiseError(Exception e) {
+            // Unexpected in these tests
+            throw new UnsupportedOperationException();
+        }
+    }
+
+    protected static class TestSinkConnector extends SinkConnector implements ConnectorTest.AssertableConnector {
+
+        public static final String VERSION = "an entirely different version";
+
+        private boolean initialized;
+        private List> taskConfigs;
+        private Map props;
+        private boolean started;
+        private boolean stopped;
+
+        @Override
+        public String version() {
+            return VERSION;
+        }
+
+        @Override
+        public void initialize(ConnectorContext ctx) {
+            super.initialize(ctx);
+            initialized = true;
+            this.taskConfigs = null;
+        }
+
+        @Override
+        public void initialize(ConnectorContext ctx, List> taskConfigs) {
+            super.initialize(ctx, taskConfigs);
+            initialized = true;
+            this.taskConfigs = taskConfigs;
+        }
+
+        @Override
+        public void start(Map props) {
+            this.props = props;
+            started = true;
+        }
+
+        @Override
+        public Class taskClass() {
+            return null;
+        }
+
+        @Override
+        public List> taskConfigs(int maxTasks) {
+            return null;
+        }
+
+        @Override
+        public void stop() {
+            stopped = true;
+        }
+
+        @Override
+        public ConfigDef config() {
+            return new ConfigDef()
+                    .define("required", ConfigDef.Type.STRING, ConfigDef.Importance.HIGH, "required docs")
+                    .define("optional", ConfigDef.Type.STRING, "defaultVal", ConfigDef.Importance.HIGH, "optional docs");
+        }
+
+        @Override
+        public void assertContext(ConnectorContext expected) {
+            assertSame(expected, context);
+            assertSame(expected, context());
+        }
+
+        @Override
+        public void assertInitialized() {
+            assertTrue(initialized);
+        }
+
+        @Override
+        public void assertTaskConfigs(List> expectedTaskConfigs) {
+            assertSame(expectedTaskConfigs, taskConfigs);
+        }
+
+        @Override
+        public void assertStarted(boolean expected) {
+            assertEquals(expected, started);
+        }
+
+        @Override
+        public void assertStopped(boolean expected) {
+            assertEquals(expected, stopped);
+        }
+
+        @Override
+        public void assertProperties(Map expected) {
+            assertSame(expected, props);
+        }
+    }
+}
\ No newline at end of file
diff --git a/connect/api/src/test/java/org/apache/kafka/connect/sink/SinkRecordTest.java b/connect/api/src/test/java/org/apache/kafka/connect/sink/SinkRecordTest.java
new file mode 100644
index 0000000000000..fbfbb322044be
--- /dev/null
+++ b/connect/api/src/test/java/org/apache/kafka/connect/sink/SinkRecordTest.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.connect.sink;
+
+import org.apache.kafka.common.record.TimestampType;
+import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.data.Values;
+import org.apache.kafka.connect.header.ConnectHeaders;
+import org.apache.kafka.connect.header.Header;
+import org.apache.kafka.connect.header.Headers;
+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.assertNotEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNotSame;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
+
+public class SinkRecordTest {
+
+    private static final String TOPIC_NAME = "myTopic";
+    private static final Integer PARTITION_NUMBER = 0;
+    private static final long KAFKA_OFFSET = 0L;
+    private static final Long KAFKA_TIMESTAMP = 0L;
+    private static final TimestampType TS_TYPE = TimestampType.CREATE_TIME;
+
+    private SinkRecord record;
+
+    @Before
+    public void beforeEach() {
+        record = new SinkRecord(TOPIC_NAME, PARTITION_NUMBER, Schema.STRING_SCHEMA, "key", Schema.BOOLEAN_SCHEMA, false, KAFKA_OFFSET,
+                                KAFKA_TIMESTAMP, TS_TYPE, null);
+    }
+
+    @Test
+    public void shouldCreateSinkRecordWithHeaders() {
+        Headers headers = new ConnectHeaders().addString("h1", "hv1").addBoolean("h2", true);
+        record = new SinkRecord(TOPIC_NAME, PARTITION_NUMBER, Schema.STRING_SCHEMA, "key", Schema.BOOLEAN_SCHEMA, false, KAFKA_OFFSET,
+                                KAFKA_TIMESTAMP, TS_TYPE, headers);
+        assertNotNull(record.headers());
+        assertSame(headers, record.headers());
+        assertFalse(record.headers().isEmpty());
+    }
+
+    @Test
+    public void shouldCreateSinkRecordWithEmptyHeaders() {
+        assertEquals(TOPIC_NAME, record.topic());
+        assertEquals(PARTITION_NUMBER, record.kafkaPartition());
+        assertEquals(Schema.STRING_SCHEMA, record.keySchema());
+        assertEquals("key", record.key());
+        assertEquals(Schema.BOOLEAN_SCHEMA, record.valueSchema());
+        assertEquals(false, record.value());
+        assertEquals(KAFKA_OFFSET, record.kafkaOffset());
+        assertEquals(KAFKA_TIMESTAMP, record.timestamp());
+        assertEquals(TS_TYPE, record.timestampType());
+        assertNotNull(record.headers());
+        assertTrue(record.headers().isEmpty());
+    }
+
+    @Test
+    public void shouldDuplicateRecordAndCloneHeaders() {
+        SinkRecord duplicate = record.newRecord(TOPIC_NAME, PARTITION_NUMBER, Schema.STRING_SCHEMA, "key", Schema.BOOLEAN_SCHEMA, false,
+                                                KAFKA_TIMESTAMP);
+
+        assertEquals(TOPIC_NAME, duplicate.topic());
+        assertEquals(PARTITION_NUMBER, duplicate.kafkaPartition());
+        assertEquals(Schema.STRING_SCHEMA, duplicate.keySchema());
+        assertEquals("key", duplicate.key());
+        assertEquals(Schema.BOOLEAN_SCHEMA, duplicate.valueSchema());
+        assertEquals(false, duplicate.value());
+        assertEquals(KAFKA_OFFSET, duplicate.kafkaOffset());
+        assertEquals(KAFKA_TIMESTAMP, duplicate.timestamp());
+        assertEquals(TS_TYPE, duplicate.timestampType());
+        assertNotNull(duplicate.headers());
+        assertTrue(duplicate.headers().isEmpty());
+        assertNotSame(record.headers(), duplicate.headers());
+        assertEquals(record.headers(), duplicate.headers());
+    }
+
+
+    @Test
+    public void shouldDuplicateRecordUsingNewHeaders() {
+        Headers newHeaders = new ConnectHeaders().addString("h3", "hv3");
+        SinkRecord duplicate = record.newRecord(TOPIC_NAME, PARTITION_NUMBER, Schema.STRING_SCHEMA, "key", Schema.BOOLEAN_SCHEMA, false,
+                                                KAFKA_TIMESTAMP, newHeaders);
+
+        assertEquals(TOPIC_NAME, duplicate.topic());
+        assertEquals(PARTITION_NUMBER, duplicate.kafkaPartition());
+        assertEquals(Schema.STRING_SCHEMA, duplicate.keySchema());
+        assertEquals("key", duplicate.key());
+        assertEquals(Schema.BOOLEAN_SCHEMA, duplicate.valueSchema());
+        assertEquals(false, duplicate.value());
+        assertEquals(KAFKA_OFFSET, duplicate.kafkaOffset());
+        assertEquals(KAFKA_TIMESTAMP, duplicate.timestamp());
+        assertEquals(TS_TYPE, duplicate.timestampType());
+        assertNotNull(duplicate.headers());
+        assertEquals(newHeaders, duplicate.headers());
+        assertSame(newHeaders, duplicate.headers());
+        assertNotSame(record.headers(), duplicate.headers());
+        assertNotEquals(record.headers(), duplicate.headers());
+    }
+
+    @Test
+    public void shouldModifyRecordHeader() {
+        assertTrue(record.headers().isEmpty());
+        record.headers().addInt("intHeader", 100);
+        assertEquals(1, record.headers().size());
+        Header header = record.headers().lastWithName("intHeader");
+        assertEquals(100, (int) Values.convertToInteger(header.schema(), header.value()));
+    }
+}
\ No newline at end of file
diff --git a/connect/api/src/test/java/org/apache/kafka/connect/source/SourceConnectorTest.java b/connect/api/src/test/java/org/apache/kafka/connect/source/SourceConnectorTest.java
new file mode 100644
index 0000000000000..9b105409fffba
--- /dev/null
+++ b/connect/api/src/test/java/org/apache/kafka/connect/source/SourceConnectorTest.java
@@ -0,0 +1,152 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.source;
+
+import java.util.List;
+import java.util.Map;
+
+import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.connect.connector.ConnectorContext;
+import org.apache.kafka.connect.connector.ConnectorTest;
+import org.apache.kafka.connect.connector.Task;
+import org.apache.kafka.connect.storage.OffsetStorageReader;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
+
+public class SourceConnectorTest extends ConnectorTest {
+
+    @Override
+    protected ConnectorContext createContext() {
+        return new TestSourceConnectorContext();
+    }
+
+    @Override
+    protected TestSourceConnector createConnector() {
+        return new TestSourceConnector();
+    }
+
+    private static class TestSourceConnectorContext implements SourceConnectorContext {
+
+        @Override
+        public void requestTaskReconfiguration() {
+            // Unexpected in these tests
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public void raiseError(Exception e) {
+            // Unexpected in these tests
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public OffsetStorageReader offsetStorageReader() {
+            return null;
+        }
+    }
+
+    private static class TestSourceConnector extends SourceConnector implements AssertableConnector {
+
+        public static final String VERSION = "an entirely different version";
+
+        private boolean initialized;
+        private List> taskConfigs;
+        private Map props;
+        private boolean started;
+        private boolean stopped;
+
+        @Override
+        public String version() {
+            return VERSION;
+        }
+
+        @Override
+        public void initialize(ConnectorContext ctx) {
+            super.initialize(ctx);
+            initialized = true;
+            this.taskConfigs = null;
+        }
+
+        @Override
+        public void initialize(ConnectorContext ctx, List> taskConfigs) {
+            super.initialize(ctx, taskConfigs);
+            initialized = true;
+            this.taskConfigs = taskConfigs;
+        }
+
+        @Override
+        public void start(Map props) {
+            this.props = props;
+            started = true;
+        }
+
+        @Override
+        public Class taskClass() {
+            return null;
+        }
+
+        @Override
+        public List> taskConfigs(int maxTasks) {
+            return null;
+        }
+
+        @Override
+        public void stop() {
+            stopped = true;
+        }
+
+        @Override
+        public ConfigDef config() {
+            return new ConfigDef()
+                    .define("required", ConfigDef.Type.STRING, ConfigDef.Importance.HIGH, "required docs")
+                    .define("optional", ConfigDef.Type.STRING, "defaultVal", ConfigDef.Importance.HIGH, "optional docs");
+        }
+
+        @Override
+        public void assertContext(ConnectorContext expected) {
+            assertSame(expected, context);
+            assertSame(expected, context());
+        }
+
+        @Override
+        public void assertInitialized() {
+            assertTrue(initialized);
+        }
+
+        @Override
+        public void assertTaskConfigs(List> expectedTaskConfigs) {
+            assertSame(expectedTaskConfigs, taskConfigs);
+        }
+
+        @Override
+        public void assertStarted(boolean expected) {
+            assertEquals(expected, started);
+        }
+
+        @Override
+        public void assertStopped(boolean expected) {
+            assertEquals(expected, stopped);
+        }
+
+        @Override
+        public void assertProperties(Map expected) {
+            assertSame(expected, props);
+        }
+    }
+}
\ No newline at end of file
diff --git a/connect/api/src/test/java/org/apache/kafka/connect/source/SourceRecordTest.java b/connect/api/src/test/java/org/apache/kafka/connect/source/SourceRecordTest.java
new file mode 100644
index 0000000000000..275109a4860c2
--- /dev/null
+++ b/connect/api/src/test/java/org/apache/kafka/connect/source/SourceRecordTest.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.connect.source;
+
+import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.data.Values;
+import org.apache.kafka.connect.header.ConnectHeaders;
+import org.apache.kafka.connect.header.Header;
+import org.apache.kafka.connect.header.Headers;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.util.Collections;
+import java.util.Map;
+
+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.assertNotSame;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
+
+public class SourceRecordTest {
+
+    private static final Map SOURCE_PARTITION = Collections.singletonMap("src", "abc");
+    private static final Map SOURCE_OFFSET = Collections.singletonMap("offset", "1");
+    private static final String TOPIC_NAME = "myTopic";
+    private static final Integer PARTITION_NUMBER = 0;
+    private static final Long KAFKA_TIMESTAMP = 0L;
+
+    private SourceRecord record;
+
+    @Before
+    public void beforeEach() {
+        record = new SourceRecord(SOURCE_PARTITION, SOURCE_OFFSET, TOPIC_NAME, PARTITION_NUMBER, Schema.STRING_SCHEMA, "key",
+                                  Schema.BOOLEAN_SCHEMA, false, KAFKA_TIMESTAMP, null);
+    }
+
+    @Test
+    public void shouldCreateSinkRecordWithHeaders() {
+        Headers headers = new ConnectHeaders().addString("h1", "hv1").addBoolean("h2", true);
+        record = new SourceRecord(SOURCE_PARTITION, SOURCE_OFFSET, TOPIC_NAME, PARTITION_NUMBER, Schema.STRING_SCHEMA, "key",
+                                  Schema.BOOLEAN_SCHEMA, false, KAFKA_TIMESTAMP, headers);
+        assertNotNull(record.headers());
+        assertSame(headers, record.headers());
+        assertFalse(record.headers().isEmpty());
+    }
+
+    @Test
+    public void shouldCreateSinkRecordWithEmtpyHeaders() {
+        assertEquals(SOURCE_PARTITION, record.sourcePartition());
+        assertEquals(SOURCE_OFFSET, record.sourceOffset());
+        assertEquals(TOPIC_NAME, record.topic());
+        assertEquals(PARTITION_NUMBER, record.kafkaPartition());
+        assertEquals(Schema.STRING_SCHEMA, record.keySchema());
+        assertEquals("key", record.key());
+        assertEquals(Schema.BOOLEAN_SCHEMA, record.valueSchema());
+        assertEquals(false, record.value());
+        assertEquals(KAFKA_TIMESTAMP, record.timestamp());
+        assertNotNull(record.headers());
+        assertTrue(record.headers().isEmpty());
+    }
+
+    @Test
+    public void shouldDuplicateRecordAndCloneHeaders() {
+        SourceRecord duplicate = record.newRecord(TOPIC_NAME, PARTITION_NUMBER, Schema.STRING_SCHEMA, "key", Schema.BOOLEAN_SCHEMA, false,
+                                                  KAFKA_TIMESTAMP);
+
+        assertEquals(SOURCE_PARTITION, duplicate.sourcePartition());
+        assertEquals(SOURCE_OFFSET, duplicate.sourceOffset());
+        assertEquals(TOPIC_NAME, duplicate.topic());
+        assertEquals(PARTITION_NUMBER, duplicate.kafkaPartition());
+        assertEquals(Schema.STRING_SCHEMA, duplicate.keySchema());
+        assertEquals("key", duplicate.key());
+        assertEquals(Schema.BOOLEAN_SCHEMA, duplicate.valueSchema());
+        assertEquals(false, duplicate.value());
+        assertEquals(KAFKA_TIMESTAMP, duplicate.timestamp());
+        assertNotNull(duplicate.headers());
+        assertTrue(duplicate.headers().isEmpty());
+        assertNotSame(record.headers(), duplicate.headers());
+        assertEquals(record.headers(), duplicate.headers());
+    }
+
+    @Test
+    public void shouldDuplicateRecordUsingNewHeaders() {
+        Headers newHeaders = new ConnectHeaders().addString("h3", "hv3");
+        SourceRecord duplicate = record.newRecord(TOPIC_NAME, PARTITION_NUMBER, Schema.STRING_SCHEMA, "key", Schema.BOOLEAN_SCHEMA, false,
+                                                  KAFKA_TIMESTAMP, newHeaders);
+
+        assertEquals(SOURCE_PARTITION, duplicate.sourcePartition());
+        assertEquals(SOURCE_OFFSET, duplicate.sourceOffset());
+        assertEquals(TOPIC_NAME, duplicate.topic());
+        assertEquals(PARTITION_NUMBER, duplicate.kafkaPartition());
+        assertEquals(Schema.STRING_SCHEMA, duplicate.keySchema());
+        assertEquals("key", duplicate.key());
+        assertEquals(Schema.BOOLEAN_SCHEMA, duplicate.valueSchema());
+        assertEquals(false, duplicate.value());
+        assertEquals(KAFKA_TIMESTAMP, duplicate.timestamp());
+        assertNotNull(duplicate.headers());
+        assertEquals(newHeaders, duplicate.headers());
+        assertSame(newHeaders, duplicate.headers());
+        assertNotSame(record.headers(), duplicate.headers());
+        assertNotEquals(record.headers(), duplicate.headers());
+    }
+
+    @Test
+    public void shouldModifyRecordHeader() {
+        assertTrue(record.headers().isEmpty());
+        record.headers().addInt("intHeader", 100);
+        assertEquals(1, record.headers().size());
+        Header header = record.headers().lastWithName("intHeader");
+        assertEquals(100, (int) Values.convertToInteger(header.schema(), header.value()));
+    }
+}
\ No newline at end of file
diff --git a/connect/api/src/test/java/org/apache/kafka/connect/storage/ConverterTypeTest.java b/connect/api/src/test/java/org/apache/kafka/connect/storage/ConverterTypeTest.java
new file mode 100644
index 0000000000000..cbaa8b23ad07c
--- /dev/null
+++ b/connect/api/src/test/java/org/apache/kafka/connect/storage/ConverterTypeTest.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.connect.storage;
+
+import org.junit.Test;
+
+import static org.junit.Assert.assertEquals;
+
+public class ConverterTypeTest {
+
+    @Test
+    public void shouldFindByName() {
+        for (ConverterType type : ConverterType.values()) {
+            assertEquals(type, ConverterType.withName(type.getName()));
+        }
+    }
+}
\ 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
new file mode 100644
index 0000000000000..58a17f57ba192
--- /dev/null
+++ b/connect/api/src/test/java/org/apache/kafka/connect/storage/SimpleHeaderConverterTest.java
@@ -0,0 +1,236 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.connect.storage;
+
+import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.data.SchemaAndValue;
+import org.apache.kafka.connect.data.SchemaBuilder;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.util.ArrayList;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+
+public class SimpleHeaderConverterTest {
+
+    private static final String TOPIC = "topic";
+    private static final String HEADER = "header";
+
+    private static final Map STRING_MAP = new LinkedHashMap<>();
+    private static final Schema STRING_MAP_SCHEMA = SchemaBuilder.map(Schema.STRING_SCHEMA, Schema.STRING_SCHEMA).schema();
+
+    private static final Map STRING_SHORT_MAP = new LinkedHashMap<>();
+    private static final Schema STRING_SHORT_MAP_SCHEMA = SchemaBuilder.map(Schema.STRING_SCHEMA, Schema.INT16_SCHEMA).schema();
+
+    private static final Map STRING_INT_MAP = new LinkedHashMap<>();
+    private static final Schema STRING_INT_MAP_SCHEMA = SchemaBuilder.map(Schema.STRING_SCHEMA, Schema.INT32_SCHEMA).schema();
+
+    private static final List INT_LIST = new ArrayList<>();
+    private static final Schema INT_LIST_SCHEMA = SchemaBuilder.array(Schema.INT32_SCHEMA).schema();
+
+    private static final List STRING_LIST = new ArrayList<>();
+    private static final Schema STRING_LIST_SCHEMA = SchemaBuilder.array(Schema.STRING_SCHEMA).schema();
+
+    static {
+        STRING_MAP.put("foo", "123");
+        STRING_MAP.put("bar", "baz");
+        STRING_SHORT_MAP.put("foo", (short) 12345);
+        STRING_SHORT_MAP.put("bar", (short) 0);
+        STRING_SHORT_MAP.put("baz", (short) -4321);
+        STRING_INT_MAP.put("foo", 1234567890);
+        STRING_INT_MAP.put("bar", 0);
+        STRING_INT_MAP.put("baz", -987654321);
+        STRING_LIST.add("foo");
+        STRING_LIST.add("bar");
+        INT_LIST.add(1234567890);
+        INT_LIST.add(-987654321);
+    }
+
+    private SimpleHeaderConverter converter;
+
+    @Before
+    public void beforeEach() {
+        converter = new SimpleHeaderConverter();
+    }
+
+    @Test
+    public void shouldConvertNullValue() {
+        assertRoundTrip(Schema.STRING_SCHEMA, null);
+        assertRoundTrip(Schema.OPTIONAL_STRING_SCHEMA, null);
+    }
+
+    @Test
+    public void shouldConvertSimpleString() {
+        assertRoundTrip(Schema.STRING_SCHEMA, "simple");
+    }
+
+    @Test
+    public void shouldConvertEmptyString() {
+        assertRoundTrip(Schema.STRING_SCHEMA, "");
+    }
+
+    @Test
+    public void shouldConvertStringWithQuotesAndOtherDelimiterCharacters() {
+        assertRoundTrip(Schema.STRING_SCHEMA, "three\"blind\\\"mice");
+        assertRoundTrip(Schema.STRING_SCHEMA, "string with delimiters: <>?,./\\=+-!@#$%^&*(){}[]|;':");
+    }
+
+    @Test
+    public void shouldConvertMapWithStringKeys() {
+        assertRoundTrip(STRING_MAP_SCHEMA, STRING_MAP);
+    }
+
+    @Test
+    public void shouldParseStringOfMapWithStringValuesWithoutWhitespaceAsMap() {
+        SchemaAndValue result = roundTrip(Schema.STRING_SCHEMA, "{\"foo\":\"123\",\"bar\":\"baz\"}");
+        assertEquals(STRING_MAP_SCHEMA, result.schema());
+        assertEquals(STRING_MAP, result.value());
+    }
+
+    @Test
+    public void shouldParseStringOfMapWithStringValuesWithWhitespaceAsMap() {
+        SchemaAndValue result = roundTrip(Schema.STRING_SCHEMA, "{ \"foo\" : \"123\", \n\"bar\" : \"baz\" } ");
+        assertEquals(STRING_MAP_SCHEMA, result.schema());
+        assertEquals(STRING_MAP, result.value());
+    }
+
+    @Test
+    public void shouldConvertMapWithStringKeysAndShortValues() {
+        assertRoundTrip(STRING_SHORT_MAP_SCHEMA, STRING_SHORT_MAP);
+    }
+
+    @Test
+    public void shouldParseStringOfMapWithShortValuesWithoutWhitespaceAsMap() {
+        SchemaAndValue result = roundTrip(Schema.STRING_SCHEMA, "{\"foo\":12345,\"bar\":0,\"baz\":-4321}");
+        assertEquals(STRING_SHORT_MAP_SCHEMA, result.schema());
+        assertEquals(STRING_SHORT_MAP, result.value());
+    }
+
+    @Test
+    public void shouldParseStringOfMapWithShortValuesWithWhitespaceAsMap() {
+        SchemaAndValue result = roundTrip(Schema.STRING_SCHEMA, " { \"foo\" :  12345 , \"bar\" : 0,  \"baz\" : -4321 }  ");
+        assertEquals(STRING_SHORT_MAP_SCHEMA, result.schema());
+        assertEquals(STRING_SHORT_MAP, result.value());
+    }
+
+    @Test
+    public void shouldConvertMapWithStringKeysAndIntegerValues() {
+        assertRoundTrip(STRING_INT_MAP_SCHEMA, STRING_INT_MAP);
+    }
+
+    @Test
+    public void shouldParseStringOfMapWithIntValuesWithoutWhitespaceAsMap() {
+        SchemaAndValue result = roundTrip(Schema.STRING_SCHEMA, "{\"foo\":1234567890,\"bar\":0,\"baz\":-987654321}");
+        assertEquals(STRING_INT_MAP_SCHEMA, result.schema());
+        assertEquals(STRING_INT_MAP, result.value());
+    }
+
+    @Test
+    public void shouldParseStringOfMapWithIntValuesWithWhitespaceAsMap() {
+        SchemaAndValue result = roundTrip(Schema.STRING_SCHEMA, " { \"foo\" :  1234567890 , \"bar\" : 0,  \"baz\" : -987654321 }  ");
+        assertEquals(STRING_INT_MAP_SCHEMA, result.schema());
+        assertEquals(STRING_INT_MAP, result.value());
+    }
+
+    @Test
+    public void shouldConvertListWithStringValues() {
+        assertRoundTrip(STRING_LIST_SCHEMA, STRING_LIST);
+    }
+
+    @Test
+    public void shouldConvertListWithIntegerValues() {
+        assertRoundTrip(INT_LIST_SCHEMA, INT_LIST);
+    }
+
+    @Test
+    public void shouldConvertMapWithStringKeysAndMixedValuesToMap() {
+        Map map = new LinkedHashMap<>();
+        map.put("foo", "bar");
+        map.put("baz", (short) 3456);
+        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
+    public void shouldConvertListWithMixedValuesToListWithoutSchema() {
+        List list = new ArrayList<>();
+        list.add("foo");
+        list.add((short) 13344);
+        SchemaAndValue result = roundTrip(null, list);
+        assertEquals(Schema.Type.ARRAY, result.schema().type());
+        assertNull(result.schema().valueSchema());
+        assertEquals(list, result.value());
+    }
+
+    @Test
+    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 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) {
+        byte[] serialized = converter.fromConnectHeader(TOPIC, HEADER, schema, input);
+        return converter.toConnectHeader(TOPIC, HEADER, serialized);
+    }
+
+    protected void assertRoundTrip(Schema schema, Object value) {
+        byte[] serialized = converter.fromConnectHeader(TOPIC, HEADER, schema, value);
+        SchemaAndValue result = converter.toConnectHeader(TOPIC, HEADER, serialized);
+
+        if (value == null) {
+            assertNull(serialized);
+            assertNull(result.schema());
+            assertNull(result.value());
+        } else {
+            assertNotNull(serialized);
+            assertEquals(value, result.value());
+            assertEquals(schema, result.schema());
+
+            byte[] serialized2 = converter.fromConnectHeader(TOPIC, HEADER, result.schema(), result.value());
+            SchemaAndValue result2 = converter.toConnectHeader(TOPIC, HEADER, serialized2);
+            assertNotNull(serialized2);
+            assertEquals(schema, result2.schema());
+            assertEquals(value, result2.value());
+            assertEquals(result, result2);
+            assertArrayEquals(serialized, serialized);
+        }
+    }
+
+}
\ No newline at end of file
diff --git a/connect/api/src/test/java/org/apache/kafka/connect/storage/StringConverterTest.java b/connect/api/src/test/java/org/apache/kafka/connect/storage/StringConverterTest.java
index e860fbb38a26d..49ffc74990a1a 100644
--- a/connect/api/src/test/java/org/apache/kafka/connect/storage/StringConverterTest.java
+++ b/connect/api/src/test/java/org/apache/kafka/connect/storage/StringConverterTest.java
@@ -79,4 +79,22 @@ public void testBytesToStringNonUtf8Encoding() throws UnsupportedEncodingExcepti
         assertEquals(Schema.OPTIONAL_STRING_SCHEMA, data.schema());
         assertEquals(SAMPLE_STRING, data.value());
     }
+
+    // Note: the header conversion methods delegates to the data conversion methods, which are tested above.
+    // The following simply verify that the delegation works.
+
+    @Test
+    public void testStringHeaderValueToBytes() throws UnsupportedEncodingException {
+        assertArrayEquals(SAMPLE_STRING.getBytes("UTF8"), converter.fromConnectHeader(TOPIC, "hdr", Schema.STRING_SCHEMA, SAMPLE_STRING));
+    }
+
+    @Test
+    public void testNonStringHeaderValueToBytes() throws UnsupportedEncodingException {
+        assertArrayEquals("true".getBytes("UTF8"), converter.fromConnectHeader(TOPIC, "hdr", Schema.BOOLEAN_SCHEMA, true));
+    }
+
+    @Test
+    public void testNullHeaderValueToBytes() {
+        assertEquals(null, converter.fromConnectHeader(TOPIC, "hdr", Schema.OPTIONAL_STRING_SCHEMA, null));
+    }
 }
diff --git a/connect/api/src/test/java/org/apache/kafka/connect/util/ConnectorUtilsTest.java b/connect/api/src/test/java/org/apache/kafka/connect/util/ConnectorUtilsTest.java
index ea53084b18103..771bba0328b44 100644
--- a/connect/api/src/test/java/org/apache/kafka/connect/util/ConnectorUtilsTest.java
+++ b/connect/api/src/test/java/org/apache/kafka/connect/util/ConnectorUtilsTest.java
@@ -55,8 +55,8 @@ public void testGroupPartitions() {
                 Arrays.asList(3),
                 Arrays.asList(4),
                 Arrays.asList(5),
-                Collections.EMPTY_LIST,
-                Collections.EMPTY_LIST), grouped);
+                Collections.emptyList(),
+                Collections.emptyList()), grouped);
     }
 
     @Test(expected = IllegalArgumentException.class)
diff --git a/connect/basic-auth-extension/src/main/java/org/apache/kafka/connect/rest/basic/auth/extension/BasicAuthSecurityRestExtension.java b/connect/basic-auth-extension/src/main/java/org/apache/kafka/connect/rest/basic/auth/extension/BasicAuthSecurityRestExtension.java
new file mode 100644
index 0000000000000..21e13c27492be
--- /dev/null
+++ b/connect/basic-auth-extension/src/main/java/org/apache/kafka/connect/rest/basic/auth/extension/BasicAuthSecurityRestExtension.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 org.apache.kafka.connect.rest.basic.auth.extension;
+
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.apache.kafka.connect.rest.ConnectRestExtension;
+import org.apache.kafka.connect.rest.ConnectRestExtensionContext;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.util.Map;
+
+/**
+ * Provides the ability to authenticate incoming BasicAuth credentials using the configured JAAS {@link
+ * javax.security.auth.spi.LoginModule}. An entry with the name {@code KafkaConnect} is expected in the JAAS config file configured in the
+ * JVM. An implementation of {@link javax.security.auth.spi.LoginModule} needs to be provided in the JAAS config file. The {@code
+ * LoginModule} implementation should configure the {@link javax.security.auth.callback.CallbackHandler} with only {@link
+ * javax.security.auth.callback.NameCallback} and {@link javax.security.auth.callback.PasswordCallback}.
+ *
+ * 
            To use this extension, one needs to add the following config in the {@code worker.properties}
+ * 
            + *     rest.extension.classes = org.apache.kafka.connect.rest.basic.auth.extension.BasicAuthSecurityRestExtension
+ * 
            
+ *
+ * 
             An example JAAS config would look as below
+ * 
            + *         KafkaConnect {
+ *              org.apache.kafka.connect.rest.basic.auth.extension.PropertyFileLoginModule required
+ *              file="/mnt/secret/credentials.properties";
+ *         };
+ *
            
+ *
+ * 
            This is a reference implementation of the {@link ConnectRestExtension} interface. It registers an implementation of {@link
+ * javax.ws.rs.container.ContainerRequestFilter} that does JAAS based authentication of incoming Basic Auth credentials. {@link
+ * ConnectRestExtension} implementations are loaded via the plugin class loader using {@link java.util.ServiceLoader} mechanism and hence
+ * the packaged jar includes {@code META-INF/services/org.apache.kafka.connect.rest.extension.ConnectRestExtension} with the entry
+ * {@code org.apache.kafka.connect.extension.auth.jaas.BasicAuthSecurityRestExtension}
+ *
+ * 
            NOTE: The implementation ships with a default {@link PropertyFileLoginModule} that helps authenticate the request against a
+ * property file. {@link PropertyFileLoginModule} is NOT intended to be used in production since the credentials are stored in PLAINTEXT. One can use
+ * this extension in production by using their own implementation of {@link javax.security.auth.spi.LoginModule} that authenticates against
+ * stores like LDAP, DB, etc.
+ */
+public class BasicAuthSecurityRestExtension implements ConnectRestExtension {
+
+    private static final Logger log = LoggerFactory.getLogger(BasicAuthSecurityRestExtension.class);
+
+    @Override
+    public void register(ConnectRestExtensionContext restPluginContext) {
+        log.trace("Registering JAAS basic auth filter");
+        restPluginContext.configurable().register(JaasBasicAuthFilter.class);
+        log.trace("Finished registering JAAS basic auth filter");
+    }
+
+    @Override
+    public void close() throws IOException {
+
+    }
+
+    @Override
+    public void configure(Map configs) {
+
+    }
+
+    @Override
+    public String version() {
+        return AppInfoParser.getVersion();
+    }
+}
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
new file mode 100644
index 0000000000000..8326239d5049a
--- /dev/null
+++ b/connect/basic-auth-extension/src/main/java/org/apache/kafka/connect/rest/basic/auth/extension/JaasBasicAuthFilter.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.connect.rest.basic.auth.extension;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.regex.Pattern;
+import javax.ws.rs.HttpMethod;
+import org.apache.kafka.common.config.ConfigException;
+import org.apache.kafka.connect.errors.ConnectException;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.util.Base64;
+
+import javax.security.auth.callback.Callback;
+import javax.security.auth.callback.CallbackHandler;
+import javax.security.auth.callback.NameCallback;
+import javax.security.auth.callback.PasswordCallback;
+import javax.security.auth.callback.UnsupportedCallbackException;
+import javax.security.auth.login.LoginContext;
+import javax.security.auth.login.LoginException;
+import javax.ws.rs.container.ContainerRequestContext;
+import javax.ws.rs.container.ContainerRequestFilter;
+import javax.ws.rs.core.Response;
+
+public class JaasBasicAuthFilter implements ContainerRequestFilter {
+
+    private static final Logger log = LoggerFactory.getLogger(JaasBasicAuthFilter.class);
+    private static final Pattern TASK_REQUEST_PATTERN = Pattern.compile("/?connectors/([^/]+)/tasks/?");
+    private static final String CONNECT_LOGIN_MODULE = "KafkaConnect";
+
+    static final String AUTHORIZATION = "Authorization";
+
+    @Override
+    public void filter(ContainerRequestContext requestContext) throws IOException {
+        if (isInternalTaskConfigRequest(requestContext)) {
+            log.trace("Skipping authentication for internal request");
+            return;
+        }
+
+        try {
+            log.debug("Authenticating request");
+            LoginContext loginContext =
+                new LoginContext(CONNECT_LOGIN_MODULE, new BasicAuthCallBackHandler(
+                    requestContext.getHeaderString(AUTHORIZATION)));
+            loginContext.login();
+        } catch (LoginException | ConfigException e) {
+            // Log at debug here in order to avoid polluting log files whenever someone mistypes their credentials
+            log.debug("Request failed authentication", e);
+            requestContext.abortWith(
+                Response.status(Response.Status.UNAUTHORIZED)
+                    .entity("User cannot access the resource.")
+                    .build());
+        }
+    }
+
+    private static boolean isInternalTaskConfigRequest(ContainerRequestContext requestContext) {
+        return requestContext.getMethod().equals(HttpMethod.POST)
+            && TASK_REQUEST_PATTERN.matcher(requestContext.getUriInfo().getPath()).matches();
+    }
+
+
+    public static class BasicAuthCallBackHandler implements CallbackHandler {
+
+        private static final String BASIC = "basic";
+        private static final char COLON = ':';
+        private static final char SPACE = ' ';
+        private String username;
+        private String password;
+
+        public BasicAuthCallBackHandler(String credentials) {
+            if (credentials == null) {
+                log.trace("No credentials were provided with the request");
+                return;
+            }
+
+            int space = credentials.indexOf(SPACE);
+            if (space <= 0) {
+                log.trace("Request credentials were malformed; no space present in value for authorization header");
+                return;
+            }
+
+            String method = credentials.substring(0, space);
+            if (!BASIC.equalsIgnoreCase(method)) {
+                log.trace("Request credentials used {} authentication, but only {} supported; ignoring", method, BASIC);
+                return;
+            }
+
+            credentials = credentials.substring(space + 1);
+            credentials = new String(Base64.getDecoder().decode(credentials),
+                                     StandardCharsets.UTF_8);
+            int i = credentials.indexOf(COLON);
+            if (i <= 0) {
+                log.trace("Request credentials were malformed; no colon present between username and password");
+                return;
+            }
+
+            username = credentials.substring(0, i);
+            password = credentials.substring(i + 1);
+        }
+
+        @Override
+        public void handle(Callback[] callbacks) throws UnsupportedCallbackException {
+            List unsupportedCallbacks = new ArrayList<>();
+            for (Callback callback : callbacks) {
+                if (callback instanceof NameCallback) {
+                    ((NameCallback) callback).setName(username);
+                } else if (callback instanceof PasswordCallback) {
+                    ((PasswordCallback) callback).setPassword(password != null
+                        ? password.toCharArray()
+                        : null
+                    );
+                } else {
+                    unsupportedCallbacks.add(callback);
+                }
+            }
+            if (!unsupportedCallbacks.isEmpty())
+                throw new ConnectException(String.format(
+                    "Unsupported callbacks %s; request authentication will fail. "
+                        + "This indicates the Connect worker was configured with a JAAS "
+                        + "LoginModule that is incompatible with the %s, and will need to be "
+                        + "corrected and restarted.",
+                    unsupportedCallbacks,
+                    BasicAuthSecurityRestExtension.class.getSimpleName()
+                ));
+        }
+    }
+}
diff --git a/connect/basic-auth-extension/src/main/java/org/apache/kafka/connect/rest/basic/auth/extension/PropertyFileLoginModule.java b/connect/basic-auth-extension/src/main/java/org/apache/kafka/connect/rest/basic/auth/extension/PropertyFileLoginModule.java
new file mode 100644
index 0000000000000..8a26dc352438c
--- /dev/null
+++ b/connect/basic-auth-extension/src/main/java/org/apache/kafka/connect/rest/basic/auth/extension/PropertyFileLoginModule.java
@@ -0,0 +1,154 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.rest.basic.auth.extension;
+
+import org.apache.kafka.common.config.ConfigException;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.file.Files;
+import java.nio.file.Paths;
+import java.util.Map;
+import java.util.Properties;
+import java.util.concurrent.ConcurrentHashMap;
+
+import javax.security.auth.Subject;
+import javax.security.auth.callback.Callback;
+import javax.security.auth.callback.CallbackHandler;
+import javax.security.auth.callback.NameCallback;
+import javax.security.auth.callback.PasswordCallback;
+import javax.security.auth.login.LoginException;
+import javax.security.auth.spi.LoginModule;
+
+/**
+ * {@link PropertyFileLoginModule} authenticates against a properties file.
+ * The credentials should be stored in the format {username}={password} in the properties file.
+ * The absolute path of the file needs to specified using the option file
+ *
+ * 
            NOTE: This implementation is NOT intended to be used in production since the credentials are stored in PLAINTEXT in the
+ * properties file.
+ */
+public class PropertyFileLoginModule implements LoginModule {
+    private static final Logger log = LoggerFactory.getLogger(PropertyFileLoginModule.class);
+
+    private CallbackHandler callbackHandler;
+    private static final String FILE_OPTIONS = "file";
+    private String fileName;
+    private boolean authenticated;
+
+    private static Map credentialPropertiesMap = new ConcurrentHashMap<>();
+
+    @Override
+    public void initialize(Subject subject, CallbackHandler callbackHandler, Map sharedState, Map options) {
+        this.callbackHandler = callbackHandler;
+        fileName = (String) options.get(FILE_OPTIONS);
+        if (fileName == null || fileName.trim().isEmpty()) {
+            throw new ConfigException("Property Credentials file must be specified");
+        }
+
+        if (!credentialPropertiesMap.containsKey(fileName)) {
+            log.trace("Opening credential properties file '{}'", fileName);
+            Properties credentialProperties = new Properties();
+            try {
+                try (InputStream inputStream = Files.newInputStream(Paths.get(fileName))) {
+                    log.trace("Parsing credential properties file '{}'", fileName);
+                    credentialProperties.load(inputStream);
+                }
+                credentialPropertiesMap.putIfAbsent(fileName, credentialProperties);
+                if (credentialProperties.isEmpty())
+                    log.warn("Credential properties file '{}' is empty; all requests will be permitted",
+                        fileName);
+            } catch (IOException e) {
+                log.error("Error loading credentials file ", e);
+                throw new ConfigException("Error loading Property Credentials file");
+            }
+        } else {
+            log.trace(
+                "Credential properties file '{}' has already been opened and parsed; will read from cached, in-memory store",
+                fileName);
+        }
+    }
+
+    @Override
+    public boolean login() throws LoginException {
+        Callback[] callbacks = configureCallbacks();
+        try {
+            log.trace("Authenticating user; invoking JAAS login callbacks");
+            callbackHandler.handle(callbacks);
+        } catch (Exception e) {
+            log.warn("Authentication failed while invoking JAAS login callbacks", e);
+            throw new LoginException(e.getMessage());
+        }
+
+        String username = ((NameCallback) callbacks[0]).getName();
+        char[] passwordChars = ((PasswordCallback) callbacks[1]).getPassword();
+        String password = passwordChars != null ? new String(passwordChars) : null;
+        Properties credentialProperties = credentialPropertiesMap.get(fileName);
+
+        if (credentialProperties.isEmpty()) {
+            log.trace("Not validating credentials for user '{}' as credential properties file '{}' is empty",
+                username,
+                fileName);
+            authenticated = true;
+        } else if (username == null) {
+            log.trace("No credentials were provided or the provided credentials were malformed");
+            authenticated = false;
+        } else if (password != null && password.equals(credentialProperties.get(username))) {
+            log.trace("Credentials provided for user '{}' match those present in the credential properties file '{}'",
+                username,
+                fileName);
+            authenticated = true;
+        } else if (!credentialProperties.containsKey(username)) {
+            log.trace("User '{}' is not present in the credential properties file '{}'",
+                username,
+                fileName);
+            authenticated = false;
+        } else {
+            log.trace("Credentials provided for user '{}' do not match those present in the credential properties file '{}'",
+                username,
+                fileName);
+            authenticated = false;
+        }
+
+        return authenticated;
+    }
+
+    @Override
+    public boolean commit() throws LoginException {
+        return authenticated;
+    }
+
+    @Override
+    public boolean abort() throws LoginException {
+        return true;
+    }
+
+    @Override
+    public boolean logout() throws LoginException {
+        return true;
+    }
+
+    private Callback[] configureCallbacks() {
+        Callback[] callbacks = new Callback[2];
+        callbacks[0] = new NameCallback("Enter user name");
+        callbacks[1] = new PasswordCallback("Enter password", false);
+        return callbacks;
+    }
+}
diff --git a/connect/basic-auth-extension/src/main/resources/META-INF/services/org.apache.kafka.connect.rest.ConnectRestExtension b/connect/basic-auth-extension/src/main/resources/META-INF/services/org.apache.kafka.connect.rest.ConnectRestExtension
new file mode 100644
index 0000000000000..ba7ae5b580d80
--- /dev/null
+++ b/connect/basic-auth-extension/src/main/resources/META-INF/services/org.apache.kafka.connect.rest.ConnectRestExtension
@@ -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.
+
+org.apache.kafka.connect.rest.basic.auth.extension.BasicAuthSecurityRestExtension
\ No newline at end of file
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
new file mode 100644
index 0000000000000..0c39c5ffe69c0
--- /dev/null
+++ b/connect/basic-auth-extension/src/test/java/org/apache/kafka/connect/rest/basic/auth/extension/JaasBasicAuthFilterTest.java
@@ -0,0 +1,221 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.kafka.connect.rest.basic.auth.extension;
+
+import javax.security.auth.callback.Callback;
+import javax.security.auth.callback.CallbackHandler;
+import javax.security.auth.callback.ChoiceCallback;
+import javax.ws.rs.HttpMethod;
+import javax.ws.rs.core.UriInfo;
+import org.apache.kafka.common.security.JaasUtils;
+import org.apache.kafka.connect.errors.ConnectException;
+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;
+
+import java.io.File;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.util.ArrayList;
+import java.util.Base64;
+import java.util.List;
+
+import javax.security.auth.login.Configuration;
+import javax.ws.rs.container.ContainerRequestContext;
+import javax.ws.rs.core.Response;
+
+import static org.powermock.api.easymock.PowerMock.replayAll;
+
+@RunWith(PowerMockRunner.class)
+@PowerMockIgnore("javax.*")
+public class JaasBasicAuthFilterTest {
+
+    @MockStrict
+    private ContainerRequestContext requestContext;
+
+    private JaasBasicAuthFilter jaasBasicAuthFilter = new JaasBasicAuthFilter();
+    private String previousJaasConfig;
+    private Configuration previousConfiguration;
+
+    @MockStrict
+    private UriInfo uriInfo;
+
+    @Before
+    public void setup() {
+        EasyMock.reset(requestContext);
+    }
+
+    @After
+    public void tearDown() {
+        if (previousJaasConfig != null) {
+            System.setProperty(JaasUtils.JAVA_LOGIN_CONFIG_PARAM, previousJaasConfig);
+        }
+        Configuration.setConfiguration(previousConfiguration);
+    }
+
+    @Test
+    public void testSuccess() throws IOException {
+        File credentialFile = File.createTempFile("credential", ".properties");
+        credentialFile.deleteOnExit();
+        List lines = new ArrayList<>();
+        lines.add("user=password");
+        lines.add("user1=password1");
+        Files.write(credentialFile.toPath(), lines, StandardCharsets.UTF_8);
+
+        setupJaasConfig("KafkaConnect", credentialFile.getPath(), true);
+        setMock("Basic", "user", "password", false);
+
+        jaasBasicAuthFilter.filter(requestContext);
+    }
+
+
+    @Test
+    public void testBadCredential() throws IOException {
+        setMock("Basic", "user1", "password", true);
+        jaasBasicAuthFilter.filter(requestContext);
+    }
+
+    @Test
+    public void testBadPassword() throws IOException {
+        setMock("Basic", "user", "password1", true);
+        jaasBasicAuthFilter.filter(requestContext);
+    }
+
+    @Test
+    public void testUnknownBearer() throws IOException {
+        setMock("Unknown", "user", "password", true);
+        jaasBasicAuthFilter.filter(requestContext);
+    }
+
+    @Test
+    public void testUnknownLoginModule() throws IOException {
+        setupJaasConfig("KafkaConnect1", "/tmp/testcrednetial", true);
+        Configuration.setConfiguration(null);
+        setMock("Basic", "user", "password", true);
+        jaasBasicAuthFilter.filter(requestContext);
+    }
+
+    @Test
+    public void testUnknownCredentialsFile() throws IOException {
+        setupJaasConfig("KafkaConnect", "/tmp/testcrednetial", true);
+        Configuration.setConfiguration(null);
+        setMock("Basic", "user", "password", true);
+        jaasBasicAuthFilter.filter(requestContext);
+    }
+
+    @Test
+    public void testEmptyCredentialsFile() throws IOException {
+        File jaasConfigFile = File.createTempFile("ks-jaas-", ".conf");
+        jaasConfigFile.deleteOnExit();
+        System.setProperty(JaasUtils.JAVA_LOGIN_CONFIG_PARAM, jaasConfigFile.getPath());
+        setupJaasConfig("KafkaConnect", "", true);
+        Configuration.setConfiguration(null);
+        setMock("Basic", "user", "password", true);
+        jaasBasicAuthFilter.filter(requestContext);
+    }
+
+    @Test
+    public void testNoFileOption() throws IOException {
+        File jaasConfigFile = File.createTempFile("ks-jaas-", ".conf");
+        jaasConfigFile.deleteOnExit();
+        System.setProperty(JaasUtils.JAVA_LOGIN_CONFIG_PARAM, jaasConfigFile.getPath());
+        setupJaasConfig("KafkaConnect", "", false);
+        Configuration.setConfiguration(null);
+        setMock("Basic", "user", "password", true);
+        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);
+    }
+
+    @Test(expected = ConnectException.class)
+    public void testUnsupportedCallback() throws Exception {
+        String authHeader = authHeader("basic", "user", "pwd");
+        CallbackHandler callbackHandler = new JaasBasicAuthFilter.BasicAuthCallBackHandler(authHeader);
+        Callback unsupportedCallback = new ChoiceCallback(
+            "You take the blue pill... the story ends, you wake up in your bed and believe whatever you want to believe. " 
+                + "You take the red pill... you stay in Wonderland, and I show you how deep the rabbit hole goes.",
+            new String[] {"blue pill", "red pill"},
+            1,
+            true
+        );
+        callbackHandler.handle(new Callback[] {unsupportedCallback});
+    }
+
+    private String authHeader(String authorization, String username, String password) {
+        return authorization + " " + Base64.getEncoder().encodeToString((username + ":" + password).getBytes());
+    }
+
+    private void setMock(String authorization, String username, String password, boolean exceptionCase) {
+        EasyMock.expect(requestContext.getMethod()).andReturn(HttpMethod.GET);
+        EasyMock.expect(requestContext.getHeaderString(JaasBasicAuthFilter.AUTHORIZATION))
+            .andReturn(authHeader(authorization, username, password));
+        if (exceptionCase) {
+            requestContext.abortWith(EasyMock.anyObject(Response.class));
+            EasyMock.expectLastCall();
+        }
+        replayAll();
+    }
+
+    private void setupJaasConfig(String loginModule, String credentialFilePath, boolean includeFileOptions) throws IOException {
+        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 ");
+        if (includeFileOptions) {
+            lines.add("file=\"" + credentialFilePath + "\"");
+        }
+        lines.add(";};");
+        Files.write(jaasConfigFile.toPath(), lines, StandardCharsets.UTF_8);
+        previousConfiguration = Configuration.getConfiguration();
+        Configuration.setConfiguration(null);
+    }
+
+}
diff --git a/connect/file/src/main/java/org/apache/kafka/connect/file/FileStreamSinkConnector.java b/connect/file/src/main/java/org/apache/kafka/connect/file/FileStreamSinkConnector.java
index 4ae7f4bb4f424..136e8998c200c 100644
--- a/connect/file/src/main/java/org/apache/kafka/connect/file/FileStreamSinkConnector.java
+++ b/connect/file/src/main/java/org/apache/kafka/connect/file/FileStreamSinkConnector.java
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.connect.file;
 
+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;
@@ -47,7 +48,8 @@ public String version() {
 
     @Override
     public void start(Map props) {
-        filename = props.get(FILE_CONFIG);
+        AbstractConfig parsedConfig = new AbstractConfig(CONFIG_DEF, props);
+        filename = parsedConfig.getString(FILE_CONFIG);
     }
 
     @Override
diff --git a/connect/file/src/main/java/org/apache/kafka/connect/file/FileStreamSinkTask.java b/connect/file/src/main/java/org/apache/kafka/connect/file/FileStreamSinkTask.java
index b663f8118bd8c..3d1d2b8543529 100644
--- a/connect/file/src/main/java/org/apache/kafka/connect/file/FileStreamSinkTask.java
+++ b/connect/file/src/main/java/org/apache/kafka/connect/file/FileStreamSinkTask.java
@@ -24,11 +24,12 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.io.FileNotFoundException;
-import java.io.FileOutputStream;
+import java.io.IOException;
 import java.io.PrintStream;
-import java.io.UnsupportedEncodingException;
 import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Paths;
+import java.nio.file.StandardOpenOption;
 import java.util.Collection;
 import java.util.Map;
 
@@ -62,10 +63,12 @@ public void start(Map props) {
             outputStream = System.out;
         } else {
             try {
-                outputStream = new PrintStream(new FileOutputStream(filename, true), false,
+                outputStream = new PrintStream(
+                    Files.newOutputStream(Paths.get(filename), StandardOpenOption.CREATE, StandardOpenOption.APPEND),
+                    false,
                     StandardCharsets.UTF_8.name());
-            } catch (FileNotFoundException | UnsupportedEncodingException e) {
-                throw new ConnectException("Couldn't find or create file for FileStreamSinkTask", e);
+            } catch (IOException e) {
+                throw new ConnectException("Couldn't find or create file '" + filename + "' for FileStreamSinkTask", e);
             }
         }
     }
diff --git a/connect/file/src/main/java/org/apache/kafka/connect/file/FileStreamSourceConnector.java b/connect/file/src/main/java/org/apache/kafka/connect/file/FileStreamSourceConnector.java
index 335fe9255197c..74b5f7c09ed23 100644
--- a/connect/file/src/main/java/org/apache/kafka/connect/file/FileStreamSourceConnector.java
+++ b/connect/file/src/main/java/org/apache/kafka/connect/file/FileStreamSourceConnector.java
@@ -16,12 +16,13 @@
  */
 package org.apache.kafka.connect.file;
 
+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.ConfigException;
 import org.apache.kafka.common.utils.AppInfoParser;
 import org.apache.kafka.connect.connector.Task;
-import org.apache.kafka.connect.errors.ConnectException;
 import org.apache.kafka.connect.source.SourceConnector;
 
 import java.util.ArrayList;
@@ -36,13 +37,19 @@
 public class FileStreamSourceConnector extends SourceConnector {
     public static final String TOPIC_CONFIG = "topic";
     public static final String FILE_CONFIG = "file";
+    public static final String TASK_BATCH_SIZE_CONFIG = "batch.size";
+
+    public static final int DEFAULT_TASK_BATCH_SIZE = 2000;
 
     private static final ConfigDef CONFIG_DEF = new ConfigDef()
         .define(FILE_CONFIG, Type.STRING, null, Importance.HIGH, "Source filename. If not specified, the standard input will be used")
-        .define(TOPIC_CONFIG, Type.STRING, Importance.HIGH, "The topic to publish data to");
+        .define(TOPIC_CONFIG, Type.LIST, Importance.HIGH, "The topic to publish data to")
+        .define(TASK_BATCH_SIZE_CONFIG, Type.INT, DEFAULT_TASK_BATCH_SIZE, Importance.LOW,
+                "The maximum number of records the Source task can read from file one time");
 
     private String filename;
     private String topic;
+    private int batchSize;
 
     @Override
     public String version() {
@@ -51,12 +58,14 @@ public String version() {
 
     @Override
     public void start(Map props) {
-        filename = props.get(FILE_CONFIG);
-        topic = props.get(TOPIC_CONFIG);
-        if (topic == null || topic.isEmpty())
-            throw new ConnectException("FileStreamSourceConnector configuration must include 'topic' setting");
-        if (topic.contains(","))
-            throw new ConnectException("FileStreamSourceConnector should only have a single topic when used as a source.");
+        AbstractConfig parsedConfig = new AbstractConfig(CONFIG_DEF, props);
+        filename = parsedConfig.getString(FILE_CONFIG);
+        List topics = parsedConfig.getList(TOPIC_CONFIG);
+        if (topics.size() != 1) {
+            throw new ConfigException("'topic' in FileStreamSourceConnector configuration requires definition of a single topic");
+        }
+        topic = topics.get(0);
+        batchSize = parsedConfig.getInt(TASK_BATCH_SIZE_CONFIG);
     }
 
     @Override
@@ -72,6 +81,7 @@ public List> taskConfigs(int maxTasks) {
         if (filename != null)
             config.put(FILE_CONFIG, filename);
         config.put(TOPIC_CONFIG, topic);
+        config.put(TASK_BATCH_SIZE_CONFIG, String.valueOf(batchSize));
         configs.add(config);
         return configs;
     }
diff --git a/connect/file/src/main/java/org/apache/kafka/connect/file/FileStreamSourceTask.java b/connect/file/src/main/java/org/apache/kafka/connect/file/FileStreamSourceTask.java
index 8edf385611a54..582889b54970b 100644
--- a/connect/file/src/main/java/org/apache/kafka/connect/file/FileStreamSourceTask.java
+++ b/connect/file/src/main/java/org/apache/kafka/connect/file/FileStreamSourceTask.java
@@ -17,12 +17,13 @@
 package org.apache.kafka.connect.file;
 
 import java.io.BufferedReader;
-import java.io.FileInputStream;
-import java.io.FileNotFoundException;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.InputStreamReader;
 import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.NoSuchFileException;
+import java.nio.file.Paths;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
@@ -50,6 +51,7 @@ public class FileStreamSourceTask extends SourceTask {
     private char[] buffer = new char[1024];
     private int offset = 0;
     private String topic = null;
+    private int batchSize = FileStreamSourceConnector.DEFAULT_TASK_BATCH_SIZE;
 
     private Long streamOffset;
 
@@ -67,16 +69,17 @@ public void start(Map props) {
             streamOffset = null;
             reader = new BufferedReader(new InputStreamReader(stream, StandardCharsets.UTF_8));
         }
+        // Missing topic or parsing error is not possible because we've parsed the config in the
+        // Connector
         topic = props.get(FileStreamSourceConnector.TOPIC_CONFIG);
-        if (topic == null)
-            throw new ConnectException("FileStreamSourceTask config missing topic setting");
+        batchSize = Integer.parseInt(props.get(FileStreamSourceConnector.TASK_BATCH_SIZE_CONFIG));
     }
 
     @Override
     public List poll() throws InterruptedException {
         if (stream == null) {
             try {
-                stream = new FileInputStream(filename);
+                stream = Files.newInputStream(Paths.get(filename));
                 Map offset = context.offsetStorageReader().offset(Collections.singletonMap(FILENAME_FIELD, filename));
                 if (offset != null) {
                     Object lastRecordedOffset = offset.get(POSITION_FIELD);
@@ -90,7 +93,7 @@ public List poll() throws InterruptedException {
                                 long skipped = stream.skip(skipLeft);
                                 skipLeft -= skipped;
                             } catch (IOException e) {
-                                log.error("Error while trying to seek to previous offset in file: ", e);
+                                log.error("Error while trying to seek to previous offset in file {}: ", filename, e);
                                 throw new ConnectException(e);
                             }
                         }
@@ -102,12 +105,15 @@ public List poll() throws InterruptedException {
                 }
                 reader = new BufferedReader(new InputStreamReader(stream, StandardCharsets.UTF_8));
                 log.debug("Opened {} for reading", logFilename());
-            } catch (FileNotFoundException e) {
+            } catch (NoSuchFileException e) {
                 log.warn("Couldn't find file {} for FileStreamSourceTask, sleeping to wait for it to be created", logFilename());
                 synchronized (this) {
                     this.wait(1000);
                 }
                 return null;
+            } catch (IOException e) {
+                log.error("Error while trying to open file {}: ", filename, e);
+                throw new ConnectException(e);
             }
         }
 
@@ -146,6 +152,10 @@ public List poll() throws InterruptedException {
                                 records = new ArrayList<>();
                             records.add(new SourceRecord(offsetKey(filename), offsetValue(streamOffset), topic, null,
                                     null, null, VALUE_SCHEMA, line, System.currentTimeMillis()));
+
+                            if (records.size() >= batchSize) {
+                                return records;
+                            }
                         }
                     } while (line != null);
                 }
diff --git a/connect/file/src/test/java/org/apache/kafka/connect/file/FileStreamSinkConnectorTest.java b/connect/file/src/test/java/org/apache/kafka/connect/file/FileStreamSinkConnectorTest.java
index c06c99144be91..2348454481648 100644
--- a/connect/file/src/test/java/org/apache/kafka/connect/file/FileStreamSinkConnectorTest.java
+++ b/connect/file/src/test/java/org/apache/kafka/connect/file/FileStreamSinkConnectorTest.java
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.connect.file;
 
+import org.apache.kafka.common.config.ConfigValue;
 import org.apache.kafka.connect.connector.ConnectorContext;
 import org.apache.kafka.connect.sink.SinkConnector;
 import org.easymock.EasyMockSupport;
@@ -49,6 +50,16 @@ public void setup() {
         sinkProperties.put(FileStreamSinkConnector.FILE_CONFIG, FILENAME);
     }
 
+    @Test
+    public void testConnectorConfigValidation() {
+        replayAll();
+        List configValues = connector.config().validate(sinkProperties);
+        for (ConfigValue val : configValues) {
+            assertEquals("Config property errors: " + val.errorMessages(), 0, val.errorMessages().size());
+        }
+        verifyAll();
+    }
+
     @Test
     public void testSinkTasks() {
         replayAll();
diff --git a/connect/file/src/test/java/org/apache/kafka/connect/file/FileStreamSinkTaskTest.java b/connect/file/src/test/java/org/apache/kafka/connect/file/FileStreamSinkTaskTest.java
index c7ec9dac572c2..a5142a13324ce 100644
--- a/connect/file/src/test/java/org/apache/kafka/connect/file/FileStreamSinkTaskTest.java
+++ b/connect/file/src/test/java/org/apache/kafka/connect/file/FileStreamSinkTaskTest.java
@@ -21,12 +21,21 @@
 import org.apache.kafka.connect.data.Schema;
 import org.apache.kafka.connect.sink.SinkRecord;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
 
+import java.io.BufferedReader;
 import java.io.ByteArrayOutputStream;
+import java.io.File;
+import java.io.IOException;
 import java.io.PrintStream;
+import java.nio.file.Files;
+import java.nio.file.Paths;
 import java.util.Arrays;
 import java.util.HashMap;
+import java.util.Map;
+import java.util.UUID;
 
 import static org.junit.Assert.assertEquals;
 
@@ -36,11 +45,17 @@ public class FileStreamSinkTaskTest {
     private ByteArrayOutputStream os;
     private PrintStream printStream;
 
+    @Rule
+    public TemporaryFolder topDir = new TemporaryFolder();
+    private String outputFile;
+
     @Before
-    public void setup() {
+    public void setup() throws Exception {
         os = new ByteArrayOutputStream();
         printStream = new PrintStream(os);
         task = new FileStreamSinkTask(printStream);
+        File outputDir = topDir.newFolder("file-stream-sink-" + UUID.randomUUID().toString());
+        outputFile = outputDir.getCanonicalPath() + "/connect.output";
     }
 
     @Test
@@ -66,4 +81,39 @@ public void testPutFlush() {
         task.flush(offsets);
         assertEquals("line1" + newLine + "line2" + newLine + "line3" + newLine, os.toString());
     }
+
+    @Test
+    public void testStart() throws IOException {
+        task = new FileStreamSinkTask();
+        Map props = new HashMap<>();
+        props.put(FileStreamSinkConnector.FILE_CONFIG, outputFile);
+        task.start(props);
+
+        HashMap offsets = new HashMap<>();
+        task.put(Arrays.asList(
+                new SinkRecord("topic1", 0, null, null, Schema.STRING_SCHEMA, "line0", 1)
+        ));
+        offsets.put(new TopicPartition("topic1", 0), new OffsetAndMetadata(1L));
+        task.flush(offsets);
+
+        int numLines = 3;
+        String[] lines = new String[numLines];
+        int i = 0;
+        try (BufferedReader reader = Files.newBufferedReader(Paths.get(outputFile))) {
+            lines[i++] = reader.readLine();
+            task.put(Arrays.asList(
+                    new SinkRecord("topic1", 0, null, null, Schema.STRING_SCHEMA, "line1", 2),
+                    new SinkRecord("topic2", 0, null, null, Schema.STRING_SCHEMA, "line2", 1)
+            ));
+            offsets.put(new TopicPartition("topic1", 0), new OffsetAndMetadata(2L));
+            offsets.put(new TopicPartition("topic2", 0), new OffsetAndMetadata(1L));
+            task.flush(offsets);
+            lines[i++] = reader.readLine();
+            lines[i++] = reader.readLine();
+        }
+
+        while (--i >= 0) {
+            assertEquals("line" + i, lines[i]);
+        }
+    }
 }
diff --git a/connect/file/src/test/java/org/apache/kafka/connect/file/FileStreamSourceConnectorTest.java b/connect/file/src/test/java/org/apache/kafka/connect/file/FileStreamSourceConnectorTest.java
index 69a94a894b285..d7efdacd2f62a 100644
--- a/connect/file/src/test/java/org/apache/kafka/connect/file/FileStreamSourceConnectorTest.java
+++ b/connect/file/src/test/java/org/apache/kafka/connect/file/FileStreamSourceConnectorTest.java
@@ -16,8 +16,9 @@
  */
 package org.apache.kafka.connect.file;
 
+import org.apache.kafka.common.config.ConfigException;
+import org.apache.kafka.common.config.ConfigValue;
 import org.apache.kafka.connect.connector.ConnectorContext;
-import org.apache.kafka.connect.errors.ConnectException;
 import org.easymock.EasyMock;
 import org.easymock.EasyMockSupport;
 import org.junit.Before;
@@ -51,6 +52,16 @@ public void setup() {
         sourceProperties.put(FileStreamSourceConnector.FILE_CONFIG, FILENAME);
     }
 
+    @Test
+    public void testConnectorConfigValidation() {
+        replayAll();
+        List configValues = connector.config().validate(sourceProperties);
+        for (ConfigValue val : configValues) {
+            assertEquals("Config property errors: " + val.errorMessages(), 0, val.errorMessages().size());
+        }
+        verifyAll();
+    }
+
     @Test
     public void testSourceTasks() {
         replayAll();
@@ -87,7 +98,7 @@ public void testSourceTasksStdin() {
         EasyMock.verify(ctx);
     }
 
-    @Test(expected = ConnectException.class)
+    @Test(expected = ConfigException.class)
     public void testMultipleSourcesInvalid() {
         sourceProperties.put(FileStreamSourceConnector.TOPIC_CONFIG, MULTIPLE_TOPICS);
         connector.start(sourceProperties);
@@ -102,4 +113,23 @@ public void testTaskClass() {
 
         EasyMock.verify(ctx);
     }
+
+    @Test(expected = ConfigException.class)
+    public void testMissingTopic() {
+        sourceProperties.remove(FileStreamSourceConnector.TOPIC_CONFIG);
+        connector.start(sourceProperties);
+    }
+
+    @Test(expected = ConfigException.class)
+    public void testBlankTopic() {
+        // Because of trimming this tests is same as testing for empty string.
+        sourceProperties.put(FileStreamSourceConnector.TOPIC_CONFIG, "     ");
+        connector.start(sourceProperties);
+    }
+
+    @Test(expected = ConfigException.class)
+    public void testInvalidBatchSize() {
+        sourceProperties.put(FileStreamSourceConnector.TASK_BATCH_SIZE_CONFIG, "abcd");
+        connector.start(sourceProperties);
+    }
 }
diff --git a/connect/file/src/test/java/org/apache/kafka/connect/file/FileStreamSourceTaskTest.java b/connect/file/src/test/java/org/apache/kafka/connect/file/FileStreamSourceTaskTest.java
index cde6c43ae9a27..feacf8f428098 100644
--- a/connect/file/src/test/java/org/apache/kafka/connect/file/FileStreamSourceTaskTest.java
+++ b/connect/file/src/test/java/org/apache/kafka/connect/file/FileStreamSourceTaskTest.java
@@ -16,7 +16,6 @@
  */
 package org.apache.kafka.connect.file;
 
-import org.apache.kafka.connect.errors.ConnectException;
 import org.apache.kafka.connect.source.SourceRecord;
 import org.apache.kafka.connect.source.SourceTaskContext;
 import org.apache.kafka.connect.storage.OffsetStorageReader;
@@ -28,8 +27,9 @@
 
 import java.io.ByteArrayInputStream;
 import java.io.File;
-import java.io.FileOutputStream;
 import java.io.IOException;
+import java.io.OutputStream;
+import java.nio.file.Files;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
@@ -55,6 +55,7 @@ public void setup() throws IOException {
         config = new HashMap<>();
         config.put(FileStreamSourceConnector.FILE_CONFIG, tempFile.getAbsolutePath());
         config.put(FileStreamSourceConnector.TOPIC_CONFIG, TOPIC);
+        config.put(FileStreamSourceConnector.TASK_BATCH_SIZE_CONFIG, String.valueOf(FileStreamSourceConnector.DEFAULT_TASK_BATCH_SIZE));
         task = new FileStreamSourceTask();
         offsetStorageReader = createMock(OffsetStorageReader.class);
         context = createMock(SourceTaskContext.class);
@@ -81,7 +82,7 @@ public void testNormalLifecycle() throws InterruptedException, IOException {
 
         task.start(config);
 
-        FileOutputStream os = new FileOutputStream(tempFile);
+        OutputStream os = Files.newOutputStream(tempFile.toPath());
         assertEquals(null, task.poll());
         os.write("partial line".getBytes());
         os.flush();
@@ -127,12 +128,28 @@ public void testNormalLifecycle() throws InterruptedException, IOException {
         task.stop();
     }
 
-    @Test(expected = ConnectException.class)
-    public void testMissingTopic() throws InterruptedException {
+    @Test
+    public void testBatchSize() throws IOException, InterruptedException {
+        expectOffsetLookupReturnNone();
         replay();
 
-        config.remove(FileStreamSourceConnector.TOPIC_CONFIG);
+        config.put(FileStreamSourceConnector.TASK_BATCH_SIZE_CONFIG, "5000");
         task.start(config);
+
+        OutputStream os = Files.newOutputStream(tempFile.toPath());
+        for (int i = 0; i < 10_000; i++) {
+            os.write("Neque porro quisquam est qui dolorem ipsum quia dolor sit amet, consectetur, adipisci velit...\n".getBytes());
+        }
+        os.flush();
+
+        List records = task.poll();
+        assertEquals(5000, records.size());
+
+        records = task.poll();
+        assertEquals(5000, records.size());
+
+        os.close();
+        task.stop();
     }
 
     @Test
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 a7b37f11832c8..1a175381e8f11 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;
@@ -23,6 +24,7 @@
 import org.apache.kafka.common.cache.Cache;
 import org.apache.kafka.common.cache.LRUCache;
 import org.apache.kafka.common.cache.SynchronizedCache;
+import org.apache.kafka.common.config.ConfigDef;
 import org.apache.kafka.common.errors.SerializationException;
 import org.apache.kafka.connect.data.SchemaBuilder;
 import org.apache.kafka.connect.data.Schema;
@@ -36,26 +38,32 @@
 import org.apache.kafka.connect.data.Date;
 import org.apache.kafka.connect.errors.DataException;
 import org.apache.kafka.connect.storage.Converter;
+import org.apache.kafka.connect.storage.ConverterType;
+import org.apache.kafka.connect.storage.HeaderConverter;
+import org.apache.kafka.connect.storage.StringConverterConfig;
 
 import java.io.IOException;
 import java.math.BigDecimal;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.EnumMap;
 import java.util.HashMap;
 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.
+ * 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}
+ * configuration option.
+ *
+ * This implementation currently does nothing with the topic names or header names.
  */
-public class JsonConverter implements Converter {
-    private static final String SCHEMAS_ENABLE_CONFIG = "schemas.enable";
-    private static final boolean SCHEMAS_ENABLE_DEFAULT = true;
-    private static final String SCHEMAS_CACHE_SIZE_CONFIG = "schemas.cache.size";
-    private static final int SCHEMAS_CACHE_SIZE_DEFAULT = 1000;
+public class JsonConverter implements Converter, HeaderConverter {
 
-    private static final HashMap TO_CONNECT_CONVERTERS = new HashMap<>();
+    private static final Map TO_CONNECT_CONVERTERS = new EnumMap<>(Schema.Type.class);
 
     static {
         TO_CONNECT_CONVERTERS.put(Schema.Type.BOOLEAN, new JsonToConnectTypeConverter() {
@@ -180,114 +188,162 @@ 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 = SCHEMAS_ENABLE_DEFAULT;
-    private int cacheSize = 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() {
+        return JsonConverterConfig.configDef();
+    }
+
+    @Override
+    public void configure(Map configs) {
+        config = new JsonConverterConfig(configs);
+
+        serializer.configure(configs, config.type() == ConverterType.KEY);
+        deserializer.configure(configs, config.type() == ConverterType.KEY);
+
+        fromConnectSchemaCache = new SynchronizedCache<>(new LRUCache<>(config.schemaCacheSize()));
+        toConnectSchemaCache = new SynchronizedCache<>(new LRUCache<>(config.schemaCacheSize()));
+    }
 
     @Override
     public void configure(Map configs, boolean isKey) {
-        Object enableConfigsVal = configs.get(SCHEMAS_ENABLE_CONFIG);
-        if (enableConfigsVal != null)
-            enableSchemas = enableConfigsVal.toString().equals("true");
-
-        serializer.configure(configs, isKey);
-        deserializer.configure(configs, isKey);
-
-        Object cacheSizeVal = configs.get(SCHEMAS_CACHE_SIZE_CONFIG);
-        if (cacheSizeVal != null)
-            cacheSize = Integer.parseInt((String) cacheSizeVal);
-        fromConnectSchemaCache = new SynchronizedCache<>(new LRUCache(cacheSize));
-        toConnectSchemaCache = new SynchronizedCache<>(new LRUCache(cacheSize));
+        Map conf = new HashMap<>(configs);
+        conf.put(StringConverterConfig.TYPE_CONFIG, isKey ? ConverterType.KEY.getName() : ConverterType.VALUE.getName());
+        configure(conf);
+    }
+
+    @Override
+    public void close() {
+        // do nothing
+    }
+
+    @Override
+    public byte[] fromConnectHeader(String topic, String headerKey, Schema schema, Object value) {
+        return fromConnectData(topic, schema, value);
+    }
+
+    @Override
+    public SchemaAndValue toConnectHeader(String topic, String headerKey, byte[] value) {
+        return toConnectData(topic, value);
     }
 
     @Override
     public byte[] fromConnectData(String topic, Schema schema, Object value) {
-        JsonNode jsonValue = enableSchemas ? convertToJsonWithEnvelope(schema, value) : convertToJsonWithoutEnvelope(schema, value);
+        if (schema == null && value == null) {
+            return null;
+        }
+
+        JsonNode jsonValue = config.schemasEnabled() ? convertToJsonWithEnvelope(schema, value) : convertToJsonWithoutEnvelope(schema, value);
         try {
             return serializer.serialize(topic, jsonValue);
         } catch (SerializationException e) {
@@ -298,37 +354,36 @@ public byte[] fromConnectData(String topic, Schema schema, Object value) {
     @Override
     public SchemaAndValue toConnectData(String topic, byte[] value) {
         JsonNode jsonValue;
+
+        // This handles a tombstone message
+        if (value == null) {
+            return SchemaAndValue.NULL;
+        }
+
         try {
             jsonValue = deserializer.deserialize(topic, value);
         } catch (SerializationException e) {
             throw new DataException("Converting byte[] to Kafka Connect data failed due to serialization error: ", e);
         }
 
-        if (enableSchemas && (jsonValue == null || !jsonValue.isObject() || jsonValue.size() != 2 || !jsonValue.has("schema") || !jsonValue.has("payload")))
+        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();
-            envelope.set("schema", null);
-            envelope.set("payload", jsonValue);
+        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;
         }
 
-        return jsonToConnect(jsonValue);
-    }
-
-    private SchemaAndValue jsonToConnect(JsonNode jsonValue) {
-        if (jsonValue == null)
-            return SchemaAndValue.NULL;
-
-        if (!jsonValue.isObject() || jsonValue.size() != 2 || !jsonValue.has(JsonSchema.ENVELOPE_SCHEMA_FIELD_NAME) || !jsonValue.has(JsonSchema.ENVELOPE_PAYLOAD_FIELD_NAME))
-            throw new DataException("JSON value converted to Kafka Connect must be in envelope containing schema");
-
         Schema schema = asConnectSchema(jsonValue.get(JsonSchema.ENVELOPE_SCHEMA_FIELD_NAME));
-        return new SchemaAndValue(schema, convertToConnect(schema, jsonValue.get(JsonSchema.ENVELOPE_PAYLOAD_FIELD_NAME)));
+        return new SchemaAndValue(
+                schema,
+                convertToConnect(schema, jsonValue.get(JsonSchema.ENVELOPE_PAYLOAD_FIELD_NAME))
+        );
     }
 
     public ObjectNode asJsonSchema(Schema schema) {
@@ -369,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());
@@ -399,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);
@@ -455,7 +510,7 @@ public Schema asConnectSchema(JsonNode jsonSchema) {
                 break;
             case JsonSchema.ARRAY_TYPE_NAME:
                 JsonNode elemSchema = jsonSchema.get(JsonSchema.ARRAY_ITEMS_FIELD_NAME);
-                if (elemSchema == null)
+                if (elemSchema == null || elemSchema.isNull())
                     throw new DataException("Array schema did not specify the element type");
                 builder = SchemaBuilder.array(asConnectSchema(elemSchema));
                 break;
@@ -545,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 {
@@ -574,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);
@@ -625,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();
@@ -637,15 +691,15 @@ 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;
                 }
                 case STRUCT: {
                     Struct struct = (Struct) value;
-                    if (struct.schema() != schema)
+                    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)));
                     }
@@ -665,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())
@@ -675,6 +729,7 @@ private static Object convertToConnect(Schema schema, JsonNode jsonValue) {
         } else {
             switch (jsonValue.getNodeType()) {
                 case NULL:
+                case MISSING:
                     // Special case. With no schema
                     return null;
                 case BOOLEAN:
@@ -697,7 +752,6 @@ private static Object convertToConnect(Schema schema, JsonNode jsonValue) {
                     break;
 
                 case BINARY:
-                case MISSING:
                 case POJO:
                 default:
                     schemaType = null;
@@ -707,15 +761,15 @@ private static Object convertToConnect(Schema schema, JsonNode jsonValue) {
 
         final JsonToConnectTypeConverter typeConverter = TO_CONNECT_CONVERTERS.get(schemaType);
         if (typeConverter == null)
-            throw new DataException("Unknown schema type: " + String.valueOf(schemaType));
+            throw new DataException("Unknown schema type: " + 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 {
@@ -723,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
new file mode 100644
index 0000000000000..efb497991cbc9
--- /dev/null
+++ b/connect/json/src/main/java/org/apache/kafka/connect/json/JsonConverterConfig.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.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;
+import org.apache.kafka.common.config.ConfigDef.Width;
+import org.apache.kafka.connect.storage.ConverterConfig;
+
+import java.util.Map;
+
+/**
+ * Configuration options for {@link JsonConverter} instances.
+ */
+public class JsonConverterConfig extends ConverterConfig {
+
+    public static final String SCHEMAS_ENABLE_CONFIG = "schemas.enable";
+    public static final boolean SCHEMAS_ENABLE_DEFAULT = true;
+    private static final String SCHEMAS_ENABLE_DOC = "Include schemas within each of the serialized values and keys.";
+    private static final String SCHEMAS_ENABLE_DISPLAY = "Enable Schemas";
+
+    public static final String SCHEMAS_CACHE_SIZE_CONFIG = "schemas.cache.size";
+    public static final int SCHEMAS_CACHE_SIZE_DEFAULT = 1000;
+    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 {
+        String group = "Schemas";
+        int orderInGroup = 0;
+        CONFIG = ConverterConfig.newConfigDef();
+        CONFIG.define(SCHEMAS_ENABLE_CONFIG, Type.BOOLEAN, SCHEMAS_ENABLE_DEFAULT, Importance.HIGH, SCHEMAS_ENABLE_DOC, group,
+                      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));
+    }
+
+    /**
+     * Return whether schemas are enabled.
+     *
+     * @return true if enabled, or false otherwise
+     */
+    public boolean schemasEnabled() {
+        return schemasEnabled;
+    }
+
+    /**
+     * Get the cache size.
+     *
+     * @return the cache size
+     */
+    public int schemaCacheSize() {
+        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 8f2171bc4bce7..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,28 +16,42 @@
  */
 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;
 
-import java.util.Map;
-
 /**
  * JSON deserializer for Jackson's JsonNode tree model. Using the tree model allows it to work with arbitrarily
  * 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));
     }
 
-    @Override
-    public void configure(Map props, boolean isKey) {
+    /**
+     * 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
@@ -54,9 +68,4 @@ public JsonNode deserialize(String topic, byte[] bytes) {
 
         return data;
     }
-
-    @Override
-    public void close() {
-
-    }
 }
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 438daa17e6b1d..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,10 +18,13 @@
 
 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.Map;
+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
@@ -34,11 +37,22 @@ public class JsonSerializer implements Serializer {
      * Default constructor needed by Kafka
      */
     public JsonSerializer() {
-
+        this(Collections.emptySet(), JsonNodeFactory.withExactBigDecimals(true));
     }
 
-    @Override
-    public void configure(Map config, boolean isKey) {
+    /**
+     * 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
@@ -52,9 +66,4 @@ public byte[] serialize(String topic, JsonNode data) {
             throw new SerializationException("Error serializing JSON message", e);
         }
     }
-
-    @Override
-    public void close() {
-    }
-
 }
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 62b52a0237195..a1ac71d4c8d57 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,18 +59,22 @@
 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() {
-        converter.configure(Collections.EMPTY_MAP, false);
+        converter.configure(Collections.emptyMap(), false);
     }
 
     // Schema metadata
@@ -172,10 +177,52 @@ public void structToConnect() {
         assertEquals(new SchemaAndValue(expectedSchema, expected), converted);
     }
 
-    @Test(expected = DataException.class)
+    @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 null should be an error -- we should *always* have the envelope
-        assertEquals(SchemaAndValue.NULL, converter.toConnectData(TOPIC, null));
+        // When schemas are enabled, trying to decode a tombstone should be an empty envelope
+        // the behavior is the same as when the json is "{ "schema": null, "payload": null }"
+        // to keep compatibility with the record
+        SchemaAndValue converted = converter.toConnectData(TOPIC, null);
+        assertEquals(SchemaAndValue.NULL, converted);
+    }
+
+    /**
+     * When schemas are disabled, empty data should be decoded to an empty envelope.
+     * This test verifies the case where `schemas.enable` configuration is set to false, and
+     * {@link JsonConverter} converts empty bytes to {@link SchemaAndValue#NULL}.
+     */
+    @Test
+    public void emptyBytesToConnect() {
+        // This characterizes the messages with empty data when Json schemas is disabled
+        Map props = Collections.singletonMap("schemas.enable", false);
+        converter.configure(props, true);
+        SchemaAndValue converted = converter.toConnectData(TOPIC, "".getBytes());
+        assertEquals(SchemaAndValue.NULL, converted);
+    }
+
+    /**
+     * When schemas are disabled, fields are mapped to Connect maps.
+     */
+    @Test
+    public void schemalessWithEmptyFieldValueToConnect() {
+        // This characterizes the messages with empty data when Json schemas is disabled
+        Map props = Collections.singletonMap("schemas.enable", false);
+        converter.configure(props, true);
+        String input = "{ \"a\": \"\", \"b\": null}";
+        SchemaAndValue converted = converter.toConnectData(TOPIC, input.getBytes());
+        Map expected = new HashMap<>();
+        expected.put("a", "");
+        expected.put("b", null);
+        assertEquals(new SchemaAndValue(null, expected), converted);
     }
 
     @Test
@@ -249,6 +296,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;
@@ -563,6 +641,20 @@ public void structToJson() {
                 converted.get(JsonSchema.ENVELOPE_PAYLOAD_FIELD_NAME));
     }
 
+    @Test
+    public void structSchemaIdentical() {
+        Schema schema = SchemaBuilder.struct().field("field1", Schema.BOOLEAN_SCHEMA)
+                                              .field("field2", Schema.STRING_SCHEMA)
+                                              .field("field3", Schema.STRING_SCHEMA)
+                                              .field("field4", Schema.BOOLEAN_SCHEMA).build();
+        Schema inputSchema = SchemaBuilder.struct().field("field1", Schema.BOOLEAN_SCHEMA)
+                                                   .field("field2", Schema.STRING_SCHEMA)
+                                                   .field("field3", Schema.STRING_SCHEMA)
+                                                   .field("field4", Schema.BOOLEAN_SCHEMA).build();
+        Struct input = new Struct(inputSchema).put("field1", true).put("field2", "string2").put("field3", "string3").put("field4", false);
+        assertStructSchemaEqual(schema, input);
+    }
+
 
     @Test
     public void decimalToJson() throws IOException {
@@ -570,11 +662,42 @@ 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 dateToJson() throws IOException {
+    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);
         calendar.setTimeZone(TimeZone.getTimeZone("UTC"));
         calendar.add(Calendar.DATE, 10000);
@@ -590,7 +713,7 @@ public void dateToJson() throws IOException {
     }
 
     @Test
-    public void timeToJson() throws IOException {
+    public void timeToJson() {
         GregorianCalendar calendar = new GregorianCalendar(1970, Calendar.JANUARY, 1, 0, 0, 0);
         calendar.setTimeZone(TimeZone.getTimeZone("UTC"));
         calendar.add(Calendar.MILLISECOND, 14400000);
@@ -606,7 +729,7 @@ public void timeToJson() throws IOException {
     }
 
     @Test
-    public void timestampToJson() throws IOException {
+    public void timestampToJson() {
         GregorianCalendar calendar = new GregorianCalendar(1970, Calendar.JANUARY, 1, 0, 0, 0);
         calendar.setTimeZone(TimeZone.getTimeZone("UTC"));
         calendar.add(Calendar.MILLISECOND, 2000000000);
@@ -682,6 +805,23 @@ public void nullSchemaAndMapNonStringKeysToJson() {
         );
     }
 
+    @Test
+    public void nullSchemaAndNullValueToJson() {
+        // This characterizes the production of tombstone messages when Json schemas is enabled
+        Map props = Collections.singletonMap("schemas.enable", true);
+        converter.configure(props, true);
+        byte[] converted = converter.fromConnectData(TOPIC, null, null);
+        assertNull(converted);
+    }
+
+    @Test
+    public void nullValueToJson() {
+        // This characterizes the production of tombstone messages when Json schemas is not enabled
+        Map props = Collections.singletonMap("schemas.enable", false);
+        converter.configure(props, true);
+        byte[] converted = converter.fromConnectData(TOPIC, null, null);
+        assertNull(converted);
+    }
 
     @Test(expected = DataException.class)
     public void mismatchSchemaJson() {
@@ -689,8 +829,6 @@ public void mismatchSchemaJson() {
         converter.fromConnectData(TOPIC, Schema.FLOAT64_SCHEMA, true);
     }
 
-
-
     @Test
     public void noSchemaToConnect() {
         Map props = Collections.singletonMap("schemas.enable", false);
@@ -735,7 +873,23 @@ public void testJsonSchemaCacheSizeFromConfigFile() throws URISyntaxException, I
 
         JsonConverter rc = new JsonConverter();
         rc.configure(workerProps, false);
+    }
+
 
+    // Note: the header conversion methods delegates to the data conversion methods, which are tested above.
+    // The following simply verify that the delegation works.
+
+    @Test
+    public void testStringHeaderToJson() {
+        JsonNode converted = parse(converter.fromConnectHeader(TOPIC, "headerName", Schema.STRING_SCHEMA, "test-string"));
+        validateEnvelope(converted);
+        assertEquals(parse("{ \"type\": \"string\", \"optional\": false }"), converted.get(JsonSchema.ENVELOPE_SCHEMA_FIELD_NAME));
+        assertEquals("test-string", converted.get(JsonSchema.ENVELOPE_PAYLOAD_FIELD_NAME).textValue());
+    }
+
+    @Test
+    public void stringHeaderToConnect() {
+        assertEquals(new SchemaAndValue(Schema.STRING_SCHEMA, "foo-bar-baz"), converter.toConnectHeader(TOPIC, "headerName", "{ \"schema\": { \"type\": \"string\" }, \"payload\": \"foo-bar-baz\" }".getBytes()));
     }
 
 
@@ -774,4 +928,9 @@ private void validateEnvelopeNullSchema(JsonNode env) {
         assertTrue(env.get(JsonSchema.ENVELOPE_SCHEMA_FIELD_NAME).isNull());
         assertTrue(env.has(JsonSchema.ENVELOPE_PAYLOAD_FIELD_NAME));
     }
+    
+    private void assertStructSchemaEqual(Schema schema, Struct struct) {
+        converter.fromConnectData(TOPIC, schema, struct);
+        assertEquals(schema, struct.schema());
+    }
 }
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..9292198ef26cb
--- /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..49da62d95c433
--- /dev/null
+++ b/connect/mirror-client/src/main/java/org/apache/kafka/connect/mirror/RemoteClusterUtils.java
@@ -0,0 +1,94 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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;
+
+
+/** 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 {
+
+    // 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..9c7c259ae2998
--- /dev/null
+++ b/connect/mirror/README.md
@@ -0,0 +1,255 @@
+
+# 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
+excluded 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.
+
+## Configuration
+The following sections target for dedicated MM2 cluster. If running MM2 in a Connect cluster, please refer to [KIP-382: MirrorMaker 2.0](https://cwiki.apache.org/confluence/display/KAFKA/KIP-382%3A+MirrorMaker+2.0) for guidance.
+ 
+### General Kafka Connect Config
+All Kafka Connect, Source Connector, Sink Connector configs, as defined in [Kafka official doc](https://kafka.apache.org/documentation/#connectconfigs), can be 
+directly used in MM2 configuration without prefix in the configuration name. As the starting point, most of these default configs may work well with the exception of `tasks.max`.
+
+In order to evenly distribute the workload across more than one MM2 instance, it is advised to set `tasks.max` at least to 2 or even larger depending on the hardware resources
+and the total number partitions to be replicated.
+
+### Kafka Connect Config for a Specific Connector
+If needed, Kafka Connect worker-level configs could be even specified "per connector", which needs to follow the format of `cluster_alias.config_name` in MM2 configuration. For example,
+ 
+    backup.ssl.truststore.location = /usr/lib/jvm/zulu-8-amd64/jre/lib/security/cacerts // SSL cert location
+    backup.security.protocol = SSL // if target cluster needs SSL to send message
+    
+### MM2 Config for a Specific Connector
+MM2 itself has many configs to control how it behaves. To override those default values, add the config name by the format of `source_cluster_alias->target_cluster_alias.config_name` in MM2 configuration. For example,
+    
+    backup->primary.enabled = false // set to false if one-way replication is desired
+    primary->backup.topics.blacklist = topics_to_blacklist
+    primary->backup.emit.heartbeats.enabled = false
+    primary->backup.sync.group.offsets = true 
+
+### Producer / Consumer / Admin Config used by MM2
+In many cases, customized values for producer or consumer configurations are needed. In order to override the default values of producer or consumer used by MM2, 
+`target_cluster_alias.producer.producer_config_name`, `source_cluster_alias.consumer.consumer_config_name` or `cluster_alias.admin.admin_config_name` are the formats to use in MM2 configuration. For example,
+
+     backup.producer.compression.type = gzip
+     backup.producer.buffer.memory = 32768
+     primary.consumer.isolation.level = read_committed
+     primary.admin.bootstrap.servers = localhost:9092
+     
+### 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..0c85f50d000ba
--- /dev/null
+++ b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/DefaultConfigPropertyFilter.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.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.common.utils.ConfigUtils;
+
+import java.util.Map;
+import java.util.regex.Pattern;
+
+/** Filters excluded property names or regexes. */
+public class DefaultConfigPropertyFilter implements ConfigPropertyFilter {
+    
+    public static final String CONFIG_PROPERTIES_EXCLUDE_CONFIG = "config.properties.exclude";
+    public static final String CONFIG_PROPERTIES_EXCLUDE_ALIAS_CONFIG = "config.properties.blacklist";
+
+    private static final String CONFIG_PROPERTIES_EXCLUDE_DOC = "List of topic configuration properties and/or regexes "
+                                                                + "that should not be replicated.";
+    public static final String CONFIG_PROPERTIES_EXCLUDE_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 excludePattern = MirrorUtils.compilePatternList(CONFIG_PROPERTIES_EXCLUDE_DEFAULT);
+
+    @Override
+    public void configure(Map props) {
+        ConfigPropertyFilterConfig config = new ConfigPropertyFilterConfig(props);
+        excludePattern = config.excludePattern();
+    }
+
+    @Override
+    public void close() {
+    }
+
+    private boolean excluded(String prop) {
+        return excludePattern != null && excludePattern.matcher(prop).matches();
+    }
+
+    @Override
+    public boolean shouldReplicateConfigProperty(String prop) {
+        return !excluded(prop);
+    }
+
+    static class ConfigPropertyFilterConfig extends AbstractConfig {
+
+        static final ConfigDef DEF = new ConfigDef()
+            .define(CONFIG_PROPERTIES_EXCLUDE_CONFIG,
+                    Type.LIST,
+                    CONFIG_PROPERTIES_EXCLUDE_DEFAULT,
+                    Importance.HIGH,
+                    CONFIG_PROPERTIES_EXCLUDE_DOC)
+            .define(CONFIG_PROPERTIES_EXCLUDE_ALIAS_CONFIG,
+                    Type.LIST,
+                    null,
+                    Importance.HIGH,
+                    "Deprecated. Use " + CONFIG_PROPERTIES_EXCLUDE_CONFIG + " instead.");
+
+        ConfigPropertyFilterConfig(Map props) {
+            super(DEF, ConfigUtils.translateDeprecatedConfigs(props, new String[][]{
+                {CONFIG_PROPERTIES_EXCLUDE_CONFIG, CONFIG_PROPERTIES_EXCLUDE_ALIAS_CONFIG}}), false);
+        }
+
+        Pattern excludePattern() {
+            return MirrorUtils.compilePatternList(getList(CONFIG_PROPERTIES_EXCLUDE_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..179067eb291a3
--- /dev/null
+++ b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/DefaultGroupFilter.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.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.common.utils.ConfigUtils;
+
+import java.util.Map;
+import java.util.regex.Pattern;
+
+/** Uses an include and exclude pattern. */
+public class DefaultGroupFilter implements GroupFilter {
+
+    public static final String GROUPS_INCLUDE_CONFIG = "groups";
+    private static final String GROUPS_INCLUDE_DOC = "List of consumer group names and/or regexes to replicate.";
+    public static final String GROUPS_INCLUDE_DEFAULT = ".*";
+
+    public static final String GROUPS_EXCLUDE_CONFIG = "groups.exclude";
+    public static final String GROUPS_EXCLUDE_CONFIG_ALIAS = "groups.blacklist";
+
+    private static final String GROUPS_EXCLUDE_DOC = "List of consumer group names and/or regexes that should not be replicated.";
+    public static final String GROUPS_EXCLUDE_DEFAULT = "console-consumer-.*, connect-.*, __.*";
+
+    private Pattern includePattern;
+    private Pattern excludePattern;
+
+    @Override
+    public void configure(Map props) {
+        GroupFilterConfig config = new GroupFilterConfig(props);
+        includePattern = config.includePattern();
+        excludePattern = config.excludePattern();
+    }
+
+    @Override
+    public void close() {
+    }
+
+    private boolean included(String group) {
+        return includePattern != null && includePattern.matcher(group).matches();
+    }
+
+    private boolean excluded(String group) {
+        return excludePattern != null && excludePattern.matcher(group).matches();
+    }
+
+    @Override
+    public boolean shouldReplicateGroup(String group) {
+        return included(group) && !excluded(group);
+    }
+
+    static class GroupFilterConfig extends AbstractConfig {
+
+        static final ConfigDef DEF = new ConfigDef()
+            .define(GROUPS_INCLUDE_CONFIG,
+                    Type.LIST,
+                    GROUPS_INCLUDE_DEFAULT,
+                    Importance.HIGH,
+                    GROUPS_INCLUDE_DOC)
+            .define(GROUPS_EXCLUDE_CONFIG,
+                    Type.LIST,
+                    GROUPS_EXCLUDE_DEFAULT,
+                    Importance.HIGH,
+                    GROUPS_EXCLUDE_DOC)
+            .define(GROUPS_EXCLUDE_CONFIG_ALIAS,
+                    Type.LIST,
+                    null,
+                    Importance.HIGH,
+                    "Deprecated. Use " + GROUPS_EXCLUDE_CONFIG + " instead.");
+
+        GroupFilterConfig(Map props) {
+            super(DEF, ConfigUtils.translateDeprecatedConfigs(props, new String[][]{
+                {GROUPS_EXCLUDE_CONFIG, GROUPS_EXCLUDE_CONFIG_ALIAS}}), false);
+        }
+
+        Pattern includePattern() {
+            return MirrorUtils.compilePatternList(getList(GROUPS_INCLUDE_CONFIG));
+        }
+
+        Pattern excludePattern() {
+            return MirrorUtils.compilePatternList(getList(GROUPS_EXCLUDE_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..f808ce82d4360
--- /dev/null
+++ b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/DefaultTopicFilter.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.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.utils.ConfigUtils;
+
+import java.util.Map;
+import java.util.regex.Pattern;
+
+/** Uses an include and exclude pattern. */
+public class DefaultTopicFilter implements TopicFilter {
+    
+    public static final String TOPICS_INCLUDE_CONFIG = "topics";
+    private static final String TOPICS_INCLUDE_DOC = "List of topics and/or regexes to replicate.";
+    public static final String TOPICS_INCLUDE_DEFAULT = ".*";
+
+    public static final String TOPICS_EXCLUDE_CONFIG = "topics.exclude";
+    public static final String TOPICS_EXCLUDE_CONFIG_ALIAS = "topics.blacklist";
+    private static final String TOPICS_EXCLUDE_DOC = "List of topics and/or regexes that should not be replicated.";
+    public static final String TOPICS_EXCLUDE_DEFAULT = ".*[\\-\\.]internal, .*\\.replica, __.*";
+
+    private Pattern includePattern;
+    private Pattern excludePattern;
+
+    @Override
+    public void configure(Map props) {
+        TopicFilterConfig config = new TopicFilterConfig(props);
+        includePattern = config.includePattern();
+        excludePattern = config.excludePattern();
+    }
+
+    @Override
+    public void close() {
+    }
+
+    private boolean included(String topic) {
+        return includePattern != null && includePattern.matcher(topic).matches();
+    }
+
+    private boolean excluded(String topic) {
+        return excludePattern != null && excludePattern.matcher(topic).matches();
+    }
+
+    @Override
+    public boolean shouldReplicateTopic(String topic) {
+        return included(topic) && !excluded(topic);
+    }
+
+    static class TopicFilterConfig extends AbstractConfig {
+
+        static final ConfigDef DEF = new ConfigDef()
+            .define(TOPICS_INCLUDE_CONFIG,
+                    Type.LIST,
+                    TOPICS_INCLUDE_DEFAULT,
+                    Importance.HIGH,
+                    TOPICS_INCLUDE_DOC)
+            .define(TOPICS_EXCLUDE_CONFIG,
+                    Type.LIST,
+                    TOPICS_EXCLUDE_DEFAULT,
+                    Importance.HIGH,
+                    TOPICS_EXCLUDE_DOC)
+            .define(TOPICS_EXCLUDE_CONFIG_ALIAS,
+                    Type.LIST,
+                    null,
+                    Importance.HIGH,
+                    "Deprecated. Use " + TOPICS_EXCLUDE_CONFIG + " instead.");
+
+        TopicFilterConfig(Map props) {
+            super(DEF, ConfigUtils.translateDeprecatedConfigs(props, new String[][]{
+                {TOPICS_EXCLUDE_CONFIG, TOPICS_EXCLUDE_CONFIG_ALIAS}}), false);
+        }
+
+        Pattern includePattern() {
+            return MirrorUtils.compilePatternList(getList(TOPICS_INCLUDE_CONFIG));
+        }
+
+        Pattern excludePattern() {
+            return MirrorUtils.compilePatternList(getList(TOPICS_EXCLUDE_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..37b6d518282b4
--- /dev/null
+++ b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorCheckpointConnector.java
@@ -0,0 +1,169 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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();
+
+    public MirrorCheckpointConnector() {
+        // nop
+    }
+
+    // visible for testing
+    MirrorCheckpointConnector(List knownConsumerGroups, MirrorConnectorConfig config) {
+        this.knownConsumerGroups = knownConsumerGroups;
+        this.config = config;
+    }
+
+    @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 the replication is disabled, known consumer group is empty, or checkpoint emission is
+        // disabled by setting 'emit.checkpoints.enabled' to false, the interval of checkpoint emission
+        // will be negative and no 'MirrorHeartbeatTask' will be created
+        if (!config.enabled() || knownConsumerGroups.isEmpty()
+                || config.emitCheckpointsInterval().isNegative()) {
+            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();
+    }
+
+    List findConsumerGroups()
+            throws InterruptedException, ExecutionException {
+        return listConsumerGroups().stream()
+                .map(x -> x.groupId())
+                .filter(this::shouldReplicate)
+                .collect(Collectors.toList());
+    }
+
+    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..8e9e764e307b1
--- /dev/null
+++ b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorCheckpointTask.java
@@ -0,0 +1,320 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.connect.mirror;
+
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.clients.admin.ConsumerGroupDescription;
+import org.apache.kafka.common.ConsumerGroupState;
+import org.apache.kafka.common.KafkaFuture;
+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.common.TopicPartition;
+import org.apache.kafka.common.utils.Utils;
+import org.apache.kafka.clients.admin.AdminClient;
+import org.apache.kafka.clients.consumer.OffsetAndMetadata;
+import org.apache.kafka.clients.producer.RecordMetadata;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.HashMap;
+import java.util.Map.Entry;
+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 Admin sourceAdminClient;
+    private Admin targetAdminClient;
+    private String sourceClusterAlias;
+    private String targetClusterAlias;
+    private String checkpointsTopic;
+    private Duration interval;
+    private Duration pollTimeout;
+    private TopicFilter topicFilter;
+    private Set consumerGroups;
+    private ReplicationPolicy replicationPolicy;
+    private OffsetSyncStore offsetSyncStore;
+    private boolean stopping;
+    private MirrorMetrics metrics;
+    private Scheduler scheduler;
+    private Map> idleConsumerGroupsOffset;
+    private Map> checkpointsPerConsumerGroup;
+    public MirrorCheckpointTask() {}
+
+    // for testing
+    MirrorCheckpointTask(String sourceClusterAlias, String targetClusterAlias,
+            ReplicationPolicy replicationPolicy, OffsetSyncStore offsetSyncStore,
+            Map> idleConsumerGroupsOffset,
+            Map> checkpointsPerConsumerGroup) {
+        this.sourceClusterAlias = sourceClusterAlias;
+        this.targetClusterAlias = targetClusterAlias;
+        this.replicationPolicy = replicationPolicy;
+        this.offsetSyncStore = offsetSyncStore;
+        this.idleConsumerGroupsOffset = idleConsumerGroupsOffset;
+        this.checkpointsPerConsumerGroup = checkpointsPerConsumerGroup;
+    }
+
+    @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();
+        offsetSyncStore = new OffsetSyncStore(config);
+        sourceAdminClient = AdminClient.create(config.sourceAdminConfig());
+        targetAdminClient = AdminClient.create(config.targetAdminConfig());
+        metrics = config.metrics();
+        idleConsumerGroupsOffset = new HashMap<>();
+        checkpointsPerConsumerGroup = new HashMap<>();
+        scheduler = new Scheduler(MirrorCheckpointTask.class, config.adminTimeout());
+        scheduler.scheduleRepeating(this::refreshIdleConsumerGroupOffset, config.syncGroupOffsetsInterval(),
+                                    "refreshing idle consumers group offsets at target cluster");
+        scheduler.scheduleRepeatingDelayed(this::syncGroupOffset, config.syncGroupOffsetsInterval(),
+                                          "sync idle consumer group offset from source to target");
+    }
+
+    @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(targetAdminClient, "target admin client");
+        Utils.closeQuietly(metrics, "metrics");
+        Utils.closeQuietly(scheduler, "scheduler");
+        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(sourceRecordsForGroup(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 sourceRecordsForGroup(String group) throws InterruptedException {
+        try {
+            long timestamp = System.currentTimeMillis();
+            List checkpoints = checkpointsForGroup(group);
+            checkpointsPerConsumerGroup.put(group, checkpoints);
+            return checkpoints.stream()
+                .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 List checkpointsForGroup(String group) throws ExecutionException, InterruptedException {
+        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
+            .collect(Collectors.toList());
+    }
+
+    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, RecordMetadata metadata) {
+        metrics.checkpointLatency(MirrorUtils.unwrapPartition(record.sourcePartition()),
+            Checkpoint.unwrapGroup(record.sourcePartition()),
+            System.currentTimeMillis() - record.timestamp());
+    }
+
+    private void refreshIdleConsumerGroupOffset() {
+        Map> consumerGroupsDesc = targetAdminClient
+            .describeConsumerGroups(consumerGroups).describedGroups();
+
+        for (String group : consumerGroups) {
+            try {
+                ConsumerGroupDescription consumerGroupDesc = consumerGroupsDesc.get(group).get();
+                ConsumerGroupState consumerGroupState = consumerGroupDesc.state();
+                // sync offset to the target cluster only if the state of current consumer group is:
+                // (1) idle: because the consumer at target is not actively consuming the mirrored topic
+                // (2) dead: the new consumer that is recently created at source and never exist at target
+                if (consumerGroupState.equals(ConsumerGroupState.EMPTY)) {
+                    idleConsumerGroupsOffset.put(group, targetAdminClient.listConsumerGroupOffsets(group)
+                        .partitionsToOffsetAndMetadata().get().entrySet().stream().collect(
+                            Collectors.toMap(e -> e.getKey(), e -> e.getValue())));
+                }
+                // new consumer upstream has state "DEAD" and will be identified during the offset sync-up
+            } catch (InterruptedException | ExecutionException e) {
+                log.error("Error querying for consumer group {} on cluster {}.", group, targetClusterAlias, e);
+            }
+        }
+    }
+
+    Map> syncGroupOffset() {
+        Map> offsetToSyncAll = new HashMap<>();
+
+        // first, sync offsets for the idle consumers at target
+        for (Entry> group : getConvertedUpstreamOffset().entrySet()) {
+            String consumerGroupId = group.getKey();
+            // for each idle consumer at target, read the checkpoints (converted upstream offset)
+            // from the pre-populated map
+            Map convertedUpstreamOffset = group.getValue();
+
+            Map offsetToSync = new HashMap<>();
+            Map targetConsumerOffset = idleConsumerGroupsOffset.get(consumerGroupId);
+            if (targetConsumerOffset == null) {
+                // this is a new consumer, just sync the offset to target
+                syncGroupOffset(consumerGroupId, convertedUpstreamOffset);
+                offsetToSyncAll.put(consumerGroupId, convertedUpstreamOffset);
+                continue;
+            }
+
+            for (Entry convertedEntry : convertedUpstreamOffset.entrySet()) {
+
+                TopicPartition topicPartition = convertedEntry.getKey();
+                OffsetAndMetadata convertedOffset = convertedUpstreamOffset.get(topicPartition);
+                if (!targetConsumerOffset.containsKey(topicPartition)) {
+                    // if is a new topicPartition from upstream, just sync the offset to target
+                    offsetToSync.put(topicPartition, convertedOffset);
+                    continue;
+                }
+
+                // if translated offset from upstream is smaller than the current consumer offset
+                // in the target, skip updating the offset for that partition
+                long latestDownstreamOffset = targetConsumerOffset.get(topicPartition).offset();
+                if (latestDownstreamOffset >= convertedOffset.offset()) {
+                    log.trace("latestDownstreamOffset {} is larger than or equal to convertedUpstreamOffset {} for "
+                        + "TopicPartition {}", latestDownstreamOffset, convertedOffset.offset(), topicPartition);
+                    continue;
+                }
+                offsetToSync.put(topicPartition, convertedOffset);
+            }
+
+            if (offsetToSync.size() == 0) {
+                log.trace("skip syncing the offset for consumer group: {}", consumerGroupId);
+                continue;
+            }
+            syncGroupOffset(consumerGroupId, offsetToSync);
+
+            offsetToSyncAll.put(consumerGroupId, offsetToSync);
+        }
+        idleConsumerGroupsOffset.clear();
+        return offsetToSyncAll;
+    }
+
+    void syncGroupOffset(String consumerGroupId, Map offsetToSync) {
+        if (targetAdminClient != null) {
+            targetAdminClient.alterConsumerGroupOffsets(consumerGroupId, offsetToSync);
+            log.trace("sync-ed the offset for consumer group: {} with {} number of offset entries",
+                      consumerGroupId, offsetToSync.size());
+        }
+    }
+
+    Map> getConvertedUpstreamOffset() {
+        Map> result = new HashMap<>();
+
+        for (Entry> entry : checkpointsPerConsumerGroup.entrySet()) {
+            String consumerId = entry.getKey();
+            Map convertedUpstreamOffset = new HashMap<>();
+            for (Checkpoint checkpoint : entry.getValue()) {
+                convertedUpstreamOffset.put(checkpoint.topicPartition(), checkpoint.offsetAndMetadata());
+            }
+            result.put(consumerId, convertedUpstreamOffset);
+        }
+        return result;
+    }
+}
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..2aa47b45de30e
--- /dev/null
+++ b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorConnectorConfig.java
@@ -0,0 +1,671 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.KafkaMetricsContext;
+import org.apache.kafka.common.metrics.MetricsReporter;
+import org.apache.kafka.common.metrics.JmxReporter;
+import org.apache.kafka.common.metrics.MetricsContext;
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.common.utils.ConfigUtils;
+import org.apache.kafka.connect.runtime.ConnectorConfig;
+import static org.apache.kafka.clients.consumer.ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG;
+import static org.apache.kafka.clients.consumer.ConsumerConfig.AUTO_OFFSET_RESET_CONFIG;
+
+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";
+    protected static final String SYNC_GROUP_OFFSETS = "sync.group.offsets";
+
+    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_INCLUDE_CONFIG;
+    public static final String TOPICS_DEFAULT = DefaultTopicFilter.TOPICS_INCLUDE_DEFAULT;
+    private static final String TOPICS_DOC = "Topics to replicate. Supports comma-separated topic names and regexes.";
+    public static final String TOPICS_EXCLUDE = DefaultTopicFilter.TOPICS_EXCLUDE_CONFIG;
+    public static final String TOPICS_EXCLUDE_ALIAS = DefaultTopicFilter.TOPICS_EXCLUDE_CONFIG_ALIAS;
+    public static final String TOPICS_EXCLUDE_DEFAULT = DefaultTopicFilter.TOPICS_EXCLUDE_DEFAULT;
+    private static final String TOPICS_EXCLUDE_DOC = "Excluded topics. Supports comma-separated topic names and regexes."
+                                                     + " Excludes take precedence over includes.";
+    public static final String GROUPS = DefaultGroupFilter.GROUPS_INCLUDE_CONFIG;
+    public static final String GROUPS_DEFAULT = DefaultGroupFilter.GROUPS_INCLUDE_DEFAULT;
+    private static final String GROUPS_DOC = "Consumer groups to replicate. Supports comma-separated group IDs and regexes.";
+    public static final String GROUPS_EXCLUDE = DefaultGroupFilter.GROUPS_EXCLUDE_CONFIG;
+    public static final String GROUPS_EXCLUDE_ALIAS = DefaultGroupFilter.GROUPS_EXCLUDE_CONFIG_ALIAS;
+
+    public static final String GROUPS_EXCLUDE_DEFAULT = DefaultGroupFilter.GROUPS_EXCLUDE_DEFAULT;
+    private static final String GROUPS_EXCLUDE_DOC = "Exclude groups. Supports comma-separated group IDs and regexes."
+                                                     + " Excludes take precedence over includes.";
+    public static final String CONFIG_PROPERTIES_EXCLUDE = DefaultConfigPropertyFilter.CONFIG_PROPERTIES_EXCLUDE_CONFIG;
+    public static final String CONFIG_PROPERTIES_EXCLUDE_ALIAS = DefaultConfigPropertyFilter.CONFIG_PROPERTIES_EXCLUDE_ALIAS_CONFIG;
+    public static final String CONFIG_PROPERTIES_EXCLUDE_DEFAULT = DefaultConfigPropertyFilter.CONFIG_PROPERTIES_EXCLUDE_DEFAULT;
+    private static final String CONFIG_PROPERTIES_EXCLUDE_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 SYNC_GROUP_OFFSETS_ENABLED = SYNC_GROUP_OFFSETS + ENABLED_SUFFIX;
+    private static final String SYNC_GROUP_OFFSETS_ENABLED_DOC = "Whether to periodically write the translated offsets to __consumer_offsets topic in target cluster, as long as no active consumers in that group are connected to the target cluster";
+    public static final boolean SYNC_GROUP_OFFSETS_ENABLED_DEFAULT = false;
+    public static final String SYNC_GROUP_OFFSETS_INTERVAL_SECONDS = SYNC_GROUP_OFFSETS + INTERVAL_SECONDS_SUFFIX;
+    private static final String SYNC_GROUP_OFFSETS_INTERVAL_SECONDS_DOC = "Frequency of consumer group offset sync.";
+    public static final long SYNC_GROUP_OFFSETS_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 SOURCE_PREFIX = MirrorMakerConfig.SOURCE_PREFIX;
+    protected static final String TARGET_PREFIX = MirrorMakerConfig.TARGET_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.";
+
+    public MirrorConnectorConfig(Map props) {
+        this(CONNECTOR_CONFIG_DEF, ConfigUtils.translateDeprecatedConfigs(props, new String[][]{
+            {TOPICS_EXCLUDE, TOPICS_EXCLUDE_ALIAS},
+            {GROUPS_EXCLUDE, GROUPS_EXCLUDE_ALIAS},
+            {CONFIG_PROPERTIES_EXCLUDE, CONFIG_PROPERTIES_EXCLUDE_ALIAS}}));
+    }
+
+    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));
+        props.putAll(originalsWithPrefix(SOURCE_PREFIX + 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.putAll(originalsWithPrefix(SOURCE_PREFIX + CONSUMER_CLIENT_PREFIX));
+        props.put(ENABLE_AUTO_COMMIT_CONFIG, "false");
+        props.putIfAbsent(AUTO_OFFSET_RESET_CONFIG, "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_PREFIX + 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_PREFIX + ADMIN_CLIENT_PREFIX));
+        return props;
+    }
+
+    List metricsReporters() {
+        List reporters = getConfiguredInstances(
+                CommonClientConfigs.METRIC_REPORTER_CLASSES_CONFIG, MetricsReporter.class);
+        JmxReporter jmxReporter = new JmxReporter();
+        jmxReporter.configure(this.originals());
+        reporters.add(jmxReporter);
+        MetricsContext metricsContext = new KafkaMetricsContext("kafka.connect.mirror");
+
+        for (MetricsReporter reporter : reporters) {
+            reporter.contextChange(metricsContext);
+        }
+
+        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);
+    }
+
+    Duration syncGroupOffsetsInterval() {
+        if (getBoolean(SYNC_GROUP_OFFSETS_ENABLED)) {
+            return Duration.ofSeconds(getLong(SYNC_GROUP_OFFSETS_INTERVAL_SECONDS));
+        } else {
+            // negative interval to disable
+            return Duration.ofMillis(-1);
+        }
+    }
+
+    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_EXCLUDE,
+                    ConfigDef.Type.LIST,
+                    TOPICS_EXCLUDE_DEFAULT,
+                    ConfigDef.Importance.HIGH,
+                    TOPICS_EXCLUDE_DOC)
+            .define(
+                    TOPICS_EXCLUDE_ALIAS,
+                    ConfigDef.Type.LIST,
+                    null,
+                    ConfigDef.Importance.HIGH,
+                    "Deprecated. Use " + TOPICS_EXCLUDE + " instead.")
+            .define(
+                    GROUPS,
+                    ConfigDef.Type.LIST,
+                    GROUPS_DEFAULT,
+                    ConfigDef.Importance.HIGH,
+                    GROUPS_DOC) 
+            .define(
+                    GROUPS_EXCLUDE,
+                    ConfigDef.Type.LIST,
+                    GROUPS_EXCLUDE_DEFAULT,
+                    ConfigDef.Importance.HIGH,
+                    GROUPS_EXCLUDE_DOC)
+            .define(
+                    GROUPS_EXCLUDE_ALIAS,
+                    ConfigDef.Type.LIST,
+                    null,
+                    ConfigDef.Importance.HIGH,
+                    "Deprecated. Use " + GROUPS_EXCLUDE + " instead.")
+            .define(
+                    CONFIG_PROPERTIES_EXCLUDE,
+                    ConfigDef.Type.LIST,
+                    CONFIG_PROPERTIES_EXCLUDE_DEFAULT,
+                    ConfigDef.Importance.HIGH,
+                    CONFIG_PROPERTIES_EXCLUDE_DOC)
+            .define(
+                    CONFIG_PROPERTIES_EXCLUDE_ALIAS,
+                    ConfigDef.Type.LIST,
+                    null,
+                    ConfigDef.Importance.HIGH,
+                    "Deprecated. Use " + CONFIG_PROPERTIES_EXCLUDE + " instead.")
+            .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(
+                    SYNC_GROUP_OFFSETS_ENABLED,
+                    ConfigDef.Type.BOOLEAN,
+                    SYNC_GROUP_OFFSETS_ENABLED_DEFAULT,
+                    ConfigDef.Importance.LOW,
+                    SYNC_GROUP_OFFSETS_ENABLED_DOC)
+            .define(
+                    SYNC_GROUP_OFFSETS_INTERVAL_SECONDS,
+                    ConfigDef.Type.LONG,
+                    SYNC_GROUP_OFFSETS_INTERVAL_SECONDS_DEFAULT,
+                    ConfigDef.Importance.LOW,
+                    SYNC_GROUP_OFFSETS_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..8b2d064f5aa1b
--- /dev/null
+++ b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorHeartbeatConnector.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 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;
+    
+    public MirrorHeartbeatConnector() {
+        // nop
+    }
+
+    // visible for testing
+    MirrorHeartbeatConnector(MirrorConnectorConfig config) {
+        this.config = config;
+    }
+
+    @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) {
+        // if the heartbeats emission is disabled by setting `emit.heartbeats.enabled` to `false`,
+        // the interval heartbeat emission will be negative and no `MirrorHeartbeatTask` will be created
+        if (config.emitHeartbeatsInterval().isNegative()) {
+            return Collections.emptyList();
+        }
+        // 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..30c497dfe1107
--- /dev/null
+++ b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorHeartbeatTask.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 org.apache.kafka.connect.mirror;
+
+import org.apache.kafka.clients.producer.RecordMetadata;
+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, RecordMetadata metadata) {
+    }
+}
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..fa73f8de3a979
--- /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());
+        Exit.addShutdownHook("mirror-maker-shutdown-hook", 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..8324e321af7c3
--- /dev/null
+++ b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorMakerConfig.java
@@ -0,0 +1,270 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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";
+
+    static final String SOURCE_CLUSTER_PREFIX = "source.cluster.";
+    static final String TARGET_CLUSTER_PREFIX = "target.cluster.";
+    static final String SOURCE_PREFIX = "source.";
+    static final String TARGET_PREFIX = "target.";
+
+    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));
+
+        Map sourceClusterProps = clusterProps(sourceAndTarget.source());
+        // attrs non prefixed with producer|consumer|admin
+        props.putAll(clusterConfigsWithPrefix(SOURCE_CLUSTER_PREFIX, sourceClusterProps));
+        // attrs prefixed with producer|consumer|admin
+        props.putAll(clientConfigsWithPrefix(SOURCE_PREFIX, sourceClusterProps));
+
+        Map targetClusterProps = clusterProps(sourceAndTarget.target());
+        props.putAll(clusterConfigsWithPrefix(TARGET_CLUSTER_PREFIX, targetClusterProps));
+        props.putAll(clientConfigsWithPrefix(TARGET_PREFIX, targetClusterProps));
+
+        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 clusterConfigsWithPrefix(String prefix, Map props) {
+        return props.entrySet().stream()
+                .filter(x -> !x.getKey().matches("(^consumer.*|^producer.*|^admin.*)"))
+                .collect(Collectors.toMap(x -> prefix + x.getKey(), x -> x.getValue()));
+    }
+
+    static Map clientConfigsWithPrefix(String prefix, Map props) {
+        return props.entrySet().stream()
+                .filter(x -> x.getKey().matches("(^consumer.*|^producer.*|^admin.*)"))
+                .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..04a92b95a20e1
--- /dev/null
+++ b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorMetrics.java
@@ -0,0 +1,208 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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;
+
+    MirrorMetrics(MirrorTaskConfig taskConfig) {
+        this.target = taskConfig.targetClusterAlias();
+        this.source = taskConfig.sourceClusterAlias();
+        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..0f6eb46dd5b4d
--- /dev/null
+++ b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorSourceConnector.java
@@ -0,0 +1,460 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.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.ArrayList;
+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 knownSourceTopicPartitions = Collections.emptyList();
+    private List knownTargetTopicPartitions = Collections.emptyList();
+    private ReplicationPolicy replicationPolicy;
+    private int replicationFactor;
+    private AdminClient sourceAdminClient;
+    private AdminClient targetAdminClient;
+
+    public MirrorSourceConnector() {
+        // nop
+    }
+
+    // visible for testing
+    MirrorSourceConnector(List knownSourceTopicPartitions, MirrorConnectorConfig config) {
+        this.knownSourceTopicPartitions = knownSourceTopicPartitions;
+        this.config = config;
+    }
+
+    // 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::computeAndCreateTopicPartitions, "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, knownSourceTopicPartitions.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
+    // since each mirrored topic has different traffic and number of partitions, to balance the load
+    // across all mirrormaker instances (workers), 'roundrobin' helps to evenly assign all
+    // topic-partition to the tasks, then the tasks are further distributed to workers.
+    // For example, 3 tasks to mirror 3 topics with 8, 2 and 2 partitions respectively.
+    // 't1' denotes 'task 1', 't0p5' denotes 'topic 0, partition 5'
+    // t1 -> [t0p0, t0p3, t0p6, t1p1]
+    // t2 -> [t0p1, t0p4, t0p7, t2p0]
+    // t3 -> [t0p2, t0p5, t1p0, t2p1]
+    @Override
+    public List> taskConfigs(int maxTasks) {
+        if (!config.enabled() || knownSourceTopicPartitions.isEmpty()) {
+            return Collections.emptyList();
+        }
+        int numTasks = Math.min(maxTasks, knownSourceTopicPartitions.size());
+        List> roundRobinByTask = new ArrayList<>(numTasks);
+        for (int i = 0; i < numTasks; i++) {
+            roundRobinByTask.add(new ArrayList<>());
+        }
+        int count = 0;
+        for (TopicPartition partition : knownSourceTopicPartitions) {
+            int index = count % numTasks;
+            roundRobinByTask.get(index).add(partition);
+            count++;
+        }
+
+        return roundRobinByTask.stream().map(config::taskConfigForTopicPartitions)
+            .collect(Collectors.toList());
+    }
+
+    @Override
+    public ConfigDef config() {
+        return MirrorConnectorConfig.CONNECTOR_CONFIG_DEF;
+    }
+
+    @Override
+    public String version() {
+        return "1";
+    }
+
+    // visible for testing
+    List findSourceTopicPartitions()
+            throws InterruptedException, ExecutionException {
+        Set topics = listTopics(sourceAdminClient).stream()
+            .filter(this::shouldReplicateTopic)
+            .collect(Collectors.toSet());
+        return describeTopics(sourceAdminClient, topics).stream()
+            .flatMap(MirrorSourceConnector::expandTopicDescription)
+            .collect(Collectors.toList());
+    }
+
+    // visible for testing
+    List findTargetTopicPartitions()
+            throws InterruptedException, ExecutionException {
+        Set topics = listTopics(targetAdminClient).stream()
+            .filter(t -> sourceAndTarget.source().equals(replicationPolicy.topicSource(t)))
+            .filter(t -> !t.equals(config.checkpointsTopic()))
+            .collect(Collectors.toSet());
+        return describeTopics(targetAdminClient, topics).stream()
+                .flatMap(MirrorSourceConnector::expandTopicDescription)
+                .collect(Collectors.toList());
+    }
+
+    // visible for testing
+    void refreshTopicPartitions()
+            throws InterruptedException, ExecutionException {
+
+        List sourceTopicPartitions = findSourceTopicPartitions();
+        List targetTopicPartitions = findTargetTopicPartitions();
+
+        Set sourceTopicPartitionsSet = new HashSet<>(sourceTopicPartitions);
+        Set knownSourceTopicPartitionsSet = new HashSet<>(knownSourceTopicPartitions);
+
+        Set upstreamTargetTopicPartitions = targetTopicPartitions.stream()
+                .map(x -> new TopicPartition(replicationPolicy.upstreamTopic(x.topic()), x.partition()))
+                .collect(Collectors.toSet());
+
+        Set missingInTarget = new HashSet<>(sourceTopicPartitions);
+        missingInTarget.removeAll(upstreamTargetTopicPartitions);
+
+        knownTargetTopicPartitions = targetTopicPartitions;
+
+        // Detect if topic-partitions were added or deleted from the source cluster
+        // or if topic-partitions are missing from the target cluster
+        if (!knownSourceTopicPartitionsSet.equals(sourceTopicPartitionsSet) || !missingInTarget.isEmpty()) {
+
+            Set newTopicPartitions = sourceTopicPartitionsSet;
+            newTopicPartitions.removeAll(knownSourceTopicPartitions);
+
+            Set deletedTopicPartitions = knownSourceTopicPartitionsSet;
+            deletedTopicPartitions.removeAll(sourceTopicPartitions);
+
+            log.info("Found {} new topic-partitions on {}. " +
+                     "Found {} deleted topic-partitions on {}. " +
+                     "Found {} topic-partitions missing on {}.",
+                     newTopicPartitions.size(), sourceAndTarget.source(),
+                     deletedTopicPartitions.size(), sourceAndTarget.source(),
+                     missingInTarget.size(), sourceAndTarget.target());
+
+            log.trace("Found new topic-partitions on {}: {}", sourceAndTarget.source(), newTopicPartitions);
+            log.trace("Found deleted topic-partitions on {}: {}", sourceAndTarget.source(), deletedTopicPartitions);
+            log.trace("Found missing topic-partitions on {}: {}", sourceAndTarget.target(), missingInTarget);
+
+            knownSourceTopicPartitions = sourceTopicPartitions;
+            computeAndCreateTopicPartitions();
+            context.requestTaskReconfiguration();
+        }
+    }
+
+    private void loadTopicPartitions()
+            throws InterruptedException, ExecutionException {
+        knownSourceTopicPartitions = findSourceTopicPartitions();
+        knownTargetTopicPartitions = findTargetTopicPartitions();
+    }
+
+    private void refreshKnownTargetTopics()
+            throws InterruptedException, ExecutionException {
+        knownTargetTopicPartitions = findTargetTopicPartitions();
+    }
+
+    private Set topicsBeingReplicated() {
+        Set knownTargetTopics = toTopics(knownTargetTopicPartitions);
+        return knownSourceTopicPartitions.stream()
+            .map(x -> x.topic())
+            .distinct()
+            .filter(x -> knownTargetTopics.contains(formatRemoteTopic(x)))
+            .collect(Collectors.toSet());
+    }
+
+    private Set toTopics(Collection tps) {
+        return tps.stream()
+                .map(x -> x.topic())
+                .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());
+    }
+
+    // visible for testing
+    void computeAndCreateTopicPartitions()
+            throws InterruptedException, ExecutionException {
+        Map partitionCounts = knownSourceTopicPartitions.stream()
+            .collect(Collectors.groupingBy(x -> x.topic(), Collectors.counting())).entrySet().stream()
+            .collect(Collectors.toMap(x -> formatRemoteTopic(x.getKey()), x -> x.getValue()));
+        Set knownTargetTopics = toTopics(knownTargetTopicPartitions);
+        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())));
+        createTopicPartitions(partitionCounts, newTopics, newPartitions);
+    }
+
+    // visible for testing
+    void createTopicPartitions(Map partitionCounts, List newTopics,
+            Map newPartitions) {
+        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 static Collection describeTopics(AdminClient adminClient, Collection topics)
+            throws InterruptedException, ExecutionException {
+        return adminClient.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..68e6441f18fc7
--- /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;
+    }
+
+    public 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..20f2ca7e2c5cc
--- /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(name + "-" + 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/main/java/org/apache/kafka/connect/mirror/formatters/CheckpointFormatter.java b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/formatters/CheckpointFormatter.java
new file mode 100644
index 0000000000000..33fe695874240
--- /dev/null
+++ b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/formatters/CheckpointFormatter.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.connect.mirror.formatters;
+
+import java.io.PrintStream;
+
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.common.MessageFormatter;
+import org.apache.kafka.connect.mirror.Checkpoint;
+
+public class CheckpointFormatter implements MessageFormatter {
+
+    @Override
+    public void writeTo(ConsumerRecord record, PrintStream output) {
+        output.println(Checkpoint.deserializeRecord(record));
+    }
+}
diff --git a/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/formatters/HeartbeatFormatter.java b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/formatters/HeartbeatFormatter.java
new file mode 100644
index 0000000000000..a193dbe153099
--- /dev/null
+++ b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/formatters/HeartbeatFormatter.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.connect.mirror.formatters;
+
+import java.io.PrintStream;
+
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.common.MessageFormatter;
+import org.apache.kafka.connect.mirror.Heartbeat;
+
+public class HeartbeatFormatter implements MessageFormatter {
+
+    @Override
+    public void writeTo(ConsumerRecord record, PrintStream output) {
+        output.println(Heartbeat.deserializeRecord(record));
+    }
+}
diff --git a/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/formatters/OffsetSyncFormatter.java b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/formatters/OffsetSyncFormatter.java
new file mode 100644
index 0000000000000..dacae600ae53d
--- /dev/null
+++ b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/formatters/OffsetSyncFormatter.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.connect.mirror.formatters;
+
+import java.io.PrintStream;
+
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.common.MessageFormatter;
+import org.apache.kafka.connect.mirror.OffsetSync;
+
+public class OffsetSyncFormatter implements MessageFormatter {
+
+    @Override
+    public void writeTo(ConsumerRecord record, PrintStream output) {
+        output.println(OffsetSync.deserializeRecord(record));
+    }
+}
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/MirrorCheckpointConnectorTest.java b/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/MirrorCheckpointConnectorTest.java
new file mode 100644
index 0000000000000..76af66ddcaea9
--- /dev/null
+++ b/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/MirrorCheckpointConnectorTest.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.clients.admin.ConsumerGroupListing;
+import org.junit.Test;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+import static org.apache.kafka.connect.mirror.TestUtils.makeProps;
+import static org.junit.Assert.assertEquals;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.spy;
+
+
+public class MirrorCheckpointConnectorTest {
+
+    @Test
+    public void testMirrorCheckpointConnectorDisabled() {
+        // disable the checkpoint emission
+        MirrorConnectorConfig config = new MirrorConnectorConfig(
+            makeProps("emit.checkpoints.enabled", "false"));
+
+        List knownConsumerGroups = new ArrayList<>();
+        knownConsumerGroups.add("consumer-group-1");
+        // MirrorCheckpointConnector as minimum to run taskConfig()
+        MirrorCheckpointConnector connector = new MirrorCheckpointConnector(knownConsumerGroups,
+                                                                            config);
+        List> output = connector.taskConfigs(1);
+        // expect no task will be created
+        assertEquals(0, output.size());
+    }
+
+    @Test
+    public void testNoConsumerGroup() {
+        MirrorConnectorConfig config = new MirrorConnectorConfig(makeProps());
+        MirrorCheckpointConnector connector = new MirrorCheckpointConnector(new ArrayList<>(), config);
+        List> output = connector.taskConfigs(1);
+        // expect no task will be created
+        assertEquals(0, output.size());
+    }
+
+    @Test
+    public void testReplicationDisabled() {
+        // disable the replication
+        MirrorConnectorConfig config = new MirrorConnectorConfig(makeProps("enabled", "false"));
+
+        List knownConsumerGroups = new ArrayList<>();
+        knownConsumerGroups.add("consumer-group-1");
+        // MirrorCheckpointConnector as minimum to run taskConfig()
+        MirrorCheckpointConnector connector = new MirrorCheckpointConnector(knownConsumerGroups, config);
+        List> output = connector.taskConfigs(1);
+        // expect no task will be created
+        assertEquals(0, output.size());
+    }
+
+    @Test
+    public void testFindConsumerGroups() throws Exception {
+        MirrorConnectorConfig config = new MirrorConnectorConfig(makeProps());
+        MirrorCheckpointConnector connector = new MirrorCheckpointConnector(Collections.emptyList(), config);
+        connector = spy(connector);
+
+        Collection groups = Arrays.asList(
+                new ConsumerGroupListing("g1", true),
+                new ConsumerGroupListing("g2", false));
+        doReturn(groups).when(connector).listConsumerGroups();
+        doReturn(true).when(connector).shouldReplicate(anyString());
+        List groupFound = connector.findConsumerGroups();
+
+        Set expectedGroups = groups.stream().map(g -> g.groupId()).collect(Collectors.toSet());
+        assertEquals(expectedGroups, new HashSet<>(groupFound));
+    }
+
+}
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..2567eb454d304
--- /dev/null
+++ b/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/MirrorCheckpointTaskTest.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.connect.mirror;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Collections;
+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, Collections.emptyMap(), Collections.emptyMap());
+        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, Collections.emptyMap(), Collections.emptyMap());
+        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());
+    }
+
+    @Test
+    public void testSyncOffset() {
+        Map> idleConsumerGroupsOffset = new HashMap<>();
+        Map> checkpointsPerConsumerGroup = new HashMap<>();
+
+        String consumer1 = "consumer1";
+        String consumer2 = "consumer2";
+
+        String topic1 = "topic1";
+        String topic2 = "topic2";
+
+        // 'c1t1' denotes consumer offsets of all partitions of topic1 for consumer1
+        Map c1t1 = new HashMap<>();
+        // 't1p0' denotes topic1, partition 0
+        TopicPartition t1p0 = new TopicPartition(topic1, 0);
+
+        c1t1.put(t1p0, new OffsetAndMetadata(100));
+
+        Map c2t2 = new HashMap<>();
+        TopicPartition t2p0 = new TopicPartition(topic2, 0);
+
+        c2t2.put(t2p0, new OffsetAndMetadata(50));
+
+        idleConsumerGroupsOffset.put(consumer1, c1t1);
+        idleConsumerGroupsOffset.put(consumer2, c2t2);
+
+        // 'cpC1T1P0' denotes 'checkpoint' of topic1, partition 0 for consumer1
+        Checkpoint cpC1T1P0 = new Checkpoint(consumer1, new TopicPartition(topic1, 0), 200, 101, "metadata");
+
+        // 'cpC2T2p0' denotes 'checkpoint' of topic2, partition 0 for consumer2
+        Checkpoint cpC2T2P0 = new Checkpoint(consumer2, new TopicPartition(topic2, 0), 100, 51, "metadata");
+
+        // 'checkpointListC1' denotes 'checkpoint' list for consumer1
+        List checkpointListC1 = new ArrayList<>();
+        checkpointListC1.add(cpC1T1P0);
+
+        // 'checkpointListC2' denotes 'checkpoint' list for consumer2
+        List checkpointListC2 = new ArrayList<>();
+        checkpointListC2.add(cpC2T2P0);
+
+        checkpointsPerConsumerGroup.put(consumer1, checkpointListC1);
+        checkpointsPerConsumerGroup.put(consumer2, checkpointListC2);
+
+        MirrorCheckpointTask mirrorCheckpointTask = new MirrorCheckpointTask("source1", "target2",
+            new DefaultReplicationPolicy(), null, idleConsumerGroupsOffset, checkpointsPerConsumerGroup);
+
+        Map> output = mirrorCheckpointTask.syncGroupOffset();
+
+        assertEquals(101, output.get(consumer1).get(t1p0).offset());
+        assertEquals(51, output.get(consumer2).get(t2p0).offset());
+    }
+}
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..69c6c2926e887
--- /dev/null
+++ b/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/MirrorConnectorConfigTest.java
@@ -0,0 +1,249 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.apache.kafka.connect.mirror.TestUtils.makeProps;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertEquals;
+
+public class MirrorConnectorConfigTest {
+
+    @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.exclude", "prop2"));
+        assertTrue(config.configPropertyFilter().shouldReplicateConfigProperty("prop1"));
+        assertFalse(config.configPropertyFilter().shouldReplicateConfigProperty("prop2"));
+    }
+
+    @Test
+    public void testConfigBackwardsCompatibility() {
+        MirrorConnectorConfig config = new MirrorConnectorConfig(
+            makeProps("config.properties.blacklist", "prop1",
+                      "groups.blacklist", "group-1",
+                      "topics.blacklist", "topic-1"));
+        assertFalse(config.configPropertyFilter().shouldReplicateConfigProperty("prop1"));
+        assertTrue(config.configPropertyFilter().shouldReplicateConfigProperty("prop2"));
+        assertFalse(config.topicFilter().shouldReplicateTopic("topic-1"));
+        assertTrue(config.topicFilter().shouldReplicateTopic("topic-2"));
+        assertFalse(config.groupFilter().shouldReplicateGroup("group-1"));
+        assertTrue(config.groupFilter().shouldReplicateGroup("group-2"));
+    }
+
+    @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)
+        ));
+    }
+
+    @Test
+    public void testSourceConsumerConfig() {
+        Map connectorProps = makeProps(
+                MirrorConnectorConfig.CONSUMER_CLIENT_PREFIX + "max.poll.interval.ms", "120000"
+        );
+        MirrorConnectorConfig config = new MirrorConnectorConfig(connectorProps);
+        Map connectorConsumerProps = config.sourceConsumerConfig();
+        Map expectedConsumerProps = new HashMap<>();
+        expectedConsumerProps.put("enable.auto.commit", "false");
+        expectedConsumerProps.put("auto.offset.reset", "earliest");
+        expectedConsumerProps.put("max.poll.interval.ms", "120000");
+        assertEquals(expectedConsumerProps, connectorConsumerProps);
+
+        // checking auto.offset.reset override works
+        connectorProps = makeProps(
+                MirrorConnectorConfig.CONSUMER_CLIENT_PREFIX + "auto.offset.reset", "latest"
+        );
+        config = new MirrorConnectorConfig(connectorProps);
+        connectorConsumerProps = config.sourceConsumerConfig();
+        expectedConsumerProps.put("auto.offset.reset", "latest");
+        expectedConsumerProps.remove("max.poll.interval.ms");
+        assertEquals(expectedConsumerProps, connectorConsumerProps);
+    }
+
+    @Test
+    public void testSourceConsumerConfigWithSourcePrefix() {
+        String prefix = MirrorConnectorConfig.SOURCE_PREFIX + MirrorConnectorConfig.CONSUMER_CLIENT_PREFIX;
+        Map connectorProps = makeProps(
+                prefix + "auto.offset.reset", "latest",
+                prefix + "max.poll.interval.ms", "100"
+        );
+        MirrorConnectorConfig config = new MirrorConnectorConfig(connectorProps);
+        Map connectorConsumerProps = config.sourceConsumerConfig();
+        Map expectedConsumerProps = new HashMap<>();
+        expectedConsumerProps.put("enable.auto.commit", "false");
+        expectedConsumerProps.put("auto.offset.reset", "latest");
+        expectedConsumerProps.put("max.poll.interval.ms", "100");
+        assertEquals(expectedConsumerProps, connectorConsumerProps);
+    }
+
+    @Test
+    public void testSourceProducerConfig() {
+        Map connectorProps = makeProps(
+                MirrorConnectorConfig.PRODUCER_CLIENT_PREFIX + "acks", "1"
+        );
+        MirrorConnectorConfig config = new MirrorConnectorConfig(connectorProps);
+        Map connectorProducerProps = config.sourceProducerConfig();
+        Map expectedProducerProps = new HashMap<>();
+        expectedProducerProps.put("acks", "1");
+        assertEquals(expectedProducerProps, connectorProducerProps);
+    }
+
+    @Test
+    public void testSourceProducerConfigWithSourcePrefix() {
+        String prefix = MirrorConnectorConfig.SOURCE_PREFIX + MirrorConnectorConfig.PRODUCER_CLIENT_PREFIX;
+        Map connectorProps = makeProps(prefix + "acks", "1");
+        MirrorConnectorConfig config = new MirrorConnectorConfig(connectorProps);
+        Map connectorProducerProps = config.sourceProducerConfig();
+        Map expectedProducerProps = new HashMap<>();
+        expectedProducerProps.put("acks", "1");
+        assertEquals(expectedProducerProps, connectorProducerProps);
+    }
+
+    @Test
+    public void testSourceAdminConfig() {
+        Map connectorProps = makeProps(
+                MirrorConnectorConfig.ADMIN_CLIENT_PREFIX +
+                        "connections.max.idle.ms", "10000"
+        );
+        MirrorConnectorConfig config = new MirrorConnectorConfig(connectorProps);
+        Map connectorAdminProps = config.sourceAdminConfig();
+        Map expectedAdminProps = new HashMap<>();
+        expectedAdminProps.put("connections.max.idle.ms", "10000");
+        assertEquals(expectedAdminProps, connectorAdminProps);
+    }
+
+    @Test
+    public void testSourceAdminConfigWithSourcePrefix() {
+        String prefix = MirrorConnectorConfig.SOURCE_PREFIX + MirrorConnectorConfig.ADMIN_CLIENT_PREFIX;
+        Map connectorProps = makeProps(prefix + "connections.max.idle.ms", "10000");
+        MirrorConnectorConfig config = new MirrorConnectorConfig(connectorProps);
+        Map connectorAdminProps = config.sourceAdminConfig();
+        Map expectedAdminProps = new HashMap<>();
+        expectedAdminProps.put("connections.max.idle.ms", "10000");
+        assertEquals(expectedAdminProps, connectorAdminProps);
+    }
+
+    @Test
+    public void testTargetAdminConfig() {
+        Map connectorProps = makeProps(
+                MirrorConnectorConfig.ADMIN_CLIENT_PREFIX +
+                        "connections.max.idle.ms", "10000"
+        );
+        MirrorConnectorConfig config = new MirrorConnectorConfig(connectorProps);
+        Map connectorAdminProps = config.targetAdminConfig();
+        Map expectedAdminProps = new HashMap<>();
+        expectedAdminProps.put("connections.max.idle.ms", "10000");
+        assertEquals(expectedAdminProps, connectorAdminProps);
+    }
+
+    @Test
+    public void testTargetAdminConfigWithSourcePrefix() {
+        String prefix = MirrorConnectorConfig.TARGET_PREFIX + MirrorConnectorConfig.ADMIN_CLIENT_PREFIX;
+        Map connectorProps = makeProps(prefix + "connections.max.idle.ms", "10000");
+        MirrorConnectorConfig config = new MirrorConnectorConfig(connectorProps);
+        Map connectorAdminProps = config.targetAdminConfig();
+        Map expectedAdminProps = new HashMap<>();
+        expectedAdminProps.put("connections.max.idle.ms", "10000");
+        assertEquals(expectedAdminProps, connectorAdminProps);
+    }
+
+}
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..d0715dc8452a9
--- /dev/null
+++ b/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/MirrorConnectorsIntegrationTest.java
@@ -0,0 +1,595 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+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.connect.util.clusters.UngracefulShutdownException;
+import org.apache.kafka.test.IntegrationTest;
+import org.apache.kafka.common.utils.Exit;
+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.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.Properties;
+import java.util.Set;
+import java.util.concurrent.TimeoutException;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import static org.apache.kafka.test.TestUtils.waitForCondition;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+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_PER_PARTITION = 10;
+    private static final int NUM_PARTITIONS = 10;
+    private static final int NUM_RECORDS_PRODUCED = NUM_PARTITIONS * NUM_RECORDS_PER_PARTITION;
+    private static final int RECORD_TRANSFER_DURATION_MS = 30_000;
+    private static final int CHECKPOINT_DURATION_MS = 20_000;
+    private static final int RECORD_CONSUME_DURATION_MS = 20_000;
+    private static final int OFFSET_SYNC_DURATION_MS = 30_000;
+
+    private volatile boolean shuttingDown;
+    private Map mm2Props;
+    private MirrorMakerConfig mm2Config;
+    private EmbeddedConnectCluster primary;
+    private EmbeddedConnectCluster backup;
+
+    private Exit.Procedure exitProcedure;
+    private Exit.Procedure haltProcedure;
+
+    @Before
+    public void setup() throws InterruptedException {
+        shuttingDown = false;
+        exitProcedure = (code, message) -> {
+            if (shuttingDown) {
+                // ignore this since we're shutting down Connect and Kafka and timing isn't always great
+                return;
+            }
+            if (code != 0) {
+                String exitMessage = "Abrupt service exit with code " + code + " and message " + message;
+                log.warn(exitMessage);
+                throw new UngracefulShutdownException(exitMessage);
+            }
+        };
+        haltProcedure = (code, message) -> {
+            if (shuttingDown) {
+                // ignore this since we're shutting down Connect and Kafka and timing isn't always great
+                return;
+            }
+            if (code != 0) {
+                String haltMessage = "Abrupt service halt with code " + code + " and message " + message;
+                log.warn(haltMessage);
+                throw new UngracefulShutdownException(haltMessage);
+            }
+        };
+        // Override the exit and halt procedure that Connect and Kafka will use. For these integration tests,
+        // we don't want to exit the JVM and instead simply want to fail the test
+        Exit.setExitProcedure(exitProcedure);
+        Exit.setHaltProcedure(haltProcedure);
+
+        Properties brokerProps = new Properties();
+        brokerProps.put("auto.create.topics.enable", "false");
+
+        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)
+                .maskExitProcedures(false)
+                .build();
+
+        backup = new EmbeddedConnectCluster.Builder()
+                .name("backup-connect-cluster")
+                .numWorkers(3)
+                .numBrokers(1)
+                .brokerProps(brokerProps)
+                .workerProps(backupWorkerProps)
+                .maskExitProcedures(false)
+                .build();
+
+        primary.start();
+        primary.assertions().assertAtLeastNumWorkersAreUp(3,
+                "Workers of primary-connect-cluster did not start in time.");
+        backup.start();
+        backup.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);
+
+        // produce to all partitions of test-topic-1
+        produceMessages(primary, "test-topic-1", "message-1-");
+        produceMessages(backup, "test-topic-1", "message-2-");
+
+        // Generate some consumer activity on both clusters to ensure the checkpoint connector always starts promptly
+        Map dummyProps = Collections.singletonMap("group.id", "consumer-group-dummy");
+        Consumer dummyConsumer = primary.kafka().createConsumerAndSubscribeTo(dummyProps, "test-topic-1");
+        consumeAllMessages(dummyConsumer);
+        dummyConsumer.close();
+        dummyConsumer = backup.kafka().createConsumerAndSubscribeTo(dummyProps, "test-topic-1");
+        consumeAllMessages(dummyConsumer);
+        dummyConsumer.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);
+    }
+
+
+    private void waitUntilMirrorMakerIsRunning(EmbeddedConnectCluster connectCluster,
+        MirrorMakerConfig mm2Config, String primary, String backup) throws InterruptedException {
+
+        connectCluster.configureConnector("MirrorSourceConnector",
+                mm2Config.connectorBaseConfig(new SourceAndTarget(primary, backup), MirrorSourceConnector.class));
+        connectCluster.configureConnector("MirrorCheckpointConnector",
+                mm2Config.connectorBaseConfig(new SourceAndTarget(primary, backup), MirrorCheckpointConnector.class));
+        connectCluster.configureConnector("MirrorHeartbeatConnector",
+                mm2Config.connectorBaseConfig(new SourceAndTarget(primary, backup), MirrorHeartbeatConnector.class));
+
+        // 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"));
+
+        for (String connector : connectorNames) {
+            connectCluster.assertions().assertConnectorAndAtLeastNumTasksAreRunning(connector, 1,
+                    "Connector " + connector + " tasks did not start in time on cluster: " + connectCluster);
+        }
+    }
+
+    @After
+    public void close() {
+        try {
+            for (String x : primary.connectors()) {
+                primary.deleteConnector(x);
+            }
+            for (String x : backup.connectors()) {
+                backup.deleteConnector(x);
+            }
+            deleteAllTopics(primary.kafka());
+            deleteAllTopics(backup.kafka());
+        } finally {
+            shuttingDown = true;
+            try {
+                try {
+                    primary.stop();
+                } finally {
+                    backup.stop();
+                }
+            } finally {
+                Exit.resetExitProcedure();
+                Exit.resetHaltProcedure();
+            }
+        }
+    }
+
+    @Test
+    public void testReplication() throws InterruptedException {
+        String consumerGroupName = "consumer-group-testReplication";
+        Map consumerProps = new HashMap() {{
+                put("group.id", consumerGroupName);
+                put("auto.offset.reset", "latest");
+            }};
+
+        // create consumers before starting the connectors so we don't need to wait for discovery
+        Consumer primaryConsumer = primary.kafka().createConsumerAndSubscribeTo(consumerProps, "test-topic-1");
+        consumeAllMessages(primaryConsumer, 0);
+        primaryConsumer.close();
+
+        Consumer backupConsumer = backup.kafka().createConsumerAndSubscribeTo(consumerProps, "test-topic-1");
+        consumeAllMessages(backupConsumer, 0);
+        backupConsumer.close();
+
+        waitUntilMirrorMakerIsRunning(backup, mm2Config, "primary", "backup");
+        waitUntilMirrorMakerIsRunning(primary, mm2Config, "backup", "primary");
+        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(consumerGroupName, "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.
+        backupConsumer = backup.kafka().createConsumer(consumerProps);
+        backupConsumer.assign(allPartitions("test-topic-1", "primary.test-topic-1"));
+        seek(backupConsumer, backupOffsets);
+        consumeAllMessages(backupConsumer, 0);
+
+        assertTrue("Consumer failedover to zero offset.", backupConsumer.position(new TopicPartition("primary.test-topic-1", 0)) > 0);
+        assertTrue("Consumer failedover beyond expected offset.", backupConsumer.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);
+
+        backupConsumer.close();
+
+        waitForCondition(() -> {
+            try {
+                return primaryClient.remoteConsumerOffsets(consumerGroupName, "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(consumerGroupName, "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(consumerGroupName, "backup",
+                Duration.ofMillis(CHECKPOINT_DURATION_MS));
+
+        primaryClient.close();
+        backupClient.close();
+
+        // Failback consumer group to primary cluster
+        primaryConsumer = primary.kafka().createConsumer(consumerProps);
+        primaryConsumer.assign(allPartitions("test-topic-1", "backup.test-topic-1"));
+        seek(primaryConsumer, primaryOffsets);
+        consumeAllMessages(primaryConsumer, 0);
+
+        assertTrue("Consumer failedback to zero upstream offset.", primaryConsumer.position(new TopicPartition("test-topic-1", 0)) > 0);
+        assertTrue("Consumer failedback to zero downstream offset.", primaryConsumer.position(new TopicPartition("backup.test-topic-1", 0)) > 0);
+        assertTrue("Consumer failedback beyond expected upstream offset.", primaryConsumer.position(
+            new TopicPartition("test-topic-1", 0)) <= NUM_RECORDS_PER_PARTITION);
+        assertTrue("Consumer failedback beyond expected downstream offset.", primaryConsumer.position(
+            new TopicPartition("backup.test-topic-1", 0)) <= NUM_RECORDS_PER_PARTITION);
+
+        Map>> messages2 = consumeAllMessages(primaryConsumer, 0);
+        // If offset translation was successful we expect no messages to be consumed after failback
+        assertEquals("Data was consumed from partitions: " + messages2.keySet() + ".", 0, messages2.size());
+        primaryConsumer.close();
+
+        // create more matching topics
+        primary.kafka().createTopic("test-topic-2", NUM_PARTITIONS);
+        backup.kafka().createTopic("test-topic-3", NUM_PARTITIONS);
+
+        produceMessages(primary, "test-topic-2", "message-3-", 1);
+        produceMessages(backup, "test-topic-3", "message-4-", 1);
+
+        assertEquals("Records were not produced to primary cluster.", NUM_RECORDS_PER_PARTITION,
+            primary.kafka().consume(NUM_RECORDS_PER_PARTITION, RECORD_TRANSFER_DURATION_MS, "test-topic-2").count());
+        assertEquals("Records were not produced to backup cluster.", NUM_RECORDS_PER_PARTITION,
+            backup.kafka().consume(NUM_RECORDS_PER_PARTITION, RECORD_TRANSFER_DURATION_MS, "test-topic-3").count());
+
+        assertEquals("New topic was not replicated to primary cluster.", NUM_RECORDS_PER_PARTITION,
+            primary.kafka().consume(NUM_RECORDS_PER_PARTITION, 2 * RECORD_TRANSFER_DURATION_MS, "backup.test-topic-3").count());
+        assertEquals("New topic was not replicated to backup cluster.", NUM_RECORDS_PER_PARTITION,
+            backup.kafka().consume(NUM_RECORDS_PER_PARTITION, 2 * RECORD_TRANSFER_DURATION_MS, "primary.test-topic-2").count());
+    }
+
+    @Test
+    public void testReplicationWithEmptyPartition() throws Exception {
+        String consumerGroupName = "consumer-group-testReplicationWithEmptyPartition";
+        Map consumerProps  = Collections.singletonMap("group.id", consumerGroupName);
+
+        // create topics
+        String topic = "test-topic-with-empty-partition";
+        primary.kafka().createTopic(topic, NUM_PARTITIONS);
+
+        // produce to all test-topic-empty's partitions *but the last one*, on the primary cluster
+        produceMessages(primary, topic, "message-1-", NUM_PARTITIONS - 1);
+        // Consume, from the primary cluster, before starting the connectors so we don't need to wait for discovery
+        int expectedRecords = NUM_RECORDS_PER_PARTITION * (NUM_PARTITIONS - 1);
+        try (Consumer consumer = primary.kafka().createConsumerAndSubscribeTo(consumerProps, topic)) {
+            consumeAllMessages(consumer, expectedRecords);
+        }
+
+        waitUntilMirrorMakerIsRunning(backup, mm2Config, "primary", "backup");
+
+        Map offsets = waitForConsumerGroupOffsetReplication(
+                Collections.singletonList("primary." + topic), consumerGroupName, false);
+
+        // check translated offset for the last partition (empty partition)
+        OffsetAndMetadata oam = offsets.get(new TopicPartition("primary." + topic, NUM_PARTITIONS - 1));
+        assertNotNull("Offset of last partition was not replicated", oam);
+        assertEquals("Offset of last partition is not zero", 0, oam.offset());
+    }
+
+    private Map waitForConsumerGroupOffsetReplication(List topics, String consumerGroupId, boolean waitForAutoSync)
+            throws InterruptedException {
+        Admin backupClient = backup.kafka().createAdminClient();
+        List tps = new ArrayList<>(NUM_PARTITIONS * topics.size());
+        for (int partitionIndex = 0; partitionIndex < NUM_PARTITIONS; partitionIndex++) {
+            for (String topic : topics) {
+                tps.add(new TopicPartition(topic, partitionIndex));
+            }
+        }
+        long expectedTotalOffsets = NUM_PARTITIONS * topics.size();
+        Map offsets = new HashMap<>();
+        waitForCondition(() -> {
+            offsets.clear();
+            Map translatedOffsets = getTranslatedOffsets(consumerGroupId);
+            Map remoteOffsets =
+                    backupClient.listConsumerGroupOffsets(consumerGroupId).partitionsToOffsetAndMetadata().get();
+            boolean isConsumerUpdated = true;
+            for (TopicPartition tp : tps) {
+                OffsetAndMetadata translatedOam = translatedOffsets.get(tp);
+                OffsetAndMetadata remoteOam = remoteOffsets.get(tp);
+                if (translatedOam != null) {
+                    offsets.put(tp, translatedOam);
+                    if (waitForAutoSync && remoteOam.offset() != translatedOam.offset())
+                        isConsumerUpdated = false;
+                }
+            }
+            return offsets.size() == expectedTotalOffsets && isConsumerUpdated;
+        }, OFFSET_SYNC_DURATION_MS, "Consumer group offset sync did not complete in time");
+        return offsets;
+    }
+
+    private Map getTranslatedOffsets(String consumerGroupId)
+            throws TimeoutException, InterruptedException {
+        return RemoteClusterUtils.translateOffsets(
+                mm2Config.clientConfig("backup").adminConfig(),
+                "primary",
+                consumerGroupId,
+                Duration.ofMillis(CHECKPOINT_DURATION_MS));
+    }
+
+    @Test
+    public void testOneWayReplicationWithAutoOffsetSync() throws InterruptedException {
+        String topic1 = "test-topic-auto-offset-sync-1";
+        String topic2 = "test-topic-auto-offset-sync-2";
+        String consumerGroupName = "consumer-group-testOneWayReplicationWithAutoOffsetSync";
+        Map consumerProps  = new HashMap() {{
+                put("group.id", consumerGroupName);
+                put("auto.offset.reset", "earliest");
+            }};
+
+        // create new topic
+        primary.kafka().createTopic(topic1, NUM_PARTITIONS);
+        backup.kafka().createTopic("primary." + topic1, 1);
+
+        // produce some records to the new topic in primary cluster
+        produceMessages(primary, topic1, "message-1-");
+        // create consumers before starting the connectors so we don't need to wait for discovery
+        try (Consumer primaryConsumer = primary.kafka().createConsumerAndSubscribeTo(consumerProps, topic1)) {
+            consumeAllMessages(primaryConsumer);
+        }
+
+        // enable automated consumer group offset sync
+        mm2Props.put("sync.group.offsets.enabled", "true");
+        mm2Props.put("sync.group.offsets.interval.seconds", "1");
+        // one way replication from primary to backup
+        mm2Props.put("backup->primary.enabled", "false");
+
+        mm2Config = new MirrorMakerConfig(mm2Props);
+
+        waitUntilMirrorMakerIsRunning(backup, mm2Config, "primary", "backup");
+
+        waitForConsumerGroupOffsetReplication(
+                Collections.singletonList("primary." + topic1), consumerGroupName, true);
+
+        // create a consumer at backup cluster with same consumer group Id to consume 1 topic
+        ConsumerRecords records = null;
+        try (Consumer backupConsumer = backup.kafka().createConsumerAndSubscribeTo(
+                consumerProps, "primary." + topic1)) {
+            records = backupConsumer.poll(Duration.ofMillis(500));
+        }
+
+        // the size of consumer record should be zero, because the offsets of the same consumer group
+        // have been automatically synchronized from primary to backup by the background job, so no
+        // more records to consume from the replicated topic by the same consumer group at backup cluster
+        assertEquals("consumer record size is not zero", 0, records.count());
+
+        // create a second topic
+        primary.kafka().createTopic(topic2, NUM_PARTITIONS);
+        backup.kafka().createTopic("primary." + topic2, 1);
+
+        // produce some records to the new topic in primary cluster
+        produceMessages(primary, topic2, "message-1-");
+
+        // create a consumer at primary cluster to consume the new topic
+        try (Consumer primaryConsumer = primary.kafka().createConsumerAndSubscribeTo(
+            consumerProps, topic2)) {
+            // we need to wait for consuming all the records for MM2 replicating the expected offsets
+            consumeAllMessages(primaryConsumer);
+        }
+
+        waitForConsumerGroupOffsetReplication(Arrays.asList("primary." + topic1, "primary." + topic2), consumerGroupName, true);
+
+        // create a consumer at backup cluster with same consumer group Id to consume old and new topic
+        try (Consumer backupConsumer = backup.kafka().createConsumerAndSubscribeTo(
+                consumerProps, "primary." + topic1, "primary." + topic2)) {
+            records = backupConsumer.poll(Duration.ofMillis(500));
+        }
+
+        // similar reasoning as above, no more records to consume by the same consumer group at backup cluster
+        assertEquals("consumer record size is not zero", 0, records.count());
+    }
+
+    private void deleteAllTopics(EmbeddedKafkaCluster cluster) {
+        Admin client = cluster.createAdminClient();
+        try {
+            client.deleteTopics(client.listTopics().names().get());
+        } catch (Throwable e) {
+        }
+    }
+
+    private void produceMessages(EmbeddedConnectCluster cluster, String topicName, String msgPrefix) {
+        produceMessages(cluster, topicName, msgPrefix, NUM_PARTITIONS);
+    }
+
+    private void produceMessages(EmbeddedConnectCluster cluster, String topicName, String msgPrefix, int numPartitions) {
+        // produce the configured number of records to all specified partitions
+        int cnt = 0;
+        for (int r = 0; r < NUM_RECORDS_PER_PARTITION; r++)
+            for (int p = 0; p < numPartitions; p++)
+                cluster.kafka().produce(topicName, p, "key", msgPrefix + cnt++);
+    }
+
+    private Map>> consumeAllMessages(Consumer consumer) throws InterruptedException {
+        return consumeAllMessages(consumer, null, null);
+    }
+
+    private Map>> consumeAllMessages(Consumer consumer, Integer expectedRecords) throws InterruptedException {
+        return consumeAllMessages(consumer, expectedRecords, null);
+    }
+
+    private Map>> consumeAllMessages(Consumer consumer, Integer numExpectedRecords, Duration timeoutDurationMs)
+            throws InterruptedException {
+        int expectedRecords = numExpectedRecords != null ? numExpectedRecords : NUM_RECORDS_PRODUCED;
+        int timeoutMs = (int) (timeoutDurationMs != null ? timeoutDurationMs.toMillis() : RECORD_CONSUME_DURATION_MS);
+
+        List> records = new ArrayList<>();
+        waitForCondition(() -> {
+            ConsumerRecords crs = consumer.poll(Duration.ofMillis(1000));
+            for (ConsumerRecord cr : crs)
+                records.add(cr);
+            assertTrue("Consumer consumed more records than expected: " + records.size() + " (expected " + expectedRecords + ").",
+                    records.size() <= expectedRecords);
+            return records.size() == expectedRecords;
+        }, timeoutMs, "Consumer could not consume all records in time.");
+
+        consumer.commitSync();
+        return records.stream().collect(Collectors.groupingBy(c -> new TopicPartition(c.topic(), c.partition())));
+    }
+
+    private void seek(Consumer consumer, Map offsets) throws InterruptedException {
+        // In case offsets are replicated faster than actual records, wait until records are replicated before seeking
+        waitForCondition(() -> {
+            boolean ready = true;
+            Map endOffsets = consumer.endOffsets(offsets.keySet());
+            for (TopicPartition tp : offsets.keySet()) {
+                if (offsets.get(tp).offset() > endOffsets.get(tp)) {
+                    ready = false;
+                    break;
+                }
+            }
+            if (!ready)
+                Thread.sleep(1000);
+            return ready;
+        }, RECORD_TRANSFER_DURATION_MS, "Records were not replicated in time.");
+        offsets.forEach(consumer::seek);
+    }
+
+    private List allPartitions(String... topics) {
+        return IntStream.range(0, NUM_PARTITIONS)
+                .boxed()
+                .flatMap(p -> Arrays.stream(topics).map(t -> new TopicPartition(t, p)))
+                .collect(Collectors.toList());
+    }
+}
diff --git a/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/MirrorHeartBeatConnectorTest.java b/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/MirrorHeartBeatConnectorTest.java
new file mode 100644
index 0000000000000..f0d8cb95a947c
--- /dev/null
+++ b/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/MirrorHeartBeatConnectorTest.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;
+
+import static org.apache.kafka.connect.mirror.TestUtils.makeProps;
+import static org.junit.Assert.assertEquals;
+import java.util.List;
+import java.util.Map;
+import org.junit.Test;
+
+public class MirrorHeartBeatConnectorTest {
+
+    @Test
+    public void testMirrorHeartbeatConnectorDisabled() {
+        // disable the heartbeat emission
+        MirrorConnectorConfig config = new MirrorConnectorConfig(
+            makeProps("emit.heartbeats.enabled", "false"));
+
+        // MirrorHeartbeatConnector as minimum to run taskConfig()
+        MirrorHeartbeatConnector connector = new MirrorHeartbeatConnector(config);
+        List> output = connector.taskConfigs(1);
+        // expect no task will be created
+        assertEquals(0, output.size());
+    }
+
+    @Test
+    public void testReplicationDisabled() {
+        // disable the replication
+        MirrorConnectorConfig config = new MirrorConnectorConfig(
+            makeProps("enabled", "false"));
+
+        // MirrorHeartbeatConnector as minimum to run taskConfig()
+        MirrorHeartbeatConnector connector = new MirrorHeartbeatConnector(config);
+        List> output = connector.taskConfigs(1);
+        // expect one task will be created, even the replication is disabled
+        assertEquals(1, output.size());
+    }
+}
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..1cba87f90f680
--- /dev/null
+++ b/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/MirrorMakerConfigTest.java
@@ -0,0 +1,281 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.exclude", "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 include should be passed through to underlying Connectors.",
+            Arrays.asList("topic-1"), connectorConfig.getList("topics"));
+        assertEquals("Groups include should be passed through to underlying Connectors.",
+            Arrays.asList("group-2"), connectorConfig.getList("groups"));
+        assertEquals("Config properties exclude should be passed through to underlying Connectors.",
+            Arrays.asList("property-3"), connectorConfig.getList("config.properties.exclude"));
+        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 testConfigBackwardsCompatibility() {
+        MirrorMakerConfig mirrorConfig = new MirrorMakerConfig(makeProps(
+            "clusters", "a, b",
+            "groups.blacklist", "group-7",
+            "topics.blacklist", "topic3",
+            "config.properties.blacklist", "property-3",
+            "topic.filter.class", DefaultTopicFilter.class.getName()));
+        SourceAndTarget sourceAndTarget = new SourceAndTarget("source", "target");
+        Map connectorProps = mirrorConfig.connectorBaseConfig(sourceAndTarget,
+                                                                              MirrorSourceConnector.class);
+        MirrorConnectorConfig connectorConfig = new MirrorConnectorConfig(connectorProps);
+        DefaultTopicFilter.TopicFilterConfig filterConfig =
+            new DefaultTopicFilter.TopicFilterConfig(connectorProps);
+
+        assertEquals("Topics exclude should be backwards compatible.",
+                     Arrays.asList("topic3"), filterConfig.getList("topics.exclude"));
+
+        assertEquals("Groups exclude should be backwards compatible.",
+                     Arrays.asList("group-7"), connectorConfig.getList("groups.exclude"));
+
+        assertEquals("Config properties exclude should be backwards compatible.",
+                     Arrays.asList("property-3"), connectorConfig.getList("config.properties.exclude"));
+
+    }
+
+    @Test
+    public void testConfigBackwardsCompatibilitySourceTarget() {
+        MirrorMakerConfig mirrorConfig = new MirrorMakerConfig(makeProps(
+            "clusters", "a, b",
+            "source->target.topics.blacklist", "topic3",
+            "source->target.groups.blacklist", "group-7",
+            "topic.filter.class", DefaultTopicFilter.class.getName()));
+        SourceAndTarget sourceAndTarget = new SourceAndTarget("source", "target");
+        Map connectorProps = mirrorConfig.connectorBaseConfig(sourceAndTarget,
+                                                                              MirrorSourceConnector.class);
+        MirrorConnectorConfig connectorConfig = new MirrorConnectorConfig(connectorProps);
+        DefaultTopicFilter.TopicFilterConfig filterConfig =
+            new DefaultTopicFilter.TopicFilterConfig(connectorProps);
+
+        assertEquals("Topics exclude should be backwards compatible.",
+                     Arrays.asList("topic3"), filterConfig.getList("topics.exclude"));
+
+        assertEquals("Groups exclude should be backwards compatible.",
+                     Arrays.asList("group-7"), connectorConfig.getList("groups.exclude"));
+    }
+
+    @Test
+    public void testIncludesTopicFilterProperties() {
+        MirrorMakerConfig mirrorConfig = new MirrorMakerConfig(makeProps(
+            "clusters", "a, b",
+            "source->target.topics", "topic1, topic2",
+            "source->target.topics.exclude", "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.exclude should be passed through to TopicFilters.",
+            Arrays.asList("topic3"), filterConfig.getList("topics.exclude"));
+    }
+
+    @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..258dc19b28770
--- /dev/null
+++ b/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/MirrorSourceConnectorTest.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.mirror;
+
+import org.apache.kafka.common.TopicPartition;
+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.connect.connector.ConnectorContext;
+import org.apache.kafka.clients.admin.ConfigEntry;
+import org.apache.kafka.clients.admin.NewTopic;
+
+import org.junit.Test;
+
+import static org.apache.kafka.connect.mirror.MirrorConnectorConfig.TASK_TOPIC_PARTITIONS;
+import static org.apache.kafka.connect.mirror.TestUtils.makeProps;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.assertFalse;
+import static org.mockito.Mockito.any;
+import static org.mockito.Mockito.doNothing;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+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 excluded properties", targetConfig.entries().stream()
+            .anyMatch(x -> x.name().equals("min.insync.replicas")));
+    }
+
+    @Test
+    public void testMirrorSourceConnectorTaskConfig() {
+        List knownSourceTopicPartitions = new ArrayList<>();
+
+        // topic `t0` has 8 partitions
+        knownSourceTopicPartitions.add(new TopicPartition("t0", 0));
+        knownSourceTopicPartitions.add(new TopicPartition("t0", 1));
+        knownSourceTopicPartitions.add(new TopicPartition("t0", 2));
+        knownSourceTopicPartitions.add(new TopicPartition("t0", 3));
+        knownSourceTopicPartitions.add(new TopicPartition("t0", 4));
+        knownSourceTopicPartitions.add(new TopicPartition("t0", 5));
+        knownSourceTopicPartitions.add(new TopicPartition("t0", 6));
+        knownSourceTopicPartitions.add(new TopicPartition("t0", 7));
+
+        // topic `t1` has 2 partitions
+        knownSourceTopicPartitions.add(new TopicPartition("t1", 0));
+        knownSourceTopicPartitions.add(new TopicPartition("t1", 1));
+
+        // topic `t2` has 2 partitions
+        knownSourceTopicPartitions.add(new TopicPartition("t2", 0));
+        knownSourceTopicPartitions.add(new TopicPartition("t2", 1));
+
+        // MirrorConnectorConfig example for test
+        MirrorConnectorConfig config = new MirrorConnectorConfig(makeProps());
+
+        // MirrorSourceConnector as minimum to run taskConfig()
+        MirrorSourceConnector connector = new MirrorSourceConnector(knownSourceTopicPartitions, config);
+
+        // distribute the topic-partition to 3 tasks by round-robin
+        List> output = connector.taskConfigs(3);
+
+        // the expected assignments over 3 tasks:
+        // t1 -> [t0p0, t0p3, t0p6, t1p1]
+        // t2 -> [t0p1, t0p4, t0p7, t2p0]
+        // t3 -> [t0p2, t0p5, t1p0, t2p1]
+
+        Map t1 = output.get(0);
+        assertEquals("t0-0,t0-3,t0-6,t1-1", t1.get(TASK_TOPIC_PARTITIONS));
+
+        Map t2 = output.get(1);
+        assertEquals("t0-1,t0-4,t0-7,t2-0", t2.get(TASK_TOPIC_PARTITIONS));
+
+        Map t3 = output.get(2);
+        assertEquals("t0-2,t0-5,t1-0,t2-1", t3.get(TASK_TOPIC_PARTITIONS));
+    }
+
+    @Test
+    public void testRefreshTopicPartitions() throws Exception {
+        MirrorSourceConnector connector = new MirrorSourceConnector(new SourceAndTarget("source", "target"),
+                new DefaultReplicationPolicy(), new DefaultTopicFilter(), new DefaultConfigPropertyFilter());
+        connector.initialize(mock(ConnectorContext.class));
+        connector = spy(connector);
+
+        List sourceTopicPartitions = Collections.singletonList(new TopicPartition("topic", 0));
+        doReturn(sourceTopicPartitions).when(connector).findSourceTopicPartitions();
+        doReturn(Collections.emptyList()).when(connector).findTargetTopicPartitions();
+        doNothing().when(connector).createTopicPartitions(any(), any(), any());
+
+        connector.refreshTopicPartitions();
+        // if target topic is not created, refreshTopicPartitions() will call createTopicPartitions() again
+        connector.refreshTopicPartitions();
+
+        Map expectedPartitionCounts = new HashMap<>();
+        expectedPartitionCounts.put("source.topic", 1L);
+        List expectedNewTopics = Arrays.asList(new NewTopic("source.topic", 1, (short) 0));
+
+        verify(connector, times(2)).computeAndCreateTopicPartitions();
+        verify(connector, times(2)).createTopicPartitions(
+                eq(expectedPartitionCounts),
+                eq(expectedNewTopics),
+                eq(Collections.emptyMap()));
+
+        List targetTopicPartitions = Collections.singletonList(new TopicPartition("source.topic", 0));
+        doReturn(targetTopicPartitions).when(connector).findTargetTopicPartitions();
+        connector.refreshTopicPartitions();
+
+        // once target topic is created, refreshTopicPartitions() will NOT call computeAndCreateTopicPartitions() again
+        verify(connector, times(2)).computeAndCreateTopicPartitions();
+    }
+
+    @Test
+    public void testRefreshTopicPartitionsTopicOnTargetFirst() throws Exception {
+        MirrorSourceConnector connector = new MirrorSourceConnector(new SourceAndTarget("source", "target"),
+                new DefaultReplicationPolicy(), new DefaultTopicFilter(), new DefaultConfigPropertyFilter());
+        connector.initialize(mock(ConnectorContext.class));
+        connector = spy(connector);
+
+        List sourceTopicPartitions = Collections.emptyList();
+        List targetTopicPartitions = Collections.singletonList(new TopicPartition("source.topic", 0));
+        doReturn(sourceTopicPartitions).when(connector).findSourceTopicPartitions();
+        doReturn(targetTopicPartitions).when(connector).findTargetTopicPartitions();
+        doNothing().when(connector).createTopicPartitions(any(), any(), any());
+
+        // partitions appearing on the target cluster should not cause reconfiguration
+        connector.refreshTopicPartitions();
+        connector.refreshTopicPartitions();
+        verify(connector, times(0)).computeAndCreateTopicPartitions();
+
+        sourceTopicPartitions = Collections.singletonList(new TopicPartition("topic", 0));
+        doReturn(sourceTopicPartitions).when(connector).findSourceTopicPartitions();
+
+        // when partitions are added to the source cluster, reconfiguration is triggered
+        connector.refreshTopicPartitions();
+        verify(connector, times(1)).computeAndCreateTopicPartitions();
+
+    }
+}
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/java/org/apache/kafka/connect/mirror/TestUtils.java b/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/TestUtils.java
new file mode 100644
index 0000000000000..3b382d6262400
--- /dev/null
+++ b/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/TestUtils.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.connect.mirror;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class TestUtils {
+
+    static 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;
+    }
+}
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 1b2f94e461377..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,22 +19,28 @@
 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.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;
 import org.apache.kafka.connect.runtime.distributed.DistributedHerder;
 import org.apache.kafka.connect.runtime.isolation.Plugins;
 import org.apache.kafka.connect.runtime.rest.RestServer;
 import org.apache.kafka.connect.storage.ConfigBackingStore;
+import org.apache.kafka.connect.storage.Converter;
 import org.apache.kafka.connect.storage.KafkaConfigBackingStore;
 import org.apache.kafka.connect.storage.KafkaOffsetBackingStore;
 import org.apache.kafka.connect.storage.KafkaStatusBackingStore;
 import org.apache.kafka.connect.storage.StatusBackingStore;
+import org.apache.kafka.connect.util.ConnectUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.net.URI;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.Map;
 
@@ -50,43 +56,74 @@
 public class ConnectDistributed {
     private static final Logger log = LoggerFactory.getLogger(ConnectDistributed.class);
 
-    public static void main(String[] args) throws Exception {
-        if (args.length < 1) {
+    private final Time time = Time.SYSTEM;
+    private final long initStart = time.hiResClockMs();
+
+    public static void main(String[] args) {
+
+        if (args.length < 1 || Arrays.asList(args).contains("--help")) {
             log.info("Usage: ConnectDistributed worker.properties");
             Exit.exit(1);
         }
 
-        Time time = Time.SYSTEM;
-        log.info("Kafka Connect distributed worker initializing ...");
-        long initStart = time.hiResClockMs();
-        WorkerInfo initInfo = new WorkerInfo();
-        initInfo.logAll();
+        try {
+            WorkerInfo initInfo = new WorkerInfo();
+            initInfo.logAll();
+
+            String workerPropsFile = args[0];
+            Map workerProps = !workerPropsFile.isEmpty() ?
+                    Utils.propsToStringMap(Utils.loadProps(workerPropsFile)) : Collections.emptyMap();
 
-        String workerPropsFile = args[0];
-        Map workerProps = !workerPropsFile.isEmpty() ?
-                Utils.propsToStringMap(Utils.loadProps(workerPropsFile)) : Collections.emptyMap();
+            ConnectDistributed connectDistributed = new ConnectDistributed();
+            Connect connect = connectDistributed.startConnect(workerProps);
 
+            // Shutdown will be triggered by Ctrl-C or via HTTP shutdown request
+            connect.awaitStop();
+
+        } catch (Throwable t) {
+            log.error("Stopping due to error", t);
+            Exit.exit(2);
+        }
+    }
+
+    public Connect startConnect(Map workerProps) {
         log.info("Scanning for plugin classes. This might take a moment ...");
         Plugins plugins = new Plugins(workerProps);
         plugins.compareAndSwapWithDelegatingLoader();
         DistributedConfig config = new DistributedConfig(workerProps);
 
+        String kafkaClusterId = ConnectUtils.lookupKafkaClusterId(config);
+        log.debug("Kafka cluster ID: {}", kafkaClusterId);
+
         RestServer rest = new RestServer(config);
+        rest.initializeServer();
+
         URI advertisedUrl = rest.advertisedUrl();
         String workerId = advertisedUrl.getHost() + ":" + advertisedUrl.getPort();
 
         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);
 
-        StatusBackingStore statusBackingStore = new KafkaStatusBackingStore(time, worker.getInternalValueConverter());
+        Worker worker = new Worker(workerId, time, plugins, config, offsetBackingStore, connectorClientConfigOverridePolicy);
+        WorkerConfigTransformer configTransformer = worker.configTransformer();
+
+        Converter internalValueConverter = worker.getInternalValueConverter();
+        StatusBackingStore statusBackingStore = new KafkaStatusBackingStore(time, internalValueConverter);
         statusBackingStore.configure(config);
 
-        ConfigBackingStore configBackingStore = new KafkaConfigBackingStore(worker.getInternalValueConverter(), config);
+        ConfigBackingStore configBackingStore = new KafkaConfigBackingStore(
+                internalValueConverter,
+                config,
+                configTransformer);
+
+        DistributedHerder herder = new DistributedHerder(config, time, worker,
+                kafkaClusterId, statusBackingStore, configBackingStore,
+                advertisedUrl.toString(), connectorClientConfigOverridePolicy);
 
-        DistributedHerder herder = new DistributedHerder(config, time, worker, statusBackingStore, configBackingStore,
-                advertisedUrl.toString());
         final Connect connect = new Connect(herder, rest);
         log.info("Kafka Connect distributed worker initialization took {}ms", time.hiResClockMs() - initStart);
         try {
@@ -94,9 +131,10 @@ public static void main(String[] args) throws Exception {
         } catch (Exception e) {
             log.error("Failed to start Connect", e);
             connect.stop();
+            Exit.exit(3);
         }
 
-        // Shutdown will be triggered by Ctrl-C or via HTTP shutdown request
-        connect.awaitStop();
+        return connect;
     }
+
 }
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 551c6b51d3b0d..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,10 +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.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;
@@ -31,6 +33,7 @@
 import org.apache.kafka.connect.runtime.standalone.StandaloneHerder;
 import org.apache.kafka.connect.storage.FileOffsetBackingStore;
 import org.apache.kafka.connect.util.Callback;
+import org.apache.kafka.connect.util.ConnectUtils;
 import org.apache.kafka.connect.util.FutureCallback;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -54,62 +57,78 @@
 public class ConnectStandalone {
     private static final Logger log = LoggerFactory.getLogger(ConnectStandalone.class);
 
-    public static void main(String[] args) throws Exception {
+    public static void main(String[] args) {
 
-        if (args.length < 2) {
+        if (args.length < 2 || Arrays.asList(args).contains("--help")) {
             log.info("Usage: ConnectStandalone worker.properties connector1.properties [connector2.properties ...]");
             Exit.exit(1);
         }
 
-        Time time = Time.SYSTEM;
-        log.info("Kafka Connect standalone worker initializing ...");
-        long initStart = time.hiResClockMs();
-        WorkerInfo initInfo = new WorkerInfo();
-        initInfo.logAll();
+        try {
+            Time time = Time.SYSTEM;
+            log.info("Kafka Connect standalone worker initializing ...");
+            long initStart = time.hiResClockMs();
+            WorkerInfo initInfo = new WorkerInfo();
+            initInfo.logAll();
 
-        String workerPropsFile = args[0];
-        Map workerProps = !workerPropsFile.isEmpty() ?
-                Utils.propsToStringMap(Utils.loadProps(workerPropsFile)) : Collections.emptyMap();
+            String workerPropsFile = args[0];
+            Map workerProps = !workerPropsFile.isEmpty() ?
+                    Utils.propsToStringMap(Utils.loadProps(workerPropsFile)) : Collections.emptyMap();
 
-        log.info("Scanning for plugin classes. This might take a moment ...");
-        Plugins plugins = new Plugins(workerProps);
-        plugins.compareAndSwapWithDelegatingLoader();
-        StandaloneConfig config = new StandaloneConfig(workerProps);
+            log.info("Scanning for plugin classes. This might take a moment ...");
+            Plugins plugins = new Plugins(workerProps);
+            plugins.compareAndSwapWithDelegatingLoader();
+            StandaloneConfig config = new StandaloneConfig(workerProps);
 
-        RestServer rest = new RestServer(config);
-        URI advertisedUrl = rest.advertisedUrl();
-        String workerId = advertisedUrl.getHost() + ":" + advertisedUrl.getPort();
+            String kafkaClusterId = ConnectUtils.lookupKafkaClusterId(config);
+            log.debug("Kafka cluster ID: {}", kafkaClusterId);
 
-        Worker worker = new Worker(workerId, time, plugins, config, new FileOffsetBackingStore());
+            RestServer rest = new RestServer(config);
+            rest.initializeServer();
 
-        Herder herder = new StandaloneHerder(worker);
-        final Connect connect = new Connect(herder, rest);
-        log.info("Kafka Connect standalone worker initialization took {}ms", time.hiResClockMs() - initStart);
+            URI advertisedUrl = rest.advertisedUrl();
+            String workerId = advertisedUrl.getHost() + ":" + advertisedUrl.getPort();
 
-        try {
-            connect.start();
-            for (final String connectorPropsFile : Arrays.copyOfRange(args, 1, args.length)) {
-                Map connectorProps = Utils.propsToStringMap(Utils.loadProps(connectorPropsFile));
-                FutureCallback> cb = new FutureCallback<>(new Callback>() {
-                    @Override
-                    public void onCompletion(Throwable error, Herder.Created info) {
-                        if (error != null)
-                            log.error("Failed to create job for {}", connectorPropsFile);
-                        else
-                            log.info("Created connector {}", info.result().name());
-                    }
-                });
-                herder.putConnectorConfig(
-                        connectorProps.get(ConnectorConfig.NAME_CONFIG),
-                        connectorProps, false, cb);
-                cb.get();
+            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, connectorClientConfigOverridePolicy);
+            final Connect connect = new Connect(herder, rest);
+            log.info("Kafka Connect standalone worker initialization took {}ms", time.hiResClockMs() - initStart);
+
+            try {
+                connect.start();
+                for (final String connectorPropsFile : Arrays.copyOfRange(args, 1, args.length)) {
+                    Map connectorProps = Utils.propsToStringMap(Utils.loadProps(connectorPropsFile));
+                    FutureCallback> cb = new FutureCallback<>(new Callback>() {
+                        @Override
+                        public void onCompletion(Throwable error, Herder.Created info) {
+                            if (error != null)
+                                log.error("Failed to create job for {}", connectorPropsFile);
+                            else
+                                log.info("Created connector {}", info.result().name());
+                        }
+                    });
+                    herder.putConnectorConfig(
+                            connectorProps.get(ConnectorConfig.NAME_CONFIG),
+                            connectorProps, false, cb);
+                    cb.get();
+                }
+            } catch (Throwable t) {
+                log.error("Stopping after connector error", t);
+                connect.stop();
+                Exit.exit(3);
             }
+
+            // Shutdown will be triggered by Ctrl-C or via HTTP shutdown request
+            connect.awaitStop();
+
         } catch (Throwable t) {
-            log.error("Stopping after connector error", t);
-            connect.stop();
+            log.error("Stopping due to error", t);
+            Exit.exit(2);
         }
-
-        // Shutdown will be triggered by Ctrl-C or via HTTP shutdown request
-        connect.awaitStop();
     }
 }
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/converters/ByteArrayConverter.java b/connect/runtime/src/main/java/org/apache/kafka/connect/converters/ByteArrayConverter.java
index 05dff27dce220..34c552e80940d 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/converters/ByteArrayConverter.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/converters/ByteArrayConverter.java
@@ -17,17 +17,33 @@
 
 package org.apache.kafka.connect.converters;
 
+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.errors.DataException;
 import org.apache.kafka.connect.storage.Converter;
+import org.apache.kafka.connect.storage.ConverterConfig;
+import org.apache.kafka.connect.storage.HeaderConverter;
 
 import java.util.Map;
 
 /**
  * Pass-through converter for raw byte data.
+ *
+ * This implementation currently does nothing with the topic names or header names.
  */
-public class ByteArrayConverter implements Converter {
+public class ByteArrayConverter implements Converter, HeaderConverter {
+
+    private static final ConfigDef CONFIG_DEF = ConverterConfig.newConfigDef();
+
+    @Override
+    public ConfigDef config() {
+        return CONFIG_DEF;
+    }
+
+    @Override
+    public void configure(Map configs) {
+    }
 
     @Override
     public void configure(Map configs, boolean isKey) {
@@ -49,4 +65,18 @@ public SchemaAndValue toConnectData(String topic, byte[] value) {
         return new SchemaAndValue(Schema.OPTIONAL_BYTES_SCHEMA, value);
     }
 
+    @Override
+    public byte[] fromConnectHeader(String topic, String headerKey, Schema schema, Object value) {
+        return fromConnectData(topic, schema, value);
+    }
+
+    @Override
+    public SchemaAndValue toConnectHeader(String topic, String headerKey, byte[] value) {
+        return toConnectData(topic, value);
+    }
+
+    @Override
+    public void close() {
+        // do nothing
+    }
 }
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/converters/DoubleConverter.java b/connect/runtime/src/main/java/org/apache/kafka/connect/converters/DoubleConverter.java
new file mode 100644
index 0000000000000..684caa1c658e3
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/converters/DoubleConverter.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.converters;
+
+import org.apache.kafka.common.serialization.DoubleDeserializer;
+import org.apache.kafka.common.serialization.DoubleSerializer;
+import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.storage.Converter;
+import org.apache.kafka.connect.storage.HeaderConverter;
+
+/**
+ * {@link Converter} and {@link HeaderConverter} implementation that only supports serializing to and deserializing from double values.
+ * It does support handling nulls. When converting from bytes to Kafka Connect format, the converter will always return an
+ * optional FLOAT64 schema.
+ * 
            
+ * This implementation currently does nothing with the topic names or header names.
+ */
+public class DoubleConverter extends NumberConverter {
+
+    public DoubleConverter() {
+        super("double", Schema.OPTIONAL_FLOAT64_SCHEMA, new DoubleSerializer(), new DoubleDeserializer());
+    }
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/converters/FloatConverter.java b/connect/runtime/src/main/java/org/apache/kafka/connect/converters/FloatConverter.java
new file mode 100644
index 0000000000000..3f92b965cec4d
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/converters/FloatConverter.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.converters;
+
+import org.apache.kafka.common.serialization.FloatDeserializer;
+import org.apache.kafka.common.serialization.FloatSerializer;
+import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.storage.Converter;
+import org.apache.kafka.connect.storage.HeaderConverter;
+
+/**
+ * {@link Converter} and {@link HeaderConverter} implementation that only supports serializing to and deserializing from float values.
+ * It does support handling nulls. When converting from bytes to Kafka Connect format, the converter will always return an
+ * optional FLOAT32 schema.
+ * 
            
+ * This implementation currently does nothing with the topic names or header names.
+ */
+public class FloatConverter extends NumberConverter {
+
+    public FloatConverter() {
+        super("float", Schema.OPTIONAL_FLOAT32_SCHEMA, new FloatSerializer(), new FloatDeserializer());
+    }
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/converters/IntegerConverter.java b/connect/runtime/src/main/java/org/apache/kafka/connect/converters/IntegerConverter.java
new file mode 100644
index 0000000000000..f5388ce35bb69
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/converters/IntegerConverter.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.converters;
+
+import org.apache.kafka.common.serialization.IntegerDeserializer;
+import org.apache.kafka.common.serialization.IntegerSerializer;
+import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.storage.Converter;
+import org.apache.kafka.connect.storage.HeaderConverter;
+
+/**
+ * {@link Converter} and {@link HeaderConverter} implementation that only supports serializing to and deserializing from integer values.
+ * It does support handling nulls. When converting from bytes to Kafka Connect format, the converter will always return an
+ * optional INT32 schema.
+ * 
            
+ * This implementation currently does nothing with the topic names or header names.
+ */
+public class IntegerConverter extends NumberConverter {
+
+    public IntegerConverter() {
+        super("integer", Schema.OPTIONAL_INT32_SCHEMA, new IntegerSerializer(), new IntegerDeserializer());
+    }
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/converters/LongConverter.java b/connect/runtime/src/main/java/org/apache/kafka/connect/converters/LongConverter.java
new file mode 100644
index 0000000000000..f91f4fad96392
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/converters/LongConverter.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.converters;
+
+import org.apache.kafka.common.serialization.LongDeserializer;
+import org.apache.kafka.common.serialization.LongSerializer;
+import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.storage.Converter;
+import org.apache.kafka.connect.storage.HeaderConverter;
+
+/**
+ * {@link Converter} and {@link HeaderConverter} implementation that only supports serializing to and deserializing from long values.
+ * It does support handling nulls. When converting from bytes to Kafka Connect format, the converter will always return an
+ * optional INT64 schema.
+ * 
            
+ * This implementation currently does nothing with the topic names or header names.
+ */
+public class LongConverter extends NumberConverter {
+
+    public LongConverter() {
+        super("long", Schema.OPTIONAL_INT64_SCHEMA, new LongSerializer(), new LongDeserializer());
+    }
+
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/converters/NumberConverter.java b/connect/runtime/src/main/java/org/apache/kafka/connect/converters/NumberConverter.java
new file mode 100644
index 0000000000000..0af4aacbbd603
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/converters/NumberConverter.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.connect.converters;
+
+import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.common.errors.SerializationException;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.serialization.Serializer;
+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;
+import org.apache.kafka.connect.storage.ConverterType;
+import org.apache.kafka.connect.storage.HeaderConverter;
+import org.apache.kafka.connect.storage.StringConverterConfig;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * {@link Converter} and {@link HeaderConverter} implementation that only supports serializing to and deserializing from number values.
+ * It does support handling nulls. When converting from bytes to Kafka Connect format, the converter will always return the specified
+ * schema.
+ * 
            
+ * This implementation currently does nothing with the topic names or header names.
+ */
+abstract class NumberConverter implements Converter, HeaderConverter {
+
+    private final Serializer serializer;
+    private final Deserializer deserializer;
+    private final String typeName;
+    private final Schema schema;
+
+    /**
+     * Create the converter.
+     *
+     * @param typeName the displayable name of the type; may not be null
+     * @param schema the optional schema to be used for all deserialized forms; may not be null
+     * @param serializer the serializer; may not be null
+     * @param deserializer the deserializer; may not be null
+     */
+    protected NumberConverter(String typeName, Schema schema, Serializer serializer, Deserializer deserializer) {
+        this.typeName = typeName;
+        this.schema = schema;
+        this.serializer = serializer;
+        this.deserializer = deserializer;
+        assert this.serializer != null;
+        assert this.deserializer != null;
+        assert this.typeName != null;
+        assert this.schema != null;
+    }
+
+    @Override
+    public ConfigDef config() {
+        return NumberConverterConfig.configDef();
+    }
+
+    @Override
+    public void configure(Map configs) {
+        NumberConverterConfig conf = new NumberConverterConfig(configs);
+        boolean isKey = conf.type() == ConverterType.KEY;
+        serializer.configure(configs, isKey);
+        deserializer.configure(configs, isKey);
+
+    }
+
+    @Override
+    public void configure(Map configs, boolean isKey) {
+        Map conf = new HashMap<>(configs);
+        conf.put(StringConverterConfig.TYPE_CONFIG, isKey ? ConverterType.KEY.getName() : ConverterType.VALUE.getName());
+        configure(conf);
+    }
+
+    @SuppressWarnings("unchecked")
+    protected T cast(Object value) {
+        return (T) value;
+    }
+
+    @Override
+    public byte[] fromConnectData(String topic, Schema schema, Object value) {
+        try {
+            return serializer.serialize(topic, value == null ? null : cast(value));
+        } catch (ClassCastException e) {
+            throw new DataException("Failed to serialize to " + typeName + " (was " + value.getClass() + "): ", e);
+        } catch (SerializationException e) {
+            throw new DataException("Failed to serialize to " + typeName + ": ", e);
+        }
+    }
+
+    @Override
+    public SchemaAndValue toConnectData(String topic, byte[] value) {
+        try {
+            return new SchemaAndValue(schema, deserializer.deserialize(topic, value));
+        } catch (SerializationException e) {
+            throw new DataException("Failed to deserialize " + typeName + ": ", e);
+        }
+    }
+
+    @Override
+    public byte[] fromConnectHeader(String topic, String headerKey, Schema schema, Object value) {
+        return fromConnectData(topic, schema, value);
+    }
+
+    @Override
+    public SchemaAndValue toConnectHeader(String topic, String headerKey, byte[] value) {
+        return toConnectData(topic, value);
+    }
+
+    @Override
+    public void close() throws IOException {
+    }
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/converters/NumberConverterConfig.java b/connect/runtime/src/main/java/org/apache/kafka/connect/converters/NumberConverterConfig.java
new file mode 100644
index 0000000000000..49ad98673621d
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/converters/NumberConverterConfig.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.converters;
+
+import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.connect.storage.ConverterConfig;
+
+import java.util.Map;
+
+/**
+ * Configuration options for instances of {@link LongConverter}, {@link IntegerConverter}, {@link ShortConverter}, {@link DoubleConverter},
+ * and {@link FloatConverter} instances.
+ */
+public class NumberConverterConfig extends ConverterConfig {
+
+    private final static ConfigDef CONFIG = ConverterConfig.newConfigDef();
+
+    public static ConfigDef configDef() {
+        return CONFIG;
+    }
+
+    public NumberConverterConfig(Map props) {
+        super(CONFIG, props);
+    }
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/converters/ShortConverter.java b/connect/runtime/src/main/java/org/apache/kafka/connect/converters/ShortConverter.java
new file mode 100644
index 0000000000000..1c455b189a3f9
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/converters/ShortConverter.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.converters;
+
+import org.apache.kafka.common.serialization.ShortDeserializer;
+import org.apache.kafka.common.serialization.ShortSerializer;
+import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.storage.Converter;
+import org.apache.kafka.connect.storage.HeaderConverter;
+
+/**
+ * {@link Converter} and {@link HeaderConverter} implementation that only supports serializing to and deserializing from short values.
+ * It does support handling nulls. When converting from bytes to Kafka Connect format, the converter will always return an
+ * optional INT16 schema.
+ * 
            
+ * This implementation currently does nothing with the topic names or header names.
+ */
+public class ShortConverter extends NumberConverter {
+
+    public ShortConverter() {
+        super("short", Schema.OPTIONAL_INT16_SCHEMA, new ShortSerializer(), new ShortDeserializer());
+    }
+}
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 fbe0ae2afb2ac..704259858e584 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,14 +16,21 @@
  */
 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;
 import org.apache.kafka.common.config.ConfigDef.Type;
+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;
+import org.apache.kafka.connect.runtime.rest.entities.ActiveTopicsInfo;
 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.ConfigKeyInfo;
@@ -46,6 +53,7 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.LinkedHashMap;
 import java.util.LinkedHashSet;
 import java.util.LinkedList;
@@ -53,6 +61,11 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+import java.util.stream.Collectors;
 
 /**
  * Abstract Herder implementation which handles connector/task lifecycle tracking. Extensions
@@ -79,19 +92,34 @@ public abstract class AbstractHerder implements Herder, TaskStatus.Listener, Con
 
     private final String workerId;
     protected final Worker worker;
+    private final String kafkaClusterId;
     protected final StatusBackingStore statusBackingStore;
     protected final ConfigBackingStore configBackingStore;
+    private final ConnectorClientConfigOverridePolicy connectorClientConfigOverridePolicy;
+    protected volatile boolean running = false;
+    private final ExecutorService connectorExecutor;
 
     private Map tempConnectors = new ConcurrentHashMap<>();
 
     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;
+        this.connectorExecutor = Executors.newCachedThreadPool();
+    }
+
+    @Override
+    public String kafkaClusterId() {
+        return kafkaClusterId;
     }
 
     protected abstract int generation();
@@ -106,6 +134,12 @@ protected void stopServices() {
         this.statusBackingStore.stop();
         this.configBackingStore.stop();
         this.worker.stop();
+        this.connectorExecutor.shutdown();
+    }
+
+    @Override
+    public boolean isRunning() {
+        return running;
     }
 
     @Override
@@ -166,10 +200,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))
@@ -194,6 +233,38 @@ public Plugins plugins() {
      */
     protected abstract Map config(String connName);
 
+    @Override
+    public void connectorConfig(String connName, Callback> callback) {
+        // Subset of connectorInfo, so piggy back on that implementation
+        connectorInfo(connName, (error, result) -> {
+            if (error != null)
+                callback.onCompletion(error, null);
+            else
+                callback.onCompletion(null, result.config());
+        });
+    }
+
+    @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);
@@ -218,6 +289,25 @@ public ConnectorStateInfo connectorStatus(String connName) {
             conf == null ? ConnectorType.UNKNOWN : connectorTypeForClass(conf.get(ConnectorConfig.CONNECTOR_CLASS_CONFIG)));
     }
 
+    @Override
+    public ActiveTopicsInfo connectorActiveTopics(String connName) {
+        Collection topics = statusBackingStore.getAllTopics(connName).stream()
+                .map(TopicStatus::topic)
+                .collect(Collectors.toList());
+        return new ActiveTopicsInfo(connName, topics);
+    }
+
+    @Override
+    public void resetConnectorActiveTopics(String connName) {
+        statusBackingStore.getAllTopics(connName).stream()
+                .forEach(status -> statusBackingStore.deleteTopic(status.connector(), status.topic()));
+    }
+
+    @Override
+    public StatusBackingStore statusBackingStore() {
+        return statusBackingStore;
+    }
+
     @Override
     public ConnectorStateInfo.TaskState taskStatus(ConnectorTaskId id) {
         TaskStatus status = statusBackingStore.get(id);
@@ -236,43 +326,183 @@ protected Map validateBasicConnectorConfig(Connector connec
     }
 
     @Override
-    public ConfigInfos validateConnectorConfig(Map connectorProps) {
+    public void validateConnectorConfig(Map connectorProps, Callback callback) {
+        validateConnectorConfig(connectorProps, callback, true);
+    }
+
+    @Override
+    public void validateConnectorConfig(Map connectorProps, Callback callback, boolean doLog) {
+        connectorExecutor.submit(() -> {
+            try {
+                ConfigInfos result = validateConnectorConfig(connectorProps, doLog);
+                callback.onCompletion(null, result);
+            } catch (Throwable t) {
+                callback.onCompletion(t, null);
+            }
+        });
+    }
+
+    ConfigInfos validateConnectorConfig(Map connectorProps, boolean doLog) {
+        if (worker.configTransformer() != null) {
+            connectorProps = worker.configTransformer().transform(connectorProps);
+        }
         String connType = connectorProps.get(ConnectorConfig.CONNECTOR_CLASS_CONFIG);
         if (connType == null)
             throw new BadRequestException("Connector config " + connectorProps + " contains no connector type");
 
-        List configValues = new ArrayList<>();
-        Map configKeys = new LinkedHashMap<>();
-        Set allGroups = new LinkedHashSet<>();
-
         Connector connector = getConnector(connType);
+        org.apache.kafka.connect.health.ConnectorType connectorType;
         ClassLoader savedLoader = plugins().compareAndSwapLoaders(connector);
         try {
-            ConfigDef baseConfigDef = (connector instanceof SourceConnector)
-                    ? SourceConnectorConfig.configDef()
-                    : SinkConnectorConfig.configDef();
+            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(
                     connector,
                     enrichedConfigDef,
                     connectorProps
             );
-            configValues.addAll(validatedConnectorConfig.values());
-            configKeys.putAll(enrichedConfigDef.configKeys());
-            allGroups.addAll(enrichedConfigDef.groups());
+            List configValues = new ArrayList<>(validatedConnectorConfig.values());
+            Map configKeys = new LinkedHashMap<>(enrichedConfigDef.configKeys());
+            Set allGroups = new LinkedHashSet<>(enrichedConfigDef.groups());
 
             // do custom connector-specific validation
-            Config config = connector.validate(connectorProps);
             ConfigDef configDef = connector.config();
+            if (null == configDef) {
+                throw new BadRequestException(
+                        String.format(
+                                "%s.config() must return a ConfigDef that is not null.",
+                                connector.getClass().getName()
+                        )
+                );
+            }
+            Config config = connector.validate(connectorProps);
+            if (null == config) {
+                throw new BadRequestException(
+                        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, doLog);
+            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;
@@ -283,8 +513,8 @@ public static ConfigInfos generateResult(String connType, Map
             String configName = configValue.name();
             configValueMap.put(configName, configValue);
             if (!configKeys.containsKey(configName)) {
-                configValue.addErrorMessage("Configuration is not defined: " + configName);
                 configInfoList.add(new ConfigInfo(null, convertConfigValue(configValue, null)));
+                errorCount += configValue.errorMessages().size();
             }
         }
 
@@ -304,7 +534,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();
 
@@ -386,7 +620,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
             );
@@ -403,4 +637,43 @@ private String trace(Throwable t) {
             return null;
         }
     }
+
+    /*
+     * Performs a reverse transformation on a set of task configs, by replacing values with variable references.
+     */
+    public static List> reverseTransform(String connName,
+                                                             ClusterConfigState configState,
+                                                             List> configs) {
+
+        // Find the config keys in the raw connector config that have variable references
+        Map rawConnConfig = configState.rawConnectorConfig(connName);
+        Set connKeysWithVariableValues = keysWithVariableValues(rawConnConfig, ConfigTransformer.DEFAULT_PATTERN);
+
+        List> result = new ArrayList<>();
+        for (Map config : configs) {
+            Map newConfig = new HashMap<>(config);
+            for (String key : connKeysWithVariableValues) {
+                if (newConfig.containsKey(key)) {
+                    newConfig.put(key, rawConnConfig.get(key));
+                }
+            }
+            result.add(newConfig);
+        }
+        return result;
+    }
+
+    // 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.find()) {
+                    keys.add(config.getKey());
+                }
+            }
+        }
+        return keys;
+    }
+
 }
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/AbstractStatus.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/AbstractStatus.java
index 00a050a310f19..dd1b94a854e6b 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/AbstractStatus.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/AbstractStatus.java
@@ -16,6 +16,8 @@
  */
 package org.apache.kafka.connect.runtime;
 
+import java.util.Objects;
+
 public abstract class AbstractStatus {
 
     public enum State {
@@ -81,12 +83,11 @@ public boolean equals(Object o) {
 
         AbstractStatus that = (AbstractStatus) o;
 
-        if (generation != that.generation) return false;
-        if (id != null ? !id.equals(that.id) : that.id != null) return false;
-        if (state != that.state) return false;
-        if (trace != null ? !trace.equals(that.trace) : that.trace != null) return false;
-        return workerId != null ? workerId.equals(that.workerId) : that.workerId == null;
-
+        return generation == that.generation
+                && Objects.equals(id, that.id)
+                && state == that.state
+                && Objects.equals(trace, that.trace)
+                && Objects.equals(workerId, that.workerId);
     }
 
     @Override
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/CloseableConnectorContext.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/CloseableConnectorContext.java
new file mode 100644
index 0000000000000..7a09a9018c8b4
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/CloseableConnectorContext.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.runtime;
+
+import org.apache.kafka.connect.connector.ConnectorContext;
+import org.apache.kafka.connect.errors.ConnectException;
+
+import java.io.Closeable;
+
+public interface CloseableConnectorContext extends ConnectorContext, Closeable {
+
+    /**
+     * Close this connector context, causing all future calls to it to throw {@link ConnectException}.
+     * This is useful to prevent zombie connector threads from making such calls after their connector
+     * instance should be shut down.
+     */
+    void close();
+}
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 846ed1a883570..80eef0369abc0 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
@@ -16,10 +16,12 @@
  */
 package org.apache.kafka.connect.runtime;
 
+import org.apache.kafka.common.utils.Exit;
 import org.apache.kafka.connect.runtime.rest.RestServer;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.net.URI;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.atomic.AtomicBoolean;
 
@@ -47,10 +49,10 @@ public Connect(Herder herder, RestServer rest) {
     public void start() {
         try {
             log.info("Kafka Connect starting");
-            Runtime.getRuntime().addShutdownHook(shutdownHook);
+            Exit.addShutdownHook("connect-shutdown-hook", shutdownHook);
 
             herder.start();
-            rest.start(herder);
+            rest.initializeResources(herder);
 
             log.info("Kafka Connect started");
         } finally {
@@ -82,6 +84,19 @@ public void awaitStop() {
         }
     }
 
+    public boolean isRunning() {
+        return herder.isRunning();
+    }
+
+    // Visible for testing
+    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 5bbe148f64463..2871bbe8be2af 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
@@ -21,16 +21,21 @@
 import org.apache.kafka.common.MetricNameTemplate;
 import org.apache.kafka.common.metrics.Gauge;
 import org.apache.kafka.common.metrics.JmxReporter;
+import org.apache.kafka.common.metrics.KafkaMetricsContext;
 import org.apache.kafka.common.metrics.MetricConfig;
 import org.apache.kafka.common.metrics.Metrics;
+import org.apache.kafka.common.metrics.MetricsContext;
 import org.apache.kafka.common.metrics.MetricsReporter;
 import org.apache.kafka.common.metrics.Sensor;
+import org.apache.kafka.common.metrics.internals.MetricsUtils;
 import org.apache.kafka.common.utils.AppInfoParser;
 import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.connect.runtime.distributed.DistributedConfig;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.LinkedHashMap;
 import java.util.List;
@@ -62,21 +67,36 @@ public class ConnectMetrics {
      * @param workerId the worker identifier; may not be null
      * @param config   the worker configuration; may not be null
      * @param time     the time; may not be null
+     * @param clusterId the Kafka cluster ID
      */
-    public ConnectMetrics(String workerId, WorkerConfig config, Time time) {
+    public ConnectMetrics(String workerId, WorkerConfig config, Time time, String clusterId) {
         this.workerId = workerId;
         this.time = time;
 
-        MetricConfig metricConfig = new MetricConfig().samples(config.getInt(CommonClientConfigs.METRICS_NUM_SAMPLES_CONFIG))
-                                                      .timeWindow(config.getLong(CommonClientConfigs.METRICS_SAMPLE_WINDOW_MS_CONFIG),
-                                                                  TimeUnit.MILLISECONDS).recordLevel(
-                        Sensor.RecordingLevel.forName(config.getString(CommonClientConfigs.METRICS_RECORDING_LEVEL_CONFIG)));
-        List reporters = config.getConfiguredInstances(CommonClientConfigs.METRIC_REPORTER_CLASSES_CONFIG,
-                                                                        MetricsReporter.class);
-        reporters.add(new JmxReporter(JMX_PREFIX));
-        this.metrics = new Metrics(metricConfig, reporters, time);
+        int numSamples = config.getInt(CommonClientConfigs.METRICS_NUM_SAMPLES_CONFIG);
+        long sampleWindowMs = config.getLong(CommonClientConfigs.METRICS_SAMPLE_WINDOW_MS_CONFIG);
+        String metricsRecordingLevel = config.getString(CommonClientConfigs.METRICS_RECORDING_LEVEL_CONFIG);
+        List reporters = config.getConfiguredInstances(CommonClientConfigs.METRIC_REPORTER_CLASSES_CONFIG, MetricsReporter.class);
+
+        MetricConfig metricConfig = new MetricConfig().samples(numSamples)
+                .timeWindow(sampleWindowMs, TimeUnit.MILLISECONDS).recordLevel(
+                        Sensor.RecordingLevel.forName(metricsRecordingLevel));
+        JmxReporter jmxReporter = new JmxReporter();
+        jmxReporter.configure(config.originals());
+        reporters.add(jmxReporter);
+
+        Map contextLabels = new HashMap<>();
+        contextLabels.putAll(config.originalsWithPrefix(CommonClientConfigs.METRICS_CONTEXT_PREFIX));
+        contextLabels.put(WorkerConfig.CONNECT_KAFKA_CLUSTER_ID, clusterId);
+        Object groupId = config.originals().get(DistributedConfig.GROUP_ID_CONFIG);
+        if (groupId != null) {
+            contextLabels.put(WorkerConfig.CONNECT_GROUP_ID, groupId);
+        }
+        MetricsContext metricsContext = new KafkaMetricsContext(JMX_PREFIX, contextLabels);
+        this.metrics = new Metrics(metricConfig, reporters, time, metricsContext);
+
         LOG.debug("Registering Connect metrics with JMX for worker '{}'", workerId);
-        AppInfoParser.registerAppInfo(JMX_PREFIX, workerId, metrics);
+        AppInfoParser.registerAppInfo(JMX_PREFIX, workerId, metrics, time.milliseconds());
     }
 
     /**
@@ -129,7 +149,7 @@ public MetricGroup group(String groupName, String... tagKeyValues) {
     }
 
     protected MetricGroupId groupId(String groupName, String... tagKeyValues) {
-        Map tags = tags(tagKeyValues);
+        Map tags = MetricsUtils.getTags(tagKeyValues);
         return new MetricGroupId(groupName, tags);
     }
 
@@ -227,7 +247,7 @@ public String toString() {
      * the {@link Metrics} class, so that the sensor names are made to be unique (based on the group name)
      * and so the sensors are removed when this group is {@link #close() closed}.
      */
-    public class MetricGroup {
+    public class MetricGroup implements AutoCloseable {
         private final MetricGroupId groupId;
         private final Set sensorNames = new HashSet<>();
         private final String sensorPrefix;
@@ -300,12 +320,7 @@ Map tags() {
         public  void addValueMetric(MetricNameTemplate nameTemplate, final LiteralSupplier supplier) {
             MetricName metricName = metricName(nameTemplate);
             if (metrics().metric(metricName) == null) {
-                metrics().addMetric(metricName, new Gauge() {
-                    @Override
-                    public T value(MetricConfig config, long now) {
-                        return supplier.metricValue(now);
-                    }
-                });
+                metrics().addMetric(metricName, (Gauge) (config, now) -> supplier.metricValue(now));
             }
         }
 
@@ -319,12 +334,7 @@ public T value(MetricConfig config, long now) {
         public  void addImmutableValueMetric(MetricNameTemplate nameTemplate, final T value) {
             MetricName metricName = metricName(nameTemplate);
             if (metrics().metric(metricName) == null) {
-                metrics().addMetric(metricName, new Gauge() {
-                    @Override
-                    public T value(MetricConfig config, long now) {
-                        return value;
-                    }
-                });
+                metrics().addMetric(metricName, (Gauge) (config, now) -> value);
             }
         }
 
@@ -424,22 +434,6 @@ public interface LiteralSupplier {
         T metricValue(long now);
     }
 
-    /**
-     * Create a set of tags using the supplied key and value pairs. The order of the tags will be kept.
-     *
-     * @param keyValue the key and value pairs for the tags; must be an even number
-     * @return the map of tags that can be supplied to the {@link Metrics} methods; never null
-     */
-    static Map tags(String... keyValue) {
-        if ((keyValue.length % 2) != 0)
-            throw new IllegalArgumentException("keyValue needs to be specified in pairs");
-        Map tags = new LinkedHashMap<>();
-        for (int i = 0; i < keyValue.length; i += 2) {
-            tags.put(keyValue[i], keyValue[i + 1]);
-        }
-        return tags;
-    }
-
     /**
      * Utility to generate the documentation for the Connect metrics.
      *
@@ -449,4 +443,4 @@ public static void main(String[] args) {
         ConnectMetricsRegistry metrics = new ConnectMetricsRegistry();
         System.out.println(Metrics.toHtmlTable(JMX_PREFIX, metrics.getAllTemplates()));
     }
-}
\ No newline at end of file
+}
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 d78576ec3eea6..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 {
@@ -34,12 +36,19 @@ public class ConnectMetricsRegistry {
     public static final String SINK_TASK_GROUP_NAME = "sink-task-metrics";
     public static final String WORKER_GROUP_NAME = "connect-worker-metrics";
     public static final String WORKER_REBALANCE_GROUP_NAME = "connect-worker-rebalance-metrics";
+    public static final String TASK_ERROR_HANDLING_GROUP_NAME = "task-error-metrics";
 
     private final List allTemplates = new ArrayList<>();
     public final MetricNameTemplate connectorStatus;
     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;
@@ -86,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;
@@ -93,6 +103,16 @@ public class ConnectMetricsRegistry {
     public final MetricNameTemplate rebalanceTimeMax;
     public final MetricNameTemplate rebalanceTimeAvg;
     public final MetricNameTemplate rebalanceTimeSinceLast;
+    public final MetricNameTemplate recordProcessingFailures;
+    public final MetricNameTemplate recordProcessingErrors;
+    public final MetricNameTemplate recordsSkipped;
+    public final MetricNameTemplate retries;
+    public final MetricNameTemplate errorsLogged;
+    public final MetricNameTemplate dlqProduceRequests;
+    public final MetricNameTemplate dlqProduceFailures;
+    public final MetricNameTemplate lastErrorTimestamp;
+
+    public Map connectorStatusMetrics;
 
     public ConnectMetricsRegistry() {
         this(new LinkedHashSet());
@@ -279,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,
@@ -294,6 +340,28 @@ public ConnectMetricsRegistry(Set tags) {
                                           "The average time in milliseconds spent by this worker to rebalance.", rebalanceTags);
         rebalanceTimeSinceLast = createTemplate("time-since-last-rebalance-ms", WORKER_REBALANCE_GROUP_NAME,
                                                 "The time in milliseconds since this worker completed the most recent rebalance.", rebalanceTags);
+
+        /***** Task Error Handling Metrics *****/
+        Set taskErrorHandlingTags = new LinkedHashSet<>(tags);
+        taskErrorHandlingTags.add(CONNECTOR_TAG_NAME);
+        taskErrorHandlingTags.add(TASK_TAG_NAME);
+
+        recordProcessingFailures = createTemplate("total-record-failures", TASK_ERROR_HANDLING_GROUP_NAME,
+                "The number of record processing failures in this task.", taskErrorHandlingTags);
+        recordProcessingErrors = createTemplate("total-record-errors", TASK_ERROR_HANDLING_GROUP_NAME,
+                "The number of record processing errors in this task. ", taskErrorHandlingTags);
+        recordsSkipped = createTemplate("total-records-skipped", TASK_ERROR_HANDLING_GROUP_NAME,
+                "The number of records skipped due to errors.", taskErrorHandlingTags);
+        retries = createTemplate("total-retries", TASK_ERROR_HANDLING_GROUP_NAME,
+                "The number of operations retried.", taskErrorHandlingTags);
+        errorsLogged = createTemplate("total-errors-logged", TASK_ERROR_HANDLING_GROUP_NAME,
+                "The number of errors that were logged.", taskErrorHandlingTags);
+        dlqProduceRequests = createTemplate("deadletterqueue-produce-requests", TASK_ERROR_HANDLING_GROUP_NAME,
+                "The number of attempted writes to the dead letter queue.", taskErrorHandlingTags);
+        dlqProduceFailures = createTemplate("deadletterqueue-produce-failures", TASK_ERROR_HANDLING_GROUP_NAME,
+                "The number of failed writes to the dead letter queue.", taskErrorHandlingTags);
+        lastErrorTimestamp = createTemplate("last-error-timestamp", TASK_ERROR_HANDLING_GROUP_NAME,
+                "The epoch timestamp when this task last encountered an error.", taskErrorHandlingTags);
     }
 
     private MetricNameTemplate createTemplate(String name, String group, String doc, Set tags) {
@@ -337,4 +405,8 @@ public String workerGroupName() {
     public String workerRebalanceGroupName() {
         return WORKER_REBALANCE_GROUP_NAME;
     }
+
+    public String taskErrorHandlingGroupName() {
+        return TASK_ERROR_HANDLING_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 0f8c39088eb86..ca9d33c1f5e05 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
@@ -22,21 +22,32 @@
 import org.apache.kafka.common.config.ConfigDef.Type;
 import org.apache.kafka.common.config.ConfigDef.Width;
 import org.apache.kafka.common.config.ConfigException;
+import org.apache.kafka.common.utils.Utils;
 import org.apache.kafka.connect.connector.ConnectRecord;
 import org.apache.kafka.connect.errors.ConnectException;
+import org.apache.kafka.connect.runtime.errors.ToleranceType;
 import org.apache.kafka.connect.runtime.isolation.PluginDesc;
 import org.apache.kafka.connect.runtime.isolation.Plugins;
 import org.apache.kafka.connect.transforms.Transformation;
+import org.apache.kafka.connect.transforms.predicates.Predicate;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
+import java.lang.reflect.Modifier;
 import java.util.ArrayList;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.HashSet;
 import java.util.LinkedHashSet;
 import java.util.List;
+import java.util.Locale;
 import java.util.Map;
+import java.util.Set;
+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;
+import static org.apache.kafka.common.config.ConfigDef.ValidString.in;
 
 /**
  * 
            
@@ -50,8 +61,12 @@
  * 
            
  */
 public class ConnectorConfig extends AbstractConfig {
+    private static final Logger log = LoggerFactory.getLogger(ConnectorConfig.class);
+
     protected static final String COMMON_GROUP = "Common";
     protected static final String TRANSFORMS_GROUP = "Transforms";
+    protected static final String PREDICATES_GROUP = "Predicates";
+    protected static final String ERROR_GROUP = "Error Handling";
 
     public static final String NAME_CONFIG = "name";
     private static final String NAME_DOC = "Globally unique name to use for this connector.";
@@ -72,6 +87,13 @@ public class ConnectorConfig extends AbstractConfig {
     public static final String VALUE_CONVERTER_CLASS_DOC = WorkerConfig.VALUE_CONVERTER_CLASS_DOC;
     public static final String VALUE_CONVERTER_CLASS_DISPLAY = "Value converter class";
 
+    public static final String HEADER_CONVERTER_CLASS_CONFIG = WorkerConfig.HEADER_CONVERTER_CLASS_CONFIG;
+    public static final String HEADER_CONVERTER_CLASS_DOC = WorkerConfig.HEADER_CONVERTER_CLASS_DOC;
+    public static final String HEADER_CONVERTER_CLASS_DISPLAY = "Header converter class";
+    // The Connector config should not have a default for the header converter, since the absence of a config property means that
+    // the worker config settings should be used. Thus, we set the default to null here.
+    public static final String HEADER_CONVERTER_CLASS_DEFAULT = null;
+
     public static final String TASKS_MAX_CONFIG = "tasks.max";
     private static final String TASKS_MAX_DOC = "Maximum number of tasks to use for this connector.";
     public static final int TASKS_MAX_DEFAULT = 1;
@@ -83,38 +105,121 @@ public class ConnectorConfig extends AbstractConfig {
     private static final String TRANSFORMS_DOC = "Aliases for the transformations to be applied to records.";
     private static final String TRANSFORMS_DISPLAY = "Transforms";
 
+    public static final String PREDICATES_CONFIG = "predicates";
+    private static final String PREDICATES_DOC = "Aliases for the predicates used by transformations.";
+    private static final String PREDICATES_DISPLAY = "Predicates";
+
+    public static final String CONFIG_RELOAD_ACTION_CONFIG = "config.action.reload";
+    private static final String CONFIG_RELOAD_ACTION_DOC =
+            "The action that Connect should take on the connector when changes in external " +
+            "configuration providers result in a change in the connector's configuration properties. " +
+            "A value of 'none' indicates that Connect will do nothing. " +
+            "A value of 'restart' indicates that Connect should restart/reload the connector with the " +
+            "updated configuration properties." +
+            "The restart may actually be scheduled in the future if the external configuration provider " +
+            "indicates that a configuration value will expire in the future.";
+
+    private static final String CONFIG_RELOAD_ACTION_DISPLAY = "Reload Action";
+    public static final String CONFIG_RELOAD_ACTION_NONE = Herder.ConfigReloadAction.NONE.name().toLowerCase(Locale.ROOT);
+    public static final String CONFIG_RELOAD_ACTION_RESTART = Herder.ConfigReloadAction.RESTART.name().toLowerCase(Locale.ROOT);
+
+    public static final String ERRORS_RETRY_TIMEOUT_CONFIG = "errors.retry.timeout";
+    public static final String ERRORS_RETRY_TIMEOUT_DISPLAY = "Retry Timeout for Errors";
+    public static final int ERRORS_RETRY_TIMEOUT_DEFAULT = 0;
+    public static final String ERRORS_RETRY_TIMEOUT_DOC = "The maximum duration in milliseconds that a failed operation " +
+            "will be reattempted. The default is 0, which means no retries will be attempted. Use -1 for infinite retries.";
+
+    public static final String ERRORS_RETRY_MAX_DELAY_CONFIG = "errors.retry.delay.max.ms";
+    public static final String ERRORS_RETRY_MAX_DELAY_DISPLAY = "Maximum Delay Between Retries for Errors";
+    public static final int ERRORS_RETRY_MAX_DELAY_DEFAULT = 60000;
+    public static final String ERRORS_RETRY_MAX_DELAY_DOC = "The maximum duration in milliseconds between consecutive retry attempts. " +
+            "Jitter will be added to the delay once this limit is reached to prevent thundering herd issues.";
+
+    public static final String ERRORS_TOLERANCE_CONFIG = "errors.tolerance";
+    public static final String ERRORS_TOLERANCE_DISPLAY = "Error Tolerance";
+    public static final ToleranceType ERRORS_TOLERANCE_DEFAULT = ToleranceType.NONE;
+    public static final String ERRORS_TOLERANCE_DOC = "Behavior for tolerating errors during connector operation. 'none' is the default value " +
+            "and signals that any error will result in an immediate connector task failure; 'all' changes the behavior to skip over problematic records.";
+
+    public static final String ERRORS_LOG_ENABLE_CONFIG = "errors.log.enable";
+    public static final String ERRORS_LOG_ENABLE_DISPLAY = "Log Errors";
+    public static final boolean ERRORS_LOG_ENABLE_DEFAULT = false;
+    public static final String ERRORS_LOG_ENABLE_DOC = "If true, write each error and the details of the failed operation and problematic record " +
+            "to the Connect application log. This is 'false' by default, so that only errors that are not tolerated are reported.";
+
+    public static final String ERRORS_LOG_INCLUDE_MESSAGES_CONFIG = "errors.log.include.messages";
+    public static final String ERRORS_LOG_INCLUDE_MESSAGES_DISPLAY = "Log Error Details";
+    public static final boolean ERRORS_LOG_INCLUDE_MESSAGES_DEFAULT = false;
+    public static final String ERRORS_LOG_INCLUDE_MESSAGES_DOC = "Whether to the include in the log the Connect record that resulted in " +
+            "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.";
+    public static final String PREDICATES_PREFIX = "predicates.";
+
     private final EnrichedConnectorConfig enrichedConfig;
     private static class EnrichedConnectorConfig extends AbstractConfig {
         EnrichedConnectorConfig(ConfigDef configDef, Map props) {
             super(configDef, props);
         }
 
+        @Override
         public Object get(String key) {
             return super.get(key);
         }
     }
 
     public static ConfigDef configDef() {
+        int orderInGroup = 0;
+        int orderInErrorGroup = 0;
         return new ConfigDef()
-                .define(NAME_CONFIG, Type.STRING, Importance.HIGH, NAME_DOC, COMMON_GROUP, 1, Width.MEDIUM, NAME_DISPLAY)
-                .define(CONNECTOR_CLASS_CONFIG, Type.STRING, Importance.HIGH, CONNECTOR_CLASS_DOC, COMMON_GROUP, 2, Width.LONG, CONNECTOR_CLASS_DISPLAY)
-                .define(TASKS_MAX_CONFIG, Type.INT, TASKS_MAX_DEFAULT, atLeast(TASKS_MIN_CONFIG), Importance.HIGH, TASKS_MAX_DOC, COMMON_GROUP, 3, Width.SHORT, TASK_MAX_DISPLAY)
-                .define(KEY_CONVERTER_CLASS_CONFIG, Type.CLASS, null, Importance.LOW, KEY_CONVERTER_CLASS_DOC, COMMON_GROUP, 4, Width.SHORT, KEY_CONVERTER_CLASS_DISPLAY)
-                .define(VALUE_CONVERTER_CLASS_CONFIG, Type.CLASS, null, Importance.LOW, VALUE_CONVERTER_CLASS_DOC, COMMON_GROUP, 5, Width.SHORT, VALUE_CONVERTER_CLASS_DISPLAY)
-                .define(TRANSFORMS_CONFIG, Type.LIST, null, new ConfigDef.Validator() {
-                    @Override
-                    public void ensureValid(String name, Object value) {
-                        if (value == null) return;
-                        final List transformAliases = (List) value;
-                        if (transformAliases.size() > new HashSet<>(transformAliases).size()) {
-                            throw new ConfigException(name, value, "Duplicate alias provided.");
-                        }
-                    }
-                }, Importance.LOW, TRANSFORMS_DOC, TRANSFORMS_GROUP, 6, Width.LONG, TRANSFORMS_DISPLAY);
+                .define(NAME_CONFIG, Type.STRING, ConfigDef.NO_DEFAULT_VALUE, nonEmptyStringWithoutControlChars(), Importance.HIGH, NAME_DOC, COMMON_GROUP, ++orderInGroup, Width.MEDIUM, NAME_DISPLAY)
+                .define(CONNECTOR_CLASS_CONFIG, Type.STRING, Importance.HIGH, CONNECTOR_CLASS_DOC, COMMON_GROUP, ++orderInGroup, Width.LONG, CONNECTOR_CLASS_DISPLAY)
+                .define(TASKS_MAX_CONFIG, Type.INT, TASKS_MAX_DEFAULT, atLeast(TASKS_MIN_CONFIG), Importance.HIGH, TASKS_MAX_DOC, COMMON_GROUP, ++orderInGroup, Width.SHORT, TASK_MAX_DISPLAY)
+                .define(KEY_CONVERTER_CLASS_CONFIG, Type.CLASS, null, Importance.LOW, KEY_CONVERTER_CLASS_DOC, COMMON_GROUP, ++orderInGroup, Width.SHORT, KEY_CONVERTER_CLASS_DISPLAY)
+                .define(VALUE_CONVERTER_CLASS_CONFIG, Type.CLASS, null, Importance.LOW, VALUE_CONVERTER_CLASS_DOC, COMMON_GROUP, ++orderInGroup, Width.SHORT, VALUE_CONVERTER_CLASS_DISPLAY)
+                .define(HEADER_CONVERTER_CLASS_CONFIG, Type.CLASS, HEADER_CONVERTER_CLASS_DEFAULT, Importance.LOW, HEADER_CONVERTER_CLASS_DOC, COMMON_GROUP, ++orderInGroup, Width.SHORT, HEADER_CONVERTER_CLASS_DISPLAY)
+                .define(TRANSFORMS_CONFIG, Type.LIST, Collections.emptyList(), aliasValidator("transformation"), Importance.LOW, TRANSFORMS_DOC, TRANSFORMS_GROUP, ++orderInGroup, Width.LONG, TRANSFORMS_DISPLAY)
+                .define(PREDICATES_CONFIG, Type.LIST, Collections.emptyList(), aliasValidator("predicate"), Importance.LOW, PREDICATES_DOC, PREDICATES_GROUP, ++orderInGroup, Width.LONG, PREDICATES_DISPLAY)
+                .define(CONFIG_RELOAD_ACTION_CONFIG, Type.STRING, CONFIG_RELOAD_ACTION_RESTART,
+                        in(CONFIG_RELOAD_ACTION_NONE, CONFIG_RELOAD_ACTION_RESTART), Importance.LOW,
+                        CONFIG_RELOAD_ACTION_DOC, COMMON_GROUP, ++orderInGroup, Width.MEDIUM, CONFIG_RELOAD_ACTION_DISPLAY)
+                .define(ERRORS_RETRY_TIMEOUT_CONFIG, Type.LONG, ERRORS_RETRY_TIMEOUT_DEFAULT, Importance.MEDIUM,
+                        ERRORS_RETRY_TIMEOUT_DOC, ERROR_GROUP, ++orderInErrorGroup, Width.MEDIUM, ERRORS_RETRY_TIMEOUT_DISPLAY)
+                .define(ERRORS_RETRY_MAX_DELAY_CONFIG, Type.LONG, ERRORS_RETRY_MAX_DELAY_DEFAULT, Importance.MEDIUM,
+                        ERRORS_RETRY_MAX_DELAY_DOC, ERROR_GROUP, ++orderInErrorGroup, Width.MEDIUM, ERRORS_RETRY_MAX_DELAY_DISPLAY)
+                .define(ERRORS_TOLERANCE_CONFIG, Type.STRING, ERRORS_TOLERANCE_DEFAULT.value(),
+                        in(ToleranceType.NONE.value(), ToleranceType.ALL.value()), Importance.MEDIUM,
+                        ERRORS_TOLERANCE_DOC, ERROR_GROUP, ++orderInErrorGroup, Width.SHORT, ERRORS_TOLERANCE_DISPLAY)
+                .define(ERRORS_LOG_ENABLE_CONFIG, Type.BOOLEAN, ERRORS_LOG_ENABLE_DEFAULT, Importance.MEDIUM,
+                        ERRORS_LOG_ENABLE_DOC, ERROR_GROUP, ++orderInErrorGroup, Width.SHORT, ERRORS_LOG_ENABLE_DISPLAY)
+                .define(ERRORS_LOG_INCLUDE_MESSAGES_CONFIG, Type.BOOLEAN, ERRORS_LOG_INCLUDE_MESSAGES_DEFAULT, Importance.MEDIUM,
+                        ERRORS_LOG_INCLUDE_MESSAGES_DOC, ERROR_GROUP, ++orderInErrorGroup, Width.SHORT, ERRORS_LOG_INCLUDE_MESSAGES_DISPLAY);
+    }
+
+    private static ConfigDef.CompositeValidator aliasValidator(String kind) {
+        return ConfigDef.CompositeValidator.of(new ConfigDef.NonNullValidator(), new ConfigDef.Validator() {
+            @SuppressWarnings("unchecked")
+            @Override
+            public void ensureValid(String name, Object value) {
+                final List aliases = (List) value;
+                if (aliases.size() > new HashSet<>(aliases).size()) {
+                    throw new ConfigException(name, value, "Duplicate alias provided.");
+                }
+            }
+
+            @Override
+            public String toString() {
+                return "unique " + kind + " aliases";
+            }
+        });
     }
 
     public ConnectorConfig(Plugins plugins) {
-        this(plugins, new HashMap());
+        this(plugins, Collections.emptyMap());
     }
 
     public ConnectorConfig(Plugins plugins, Map props) {
@@ -134,26 +239,61 @@ public Object get(String key) {
         return enrichedConfig.get(key);
     }
 
+    public long errorRetryTimeout() {
+        return getLong(ERRORS_RETRY_TIMEOUT_CONFIG);
+    }
+
+    public long errorMaxDelayInMillis() {
+        return getLong(ERRORS_RETRY_MAX_DELAY_CONFIG);
+    }
+
+    public ToleranceType errorToleranceType() {
+        String tolerance = getString(ERRORS_TOLERANCE_CONFIG);
+        for (ToleranceType type: ToleranceType.values()) {
+            if (type.name().equalsIgnoreCase(tolerance)) {
+                return type;
+            }
+        }
+        return ERRORS_TOLERANCE_DEFAULT;
+    }
+
+    public boolean enableErrorLog() {
+        return getBoolean(ERRORS_LOG_ENABLE_CONFIG);
+    }
+
+    public boolean includeRecordDetailsInErrorLog() {
+        return getBoolean(ERRORS_LOG_INCLUDE_MESSAGES_CONFIG);
+    }
+
     /**
      * Returns the initialized list of {@link Transformation} which are specified in {@link #TRANSFORMS_CONFIG}.
      */
     public > List> transformations() {
         final List transformAliases = getList(TRANSFORMS_CONFIG);
-        if (transformAliases == null || transformAliases.isEmpty()) {
-            return Collections.emptyList();
-        }
 
         final List> transformations = new ArrayList<>(transformAliases.size());
         for (String alias : transformAliases) {
             final String prefix = TRANSFORMS_CONFIG + "." + alias + ".";
-            final Transformation transformation;
+
             try {
-                transformation = getClass(prefix + "type").asSubclass(Transformation.class).newInstance();
+                @SuppressWarnings("unchecked")
+                final Transformation transformation = Utils.newInstance(getClass(prefix + "type"), Transformation.class);
+                Map configs = originalsWithPrefix(prefix);
+                Object predicateAlias = configs.remove(PredicatedTransformation.PREDICATE_CONFIG);
+                Object negate = configs.remove(PredicatedTransformation.NEGATE_CONFIG);
+                transformation.configure(configs);
+                if (predicateAlias != null) {
+                    String predicatePrefix = PREDICATES_PREFIX + predicateAlias + ".";
+                    @SuppressWarnings("unchecked")
+                    Predicate predicate = Utils.newInstance(getClass(predicatePrefix + "type"), Predicate.class);
+                    predicate.configure(originalsWithPrefix(predicatePrefix));
+                    transformations.add(new PredicatedTransformation<>(predicate, negate == null ? false : Boolean.parseBoolean(negate.toString()), transformation));
+                } else {
+                    transformations.add(transformation);
+                }
             } catch (Exception e) {
                 throw new ConnectException(e);
             }
-            transformation.configure(originalsWithPrefix(prefix));
-            transformations.add(transformation);
         }
 
         return transformations;
@@ -164,91 +304,250 @@ public > List> transformations() {
      * 
            
      * {@code requireFullConfig} specifies whether required config values that are missing should cause an exception to be thrown.
      */
+    @SuppressWarnings({"rawtypes", "unchecked"})
     public static ConfigDef enrich(Plugins plugins, ConfigDef baseConfigDef, Map props, boolean requireFullConfig) {
-        Object transformAliases = ConfigDef.parseType(TRANSFORMS_CONFIG, props.get(TRANSFORMS_CONFIG), Type.LIST);
-        if (!(transformAliases instanceof List)) {
-            return baseConfigDef;
-        }
-
         ConfigDef newDef = new ConfigDef(baseConfigDef);
-        LinkedHashSet uniqueTransformAliases = new LinkedHashSet<>((List) transformAliases);
-        for (Object o : uniqueTransformAliases) {
-            if (!(o instanceof String)) {
-                throw new ConfigException("Item in " + TRANSFORMS_CONFIG + " property is not of "
-                        + "type String");
+        new EnrichablePlugin>("Transformation", TRANSFORMS_CONFIG, TRANSFORMS_GROUP, (Class) Transformation.class,
+                props, requireFullConfig) {
+            @SuppressWarnings("rawtypes")
+            @Override
+            protected Set>> plugins() {
+                return (Set) plugins.transformations();
             }
-            String alias = (String) o;
-            final String prefix = TRANSFORMS_CONFIG + "." + alias + ".";
-            final String group = TRANSFORMS_GROUP + ": " + alias;
-            int orderInGroup = 0;
-
-            final String transformationTypeConfig = prefix + "type";
-            final ConfigDef.Validator typeValidator = new ConfigDef.Validator() {
-                @Override
-                public void ensureValid(String name, Object value) {
-                    getConfigDefFromTransformation(transformationTypeConfig, (Class) value);
-                }
-            };
-            newDef.define(transformationTypeConfig, Type.CLASS, ConfigDef.NO_DEFAULT_VALUE, typeValidator, Importance.HIGH,
-                    "Class for the '" + alias + "' transformation.", group, orderInGroup++, Width.LONG, "Transformation type for " + alias,
-                    Collections.emptyList(), new TransformationClassRecommender(plugins));
 
-            final ConfigDef transformationConfigDef;
-            try {
-                final String className = props.get(transformationTypeConfig);
-                final Class cls = (Class) ConfigDef.parseType(transformationTypeConfig, className, Type.CLASS);
-                transformationConfigDef = getConfigDefFromTransformation(transformationTypeConfig, cls);
-            } catch (ConfigException e) {
-                if (requireFullConfig) {
-                    throw e;
-                } else {
-                    continue;
+            @Override
+            protected ConfigDef initialConfigDef() {
+                // All Transformations get these config parameters implicitly
+                return super.initialConfigDef()
+                        .define(PredicatedTransformation.PREDICATE_CONFIG, Type.STRING, "", Importance.MEDIUM,
+                                "The alias of a predicate used to determine whether to apply this transformation.")
+                        .define(PredicatedTransformation.NEGATE_CONFIG, Type.BOOLEAN, false, Importance.MEDIUM,
+                                "Whether the configured predicate should be negated.");
+            }
+
+            @Override
+            protected Stream> configDefsForClass(String typeConfig) {
+                return super.configDefsForClass(typeConfig)
+                    .filter(entry -> {
+                        // The implicit parameters mask any from the transformer with the same name
+                        if (PredicatedTransformation.PREDICATE_CONFIG.equals(entry.getKey())
+                                || PredicatedTransformation.NEGATE_CONFIG.equals(entry.getKey())) {
+                            log.warn("Transformer config {} is masked by implicit config of that name",
+                                    entry.getKey());
+                            return false;
+                        } else {
+                            return true;
+                        }
+                    });
+            }
+
+            @Override
+            protected ConfigDef config(Transformation transformation) {
+                return transformation.config();
+            }
+
+            @Override
+            protected void validateProps(String prefix) {
+                String prefixedNegate = prefix + PredicatedTransformation.NEGATE_CONFIG;
+                String prefixedPredicate = prefix + PredicatedTransformation.PREDICATE_CONFIG;
+                if (props.containsKey(prefixedNegate) &&
+                        !props.containsKey(prefixedPredicate)) {
+                    throw new ConfigException("Config '" + prefixedNegate + "' was provided " +
+                            "but there is no config '" + prefixedPredicate + "' defining a predicate to be negated.");
                 }
             }
+        }.enrich(newDef);
 
-            newDef.embed(prefix, group, orderInGroup, transformationConfigDef);
-        }
+        new EnrichablePlugin>("Predicate", PREDICATES_CONFIG, PREDICATES_GROUP,
+                (Class) Predicate.class, props, requireFullConfig) {
+            @Override
+            protected Set>> plugins() {
+                return (Set) plugins.predicates();
+            }
 
+            @Override
+            protected ConfigDef config(Predicate predicate) {
+                return predicate.config();
+            }
+        }.enrich(newDef);
         return newDef;
     }
 
     /**
-     * Return {@link ConfigDef} from {@code transformationCls}, which is expected to be a non-null {@code Class},
-     * by instantiating it and invoking {@link Transformation#config()}.
+     * An abstraction over "enrichable plugins" ({@link Transformation}s and {@link Predicate}s) used for computing the
+     * contribution to a Connectors ConfigDef.
+     *
+     * This is not entirely elegant because
+     * although they basically use the same "alias prefix" configuration idiom there are some differences.
+     * The abstract method pattern is used to cope with this.
+     * @param  The type of plugin (either {@code Transformation} or {@code Predicate}).
      */
-    static ConfigDef getConfigDefFromTransformation(String key, Class transformationCls) {
-        if (transformationCls == null || !Transformation.class.isAssignableFrom(transformationCls)) {
-            throw new ConfigException(key, String.valueOf(transformationCls), "Not a Transformation");
+    static abstract class EnrichablePlugin {
+
+        private final String aliasKind;
+        private final String aliasConfig;
+        private final String aliasGroup;
+        private final Class baseClass;
+        private final Map props;
+        private final boolean requireFullConfig;
+
+        public EnrichablePlugin(
+                String aliasKind,
+                String aliasConfig, String aliasGroup, Class baseClass,
+                Map props, boolean requireFullConfig) {
+            this.aliasKind = aliasKind;
+            this.aliasConfig = aliasConfig;
+            this.aliasGroup = aliasGroup;
+            this.baseClass = baseClass;
+            this.props = props;
+            this.requireFullConfig = requireFullConfig;
         }
-        try {
-            return (transformationCls.asSubclass(Transformation.class).newInstance()).config();
-        } catch (Exception e) {
-            throw new ConfigException(key, String.valueOf(transformationCls), "Error getting config definition from Transformation: " + e.getMessage());
+
+        /** Add the configs for this alias to the given {@code ConfigDef}. */
+        void enrich(ConfigDef newDef) {
+            Object aliases = ConfigDef.parseType(aliasConfig, props.get(aliasConfig), Type.LIST);
+            if (!(aliases instanceof List)) {
+                return;
+            }
+
+            LinkedHashSet uniqueAliases = new LinkedHashSet<>((List) aliases);
+            for (Object o : uniqueAliases) {
+                if (!(o instanceof String)) {
+                    throw new ConfigException("Item in " + aliasConfig + " property is not of "
+                            + "type String");
+                }
+                String alias = (String) o;
+                final String prefix = aliasConfig + "." + alias + ".";
+                final String group = aliasGroup + ": " + alias;
+                int orderInGroup = 0;
+
+                final String typeConfig = prefix + "type";
+                final ConfigDef.Validator typeValidator = ConfigDef.LambdaValidator.with(
+                    (String name, Object value) -> {
+                        validateProps(prefix);
+                        getConfigDefFromConfigProvidingClass(typeConfig, (Class) value);
+                    },
+                    () -> "valid configs for " + alias + " " + aliasKind.toLowerCase(Locale.ENGLISH));
+                newDef.define(typeConfig, Type.CLASS, ConfigDef.NO_DEFAULT_VALUE, typeValidator, Importance.HIGH,
+                        "Class for the '" + alias + "' " + aliasKind.toLowerCase(Locale.ENGLISH) + ".", group, orderInGroup++, Width.LONG,
+                        baseClass.getSimpleName() + " type for " + alias,
+                        Collections.emptyList(), new ClassRecommender());
+
+                final ConfigDef configDef = populateConfigDef(typeConfig);
+                if (configDef == null) continue;
+                newDef.embed(prefix, group, orderInGroup, configDef);
+            }
         }
-    }
 
-    /**
-     * Recommend bundled transformations.
-     */
-    static final class TransformationClassRecommender implements ConfigDef.Recommender {
-        private final Plugins plugins;
+        /** Subclasses can add extra validation of the {@link #props}. */
+        protected void validateProps(String prefix) { }
+
+        /**
+         * Populates the ConfigDef according to the configs returned from {@code configs()} method of class
+         * named in the {@code ...type} parameter of the {@code props}.
+         */
+        protected ConfigDef populateConfigDef(String typeConfig) {
+            final ConfigDef configDef = initialConfigDef();
+            try {
+                configDefsForClass(typeConfig)
+                        .forEach(entry -> configDef.define(entry.getValue()));
+
+            } catch (ConfigException e) {
+                if (requireFullConfig) {
+                    throw e;
+                } else {
+                    return null;
+                }
+            }
+            return configDef;
+        }
 
-        TransformationClassRecommender(Plugins plugins) {
-            this.plugins = plugins;
+        /**
+         * Return a stream of configs provided by the {@code configs()} method of class
+         * named in the {@code ...type} parameter of the {@code props}.
+         */
+        protected Stream> configDefsForClass(String typeConfig) {
+            final Class cls = (Class) ConfigDef.parseType(typeConfig, props.get(typeConfig), Type.CLASS);
+            return getConfigDefFromConfigProvidingClass(typeConfig, cls)
+                    .configKeys().entrySet().stream();
         }
 
-        @Override
-        public List validValues(String name, Map parsedConfig) {
-            List transformationPlugins = new ArrayList<>();
-            for (PluginDesc plugin : plugins.transformations()) {
-                transformationPlugins.add(plugin.pluginClass());
+        /** Get an initial ConfigDef */
+        protected ConfigDef initialConfigDef() {
+            return new ConfigDef();
+        }
+
+        /**
+         * Return {@link ConfigDef} from {@code cls}, which is expected to be a non-null {@code Class},
+         * by instantiating it and invoking {@link #config(T)}.
+         * @param key
+         * @param cls The subclass of the baseclass.
+         */
+        ConfigDef getConfigDefFromConfigProvidingClass(String key, Class cls) {
+            if (cls == null || !baseClass.isAssignableFrom(cls)) {
+                throw new ConfigException(key, String.valueOf(cls), "Not a " + baseClass.getSimpleName());
+            }
+            if (Modifier.isAbstract(cls.getModifiers())) {
+                String childClassNames = Stream.of(cls.getClasses())
+                        .filter(cls::isAssignableFrom)
+                        .filter(c -> !Modifier.isAbstract(c.getModifiers()))
+                        .filter(c -> Modifier.isPublic(c.getModifiers()))
+                        .map(Class::getName)
+                        .collect(Collectors.joining(", "));
+                String message = childClassNames.trim().isEmpty() ?
+                        aliasKind + " is abstract and cannot be created." :
+                        aliasKind + " is abstract and cannot be created. Did you mean " + childClassNames + "?";
+                throw new ConfigException(key, String.valueOf(cls), message);
             }
-            return Collections.unmodifiableList(transformationPlugins);
+            T transformation;
+            try {
+                transformation = Utils.newInstance(cls, baseClass);
+            } catch (Exception e) {
+                throw new ConfigException(key, String.valueOf(cls), "Error getting config definition from " + baseClass.getSimpleName() + ": " + e.getMessage());
+            }
+            ConfigDef configDef = config(transformation);
+            if (null == configDef) {
+                throw new ConnectException(
+                    String.format(
+                        "%s.config() must return a ConfigDef that is not null.",
+                        cls.getName()
+                    )
+                );
+            }
+            return configDef;
         }
 
-        @Override
-        public boolean visible(String name, Map parsedConfig) {
-            return true;
+        /**
+         * Get the ConfigDef from the given entity.
+         * This is necessary because there's no abstraction across {@link Transformation#config()} and
+         * {@link Predicate#config()}.
+         */
+        protected abstract ConfigDef config(T t);
+
+        /**
+         * The transformation or predicate plugins (as appropriate for T) to be used
+         * for the {@link ClassRecommender}.
+         */
+        protected abstract Set> plugins();
+
+        /**
+         * Recommend bundled transformations or predicates.
+         */
+        final class ClassRecommender implements ConfigDef.Recommender {
+
+            @Override
+            public List validValues(String name, Map parsedConfig) {
+                List result = new ArrayList<>();
+                for (PluginDesc plugin : plugins()) {
+                    result.add(plugin.pluginClass());
+                }
+                return Collections.unmodifiableList(result);
+            }
+
+            @Override
+            public boolean visible(String name, Map parsedConfig) {
+                return true;
+            }
         }
     }
 
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 5dfb808f764b9..c5fefe9562da9 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,10 +17,13 @@
 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.ActiveTopicsInfo;
 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;
 import org.apache.kafka.connect.runtime.rest.entities.TaskInfo;
+import org.apache.kafka.connect.storage.StatusBackingStore;
 import org.apache.kafka.connect.util.Callback;
 import org.apache.kafka.connect.util.ConnectorTaskId;
 
@@ -56,6 +59,8 @@ public interface Herder {
 
     void stop();
 
+    boolean isRunning();
+
     /**
      * Get a list of connectors currently running in this cluster. This is a full list of connectors in the cluster gathered
      * from the current configuration. However, note
@@ -120,8 +125,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, InternalRequestSignature requestSignature);
+
+    /**
+     * Get a list of connectors currently running in this cluster.
+     * @returns A list of connector names
      */
-    void putTaskConfigs(String connName, List> configs, Callback callback);
+    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.
@@ -129,6 +148,28 @@ public interface Herder {
      */
     ConnectorStateInfo connectorStatus(String connName);
 
+    /**
+     * Lookup the set of topics currently used by a connector.
+     *
+     * @param connName name of the connector
+     * @return the set of active topics
+     */
+    ActiveTopicsInfo connectorActiveTopics(String connName);
+
+    /**
+     * Request to asynchronously reset the active topics for the named connector.
+     *
+     * @param connName name of the connector
+     */
+    void resetConnectorActiveTopics(String connName);
+
+    /**
+     * Return a reference to the status backing store used by this herder.
+     *
+     * @return the status backing store used by this herder
+     */
+    StatusBackingStore statusBackingStore();
+
     /**
      * Lookup the status of the a task.
      * @param id id of the task
@@ -138,8 +179,20 @@ public interface Herder {
     /**
      * Validate the provided connector config values against the configuration definition.
      * @param connectorConfig the provided connector config values
+     * @param callback the callback to invoke after validation has completed (successfully or not)
+     */
+    void validateConnectorConfig(Map connectorConfig, Callback callback);
+
+    /**
+     * Validate the provided connector config values against the configuration definition.
+     * @param connectorConfig the provided connector config values
+     * @param callback the callback to invoke after validation has completed (successfully or not)
+     * @param doLog if true log all the connector configurations at INFO level; if false, no connector configurations are logged.
+     *              Note that logging of configuration is not necessary in every endpoint that uses this method.
      */
-    ConfigInfos validateConnectorConfig(Map connectorConfig);
+    default void validateConnectorConfig(Map connectorConfig, Callback callback, boolean doLog) {
+        validateConnectorConfig(connectorConfig, callback);
+    }
 
     /**
      * Restart the task with the given id.
@@ -155,6 +208,15 @@ public interface Herder {
      */
     void restartConnector(String connName, Callback cb);
 
+    /**
+     * Restart the connector.
+     * @param delayMs delay before restart
+     * @param connName name of the connector
+     * @param cb callback to invoke upon completion
+     * @returns The id of the request
+     */
+    HerderRequest restartConnector(long delayMs, String connName, Callback cb);
+
     /**
      * Pause the connector. This call will asynchronously suspend processing by the connector and all
      * of its tasks.
@@ -176,6 +238,17 @@ public interface Herder {
      */
     Plugins plugins();
 
+    /**
+     * 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();
+
+    enum ConfigReloadAction {
+        NONE,
+        RESTART
+    }
+
     class Created {
         private final boolean created;
         private final T result;
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/HerderConnectorContext.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/HerderConnectorContext.java
index 04c7ad194b03d..60092ba502ac0 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/HerderConnectorContext.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/HerderConnectorContext.java
@@ -16,23 +16,35 @@
  */
 package org.apache.kafka.connect.runtime;
 
-import org.apache.kafka.connect.connector.ConnectorContext;
+import org.apache.kafka.connect.errors.ConnectException;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /**
  * ConnectorContext for use with a Herder
  */
-public class HerderConnectorContext implements ConnectorContext {
+public class HerderConnectorContext implements CloseableConnectorContext {
+
+    private static final Logger log = LoggerFactory.getLogger(HerderConnectorContext.class);
 
     private final AbstractHerder herder;
     private final String connectorName;
+    private volatile boolean closed;
 
     public HerderConnectorContext(AbstractHerder herder, String connectorName) {
         this.herder = herder;
         this.connectorName = connectorName;
+        this.closed = false;
     }
 
     @Override
     public void requestTaskReconfiguration() {
+        if (closed) {
+            throw new ConnectException("The request for task reconfiguration has been rejected " 
+                    + "because this instance of the connector '" + connectorName + "' has already " 
+                    + "been shut down.");
+        }
+
         // Local herder runs in memory in this process
         // Distributed herder will forward the request to the leader if needed
         herder.requestTaskReconfiguration(connectorName);
@@ -40,6 +52,18 @@ public void requestTaskReconfiguration() {
 
     @Override
     public void raiseError(Exception e) {
+        if (closed) {
+            log.warn("Connector {} attempted to raise error after shutdown:", connectorName, e);
+            throw new ConnectException("The request to fail the connector has been rejected " 
+                    + "because this instance of the connector '" + connectorName + "' has already " 
+                    + "been shut down.");
+        }
+
         herder.onFailure(connectorName, e);
     }
+
+    @Override
+    public void close() {
+        closed = true;
+    }
 }
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/HerderRequest.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/HerderRequest.java
new file mode 100644
index 0000000000000..627da4df823a3
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/HerderRequest.java
@@ -0,0 +1,21 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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;
+
+public interface HerderRequest {
+    void cancel();
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/InternalSinkRecord.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/InternalSinkRecord.java
new file mode 100644
index 0000000000000..69554ffb306b5
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/InternalSinkRecord.java
@@ -0,0 +1,82 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.kafka.connect.runtime;
+
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.common.record.TimestampType;
+import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.header.Header;
+import org.apache.kafka.connect.sink.SinkRecord;
+
+/**
+ * A specialization of {@link SinkRecord} that allows a {@link WorkerSinkTask} to track the
+ * original {@link ConsumerRecord} for each {@link SinkRecord}. It is used internally and not
+ * exposed to connectors.
+ */
+public class InternalSinkRecord extends SinkRecord {
+
+    private final ConsumerRecord originalRecord;
+
+    public InternalSinkRecord(ConsumerRecord originalRecord, SinkRecord record) {
+        super(record.topic(), record.kafkaPartition(), record.keySchema(), record.key(),
+            record.valueSchema(), record.value(), record.kafkaOffset(), record.timestamp(),
+            record.timestampType(), record.headers());
+        this.originalRecord = originalRecord;
+    }
+
+    protected InternalSinkRecord(ConsumerRecord originalRecord, String topic,
+                                 int partition, Schema keySchema, Object key, Schema valueSchema,
+                                 Object value, long kafkaOffset, Long timestamp,
+                                 TimestampType timestampType, Iterable
             headers) {
+        super(topic, partition, keySchema, key, valueSchema, value, kafkaOffset, timestamp, timestampType, headers);
+        this.originalRecord = originalRecord;
+    }
+
+    @Override
+    public SinkRecord newRecord(String topic, Integer kafkaPartition, Schema keySchema, Object key,
+                                Schema valueSchema, Object value, Long timestamp,
+                                Iterable
             headers) {
+        return new InternalSinkRecord(originalRecord, topic, kafkaPartition, keySchema, key,
+            valueSchema, value, kafkaOffset(), timestamp, timestampType(), headers());
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        return super.equals(o);
+    }
+
+    @Override
+    public int hashCode() {
+        return super.hashCode();
+    }
+
+    @Override
+    public String toString() {
+        return super.toString();
+    }
+
+    /**
+    * Return the original consumer record that this sink record represents.
+    *
+    * @return the original consumer record; never null
+    */
+    public ConsumerRecord originalRecord() {
+        return originalRecord;
+    }
+}
+
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/PredicatedTransformation.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/PredicatedTransformation.java
new file mode 100644
index 0000000000000..d61772f9f890c
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/PredicatedTransformation.java
@@ -0,0 +1,81 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.util.Map;
+
+import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.common.utils.Utils;
+import org.apache.kafka.connect.connector.ConnectRecord;
+import org.apache.kafka.connect.errors.ConnectException;
+import org.apache.kafka.connect.transforms.Transformation;
+import org.apache.kafka.connect.transforms.predicates.Predicate;
+
+/**
+ * Decorator for a {@link Transformation} which applies the delegate only when a
+ * {@link Predicate} is true (or false, according to {@code negate}).
+ * @param 
+ */
+class PredicatedTransformation> implements Transformation {
+
+    static final String PREDICATE_CONFIG = "predicate";
+    static final String NEGATE_CONFIG = "negate";
+    Predicate predicate;
+    Transformation delegate;
+    boolean negate;
+
+    PredicatedTransformation(Predicate predicate, boolean negate, Transformation delegate) {
+        this.predicate = predicate;
+        this.negate = negate;
+        this.delegate = delegate;
+    }
+
+    @Override
+    public void configure(Map configs) {
+        throw new ConnectException(PredicatedTransformation.class.getName() + ".configure() " +
+                "should never be called directly.");
+    }
+
+    @Override
+    public R apply(R record) {
+        if (negate ^ predicate.test(record)) {
+            return delegate.apply(record);
+        }
+        return record;
+    }
+
+    @Override
+    public ConfigDef config() {
+        throw new ConnectException(PredicatedTransformation.class.getName() + ".config() " +
+                "should never be called directly.");
+    }
+
+    @Override
+    public void close() {
+        Utils.closeQuietly(delegate, "predicated");
+        Utils.closeQuietly(predicate, "predicate");
+    }
+
+    @Override
+    public String toString() {
+        return "PredicatedTransformation{" +
+                "predicate=" + predicate +
+                ", delegate=" + delegate +
+                ", negate=" + negate +
+                '}';
+    }
+}
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 cf5564c25c245..415d46f6f8694 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
@@ -17,11 +17,18 @@
 package org.apache.kafka.connect.runtime;
 
 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
@@ -34,16 +41,40 @@ public class SinkConnectorConfig extends ConnectorConfig {
     public static final String TOPICS_DEFAULT = "";
     private static final String TOPICS_DISPLAY = "Topics";
 
-    private static final String TOPICS_REGEX_CONFIG = SinkTask.TOPICS_REGEX_CONFIG;
+    public static final String TOPICS_REGEX_CONFIG = SinkTask.TOPICS_REGEX_CONFIG;
     private static final String TOPICS_REGEX_DOC = "Regular expression giving topics to consume. " +
         "Under the hood, the regex is compiled to a java.util.regex.Pattern. " +
         "Only one of " + TOPICS_CONFIG + " or " + TOPICS_REGEX_CONFIG + " should be specified.";
     public static final String TOPICS_REGEX_DEFAULT = "";
     private static final String TOPICS_REGEX_DISPLAY = "Topics regex";
 
+    public static final String DLQ_PREFIX = "errors.deadletterqueue.";
+
+    public static final String DLQ_TOPIC_NAME_CONFIG = DLQ_PREFIX + "topic.name";
+    public static final String DLQ_TOPIC_NAME_DOC = "The name of the topic to be used as the dead letter queue (DLQ) for messages that " +
+        "result in an error when processed by this sink connector, or its transformations or converters. The topic name is blank by default, " +
+        "which means that no messages are to be recorded in the DLQ.";
+    public static final String DLQ_TOPIC_DEFAULT = "";
+    private static final String DLQ_TOPIC_DISPLAY = "Dead Letter Queue Topic Name";
+
+    public static final String DLQ_TOPIC_REPLICATION_FACTOR_CONFIG = DLQ_PREFIX + "topic.replication.factor";
+    private static final String DLQ_TOPIC_REPLICATION_FACTOR_CONFIG_DOC = "Replication factor used to create the dead letter queue topic when it doesn't already exist.";
+    public static final short DLQ_TOPIC_REPLICATION_FACTOR_CONFIG_DEFAULT = 3;
+    private static final String DLQ_TOPIC_REPLICATION_FACTOR_CONFIG_DISPLAY = "Dead Letter Queue Topic Replication Factor";
+
+    public static final String DLQ_CONTEXT_HEADERS_ENABLE_CONFIG = DLQ_PREFIX + "context.headers.enable";
+    public static final boolean DLQ_CONTEXT_HEADERS_ENABLE_DEFAULT = false;
+    public static final String DLQ_CONTEXT_HEADERS_ENABLE_DOC = "If true, add headers containing error context to the messages " +
+            "written to the dead letter queue. To avoid clashing with headers from the original record, all error context header " +
+            "keys, all error context header keys will start with __connect.errors.";
+    private static final String DLQ_CONTEXT_HEADERS_ENABLE_DISPLAY = "Enable Error Context Headers";
+
     static ConfigDef config = ConnectorConfig.configDef()
         .define(TOPICS_CONFIG, ConfigDef.Type.LIST, TOPICS_DEFAULT, ConfigDef.Importance.HIGH, TOPICS_DOC, COMMON_GROUP, 4, ConfigDef.Width.LONG, TOPICS_DISPLAY)
-        .define(TOPICS_REGEX_CONFIG, ConfigDef.Type.STRING, TOPICS_REGEX_DEFAULT, new RegexValidator(), ConfigDef.Importance.HIGH, TOPICS_REGEX_DOC, COMMON_GROUP, 4, ConfigDef.Width.LONG, TOPICS_REGEX_DISPLAY);
+        .define(TOPICS_REGEX_CONFIG, ConfigDef.Type.STRING, TOPICS_REGEX_DEFAULT, new RegexValidator(), ConfigDef.Importance.HIGH, TOPICS_REGEX_DOC, COMMON_GROUP, 4, ConfigDef.Width.LONG, TOPICS_REGEX_DISPLAY)
+        .define(DLQ_TOPIC_NAME_CONFIG, ConfigDef.Type.STRING, DLQ_TOPIC_DEFAULT, Importance.MEDIUM, DLQ_TOPIC_NAME_DOC, ERROR_GROUP, 6, ConfigDef.Width.MEDIUM, DLQ_TOPIC_DISPLAY)
+        .define(DLQ_TOPIC_REPLICATION_FACTOR_CONFIG, ConfigDef.Type.SHORT, DLQ_TOPIC_REPLICATION_FACTOR_CONFIG_DEFAULT, Importance.MEDIUM, DLQ_TOPIC_REPLICATION_FACTOR_CONFIG_DOC, ERROR_GROUP, 7, ConfigDef.Width.MEDIUM, DLQ_TOPIC_REPLICATION_FACTOR_CONFIG_DISPLAY)
+        .define(DLQ_CONTEXT_HEADERS_ENABLE_CONFIG, ConfigDef.Type.BOOLEAN, DLQ_CONTEXT_HEADERS_ENABLE_DEFAULT, Importance.MEDIUM, DLQ_CONTEXT_HEADERS_ENABLE_DOC, ERROR_GROUP, 8, ConfigDef.Width.MEDIUM, DLQ_CONTEXT_HEADERS_ENABLE_DISPLAY);
 
     public static ConfigDef configDef() {
         return config;
@@ -52,4 +83,92 @@ public static ConfigDef configDef() {
     public SinkConnectorConfig(Plugins plugins, Map props) {
         super(plugins, config, props);
     }
+
+    /**
+     * Throw an exception if the passed-in properties do not constitute a valid sink.
+     * @param props sink configuration properties
+     */
+    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 +
+                " are mutually exclusive options, but both are set.");
+        }
+
+        if (!hasTopicsConfig && !hasTopicsRegexConfig) {
+            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) {
+        String topicsStr = props.get(TOPICS_CONFIG);
+        return topicsStr != null && !topicsStr.trim().isEmpty();
+    }
+
+    public static boolean hasTopicsRegexConfig(Map props) {
+        String topicsRegexStr = props.get(TOPICS_REGEX_CONFIG);
+        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);
+    }
+
+    public short dlqTopicReplicationFactor() {
+        return getShort(DLQ_TOPIC_REPLICATION_FACTOR_CONFIG);
+    }
+
+    public boolean isDlqContextHeadersEnabled() {
+        return getBoolean(DLQ_CONTEXT_HEADERS_ENABLE_CONFIG);
+    }
+
+    public boolean enableErrantRecordReporter() {
+        String dqlTopic = dlqTopicName();
+        return !dqlTopic.isEmpty() || enableErrorLog();
+    }
+
+    public static void main(String[] args) {
+        System.out.println(config.toHtml(4, config -> "sinkconnectorconfigs_" + config));
+    }
 }
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 ab8fd01c035d3..7cf5d67899e61 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
@@ -16,20 +16,171 @@
  */
 package org.apache.kafka.connect.runtime;
 
+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.connect.runtime.isolation.Plugins;
 
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
 import java.util.Map;
+import java.util.stream.Collectors;
+
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.DEFAULT_TOPIC_CREATION_GROUP;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.DEFAULT_TOPIC_CREATION_PREFIX;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.EXCLUDE_REGEX_CONFIG;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.INCLUDE_REGEX_CONFIG;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.PARTITIONS_CONFIG;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.REPLICATION_FACTOR_CONFIG;
 
 public class SourceConnectorConfig extends ConnectorConfig {
 
-    private static ConfigDef config = ConnectorConfig.configDef();
+    protected static final String TOPIC_CREATION_GROUP = "Topic Creation";
+
+    public static final String TOPIC_CREATION_PREFIX = "topic.creation.";
+
+    public static final String TOPIC_CREATION_GROUPS_CONFIG = TOPIC_CREATION_PREFIX + "groups";
+    private static final String TOPIC_CREATION_GROUPS_DOC = "Groups of configurations for topics "
+            + "created by source connectors";
+    private static final String TOPIC_CREATION_GROUPS_DISPLAY = "Topic Creation Groups";
+
+    private static class EnrichedSourceConnectorConfig extends ConnectorConfig {
+        EnrichedSourceConnectorConfig(Plugins plugins, ConfigDef configDef, Map props) {
+            super(plugins, configDef, props);
+        }
+
+        @Override
+        public Object get(String key) {
+            return super.get(key);
+        }
+    }
+
+    private static ConfigDef config = SourceConnectorConfig.configDef();
+    private final EnrichedSourceConnectorConfig enrichedSourceConfig;
 
     public static ConfigDef configDef() {
-        return config;
+        int orderInGroup = 0;
+        return new ConfigDef(ConnectorConfig.configDef())
+                .define(TOPIC_CREATION_GROUPS_CONFIG, ConfigDef.Type.LIST, Collections.emptyList(),
+                        ConfigDef.CompositeValidator.of(new ConfigDef.NonNullValidator(), ConfigDef.LambdaValidator.with(
+                            (name, value) -> {
+                                List groupAliases = (List) value;
+                                if (groupAliases.size() > new HashSet<>(groupAliases).size()) {
+                                    throw new ConfigException(name, value, "Duplicate alias provided.");
+                                }
+                            },
+                            () -> "unique topic creation groups")),
+                        ConfigDef.Importance.LOW, TOPIC_CREATION_GROUPS_DOC, TOPIC_CREATION_GROUP,
+                        ++orderInGroup, ConfigDef.Width.LONG, TOPIC_CREATION_GROUPS_DISPLAY);
     }
 
-    public SourceConnectorConfig(Plugins plugins, Map props) {
+    public static ConfigDef embedDefaultGroup(ConfigDef baseConfigDef) {
+        String defaultGroup = "default";
+        ConfigDef newDefaultDef = new ConfigDef(baseConfigDef);
+        newDefaultDef.embed(DEFAULT_TOPIC_CREATION_PREFIX, defaultGroup, 0, TopicCreationConfig.defaultGroupConfigDef());
+        return newDefaultDef;
+    }
+
+    /**
+     * Returns an enriched {@link ConfigDef} building upon the {@code ConfigDef}, using the current configuration specified in {@code props} as an input.
+     *
+     * @param baseConfigDef the base configuration definition to be enriched
+     * @param props the non parsed configuration properties
+     * @return the enriched configuration definition
+     */
+    public static ConfigDef enrich(ConfigDef baseConfigDef, Map props, AbstractConfig defaultGroupConfig) {
+        List topicCreationGroups = new ArrayList<>();
+        Object aliases = ConfigDef.parseType(TOPIC_CREATION_GROUPS_CONFIG, props.get(TOPIC_CREATION_GROUPS_CONFIG), ConfigDef.Type.LIST);
+        if (aliases instanceof List) {
+            topicCreationGroups.addAll((List) aliases);
+        }
+
+        ConfigDef newDef = new ConfigDef(baseConfigDef);
+        String defaultGroupPrefix = TOPIC_CREATION_PREFIX + DEFAULT_TOPIC_CREATION_GROUP + ".";
+        short defaultGroupReplicationFactor = defaultGroupConfig.getShort(defaultGroupPrefix + REPLICATION_FACTOR_CONFIG);
+        int defaultGroupPartitions = defaultGroupConfig.getInt(defaultGroupPrefix + PARTITIONS_CONFIG);
+        topicCreationGroups.stream().distinct().forEach(group -> {
+            if (!(group instanceof String)) {
+                throw new ConfigException("Item in " + TOPIC_CREATION_GROUPS_CONFIG + " property is not of type String");
+            }
+            String alias = (String) group;
+            String prefix = TOPIC_CREATION_PREFIX + alias + ".";
+            String configGroup = TOPIC_CREATION_GROUP + ": " + alias;
+            newDef.embed(prefix, configGroup, 0,
+                    TopicCreationConfig.configDef(configGroup, defaultGroupReplicationFactor, defaultGroupPartitions));
+        });
+        return newDef;
+    }
+
+    public SourceConnectorConfig(Plugins plugins, Map props, boolean createTopics) {
         super(plugins, config, props);
+        if (createTopics && props.entrySet().stream().anyMatch(e -> e.getKey().startsWith(TOPIC_CREATION_PREFIX))) {
+            ConfigDef defaultConfigDef = embedDefaultGroup(config);
+            // This config is only used to set default values for partitions and replication
+            // factor from the default group and otherwise it remains unused
+            AbstractConfig defaultGroup = new AbstractConfig(defaultConfigDef, props, false);
+
+            // If the user has added regex of include or exclude patterns in the default group,
+            // they should be ignored.
+            Map propsWithoutRegexForDefaultGroup = new HashMap<>(props);
+            propsWithoutRegexForDefaultGroup.entrySet()
+                    .removeIf(e -> e.getKey().equals(DEFAULT_TOPIC_CREATION_PREFIX + INCLUDE_REGEX_CONFIG)
+                            || e.getKey().equals(DEFAULT_TOPIC_CREATION_PREFIX + EXCLUDE_REGEX_CONFIG));
+            enrichedSourceConfig = new EnrichedSourceConnectorConfig(plugins,
+                    enrich(defaultConfigDef, props, defaultGroup),
+                    propsWithoutRegexForDefaultGroup);
+        } else {
+            enrichedSourceConfig = null;
+        }
+    }
+
+    @Override
+    public Object get(String key) {
+        return enrichedSourceConfig != null ? enrichedSourceConfig.get(key) : super.get(key);
+    }
+
+    /**
+     * Returns whether this configuration uses topic creation properties.
+     *
+     * @return true if the configuration should be validated and used for topic creation; false otherwise
+     */
+    public boolean usesTopicCreation() {
+        return enrichedSourceConfig != null;
+    }
+
+    public List topicCreationInclude(String group) {
+        return getList(TOPIC_CREATION_PREFIX + group + '.' + INCLUDE_REGEX_CONFIG);
+    }
+
+    public List topicCreationExclude(String group) {
+        return getList(TOPIC_CREATION_PREFIX + group + '.' + EXCLUDE_REGEX_CONFIG);
+    }
+
+    public Short topicCreationReplicationFactor(String group) {
+        return getShort(TOPIC_CREATION_PREFIX + group + '.' + REPLICATION_FACTOR_CONFIG);
+    }
+
+    public Integer topicCreationPartitions(String group) {
+        return getInt(TOPIC_CREATION_PREFIX + group + '.' + PARTITIONS_CONFIG);
+    }
+
+    public Map topicCreationOtherConfigs(String group) {
+        if (enrichedSourceConfig == null) {
+            return Collections.emptyMap();
+        }
+        return enrichedSourceConfig.originalsWithPrefix(TOPIC_CREATION_PREFIX + group + '.').entrySet().stream()
+                .filter(e -> {
+                    String key = e.getKey();
+                    return !(INCLUDE_REGEX_CONFIG.equals(key) || EXCLUDE_REGEX_CONFIG.equals(key)
+                            || REPLICATION_FACTOR_CONFIG.equals(key) || PARTITIONS_CONFIG.equals(key));
+                })
+                .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+    }
+
+    public static void main(String[] args) {
+        System.out.println(config.toHtml(4, config -> "sourceconnectorconfigs_" + config));
     }
 }
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..012c7074ad900 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,8 @@
 
 import org.apache.kafka.connect.errors.ConnectException;
 import org.apache.kafka.connect.util.ConnectorTaskId;
+import org.apache.kafka.connect.util.LoggingContext;
+import org.apache.kafka.common.utils.ThreadUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -59,7 +61,8 @@ class SourceTaskOffsetCommitter {
     }
 
     public SourceTaskOffsetCommitter(WorkerConfig config) {
-        this(config, Executors.newSingleThreadScheduledExecutor(),
+        this(config, Executors.newSingleThreadScheduledExecutor(ThreadUtils.createThreadFactory(
+                SourceTaskOffsetCommitter.class.getSimpleName() + "-%d", false)),
                 new ConcurrentHashMap>());
     }
 
@@ -79,7 +82,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 +95,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/TopicCreationConfig.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/TopicCreationConfig.java
new file mode 100644
index 0000000000000..e5c5f15704e0b
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/TopicCreationConfig.java
@@ -0,0 +1,142 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.common.config.ConfigDef;
+import org.apache.kafka.common.config.ConfigException;
+import org.apache.kafka.connect.util.TopicAdmin;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.regex.Pattern;
+import java.util.regex.PatternSyntaxException;
+
+public class TopicCreationConfig {
+
+    public static final String DEFAULT_TOPIC_CREATION_PREFIX = "topic.creation.default.";
+    public static final String DEFAULT_TOPIC_CREATION_GROUP = "default";
+
+    public static final String INCLUDE_REGEX_CONFIG = "include";
+    private static final String INCLUDE_REGEX_DOC = "A list of regular expression literals "
+            + "used to match the topic names used by the source connector. This list is used "
+            + "to include topics that should be created using the topic settings defined by this group.";
+
+    public static final String EXCLUDE_REGEX_CONFIG = "exclude";
+    private static final String EXCLUDE_REGEX_DOC = "A list of regular expression literals "
+            + "used to match the topic names used by the source connector. This list is used "
+            + "to exclude topics from being created with the topic settings defined by this group. "
+            + "Note that exclusion rules have precedent and override any inclusion rules for the topics.";
+
+    public static final String REPLICATION_FACTOR_CONFIG = "replication.factor";
+    private static final String REPLICATION_FACTOR_DOC = "The replication factor for new topics "
+            + "created for this connector using this group. This value may be -1 to use the broker's"
+            + "default replication factor, or may be a positive number not larger than the number of "
+            + "brokers in the Kafka cluster. A value larger than the number of brokers in the Kafka cluster "
+            + "will result in an error when the new topic is created. For the default group this configuration "
+            + "is required. For any other group defined in topic.creation.groups this config is "
+            + "optional and if it's missing it gets the value of the default group";
+
+    public static final String PARTITIONS_CONFIG = "partitions";
+    private static final String PARTITIONS_DOC = "The number of partitions new topics created for "
+            + "this connector. This value may be -1 to use the broker's default number of partitions, "
+            + "or a positive number representing the desired number of partitions. "
+            + "For the default group this configuration is required. For any "
+            + "other group defined in topic.creation.groups this config is optional and if it's "
+            + "missing it gets the value of the default group";
+
+    public static final ConfigDef.Validator REPLICATION_FACTOR_VALIDATOR = ConfigDef.LambdaValidator.with(
+        (name, value) -> validateReplicationFactor(name, (short) value),
+        () -> "Positive number not larger than the number of brokers in the Kafka cluster, or -1 to use the broker's default"
+    );
+    public static final ConfigDef.Validator PARTITIONS_VALIDATOR = ConfigDef.LambdaValidator.with(
+        (name, value) -> validatePartitions(name, (int) value),
+        () -> "Positive number, or -1 to use the broker's default"
+    );
+    @SuppressWarnings("unchecked")
+    public static final ConfigDef.Validator REGEX_VALIDATOR = ConfigDef.LambdaValidator.with(
+        (name, value) -> {
+            try {
+                ((List) value).forEach(Pattern::compile);
+            } catch (PatternSyntaxException e) {
+                throw new ConfigException(name, value,
+                        "Syntax error in regular expression: " + e.getMessage());
+            }
+        },
+        () -> "Positive number, or -1 to use the broker's default"
+    );
+
+    private static void validatePartitions(String configName, int factor) {
+        if (factor != TopicAdmin.NO_PARTITIONS && factor < 1) {
+            throw new ConfigException(configName, factor,
+                    "Number of partitions must be positive, or -1 to use the broker's default");
+        }
+    }
+
+    private static void validateReplicationFactor(String configName, short factor) {
+        if (factor != TopicAdmin.NO_REPLICATION_FACTOR && factor < 1) {
+            throw new ConfigException(configName, factor,
+                    "Replication factor must be positive and not larger than the number of brokers in the Kafka cluster, or -1 to use the broker's default");
+        }
+    }
+
+    public static ConfigDef configDef(String group, short defaultReplicationFactor, int defaultParitionCount) {
+        int orderInGroup = 0;
+        ConfigDef configDef = new ConfigDef();
+        configDef
+                .define(INCLUDE_REGEX_CONFIG, ConfigDef.Type.LIST, Collections.emptyList(),
+                        REGEX_VALIDATOR, ConfigDef.Importance.LOW,
+                        INCLUDE_REGEX_DOC, group, ++orderInGroup, ConfigDef.Width.LONG,
+                        "Inclusion Topic Pattern for " + group)
+                .define(EXCLUDE_REGEX_CONFIG, ConfigDef.Type.LIST, Collections.emptyList(),
+                        REGEX_VALIDATOR, ConfigDef.Importance.LOW,
+                        EXCLUDE_REGEX_DOC, group, ++orderInGroup, ConfigDef.Width.LONG,
+                        "Exclusion Topic Pattern for " + group)
+                .define(REPLICATION_FACTOR_CONFIG, ConfigDef.Type.SHORT,
+                        defaultReplicationFactor, REPLICATION_FACTOR_VALIDATOR,
+                        ConfigDef.Importance.LOW, REPLICATION_FACTOR_DOC, group, ++orderInGroup,
+                        ConfigDef.Width.LONG, "Replication Factor for Topics in " + group)
+                .define(PARTITIONS_CONFIG, ConfigDef.Type.INT,
+                        defaultParitionCount, PARTITIONS_VALIDATOR,
+                        ConfigDef.Importance.LOW, PARTITIONS_DOC, group, ++orderInGroup,
+                        ConfigDef.Width.LONG, "Partition Count for Topics in " + group);
+        return configDef;
+    }
+
+    public static ConfigDef defaultGroupConfigDef() {
+        int orderInGroup = 0;
+        ConfigDef configDef = new ConfigDef();
+        configDef
+                .define(INCLUDE_REGEX_CONFIG, ConfigDef.Type.LIST, ".*",
+                        new ConfigDef.NonNullValidator(), ConfigDef.Importance.LOW,
+                        INCLUDE_REGEX_DOC, DEFAULT_TOPIC_CREATION_GROUP, ++orderInGroup, ConfigDef.Width.LONG,
+                        "Inclusion Topic Pattern for " + DEFAULT_TOPIC_CREATION_GROUP)
+                .define(EXCLUDE_REGEX_CONFIG, ConfigDef.Type.LIST, Collections.emptyList(),
+                        new ConfigDef.NonNullValidator(), ConfigDef.Importance.LOW,
+                        EXCLUDE_REGEX_DOC, DEFAULT_TOPIC_CREATION_GROUP, ++orderInGroup, ConfigDef.Width.LONG,
+                        "Exclusion Topic Pattern for " + DEFAULT_TOPIC_CREATION_GROUP)
+                .define(REPLICATION_FACTOR_CONFIG, ConfigDef.Type.SHORT,
+                        ConfigDef.NO_DEFAULT_VALUE, REPLICATION_FACTOR_VALIDATOR,
+                        ConfigDef.Importance.LOW, REPLICATION_FACTOR_DOC, DEFAULT_TOPIC_CREATION_GROUP, ++orderInGroup,
+                        ConfigDef.Width.LONG, "Replication Factor for Topics in " + DEFAULT_TOPIC_CREATION_GROUP)
+                .define(PARTITIONS_CONFIG, ConfigDef.Type.INT,
+                        ConfigDef.NO_DEFAULT_VALUE, PARTITIONS_VALIDATOR,
+                        ConfigDef.Importance.LOW, PARTITIONS_DOC, DEFAULT_TOPIC_CREATION_GROUP, ++orderInGroup,
+                        ConfigDef.Width.LONG, "Partition Count for Topics in " + DEFAULT_TOPIC_CREATION_GROUP);
+        return configDef;
+    }
+
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/TopicStatus.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/TopicStatus.java
new file mode 100644
index 0000000000000..16dcd80b43f84
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/TopicStatus.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.connect.runtime;
+
+import org.apache.kafka.connect.util.ConnectorTaskId;
+
+import java.util.Objects;
+
+/**
+ * Represents the metadata that is stored as the value of the record that is stored in the
+ * {@link org.apache.kafka.connect.storage.StatusBackingStore#put(TopicStatus)},
+ */
+public class TopicStatus {
+    private final String topic;
+    private final String connector;
+    private final int task;
+    private final long discoverTimestamp;
+
+    public TopicStatus(String topic, ConnectorTaskId task, long discoverTimestamp) {
+        this(topic, task.connector(), task.task(), discoverTimestamp);
+    }
+
+    public TopicStatus(String topic, String connector, int task, long discoverTimestamp) {
+        this.topic = Objects.requireNonNull(topic);
+        this.connector = Objects.requireNonNull(connector);
+        this.task = task;
+        this.discoverTimestamp = discoverTimestamp;
+    }
+
+    /**
+     * Get the name of the topic.
+     *
+     * @return the topic name; never null
+     */
+    public String topic() {
+        return topic;
+    }
+
+    /**
+     * Get the name of the connector.
+     *
+     * @return the connector name; never null
+     */
+    public String connector() {
+        return connector;
+    }
+
+    /**
+     * Get the ID of the task that stored the topic status.
+     *
+     * @return the task ID
+     */
+    public int task() {
+        return task;
+    }
+
+    /**
+     * Get a timestamp that represents when this topic was discovered as being actively used by
+     * this connector.
+     *
+     * @return the discovery timestamp
+     */
+    public long discoverTimestamp() {
+        return discoverTimestamp;
+    }
+
+    @Override
+    public String toString() {
+        return "TopicStatus{" +
+                "topic='" + topic + '\'' +
+                ", connector='" + connector + '\'' +
+                ", task=" + task +
+                ", discoverTimestamp=" + discoverTimestamp +
+                '}';
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (!(o instanceof TopicStatus)) {
+            return false;
+        }
+        TopicStatus that = (TopicStatus) o;
+        return task == that.task &&
+                discoverTimestamp == that.discoverTimestamp &&
+                topic.equals(that.topic) &&
+                connector.equals(that.connector);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(topic, connector, task, discoverTimestamp);
+    }
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/TransformationChain.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/TransformationChain.java
index e1d8b1f2e856a..5027cb534d79f 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/TransformationChain.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/TransformationChain.java
@@ -17,31 +17,45 @@
 package org.apache.kafka.connect.runtime;
 
 import org.apache.kafka.connect.connector.ConnectRecord;
+import org.apache.kafka.connect.runtime.errors.RetryWithToleranceOperator;
+import org.apache.kafka.connect.runtime.errors.Stage;
 import org.apache.kafka.connect.transforms.Transformation;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
-import java.util.Collections;
 import java.util.List;
 import java.util.Objects;
+import java.util.StringJoiner;
 
-public class TransformationChain> {
+public class TransformationChain> implements AutoCloseable {
+    private static final Logger log = LoggerFactory.getLogger(TransformationChain.class);
 
     private final List> transformations;
+    private final RetryWithToleranceOperator retryWithToleranceOperator;
 
-    public TransformationChain(List> transformations) {
+    public TransformationChain(List> transformations, RetryWithToleranceOperator retryWithToleranceOperator) {
         this.transformations = transformations;
+        this.retryWithToleranceOperator = retryWithToleranceOperator;
     }
 
     public R apply(R record) {
         if (transformations.isEmpty()) return record;
 
-        for (Transformation transformation : transformations) {
-            record = transformation.apply(record);
+        for (final Transformation transformation : transformations) {
+            final R current = record;
+
+            log.trace("Applying transformation {} to {}",
+                transformation.getClass().getName(), record);
+            // execute the operation
+            record = retryWithToleranceOperator.execute(() -> transformation.apply(current), Stage.TRANSFORMATION, transformation.getClass());
+
             if (record == null) break;
         }
 
         return record;
     }
 
+    @Override
     public void close() {
         for (Transformation transformation : transformations) {
             transformation.close();
@@ -61,8 +75,11 @@ public int hashCode() {
         return Objects.hash(transformations);
     }
 
-    public static > TransformationChain noOp() {
-        return new TransformationChain(Collections.>emptyList());
+    public String toString() {
+        StringJoiner chain = new StringJoiner(", ", getClass().getName() + "{", "}");
+        for (Transformation transformation : transformations) {
+            chain.add(transformation.getClass().getName());
+        }
+        return chain.toString();
     }
-
 }
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 992825ca23bc3..45a93dfc4650b 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,31 +16,54 @@
  */
 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;
+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.WorkerErrantRecordReporter;
 import org.apache.kafka.connect.runtime.isolation.Plugins;
+import org.apache.kafka.connect.runtime.isolation.Plugins.ClassLoaderUsage;
 import org.apache.kafka.connect.sink.SinkRecord;
 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.ConnectUtils;
+import org.apache.kafka.connect.util.Callback;
 import org.apache.kafka.connect.util.ConnectorTaskId;
+import org.apache.kafka.connect.util.LoggingContext;
+import org.apache.kafka.connect.util.SinkUtils;
+import org.apache.kafka.connect.util.TopicAdmin;
+import org.apache.kafka.connect.util.TopicCreationGroup;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -55,7 +78,8 @@
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
-
+import java.util.concurrent.TimeUnit;
+import java.util.stream.Collectors;
 
 /**
  * 
            
@@ -67,72 +91,101 @@
  * 
            
  */
 public class Worker {
+
+    public static final long CONNECTOR_GRACEFUL_SHUTDOWN_TIMEOUT_MS = TimeUnit.SECONDS.toMillis(5);
+
     private static final Logger log = LoggerFactory.getLogger(Worker.class);
 
+    protected Herder herder;
     private final ExecutorService executor;
     private final Time time;
     private final String workerId;
+    //kafka cluster id
+    private final String kafkaClusterId;
     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;
     private final OffsetBackingStore offsetBackingStore;
-    private final Map producerProps;
 
     private final ConcurrentMap connectors = new ConcurrentHashMap<>();
     private final ConcurrentMap tasks = new ConcurrentHashMap<>();
     private SourceTaskOffsetCommitter sourceTaskOffsetCommitter;
+    private final WorkerConfigTransformer workerConfigTransformer;
+    private final ConnectorClientConfigOverridePolicy connectorClientConfigOverridePolicy;
 
     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.kafkaClusterId = ConnectUtils.lookupKafkaClusterId(config);
+        this.metrics = new ConnectMetrics(workerId, config, time, kafkaClusterId);
+        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.
         this.internalKeyConverter = plugins.newConverter(
-                config.getClass(WorkerConfig.INTERNAL_KEY_CONVERTER_CLASS_CONFIG).getName(),
-                config
+                config,
+                WorkerConfig.INTERNAL_KEY_CONVERTER_CLASS_CONFIG,
+                ClassLoaderUsage.PLUGINS
         );
-        this.internalKeyConverter.configure(
-                config.originalsWithPrefix("internal.key.converter."),
-                true);
         this.internalValueConverter = plugins.newConverter(
-                config.getClass(WorkerConfig.INTERNAL_VALUE_CONVERTER_CLASS_CONFIG).getName(),
-                config
-        );
-        this.internalValueConverter.configure(
-                config.originalsWithPrefix("internal.value.converter."),
-                false
+                config,
+                WorkerConfig.INTERNAL_VALUE_CONVERTER_CLASS_CONFIG,
+                ClassLoaderUsage.PLUGINS
         );
 
         this.offsetBackingStore = offsetBackingStore;
         this.offsetBackingStore.configure(config);
 
-        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.
-        producerProps.put(ProducerConfig.REQUEST_TIMEOUT_MS_CONFIG, Integer.toString(Integer.MAX_VALUE));
-        producerProps.put(ProducerConfig.RETRIES_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");
-        // User-specified overrides
-        producerProps.putAll(config.originalsWithPrefix("producer."));
+        this.workerConfigTransformer = initConfigTransformer();
+
+    }
+
+    private WorkerConfigTransformer initConfigTransformer() {
+        final List providerNames = config.getList(WorkerConfig.CONFIG_PROVIDERS_CONFIG);
+        Map providerMap = new HashMap<>();
+        for (String providerName : providerNames) {
+            ConfigProvider configProvider = plugins.newConfigProvider(
+                    config,
+                    WorkerConfig.CONFIG_PROVIDERS_CONFIG + "." + providerName,
+                    ClassLoaderUsage.PLUGINS
+            );
+            providerMap.put(providerName, configProvider);
+        }
+        return new WorkerConfigTransformer(this, providerMap);
+    }
+
+    public WorkerConfigTransformer configTransformer() {
+        return workerConfigTransformer;
+    }
+
+    protected Herder herder() {
+        return herder;
     }
 
     /**
@@ -144,6 +197,8 @@ public void start() {
         offsetBackingStore.start();
         sourceTaskOffsetCommitter = new SourceTaskOffsetCommitter(config);
 
+        connectorStatusMetricsGroup = new ConnectorStatusMetricsGroup(metrics, tasks, herder);
+
         log.info("Worker started");
     }
 
@@ -158,7 +213,7 @@ public void stop() {
 
         if (!connectors.isEmpty()) {
             log.warn("Shutting down connectors {} uncleanly; herder should have shut down connectors before the Worker is stopped", connectors.keySet());
-            stopConnectors();
+            stopAndAwaitConnectors();
         }
 
         if (!tasks.isEmpty()) {
@@ -175,6 +230,9 @@ public void stop() {
         log.info("Worker stopped");
 
         workerMetricsGroup.close();
+        connectorStatusMetricsGroup.close();
+
+        workerConfigTransformer.close();
     }
 
     /**
@@ -185,48 +243,73 @@ public void stop() {
      * @param ctx the connector runtime context.
      * @param statusListener a listener for the runtime status transitions of the connector.
      * @param initialState the initial state of the connector.
-     * @return true if the connector started successfully.
+     * @param onConnectorStateChange invoked when the initial state change of the connector is completed
      */
-    public boolean startConnector(
+    public void startConnector(
             String connName,
             Map connProps,
-            ConnectorContext ctx,
+            CloseableConnectorContext ctx,
             ConnectorStatus.Listener statusListener,
-            TargetState initialState
+            TargetState initialState,
+            Callback onConnectorStateChange
     ) {
-        if (connectors.containsKey(connName))
-            throw new ConnectException("Connector with name " + connName + " already exists");
+        try (LoggingContext loggingContext = LoggingContext.forConnector(connName)) {
+            if (connectors.containsKey(connName)) {
+                onConnectorStateChange.onCompletion(
+                        new ConnectException("Connector with name " + connName + " already exists"),
+                        null);
+                return;
+            }
 
-        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;
-        }
+            final WorkerConnector workerConnector;
+            ClassLoader savedLoader = plugins.currentThreadLoader();
+            try {
+                // By the time we arrive here, CONNECTOR_CLASS_CONFIG has been validated already
+                // Getting this value from the unparsed map will allow us to instantiate the
+                // right config (source or sink)
+                final String connClass = connProps.get(ConnectorConfig.CONNECTOR_CLASS_CONFIG);
+                ClassLoader connectorLoader = plugins.delegatingLoader().connectorLoader(connClass);
+                savedLoader = Plugins.compareAndSwapLoaders(connectorLoader);
+
+                log.info("Creating connector {} of type {}", connName, connClass);
+                final Connector connector = plugins.newConnector(connClass);
+                final ConnectorConfig connConfig = ConnectUtils.isSinkConnector(connector)
+                        ? new SinkConnectorConfig(plugins, connProps)
+                        : new SourceConnectorConfig(plugins, connProps, config.topicCreationEnable());
+
+                final OffsetStorageReader offsetReader = new OffsetStorageReaderImpl(
+                        offsetBackingStore, connName, internalKeyConverter, internalValueConverter);
+                workerConnector = new WorkerConnector(
+                        connName, connector, connConfig, ctx, metrics, statusListener, offsetReader, connectorLoader);
+                log.info("Instantiated connector {} with version {} of type {}", connName, connector.version(), connector.getClass());
+                workerConnector.transitionTo(initialState, onConnectorStateChange);
+                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);
+                onConnectorStateChange.onCompletion(t, null);
+                return;
+            }
 
-        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) {
+                onConnectorStateChange.onCompletion(
+                        new ConnectException("Connector with name " + connName + " already exists"),
+                        null);
+                // Don't need to do any cleanup of the WorkerConnector instance (such as calling
+                // shutdown() on it) here because it hasn't actually started running yet
+                return;
+            }
+
+            executor.submit(workerConnector);
 
-        log.info("Finished creating connector {}", connName);
-        workerMetricsGroup.recordConnectorStartupSuccess();
-        return true;
+            log.info("Finished creating connector {}", connName);
+            workerMetricsGroup.recordConnectorStartupSuccess();
+        }
     }
 
     /**
@@ -243,7 +326,7 @@ public boolean isSinkConnector(String connName) {
 
         ClassLoader savedLoader = plugins.currentThreadLoader();
         try {
-            savedLoader = plugins.compareAndSwapLoaders(workerConnector.connector());
+            savedLoader = Plugins.compareAndSwapLoaders(workerConnector.loader());
             return workerConnector.isSinkConnector();
         } finally {
             Plugins.compareAndSwapLoaders(savedLoader);
@@ -257,72 +340,130 @@ 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(workerConnector.loader());
+                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;
     }
 
-    private void stopConnectors() {
-        // Herder is responsible for stopping connectors. This is an internal method to sequentially
-        // stop connectors that have not explicitly been stopped.
-        for (String connector: connectors.keySet())
-            stopConnector(connector);
-    }
-
     /**
      * Stop a connector managed by this worker.
      *
      * @param connName the connector name.
-     * @return true if the connector belonged to this worker and was successfully stopped.
      */
-    public boolean stopConnector(String connName) {
-        log.info("Stopping connector {}", connName);
+    private void stopConnector(String connName) {
+        try (LoggingContext loggingContext = LoggingContext.forConnector(connName)) {
+            WorkerConnector workerConnector = connectors.get(connName);
+            log.info("Stopping connector {}", connName);
+
+            if (workerConnector == null) {
+                log.warn("Ignoring stop request for unowned connector {}", connName);
+                return;
+            }
 
-        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.loader());
+                workerConnector.shutdown();
+            } finally {
+                Plugins.compareAndSwapLoaders(savedLoader);
+            }
         }
+    }
 
-        ClassLoader savedLoader = plugins.currentThreadLoader();
-        try {
-            savedLoader = plugins.compareAndSwapLoaders(workerConnector.connector());
-            workerConnector.shutdown();
-        } finally {
-            Plugins.compareAndSwapLoaders(savedLoader);
+    private void stopConnectors(Collection ids) {
+        // Herder is responsible for stopping connectors. This is an internal method to sequentially
+        // stop connectors that have not explicitly been stopped.
+        for (String connector: ids)
+            stopConnector(connector);
+    }
+
+    private void awaitStopConnector(String connName, long timeout) {
+        try (LoggingContext loggingContext = LoggingContext.forConnector(connName)) {
+            WorkerConnector connector = connectors.remove(connName);
+            if (connector == null) {
+                log.warn("Ignoring await stop request for non-present connector {}", connName);
+                return;
+            }
+
+            if (!connector.awaitShutdown(timeout)) {
+                log.error("Connector ‘{}’ failed to properly shut down, has become unresponsive, and "
+                        + "may be consuming external resources. Correct the configuration for "
+                        + "this connector or remove the connector. After fixing the connector, it "
+                        + "may be necessary to restart this worker to release any consumed "
+                        + "resources.", connName);
+                connector.cancel();
+            } else {
+                log.debug("Graceful stop of connector {} succeeded.", connName);
+            }
+        }
+    }
+
+    private void awaitStopConnectors(Collection ids) {
+        long now = time.milliseconds();
+        long deadline = now + CONNECTOR_GRACEFUL_SHUTDOWN_TIMEOUT_MS;
+        for (String id : ids) {
+            long remaining = Math.max(0, deadline - time.milliseconds());
+            awaitStopConnector(id, remaining);
         }
+    }
+
+    /**
+     * Stop asynchronously all the worker's connectors and await their termination.
+     */
+    public void stopAndAwaitConnectors() {
+        stopAndAwaitConnectors(new ArrayList<>(connectors.keySet()));
+    }
 
-        log.info("Stopped connector {}", connName);
-        return true;
+    /**
+     * Stop asynchronously a collection of connectors that belong to this worker and await their
+     * termination.
+     *
+     * @param ids the collection of connectors to be stopped.
+     */
+    public void stopAndAwaitConnectors(Collection ids) {
+        stopConnectors(ids);
+        awaitStopConnectors(ids);
+    }
+
+    /**
+     * Stop a connector that belongs to this worker and await its termination.
+     *
+     * @param connName the name of the connector to be stopped.
+     */
+    public void stopAndAwaitConnector(String connName) {
+        stopConnector(connName);
+        awaitStopConnectors(Collections.singletonList(connName));
     }
 
     /**
@@ -357,124 +498,352 @@ public boolean isRunning(String connName) {
      */
     public boolean startTask(
             ConnectorTaskId id,
+            ClusterConfigState configState,
             Map connProps,
             Map taskProps,
             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, and if not found the
-            // plugin class loader delegates loading to the delegating classloader.
-            Converter keyConverter = connConfig.getConfiguredInstance(WorkerConfig.KEY_CONVERTER_CLASS_CONFIG, Converter.class);
-            if (keyConverter != null)
-                keyConverter.configure(connConfig.originalsWithPrefix("key.converter."), true);
-            else {
-                Converter defaultKeyConverter = plugins.newConverter(
-                        config.getClass(WorkerConfig.KEY_CONVERTER_CLASS_CONFIG).getName(),
-                        config
-                );
-                defaultKeyConverter.configure(config.originalsWithPrefix("key.converter."), true);
-                keyConverter = defaultKeyConverter;
-            }
-            Converter valueConverter = connConfig.getConfiguredInstance(WorkerConfig.VALUE_CONVERTER_CLASS_CONFIG, Converter.class);
-            if (valueConverter != null)
-                valueConverter.configure(connConfig.originalsWithPrefix("value.converter."), false);
-            else {
-                Converter defaultValueConverter = plugins.newConverter(
-                        config.getClass(WorkerConfig.VALUE_CONVERTER_CLASS_CONFIG).getName(),
-                        config
-                );
-                defaultValueConverter.configure(config.originalsWithPrefix("value.converter."), false);
-                valueConverter = defaultValueConverter;
-            }
+        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(connConfig, id, task, statusListener, initialState, keyConverter, valueConverter, 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(ConnectorConfig connConfig,
+    private WorkerTask buildWorkerTask(ClusterConfigState configState,
+                                       ConnectorConfig connConfig,
                                        ConnectorTaskId id,
                                        Task task,
                                        TaskStatus.Listener statusListener,
                                        TargetState initialState,
                                        Converter keyConverter,
                                        Converter valueConverter,
+                                       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);
+
         // Decide which type of worker task we need based on the type of task.
         if (task instanceof SourceTask) {
-            TransformationChain transformationChain = new TransformationChain<>(connConfig.transformations());
-            OffsetStorageReader offsetReader = new OffsetStorageReaderImpl(offsetBackingStore, id.connector(),
+            SourceConnectorConfig sourceConfig = new SourceConnectorConfig(plugins,
+                    connConfig.originalsStrings(), config.topicCreationEnable());
+            retryWithToleranceOperator.reporters(sourceTaskReporters(id, sourceConfig, errorHandlingMetrics));
+            TransformationChain transformationChain = new TransformationChain<>(sourceConfig.transformations(), retryWithToleranceOperator);
+            log.info("Initializing: {}", transformationChain);
+            CloseableOffsetStorageReader offsetReader = new OffsetStorageReaderImpl(offsetBackingStore, id.connector(),
                     internalKeyConverter, internalValueConverter);
             OffsetStorageWriter offsetWriter = new OffsetStorageWriter(offsetBackingStore, id.connector(),
                     internalKeyConverter, internalValueConverter);
+            Map producerProps = producerConfigs(id, "connector-producer-" + id, config, sourceConfig, connectorClass,
+                                                                connectorClientConfigOverridePolicy, kafkaClusterId);
             KafkaProducer producer = new KafkaProducer<>(producerProps);
-            return new WorkerSourceTask(id, (SourceTask) task, statusListener, initialState, keyConverter,
-                    valueConverter, transformationChain, producer, offsetReader, offsetWriter, config, metrics, loader, time);
+            TopicAdmin admin;
+            Map topicCreationGroups;
+            if (config.topicCreationEnable() && sourceConfig.usesTopicCreation()) {
+                Map adminProps = adminConfigs(id, "connector-adminclient-" + id, config,
+                        sourceConfig, connectorClass, connectorClientConfigOverridePolicy, kafkaClusterId);
+                admin = new TopicAdmin(adminProps);
+                topicCreationGroups = TopicCreationGroup.configuredGroups(sourceConfig);
+            } else {
+                admin = null;
+                topicCreationGroups = null;
+            }
+
+            // Note we pass the configState as it performs dynamic transformations under the covers
+            return new WorkerSourceTask(id, (SourceTask) task, statusListener, initialState, keyConverter, valueConverter,
+                    headerConverter, transformationChain, producer, admin, topicCreationGroups,
+                    offsetReader, offsetWriter, config, configState, metrics, loader, time, retryWithToleranceOperator, herder.statusBackingStore());
         } else if (task instanceof SinkTask) {
-            TransformationChain transformationChain = new TransformationChain<>(connConfig.transformations());
-            return new WorkerSinkTask(id, (SinkTask) task, statusListener, initialState, config, metrics, keyConverter,
-                    valueConverter, transformationChain, loader, time);
+            TransformationChain transformationChain = new TransformationChain<>(connConfig.transformations(), retryWithToleranceOperator);
+            log.info("Initializing: {}", transformationChain);
+            SinkConnectorConfig sinkConfig = new SinkConnectorConfig(plugins, connConfig.originalsStrings());
+            retryWithToleranceOperator.reporters(sinkTaskReporters(id, sinkConfig, errorHandlingMetrics, connectorClass));
+            WorkerErrantRecordReporter workerErrantRecordReporter = createWorkerErrantRecordReporter(sinkConfig, retryWithToleranceOperator,
+                    keyConverter, valueConverter, headerConverter);
+
+            Map consumerProps = consumerConfigs(id, config, connConfig, connectorClass, connectorClientConfigOverridePolicy, kafkaClusterId);
+            KafkaConsumer consumer = new KafkaConsumer<>(consumerProps);
+
+            return new WorkerSinkTask(id, (SinkTask) task, statusListener, initialState, config, configState, metrics, keyConverter,
+                                      valueConverter, headerConverter, transformationChain, consumer, loader, time,
+                                      retryWithToleranceOperator, workerErrantRecordReporter, herder.statusBackingStore());
         } 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");
         }
     }
 
-    private void stopTask(ConnectorTaskId taskId) {
-        WorkerTask task = tasks.get(taskId);
-        if (task == null) {
-            log.warn("Ignoring stop request for unowned task {}", taskId);
-            return;
+    static Map producerConfigs(ConnectorTaskId id,
+                                               String defaultClientId,
+                                               WorkerConfig config,
+                                               ConnectorConfig connConfig,
+                                               Class  connectorClass,
+                                               ConnectorClientConfigOverridePolicy connectorClientConfigOverridePolicy,
+                                               String clusterId) {
+        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 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."));
+        //add client metrics.context properties
+        ConnectUtils.addMetricsContextProperties(producerProps, config, clusterId);
+
+        // 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,
+                                               ConnectorConfig connConfig,
+                                               Class connectorClass,
+                                               ConnectorClientConfigOverridePolicy connectorClientConfigOverridePolicy,
+                                               String clusterId) {
+        // 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");
+        consumerProps.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
+        consumerProps.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.ByteArrayDeserializer");
+        consumerProps.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.ByteArrayDeserializer");
+
+        consumerProps.putAll(config.originalsWithPrefix("consumer."));
+        //add client metrics.context properties
+        ConnectUtils.addMetricsContextProperties(consumerProps, config, clusterId);
+        // 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,
+                                            String defaultClientId,
+                                            WorkerConfig config,
+                                            ConnectorConfig connConfig,
+                                            Class connectorClass,
+                                            ConnectorClientConfigOverridePolicy connectorClientConfigOverridePolicy,
+                                            String clusterId) {
+        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.put(AdminClientConfig.CLIENT_ID_CONFIG, defaultClientId);
+        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);
+
+        //add client metrics.context properties
+        ConnectUtils.addMetricsContextProperties(adminProps, config, clusterId);
+
+        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;
+    }
 
-        log.info("Stopping task {}", task.id());
-        if (task instanceof WorkerSourceTask)
-            sourceTaskOffsetCommitter.remove(task.id());
+    ErrorHandlingMetrics errorHandlingMetrics(ConnectorTaskId id) {
+        return new ErrorHandlingMetrics(id, metrics);
+    }
 
-        ClassLoader savedLoader = plugins.currentThreadLoader();
-        try {
-            savedLoader = Plugins.compareAndSwapLoaders(task.loader());
-            task.stop();
-        } finally {
-            Plugins.compareAndSwapLoaders(savedLoader);
+    private List sinkTaskReporters(ConnectorTaskId id, SinkConnectorConfig connConfig,
+                                                  ErrorHandlingMetrics errorHandlingMetrics,
+                                                  Class connectorClass) {
+        ArrayList reporters = new ArrayList<>();
+        LogReporter logReporter = new LogReporter(id, connConfig, errorHandlingMetrics);
+        reporters.add(logReporter);
+
+        // check if topic for dead letter queue exists
+        String topic = connConfig.dlqTopicName();
+        if (topic != null && !topic.isEmpty()) {
+            Map producerProps = producerConfigs(id, "connector-dlq-producer-" + id, config, connConfig, connectorClass,
+                                                                connectorClientConfigOverridePolicy, kafkaClusterId);
+            Map adminProps = adminConfigs(id, "connector-dlq-adminclient-", config, connConfig, connectorClass, connectorClientConfigOverridePolicy, kafkaClusterId);
+            DeadLetterQueueReporter reporter = DeadLetterQueueReporter.createAndSetup(adminProps, id, connConfig, producerProps, errorHandlingMetrics);
+
+            reporters.add(reporter);
+        }
+
+        return reporters;
+    }
+
+    private List sourceTaskReporters(ConnectorTaskId id, ConnectorConfig connConfig,
+                                                      ErrorHandlingMetrics errorHandlingMetrics) {
+        List reporters = new ArrayList<>();
+        LogReporter logReporter = new LogReporter(id, connConfig, errorHandlingMetrics);
+        reporters.add(logReporter);
+
+        return reporters;
+    }
+
+    private WorkerErrantRecordReporter createWorkerErrantRecordReporter(
+        SinkConnectorConfig connConfig,
+        RetryWithToleranceOperator retryWithToleranceOperator,
+        Converter keyConverter,
+        Converter valueConverter,
+        HeaderConverter headerConverter
+    ) {
+        // check if errant record reporter topic is configured
+        if (connConfig.enableErrantRecordReporter()) {
+            return new WorkerErrantRecordReporter(retryWithToleranceOperator, keyConverter, valueConverter, headerConverter);
+        }
+        return null;
+    }
+
+    private void stopTask(ConnectorTaskId taskId) {
+        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());
+
+            ClassLoader savedLoader = plugins.currentThreadLoader();
+            try {
+                savedLoader = Plugins.compareAndSwapLoaders(task.loader());
+                task.stop();
+            } finally {
+                Plugins.compareAndSwapLoaders(savedLoader);
+            }
         }
     }
 
@@ -487,15 +856,25 @@ 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();
+            } else {
+                log.debug("Graceful stop of task {} succeeded.", task.id());
+            }
 
-        if (!task.awaitStop(timeout)) {
-            log.error("Graceful stop of task {} failed.", task.id());
-            task.cancel();
+            try {
+                task.removeMetrics();
+            } finally {
+                connectorStatusMetricsGroup.recordTaskRemoved(taskId);
+            }
         }
     }
 
@@ -558,6 +937,16 @@ public String workerId() {
         return workerId;
     }
 
+    /**
+     * Returns whether this worker is configured to allow source connectors to create the topics
+     * that they use with custom configurations, if these topics don't already exist.
+     *
+     * @return true if topic creation by source connectors is allowed; false otherwise
+     */
+    public boolean isTopicCreationEnabled() {
+        return config.topicCreationEnable();
+    }
+
     /**
      * Get the {@link ConnectMetrics} that uses Kafka Metrics and manages the JMX reporter.
      * @return the Connect-specific metrics; never null
@@ -566,47 +955,114 @@ public ConnectMetrics metrics() {
         return metrics;
     }
 
-    public void setTargetState(String connName, TargetState state) {
+    public void setTargetState(String connName, TargetState state, Callback stateChangeCallback) {
         log.info("Setting connector {} state to {}", connName, state);
 
         WorkerConnector workerConnector = connectors.get(connName);
         if (workerConnector != null) {
             ClassLoader connectorLoader =
                     plugins.delegatingLoader().connectorLoader(workerConnector.connector());
-            transitionTo(workerConnector, state, connectorLoader);
+            executeStateTransition(
+                () -> workerConnector.transitionTo(state, stateChangeCallback),
+                connectorLoader);
         }
 
         for (Map.Entry taskEntry : tasks.entrySet()) {
             if (taskEntry.getKey().connector().equals(connName)) {
                 WorkerTask workerTask = taskEntry.getValue();
-                transitionTo(workerTask, state, workerTask.loader());
+                executeStateTransition(() -> workerTask.transitionTo(state), workerTask.loader);
             }
         }
     }
 
-    private void transitionTo(Object connectorOrTask, TargetState state, ClassLoader loader) {
+    private void executeStateTransition(Runnable stateTransition, ClassLoader loader) {
         ClassLoader savedLoader = plugins.currentThreadLoader();
         try {
             savedLoader = Plugins.compareAndSwapLoaders(loader);
-            if (connectorOrTask instanceof WorkerConnector) {
-                ((WorkerConnector) connectorOrTask).transitionTo(state);
-            } else if (connectorOrTask instanceof WorkerTask) {
-                ((WorkerTask) connectorOrTask).transitionTo(state);
-            } else {
-                throw new ConnectException(
-                        "Request for state transition on an object that is neither a "
-                                + "WorkerConnector nor a WorkerTask: "
-                                + connectorOrTask.getClass());
-            }
+            stateTransition.run();
         } finally {
             Plugins.compareAndSwapLoaders(savedLoader);
         }
     }
 
+    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;
@@ -622,18 +1078,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);
@@ -642,13 +1088,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);
@@ -657,13 +1103,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 dfae761ec9bd8..a137b2d8f6adc 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
@@ -16,25 +16,47 @@
  */
 package org.apache.kafka.connect.runtime;
 
+import org.apache.kafka.clients.ClientDnsLookup;
 import org.apache.kafka.clients.CommonClientConfigs;
 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.ConfigException;
+import org.apache.kafka.common.config.internals.BrokerSecurityConfigs;
 import org.apache.kafka.common.metrics.Sensor;
+import org.apache.kafka.connect.json.JsonConverter;
+import org.apache.kafka.connect.json.JsonConverterConfig;
+import org.apache.kafka.connect.storage.Converter;
+import org.apache.kafka.connect.storage.SimpleHeaderConverter;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
+import java.util.regex.Pattern;
+
+import org.eclipse.jetty.util.StringUtil;
 
 import static org.apache.kafka.common.config.ConfigDef.Range.atLeast;
 import static org.apache.kafka.common.config.ConfigDef.ValidString.in;
+import static org.apache.kafka.connect.runtime.SourceConnectorConfig.TOPIC_CREATION_PREFIX;
 
 /**
  * Common base class providing configuration for Kafka Connect workers, whether standalone or distributed.
  */
 public class WorkerConfig extends AbstractConfig {
+    private static final Logger log = LoggerFactory.getLogger(WorkerConfig.class);
+
+    private static final Pattern COMMA_WITH_WHITESPACE = Pattern.compile("\\s*,\\s*");
+    private static final Collection HEADER_ACTIONS = Collections.unmodifiableList(
+            Arrays.asList("set", "add", "setDate", "addDate")
+    );
 
     public static final String BOOTSTRAP_SERVERS_CONFIG = "bootstrap.servers";
     public static final String BOOTSTRAP_SERVERS_DOC
@@ -48,6 +70,9 @@ public class WorkerConfig extends AbstractConfig {
             + "than one, though, in case a server is down).";
     public static final String BOOTSTRAP_SERVERS_DEFAULT = "localhost:9092";
 
+    public static final String CLIENT_DNS_LOOKUP_CONFIG = CommonClientConfigs.CLIENT_DNS_LOOKUP_CONFIG;
+    public static final String CLIENT_DNS_LOOKUP_DOC = CommonClientConfigs.CLIENT_DNS_LOOKUP_DOC;
+
     public static final String KEY_CONVERTER_CLASS_CONFIG = "key.converter";
     public static final String KEY_CONVERTER_CLASS_DOC =
             "Converter class used to convert between Kafka Connect format and the serialized form that is written to Kafka." +
@@ -62,6 +87,19 @@ public class WorkerConfig extends AbstractConfig {
                     " independent of connectors it allows any connector to work with any serialization format." +
                     " Examples of common formats include JSON and Avro.";
 
+    public static final String HEADER_CONVERTER_CLASS_CONFIG = "header.converter";
+    public static final String HEADER_CONVERTER_CLASS_DOC =
+            "HeaderConverter class used to convert between Kafka Connect format and the serialized form that is written to Kafka." +
+                    " This controls the format of the header values in messages written to or read from Kafka, and since this is" +
+                    " independent of connectors it allows any connector to work with any serialization format." +
+                    " Examples of common formats include JSON and Avro. By default, the SimpleHeaderConverter is used to serialize" +
+                    " header values to strings and deserialize them by inferring the schemas.";
+    public static final String HEADER_CONVERTER_CLASS_DEFAULT = SimpleHeaderConverter.class.getName();
+
+    /**
+     * @deprecated As of 2.0.0
+     */
+    @Deprecated
     public static final String INTERNAL_KEY_CONVERTER_CLASS_CONFIG = "internal.key.converter";
     public static final String INTERNAL_KEY_CONVERTER_CLASS_DOC =
             "Converter class used to convert between Kafka Connect format and the serialized form that is written to Kafka." +
@@ -69,8 +107,13 @@ public class WorkerConfig extends AbstractConfig {
                     " independent of connectors it allows any connector to work with any serialization format." +
                     " Examples of common formats include JSON and Avro." +
                     " This setting controls the format used for internal bookkeeping data used by the framework, such as" +
-                    " configs and offsets, so users can typically use any functioning Converter implementation.";
+                    " configs and offsets, so users can typically use any functioning Converter implementation." +
+                    " Deprecated; will be removed in an upcoming version.";
 
+    /**
+     * @deprecated As of 2.0.0
+     */
+    @Deprecated
     public static final String INTERNAL_VALUE_CONVERTER_CLASS_CONFIG = "internal.value.converter";
     public static final String INTERNAL_VALUE_CONVERTER_CLASS_DOC =
             "Converter class used to convert between Kafka Connect format and the serialized form that is written to Kafka." +
@@ -78,7 +121,10 @@ public class WorkerConfig extends AbstractConfig {
                     " independent of connectors it allows any connector to work with any serialization format." +
                     " Examples of common formats include JSON and Avro." +
                     " This setting controls the format used for internal bookkeeping data used by the framework, such as" +
-                    " configs and offsets, so users can typically use any functioning Converter implementation.";
+                    " configs and offsets, so users can typically use any functioning Converter implementation." +
+                    " Deprecated; will be removed in an upcoming version.";
+
+    private static final Class INTERNAL_CONVERTER_DEFAULT = JsonConverter.class;
 
     public static final String TASK_SHUTDOWN_GRACEFUL_TIMEOUT_MS_CONFIG
             = "task.shutdown.graceful.timeout.ms";
@@ -99,15 +145,30 @@ public class WorkerConfig extends AbstractConfig {
             + "data to be committed in a future attempt.";
     public static final long OFFSET_COMMIT_TIMEOUT_MS_DEFAULT = 5000L;
 
+    /**
+     * @deprecated As of 1.1.0.
+     */
+    @Deprecated
     public static final String REST_HOST_NAME_CONFIG = "rest.host.name";
     private static final String REST_HOST_NAME_DOC
             = "Hostname for the REST API. If this is set, it will only bind to this interface.";
 
+    /**
+     * @deprecated As of 1.1.0.
+     */
+    @Deprecated
     public static final String REST_PORT_CONFIG = "rest.port";
     private static final String REST_PORT_DOC
             = "Port for the REST API to listen on.";
     public static final int REST_PORT_DEFAULT = 8083;
 
+    public static final String LISTENERS_CONFIG = "listeners";
+    private static final String LISTENERS_DOC
+            = "List of comma-separated URIs the REST API will listen on. The supported protocols are HTTP and HTTPS.\n" +
+            " Specify hostname as 0.0.0.0 to bind to all interfaces.\n" +
+            " Leave hostname empty to bind to default interface.\n" +
+            " Examples of legal listener lists: HTTP://myhost:8083,HTTPS://myhost:8084";
+
     public static final String REST_ADVERTISED_HOST_NAME_CONFIG = "rest.advertised.host.name";
     private static final String REST_ADVERTISED_HOST_NAME_DOC
             = "If this is set, this is the hostname that will be given out to other workers to connect to.";
@@ -116,6 +177,10 @@ public class WorkerConfig extends AbstractConfig {
     private static final String REST_ADVERTISED_PORT_DOC
             = "If this is set, this is the port that will be given out to other workers to connect to.";
 
+    public static final String REST_ADVERTISED_LISTENER_CONFIG = "rest.advertised.listener";
+    private static final String REST_ADVERTISED_LISTENER_DOC
+            = "Sets the advertised listener (HTTP or HTTPS) which will be given to other workers to use.";
+
     public static final String ACCESS_CONTROL_ALLOW_ORIGIN_CONFIG = "access.control.allow.origin";
     protected static final String ACCESS_CONTROL_ALLOW_ORIGIN_DOC =
             "Value to set the Access-Control-Allow-Origin header to for REST API requests." +
@@ -130,6 +195,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"
@@ -140,13 +213,63 @@ 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 =
+            "Comma-separated names of ConfigProvider classes, loaded and used "
+            + "in the order specified. Implementing the interface  "
+            + "ConfigProvider allows you to replace variable references in connector configurations, "
+            + "such as for externalized secrets. ";
+
+    public static final String REST_EXTENSION_CLASSES_CONFIG = "rest.extension.classes";
+    protected static final String REST_EXTENSION_CLASSES_DOC =
+            "Comma-separated names of ConnectRestExtension classes, loaded and called "
+            + "in the order specified. Implementing the interface  "
+            + "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;
     public static final String METRIC_REPORTER_CLASSES_CONFIG = CommonClientConfigs.METRIC_REPORTER_CLASSES_CONFIG;
 
+    public static final String TOPIC_TRACKING_ENABLE_CONFIG = "topic.tracking.enable";
+    protected static final String TOPIC_TRACKING_ENABLE_DOC = "Enable tracking the set of active "
+            + "topics per connector during runtime.";
+    protected static final boolean TOPIC_TRACKING_ENABLE_DEFAULT = true;
+
+    public static final String TOPIC_TRACKING_ALLOW_RESET_CONFIG = "topic.tracking.allow.reset";
+    protected static final String TOPIC_TRACKING_ALLOW_RESET_DOC = "If set to true, it allows "
+            + "user requests to reset the set of active topics per connector.";
+    protected static final boolean TOPIC_TRACKING_ALLOW_RESET_DEFAULT = true;
+
+    public static final String CONNECT_KAFKA_CLUSTER_ID = "connect.kafka.cluster.id";
+    public static final String CONNECT_GROUP_ID = "connect.group.id";
+
+    public static final String TOPIC_CREATION_ENABLE_CONFIG = "topic.creation.enable";
+    protected static final String TOPIC_CREATION_ENABLE_DOC = "Whether to allow "
+            + "automatic creation of topics used by source connectors, when source connectors "
+            + "are configured with `" + TOPIC_CREATION_PREFIX + "` properties. Each task will use an "
+            + "admin client to create its topics and will not depend on the Kafka brokers "
+            + "to create topics automatically.";
+    protected static final boolean TOPIC_CREATION_ENABLE_DEFAULT = true;
+
+    public static final String RESPONSE_HTTP_HEADERS_CONFIG = "response.http.headers.config";
+    protected static final String RESPONSE_HTTP_HEADERS_DOC = "Rules for REST API HTTP response headers";
+    protected static final String RESPONSE_HTTP_HEADERS_DEFAULT = "";
+
     /**
      * Get a basic ConfigDef for a WorkerConfig. This includes all the common settings. Subclasses can use this to
      * bootstrap their own ConfigDef.
@@ -156,13 +279,21 @@ protected static ConfigDef baseConfigDef() {
         return new ConfigDef()
                 .define(BOOTSTRAP_SERVERS_CONFIG, Type.LIST, BOOTSTRAP_SERVERS_DEFAULT,
                         Importance.HIGH, BOOTSTRAP_SERVERS_DOC)
+                .define(CLIENT_DNS_LOOKUP_CONFIG,
+                        Type.STRING,
+                        ClientDnsLookup.USE_ALL_DNS_IPS.toString(),
+                        in(ClientDnsLookup.DEFAULT.toString(),
+                           ClientDnsLookup.USE_ALL_DNS_IPS.toString(),
+                           ClientDnsLookup.RESOLVE_CANONICAL_BOOTSTRAP_SERVERS_ONLY.toString()),
+                        Importance.MEDIUM,
+                        CLIENT_DNS_LOOKUP_DOC)
                 .define(KEY_CONVERTER_CLASS_CONFIG, Type.CLASS,
                         Importance.HIGH, KEY_CONVERTER_CLASS_DOC)
                 .define(VALUE_CONVERTER_CLASS_CONFIG, Type.CLASS,
                         Importance.HIGH, VALUE_CONVERTER_CLASS_DOC)
-                .define(INTERNAL_KEY_CONVERTER_CLASS_CONFIG, Type.CLASS,
+                .define(INTERNAL_KEY_CONVERTER_CLASS_CONFIG, Type.CLASS, INTERNAL_CONVERTER_DEFAULT,
                         Importance.LOW, INTERNAL_KEY_CONVERTER_CLASS_DOC)
-                .define(INTERNAL_VALUE_CONVERTER_CLASS_CONFIG, Type.CLASS,
+                .define(INTERNAL_VALUE_CONVERTER_CLASS_CONFIG, Type.CLASS, INTERNAL_CONVERTER_DEFAULT,
                         Importance.LOW, INTERNAL_VALUE_CONVERTER_CLASS_DOC)
                 .define(TASK_SHUTDOWN_GRACEFUL_TIMEOUT_MS_CONFIG, Type.LONG,
                         TASK_SHUTDOWN_GRACEFUL_TIMEOUT_MS_DEFAULT, Importance.LOW,
@@ -173,8 +304,10 @@ protected static ConfigDef baseConfigDef() {
                         Importance.LOW, OFFSET_COMMIT_TIMEOUT_MS_DOC)
                 .define(REST_HOST_NAME_CONFIG, Type.STRING, null, Importance.LOW, REST_HOST_NAME_DOC)
                 .define(REST_PORT_CONFIG, Type.INT, REST_PORT_DEFAULT, Importance.LOW, REST_PORT_DOC)
+                .define(LISTENERS_CONFIG, Type.LIST, null, Importance.LOW, LISTENERS_DOC)
                 .define(REST_ADVERTISED_HOST_NAME_CONFIG, Type.STRING,  null, Importance.LOW, REST_ADVERTISED_HOST_NAME_DOC)
                 .define(REST_ADVERTISED_PORT_CONFIG, Type.INT,  null, Importance.LOW, REST_ADVERTISED_PORT_DOC)
+                .define(REST_ADVERTISED_LISTENER_CONFIG, Type.STRING,  null, Importance.LOW, REST_ADVERTISED_LISTENER_DOC)
                 .define(ACCESS_CONTROL_ALLOW_ORIGIN_CONFIG, Type.STRING,
                         ACCESS_CONTROL_ALLOW_ORIGIN_DEFAULT, Importance.LOW,
                         ACCESS_CONTROL_ALLOW_ORIGIN_DOC)
@@ -199,7 +332,108 @@ protected static ConfigDef baseConfigDef() {
                         CommonClientConfigs.METRICS_RECORDING_LEVEL_DOC)
                 .define(METRIC_REPORTER_CLASSES_CONFIG, Type.LIST,
                         "", Importance.LOW,
-                        CommonClientConfigs.METRIC_REPORTER_CLASSES_DOC);
+                        CommonClientConfigs.METRIC_REPORTER_CLASSES_DOC)
+                .define(BrokerSecurityConfigs.SSL_CLIENT_AUTH_CONFIG,
+                        ConfigDef.Type.STRING, "none", ConfigDef.Importance.LOW, BrokerSecurityConfigs.SSL_CLIENT_AUTH_DOC)
+                .define(HEADER_CONVERTER_CLASS_CONFIG, Type.CLASS,
+                        HEADER_CONVERTER_CLASS_DEFAULT,
+                        Importance.LOW, HEADER_CONVERTER_CLASS_DOC)
+                .define(CONFIG_PROVIDERS_CONFIG, Type.LIST,
+                        Collections.emptyList(),
+                        Importance.LOW, CONFIG_PROVIDERS_DOC)
+                .define(REST_EXTENSION_CLASSES_CONFIG, Type.LIST, "",
+                        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)
+                .define(TOPIC_TRACKING_ENABLE_CONFIG, Type.BOOLEAN, TOPIC_TRACKING_ENABLE_DEFAULT,
+                        Importance.LOW, TOPIC_TRACKING_ENABLE_DOC)
+                .define(TOPIC_TRACKING_ALLOW_RESET_CONFIG, Type.BOOLEAN, TOPIC_TRACKING_ALLOW_RESET_DEFAULT,
+                        Importance.LOW, TOPIC_TRACKING_ALLOW_RESET_DOC)
+                .define(TOPIC_CREATION_ENABLE_CONFIG, Type.BOOLEAN, TOPIC_CREATION_ENABLE_DEFAULT, Importance.LOW,
+                        TOPIC_CREATION_ENABLE_DOC)
+                .define(RESPONSE_HTTP_HEADERS_CONFIG, Type.STRING, RESPONSE_HTTP_HEADERS_DEFAULT,
+                        new ResponseHttpHeadersValidator(), Importance.LOW, RESPONSE_HTTP_HEADERS_DOC)
+                // security support
+                .withClientSslSupport();
+    }
+
+    private void logInternalConverterDeprecationWarnings(Map props) {
+        String[] deprecatedConfigs = new String[] {
+            INTERNAL_KEY_CONVERTER_CLASS_CONFIG,
+            INTERNAL_VALUE_CONVERTER_CLASS_CONFIG
+        };
+        for (String config : deprecatedConfigs) {
+            if (props.containsKey(config)) {
+                Class internalConverterClass = getClass(config);
+                logDeprecatedProperty(config, internalConverterClass.getCanonicalName(), INTERNAL_CONVERTER_DEFAULT.getCanonicalName(), null);
+                if (internalConverterClass.equals(INTERNAL_CONVERTER_DEFAULT)) {
+                    // log the properties for this converter ...
+                    for (Map.Entry propEntry : originalsWithPrefix(config + ".").entrySet()) {
+                        String prop = propEntry.getKey();
+                        String propValue = propEntry.getValue().toString();
+                        String defaultValue = JsonConverterConfig.SCHEMAS_ENABLE_CONFIG.equals(prop) ? "false" : null;
+                        logDeprecatedProperty(config + "." + prop, propValue, defaultValue, config);
+                    }
+                }
+            }
+        }
+    }
+
+    private void logDeprecatedProperty(String propName, String propValue, String defaultValue, String prefix) {
+        String prefixNotice = prefix != null
+            ? " (along with all configuration for '" + prefix + "')"
+            : "";
+        if (defaultValue != null && defaultValue.equalsIgnoreCase(propValue)) {
+            log.info(
+                "Worker configuration property '{}'{} is deprecated and may be removed in an upcoming release. "
+                    + "The specified value '{}' matches the default, so this property can be safely removed from the worker configuration.",
+                propName,
+                prefixNotice,
+                propValue
+            );
+        } else if (defaultValue != null) {
+            log.warn(
+                "Worker configuration property '{}'{} is deprecated and may be removed in an upcoming release. "
+                    + "The specified value '{}' does NOT match the default and recommended value '{}'.",
+                propName,
+                prefixNotice,
+                propValue,
+                defaultValue
+            );
+        } else {
+            log.warn(
+                "Worker configuration property '{}'{} is deprecated and may be removed in an upcoming release.",
+                propName,
+                prefixNotice
+            );
+        }
+    }
+
+    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;
+    }
+
+    public boolean topicCreationEnable() {
+        return getBoolean(TOPIC_CREATION_ENABLE_CONFIG);
     }
 
     @Override
@@ -211,10 +445,101 @@ public static List pluginLocations(Map props) {
         String locationList = props.get(WorkerConfig.PLUGIN_PATH_CONFIG);
         return locationList == null
                          ? new ArrayList()
-                         : Arrays.asList(locationList.trim().split("\\s*,\\s*", -1));
+                         : Arrays.asList(COMMA_WITH_WHITESPACE.split(locationList.trim(), -1));
     }
 
     public WorkerConfig(ConfigDef definition, Map props) {
         super(definition, props);
+        logInternalConverterDeprecationWarnings(props);
+        logPluginPathConfigProviderWarning(props);
+    }
+
+    // Visible for testing
+    static void validateHttpResponseHeaderConfig(String config) {
+        try {
+            // validate format
+            String[] configTokens = config.trim().split("\\s+", 2);
+            if (configTokens.length != 2) {
+                throw new ConfigException(String.format("Invalid format of header config '%s\'. "
+                        + "Expected: '[ation] [header name]:[header value]'", config));
+            }
+
+            // validate action
+            String method = configTokens[0].trim();
+            validateHeaderConfigAction(method);
+
+            // validate header name and header value pair
+            String header = configTokens[1];
+            String[] headerTokens = header.trim().split(":");
+            if (headerTokens.length != 2) {
+                throw new ConfigException(
+                        String.format("Invalid format of header name and header value pair '%s'. "
+                                + "Expected: '[header name]:[header value]'", header));
+            }
+
+            // validate header name
+            String headerName = headerTokens[0].trim();
+            if (headerName.isEmpty() || headerName.matches(".*\\s+.*")) {
+                throw new ConfigException(String.format("Invalid header name '%s'. "
+                        + "The '[header name]' cannot contain whitespace", headerName));
+            }
+        } catch (ArrayIndexOutOfBoundsException e) {
+            throw new ConfigException(String.format("Invalid header config '%s'.", config), e);
+        }
+    }
+
+    // Visible for testing
+    static void validateHeaderConfigAction(String action) {
+        if (!HEADER_ACTIONS.stream().anyMatch(action::equalsIgnoreCase)) {
+            throw new ConfigException(String.format("Invalid header config action: '%s'. "
+                    + "Expected one of %s", action, HEADER_ACTIONS));
+        }
+    }
+
+    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.");
+                }
+            }
+        }
+    }
+
+    private static class ResponseHttpHeadersValidator implements ConfigDef.Validator {
+        @Override
+        public void ensureValid(String name, Object value) {
+            String strValue = (String) value;
+            if (strValue == null || strValue.trim().isEmpty()) {
+                return;
+            }
+
+            String[] configs = StringUtil.csvSplit(strValue); // handles and removed surrounding quotes
+            Arrays.stream(configs).forEach(WorkerConfig::validateHttpResponseHeaderConfig);
+        }
+
+        @Override
+        public String toString() {
+            return "Comma-separated header rules, where each header rule is of the form "
+                    + "'[action] [header name]:[header value]' and optionally surrounded by double quotes "
+                    + "if any part of a header rule contains a comma";
+        }
     }
 }
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
new file mode 100644
index 0000000000000..318626bd5cade
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerConfigTransformer.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.connect.runtime;
+
+import org.apache.kafka.common.config.ConfigDef;
+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;
+
+import java.util.Locale;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+
+/**
+ * A wrapper class to perform configuration transformations and schedule reloads for any
+ * retrieved TTL values.
+ */
+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);
+    }
+
+    public Map transform(Map configs) {
+        return transform(null, configs);
+    }
+
+    public Map transform(String connectorName, Map configs) {
+        if (configs == null) return null;
+        ConfigTransformerResult result = configTransformer.transform(configs);
+        if (connectorName != null) {
+            String key = ConnectorConfig.CONFIG_RELOAD_ACTION_CONFIG;
+            String action = (String) ConfigDef.parseType(key, configs.get(key), ConfigDef.Type.STRING);
+            if (action == null) {
+                // The default action is "restart".
+                action = ConnectorConfig.CONFIG_RELOAD_ACTION_RESTART;
+            }
+            ConfigReloadAction reloadAction = ConfigReloadAction.valueOf(action.toUpperCase(Locale.ROOT));
+            if (reloadAction == ConfigReloadAction.RESTART) {
+                scheduleReload(connectorName, result.ttls());
+            }
+        }
+        return result.data();
+    }
+
+    private void scheduleReload(String connectorName, Map ttls) {
+        for (Map.Entry entry : ttls.entrySet()) {
+            scheduleReload(connectorName, entry.getKey(), entry.getValue());
+        }
+    }
+
+    private void scheduleReload(String connectorName, String path, long ttl) {
+        Map connectorRequests = requests.get(connectorName);
+        if (connectorRequests == null) {
+            connectorRequests = new ConcurrentHashMap<>();
+            requests.put(connectorName, connectorRequests);
+        } else {
+            HerderRequest previousRequest = connectorRequests.get(path);
+            if (previousRequest != null) {
+                // Delete previous request for ttl which is now stale
+                previousRequest.cancel();
+            }
+        }
+        log.info("Scheduling a restart of connector {} in {} ms", connectorName, ttl);
+        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 9e65cd2d80f61..7cfb4f434b90d 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,16 +18,24 @@
 
 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.errors.ConnectException;
 import org.apache.kafka.connect.runtime.ConnectMetrics.MetricGroup;
-import org.apache.kafka.connect.sink.SinkConnector;
-import org.apache.kafka.connect.source.SourceConnector;
+import org.apache.kafka.connect.runtime.isolation.Plugins;
+import org.apache.kafka.connect.sink.SinkConnectorContext;
+import org.apache.kafka.connect.source.SourceConnectorContext;
+import org.apache.kafka.connect.storage.OffsetStorageReader;
+import org.apache.kafka.connect.util.Callback;
+import org.apache.kafka.connect.util.ConnectUtils;
+import org.apache.kafka.connect.util.LoggingContext;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.util.Locale;
 import java.util.Map;
 import java.util.Objects;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
 
 /**
  * Container for connectors which is responsible for managing their lifecycle (e.g. handling startup,
@@ -41,8 +49,9 @@
  *
  * Note that this class is NOT thread-safe.
  */
-public class WorkerConnector {
+public class WorkerConnector implements Runnable {
     private static final Logger log = LoggerFactory.getLogger(WorkerConnector.class);
+    private static final String THREAD_NAME_PREFIX = "connector-thread-";
 
     private enum State {
         INIT,    // initial state before startup
@@ -52,52 +61,119 @@ private enum State {
     }
 
     private final String connName;
+    private final Map config;
     private final ConnectorStatus.Listener statusListener;
-    private final ConnectorContext ctx;
+    private final ClassLoader loader;
+    private final CloseableConnectorContext ctx;
     private final Connector connector;
     private final ConnectorMetricsGroup metrics;
+    private final AtomicReference pendingTargetStateChange;
+    private final AtomicReference> pendingStateChangeCallback;
+    private final CountDownLatch shutdownLatch;
+    private volatile boolean stopping;  // indicates whether the Worker has asked the connector to stop
+    private volatile boolean cancelled; // indicates whether the Worker has cancelled the connector (e.g. because of slow shutdown)
 
-    private Map config;
     private State state;
+    private final OffsetStorageReader offsetStorageReader;
 
     public WorkerConnector(String connName,
                            Connector connector,
-                           ConnectorContext ctx,
+                           ConnectorConfig connectorConfig,
+                           CloseableConnectorContext ctx,
                            ConnectMetrics metrics,
-                           ConnectorStatus.Listener statusListener) {
+                           ConnectorStatus.Listener statusListener,
+                           OffsetStorageReader offsetStorageReader,
+                           ClassLoader loader) {
         this.connName = connName;
+        this.config = connectorConfig.originalsStrings();
+        this.loader = loader;
         this.ctx = ctx;
         this.connector = connector;
         this.state = State.INIT;
         this.metrics = new ConnectorMetricsGroup(metrics, AbstractStatus.State.UNASSIGNED, statusListener);
         this.statusListener = this.metrics;
+        this.offsetStorageReader = offsetStorageReader;
+        this.pendingTargetStateChange = new AtomicReference<>();
+        this.pendingStateChangeCallback = new AtomicReference<>();
+        this.shutdownLatch = new CountDownLatch(1);
+        this.stopping = false;
+        this.cancelled = false;
     }
 
-    public void initialize(ConnectorConfig connectorConfig) {
-        try {
-            this.config = connectorConfig.originalsStrings();
-            log.debug("{} Initializing connector {} with config {}", this, connName, config);
+    public ClassLoader loader() {
+        return loader;
+    }
 
-            connector.initialize(new ConnectorContext() {
-                @Override
-                public void requestTaskReconfiguration() {
-                    ctx.requestTaskReconfiguration();
-                }
+    @Override
+    public void run() {
+        // Clear all MDC parameters, in case this thread is being reused
+        LoggingContext.clear();
+
+        try (LoggingContext loggingContext = LoggingContext.forConnector(connName)) {
+            ClassLoader savedLoader = Plugins.compareAndSwapLoaders(loader);
+            String savedName = Thread.currentThread().getName();
+            try {
+                Thread.currentThread().setName(THREAD_NAME_PREFIX + connName);
+                doRun();
+            } finally {
+                Thread.currentThread().setName(savedName);
+                Plugins.compareAndSwapLoaders(savedLoader);
+            }
+        } finally {
+            // In the rare case of an exception being thrown outside the doRun() method, or an
+            // uncaught one being thrown from within it, mark the connector as shut down to avoid
+            // unnecessarily blocking and eventually timing out during awaitShutdown
+            shutdownLatch.countDown();
+        }
+    }
 
-                @Override
-                public void raiseError(Exception e) {
-                    log.error("{} Connector raised an error", this, e);
-                    onFailure(e);
-                    ctx.raiseError(e);
+    void doRun() {
+        initialize();
+        while (!stopping) {
+            TargetState newTargetState;
+            Callback stateChangeCallback;
+            synchronized (this) {
+                newTargetState = pendingTargetStateChange.getAndSet(null);
+                stateChangeCallback = pendingStateChangeCallback.getAndSet(null);
+            }
+            if (newTargetState != null && !stopping) {
+                doTransitionTo(newTargetState, stateChangeCallback);
+            }
+            synchronized (this) {
+                if (pendingTargetStateChange.get() != null || stopping) {
+                    // An update occurred before we entered the synchronized block; no big deal,
+                    // just start the loop again until we've handled everything
+                } else {
+                    try {
+                        wait();
+                    } catch (InterruptedException e) {
+                        // We'll pick up any potential state changes at the top of the loop
+                    }
                 }
-            });
+            }
+        }
+        doShutdown();
+    }
+
+    void initialize() {
+        try {
+            if (!isSourceConnector() && !isSinkConnector()) {
+                throw new ConnectException("Connector implementations must be a subclass of either SourceConnector or SinkConnector");
+            }
+            log.debug("{} Initializing connector {}", this, connName);
+            if (isSinkConnector()) {
+                SinkConnectorConfig.validate(config);
+                connector.initialize(new WorkerSinkConnectorContext());
+            } else {
+                connector.initialize(new WorkerSourceConnectorContext(offsetStorageReader));
+            }
         } catch (Throwable t) {
             log.error("{} Error initializing connector", this, t);
             onFailure(t);
         }
     }
 
-    private boolean doStart() {
+    private boolean doStart() throws Throwable {
         try {
             switch (state) {
                 case STARTED:
@@ -115,7 +191,7 @@ private boolean doStart() {
         } catch (Throwable t) {
             log.error("{} Error while starting connector", this, t);
             onFailure(t);
-            return false;
+            throw t;
         }
     }
 
@@ -124,12 +200,12 @@ private void onFailure(Throwable t) {
         this.state = State.FAILED;
     }
 
-    private void resume() {
+    private void resume() throws Throwable {
         if (doStart())
             statusListener.onResume(connName);
     }
 
-    private void start() {
+    private void start() throws Throwable {
         if (doStart())
             statusListener.onStartup(connName);
     }
@@ -138,6 +214,7 @@ public boolean isRunning() {
         return state == State.STARTED;
     }
 
+    @SuppressWarnings("fallthrough")
     private void pause() {
         try {
             switch (state) {
@@ -163,27 +240,106 @@ private void pause() {
         }
     }
 
-    public void shutdown() {
+    /**
+     * Stop this connector. This method does not block, it only triggers shutdown. Use
+     * #{@link #awaitShutdown} to block until completion.
+     */
+    public synchronized void shutdown() {
+        log.info("Scheduled shutdown for {}", this);
+        stopping = true;
+        notify();
+    }
+
+    void doShutdown() {
         try {
+            TargetState preEmptedState = pendingTargetStateChange.getAndSet(null);
+            Callback stateChangeCallback = pendingStateChangeCallback.getAndSet(null);
+            if (stateChangeCallback != null) {
+                stateChangeCallback.onCompletion(
+                        new ConnectException(
+                                "Could not begin changing connector state to " + preEmptedState.name()
+                                    + " as the connector has been scheduled for shutdown"),
+                        null);
+            }
             if (state == State.STARTED)
                 connector.stop();
             this.state = State.STOPPED;
             statusListener.onShutdown(connName);
+            log.info("Completed shutdown for {}", this);
         } catch (Throwable t) {
             log.error("{} Error while shutting down connector", this, t);
-            this.state = State.FAILED;
+            state = State.FAILED;
             statusListener.onFailure(connName, t);
         } finally {
+            ctx.close();
             metrics.close();
         }
     }
 
-    public void transitionTo(TargetState targetState) {
+    public synchronized void cancel() {
+        // Proactively update the status of the connector to UNASSIGNED since this connector
+        // instance is being abandoned and we won't update the status on its behalf any more
+        // after this since a new instance may be started soon
+        statusListener.onShutdown(connName);
+        ctx.close();
+        cancelled = true;
+    }
+
+    /**
+     * Wait for this connector to finish shutting down.
+     *
+     * @param timeoutMs time in milliseconds to await shutdown
+     * @return true if successful, false if the timeout was reached
+     */
+    public boolean awaitShutdown(long timeoutMs) {
+        try {
+            return shutdownLatch.await(timeoutMs, TimeUnit.MILLISECONDS);
+        } catch (InterruptedException e) {
+            return false;
+        }
+    }
+
+    public void transitionTo(TargetState targetState, Callback stateChangeCallback) {
+        Callback preEmptedStateChangeCallback;
+        TargetState preEmptedState;
+        synchronized (this) {
+            preEmptedStateChangeCallback = pendingStateChangeCallback.getAndSet(stateChangeCallback);
+            preEmptedState = pendingTargetStateChange.getAndSet(targetState);
+            notify();
+        }
+        if (preEmptedStateChangeCallback != null) {
+            preEmptedStateChangeCallback.onCompletion(
+                    new ConnectException(
+                            "Could not begin changing connector state to " + preEmptedState.name()
+                                + " before another request to change state was made;"
+                                + " the new request (which is to change the state to " + targetState.name()
+                                + ") has pre-empted this one"),
+                    null
+            );
+        }
+    }
+
+    void doTransitionTo(TargetState targetState, Callback stateChangeCallback) {
         if (state == State.FAILED) {
-            log.warn("{} Cannot transition connector to {} since it has failed", this, targetState);
+            stateChangeCallback.onCompletion(
+                    new ConnectException(this + " Cannot transition connector to " + targetState + " since it has failed"),
+                    null);
             return;
         }
 
+        try {
+            doTransitionTo(targetState);
+            stateChangeCallback.onCompletion(null, targetState);
+        } catch (Throwable t) {
+            stateChangeCallback.onCompletion(
+                    new ConnectException(
+                            "Failed to transition connector " + connName + " to state " + targetState,
+                            t),
+                    null);
+        }
+    }
+
+    private void doTransitionTo(TargetState targetState) throws Throwable {
         log.debug("{} Transition connector to {}", this, targetState);
         if (targetState == TargetState.PAUSED) {
             pause();
@@ -198,11 +354,11 @@ public void transitionTo(TargetState targetState) {
     }
 
     public boolean isSinkConnector() {
-        return SinkConnector.class.isAssignableFrom(connector.getClass());
+        return ConnectUtils.isSinkConnector(connector);
     }
 
     public boolean isSourceConnector() {
-        return SourceConnector.class.isAssignableFrom(connector.getClass());
+        return ConnectUtils.isSourceConnector(connector);
     }
 
     protected String connectorType() {
@@ -228,7 +384,7 @@ public String toString() {
                        '}';
     }
 
-    class ConnectorMetricsGroup implements ConnectorStatus.Listener {
+    class ConnectorMetricsGroup implements ConnectorStatus.Listener, AutoCloseable {
         /**
          * Use {@link AbstractStatus.State} since it has all of the states we want,
          * unlike {@link WorkerConnector.State}.
@@ -247,16 +403,13 @@ public ConnectorMetricsGroup(ConnectMetrics connectMetrics, AbstractStatus.State
             ConnectMetricsRegistry registry = connectMetrics.registry();
             this.metricGroup = connectMetrics.group(registry.connectorGroupName(),
                     registry.connectorTagName(), connName);
+            // prevent collisions by removing any previously created metrics in this group.
+            metricGroup.close();
 
             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() {
@@ -266,31 +419,51 @@ public void close() {
         @Override
         public void onStartup(String connector) {
             state = AbstractStatus.State.RUNNING;
-            delegate.onStartup(connector);
+            synchronized (this) {
+                if (!cancelled) {
+                    delegate.onStartup(connector);
+                }
+            }
         }
 
         @Override
         public void onShutdown(String connector) {
             state = AbstractStatus.State.UNASSIGNED;
-            delegate.onShutdown(connector);
+            synchronized (this) {
+                if (!cancelled) {
+                    delegate.onShutdown(connector);
+                }
+            }
         }
 
         @Override
         public void onPause(String connector) {
             state = AbstractStatus.State.PAUSED;
-            delegate.onPause(connector);
+            synchronized (this) {
+                if (!cancelled) {
+                    delegate.onPause(connector);
+                }
+            }
         }
 
         @Override
         public void onResume(String connector) {
             state = AbstractStatus.State.RUNNING;
-            delegate.onResume(connector);
+            synchronized (this) {
+                if (!cancelled) {
+                    delegate.onResume(connector);
+                }
+            }
         }
 
         @Override
         public void onFailure(String connector, Throwable cause) {
             state = AbstractStatus.State.FAILED;
-            delegate.onFailure(connector, cause);
+            synchronized (this) {
+                if (!cancelled) {
+                    delegate.onFailure(connector, cause);
+                }
+            }
         }
 
         @Override
@@ -319,4 +492,36 @@ protected MetricGroup metricGroup() {
             return metricGroup;
         }
     }
+
+    private abstract class WorkerConnectorContext implements ConnectorContext {
+
+        @Override
+        public void requestTaskReconfiguration() {
+            WorkerConnector.this.ctx.requestTaskReconfiguration();
+        }
+
+        @Override
+        public void raiseError(Exception e) {
+            log.error("{} Connector raised an error", WorkerConnector.this, e);
+            onFailure(e);
+            WorkerConnector.this.ctx.raiseError(e);
+        }
+    }
+
+    private class WorkerSinkConnectorContext extends WorkerConnectorContext implements SinkConnectorContext {
+    }
+
+    private class WorkerSourceConnectorContext extends WorkerConnectorContext implements SourceConnectorContext {
+
+        private final OffsetStorageReader offsetStorageReader;
+
+        WorkerSourceConnectorContext(OffsetStorageReader offsetStorageReader) {
+            this.offsetStorageReader = offsetStorageReader;
+        }
+
+        @Override
+        public OffsetStorageReader offsetStorageReader() {
+            return offsetStorageReader;
+        }
+    }
 }
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 05ace587249c0..f95c6a71759d0 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
@@ -16,7 +16,6 @@
  */
 package org.apache.kafka.connect.runtime;
 
-import org.apache.kafka.clients.consumer.ConsumerConfig;
 import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.consumer.ConsumerRecords;
@@ -25,31 +24,38 @@
 import org.apache.kafka.clients.consumer.OffsetCommitCallback;
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.TopicPartition;
-import org.apache.kafka.common.config.ConfigException;
 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;
+import org.apache.kafka.connect.header.ConnectHeaders;
+import org.apache.kafka.connect.header.Headers;
 import org.apache.kafka.connect.runtime.ConnectMetrics.MetricGroup;
+import org.apache.kafka.connect.runtime.distributed.ClusterConfigState;
+import org.apache.kafka.connect.runtime.errors.RetryWithToleranceOperator;
+import org.apache.kafka.connect.runtime.errors.Stage;
+import org.apache.kafka.connect.runtime.errors.WorkerErrantRecordReporter;
 import org.apache.kafka.connect.sink.SinkRecord;
 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.StatusBackingStore;
 import org.apache.kafka.connect.util.ConnectUtils;
 import org.apache.kafka.connect.util.ConnectorTaskId;
-import org.apache.kafka.connect.util.SinkUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.time.Duration;
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
@@ -57,6 +63,7 @@
 import java.util.regex.Pattern;
 
 import static java.util.Collections.singleton;
+import static org.apache.kafka.connect.runtime.WorkerConfig.TOPIC_TRACKING_ENABLE_CONFIG;
 
 /**
  * WorkerTask that uses a SinkTask to export data from Kafka.
@@ -66,12 +73,14 @@ class WorkerSinkTask extends WorkerTask {
 
     private final WorkerConfig workerConfig;
     private final SinkTask task;
+    private final ClusterConfigState configState;
     private Map taskConfig;
-    private final Time time;
     private final Converter keyConverter;
     private final Converter valueConverter;
+    private final HeaderConverter headerConverter;
     private final TransformationChain transformationChain;
     private final SinkTaskMetricsGroup sinkTaskMetricsGroup;
+    private final boolean isTopicTrackingEnabled;
     private KafkaConsumer consumer;
     private WorkerSinkTaskContext context;
     private final List messageBatch;
@@ -85,26 +94,36 @@ class WorkerSinkTask extends WorkerTask {
     private int commitFailures;
     private boolean pausedForRedelivery;
     private boolean committing;
+    private boolean taskStopped;
+    private final WorkerErrantRecordReporter workerErrantRecordReporter;
 
     public WorkerSinkTask(ConnectorTaskId id,
                           SinkTask task,
                           TaskStatus.Listener statusListener,
                           TargetState initialState,
                           WorkerConfig workerConfig,
+                          ClusterConfigState configState,
                           ConnectMetrics connectMetrics,
                           Converter keyConverter,
                           Converter valueConverter,
+                          HeaderConverter headerConverter,
                           TransformationChain transformationChain,
+                          KafkaConsumer consumer,
                           ClassLoader loader,
-                          Time time) {
-        super(id, statusListener, initialState, loader, connectMetrics);
+                          Time time,
+                          RetryWithToleranceOperator retryWithToleranceOperator,
+                          WorkerErrantRecordReporter workerErrantRecordReporter,
+                          StatusBackingStore statusBackingStore) {
+        super(id, statusListener, initialState, loader, connectMetrics,
+                retryWithToleranceOperator, time, statusBackingStore);
 
         this.workerConfig = workerConfig;
         this.task = task;
+        this.configState = configState;
         this.keyConverter = keyConverter;
         this.valueConverter = valueConverter;
+        this.headerConverter = headerConverter;
         this.transformationChain = transformationChain;
-        this.time = time;
         this.messageBatch = new ArrayList<>();
         this.currentOffsets = new HashMap<>();
         this.origOffsets = new HashMap<>();
@@ -118,14 +137,17 @@ public WorkerSinkTask(ConnectorTaskId id,
         this.commitFailures = 0;
         this.sinkTaskMetricsGroup = new SinkTaskMetricsGroup(id, connectMetrics);
         this.sinkTaskMetricsGroup.recordOffsetSequenceNumber(commitSeqno);
+        this.consumer = consumer;
+        this.isTopicTrackingEnabled = workerConfig.getBoolean(TOPIC_TRACKING_ENABLE_CONFIG);
+        this.taskStopped = false;
+        this.workerErrantRecordReporter = workerErrantRecordReporter;
     }
 
     @Override
     public void initialize(TaskConfig taskConfig) {
         try {
             this.taskConfig = taskConfig.originalsStrings();
-            this.consumer = createConsumer();
-            this.context = new WorkerSinkTaskContext(consumer);
+            this.context = new WorkerSinkTaskContext(consumer, this, configState);
         } catch (Throwable t) {
             log.error("{} Task failed initialization and will not be started.", this, t);
             onFailure(t);
@@ -143,15 +165,24 @@ public void stop() {
     protected void close() {
         // FIXME Kafka needs to add a timeout parameter here for us to properly obey the timeout
         // passed in
-        task.stop();
-        if (consumer != null)
-            consumer.close();
-        transformationChain.close();
+        try {
+            task.stop();
+        } catch (Throwable t) {
+            log.warn("Could not stop task", t);
+        }
+        taskStopped = true;
+        Utils.closeQuietly(consumer, "consumer");
+        Utils.closeQuietly(transformationChain, "transformation chain");
+        Utils.closeQuietly(retryWithToleranceOperator, "retry operator");
     }
 
     @Override
-    protected void releaseResources() {
-        sinkTaskMetricsGroup.close();
+    public void removeMetrics() {
+        try {
+            sinkTaskMetricsGroup.close();
+        } finally {
+            super.removeMetrics();
+        }
     }
 
     @Override
@@ -163,13 +194,14 @@ 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();
+        } catch (WakeupException e) {
+            log.trace("Consumer woken up during initial offset commit attempt, " 
+                + "but succeeded during a later attempt");
         }
     }
 
@@ -182,7 +214,7 @@ protected void iteration() {
             // Maybe commit
             if (!committing && (context.isCommitRequested() || now >= nextCommit)) {
                 commitOffsets(now, false);
-                nextCommit += offsetCommitIntervalMs;
+                nextCommit = now + offsetCommitIntervalMs;
                 context.clearCommitRequest();
             }
 
@@ -259,27 +291,14 @@ public int commitFailures() {
      * Initializes and starts the SinkTask.
      */
     protected void initializeAndStart() {
-        String topicsStr = taskConfig.get(SinkTask.TOPICS_CONFIG);
-        boolean topicsStrPresent = topicsStr != null && !topicsStr.trim().isEmpty();
-
-        String topicsRegexStr = taskConfig.get(SinkTask.TOPICS_REGEX_CONFIG);
-        boolean topicsRegexStrPresent = topicsRegexStr != null && !topicsRegexStr.trim().isEmpty();
-
-        if (topicsStrPresent && topicsRegexStrPresent) {
-            throw new ConfigException(SinkTask.TOPICS_CONFIG + " and " + SinkTask.TOPICS_REGEX_CONFIG +
-                " are mutually exclusive options, but both are set.");
-        }
+        SinkConnectorConfig.validate(taskConfig);
 
-        if (!topicsStrPresent && !topicsRegexStrPresent) {
-            throw new ConfigException("Must configure one of " +
-                SinkTask.TOPICS_CONFIG + " or " + SinkTask.TOPICS_REGEX_CONFIG);
-        }
-
-        if (topicsStrPresent) {
-            String[] topics = topicsStr.split(",");
-            consumer.subscribe(Arrays.asList(topics), new HandleRebalance());
-            log.debug("{} Initializing and starting task for topics {}", this, topics);
+        if (SinkConnectorConfig.hasTopicsConfig(taskConfig)) {
+            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);
             consumer.subscribe(pattern, new HandleRebalance());
             log.debug("{} Initializing and starting task for topics regex {}", this, topicsRegexStr);
@@ -316,7 +335,7 @@ boolean isCommitting() {
     }
 
     private void doCommitSync(Map offsets, int seqno) {
-        log.info("{} Committing offsets synchronously using sequence number {}: {}", this, seqno, offsets);
+        log.debug("{} Committing offsets synchronously using sequence number {}: {}", this, seqno, offsets);
         try {
             consumer.commitSync(offsets);
             onCommitCompleted(null, seqno, offsets);
@@ -330,7 +349,7 @@ private void doCommitSync(Map offsets, int se
     }
 
     private void doCommitAsync(Map offsets, final int seqno) {
-        log.info("{} Committing offsets asynchronously using sequence number {}: {}", this, seqno, offsets);
+        log.debug("{} Committing offsets asynchronously using sequence number {}: {}", this, seqno, offsets);
         OffsetCommitCallback cb = new OffsetCommitCallback() {
             @Override
             public void onComplete(Map offsets, Exception error) {
@@ -353,6 +372,12 @@ private void doCommit(Map offsets, boolean cl
     }
 
     private void commitOffsets(long now, boolean closing) {
+        if (workerErrantRecordReporter != null) {
+            log.trace("Awaiting all reported errors to be completed");
+            workerErrantRecordReporter.awaitAllFutures();
+            log.trace("Completed all reported errors");
+        }
+
         if (currentOffsets.isEmpty())
             return;
 
@@ -429,7 +454,7 @@ public String toString() {
     }
 
     private ConsumerRecords pollConsumer(long timeoutMs) {
-        ConsumerRecords msgs = consumer.poll(timeoutMs);
+        ConsumerRecords msgs = consumer.poll(Duration.ofMillis(timeoutMs));
 
         // Exceptions raised from the task during a rebalance should be rethrown to stop the worker
         if (rebalanceException != null) {
@@ -442,62 +467,106 @@ private ConsumerRecords pollConsumer(long timeoutMs) {
         return msgs;
     }
 
-    private KafkaConsumer createConsumer() {
-        // Include any unknown worker configs so consumer configs can be set globally on the worker
-        // and through to the task
-        Map props = new HashMap<>();
-
-        props.put(ConsumerConfig.GROUP_ID_CONFIG, SinkUtils.consumerGroupId(id.connector()));
-        props.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG,
-                Utils.join(workerConfig.getList(WorkerConfig.BOOTSTRAP_SERVERS_CONFIG), ","));
-        props.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false");
-        props.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
-        props.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.ByteArrayDeserializer");
-        props.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.ByteArrayDeserializer");
-
-        props.putAll(workerConfig.originalsWithPrefix("consumer."));
-
-        KafkaConsumer newConsumer;
-        try {
-            newConsumer = new KafkaConsumer<>(props);
-        } catch (Throwable t) {
-            throw new ConnectException("Failed to create consumer", t);
-        }
-
-        return newConsumer;
-    }
-
     private void convertMessages(ConsumerRecords msgs) {
         origOffsets.clear();
         for (ConsumerRecord msg : msgs) {
             log.trace("{} Consuming and converting message in topic '{}' partition {} at offset {} and timestamp {}",
                     this, msg.topic(), msg.partition(), msg.offset(), msg.timestamp());
-            SchemaAndValue keyAndSchema = keyConverter.toConnectData(msg.topic(), msg.key());
-            SchemaAndValue valueAndSchema = valueConverter.toConnectData(msg.topic(), msg.value());
-            Long timestamp = ConnectUtils.checkAndConvertTimestamp(msg.timestamp());
-            SinkRecord origRecord = new SinkRecord(msg.topic(), msg.partition(),
-                    keyAndSchema.schema(), keyAndSchema.value(),
-                    valueAndSchema.schema(), valueAndSchema.value(),
-                    msg.offset(),
-                    timestamp,
-                    msg.timestampType());
-            log.trace("{} Applying transformations to record in topic '{}' partition {} at offset {} and timestamp {} with key {} and value {}",
-                    this, msg.topic(), msg.partition(), msg.offset(), timestamp, keyAndSchema.value(), valueAndSchema.value());
-            SinkRecord transRecord = transformationChain.apply(origRecord);
+
+            retryWithToleranceOperator.consumerRecord(msg);
+
+            SinkRecord transRecord = convertAndTransformRecord(msg);
+
             origOffsets.put(
-                    new TopicPartition(origRecord.topic(), origRecord.kafkaPartition()),
-                    new OffsetAndMetadata(origRecord.kafkaOffset() + 1)
+                    new TopicPartition(msg.topic(), msg.partition()),
+                    new OffsetAndMetadata(msg.offset() + 1)
             );
             if (transRecord != null) {
                 messageBatch.add(transRecord);
             } else {
-                log.trace("{} Transformations returned null, so dropping record in topic '{}' partition {} at offset {} and timestamp {} with key {} and value {}",
-                        this, msg.topic(), msg.partition(), msg.offset(), timestamp, keyAndSchema.value(), valueAndSchema.value());
+                log.trace(
+                        "{} Converters and transformations returned null, possibly because of too many retries, so " +
+                                "dropping record in topic '{}' partition {} at offset {}",
+                        this, msg.topic(), msg.partition(), msg.offset()
+                );
             }
         }
         sinkTaskMetricsGroup.recordConsumedOffsets(origOffsets);
     }
 
+    private SinkRecord convertAndTransformRecord(final ConsumerRecord msg) {
+        SchemaAndValue keyAndSchema = retryWithToleranceOperator.execute(() -> convertKey(msg),
+                Stage.KEY_CONVERTER, keyConverter.getClass());
+
+        SchemaAndValue valueAndSchema = retryWithToleranceOperator.execute(() -> convertValue(msg),
+                Stage.VALUE_CONVERTER, valueConverter.getClass());
+
+        Headers headers = retryWithToleranceOperator.execute(() -> convertHeadersFor(msg), Stage.HEADER_CONVERTER, headerConverter.getClass());
+
+        if (retryWithToleranceOperator.failed()) {
+            return null;
+        }
+
+        Long timestamp = ConnectUtils.checkAndConvertTimestamp(msg.timestamp());
+        SinkRecord origRecord = new SinkRecord(msg.topic(), msg.partition(),
+                keyAndSchema.schema(), keyAndSchema.value(),
+                valueAndSchema.schema(), valueAndSchema.value(),
+                msg.offset(),
+                timestamp,
+                msg.timestampType(),
+                headers);
+        log.trace("{} Applying transformations to record in topic '{}' partition {} at offset {} and timestamp {} with key {} and value {}",
+                this, msg.topic(), msg.partition(), msg.offset(), timestamp, keyAndSchema.value(), valueAndSchema.value());
+        if (isTopicTrackingEnabled) {
+            recordActiveTopic(origRecord.topic());
+        }
+
+        // Apply the transformations
+        SinkRecord transformedRecord = transformationChain.apply(origRecord);
+        if (transformedRecord == null) {
+            return null;
+        }
+        // Error reporting will need to correlate each sink record with the original consumer record
+        return new InternalSinkRecord(msg, transformedRecord);
+    }
+
+    private SchemaAndValue convertKey(ConsumerRecord msg) {
+        try {
+            return keyConverter.toConnectData(msg.topic(), msg.headers(), msg.key());
+        } catch (Exception e) {
+            log.error("{} Error converting message key in topic '{}' partition {} at offset {} and timestamp {}: {}",
+                    this, msg.topic(), msg.partition(), msg.offset(), msg.timestamp(), e.getMessage(), e);
+            throw e;
+        }
+    }
+
+    private SchemaAndValue convertValue(ConsumerRecord msg) {
+        try {
+            return valueConverter.toConnectData(msg.topic(), msg.headers(), msg.value());
+        } catch (Exception e) {
+            log.error("{} Error converting message value in topic '{}' partition {} at offset {} and timestamp {}: {}",
+                    this, msg.topic(), msg.partition(), msg.offset(), msg.timestamp(), e.getMessage(), e);
+            throw e;
+        }
+    }
+
+    private Headers convertHeadersFor(ConsumerRecord record) {
+        Headers result = new ConnectHeaders();
+        org.apache.kafka.common.header.Headers recordHeaders = record.headers();
+        if (recordHeaders != null) {
+            String topic = record.topic();
+            for (org.apache.kafka.common.header.Header recordHeader : recordHeaders) {
+                SchemaAndValue schemaAndValue = headerConverter.toConnectHeader(topic, recordHeader.key(), recordHeader.value());
+                result.add(recordHeader.key(), schemaAndValue);
+            }
+        }
+        return result;
+    }
+
+    protected WorkerErrantRecordReporter workerErrantRecordReporter() {
+        return workerErrantRecordReporter;
+    }
+
     private void resumeAll() {
         for (TopicPartition tp : consumer.assignment())
             if (!context.pausedPartitions().contains(tp))
@@ -515,6 +584,12 @@ private void deliverMessages() {
             log.trace("{} Delivering batch of {} messages to task", this, messageBatch.size());
             long start = time.milliseconds();
             task.put(new ArrayList<>(messageBatch));
+            // if errors raised from the operator were swallowed by the task implementation, an
+            // exception needs to be thrown to kill the task indicating the tolerance was exceeded
+            if (retryWithToleranceOperator.failed() && !retryWithToleranceOperator.withinToleranceLimits()) {
+                throw new ConnectException("Tolerance exceeded in error handler",
+                    retryWithToleranceOperator.error());
+            }
             recordBatch(messageBatch.size());
             sinkTaskMetricsGroup.recordPut(time.milliseconds() - start);
             currentOffsets.putAll(origOffsets);
@@ -535,7 +610,7 @@ private void deliverMessages() {
             // Let this exit normally, the batch will be reprocessed on the next loop.
         } catch (Throwable t) {
             log.error("{} Task threw an uncaught and unrecoverable exception. Task is being killed and will not "
-                    + "recover until manually restarted.", this, t);
+                    + "recover until manually restarted. Error: {}", this, t.getMessage(), t);
             throw new ConnectException("Exiting WorkerSinkTask due to unrecoverable exception.", t);
         }
     }
@@ -591,17 +666,22 @@ SinkTaskMetricsGroup sinkTaskMetricsGroup() {
         return sinkTaskMetricsGroup;
     }
 
+    // Visible for testing
+    long getNextCommit() {
+        return nextCommit;
+    }
+
     private class HandleRebalance implements ConsumerRebalanceListener {
         @Override
         public void onPartitionsAssigned(Collection partitions) {
-            log.debug("{} Partitions assigned", WorkerSinkTask.this);
+            log.debug("{} Partitions assigned {}", WorkerSinkTask.this, partitions);
             lastCommittedOffsets = new HashMap<>();
             currentOffsets = new HashMap<>();
             for (TopicPartition tp : partitions) {
                 long pos = consumer.position(tp);
                 lastCommittedOffsets.put(tp, new OffsetAndMetadata(pos));
                 currentOffsets.put(tp, new OffsetAndMetadata(pos));
-                log.debug("{} Assigned topic partition {} with offset {}", this, tp, pos);
+                log.debug("{} Assigned topic partition {} with offset {}", WorkerSinkTask.this, tp, pos);
             }
             sinkTaskMetricsGroup.assignedOffsets(currentOffsets);
 
@@ -635,6 +715,10 @@ else if (!context.pausedPartitions().isEmpty())
 
         @Override
         public void onPartitionsRevoked(Collection partitions) {
+            if (taskStopped) {
+                log.trace("Skipping partition revocation callback as task has already been stopped");
+                return;
+            }
             log.debug("{} Partitions revoked", WorkerSinkTask.this);
             try {
                 closePartitions();
@@ -674,35 +758,37 @@ public SinkTaskMetricsGroup(ConnectorTaskId id, ConnectMetrics connectMetrics) {
             metricGroup = connectMetrics
                                   .group(registry.sinkTaskGroupName(), registry.connectorTagName(), id.connector(), registry.taskTagName(),
                                          Integer.toString(id.task()));
+            // prevent collisions by removing any previously created metrics in this group.
+            metricGroup.close();
 
-            sinkRecordRead = metricGroup.metrics().sensor("sink-record-read");
+            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.metrics().sensor("sink-record-send");
+            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.metrics().sensor("sink-record-active-count");
+            sinkRecordActiveCount = metricGroup.sensor("sink-record-active-count");
             sinkRecordActiveCount.add(metricGroup.metricName(registry.sinkRecordActiveCount), new Value());
             sinkRecordActiveCount.add(metricGroup.metricName(registry.sinkRecordActiveCountMax), new Max());
             sinkRecordActiveCount.add(metricGroup.metricName(registry.sinkRecordActiveCountAvg), new Avg());
 
-            partitionCount = metricGroup.metrics().sensor("partition-count");
+            partitionCount = metricGroup.sensor("partition-count");
             partitionCount.add(metricGroup.metricName(registry.sinkRecordPartitionCount), new Value());
 
-            offsetSeqNum = metricGroup.metrics().sensor("offset-seq-number");
+            offsetSeqNum = metricGroup.sensor("offset-seq-number");
             offsetSeqNum.add(metricGroup.metricName(registry.sinkRecordOffsetCommitSeqNum), new Value());
 
-            offsetCompletion = metricGroup.metrics().sensor("offset-commit-completion");
+            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.metrics().sensor("offset-commit-completion-skip");
+            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.metrics().sensor("put-batch-time");
+            putBatchTime = metricGroup.sensor("put-batch-time");
             putBatchTime.add(metricGroup.metricName(registry.sinkRecordPutBatchTimeMax), new Max());
             putBatchTime.add(metricGroup.metricName(registry.sinkRecordPutBatchTimeAvg), new Avg());
         }
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerSinkTaskContext.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerSinkTaskContext.java
index 64c8fff52d670..724b02ec56724 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerSinkTaskContext.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerSinkTaskContext.java
@@ -19,35 +19,55 @@
 import org.apache.kafka.clients.consumer.KafkaConsumer;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.connect.errors.IllegalWorkerStateException;
+import org.apache.kafka.connect.runtime.distributed.ClusterConfigState;
+import org.apache.kafka.connect.sink.ErrantRecordReporter;
 import org.apache.kafka.connect.sink.SinkTaskContext;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
 
 public class WorkerSinkTaskContext implements SinkTaskContext {
+
+    private final Logger log = LoggerFactory.getLogger(getClass());
     private Map offsets;
     private long timeoutMs;
     private KafkaConsumer consumer;
+    private final WorkerSinkTask sinkTask;
+    private final ClusterConfigState configState;
     private final Set pausedPartitions;
     private boolean commitRequested;
 
-    public WorkerSinkTaskContext(KafkaConsumer consumer) {
+    public WorkerSinkTaskContext(KafkaConsumer consumer,
+                                 WorkerSinkTask sinkTask,
+                                 ClusterConfigState configState) {
         this.offsets = new HashMap<>();
         this.timeoutMs = -1L;
         this.consumer = consumer;
+        this.sinkTask = sinkTask;
+        this.configState = configState;
         this.pausedPartitions = new HashSet<>();
     }
 
+    @Override
+    public Map configs() {
+        return configState.taskConfig(sinkTask.id());
+    }
+
     @Override
     public void offset(Map offsets) {
+        log.debug("{} Setting offsets for topic partitions {}", this, offsets);
         this.offsets.putAll(offsets);
     }
 
     @Override
     public void offset(TopicPartition tp, long offset) {
+        log.debug("{} Setting offset for topic partition {} to {}", this, tp, offset);
         offsets.put(tp, offset);
     }
 
@@ -65,6 +85,7 @@ public Map offsets() {
 
     @Override
     public void timeout(long timeoutMs) {
+        log.debug("{} Setting timeout to {} ms", this, timeoutMs);
         this.timeoutMs = timeoutMs;
     }
 
@@ -90,9 +111,13 @@ public void pause(TopicPartition... partitions) {
             throw new IllegalWorkerStateException("SinkTaskContext may not be used to pause consumption until the task is initialized");
         }
         try {
-            for (TopicPartition partition : partitions)
-                pausedPartitions.add(partition);
-            consumer.pause(Arrays.asList(partitions));
+            Collections.addAll(pausedPartitions, partitions);
+            if (sinkTask.shouldPause()) {
+                log.debug("{} Connector is paused, so not pausing consumer's partitions {}", this, partitions);
+            } else {
+                consumer.pause(Arrays.asList(partitions));
+                log.debug("{} Pausing partitions {}. Connector is not paused.", this, partitions);
+            }
         } catch (IllegalStateException e) {
             throw new IllegalWorkerStateException("SinkTasks may not pause partitions that are not currently assigned to them.", e);
         }
@@ -104,9 +129,13 @@ public void resume(TopicPartition... partitions) {
             throw new IllegalWorkerStateException("SinkTaskContext may not be used to resume consumption until the task is initialized");
         }
         try {
-            for (TopicPartition partition : partitions)
-                pausedPartitions.remove(partition);
-            consumer.resume(Arrays.asList(partitions));
+            pausedPartitions.removeAll(Arrays.asList(partitions));
+            if (sinkTask.shouldPause()) {
+                log.debug("{} Connector is paused, so not resuming consumer's partitions {}", this, partitions);
+            } else {
+                consumer.resume(Arrays.asList(partitions));
+                log.debug("{} Resuming partitions: {}", this, partitions);
+            }
         } catch (IllegalStateException e) {
             throw new IllegalWorkerStateException("SinkTasks may not resume partitions that are not currently assigned to them.", e);
         }
@@ -118,6 +147,7 @@ public Set pausedPartitions() {
 
     @Override
     public void requestCommit() {
+        log.debug("{} Requesting commit", this);
         commitRequested = true;
     }
 
@@ -129,4 +159,15 @@ public void clearCommitRequest() {
         commitRequested = false;
     }
 
+    @Override
+    public ErrantRecordReporter errantRecordReporter() {
+        return sinkTask.workerErrantRecordReporter();
+    }
+
+    @Override
+    public String toString() {
+        return "WorkerSinkTaskContext{" +
+               "id=" + sinkTask.id +
+               '}';
+    }
 }
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 9072cd47c812a..bc7df97476b30 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
@@ -16,31 +16,45 @@
  */
 package org.apache.kafka.connect.runtime;
 
-import org.apache.kafka.clients.producer.Callback;
+import org.apache.kafka.clients.admin.NewTopic;
+import org.apache.kafka.clients.admin.TopicDescription;
 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.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;
+import org.apache.kafka.connect.runtime.distributed.ClusterConfigState;
+import org.apache.kafka.connect.runtime.errors.RetryWithToleranceOperator;
+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.OffsetStorageReader;
+import org.apache.kafka.connect.storage.HeaderConverter;
 import org.apache.kafka.connect.storage.OffsetStorageWriter;
+import org.apache.kafka.connect.storage.StatusBackingStore;
 import org.apache.kafka.connect.util.ConnectUtils;
 import org.apache.kafka.connect.util.ConnectorTaskId;
+import org.apache.kafka.connect.util.TopicAdmin;
+import org.apache.kafka.connect.util.TopicCreation;
+import org.apache.kafka.connect.util.TopicCreationGroup;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.time.Duration;
 import java.util.IdentityHashMap;
 import java.util.List;
 import java.util.Map;
@@ -49,6 +63,9 @@
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static org.apache.kafka.connect.runtime.WorkerConfig.TOPIC_TRACKING_ENABLE_CONFIG;
 
 /**
  * WorkerTask that uses a SourceTask to ingest data into Kafka.
@@ -60,14 +77,19 @@ class WorkerSourceTask extends WorkerTask {
 
     private final WorkerConfig workerConfig;
     private final SourceTask task;
+    private final ClusterConfigState configState;
     private final Converter keyConverter;
     private final Converter valueConverter;
+    private final HeaderConverter headerConverter;
     private final TransformationChain transformationChain;
-    private KafkaProducer producer;
-    private final OffsetStorageReader offsetReader;
+    private final KafkaProducer producer;
+    private final TopicAdmin admin;
+    private final CloseableOffsetStorageReader offsetReader;
     private final OffsetStorageWriter offsetWriter;
-    private final Time time;
     private final SourceTaskMetricsGroup sourceTaskMetricsGroup;
+    private final AtomicReference producerSendException;
+    private final boolean isTopicTrackingEnabled;
+    private final TopicCreation topicCreation;
 
     private List toSend;
     private boolean lastSendFailed; // Whether the last send failed *synchronously*, i.e. never made it into the producer's RecordAccumulator
@@ -80,8 +102,7 @@ class WorkerSourceTask extends WorkerTask {
     private CountDownLatch stopRequestedLatch;
 
     private Map taskConfig;
-    private boolean finishedStart = false;
-    private boolean startedShutdownBeforeStartCompleted = false;
+    private boolean started = false;
 
     public WorkerSourceTask(ConnectorTaskId id,
                             SourceTask task,
@@ -89,25 +110,35 @@ public WorkerSourceTask(ConnectorTaskId id,
                             TargetState initialState,
                             Converter keyConverter,
                             Converter valueConverter,
+                            HeaderConverter headerConverter,
                             TransformationChain transformationChain,
                             KafkaProducer producer,
-                            OffsetStorageReader offsetReader,
+                            TopicAdmin admin,
+                            Map topicGroups,
+                            CloseableOffsetStorageReader offsetReader,
                             OffsetStorageWriter offsetWriter,
                             WorkerConfig workerConfig,
+                            ClusterConfigState configState,
                             ConnectMetrics connectMetrics,
                             ClassLoader loader,
-                            Time time) {
-        super(id, statusListener, initialState, loader, connectMetrics);
+                            Time time,
+                            RetryWithToleranceOperator retryWithToleranceOperator,
+                            StatusBackingStore statusBackingStore) {
+
+        super(id, statusListener, initialState, loader, connectMetrics,
+                retryWithToleranceOperator, time, statusBackingStore);
 
         this.workerConfig = workerConfig;
         this.task = task;
+        this.configState = configState;
         this.keyConverter = keyConverter;
         this.valueConverter = valueConverter;
+        this.headerConverter = headerConverter;
         this.transformationChain = transformationChain;
         this.producer = producer;
+        this.admin = admin;
         this.offsetReader = offsetReader;
         this.offsetWriter = offsetWriter;
-        this.time = time;
 
         this.toSend = null;
         this.lastSendFailed = false;
@@ -116,6 +147,9 @@ public WorkerSourceTask(ConnectorTaskId id,
         this.flushing = false;
         this.stopRequestedLatch = new CountDownLatch(1);
         this.sourceTaskMetricsGroup = new SourceTaskMetricsGroup(id, connectMetrics);
+        this.producerSendException = new AtomicReference<>();
+        this.isTopicTrackingEnabled = workerConfig.getBoolean(TOPIC_TRACKING_ENABLE_CONFIG);
+        this.topicCreation = TopicCreation.newTopicCreation(workerConfig, topicGroups);
     }
 
     @Override
@@ -128,42 +162,67 @@ public void initialize(TaskConfig taskConfig) {
         }
     }
 
+    @Override
     protected void close() {
-        producer.close(30, TimeUnit.SECONDS);
-        transformationChain.close();
+        if (started) {
+            try {
+                task.stop();
+            } catch (Throwable t) {
+                log.warn("Could not stop task", t);
+            }
+        }
+        if (producer != null) {
+            try {
+                producer.close(Duration.ofSeconds(30));
+            } catch (Throwable t) {
+                log.warn("Could not close producer", t);
+            }
+        }
+        if (admin != null) {
+            try {
+                admin.close(Duration.ofSeconds(30));
+            } catch (Throwable t) {
+                log.warn("Failed to close admin client on time", t);
+            }
+        }
+        Utils.closeQuietly(transformationChain, "transformation chain");
+        Utils.closeQuietly(retryWithToleranceOperator, "retry operator");
+    }
+
+    @Override
+    public void removeMetrics() {
+        try {
+            sourceTaskMetricsGroup.close();
+        } finally {
+            super.removeMetrics();
+        }
     }
 
     @Override
-    protected void releaseResources() {
-        sourceTaskMetricsGroup.close();
+    public void cancel() {
+        super.cancel();
+        offsetReader.close();
     }
 
     @Override
     public void stop() {
         super.stop();
         stopRequestedLatch.countDown();
-        synchronized (this) {
-            if (finishedStart)
-                task.stop();
-            else
-                startedShutdownBeforeStartCompleted = true;
-        }
     }
 
     @Override
     public void execute() {
         try {
-            task.initialize(new WorkerSourceTaskContext(offsetReader));
+            // If we try to start the task at all by invoking initialize, then count this as
+            // "started" and expect a subsequent call to the task's stop() method
+            // to properly clean up any resources allocated by its initialize() or 
+            // start() methods. If the task throws an exception during stop(),
+            // the worst thing that happens is another exception gets logged for an already-
+            // failed task
+            started = true;
+            task.initialize(new WorkerSourceTaskContext(offsetReader, this, configState));
             task.start(taskConfig);
             log.info("{} Source task finished initialization and start", this);
-            synchronized (this) {
-                if (startedShutdownBeforeStartCompleted) {
-                    task.stop();
-                    return;
-                }
-                finishedStart = true;
-            }
-
             while (!isStopping()) {
                 if (shouldPause()) {
                     onPause();
@@ -173,17 +232,19 @@ 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();
-                    toSend = task.poll();
+                    toSend = poll();
                     if (toSend != null) {
                         recordPollReturned(toSend.size(), time.milliseconds() - start);
                     }
                 }
                 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);
             }
@@ -198,6 +259,53 @@ 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 | 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;
+        }
+    }
+
+    /**
+     * Convert the source record into a producer record.
+     *
+     * @param record the transformed record
+     * @return the producer record which can sent over to Kafka. A null is returned if the input is null or
+     * if an error was encountered during any of the converter stages.
+     */
+    private ProducerRecord convertTransformedRecord(SourceRecord record) {
+        if (record == null) {
+            return null;
+        }
+
+        RecordHeaders headers = retryWithToleranceOperator.execute(() -> convertHeaderFor(record), Stage.HEADER_CONVERTER, headerConverter.getClass());
+
+        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(), headers, record.valueSchema(), record.value()),
+                Stage.VALUE_CONVERTER, valueConverter.getClass());
+
+        if (retryWithToleranceOperator.failed()) {
+            return null;
+        }
+
+        return new ProducerRecord<>(record.topic(), record.kafkaPartition(),
+                ConnectUtils.checkAndConvertTimestamp(record.timestamp()), key, value, headers);
+    }
+
     /**
      * Try to send a batch of records. If a send fails and is retriable, this saves the remainder of the batch so it can
      * be retried after backing off. If a send fails and is not retriable, this will throw a ConnectException.
@@ -206,21 +314,21 @@ public void execute() {
     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) {
-            final SourceRecord record = transformationChain.apply(preTransformRecord);
+            maybeThrowProducerSendException();
 
-            if (record == null) {
+            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;
             }
 
-            byte[] key = keyConverter.fromConnectData(record.topic(), record.keySchema(), record.key());
-            byte[] value = valueConverter.fromConnectData(record.topic(), record.valueSchema(), record.value());
-            final ProducerRecord producerRecord = new ProducerRecord<>(record.topic(), record.kafkaPartition(),
-                    ConnectUtils.checkAndConvertTimestamp(record.timestamp()), key, value);
-            log.trace("{} Appending record with key {}, value {}", this, record.key(), record.value());
+            log.trace("{} Appending record to the topic {} with key {}, value {}", this, record.topic(), record.key(), record.value());
             // We need this queued first since the callback could happen immediately (even synchronously in some cases).
             // Because of this we need to be careful about handling retries -- we always save the previously attempted
             // record as part of toSend and need to use a flag to track whether we should actually add it to the outstanding
@@ -237,38 +345,41 @@ private boolean sendRecords() {
                 }
             }
             try {
+                maybeCreateTopic(record.topic());
                 final String topic = producerRecord.topic();
                 producer.send(
-                        producerRecord,
-                        new Callback() {
-                            @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 {}: {}", this, topic, e);
-                                    log.debug("{} Failed record: {}", this, preTransformRecord);
-                                } else {
-                                    log.trace("{} Wrote record successfully: topic {} partition {} offset {}",
-                                            this,
-                                            recordMetadata.topic(), recordMetadata.partition(),
-                                            recordMetadata.offset());
-                                    commitTaskRecord(preTransformRecord);
-                                }
-                                recordSent(producerRecord);
-                                counter.completeRecord();
+                    producerRecord,
+                    (recordMetadata, e) -> {
+                        if (e != null) {
+                            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, recordMetadata);
+                            if (isTopicTrackingEnabled) {
+                                recordActiveTopic(producerRecord.topic());
                             }
-                        });
+                        }
+                    });
                 lastSendFailed = false;
-            } catch (RetriableException e) {
-                log.warn("{} Failed to send {}, backing off before retrying:", this, producerRecord, e);
+            } catch (RetriableException | org.apache.kafka.common.errors.RetriableException e) {
+                log.warn("{} Failed to send record to topic '{}' and partition '{}'. Backing off before retrying: ",
+                        this, producerRecord.topic(), producerRecord.partition(), e);
                 toSend = toSend.subList(processed, toSend.size());
                 lastSendFailed = true;
                 counter.retryRemaining();
                 return false;
+            } catch (ConnectException e) {
+                log.warn("{} Failed to send record to topic '{}' and partition '{}' due to an unrecoverable exception: ",
+                        this, producerRecord.topic(), producerRecord.partition(), e);
+                log.warn("{} Failed to send {} with unrecoverable exception: ", this, producerRecord, e);
+                throw e;
             } catch (KafkaException e) {
                 throw new ConnectException("Unrecoverable exception trying to send", e);
             }
@@ -278,9 +389,57 @@ public void onCompletion(RecordMetadata recordMetadata, Exception e) {
         return true;
     }
 
-    private void commitTaskRecord(SourceRecord record) {
+    // Due to transformations that may change the destination topic of a record (such as
+    // RegexRouter) topic creation can not be batched for multiple topics
+    private void maybeCreateTopic(String topic) {
+        if (!topicCreation.isTopicCreationRequired(topic)) {
+            log.trace("Topic creation by the connector is disabled or the topic {} was previously created." +
+                "If auto.create.topics.enable is enabled on the broker, " +
+                "the topic will be created with default settings", topic);
+            return;
+        }
+        log.info("The task will send records to topic '{}' for the first time. Checking "
+                + "whether topic exists", topic);
+        Map existing = admin.describeTopics(topic);
+        if (!existing.isEmpty()) {
+            log.info("Topic '{}' already exists.", topic);
+            topicCreation.addTopic(topic);
+            return;
+        }
+
+        log.info("Creating topic '{}'", topic);
+        TopicCreationGroup topicGroup = topicCreation.findFirstGroup(topic);
+        log.debug("Topic '{}' matched topic creation group: {}", topic, topicGroup);
+        NewTopic newTopic = topicGroup.newTopic(topic);
+
+        if (admin.createTopic(newTopic)) {
+            topicCreation.addTopic(topic);
+            log.info("Created topic '{}' using creation group {}", newTopic, topicGroup);
+        } else {
+            log.warn("Request to create new topic '{}' failed", topic);
+            throw new ConnectException("Task failed to create new topic " + newTopic + ". Ensure "
+                    + "that the task is authorized to create topics or that the topic exists and "
+                    + "restart the task");
+        }
+    }
+
+    private RecordHeaders convertHeaderFor(SourceRecord record) {
+        Headers headers = record.headers();
+        RecordHeaders result = new RecordHeaders();
+        if (headers != null) {
+            String topic = record.topic();
+            for (Header header : headers) {
+                String key = header.key();
+                byte[] rawHeader = headerConverter.fromConnectHeader(topic, key, header.schema(), header.value());
+                result.add(key, rawHeader);
+            }
+        }
+        return result;
+    }
+
+    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);
         }
@@ -303,7 +462,7 @@ private synchronized void recordSent(final ProducerRecord record
     public boolean commitOffsets() {
         long commitTimeoutMs = workerConfig.getLong(WorkerConfig.OFFSET_COMMIT_TIMEOUT_MS_CONFIG);
 
-        log.info("{} Committing offsets", this);
+        log.debug("{} Committing offsets", this);
 
         long started = time.milliseconds();
         long timeout = started + commitTimeoutMs;
@@ -359,14 +518,11 @@ public boolean commitOffsets() {
         }
 
         // Now we can actually flush the offsets to user storage.
-        Future flushFuture = offsetWriter.doFlush(new org.apache.kafka.connect.util.Callback() {
-            @Override
-            public void onCompletion(Throwable error, Void result) {
-                if (error != null) {
-                    log.error("{} Failed to flush offsets to storage: ", this, error);
-                } else {
-                    log.trace("{} Finished flushing offsets to storage", this);
-                }
+        Future flushFuture = offsetWriter.doFlush((error, result) -> {
+            if (error != null) {
+                log.error("{} Failed to flush offsets to storage: ", WorkerSourceTask.this, error);
+            } else {
+                log.trace("{} Finished flushing offsets to storage", WorkerSourceTask.this);
             }
         });
         // Very rare case: offsets were unserializable and we finished immediately, unable to store
@@ -402,7 +558,7 @@ public void onCompletion(Throwable error, Void result) {
         finishSuccessfulFlush();
         long durationMillis = time.milliseconds() - started;
         recordCommitSuccess(durationMillis);
-        log.info("{} Finished commitOffsets successfully in {} ms",
+        log.debug("{} Finished commitOffsets successfully in {} ms",
                 this, durationMillis);
 
         commitSourceTask();
@@ -494,20 +650,22 @@ public SourceTaskMetricsGroup(ConnectorTaskId id, ConnectMetrics connectMetrics)
             metricGroup = connectMetrics.group(registry.sourceTaskGroupName(),
                     registry.connectorTagName(), id.connector(),
                     registry.taskTagName(), Integer.toString(id.task()));
+            // remove any previously created metrics in this group to prevent collisions.
+            metricGroup.close();
 
             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());
             pollTime.add(metricGroup.metricName(registry.sourceRecordPollBatchTimeAvg), new Avg());
 
-            sourceRecordActiveCount = metricGroup.metrics().sensor("sink-record-active-count");
+            sourceRecordActiveCount = metricGroup.sensor("source-record-active-count");
             sourceRecordActiveCount.add(metricGroup.metricName(registry.sourceRecordActiveCount), new Value());
             sourceRecordActiveCount.add(metricGroup.metricName(registry.sourceRecordActiveCountMax), new Max());
             sourceRecordActiveCount.add(metricGroup.metricName(registry.sourceRecordActiveCountAvg), new Avg());
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerSourceTaskContext.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerSourceTaskContext.java
index 8f60e57e005a7..fe1409b282aa0 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerSourceTaskContext.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerSourceTaskContext.java
@@ -16,15 +16,29 @@
  */
 package org.apache.kafka.connect.runtime;
 
+import org.apache.kafka.connect.runtime.distributed.ClusterConfigState;
 import org.apache.kafka.connect.source.SourceTaskContext;
 import org.apache.kafka.connect.storage.OffsetStorageReader;
 
+import java.util.Map;
+
 public class WorkerSourceTaskContext implements SourceTaskContext {
 
     private final OffsetStorageReader reader;
+    private final WorkerSourceTask task;
+    private final ClusterConfigState configState;
 
-    public WorkerSourceTaskContext(OffsetStorageReader reader) {
+    public WorkerSourceTaskContext(OffsetStorageReader reader,
+                                   WorkerSourceTask task,
+                                   ClusterConfigState configState) {
         this.reader = reader;
+        this.task = task;
+        this.configState = configState;
+    }
+
+    @Override
+    public Map configs() {
+        return configState.taskConfig(task.id());
     }
 
     @Override
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 ec069245b3da3..0bd76300dca5e 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
@@ -28,8 +28,11 @@
 import org.apache.kafka.connect.runtime.AbstractStatus.State;
 import org.apache.kafka.connect.runtime.ConnectMetrics.LiteralSupplier;
 import org.apache.kafka.connect.runtime.ConnectMetrics.MetricGroup;
+import org.apache.kafka.connect.runtime.errors.RetryWithToleranceOperator;
 import org.apache.kafka.connect.runtime.isolation.Plugins;
+import org.apache.kafka.connect.storage.StatusBackingStore;
 import org.apache.kafka.connect.util.ConnectorTaskId;
+import org.apache.kafka.connect.util.LoggingContext;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -50,21 +53,29 @@
  */
 abstract class WorkerTask implements Runnable {
     private static final Logger log = LoggerFactory.getLogger(WorkerTask.class);
+    private static final String THREAD_NAME_PREFIX = "task-thread-";
 
     protected final ConnectorTaskId id;
     private final TaskStatus.Listener statusListener;
     protected final ClassLoader loader;
+    protected final StatusBackingStore statusBackingStore;
+    protected final Time time;
     private final CountDownLatch shutdownLatch = new CountDownLatch(1);
     private final TaskMetricsGroup taskMetricsGroup;
     private volatile TargetState targetState;
     private volatile boolean stopping;   // indicates whether the Worker has asked the task to stop
     private volatile boolean cancelled;  // indicates whether the Worker has cancelled the task (e.g. because of slow shutdown)
 
+    protected final RetryWithToleranceOperator retryWithToleranceOperator;
+
     public WorkerTask(ConnectorTaskId id,
                       TaskStatus.Listener statusListener,
                       TargetState initialState,
                       ClassLoader loader,
-                      ConnectMetrics connectMetrics) {
+                      ConnectMetrics connectMetrics,
+                      RetryWithToleranceOperator retryWithToleranceOperator,
+                      Time time,
+                      StatusBackingStore statusBackingStore) {
         this.id = id;
         this.taskMetricsGroup = new TaskMetricsGroup(this.id, connectMetrics, statusListener);
         this.statusListener = taskMetricsGroup;
@@ -73,6 +84,9 @@ public WorkerTask(ConnectorTaskId id,
         this.stopping = false;
         this.cancelled = false;
         this.taskMetricsGroup.recordState(this.targetState);
+        this.retryWithToleranceOperator = retryWithToleranceOperator;
+        this.time = time;
+        this.statusBackingStore = statusBackingStore;
     }
 
     public ConnectorTaskId id() {
@@ -130,16 +144,17 @@ public boolean awaitStop(long timeoutMs) {
         }
     }
 
+    /**
+     * Remove all metrics published by this task.
+     */
+    public void removeMetrics() {
+        taskMetricsGroup.close();
+    }
+
     protected abstract void execute();
 
     protected abstract void close();
 
-    /**
-     * Method called when this worker task has been completely closed, and when the subclass should clean up
-     * all resources.
-     */
-    protected abstract void releaseResources();
-
     protected boolean isStopping() {
         return stopping;
     }
@@ -169,8 +184,7 @@ private void doRun() throws InterruptedException {
 
             execute();
         } catch (Throwable t) {
-            log.error("{} Task threw an uncaught and unrecoverable exception", this, t);
-            log.error("{} Task is being killed and will not recover until manually restarted", this);
+            log.error("{} Task threw an uncaught and unrecoverable exception. Task is being killed and will not recover until manually restarted", this, t);
             throw t;
         } finally {
             doClose();
@@ -209,25 +223,25 @@ protected synchronized void onResume() {
 
     @Override
     public void run() {
-        ClassLoader savedLoader = Plugins.compareAndSwapLoaders(loader);
-        try {
-            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(THREAD_NAME_PREFIX + id);
+                doRun();
+                onShutdown();
+            } catch (Throwable t) {
+                onFailure(t);
+
+                if (t instanceof Error)
+                    throw (Error) t;
+            } finally {
+                Thread.currentThread().setName(savedName);
                 Plugins.compareAndSwapLoaders(savedLoader);
                 shutdownLatch.countDown();
-            } finally {
-                try {
-                    releaseResources();
-                } finally {
-                    taskMetricsGroup.close();
-                }
             }
         }
     }
@@ -264,13 +278,27 @@ public void transitionTo(TargetState state) {
         }
     }
 
+    /**
+     * Include this topic to the set of active topics for the connector that this worker task
+     * is running. This information is persisted in the status backing store used by this worker.
+     *
+     * @param topic the topic to mark as active for this connector
+     */
+    protected void recordActiveTopic(String topic) {
+        if (statusBackingStore.getTopic(id.connector(), topic) != null) {
+            // The topic is already recorded as active. No further action is required.
+            return;
+        }
+        statusBackingStore.put(new TopicStatus(topic, id, time.milliseconds()));
+    }
+
     /**
      * Record that offsets have been committed.
      *
      * @param duration the length of time in milliseconds for the commit attempt to complete
      */
     protected void recordCommitSuccess(long duration) {
-        taskMetricsGroup.recordCommit(duration, true, null);
+        taskMetricsGroup.recordCommit(duration, null);
     }
 
     /**
@@ -280,7 +308,7 @@ protected void recordCommitSuccess(long duration) {
      * @param error the unexpected error that occurred; may be null in the case of timeouts or interruptions
      */
     protected void recordCommitFailure(long duration, Throwable error) {
-        taskMetricsGroup.recordCommit(duration, false, error);
+        taskMetricsGroup.recordCommit(duration, error);
     }
 
     /**
@@ -313,6 +341,8 @@ public TaskMetricsGroup(ConnectorTaskId id, ConnectMetrics connectMetrics, TaskS
             metricGroup = connectMetrics.group(registry.taskGroupName(),
                     registry.connectorTagName(), id.connector(),
                     registry.taskTagName(), Integer.toString(id.task()));
+            // prevent collisions by removing any previously created metrics in this group.
+            metricGroup.close();
 
             metricGroup.addValueMetric(registry.taskStatus, new LiteralSupplier() {
                 @Override
@@ -355,8 +385,8 @@ void close() {
             metricGroup.close();
         }
 
-        void recordCommit(long duration, boolean success, Throwable error) {
-            if (success) {
+        void recordCommit(long duration, Throwable error) {
+            if (error == null) {
                 commitTime.record(duration);
                 commitAttempts.record(1.0d);
             } else {
@@ -398,6 +428,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:
@@ -419,4 +455,4 @@ protected MetricGroup metricGroup() {
             return metricGroup;
         }
     }
-}
\ No newline at end of file
+}
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 cac71ddc73b02..717120d8508ee 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
@@ -16,14 +16,17 @@
  */
 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;
 
 import java.util.ArrayList;
 import java.util.Collections;
-import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
 import java.util.TreeMap;
 
@@ -34,31 +37,55 @@ public class ClusterConfigState {
     public static final long NO_OFFSET = -1;
     public static final ClusterConfigState EMPTY = new ClusterConfigState(
             NO_OFFSET,
-            Collections.emptyMap(),
-            Collections.>emptyMap(),
-            Collections.emptyMap(),
-            Collections.>emptyMap(),
-            Collections.emptySet());
+            null,
+            Collections.emptyMap(),
+            Collections.emptyMap(),
+            Collections.emptyMap(),
+            Collections.emptyMap(),
+            Collections.emptySet());
 
     private final long offset;
+    private final SessionKey sessionKey;
     private final Map connectorTaskCounts;
     private final Map> connectorConfigs;
     private final Map connectorTargetStates;
     private final Map> taskConfigs;
     private final Set inconsistentConnectors;
+    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,
+                taskConfigs,
+                inconsistentConnectors,
+                null);
+    }
+
+    public ClusterConfigState(long offset,
+                              SessionKey sessionKey,
+                              Map connectorTaskCounts,
+                              Map> connectorConfigs,
+                              Map connectorTargetStates,
+                              Map> taskConfigs,
+                              Set inconsistentConnectors,
+                              WorkerConfigTransformer configTransformer) {
         this.offset = offset;
+        this.sessionKey = sessionKey;
         this.connectorTaskCounts = connectorTaskCounts;
         this.connectorConfigs = connectorConfigs;
         this.connectorTargetStates = connectorTargetStates;
         this.taskConfigs = taskConfigs;
         this.inconsistentConnectors = inconsistentConnectors;
+        this.configTransformer = configTransformer;
     }
 
     /**
@@ -70,6 +97,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
@@ -87,11 +122,22 @@ public Set connectors() {
     }
 
     /**
-     * Get the configuration for a connector.
+     * Get the configuration for a connector.  The configuration will have been transformed by
+     * {@link org.apache.kafka.common.config.ConfigTransformer} by having all variable
+     * references replaced with the current values from external instances of
+     * {@link ConfigProvider}, and may include secrets.
      * @param connector name of the connector
      * @return a map containing configuration parameters
      */
     public Map connectorConfig(String connector) {
+        Map configs = connectorConfigs.get(connector);
+        if (configTransformer != null) {
+            configs = configTransformer.transform(connector, configs);
+        }
+        return configs;
+    }
+
+    public Map rawConnectorConfig(String connector) {
         return connectorConfigs.get(connector);
     }
 
@@ -105,26 +151,45 @@ public TargetState targetState(String connector) {
     }
 
     /**
-     * Get the configuration for a task.
+     * Get the configuration for a task.  The configuration will have been transformed by
+     * {@link org.apache.kafka.common.config.ConfigTransformer} by having all variable
+     * references replaced with the current values from external instances of
+     * {@link ConfigProvider}, and may include secrets.
      * @param task id of the task
      * @return a map containing configuration parameters
      */
     public Map taskConfig(ConnectorTaskId task) {
+        Map configs = taskConfigs.get(task);
+        if (configTransformer != null) {
+            configs = configTransformer.transform(task.connector(), configs);
+        }
+        return configs;
+    }
+
+    public Map rawTaskConfig(ConnectorTaskId task) {
         return taskConfigs.get(task);
     }
 
     /**
-     * Get all task configs for a connector.
+     * Get all task configs for a connector.  The configurations will have been transformed by
+     * {@link org.apache.kafka.common.config.ConfigTransformer} by having all variable
+     * references replaced with the current values from external instances of
+     * {@link ConfigProvider}, and may include secrets.
      * @param connector name of the connector
      * @return a list of task configurations
      */
     public List> allTaskConfigs(String connector) {
         Map> taskConfigs = new TreeMap<>();
         for (Map.Entry> taskConfigEntry : this.taskConfigs.entrySet()) {
-            if (taskConfigEntry.getKey().connector().equals(connector))
-                taskConfigs.put(taskConfigEntry.getKey().task(), taskConfigEntry.getValue());
+            if (taskConfigEntry.getKey().connector().equals(connector)) {
+                Map configs = taskConfigEntry.getValue();
+                if (configTransformer != null) {
+                    configs = configTransformer.transform(connector, configs);
+                }
+                taskConfigs.put(taskConfigEntry.getKey().task(), configs);
+            }
         }
-        return new LinkedList<>(taskConfigs.values());
+        return Collections.unmodifiableList(new ArrayList<>(taskConfigs.values()));
     }
 
     /**
@@ -143,19 +208,21 @@ public int taskCount(String connectorName) {
      * @return the current set of connector task IDs
      */
     public List tasks(String connectorName) {
-        if (inconsistentConnectors.contains(connectorName))
+        if (inconsistentConnectors.contains(connectorName)) {
             return Collections.emptyList();
+        }
 
         Integer numTasks = connectorTaskCounts.get(connectorName);
-        if (numTasks == null)
+        if (numTasks == null) {
             return Collections.emptyList();
+        }
 
-        List taskIds = new ArrayList<>();
+        List taskIds = new ArrayList<>(numTasks);
         for (int taskIndex = 0; taskIndex < numTasks; taskIndex++) {
             ConnectorTaskId taskId = new ConnectorTaskId(connectorName, taskIndex);
             taskIds.add(taskId);
         }
-        return taskIds;
+        return Collections.unmodifiableList(taskIds);
     }
 
     /**
@@ -178,10 +245,39 @@ public Set inconsistentConnectors() {
     public String toString() {
         return "ClusterConfigState{" +
                 "offset=" + offset +
+                ", sessionKey=" + (sessionKey != null ? "[hidden]" : "null") +
                 ", connectorTaskCounts=" + connectorTaskCounts +
                 ", connectorConfigs=" + connectorConfigs +
                 ", taskConfigs=" + taskConfigs +
                 ", inconsistentConnectors=" + inconsistentConnectors +
                 '}';
     }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        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) &&
+                Objects.equals(taskConfigs, that.taskConfigs) &&
+                Objects.equals(inconsistentConnectors, that.inconsistentConnectors) &&
+                Objects.equals(configTransformer, that.configTransformer);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(
+                offset,
+                sessionKey,
+                connectorTaskCounts,
+                connectorConfigs,
+                connectorTargetStates,
+                taskConfigs,
+                inconsistentConnectors,
+                configTransformer);
+    }
 }
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..c167e80991fe5 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,25 +383,15 @@ 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);
-                if (connectorTasks == null) {
-                    connectorTasks = new ArrayList<>();
-                    taskMap.put(connectorId, connectorTasks);
-                }
-                connectorTasks.add(CONNECTOR_TASK);
+                taskMap.computeIfAbsent(connectorId, key -> new ArrayList<>()).add(CONNECTOR_TASK);
             }
             for (ConnectorTaskId taskId : taskIds) {
                 String connectorId = taskId.connector();
-                List connectorTasks = taskMap.get(connectorId);
-                if (connectorTasks == null) {
-                    connectorTasks = new ArrayList<>();
-                    taskMap.put(connectorId, connectorTasks);
-                }
-                connectorTasks.add(taskId.task());
+                taskMap.computeIfAbsent(connectorId, key -> new ArrayList<>()).add(taskId.task());
             }
             return taskMap;
         }
@@ -264,5 +404,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..0823fbcc30ad3 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,14 +18,31 @@
 
 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.config.TopicConfig;
+import org.apache.kafka.common.utils.Utils;
 import org.apache.kafka.connect.runtime.WorkerConfig;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
+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;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.PARTITIONS_VALIDATOR;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.REPLICATION_FACTOR_VALIDATOR;
 
 public class DistributedConfig extends WorkerConfig {
-    private static final ConfigDef CONFIG;
+
+    private static final Logger log = LoggerFactory.getLogger(DistributedConfig.class);
 
     /*
      * NOTE: DO NOT CHANGE EITHER CONFIG STRINGS OR THEIR JAVA VARIABLE NAMES AS
@@ -35,13 +52,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 +69,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 +79,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
@@ -84,195 +98,397 @@ public class DistributedConfig extends WorkerConfig {
             " fails to catch up within worker.sync.timeout.ms, leave the Connect cluster for this long before rejoining.";
     public static final int WORKER_UNSYNC_BACKOFF_MS_DEFAULT = 5 * 60 * 1000;
 
+    public static final String CONFIG_STORAGE_PREFIX = "config.storage.";
+    public static final String OFFSET_STORAGE_PREFIX = "offset.storage.";
+    public static final String STATUS_STORAGE_PREFIX = "status.storage.";
+    public static final String TOPIC_SUFFIX = "topic";
+    public static final String PARTITIONS_SUFFIX = "partitions";
+    public static final String REPLICATION_FACTOR_SUFFIX = "replication.factor";
+
     /**
      * offset.storage.topic
      */
-    public static final String OFFSET_STORAGE_TOPIC_CONFIG = "offset.storage.topic";
+    public static final String OFFSET_STORAGE_TOPIC_CONFIG = OFFSET_STORAGE_PREFIX + TOPIC_SUFFIX;
     private static final String OFFSET_STORAGE_TOPIC_CONFIG_DOC = "The name of the Kafka topic where connector offsets are stored";
 
     /**
      * offset.storage.partitions
      */
-    public static final String OFFSET_STORAGE_PARTITIONS_CONFIG = "offset.storage.partitions";
+    public static final String OFFSET_STORAGE_PARTITIONS_CONFIG = OFFSET_STORAGE_PREFIX + PARTITIONS_SUFFIX;
     private static final String OFFSET_STORAGE_PARTITIONS_CONFIG_DOC = "The number of partitions used when creating the offset storage topic";
 
     /**
      * offset.storage.replication.factor
      */
-    public static final String OFFSET_STORAGE_REPLICATION_FACTOR_CONFIG = "offset.storage.replication.factor";
+    public static final String OFFSET_STORAGE_REPLICATION_FACTOR_CONFIG = OFFSET_STORAGE_PREFIX + REPLICATION_FACTOR_SUFFIX;
     private static final String OFFSET_STORAGE_REPLICATION_FACTOR_CONFIG_DOC = "Replication factor used when creating the offset storage topic";
 
     /**
      * config.storage.topic
      */
-    public static final String CONFIG_TOPIC_CONFIG = "config.storage.topic";
+    public static final String CONFIG_TOPIC_CONFIG = CONFIG_STORAGE_PREFIX + TOPIC_SUFFIX;
     private static final String CONFIG_TOPIC_CONFIG_DOC = "The name of the Kafka topic where connector configurations are stored";
 
     /**
      * config.storage.replication.factor
      */
-    public static final String CONFIG_STORAGE_REPLICATION_FACTOR_CONFIG = "config.storage.replication.factor";
+    public static final String CONFIG_STORAGE_REPLICATION_FACTOR_CONFIG = CONFIG_STORAGE_PREFIX + REPLICATION_FACTOR_SUFFIX;
     private static final String CONFIG_STORAGE_REPLICATION_FACTOR_CONFIG_DOC = "Replication factor used when creating the configuration storage topic";
 
     /**
      * status.storage.topic
      */
-    public static final String STATUS_STORAGE_TOPIC_CONFIG = "status.storage.topic";
+    public static final String STATUS_STORAGE_TOPIC_CONFIG = STATUS_STORAGE_PREFIX + TOPIC_SUFFIX;
     public static final String STATUS_STORAGE_TOPIC_CONFIG_DOC = "The name of the Kafka topic where connector and task status are stored";
 
     /**
      * status.storage.partitions
      */
-    public static final String STATUS_STORAGE_PARTITIONS_CONFIG = "status.storage.partitions";
+    public static final String STATUS_STORAGE_PARTITIONS_CONFIG = STATUS_STORAGE_PREFIX + PARTITIONS_SUFFIX;
     private static final String STATUS_STORAGE_PARTITIONS_CONFIG_DOC = "The number of partitions used when creating the status storage topic";
 
     /**
      * status.storage.replication.factor
      */
-    public static final String STATUS_STORAGE_REPLICATION_FACTOR_CONFIG = "status.storage.replication.factor";
+    public static final String STATUS_STORAGE_REPLICATION_FACTOR_CONFIG = STATUS_STORAGE_PREFIX + REPLICATION_FACTOR_SUFFIX;
     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.SOCKET_CONNECTION_SETUP_TIMEOUT_MS_CONFIG,
+                    ConfigDef.Type.LONG,
+                    CommonClientConfigs.DEFAULT_SOCKET_CONNECTION_SETUP_TIMEOUT_MS,
+                    atLeast(0L),
+                    ConfigDef.Importance.LOW,
+                    CommonClientConfigs.SOCKET_CONNECTION_SETUP_TIMEOUT_MS_DOC)
+            .define(CommonClientConfigs.SOCKET_CONNECTION_SETUP_TIMEOUT_MAX_MS_CONFIG,
+                    ConfigDef.Type.LONG,
+                    CommonClientConfigs.DEFAULT_SOCKET_CONNECTION_SETUP_TIMEOUT_MAX_MS,
+                    atLeast(0L),
+                    ConfigDef.Importance.LOW,
+                    CommonClientConfigs.SOCKET_CONNECTION_SETUP_TIMEOUT_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,
+                    PARTITIONS_VALIDATOR,
+                    ConfigDef.Importance.LOW,
+                    OFFSET_STORAGE_PARTITIONS_CONFIG_DOC)
+            .define(OFFSET_STORAGE_REPLICATION_FACTOR_CONFIG,
+                    ConfigDef.Type.SHORT,
+                    (short) 3,
+                    REPLICATION_FACTOR_VALIDATOR,
+                    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,
+                    REPLICATION_FACTOR_VALIDATOR,
+                    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,
+                    PARTITIONS_VALIDATOR,
+                    ConfigDef.Importance.LOW,
+                    STATUS_STORAGE_PARTITIONS_CONFIG_DOC)
+            .define(STATUS_STORAGE_REPLICATION_FACTOR_CONFIG,
+                    ConfigDef.Type.SHORT,
+                    (short) 3,
+                    REPLICATION_FACTOR_VALIDATOR,
+                    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(4, config -> "connectconfigs_" + config));
+    }
+
+    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 Map topicSettings(String prefix) {
+        Map result = originalsWithPrefix(prefix);
+        if (CONFIG_STORAGE_PREFIX.equals(prefix) && result.containsKey(PARTITIONS_SUFFIX)) {
+            log.warn("Ignoring '{}{}={}' setting, since config topic partitions is always 1", prefix, PARTITIONS_SUFFIX, result.get("partitions"));
+        }
+        Object removedPolicy = result.remove(TopicConfig.CLEANUP_POLICY_CONFIG);
+        if (removedPolicy != null) {
+            log.warn("Ignoring '{}cleanup.policy={}' setting, since compaction is always used", prefix, removedPolicy);
+        }
+        result.remove(TOPIC_SUFFIX);
+        result.remove(REPLICATION_FACTOR_SUFFIX);
+        result.remove(PARTITIONS_SUFFIX);
+        return result;
+    }
+
+    public Map configStorageTopicSettings() {
+        return topicSettings(CONFIG_STORAGE_PREFIX);
+    }
+
+    public Map offsetStorageTopicSettings() {
+        return topicSettings(OFFSET_STORAGE_PREFIX);
+    }
+
+    public Map statusStorageTopicSettings() {
+        return topicSettings(STATUS_STORAGE_PREFIX);
+    }
+
+    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 a1cc56a487b9f..6299eb6199d5b 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,35 +16,45 @@
  */
 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.ThreadUtils;
 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;
 import org.apache.kafka.connect.runtime.AbstractHerder;
+import org.apache.kafka.connect.runtime.CloseableConnectorContext;
 import org.apache.kafka.connect.runtime.ConnectMetrics;
 import org.apache.kafka.connect.runtime.ConnectMetrics.LiteralSupplier;
 import org.apache.kafka.connect.runtime.ConnectMetrics.MetricGroup;
 import org.apache.kafka.connect.runtime.ConnectMetricsRegistry;
 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;
@@ -52,8 +62,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;
@@ -69,12 +81,19 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.LinkedBlockingDeque;
-import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.ThreadPoolExecutor;
 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.WorkerConfig.TOPIC_TRACKING_ENABLE_CONFIG;
+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;
 
 /**
  * 
            
@@ -106,10 +125,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();
 
@@ -120,39 +143,58 @@ 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;
     private final ExecutorService startAndStopExecutor;
     private final WorkerGroupMember member;
     private final AtomicBoolean stopping;
+    private final boolean isTopicTrackingEnabled;
 
     // 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;
+    // visible for testing
+    protected ClusterConfigState configState;
 
     // To handle most external requests, like creating or destroying a connector, we can use a generic request where
     // the caller specifies all the code that should be executed.
-    final NavigableSet requests = new ConcurrentSkipListSet<>();
+    final NavigableSet requests = new ConcurrentSkipListSet<>();
     // 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 volatile SecretKey sessionKey;
+    private volatile long keyExpiration;
+    private short currentProtocolVersion;
+    private short backoffRetries;
+
+    private final DistributedConfig config;
 
     public DistributedHerder(DistributedConfig config,
                              Time time,
                              Worker worker,
+                             String kafkaClusterId,
                              StatusBackingStore statusBackingStore,
                              ConfigBackingStore configBackingStore,
-                             String restUrl) {
-        this(config, worker, worker.workerId(), 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());
     }
 
@@ -160,13 +202,15 @@ public DistributedHerder(DistributedConfig config,
     DistributedHerder(DistributedConfig config,
                       Worker worker,
                       String workerId,
+                      String kafkaClusterId,
                       StatusBackingStore statusBackingStore,
                       ConfigBackingStore configBackingStore,
                       WorkerGroupMember member,
                       String restUrl,
                       ConnectMetrics metrics,
-                      Time time) {
-        super(worker, workerId, statusBackingStore, configBackingStore);
+                      Time time,
+                      ConnectorClientConfigOverridePolicy connectorClientConfigOverridePolicy) {
+        super(worker, workerId, kafkaClusterId, statusBackingStore, configBackingStore, connectorClientConfigOverridePolicy);
 
         this.time = time;
         this.herderMetrics = new HerderMetrics(metrics);
@@ -174,22 +218,61 @@ 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.herderExecutor = new ThreadPoolExecutor(1, 1, 0L, TimeUnit.MILLISECONDS, new LinkedBlockingDeque(1),
-                new ThreadFactory() {
-                    @Override
-                    public Thread newThread(Runnable herder) {
-                        return new Thread(herder, "DistributedHerder");
-                    }
-                });
-        this.forwardRequestExecutor = Executors.newSingleThreadExecutor();
-        this.startAndStopExecutor = Executors.newFixedThreadPool(START_STOP_THREAD_POOL_SIZE);
+        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();
+        this.isTopicTrackingEnabled = config.getBoolean(TOPIC_TRACKING_ENABLE_CONFIG);
+
+        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),
+                ThreadUtils.createThreadFactory(
+                        this.getClass().getSimpleName() + "-" + clientId + "-%d", false));
+
+        this.forwardRequestExecutor = Executors.newFixedThreadPool(1,
+                ThreadUtils.createThreadFactory(
+                        "ForwardRequestExecutor-" + clientId + "-%d", false));
+        this.startAndStopExecutor = Executors.newFixedThreadPool(START_STOP_THREAD_POOL_SIZE,
+                ThreadUtils.createThreadFactory(
+                        "StartAndStopExecutor-" + clientId + "-%d", false));
+        this.config = config;
 
         stopping = new AtomicBoolean(false);
         configState = ClusterConfigState.EMPTY;
         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
@@ -205,6 +288,7 @@ public void run() {
             startServices();
 
             log.info("Herder started");
+            running = true;
 
             while (!stopping.get()) {
                 tick();
@@ -217,6 +301,8 @@ public void run() {
         } catch (Throwable t) {
             log.error("Uncaught exception in herder work thread, exiting: ", t);
             Exit.exit(1);
+        } finally {
+            running = false;
         }
     }
 
@@ -230,9 +316,15 @@ public void tick() {
         try {
             // if we failed to read to end of log before, we need to make sure the issue was resolved before joining group
             // Joining and immediately leaving for failure to read configs is exceedingly impolite
-            if (!canReadConfigs && !readConfigToEnd(workerSyncTimeoutMs))
-                return; // Safe to return and tick immediately because readConfigToEnd will do the backoff for us
+            if (!canReadConfigs) {
+                if (readConfigToEnd(workerSyncTimeoutMs)) {
+                    canReadConfigs = true;
+                } else {
+                    return; // Safe to return and tick immediately because readConfigToEnd will do the backoff for us
+                }
+            }
 
+            log.debug("Ensuring group membership is still active");
             member.ensureActive();
             // Ensure we're in a good state in our group. If not restart and everything should be setup to rejoin
             if (!handleRebalanceCompleted()) return;
@@ -240,14 +332,31 @@ public void tick() {
             // May be due to a request from another thread, or might be stopping. If the latter, we need to check the
             // flag immediately. If the former, we need to re-run the ensureActive call since we can't handle requests
             // unless we're in the group.
+            log.trace("Woken up while ensure group membership is still active");
             return;
         }
 
+        long now = time.milliseconds();
+
+        if (checkForKeyRotation(now)) {
+            log.debug("Distributing new session key");
+            keyExpiration = Long.MAX_VALUE;
+            configBackingStore.putSessionKey(new SessionKey(
+                keyGenerator.generateKey(),
+                now
+            ));
+        }
+
         // Process any external requests
-        final long now = time.milliseconds();
+        // TODO: Some of these can be performed concurrently or even optimized away entirely.
+        //       For example, if three different connectors are slated to be restarted, it's fine to
+        //       restart all three at the same time instead.
+        //       Another example: if multiple configurations are submitted for the same connector,
+        //       the only one that actually has to be written to the config topic is the
+        //       most-recently one.
         long nextRequestTimeoutMs = Long.MAX_VALUE;
         while (true) {
-            final HerderRequest next = peekWithoutException();
+            final DistributedHerderRequest next = peekWithoutException();
             if (next == null) {
                 break;
             } else if (now >= next.at) {
@@ -265,79 +374,239 @@ 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 (connectorTargetStateChangesCopy.get() != null) {
+                processTargetStateChanges(connectorTargetStateChangesCopy.get());
+            }
+        } else {
+            shouldReturn = updateConfigsWithIncrementalCooperative(connectorConfigUpdatesCopy,
+                    connectorTargetStateChangesCopy, taskConfigUpdatesCopy);
+
+            if (connectorConfigUpdatesCopy.get() != null) {
+                processConnectorConfigUpdates(connectorConfigUpdatesCopy.get());
             }
-        }
 
-        if (connectorConfigUpdatesCopy != null)
-            processConnectorConfigUpdates(connectorConfigUpdatesCopy);
+            if (connectorTargetStateChangesCopy.get() != null) {
+                processTargetStateChanges(connectorTargetStateChangesCopy.get());
+            }
 
-        if (connectorTargetStateChangesCopy != null)
-            processTargetStateChanges(connectorTargetStateChangesCopy);
+            if (taskConfigUpdatesCopy.get() != null) {
+                processTaskConfigUpdatesWithIncrementalCooperative(taskConfigUpdatesCopy.get());
+            }
+
+            if (shouldReturn) {
+                return;
+            }
+        }
 
         // Let the group take any actions it needs to
         try {
+            log.trace("Polling for group activity; will wait for {}ms or until poll is interrupted by "
+                    + "either config backing store updates or a new external request",
+                    nextRequestTimeoutMs);
             member.poll(nextRequestTimeoutMs);
             // Ensure we're in a good state in our group. If not restart and everything should be setup to rejoin
             handleRebalanceCompleted();
         } catch (WakeupException e) { // FIXME should not be WakeupException
+            log.trace("Woken up while polling for group activity");
             // Ignore. Just indicates we need to check the exit flag, for requested actions, etc.
         }
     }
 
+    private boolean checkForKeyRotation(long now) {
+        SecretKey key;
+        long expiration;
+        synchronized (this) {
+            key = sessionKey;
+            expiration = keyExpiration;
+        }
+
+        if (internalRequestValidationEnabled()) {
+            if (isLeader()) {
+                if (key == null) {
+                    log.debug("Internal request signing is enabled but no session key has been distributed yet. "
+                        + "Distributing new key now.");
+                    return true;
+                } else if (expiration <= now) {
+                    log.debug("Existing key has expired. Distributing new key now.");
+                    return true;
+                } else if (!key.getAlgorithm().equals(keyGenerator.getAlgorithm())
+                        || key.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 (key == 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()) {
+            log.trace("Handling config updates with eager rebalancing");
+            // 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<>();
+                }
+            }
+        } else {
+            log.trace("Skipping config updates with eager rebalancing "
+                + "since no config rebalance is required "
+                + "and there are no connector config, task config, or target state changes pending");
+        }
+        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()) {
+            log.trace("Handling config updates with incremental cooperative rebalancing");
+            // 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<>();
+            }
+        } else {
+            log.trace("Skipping config updates with incremental cooperative rebalancing " 
+                + "since no config rebalance is required " 
+                + "and there are no connector config, task config, or target state changes pending");
+        }
+        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.
         Set localConnectors = assignment == null ? Collections.emptySet() : new HashSet<>(assignment.connectors());
+        log.trace("Processing connector config updates; "
+                + "currently-owned connectors are {}, and to-be-updated connectors are {}",
+                localConnectors,
+                connectorConfigUpdates);
         for (String connectorName : connectorConfigUpdates) {
-            if (!localConnectors.contains(connectorName))
+            if (!localConnectors.contains(connectorName)) {
+                log.trace("Skipping config update for connector {} as it is not owned by this worker",
+                        connectorName);
                 continue;
+            }
             boolean remains = configState.contains(connectorName);
             log.info("Handling connector-only config update by {} connector {}",
                     remains ? "restarting" : "stopping", connectorName);
-            worker.stopConnector(connectorName);
+            worker.stopAndAwaitConnector(connectorName);
             // The update may be a deletion, so verify we actually need to restart the connector
-            if (remains)
-                startConnector(connectorName);
+            if (remains) {
+                startConnector(connectorName, (error, result) -> {
+                    if (error != null) {
+                        log.error("Failed to start connector '" + connectorName + "'", error);
+                    }
+                });
+            }
         }
     }
 
     private void processTargetStateChanges(Set connectorTargetStateChanges) {
+        log.trace("Processing target state updates; "
+                + "currently-known connectors are {}, and to-be-updated connectors are {}",
+                configState.connectors(), connectorTargetStateChanges);
         for (String connector : connectorTargetStateChanges) {
             TargetState targetState = configState.targetState(connector);
             if (!configState.connectors().contains(connector)) {
@@ -347,15 +616,38 @@ private void processTargetStateChanges(Set connectorTargetStateChanges)
 
             // we must propagate the state change to the worker so that the connector's
             // tasks can transition to the new target state
-            worker.setTargetState(connector, targetState);
-
-            // 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);
+            worker.setTargetState(connector, targetState, (error, newState) -> {
+                if (error != null) {
+                    log.error("Failed to transition connector to target state", error);
+                    return;
+                }
+                // 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 (newState == TargetState.STARTED) {
+                    requestTaskReconfiguration(connector);
+                }
+            });
         }
     }
 
+    private void processTaskConfigUpdatesWithIncrementalCooperative(Set taskConfigUpdates) {
+        Set localTasks = assignment == null
+                                          ? Collections.emptySet()
+                                          : new HashSet<>(assignment.tasks());
+        log.trace("Processing task config updates with incremental cooperative rebalance protocol; "
+                + "currently-owned tasks are {}, and to-be-updated tasks are {}",
+                localTasks, taskConfigUpdates);
+        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) {
@@ -374,7 +666,7 @@ public void halt() {
 
             // Explicitly fail any outstanding requests so they actually get a response and get an
             // understandable reason for their failure.
-            HerderRequest request = requests.pollFirst();
+            DistributedHerderRequest request = requests.pollFirst();
             while (request != null) {
                 request.callback().onCompletion(new ConnectException("Worker is shutting down"), null);
                 request = requests.pollFirst();
@@ -398,15 +690,16 @@ public void stop() {
             forwardRequestExecutor.shutdown();
             startAndStopExecutor.shutdown();
 
-            if (!forwardRequestExecutor.awaitTermination(10000L, TimeUnit.MILLISECONDS))
+            if (!forwardRequestExecutor.awaitTermination(FORWARD_REQUEST_SHUTDOWN_TIMEOUT_MS, TimeUnit.MILLISECONDS))
                 forwardRequestExecutor.shutdownNow();
-            if (!startAndStopExecutor.awaitTermination(1000L, TimeUnit.MILLISECONDS))
+            if (!startAndStopExecutor.awaitTermination(START_AND_STOP_SHUTDOWN_TIMEOUT_MS, TimeUnit.MILLISECONDS))
                 startAndStopExecutor.shutdownNow();
         } catch (InterruptedException e) {
             // ignore
         }
 
         log.info("Herder stopped");
+        running = false;
     }
 
     @Override
@@ -442,10 +735,7 @@ public Void call() throws Exception {
                         if (!configState.contains(connName)) {
                             callback.onCompletion(new NotFoundException("Connector " + connName + " not found"), null);
                         } else {
-                            Map config = configState.connectorConfig(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;
                     }
@@ -461,17 +751,8 @@ protected Map config(String connName) {
 
     @Override
     public void connectorConfig(String connName, final Callback> callback) {
-        // Subset of connectorInfo, so piggy back on that implementation
         log.trace("Submitting connector config read request {}", connName);
-        connectorInfo(connName, new Callback() {
-            @Override
-            public void onCompletion(Throwable error, ConnectorInfo result) {
-                if (error != null)
-                    callback.onCompletion(error, null);
-                else
-                    callback.onCompletion(null, result.config());
-            }
-        });
+        super.connectorConfig(connName, callback);
     }
 
     @Override
@@ -482,7 +763,7 @@ public void deleteConnectorConfig(final String connName, final Callback
                 new Callable() {
                     @Override
                     public Void call() throws Exception {
-                        if (maybeAddConfigErrors(validateConnectorConfig(config), callback)) {
-                            return null;
-                        }
-
-                        log.trace("Handling connector config request {}", connName);
-                        if (!isLeader()) {
-                            callback.onCompletion(new NotLeaderException("Only the leader can set connector configs.", leaderUrl()), null);
-                            return null;
-                        }
-
-                        boolean exists = configState.contains(connName);
-                        if (!allowReplace && exists) {
-                            callback.onCompletion(new AlreadyExistsException("Connector " + connName + " already exists"), null);
-                            return null;
-                        }
-
-                        log.trace("Submitting connector config {} {} {}", connName, allowReplace, configState.connectors());
-                        configBackingStore.putConnectorConfig(connName, config);
+                        validateConnectorConfig(config, (error, configInfos) -> {
+                            if (error != null) {
+                                callback.onCompletion(error, null);
+                                return;
+                            }
 
-                        // Note that we use the updated connector config despite the fact that we don't have an updated
-                        // snapshot yet. The existing task info should still be accurate.
-                        Map map = configState.connectorConfig(connName);
-                        ConnectorInfo info = new ConnectorInfo(connName, config, configState.tasks(connName),
-                            map == null ? null : connectorTypeForClass(map.get(ConnectorConfig.CONNECTOR_CLASS_CONFIG)));
-                        callback.onCompletion(null, new Created<>(!exists, info));
+                            // Complete the connector config write via another herder request in order to
+                            // perform the write to the backing store (or forward to the leader) during
+                            // the "external request" portion of the tick loop
+                            addRequest(
+                                    new Callable() {
+                                        @Override
+                                        public Void call() {
+                                            if (maybeAddConfigErrors(configInfos, callback)) {
+                                                return null;
+                                            }
+
+                                            log.trace("Handling connector config request {}", connName);
+                                            if (!isLeader()) {
+                                                callback.onCompletion(new NotLeaderException("Only the leader can set connector configs.", leaderUrl()), null);
+                                                return null;
+                                            }
+                                            boolean exists = configState.contains(connName);
+                                            if (!allowReplace && exists) {
+                                                callback.onCompletion(new AlreadyExistsException("Connector " + connName + " already exists"), null);
+                                                return null;
+                                            }
+
+                                            log.trace("Submitting connector config {} {} {}", connName, allowReplace, configState.connectors());
+                                            configBackingStore.putConnectorConfig(connName, config);
+
+                                            // Note that we use the updated connector config despite the fact that we don't have an updated
+                                            // snapshot yet. The existing task info should still be accurate.
+                                            ConnectorInfo info = new ConnectorInfo(connName, config, configState.tasks(connName),
+                                                // validateConnectorConfig have checked the existence of CONNECTOR_CLASS_CONFIG
+                                                connectorTypeForClass(config.get(ConnectorConfig.CONNECTOR_CLASS_CONFIG)));
+                                            callback.onCompletion(null, new Created<>(!exists, info));
+                                            return null;
+                                        }
+                                    },
+                                    forwardErrorCallback(callback)
+                            );
+                        });
                         return null;
                     }
                 },
@@ -562,22 +861,19 @@ public void requestTaskReconfiguration(final String connName) {
         log.trace("Submitting connector task reconfiguration request {}", connName);
 
         addRequest(
-                new Callable() {
-                    @Override
-                    public Void call() throws Exception {
-                        reconfigureConnectorTasksWithRetry(connName);
-                        return null;
-                    }
-                },
-                new Callback() {
-                    @Override
-                    public void onCompletion(Throwable error, Void result) {
-                        if (error != null) {
-                            log.error("Unexpected error during task reconfiguration: ", error);
-                            log.error("Task reconfiguration for {} failed unexpectedly, this connector will not be properly reconfigured unless manually triggered.", connName);
-                        }
-                    }
+            new Callable() {
+                @Override
+                public Void call() throws Exception {
+                    reconfigureConnectorTasksWithRetry(time.milliseconds(), connName);
+                    return null;
+                }
+            },
+            (error, result) -> {
+                if (error != null) {
+                    log.error("Unexpected error during task reconfiguration: ", error);
+                    log.error("Task reconfiguration for {} failed unexpectedly, this connector will not be properly reconfigured unless manually triggered.", connName);
                 }
+            }
         );
     }
 
@@ -598,7 +894,7 @@ public Void call() throws Exception {
                             List result = new ArrayList<>();
                             for (int i = 0; i < configState.taskCount(connName); i++) {
                                 ConnectorTaskId id = new ConnectorTaskId(connName, i);
-                                result.add(new TaskInfo(id, configState.taskConfig(id)));
+                                result.add(new TaskInfo(id, configState.rawTaskConfig(id)));
                             }
                             callback.onCompletion(null, result);
                         }
@@ -610,8 +906,34 @@ 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 {
+                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() {
@@ -634,7 +956,12 @@ else if (!configState.contains(connName))
 
     @Override
     public void restartConnector(final String connName, final Callback callback) {
-        addRequest(new Callable() {
+        restartConnector(0, connName, callback);
+    }
+
+    @Override
+    public HerderRequest restartConnector(final long delayMs, final String connName, final Callback callback) {
+        return addRequest(delayMs, new Callable() {
             @Override
             public Void call() throws Exception {
                 if (checkRebalanceNeeded(callback))
@@ -647,11 +974,8 @@ public Void call() throws Exception {
 
                 if (assignment.connectors().contains(connName)) {
                     try {
-                        worker.stopConnector(connName);
-                        if (startConnector(connName))
-                            callback.onCompletion(null, null);
-                        else
-                            callback.onCompletion(new ConnectException("Failed to start connector: " + connName), null);
+                        worker.stopAndAwaitConnector(connName);
+                        startConnector(connName, callback);
                     } catch (Throwable t) {
                         callback.onCompletion(t, null);
                     }
@@ -708,7 +1032,6 @@ public int generation() {
         return generation;
     }
 
-
     // Should only be called from work thread, so synchronization should not be needed
     private boolean isLeader() {
         return assignment != null && member.memberId().equals(assignment.leader());
@@ -725,13 +1048,16 @@ 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;
+        }
+        log.debug("Handling completed but unresolved rebalance");
 
         // We need to handle a variety of cases after a rebalance:
         // 1. Assignment failed
@@ -764,12 +1090,23 @@ 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,
             // in which case giving back control to the main loop will prevent hanging around indefinitely after getting kicked out of the group.
             // We also indicate to the main loop that we failed to readConfigs so it will check that the issue was resolved before trying to join the group
-            if (!readConfigToEnd(workerSyncTimeoutMs)) {
+            if (readConfigToEnd(workerSyncTimeoutMs)) {
+                canReadConfigs = true;
+            } else {
                 canReadConfigs = false;
                 needsRejoin = true;
             }
@@ -795,6 +1132,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;
     }
 
@@ -809,19 +1153,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) {
@@ -834,23 +1200,58 @@ 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(
                 taskId,
+                configState,
                 configState.connectorConfig(taskId.connector()),
                 configState.taskConfig(taskId),
                 this,
@@ -886,20 +1287,39 @@ public Void call() throws Exception {
 
     // Helper for starting a connector with the given name, which will extract & parse the config, generate connector
     // context and add to the worker. This needs to be called from within the main worker thread for this herder.
-    private boolean startConnector(String connectorName) {
+    // The callback is invoked after the connector has finished startup and generated task configs, or failed in the process.
+    private void startConnector(String connectorName, Callback callback) {
         log.info("Starting connector {}", connectorName);
         final Map configProps = configState.connectorConfig(connectorName);
-        final ConnectorContext ctx = new HerderConnectorContext(this, connectorName);
+        final CloseableConnectorContext ctx = new HerderConnectorContext(this, connectorName);
         final TargetState initialState = configState.targetState(connectorName);
-        boolean started = worker.startConnector(connectorName, configProps, ctx, this, initialState);
-
-        // Immediately request configuration since this could be a brand new connector. However, also only update those
-        // 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);
+        final Callback onInitialStateChange = (error, newState) -> {
+            if (error != null) {
+                callback.onCompletion(new ConnectException("Failed to start connector: " + connectorName), null);
+                return;
+            }
 
-        return started;
+            // Use newState here in case the connector has been paused right after being created
+            if (newState == TargetState.STARTED) {
+                addRequest(
+                        new Callable() {
+                            @Override
+                            public Void call() {
+                                // Request configuration since this could be a brand new connector. However, also only update those
+                                // task configs if they are actually different from the existing ones to avoid unnecessary updates when this is
+                                // just restoring an existing connector.
+                                reconfigureConnectorTasksWithRetry(time.milliseconds(), connectorName);
+                                callback.onCompletion(null, null);
+                                return null;
+                            }
+                        },
+                        forwardErrorCallback(callback)
+                );
+            } else {
+                callback.onCompletion(null, null);
+            }
+        };
+        worker.startConnector(connectorName, configProps, ctx, this, initialState, onInitialStateChange);
     }
 
     private Callable getConnectorStartingCallable(final String connectorName) {
@@ -907,10 +1327,13 @@ private Callable getConnectorStartingCallable(final String connectorName)
             @Override
             public Void call() throws Exception {
                 try {
-                    startConnector(connectorName);
+                    startConnector(connectorName, (error, result) -> {
+                        if (error != null) {
+                            log.error("Failed to start connector '" + connectorName + "'", error);
+                        }
+                    });
                 } catch (Throwable t) {
-                    log.error("Couldn't instantiate connector " + connectorName + " because it has an invalid connector " +
-                            "configuration. This connector will not execute until reconfigured.", t);
+                    log.error("Unexpected error while trying to start connector " + connectorName, t);
                     onFailure(connectorName, t);
                 }
                 return null;
@@ -923,7 +1346,7 @@ private Callable getConnectorStoppingCallable(final String connectorName)
             @Override
             public Void call() throws Exception {
                 try {
-                    worker.stopConnector(connectorName);
+                    worker.stopAndAwaitConnector(connectorName);
                 } catch (Throwable t) {
                     log.error("Failed to shut down connector " + connectorName, t);
                 }
@@ -932,28 +1355,34 @@ 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) {
                 // If we encountered an error, we don't have much choice but to just retry. If we don't, we could get
                 // stuck with a connector that thinks it has generated tasks, but wasn't actually successful and therefore
-                // never makes progress. The retry has to run through a HerderRequest since this callback could be happening
+                // 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", connName);
+                    } else {
+                        log.error("Failed to reconfigure connector's tasks ({}), retrying after backoff:", connName, 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);
+                                    }
                                 }
                             }
                     );
@@ -962,6 +1391,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) {
@@ -977,7 +1415,7 @@ private void reconfigureConnector(final String connName, final Callback cb
             if (worker.isSinkConnector(connName)) {
                 connConfig = new SinkConnectorConfig(plugins(), configs);
             } else {
-                connConfig = new SourceConnectorConfig(plugins(), configs);
+                connConfig = new SourceConnectorConfig(plugins(), configs, worker.isTopicCreationEnabled());
             }
 
             final List> taskProps = worker.connectorTaskConfigs(connName, connConfig);
@@ -998,23 +1436,29 @@ private void reconfigureConnector(final String connName, final Callback cb
                 }
             }
             if (changed) {
+                List> rawTaskProps = reverseTransform(connName, configState, taskProps);
                 if (isLeader()) {
-                    configBackingStore.putTaskConfigs(connName, taskProps);
+                    configBackingStore.putTaskConfigs(connName, rawTaskProps);
                     cb.onCompletion(null, null);
                 } else {
                     // We cannot forward the request on the same thread because this reconfiguration can happen as a result of connector
                     // addition or removal. If we blocked waiting for the response from leader, we may be kicked out of the worker group.
-                    forwardRequestExecutor.submit(new Runnable() {
-                        @Override
-                        public void run() {
-                            try {
-                                String reconfigUrl = RestServer.urlJoin(leaderUrl(), "/connectors/" + connName + "/tasks");
-                                RestServer.httpRequest(reconfigUrl, "POST", taskProps, null);
-                                cb.onCompletion(null, null);
-                            } catch (ConnectException e) {
-                                log.error("Request to leader to reconfigure connector tasks failed", e);
-                                cb.onCompletion(e, null);
+                    forwardRequestExecutor.submit(() -> {
+                        try {
+                            String leaderUrl = leaderUrl();
+                            if (leaderUrl == null || leaderUrl.trim().isEmpty()) {
+                                cb.onCompletion(new ConnectException("Request to leader to " +
+                                        "reconfigure connector tasks failed " +
+                                        "because the URL of the leader's REST interface is empty!"), null);
+                                return;
                             }
+                            String reconfigUrl = RestServer.urlJoin(leaderUrl, "/connectors/" + connName + "/tasks");
+                            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);
+                            cb.onCompletion(e, null);
                         }
                     });
                 }
@@ -1034,19 +1478,27 @@ private boolean checkRebalanceNeeded(Callback callback) {
         return false;
     }
 
-    HerderRequest addRequest(Callable action, Callback callback) {
+    DistributedHerderRequest addRequest(Callable action, Callback callback) {
         return addRequest(0, action, callback);
     }
 
-    HerderRequest addRequest(long delayMs, Callable action, Callback callback) {
-        HerderRequest req = new HerderRequest(time.milliseconds() + delayMs, requestSeqNum.incrementAndGet(), action, callback);
+    DistributedHerderRequest addRequest(long delayMs, Callable action, Callback callback) {
+        DistributedHerderRequest req = new DistributedHerderRequest(time.milliseconds() + delayMs, requestSeqNum.incrementAndGet(), action, callback);
         requests.add(req);
         if (peekWithoutException() == req)
             member.wakeup();
         return req;
     }
 
-    private HerderRequest peekWithoutException() {
+    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();
         } catch (NoSuchElementException e) {
@@ -1089,12 +1541,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();
         }
@@ -1108,15 +1565,30 @@ 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 leader and won't be responsible for
+                // tracking expiration and distributing new keys themselves
+                if (isLeader() && keyRotationIntervalMs > 0) {
+                    DistributedHerder.this.keyExpiration = sessionKey.creationTimestamp() + keyRotationIntervalMs;
+                }
+            }
+        }
     }
 
-    static class HerderRequest implements Comparable {
+    class DistributedHerderRequest implements HerderRequest, Comparable {
         private final long at;
         private final long seq;
         private final Callable action;
         private final Callback callback;
 
-        public HerderRequest(long at, long seq, Callable action, Callback callback) {
+        public DistributedHerderRequest(long at, long seq, Callable action, Callback callback) {
             this.at = at;
             this.seq = seq;
             this.action = action;
@@ -1132,7 +1604,12 @@ public Callback callback() {
         }
 
         @Override
-        public int compareTo(HerderRequest o) {
+        public void cancel() {
+            DistributedHerder.this.requests.remove(this);
+        }
+
+        @Override
+        public int compareTo(DistributedHerderRequest o) {
             final int cmp = Long.compare(at, o.at);
             return cmp == 0 ? Long.compare(seq, o.seq) : cmp;
         }
@@ -1140,9 +1617,9 @@ public int compareTo(HerderRequest o) {
         @Override
         public boolean equals(Object o) {
             if (this == o) return true;
-            if (!(o instanceof HerderRequest))
+            if (!(o instanceof DistributedHerderRequest))
                 return false;
-            HerderRequest other = (HerderRequest) o;
+            DistributedHerderRequest other = (DistributedHerderRequest) o;
             return compareTo(other) == 0;
         }
 
@@ -1152,13 +1629,10 @@ public int hashCode() {
         }
     }
 
-    private static final Callback forwardErrorCallback(final Callback callback) {
-        return new Callback() {
-            @Override
-            public void onCompletion(Throwable error, Void result) {
-                if (error != null)
-                    callback.onCompletion(error, null);
-            }
+    private static Callback forwardErrorCallback(final Callback callback) {
+        return (error, result) -> {
+            if (error != null)
+                callback.onCompletion(error, null);
         };
     }
 
@@ -1173,32 +1647,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
@@ -1208,16 +1737,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));
@@ -1231,18 +1754,44 @@ public void onRevoked(String leader, Collection connectors, Collection connectors, Collection tasks) {
+            connectors.stream()
+                    .filter(connectorName -> !configState.contains(connectorName))
+                    .forEach(DistributedHerder.this::resetConnectorActiveTopics);
+            tasks.stream()
+                    .map(ConnectorTaskId::connector)
+                    .distinct()
+                    .filter(connectorName -> !configState.contains(connectorName))
+                    .forEach(DistributedHerder.this::resetConnectorActiveTopics);
+        }
     }
 
     class HerderMetrics {
@@ -1257,6 +1806,12 @@ public HerderMetrics(ConnectMetrics connectMetrics) {
             ConnectMetricsRegistry registry = connectMetrics.registry();
             metricGroup = connectMetrics.group(registry.workerRebalanceGroupName());
 
+            metricGroup.addValueMetric(registry.connectProtocol, new LiteralSupplier() {
+                @Override
+                public String metricValue(long now) {
+                    return ConnectProtocolCompatibility.fromProtocolVersion(member.currentProtocolVersion()).name();
+                }
+            });
             metricGroup.addValueMetric(registry.leaderName, new LiteralSupplier() {
                 @Override
                 public String metricValue(long now) {
@@ -1277,7 +1832,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 600772851f0f4..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,14 +18,14 @@
 
 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;
 import org.apache.kafka.connect.storage.ConfigBackingStore;
 import org.apache.kafka.connect.util.ConnectorTaskId;
 import org.slf4j.Logger;
@@ -35,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;
@@ -91,8 +96,15 @@ 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
     public void requestRejoin() {
         rejoinRequested = true;
     }
@@ -102,6 +114,12 @@ public String protocolType() {
         return "connect";
     }
 
+    // expose for tests
+    @Override
+    protected synchronized boolean ensureCoordinatorReady(final Timer timer) {
+        return super.ensureCoordinatorReady(timer);
+    }
+
     public void poll(long timeout) {
         // poll for io until the timeout expires
         final long start = time.milliseconds();
@@ -110,11 +128,24 @@ public void poll(long timeout) {
 
         do {
             if (coordinatorUnknown()) {
-                ensureCoordinatorReady();
+                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();
             }
 
-            if (needRejoin()) {
+            if (rejoinNeededOrPending()) {
                 ensureActiveGroup();
                 now = time.milliseconds();
             }
@@ -126,7 +157,8 @@ public void poll(long timeout) {
 
             // Note that because the network client is shared with the background heartbeat thread,
             // we do not want to block in poll longer than the time to the next heartbeat.
-            client.poll(Math.min(Math.max(0, remaining), timeToNextHeartbeat(now)));
+            long pollTimeout = Math.min(Math.max(0, remaining), timeToNextHeartbeat(now));
+            client.poll(time.timer(pollTimeout));
 
             now = time.milliseconds();
             elapsed = now - start;
@@ -135,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 needRejoin() {
-        return super.needRejoin() || (assignmentSnapshot == null || assignmentSnapshot.failed()) || rejoinRequested;
+    protected boolean rejoinNeededOrPending() {
+        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 (needRejoin() || !isLeader())
+        if (rejoinNeededOrPending() || !isLeader())
             return null;
-        return leaderState.ownerUrl(connector);
+        return leaderState().ownerUrl(connector);
     }
 
     public String ownerUrl(ConnectorTaskId task) {
-        if (needRejoin() || !isLeader())
+        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 {
@@ -315,14 +351,24 @@ public WorkerCoordinatorMetrics(Metrics metrics, String metricGrpPrefix) {
             this.metricGrpName = metricGrpPrefix + "-coordinator-metrics";
 
             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();
                 }
             };
 
@@ -335,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);
@@ -351,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);
@@ -380,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 525ce7e2de93d..2a9a9f2a2bac6 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
@@ -17,14 +17,18 @@
 package org.apache.kafka.connect.runtime.distributed;
 
 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.Metadata;
 import org.apache.kafka.clients.NetworkClient;
 import org.apache.kafka.clients.consumer.internals.ConsumerNetworkClient;
-import org.apache.kafka.common.Cluster;
+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;
+import org.apache.kafka.common.metrics.MetricsContext;
+import org.apache.kafka.common.metrics.KafkaMetricsContext;
 import org.apache.kafka.common.metrics.MetricConfig;
 import org.apache.kafka.common.metrics.Metrics;
 import org.apache.kafka.common.metrics.MetricsReporter;
@@ -33,17 +37,20 @@
 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.runtime.WorkerConfig;
 import org.apache.kafka.connect.storage.ConfigBackingStore;
+import org.apache.kafka.connect.util.ConnectUtils;
 import org.apache.kafka.connect.util.ConnectorTaskId;
 import org.slf4j.Logger;
 
 import java.net.InetSocketAddress;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.LinkedHashMap;
 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;
 
 /**
@@ -53,8 +60,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;
@@ -72,15 +77,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<>();
@@ -88,15 +90,29 @@ public WorkerGroupMember(DistributedConfig config,
             MetricConfig metricConfig = new MetricConfig().samples(config.getInt(CommonClientConfigs.METRICS_NUM_SAMPLES_CONFIG))
                     .timeWindow(config.getLong(CommonClientConfigs.METRICS_SAMPLE_WINDOW_MS_CONFIG), TimeUnit.MILLISECONDS)
                     .tags(metricsTags);
-            List reporters = config.getConfiguredInstances(CommonClientConfigs.METRIC_REPORTER_CLASSES_CONFIG, MetricsReporter.class);
-            reporters.add(new JmxReporter(JMX_PREFIX));
-            this.metrics = new Metrics(metricConfig, reporters, time);
+            List reporters = config.getConfiguredInstances(CommonClientConfigs.METRIC_REPORTER_CLASSES_CONFIG,
+                    MetricsReporter.class,
+                    Collections.singletonMap(CommonClientConfigs.CLIENT_ID_CONFIG, clientId));
+            JmxReporter jmxReporter = new JmxReporter();
+            jmxReporter.configure(config.originals());
+            reporters.add(jmxReporter);
+
+            Map contextLabels = new HashMap<>();
+            contextLabels.putAll(config.originalsWithPrefix(CommonClientConfigs.METRICS_CONTEXT_PREFIX));
+            contextLabels.put(WorkerConfig.CONNECT_KAFKA_CLUSTER_ID, ConnectUtils.lookupKafkaClusterId(config));
+            contextLabels.put(WorkerConfig.CONNECT_GROUP_ID, config.getString(DistributedConfig.GROUP_ID_CONFIG));
+            MetricsContext metricsContext = new KafkaMetricsContext(JMX_PREFIX, contextLabels);
+
+            this.metrics = new Metrics(metricConfig, reporters, time, metricsContext);
             this.retryBackoffMs = config.getLong(CommonClientConfigs.RETRY_BACKOFF_MS_CONFIG);
-            this.metadata = new Metadata(retryBackoffMs, config.getLong(CommonClientConfigs.METADATA_MAX_AGE_CONFIG), true);
-            List addresses = ClientUtils.parseAndValidateAddresses(config.getList(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG));
-            this.metadata.update(Cluster.bootstrap(addresses), Collections.emptySet(), 0);
+            this.metadata = new Metadata(retryBackoffMs, config.getLong(CommonClientConfigs.METADATA_MAX_AGE_CONFIG),
+                    logContext, new ClusterResourceListeners());
+            List addresses = ClientUtils.parseAndValidateAddresses(
+                    config.getList(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG),
+                    config.getString(CommonClientConfigs.CLIENT_DNS_LOOKUP_CONFIG));
+            this.metadata.bootstrap(addresses);
             String metricGrpPrefix = "connect";
-            ChannelBuilder channelBuilder = ClientUtils.createChannelBuilder(config);
+            ChannelBuilder channelBuilder = ClientUtils.createChannelBuilder(config, time, logContext);
             NetworkClient netClient = new NetworkClient(
                     new Selector(config.getLong(CommonClientConfigs.CONNECTIONS_MAX_IDLE_MS_CONFIG), metrics, time, metricGrpPrefix, channelBuilder, logContext),
                     this.metadata,
@@ -107,6 +123,9 @@ public WorkerGroupMember(DistributedConfig config,
                     config.getInt(CommonClientConfigs.SEND_BUFFER_CONFIG),
                     config.getInt(CommonClientConfigs.RECEIVE_BUFFER_CONFIG),
                     config.getInt(CommonClientConfigs.REQUEST_TIMEOUT_MS_CONFIG),
+                    config.getLong(CommonClientConfigs.SOCKET_CONNECTION_SETUP_TIMEOUT_MS_CONFIG),
+                    config.getLong(CommonClientConfigs.SOCKET_CONNECTION_SETUP_TIMEOUT_MAX_MS_CONFIG),
+                    ClientDnsLookup.forConfig(config.getString(CommonClientConfigs.CLIENT_DNS_LOOKUP_CONFIG)),
                     time,
                     true,
                     new ApiVersions(),
@@ -120,21 +139,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
@@ -150,6 +167,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);
     }
@@ -182,8 +203,8 @@ public void requestRejoin() {
         coordinator.requestRejoin();
     }
 
-    public void maybeLeaveGroup() {
-        coordinator.maybeLeaveGroup();
+    public void maybeLeaveGroup(String leaveReason) {
+        coordinator.maybeLeaveGroup(leaveReason);
     }
 
     public String ownerUrl(String connector) {
@@ -194,17 +215,35 @@ 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());
         else
             log.debug("The Connect group member has stopped.");
     }
+
+    // Visible for testing
+    Metrics metrics() {
+        return this.metrics;
+    }
 }
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
new file mode 100644
index 0000000000000..e07a5d01184a1
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/DeadLetterQueueReporter.java
@@ -0,0 +1,216 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.connect.runtime.errors;
+
+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;
+import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.clients.producer.RecordMetadata;
+import org.apache.kafka.common.errors.TopicExistsException;
+import org.apache.kafka.common.header.Headers;
+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.util.ConnectorTaskId;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.PrintStream;
+import java.nio.charset.StandardCharsets;
+import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Future;
+
+import static java.util.Collections.singleton;
+
+/**
+ * Write the original consumed record into a dead letter queue. The dead letter queue is a Kafka topic located
+ * on the same cluster used by the worker to maintain internal topics. Each connector is typically configured
+ * with its own Kafka topic dead letter queue. By default, the topic name is not set, and if the
+ * connector config doesn't specify one, this feature is disabled.
+ */
+public class DeadLetterQueueReporter implements ErrorReporter {
+
+    private static final Logger log = LoggerFactory.getLogger(DeadLetterQueueReporter.class);
+
+    private static final int DLQ_NUM_DESIRED_PARTITIONS = 1;
+
+    public static final String HEADER_PREFIX = "__connect.errors.";
+    public static final String ERROR_HEADER_ORIG_TOPIC = HEADER_PREFIX + "topic";
+    public static final String ERROR_HEADER_ORIG_PARTITION = HEADER_PREFIX + "partition";
+    public static final String ERROR_HEADER_ORIG_OFFSET = HEADER_PREFIX + "offset";
+    public static final String ERROR_HEADER_CONNECTOR_NAME = HEADER_PREFIX + "connector.name";
+    public static final String ERROR_HEADER_TASK_ID = HEADER_PREFIX + "task.id";
+    public static final String ERROR_HEADER_STAGE = HEADER_PREFIX + "stage";
+    public static final String ERROR_HEADER_EXECUTING_CLASS = HEADER_PREFIX + "class.name";
+    public static final String ERROR_HEADER_EXCEPTION = HEADER_PREFIX + "exception.class.name";
+    public static final String ERROR_HEADER_EXCEPTION_MESSAGE = HEADER_PREFIX + "exception.message";
+    public static final String ERROR_HEADER_EXCEPTION_STACK_TRACE = HEADER_PREFIX + "exception.stacktrace";
+
+    private final SinkConnectorConfig connConfig;
+    private final ConnectorTaskId connectorTaskId;
+    private final ErrorHandlingMetrics errorHandlingMetrics;
+    private final String dlqTopicName;
+
+    private KafkaProducer kafkaProducer;
+
+    public static DeadLetterQueueReporter createAndSetup(Map adminProps,
+                                                         ConnectorTaskId id,
+                                                         SinkConnectorConfig sinkConfig, Map producerProps,
+                                                         ErrorHandlingMetrics errorHandlingMetrics) {
+        String topic = sinkConfig.dlqTopicName();
+
+        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());
+                admin.createTopics(singleton(schemaTopicRequest)).all().get();
+            }
+        } catch (InterruptedException e) {
+            throw new ConnectException("Could not initialize dead letter queue with topic=" + topic, e);
+        } catch (ExecutionException e) {
+            if (!(e.getCause() instanceof TopicExistsException)) {
+                throw new ConnectException("Could not initialize dead letter queue with topic=" + topic, e);
+            }
+        }
+
+        KafkaProducer dlqProducer = new KafkaProducer<>(producerProps);
+        return new DeadLetterQueueReporter(dlqProducer, sinkConfig, id, errorHandlingMetrics);
+    }
+
+    /**
+     * Initialize the dead letter queue reporter with a {@link KafkaProducer}.
+     *
+     * @param kafkaProducer a Kafka Producer to produce the original consumed records.
+     */
+    // Visible for testing
+    DeadLetterQueueReporter(KafkaProducer kafkaProducer, SinkConnectorConfig connConfig,
+                            ConnectorTaskId id, ErrorHandlingMetrics errorHandlingMetrics) {
+        Objects.requireNonNull(kafkaProducer);
+        Objects.requireNonNull(connConfig);
+        Objects.requireNonNull(id);
+        Objects.requireNonNull(errorHandlingMetrics);
+
+        this.kafkaProducer = kafkaProducer;
+        this.connConfig = connConfig;
+        this.connectorTaskId = id;
+        this.errorHandlingMetrics = errorHandlingMetrics;
+        this.dlqTopicName = connConfig.dlqTopicName().trim();
+    }
+
+    /**
+     * Write the raw records into a Kafka topic and return the producer future.
+     *
+     * @param context processing context containing the raw record at {@link ProcessingContext#consumerRecord()}.
+     * @return the future associated with the writing of this record; never null
+     */
+    public Future report(ProcessingContext context) {
+        if (dlqTopicName.isEmpty()) {
+            return CompletableFuture.completedFuture(null);
+        }
+        errorHandlingMetrics.recordDeadLetterQueueProduceRequest();
+
+        ConsumerRecord originalMessage = context.consumerRecord();
+        if (originalMessage == null) {
+            errorHandlingMetrics.recordDeadLetterQueueProduceFailed();
+            return CompletableFuture.completedFuture(null);
+        }
+
+        ProducerRecord producerRecord;
+        if (originalMessage.timestamp() == RecordBatch.NO_TIMESTAMP) {
+            producerRecord = new ProducerRecord<>(dlqTopicName, null,
+                    originalMessage.key(), originalMessage.value(), originalMessage.headers());
+        } else {
+            producerRecord = new ProducerRecord<>(dlqTopicName, null, originalMessage.timestamp(),
+                    originalMessage.key(), originalMessage.value(), originalMessage.headers());
+        }
+
+        if (connConfig.isDlqContextHeadersEnabled()) {
+            populateContextHeaders(producerRecord, context);
+        }
+
+        return this.kafkaProducer.send(producerRecord, (metadata, exception) -> {
+            if (exception != null) {
+                log.error("Could not produce message to dead letter queue. topic=" + dlqTopicName, exception);
+                errorHandlingMetrics.recordDeadLetterQueueProduceFailed();
+            }
+        });
+    }
+
+    // Visible for testing
+    void populateContextHeaders(ProducerRecord producerRecord, ProcessingContext context) {
+        Headers headers = producerRecord.headers();
+        if (context.consumerRecord() != null) {
+            headers.add(ERROR_HEADER_ORIG_TOPIC, toBytes(context.consumerRecord().topic()));
+            headers.add(ERROR_HEADER_ORIG_PARTITION, toBytes(context.consumerRecord().partition()));
+            headers.add(ERROR_HEADER_ORIG_OFFSET, toBytes(context.consumerRecord().offset()));
+        }
+
+        headers.add(ERROR_HEADER_CONNECTOR_NAME, toBytes(connectorTaskId.connector()));
+        headers.add(ERROR_HEADER_TASK_ID, toBytes(String.valueOf(connectorTaskId.task())));
+        headers.add(ERROR_HEADER_STAGE, toBytes(context.stage().name()));
+        headers.add(ERROR_HEADER_EXECUTING_CLASS, toBytes(context.executingClass().getName()));
+        if (context.error() != null) {
+            headers.add(ERROR_HEADER_EXCEPTION, toBytes(context.error().getClass().getName()));
+            headers.add(ERROR_HEADER_EXCEPTION_MESSAGE, toBytes(context.error().getMessage()));
+            byte[] trace;
+            if ((trace = stacktrace(context.error())) != null) {
+                headers.add(ERROR_HEADER_EXCEPTION_STACK_TRACE, trace);
+            }
+        }
+    }
+
+    private byte[] stacktrace(Throwable error) {
+        ByteArrayOutputStream bos = new ByteArrayOutputStream();
+        try {
+            PrintStream stream = new PrintStream(bos, true, "UTF-8");
+            error.printStackTrace(stream);
+            bos.close();
+            return bos.toByteArray();
+        } catch (IOException e) {
+            log.error("Could not serialize stacktrace.", e);
+        }
+        return null;
+    }
+
+    private byte[] toBytes(int value) {
+        return toBytes(String.valueOf(value));
+    }
+
+    private byte[] toBytes(long value) {
+        return toBytes(String.valueOf(value));
+    }
+
+    private byte[] toBytes(String value) {
+        if (value != null) {
+            return value.getBytes(StandardCharsets.UTF_8);
+        } else {
+            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
new file mode 100644
index 0000000000000..419bea97f1200
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/ErrorHandlingMetrics.java
@@ -0,0 +1,141 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.errors;
+
+import org.apache.kafka.common.metrics.Sensor;
+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;
+import org.apache.kafka.connect.runtime.ConnectMetricsRegistry;
+import org.apache.kafka.connect.util.ConnectorTaskId;
+
+/**
+ * Contains various sensors used for monitoring errors.
+ */
+public class ErrorHandlingMetrics {
+
+    private final Time time = new SystemTime();
+
+    private final ConnectMetrics.MetricGroup metricGroup;
+
+    // metrics
+    private final Sensor recordProcessingFailures;
+    private final Sensor recordProcessingErrors;
+    private final Sensor recordsSkipped;
+    private final Sensor retries;
+    private final Sensor errorsLogged;
+    private final Sensor dlqProduceRequests;
+    private final Sensor dlqProduceFailures;
+    private long lastErrorTime = 0;
+
+    public ErrorHandlingMetrics(ConnectorTaskId id, ConnectMetrics connectMetrics) {
+
+        ConnectMetricsRegistry registry = connectMetrics.registry();
+        metricGroup = connectMetrics.group(registry.taskErrorHandlingGroupName(),
+                registry.connectorTagName(), id.connector(), registry.taskTagName(), Integer.toString(id.task()));
+
+        // prevent collisions by removing any previously created metrics in this group.
+        metricGroup.close();
+
+        recordProcessingFailures = metricGroup.sensor("total-record-failures");
+        recordProcessingFailures.add(metricGroup.metricName(registry.recordProcessingFailures), new CumulativeSum());
+
+        recordProcessingErrors = metricGroup.sensor("total-record-errors");
+        recordProcessingErrors.add(metricGroup.metricName(registry.recordProcessingErrors), new CumulativeSum());
+
+        recordsSkipped = metricGroup.sensor("total-records-skipped");
+        recordsSkipped.add(metricGroup.metricName(registry.recordsSkipped), new CumulativeSum());
+
+        retries = metricGroup.sensor("total-retries");
+        retries.add(metricGroup.metricName(registry.retries), new CumulativeSum());
+
+        errorsLogged = metricGroup.sensor("total-errors-logged");
+        errorsLogged.add(metricGroup.metricName(registry.errorsLogged), new CumulativeSum());
+
+        dlqProduceRequests = metricGroup.sensor("deadletterqueue-produce-requests");
+        dlqProduceRequests.add(metricGroup.metricName(registry.dlqProduceRequests), new CumulativeSum());
+
+        dlqProduceFailures = metricGroup.sensor("deadletterqueue-produce-failures");
+        dlqProduceFailures.add(metricGroup.metricName(registry.dlqProduceFailures), new CumulativeSum());
+
+        metricGroup.addValueMetric(registry.lastErrorTimestamp, now -> lastErrorTime);
+    }
+
+    /**
+     * Increment the number of failed operations (retriable and non-retriable).
+     */
+    public void recordFailure() {
+        recordProcessingFailures.record();
+    }
+
+    /**
+     * Increment the number of operations which could not be successfully executed.
+     */
+    public void recordError() {
+        recordProcessingErrors.record();
+    }
+
+    /**
+     * Increment the number of records skipped.
+     */
+    public void recordSkipped() {
+        recordsSkipped.record();
+    }
+
+    /**
+     * The number of retries made while executing operations.
+     */
+    public void recordRetry() {
+        retries.record();
+    }
+
+    /**
+     * The number of errors logged by the {@link LogReporter}.
+     */
+    public void recordErrorLogged() {
+        errorsLogged.record();
+    }
+
+    /**
+     * The number of produce requests to the {@link DeadLetterQueueReporter}.
+     */
+    public void recordDeadLetterQueueProduceRequest() {
+        dlqProduceRequests.record();
+    }
+
+    /**
+     * The number of produce requests to the {@link DeadLetterQueueReporter} which failed to be successfully produced into Kafka.
+     */
+    public void recordDeadLetterQueueProduceFailed() {
+        dlqProduceFailures.record();
+    }
+
+    /**
+     * Record the time of error.
+     */
+    public void recordErrorTimestamp() {
+        this.lastErrorTime = time.milliseconds();
+    }
+
+    /**
+     * @return the metric group for this class.
+     */
+    public ConnectMetrics.MetricGroup metricGroup() {
+        return metricGroup;
+    }
+}
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
new file mode 100644
index 0000000000000..f9bc2f2360daa
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/ErrorReporter.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.runtime.errors;
+
+import org.apache.kafka.clients.producer.RecordMetadata;
+
+import java.util.concurrent.Future;
+
+/**
+ * Report an error using the information contained in the {@link ProcessingContext}.
+ */
+public interface ErrorReporter extends AutoCloseable {
+
+    /**
+     * Report an error and return the producer future.
+     *
+     * @param context the processing context (cannot be null).
+     * @return future result from the producer sending a record to Kafka.
+     */
+    Future report(ProcessingContext context);
+
+    @Override
+    default void close() { }
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/LogReporter.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/LogReporter.java
new file mode 100644
index 0000000000000..cf9db2c44d329
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/LogReporter.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.runtime.errors;
+
+import org.apache.kafka.clients.producer.RecordMetadata;
+import org.apache.kafka.connect.runtime.ConnectorConfig;
+import org.apache.kafka.connect.util.ConnectorTaskId;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.Objects;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Future;
+
+/**
+ * Writes errors and their context to application logs.
+ */
+public class LogReporter implements ErrorReporter {
+
+    private static final Logger log = LoggerFactory.getLogger(LogReporter.class);
+    private static final Future COMPLETED = CompletableFuture.completedFuture(null);
+
+    private final ConnectorTaskId id;
+    private final ConnectorConfig connConfig;
+    private final ErrorHandlingMetrics errorHandlingMetrics;
+
+    public LogReporter(ConnectorTaskId id, ConnectorConfig connConfig, ErrorHandlingMetrics errorHandlingMetrics) {
+        Objects.requireNonNull(id);
+        Objects.requireNonNull(connConfig);
+        Objects.requireNonNull(errorHandlingMetrics);
+
+        this.id = id;
+        this.connConfig = connConfig;
+        this.errorHandlingMetrics = errorHandlingMetrics;
+    }
+
+    /**
+     * Log error context.
+     *
+     * @param context the processing context.
+     */
+    @Override
+    public Future report(ProcessingContext context) {
+        if (!connConfig.enableErrorLog()) {
+            return COMPLETED;
+        }
+
+        if (!context.failed()) {
+            return COMPLETED;
+        }
+
+        log.error(message(context), context.error());
+        errorHandlingMetrics.recordErrorLogged();
+        return COMPLETED;
+    }
+
+    // Visible for testing
+    String message(ProcessingContext context) {
+        return String.format("Error encountered in task %s. %s", String.valueOf(id),
+                context.toString(connConfig.includeRecordDetailsInErrorLog()));
+    }
+
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/Operation.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/Operation.java
new file mode 100644
index 0000000000000..3e0f792c096a6
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/Operation.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.connect.runtime.errors;
+
+import java.util.concurrent.Callable;
+
+/**
+ * A recoverable operation evaluated in the connector pipeline.
+ *
+ * @param  return type of the result of the operation.
+ */
+public interface Operation extends Callable {
+
+}
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
new file mode 100644
index 0000000000000..0ddf894e57d93
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/ProcessingContext.java
@@ -0,0 +1,253 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.errors;
+
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.producer.RecordMetadata;
+import org.apache.kafka.common.record.TimestampType;
+import org.apache.kafka.connect.errors.ConnectException;
+import org.apache.kafka.connect.runtime.errors.WorkerErrantRecordReporter.ErrantRecordFuture;
+import org.apache.kafka.connect.source.SourceRecord;
+
+import java.util.Collection;
+import java.util.Collections;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Objects;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Future;
+import java.util.stream.Collectors;
+
+/**
+ * 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 implements AutoCloseable {
+
+    private Collection reporters = Collections.emptyList();
+
+    private ConsumerRecord consumedMessage;
+    private SourceRecord sourceRecord;
+
+    /**
+     * The following fields need to be reset every time a new record is seen.
+     */
+
+    private Stage position;
+    private Class klass;
+    private int attempt;
+    private Throwable error;
+
+    /**
+     * Reset the internal fields before executing operations on a new record.
+     */
+    private void reset() {
+        attempt = 0;
+        position = null;
+        klass = null;
+        error = null;
+    }
+
+    /**
+     * Set the record consumed from Kafka in a sink connector.
+     *
+     * @param consumedMessage the record
+     */
+    public void consumerRecord(ConsumerRecord consumedMessage) {
+        this.consumedMessage = consumedMessage;
+        reset();
+    }
+
+    /**
+     * @return the record consumed from Kafka. could be null
+     */
+    public ConsumerRecord consumerRecord() {
+        return consumedMessage;
+    }
+
+    /**
+     * @return the source record being processed.
+     */
+    public SourceRecord sourceRecord() {
+        return sourceRecord;
+    }
+
+    /**
+     * Set the source record being processed in the connect pipeline.
+     *
+     * @param record the source record
+     */
+    public void sourceRecord(SourceRecord record) {
+        this.sourceRecord = record;
+        reset();
+    }
+
+    /**
+     * Set the stage in the connector pipeline which is currently executing.
+     *
+     * @param position the stage
+     */
+    public void position(Stage position) {
+        this.position = position;
+    }
+
+    /**
+     * @return the stage in the connector pipeline which is currently executing.
+     */
+    public Stage stage() {
+        return position;
+    }
+
+    /**
+     * @return the class which is going to execute the current operation.
+     */
+    public Class executingClass() {
+        return klass;
+    }
+
+    /**
+     * @param klass set the class which is currently executing.
+     */
+    public void executingClass(Class klass) {
+        this.klass = klass;
+    }
+
+    /**
+     * A helper method to set both the stage and the class.
+     *
+     * @param stage the stage
+     * @param klass the class which will execute the operation in this stage.
+     */
+    public void currentContext(Stage stage, Class klass) {
+        position(stage);
+        executingClass(klass);
+    }
+
+    /**
+     * Report errors. Should be called only if an error was encountered while executing the operation.
+     *
+     * @return a errant record future that potentially aggregates the producer futures
+     */
+    public Future report() {
+        if (reporters.size() == 1) {
+            return new ErrantRecordFuture(Collections.singletonList(reporters.iterator().next().report(this)));
+        }
+
+        List> futures = reporters.stream()
+                .map(r -> r.report(this))
+                .filter(Future::isDone)
+                .collect(Collectors.toCollection(LinkedList::new));
+        if (futures.isEmpty()) {
+            return CompletableFuture.completedFuture(null);
+        }
+        return new ErrantRecordFuture(futures);
+    }
+
+    @Override
+    public String toString() {
+        return toString(false);
+    }
+
+    public String toString(boolean includeMessage) {
+        StringBuilder builder = new StringBuilder();
+        builder.append("Executing stage '");
+        builder.append(stage().name());
+        builder.append("' with class '");
+        builder.append(executingClass() == null ? "null" : executingClass().getName());
+        builder.append('\'');
+        if (includeMessage && sourceRecord() != null) {
+            builder.append(", where source record is = ");
+            builder.append(sourceRecord());
+        } else if (includeMessage && consumerRecord() != null) {
+            ConsumerRecord msg = consumerRecord();
+            builder.append(", where consumed record is ");
+            builder.append("{topic='").append(msg.topic()).append('\'');
+            builder.append(", partition=").append(msg.partition());
+            builder.append(", offset=").append(msg.offset());
+            if (msg.timestampType() == TimestampType.CREATE_TIME || msg.timestampType() == TimestampType.LOG_APPEND_TIME) {
+                builder.append(", timestamp=").append(msg.timestamp());
+                builder.append(", timestampType=").append(msg.timestampType());
+            }
+            builder.append("}");
+        }
+        builder.append('.');
+        return builder.toString();
+    }
+
+    /**
+     * @param attempt the number of attempts made to execute the current operation.
+     */
+    public void attempt(int attempt) {
+        this.attempt = attempt;
+    }
+
+    /**
+     * @return the number of attempts made to execute the current operation.
+     */
+    public int attempt() {
+        return attempt;
+    }
+
+    /**
+     * @return the error (if any) which was encountered while processing the current stage.
+     */
+    public Throwable error() {
+        return error;
+    }
+
+    /**
+     * The error (if any) which was encountered while processing the current stage.
+     *
+     * @param error the error
+     */
+    public void error(Throwable error) {
+        this.error = error;
+    }
+
+    /**
+     * @return true, if the last operation encountered an error; false otherwise
+     */
+    public boolean failed() {
+        return error() != null;
+    }
+
+    /**
+     * Set the error reporters for this connector.
+     *
+     * @param reporters the error reporters (should not be null).
+     */
+    public void reporters(Collection reporters) {
+        Objects.requireNonNull(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
new file mode 100644
index 0000000000000..ce4c1e27dda09
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/RetryWithToleranceOperator.java
@@ -0,0 +1,314 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.errors;
+
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.common.config.ConfigException;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.connect.errors.ConnectException;
+import org.apache.kafka.connect.errors.RetriableException;
+import org.apache.kafka.connect.runtime.ConnectorConfig;
+import org.apache.kafka.connect.source.SourceRecord;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.Future;
+import java.util.concurrent.ThreadLocalRandom;
+
+/**
+ * Attempt to recover a failed operation with retries and tolerance limits.
+ * 
            
+ *
+ * A retry is attempted if the operation throws a {@link RetriableException}. Retries are accompanied by exponential backoffs, starting with
+ * {@link #RETRIES_DELAY_MIN_MS}, up to what is specified with {@link ConnectorConfig#errorMaxDelayInMillis()}.
+ * Including the first attempt and future retries, the total time taken to evaluate the operation should be within
+ * {@link ConnectorConfig#errorMaxDelayInMillis()} millis.
+ * 
            
+ *
+ * This executor will tolerate failures, as specified by {@link ConnectorConfig#errorToleranceType()}.
+ * For transformations and converters, all exceptions are tolerated. For others operations, only {@link RetriableException} are tolerated.
+ * 
            
+ *
+ * There are three outcomes to executing an operation. It might succeed, in which case the result is returned to the caller.
+ * If it fails, this class does one of these two things: (1) if the failure occurred due to a tolerable exception, then
+ * set appropriate error reason in the {@link ProcessingContext} and return null, or (2) if the exception is not tolerated,
+ * then it is wrapped into a ConnectException and rethrown to the caller.
+ * 
            
+ *
+ * Instances of this class are thread safe.
+ * 
            
+ */
+public class RetryWithToleranceOperator implements AutoCloseable {
+
+    private static final Logger log = LoggerFactory.getLogger(RetryWithToleranceOperator.class);
+
+    public static final long RETRIES_DELAY_MIN_MS = 300;
+
+    private static final Map> TOLERABLE_EXCEPTIONS = new HashMap<>();
+    static {
+        TOLERABLE_EXCEPTIONS.put(Stage.TRANSFORMATION, Exception.class);
+        TOLERABLE_EXCEPTIONS.put(Stage.HEADER_CONVERTER, Exception.class);
+        TOLERABLE_EXCEPTIONS.put(Stage.KEY_CONVERTER, Exception.class);
+        TOLERABLE_EXCEPTIONS.put(Stage.VALUE_CONVERTER, Exception.class);
+    }
+
+    private final long errorRetryTimeout;
+    private final long errorMaxDelayInMillis;
+    private final ToleranceType errorToleranceType;
+
+    private long totalFailures = 0;
+    private final Time time;
+    private ErrorHandlingMetrics errorHandlingMetrics;
+
+    protected final ProcessingContext context;
+
+    public RetryWithToleranceOperator(long errorRetryTimeout, long errorMaxDelayInMillis,
+                                      ToleranceType toleranceType, Time time) {
+        this(errorRetryTimeout, errorMaxDelayInMillis, toleranceType, time, new ProcessingContext());
+    }
+
+    RetryWithToleranceOperator(long errorRetryTimeout, long errorMaxDelayInMillis,
+                               ToleranceType toleranceType, Time time,
+                               ProcessingContext context) {
+        this.errorRetryTimeout = errorRetryTimeout;
+        this.errorMaxDelayInMillis = errorMaxDelayInMillis;
+        this.errorToleranceType = toleranceType;
+        this.time = time;
+        this.context = context;
+    }
+
+    public synchronized Future executeFailed(Stage stage, Class executingClass,
+                                      ConsumerRecord consumerRecord,
+                                      Throwable error) {
+
+        markAsFailed();
+        context.consumerRecord(consumerRecord);
+        context.currentContext(stage, executingClass);
+        context.error(error);
+        errorHandlingMetrics.recordFailure();
+        Future errantRecordFuture = context.report();
+        if (!withinToleranceLimits()) {
+            errorHandlingMetrics.recordError();
+            throw new ConnectException("Tolerance exceeded in error handler", error);
+        }
+        return errantRecordFuture;
+    }
+
+    /**
+     * Execute the recoverable operation. If the operation is already in a failed state, then simply return
+     * with the existing failure.
+     *
+     * @param operation the recoverable operation
+     * @param  return type of the result of the operation.
+     * @return result of the operation
+     */
+    public synchronized  V execute(Operation operation, Stage stage, Class executingClass) {
+        context.currentContext(stage, executingClass);
+
+        if (context.failed()) {
+            log.debug("ProcessingContext is already in failed state. Ignoring requested operation.");
+            return null;
+        }
+
+        try {
+            Class ex = TOLERABLE_EXCEPTIONS.getOrDefault(context.stage(), RetriableException.class);
+            return execAndHandleError(operation, ex);
+        } finally {
+            if (context.failed()) {
+                errorHandlingMetrics.recordError();
+                context.report();
+            }
+        }
+    }
+
+    /**
+     * Attempt to execute an operation. Retry if a {@link RetriableException} is raised. Re-throw everything else.
+     *
+     * @param operation the operation to be executed.
+     * @param  the return type of the result of the operation.
+     * @return the result of the operation.
+     * @throws Exception rethrow if a non-retriable Exception is thrown by the operation
+     */
+    protected  V execAndRetry(Operation operation) throws Exception {
+        int attempt = 0;
+        long startTime = time.milliseconds();
+        long deadline = startTime + errorRetryTimeout;
+        do {
+            try {
+                attempt++;
+                return operation.call();
+            } catch (RetriableException e) {
+                log.trace("Caught a retriable exception while executing {} operation with {}", context.stage(), context.executingClass());
+                errorHandlingMetrics.recordFailure();
+                if (checkRetry(startTime)) {
+                    backoff(attempt, deadline);
+                    if (Thread.currentThread().isInterrupted()) {
+                        log.trace("Thread was interrupted. Marking operation as failed.");
+                        context.error(e);
+                        return null;
+                    }
+                    errorHandlingMetrics.recordRetry();
+                } else {
+                    log.trace("Can't retry. start={}, attempt={}, deadline={}", startTime, attempt, deadline);
+                    context.error(e);
+                    return null;
+                }
+            } finally {
+                context.attempt(attempt);
+            }
+        } while (true);
+    }
+
+    /**
+     * Execute a given operation multiple times (if needed), and tolerate certain exceptions.
+     *
+     * @param operation the operation to be executed.
+     * @param tolerated the class of exceptions which can be tolerated.
+     * @param  The return type of the result of the operation.
+     * @return the result of the operation
+     */
+    // Visible for testing
+    protected  V execAndHandleError(Operation operation, Class tolerated) {
+        try {
+            V result = execAndRetry(operation);
+            if (context.failed()) {
+                markAsFailed();
+                errorHandlingMetrics.recordSkipped();
+            }
+            return result;
+        } catch (Exception e) {
+            errorHandlingMetrics.recordFailure();
+            markAsFailed();
+            context.error(e);
+
+            if (!tolerated.isAssignableFrom(e.getClass())) {
+                throw new ConnectException("Unhandled exception in error handler", e);
+            }
+
+            if (!withinToleranceLimits()) {
+                throw new ConnectException("Tolerance exceeded in error handler", e);
+            }
+
+            errorHandlingMetrics.recordSkipped();
+            return null;
+        }
+    }
+
+    // Visible for testing
+    void markAsFailed() {
+        errorHandlingMetrics.recordErrorTimestamp();
+        totalFailures++;
+    }
+
+    @SuppressWarnings("fallthrough")
+    public synchronized boolean withinToleranceLimits() {
+        switch (errorToleranceType) {
+            case NONE:
+                if (totalFailures > 0) return false;
+            case ALL:
+                return true;
+            default:
+                throw new ConfigException("Unknown tolerance type: {}", errorToleranceType);
+        }
+    }
+
+    // Visible for testing
+    boolean checkRetry(long startTime) {
+        return (time.milliseconds() - startTime) < errorRetryTimeout;
+    }
+
+    // Visible for testing
+    void backoff(int attempt, long deadline) {
+        int numRetry = attempt - 1;
+        long delay = RETRIES_DELAY_MIN_MS << numRetry;
+        if (delay > errorMaxDelayInMillis) {
+            delay = ThreadLocalRandom.current().nextLong(errorMaxDelayInMillis);
+        }
+        if (delay + time.milliseconds() > deadline) {
+            delay = deadline - time.milliseconds();
+        }
+        log.debug("Sleeping for {} millis", delay);
+        time.sleep(delay);
+    }
+
+    public synchronized void metrics(ErrorHandlingMetrics errorHandlingMetrics) {
+        this.errorHandlingMetrics = errorHandlingMetrics;
+    }
+
+    @Override
+    public String toString() {
+        return "RetryWithToleranceOperator{" +
+                "errorRetryTimeout=" + errorRetryTimeout +
+                ", errorMaxDelayInMillis=" + errorMaxDelayInMillis +
+                ", errorToleranceType=" + errorToleranceType +
+                ", totalFailures=" + totalFailures +
+                ", time=" + time +
+                ", context=" + context +
+                '}';
+    }
+
+    /**
+     * Set the error reporters for this connector.
+     *
+     * @param reporters the error reporters (should not be null).
+     */
+    public synchronized void reporters(List reporters) {
+        this.context.reporters(reporters);
+    }
+
+    /**
+     * Set the source record being processed in the connect pipeline.
+     *
+     * @param preTransformRecord the source record
+     */
+    public synchronized void sourceRecord(SourceRecord preTransformRecord) {
+        this.context.sourceRecord(preTransformRecord);
+    }
+
+    /**
+     * Set the record consumed from Kafka in a sink connector.
+     *
+     * @param consumedMessage the record
+     */
+    public synchronized void consumerRecord(ConsumerRecord consumedMessage) {
+        this.context.consumerRecord(consumedMessage);
+    }
+
+    /**
+     * @return true, if the last operation encountered an error; false otherwise
+     */
+    public synchronized boolean failed() {
+        return this.context.failed();
+    }
+
+    /**
+     * Returns the error encountered when processing the current stage.
+     *
+     * @return the error encountered when processing the current stage
+     */
+    public synchronized Throwable error() {
+        return this.context.error();
+    }
+
+    @Override
+    public synchronized void close() {
+        this.context.close();
+    }
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/Stage.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/Stage.java
new file mode 100644
index 0000000000000..b9aa1f21bc68d
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/Stage.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.connect.runtime.errors;
+
+/**
+ * A logical stage in a Connect pipeline.
+ */
+public enum Stage {
+
+    /**
+     * When calling the poll() method on a SourceConnector
+     */
+    TASK_POLL,
+
+    /**
+     * When calling the put() method on a SinkConnector
+     */
+    TASK_PUT,
+
+    /**
+     * When running any transformation operation on a record
+     */
+    TRANSFORMATION,
+
+    /**
+     * When using the key converter to serialize/deserialize keys in ConnectRecords
+     */
+    KEY_CONVERTER,
+
+    /**
+     * When using the value converter to serialize/deserialize values in ConnectRecords
+     */
+    VALUE_CONVERTER,
+
+    /**
+     * When using the header converter to serialize/deserialize headers in ConnectRecords
+     */
+    HEADER_CONVERTER,
+
+    /**
+     * When producing to Kafka topic
+     */
+    KAFKA_PRODUCE,
+
+    /**
+     * When consuming from a Kafka topic
+     */
+    KAFKA_CONSUME
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/ToleranceType.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/ToleranceType.java
new file mode 100644
index 0000000000000..dd40a60648a34
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/ToleranceType.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.runtime.errors;
+
+import java.util.Locale;
+
+/**
+ * The different levels of error tolerance.
+ */
+public enum ToleranceType {
+
+    /**
+     * Tolerate no errors.
+     */
+    NONE,
+
+    /**
+     * Tolerate all errors.
+     */
+    ALL;
+
+    public String value() {
+        return name().toLowerCase(Locale.ROOT);
+    }
+
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/WorkerErrantRecordReporter.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/WorkerErrantRecordReporter.java
new file mode 100644
index 0000000000000..b709e28c5edfe
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/WorkerErrantRecordReporter.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.runtime.errors;
+
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.producer.RecordMetadata;
+import org.apache.kafka.common.header.internals.RecordHeaders;
+import org.apache.kafka.connect.errors.ConnectException;
+import org.apache.kafka.connect.header.Header;
+import org.apache.kafka.connect.runtime.InternalSinkRecord;
+import org.apache.kafka.connect.sink.ErrantRecordReporter;
+import org.apache.kafka.connect.sink.SinkRecord;
+
+import org.apache.kafka.connect.sink.SinkTask;
+import org.apache.kafka.connect.storage.Converter;
+import org.apache.kafka.connect.storage.HeaderConverter;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.LinkedList;
+import java.util.List;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+
+public class WorkerErrantRecordReporter implements ErrantRecordReporter {
+
+    private static final Logger log = LoggerFactory.getLogger(WorkerErrantRecordReporter.class);
+
+    private final RetryWithToleranceOperator retryWithToleranceOperator;
+    private final Converter keyConverter;
+    private final Converter valueConverter;
+    private final HeaderConverter headerConverter;
+
+    // Visible for testing
+    protected final LinkedList> futures;
+
+    public WorkerErrantRecordReporter(
+        RetryWithToleranceOperator retryWithToleranceOperator,
+        Converter keyConverter,
+        Converter valueConverter,
+        HeaderConverter headerConverter
+    ) {
+        this.retryWithToleranceOperator = retryWithToleranceOperator;
+        this.keyConverter = keyConverter;
+        this.valueConverter = valueConverter;
+        this.headerConverter = headerConverter;
+        this.futures = new LinkedList<>();
+    }
+
+    @Override
+    public Future report(SinkRecord record, Throwable error) {
+        ConsumerRecord consumerRecord;
+
+        // Most of the records will be an internal sink record, but the task could potentially
+        // report modified or new records, so handle both cases
+        if (record instanceof InternalSinkRecord) {
+            consumerRecord = ((InternalSinkRecord) record).originalRecord();
+        } else {
+            // Generate a new consumer record from the modified sink record. We prefer
+            // to send the original consumer record (pre-transformed) to the DLQ,
+            // but in this case we don't have one and send the potentially transformed
+            // record instead
+            String topic = record.topic();
+            byte[] key = keyConverter.fromConnectData(topic, record.keySchema(), record.key());
+            byte[] value = valueConverter.fromConnectData(topic,
+                record.valueSchema(), record.value());
+
+            RecordHeaders headers = new RecordHeaders();
+            if (record.headers() != null) {
+                for (Header header : record.headers()) {
+                    String headerKey = header.key();
+                    byte[] rawHeader = headerConverter.fromConnectHeader(topic, headerKey,
+                        header.schema(), header.value());
+                    headers.add(headerKey, rawHeader);
+                }
+            }
+
+            int keyLength = key != null ? key.length : -1;
+            int valLength = value != null ? value.length : -1;
+
+            consumerRecord = new ConsumerRecord<>(record.topic(), record.kafkaPartition(),
+                record.kafkaOffset(), record.timestamp(), record.timestampType(), -1L, keyLength,
+                valLength, key, value, headers);
+        }
+
+        Future future = retryWithToleranceOperator.executeFailed(Stage.TASK_PUT, SinkTask.class, consumerRecord, error);
+
+        if (!future.isDone()) {
+            futures.add(future);
+        }
+        return future;
+    }
+
+    /**
+     * Gets all futures returned by the sink records sent to Kafka by the errant
+     * record reporter. This function is intended to be used to block on all the errant record
+     * futures.
+     */
+    public void awaitAllFutures() {
+        Future future;
+        while ((future = futures.poll()) != null) {
+            try {
+                future.get();
+            } catch (InterruptedException | ExecutionException e) {
+                log.error("Encountered an error while awaiting an errant record future's completion.");
+                throw new ConnectException(e);
+            }
+        }
+    }
+
+    /**
+     * Wrapper class to aggregate producer futures and abstract away the record metadata from the
+     * Connect user.
+     */
+    public static class ErrantRecordFuture implements Future {
+
+        private final List> futures;
+
+        public ErrantRecordFuture(List> producerFutures) {
+            futures = producerFutures;
+        }
+
+        public boolean cancel(boolean mayInterruptIfRunning) {
+            throw new UnsupportedOperationException("Reporting an errant record cannot be cancelled.");
+        }
+
+        public boolean isCancelled() {
+            return false;
+        }
+
+        public boolean isDone() {
+            return futures.stream().allMatch(Future::isDone);
+        }
+
+        public Void get() throws InterruptedException, ExecutionException {
+            for (Future future: futures) {
+                future.get();
+            }
+            return null;
+        }
+
+        public Void get(long timeout, TimeUnit unit)
+            throws InterruptedException, ExecutionException, TimeoutException {
+            for (Future future: futures) {
+                future.get(timeout, unit);
+            }
+            return null;
+        }
+    }
+}
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
new file mode 100644
index 0000000000000..6b7285df50b97
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/health/ConnectClusterStateImpl.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.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.Herder;
+import org.apache.kafka.connect.runtime.rest.entities.ConnectorStateInfo;
+import org.apache.kafka.connect.util.FutureCallback;
+
+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;
+
+    public ConnectClusterStateImpl(
+        long connectorsTimeoutMs,
+        ConnectClusterDetails clusterDetails,
+        Herder herder
+    ) {
+        this.herderRequestTimeoutMs = connectorsTimeoutMs;
+        this.clusterDetails = clusterDetails;
+        this.herder = herder;
+    }
+
+    @Override
+    public Collection 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 = herder.connectorStatus(connName);
+        ConnectorState connectorState = new ConnectorState(
+            state.connector().state(),
+            state.connector().workerId(),
+            state.connector().trace()
+        );
+        Map taskStates = taskStates(state.tasks());
+        ConnectorHealth connectorHealth = new ConnectorHealth(
+            connName,
+            connectorState,
+            taskStates,
+            ConnectorType.valueOf(state.type().name())
+        );
+        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<>();
+
+        for (ConnectorStateInfo.TaskState state : states) {
+            taskStates.put(
+                state.id(),
+                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 0c133b9a8918a..86a53c75ced58 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
@@ -16,10 +16,19 @@
  */
 package org.apache.kafka.connect.runtime.isolation;
 
+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;
 import org.apache.kafka.connect.transforms.Transformation;
+import org.apache.kafka.connect.transforms.predicates.Predicate;
+import org.reflections.Configuration;
 import org.reflections.Reflections;
+import org.reflections.ReflectionsException;
+import org.reflections.scanners.SubTypesScanner;
 import org.reflections.util.ClasspathHelper;
 import org.reflections.util.ConfigurationBuilder;
 import org.slf4j.Logger;
@@ -39,6 +48,8 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
+import java.util.Collections;
+import java.util.Enumeration;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
@@ -49,31 +60,52 @@
 import java.util.SortedSet;
 import java.util.TreeMap;
 import java.util.TreeSet;
+import java.util.stream.Collectors;
 
 public class DelegatingClassLoader extends URLClassLoader {
     private static final Logger log = LoggerFactory.getLogger(DelegatingClassLoader.class);
+    private static final String CLASSPATH_NAME = "classpath";
+    private static final String UNDEFINED_VERSION = "undefined";
 
     private final Map, ClassLoader>> pluginLoaders;
     private final Map aliases;
     private final SortedSet> connectors;
     private final SortedSet> converters;
+    private final SortedSet> headerConverters;
     private final SortedSet> transformations;
+    private final SortedSet> predicates;
+    private final SortedSet> configProviders;
+    private final SortedSet> restExtensions;
+    private final SortedSet> connectorClientConfigPolicies;
     private final List pluginPaths;
-    private final Map activePaths;
+
+    private static final String MANIFEST_PREFIX = "META-INF/services/";
+    private static final Class[] SERVICE_LOADER_PLUGINS = new Class[] {ConnectRestExtension.class, ConfigProvider.class};
+    private static final Set PLUGIN_MANIFEST_FILES =
+        Arrays.stream(SERVICE_LOADER_PLUGINS).map(serviceLoaderPlugin -> MANIFEST_PREFIX + serviceLoaderPlugin.getName())
+            .collect(Collectors.toSet());
 
     public DelegatingClassLoader(List pluginPaths, ClassLoader parent) {
         super(new URL[0], parent);
         this.pluginPaths = pluginPaths;
         this.pluginLoaders = new HashMap<>();
         this.aliases = new HashMap<>();
-        this.activePaths = new HashMap<>();
         this.connectors = new TreeSet<>();
         this.converters = new TreeSet<>();
+        this.headerConverters = new TreeSet<>();
         this.transformations = new TreeSet<>();
+        this.predicates = new TreeSet<>();
+        this.configProviders = new TreeSet<>();
+        this.restExtensions = new TreeSet<>();
+        this.connectorClientConfigPolicies = new TreeSet<>();
     }
 
     public DelegatingClassLoader(List pluginPaths) {
-        this(pluginPaths, ClassLoader.getSystemClassLoader());
+        // Use as parent the classloader that loaded this class. In most cases this will be the
+        // System classloader. But this choice here provides additional flexibility in managed
+        // environments that control classloading differently (OSGi, Spring and others) and don't
+        // depend on the System classloader to load Connect's classes.
+        this(pluginPaths, DelegatingClassLoader.class.getClassLoader());
     }
 
     public Set> connectors() {
@@ -84,29 +116,65 @@ public Set> converters() {
         return converters;
     }
 
+    public Set> headerConverters() {
+        return headerConverters;
+    }
+
     public Set> transformations() {
         return transformations;
     }
 
+    public Set> predicates() {
+        return predicates;
+    }
+
+    public Set> configProviders() {
+        return configProviders;
+    }
+
+    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());
     }
 
     public ClassLoader connectorLoader(String connectorClassOrAlias) {
-        log.debug("Getting plugin class loader for connector: '{}'", connectorClassOrAlias);
         String fullName = aliases.containsKey(connectorClassOrAlias)
                           ? aliases.get(connectorClassOrAlias)
                           : connectorClassOrAlias;
-        SortedMap, ClassLoader> inner = pluginLoaders.get(fullName);
-        if (inner == null) {
-            log.error(
-                    "Plugin class loader for connector: '{}' was not found. Returning: {}",
-                    connectorClassOrAlias,
-                    this
-            );
-            return this;
-        }
-        return inner.get(inner.lastKey());
+        ClassLoader classLoader = pluginClassLoader(fullName);
+        if (classLoader == null) classLoader = this;
+        log.debug(
+            "Getting plugin class loader: '{}' for connector: {}",
+            classLoader,
+            connectorClassOrAlias
+        );
+        return classLoader;
     }
 
     private static PluginClassLoader newPluginClassLoader(
@@ -114,13 +182,8 @@ private static PluginClassLoader newPluginClassLoader(
             final URL[] urls,
             final ClassLoader parent
     ) {
-        return (PluginClassLoader) AccessController.doPrivileged(
-                new PrivilegedAction() {
-                    @Override
-                    public Object run() {
-                        return new PluginClassLoader(pluginLocation, urls, parent);
-                    }
-                }
+        return AccessController.doPrivileged(
+                (PrivilegedAction) () -> new PluginClassLoader(pluginLocation, urls, parent)
         );
     }
 
@@ -139,10 +202,23 @@ private  void addPlugins(Collection> plugins, ClassLoader loade
     }
 
     protected void initLoaders() {
-        String path = null;
+        for (String configPath : pluginPaths) {
+            initPluginLoader(configPath);
+        }
+        // Finally add parent/system loader.
+        initPluginLoader(CLASSPATH_NAME);
+        addAllAliases();
+    }
+
+    private void initPluginLoader(String path) {
         try {
-            for (String configPath : pluginPaths) {
-                path = configPath;
+            if (CLASSPATH_NAME.equals(path)) {
+                scanUrlsAndAddPlugins(
+                        getParent(),
+                        ClasspathHelper.forJavaClassPath().toArray(new URL[0]),
+                        null
+                );
+            } else {
                 Path pluginPath = Paths.get(path).toAbsolutePath();
                 // Update for exception handling
                 path = pluginPath.toString();
@@ -156,14 +232,6 @@ protected void initLoaders() {
                     registerPlugin(pluginPath);
                 }
             }
-
-            path = "classpath";
-            // Finally add parent/system loader.
-            scanUrlsAndAddPlugins(
-                    getParent(),
-                    ClasspathHelper.forJavaClassPath().toArray(new URL[0]),
-                    null
-            );
         } catch (InvalidPathException | MalformedURLException e) {
             log.error("Invalid path in plugin path: {}. Ignoring.", path, e);
         } catch (IOException e) {
@@ -171,7 +239,6 @@ protected void initLoaders() {
         } catch (InstantiationException | IllegalAccessException e) {
             log.error("Could not instantiate plugins in: {}. Ignoring: {}", path, e);
         }
-        addAllAliases();
     }
 
     private void registerPlugin(Path pluginLocation)
@@ -201,16 +268,22 @@ private void scanUrlsAndAddPlugins(
         PluginScanResult plugins = scanPluginPath(loader, urls);
         log.info("Registered loader: {}", loader);
         if (!plugins.isEmpty()) {
-            if (loader instanceof PluginClassLoader) {
-                activePaths.put(pluginLocation, (PluginClassLoader) loader);
-            }
-
             addPlugins(plugins.connectors(), loader);
             connectors.addAll(plugins.connectors());
             addPlugins(plugins.converters(), loader);
             converters.addAll(plugins.converters());
+            addPlugins(plugins.headerConverters(), loader);
+            headerConverters.addAll(plugins.headerConverters());
             addPlugins(plugins.transformations(), loader);
             transformations.addAll(plugins.transformations());
+            addPlugins(plugins.predicates(), loader);
+            predicates.addAll(plugins.predicates());
+            addPlugins(plugins.configProviders(), loader);
+            configProviders.addAll(plugins.configProviders());
+            addPlugins(plugins.restExtensions(), loader);
+            restExtensions.addAll(plugins.restExtensions());
+            addPlugins(plugins.connectorClientConfigPolicies(), loader);
+            connectorClientConfigPolicies.addAll(plugins.connectorClientConfigPolicies());
         }
 
         loadJdbcDrivers(loader);
@@ -221,6 +294,7 @@ private void loadJdbcDrivers(final ClassLoader loader) {
         // implementing the java.sql.Driver interface.
         AccessController.doPrivileged(
                 new PrivilegedAction() {
+                    @Override
                     public Void run() {
                         ServiceLoader loadedDrivers = ServiceLoader.load(
                                 Driver.class,
@@ -255,12 +329,19 @@ private PluginScanResult scanPluginPath(
         ConfigurationBuilder builder = new ConfigurationBuilder();
         builder.setClassLoaders(new ClassLoader[]{loader});
         builder.addUrls(urls);
-        Reflections reflections = new Reflections(builder);
+        builder.setScanners(new SubTypesScanner());
+        builder.useParallelExecutor();
+        Reflections reflections = new InternalReflections(builder);
 
         return new PluginScanResult(
                 getPluginDesc(reflections, Connector.class, loader),
                 getPluginDesc(reflections, Converter.class, loader),
-                getPluginDesc(reflections, Transformation.class, loader)
+                getPluginDesc(reflections, HeaderConverter.class, loader),
+                getPluginDesc(reflections, Transformation.class, loader),
+                getPluginDesc(reflections, Predicate.class, loader),
+                getServiceLoaderPluginDesc(ConfigProvider.class, loader),
+                getServiceLoaderPluginDesc(ConnectRestExtension.class, loader),
+                getServiceLoaderPluginDesc(ConnectorClientConfigOverridePolicy.class, loader)
         );
     }
 
@@ -269,43 +350,58 @@ private  Collection> getPluginDesc(
             Class klass,
             ClassLoader loader
     ) throws InstantiationException, IllegalAccessException {
-        Set> plugins = reflections.getSubTypesOf(klass);
+        Set> plugins;
+        try {
+            plugins = reflections.getSubTypesOf(klass);
+        } catch (ReflectionsException e) {
+            log.debug("Reflections scanner could not find any classes for URLs: " +
+                    reflections.getConfiguration().getUrls(), e);
+            return Collections.emptyList();
+        }
 
         Collection> result = new ArrayList<>();
         for (Class plugin : plugins) {
             if (PluginUtils.isConcrete(plugin)) {
-                // Temporary workaround until all the plugins are versioned.
-                if (Connector.class.isAssignableFrom(plugin)) {
-                    result.add(
-                            new PluginDesc<>(
-                                    plugin,
-                                    ((Connector) plugin.newInstance()).version(),
-                                    loader
-                            )
-                    );
-                } else {
-                    result.add(new PluginDesc<>(plugin, "undefined", loader));
-                }
+                result.add(new PluginDesc<>(plugin, versionFor(plugin), loader));
+            } else {
+                log.debug("Skipping {} as it is not concrete implementation", plugin);
             }
         }
         return result;
     }
 
-    @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);
+    @SuppressWarnings("unchecked")
+    private  Collection> getServiceLoaderPluginDesc(Class klass, ClassLoader loader) {
+        ClassLoader savedLoader = Plugins.compareAndSwapLoaders(loader);
+        Collection> result = new ArrayList<>();
+        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;
+    }
+
+    private static   String versionFor(T pluginImpl) {
+        return pluginImpl instanceof Versioned ? ((Versioned) pluginImpl).version() : UNDEFINED_VERSION;
+    }
+
+    private static  String versionFor(Class pluginKlass) throws IllegalAccessException, InstantiationException {
+        // Temporary workaround until all the plugins are versioned.
+        return Connector.class.isAssignableFrom(pluginKlass) ? versionFor(pluginKlass.newInstance()) : UNDEFINED_VERSION;
+    }
 
+    @Override
+    protected Class loadClass(String name, boolean resolve) throws ClassNotFoundException {
         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);
@@ -314,7 +410,11 @@ protected Class loadClass(String name, boolean resolve) throws ClassNotFoundE
     private void addAllAliases() {
         addAliases(connectors);
         addAliases(converters);
+        addAliases(headerConverters);
         addAliases(transformations);
+        addAliases(predicates);
+        addAliases(restExtensions);
+        addAliases(connectorClientConfigPolicies);
     }
 
     private  void addAliases(Collection> plugins) {
@@ -337,4 +437,52 @@ private  void addAliases(Collection> plugins) {
             }
         }
     }
+
+    private static class InternalReflections extends Reflections {
+
+        public InternalReflections(Configuration configuration) {
+            super(configuration);
+        }
+
+        // When Reflections is used for parallel scans, it has a bug where it propagates ReflectionsException
+        // as RuntimeException.  Override the scan behavior to emulate the singled-threaded logic.
+        @Override
+        protected void scan(URL url) {
+            try {
+                super.scan(url);
+            } catch (ReflectionsException e) {
+                Logger log = Reflections.log;
+                if (log != null && log.isWarnEnabled()) {
+                    log.warn("could not create Vfs.Dir from url. ignoring the exception and continuing", e);
+                }
+            }
+        }
+    }
+
+    @Override
+    public URL getResource(String name) {
+        if (serviceLoaderManifestForPlugin(name)) {
+            // Default implementation of getResource searches the parent class loader and if not available/found, its own URL paths.
+            // This will enable thePluginClassLoader to limit its resource search only to its own URL paths.
+            return null;
+        } else {
+            return super.getResource(name);
+        }
+    }
+
+    @Override
+    public Enumeration getResources(String name) throws IOException {
+        if (serviceLoaderManifestForPlugin(name)) {
+            // Default implementation of getResources searches the parent class loader and and also its own URL paths. This will enable the
+            // PluginClassLoader to limit its resource search to only its own URL paths.
+            return null;
+        } else {
+            return super.getResources(name);
+        }
+    }
+
+    //Visible for testing
+    static boolean serviceLoaderManifestForPlugin(String name) {
+        return PLUGIN_MANIFEST_FILES.contains(name);
+    }
 }
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/PluginClassLoader.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/PluginClassLoader.java
index 780ebd09bfcbf..9a4fed22a90e2 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/PluginClassLoader.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/PluginClassLoader.java
@@ -22,20 +22,56 @@
 import java.net.URL;
 import java.net.URLClassLoader;
 
+/**
+ * A custom classloader dedicated to loading Connect plugin classes in classloading isolation.
+ * 
            
+ * Under the current scheme for classloading isolation in Connect, a plugin classloader loads the
+ * classes that it finds in its urls. For classes that are either not found or are not supposed to
+ * be loaded in isolation, this plugin classloader delegates their loading to its parent. This makes
+ * this classloader a child-first classloader.
+ * 
            
+ * This class is thread-safe.
+ */
 public class PluginClassLoader extends URLClassLoader {
     private static final Logger log = LoggerFactory.getLogger(PluginClassLoader.class);
     private final URL pluginLocation;
 
+    static {
+        ClassLoader.registerAsParallelCapable();
+    }
+
+    /**
+     * Constructor that accepts a specific classloader as parent.
+     *
+     * @param pluginLocation the top-level location of the plugin to be loaded in isolation by this
+     * classloader.
+     * @param urls the list of urls from which to load classes and resources for this plugin.
+     * @param parent the parent classloader to be used for delegation for classes that were
+     * not found or should not be loaded in isolation by this classloader.
+     */
     public PluginClassLoader(URL pluginLocation, URL[] urls, ClassLoader parent) {
         super(urls, parent);
         this.pluginLocation = pluginLocation;
     }
 
+    /**
+     * Constructor that defines the system classloader as parent of this plugin classloader.
+     *
+     * @param pluginLocation the top-level location of the plugin to be loaded in isolation by this
+     * classloader.
+     * @param urls the list of urls from which to load classes and resources for this plugin.
+     */
     public PluginClassLoader(URL pluginLocation, URL[] urls) {
         super(urls);
         this.pluginLocation = pluginLocation;
     }
 
+    /**
+     * Returns the top-level location of the classes and dependencies required by the plugin that
+     * is loaded by this classloader.
+     *
+     * @return the plugin location.
+     */
     public String location() {
         return pluginLocation.toString();
     }
@@ -45,25 +81,32 @@ public String toString() {
         return "PluginClassLoader{pluginLocation=" + pluginLocation + "}";
     }
 
+    // This method needs to be thread-safe because it is supposed to be called by multiple
+    // Connect tasks. While findClass is thread-safe, defineClass called within loadClass of the
+    // base method is not. More on multithreaded classloaders in:
+    // https://docs.oracle.com/javase/7/docs/technotes/guides/lang/cl-mt.html
     @Override
-    protected Class loadClass(String name, boolean resolve) throws ClassNotFoundException {
-        Class klass = findLoadedClass(name);
-        if (klass == null) {
-            try {
-                if (PluginUtils.shouldLoadInIsolation(name)) {
-                    klass = findClass(name);
+    protected synchronized Class loadClass(String name, boolean resolve)
+            throws ClassNotFoundException {
+        synchronized (getClassLoadingLock(name)) {
+            Class klass = findLoadedClass(name);
+            if (klass == null) {
+                try {
+                    if (PluginUtils.shouldLoadInIsolation(name)) {
+                        klass = findClass(name);
+                    }
+                } catch (ClassNotFoundException e) {
+                    // Not found in loader's path. Search in parents.
+                    log.trace("Class '{}' not found. Delegating to parent", name);
                 }
-            } catch (ClassNotFoundException e) {
-                // Not found in loader's path. Search in parents.
-                log.trace("Class '{}' not found. Delegating to parent", name);
             }
+            if (klass == null) {
+                klass = super.loadClass(name, false);
+            }
+            if (resolve) {
+                resolveClass(klass);
+            }
+            return klass;
         }
-        if (klass == null) {
-            klass = super.loadClass(name, false);
-        }
-        if (resolve) {
-            resolveClass(klass);
-        }
-        return klass;
     }
 }
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 f3d2f21e074fe..ac42dbb98a1fd 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
@@ -16,25 +16,52 @@
  */
 package org.apache.kafka.connect.runtime.isolation;
 
+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 org.apache.kafka.connect.transforms.predicates.Predicate;
 
+import java.util.Arrays;
 import java.util.Collection;
+import java.util.List;
 
 public class PluginScanResult {
     private final Collection> connectors;
     private final Collection> converters;
+    private final Collection> headerConverters;
     private final Collection> transformations;
+    private final Collection> predicates;
+    private final Collection> configProviders;
+    private final Collection> restExtensions;
+    private final Collection> connectorClientConfigPolicies;
+
+    private final List allPlugins;
 
     public PluginScanResult(
             Collection> connectors,
             Collection> converters,
-            Collection> transformations
+            Collection> headerConverters,
+            Collection> transformations,
+            Collection> predicates,
+            Collection> configProviders,
+            Collection> restExtensions,
+            Collection> connectorClientConfigPolicies
     ) {
         this.connectors = connectors;
         this.converters = converters;
+        this.headerConverters = headerConverters;
         this.transformations = transformations;
+        this.predicates = predicates;
+        this.configProviders = configProviders;
+        this.restExtensions = restExtensions;
+        this.connectorClientConfigPolicies = connectorClientConfigPolicies;
+        this.allPlugins =
+            Arrays.asList(connectors, converters, headerConverters, transformations, configProviders,
+                          connectorClientConfigPolicies);
     }
 
     public Collection> connectors() {
@@ -45,11 +72,35 @@ public Collection> converters() {
         return converters;
     }
 
+    public Collection> headerConverters() {
+        return headerConverters;
+    }
+
     public Collection> transformations() {
         return transformations;
     }
 
+    public Collection> predicates() {
+        return predicates;
+    }
+
+    public Collection> configProviders() {
+        return configProviders;
+    }
+
+    public Collection> restExtensions() {
+        return restExtensions;
+    }
+
+    public Collection> connectorClientConfigPolicies() {
+        return connectorClientConfigPolicies;
+    }
+
     public boolean isEmpty() {
-        return connectors().isEmpty() && converters().isEmpty() && transformations().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 5649213a2f33d..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
@@ -16,7 +16,10 @@
  */
 package org.apache.kafka.connect.runtime.isolation;
 
+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;
 import org.apache.kafka.connect.storage.Converter;
@@ -30,6 +33,9 @@ public enum PluginType {
     CONNECTOR(Connector.class),
     CONVERTER(Converter.class),
     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 9ef7c3a40e33f..8b21e7324af76 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
@@ -34,6 +34,8 @@
 import java.util.List;
 import java.util.Locale;
 import java.util.Set;
+import java.util.TreeSet;
+import java.util.regex.Pattern;
 
 /**
  * Connect plugin utility methods.
@@ -42,7 +44,7 @@ public class PluginUtils {
     private static final Logger log = LoggerFactory.getLogger(PluginUtils.class);
 
     // Be specific about javax packages and exclude those existing in Java SE and Java EE libraries.
-    private static final String BLACKLIST = "^(?:"
+    private static final Pattern EXCLUDE = Pattern.compile("^(?:"
             + "java"
             + "|javax\\.accessibility"
             + "|javax\\.activation"
@@ -117,23 +119,33 @@ public class PluginUtils {
             + "|org\\.omg\\.stub\\.java\\.rmi"
             + "|org\\.w3c\\.dom"
             + "|org\\.xml\\.sax"
-            + "|org\\.apache\\.kafka\\.common"
-            + "|org\\.apache\\.kafka\\.connect"
+            + "|org\\.apache\\.kafka"
             + "|org\\.slf4j"
-            + ")\\..*$";
+            + ")\\..*$");
 
-    private static final String WHITELIST = "^org\\.apache\\.kafka\\.connect\\.(?:"
-            + "transforms\\.(?!Transformation$).*"
+    // 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 INCLUDE pattern), then this base interface or class needs to be excluded in the
+    // regular expression pattern
+    private static final Pattern INCLUDE = Pattern.compile("^org\\.apache\\.kafka\\.(?:connect\\.(?:"
+            + "transforms\\.(?!Transformation|predicates\\.Predicate$).*"
             + "|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$).*"
+            + ")$");
 
     private static final DirectoryStream.Filter PLUGIN_PATH_FILTER = new DirectoryStream
             .Filter() {
         @Override
-        public boolean accept(Path path) throws IOException {
+        public boolean accept(Path path) {
             return Files.isDirectory(path) || isArchive(path) || isClassFile(path);
         }
     };
@@ -146,7 +158,7 @@ public boolean accept(Path path) throws IOException {
      * @return true if this class should be loaded in isolation, false otherwise.
      */
     public static boolean shouldLoadInIsolation(String name) {
-        return !(name.matches(BLACKLIST) && !name.matches(WHITELIST));
+        return !(EXCLUDE.matcher(name).matches() && !INCLUDE.matcher(name).matches());
     }
 
     /**
@@ -208,7 +220,7 @@ public static List pluginLocations(Path topPath) throws IOException {
      */
     public static List pluginUrls(Path topPath) throws IOException {
         boolean containsClassFiles = false;
-        Set archives = new HashSet<>();
+        Set archives = new TreeSet<>();
         LinkedList dfs = new LinkedList<>();
         Set visited = new HashSet<>();
 
@@ -232,16 +244,29 @@ public static List pluginUrls(Path topPath) throws IOException {
 
                 Path adjacent = neighbors.next();
                 if (Files.isSymbolicLink(adjacent)) {
-                    Path symlink = Files.readSymbolicLink(adjacent);
-                    // if symlink is absolute resolve() returns the absolute symlink itself
-                    Path parent = adjacent.getParent();
-                    if (parent == null) {
-                        continue;
-                    }
-                    Path absolute = parent.resolve(symlink).toRealPath();
-                    if (Files.exists(absolute)) {
-                        adjacent = absolute;
-                    } else {
+                    try {
+                        Path symlink = Files.readSymbolicLink(adjacent);
+                        // if symlink is absolute resolve() returns the absolute symlink itself
+                        Path parent = adjacent.getParent();
+                        if (parent == null) {
+                            continue;
+                        }
+                        Path absolute = parent.resolve(symlink).toRealPath();
+                        if (Files.exists(absolute)) {
+                            adjacent = absolute;
+                        } else {
+                            continue;
+                        }
+                    } catch (IOException e) {
+                        // See https://issues.apache.org/jira/browse/KAFKA-6288 for a reported
+                        // failure. Such a failure at this stage is not easily reproducible and
+                        // therefore an exception is caught and ignored after issuing a
+                        // warning. This allows class scanning to continue for non-broken plugins.
+                        log.warn(
+                                "Resolving symbolic link '{}' failed. Ignoring this path.",
+                                adjacent,
+                                e
+                        );
                         continue;
                     }
                 }
@@ -341,8 +366,8 @@ private static String prunePluginName(PluginDesc plugin, String suffix) {
     }
 
     private static class DirectoryEntry {
-        DirectoryStream stream;
-        Iterator iterator;
+        final DirectoryStream stream;
+        final Iterator iterator;
 
         DirectoryEntry(DirectoryStream stream) {
             this.stream = stream;
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 f7a5553126ff7..d507059eacc8c 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
@@ -18,14 +18,22 @@
 
 import org.apache.kafka.common.Configurable;
 import org.apache.kafka.common.config.AbstractConfig;
+import org.apache.kafka.common.config.provider.ConfigProvider;
 import org.apache.kafka.common.utils.Utils;
+import org.apache.kafka.connect.components.Versioned;
 import org.apache.kafka.connect.connector.ConnectRecord;
 import org.apache.kafka.connect.connector.Connector;
 import org.apache.kafka.connect.connector.Task;
 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.runtime.WorkerConfig;
 import org.apache.kafka.connect.storage.Converter;
+import org.apache.kafka.connect.storage.ConverterConfig;
+import org.apache.kafka.connect.storage.ConverterType;
+import org.apache.kafka.connect.storage.HeaderConverter;
 import org.apache.kafka.connect.transforms.Transformation;
+import org.apache.kafka.connect.transforms.predicates.Predicate;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -38,6 +46,12 @@
 import java.util.Set;
 
 public class Plugins {
+
+    public enum ClassLoaderUsage {
+        CURRENT_CLASSLOADER,
+        PLUGINS
+    }
+
     private static final Logger log = LoggerFactory.getLogger(Plugins.class);
     private final DelegatingClassLoader delegatingLoader;
 
@@ -48,13 +62,8 @@ public Plugins(Map props) {
     }
 
     private static DelegatingClassLoader newDelegatingClassLoader(final List paths) {
-        return (DelegatingClassLoader) AccessController.doPrivileged(
-                new PrivilegedAction() {
-                    @Override
-                    public Object run() {
-                        return new DelegatingClassLoader(paths);
-                    }
-                }
+        return AccessController.doPrivileged(
+                (PrivilegedAction) () -> new DelegatingClassLoader(paths)
         );
     }
 
@@ -63,19 +72,35 @@ 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);
         }
     }
 
-    protected static  T newConfiguredPlugin(AbstractConfig config, Class klass) {
-        T plugin = Utils.newInstance(klass);
-        if (plugin instanceof Configurable) {
-            ((Configurable) plugin).configure(config.originals());
+    @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;
         }
-        return plugin;
+        throw new ConnectException(
+            "Failed to find any class that implements " + pluginClass.getSimpleName()
+                + " for the config "
+                + propertyName + ", available classes are: "
+                + pluginNames(plugins)
+        );
     }
 
     @SuppressWarnings("unchecked")
@@ -96,6 +121,12 @@ protected static  Class pluginClass(
         );
     }
 
+    @SuppressWarnings("deprecation")
+    protected static boolean isInternalConverter(String classPropertyName) {
+        return classPropertyName.equals(WorkerConfig.INTERNAL_KEY_CONVERTER_CLASS_CONFIG)
+            || classPropertyName.equals(WorkerConfig.INTERNAL_VALUE_CONVERTER_CLASS_CONFIG);
+    }
+
     public static ClassLoader compareAndSwapLoaders(ClassLoader loader) {
         ClassLoader current = Thread.currentThread().getContextClassLoader();
         if (!current.equals(loader)) {
@@ -137,7 +168,20 @@ public Set> transformations() {
         return delegatingLoader.transformations();
     }
 
+    public Set> predicates() {
+        return delegatingLoader.predicates();
+    }
+
+    public Set> configProviders() {
+        return delegatingLoader.configProviders();
+    }
+
     public Connector newConnector(String connectorClassOrAlias) {
+        Class klass = connectorClass(connectorClassOrAlias);
+        return newPlugin(klass);
+    }
+
+    public Class connectorClass(String connectorClassOrAlias) {
         Class klass;
         try {
             klass = pluginClass(
@@ -177,34 +221,242 @@ public Connector newConnector(String connectorClassOrAlias) {
             PluginDesc entry = matches.get(0);
             klass = entry.pluginClass();
         }
-        return newPlugin(klass);
+        return klass;
     }
 
     public Task newTask(Class taskClass) {
         return newPlugin(taskClass);
     }
 
-    public Converter newConverter(String converterClassOrAlias) {
-        return newConverter(converterClassOrAlias, null);
+    /**
+     * If the given configuration defines a {@link Converter} using the named configuration property, return a new configured instance.
+     *
+     * @param config             the configuration containing the {@link Converter}'s configuration; may not be null
+     * @param classPropertyName  the name of the property that contains the name of the {@link Converter} class; may not be null
+     * @param classLoaderUsage   which classloader should be used
+     * @return the instantiated and configured {@link Converter}; null if the configuration did not define the specified property
+     * @throws ConnectException if the {@link Converter} implementation class could not be found
+     */
+    public Converter newConverter(AbstractConfig config, String classPropertyName, ClassLoaderUsage classLoaderUsage) {
+        if (!config.originals().containsKey(classPropertyName) && !isInternalConverter(classPropertyName)) {
+            // This configuration does not define the converter via the specified property name, and
+            // it does not represent an internal converter (which has a default available)
+            return 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.
+                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();
+                try {
+                    klass = pluginClass(delegatingLoader, converterClassOrAlias, Converter.class);
+                } catch (ClassNotFoundException e) {
+                    throw new ConnectException(
+                            "Failed to find any class that implements Converter and which name matches "
+                            + converterClassOrAlias + ", available converters are: "
+                            + pluginNames(delegatingLoader.converters())
+                    );
+                }
+                break;
+        }
+        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 ...
+        @SuppressWarnings("deprecation")
+        final boolean isKeyConverter = WorkerConfig.KEY_CONVERTER_CLASS_CONFIG.equals(classPropertyName)
+                                     || WorkerConfig.INTERNAL_KEY_CONVERTER_CLASS_CONFIG.equals(classPropertyName);
+
+        // Configure the Converter using only the old configuration mechanism ...
+        String configPrefix = classPropertyName + ".";
+        Map converterConfig = config.originalsWithPrefix(configPrefix);
+        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 (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)) {
+                converterConfig.put(JsonConverterConfig.SCHEMAS_ENABLE_CONFIG, false);
+            }
+        }
+
+        Converter plugin;
+        ClassLoader savedLoader = compareAndSwapLoaders(klass.getClassLoader());
+        try {
+            plugin = newPlugin(klass);
+            plugin.configure(converterConfig, isKeyConverter);
+        } finally {
+            compareAndSwapLoaders(savedLoader);
+        }
+        return plugin;
     }
 
-    public Converter newConverter(String converterClassOrAlias, AbstractConfig config) {
-        Class klass;
+    /**
+     * If the given configuration defines a {@link HeaderConverter} using the named configuration property, return a new configured
+     * instance.
+     *
+     * @param config             the configuration containing the {@link Converter}'s configuration; may not be null
+     * @param classPropertyName  the name of the property that contains the name of the {@link Converter} class; may not be null
+     * @param classLoaderUsage   which classloader should be used
+     * @return the instantiated and configured {@link HeaderConverter}; null if the configuration did not define the specified property
+     * @throws ConnectException if the {@link HeaderConverter} implementation class could not be found
+     */
+    public HeaderConverter newHeaderConverter(AbstractConfig config, String classPropertyName, ClassLoaderUsage classLoaderUsage) {
+        Class klass = null;
+        switch (classLoaderUsage) {
+            case CURRENT_CLASSLOADER:
+                if (!config.originals().containsKey(classPropertyName)) {
+                    // This connector configuration does not define the header converter via the specified property name
+                    return null;
+                }
+                // 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(...)
+                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();
+                try {
+                    klass = pluginClass(
+                            delegatingLoader,
+                            converterClassOrAlias,
+                            HeaderConverter.class
+                    );
+                } catch (ClassNotFoundException e) {
+                    throw new ConnectException(
+                            "Failed to find any class that implements HeaderConverter and which name matches "
+                                    + converterClassOrAlias
+                                    + ", available header converters are: "
+                                    + pluginNames(delegatingLoader.headerConverters())
+                    );
+                }
+        }
+        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 keys:{}{}", System.lineSeparator(), converterConfig.keySet());
+
+        HeaderConverter plugin;
+        ClassLoader savedLoader = compareAndSwapLoaders(klass.getClassLoader());
         try {
-            klass = pluginClass(
-                    delegatingLoader,
-                    converterClassOrAlias,
-                    Converter.class
-            );
+            plugin = newPlugin(klass);
+            plugin.configure(converterConfig);
+        } finally {
+            compareAndSwapLoaders(savedLoader);
+        }
+        return plugin;
+    }
+
+    public ConfigProvider newConfigProvider(AbstractConfig config, String providerPrefix, ClassLoaderUsage classLoaderUsage) {
+        String classPropertyName = providerPrefix + ".class";
+        Map originalConfig = config.originalsStrings();
+        if (!originalConfig.containsKey(classPropertyName)) {
+            // This configuration does not define the config provider via the specified property name
+            return null;
+        }
+        Class klass = null;
+        switch (classLoaderUsage) {
+            case CURRENT_CLASSLOADER:
+                // Attempt to load first with the current classloader, and plugins as a fallback.
+                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);
+                try {
+                    klass = pluginClass(delegatingLoader, configProviderClassOrAlias, ConfigProvider.class);
+                } catch (ClassNotFoundException e) {
+                    throw new ConnectException(
+                            "Failed to find any class that implements ConfigProvider and which name matches "
+                                    + configProviderClassOrAlias + ", available ConfigProviders are: "
+                                    + pluginNames(delegatingLoader.configProviders())
+                    );
+                }
+                break;
+        }
+        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);
+
+        ConfigProvider plugin;
+        ClassLoader savedLoader = compareAndSwapLoaders(klass.getClassLoader());
+        try {
+            plugin = newPlugin(klass);
+            plugin.configure(configProviderConfig);
+        } finally {
+            compareAndSwapLoaders(savedLoader);
+        }
+        return plugin;
+    }
+
+    /**
+     * If the given class names are available in the classloader, return a list of new configured
+     * instances. If the instances implement {@link Configurable}, they are configured with provided {@param config}
+     *
+     * @param klassNames         the list of class names of plugins that needs to instantiated and configured
+     * @param config             the configuration containing the {@link org.apache.kafka.connect.runtime.Worker}'s configuration; may not be {@code null}
+     * @param pluginKlass        the type of the plugin class that is being instantiated
+     * @return the instantiated and configured list of plugins of type ; empty list if the {@param klassNames} is {@code null} or empty
+     * @throws ConnectException if the implementation class could not be found
+     */
+    public  List newPlugins(List klassNames, AbstractConfig config, Class pluginKlass) {
+        List plugins = new ArrayList<>();
+        if (klassNames != null) {
+            for (String klassName : klassNames) {
+                plugins.add(newPlugin(klassName, config, pluginKlass));
+            }
+        }
+        return plugins;
+    }
+
+    public  T newPlugin(String klassName, AbstractConfig config, Class pluginKlass) {
+        T plugin;
+        Class klass;
+        try {
+            klass = pluginClass(delegatingLoader, klassName, pluginKlass);
         } catch (ClassNotFoundException e) {
-            throw new ConnectException(
-                    "Failed to find any class that implements Converter and which name matches "
-                            + converterClassOrAlias
-                            + ", available connectors are: "
-                            + pluginNames(delegatingLoader.converters())
-            );
+            String msg = String.format("Failed to find any class that implements %s and which "
+                                       + "name matches %s", pluginKlass, klassName);
+            throw new ConnectException(msg);
+        }
+        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());
+            }
+        } finally {
+            compareAndSwapLoaders(savedLoader);
         }
-        return config != null ? newConfiguredPlugin(config, klass) : newPlugin(klass);
+        return plugin;
     }
 
     public > Transformation newTranformations(
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/ConnectRestConfigurable.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/ConnectRestConfigurable.java
new file mode 100644
index 0000000000000..0d5cbd6333b48
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/ConnectRestConfigurable.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.connect.runtime.rest;
+
+import org.glassfish.jersey.server.ResourceConfig;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.Map;
+import java.util.Objects;
+
+import javax.ws.rs.core.Configurable;
+import javax.ws.rs.core.Configuration;
+
+/**
+ * The implementation delegates to {@link ResourceConfig} so that we can handle duplicate
+ * registrations deterministically by not re-registering them again.
+ */
+public class ConnectRestConfigurable implements Configurable {
+
+    private static final Logger log = LoggerFactory.getLogger(ConnectRestConfigurable.class);
+
+    private static final boolean ALLOWED_TO_REGISTER = true;
+    private static final boolean NOT_ALLOWED_TO_REGISTER = false;
+
+    private ResourceConfig resourceConfig;
+
+    public ConnectRestConfigurable(ResourceConfig resourceConfig) {
+        Objects.requireNonNull(resourceConfig, "ResourceConfig can't be null");
+        this.resourceConfig = resourceConfig;
+    }
+
+
+    @Override
+    public Configuration getConfiguration() {
+        return resourceConfig.getConfiguration();
+    }
+
+    @Override
+    public ResourceConfig property(String name, Object value) {
+        return resourceConfig.property(name, value);
+    }
+
+    @Override
+    public ResourceConfig register(Object component) {
+        if (allowedToRegister(component)) {
+            resourceConfig.register(component);
+        }
+        return resourceConfig;
+    }
+
+    @Override
+    public ResourceConfig register(Object component, int priority) {
+        if (allowedToRegister(component)) {
+            resourceConfig.register(component, priority);
+        }
+        return resourceConfig;
+    }
+
+    @Override
+    public ResourceConfig register(Object component, Map, Integer> contracts) {
+        if (allowedToRegister(component)) {
+            resourceConfig.register(component, contracts);
+        }
+        return resourceConfig;
+    }
+
+    @Override
+    public ResourceConfig register(Object component, Class... contracts) {
+        if (allowedToRegister(component)) {
+            resourceConfig.register(component, contracts);
+        }
+        return resourceConfig;
+    }
+
+    @Override
+    public ResourceConfig register(Class componentClass, Map, Integer> contracts) {
+        if (allowedToRegister(componentClass)) {
+            resourceConfig.register(componentClass, contracts);
+        }
+        return resourceConfig;
+    }
+
+    @Override
+    public ResourceConfig register(Class componentClass, Class... contracts) {
+        if (allowedToRegister(componentClass)) {
+            resourceConfig.register(componentClass, contracts);
+        }
+        return resourceConfig;
+    }
+
+    @Override
+    public ResourceConfig register(Class componentClass, int priority) {
+        if (allowedToRegister(componentClass)) {
+            resourceConfig.register(componentClass, priority);
+        }
+        return resourceConfig;
+    }
+
+    @Override
+    public ResourceConfig register(Class componentClass) {
+        if (allowedToRegister(componentClass)) {
+            resourceConfig.register(componentClass);
+        }
+        return resourceConfig;
+    }
+
+    private boolean allowedToRegister(Object component) {
+        if (resourceConfig.isRegistered(component)) {
+            log.warn("The resource {} is already registered", component);
+            return NOT_ALLOWED_TO_REGISTER;
+        }
+        return ALLOWED_TO_REGISTER;
+    }
+
+    private boolean allowedToRegister(Class componentClass) {
+        if (resourceConfig.isRegistered(componentClass)) {
+            log.warn("The resource {} is already registered", componentClass);
+            return NOT_ALLOWED_TO_REGISTER;
+        }
+        return ALLOWED_TO_REGISTER;
+    }
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/ConnectRestExtensionContextImpl.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/ConnectRestExtensionContextImpl.java
new file mode 100644
index 0000000000000..cdf282fc90768
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/ConnectRestExtensionContextImpl.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.runtime.rest;
+
+import org.apache.kafka.connect.health.ConnectClusterState;
+import org.apache.kafka.connect.rest.ConnectRestExtensionContext;
+
+import javax.ws.rs.core.Configurable;
+
+public class ConnectRestExtensionContextImpl implements ConnectRestExtensionContext {
+
+    private Configurable configurable;
+    private ConnectClusterState clusterState;
+
+    public ConnectRestExtensionContextImpl(
+        Configurable configurable,
+        ConnectClusterState clusterState
+    ) {
+        this.configurable = configurable;
+        this.clusterState = clusterState;
+    }
+
+    @Override
+    public Configurable configurable() {
+        return configurable;
+    }
+
+    @Override
+    public ConnectClusterState clusterState() {
+        return clusterState;
+    }
+}
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
new file mode 100644
index 0000000000000..b38125971656c
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/RestClient.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.runtime.rest;
+
+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;
+import org.apache.kafka.connect.runtime.rest.util.SSLUtils;
+import org.eclipse.jetty.client.HttpClient;
+import org.eclipse.jetty.client.api.ContentResponse;
+import org.eclipse.jetty.client.api.Request;
+import org.eclipse.jetty.client.util.StringContentProvider;
+import org.eclipse.jetty.http.HttpField;
+import org.eclipse.jetty.http.HttpFields;
+import org.eclipse.jetty.http.HttpStatus;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.ws.rs.core.Response;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeoutException;
+
+public class RestClient {
+    private static final Logger log = LoggerFactory.getLogger(RestClient.class);
+    private static final ObjectMapper JSON_SERDE = new ObjectMapper();
+
+    /**
+     * 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.
+     * @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) {
+        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.createClientSideSslContextFactory(config));
+        } else {
+            client = new HttpClient();
+        }
+
+        client.setFollowRedirects(false);
+
+        try {
+            client.start();
+        } catch (Exception e) {
+            log.error("Failed to start RestClient: ", e);
+            throw new ConnectRestException(Response.Status.INTERNAL_SERVER_ERROR, "Failed to start RestClient: " + e.getMessage(), e);
+        }
+
+        try {
+            String serializedBody = requestBodyData == null ? null : JSON_SERDE.writeValueAsString(requestBodyData);
+            log.trace("Sending {} with input {} to {}", method, serializedBody, url);
+
+            Request req = client.newRequest(url);
+            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();
+
+            int responseCode = res.getStatus();
+            log.debug("Request's response code: {}", responseCode);
+            if (responseCode == HttpStatus.NO_CONTENT_204) {
+                return new HttpResponse<>(responseCode, convertHttpFieldsToMap(res.getHeaders()), null);
+            } else if (responseCode >= 400) {
+                ErrorMessage errorMessage = JSON_SERDE.readValue(res.getContentAsString(), ErrorMessage.class);
+                throw new ConnectRestException(responseCode, errorMessage.errorCode(), errorMessage.message());
+            } else if (responseCode >= 200 && responseCode < 300) {
+                T result = JSON_SERDE.readValue(res.getContentAsString(), responseFormat);
+                return new HttpResponse<>(responseCode, convertHttpFieldsToMap(res.getHeaders()), result);
+            } else {
+                throw new ConnectRestException(Response.Status.INTERNAL_SERVER_ERROR,
+                        Response.Status.INTERNAL_SERVER_ERROR.getStatusCode(),
+                        "Unexpected status code when handling forwarded request: " + responseCode);
+            }
+        } catch (IOException | InterruptedException | TimeoutException | ExecutionException e) {
+            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 {
+            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);
+            }
+        }
+    }
+
+    /**
+     * Convert response parameters from Jetty format (HttpFields)
+     * @param httpFields
+     * @return
+     */
+    private static Map convertHttpFieldsToMap(HttpFields httpFields) {
+        Map headers = new HashMap();
+
+        if (httpFields == null || httpFields.size() == 0)
+            return headers;
+
+        for (HttpField field : httpFields) {
+            headers.put(field.getName(), field.getValue());
+        }
+
+        return headers;
+    }
+
+    public static class HttpResponse {
+        private int status;
+        private Map headers;
+        private T body;
+
+        public HttpResponse(int status, Map headers, T body) {
+            this.status = status;
+            this.headers = headers;
+            this.body = body;
+        }
+
+        public int status() {
+            return status;
+        }
+
+        public Map headers() {
+            return headers;
+        }
+
+        public T body() {
+            return body;
+        }
+    }
+}
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 7afb4e80572ee..408f72e1dd91c 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
@@ -16,51 +16,57 @@
  */
 package org.apache.kafka.connect.runtime.rest;
 
-import com.fasterxml.jackson.core.type.TypeReference;
-import com.fasterxml.jackson.databind.ObjectMapper;
 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.Herder;
 import org.apache.kafka.connect.runtime.WorkerConfig;
-import org.apache.kafka.connect.runtime.rest.entities.ErrorMessage;
+import org.apache.kafka.connect.runtime.health.ConnectClusterDetailsImpl;
+import org.apache.kafka.connect.runtime.health.ConnectClusterStateImpl;
 import org.apache.kafka.connect.runtime.rest.errors.ConnectExceptionMapper;
-import org.apache.kafka.connect.runtime.rest.errors.ConnectRestException;
 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;
 import org.eclipse.jetty.servlet.ServletContextHandler;
 import org.eclipse.jetty.servlet.ServletHolder;
 import org.eclipse.jetty.servlets.CrossOriginFilter;
+import org.eclipse.jetty.servlets.HeaderFilter;
+import org.eclipse.jetty.util.ssl.SslContextFactory;
 import org.glassfish.jersey.server.ResourceConfig;
+import org.glassfish.jersey.server.ServerProperties;
 import org.glassfish.jersey.servlet.ServletContainer;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import javax.servlet.DispatcherType;
+import javax.ws.rs.core.UriBuilder;
 import java.io.IOException;
-import java.io.InputStream;
-import java.io.OutputStream;
-import java.net.HttpURLConnection;
 import java.net.URI;
-import java.net.URL;
-import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.Collections;
 import java.util.EnumSet;
 import java.util.List;
-import java.util.Map;
+import java.util.Locale;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 
-import javax.servlet.DispatcherType;
-import javax.ws.rs.core.Response;
-import javax.ws.rs.core.UriBuilder;
+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.
@@ -68,50 +74,205 @@
 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;
 
-    private static final ObjectMapper JSON_SERDE = new ObjectMapper();
+    private static final String PROTOCOL_HTTP = "http";
+    private static final String PROTOCOL_HTTPS = "https";
 
     private final WorkerConfig config;
+    private ContextHandlerCollection handlers;
     private Server jettyServer;
 
+    private List connectRestExtensions = Collections.emptyList();
+
     /**
      * Create a REST server for this herder using the specified configs.
      */
     public RestServer(WorkerConfig config) {
         this.config = config;
 
-        // To make the advertised port available immediately, we need to do some configuration here
-        String hostname = config.getString(WorkerConfig.REST_HOST_NAME_CONFIG);
-        Integer port = config.getInt(WorkerConfig.REST_PORT_CONFIG);
+        List listeners = parseListeners();
+        List adminListeners = config.getList(WorkerConfig.ADMIN_LISTENERS_CONFIG);
 
         jettyServer = new Server();
+        handlers = new ContextHandlerCollection();
+
+        createConnectors(listeners, adminListeners);
+    }
+
+    @SuppressWarnings("deprecation")
+    List parseListeners() {
+        List listeners = config.getList(WorkerConfig.LISTENERS_CONFIG);
+        if (listeners == null || listeners.size() == 0) {
+            String hostname = config.getString(WorkerConfig.REST_HOST_NAME_CONFIG);
+
+            if (hostname == null)
+                hostname = "";
+
+            listeners = Collections.singletonList(String.format("%s://%s:%d", PROTOCOL_HTTP, hostname, config.getInt(WorkerConfig.REST_PORT_CONFIG)));
+        }
+
+        return listeners;
+    }
+
+    /**
+     * Adds Jetty connector for each configured listener
+     */
+    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);
+            }
+        }
+
+        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 regular (non-admin) Jetty connector according to configuration
+     */
+    public Connector createConnector(String listener) {
+        return createConnector(listener, false);
+    }
 
-        ServerConnector connector = new ServerConnector(jettyServer);
-        if (hostname != null && !hostname.isEmpty())
+    /**
+     * Creates Jetty connector according to configuration
+     */
+    public Connector createConnector(String listener, boolean isAdmin) {
+        Matcher listenerMatcher = LISTENER_PATTERN.matcher(listener);
+
+        if (!listenerMatcher.matches())
+            throw new ConfigException("Listener doesn't have the right format (protocol://hostname:port).");
+
+        String protocol = listenerMatcher.group(1).toLowerCase(Locale.ENGLISH);
+
+        if (!PROTOCOL_HTTP.equals(protocol) && !PROTOCOL_HTTPS.equals(protocol))
+            throw new ConfigException(String.format("Listener protocol must be either \"%s\" or \"%s\".", PROTOCOL_HTTP, PROTOCOL_HTTPS));
+
+        String hostname = listenerMatcher.group(2);
+        int port = Integer.parseInt(listenerMatcher.group(3));
+
+        ServerConnector connector;
+
+        if (PROTOCOL_HTTPS.equals(protocol)) {
+            SslContextFactory ssl;
+            if (isAdmin) {
+                ssl = SSLUtils.createServerSideSslContextFactory(config, ADMIN_LISTENERS_HTTPS_CONFIGS_PREFIX);
+            } else {
+                ssl = SSLUtils.createServerSideSslContextFactory(config);
+            }
+            connector = new ServerConnector(jettyServer, ssl);
+            if (!isAdmin) {
+                connector.setName(String.format("%s_%s%d", PROTOCOL_HTTPS, hostname, port));
+            }
+        } else {
+            connector = new ServerConnector(jettyServer);
+            if (!isAdmin) {
+                connector.setName(String.format("%s_%s%d", PROTOCOL_HTTP, hostname, port));
+            }
+        }
+
+        if (isAdmin) {
+            connector.setName(ADMIN_SERVER_CONNECTOR_NAME);
+        }
+
+        if (!hostname.isEmpty())
             connector.setHost(hostname);
+
         connector.setPort(port);
-        jettyServer.setConnectors(new Connector[]{connector});
+
+        return connector;
+    }
+
+    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 start(Herder herder) {
-        log.info("Starting REST server");
+    public void initializeResources(Herder herder) {
+        log.info("Initializing REST resources");
 
         ResourceConfig resourceConfig = new ResourceConfig();
         resourceConfig.register(new JacksonJsonProvider());
 
-        resourceConfig.register(RootResource.class);
-        resourceConfig.register(new ConnectorsResource(herder));
+        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(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()) {
@@ -125,41 +286,59 @@ public void start(Herder herder) {
             context.addFilter(filterHolder, "/*", EnumSet.of(DispatcherType.REQUEST));
         }
 
+        String headerConfig = config.getString(WorkerConfig.RESPONSE_HTTP_HEADERS_CONFIG);
+        if (headerConfig != null && !headerConfig.trim().isEmpty()) {
+            configureHttpResponsHeaderFilter(context);
+        }
+
         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() {
+        return jettyServer.getURI();
     }
 
     public void stop() {
         log.info("Stopping REST server");
 
         try {
+            for (ConnectRestExtension connectRestExtension : connectRestExtensions) {
+                try {
+                    connectRestExtension.close();
+                } catch (IOException e) {
+                    log.warn("Error while invoking close on " + connectRestExtension.getClass(), e);
+                }
+            }
             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");
@@ -167,107 +346,129 @@ 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());
+
+        String advertisedSecurityProtocol = determineAdvertisedProtocol();
+        ServerConnector serverConnector = findConnector(advertisedSecurityProtocol);
+        builder.scheme(advertisedSecurityProtocol);
+
         String advertisedHostname = config.getString(WorkerConfig.REST_ADVERTISED_HOST_NAME_CONFIG);
         if (advertisedHostname != null && !advertisedHostname.isEmpty())
             builder.host(advertisedHostname);
+        else if (serverConnector != null && serverConnector.getHost() != null && serverConnector.getHost().length() > 0)
+            builder.host(serverConnector.getHost());
+
         Integer advertisedPort = config.getInt(WorkerConfig.REST_ADVERTISED_PORT_CONFIG);
         if (advertisedPort != null)
             builder.port(advertisedPort);
-        else
-            builder.port(config.getInt(WorkerConfig.REST_PORT_CONFIG));
+        else if (serverConnector != null && serverConnector.getPort() > 0)
+            builder.port(serverConnector.getPort());
+
+        log.info("Advertised URI: {}", builder.build());
+
         return builder.build();
     }
 
     /**
-     * @param url               HTTP connection will be established with this url.
-     * @param method            HTTP method ("GET", "POST", "PUT", etc.)
-     * @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.
+     * @return the admin url for this worker. can be null if admin endpoints are disabled.
      */
-    public static  HttpResponse httpRequest(String url, String method, Object requestBodyData,
-                                    TypeReference responseFormat) {
-        HttpURLConnection connection = null;
-        try {
-            String serializedBody = requestBodyData == null ? null : JSON_SERDE.writeValueAsString(requestBodyData);
-            log.debug("Sending {} with input {} to {}", method, serializedBody, url);
-
-            connection = (HttpURLConnection) new URL(url).openConnection();
-            connection.setRequestMethod(method);
-
-            connection.setRequestProperty("User-Agent", "kafka-connect");
-            connection.setRequestProperty("Accept", "application/json");
+    public URI adminUrl() {
+        ServerConnector adminConnector = null;
+        for (Connector connector : jettyServer.getConnectors()) {
+            if (ADMIN_SERVER_CONNECTOR_NAME.equals(connector.getName()))
+                adminConnector = (ServerConnector) connector;
+        }
 
-            // connection.getResponseCode() implicitly calls getInputStream, so always set to true.
-            // On the other hand, leaving this out breaks nothing.
-            connection.setDoInput(true);
+        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;
+            }
+        }
 
-            connection.setUseCaches(false);
+        UriBuilder builder = UriBuilder.fromUri(jettyServer.getURI());
+        builder.port(adminConnector.getLocalPort());
 
-            if (requestBodyData != null) {
-                connection.setRequestProperty("Content-Type", "application/json");
-                connection.setDoOutput(true);
+        return builder.build();
+    }
 
-                OutputStream os = connection.getOutputStream();
-                os.write(serializedBody.getBytes(StandardCharsets.UTF_8));
-                os.flush();
-                os.close();
-            }
+    String determineAdvertisedProtocol() {
+        String advertisedSecurityProtocol = config.getString(WorkerConfig.REST_ADVERTISED_LISTENER_CONFIG);
+        if (advertisedSecurityProtocol == null) {
+            String listeners = (String) config.originals().get(WorkerConfig.LISTENERS_CONFIG);
+
+            if (listeners == null)
+                return PROTOCOL_HTTP;
+            else
+                listeners = listeners.toLowerCase(Locale.ENGLISH);
+
+            if (listeners.contains(String.format("%s://", PROTOCOL_HTTP)))
+                return PROTOCOL_HTTP;
+            else if (listeners.contains(String.format("%s://", PROTOCOL_HTTPS)))
+                return PROTOCOL_HTTPS;
+            else
+                return PROTOCOL_HTTP;
+        } else {
+            return advertisedSecurityProtocol.toLowerCase(Locale.ENGLISH);
+        }
+    }
 
-            int responseCode = connection.getResponseCode();
-            if (responseCode == HttpURLConnection.HTTP_NO_CONTENT) {
-                return new HttpResponse<>(responseCode, connection.getHeaderFields(), null);
-            } else if (responseCode >= 400) {
-                InputStream es = connection.getErrorStream();
-                ErrorMessage errorMessage = JSON_SERDE.readValue(es, ErrorMessage.class);
-                es.close();
-                throw new ConnectRestException(responseCode, errorMessage.errorCode(), errorMessage.message());
-            } else if (responseCode >= 200 && responseCode < 300) {
-                InputStream is = connection.getInputStream();
-                T result = JSON_SERDE.readValue(is, responseFormat);
-                is.close();
-                return new HttpResponse<>(responseCode, connection.getHeaderFields(), result);
-            } else {
-                throw new ConnectRestException(Response.Status.INTERNAL_SERVER_ERROR,
-                        Response.Status.INTERNAL_SERVER_ERROR.getStatusCode(),
-                        "Unexpected status code when handling forwarded request: " + responseCode);
-            }
-        } catch (IOException e) {
-            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 (connection != null)
-                connection.disconnect();
+    /**
+     * 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()) {
+            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;
     }
 
-    public static class HttpResponse {
-        private int status;
-        private Map> headers;
-        private T body;
+    void registerRestExtensions(Herder herder, ResourceConfig resourceConfig) {
+        connectRestExtensions = herder.plugins().newPlugins(
+            config.getList(WorkerConfig.REST_EXTENSION_CLASSES_CONFIG),
+            config, ConnectRestExtension.class);
 
-        public HttpResponse(int status, Map> headers, T body) {
-            this.status = status;
-            this.headers = headers;
-            this.body = body;
-        }
+        long herderRequestTimeoutMs = ConnectorsResource.REQUEST_TIMEOUT_MS;
 
-        public int status() {
-            return status;
-        }
+        Integer rebalanceTimeoutMs = config.getRebalanceTimeout();
 
-        public Map> headers() {
-            return headers;
+        if (rebalanceTimeoutMs != null) {
+            herderRequestTimeoutMs = Math.min(herderRequestTimeoutMs, rebalanceTimeoutMs.longValue());
         }
 
-        public T body() {
-            return body;
+        ConnectClusterDetails connectClusterDetails = new ConnectClusterDetailsImpl(
+            herder.kafkaClusterId()
+        );
+
+        ConnectRestExtensionContext connectRestExtensionContext =
+            new ConnectRestExtensionContextImpl(
+                new ConnectRestConfigurable(resourceConfig),
+                new ConnectClusterStateImpl(herderRequestTimeoutMs, connectClusterDetails, herder)
+            );
+        for (ConnectRestExtension connectRestExtension : connectRestExtensions) {
+            connectRestExtension.register(connectRestExtensionContext);
         }
+
     }
 
     public static String urlJoin(String base, String path) {
@@ -277,4 +478,14 @@ public static String urlJoin(String base, String path) {
             return base + path;
     }
 
-}
\ No newline at end of file
+    /**
+     * Register header filter to ServletContextHandler.
+     * @param context The serverlet context handler
+     */
+    protected void configureHttpResponsHeaderFilter(ServletContextHandler context) {
+        String headerConfig = config.getString(WorkerConfig.RESPONSE_HTTP_HEADERS_CONFIG);
+        FilterHolder headerFilterHolder = new FilterHolder(HeaderFilter.class);
+        headerFilterHolder.setInitParameter("headerConfig", headerConfig);
+        context.addFilter(headerFilterHolder, "/*", EnumSet.of(DispatcherType.REQUEST));
+    }
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/entities/ActiveTopicsInfo.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/entities/ActiveTopicsInfo.java
new file mode 100644
index 0000000000000..b43c5aa88da84
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/entities/ActiveTopicsInfo.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.rest.entities;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+import java.util.Collection;
+
+public class ActiveTopicsInfo {
+    private final String connector;
+    private final Collection topics;
+
+    @JsonCreator
+    public ActiveTopicsInfo(String connector, @JsonProperty("topics") Collection topics) {
+        this.connector = connector;
+        this.topics = topics;
+    }
+
+    public String connector() {
+        return connector;
+    }
+
+    @JsonProperty
+    public Collection topics() {
+        return topics;
+    }
+
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/entities/ConnectorInfo.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/entities/ConnectorInfo.java
index 9a10d7488770f..f36ee74e980be 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/entities/ConnectorInfo.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/entities/ConnectorInfo.java
@@ -71,12 +71,13 @@ public boolean equals(Object o) {
         ConnectorInfo that = (ConnectorInfo) o;
         return Objects.equals(name, that.name) &&
                 Objects.equals(config, that.config) &&
-                Objects.equals(tasks, that.tasks);
+                Objects.equals(tasks, that.tasks) &&
+                Objects.equals(type, that.type);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(name, config, tasks);
+        return Objects.hash(name, config, tasks, type);
     }
 
 }
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/entities/ConnectorStateInfo.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/entities/ConnectorStateInfo.java
index 80192ca064489..6280473af964d 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/entities/ConnectorStateInfo.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/entities/ConnectorStateInfo.java
@@ -90,7 +90,10 @@ public String trace() {
     }
 
     public static class ConnectorState extends AbstractState {
-        public ConnectorState(String state, String worker, String msg) {
+        @JsonCreator
+        public ConnectorState(@JsonProperty("state") String state,
+                              @JsonProperty("worker_id") String worker,
+                              @JsonProperty("msg") String msg) {
             super(state, worker, msg);
         }
     }
@@ -98,7 +101,11 @@ public ConnectorState(String state, String worker, String msg) {
     public static class TaskState extends AbstractState implements Comparable {
         private final int id;
 
-        public TaskState(int id, String state, String worker, String msg) {
+        @JsonCreator
+        public TaskState(@JsonProperty("id") int id,
+                         @JsonProperty("state") String state,
+                         @JsonProperty("worker_id") String worker,
+                         @JsonProperty("msg") String msg) {
             super(state, worker, msg);
             this.id = id;
         }
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 25ce7311bac38..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,16 +16,26 @@
  */
 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;
 
 public class ServerInfo {
-    private String version;
-    private String commit;
+    private final String version;
+    private final String commit;
+    private final String kafkaClusterId;
 
-    public ServerInfo() {
-        version = AppInfoParser.getVersion();
-        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
@@ -37,4 +47,9 @@ public String version() {
     public String commit() {
         return commit;
     }
+
+    @JsonProperty("kafka_cluster_id")
+    public String clusterId() {
+        return kafkaClusterId;
+    }
 }
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 24eb93b8c0d5f..1f6161e2ebfb8 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
@@ -22,12 +22,14 @@
 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;
+import org.apache.kafka.connect.runtime.rest.errors.ConnectRestException;
 import org.apache.kafka.connect.tools.MockConnector;
 import org.apache.kafka.connect.tools.MockSinkConnector;
 import org.apache.kafka.connect.tools.MockSourceConnector;
 import org.apache.kafka.connect.tools.SchemaSourceConnector;
 import org.apache.kafka.connect.tools.VerifiableSinkConnector;
 import org.apache.kafka.connect.tools.VerifiableSourceConnector;
+import org.apache.kafka.connect.util.FutureCallback;
 
 import javax.ws.rs.BadRequestException;
 import javax.ws.rs.Consumes;
@@ -37,11 +39,14 @@
 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.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 
 @Path("/connector-plugins")
 @Produces(MediaType.APPLICATION_JSON)
@@ -78,7 +83,19 @@ public ConfigInfos validateConfigs(
             );
         }
 
-        return herder.validateConnectorConfig(connectorConfig);
+        // the validated configs don't need to be logged
+        FutureCallback validationCallback = new FutureCallback<>();
+        herder.validateConnectorConfig(connectorConfig, validationCallback, false);
+
+        try {
+            return validationCallback.get(ConnectorsResource.REQUEST_TIMEOUT_MS, TimeUnit.SECONDS);
+        } catch (TimeoutException e) {
+            // This timeout is for the operation itself. None of the timeout error codes are relevant, so internal server
+            // error is the best option
+            throw new ConnectRestException(Response.Status.INTERNAL_SERVER_ERROR.getStatusCode(), "Request timed out");
+        } catch (InterruptedException e) {
+            throw new ConnectRestException(Response.Status.INTERNAL_SERVER_ERROR.getStatusCode(), "Request interrupted");
+        }
     }
 
     @GET
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 e681a680acccf..b3219b2204bb3 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,11 +17,19 @@
 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.WorkerConfig;
 import org.apache.kafka.connect.runtime.distributed.RebalanceNeededException;
 import org.apache.kafka.connect.runtime.distributed.RequestTargetException;
-import org.apache.kafka.connect.runtime.rest.RestServer;
+import org.apache.kafka.connect.runtime.rest.InternalRequestSignature;
+import org.apache.kafka.connect.runtime.rest.RestClient;
+import org.apache.kafka.connect.runtime.rest.entities.ActiveTopicsInfo;
 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.CreateConnectorRequest;
@@ -43,61 +51,117 @@
 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.Collections;
+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;
 
+import static org.apache.kafka.connect.runtime.WorkerConfig.TOPIC_TRACKING_ALLOW_RESET_CONFIG;
+import static org.apache.kafka.connect.runtime.WorkerConfig.TOPIC_TRACKING_ENABLE_CONFIG;
+
 @Path("/connectors")
 @Produces(MediaType.APPLICATION_JSON)
 @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;
+    // Mutable for integration testing; otherwise, some tests would take at least REQUEST_TIMEOUT_MS
+    // to run
+    private static long requestTimeoutMs = REQUEST_TIMEOUT_MS;
 
     private final Herder herder;
+    private final WorkerConfig config;
     @javax.ws.rs.core.Context
     private ServletContext context;
+    private final boolean isTopicTrackingDisabled;
+    private final boolean isTopicTrackingResetDisabled;
 
-    public ConnectorsResource(Herder herder) {
+    public ConnectorsResource(Herder herder, WorkerConfig config) {
         this.herder = herder;
+        this.config = config;
+        isTopicTrackingDisabled = !config.getBoolean(TOPIC_TRACKING_ENABLE_CONFIG);
+        isTopicTrackingResetDisabled = !config.getBoolean(TOPIC_TRACKING_ALLOW_RESET_CONFIG);
+    }
+
+    // For testing purposes only
+    public static void setRequestTimeout(long requestTimeoutMs) {
+        ConnectorsResource.requestTimeoutMs = requestTimeoutMs;
+    }
+
+    public static void resetRequestTimeout() {
+        ConnectorsResource.requestTimeoutMs = REQUEST_TIMEOUT_MS;
     }
 
     @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
+    ) {
+        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 expansion 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 {
-        String name = createRequest.name();
-        if (name.contains("/")) {
-            throw new BadRequestException("connector name should not contain '/'");
-        }
+        // 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
+        // allows null values)
+        String name = createRequest.name() == null ? "" : createRequest.name().trim();
+
         Map configs = createRequest.config();
-        if (!configs.containsKey(ConnectorConfig.NAME_CONFIG))
-            configs.put(ConnectorConfig.NAME_CONFIG, name);
+        checkAndPutConnectorConfigName(name, configs);
 
         FutureCallback> cb = new FutureCallback<>();
         herder.putConnectorConfig(name, configs, false, cb);
-        Herder.Created info = completeOrForwardRequest(cb, "/connectors", "POST", createRequest,
+        Herder.Created info = completeOrForwardRequest(cb, "/connectors", "POST", headers, createRequest,
                 new TypeReference() { }, new CreatedConnectorInfoTranslator(), forward);
 
         URI location = UriBuilder.fromUri("/connectors").path(name).build();
@@ -107,44 +171,67 @@ 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);
+        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);
+        return completeOrForwardRequest(cb, "/connectors/" + connector + "/config", "GET", headers, null, forward);
     }
 
     @GET
     @Path("/{connector}/status")
-    public ConnectorStateInfo getConnectorStatus(final @PathParam("connector") String connector) throws Throwable {
+    public ConnectorStateInfo getConnectorStatus(final @PathParam("connector") String connector) {
         return herder.connectorStatus(connector);
     }
 
+    @GET
+    @Path("/{connector}/topics")
+    public Response getConnectorActiveTopics(final @PathParam("connector") String connector) {
+        if (isTopicTrackingDisabled) {
+            throw new ConnectRestException(Response.Status.FORBIDDEN.getStatusCode(),
+                    "Topic tracking is disabled.");
+        }
+        ActiveTopicsInfo info = herder.connectorActiveTopics(connector);
+        return Response.ok(Collections.singletonMap(info.connector(), info)).build();
+    }
+
+    @PUT
+    @Path("/{connector}/topics/reset")
+    public Response resetConnectorActiveTopics(final @PathParam("connector") String connector, final @Context HttpHeaders headers) {
+        if (isTopicTrackingDisabled) {
+            throw new ConnectRestException(Response.Status.FORBIDDEN.getStatusCode(),
+                    "Topic tracking is disabled.");
+        }
+        if (isTopicTrackingResetDisabled) {
+            throw new ConnectRestException(Response.Status.FORBIDDEN.getStatusCode(),
+                    "Topic tracking reset is disabled.");
+        }
+        herder.resetConnectorActiveTopics(connector);
+        return Response.accepted().build();
+    }
+
     @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<>();
-        String includedName = connectorConfig.get(ConnectorConfig.NAME_CONFIG);
-        if (includedName != null) {
-            if (!includedName.equals(connector))
-                throw new BadRequestException("Connector name configuration (" + includedName + ") doesn't match connector name in the URL (" + connector + ")");
-        } else {
-            connectorConfig.put(ConnectorConfig.NAME_CONFIG, connector);
-        }
+        checkAndPutConnectorConfigName(connector, connectorConfig);
 
         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();
@@ -158,15 +245,16 @@ 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);
+        completeOrForwardRequest(cb, "/connectors/" + connector + "/restart", "POST", headers, null, forward);
     }
 
     @PUT
     @Path("/{connector}/pause")
-    public Response pauseConnector(@PathParam("connector") String connector) {
+    public Response pauseConnector(@PathParam("connector") String connector, final @Context HttpHeaders headers) {
         herder.pauseConnector(connector);
         return Response.accepted().build();
     }
@@ -181,27 +269,31 @@ public Response resumeConnector(@PathParam("connector") String connector) {
     @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>() {
+        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 @PathParam("task") Integer task) throws Throwable {
+                                                      final @Context HttpHeaders headers,
+                                                      final @PathParam("task") Integer task) {
         return herder.taskStatus(new ConnectorTaskId(connector, task));
     }
 
@@ -209,20 +301,34 @@ public ConnectorStateInfo.TaskState getTaskStatus(final @PathParam("connector")
     @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);
+        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);
+        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
+    // object. Throw BadRequestException on mismatch, otherwise put connectorName in config
+    private void checkAndPutConnectorConfigName(String connectorName, Map connectorConfig) {
+        String includedName = connectorConfig.get(ConnectorConfig.NAME_CONFIG);
+        if (includedName != null) {
+            if (!includedName.equals(connectorName))
+                throw new BadRequestException("Connector name configuration (" + includedName + ") doesn't match connector name in the URL (" + connectorName + ")");
+        } else {
+            connectorConfig.put(ConnectorConfig.NAME_CONFIG, connectorName);
+        }
     }
 
     // Wait for a FutureCallback to complete. If it succeeds, return the parsed response. If it fails, try to forward the
@@ -230,12 +336,13 @@ public void destroyConnector(final @PathParam("connector") String connector,
     private  T completeOrForwardRequest(FutureCallback cb,
                                               String path,
                                               String method,
+                                              HttpHeaders headers,
                                               Object body,
                                               TypeReference resultType,
                                               Translator translator,
                                               Boolean forward) throws Throwable {
         try {
-            return cb.get(REQUEST_TIMEOUT_MS, TimeUnit.MILLISECONDS);
+            return cb.get(requestTimeoutMs, TimeUnit.MILLISECONDS);
         } catch (ExecutionException e) {
             Throwable cause = e.getCause();
 
@@ -246,13 +353,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(RestServer.httpRequest(forwardUrl, method, body, resultType));
+                    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.
@@ -274,30 +388,30 @@ 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 {
-        T translate(RestServer.HttpResponse response);
+        T translate(RestClient.HttpResponse response);
     }
 
     private static class IdentityTranslator implements Translator {
         @Override
-        public T translate(RestServer.HttpResponse response) {
+        public T translate(RestClient.HttpResponse response) {
             return response.body();
         }
     }
 
     private static class CreatedConnectorInfoTranslator implements Translator, ConnectorInfo> {
         @Override
-        public Herder.Created translate(RestServer.HttpResponse response) {
+        public Herder.Created translate(RestClient.HttpResponse response) {
             boolean created = response.status() == 201;
             return new Herder.Created<>(created, response.body());
         }
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 9d94001ecfa04..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,6 +16,7 @@
  */
 package org.apache.kafka.connect.runtime.rest.resources;
 
+import org.apache.kafka.connect.runtime.Herder;
 import org.apache.kafka.connect.runtime.rest.entities.ServerInfo;
 
 import javax.ws.rs.GET;
@@ -27,9 +28,15 @@
 @Produces(MediaType.APPLICATION_JSON)
 public class RootResource {
 
+    private final Herder herder;
+
+    public RootResource(Herder herder) {
+        this.herder = herder;
+    }
+
     @GET
     @Path("/")
     public ServerInfo serverInfo() {
-        return new ServerInfo();
+        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
new file mode 100644
index 0000000000000..6b391d96adaa7
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/rest/util/SSLUtils.java
@@ -0,0 +1,171 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.util;
+
+import org.apache.kafka.common.config.SslConfigs;
+import org.apache.kafka.common.config.internals.BrokerSecurityConfigs;
+import org.apache.kafka.common.config.types.Password;
+import org.apache.kafka.connect.runtime.WorkerConfig;
+import org.eclipse.jetty.util.ssl.SslContextFactory;
+
+import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
+import java.util.regex.Pattern;
+
+/**
+ * Helper class for setting up SSL for RestServer and RestClient
+ */
+public class SSLUtils {
+
+    private static final Pattern COMMA_WITH_WHITESPACE = Pattern.compile("\\s*,\\s*");
+
+
+    /**
+     * Configures SSL/TLS for HTTPS Jetty Server using configs with the given prefix
+     */
+    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
+     */
+    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.");
+
+        final SslContextFactory.Client ssl = new SslContextFactory.Client();
+
+        configureSslContextFactoryKeyStore(ssl, sslConfigValues);
+        configureSslContextFactoryTrustStore(ssl, sslConfigValues);
+        configureSslContextFactoryAlgorithms(ssl, sslConfigValues);
+        configureSslContextFactoryEndpointIdentification(ssl, sslConfigValues);
+
+        return ssl;
+    }
+
+    /**
+     * Configures KeyStore related settings in SslContextFactory
+     */
+    protected static void configureSslContextFactoryKeyStore(SslContextFactory ssl, Map sslConfigValues) {
+        ssl.setKeyStoreType((String) getOrDefault(sslConfigValues, SslConfigs.SSL_KEYSTORE_TYPE_CONFIG, SslConfigs.DEFAULT_SSL_KEYSTORE_TYPE));
+
+        String sslKeystoreLocation = (String) sslConfigValues.get(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG);
+        if (sslKeystoreLocation != null)
+            ssl.setKeyStorePath(sslKeystoreLocation);
+
+        Password sslKeystorePassword = (Password) sslConfigValues.get(SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG);
+        if (sslKeystorePassword != null)
+            ssl.setKeyStorePassword(sslKeystorePassword.value());
+
+        Password sslKeyPassword = (Password) sslConfigValues.get(SslConfigs.SSL_KEY_PASSWORD_CONFIG);
+        if (sslKeyPassword != null)
+            ssl.setKeyManagerPassword(sslKeyPassword.value());
+    }
+
+    protected static Object getOrDefault(Map configMap, String key, Object defaultValue) {
+        if (configMap.containsKey(key))
+            return configMap.get(key);
+
+        return defaultValue;
+    }
+
+    /**
+     * Configures TrustStore related settings in SslContextFactory
+     */
+    protected static void configureSslContextFactoryTrustStore(SslContextFactory ssl, Map sslConfigValues) {
+        ssl.setTrustStoreType((String) getOrDefault(sslConfigValues, SslConfigs.SSL_TRUSTSTORE_TYPE_CONFIG, SslConfigs.DEFAULT_SSL_TRUSTSTORE_TYPE));
+
+        String sslTruststoreLocation = (String) sslConfigValues.get(SslConfigs.SSL_TRUSTSTORE_LOCATION_CONFIG);
+        if (sslTruststoreLocation != null)
+            ssl.setTrustStorePath(sslTruststoreLocation);
+
+        Password sslTruststorePassword = (Password) sslConfigValues.get(SslConfigs.SSL_TRUSTSTORE_PASSWORD_CONFIG);
+        if (sslTruststorePassword != null)
+            ssl.setTrustStorePassword(sslTruststorePassword.value());
+    }
+
+    /**
+     * Configures Protocol, Algorithm and Provider related settings in SslContextFactory
+     */
+    @SuppressWarnings("unchecked")
+    protected static void configureSslContextFactoryAlgorithms(SslContextFactory ssl, Map sslConfigValues) {
+        List sslEnabledProtocols = (List) getOrDefault(sslConfigValues, SslConfigs.SSL_ENABLED_PROTOCOLS_CONFIG, Arrays.asList(COMMA_WITH_WHITESPACE.split(SslConfigs.DEFAULT_SSL_ENABLED_PROTOCOLS)));
+        ssl.setIncludeProtocols(sslEnabledProtocols.toArray(new String[sslEnabledProtocols.size()]));
+
+        String sslProvider = (String) sslConfigValues.get(SslConfigs.SSL_PROVIDER_CONFIG);
+        if (sslProvider != null)
+            ssl.setProvider(sslProvider);
+
+        ssl.setProtocol((String) getOrDefault(sslConfigValues, SslConfigs.SSL_PROTOCOL_CONFIG, SslConfigs.DEFAULT_SSL_PROTOCOL));
+
+        List sslCipherSuites = (List) sslConfigValues.get(SslConfigs.SSL_CIPHER_SUITES_CONFIG);
+        if (sslCipherSuites != null)
+            ssl.setIncludeCipherSuites(sslCipherSuites.toArray(new String[sslCipherSuites.size()]));
+
+        ssl.setKeyManagerFactoryAlgorithm((String) getOrDefault(sslConfigValues, SslConfigs.SSL_KEYMANAGER_ALGORITHM_CONFIG, SslConfigs.DEFAULT_SSL_KEYMANGER_ALGORITHM));
+
+        String sslSecureRandomImpl = (String) sslConfigValues.get(SslConfigs.SSL_SECURE_RANDOM_IMPLEMENTATION_CONFIG);
+        if (sslSecureRandomImpl != null)
+            ssl.setSecureRandomAlgorithm(sslSecureRandomImpl);
+
+        ssl.setTrustManagerFactoryAlgorithm((String) getOrDefault(sslConfigValues, SslConfigs.SSL_TRUSTMANAGER_ALGORITHM_CONFIG, SslConfigs.DEFAULT_SSL_TRUSTMANAGER_ALGORITHM));
+    }
+
+    /**
+     * Configures Protocol, Algorithm and Provider related settings in SslContextFactory
+     */
+    protected static void configureSslContextFactoryEndpointIdentification(SslContextFactory ssl, Map sslConfigValues) {
+        String sslEndpointIdentificationAlg = (String) sslConfigValues.get(SslConfigs.SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG);
+        if (sslEndpointIdentificationAlg != null)
+            ssl.setEndpointIdentificationAlgorithm(sslEndpointIdentificationAlg);
+    }
+
+    /**
+     * Configures Authentication related settings in SslContextFactory
+     */
+    protected static void configureSslContextFactoryAuthentication(SslContextFactory.Server ssl, Map sslConfigValues) {
+        String sslClientAuth = (String) getOrDefault(sslConfigValues, BrokerSecurityConfigs.SSL_CLIENT_AUTH_CONFIG, "none");
+        switch (sslClientAuth) {
+            case "requested":
+                ssl.setWantClientAuth(true);
+                break;
+            case "required":
+                ssl.setNeedClientAuth(true);
+                break;
+            default:
+                ssl.setNeedClientAuth(false);
+                ssl.setWantClientAuth(false);
+        }
+    }
+}
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 e9ec0f9092f45..fc351f708d360 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,17 +16,22 @@
  */
 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;
 import org.apache.kafka.connect.runtime.AbstractHerder;
 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.ConfigInfos;
 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;
@@ -42,6 +47,12 @@
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledFuture;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicLong;
 
 
 /**
@@ -50,39 +61,62 @@
 public class StandaloneHerder extends AbstractHerder {
     private static final Logger log = LoggerFactory.getLogger(StandaloneHerder.class);
 
+    private final AtomicLong requestSeqNum = new AtomicLong();
+    private final ScheduledExecutorService requestExecutorService;
+
     private ClusterConfigState configState;
 
-    public StandaloneHerder(Worker worker) {
-        this(worker, worker.workerId(), new MemoryStatusBackingStore(), new MemoryConfigBackingStore());
+    public StandaloneHerder(Worker worker, String kafkaClusterId,
+                            ConnectorClientConfigOverridePolicy connectorClientConfigOverridePolicy) {
+        this(worker,
+                worker.workerId(),
+                kafkaClusterId,
+                new MemoryStatusBackingStore(),
+                new MemoryConfigBackingStore(worker.configTransformer()),
+             connectorClientConfigOverridePolicy);
     }
 
     // visible for testing
     StandaloneHerder(Worker worker,
                      String workerId,
+                     String kafkaClusterId,
                      StatusBackingStore statusBackingStore,
-                     MemoryConfigBackingStore configBackingStore) {
-        super(worker, workerId, 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());
     }
 
+    @Override
     public synchronized void start() {
         log.info("Herder starting");
         startServices();
+        running = true;
         log.info("Herder started");
     }
 
+    @Override
     public synchronized void stop() {
         log.info("Herder stopping");
+        requestExecutorService.shutdown();
+        try {
+            if (!requestExecutorService.awaitTermination(30, TimeUnit.SECONDS))
+                requestExecutorService.shutdownNow();
+        } catch (InterruptedException e) {
+            // ignore
+        }
 
         // 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);
+            worker.stopAndAwaitConnector(connName);
         }
         stopServices();
+        running = false;
         log.info("Herder stopped");
     }
 
@@ -93,12 +127,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;
@@ -106,34 +140,19 @@ public synchronized void connectorInfo(String connName, Callback
         callback.onCompletion(null, connectorInfo);
     }
 
-    private ConnectorInfo createConnectorInfo(String connector) {
+    private synchronized ConnectorInfo createConnectorInfo(String connector) {
         if (!configState.contains(connector))
             return null;
-        Map config = configState.connectorConfig(connector);
+        Map config = configState.rawConnectorConfig(connector);
         return new ConnectorInfo(connector, config, configState.tasks(connector),
             connectorTypeForClass(config.get(ConnectorConfig.CONNECTOR_CLASS_CONFIG)));
     }
 
     @Override
-    protected Map config(String connName) {
+    protected synchronized Map config(String connName) {
         return configState.connectorConfig(connName);
     }
 
-    @Override
-    public void connectorConfig(String connName, final Callback> callback) {
-        // Subset of connectorInfo, so piggy back on that implementation
-        connectorInfo(connName, new Callback() {
-            @Override
-            public void onCompletion(Throwable error, ConnectorInfo result) {
-                if (error != null) {
-                    callback.onCompletion(error, null);
-                    return;
-                }
-                callback.onCompletion(null, result.config());
-            }
-        });
-    }
-
     @Override
     public synchronized void deleteConnectorConfig(String connName, Callback> callback) {
         try {
@@ -144,7 +163,7 @@ public synchronized void deleteConnectorConfig(String connName, Callback(false, null));
@@ -160,30 +179,58 @@ public synchronized void putConnectorConfig(String connName,
                                                 boolean allowReplace,
                                                 final Callback> callback) {
         try {
-            if (maybeAddConfigErrors(validateConnectorConfig(config), callback)) {
+            validateConnectorConfig(config, (error, configInfos) -> {
+                if (error != null) {
+                    callback.onCompletion(error, null);
+                    return;
+                }
+
+                requestExecutorService.submit(
+                    () -> putConnectorConfig(connName, config, allowReplace, callback, configInfos)
+                );
+            });
+        } catch (Throwable t) {
+            callback.onCompletion(t, null);
+        }
+    }
+
+    private synchronized void putConnectorConfig(String connName,
+                                                 final Map config,
+                                                 boolean allowReplace,
+                                                 final Callback> callback,
+                                                 ConfigInfos configInfos) {
+        try {
+            if (maybeAddConfigErrors(configInfos, callback)) {
                 return;
             }
 
-            boolean created = false;
+            final boolean created;
             if (configState.contains(connName)) {
                 if (!allowReplace) {
                     callback.onCompletion(new AlreadyExistsException("Connector " + connName + " already exists"), null);
                     return;
                 }
-                worker.stopConnector(connName);
+                worker.stopAndAwaitConnector(connName);
+                created = false;
             } else {
                 created = true;
             }
 
-            if (!startConnector(config)) {
-                callback.onCompletion(new ConnectException("Failed to start connector: " + connName), null);
-                return;
-            }
+            configBackingStore.putConnectorConfig(connName, config);
 
-            updateConnectorTasks(connName);
-            callback.onCompletion(null, new Created<>(created, createConnectorInfo(connName)));
-        } catch (ConnectException e) {
-            callback.onCompletion(e, null);
+            startConnector(connName, (error, result) -> {
+                if (error != null) {
+                    callback.onCompletion(error, null);
+                    return;
+                }
+
+                requestExecutorService.submit(() -> {
+                    updateConnectorTasks(connName);
+                    callback.onCompletion(null, new Created<>(created, createConnectorInfo(connName)));
+                });
+            });
+        } catch (Throwable t) {
+            callback.onCompletion(t, null);
         }
     }
 
@@ -205,12 +252,12 @@ public synchronized void taskConfigs(String connName, Callback> c
 
         List result = new ArrayList<>();
         for (ConnectorTaskId taskId : configState.tasks(connName))
-            result.add(new TaskInfo(taskId, configState.taskConfig(taskId)));
+            result.add(new TaskInfo(taskId, configState.rawTaskConfig(taskId)));
         callback.onCompletion(null, result);
     }
 
     @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.");
     }
 
@@ -226,7 +273,7 @@ public synchronized void restartTask(ConnectorTaskId taskId, Callback cb)
 
         TargetState targetState = configState.targetState(taskId.connector());
         worker.stopAndAwaitTask(taskId);
-        if (worker.startTask(taskId, connConfigProps, taskConfigProps, this, targetState))
+        if (worker.startTask(taskId, configState, connConfigProps, taskConfigProps, this, targetState))
             cb.onCompletion(null, null);
         else
             cb.onCompletion(new ConnectException("Failed to start task: " + taskId), null);
@@ -237,19 +284,23 @@ public synchronized void restartConnector(String connName, Callback cb) {
         if (!configState.contains(connName))
             cb.onCompletion(new NotFoundException("Connector " + connName + " not found", null), null);
 
-        Map config = configState.connectorConfig(connName);
-        worker.stopConnector(connName);
-        if (startConnector(config))
-            cb.onCompletion(null, null);
-        else
-            cb.onCompletion(new ConnectException("Failed to start connector: " + connName), null);
+        worker.stopAndAwaitConnector(connName);
+
+        startConnector(connName, (error, result) -> cb.onCompletion(error, null));
     }
 
-    private boolean startConnector(Map connectorProps) {
-        String connName = connectorProps.get(ConnectorConfig.NAME_CONFIG);
-        configBackingStore.putConnectorConfig(connName, connectorProps);
+    @Override
+    public synchronized HerderRequest restartConnector(long delayMs, final String connName, final Callback cb) {
+        ScheduledFuture future = requestExecutorService.schedule(
+            () -> restartConnector(connName, cb), delayMs, TimeUnit.MILLISECONDS);
+
+        return new StandaloneHerderRequest(requestSeqNum.incrementAndGet(), future);
+    }
+
+    private void startConnector(String connName, Callback onStart) {
+        Map connConfigs = configState.connectorConfig(connName);
         TargetState targetState = configState.targetState(connName);
-        return worker.startConnector(connName, connectorProps, new HerderConnectorContext(this, connName), this, targetState);
+        worker.startConnector(connName, connConfigs, new HerderConnectorContext(this, connName), this, targetState, onStart);
     }
 
     private List> recomputeTaskConfigs(String connName) {
@@ -257,7 +308,7 @@ private List> recomputeTaskConfigs(String connName) {
 
         ConnectorConfig connConfig = worker.isSinkConnector(connName) ?
             new SinkConnectorConfig(plugins(), config) :
-            new SourceConnectorConfig(plugins(), config);
+            new SourceConnectorConfig(plugins(), config, worker.isTopicCreationEnabled());
 
         return worker.connectorTaskConfigs(connName, connConfig);
     }
@@ -267,7 +318,7 @@ private void createConnectorTasks(String connName, TargetState initialState) {
 
         for (ConnectorTaskId taskId : configState.tasks(connName)) {
             Map taskConfigMap = configState.taskConfig(taskId);
-            worker.startTask(taskId, connConfigs, taskConfigMap, this, initialState);
+            worker.startTask(taskId, configState, connConfigs, taskConfigMap, this, initialState);
         }
     }
 
@@ -276,12 +327,13 @@ private void removeConnectorTasks(String connName) {
         if (!tasks.isEmpty()) {
             worker.stopAndAwaitTasks(tasks);
             configBackingStore.removeTaskConfigs(connName);
+            tasks.forEach(this::onDeletion);
         }
     }
 
     private void updateConnectorTasks(String connName) {
         if (!worker.isRunning(connName)) {
-            log.info("Skipping reconfiguration of connector {} since it is not running", connName);
+            log.info("Skipping update of connector {} since it is not running", connName);
             return;
         }
 
@@ -290,7 +342,8 @@ private void updateConnectorTasks(String connName) {
 
         if (!newTaskConfigs.equals(oldTaskConfigs)) {
             removeConnectorTasks(connName);
-            configBackingStore.putTaskConfigs(connName, newTaskConfigs);
+            List> rawTaskConfigs = reverseTransform(connName, configState, newTaskConfigs);
+            configBackingStore.putTaskConfigs(connName, rawTaskConfigs);
             createConnectorTasks(connName, configState.targetState(connName));
         }
     }
@@ -332,11 +385,53 @@ public void onConnectorTargetStateChange(String connector) {
             synchronized (StandaloneHerder.this) {
                 configState = configBackingStore.snapshot();
                 TargetState targetState = configState.targetState(connector);
-                worker.setTargetState(connector, targetState);
-                if (targetState == TargetState.STARTED)
-                    updateConnectorTasks(connector);
+                worker.setTargetState(connector, targetState, (error, newState) -> {
+                    if (error != null) {
+                        log.error("Failed to transition connector {} to target state {}", connector, targetState, error);
+                        return;
+                    }
+
+                    if (newState == TargetState.STARTED) {
+                        requestExecutorService.submit(() -> {
+                            updateConnectorTasks(connector);
+                        });
+                    }
+                });
             }
         }
+
+        @Override
+        public void onSessionKeyUpdate(SessionKey sessionKey) {
+            // no-op
+        }
     }
 
+    static class StandaloneHerderRequest implements HerderRequest {
+        private final long seq;
+        private final ScheduledFuture future;
+
+        public StandaloneHerderRequest(long seq, ScheduledFuture future) {
+            this.seq = seq;
+            this.future = future;
+        }
+
+        @Override
+        public void cancel() {
+            future.cancel(false);
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (!(o instanceof StandaloneHerderRequest))
+                return false;
+            StandaloneHerderRequest other = (StandaloneHerderRequest) o;
+            return seq == other.seq;
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(seq);
+        }
+    }
 }
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/FileOffsetBackingStore.java b/connect/runtime/src/main/java/org/apache/kafka/connect/storage/FileOffsetBackingStore.java
index d868f625ffb3a..8f828fbcf0d99 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/storage/FileOffsetBackingStore.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/storage/FileOffsetBackingStore.java
@@ -25,12 +25,11 @@
 
 import java.io.EOFException;
 import java.io.File;
-import java.io.FileInputStream;
-import java.io.FileNotFoundException;
-import java.io.FileOutputStream;
 import java.io.IOException;
 import java.io.ObjectOutputStream;
 import java.nio.ByteBuffer;
+import java.nio.file.Files;
+import java.nio.file.NoSuchFileException;
 import java.util.HashMap;
 import java.util.Map;
 
@@ -69,7 +68,7 @@ public synchronized void stop() {
 
     @SuppressWarnings("unchecked")
     private void load() {
-        try (SafeObjectInputStream is = new SafeObjectInputStream(new FileInputStream(file))) {
+        try (SafeObjectInputStream is = new SafeObjectInputStream(Files.newInputStream(file.toPath()))) {
             Object obj = is.readObject();
             if (!(obj instanceof HashMap))
                 throw new ConnectException("Expected HashMap but found " + obj.getClass());
@@ -80,17 +79,17 @@ private void load() {
                 ByteBuffer value = (mapEntry.getValue() != null) ? ByteBuffer.wrap(mapEntry.getValue()) : null;
                 data.put(key, value);
             }
-        } catch (FileNotFoundException | EOFException e) {
-            // FileNotFoundException: Ignore, may be new.
+        } catch (NoSuchFileException | EOFException e) {
+            // NoSuchFileException: Ignore, may be new.
             // EOFException: Ignore, this means the file was missing or corrupt
         } catch (IOException | ClassNotFoundException e) {
             throw new ConnectException(e);
         }
     }
 
+    @Override
     protected void save() {
-        try {
-            ObjectOutputStream os = new ObjectOutputStream(new FileOutputStream(file));
+        try (ObjectOutputStream os = new ObjectOutputStream(Files.newOutputStream(file.toPath()))) {
             Map raw = new HashMap<>();
             for (Map.Entry mapEntry : data.entrySet()) {
                 byte[] key = (mapEntry.getKey() != null) ? mapEntry.getKey().array() : null;
@@ -98,7 +97,6 @@ protected void save() {
                 raw.put(key, value);
             }
             os.writeObject(raw);
-            os.close();
         } catch (IOException e) {
             throw new ConnectException(e);
         }
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 dd69c300b69b3..fbcc35bb4bcb4 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
@@ -21,6 +21,7 @@
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.producer.ProducerConfig;
 import org.apache.kafka.common.config.ConfigException;
+import org.apache.kafka.common.config.TopicConfig;
 import org.apache.kafka.common.serialization.ByteArrayDeserializer;
 import org.apache.kafka.common.serialization.ByteArraySerializer;
 import org.apache.kafka.common.serialization.StringDeserializer;
@@ -33,19 +34,25 @@
 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;
 import org.apache.kafka.connect.runtime.distributed.ClusterConfigState;
 import org.apache.kafka.connect.runtime.distributed.DistributedConfig;
 import org.apache.kafka.connect.util.Callback;
+import org.apache.kafka.connect.util.ConnectUtils;
 import org.apache.kafka.connect.util.ConnectorTaskId;
 import org.apache.kafka.connect.util.KafkaBasedLog;
 import org.apache.kafka.connect.util.TopicAdmin;
 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.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
@@ -175,6 +182,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
@@ -189,6 +198,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;
 
@@ -215,23 +231,28 @@ 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<>();
 
     private final Map connectorTargetStates = new HashMap<>();
 
-    public KafkaConfigBackingStore(Converter converter, WorkerConfig config) {
+    private final WorkerConfigTransformer configTransformer;
+
+    public KafkaConfigBackingStore(Converter converter, WorkerConfig config, WorkerConfigTransformer configTransformer) {
         this.lock = new Object();
         this.started = false;
         this.converter = converter;
         this.offset = -1;
 
         this.topic = config.getString(DistributedConfig.CONFIG_TOPIC_CONFIG);
-        if (this.topic.equals(""))
+        if (this.topic == null || this.topic.trim().length() == 0)
             throw new ConfigException("Must specify topic for connector configuration.");
 
         configLog = setupAndCreateKafkaBasedLog(this.topic, config);
+        this.configTransformer = configTransformer;
     }
 
     @Override
@@ -245,6 +266,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");
     }
@@ -262,15 +293,17 @@ 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),
                     new HashMap<>(taskConfigs),
-                    new HashSet<>(inconsistent)
+                    new HashSet<>(inconsistent),
+                    configTransformer
             );
         }
     }
@@ -291,7 +324,7 @@ public boolean contains(String connector) {
      */
     @Override
     public void putConnectorConfig(String connector, Map properties) {
-        log.debug("Writing connector configuration {} for connector {} configuration", properties, connector);
+        log.debug("Writing connector configuration for connector '{}'", connector);
         Struct connectConfig = new Struct(CONNECTOR_CONFIGURATION_V0);
         connectConfig.put("properties", properties);
         byte[] serializedConfig = converter.fromConnectData(topic, CONNECTOR_CONFIGURATION_V0, connectConfig);
@@ -304,7 +337,7 @@ public void putConnectorConfig(String connector, Map properties)
      */
     @Override
     public void removeConnectorConfig(String connector) {
-        log.debug("Removing connector configuration for connector {}", connector);
+        log.debug("Removing connector configuration for connector '{}'", connector);
         try {
             configLog.send(CONNECTOR_KEY(connector), null);
             configLog.send(TARGET_STATE_KEY(connector), null);
@@ -358,7 +391,7 @@ public void putTaskConfigs(String connector, List> configs)
             Struct connectConfig = new Struct(TASK_CONFIGURATION_V0);
             connectConfig.put("properties", taskConfig);
             byte[] serializedConfig = converter.fromConnectData(topic, TASK_CONFIGURATION_V0, connectConfig);
-            log.debug("Writing configuration for task " + index + " configuration: " + taskConfig);
+            log.debug("Writing configuration for connector '{}' task {}", connector, index);
             ConnectorTaskId connectorTaskId = new ConnectorTaskId(connector, index);
             configLog.send(TASK_KEY(connectorTaskId), serializedConfig);
             index++;
@@ -375,7 +408,7 @@ public void putTaskConfigs(String connector, List> configs)
             Struct connectConfig = new Struct(CONNECTOR_TASKS_COMMIT_V0);
             connectConfig.put("tasks", taskCount);
             byte[] serializedConfig = converter.fromConnectData(topic, CONNECTOR_TASKS_COMMIT_V0, connectConfig);
-            log.debug("Writing commit for connector " + connector + " with " + taskCount + " tasks.");
+            log.debug("Writing commit for connector '{}' with {} tasks.", connector, taskCount);
             configLog.send(COMMIT_TASKS_KEY(connector), serializedConfig);
 
             // Read to end to ensure all the commit messages have been written
@@ -404,25 +437,49 @@ 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 producerProps = new HashMap<>();
-        producerProps.putAll(config.originals());
+        String clusterId = ConnectUtils.lookupKafkaClusterId(config);
+        Map originals = config.originals();
+        Map producerProps = new HashMap<>(originals);
         producerProps.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, StringSerializer.class.getName());
         producerProps.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, ByteArraySerializer.class.getName());
-        producerProps.put(ProducerConfig.RETRIES_CONFIG, Integer.MAX_VALUE);
+        producerProps.put(ProducerConfig.DELIVERY_TIMEOUT_MS_CONFIG, Integer.MAX_VALUE);
+        ConnectUtils.addMetricsContextProperties(producerProps, config, clusterId);
 
-        Map consumerProps = new HashMap<>();
-        consumerProps.putAll(config.originals());
+        Map consumerProps = new HashMap<>(originals);
         consumerProps.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class.getName());
         consumerProps.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class.getName());
-
-        Map adminProps = new HashMap<>(config.originals());
-        NewTopic topicDescription = TopicAdmin.defineTopic(topic).
-                compacted().
-                partitions(1).
-                replicationFactor(config.getShort(DistributedConfig.CONFIG_STORAGE_REPLICATION_FACTOR_CONFIG)).
-                build();
+        ConnectUtils.addMetricsContextProperties(consumerProps, config, clusterId);
+
+        Map adminProps = new HashMap<>(originals);
+        ConnectUtils.addMetricsContextProperties(adminProps, config, clusterId);
+        Map topicSettings = config instanceof DistributedConfig
+                                            ? ((DistributedConfig) config).configStorageTopicSettings()
+                                            : Collections.emptyMap();
+        NewTopic topicDescription = TopicAdmin.defineTopic(topic)
+                .config(topicSettings) // first so that we override user-supplied settings as needed
+                .compacted()
+                .partitions(1)
+                .replicationFactor(config.getShort(DistributedConfig.CONFIG_STORAGE_REPLICATION_FACTOR_CONFIG))
+                .build();
 
         return createKafkaBasedLog(topic, producerProps, consumerProps, new ConsumeCallback(), topicDescription, adminProps);
     }
@@ -434,8 +491,16 @@ private KafkaBasedLog createKafkaBasedLog(String topic, Map newTopics = admin.createTopics(topicDescription);
+                    if (!newTopics.contains(topic)) {
+                        // It already existed, so check that the topic cleanup policy is compact only and not delete
+                        log.debug("Using admin client to check cleanup policy of '{}' topic is '{}'", topic, TopicConfig.CLEANUP_POLICY_COMPACT);
+                        admin.verifyTopicCleanupPolicyOnlyCompact(topic,
+                                DistributedConfig.CONFIG_TOPIC_CONFIG, "connector configurations");
+                    }
                 }
             }
         };
@@ -483,17 +548,17 @@ public void onCompletion(Throwable error, ConsumerRecord record)
                         }
                         Object targetState = ((Map) value.value()).get("state");
                         if (!(targetState instanceof String)) {
-                            log.error("Invalid data for target state for connector ({}): 'state' field should be a Map but is {}",
+                            log.error("Invalid data for target state for connector '{}': 'state' field should be a Map but is {}",
                                     connectorName, targetState == null ? null : targetState.getClass());
                             return;
                         }
 
                         try {
                             TargetState state = TargetState.valueOf((String) targetState);
-                            log.debug("Setting target state for connector {} to {}", connectorName, targetState);
+                            log.debug("Setting target state for connector '{}' to {}", connectorName, targetState);
                             connectorTargetStates.put(connectorName, state);
                         } catch (IllegalArgumentException e) {
-                            log.error("Invalid target state for connector ({}): {}", connectorName, targetState);
+                            log.error("Invalid target state for connector '{}': {}", connectorName, targetState);
                             return;
                         }
                     }
@@ -510,22 +575,24 @@ public void onCompletion(Throwable error, ConsumerRecord record)
                 synchronized (lock) {
                     if (value.value() == null) {
                         // Connector deletion will be written as a null value
-                        log.info("Removed connector " + connectorName + " due to null configuration. This is usually intentional and does not indicate an issue.");
+                        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
                         if (!(value.value() instanceof Map)) {
-                            log.error("Found connector configuration (" + record.key() + ") in wrong format: " + value.value().getClass());
+                            log.error("Found configuration for connector '{}' in wrong format: {}", record.key(), value.value().getClass());
                             return;
                         }
                         Object newConnectorConfig = ((Map) value.value()).get("properties");
                         if (!(newConnectorConfig instanceof Map)) {
-                            log.error("Invalid data for connector config ({}): properties field should be a Map but is {}", connectorName,
-                                    newConnectorConfig == null ? null : newConnectorConfig.getClass());
+                            log.error("Invalid data for config for connector '{}': 'properties' field should be a Map but is {}",
+                                      connectorName, newConnectorConfig == null ? null : newConnectorConfig.getClass());
                             return;
                         }
-                        log.debug("Updating configuration for connector " + connectorName + " configuration: " + newConnectorConfig);
+                        log.debug("Updating configuration for connector '{}'", connectorName);
                         connectorConfigs.put(connectorName, (Map) newConnectorConfig);
 
                         // Set the initial state of the connector to STARTED, which ensures that any connectors
@@ -544,17 +611,21 @@ public void onCompletion(Throwable error, ConsumerRecord record)
                 synchronized (lock) {
                     ConnectorTaskId taskId = parseTaskId(record.key());
                     if (taskId == null) {
-                        log.error("Ignoring task configuration because " + record.key() + " couldn't be parsed as a task config key");
+                        log.error("Ignoring task configuration because {} couldn't be parsed as a task config key", record.key());
+                        return;
+                    }
+                    if (value.value() == null) {
+                        log.error("Ignoring task configuration for task {} because it is unexpectedly null", taskId);
                         return;
                     }
                     if (!(value.value() instanceof Map)) {
-                        log.error("Ignoring task configuration for task " + taskId + " because it is in the wrong format: " + value.value());
+                        log.error("Ignoring task configuration for task {} because the value is not a Map but is {}", taskId, value.value().getClass());
                         return;
                     }
 
                     Object newTaskConfig = ((Map) value.value()).get("properties");
                     if (!(newTaskConfig instanceof Map)) {
-                        log.error("Invalid data for task config (" + taskId + "): properties filed should be a Map but is " + newTaskConfig.getClass());
+                        log.error("Invalid data for config of task {} 'properties' field should be a Map but is {}", taskId, newTaskConfig.getClass());
                         return;
                     }
 
@@ -563,7 +634,7 @@ public void onCompletion(Throwable error, ConsumerRecord record)
                         deferred = new HashMap<>();
                         deferredTaskUpdates.put(taskId.connector(), deferred);
                     }
-                    log.debug("Storing new config for task " + taskId + " this will wait for a commit message before the new config will take effect. New config: " + newTaskConfig);
+                    log.debug("Storing new config for task {}; this will wait for a commit message before the new config will take effect.", taskId);
                     deferred.put(taskId, (Map) newTaskConfig);
                 }
             } else if (record.key().startsWith(COMMIT_TASKS_PREFIX)) {
@@ -592,7 +663,7 @@ public void onCompletion(Throwable error, ConsumerRecord record)
                     // resolve this (i.e., get the connector to recommit its configuration). This inconsistent state is
                     // exposed in the snapshots provided via ClusterConfigState so they are easy to handle.
                     if (!(value.value() instanceof Map)) { // Schema-less, so we get maps instead of structs
-                        log.error("Ignoring connector tasks configuration commit for connector " + connectorName + " because it is in the wrong format: " + value.value());
+                        log.error("Ignoring connector tasks configuration commit for connector '{}' because it is in the wrong format: {}", connectorName, value.value());
                         return;
                     }
                     Map> deferred = deferredTaskUpdates.get(connectorName);
@@ -607,12 +678,12 @@ public void onCompletion(Throwable error, ConsumerRecord record)
                         // historical data, in which case we would not have applied any updates yet and there will be no
                         // task config data already committed for the connector, so we shouldn't have to clear any data
                         // out. All we need to do is add the flag marking it inconsistent.
-                        log.debug("We have an incomplete set of task configs for connector " + connectorName + " probably due to compaction. So we are not doing anything with the new configuration.");
+                        log.debug("We have an incomplete set of task configs for connector '{}' probably due to compaction. So we are not doing anything with the new configuration.", connectorName);
                         inconsistent.add(connectorName);
                     } else {
                         if (deferred != null) {
                             taskConfigs.putAll(deferred);
-                            updatedTasks.addAll(taskConfigs.keySet());
+                            updatedTasks.addAll(deferred.keySet());
                         }
                         inconsistent.remove(connectorName);
                     }
@@ -627,8 +698,45 @@ public void onCompletion(Throwable error, ConsumerRecord record)
 
                 if (started)
                     updateListener.onTaskConfigUpdate(updatedTasks);
+            } else if (record.key().equals(SESSION_KEY_KEY)) {
+                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());
+                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 cb4f0897887e2..8408f99c10576 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
@@ -22,12 +22,14 @@
 import org.apache.kafka.clients.producer.ProducerConfig;
 import org.apache.kafka.clients.producer.RecordMetadata;
 import org.apache.kafka.common.config.ConfigException;
+import org.apache.kafka.common.config.TopicConfig;
 import org.apache.kafka.common.serialization.ByteArrayDeserializer;
 import org.apache.kafka.common.serialization.ByteArraySerializer;
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.connect.runtime.WorkerConfig;
 import org.apache.kafka.connect.runtime.distributed.DistributedConfig;
 import org.apache.kafka.connect.util.Callback;
+import org.apache.kafka.connect.util.ConnectUtils;
 import org.apache.kafka.connect.util.ConvertingFutureCallback;
 import org.apache.kafka.connect.util.KafkaBasedLog;
 import org.apache.kafka.connect.util.TopicAdmin;
@@ -36,8 +38,10 @@
 
 import java.nio.ByteBuffer;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
@@ -62,28 +66,35 @@ public class KafkaOffsetBackingStore implements OffsetBackingStore {
     @Override
     public void configure(final WorkerConfig config) {
         String topic = config.getString(DistributedConfig.OFFSET_STORAGE_TOPIC_CONFIG);
-        if (topic.equals(""))
+        if (topic == null || topic.trim().length() == 0)
             throw new ConfigException("Offset storage topic must be specified");
 
+        String clusterId = ConnectUtils.lookupKafkaClusterId(config);
         data = new HashMap<>();
 
-        Map producerProps = new HashMap<>();
-        producerProps.putAll(config.originals());
+        Map originals = config.originals();
+        Map producerProps = new HashMap<>(originals);
         producerProps.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, ByteArraySerializer.class.getName());
         producerProps.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, ByteArraySerializer.class.getName());
-        producerProps.put(ProducerConfig.RETRIES_CONFIG, Integer.MAX_VALUE);
+        producerProps.put(ProducerConfig.DELIVERY_TIMEOUT_MS_CONFIG, Integer.MAX_VALUE);
+        ConnectUtils.addMetricsContextProperties(producerProps, config, clusterId);
 
-        Map consumerProps = new HashMap<>();
-        consumerProps.putAll(config.originals());
+        Map consumerProps = new HashMap<>(originals);
         consumerProps.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class.getName());
         consumerProps.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class.getName());
-
-        Map adminProps = new HashMap<>(config.originals());
-        NewTopic topicDescription = TopicAdmin.defineTopic(topic).
-                compacted().
-                partitions(config.getInt(DistributedConfig.OFFSET_STORAGE_PARTITIONS_CONFIG)).
-                replicationFactor(config.getShort(DistributedConfig.OFFSET_STORAGE_REPLICATION_FACTOR_CONFIG)).
-                build();
+        ConnectUtils.addMetricsContextProperties(consumerProps, config, clusterId);
+
+        Map adminProps = new HashMap<>(originals);
+        ConnectUtils.addMetricsContextProperties(adminProps, config, clusterId);
+        Map topicSettings = config instanceof DistributedConfig
+                                            ? ((DistributedConfig) config).offsetStorageTopicSettings()
+                                            : Collections.emptyMap();
+        NewTopic topicDescription = TopicAdmin.defineTopic(topic)
+                .config(topicSettings) // first so that we override user-supplied settings as needed
+                .compacted()
+                .partitions(config.getInt(DistributedConfig.OFFSET_STORAGE_PARTITIONS_CONFIG))
+                .replicationFactor(config.getShort(DistributedConfig.OFFSET_STORAGE_REPLICATION_FACTOR_CONFIG))
+                .build();
 
         offsetLog = createKafkaBasedLog(topic, producerProps, consumerProps, consumedCallback, topicDescription, adminProps);
     }
@@ -95,8 +106,16 @@ private KafkaBasedLog createKafkaBasedLog(String topic, Map newTopics = admin.createTopics(topicDescription);
+                    if (!newTopics.contains(topic)) {
+                        // It already existed, so check that the topic cleanup policy is compact only and not delete
+                        log.debug("Using admin client to check cleanup policy for '{}' topic is '{}'", topic, TopicConfig.CLEANUP_POLICY_COMPACT);
+                        admin.verifyTopicCleanupPolicyOnlyCompact(topic,
+                                DistributedConfig.OFFSET_STORAGE_TOPIC_CONFIG, "source connector offsets");
+                    }
                 }
             }
         };
@@ -118,9 +137,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 +248,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/KafkaStatusBackingStore.java b/connect/runtime/src/main/java/org/apache/kafka/connect/storage/KafkaStatusBackingStore.java
index 11f951a73c348..5d6057d799f26 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/storage/KafkaStatusBackingStore.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/storage/KafkaStatusBackingStore.java
@@ -22,6 +22,7 @@
 import org.apache.kafka.clients.producer.ProducerConfig;
 import org.apache.kafka.clients.producer.RecordMetadata;
 import org.apache.kafka.common.config.ConfigException;
+import org.apache.kafka.common.config.TopicConfig;
 import org.apache.kafka.common.errors.RetriableException;
 import org.apache.kafka.common.serialization.ByteArrayDeserializer;
 import org.apache.kafka.common.serialization.ByteArraySerializer;
@@ -36,9 +37,11 @@
 import org.apache.kafka.connect.runtime.AbstractStatus;
 import org.apache.kafka.connect.runtime.ConnectorStatus;
 import org.apache.kafka.connect.runtime.TaskStatus;
+import org.apache.kafka.connect.runtime.TopicStatus;
 import org.apache.kafka.connect.runtime.WorkerConfig;
 import org.apache.kafka.connect.runtime.distributed.DistributedConfig;
 import org.apache.kafka.connect.util.Callback;
+import org.apache.kafka.connect.util.ConnectUtils;
 import org.apache.kafka.connect.util.ConnectorTaskId;
 import org.apache.kafka.connect.util.KafkaBasedLog;
 import org.apache.kafka.connect.util.Table;
@@ -49,11 +52,15 @@
 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.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
 
 /**
  * StatusBackingStore implementation which uses a compacted topic for storage
@@ -80,14 +87,22 @@
 public class KafkaStatusBackingStore implements StatusBackingStore {
     private static final Logger log = LoggerFactory.getLogger(KafkaStatusBackingStore.class);
 
-    private static final String TASK_STATUS_PREFIX = "status-task-";
-    private static final String CONNECTOR_STATUS_PREFIX = "status-connector-";
+    public static final String TASK_STATUS_PREFIX = "status-task-";
+    public static final String CONNECTOR_STATUS_PREFIX = "status-connector-";
+    public static final String TOPIC_STATUS_PREFIX = "status-topic-";
+    public static final String TOPIC_STATUS_SEPARATOR = ":connector-";
 
     public static final String STATE_KEY_NAME = "state";
     public static final String TRACE_KEY_NAME = "trace";
     public static final String WORKER_ID_KEY_NAME = "worker_id";
     public static final String GENERATION_KEY_NAME = "generation";
 
+    public static final String TOPIC_STATE_KEY = "topic";
+    public static final String TOPIC_NAME_KEY = "name";
+    public static final String TOPIC_CONNECTOR_KEY = "connector";
+    public static final String TOPIC_TASK_KEY = "task";
+    public static final String TOPIC_DISCOVER_TIMESTAMP_KEY = "discoverTimestamp";
+
     private static final Schema STATUS_SCHEMA_V0 = SchemaBuilder.struct()
             .field(STATE_KEY_NAME, Schema.STRING_SCHEMA)
             .field(TRACE_KEY_NAME, SchemaBuilder.string().optional().build())
@@ -95,12 +110,26 @@ public class KafkaStatusBackingStore implements StatusBackingStore {
             .field(GENERATION_KEY_NAME, Schema.INT32_SCHEMA)
             .build();
 
+    private static final Schema TOPIC_STATUS_VALUE_SCHEMA_V0 = SchemaBuilder.struct()
+            .field(TOPIC_NAME_KEY, Schema.STRING_SCHEMA)
+            .field(TOPIC_CONNECTOR_KEY, Schema.STRING_SCHEMA)
+            .field(TOPIC_TASK_KEY, Schema.INT32_SCHEMA)
+            .field(TOPIC_DISCOVER_TIMESTAMP_KEY, Schema.INT64_SCHEMA)
+            .build();
+
+    private static final Schema TOPIC_STATUS_SCHEMA_V0 = SchemaBuilder.map(
+            Schema.STRING_SCHEMA,
+            TOPIC_STATUS_VALUE_SCHEMA_V0
+    ).build();
+
     private final Time time;
     private final Converter converter;
-    private final Table> tasks;
-    private final Map> connectors;
+    //visible for testing
+    protected final Table> tasks;
+    protected final Map> connectors;
+    protected final ConcurrentMap> topics;
 
-    private String topic;
+    private String statusTopic;
     private KafkaBasedLog kafkaLog;
     private int generation;
 
@@ -109,38 +138,47 @@ public KafkaStatusBackingStore(Time time, Converter converter) {
         this.converter = converter;
         this.tasks = new Table<>();
         this.connectors = new HashMap<>();
+        this.topics = new ConcurrentHashMap<>();
     }
 
     // visible for testing
-    KafkaStatusBackingStore(Time time, Converter converter, String topic, KafkaBasedLog kafkaLog) {
+    KafkaStatusBackingStore(Time time, Converter converter, String statusTopic, KafkaBasedLog kafkaLog) {
         this(time, converter);
         this.kafkaLog = kafkaLog;
-        this.topic = topic;
+        this.statusTopic = statusTopic;
     }
 
     @Override
     public void configure(final WorkerConfig config) {
-        this.topic = config.getString(DistributedConfig.STATUS_STORAGE_TOPIC_CONFIG);
-        if (topic.equals(""))
+        this.statusTopic = config.getString(DistributedConfig.STATUS_STORAGE_TOPIC_CONFIG);
+        if (this.statusTopic == null || this.statusTopic.trim().length() == 0)
             throw new ConfigException("Must specify topic for connector status.");
 
-        Map producerProps = new HashMap<>();
-        producerProps.putAll(config.originals());
+        String clusterId = ConnectUtils.lookupKafkaClusterId(config);
+        Map originals = config.originals();
+        Map producerProps = new HashMap<>(originals);
         producerProps.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, StringSerializer.class.getName());
         producerProps.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, ByteArraySerializer.class.getName());
         producerProps.put(ProducerConfig.RETRIES_CONFIG, 0); // we handle retries in this class
+        ConnectUtils.addMetricsContextProperties(producerProps, config, clusterId);
 
-        Map consumerProps = new HashMap<>();
-        consumerProps.putAll(config.originals());
+        Map consumerProps = new HashMap<>(originals);
         consumerProps.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class.getName());
         consumerProps.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class.getName());
+        ConnectUtils.addMetricsContextProperties(consumerProps, config, clusterId);
+
+        Map adminProps = new HashMap<>(originals);
+        ConnectUtils.addMetricsContextProperties(adminProps, config, clusterId);
 
-        Map adminProps = new HashMap<>(config.originals());
-        NewTopic topicDescription = TopicAdmin.defineTopic(topic).
-                compacted().
-                partitions(config.getInt(DistributedConfig.STATUS_STORAGE_PARTITIONS_CONFIG)).
-                replicationFactor(config.getShort(DistributedConfig.STATUS_STORAGE_REPLICATION_FACTOR_CONFIG)).
-                build();
+        Map topicSettings = config instanceof DistributedConfig
+                                            ? ((DistributedConfig) config).statusStorageTopicSettings()
+                                            : Collections.emptyMap();
+        NewTopic topicDescription = TopicAdmin.defineTopic(statusTopic)
+                .config(topicSettings) // first so that we override user-supplied settings as needed
+                .compacted()
+                .partitions(config.getInt(DistributedConfig.STATUS_STORAGE_PARTITIONS_CONFIG))
+                .replicationFactor(config.getShort(DistributedConfig.STATUS_STORAGE_REPLICATION_FACTOR_CONFIG))
+                .build();
 
         Callback> readCallback = new Callback>() {
             @Override
@@ -148,7 +186,7 @@ public void onCompletion(Throwable error, ConsumerRecord record)
                 read(record);
             }
         };
-        this.kafkaLog = createKafkaBasedLog(topic, producerProps, consumerProps, readCallback, topicDescription, adminProps);
+        this.kafkaLog = createKafkaBasedLog(statusTopic, producerProps, consumerProps, readCallback, topicDescription, adminProps);
     }
 
     private KafkaBasedLog createKafkaBasedLog(String topic, Map producerProps,
@@ -158,8 +196,16 @@ private KafkaBasedLog createKafkaBasedLog(String topic, Map newTopics = admin.createTopics(topicDescription);
+                    if (!newTopics.contains(topic)) {
+                        // It already existed, so check that the topic cleanup policy is compact only and not delete
+                        log.debug("Using admin client to check cleanup policy of '{}' topic is '{}'", topic, TopicConfig.CLEANUP_POLICY_COMPACT);
+                        admin.verifyTopicCleanupPolicyOnlyCompact(topic,
+                                DistributedConfig.STATUS_STORAGE_TOPIC_CONFIG, "connector and task statuses");
+                    }
                 }
             }
         };
@@ -199,13 +245,18 @@ public void putSafe(final TaskStatus status) {
         sendTaskStatus(status, true);
     }
 
+    @Override
+    public void put(final TopicStatus status) {
+        sendTopicStatus(status.connector(), status.topic(), status);
+    }
+
     @Override
     public void flush() {
         kafkaLog.flush();
     }
 
     private void sendConnectorStatus(final ConnectorStatus status, boolean safeWrite) {
-        String connector  = status.id();
+        String connector = status.id();
         CacheEntry entry = getOrAdd(connector);
         String key = CONNECTOR_STATUS_PREFIX + connector;
         send(key, status, entry, safeWrite);
@@ -218,6 +269,25 @@ private void sendTaskStatus(final TaskStatus status, boolean safeWrite) {
         send(key, status, entry, safeWrite);
     }
 
+    private void sendTopicStatus(final String connector, final String topic, final TopicStatus status) {
+        String key = TOPIC_STATUS_PREFIX + topic + TOPIC_STATUS_SEPARATOR + connector;
+
+        final byte[] value = serializeTopicStatus(status);
+
+        kafkaLog.send(key, value, new org.apache.kafka.clients.producer.Callback() {
+            @Override
+            public void onCompletion(RecordMetadata metadata, Exception exception) {
+                if (exception == null) return;
+                // TODO: retry more gracefully and not forever
+                if (exception instanceof RetriableException) {
+                    kafkaLog.send(key, value, this);
+                } else {
+                    log.error("Failed to write status update", exception);
+                }
+            }
+        });
+    }
+
     private  void send(final String key,
                                                  final V status,
                                                  final CacheEntry entry,
@@ -235,18 +305,17 @@ private  void send(final String key,
         kafkaLog.send(key, value, new org.apache.kafka.clients.producer.Callback() {
             @Override
             public void onCompletion(RecordMetadata metadata, Exception exception) {
-                if (exception != null) {
-                    if (exception instanceof RetriableException) {
-                        synchronized (KafkaStatusBackingStore.this) {
-                            if (entry.isDeleted()
-                                    || status.generation() != generation
-                                    || (safeWrite && !entry.canWriteSafely(status, sequence)))
-                                return;
-                        }
-                        kafkaLog.send(key, value, this);
-                    } else {
-                        log.error("Failed to write status update", exception);
+                if (exception == null) return;
+                if (exception instanceof RetriableException) {
+                    synchronized (KafkaStatusBackingStore.this) {
+                        if (entry.isDeleted()
+                            || status.generation() != generation
+                            || (safeWrite && !entry.canWriteSafely(status, sequence)))
+                            return;
                     }
+                    kafkaLog.send(key, value, this);
+                } else {
+                    log.error("Failed to write status update", exception);
                 }
             }
         });
@@ -288,6 +357,14 @@ private synchronized void remove(ConnectorTaskId id) {
             removed.delete();
     }
 
+    private void removeTopic(String topic, String connector) {
+        ConcurrentMap activeTopics = topics.get(connector);
+        if (activeTopics == null) {
+            return;
+        }
+        activeTopics.remove(topic);
+    }
+
     @Override
     public synchronized TaskStatus get(ConnectorTaskId id) {
         CacheEntry entry = tasks.get(id.connector(), id.task());
@@ -311,6 +388,25 @@ public synchronized Collection getAll(String connector) {
         return res;
     }
 
+    @Override
+    public TopicStatus getTopic(String connector, String topic) {
+        ConcurrentMap activeTopics = topics.get(Objects.requireNonNull(connector));
+        return activeTopics != null ? activeTopics.get(Objects.requireNonNull(topic)) : null;
+    }
+
+    @Override
+    public Collection getAllTopics(String connector) {
+        ConcurrentMap activeTopics = topics.get(Objects.requireNonNull(connector));
+        return activeTopics != null
+               ? Collections.unmodifiableCollection(Objects.requireNonNull(activeTopics.values()))
+               : Collections.emptySet();
+    }
+
+    @Override
+    public void deleteTopic(String connector, String topic) {
+        sendTopicStatus(Objects.requireNonNull(connector), Objects.requireNonNull(topic), null);
+    }
+
     @Override
     public synchronized Set connectors() {
         return new HashSet<>(connectors.keySet());
@@ -318,7 +414,7 @@ public synchronized Set connectors() {
 
     private ConnectorStatus parseConnectorStatus(String connector, byte[] data) {
         try {
-            SchemaAndValue schemaAndValue = converter.toConnectData(topic, data);
+            SchemaAndValue schemaAndValue = converter.toConnectData(statusTopic, data);
             if (!(schemaAndValue.value() instanceof Map)) {
                 log.error("Invalid connector status type {}", schemaAndValue.value().getClass());
                 return null;
@@ -339,9 +435,9 @@ private ConnectorStatus parseConnectorStatus(String connector, byte[] data) {
 
     private TaskStatus parseTaskStatus(ConnectorTaskId taskId, byte[] data) {
         try {
-            SchemaAndValue schemaAndValue = converter.toConnectData(topic, data);
+            SchemaAndValue schemaAndValue = converter.toConnectData(statusTopic, data);
             if (!(schemaAndValue.value() instanceof Map)) {
-                log.error("Invalid connector status type {}", schemaAndValue.value().getClass());
+                log.error("Invalid task status type {}", schemaAndValue.value().getClass());
                 return null;
             }
             @SuppressWarnings("unchecked")
@@ -357,6 +453,31 @@ private TaskStatus parseTaskStatus(ConnectorTaskId taskId, byte[] data) {
         }
     }
 
+    protected TopicStatus parseTopicStatus(byte[] data) {
+        try {
+            SchemaAndValue schemaAndValue = converter.toConnectData(statusTopic, data);
+            if (!(schemaAndValue.value() instanceof Map)) {
+                log.error("Invalid topic status value {}", schemaAndValue.value());
+                return null;
+            }
+            @SuppressWarnings("unchecked")
+            Object innerValue = ((Map) schemaAndValue.value()).get(TOPIC_STATE_KEY);
+            if (!(innerValue instanceof Map)) {
+                log.error("Invalid topic status value {} for field {}", innerValue, TOPIC_STATE_KEY);
+                return null;
+            }
+            @SuppressWarnings("unchecked")
+            Map topicStatusMetadata = (Map) innerValue;
+            return new TopicStatus((String) topicStatusMetadata.get(TOPIC_NAME_KEY),
+                    (String) topicStatusMetadata.get(TOPIC_CONNECTOR_KEY),
+                    ((Long) topicStatusMetadata.get(TOPIC_TASK_KEY)).intValue(),
+                    (long) topicStatusMetadata.get(TOPIC_DISCOVER_TIMESTAMP_KEY));
+        } catch (Exception e) {
+            log.error("Failed to deserialize topic status", e);
+            return null;
+        }
+    }
+
     private byte[] serialize(AbstractStatus status) {
         Struct struct = new Struct(STATUS_SCHEMA_V0);
         struct.put(STATE_KEY_NAME, status.state().name());
@@ -364,7 +485,24 @@ private byte[] serialize(AbstractStatus status) {
             struct.put(TRACE_KEY_NAME, status.trace());
         struct.put(WORKER_ID_KEY_NAME, status.workerId());
         struct.put(GENERATION_KEY_NAME, status.generation());
-        return converter.fromConnectData(topic, STATUS_SCHEMA_V0, struct);
+        return converter.fromConnectData(statusTopic, STATUS_SCHEMA_V0, struct);
+    }
+
+    //visible for testing
+    protected byte[] serializeTopicStatus(TopicStatus status) {
+        if (status == null) {
+            // This should send a tombstone record that will represent delete
+            return null;
+        }
+        Struct struct = new Struct(TOPIC_STATUS_VALUE_SCHEMA_V0);
+        struct.put(TOPIC_NAME_KEY, status.topic());
+        struct.put(TOPIC_CONNECTOR_KEY, status.connector());
+        struct.put(TOPIC_TASK_KEY, status.task());
+        struct.put(TOPIC_DISCOVER_TIMESTAMP_KEY, status.discoverTimestamp());
+        return converter.fromConnectData(
+                statusTopic,
+                TOPIC_STATUS_SCHEMA_V0,
+                Collections.singletonMap(TOPIC_STATE_KEY, struct));
     }
 
     private String parseConnectorStatusKey(String key) {
@@ -435,6 +573,50 @@ private void readTaskStatus(String key, byte[] value) {
         }
     }
 
+    private void readTopicStatus(String key, byte[] value) {
+        int delimiterPos = key.indexOf(':');
+        int beginPos = TOPIC_STATUS_PREFIX.length();
+        if (beginPos > delimiterPos) {
+            log.warn("Discarding record with invalid topic status key {}", key);
+            return;
+        }
+
+        String topic = key.substring(beginPos, delimiterPos);
+        if (topic.isEmpty()) {
+            log.warn("Discarding record with invalid topic status key containing empty topic {}", key);
+            return;
+        }
+
+        beginPos = delimiterPos + TOPIC_STATUS_SEPARATOR.length();
+        int endPos = key.length();
+        if (beginPos > endPos) {
+            log.warn("Discarding record with invalid topic status key {}", key);
+            return;
+        }
+
+        String connector = key.substring(beginPos);
+        if (connector.isEmpty()) {
+            log.warn("Discarding record with invalid topic status key containing empty connector {}", key);
+            return;
+        }
+
+        if (value == null) {
+            log.trace("Removing status for topic {} and connector {}", topic, connector);
+            removeTopic(topic, connector);
+            return;
+        }
+
+        TopicStatus status = parseTopicStatus(value);
+        if (status == null) {
+            log.warn("Failed to parse topic status with key {}", key);
+            return;
+        }
+
+        log.trace("Received topic status update {}", status);
+        topics.computeIfAbsent(connector, k -> new ConcurrentHashMap<>())
+                .put(topic, status);
+    }
+
     // visible for testing
     void read(ConsumerRecord record) {
         String key = record.key();
@@ -442,6 +624,8 @@ void read(ConsumerRecord record) {
             readConnectorStatus(key, record.value());
         } else if (key.startsWith(TASK_STATUS_PREFIX)) {
             readTaskStatus(key, record.value());
+        } else if (key.startsWith(TOPIC_STATUS_PREFIX)) {
+            readTopicStatus(key, record.value());
         } else {
             log.warn("Discarding record with invalid key {}", key);
         }
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 25891f52354f6..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,7 +16,9 @@
  */
 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;
 import org.apache.kafka.connect.util.ConnectorTaskId;
 
@@ -32,6 +34,14 @@ public class MemoryConfigBackingStore implements ConfigBackingStore {
 
     private Map connectors = new HashMap<>();
     private UpdateListener updateListener;
+    private WorkerConfigTransformer configTransformer;
+
+    public MemoryConfigBackingStore() {
+    }
+
+    public MemoryConfigBackingStore(WorkerConfigTransformer configTransformer) {
+        this.configTransformer = configTransformer;
+    }
 
     @Override
     public synchronized void start() {
@@ -59,11 +69,13 @@ public synchronized ClusterConfigState snapshot() {
 
         return new ClusterConfigState(
                 ClusterConfigState.NO_OFFSET,
+                null,
                 connectorTaskCounts,
                 connectorConfigs,
                 connectorTargetStates,
                 taskConfigs,
-                Collections.emptySet());
+                Collections.emptySet(),
+                configTransformer);
     }
 
     @Override
@@ -133,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..ccec12a864b2b 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
@@ -19,6 +19,7 @@
 import org.apache.kafka.connect.errors.ConnectException;
 import org.apache.kafka.connect.runtime.WorkerConfig;
 import org.apache.kafka.connect.util.Callback;
+import org.apache.kafka.common.utils.ThreadUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -53,7 +54,8 @@ public void configure(WorkerConfig config) {
 
     @Override
     public void start() {
-        executor = Executors.newSingleThreadExecutor();
+        executor = Executors.newFixedThreadPool(1, ThreadUtils.createThreadFactory(
+                this.getClass().getSimpleName() + "-%d", false));
     }
 
     @Override
@@ -75,9 +77,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 +85,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/MemoryStatusBackingStore.java b/connect/runtime/src/main/java/org/apache/kafka/connect/storage/MemoryStatusBackingStore.java
index 39d098dd1c59a..fbd7048ed26f3 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/storage/MemoryStatusBackingStore.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/storage/MemoryStatusBackingStore.java
@@ -18,23 +18,30 @@
 
 import org.apache.kafka.connect.runtime.ConnectorStatus;
 import org.apache.kafka.connect.runtime.TaskStatus;
+import org.apache.kafka.connect.runtime.TopicStatus;
 import org.apache.kafka.connect.runtime.WorkerConfig;
 import org.apache.kafka.connect.util.ConnectorTaskId;
 import org.apache.kafka.connect.util.Table;
 
 import java.util.Collection;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
 
 public class MemoryStatusBackingStore implements StatusBackingStore {
     private final Table tasks;
     private final Map connectors;
+    private final ConcurrentMap> topics;
 
     public MemoryStatusBackingStore() {
         this.tasks = new Table<>();
         this.connectors = new HashMap<>();
+        this.topics = new ConcurrentHashMap<>();
     }
 
     @Override
@@ -78,6 +85,12 @@ public synchronized void putSafe(TaskStatus status) {
         put(status);
     }
 
+    @Override
+    public void put(final TopicStatus status) {
+        topics.computeIfAbsent(status.connector(), k -> new ConcurrentHashMap<>())
+                .put(status.topic(), status);
+    }
+
     @Override
     public synchronized TaskStatus get(ConnectorTaskId id) {
         return tasks.get(id.connector(), id.task());
@@ -93,6 +106,28 @@ public synchronized Collection getAll(String connector) {
         return new HashSet<>(tasks.row(connector).values());
     }
 
+    @Override
+    public TopicStatus getTopic(String connector, String topic) {
+        ConcurrentMap activeTopics = topics.get(Objects.requireNonNull(connector));
+        return activeTopics != null ? activeTopics.get(Objects.requireNonNull(topic)) : null;
+    }
+
+    @Override
+    public Collection getAllTopics(String connector) {
+        ConcurrentMap activeTopics = topics.get(Objects.requireNonNull(connector));
+        return activeTopics != null
+               ? Collections.unmodifiableCollection(activeTopics.values())
+               : Collections.emptySet();
+    }
+
+    @Override
+    public void deleteTopic(String connector, String topic) {
+        ConcurrentMap activeTopics = topics.get(Objects.requireNonNull(connector));
+        if (activeTopics != null) {
+            activeTopics.remove(Objects.requireNonNull(topic));
+        }
+    }
+
     @Override
     public synchronized Set connectors() {
         return new HashSet<>(connectors.keySet());
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 53cc1a3a6f704..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
@@ -24,21 +24,29 @@
 import java.nio.ByteBuffer;
 import java.util.Arrays;
 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) {
@@ -46,11 +54,13 @@ 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
     public  Map offset(Map partition) {
-        return offsets(Arrays.asList(partition)).get(partition);
+        return offsets(Collections.singletonList(partition)).get(partition);
     }
 
     @Override
@@ -75,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);
@@ -107,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/storage/OffsetStorageWriter.java b/connect/runtime/src/main/java/org/apache/kafka/connect/storage/OffsetStorageWriter.java
index 3239b6734cd9b..c360d4373b083 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/storage/OffsetStorageWriter.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/storage/OffsetStorageWriter.java
@@ -89,6 +89,7 @@ public OffsetStorageWriter(OffsetBackingStore backingStore,
      * @param partition the partition to store an offset for
      * @param offset the offset
      */
+    @SuppressWarnings("unchecked")
     public synchronized void offset(Map partition, Map offset) {
         data.put((Map) partition, (Map) offset);
     }
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/storage/StatusBackingStore.java b/connect/runtime/src/main/java/org/apache/kafka/connect/storage/StatusBackingStore.java
index 63294a47525c5..0250932df14af 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/storage/StatusBackingStore.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/storage/StatusBackingStore.java
@@ -18,6 +18,7 @@
 
 import org.apache.kafka.connect.runtime.ConnectorStatus;
 import org.apache.kafka.connect.runtime.TaskStatus;
+import org.apache.kafka.connect.runtime.TopicStatus;
 import org.apache.kafka.connect.runtime.WorkerConfig;
 import org.apache.kafka.connect.util.ConnectorTaskId;
 
@@ -66,6 +67,12 @@ public interface StatusBackingStore {
      */
     void putSafe(TaskStatus status);
 
+    /**
+     * Set the state of a connector's topic to the given value.
+     * @param status the status of the topic used by a connector
+     */
+    void put(TopicStatus status);
+
     /**
      * Get the current state of the task.
      * @param id the id of the task
@@ -87,6 +94,28 @@ public interface StatusBackingStore {
      */
     Collection getAll(String connector);
 
+    /**
+     * Get the status of a connector's topic if the connector is actively using this topic
+     * @param connector the connector name; never null
+     * @param topic the topic name; never null
+     * @return the state or null if there is none
+     */
+    TopicStatus getTopic(String connector, String topic);
+
+    /**
+     * Get the states of all topics that a connector is using.
+     * @param connector the connector name; never null
+     * @return a collection of topic states or an empty collection if there is none
+     */
+    Collection getAllTopics(String connector);
+
+    /**
+     * Delete this topic from the connector's set of active topics
+     * @param connector the connector name; never null
+     * @param topic the topic name; never null
+     */
+    void deleteTopic(String connector, String topic);
+
     /**
      * Get all cached connectors.
      * @return the set of connector names
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/tools/PredicateDoc.java b/connect/runtime/src/main/java/org/apache/kafka/connect/tools/PredicateDoc.java
new file mode 100644
index 0000000000000..d4399d6cb9ac1
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/tools/PredicateDoc.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 org.apache.kafka.connect.tools;
+
+import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.connect.runtime.isolation.Plugins;
+import org.apache.kafka.connect.transforms.predicates.Predicate;
+
+import java.io.PrintStream;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.List;
+import java.util.stream.Collectors;
+
+public class PredicateDoc {
+
+    private static final class DocInfo {
+        final String predicateName;
+        final String overview;
+        final ConfigDef configDef;
+
+        private 
            > DocInfo(Class
             predicateClass, String overview, ConfigDef configDef) {
+            this.predicateName = predicateClass.getName();
+            this.overview = overview;
+            this.configDef = configDef;
+        }
+    }
+
+    private static final List PREDICATES;
+    static {
+        List collect = new Plugins(Collections.emptyMap()).predicates().stream()
+            .map(p -> {
+                try {
+                    String overviewDoc = (String) p.pluginClass().getDeclaredField("OVERVIEW_DOC").get(null);
+                    ConfigDef configDef = (ConfigDef) p.pluginClass().getDeclaredField("CONFIG_DEF").get(null);
+                    return new DocInfo(p.pluginClass(), overviewDoc, configDef);
+                } catch (ReflectiveOperationException e) {
+                    throw new RuntimeException("Predicate class " + p.pluginClass().getName() + " lacks either a `public static final String OVERVIEW_DOC` or `public static final ConfigDef CONFIG_DEF`");
+                }
+            })
+            .collect(Collectors.toList());
+        collect.sort(Comparator.comparing(docInfo -> docInfo.predicateName));
+        PREDICATES = collect;
+    }
+
+    private static void printPredicateHtml(PrintStream out, DocInfo docInfo) {
+        out.println("
            ");
+
+        out.print("
            ");
+        out.print(docInfo.predicateName);
+        out.println("
            ");
+
+        out.println(docInfo.overview);
+
+        out.println("
            ");
+
+        out.println(docInfo.configDef.toHtml(6, key -> docInfo.predicateName + "_" + key));
+
+        out.println("
            ");
+    }
+
+    private static void printHtml(PrintStream out) {
+        for (final DocInfo docInfo : PREDICATES) {
+            printPredicateHtml(out, docInfo);
+        }
+    }
+
+    public static void main(String... args) {
+        printHtml(System.out);
+    }
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/tools/SchemaSourceConnector.java b/connect/runtime/src/main/java/org/apache/kafka/connect/tools/SchemaSourceConnector.java
index eeee0f6ca8d32..06379d5f2811c 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/tools/SchemaSourceConnector.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/tools/SchemaSourceConnector.java
@@ -48,9 +48,9 @@ public Class taskClass() {
     @Override
     public List> taskConfigs(int maxTasks) {
         ArrayList> configs = new ArrayList<>();
-        for (Integer i = 0; i < maxTasks; i++) {
+        for (int i = 0; i < maxTasks; i++) {
             Map props = new HashMap<>(config);
-            props.put(SchemaSourceTask.ID_CONFIG, i.toString());
+            props.put(SchemaSourceTask.ID_CONFIG, String.valueOf(i));
             configs.add(props);
         }
         return configs;
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/tools/SchemaSourceTask.java b/connect/runtime/src/main/java/org/apache/kafka/connect/tools/SchemaSourceTask.java
index 6a51b52cead12..6fde784e7861b 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/tools/SchemaSourceTask.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/tools/SchemaSourceTask.java
@@ -26,8 +26,6 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
@@ -82,6 +80,7 @@ public class SchemaSourceTask extends SourceTask {
         .field("seqno", Schema.INT64_SCHEMA)
         .build();
 
+    @Override
     public String version() {
         return new SchemaSourceConnector().version();
     }
@@ -154,13 +153,12 @@ public List poll() throws InterruptedException {
             }
 
             System.out.println("{\"task\": " + id + ", \"seqno\": " + seqno + "}");
-            List result = Arrays.asList(srcRecord);
             seqno++;
             count++;
-            return result;
+            return Collections.singletonList(srcRecord);
         } else {
             throttler.throttle();
-            return new ArrayList<>();
+            return Collections.emptyList();
         }
     }
 
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..82c2663714c47 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
@@ -19,6 +19,7 @@
 import org.apache.kafka.common.config.ConfigDef;
 import org.apache.kafka.connect.transforms.Cast;
 import org.apache.kafka.connect.transforms.ExtractField;
+import org.apache.kafka.connect.transforms.Filter;
 import org.apache.kafka.connect.transforms.Flatten;
 import org.apache.kafka.connect.transforms.HoistField;
 import org.apache.kafka.connect.transforms.InsertField;
@@ -60,7 +61,8 @@ private DocInfo(String transformationName, String overview, ConfigDef configDef)
             new DocInfo(RegexRouter.class.getName(), RegexRouter.OVERVIEW_DOC, RegexRouter.CONFIG_DEF),
             new DocInfo(Flatten.class.getName(), Flatten.OVERVIEW_DOC, Flatten.CONFIG_DEF),
             new DocInfo(Cast.class.getName(), Cast.OVERVIEW_DOC, Cast.CONFIG_DEF),
-            new DocInfo(TimestampConverter.class.getName(), TimestampConverter.OVERVIEW_DOC, TimestampConverter.CONFIG_DEF)
+            new DocInfo(TimestampConverter.class.getName(), TimestampConverter.OVERVIEW_DOC, TimestampConverter.CONFIG_DEF),
+            new DocInfo(Filter.class.getName(), Filter.OVERVIEW_DOC, Filter.CONFIG_DEF)
     );
 
     private static void printTransformationHtml(PrintStream out, DocInfo docInfo) {
@@ -74,18 +76,18 @@ private static void printTransformationHtml(PrintStream out, DocInfo docInfo) {
 
         out.println("
            ");
 
-        out.println(docInfo.configDef.toHtmlTable());
+        out.println(docInfo.configDef.toHtml(6, key -> docInfo.transformationName + "_"  + key));
 
         out.println("");
     }
 
-    private static void printHtml(PrintStream out) throws NoSuchFieldException, IllegalAccessException, InstantiationException {
+    private static void printHtml(PrintStream out) {
         for (final DocInfo docInfo : TRANSFORMATIONS) {
             printTransformationHtml(out, docInfo);
         }
     }
 
-    public static void main(String... args) throws Exception {
+    public static void main(String... args) {
         printHtml(System.out);
     }
 
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/tools/VerifiableSinkConnector.java b/connect/runtime/src/main/java/org/apache/kafka/connect/tools/VerifiableSinkConnector.java
index 12bcb353cd875..55f95a3448c67 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/tools/VerifiableSinkConnector.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/tools/VerifiableSinkConnector.java
@@ -50,9 +50,9 @@ public Class taskClass() {
     @Override
     public List> taskConfigs(int maxTasks) {
         ArrayList> configs = new ArrayList<>();
-        for (Integer i = 0; i < maxTasks; i++) {
+        for (int i = 0; i < maxTasks; i++) {
             Map props = new HashMap<>(config);
-            props.put(VerifiableSinkTask.ID_CONFIG, i.toString());
+            props.put(VerifiableSinkTask.ID_CONFIG, String.valueOf(i));
             configs.add(props);
         }
         return configs;
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/tools/VerifiableSourceConnector.java b/connect/runtime/src/main/java/org/apache/kafka/connect/tools/VerifiableSourceConnector.java
index cd7e3105457cc..6262cc3b0bb1e 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/tools/VerifiableSourceConnector.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/tools/VerifiableSourceConnector.java
@@ -50,9 +50,9 @@ public Class taskClass() {
     @Override
     public List> taskConfigs(int maxTasks) {
         ArrayList> configs = new ArrayList<>();
-        for (Integer i = 0; i < maxTasks; i++) {
+        for (int i = 0; i < maxTasks; i++) {
             Map props = new HashMap<>(config);
-            props.put(VerifiableSourceTask.ID_CONFIG, i.toString());
+            props.put(VerifiableSourceTask.ID_CONFIG, String.valueOf(i));
             configs.add(props);
         }
         return configs;
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/tools/VerifiableSourceTask.java b/connect/runtime/src/main/java/org/apache/kafka/connect/tools/VerifiableSourceTask.java
index dfd8bac5aa06d..8afbdffa167fd 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/tools/VerifiableSourceTask.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/tools/VerifiableSourceTask.java
@@ -19,6 +19,7 @@
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import org.apache.kafka.tools.ThroughputThrottler;
+import org.apache.kafka.clients.producer.RecordMetadata;
 import org.apache.kafka.connect.data.Schema;
 import org.apache.kafka.connect.errors.ConnectException;
 import org.apache.kafka.connect.source.SourceRecord;
@@ -26,7 +27,6 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
@@ -115,13 +115,13 @@ public List poll() throws InterruptedException {
 
         Map ccOffset = Collections.singletonMap(SEQNO_FIELD, seqno);
         SourceRecord srcRecord = new SourceRecord(partition, ccOffset, topic, Schema.INT32_SCHEMA, id, Schema.INT64_SCHEMA, seqno);
-        List result = Arrays.asList(srcRecord);
+        List result = Collections.singletonList(srcRecord);
         seqno++;
         return result;
     }
 
     @Override
-    public void commitRecord(SourceRecord record) throws InterruptedException {
+    public void commitRecord(SourceRecord record, RecordMetadata metadata) throws InterruptedException {
         Map data = new HashMap<>();
         data.put("name", name);
         data.put("task", id);
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 913ae1f10583b..c1b83eb648828 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,10 +16,26 @@
  */
 package org.apache.kafka.connect.util;
 
-import org.apache.kafka.common.record.InvalidRecordException;
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.common.KafkaFuture;
+import org.apache.kafka.common.InvalidRecordException;
 import org.apache.kafka.common.record.RecordBatch;
+import org.apache.kafka.connect.connector.Connector;
+import org.apache.kafka.connect.errors.ConnectException;
+import org.apache.kafka.connect.runtime.WorkerConfig;
+import org.apache.kafka.connect.runtime.distributed.DistributedConfig;
+import org.apache.kafka.connect.sink.SinkConnector;
+import org.apache.kafka.connect.source.SourceConnector;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.Map;
+import java.util.concurrent.ExecutionException;
 
 public final class ConnectUtils {
+    private static final Logger log = LoggerFactory.getLogger(ConnectUtils.class);
+
     public static Long checkAndConvertTimestamp(Long timestamp) {
         if (timestamp == null || timestamp >= 0)
             return timestamp;
@@ -28,4 +44,50 @@ else if (timestamp == RecordBatch.NO_TIMESTAMP)
         else
             throw new InvalidRecordException(String.format("Invalid record timestamp %d", timestamp));
     }
+
+    public static String lookupKafkaClusterId(WorkerConfig config) {
+        log.info("Creating Kafka admin client");
+        try (Admin adminClient = Admin.create(config.originals())) {
+            return lookupKafkaClusterId(adminClient);
+        }
+    }
+
+    static String lookupKafkaClusterId(Admin adminClient) {
+        log.debug("Looking up Kafka cluster ID");
+        try {
+            KafkaFuture clusterIdFuture = adminClient.describeCluster().clusterId();
+            if (clusterIdFuture == null) {
+                log.info("Kafka cluster version is too old to return cluster ID");
+                return null;
+            }
+            log.debug("Fetching Kafka cluster ID");
+            String kafkaClusterId = clusterIdFuture.get();
+            log.info("Kafka cluster ID: {}", kafkaClusterId);
+            return kafkaClusterId;
+        } catch (InterruptedException e) {
+            throw new ConnectException("Unexpectedly interrupted when looking up Kafka cluster info", e);
+        } catch (ExecutionException e) {
+            throw new ConnectException("Failed to connect to and describe Kafka cluster. "
+                                       + "Check worker's broker connection and security properties.", e);
+        }
+    }
+
+    public static void addMetricsContextProperties(Map prop, WorkerConfig config, String clusterId) {
+        //add all properties predefined with "metrics.context."
+        prop.putAll(config.originalsWithPrefix(CommonClientConfigs.METRICS_CONTEXT_PREFIX, false));
+        //add connect properties
+        prop.put(CommonClientConfigs.METRICS_CONTEXT_PREFIX + WorkerConfig.CONNECT_KAFKA_CLUSTER_ID, clusterId);
+        Object groupId = config.originals().get(DistributedConfig.GROUP_ID_CONFIG);
+        if (groupId != null) {
+            prop.put(CommonClientConfigs.METRICS_CONTEXT_PREFIX + WorkerConfig.CONNECT_GROUP_ID, groupId);
+        }
+    }
+
+    public static boolean isSinkConnector(Connector connector) {
+        return SinkConnector.class.isAssignableFrom(connector.getClass());
+    }
+
+    public static boolean isSourceConnector(Connector connector) {
+        return SourceConnector.class.isAssignableFrom(connector.getClass());
+    }
 }
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 de1ceb3be1006..5248715aa6293 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,13 +28,14 @@
 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;
 import org.apache.kafka.connect.errors.ConnectException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.time.Duration;
 import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Iterator;
@@ -43,6 +44,7 @@
 import java.util.Queue;
 import java.util.Set;
 import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
 
 
 /**
@@ -68,10 +70,12 @@
  */
 public class KafkaBasedLog {
     private static final Logger log = LoggerFactory.getLogger(KafkaBasedLog.class);
-    private static final long CREATE_TOPIC_TIMEOUT_MS = 30000;
+    private static final long CREATE_TOPIC_TIMEOUT_NS = TimeUnit.SECONDS.toNanos(30);
+    private static final long MAX_SLEEP_MS = TimeUnit.SECONDS.toMillis(1);
 
     private Time time;
     private final String topic;
+    private int partitionCount;
     private final Map producerConfigs;
     private final Map consumerConfigs;
     private final Callback> consumedCallback;
@@ -130,11 +134,13 @@ public void start() {
         List partitions = new ArrayList<>();
 
         // We expect that the topics will have been created either manually by the user or automatically by the herder
-        List partitionInfos = null;
-        long started = time.milliseconds();
-        while (partitionInfos == null && time.milliseconds() - started < CREATE_TOPIC_TIMEOUT_MS) {
+        List partitionInfos = consumer.partitionsFor(topic);
+        long started = time.nanoseconds();
+        long sleepMs = 100;
+        while (partitionInfos == null && time.nanoseconds() - started < CREATE_TOPIC_TIMEOUT_NS) {
+            time.sleep(sleepMs);
+            sleepMs = Math.min(2 * sleepMs, MAX_SLEEP_MS);
             partitionInfos = consumer.partitionsFor(topic);
-            Utils.sleep(Math.min(time.milliseconds() - started, 1000));
         }
         if (partitionInfos == null)
             throw new ConnectException("Could not look up partition metadata for offset backing store topic in" +
@@ -143,8 +149,13 @@ 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
+        // when a 'group.id' is specified (if offsets happen to have been committed unexpectedly).
+        consumer.seekToBeginning(partitions);
+
         readToLogEnd();
 
         thread = new WorkThread();
@@ -232,6 +243,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
@@ -253,7 +267,7 @@ private Consumer createConsumer() {
 
     private void poll(long timeoutMs) {
         try {
-            ConsumerRecords records = consumer.poll(timeoutMs);
+            ConsumerRecords records = consumer.poll(Duration.ofMillis(timeoutMs));
             for (ConsumerRecord record : records)
                 consumedCallback.onCompletion(null, record);
         } catch (WakeupException e) {
@@ -275,9 +289,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;
                 }
@@ -307,6 +327,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.
@@ -335,4 +359,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/ReflectionsUtil.java b/connect/runtime/src/main/java/org/apache/kafka/connect/util/ReflectionsUtil.java
deleted file mode 100644
index 61d0e35c6288c..0000000000000
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/util/ReflectionsUtil.java
+++ /dev/null
@@ -1,92 +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.util;
-
-import org.reflections.vfs.Vfs;
-import org.reflections.vfs.Vfs.Dir;
-import org.reflections.vfs.Vfs.File;
-import org.reflections.vfs.Vfs.UrlType;
-
-import java.net.URL;
-import java.util.Arrays;
-import java.util.Collections;
-import java.util.LinkedList;
-import java.util.List;
-
-/**
- * CLASSPATH on OSX contains .mar and .jnilib file extensions. Vfs used by Reflections does not recognize
- * urls with those extensions and log WARNs when scan them. Those WARNs can be eliminated by registering
- * URL types before using reflection.
- */
-public class ReflectionsUtil {
-
-    private static final String FILE_PROTOCOL = "file";
-    private static final List ENDINGS = Arrays.asList(".mar", ".jnilib", "*");
-
-    public static void registerUrlTypes() {
-        final List urlTypes = new LinkedList<>();
-        urlTypes.add(new EmptyUrlType(ENDINGS));
-        urlTypes.addAll(Arrays.asList(Vfs.DefaultUrlTypes.values()));
-        Vfs.setDefaultURLTypes(urlTypes);
-    }
-
-    private static class EmptyUrlType implements UrlType {
-
-        private final List endings;
-
-        private EmptyUrlType(final List endings) {
-            this.endings = endings;
-        }
-
-        public boolean matches(URL url) {
-            final String protocol = url.getProtocol();
-            final String externalForm = url.toExternalForm();
-            if (!protocol.equals(FILE_PROTOCOL)) {
-                return false;
-            }
-            for (String ending : endings) {
-                if (externalForm.endsWith(ending)) {
-                    return true;
-                }
-            }
-            return false;
-        }
-
-        public Dir createDir(final URL url) throws Exception {
-            return emptyVfsDir(url);
-        }
-
-        private static Dir emptyVfsDir(final URL url) {
-            return new Dir() {
-                @Override
-                public String getPath() {
-                    return url.toExternalForm();
-                }
-
-                @Override
-                public Iterable getFiles() {
-                    return Collections.emptyList();
-                }
-
-                @Override
-                public void close() {
-
-                }
-            };
-        }
-    }
-}
\ No newline at end of file
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/util/SafeObjectInputStream.java b/connect/runtime/src/main/java/org/apache/kafka/connect/util/SafeObjectInputStream.java
index b5c4c2559a4dc..0ad3889b5f0ea 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/util/SafeObjectInputStream.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/util/SafeObjectInputStream.java
@@ -52,14 +52,14 @@ public SafeObjectInputStream(InputStream in) throws IOException {
     protected Class resolveClass(ObjectStreamClass desc) throws IOException, ClassNotFoundException {
         String name = desc.getName();
 
-        if (isBlacklisted(name)) {
+        if (isBlocked(name)) {
             throw new SecurityException("Illegal type to deserialize: prevented for security reasons");
         }
 
         return super.resolveClass(desc);
     }
 
-    private boolean isBlacklisted(String name) {
+    private boolean isBlocked(String name) {
         for (String list : DEFAULT_NO_DESERIALIZE_CLASS_NAMES) {
             if (name.endsWith(list)) {
                 return true;
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/util/Table.java b/connect/runtime/src/main/java/org/apache/kafka/connect/util/Table.java
index bf7ca15761236..1b7131a51445d 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/util/Table.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/util/Table.java
@@ -62,4 +62,7 @@ public Map row(R row) {
         return Collections.unmodifiableMap(columns);
     }
 
+    public boolean isEmpty() {
+        return table.isEmpty();
+    }
 }
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 5da4f2d00d0db..615e7a35d750e 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,46 +16,66 @@
  */
 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.Config;
+import org.apache.kafka.clients.admin.ConfigEntry;
 import org.apache.kafka.clients.admin.CreateTopicsOptions;
+import org.apache.kafka.clients.admin.DescribeConfigsOptions;
+import org.apache.kafka.clients.admin.DescribeTopicsOptions;
 import org.apache.kafka.clients.admin.NewTopic;
+import org.apache.kafka.clients.admin.TopicDescription;
 import org.apache.kafka.common.KafkaFuture;
+import org.apache.kafka.common.config.ConfigException;
+import org.apache.kafka.common.config.ConfigResource;
+import org.apache.kafka.common.config.TopicConfig;
+import org.apache.kafka.common.errors.ClusterAuthorizationException;
+import org.apache.kafka.common.errors.InvalidConfigurationException;
 import org.apache.kafka.common.errors.TimeoutException;
+import org.apache.kafka.common.errors.TopicAuthorizationException;
 import org.apache.kafka.common.errors.TopicExistsException;
+import org.apache.kafka.common.errors.UnknownTopicOrPartitionException;
 import org.apache.kafka.common.errors.UnsupportedVersionException;
 import org.apache.kafka.common.utils.Utils;
 import org.apache.kafka.connect.errors.ConnectException;
+import org.apache.kafka.connect.errors.RetriableException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.time.Duration;
+import java.util.Arrays;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
+import java.util.stream.Collectors;
 
 /**
- * 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 {
 
-    private static final String CLEANUP_POLICY_CONFIG = "cleanup.policy";
-    private static final String CLEANUP_POLICY_COMPACT = "compact";
+    public static final int NO_PARTITIONS = -1;
+    public static final short NO_REPLICATION_FACTOR = -1;
 
-    private static final String MIN_INSYNC_REPLICAS_CONFIG = "min.insync.replicas";
-
-    private static final String UNCLEAN_LEADER_ELECTION_ENABLE_CONFIG = "unclean.leader.election.enable";
+    private static final String CLEANUP_POLICY_CONFIG = TopicConfig.CLEANUP_POLICY_CONFIG;
+    private static final String CLEANUP_POLICY_COMPACT = TopicConfig.CLEANUP_POLICY_COMPACT;
+    private static final String MIN_INSYNC_REPLICAS_CONFIG = TopicConfig.MIN_IN_SYNC_REPLICAS_CONFIG;
+    private static final String UNCLEAN_LEADER_ELECTION_ENABLE_CONFIG = TopicConfig.UNCLEAN_LEADER_ELECTION_ENABLE_CONFIG;
 
     /**
      * A builder of {@link NewTopic} instances.
      */
     public static class NewTopicBuilder {
-        private String name;
-        private int numPartitions;
-        private short replicationFactor;
-        private Map configs = new HashMap<>();
+        private final String name;
+        private int numPartitions = NO_PARTITIONS;
+        private short replicationFactor = NO_REPLICATION_FACTOR;
+        private final Map configs = new HashMap<>();
 
         NewTopicBuilder(String name) {
             this.name = name;
@@ -64,7 +84,8 @@ public static class NewTopicBuilder {
         /**
          * Specify the desired number of partitions for the topic.
          *
-         * @param numPartitions the desired number of partitions; must be positive
+         * @param numPartitions the desired number of partitions; must be positive, or -1 to
+         *                      signify using the broker's default
          * @return this builder to allow methods to be chained; never null
          */
         public NewTopicBuilder partitions(int numPartitions) {
@@ -72,10 +93,22 @@ public NewTopicBuilder partitions(int numPartitions) {
             return this;
         }
 
+        /**
+         * Specify the topic's number of partition should be the broker configuration for
+         * {@code num.partitions}.
+         *
+         * @return this builder to allow methods to be chained; never null
+         */
+        public NewTopicBuilder defaultPartitions() {
+            this.numPartitions = NO_PARTITIONS;
+            return this;
+        }
+
         /**
          * Specify the desired replication factor for the topic.
          *
-         * @param replicationFactor the desired replication factor; must be positive
+         * @param replicationFactor the desired replication factor; must be positive, or -1 to
+         *                          signify using the broker's default
          * @return this builder to allow methods to be chained; never null
          */
         public NewTopicBuilder replicationFactor(short replicationFactor) {
@@ -83,6 +116,17 @@ public NewTopicBuilder replicationFactor(short replicationFactor) {
             return this;
         }
 
+        /**
+         * Specify the replication factor for the topic should be the broker configurations for
+         * {@code default.replication.factor}.
+         *
+         * @return this builder to allow methods to be chained; never null
+         */
+        public NewTopicBuilder defaultReplicationFactor() {
+            this.replicationFactor = NO_REPLICATION_FACTOR;
+            return this;
+        }
+
         /**
          * Specify that the topic should be compacted.
          *
@@ -140,7 +184,11 @@ public NewTopicBuilder config(Map configs) {
          * @return the topic description; never null
          */
         public NewTopic build() {
-            return new NewTopic(name, numPartitions, replicationFactor).configs(configs);
+            return new NewTopic(
+                    name,
+                    Optional.of(numPartitions),
+                    Optional.of(replicationFactor)
+            ).configs(configs);
         }
     }
 
@@ -156,21 +204,28 @@ 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;
+    private final boolean logCreation;
 
     /**
      * 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(adminConfig, adminClient, true);
+    }
+
+    // visible for testing
+    TopicAdmin(Map adminConfig, Admin adminClient, boolean logCreation) {
         this.admin = adminClient;
         this.adminConfig = adminConfig != null ? adminConfig : Collections.emptyMap();
+        this.logCreation = logCreation;
     }
 
    /**
@@ -225,19 +280,38 @@ public Set createTopics(NewTopic... topics) {
             String topic = entry.getKey();
             try {
                 entry.getValue().get();
-                log.info("Created topic {} on brokers at {}", topicsByName.get(topic), bootstrapServers);
+                if (logCreation) {
+                    log.info("Created topic {} on brokers at {}", topicsByName.get(topic), bootstrapServers);
+                }
                 newlyCreatedTopicNames.add(topic);
             } catch (ExecutionException e) {
                 Throwable cause = e.getCause();
-                if (e.getCause() instanceof TopicExistsException) {
+                if (cause instanceof TopicExistsException) {
                     log.debug("Found existing topic '{}' on the brokers at {}", topic, bootstrapServers);
                     continue;
                 }
                 if (cause instanceof UnsupportedVersionException) {
-                    log.debug("Unable to use Kafka admin client to create topic descriptions for '{}' using the brokers at {}," +
-                                      "falling back to assume topic(s) exist or will be auto-created by the broker", topicNameList, bootstrapServers);
+                    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) '{}' 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 InvalidConfigurationException) {
+                    throw new ConnectException("Unable to create topic(s) '" + topicNameList + "': " + cause.getMessage(),
+                            cause);
+                }
                 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 + "'." +
@@ -253,11 +327,227 @@ public Set createTopics(NewTopic... topics) {
         return newlyCreatedTopicNames;
     }
 
+    /**
+     * Attempt to fetch the descriptions of the given topics
+     * Apache Kafka added support for describing topics in 0.10.0.0, so this method works as expected with that and later versions.
+     * With brokers older than 0.10.0.0, this method is unable to describe topics and always returns an empty set.
+     *
+     * @param topics the topics to describe
+     * @return a map of topic names to topic descriptions of the topics that were requested; never null but possibly empty
+     * @throws RetriableException if a retriable error occurs, the operation takes too long, or the
+     * thread is interrupted while attempting to perform this operation
+     * @throws ConnectException if a non retriable error occurs
+     */
+    public Map describeTopics(String... topics) {
+        if (topics == null) {
+            return Collections.emptyMap();
+        }
+        String bootstrapServers = bootstrapServers();
+        String topicNameList = String.join(", ", topics);
+
+        Map> newResults =
+                admin.describeTopics(Arrays.asList(topics), new DescribeTopicsOptions()).values();
+
+        // Iterate over each future so that we can handle individual failures like when some topics don't exist
+        Map existingTopics = new HashMap<>();
+        newResults.forEach((topic, desc) -> {
+            try {
+                existingTopics.put(topic, desc.get());
+            } catch (ExecutionException e) {
+                Throwable cause = e.getCause();
+                if (cause instanceof UnknownTopicOrPartitionException) {
+                    log.debug("Topic '{}' does not exist on the brokers at {}", topic, bootstrapServers);
+                    return;
+                }
+                if (cause instanceof ClusterAuthorizationException || cause instanceof TopicAuthorizationException) {
+                    String msg = String.format("Not authorized to describe topic(s) '%s' on the brokers %s",
+                            topicNameList, bootstrapServers);
+                    throw new ConnectException(msg, cause);
+                }
+                if (cause instanceof UnsupportedVersionException) {
+                    String msg = String.format("Unable to describe topic(s) '%s' since the brokers "
+                                    + "at %s do not support the DescribeTopics API.",
+                            topicNameList, bootstrapServers);
+                    throw new ConnectException(msg, cause);
+                }
+                if (cause instanceof TimeoutException) {
+                    // Timed out waiting for the operation to complete
+                    throw new RetriableException("Timed out while describing topics '" + topicNameList + "'", cause);
+                }
+                throw new ConnectException("Error while attempting to describe topics '" + topicNameList + "'", e);
+            } catch (InterruptedException e) {
+                Thread.interrupted();
+                throw new RetriableException("Interrupted while attempting to describe topics '" + topicNameList + "'", e);
+            }
+        });
+        return existingTopics;
+    }
+
+    /**
+     * Verify the named topic uses only compaction for the cleanup policy.
+     *
+     * @param topic             the name of the topic
+     * @param workerTopicConfig the name of the worker configuration that specifies the topic name
+     * @return true if the admin client could be used to verify the topic setting, or false if
+     *         the verification could not be performed, likely because the admin client principal
+     *         did not have the required permissions or because the broker was older than 0.11.0.0
+     * @throws ConfigException if the actual topic setting did not match the required setting
+     */
+    public boolean verifyTopicCleanupPolicyOnlyCompact(String topic, String workerTopicConfig,
+            String topicPurpose) {
+        Set cleanupPolicies = topicCleanupPolicy(topic);
+        if (cleanupPolicies.isEmpty()) {
+            log.info("Unable to use admin client to verify the cleanup policy of '{}' "
+                      + "topic is '{}', either because the broker is an older "
+                      + "version or because the Kafka principal used for Connect "
+                      + "internal topics does not have the required permission to "
+                      + "describe topic configurations.", topic, TopicConfig.CLEANUP_POLICY_COMPACT);
+            return false;
+        }
+        Set expectedPolicies = Collections.singleton(TopicConfig.CLEANUP_POLICY_COMPACT);
+        if (!cleanupPolicies.equals(expectedPolicies)) {
+            String expectedPolicyStr = String.join(",", expectedPolicies);
+            String cleanupPolicyStr = String.join(",", cleanupPolicies);
+            String msg = String.format("Topic '%s' supplied via the '%s' property is required "
+                    + "to have '%s=%s' to guarantee consistency and durability of "
+                    + "%s, but found the topic currently has '%s=%s'. Continuing would likely "
+                    + "result in eventually losing %s and problems restarting this Connect "
+                    + "cluster in the future. Change the '%s' property in the "
+                    + "Connect worker configurations to use a topic with '%s=%s'.",
+                    topic, workerTopicConfig, TopicConfig.CLEANUP_POLICY_CONFIG, expectedPolicyStr,
+                    topicPurpose, TopicConfig.CLEANUP_POLICY_CONFIG, cleanupPolicyStr, topicPurpose,
+                    workerTopicConfig, TopicConfig.CLEANUP_POLICY_CONFIG, expectedPolicyStr);
+            throw new ConfigException(msg);
+        }
+        return true;
+    }
+
+    /**
+     * Get the cleanup policy for a topic.
+     *
+     * @param topic the name of the topic
+     * @return the set of cleanup policies set for the topic; may be empty if the topic does not
+     *         exist or the topic's cleanup policy could not be retrieved
+     */
+    public Set topicCleanupPolicy(String topic) {
+        Config topicConfig = describeTopicConfig(topic);
+        if (topicConfig == null) {
+            // The topic must not exist
+            log.debug("Unable to find topic '{}' when getting cleanup policy", topic);
+            return Collections.emptySet();
+        }
+        ConfigEntry entry = topicConfig.get(CLEANUP_POLICY_CONFIG);
+        if (entry != null && entry.value() != null) {
+            String policyStr = entry.value();
+            log.debug("Found cleanup.policy={} for topic '{}'", policyStr, topic);
+            return Arrays.stream(policyStr.split(","))
+                         .map(String::trim)
+                         .filter(s -> !s.isEmpty())
+                         .map(String::toLowerCase)
+                         .collect(Collectors.toSet());
+        }
+        // This is unexpected, as the topic config should include the cleanup.policy even if
+        // the topic settings don't override the broker's log.cleanup.policy. But just to be safe.
+        log.debug("Found no cleanup.policy for topic '{}'", topic);
+        return Collections.emptySet();
+    }
+
+    /**
+     * Attempt to fetch the topic configuration for the given topic.
+     * Apache Kafka added support for describing topic configurations in 0.11.0.0, so this method
+     * works as expected with that and later versions. With brokers older than 0.11.0.0, this method
+     * is unable get the topic configurations and always returns a null value.
+     *
+     * 
            If the topic does not exist, a null value is returned.
+     *
+     * @param topic the name of the topic for which the topic configuration should be obtained
+     * @return the topic configuration if the topic exists, or null if the topic did not exist
+     * @throws RetriableException if a retriable error occurs, the operation takes too long, or the
+     *         thread is interrupted while attempting to perform this operation
+     * @throws ConnectException if a non retriable error occurs
+     */
+    public Config describeTopicConfig(String topic) {
+        return describeTopicConfigs(topic).get(topic);
+    }
+
+    /**
+     * Attempt to fetch the topic configurations for the given topics.
+     * Apache Kafka added support for describing topic configurations in 0.11.0.0, so this method
+     * works as expected with that and later versions. With brokers older than 0.11.0.0, this method
+     * is unable get the topic configurations and always returns an empty set.
+     *
+     * 
            An entry with a null Config is placed into the resulting map for any topic that does
+     * not exist on the brokers.
+     *
+     * @param topicNames the topics to obtain configurations
+     * @return the map of topic configurations for each existing topic, or an empty map if none
+     * of the topics exist
+     * @throws RetriableException if a retriable error occurs, the operation takes too long, or the
+     *         thread is interrupted while attempting to perform this operation
+     * @throws ConnectException if a non retriable error occurs
+     */
+    public Map describeTopicConfigs(String... topicNames) {
+        if (topicNames == null) {
+            return Collections.emptyMap();
+        }
+        Collection topics = Arrays.stream(topicNames)
+                                          .filter(Objects::nonNull)
+                                          .map(String::trim)
+                                          .filter(s -> !s.isEmpty())
+                                          .collect(Collectors.toList());
+        if (topics.isEmpty()) {
+            return Collections.emptyMap();
+        }
+        String bootstrapServers = bootstrapServers();
+        String topicNameList = topics.stream().collect(Collectors.joining(", "));
+        Collection resources = topics.stream()
+                                                     .map(t -> new ConfigResource(ConfigResource.Type.TOPIC, t))
+                                                     .collect(Collectors.toList());
+
+        Map> newResults = admin.describeConfigs(resources, new DescribeConfigsOptions()).values();
+
+        // Iterate over each future so that we can handle individual failures like when some topics don't exist
+        Map result = new HashMap<>();
+        newResults.forEach((resource, configs) -> {
+            String topic = resource.name();
+            try {
+                result.put(topic, configs.get());
+            } catch (ExecutionException e) {
+                Throwable cause = e.getCause();
+                if (cause instanceof UnknownTopicOrPartitionException) {
+                    log.debug("Topic '{}' does not exist on the brokers at {}", topic, bootstrapServers);
+                    result.put(topic, null);
+                } else if (cause instanceof ClusterAuthorizationException || cause instanceof TopicAuthorizationException) {
+                    log.debug("Not authorized to describe topic config for topic '{}' on brokers at {}", topic, bootstrapServers);
+                } else if (cause instanceof UnsupportedVersionException) {
+                    log.debug("API to describe topic config for topic '{}' is unsupported on brokers at {}", topic, bootstrapServers);
+                } else if (cause instanceof TimeoutException) {
+                    String msg = String.format("Timed out while waiting to describe topic config for topic '%s' on brokers at %s",
+                            topic, bootstrapServers);
+                    throw new RetriableException(msg, e);
+                } else {
+                    String msg = String.format("Error while attempting to describe topic config for topic '%s' on brokers at %s",
+                            topic, bootstrapServers);
+                    throw new ConnectException(msg, e);
+                }
+            } catch (InterruptedException e) {
+                Thread.interrupted();
+                String msg = String.format("Interrupted while attempting to describe topic configs '%s'", topicNameList);
+                throw new RetriableException(msg, e);
+            }
+        });
+        return result;
+    }
+
     @Override
     public void close() {
         admin.close();
     }
 
+    public void close(Duration timeout) {
+        admin.close(timeout);
+    }
+
     private String bootstrapServers() {
         Object servers = adminConfig.get(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG);
         return servers != null ? servers.toString() : "";
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/util/TopicCreation.java b/connect/runtime/src/main/java/org/apache/kafka/connect/util/TopicCreation.java
new file mode 100644
index 0000000000000..f914ed92dc283
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/util/TopicCreation.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.util;
+
+import org.apache.kafka.connect.runtime.WorkerConfig;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.LinkedHashMap;
+import java.util.Map;
+import java.util.Set;
+
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.DEFAULT_TOPIC_CREATION_GROUP;
+
+/**
+ * Utility to be used by worker source tasks in order to create topics, if topic creation is
+ * enabled for source connectors at the worker and the connector configurations.
+ */
+public class TopicCreation {
+    private static final Logger log = LoggerFactory.getLogger(TopicCreation.class);
+    private static final TopicCreation EMPTY =
+            new TopicCreation(false, null, Collections.emptyMap(), Collections.emptySet());
+
+    private final boolean isTopicCreationEnabled;
+    private final TopicCreationGroup defaultTopicGroup;
+    private final Map topicGroups;
+    private final Set topicCache;
+
+    protected TopicCreation(boolean isTopicCreationEnabled,
+                            TopicCreationGroup defaultTopicGroup,
+                            Map topicGroups,
+                            Set topicCache) {
+        this.isTopicCreationEnabled = isTopicCreationEnabled;
+        this.defaultTopicGroup = defaultTopicGroup;
+        this.topicGroups = topicGroups;
+        this.topicCache = topicCache;
+    }
+
+    public static TopicCreation newTopicCreation(WorkerConfig workerConfig,
+            Map topicGroups) {
+        if (!workerConfig.topicCreationEnable() || topicGroups == null) {
+            return EMPTY;
+        }
+        Map groups = new LinkedHashMap<>(topicGroups);
+        groups.remove(DEFAULT_TOPIC_CREATION_GROUP);
+        return new TopicCreation(true, topicGroups.get(DEFAULT_TOPIC_CREATION_GROUP), groups, new HashSet<>());
+    }
+
+    /**
+     * Return an instance of this utility that represents what the state of the internal data
+     * structures should be when topic creation is disabled.
+     *
+     * @return the utility when topic creation is disabled
+     */
+    public static TopicCreation empty() {
+        return EMPTY;
+    }
+
+    /**
+     * Check whether topic creation is enabled for this utility instance. This is state is set at
+     * instantiation time and remains unchanged for the lifetime of every {@link TopicCreation}
+     * object.
+     *
+     * @return true if topic creation is enabled; false otherwise
+     */
+    public boolean isTopicCreationEnabled() {
+        return isTopicCreationEnabled;
+    }
+
+    /**
+     * Check whether topic creation may be required for a specific topic name.
+     *
+     * @return true if topic creation is enabled and the topic name is not in the topic cache;
+     * false otherwise
+     */
+    public boolean isTopicCreationRequired(String topic) {
+        return isTopicCreationEnabled && !topicCache.contains(topic);
+    }
+
+    /**
+     * Return the default topic creation group. This group is always defined when topic creation is
+     * enabled but is {@code null} if topic creation is disabled.
+     *
+     * @return the default topic creation group if topic creation is enabled; {@code null} otherwise
+     */
+    public TopicCreationGroup defaultTopicGroup() {
+        return defaultTopicGroup;
+    }
+
+    /**
+     * Return the topic creation groups defined for a source connector as a map of topic creation
+     * group name to topic creation group instance. This map maintains all the optionally defined
+     * groups besides the default group which is defined for any connector when topic creation is
+     * enabled.
+     *
+     * @return the map of all the topic creation groups besides the default group; may be empty
+     * but not {@code null}
+     */
+    public Map topicGroups() {
+        return topicGroups;
+    }
+
+    /**
+     * Inform this utility instance that a topic has been created and its creation will no
+     * longer be required. After {@link #addTopic(String)} is called for a give {@param topic}
+     * any subsequent calls to {@link #isTopicCreationRequired} will return {@code false} for the
+     * same topic.
+     *
+     * @param topic the topic name to mark as created
+     */
+    public void addTopic(String topic) {
+        if (isTopicCreationEnabled) {
+            topicCache.add(topic);
+        }
+    }
+
+    /**
+     * Get the first topic creation group that is configured to match the given {@param topic}
+     * name. If topic creation is enabled, any topic should match at least the default topic
+     * creation group.
+     *
+     * @param topic the topic name to match against group configurations
+     *
+     * @return the first group that matches the given topic
+     */
+    public TopicCreationGroup findFirstGroup(String topic) {
+        return topicGroups.values().stream()
+                .filter(group -> group.matches(topic))
+                .findFirst()
+                .orElse(defaultTopicGroup);
+    }
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/util/TopicCreationGroup.java b/connect/runtime/src/main/java/org/apache/kafka/connect/util/TopicCreationGroup.java
new file mode 100644
index 0000000000000..11970339ab845
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/util/TopicCreationGroup.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.util;
+
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.clients.admin.NewTopic;
+import org.apache.kafka.connect.runtime.SourceConnectorConfig;
+
+import java.util.Collections;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.regex.Pattern;
+
+import static org.apache.kafka.connect.runtime.SourceConnectorConfig.TOPIC_CREATION_GROUPS_CONFIG;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.DEFAULT_TOPIC_CREATION_GROUP;
+
+/**
+ * Utility to simplify creating and managing topics via the {@link Admin}.
+ */
+public class TopicCreationGroup {
+    private final String name;
+    private final Pattern inclusionPattern;
+    private final Pattern exclusionPattern;
+    private final int numPartitions;
+    private final short replicationFactor;
+    private final Map otherConfigs;
+
+    protected TopicCreationGroup(String group, SourceConnectorConfig config) {
+        this.name = group;
+        this.inclusionPattern = Pattern.compile(String.join("|", config.topicCreationInclude(group)));
+        this.exclusionPattern = Pattern.compile(String.join("|", config.topicCreationExclude(group)));
+        this.numPartitions = config.topicCreationPartitions(group);
+        this.replicationFactor = config.topicCreationReplicationFactor(group);
+        this.otherConfigs = config.topicCreationOtherConfigs(group);
+    }
+
+    /**
+     * Parses the configuration of a source connector and returns the topic creation groups
+     * defined in the given configuration as a map of group names to {@link TopicCreation} objects.
+     *
+     * @param config the source connector configuration
+     *
+     * @return the map of topic creation groups; may be empty but not {@code null}
+     */
+    public static Map configuredGroups(SourceConnectorConfig config) {
+        if (!config.usesTopicCreation()) {
+            return Collections.emptyMap();
+        }
+        List groupNames = config.getList(TOPIC_CREATION_GROUPS_CONFIG);
+        Map groups = new LinkedHashMap<>();
+        for (String group : groupNames) {
+            groups.put(group, new TopicCreationGroup(group, config));
+        }
+        // Even if there was a group called 'default' in the config, it will be overridden here.
+        // Order matters for all the topic groups besides the default, since it will be
+        // removed from this collection by the Worker
+        groups.put(DEFAULT_TOPIC_CREATION_GROUP, new TopicCreationGroup(DEFAULT_TOPIC_CREATION_GROUP, config));
+        return groups;
+    }
+
+    /**
+     * Return the name of the topic creation group.
+     *
+     * @return the name of the topic creation group
+     */
+    public String name() {
+        return name;
+    }
+
+    /**
+     * Answer whether this topic creation group is configured to allow the creation of the given
+     * {@param topic} name.
+     *
+     * @param topic the topic name to check against the groups configuration
+     *
+     * @return true if the topic name matches the inclusion regex and does
+     * not match the exclusion regex of this group's configuration; false otherwise
+     */
+    public boolean matches(String topic) {
+        return !exclusionPattern.matcher(topic).matches() && inclusionPattern.matcher(topic)
+                .matches();
+    }
+
+    /**
+     * Return the description for a new topic with the given {@param topic} name with the topic
+     * settings defined for this topic creation group.
+     *
+     * @param topic the name of the topic to be created
+     *
+     * @return the topic description of the given topic with settings of this topic creation group
+     */
+    public NewTopic newTopic(String topic) {
+        TopicAdmin.NewTopicBuilder builder = new TopicAdmin.NewTopicBuilder(topic);
+        return builder.partitions(numPartitions)
+                .replicationFactor(replicationFactor)
+                .config(otherConfigs)
+                .build();
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (!(o instanceof TopicCreationGroup)) {
+            return false;
+        }
+        TopicCreationGroup that = (TopicCreationGroup) o;
+        return Objects.equals(name, that.name)
+                && numPartitions == that.numPartitions
+                && replicationFactor == that.replicationFactor
+                && Objects.equals(inclusionPattern.pattern(), that.inclusionPattern.pattern())
+                && Objects.equals(exclusionPattern.pattern(), that.exclusionPattern.pattern())
+                && Objects.equals(otherConfigs, that.otherConfigs);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(name, numPartitions, replicationFactor, inclusionPattern.pattern(),
+                exclusionPattern.pattern(), otherConfigs
+        );
+    }
+
+    @Override
+    public String toString() {
+        return "TopicCreationGroup{" +
+                "name='" + name + '\'' +
+                ", inclusionPattern=" + inclusionPattern +
+                ", exclusionPattern=" + exclusionPattern +
+                ", numPartitions=" + numPartitions +
+                ", replicationFactor=" + replicationFactor +
+                ", otherConfigs=" + otherConfigs +
+                '}';
+    }
+}
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/converters/DoubleConverterTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/converters/DoubleConverterTest.java
new file mode 100644
index 0000000000000..acc3ddedc6815
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/converters/DoubleConverterTest.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.converters;
+
+import org.apache.kafka.common.serialization.DoubleSerializer;
+import org.apache.kafka.common.serialization.Serializer;
+import org.apache.kafka.connect.data.Schema;
+
+public class DoubleConverterTest extends NumberConverterTest {
+
+    public Double[] samples() {
+        return new Double[]{Double.MIN_VALUE, 1234.31, Double.MAX_VALUE};
+    }
+
+    @Override
+    protected Schema schema() {
+        return Schema.OPTIONAL_FLOAT64_SCHEMA;
+    }
+
+    @Override
+    protected NumberConverter createConverter() {
+        return new DoubleConverter();
+    }
+
+    @Override
+    protected Serializer createSerializer() {
+        return new DoubleSerializer();
+    }
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/converters/FloatConverterTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/converters/FloatConverterTest.java
new file mode 100644
index 0000000000000..e95ff56d9cfc8
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/converters/FloatConverterTest.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.converters;
+
+import org.apache.kafka.common.serialization.FloatSerializer;
+import org.apache.kafka.common.serialization.Serializer;
+import org.apache.kafka.connect.data.Schema;
+
+public class FloatConverterTest extends NumberConverterTest {
+
+    public Float[] samples() {
+        return new Float[]{Float.MIN_VALUE, 1234.31f, Float.MAX_VALUE};
+    }
+
+    @Override
+    protected Schema schema() {
+        return Schema.OPTIONAL_FLOAT32_SCHEMA;
+    }
+
+    @Override
+    protected NumberConverter createConverter() {
+        return new FloatConverter();
+    }
+
+    @Override
+    protected Serializer createSerializer() {
+        return new FloatSerializer();
+    }
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/converters/IntegerConverterTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/converters/IntegerConverterTest.java
new file mode 100644
index 0000000000000..0c9ed28d9da46
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/converters/IntegerConverterTest.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.converters;
+
+import org.apache.kafka.common.serialization.IntegerSerializer;
+import org.apache.kafka.common.serialization.Serializer;
+import org.apache.kafka.connect.data.Schema;
+
+public class IntegerConverterTest extends NumberConverterTest {
+
+    public Integer[] samples() {
+        return new Integer[]{Integer.MIN_VALUE, 1234, Integer.MAX_VALUE};
+    }
+
+    @Override
+    protected Schema schema() {
+        return Schema.OPTIONAL_INT32_SCHEMA;
+    }
+
+    @Override
+    protected NumberConverter createConverter() {
+        return new IntegerConverter();
+    }
+
+    @Override
+    protected Serializer createSerializer() {
+        return new IntegerSerializer();
+    }
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/converters/LongConverterTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/converters/LongConverterTest.java
new file mode 100644
index 0000000000000..35d26b70e892c
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/converters/LongConverterTest.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.converters;
+
+import org.apache.kafka.common.serialization.LongSerializer;
+import org.apache.kafka.common.serialization.Serializer;
+import org.apache.kafka.connect.data.Schema;
+
+public class LongConverterTest extends NumberConverterTest {
+
+    public Long[] samples() {
+        return new Long[]{Long.MIN_VALUE, 1234L, Long.MAX_VALUE};
+    }
+
+    @Override
+    protected Schema schema() {
+        return Schema.OPTIONAL_INT64_SCHEMA;
+    }
+
+    @Override
+    protected NumberConverter createConverter() {
+        return new LongConverter();
+    }
+
+    @Override
+    protected Serializer createSerializer() {
+        return new LongSerializer();
+    }
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/converters/NumberConverterTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/converters/NumberConverterTest.java
new file mode 100644
index 0000000000000..2bd07325b8a02
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/converters/NumberConverterTest.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.connect.converters;
+
+import org.apache.kafka.common.serialization.Serializer;
+import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.data.SchemaAndValue;
+import org.apache.kafka.connect.errors.DataException;
+import org.junit.Before;
+import org.junit.Test;
+
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+
+public abstract class NumberConverterTest {
+    private static final String TOPIC = "topic";
+    private static final String HEADER_NAME = "header";
+
+    private T[] samples;
+    private Schema schema;
+    private NumberConverter converter;
+    private Serializer serializer;
+
+    protected abstract T[] samples();
+
+    protected abstract NumberConverter createConverter();
+
+    protected abstract Serializer createSerializer();
+
+    protected abstract Schema schema();
+
+    @Before
+    public void setup() {
+        converter = createConverter();
+        serializer = createSerializer();
+        schema = schema();
+        samples = samples();
+    }
+
+    @Test
+    public void testConvertingSamplesToAndFromBytes() throws UnsupportedOperationException {
+        for (T sample : samples) {
+            byte[] expected = serializer.serialize(TOPIC, sample);
+
+            // Data conversion
+            assertArrayEquals(expected, converter.fromConnectData(TOPIC, schema, sample));
+            SchemaAndValue data = converter.toConnectData(TOPIC, expected);
+            assertEquals(schema, data.schema());
+            assertEquals(sample, data.value());
+
+            // Header conversion
+            assertArrayEquals(expected, converter.fromConnectHeader(TOPIC, HEADER_NAME, schema, sample));
+            data = converter.toConnectHeader(TOPIC, HEADER_NAME, expected);
+            assertEquals(schema, data.schema());
+            assertEquals(sample, data.value());
+        }
+    }
+
+    @Test(expected = DataException.class)
+    public void testDeserializingDataWithTooManyBytes() {
+        converter.toConnectData(TOPIC, new byte[10]);
+    }
+
+    @Test(expected = DataException.class)
+    public void testDeserializingHeaderWithTooManyBytes() {
+        converter.toConnectHeader(TOPIC, HEADER_NAME, new byte[10]);
+    }
+
+    @Test(expected = DataException.class)
+    public void testSerializingIncorrectType() {
+        converter.fromConnectData(TOPIC, schema, "not a valid number");
+    }
+
+    @Test(expected = DataException.class)
+    public void testSerializingIncorrectHeader() {
+        converter.fromConnectHeader(TOPIC, HEADER_NAME, schema, "not a valid number");
+    }
+
+    @Test
+    public void testNullToBytes() {
+        assertEquals(null, converter.fromConnectData(TOPIC, schema, null));
+    }
+
+    @Test
+    public void testBytesNullToNumber() {
+        SchemaAndValue data = converter.toConnectData(TOPIC, null);
+        assertEquals(schema(), data.schema());
+        assertNull(data.value());
+    }
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/converters/ShortConverterTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/converters/ShortConverterTest.java
new file mode 100644
index 0000000000000..d1237c9ca0442
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/converters/ShortConverterTest.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.connect.converters;
+
+import org.apache.kafka.common.serialization.Serializer;
+import org.apache.kafka.common.serialization.ShortSerializer;
+import org.apache.kafka.connect.data.Schema;
+
+public class ShortConverterTest extends NumberConverterTest {
+
+    public Short[] samples() {
+        return new Short[]{Short.MIN_VALUE, 123, Short.MAX_VALUE};
+    }
+
+    @Override
+    protected Schema schema() {
+        return Schema.OPTIONAL_INT16_SCHEMA;
+    }
+
+    @Override
+    protected NumberConverter createConverter() {
+        return new ShortConverter();
+    }
+
+    @Override
+    protected Serializer createSerializer() {
+        return new ShortSerializer();
+    }
+}
+
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/BlockingConnectorTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/BlockingConnectorTest.java
new file mode 100644
index 0000000000000..abc9a938c6606
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/BlockingConnectorTest.java
@@ -0,0 +1,799 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.consumer.OffsetAndMetadata;
+import org.apache.kafka.clients.producer.RecordMetadata;
+import org.apache.kafka.common.TopicPartition;
+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.connect.connector.Connector;
+import org.apache.kafka.connect.connector.ConnectorContext;
+import org.apache.kafka.connect.connector.Task;
+import org.apache.kafka.connect.runtime.Worker;
+import org.apache.kafka.connect.runtime.rest.errors.ConnectRestException;
+import org.apache.kafka.connect.runtime.rest.resources.ConnectorsResource;
+import org.apache.kafka.connect.sink.SinkConnector;
+import org.apache.kafka.connect.sink.SinkRecord;
+import org.apache.kafka.connect.sink.SinkTask;
+import org.apache.kafka.connect.sink.SinkTaskContext;
+import org.apache.kafka.connect.source.SourceConnector;
+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.util.clusters.EmbeddedConnectCluster;
+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.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Properties;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import static org.apache.kafka.connect.runtime.ConnectorConfig.CONNECTOR_CLASS_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.test.TestUtils.waitForCondition;
+import static org.junit.Assert.assertThrows;
+
+/**
+ * Tests situations during which certain connector operations, such as start, validation,
+ * configuration and others, take longer than expected.
+ */
+@Category(IntegrationTest.class)
+public class BlockingConnectorTest {
+
+    private static final Logger log = LoggerFactory.getLogger(BlockingConnectorTest.class);
+
+    private static final int NUM_WORKERS = 1;
+    private static final String BLOCKING_CONNECTOR_NAME = "blocking-connector";
+    private static final String NORMAL_CONNECTOR_NAME = "normal-connector";
+    private static final String TEST_TOPIC = "normal-topic";
+    private static final int NUM_RECORDS_PRODUCED = 100;
+    private static final long CONNECT_WORKER_STARTUP_TIMEOUT = TimeUnit.SECONDS.toMillis(30);
+    private static final long RECORD_TRANSFER_DURATION_MS = TimeUnit.SECONDS.toMillis(30);
+    private static final long REST_REQUEST_TIMEOUT = Worker.CONNECTOR_GRACEFUL_SHUTDOWN_TIMEOUT_MS * 2;
+
+    private static final String CONNECTOR_INITIALIZE = "Connector::initialize";
+    private static final String CONNECTOR_INITIALIZE_WITH_TASK_CONFIGS = "Connector::initializeWithTaskConfigs";
+    private static final String CONNECTOR_START = "Connector::start";
+    private static final String CONNECTOR_RECONFIGURE = "Connector::reconfigure";
+    private static final String CONNECTOR_TASK_CLASS = "Connector::taskClass";
+    private static final String CONNECTOR_TASK_CONFIGS = "Connector::taskConfigs";
+    private static final String CONNECTOR_STOP = "Connector::stop";
+    private static final String CONNECTOR_VALIDATE = "Connector::validate";
+    private static final String CONNECTOR_CONFIG = "Connector::config";
+    private static final String CONNECTOR_VERSION = "Connector::version";
+    private static final String TASK_START = "Task::start";
+    private static final String TASK_STOP = "Task::stop";
+    private static final String TASK_VERSION = "Task::version";
+    private static final String SINK_TASK_INITIALIZE = "SinkTask::initialize";
+    private static final String SINK_TASK_PUT = "SinkTask::put";
+    private static final String SINK_TASK_FLUSH = "SinkTask::flush";
+    private static final String SINK_TASK_PRE_COMMIT = "SinkTask::preCommit";
+    private static final String SINK_TASK_OPEN = "SinkTask::open";
+    private static final String SINK_TASK_ON_PARTITIONS_ASSIGNED = "SinkTask::onPartitionsAssigned";
+    private static final String SINK_TASK_CLOSE = "SinkTask::close";
+    private static final String SINK_TASK_ON_PARTITIONS_REVOKED = "SinkTask::onPartitionsRevoked";
+    private static final String SOURCE_TASK_INITIALIZE = "SourceTask::initialize";
+    private static final String SOURCE_TASK_POLL = "SourceTask::poll";
+    private static final String SOURCE_TASK_COMMIT = "SourceTask::commit";
+    private static final String SOURCE_TASK_COMMIT_RECORD = "SourceTask::commitRecord";
+    private static final String SOURCE_TASK_COMMIT_RECORD_WITH_METADATA = "SourceTask::commitRecordWithMetadata";
+
+    private EmbeddedConnectCluster connect;
+    private ConnectorHandle normalConnectorHandle;
+
+    @Before
+    public void setup() throws Exception {
+        // Artificially reduce the REST request timeout so that these don't take forever
+        ConnectorsResource.setRequestTimeout(REST_REQUEST_TIMEOUT);
+        // build a Connect cluster backed by Kafka and Zk
+        connect = new EmbeddedConnectCluster.Builder()
+                .name("connect-cluster")
+                .numWorkers(NUM_WORKERS)
+                .numBrokers(1)
+                .workerProps(new HashMap<>())
+                .brokerProps(new Properties())
+                .build();
+
+        // start the clusters
+        connect.start();
+
+        // wait for the Connect REST API to become available. necessary because of the reduced REST
+        // request timeout; otherwise, we may get an unexpected 500 with our first real REST request
+        // if the worker is still getting on its feet.
+        waitForCondition(
+            () -> connect.requestGet(connect.endpointForResource("connectors/nonexistent")).getStatus() == 404,
+            CONNECT_WORKER_STARTUP_TIMEOUT,
+            "Worker did not complete startup in time"
+        );
+    }
+
+    @After
+    public void close() {
+        // stop all Connect, Kafka and Zk threads.
+        connect.stop();
+        ConnectorsResource.resetRequestTimeout();
+        Block.resetBlockLatch();
+    }
+
+    @Test
+    public void testBlockInConnectorValidate() throws Exception {
+        log.info("Starting test testBlockInConnectorValidate");
+        assertThrows(ConnectRestException.class, () -> createConnectorWithBlock(ValidateBlockingConnector.class, CONNECTOR_VALIDATE));
+        // Will NOT assert that connector has failed, since the request should fail before it's even created
+
+        // Connector should already be blocked so this should return immediately, but check just to
+        // make sure that it actually did block
+        Block.waitForBlock();
+
+        createNormalConnector();
+        verifyNormalConnector();
+    }
+
+    @Test
+    public void testBlockInConnectorConfig() throws Exception {
+        log.info("Starting test testBlockInConnectorConfig");
+        assertThrows(ConnectRestException.class, () -> createConnectorWithBlock(ConfigBlockingConnector.class, CONNECTOR_CONFIG));
+        // Will NOT assert that connector has failed, since the request should fail before it's even created
+
+        // Connector should already be blocked so this should return immediately, but check just to
+        // make sure that it actually did block
+        Block.waitForBlock();
+
+        createNormalConnector();
+        verifyNormalConnector();
+    }
+
+    @Test
+    public void testBlockInConnectorInitialize() throws Exception {
+        log.info("Starting test testBlockInConnectorInitialize");
+        createConnectorWithBlock(InitializeBlockingConnector.class, CONNECTOR_INITIALIZE);
+        Block.waitForBlock();
+
+        createNormalConnector();
+        verifyNormalConnector();
+    }
+
+    @Test
+    public void testBlockInConnectorStart() throws Exception {
+        log.info("Starting test testBlockInConnectorStart");
+        createConnectorWithBlock(BlockingConnector.class, CONNECTOR_START);
+        Block.waitForBlock();
+
+        createNormalConnector();
+        verifyNormalConnector();
+    }
+
+    @Test
+    public void testBlockInConnectorStop() throws Exception {
+        log.info("Starting test testBlockInConnectorStop");
+        createConnectorWithBlock(BlockingConnector.class, CONNECTOR_STOP);
+        waitForConnectorStart(BLOCKING_CONNECTOR_NAME);
+        connect.deleteConnector(BLOCKING_CONNECTOR_NAME);
+        Block.waitForBlock();
+
+        createNormalConnector();
+        verifyNormalConnector();
+    }
+
+    @Test
+    public void testBlockInSourceTaskStart() throws Exception {
+        log.info("Starting test testBlockInSourceTaskStart");
+        createConnectorWithBlock(BlockingSourceConnector.class, TASK_START);
+        Block.waitForBlock();
+
+        createNormalConnector();
+        verifyNormalConnector();
+    }
+
+    @Test
+    public void testBlockInSourceTaskStop() throws Exception {
+        log.info("Starting test testBlockInSourceTaskStop");
+        createConnectorWithBlock(BlockingSourceConnector.class, TASK_STOP);
+        waitForConnectorStart(BLOCKING_CONNECTOR_NAME);
+        connect.deleteConnector(BLOCKING_CONNECTOR_NAME);
+        Block.waitForBlock();
+
+        createNormalConnector();
+        verifyNormalConnector();
+    }
+
+    @Test
+    public void testBlockInSinkTaskStart() throws Exception {
+        log.info("Starting test testBlockInSinkTaskStart");
+        createConnectorWithBlock(BlockingSinkConnector.class, TASK_START);
+        Block.waitForBlock();
+
+        createNormalConnector();
+        verifyNormalConnector();
+    }
+
+    @Test
+    public void testBlockInSinkTaskStop() throws Exception {
+        log.info("Starting test testBlockInSinkTaskStop");
+        createConnectorWithBlock(BlockingSinkConnector.class, TASK_STOP);
+        waitForConnectorStart(BLOCKING_CONNECTOR_NAME);
+        connect.deleteConnector(BLOCKING_CONNECTOR_NAME);
+        Block.waitForBlock();
+
+        createNormalConnector();
+        verifyNormalConnector();
+    }
+
+    @Test
+    public void testWorkerRestartWithBlockInConnectorStart() throws Exception {
+        log.info("Starting test testWorkerRestartWithBlockInConnectorStart");
+        createConnectorWithBlock(BlockingConnector.class, CONNECTOR_START);
+        // First instance of the connector should block on startup
+        Block.waitForBlock();
+        createNormalConnector();
+        connect.removeWorker();
+
+        connect.addWorker();
+        // After stopping the only worker and restarting it, a new instance of the blocking
+        // connector should be created and we can ensure that it blocks again
+        Block.waitForBlock();
+        verifyNormalConnector();
+    }
+
+    @Test
+    public void testWorkerRestartWithBlockInConnectorStop() throws Exception {
+        log.info("Starting test testWorkerRestartWithBlockInConnectorStop");
+        createConnectorWithBlock(BlockingConnector.class, CONNECTOR_STOP);
+        waitForConnectorStart(BLOCKING_CONNECTOR_NAME);
+        createNormalConnector();
+        waitForConnectorStart(NORMAL_CONNECTOR_NAME);
+        connect.removeWorker();
+        Block.waitForBlock();
+
+        connect.addWorker();
+        waitForConnectorStart(BLOCKING_CONNECTOR_NAME);
+        verifyNormalConnector();
+    }
+
+    private void createConnectorWithBlock(Class connectorClass, String block) {
+        Map props = new HashMap<>();
+        props.put(CONNECTOR_CLASS_CONFIG, connectorClass.getName());
+        props.put(TASKS_MAX_CONFIG, "1");
+        props.put(TOPICS_CONFIG, "t1"); // Required for sink connectors
+        props.put(Block.BLOCK_CONFIG, Objects.requireNonNull(block));
+        log.info("Creating blocking connector of type {} with block in {}", connectorClass.getSimpleName(), block);
+        try {
+            connect.configureConnector(BLOCKING_CONNECTOR_NAME, props);
+        } catch (RuntimeException e) {
+            log.info("Failed to create connector", e);
+            throw e;
+        }
+    }
+
+    private void createNormalConnector() {
+        connect.kafka().createTopic(TEST_TOPIC, 3);
+
+        normalConnectorHandle = RuntimeHandles.get().connectorHandle(NORMAL_CONNECTOR_NAME);
+        normalConnectorHandle.expectedRecords(NUM_RECORDS_PRODUCED);
+        normalConnectorHandle.expectedCommits(NUM_RECORDS_PRODUCED);
+
+        Map props = new HashMap<>();
+        props.put(CONNECTOR_CLASS_CONFIG, MonitorableSourceConnector.class.getName());
+        props.put(TASKS_MAX_CONFIG, "1");
+        props.put(MonitorableSourceConnector.TOPIC_CONFIG, TEST_TOPIC);
+        log.info("Creating normal connector");
+        try {
+            connect.configureConnector(NORMAL_CONNECTOR_NAME, props);
+        } catch (RuntimeException e) {
+            log.info("Failed to create connector", e);
+            throw e;
+        }
+    }
+
+    private void waitForConnectorStart(String connector) throws InterruptedException {
+        connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(
+            connector,
+            0,
+            String.format(
+                "Failed to observe transition to 'RUNNING' state for connector '%s' in time",
+                connector
+            )
+        );
+    }
+
+    private void verifyNormalConnector() throws InterruptedException {
+        waitForConnectorStart(NORMAL_CONNECTOR_NAME);
+        normalConnectorHandle.awaitRecords(RECORD_TRANSFER_DURATION_MS);
+        normalConnectorHandle.awaitCommits(RECORD_TRANSFER_DURATION_MS);
+    }
+
+    private static class Block {
+        private static CountDownLatch blockLatch;
+
+        private final String block;
+
+        public static final String BLOCK_CONFIG = "block";
+
+        private static ConfigDef config() {
+            return new ConfigDef()
+                .define(
+                    BLOCK_CONFIG,
+                    ConfigDef.Type.STRING,
+                    "",
+                    ConfigDef.Importance.MEDIUM,
+                    "Where to block indefinitely, e.g., 'Connector::start', 'Connector::initialize', " 
+                        + "'Connector::taskConfigs', 'Task::version', 'SinkTask::put', 'SourceTask::poll'"
+                );
+        }
+
+        public static void waitForBlock() throws InterruptedException {
+            synchronized (Block.class) {
+                if (blockLatch == null) {
+                    throw new IllegalArgumentException("No connector has been created yet");
+                }
+            }
+
+            log.debug("Waiting for connector to block");
+            blockLatch.await();
+            log.debug("Connector should now be blocked");
+        }
+
+        // Note that there is only ever at most one global block latch at a time, which makes tests that
+        // use blocks in multiple places impossible. If necessary, this can be addressed in the future by
+        // adding support for multiple block latches at a time, possibly identifiable by a connector/task
+        // ID, the location of the expected block, or both.
+        public static void resetBlockLatch() {
+            synchronized (Block.class) {
+                if (blockLatch != null) {
+                    blockLatch.countDown();
+                    blockLatch = null;
+                }
+            }
+        }
+
+        public Block(Map props) {
+            this(new AbstractConfig(config(), props).getString(BLOCK_CONFIG));
+        }
+
+        public Block(String block) {
+            this.block = block;
+            synchronized (Block.class) {
+                if (blockLatch != null) {
+                    blockLatch.countDown();
+                }
+                blockLatch = new CountDownLatch(1);
+            }
+        }
+
+        public Map taskConfig() {
+            return Collections.singletonMap(BLOCK_CONFIG, block);
+        }
+
+        public void maybeBlockOn(String block) {
+            if (block.equals(this.block)) {
+                log.info("Will block on {}", block);
+                blockLatch.countDown();
+                while (true) {
+                    try {
+                        Thread.sleep(Long.MAX_VALUE);
+                    } catch (InterruptedException e) {
+                        // No-op. Just keep blocking.
+                    }
+                }
+            } else {
+                log.debug("Will not block on {}", block);
+            }
+        }
+    }
+
+    // Used to test blocks in Connector (as opposed to Task) methods
+    public static class BlockingConnector extends SourceConnector {
+
+        private Block block;
+
+        // No-args constructor required by the framework
+        public BlockingConnector() {
+            this(null);
+        }
+
+        protected BlockingConnector(String block) {
+            this.block = new Block(block);
+        }
+
+        @Override
+        public void initialize(ConnectorContext ctx) {
+            block.maybeBlockOn(CONNECTOR_INITIALIZE);
+            super.initialize(ctx);
+        }
+
+        @Override
+        public void initialize(ConnectorContext ctx, List> taskConfigs) {
+            block.maybeBlockOn(CONNECTOR_INITIALIZE_WITH_TASK_CONFIGS);
+            super.initialize(ctx, taskConfigs);
+        }
+
+        @Override
+        public void start(Map props) {
+            this.block = new Block(props);
+            block.maybeBlockOn(CONNECTOR_START);
+        }
+
+        @Override
+        public void reconfigure(Map props) {
+            block.maybeBlockOn(CONNECTOR_RECONFIGURE);
+            super.reconfigure(props);
+        }
+
+        @Override
+        public Class taskClass() {
+            block.maybeBlockOn(CONNECTOR_TASK_CLASS);
+            return BlockingTask.class;
+        }
+
+        @Override
+        public List> taskConfigs(int maxTasks) {
+            block.maybeBlockOn(CONNECTOR_TASK_CONFIGS);
+            return Collections.singletonList(Collections.emptyMap());
+        }
+
+        @Override
+        public void stop() {
+            block.maybeBlockOn(CONNECTOR_STOP);
+        }
+
+        @Override
+        public Config validate(Map connectorConfigs) {
+            block.maybeBlockOn(CONNECTOR_VALIDATE);
+            return super.validate(connectorConfigs);
+        }
+
+        @Override
+        public ConfigDef config() {
+            block.maybeBlockOn(CONNECTOR_CONFIG);
+            return Block.config();
+        }
+
+        @Override
+        public String version() {
+            block.maybeBlockOn(CONNECTOR_VERSION);
+            return "0.0.0";
+        }
+
+        public static class BlockingTask extends SourceTask {
+            @Override
+            public void start(Map props) {
+            }
+
+            @Override
+            public List poll() {
+                return null;
+            }
+
+            @Override
+            public void stop() {
+            }
+
+            @Override
+            public String version() {
+                return "0.0.0";
+            }
+        }
+    }
+
+    // Some methods are called before Connector::start, so we use this as a workaround
+    public static class InitializeBlockingConnector extends BlockingConnector {
+        public InitializeBlockingConnector() {
+            super(CONNECTOR_INITIALIZE);
+        }
+    }
+
+    public static class ConfigBlockingConnector extends BlockingConnector {
+        public ConfigBlockingConnector() {
+            super(CONNECTOR_CONFIG);
+        }
+    }
+
+    public static class ValidateBlockingConnector extends BlockingConnector {
+        public ValidateBlockingConnector() {
+            super(CONNECTOR_VALIDATE);
+        }
+    }
+
+    // Used to test blocks in SourceTask methods
+    public static class BlockingSourceConnector extends SourceConnector {
+
+        private Map props;
+        private final Class taskClass;
+
+        // No-args constructor required by the framework
+        public BlockingSourceConnector() {
+            this(BlockingSourceTask.class);
+        }
+
+        protected BlockingSourceConnector(Class taskClass) {
+            this.taskClass = taskClass;
+        }
+
+        @Override
+        public void start(Map props) {
+            this.props = props;
+        }
+
+        @Override
+        public Class taskClass() {
+            return taskClass;
+        }
+
+        @Override
+        public List> taskConfigs(int maxTasks) {
+            return IntStream.range(0, maxTasks)
+                .mapToObj(i -> new HashMap<>(props))
+                .collect(Collectors.toList());
+        }
+
+        @Override
+        public void stop() {
+        }
+
+        @Override
+        public Config validate(Map connectorConfigs) {
+            return super.validate(connectorConfigs);
+        }
+
+        @Override
+        public ConfigDef config() {
+            return Block.config();
+        }
+
+        @Override
+        public String version() {
+            return "0.0.0";
+        }
+
+        public static class BlockingSourceTask extends SourceTask {
+            private Block block;
+
+            // No-args constructor required by the framework
+            public BlockingSourceTask() {
+                this(null);
+            }
+
+            protected BlockingSourceTask(String block) {
+                this.block = new Block(block);
+            }
+
+            @Override
+            public void start(Map props) {
+                this.block = new Block(props);
+                block.maybeBlockOn(TASK_START);
+            }
+
+            @Override
+            public List poll() {
+                block.maybeBlockOn(SOURCE_TASK_POLL);
+                return null;
+            }
+
+            @Override
+            public void stop() {
+                block.maybeBlockOn(TASK_STOP);
+            }
+
+            @Override
+            public String version() {
+                block.maybeBlockOn(TASK_VERSION);
+                return "0.0.0";
+            }
+
+            @Override
+            public void initialize(SourceTaskContext context) {
+                block.maybeBlockOn(SOURCE_TASK_INITIALIZE);
+                super.initialize(context);
+            }
+
+            @Override
+            public void commit() throws InterruptedException {
+                block.maybeBlockOn(SOURCE_TASK_COMMIT);
+                super.commit();
+            }
+
+            @Override
+            @SuppressWarnings("deprecation")
+            public void commitRecord(SourceRecord record) throws InterruptedException {
+                block.maybeBlockOn(SOURCE_TASK_COMMIT_RECORD);
+                super.commitRecord(record);
+            }
+
+            @Override
+            public void commitRecord(SourceRecord record, RecordMetadata metadata) throws InterruptedException {
+                block.maybeBlockOn(SOURCE_TASK_COMMIT_RECORD_WITH_METADATA);
+                super.commitRecord(record, metadata);
+            }
+        }
+    }
+
+    public static class TaskInitializeBlockingSourceConnector extends BlockingSourceConnector {
+        public TaskInitializeBlockingSourceConnector() {
+            super(InitializeBlockingSourceTask.class);
+        }
+
+        public static class InitializeBlockingSourceTask extends BlockingSourceTask {
+            public InitializeBlockingSourceTask() {
+                super(SOURCE_TASK_INITIALIZE);
+            }
+        }
+    }
+
+    // Used to test blocks in SinkTask methods
+    public static class BlockingSinkConnector extends SinkConnector {
+
+        private Map props;
+        private final Class taskClass;
+
+        // No-args constructor required by the framework
+        public BlockingSinkConnector() {
+            this(BlockingSinkTask.class);
+        }
+
+        protected BlockingSinkConnector(Class taskClass) {
+            this.taskClass = taskClass;
+        }
+
+        @Override
+        public void start(Map props) {
+            this.props = props;
+        }
+
+        @Override
+        public Class taskClass() {
+            return taskClass;
+        }
+
+        @Override
+        public List> taskConfigs(int maxTasks) {
+            return IntStream.rangeClosed(0, maxTasks)
+                .mapToObj(i -> new HashMap<>(props))
+                .collect(Collectors.toList());
+        }
+
+        @Override
+        public void stop() {
+        }
+
+        @Override
+        public Config validate(Map connectorConfigs) {
+            return super.validate(connectorConfigs);
+        }
+
+        @Override
+        public ConfigDef config() {
+            return Block.config();
+        }
+
+        @Override
+        public String version() {
+            return "0.0.0";
+        }
+
+        public static class BlockingSinkTask extends SinkTask {
+            private Block block;
+
+            // No-args constructor required by the framework
+            public BlockingSinkTask() {
+                this(null);
+            }
+
+            protected BlockingSinkTask(String block) {
+                this.block = new Block(block);
+            }
+
+            @Override
+            public void start(Map props) {
+                this.block = new Block(props);
+                block.maybeBlockOn(TASK_START);
+            }
+
+            @Override
+            public void put(Collection records) {
+                block.maybeBlockOn(SINK_TASK_PUT);
+            }
+
+            @Override
+            public void stop() {
+                block.maybeBlockOn(TASK_STOP);
+            }
+
+            @Override
+            public String version() {
+                block.maybeBlockOn(TASK_VERSION);
+                return "0.0.0";
+            }
+
+            @Override
+            public void initialize(SinkTaskContext context) {
+                block.maybeBlockOn(SINK_TASK_INITIALIZE);
+                super.initialize(context);
+            }
+
+            @Override
+            public void flush(Map currentOffsets) {
+                block.maybeBlockOn(SINK_TASK_FLUSH);
+                super.flush(currentOffsets);
+            }
+
+            @Override
+            public Map preCommit(Map currentOffsets) {
+                block.maybeBlockOn(SINK_TASK_PRE_COMMIT);
+                return super.preCommit(currentOffsets);
+            }
+
+            @Override
+            public void open(Collection partitions) {
+                block.maybeBlockOn(SINK_TASK_OPEN);
+                super.open(partitions);
+            }
+
+            @Override
+            @SuppressWarnings("deprecation")
+            public void onPartitionsAssigned(Collection partitions) {
+                block.maybeBlockOn(SINK_TASK_ON_PARTITIONS_ASSIGNED);
+                super.onPartitionsAssigned(partitions);
+            }
+
+            @Override
+            public void close(Collection partitions) {
+                block.maybeBlockOn(SINK_TASK_CLOSE);
+                super.close(partitions);
+            }
+
+            @Override
+            @SuppressWarnings("deprecation")
+            public void onPartitionsRevoked(Collection partitions) {
+                block.maybeBlockOn(SINK_TASK_ON_PARTITIONS_REVOKED);
+                super.onPartitionsRevoked(partitions);
+            }
+        }
+    }
+
+    public static class TaskInitializeBlockingSinkConnector extends BlockingSinkConnector {
+        public TaskInitializeBlockingSinkConnector() {
+            super(InitializeBlockingSinkTask.class);
+        }
+
+        public static class InitializeBlockingSinkTask extends BlockingSinkTask {
+            public InitializeBlockingSinkTask() {
+                super(SINK_TASK_INITIALIZE);
+            }
+        }
+    }
+
+    // We don't declare a class here that blocks in the version() method since that method is used
+    // in plugin path scanning. Until/unless plugin path scanning is altered to not block completely
+    // on connectors' version() methods, we can't even declare a class that does that without
+    // causing the workers in this test to hang on startup.
+}
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..bd981832e1174
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ConnectWorkerIntegrationTest.java
@@ -0,0 +1,305 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.integration.MonitorableSourceConnector.TOPIC_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.TopicCreationConfig.DEFAULT_TOPIC_CREATION_PREFIX;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.PARTITIONS_CONFIG;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.REPLICATION_FACTOR_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 int NUM_TASKS = 4;
+    private static final String CONNECTOR_NAME = "simple-source";
+    private static final String TOPIC_NAME = "test-topic";
+
+    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();
+
+        // create test topic
+        connect.kafka().createTopic(TOPIC_NAME, NUM_TOPIC_PARTITIONS);
+
+        // set up props for the source connector
+        Map props = defaultSourceConnectorProps(TOPIC_NAME);
+
+        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, NUM_TASKS,
+                "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, NUM_TASKS,
+                "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;
+
+        // setup up props for the source connector
+        Map props = defaultSourceConnectorProps(TOPIC_NAME);
+        // Properties for the source connector. The task should fail at startup due to the bad broker address.
+        props.put(TASKS_MAX_CONFIG, Objects.toString(numTasks));
+        props.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, props);
+
+        connect.assertions().assertConnectorIsRunningAndTasksHaveFailed(CONNECTOR_NAME, numTasks,
+                "Connector tasks did not fail in time");
+
+        // Reconfigure the connector without the bad broker address.
+        props.remove(CONNECTOR_CLIENT_PRODUCER_OVERRIDES_PREFIX + BOOTSTRAP_SERVERS_CONFIG);
+        connect.configureConnector(CONNECTOR_NAME, props);
+
+        // 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(TOPIC_NAME, NUM_TOPIC_PARTITIONS);
+
+        // set up props for the source connector
+        Map props = defaultSourceConnectorProps(TOPIC_NAME);
+
+        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 props = defaultSourceConnectorProps(TOPIC_NAME);
+        props.put(CONNECTOR_CLASS_CONFIG, MonitorableSourceConnector.class.getSimpleName());
+
+        // start the connector with only one task
+        int initialNumTasks = 1;
+        props.put(TASKS_MAX_CONFIG, String.valueOf(initialNumTasks));
+        connect.configureConnector(CONNECTOR_NAME, props);
+        connect.assertions().assertConnectorAndExactlyNumTasksAreRunning(CONNECTOR_NAME,
+                initialNumTasks, "Connector tasks did not start in time");
+
+        // then reconfigure it to use more tasks
+        int increasedNumTasks = 5;
+        props.put(TASKS_MAX_CONFIG, String.valueOf(increasedNumTasks));
+        connect.configureConnector(CONNECTOR_NAME, props);
+        connect.assertions().assertConnectorAndExactlyNumTasksAreRunning(CONNECTOR_NAME,
+                increasedNumTasks, "Connector task statuses did not update in time.");
+
+        // then reconfigure it to use fewer tasks
+        int decreasedNumTasks = 3;
+        props.put(TASKS_MAX_CONFIG, String.valueOf(decreasedNumTasks));
+        connect.configureConnector(CONNECTOR_NAME, props);
+        connect.assertions().assertConnectorAndExactlyNumTasksAreRunning(CONNECTOR_NAME,
+                decreasedNumTasks, "Connector task statuses did not update in time.");
+    }
+
+    private Map defaultSourceConnectorProps(String topic) {
+        // 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(TOPIC_CONFIG, topic);
+        props.put("throughput", String.valueOf(10));
+        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());
+        props.put(DEFAULT_TOPIC_CREATION_PREFIX + REPLICATION_FACTOR_CONFIG, String.valueOf(1));
+        props.put(DEFAULT_TOPIC_CREATION_PREFIX + PARTITIONS_CONFIG, String.valueOf(1));
+        return props;
+    }
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ConnectorClientPolicyIntegrationTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ConnectorClientPolicyIntegrationTest.java
new file mode 100644
index 0000000000000..db324fa5a01e0
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ConnectorClientPolicyIntegrationTest.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 ConnectorClientPolicyIntegrationTest {
+
+    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
new file mode 100644
index 0000000000000..06bc37352e455
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ConnectorHandle.java
@@ -0,0 +1,343 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.DataException;
+import org.apache.kafka.connect.sink.SinkRecord;
+import org.slf4j.Logger;
+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.function.Consumer;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+/**
+ * A handle to a connector executing in a Connect cluster.
+ */
+public class ConnectorHandle {
+
+    private static final Logger log = LoggerFactory.getLogger(ConnectorHandle.class);
+
+    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;
+    }
+
+    /**
+     * Get or create a task handle for a given task id. The task need not be created when this method is called. If the
+     * handle is called before the task is created, the task will bind to the handle once it starts (or restarts).
+     *
+     * @param taskId the task id
+     * @return a non-null {@link TaskHandle}
+     */
+    public TaskHandle taskHandle(String taskId) {
+        return taskHandle(taskId, null);
+    }
+
+    /**
+     * Get or create a task handle for a given task id. The task need not be created when this method is called. If the
+     * handle is called before the task is created, the task will bind to the handle once it starts (or restarts).
+     *
+     * @param taskId the task id
+     * @param consumer A callback invoked when a sink task processes a record.
+     * @return a non-null {@link TaskHandle}
+     */
+    public TaskHandle taskHandle(String taskId, Consumer consumer) {
+        return taskHandles.computeIfAbsent(taskId, k -> new TaskHandle(this, taskId, consumer));
+    }
+
+    /**
+     * 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();
+    }
+
+    /**
+     * Delete the task handle for this task id.
+     *
+     * @param taskId the task id.
+     */
+    public void deleteTask(String taskId) {
+        log.info("Removing handle for {} task in connector {}", taskId, connectorName);
+        taskHandles.remove(taskId);
+    }
+
+    /**
+     * 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 expected number of commits
+     */
+    public void expectedCommits(int expected) {
+        expectedCommits = expected;
+        recordsToCommitLatch = new CountDownLatch(expected);
+    }
+
+    /**
+     * Record a message arrival at the connector.
+     */
+    public void record() {
+        if (recordsRemainingLatch != null) {
+            recordsRemainingLatch.countDown();
+        }
+    }
+
+    /**
+     * 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 timeout max duration to wait for records
+     * @throws InterruptedException if another threads interrupts this one while waiting for records
+     */
+    public void awaitRecords(long timeout) throws InterruptedException {
+        if (recordsRemainingLatch == null || expectedRecords < 0) {
+            throw new IllegalStateException("expectedRecords() was not set for this connector?");
+        }
+        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,
+                    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{" +
+                "connectorName='" + connectorName + '\'' +
+                '}';
+    }
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ConnectorTopicsIntegrationTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ConnectorTopicsIntegrationTest.java
new file mode 100644
index 0000000000000..75374a94819ac
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ConnectorTopicsIntegrationTest.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.integration;
+
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.connect.runtime.distributed.DistributedConfig;
+import org.apache.kafka.connect.runtime.rest.errors.ConnectRestException;
+import org.apache.kafka.connect.storage.KafkaStatusBackingStore;
+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.Test;
+import org.junit.experimental.categories.Category;
+
+import java.nio.charset.StandardCharsets;
+import java.time.Duration;
+import java.util.Arrays;
+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 java.util.Optional;
+import java.util.Properties;
+import java.util.Set;
+import java.util.stream.Collectors;
+import java.util.stream.StreamSupport;
+
+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.TopicCreationConfig.DEFAULT_TOPIC_CREATION_PREFIX;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.PARTITIONS_CONFIG;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.REPLICATION_FACTOR_CONFIG;
+import static org.apache.kafka.connect.runtime.WorkerConfig.CONNECTOR_CLIENT_POLICY_CLASS_CONFIG;
+import static org.apache.kafka.connect.runtime.WorkerConfig.TOPIC_TRACKING_ALLOW_RESET_CONFIG;
+import static org.apache.kafka.connect.runtime.WorkerConfig.TOPIC_TRACKING_ENABLE_CONFIG;
+import static org.apache.kafka.connect.sink.SinkConnector.TOPICS_CONFIG;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+
+/**
+ * Integration test for the endpoints that offer topic tracking of a connector's active
+ * topics.
+ */
+@Category(IntegrationTest.class)
+public class ConnectorTopicsIntegrationTest {
+
+    private static final int NUM_WORKERS = 5;
+    private static final int NUM_TASKS = 1;
+    private static final String FOO_TOPIC = "foo-topic";
+    private static final String FOO_CONNECTOR = "foo-source";
+    private static final String BAR_TOPIC = "bar-topic";
+    private static final String BAR_CONNECTOR = "bar-source";
+    private static final String SINK_CONNECTOR = "baz-sink";
+    private static final int NUM_TOPIC_PARTITIONS = 3;
+
+    private EmbeddedConnectCluster.Builder connectBuilder;
+    private EmbeddedConnectCluster connect;
+    Map workerProps = new HashMap<>();
+    Properties brokerProps = new Properties();
+
+    @Before
+    public void setup() {
+        // setup Connect worker properties
+        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();
+    }
+
+    @Test
+    public void testGetActiveTopics() throws InterruptedException {
+        connect = connectBuilder.build();
+        // start the clusters
+        connect.start();
+
+        // create test topic
+        connect.kafka().createTopic(FOO_TOPIC, NUM_TOPIC_PARTITIONS);
+        connect.kafka().createTopic(BAR_TOPIC, NUM_TOPIC_PARTITIONS);
+
+        connect.assertions().assertAtLeastNumWorkersAreUp(NUM_WORKERS, "Initial group of workers did not start in time.");
+
+        connect.assertions().assertConnectorActiveTopics(FOO_CONNECTOR, Collections.emptyList(),
+                "Active topic set is not empty for connector: " + FOO_CONNECTOR);
+
+        // start a source connector
+        connect.configureConnector(FOO_CONNECTOR, defaultSourceConnectorProps(FOO_TOPIC));
+
+        connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(FOO_CONNECTOR, NUM_TASKS,
+                "Connector tasks did not start in time.");
+
+        connect.assertions().assertConnectorActiveTopics(FOO_CONNECTOR, Collections.singletonList(FOO_TOPIC),
+                "Active topic set is not: " + Collections.singletonList(FOO_TOPIC) + " for connector: " + FOO_CONNECTOR);
+
+        // start another source connector
+        connect.configureConnector(BAR_CONNECTOR, defaultSourceConnectorProps(BAR_TOPIC));
+
+        connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(BAR_CONNECTOR, NUM_TASKS,
+                "Connector tasks did not start in time.");
+
+        connect.assertions().assertConnectorActiveTopics(BAR_CONNECTOR, Collections.singletonList(BAR_TOPIC),
+                "Active topic set is not: " + Collections.singletonList(BAR_TOPIC) + " for connector: " + BAR_CONNECTOR);
+
+        // start a sink connector
+        connect.configureConnector(SINK_CONNECTOR, defaultSinkConnectorProps(FOO_TOPIC, BAR_TOPIC));
+
+        connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(SINK_CONNECTOR, NUM_TASKS,
+                "Connector tasks did not start in time.");
+
+        connect.assertions().assertConnectorActiveTopics(SINK_CONNECTOR, Arrays.asList(FOO_TOPIC, BAR_TOPIC),
+                "Active topic set is not: " + Arrays.asList(FOO_TOPIC, BAR_TOPIC) + " for connector: " + SINK_CONNECTOR);
+
+        // deleting a connector resets its active topics
+        connect.deleteConnector(BAR_CONNECTOR);
+
+        connect.assertions().assertConnectorAndTasksAreStopped(BAR_CONNECTOR,
+                "Connector tasks did not stop in time.");
+
+        connect.assertions().assertConnectorActiveTopics(BAR_CONNECTOR, Collections.emptyList(),
+                "Active topic set is not empty for deleted connector: " + BAR_CONNECTOR);
+
+        // Unfortunately there's currently no easy way to know when the consumer caught up with
+        // the last records that the producer of the stopped connector managed to produce.
+        // Repeated runs show that this amount of time is sufficient for the consumer to catch up.
+        Thread.sleep(5000);
+
+        // reset active topics for the sink connector after one of the topics has become idle
+        connect.resetConnectorTopics(SINK_CONNECTOR);
+
+        connect.assertions().assertConnectorActiveTopics(SINK_CONNECTOR, Collections.singletonList(FOO_TOPIC),
+                "Active topic set is not: " + Collections.singletonList(FOO_TOPIC) + " for connector: " + SINK_CONNECTOR);
+    }
+
+    @Test
+    public void testTopicTrackingResetIsDisabled() throws InterruptedException {
+        workerProps.put(TOPIC_TRACKING_ALLOW_RESET_CONFIG, "false");
+        connect = connectBuilder.build();
+        // start the clusters
+        connect.start();
+
+        // create test topic
+        connect.kafka().createTopic(FOO_TOPIC, NUM_TOPIC_PARTITIONS);
+        connect.kafka().createTopic(BAR_TOPIC, NUM_TOPIC_PARTITIONS);
+
+        connect.assertions().assertAtLeastNumWorkersAreUp(NUM_WORKERS, "Initial group of workers did not start in time.");
+
+        connect.assertions().assertConnectorActiveTopics(FOO_CONNECTOR, Collections.emptyList(),
+                "Active topic set is not empty for connector: " + FOO_CONNECTOR);
+
+        // start a source connector
+        connect.configureConnector(FOO_CONNECTOR, defaultSourceConnectorProps(FOO_TOPIC));
+
+        connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(FOO_CONNECTOR, NUM_TASKS,
+                "Connector tasks did not start in time.");
+
+        connect.assertions().assertConnectorActiveTopics(FOO_CONNECTOR, Collections.singletonList(FOO_TOPIC),
+                "Active topic set is not: " + Collections.singletonList(FOO_TOPIC) + " for connector: " + FOO_CONNECTOR);
+
+        // start a sink connector
+        connect.configureConnector(SINK_CONNECTOR, defaultSinkConnectorProps(FOO_TOPIC));
+
+        connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(SINK_CONNECTOR, NUM_TASKS,
+                "Connector tasks did not start in time.");
+
+        connect.assertions().assertConnectorActiveTopics(SINK_CONNECTOR, Arrays.asList(FOO_TOPIC),
+                "Active topic set is not: " + Arrays.asList(FOO_TOPIC) + " for connector: " + SINK_CONNECTOR);
+
+        // deleting a connector resets its active topics
+        connect.deleteConnector(FOO_CONNECTOR);
+
+        connect.assertions().assertConnectorAndTasksAreStopped(FOO_CONNECTOR,
+                "Connector tasks did not stop in time.");
+
+        connect.assertions().assertConnectorActiveTopics(FOO_CONNECTOR, Collections.emptyList(),
+                "Active topic set is not empty for deleted connector: " + FOO_CONNECTOR);
+
+        // Unfortunately there's currently no easy way to know when the consumer caught up with
+        // the last records that the producer of the stopped connector managed to produce.
+        // Repeated runs show that this amount of time is sufficient for the consumer to catch up.
+        Thread.sleep(5000);
+
+        // resetting active topics for the sink connector won't work when the config is disabled
+        Exception e = assertThrows(ConnectRestException.class, () -> connect.resetConnectorTopics(SINK_CONNECTOR));
+        assertTrue(e.getMessage().contains("Topic tracking reset is disabled."));
+
+        connect.assertions().assertConnectorActiveTopics(SINK_CONNECTOR, Collections.singletonList(FOO_TOPIC),
+                "Active topic set is not: " + Collections.singletonList(FOO_TOPIC) + " for connector: " + SINK_CONNECTOR);
+    }
+
+    @Test
+    public void testTopicTrackingIsDisabled() throws InterruptedException {
+        workerProps.put(TOPIC_TRACKING_ENABLE_CONFIG, "false");
+        connect = connectBuilder.build();
+        // start the clusters
+        connect.start();
+
+        // create test topic
+        connect.kafka().createTopic(FOO_TOPIC, NUM_TOPIC_PARTITIONS);
+        connect.kafka().createTopic(BAR_TOPIC, NUM_TOPIC_PARTITIONS);
+
+        connect.assertions().assertAtLeastNumWorkersAreUp(NUM_WORKERS, "Initial group of workers did not start in time.");
+
+        // start a source connector
+        connect.configureConnector(FOO_CONNECTOR, defaultSourceConnectorProps(FOO_TOPIC));
+        connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(FOO_CONNECTOR, NUM_TASKS,
+                "Connector tasks did not start in time.");
+
+        // resetting active topics for the sink connector won't work when the config is disabled
+        Exception e = assertThrows(ConnectRestException.class, () -> connect.resetConnectorTopics(SINK_CONNECTOR));
+        assertTrue(e.getMessage().contains("Topic tracking is disabled."));
+
+        e = assertThrows(ConnectRestException.class, () -> connect.connectorTopics(SINK_CONNECTOR));
+        assertTrue(e.getMessage().contains("Topic tracking is disabled."));
+
+        // Wait for tasks to produce a few records
+        Thread.sleep(5000);
+
+        assertNoTopicStatusInStatusTopic();
+    }
+
+    public void assertNoTopicStatusInStatusTopic() {
+        String statusTopic = workerProps.get(DistributedConfig.STATUS_STORAGE_TOPIC_CONFIG);
+        Consumer verifiableConsumer = connect.kafka().createConsumer(
+                Collections.singletonMap("group.id", "verifiable-consumer-group-0"));
+
+        List partitions =
+                Optional.ofNullable(verifiableConsumer.partitionsFor(statusTopic))
+                .orElseThrow(() -> new AssertionError("Unable to retrieve partitions info for status topic"))
+                .stream()
+                .map(info -> new TopicPartition(info.topic(), info.partition()))
+                .collect(Collectors.toList());
+        verifiableConsumer.assign(partitions);
+
+        // Based on the implementation of {@link org.apache.kafka.connect.util.KafkaBasedLog#readToLogEnd}
+        Set assignment = verifiableConsumer.assignment();
+        verifiableConsumer.seekToBeginning(assignment);
+        Map endOffsets = verifiableConsumer.endOffsets(assignment);
+        while (!endOffsets.isEmpty()) {
+            Iterator> it = endOffsets.entrySet().iterator();
+            while (it.hasNext()) {
+                Map.Entry entry = it.next();
+                if (verifiableConsumer.position(entry.getKey()) >= entry.getValue())
+                    it.remove();
+                else {
+                    try {
+                        StreamSupport.stream(verifiableConsumer.poll(Duration.ofMillis(Integer.MAX_VALUE)).spliterator(), false)
+                                .map(ConsumerRecord::key)
+                                .filter(Objects::nonNull)
+                                .filter(key -> new String(key, StandardCharsets.UTF_8).startsWith(KafkaStatusBackingStore.TOPIC_STATUS_PREFIX))
+                                .findFirst()
+                                .ifPresent(key -> {
+                                    throw new AssertionError("Found unexpected key: " + new String(key, StandardCharsets.UTF_8) + " in status topic");
+                                });
+                    } catch (KafkaException e) {
+                        throw new AssertionError("Error while reading to the end of status topic", e);
+                    }
+                    break;
+                }
+            }
+        }
+    }
+
+    private Map defaultSourceConnectorProps(String topic) {
+        // 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(TOPIC_CONFIG, topic);
+        props.put("throughput", String.valueOf(10));
+        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());
+        props.put(DEFAULT_TOPIC_CREATION_PREFIX + REPLICATION_FACTOR_CONFIG, String.valueOf(1));
+        props.put(DEFAULT_TOPIC_CREATION_PREFIX + PARTITIONS_CONFIG, String.valueOf(1));
+        return props;
+    }
+
+    private Map defaultSinkConnectorProps(String... topics) {
+        // setup up props for the sink connector
+        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, String.join(",", topics));
+        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/ErrantRecordSinkConnector.java b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ErrantRecordSinkConnector.java
new file mode 100644
index 0000000000000..0fe2f88083883
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ErrantRecordSinkConnector.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.connect.integration;
+
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.connect.connector.Task;
+import org.apache.kafka.connect.sink.ErrantRecordReporter;
+import org.apache.kafka.connect.sink.SinkRecord;
+
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.Map;
+
+public class ErrantRecordSinkConnector extends MonitorableSinkConnector {
+
+    @Override
+    public Class taskClass() {
+        return ErrantRecordSinkTask.class;
+    }
+
+    public static class ErrantRecordSinkTask extends MonitorableSinkTask {
+        private ErrantRecordReporter reporter;
+
+        public ErrantRecordSinkTask() {
+            super();
+        }
+
+        @Override
+        public void start(Map props) {
+            super.start(props);
+            reporter = context.errantRecordReporter();
+        }
+
+        @Override
+        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);
+                reporter.report(rec, new Throwable());
+            }
+        }
+    }
+}
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
new file mode 100644
index 0000000000000..b6211edc6b39c
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ErrorHandlingIntegrationTest.java
@@ -0,0 +1,328 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.common.header.Headers;
+import org.apache.kafka.connect.connector.ConnectRecord;
+import org.apache.kafka.connect.errors.RetriableException;
+import org.apache.kafka.connect.runtime.rest.entities.ConnectorStateInfo;
+import org.apache.kafka.connect.storage.StringConverter;
+import org.apache.kafka.connect.transforms.Transformation;
+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.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 org.apache.kafka.connect.runtime.ConnectorConfig.CONNECTOR_CLASS_CONFIG;
+import static org.apache.kafka.connect.runtime.ConnectorConfig.ERRORS_LOG_ENABLE_CONFIG;
+import static org.apache.kafka.connect.runtime.ConnectorConfig.ERRORS_LOG_INCLUDE_MESSAGES_CONFIG;
+import static org.apache.kafka.connect.runtime.ConnectorConfig.ERRORS_RETRY_TIMEOUT_CONFIG;
+import static org.apache.kafka.connect.runtime.ConnectorConfig.ERRORS_TOLERANCE_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.TRANSFORMS_CONFIG;
+import static org.apache.kafka.connect.runtime.ConnectorConfig.VALUE_CONVERTER_CLASS_CONFIG;
+import static org.apache.kafka.connect.runtime.SinkConnectorConfig.DLQ_CONTEXT_HEADERS_ENABLE_CONFIG;
+import static org.apache.kafka.connect.runtime.SinkConnectorConfig.DLQ_TOPIC_NAME_CONFIG;
+import static org.apache.kafka.connect.runtime.SinkConnectorConfig.DLQ_TOPIC_REPLICATION_FACTOR_CONFIG;
+import static org.apache.kafka.connect.runtime.SinkConnectorConfig.TOPICS_CONFIG;
+import static org.apache.kafka.connect.runtime.errors.DeadLetterQueueReporter.ERROR_HEADER_EXCEPTION;
+import static org.apache.kafka.connect.runtime.errors.DeadLetterQueueReporter.ERROR_HEADER_EXCEPTION_MESSAGE;
+import static org.apache.kafka.connect.runtime.errors.DeadLetterQueueReporter.ERROR_HEADER_ORIG_TOPIC;
+import static org.apache.kafka.test.TestUtils.waitForCondition;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+
+/**
+ * Integration test for the different error handling policies in Connect (namely, retry policies, skipping bad records,
+ * and dead letter queues).
+ */
+@Category(IntegrationTest.class)
+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";
+    private static final int NUM_RECORDS_PRODUCED = 20;
+    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 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 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);
+    }
+
+    @After
+    public void close() {
+        RuntimeHandles.get().deleteConnector(CONNECTOR_NAME);
+        connect.stop();
+    }
+
+    @Test
+    public void testSkipRetryAndDLQWithHeaders() throws Exception {
+        // create test topic
+        connect.kafka().createTopic("test-topic");
+
+        // setup connector config
+        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());
+        props.put(TRANSFORMS_CONFIG, "failing_transform");
+        props.put("transforms.failing_transform.type", FaultyPassthrough.class.getName());
+
+        // log all errors, along with message metadata
+        props.put(ERRORS_LOG_ENABLE_CONFIG, "true");
+        props.put(ERRORS_LOG_INCLUDE_MESSAGES_CONFIG, "true");
+
+        // produce bad messages into dead letter queue
+        props.put(DLQ_TOPIC_NAME_CONFIG, DLQ_TOPIC);
+        props.put(DLQ_CONTEXT_HEADERS_ENABLE_CONFIG, "true");
+        props.put(DLQ_TOPIC_REPLICATION_FACTOR_CONFIG, "1");
+
+        // tolerate all erros
+        props.put(ERRORS_TOLERANCE_CONFIG, "all");
+
+        // retry for up to one second
+        props.put(ERRORS_RETRY_TIMEOUT_CONFIG, "1000");
+
+        // set expected records to successfully reach the task
+        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,
+                "Connector task was not assigned a partition.");
+
+        // produce some strings into test topic
+        for (int i = 0; i < NUM_RECORDS_PRODUCED; i++) {
+            connect.kafka().produce("test-topic", "key-" + i, "value-" + i);
+        }
+
+        // consume all records from test topic
+        log.info("Consuming records from test topic");
+        int i = 0;
+        for (ConsumerRecord rec : connect.kafka().consume(NUM_RECORDS_PRODUCED, CONSUME_MAX_DURATION_MS, "test-topic")) {
+            String k = new String(rec.key());
+            String v = new String(rec.value());
+            log.debug("Consumed record (key='{}', value='{}') from topic {}", k, v, rec.topic());
+            assertEquals("Unexpected key", k, "key-" + i);
+            assertEquals("Unexpected value", v, "value-" + i);
+            i++;
+        }
+
+        // wait for records to reach the task
+        connectorHandle.taskHandle(TASK_ID).awaitRecords(CONSUME_MAX_DURATION_MS);
+
+        // consume failed records from dead letter queue topic
+        log.info("Consuming records from test topic");
+        ConsumerRecords messages = connect.kafka().consume(EXPECTED_INCORRECT_RECORDS, CONSUME_MAX_DURATION_MS, DLQ_TOPIC);
+        for (ConsumerRecord recs : messages) {
+            log.debug("Consumed record (key={}, value={}) from dead letter queue topic {}",
+                    new String(recs.key()), new String(recs.value()), DLQ_TOPIC);
+            assertTrue(recs.headers().toArray().length > 0);
+            assertValue("test-topic", recs.headers(), ERROR_HEADER_ORIG_TOPIC);
+            assertValue(RetriableException.class.getName(), recs.headers(), ERROR_HEADER_EXCEPTION);
+            assertValue("Error when value='value-7'", recs.headers(), ERROR_HEADER_EXCEPTION_MESSAGE);
+        }
+
+        connect.deleteConnector(CONNECTOR_NAME);
+        connect.assertions().assertConnectorAndTasksAreStopped(CONNECTOR_NAME,
+                "Connector tasks did not stop in time.");
+
+    }
+
+    @Test
+    public void testErrantRecordReporter() throws Exception {
+        // create test topic
+        connect.kafka().createTopic("test-topic");
+
+        // setup connector config
+        Map props = new HashMap<>();
+        props.put(CONNECTOR_CLASS_CONFIG, ErrantRecordSinkConnector.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());
+
+        // log all errors, along with message metadata
+        props.put(ERRORS_LOG_ENABLE_CONFIG, "true");
+        props.put(ERRORS_LOG_INCLUDE_MESSAGES_CONFIG, "true");
+
+        // produce bad messages into dead letter queue
+        props.put(DLQ_TOPIC_NAME_CONFIG, DLQ_TOPIC);
+        props.put(DLQ_CONTEXT_HEADERS_ENABLE_CONFIG, "true");
+        props.put(DLQ_TOPIC_REPLICATION_FACTOR_CONFIG, "1");
+
+        // tolerate all erros
+        props.put(ERRORS_TOLERANCE_CONFIG, "all");
+
+        // retry for up to one second
+        props.put(ERRORS_RETRY_TIMEOUT_CONFIG, "1000");
+
+        // set expected records to successfully reach the task
+        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,
+            "Connector task was not assigned a partition.");
+
+        // produce some strings into test topic
+        for (int i = 0; i < NUM_RECORDS_PRODUCED; i++) {
+            connect.kafka().produce("test-topic", "key-" + i, "value-" + i);
+        }
+
+        // consume all records from test topic
+        log.info("Consuming records from test topic");
+        int i = 0;
+        for (ConsumerRecord rec : connect.kafka().consume(NUM_RECORDS_PRODUCED, CONSUME_MAX_DURATION_MS, "test-topic")) {
+            String k = new String(rec.key());
+            String v = new String(rec.value());
+            log.debug("Consumed record (key='{}', value='{}') from topic {}", k, v, rec.topic());
+            assertEquals("Unexpected key", k, "key-" + i);
+            assertEquals("Unexpected value", v, "value-" + i);
+            i++;
+        }
+
+        // wait for records to reach the task
+        connectorHandle.taskHandle(TASK_ID).awaitRecords(CONSUME_MAX_DURATION_MS);
+
+        // consume failed records from dead letter queue topic
+        log.info("Consuming records from test topic");
+        ConsumerRecords messages = connect.kafka().consume(EXPECTED_INCORRECT_RECORDS, CONSUME_MAX_DURATION_MS, DLQ_TOPIC);
+
+        connect.deleteConnector(CONNECTOR_NAME);
+        connect.assertions().assertConnectorAndTasksAreStopped(CONNECTOR_NAME,
+            "Connector tasks did not stop in time.");
+    }
+
+    /**
+     * Check if a partition was assigned to each task. This method swallows exceptions since it is invoked from a
+     * {@link org.apache.kafka.test.TestUtils#waitForCondition} that will throw an error if this method continued
+     * to return false after the specified duration has elapsed.
+     *
+     * @return true if each task was assigned a partition each, false if this was not true or an error occurred when
+     * executing this operation.
+     */
+    private boolean checkForPartitionAssignment() {
+        try {
+            ConnectorStateInfo info = connect.connectorStatus(CONNECTOR_NAME);
+            return info != null && info.tasks().size() == NUM_TASKS
+                    && connectorHandle.taskHandle(TASK_ID).partitionsAssigned() == 1;
+        }  catch (Exception e) {
+            // Log the exception and return that the partitions were not assigned
+            log.error("Could not check connector state info.", e);
+            return false;
+        }
+    }
+
+    private void assertValue(String expected, Headers headers, String headerKey) {
+        byte[] actual = headers.lastHeader(headerKey).value();
+        if (expected == null && actual == null) {
+            return;
+        }
+        if (expected == null || actual == null) {
+            fail();
+        }
+        assertEquals(expected, new String(actual));
+    }
+
+    public static class FaultyPassthrough> implements Transformation {
+
+        static final ConfigDef CONFIG_DEF = new ConfigDef();
+
+        /**
+         * An arbitrary id which causes this transformation to fail with a {@link RetriableException}, but succeeds
+         * on subsequent attempt.
+         */
+        static final int BAD_RECORD_VAL_RETRIABLE = 4;
+
+        /**
+         * An arbitrary id which causes this transformation to fail with a {@link RetriableException}.
+         */
+        static final int BAD_RECORD_VAL = 7;
+
+        private boolean shouldFail = true;
+
+        @Override
+        public R apply(R record) {
+            String badValRetriable = "value-" + BAD_RECORD_VAL_RETRIABLE;
+            if (badValRetriable.equals(record.value()) && shouldFail) {
+                shouldFail = false;
+                throw new RetriableException("Error when value='" + badValRetriable
+                        + "'. A reattempt with this record will succeed.");
+            }
+            String badVal = "value-" + BAD_RECORD_VAL;
+            if (badVal.equals(record.value())) {
+                throw new RetriableException("Error when value='" + badVal + "'");
+            }
+            return record;
+        }
+
+        @Override
+        public ConfigDef config() {
+            return CONFIG_DEF;
+        }
+
+        @Override
+        public void close() {
+        }
+
+        @Override
+        public void configure(Map configs) {
+        }
+    }
+}
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
new file mode 100644
index 0000000000000..6f8d8a1596783
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ExampleConnectIntegrationTest.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.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.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;
+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.TopicCreationConfig.DEFAULT_TOPIC_CREATION_PREFIX;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.PARTITIONS_CONFIG;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.REPLICATION_FACTOR_CONFIG;
+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.
+ * 
            
+ * The following test configures and executes up a sink connector pipeline in a worker, produces messages into
+ * the source topic-partitions, and demonstrates how to check the overall behavior of the pipeline.
+ */
+@Category(IntegrationTest.class)
+public class ExampleConnectIntegrationTest {
+
+    private static final Logger log = LoggerFactory.getLogger(ExampleConnectIntegrationTest.class);
+
+    private static final int NUM_RECORDS_PRODUCED = 2000;
+    private static final int NUM_TOPIC_PARTITIONS = 3;
+    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 static final String SINK_CONNECTOR_CLASS_NAME = MonitorableSinkConnector.class.getSimpleName();
+    private static final String SOURCE_CONNECTOR_CLASS_NAME = MonitorableSourceConnector.class.getSimpleName();
+
+    private EmbeddedConnectCluster connect;
+    private ConnectorHandle connectorHandle;
+
+    @Rule
+    public TestRule watcher = ConnectIntegrationTestUtils.newTestWatcher(log);
+
+    @Before
+    public void setup() {
+        // setup Connect worker properties
+        Map exampleWorkerProps = new HashMap<>();
+        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 Kafka and Zk
+        connect = new EmbeddedConnectCluster.Builder()
+                .name("connect-cluster")
+                .numWorkers(NUM_WORKERS)
+                .numBrokers(1)
+                .workerProps(exampleWorkerProps)
+                .brokerProps(exampleBrokerProps)
+                .build();
+
+        // start the clusters
+        connect.start();
+
+        // get a handle to the connector
+        connectorHandle = RuntimeHandles.get().connectorHandle(CONNECTOR_NAME);
+    }
+
+    @After
+    public void close() {
+        // delete connector handle
+        RuntimeHandles.get().deleteConnector(CONNECTOR_NAME);
+
+        // stop all Connect, Kafka and Zk threads.
+        connect.stop();
+    }
+
+    /**
+     * 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 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, SINK_CONNECTOR_CLASS_NAME);
+        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());
+
+        // 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);
+
+        // validate the intended connector configuration, a config that errors
+        connect.assertions().assertExactlyNumErrorsOnConnectorConfigValidation(SINK_CONNECTOR_CLASS_NAME, props, 1,
+            "Validating connector configuration produced an unexpected number or errors.");
+
+        // add missing configuration to make the config valid
+        props.put("name", CONNECTOR_NAME);
+
+        // validate the intended connector configuration, a valid config
+        connect.assertions().assertExactlyNumErrorsOnConnectorConfigValidation(SINK_CONNECTOR_CLASS_NAME, props, 0,
+            "Validating connector configuration produced an unexpected number or errors.");
+
+        // start a sink connector
+        connect.configureConnector(CONNECTOR_NAME, props);
+
+        waitForCondition(this::checkForPartitionAssignment,
+                CONNECTOR_SETUP_DURATION_MS,
+                "Connector tasks were not assigned a partition each.");
+
+        // produce some messages into source topic partitions
+        for (int i = 0; i < NUM_RECORDS_PRODUCED; i++) {
+            connect.kafka().produce("test-topic", i % NUM_TOPIC_PARTITIONS, "key", "simple-message-value-" + i);
+        }
+
+        // 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, RECORD_TRANSFER_DURATION_MS, "test-topic").count());
+
+        // wait for the connector tasks to consume all records.
+        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 source connector
+        Map props = new HashMap<>();
+        props.put(CONNECTOR_CLASS_CONFIG, SOURCE_CONNECTOR_CLASS_NAME);
+        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());
+        props.put(DEFAULT_TOPIC_CREATION_PREFIX + REPLICATION_FACTOR_CONFIG, String.valueOf(1));
+        props.put(DEFAULT_TOPIC_CREATION_PREFIX + PARTITIONS_CONFIG, String.valueOf(1));
+
+        // 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);
+
+        // validate the intended connector configuration, a config that errors
+        connect.assertions().assertExactlyNumErrorsOnConnectorConfigValidation(SOURCE_CONNECTOR_CLASS_NAME, props, 1,
+            "Validating connector configuration produced an unexpected number or errors.");
+
+        // add missing configuration to make the config valid
+        props.put("name", CONNECTOR_NAME);
+
+        // validate the intended connector configuration, a valid config
+        connect.assertions().assertExactlyNumErrorsOnConnectorConfigValidation(SOURCE_CONNECTOR_CLASS_NAME, props, 0,
+            "Validating connector configuration produced an unexpected number or errors.");
+
+        // 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);
+    }
+
+    /**
+     * Check if a partition was assigned to each task. This method swallows exceptions since it is invoked from a
+     * {@link org.apache.kafka.test.TestUtils#waitForCondition} that will throw an error if this method continued
+     * to return false after the specified duration has elapsed.
+     *
+     * @return true if each task was assigned a partition each, false if this was not true or an error occurred when
+     * executing this operation.
+     */
+    private boolean checkForPartitionAssignment() {
+        try {
+            ConnectorStateInfo info = connect.connectorStatus(CONNECTOR_NAME);
+            return info != null && info.tasks().size() == NUM_TASKS
+                    && connectorHandle.tasks().stream().allMatch(th -> th.partitionsAssigned() == 1);
+        } catch (Exception e) {
+            // Log the exception and return that the partitions were not assigned
+            log.error("Could not check connector state info.", e);
+            return false;
+        }
+    }
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/InternalTopicsIntegrationTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/InternalTopicsIntegrationTest.java
new file mode 100644
index 0000000000000..d73d1c4ed069e
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/InternalTopicsIntegrationTest.java
@@ -0,0 +1,315 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.HashMap;
+import java.util.Map;
+import java.util.Properties;
+
+import org.apache.kafka.common.config.TopicConfig;
+import org.apache.kafka.connect.runtime.distributed.DistributedConfig;
+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.Test;
+import org.junit.experimental.categories.Category;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import static org.junit.Assert.assertFalse;
+
+/**
+ * Integration test for the creation of internal topics.
+ */
+@Category(IntegrationTest.class)
+public class InternalTopicsIntegrationTest {
+
+    private static final Logger log = LoggerFactory.getLogger(InternalTopicsIntegrationTest.class);
+
+    private EmbeddedConnectCluster connect;
+    Map workerProps = new HashMap<>();
+    Properties brokerProps = new Properties();
+
+    @Before
+    public void setup() {
+        // setup Kafka broker properties
+        brokerProps.put("auto.create.topics.enable", String.valueOf(false));
+    }
+
+    @After
+    public void close() {
+        // stop all Connect, Kafka and Zk threads.
+        connect.stop();
+    }
+
+    @Test
+    public void testCreateInternalTopicsWithDefaultSettings() throws InterruptedException {
+        int numWorkers = 1;
+        int numBrokers = 3;
+        connect = new EmbeddedConnectCluster.Builder().name("connect-cluster-1")
+                                                      .workerProps(workerProps)
+                                                      .numWorkers(numWorkers)
+                                                      .numBrokers(numBrokers)
+                                                      .brokerProps(brokerProps)
+                                                      .build();
+
+        // Start the Connect cluster
+        connect.start();
+        connect.assertions().assertExactlyNumBrokersAreUp(numBrokers, "Brokers did not start in time.");
+        connect.assertions().assertExactlyNumWorkersAreUp(numWorkers, "Worker did not start in time.");
+        log.info("Completed startup of {} Kafka brokers and {} Connect workers", numBrokers, numWorkers);
+
+        // Check the topics
+        log.info("Verifying the internal topics for Connect");
+        connect.assertions().assertTopicsExist(configTopic(), offsetTopic(), statusTopic());
+        assertInternalTopicSettings();
+
+        // Remove the Connect worker
+        log.info("Stopping the Connect worker");
+        connect.removeWorker();
+
+        // And restart
+        log.info("Starting the Connect worker");
+        connect.startConnect();
+
+        // Check the topics
+        log.info("Verifying the internal topics for Connect");
+        connect.assertions().assertTopicsExist(configTopic(), offsetTopic(), statusTopic());
+        assertInternalTopicSettings();
+    }
+
+    @Test
+    public void testCreateInternalTopicsWithFewerReplicasThanBrokers() throws InterruptedException {
+        workerProps.put(DistributedConfig.CONFIG_STORAGE_REPLICATION_FACTOR_CONFIG, "1");
+        workerProps.put(DistributedConfig.OFFSET_STORAGE_REPLICATION_FACTOR_CONFIG, "2");
+        workerProps.put(DistributedConfig.STATUS_STORAGE_REPLICATION_FACTOR_CONFIG, "1");
+        int numWorkers = 1;
+        int numBrokers = 2;
+        connect = new EmbeddedConnectCluster.Builder().name("connect-cluster-1")
+                                                      .workerProps(workerProps)
+                                                      .numWorkers(numWorkers)
+                                                      .numBrokers(numBrokers)
+                                                      .brokerProps(brokerProps)
+                                                      .build();
+
+        // Start the Connect cluster
+        connect.start();
+        connect.assertions().assertExactlyNumBrokersAreUp(numBrokers, "Broker did not start in time.");
+        connect.assertions().assertAtLeastNumWorkersAreUp(numWorkers, "Worker did not start in time.");
+        log.info("Completed startup of {} Kafka brokers and {} Connect workers", numBrokers, numWorkers);
+
+        // Check the topics
+        log.info("Verifying the internal topics for Connect");
+        connect.assertions().assertTopicsExist(configTopic(), offsetTopic(), statusTopic());
+        assertInternalTopicSettings();
+    }
+
+    @Test
+    public void testFailToCreateInternalTopicsWithMoreReplicasThanBrokers() throws InterruptedException {
+        workerProps.put(DistributedConfig.CONFIG_STORAGE_REPLICATION_FACTOR_CONFIG, "3");
+        workerProps.put(DistributedConfig.OFFSET_STORAGE_REPLICATION_FACTOR_CONFIG, "2");
+        workerProps.put(DistributedConfig.STATUS_STORAGE_REPLICATION_FACTOR_CONFIG, "1");
+        int numWorkers = 1;
+        int numBrokers = 1;
+        connect = new EmbeddedConnectCluster.Builder().name("connect-cluster-1")
+                                                      .workerProps(workerProps)
+                                                      .numWorkers(numWorkers)
+                                                      .numBrokers(numBrokers)
+                                                      .brokerProps(brokerProps)
+                                                      .build();
+
+        // Start the brokers and Connect, but Connect should fail to create config and offset topic
+        connect.start();
+        connect.assertions().assertExactlyNumBrokersAreUp(numBrokers, "Broker did not start in time.");
+        log.info("Completed startup of {} Kafka broker. Expected Connect worker to fail", numBrokers);
+
+        // Verify that the offset and config topic don't exist;
+        // the status topic may have been created if timing was right but we don't care
+        log.info("Verifying the internal topics for Connect");
+        connect.assertions().assertTopicsDoNotExist(configTopic(), offsetTopic());
+    }
+
+    @Test
+    public void testFailToStartWhenInternalTopicsAreNotCompacted() throws InterruptedException {
+        // Change the broker default cleanup policy to something Connect doesn't like
+        brokerProps.put("log." + TopicConfig.CLEANUP_POLICY_CONFIG, TopicConfig.CLEANUP_POLICY_DELETE);
+        // Start out using the improperly configured topics
+        workerProps.put(DistributedConfig.CONFIG_TOPIC_CONFIG, "bad-config");
+        workerProps.put(DistributedConfig.OFFSET_STORAGE_TOPIC_CONFIG, "bad-offset");
+        workerProps.put(DistributedConfig.STATUS_STORAGE_TOPIC_CONFIG, "bad-status");
+        workerProps.put(DistributedConfig.CONFIG_STORAGE_REPLICATION_FACTOR_CONFIG, "1");
+        workerProps.put(DistributedConfig.OFFSET_STORAGE_REPLICATION_FACTOR_CONFIG, "1");
+        workerProps.put(DistributedConfig.STATUS_STORAGE_REPLICATION_FACTOR_CONFIG, "1");
+        int numWorkers = 0;
+        int numBrokers = 1;
+        connect = new EmbeddedConnectCluster.Builder().name("connect-cluster-1")
+                                                      .workerProps(workerProps)
+                                                      .numWorkers(numWorkers)
+                                                      .numBrokers(numBrokers)
+                                                      .brokerProps(brokerProps)
+                                                      .build();
+
+        // Start the brokers but not Connect
+        log.info("Starting {} Kafka brokers, but no Connect workers yet", numBrokers);
+        connect.start();
+        connect.assertions().assertExactlyNumBrokersAreUp(numBrokers, "Broker did not start in time.");
+        log.info("Completed startup of {} Kafka broker. Expected Connect worker to fail", numBrokers);
+
+        // Create the good topics
+        connect.kafka().createTopic("good-config", 1, 1, compactCleanupPolicy());
+        connect.kafka().createTopic("good-offset", 1, 1, compactCleanupPolicy());
+        connect.kafka().createTopic("good-status", 1, 1, compactCleanupPolicy());
+
+        // Create the poorly-configured topics
+        connect.kafka().createTopic("bad-config", 1, 1, deleteCleanupPolicy());
+        connect.kafka().createTopic("bad-offset", 1, 1, compactAndDeleteCleanupPolicy());
+        connect.kafka().createTopic("bad-status", 1, 1, noTopicSettings());
+
+        // Check the topics
+        log.info("Verifying the internal topics for Connect were manually created");
+        connect.assertions().assertTopicsExist("good-config", "good-offset", "good-status", "bad-config", "bad-offset", "bad-status");
+
+        // Try to start one worker, with three bad topics
+        WorkerHandle worker = connect.addWorker(); // should have failed to start before returning
+        assertFalse(worker.isRunning());
+        assertFalse(connect.allWorkersRunning());
+        assertFalse(connect.anyWorkersRunning());
+        connect.removeWorker(worker);
+
+        // We rely upon the fact that we can change the worker properties before the workers are started
+        workerProps.put(DistributedConfig.CONFIG_TOPIC_CONFIG, "good-config");
+
+        // Try to start one worker, with two bad topics remaining
+        worker = connect.addWorker(); // should have failed to start before returning
+        assertFalse(worker.isRunning());
+        assertFalse(connect.allWorkersRunning());
+        assertFalse(connect.anyWorkersRunning());
+        connect.removeWorker(worker);
+
+        // We rely upon the fact that we can change the worker properties before the workers are started
+        workerProps.put(DistributedConfig.OFFSET_STORAGE_TOPIC_CONFIG, "good-offset");
+
+        // Try to start one worker, with one bad topic remaining
+        worker = connect.addWorker(); // should have failed to start before returning
+        assertFalse(worker.isRunning());
+        assertFalse(connect.allWorkersRunning());
+        assertFalse(connect.anyWorkersRunning());
+        connect.removeWorker(worker);
+        // We rely upon the fact that we can change the worker properties before the workers are started
+        workerProps.put(DistributedConfig.STATUS_STORAGE_TOPIC_CONFIG, "good-status");
+
+        // Try to start one worker, now using all good internal topics
+        connect.addWorker();
+        connect.assertions().assertAtLeastNumWorkersAreUp(1, "Worker did not start in time.");
+    }
+
+    @Test
+    public void testStartWhenInternalTopicsCreatedManuallyWithCompactForBrokersDefaultCleanupPolicy() throws InterruptedException {
+        // Change the broker default cleanup policy to compact, which is good for Connect
+        brokerProps.put("log." + TopicConfig.CLEANUP_POLICY_CONFIG, TopicConfig.CLEANUP_POLICY_COMPACT);
+        // Start out using the properly configured topics
+        workerProps.put(DistributedConfig.CONFIG_TOPIC_CONFIG, "config-topic");
+        workerProps.put(DistributedConfig.OFFSET_STORAGE_TOPIC_CONFIG, "offset-topic");
+        workerProps.put(DistributedConfig.STATUS_STORAGE_TOPIC_CONFIG, "status-topic");
+        workerProps.put(DistributedConfig.CONFIG_STORAGE_REPLICATION_FACTOR_CONFIG, "1");
+        workerProps.put(DistributedConfig.OFFSET_STORAGE_REPLICATION_FACTOR_CONFIG, "1");
+        workerProps.put(DistributedConfig.STATUS_STORAGE_REPLICATION_FACTOR_CONFIG, "1");
+        int numWorkers = 0;
+        int numBrokers = 1;
+        connect = new EmbeddedConnectCluster.Builder().name("connect-cluster-1")
+                                                      .workerProps(workerProps)
+                                                      .numWorkers(numWorkers)
+                                                      .numBrokers(numBrokers)
+                                                      .brokerProps(brokerProps)
+                                                      .build();
+
+        // Start the brokers but not Connect
+        log.info("Starting {} Kafka brokers, but no Connect workers yet", numBrokers);
+        connect.start();
+        connect.assertions().assertExactlyNumBrokersAreUp(numBrokers, "Broker did not start in time.");
+        log.info("Completed startup of {} Kafka broker. Expected Connect worker to fail", numBrokers);
+
+        // Create the valid internal topics w/o topic settings, so these will use the broker's
+        // broker's log.cleanup.policy=compact (set above)
+        connect.kafka().createTopic("config-topic", 1, 1, noTopicSettings());
+        connect.kafka().createTopic("offset-topic", 1, 1, noTopicSettings());
+        connect.kafka().createTopic("status-topic", 1, 1, noTopicSettings());
+
+        // Check the topics
+        log.info("Verifying the internal topics for Connect were manually created");
+        connect.assertions().assertTopicsExist("config-topic", "offset-topic", "status-topic");
+
+        // Try to start one worker using valid internal topics
+        connect.addWorker();
+        connect.assertions().assertAtLeastNumWorkersAreUp(1, "Worker did not start in time.");
+    }
+
+    protected Map compactCleanupPolicy() {
+        return Collections.singletonMap(TopicConfig.CLEANUP_POLICY_CONFIG, TopicConfig.CLEANUP_POLICY_COMPACT);
+    }
+
+    protected Map deleteCleanupPolicy() {
+        return Collections.singletonMap(TopicConfig.CLEANUP_POLICY_CONFIG, TopicConfig.CLEANUP_POLICY_DELETE);
+    }
+
+    protected Map noTopicSettings() {
+        return Collections.emptyMap();
+    }
+
+    protected Map compactAndDeleteCleanupPolicy() {
+        Map config = new HashMap<>();
+        config.put(TopicConfig.CLEANUP_POLICY_CONFIG, TopicConfig.CLEANUP_POLICY_DELETE + "," + TopicConfig.CLEANUP_POLICY_COMPACT);
+        return config;
+    }
+
+    protected void assertInternalTopicSettings() throws InterruptedException {
+        DistributedConfig config = new DistributedConfig(workerProps);
+        connect.assertions().assertTopicSettings(
+                configTopic(),
+                config.getShort(DistributedConfig.CONFIG_STORAGE_REPLICATION_FACTOR_CONFIG),
+                1,
+                "Config topic does not have the expected settings"
+        );
+        connect.assertions().assertTopicSettings(
+                statusTopic(),
+                config.getShort(DistributedConfig.STATUS_STORAGE_REPLICATION_FACTOR_CONFIG),
+                config.getInt(DistributedConfig.STATUS_STORAGE_PARTITIONS_CONFIG),
+                "Status topic does not have the expected settings"
+        );
+        connect.assertions().assertTopicSettings(
+                offsetTopic(),
+                config.getShort(DistributedConfig.OFFSET_STORAGE_REPLICATION_FACTOR_CONFIG),
+                config.getInt(DistributedConfig.OFFSET_STORAGE_PARTITIONS_CONFIG),
+                "Offset topic does not have the expected settings"
+        );
+    }
+
+    protected String configTopic() {
+        return workerProps.get(DistributedConfig.CONFIG_TOPIC_CONFIG);
+    }
+
+    protected String offsetTopic() {
+        return workerProps.get(DistributedConfig.OFFSET_STORAGE_TOPIC_CONFIG);
+    }
+
+    protected String statusTopic() {
+        return workerProps.get(DistributedConfig.STATUS_STORAGE_TOPIC_CONFIG);
+    }
+}
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
new file mode 100644
index 0000000000000..7b9afa4419008
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/MonitorableSinkConnector.java
@@ -0,0 +1,161 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.consumer.OffsetAndMetadata;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.connect.connector.Task;
+import org.apache.kafka.connect.runtime.TestSinkConnector;
+import org.apache.kafka.connect.sink.SinkRecord;
+import org.apache.kafka.connect.sink.SinkTask;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+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 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.
+ */
+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
+    public Class taskClass() {
+        return MonitorableSinkTask.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() {
+        return new ConfigDef();
+    }
+
+    public static class MonitorableSinkTask extends SinkTask {
+
+        private String connectorName;
+        private String taskId;
+        TaskHandle taskHandle;
+        Set assignments;
+        Map committedOffsets;
+        Map> cachedTopicPartitions;
+
+        public MonitorableSinkTask() {
+            this.assignments = new HashSet<>();
+            this.committedOffsets = new HashMap<>();
+            this.cachedTopicPartitions = new HashMap<>();
+        }
+
+        @Override
+        public String version() {
+            return "unknown";
+        }
+
+        @Override
+        public void start(Map props) {
+            taskId = props.get("task.id");
+            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());
+            assignments.addAll(partitions);
+            taskHandle.partitionsAssigned(partitions.size());
+        }
+
+        @Override
+        public void put(Collection records) {
+            for (SinkRecord rec : records) {
+                taskHandle.record(rec);
+                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() {
+            log.info("Stopped {} task {}", this.getClass().getSimpleName(), taskId);
+            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..c11b9bd79d0fb
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/MonitorableSourceConnector.java
@@ -0,0 +1,169 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.header.ConnectHeaders;
+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,
+                                null,
+                                new ConnectHeaders().addLong("header-" + seqno, 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..2914434362ee5
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/RebalanceSourceConnectorsIntegrationTest.java
@@ -0,0 +1,316 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.TopicCreationConfig.DEFAULT_TOPIC_CREATION_PREFIX;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.PARTITIONS_CONFIG;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.REPLICATION_FACTOR_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 = defaultSourceConnectorProps(TOPIC_NAME);
+
+        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 = defaultSourceConnectorProps(TOPIC_NAME);
+
+        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 = defaultSourceConnectorProps(TOPIC_NAME);
+
+        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 = defaultSourceConnectorProps(TOPIC_NAME);
+
+        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 = defaultSourceConnectorProps(TOPIC_NAME);
+
+        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 Map defaultSourceConnectorProps(String topic) {
+        // 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(TOPIC_CONFIG, topic);
+        props.put("throughput", String.valueOf(10));
+        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());
+        props.put(DEFAULT_TOPIC_CREATION_PREFIX + REPLICATION_FACTOR_CONFIG, String.valueOf(1));
+        props.put(DEFAULT_TOPIC_CREATION_PREFIX + PARTITIONS_CONFIG, String.valueOf(1));
+        return props;
+    }
+
+    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/RuntimeHandles.java b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/RuntimeHandles.java
new file mode 100644
index 0000000000000..c9900f3a7fb62
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/RuntimeHandles.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.connect.integration;
+
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+/**
+ * A singleton class which provides a shared class for {@link ConnectorHandle}s and {@link TaskHandle}s that are
+ * required for integration tests.
+ */
+public class RuntimeHandles {
+
+    private static final RuntimeHandles INSTANCE = new RuntimeHandles();
+
+    private final Map connectorHandles = new ConcurrentHashMap<>();
+
+    private RuntimeHandles() {
+    }
+
+    /**
+     * @return the shared {@link RuntimeHandles} instance.
+     */
+    public static RuntimeHandles get() {
+        return INSTANCE;
+    }
+
+    /**
+     * Get or create a connector handle for a given connector name. The connector need not be running at the time
+     * this method is called. Once the connector is created, it will bind to this handle. Binding happens with the
+     * connectorName.
+     *
+     * @param connectorName the name of the connector
+     * @return a non-null {@link ConnectorHandle}
+     */
+    public ConnectorHandle connectorHandle(String connectorName) {
+        return connectorHandles.computeIfAbsent(connectorName, k -> new ConnectorHandle(connectorName));
+    }
+
+    /**
+     * Delete the connector handle for this connector name.
+     *
+     * @param connectorName name of the connector
+     */
+    public void deleteConnector(String connectorName) {
+        connectorHandles.remove(connectorName);
+    }
+
+}
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/SourceConnectorsIntegrationTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/SourceConnectorsIntegrationTest.java
new file mode 100644
index 0000000000000..b35b072080257
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/SourceConnectorsIntegrationTest.java
@@ -0,0 +1,239 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.SourceConnectorConfig;
+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.Test;
+import org.junit.experimental.categories.Category;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Properties;
+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.NAME_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.SourceConnectorConfig.TOPIC_CREATION_GROUPS_CONFIG;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.DEFAULT_TOPIC_CREATION_PREFIX;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.EXCLUDE_REGEX_CONFIG;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.INCLUDE_REGEX_CONFIG;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.PARTITIONS_CONFIG;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.REPLICATION_FACTOR_CONFIG;
+import static org.apache.kafka.connect.runtime.WorkerConfig.CONNECTOR_CLIENT_POLICY_CLASS_CONFIG;
+import static org.apache.kafka.connect.runtime.WorkerConfig.TOPIC_CREATION_ENABLE_CONFIG;
+import static org.apache.kafka.connect.util.clusters.EmbeddedConnectCluster.DEFAULT_NUM_BROKERS;
+
+/**
+ * Integration test for source connectors with a focus on topic creation with custom properties by
+ * the connector tasks.
+ */
+@Category(IntegrationTest.class)
+public class SourceConnectorsIntegrationTest {
+
+    private static final int NUM_WORKERS = 3;
+    private static final int NUM_TASKS = 1;
+    private static final String FOO_TOPIC = "foo-topic";
+    private static final String FOO_CONNECTOR = "foo-source";
+    private static final String BAR_TOPIC = "bar-topic";
+    private static final String BAR_CONNECTOR = "bar-source";
+    private static final String FOO_GROUP = "foo";
+    private static final String BAR_GROUP = "bar";
+    private static final int DEFAULT_REPLICATION_FACTOR = DEFAULT_NUM_BROKERS;
+    private static final int DEFAULT_PARTITIONS = 1;
+    private static final int FOO_GROUP_REPLICATION_FACTOR = DEFAULT_NUM_BROKERS;
+    private static final int FOO_GROUP_PARTITIONS = 9;
+
+    private EmbeddedConnectCluster.Builder connectBuilder;
+    private EmbeddedConnectCluster connect;
+    Map workerProps = new HashMap<>();
+    Properties brokerProps = new Properties();
+
+    @Before
+    public void setup() {
+        // setup Connect worker properties
+        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();
+    }
+
+    @Test
+    public void testTopicsAreCreatedWhenAutoCreateTopicsIsEnabledAtTheBroker() throws InterruptedException {
+        brokerProps.put("auto.create.topics.enable", String.valueOf(true));
+        workerProps.put(TOPIC_CREATION_ENABLE_CONFIG, String.valueOf(false));
+        connect = connectBuilder.brokerProps(brokerProps).workerProps(workerProps).build();
+        // start the clusters
+        connect.start();
+
+        connect.assertions().assertAtLeastNumWorkersAreUp(NUM_WORKERS, "Initial group of workers did not start in time.");
+
+        Map fooProps = sourceConnectorPropsWithGroups(FOO_TOPIC);
+
+        // start a source connector
+        connect.configureConnector(FOO_CONNECTOR, fooProps);
+        fooProps.put(NAME_CONFIG, FOO_CONNECTOR);
+
+        connect.assertions().assertExactlyNumErrorsOnConnectorConfigValidation(fooProps.get(CONNECTOR_CLASS_CONFIG), fooProps, 0,
+                "Validating connector configuration produced an unexpected number or errors.");
+
+        connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(FOO_CONNECTOR, NUM_TASKS,
+                "Connector tasks did not start in time.");
+
+        connect.assertions().assertTopicsExist(FOO_TOPIC);
+        connect.assertions().assertTopicSettings(FOO_TOPIC, DEFAULT_REPLICATION_FACTOR,
+                DEFAULT_PARTITIONS, "Topic " + FOO_TOPIC + " does not have the expected settings");
+    }
+
+    @Test
+    public void testTopicsAreCreatedWhenTopicCreationIsEnabled() throws InterruptedException {
+        connect = connectBuilder.build();
+        // start the clusters
+        connect.start();
+
+        connect.assertions().assertAtLeastNumWorkersAreUp(NUM_WORKERS, "Initial group of workers did not start in time.");
+
+        Map fooProps = sourceConnectorPropsWithGroups(FOO_TOPIC);
+
+        // start a source connector
+        connect.configureConnector(FOO_CONNECTOR, fooProps);
+        fooProps.put(NAME_CONFIG, FOO_CONNECTOR);
+
+        connect.assertions().assertExactlyNumErrorsOnConnectorConfigValidation(fooProps.get(CONNECTOR_CLASS_CONFIG), fooProps, 0,
+                "Validating connector configuration produced an unexpected number or errors.");
+
+        connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(FOO_CONNECTOR, NUM_TASKS,
+                "Connector tasks did not start in time.");
+
+        connect.assertions().assertTopicsExist(FOO_TOPIC);
+        connect.assertions().assertTopicSettings(FOO_TOPIC, FOO_GROUP_REPLICATION_FACTOR,
+                FOO_GROUP_PARTITIONS, "Topic " + FOO_TOPIC + " does not have the expected settings");
+    }
+
+    @Test
+    public void testSwitchingToTopicCreationEnabled() throws InterruptedException {
+        workerProps.put(TOPIC_CREATION_ENABLE_CONFIG, String.valueOf(false));
+        connect = connectBuilder.build();
+        // start the clusters
+        connect.start();
+
+        connect.kafka().createTopic(BAR_TOPIC, DEFAULT_PARTITIONS, DEFAULT_REPLICATION_FACTOR, Collections.emptyMap());
+
+        connect.assertions().assertTopicsExist(BAR_TOPIC);
+        connect.assertions().assertTopicSettings(BAR_TOPIC, DEFAULT_REPLICATION_FACTOR,
+                DEFAULT_PARTITIONS, "Topic " + BAR_TOPIC + " does not have the expected settings");
+
+        connect.assertions().assertAtLeastNumWorkersAreUp(NUM_WORKERS, "Initial group of workers did not start in time.");
+
+        Map barProps = defaultSourceConnectorProps(BAR_TOPIC);
+        // start a source connector with topic creation properties
+        connect.configureConnector(BAR_CONNECTOR, barProps);
+        barProps.put(NAME_CONFIG, BAR_CONNECTOR);
+
+        Map fooProps = sourceConnectorPropsWithGroups(FOO_TOPIC);
+        // start a source connector without topic creation properties
+        connect.configureConnector(FOO_CONNECTOR, fooProps);
+        fooProps.put(NAME_CONFIG, FOO_CONNECTOR);
+
+        connect.assertions().assertExactlyNumErrorsOnConnectorConfigValidation(fooProps.get(CONNECTOR_CLASS_CONFIG), fooProps, 0,
+                "Validating connector configuration produced an unexpected number or errors.");
+
+        connect.assertions().assertExactlyNumErrorsOnConnectorConfigValidation(barProps.get(CONNECTOR_CLASS_CONFIG), barProps, 0,
+                "Validating connector configuration produced an unexpected number or errors.");
+
+        connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(FOO_CONNECTOR, NUM_TASKS,
+                "Connector tasks did not start in time.");
+
+        connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(BAR_CONNECTOR, NUM_TASKS,
+                "Connector tasks did not start in time.");
+
+        connect.assertions().assertTopicsExist(BAR_TOPIC);
+        connect.assertions().assertTopicSettings(BAR_TOPIC, DEFAULT_REPLICATION_FACTOR,
+                DEFAULT_PARTITIONS, "Topic " + BAR_TOPIC + " does not have the expected settings");
+
+        connect.assertions().assertTopicsDoNotExist(FOO_TOPIC);
+
+        connect.activeWorkers().forEach(w -> connect.removeWorker(w));
+
+        workerProps.put(TOPIC_CREATION_ENABLE_CONFIG, String.valueOf(true));
+
+        IntStream.range(0, 3).forEach(i -> connect.addWorker());
+        connect.assertions().assertAtLeastNumWorkersAreUp(NUM_WORKERS, "Initial group of workers did not start in time.");
+
+        connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(FOO_CONNECTOR, NUM_TASKS,
+                "Connector tasks did not start in time.");
+
+        connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(BAR_CONNECTOR, NUM_TASKS,
+                "Connector tasks did not start in time.");
+
+        connect.assertions().assertTopicsExist(FOO_TOPIC);
+        connect.assertions().assertTopicSettings(FOO_TOPIC, FOO_GROUP_REPLICATION_FACTOR,
+                FOO_GROUP_PARTITIONS, "Topic " + FOO_TOPIC + " does not have the expected settings");
+        connect.assertions().assertTopicsExist(BAR_TOPIC);
+        connect.assertions().assertTopicSettings(BAR_TOPIC, DEFAULT_REPLICATION_FACTOR,
+                DEFAULT_PARTITIONS, "Topic " + BAR_TOPIC + " does not have the expected settings");
+    }
+
+    private Map defaultSourceConnectorProps(String topic) {
+        // 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(TOPIC_CONFIG, topic);
+        props.put("throughput", String.valueOf(10));
+        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());
+        return props;
+    }
+
+    private Map sourceConnectorPropsWithGroups(String topic) {
+        // setup up props for the source connector
+        Map props = defaultSourceConnectorProps(topic);
+        props.put(TOPIC_CREATION_GROUPS_CONFIG, String.join(",", FOO_GROUP, BAR_GROUP));
+        props.put(DEFAULT_TOPIC_CREATION_PREFIX + REPLICATION_FACTOR_CONFIG, String.valueOf(DEFAULT_REPLICATION_FACTOR));
+        props.put(DEFAULT_TOPIC_CREATION_PREFIX + PARTITIONS_CONFIG, String.valueOf(DEFAULT_PARTITIONS));
+        props.put(SourceConnectorConfig.TOPIC_CREATION_PREFIX + FOO_GROUP + "." + INCLUDE_REGEX_CONFIG, FOO_TOPIC);
+        props.put(SourceConnectorConfig.TOPIC_CREATION_PREFIX + FOO_GROUP + "." + EXCLUDE_REGEX_CONFIG, BAR_TOPIC);
+        props.put(SourceConnectorConfig.TOPIC_CREATION_PREFIX + FOO_GROUP + "." + PARTITIONS_CONFIG,
+                String.valueOf(FOO_GROUP_PARTITIONS));
+        return props;
+    }
+}
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
new file mode 100644
index 0000000000000..b63c08bf36811
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/TaskHandle.java
@@ -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 org.apache.kafka.connect.integration;
+
+import org.apache.kafka.connect.errors.DataException;
+import org.apache.kafka.connect.sink.SinkRecord;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Consumer;
+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
+ * by the task using so far, or waiting for partitions to be assigned to the task.
+ */
+public class TaskHandle {
+
+    private static final Logger log = LoggerFactory.getLogger(TaskHandle.class);
+
+    private final String taskId;
+    private final ConnectorHandle connectorHandle;
+    private final AtomicInteger partitionsAssigned = new AtomicInteger(0);
+    private final StartAndStopCounter startAndStopCounter = new StartAndStopCounter();
+    private final Consumer consumer;
+
+    private CountDownLatch recordsRemainingLatch;
+    private CountDownLatch recordsToCommitLatch;
+    private int expectedRecords = -1;
+    private int expectedCommits = -1;
+
+    public TaskHandle(ConnectorHandle connectorHandle, String taskId, Consumer consumer) {
+        this.taskId = taskId;
+        this.connectorHandle = connectorHandle;
+        this.consumer = consumer;
+    }
+
+    public void record() {
+        record(null);
+    }
+
+    /**
+     * Record a message arrival at the task and the connector overall.
+     */
+    public void record(SinkRecord record) {
+        if (consumer != null && record != null) {
+            consumer.accept(record);
+        }
+        if (recordsRemainingLatch != null) {
+            recordsRemainingLatch.countDown();
+        }
+        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 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 expectedCommits(int expected) {
+        expectedRecords = expected;
+        recordsToCommitLatch = new CountDownLatch(expected);
+    }
+
+    /**
+     * Set the number of partitions assigned to this task.
+     *
+     * @param numPartitions number of partitions
+     */
+    public void partitionsAssigned(int numPartitions) {
+        partitionsAssigned.set(numPartitions);
+    }
+
+    /**
+     * @return the number of topic partitions assigned to this task.
+     */
+    public int partitionsAssigned() {
+        return partitionsAssigned.get();
+    }
+
+    /**
+     * Wait up to the specified number of milliseconds for this task to meet the expected number of
+     * records as defined by {@code expectedRecords}.
+     *
+     * @param timeoutMillis number of milliseconds to wait for records
+     * @throws InterruptedException if another threads interrupts this one while waiting for records
+     */
+    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(timeout, unit)) {
+            String msg = String.format(
+                    "Insufficient records seen by task %s in %d millis. Records expected=%d, actual=%d",
+                    taskId,
+                    unit.toMillis(timeout),
+                    expectedRecords,
+                    expectedRecords - recordsRemainingLatch.getCount());
+            throw new DataException(msg);
+        }
+        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
+    public String toString() {
+        return "Handle{" +
+                "taskId='" + taskId + '\'' +
+                '}';
+    }
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/TransformationIntegrationTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/TransformationIntegrationTest.java
new file mode 100644
index 0000000000000..379172ad162a1
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/TransformationIntegrationTest.java
@@ -0,0 +1,327 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.consumer.ConsumerRecord;
+import org.apache.kafka.connect.storage.StringConverter;
+import org.apache.kafka.connect.transforms.Filter;
+import org.apache.kafka.connect.transforms.predicates.HasHeaderKey;
+import org.apache.kafka.connect.transforms.predicates.RecordIsTombstone;
+import org.apache.kafka.connect.transforms.predicates.TopicNameMatches;
+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.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.Properties;
+import java.util.concurrent.TimeUnit;
+
+import static java.util.Collections.singletonMap;
+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.PREDICATES_CONFIG;
+import static org.apache.kafka.connect.runtime.ConnectorConfig.TASKS_MAX_CONFIG;
+import static org.apache.kafka.connect.runtime.ConnectorConfig.TRANSFORMS_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.TopicCreationConfig.DEFAULT_TOPIC_CREATION_PREFIX;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.PARTITIONS_CONFIG;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.REPLICATION_FACTOR_CONFIG;
+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.assertNotNull;
+
+/**
+ * An integration test for connectors with transformations
+ */
+@Category(IntegrationTest.class)
+public class TransformationIntegrationTest {
+
+    private static final Logger log = LoggerFactory.getLogger(TransformationIntegrationTest.class);
+
+    private static final int NUM_RECORDS_PRODUCED = 2000;
+    private static final int NUM_TOPIC_PARTITIONS = 3;
+    private static final long RECORD_TRANSFER_DURATION_MS = TimeUnit.SECONDS.toMillis(30);
+    private static final long OBSERVED_RECORDS_DURATION_MS = TimeUnit.SECONDS.toMillis(60);
+    private static final int NUM_TASKS = 1;
+    private static final int NUM_WORKERS = 3;
+    private static final String CONNECTOR_NAME = "simple-conn";
+    private static final String SINK_CONNECTOR_CLASS_NAME = MonitorableSinkConnector.class.getSimpleName();
+    private static final String SOURCE_CONNECTOR_CLASS_NAME = MonitorableSourceConnector.class.getSimpleName();
+
+    private EmbeddedConnectCluster connect;
+    private ConnectorHandle connectorHandle;
+
+    @Before
+    public void setup() {
+        // setup Connect worker properties
+        Map workerProps = new HashMap<>();
+        workerProps.put(OFFSET_COMMIT_INTERVAL_MS_CONFIG, String.valueOf(5_000));
+
+        // setup Kafka broker properties
+        Properties brokerProps = new Properties();
+        // This is required because tests in this class also test per-connector topic creation with transformations
+        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();
+
+        // get a handle to the connector
+        connectorHandle = RuntimeHandles.get().connectorHandle(CONNECTOR_NAME);
+    }
+
+    @After
+    public void close() {
+        // delete connector handle
+        RuntimeHandles.get().deleteConnector(CONNECTOR_NAME);
+
+        // stop all Connect, Kafka and Zk threads.
+        connect.stop();
+    }
+
+    /**
+     * Test the {@link Filter} transformer with a
+     * {@link TopicNameMatches} predicate on a sink connector.
+     */
+    @Test
+    public void testFilterOnTopicNameWithSinkConnector() throws Exception {
+        assertConnectReady();
+
+        Map observedRecords = observeRecords();
+
+        // create test topics
+        String fooTopic = "foo-topic";
+        String barTopic = "bar-topic";
+        int numFooRecords = NUM_RECORDS_PRODUCED;
+        int numBarRecords = NUM_RECORDS_PRODUCED;
+        connect.kafka().createTopic(fooTopic, NUM_TOPIC_PARTITIONS);
+        connect.kafka().createTopic(barTopic, NUM_TOPIC_PARTITIONS);
+
+        // setup up props for the sink connector
+        Map props = new HashMap<>();
+        props.put("name", CONNECTOR_NAME);
+        props.put(CONNECTOR_CLASS_CONFIG, SINK_CONNECTOR_CLASS_NAME);
+        props.put(TASKS_MAX_CONFIG, String.valueOf(NUM_TASKS));
+        props.put(TOPICS_CONFIG, String.join(",", fooTopic, barTopic));
+        props.put(KEY_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
+        props.put(VALUE_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
+        props.put(TRANSFORMS_CONFIG, "filter");
+        props.put(TRANSFORMS_CONFIG + ".filter.type", Filter.class.getSimpleName());
+        props.put(TRANSFORMS_CONFIG + ".filter.predicate", "barPredicate");
+        props.put(PREDICATES_CONFIG, "barPredicate");
+        props.put(PREDICATES_CONFIG + ".barPredicate.type", TopicNameMatches.class.getSimpleName());
+        props.put(PREDICATES_CONFIG + ".barPredicate.pattern", "bar-.*");
+
+        // expect all records to be consumed by the connector
+        connectorHandle.expectedRecords(numFooRecords);
+
+        // expect all records to be consumed by the connector
+        connectorHandle.expectedCommits(numFooRecords);
+
+        // start a sink connector
+        connect.configureConnector(CONNECTOR_NAME, props);
+        assertConnectorRunning();
+
+        // produce some messages into source topic partitions
+        for (int i = 0; i < numBarRecords; i++) {
+            connect.kafka().produce(barTopic, i % NUM_TOPIC_PARTITIONS, "key", "simple-message-value-" + i);
+        }
+        for (int i = 0; i < numFooRecords; i++) {
+            connect.kafka().produce(fooTopic, i % NUM_TOPIC_PARTITIONS, "key", "simple-message-value-" + i);
+        }
+
+        // consume all records from the source topic or fail, to ensure that they were correctly produced.
+        assertEquals("Unexpected number of records consumed", numFooRecords,
+                connect.kafka().consume(numFooRecords, RECORD_TRANSFER_DURATION_MS, fooTopic).count());
+        assertEquals("Unexpected number of records consumed", numBarRecords,
+                connect.kafka().consume(numBarRecords, RECORD_TRANSFER_DURATION_MS, barTopic).count());
+
+        // wait for the connector tasks to consume all records.
+        connectorHandle.awaitRecords(RECORD_TRANSFER_DURATION_MS);
+
+        // wait for the connector tasks to commit all records.
+        connectorHandle.awaitCommits(RECORD_TRANSFER_DURATION_MS);
+
+        // Assert that we didn't see any baz
+        Map expectedRecordCounts = singletonMap(fooTopic, Long.valueOf(numFooRecords));
+        assertObservedRecords(observedRecords, expectedRecordCounts);
+
+        // delete connector
+        connect.deleteConnector(CONNECTOR_NAME);
+    }
+
+    private void assertConnectReady() throws InterruptedException {
+        connect.assertions().assertExactlyNumBrokersAreUp(1, "Brokers did not start in time.");
+        connect.assertions().assertExactlyNumWorkersAreUp(NUM_WORKERS, "Worker did not start in time.");
+        log.info("Completed startup of {} Kafka brokers and {} Connect workers", 1, NUM_WORKERS);
+    }
+
+    private void assertConnectorRunning() throws InterruptedException {
+        connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(CONNECTOR_NAME, NUM_TASKS,
+                "Connector tasks did not start in time.");
+    }
+
+    private void assertObservedRecords(Map observedRecords, Map expectedRecordCounts) throws InterruptedException {
+        waitForCondition(() -> expectedRecordCounts.equals(observedRecords),
+            OBSERVED_RECORDS_DURATION_MS,
+            () -> "The observed records should be " + expectedRecordCounts + " but was " + observedRecords);
+    }
+
+    private Map observeRecords() {
+        Map observedRecords = new HashMap<>();
+        // record all the record we see
+        connectorHandle.taskHandle(CONNECTOR_NAME + "-0",
+            record -> observedRecords.compute(record.topic(),
+                (key, value) -> value == null ? 1 : value + 1));
+        return observedRecords;
+    }
+
+    /**
+     * Test the {@link Filter} transformer with a
+     * {@link RecordIsTombstone} predicate on a sink connector.
+     */
+    @Test
+    public void testFilterOnTombstonesWithSinkConnector() throws Exception {
+        assertConnectReady();
+
+        Map observedRecords = observeRecords();
+
+        // create test topics
+        String topic = "foo-topic";
+        int numRecords = NUM_RECORDS_PRODUCED;
+        connect.kafka().createTopic(topic, NUM_TOPIC_PARTITIONS);
+
+        // setup up props for the sink connector
+        Map props = new HashMap<>();
+        props.put("name", CONNECTOR_NAME);
+        props.put(CONNECTOR_CLASS_CONFIG, SINK_CONNECTOR_CLASS_NAME);
+        props.put(TASKS_MAX_CONFIG, String.valueOf(NUM_TASKS));
+        props.put(TOPICS_CONFIG, String.join(",", topic));
+        props.put(KEY_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
+        props.put(VALUE_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
+        props.put(TRANSFORMS_CONFIG, "filter");
+        props.put(TRANSFORMS_CONFIG + ".filter.type", Filter.class.getSimpleName());
+        props.put(TRANSFORMS_CONFIG + ".filter.predicate", "barPredicate");
+        props.put(PREDICATES_CONFIG, "barPredicate");
+        props.put(PREDICATES_CONFIG + ".barPredicate.type", RecordIsTombstone.class.getSimpleName());
+
+        // expect only half the records to be consumed by the connector
+        connectorHandle.expectedCommits(numRecords);
+        connectorHandle.expectedRecords(numRecords / 2);
+
+        // start a sink connector
+        connect.configureConnector(CONNECTOR_NAME, props);
+        assertConnectorRunning();
+
+        // produce some messages into source topic partitions
+        for (int i = 0; i < numRecords; i++) {
+            connect.kafka().produce(topic, i % NUM_TOPIC_PARTITIONS, "key", i % 2 == 0 ? "simple-message-value-" + i : null);
+        }
+
+        // consume all records from the source topic or fail, to ensure that they were correctly produced.
+        assertEquals("Unexpected number of records consumed", numRecords,
+                connect.kafka().consume(numRecords, RECORD_TRANSFER_DURATION_MS, topic).count());
+
+        // wait for the connector tasks to consume all records.
+        connectorHandle.awaitRecords(RECORD_TRANSFER_DURATION_MS);
+
+        // wait for the connector tasks to commit all records.
+        connectorHandle.awaitCommits(RECORD_TRANSFER_DURATION_MS);
+
+        Map expectedRecordCounts = singletonMap(topic, Long.valueOf(numRecords / 2));
+        assertObservedRecords(observedRecords, expectedRecordCounts);
+
+        // delete connector
+        connect.deleteConnector(CONNECTOR_NAME);
+    }
+
+    /**
+     * Test the {@link Filter} transformer with a {@link HasHeaderKey} predicate on a source connector.
+     * Note that this test uses topic creation configs to allow the source connector to create
+     * the topic when it tries to produce the first source record, instead of requiring the topic
+     * to exist before the connector starts.
+     */
+    @Test
+    public void testFilterOnHasHeaderKeyWithSourceConnectorAndTopicCreation() throws Exception {
+        assertConnectReady();
+
+        // setup up props for the sink connector
+        Map props = new HashMap<>();
+        props.put("name", CONNECTOR_NAME);
+        props.put(CONNECTOR_CLASS_CONFIG, SOURCE_CONNECTOR_CLASS_NAME);
+        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());
+        props.put(TRANSFORMS_CONFIG, "filter");
+        props.put(TRANSFORMS_CONFIG + ".filter.type", Filter.class.getSimpleName());
+        props.put(TRANSFORMS_CONFIG + ".filter.predicate", "headerPredicate");
+        props.put(TRANSFORMS_CONFIG + ".filter.negate", "true");
+        props.put(PREDICATES_CONFIG, "headerPredicate");
+        props.put(PREDICATES_CONFIG + ".headerPredicate.type", HasHeaderKey.class.getSimpleName());
+        props.put(PREDICATES_CONFIG + ".headerPredicate.name", "header-8");
+        // custom topic creation is used, so there's no need to proactively create the test topic
+        props.put(DEFAULT_TOPIC_CREATION_PREFIX + REPLICATION_FACTOR_CONFIG, String.valueOf(-1));
+        props.put(DEFAULT_TOPIC_CREATION_PREFIX + PARTITIONS_CONFIG, String.valueOf(NUM_TOPIC_PARTITIONS));
+
+        // 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);
+
+        // validate the intended connector configuration, a valid config
+        connect.assertions().assertExactlyNumErrorsOnConnectorConfigValidation(SOURCE_CONNECTOR_CLASS_NAME, props, 0,
+                "Validating connector configuration produced an unexpected number or errors.");
+
+        // start a source connector
+        connect.configureConnector(CONNECTOR_NAME, props);
+        assertConnectorRunning();
+
+        // 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
+        for (ConsumerRecord record : connect.kafka().consume(1, RECORD_TRANSFER_DURATION_MS, "test-topic")) {
+            assertNotNull("Expected header to exist",
+                    record.headers().lastHeader("header-8"));
+        }
+
+        // delete connector
+        connect.deleteConnector(CONNECTOR_NAME);
+    }
+}
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 ac9c312823cfa..3943f9d9d914f 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,25 +16,47 @@
  */
 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.ConfigValue;
+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.ConfigValueInfo;
 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.errors.BadRequestException;
+import org.apache.kafka.connect.source.SourceConnector;
+import org.apache.kafka.connect.source.SourceTask;
 import org.apache.kafka.connect.storage.ConfigBackingStore;
 import org.apache.kafka.connect.storage.StatusBackingStore;
 import org.apache.kafka.connect.transforms.Transformation;
+import org.apache.kafka.connect.transforms.predicates.Predicate;
 import org.apache.kafka.connect.util.ConnectorTaskId;
 import org.easymock.Capture;
 import org.easymock.EasyMock;
-import org.easymock.EasyMockSupport;
 import org.easymock.IAnswer;
 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.PrepareForTest;
+import org.powermock.modules.junit4.PowerMockRunner;
 
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
@@ -42,30 +64,144 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
-
+import java.util.function.Function;
+import java.util.stream.Collectors;
+
+import static org.apache.kafka.connect.runtime.AbstractHerder.keysWithVariableValues;
+import static org.junit.Assert.assertNull;
+import static org.powermock.api.easymock.PowerMock.verifyAll;
+import static org.powermock.api.easymock.PowerMock.replayAll;
+import static org.easymock.EasyMock.strictMock;
+import static org.easymock.EasyMock.partialMockBuilder;
 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 AbstractHerderTest extends EasyMockSupport {
-    private final Worker worker = strictMock(Worker.class);
+@RunWith(PowerMockRunner.class)
+@PrepareForTest({AbstractHerder.class})
+public class AbstractHerderTest {
+
+    private static final String CONN1 = "sourceA";
+    private static final ConnectorTaskId TASK0 = new ConnectorTaskId(CONN1, 0);
+    private static final ConnectorTaskId TASK1 = new ConnectorTaskId(CONN1, 1);
+    private static final ConnectorTaskId TASK2 = new ConnectorTaskId(CONN1, 2);
+    private static final Integer MAX_TASKS = 3;
+    private static final Map CONN1_CONFIG = new HashMap<>();
+    private static final String TEST_KEY = "testKey";
+    private static final String TEST_KEY2 = "testKey2";
+    private static final String TEST_KEY3 = "testKey3";
+    private static final String TEST_VAL = "testVal";
+    private static final String TEST_VAL2 = "testVal2";
+    private static final String TEST_REF = "${file:/tmp/somefile.txt:somevar}";
+    private static final String TEST_REF2 = "${file:/tmp/somefile2.txt:somevar2}";
+    private static final String TEST_REF3 = "${file:/tmp/somefile3.txt:somevar3}";
+    static {
+        CONN1_CONFIG.put(ConnectorConfig.NAME_CONFIG, CONN1);
+        CONN1_CONFIG.put(ConnectorConfig.TASKS_MAX_CONFIG, MAX_TASKS.toString());
+        CONN1_CONFIG.put(SinkConnectorConfig.TOPICS_CONFIG, "foo,bar");
+        CONN1_CONFIG.put(ConnectorConfig.CONNECTOR_CLASS_CONFIG, BogusSourceConnector.class.getName());
+        CONN1_CONFIG.put(TEST_KEY, TEST_REF);
+        CONN1_CONFIG.put(TEST_KEY2, TEST_REF2);
+        CONN1_CONFIG.put(TEST_KEY3, TEST_REF3);
+    }
+    private static final Map TASK_CONFIG = new HashMap<>();
+    static {
+        TASK_CONFIG.put(TaskConfig.TASK_CLASS_CONFIG, BogusSourceTask.class.getName());
+        TASK_CONFIG.put(TEST_KEY, TEST_REF);
+    }
+    private static final List> TASK_CONFIGS = new ArrayList<>();
+    static {
+        TASK_CONFIGS.add(TASK_CONFIG);
+        TASK_CONFIGS.add(TASK_CONFIG);
+        TASK_CONFIGS.add(TASK_CONFIG);
+    }
+    private static final HashMap> TASK_CONFIGS_MAP = new HashMap<>();
+    static {
+        TASK_CONFIGS_MAP.put(TASK0, TASK_CONFIG);
+        TASK_CONFIGS_MAP.put(TASK1, TASK_CONFIG);
+        TASK_CONFIGS_MAP.put(TASK2, TASK_CONFIG);
+    }
+    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, null, Collections.singletonMap(CONN1, 3),
+            Collections.singletonMap(CONN1, CONN1_CONFIG), Collections.singletonMap(CONN1, TargetState.STARTED),
+            Collections.emptyMap(), Collections.emptySet());
+
     private final String workerId = "workerId";
+    private final String kafkaClusterId = "I4ZmrWqfT2e-upky_4fdPA";
     private final int generation = 5;
     private final String connector = "connector";
-    private final Plugins plugins = strictMock(Plugins.class);
-    private final ClassLoader classLoader = strictMock(ClassLoader.class);
+    private final ConnectorClientConfigOverridePolicy noneConnectorClientConfigOverridePolicy = new NoneConnectorClientConfigOverridePolicy();
+
+    @MockStrict private Worker worker;
+    @MockStrict private WorkerConfigTransformer transformer;
+    @MockStrict private Plugins plugins;
+    @MockStrict private ClassLoader classLoader;
+    @MockStrict private ConfigBackingStore configStore;
+    @MockStrict private StatusBackingStore statusStore;
 
     @Test
-    public void connectorStatus() {
+    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();
 
-        ConfigBackingStore configStore = strictMock(ConfigBackingStore.class);
-        StatusBackingStore statusStore = strictMock(StatusBackingStore.class);
+        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, StatusBackingStore.class, ConfigBackingStore.class)
-                .withArgs(worker, workerId, 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();
 
@@ -95,7 +231,7 @@ public void connectorStatus() {
         assertEquals("UNASSIGNED", taskState.state());
         assertEquals(workerId, taskState.workerId());
 
-        verifyAll();
+        PowerMock.verifyAll();
     }
 
     @Test
@@ -103,12 +239,10 @@ public void taskStatus() {
         ConnectorTaskId taskId = new ConnectorTaskId("connector", 0);
         String workerId = "workerId";
 
-        ConfigBackingStore configStore = strictMock(ConfigBackingStore.class);
-        StatusBackingStore statusStore = strictMock(StatusBackingStore.class);
-
         AbstractHerder herder = partialMockBuilder(AbstractHerder.class)
-                .withConstructor(Worker.class, String.class, StatusBackingStore.class, ConfigBackingStore.class)
-                .withArgs(worker, workerId, 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();
 
@@ -140,43 +274,92 @@ public TaskStatus answer() throws Throwable {
 
 
     @Test(expected = BadRequestException.class)
-    public void testConfigValidationEmptyConfig() {
-        AbstractHerder herder = createConfigValidationHerder(TestSourceConnector.class);
+    public void testConfigValidationEmptyConfig() throws Throwable {
+        AbstractHerder herder = createConfigValidationHerder(TestSourceConnector.class, noneConnectorClientConfigOverridePolicy);
         replayAll();
 
-        herder.validateConnectorConfig(new HashMap());
+        herder.validateConnectorConfig(Collections.emptyMap(), false);
 
         verifyAll();
     }
 
     @Test()
-    public void testConfigValidationMissingName() {
-        AbstractHerder herder = createConfigValidationHerder(TestSourceConnector.class);
+    public void testConfigValidationMissingName() throws Throwable {
+        AbstractHerder herder = createConfigValidationHerder(TestSourceConnector.class, noneConnectorClientConfigOverridePolicy);
         replayAll();
 
         Map config = Collections.singletonMap(ConnectorConfig.CONNECTOR_CLASS_CONFIG, TestSourceConnector.class.getName());
-        ConfigInfos result = herder.validateConnectorConfig(config);
+        ConfigInfos result = herder.validateConnectorConfig(config, false);
 
         // We expect there to be errors due to the missing name and .... 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());
-        assertEquals(Arrays.asList(ConnectorConfig.COMMON_GROUP, ConnectorConfig.TRANSFORMS_GROUP), result.groups());
+        assertEquals(Arrays.asList(ConnectorConfig.COMMON_GROUP, ConnectorConfig.TRANSFORMS_GROUP,
+                ConnectorConfig.PREDICATES_GROUP, ConnectorConfig.ERROR_GROUP, SourceConnectorConfig.TOPIC_CREATION_GROUP), result.groups());
         assertEquals(2, result.errorCount());
-        // Base connector config has 6 fields, connector's configs add 2
-        assertEquals(8, result.values().size());
+        Map infos = result.values().stream()
+                .collect(Collectors.toMap(info -> info.configKey().name(), Function.identity()));
+        // Base connector config has 14 fields, connector's configs add 2
+        assertEquals(17, infos.size());
         // Missing name should generate an error
-        assertEquals(ConnectorConfig.NAME_CONFIG, result.values().get(0).configValue().name());
-        assertEquals(1, result.values().get(0).configValue().errors().size());
+        assertEquals(ConnectorConfig.NAME_CONFIG,
+                infos.get(ConnectorConfig.NAME_CONFIG).configValue().name());
+        assertEquals(1, infos.get(ConnectorConfig.NAME_CONFIG).configValue().errors().size());
         // "required" config from connector should generate an error
-        assertEquals("required", result.values().get(6).configValue().name());
-        assertEquals(1, result.values().get(6).configValue().errors().size());
+        assertEquals("required", infos.get("required").configValue().name());
+        assertEquals(1, infos.get("required").configValue().errors().size());
+
+        verifyAll();
+    }
+
+    @Test(expected = ConfigException.class)
+    public void testConfigValidationInvalidTopics() throws Throwable {
+        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,topic2");
+        config.put(SinkConnectorConfig.TOPICS_REGEX_CONFIG, "topic.*");
+
+        herder.validateConnectorConfig(config, false);
+
+        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, false);
+
+        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, false);
 
         verifyAll();
     }
 
     @Test()
-    public void testConfigValidationTransformsExtendResults() {
-        AbstractHerder herder = createConfigValidationHerder(TestSourceConnector.class);
+    public void testConfigValidationTransformsExtendResults() throws Throwable {
+        AbstractHerder herder = createConfigValidationHerder(TestSourceConnector.class, noneConnectorClientConfigOverridePolicy);
 
         // 2 transform aliases defined -> 2 plugin lookups
         Set> transformations = new HashSet<>();
@@ -193,7 +376,7 @@ public void testConfigValidationTransformsExtendResults() {
         config.put(ConnectorConfig.TRANSFORMS_CONFIG, "xformA,xformB");
         config.put(ConnectorConfig.TRANSFORMS_CONFIG + ".xformA.type", SampleTransformation.class.getName());
         config.put("required", "value"); // connector required config
-        ConfigInfos result = herder.validateConnectorConfig(config);
+        ConfigInfos result = herder.validateConnectorConfig(config, false);
         assertEquals(herder.connectorTypeForClass(config.get(ConnectorConfig.CONNECTOR_CLASS_CONFIG)), ConnectorType.SOURCE);
 
         // We expect there to be errors due to the missing name and .... Note that these assertions depend heavily on
@@ -203,43 +386,431 @@ public void testConfigValidationTransformsExtendResults() {
         List expectedGroups = Arrays.asList(
                 ConnectorConfig.COMMON_GROUP,
                 ConnectorConfig.TRANSFORMS_GROUP,
+                ConnectorConfig.PREDICATES_GROUP,
+                ConnectorConfig.ERROR_GROUP,
+                SourceConnectorConfig.TOPIC_CREATION_GROUP,
                 "Transforms: xformA",
                 "Transforms: xformB"
         );
         assertEquals(expectedGroups, result.groups());
         assertEquals(2, result.errorCount());
-        // Base connector config has 6 fields, connector's configs add 2, 2 type fields from the transforms, and
-        // 1 from the valid transformation's config
-        assertEquals(11, result.values().size());
+        Map infos = result.values().stream()
+                .collect(Collectors.toMap(info -> info.configKey().name(), Function.identity()));
+        assertEquals(22, infos.size());
         // Should get 2 type fields from the transforms, first adds its own config since it has a valid class
-        assertEquals("transforms.xformA.type", result.values().get(6).configValue().name());
-        assertTrue(result.values().get(6).configValue().errors().isEmpty());
-        assertEquals("transforms.xformA.subconfig", result.values().get(7).configValue().name());
-        assertEquals("transforms.xformB.type", result.values().get(8).configValue().name());
-        assertFalse(result.values().get(8).configValue().errors().isEmpty());
+        assertEquals("transforms.xformA.type",
+                infos.get("transforms.xformA.type").configValue().name());
+        assertTrue(infos.get("transforms.xformA.type").configValue().errors().isEmpty());
+        assertEquals("transforms.xformA.subconfig",
+                infos.get("transforms.xformA.subconfig").configValue().name());
+        assertEquals("transforms.xformB.type", infos.get("transforms.xformB.type").configValue().name());
+        assertFalse(infos.get("transforms.xformB.type").configValue().errors().isEmpty());
 
         verifyAll();
     }
 
-    private AbstractHerder createConfigValidationHerder(Class connectorClass) {
+    @Test()
+    public void testConfigValidationPredicatesExtendResults() {
+        AbstractHerder herder = createConfigValidationHerder(TestSourceConnector.class, noneConnectorClientConfigOverridePolicy);
+
+        // 2 transform aliases defined -> 2 plugin lookups
+        Set> transformations = new HashSet<>();
+        transformations.add(new PluginDesc(SampleTransformation.class, "1.0", classLoader));
+        EasyMock.expect(plugins.transformations()).andReturn(transformations).times(1);
+
+        Set> predicates = new HashSet<>();
+        predicates.add(new PluginDesc(SamplePredicate.class, "1.0", classLoader));
+        EasyMock.expect(plugins.predicates()).andReturn(predicates).times(2);
+
+        replayAll();
+
+        // Define 2 transformations. One has a class defined and so can get embedded configs, the other is missing
+        // class info that should generate an error.
+        Map config = new HashMap<>();
+        config.put(ConnectorConfig.CONNECTOR_CLASS_CONFIG, TestSourceConnector.class.getName());
+        config.put(ConnectorConfig.NAME_CONFIG, "connector-name");
+        config.put(ConnectorConfig.TRANSFORMS_CONFIG, "xformA");
+        config.put(ConnectorConfig.TRANSFORMS_CONFIG + ".xformA.type", SampleTransformation.class.getName());
+        config.put(ConnectorConfig.TRANSFORMS_CONFIG + ".xformA.predicate", "predX");
+        config.put(ConnectorConfig.PREDICATES_CONFIG, "predX,predY");
+        config.put(ConnectorConfig.PREDICATES_CONFIG + ".predX.type", SamplePredicate.class.getName());
+        config.put("required", "value"); // connector required config
+        ConfigInfos result = herder.validateConnectorConfig(config, false);
+        assertEquals(herder.connectorTypeForClass(config.get(ConnectorConfig.CONNECTOR_CLASS_CONFIG)), ConnectorType.SOURCE);
+
+        // We expect there to be errors due to the missing name and .... 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.PREDICATES_GROUP,
+                ConnectorConfig.ERROR_GROUP,
+                SourceConnectorConfig.TOPIC_CREATION_GROUP,
+                "Transforms: xformA",
+                "Predicates: predX",
+                "Predicates: predY"
+        );
+        assertEquals(expectedGroups, result.groups());
+        assertEquals(2, result.errorCount());
+        Map infos = result.values().stream()
+                .collect(Collectors.toMap(info -> info.configKey().name(), Function.identity()));
+        assertEquals(24, infos.size());
+        // Should get 2 type fields from the transforms, first adds its own config since it has a valid class
+        assertEquals("transforms.xformA.type",
+                infos.get("transforms.xformA.type").configValue().name());
+        assertTrue(infos.get("transforms.xformA.type").configValue().errors().isEmpty());
+        assertEquals("transforms.xformA.subconfig",
+                infos.get("transforms.xformA.subconfig").configValue().name());
+        assertEquals("transforms.xformA.predicate",
+                infos.get("transforms.xformA.predicate").configValue().name());
+        assertTrue(infos.get("transforms.xformA.predicate").configValue().errors().isEmpty());
+        assertEquals("transforms.xformA.negate",
+                infos.get("transforms.xformA.negate").configValue().name());
+        assertTrue(infos.get("transforms.xformA.negate").configValue().errors().isEmpty());
+        assertEquals("predicates.predX.type",
+                infos.get("predicates.predX.type").configValue().name());
+        assertEquals("predicates.predX.predconfig",
+                infos.get("predicates.predX.predconfig").configValue().name());
+        assertEquals("predicates.predY.type",
+                infos.get("predicates.predY.type").configValue().name());
+        assertFalse(
+                infos.get("predicates.predY.type").configValue().errors().isEmpty());
+
+        verifyAll();
+    }
+
+    @Test()
+    public void testConfigValidationPrincipalOnlyOverride() throws Throwable {
+        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, false);
+        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.PREDICATES_GROUP,
+            ConnectorConfig.ERROR_GROUP,
+            SourceConnectorConfig.TOPIC_CREATION_GROUP
+        );
+        assertEquals(expectedGroups, result.groups());
+        assertEquals(1, result.errorCount());
+        // Base connector config has 14 fields, connector's configs add 2, and 2 producer overrides
+        assertEquals(19, 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() throws Throwable {
+        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, false);
+        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
+        Map newTaskConfig = new HashMap<>();
+        newTaskConfig.put(TaskConfig.TASK_CLASS_CONFIG, BogusSourceTask.class.getName());
+        newTaskConfig.put(TEST_KEY, TEST_VAL);
+        newTaskConfig.put(TEST_KEY2, TEST_VAL2);
+        List> newTaskConfigs = new ArrayList<>();
+        newTaskConfigs.add(newTaskConfig);
+
+        // The SNAPSHOT has a task config with TEST_KEY and TEST_REF
+        List> reverseTransformed = AbstractHerder.reverseTransform(CONN1, SNAPSHOT, newTaskConfigs);
+        assertEquals(TEST_REF, reverseTransformed.get(0).get(TEST_KEY));
+
+        // The SNAPSHOT has no task configs but does have a connector config with TEST_KEY2 and TEST_REF2
+        reverseTransformed = AbstractHerder.reverseTransform(CONN1, SNAPSHOT_NO_TASKS, newTaskConfigs);
+        assertEquals(TEST_REF2, reverseTransformed.get(0).get(TEST_KEY2));
+
+        // The reverseTransformed result should not have TEST_KEY3 since newTaskConfigs does not have TEST_KEY3
+        reverseTransformed = AbstractHerder.reverseTransform(CONN1, SNAPSHOT_NO_TASKS, newTaskConfigs);
+        assertFalse(reverseTransformed.get(0).containsKey(TEST_KEY3));
+    }
+
+    @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);
+    }
+
+    @Test
+    public void testGenerateResultWithConfigValuesAllUsingConfigKeysAndWithNoErrors() {
+        String name = "com.acme.connector.MyConnector";
+        Map keys = new HashMap<>();
+        addConfigKey(keys, "config.a1", null);
+        addConfigKey(keys, "config.b1", "group B");
+        addConfigKey(keys, "config.b2", "group B");
+        addConfigKey(keys, "config.c1", "group C");
+
+        List groups = Arrays.asList("groupB", "group C");
+        List values = new ArrayList<>();
+        addValue(values, "config.a1", "value.a1");
+        addValue(values, "config.b1", "value.b1");
+        addValue(values, "config.b2", "value.b2");
+        addValue(values, "config.c1", "value.c1");
+
+        ConfigInfos infos = AbstractHerder.generateResult(name, keys, values, groups);
+        assertEquals(name, infos.name());
+        assertEquals(groups, infos.groups());
+        assertEquals(values.size(), infos.values().size());
+        assertEquals(0, infos.errorCount());
+        assertInfoKey(infos, "config.a1", null);
+        assertInfoKey(infos, "config.b1", "group B");
+        assertInfoKey(infos, "config.b2", "group B");
+        assertInfoKey(infos, "config.c1", "group C");
+        assertInfoValue(infos, "config.a1", "value.a1");
+        assertInfoValue(infos, "config.b1", "value.b1");
+        assertInfoValue(infos, "config.b2", "value.b2");
+        assertInfoValue(infos, "config.c1", "value.c1");
+    }
+
+    @Test
+    public void testGenerateResultWithConfigValuesAllUsingConfigKeysAndWithSomeErrors() {
+        String name = "com.acme.connector.MyConnector";
+        Map keys = new HashMap<>();
+        addConfigKey(keys, "config.a1", null);
+        addConfigKey(keys, "config.b1", "group B");
+        addConfigKey(keys, "config.b2", "group B");
+        addConfigKey(keys, "config.c1", "group C");
+
+        List groups = Arrays.asList("groupB", "group C");
+        List values = new ArrayList<>();
+        addValue(values, "config.a1", "value.a1");
+        addValue(values, "config.b1", "value.b1");
+        addValue(values, "config.b2", "value.b2");
+        addValue(values, "config.c1", "value.c1", "error c1");
+
+        ConfigInfos infos = AbstractHerder.generateResult(name, keys, values, groups);
+        assertEquals(name, infos.name());
+        assertEquals(groups, infos.groups());
+        assertEquals(values.size(), infos.values().size());
+        assertEquals(1, infos.errorCount());
+        assertInfoKey(infos, "config.a1", null);
+        assertInfoKey(infos, "config.b1", "group B");
+        assertInfoKey(infos, "config.b2", "group B");
+        assertInfoKey(infos, "config.c1", "group C");
+        assertInfoValue(infos, "config.a1", "value.a1");
+        assertInfoValue(infos, "config.b1", "value.b1");
+        assertInfoValue(infos, "config.b2", "value.b2");
+        assertInfoValue(infos, "config.c1", "value.c1", "error c1");
+    }
+
+    @Test
+    public void testGenerateResultWithConfigValuesMoreThanConfigKeysAndWithSomeErrors() {
+        String name = "com.acme.connector.MyConnector";
+        Map keys = new HashMap<>();
+        addConfigKey(keys, "config.a1", null);
+        addConfigKey(keys, "config.b1", "group B");
+        addConfigKey(keys, "config.b2", "group B");
+        addConfigKey(keys, "config.c1", "group C");
+
+        List groups = Arrays.asList("groupB", "group C");
+        List values = new ArrayList<>();
+        addValue(values, "config.a1", "value.a1");
+        addValue(values, "config.b1", "value.b1");
+        addValue(values, "config.b2", "value.b2");
+        addValue(values, "config.c1", "value.c1", "error c1");
+        addValue(values, "config.extra1", "value.extra1");
+        addValue(values, "config.extra2", "value.extra2", "error extra2");
+
+        ConfigInfos infos = AbstractHerder.generateResult(name, keys, values, groups);
+        assertEquals(name, infos.name());
+        assertEquals(groups, infos.groups());
+        assertEquals(values.size(), infos.values().size());
+        assertEquals(2, infos.errorCount());
+        assertInfoKey(infos, "config.a1", null);
+        assertInfoKey(infos, "config.b1", "group B");
+        assertInfoKey(infos, "config.b2", "group B");
+        assertInfoKey(infos, "config.c1", "group C");
+        assertNoInfoKey(infos, "config.extra1");
+        assertNoInfoKey(infos, "config.extra2");
+        assertInfoValue(infos, "config.a1", "value.a1");
+        assertInfoValue(infos, "config.b1", "value.b1");
+        assertInfoValue(infos, "config.b2", "value.b2");
+        assertInfoValue(infos, "config.c1", "value.c1", "error c1");
+        assertInfoValue(infos, "config.extra1", "value.extra1");
+        assertInfoValue(infos, "config.extra2", "value.extra2", "error extra2");
+    }
+
+    @Test
+    public void testGenerateResultWithConfigValuesWithNoConfigKeysAndWithSomeErrors() {
+        String name = "com.acme.connector.MyConnector";
+        Map keys = new HashMap<>();
+
+        List groups = new ArrayList<>();
+        List values = new ArrayList<>();
+        addValue(values, "config.a1", "value.a1");
+        addValue(values, "config.b1", "value.b1");
+        addValue(values, "config.b2", "value.b2");
+        addValue(values, "config.c1", "value.c1", "error c1");
+        addValue(values, "config.extra1", "value.extra1");
+        addValue(values, "config.extra2", "value.extra2", "error extra2");
+
+        ConfigInfos infos = AbstractHerder.generateResult(name, keys, values, groups);
+        assertEquals(name, infos.name());
+        assertEquals(groups, infos.groups());
+        assertEquals(values.size(), infos.values().size());
+        assertEquals(2, infos.errorCount());
+        assertNoInfoKey(infos, "config.a1");
+        assertNoInfoKey(infos, "config.b1");
+        assertNoInfoKey(infos, "config.b2");
+        assertNoInfoKey(infos, "config.c1");
+        assertNoInfoKey(infos, "config.extra1");
+        assertNoInfoKey(infos, "config.extra2");
+        assertInfoValue(infos, "config.a1", "value.a1");
+        assertInfoValue(infos, "config.b1", "value.b1");
+        assertInfoValue(infos, "config.b2", "value.b2");
+        assertInfoValue(infos, "config.c1", "value.c1", "error c1");
+        assertInfoValue(infos, "config.extra1", "value.extra1");
+        assertInfoValue(infos, "config.extra2", "value.extra2", "error extra2");
+    }
+
+    protected void addConfigKey(Map keys, String name, String group) {
+        keys.put(name, new ConfigDef.ConfigKey(name, ConfigDef.Type.STRING, null, null,
+                ConfigDef.Importance.HIGH, "doc", group, 10,
+                ConfigDef.Width.MEDIUM, "display name", Collections.emptyList(), null, false));
+    }
+
+    protected void addValue(List values, String name, String value, String...errors) {
+        values.add(new ConfigValue(name, value, new ArrayList<>(), Arrays.asList(errors)));
+    }
+
+    protected void assertInfoKey(ConfigInfos infos, String name, String group) {
+        ConfigInfo info = findInfo(infos, name);
+        assertEquals(name, info.configKey().name());
+        assertEquals(group, info.configKey().group());
+    }
+
+    protected void assertNoInfoKey(ConfigInfos infos, String name) {
+        ConfigInfo info = findInfo(infos, name);
+        assertNull(info.configKey());
+    }
+
+    protected void assertInfoValue(ConfigInfos infos, String name, String value, String...errors) {
+        ConfigValueInfo info = findInfo(infos, name).configValue();
+        assertEquals(name, info.name());
+        assertEquals(value, info.value());
+        assertEquals(Arrays.asList(errors), info.errors());
+    }
+
+    protected ConfigInfo findInfo(ConfigInfos infos, String name) {
+        return infos.values()
+                    .stream()
+                    .filter(i -> i.configValue().name().equals(name))
+                    .findFirst()
+                    .orElse(null);
+    }
+
+    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, StatusBackingStore.class, ConfigBackingStore.class)
-                .withArgs(worker, workerId, 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);
 
         // Call to validateConnectorConfig
+        EasyMock.expect(worker.configTransformer()).andReturn(transformer).times(2);
+        final Capture> configCapture = EasyMock.newCapture();
+        EasyMock.expect(transformer.transform(EasyMock.capture(configCapture))).andAnswer(configCapture::getValue);
         EasyMock.expect(worker.getPlugins()).andStubReturn(plugins);
         final Connector connector;
         try {
-            connector = connectorClass.newInstance();
-        } catch (InstantiationException | IllegalAccessException e) {
+            connector = connectorClass.getConstructor().newInstance();
+        } catch (ReflectiveOperationException e) {
             throw new RuntimeException("Couldn't create connector", e);
         }
         EasyMock.expect(plugins.newConnector(connectorClass.getName())).andReturn(connector);
@@ -261,12 +832,47 @@ public R apply(R record) {
         @Override
         public ConfigDef config() {
             return new ConfigDef()
-                    .define("subconfig", ConfigDef.Type.STRING, "default", ConfigDef.Importance.LOW, "docs");
+                           .define("subconfig", ConfigDef.Type.STRING, "default", ConfigDef.Importance.LOW, "docs");
+        }
+
+        @Override
+        public void close() {
+
+        }
+    }
+
+    public static class SamplePredicate> implements Predicate {
+
+        @Override
+        public ConfigDef config() {
+            return new ConfigDef()
+                    .define("predconfig", ConfigDef.Type.STRING, "default", ConfigDef.Importance.LOW, "docs");
+        }
+
+        @Override
+        public boolean test(R record) {
+            return false;
         }
 
         @Override
         public void close() {
 
         }
+
+        @Override
+        public void configure(Map configs) {
+
+        }
+    }
+
+    // We need to use a real class here due to some issue with mocking java.lang.Class
+    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/ConnectMetricsTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/ConnectMetricsTest.java
index 2de7cb6e10775..a5fd53b57ea9f 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/ConnectMetricsTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/ConnectMetricsTest.java
@@ -16,13 +16,18 @@
  */
 package org.apache.kafka.connect.runtime;
 
+import org.apache.kafka.common.MetricName;
+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.utils.MockTime;
 import org.apache.kafka.connect.runtime.ConnectMetrics.MetricGroup;
 import org.apache.kafka.connect.runtime.ConnectMetrics.MetricGroupId;
-import org.apache.kafka.connect.util.MockTime;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 
@@ -32,6 +37,7 @@
 import static org.junit.Assert.assertNotSame;
 import static org.junit.Assert.assertSame;
 
+@SuppressWarnings("deprecation")
 public class ConnectMetricsTest {
 
     private static final Map DEFAULT_WORKER_CONFIG = new HashMap<>();
@@ -47,7 +53,7 @@ public class ConnectMetricsTest {
 
     @Before
     public void setUp() {
-        metrics = new ConnectMetrics("worker1", new WorkerConfig(WorkerConfig.baseConfigDef(), DEFAULT_WORKER_CONFIG), new MockTime());
+        metrics = new ConnectMetrics("worker1", new WorkerConfig(WorkerConfig.baseConfigDef(), DEFAULT_WORKER_CONFIG), new MockTime(), "cluster-1");
     }
 
     @After
@@ -61,19 +67,6 @@ public void testKafkaMetricsNotNull() {
         assertNotNull(metrics.metrics());
     }
 
-    @Test
-    public void testCreatingTags() {
-        Map tags = ConnectMetrics.tags("k1", "v1", "k2", "v2");
-        assertEquals("v1", tags.get("k1"));
-        assertEquals("v2", tags.get("k2"));
-        assertEquals(2, tags.size());
-    }
-
-    @Test(expected = IllegalArgumentException.class)
-    public void testCreatingTagsWithOddNumberOfTags() {
-        ConnectMetrics.tags("k1", "v1", "k2", "v2", "extra");
-    }
-
     @Test(expected = IllegalArgumentException.class)
     public void testGettingGroupWithOddNumberOfTags() {
         metrics.group("name", "k1", "v1", "k2", "v2", "extra");
@@ -136,4 +129,40 @@ public void testMetricGroupIdWithoutTags() {
         assertNotNull(id1.tags());
         assertNotNull(id2.tags());
     }
-}
\ No newline at end of file
+
+    @Test
+    public void testRecreateWithClose() {
+        final Sensor originalSensor = addToGroup(metrics, false);
+        final Sensor recreatedSensor = addToGroup(metrics, true);
+        // because we closed the metricGroup, we get a brand-new sensor
+        assertNotSame(originalSensor, recreatedSensor);
+    }
+
+    @Test
+    public void testRecreateWithoutClose() {
+        final Sensor originalSensor = addToGroup(metrics, false);
+        final Sensor recreatedSensor = addToGroup(metrics, false);
+        // since we didn't close the group, the second addToGroup is idempotent
+        assertSame(originalSensor, recreatedSensor);
+    }
+
+    private Sensor addToGroup(ConnectMetrics connectMetrics, boolean shouldClose) {
+        ConnectMetricsRegistry registry = connectMetrics.registry();
+        ConnectMetrics.MetricGroup metricGroup = connectMetrics.group(registry.taskGroupName(),
+                registry.connectorTagName(), "conn_name");
+
+        if (shouldClose) {
+            metricGroup.close();
+        }
+
+        Sensor sensor = metricGroup.sensor("my_sensor");
+        sensor.add(metricName("x1"), new Max());
+        sensor.add(metricName("y2"), new Avg());
+
+        return sensor;
+    }
+
+    static MetricName metricName(String name) {
+        return new MetricName(name, "test_group", "metrics for testing", Collections.emptyMap());
+    }
+}
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 f8c4fd690dc52..c400e48e86d7d 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
@@ -23,6 +23,7 @@
 import org.apache.kafka.connect.runtime.isolation.PluginDesc;
 import org.apache.kafka.connect.runtime.isolation.Plugins;
 import org.apache.kafka.connect.transforms.Transformation;
+import org.apache.kafka.connect.transforms.predicates.Predicate;
 import org.junit.Test;
 
 import java.util.Collections;
@@ -32,12 +33,15 @@
 import java.util.Set;
 
 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 ConnectorConfigTest> {
 
-    private static final Plugins MOCK_PLUGINS = new Plugins(new HashMap()) {
+    public static final Plugins MOCK_PLUGINS = new Plugins(new HashMap()) {
         @Override
         public Set> transformations() {
             return Collections.emptySet();
@@ -81,33 +85,45 @@ public void noTransforms() {
         new ConnectorConfig(MOCK_PLUGINS, props);
     }
 
-    @Test(expected = ConfigException.class)
+    @Test
     public void danglingTransformAlias() {
         Map props = new HashMap<>();
         props.put("name", "test");
         props.put("connector.class", TestConnector.class.getName());
         props.put("transforms", "dangler");
-        new ConnectorConfig(MOCK_PLUGINS, props);
+        ConfigException e = assertThrows(ConfigException.class, () -> new ConnectorConfig(MOCK_PLUGINS, props));
+        assertTrue(e.getMessage().contains("Not a Transformation"));
     }
 
-    @Test(expected = ConfigException.class)
+    @Test
+    public void emptyConnectorName() {
+        Map props = new HashMap<>();
+        props.put("name", "");
+        props.put("connector.class", TestConnector.class.getName());
+        ConfigException e = assertThrows(ConfigException.class, () -> new ConnectorConfig(MOCK_PLUGINS, props));
+        assertTrue(e.getMessage().contains("String may not be empty"));
+    }
+
+    @Test
     public void wrongTransformationType() {
         Map props = new HashMap<>();
         props.put("name", "test");
         props.put("connector.class", TestConnector.class.getName());
         props.put("transforms", "a");
         props.put("transforms.a.type", "uninstantiable");
-        new ConnectorConfig(MOCK_PLUGINS, props);
+        ConfigException e = assertThrows(ConfigException.class, () -> new ConnectorConfig(MOCK_PLUGINS, props));
+        assertTrue(e.getMessage().contains("Class uninstantiable could not be found"));
     }
 
-    @Test(expected = ConfigException.class)
+    @Test
     public void unconfiguredTransform() {
         Map props = new HashMap<>();
         props.put("name", "test");
         props.put("connector.class", TestConnector.class.getName());
         props.put("transforms", "a");
         props.put("transforms.a.type", SimpleTransformation.class.getName());
-        new ConnectorConfig(MOCK_PLUGINS, props);
+        ConfigException e = assertThrows(ConfigException.class, () -> new ConnectorConfig(MOCK_PLUGINS, props));
+        assertTrue(e.getMessage().contains("Missing required configuration \"transforms.a.magic.number\" which"));
     }
 
     @Test
@@ -118,12 +134,8 @@ public void misconfiguredTransform() {
         props.put("transforms", "a");
         props.put("transforms.a.type", SimpleTransformation.class.getName());
         props.put("transforms.a.magic.number", "40");
-        try {
-            new ConnectorConfig(MOCK_PLUGINS, props);
-            fail();
-        } catch (ConfigException e) {
-            assertTrue(e.getMessage().contains("Value must be at least 42"));
-        }
+        ConfigException e = assertThrows(ConfigException.class, () -> new ConnectorConfig(MOCK_PLUGINS, props));
+        assertTrue(e.getMessage().contains("Value must be at least 42"));
     }
 
     @Test
@@ -168,5 +180,308 @@ public void multipleTransforms() {
         assertEquals(42, ((SimpleTransformation) transformations.get(0)).magicNumber);
         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())
+            );
+        }
+    }
+
+    @Test
+    public void wrongPredicateType() {
+        Map props = new HashMap<>();
+        props.put("name", "test");
+        props.put("connector.class", TestConnector.class.getName());
+        props.put("transforms", "a");
+        props.put("transforms.a.type", SimpleTransformation.class.getName());
+        props.put("transforms.a.magic.number", "42");
+        props.put("transforms.a.predicate", "my-pred");
+        props.put("predicates", "my-pred");
+        props.put("predicates.my-pred.type", TestConnector.class.getName());
+        ConfigException e = assertThrows(ConfigException.class, () -> new ConnectorConfig(MOCK_PLUGINS, props));
+        assertTrue(e.getMessage().contains("Not a Predicate"));
+    }
+
+    @Test
+    public void singleConditionalTransform() {
+        Map props = new HashMap<>();
+        props.put("name", "test");
+        props.put("connector.class", TestConnector.class.getName());
+        props.put("transforms", "a");
+        props.put("transforms.a.type", SimpleTransformation.class.getName());
+        props.put("transforms.a.magic.number", "42");
+        props.put("transforms.a.predicate", "my-pred");
+        props.put("transforms.a.negate", "true");
+        props.put("predicates", "my-pred");
+        props.put("predicates.my-pred.type", TestPredicate.class.getName());
+        props.put("predicates.my-pred.int", "84");
+        assertPredicatedTransform(props, true);
+    }
+
+    @Test
+    public void predicateNegationDefaultsToFalse() {
+        Map props = new HashMap<>();
+        props.put("name", "test");
+        props.put("connector.class", TestConnector.class.getName());
+        props.put("transforms", "a");
+        props.put("transforms.a.type", SimpleTransformation.class.getName());
+        props.put("transforms.a.magic.number", "42");
+        props.put("transforms.a.predicate", "my-pred");
+        props.put("predicates", "my-pred");
+        props.put("predicates.my-pred.type", TestPredicate.class.getName());
+        props.put("predicates.my-pred.int", "84");
+        assertPredicatedTransform(props, false);
+    }
+
+    @Test
+    public void abstractPredicate() {
+        Map props = new HashMap<>();
+        props.put("name", "test");
+        props.put("connector.class", TestConnector.class.getName());
+        props.put("transforms", "a");
+        props.put("transforms.a.type", SimpleTransformation.class.getName());
+        props.put("transforms.a.magic.number", "42");
+        props.put("transforms.a.predicate", "my-pred");
+        props.put("predicates", "my-pred");
+        props.put("predicates.my-pred.type", AbstractTestPredicate.class.getName());
+        props.put("predicates.my-pred.int", "84");
+        ConfigException e = assertThrows(ConfigException.class, () -> new ConnectorConfig(MOCK_PLUGINS, props));
+        assertTrue(e.getMessage().contains("Predicate is abstract and cannot be created"));
+    }
+
+    private void assertPredicatedTransform(Map props, boolean expectedNegated) {
+        final ConnectorConfig config = new ConnectorConfig(MOCK_PLUGINS, props);
+        final List> transformations = config.transformations();
+        assertEquals(1, transformations.size());
+        assertTrue(transformations.get(0) instanceof PredicatedTransformation);
+        PredicatedTransformation predicated = (PredicatedTransformation) transformations.get(0);
+
+        assertEquals(expectedNegated, predicated.negate);
+
+        assertTrue(predicated.delegate instanceof ConnectorConfigTest.SimpleTransformation);
+        assertEquals(42, ((SimpleTransformation) predicated.delegate).magicNumber);
+
+        assertTrue(predicated.predicate instanceof ConnectorConfigTest.TestPredicate);
+        assertEquals(84, ((TestPredicate) predicated.predicate).param);
+
+        predicated.close();
+
+        assertEquals(0, ((SimpleTransformation) predicated.delegate).magicNumber);
+        assertEquals(0, ((TestPredicate) predicated.predicate).param);
+    }
+
+    @Test
+    public void misconfiguredPredicate() {
+        Map props = new HashMap<>();
+        props.put("name", "test");
+        props.put("connector.class", TestConnector.class.getName());
+        props.put("transforms", "a");
+        props.put("transforms.a.type", SimpleTransformation.class.getName());
+        props.put("transforms.a.magic.number", "42");
+        props.put("transforms.a.predicate", "my-pred");
+        props.put("transforms.a.negate", "true");
+        props.put("predicates", "my-pred");
+        props.put("predicates.my-pred.type", TestPredicate.class.getName());
+        props.put("predicates.my-pred.int", "79");
+        try {
+            new ConnectorConfig(MOCK_PLUGINS, props);
+            fail();
+        } catch (ConfigException e) {
+            assertTrue(e.getMessage().contains("Value must be at least 80"));
+        }
+    }
+
+    @Test
+    public void missingPredicateAliasProperty() {
+        Map props = new HashMap<>();
+        props.put("name", "test");
+        props.put("connector.class", TestConnector.class.getName());
+        props.put("transforms", "a");
+        props.put("transforms.a.type", SimpleTransformation.class.getName());
+        props.put("transforms.a.magic.number", "42");
+        props.put("transforms.a.predicate", "my-pred");
+        // technically not needed
+        //props.put("predicates", "my-pred");
+        props.put("predicates.my-pred.type", TestPredicate.class.getName());
+        props.put("predicates.my-pred.int", "84");
+        new ConnectorConfig(MOCK_PLUGINS, props);
+    }
+
+    @Test
+    public void missingPredicateConfig() {
+        Map props = new HashMap<>();
+        props.put("name", "test");
+        props.put("connector.class", TestConnector.class.getName());
+        props.put("transforms", "a");
+        props.put("transforms.a.type", SimpleTransformation.class.getName());
+        props.put("transforms.a.magic.number", "42");
+        props.put("transforms.a.predicate", "my-pred");
+        props.put("predicates", "my-pred");
+        //props.put("predicates.my-pred.type", TestPredicate.class.getName());
+        //props.put("predicates.my-pred.int", "84");
+        ConfigException e = assertThrows(ConfigException.class, () -> new ConnectorConfig(MOCK_PLUGINS, props));
+        assertTrue(e.getMessage().contains("Not a Predicate"));
+    }
+
+    @Test
+    public void negatedButNoPredicate() {
+        Map props = new HashMap<>();
+        props.put("name", "test");
+        props.put("connector.class", TestConnector.class.getName());
+        props.put("transforms", "a");
+        props.put("transforms.a.type", SimpleTransformation.class.getName());
+        props.put("transforms.a.magic.number", "42");
+        props.put("transforms.a.negate", "true");
+        ConfigException e = assertThrows(ConfigException.class, () -> new ConnectorConfig(MOCK_PLUGINS, props));
+        assertTrue(e.getMessage().contains("there is no config 'transforms.a.predicate' defining a predicate to be negated"));
+    }
+
+    public static class TestPredicate> implements Predicate  {
+
+        int param;
+
+        public TestPredicate() { }
+
+        @Override
+        public ConfigDef config() {
+            return new ConfigDef().define("int", ConfigDef.Type.INT, 80, ConfigDef.Range.atLeast(80), ConfigDef.Importance.MEDIUM,
+                    "A test parameter");
+        }
+
+        @Override
+        public boolean test(R record) {
+            return false;
+        }
+
+        @Override
+        public void close() {
+            param = 0;
+        }
+
+        @Override
+        public void configure(Map configs) {
+            param = Integer.parseInt((String) configs.get("int"));
+        }
+    }
+
+    public static abstract class AbstractTestPredicate> implements Predicate  {
+
+        public AbstractTestPredicate() { }
+
+    }
+
+    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 {
+
+        }
+    }
+
+    @Test
+    public void testEnrichedConfigDef() {
+        String alias = "hdt";
+        String prefix = ConnectorConfig.TRANSFORMS_CONFIG + "." + alias + ".";
+        Map props = new HashMap<>();
+        props.put(ConnectorConfig.TRANSFORMS_CONFIG, alias);
+        props.put(prefix + "type", HasDuplicateConfigTransformation.class.getName());
+        ConfigDef def = ConnectorConfig.enrich(MOCK_PLUGINS, new ConfigDef(), props, false);
+        assertEnrichedConfigDef(def, prefix, HasDuplicateConfigTransformation.MUST_EXIST_KEY, ConfigDef.Type.BOOLEAN);
+        assertEnrichedConfigDef(def, prefix, PredicatedTransformation.PREDICATE_CONFIG, ConfigDef.Type.STRING);
+        assertEnrichedConfigDef(def, prefix, PredicatedTransformation.NEGATE_CONFIG, ConfigDef.Type.BOOLEAN);
+    }
+
+    private static void assertEnrichedConfigDef(ConfigDef def, String prefix, String keyName, ConfigDef.Type expectedType) {
+        assertNull(def.configKeys().get(keyName));
+        ConfigDef.ConfigKey configKey = def.configKeys().get(prefix + keyName);
+        assertNotNull(prefix + keyName + "' config must be present", configKey);
+        assertEquals(prefix + keyName + "' config should be a " + expectedType, expectedType, configKey.type);
+    }
+
+    public static class HasDuplicateConfigTransformation> implements Transformation {
+        private static final String MUST_EXIST_KEY = "must.exist.key";
+        private static final ConfigDef CONFIG_DEF = new ConfigDef()
+                // this configDef is duplicate. It should be removed automatically so as to avoid duplicate config error.
+                .define(PredicatedTransformation.PREDICATE_CONFIG, ConfigDef.Type.INT, ConfigDef.NO_DEFAULT_VALUE, ConfigDef.Importance.MEDIUM, "fake")
+                // this configDef is duplicate. It should be removed automatically so as to avoid duplicate config error.
+                .define(PredicatedTransformation.NEGATE_CONFIG, ConfigDef.Type.INT, 123, ConfigDef.Importance.MEDIUM, "fake")
+                // this configDef should appear if above duplicate configDef is removed without any error
+                .define(MUST_EXIST_KEY, ConfigDef.Type.BOOLEAN, true, ConfigDef.Importance.MEDIUM, "this key must exist");
+
+        @Override
+        public R apply(R record) {
+            return record;
+        }
+
+        @Override
+        public ConfigDef config() {
+            return CONFIG_DEF;
+        }
+
+        @Override
+        public void close() {
+        }
+
+        @Override
+        public void configure(Map configs) {
+        }
+    }
 }
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
new file mode 100644
index 0000000000000..70bbfc6590cd1
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/ErrorHandlingTaskTest.java
@@ -0,0 +1,639 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.admin.NewTopic;
+import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.clients.producer.KafkaProducer;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.connect.connector.ConnectRecord;
+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.RetriableException;
+import org.apache.kafka.connect.integration.MonitorableSourceConnector;
+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;
+import org.apache.kafka.connect.runtime.isolation.PluginClassLoader;
+import org.apache.kafka.connect.runtime.isolation.Plugins;
+import org.apache.kafka.connect.runtime.standalone.StandaloneConfig;
+import org.apache.kafka.connect.sink.SinkConnector;
+import org.apache.kafka.connect.sink.SinkRecord;
+import org.apache.kafka.connect.sink.SinkTask;
+import org.apache.kafka.connect.source.SourceRecord;
+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.OffsetStorageReaderImpl;
+import org.apache.kafka.connect.storage.OffsetStorageWriter;
+import org.apache.kafka.connect.storage.StatusBackingStore;
+import org.apache.kafka.connect.storage.StringConverter;
+import org.apache.kafka.connect.transforms.Transformation;
+import org.apache.kafka.connect.transforms.util.SimpleConfig;
+import org.apache.kafka.connect.util.ConnectorTaskId;
+import org.apache.kafka.connect.util.TopicAdmin;
+import org.easymock.Capture;
+import org.easymock.EasyMock;
+import org.easymock.IExpectationSetters;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.powermock.api.easymock.PowerMock;
+import org.powermock.api.easymock.annotation.Mock;
+import org.powermock.core.classloader.annotations.PowerMockIgnore;
+import org.powermock.core.classloader.annotations.PrepareForTest;
+import org.powermock.modules.junit4.PowerMockRunner;
+import org.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.Map;
+
+import static java.util.Collections.emptyMap;
+import static java.util.Collections.singletonList;
+import static org.apache.kafka.common.utils.Time.SYSTEM;
+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.TOPIC_CREATION_ENABLE_CONFIG;
+import static org.junit.Assert.assertEquals;
+
+@RunWith(PowerMockRunner.class)
+@PrepareForTest({WorkerSinkTask.class, WorkerSourceTask.class})
+@PowerMockIgnore("javax.management.*")
+public class ErrorHandlingTaskTest {
+
+    private static final String TOPIC = "test";
+    private static final int PARTITION1 = 12;
+    private static final int PARTITION2 = 13;
+    private static final long FIRST_OFFSET = 45;
+
+    @Mock Plugins plugins;
+
+    private static final Map TASK_PROPS = new HashMap<>();
+
+    static {
+        TASK_PROPS.put(SinkConnector.TOPICS_CONFIG, TOPIC);
+        TASK_PROPS.put(TaskConfig.TASK_CLASS_CONFIG, TestSinkTask.class.getName());
+    }
+
+    public static final long OPERATOR_RETRY_TIMEOUT_MILLIS = 60000;
+    public static final long OPERATOR_RETRY_MAX_DELAY_MILLIS = 5000;
+    public static final ToleranceType OPERATOR_TOLERANCE_TYPE = ToleranceType.ALL;
+
+    private static final TaskConfig TASK_CONFIG = new TaskConfig(TASK_PROPS);
+
+    private ConnectorTaskId taskId = new ConnectorTaskId("job", 0);
+    private TargetState initialState = TargetState.STARTED;
+    private Time time;
+    private MockConnectMetrics metrics;
+    @SuppressWarnings("unused")
+    @Mock
+    private SinkTask sinkTask;
+    @SuppressWarnings("unused")
+    @Mock
+    private SourceTask sourceTask;
+    private Capture sinkTaskContext = EasyMock.newCapture();
+    private WorkerConfig workerConfig;
+    private SourceConnectorConfig sourceConfig;
+    @Mock
+    private PluginClassLoader pluginLoader;
+    @SuppressWarnings("unused")
+    @Mock
+    private HeaderConverter headerConverter;
+    private WorkerSinkTask workerSinkTask;
+    private WorkerSourceTask workerSourceTask;
+    @SuppressWarnings("unused")
+    @Mock
+    private KafkaConsumer consumer;
+    @SuppressWarnings("unused")
+    @Mock
+    private KafkaProducer producer;
+    @SuppressWarnings("unused")
+    @Mock private TopicAdmin admin;
+
+    @Mock
+    OffsetStorageReaderImpl offsetReader;
+    @Mock
+    OffsetStorageWriter offsetWriter;
+
+    private Capture rebalanceListener = EasyMock.newCapture();
+    @SuppressWarnings("unused")
+    @Mock
+    private TaskStatus.Listener statusListener;
+    @SuppressWarnings("unused")
+    @Mock private StatusBackingStore statusBackingStore;
+
+    private ErrorHandlingMetrics errorHandlingMetrics;
+
+    // when this test becomes parameterized, this variable will be a test parameter
+    public boolean enableTopicCreation = false;
+
+    @Before
+    public void setup() {
+        time = new MockTime(0, 0, 0);
+        metrics = new MockConnectMetrics();
+        Map workerProps = new HashMap<>();
+        workerProps.put("key.converter", "org.apache.kafka.connect.json.JsonConverter");
+        workerProps.put("value.converter", "org.apache.kafka.connect.json.JsonConverter");
+        workerProps.put("internal.key.converter", "org.apache.kafka.connect.json.JsonConverter");
+        workerProps.put("internal.value.converter", "org.apache.kafka.connect.json.JsonConverter");
+        workerProps.put("internal.key.converter.schemas.enable", "false");
+        workerProps.put("internal.value.converter.schemas.enable", "false");
+        workerProps.put("offset.storage.file.filename", "/tmp/connect.offsets");
+        workerProps.put(TOPIC_CREATION_ENABLE_CONFIG, String.valueOf(enableTopicCreation));
+        pluginLoader = PowerMock.createMock(PluginClassLoader.class);
+        workerConfig = new StandaloneConfig(workerProps);
+        sourceConfig = new SourceConnectorConfig(plugins, sourceConnectorProps(TOPIC), true);
+        errorHandlingMetrics = new ErrorHandlingMetrics(taskId, metrics);
+    }
+
+    private Map sourceConnectorProps(String topic) {
+        // setup up props for the source connector
+        Map props = new HashMap<>();
+        props.put("name", "foo-connector");
+        props.put(CONNECTOR_CLASS_CONFIG, MonitorableSourceConnector.class.getSimpleName());
+        props.put(TASKS_MAX_CONFIG, String.valueOf(1));
+        props.put(TOPIC_CONFIG, topic);
+        props.put(KEY_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
+        props.put(VALUE_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
+        return props;
+    }
+
+    @After
+    public void tearDown() {
+        if (metrics != null) {
+            metrics.stop();
+        }
+    }
+
+    @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);
+
+        expectClose();
+
+        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);
+
+        expectClose();
+
+        // 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<>();
+        reportProps.put(ConnectorConfig.ERRORS_LOG_ENABLE_CONFIG, "true");
+        reportProps.put(ConnectorConfig.ERRORS_LOG_INCLUDE_MESSAGES_CONFIG, "true");
+        LogReporter reporter = new LogReporter(taskId, connConfig(reportProps), errorHandlingMetrics);
+
+        RetryWithToleranceOperator retryWithToleranceOperator = operator();
+        retryWithToleranceOperator.metrics(errorHandlingMetrics);
+        retryWithToleranceOperator.reporters(singletonList(reporter));
+        createSinkTask(initialState, retryWithToleranceOperator);
+
+        expectInitializeTask();
+        expectTaskGetTopic(true);
+
+        // valid json
+        ConsumerRecord record1 = new ConsumerRecord<>(TOPIC, PARTITION1, FIRST_OFFSET, null, "{\"a\": 10}".getBytes());
+        // bad json
+        ConsumerRecord record2 = new ConsumerRecord<>(TOPIC, PARTITION2, FIRST_OFFSET, null, "{\"a\" 10}".getBytes());
+
+        EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andReturn(records(record1));
+        EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andReturn(records(record2));
+
+        sinkTask.put(EasyMock.anyObject());
+        EasyMock.expectLastCall().times(2);
+
+        PowerMock.replayAll();
+
+        workerSinkTask.initialize(TASK_CONFIG);
+        workerSinkTask.initializeAndStart();
+        workerSinkTask.iteration();
+
+        workerSinkTask.iteration();
+
+        // two records were consumed from Kafka
+        assertSinkMetricValue("sink-record-read-total", 2.0);
+        // only one was written to the task
+        assertSinkMetricValue("sink-record-send-total", 1.0);
+        // one record completely failed (converter issues)
+        assertErrorHandlingMetricValue("total-record-errors", 1.0);
+        // 2 failures in the transformation, and 1 in the converter
+        assertErrorHandlingMetricValue("total-record-failures", 3.0);
+        // one record completely failed (converter issues), and thus was skipped
+        assertErrorHandlingMetricValue("total-records-skipped", 1.0);
+
+        PowerMock.verifyAll();
+    }
+
+    private RetryWithToleranceOperator operator() {
+        return new RetryWithToleranceOperator(OPERATOR_RETRY_TIMEOUT_MILLIS, OPERATOR_RETRY_MAX_DELAY_MILLIS, OPERATOR_TOLERANCE_TYPE, SYSTEM);
+    }
+
+    @Test
+    public void testErrorHandlingInSourceTasks() throws Exception {
+        Map reportProps = new HashMap<>();
+        reportProps.put(ConnectorConfig.ERRORS_LOG_ENABLE_CONFIG, "true");
+        reportProps.put(ConnectorConfig.ERRORS_LOG_INCLUDE_MESSAGES_CONFIG, "true");
+        LogReporter reporter = new LogReporter(taskId, connConfig(reportProps), errorHandlingMetrics);
+
+        RetryWithToleranceOperator retryWithToleranceOperator = operator();
+        retryWithToleranceOperator.metrics(errorHandlingMetrics);
+        retryWithToleranceOperator.reporters(singletonList(reporter));
+        createSourceTask(initialState, retryWithToleranceOperator);
+
+        // valid json
+        Schema valSchema = SchemaBuilder.struct().field("val", Schema.INT32_SCHEMA).build();
+        Struct struct1 = new Struct(valSchema).put("val", 1234);
+        SourceRecord record1 = new SourceRecord(emptyMap(), emptyMap(), TOPIC, PARTITION1, valSchema, struct1);
+        Struct struct2 = new Struct(valSchema).put("val", 6789);
+        SourceRecord record2 = new SourceRecord(emptyMap(), emptyMap(), TOPIC, PARTITION1, valSchema, struct2);
+
+        EasyMock.expect(workerSourceTask.isStopping()).andReturn(false);
+        EasyMock.expect(workerSourceTask.isStopping()).andReturn(false);
+        EasyMock.expect(workerSourceTask.isStopping()).andReturn(true);
+
+        EasyMock.expect(workerSourceTask.commitOffsets()).andReturn(true);
+
+        offsetWriter.offset(EasyMock.anyObject(), EasyMock.anyObject());
+        EasyMock.expectLastCall().times(2);
+        sourceTask.initialize(EasyMock.anyObject());
+        EasyMock.expectLastCall();
+
+        sourceTask.start(EasyMock.anyObject());
+        EasyMock.expectLastCall();
+
+        EasyMock.expect(sourceTask.poll()).andReturn(singletonList(record1));
+        EasyMock.expect(sourceTask.poll()).andReturn(singletonList(record2));
+        expectTopicCreation(TOPIC);
+        EasyMock.expect(producer.send(EasyMock.anyObject(), EasyMock.anyObject())).andReturn(null).times(2);
+
+        PowerMock.replayAll();
+
+        workerSourceTask.initialize(TASK_CONFIG);
+        workerSourceTask.execute();
+
+        // two records were consumed from Kafka
+        assertSourceMetricValue("source-record-poll-total", 2.0);
+        // only one was written to the task
+        assertSourceMetricValue("source-record-write-total", 0.0);
+        // one record completely failed (converter issues)
+        assertErrorHandlingMetricValue("total-record-errors", 0.0);
+        // 2 failures in the transformation, and 1 in the converter
+        assertErrorHandlingMetricValue("total-record-failures", 4.0);
+        // one record completely failed (converter issues), and thus was skipped
+        assertErrorHandlingMetricValue("total-records-skipped", 0.0);
+
+        PowerMock.verifyAll();
+    }
+
+    private ConnectorConfig connConfig(Map connProps) {
+        Map props = new HashMap<>();
+        props.put(ConnectorConfig.NAME_CONFIG, "test");
+        props.put(ConnectorConfig.CONNECTOR_CLASS_CONFIG, SinkTask.class.getName());
+        props.putAll(connProps);
+        return new ConnectorConfig(plugins, props);
+    }
+
+    @Test
+    public void testErrorHandlingInSourceTasksWthBadConverter() throws Exception {
+        Map reportProps = new HashMap<>();
+        reportProps.put(ConnectorConfig.ERRORS_LOG_ENABLE_CONFIG, "true");
+        reportProps.put(ConnectorConfig.ERRORS_LOG_INCLUDE_MESSAGES_CONFIG, "true");
+        LogReporter reporter = new LogReporter(taskId, connConfig(reportProps), errorHandlingMetrics);
+
+        RetryWithToleranceOperator retryWithToleranceOperator = operator();
+        retryWithToleranceOperator.metrics(errorHandlingMetrics);
+        retryWithToleranceOperator.reporters(singletonList(reporter));
+        createSourceTask(initialState, retryWithToleranceOperator, badConverter());
+
+        // valid json
+        Schema valSchema = SchemaBuilder.struct().field("val", Schema.INT32_SCHEMA).build();
+        Struct struct1 = new Struct(valSchema).put("val", 1234);
+        SourceRecord record1 = new SourceRecord(emptyMap(), emptyMap(), TOPIC, PARTITION1, valSchema, struct1);
+        Struct struct2 = new Struct(valSchema).put("val", 6789);
+        SourceRecord record2 = new SourceRecord(emptyMap(), emptyMap(), TOPIC, PARTITION1, valSchema, struct2);
+
+        EasyMock.expect(workerSourceTask.isStopping()).andReturn(false);
+        EasyMock.expect(workerSourceTask.isStopping()).andReturn(false);
+        EasyMock.expect(workerSourceTask.isStopping()).andReturn(true);
+
+        EasyMock.expect(workerSourceTask.commitOffsets()).andReturn(true);
+
+        offsetWriter.offset(EasyMock.anyObject(), EasyMock.anyObject());
+        EasyMock.expectLastCall().times(2);
+        sourceTask.initialize(EasyMock.anyObject());
+        EasyMock.expectLastCall();
+
+        sourceTask.start(EasyMock.anyObject());
+        EasyMock.expectLastCall();
+
+        EasyMock.expect(sourceTask.poll()).andReturn(singletonList(record1));
+        EasyMock.expect(sourceTask.poll()).andReturn(singletonList(record2));
+        expectTopicCreation(TOPIC);
+        EasyMock.expect(producer.send(EasyMock.anyObject(), EasyMock.anyObject())).andReturn(null).times(2);
+
+        PowerMock.replayAll();
+
+        workerSourceTask.initialize(TASK_CONFIG);
+        workerSourceTask.execute();
+
+        // two records were consumed from Kafka
+        assertSourceMetricValue("source-record-poll-total", 2.0);
+        // only one was written to the task
+        assertSourceMetricValue("source-record-write-total", 0.0);
+        // one record completely failed (converter issues)
+        assertErrorHandlingMetricValue("total-record-errors", 0.0);
+        // 2 failures in the transformation, and 1 in the converter
+        assertErrorHandlingMetricValue("total-record-failures", 8.0);
+        // one record completely failed (converter issues), and thus was skipped
+        assertErrorHandlingMetricValue("total-records-skipped", 0.0);
+
+        PowerMock.verifyAll();
+    }
+
+    private void assertSinkMetricValue(String name, double expected) {
+        ConnectMetrics.MetricGroup sinkTaskGroup = workerSinkTask.sinkTaskMetricsGroup().metricGroup();
+        double measured = metrics.currentMetricValueAsDouble(sinkTaskGroup, name);
+        assertEquals(expected, measured, 0.001d);
+    }
+
+    private void assertSourceMetricValue(String name, double expected) {
+        ConnectMetrics.MetricGroup sinkTaskGroup = workerSourceTask.sourceTaskMetricsGroup().metricGroup();
+        double measured = metrics.currentMetricValueAsDouble(sinkTaskGroup, name);
+        assertEquals(expected, measured, 0.001d);
+    }
+
+    private void assertErrorHandlingMetricValue(String name, double expected) {
+        ConnectMetrics.MetricGroup sinkTaskGroup = errorHandlingMetrics.metricGroup();
+        double measured = metrics.currentMetricValueAsDouble(sinkTaskGroup, name);
+        assertEquals(expected, measured, 0.001d);
+    }
+
+    private void expectInitializeTask() throws Exception {
+        consumer.subscribe(EasyMock.eq(singletonList(TOPIC)), EasyMock.capture(rebalanceListener));
+        PowerMock.expectLastCall();
+
+        sinkTask.initialize(EasyMock.capture(sinkTaskContext));
+        PowerMock.expectLastCall();
+        sinkTask.start(TASK_PROPS);
+        PowerMock.expectLastCall();
+    }
+
+    private void expectTaskGetTopic(boolean anyTimes) {
+        final Capture connectorCapture = EasyMock.newCapture();
+        final Capture topicCapture = EasyMock.newCapture();
+        IExpectationSetters expect = EasyMock.expect(statusBackingStore.getTopic(
+                EasyMock.capture(connectorCapture),
+                EasyMock.capture(topicCapture)));
+        if (anyTimes) {
+            expect.andStubAnswer(() -> new TopicStatus(
+                    topicCapture.getValue(),
+                    new ConnectorTaskId(connectorCapture.getValue(), 0),
+                    Time.SYSTEM.milliseconds()));
+        } else {
+            expect.andAnswer(() -> new TopicStatus(
+                    topicCapture.getValue(),
+                    new ConnectorTaskId(connectorCapture.getValue(), 0),
+                    Time.SYSTEM.milliseconds()));
+        }
+        if (connectorCapture.hasCaptured() && topicCapture.hasCaptured()) {
+            assertEquals("job", connectorCapture.getValue());
+            assertEquals(TOPIC, topicCapture.getValue());
+        }
+    }
+
+    private void expectClose() {
+        producer.close(EasyMock.anyObject(Duration.class));
+        EasyMock.expectLastCall();
+
+        admin.close(EasyMock.anyObject(Duration.class));
+        EasyMock.expectLastCall();
+    }
+
+    private void expectTopicCreation(String topic) {
+        if (workerConfig.topicCreationEnable()) {
+            EasyMock.expect(admin.describeTopics(topic)).andReturn(Collections.emptyMap());
+
+            Capture newTopicCapture = EasyMock.newCapture();
+            EasyMock.expect(admin.createTopic(EasyMock.capture(newTopicCapture))).andReturn(true);
+        }
+    }
+
+    private void createSinkTask(TargetState initialState, RetryWithToleranceOperator retryWithToleranceOperator) {
+        JsonConverter converter = new JsonConverter();
+        Map oo = workerConfig.originalsWithPrefix("value.converter.");
+        oo.put("converter.type", "value");
+        oo.put("schemas.enable", "false");
+        converter.configure(oo);
+
+        TransformationChain sinkTransforms = new TransformationChain<>(singletonList(new FaultyPassthrough()), retryWithToleranceOperator);
+
+        workerSinkTask = new WorkerSinkTask(
+            taskId, sinkTask, statusListener, initialState, workerConfig,
+            ClusterConfigState.EMPTY, metrics, converter, converter,
+            headerConverter, sinkTransforms, consumer, pluginLoader, time,
+                retryWithToleranceOperator, null, statusBackingStore);
+    }
+
+    private void createSourceTask(TargetState initialState, RetryWithToleranceOperator retryWithToleranceOperator) {
+        JsonConverter converter = new JsonConverter();
+        Map oo = workerConfig.originalsWithPrefix("value.converter.");
+        oo.put("converter.type", "value");
+        oo.put("schemas.enable", "false");
+        converter.configure(oo);
+
+        createSourceTask(initialState, retryWithToleranceOperator, converter);
+    }
+
+    private Converter badConverter() {
+        FaultyConverter converter = new FaultyConverter();
+        Map oo = workerConfig.originalsWithPrefix("value.converter.");
+        oo.put("converter.type", "value");
+        oo.put("schemas.enable", "false");
+        converter.configure(oo);
+        return converter;
+    }
+
+    private void createSourceTask(TargetState initialState, RetryWithToleranceOperator retryWithToleranceOperator, Converter converter) {
+        TransformationChain sourceTransforms = new TransformationChain<>(singletonList(new FaultyPassthrough()), retryWithToleranceOperator);
+
+        workerSourceTask = PowerMock.createPartialMock(
+                WorkerSourceTask.class, new String[]{"commitOffsets", "isStopping"},
+                taskId, sourceTask, statusListener, initialState, converter, converter, headerConverter, sourceTransforms,
+                producer, admin, null,
+                offsetReader, offsetWriter, workerConfig,
+                ClusterConfigState.EMPTY, metrics, pluginLoader, time, retryWithToleranceOperator,
+                statusBackingStore);
+    }
+
+    private ConsumerRecords records(ConsumerRecord record) {
+        return new ConsumerRecords<>(Collections.singletonMap(
+                new TopicPartition(record.topic(), record.partition()), singletonList(record)));
+    }
+
+    private abstract static class TestSinkTask extends SinkTask {
+    }
+
+    static class FaultyConverter extends JsonConverter {
+        private static final Logger log = LoggerFactory.getLogger(FaultyConverter.class);
+        private int invocations = 0;
+
+        public byte[] fromConnectData(String topic, Schema schema, Object value) {
+            if (value == null) {
+                return super.fromConnectData(topic, schema, null);
+            }
+            invocations++;
+            if (invocations % 3 == 0) {
+                log.debug("Succeeding record: {} where invocations={}", value, invocations);
+                return super.fromConnectData(topic, schema, value);
+            } else {
+                log.debug("Failing record: {} at invocations={}", value, invocations);
+                throw new RetriableException("Bad invocations " + invocations + " for mod 3");
+            }
+        }
+    }
+
+    static class FaultyPassthrough> implements Transformation {
+
+        private static final Logger log = LoggerFactory.getLogger(FaultyPassthrough.class);
+
+        private static final String MOD_CONFIG = "mod";
+        private static final int MOD_CONFIG_DEFAULT = 3;
+
+        public static final ConfigDef CONFIG_DEF = new ConfigDef()
+                .define(MOD_CONFIG, ConfigDef.Type.INT, MOD_CONFIG_DEFAULT, ConfigDef.Importance.MEDIUM, "Pass records without failure only if timestamp % mod == 0");
+
+        private int mod = MOD_CONFIG_DEFAULT;
+
+        private int invocations = 0;
+
+        @Override
+        public R apply(R record) {
+            invocations++;
+            if (invocations % mod == 0) {
+                log.debug("Succeeding record: {} where invocations={}", record, invocations);
+                return record;
+            } else {
+                log.debug("Failing record: {} at invocations={}", record, invocations);
+                throw new RetriableException("Bad invocations " + invocations + " for mod " + mod);
+            }
+        }
+
+        @Override
+        public ConfigDef config() {
+            return CONFIG_DEF;
+        }
+
+        @Override
+        public void close() {
+            log.info("Shutting down transform");
+        }
+
+        @Override
+        public void configure(Map configs) {
+            final SimpleConfig config = new SimpleConfig(CONFIG_DEF, configs);
+            mod = Math.max(config.getInt(MOD_CONFIG), 2);
+            log.info("Configuring {}. Setting mod to {}", this.getClass(), mod);
+        }
+    }
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/ErrorHandlingTaskWithTopicCreationTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/ErrorHandlingTaskWithTopicCreationTest.java
new file mode 100644
index 0000000000000..9c115ac517821
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/ErrorHandlingTaskWithTopicCreationTest.java
@@ -0,0 +1,653 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.admin.NewTopic;
+import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.clients.producer.KafkaProducer;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.connect.connector.ConnectRecord;
+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.RetriableException;
+import org.apache.kafka.connect.integration.MonitorableSourceConnector;
+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;
+import org.apache.kafka.connect.runtime.errors.WorkerErrantRecordReporter;
+import org.apache.kafka.connect.runtime.isolation.PluginClassLoader;
+import org.apache.kafka.connect.runtime.isolation.Plugins;
+import org.apache.kafka.connect.runtime.standalone.StandaloneConfig;
+import org.apache.kafka.connect.sink.SinkConnector;
+import org.apache.kafka.connect.sink.SinkRecord;
+import org.apache.kafka.connect.sink.SinkTask;
+import org.apache.kafka.connect.source.SourceRecord;
+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.OffsetStorageReaderImpl;
+import org.apache.kafka.connect.storage.OffsetStorageWriter;
+import org.apache.kafka.connect.storage.StatusBackingStore;
+import org.apache.kafka.connect.storage.StringConverter;
+import org.apache.kafka.connect.transforms.Transformation;
+import org.apache.kafka.connect.transforms.util.SimpleConfig;
+import org.apache.kafka.connect.util.ConnectorTaskId;
+import org.apache.kafka.connect.util.TopicAdmin;
+import org.apache.kafka.connect.util.TopicCreationGroup;
+import org.easymock.Capture;
+import org.easymock.EasyMock;
+import org.easymock.IExpectationSetters;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.powermock.api.easymock.PowerMock;
+import org.powermock.api.easymock.annotation.Mock;
+import org.powermock.core.classloader.annotations.PowerMockIgnore;
+import org.powermock.core.classloader.annotations.PrepareForTest;
+import org.powermock.modules.junit4.PowerMockRunner;
+import org.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.Map;
+
+import static java.util.Collections.emptyMap;
+import static java.util.Collections.singletonList;
+import static org.apache.kafka.common.utils.Time.SYSTEM;
+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.SourceConnectorConfig.TOPIC_CREATION_GROUPS_CONFIG;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.DEFAULT_TOPIC_CREATION_PREFIX;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.INCLUDE_REGEX_CONFIG;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.PARTITIONS_CONFIG;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.REPLICATION_FACTOR_CONFIG;
+import static org.apache.kafka.connect.runtime.WorkerConfig.TOPIC_CREATION_ENABLE_CONFIG;
+import static org.junit.Assert.assertEquals;
+
+@RunWith(PowerMockRunner.class)
+@PrepareForTest({WorkerSinkTask.class, WorkerSourceTask.class})
+@PowerMockIgnore("javax.management.*")
+public class ErrorHandlingTaskWithTopicCreationTest {
+
+    private static final String TOPIC = "test";
+    private static final int PARTITION1 = 12;
+    private static final int PARTITION2 = 13;
+    private static final long FIRST_OFFSET = 45;
+
+    @Mock Plugins plugins;
+
+    private static final Map TASK_PROPS = new HashMap<>();
+
+    static {
+        TASK_PROPS.put(SinkConnector.TOPICS_CONFIG, TOPIC);
+        TASK_PROPS.put(TaskConfig.TASK_CLASS_CONFIG, TestSinkTask.class.getName());
+    }
+
+    public static final long OPERATOR_RETRY_TIMEOUT_MILLIS = 60000;
+    public static final long OPERATOR_RETRY_MAX_DELAY_MILLIS = 5000;
+    public static final ToleranceType OPERATOR_TOLERANCE_TYPE = ToleranceType.ALL;
+
+    private static final TaskConfig TASK_CONFIG = new TaskConfig(TASK_PROPS);
+
+    private ConnectorTaskId taskId = new ConnectorTaskId("job", 0);
+    private TargetState initialState = TargetState.STARTED;
+    private Time time;
+    private MockConnectMetrics metrics;
+    @SuppressWarnings("unused")
+    @Mock
+    private SinkTask sinkTask;
+    @SuppressWarnings("unused")
+    @Mock
+    private SourceTask sourceTask;
+    private Capture sinkTaskContext = EasyMock.newCapture();
+    private WorkerConfig workerConfig;
+    private SourceConnectorConfig sourceConfig;
+    @Mock
+    private PluginClassLoader pluginLoader;
+    @SuppressWarnings("unused")
+    @Mock
+    private HeaderConverter headerConverter;
+    private WorkerSinkTask workerSinkTask;
+    private WorkerSourceTask workerSourceTask;
+    @SuppressWarnings("unused")
+    @Mock
+    private KafkaConsumer consumer;
+    @SuppressWarnings("unused")
+    @Mock
+    private KafkaProducer producer;
+    @SuppressWarnings("unused")
+    @Mock private TopicAdmin admin;
+
+    @Mock
+    OffsetStorageReaderImpl offsetReader;
+    @Mock
+    OffsetStorageWriter offsetWriter;
+
+    private Capture rebalanceListener = EasyMock.newCapture();
+    @SuppressWarnings("unused")
+    @Mock
+    private TaskStatus.Listener statusListener;
+    @SuppressWarnings("unused")
+    @Mock private StatusBackingStore statusBackingStore;
+
+    @Mock
+    private WorkerErrantRecordReporter workerErrantRecordReporter;
+
+    private ErrorHandlingMetrics errorHandlingMetrics;
+
+    // when this test becomes parameterized, this variable will be a test parameter
+    public boolean enableTopicCreation = true;
+
+    @Before
+    public void setup() {
+        time = new MockTime(0, 0, 0);
+        metrics = new MockConnectMetrics();
+        Map workerProps = new HashMap<>();
+        workerProps.put("key.converter", "org.apache.kafka.connect.json.JsonConverter");
+        workerProps.put("value.converter", "org.apache.kafka.connect.json.JsonConverter");
+        workerProps.put("internal.key.converter", "org.apache.kafka.connect.json.JsonConverter");
+        workerProps.put("internal.value.converter", "org.apache.kafka.connect.json.JsonConverter");
+        workerProps.put("internal.key.converter.schemas.enable", "false");
+        workerProps.put("internal.value.converter.schemas.enable", "false");
+        workerProps.put("offset.storage.file.filename", "/tmp/connect.offsets");
+        workerProps.put(TOPIC_CREATION_ENABLE_CONFIG, String.valueOf(enableTopicCreation));
+        pluginLoader = PowerMock.createMock(PluginClassLoader.class);
+        workerConfig = new StandaloneConfig(workerProps);
+        sourceConfig = new SourceConnectorConfig(plugins, sourceConnectorPropsWithGroups(TOPIC), true);
+        errorHandlingMetrics = new ErrorHandlingMetrics(taskId, metrics);
+    }
+
+    private Map sourceConnectorPropsWithGroups(String topic) {
+        // setup up props for the source connector
+        Map props = new HashMap<>();
+        props.put("name", "foo-connector");
+        props.put(CONNECTOR_CLASS_CONFIG, MonitorableSourceConnector.class.getSimpleName());
+        props.put(TASKS_MAX_CONFIG, String.valueOf(1));
+        props.put(TOPIC_CONFIG, topic);
+        props.put(KEY_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
+        props.put(VALUE_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
+        props.put(TOPIC_CREATION_GROUPS_CONFIG, String.join(",", "foo", "bar"));
+        props.put(DEFAULT_TOPIC_CREATION_PREFIX + REPLICATION_FACTOR_CONFIG, String.valueOf(1));
+        props.put(DEFAULT_TOPIC_CREATION_PREFIX + PARTITIONS_CONFIG, String.valueOf(1));
+        props.put(SourceConnectorConfig.TOPIC_CREATION_PREFIX + "foo" + "." + INCLUDE_REGEX_CONFIG, topic);
+        return props;
+    }
+
+    @After
+    public void tearDown() {
+        if (metrics != null) {
+            metrics.stop();
+        }
+    }
+
+    @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);
+
+        expectClose();
+
+        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);
+
+        expectClose();
+
+        // 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<>();
+        reportProps.put(ConnectorConfig.ERRORS_LOG_ENABLE_CONFIG, "true");
+        reportProps.put(ConnectorConfig.ERRORS_LOG_INCLUDE_MESSAGES_CONFIG, "true");
+        LogReporter reporter = new LogReporter(taskId, connConfig(reportProps), errorHandlingMetrics);
+
+        RetryWithToleranceOperator retryWithToleranceOperator = operator();
+        retryWithToleranceOperator.metrics(errorHandlingMetrics);
+        retryWithToleranceOperator.reporters(singletonList(reporter));
+        createSinkTask(initialState, retryWithToleranceOperator);
+
+        expectInitializeTask();
+        expectTaskGetTopic(true);
+
+        // valid json
+        ConsumerRecord record1 = new ConsumerRecord<>(TOPIC, PARTITION1, FIRST_OFFSET, null, "{\"a\": 10}".getBytes());
+        // bad json
+        ConsumerRecord record2 = new ConsumerRecord<>(TOPIC, PARTITION2, FIRST_OFFSET, null, "{\"a\" 10}".getBytes());
+
+        EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andReturn(records(record1));
+        EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andReturn(records(record2));
+
+        sinkTask.put(EasyMock.anyObject());
+        EasyMock.expectLastCall().times(2);
+
+        PowerMock.replayAll();
+
+        workerSinkTask.initialize(TASK_CONFIG);
+        workerSinkTask.initializeAndStart();
+        workerSinkTask.iteration();
+
+        workerSinkTask.iteration();
+
+        // two records were consumed from Kafka
+        assertSinkMetricValue("sink-record-read-total", 2.0);
+        // only one was written to the task
+        assertSinkMetricValue("sink-record-send-total", 1.0);
+        // one record completely failed (converter issues)
+        assertErrorHandlingMetricValue("total-record-errors", 1.0);
+        // 2 failures in the transformation, and 1 in the converter
+        assertErrorHandlingMetricValue("total-record-failures", 3.0);
+        // one record completely failed (converter issues), and thus was skipped
+        assertErrorHandlingMetricValue("total-records-skipped", 1.0);
+
+        PowerMock.verifyAll();
+    }
+
+    private RetryWithToleranceOperator operator() {
+        return new RetryWithToleranceOperator(OPERATOR_RETRY_TIMEOUT_MILLIS, OPERATOR_RETRY_MAX_DELAY_MILLIS, OPERATOR_TOLERANCE_TYPE, SYSTEM);
+    }
+
+    @Test
+    public void testErrorHandlingInSourceTasks() throws Exception {
+        Map reportProps = new HashMap<>();
+        reportProps.put(ConnectorConfig.ERRORS_LOG_ENABLE_CONFIG, "true");
+        reportProps.put(ConnectorConfig.ERRORS_LOG_INCLUDE_MESSAGES_CONFIG, "true");
+        LogReporter reporter = new LogReporter(taskId, connConfig(reportProps), errorHandlingMetrics);
+
+        RetryWithToleranceOperator retryWithToleranceOperator = operator();
+        retryWithToleranceOperator.metrics(errorHandlingMetrics);
+        retryWithToleranceOperator.reporters(singletonList(reporter));
+        createSourceTask(initialState, retryWithToleranceOperator);
+
+        // valid json
+        Schema valSchema = SchemaBuilder.struct().field("val", Schema.INT32_SCHEMA).build();
+        Struct struct1 = new Struct(valSchema).put("val", 1234);
+        SourceRecord record1 = new SourceRecord(emptyMap(), emptyMap(), TOPIC, PARTITION1, valSchema, struct1);
+        Struct struct2 = new Struct(valSchema).put("val", 6789);
+        SourceRecord record2 = new SourceRecord(emptyMap(), emptyMap(), TOPIC, PARTITION1, valSchema, struct2);
+
+        EasyMock.expect(workerSourceTask.isStopping()).andReturn(false);
+        EasyMock.expect(workerSourceTask.isStopping()).andReturn(false);
+        EasyMock.expect(workerSourceTask.isStopping()).andReturn(true);
+
+        EasyMock.expect(workerSourceTask.commitOffsets()).andReturn(true);
+
+        offsetWriter.offset(EasyMock.anyObject(), EasyMock.anyObject());
+        EasyMock.expectLastCall().times(2);
+        sourceTask.initialize(EasyMock.anyObject());
+        EasyMock.expectLastCall();
+
+        sourceTask.start(EasyMock.anyObject());
+        EasyMock.expectLastCall();
+
+        EasyMock.expect(sourceTask.poll()).andReturn(singletonList(record1));
+        EasyMock.expect(sourceTask.poll()).andReturn(singletonList(record2));
+        expectTopicDoesNotExist(TOPIC);
+        EasyMock.expect(producer.send(EasyMock.anyObject(), EasyMock.anyObject())).andReturn(null).times(2);
+
+        PowerMock.replayAll();
+
+        workerSourceTask.initialize(TASK_CONFIG);
+        workerSourceTask.execute();
+
+        // two records were consumed from Kafka
+        assertSourceMetricValue("source-record-poll-total", 2.0);
+        // only one was written to the task
+        assertSourceMetricValue("source-record-write-total", 0.0);
+        // one record completely failed (converter issues)
+        assertErrorHandlingMetricValue("total-record-errors", 0.0);
+        // 2 failures in the transformation, and 1 in the converter
+        assertErrorHandlingMetricValue("total-record-failures", 4.0);
+        // one record completely failed (converter issues), and thus was skipped
+        assertErrorHandlingMetricValue("total-records-skipped", 0.0);
+
+        PowerMock.verifyAll();
+    }
+
+    private ConnectorConfig connConfig(Map connProps) {
+        Map props = new HashMap<>();
+        props.put(ConnectorConfig.NAME_CONFIG, "test");
+        props.put(ConnectorConfig.CONNECTOR_CLASS_CONFIG, SinkTask.class.getName());
+        props.putAll(connProps);
+        return new ConnectorConfig(plugins, props);
+    }
+
+    @Test
+    public void testErrorHandlingInSourceTasksWthBadConverter() throws Exception {
+        Map reportProps = new HashMap<>();
+        reportProps.put(ConnectorConfig.ERRORS_LOG_ENABLE_CONFIG, "true");
+        reportProps.put(ConnectorConfig.ERRORS_LOG_INCLUDE_MESSAGES_CONFIG, "true");
+        LogReporter reporter = new LogReporter(taskId, connConfig(reportProps), errorHandlingMetrics);
+
+        RetryWithToleranceOperator retryWithToleranceOperator = operator();
+        retryWithToleranceOperator.metrics(errorHandlingMetrics);
+        retryWithToleranceOperator.reporters(singletonList(reporter));
+        createSourceTask(initialState, retryWithToleranceOperator, badConverter());
+
+        // valid json
+        Schema valSchema = SchemaBuilder.struct().field("val", Schema.INT32_SCHEMA).build();
+        Struct struct1 = new Struct(valSchema).put("val", 1234);
+        SourceRecord record1 = new SourceRecord(emptyMap(), emptyMap(), TOPIC, PARTITION1, valSchema, struct1);
+        Struct struct2 = new Struct(valSchema).put("val", 6789);
+        SourceRecord record2 = new SourceRecord(emptyMap(), emptyMap(), TOPIC, PARTITION1, valSchema, struct2);
+
+        EasyMock.expect(workerSourceTask.isStopping()).andReturn(false);
+        EasyMock.expect(workerSourceTask.isStopping()).andReturn(false);
+        EasyMock.expect(workerSourceTask.isStopping()).andReturn(true);
+
+        EasyMock.expect(workerSourceTask.commitOffsets()).andReturn(true);
+
+        offsetWriter.offset(EasyMock.anyObject(), EasyMock.anyObject());
+        EasyMock.expectLastCall().times(2);
+        sourceTask.initialize(EasyMock.anyObject());
+        EasyMock.expectLastCall();
+
+        sourceTask.start(EasyMock.anyObject());
+        EasyMock.expectLastCall();
+
+        EasyMock.expect(sourceTask.poll()).andReturn(singletonList(record1));
+        EasyMock.expect(sourceTask.poll()).andReturn(singletonList(record2));
+        expectTopicDoesNotExist(TOPIC);
+        EasyMock.expect(producer.send(EasyMock.anyObject(), EasyMock.anyObject())).andReturn(null).times(2);
+
+        PowerMock.replayAll();
+
+        workerSourceTask.initialize(TASK_CONFIG);
+        workerSourceTask.execute();
+
+        // two records were consumed from Kafka
+        assertSourceMetricValue("source-record-poll-total", 2.0);
+        // only one was written to the task
+        assertSourceMetricValue("source-record-write-total", 0.0);
+        // one record completely failed (converter issues)
+        assertErrorHandlingMetricValue("total-record-errors", 0.0);
+        // 2 failures in the transformation, and 1 in the converter
+        assertErrorHandlingMetricValue("total-record-failures", 8.0);
+        // one record completely failed (converter issues), and thus was skipped
+        assertErrorHandlingMetricValue("total-records-skipped", 0.0);
+
+        PowerMock.verifyAll();
+    }
+
+    private void assertSinkMetricValue(String name, double expected) {
+        ConnectMetrics.MetricGroup sinkTaskGroup = workerSinkTask.sinkTaskMetricsGroup().metricGroup();
+        double measured = metrics.currentMetricValueAsDouble(sinkTaskGroup, name);
+        assertEquals(expected, measured, 0.001d);
+    }
+
+    private void assertSourceMetricValue(String name, double expected) {
+        ConnectMetrics.MetricGroup sinkTaskGroup = workerSourceTask.sourceTaskMetricsGroup().metricGroup();
+        double measured = metrics.currentMetricValueAsDouble(sinkTaskGroup, name);
+        assertEquals(expected, measured, 0.001d);
+    }
+
+    private void assertErrorHandlingMetricValue(String name, double expected) {
+        ConnectMetrics.MetricGroup sinkTaskGroup = errorHandlingMetrics.metricGroup();
+        double measured = metrics.currentMetricValueAsDouble(sinkTaskGroup, name);
+        assertEquals(expected, measured, 0.001d);
+    }
+
+    private void expectInitializeTask() throws Exception {
+        consumer.subscribe(EasyMock.eq(singletonList(TOPIC)), EasyMock.capture(rebalanceListener));
+        PowerMock.expectLastCall();
+
+        sinkTask.initialize(EasyMock.capture(sinkTaskContext));
+        PowerMock.expectLastCall();
+        sinkTask.start(TASK_PROPS);
+        PowerMock.expectLastCall();
+    }
+
+    private void expectTaskGetTopic(boolean anyTimes) {
+        final Capture connectorCapture = EasyMock.newCapture();
+        final Capture topicCapture = EasyMock.newCapture();
+        IExpectationSetters expect = EasyMock.expect(statusBackingStore.getTopic(
+                EasyMock.capture(connectorCapture),
+                EasyMock.capture(topicCapture)));
+        if (anyTimes) {
+            expect.andStubAnswer(() -> new TopicStatus(
+                    topicCapture.getValue(),
+                    new ConnectorTaskId(connectorCapture.getValue(), 0),
+                    Time.SYSTEM.milliseconds()));
+        } else {
+            expect.andAnswer(() -> new TopicStatus(
+                    topicCapture.getValue(),
+                    new ConnectorTaskId(connectorCapture.getValue(), 0),
+                    Time.SYSTEM.milliseconds()));
+        }
+        if (connectorCapture.hasCaptured() && topicCapture.hasCaptured()) {
+            assertEquals("job", connectorCapture.getValue());
+            assertEquals(TOPIC, topicCapture.getValue());
+        }
+    }
+
+    private void expectClose() {
+        producer.close(EasyMock.anyObject(Duration.class));
+        EasyMock.expectLastCall();
+
+        admin.close(EasyMock.anyObject(Duration.class));
+        EasyMock.expectLastCall();
+    }
+
+    private void expectTopicDoesNotExist(String topic) {
+        if (workerConfig.topicCreationEnable()) {
+            EasyMock.expect(admin.describeTopics(topic)).andReturn(Collections.emptyMap());
+
+            Capture newTopicCapture = EasyMock.newCapture();
+            EasyMock.expect(admin.createTopic(EasyMock.capture(newTopicCapture))).andReturn(true);
+        }
+    }
+
+    private void createSinkTask(TargetState initialState, RetryWithToleranceOperator retryWithToleranceOperator) {
+        JsonConverter converter = new JsonConverter();
+        Map oo = workerConfig.originalsWithPrefix("value.converter.");
+        oo.put("converter.type", "value");
+        oo.put("schemas.enable", "false");
+        converter.configure(oo);
+
+        TransformationChain sinkTransforms = new TransformationChain<>(singletonList(new FaultyPassthrough()), retryWithToleranceOperator);
+
+        workerSinkTask = new WorkerSinkTask(
+            taskId, sinkTask, statusListener, initialState, workerConfig,
+            ClusterConfigState.EMPTY, metrics, converter, converter,
+            headerConverter, sinkTransforms, consumer, pluginLoader, time,
+                retryWithToleranceOperator, workerErrantRecordReporter, statusBackingStore);
+    }
+
+    private void createSourceTask(TargetState initialState, RetryWithToleranceOperator retryWithToleranceOperator) {
+        JsonConverter converter = new JsonConverter();
+        Map oo = workerConfig.originalsWithPrefix("value.converter.");
+        oo.put("converter.type", "value");
+        oo.put("schemas.enable", "false");
+        converter.configure(oo);
+
+        createSourceTask(initialState, retryWithToleranceOperator, converter);
+    }
+
+    private Converter badConverter() {
+        FaultyConverter converter = new FaultyConverter();
+        Map oo = workerConfig.originalsWithPrefix("value.converter.");
+        oo.put("converter.type", "value");
+        oo.put("schemas.enable", "false");
+        converter.configure(oo);
+        return converter;
+    }
+
+    private void createSourceTask(TargetState initialState, RetryWithToleranceOperator retryWithToleranceOperator, Converter converter) {
+        TransformationChain sourceTransforms = new TransformationChain<>(singletonList(new FaultyPassthrough()), retryWithToleranceOperator);
+
+        workerSourceTask = PowerMock.createPartialMock(
+                WorkerSourceTask.class, new String[]{"commitOffsets", "isStopping"},
+                taskId, sourceTask, statusListener, initialState, converter, converter, headerConverter, sourceTransforms,
+                producer, admin, TopicCreationGroup.configuredGroups(sourceConfig),
+                offsetReader, offsetWriter, workerConfig,
+                ClusterConfigState.EMPTY, metrics, pluginLoader, time, retryWithToleranceOperator,
+                statusBackingStore);
+    }
+
+    private ConsumerRecords records(ConsumerRecord record) {
+        return new ConsumerRecords<>(Collections.singletonMap(
+                new TopicPartition(record.topic(), record.partition()), singletonList(record)));
+    }
+
+    private abstract static class TestSinkTask extends SinkTask {
+    }
+
+    static class FaultyConverter extends JsonConverter {
+        private static final Logger log = LoggerFactory.getLogger(FaultyConverter.class);
+        private int invocations = 0;
+
+        public byte[] fromConnectData(String topic, Schema schema, Object value) {
+            if (value == null) {
+                return super.fromConnectData(topic, schema, null);
+            }
+            invocations++;
+            if (invocations % 3 == 0) {
+                log.debug("Succeeding record: {} where invocations={}", value, invocations);
+                return super.fromConnectData(topic, schema, value);
+            } else {
+                log.debug("Failing record: {} at invocations={}", value, invocations);
+                throw new RetriableException("Bad invocations " + invocations + " for mod 3");
+            }
+        }
+    }
+
+    static class FaultyPassthrough> implements Transformation {
+
+        private static final Logger log = LoggerFactory.getLogger(FaultyPassthrough.class);
+
+        private static final String MOD_CONFIG = "mod";
+        private static final int MOD_CONFIG_DEFAULT = 3;
+
+        public static final ConfigDef CONFIG_DEF = new ConfigDef()
+                .define(MOD_CONFIG, ConfigDef.Type.INT, MOD_CONFIG_DEFAULT, ConfigDef.Importance.MEDIUM, "Pass records without failure only if timestamp % mod == 0");
+
+        private int mod = MOD_CONFIG_DEFAULT;
+
+        private int invocations = 0;
+
+        @Override
+        public R apply(R record) {
+            invocations++;
+            if (invocations % mod == 0) {
+                log.debug("Succeeding record: {} where invocations={}", record, invocations);
+                return record;
+            } else {
+                log.debug("Failing record: {} at invocations={}", record, invocations);
+                throw new RetriableException("Bad invocations " + invocations + " for mod " + mod);
+            }
+        }
+
+        @Override
+        public ConfigDef config() {
+            return CONFIG_DEF;
+        }
+
+        @Override
+        public void close() {
+            log.info("Shutting down transform");
+        }
+
+        @Override
+        public void configure(Map configs) {
+            final SimpleConfig config = new SimpleConfig(CONFIG_DEF, configs);
+            mod = Math.max(config.getInt(MOD_CONFIG), 2);
+            log.info("Configuring {}. Setting mod to {}", this.getClass(), mod);
+        }
+    }
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/MockConnectMetrics.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/MockConnectMetrics.java
index f1df140f2102c..38bd4017fcb78 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/MockConnectMetrics.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/MockConnectMetrics.java
@@ -19,8 +19,9 @@
 import org.apache.kafka.clients.CommonClientConfigs;
 import org.apache.kafka.common.MetricName;
 import org.apache.kafka.common.metrics.KafkaMetric;
+import org.apache.kafka.common.metrics.MetricsContext;
 import org.apache.kafka.common.metrics.MetricsReporter;
-import org.apache.kafka.connect.util.MockTime;
+import org.apache.kafka.common.utils.MockTime;
 
 import java.util.HashMap;
 import java.util.List;
@@ -40,6 +41,7 @@
  * If the same metric is created a second time (e.g., a worker task is re-created), the new metric will replace
  * the previous metric in the custom reporter.
  */
+@SuppressWarnings("deprecation")
 public class MockConnectMetrics extends ConnectMetrics {
 
     private static final Map DEFAULT_WORKER_CONFIG = new HashMap<>();
@@ -56,12 +58,8 @@ public MockConnectMetrics() {
         this(new MockTime());
     }
 
-    public MockConnectMetrics(org.apache.kafka.common.utils.MockTime time) {
-        super("mock", new WorkerConfig(WorkerConfig.baseConfigDef(), DEFAULT_WORKER_CONFIG), time);
-    }
-
     public MockConnectMetrics(MockTime time) {
-        super("mock", new WorkerConfig(WorkerConfig.baseConfigDef(), DEFAULT_WORKER_CONFIG), time);
+        super("mock", new WorkerConfig(WorkerConfig.baseConfigDef(), DEFAULT_WORKER_CONFIG), time, "cluster-1");
     }
 
     @Override
@@ -157,6 +155,8 @@ public static String currentMetricValueAsString(ConnectMetrics metrics, MetricGr
     public static class MockMetricsReporter implements MetricsReporter {
         private Map metricsByName = new HashMap<>();
 
+        private MetricsContext metricsContext;
+
         public MockMetricsReporter() {
         }
 
@@ -197,5 +197,14 @@ public Object currentMetricValue(MetricName metricName) {
             KafkaMetric metric = metricsByName.get(metricName);
             return metric != null ? metric.metricValue() : null;
         }
+
+        @Override
+        public void contextChange(MetricsContext metricsContext) {
+            this.metricsContext = metricsContext;
+        }
+
+        public MetricsContext getMetricsContext() {
+            return this.metricsContext;
+        }
     }
 }
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/PredicatedTransformationTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/PredicatedTransformationTest.java
new file mode 100644
index 0000000000000..75542e7352881
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/PredicatedTransformationTest.java
@@ -0,0 +1,126 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.util.Map;
+
+import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.connect.source.SourceRecord;
+import org.apache.kafka.connect.transforms.Transformation;
+import org.apache.kafka.connect.transforms.predicates.Predicate;
+import org.junit.Test;
+
+import static java.util.Collections.singletonMap;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+public class PredicatedTransformationTest {
+
+    private final SourceRecord initial = new SourceRecord(singletonMap("initial", 1), null, null, null, null);
+    private final SourceRecord transformed = new SourceRecord(singletonMap("transformed", 2), null, null, null, null);
+
+    @Test
+    public void apply() {
+        applyAndAssert(true, false, transformed);
+        applyAndAssert(true, true, initial);
+        applyAndAssert(false, false, initial);
+        applyAndAssert(false, true, transformed);
+    }
+
+    private void applyAndAssert(boolean predicateResult, boolean negate,
+                                SourceRecord expectedResult) {
+        class TestTransformation implements Transformation {
+
+            private boolean closed = false;
+            private SourceRecord transformedRecord;
+
+            private TestTransformation(SourceRecord transformedRecord) {
+                this.transformedRecord = transformedRecord;
+            }
+
+            @Override
+            public SourceRecord apply(SourceRecord record) {
+                return transformedRecord;
+            }
+
+            @Override
+            public ConfigDef config() {
+                return null;
+            }
+
+            @Override
+            public void close() {
+                closed = true;
+            }
+
+            @Override
+            public void configure(Map configs) {
+
+            }
+
+            private void assertClosed() {
+                assertTrue("Transformer should be closed", closed);
+            }
+        }
+
+        class TestPredicate implements Predicate {
+
+            private boolean testResult;
+            private boolean closed = false;
+
+            private TestPredicate(boolean testResult) {
+                this.testResult = testResult;
+            }
+
+            @Override
+            public ConfigDef config() {
+                return null;
+            }
+
+            @Override
+            public boolean test(SourceRecord record) {
+                return testResult;
+            }
+
+            @Override
+            public void close() {
+                closed = true;
+            }
+
+            @Override
+            public void configure(Map configs) {
+
+            }
+
+            private void assertClosed() {
+                assertTrue("Predicate should be closed", closed);
+            }
+        }
+        TestPredicate predicate = new TestPredicate(predicateResult);
+        TestTransformation predicatedTransform = new TestTransformation(transformed);
+        PredicatedTransformation pt = new PredicatedTransformation<>(
+                predicate,
+                negate,
+                predicatedTransform);
+
+        assertEquals(expectedResult, pt.apply(initial));
+
+        pt.close();
+        predicate.assertClosed();
+        predicatedTransform.assertClosed();
+    }
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/SourceConnectorConfigTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/SourceConnectorConfigTest.java
new file mode 100644
index 0000000000000..1972b62e81113
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/SourceConnectorConfigTest.java
@@ -0,0 +1,150 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.common.config.ConfigException;
+import org.junit.Test;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.TimeUnit;
+import java.util.stream.Collectors;
+
+import static org.apache.kafka.common.config.TopicConfig.CLEANUP_POLICY_COMPACT;
+import static org.apache.kafka.common.config.TopicConfig.CLEANUP_POLICY_CONFIG;
+import static org.apache.kafka.common.config.TopicConfig.COMPRESSION_TYPE_CONFIG;
+import static org.apache.kafka.common.config.TopicConfig.RETENTION_MS_CONFIG;
+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.ConnectorConfigTest.MOCK_PLUGINS;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.DEFAULT_TOPIC_CREATION_GROUP;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.DEFAULT_TOPIC_CREATION_PREFIX;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.PARTITIONS_CONFIG;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.REPLICATION_FACTOR_CONFIG;
+import static org.hamcrest.CoreMatchers.containsString;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+
+public class SourceConnectorConfigTest {
+
+    private static final String FOO_CONNECTOR = "foo-source";
+    private static final short DEFAULT_REPLICATION_FACTOR = -1;
+    private static final int DEFAULT_PARTITIONS = -1;
+
+    public Map defaultConnectorProps() {
+        Map props = new HashMap<>();
+        props.put(NAME_CONFIG, FOO_CONNECTOR);
+        props.put(CONNECTOR_CLASS_CONFIG, ConnectorConfigTest.TestConnector.class.getName());
+        return props;
+    }
+
+    public Map defaultConnectorPropsWithTopicCreation() {
+        Map props = defaultConnectorProps();
+        props.put(DEFAULT_TOPIC_CREATION_PREFIX + REPLICATION_FACTOR_CONFIG, String.valueOf(DEFAULT_REPLICATION_FACTOR));
+        props.put(DEFAULT_TOPIC_CREATION_PREFIX + PARTITIONS_CONFIG, String.valueOf(DEFAULT_PARTITIONS));
+        return props;
+    }
+
+    @Test
+    public void noTopicCreation() {
+        Map props = defaultConnectorProps();
+        SourceConnectorConfig config = new SourceConnectorConfig(MOCK_PLUGINS, props, false);
+        assertFalse(config.usesTopicCreation());
+    }
+
+    @Test
+    public void shouldNotAllowZeroPartitionsOrReplicationFactor() {
+        Map props = defaultConnectorPropsWithTopicCreation();
+        props.put(DEFAULT_TOPIC_CREATION_PREFIX + PARTITIONS_CONFIG, String.valueOf(0));
+        Exception e = assertThrows(ConfigException.class, () -> new SourceConnectorConfig(MOCK_PLUGINS, props, true));
+        assertThat(e.getMessage(), containsString("Number of partitions must be positive, or -1"));
+
+        props.put(DEFAULT_TOPIC_CREATION_PREFIX + PARTITIONS_CONFIG, String.valueOf(DEFAULT_PARTITIONS));
+        props.put(DEFAULT_TOPIC_CREATION_PREFIX + REPLICATION_FACTOR_CONFIG, String.valueOf(0));
+
+        e = assertThrows(ConfigException.class, () -> new SourceConnectorConfig(MOCK_PLUGINS, props, true));
+        assertThat(e.getMessage(), containsString("Replication factor must be positive and not "
+                + "larger than the number of brokers in the Kafka cluster, or -1 to use the "
+                + "broker's default"));
+    }
+
+    @Test
+    public void shouldNotAllowPartitionsOrReplicationFactorLessThanNegativeOne() {
+        Map props = defaultConnectorPropsWithTopicCreation();
+        for (int i = -2; i > -100; --i) {
+            props.put(DEFAULT_TOPIC_CREATION_PREFIX + PARTITIONS_CONFIG, String.valueOf(i));
+            props.put(DEFAULT_TOPIC_CREATION_PREFIX + REPLICATION_FACTOR_CONFIG, String.valueOf(DEFAULT_REPLICATION_FACTOR));
+            Exception e = assertThrows(ConfigException.class, () -> new SourceConnectorConfig(MOCK_PLUGINS, props, true));
+            assertThat(e.getMessage(), containsString("Number of partitions must be positive, or -1"));
+
+            props.put(DEFAULT_TOPIC_CREATION_PREFIX + PARTITIONS_CONFIG, String.valueOf(DEFAULT_PARTITIONS));
+            props.put(DEFAULT_TOPIC_CREATION_PREFIX + REPLICATION_FACTOR_CONFIG, String.valueOf(i));
+            e = assertThrows(ConfigException.class, () -> new SourceConnectorConfig(MOCK_PLUGINS, props, true));
+            assertThat(e.getMessage(), containsString("Replication factor must be positive and not "
+                    + "larger than the number of brokers in the Kafka cluster, or -1 to use the "
+                    + "broker's default"));
+        }
+    }
+
+    @Test
+    public void shouldAllowNegativeOneAndPositiveForReplicationFactor() {
+        Map props = defaultConnectorPropsWithTopicCreation();
+        SourceConnectorConfig config = new SourceConnectorConfig(MOCK_PLUGINS, props, true);
+        assertTrue(config.usesTopicCreation());
+
+        for (int i = 1; i <= 100; ++i) {
+            props.put(DEFAULT_TOPIC_CREATION_PREFIX + PARTITIONS_CONFIG, String.valueOf(i));
+            props.put(DEFAULT_TOPIC_CREATION_PREFIX + REPLICATION_FACTOR_CONFIG, String.valueOf(DEFAULT_REPLICATION_FACTOR));
+            config = new SourceConnectorConfig(MOCK_PLUGINS, props, true);
+            assertTrue(config.usesTopicCreation());
+
+            props.put(DEFAULT_TOPIC_CREATION_PREFIX + PARTITIONS_CONFIG, String.valueOf(DEFAULT_PARTITIONS));
+            props.put(DEFAULT_TOPIC_CREATION_PREFIX + REPLICATION_FACTOR_CONFIG, String.valueOf(i));
+            config = new SourceConnectorConfig(MOCK_PLUGINS, props, true);
+            assertTrue(config.usesTopicCreation());
+        }
+    }
+
+    @Test
+    public void shouldAllowSettingTopicProperties() {
+        Map topicProps = new HashMap<>();
+        topicProps.put(CLEANUP_POLICY_CONFIG, CLEANUP_POLICY_COMPACT);
+        topicProps.put(COMPRESSION_TYPE_CONFIG, "lz4");
+        topicProps.put(RETENTION_MS_CONFIG, String.valueOf(TimeUnit.DAYS.toMillis(30)));
+
+        Map props = defaultConnectorPropsWithTopicCreation();
+        topicProps.forEach((k, v) -> props.put(DEFAULT_TOPIC_CREATION_PREFIX + k, v));
+
+        SourceConnectorConfig config = new SourceConnectorConfig(MOCK_PLUGINS, props, true);
+        assertEquals(topicProps,
+                convertToStringValues(config.topicCreationOtherConfigs(DEFAULT_TOPIC_CREATION_GROUP)));
+    }
+
+    private static Map convertToStringValues(Map config) {
+        // null values are not allowed
+        return config.entrySet().stream()
+                .collect(Collectors.toMap(Map.Entry::getKey, e -> {
+                    Objects.requireNonNull(e.getValue());
+                    return e.getValue().toString();
+                }));
+    }
+}
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 99f4dedd29639..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
@@ -26,11 +26,9 @@
 import org.junit.runner.RunWith;
 import org.powermock.api.easymock.PowerMock;
 import org.powermock.api.easymock.annotation.Mock;
-import org.powermock.core.classloader.annotations.PrepareForTest;
 import org.powermock.modules.junit4.PowerMockRunner;
 import org.powermock.reflect.Whitebox;
 import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
 import java.util.HashMap;
 import java.util.Map;
@@ -40,20 +38,24 @@
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
 
+import static java.util.Collections.singletonMap;
 import static org.easymock.EasyMock.eq;
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
 @RunWith(PowerMockRunner.class)
-@PrepareForTest({SourceTaskOffsetCommitter.class, LoggerFactory.class})
 public class SourceTaskOffsetCommitterTest extends ThreadedTest {
-    @Mock
-    private ScheduledExecutorService executor;
-    @Mock
-    private ConcurrentHashMap committers;
-    @Mock
-    private Logger mockLog;
+
+    private final ConcurrentHashMap> committers = new ConcurrentHashMap<>();
+
+    @Mock private ScheduledExecutorService executor;
+    @Mock private Logger mockLog;
+    @Mock private ScheduledFuture commitFuture;
+    @Mock private ScheduledFuture taskFuture;
+    @Mock private ConnectorTaskId taskId;
+    @Mock private WorkerSourceTask task;
 
     private SourceTaskOffsetCommitter committer;
 
@@ -77,26 +79,22 @@ public void setup() {
         Whitebox.setInternalState(SourceTaskOffsetCommitter.class, "log", mockLog);
     }
 
+    @SuppressWarnings("unchecked")
     @Test
-    public void testSchedule() throws Exception {
+    public void testSchedule() {
         Capture taskWrapper = EasyMock.newCapture();
 
-        ScheduledFuture commitFuture = PowerMock.createMock(ScheduledFuture.class);
         EasyMock.expect(executor.scheduleWithFixedDelay(
                 EasyMock.capture(taskWrapper), eq(DEFAULT_OFFSET_COMMIT_INTERVAL_MS),
                 eq(DEFAULT_OFFSET_COMMIT_INTERVAL_MS), eq(TimeUnit.MILLISECONDS))
-        ).andReturn(commitFuture);
-
-        ConnectorTaskId taskId = PowerMock.createMock(ConnectorTaskId.class);
-        WorkerSourceTask task = PowerMock.createMock(WorkerSourceTask.class);
-
-        EasyMock.expect(committers.put(taskId, commitFuture)).andReturn(null);
+        ).andReturn((ScheduledFuture) commitFuture);
 
         PowerMock.replayAll();
 
         committer.schedule(taskId, task);
         assertTrue(taskWrapper.hasCaptured());
         assertNotNull(taskWrapper.getValue());
+        assertEquals(singletonMap(taskId, commitFuture), committers);
 
         PowerMock.verifyAll();
     }
@@ -135,52 +133,58 @@ public void testClose() throws Exception {
 
     @Test
     public void testRemove() throws Exception {
-        ConnectorTaskId taskId = PowerMock.createMock(ConnectorTaskId.class);
-        ScheduledFuture task = PowerMock.createMock(ScheduledFuture.class);
-
         // Try to remove a non-existing task
-        EasyMock.expect(committers.remove(taskId)).andReturn(null);
         PowerMock.replayAll();
 
+        assertTrue(committers.isEmpty());
         committer.remove(taskId);
+        assertTrue(committers.isEmpty());
 
         PowerMock.verifyAll();
         PowerMock.resetAll();
 
         // Try to remove an existing task
-        EasyMock.expect(committers.remove(taskId)).andReturn(task);
-        EasyMock.expect(task.cancel(eq(false))).andReturn(false);
-        EasyMock.expect(task.isDone()).andReturn(false);
-        EasyMock.expect(task.get()).andReturn(null);
+        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);
         committer.remove(taskId);
+        assertTrue(committers.isEmpty());
 
         PowerMock.verifyAll();
         PowerMock.resetAll();
 
         // Try to remove a cancelled task
-        EasyMock.expect(committers.remove(taskId)).andReturn(task);
-        EasyMock.expect(task.cancel(eq(false))).andReturn(false);
-        EasyMock.expect(task.isDone()).andReturn(false);
-        EasyMock.expect(task.get()).andThrow(new CancellationException());
+        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();
 
+        committers.put(taskId, taskFuture);
         committer.remove(taskId);
+        assertTrue(committers.isEmpty());
 
         PowerMock.verifyAll();
         PowerMock.resetAll();
 
         // Try to remove an interrupted task
-        EasyMock.expect(committers.remove(taskId)).andReturn(task);
-        EasyMock.expect(task.cancel(eq(false))).andReturn(false);
-        EasyMock.expect(task.isDone()).andReturn(false);
-        EasyMock.expect(task.get()).andThrow(new InterruptedException());
+        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 {
+            committers.put(taskId, taskFuture);
             committer.remove(taskId);
             fail("Expected ConnectException to be raised");
         } catch (ConnectException e) {
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/StateTrackerTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/StateTrackerTest.java
index 7423854625586..2d9f02c26ba51 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/StateTrackerTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/StateTrackerTest.java
@@ -17,7 +17,7 @@
 package org.apache.kafka.connect.runtime;
 
 import org.apache.kafka.connect.runtime.AbstractStatus.State;
-import org.apache.kafka.connect.util.MockTime;
+import org.apache.kafka.common.utils.MockTime;
 import org.junit.Before;
 import org.junit.Test;
 
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/TransformationConfigTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/TransformationConfigTest.java
index 5cd0bcb197cba..1d63d7db55f34 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/TransformationConfigTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/TransformationConfigTest.java
@@ -120,7 +120,7 @@ public void testEmbeddedConfigMaskField() {
         connProps.put("transforms", "example");
         connProps.put("transforms.example.type", MaskField.Value.class.getName());
         connProps.put("transforms.example.fields", "field");
-
+        connProps.put("transforms.example.replacement", "nothing");
 
         Plugins plugins = null; // Safe when we're only constructing the config
         new ConnectorConfig(plugins, connProps);
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..1eeb13ea2d5d2
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerConfigTest.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;
+
+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 java.util.List;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.assertThrows;
+
+public class WorkerConfigTest {
+    private static final List VALID_HEADER_CONFIGS = Arrays.asList(
+            "add \t Cache-Control: no-cache, no-store, must-revalidate",
+            "add \r X-XSS-Protection: 1; mode=block",
+            "\n add Strict-Transport-Security: max-age=31536000; includeSubDomains",
+            "AdD   Strict-Transport-Security:  \r  max-age=31536000;  includeSubDomains",
+            "AdD \t Strict-Transport-Security : \n   max-age=31536000;  includeSubDomains",
+            "add X-Content-Type-Options: \r nosniff",
+            "Set \t X-Frame-Options: \t Deny\n ",
+            "seT \t X-Cache-Info: \t not cacheable\n ",
+            "seTDate \t Expires: \r 31540000000",
+            "adDdate \n Last-Modified: \t 0"
+    );
+
+    private static final List INVALID_HEADER_CONFIGS = Arrays.asList(
+            "set \t",
+            "badaction \t X-Frame-Options:DENY",
+            "set add X-XSS-Protection:1",
+            "addX-XSS-Protection",
+            "X-XSS-Protection:",
+            "add set X-XSS-Protection: 1",
+            "add X-XSS-Protection:1 X-XSS-Protection:1 ",
+            "add X-XSS-Protection",
+            "set X-Frame-Options:DENY, add  :no-cache, no-store, must-revalidate "
+    );
+
+    @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);
+    }
+
+    @Test
+    public void testInvalidHeaderConfigs() {
+        for (String config : INVALID_HEADER_CONFIGS) {
+            assertInvalidHeaderConfig(config);
+        }
+    }
+
+    @Test
+    public void testValidHeaderConfigs() {
+        for (String config : VALID_HEADER_CONFIGS) {
+            assertValidHeaderConfig(config);
+        }
+    }
+
+    private void assertInvalidHeaderConfig(String config) {
+        assertThrows(ConfigException.class, () -> WorkerConfig.validateHttpResponseHeaderConfig(config));
+    }
+
+    private void assertValidHeaderConfig(String config) {
+        WorkerConfig.validateHttpResponseHeaderConfig(config);
+    }
+
+    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
new file mode 100644
index 0000000000000..6f4bda66904d7
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerConfigTransformerTest.java
@@ -0,0 +1,154 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.common.config.ConfigChangeCallback;
+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;
+import org.powermock.api.easymock.PowerMock;
+import org.powermock.api.easymock.annotation.Mock;
+import org.powermock.modules.junit4.PowerMockRunner;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Set;
+
+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;
+
+@RunWith(PowerMockRunner.class)
+public class WorkerConfigTransformerTest {
+
+    public static final String MY_KEY = "myKey";
+    public static final String MY_CONNECTOR = "myConnector";
+    public static final String TEST_KEY = "testKey";
+    public static final String TEST_PATH = "testPath";
+    public static final String TEST_KEY_WITH_TTL = "testKeyWithTTL";
+    public static final String TEST_KEY_WITH_LONGER_TTL = "testKeyWithLongerTTL";
+    public static final String TEST_RESULT = "testResult";
+    public static final String TEST_RESULT_WITH_TTL = "testResultWithTTL";
+    public static final String TEST_RESULT_WITH_LONGER_TTL = "testResultWithLongerTTL";
+
+    @Mock private Herder herder;
+    @Mock private Worker worker;
+    @Mock private HerderRequest requestId;
+    private WorkerConfigTransformer configTransformer;
+
+    @Before
+    public void setup() {
+        worker = PowerMock.createMock(Worker.class);
+        herder = PowerMock.createMock(Herder.class);
+        configTransformer = new WorkerConfigTransformer(worker, Collections.singletonMap("test", new TestConfigProvider()));
+    }
+
+    @Test
+    public void testReplaceVariable() {
+        Map result = configTransformer.transform(MY_CONNECTOR, Collections.singletonMap(MY_KEY, "${test:testPath:testKey}"));
+        assertEquals(TEST_RESULT, result.get(MY_KEY));
+    }
+
+    @Test
+    public void testReplaceVariableWithTTL() {
+        EasyMock.expect(worker.herder()).andReturn(herder);
+
+        replayAll();
+
+        Map props = new HashMap<>();
+        props.put(MY_KEY, "${test:testPath:testKeyWithTTL}");
+        props.put(CONFIG_RELOAD_ACTION_CONFIG, CONFIG_RELOAD_ACTION_NONE);
+        Map result = configTransformer.transform(MY_CONNECTOR, props);
+    }
+
+    @Test
+    public void testReplaceVariableWithTTLAndScheduleRestart() {
+        EasyMock.expect(worker.herder()).andReturn(herder);
+        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}"));
+        assertEquals(TEST_RESULT_WITH_TTL, result.get(MY_KEY));
+    }
+
+    @Test
+    public void testReplaceVariableWithTTLFirstCancelThenScheduleRestart() {
+        EasyMock.expect(worker.herder()).andReturn(herder);
+        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(eq(10L), eq(MY_CONNECTOR), notNull())).andReturn(requestId);
+
+        replayAll();
+
+        Map result = configTransformer.transform(MY_CONNECTOR, Collections.singletonMap(MY_KEY, "${test:testPath:testKeyWithTTL}"));
+        assertEquals(TEST_RESULT_WITH_TTL, result.get(MY_KEY));
+
+        result = configTransformer.transform(MY_CONNECTOR, Collections.singletonMap(MY_KEY, "${test:testPath:testKeyWithLongerTTL}"));
+        assertEquals(TEST_RESULT_WITH_LONGER_TTL, result.get(MY_KEY));
+    }
+
+    @Test
+    public void testTransformNullConfiguration() {
+        assertNull(configTransformer.transform(MY_CONNECTOR, null));
+    }
+
+    public static class TestConfigProvider implements ConfigProvider {
+
+        public void configure(Map configs) {
+        }
+
+        public ConfigData get(String path) {
+            return null;
+        }
+
+        public ConfigData get(String path, Set keys) {
+            if (path.equals(TEST_PATH)) {
+                if (keys.contains(TEST_KEY)) {
+                    return new ConfigData(Collections.singletonMap(TEST_KEY, TEST_RESULT));
+                } else if (keys.contains(TEST_KEY_WITH_TTL)) {
+                    return new ConfigData(Collections.singletonMap(TEST_KEY_WITH_TTL, TEST_RESULT_WITH_TTL), 1L);
+                } else if (keys.contains(TEST_KEY_WITH_LONGER_TTL)) {
+                    return new ConfigData(Collections.singletonMap(TEST_KEY_WITH_LONGER_TTL, TEST_RESULT_WITH_LONGER_TTL), 10L);
+                }
+            }
+            return new ConfigData(Collections.emptyMap());
+        }
+
+        public void subscribe(String path, Set keys, ConfigChangeCallback callback) {
+            throw new UnsupportedOperationException();
+        }
+
+        public void unsubscribe(String path, Set keys) {
+            throw new UnsupportedOperationException();
+        }
+
+        public void close() {
+        }
+    }
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerConnectorTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerConnectorTest.java
index 10c413d23d600..f99b4c1067aff 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerConnectorTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerConnectorTest.java
@@ -17,9 +17,16 @@
 package org.apache.kafka.connect.runtime;
 
 import org.apache.kafka.connect.connector.Connector;
-import org.apache.kafka.connect.connector.ConnectorContext;
+import org.apache.kafka.connect.errors.ConnectException;
 import org.apache.kafka.connect.runtime.ConnectMetrics.MetricGroup;
 import org.apache.kafka.connect.runtime.isolation.Plugins;
+import org.apache.kafka.connect.sink.SinkConnector;
+import org.apache.kafka.connect.sink.SinkConnectorContext;
+import org.apache.kafka.connect.source.SourceConnector;
+import org.apache.kafka.connect.source.SourceConnectorContext;
+import org.apache.kafka.connect.storage.OffsetStorageReader;
+import org.easymock.Capture;
+import org.apache.kafka.connect.util.Callback;
 import org.easymock.EasyMock;
 import org.easymock.EasyMockRunner;
 import org.easymock.EasyMockSupport;
@@ -47,14 +54,19 @@ public class WorkerConnectorTest extends EasyMockSupport {
     static {
         CONFIG.put(ConnectorConfig.CONNECTOR_CLASS_CONFIG, TestConnector.class.getName());
         CONFIG.put(ConnectorConfig.NAME_CONFIG, CONNECTOR);
+        CONFIG.put(SinkConnectorConfig.TOPICS_CONFIG, "my-topic");
     }
     public ConnectorConfig connectorConfig;
     public MockConnectMetrics metrics;
 
     @Mock Plugins plugins;
+    @Mock SourceConnector sourceConnector;
+    @Mock SinkConnector sinkConnector;
     @Mock Connector connector;
-    @Mock ConnectorContext ctx;
+    @Mock CloseableConnectorContext ctx;
     @Mock ConnectorStatus.Listener listener;
+    @Mock OffsetStorageReader offsetStorageReader;
+    @Mock ClassLoader classLoader;
 
     @Before
     public void setup() {
@@ -68,13 +80,14 @@ public void tearDown() {
     }
 
     @Test
-    public void testInitializeFailure() {
+    public void testInitializeFailure() throws InterruptedException {
         RuntimeException exception = new RuntimeException();
+        connector = sourceConnector;
 
         connector.version();
         expectLastCall().andReturn(VERSION);
 
-        connector.initialize(EasyMock.notNull(ConnectorContext.class));
+        connector.initialize(EasyMock.notNull(SourceConnectorContext.class));
         expectLastCall().andThrow(exception);
 
         listener.onFailure(CONNECTOR, exception);
@@ -83,13 +96,17 @@ public void testInitializeFailure() {
         listener.onShutdown(CONNECTOR);
         expectLastCall();
 
+        ctx.close();
+        expectLastCall();
+
         replayAll();
 
-        WorkerConnector workerConnector = new WorkerConnector(CONNECTOR, connector, ctx, metrics, listener);
+        WorkerConnector workerConnector = new WorkerConnector(CONNECTOR, connector, connectorConfig, ctx, metrics, listener, offsetStorageReader, classLoader);
 
-        workerConnector.initialize(connectorConfig);
+        workerConnector.initialize();
         assertFailedMetric(workerConnector);
         workerConnector.shutdown();
+        workerConnector.doShutdown();
         assertStoppedMetric(workerConnector);
 
         verifyAll();
@@ -98,11 +115,12 @@ public void testInitializeFailure() {
     @Test
     public void testFailureIsFinalState() {
         RuntimeException exception = new RuntimeException();
+        connector = sinkConnector;
 
         connector.version();
         expectLastCall().andReturn(VERSION);
 
-        connector.initialize(EasyMock.notNull(ConnectorContext.class));
+        connector.initialize(EasyMock.notNull(SinkConnectorContext.class));
         expectLastCall().andThrow(exception);
 
         listener.onFailure(CONNECTOR, exception);
@@ -113,15 +131,23 @@ public void testFailureIsFinalState() {
         listener.onShutdown(CONNECTOR);
         expectLastCall();
 
+        ctx.close();
+        expectLastCall();
+
+        Callback onStateChange = createStrictMock(Callback.class);
+        onStateChange.onCompletion(EasyMock.anyObject(Exception.class), EasyMock.isNull());
+        expectLastCall();
+
         replayAll();
 
-        WorkerConnector workerConnector = new WorkerConnector(CONNECTOR, connector, ctx, metrics, listener);
+        WorkerConnector workerConnector = new WorkerConnector(CONNECTOR, connector, connectorConfig, ctx, metrics, listener, offsetStorageReader, classLoader);
 
-        workerConnector.initialize(connectorConfig);
+        workerConnector.initialize();
         assertFailedMetric(workerConnector);
-        workerConnector.transitionTo(TargetState.STARTED);
+        workerConnector.doTransitionTo(TargetState.STARTED, onStateChange);
         assertFailedMetric(workerConnector);
         workerConnector.shutdown();
+        workerConnector.doShutdown();
         assertStoppedMetric(workerConnector);
 
         verifyAll();
@@ -129,10 +155,11 @@ public void testFailureIsFinalState() {
 
     @Test
     public void testStartupAndShutdown() {
+        connector = sourceConnector;
         connector.version();
         expectLastCall().andReturn(VERSION);
 
-        connector.initialize(EasyMock.notNull(ConnectorContext.class));
+        connector.initialize(EasyMock.notNull(SourceConnectorContext.class));
         expectLastCall();
 
         connector.start(CONFIG);
@@ -147,15 +174,23 @@ public void testStartupAndShutdown() {
         listener.onShutdown(CONNECTOR);
         expectLastCall();
 
+        ctx.close();
+        expectLastCall();
+
+        Callback onStateChange = createStrictMock(Callback.class);
+        onStateChange.onCompletion(EasyMock.isNull(), EasyMock.eq(TargetState.STARTED));
+        expectLastCall();
+
         replayAll();
 
-        WorkerConnector workerConnector = new WorkerConnector(CONNECTOR, connector, ctx, metrics, listener);
+        WorkerConnector workerConnector = new WorkerConnector(CONNECTOR, connector, connectorConfig, ctx, metrics, listener, offsetStorageReader, classLoader);
 
-        workerConnector.initialize(connectorConfig);
-        assertInitializedMetric(workerConnector);
-        workerConnector.transitionTo(TargetState.STARTED);
+        workerConnector.initialize();
+        assertInitializedSourceMetric(workerConnector);
+        workerConnector.doTransitionTo(TargetState.STARTED, onStateChange);
         assertRunningMetric(workerConnector);
         workerConnector.shutdown();
+        workerConnector.doShutdown();
         assertStoppedMetric(workerConnector);
 
         verifyAll();
@@ -163,10 +198,11 @@ public void testStartupAndShutdown() {
 
     @Test
     public void testStartupAndPause() {
+        connector = sinkConnector;
         connector.version();
         expectLastCall().andReturn(VERSION);
 
-        connector.initialize(EasyMock.notNull(ConnectorContext.class));
+        connector.initialize(EasyMock.notNull(SinkConnectorContext.class));
         expectLastCall();
 
         connector.start(CONFIG);
@@ -184,17 +220,27 @@ public void testStartupAndPause() {
         listener.onShutdown(CONNECTOR);
         expectLastCall();
 
+        ctx.close();
+        expectLastCall();
+
+        Callback onStateChange = createStrictMock(Callback.class);
+        onStateChange.onCompletion(EasyMock.isNull(), EasyMock.eq(TargetState.STARTED));
+        expectLastCall();
+        onStateChange.onCompletion(EasyMock.isNull(), EasyMock.eq(TargetState.PAUSED));
+        expectLastCall();
+
         replayAll();
 
-        WorkerConnector workerConnector = new WorkerConnector(CONNECTOR, connector, ctx, metrics, listener);
+        WorkerConnector workerConnector = new WorkerConnector(CONNECTOR, connector, connectorConfig, ctx, metrics, listener, offsetStorageReader, classLoader);
 
-        workerConnector.initialize(connectorConfig);
-        assertInitializedMetric(workerConnector);
-        workerConnector.transitionTo(TargetState.STARTED);
+        workerConnector.initialize();
+        assertInitializedSinkMetric(workerConnector);
+        workerConnector.doTransitionTo(TargetState.STARTED, onStateChange);
         assertRunningMetric(workerConnector);
-        workerConnector.transitionTo(TargetState.PAUSED);
+        workerConnector.doTransitionTo(TargetState.PAUSED, onStateChange);
         assertPausedMetric(workerConnector);
         workerConnector.shutdown();
+        workerConnector.doShutdown();
         assertStoppedMetric(workerConnector);
 
         verifyAll();
@@ -202,10 +248,11 @@ public void testStartupAndPause() {
 
     @Test
     public void testOnResume() {
+        connector = sourceConnector;
         connector.version();
         expectLastCall().andReturn(VERSION);
 
-        connector.initialize(EasyMock.notNull(ConnectorContext.class));
+        connector.initialize(EasyMock.notNull(SourceConnectorContext.class));
         expectLastCall();
 
         listener.onPause(CONNECTOR);
@@ -223,17 +270,27 @@ public void testOnResume() {
         listener.onShutdown(CONNECTOR);
         expectLastCall();
 
+        ctx.close();
+        expectLastCall();
+
+        Callback onStateChange = createStrictMock(Callback.class);
+        onStateChange.onCompletion(EasyMock.isNull(), EasyMock.eq(TargetState.PAUSED));
+        expectLastCall();
+        onStateChange.onCompletion(EasyMock.isNull(), EasyMock.eq(TargetState.STARTED));
+        expectLastCall();
+
         replayAll();
 
-        WorkerConnector workerConnector = new WorkerConnector(CONNECTOR, connector, ctx, metrics, listener);
+        WorkerConnector workerConnector = new WorkerConnector(CONNECTOR, connector, connectorConfig, ctx, metrics, listener, offsetStorageReader, classLoader);
 
-        workerConnector.initialize(connectorConfig);
-        assertInitializedMetric(workerConnector);
-        workerConnector.transitionTo(TargetState.PAUSED);
+        workerConnector.initialize();
+        assertInitializedSourceMetric(workerConnector);
+        workerConnector.doTransitionTo(TargetState.PAUSED, onStateChange);
         assertPausedMetric(workerConnector);
-        workerConnector.transitionTo(TargetState.STARTED);
+        workerConnector.doTransitionTo(TargetState.STARTED, onStateChange);
         assertRunningMetric(workerConnector);
         workerConnector.shutdown();
+        workerConnector.doShutdown();
         assertStoppedMetric(workerConnector);
 
         verifyAll();
@@ -241,10 +298,11 @@ public void testOnResume() {
 
     @Test
     public void testStartupPaused() {
+        connector = sinkConnector;
         connector.version();
         expectLastCall().andReturn(VERSION);
 
-        connector.initialize(EasyMock.notNull(ConnectorContext.class));
+        connector.initialize(EasyMock.notNull(SinkConnectorContext.class));
         expectLastCall();
 
         // connector never gets started
@@ -255,15 +313,23 @@ public void testStartupPaused() {
         listener.onShutdown(CONNECTOR);
         expectLastCall();
 
+        ctx.close();
+        expectLastCall();
+
+        Callback onStateChange = createStrictMock(Callback.class);
+        onStateChange.onCompletion(EasyMock.isNull(), EasyMock.eq(TargetState.PAUSED));
+        expectLastCall();
+
         replayAll();
 
-        WorkerConnector workerConnector = new WorkerConnector(CONNECTOR, connector, ctx, metrics, listener);
+        WorkerConnector workerConnector = new WorkerConnector(CONNECTOR, connector, connectorConfig, ctx, metrics, listener, offsetStorageReader, classLoader);
 
-        workerConnector.initialize(connectorConfig);
-        assertInitializedMetric(workerConnector);
-        workerConnector.transitionTo(TargetState.PAUSED);
+        workerConnector.initialize();
+        assertInitializedSinkMetric(workerConnector);
+        workerConnector.doTransitionTo(TargetState.PAUSED, onStateChange);
         assertPausedMetric(workerConnector);
         workerConnector.shutdown();
+        workerConnector.doShutdown();
         assertStoppedMetric(workerConnector);
 
         verifyAll();
@@ -273,10 +339,11 @@ public void testStartupPaused() {
     public void testStartupFailure() {
         RuntimeException exception = new RuntimeException();
 
+        connector = sinkConnector;
         connector.version();
         expectLastCall().andReturn(VERSION);
 
-        connector.initialize(EasyMock.notNull(ConnectorContext.class));
+        connector.initialize(EasyMock.notNull(SinkConnectorContext.class));
         expectLastCall();
 
         connector.start(CONFIG);
@@ -288,15 +355,23 @@ public void testStartupFailure() {
         listener.onShutdown(CONNECTOR);
         expectLastCall();
 
+        ctx.close();
+        expectLastCall();
+
+        Callback onStateChange = createStrictMock(Callback.class);
+        onStateChange.onCompletion(EasyMock.anyObject(Exception.class), EasyMock.isNull());
+        expectLastCall();
+
         replayAll();
 
-        WorkerConnector workerConnector = new WorkerConnector(CONNECTOR, connector, ctx, metrics, listener);
+        WorkerConnector workerConnector = new WorkerConnector(CONNECTOR, connector, connectorConfig, ctx, metrics, listener, offsetStorageReader, classLoader);
 
-        workerConnector.initialize(connectorConfig);
-        assertInitializedMetric(workerConnector);
-        workerConnector.transitionTo(TargetState.STARTED);
+        workerConnector.initialize();
+        assertInitializedSinkMetric(workerConnector);
+        workerConnector.doTransitionTo(TargetState.STARTED, onStateChange);
         assertFailedMetric(workerConnector);
         workerConnector.shutdown();
+        workerConnector.doShutdown();
         assertStoppedMetric(workerConnector);
 
         verifyAll();
@@ -305,11 +380,12 @@ public void testStartupFailure() {
     @Test
     public void testShutdownFailure() {
         RuntimeException exception = new RuntimeException();
+        connector = sourceConnector;
 
         connector.version();
         expectLastCall().andReturn(VERSION);
 
-        connector.initialize(EasyMock.notNull(ConnectorContext.class));
+        connector.initialize(EasyMock.notNull(SourceConnectorContext.class));
         expectLastCall();
 
         connector.start(CONFIG);
@@ -321,18 +397,26 @@ public void testShutdownFailure() {
         connector.stop();
         expectLastCall().andThrow(exception);
 
+        Callback onStateChange = createStrictMock(Callback.class);
+        onStateChange.onCompletion(EasyMock.isNull(), EasyMock.eq(TargetState.STARTED));
+        expectLastCall();
+
         listener.onFailure(CONNECTOR, exception);
         expectLastCall();
 
+        ctx.close();
+        expectLastCall();
+
         replayAll();
 
-        WorkerConnector workerConnector = new WorkerConnector(CONNECTOR, connector, ctx, metrics, listener);
+        WorkerConnector workerConnector = new WorkerConnector(CONNECTOR, connector, connectorConfig, ctx, metrics, listener, offsetStorageReader, classLoader);
 
-        workerConnector.initialize(connectorConfig);
-        assertInitializedMetric(workerConnector);
-        workerConnector.transitionTo(TargetState.STARTED);
+        workerConnector.initialize();
+        assertInitializedSourceMetric(workerConnector);
+        workerConnector.doTransitionTo(TargetState.STARTED, onStateChange);
         assertRunningMetric(workerConnector);
         workerConnector.shutdown();
+        workerConnector.doShutdown();
         assertFailedMetric(workerConnector);
 
         verifyAll();
@@ -340,10 +424,11 @@ public void testShutdownFailure() {
 
     @Test
     public void testTransitionStartedToStarted() {
+        connector = sourceConnector;
         connector.version();
         expectLastCall().andReturn(VERSION);
 
-        connector.initialize(EasyMock.notNull(ConnectorContext.class));
+        connector.initialize(EasyMock.notNull(SourceConnectorContext.class));
         expectLastCall();
 
         connector.start(CONFIG);
@@ -359,17 +444,25 @@ public void testTransitionStartedToStarted() {
         listener.onShutdown(CONNECTOR);
         expectLastCall();
 
+        ctx.close();
+        expectLastCall();
+
+        Callback onStateChange = createStrictMock(Callback.class);
+        onStateChange.onCompletion(EasyMock.isNull(), EasyMock.eq(TargetState.STARTED));
+        expectLastCall().times(2);
+
         replayAll();
 
-        WorkerConnector workerConnector = new WorkerConnector(CONNECTOR, connector, ctx, metrics, listener);
+        WorkerConnector workerConnector = new WorkerConnector(CONNECTOR, connector, connectorConfig, ctx, metrics, listener, offsetStorageReader, classLoader);
 
-        workerConnector.initialize(connectorConfig);
-        assertInitializedMetric(workerConnector);
-        workerConnector.transitionTo(TargetState.STARTED);
+        workerConnector.initialize();
+        assertInitializedSourceMetric(workerConnector);
+        workerConnector.doTransitionTo(TargetState.STARTED, onStateChange);
         assertRunningMetric(workerConnector);
-        workerConnector.transitionTo(TargetState.STARTED);
+        workerConnector.doTransitionTo(TargetState.STARTED, onStateChange);
         assertRunningMetric(workerConnector);
         workerConnector.shutdown();
+        workerConnector.doShutdown();
         assertStoppedMetric(workerConnector);
 
         verifyAll();
@@ -377,10 +470,11 @@ public void testTransitionStartedToStarted() {
 
     @Test
     public void testTransitionPausedToPaused() {
+        connector = sourceConnector;
         connector.version();
         expectLastCall().andReturn(VERSION);
 
-        connector.initialize(EasyMock.notNull(ConnectorContext.class));
+        connector.initialize(EasyMock.notNull(SourceConnectorContext.class));
         expectLastCall();
 
         connector.start(CONFIG);
@@ -398,24 +492,55 @@ public void testTransitionPausedToPaused() {
         listener.onShutdown(CONNECTOR);
         expectLastCall();
 
+        ctx.close();
+        expectLastCall();
+
+        Callback onStateChange = createStrictMock(Callback.class);
+        onStateChange.onCompletion(EasyMock.isNull(), EasyMock.eq(TargetState.STARTED));
+        expectLastCall();
+        onStateChange.onCompletion(EasyMock.isNull(), EasyMock.eq(TargetState.PAUSED));
+        expectLastCall().times(2);
+
         replayAll();
 
-        WorkerConnector workerConnector = new WorkerConnector(CONNECTOR, connector, ctx, metrics, listener);
+        WorkerConnector workerConnector = new WorkerConnector(CONNECTOR, connector, connectorConfig, ctx, metrics, listener, offsetStorageReader, classLoader);
 
-        workerConnector.initialize(connectorConfig);
-        assertInitializedMetric(workerConnector);
-        workerConnector.transitionTo(TargetState.STARTED);
+        workerConnector.initialize();
+        assertInitializedSourceMetric(workerConnector);
+        workerConnector.doTransitionTo(TargetState.STARTED, onStateChange);
         assertRunningMetric(workerConnector);
-        workerConnector.transitionTo(TargetState.PAUSED);
+        workerConnector.doTransitionTo(TargetState.PAUSED, onStateChange);
         assertPausedMetric(workerConnector);
-        workerConnector.transitionTo(TargetState.PAUSED);
+        workerConnector.doTransitionTo(TargetState.PAUSED, onStateChange);
         assertPausedMetric(workerConnector);
         workerConnector.shutdown();
+        workerConnector.doShutdown();
         assertStoppedMetric(workerConnector);
 
         verifyAll();
     }
 
+    @Test
+    public void testFailConnectorThatIsNeitherSourceNorSink() {
+        connector.version();
+        expectLastCall().andReturn(VERSION);
+
+        Capture exceptionCapture = Capture.newInstance();
+        listener.onFailure(EasyMock.eq(CONNECTOR), EasyMock.capture(exceptionCapture));
+        expectLastCall();
+
+        replayAll();
+
+        WorkerConnector workerConnector = new WorkerConnector(CONNECTOR, connector, connectorConfig, ctx, metrics, listener, offsetStorageReader, classLoader);
+
+        workerConnector.initialize();
+        Throwable e = exceptionCapture.getValue();
+        assertTrue(e instanceof ConnectException);
+        assertTrue(e.getMessage().contains("must be a subclass of"));
+
+        verifyAll();
+    }
+
     protected void assertFailedMetric(WorkerConnector workerConnector) {
         assertFalse(workerConnector.metrics().isUnassigned());
         assertTrue(workerConnector.metrics().isFailed());
@@ -444,7 +569,15 @@ protected void assertStoppedMetric(WorkerConnector workerConnector) {
         assertFalse(workerConnector.metrics().isRunning());
     }
 
-    protected void assertInitializedMetric(WorkerConnector workerConnector) {
+    protected void assertInitializedSinkMetric(WorkerConnector workerConnector) {
+        assertInitializedMetric(workerConnector, "sink");
+    }
+
+    protected void assertInitializedSourceMetric(WorkerConnector workerConnector) {
+        assertInitializedMetric(workerConnector, "source");
+    }
+
+    protected void assertInitializedMetric(WorkerConnector workerConnector, String expectedType) {
         assertTrue(workerConnector.metrics().isUnassigned());
         assertFalse(workerConnector.metrics().isFailed());
         assertFalse(workerConnector.metrics().isPaused());
@@ -454,12 +587,11 @@ protected void assertInitializedMetric(WorkerConnector workerConnector) {
         String type = metrics.currentMetricValueAsString(metricGroup, "connector-type");
         String clazz = metrics.currentMetricValueAsString(metricGroup, "connector-class");
         String version = metrics.currentMetricValueAsString(metricGroup, "connector-version");
-        assertEquals(type, "unknown");
+        assertEquals(expectedType, type);
         assertNotNull(clazz);
         assertEquals(VERSION, version);
     }
 
     private static abstract class TestConnector extends Connector {
     }
-
 }
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 48d87405dc12b..c2beb666b6fce 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,32 +16,47 @@
  */
 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;
 import org.apache.kafka.clients.consumer.KafkaConsumer;
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.clients.consumer.OffsetCommitCallback;
+import org.apache.kafka.common.MetricName;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.errors.WakeupException;
+import org.apache.kafka.common.header.Header;
+import org.apache.kafka.common.header.Headers;
+import org.apache.kafka.common.header.internals.RecordHeaders;
 import org.apache.kafka.common.record.RecordBatch;
 import org.apache.kafka.common.record.TimestampType;
+import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.connect.data.Schema;
 import org.apache.kafka.connect.data.SchemaAndValue;
+import org.apache.kafka.connect.errors.ConnectException;
 import org.apache.kafka.connect.errors.RetriableException;
 import org.apache.kafka.connect.runtime.ConnectMetrics.MetricGroup;
+import org.apache.kafka.connect.runtime.distributed.ClusterConfigState;
+import org.apache.kafka.connect.runtime.WorkerSinkTask.SinkTaskMetricsGroup;
+import org.apache.kafka.connect.runtime.errors.RetryWithToleranceOperatorTest;
 import org.apache.kafka.connect.runtime.isolation.PluginClassLoader;
 import org.apache.kafka.connect.runtime.standalone.StandaloneConfig;
 import org.apache.kafka.connect.sink.SinkConnector;
 import org.apache.kafka.connect.sink.SinkRecord;
 import org.apache.kafka.connect.sink.SinkTask;
 import org.apache.kafka.connect.storage.Converter;
+import org.apache.kafka.connect.storage.HeaderConverter;
+import org.apache.kafka.connect.storage.StatusBackingStore;
+import org.apache.kafka.connect.storage.StringConverter;
 import org.apache.kafka.connect.util.ConnectorTaskId;
-import org.apache.kafka.connect.util.MockTime;
 import org.easymock.Capture;
 import org.easymock.CaptureType;
 import org.easymock.EasyMock;
 import org.easymock.IAnswer;
+import org.easymock.IExpectationSetters;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
@@ -53,6 +68,7 @@
 import org.powermock.modules.junit4.PowerMockRunner;
 import org.powermock.reflect.Whitebox;
 
+import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
@@ -71,9 +87,11 @@
 
 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;
+import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
@@ -107,6 +125,7 @@ public class WorkerSinkTaskTest {
     private static final TaskConfig TASK_CONFIG = new TaskConfig(TASK_PROPS);
 
     private ConnectorTaskId taskId = new ConnectorTaskId("job", 0);
+    private ConnectorTaskId taskId1 = new ConnectorTaskId("job", 1);
     private TargetState initialState = TargetState.STARTED;
     private MockTime time;
     private WorkerSinkTask workerTask;
@@ -122,10 +141,14 @@ public class WorkerSinkTaskTest {
     @Mock
     private Converter valueConverter;
     @Mock
+    private HeaderConverter headerConverter;
+    @Mock
     private TransformationChain transformationChain;
     @Mock
     private TaskStatus.Listener statusListener;
     @Mock
+    private StatusBackingStore statusBackingStore;
+    @Mock
     private KafkaConsumer consumer;
     private Capture rebalanceListener = EasyMock.newCapture();
     private Capture topicsRegex = EasyMock.newCapture();
@@ -152,9 +175,15 @@ public void setUp() {
     }
 
     private void createTask(TargetState initialState) {
-        workerTask = PowerMock.createPartialMock(
-                WorkerSinkTask.class, new String[]{"createConsumer"},
-                taskId, sinkTask, statusListener, initialState, workerConfig, metrics, keyConverter, valueConverter, transformationChain, pluginLoader, time);
+        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,
+            transformationChain, consumer, pluginLoader, time,
+            RetryWithToleranceOperatorTest.NOOP_OPERATOR, null, statusBackingStore);
     }
 
     @After
@@ -167,6 +196,7 @@ public void testStartPaused() throws Exception {
         createTask(TargetState.PAUSED);
 
         expectInitializeTask();
+        expectTaskGetTopic(true);
         expectPollInitialAssignment();
 
         Set partitions = new HashSet<>(asList(TOPIC_PARTITION, TOPIC_PARTITION2));
@@ -185,7 +215,7 @@ public void testStartPaused() throws Exception {
         assertTaskMetricValue("status", "paused");
         assertTaskMetricValue("running-ratio", 0.0);
         assertTaskMetricValue("pause-ratio", 1.0);
-        assertTaskMetricValue("offset-commit-max-time-ms", Double.NEGATIVE_INFINITY);
+        assertTaskMetricValue("offset-commit-max-time-ms", Double.NaN);
 
         PowerMock.verifyAll();
     }
@@ -195,6 +225,7 @@ public void testPause() throws Exception {
         createTask(initialState);
 
         expectInitializeTask();
+        expectTaskGetTopic(true);
         expectPollInitialAssignment();
 
         expectConsumerPoll(1);
@@ -259,7 +290,7 @@ public void testPause() throws Exception {
         assertTaskMetricValue("pause-ratio", 0.0);
         assertTaskMetricValue("batch-size-max", 1.0);
         assertTaskMetricValue("batch-size-avg", 0.5);
-        assertTaskMetricValue("offset-commit-max-time-ms", Double.NEGATIVE_INFINITY);
+        assertTaskMetricValue("offset-commit-max-time-ms", Double.NaN);
         assertTaskMetricValue("offset-commit-failure-percentage", 0.0);
         assertTaskMetricValue("offset-commit-success-percentage", 0.0);
 
@@ -284,11 +315,62 @@ public void testPause() throws Exception {
         PowerMock.verifyAll();
     }
 
+    @Test
+    public void testShutdown() throws Exception {
+        createTask(initialState);
+
+        expectInitializeTask();
+        expectTaskGetTopic(true);
+
+        // first iteration
+        expectPollInitialAssignment();
+
+        // second iteration
+        EasyMock.expect(sinkTask.preCommit(EasyMock.anyObject())).andReturn(Collections.emptyMap());
+        expectConsumerPoll(1);
+        expectConversionAndTransformation(1);
+        sinkTask.put(EasyMock.>anyObject());
+        EasyMock.expectLastCall();
+
+        // WorkerSinkTask::stop
+        consumer.wakeup();
+        PowerMock.expectLastCall();
+        sinkTask.stop();
+        PowerMock.expectLastCall();
+
+        // WorkerSinkTask::close
+        consumer.close();
+        PowerMock.expectLastCall().andAnswer(new IAnswer() {
+            @Override
+            public Object answer() throws Throwable {
+                rebalanceListener.getValue().onPartitionsRevoked(
+                    asList(TOPIC_PARTITION, TOPIC_PARTITION2)
+                );
+                return null;
+            }
+        });
+        transformationChain.close();
+        PowerMock.expectLastCall();
+
+        PowerMock.replayAll();
+
+        workerTask.initialize(TASK_CONFIG);
+        workerTask.initializeAndStart();
+        workerTask.iteration();
+        sinkTaskContext.getValue().requestCommit(); // Force an offset commit
+        workerTask.iteration();
+        workerTask.stop();
+        workerTask.close();
+
+        PowerMock.verifyAll();
+    }
+
     @Test
     public void testPollRedelivery() throws Exception {
         createTask(initialState);
 
         expectInitializeTask();
+        expectTaskGetTopic(true);
         expectPollInitialAssignment();
 
         // If a retriable exception is thrown, we should redeliver the same batch, pausing the consumer in the meantime
@@ -337,7 +419,7 @@ public void testPollRedelivery() throws Exception {
         assertTaskMetricValue("pause-ratio", 0.0);
         assertTaskMetricValue("batch-size-max", 0.0);
         assertTaskMetricValue("batch-size-avg", 0.0);
-        assertTaskMetricValue("offset-commit-max-time-ms", Double.NEGATIVE_INFINITY);
+        assertTaskMetricValue("offset-commit-max-time-ms", Double.NaN);
         assertTaskMetricValue("offset-commit-failure-percentage", 0.0);
         assertTaskMetricValue("offset-commit-success-percentage", 0.0);
 
@@ -365,6 +447,7 @@ public void testErrorInRebalancePartitionRevocation() throws Exception {
         createTask(initialState);
 
         expectInitializeTask();
+        expectTaskGetTopic(true);
         expectPollInitialAssignment();
         expectRebalanceRevocationError(exception);
 
@@ -390,6 +473,7 @@ public void testErrorInRebalancePartitionAssignment() throws Exception {
         createTask(initialState);
 
         expectInitializeTask();
+        expectTaskGetTopic(true);
         expectPollInitialAssignment();
         expectRebalanceAssignmentError(exception);
 
@@ -413,7 +497,7 @@ public void testWakeupInCommitSyncCausesRetry() throws Exception {
         createTask(initialState);
 
         expectInitializeTask();
-
+        expectTaskGetTopic(true);
         expectPollInitialAssignment();
 
         expectConsumerPoll(1);
@@ -446,7 +530,7 @@ public void testWakeupInCommitSyncCausesRetry() throws Exception {
         sinkTask.open(partitions);
         EasyMock.expectLastCall();
 
-        EasyMock.expect(consumer.poll(EasyMock.anyLong())).andAnswer(
+        EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer(
                 new IAnswer>() {
                     @Override
                     public ConsumerRecords answer() throws Throwable {
@@ -503,11 +587,65 @@ public ConsumerRecords answer() throws Throwable {
     }
 
     @Test
-    public void testRequestCommit() throws Exception {
+    public void testWakeupNotThrownDuringShutdown() throws Exception {
         createTask(initialState);
 
         expectInitializeTask();
+        expectTaskGetTopic(true);
+        expectPollInitialAssignment();
+
+        expectConsumerPoll(1);
+        expectConversionAndTransformation(1);
+        sinkTask.put(EasyMock.>anyObject());
+        EasyMock.expectLastCall();
 
+        EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer(new IAnswer>() {
+            @Override
+            public ConsumerRecords answer() throws Throwable {
+                // stop the task during its second iteration
+                workerTask.stop();
+                return new ConsumerRecords<>(Collections.emptyMap());
+            }
+        });
+        consumer.wakeup();
+        EasyMock.expectLastCall();
+
+        sinkTask.put(EasyMock.eq(Collections.emptyList()));
+        EasyMock.expectLastCall();
+
+        final Map offsets = new HashMap<>();
+        offsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 1));
+        offsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET));
+        sinkTask.preCommit(offsets);
+        EasyMock.expectLastCall().andReturn(offsets);
+
+        sinkTask.close(EasyMock.anyObject());
+        PowerMock.expectLastCall();
+
+        // fail the first time
+        consumer.commitSync(EasyMock.eq(offsets));
+        EasyMock.expectLastCall().andThrow(new WakeupException());
+
+        // and succeed the second time
+        consumer.commitSync(EasyMock.eq(offsets));
+        EasyMock.expectLastCall();
+
+        PowerMock.replayAll();
+
+        workerTask.initialize(TASK_CONFIG);
+        workerTask.execute();
+
+        assertEquals(0, workerTask.commitFailures());
+
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testRequestCommit() throws Exception {
+        createTask(initialState);
+
+        expectInitializeTask();
+        expectTaskGetTopic(true);
         expectPollInitialAssignment();
 
         expectConsumerPoll(1);
@@ -564,6 +702,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"));
@@ -573,6 +715,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);
@@ -601,6 +751,7 @@ public void testPreCommit() throws Exception {
         createTask(initialState);
 
         expectInitializeTask();
+        expectTaskGetTopic(true);
 
         // iter 1
         expectPollInitialAssignment();
@@ -668,6 +819,7 @@ public void testIgnoredCommit() throws Exception {
         createTask(initialState);
 
         expectInitializeTask();
+        expectTaskGetTopic(true);
 
         // iter 1
         expectPollInitialAssignment();
@@ -718,6 +870,7 @@ public void testLongRunningCommitWithoutTimeout() throws Exception {
         createTask(initialState);
 
         expectInitializeTask();
+        expectTaskGetTopic(true);
 
         // iter 1
         expectPollInitialAssignment();
@@ -809,6 +962,90 @@ public void run() {
         PowerMock.verifyAll();
     }
 
+    @Test
+    public void testSinkTasksHandleCloseErrors() throws Exception {
+        createTask(initialState);
+        expectInitializeTask();
+        expectTaskGetTopic(true);
+
+        // 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();
+        expectTaskGetTopic(true);
+
+        // 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.
@@ -817,6 +1054,7 @@ public void testCommitWithOutOfOrderCallback() throws Exception {
         createTask(initialState);
 
         expectInitializeTask();
+        expectTaskGetTopic(true);
 
         // iter 1
         expectPollInitialAssignment();
@@ -881,7 +1119,7 @@ public void run() {
         // Expect the next poll to discover and perform the rebalance, THEN complete the previous callback handler,
         // and then return one record for TP1 and one for TP3.
         final AtomicBoolean rebalanced = new AtomicBoolean();
-        EasyMock.expect(consumer.poll(EasyMock.anyLong())).andAnswer(
+        EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer(
                 new IAnswer>() {
                     @Override
                     public ConsumerRecords answer() throws Throwable {
@@ -1037,6 +1275,7 @@ public void testDeliveryWithMutatingTransform() throws Exception {
         createTask(initialState);
 
         expectInitializeTask();
+        expectTaskGetTopic(true);
 
         expectPollInitialAssignment();
 
@@ -1091,6 +1330,7 @@ public void testMissingTimestampPropagation() throws Exception {
         createTask(initialState);
 
         expectInitializeTask();
+        expectTaskGetTopic(true);
         expectPollInitialAssignment();
         expectConsumerPoll(1, RecordBatch.NO_TIMESTAMP, TimestampType.CREATE_TIME);
         expectConversionAndTransformation(1);
@@ -1123,6 +1363,7 @@ public void testTimestampPropagation() throws Exception {
         createTask(initialState);
 
         expectInitializeTask();
+        expectTaskGetTopic(true);
         expectPollInitialAssignment();
         expectConsumerPoll(1, timestamp, timestampType);
         expectConversionAndTransformation(1);
@@ -1154,7 +1395,6 @@ public void testTopicsRegex() throws Exception {
 
         createTask(TargetState.PAUSED);
 
-        PowerMock.expectPrivate(workerTask, "createConsumer").andReturn(consumer);
         consumer.subscribe(EasyMock.capture(topicsRegex), EasyMock.capture(rebalanceListener));
         PowerMock.expectLastCall();
 
@@ -1180,8 +1420,149 @@ public void testTopicsRegex() throws Exception {
         PowerMock.verifyAll();
     }
 
+    @Test
+    public void testMetricsGroup() {
+        SinkTaskMetricsGroup group = new SinkTaskMetricsGroup(taskId, metrics);
+        SinkTaskMetricsGroup group1 = new SinkTaskMetricsGroup(taskId1, metrics);
+        for (int i = 0; i != 10; ++i) {
+            group.recordRead(1);
+            group.recordSend(2);
+            group.recordPut(3);
+            group.recordPartitionCount(4);
+            group.recordOffsetSequenceNumber(5);
+        }
+        Map committedOffsets = new HashMap<>();
+        committedOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 1));
+        group.recordCommittedOffsets(committedOffsets);
+        Map consumedOffsets = new HashMap<>();
+        consumedOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 10));
+        group.recordConsumedOffsets(consumedOffsets);
+
+        for (int i = 0; i != 20; ++i) {
+            group1.recordRead(1);
+            group1.recordSend(2);
+            group1.recordPut(30);
+            group1.recordPartitionCount(40);
+            group1.recordOffsetSequenceNumber(50);
+        }
+        committedOffsets = new HashMap<>();
+        committedOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET + 2));
+        committedOffsets.put(TOPIC_PARTITION3, new OffsetAndMetadata(FIRST_OFFSET + 3));
+        group1.recordCommittedOffsets(committedOffsets);
+        consumedOffsets = new HashMap<>();
+        consumedOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET + 20));
+        consumedOffsets.put(TOPIC_PARTITION3, new OffsetAndMetadata(FIRST_OFFSET + 30));
+        group1.recordConsumedOffsets(consumedOffsets);
+
+        assertEquals(0.333, metrics.currentMetricValueAsDouble(group.metricGroup(), "sink-record-read-rate"), 0.001d);
+        assertEquals(0.667, metrics.currentMetricValueAsDouble(group.metricGroup(), "sink-record-send-rate"), 0.001d);
+        assertEquals(9, metrics.currentMetricValueAsDouble(group.metricGroup(), "sink-record-active-count"), 0.001d);
+        assertEquals(4, metrics.currentMetricValueAsDouble(group.metricGroup(), "partition-count"), 0.001d);
+        assertEquals(5, metrics.currentMetricValueAsDouble(group.metricGroup(), "offset-commit-seq-no"), 0.001d);
+        assertEquals(3, metrics.currentMetricValueAsDouble(group.metricGroup(), "put-batch-max-time-ms"), 0.001d);
+
+        // Close the group
+        group.close();
+
+        for (MetricName metricName : group.metricGroup().metrics().metrics().keySet()) {
+            // Metrics for this group should no longer exist
+            assertFalse(group.metricGroup().groupId().includes(metricName));
+        }
+        // Sensors for this group should no longer exist
+        assertNull(group.metricGroup().metrics().getSensor("source-record-poll"));
+        assertNull(group.metricGroup().metrics().getSensor("source-record-write"));
+        assertNull(group.metricGroup().metrics().getSensor("poll-batch-time"));
+
+        assertEquals(0.667, metrics.currentMetricValueAsDouble(group1.metricGroup(), "sink-record-read-rate"), 0.001d);
+        assertEquals(1.333, metrics.currentMetricValueAsDouble(group1.metricGroup(), "sink-record-send-rate"), 0.001d);
+        assertEquals(45, metrics.currentMetricValueAsDouble(group1.metricGroup(), "sink-record-active-count"), 0.001d);
+        assertEquals(40, metrics.currentMetricValueAsDouble(group1.metricGroup(), "partition-count"), 0.001d);
+        assertEquals(50, metrics.currentMetricValueAsDouble(group1.metricGroup(), "offset-commit-seq-no"), 0.001d);
+        assertEquals(30, metrics.currentMetricValueAsDouble(group1.metricGroup(), "put-batch-max-time-ms"), 0.001d);
+    }
+
+    @Test
+    public void testHeaders() throws Exception {
+        Headers headers = new RecordHeaders();
+        headers.add("header_key", "header_value".getBytes());
+
+        createTask(initialState);
+
+        expectInitializeTask();
+        expectTaskGetTopic(true);
+        expectPollInitialAssignment();
+
+        expectConsumerPoll(1, headers);
+        expectConversionAndTransformation(1, null, headers);
+        sinkTask.put(EasyMock.>anyObject());
+        EasyMock.expectLastCall();
+
+        PowerMock.replayAll();
+
+        workerTask.initialize(TASK_CONFIG);
+        workerTask.initializeAndStart();
+        workerTask.iteration(); // iter 1 -- initial assignment
+        workerTask.iteration(); // iter 2 -- deliver 1 record
+
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testHeadersWithCustomConverter() throws Exception {
+        StringConverter stringConverter = new StringConverter();
+        TestConverterWithHeaders testConverter = new TestConverterWithHeaders();
+
+        createTask(initialState, stringConverter, testConverter, stringConverter);
+
+        expectInitializeTask();
+        expectTaskGetTopic(true);
+        expectPollInitialAssignment();
+
+        String keyA = "a";
+        String valueA = "Árvíztűrő tükörfúrógép";
+        Headers headersA = new RecordHeaders();
+        String encodingA = "latin2";
+        headersA.add("encoding", encodingA.getBytes());
+
+        String keyB = "b";
+        String valueB = "Тестовое сообщение";
+        Headers headersB = new RecordHeaders();
+        String encodingB = "koi8_r";
+        headersB.add("encoding", encodingB.getBytes());
+
+        expectConsumerPoll(Arrays.asList(
+            new ConsumerRecord<>(TOPIC, PARTITION, FIRST_OFFSET + recordsReturnedTp1 + 1, RecordBatch.NO_TIMESTAMP, TimestampType.NO_TIMESTAMP_TYPE,
+                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 {
-        PowerMock.expectPrivate(workerTask, "createConsumer").andReturn(consumer);
         consumer.subscribe(EasyMock.eq(asList(TOPIC)), EasyMock.capture(rebalanceListener));
         PowerMock.expectLastCall();
 
@@ -1200,7 +1581,7 @@ private void expectRebalanceRevocationError(RuntimeException e) {
         sinkTask.preCommit(EasyMock.>anyObject());
         EasyMock.expectLastCall().andReturn(Collections.emptyMap());
 
-        EasyMock.expect(consumer.poll(EasyMock.anyLong())).andAnswer(
+        EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer(
                 new IAnswer>() {
                     @Override
                     public ConsumerRecords answer() throws Throwable {
@@ -1225,7 +1606,7 @@ private void expectRebalanceAssignmentError(RuntimeException e) {
         sinkTask.open(partitions);
         EasyMock.expectLastCall().andThrow(e);
 
-        EasyMock.expect(consumer.poll(EasyMock.anyLong())).andAnswer(
+        EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer(
                 new IAnswer>() {
                     @Override
                     public ConsumerRecords answer() throws Throwable {
@@ -1242,7 +1623,7 @@ private void expectPollInitialAssignment() {
         sinkTask.open(partitions);
         EasyMock.expectLastCall();
 
-        EasyMock.expect(consumer.poll(EasyMock.anyLong())).andAnswer(new IAnswer>() {
+        EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer(new IAnswer>() {
             @Override
             public ConsumerRecords answer() throws Throwable {
                 rebalanceListener.getValue().onPartitionsAssigned(partitions);
@@ -1259,21 +1640,29 @@ public ConsumerRecords answer() throws Throwable {
     private void expectConsumerWakeup() {
         consumer.wakeup();
         EasyMock.expectLastCall();
-        EasyMock.expect(consumer.poll(EasyMock.anyLong())).andThrow(new WakeupException());
+        EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andThrow(new WakeupException());
     }
 
     private void expectConsumerPoll(final int numMessages) {
-        expectConsumerPoll(numMessages, RecordBatch.NO_TIMESTAMP, TimestampType.NO_TIMESTAMP_TYPE);
+        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) {
-        EasyMock.expect(consumer.poll(EasyMock.anyLong())).andAnswer(
+        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 ?
@@ -1284,14 +1673,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() {
@@ -1306,13 +1721,36 @@ public SinkRecord answer() {
                                        origRecord.key(),
                                        origRecord.valueSchema(),
                                        origRecord.value(),
-                                       origRecord.timestamp()
+                                       origRecord.timestamp(),
+                                       origRecord.headers()
                                )
                                : origRecord;
                     }
                 }).times(numMessages);
     }
 
+    private void expectTaskGetTopic(boolean anyTimes) {
+        final Capture connectorCapture = EasyMock.newCapture();
+        final Capture topicCapture = EasyMock.newCapture();
+        IExpectationSetters expect = EasyMock.expect(statusBackingStore.getTopic(
+                EasyMock.capture(connectorCapture),
+                EasyMock.capture(topicCapture)));
+        if (anyTimes) {
+            expect.andStubAnswer(() -> new TopicStatus(
+                    topicCapture.getValue(),
+                    new ConnectorTaskId(connectorCapture.getValue(), 0),
+                    Time.SYSTEM.milliseconds()));
+        } else {
+            expect.andAnswer(() -> new TopicStatus(
+                    topicCapture.getValue(),
+                    new ConnectorTaskId(connectorCapture.getValue(), 0),
+                    Time.SYSTEM.milliseconds()));
+        }
+        if (connectorCapture.hasCaptured() && topicCapture.hasCaptured()) {
+            assertEquals("job", connectorCapture.getValue());
+            assertEquals(TOPIC, topicCapture.getValue());
+        }
+    }
 
     private void assertSinkMetricValue(String name, double expected) {
         MetricGroup sinkTaskGroup = workerTask.sinkTaskMetricsGroup().metricGroup();
@@ -1333,7 +1771,7 @@ private void assertTaskMetricValue(String name, String expected) {
     }
 
     private void printMetrics() {
-        System.out.println("");
+        System.out.println();
         sinkMetricValue("sink-record-read-rate");
         sinkMetricValue("sink-record-read-total");
         sinkMetricValue("sink-record-send-rate");
@@ -1389,7 +1827,10 @@ private void assertMetrics(int minimumPollCountExpected) {
         double sendTotal = metrics.currentMetricValueAsDouble(sinkTaskGroup, "sink-record-send-total");
     }
 
-    private abstract static class TestSinkTask extends SinkTask  {
+    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 ce297575b08e6..7a82b20de278a 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,19 +23,24 @@
 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;
 import org.apache.kafka.connect.data.SchemaAndValue;
 import org.apache.kafka.connect.errors.ConnectException;
+import org.apache.kafka.connect.runtime.distributed.ClusterConfigState;
+import org.apache.kafka.connect.runtime.errors.RetryWithToleranceOperatorTest;
 import org.apache.kafka.connect.runtime.isolation.PluginClassLoader;
 import org.apache.kafka.connect.runtime.standalone.StandaloneConfig;
 import org.apache.kafka.connect.sink.SinkConnector;
 import org.apache.kafka.connect.sink.SinkRecord;
 import org.apache.kafka.connect.sink.SinkTask;
 import org.apache.kafka.connect.storage.Converter;
+import org.apache.kafka.connect.storage.HeaderConverter;
+import org.apache.kafka.connect.storage.StatusBackingStore;
 import org.apache.kafka.connect.util.ConnectorTaskId;
-import org.apache.kafka.connect.util.MockTime;
+import org.apache.kafka.common.utils.MockTime;
 import org.apache.kafka.connect.util.ThreadedTest;
 import org.easymock.Capture;
 import org.easymock.CaptureType;
@@ -52,6 +57,7 @@
 import org.powermock.modules.junit4.PowerMockRunner;
 import org.powermock.reflect.Whitebox;
 
+import java.time.Duration;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
@@ -108,11 +114,13 @@ public class WorkerSinkTaskThreadedTest extends ThreadedTest {
     private PluginClassLoader pluginLoader;
     @Mock private Converter keyConverter;
     @Mock private Converter valueConverter;
-    @Mock private TransformationChain transformationChain;
+    @Mock private HeaderConverter headerConverter;
+    @Mock private TransformationChain transformationChain;
     private WorkerSinkTask workerTask;
     @Mock private KafkaConsumer consumer;
     private Capture rebalanceListener = EasyMock.newCapture();
     @Mock private TaskStatus.Listener statusListener;
+    @Mock private StatusBackingStore statusBackingStore;
 
     private long recordsReturned;
 
@@ -132,10 +140,11 @@ public void setup() {
         workerProps.put("offset.storage.file.filename", "/tmp/connect.offsets");
         pluginLoader = PowerMock.createMock(PluginClassLoader.class);
         workerConfig = new StandaloneConfig(workerProps);
-        workerTask = PowerMock.createPartialMock(
-                WorkerSinkTask.class, new String[]{"createConsumer"},
-                taskId, sinkTask, statusListener, initialState, workerConfig, metrics, keyConverter,
-                valueConverter, TransformationChain.noOp(), pluginLoader, time);
+        workerTask = new WorkerSinkTask(
+                taskId, sinkTask, statusListener, initialState, workerConfig, ClusterConfigState.EMPTY, metrics, keyConverter,
+                valueConverter, headerConverter,
+                new TransformationChain<>(Collections.emptyList(), RetryWithToleranceOperatorTest.NOOP_OPERATOR),
+                consumer, pluginLoader, time, RetryWithToleranceOperatorTest.NOOP_OPERATOR, null, statusBackingStore);
 
         recordsReturned = 0;
     }
@@ -148,6 +157,7 @@ public void tearDown() {
     @Test
     public void testPollsInBackground() throws Exception {
         expectInitializeTask();
+        expectTaskGetTopic(true);
         expectPollInitialAssignment();
 
         Capture> capturedRecords = expectPolls(1L);
@@ -177,7 +187,9 @@ public void testPollsInBackground() throws Exception {
             for (SinkRecord rec : recs) {
                 SinkRecord referenceSinkRecord
                         = new SinkRecord(TOPIC, PARTITION, KEY_SCHEMA, KEY, VALUE_SCHEMA, VALUE, FIRST_OFFSET + offset, TIMESTAMP, TIMESTAMP_TYPE);
-                assertEquals(referenceSinkRecord, rec);
+                InternalSinkRecord referenceInternalSinkRecord =
+                    new InternalSinkRecord(null, referenceSinkRecord);
+                assertEquals(referenceInternalSinkRecord, rec);
                 offset++;
             }
         }
@@ -188,6 +200,7 @@ public void testPollsInBackground() throws Exception {
     @Test
     public void testCommit() throws Exception {
         expectInitializeTask();
+        expectTaskGetTopic(true);
         expectPollInitialAssignment();
 
         // Make each poll() take the offset commit interval
@@ -221,6 +234,7 @@ public void testCommit() throws Exception {
     @Test
     public void testCommitFailure() throws Exception {
         expectInitializeTask();
+        expectTaskGetTopic(true);
         expectPollInitialAssignment();
 
         Capture> capturedRecords = expectPolls(WorkerConfig.OFFSET_COMMIT_INTERVAL_MS_DEFAULT);
@@ -260,6 +274,7 @@ public void testCommitSuccessFollowedByFailure() throws Exception {
         // Validate that we rewind to the correct offsets if a task's preCommit() method throws an exception
 
         expectInitializeTask();
+        expectTaskGetTopic(true);
         expectPollInitialAssignment();
         Capture> capturedRecords = expectPolls(WorkerConfig.OFFSET_COMMIT_INTERVAL_MS_DEFAULT);
         expectOffsetCommit(1L, null, null, 0, true);
@@ -298,6 +313,7 @@ public void testCommitSuccessFollowedByFailure() throws Exception {
     @Test
     public void testCommitConsumerFailure() throws Exception {
         expectInitializeTask();
+        expectTaskGetTopic(true);
         expectPollInitialAssignment();
 
         Capture> capturedRecords
@@ -329,6 +345,7 @@ public void testCommitConsumerFailure() throws Exception {
     @Test
     public void testCommitTimeout() throws Exception {
         expectInitializeTask();
+        expectTaskGetTopic(true);
         expectPollInitialAssignment();
 
         // Cut down amount of time to pass in each poll so we trigger exactly 1 offset commit
@@ -366,6 +383,7 @@ public void testAssignmentPauseResume() throws Exception {
         // Just validate that the calls are passed through to the consumer, and that where appropriate errors are
         // converted
         expectInitializeTask();
+        expectTaskGetTopic(true);
 
         expectPollInitialAssignment();
         expectOnePoll().andAnswer(new IAnswer() {
@@ -434,6 +452,7 @@ public Object answer() throws Throwable {
     @Test
     public void testRewind() throws Exception {
         expectInitializeTask();
+        expectTaskGetTopic(true);
         expectPollInitialAssignment();
 
         final long startOffset = 40L;
@@ -477,6 +496,7 @@ public Object answer() throws Throwable {
     @Test
     public void testRewindOnRebalanceDuringPoll() throws Exception {
         expectInitializeTask();
+        expectTaskGetTopic(true);
         expectPollInitialAssignment();
 
         expectRebalanceDuringPoll().andAnswer(new IAnswer() {
@@ -502,7 +522,6 @@ public Object answer() throws Throwable {
     }
 
     private void expectInitializeTask() throws Exception {
-        PowerMock.expectPrivate(workerTask, "createConsumer").andReturn(consumer);
 
         consumer.subscribe(EasyMock.eq(Arrays.asList(TOPIC)), EasyMock.capture(rebalanceListener));
         PowerMock.expectLastCall();
@@ -519,7 +538,7 @@ private void expectPollInitialAssignment() throws Exception {
         sinkTask.open(partitions);
         EasyMock.expectLastCall();
 
-        EasyMock.expect(consumer.poll(EasyMock.anyLong())).andAnswer(new IAnswer>() {
+        EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer(new IAnswer>() {
             @Override
             public ConsumerRecords answer() throws Throwable {
                 rebalanceListener.getValue().onPartitionsAssigned(partitions);
@@ -551,7 +570,7 @@ private void expectStopTask() throws Exception {
     private Capture> expectPolls(final long pollDelayMs) throws Exception {
         // Stub out all the consumer stream/iterator responses, which we just want to verify occur,
         // but don't care about the exact details here.
-        EasyMock.expect(consumer.poll(EasyMock.anyLong())).andStubAnswer(
+        EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andStubAnswer(
                 new IAnswer>() {
                     @Override
                     public ConsumerRecords answer() throws Throwable {
@@ -567,16 +586,12 @@ 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(new IAnswer() {
-            @Override
-            public SinkRecord answer() {
-                return recordCapture.getValue();
-            }
-        }).anyTimes();
+        EasyMock.expect(transformationChain.apply(EasyMock.capture(recordCapture))).andAnswer(
+            (IAnswer) () -> recordCapture.getValue()).anyTimes();
 
         Capture> capturedRecords = EasyMock.newCapture(CaptureType.ALL);
         sinkTask.put(EasyMock.capture(capturedRecords));
@@ -589,7 +604,7 @@ private IExpectationSetters expectOnePoll() {
         // Currently the SinkTask's put() method will not be invoked unless we provide some data, so instead of
         // returning empty data, we return one record. The expectation is that the data will be ignored by the
         // response behavior specified using the return value of this method.
-        EasyMock.expect(consumer.poll(EasyMock.anyLong())).andAnswer(
+        EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer(
                 new IAnswer>() {
                     @Override
                     public ConsumerRecords answer() throws Throwable {
@@ -605,8 +620,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();
     }
@@ -619,7 +634,7 @@ private IExpectationSetters expectRebalanceDuringPoll() throws Exception
         final Map offsets = new HashMap<>();
         offsets.put(TOPIC_PARTITION, startOffset);
 
-        EasyMock.expect(consumer.poll(EasyMock.anyLong())).andAnswer(
+        EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer(
                 new IAnswer>() {
                     @Override
                     public ConsumerRecords answer() throws Throwable {
@@ -650,8 +665,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();
     }
@@ -693,7 +708,33 @@ public Object answer() throws Throwable {
         return capturedCallback;
     }
 
-    private static abstract class TestSinkTask extends SinkTask {
+    private void expectTaskGetTopic(boolean anyTimes) {
+        final Capture connectorCapture = EasyMock.newCapture();
+        final Capture topicCapture = EasyMock.newCapture();
+        IExpectationSetters expect = EasyMock.expect(statusBackingStore.getTopic(
+                EasyMock.capture(connectorCapture),
+                EasyMock.capture(topicCapture)));
+        if (anyTimes) {
+            expect.andStubAnswer(() -> new TopicStatus(
+                    topicCapture.getValue(),
+                    new ConnectorTaskId(connectorCapture.getValue(), 0),
+                    Time.SYSTEM.milliseconds()));
+        } else {
+            expect.andAnswer(() -> new TopicStatus(
+                    topicCapture.getValue(),
+                    new ConnectorTaskId(connectorCapture.getValue(), 0),
+                    Time.SYSTEM.milliseconds()));
+        }
+        if (connectorCapture.hasCaptured() && topicCapture.hasCaptured()) {
+            assertEquals("job", connectorCapture.getValue());
+            assertEquals(TOPIC, topicCapture.getValue());
+        }
+    }
+
+    private 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 4f0d24357e760..ff61ea3dc7cff 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,26 +16,42 @@
  */
 package org.apache.kafka.connect.runtime;
 
+import org.apache.kafka.clients.admin.NewTopic;
 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.InvalidRecordException;
+import org.apache.kafka.common.MetricName;
 import org.apache.kafka.common.TopicPartition;
-import org.apache.kafka.common.record.InvalidRecordException;
+import org.apache.kafka.common.errors.TopicAuthorizationException;
+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.utils.Time;
 import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.data.SchemaAndValue;
+import org.apache.kafka.connect.errors.ConnectException;
+import org.apache.kafka.connect.header.ConnectHeaders;
+import org.apache.kafka.connect.integration.MonitorableSourceConnector;
 import org.apache.kafka.connect.runtime.ConnectMetrics.MetricGroup;
 import org.apache.kafka.connect.runtime.WorkerSourceTask.SourceTaskMetricsGroup;
+import org.apache.kafka.connect.runtime.distributed.ClusterConfigState;
+import org.apache.kafka.connect.runtime.errors.RetryWithToleranceOperatorTest;
 import org.apache.kafka.connect.runtime.isolation.Plugins;
 import org.apache.kafka.connect.runtime.standalone.StandaloneConfig;
 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.OffsetStorageReader;
+import org.apache.kafka.connect.storage.HeaderConverter;
 import org.apache.kafka.connect.storage.OffsetStorageWriter;
+import org.apache.kafka.connect.storage.StatusBackingStore;
+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;
+import org.apache.kafka.connect.util.TopicAdmin;
 import org.easymock.Capture;
 import org.easymock.EasyMock;
 import org.easymock.IAnswer;
@@ -50,6 +66,8 @@
 import org.powermock.modules.junit4.PowerMockRunner;
 import org.powermock.reflect.Whitebox;
 
+import java.nio.ByteBuffer;
+import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -64,11 +82,19 @@
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicInteger;
 
+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.TOPIC_CREATION_ENABLE_CONFIG;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 
-@PowerMockIgnore("javax.management.*")
+@PowerMockIgnore({"javax.management.*",
+                  "org.apache.log4j.*"})
 @RunWith(PowerMockRunner.class)
 public class WorkerSourceTaskTest extends ThreadedTest {
     private static final String TOPIC = "topic";
@@ -87,19 +113,25 @@ public class WorkerSourceTaskTest extends ThreadedTest {
 
     private ExecutorService executor = Executors.newSingleThreadExecutor();
     private ConnectorTaskId taskId = new ConnectorTaskId("job", 0);
+    private ConnectorTaskId taskId1 = new ConnectorTaskId("job", 1);
     private WorkerConfig config;
+    private SourceConnectorConfig sourceConfig;
     private Plugins plugins;
     private MockConnectMetrics metrics;
     @Mock private SourceTask sourceTask;
     @Mock private Converter keyConverter;
     @Mock private Converter valueConverter;
+    @Mock private HeaderConverter headerConverter;
     @Mock private TransformationChain transformationChain;
     @Mock private KafkaProducer producer;
-    @Mock private OffsetStorageReader offsetReader;
+    @Mock private TopicAdmin admin;
+    @Mock private CloseableOffsetStorageReader offsetReader;
     @Mock private OffsetStorageWriter offsetWriter;
+    @Mock private ClusterConfigState clusterConfigState;
     private WorkerSourceTask workerTask;
     @Mock private Future sendFuture;
     @MockStrict private TaskStatus.Listener statusListener;
+    @Mock private StatusBackingStore statusBackingStore;
 
     private Capture producerCallbacks;
 
@@ -113,6 +145,9 @@ public class WorkerSourceTaskTest extends ThreadedTest {
             new SourceRecord(PARTITION, OFFSET, "topic", null, KEY_SCHEMA, KEY, RECORD_SCHEMA, RECORD)
     );
 
+    // when this test becomes parameterized, this variable will be a test parameter
+    public boolean enableTopicCreation = false;
+
     @Override
     public void setup() {
         super.setup();
@@ -124,12 +159,26 @@ public void setup() {
         workerProps.put("internal.key.converter.schemas.enable", "false");
         workerProps.put("internal.value.converter.schemas.enable", "false");
         workerProps.put("offset.storage.file.filename", "/tmp/connect.offsets");
+        workerProps.put(TOPIC_CREATION_ENABLE_CONFIG, String.valueOf(enableTopicCreation));
         plugins = new Plugins(workerProps);
         config = new StandaloneConfig(workerProps);
+        sourceConfig = new SourceConnectorConfig(plugins, sourceConnectorProps(TOPIC), true);
         producerCallbacks = EasyMock.newCapture();
         metrics = new MockConnectMetrics();
     }
 
+    private Map sourceConnectorProps(String topic) {
+        // setup up props for the source connector
+        Map props = new HashMap<>();
+        props.put("name", "foo-connector");
+        props.put(CONNECTOR_CLASS_CONFIG, MonitorableSourceConnector.class.getSimpleName());
+        props.put(TASKS_MAX_CONFIG, String.valueOf(1));
+        props.put(TOPIC_CONFIG, topic);
+        props.put(KEY_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
+        props.put(VALUE_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
+        return props;
+    }
+
     @After
     public void tearDown() {
         if (metrics != null) metrics.stop();
@@ -140,8 +189,14 @@ private void createWorkerTask() {
     }
 
     private void createWorkerTask(TargetState initialState) {
-        workerTask = new WorkerSourceTask(taskId, sourceTask, statusListener, initialState, keyConverter, valueConverter, transformationChain,
-                producer, offsetReader, offsetWriter, config, metrics, plugins.delegatingLoader(), Time.SYSTEM);
+        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, admin, null,
+                offsetReader, offsetWriter, config, clusterConfigState, metrics, plugins.delegatingLoader(), Time.SYSTEM,
+                RetryWithToleranceOperatorTest.NOOP_OPERATOR, statusBackingStore);
     }
 
     @Test
@@ -159,11 +214,7 @@ public Void answer() throws Throwable {
             }
         });
 
-        producer.close(EasyMock.anyLong(), EasyMock.anyObject(TimeUnit.class));
-        EasyMock.expectLastCall();
-
-        transformationChain.close();
-        EasyMock.expectLastCall();
+        expectClose();
 
         statusListener.onShutdown(taskId);
         EasyMock.expectLastCall();
@@ -197,6 +248,8 @@ public void testPause() throws Exception {
         CountDownLatch pollLatch = expectPolls(10, count);
         // In this test, we don't flush, so nothing goes any further than the offset writer
 
+        expectTopicCreation(TOPIC);
+
         statusListener.onPause(taskId);
         EasyMock.expectLastCall();
 
@@ -207,11 +260,7 @@ public void testPause() throws Exception {
         statusListener.onShutdown(taskId);
         EasyMock.expectLastCall();
 
-        producer.close(EasyMock.anyLong(), EasyMock.anyObject(TimeUnit.class));
-        EasyMock.expectLastCall();
-
-        transformationChain.close();
-        EasyMock.expectLastCall();
+        expectClose();
 
         PowerMock.replayAll();
 
@@ -249,6 +298,8 @@ public void testPollsInBackground() throws Exception {
         final CountDownLatch pollLatch = expectPolls(10);
         // In this test, we don't flush, so nothing goes any further than the offset writer
 
+        expectTopicCreation(TOPIC);
+
         sourceTask.stop();
         EasyMock.expectLastCall();
         expectOffsetFlush(true);
@@ -256,11 +307,7 @@ public void testPollsInBackground() throws Exception {
         statusListener.onShutdown(taskId);
         EasyMock.expectLastCall();
 
-        producer.close(EasyMock.anyLong(), EasyMock.anyObject(TimeUnit.class));
-        EasyMock.expectLastCall();
-
-        transformationChain.close();
-        EasyMock.expectLastCall();
+        expectClose();
 
         PowerMock.replayAll();
 
@@ -305,11 +352,47 @@ public List answer() throws Throwable {
         EasyMock.expectLastCall();
         expectOffsetFlush(true);
 
-        producer.close(EasyMock.anyLong(), EasyMock.anyObject(TimeUnit.class));
+        expectClose();
+
+        PowerMock.replayAll();
+
+        workerTask.initialize(TASK_CONFIG);
+        Future taskFuture = executor.submit(workerTask);
+
+        assertTrue(awaitLatch(pollLatch));
+        workerTask.stop();
+        assertTrue(workerTask.awaitStop(1000));
+
+        taskFuture.get();
+        assertPollMetrics(0);
+
+        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();
 
-        transformationChain.close();
+        // 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();
+
+        expectClose();
 
         PowerMock.replayAll();
 
@@ -317,6 +400,7 @@ public List answer() throws Throwable {
         Future taskFuture = executor.submit(workerTask);
 
         assertTrue(awaitLatch(pollLatch));
+        assertTrue(workerTask.commitOffsets());
         workerTask.stop();
         assertTrue(workerTask.awaitStop(1000));
 
@@ -342,6 +426,8 @@ public void testCommit() throws Exception {
         final CountDownLatch pollLatch = expectPolls(1);
         expectOffsetFlush(true);
 
+        expectTopicCreation(TOPIC);
+
         sourceTask.stop();
         EasyMock.expectLastCall();
         expectOffsetFlush(true);
@@ -349,11 +435,7 @@ public void testCommit() throws Exception {
         statusListener.onShutdown(taskId);
         EasyMock.expectLastCall();
 
-        producer.close(EasyMock.anyLong(), EasyMock.anyObject(TimeUnit.class));
-        EasyMock.expectLastCall();
-
-        transformationChain.close();
-        EasyMock.expectLastCall();
+        expectClose();
 
         PowerMock.replayAll();
 
@@ -387,6 +469,8 @@ public void testCommitFailure() throws Exception {
         final CountDownLatch pollLatch = expectPolls(1);
         expectOffsetFlush(true);
 
+        expectTopicCreation(TOPIC);
+
         sourceTask.stop();
         EasyMock.expectLastCall();
         expectOffsetFlush(false);
@@ -394,11 +478,7 @@ public void testCommitFailure() throws Exception {
         statusListener.onShutdown(taskId);
         EasyMock.expectLastCall();
 
-        producer.close(EasyMock.anyLong(), EasyMock.anyObject(TimeUnit.class));
-        EasyMock.expectLastCall();
-
-        transformationChain.close();
-        EasyMock.expectLastCall();
+        expectClose();
 
         PowerMock.replayAll();
 
@@ -426,6 +506,8 @@ public void testSendRecordsConvertsData() throws Exception {
 
         Capture> sent = expectSendRecordAnyTimes();
 
+        expectTopicCreation(TOPIC);
+
         PowerMock.replayAll();
 
         Whitebox.setInternalState(workerTask, "toSend", records);
@@ -448,6 +530,8 @@ public void testSendRecordsPropagatesTimestamp() throws Exception {
 
         Capture> sent = expectSendRecordAnyTimes();
 
+        expectTopicCreation(TOPIC);
+
         PowerMock.replayAll();
 
         Whitebox.setInternalState(workerTask, "toSend", records);
@@ -488,6 +572,8 @@ public void testSendRecordsNoTimestamp() throws Exception {
 
         Capture> sent = expectSendRecordAnyTimes();
 
+        expectTopicCreation(TOPIC);
+
         PowerMock.replayAll();
 
         Whitebox.setInternalState(workerTask, "toSend", records);
@@ -506,6 +592,8 @@ public void testSendRecordsRetries() throws Exception {
         SourceRecord record2 = new SourceRecord(PARTITION, OFFSET, "topic", 2, KEY_SCHEMA, KEY, RECORD_SCHEMA, RECORD);
         SourceRecord record3 = new SourceRecord(PARTITION, OFFSET, "topic", 3, KEY_SCHEMA, KEY, RECORD_SCHEMA, RECORD);
 
+        expectTopicCreation(TOPIC);
+
         // First round
         expectSendRecordOnce(false);
         // Any Producer retriable exception should work here
@@ -531,6 +619,23 @@ 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);
+
+        expectTopicCreation(TOPIC);
+
+        expectSendRecordProducerCallbackFail();
+
+        PowerMock.replayAll();
+
+        Whitebox.setInternalState(workerTask, "toSend", Arrays.asList(record1, record2));
+        Whitebox.invokeMethod(workerTask, "sendRecords");
+    }
+
     @Test
     public void testSendRecordsTaskCommitRecordFail() throws Exception {
         createWorkerTask();
@@ -540,6 +645,8 @@ public void testSendRecordsTaskCommitRecordFail() throws Exception {
         SourceRecord record2 = new SourceRecord(PARTITION, OFFSET, "topic", 2, KEY_SCHEMA, KEY, RECORD_SCHEMA, RECORD);
         SourceRecord record3 = new SourceRecord(PARTITION, OFFSET, "topic", 3, KEY_SCHEMA, KEY, RECORD_SCHEMA, RECORD);
 
+        expectTopicCreation(TOPIC);
+
         // Source task commit record failure will not cause the task to abort
         expectSendRecordOnce(false);
         expectSendRecordTaskCommitRecordFail(false, false);
@@ -584,11 +691,7 @@ public Object answer() throws Throwable {
         statusListener.onShutdown(taskId);
         EasyMock.expectLastCall();
 
-        producer.close(EasyMock.anyLong(), EasyMock.anyObject(TimeUnit.class));
-        EasyMock.expectLastCall();
-
-        transformationChain.close();
-        EasyMock.expectLastCall();
+        expectClose();
 
         PowerMock.replayAll();
 
@@ -608,19 +711,160 @@ 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);
+        SourceTaskMetricsGroup group1 = new SourceTaskMetricsGroup(taskId1, metrics);
         for (int i = 0; i != 10; ++i) {
             group.recordPoll(100, 1000 + i * 100);
             group.recordWrite(10);
         }
+        for (int i = 0; i != 20; ++i) {
+            group1.recordPoll(100, 1000 + i * 100);
+            group1.recordWrite(10);
+        }
         assertEquals(1900.0, metrics.currentMetricValueAsDouble(group.metricGroup(), "poll-batch-max-time-ms"), 0.001d);
         assertEquals(1450.0, metrics.currentMetricValueAsDouble(group.metricGroup(), "poll-batch-avg-time-ms"), 0.001d);
         assertEquals(33.333, metrics.currentMetricValueAsDouble(group.metricGroup(), "source-record-poll-rate"), 0.001d);
         assertEquals(1000, metrics.currentMetricValueAsDouble(group.metricGroup(), "source-record-poll-total"), 0.001d);
         assertEquals(3.3333, metrics.currentMetricValueAsDouble(group.metricGroup(), "source-record-write-rate"), 0.001d);
         assertEquals(100, metrics.currentMetricValueAsDouble(group.metricGroup(), "source-record-write-total"), 0.001d);
+        assertEquals(900.0, metrics.currentMetricValueAsDouble(group.metricGroup(), "source-record-active-count"), 0.001d);
+
+        // Close the group
+        group.close();
+
+        for (MetricName metricName : group.metricGroup().metrics().metrics().keySet()) {
+            // Metrics for this group should no longer exist
+            assertFalse(group.metricGroup().groupId().includes(metricName));
+        }
+        // Sensors for this group should no longer exist
+        assertNull(group.metricGroup().metrics().getSensor("sink-record-read"));
+        assertNull(group.metricGroup().metrics().getSensor("sink-record-send"));
+        assertNull(group.metricGroup().metrics().getSensor("sink-record-active-count"));
+        assertNull(group.metricGroup().metrics().getSensor("partition-count"));
+        assertNull(group.metricGroup().metrics().getSensor("offset-seq-number"));
+        assertNull(group.metricGroup().metrics().getSensor("offset-commit-completion"));
+        assertNull(group.metricGroup().metrics().getSensor("offset-commit-completion-skip"));
+        assertNull(group.metricGroup().metrics().getSensor("put-batch-time"));
+
+        assertEquals(2900.0, metrics.currentMetricValueAsDouble(group1.metricGroup(), "poll-batch-max-time-ms"), 0.001d);
+        assertEquals(1950.0, metrics.currentMetricValueAsDouble(group1.metricGroup(), "poll-batch-avg-time-ms"), 0.001d);
+        assertEquals(66.667, metrics.currentMetricValueAsDouble(group1.metricGroup(), "source-record-poll-rate"), 0.001d);
+        assertEquals(2000, metrics.currentMetricValueAsDouble(group1.metricGroup(), "source-record-poll-total"), 0.001d);
+        assertEquals(6.667, metrics.currentMetricValueAsDouble(group1.metricGroup(), "source-record-write-rate"), 0.001d);
+        assertEquals(200, metrics.currentMetricValueAsDouble(group1.metricGroup(), "source-record-write-total"), 0.001d);
+        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));
+
+        expectTopicCreation(TOPIC);
+
+        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));
+
+        expectTopicCreation(TOPIC);
+
+        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 {
@@ -649,7 +893,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);
@@ -669,17 +913,34 @@ 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());
+    }
+
+    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();
@@ -695,15 +956,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();
                 }
@@ -715,19 +980,33 @@ 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);
+            expectTaskGetTopic(anyTimes);
+        }
 
         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
@@ -753,8 +1032,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"));
@@ -764,6 +1043,29 @@ private void expectTaskCommitRecord(boolean anyTimes, boolean succeed) throws In
         }
     }
 
+    private void expectTaskGetTopic(boolean anyTimes) {
+        final Capture connectorCapture = EasyMock.newCapture();
+        final Capture topicCapture = EasyMock.newCapture();
+        IExpectationSetters expect = EasyMock.expect(statusBackingStore.getTopic(
+                EasyMock.capture(connectorCapture),
+                EasyMock.capture(topicCapture)));
+        if (anyTimes) {
+            expect.andStubAnswer(() -> new TopicStatus(
+                    topicCapture.getValue(),
+                    new ConnectorTaskId(connectorCapture.getValue(), 0),
+                    Time.SYSTEM.milliseconds()));
+        } else {
+            expect.andAnswer(() -> new TopicStatus(
+                    topicCapture.getValue(),
+                    new ConnectorTaskId(connectorCapture.getValue(), 0),
+                    Time.SYSTEM.milliseconds()));
+        }
+        if (connectorCapture.hasCaptured() && topicCapture.hasCaptured()) {
+            assertEquals("job", connectorCapture.getValue());
+            assertEquals(TOPIC, topicCapture.getValue());
+        }
+    }
+
     private boolean awaitLatch(CountDownLatch latch) {
         try {
             return latch.await(5000, TimeUnit.MILLISECONDS);
@@ -820,7 +1122,7 @@ private void assertPollMetrics(int minimumPollCountExpected) {
         if (minimumPollCountExpected > 0) {
             assertTrue(pollBatchTimeMax >= 0.0d);
         }
-        assertTrue(pollBatchTimeAvg >= 0.0d);
+        assertTrue(Double.isNaN(pollBatchTimeAvg) || pollBatchTimeAvg > 0.0d);
         double activeCount = metrics.currentMetricValueAsDouble(sourceTaskGroup, "source-record-active-count");
         double activeCountMax = metrics.currentMetricValueAsDouble(sourceTaskGroup, "source-record-active-count-max");
         assertEquals(0, activeCount, 0.000001d);
@@ -829,7 +1131,30 @@ private void assertPollMetrics(int minimumPollCountExpected) {
         }
     }
 
+    private RecordHeaders emptyHeaders() {
+        return new RecordHeaders();
+    }
+
     private abstract static class TestSourceTask extends SourceTask {
     }
 
+    private void expectClose() {
+        producer.close(EasyMock.anyObject(Duration.class));
+        EasyMock.expectLastCall();
+
+        admin.close(EasyMock.anyObject(Duration.class));
+        EasyMock.expectLastCall();
+
+        transformationChain.close();
+        EasyMock.expectLastCall();
+    }
+
+    private void expectTopicCreation(String topic) {
+        if (config.topicCreationEnable()) {
+            EasyMock.expect(admin.describeTopics(topic)).andReturn(Collections.emptyMap());
+
+            Capture newTopicCapture = EasyMock.newCapture();
+            EasyMock.expect(admin.createTopic(EasyMock.capture(newTopicCapture))).andReturn(true);
+        }
+    }
 }
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSourceTaskWithTopicCreationTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSourceTaskWithTopicCreationTest.java
new file mode 100644
index 0000000000000..06862d13f01b0
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSourceTaskWithTopicCreationTest.java
@@ -0,0 +1,1437 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.admin.NewTopic;
+import org.apache.kafka.clients.admin.TopicDescription;
+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.InvalidRecordException;
+import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.MetricName;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.TopicPartitionInfo;
+import org.apache.kafka.common.errors.InvalidTopicException;
+import org.apache.kafka.common.errors.TopicAuthorizationException;
+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.utils.Time;
+import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.data.SchemaAndValue;
+import org.apache.kafka.connect.errors.ConnectException;
+import org.apache.kafka.connect.errors.RetriableException;
+import org.apache.kafka.connect.header.ConnectHeaders;
+import org.apache.kafka.connect.integration.MonitorableSourceConnector;
+import org.apache.kafka.connect.runtime.ConnectMetrics.MetricGroup;
+import org.apache.kafka.connect.runtime.WorkerSourceTask.SourceTaskMetricsGroup;
+import org.apache.kafka.connect.runtime.distributed.ClusterConfigState;
+import org.apache.kafka.connect.runtime.errors.RetryWithToleranceOperatorTest;
+import org.apache.kafka.connect.runtime.isolation.Plugins;
+import org.apache.kafka.connect.runtime.standalone.StandaloneConfig;
+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.OffsetStorageWriter;
+import org.apache.kafka.connect.storage.StatusBackingStore;
+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;
+import org.apache.kafka.connect.util.TopicAdmin;
+import org.apache.kafka.connect.util.TopicCreationGroup;
+import org.easymock.Capture;
+import org.easymock.EasyMock;
+import org.easymock.IAnswer;
+import org.easymock.IExpectationSetters;
+import org.junit.After;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.powermock.api.easymock.PowerMock;
+import org.powermock.api.easymock.annotation.Mock;
+import org.powermock.api.easymock.annotation.MockStrict;
+import org.powermock.core.classloader.annotations.PowerMockIgnore;
+import org.powermock.modules.junit4.PowerMockRunner;
+import org.powermock.reflect.Whitebox;
+
+import java.nio.ByteBuffer;
+import java.time.Duration;
+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.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.TimeoutException;
+import java.util.concurrent.atomic.AtomicInteger;
+
+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.SourceConnectorConfig.TOPIC_CREATION_GROUPS_CONFIG;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.DEFAULT_TOPIC_CREATION_PREFIX;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.EXCLUDE_REGEX_CONFIG;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.INCLUDE_REGEX_CONFIG;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.PARTITIONS_CONFIG;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.REPLICATION_FACTOR_CONFIG;
+import static org.apache.kafka.connect.runtime.WorkerConfig.TOPIC_CREATION_ENABLE_CONFIG;
+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;
+
+@PowerMockIgnore({"javax.management.*",
+                  "org.apache.log4j.*"})
+@RunWith(PowerMockRunner.class)
+public class WorkerSourceTaskWithTopicCreationTest extends ThreadedTest {
+    private static final String TOPIC = "topic";
+    private static final String OTHER_TOPIC = "other-topic";
+    private static final Map PARTITION = Collections.singletonMap("key", "partition".getBytes());
+    private static final Map OFFSET = Collections.singletonMap("key", 12);
+
+    // Connect-format data
+    private static final Schema KEY_SCHEMA = Schema.INT32_SCHEMA;
+    private static final Integer KEY = -1;
+    private static final Schema RECORD_SCHEMA = Schema.INT64_SCHEMA;
+    private static final Long RECORD = 12L;
+    // Serialized data. The actual format of this data doesn't matter -- we just want to see that the right version
+    // is used in the right place.
+    private static final byte[] SERIALIZED_KEY = "converted-key".getBytes();
+    private static final byte[] SERIALIZED_RECORD = "converted-record".getBytes();
+
+    private ExecutorService executor = Executors.newSingleThreadExecutor();
+    private ConnectorTaskId taskId = new ConnectorTaskId("job", 0);
+    private ConnectorTaskId taskId1 = new ConnectorTaskId("job", 1);
+    private WorkerConfig config;
+    private SourceConnectorConfig sourceConfig;
+    private Plugins plugins;
+    private MockConnectMetrics metrics;
+    @Mock private SourceTask sourceTask;
+    @Mock private Converter keyConverter;
+    @Mock private Converter valueConverter;
+    @Mock private HeaderConverter headerConverter;
+    @Mock private TransformationChain transformationChain;
+    @Mock private KafkaProducer producer;
+    @Mock private TopicAdmin admin;
+    @Mock private CloseableOffsetStorageReader offsetReader;
+    @Mock private OffsetStorageWriter offsetWriter;
+    @Mock private ClusterConfigState clusterConfigState;
+    private WorkerSourceTask workerTask;
+    @Mock private Future sendFuture;
+    @MockStrict private TaskStatus.Listener statusListener;
+    @Mock private StatusBackingStore statusBackingStore;
+
+    private Capture producerCallbacks;
+
+    private static final Map TASK_PROPS = new HashMap<>();
+    static {
+        TASK_PROPS.put(TaskConfig.TASK_CLASS_CONFIG, TestSourceTask.class.getName());
+    }
+    private static final TaskConfig TASK_CONFIG = new TaskConfig(TASK_PROPS);
+
+    private static final List RECORDS = Arrays.asList(
+            new SourceRecord(PARTITION, OFFSET, TOPIC, null, KEY_SCHEMA, KEY, RECORD_SCHEMA, RECORD)
+    );
+
+    // when this test becomes parameterized, this variable will be a test parameter
+    public boolean enableTopicCreation = true;
+
+    @Override
+    public void setup() {
+        super.setup();
+        Map workerProps = workerProps();
+        plugins = new Plugins(workerProps);
+        config = new StandaloneConfig(workerProps);
+        sourceConfig = new SourceConnectorConfig(plugins, sourceConnectorPropsWithGroups(TOPIC), true);
+        producerCallbacks = EasyMock.newCapture();
+        metrics = new MockConnectMetrics();
+    }
+
+    private Map workerProps() {
+        Map props = new HashMap<>();
+        props.put("key.converter", "org.apache.kafka.connect.json.JsonConverter");
+        props.put("value.converter", "org.apache.kafka.connect.json.JsonConverter");
+        props.put("internal.key.converter", "org.apache.kafka.connect.json.JsonConverter");
+        props.put("internal.value.converter", "org.apache.kafka.connect.json.JsonConverter");
+        props.put("internal.key.converter.schemas.enable", "false");
+        props.put("internal.value.converter.schemas.enable", "false");
+        props.put("offset.storage.file.filename", "/tmp/connect.offsets");
+        props.put(TOPIC_CREATION_ENABLE_CONFIG, String.valueOf(enableTopicCreation));
+        return props;
+    }
+
+    private Map sourceConnectorPropsWithGroups(String topic) {
+        // setup up props for the source connector
+        Map props = new HashMap<>();
+        props.put("name", "foo-connector");
+        props.put(CONNECTOR_CLASS_CONFIG, MonitorableSourceConnector.class.getSimpleName());
+        props.put(TASKS_MAX_CONFIG, String.valueOf(1));
+        props.put(TOPIC_CONFIG, topic);
+        props.put(KEY_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
+        props.put(VALUE_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
+        props.put(TOPIC_CREATION_GROUPS_CONFIG, String.join(",", "foo", "bar"));
+        props.put(DEFAULT_TOPIC_CREATION_PREFIX + REPLICATION_FACTOR_CONFIG, String.valueOf(1));
+        props.put(DEFAULT_TOPIC_CREATION_PREFIX + PARTITIONS_CONFIG, String.valueOf(1));
+        props.put(SourceConnectorConfig.TOPIC_CREATION_PREFIX + "foo" + "." + INCLUDE_REGEX_CONFIG, topic);
+        props.put(SourceConnectorConfig.TOPIC_CREATION_PREFIX + "bar" + "." + INCLUDE_REGEX_CONFIG, ".*");
+        props.put(SourceConnectorConfig.TOPIC_CREATION_PREFIX + "bar" + "." + EXCLUDE_REGEX_CONFIG, topic);
+        return props;
+    }
+
+    @After
+    public void tearDown() {
+        if (metrics != null) metrics.stop();
+    }
+
+    private void createWorkerTask() {
+        createWorkerTask(TargetState.STARTED);
+    }
+
+    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, admin, TopicCreationGroup.configuredGroups(sourceConfig),
+                offsetReader, offsetWriter, config, clusterConfigState, metrics, plugins.delegatingLoader(), Time.SYSTEM,
+                RetryWithToleranceOperatorTest.NOOP_OPERATOR, statusBackingStore);
+    }
+
+    @Test
+    public void testStartPaused() throws Exception {
+        final CountDownLatch pauseLatch = new CountDownLatch(1);
+
+        createWorkerTask(TargetState.PAUSED);
+
+        statusListener.onPause(taskId);
+        EasyMock.expectLastCall().andAnswer(new IAnswer() {
+            @Override
+            public Void answer() throws Throwable {
+                pauseLatch.countDown();
+                return null;
+            }
+        });
+
+        expectClose();
+
+        statusListener.onShutdown(taskId);
+        EasyMock.expectLastCall();
+
+        PowerMock.replayAll();
+
+        workerTask.initialize(TASK_CONFIG);
+        Future taskFuture = executor.submit(workerTask);
+
+        assertTrue(pauseLatch.await(5, TimeUnit.SECONDS));
+        workerTask.stop();
+        assertTrue(workerTask.awaitStop(1000));
+
+        taskFuture.get();
+
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testPause() throws Exception {
+        createWorkerTask();
+
+        sourceTask.initialize(EasyMock.anyObject(SourceTaskContext.class));
+        EasyMock.expectLastCall();
+        sourceTask.start(TASK_PROPS);
+        EasyMock.expectLastCall();
+        statusListener.onStartup(taskId);
+        EasyMock.expectLastCall();
+
+        AtomicInteger count = new AtomicInteger(0);
+        CountDownLatch pollLatch = expectPolls(10, count);
+        // In this test, we don't flush, so nothing goes any further than the offset writer
+
+        expectTopicCreation(TOPIC);
+
+        statusListener.onPause(taskId);
+        EasyMock.expectLastCall();
+
+        sourceTask.stop();
+        EasyMock.expectLastCall();
+        expectOffsetFlush(true);
+
+        statusListener.onShutdown(taskId);
+        EasyMock.expectLastCall();
+
+        expectClose();
+
+        PowerMock.replayAll();
+
+        workerTask.initialize(TASK_CONFIG);
+        Future taskFuture = executor.submit(workerTask);
+        assertTrue(awaitLatch(pollLatch));
+
+        workerTask.transitionTo(TargetState.PAUSED);
+
+        int priorCount = count.get();
+        Thread.sleep(100);
+
+        // since the transition is observed asynchronously, the count could be off by one loop iteration
+        assertTrue(count.get() - priorCount <= 1);
+
+        workerTask.stop();
+        assertTrue(workerTask.awaitStop(1000));
+
+        taskFuture.get();
+
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testPollsInBackground() throws Exception {
+        createWorkerTask();
+
+        sourceTask.initialize(EasyMock.anyObject(SourceTaskContext.class));
+        EasyMock.expectLastCall();
+        sourceTask.start(TASK_PROPS);
+        EasyMock.expectLastCall();
+        statusListener.onStartup(taskId);
+        EasyMock.expectLastCall();
+
+        final CountDownLatch pollLatch = expectPolls(10);
+        // In this test, we don't flush, so nothing goes any further than the offset writer
+
+        expectTopicCreation(TOPIC);
+
+        sourceTask.stop();
+        EasyMock.expectLastCall();
+        expectOffsetFlush(true);
+
+        statusListener.onShutdown(taskId);
+        EasyMock.expectLastCall();
+
+        expectClose();
+
+        PowerMock.replayAll();
+
+        workerTask.initialize(TASK_CONFIG);
+        Future taskFuture = executor.submit(workerTask);
+
+        assertTrue(awaitLatch(pollLatch));
+        workerTask.stop();
+        assertTrue(workerTask.awaitStop(1000));
+
+        taskFuture.get();
+        assertPollMetrics(10);
+
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testFailureInPoll() throws Exception {
+        createWorkerTask();
+
+        sourceTask.initialize(EasyMock.anyObject(SourceTaskContext.class));
+        EasyMock.expectLastCall();
+        sourceTask.start(TASK_PROPS);
+        EasyMock.expectLastCall();
+        statusListener.onStartup(taskId);
+        EasyMock.expectLastCall();
+
+        final CountDownLatch pollLatch = new CountDownLatch(1);
+        final RuntimeException exception = new RuntimeException();
+        EasyMock.expect(sourceTask.poll()).andAnswer(new IAnswer>() {
+            @Override
+            public List answer() throws Throwable {
+                pollLatch.countDown();
+                throw exception;
+            }
+        });
+
+        statusListener.onFailure(taskId, exception);
+        EasyMock.expectLastCall();
+
+        sourceTask.stop();
+        EasyMock.expectLastCall();
+        expectOffsetFlush(true);
+
+        expectClose();
+
+        PowerMock.replayAll();
+
+        workerTask.initialize(TASK_CONFIG);
+        Future taskFuture = executor.submit(workerTask);
+
+        assertTrue(awaitLatch(pollLatch));
+        workerTask.stop();
+        assertTrue(workerTask.awaitStop(1000));
+
+        taskFuture.get();
+        assertPollMetrics(0);
+
+        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();
+
+        expectClose();
+
+        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
+        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 = expectPolls(1);
+        expectOffsetFlush(true);
+
+        expectTopicCreation(TOPIC);
+
+        sourceTask.stop();
+        EasyMock.expectLastCall();
+        expectOffsetFlush(true);
+
+        statusListener.onShutdown(taskId);
+        EasyMock.expectLastCall();
+
+        expectClose();
+
+        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(1);
+
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testCommitFailure() throws Exception {
+        // Test that the task commits properly when prompted
+        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 = expectPolls(1);
+        expectOffsetFlush(true);
+
+        expectTopicCreation(TOPIC);
+
+        sourceTask.stop();
+        EasyMock.expectLastCall();
+        expectOffsetFlush(false);
+
+        statusListener.onShutdown(taskId);
+        EasyMock.expectLastCall();
+
+        expectClose();
+
+        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(1);
+
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testSendRecordsConvertsData() throws Exception {
+        createWorkerTask();
+
+        List records = new ArrayList<>();
+        // Can just use the same record for key and value
+        records.add(new SourceRecord(PARTITION, OFFSET, TOPIC, null, KEY_SCHEMA, KEY, RECORD_SCHEMA, RECORD));
+
+        Capture> sent = expectSendRecordAnyTimes();
+
+        expectTopicCreation(TOPIC);
+
+        PowerMock.replayAll();
+
+        Whitebox.setInternalState(workerTask, "toSend", records);
+        Whitebox.invokeMethod(workerTask, "sendRecords");
+        assertEquals(SERIALIZED_KEY, sent.getValue().key());
+        assertEquals(SERIALIZED_RECORD, sent.getValue().value());
+
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testSendRecordsPropagatesTimestamp() throws Exception {
+        final Long timestamp = System.currentTimeMillis();
+
+        createWorkerTask();
+
+        List records = Collections.singletonList(
+                new SourceRecord(PARTITION, OFFSET, TOPIC, null, KEY_SCHEMA, KEY, RECORD_SCHEMA, RECORD, timestamp)
+        );
+
+        Capture> sent = expectSendRecordAnyTimes();
+
+        expectTopicCreation(TOPIC);
+
+        PowerMock.replayAll();
+
+        Whitebox.setInternalState(workerTask, "toSend", records);
+        Whitebox.invokeMethod(workerTask, "sendRecords");
+        assertEquals(timestamp, sent.getValue().timestamp());
+
+        PowerMock.verifyAll();
+    }
+
+    @Test(expected = InvalidRecordException.class)
+    public void testSendRecordsCorruptTimestamp() throws Exception {
+        final Long timestamp = -3L;
+        createWorkerTask();
+
+        List records = Collections.singletonList(
+                new SourceRecord(PARTITION, OFFSET, TOPIC, null, KEY_SCHEMA, KEY, RECORD_SCHEMA, RECORD, timestamp)
+        );
+
+        Capture> sent = expectSendRecordAnyTimes();
+
+        PowerMock.replayAll();
+
+        Whitebox.setInternalState(workerTask, "toSend", records);
+        Whitebox.invokeMethod(workerTask, "sendRecords");
+        assertEquals(null, sent.getValue().timestamp());
+
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testSendRecordsNoTimestamp() throws Exception {
+        final Long timestamp = -1L;
+        createWorkerTask();
+
+        List records = Collections.singletonList(
+                new SourceRecord(PARTITION, OFFSET, TOPIC, null, KEY_SCHEMA, KEY, RECORD_SCHEMA, RECORD, timestamp)
+        );
+
+        Capture> sent = expectSendRecordAnyTimes();
+
+        expectTopicCreation(TOPIC);
+
+        PowerMock.replayAll();
+
+        Whitebox.setInternalState(workerTask, "toSend", records);
+        Whitebox.invokeMethod(workerTask, "sendRecords");
+        assertEquals(null, sent.getValue().timestamp());
+
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testSendRecordsRetries() throws Exception {
+        createWorkerTask();
+
+        // Differentiate only by Kafka partition so we can reuse conversion expectations
+        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);
+        SourceRecord record3 = new SourceRecord(PARTITION, OFFSET, TOPIC, 3, KEY_SCHEMA, KEY, RECORD_SCHEMA, RECORD);
+
+        expectTopicCreation(TOPIC);
+
+        // First round
+        expectSendRecordOnce(false);
+        // Any Producer retriable exception should work here
+        expectSendRecordSyncFailure(new org.apache.kafka.common.errors.TimeoutException("retriable sync failure"));
+
+        // Second round
+        expectSendRecordOnce(true);
+        expectSendRecordOnce(false);
+
+        PowerMock.replayAll();
+
+        // Try to send 3, make first pass, second fail. Should save last two
+        Whitebox.setInternalState(workerTask, "toSend", Arrays.asList(record1, record2, record3));
+        Whitebox.invokeMethod(workerTask, "sendRecords");
+        assertEquals(true, Whitebox.getInternalState(workerTask, "lastSendFailed"));
+        assertEquals(Arrays.asList(record2, record3), Whitebox.getInternalState(workerTask, "toSend"));
+
+        // Next they all succeed
+        Whitebox.invokeMethod(workerTask, "sendRecords");
+        assertEquals(false, Whitebox.getInternalState(workerTask, "lastSendFailed"));
+        assertNull(Whitebox.getInternalState(workerTask, "toSend"));
+
+        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);
+
+        expectTopicCreation(TOPIC);
+
+        expectSendRecordProducerCallbackFail();
+
+        PowerMock.replayAll();
+
+        Whitebox.setInternalState(workerTask, "toSend", Arrays.asList(record1, record2));
+        Whitebox.invokeMethod(workerTask, "sendRecords");
+    }
+
+    @Test(expected = ConnectException.class)
+    public void testSendRecordsProducerSendFailsImmediately() 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);
+
+        expectPreliminaryCalls();
+        expectTopicCreation(TOPIC);
+
+        EasyMock.expect(producer.send(EasyMock.anyObject(), EasyMock.anyObject()))
+                .andThrow(new KafkaException("Producer closed while send in progress", new InvalidTopicException(TOPIC)));
+
+        PowerMock.replayAll();
+
+        Whitebox.setInternalState(workerTask, "toSend", Arrays.asList(record1, record2));
+        Whitebox.invokeMethod(workerTask, "sendRecords");
+    }
+
+    @Test
+    public void testSendRecordsTaskCommitRecordFail() throws Exception {
+        createWorkerTask();
+
+        // Differentiate only by Kafka partition so we can reuse conversion expectations
+        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);
+        SourceRecord record3 = new SourceRecord(PARTITION, OFFSET, TOPIC, 3, KEY_SCHEMA, KEY, RECORD_SCHEMA, RECORD);
+
+        expectTopicCreation(TOPIC);
+
+        // Source task commit record failure will not cause the task to abort
+        expectSendRecordOnce(false);
+        expectSendRecordTaskCommitRecordFail(false, false);
+        expectSendRecordOnce(false);
+
+        PowerMock.replayAll();
+
+        Whitebox.setInternalState(workerTask, "toSend", Arrays.asList(record1, record2, record3));
+        Whitebox.invokeMethod(workerTask, "sendRecords");
+        assertEquals(false, Whitebox.getInternalState(workerTask, "lastSendFailed"));
+        assertNull(Whitebox.getInternalState(workerTask, "toSend"));
+
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testSlowTaskStart() throws Exception {
+        final CountDownLatch startupLatch = new CountDownLatch(1);
+        final CountDownLatch finishStartupLatch = new CountDownLatch(1);
+
+        createWorkerTask();
+
+        sourceTask.initialize(EasyMock.anyObject(SourceTaskContext.class));
+        EasyMock.expectLastCall();
+        sourceTask.start(TASK_PROPS);
+        EasyMock.expectLastCall().andAnswer(new IAnswer() {
+            @Override
+            public Object answer() throws Throwable {
+                startupLatch.countDown();
+                assertTrue(awaitLatch(finishStartupLatch));
+                return null;
+            }
+        });
+
+        statusListener.onStartup(taskId);
+        EasyMock.expectLastCall();
+
+        sourceTask.stop();
+        EasyMock.expectLastCall();
+        expectOffsetFlush(true);
+
+        statusListener.onShutdown(taskId);
+        EasyMock.expectLastCall();
+
+        expectClose();
+
+        PowerMock.replayAll();
+
+        workerTask.initialize(TASK_CONFIG);
+        Future workerTaskFuture = executor.submit(workerTask);
+
+        // Stopping immediately while the other thread has work to do should result in no polling, no offset commits,
+        // exiting the work thread immediately, and the stop() method will be invoked in the background thread since it
+        // cannot be invoked immediately in the thread trying to stop the task.
+        assertTrue(awaitLatch(startupLatch));
+        workerTask.stop();
+        finishStartupLatch.countDown();
+        assertTrue(workerTask.awaitStop(1000));
+
+        workerTaskFuture.get();
+
+        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);
+        SourceTaskMetricsGroup group1 = new SourceTaskMetricsGroup(taskId1, metrics);
+        for (int i = 0; i != 10; ++i) {
+            group.recordPoll(100, 1000 + i * 100);
+            group.recordWrite(10);
+        }
+        for (int i = 0; i != 20; ++i) {
+            group1.recordPoll(100, 1000 + i * 100);
+            group1.recordWrite(10);
+        }
+        assertEquals(1900.0, metrics.currentMetricValueAsDouble(group.metricGroup(), "poll-batch-max-time-ms"), 0.001d);
+        assertEquals(1450.0, metrics.currentMetricValueAsDouble(group.metricGroup(), "poll-batch-avg-time-ms"), 0.001d);
+        assertEquals(33.333, metrics.currentMetricValueAsDouble(group.metricGroup(), "source-record-poll-rate"), 0.001d);
+        assertEquals(1000, metrics.currentMetricValueAsDouble(group.metricGroup(), "source-record-poll-total"), 0.001d);
+        assertEquals(3.3333, metrics.currentMetricValueAsDouble(group.metricGroup(), "source-record-write-rate"), 0.001d);
+        assertEquals(100, metrics.currentMetricValueAsDouble(group.metricGroup(), "source-record-write-total"), 0.001d);
+        assertEquals(900.0, metrics.currentMetricValueAsDouble(group.metricGroup(), "source-record-active-count"), 0.001d);
+
+        // Close the group
+        group.close();
+
+        for (MetricName metricName : group.metricGroup().metrics().metrics().keySet()) {
+            // Metrics for this group should no longer exist
+            assertFalse(group.metricGroup().groupId().includes(metricName));
+        }
+        // Sensors for this group should no longer exist
+        assertNull(group.metricGroup().metrics().getSensor("sink-record-read"));
+        assertNull(group.metricGroup().metrics().getSensor("sink-record-send"));
+        assertNull(group.metricGroup().metrics().getSensor("sink-record-active-count"));
+        assertNull(group.metricGroup().metrics().getSensor("partition-count"));
+        assertNull(group.metricGroup().metrics().getSensor("offset-seq-number"));
+        assertNull(group.metricGroup().metrics().getSensor("offset-commit-completion"));
+        assertNull(group.metricGroup().metrics().getSensor("offset-commit-completion-skip"));
+        assertNull(group.metricGroup().metrics().getSensor("put-batch-time"));
+
+        assertEquals(2900.0, metrics.currentMetricValueAsDouble(group1.metricGroup(), "poll-batch-max-time-ms"), 0.001d);
+        assertEquals(1950.0, metrics.currentMetricValueAsDouble(group1.metricGroup(), "poll-batch-avg-time-ms"), 0.001d);
+        assertEquals(66.667, metrics.currentMetricValueAsDouble(group1.metricGroup(), "source-record-poll-rate"), 0.001d);
+        assertEquals(2000, metrics.currentMetricValueAsDouble(group1.metricGroup(), "source-record-poll-total"), 0.001d);
+        assertEquals(6.667, metrics.currentMetricValueAsDouble(group1.metricGroup(), "source-record-write-rate"), 0.001d);
+        assertEquals(200, metrics.currentMetricValueAsDouble(group1.metricGroup(), "source-record-write-total"), 0.001d);
+        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));
+
+        expectTopicCreation(TOPIC);
+
+        Capture> sent = expectSendRecord(TOPIC, 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));
+
+        expectTopicCreation(TOPIC);
+
+        Capture> sentRecordA = expectSendRecord(TOPIC, false, false, true, true, false, null);
+        Capture> sentRecordB = expectSendRecord(TOPIC, 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();
+    }
+
+    @Test
+    public void testTopicCreateWhenTopicExists() 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);
+
+        expectPreliminaryCalls();
+        TopicPartitionInfo topicPartitionInfo = new TopicPartitionInfo(0, null, Collections.emptyList(), Collections.emptyList());
+        TopicDescription topicDesc = new TopicDescription(TOPIC, false, Collections.singletonList(topicPartitionInfo));
+        EasyMock.expect(admin.describeTopics(TOPIC)).andReturn(Collections.singletonMap(TOPIC, topicDesc));
+
+        expectSendRecordTaskCommitRecordSucceed(false, false);
+        expectSendRecordTaskCommitRecordSucceed(false, false);
+
+        PowerMock.replayAll();
+
+        Whitebox.setInternalState(workerTask, "toSend", Arrays.asList(record1, record2));
+        Whitebox.invokeMethod(workerTask, "sendRecords");
+    }
+
+    @Test
+    public void testSendRecordsTopicDescribeRetries() 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);
+
+        expectPreliminaryCalls();
+        // First round - call to describe the topic times out
+        EasyMock.expect(admin.describeTopics(TOPIC))
+                .andThrow(new RetriableException(new TimeoutException("timeout")));
+
+        // Second round - calls to describe and create succeed
+        expectTopicCreation(TOPIC);
+        // Exactly two records are sent
+        expectSendRecordTaskCommitRecordSucceed(false, false);
+        expectSendRecordTaskCommitRecordSucceed(false, false);
+
+        PowerMock.replayAll();
+
+        Whitebox.setInternalState(workerTask, "toSend", Arrays.asList(record1, record2));
+        Whitebox.invokeMethod(workerTask, "sendRecords");
+        assertEquals(true, Whitebox.getInternalState(workerTask, "lastSendFailed"));
+        assertEquals(Arrays.asList(record1, record2), Whitebox.getInternalState(workerTask, "toSend"));
+
+        // Next they all succeed
+        Whitebox.invokeMethod(workerTask, "sendRecords");
+        assertEquals(false, Whitebox.getInternalState(workerTask, "lastSendFailed"));
+        assertNull(Whitebox.getInternalState(workerTask, "toSend"));
+    }
+
+    @Test
+    public void testSendRecordsTopicCreateRetries() 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);
+
+        // First call to describe the topic times out
+        expectPreliminaryCalls();
+        EasyMock.expect(admin.describeTopics(TOPIC)).andReturn(Collections.emptyMap());
+        Capture newTopicCapture = EasyMock.newCapture();
+        EasyMock.expect(admin.createTopic(EasyMock.capture(newTopicCapture)))
+                .andThrow(new RetriableException(new TimeoutException("timeout")));
+
+        // Second round
+        expectTopicCreation(TOPIC);
+        expectSendRecordTaskCommitRecordSucceed(false, false);
+        expectSendRecordTaskCommitRecordSucceed(false, false);
+
+        PowerMock.replayAll();
+
+        Whitebox.setInternalState(workerTask, "toSend", Arrays.asList(record1, record2));
+        Whitebox.invokeMethod(workerTask, "sendRecords");
+        assertEquals(true, Whitebox.getInternalState(workerTask, "lastSendFailed"));
+        assertEquals(Arrays.asList(record1, record2), Whitebox.getInternalState(workerTask, "toSend"));
+
+        // Next they all succeed
+        Whitebox.invokeMethod(workerTask, "sendRecords");
+        assertEquals(false, Whitebox.getInternalState(workerTask, "lastSendFailed"));
+        assertNull(Whitebox.getInternalState(workerTask, "toSend"));
+    }
+
+    @Test
+    public void testSendRecordsTopicDescribeRetriesMidway() throws Exception {
+        createWorkerTask();
+
+        // Differentiate only by Kafka partition so we can reuse conversion expectations
+        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);
+        SourceRecord record3 = new SourceRecord(PARTITION, OFFSET, OTHER_TOPIC, 3, KEY_SCHEMA, KEY, RECORD_SCHEMA, RECORD);
+
+        // First round
+        expectPreliminaryCalls(OTHER_TOPIC);
+        expectTopicCreation(TOPIC);
+        expectSendRecordTaskCommitRecordSucceed(false, false);
+        expectSendRecordTaskCommitRecordSucceed(false, false);
+
+        // First call to describe the topic times out
+        EasyMock.expect(admin.describeTopics(OTHER_TOPIC))
+                .andThrow(new RetriableException(new TimeoutException("timeout")));
+
+        // Second round
+        expectTopicCreation(OTHER_TOPIC);
+        expectSendRecord(OTHER_TOPIC, false, true, true, true, true, emptyHeaders());
+
+        PowerMock.replayAll();
+
+        // Try to send 3, make first pass, second fail. Should save last two
+        Whitebox.setInternalState(workerTask, "toSend", Arrays.asList(record1, record2, record3));
+        Whitebox.invokeMethod(workerTask, "sendRecords");
+        assertEquals(true, Whitebox.getInternalState(workerTask, "lastSendFailed"));
+        assertEquals(Arrays.asList(record3), Whitebox.getInternalState(workerTask, "toSend"));
+
+        // Next they all succeed
+        Whitebox.invokeMethod(workerTask, "sendRecords");
+        assertEquals(false, Whitebox.getInternalState(workerTask, "lastSendFailed"));
+        assertNull(Whitebox.getInternalState(workerTask, "toSend"));
+
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testSendRecordsTopicCreateRetriesMidway() throws Exception {
+        createWorkerTask();
+
+        // Differentiate only by Kafka partition so we can reuse conversion expectations
+        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);
+        SourceRecord record3 = new SourceRecord(PARTITION, OFFSET, OTHER_TOPIC, 3, KEY_SCHEMA, KEY, RECORD_SCHEMA, RECORD);
+
+        // First round
+        expectPreliminaryCalls(OTHER_TOPIC);
+        expectTopicCreation(TOPIC);
+        expectSendRecordTaskCommitRecordSucceed(false, false);
+        expectSendRecordTaskCommitRecordSucceed(false, false);
+
+        EasyMock.expect(admin.describeTopics(OTHER_TOPIC)).andReturn(Collections.emptyMap());
+        // First call to create the topic times out
+        Capture newTopicCapture = EasyMock.newCapture();
+        EasyMock.expect(admin.createTopic(EasyMock.capture(newTopicCapture)))
+                .andThrow(new RetriableException(new TimeoutException("timeout")));
+
+        // Second round
+        expectTopicCreation(OTHER_TOPIC);
+        expectSendRecord(OTHER_TOPIC, false, true, true, true, true, emptyHeaders());
+
+        PowerMock.replayAll();
+
+        // Try to send 3, make first pass, second fail. Should save last two
+        Whitebox.setInternalState(workerTask, "toSend", Arrays.asList(record1, record2, record3));
+        Whitebox.invokeMethod(workerTask, "sendRecords");
+        assertEquals(true, Whitebox.getInternalState(workerTask, "lastSendFailed"));
+        assertEquals(Arrays.asList(record3), Whitebox.getInternalState(workerTask, "toSend"));
+
+        // Next they all succeed
+        Whitebox.invokeMethod(workerTask, "sendRecords");
+        assertEquals(false, Whitebox.getInternalState(workerTask, "lastSendFailed"));
+        assertNull(Whitebox.getInternalState(workerTask, "toSend"));
+
+        PowerMock.verifyAll();
+    }
+
+    @Test(expected = ConnectException.class)
+    public void testTopicDescribeFails() 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);
+
+        expectPreliminaryCalls();
+        EasyMock.expect(admin.describeTopics(TOPIC))
+                .andThrow(new ConnectException(new TopicAuthorizationException("unauthorized")));
+
+        PowerMock.replayAll();
+
+        Whitebox.setInternalState(workerTask, "toSend", Arrays.asList(record1, record2));
+        Whitebox.invokeMethod(workerTask, "sendRecords");
+    }
+
+    @Test(expected = ConnectException.class)
+    public void testTopicCreateFails() 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);
+
+        expectPreliminaryCalls();
+        EasyMock.expect(admin.describeTopics(TOPIC)).andReturn(Collections.emptyMap());
+
+        Capture newTopicCapture = EasyMock.newCapture();
+        EasyMock.expect(admin.createTopic(EasyMock.capture(newTopicCapture)))
+                .andThrow(new ConnectException(new TopicAuthorizationException("unauthorized")));
+
+        PowerMock.replayAll();
+
+        Whitebox.setInternalState(workerTask, "toSend", Arrays.asList(record1, record2));
+        Whitebox.invokeMethod(workerTask, "sendRecords");
+        assertNotNull(newTopicCapture.getValue());
+    }
+
+    @Test(expected = ConnectException.class)
+    public void testTopicCreateFailsWithExceptionWhenCreateReturnsFalse() 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);
+
+        expectPreliminaryCalls();
+        EasyMock.expect(admin.describeTopics(TOPIC)).andReturn(Collections.emptyMap());
+
+        Capture newTopicCapture = EasyMock.newCapture();
+        EasyMock.expect(admin.createTopic(EasyMock.capture(newTopicCapture))).andReturn(false);
+
+        PowerMock.replayAll();
+
+        Whitebox.setInternalState(workerTask, "toSend", Arrays.asList(record1, record2));
+        Whitebox.invokeMethod(workerTask, "sendRecords");
+        assertNotNull(newTopicCapture.getValue());
+    }
+
+    private void expectPreliminaryCalls() {
+        expectPreliminaryCalls(TOPIC);
+    }
+
+    private void expectPreliminaryCalls(String topic) {
+        expectConvertHeadersAndKeyValue(topic, true, emptyHeaders());
+        expectApplyTransformationChain(false);
+        offsetWriter.offset(PARTITION, OFFSET);
+        PowerMock.expectLastCall();
+    }
+
+    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
+        // 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 RECORDS;
+                    }
+                });
+        // Fallout of the poll() call
+        expectSendRecordAnyTimes();
+        return latch;
+    }
+
+    private CountDownLatch expectPolls(int count) throws InterruptedException {
+        return expectPolls(count, new AtomicInteger());
+    }
+
+    @SuppressWarnings("unchecked")
+    private void expectSendRecordSyncFailure(Throwable error) throws InterruptedException {
+        expectConvertHeadersAndKeyValue(false);
+        expectApplyTransformationChain(false);
+
+        offsetWriter.offset(PARTITION, OFFSET);
+        PowerMock.expectLastCall();
+
+        EasyMock.expect(
+                producer.send(EasyMock.anyObject(ProducerRecord.class),
+                        EasyMock.anyObject(org.apache.kafka.clients.producer.Callback.class)))
+                .andThrow(error);
+    }
+
+    private Capture> expectSendRecordAnyTimes() throws InterruptedException {
+        return expectSendRecordTaskCommitRecordSucceed(true, false);
+    }
+
+    private Capture> expectSendRecordOnce(boolean isRetry) throws InterruptedException {
+        return expectSendRecordTaskCommitRecordSucceed(false, isRetry);
+    }
+
+    private Capture> expectSendRecordProducerCallbackFail() throws InterruptedException {
+        return expectSendRecord(TOPIC, false, false, false, false, true, emptyHeaders());
+    }
+
+    private Capture> expectSendRecordTaskCommitRecordSucceed(boolean anyTimes, boolean isRetry) throws InterruptedException {
+        return expectSendRecord(TOPIC, anyTimes, isRetry, true, true, true, emptyHeaders());
+    }
+
+    private Capture> expectSendRecordTaskCommitRecordFail(boolean anyTimes, boolean isRetry) throws InterruptedException {
+        return expectSendRecord(TOPIC, anyTimes, isRetry, true, false, true, emptyHeaders());
+    }
+
+    private Capture> expectSendRecord(boolean anyTimes, boolean isRetry, boolean succeed) throws InterruptedException {
+        return expectSendRecord(TOPIC, anyTimes, isRetry, succeed, true, true, emptyHeaders());
+    }
+
+    private Capture> expectSendRecord(
+            String topic,
+            boolean anyTimes,
+            boolean isRetry,
+            boolean sendSuccess,
+            boolean commitSuccess,
+            boolean isMockedConverters,
+            Headers headers
+    ) throws InterruptedException {
+        if (isMockedConverters) {
+            expectConvertHeadersAndKeyValue(topic, anyTimes, headers);
+        }
+
+        expectApplyTransformationChain(anyTimes);
+
+        Capture> sent = EasyMock.newCapture();
+
+        // 1. Offset data is passed to the offset storage.
+        if (!isRetry) {
+            offsetWriter.offset(PARTITION, OFFSET);
+            if (anyTimes)
+                PowerMock.expectLastCall().anyTimes();
+            else
+                PowerMock.expectLastCall();
+        }
+
+        // 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)));
+        IAnswer> expectResponse = new IAnswer>() {
+            @Override
+            public Future answer() throws Throwable {
+                synchronized (producerCallbacks) {
+                    for (org.apache.kafka.clients.producer.Callback cb : producerCallbacks.getValues()) {
+                        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();
+                }
+                return sendFuture;
+            }
+        };
+        if (anyTimes)
+            expect.andStubAnswer(expectResponse);
+        else
+            expect.andAnswer(expectResponse);
+
+        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);
+            expectTaskGetTopic(anyTimes);
+        }
+
+        return sent;
+    }
+
+    private void expectConvertHeadersAndKeyValue(boolean anyTimes) {
+        expectConvertHeadersAndKeyValue(TOPIC, anyTimes, emptyHeaders());
+    }
+
+    private void expectConvertHeadersAndKeyValue(String topic, 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, headers, RECORD_SCHEMA, RECORD));
+        if (anyTimes)
+            convertValueExpect.andStubReturn(SERIALIZED_RECORD);
+        else
+            convertValueExpect.andReturn(SERIALIZED_RECORD);
+    }
+
+    private void expectApplyTransformationChain(boolean anyTimes) {
+        final Capture recordCapture = EasyMock.newCapture();
+        IExpectationSetters convertKeyExpect = EasyMock.expect(transformationChain.apply(EasyMock.capture(recordCapture)));
+        if (anyTimes)
+            convertKeyExpect.andStubAnswer(new IAnswer() {
+                @Override
+                public SourceRecord answer() {
+                    return recordCapture.getValue();
+                }
+            });
+        else
+            convertKeyExpect.andAnswer(new IAnswer() {
+                @Override
+                public SourceRecord answer() {
+                    return recordCapture.getValue();
+                }
+            });
+    }
+
+    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"));
+        }
+        if (anyTimes) {
+            expect.anyTimes();
+        }
+    }
+
+    private void expectTaskGetTopic(boolean anyTimes) {
+        final Capture connectorCapture = EasyMock.newCapture();
+        final Capture topicCapture = EasyMock.newCapture();
+        IExpectationSetters expect = EasyMock.expect(statusBackingStore.getTopic(
+                EasyMock.capture(connectorCapture),
+                EasyMock.capture(topicCapture)));
+        if (anyTimes) {
+            expect.andStubAnswer(() -> new TopicStatus(
+                    topicCapture.getValue(),
+                    new ConnectorTaskId(connectorCapture.getValue(), 0),
+                    Time.SYSTEM.milliseconds()));
+        } else {
+            expect.andAnswer(() -> new TopicStatus(
+                    topicCapture.getValue(),
+                    new ConnectorTaskId(connectorCapture.getValue(), 0),
+                    Time.SYSTEM.milliseconds()));
+        }
+        if (connectorCapture.hasCaptured() && topicCapture.hasCaptured()) {
+            assertEquals("job", connectorCapture.getValue());
+            assertEquals(TOPIC, topicCapture.getValue());
+        }
+    }
+
+    private boolean awaitLatch(CountDownLatch latch) {
+        try {
+            return latch.await(5000, TimeUnit.MILLISECONDS);
+        } catch (InterruptedException e) {
+            // ignore
+        }
+        return false;
+    }
+
+    @SuppressWarnings("unchecked")
+    private void expectOffsetFlush(boolean succeed) throws Exception {
+        EasyMock.expect(offsetWriter.beginFlush()).andReturn(true);
+        Future flushFuture = PowerMock.createMock(Future.class);
+        EasyMock.expect(offsetWriter.doFlush(EasyMock.anyObject(Callback.class))).andReturn(flushFuture);
+        // Should throw for failure
+        IExpectationSetters futureGetExpect = EasyMock.expect(
+                flushFuture.get(EasyMock.anyLong(), EasyMock.anyObject(TimeUnit.class)));
+        if (succeed) {
+            sourceTask.commit();
+            EasyMock.expectLastCall();
+            futureGetExpect.andReturn(null);
+        } else {
+            futureGetExpect.andThrow(new TimeoutException());
+            offsetWriter.cancelFlush();
+            PowerMock.expectLastCall();
+        }
+    }
+
+    private void assertPollMetrics(int minimumPollCountExpected) {
+        MetricGroup sourceTaskGroup = workerTask.sourceTaskMetricsGroup().metricGroup();
+        MetricGroup taskGroup = workerTask.taskMetricsGroup().metricGroup();
+        double pollRate = metrics.currentMetricValueAsDouble(sourceTaskGroup, "source-record-poll-rate");
+        double pollTotal = metrics.currentMetricValueAsDouble(sourceTaskGroup, "source-record-poll-total");
+        if (minimumPollCountExpected > 0) {
+            assertEquals(RECORDS.size(), metrics.currentMetricValueAsDouble(taskGroup, "batch-size-max"), 0.000001d);
+            assertEquals(RECORDS.size(), metrics.currentMetricValueAsDouble(taskGroup, "batch-size-avg"), 0.000001d);
+            assertTrue(pollRate > 0.0d);
+        } else {
+            assertTrue(pollRate == 0.0d);
+        }
+        assertTrue(pollTotal >= minimumPollCountExpected);
+
+        double writeRate = metrics.currentMetricValueAsDouble(sourceTaskGroup, "source-record-write-rate");
+        double writeTotal = metrics.currentMetricValueAsDouble(sourceTaskGroup, "source-record-write-total");
+        if (minimumPollCountExpected > 0) {
+            assertTrue(writeRate > 0.0d);
+        } else {
+            assertTrue(writeRate == 0.0d);
+        }
+        assertTrue(writeTotal >= minimumPollCountExpected);
+
+        double pollBatchTimeMax = metrics.currentMetricValueAsDouble(sourceTaskGroup, "poll-batch-max-time-ms");
+        double pollBatchTimeAvg = metrics.currentMetricValueAsDouble(sourceTaskGroup, "poll-batch-avg-time-ms");
+        if (minimumPollCountExpected > 0) {
+            assertTrue(pollBatchTimeMax >= 0.0d);
+        }
+        assertTrue(Double.isNaN(pollBatchTimeAvg) || pollBatchTimeAvg > 0.0d);
+        double activeCount = metrics.currentMetricValueAsDouble(sourceTaskGroup, "source-record-active-count");
+        double activeCountMax = metrics.currentMetricValueAsDouble(sourceTaskGroup, "source-record-active-count-max");
+        assertEquals(0, activeCount, 0.000001d);
+        if (minimumPollCountExpected > 0) {
+            assertEquals(RECORDS.size(), activeCountMax, 0.000001d);
+        }
+    }
+
+    private RecordHeaders emptyHeaders() {
+        return new RecordHeaders();
+    }
+
+    private abstract static class TestSourceTask extends SourceTask {
+    }
+
+    private void expectClose() {
+        producer.close(EasyMock.anyObject(Duration.class));
+        EasyMock.expectLastCall();
+
+        admin.close(EasyMock.anyObject(Duration.class));
+        EasyMock.expectLastCall();
+
+        transformationChain.close();
+        EasyMock.expectLastCall();
+    }
+
+    private void expectTopicCreation(String topic) {
+        if (config.topicCreationEnable()) {
+            EasyMock.expect(admin.describeTopics(topic)).andReturn(Collections.emptyMap());
+            Capture newTopicCapture = EasyMock.newCapture();
+            EasyMock.expect(admin.createTopic(EasyMock.capture(newTopicCapture))).andReturn(true);
+        }
+    }
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerTaskTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerTaskTest.java
index 516b71ab3fcbd..c26a88c0a68f0 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerTaskTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerTaskTest.java
@@ -16,16 +16,25 @@
  */
 package org.apache.kafka.connect.runtime;
 
+import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.connect.errors.ConnectException;
 import org.apache.kafka.connect.runtime.WorkerTask.TaskMetricsGroup;
+import org.apache.kafka.connect.runtime.errors.RetryWithToleranceOperator;
+import org.apache.kafka.connect.runtime.errors.RetryWithToleranceOperatorTest;
 import org.apache.kafka.connect.sink.SinkTask;
+import org.apache.kafka.connect.storage.StatusBackingStore;
 import org.apache.kafka.connect.util.ConnectorTaskId;
-import org.apache.kafka.connect.util.MockTime;
+import org.apache.kafka.common.utils.MockTime;
 import org.easymock.EasyMock;
 import org.easymock.IAnswer;
+import org.easymock.Mock;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.powermock.core.classloader.annotations.PowerMockIgnore;
+import org.powermock.core.classloader.annotations.PrepareForTest;
+import org.powermock.modules.junit4.PowerMockRunner;
 
 import java.util.HashMap;
 import java.util.Map;
@@ -37,6 +46,9 @@
 import static org.easymock.EasyMock.verify;
 import static org.junit.Assert.assertEquals;
 
+@RunWith(PowerMockRunner.class)
+@PrepareForTest({WorkerTask.class})
+@PowerMockIgnore("javax.management.*")
 public class WorkerTaskTest {
 
     private static final Map TASK_PROPS = new HashMap<>();
@@ -46,10 +58,16 @@ public class WorkerTaskTest {
     private static final TaskConfig TASK_CONFIG = new TaskConfig(TASK_PROPS);
 
     private ConnectMetrics metrics;
+    @Mock private TaskStatus.Listener statusListener;
+    @Mock private ClassLoader loader;
+    RetryWithToleranceOperator retryWithToleranceOperator;
+    @Mock
+    StatusBackingStore statusBackingStore;
 
     @Before
     public void setup() {
         metrics = new MockConnectMetrics();
+        retryWithToleranceOperator = RetryWithToleranceOperatorTest.NOOP_OPERATOR;
     }
 
     @After
@@ -61,18 +79,19 @@ public void tearDown() {
     public void standardStartup() {
         ConnectorTaskId taskId = new ConnectorTaskId("foo", 0);
 
-        TaskStatus.Listener statusListener = EasyMock.createMock(TaskStatus.Listener.class);
-        ClassLoader loader = EasyMock.createMock(ClassLoader.class);
-
         WorkerTask workerTask = partialMockBuilder(WorkerTask.class)
                 .withConstructor(
                         ConnectorTaskId.class,
                         TaskStatus.Listener.class,
                         TargetState.class,
                         ClassLoader.class,
-                        ConnectMetrics.class
+                        ConnectMetrics.class,
+                        RetryWithToleranceOperator.class,
+                        Time.class,
+                        StatusBackingStore.class
                 )
-                .withArgs(taskId, statusListener, TargetState.STARTED, loader, metrics)
+                .withArgs(taskId, statusListener, TargetState.STARTED, loader, metrics,
+                        retryWithToleranceOperator, Time.SYSTEM, statusBackingStore)
                 .addMockedMethod("initialize")
                 .addMockedMethod("execute")
                 .addMockedMethod("close")
@@ -90,9 +109,6 @@ public void standardStartup() {
         workerTask.close();
         expectLastCall();
 
-        workerTask.releaseResources();
-        EasyMock.expectLastCall();
-
         statusListener.onShutdown(taskId);
         expectLastCall();
 
@@ -110,18 +126,19 @@ public void standardStartup() {
     public void stopBeforeStarting() {
         ConnectorTaskId taskId = new ConnectorTaskId("foo", 0);
 
-        TaskStatus.Listener statusListener = EasyMock.createMock(TaskStatus.Listener.class);
-        ClassLoader loader = EasyMock.createMock(ClassLoader.class);
-
         WorkerTask workerTask = partialMockBuilder(WorkerTask.class)
                 .withConstructor(
                         ConnectorTaskId.class,
                         TaskStatus.Listener.class,
                         TargetState.class,
                         ClassLoader.class,
-                        ConnectMetrics.class
+                        ConnectMetrics.class,
+                        RetryWithToleranceOperator.class,
+                        Time.class,
+                        StatusBackingStore.class
                 )
-                .withArgs(taskId, statusListener, TargetState.STARTED, loader, metrics)
+                .withArgs(taskId, statusListener, TargetState.STARTED, loader, metrics,
+                        retryWithToleranceOperator, Time.SYSTEM, statusBackingStore)
                 .addMockedMethod("initialize")
                 .addMockedMethod("execute")
                 .addMockedMethod("close")
@@ -133,9 +150,6 @@ public void stopBeforeStarting() {
         workerTask.close();
         EasyMock.expectLastCall();
 
-        workerTask.releaseResources();
-        EasyMock.expectLastCall();
-
         replay(workerTask);
 
         workerTask.initialize(TASK_CONFIG);
@@ -152,18 +166,19 @@ public void stopBeforeStarting() {
     public void cancelBeforeStopping() throws Exception {
         ConnectorTaskId taskId = new ConnectorTaskId("foo", 0);
 
-        TaskStatus.Listener statusListener = EasyMock.createMock(TaskStatus.Listener.class);
-        ClassLoader loader = EasyMock.createMock(ClassLoader.class);
-
         WorkerTask workerTask = partialMockBuilder(WorkerTask.class)
                 .withConstructor(
                         ConnectorTaskId.class,
                         TaskStatus.Listener.class,
                         TargetState.class,
                         ClassLoader.class,
-                        ConnectMetrics.class
+                        ConnectMetrics.class,
+                        RetryWithToleranceOperator.class,
+                        Time.class,
+                        StatusBackingStore.class
                 )
-                .withArgs(taskId, statusListener, TargetState.STARTED, loader, metrics)
+                .withArgs(taskId, statusListener, TargetState.STARTED, loader, metrics,
+                        retryWithToleranceOperator, Time.SYSTEM, statusBackingStore)
                 .addMockedMethod("initialize")
                 .addMockedMethod("execute")
                 .addMockedMethod("close")
@@ -198,9 +213,6 @@ public Void answer() throws Throwable {
         workerTask.close();
         expectLastCall();
 
-        workerTask.releaseResources();
-        EasyMock.expectLastCall();
-
         // there should be no call to onShutdown()
 
         replay(workerTask);
@@ -220,7 +232,6 @@ public Void answer() throws Throwable {
     public void updateMetricsOnListenerEventsForStartupPauseResumeAndShutdown() {
         ConnectorTaskId taskId = new ConnectorTaskId("foo", 0);
         ConnectMetrics metrics = new MockConnectMetrics();
-        TaskStatus.Listener statusListener = EasyMock.createMock(TaskStatus.Listener.class);
         TaskMetricsGroup group = new TaskMetricsGroup(taskId, metrics, statusListener);
 
         statusListener.onStartup(taskId);
@@ -255,7 +266,6 @@ public void updateMetricsOnListenerEventsForStartupPauseResumeAndFailure() {
         MockConnectMetrics metrics = new MockConnectMetrics();
         MockTime time = metrics.time();
         ConnectException error = new ConnectException("error");
-        TaskStatus.Listener statusListener = EasyMock.createMock(TaskStatus.Listener.class);
         TaskMetricsGroup group = new TaskMetricsGroup(taskId, metrics, statusListener);
 
         statusListener.onStartup(taskId);
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 e78ccc8f8c760..15933a7c3c9c2 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
@@ -17,68 +17,118 @@
 package org.apache.kafka.connect.runtime;
 
 import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
 import org.apache.kafka.clients.producer.KafkaProducer;
+import org.apache.kafka.clients.producer.ProducerConfig;
+import org.apache.kafka.common.Configurable;
+import org.apache.kafka.common.MetricName;
+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.metrics.MetricsReporter;
+import org.apache.kafka.common.metrics.stats.Avg;
+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;
 import org.apache.kafka.connect.json.JsonConverter;
+import org.apache.kafka.connect.json.JsonConverterConfig;
 import org.apache.kafka.connect.runtime.ConnectMetrics.MetricGroup;
 import org.apache.kafka.connect.runtime.MockConnectMetrics.MockMetricsReporter;
+import org.apache.kafka.connect.runtime.distributed.ClusterConfigState;
+import org.apache.kafka.connect.runtime.errors.RetryWithToleranceOperator;
 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.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.SinkConnector;
 import org.apache.kafka.connect.sink.SinkTask;
+import org.apache.kafka.connect.source.SourceConnector;
 import org.apache.kafka.connect.source.SourceRecord;
 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.OffsetBackingStore;
 import org.apache.kafka.connect.storage.OffsetStorageReader;
 import org.apache.kafka.connect.storage.OffsetStorageWriter;
+import org.apache.kafka.connect.storage.StatusBackingStore;
+import org.apache.kafka.connect.util.ConnectUtils;
 import org.apache.kafka.connect.util.ConnectorTaskId;
-import org.apache.kafka.connect.util.MockTime;
+import org.apache.kafka.connect.util.FutureCallback;
 import org.apache.kafka.connect.util.ThreadedTest;
-import org.easymock.Capture;
+import org.apache.kafka.connect.util.TopicAdmin;
+import org.apache.kafka.connect.util.TopicCreationGroup;
 import org.easymock.EasyMock;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.powermock.api.easymock.PowerMock;
 import org.powermock.api.easymock.annotation.Mock;
+import org.powermock.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;
 import org.powermock.modules.junit4.PowerMockRunner;
 
+import javax.management.MBeanServer;
+import javax.management.ObjectInstance;
+import javax.management.ObjectName;
+import java.lang.management.ManagementFactory;
 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.UUID;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.TimeUnit;
+
+import static org.apache.kafka.connect.runtime.WorkerConfig.TOPIC_CREATION_ENABLE_CONFIG;
+import static org.apache.kafka.connect.runtime.errors.RetryWithToleranceOperatorTest.NOOP_OPERATOR;
+import static org.easymock.EasyMock.anyObject;
+import static org.easymock.EasyMock.eq;
+import static org.easymock.EasyMock.expectLastCall;
+import static org.hamcrest.CoreMatchers.instanceOf;
+import static org.hamcrest.MatcherAssert.assertThat;
 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)
-@PrepareForTest({Worker.class, Plugins.class})
+@PrepareForTest({Worker.class, Plugins.class, ConnectUtils.class})
 @PowerMockIgnore("javax.management.*")
 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 static final String CLUSTER_ID = "test-cluster";
+    private final ConnectorClientConfigOverridePolicy noneConnectorClientConfigOverridePolicy = new NoneConnectorClientConfigOverridePolicy();
+    private final ConnectorClientConfigOverridePolicy allConnectorClientConfigOverridePolicy = new AllConnectorClientConfigOverridePolicy();
 
+    private Map workerProps = new HashMap<>();
     private WorkerConfig config;
     private Worker worker;
 
+    private Map defaultProducerConfigs = new HashMap<>();
+    private Map defaultConsumerConfigs = new HashMap<>();
+
     @Mock
     private Plugins plugins;
     @Mock
@@ -92,86 +142,137 @@ public class WorkerTest extends ThreadedTest {
     @MockStrict
     private ConnectorStatus.Listener connectorStatusListener;
 
+    @Mock private Herder herder;
+    @Mock private StatusBackingStore statusBackingStore;
+    @Mock private SourceConnector sourceConnector;
+    @Mock private SinkConnector sinkConnector;
+    @Mock private CloseableConnectorContext ctx;
+    @Mock private TestSourceTask task;
+    @Mock private WorkerSourceTask workerTask;
+    @Mock private Converter keyConverter;
+    @Mock private Converter valueConverter;
+    @Mock private Converter taskKeyConverter;
+    @Mock private Converter taskValueConverter;
+    @Mock private HeaderConverter taskHeaderConverter;
+    @Mock private ExecutorService executorService;
+    @MockNice private ConnectorConfig connectorConfig;
+    private String mockFileProviderTestId;
+    private Map connectorProps;
+
+    // when this test becomes parameterized, this variable will be a test parameter
+    public boolean enableTopicCreation = false;
+
     @Before
     public void setup() {
         super.setup();
-
-        Map workerProps = new HashMap<>();
         workerProps.put("key.converter", "org.apache.kafka.connect.json.JsonConverter");
         workerProps.put("value.converter", "org.apache.kafka.connect.json.JsonConverter");
-        workerProps.put("internal.key.converter", "org.apache.kafka.connect.json.JsonConverter");
-        workerProps.put("internal.value.converter", "org.apache.kafka.connect.json.JsonConverter");
-        workerProps.put("internal.key.converter.schemas.enable", "false");
-        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);
+        workerProps.put(TOPIC_CREATION_ENABLE_CONFIG, String.valueOf(enableTopicCreation));
         config = new StandaloneConfig(workerProps);
 
+        defaultProducerConfigs.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092");
+        defaultProducerConfigs.put(
+            ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.ByteArraySerializer");
+        defaultProducerConfigs.put(
+            ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.ByteArraySerializer");
+        defaultProducerConfigs.put(ProducerConfig.REQUEST_TIMEOUT_MS_CONFIG, Integer.toString(Integer.MAX_VALUE));
+        defaultProducerConfigs.put(ProducerConfig.MAX_BLOCK_MS_CONFIG, Long.toString(Long.MAX_VALUE));
+        defaultProducerConfigs.put(ProducerConfig.ACKS_CONFIG, "all");
+        defaultProducerConfigs.put(ProducerConfig.MAX_IN_FLIGHT_REQUESTS_PER_CONNECTION, "1");
+        defaultProducerConfigs.put(ProducerConfig.DELIVERY_TIMEOUT_MS_CONFIG, Integer.toString(Integer.MAX_VALUE));
+
+        defaultConsumerConfigs.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092");
+        defaultConsumerConfigs.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false");
+        defaultConsumerConfigs.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
+        defaultConsumerConfigs
+            .put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.ByteArrayDeserializer");
+        defaultConsumerConfigs
+            .put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.ByteArrayDeserializer");
+
+        // Some common defaults. They might change on individual tests
+        connectorProps = anyConnectorConfigMap();
         PowerMock.mockStatic(Plugins.class);
     }
 
     @Test
-    public void testStartAndStopConnector() throws Exception {
+    public void testStartAndStopConnector() throws Throwable {
         expectConverters();
         expectStartStorage();
 
-        // Create
-        Connector connector = PowerMock.createMock(Connector.class);
-        ConnectorContext ctx = PowerMock.createMock(ConnectorContext.class);
+        final String connectorClass = WorkerTestConnector.class.getName();
 
+        // Create
         EasyMock.expect(plugins.currentThreadLoader()).andReturn(delegatingLoader).times(2);
-        EasyMock.expect(plugins.newConnector(WorkerTestConnector.class.getName()))
-                .andReturn(connector);
-        EasyMock.expect(connector.version()).andReturn("1.0");
+        EasyMock.expect(plugins.delegatingLoader()).andReturn(delegatingLoader);
+        EasyMock.expect(delegatingLoader.connectorLoader(connectorClass)).andReturn(pluginLoader);
+        EasyMock.expect(plugins.newConnector(connectorClass))
+                .andReturn(sourceConnector);
+        EasyMock.expect(sourceConnector.version()).andReturn("1.0");
 
-        Map props = new HashMap<>();
-        props.put(SinkConnectorConfig.TOPICS_CONFIG, "foo,bar");
-        props.put(ConnectorConfig.TASKS_MAX_CONFIG, "1");
-        props.put(ConnectorConfig.NAME_CONFIG, CONNECTOR_ID);
-        props.put(ConnectorConfig.CONNECTOR_CLASS_CONFIG, WorkerTestConnector.class.getName());
+        connectorProps.put(ConnectorConfig.CONNECTOR_CLASS_CONFIG, connectorClass);
 
-        EasyMock.expect(connector.version()).andReturn("1.0");
+        EasyMock.expect(sourceConnector.version()).andReturn("1.0");
 
-        EasyMock.expect(plugins.compareAndSwapLoaders(connector))
+        expectFileConfigProvider();
+        EasyMock.expect(Plugins.compareAndSwapLoaders(pluginLoader))
                 .andReturn(delegatingLoader)
-                .times(2);
-        connector.initialize(EasyMock.anyObject(ConnectorContext.class));
+                .times(3);
+        sourceConnector.initialize(anyObject(ConnectorContext.class));
         EasyMock.expectLastCall();
-        connector.start(props);
+        sourceConnector.start(connectorProps);
         EasyMock.expectLastCall();
 
         EasyMock.expect(Plugins.compareAndSwapLoaders(delegatingLoader))
-                .andReturn(pluginLoader).times(2);
+                .andReturn(pluginLoader).times(3);
 
         connectorStatusListener.onStartup(CONNECTOR_ID);
         EasyMock.expectLastCall();
 
         // Remove
-        connector.stop();
+        sourceConnector.stop();
         EasyMock.expectLastCall();
 
         connectorStatusListener.onShutdown(CONNECTOR_ID);
         EasyMock.expectLastCall();
 
+        ctx.close();
+        expectLastCall();
+
         expectStopStorage();
+        expectClusterId();
 
         PowerMock.replayAll();
 
-        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore);
+        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore, noneConnectorClientConfigOverridePolicy);
+        worker.herder = herder;
         worker.start();
 
         assertEquals(Collections.emptySet(), worker.connectorNames());
-        worker.startConnector(CONNECTOR_ID, props, ctx, connectorStatusListener, TargetState.STARTED);
+
+        FutureCallback onFirstStart = new FutureCallback<>();
+        worker.startConnector(CONNECTOR_ID, connectorProps, ctx, connectorStatusListener, TargetState.STARTED, onFirstStart);
+        // Wait for the connector to actually start
+        assertEquals(TargetState.STARTED, onFirstStart.get(1000, TimeUnit.MILLISECONDS));
         assertEquals(new HashSet<>(Arrays.asList(CONNECTOR_ID)), worker.connectorNames());
+
+        FutureCallback onSecondStart = new FutureCallback<>();
+        worker.startConnector(CONNECTOR_ID, connectorProps, ctx, connectorStatusListener, TargetState.STARTED, onSecondStart);
         try {
-            worker.startConnector(CONNECTOR_ID, props, ctx, connectorStatusListener, TargetState.STARTED);
-            fail("Should have thrown exception when trying to add connector with same name.");
-        } catch (ConnectException e) {
-            // expected
+            onSecondStart.get(0, TimeUnit.MILLISECONDS);
+            fail("Should have failed while trying to start second connector with same name");
+        } catch (ExecutionException e) {
+            assertThat(e.getCause(), instanceOf(ConnectException.class));
         }
+
         assertStatistics(worker, 1, 0);
         assertStartupStatistics(worker, 1, 0, 0, 0);
-        worker.stopConnector(CONNECTOR_ID);
+        worker.stopAndAwaitConnector(CONNECTOR_ID);
         assertStatistics(worker, 0, 0);
         assertStartupStatistics(worker, 1, 0, 0, 0);
         assertEquals(Collections.emptySet(), worker.connectorNames());
@@ -180,27 +281,36 @@ public void testStartAndStopConnector() throws Exception {
         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() throws Exception {
         expectConverters();
         expectStartStorage();
+        expectFileConfigProvider();
 
-        ConnectorContext ctx = PowerMock.createMock(ConnectorContext.class);
-
-        Map props = new HashMap<>();
-        props.put(SinkConnectorConfig.TOPICS_CONFIG, "foo,bar");
-        props.put(ConnectorConfig.TASKS_MAX_CONFIG, "1");
-        props.put(ConnectorConfig.NAME_CONFIG, CONNECTOR_ID);
-        props.put(ConnectorConfig.CONNECTOR_CLASS_CONFIG, "java.util.HashMap"); // Bad connector class name
+        final String nonConnectorClass = "java.util.HashMap";
+        connectorProps.put(ConnectorConfig.CONNECTOR_CLASS_CONFIG, nonConnectorClass); // Bad connector class name
 
+        Exception exception = new ConnectException("Failed to find Connector");
         EasyMock.expect(plugins.currentThreadLoader()).andReturn(delegatingLoader);
-        EasyMock.expect(plugins.newConnector(EasyMock.anyString()))
-                .andThrow(new ConnectException("Failed to find Connector"));
-
+        EasyMock.expect(plugins.delegatingLoader()).andReturn(delegatingLoader);
+        EasyMock.expect(delegatingLoader.connectorLoader(nonConnectorClass)).andReturn(delegatingLoader);
         EasyMock.expect(Plugins.compareAndSwapLoaders(delegatingLoader))
-                .andReturn(pluginLoader);
+                .andReturn(delegatingLoader).times(2);
+        EasyMock.expect(plugins.newConnector(EasyMock.anyString()))
+                .andThrow(exception);
 
         connectorStatusListener.onFailure(
                 EasyMock.eq(CONNECTOR_ID),
@@ -208,20 +318,30 @@ public void testStartConnectorFailure() throws Exception {
         );
         EasyMock.expectLastCall();
 
+        expectClusterId();
+
         PowerMock.replayAll();
 
-        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore);
+        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore, noneConnectorClientConfigOverridePolicy);
+        worker.herder = herder;
         worker.start();
 
         assertStatistics(worker, 0, 0);
-        assertFalse(worker.startConnector(CONNECTOR_ID, props, ctx, connectorStatusListener, TargetState.STARTED));
+        FutureCallback onStart = new FutureCallback<>();
+        worker.startConnector(CONNECTOR_ID, connectorProps, ctx, connectorStatusListener, TargetState.STARTED, onStart);
+        try {
+            onStart.get(0, TimeUnit.MILLISECONDS);
+            fail("Should have failed to start connector");
+        } catch (ExecutionException e) {
+            assertEquals(exception, e.getCause());
+        }
 
         assertStartupStatistics(worker, 1, 1, 0, 0);
         assertEquals(Collections.emptySet(), worker.connectorNames());
 
         assertStatistics(worker, 0, 0);
         assertStartupStatistics(worker, 1, 1, 0, 0);
-        assertFalse(worker.stopConnector(CONNECTOR_ID));
+        worker.stopAndAwaitConnector(CONNECTOR_ID);
         assertStatistics(worker, 0, 0);
         assertStartupStatistics(worker, 1, 1, 0, 0);
 
@@ -229,62 +349,68 @@ public void testStartConnectorFailure() throws Exception {
     }
 
     @Test
-    public void testAddConnectorByAlias() throws Exception {
+    public void testAddConnectorByAlias() throws Throwable {
         expectConverters();
         expectStartStorage();
+        expectFileConfigProvider();
 
-        // Create
-        Connector connector = PowerMock.createMock(Connector.class);
-        ConnectorContext ctx = PowerMock.createMock(ConnectorContext.class);
+        final String connectorAlias = "WorkerTestConnector";
 
         EasyMock.expect(plugins.currentThreadLoader()).andReturn(delegatingLoader).times(2);
-        EasyMock.expect(plugins.newConnector("WorkerTestConnector")).andReturn(connector);
-        EasyMock.expect(connector.version()).andReturn("1.0");
+        EasyMock.expect(plugins.delegatingLoader()).andReturn(delegatingLoader);
+        EasyMock.expect(delegatingLoader.connectorLoader(connectorAlias)).andReturn(pluginLoader);
+        EasyMock.expect(plugins.newConnector(connectorAlias)).andReturn(sinkConnector);
+        EasyMock.expect(sinkConnector.version()).andReturn("1.0");
 
-        Map props = new HashMap<>();
-        props.put(SinkConnectorConfig.TOPICS_CONFIG, "foo,bar");
-        props.put(ConnectorConfig.TASKS_MAX_CONFIG, "1");
-        props.put(ConnectorConfig.NAME_CONFIG, CONNECTOR_ID);
-        props.put(ConnectorConfig.CONNECTOR_CLASS_CONFIG, "WorkerTestConnector");
+        connectorProps.put(ConnectorConfig.CONNECTOR_CLASS_CONFIG, connectorAlias);
+        connectorProps.put(SinkConnectorConfig.TOPICS_CONFIG, "gfieyls, wfru");
 
-        EasyMock.expect(connector.version()).andReturn("1.0");
-        EasyMock.expect(plugins.compareAndSwapLoaders(connector))
+        EasyMock.expect(sinkConnector.version()).andReturn("1.0");
+        EasyMock.expect(Plugins.compareAndSwapLoaders(pluginLoader))
                 .andReturn(delegatingLoader)
-                .times(2);
-        connector.initialize(EasyMock.anyObject(ConnectorContext.class));
+                .times(3);
+        sinkConnector.initialize(anyObject(ConnectorContext.class));
         EasyMock.expectLastCall();
-        connector.start(props);
+        sinkConnector.start(connectorProps);
         EasyMock.expectLastCall();
 
         EasyMock.expect(Plugins.compareAndSwapLoaders(delegatingLoader))
                 .andReturn(pluginLoader)
-                .times(2);
+                .times(3);
 
         connectorStatusListener.onStartup(CONNECTOR_ID);
         EasyMock.expectLastCall();
 
         // Remove
-        connector.stop();
+        sinkConnector.stop();
         EasyMock.expectLastCall();
 
         connectorStatusListener.onShutdown(CONNECTOR_ID);
         EasyMock.expectLastCall();
 
+        ctx.close();
+        expectLastCall();
+
         expectStopStorage();
+        expectClusterId();
 
         PowerMock.replayAll();
 
-        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore);
+        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore, noneConnectorClientConfigOverridePolicy);
+        worker.herder = herder;
         worker.start();
 
         assertStatistics(worker, 0, 0);
         assertEquals(Collections.emptySet(), worker.connectorNames());
-        worker.startConnector(CONNECTOR_ID, props, ctx, connectorStatusListener, TargetState.STARTED);
+        FutureCallback onStart = new FutureCallback<>();
+        worker.startConnector(CONNECTOR_ID, connectorProps, ctx, connectorStatusListener, TargetState.STARTED, onStart);
+        // Wait for the connector to actually start
+        assertEquals(TargetState.STARTED, onStart.get(1000, TimeUnit.MILLISECONDS));
         assertEquals(new HashSet<>(Arrays.asList(CONNECTOR_ID)), worker.connectorNames());
         assertStatistics(worker, 1, 0);
         assertStartupStatistics(worker, 1, 0, 0, 0);
 
-        worker.stopConnector(CONNECTOR_ID);
+        worker.stopAndAwaitConnector(CONNECTOR_ID);
         assertStatistics(worker, 0, 0);
         assertStartupStatistics(worker, 1, 0, 0, 0);
         assertEquals(Collections.emptySet(), worker.connectorNames());
@@ -297,61 +423,67 @@ public void testAddConnectorByAlias() throws Exception {
     }
 
     @Test
-    public void testAddConnectorByShortAlias() throws Exception {
+    public void testAddConnectorByShortAlias() throws Throwable {
         expectConverters();
         expectStartStorage();
+        expectFileConfigProvider();
 
-        // Create
-        Connector connector = PowerMock.createMock(Connector.class);
-        ConnectorContext ctx = PowerMock.createMock(ConnectorContext.class);
+        final String shortConnectorAlias = "WorkerTest";
 
         EasyMock.expect(plugins.currentThreadLoader()).andReturn(delegatingLoader).times(2);
-        EasyMock.expect(plugins.newConnector("WorkerTest")).andReturn(connector);
-        EasyMock.expect(connector.version()).andReturn("1.0");
+        EasyMock.expect(plugins.delegatingLoader()).andReturn(delegatingLoader);
+        EasyMock.expect(delegatingLoader.connectorLoader(shortConnectorAlias)).andReturn(pluginLoader);
+        EasyMock.expect(plugins.newConnector(shortConnectorAlias)).andReturn(sinkConnector);
+        EasyMock.expect(sinkConnector.version()).andReturn("1.0");
 
-        Map props = new HashMap<>();
-        props.put(SinkConnectorConfig.TOPICS_CONFIG, "foo,bar");
-        props.put(ConnectorConfig.TASKS_MAX_CONFIG, "1");
-        props.put(ConnectorConfig.NAME_CONFIG, CONNECTOR_ID);
-        props.put(ConnectorConfig.CONNECTOR_CLASS_CONFIG, "WorkerTest");
+        connectorProps.put(ConnectorConfig.CONNECTOR_CLASS_CONFIG, shortConnectorAlias);
+        connectorProps.put(SinkConnectorConfig.TOPICS_CONFIG, "gfieyls, wfru");
 
-        EasyMock.expect(connector.version()).andReturn("1.0");
-        EasyMock.expect(plugins.compareAndSwapLoaders(connector))
+        EasyMock.expect(sinkConnector.version()).andReturn("1.0");
+        EasyMock.expect(Plugins.compareAndSwapLoaders(pluginLoader))
                 .andReturn(delegatingLoader)
-                .times(2);
-        connector.initialize(EasyMock.anyObject(ConnectorContext.class));
+                .times(3);
+        sinkConnector.initialize(anyObject(ConnectorContext.class));
         EasyMock.expectLastCall();
-        connector.start(props);
+        sinkConnector.start(connectorProps);
         EasyMock.expectLastCall();
 
         EasyMock.expect(Plugins.compareAndSwapLoaders(delegatingLoader))
                 .andReturn(pluginLoader)
-                .times(2);
+                .times(3);
 
         connectorStatusListener.onStartup(CONNECTOR_ID);
         EasyMock.expectLastCall();
 
         // Remove
-        connector.stop();
+        sinkConnector.stop();
         EasyMock.expectLastCall();
 
         connectorStatusListener.onShutdown(CONNECTOR_ID);
         EasyMock.expectLastCall();
 
+        ctx.close();
+        expectLastCall();
+
         expectStopStorage();
+        expectClusterId();
 
         PowerMock.replayAll();
 
-        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore);
+        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore, noneConnectorClientConfigOverridePolicy);
+        worker.herder = herder;
         worker.start();
 
         assertStatistics(worker, 0, 0);
         assertEquals(Collections.emptySet(), worker.connectorNames());
-        worker.startConnector(CONNECTOR_ID, props, ctx, connectorStatusListener, TargetState.STARTED);
+        FutureCallback onStart = new FutureCallback<>();
+        worker.startConnector(CONNECTOR_ID, connectorProps, ctx, connectorStatusListener, TargetState.STARTED, onStart);
+        // Wait for the connector to actually start
+        assertEquals(TargetState.STARTED, onStart.get(1000, TimeUnit.MILLISECONDS));
         assertEquals(new HashSet<>(Arrays.asList(CONNECTOR_ID)), worker.connectorNames());
         assertStatistics(worker, 1, 0);
 
-        worker.stopConnector(CONNECTOR_ID);
+        worker.stopAndAwaitConnector(CONNECTOR_ID);
         assertStatistics(worker, 0, 0);
         assertEquals(Collections.emptySet(), worker.connectorNames());
         // Nothing should be left, so this should effectively be a nop
@@ -365,85 +497,98 @@ public void testAddConnectorByShortAlias() throws Exception {
     public void testStopInvalidConnector() {
         expectConverters();
         expectStartStorage();
+        expectFileConfigProvider();
+        expectClusterId();
 
         PowerMock.replayAll();
 
-        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore);
+        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore, noneConnectorClientConfigOverridePolicy);
+        worker.herder = herder;
         worker.start();
 
-        worker.stopConnector(CONNECTOR_ID);
+        worker.stopAndAwaitConnector(CONNECTOR_ID);
 
         PowerMock.verifyAll();
     }
 
     @Test
-    public void testReconfigureConnectorTasks() throws Exception {
+    public void testReconfigureConnectorTasks() throws Throwable {
         expectConverters();
         expectStartStorage();
+        expectFileConfigProvider();
 
-        // Create
-        Connector connector = PowerMock.createMock(Connector.class);
-        ConnectorContext ctx = PowerMock.createMock(ConnectorContext.class);
+        final String connectorClass = WorkerTestConnector.class.getName();
 
         EasyMock.expect(plugins.currentThreadLoader()).andReturn(delegatingLoader).times(3);
-        EasyMock.expect(plugins.newConnector(WorkerTestConnector.class.getName()))
-                .andReturn(connector);
-        EasyMock.expect(connector.version()).andReturn("1.0");
+        EasyMock.expect(plugins.delegatingLoader()).andReturn(delegatingLoader).times(1);
+        EasyMock.expect(delegatingLoader.connectorLoader(connectorClass)).andReturn(pluginLoader);
+        EasyMock.expect(plugins.newConnector(connectorClass))
+                .andReturn(sinkConnector);
+        EasyMock.expect(sinkConnector.version()).andReturn("1.0");
 
-        Map props = new HashMap<>();
-        props.put(SinkConnectorConfig.TOPICS_CONFIG, "foo,bar");
-        props.put(ConnectorConfig.TASKS_MAX_CONFIG, "1");
-        props.put(ConnectorConfig.NAME_CONFIG, CONNECTOR_ID);
-        props.put(ConnectorConfig.CONNECTOR_CLASS_CONFIG, WorkerTestConnector.class.getName());
+        connectorProps.put(SinkConnectorConfig.TOPICS_CONFIG, "foo,bar");
+        connectorProps.put(ConnectorConfig.CONNECTOR_CLASS_CONFIG, connectorClass);
 
-        EasyMock.expect(connector.version()).andReturn("1.0");
-        EasyMock.expect(plugins.compareAndSwapLoaders(connector))
+        EasyMock.expect(sinkConnector.version()).andReturn("1.0");
+        EasyMock.expect(Plugins.compareAndSwapLoaders(pluginLoader))
                 .andReturn(delegatingLoader)
-                .times(3);
-        connector.initialize(EasyMock.anyObject(ConnectorContext.class));
+                .times(4);
+        sinkConnector.initialize(anyObject(ConnectorContext.class));
         EasyMock.expectLastCall();
-        connector.start(props);
+        sinkConnector.start(connectorProps);
         EasyMock.expectLastCall();
 
         EasyMock.expect(Plugins.compareAndSwapLoaders(delegatingLoader))
                 .andReturn(pluginLoader)
-                .times(3);
+                .times(4);
 
         connectorStatusListener.onStartup(CONNECTOR_ID);
         EasyMock.expectLastCall();
 
         // Reconfigure
-        EasyMock.>expect(connector.taskClass()).andReturn(TestSourceTask.class);
+        EasyMock.>expect(sinkConnector.taskClass()).andReturn(TestSourceTask.class);
         Map taskProps = new HashMap<>();
         taskProps.put("foo", "bar");
-        EasyMock.expect(connector.taskConfigs(2)).andReturn(Arrays.asList(taskProps, taskProps));
+        EasyMock.expect(sinkConnector.taskConfigs(2)).andReturn(Arrays.asList(taskProps, taskProps));
 
         // Remove
-        connector.stop();
+        sinkConnector.stop();
         EasyMock.expectLastCall();
 
         connectorStatusListener.onShutdown(CONNECTOR_ID);
         EasyMock.expectLastCall();
 
+        ctx.close();
+        expectLastCall();
+
         expectStopStorage();
+        expectClusterId();
 
         PowerMock.replayAll();
 
-        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore);
+        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore, noneConnectorClientConfigOverridePolicy);
+        worker.herder = herder;
         worker.start();
 
         assertStatistics(worker, 0, 0);
         assertEquals(Collections.emptySet(), worker.connectorNames());
-        worker.startConnector(CONNECTOR_ID, props, ctx, connectorStatusListener, TargetState.STARTED);
+        FutureCallback onFirstStart = new FutureCallback<>();
+        worker.startConnector(CONNECTOR_ID, connectorProps, ctx, connectorStatusListener, TargetState.STARTED, onFirstStart);
+        // Wait for the connector to actually start
+        assertEquals(TargetState.STARTED, onFirstStart.get(1000, TimeUnit.MILLISECONDS));
         assertStatistics(worker, 1, 0);
         assertEquals(new HashSet<>(Arrays.asList(CONNECTOR_ID)), worker.connectorNames());
+
+        FutureCallback onSecondStart = new FutureCallback<>();
+        worker.startConnector(CONNECTOR_ID, connectorProps, ctx, connectorStatusListener, TargetState.STARTED, onSecondStart);
         try {
-            worker.startConnector(CONNECTOR_ID, props, ctx, connectorStatusListener, TargetState.STARTED);
-            fail("Should have thrown exception when trying to add connector with same name.");
-        } catch (ConnectException e) {
-            // expected
+            onSecondStart.get(0, TimeUnit.MILLISECONDS);
+            fail("Should have failed while trying to start second connector with same name");
+        } catch (ExecutionException e) {
+            assertThat(e.getCause(), instanceOf(ConnectException.class));
         }
-        Map connProps = new HashMap<>(props);
+
+        Map connProps = new HashMap<>(connectorProps);
         connProps.put(ConnectorConfig.TASKS_MAX_CONFIG, "2");
         ConnectorConfig connConfig = new SinkConnectorConfig(plugins, connProps);
         List> taskConfigs = worker.connectorTaskConfigs(CONNECTOR_ID, connConfig);
@@ -456,7 +601,7 @@ public void testReconfigureConnectorTasks() throws Exception {
         assertEquals(expectedTaskProps, taskConfigs.get(1));
         assertStatistics(worker, 1, 0);
         assertStartupStatistics(worker, 1, 0, 0, 0);
-        worker.stopConnector(CONNECTOR_ID);
+        worker.stopAndAwaitConnector(CONNECTOR_ID);
         assertStatistics(worker, 0, 0);
         assertStartupStatistics(worker, 1, 0, 0, 0);
         assertEquals(Collections.emptySet(), worker.connectorNames());
@@ -467,34 +612,16 @@ public void testReconfigureConnectorTasks() throws Exception {
         PowerMock.verifyAll();
     }
 
-
     @Test
     public void testAddRemoveTask() throws Exception {
-        expectConverters(true);
+        expectConverters();
         expectStartStorage();
+        expectFileConfigProvider();
 
-        // Create
-        TestSourceTask task = PowerMock.createMock(TestSourceTask.class);
-        WorkerSourceTask workerTask = PowerMock.createMock(WorkerSourceTask.class);
         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),
-                EasyMock.anyObject(TaskStatus.Listener.class),
-                EasyMock.eq(TargetState.STARTED),
-                EasyMock.anyObject(JsonConverter.class),
-                EasyMock.anyObject(JsonConverter.class),
-                EasyMock.eq(TransformationChain.noOp()),
-                EasyMock.anyObject(KafkaProducer.class),
-                EasyMock.anyObject(OffsetStorageReader.class),
-                EasyMock.anyObject(OffsetStorageWriter.class),
-                EasyMock.eq(config),
-                EasyMock.anyObject(ConnectMetrics.class),
-                EasyMock.anyObject(ClassLoader.class),
-                EasyMock.anyObject(Time.class))
-                .andReturn(workerTask);
+        expectNewWorkerTask();
         Map origProps = new HashMap<>();
         origProps.put(TaskConfig.TASK_CLASS_CONFIG, TestSourceTask.class.getName());
 
@@ -507,19 +634,20 @@ public void testAddRemoveTask() throws Exception {
 
         workerTask.initialize(taskConfig);
         EasyMock.expectLastCall();
-        // We should expect this call, but the pluginLoader being swapped in is only mocked.
-        // Serializers for the Producer that the task generates. These are loaded while the PluginClassLoader is active
-        // and then delegated to the system classloader. This is only called once due to caching
-        // EasyMock.expect(pluginLoader.loadClass(ByteArraySerializer.class.getName()))
-        //        .andReturn((Class) ByteArraySerializer.class);
 
-        workerTask.run();
-        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);
 
@@ -527,23 +655,30 @@ 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();
         EasyMock.expect(workerTask.awaitStop(EasyMock.anyLong())).andStubReturn(true);
         EasyMock.expectLastCall();
 
+        workerTask.removeMetrics();
+        EasyMock.expectLastCall();
+
         expectStopStorage();
+        expectClusterId();
 
         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.herder = herder;
         worker.start();
         assertStatistics(worker, 0, 0);
         assertStartupStatistics(worker, 0, 0, 0, 0);
         assertEquals(Collections.emptySet(), worker.taskIds());
-        worker.startTask(TASK_ID, anyConnectorConfigMap(), origProps, taskStatusListener, TargetState.STARTED);
+        worker.startTask(TASK_ID, ClusterConfigState.EMPTY, anyConnectorConfigMap(), origProps, taskStatusListener, TargetState.STARTED);
         assertStatistics(worker, 0, 1);
         assertStartupStatistics(worker, 0, 0, 1, 0);
         assertEquals(new HashSet<>(Arrays.asList(TASK_ID)), worker.taskIds());
@@ -560,9 +695,167 @@ public void testAddRemoveTask() throws Exception {
     }
 
     @Test
-    public void testStartTaskFailure() throws Exception {
+    public void testTaskStatusMetricsStatuses() throws Exception {
+        expectConverters();
+        expectStartStorage();
+        expectFileConfigProvider();
+
+        EasyMock.expect(workerTask.id()).andStubReturn(TASK_ID);
+
+        EasyMock.expect(plugins.currentThreadLoader()).andReturn(delegatingLoader).times(2);
+        expectNewWorkerTask();
+        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();
+
+        workerTask.removeMetrics();
+        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();
+
+        expectClusterId();
+
+        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();
+
+        expectClusterId();
+
+        PowerMock.replayAll();
+
+        worker = new Worker(WORKER_ID,
+            new MockTime(),
+            plugins,
+            config,
+            offsetBackingStore,
+            noneConnectorClientConfigOverridePolicy);
+        worker.herder = herder;
+
+        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");
@@ -587,14 +880,17 @@ public void testStartTaskFailure() throws Exception {
         taskStatusListener.onFailure(EasyMock.eq(TASK_ID), EasyMock.anyObject());
         EasyMock.expectLastCall();
 
+        expectClusterId();
+
         PowerMock.replayAll();
 
-        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore);
+        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore, noneConnectorClientConfigOverridePolicy);
+        worker.herder = herder;
         worker.start();
         assertStatistics(worker, 0, 0);
         assertStartupStatistics(worker, 0, 0, 0, 0);
 
-        assertFalse(worker.startTask(TASK_ID, anyConnectorConfigMap(), origProps, taskStatusListener, TargetState.STARTED));
+        assertFalse(worker.startTask(TASK_ID, ClusterConfigState.EMPTY, anyConnectorConfigMap(), origProps, taskStatusListener, TargetState.STARTED));
         assertStartupStatistics(worker, 0, 0, 1, 1);
 
         assertStatistics(worker, 0, 0);
@@ -606,31 +902,14 @@ public void testStartTaskFailure() throws Exception {
 
     @Test
     public void testCleanupTasksOnStop() throws Exception {
-        expectConverters(true);
+        expectConverters();
         expectStartStorage();
+        expectFileConfigProvider();
 
-        // Create
-        TestSourceTask task = PowerMock.createMock(TestSourceTask.class);
-        WorkerSourceTask workerTask = PowerMock.createMock(WorkerSourceTask.class);
         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),
-                EasyMock.anyObject(TaskStatus.Listener.class),
-                EasyMock.eq(TargetState.STARTED),
-                EasyMock.anyObject(JsonConverter.class),
-                EasyMock.anyObject(JsonConverter.class),
-                EasyMock.eq(TransformationChain.noOp()),
-                EasyMock.anyObject(KafkaProducer.class),
-                EasyMock.anyObject(OffsetStorageReader.class),
-                EasyMock.anyObject(OffsetStorageWriter.class),
-                EasyMock.anyObject(WorkerConfig.class),
-                EasyMock.anyObject(ConnectMetrics.class),
-                EasyMock.eq(pluginLoader),
-                EasyMock.anyObject(Time.class))
-                .andReturn(workerTask);
+        expectNewWorkerTask();
         Map origProps = new HashMap<>();
         origProps.put(TaskConfig.TASK_CLASS_CONFIG, TestSourceTask.class.getName());
 
@@ -643,14 +922,19 @@ public void testCleanupTasksOnStop() throws Exception {
 
         workerTask.initialize(taskConfig);
         EasyMock.expectLastCall();
-        // We should expect this call, but the pluginLoader being swapped in is only mocked.
-        // Serializers for the Producer that the task generates. These are loaded while the PluginClassLoader is active
-        // and then delegated to the system classloader. This is only called once due to caching
-        // EasyMock.expect(pluginLoader.loadClass(ByteArraySerializer.class.getName()))
-        //        .andReturn((Class) ByteArraySerializer.class);
 
-        workerTask.run();
-        EasyMock.expectLastCall();
+        // Expect that the worker will create converters and will not initially find them using the current classloader ...
+        assertNotNull(taskKeyConverter);
+        assertNotNull(taskValueConverter);
+        assertNotNull(taskHeaderConverter);
+        expectTaskKeyConverters(ClassLoaderUsage.CURRENT_CLASSLOADER, null);
+        expectTaskKeyConverters(ClassLoaderUsage.PLUGINS, taskKeyConverter);
+        expectTaskValueConverters(ClassLoaderUsage.CURRENT_CLASSLOADER, null);
+        expectTaskValueConverters(ClassLoaderUsage.PLUGINS, taskValueConverter);
+        expectTaskHeaderConverter(ClassLoaderUsage.CURRENT_CLASSLOADER, null);
+        expectTaskHeaderConverter(ClassLoaderUsage.PLUGINS, taskHeaderConverter);
+
+        EasyMock.expect(executorService.submit(workerTask)).andReturn(null);
 
         EasyMock.expect(plugins.delegatingLoader()).andReturn(delegatingLoader);
         EasyMock.expect(delegatingLoader.connectorLoader(WorkerTestConnector.class.getName()))
@@ -663,7 +947,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();
@@ -672,14 +957,20 @@ public void testCleanupTasksOnStop() throws Exception {
         // Note that in this case we *do not* commit offsets since it's an unclean shutdown
         EasyMock.expectLastCall();
 
+        workerTask.removeMetrics();
+        EasyMock.expectLastCall();
+
         expectStopStorage();
+        expectClusterId();
 
         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.herder = herder;
         worker.start();
         assertStatistics(worker, 0, 0);
-        worker.startTask(TASK_ID, anyConnectorConfigMap(), origProps, taskStatusListener, TargetState.STARTED);
+        worker.startTask(TASK_ID, ClusterConfigState.EMPTY, anyConnectorConfigMap(), origProps, taskStatusListener, TargetState.STARTED);
         assertStatistics(worker, 0, 1);
         worker.stop();
         assertStatistics(worker, 0, 0);
@@ -691,31 +982,12 @@ public void testCleanupTasksOnStop() throws Exception {
     public void testConverterOverrides() throws Exception {
         expectConverters();
         expectStartStorage();
+        expectFileConfigProvider();
 
-        TestSourceTask task = PowerMock.createMock(TestSourceTask.class);
-        WorkerSourceTask workerTask = PowerMock.createMock(WorkerSourceTask.class);
         EasyMock.expect(workerTask.id()).andStubReturn(TASK_ID);
 
-        Capture keyConverter = EasyMock.newCapture();
-        Capture valueConverter = EasyMock.newCapture();
-
         EasyMock.expect(plugins.currentThreadLoader()).andReturn(delegatingLoader).times(2);
-        PowerMock.expectNew(
-                WorkerSourceTask.class, EasyMock.eq(TASK_ID),
-                EasyMock.eq(task),
-                EasyMock.anyObject(TaskStatus.Listener.class),
-                EasyMock.eq(TargetState.STARTED),
-                EasyMock.capture(keyConverter),
-                EasyMock.capture(valueConverter),
-                EasyMock.eq(TransformationChain.noOp()),
-                EasyMock.anyObject(KafkaProducer.class),
-                EasyMock.anyObject(OffsetStorageReader.class),
-                EasyMock.anyObject(OffsetStorageWriter.class),
-                EasyMock.anyObject(WorkerConfig.class),
-                EasyMock.anyObject(ConnectMetrics.class),
-                EasyMock.eq(pluginLoader),
-                EasyMock.anyObject(Time.class))
-                .andReturn(workerTask);
+        expectNewWorkerTask();
         Map origProps = new HashMap<>();
         origProps.put(TaskConfig.TASK_CLASS_CONFIG, TestSourceTask.class.getName());
 
@@ -728,14 +1000,19 @@ public void testConverterOverrides() throws Exception {
 
         workerTask.initialize(taskConfig);
         EasyMock.expectLastCall();
-        // We should expect this call, but the pluginLoader being swapped in is only mocked.
-        // Serializers for the Producer that the task generates. These are loaded while the PluginClassLoader is active
-        // and then delegated to the system classloader. This is only called once due to caching
-        // EasyMock.expect(pluginLoader.loadClass(ByteArraySerializer.class.getName()))
-        //        .andReturn((Class) ByteArraySerializer.class);
 
-        workerTask.run();
-        EasyMock.expectLastCall();
+        // Expect that the worker will create converters and will not initially find them using the current classloader ...
+        assertNotNull(taskKeyConverter);
+        assertNotNull(taskValueConverter);
+        assertNotNull(taskHeaderConverter);
+        expectTaskKeyConverters(ClassLoaderUsage.CURRENT_CLASSLOADER, null);
+        expectTaskKeyConverters(ClassLoaderUsage.PLUGINS, taskKeyConverter);
+        expectTaskValueConverters(ClassLoaderUsage.CURRENT_CLASSLOADER, null);
+        expectTaskValueConverters(ClassLoaderUsage.PLUGINS, taskValueConverter);
+        expectTaskHeaderConverter(ClassLoaderUsage.CURRENT_CLASSLOADER, null);
+        expectTaskHeaderConverter(ClassLoaderUsage.PLUGINS, taskHeaderConverter);
+
+        EasyMock.expect(executorService.submit(workerTask)).andReturn(null);
 
         EasyMock.expect(plugins.delegatingLoader()).andReturn(delegatingLoader);
         EasyMock.expect(delegatingLoader.connectorLoader(WorkerTestConnector.class.getName()))
@@ -748,6 +1025,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();
@@ -755,20 +1034,26 @@ public void testConverterOverrides() throws Exception {
         EasyMock.expect(workerTask.awaitStop(EasyMock.anyLong())).andStubReturn(true);
         EasyMock.expectLastCall();
 
+        workerTask.removeMetrics();
+        EasyMock.expectLastCall();
+
         expectStopStorage();
+        expectClusterId();
 
         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.herder = herder;
         worker.start();
         assertStatistics(worker, 0, 0);
         assertEquals(Collections.emptySet(), worker.taskIds());
         Map connProps = anyConnectorConfigMap();
         connProps.put(ConnectorConfig.KEY_CONVERTER_CLASS_CONFIG, TestConverter.class.getName());
         connProps.put("key.converter.extra.config", "foo");
-        connProps.put(ConnectorConfig.VALUE_CONVERTER_CLASS_CONFIG, TestConverter.class.getName());
+        connProps.put(ConnectorConfig.VALUE_CONVERTER_CLASS_CONFIG, TestConfigurableConverter.class.getName());
         connProps.put("value.converter.extra.config", "bar");
-        worker.startTask(TASK_ID, connProps, origProps, taskStatusListener, TargetState.STARTED);
+        worker.startTask(TASK_ID, ClusterConfigState.EMPTY, connProps, origProps, taskStatusListener, TargetState.STARTED);
         assertStatistics(worker, 0, 1);
         assertEquals(new HashSet<>(Arrays.asList(TASK_ID)), worker.taskIds());
         worker.stopAndAwaitTask(TASK_ID);
@@ -778,14 +1063,272 @@ public void testConverterOverrides() throws Exception {
         worker.stop();
         assertStatistics(worker, 0, 0);
 
-        // Validate extra configs got passed through to overridden converters
-        assertEquals("foo", keyConverter.getValue().configs.get("extra.config"));
-        assertEquals("bar", valueConverter.getValue().configs.get("extra.config"));
+        // We've mocked the Plugin.newConverter method, so we don't currently configure the converters
 
         PowerMock.verifyAll();
     }
 
+    @Test
+    public void testProducerConfigsWithoutOverrides() {
+        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");
+        expectedConfigs.put("metrics.context.connect.kafka.cluster.id", CLUSTER_ID);
+        assertEquals(expectedConfigs,
+                     Worker.producerConfigs(TASK_ID, "connector-producer-" + TASK_ID, config, connectorConfig, null, noneConnectorClientConfigOverridePolicy, CLUSTER_ID));
+    }
+
+    @Test
+    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");
+        expectedConfigs.put("client.id", "producer-test-id");
+        expectedConfigs.put("metrics.context.connect.kafka.cluster.id", CLUSTER_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, CLUSTER_ID));
+    }
+
+    @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");
+        expectedConfigs.put("metrics.context.connect.kafka.cluster.id", CLUSTER_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, CLUSTER_ID));
+    }
+
+    @Test
+    public void testConsumerConfigsWithoutOverrides() {
+        Map expectedConfigs = new HashMap<>(defaultConsumerConfigs);
+        expectedConfigs.put("group.id", "connect-test");
+        expectedConfigs.put("client.id", "connector-consumer-test-1");
+        expectedConfigs.put("metrics.context.connect.kafka.cluster.id", CLUSTER_ID);
+
+        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, CLUSTER_ID));
+    }
+
+    @Test
+    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");
+        expectedConfigs.put("client.id", "consumer-test-id");
+        expectedConfigs.put("metrics.context.connect.kafka.cluster.id", CLUSTER_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, CLUSTER_ID));
+
+    }
+
+    @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");
+        expectedConfigs.put("metrics.context.connect.kafka.cluster.id", CLUSTER_ID);
+
+        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, CLUSTER_ID));
+    }
+
+    @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, CLUSTER_ID);
+    }
+
+    @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");
+        //we added a config on the fly
+        expectedConfigs.put("metrics.context.connect.kafka.cluster.id", CLUSTER_ID);
+
+        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, CLUSTER_ID));
+    }
+
+    @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, CLUSTER_ID);
+
+    }
+
+    @Test
+    public void testWorkerMetrics() throws Exception {
+        expectConverters();
+        expectStartStorage();
+        expectFileConfigProvider();
+
+        // Create
+        EasyMock.expect(plugins.currentThreadLoader()).andReturn(delegatingLoader).times(2);
+        EasyMock.expect(plugins.newConnector(WorkerTestConnector.class.getName()))
+                .andReturn(sourceConnector);
+        EasyMock.expect(sourceConnector.version()).andReturn("1.0");
+
+        Map props = new HashMap<>();
+        props.put(SinkConnectorConfig.TOPICS_CONFIG, "foo,bar");
+        props.put(ConnectorConfig.TASKS_MAX_CONFIG, "1");
+        props.put(ConnectorConfig.NAME_CONFIG, CONNECTOR_ID);
+        props.put(ConnectorConfig.CONNECTOR_CLASS_CONFIG, WorkerTestConnector.class.getName());
+
+        EasyMock.expect(sourceConnector.version()).andReturn("1.0");
+
+        EasyMock.expect(plugins.compareAndSwapLoaders(sourceConnector))
+                .andReturn(delegatingLoader)
+                .times(2);
+        sourceConnector.initialize(anyObject(ConnectorContext.class));
+        EasyMock.expectLastCall();
+        sourceConnector.start(props);
+        EasyMock.expectLastCall();
+
+        EasyMock.expect(Plugins.compareAndSwapLoaders(delegatingLoader))
+                .andReturn(pluginLoader).times(2);
+
+        connectorStatusListener.onStartup(CONNECTOR_ID);
+        EasyMock.expectLastCall();
+
+        // Remove
+        sourceConnector.stop();
+        EasyMock.expectLastCall();
+
+        connectorStatusListener.onShutdown(CONNECTOR_ID);
+        EasyMock.expectLastCall();
+
+        expectStopStorage();
+        expectClusterId();
+
+        PowerMock.replayAll();
+
+        Worker worker = new Worker("worker-1",
+                Time.SYSTEM,
+                plugins,
+                config,
+                offsetBackingStore,
+                noneConnectorClientConfigOverridePolicy
+                );
+        MetricName name = worker.metrics().metrics().metricName("test.avg", "grp1");
+        worker.metrics().metrics().addMetric(name, new Avg());
+        MBeanServer server = ManagementFactory.getPlatformMBeanServer();
+        Set ret = server.queryMBeans(null, null);
+
+        List list = worker.metrics().metrics().reporters();
+        for (MetricsReporter reporter : list) {
+            if (reporter instanceof MockMetricsReporter) {
+                MockMetricsReporter mockMetricsReporter = (MockMetricsReporter) reporter;
+                //verify connect cluster is set in MetricsContext
+                assertEquals(CLUSTER_ID, mockMetricsReporter.getMetricsContext().contextLabels().get(WorkerConfig.CONNECT_KAFKA_CLUSTER_ID));
+            }
+        }
+        //verify metric is created with correct jmx prefix
+        assertNotNull(server.getObjectInstance(new ObjectName("kafka.connect:type=grp1")));
+    }
+
+    private void assertStatusMetrics(long expected, String metricName) {
+        MetricGroup statusMetrics = worker.connectorStatusMetricsGroup().metricGroup(TASK_ID.connector());
+        if (expected == 0L) {
+            assertNull(statusMetrics);
+            return;
+        }
+        assertEquals(expected, MockConnectMetrics.currentMetricValue(worker.metrics(), statusMetrics, metricName));
+    }
+
     private void assertStatistics(Worker worker, int connectors, int tasks) {
+        assertStatusMetrics(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);
@@ -821,10 +1364,12 @@ private void assertStartupStatistics(Worker worker, int connectorStartupAttempts
     }
 
     private void expectStartStorage() {
-        offsetBackingStore.configure(EasyMock.anyObject(WorkerConfig.class));
+        offsetBackingStore.configure(anyObject(WorkerConfig.class));
         EasyMock.expectLastCall();
         offsetBackingStore.start();
         EasyMock.expectLastCall();
+        EasyMock.expect(herder.statusBackingStore())
+                .andReturn(statusBackingStore).anyTimes();
     }
 
     private void expectStopStorage() {
@@ -840,27 +1385,16 @@ private void expectConverters(Boolean expectDefaultConverters) {
         expectConverters(JsonConverter.class, expectDefaultConverters);
     }
 
+    @SuppressWarnings("deprecation")
     private void expectConverters(Class converterClass, Boolean expectDefaultConverters) {
         // As default converters are instantiated when a task starts, they are expected only if the `startTask` method is called
         if (expectDefaultConverters) {
-            // connector default
-            Converter keyConverter = PowerMock.createMock(converterClass);
-            Converter valueConverter = PowerMock.createMock(converterClass);
 
             // Instantiate and configure default
-            EasyMock.expect(plugins.newConverter(JsonConverter.class.getName(), config))
+            EasyMock.expect(plugins.newConverter(config, WorkerConfig.KEY_CONVERTER_CLASS_CONFIG, ClassLoaderUsage.PLUGINS))
                     .andReturn(keyConverter);
-            keyConverter.configure(
-                    EasyMock.>anyObject(),
-                    EasyMock.anyBoolean()
-            );
-            EasyMock.expectLastCall();
-            EasyMock.expect(plugins.newConverter(JsonConverter.class.getName(), config))
+            EasyMock.expect(plugins.newConverter(config, WorkerConfig.VALUE_CONVERTER_CLASS_CONFIG, ClassLoaderUsage.PLUGINS))
                     .andReturn(valueConverter);
-            valueConverter.configure(
-                    EasyMock.>anyObject(),
-                    EasyMock.anyBoolean()
-            );
             EasyMock.expectLastCall();
         }
 
@@ -869,22 +1403,50 @@ private void expectConverters(Class converterClass, Boolean
         Converter internalValueConverter = PowerMock.createMock(converterClass);
 
         // Instantiate and configure internal
-        EasyMock.expect(plugins.newConverter(JsonConverter.class.getName(), config))
-                .andReturn(internalKeyConverter);
-        internalKeyConverter.configure(
-                EasyMock.>anyObject(),
-                EasyMock.anyBoolean()
-        );
-        EasyMock.expectLastCall();
-        EasyMock.expect(plugins.newConverter(JsonConverter.class.getName(), config))
-                .andReturn(internalValueConverter);
-        internalValueConverter.configure(
-                EasyMock.>anyObject(),
-                EasyMock.anyBoolean()
-        );
+        EasyMock.expect(
+                plugins.newConverter(
+                        config,
+                        WorkerConfig.INTERNAL_KEY_CONVERTER_CLASS_CONFIG,
+                        ClassLoaderUsage.PLUGINS
+                )
+        ).andReturn(internalKeyConverter);
+        EasyMock.expect(
+                plugins.newConverter(
+                        config,
+                        WorkerConfig.INTERNAL_VALUE_CONVERTER_CLASS_CONFIG,
+                        ClassLoaderUsage.PLUGINS
+                )
+        ).andReturn(internalValueConverter);
         EasyMock.expectLastCall();
     }
 
+    private void expectTaskKeyConverters(ClassLoaderUsage classLoaderUsage, Converter returning) {
+        EasyMock.expect(
+                plugins.newConverter(
+                        anyObject(AbstractConfig.class),
+                        eq(WorkerConfig.KEY_CONVERTER_CLASS_CONFIG),
+                        eq(classLoaderUsage)))
+                .andReturn(returning);
+    }
+
+    private void expectTaskValueConverters(ClassLoaderUsage classLoaderUsage, Converter returning) {
+        EasyMock.expect(
+                plugins.newConverter(
+                        anyObject(AbstractConfig.class),
+                        eq(WorkerConfig.VALUE_CONVERTER_CLASS_CONFIG),
+                        eq(classLoaderUsage)))
+                .andReturn(returning);
+    }
+
+    private void expectTaskHeaderConverter(ClassLoaderUsage classLoaderUsage, HeaderConverter returning) {
+        EasyMock.expect(
+                plugins.newHeaderConverter(
+                        anyObject(AbstractConfig.class),
+                        eq(WorkerConfig.HEADER_CONVERTER_CLASS_CONFIG),
+                        eq(classLoaderUsage)))
+                .andReturn(returning);
+    }
+
     private Map anyConnectorConfigMap() {
         Map props = new HashMap<>();
         props.put(ConnectorConfig.NAME_CONFIG, CONNECTOR_ID);
@@ -893,8 +1455,37 @@ private Map anyConnectorConfigMap() {
         return props;
     }
 
+    private void expectClusterId() {
+        PowerMock.mockStaticPartial(ConnectUtils.class, "lookupKafkaClusterId");
+        EasyMock.expect(ConnectUtils.lookupKafkaClusterId(EasyMock.anyObject())).andReturn("test-cluster").anyTimes();
+    }
+
+    private void expectNewWorkerTask() throws Exception {
+        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(TopicAdmin.class),
+                EasyMock.>anyObject(),
+                anyObject(OffsetStorageReader.class),
+                anyObject(OffsetStorageWriter.class),
+                EasyMock.eq(config),
+                anyObject(ClusterConfigState.class),
+                anyObject(ConnectMetrics.class),
+                EasyMock.eq(pluginLoader),
+                anyObject(Time.class),
+                anyObject(RetryWithToleranceOperator.class),
+                anyObject(StatusBackingStore.class))
+                .andReturn(workerTask);
+    }
     /* Name here needs to be unique as we are testing the aliasing mechanism */
-    public static class WorkerTestConnector extends Connector {
+    public static class WorkerTestConnector extends SourceConnector {
 
         private static final ConfigDef CONFIG_DEF  = new ConfigDef()
             .define("configName", ConfigDef.Type.STRING, ConfigDef.Importance.HIGH, "Test configName.");
@@ -971,4 +1562,33 @@ public SchemaAndValue toConnectData(String topic, byte[] value) {
             return null;
         }
     }
+
+    public static class TestConfigurableConverter implements Converter, Configurable {
+        public Map configs;
+
+        public ConfigDef config() {
+            return JsonConverterConfig.configDef();
+        }
+
+        @Override
+        public void configure(Map configs) {
+            this.configs = configs;
+            new JsonConverterConfig(configs); // requires the `converter.type` config be set
+        }
+
+        @Override
+        public void configure(Map configs, boolean isKey) {
+            this.configs = configs;
+        }
+
+        @Override
+        public byte[] fromConnectData(String topic, Schema schema, Object value) {
+            return new byte[0];
+        }
+
+        @Override
+        public SchemaAndValue toConnectData(String topic, byte[] value) {
+            return null;
+        }
+    }
 }
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/WorkerWithTopicCreationTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerWithTopicCreationTest.java
new file mode 100644
index 0000000000000..763ef36e2f851
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerWithTopicCreationTest.java
@@ -0,0 +1,1508 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.producer.KafkaProducer;
+import org.apache.kafka.clients.producer.ProducerConfig;
+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.MockFileConfigProvider;
+import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.common.utils.Time;
+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;
+import org.apache.kafka.connect.json.JsonConverter;
+import org.apache.kafka.connect.json.JsonConverterConfig;
+import org.apache.kafka.connect.runtime.ConnectMetrics.MetricGroup;
+import org.apache.kafka.connect.runtime.MockConnectMetrics.MockMetricsReporter;
+import org.apache.kafka.connect.runtime.distributed.ClusterConfigState;
+import org.apache.kafka.connect.runtime.errors.RetryWithToleranceOperator;
+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.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.SourceConnector;
+import org.apache.kafka.connect.source.SourceRecord;
+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.OffsetBackingStore;
+import org.apache.kafka.connect.storage.OffsetStorageReader;
+import org.apache.kafka.connect.storage.OffsetStorageWriter;
+import org.apache.kafka.connect.storage.StatusBackingStore;
+import org.apache.kafka.connect.util.ConnectUtils;
+import org.apache.kafka.connect.util.ConnectorTaskId;
+import org.apache.kafka.connect.util.FutureCallback;
+import org.apache.kafka.connect.util.ThreadedTest;
+import org.apache.kafka.connect.util.TopicAdmin;
+import org.apache.kafka.connect.util.TopicCreationGroup;
+import org.easymock.EasyMock;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.powermock.api.easymock.PowerMock;
+import org.powermock.api.easymock.annotation.Mock;
+import org.powermock.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;
+import org.powermock.modules.junit4.PowerMockRunner;
+
+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.UUID;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.TimeUnit;
+
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.DEFAULT_TOPIC_CREATION_PREFIX;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.PARTITIONS_CONFIG;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.REPLICATION_FACTOR_CONFIG;
+import static org.apache.kafka.connect.runtime.WorkerConfig.TOPIC_CREATION_ENABLE_CONFIG;
+import static org.apache.kafka.connect.runtime.errors.RetryWithToleranceOperatorTest.NOOP_OPERATOR;
+import static org.easymock.EasyMock.anyObject;
+import static org.easymock.EasyMock.eq;
+import static org.easymock.EasyMock.expectLastCall;
+import static org.hamcrest.CoreMatchers.instanceOf;
+import static org.hamcrest.MatcherAssert.assertThat;
+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)
+@PrepareForTest({Worker.class, Plugins.class, ConnectUtils.class})
+@PowerMockIgnore("javax.management.*")
+public class WorkerWithTopicCreationTest 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 static final String CLUSTER_ID = "test-cluster";
+    private final ConnectorClientConfigOverridePolicy noneConnectorClientConfigOverridePolicy = new NoneConnectorClientConfigOverridePolicy();
+    private final ConnectorClientConfigOverridePolicy allConnectorClientConfigOverridePolicy = new AllConnectorClientConfigOverridePolicy();
+
+    private Map workerProps = new HashMap<>();
+    private WorkerConfig config;
+    private Worker worker;
+
+    private Map defaultProducerConfigs = new HashMap<>();
+    private Map defaultConsumerConfigs = new HashMap<>();
+
+    @Mock
+    private Plugins plugins;
+    @Mock
+    private PluginClassLoader pluginLoader;
+    @Mock
+    private DelegatingClassLoader delegatingLoader;
+    @Mock
+    private OffsetBackingStore offsetBackingStore;
+    @MockStrict
+    private TaskStatus.Listener taskStatusListener;
+    @MockStrict
+    private ConnectorStatus.Listener connectorStatusListener;
+
+    @Mock private Herder herder;
+    @Mock private StatusBackingStore statusBackingStore;
+    @Mock private SourceConnector connector;
+    @Mock private CloseableConnectorContext ctx;
+    @Mock private TestSourceTask task;
+    @Mock private WorkerSourceTask workerTask;
+    @Mock private Converter keyConverter;
+    @Mock private Converter valueConverter;
+    @Mock private Converter taskKeyConverter;
+    @Mock private Converter taskValueConverter;
+    @Mock private HeaderConverter taskHeaderConverter;
+    @Mock private ExecutorService executorService;
+    @MockNice private ConnectorConfig connectorConfig;
+    private String mockFileProviderTestId;
+    private Map connectorProps;
+
+    // when this test becomes parameterized, this variable will be a test parameter
+    public boolean enableTopicCreation = true;
+
+    @Before
+    public void setup() {
+        super.setup();
+        workerProps.put("key.converter", "org.apache.kafka.connect.json.JsonConverter");
+        workerProps.put("value.converter", "org.apache.kafka.connect.json.JsonConverter");
+        workerProps.put("offset.storage.file.filename", "/tmp/connect.offsets");
+        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);
+        workerProps.put(TOPIC_CREATION_ENABLE_CONFIG, String.valueOf(enableTopicCreation));
+        config = new StandaloneConfig(workerProps);
+
+        defaultProducerConfigs.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092");
+        defaultProducerConfigs.put(
+            ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.ByteArraySerializer");
+        defaultProducerConfigs.put(
+            ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.ByteArraySerializer");
+        defaultProducerConfigs.put(ProducerConfig.REQUEST_TIMEOUT_MS_CONFIG, Integer.toString(Integer.MAX_VALUE));
+        defaultProducerConfigs.put(ProducerConfig.MAX_BLOCK_MS_CONFIG, Long.toString(Long.MAX_VALUE));
+        defaultProducerConfigs.put(ProducerConfig.ACKS_CONFIG, "all");
+        defaultProducerConfigs.put(ProducerConfig.MAX_IN_FLIGHT_REQUESTS_PER_CONNECTION, "1");
+        defaultProducerConfigs.put(ProducerConfig.DELIVERY_TIMEOUT_MS_CONFIG, Integer.toString(Integer.MAX_VALUE));
+
+        defaultConsumerConfigs.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092");
+        defaultConsumerConfigs.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false");
+        defaultConsumerConfigs.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
+        defaultConsumerConfigs
+            .put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.ByteArrayDeserializer");
+        defaultConsumerConfigs
+            .put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.ByteArrayDeserializer");
+
+        // Some common defaults. They might change on individual tests
+        connectorProps = anyConnectorConfigMap();
+        PowerMock.mockStatic(Plugins.class);
+    }
+
+    @Test
+    public void testStartAndStopConnector() throws Throwable {
+        expectConverters();
+        expectStartStorage();
+
+        final String connectorClass = WorkerTest.WorkerTestConnector.class.getName();
+
+        // Create
+        EasyMock.expect(plugins.currentThreadLoader()).andReturn(delegatingLoader).times(2);
+        EasyMock.expect(plugins.delegatingLoader()).andReturn(delegatingLoader);
+        EasyMock.expect(delegatingLoader.connectorLoader(connectorClass)).andReturn(pluginLoader);
+        EasyMock.expect(plugins.newConnector(connectorClass))
+            .andReturn(connector);
+        EasyMock.expect(connector.version()).andReturn("1.0");
+
+        connectorProps.put(ConnectorConfig.CONNECTOR_CLASS_CONFIG, connectorClass);
+
+        EasyMock.expect(connector.version()).andReturn("1.0");
+
+        expectFileConfigProvider();
+        EasyMock.expect(Plugins.compareAndSwapLoaders(pluginLoader))
+            .andReturn(delegatingLoader)
+            .times(3);
+        connector.initialize(anyObject(ConnectorContext.class));
+        EasyMock.expectLastCall();
+        connector.start(connectorProps);
+        EasyMock.expectLastCall();
+
+        EasyMock.expect(Plugins.compareAndSwapLoaders(delegatingLoader))
+            .andReturn(pluginLoader).times(3);
+
+        connectorStatusListener.onStartup(CONNECTOR_ID);
+        EasyMock.expectLastCall();
+
+        // Remove
+        connector.stop();
+        EasyMock.expectLastCall();
+
+        connectorStatusListener.onShutdown(CONNECTOR_ID);
+        EasyMock.expectLastCall();
+
+        ctx.close();
+        expectLastCall();
+
+        expectStopStorage();
+        expectClusterId();
+
+        PowerMock.replayAll();
+
+        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore, noneConnectorClientConfigOverridePolicy);
+        worker.herder = herder;
+        worker.start();
+
+        assertEquals(Collections.emptySet(), worker.connectorNames());
+
+        FutureCallback onFirstStart = new FutureCallback<>();
+        worker.startConnector(CONNECTOR_ID, connectorProps, ctx, connectorStatusListener, TargetState.STARTED, onFirstStart);
+        // Wait for the connector to actually start
+        assertEquals(TargetState.STARTED, onFirstStart.get(1000, TimeUnit.MILLISECONDS));
+        assertEquals(new HashSet<>(Arrays.asList(CONNECTOR_ID)), worker.connectorNames());
+
+        FutureCallback onSecondStart = new FutureCallback<>();
+        worker.startConnector(CONNECTOR_ID, connectorProps, ctx, connectorStatusListener, TargetState.STARTED, onSecondStart);
+        try {
+            onSecondStart.get(0, TimeUnit.MILLISECONDS);
+            fail("Should have failed while trying to start second connector with same name");
+        } catch (ExecutionException e) {
+            assertThat(e.getCause(), instanceOf(ConnectException.class));
+        }
+
+        assertStatistics(worker, 1, 0);
+        assertStartupStatistics(worker, 1, 0, 0, 0);
+        worker.stopAndAwaitConnector(CONNECTOR_ID);
+        assertStatistics(worker, 0, 0);
+        assertStartupStatistics(worker, 1, 0, 0, 0);
+        assertEquals(Collections.emptySet(), worker.connectorNames());
+        // Nothing should be left, so this should effectively be a nop
+        worker.stop();
+        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() throws Exception {
+        expectConverters();
+        expectStartStorage();
+        expectFileConfigProvider();
+
+        final String nonConnectorClass = "java.util.HashMap";
+        connectorProps.put(ConnectorConfig.CONNECTOR_CLASS_CONFIG, nonConnectorClass); // Bad connector class name
+
+        Exception exception = new ConnectException("Failed to find Connector");
+        EasyMock.expect(plugins.currentThreadLoader()).andReturn(delegatingLoader);
+        EasyMock.expect(plugins.delegatingLoader()).andReturn(delegatingLoader);
+        EasyMock.expect(delegatingLoader.connectorLoader(nonConnectorClass)).andReturn(delegatingLoader);
+        EasyMock.expect(Plugins.compareAndSwapLoaders(delegatingLoader))
+            .andReturn(delegatingLoader).times(2);
+        EasyMock.expect(plugins.newConnector(EasyMock.anyString()))
+            .andThrow(exception);
+
+        connectorStatusListener.onFailure(
+            EasyMock.eq(CONNECTOR_ID),
+            EasyMock.anyObject()
+        );
+        EasyMock.expectLastCall();
+
+        expectClusterId();
+
+        PowerMock.replayAll();
+
+        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore, noneConnectorClientConfigOverridePolicy);
+        worker.herder = herder;
+        worker.start();
+
+        assertStatistics(worker, 0, 0);
+        FutureCallback onStart = new FutureCallback<>();
+        worker.startConnector(CONNECTOR_ID, connectorProps, ctx, connectorStatusListener, TargetState.STARTED, onStart);
+        try {
+            onStart.get(0, TimeUnit.MILLISECONDS);
+            fail("Should have failed to start connector");
+        } catch (ExecutionException e) {
+            assertEquals(exception, e.getCause());
+        }
+
+        assertStartupStatistics(worker, 1, 1, 0, 0);
+        assertEquals(Collections.emptySet(), worker.connectorNames());
+
+        assertStatistics(worker, 0, 0);
+        assertStartupStatistics(worker, 1, 1, 0, 0);
+        worker.stopAndAwaitConnector(CONNECTOR_ID);
+        assertStatistics(worker, 0, 0);
+        assertStartupStatistics(worker, 1, 1, 0, 0);
+
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testAddConnectorByAlias() throws Throwable {
+        expectConverters();
+        expectStartStorage();
+        expectFileConfigProvider();
+
+        final String connectorAlias = "WorkerTestConnector";
+
+        EasyMock.expect(plugins.currentThreadLoader()).andReturn(delegatingLoader).times(2);
+        EasyMock.expect(plugins.delegatingLoader()).andReturn(delegatingLoader);
+        EasyMock.expect(delegatingLoader.connectorLoader(connectorAlias)).andReturn(pluginLoader);
+        EasyMock.expect(plugins.newConnector(connectorAlias)).andReturn(connector);
+        EasyMock.expect(connector.version()).andReturn("1.0");
+
+        connectorProps.put(ConnectorConfig.CONNECTOR_CLASS_CONFIG, connectorAlias);
+
+        EasyMock.expect(connector.version()).andReturn("1.0");
+        EasyMock.expect(Plugins.compareAndSwapLoaders(pluginLoader))
+            .andReturn(delegatingLoader)
+            .times(3);
+        connector.initialize(anyObject(ConnectorContext.class));
+        EasyMock.expectLastCall();
+        connector.start(connectorProps);
+        EasyMock.expectLastCall();
+
+        EasyMock.expect(Plugins.compareAndSwapLoaders(delegatingLoader))
+            .andReturn(pluginLoader)
+            .times(3);
+
+        connectorStatusListener.onStartup(CONNECTOR_ID);
+        EasyMock.expectLastCall();
+
+        // Remove
+        connector.stop();
+        EasyMock.expectLastCall();
+
+        connectorStatusListener.onShutdown(CONNECTOR_ID);
+        EasyMock.expectLastCall();
+
+        ctx.close();
+        expectLastCall();
+
+        expectStopStorage();
+        expectClusterId();
+
+        PowerMock.replayAll();
+
+        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore, noneConnectorClientConfigOverridePolicy);
+        worker.herder = herder;
+        worker.start();
+
+        assertStatistics(worker, 0, 0);
+        assertEquals(Collections.emptySet(), worker.connectorNames());
+        FutureCallback onStart = new FutureCallback<>();
+        worker.startConnector(CONNECTOR_ID, connectorProps, ctx, connectorStatusListener, TargetState.STARTED, onStart);
+        // Wait for the connector to actually start
+        assertEquals(TargetState.STARTED, onStart.get(1000, TimeUnit.MILLISECONDS));
+        assertEquals(new HashSet<>(Arrays.asList(CONNECTOR_ID)), worker.connectorNames());
+        assertStatistics(worker, 1, 0);
+        assertStartupStatistics(worker, 1, 0, 0, 0);
+
+        worker.stopAndAwaitConnector(CONNECTOR_ID);
+        assertStatistics(worker, 0, 0);
+        assertStartupStatistics(worker, 1, 0, 0, 0);
+        assertEquals(Collections.emptySet(), worker.connectorNames());
+        // Nothing should be left, so this should effectively be a nop
+        worker.stop();
+        assertStatistics(worker, 0, 0);
+        assertStartupStatistics(worker, 1, 0, 0, 0);
+
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testAddConnectorByShortAlias() throws Throwable {
+        expectConverters();
+        expectStartStorage();
+        expectFileConfigProvider();
+
+        final String shortConnectorAlias = "WorkerTest";
+
+        EasyMock.expect(plugins.currentThreadLoader()).andReturn(delegatingLoader).times(2);
+        EasyMock.expect(plugins.delegatingLoader()).andReturn(delegatingLoader);
+        EasyMock.expect(delegatingLoader.connectorLoader(shortConnectorAlias)).andReturn(pluginLoader);
+        EasyMock.expect(plugins.newConnector(shortConnectorAlias)).andReturn(connector);
+        EasyMock.expect(connector.version()).andReturn("1.0");
+
+        connectorProps.put(ConnectorConfig.CONNECTOR_CLASS_CONFIG, shortConnectorAlias);
+
+        EasyMock.expect(connector.version()).andReturn("1.0");
+        EasyMock.expect(Plugins.compareAndSwapLoaders(pluginLoader))
+            .andReturn(delegatingLoader)
+            .times(3);
+        connector.initialize(anyObject(ConnectorContext.class));
+        EasyMock.expectLastCall();
+        connector.start(connectorProps);
+        EasyMock.expectLastCall();
+
+        EasyMock.expect(Plugins.compareAndSwapLoaders(delegatingLoader))
+            .andReturn(pluginLoader)
+            .times(3);
+
+        connectorStatusListener.onStartup(CONNECTOR_ID);
+        EasyMock.expectLastCall();
+
+        // Remove
+        connector.stop();
+        EasyMock.expectLastCall();
+
+        connectorStatusListener.onShutdown(CONNECTOR_ID);
+        EasyMock.expectLastCall();
+
+        ctx.close();
+        expectLastCall();
+
+        expectStopStorage();
+        expectClusterId();
+
+        PowerMock.replayAll();
+
+        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore, noneConnectorClientConfigOverridePolicy);
+        worker.herder = herder;
+        worker.start();
+
+        assertStatistics(worker, 0, 0);
+        assertEquals(Collections.emptySet(), worker.connectorNames());
+        FutureCallback onStart = new FutureCallback<>();
+        worker.startConnector(CONNECTOR_ID, connectorProps, ctx, connectorStatusListener, TargetState.STARTED, onStart);
+        // Wait for the connector to actually start
+        assertEquals(TargetState.STARTED, onStart.get(1000, TimeUnit.MILLISECONDS));
+        assertEquals(new HashSet<>(Arrays.asList(CONNECTOR_ID)), worker.connectorNames());
+        assertStatistics(worker, 1, 0);
+
+        worker.stopAndAwaitConnector(CONNECTOR_ID);
+        assertStatistics(worker, 0, 0);
+        assertEquals(Collections.emptySet(), worker.connectorNames());
+        // Nothing should be left, so this should effectively be a nop
+        worker.stop();
+        assertStatistics(worker, 0, 0);
+
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testStopInvalidConnector() {
+        expectConverters();
+        expectStartStorage();
+        expectFileConfigProvider();
+        expectClusterId();
+
+        PowerMock.replayAll();
+
+        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore, noneConnectorClientConfigOverridePolicy);
+        worker.herder = herder;
+        worker.start();
+
+        worker.stopAndAwaitConnector(CONNECTOR_ID);
+
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testReconfigureConnectorTasks() throws Throwable {
+        expectConverters();
+        expectStartStorage();
+        expectFileConfigProvider();
+
+        final String connectorClass = WorkerTest.WorkerTestConnector.class.getName();
+
+        EasyMock.expect(plugins.currentThreadLoader()).andReturn(delegatingLoader).times(3);
+        EasyMock.expect(plugins.delegatingLoader()).andReturn(delegatingLoader).times(1);
+        EasyMock.expect(delegatingLoader.connectorLoader(connectorClass)).andReturn(pluginLoader);
+        EasyMock.expect(plugins.newConnector(connectorClass))
+            .andReturn(connector);
+        EasyMock.expect(connector.version()).andReturn("1.0");
+
+        connectorProps.put(SinkConnectorConfig.TOPICS_CONFIG, "foo,bar");
+        connectorProps.put(ConnectorConfig.CONNECTOR_CLASS_CONFIG, connectorClass);
+
+        EasyMock.expect(connector.version()).andReturn("1.0");
+        EasyMock.expect(Plugins.compareAndSwapLoaders(pluginLoader))
+            .andReturn(delegatingLoader)
+            .times(4);
+        connector.initialize(anyObject(ConnectorContext.class));
+        EasyMock.expectLastCall();
+        connector.start(connectorProps);
+        EasyMock.expectLastCall();
+
+        EasyMock.expect(Plugins.compareAndSwapLoaders(delegatingLoader))
+            .andReturn(pluginLoader)
+            .times(4);
+
+        connectorStatusListener.onStartup(CONNECTOR_ID);
+        EasyMock.expectLastCall();
+
+        // Reconfigure
+        EasyMock.>expect(connector.taskClass()).andReturn(TestSourceTask.class);
+        Map taskProps = new HashMap<>();
+        taskProps.put("foo", "bar");
+        EasyMock.expect(connector.taskConfigs(2)).andReturn(Arrays.asList(taskProps, taskProps));
+
+        // Remove
+        connector.stop();
+        EasyMock.expectLastCall();
+
+        connectorStatusListener.onShutdown(CONNECTOR_ID);
+        EasyMock.expectLastCall();
+
+        ctx.close();
+        expectLastCall();
+
+        expectStopStorage();
+        expectClusterId();
+
+        PowerMock.replayAll();
+
+        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore, noneConnectorClientConfigOverridePolicy);
+        worker.herder = herder;
+        worker.start();
+
+        assertStatistics(worker, 0, 0);
+        assertEquals(Collections.emptySet(), worker.connectorNames());
+        FutureCallback onFirstStart = new FutureCallback<>();
+        worker.startConnector(CONNECTOR_ID, connectorProps, ctx, connectorStatusListener, TargetState.STARTED, onFirstStart);
+        // Wait for the connector to actually start
+        assertEquals(TargetState.STARTED, onFirstStart.get(1000, TimeUnit.MILLISECONDS));
+        assertStatistics(worker, 1, 0);
+        assertEquals(new HashSet<>(Arrays.asList(CONNECTOR_ID)), worker.connectorNames());
+
+        FutureCallback onSecondStart = new FutureCallback<>();
+        worker.startConnector(CONNECTOR_ID, connectorProps, ctx, connectorStatusListener, TargetState.STARTED, onSecondStart);
+        try {
+            onSecondStart.get(0, TimeUnit.MILLISECONDS);
+            fail("Should have failed while trying to start second connector with same name");
+        } catch (ExecutionException e) {
+            assertThat(e.getCause(), instanceOf(ConnectException.class));
+        }
+
+        Map connProps = new HashMap<>(connectorProps);
+        connProps.put(ConnectorConfig.TASKS_MAX_CONFIG, "2");
+        ConnectorConfig connConfig = new SinkConnectorConfig(plugins, connProps);
+        List> taskConfigs = worker.connectorTaskConfigs(CONNECTOR_ID, connConfig);
+        Map expectedTaskProps = new HashMap<>();
+        expectedTaskProps.put("foo", "bar");
+        expectedTaskProps.put(TaskConfig.TASK_CLASS_CONFIG, TestSourceTask.class.getName());
+        expectedTaskProps.put(SinkTask.TOPICS_CONFIG, "foo,bar");
+        assertEquals(2, taskConfigs.size());
+        assertEquals(expectedTaskProps, taskConfigs.get(0));
+        assertEquals(expectedTaskProps, taskConfigs.get(1));
+        assertStatistics(worker, 1, 0);
+        assertStartupStatistics(worker, 1, 0, 0, 0);
+        worker.stopAndAwaitConnector(CONNECTOR_ID);
+        assertStatistics(worker, 0, 0);
+        assertStartupStatistics(worker, 1, 0, 0, 0);
+        assertEquals(Collections.emptySet(), worker.connectorNames());
+        // Nothing should be left, so this should effectively be a nop
+        worker.stop();
+        assertStatistics(worker, 0, 0);
+
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testAddRemoveTask() throws Exception {
+        expectConverters();
+        expectStartStorage();
+        expectFileConfigProvider();
+
+        EasyMock.expect(workerTask.id()).andStubReturn(TASK_ID);
+
+        EasyMock.expect(plugins.currentThreadLoader()).andReturn(delegatingLoader).times(2);
+        expectNewWorkerTask();
+        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(workerTask.loader()).andReturn(pluginLoader);
+
+        EasyMock.expect(Plugins.compareAndSwapLoaders(delegatingLoader)).andReturn(pluginLoader)
+                .times(2);
+        plugins.connectorClass(WorkerTestConnector.class.getName());
+        EasyMock.expectLastCall().andReturn(WorkerTestConnector.class);
+        // Remove
+        workerTask.stop();
+        EasyMock.expectLastCall();
+        EasyMock.expect(workerTask.awaitStop(EasyMock.anyLong())).andStubReturn(true);
+        EasyMock.expectLastCall();
+
+        workerTask.removeMetrics();
+        EasyMock.expectLastCall();
+
+        expectStopStorage();
+        expectClusterId();
+
+        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);
+        assertStatistics(worker, 0, 1);
+        assertStartupStatistics(worker, 0, 0, 1, 0);
+        assertEquals(new HashSet<>(Arrays.asList(TASK_ID)), worker.taskIds());
+        worker.stopAndAwaitTask(TASK_ID);
+        assertStatistics(worker, 0, 0);
+        assertStartupStatistics(worker, 0, 0, 1, 0);
+        assertEquals(Collections.emptySet(), worker.taskIds());
+        // Nothing should be left, so this should effectively be a nop
+        worker.stop();
+        assertStatistics(worker, 0, 0);
+        assertStartupStatistics(worker, 0, 0, 1, 0);
+
+        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);
+        expectNewWorkerTask();
+        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();
+
+        workerTask.removeMetrics();
+        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();
+        expectClusterId();
+
+        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();
+        expectClusterId();
+
+        PowerMock.replayAll();
+
+        worker = new Worker(WORKER_ID,
+            new MockTime(),
+            plugins,
+            config,
+            offsetBackingStore,
+            noneConnectorClientConfigOverridePolicy);
+        worker.herder = herder;
+
+        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");
+
+        EasyMock.expect(plugins.currentThreadLoader()).andReturn(delegatingLoader);
+        EasyMock.expect(plugins.delegatingLoader()).andReturn(delegatingLoader);
+        EasyMock.expect(delegatingLoader.connectorLoader(WorkerTestConnector.class.getName()))
+                .andReturn(pluginLoader);
+
+        // We would normally expect this since the plugin loader would have been swapped in. However, since we mock out
+        // all classloader changes, the call actually goes to the normal default classloader. However, this works out
+        // fine since we just wanted a ClassNotFoundException anyway.
+        // EasyMock.expect(pluginLoader.loadClass(origProps.get(TaskConfig.TASK_CLASS_CONFIG)))
+        //        .andThrow(new ClassNotFoundException());
+
+        EasyMock.expect(Plugins.compareAndSwapLoaders(pluginLoader))
+                .andReturn(delegatingLoader);
+
+        EasyMock.expect(Plugins.compareAndSwapLoaders(delegatingLoader))
+                .andReturn(pluginLoader);
+
+        taskStatusListener.onFailure(EasyMock.eq(TASK_ID), EasyMock.anyObject());
+        EasyMock.expectLastCall();
+        expectClusterId();
+
+        PowerMock.replayAll();
+
+        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore, noneConnectorClientConfigOverridePolicy);
+        worker.herder = herder;
+        worker.start();
+        assertStatistics(worker, 0, 0);
+        assertStartupStatistics(worker, 0, 0, 0, 0);
+
+        assertFalse(worker.startTask(TASK_ID, ClusterConfigState.EMPTY, anyConnectorConfigMap(), origProps, taskStatusListener, TargetState.STARTED));
+        assertStartupStatistics(worker, 0, 0, 1, 1);
+
+        assertStatistics(worker, 0, 0);
+        assertStartupStatistics(worker, 0, 0, 1, 1);
+        assertEquals(Collections.emptySet(), worker.taskIds());
+
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testCleanupTasksOnStop() throws Exception {
+        expectConverters();
+        expectStartStorage();
+        expectFileConfigProvider();
+
+        EasyMock.expect(workerTask.id()).andStubReturn(TASK_ID);
+
+        EasyMock.expect(plugins.currentThreadLoader()).andReturn(delegatingLoader).times(2);
+        expectNewWorkerTask();
+        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 not initially find them using the current classloader ...
+        assertNotNull(taskKeyConverter);
+        assertNotNull(taskValueConverter);
+        assertNotNull(taskHeaderConverter);
+        expectTaskKeyConverters(ClassLoaderUsage.CURRENT_CLASSLOADER, null);
+        expectTaskKeyConverters(ClassLoaderUsage.PLUGINS, taskKeyConverter);
+        expectTaskValueConverters(ClassLoaderUsage.CURRENT_CLASSLOADER, null);
+        expectTaskValueConverters(ClassLoaderUsage.PLUGINS, taskValueConverter);
+        expectTaskHeaderConverter(ClassLoaderUsage.CURRENT_CLASSLOADER, null);
+        expectTaskHeaderConverter(ClassLoaderUsage.PLUGINS, 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(workerTask.loader()).andReturn(pluginLoader);
+
+        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();
+
+        EasyMock.expect(workerTask.awaitStop(EasyMock.anyLong())).andReturn(true);
+        // Note that in this case we *do not* commit offsets since it's an unclean shutdown
+        EasyMock.expectLastCall();
+
+        workerTask.removeMetrics();
+        EasyMock.expectLastCall();
+
+        expectStopStorage();
+        expectClusterId();
+
+        PowerMock.replayAll();
+
+        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore, executorService,
+                            noneConnectorClientConfigOverridePolicy);
+        worker.herder = herder;
+        worker.start();
+        assertStatistics(worker, 0, 0);
+        worker.startTask(TASK_ID, ClusterConfigState.EMPTY, anyConnectorConfigMap(), origProps, taskStatusListener, TargetState.STARTED);
+        assertStatistics(worker, 0, 1);
+        worker.stop();
+        assertStatistics(worker, 0, 0);
+
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testConverterOverrides() throws Exception {
+        expectConverters();
+        expectStartStorage();
+        expectFileConfigProvider();
+
+        EasyMock.expect(workerTask.id()).andStubReturn(TASK_ID);
+
+        EasyMock.expect(plugins.currentThreadLoader()).andReturn(delegatingLoader).times(2);
+        expectNewWorkerTask();
+        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 not initially find them using the current classloader ...
+        assertNotNull(taskKeyConverter);
+        assertNotNull(taskValueConverter);
+        assertNotNull(taskHeaderConverter);
+        expectTaskKeyConverters(ClassLoaderUsage.CURRENT_CLASSLOADER, null);
+        expectTaskKeyConverters(ClassLoaderUsage.PLUGINS, taskKeyConverter);
+        expectTaskValueConverters(ClassLoaderUsage.CURRENT_CLASSLOADER, null);
+        expectTaskValueConverters(ClassLoaderUsage.PLUGINS, taskValueConverter);
+        expectTaskHeaderConverter(ClassLoaderUsage.CURRENT_CLASSLOADER, null);
+        expectTaskHeaderConverter(ClassLoaderUsage.PLUGINS, 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(workerTask.loader()).andReturn(pluginLoader);
+
+        EasyMock.expect(Plugins.compareAndSwapLoaders(delegatingLoader)).andReturn(pluginLoader)
+                .times(2);
+        plugins.connectorClass(WorkerTestConnector.class.getName());
+        EasyMock.expectLastCall().andReturn(WorkerTestConnector.class);
+
+        // Remove
+        workerTask.stop();
+        EasyMock.expectLastCall();
+        EasyMock.expect(workerTask.awaitStop(EasyMock.anyLong())).andStubReturn(true);
+        EasyMock.expectLastCall();
+
+        workerTask.removeMetrics();
+        EasyMock.expectLastCall();
+
+        expectStopStorage();
+        expectClusterId();
+
+        PowerMock.replayAll();
+
+        worker = new Worker(WORKER_ID, new MockTime(), plugins, config, offsetBackingStore, executorService,
+                            noneConnectorClientConfigOverridePolicy);
+        worker.herder = herder;
+        worker.start();
+        assertStatistics(worker, 0, 0);
+        assertEquals(Collections.emptySet(), worker.taskIds());
+        Map connProps = anyConnectorConfigMap();
+        connProps.put(ConnectorConfig.KEY_CONVERTER_CLASS_CONFIG, TestConverter.class.getName());
+        connProps.put("key.converter.extra.config", "foo");
+        connProps.put(ConnectorConfig.VALUE_CONVERTER_CLASS_CONFIG, TestConfigurableConverter.class.getName());
+        connProps.put("value.converter.extra.config", "bar");
+        worker.startTask(TASK_ID, ClusterConfigState.EMPTY, connProps, origProps, taskStatusListener, TargetState.STARTED);
+        assertStatistics(worker, 0, 1);
+        assertEquals(new HashSet<>(Arrays.asList(TASK_ID)), worker.taskIds());
+        worker.stopAndAwaitTask(TASK_ID);
+        assertStatistics(worker, 0, 0);
+        assertEquals(Collections.emptySet(), worker.taskIds());
+        // Nothing should be left, so this should effectively be a nop
+        worker.stop();
+        assertStatistics(worker, 0, 0);
+
+        // We've mocked the Plugin.newConverter method, so we don't currently configure the converters
+
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testProducerConfigsWithoutOverrides() {
+        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");
+        expectedConfigs.put("metrics.context.connect.kafka.cluster.id", "test-cluster");
+        assertEquals(expectedConfigs,
+                     Worker.producerConfigs(TASK_ID, "connector-producer-" + TASK_ID, config, connectorConfig, null, noneConnectorClientConfigOverridePolicy, CLUSTER_ID));
+    }
+
+    @Test
+    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");
+        expectedConfigs.put("client.id", "producer-test-id");
+        expectedConfigs.put("metrics.context.connect.kafka.cluster.id", "test-cluster");
+        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, CLUSTER_ID));
+    }
+
+    @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");
+        expectedConfigs.put("metrics.context.connect.kafka.cluster.id", "test-cluster");
+        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, CLUSTER_ID));
+    }
+
+    @Test
+    public void testConsumerConfigsWithoutOverrides() {
+        Map expectedConfigs = new HashMap<>(defaultConsumerConfigs);
+        expectedConfigs.put("group.id", "connect-test");
+        expectedConfigs.put("client.id", "connector-consumer-test-1");
+        expectedConfigs.put("metrics.context.connect.kafka.cluster.id", "test-cluster");
+        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, CLUSTER_ID));
+    }
+
+    @Test
+    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");
+        expectedConfigs.put("client.id", "consumer-test-id");
+        expectedConfigs.put("metrics.context.connect.kafka.cluster.id", "test-cluster");
+        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, CLUSTER_ID));
+
+    }
+
+    @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");
+        expectedConfigs.put("metrics.context.connect.kafka.cluster.id", "test-cluster");
+        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, CLUSTER_ID));
+    }
+
+    @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, CLUSTER_ID);
+    }
+
+    @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");
+        expectedConfigs.put("metrics.context.connect.kafka.cluster.id", "test-cluster");
+
+        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, CLUSTER_ID));
+    }
+
+    @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, CLUSTER_ID);
+    }
+
+    private void assertStatusMetrics(long expected, String metricName) {
+        MetricGroup statusMetrics = worker.connectorStatusMetricsGroup().metricGroup(TASK_ID.connector());
+        if (expected == 0L) {
+            assertNull(statusMetrics);
+            return;
+        }
+        assertEquals(expected, MockConnectMetrics.currentMetricValue(worker.metrics(), statusMetrics, metricName));
+    }
+
+    private void assertStatistics(Worker worker, int connectors, int tasks) {
+        assertStatusMetrics(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);
+        assertEquals(tasks, MockConnectMetrics.currentMetricValueAsDouble(worker.metrics(), workerMetrics, "task-count"), 0.0001d);
+    }
+
+    private void assertStartupStatistics(Worker worker, int connectorStartupAttempts, int connectorStartupFailures, int taskStartupAttempts, int taskStartupFailures) {
+        double connectStartupSuccesses = connectorStartupAttempts - connectorStartupFailures;
+        double taskStartupSuccesses = taskStartupAttempts - taskStartupFailures;
+        double connectStartupSuccessPct = 0.0d;
+        double connectStartupFailurePct = 0.0d;
+        double taskStartupSuccessPct = 0.0d;
+        double taskStartupFailurePct = 0.0d;
+        if (connectorStartupAttempts != 0) {
+            connectStartupSuccessPct = connectStartupSuccesses / connectorStartupAttempts;
+            connectStartupFailurePct = (double) connectorStartupFailures / connectorStartupAttempts;
+        }
+        if (taskStartupAttempts != 0) {
+            taskStartupSuccessPct = taskStartupSuccesses / taskStartupAttempts;
+            taskStartupFailurePct = (double) taskStartupFailures / taskStartupAttempts;
+        }
+        MetricGroup workerMetrics = worker.workerMetricsGroup().metricGroup();
+        assertEquals(connectorStartupAttempts, MockConnectMetrics.currentMetricValueAsDouble(worker.metrics(), workerMetrics, "connector-startup-attempts-total"), 0.0001d);
+        assertEquals(connectStartupSuccesses, MockConnectMetrics.currentMetricValueAsDouble(worker.metrics(), workerMetrics, "connector-startup-success-total"), 0.0001d);
+        assertEquals(connectorStartupFailures, MockConnectMetrics.currentMetricValueAsDouble(worker.metrics(), workerMetrics, "connector-startup-failure-total"), 0.0001d);
+        assertEquals(connectStartupSuccessPct, MockConnectMetrics.currentMetricValueAsDouble(worker.metrics(), workerMetrics, "connector-startup-success-percentage"), 0.0001d);
+        assertEquals(connectStartupFailurePct, MockConnectMetrics.currentMetricValueAsDouble(worker.metrics(), workerMetrics, "connector-startup-failure-percentage"), 0.0001d);
+        assertEquals(taskStartupAttempts, MockConnectMetrics.currentMetricValueAsDouble(worker.metrics(), workerMetrics, "task-startup-attempts-total"), 0.0001d);
+        assertEquals(taskStartupSuccesses, MockConnectMetrics.currentMetricValueAsDouble(worker.metrics(), workerMetrics, "task-startup-success-total"), 0.0001d);
+        assertEquals(taskStartupFailures, MockConnectMetrics.currentMetricValueAsDouble(worker.metrics(), workerMetrics, "task-startup-failure-total"), 0.0001d);
+        assertEquals(taskStartupSuccessPct, MockConnectMetrics.currentMetricValueAsDouble(worker.metrics(), workerMetrics, "task-startup-success-percentage"), 0.0001d);
+        assertEquals(taskStartupFailurePct, MockConnectMetrics.currentMetricValueAsDouble(worker.metrics(), workerMetrics, "task-startup-failure-percentage"), 0.0001d);
+    }
+
+    private void expectStartStorage() {
+        offsetBackingStore.configure(anyObject(WorkerConfig.class));
+        EasyMock.expectLastCall();
+        offsetBackingStore.start();
+        EasyMock.expectLastCall();
+        EasyMock.expect(herder.statusBackingStore())
+                .andReturn(statusBackingStore).anyTimes();
+    }
+
+    private void expectStopStorage() {
+        offsetBackingStore.stop();
+        EasyMock.expectLastCall();
+    }
+
+    private void expectConverters() {
+        expectConverters(JsonConverter.class, false);
+    }
+
+    private void expectConverters(Boolean expectDefaultConverters) {
+        expectConverters(JsonConverter.class, expectDefaultConverters);
+    }
+
+    @SuppressWarnings("deprecation")
+    private void expectConverters(Class converterClass, Boolean expectDefaultConverters) {
+        // As default converters are instantiated when a task starts, they are expected only if the `startTask` method is called
+        if (expectDefaultConverters) {
+
+            // Instantiate and configure default
+            EasyMock.expect(plugins.newConverter(config, WorkerConfig.KEY_CONVERTER_CLASS_CONFIG, ClassLoaderUsage.PLUGINS))
+                    .andReturn(keyConverter);
+            EasyMock.expect(plugins.newConverter(config, WorkerConfig.VALUE_CONVERTER_CLASS_CONFIG, ClassLoaderUsage.PLUGINS))
+                    .andReturn(valueConverter);
+            EasyMock.expectLastCall();
+        }
+
+        //internal
+        Converter internalKeyConverter = PowerMock.createMock(converterClass);
+        Converter internalValueConverter = PowerMock.createMock(converterClass);
+
+        // Instantiate and configure internal
+        EasyMock.expect(
+                plugins.newConverter(
+                        config,
+                        WorkerConfig.INTERNAL_KEY_CONVERTER_CLASS_CONFIG,
+                        ClassLoaderUsage.PLUGINS
+                )
+        ).andReturn(internalKeyConverter);
+        EasyMock.expect(
+                plugins.newConverter(
+                        config,
+                        WorkerConfig.INTERNAL_VALUE_CONVERTER_CLASS_CONFIG,
+                        ClassLoaderUsage.PLUGINS
+                )
+        ).andReturn(internalValueConverter);
+        EasyMock.expectLastCall();
+    }
+
+    private void expectTaskKeyConverters(ClassLoaderUsage classLoaderUsage, Converter returning) {
+        EasyMock.expect(
+                plugins.newConverter(
+                        anyObject(AbstractConfig.class),
+                        eq(WorkerConfig.KEY_CONVERTER_CLASS_CONFIG),
+                        eq(classLoaderUsage)))
+                .andReturn(returning);
+    }
+
+    private void expectTaskValueConverters(ClassLoaderUsage classLoaderUsage, Converter returning) {
+        EasyMock.expect(
+                plugins.newConverter(
+                        anyObject(AbstractConfig.class),
+                        eq(WorkerConfig.VALUE_CONVERTER_CLASS_CONFIG),
+                        eq(classLoaderUsage)))
+                .andReturn(returning);
+    }
+
+    private void expectTaskHeaderConverter(ClassLoaderUsage classLoaderUsage, HeaderConverter returning) {
+        EasyMock.expect(
+                plugins.newHeaderConverter(
+                        anyObject(AbstractConfig.class),
+                        eq(WorkerConfig.HEADER_CONVERTER_CLASS_CONFIG),
+                        eq(classLoaderUsage)))
+                .andReturn(returning);
+    }
+
+    private Map anyConnectorConfigMap() {
+        Map props = new HashMap<>();
+        props.put(ConnectorConfig.NAME_CONFIG, CONNECTOR_ID);
+        props.put(ConnectorConfig.CONNECTOR_CLASS_CONFIG, WorkerTestConnector.class.getName());
+        props.put(ConnectorConfig.TASKS_MAX_CONFIG, "1");
+        props.put(DEFAULT_TOPIC_CREATION_PREFIX + REPLICATION_FACTOR_CONFIG, String.valueOf(1));
+        props.put(DEFAULT_TOPIC_CREATION_PREFIX + PARTITIONS_CONFIG, String.valueOf(1));
+        return props;
+    }
+
+    private void expectNewWorkerTask() throws Exception {
+        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(TopicAdmin.class),
+                EasyMock.>anyObject(),
+                anyObject(OffsetStorageReader.class),
+                anyObject(OffsetStorageWriter.class),
+                EasyMock.eq(config),
+                anyObject(ClusterConfigState.class),
+                anyObject(ConnectMetrics.class),
+                EasyMock.eq(pluginLoader),
+                anyObject(Time.class),
+                anyObject(RetryWithToleranceOperator.class),
+                anyObject(StatusBackingStore.class))
+                .andReturn(workerTask);
+    }
+
+    private void expectClusterId() {
+        PowerMock.mockStaticPartial(ConnectUtils.class, "lookupKafkaClusterId");
+        EasyMock.expect(ConnectUtils.lookupKafkaClusterId(EasyMock.anyObject())).andReturn("test-cluster").anyTimes();
+    }
+
+    /* Name here needs to be unique as we are testing the aliasing mechanism */
+    public static class WorkerTestConnector extends SourceConnector {
+
+        private static final ConfigDef CONFIG_DEF  = new ConfigDef()
+            .define("configName", ConfigDef.Type.STRING, ConfigDef.Importance.HIGH, "Test configName.");
+
+        @Override
+        public String version() {
+            return "1.0";
+        }
+
+        @Override
+        public void start(Map props) {
+
+        }
+
+        @Override
+        public Class taskClass() {
+            return null;
+        }
+
+        @Override
+        public List> taskConfigs(int maxTasks) {
+            return null;
+        }
+
+        @Override
+        public void stop() {
+
+        }
+
+        @Override
+        public ConfigDef config() {
+            return CONFIG_DEF;
+        }
+    }
+
+    private static class TestSourceTask extends SourceTask {
+        public TestSourceTask() {
+        }
+
+        @Override
+        public String version() {
+            return "1.0";
+        }
+
+        @Override
+        public void start(Map props) {
+        }
+
+        @Override
+        public List poll() throws InterruptedException {
+            return null;
+        }
+
+        @Override
+        public void stop() {
+        }
+    }
+
+    public static class TestConverter implements Converter {
+        public Map configs;
+
+        @Override
+        public void configure(Map configs, boolean isKey) {
+            this.configs = configs;
+        }
+
+        @Override
+        public byte[] fromConnectData(String topic, Schema schema, Object value) {
+            return new byte[0];
+        }
+
+        @Override
+        public SchemaAndValue toConnectData(String topic, byte[] value) {
+            return null;
+        }
+    }
+
+    public static class TestConfigurableConverter implements Converter, Configurable {
+        public Map configs;
+
+        public ConfigDef config() {
+            return JsonConverterConfig.configDef();
+        }
+
+        @Override
+        public void configure(Map configs) {
+            this.configs = configs;
+            new JsonConverterConfig(configs); // requires the `converter.type` config be set
+        }
+
+        @Override
+        public void configure(Map configs, boolean isKey) {
+            this.configs = configs;
+        }
+
+        @Override
+        public byte[] fromConnectData(String topic, Schema schema, Object value) {
+            return new byte[0];
+        }
+
+        @Override
+        public SchemaAndValue toConnectData(String topic, byte[] value) {
+            return null;
+        }
+    }
+}
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..1e347cbd12cb9
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/distributed/DistributedConfigTest.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.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.assertNotEquals;
+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));
+        }
+    }
+
+    @Test
+    public void shouldAllowNegativeOneAndPositiveForPartitions() {
+        Map settings = configs();
+        settings.put(DistributedConfig.OFFSET_STORAGE_PARTITIONS_CONFIG, "-1");
+        settings.put(DistributedConfig.STATUS_STORAGE_PARTITIONS_CONFIG, "-1");
+        new DistributedConfig(configs());
+        settings.remove(DistributedConfig.OFFSET_STORAGE_PARTITIONS_CONFIG);
+        settings.remove(DistributedConfig.STATUS_STORAGE_PARTITIONS_CONFIG);
+
+        for (int i = 1; i != 100; ++i) {
+            settings.put(DistributedConfig.OFFSET_STORAGE_PARTITIONS_CONFIG, Integer.toString(i));
+            new DistributedConfig(settings);
+            settings.remove(DistributedConfig.OFFSET_STORAGE_PARTITIONS_CONFIG);
+
+            settings.put(DistributedConfig.STATUS_STORAGE_PARTITIONS_CONFIG, Integer.toString(i));
+            new DistributedConfig(settings);
+        }
+    }
+
+    @Test
+    public void shouldNotAllowZeroPartitions() {
+        Map settings = configs();
+        settings.put(DistributedConfig.OFFSET_STORAGE_PARTITIONS_CONFIG, "0");
+        assertThrows(ConfigException.class, () -> new DistributedConfig(settings));
+        settings.remove(DistributedConfig.OFFSET_STORAGE_PARTITIONS_CONFIG);
+
+        settings.put(DistributedConfig.STATUS_STORAGE_PARTITIONS_CONFIG, "0");
+        assertThrows(ConfigException.class, () -> new DistributedConfig(settings));
+    }
+
+    @Test
+    public void shouldNotAllowNegativePartitionsLessThanNegativeOne() {
+        Map settings = configs();
+        for (int i = -2; i > -100; --i) {
+            settings.put(DistributedConfig.OFFSET_STORAGE_PARTITIONS_CONFIG, Integer.toString(i));
+            assertThrows(ConfigException.class, () -> new DistributedConfig(settings));
+            settings.remove(DistributedConfig.OFFSET_STORAGE_PARTITIONS_CONFIG);
+
+            settings.put(DistributedConfig.STATUS_STORAGE_PARTITIONS_CONFIG, Integer.toString(i));
+            assertThrows(ConfigException.class, () -> new DistributedConfig(settings));
+        }
+    }
+
+    @Test
+    public void shouldAllowNegativeOneAndPositiveForReplicationFactor() {
+        Map settings = configs();
+        settings.put(DistributedConfig.CONFIG_STORAGE_REPLICATION_FACTOR_CONFIG, "-1");
+        settings.put(DistributedConfig.OFFSET_STORAGE_REPLICATION_FACTOR_CONFIG, "-1");
+        settings.put(DistributedConfig.STATUS_STORAGE_REPLICATION_FACTOR_CONFIG, "-1");
+        new DistributedConfig(configs());
+        settings.remove(DistributedConfig.CONFIG_STORAGE_REPLICATION_FACTOR_CONFIG);
+        settings.remove(DistributedConfig.OFFSET_STORAGE_PARTITIONS_CONFIG);
+        settings.remove(DistributedConfig.STATUS_STORAGE_PARTITIONS_CONFIG);
+
+        for (int i = 1; i != 100; ++i) {
+            settings.put(DistributedConfig.CONFIG_STORAGE_REPLICATION_FACTOR_CONFIG, Integer.toString(i));
+            new DistributedConfig(settings);
+            settings.remove(DistributedConfig.CONFIG_STORAGE_REPLICATION_FACTOR_CONFIG);
+
+            settings.put(DistributedConfig.OFFSET_STORAGE_PARTITIONS_CONFIG, Integer.toString(i));
+            new DistributedConfig(settings);
+            settings.remove(DistributedConfig.OFFSET_STORAGE_PARTITIONS_CONFIG);
+
+            settings.put(DistributedConfig.STATUS_STORAGE_PARTITIONS_CONFIG, Integer.toString(i));
+            new DistributedConfig(settings);
+        }
+    }
+
+    @Test
+    public void shouldNotAllowZeroReplicationFactor() {
+        Map settings = configs();
+        settings.put(DistributedConfig.CONFIG_STORAGE_REPLICATION_FACTOR_CONFIG, "0");
+        assertThrows(ConfigException.class, () -> new DistributedConfig(settings));
+        settings.remove(DistributedConfig.CONFIG_STORAGE_REPLICATION_FACTOR_CONFIG);
+
+        settings.put(DistributedConfig.OFFSET_STORAGE_REPLICATION_FACTOR_CONFIG, "0");
+        assertThrows(ConfigException.class, () -> new DistributedConfig(settings));
+        settings.remove(DistributedConfig.OFFSET_STORAGE_REPLICATION_FACTOR_CONFIG);
+
+        settings.put(DistributedConfig.STATUS_STORAGE_REPLICATION_FACTOR_CONFIG, "0");
+        assertThrows(ConfigException.class, () -> new DistributedConfig(settings));
+    }
+
+    @Test
+    public void shouldNotAllowNegativeReplicationFactorLessThanNegativeOne() {
+        Map settings = configs();
+        for (int i = -2; i > -100; --i) {
+            settings.put(DistributedConfig.CONFIG_STORAGE_REPLICATION_FACTOR_CONFIG, Integer.toString(i));
+            assertThrows(ConfigException.class, () -> new DistributedConfig(settings));
+            settings.remove(DistributedConfig.CONFIG_STORAGE_REPLICATION_FACTOR_CONFIG);
+
+            settings.put(DistributedConfig.OFFSET_STORAGE_REPLICATION_FACTOR_CONFIG, Integer.toString(i));
+            assertThrows(ConfigException.class, () -> new DistributedConfig(settings));
+            settings.remove(DistributedConfig.OFFSET_STORAGE_REPLICATION_FACTOR_CONFIG);
+
+            settings.put(DistributedConfig.STATUS_STORAGE_REPLICATION_FACTOR_CONFIG, Integer.toString(i));
+            assertThrows(ConfigException.class, () -> new DistributedConfig(settings));
+        }
+    }
+
+    @Test
+    public void shouldAllowSettingConfigTopicSettings() {
+        Map topicSettings = new HashMap<>();
+        topicSettings.put("foo", "foo value");
+        topicSettings.put("bar", "bar value");
+        topicSettings.put("baz.bim", "100");
+        Map settings = configs();
+        topicSettings.entrySet().forEach(e -> {
+            settings.put(DistributedConfig.CONFIG_STORAGE_PREFIX + e.getKey(), e.getValue());
+        });
+        DistributedConfig config = new DistributedConfig(settings);
+        assertEquals(topicSettings, config.configStorageTopicSettings());
+    }
+
+    @Test
+    public void shouldAllowSettingOffsetTopicSettings() {
+        Map topicSettings = new HashMap<>();
+        topicSettings.put("foo", "foo value");
+        topicSettings.put("bar", "bar value");
+        topicSettings.put("baz.bim", "100");
+        Map settings = configs();
+        topicSettings.entrySet().forEach(e -> {
+            settings.put(DistributedConfig.OFFSET_STORAGE_PREFIX + e.getKey(), e.getValue());
+        });
+        DistributedConfig config = new DistributedConfig(settings);
+        assertEquals(topicSettings, config.offsetStorageTopicSettings());
+    }
+
+    @Test
+    public void shouldAllowSettingStatusTopicSettings() {
+        Map topicSettings = new HashMap<>();
+        topicSettings.put("foo", "foo value");
+        topicSettings.put("bar", "bar value");
+        topicSettings.put("baz.bim", "100");
+        Map settings = configs();
+        topicSettings.entrySet().forEach(e -> {
+            settings.put(DistributedConfig.STATUS_STORAGE_PREFIX + e.getKey(), e.getValue());
+        });
+        DistributedConfig config = new DistributedConfig(settings);
+        assertEquals(topicSettings, config.statusStorageTopicSettings());
+    }
+
+    @Test
+    public void shouldRemoveCompactionFromConfigTopicSettings() {
+        Map expectedTopicSettings = new HashMap<>();
+        expectedTopicSettings.put("foo", "foo value");
+        expectedTopicSettings.put("bar", "bar value");
+        expectedTopicSettings.put("baz.bim", "100");
+        Map topicSettings = new HashMap<>(expectedTopicSettings);
+        topicSettings.put("cleanup.policy", "something-else");
+        topicSettings.put("partitions", "3");
+
+        Map settings = configs();
+        topicSettings.forEach((k, v) -> {
+            settings.put(DistributedConfig.CONFIG_STORAGE_PREFIX + k, v);
+        });
+        DistributedConfig config = new DistributedConfig(settings);
+        Map actual = config.configStorageTopicSettings();
+        assertEquals(expectedTopicSettings, actual);
+        assertNotEquals(topicSettings, actual);
+    }
+
+    @Test
+    public void shouldRemoveCompactionFromOffsetTopicSettings() {
+        Map expectedTopicSettings = new HashMap<>();
+        expectedTopicSettings.put("foo", "foo value");
+        expectedTopicSettings.put("bar", "bar value");
+        expectedTopicSettings.put("baz.bim", "100");
+        Map topicSettings = new HashMap<>(expectedTopicSettings);
+        topicSettings.put("cleanup.policy", "something-else");
+
+        Map settings = configs();
+        topicSettings.forEach((k, v) -> {
+            settings.put(DistributedConfig.OFFSET_STORAGE_PREFIX + k, v);
+        });
+        DistributedConfig config = new DistributedConfig(settings);
+        Map actual = config.offsetStorageTopicSettings();
+        assertEquals(expectedTopicSettings, actual);
+        assertNotEquals(topicSettings, actual);
+    }
+
+    @Test
+    public void shouldRemoveCompactionFromStatusTopicSettings() {
+        Map expectedTopicSettings = new HashMap<>();
+        expectedTopicSettings.put("foo", "foo value");
+        expectedTopicSettings.put("bar", "bar value");
+        expectedTopicSettings.put("baz.bim", "100");
+        Map topicSettings = new HashMap<>(expectedTopicSettings);
+        topicSettings.put("cleanup.policy", "something-else");
+
+        Map settings = configs();
+        topicSettings.forEach((k, v) -> {
+            settings.put(DistributedConfig.STATUS_STORAGE_PREFIX + k, v);
+        });
+        DistributedConfig config = new DistributedConfig(settings);
+        Map actual = config.statusStorageTopicSettings();
+        assertEquals(expectedTopicSettings, actual);
+        assertNotEquals(topicSettings, actual);
+    }
+}
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 d41ccbebae00d..8cf485ac9fa9f 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
@@ -17,14 +17,14 @@
 package org.apache.kafka.connect.runtime.distributed;
 
 import org.apache.kafka.clients.CommonClientConfigs;
-import org.apache.kafka.common.config.Config;
-import org.apache.kafka.common.config.ConfigDef;
 import org.apache.kafka.common.config.ConfigValue;
 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.CloseableConnectorContext;
 import org.apache.kafka.connect.runtime.ConnectMetrics.MetricGroup;
 import org.apache.kafka.connect.runtime.ConnectorConfig;
 import org.apache.kafka.connect.runtime.Herder;
@@ -32,16 +32,21 @@
 import org.apache.kafka.connect.runtime.SinkConnectorConfig;
 import org.apache.kafka.connect.runtime.TargetState;
 import org.apache.kafka.connect.runtime.TaskConfig;
+import org.apache.kafka.connect.runtime.TopicStatus;
 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.distributed.DistributedHerder.HerderMetrics;
 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.ConfigInfos;
 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;
@@ -62,26 +67,38 @@
 import org.powermock.core.classloader.annotations.PowerMockIgnore;
 import org.powermock.core.classloader.annotations.PrepareForTest;
 import org.powermock.modules.junit4.PowerMockRunner;
+import org.powermock.reflect.Whitebox;
 
 import java.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.concurrent.ExecutionException;
+import java.util.concurrent.ThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
 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.assertNotNull;
 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 {
@@ -105,23 +122,32 @@ public class DistributedHerderTest {
     private static final ConnectorTaskId TASK2 = new ConnectorTaskId(CONN1, 2);
     private static final Integer MAX_TASKS = 3;
     private static final Map CONN1_CONFIG = new HashMap<>();
+    private static final String FOO_TOPIC = "foo";
+    private static final String BAR_TOPIC = "bar";
+    private static final String BAZ_TOPIC = "baz";
     static {
         CONN1_CONFIG.put(ConnectorConfig.NAME_CONFIG, CONN1);
         CONN1_CONFIG.put(ConnectorConfig.TASKS_MAX_CONFIG, MAX_TASKS.toString());
-        CONN1_CONFIG.put(SinkConnectorConfig.TOPICS_CONFIG, "foo,bar");
+        CONN1_CONFIG.put(SinkConnectorConfig.TOPICS_CONFIG, String.join(",", FOO_TOPIC, BAR_TOPIC));
         CONN1_CONFIG.put(ConnectorConfig.CONNECTOR_CLASS_CONFIG, BogusSourceConnector.class.getName());
     }
     private static final Map CONN1_CONFIG_UPDATED = new HashMap<>(CONN1_CONFIG);
     static {
-        CONN1_CONFIG_UPDATED.put(SinkConnectorConfig.TOPICS_CONFIG, "foo,bar,baz");
+        CONN1_CONFIG_UPDATED.put(SinkConnectorConfig.TOPICS_CONFIG, String.join(",", FOO_TOPIC, BAR_TOPIC, BAZ_TOPIC));
     }
+    private static final ConfigInfos CONN1_CONFIG_INFOS =
+        new ConfigInfos(CONN1, 0, Collections.emptyList(), Collections.emptyList());
     private static final Map CONN2_CONFIG = new HashMap<>();
     static {
         CONN2_CONFIG.put(ConnectorConfig.NAME_CONFIG, CONN2);
         CONN2_CONFIG.put(ConnectorConfig.TASKS_MAX_CONFIG, MAX_TASKS.toString());
-        CONN2_CONFIG.put(SinkConnectorConfig.TOPICS_CONFIG, "foo,bar");
+        CONN2_CONFIG.put(SinkConnectorConfig.TOPICS_CONFIG, String.join(",", FOO_TOPIC, BAR_TOPIC));
         CONN2_CONFIG.put(ConnectorConfig.CONNECTOR_CLASS_CONFIG, BogusSourceConnector.class.getName());
     }
+    private static final ConfigInfos CONN2_CONFIG_INFOS =
+        new ConfigInfos(CONN2, 0, Collections.emptyList(), Collections.emptyList());
+    private static final ConfigInfos CONN2_INVALID_CONFIG_INFOS =
+        new ConfigInfos(CONN2, 1, Collections.emptyList(), Collections.emptyList());
     private static final Map TASK_CONFIG = new HashMap<>();
     static {
         TASK_CONFIG.put(TaskConfig.TASK_CLASS_CONFIG, BogusSourceTask.class.getName());
@@ -138,17 +164,20 @@ 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());
 
     private static final String WORKER_ID = "localhost:8083";
+    private static final String KAFKA_CLUSTER_ID = "I4ZmrWqfT2e-upky_4fdPA";
+    private static final Runnable EMPTY_RUNNABLE = () -> {
+    };
 
     @Mock private ConfigBackingStore configBackingStore;
     @Mock private StatusBackingStore statusBackingStore;
@@ -157,19 +186,21 @@ public class DistributedHerderTest {
     private DistributedHerder herder;
     private MockConnectMetrics metrics;
     @Mock private Worker worker;
+    @Mock private WorkerConfigTransformer transformer;
     @Mock private Callback> putConnectorCallback;
-    @Mock
-    private Plugins plugins;
-    @Mock
-    private PluginClassLoader pluginLoader;
-    @Mock
-    private DelegatingClassLoader delegatingLoader;
+    @Mock private Plugins plugins;
+    @Mock private PluginClassLoader pluginLoader;
+    @Mock private DelegatingClassLoader delegatingLoader;
 
     private ConfigBackingStore.UpdateListener configUpdateListener;
     private WorkerRebalanceListener rebalanceListener;
 
     private SinkConnectorConfig conn1SinkConfig;
     private SinkConnectorConfig conn1SinkConfigUpdated;
+    private short connectProtocolVersion;
+    private final ConnectorClientConfigOverridePolicy
+        noneConnectorClientConfigOverridePolicy = new NoneConnectorClientConfigOverridePolicy();
+
 
     @Before
     public void setUp() throws Exception {
@@ -178,11 +209,16 @@ public void setUp() throws Exception {
         worker = PowerMock.createMock(Worker.class);
         EasyMock.expect(worker.isSinkConnector(CONN1)).andStubReturn(Boolean.TRUE);
 
-        herder = PowerMock.createPartialMock(DistributedHerder.class, new String[]{"backoff", "connectorTypeForClass", "updateDeletedConnectorStatus"},
-                new DistributedConfig(HERDER_CONFIG), worker, WORKER_ID, statusBackingStore, configBackingStore, member, MEMBER_URL, metrics, time);
+        // Default to the old protocol unless specified otherwise
+        connectProtocolVersion = CONNECT_PROTOCOL_V0;
+
+        herder = PowerMock.createPartialMock(DistributedHerder.class,
+                new String[]{"connectorTypeForClass", "updateDeletedConnectorStatus", "updateDeletedTaskStatus", "validateConnectorConfig"},
+                new DistributedConfig(HERDER_CONFIG), worker, WORKER_ID, KAFKA_CLUSTER_ID,
+                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);
@@ -191,6 +227,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
@@ -202,16 +239,26 @@ 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);
-        worker.startConnector(EasyMock.eq(CONN1), EasyMock.>anyObject(), EasyMock.anyObject(),
-                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
-        PowerMock.expectLastCall().andReturn(true);
+        Capture> onStart = newCapture();
+        worker.startConnector(EasyMock.eq(CONN1), EasyMock.>anyObject(), EasyMock.anyObject(),
+                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED), capture(onStart));
+        PowerMock.expectLastCall().andAnswer(new IAnswer() {
+            @Override
+            public Boolean answer() throws Throwable {
+                onStart.getValue().onCompletion(null, TargetState.STARTED);
+                return true;
+            }
+        });
+        member.wakeup();
+        PowerMock.expectLastCall();
         EasyMock.expect(worker.isRunning(CONN1)).andReturn(true);
-
+        PowerMock.expectLastCall();
         EasyMock.expect(worker.connectorTaskConfigs(CONN1, conn1SinkConfig)).andReturn(TASK_CONFIGS);
-        worker.startTask(EasyMock.eq(TASK1), EasyMock.>anyObject(), EasyMock.>anyObject(),
+        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());
@@ -230,15 +277,25 @@ 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);
-        worker.startConnector(EasyMock.eq(CONN1), EasyMock.>anyObject(), EasyMock.anyObject(),
-                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
-        PowerMock.expectLastCall().andReturn(true);
+        Capture> onFirstStart = newCapture();
+        worker.startConnector(EasyMock.eq(CONN1), EasyMock.>anyObject(), EasyMock.anyObject(),
+                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED), capture(onFirstStart));
+        PowerMock.expectLastCall().andAnswer(new IAnswer() {
+            @Override
+            public Boolean answer() throws Throwable {
+                onFirstStart.getValue().onCompletion(null, TargetState.STARTED);
+                return true;
+            }
+        });
+        member.wakeup();
+        PowerMock.expectLastCall();
         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(),
+        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());
@@ -249,9 +306,18 @@ public void testRebalance() throws Exception {
                 1, Arrays.asList(CONN1), Arrays.asList());
 
         // 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);
+        Capture> onSecondStart = newCapture();
+        worker.startConnector(EasyMock.eq(CONN1), EasyMock.>anyObject(), EasyMock.anyObject(),
+                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED), capture(onSecondStart));
+        PowerMock.expectLastCall().andAnswer(new IAnswer() {
+            @Override
+            public Boolean answer() throws Throwable {
+                onSecondStart.getValue().onCompletion(null, TargetState.STARTED);
+                return true;
+            }
+        });
+        member.wakeup();
+        PowerMock.expectLastCall();
         EasyMock.expect(worker.isRunning(CONN1)).andReturn(true);
         EasyMock.expect(worker.connectorTaskConfigs(CONN1, conn1SinkConfig)).andReturn(TASK_CONFIGS);
         member.poll(EasyMock.anyInt());
@@ -273,19 +339,192 @@ 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
+        Capture> onStart = newCapture();
+        worker.startConnector(EasyMock.eq(CONN1), EasyMock.anyObject(), EasyMock.anyObject(),
+                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED), capture(onStart));
+        PowerMock.expectLastCall().andAnswer(new IAnswer() {
+            @Override
+            public Boolean answer() throws Throwable {
+                onStart.getValue().onCompletion(null, TargetState.STARTED);
+                return true;
+            }
+        });
+        member.wakeup();
+        PowerMock.expectLastCall();
+        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();
+
+        // 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();
+
+        herder.configState = SNAPSHOT;
+        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
+        Capture> onStart = newCapture();
+        worker.startConnector(EasyMock.eq(CONN1), EasyMock.anyObject(), EasyMock.anyObject(),
+                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED), capture(onStart));
+        PowerMock.expectLastCall().andAnswer(new IAnswer() {
+            @Override
+            public Boolean answer() throws Throwable {
+                onStart.getValue().onCompletion(null, TargetState.STARTED);
+                return true;
+            }
+        });
+        member.wakeup();
+        PowerMock.expectLastCall();
+        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);
-        worker.startConnector(EasyMock.eq(CONN1), EasyMock.>anyObject(), EasyMock.anyObject(),
-                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
-        PowerMock.expectLastCall().andReturn(true);
+        Capture> onFirstStart = newCapture();
+        worker.startConnector(EasyMock.eq(CONN1), EasyMock.>anyObject(), EasyMock.anyObject(),
+                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED), capture(onFirstStart));
+        PowerMock.expectLastCall().andAnswer(new IAnswer() {
+            @Override
+            public Boolean answer() throws Throwable {
+                onFirstStart.getValue().onCompletion(null, TargetState.STARTED);
+                return true;
+            }
+        });
+        member.wakeup();
+        PowerMock.expectLastCall();
         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(),
+        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());
@@ -295,9 +534,18 @@ public void testRebalanceFailedConnector() throws Exception {
                 1, Arrays.asList(CONN1), Arrays.asList());
 
         // 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);
+        Capture> onSecondStart = newCapture();
+        worker.startConnector(EasyMock.eq(CONN1), EasyMock.>anyObject(), EasyMock.anyObject(),
+                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED), capture(onSecondStart));
+        PowerMock.expectLastCall().andAnswer(new IAnswer() {
+            @Override
+            public Boolean answer() throws Throwable {
+                onSecondStart.getValue().onCompletion(null, TargetState.STARTED);
+                return true;
+            }
+        });
+        member.wakeup();
+        PowerMock.expectLastCall();
         EasyMock.expect(worker.isRunning(CONN1)).andReturn(false);
 
         // worker is not running, so we should see no call to connectorTaskConfigs()
@@ -321,8 +569,8 @@ public void testRebalanceFailedConnector() throws Exception {
     @Test
     public void testHaltCleansUpWorker() {
         EasyMock.expect(worker.connectorNames()).andReturn(Collections.singleton(CONN1));
-        worker.stopConnector(CONN1);
-        PowerMock.expectLastCall().andReturn(true);
+        worker.stopAndAwaitConnector(CONN1);
+        PowerMock.expectLastCall();
         EasyMock.expect(worker.taskIds()).andReturn(Collections.singleton(TASK1));
         worker.stopAndAwaitTask(TASK1);
         PowerMock.expectLastCall();
@@ -345,20 +593,24 @@ 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);
 
         member.wakeup();
         PowerMock.expectLastCall();
 
-        // config validation
-        Connector connectorMock = PowerMock.createMock(Connector.class);
-        EasyMock.expect(worker.getPlugins()).andReturn(plugins).times(3);
-        EasyMock.expect(plugins.compareAndSwapLoaders(connectorMock)).andReturn(delegatingLoader);
-        EasyMock.expect(plugins.newConnector(EasyMock.anyString())).andReturn(connectorMock);
-        EasyMock.expect(connectorMock.config()).andReturn(new ConfigDef());
-        EasyMock.expect(connectorMock.validate(CONN2_CONFIG)).andReturn(new Config(Collections.emptyList()));
-        EasyMock.expect(Plugins.compareAndSwapLoaders(delegatingLoader)).andReturn(pluginLoader);
+        // mock the actual validation since its asynchronous nature is difficult to test and should
+        // be covered sufficiently by the unit tests for the AbstractHerder class
+        Capture> validateCallback = newCapture();
+        herder.validateConnectorConfig(EasyMock.eq(CONN2_CONFIG), capture(validateCallback));
+        PowerMock.expectLastCall().andAnswer(new IAnswer() {
+            @Override
+            public Void answer() throws Throwable {
+                validateCallback.getValue().onCompletion(null, CONN2_CONFIG_INFOS);
+                return null;
+            }
+        });
 
         // CONN2 is new, should succeed
         configBackingStore.putConnectorConfig(CONN2, CONN2_CONFIG);
@@ -369,11 +621,25 @@ public void testCreateConnector() throws Exception {
         PowerMock.expectLastCall();
         member.poll(EasyMock.anyInt());
         PowerMock.expectLastCall();
+        // These will occur just before/during the second tick
+        member.wakeup();
+        PowerMock.expectLastCall();
+        member.ensureActive();
+        PowerMock.expectLastCall();
+        member.poll(EasyMock.anyInt());
+        PowerMock.expectLastCall();
+
         // No immediate action besides this -- change will be picked up via the config log
 
         PowerMock.replayAll();
 
         herder.putConnectorConfig(CONN2, CONN2_CONFIG, false, putConnectorCallback);
+        // First tick runs the initial herder request, which issues an asynchronous request for
+        // connector validation
+        herder.tick();
+
+        // Once that validation is complete, another request is added to the herder request queue
+        // for actually performing the config write; this tick is for that request
         herder.tick();
 
         time.sleep(1000L);
@@ -383,8 +649,9 @@ public void testCreateConnector() throws Exception {
     }
 
     @Test
-    public void testCreateConnectorFailedBasicValidation() throws Exception {
+    public void testCreateConnectorFailedValidation() throws Exception {
         EasyMock.expect(member.memberId()).andStubReturn("leader");
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
         expectRebalance(1, Collections.emptyList(), Collections.emptyList());
         expectPostRebalanceCatchup(SNAPSHOT);
 
@@ -394,79 +661,41 @@ public void testCreateConnectorFailedBasicValidation() throws Exception {
         member.wakeup();
         PowerMock.expectLastCall();
 
-        // config validation
-        Connector connectorMock = PowerMock.createMock(Connector.class);
-        EasyMock.expect(worker.getPlugins()).andReturn(plugins).times(3);
-        EasyMock.expect(plugins.compareAndSwapLoaders(connectorMock)).andReturn(delegatingLoader);
-        EasyMock.expect(plugins.newConnector(EasyMock.anyString())).andReturn(connectorMock);
-
-        EasyMock.expect(connectorMock.config()).andStubReturn(new ConfigDef());
-        ConfigValue validatedValue = new ConfigValue("foo.bar");
-        EasyMock.expect(connectorMock.validate(config)).andReturn(new Config(singletonList(validatedValue)));
-
-        EasyMock.expect(Plugins.compareAndSwapLoaders(delegatingLoader)).andReturn(pluginLoader);
-
-        // CONN2 creation should fail
+        // mock the actual validation since its asynchronous nature is difficult to test and should
+        // be covered sufficiently by the unit tests for the AbstractHerder class
+        Capture> validateCallback = newCapture();
+        herder.validateConnectorConfig(EasyMock.eq(config), capture(validateCallback));
+        PowerMock.expectLastCall().andAnswer(new IAnswer() {
+            @Override
+            public Void answer() throws Throwable {
+                // CONN2 creation should fail
+                validateCallback.getValue().onCompletion(null, CONN2_INVALID_CONFIG_INFOS);
+                return null;
+            }
+        });
 
-        Capture error = EasyMock.newCapture();
-        putConnectorCallback.onCompletion(EasyMock.capture(error), EasyMock.>isNull());
+        Capture error = newCapture();
+        putConnectorCallback.onCompletion(capture(error), EasyMock.>isNull());
         PowerMock.expectLastCall();
 
         member.poll(EasyMock.anyInt());
         PowerMock.expectLastCall();
-        // No immediate action besides this -- change will be picked up via the config log
-
-        PowerMock.replayAll();
-
-        herder.putConnectorConfig(CONN2, config, false, putConnectorCallback);
-        herder.tick();
-
-        assertTrue(error.hasCaptured());
-        assertTrue(error.getValue() instanceof BadRequestException);
-
-        time.sleep(1000L);
-        assertStatistics(3, 1, 100, 1000L);
-
-        PowerMock.verifyAll();
-    }
-
-    @Test
-    public void testCreateConnectorFailedCustomValidation() throws Exception {
-        EasyMock.expect(member.memberId()).andStubReturn("leader");
-        expectRebalance(1, Collections.emptyList(), Collections.emptyList());
-        expectPostRebalanceCatchup(SNAPSHOT);
 
+        // These will occur just before/during the second tick
         member.wakeup();
         PowerMock.expectLastCall();
 
-        // config validation
-        Connector connectorMock = PowerMock.createMock(Connector.class);
-        EasyMock.expect(worker.getPlugins()).andReturn(plugins).times(3);
-        EasyMock.expect(plugins.compareAndSwapLoaders(connectorMock)).andReturn(delegatingLoader);
-        EasyMock.expect(plugins.newConnector(EasyMock.anyString())).andReturn(connectorMock);
-
-        ConfigDef configDef = new ConfigDef();
-        configDef.define("foo.bar", ConfigDef.Type.STRING, ConfigDef.Importance.HIGH, "foo.bar doc");
-        EasyMock.expect(connectorMock.config()).andReturn(configDef);
-
-        ConfigValue validatedValue = new ConfigValue("foo.bar");
-        validatedValue.addErrorMessage("Failed foo.bar validation");
-        EasyMock.expect(connectorMock.validate(CONN2_CONFIG)).andReturn(new Config(singletonList(validatedValue)));
-        EasyMock.expect(Plugins.compareAndSwapLoaders(delegatingLoader)).andReturn(pluginLoader);
-
-        // CONN2 creation should fail
-
-        Capture error = EasyMock.newCapture();
-        putConnectorCallback.onCompletion(EasyMock.capture(error), EasyMock.>isNull());
+        member.ensureActive();
         PowerMock.expectLastCall();
-
         member.poll(EasyMock.anyInt());
         PowerMock.expectLastCall();
+
         // No immediate action besides this -- change will be picked up via the config log
 
         PowerMock.replayAll();
 
-        herder.putConnectorConfig(CONN2, CONN2_CONFIG, false, putConnectorCallback);
+        herder.putConnectorConfig(CONN2, config, false, putConnectorCallback);
+        herder.tick();
         herder.tick();
 
         assertTrue(error.hasCaptured());
@@ -478,57 +707,41 @@ public void testCreateConnectorFailedCustomValidation() throws Exception {
         PowerMock.verifyAll();
     }
 
+    @SuppressWarnings("unchecked")
     @Test
     public void testConnectorNameConflictsWithWorkerGroupId() throws Exception {
-        EasyMock.expect(member.memberId()).andStubReturn("leader");
-        expectRebalance(1, Collections.emptyList(), Collections.emptyList());
-        expectPostRebalanceCatchup(SNAPSHOT);
-
-        member.wakeup();
-        PowerMock.expectLastCall();
-
         Map config = new HashMap<>(CONN2_CONFIG);
         config.put(ConnectorConfig.NAME_CONFIG, "test-group");
 
-        // config validation
         Connector connectorMock = PowerMock.createMock(SinkConnector.class);
-        EasyMock.expect(worker.getPlugins()).andReturn(plugins).times(3);
-        EasyMock.expect(plugins.compareAndSwapLoaders(connectorMock)).andReturn(delegatingLoader);
-        EasyMock.expect(plugins.newConnector(EasyMock.anyString())).andReturn(connectorMock);
-        EasyMock.expect(connectorMock.config()).andReturn(new ConfigDef());
-        EasyMock.expect(connectorMock.validate(config)).andReturn(new Config(Collections.emptyList()));
-        EasyMock.expect(Plugins.compareAndSwapLoaders(delegatingLoader)).andReturn(pluginLoader);
 
         // CONN2 creation should fail because the worker group id (connect-test-group) conflicts with
         // the consumer group id we would use for this sink
+        Map validatedConfigs =
+            herder.validateBasicConnectorConfig(connectorMock, ConnectorConfig.configDef(), config);
 
-        Capture error = EasyMock.newCapture();
-        putConnectorCallback.onCompletion(EasyMock.capture(error), EasyMock.isNull(Herder.Created.class));
-        PowerMock.expectLastCall();
-
-        member.poll(EasyMock.anyInt());
-        PowerMock.expectLastCall();
-        // No immediate action besides this -- change will be picked up via the config log
-
-        PowerMock.replayAll();
-
-        herder.putConnectorConfig(CONN2, config, false, putConnectorCallback);
-        herder.tick();
-
-        assertTrue(error.hasCaptured());
-        assertTrue(error.getValue() instanceof BadRequestException);
-
-        time.sleep(1000L);
-        assertStatistics(3, 1, 100, 1000L);
-
-        PowerMock.verifyAll();
+        ConfigValue nameConfig = validatedConfigs.get(ConnectorConfig.NAME_CONFIG);
+        assertNotNull(nameConfig.errorMessages());
+        assertFalse(nameConfig.errorMessages().isEmpty());
     }
 
     @Test
     public void testCreateConnectorAlreadyExists() throws Exception {
         EasyMock.expect(member.memberId()).andStubReturn("leader");
-        EasyMock.expect(worker.getPlugins()).andReturn(plugins);
-        EasyMock.expect(plugins.newConnector(EasyMock.anyString())).andReturn(null);
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
+
+        // mock the actual validation since its asynchronous nature is difficult to test and should
+        // be covered sufficiently by the unit tests for the AbstractHerder class
+        Capture> validateCallback = newCapture();
+        herder.validateConnectorConfig(EasyMock.eq(CONN1_CONFIG), capture(validateCallback));
+        PowerMock.expectLastCall().andAnswer(new IAnswer() {
+            @Override
+            public Void answer() throws Throwable {
+                validateCallback.getValue().onCompletion(null, CONN1_CONFIG_INFOS);
+                return null;
+            }
+        });
+
         expectRebalance(1, Collections.emptyList(), Collections.emptyList());
         expectPostRebalanceCatchup(SNAPSHOT);
 
@@ -539,12 +752,22 @@ public void testCreateConnectorAlreadyExists() throws Exception {
         PowerMock.expectLastCall();
         member.poll(EasyMock.anyInt());
         PowerMock.expectLastCall();
+
+        // These will occur just before/during the second tick
+        member.wakeup();
+        PowerMock.expectLastCall();
+        member.ensureActive();
+        PowerMock.expectLastCall();
+        member.poll(EasyMock.anyInt());
+        PowerMock.expectLastCall();
+
         // No immediate action besides this -- change will be picked up via the config log
 
         PowerMock.replayAll();
 
         herder.putConnectorConfig(CONN1, CONN1_CONFIG, false, putConnectorCallback);
         herder.tick();
+        herder.tick();
 
         time.sleep(1000L);
         assertStatistics(3, 1, 100, 1000L);
@@ -555,13 +778,21 @@ 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());
         expectPostRebalanceCatchup(SNAPSHOT);
-        worker.startConnector(EasyMock.eq(CONN1), EasyMock.>anyObject(), EasyMock.anyObject(),
-                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
-        PowerMock.expectLastCall().andReturn(true);
+        Capture> onStart = newCapture();
+        worker.startConnector(EasyMock.eq(CONN1), EasyMock.>anyObject(), EasyMock.anyObject(),
+                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED), capture(onStart));
+        PowerMock.expectLastCall().andAnswer(new IAnswer() {
+            @Override
+            public Boolean answer() throws Throwable {
+                onStart.getValue().onCompletion(null, TargetState.STARTED);
+                return true;
+            }
+        });
         EasyMock.expect(worker.isRunning(CONN1)).andReturn(true);
         EasyMock.expect(worker.connectorTaskConfigs(CONN1, conn1SinkConfig)).andReturn(TASK_CONFIGS);
 
@@ -574,15 +805,37 @@ public void testDestroyConnector() throws Exception {
         PowerMock.expectLastCall();
         member.poll(EasyMock.anyInt());
         PowerMock.expectLastCall();
-        // No immediate action besides this -- change will be picked up via the config log
 
+        // The change eventually is reflected to the config topic and the deleted connector and
+        // tasks are revoked
+        member.wakeup();
+        PowerMock.expectLastCall();
+        TopicStatus fooStatus = new TopicStatus(FOO_TOPIC, CONN1, 0, time.milliseconds());
+        TopicStatus barStatus = new TopicStatus(BAR_TOPIC, CONN1, 0, time.milliseconds());
+        EasyMock.expect(statusBackingStore.getAllTopics(EasyMock.eq(CONN1))).andReturn(new HashSet<>(Arrays.asList(fooStatus, barStatus))).times(2);
+        statusBackingStore.deleteTopic(EasyMock.eq(CONN1), EasyMock.eq(FOO_TOPIC));
+        PowerMock.expectLastCall().times(2);
+        statusBackingStore.deleteTopic(EasyMock.eq(CONN1), EasyMock.eq(BAR_TOPIC));
+        PowerMock.expectLastCall().times(2);
+        expectRebalance(Arrays.asList(CONN1), Arrays.asList(TASK1),
+                ConnectProtocol.Assignment.NO_ERROR, 2,
+                Collections.emptyList(), Collections.emptyList(), 0);
+        expectPostRebalanceCatchup(ClusterConfigState.EMPTY);
+        member.requestRejoin();
+        PowerMock.expectLastCall();
         PowerMock.replayAll();
 
         herder.deleteConnectorConfig(CONN1, putConnectorCallback);
         herder.tick();
 
         time.sleep(1000L);
-        assertStatistics(3, 1, 100, 1000L);
+        assertStatistics("leaderUrl", false, 3, 1, 100, 1000L);
+
+        configUpdateListener.onConnectorConfigRemove(CONN1); // read updated config that removes the connector
+        herder.configState = ClusterConfigState.EMPTY;
+        herder.tick();
+        time.sleep(1000L);
+        assertStatistics("leaderUrl", true, 3, 1, 100, 2100L);
 
         PowerMock.verifyAll();
     }
@@ -593,14 +846,24 @@ 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);
         member.poll(EasyMock.anyInt());
         PowerMock.expectLastCall();
-        worker.startConnector(EasyMock.eq(CONN1), EasyMock.>anyObject(), EasyMock.anyObject(),
-                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
-        PowerMock.expectLastCall().andReturn(true);
+        Capture> onFirstStart = newCapture();
+        worker.startConnector(EasyMock.eq(CONN1), EasyMock.>anyObject(), EasyMock.anyObject(),
+                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED), capture(onFirstStart));
+        PowerMock.expectLastCall().andAnswer(new IAnswer() {
+            @Override
+            public Boolean answer() throws Throwable {
+                onFirstStart.getValue().onCompletion(null, TargetState.STARTED);
+                return true;
+            }
+        });
+        member.wakeup();
+        PowerMock.expectLastCall();
         EasyMock.expect(worker.isRunning(CONN1)).andReturn(true);
 
         // now handle the connector restart
@@ -611,12 +874,21 @@ public void testRestartConnector() throws Exception {
         member.poll(EasyMock.anyInt());
         PowerMock.expectLastCall();
 
-        worker.stopConnector(CONN1);
-        PowerMock.expectLastCall().andReturn(true);
+        worker.stopAndAwaitConnector(CONN1);
+        PowerMock.expectLastCall();
         EasyMock.expect(worker.getPlugins()).andReturn(plugins);
-        worker.startConnector(EasyMock.eq(CONN1), EasyMock.>anyObject(), EasyMock.anyObject(),
-                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
-        PowerMock.expectLastCall().andReturn(true);
+        Capture> onSecondStart = newCapture();
+        worker.startConnector(EasyMock.eq(CONN1), EasyMock.>anyObject(), EasyMock.anyObject(),
+                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED), capture(onSecondStart));
+        PowerMock.expectLastCall().andAnswer(new IAnswer() {
+            @Override
+            public Boolean answer() throws Throwable {
+                onSecondStart.getValue().onCompletion(null, TargetState.STARTED);
+                return true;
+            }
+        });
+        member.wakeup();
+        PowerMock.expectLastCall();
         EasyMock.expect(worker.isRunning(CONN1)).andReturn(true);
 
         PowerMock.replayAll();
@@ -634,12 +906,12 @@ 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());
         PowerMock.expectLastCall();
 
-
         // now handle the connector restart
         member.wakeup();
         PowerMock.expectLastCall();
@@ -668,6 +940,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());
@@ -702,6 +975,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());
@@ -749,11 +1023,12 @@ 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());
         PowerMock.expectLastCall();
-        worker.startTask(EasyMock.eq(TASK0), EasyMock.>anyObject(), EasyMock.>anyObject(),
+        worker.startTask(EasyMock.eq(TASK0), EasyMock.anyObject(), EasyMock.>anyObject(), EasyMock.>anyObject(),
                 EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
         PowerMock.expectLastCall().andReturn(true);
 
@@ -767,7 +1042,7 @@ public void testRestartTask() throws Exception {
 
         worker.stopAndAwaitTask(TASK0);
         PowerMock.expectLastCall();
-        worker.startTask(EasyMock.eq(TASK0), EasyMock.>anyObject(), EasyMock.>anyObject(),
+        worker.startTask(EasyMock.eq(TASK0), EasyMock.anyObject(), EasyMock.>anyObject(), EasyMock.>anyObject(),
                 EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
         PowerMock.expectLastCall().andReturn(true);
 
@@ -786,6 +1061,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());
@@ -816,11 +1092,11 @@ public void testRestartUnknownTask() throws Exception {
     }
 
     @Test
-    public void testRequestProcessingOrder() throws Exception {
-        final DistributedHerder.HerderRequest req1 = herder.addRequest(100, null, null);
-        final DistributedHerder.HerderRequest req2 = herder.addRequest(10, null, null);
-        final DistributedHerder.HerderRequest req3 = herder.addRequest(200, null, null);
-        final DistributedHerder.HerderRequest req4 = herder.addRequest(200, null, null);
+    public void testRequestProcessingOrder() {
+        final DistributedHerder.DistributedHerderRequest req1 = herder.addRequest(100, null, null);
+        final DistributedHerder.DistributedHerderRequest req2 = herder.addRequest(10, null, null);
+        final DistributedHerder.DistributedHerderRequest req3 = herder.addRequest(200, null, null);
+        final DistributedHerder.DistributedHerderRequest req4 = herder.addRequest(200, null, null);
 
         assertEquals(req2, herder.requests.pollFirst()); // lowest delay
         assertEquals(req1, herder.requests.pollFirst()); // next lowest delay
@@ -832,6 +1108,7 @@ public void testRequestProcessingOrder() throws Exception {
     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());
@@ -866,6 +1143,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());
@@ -904,6 +1182,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());
@@ -912,6 +1191,7 @@ public void testConnectorConfigAdded() {
 
         // apply config
         member.wakeup();
+        PowerMock.expectLastCall();
         member.ensureActive();
         PowerMock.expectLastCall();
         // Checks for config updates and starts rebalance
@@ -921,9 +1201,19 @@ public void testConnectorConfigAdded() {
         // Performs rebalance and gets new assignment
         expectRebalance(Collections.emptyList(), Collections.emptyList(),
                 ConnectProtocol.Assignment.NO_ERROR, 1, Arrays.asList(CONN1), Collections.emptyList());
-        worker.startConnector(EasyMock.eq(CONN1), EasyMock.>anyObject(), EasyMock.anyObject(),
-                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
-        PowerMock.expectLastCall().andReturn(true);
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
+        Capture> onStart = newCapture();
+        worker.startConnector(EasyMock.eq(CONN1), EasyMock.>anyObject(), EasyMock.anyObject(),
+                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED), capture(onStart));
+        PowerMock.expectLastCall().andAnswer(new IAnswer() {
+            @Override
+            public Boolean answer() throws Throwable {
+                onStart.getValue().onCompletion(null, TargetState.STARTED);
+                return true;
+            }
+        });
+        member.wakeup();
+        PowerMock.expectLastCall();
         EasyMock.expect(worker.getPlugins()).andReturn(plugins);
         EasyMock.expect(worker.isRunning(CONN1)).andReturn(true);
         EasyMock.expect(worker.connectorTaskConfigs(CONN1, conn1SinkConfig)).andReturn(TASK_CONFIGS);
@@ -945,14 +1235,25 @@ 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);
-        worker.startConnector(EasyMock.eq(CONN1), EasyMock.>anyObject(), EasyMock.anyObject(),
-                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
-        PowerMock.expectLastCall().andReturn(true);
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
+        Capture> onFirstStart = newCapture();
+        worker.startConnector(EasyMock.eq(CONN1), EasyMock.>anyObject(), EasyMock.anyObject(),
+                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED), capture(onFirstStart));
+        PowerMock.expectLastCall().andAnswer(new IAnswer() {
+            @Override
+            public Boolean answer() throws Throwable {
+                onFirstStart.getValue().onCompletion(null, TargetState.STARTED);
+                return true;
+            }
+        });
+        member.wakeup();
+        PowerMock.expectLastCall();
         EasyMock.expect(worker.getPlugins()).andReturn(plugins);
         EasyMock.expect(worker.isRunning(CONN1)).andReturn(true);
         EasyMock.expect(worker.connectorTaskConfigs(CONN1, conn1SinkConfig)).andReturn(TASK_CONFIGS);
@@ -961,25 +1262,43 @@ public void testConnectorConfigUpdate() throws Exception {
 
         // apply config
         member.wakeup();
+        PowerMock.expectLastCall();
         member.ensureActive();
         PowerMock.expectLastCall();
         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);
-        worker.startConnector(EasyMock.eq(CONN1), EasyMock.>anyObject(), EasyMock.anyObject(),
-                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
-        PowerMock.expectLastCall().andReturn(true);
+        worker.stopAndAwaitConnector(CONN1);
+        PowerMock.expectLastCall();
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
+        Capture> onSecondStart = newCapture();
+        worker.startConnector(EasyMock.eq(CONN1), EasyMock.>anyObject(), EasyMock.anyObject(),
+                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED), capture(onSecondStart));
+        PowerMock.expectLastCall().andAnswer(new IAnswer() {
+            @Override
+            public Boolean answer() throws Throwable {
+                onSecondStart.getValue().onCompletion(null, TargetState.STARTED);
+                return true;
+            }
+        });
+        member.wakeup();
+        PowerMock.expectLastCall();
         EasyMock.expect(worker.getPlugins()).andReturn(plugins);
         EasyMock.expect(worker.isRunning(CONN1)).andReturn(true);
         EasyMock.expect(worker.connectorTaskConfigs(CONN1, conn1SinkConfig)).andReturn(TASK_CONFIGS);
         member.poll(EasyMock.anyInt());
         PowerMock.expectLastCall();
 
+        // These will occur just before/during the third tick
+        member.ensureActive();
+        PowerMock.expectLastCall();
+        member.poll(EasyMock.anyInt());
+        PowerMock.expectLastCall();
+
         PowerMock.replayAll();
 
         herder.tick(); // join
         configUpdateListener.onConnectorConfigUpdate(CONN1); // read updated config
         herder.tick(); // apply config
+        herder.tick();
 
         PowerMock.verifyAll();
     }
@@ -989,14 +1308,25 @@ 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);
-        worker.startConnector(EasyMock.eq(CONN1), EasyMock.>anyObject(), EasyMock.anyObject(),
-                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
-        PowerMock.expectLastCall().andReturn(true);
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
+        Capture> onStart = newCapture();
+        worker.startConnector(EasyMock.eq(CONN1), EasyMock.>anyObject(), EasyMock.anyObject(),
+                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED), capture(onStart));
+        PowerMock.expectLastCall().andAnswer(new IAnswer() {
+            @Override
+            public Boolean answer() throws Throwable {
+                onStart.getValue().onCompletion(null, TargetState.STARTED);
+                return true;
+            }
+        });
+        member.wakeup();
+        PowerMock.expectLastCall();
         EasyMock.expect(worker.getPlugins()).andReturn(plugins);
         EasyMock.expect(worker.isRunning(CONN1)).andReturn(true);
         EasyMock.expect(worker.connectorTaskConfigs(CONN1, conn1SinkConfig)).andReturn(TASK_CONFIGS);
@@ -1005,15 +1335,29 @@ public void testConnectorPaused() throws Exception {
 
         // handle the state change
         member.wakeup();
+        PowerMock.expectLastCall();
         member.ensureActive();
         PowerMock.expectLastCall();
 
         EasyMock.expect(configBackingStore.snapshot()).andReturn(SNAPSHOT_PAUSED_CONN1);
         PowerMock.expectLastCall();
 
-        worker.setTargetState(CONN1, TargetState.PAUSED);
+        Capture> onPause = newCapture();
+        worker.setTargetState(EasyMock.eq(CONN1), EasyMock.eq(TargetState.PAUSED), capture(onPause));
+        PowerMock.expectLastCall().andAnswer(new IAnswer() {
+            @Override
+            public Void answer() throws Throwable {
+                onStart.getValue().onCompletion(null, TargetState.PAUSED);
+                return null;
+            }
+        });
+
+        member.poll(EasyMock.anyInt());
         PowerMock.expectLastCall();
 
+        // These will occur just before/during the third tick
+        member.ensureActive();
+        PowerMock.expectLastCall();
         member.poll(EasyMock.anyInt());
         PowerMock.expectLastCall();
 
@@ -1022,6 +1366,7 @@ public void testConnectorPaused() throws Exception {
         herder.tick(); // join
         configUpdateListener.onConnectorTargetStateChange(CONN1); // state changes to paused
         herder.tick(); // worker should apply the state change
+        herder.tick();
 
         PowerMock.verifyAll();
     }
@@ -1029,14 +1374,23 @@ 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);
-        worker.startConnector(EasyMock.eq(CONN1), EasyMock.>anyObject(), EasyMock.anyObject(),
-                EasyMock.eq(herder), EasyMock.eq(TargetState.PAUSED));
-        PowerMock.expectLastCall().andReturn(true);
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
+        Capture> onStart = newCapture();
+        worker.startConnector(EasyMock.eq(CONN1), EasyMock.>anyObject(), EasyMock.anyObject(),
+                EasyMock.eq(herder), EasyMock.eq(TargetState.PAUSED), capture(onStart));
+        PowerMock.expectLastCall().andAnswer(new IAnswer() {
+            @Override
+            public Boolean answer() throws Throwable {
+                onStart.getValue().onCompletion(null, TargetState.PAUSED);
+                return true;
+            }
+        });
         EasyMock.expect(worker.getPlugins()).andReturn(plugins);
 
         member.poll(EasyMock.anyInt());
@@ -1044,6 +1398,7 @@ public void testConnectorResumed() throws Exception {
 
         // handle the state change
         member.wakeup();
+        PowerMock.expectLastCall();
         member.ensureActive();
         PowerMock.expectLastCall();
 
@@ -1054,9 +1409,24 @@ public void testConnectorResumed() throws Exception {
         EasyMock.expect(worker.isRunning(CONN1)).andReturn(true);
         EasyMock.expect(worker.connectorTaskConfigs(CONN1, conn1SinkConfig)).andReturn(TASK_CONFIGS);
 
-        worker.setTargetState(CONN1, TargetState.STARTED);
+        Capture> onResume = newCapture();
+        worker.setTargetState(EasyMock.eq(CONN1), EasyMock.eq(TargetState.STARTED), capture(onResume));
+        PowerMock.expectLastCall().andAnswer(new IAnswer() {
+            @Override
+            public Void answer() throws Throwable {
+                onResume.getValue().onCompletion(null, TargetState.STARTED);
+                return null;
+            }
+        });
+        member.wakeup();
+        PowerMock.expectLastCall();
+
+        member.poll(EasyMock.anyInt());
         PowerMock.expectLastCall();
 
+        // These will occur just before/during the third tick
+        member.ensureActive();
+        PowerMock.expectLastCall();
         member.poll(EasyMock.anyInt());
         PowerMock.expectLastCall();
 
@@ -1065,6 +1435,7 @@ public void testConnectorResumed() throws Exception {
         herder.tick(); // join
         configUpdateListener.onConnectorTargetStateChange(CONN1); // state changes to started
         herder.tick(); // apply state change
+        herder.tick();
 
         PowerMock.verifyAll();
     }
@@ -1072,12 +1443,13 @@ 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
         expectRebalance(1, Collections.emptyList(), singletonList(TASK0));
         expectPostRebalanceCatchup(SNAPSHOT);
-        worker.startTask(EasyMock.eq(TASK0), EasyMock.>anyObject(), EasyMock.>anyObject(),
+        worker.startTask(EasyMock.eq(TASK0), EasyMock.anyObject(), EasyMock.>anyObject(), EasyMock.>anyObject(),
                 EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
         PowerMock.expectLastCall().andReturn(true);
         member.poll(EasyMock.anyInt());
@@ -1085,6 +1457,7 @@ public void testUnknownConnectorPaused() throws Exception {
 
         // state change is ignored since we have no target state
         member.wakeup();
+        PowerMock.expectLastCall();
         member.ensureActive();
         PowerMock.expectLastCall();
 
@@ -1109,12 +1482,13 @@ 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
         expectRebalance(1, Collections.emptyList(), singletonList(TASK0));
         expectPostRebalanceCatchup(SNAPSHOT);
-        worker.startTask(EasyMock.eq(TASK0), EasyMock.>anyObject(), EasyMock.>anyObject(),
+        worker.startTask(EasyMock.eq(TASK0), EasyMock.anyObject(), EasyMock.>anyObject(), EasyMock.>anyObject(),
                 EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
         PowerMock.expectLastCall().andReturn(true);
         member.poll(EasyMock.anyInt());
@@ -1122,13 +1496,23 @@ public void testConnectorPausedRunningTaskOnly() throws Exception {
 
         // handle the state change
         member.wakeup();
+        PowerMock.expectLastCall();
         member.ensureActive();
         PowerMock.expectLastCall();
 
         EasyMock.expect(configBackingStore.snapshot()).andReturn(SNAPSHOT_PAUSED_CONN1);
         PowerMock.expectLastCall();
 
-        worker.setTargetState(CONN1, TargetState.PAUSED);
+        Capture> onPause = newCapture();
+        worker.setTargetState(EasyMock.eq(CONN1), EasyMock.eq(TargetState.PAUSED), capture(onPause));
+        PowerMock.expectLastCall().andAnswer(new IAnswer() {
+            @Override
+            public Void answer() throws Throwable {
+                onPause.getValue().onCompletion(null, TargetState.STARTED);
+                return null;
+            }
+        });
+        member.wakeup();
         PowerMock.expectLastCall();
 
         member.poll(EasyMock.anyInt());
@@ -1149,12 +1533,13 @@ 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
         expectRebalance(1, Collections.emptyList(), singletonList(TASK0));
         expectPostRebalanceCatchup(SNAPSHOT_PAUSED_CONN1);
-        worker.startTask(EasyMock.eq(TASK0), EasyMock.>anyObject(), EasyMock.>anyObject(),
+        worker.startTask(EasyMock.eq(TASK0), EasyMock.anyObject(), EasyMock.>anyObject(), EasyMock.>anyObject(),
                 EasyMock.eq(herder), EasyMock.eq(TargetState.PAUSED));
         PowerMock.expectLastCall().andReturn(true);
         member.poll(EasyMock.anyInt());
@@ -1162,13 +1547,23 @@ public void testConnectorResumedRunningTaskOnly() throws Exception {
 
         // handle the state change
         member.wakeup();
+        PowerMock.expectLastCall();
         member.ensureActive();
         PowerMock.expectLastCall();
 
         EasyMock.expect(configBackingStore.snapshot()).andReturn(SNAPSHOT);
         PowerMock.expectLastCall();
 
-        worker.setTargetState(CONN1, TargetState.STARTED);
+        Capture> onStart = newCapture();
+        worker.setTargetState(EasyMock.eq(CONN1), EasyMock.eq(TargetState.STARTED), capture(onStart));
+        PowerMock.expectLastCall().andAnswer(new IAnswer() {
+            @Override
+            public Void answer() throws Throwable {
+                onStart.getValue().onCompletion(null, TargetState.STARTED);
+                return null;
+            }
+        });
+        member.wakeup();
         PowerMock.expectLastCall();
 
         EasyMock.expect(worker.isRunning(CONN1)).andReturn(false);
@@ -1176,11 +1571,18 @@ public void testConnectorResumedRunningTaskOnly() throws Exception {
         member.poll(EasyMock.anyInt());
         PowerMock.expectLastCall();
 
+        // These will occur just before/during the third tick
+        member.ensureActive();
+        PowerMock.expectLastCall();
+        member.poll(EasyMock.anyInt());
+        PowerMock.expectLastCall();
+
         PowerMock.replayAll();
 
         herder.tick(); // join
         configUpdateListener.onConnectorTargetStateChange(CONN1); // state changes to paused
         herder.tick(); // apply state change
+        herder.tick();
 
         PowerMock.verifyAll();
     }
@@ -1189,6 +1591,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());
@@ -1197,6 +1600,7 @@ public void testTaskConfigAdded() {
 
         // apply config
         member.wakeup();
+        PowerMock.expectLastCall();
         member.ensureActive();
         PowerMock.expectLastCall();
         // Checks for config updates and starts rebalance
@@ -1207,7 +1611,7 @@ public void testTaskConfigAdded() {
         expectRebalance(Collections.emptyList(), Collections.emptyList(),
                 ConnectProtocol.Assignment.NO_ERROR, 1, Collections.emptyList(),
                 Arrays.asList(TASK0));
-        worker.startTask(EasyMock.eq(TASK0), EasyMock.>anyObject(), EasyMock.>anyObject(),
+        worker.startTask(EasyMock.eq(TASK0), EasyMock.anyObject(), EasyMock.>anyObject(), EasyMock.>anyObject(),
                 EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
         PowerMock.expectLastCall().andReturn(true);
         member.poll(EasyMock.anyInt());
@@ -1227,43 +1631,263 @@ 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);
 
-        worker.startConnector(EasyMock.eq(CONN1), EasyMock.>anyObject(), EasyMock.anyObject(),
-                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
-        PowerMock.expectLastCall().andReturn(true);
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
+        Capture> onStart = newCapture();
+        worker.startConnector(EasyMock.eq(CONN1), EasyMock.>anyObject(), EasyMock.anyObject(),
+                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED), capture(onStart));
+        PowerMock.expectLastCall().andAnswer(new IAnswer() {
+            @Override
+            public Boolean answer() throws Throwable {
+                onStart.getValue().onCompletion(null, TargetState.STARTED);
+                return true;
+            }
+        });
+        member.wakeup();
+        PowerMock.expectLastCall();
         EasyMock.expect(worker.getPlugins()).andReturn(plugins);
         EasyMock.expect(worker.connectorTaskConfigs(CONN1, conn1SinkConfig)).andReturn(TASK_CONFIGS);
-        worker.startTask(EasyMock.eq(TASK1), EasyMock.>anyObject(), EasyMock.>anyObject(),
+        worker.startTask(EasyMock.eq(TASK1), EasyMock.anyObject(), EasyMock.>anyObject(), EasyMock.>anyObject(),
                 EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
         PowerMock.expectLastCall().andReturn(true);
         EasyMock.expect(worker.isRunning(CONN1)).andReturn(true);
         member.poll(EasyMock.anyInt());
         PowerMock.expectLastCall();
 
+        // one more tick, to make sure we don't keep trying to read to the config topic unnecessarily
+        expectRebalance(1, Collections.emptyList(), Collections.emptyList());
+        member.poll(EasyMock.anyInt());
+        PowerMock.expectLastCall();
+
         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);
 
+        // tick once more to ensure that the successful read to the end of the config topic was 
+        // tracked and no further unnecessary attempts were made
+        herder.tick();
+
+        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);
+        Capture> onStart = newCapture();
+        worker.startConnector(EasyMock.eq(CONN1), EasyMock.anyObject(), EasyMock.anyObject(),
+                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED), capture(onStart));
+        PowerMock.expectLastCall().andAnswer(new IAnswer() {
+            @Override
+            public Boolean answer() throws Throwable {
+                onStart.getValue().onCompletion(null, TargetState.STARTED);
+                return true;
+            }
+        });
+        member.wakeup();
+        PowerMock.expectLastCall();
+        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
+        Capture> onStart = newCapture();
+        worker.startConnector(EasyMock.eq(CONN1), EasyMock.anyObject(), EasyMock.anyObject(),
+                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED), capture(onStart));
+        PowerMock.expectLastCall().andAnswer(new IAnswer() {
+            @Override
+            public Boolean answer() throws Throwable {
+                onStart.getValue().onCompletion(null, TargetState.STARTED);
+                return true;
+            }
+        });
+        member.wakeup();
+        PowerMock.expectLastCall();
+        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();
     }
@@ -1271,9 +1895,20 @@ 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());
-        expectPostRebalanceCatchup(SNAPSHOT);
+        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, null, Collections.singletonMap(CONN1, 3),
+            Collections.singletonMap(CONN1, CONN1_CONFIG), Collections.singletonMap(CONN1, TargetState.STARTED),
+            TASK_CONFIGS_MAP, Collections.emptySet(), configTransformer);
+
+        expectPostRebalanceCatchup(snapshotWithTransform);
 
 
         member.wakeup();
@@ -1283,6 +1918,7 @@ public void testAccessors() throws Exception {
         PowerMock.expectLastCall();
 
 
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
         PowerMock.replayAll();
 
         FutureCallback> listConnectorsCb = new FutureCallback<>();
@@ -1318,9 +1954,17 @@ public void testPutConnectorConfig() throws Exception {
         EasyMock.expect(member.memberId()).andStubReturn("leader");
         expectRebalance(1, Arrays.asList(CONN1), Collections.emptyList());
         expectPostRebalanceCatchup(SNAPSHOT);
-        worker.startConnector(EasyMock.eq(CONN1), EasyMock.>anyObject(), EasyMock.anyObject(),
-                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
-        PowerMock.expectLastCall().andReturn(true);
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
+        Capture> onFirstStart = newCapture();
+        worker.startConnector(EasyMock.eq(CONN1), EasyMock.>anyObject(), EasyMock.anyObject(),
+                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED), capture(onFirstStart));
+        PowerMock.expectLastCall().andAnswer(new IAnswer() {
+            @Override
+            public Boolean answer() throws Throwable {
+                onFirstStart.getValue().onCompletion(null, TargetState.STARTED);
+                return true;
+            }
+        });
         EasyMock.expect(worker.isRunning(CONN1)).andReturn(true);
         EasyMock.expect(worker.connectorTaskConfigs(CONN1, conn1SinkConfig)).andReturn(TASK_CONFIGS);
 
@@ -1334,15 +1978,18 @@ public void testPutConnectorConfig() throws Exception {
         member.ensureActive();
         PowerMock.expectLastCall();
 
-        // config validation
-        Connector connectorMock = PowerMock.createMock(Connector.class);
         EasyMock.expect(worker.getPlugins()).andReturn(plugins).anyTimes();
-        EasyMock.expect(plugins.compareAndSwapLoaders(connectorMock)).andReturn(delegatingLoader);
-        EasyMock.expect(plugins.newConnector(EasyMock.anyString())).andReturn(connectorMock);
-        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);
 
+        Capture> validateCallback = newCapture();
+        herder.validateConnectorConfig(EasyMock.eq(CONN1_CONFIG_UPDATED), capture(validateCallback));
+        PowerMock.expectLastCall().andAnswer(new IAnswer() {
+            @Override
+            public Void answer() throws Throwable {
+                validateCallback.getValue().onCompletion(null, CONN1_CONFIG_INFOS);
+                return null;
+            }
+        });
         configBackingStore.putConnectorConfig(CONN1, CONN1_CONFIG_UPDATED);
         PowerMock.expectLastCall().andAnswer(new IAnswer() {
             @Override
@@ -1354,12 +2001,20 @@ 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);
-        worker.stopConnector(CONN1);
-        PowerMock.expectLastCall().andReturn(true);
-        worker.startConnector(EasyMock.eq(CONN1), EasyMock.>anyObject(), EasyMock.anyObject(),
-                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
-        PowerMock.expectLastCall().andReturn(true);
+        EasyMock.expect(configBackingStore.snapshot()).andReturn(SNAPSHOT_UPDATED_CONN1_CONFIG).times(2);
+        worker.stopAndAwaitConnector(CONN1);
+        PowerMock.expectLastCall();
+        EasyMock.expect(member.currentProtocolVersion()).andStubReturn(CONNECT_PROTOCOL_V0);
+        Capture> onSecondStart = newCapture();
+        worker.startConnector(EasyMock.eq(CONN1), EasyMock.>anyObject(), EasyMock.anyObject(),
+                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED), capture(onSecondStart));
+        PowerMock.expectLastCall().andAnswer(new IAnswer() {
+            @Override
+            public Boolean answer() throws Throwable {
+                onSecondStart.getValue().onCompletion(null, TargetState.STARTED);
+                return true;
+            }
+        });
         EasyMock.expect(worker.isRunning(CONN1)).andReturn(true);
         EasyMock.expect(worker.connectorTaskConfigs(CONN1, conn1SinkConfigUpdated)).andReturn(TASK_CONFIGS);
 
@@ -1396,19 +2051,159 @@ public Object answer() throws Throwable {
         herder.tick();
         assertTrue(connectorConfigCb.isDone());
         assertEquals(CONN1_CONFIG_UPDATED, connectorConfigCb.get());
+
+        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(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(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(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 testInconsistentConfigs() throws Exception {
+    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
     }
 
+
+    @Test
+    public void testThreadNames() {
+        assertTrue(Whitebox.getInternalState(herder, "herderExecutor").
+                getThreadFactory().newThread(EMPTY_RUNNABLE).getName().startsWith(DistributedHerder.class.getSimpleName()));
+
+        assertTrue(Whitebox.getInternalState(herder, "forwardRequestExecutor").
+                getThreadFactory().newThread(EMPTY_RUNNABLE).getName().startsWith("ForwardRequestExecutor"));
+
+        assertTrue(Whitebox.getInternalState(herder, "startAndStopExecutor").
+                getThreadFactory().newThread(EMPTY_RUNNABLE).getName().startsWith("StartAndStopExecutor"));
+    }
+
     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.
@@ -1418,33 +2213,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);
+                worker.stopAndAwaitConnector(connector);
+                PowerMock.expectLastCall();
             }
         }
 
-        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();
         }
@@ -1481,8 +2299,8 @@ private void assertStatistics(String expectedLeader, boolean isRebalancing, int
         assertEquals(isRebalancing ? 1.0d : 0.0d, rebalancing, 0.0001d);
         assertEquals(millisSinceLastRebalance, rebalanceTimeSinceLast, 0.0001d);
         if (rebalanceTime <= 0L) {
-            assertEquals(Double.NEGATIVE_INFINITY, rebalanceTimeMax, 0.0001d);
-            assertEquals(0.0d, rebalanceTimeAvg, 0.0001d);
+            assertEquals(Double.NaN, rebalanceTimeMax, 0.0001d);
+            assertEquals(Double.NaN, rebalanceTimeAvg, 0.0001d);
         } else {
             assertEquals(rebalanceTime, rebalanceTimeMax, 0.0001d);
             assertEquals(rebalanceTime, rebalanceTimeAvg, 0.0001d);
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..a44ac550d973f
--- /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.requests.RequestTestUtils;
+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.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(RequestTestUtils.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 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 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 to 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 154b1df5e73e3..2ac05c09d377b 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,44 +16,56 @@
  */
 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.Cluster;
 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.RequestTestUtils;
 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;
-import org.apache.kafka.test.TestUtils;
 import org.easymock.EasyMock;
 import org.easymock.Mock;
 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";
@@ -74,51 +86,70 @@ public class WorkerCoordinatorTest {
     private long retryBackoffMs = 100;
     private MockTime time;
     private MockClient client;
-    private Cluster cluster = TestUtils.singletonCluster("topic", 1);
-    private Node node = cluster.nodes().get(0);
+    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 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 loggerFactory = new LogContext();
+        LogContext logContext = new LogContext();
 
         this.time = new MockTime();
-        this.client = new MockClient(time);
-        this.metadata = new Metadata(0, Long.MAX_VALUE, true);
-        this.metadata.update(cluster, Collections.emptySet(), time.milliseconds());
-        this.consumerClient = new ConsumerNetworkClient(loggerFactory, client, metadata, time, 100, 1000, heartbeatIntervalMs);
+        this.metadata = new Metadata(0, Long.MAX_VALUE, logContext, new ClusterResourceListeners());
+        this.client = new MockClient(time, metadata);
+        this.client.updateMetadata(RequestTestUtils.metadataUpdateWith(1, Collections.singletonMap("topic", 1)));
+        this.node = metadata.fetch().nodes().get(0);
+        this.consumerClient = new ConsumerNetworkClient(logContext, client, metadata, time, 100, 1000, heartbeatIntervalMs);
         this.metrics = new Metrics(time);
         this.rebalanceListener = new MockRebalanceListener();
         this.configStorage = PowerMock.createMock(KafkaConfigBackingStore.class);
-
-        client.setNode(node);
-
-        this.coordinator = new WorkerCoordinator(
-                loggerFactory,
-                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),
@@ -141,6 +172,7 @@ public void setup() {
         configState2TaskConfigs.put(taskId2x0, new HashMap());
         configState2 = new ClusterConfigState(
                 2L,
+                null,
                 configState2ConnectorTaskCounts,
                 configState2ConnectorConfigs,
                 configState2TargetStates,
@@ -166,6 +198,7 @@ public void setup() {
         configStateSingleTaskConnectorsTaskConfigs.put(taskId3x0, new HashMap());
         configStateSingleTaskConnectors = new ClusterConfigState(
                 2L,
+                null,
                 configStateSingleTaskConnectorsConnectorTaskCounts,
                 configStateSingleTaskConnectorsConnectorConfigs,
                 configStateSingleTaskConnectorsTargetStates,
@@ -188,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();
@@ -207,8 +243,8 @@ public void testNormalJoinGroupLeader() {
 
         final String consumerId = "leader";
 
-        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        client.prepareResponse(FindCoordinatorResponse.prepareResponse(Errors.NONE, node));
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         // normal join group
         Map memberConfigOffsets = new HashMap<>();
@@ -219,15 +255,15 @@ 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));
         coordinator.ensureActiveGroup();
 
-        assertFalse(coordinator.needRejoin());
+        assertFalse(coordinator.rejoinNeededOrPending());
         assertEquals(0, rebalanceListener.revokedCount);
         assertEquals(1, rebalanceListener.assignedCount);
         assertFalse(rebalanceListener.assignment.failed());
@@ -247,8 +283,8 @@ public void testNormalJoinGroupFollower() {
 
         final String memberId = "member";
 
-        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        client.prepareResponse(FindCoordinatorResponse.prepareResponse(Errors.NONE, node));
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         // normal join group
         client.prepareResponse(joinGroupFollowerResponse(1, memberId, "leader", Errors.NONE));
@@ -256,15 +292,15 @@ 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));
         coordinator.ensureActiveGroup();
 
-        assertFalse(coordinator.needRejoin());
+        assertFalse(coordinator.rejoinNeededOrPending());
         assertEquals(0, rebalanceListener.revokedCount);
         assertEquals(1, rebalanceListener.assignedCount);
         assertFalse(rebalanceListener.assignment.failed());
@@ -288,8 +324,8 @@ public void testJoinLeaderCannotAssign() {
 
         final String memberId = "member";
 
-        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        client.prepareResponse(FindCoordinatorResponse.prepareResponse(Errors.NONE, node));
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         // config mismatch results in assignment error
         client.prepareResponse(joinGroupFollowerResponse(1, memberId, "leader", Errors.NONE));
@@ -297,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,
@@ -319,8 +355,8 @@ public void testRejoinGroup() {
 
         PowerMock.replayAll();
 
-        client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE));
-        coordinator.ensureCoordinatorReady();
+        client.prepareResponse(FindCoordinatorResponse.prepareResponse(Errors.NONE, node));
+        coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
 
         // join the group once
         client.prepareResponse(joinGroupFollowerResponse(1, "consumer", "leader", Errors.NONE));
@@ -366,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"));
@@ -402,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"));
@@ -438,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
@@ -463,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;
@@ -504,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/distributed/WorkerGroupMemberTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/distributed/WorkerGroupMemberTest.java
new file mode 100644
index 0000000000000..05cd01734fef8
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/distributed/WorkerGroupMemberTest.java
@@ -0,0 +1,103 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.CommonClientConfigs;
+import org.apache.kafka.common.MetricName;
+import org.apache.kafka.common.metrics.MetricsReporter;
+import org.apache.kafka.common.metrics.stats.Avg;
+import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.connect.runtime.MockConnectMetrics;
+import org.apache.kafka.connect.runtime.WorkerConfig;
+import org.apache.kafka.connect.storage.ConfigBackingStore;
+import org.apache.kafka.connect.storage.StatusBackingStore;
+import org.apache.kafka.connect.util.ConnectUtils;
+import org.easymock.EasyMock;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.powermock.api.easymock.PowerMock;
+import org.powermock.api.easymock.annotation.Mock;
+import org.powermock.core.classloader.annotations.PowerMockIgnore;
+import org.powermock.core.classloader.annotations.PrepareForTest;
+import org.powermock.modules.junit4.PowerMockRunner;
+
+import javax.management.MBeanServer;
+import javax.management.ObjectName;
+import java.lang.management.ManagementFactory;
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+
+@RunWith(PowerMockRunner.class)
+@PrepareForTest({ConnectUtils.class})
+@PowerMockIgnore({"javax.management.*", "javax.crypto.*"})
+public class WorkerGroupMemberTest {
+    @Mock
+    private ConfigBackingStore configBackingStore;
+    @Mock
+    private StatusBackingStore statusBackingStore;
+
+    @Test
+    public void testMetrics() throws Exception {
+        WorkerGroupMember member;
+        Map workerProps = new HashMap<>();
+        workerProps.put("key.converter", "org.apache.kafka.connect.json.JsonConverter");
+        workerProps.put("value.converter", "org.apache.kafka.connect.json.JsonConverter");
+        workerProps.put("offset.storage.file.filename", "/tmp/connect.offsets");
+        workerProps.put("group.id", "group-1");
+        workerProps.put("offset.storage.topic", "topic-1");
+        workerProps.put("config.storage.topic", "topic-1");
+        workerProps.put("status.storage.topic", "topic-1");
+        workerProps.put(CommonClientConfigs.METRIC_REPORTER_CLASSES_CONFIG, MockConnectMetrics.MockMetricsReporter.class.getName());
+        DistributedConfig config = new DistributedConfig(workerProps);
+
+
+        LogContext logContext = new LogContext("[Worker clientId=client-1 + groupId= group-1]");
+
+        expectClusterId();
+
+        member = new WorkerGroupMember(config, "", configBackingStore,
+        null, Time.SYSTEM, "client-1", logContext);
+
+        boolean entered = false;
+        for (MetricsReporter reporter : member.metrics().reporters()) {
+            if (reporter instanceof MockConnectMetrics.MockMetricsReporter) {
+                entered = true;
+                MockConnectMetrics.MockMetricsReporter mockMetricsReporter = (MockConnectMetrics.MockMetricsReporter) reporter;
+                assertEquals("cluster-1", mockMetricsReporter.getMetricsContext().contextLabels().get(WorkerConfig.CONNECT_KAFKA_CLUSTER_ID));
+                assertEquals("group-1", mockMetricsReporter.getMetricsContext().contextLabels().get(WorkerConfig.CONNECT_GROUP_ID));
+            }
+        }
+        assertTrue("Failed to verify MetricsReporter", entered);
+
+        MetricName name = member.metrics().metricName("test.avg", "grp1");
+        member.metrics().addMetric(name, new Avg());
+        MBeanServer server = ManagementFactory.getPlatformMBeanServer();
+        //verify metric exists with correct prefix
+        assertNotNull(server.getObjectInstance(new ObjectName("kafka.connect:type=grp1,client-id=client-1")));
+    }
+    private void expectClusterId() {
+        PowerMock.mockStaticPartial(ConnectUtils.class, "lookupKafkaClusterId");
+        EasyMock.expect(ConnectUtils.lookupKafkaClusterId(EasyMock.anyObject())).andReturn("cluster-1").anyTimes();
+        PowerMock.replay(ConnectUtils.class);
+    }
+
+}
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
new file mode 100644
index 0000000000000..85f69486d29a8
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/errors/ErrorReporterTest.java
@@ -0,0 +1,366 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.errors;
+
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+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.header.Header;
+import org.apache.kafka.connect.errors.ConnectException;
+import org.apache.kafka.connect.json.JsonConverter;
+import org.apache.kafka.connect.runtime.ConnectMetrics;
+import org.apache.kafka.connect.runtime.ConnectorConfig;
+import org.apache.kafka.connect.runtime.MockConnectMetrics;
+import org.apache.kafka.connect.runtime.SinkConnectorConfig;
+import org.apache.kafka.connect.runtime.isolation.Plugins;
+import org.apache.kafka.connect.sink.SinkTask;
+import org.apache.kafka.connect.transforms.Transformation;
+import org.apache.kafka.connect.util.ConnectorTaskId;
+import org.easymock.EasyMock;
+import org.easymock.Mock;
+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.core.classloader.annotations.PowerMockIgnore;
+import org.powermock.modules.junit4.PowerMockRunner;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Future;
+
+import static java.util.Collections.emptyMap;
+import static java.util.Collections.singletonMap;
+import static org.apache.kafka.connect.runtime.errors.DeadLetterQueueReporter.ERROR_HEADER_CONNECTOR_NAME;
+import static org.apache.kafka.connect.runtime.errors.DeadLetterQueueReporter.ERROR_HEADER_EXCEPTION;
+import static org.apache.kafka.connect.runtime.errors.DeadLetterQueueReporter.ERROR_HEADER_EXCEPTION_MESSAGE;
+import static org.apache.kafka.connect.runtime.errors.DeadLetterQueueReporter.ERROR_HEADER_EXCEPTION_STACK_TRACE;
+import static org.apache.kafka.connect.runtime.errors.DeadLetterQueueReporter.ERROR_HEADER_EXECUTING_CLASS;
+import static org.apache.kafka.connect.runtime.errors.DeadLetterQueueReporter.ERROR_HEADER_ORIG_OFFSET;
+import static org.apache.kafka.connect.runtime.errors.DeadLetterQueueReporter.ERROR_HEADER_ORIG_PARTITION;
+import static org.apache.kafka.connect.runtime.errors.DeadLetterQueueReporter.ERROR_HEADER_ORIG_TOPIC;
+import static org.apache.kafka.connect.runtime.errors.DeadLetterQueueReporter.ERROR_HEADER_STAGE;
+import static org.apache.kafka.connect.runtime.errors.DeadLetterQueueReporter.ERROR_HEADER_TASK_ID;
+import static org.easymock.EasyMock.replay;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+
+@RunWith(PowerMockRunner.class)
+@PowerMockIgnore("javax.management.*")
+public class ErrorReporterTest {
+
+    private static final String TOPIC = "test-topic";
+    private static final String DLQ_TOPIC = "test-topic-errors";
+    private static final ConnectorTaskId TASK_ID = new ConnectorTaskId("job", 0);
+
+    @Mock
+    KafkaProducer producer;
+
+    @Mock
+    Future metadata;
+
+    @Mock
+    Plugins plugins;
+
+    private ErrorHandlingMetrics errorHandlingMetrics;
+    private MockConnectMetrics metrics;
+
+    @Before
+    public void setup() {
+        metrics = new MockConnectMetrics();
+        errorHandlingMetrics = new ErrorHandlingMetrics(new ConnectorTaskId("connector-", 1), metrics);
+    }
+
+    @After
+    public void tearDown() {
+        if (metrics != null) {
+            metrics.stop();
+        }
+    }
+
+    @Test(expected = NullPointerException.class)
+    public void initializeDLQWithNullMetrics() {
+        new DeadLetterQueueReporter(producer, config(emptyMap()), TASK_ID, null);
+    }
+
+    @Test
+    public void testDLQConfigWithEmptyTopicName() {
+        DeadLetterQueueReporter deadLetterQueueReporter = new DeadLetterQueueReporter(
+                producer, config(emptyMap()), TASK_ID, errorHandlingMetrics);
+
+        ProcessingContext context = processingContext();
+
+        EasyMock.expect(producer.send(EasyMock.anyObject(), EasyMock.anyObject())).andThrow(new RuntimeException());
+        replay(producer);
+
+        // since topic name is empty, this method should be a NOOP.
+        // if it attempts to log to the DLQ via the producer, the send mock will throw a RuntimeException.
+        deadLetterQueueReporter.report(context);
+    }
+
+    @Test
+    public void testDLQConfigWithValidTopicName() {
+        DeadLetterQueueReporter deadLetterQueueReporter = new DeadLetterQueueReporter(
+                producer, config(singletonMap(SinkConnectorConfig.DLQ_TOPIC_NAME_CONFIG, DLQ_TOPIC)), TASK_ID, errorHandlingMetrics);
+
+        ProcessingContext context = processingContext();
+
+        EasyMock.expect(producer.send(EasyMock.anyObject(), EasyMock.anyObject())).andReturn(metadata);
+        replay(producer);
+
+        deadLetterQueueReporter.report(context);
+
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testReportDLQTwice() {
+        DeadLetterQueueReporter deadLetterQueueReporter = new DeadLetterQueueReporter(
+                producer, config(singletonMap(SinkConnectorConfig.DLQ_TOPIC_NAME_CONFIG, DLQ_TOPIC)), TASK_ID, errorHandlingMetrics);
+
+        ProcessingContext context = processingContext();
+
+        EasyMock.expect(producer.send(EasyMock.anyObject(), EasyMock.anyObject())).andReturn(metadata).times(2);
+        replay(producer);
+
+        deadLetterQueueReporter.report(context);
+        deadLetterQueueReporter.report(context);
+
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testDLQReportAndReturnFuture() {
+        DeadLetterQueueReporter deadLetterQueueReporter = new DeadLetterQueueReporter(
+            producer, config(singletonMap(SinkConnectorConfig.DLQ_TOPIC_NAME_CONFIG, DLQ_TOPIC)), TASK_ID, errorHandlingMetrics);
+
+        ProcessingContext context = processingContext();
+
+        EasyMock.expect(producer.send(EasyMock.anyObject(), EasyMock.anyObject())).andReturn(metadata);
+        replay(producer);
+
+        deadLetterQueueReporter.report(context);
+
+        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);
+
+        ProcessingContext context = processingContext();
+        context.error(new RuntimeException());
+
+        // reporting a context without an error should not cause any errors.
+        logReporter.report(context);
+        assertErrorHandlingMetricValue("total-errors-logged", 0.0);
+    }
+
+    @Test
+    public void testLogOnEnabledLogReporter() {
+        LogReporter logReporter = new LogReporter(TASK_ID, config(singletonMap(ConnectorConfig.ERRORS_LOG_ENABLE_CONFIG, "true")), errorHandlingMetrics);
+
+        ProcessingContext context = processingContext();
+        context.error(new RuntimeException());
+
+        // reporting a context without an error should not cause any errors.
+        logReporter.report(context);
+        assertErrorHandlingMetricValue("total-errors-logged", 1.0);
+    }
+
+    @Test
+    public void testLogMessageWithNoRecords() {
+        LogReporter logReporter = new LogReporter(TASK_ID, config(singletonMap(ConnectorConfig.ERRORS_LOG_ENABLE_CONFIG, "true")), errorHandlingMetrics);
+
+        ProcessingContext context = processingContext();
+
+        String msg = logReporter.message(context);
+        assertEquals("Error encountered in task job-0. Executing stage 'KEY_CONVERTER' with class " +
+                "'org.apache.kafka.connect.json.JsonConverter'.", msg);
+    }
+
+    @Test
+    public void testLogMessageWithSinkRecords() {
+        Map props = new HashMap<>();
+        props.put(ConnectorConfig.ERRORS_LOG_ENABLE_CONFIG, "true");
+        props.put(ConnectorConfig.ERRORS_LOG_INCLUDE_MESSAGES_CONFIG, "true");
+
+        LogReporter logReporter = new LogReporter(TASK_ID, config(props), errorHandlingMetrics);
+
+        ProcessingContext context = processingContext();
+
+        String msg = logReporter.message(context);
+        assertEquals("Error encountered in task job-0. Executing stage 'KEY_CONVERTER' with class " +
+                "'org.apache.kafka.connect.json.JsonConverter', where consumed record is {topic='test-topic', " +
+                "partition=5, offset=100}.", msg);
+    }
+
+    @Test
+    public void testLogReportAndReturnFuture() {
+        Map props = new HashMap<>();
+        props.put(ConnectorConfig.ERRORS_LOG_ENABLE_CONFIG, "true");
+        props.put(ConnectorConfig.ERRORS_LOG_INCLUDE_MESSAGES_CONFIG, "true");
+
+        LogReporter logReporter = new LogReporter(TASK_ID, config(props), errorHandlingMetrics);
+
+        ProcessingContext context = processingContext();
+
+        String msg = logReporter.message(context);
+        assertEquals("Error encountered in task job-0. Executing stage 'KEY_CONVERTER' with class " +
+            "'org.apache.kafka.connect.json.JsonConverter', where consumed record is {topic='test-topic', " +
+            "partition=5, offset=100}.", msg);
+
+        Future future = logReporter.report(context);
+        assertTrue(future instanceof CompletableFuture);
+    }
+
+    @Test
+    public void testSetDLQConfigs() {
+        SinkConnectorConfig configuration = config(singletonMap(SinkConnectorConfig.DLQ_TOPIC_NAME_CONFIG, DLQ_TOPIC));
+        assertEquals(configuration.dlqTopicName(), DLQ_TOPIC);
+
+        configuration = config(singletonMap(SinkConnectorConfig.DLQ_TOPIC_REPLICATION_FACTOR_CONFIG, "7"));
+        assertEquals(configuration.dlqTopicReplicationFactor(), 7);
+    }
+
+    @Test
+    public void testDlqHeaderConsumerRecord() {
+        Map props = new HashMap<>();
+        props.put(SinkConnectorConfig.DLQ_TOPIC_NAME_CONFIG, DLQ_TOPIC);
+        props.put(SinkConnectorConfig.DLQ_CONTEXT_HEADERS_ENABLE_CONFIG, "true");
+        DeadLetterQueueReporter deadLetterQueueReporter = new DeadLetterQueueReporter(producer, config(props), TASK_ID, errorHandlingMetrics);
+
+        ProcessingContext context = new ProcessingContext();
+        context.consumerRecord(new ConsumerRecord<>("source-topic", 7, 10, "source-key".getBytes(), "source-value".getBytes()));
+        context.currentContext(Stage.TRANSFORMATION, Transformation.class);
+        context.error(new ConnectException("Test Exception"));
+
+        ProducerRecord producerRecord = new ProducerRecord<>(DLQ_TOPIC, "source-key".getBytes(), "source-value".getBytes());
+
+        deadLetterQueueReporter.populateContextHeaders(producerRecord, context);
+        assertEquals("source-topic", headerValue(producerRecord, ERROR_HEADER_ORIG_TOPIC));
+        assertEquals("7", headerValue(producerRecord, ERROR_HEADER_ORIG_PARTITION));
+        assertEquals("10", headerValue(producerRecord, ERROR_HEADER_ORIG_OFFSET));
+        assertEquals(TASK_ID.connector(), headerValue(producerRecord, ERROR_HEADER_CONNECTOR_NAME));
+        assertEquals(String.valueOf(TASK_ID.task()), headerValue(producerRecord, ERROR_HEADER_TASK_ID));
+        assertEquals(Stage.TRANSFORMATION.name(), headerValue(producerRecord, ERROR_HEADER_STAGE));
+        assertEquals(Transformation.class.getName(), headerValue(producerRecord, ERROR_HEADER_EXECUTING_CLASS));
+        assertEquals(ConnectException.class.getName(), headerValue(producerRecord, ERROR_HEADER_EXCEPTION));
+        assertEquals("Test Exception", headerValue(producerRecord, ERROR_HEADER_EXCEPTION_MESSAGE));
+        assertTrue(headerValue(producerRecord, ERROR_HEADER_EXCEPTION_STACK_TRACE).length() > 0);
+        assertTrue(headerValue(producerRecord, ERROR_HEADER_EXCEPTION_STACK_TRACE).startsWith("org.apache.kafka.connect.errors.ConnectException: Test Exception"));
+    }
+
+    @Test
+    public void testDlqHeaderOnNullExceptionMessage() {
+        Map props = new HashMap<>();
+        props.put(SinkConnectorConfig.DLQ_TOPIC_NAME_CONFIG, DLQ_TOPIC);
+        props.put(SinkConnectorConfig.DLQ_CONTEXT_HEADERS_ENABLE_CONFIG, "true");
+        DeadLetterQueueReporter deadLetterQueueReporter = new DeadLetterQueueReporter(producer, config(props), TASK_ID, errorHandlingMetrics);
+
+        ProcessingContext context = new ProcessingContext();
+        context.consumerRecord(new ConsumerRecord<>("source-topic", 7, 10, "source-key".getBytes(), "source-value".getBytes()));
+        context.currentContext(Stage.TRANSFORMATION, Transformation.class);
+        context.error(new NullPointerException());
+
+        ProducerRecord producerRecord = new ProducerRecord<>(DLQ_TOPIC, "source-key".getBytes(), "source-value".getBytes());
+
+        deadLetterQueueReporter.populateContextHeaders(producerRecord, context);
+        assertEquals("source-topic", headerValue(producerRecord, ERROR_HEADER_ORIG_TOPIC));
+        assertEquals("7", headerValue(producerRecord, ERROR_HEADER_ORIG_PARTITION));
+        assertEquals("10", headerValue(producerRecord, ERROR_HEADER_ORIG_OFFSET));
+        assertEquals(TASK_ID.connector(), headerValue(producerRecord, ERROR_HEADER_CONNECTOR_NAME));
+        assertEquals(String.valueOf(TASK_ID.task()), headerValue(producerRecord, ERROR_HEADER_TASK_ID));
+        assertEquals(Stage.TRANSFORMATION.name(), headerValue(producerRecord, ERROR_HEADER_STAGE));
+        assertEquals(Transformation.class.getName(), headerValue(producerRecord, ERROR_HEADER_EXECUTING_CLASS));
+        assertEquals(NullPointerException.class.getName(), headerValue(producerRecord, ERROR_HEADER_EXCEPTION));
+        assertNull(producerRecord.headers().lastHeader(ERROR_HEADER_EXCEPTION_MESSAGE).value());
+        assertTrue(headerValue(producerRecord, ERROR_HEADER_EXCEPTION_STACK_TRACE).length() > 0);
+        assertTrue(headerValue(producerRecord, ERROR_HEADER_EXCEPTION_STACK_TRACE).startsWith("java.lang.NullPointerException"));
+    }
+
+    @Test
+    public void testDlqHeaderIsAppended() {
+        Map props = new HashMap<>();
+        props.put(SinkConnectorConfig.DLQ_TOPIC_NAME_CONFIG, DLQ_TOPIC);
+        props.put(SinkConnectorConfig.DLQ_CONTEXT_HEADERS_ENABLE_CONFIG, "true");
+        DeadLetterQueueReporter deadLetterQueueReporter = new DeadLetterQueueReporter(producer, config(props), TASK_ID, errorHandlingMetrics);
+
+        ProcessingContext context = new ProcessingContext();
+        context.consumerRecord(new ConsumerRecord<>("source-topic", 7, 10, "source-key".getBytes(), "source-value".getBytes()));
+        context.currentContext(Stage.TRANSFORMATION, Transformation.class);
+        context.error(new ConnectException("Test Exception"));
+
+        ProducerRecord producerRecord = new ProducerRecord<>(DLQ_TOPIC, "source-key".getBytes(), "source-value".getBytes());
+        producerRecord.headers().add(ERROR_HEADER_ORIG_TOPIC, "dummy".getBytes());
+
+        deadLetterQueueReporter.populateContextHeaders(producerRecord, context);
+        int appearances = 0;
+        for (Header header: producerRecord.headers()) {
+            if (ERROR_HEADER_ORIG_TOPIC.equalsIgnoreCase(header.key())) {
+                appearances++;
+            }
+        }
+
+        assertEquals("source-topic", headerValue(producerRecord, ERROR_HEADER_ORIG_TOPIC));
+        assertEquals(2, appearances);
+    }
+
+    private String headerValue(ProducerRecord producerRecord, String headerSuffix) {
+        return new String(producerRecord.headers().lastHeader(headerSuffix).value());
+    }
+
+    private ProcessingContext processingContext() {
+        ProcessingContext context = new ProcessingContext();
+        context.consumerRecord(new ConsumerRecord<>(TOPIC, 5, 100, new byte[]{'a', 'b'}, new byte[]{'x'}));
+        context.currentContext(Stage.KEY_CONVERTER, JsonConverter.class);
+        return context;
+    }
+
+    private SinkConnectorConfig config(Map configProps) {
+        Map props = new HashMap<>();
+        props.put(ConnectorConfig.NAME_CONFIG, "test");
+        props.put(ConnectorConfig.CONNECTOR_CLASS_CONFIG, SinkTask.class.getName());
+        props.putAll(configProps);
+        return new SinkConnectorConfig(plugins, props);
+    }
+
+    private void assertErrorHandlingMetricValue(String name, double expected) {
+        ConnectMetrics.MetricGroup sinkTaskGroup = errorHandlingMetrics.metricGroup();
+        double measured = metrics.currentMetricValueAsDouble(sinkTaskGroup, name);
+        assertEquals(expected, measured, 0.001d);
+    }
+
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/errors/RetryWithToleranceOperatorTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/errors/RetryWithToleranceOperatorTest.java
new file mode 100644
index 0000000000000..9e44321567790
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/errors/RetryWithToleranceOperatorTest.java
@@ -0,0 +1,456 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.errors;
+
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.common.metrics.Sensor;
+import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.common.utils.SystemTime;
+import org.apache.kafka.connect.errors.ConnectException;
+import org.apache.kafka.connect.errors.RetriableException;
+import org.apache.kafka.connect.runtime.ConnectMetrics;
+import org.apache.kafka.connect.runtime.ConnectorConfig;
+import org.apache.kafka.connect.runtime.WorkerConfig;
+import org.apache.kafka.connect.runtime.isolation.Plugins;
+import org.apache.kafka.connect.runtime.isolation.PluginsTest.TestConverter;
+import org.apache.kafka.connect.runtime.isolation.PluginsTest.TestableWorkerConfig;
+import org.apache.kafka.connect.sink.SinkTask;
+import org.apache.kafka.connect.util.ConnectorTaskId;
+import org.easymock.EasyMock;
+import org.easymock.Mock;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.powermock.api.easymock.PowerMock;
+import org.powermock.core.classloader.annotations.PowerMockIgnore;
+import org.powermock.core.classloader.annotations.PrepareForTest;
+import org.powermock.modules.junit4.PowerMockRunner;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+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.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import static java.util.Collections.emptyMap;
+import static java.util.Collections.singletonMap;
+import static org.apache.kafka.common.utils.Time.SYSTEM;
+import static org.apache.kafka.connect.runtime.ConnectorConfig.ERRORS_RETRY_MAX_DELAY_CONFIG;
+import static org.apache.kafka.connect.runtime.ConnectorConfig.ERRORS_RETRY_MAX_DELAY_DEFAULT;
+import static org.apache.kafka.connect.runtime.ConnectorConfig.ERRORS_RETRY_TIMEOUT_CONFIG;
+import static org.apache.kafka.connect.runtime.ConnectorConfig.ERRORS_RETRY_TIMEOUT_DEFAULT;
+import static org.apache.kafka.connect.runtime.ConnectorConfig.ERRORS_TOLERANCE_CONFIG;
+import static org.apache.kafka.connect.runtime.ConnectorConfig.ERRORS_TOLERANCE_DEFAULT;
+import static org.apache.kafka.connect.runtime.errors.ToleranceType.ALL;
+import static org.apache.kafka.connect.runtime.errors.ToleranceType.NONE;
+import static org.easymock.EasyMock.replay;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+
+@RunWith(PowerMockRunner.class)
+@PrepareForTest({ProcessingContext.class})
+@PowerMockIgnore("javax.management.*")
+public class RetryWithToleranceOperatorTest {
+
+    public static final RetryWithToleranceOperator NOOP_OPERATOR = new RetryWithToleranceOperator(
+            ERRORS_RETRY_TIMEOUT_DEFAULT, ERRORS_RETRY_MAX_DELAY_DEFAULT, NONE, SYSTEM);
+    static {
+        Map properties = new HashMap<>();
+        properties.put(CommonClientConfigs.METRICS_NUM_SAMPLES_CONFIG, Objects.toString(2));
+        properties.put(CommonClientConfigs.METRICS_SAMPLE_WINDOW_MS_CONFIG, Objects.toString(3000));
+        properties.put(CommonClientConfigs.METRICS_RECORDING_LEVEL_CONFIG, Sensor.RecordingLevel.INFO.toString());
+
+        // define required properties
+        properties.put(WorkerConfig.KEY_CONVERTER_CLASS_CONFIG, TestConverter.class.getName());
+        properties.put(WorkerConfig.VALUE_CONVERTER_CLASS_CONFIG, TestConverter.class.getName());
+
+        NOOP_OPERATOR.metrics(new ErrorHandlingMetrics(
+            new ConnectorTaskId("noop-connector", -1),
+            new ConnectMetrics("noop-worker", new TestableWorkerConfig(properties), new SystemTime(), "test-cluster"))
+        );
+    }
+
+    @SuppressWarnings("unused")
+    @Mock
+    private Operation mockOperation;
+
+    @Mock
+    private ConsumerRecord consumerRecord;
+
+    @Mock
+    ErrorHandlingMetrics errorHandlingMetrics;
+
+    @Mock
+    Plugins plugins;
+
+    @Test
+    public void testExecuteFailed() {
+        RetryWithToleranceOperator retryWithToleranceOperator = new RetryWithToleranceOperator(0,
+            ERRORS_RETRY_MAX_DELAY_DEFAULT, ALL, SYSTEM);
+        retryWithToleranceOperator.metrics(errorHandlingMetrics);
+
+        retryWithToleranceOperator.executeFailed(Stage.TASK_PUT,
+            SinkTask.class, consumerRecord, new Throwable());
+    }
+
+    @Test(expected = ConnectException.class)
+    public void testExecuteFailedNoTolerance() {
+        RetryWithToleranceOperator retryWithToleranceOperator = new RetryWithToleranceOperator(0,
+            ERRORS_RETRY_MAX_DELAY_DEFAULT, NONE, SYSTEM);
+        retryWithToleranceOperator.metrics(errorHandlingMetrics);
+
+        retryWithToleranceOperator.executeFailed(Stage.TASK_PUT,
+            SinkTask.class, consumerRecord, new Throwable());
+    }
+
+    @Test
+    public void testHandleExceptionInTransformations() {
+        testHandleExceptionInStage(Stage.TRANSFORMATION, new Exception());
+    }
+
+    @Test
+    public void testHandleExceptionInHeaderConverter() {
+        testHandleExceptionInStage(Stage.HEADER_CONVERTER, new Exception());
+    }
+
+    @Test
+    public void testHandleExceptionInValueConverter() {
+        testHandleExceptionInStage(Stage.VALUE_CONVERTER, new Exception());
+    }
+
+    @Test
+    public void testHandleExceptionInKeyConverter() {
+        testHandleExceptionInStage(Stage.KEY_CONVERTER, new Exception());
+    }
+
+    @Test
+    public void testHandleExceptionInTaskPut() {
+        testHandleExceptionInStage(Stage.TASK_PUT, new org.apache.kafka.connect.errors.RetriableException("Test"));
+    }
+
+    @Test
+    public void testHandleExceptionInTaskPoll() {
+        testHandleExceptionInStage(Stage.TASK_POLL, new org.apache.kafka.connect.errors.RetriableException("Test"));
+    }
+
+    @Test(expected = ConnectException.class)
+    public void testThrowExceptionInTaskPut() {
+        testHandleExceptionInStage(Stage.TASK_PUT, new Exception());
+    }
+
+    @Test(expected = ConnectException.class)
+    public void testThrowExceptionInTaskPoll() {
+        testHandleExceptionInStage(Stage.TASK_POLL, new Exception());
+    }
+
+    @Test(expected = ConnectException.class)
+    public void testThrowExceptionInKafkaConsume() {
+        testHandleExceptionInStage(Stage.KAFKA_CONSUME, new Exception());
+    }
+
+    @Test(expected = ConnectException.class)
+    public void testThrowExceptionInKafkaProduce() {
+        testHandleExceptionInStage(Stage.KAFKA_PRODUCE, new Exception());
+    }
+
+    private void testHandleExceptionInStage(Stage type, Exception ex) {
+        RetryWithToleranceOperator retryWithToleranceOperator = setupExecutor();
+        retryWithToleranceOperator.execute(new ExceptionThrower(ex), type, ExceptionThrower.class);
+        assertTrue(retryWithToleranceOperator.failed());
+        PowerMock.verifyAll();
+    }
+
+    private RetryWithToleranceOperator setupExecutor() {
+        RetryWithToleranceOperator retryWithToleranceOperator = new RetryWithToleranceOperator(0, ERRORS_RETRY_MAX_DELAY_DEFAULT, ALL, SYSTEM);
+        retryWithToleranceOperator.metrics(errorHandlingMetrics);
+        return retryWithToleranceOperator;
+    }
+
+    @Test
+    public void testExecAndHandleRetriableErrorOnce() throws Exception {
+        execAndHandleRetriableError(1, 300, new RetriableException("Test"));
+    }
+
+    @Test
+    public void testExecAndHandleRetriableErrorThrice() throws Exception {
+        execAndHandleRetriableError(3, 2100, new RetriableException("Test"));
+    }
+
+    @Test
+    public void testExecAndHandleNonRetriableErrorOnce() throws Exception {
+        execAndHandleNonRetriableError(1, 0, new Exception("Non Retriable Test"));
+    }
+
+    @Test
+    public void testExecAndHandleNonRetriableErrorThrice() throws Exception {
+        execAndHandleNonRetriableError(3, 0, new Exception("Non Retriable Test"));
+    }
+
+    public void execAndHandleRetriableError(int numRetriableExceptionsThrown, long expectedWait, Exception e) throws Exception {
+        MockTime time = new MockTime(0, 0, 0);
+        RetryWithToleranceOperator retryWithToleranceOperator = new RetryWithToleranceOperator(6000, ERRORS_RETRY_MAX_DELAY_DEFAULT, ALL, time);
+        retryWithToleranceOperator.metrics(errorHandlingMetrics);
+
+        EasyMock.expect(mockOperation.call()).andThrow(e).times(numRetriableExceptionsThrown);
+        EasyMock.expect(mockOperation.call()).andReturn("Success");
+
+        replay(mockOperation);
+
+        String result = retryWithToleranceOperator.execAndHandleError(mockOperation, Exception.class);
+        assertFalse(retryWithToleranceOperator.failed());
+        assertEquals("Success", result);
+        assertEquals(expectedWait, time.hiResClockMs());
+
+        PowerMock.verifyAll();
+    }
+
+    public void execAndHandleNonRetriableError(int numRetriableExceptionsThrown, long expectedWait, Exception e) throws Exception {
+        MockTime time = new MockTime(0, 0, 0);
+        RetryWithToleranceOperator retryWithToleranceOperator = new RetryWithToleranceOperator(6000, ERRORS_RETRY_MAX_DELAY_DEFAULT, ALL, time);
+        retryWithToleranceOperator.metrics(errorHandlingMetrics);
+
+        EasyMock.expect(mockOperation.call()).andThrow(e).times(numRetriableExceptionsThrown);
+        EasyMock.expect(mockOperation.call()).andReturn("Success");
+
+        replay(mockOperation);
+
+        String result = retryWithToleranceOperator.execAndHandleError(mockOperation, Exception.class);
+        assertTrue(retryWithToleranceOperator.failed());
+        assertNull(result);
+        assertEquals(expectedWait, time.hiResClockMs());
+
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testCheckRetryLimit() {
+        MockTime time = new MockTime(0, 0, 0);
+        RetryWithToleranceOperator retryWithToleranceOperator = new RetryWithToleranceOperator(500, 100, NONE, time);
+
+        time.setCurrentTimeMs(100);
+        assertTrue(retryWithToleranceOperator.checkRetry(0));
+
+        time.setCurrentTimeMs(200);
+        assertTrue(retryWithToleranceOperator.checkRetry(0));
+
+        time.setCurrentTimeMs(400);
+        assertTrue(retryWithToleranceOperator.checkRetry(0));
+
+        time.setCurrentTimeMs(499);
+        assertTrue(retryWithToleranceOperator.checkRetry(0));
+
+        time.setCurrentTimeMs(501);
+        assertFalse(retryWithToleranceOperator.checkRetry(0));
+
+        time.setCurrentTimeMs(600);
+        assertFalse(retryWithToleranceOperator.checkRetry(0));
+    }
+
+    @Test
+    public void testBackoffLimit() {
+        MockTime time = new MockTime(0, 0, 0);
+        RetryWithToleranceOperator retryWithToleranceOperator = new RetryWithToleranceOperator(5, 5000, NONE, time);
+
+        long prevTs = time.hiResClockMs();
+        retryWithToleranceOperator.backoff(1, 5000);
+        assertEquals(300, time.hiResClockMs() - prevTs);
+
+        prevTs = time.hiResClockMs();
+        retryWithToleranceOperator.backoff(2, 5000);
+        assertEquals(600, time.hiResClockMs() - prevTs);
+
+        prevTs = time.hiResClockMs();
+        retryWithToleranceOperator.backoff(3, 5000);
+        assertEquals(1200, time.hiResClockMs() - prevTs);
+
+        prevTs = time.hiResClockMs();
+        retryWithToleranceOperator.backoff(4, 5000);
+        assertEquals(2400, time.hiResClockMs() - prevTs);
+
+        prevTs = time.hiResClockMs();
+        retryWithToleranceOperator.backoff(5, 5000);
+        assertEquals(500, time.hiResClockMs() - prevTs);
+
+        prevTs = time.hiResClockMs();
+        retryWithToleranceOperator.backoff(6, 5000);
+        assertEquals(0, time.hiResClockMs() - prevTs);
+
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testToleranceLimit() {
+        RetryWithToleranceOperator retryWithToleranceOperator = new RetryWithToleranceOperator(ERRORS_RETRY_TIMEOUT_DEFAULT, ERRORS_RETRY_MAX_DELAY_DEFAULT, NONE, SYSTEM);
+        retryWithToleranceOperator.metrics(errorHandlingMetrics);
+        retryWithToleranceOperator.markAsFailed();
+        assertFalse("should not tolerate any errors", retryWithToleranceOperator.withinToleranceLimits());
+
+        retryWithToleranceOperator = new RetryWithToleranceOperator(ERRORS_RETRY_TIMEOUT_DEFAULT, ERRORS_RETRY_MAX_DELAY_DEFAULT, ALL, SYSTEM);
+        retryWithToleranceOperator.metrics(errorHandlingMetrics);
+        retryWithToleranceOperator.markAsFailed();
+        retryWithToleranceOperator.markAsFailed();
+        assertTrue("should tolerate all errors", retryWithToleranceOperator.withinToleranceLimits());
+
+        retryWithToleranceOperator = new RetryWithToleranceOperator(ERRORS_RETRY_TIMEOUT_DEFAULT, ERRORS_RETRY_MAX_DELAY_DEFAULT, NONE, SYSTEM);
+        assertTrue("no tolerance is within limits if no failures", retryWithToleranceOperator.withinToleranceLimits());
+    }
+
+    @Test
+    public void testDefaultConfigs() {
+        ConnectorConfig configuration = config(emptyMap());
+        assertEquals(configuration.errorRetryTimeout(), ERRORS_RETRY_TIMEOUT_DEFAULT);
+        assertEquals(configuration.errorMaxDelayInMillis(), ERRORS_RETRY_MAX_DELAY_DEFAULT);
+        assertEquals(configuration.errorToleranceType(), ERRORS_TOLERANCE_DEFAULT);
+
+        PowerMock.verifyAll();
+    }
+
+    ConnectorConfig config(Map connProps) {
+        Map props = new HashMap<>();
+        props.put(ConnectorConfig.NAME_CONFIG, "test");
+        props.put(ConnectorConfig.CONNECTOR_CLASS_CONFIG, SinkTask.class.getName());
+        props.putAll(connProps);
+        return new ConnectorConfig(plugins, props);
+    }
+
+    @Test
+    public void testSetConfigs() {
+        ConnectorConfig configuration;
+        configuration = config(singletonMap(ERRORS_RETRY_TIMEOUT_CONFIG, "100"));
+        assertEquals(configuration.errorRetryTimeout(), 100);
+
+        configuration = config(singletonMap(ERRORS_RETRY_MAX_DELAY_CONFIG, "100"));
+        assertEquals(configuration.errorMaxDelayInMillis(), 100);
+
+        configuration = config(singletonMap(ERRORS_TOLERANCE_CONFIG, "none"));
+        assertEquals(configuration.errorToleranceType(), ToleranceType.NONE);
+
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testThreadSafety() throws Throwable {
+        long runtimeMs = 5_000;
+        int numThreads = 10;
+        // Check that multiple threads using RetryWithToleranceOperator concurrently
+        // can't corrupt the state of the ProcessingContext
+        AtomicReference failed = new AtomicReference<>(null);
+        RetryWithToleranceOperator retryWithToleranceOperator = new RetryWithToleranceOperator(0,
+                ERRORS_RETRY_MAX_DELAY_DEFAULT, ALL, SYSTEM, new ProcessingContext() {
+                    private AtomicInteger count = new AtomicInteger();
+                    private AtomicInteger attempt = new AtomicInteger();
+
+                    @Override
+                    public void error(Throwable error) {
+                        if (count.getAndIncrement() > 0) {
+                            failed.compareAndSet(null, new AssertionError("Concurrent call to error()"));
+                        }
+                        super.error(error);
+                    }
+
+                    @Override
+                    public Future report() {
+                        if (count.getAndSet(0) > 1) {
+                            failed.compareAndSet(null, new AssertionError("Concurrent call to error() in report()"));
+                        }
+
+                        return super.report();
+                    }
+
+                    @Override
+                    public void currentContext(Stage stage, Class klass) {
+                        this.attempt.set(0);
+                        super.currentContext(stage, klass);
+                    }
+
+                    @Override
+                    public void attempt(int attempt) {
+                        if (!this.attempt.compareAndSet(attempt - 1, attempt)) {
+                            failed.compareAndSet(null, new AssertionError(
+                                    "Concurrent call to attempt(): Attempts should increase monotonically " +
+                                            "within the scope of a given currentContext()"));
+                        }
+                        super.attempt(attempt);
+                    }
+                });
+        retryWithToleranceOperator.metrics(errorHandlingMetrics);
+
+        ExecutorService pool = Executors.newFixedThreadPool(numThreads);
+        List> futures = IntStream.range(0, numThreads).boxed()
+                .map(id ->
+                        pool.submit(() -> {
+                            long t0 = System.currentTimeMillis();
+                            long i = 0;
+                            while (true) {
+                                if (++i % 10000 == 0 && System.currentTimeMillis() > t0 + runtimeMs) {
+                                    break;
+                                }
+                                if (failed.get() != null) {
+                                    break;
+                                }
+                                try {
+                                    if (id < numThreads / 2) {
+                                        retryWithToleranceOperator.executeFailed(Stage.TASK_PUT,
+                                                SinkTask.class, consumerRecord, new Throwable()).get();
+                                    } else {
+                                        retryWithToleranceOperator.execute(() -> null, Stage.TRANSFORMATION,
+                                                SinkTask.class);
+                                    }
+                                } catch (Exception e) {
+                                    failed.compareAndSet(null, e);
+                                }
+                            }
+                        }))
+                .collect(Collectors.toList());
+        pool.shutdown();
+        pool.awaitTermination((long) (1.5 * runtimeMs), TimeUnit.MILLISECONDS);
+        futures.forEach(future -> {
+            try {
+                future.get();
+            } catch (Exception e) {
+                failed.compareAndSet(null, e);
+            }
+        });
+        Throwable exception = failed.get();
+        if (exception != null) {
+            throw exception;
+        }
+    }
+
+
+    private static class ExceptionThrower implements Operation {
+        private Exception e;
+
+        public ExceptionThrower(Exception e) {
+            this.e = e;
+        }
+
+        @Override
+        public Object call() throws Exception {
+            throw e;
+        }
+    }
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/errors/WorkerErrantRecordReporterTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/errors/WorkerErrantRecordReporterTest.java
new file mode 100644
index 0000000000000..07a4f9e8f1197
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/errors/WorkerErrantRecordReporterTest.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.errors;
+
+import org.apache.kafka.connect.sink.SinkRecord;
+import org.apache.kafka.connect.storage.Converter;
+import org.apache.kafka.connect.storage.HeaderConverter;
+import org.easymock.Mock;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.powermock.core.classloader.annotations.PowerMockIgnore;
+import org.powermock.modules.junit4.PowerMockRunner;
+
+import java.util.concurrent.CompletableFuture;
+
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+
+@RunWith(PowerMockRunner.class)
+@PowerMockIgnore("javax.management.*")
+public class WorkerErrantRecordReporterTest {
+
+    private WorkerErrantRecordReporter reporter;
+
+    @Mock
+    private RetryWithToleranceOperator retryWithToleranceOperator;
+
+    @Mock
+    private Converter converter;
+
+    @Mock
+    private HeaderConverter headerConverter;
+
+    @Mock
+    private SinkRecord record;
+
+    @Before
+    public void setup() {
+        reporter = new WorkerErrantRecordReporter(
+            retryWithToleranceOperator,
+            converter,
+            converter,
+            headerConverter
+        );
+    }
+
+    @Test
+    public void testGetAllFutures() {
+        assertTrue(reporter.futures.isEmpty());
+        for (int i = 0; i < 4; i++) {
+            reporter.futures.add(CompletableFuture.completedFuture(null));
+        }
+        assertFalse(reporter.futures.isEmpty());
+        reporter.awaitAllFutures();
+        assertTrue(reporter.futures.isEmpty());
+    }
+}
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
new file mode 100644
index 0000000000000..757a43752d225
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/isolation/DelegatingClassLoaderTest.java
@@ -0,0 +1,131 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.kafka.connect.runtime.isolation;
+
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+
+import java.io.File;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.Collections;
+
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+
+public class DelegatingClassLoaderTest {
+
+    @Rule
+    public TemporaryFolder pluginDir = new TemporaryFolder();
+
+    @Test
+    public void testPermittedManifestResources() {
+        assertTrue(
+            DelegatingClassLoader.serviceLoaderManifestForPlugin("META-INF/services/org.apache.kafka.connect.rest.ConnectRestExtension"));
+        assertTrue(
+            DelegatingClassLoader.serviceLoaderManifestForPlugin("META-INF/services/org.apache.kafka.common.config.provider.ConfigProvider"));
+    }
+
+    @Test
+    public void testOtherResources() {
+        assertFalse(
+            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));
+        }
+    }
+
+    @Test
+    public void testLoadingInvalidUberJar() throws Exception {
+        pluginDir.newFile("invalid.jar");
+
+        DelegatingClassLoader classLoader = new DelegatingClassLoader(
+            Collections.singletonList(pluginDir.getRoot().getAbsolutePath()));
+        classLoader.initLoaders();
+    }
+
+    @Test
+    public void testLoadingPluginDirContainsInvalidJarsOnly() throws Exception {
+        pluginDir.newFolder("my-plugin");
+        pluginDir.newFile("my-plugin/invalid.jar");
+
+        DelegatingClassLoader classLoader = new DelegatingClassLoader(
+            Collections.singletonList(pluginDir.getRoot().getAbsolutePath()));
+        classLoader.initLoaders();
+    }
+
+    @Test
+    public void testLoadingNoPlugins() throws Exception {
+        DelegatingClassLoader classLoader = new DelegatingClassLoader(
+            Collections.singletonList(pluginDir.getRoot().getAbsolutePath()));
+        classLoader.initLoaders();
+    }
+
+    @Test
+    public void testLoadingPluginDirEmpty() throws Exception {
+        pluginDir.newFolder("my-plugin");
+
+        DelegatingClassLoader classLoader = new DelegatingClassLoader(
+            Collections.singletonList(pluginDir.getRoot().getAbsolutePath()));
+        classLoader.initLoaders();
+    }
+
+    @Test
+    public void testLoadingMixOfValidAndInvalidPlugins() throws Exception {
+        TestPlugins.assertAvailable();
+
+        pluginDir.newFile("invalid.jar");
+        pluginDir.newFolder("my-plugin");
+        pluginDir.newFile("my-plugin/invalid.jar");
+        Path pluginPath = this.pluginDir.getRoot().toPath();
+
+        for (String sourceJar : TestPlugins.pluginPath()) {
+            Path source = new File(sourceJar).toPath();
+            Files.copy(source, pluginPath.resolve(source.getFileName()));
+        }
+
+        DelegatingClassLoader classLoader = new DelegatingClassLoader(
+            Collections.singletonList(pluginDir.getRoot().getAbsolutePath()));
+        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/PluginDescTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/isolation/PluginDescTest.java
index 3d8f7035953e9..950aa9021017c 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/isolation/PluginDescTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/isolation/PluginDescTest.java
@@ -49,7 +49,7 @@ public void setUp() throws Exception {
     }
 
     @Test
-    public void testRegularPluginDesc() throws Exception {
+    public void testRegularPluginDesc() {
         PluginDesc connectorDesc = new PluginDesc<>(
                 Connector.class,
                 regularVersion,
@@ -76,7 +76,7 @@ public void testRegularPluginDesc() throws Exception {
     }
 
     @Test
-    public void testPluginDescWithSystemClassLoader() throws Exception {
+    public void testPluginDescWithSystemClassLoader() {
         String location = "classpath";
         PluginDesc connectorDesc = new PluginDesc<>(
                 SinkConnector.class,
@@ -104,7 +104,7 @@ public void testPluginDescWithSystemClassLoader() throws Exception {
     }
 
     @Test
-    public void testPluginDescWithNullVersion() throws Exception {
+    public void testPluginDescWithNullVersion() {
         String nullVersion = "null";
         PluginDesc connectorDesc = new PluginDesc<>(
                 SourceConnector.class,
@@ -130,7 +130,7 @@ public void testPluginDescWithNullVersion() throws Exception {
     }
 
     @Test
-    public void testPluginDescEquality() throws Exception {
+    public void testPluginDescEquality() {
         PluginDesc connectorDescPluginPath = new PluginDesc<>(
                 Connector.class,
                 snaphotVersion,
@@ -177,7 +177,7 @@ public void testPluginDescEquality() throws Exception {
     }
 
     @Test
-    public void testPluginDescComparison() throws Exception {
+    public void testPluginDescComparison() {
         PluginDesc connectorDescPluginPath = new PluginDesc<>(
                 Connector.class,
                 regularVersion,
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 f9532a6d8cb62..19766989fcf46 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
@@ -26,6 +26,7 @@
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 
@@ -44,7 +45,7 @@ public void setUp() throws Exception {
     }
 
     @Test
-    public void testJavaLibraryClasses() throws Exception {
+    public void testJavaLibraryClasses() {
         assertFalse(PluginUtils.shouldLoadInIsolation("java."));
         assertFalse(PluginUtils.shouldLoadInIsolation("java.lang.Object"));
         assertFalse(PluginUtils.shouldLoadInIsolation("java.lang.String"));
@@ -63,13 +64,13 @@ public void testJavaLibraryClasses() throws Exception {
     }
 
     @Test
-    public void testThirdPartyClasses() throws Exception {
+    public void testThirdPartyClasses() {
         assertFalse(PluginUtils.shouldLoadInIsolation("org.slf4j."));
         assertFalse(PluginUtils.shouldLoadInIsolation("org.slf4j.LoggerFactory"));
     }
 
     @Test
-    public void testConnectFrameworkClasses() throws Exception {
+    public void testKafkaDependencyClasses() {
         assertFalse(PluginUtils.shouldLoadInIsolation("org.apache.kafka.common."));
         assertFalse(PluginUtils.shouldLoadInIsolation(
                 "org.apache.kafka.common.config.AbstractConfig")
@@ -80,61 +81,292 @@ public void testConnectFrameworkClasses() throws Exception {
         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")
+                "org.apache.kafka.clients.producer.ProducerConfig")
         );
         assertFalse(PluginUtils.shouldLoadInIsolation(
-                "org.apache.kafka.connect.source.SourceConnector")
+                "org.apache.kafka.clients.consumer.ConsumerConfig")
         );
         assertFalse(PluginUtils.shouldLoadInIsolation(
-                "org.apache.kafka.connect.sink.SinkConnector")
+                "org.apache.kafka.clients.admin.KafkaAdminClient")
         );
-        assertFalse(PluginUtils.shouldLoadInIsolation("org.apache.kafka.connect.connector.Task"));
-        assertFalse(PluginUtils.shouldLoadInIsolation(
-                "org.apache.kafka.connect.source.SourceTask")
+    }
+
+    @Test
+    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.transforms.predicates.Predicate",
+            "org.apache.kafka.connect.util.",
+            "org.apache.kafka.connect.util.ConnectorUtils"
         );
-        assertFalse(PluginUtils.shouldLoadInIsolation("org.apache.kafka.connect.sink.SinkTask"));
-        assertFalse(PluginUtils.shouldLoadInIsolation(
-                "org.apache.kafka.connect.transforms.Transformation")
+        // 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."
         );
-        assertFalse(PluginUtils.shouldLoadInIsolation(
-                "org.apache.kafka.connect.storage.Converter")
+        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"
         );
-        assertFalse(PluginUtils.shouldLoadInIsolation(
-                "org.apache.kafka.connect.storage.OffsetBackingStore")
+        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",
+            "org.apache.kafka.connect.transforms.predicates.",
+            "org.apache.kafka.connect.transforms.predicates.HasHeaderKey",
+            "org.apache.kafka.connect.transforms.predicates.RecordIsTombstone",
+            "org.apache.kafka.connect.transforms.predicates.TopicNameMatches"
         );
+        for (String clazz : transformsClasses) {
+            assertTrue(
+                clazz + " from 'transforms' is not loaded in isolation but should be",
+                PluginUtils.shouldLoadInIsolation(clazz)
+            );
+        }
     }
 
     @Test
-    public void testAllowedConnectFrameworkClasses() throws Exception {
-        assertTrue(PluginUtils.shouldLoadInIsolation("org.apache.kafka.connect.transforms."));
-        assertTrue(PluginUtils.shouldLoadInIsolation(
-                "org.apache.kafka.connect.transforms.ExtractField")
+    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"
         );
-        assertTrue(PluginUtils.shouldLoadInIsolation(
-                "org.apache.kafka.connect.transforms.ExtractField$Key")
+        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"
         );
-        assertTrue(PluginUtils.shouldLoadInIsolation("org.apache.kafka.connect.json."));
-        assertTrue(PluginUtils.shouldLoadInIsolation(
-                "org.apache.kafka.connect.json.JsonConverter")
+        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.json.JsonConverter$21")
+                "org.apache.kafka.connect.mirror.MirrorSourceTask")
         );
-        assertTrue(PluginUtils.shouldLoadInIsolation("org.apache.kafka.connect.file."));
         assertTrue(PluginUtils.shouldLoadInIsolation(
-                "org.apache.kafka.connect.file.FileStreamSourceTask")
+                "org.apache.kafka.connect.mirror.MirrorSourceConnector")
         );
-        assertTrue(PluginUtils.shouldLoadInIsolation(
-                "org.apache.kafka.connect.file.FileStreamSinkConnector")
+    }
+
+    @Test
+    public void testClientConfigProvider() {
+        assertFalse(PluginUtils.shouldLoadInIsolation(
+                "org.apache.kafka.common.config.provider.ConfigProvider")
         );
-        assertTrue(PluginUtils.shouldLoadInIsolation("org.apache.kafka.connect.converters."));
         assertTrue(PluginUtils.shouldLoadInIsolation(
-                "org.apache.kafka.connect.converters.ByteArrayConverter")
+                "org.apache.kafka.common.config.provider.FileConfigProvider")
         );
         assertTrue(PluginUtils.shouldLoadInIsolation(
-                "org.apache.kafka.connect.storage.StringConverter")
+                "org.apache.kafka.common.config.provider.FutureConfigProvider")
         );
     }
 
@@ -158,6 +390,22 @@ public void testPluginUrlsWithJars() throws Exception {
         assertUrls(expectedUrls, PluginUtils.pluginUrls(pluginPath));
     }
 
+    @Test
+    public void testOrderOfPluginUrlsWithJars() throws Exception {
+        createBasicDirectoryLayout();
+        // Here this method is just used to create the files. The result is not used.
+        createBasicExpectedUrls();
+
+        List actual = PluginUtils.pluginUrls(pluginPath);
+        // 'simple-transform.jar' is created first. In many cases, without sorting within the
+        // PluginUtils, this jar will be placed before 'another-transform.jar'. However this is
+        // not guaranteed because a DirectoryStream does not maintain a certain order in its
+        // results. Besides this test case, sorted order in every call to assertUrls below.
+        int i = Arrays.toString(actual.toArray()).indexOf("another-transform.jar");
+        int j = Arrays.toString(actual.toArray()).indexOf("simple-transform.jar");
+        assertTrue(i < j);
+    }
+
     @Test
     public void testPluginUrlsWithZips() throws Exception {
         createBasicDirectoryLayout();
@@ -262,9 +510,8 @@ private List createBasicExpectedUrls() throws IOException {
     }
 
     private void assertUrls(List expected, List actual) {
-        List actualCopy = new ArrayList<>(actual);
         Collections.sort(expected);
-        Collections.sort(actualCopy);
-        assertEquals(expected, actualCopy);
+        // not sorting 'actual' because it should be returned sorted from withing the PluginUtils.
+        assertEquals(expected, actual);
     }
 }
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
new file mode 100644
index 0000000000000..4d304c710d838
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/isolation/PluginsTest.java
@@ -0,0 +1,482 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.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;
+import org.apache.kafka.connect.rest.ConnectRestExtensionContext;
+import org.apache.kafka.connect.runtime.WorkerConfig;
+import org.apache.kafka.connect.runtime.isolation.Plugins.ClassLoaderUsage;
+import org.apache.kafka.connect.storage.Converter;
+import org.apache.kafka.connect.storage.ConverterConfig;
+import org.apache.kafka.connect.storage.ConverterType;
+import org.apache.kafka.connect.storage.HeaderConverter;
+import org.apache.kafka.connect.storage.SimpleHeaderConverter;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.io.IOException;
+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.assertNull;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
+
+public class PluginsTest {
+
+    private Plugins plugins;
+    private Map props;
+    private AbstractConfig config;
+    private TestConverter converter;
+    private TestHeaderConverter headerConverter;
+    private TestInternalConverter internalConverter;
+
+    @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());
+        props.put("key.converter." + JsonConverterConfig.SCHEMAS_ENABLE_CONFIG, "true");
+        props.put("value.converter." + JsonConverterConfig.SCHEMAS_ENABLE_CONFIG, "true");
+        props.put("key.converter.extra.config", "foo1");
+        props.put("value.converter.extra.config", "foo2");
+        props.put(WorkerConfig.INTERNAL_KEY_CONVERTER_CLASS_CONFIG, TestInternalConverter.class.getName());
+        props.put(WorkerConfig.INTERNAL_VALUE_CONVERTER_CLASS_CONFIG, TestInternalConverter.class.getName());
+        props.put("internal.key.converter.extra.config", "bar1");
+        props.put("internal.value.converter.extra.config", "bar2");
+        props.put(WorkerConfig.HEADER_CONVERTER_CLASS_CONFIG, TestHeaderConverter.class.getName());
+        props.put("header.converter.extra.config", "baz");
+
+        createConfig();
+    }
+
+    protected void createConfig() {
+        this.config = new TestableWorkerConfig(props);
+    }
+
+    @Test
+    public void shouldInstantiateAndConfigureConverters() {
+        instantiateAndConfigureConverter(WorkerConfig.KEY_CONVERTER_CLASS_CONFIG, ClassLoaderUsage.CURRENT_CLASSLOADER);
+        // Validate extra configs got passed through to overridden converters
+        assertEquals("true", converter.configs.get(JsonConverterConfig.SCHEMAS_ENABLE_CONFIG));
+        assertEquals("foo1", converter.configs.get("extra.config"));
+
+        instantiateAndConfigureConverter(WorkerConfig.VALUE_CONVERTER_CLASS_CONFIG, ClassLoaderUsage.PLUGINS);
+        // Validate extra configs got passed through to overridden converters
+        assertEquals("true", converter.configs.get(JsonConverterConfig.SCHEMAS_ENABLE_CONFIG));
+        assertEquals("foo2", converter.configs.get("extra.config"));
+    }
+
+    @SuppressWarnings("deprecation")
+    @Test
+    public void shouldInstantiateAndConfigureInternalConverters() {
+        instantiateAndConfigureInternalConverter(WorkerConfig.INTERNAL_KEY_CONVERTER_CLASS_CONFIG, ClassLoaderUsage.CURRENT_CLASSLOADER);
+        // Validate schemas.enable is defaulted to false for internal converter
+        assertEquals(false, internalConverter.configs.get(JsonConverterConfig.SCHEMAS_ENABLE_CONFIG));
+        // Validate internal converter properties can still be set
+        assertEquals("bar1", internalConverter.configs.get("extra.config"));
+
+        instantiateAndConfigureInternalConverter(WorkerConfig.INTERNAL_VALUE_CONVERTER_CLASS_CONFIG, ClassLoaderUsage.PLUGINS);
+        // Validate schemas.enable is defaulted to false for internal converter
+        assertEquals(false, internalConverter.configs.get(JsonConverterConfig.SCHEMAS_ENABLE_CONFIG));
+        // Validate internal converter properties can still be set
+        assertEquals("bar2", internalConverter.configs.get("extra.config"));
+    }
+
+    @Test
+    public void shouldInstantiateAndConfigureExplicitlySetHeaderConverterWithCurrentClassLoader() {
+        assertNotNull(props.get(WorkerConfig.HEADER_CONVERTER_CLASS_CONFIG));
+        HeaderConverter headerConverter = plugins.newHeaderConverter(config,
+                                                                     WorkerConfig.HEADER_CONVERTER_CLASS_CONFIG,
+                                                                     ClassLoaderUsage.CURRENT_CLASSLOADER);
+        assertNotNull(headerConverter);
+        assertTrue(headerConverter instanceof TestHeaderConverter);
+        this.headerConverter = (TestHeaderConverter) headerConverter;
+
+        // Validate extra configs got passed through to overridden converters
+        assertConverterType(ConverterType.HEADER, this.headerConverter.configs);
+        assertEquals("baz", this.headerConverter.configs.get("extra.config"));
+
+        headerConverter = plugins.newHeaderConverter(config,
+                                                     WorkerConfig.HEADER_CONVERTER_CLASS_CONFIG,
+                                                     ClassLoaderUsage.PLUGINS);
+        assertNotNull(headerConverter);
+        assertTrue(headerConverter instanceof TestHeaderConverter);
+        this.headerConverter = (TestHeaderConverter) headerConverter;
+
+        // Validate extra configs got passed through to overridden converters
+        assertConverterType(ConverterType.HEADER, this.headerConverter.configs);
+        assertEquals("baz", this.headerConverter.configs.get("extra.config"));
+    }
+
+    @Test
+    public void shouldInstantiateAndConfigureConnectRestExtension() {
+        props.put(WorkerConfig.REST_EXTENSION_CLASSES_CONFIG,
+                  TestConnectRestExtension.class.getName());
+        createConfig();
+
+        List connectRestExtensions =
+            plugins.newPlugins(config.getList(WorkerConfig.REST_EXTENSION_CLASSES_CONFIG),
+                               config,
+                               ConnectRestExtension.class);
+        assertNotNull(connectRestExtensions);
+        assertEquals("One Rest Extension expected", 1, connectRestExtensions.size());
+        assertNotNull(connectRestExtensions.get(0));
+        assertTrue("Should be instance of TestConnectRestExtension",
+                   connectRestExtensions.get(0) instanceof TestConnectRestExtension);
+        assertNotNull(((TestConnectRestExtension) connectRestExtensions.get(0)).configs);
+        assertEquals(config.originals(),
+                     ((TestConnectRestExtension) connectRestExtensions.get(0)).configs);
+    }
+
+    @Test
+    public void shouldInstantiateAndConfigureDefaultHeaderConverter() {
+        props.remove(WorkerConfig.HEADER_CONVERTER_CLASS_CONFIG);
+        createConfig();
+
+        // Because it's not explicitly set on the supplied configuration, the logic to use the current classloader for the connector
+        // will exit immediately, and so this method always returns null
+        HeaderConverter headerConverter = plugins.newHeaderConverter(config,
+                                                                     WorkerConfig.HEADER_CONVERTER_CLASS_CONFIG,
+                                                                     ClassLoaderUsage.CURRENT_CLASSLOADER);
+        assertNull(headerConverter);
+        // But we should always find it (or the worker's default) when using the plugins classloader ...
+        headerConverter = plugins.newHeaderConverter(config,
+                                                     WorkerConfig.HEADER_CONVERTER_CLASS_CONFIG,
+                                                     ClassLoaderUsage.PLUGINS);
+        assertNotNull(headerConverter);
+        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);
+    }
+
+    protected void instantiateAndConfigureHeaderConverter(String configPropName) {
+        headerConverter = (TestHeaderConverter) plugins.newHeaderConverter(config, configPropName, ClassLoaderUsage.CURRENT_CLASSLOADER);
+        assertNotNull(headerConverter);
+    }
+
+    protected void instantiateAndConfigureInternalConverter(String configPropName, ClassLoaderUsage classLoaderUsage) {
+        internalConverter = (TestInternalConverter) plugins.newConverter(config, configPropName, classLoaderUsage);
+        assertNotNull(internalConverter);
+    }
+
+    protected void assertConverterType(ConverterType type, Map props) {
+        assertEquals(type.getName(), props.get(ConverterConfig.TYPE_CONFIG));
+    }
+
+    public static class TestableWorkerConfig extends WorkerConfig {
+        public TestableWorkerConfig(Map props) {
+            super(WorkerConfig.baseConfigDef(), props);
+        }
+    }
+
+    public static class TestConverter implements Converter, Configurable {
+        public Map configs;
+
+        public ConfigDef config() {
+            return JsonConverterConfig.configDef();
+        }
+
+        @Override
+        public void configure(Map configs) {
+            this.configs = configs;
+            new JsonConverterConfig(configs); // requires the `converter.type` config be set
+        }
+
+        @Override
+        public void configure(Map configs, boolean isKey) {
+            this.configs = configs;
+        }
+
+        @Override
+        public byte[] fromConnectData(String topic, Schema schema, Object value) {
+            return new byte[0];
+        }
+
+        @Override
+        public SchemaAndValue toConnectData(String topic, byte[] value) {
+            return null;
+        }
+    }
+
+    public static class TestHeaderConverter implements HeaderConverter {
+        public Map configs;
+
+        @Override
+        public ConfigDef config() {
+            return JsonConverterConfig.configDef();
+        }
+
+        @Override
+        public void configure(Map configs) {
+            this.configs = configs;
+            new JsonConverterConfig(configs); // requires the `converter.type` config be set
+        }
+
+        @Override
+        public byte[] fromConnectHeader(String topic, String headerKey, Schema schema, Object value) {
+            return new byte[0];
+        }
+
+        @Override
+        public SchemaAndValue toConnectHeader(String topic, String headerKey, byte[] value) {
+            return null;
+        }
+
+        @Override
+        public void close() throws IOException {
+        }
+    }
+
+
+    public static class TestConnectRestExtension implements ConnectRestExtension {
+
+        public Map configs;
+
+        @Override
+        public void register(ConnectRestExtensionContext restPluginContext) {
+        }
+
+        @Override
+        public void close() throws IOException {
+        }
+
+        @Override
+        public void configure(Map configs) {
+            this.configs = configs;
+        }
+
+        @Override
+        public String version() {
+            return "test";
+        }
+    }
+
+    public static class TestInternalConverter extends JsonConverter {
+        public Map configs;
+
+        public void configure(Map configs) {
+            this.configs = configs;
+            super.configure(configs);
+        }
+    }
+}
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 8c533721b762f..0c81ddd739411 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,146 +16,473 @@
  */
 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.WorkerConfig;
+import org.apache.kafka.connect.runtime.distributed.DistributedConfig;
+import org.apache.kafka.connect.runtime.isolation.Plugins;
 import org.apache.kafka.connect.runtime.standalone.StandaloneConfig;
-import org.apache.kafka.connect.util.Callback;
-import org.easymock.Capture;
 import org.easymock.EasyMock;
-import org.easymock.IAnswer;
 import org.junit.After;
+import org.junit.Assert;
 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;
+import org.slf4j.LoggerFactory;
 
-import java.net.URI;
-import java.net.URISyntaxException;
+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 javax.ws.rs.client.Client;
-import javax.ws.rs.client.ClientBuilder;
-import javax.ws.rs.client.Invocation;
-import javax.ws.rs.client.WebTarget;
-import javax.ws.rs.core.Response;
-
+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.*", "javax.crypto.*"})
 public class RestServerTest {
-
     @MockStrict
     private Herder herder;
+    @MockStrict
+    private Plugins plugins;
     private RestServer server;
 
+    protected static final String KAFKA_CLUSTER_ID = "Xbafgnagvar";
+
     @After
     public void tearDown() {
-        server.stop();
+        if (server != null) {
+            server.stop();
+        }
     }
 
+    @SuppressWarnings("deprecation")
     private Map baseWorkerProps() {
         Map workerProps = new HashMap<>();
-        workerProps.put("key.converter", "org.apache.kafka.connect.json.JsonConverter");
-        workerProps.put("value.converter", "org.apache.kafka.connect.json.JsonConverter");
-        workerProps.put("internal.key.converter", "org.apache.kafka.connect.json.JsonConverter");
-        workerProps.put("internal.value.converter", "org.apache.kafka.connect.json.JsonConverter");
-        workerProps.put("internal.key.converter.schemas.enable", "false");
-        workerProps.put("internal.value.converter.schemas.enable", "false");
-        workerProps.put("offset.storage.file.filename", "/tmp/connect.offsets");
+        workerProps.put(DistributedConfig.STATUS_STORAGE_TOPIC_CONFIG, "status-topic");
+        workerProps.put(DistributedConfig.CONFIG_TOPIC_CONFIG, "config-topic");
+        workerProps.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092");
+        workerProps.put(DistributedConfig.GROUP_ID_CONFIG, "connect-test-group");
+        workerProps.put(WorkerConfig.KEY_CONVERTER_CLASS_CONFIG, "org.apache.kafka.connect.json.JsonConverter");
+        workerProps.put(WorkerConfig.VALUE_CONVERTER_CLASS_CONFIG, "org.apache.kafka.connect.json.JsonConverter");
+        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;
     }
 
     @Test
-    public void testCORSEnabled() {
+    public void testCORSEnabled() throws IOException {
         checkCORSRequest("*", "http://bar.com", "http://bar.com", "PUT");
     }
 
     @Test
-    public void testCORSDisabled() {
+    public void testCORSDisabled() throws IOException {
         checkCORSRequest("", "http://bar.com", null, null);
     }
 
-    public void checkCORSRequest(String corsDomain, String origin, String expectedHeader, String method) {
-        // To be able to set the Origin, we need to toggle this flag
-        System.setProperty("sun.net.http.allowRestrictedHeaders", "true");
+    @SuppressWarnings("deprecation")
+    @Test
+    public void testParseListeners() {
+        // Use listeners field
+        Map configMap = new HashMap<>(baseWorkerProps());
+        configMap.put(WorkerConfig.LISTENERS_CONFIG, "http://localhost:8080,https://localhost:8443");
+        DistributedConfig config = new DistributedConfig(configMap);
+
+        server = new RestServer(config);
+        Assert.assertArrayEquals(new String[] {"http://localhost:8080", "https://localhost:8443"}, server.parseListeners().toArray());
+
+        // 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);
+        server = new RestServer(config);
+        Assert.assertArrayEquals(new String[] {"http://my-hostname:8080"}, server.parseListeners().toArray());
+    }
+
+    @SuppressWarnings("deprecation")
+    @Test
+    public void testAdvertisedUri() {
+        // 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);
 
-        final Capture>> connectorsCallback = EasyMock.newCapture();
-        herder.connectors(EasyMock.capture(connectorsCallback));
-        PowerMock.expectLastCall().andAnswer(new IAnswer() {
-            @Override
-            public Object answer() throws Throwable {
-                connectorsCallback.getValue().onCompletion(null, Arrays.asList("a", "b"));
-                return null;
-            }
-        });
+        server = new RestServer(config);
+        Assert.assertEquals("http://localhost:8080/", server.advertisedUrl().toString());
 
+        // Advertised URI from listeners with protocol
+        configMap = new HashMap<>(baseWorkerProps());
+        configMap.put(WorkerConfig.LISTENERS_CONFIG, "http://localhost:8080,https://localhost:8443");
+        configMap.put(WorkerConfig.REST_ADVERTISED_LISTENER_CONFIG, "https");
+        config = new DistributedConfig(configMap);
+
+        server = new RestServer(config);
+        Assert.assertEquals("https://localhost:8443/", server.advertisedUrl().toString());
+
+        // Advertised URI from listeners with only SSL available
+        configMap = new HashMap<>(baseWorkerProps());
+        configMap.put(WorkerConfig.LISTENERS_CONFIG, "https://localhost:8443");
+        config = new DistributedConfig(configMap);
+
+        server = new RestServer(config);
+        Assert.assertEquals("https://localhost:8443/", server.advertisedUrl().toString());
+
+        // Listener is overriden by advertised values
+        configMap = new HashMap<>(baseWorkerProps());
+        configMap.put(WorkerConfig.LISTENERS_CONFIG, "https://localhost:8443");
+        configMap.put(WorkerConfig.REST_ADVERTISED_LISTENER_CONFIG, "http");
+        configMap.put(WorkerConfig.REST_ADVERTISED_HOST_NAME_CONFIG, "somehost");
+        configMap.put(WorkerConfig.REST_ADVERTISED_PORT_CONFIG, "10000");
+        config = new DistributedConfig(configMap);
+
+        server = new RestServer(config);
+        Assert.assertEquals("http://somehost:10000/", server.advertisedUrl().toString());
+
+        // 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
+    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,
+            ConnectRestExtension.class))
+            .andStubReturn(Collections.emptyList());
         PowerMock.replayAll();
 
+        server = new RestServer(workerConfig);
+        server.initializeServer();
+        server.initializeResources(herder);
+
+        HttpOptions request = new HttpOptions("/connectors");
+        request.addHeader("Content-Type", MediaType.WILDCARD);
+        CloseableHttpClient httpClient = HttpClients.createMinimal();
+        HttpHost httpHost = new HttpHost(
+            server.advertisedUrl().getHost(),
+            server.advertisedUrl().getPort()
+        );
+        CloseableHttpResponse response = httpClient.execute(httpHost, request);
+        Assert.assertEquals(MediaType.TEXT_PLAIN, response.getEntity().getContentType().getValue());
+        ByteArrayOutputStream baos = new ByteArrayOutputStream();
+        response.getEntity().writeTo(baos);
+        Assert.assertArrayEquals(
+            request.getAllowedMethods(response).toArray(),
+            new String(baos.toByteArray(), StandardCharsets.UTF_8).split(", ")
+        );
+        PowerMock.verifyAll();
+    }
+
+    public void checkCORSRequest(String corsDomain, String origin, String expectedHeader, String method)
+        throws IOException {
         Map workerProps = baseWorkerProps();
         workerProps.put(WorkerConfig.ACCESS_CONTROL_ALLOW_ORIGIN_CONFIG, corsDomain);
         workerProps.put(WorkerConfig.ACCESS_CONTROL_ALLOW_METHODS_CONFIG, method);
-        WorkerConfig workerConfig = new StandaloneConfig(workerProps);
+        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());
+
+        EasyMock.expect(herder.connectors()).andReturn(Arrays.asList("a", "b"));
+
+        PowerMock.replayAll();
+
         server = new RestServer(workerConfig);
-        server.start(herder);
-
-        Response response = request("/connectors")
-                .header("Referer", origin + "/page")
-                .header("Origin", origin)
-                .get();
-        assertEquals(200, response.getStatus());
-
-        assertEquals(expectedHeader, response.getHeaderString("Access-Control-Allow-Origin"));
-
-        response = request("/connector-plugins/FileStreamSource/validate")
-            .header("Referer", origin + "/page")
-            .header("Origin", origin)
-            .header("Access-Control-Request-Method", method)
-            .options();
-        assertEquals(404, response.getStatus());
-        assertEquals(expectedHeader, response.getHeaderString("Access-Control-Allow-Origin"));
-        assertEquals(method, response.getHeaderString("Access-Control-Allow-Methods"));
+        server.initializeServer();
+        server.initializeResources(herder);
+        HttpRequest request = new HttpGet("/connectors");
+        request.addHeader("Referer", origin + "/page");
+        request.addHeader("Origin", origin);
+        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());
+
+        if (expectedHeader != null) {
+            Assert.assertEquals(expectedHeader,
+                response.getFirstHeader("Access-Control-Allow-Origin").getValue());
+        }
+
+        request = new HttpOptions("/connector-plugins/FileStreamSource/validate");
+        request.addHeader("Referer", origin + "/page");
+        request.addHeader("Origin", origin);
+        request.addHeader("Access-Control-Request-Method", method);
+        response = httpClient.execute(httpHost, request);
+        Assert.assertEquals(404, response.getStatusLine().getStatusCode());
+        if (expectedHeader != null) {
+            Assert.assertEquals(expectedHeader,
+                response.getFirstHeader("Access-Control-Allow-Origin").getValue());
+        }
+        if (method != null) {
+            Assert.assertEquals(method,
+                response.getFirstHeader("Access-Control-Allow-Methods").getValue());
+        }
         PowerMock.verifyAll();
     }
 
-    protected Invocation.Builder request(String path) {
-        return request(path, null, null, null);
+    @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());
     }
 
-    protected Invocation.Builder request(String path, Map queryParams) {
-        return request(path, null, null, queryParams);
+    @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");
     }
 
-    protected Invocation.Builder request(String path, String templateName, Object templateValue) {
-        return request(path, templateName, templateValue, null);
+    @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());
     }
 
-    protected Invocation.Builder request(String path, String templateName, Object templateValue,
-                                         Map queryParams) {
-        Client client = ClientBuilder.newClient();
-        WebTarget target;
-        URI pathUri = null;
+    @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());
+    }
+
+    @Test
+    public void testValidCustomizedHttpResponseHeaders() throws IOException  {
+        String headerConfig =
+                "add X-XSS-Protection: 1; mode=block, \"add Cache-Control: no-cache, no-store, must-revalidate\"";
+        Map expectedHeaders = new HashMap<>();
+        expectedHeaders.put("X-XSS-Protection", "1; mode=block");
+        expectedHeaders.put("Cache-Control", "no-cache, no-store, must-revalidate");
+        checkCustomizedHttpResponseHeaders(headerConfig, expectedHeaders);
+    }
+
+    @Test
+    public void testDefaultCustomizedHttpResponseHeaders() throws IOException  {
+        String headerConfig = "";
+        Map expectedHeaders = new HashMap<>();
+        checkCustomizedHttpResponseHeaders(headerConfig, expectedHeaders);
+    }
+
+    private void checkCustomizedHttpResponseHeaders(String headerConfig, Map expectedHeaders)
+            throws IOException  {
+        Map workerProps = baseWorkerProps();
+        workerProps.put("offset.storage.file.filename", "/tmp");
+        workerProps.put(WorkerConfig.RESPONSE_HTTP_HEADERS_CONFIG, headerConfig);
+        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());
+
+        EasyMock.expect(herder.connectors()).andReturn(Arrays.asList("a", "b"));
+
+        PowerMock.replayAll();
+
+        server = new RestServer(workerConfig);
         try {
-            pathUri = new URI(path);
-        } catch (URISyntaxException e) {
-            // Ignore, use restConnect and assume this is a valid path part
-        }
-        if (pathUri != null && pathUri.isAbsolute()) {
-            target = client.target(path);
-        } else {
-            target = client.target(server.advertisedUrl()).path(path);
-        }
-        if (templateName != null && templateValue != null) {
-            target = target.resolveTemplate(templateName, templateValue);
-        }
-        if (queryParams != null) {
-            for (Map.Entry queryParam : queryParams.entrySet()) {
-                target = target.queryParam(queryParam.getKey(), queryParam.getValue());
+            server.initializeServer();
+            server.initializeResources(herder);
+            HttpRequest request = new HttpGet("/connectors");
+            try (CloseableHttpClient httpClient = HttpClients.createMinimal()) {
+                HttpHost httpHost = new HttpHost(server.advertisedUrl().getHost(), server.advertisedUrl().getPort());
+                try (CloseableHttpResponse response = httpClient.execute(httpHost, request)) {
+                    Assert.assertEquals(200, response.getStatusLine().getStatusCode());
+                    if (!headerConfig.isEmpty()) {
+                        expectedHeaders.forEach((k, v) ->
+                                Assert.assertEquals(response.getFirstHeader(k).getValue(), v));
+                    } else {
+                        Assert.assertNull(response.getFirstHeader("X-Frame-Options"));
+                    }
+                }
             }
+        } finally {
+            server.stop();
+            server = null;
         }
-        return target.request();
+    }
+
+    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 ff9043e87fd58..c3e5d82171641 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;
@@ -32,10 +33,11 @@
 import org.apache.kafka.connect.runtime.Herder;
 import org.apache.kafka.connect.runtime.TestSinkConnector;
 import org.apache.kafka.connect.runtime.TestSourceConnector;
+import org.apache.kafka.connect.runtime.WorkerConfig;
 import org.apache.kafka.connect.runtime.isolation.PluginClassLoader;
 import org.apache.kafka.connect.runtime.isolation.PluginDesc;
 import org.apache.kafka.connect.runtime.isolation.Plugins;
-import org.apache.kafka.connect.runtime.rest.RestServer;
+import org.apache.kafka.connect.runtime.rest.RestClient;
 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.ConfigKeyInfo;
@@ -50,6 +52,8 @@
 import org.apache.kafka.connect.tools.SchemaSourceConnector;
 import org.apache.kafka.connect.tools.VerifiableSinkConnector;
 import org.apache.kafka.connect.tools.VerifiableSourceConnector;
+import org.apache.kafka.connect.util.Callback;
+import org.easymock.Capture;
 import org.easymock.EasyMock;
 import org.easymock.IAnswer;
 import org.junit.Before;
@@ -78,7 +82,7 @@
 import static org.junit.Assert.assertTrue;
 
 @RunWith(PowerMockRunner.class)
-@PrepareForTest(RestServer.class)
+@PrepareForTest(RestClient.class)
 @PowerMockIgnore("javax.management.*")
 public class ConnectorPluginsResourceTest {
 
@@ -158,11 +162,14 @@ public class ConnectorPluginsResourceTest {
 
         try {
             for (Class klass : abstractConnectorClasses) {
-                CONNECTOR_PLUGINS.add(
-                        new MockConnectorPluginDesc((Class) klass, "0.0.0"));
+                @SuppressWarnings("unchecked")
+                MockConnectorPluginDesc pluginDesc = new MockConnectorPluginDesc((Class) klass, "0.0.0");
+                CONNECTOR_PLUGINS.add(pluginDesc);
             }
             for (Class klass : connectorClasses) {
-                CONNECTOR_PLUGINS.add(new MockConnectorPluginDesc((Class) klass));
+                @SuppressWarnings("unchecked")
+                MockConnectorPluginDesc pluginDesc = new MockConnectorPluginDesc((Class) klass);
+                CONNECTOR_PLUGINS.add(pluginDesc);
             }
         } catch (Exception e) {
             throw new RuntimeException(e);
@@ -177,8 +184,8 @@ public class ConnectorPluginsResourceTest {
 
     @Before
     public void setUp() throws Exception {
-        PowerMock.mockStatic(RestServer.class,
-                             RestServer.class.getMethod("httpRequest", String.class, String.class, Object.class, TypeReference.class));
+        PowerMock.mockStatic(RestClient.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);
@@ -193,31 +200,31 @@ private void expectPlugins() {
 
     @Test
     public void testValidateConfigWithSingleErrorDueToMissingConnectorClassname() throws Throwable {
-        herder.validateConnectorConfig(EasyMock.eq(partialProps));
-
-        PowerMock.expectLastCall().andAnswer(new IAnswer() {
-            @Override
-            public ConfigInfos answer() {
-                ConfigDef connectorConfigDef = ConnectorConfig.configDef();
-                List connectorConfigValues = connectorConfigDef.validate(partialProps);
-
-                Connector connector = new ConnectorPluginsResourceTestConnector();
-                Config config = connector.validate(partialProps);
-                ConfigDef configDef = connector.config();
-                Map configKeys = configDef.configKeys();
-                List configValues = config.configValues();
-
-                Map resultConfigKeys = new HashMap<>(configKeys);
-                resultConfigKeys.putAll(connectorConfigDef.configKeys());
-                configValues.addAll(connectorConfigValues);
-
-                return AbstractHerder.generateResult(
-                    ConnectorPluginsResourceTestConnector.class.getName(),
-                    resultConfigKeys,
-                    configValues,
-                    Collections.singletonList("Test")
-                );
-            }
+        Capture> configInfosCallback = EasyMock.newCapture();
+        herder.validateConnectorConfig(EasyMock.eq(partialProps), EasyMock.capture(configInfosCallback), EasyMock.anyBoolean());
+
+        PowerMock.expectLastCall().andAnswer((IAnswer) () -> {
+            ConfigDef connectorConfigDef = ConnectorConfig.configDef();
+            List connectorConfigValues = connectorConfigDef.validate(partialProps);
+
+            Connector connector = new ConnectorPluginsResourceTestConnector();
+            Config config = connector.validate(partialProps);
+            ConfigDef configDef = connector.config();
+            Map configKeys = configDef.configKeys();
+            List configValues = config.configValues();
+
+            Map resultConfigKeys = new HashMap<>(configKeys);
+            resultConfigKeys.putAll(connectorConfigDef.configKeys());
+            configValues.addAll(connectorConfigValues);
+
+            ConfigInfos configInfos = AbstractHerder.generateResult(
+                ConnectorPluginsResourceTestConnector.class.getName(),
+                resultConfigKeys,
+                configValues,
+                Collections.singletonList("Test")
+            );
+            configInfosCallback.getValue().onCompletion(null, configInfos);
+            return null;
         });
 
         PowerMock.replayAll();
@@ -241,31 +248,31 @@ public ConfigInfos answer() {
 
     @Test
     public void testValidateConfigWithSimpleName() throws Throwable {
-        herder.validateConnectorConfig(EasyMock.eq(props));
+        Capture> configInfosCallback = EasyMock.newCapture();
+        herder.validateConnectorConfig(EasyMock.eq(props), EasyMock.capture(configInfosCallback), EasyMock.anyBoolean());
 
-        PowerMock.expectLastCall().andAnswer(new IAnswer() {
-            @Override
-            public ConfigInfos answer() {
-                ConfigDef connectorConfigDef = ConnectorConfig.configDef();
-                List connectorConfigValues = connectorConfigDef.validate(props);
+        PowerMock.expectLastCall().andAnswer((IAnswer) () -> {
+            ConfigDef connectorConfigDef = ConnectorConfig.configDef();
+            List connectorConfigValues = connectorConfigDef.validate(props);
 
-                Connector connector = new ConnectorPluginsResourceTestConnector();
-                Config config = connector.validate(props);
-                ConfigDef configDef = connector.config();
-                Map configKeys = configDef.configKeys();
-                List configValues = config.configValues();
+            Connector connector = new ConnectorPluginsResourceTestConnector();
+            Config config = connector.validate(props);
+            ConfigDef configDef = connector.config();
+            Map configKeys = configDef.configKeys();
+            List configValues = config.configValues();
 
-                Map resultConfigKeys = new HashMap<>(configKeys);
-                resultConfigKeys.putAll(connectorConfigDef.configKeys());
-                configValues.addAll(connectorConfigValues);
+            Map resultConfigKeys = new HashMap<>(configKeys);
+            resultConfigKeys.putAll(connectorConfigDef.configKeys());
+            configValues.addAll(connectorConfigValues);
 
-                return AbstractHerder.generateResult(
+            ConfigInfos configInfos = AbstractHerder.generateResult(
                     ConnectorPluginsResourceTestConnector.class.getName(),
                     resultConfigKeys,
                     configValues,
                     Collections.singletonList("Test")
-                );
-            }
+            );
+            configInfosCallback.getValue().onCompletion(null, configInfos);
+            return null;
         });
 
         PowerMock.replayAll();
@@ -285,31 +292,31 @@ public ConfigInfos answer() {
 
     @Test
     public void testValidateConfigWithAlias() throws Throwable {
-        herder.validateConnectorConfig(EasyMock.eq(props));
+        Capture> configInfosCallback = EasyMock.newCapture();
+        herder.validateConnectorConfig(EasyMock.eq(props), EasyMock.capture(configInfosCallback), EasyMock.anyBoolean());
 
-        PowerMock.expectLastCall().andAnswer(new IAnswer() {
-            @Override
-            public ConfigInfos answer() {
-                ConfigDef connectorConfigDef = ConnectorConfig.configDef();
-                List connectorConfigValues = connectorConfigDef.validate(props);
+        PowerMock.expectLastCall().andAnswer((IAnswer) () -> {
+            ConfigDef connectorConfigDef = ConnectorConfig.configDef();
+            List connectorConfigValues = connectorConfigDef.validate(props);
 
-                Connector connector = new ConnectorPluginsResourceTestConnector();
-                Config config = connector.validate(props);
-                ConfigDef configDef = connector.config();
-                Map configKeys = configDef.configKeys();
-                List configValues = config.configValues();
+            Connector connector = new ConnectorPluginsResourceTestConnector();
+            Config config = connector.validate(props);
+            ConfigDef configDef = connector.config();
+            Map configKeys = configDef.configKeys();
+            List configValues = config.configValues();
 
-                Map resultConfigKeys = new HashMap<>(configKeys);
-                resultConfigKeys.putAll(connectorConfigDef.configKeys());
-                configValues.addAll(connectorConfigValues);
+            Map resultConfigKeys = new HashMap<>(configKeys);
+            resultConfigKeys.putAll(connectorConfigDef.configKeys());
+            configValues.addAll(connectorConfigValues);
 
-                return AbstractHerder.generateResult(
+            ConfigInfos configInfos = AbstractHerder.generateResult(
                     ConnectorPluginsResourceTestConnector.class.getName(),
                     resultConfigKeys,
                     configValues,
                     Collections.singletonList("Test")
-                );
-            }
+            );
+            configInfosCallback.getValue().onCompletion(null, configInfos);
+            return null;
         });
 
         PowerMock.replayAll();
@@ -329,31 +336,30 @@ public ConfigInfos answer() {
 
     @Test(expected = BadRequestException.class)
     public void testValidateConfigWithNonExistentName() throws Throwable {
-        herder.validateConnectorConfig(EasyMock.eq(props));
-
-        PowerMock.expectLastCall().andAnswer(new IAnswer() {
-            @Override
-            public ConfigInfos answer() {
-                ConfigDef connectorConfigDef = ConnectorConfig.configDef();
-                List connectorConfigValues = connectorConfigDef.validate(props);
-
-                Connector connector = new ConnectorPluginsResourceTestConnector();
-                Config config = connector.validate(props);
-                ConfigDef configDef = connector.config();
-                Map configKeys = configDef.configKeys();
-                List configValues = config.configValues();
-
-                Map resultConfigKeys = new HashMap<>(configKeys);
-                resultConfigKeys.putAll(connectorConfigDef.configKeys());
-                configValues.addAll(connectorConfigValues);
-
-                return AbstractHerder.generateResult(
+        Capture> configInfosCallback = EasyMock.newCapture();
+        herder.validateConnectorConfig(EasyMock.eq(props), EasyMock.capture(configInfosCallback), EasyMock.anyBoolean());
+        PowerMock.expectLastCall().andAnswer((IAnswer) () -> {
+            ConfigDef connectorConfigDef = ConnectorConfig.configDef();
+            List connectorConfigValues = connectorConfigDef.validate(props);
+
+            Connector connector = new ConnectorPluginsResourceTestConnector();
+            Config config = connector.validate(props);
+            ConfigDef configDef = connector.config();
+            Map configKeys = configDef.configKeys();
+            List configValues = config.configValues();
+
+            Map resultConfigKeys = new HashMap<>(configKeys);
+            resultConfigKeys.putAll(connectorConfigDef.configKeys());
+            configValues.addAll(connectorConfigValues);
+
+            ConfigInfos configInfos = AbstractHerder.generateResult(
                     ConnectorPluginsResourceTestConnector.class.getName(),
                     resultConfigKeys,
                     configValues,
                     Collections.singletonList("Test")
-                );
-            }
+            );
+            configInfosCallback.getValue().onCompletion(null, configInfos);
+            return null;
         });
 
         PowerMock.replayAll();
@@ -369,31 +375,31 @@ public ConfigInfos answer() {
 
     @Test(expected = BadRequestException.class)
     public void testValidateConfigWithNonExistentAlias() throws Throwable {
-        herder.validateConnectorConfig(EasyMock.eq(props));
+        Capture> configInfosCallback = EasyMock.newCapture();
+        herder.validateConnectorConfig(EasyMock.eq(props), EasyMock.capture(configInfosCallback), EasyMock.anyBoolean());
 
-        PowerMock.expectLastCall().andAnswer(new IAnswer() {
-            @Override
-            public ConfigInfos answer() {
-                ConfigDef connectorConfigDef = ConnectorConfig.configDef();
-                List connectorConfigValues = connectorConfigDef.validate(props);
+        PowerMock.expectLastCall().andAnswer((IAnswer) () -> {
+            ConfigDef connectorConfigDef = ConnectorConfig.configDef();
+            List connectorConfigValues = connectorConfigDef.validate(props);
 
-                Connector connector = new ConnectorPluginsResourceTestConnector();
-                Config config = connector.validate(props);
-                ConfigDef configDef = connector.config();
-                Map configKeys = configDef.configKeys();
-                List configValues = config.configValues();
+            Connector connector = new ConnectorPluginsResourceTestConnector();
+            Config config = connector.validate(props);
+            ConfigDef configDef = connector.config();
+            Map configKeys = configDef.configKeys();
+            List configValues = config.configValues();
 
-                Map resultConfigKeys = new HashMap<>(configKeys);
-                resultConfigKeys.putAll(connectorConfigDef.configKeys());
-                configValues.addAll(connectorConfigValues);
+            Map resultConfigKeys = new HashMap<>(configKeys);
+            resultConfigKeys.putAll(connectorConfigDef.configKeys());
+            configValues.addAll(connectorConfigValues);
 
-                return AbstractHerder.generateResult(
+            ConfigInfos configInfos = AbstractHerder.generateResult(
                     ConnectorPluginsResourceTestConnector.class.getName(),
                     resultConfigKeys,
                     configValues,
                     Collections.singletonList("Test")
-                );
-            }
+            );
+            configInfosCallback.getValue().onCompletion(null, configInfos);
+            return null;
         });
 
         PowerMock.replayAll();
@@ -489,7 +495,7 @@ public MockConnectorPluginDesc(Class klass, String version)
         public MockConnectorPluginDesc(Class klass) throws Exception {
             super(
                     klass,
-                    klass.newInstance().version(),
+                    klass.getConstructor().newInstance().version(),
                     new MockPluginClassLoader(null, new URL[0])
             );
         }
@@ -545,7 +551,7 @@ private static class IntegerRecommender implements Recommender {
 
         @Override
         public List validValues(String name, Map parsedConfig) {
-            return Arrays.asList(1, 2, 3);
+            return Arrays.asList(1, 2, 3);
         }
 
         @Override
@@ -557,7 +563,7 @@ public boolean visible(String name, Map parsedConfig) {
     private static class ListRecommender implements Recommender {
         @Override
         public List validValues(String name, Map parsedConfig) {
-            return Arrays.asList("a", "b", "c");
+            return Arrays.asList("a", "b", "c");
         }
 
         @Override
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 cb86143c207a7..677f830e78384 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,18 +18,26 @@
 
 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.WorkerConfig;
 import org.apache.kafka.connect.runtime.distributed.NotAssignedException;
 import org.apache.kafka.connect.runtime.distributed.NotLeaderException;
-import org.apache.kafka.connect.runtime.rest.RestServer;
+import org.apache.kafka.connect.runtime.rest.InternalRequestSignature;
+import org.apache.kafka.connect.runtime.rest.RestClient;
+import org.apache.kafka.connect.runtime.rest.entities.ActiveTopicsInfo;
 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;
@@ -46,31 +54,47 @@
 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.Response;
+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;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
+import static org.apache.kafka.connect.runtime.WorkerConfig.TOPIC_TRACKING_ALLOW_RESET_CONFIG;
+import static org.apache.kafka.connect.runtime.WorkerConfig.TOPIC_TRACKING_ENABLE_CONFIG;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
 
 @RunWith(PowerMockRunner.class)
-@PrepareForTest(RestServer.class)
-@PowerMockIgnore("javax.management.*")
+@PrepareForTest(RestClient.class)
+@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
     // URL construction properly, avoiding //, which will mess up routing in the REST server
     private static final String LEADER_URL = "http://leader:8083/";
     private static final String CONNECTOR_NAME = "test";
-    private static final String CONNECTOR_NAME_SPECIAL_CHARS = "t?a=b&c=d\rx=1.1\n>< \t`'\" x%y+z!#$&'()*+,:;=?@[]";
+    private static final String CONNECTOR_NAME_SPECIAL_CHARS = "ta/b&c=d//\\rx=1þ.1>< `'\" x%y+z!ሴ#$&'(æ)*+,:;=?ñ@[]ÿ";
+    private static final String CONNECTOR_NAME_CONTROL_SEQUENCES1 = "ta/b&c=drx=1\n.1>< `'\" x%y+z!#$&'()*+,:;=?@[]";
     private static final String CONNECTOR2_NAME = "test2";
+    private static final String CONNECTOR_NAME_ALL_WHITESPACES = "   \t\n  \b";
+    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");
@@ -81,6 +105,24 @@ public class ConnectorsResourceTest {
         CONNECTOR_CONFIG.put("name", CONNECTOR_NAME);
         CONNECTOR_CONFIG.put("sample_config", "test_config");
     }
+
+    private static final Map CONNECTOR_CONFIG_CONTROL_SEQUENCES = new HashMap<>();
+    static {
+        CONNECTOR_CONFIG_CONTROL_SEQUENCES.put("name", CONNECTOR_NAME_CONTROL_SEQUENCES1);
+        CONNECTOR_CONFIG_CONTROL_SEQUENCES.put("sample_config", "test_config");
+    }
+
+    private static final Map CONNECTOR_CONFIG_WITHOUT_NAME = new HashMap<>();
+    static {
+        CONNECTOR_CONFIG_WITHOUT_NAME.put("sample_config", "test_config");
+    }
+
+    private static final Map CONNECTOR_CONFIG_WITH_EMPTY_NAME = new HashMap<>();
+
+    static {
+        CONNECTOR_CONFIG_WITH_EMPTY_NAME.put(ConnectorConfig.NAME_CONFIG, "");
+        CONNECTOR_CONFIG_WITH_EMPTY_NAME.put("sample_config", "test_config");
+    }
     private static final List CONNECTOR_TASK_NAMES = Arrays.asList(
             new ConnectorTaskId(CONNECTOR_NAME, 0),
             new ConnectorTaskId(CONNECTOR_NAME, 1)
@@ -96,26 +138,46 @@ public class ConnectorsResourceTest {
         TASK_INFOS.add(new TaskInfo(new ConnectorTaskId(CONNECTOR_NAME, 1), TASK_CONFIGS.get(1)));
     }
 
+    private static final Set CONNECTOR_ACTIVE_TOPICS = new HashSet<>(
+            Arrays.asList("foo_topic", "bar_topic"));
+    private static final Set CONNECTOR2_ACTIVE_TOPICS = new HashSet<>(
+            Arrays.asList("foo_topic", "baz_topic"));
+
     @Mock
     private Herder herder;
     private ConnectorsResource connectorsResource;
+    private UriInfo forward;
+    @Mock
+    private WorkerConfig workerConfig;
 
     @Before
     public void setUp() throws NoSuchMethodException {
-        PowerMock.mockStatic(RestServer.class,
-                RestServer.class.getMethod("httpRequest", String.class, String.class, Object.class, TypeReference.class));
-        connectorsResource = new ConnectorsResource(herder);
+        PowerMock.mockStatic(RestClient.class,
+                RestClient.class.getMethod("httpRequest", String.class, String.class, HttpHeaders.class, Object.class, TypeReference.class, WorkerConfig.class));
+        EasyMock.expect(workerConfig.getBoolean(TOPIC_TRACKING_ENABLE_CONFIG)).andReturn(true);
+        EasyMock.expect(workerConfig.getBoolean(TOPIC_TRACKING_ALLOW_RESET_CONFIG)).andReturn(true);
+        PowerMock.replay(workerConfig);
+        connectorsResource = new ConnectorsResource(herder, workerConfig);
+        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) {
+        Map result = new HashMap<>(mapToClone);
+        return result;
     }
 
     @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));
 
@@ -123,36 +185,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(RestServer.httpRequest(EasyMock.eq("http://leader:8083/connectors?forward=false"), EasyMock.eq("GET"),
-                EasyMock.isNull(), EasyMock.anyObject(TypeReference.class)))
-                .andReturn(new RestServer.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));
@@ -164,7 +297,7 @@ public void testCreateConnector() throws Throwable {
 
         PowerMock.replayAll();
 
-        connectorsResource.createConnector(FORWARD, body);
+        connectorsResource.createConnector(FORWARD, NULL_HEADERS, body);
 
         PowerMock.verifyAll();
     }
@@ -177,17 +310,55 @@ 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(RestServer.httpRequest(EasyMock.eq("http://leader:8083/connectors?forward=false"), EasyMock.eq("POST"), EasyMock.eq(body), EasyMock.anyObject()))
-                .andReturn(new RestServer.HttpResponse<>(201, new HashMap>(), new ConnectorInfo(CONNECTOR_NAME, CONNECTOR_CONFIG, CONNECTOR_TASK_NAMES,
+        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)
@@ -200,25 +371,64 @@ public void testCreateConnectorExists() throws Throwable {
 
         PowerMock.replayAll();
 
-        connectorsResource.createConnector(FORWARD, body);
+        connectorsResource.createConnector(FORWARD, NULL_HEADERS, body);
 
         PowerMock.verifyAll();
     }
 
-    @Test(expected = BadRequestException.class)
-    public void testCreateConnectorWithASlashInItsName() throws Throwable {
-        String badConnectorName = CONNECTOR_NAME + "/" + "test";
+    @Test
+    public void testCreateConnectorNameTrimWhitespaces() throws Throwable {
+        // Clone CONNECTOR_CONFIG_WITHOUT_NAME Map, as createConnector changes it (puts the name in it) and this
+        // will affect later tests
+        Map inputConfig = getConnectorConfig(CONNECTOR_CONFIG_WITHOUT_NAME);
+        final CreateConnectorRequest bodyIn = new CreateConnectorRequest(CONNECTOR_NAME_PADDING_WHITESPACES, inputConfig);
+        final CreateConnectorRequest bodyOut = new CreateConnectorRequest(CONNECTOR_NAME, CONNECTOR_CONFIG);
 
-        CreateConnectorRequest body = new CreateConnectorRequest(badConnectorName, Collections.singletonMap(ConnectorConfig.NAME_CONFIG, badConnectorName));
+        final Capture>> cb = Capture.newInstance();
+        herder.putConnectorConfig(EasyMock.eq(bodyOut.name()), EasyMock.eq(bodyOut.config()), EasyMock.eq(false), EasyMock.capture(cb));
+        expectAndCallbackResult(cb, new Herder.Created<>(true, new ConnectorInfo(bodyOut.name(), bodyOut.config(), CONNECTOR_TASK_NAMES, ConnectorType.SOURCE)));
+
+        PowerMock.replayAll();
+
+        connectorsResource.createConnector(FORWARD, NULL_HEADERS, bodyIn);
+
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testCreateConnectorNameAllWhitespaces() throws Throwable {
+        // Clone CONNECTOR_CONFIG_WITHOUT_NAME Map, as createConnector changes it (puts the name in it) and this
+        // will affect later tests
+        Map inputConfig = getConnectorConfig(CONNECTOR_CONFIG_WITHOUT_NAME);
+        final CreateConnectorRequest bodyIn = new CreateConnectorRequest(CONNECTOR_NAME_ALL_WHITESPACES, inputConfig);
+        final CreateConnectorRequest bodyOut = new CreateConnectorRequest("", CONNECTOR_CONFIG_WITH_EMPTY_NAME);
 
         final Capture>> cb = Capture.newInstance();
-        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)));
+        herder.putConnectorConfig(EasyMock.eq(bodyOut.name()), EasyMock.eq(bodyOut.config()), EasyMock.eq(false), EasyMock.capture(cb));
+        expectAndCallbackResult(cb, new Herder.Created<>(true, new ConnectorInfo(bodyOut.name(), bodyOut.config(), CONNECTOR_TASK_NAMES, ConnectorType.SOURCE)));
 
         PowerMock.replayAll();
 
-        connectorsResource.createConnector(FORWARD, body);
+        connectorsResource.createConnector(FORWARD, NULL_HEADERS, bodyIn);
+
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testCreateConnectorNoName() throws Throwable {
+        // Clone CONNECTOR_CONFIG_WITHOUT_NAME Map, as createConnector changes it (puts the name in it) and this
+        // will affect later tests
+        Map inputConfig = getConnectorConfig(CONNECTOR_CONFIG_WITHOUT_NAME);
+        final CreateConnectorRequest bodyIn = new CreateConnectorRequest(null, inputConfig);
+        final CreateConnectorRequest bodyOut = new CreateConnectorRequest("", CONNECTOR_CONFIG_WITH_EMPTY_NAME);
+
+        final Capture>> cb = Capture.newInstance();
+        herder.putConnectorConfig(EasyMock.eq(bodyOut.name()), EasyMock.eq(bodyOut.config()), EasyMock.eq(false), EasyMock.capture(cb));
+        expectAndCallbackResult(cb, new Herder.Created<>(true, new ConnectorInfo(bodyOut.name(), bodyOut.config(), CONNECTOR_TASK_NAMES, ConnectorType.SOURCE)));
+
+        PowerMock.replayAll();
+
+        connectorsResource.createConnector(FORWARD, NULL_HEADERS, bodyIn);
 
         PowerMock.verifyAll();
     }
@@ -231,7 +441,7 @@ public void testDeleteConnector() throws Throwable {
 
         PowerMock.replayAll();
 
-        connectorsResource.destroyConnector(CONNECTOR_NAME, FORWARD);
+        connectorsResource.destroyConnector(CONNECTOR_NAME, NULL_HEADERS, FORWARD);
 
         PowerMock.verifyAll();
     }
@@ -242,12 +452,12 @@ public void testDeleteConnectorNotLeader() throws Throwable {
         herder.deleteConnectorConfig(EasyMock.eq(CONNECTOR_NAME), EasyMock.capture(cb));
         expectAndCallbackNotLeaderException(cb);
         // Should forward request
-        EasyMock.expect(RestServer.httpRequest("http://leader:8083/connectors/" + CONNECTOR_NAME + "?forward=false", "DELETE", null, null))
-                .andReturn(new RestServer.HttpResponse<>(204, new HashMap>(), null));
+        EasyMock.expect(RestClient.httpRequest("http://leader:8083/connectors/" + CONNECTOR_NAME + "?forward=false", "DELETE", NULL_HEADERS, null, null, workerConfig))
+                .andReturn(new RestClient.HttpResponse<>(204, new HashMap<>(), null));
 
         PowerMock.replayAll();
 
-        connectorsResource.destroyConnector(CONNECTOR_NAME, FORWARD);
+        connectorsResource.destroyConnector(CONNECTOR_NAME, NULL_HEADERS, FORWARD);
 
         PowerMock.verifyAll();
     }
@@ -261,7 +471,7 @@ public void testDeleteConnectorNotFound() throws Throwable {
 
         PowerMock.replayAll();
 
-        connectorsResource.destroyConnector(CONNECTOR_NAME, FORWARD);
+        connectorsResource.destroyConnector(CONNECTOR_NAME, NULL_HEADERS, FORWARD);
 
         PowerMock.verifyAll();
     }
@@ -275,7 +485,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);
 
@@ -290,7 +500,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();
@@ -304,7 +514,7 @@ public void testGetConnectorConfigConnectorNotFound() throws Throwable {
 
         PowerMock.replayAll();
 
-        connectorsResource.getConnectorConfig(CONNECTOR_NAME, FORWARD);
+        connectorsResource.getConnectorConfig(CONNECTOR_NAME, NULL_HEADERS, FORWARD);
 
         PowerMock.verifyAll();
     }
@@ -318,7 +528,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();
     }
@@ -334,13 +544,31 @@ 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);
 
         PowerMock.verifyAll();
     }
 
+    @Test
+    public void testCreateConnectorWithControlSequenceInName() throws Throwable {
+        CreateConnectorRequest body = new CreateConnectorRequest(CONNECTOR_NAME_CONTROL_SEQUENCES1, Collections.singletonMap(ConnectorConfig.NAME_CONFIG, CONNECTOR_NAME_CONTROL_SEQUENCES1));
+
+        final Capture>> cb = Capture.newInstance();
+        herder.putConnectorConfig(EasyMock.eq(CONNECTOR_NAME_CONTROL_SEQUENCES1), EasyMock.eq(body.config()), EasyMock.eq(false), EasyMock.capture(cb));
+        expectAndCallbackResult(cb, new Herder.Created<>(true, new ConnectorInfo(CONNECTOR_NAME_CONTROL_SEQUENCES1, CONNECTOR_CONFIG,
+                CONNECTOR_TASK_NAMES, ConnectorType.SOURCE)));
+
+        PowerMock.replayAll();
+
+        String rspLocation = connectorsResource.createConnector(FORWARD, NULL_HEADERS, body).getLocation().toString();
+        String decoded = new URI(rspLocation).getPath();
+        Assert.assertEquals("/connectors/" + CONNECTOR_NAME_CONTROL_SEQUENCES1, decoded);
+
+        PowerMock.verifyAll();
+    }
+
     @Test
     public void testPutConnectorConfigWithSpecialCharsInName() throws Throwable {
         final Capture>> cb = Capture.newInstance();
@@ -351,18 +579,43 @@ 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);
 
         PowerMock.verifyAll();
     }
 
+    @Test
+    public void testPutConnectorConfigWithControlSequenceInName() throws Throwable {
+        final Capture>> cb = Capture.newInstance();
+
+        herder.putConnectorConfig(EasyMock.eq(CONNECTOR_NAME_CONTROL_SEQUENCES1), EasyMock.eq(CONNECTOR_CONFIG_CONTROL_SEQUENCES), EasyMock.eq(true), EasyMock.capture(cb));
+        expectAndCallbackResult(cb, new Herder.Created<>(true, new ConnectorInfo(CONNECTOR_NAME_CONTROL_SEQUENCES1, CONNECTOR_CONFIG_CONTROL_SEQUENCES, CONNECTOR_TASK_NAMES,
+                ConnectorType.SINK)));
+
+        PowerMock.replayAll();
+
+        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);
+
+        PowerMock.verifyAll();
+    }
+
     @Test(expected = BadRequestException.class)
     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)
+    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, NULL_HEADERS, request);
     }
 
     @Test
@@ -373,7 +626,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();
@@ -387,33 +640,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();
     }
@@ -426,7 +728,7 @@ public void testRestartConnectorNotFound() throws Throwable {
 
         PowerMock.replayAll();
 
-        connectorsResource.restartConnector(CONNECTOR_NAME, FORWARD);
+        connectorsResource.restartConnector(CONNECTOR_NAME, NULL_HEADERS, FORWARD);
 
         PowerMock.verifyAll();
     }
@@ -437,13 +739,13 @@ public void testRestartConnectorLeaderRedirect() throws Throwable {
         herder.restartConnector(EasyMock.eq(CONNECTOR_NAME), EasyMock.capture(cb));
         expectAndCallbackNotLeaderException(cb);
 
-        EasyMock.expect(RestServer.httpRequest(EasyMock.eq("http://leader:8083/connectors/" + CONNECTOR_NAME + "/restart?forward=true"),
-                EasyMock.eq("POST"), EasyMock.isNull(), EasyMock.anyObject()))
-                .andReturn(new RestServer.HttpResponse<>(202, new HashMap>(), null));
+        EasyMock.expect(RestClient.httpRequest(EasyMock.eq("http://leader:8083/connectors/" + CONNECTOR_NAME + "/restart?forward=true"),
+                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();
     }
@@ -455,13 +757,13 @@ public void testRestartConnectorOwnerRedirect() throws Throwable {
         String ownerUrl = "http://owner:8083";
         expectAndCallbackException(cb, new NotAssignedException("not owner test", ownerUrl));
 
-        EasyMock.expect(RestServer.httpRequest(EasyMock.eq("http://owner:8083/connectors/" + CONNECTOR_NAME + "/restart?forward=false"),
-                EasyMock.eq("POST"), EasyMock.isNull(), EasyMock.anyObject()))
-                .andReturn(new RestServer.HttpResponse<>(202, new HashMap>(), null));
+        EasyMock.expect(RestClient.httpRequest(EasyMock.eq("http://owner:8083/connectors/" + CONNECTOR_NAME + "/restart?forward=false"),
+                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();
     }
@@ -475,7 +777,7 @@ public void testRestartTaskNotFound() throws Throwable {
 
         PowerMock.replayAll();
 
-        connectorsResource.restartTask(CONNECTOR_NAME, 0, FORWARD);
+        connectorsResource.restartTask(CONNECTOR_NAME, 0, NULL_HEADERS, FORWARD);
 
         PowerMock.verifyAll();
     }
@@ -488,13 +790,13 @@ public void testRestartTaskLeaderRedirect() throws Throwable {
         herder.restartTask(EasyMock.eq(taskId), EasyMock.capture(cb));
         expectAndCallbackNotLeaderException(cb);
 
-        EasyMock.expect(RestServer.httpRequest(EasyMock.eq("http://leader:8083/connectors/" + CONNECTOR_NAME + "/tasks/0/restart?forward=true"),
-                EasyMock.eq("POST"), EasyMock.isNull(), EasyMock.anyObject()))
-                .andReturn(new RestServer.HttpResponse<>(202, new HashMap>(), null));
+        EasyMock.expect(RestClient.httpRequest(EasyMock.eq("http://leader:8083/connectors/" + CONNECTOR_NAME + "/tasks/0/restart?forward=true"),
+                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();
     }
@@ -508,17 +810,121 @@ public void testRestartTaskOwnerRedirect() throws Throwable {
         String ownerUrl = "http://owner:8083";
         expectAndCallbackException(cb, new NotAssignedException("not owner test", ownerUrl));
 
-        EasyMock.expect(RestServer.httpRequest(EasyMock.eq("http://owner:8083/connectors/" + CONNECTOR_NAME + "/tasks/0/restart?forward=false"),
-                EasyMock.eq("POST"), EasyMock.isNull(), EasyMock.anyObject()))
-                .andReturn(new RestServer.HttpResponse<>(202, new HashMap>(), null));
+        EasyMock.expect(RestClient.httpRequest(EasyMock.eq("http://owner:8083/connectors/" + CONNECTOR_NAME + "/tasks/0/restart?forward=false"),
+                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 testConnectorActiveTopicsWithTopicTrackingDisabled() {
+        PowerMock.reset(workerConfig);
+        EasyMock.expect(workerConfig.getBoolean(TOPIC_TRACKING_ENABLE_CONFIG)).andReturn(false);
+        EasyMock.expect(workerConfig.getBoolean(TOPIC_TRACKING_ALLOW_RESET_CONFIG)).andReturn(false);
+        PowerMock.replay(workerConfig);
+        connectorsResource = new ConnectorsResource(herder, workerConfig);
+        PowerMock.replayAll();
+
+        Exception e = assertThrows(ConnectRestException.class,
+            () -> connectorsResource.getConnectorActiveTopics(CONNECTOR_NAME));
+        assertEquals("Topic tracking is disabled.", e.getMessage());
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testResetConnectorActiveTopicsWithTopicTrackingDisabled() {
+        PowerMock.reset(workerConfig);
+        EasyMock.expect(workerConfig.getBoolean(TOPIC_TRACKING_ENABLE_CONFIG)).andReturn(false);
+        EasyMock.expect(workerConfig.getBoolean(TOPIC_TRACKING_ALLOW_RESET_CONFIG)).andReturn(true);
+        HttpHeaders headers = EasyMock.mock(HttpHeaders.class);
+        PowerMock.replay(workerConfig);
+        connectorsResource = new ConnectorsResource(herder, workerConfig);
+        PowerMock.replayAll();
+
+        Exception e = assertThrows(ConnectRestException.class,
+            () -> connectorsResource.resetConnectorActiveTopics(CONNECTOR_NAME, headers));
+        assertEquals("Topic tracking is disabled.", e.getMessage());
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testResetConnectorActiveTopicsWithTopicTrackingEnabled() {
+        PowerMock.reset(workerConfig);
+        EasyMock.expect(workerConfig.getBoolean(TOPIC_TRACKING_ENABLE_CONFIG)).andReturn(true);
+        EasyMock.expect(workerConfig.getBoolean(TOPIC_TRACKING_ALLOW_RESET_CONFIG)).andReturn(false);
+        HttpHeaders headers = EasyMock.mock(HttpHeaders.class);
+        PowerMock.replay(workerConfig);
+        connectorsResource = new ConnectorsResource(herder, workerConfig);
+        PowerMock.replayAll();
+
+        Exception e = assertThrows(ConnectRestException.class,
+            () -> connectorsResource.resetConnectorActiveTopics(CONNECTOR_NAME, headers));
+        assertEquals("Topic tracking reset is disabled.", e.getMessage());
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testConnectorActiveTopics() {
+        PowerMock.reset(workerConfig);
+        EasyMock.expect(workerConfig.getBoolean(TOPIC_TRACKING_ENABLE_CONFIG)).andReturn(true);
+        EasyMock.expect(workerConfig.getBoolean(TOPIC_TRACKING_ALLOW_RESET_CONFIG)).andReturn(true);
+        EasyMock.expect(herder.connectorActiveTopics(CONNECTOR_NAME))
+                .andReturn(new ActiveTopicsInfo(CONNECTOR_NAME, CONNECTOR_ACTIVE_TOPICS));
+        PowerMock.replay(workerConfig);
+        connectorsResource = new ConnectorsResource(herder, workerConfig);
+        PowerMock.replayAll();
+
+        Response response = connectorsResource.getConnectorActiveTopics(CONNECTOR_NAME);
+        assertEquals(Response.Status.OK.getStatusCode(), response.getStatus());
+        Map> body = (Map>) response.getEntity();
+        assertEquals(CONNECTOR_NAME, ((ActiveTopicsInfo) body.get(CONNECTOR_NAME)).connector());
+        assertEquals(new HashSet<>(CONNECTOR_ACTIVE_TOPICS),
+                ((ActiveTopicsInfo) body.get(CONNECTOR_NAME)).topics());
+        PowerMock.verifyAll();
+    }
+
+    @Test
+    public void testResetConnectorActiveTopics() {
+        PowerMock.reset(workerConfig);
+        EasyMock.expect(workerConfig.getBoolean(TOPIC_TRACKING_ENABLE_CONFIG)).andReturn(true);
+        EasyMock.expect(workerConfig.getBoolean(TOPIC_TRACKING_ALLOW_RESET_CONFIG)).andReturn(true);
+        HttpHeaders headers = EasyMock.mock(HttpHeaders.class);
+        herder.resetConnectorActiveTopics(CONNECTOR_NAME);
+        EasyMock.expectLastCall();
+        PowerMock.replay(workerConfig);
+        connectorsResource = new ConnectorsResource(herder, workerConfig);
+        PowerMock.replayAll();
+
+        Response response = connectorsResource.resetConnectorActiveTopics(CONNECTOR_NAME, headers);
+        assertEquals(Response.Status.ACCEPTED.getStatusCode(), response.getStatus());
+        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
new file mode 100644
index 0000000000000..4e928a370372f
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/rest/resources/RootResourceTest.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.connect.runtime.rest.resources;
+
+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.rest.entities.ServerInfo;
+import org.easymock.EasyMock;
+import org.easymock.EasyMockRunner;
+import org.easymock.EasyMockSupport;
+import org.easymock.Mock;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import static org.junit.Assert.assertEquals;
+
+@RunWith(EasyMockRunner.class)
+public class RootResourceTest extends EasyMockSupport {
+
+    @Mock
+    private Herder herder;
+    private RootResource rootResource;
+
+    @Before
+    public void setUp() {
+        rootResource = new RootResource(herder);
+    }
+
+    @Test
+    public void testRootGet() {
+        EasyMock.expect(herder.kafkaClusterId()).andReturn(MockAdminClient.DEFAULT_CLUSTER_ID);
+
+        replayAll();
+
+        ServerInfo info = rootResource.serverInfo();
+        assertEquals(AppInfoParser.getVersion(), info.version());
+        assertEquals(AppInfoParser.getCommitId(), info.commit());
+        assertEquals(MockAdminClient.DEFAULT_CLUSTER_ID, info.clusterId());
+
+        verifyAll();
+    }
+}
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
new file mode 100644
index 0000000000000..8959a6c6e0b1f
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/rest/util/SSLUtilsTest.java
@@ -0,0 +1,196 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.util;
+
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.common.config.SslConfigs;
+import org.apache.kafka.connect.runtime.WorkerConfig;
+import org.apache.kafka.connect.runtime.distributed.DistributedConfig;
+import org.apache.kafka.connect.runtime.standalone.StandaloneConfig;
+import org.eclipse.jetty.util.ssl.SslContextFactory;
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.Map;
+
+@SuppressWarnings("deprecation")
+public class SSLUtilsTest {
+    private static final Map DEFAULT_CONFIG = new HashMap<>();
+    static {
+        // The WorkerConfig base class has some required settings without defaults
+        DEFAULT_CONFIG.put(DistributedConfig.STATUS_STORAGE_TOPIC_CONFIG, "status-topic");
+        DEFAULT_CONFIG.put(DistributedConfig.CONFIG_TOPIC_CONFIG, "config-topic");
+        DEFAULT_CONFIG.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092");
+        DEFAULT_CONFIG.put(DistributedConfig.GROUP_ID_CONFIG, "connect-test-group");
+        DEFAULT_CONFIG.put(WorkerConfig.KEY_CONVERTER_CLASS_CONFIG, "org.apache.kafka.connect.json.JsonConverter");
+        DEFAULT_CONFIG.put(WorkerConfig.VALUE_CONVERTER_CLASS_CONFIG, "org.apache.kafka.connect.json.JsonConverter");
+        DEFAULT_CONFIG.put(WorkerConfig.INTERNAL_KEY_CONVERTER_CLASS_CONFIG, "org.apache.kafka.connect.json.JsonConverter");
+        DEFAULT_CONFIG.put(WorkerConfig.INTERNAL_VALUE_CONVERTER_CLASS_CONFIG, "org.apache.kafka.connect.json.JsonConverter");
+        DEFAULT_CONFIG.put(DistributedConfig.OFFSET_STORAGE_TOPIC_CONFIG, "connect-offsets");
+    }
+
+    @Test
+    public void testGetOrDefault() {
+        String existingKey = "exists";
+        String missingKey = "missing";
+        String value = "value";
+        String defaultValue = "default";
+        Map map = new HashMap<>();
+        map.put("exists", "value");
+
+        Assert.assertEquals(SSLUtils.getOrDefault(map, existingKey, defaultValue), value);
+        Assert.assertEquals(SSLUtils.getOrDefault(map, missingKey, defaultValue), defaultValue);
+    }
+
+    @Test
+    public void testCreateServerSideSslContextFactory() {
+        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.createServerSideSslContextFactory(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.assertTrue(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 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");
+        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.createClientSideSslContextFactory(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());
+    }
+}
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 18d2739f13ee1..837ae196c0254 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
@@ -20,12 +20,14 @@
 import org.apache.kafka.common.config.ConfigDef;
 import org.apache.kafka.common.config.ConfigValue;
 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;
 import org.apache.kafka.connect.runtime.AbstractStatus;
+import org.apache.kafka.connect.runtime.CloseableConnectorContext;
 import org.apache.kafka.connect.runtime.ConnectorConfig;
 import org.apache.kafka.connect.runtime.ConnectorStatus;
 import org.apache.kafka.connect.runtime.Herder;
@@ -36,7 +38,9 @@
 import org.apache.kafka.connect.runtime.TaskConfig;
 import org.apache.kafka.connect.runtime.TaskStatus;
 import org.apache.kafka.connect.runtime.Worker;
+import org.apache.kafka.connect.runtime.WorkerConfigTransformer;
 import org.apache.kafka.connect.runtime.WorkerConnector;
+import org.apache.kafka.connect.runtime.distributed.ClusterConfigState;
 import org.apache.kafka.connect.runtime.isolation.DelegatingClassLoader;
 import org.apache.kafka.connect.runtime.isolation.PluginClassLoader;
 import org.apache.kafka.connect.runtime.isolation.Plugins;
@@ -55,6 +59,7 @@
 import org.apache.kafka.connect.util.FutureCallback;
 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;
@@ -68,6 +73,7 @@
 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.concurrent.ExecutionException;
@@ -76,7 +82,11 @@
 import static java.util.Collections.singleton;
 import static java.util.Collections.singletonList;
 import static java.util.Collections.singletonMap;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.DEFAULT_TOPIC_CREATION_PREFIX;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.PARTITIONS_CONFIG;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.REPLICATION_FACTOR_CONFIG;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
@@ -89,6 +99,7 @@ public class StandaloneHerderTest {
     private static final String TOPICS_LIST_STR = "topic1,topic2";
     private static final int DEFAULT_MAX_TASKS = 1;
     private static final String WORKER_ID = "localhost:8083";
+    private static final String KAFKA_CLUSTER_ID = "I4ZmrWqfT2e-upky_4fdPA";
 
     private enum SourceSink {
         SOURCE, SINK
@@ -98,24 +109,32 @@ private enum SourceSink {
 
     private Connector connector;
     @Mock protected Worker worker;
+    @Mock protected WorkerConfigTransformer transformer;
     @Mock private Plugins plugins;
     @Mock
     private PluginClassLoader pluginLoader;
     @Mock
     private DelegatingClassLoader delegatingLoader;
-    @Mock protected Callback> createCallback;
+    protected FutureCallback> 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, statusBackingStore, new MemoryConfigBackingStore());
+        herder = PowerMock.createPartialMock(StandaloneHerder.class, new String[]{"connectorTypeForClass"/*, "validateConnectorConfig"*/},
+            worker, WORKER_ID, KAFKA_CLUSTER_ID, statusBackingStore, new MemoryConfigBackingStore(transformer), noneConnectorClientConfigOverridePolicy);
+        createCallback = new FutureCallback<>();
         plugins = PowerMock.createMock(Plugins.class);
         pluginLoader = PowerMock.createMock(PluginClassLoader.class);
         delegatingLoader = PowerMock.createMock(DelegatingClassLoader.class);
         PowerMock.mockStatic(Plugins.class);
         PowerMock.mockStatic(WorkerConnector.class);
+        Capture> configCapture = Capture.newInstance();
+        EasyMock.expect(transformer.transform(EasyMock.eq(CONNECTOR_NAME), EasyMock.capture(configCapture))).andAnswer(configCapture::getValue).anyTimes();
     }
 
     @Test
@@ -124,24 +143,30 @@ public void testCreateSourceConnector() throws Exception {
         expectAdd(SourceSink.SOURCE);
 
         Map config = connectorConfig(SourceSink.SOURCE);
-        expectConfigValidation(config);
+        Connector connectorMock = PowerMock.createMock(SourceConnector.class);
+        expectConfigValidation(connectorMock, true, config);
 
         PowerMock.replayAll();
 
         herder.putConnectorConfig(CONNECTOR_NAME, config, false, createCallback);
+        Herder.Created connectorInfo = createCallback.get(1000L, TimeUnit.SECONDS);
+        assertEquals(createdInfo(SourceSink.SOURCE), connectorInfo.result());
 
         PowerMock.verifyAll();
     }
 
-    @Test
-    public void testCreateConnectorFailedBasicValidation() throws Exception {
+    @Test(expected = BadRequestException.class)
+    public void testCreateConnectorFailedValidation() throws Throwable {
         // Basic validation should be performed and return an error, but should still evaluate the connector's config
         connector = PowerMock.createMock(BogusSourceConnector.class);
 
         Map config = connectorConfig(SourceSink.SOURCE);
         config.remove(ConnectorConfig.NAME_CONFIG);
 
-        Connector connectorMock = PowerMock.createMock(Connector.class);
+        Connector connectorMock = PowerMock.createMock(SourceConnector.class);
+        EasyMock.expect(worker.configTransformer()).andReturn(transformer).times(2);
+        final Capture> configCapture = EasyMock.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);
         EasyMock.expect(plugins.newConnector(EasyMock.anyString())).andReturn(connectorMock);
@@ -152,69 +177,56 @@ public void testCreateConnectorFailedBasicValidation() throws Exception {
         EasyMock.expect(connectorMock.validate(config)).andReturn(new Config(singletonList(validatedValue)));
         EasyMock.expect(Plugins.compareAndSwapLoaders(delegatingLoader)).andReturn(pluginLoader);
 
-        createCallback.onCompletion(EasyMock.anyObject(), EasyMock.>isNull());
-        PowerMock.expectLastCall();
-
-        PowerMock.replayAll();
-
-        herder.putConnectorConfig(CONNECTOR_NAME, config, false, createCallback);
-
-        PowerMock.verifyAll();
-    }
-
-    @Test
-    public void testCreateConnectorFailedCustomValidation() throws Exception {
-        connector = PowerMock.createMock(BogusSourceConnector.class);
-
-        Connector connectorMock = PowerMock.createMock(Connector.class);
-        EasyMock.expect(worker.getPlugins()).andReturn(plugins).times(3);
-        EasyMock.expect(plugins.compareAndSwapLoaders(connectorMock)).andReturn(delegatingLoader);
-        EasyMock.expect(plugins.newConnector(EasyMock.anyString())).andReturn(connectorMock);
-
-        ConfigDef configDef = new ConfigDef();
-        configDef.define("foo.bar", ConfigDef.Type.STRING, ConfigDef.Importance.HIGH, "foo.bar doc");
-        EasyMock.expect(connectorMock.config()).andReturn(configDef);
-
-        ConfigValue validatedValue = new ConfigValue("foo.bar");
-        validatedValue.addErrorMessage("Failed foo.bar validation");
-        Map config = connectorConfig(SourceSink.SOURCE);
-        EasyMock.expect(connectorMock.validate(config)).andReturn(new Config(singletonList(validatedValue)));
-        EasyMock.expect(Plugins.compareAndSwapLoaders(delegatingLoader)).andReturn(pluginLoader);
-
-        createCallback.onCompletion(EasyMock.anyObject(), EasyMock.>isNull());
-        PowerMock.expectLastCall();
-
         PowerMock.replayAll();
 
         herder.putConnectorConfig(CONNECTOR_NAME, config, false, createCallback);
 
-        PowerMock.verifyAll();
+        try {
+            createCallback.get(1000L, TimeUnit.SECONDS);
+        } catch (ExecutionException e) {
+            assertNotNull(e.getCause());
+            throw e.getCause();
+        } finally {
+            PowerMock.verifyAll();
+        }
     }
 
-    @Test
-    public void testCreateConnectorAlreadyExists() throws Exception {
+    @Test(expected = AlreadyExistsException.class)
+    public void testCreateConnectorAlreadyExists() throws Throwable {
         connector = PowerMock.createMock(BogusSourceConnector.class);
         // First addition should succeed
         expectAdd(SourceSink.SOURCE);
 
         Map config = connectorConfig(SourceSink.SOURCE);
-        Connector connectorMock = PowerMock.createMock(Connector.class);
+        Connector connectorMock = PowerMock.createMock(SourceConnector.class);
         expectConfigValidation(connectorMock, true, config, config);
 
+        EasyMock.expect(worker.configTransformer()).andReturn(transformer).times(2);
+        final Capture> configCapture = EasyMock.newCapture();
+        EasyMock.expect(transformer.transform(EasyMock.capture(configCapture))).andAnswer(configCapture::getValue);
         EasyMock.expect(worker.getPlugins()).andReturn(plugins).times(2);
         EasyMock.expect(plugins.compareAndSwapLoaders(connectorMock)).andReturn(delegatingLoader);
         // No new connector is created
         EasyMock.expect(Plugins.compareAndSwapLoaders(delegatingLoader)).andReturn(pluginLoader);
-        // Second should fail
-        createCallback.onCompletion(EasyMock.anyObject(), EasyMock.>isNull());
-        PowerMock.expectLastCall();
 
         PowerMock.replayAll();
 
         herder.putConnectorConfig(CONNECTOR_NAME, config, false, createCallback);
-        herder.putConnectorConfig(CONNECTOR_NAME, config, false, createCallback);
+        Herder.Created connectorInfo = createCallback.get(1000L, TimeUnit.SECONDS);
+        assertEquals(createdInfo(SourceSink.SOURCE), connectorInfo.result());
 
-        PowerMock.verifyAll();
+        // Second should fail
+        FutureCallback> failedCreateCallback = new FutureCallback<>();
+        herder.putConnectorConfig(CONNECTOR_NAME, config, false, failedCreateCallback);
+
+        try {
+            failedCreateCallback.get(1000L, TimeUnit.SECONDS);
+        } catch (ExecutionException e) {
+            assertNotNull(e.getCause());
+            throw e.getCause();
+        } finally {
+            PowerMock.verifyAll();
+        }
     }
 
     @Test
@@ -223,10 +235,13 @@ public void testCreateSinkConnector() throws Exception {
         expectAdd(SourceSink.SINK);
 
         Map config = connectorConfig(SourceSink.SINK);
-        expectConfigValidation(config);
+        Connector connectorMock = PowerMock.createMock(SinkConnector.class);
+        expectConfigValidation(connectorMock, true, config);
         PowerMock.replayAll();
 
         herder.putConnectorConfig(CONNECTOR_NAME, config, false, createCallback);
+        Herder.Created connectorInfo = createCallback.get(1000L, TimeUnit.SECONDS);
+        assertEquals(createdInfo(SourceSink.SINK), connectorInfo.result());
 
         PowerMock.verifyAll();
     }
@@ -237,25 +252,30 @@ public void testDestroyConnector() throws Exception {
         expectAdd(SourceSink.SOURCE);
 
         Map config = connectorConfig(SourceSink.SOURCE);
-        expectConfigValidation(config);
+        Connector connectorMock = PowerMock.createMock(SourceConnector.class);
+        expectConfigValidation(connectorMock, true, config);
 
         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();
 
         PowerMock.replayAll();
 
         herder.putConnectorConfig(CONNECTOR_NAME, config, false, createCallback);
-        FutureCallback> futureCb = new FutureCallback<>();
-        herder.deleteConnectorConfig(CONNECTOR_NAME, futureCb);
-        futureCb.get(1000L, TimeUnit.MILLISECONDS);
+        Herder.Created connectorInfo = createCallback.get(1000L, TimeUnit.SECONDS);
+        assertEquals(createdInfo(SourceSink.SOURCE), connectorInfo.result());
+
+        FutureCallback> deleteCallback = new FutureCallback<>();
+        herder.deleteConnectorConfig(CONNECTOR_NAME, deleteCallback);
+        deleteCallback.get(1000L, TimeUnit.MILLISECONDS);
 
         // Second deletion should fail since the connector is gone
-        futureCb = new FutureCallback<>();
-        herder.deleteConnectorConfig(CONNECTOR_NAME, futureCb);
+        FutureCallback> failedDeleteCallback = new FutureCallback<>();
+        herder.deleteConnectorConfig(CONNECTOR_NAME, failedDeleteCallback);
         try {
-            futureCb.get(1000L, TimeUnit.MILLISECONDS);
+            failedDeleteCallback.get(1000L, TimeUnit.MILLISECONDS);
             fail("Should have thrown NotFoundException");
         } catch (ExecutionException e) {
             assertTrue(e.getCause() instanceof NotFoundException);
@@ -269,22 +289,32 @@ public void testRestartConnector() throws Exception {
         expectAdd(SourceSink.SOURCE);
 
         Map config = connectorConfig(SourceSink.SOURCE);
-        expectConfigValidation(config);
+        Connector connectorMock = PowerMock.createMock(SourceConnector.class);
+        expectConfigValidation(connectorMock, true, config);
 
-        worker.stopConnector(CONNECTOR_NAME);
-        EasyMock.expectLastCall().andReturn(true);
+        worker.stopAndAwaitConnector(CONNECTOR_NAME);
+        EasyMock.expectLastCall();
 
-        worker.startConnector(EasyMock.eq(CONNECTOR_NAME), EasyMock.eq(config),
-                EasyMock.anyObject(HerderConnectorContext.class), EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
-        EasyMock.expectLastCall().andReturn(true);
+        Capture> onStart = EasyMock.newCapture();
+        worker.startConnector(EasyMock.eq(CONNECTOR_NAME), EasyMock.eq(config), EasyMock.anyObject(HerderConnectorContext.class),
+            EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED), EasyMock.capture(onStart));
+        EasyMock.expectLastCall().andAnswer(new IAnswer() {
+            @Override
+            public Boolean answer() throws Throwable {
+                onStart.getValue().onCompletion(null, TargetState.STARTED);
+                return true;
+            }
+        });
 
         PowerMock.replayAll();
 
         herder.putConnectorConfig(CONNECTOR_NAME, config, false, createCallback);
+        Herder.Created connectorInfo = createCallback.get(1000L, TimeUnit.SECONDS);
+        assertEquals(createdInfo(SourceSink.SOURCE), connectorInfo.result());
 
-        FutureCallback cb = new FutureCallback<>();
-        herder.restartConnector(CONNECTOR_NAME, cb);
-        cb.get(1000L, TimeUnit.MILLISECONDS);
+        FutureCallback restartCallback = new FutureCallback<>();
+        herder.restartConnector(CONNECTOR_NAME, restartCallback);
+        restartCallback.get(1000L, TimeUnit.MILLISECONDS);
 
         PowerMock.verifyAll();
     }
@@ -294,26 +324,37 @@ public void testRestartConnectorFailureOnStart() throws Exception {
         expectAdd(SourceSink.SOURCE);
 
         Map config = connectorConfig(SourceSink.SOURCE);
-        expectConfigValidation(config);
+        Connector connectorMock = PowerMock.createMock(SourceConnector.class);
+        expectConfigValidation(connectorMock, true, config);
 
-        worker.stopConnector(CONNECTOR_NAME);
-        EasyMock.expectLastCall().andReturn(true);
+        worker.stopAndAwaitConnector(CONNECTOR_NAME);
+        EasyMock.expectLastCall();
 
-        worker.startConnector(EasyMock.eq(CONNECTOR_NAME), EasyMock.eq(config),
-                EasyMock.anyObject(HerderConnectorContext.class), EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
-        EasyMock.expectLastCall().andReturn(false);
+        Capture> onStart = EasyMock.newCapture();
+        worker.startConnector(EasyMock.eq(CONNECTOR_NAME), EasyMock.eq(config), EasyMock.anyObject(HerderConnectorContext.class),
+            EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED), EasyMock.capture(onStart));
+        Exception exception = new ConnectException("Failed to start connector");
+        EasyMock.expectLastCall().andAnswer(new IAnswer() {
+            @Override
+            public Boolean answer() throws Throwable {
+                onStart.getValue().onCompletion(exception, null);
+                return true;
+            }
+        });
 
         PowerMock.replayAll();
 
         herder.putConnectorConfig(CONNECTOR_NAME, config, false, createCallback);
+        Herder.Created connectorInfo = createCallback.get(1000L, TimeUnit.SECONDS);
+        assertEquals(createdInfo(SourceSink.SOURCE), connectorInfo.result());
 
-        FutureCallback cb = new FutureCallback<>();
-        herder.restartConnector(CONNECTOR_NAME, cb);
+        FutureCallback restartCallback = new FutureCallback<>();
+        herder.restartConnector(CONNECTOR_NAME, restartCallback);
         try {
-            cb.get(1000L, TimeUnit.MILLISECONDS);
+            restartCallback.get(1000L, TimeUnit.MILLISECONDS);
             fail();
-        } catch (ExecutionException exception) {
-            assertEquals(ConnectException.class, exception.getCause().getClass());
+        } catch (ExecutionException e) {
+            assertEquals(exception, e.getCause());
         }
 
         PowerMock.verifyAll();
@@ -325,21 +366,34 @@ public void testRestartTask() throws Exception {
         expectAdd(SourceSink.SOURCE);
 
         Map connectorConfig = connectorConfig(SourceSink.SOURCE);
-        expectConfigValidation(connectorConfig);
+        Connector connectorMock = PowerMock.createMock(SourceConnector.class);
+        expectConfigValidation(connectorMock, true, connectorConfig);
 
         worker.stopAndAwaitTask(taskId);
         EasyMock.expectLastCall();
 
-        worker.startTask(taskId, connectorConfig, taskConfig(SourceSink.SOURCE), herder, TargetState.STARTED);
+        ClusterConfigState configState = new ClusterConfigState(
+                -1,
+                null,
+                Collections.singletonMap(CONNECTOR_NAME, 1),
+                Collections.singletonMap(CONNECTOR_NAME, connectorConfig),
+                Collections.singletonMap(CONNECTOR_NAME, TargetState.STARTED),
+                Collections.singletonMap(taskId, taskConfig(SourceSink.SOURCE)),
+                new HashSet<>(),
+                transformer);
+        worker.startTask(taskId, configState, connectorConfig, taskConfig(SourceSink.SOURCE), herder, TargetState.STARTED);
         EasyMock.expectLastCall().andReturn(true);
 
         PowerMock.replayAll();
 
         herder.putConnectorConfig(CONNECTOR_NAME, connectorConfig, false, createCallback);
+        Herder.Created connectorInfo = createCallback.get(1000L, TimeUnit.SECONDS);
+        createCallback.get(1000L, TimeUnit.SECONDS);
+        assertEquals(createdInfo(SourceSink.SOURCE), connectorInfo.result());
 
-        FutureCallback cb = new FutureCallback<>();
-        herder.restartTask(taskId, cb);
-        cb.get(1000L, TimeUnit.MILLISECONDS);
+        FutureCallback restartTaskCallback = new FutureCallback<>();
+        herder.restartTask(taskId, restartTaskCallback);
+        restartTaskCallback.get(1000L, TimeUnit.MILLISECONDS);
 
         PowerMock.verifyAll();
     }
@@ -350,17 +404,29 @@ public void testRestartTaskFailureOnStart() throws Exception {
         expectAdd(SourceSink.SOURCE);
 
         Map connectorConfig = connectorConfig(SourceSink.SOURCE);
-        expectConfigValidation(connectorConfig);
+        Connector connectorMock = PowerMock.createMock(SourceConnector.class);
+        expectConfigValidation(connectorMock, true, connectorConfig);
 
         worker.stopAndAwaitTask(taskId);
         EasyMock.expectLastCall();
 
-        worker.startTask(taskId, connectorConfig, taskConfig(SourceSink.SOURCE), herder, TargetState.STARTED);
+        ClusterConfigState configState = new ClusterConfigState(
+                -1,
+                null,
+                Collections.singletonMap(CONNECTOR_NAME, 1),
+                Collections.singletonMap(CONNECTOR_NAME, connectorConfig),
+                Collections.singletonMap(CONNECTOR_NAME, TargetState.STARTED),
+                Collections.singletonMap(new ConnectorTaskId(CONNECTOR_NAME, 0), taskConfig(SourceSink.SOURCE)),
+                new HashSet<>(),
+                transformer);
+        worker.startTask(taskId, configState, connectorConfig, taskConfig(SourceSink.SOURCE), herder, TargetState.STARTED);
         EasyMock.expectLastCall().andReturn(false);
 
         PowerMock.replayAll();
 
         herder.putConnectorConfig(CONNECTOR_NAME, connectorConfig, false, createCallback);
+        Herder.Created connectorInfo = createCallback.get(1000L, TimeUnit.MILLISECONDS);
+        assertEquals(createdInfo(SourceSink.SOURCE), connectorInfo.result());
 
         FutureCallback cb = new FutureCallback<>();
         herder.restartTask(taskId, cb);
@@ -380,11 +446,14 @@ public void testCreateAndStop() throws Exception {
         expectAdd(SourceSink.SOURCE);
 
         Map connectorConfig = connectorConfig(SourceSink.SOURCE);
-        expectConfigValidation(connectorConfig);
+        Connector connectorMock = PowerMock.createMock(SourceConnector.class);
+        expectConfigValidation(connectorMock, true, connectorConfig);
 
         // 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();
@@ -393,6 +462,9 @@ public void testCreateAndStop() throws Exception {
         PowerMock.replayAll();
 
         herder.putConnectorConfig(CONNECTOR_NAME, connectorConfig, false, createCallback);
+        Herder.Created connectorInfo = createCallback.get(1000L, TimeUnit.SECONDS);
+        assertEquals(createdInfo(SourceSink.SOURCE), connectorInfo.result());
+
         herder.stop();
 
         PowerMock.verifyAll();
@@ -401,6 +473,7 @@ public void testCreateAndStop() throws Exception {
     @Test
     public void testAccessors() throws Exception {
         Map connConfig = connectorConfig(SourceSink.SOURCE);
+        System.out.println(connConfig);
 
         Callback> listConnectorsCb = PowerMock.createMock(Callback.class);
         Callback connectorInfoCb = PowerMock.createMock(Callback.class);
@@ -420,7 +493,7 @@ public void testAccessors() throws Exception {
         // Create connector
         connector = PowerMock.createMock(BogusSourceConnector.class);
         expectAdd(SourceSink.SOURCE);
-        expectConfigValidation(connConfig);
+        expectConfigValidation(connector, true, connConfig);
 
         // Validate accessors with 1 connector
         listConnectorsCb.onCompletion(null, singleton(CONNECTOR_NAME));
@@ -446,6 +519,15 @@ public void testAccessors() throws Exception {
         herder.taskConfigs(CONNECTOR_NAME, taskConfigsCb);
 
         herder.putConnectorConfig(CONNECTOR_NAME, connConfig, false, createCallback);
+        Herder.Created connectorInfo = createCallback.get(1000L, TimeUnit.SECONDS);
+        assertEquals(createdInfo(SourceSink.SOURCE), connectorInfo.result());
+
+        EasyMock.reset(transformer);
+        EasyMock.expect(transformer.transform(EasyMock.eq(CONNECTOR_NAME), EasyMock.anyObject()))
+            .andThrow(new AssertionError("Config transformation should not occur when requesting connector or task info"))
+            .anyTimes();
+        EasyMock.replay(transformer);
+
         herder.connectors(listConnectorsCb);
         herder.connectorInfo(CONNECTOR_NAME, connectorInfoCb);
         herder.connectorConfig(CONNECTOR_NAME, connectorConfigCb);
@@ -461,35 +543,39 @@ public void testPutConnectorConfig() throws Exception {
         newConnConfig.put("foo", "bar");
 
         Callback> connectorConfigCb = PowerMock.createMock(Callback.class);
-        Callback> putConnectorConfigCb = PowerMock.createMock(Callback.class);
+        // Callback> putConnectorConfigCb = PowerMock.createMock(Callback.class);
 
         // Create
         connector = PowerMock.createMock(BogusSourceConnector.class);
         expectAdd(SourceSink.SOURCE);
-        Connector connectorMock = PowerMock.createMock(Connector.class);
+        Connector connectorMock = PowerMock.createMock(SourceConnector.class);
         expectConfigValidation(connectorMock, true, connConfig);
 
         // Should get first config
         connectorConfigCb.onCompletion(null, connConfig);
         EasyMock.expectLastCall();
         // Update config, which requires stopping and restarting
-        worker.stopConnector(CONNECTOR_NAME);
-        EasyMock.expectLastCall().andReturn(true);
+        worker.stopAndAwaitConnector(CONNECTOR_NAME);
+        EasyMock.expectLastCall();
         Capture> capturedConfig = EasyMock.newCapture();
-        worker.startConnector(EasyMock.eq(CONNECTOR_NAME), EasyMock.capture(capturedConfig), EasyMock.anyObject(),
-                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
-        EasyMock.expectLastCall().andReturn(true);
+        Capture> onStart = EasyMock.newCapture();
+        worker.startConnector(EasyMock.eq(CONNECTOR_NAME), EasyMock.capture(capturedConfig), EasyMock.anyObject(),
+                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED), EasyMock.capture(onStart));
+        EasyMock.expectLastCall().andAnswer(new IAnswer() {
+            @Override
+            public Boolean answer() throws Throwable {
+                onStart.getValue().onCompletion(null, TargetState.STARTED);
+                return true;
+            }
+        });
         EasyMock.expect(worker.isRunning(CONNECTOR_NAME)).andReturn(true);
+        EasyMock.expect(worker.isTopicCreationEnabled()).andReturn(true);
         // Generate same task config, which should result in no additional action to restart tasks
-        EasyMock.expect(worker.connectorTaskConfigs(CONNECTOR_NAME, new SourceConnectorConfig(plugins, newConnConfig)))
+        EasyMock.expect(worker.connectorTaskConfigs(CONNECTOR_NAME, new SourceConnectorConfig(plugins, newConnConfig, true)))
                 .andReturn(singletonList(taskConfig(SourceSink.SOURCE)));
         worker.isSinkConnector(CONNECTOR_NAME);
         EasyMock.expectLastCall().andReturn(false);
-        ConnectorInfo newConnInfo = new ConnectorInfo(CONNECTOR_NAME, newConnConfig, Arrays.asList(new ConnectorTaskId(CONNECTOR_NAME, 0)),
-            ConnectorType.SOURCE);
-        putConnectorConfigCb.onCompletion(null, new Herder.Created<>(false, newConnInfo));
-        EasyMock.expectLastCall();
-        // Should get new config
+
         expectConfigValidation(connectorMock, false, newConnConfig);
         connectorConfigCb.onCompletion(null, newConnConfig);
         EasyMock.expectLastCall();
@@ -498,13 +584,22 @@ public void testPutConnectorConfig() throws Exception {
         PowerMock.replayAll();
 
         herder.putConnectorConfig(CONNECTOR_NAME, connConfig, false, createCallback);
+        Herder.Created connectorInfo = createCallback.get(1000L, TimeUnit.SECONDS);
+        assertEquals(createdInfo(SourceSink.SOURCE), connectorInfo.result());
+
         herder.connectorConfig(CONNECTOR_NAME, connectorConfigCb);
-        herder.putConnectorConfig(CONNECTOR_NAME, newConnConfig, true, putConnectorConfigCb);
+
+        FutureCallback> reconfigureCallback = new FutureCallback<>();
+        herder.putConnectorConfig(CONNECTOR_NAME, newConnConfig, true, reconfigureCallback);
+        Herder.Created newConnectorInfo = reconfigureCallback.get(1000L, TimeUnit.SECONDS);
+        ConnectorInfo newConnInfo = new ConnectorInfo(CONNECTOR_NAME, newConnConfig, Arrays.asList(new ConnectorTaskId(CONNECTOR_NAME, 0)),
+            ConnectorType.SOURCE);
+        assertEquals(newConnInfo, newConnectorInfo.result());
+
         assertEquals("bar", capturedConfig.getValue().get("foo"));
         herder.connectorConfig(CONNECTOR_NAME, connectorConfigCb);
 
         PowerMock.verifyAll();
-
     }
 
     @Test(expected = UnsupportedOperationException.class)
@@ -515,17 +610,19 @@ public void testPutTaskConfigs() {
 
         herder.putTaskConfigs(CONNECTOR_NAME,
                 Arrays.asList(singletonMap("config", "value")),
-                cb);
+                cb,
+                null);
 
         PowerMock.verifyAll();
     }
 
     @Test
-    public void testCorruptConfig() {
+    public void testCorruptConfig() throws Throwable {
         Map config = new HashMap<>();
         config.put(ConnectorConfig.NAME_CONFIG, CONNECTOR_NAME);
         config.put(ConnectorConfig.CONNECTOR_CLASS_CONFIG, BogusSinkConnector.class.getName());
-        Connector connectorMock = PowerMock.createMock(Connector.class);
+        config.put(SinkConnectorConfig.TOPICS_CONFIG, TOPICS_LIST_STR);
+        Connector connectorMock = PowerMock.createMock(SinkConnector.class);
         String error = "This is an error in your config!";
         List errors = new ArrayList<>(singletonList(error));
         String key = "foo.invalid.key";
@@ -536,47 +633,58 @@ public void testCorruptConfig() {
         );
         ConfigDef configDef = new ConfigDef();
         configDef.define(key, ConfigDef.Type.STRING, ConfigDef.Importance.HIGH, "");
+        EasyMock.expect(worker.configTransformer()).andReturn(transformer).times(2);
+        final Capture> configCapture = EasyMock.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);
         EasyMock.expect(worker.getPlugins()).andStubReturn(plugins);
         EasyMock.expect(plugins.newConnector(EasyMock.anyString())).andReturn(connectorMock);
         EasyMock.expect(connectorMock.config()).andStubReturn(configDef);
         EasyMock.expect(Plugins.compareAndSwapLoaders(delegatingLoader)).andReturn(pluginLoader);
-        Callback> callback = PowerMock.createMock(Callback.class);
-        Capture capture = Capture.newInstance();
-        callback.onCompletion(
-            EasyMock.capture(capture), EasyMock.isNull(Herder.Created.class)
-        );
 
         PowerMock.replayAll();
 
-        herder.putConnectorConfig(CONNECTOR_NAME, config, true, callback);
-        assertEquals(
-            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`"
-        );
+        herder.putConnectorConfig(CONNECTOR_NAME, config, true, createCallback);
+        try {
+            createCallback.get(1000L, TimeUnit.SECONDS);
+            fail("Should have failed to configure connector");
+        } catch (ExecutionException e) {
+            assertNotNull(e.getCause());
+            Throwable cause = e.getCause();
+            assertTrue(cause instanceof BadRequestException);
+            assertEquals(
+                cause.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 `/connector-plugins/{connectorType}/config/validate`"
+            );
+        }
 
         PowerMock.verifyAll();
     }
 
-    private void expectAdd(SourceSink sourceSink) throws Exception {
-
+    private void expectAdd(SourceSink sourceSink) {
         Map connectorProps = connectorConfig(sourceSink);
         ConnectorConfig connConfig = sourceSink == SourceSink.SOURCE ?
-            new SourceConnectorConfig(plugins, connectorProps) :
+            new SourceConnectorConfig(plugins, connectorProps, true) :
             new SinkConnectorConfig(plugins, connectorProps);
 
+        Capture> onStart = EasyMock.newCapture();
         worker.startConnector(EasyMock.eq(CONNECTOR_NAME), EasyMock.eq(connectorProps), EasyMock.anyObject(HerderConnectorContext.class),
-                              EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED));
-        EasyMock.expectLastCall().andReturn(true);
+                EasyMock.eq(herder), EasyMock.eq(TargetState.STARTED), EasyMock.capture(onStart));
+        // EasyMock.expectLastCall().andReturn(true);
+        EasyMock.expectLastCall().andAnswer(new IAnswer() {
+            @Override
+            public Boolean answer() throws Throwable {
+                onStart.getValue().onCompletion(null, TargetState.STARTED);
+                return true;
+            }
+        });
         EasyMock.expect(worker.isRunning(CONNECTOR_NAME)).andReturn(true);
-
-        ConnectorInfo connInfo = new ConnectorInfo(CONNECTOR_NAME, connectorProps, Arrays.asList(new ConnectorTaskId(CONNECTOR_NAME, 0)),
-            SourceSink.SOURCE == sourceSink ? ConnectorType.SOURCE : ConnectorType.SINK);
-        createCallback.onCompletion(null, new Herder.Created<>(true, connInfo));
-        EasyMock.expectLastCall();
+        if (sourceSink == SourceSink.SOURCE) {
+            EasyMock.expect(worker.isTopicCreationEnabled()).andReturn(true);
+        }
 
         // And we should instantiate the tasks. For a sink task, we should see added properties for the input topic partitions
 
@@ -585,23 +693,38 @@ private void expectAdd(SourceSink sourceSink) throws Exception {
         EasyMock.expect(worker.connectorTaskConfigs(CONNECTOR_NAME, connConfig))
             .andReturn(singletonList(generatedTaskProps));
 
-        worker.startTask(new ConnectorTaskId(CONNECTOR_NAME, 0), connectorConfig(sourceSink), generatedTaskProps, herder, TargetState.STARTED);
+        ClusterConfigState configState = new ClusterConfigState(
+                -1,
+                null,
+                Collections.singletonMap(CONNECTOR_NAME, 1),
+                Collections.singletonMap(CONNECTOR_NAME, connectorConfig(sourceSink)),
+                Collections.singletonMap(CONNECTOR_NAME, TargetState.STARTED),
+                Collections.singletonMap(new ConnectorTaskId(CONNECTOR_NAME, 0), generatedTaskProps),
+                new HashSet<>(),
+                transformer);
+        worker.startTask(new ConnectorTaskId(CONNECTOR_NAME, 0), configState, connectorConfig(sourceSink), generatedTaskProps, herder, TargetState.STARTED);
         EasyMock.expectLastCall().andReturn(true);
 
         EasyMock.expect(herder.connectorTypeForClass(BogusSourceConnector.class.getName()))
             .andReturn(ConnectorType.SOURCE).anyTimes();
         EasyMock.expect(herder.connectorTypeForClass(BogusSinkConnector.class.getName()))
-        .andReturn(ConnectorType.SINK).anyTimes();
+            .andReturn(ConnectorType.SINK).anyTimes();
         worker.isSinkConnector(CONNECTOR_NAME);
         PowerMock.expectLastCall().andReturn(sourceSink == SourceSink.SINK);
     }
 
+    private ConnectorInfo createdInfo(SourceSink sourceSink) {
+        return new ConnectorInfo(CONNECTOR_NAME, connectorConfig(sourceSink),
+            Arrays.asList(new ConnectorTaskId(CONNECTOR_NAME, 0)),
+            SourceSink.SOURCE == sourceSink ? ConnectorType.SOURCE : ConnectorType.SINK);
+    }
+
     private void expectStop() {
         ConnectorTaskId task = new ConnectorTaskId(CONNECTOR_NAME, 0);
         worker.stopAndAwaitTasks(singletonList(task));
         EasyMock.expectLastCall();
-        worker.stopConnector(CONNECTOR_NAME);
-        EasyMock.expectLastCall().andReturn(true);
+        worker.stopAndAwaitConnector(CONNECTOR_NAME);
+        EasyMock.expectLastCall();
     }
 
     private void expectDestroy() {
@@ -614,8 +737,12 @@ private static Map connectorConfig(SourceSink sourceSink) {
         Class connectorClass = sourceSink == SourceSink.SINK ? BogusSinkConnector.class : BogusSourceConnector.class;
         props.put(ConnectorConfig.CONNECTOR_CLASS_CONFIG, connectorClass.getName());
         props.put(ConnectorConfig.TASKS_MAX_CONFIG, "1");
-        if (sourceSink == SourceSink.SINK)
+        if (sourceSink == SourceSink.SINK) {
             props.put(SinkTask.TOPICS_CONFIG, TOPICS_LIST_STR);
+        } else {
+            props.put(DEFAULT_TOPIC_CREATION_PREFIX + REPLICATION_FACTOR_CONFIG, String.valueOf(1));
+            props.put(DEFAULT_TOPIC_CREATION_PREFIX + PARTITIONS_CONFIG, String.valueOf(1));
+        }
         return props;
     }
 
@@ -630,18 +757,15 @@ private static Map taskConfig(SourceSink sourceSink) {
         return generatedTaskProps;
     }
 
-
-    private void expectConfigValidation(Map ... configs) {
-        Connector connectorMock = PowerMock.createMock(Connector.class);
-        expectConfigValidation(connectorMock, true, configs);
-    }
-
     private void expectConfigValidation(
             Connector connectorMock,
             boolean shouldCreateConnector,
             Map... configs
     ) {
         // config validation
+        EasyMock.expect(worker.configTransformer()).andReturn(transformer).times(2);
+        final Capture> configCapture = EasyMock.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);
         if (shouldCreateConnector) {
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 c6b61b4d2045c..460fb0d6c38a8 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
@@ -22,7 +22,9 @@
 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.modules.junit4.PowerMockRunner;
 
 import java.io.File;
 import java.io.IOException;
@@ -30,9 +32,12 @@
 import java.util.Arrays;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.concurrent.ThreadPoolExecutor;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
 
+@RunWith(PowerMockRunner.class)
 public class FileOffsetBackingStoreTest {
 
     FileOffsetBackingStore store;
@@ -41,12 +46,15 @@ public class FileOffsetBackingStoreTest {
     File tempFile;
 
     private static Map firstSet = new HashMap<>();
+    private static final Runnable EMPTY_RUNNABLE = () -> {
+    };
 
     static {
         firstSet.put(buffer("key"), buffer("value"));
         firstSet.put(null, null);
     }
 
+    @SuppressWarnings("deprecation")
     @Before
     public void setup() throws IOException {
         store = new FileOffsetBackingStore();
@@ -70,12 +78,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")));
 
@@ -85,7 +92,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();
@@ -95,12 +101,18 @@ 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();
     }
 
+    @Test
+    public void testThreadName() {
+        assertTrue(((ThreadPoolExecutor) store.executor).getThreadFactory()
+                .newThread(EMPTY_RUNNABLE).getName().startsWith(FileOffsetBackingStore.class.getSimpleName()));
+    }
+
     private static ByteBuffer buffer(String v) {
         return ByteBuffer.wrap(v.getBytes());
     }
@@ -112,12 +124,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 aac1b78c9184d..a46f6fefcfe73 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;
@@ -30,6 +31,7 @@
 import org.apache.kafka.connect.runtime.distributed.ClusterConfigState;
 import org.apache.kafka.connect.runtime.distributed.DistributedConfig;
 import org.apache.kafka.connect.util.Callback;
+import org.apache.kafka.connect.util.ConnectUtils;
 import org.apache.kafka.connect.util.ConnectorTaskId;
 import org.apache.kafka.connect.util.KafkaBasedLog;
 import org.apache.kafka.connect.util.TestFuture;
@@ -57,13 +59,15 @@
 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.*")
-@SuppressWarnings("unchecked")
+@PrepareForTest({KafkaConfigBackingStore.class, ConnectUtils.class})
+@PowerMockIgnore({"javax.management.*", "javax.crypto.*"})
+@SuppressWarnings({"unchecked", "deprecation"})
 public class KafkaConfigBackingStoreTest {
     private static final String TOPIC = "connect-configs";
     private static final short TOPIC_REPLICATION_FACTOR = 5;
@@ -146,7 +150,11 @@ public class KafkaConfigBackingStoreTest {
 
     @Before
     public void setUp() {
-        configStorage = PowerMock.createPartialMock(KafkaConfigBackingStore.class, new String[]{"createKafkaBasedLog"}, converter, DEFAULT_DISTRIBUTED_CONFIG);
+        PowerMock.mockStaticPartial(ConnectUtils.class, "lookupKafkaClusterId");
+        EasyMock.expect(ConnectUtils.lookupKafkaClusterId(EasyMock.anyObject())).andReturn("test-cluster").anyTimes();
+        PowerMock.replay(ConnectUtils.class);
+
+        configStorage = PowerMock.createPartialMock(KafkaConfigBackingStore.class, new String[]{"createKafkaBasedLog"}, converter, DEFAULT_DISTRIBUTED_CONFIG, null);
         Whitebox.setInternalState(configStorage, "configLog", storeLog);
         configStorage.setUpdateListener(configUpdateListener);
     }
@@ -154,12 +162,15 @@ public void setUp() {
     @Test
     public void testStartStop() throws Exception {
         expectConfigure();
-        expectStart(Collections.EMPTY_LIST, Collections.EMPTY_MAP);
+        expectStart(Collections.emptyList(), Collections.emptyMap());
+        expectPartitionCount(1);
         expectStop();
-
         PowerMock.replayAll();
 
-        configStorage.setupAndCreateKafkaBasedLog(TOPIC, DEFAULT_DISTRIBUTED_CONFIG);
+        Map settings = new HashMap<>(DEFAULT_CONFIG_STORAGE_PROPS);
+        settings.put("config.storage.min.insync.replicas", "3");
+        settings.put("config.storage.max.message.bytes", "1001");
+        configStorage.setupAndCreateKafkaBasedLog(TOPIC, new DistributedConfig(settings));
 
         assertEquals(TOPIC, capturedTopic.getValue());
         assertEquals("org.apache.kafka.common.serialization.StringSerializer", capturedProducerProps.getValue().get(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG));
@@ -170,6 +181,8 @@ public void testStartStop() throws Exception {
         assertEquals(TOPIC, capturedNewTopic.getValue().name());
         assertEquals(1, capturedNewTopic.getValue().numPartitions());
         assertEquals(TOPIC_REPLICATION_FACTOR, capturedNewTopic.getValue().replicationFactor());
+        assertEquals("3", capturedNewTopic.getValue().configs().get("min.insync.replicas"));
+        assertEquals("1001", capturedNewTopic.getValue().configs().get("max.message.bytes"));
         configStorage.start();
         configStorage.stop();
 
@@ -179,7 +192,7 @@ public void testStartStop() throws Exception {
     @Test
     public void testPutConnectorConfig() throws Exception {
         expectConfigure();
-        expectStart(Collections.EMPTY_LIST, Collections.EMPTY_MAP);
+        expectStart(Collections.emptyList(), Collections.emptyMap());
 
         expectConvertWriteAndRead(
                 CONNECTOR_CONFIG_KEYS.get(0), KafkaConfigBackingStore.CONNECTOR_CONFIGURATION_V0, CONFIGS_SERIALIZED.get(0),
@@ -198,6 +211,7 @@ public void testPutConnectorConfig() throws Exception {
         configUpdateListener.onConnectorConfigRemove(CONNECTOR_IDS.get(1));
         EasyMock.expectLastCall();
 
+        expectPartitionCount(1);
         expectStop();
 
         PowerMock.replayAll();
@@ -241,10 +255,10 @@ public void testPutConnectorConfig() throws Exception {
     @Test
     public void testPutTaskConfigs() throws Exception {
         expectConfigure();
-        expectStart(Collections.EMPTY_LIST, Collections.EMPTY_MAP);
+        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());
+        expectReadToEnd(new LinkedHashMap<>());
         expectConvertWriteRead(
                 TASK_CONFIG_KEYS.get(0), KafkaConfigBackingStore.TASK_CONFIGURATION_V0, CONFIGS_SERIALIZED.get(0),
                 "properties", SAMPLE_CONFIGS.get(0));
@@ -266,6 +280,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();
@@ -274,7 +289,7 @@ public void testPutTaskConfigs() throws Exception {
         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.EMPTY_LIST);
+        whiteboxAddConnector(CONNECTOR_IDS.get(0), SAMPLE_CONFIGS.get(0), Collections.emptyList());
 
         // Null before writing
         ClusterConfigState configState = configStorage.snapshot();
@@ -302,13 +317,104 @@ 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();
-        expectStart(Collections.EMPTY_LIST, Collections.EMPTY_MAP);
+        expectStart(Collections.emptyList(), Collections.emptyMap());
 
         // Task configs should read to end, write to the log, read to end, write root.
-        expectReadToEnd(new LinkedHashMap());
+        expectReadToEnd(new LinkedHashMap<>());
         expectConvertWriteRead(
             COMMIT_TASKS_CONFIG_KEYS.get(0), KafkaConfigBackingStore.CONNECTOR_TASKS_COMMIT_V0, CONFIGS_SERIALIZED.get(0),
             "tasks", 0); // We have 0 tasks
@@ -321,6 +427,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();
@@ -329,7 +436,7 @@ public void testPutTaskConfigsZeroTasks() throws Exception {
         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.EMPTY_LIST);
+        whiteboxAddConnector(CONNECTOR_IDS.get(0), SAMPLE_CONFIGS.get(0), Collections.emptyList());
 
         // Null before writing
         ClusterConfigState configState = configStorage.snapshot();
@@ -374,6 +481,7 @@ public void testRestoreTargetState() throws Exception {
 
         // Shouldn't see any callbacks since this is during startup
 
+        expectPartitionCount(1);
         expectStop();
 
         PowerMock.replayAll();
@@ -416,6 +524,7 @@ public void testBackgroundUpdateTargetState() throws Exception {
         configUpdateListener.onConnectorTargetStateChange(CONNECTOR_IDS.get(0));
         EasyMock.expectLastCall();
 
+        expectPartitionCount(1);
         expectStop();
 
         PowerMock.replayAll();
@@ -447,7 +556,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 +575,7 @@ public void testBackgroundConnectorDeletion() throws Exception {
         configUpdateListener.onConnectorConfigRemove(CONNECTOR_IDS.get(0));
         EasyMock.expectLastCall();
 
+        expectPartitionCount(1);
         expectStop();
 
         PowerMock.replayAll();
@@ -476,8 +586,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 +621,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 +669,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 +724,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 +772,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,13 +820,14 @@ 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());
+        expectReadToEnd(new LinkedHashMap<>());
         expectConvertWriteRead(
                 TASK_CONFIG_KEYS.get(0), KafkaConfigBackingStore.TASK_CONFIGURATION_V0, CONFIGS_SERIALIZED.get(0),
                 "properties", SAMPLE_CONFIGS.get(0));
-        expectReadToEnd(new LinkedHashMap());
+        expectReadToEnd(new LinkedHashMap<>());
         expectConvertWriteRead(
                 COMMIT_TASKS_CONFIG_KEYS.get(0), KafkaConfigBackingStore.CONNECTOR_TASKS_COMMIT_V0, CONFIGS_SERIALIZED.get(2),
                 "tasks", 1); // Updated to just 1 task
@@ -727,7 +851,7 @@ public void testPutTaskConfigsDoesNotResolveAllInconsistencies() throws Exceptio
         assertEquals(6, configState.offset()); // Should always be next to be read, not last committed
         assertEquals(Arrays.asList(CONNECTOR_IDS.get(0)), new ArrayList<>(configState.connectors()));
         // Inconsistent data should leave us with no tasks listed for the connector and an entry in the inconsistent list
-        assertEquals(Collections.EMPTY_LIST, configState.tasks(CONNECTOR_IDS.get(0)));
+        assertEquals(Collections.emptyList(), configState.tasks(CONNECTOR_IDS.get(0)));
         // Both TASK_CONFIG_STRUCTS[0] -> SAMPLE_CONFIGS[0]
         assertNull(configState.taskConfig(TASK_IDS.get(0)));
         assertNull(configState.taskConfig(TASK_IDS.get(1)));
@@ -751,6 +875,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,13 +899,18 @@ 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) throws Exception {
+                             final Map deserializations) {
         storeLog.start();
         PowerMock.expectLastCall().andAnswer(new IAnswer() {
             @Override
-            public Object answer() throws Throwable {
+            public Object answer() {
                 for (ConsumerRecord rec : preexistingRecords)
                     capturedConsumedCallback.getValue().onCompletion(null, rec);
                 return null;
@@ -827,7 +972,7 @@ private void expectReadToEnd(final LinkedHashMap serializedConfi
         EasyMock.expect(storeLog.readToEnd())
                 .andAnswer(new IAnswer>() {
                     @Override
-                    public Future answer() throws Throwable {
+                    public Future answer() {
                         TestFuture future = new TestFuture();
                         for (Map.Entry entry : serializedConfigs.entrySet())
                             capturedConsumedCallback.getValue().onCompletion(null, new ConsumerRecord<>(TOPIC, 0, logOffset++, 0L, TimestampType.CREATE_TIME, 0L, 0, 0, entry.getKey(), entry.getValue()));
@@ -877,5 +1022,4 @@ private Map structToMap(Struct struct) {
             result.put(field.name(), struct.get(field));
         return result;
     }
-
 }
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 8cd2f0b0560aa..0d0194fc9415e 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
@@ -25,6 +25,7 @@
 import org.apache.kafka.common.record.TimestampType;
 import org.apache.kafka.connect.runtime.distributed.DistributedConfig;
 import org.apache.kafka.connect.util.Callback;
+import org.apache.kafka.connect.util.ConnectUtils;
 import org.apache.kafka.connect.util.KafkaBasedLog;
 import org.easymock.Capture;
 import org.easymock.EasyMock;
@@ -58,9 +59,9 @@
 import static org.junit.Assert.fail;
 
 @RunWith(PowerMockRunner.class)
-@PrepareForTest(KafkaOffsetBackingStore.class)
-@PowerMockIgnore("javax.management.*")
-@SuppressWarnings("unchecked")
+@PrepareForTest({KafkaOffsetBackingStore.class, ConnectUtils.class})
+@PowerMockIgnore({"javax.management.*", "javax.crypto.*"})
+@SuppressWarnings({"unchecked", "deprecation"})
 public class KafkaOffsetBackingStoreTest {
     private static final String TOPIC = "connect-offsets";
     private static final short TOPIC_PARTITIONS = 2;
@@ -111,18 +112,22 @@ public class KafkaOffsetBackingStoreTest {
 
     @Before
     public void setUp() throws Exception {
-        store = PowerMock.createPartialMockAndInvokeDefaultConstructor(KafkaOffsetBackingStore.class, new String[]{"createKafkaBasedLog"});
+        store = PowerMock.createPartialMockAndInvokeDefaultConstructor(KafkaOffsetBackingStore.class, "createKafkaBasedLog");
     }
 
     @Test
     public void testStartStop() throws Exception {
         expectConfigure();
-        expectStart(Collections.EMPTY_LIST);
+        expectStart(Collections.emptyList());
         expectStop();
+        expectClusterId();
 
         PowerMock.replayAll();
 
-        store.configure(DEFAULT_DISTRIBUTED_CONFIG);
+        Map settings = new HashMap<>(DEFAULT_PROPS);
+        settings.put("offset.storage.min.insync.replicas", "3");
+        settings.put("offset.storage.max.message.bytes", "1001");
+        store.configure(new DistributedConfig(settings));
         assertEquals(TOPIC, capturedTopic.getValue());
         assertEquals("org.apache.kafka.common.serialization.ByteArraySerializer", capturedProducerProps.getValue().get(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG));
         assertEquals("org.apache.kafka.common.serialization.ByteArraySerializer", capturedProducerProps.getValue().get(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG));
@@ -149,6 +154,7 @@ public void testReloadOnStart() throws Exception {
                 new ConsumerRecord<>(TOPIC, 1, 1, 0L, TimestampType.CREATE_TIME, 0L, 0, 0, TP1_KEY.array(), TP1_VALUE_NEW.array())
         ));
         expectStop();
+        expectClusterId();
 
         PowerMock.replayAll();
 
@@ -166,18 +172,15 @@ public void testReloadOnStart() throws Exception {
     @Test
     public void testGetSet() throws Exception {
         expectConfigure();
-        expectStart(Collections.EMPTY_LIST);
+        expectStart(Collections.emptyList());
         expectStop();
 
         // First get() against an empty store
         final Capture> firstGetReadToEndCallback = EasyMock.newCapture();
         storeLog.readToEnd(EasyMock.capture(firstGetReadToEndCallback));
-        PowerMock.expectLastCall().andAnswer(new IAnswer() {
-            @Override
-            public Object answer() throws Throwable {
-                firstGetReadToEndCallback.getValue().onCompletion(null, null);
-                return null;
-            }
+        PowerMock.expectLastCall().andAnswer(() -> {
+            firstGetReadToEndCallback.getValue().onCompletion(null, null);
+            return null;
         });
 
         // Set offsets
@@ -214,23 +217,17 @@ public Object answer() throws Throwable {
             }
         });
 
+        expectClusterId();
         PowerMock.replayAll();
 
         store.configure(DEFAULT_DISTRIBUTED_CONFIG);
         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<>();
@@ -253,28 +250,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();
 
@@ -284,7 +267,7 @@ public void onCompletion(Throwable error, Map result) {
     @Test
     public void testGetSetNull() throws Exception {
         expectConfigure();
-        expectStart(Collections.EMPTY_LIST);
+        expectStart(Collections.emptyList());
 
         // Set offsets
         Capture callback0 = EasyMock.newCapture();
@@ -297,17 +280,15 @@ public void testGetSetNull() throws Exception {
         // Second get() should get the produced data and return the new values
         final Capture> secondGetReadToEndCallback = EasyMock.newCapture();
         storeLog.readToEnd(EasyMock.capture(secondGetReadToEndCallback));
-        PowerMock.expectLastCall().andAnswer(new IAnswer() {
-            @Override
-            public Object answer() throws Throwable {
-                capturedConsumedCallback.getValue().onCompletion(null, new ConsumerRecord<>(TOPIC, 0, 0, 0L, TimestampType.CREATE_TIME, 0L, 0, 0, (byte[]) null, TP0_VALUE.array()));
-                capturedConsumedCallback.getValue().onCompletion(null, new ConsumerRecord<>(TOPIC, 1, 0, 0L, TimestampType.CREATE_TIME, 0L, 0, 0, TP1_KEY.array(), (byte[]) null));
-                secondGetReadToEndCallback.getValue().onCompletion(null, null);
-                return null;
-            }
+        PowerMock.expectLastCall().andAnswer(() -> {
+            capturedConsumedCallback.getValue().onCompletion(null, new ConsumerRecord<>(TOPIC, 0, 0, 0L, TimestampType.CREATE_TIME, 0L, 0, 0, (byte[]) null, TP0_VALUE.array()));
+            capturedConsumedCallback.getValue().onCompletion(null, new ConsumerRecord<>(TOPIC, 1, 0, 0L, TimestampType.CREATE_TIME, 0L, 0, 0, TP1_KEY.array(), (byte[]) null));
+            secondGetReadToEndCallback.getValue().onCompletion(null, null);
+            return null;
         });
 
         expectStop();
+        expectClusterId();
 
         PowerMock.replayAll();
 
@@ -335,16 +316,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();
 
@@ -354,7 +328,7 @@ public void onCompletion(Throwable error, Map result) {
     @Test
     public void testSetFailure() throws Exception {
         expectConfigure();
-        expectStart(Collections.EMPTY_LIST);
+        expectStart(Collections.emptyList());
         expectStop();
 
         // Set offsets
@@ -368,6 +342,8 @@ public void testSetFailure() throws Exception {
         storeLog.send(EasyMock.aryEq(TP2_KEY.array()), EasyMock.aryEq(TP2_VALUE.array()), EasyMock.capture(callback2));
         PowerMock.expectLastCall();
 
+        expectClusterId();
+
         PowerMock.replayAll();
 
         store.configure(DEFAULT_DISTRIBUTED_CONFIG);
@@ -435,6 +411,11 @@ private void expectStop() {
         PowerMock.expectLastCall();
     }
 
+    private void expectClusterId() {
+        PowerMock.mockStaticPartial(ConnectUtils.class, "lookupKafkaClusterId");
+        EasyMock.expect(ConnectUtils.lookupKafkaClusterId(EasyMock.anyObject())).andReturn("test-cluster").anyTimes();
+    }
+
     private static ByteBuffer buffer(String v) {
         return ByteBuffer.wrap(v.getBytes());
     }
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/storage/KafkaStatusBackingStoreFormatTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/storage/KafkaStatusBackingStoreFormatTest.java
new file mode 100644
index 0000000000000..53f481ac9a246
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/storage/KafkaStatusBackingStoreFormatTest.java
@@ -0,0 +1,302 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.connect.storage;
+
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.producer.Callback;
+import org.apache.kafka.common.errors.TimeoutException;
+import org.apache.kafka.common.errors.UnknownServerException;
+import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.connect.json.JsonConverter;
+import org.apache.kafka.connect.runtime.TopicStatus;
+import org.apache.kafka.connect.util.ConnectorTaskId;
+import org.apache.kafka.connect.util.KafkaBasedLog;
+import org.easymock.Capture;
+import org.easymock.EasyMockSupport;
+import org.easymock.Mock;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.powermock.modules.junit4.PowerMockRunner;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.concurrent.ConcurrentHashMap;
+
+import static org.apache.kafka.connect.json.JsonConverterConfig.SCHEMAS_ENABLE_CONFIG;
+import static org.apache.kafka.connect.storage.KafkaStatusBackingStore.CONNECTOR_STATUS_PREFIX;
+import static org.apache.kafka.connect.storage.KafkaStatusBackingStore.TASK_STATUS_PREFIX;
+import static org.apache.kafka.connect.storage.KafkaStatusBackingStore.TOPIC_STATUS_PREFIX;
+import static org.apache.kafka.connect.storage.KafkaStatusBackingStore.TOPIC_STATUS_SEPARATOR;
+import static org.easymock.EasyMock.capture;
+import static org.easymock.EasyMock.eq;
+import static org.easymock.EasyMock.expectLastCall;
+import static org.easymock.EasyMock.newCapture;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+@RunWith(PowerMockRunner.class)
+public class KafkaStatusBackingStoreFormatTest extends EasyMockSupport {
+
+    private static final String STATUS_TOPIC = "status-topic";
+    private static final String FOO_TOPIC = "foo-topic";
+    private static final String FOO_CONNECTOR = "foo-source";
+    private static final String BAR_TOPIC = "bar-topic";
+
+    private Time time;
+    private KafkaStatusBackingStore store;
+    private JsonConverter converter;
+    @Mock
+    private KafkaBasedLog kafkaBasedLog;
+
+    @Before
+    public void setup() {
+        time = new MockTime();
+        converter = new JsonConverter();
+        converter.configure(Collections.singletonMap(SCHEMAS_ENABLE_CONFIG, false), false);
+        store = new KafkaStatusBackingStore(new MockTime(), converter, STATUS_TOPIC, kafkaBasedLog);
+    }
+
+    @Test
+    public void readInvalidStatus() {
+        String key = "status-unknown";
+        byte[] value = new byte[0];
+        ConsumerRecord statusRecord = new ConsumerRecord<>(STATUS_TOPIC, 0, 0, key, value);
+        assertTrue(store.connectors().isEmpty());
+        assertTrue(store.tasks.isEmpty());
+        assertTrue(store.topics.isEmpty());
+        store.read(statusRecord);
+        assertTrue(store.connectors().isEmpty());
+        assertTrue(store.tasks.isEmpty());
+        assertTrue(store.topics.isEmpty());
+
+        key = CONNECTOR_STATUS_PREFIX;
+        statusRecord = new ConsumerRecord<>(STATUS_TOPIC, 0, 0, key, value);
+        assertTrue(store.connectors().isEmpty());
+        store.read(statusRecord);
+        assertTrue(store.connectors().isEmpty());
+
+        key = TASK_STATUS_PREFIX;
+        statusRecord = new ConsumerRecord<>(STATUS_TOPIC, 0, 0, key, value);
+        assertTrue(store.tasks.isEmpty());
+        store.read(statusRecord);
+        assertTrue(store.tasks.isEmpty());
+
+        key = TASK_STATUS_PREFIX + FOO_CONNECTOR + "-#";
+        statusRecord = new ConsumerRecord<>(STATUS_TOPIC, 0, 0, key, value);
+        assertTrue(store.tasks.isEmpty());
+        store.read(statusRecord);
+        assertTrue(store.tasks.isEmpty());
+
+        key = TOPIC_STATUS_PREFIX;
+        statusRecord = new ConsumerRecord<>(STATUS_TOPIC, 0, 0, key, value);
+        assertTrue(store.topics.isEmpty());
+        store.read(statusRecord);
+        assertTrue(store.topics.isEmpty());
+
+        key = TOPIC_STATUS_PREFIX + TOPIC_STATUS_SEPARATOR;
+        statusRecord = new ConsumerRecord<>(STATUS_TOPIC, 0, 0, key, value);
+        assertTrue(store.topics.isEmpty());
+        store.read(statusRecord);
+        assertTrue(store.topics.isEmpty());
+
+        key = TOPIC_STATUS_PREFIX + FOO_TOPIC + ":";
+        statusRecord = new ConsumerRecord<>(STATUS_TOPIC, 0, 0, key, value);
+        assertTrue(store.topics.isEmpty());
+        store.read(statusRecord);
+        assertTrue(store.topics.isEmpty());
+
+        key = TOPIC_STATUS_PREFIX + FOO_TOPIC + TOPIC_STATUS_SEPARATOR;
+        statusRecord = new ConsumerRecord<>(STATUS_TOPIC, 0, 0, key, value);
+        assertTrue(store.topics.isEmpty());
+        store.read(statusRecord);
+        assertTrue(store.topics.isEmpty());
+    }
+
+    @Test
+    public void readInvalidStatusValue() {
+        String key = CONNECTOR_STATUS_PREFIX + FOO_CONNECTOR;
+        byte[] value = "invalid".getBytes();
+        ConsumerRecord statusRecord = new ConsumerRecord<>(STATUS_TOPIC, 0, 0, key, value);
+        assertTrue(store.connectors().isEmpty());
+        store.read(statusRecord);
+        assertTrue(store.connectors().isEmpty());
+
+        key = TASK_STATUS_PREFIX + FOO_CONNECTOR + "-0";
+        statusRecord = new ConsumerRecord<>(STATUS_TOPIC, 0, 0, key, value);
+        assertTrue(store.tasks.isEmpty());
+        store.read(statusRecord);
+        assertTrue(store.tasks.isEmpty());
+
+        key = TOPIC_STATUS_PREFIX + FOO_TOPIC + TOPIC_STATUS_SEPARATOR + FOO_CONNECTOR;
+        statusRecord = new ConsumerRecord<>(STATUS_TOPIC, 0, 0, key, value);
+        assertTrue(store.topics.isEmpty());
+        store.read(statusRecord);
+        assertTrue(store.topics.isEmpty());
+    }
+
+    @Test
+    public void readTopicStatus() {
+        TopicStatus topicStatus = new TopicStatus(FOO_TOPIC, new ConnectorTaskId(FOO_CONNECTOR, 0), Time.SYSTEM.milliseconds());
+        String key = TOPIC_STATUS_PREFIX + FOO_TOPIC + TOPIC_STATUS_SEPARATOR + FOO_CONNECTOR;
+        byte[] value = store.serializeTopicStatus(topicStatus);
+        ConsumerRecord statusRecord = new ConsumerRecord<>(STATUS_TOPIC, 0, 0, key, value);
+        store.read(statusRecord);
+        assertTrue(store.topics.containsKey(FOO_CONNECTOR));
+        assertTrue(store.topics.get(FOO_CONNECTOR).containsKey(FOO_TOPIC));
+        assertEquals(topicStatus, store.topics.get(FOO_CONNECTOR).get(FOO_TOPIC));
+    }
+
+    @Test
+    public void deleteTopicStatus() {
+        TopicStatus topicStatus = new TopicStatus("foo", new ConnectorTaskId("bar", 0), Time.SYSTEM.milliseconds());
+        store.topics.computeIfAbsent("bar", k -> new ConcurrentHashMap<>()).put("foo", topicStatus);
+        assertTrue(store.topics.containsKey("bar"));
+        assertTrue(store.topics.get("bar").containsKey("foo"));
+        assertEquals(topicStatus, store.topics.get("bar").get("foo"));
+        // should return null
+        byte[] value = store.serializeTopicStatus(null);
+        ConsumerRecord statusRecord = new ConsumerRecord<>(STATUS_TOPIC, 0, 0, "status-topic-foo:connector-bar", value);
+        store.read(statusRecord);
+        assertTrue(store.topics.containsKey("bar"));
+        assertFalse(store.topics.get("bar").containsKey("foo"));
+        assertEquals(Collections.emptyMap(), store.topics.get("bar"));
+    }
+
+    @Test
+    public void putTopicState() {
+        TopicStatus topicStatus = new TopicStatus(FOO_TOPIC, new ConnectorTaskId(FOO_CONNECTOR, 0), time.milliseconds());
+        String key = TOPIC_STATUS_PREFIX + FOO_TOPIC + TOPIC_STATUS_SEPARATOR + FOO_CONNECTOR;
+        Capture valueCapture = newCapture();
+        Capture callbackCapture = newCapture();
+        kafkaBasedLog.send(eq(key), capture(valueCapture), capture(callbackCapture));
+        expectLastCall()
+                .andAnswer(() -> {
+                    callbackCapture.getValue().onCompletion(null, null);
+                    return null;
+                });
+        replayAll();
+
+        store.put(topicStatus);
+        // check capture state
+        assertEquals(topicStatus, store.parseTopicStatus(valueCapture.getValue()));
+        // state is not visible until read back from the log
+        assertEquals(null, store.getTopic(FOO_CONNECTOR, FOO_TOPIC));
+
+        ConsumerRecord statusRecord = new ConsumerRecord<>(STATUS_TOPIC, 0, 0, key, valueCapture.getValue());
+        store.read(statusRecord);
+        assertEquals(topicStatus, store.getTopic(FOO_CONNECTOR, FOO_TOPIC));
+        assertEquals(new HashSet<>(Collections.singletonList(topicStatus)), new HashSet<>(store.getAllTopics(FOO_CONNECTOR)));
+
+        verifyAll();
+    }
+
+    @Test
+    public void putTopicStateRetriableFailure() {
+        TopicStatus topicStatus = new TopicStatus(FOO_TOPIC, new ConnectorTaskId(FOO_CONNECTOR, 0), time.milliseconds());
+        String key = TOPIC_STATUS_PREFIX + FOO_TOPIC + TOPIC_STATUS_SEPARATOR + FOO_CONNECTOR;
+        Capture valueCapture = newCapture();
+        Capture callbackCapture = newCapture();
+        kafkaBasedLog.send(eq(key), capture(valueCapture), capture(callbackCapture));
+        expectLastCall()
+                .andAnswer(() -> {
+                    callbackCapture.getValue().onCompletion(null, new TimeoutException());
+                    return null;
+                })
+                .andAnswer(() -> {
+                    callbackCapture.getValue().onCompletion(null, null);
+                    return null;
+                });
+
+        replayAll();
+        store.put(topicStatus);
+
+        // check capture state
+        assertEquals(topicStatus, store.parseTopicStatus(valueCapture.getValue()));
+        // state is not visible until read back from the log
+        assertEquals(null, store.getTopic(FOO_CONNECTOR, FOO_TOPIC));
+
+        verifyAll();
+    }
+
+    @Test
+    public void putTopicStateNonRetriableFailure() {
+        TopicStatus topicStatus = new TopicStatus(FOO_TOPIC, new ConnectorTaskId(FOO_CONNECTOR, 0), time.milliseconds());
+        String key = TOPIC_STATUS_PREFIX + FOO_TOPIC + TOPIC_STATUS_SEPARATOR + FOO_CONNECTOR;
+        Capture valueCapture = newCapture();
+        Capture callbackCapture = newCapture();
+        kafkaBasedLog.send(eq(key), capture(valueCapture), capture(callbackCapture));
+        expectLastCall()
+                .andAnswer(() -> {
+                    callbackCapture.getValue().onCompletion(null, new UnknownServerException());
+                    return null;
+                });
+
+        replayAll();
+
+        // the error is logged and ignored
+        store.put(topicStatus);
+
+        // check capture state
+        assertEquals(topicStatus, store.parseTopicStatus(valueCapture.getValue()));
+        // state is not visible until read back from the log
+        assertEquals(null, store.getTopic(FOO_CONNECTOR, FOO_TOPIC));
+
+        verifyAll();
+    }
+
+    @Test
+    public void putTopicStateShouldOverridePreviousState() {
+        TopicStatus firstTopicStatus = new TopicStatus(FOO_TOPIC, new ConnectorTaskId(FOO_CONNECTOR,
+                0), time.milliseconds());
+        time.sleep(1000);
+        TopicStatus secondTopicStatus = new TopicStatus(BAR_TOPIC, new ConnectorTaskId(FOO_CONNECTOR,
+                0), time.milliseconds());
+        String firstKey = TOPIC_STATUS_PREFIX + FOO_TOPIC + TOPIC_STATUS_SEPARATOR + FOO_CONNECTOR;
+        String secondKey = TOPIC_STATUS_PREFIX + BAR_TOPIC + TOPIC_STATUS_SEPARATOR + FOO_CONNECTOR;
+        Capture valueCapture = newCapture();
+        Capture callbackCapture = newCapture();
+        kafkaBasedLog.send(eq(secondKey), capture(valueCapture), capture(callbackCapture));
+        expectLastCall()
+                .andAnswer(() -> {
+                    callbackCapture.getValue().onCompletion(null, null);
+                    // The second status record is read soon after it's persisted in the status topic
+                    ConsumerRecord statusRecord = new ConsumerRecord<>(STATUS_TOPIC, 0, 0, secondKey, valueCapture.getValue());
+                    store.read(statusRecord);
+                    return null;
+                });
+        replayAll();
+
+        byte[] value = store.serializeTopicStatus(firstTopicStatus);
+        ConsumerRecord statusRecord = new ConsumerRecord<>(STATUS_TOPIC, 0, 0, firstKey, value);
+        store.read(statusRecord);
+        store.put(secondTopicStatus);
+
+        // check capture state
+        assertEquals(secondTopicStatus, store.parseTopicStatus(valueCapture.getValue()));
+        assertEquals(firstTopicStatus, store.getTopic(FOO_CONNECTOR, FOO_TOPIC));
+        assertEquals(secondTopicStatus, store.getTopic(FOO_CONNECTOR, BAR_TOPIC));
+        assertEquals(new HashSet<>(Arrays.asList(firstTopicStatus, secondTopicStatus)),
+                new HashSet<>(store.getAllTopics(FOO_CONNECTOR)));
+
+        verifyAll();
+    }
+
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/storage/KafkaStatusBackingStoreTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/storage/KafkaStatusBackingStoreTest.java
index e2f5a406c195e..bd8e7a512a3cb 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/storage/KafkaStatusBackingStoreTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/storage/KafkaStatusBackingStoreTest.java
@@ -18,6 +18,7 @@
 
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.producer.Callback;
+import org.apache.kafka.common.config.ConfigException;
 import org.apache.kafka.common.errors.TimeoutException;
 import org.apache.kafka.common.errors.UnknownServerException;
 import org.apache.kafka.common.record.TimestampType;
@@ -27,15 +28,23 @@
 import org.apache.kafka.connect.data.Struct;
 import org.apache.kafka.connect.runtime.ConnectorStatus;
 import org.apache.kafka.connect.runtime.TaskStatus;
+import org.apache.kafka.connect.runtime.WorkerConfig;
+import org.apache.kafka.connect.runtime.distributed.DistributedConfig;
 import org.apache.kafka.connect.util.ConnectorTaskId;
 import org.apache.kafka.connect.util.KafkaBasedLog;
 import org.easymock.Capture;
 import org.easymock.EasyMock;
 import org.easymock.EasyMockSupport;
 import org.easymock.IAnswer;
+import org.easymock.Mock;
+import org.junit.Before;
 import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.powermock.modules.junit4.PowerMockRunner;
 
+import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.Map;
 
 import static org.easymock.EasyMock.anyObject;
@@ -45,8 +54,11 @@
 import static org.easymock.EasyMock.expectLastCall;
 import static org.easymock.EasyMock.newCapture;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
 
 @SuppressWarnings("unchecked")
+@RunWith(PowerMockRunner.class)
 public class KafkaStatusBackingStoreTest extends EasyMockSupport {
 
     private static final String STATUS_TOPIC = "status-topic";
@@ -54,12 +66,34 @@ public class KafkaStatusBackingStoreTest extends EasyMockSupport {
     private static final String CONNECTOR = "conn";
     private static final ConnectorTaskId TASK = new ConnectorTaskId(CONNECTOR, 0);
 
+    private KafkaStatusBackingStore store;
+    @Mock
+    Converter converter;
+    @Mock
+    private KafkaBasedLog kafkaBasedLog;
+    @Mock
+    WorkerConfig workerConfig;
+
+    @Before
+    public void setup() {
+        store = new KafkaStatusBackingStore(new MockTime(), converter, STATUS_TOPIC, kafkaBasedLog);
+    }
+
     @Test
-    public void putConnectorState() {
-        KafkaBasedLog kafkaBasedLog = mock(KafkaBasedLog.class);
-        Converter converter = mock(Converter.class);
-        KafkaStatusBackingStore store = new KafkaStatusBackingStore(new MockTime(), converter, STATUS_TOPIC, kafkaBasedLog);
+    public void misconfigurationOfStatusBackingStore() {
+        expect(workerConfig.getString(DistributedConfig.STATUS_STORAGE_TOPIC_CONFIG)).andReturn(null);
+        expect(workerConfig.getString(DistributedConfig.STATUS_STORAGE_TOPIC_CONFIG)).andReturn("   ");
+        replayAll();
 
+        Exception e = assertThrows(ConfigException.class, () -> store.configure(workerConfig));
+        assertEquals("Must specify topic for connector status.", e.getMessage());
+        e = assertThrows(ConfigException.class, () -> store.configure(workerConfig));
+        assertEquals("Must specify topic for connector status.", e.getMessage());
+        verifyAll();
+    }
+
+    @Test
+    public void putConnectorState() {
         byte[] value = new byte[0];
         expect(converter.fromConnectData(eq(STATUS_TOPIC), anyObject(Schema.class), anyObject(Struct.class)))
                 .andStubReturn(value);
@@ -87,10 +121,6 @@ public Void answer() throws Throwable {
 
     @Test
     public void putConnectorStateRetriableFailure() {
-        KafkaBasedLog kafkaBasedLog = mock(KafkaBasedLog.class);
-        Converter converter = mock(Converter.class);
-        KafkaStatusBackingStore store = new KafkaStatusBackingStore(new MockTime(), converter, STATUS_TOPIC, kafkaBasedLog);
-
         byte[] value = new byte[0];
         expect(converter.fromConnectData(eq(STATUS_TOPIC), anyObject(Schema.class), anyObject(Struct.class)))
                 .andStubReturn(value);
@@ -125,10 +155,6 @@ public Void answer() throws Throwable {
 
     @Test
     public void putConnectorStateNonRetriableFailure() {
-        KafkaBasedLog kafkaBasedLog = mock(KafkaBasedLog.class);
-        Converter converter = mock(Converter.class);
-        KafkaStatusBackingStore store = new KafkaStatusBackingStore(new MockTime(), converter, STATUS_TOPIC, kafkaBasedLog);
-
         byte[] value = new byte[0];
         expect(converter.fromConnectData(eq(STATUS_TOPIC), anyObject(Schema.class), anyObject(Struct.class)))
                 .andStubReturn(value);
@@ -160,10 +186,6 @@ public void putSafeConnectorIgnoresStaleStatus() {
         byte[] value = new byte[0];
         String otherWorkerId = "anotherhost:8083";
 
-        KafkaBasedLog kafkaBasedLog = mock(KafkaBasedLog.class);
-        Converter converter = mock(Converter.class);
-        KafkaStatusBackingStore store = new KafkaStatusBackingStore(new MockTime(), converter, STATUS_TOPIC, kafkaBasedLog);
-
         // the persisted came from a different host and has a newer generation
         Map statusMap = new HashMap<>();
         statusMap.put("worker_id", otherWorkerId);
@@ -188,10 +210,6 @@ public void putSafeConnectorIgnoresStaleStatus() {
 
     @Test
     public void putSafeWithNoPreviousValueIsPropagated() {
-        final Converter converter = mock(Converter.class);
-        final KafkaBasedLog kafkaBasedLog = mock(KafkaBasedLog.class);
-        final KafkaStatusBackingStore store = new KafkaStatusBackingStore(new MockTime(), converter, STATUS_TOPIC, kafkaBasedLog);
-
         final byte[] value = new byte[0];
 
         final Capture statusValueStruct = newCapture();
@@ -217,10 +235,6 @@ public void putSafeWithNoPreviousValueIsPropagated() {
     public void putSafeOverridesValueSetBySameWorker() {
         final byte[] value = new byte[0];
 
-        KafkaBasedLog kafkaBasedLog = mock(KafkaBasedLog.class);
-        Converter converter = mock(Converter.class);
-        final KafkaStatusBackingStore store = new KafkaStatusBackingStore(new MockTime(), converter, STATUS_TOPIC, kafkaBasedLog);
-
         // the persisted came from the same host, but has a newer generation
         Map firstStatusRead = new HashMap<>();
         firstStatusRead.put("worker_id", WORKER_ID);
@@ -267,10 +281,6 @@ public void putConnectorStateShouldOverride() {
         final byte[] value = new byte[0];
         String otherWorkerId = "anotherhost:8083";
 
-        KafkaBasedLog kafkaBasedLog = mock(KafkaBasedLog.class);
-        Converter converter = mock(Converter.class);
-        final KafkaStatusBackingStore store = new KafkaStatusBackingStore(new MockTime(), converter, STATUS_TOPIC, kafkaBasedLog);
-
         // the persisted came from a different host and has a newer generation
         Map firstStatusRead = new HashMap<>();
         firstStatusRead.put("worker_id", otherWorkerId);
@@ -315,10 +325,6 @@ public Void answer() throws Throwable {
     public void readConnectorState() {
         byte[] value = new byte[0];
 
-        KafkaBasedLog kafkaBasedLog = mock(KafkaBasedLog.class);
-        Converter converter = mock(Converter.class);
-        KafkaStatusBackingStore store = new KafkaStatusBackingStore(new MockTime(), converter, STATUS_TOPIC, kafkaBasedLog);
-
         Map statusMap = new HashMap<>();
         statusMap.put("worker_id", WORKER_ID);
         statusMap.put("state", "RUNNING");
@@ -339,10 +345,6 @@ public void readConnectorState() {
 
     @Test
     public void putTaskState() {
-        KafkaBasedLog kafkaBasedLog = mock(KafkaBasedLog.class);
-        Converter converter = mock(Converter.class);
-        KafkaStatusBackingStore store = new KafkaStatusBackingStore(new MockTime(), converter, STATUS_TOPIC, kafkaBasedLog);
-
         byte[] value = new byte[0];
         expect(converter.fromConnectData(eq(STATUS_TOPIC), anyObject(Schema.class), anyObject(Struct.class)))
                 .andStubReturn(value);
@@ -372,10 +374,6 @@ public Void answer() throws Throwable {
     public void readTaskState() {
         byte[] value = new byte[0];
 
-        KafkaBasedLog kafkaBasedLog = mock(KafkaBasedLog.class);
-        Converter converter = mock(Converter.class);
-        KafkaStatusBackingStore store = new KafkaStatusBackingStore(new MockTime(), converter, STATUS_TOPIC, kafkaBasedLog);
-
         Map statusMap = new HashMap<>();
         statusMap.put("worker_id", WORKER_ID);
         statusMap.put("state", "RUNNING");
@@ -394,6 +392,69 @@ public void readTaskState() {
         verifyAll();
     }
 
+    @Test
+    public void deleteConnectorState() {
+        final byte[] value = new byte[0];
+        Map statusMap = new HashMap<>();
+        statusMap.put("worker_id", WORKER_ID);
+        statusMap.put("state", "RUNNING");
+        statusMap.put("generation", 0L);
+
+        converter.fromConnectData(eq(STATUS_TOPIC), anyObject(Schema.class), anyObject(Struct.class));
+        EasyMock.expectLastCall().andReturn(value);
+        kafkaBasedLog.send(eq("status-connector-" + CONNECTOR), eq(value), anyObject(Callback.class));
+        expectLastCall();
+
+        converter.fromConnectData(eq(STATUS_TOPIC), anyObject(Schema.class), anyObject(Struct.class));
+        EasyMock.expectLastCall().andReturn(value);
+        kafkaBasedLog.send(eq("status-task-conn-0"), eq(value), anyObject(Callback.class));
+        expectLastCall();
+
+        expect(converter.toConnectData(STATUS_TOPIC, value)).andReturn(new SchemaAndValue(null, statusMap));
+
+        replayAll();
+
+        ConnectorStatus connectorStatus = new ConnectorStatus(CONNECTOR, ConnectorStatus.State.RUNNING, WORKER_ID, 0);
+        store.put(connectorStatus);
+        TaskStatus taskStatus = new TaskStatus(TASK, TaskStatus.State.RUNNING, WORKER_ID, 0);
+        store.put(taskStatus);
+        store.read(consumerRecord(0, "status-task-conn-0", value));
+
+        assertEquals(new HashSet<>(Collections.singletonList(CONNECTOR)), store.connectors());
+        assertEquals(new HashSet<>(Collections.singletonList(taskStatus)), new HashSet<>(store.getAll(CONNECTOR)));
+        store.read(consumerRecord(0, "status-connector-conn", null));
+        assertTrue(store.connectors().isEmpty());
+        assertTrue(store.getAll(CONNECTOR).isEmpty());
+        verifyAll();
+    }
+
+    @Test
+    public void deleteTaskState() {
+        final byte[] value = new byte[0];
+        Map statusMap = new HashMap<>();
+        statusMap.put("worker_id", WORKER_ID);
+        statusMap.put("state", "RUNNING");
+        statusMap.put("generation", 0L);
+
+        converter.fromConnectData(eq(STATUS_TOPIC), anyObject(Schema.class), anyObject(Struct.class));
+        EasyMock.expectLastCall().andReturn(value);
+        kafkaBasedLog.send(eq("status-task-conn-0"), eq(value), anyObject(Callback.class));
+        expectLastCall();
+
+        expect(converter.toConnectData(STATUS_TOPIC, value)).andReturn(new SchemaAndValue(null, statusMap));
+
+        replayAll();
+
+        TaskStatus taskStatus = new TaskStatus(TASK, TaskStatus.State.RUNNING, WORKER_ID, 0);
+        store.put(taskStatus);
+        store.read(consumerRecord(0, "status-task-conn-0", value));
+
+        assertEquals(new HashSet<>(Collections.singletonList(taskStatus)), new HashSet<>(store.getAll(CONNECTOR)));
+        store.read(consumerRecord(0, "status-task-conn-0", null));
+        assertTrue(store.getAll(CONNECTOR).isEmpty());
+        verifyAll();
+    }
+
     private static ConsumerRecord consumerRecord(long offset, String key, byte[] value) {
         return new ConsumerRecord<>(STATUS_TOPIC, 0, offset, System.currentTimeMillis(),
                 TimestampType.CREATE_TIME, 0L, 0, 0, key, value);
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/storage/OffsetStorageWriterTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/storage/OffsetStorageWriterTest.java
index 5a671ca6a733d..3aad0d5314b3f 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/storage/OffsetStorageWriterTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/storage/OffsetStorageWriterTest.java
@@ -165,7 +165,7 @@ public void testFlushFailureReplacesOffsets() throws Exception {
     }
 
     @Test(expected = ConnectException.class)
-    public void testAlreadyFlushing() throws Exception {
+    public void testAlreadyFlushing() {
         @SuppressWarnings("unchecked")
         final Callback callback = PowerMock.createMock(Callback.class);
         // Trigger the send, but don't invoke the callback so we'll still be mid-flush
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/util/ConnectUtilsTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/util/ConnectUtilsTest.java
new file mode 100644
index 0000000000000..f92143a6e0b08
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/util/ConnectUtilsTest.java
@@ -0,0 +1,103 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.clients.CommonClientConfigs;
+import org.apache.kafka.clients.admin.MockAdminClient;
+import org.apache.kafka.common.Node;
+import org.apache.kafka.connect.errors.ConnectException;
+import org.apache.kafka.connect.runtime.WorkerConfig;
+import org.apache.kafka.connect.runtime.distributed.DistributedConfig;
+import org.apache.kafka.connect.runtime.standalone.StandaloneConfig;
+import org.junit.Test;
+
+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.assertNull;
+
+public class ConnectUtilsTest {
+
+    @Test
+    public void testLookupKafkaClusterId() {
+        final Node broker1 = new Node(0, "dummyHost-1", 1234);
+        final Node broker2 = new Node(1, "dummyHost-2", 1234);
+        List cluster = Arrays.asList(broker1, broker2);
+        MockAdminClient adminClient = new MockAdminClient.Builder().
+            brokers(cluster).build();
+        assertEquals(MockAdminClient.DEFAULT_CLUSTER_ID, ConnectUtils.lookupKafkaClusterId(adminClient));
+    }
+
+    @Test
+    public void testLookupNullKafkaClusterId() {
+        final Node broker1 = new Node(0, "dummyHost-1", 1234);
+        final Node broker2 = new Node(1, "dummyHost-2", 1234);
+        List cluster = Arrays.asList(broker1, broker2);
+        MockAdminClient adminClient = new MockAdminClient.Builder().
+            brokers(cluster).clusterId(null).build();
+        assertNull(ConnectUtils.lookupKafkaClusterId(adminClient));
+    }
+
+    @Test(expected = ConnectException.class)
+    public void testLookupKafkaClusterIdTimeout() {
+        final Node broker1 = new Node(0, "dummyHost-1", 1234);
+        final Node broker2 = new Node(1, "dummyHost-2", 1234);
+        List cluster = Arrays.asList(broker1, broker2);
+        MockAdminClient adminClient = new MockAdminClient.Builder().
+            brokers(cluster).build();
+        adminClient.timeoutNextRequest(1);
+
+        ConnectUtils.lookupKafkaClusterId(adminClient);
+    }
+
+    @Test
+    public void testAddMetricsContextPropertiesDistributed() {
+        Map props = new HashMap<>();
+        props.put(DistributedConfig.GROUP_ID_CONFIG, "connect-cluster");
+        props.put(DistributedConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092");
+        props.put(DistributedConfig.CONFIG_TOPIC_CONFIG, "connect-configs");
+        props.put(DistributedConfig.OFFSET_STORAGE_TOPIC_CONFIG, "connect-offsets");
+        props.put(DistributedConfig.STATUS_STORAGE_TOPIC_CONFIG, "connect-status");
+        props.put(DistributedConfig.KEY_CONVERTER_CLASS_CONFIG, "org.apache.kafka.connect.json.JsonConverter");
+        props.put(DistributedConfig.VALUE_CONVERTER_CLASS_CONFIG, "org.apache.kafka.connect.json.JsonConverter");
+        DistributedConfig config = new DistributedConfig(props);
+
+        Map prop = new HashMap<>();
+        ConnectUtils.addMetricsContextProperties(prop, config, "cluster-1");
+        assertEquals("connect-cluster", prop.get(CommonClientConfigs.METRICS_CONTEXT_PREFIX + WorkerConfig.CONNECT_GROUP_ID));
+        assertEquals("cluster-1", prop.get(CommonClientConfigs.METRICS_CONTEXT_PREFIX + WorkerConfig.CONNECT_KAFKA_CLUSTER_ID));
+    }
+
+    @Test
+    public void testAddMetricsContextPropertiesStandalone() {
+        Map props = new HashMap<>();
+        props.put(StandaloneConfig.OFFSET_STORAGE_FILE_FILENAME_CONFIG, "offsetStorageFile");
+        props.put(StandaloneConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092");
+        props.put(StandaloneConfig.KEY_CONVERTER_CLASS_CONFIG, "org.apache.kafka.connect.json.JsonConverter");
+        props.put(StandaloneConfig.VALUE_CONVERTER_CLASS_CONFIG, "org.apache.kafka.connect.json.JsonConverter");
+        StandaloneConfig config = new StandaloneConfig(props);
+
+        Map prop = new HashMap<>();
+        ConnectUtils.addMetricsContextProperties(prop, config, "cluster-1");
+        assertEquals(null, prop.get(CommonClientConfigs.METRICS_CONTEXT_PREFIX + WorkerConfig.CONNECT_GROUP_ID));
+        assertEquals("cluster-1", prop.get(CommonClientConfigs.METRICS_CONTEXT_PREFIX + WorkerConfig.CONNECT_KAFKA_CLUSTER_ID));
+
+    }
+}
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 b2c164dc1c945..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,10 +30,12 @@
 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;
 import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.common.utils.MockTime;
 import org.easymock.Capture;
 import org.easymock.EasyMock;
 import org.junit.Before;
@@ -130,8 +132,9 @@ public void onCompletion(Throwable error, ConsumerRecord record)
         }
     };
 
+    @SuppressWarnings("unchecked")
     @Before
-    public void setUp() throws Exception {
+    public void setUp() {
         store = PowerMock.createPartialMock(KafkaBasedLog.class, new String[]{"createConsumer", "createProducer"},
                 TOPIC, PRODUCER_PROPS, CONSUMER_PROPS, consumedCallback, time, initializer);
         consumer = new MockConsumer<>(OffsetResetStrategy.EARLIEST);
@@ -368,7 +371,7 @@ public void run() {
     }
 
     @Test
-    public void testConsumerError() throws Exception {
+    public void testPollConsumerError() throws Exception {
         expectStart();
         expectStop();
 
@@ -386,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());
                     }
                 });
 
@@ -421,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();
@@ -483,4 +557,4 @@ private static ByteBuffer buffer(String v) {
         return ByteBuffer.wrap(v.getBytes());
     }
 
-}
\ No newline at end of file
+}
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/MockTime.java b/connect/runtime/src/test/java/org/apache/kafka/connect/util/MockTime.java
deleted file mode 100644
index 554c5f982fad6..0000000000000
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/util/MockTime.java
+++ /dev/null
@@ -1,55 +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.util;
-
-import org.apache.kafka.common.utils.Time;
-
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicLong;
-
-/**
- * A clock that you can manually advance by calling sleep
- */
-public class MockTime implements Time {
-
-    private final AtomicLong nanos;
-
-    public MockTime() {
-        this.nanos = new AtomicLong(System.nanoTime());
-    }
-
-    @Override
-    public long milliseconds() {
-        return TimeUnit.MILLISECONDS.convert(this.nanos.get(), TimeUnit.NANOSECONDS);
-    }
-
-    @Override
-    public long hiResClockMs() {
-        return TimeUnit.NANOSECONDS.toMillis(nanos.get());
-    }
-
-    @Override
-    public long nanoseconds() {
-        return nanos.get();
-    }
-
-    @Override
-    public void sleep(long ms) {
-        this.nanos.addAndGet(TimeUnit.NANOSECONDS.convert(ms, TimeUnit.MILLISECONDS));
-    }
-
-}
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 0a61d3e9042c8..e4b456df61663 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
@@ -17,24 +17,54 @@
 package org.apache.kafka.connect.util;
 
 import org.apache.kafka.clients.NodeApiVersions;
-import org.apache.kafka.clients.admin.MockKafkaAdminClientEnv;
+import org.apache.kafka.clients.admin.AdminClientUnitTestEnv;
+import org.apache.kafka.clients.admin.Config;
+import org.apache.kafka.clients.admin.DescribeTopicsResult;
+import org.apache.kafka.clients.admin.MockAdminClient;
 import org.apache.kafka.clients.admin.NewTopic;
+import org.apache.kafka.clients.admin.TopicDescription;
 import org.apache.kafka.common.Cluster;
+import org.apache.kafka.common.KafkaFuture;
 import org.apache.kafka.common.Node;
 import org.apache.kafka.common.PartitionInfo;
+import org.apache.kafka.common.TopicPartitionInfo;
+import org.apache.kafka.common.config.ConfigException;
+import org.apache.kafka.common.config.ConfigResource;
+import org.apache.kafka.common.config.TopicConfig;
+import org.apache.kafka.common.errors.ClusterAuthorizationException;
+import org.apache.kafka.common.errors.TopicAuthorizationException;
+import org.apache.kafka.common.errors.UnsupportedVersionException;
+import org.apache.kafka.common.message.CreateTopicsResponseData;
+import org.apache.kafka.common.message.CreateTopicsResponseData.CreatableTopicResult;
+import org.apache.kafka.common.message.DescribeConfigsResponseData;
+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.requests.ApiError;
 import org.apache.kafka.common.requests.CreateTopicsResponse;
+import org.apache.kafka.common.requests.DescribeConfigsResponse;
+import org.apache.kafka.common.requests.MetadataResponse;
+import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.connect.errors.ConnectException;
 import org.junit.Test;
 
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
+import static org.apache.kafka.common.message.MetadataResponseData.MetadataResponseTopic;
 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;
 
 public class TopicAdminTest {
 
@@ -44,13 +74,11 @@ public class TopicAdminTest {
      * create no topics, and return false.
      */
     @Test
-    public void returnNullWithApiVersionMismatch() {
+    public void returnNullWithApiVersionMismatchOnCreate() {
         final NewTopic newTopic = TopicAdmin.defineTopic("myTopic").partitions(1).compacted().build();
         Cluster cluster = createCluster(1);
-        try (MockKafkaAdminClientEnv env = new MockKafkaAdminClientEnv(cluster)) {
-            env.kafkaClient().setNode(cluster.controller());
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(new MockTime(), cluster)) {
             env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
-            env.kafkaClient().prepareMetadataUpdate(env.cluster(), Collections.emptySet());
             env.kafkaClient().prepareResponse(createTopicResponseWithUnsupportedVersion(newTopic));
             TopicAdmin admin = new TopicAdmin(null, env.adminClient());
             boolean created = admin.createTopic(newTopic);
@@ -58,15 +86,30 @@ public void returnNullWithApiVersionMismatch() {
         }
     }
 
+    /**
+     * 0.11.0.0 clients can talk with older brokers, but the DESCRIBE_TOPIC API was added in 0.10.0.0. That means,
+     * if our TopicAdmin talks to a pre 0.10.0 broker, it should receive an UnsupportedVersionException, should
+     * create no topics, and return false.
+     */
     @Test
-    public void shouldNotCreateTopicWhenItAlreadyExists() {
-        NewTopic newTopic = TopicAdmin.defineTopic("myTopic").partitions(1).compacted().build();
+    public void throwsWithApiVersionMismatchOnDescribe() {
+        final NewTopic newTopic = TopicAdmin.defineTopic("myTopic").partitions(1).compacted().build();
         Cluster cluster = createCluster(1);
-        try (MockKafkaAdminClientEnv env = new MockKafkaAdminClientEnv(cluster)) {
-            env.kafkaClient().setNode(cluster.controller());
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(new MockTime(), cluster)) {
             env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
-            env.kafkaClient().prepareMetadataUpdate(env.cluster(), Collections.emptySet());
-            env.kafkaClient().prepareResponse(createTopicResponseWithAlreadyExists(newTopic));
+            env.kafkaClient().prepareResponse(describeTopicResponseWithUnsupportedVersion(newTopic));
+            TopicAdmin admin = new TopicAdmin(null, env.adminClient());
+            Exception e = assertThrows(ConnectException.class, () -> admin.describeTopics(newTopic.name()));
+            assertTrue(e.getCause() instanceof UnsupportedVersionException);
+        }
+    }
+
+    @Test
+    public void returnNullWithClusterAuthorizationFailureOnCreate() {
+        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(createTopicResponseWithClusterAuthorizationException(newTopic));
             TopicAdmin admin = new TopicAdmin(null, env.adminClient());
             boolean created = admin.createTopic(newTopic);
             assertFalse(created);
@@ -74,17 +117,113 @@ public void shouldNotCreateTopicWhenItAlreadyExists() {
     }
 
     @Test
-    public void shouldCreateTopicWhenItDoesNotExist() {
-        NewTopic newTopic = TopicAdmin.defineTopic("myTopic").partitions(1).compacted().build();
+    public void throwsWithClusterAuthorizationFailureOnDescribe() {
+        final NewTopic newTopic = TopicAdmin.defineTopic("myTopic").partitions(1).compacted().build();
         Cluster cluster = createCluster(1);
-        try (MockKafkaAdminClientEnv env = new MockKafkaAdminClientEnv(cluster)) {
-            env.kafkaClient().setNode(cluster.controller());
-            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
-            env.kafkaClient().prepareMetadataUpdate(env.cluster(), Collections.emptySet());
-            env.kafkaClient().prepareResponse(createTopicResponse(newTopic));
+        try (AdminClientUnitTestEnv env = new AdminClientUnitTestEnv(new MockTime(), cluster)) {
+            env.kafkaClient().prepareResponse(describeTopicResponseWithClusterAuthorizationException(newTopic));
+            TopicAdmin admin = new TopicAdmin(null, env.adminClient());
+            Exception e = assertThrows(ConnectException.class, () -> admin.describeTopics(newTopic.name()));
+            assertTrue(e.getCause() instanceof ClusterAuthorizationException);
+        }
+    }
+
+    @Test
+    public void returnNullWithTopicAuthorizationFailureOnCreate() {
+        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);
-            assertTrue(created);
+            assertFalse(created);
+        }
+    }
+
+    @Test
+    public void throwsWithTopicAuthorizationFailureOnDescribe() {
+        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(describeTopicResponseWithTopicAuthorizationException(newTopic));
+            TopicAdmin admin = new TopicAdmin(null, env.adminClient());
+            Exception e = assertThrows(ConnectException.class, () -> admin.describeTopics(newTopic.name()));
+            assertTrue(e.getCause() instanceof TopicAuthorizationException);
+        }
+    }
+
+    @Test
+    public void shouldNotCreateTopicWhenItAlreadyExists() {
+        NewTopic newTopic = TopicAdmin.defineTopic("myTopic").partitions(1).compacted().build();
+        Cluster cluster = createCluster(1);
+        try (MockAdminClient mockAdminClient = new MockAdminClient(cluster.nodes(), cluster.nodeById(0))) {
+            TopicPartitionInfo topicPartitionInfo = new TopicPartitionInfo(0, cluster.nodeById(0), cluster.nodes(), Collections.emptyList());
+            mockAdminClient.addTopic(false, "myTopic", Collections.singletonList(topicPartitionInfo), null);
+            TopicAdmin admin = new TopicAdmin(null, mockAdminClient);
+            assertFalse(admin.createTopic(newTopic));
+        }
+    }
+
+    @Test
+    public void shouldCreateTopicWithPartitionsWhenItDoesNotExist() {
+        for (int numBrokers = 1; numBrokers < 10; ++numBrokers) {
+            int expectedReplicas = Math.min(3, numBrokers);
+            int maxDefaultRf = Math.min(numBrokers, 5);
+            for (int numPartitions = 1; numPartitions < 30; ++numPartitions) {
+                NewTopic newTopic = TopicAdmin.defineTopic("myTopic").partitions(numPartitions).compacted().build();
+
+                // Try clusters with no default replication factor or default partitions
+                assertTopicCreation(numBrokers, newTopic, null, null, expectedReplicas, numPartitions);
+
+                // Try clusters with different default partitions
+                for (int defaultPartitions = 1; defaultPartitions < 20; ++defaultPartitions) {
+                    assertTopicCreation(numBrokers, newTopic, defaultPartitions, null, expectedReplicas, numPartitions);
+                }
+
+                // Try clusters with different default replication factors
+                for (int defaultRF = 1; defaultRF < maxDefaultRf; ++defaultRF) {
+                    assertTopicCreation(numBrokers, newTopic, null, defaultRF, defaultRF, numPartitions);
+                }
+            }
+        }
+    }
+
+    @Test
+    public void shouldCreateTopicWithReplicationFactorWhenItDoesNotExist() {
+        for (int numBrokers = 1; numBrokers < 10; ++numBrokers) {
+            int maxRf = Math.min(numBrokers, 5);
+            int maxDefaultRf = Math.min(numBrokers, 5);
+            for (short rf = 1; rf < maxRf; ++rf) {
+                NewTopic newTopic = TopicAdmin.defineTopic("myTopic").replicationFactor(rf).compacted().build();
+
+                // Try clusters with no default replication factor or default partitions
+                assertTopicCreation(numBrokers, newTopic, null, null, rf, 1);
+
+                // Try clusters with different default partitions
+                for (int numPartitions = 1; numPartitions < 30; ++numPartitions) {
+                    assertTopicCreation(numBrokers, newTopic, numPartitions, null, rf, numPartitions);
+                }
+
+                // Try clusters with different default replication factors
+                for (int defaultRF = 1; defaultRF < maxDefaultRf; ++defaultRF) {
+                    assertTopicCreation(numBrokers, newTopic, null, defaultRF, rf, 1);
+                }
+            }
+        }
+    }
+
+    @Test
+    public void shouldCreateTopicWithDefaultPartitionsAndReplicationFactorWhenItDoesNotExist() {
+        NewTopic newTopic = TopicAdmin.defineTopic("my-topic")
+                                      .defaultPartitions()
+                                      .defaultReplicationFactor()
+                                      .compacted()
+                                      .build();
+
+        for (int numBrokers = 1; numBrokers < 10; ++numBrokers) {
+            int expectedReplicas = Math.min(3, numBrokers);
+            assertTopicCreation(numBrokers, newTopic, null, null, expectedReplicas, 1);
+            assertTopicCreation(numBrokers, newTopic, 30, null, expectedReplicas, 30);
         }
     }
 
@@ -93,12 +232,7 @@ public void shouldCreateOneTopicWhenProvidedMultipleDefinitionsWithSameTopicName
         NewTopic newTopic1 = TopicAdmin.defineTopic("myTopic").partitions(1).compacted().build();
         NewTopic newTopic2 = TopicAdmin.defineTopic("myTopic").partitions(1).compacted().build();
         Cluster cluster = createCluster(1);
-        try (MockKafkaAdminClientEnv env = new MockKafkaAdminClientEnv(cluster)) {
-            env.kafkaClient().setNode(cluster.controller());
-            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
-            env.kafkaClient().prepareMetadataUpdate(env.cluster(), Collections.emptySet());
-            env.kafkaClient().prepareResponse(createTopicResponse(newTopic1));
-            TopicAdmin admin = new TopicAdmin(null, env.adminClient());
+        try (TopicAdmin admin = new TopicAdmin(null, new MockAdminClient(cluster.nodes(), cluster.nodeById(0)))) {
             Set newTopicNames = admin.createTopics(newTopic1, newTopic2);
             assertEquals(1, newTopicNames.size());
             assertEquals(newTopic2.name(), newTopicNames.iterator().next());
@@ -106,21 +240,223 @@ public void shouldCreateOneTopicWhenProvidedMultipleDefinitionsWithSameTopicName
     }
 
     @Test
-    public void shouldReturnFalseWhenSuppliedNullTopicDescription() {
+    public void createShouldReturnFalseWhenSuppliedNullTopicDescription() {
         Cluster cluster = createCluster(1);
-        try (MockKafkaAdminClientEnv env = new MockKafkaAdminClientEnv(cluster)) {
-            env.kafkaClient().setNode(cluster.controller());
-            env.kafkaClient().setNodeApiVersions(NodeApiVersions.create());
-            env.kafkaClient().prepareMetadataUpdate(env.cluster(), Collections.emptySet());
-            TopicAdmin admin = new TopicAdmin(null, env.adminClient());
+        try (TopicAdmin admin = new TopicAdmin(null, new MockAdminClient(cluster.nodes(), cluster.nodeById(0)))) {
             boolean created = admin.createTopic(null);
             assertFalse(created);
         }
     }
 
+    @Test
+    public void describeShouldReturnEmptyWhenTopicDoesNotExist() {
+        NewTopic newTopic = TopicAdmin.defineTopic("myTopic").partitions(1).compacted().build();
+        Cluster cluster = createCluster(1);
+        try (TopicAdmin admin = new TopicAdmin(null, new MockAdminClient(cluster.nodes(), cluster.nodeById(0)))) {
+            assertTrue(admin.describeTopics(newTopic.name()).isEmpty());
+        }
+    }
+
+    @Test
+    public void describeShouldReturnTopicDescriptionWhenTopicExists() {
+        String topicName = "myTopic";
+        NewTopic newTopic = TopicAdmin.defineTopic(topicName).partitions(1).compacted().build();
+        Cluster cluster = createCluster(1);
+        try (MockAdminClient mockAdminClient = new MockAdminClient(cluster.nodes(), cluster.nodeById(0))) {
+            TopicPartitionInfo topicPartitionInfo = new TopicPartitionInfo(0, cluster.nodeById(0), cluster.nodes(), Collections.emptyList());
+            mockAdminClient.addTopic(false, topicName, Collections.singletonList(topicPartitionInfo), null);
+            TopicAdmin admin = new TopicAdmin(null, mockAdminClient);
+            Map desc = admin.describeTopics(newTopic.name());
+            assertFalse(desc.isEmpty());
+            TopicDescription topicDesc = new TopicDescription(topicName, false, Collections.singletonList(topicPartitionInfo));
+            assertEquals(desc.get("myTopic"), topicDesc);
+        }
+    }
+
+    @Test
+    public void describeTopicConfigShouldReturnEmptyMapWhenNoTopicsAreSpecified() {
+        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(describeConfigsResponseWithUnsupportedVersion(newTopic));
+            TopicAdmin admin = new TopicAdmin(null, env.adminClient());
+            Map results = admin.describeTopicConfigs();
+            assertTrue(results.isEmpty());
+        }
+    }
+
+    @Test
+    public void describeTopicConfigShouldReturnEmptyMapWhenUnsupportedVersionFailure() {
+        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(describeConfigsResponseWithUnsupportedVersion(newTopic));
+            TopicAdmin admin = new TopicAdmin(null, env.adminClient());
+            Map results = admin.describeTopicConfigs(newTopic.name());
+            assertTrue(results.isEmpty());
+        }
+    }
+
+    @Test
+    public void describeTopicConfigShouldReturnEmptyMapWhenClusterAuthorizationFailure() {
+        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(describeConfigsResponseWithClusterAuthorizationException(newTopic));
+            TopicAdmin admin = new TopicAdmin(null, env.adminClient());
+            Map results = admin.describeTopicConfigs(newTopic.name());
+            assertTrue(results.isEmpty());
+        }
+    }
+
+    @Test
+    public void describeTopicConfigShouldReturnEmptyMapWhenTopicAuthorizationFailure() {
+        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(describeConfigsResponseWithTopicAuthorizationException(newTopic));
+            TopicAdmin admin = new TopicAdmin(null, env.adminClient());
+            Map results = admin.describeTopicConfigs(newTopic.name());
+            assertTrue(results.isEmpty());
+        }
+    }
+
+    @Test
+    public void describeTopicConfigShouldReturnMapWithNullValueWhenTopicDoesNotExist() {
+        NewTopic newTopic = TopicAdmin.defineTopic("myTopic").partitions(1).compacted().build();
+        Cluster cluster = createCluster(1);
+        try (TopicAdmin admin = new TopicAdmin(null, new MockAdminClient(cluster.nodes(), cluster.nodeById(0)))) {
+            Map results = admin.describeTopicConfigs(newTopic.name());
+            assertFalse(results.isEmpty());
+            assertEquals(1, results.size());
+            assertNull(results.get("myTopic"));
+        }
+    }
+
+    @Test
+    public void describeTopicConfigShouldReturnTopicConfigWhenTopicExists() {
+        String topicName = "myTopic";
+        NewTopic newTopic = TopicAdmin.defineTopic(topicName)
+                                      .config(Collections.singletonMap("foo", "bar"))
+                                      .partitions(1)
+                                      .compacted()
+                                      .build();
+        Cluster cluster = createCluster(1);
+        try (MockAdminClient mockAdminClient = new MockAdminClient(cluster.nodes(), cluster.nodeById(0))) {
+            TopicPartitionInfo topicPartitionInfo = new TopicPartitionInfo(0, cluster.nodeById(0), cluster.nodes(), Collections.emptyList());
+            mockAdminClient.addTopic(false, topicName, Collections.singletonList(topicPartitionInfo), null);
+            TopicAdmin admin = new TopicAdmin(null, mockAdminClient);
+            Map result = admin.describeTopicConfigs(newTopic.name());
+            assertFalse(result.isEmpty());
+            assertEquals(1, result.size());
+            Config config = result.get("myTopic");
+            assertNotNull(config);
+            config.entries().forEach(entry -> {
+                assertEquals(newTopic.configs().get(entry.name()), entry.value());
+            });
+        }
+    }
+
+    @Test
+    public void verifyingTopicCleanupPolicyShouldReturnFalseWhenBrokerVersionIsUnsupported() {
+        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(describeConfigsResponseWithUnsupportedVersion(newTopic));
+            TopicAdmin admin = new TopicAdmin(null, env.adminClient());
+            boolean result = admin.verifyTopicCleanupPolicyOnlyCompact("myTopic", "worker.topic", "purpose");
+            assertFalse(result);
+        }
+    }
+
+    @Test
+    public void verifyingTopicCleanupPolicyShouldReturnFalseWhenClusterAuthorizationError() {
+        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(describeConfigsResponseWithClusterAuthorizationException(newTopic));
+            TopicAdmin admin = new TopicAdmin(null, env.adminClient());
+            boolean result = admin.verifyTopicCleanupPolicyOnlyCompact("myTopic", "worker.topic", "purpose");
+            assertFalse(result);
+        }
+    }
+
+    @Test
+    public void verifyingTopicCleanupPolicyShouldReturnFalseWhenTopicAuthorizationError() {
+        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(describeConfigsResponseWithTopicAuthorizationException(newTopic));
+            TopicAdmin admin = new TopicAdmin(null, env.adminClient());
+            boolean result = admin.verifyTopicCleanupPolicyOnlyCompact("myTopic", "worker.topic", "purpose");
+            assertFalse(result);
+        }
+    }
+
+    @Test
+    public void verifyingTopicCleanupPolicyShouldReturnTrueWhenTopicHasCorrectPolicy() {
+        String topicName = "myTopic";
+        Map topicConfigs = Collections.singletonMap("cleanup.policy", "compact");
+        Cluster cluster = createCluster(1);
+        try (MockAdminClient mockAdminClient = new MockAdminClient(cluster.nodes(), cluster.nodeById(0))) {
+            TopicPartitionInfo topicPartitionInfo = new TopicPartitionInfo(0, cluster.nodeById(0), cluster.nodes(), Collections.emptyList());
+            mockAdminClient.addTopic(false, topicName, Collections.singletonList(topicPartitionInfo), topicConfigs);
+            TopicAdmin admin = new TopicAdmin(null, mockAdminClient);
+            boolean result = admin.verifyTopicCleanupPolicyOnlyCompact("myTopic", "worker.topic", "purpose");
+            assertTrue(result);
+        }
+    }
+
+    @Test
+    public void verifyingTopicCleanupPolicyShouldFailWhenTopicHasDeletePolicy() {
+        String topicName = "myTopic";
+        Map topicConfigs = Collections.singletonMap("cleanup.policy", "delete");
+        Cluster cluster = createCluster(1);
+        try (MockAdminClient mockAdminClient = new MockAdminClient(cluster.nodes(), cluster.nodeById(0))) {
+            TopicPartitionInfo topicPartitionInfo = new TopicPartitionInfo(0, cluster.nodeById(0), cluster.nodes(), Collections.emptyList());
+            mockAdminClient.addTopic(false, topicName, Collections.singletonList(topicPartitionInfo), topicConfigs);
+            TopicAdmin admin = new TopicAdmin(null, mockAdminClient);
+            ConfigException e = assertThrows(ConfigException.class, () -> {
+                admin.verifyTopicCleanupPolicyOnlyCompact("myTopic", "worker.topic", "purpose");
+            });
+            assertTrue(e.getMessage().contains("to guarantee consistency and durability"));
+        }
+    }
+
+    @Test
+    public void verifyingTopicCleanupPolicyShouldFailWhenTopicHasDeleteAndCompactPolicy() {
+        String topicName = "myTopic";
+        Map topicConfigs = Collections.singletonMap("cleanup.policy", "delete,compact");
+        Cluster cluster = createCluster(1);
+        try (MockAdminClient mockAdminClient = new MockAdminClient(cluster.nodes(), cluster.nodeById(0))) {
+            TopicPartitionInfo topicPartitionInfo = new TopicPartitionInfo(0, cluster.nodeById(0), cluster.nodes(), Collections.emptyList());
+            mockAdminClient.addTopic(false, topicName, Collections.singletonList(topicPartitionInfo), topicConfigs);
+            TopicAdmin admin = new TopicAdmin(null, mockAdminClient);
+            ConfigException e = assertThrows(ConfigException.class, () -> {
+                admin.verifyTopicCleanupPolicyOnlyCompact("myTopic", "worker.topic", "purpose");
+            });
+            assertTrue(e.getMessage().contains("to guarantee consistency and durability"));
+        }
+    }
+
+    @Test
+    public void verifyingGettingTopicCleanupPolicies() {
+        String topicName = "myTopic";
+        Map topicConfigs = Collections.singletonMap("cleanup.policy", "compact");
+        Cluster cluster = createCluster(1);
+        try (MockAdminClient mockAdminClient = new MockAdminClient(cluster.nodes(), cluster.nodeById(0))) {
+            TopicPartitionInfo topicPartitionInfo = new TopicPartitionInfo(0, cluster.nodeById(0), cluster.nodes(), Collections.emptyList());
+            mockAdminClient.addTopic(false, topicName, Collections.singletonList(topicPartitionInfo), topicConfigs);
+            TopicAdmin admin = new TopicAdmin(null, mockAdminClient);
+            Set policies = admin.topicCleanupPolicy("myTopic");
+            assertEquals(1, policies.size());
+            assertEquals(TopicConfig.CLEANUP_POLICY_COMPACT, policies.iterator().next());
+        }
+    }
+
     private Cluster createCluster(int numNodes) {
         HashMap nodes = new HashMap<>();
-        for (int i = 0; i != numNodes; ++i) {
+        for (int i = 0; i < numNodes; ++i) {
             nodes.put(i, new Node(i, "localhost", 8121 + i));
         }
         Cluster cluster = new Cluster("mockClusterId", nodes.values(),
@@ -129,24 +465,125 @@ private Cluster createCluster(int numNodes) {
         return cluster;
     }
 
-    private CreateTopicsResponse createTopicResponse(NewTopic... topics) {
-        return createTopicResponse(new ApiError(Errors.NONE, ""), topics);
+    private CreateTopicsResponse createTopicResponseWithUnsupportedVersion(NewTopic... topics) {
+        return createTopicResponse(new ApiError(Errors.UNSUPPORTED_VERSION, "This version of the API is not supported"), topics);
     }
 
-    private CreateTopicsResponse createTopicResponseWithAlreadyExists(NewTopic... topics) {
-        return createTopicResponse(new ApiError(Errors.TOPIC_ALREADY_EXISTS, "Topic already exists"), topics);
+    private CreateTopicsResponse createTopicResponseWithClusterAuthorizationException(NewTopic... topics) {
+        return createTopicResponse(new ApiError(Errors.CLUSTER_AUTHORIZATION_FAILED, "Not authorized to create topic(s)"), topics);
     }
 
-    private CreateTopicsResponse createTopicResponseWithUnsupportedVersion(NewTopic... topics) {
-        return createTopicResponse(new ApiError(Errors.UNSUPPORTED_VERSION, "This version of the API is not supported"), 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, "");
-        Map topicResults = new HashMap<>();
+        CreateTopicsResponseData response = new CreateTopicsResponseData();
+        for (NewTopic topic : topics) {
+            response.topics().add(new CreatableTopicResult().
+                setName(topic.name()).
+                setErrorCode(error.error().code()).
+                setErrorMessage(error.message()));
+        }
+        return new CreateTopicsResponse(response);
+    }
+
+    protected void assertTopicCreation(
+            int brokers,
+            NewTopic newTopic,
+            Integer defaultPartitions,
+            Integer defaultReplicationFactor,
+            int expectedReplicas,
+            int expectedPartitions
+    ) {
+        Cluster cluster = createCluster(brokers);
+        MockAdminClient.Builder clientBuilder = MockAdminClient.create();
+        if (defaultPartitions != null) {
+            clientBuilder.defaultPartitions(defaultPartitions.shortValue());
+        }
+        if (defaultReplicationFactor != null) {
+            clientBuilder.defaultReplicationFactor(defaultReplicationFactor.intValue());
+        }
+        clientBuilder.brokers(cluster.nodes());
+        clientBuilder.controller(0);
+        try (MockAdminClient admin = clientBuilder.build()) {
+            TopicAdmin topicClient = new TopicAdmin(null, admin, false);
+            assertTrue(topicClient.createTopic(newTopic));
+            assertTopic(admin, newTopic.name(), expectedPartitions, expectedReplicas);
+        }
+    }
+
+    protected void assertTopic(MockAdminClient admin, String topicName, int expectedPartitions, int expectedReplicas) {
+        TopicDescription desc = null;
+        try {
+            desc = topicDescription(admin, topicName);
+        } catch (Throwable t) {
+            fail("Failed to find topic description for topic '" + topicName + "'");
+        }
+        assertEquals(expectedPartitions, desc.partitions().size());
+        for (TopicPartitionInfo tp : desc.partitions()) {
+            assertEquals(expectedReplicas, tp.replicas().size());
+        }
+    }
+
+    protected TopicDescription topicDescription(MockAdminClient admin, String topicName)
+            throws ExecutionException, InterruptedException {
+        DescribeTopicsResult result = admin.describeTopics(Collections.singleton(topicName));
+        Map> byName = result.values();
+        return byName.get(topicName).get();
+    }
+
+    private MetadataResponse describeTopicResponseWithUnsupportedVersion(NewTopic... topics) {
+        return describeTopicResponse(new ApiError(Errors.UNSUPPORTED_VERSION, "This version of the API is not supported"), topics);
+    }
+
+    private MetadataResponse describeTopicResponseWithClusterAuthorizationException(NewTopic... topics) {
+        return describeTopicResponse(new ApiError(Errors.CLUSTER_AUTHORIZATION_FAILED, "Not authorized to create topic(s)"), topics);
+    }
+
+    private MetadataResponse describeTopicResponseWithTopicAuthorizationException(NewTopic... topics) {
+        return describeTopicResponse(new ApiError(Errors.TOPIC_AUTHORIZATION_FAILED, "Not authorized to create topic(s)"), topics);
+    }
+
+    private MetadataResponse describeTopicResponse(ApiError error, NewTopic... topics) {
+        if (error == null) error = new ApiError(Errors.NONE, "");
+        MetadataResponseData response = new MetadataResponseData();
         for (NewTopic topic : topics) {
-            topicResults.put(topic.name(), error);
+            response.topics().add(new MetadataResponseTopic()
+                    .setName(topic.name())
+                    .setErrorCode(error.error().code()));
         }
-        return new CreateTopicsResponse(topicResults);
+        return new MetadataResponse(response, ApiKeys.METADATA.latestVersion());
+    }
+
+    private DescribeConfigsResponse describeConfigsResponseWithUnsupportedVersion(NewTopic... topics) {
+        return describeConfigsResponse(new ApiError(Errors.UNSUPPORTED_VERSION, "This version of the API is not supported"), topics);
+    }
+
+    private DescribeConfigsResponse describeConfigsResponseWithClusterAuthorizationException(NewTopic... topics) {
+        return describeConfigsResponse(new ApiError(Errors.CLUSTER_AUTHORIZATION_FAILED, "Not authorized to create topic(s)"), topics);
     }
+
+    private DescribeConfigsResponse describeConfigsResponseWithTopicAuthorizationException(NewTopic... topics) {
+        return describeConfigsResponse(new ApiError(Errors.TOPIC_AUTHORIZATION_FAILED, "Not authorized to create topic(s)"), topics);
+    }
+
+    private DescribeConfigsResponse describeConfigsResponse(ApiError error, NewTopic... topics) {
+        List results = Stream.of(topics)
+                .map(topic -> new DescribeConfigsResponseData.DescribeConfigsResult()
+                        .setErrorCode(error.error().code())
+                        .setErrorMessage(error.message())
+                        .setResourceType(ConfigResource.Type.TOPIC.id())
+                        .setResourceName(topic.name())
+                        .setConfigs(topic.configs().entrySet()
+                                .stream()
+                                .map(e -> new DescribeConfigsResponseData.DescribeConfigsResourceResult()
+                                        .setName(e.getKey())
+                                        .setValue(e.getValue()))
+                                .collect(Collectors.toList())))
+                .collect(Collectors.toList());
+        return new DescribeConfigsResponse(new DescribeConfigsResponseData().setThrottleTimeMs(1000).setResults(results));
+    }
+
 }
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/util/TopicCreationTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/util/TopicCreationTest.java
new file mode 100644
index 0000000000000..feb0e5f5ac9a5
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/util/TopicCreationTest.java
@@ -0,0 +1,638 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.clients.admin.NewTopic;
+import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.runtime.SourceConnectorConfig;
+import org.apache.kafka.connect.runtime.WorkerConfig;
+import org.apache.kafka.connect.runtime.distributed.DistributedConfig;
+import org.apache.kafka.connect.source.SourceRecord;
+import org.apache.kafka.connect.storage.StringConverter;
+import org.apache.kafka.connect.transforms.Cast;
+import org.apache.kafka.connect.transforms.RegexRouter;
+import org.apache.kafka.connect.transforms.Transformation;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.TimeUnit;
+
+import static org.apache.kafka.common.config.TopicConfig.CLEANUP_POLICY_COMPACT;
+import static org.apache.kafka.common.config.TopicConfig.CLEANUP_POLICY_CONFIG;
+import static org.apache.kafka.common.config.TopicConfig.COMPRESSION_TYPE_CONFIG;
+import static org.apache.kafka.common.config.TopicConfig.RETENTION_MS_CONFIG;
+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.ConnectorConfigTest.MOCK_PLUGINS;
+import static org.apache.kafka.connect.runtime.SourceConnectorConfig.TOPIC_CREATION_GROUPS_CONFIG;
+import static org.apache.kafka.connect.runtime.SourceConnectorConfig.TOPIC_CREATION_PREFIX;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.DEFAULT_TOPIC_CREATION_GROUP;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.DEFAULT_TOPIC_CREATION_PREFIX;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.EXCLUDE_REGEX_CONFIG;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.INCLUDE_REGEX_CONFIG;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.PARTITIONS_CONFIG;
+import static org.apache.kafka.connect.runtime.TopicCreationConfig.REPLICATION_FACTOR_CONFIG;
+import static org.apache.kafka.connect.runtime.WorkerConfig.BOOTSTRAP_SERVERS_CONFIG;
+import static org.apache.kafka.connect.runtime.WorkerConfig.KEY_CONVERTER_CLASS_CONFIG;
+import static org.apache.kafka.connect.runtime.WorkerConfig.TOPIC_CREATION_ENABLE_CONFIG;
+import static org.apache.kafka.connect.runtime.WorkerConfig.VALUE_CONVERTER_CLASS_CONFIG;
+import static org.apache.kafka.connect.runtime.distributed.DistributedConfig.CONFIG_TOPIC_CONFIG;
+import static org.apache.kafka.connect.runtime.distributed.DistributedConfig.GROUP_ID_CONFIG;
+import static org.apache.kafka.connect.runtime.distributed.DistributedConfig.OFFSET_STORAGE_TOPIC_CONFIG;
+import static org.apache.kafka.connect.runtime.distributed.DistributedConfig.STATUS_STORAGE_TOPIC_CONFIG;
+import static org.hamcrest.CoreMatchers.hasItem;
+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.assertFalse;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+
+public class TopicCreationTest {
+
+    private static final String FOO_CONNECTOR = "foo-source";
+    private static final String FOO_GROUP = "foo";
+    private static final String FOO_TOPIC = "foo-topic";
+    private static final String FOO_REGEX = ".*foo.*";
+
+    private static final String BAR_GROUP = "bar";
+    private static final String BAR_TOPIC = "bar-topic";
+    private static final String BAR_REGEX = ".*bar.*";
+
+    private static final short DEFAULT_REPLICATION_FACTOR = -1;
+    private static final int DEFAULT_PARTITIONS = -1;
+
+    Map workerProps;
+    WorkerConfig workerConfig;
+    Map sourceProps;
+    SourceConnectorConfig sourceConfig;
+
+    @Before
+    public void setup() {
+        workerProps = defaultWorkerProps();
+        workerConfig = new DistributedConfig(workerProps);
+    }
+
+    public Map defaultWorkerProps() {
+        Map props = new HashMap<>();
+        props.put(GROUP_ID_CONFIG, "connect-cluster");
+        props.put(BOOTSTRAP_SERVERS_CONFIG, "localhost:9092");
+        props.put(CONFIG_TOPIC_CONFIG, "connect-configs");
+        props.put(OFFSET_STORAGE_TOPIC_CONFIG, "connect-offsets");
+        props.put(STATUS_STORAGE_TOPIC_CONFIG, "connect-status");
+        props.put(KEY_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
+        props.put(VALUE_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
+        props.put(TOPIC_CREATION_ENABLE_CONFIG, String.valueOf(true));
+        return props;
+    }
+
+    public Map defaultConnectorProps() {
+        Map props = new HashMap<>();
+        props.put(NAME_CONFIG, FOO_CONNECTOR);
+        props.put(CONNECTOR_CLASS_CONFIG, "TestConnector");
+        return props;
+    }
+
+    public Map defaultConnectorPropsWithTopicCreation() {
+        Map props = defaultConnectorProps();
+        props.put(DEFAULT_TOPIC_CREATION_PREFIX + REPLICATION_FACTOR_CONFIG, String.valueOf(DEFAULT_REPLICATION_FACTOR));
+        props.put(DEFAULT_TOPIC_CREATION_PREFIX + PARTITIONS_CONFIG, String.valueOf(DEFAULT_PARTITIONS));
+        return props;
+    }
+
+    @Test
+    public void testTopicCreationWhenTopicCreationIsEnabled() {
+        sourceProps = defaultConnectorPropsWithTopicCreation();
+        sourceProps.put(TOPIC_CREATION_GROUPS_CONFIG, String.join(",", FOO_GROUP, BAR_GROUP));
+        sourceConfig = new SourceConnectorConfig(MOCK_PLUGINS, sourceProps, true);
+
+        Map groups = TopicCreationGroup.configuredGroups(sourceConfig);
+        TopicCreation topicCreation = TopicCreation.newTopicCreation(workerConfig, groups);
+
+        assertTrue(topicCreation.isTopicCreationEnabled());
+        assertTrue(topicCreation.isTopicCreationRequired(FOO_TOPIC));
+        assertThat(topicCreation.defaultTopicGroup(), is(groups.get(DEFAULT_TOPIC_CREATION_GROUP)));
+        assertEquals(2, topicCreation.topicGroups().size());
+        assertThat(topicCreation.topicGroups().keySet(), hasItems(FOO_GROUP, BAR_GROUP));
+        assertEquals(topicCreation.defaultTopicGroup(), topicCreation.findFirstGroup(FOO_TOPIC));
+        topicCreation.addTopic(FOO_TOPIC);
+        assertFalse(topicCreation.isTopicCreationRequired(FOO_TOPIC));
+    }
+
+    @Test
+    public void testTopicCreationWhenTopicCreationIsDisabled() {
+        workerProps.put(TOPIC_CREATION_ENABLE_CONFIG, String.valueOf(false));
+        workerConfig = new DistributedConfig(workerProps);
+        sourceProps = defaultConnectorPropsWithTopicCreation();
+        sourceConfig = new SourceConnectorConfig(MOCK_PLUGINS, sourceProps, true);
+
+        TopicCreation topicCreation = TopicCreation.newTopicCreation(workerConfig,
+                TopicCreationGroup.configuredGroups(sourceConfig));
+
+        assertFalse(topicCreation.isTopicCreationEnabled());
+        assertFalse(topicCreation.isTopicCreationRequired(FOO_TOPIC));
+        assertNull(topicCreation.defaultTopicGroup());
+        assertThat(topicCreation.topicGroups(), is(Collections.emptyMap()));
+        assertNull(topicCreation.findFirstGroup(FOO_TOPIC));
+        topicCreation.addTopic(FOO_TOPIC);
+        assertFalse(topicCreation.isTopicCreationRequired(FOO_TOPIC));
+    }
+
+    @Test
+    public void testEmptyTopicCreation() {
+        TopicCreation topicCreation = TopicCreation.newTopicCreation(workerConfig, null);
+
+        assertEquals(TopicCreation.empty(), topicCreation);
+        assertFalse(topicCreation.isTopicCreationEnabled());
+        assertFalse(topicCreation.isTopicCreationRequired(FOO_TOPIC));
+        assertNull(topicCreation.defaultTopicGroup());
+        assertEquals(0, topicCreation.topicGroups().size());
+        assertThat(topicCreation.topicGroups(), is(Collections.emptyMap()));
+        assertNull(topicCreation.findFirstGroup(FOO_TOPIC));
+        topicCreation.addTopic(FOO_TOPIC);
+        assertFalse(topicCreation.isTopicCreationRequired(FOO_TOPIC));
+    }
+
+    @Test
+    public void withDefaultTopicCreation() {
+        sourceProps = defaultConnectorPropsWithTopicCreation();
+        // Setting here but they should be ignored for the default group
+        sourceProps.put(TOPIC_CREATION_PREFIX + DEFAULT_TOPIC_CREATION_GROUP + "." + INCLUDE_REGEX_CONFIG, FOO_REGEX);
+        sourceProps.put(TOPIC_CREATION_PREFIX + DEFAULT_TOPIC_CREATION_GROUP + "." + EXCLUDE_REGEX_CONFIG, BAR_REGEX);
+
+        // verify config creation
+        sourceConfig = new SourceConnectorConfig(MOCK_PLUGINS, sourceProps, true);
+        assertTrue(sourceConfig.usesTopicCreation());
+        assertEquals(DEFAULT_REPLICATION_FACTOR, (short) sourceConfig.topicCreationReplicationFactor(DEFAULT_TOPIC_CREATION_GROUP));
+        assertEquals(DEFAULT_PARTITIONS, (int) sourceConfig.topicCreationPartitions(DEFAULT_TOPIC_CREATION_GROUP));
+        assertThat(sourceConfig.topicCreationInclude(DEFAULT_TOPIC_CREATION_GROUP), is(Collections.singletonList(".*")));
+        assertThat(sourceConfig.topicCreationExclude(DEFAULT_TOPIC_CREATION_GROUP), is(Collections.emptyList()));
+        assertThat(sourceConfig.topicCreationOtherConfigs(DEFAULT_TOPIC_CREATION_GROUP), is(Collections.emptyMap()));
+
+        // verify topic creation group is instantiated correctly
+        Map groups = TopicCreationGroup.configuredGroups(sourceConfig);
+        assertEquals(1, groups.size());
+        assertThat(groups.keySet(), hasItem(DEFAULT_TOPIC_CREATION_GROUP));
+
+        // verify topic creation
+        TopicCreation topicCreation = TopicCreation.newTopicCreation(workerConfig, groups);
+        TopicCreationGroup group = topicCreation.defaultTopicGroup();
+        // Default group will match all topics besides empty string
+        assertTrue(group.matches(" "));
+        assertTrue(group.matches(FOO_TOPIC));
+        assertEquals(DEFAULT_TOPIC_CREATION_GROUP, group.name());
+        assertTrue(topicCreation.isTopicCreationEnabled());
+        assertTrue(topicCreation.isTopicCreationRequired(FOO_TOPIC));
+        assertThat(topicCreation.topicGroups(), is(Collections.emptyMap()));
+        assertEquals(topicCreation.defaultTopicGroup(), topicCreation.findFirstGroup(FOO_TOPIC));
+        topicCreation.addTopic(FOO_TOPIC);
+        assertFalse(topicCreation.isTopicCreationRequired(FOO_TOPIC));
+
+        // verify new topic properties
+        NewTopic topicSpec = topicCreation.findFirstGroup(FOO_TOPIC).newTopic(FOO_TOPIC);
+        assertEquals(FOO_TOPIC, topicSpec.name());
+        assertEquals(DEFAULT_REPLICATION_FACTOR, topicSpec.replicationFactor());
+        assertEquals(DEFAULT_PARTITIONS, topicSpec.numPartitions());
+        assertThat(topicSpec.configs(), is(Collections.emptyMap()));
+    }
+
+    @Test
+    public void topicCreationWithDefaultGroupAndCustomProps() {
+        short replicas = 3;
+        int partitions = 5;
+        long retentionMs = TimeUnit.DAYS.toMillis(30);
+        String compressionType = "lz4";
+        Map topicProps = new HashMap<>();
+        topicProps.put(COMPRESSION_TYPE_CONFIG, compressionType);
+        topicProps.put(RETENTION_MS_CONFIG, String.valueOf(retentionMs));
+
+        sourceProps = defaultConnectorPropsWithTopicCreation();
+        sourceProps.put(DEFAULT_TOPIC_CREATION_PREFIX + REPLICATION_FACTOR_CONFIG, String.valueOf(replicas));
+        sourceProps.put(DEFAULT_TOPIC_CREATION_PREFIX + PARTITIONS_CONFIG, String.valueOf(partitions));
+        topicProps.forEach((k, v) -> sourceProps.put(DEFAULT_TOPIC_CREATION_PREFIX + k, v));
+        // Setting here but they should be ignored for the default group
+        sourceProps.put(TOPIC_CREATION_PREFIX + DEFAULT_TOPIC_CREATION_GROUP + "." + INCLUDE_REGEX_CONFIG, FOO_REGEX);
+        sourceProps.put(TOPIC_CREATION_PREFIX + DEFAULT_TOPIC_CREATION_GROUP + "." + EXCLUDE_REGEX_CONFIG, BAR_REGEX);
+
+        // verify config creation
+        sourceConfig = new SourceConnectorConfig(MOCK_PLUGINS, sourceProps, true);
+        assertTrue(sourceConfig.usesTopicCreation());
+        assertEquals(replicas, (short) sourceConfig.topicCreationReplicationFactor(DEFAULT_TOPIC_CREATION_GROUP));
+        assertEquals(partitions, (int) sourceConfig.topicCreationPartitions(DEFAULT_TOPIC_CREATION_GROUP));
+        assertThat(sourceConfig.topicCreationInclude(DEFAULT_TOPIC_CREATION_GROUP), is(Collections.singletonList(".*")));
+        assertThat(sourceConfig.topicCreationExclude(DEFAULT_TOPIC_CREATION_GROUP), is(Collections.emptyList()));
+        assertThat(sourceConfig.topicCreationOtherConfigs(DEFAULT_TOPIC_CREATION_GROUP), is(topicProps));
+
+        // verify topic creation group is instantiated correctly
+        Map groups = TopicCreationGroup.configuredGroups(sourceConfig);
+        assertEquals(1, groups.size());
+        assertThat(groups.keySet(), hasItem(DEFAULT_TOPIC_CREATION_GROUP));
+
+        // verify topic creation
+        TopicCreation topicCreation = TopicCreation.newTopicCreation(workerConfig, groups);
+        TopicCreationGroup group = topicCreation.defaultTopicGroup();
+        // Default group will match all topics besides empty string
+        assertTrue(group.matches(" "));
+        assertTrue(group.matches(FOO_TOPIC));
+        assertEquals(DEFAULT_TOPIC_CREATION_GROUP, group.name());
+        assertTrue(topicCreation.isTopicCreationEnabled());
+        assertTrue(topicCreation.isTopicCreationRequired(FOO_TOPIC));
+        assertThat(topicCreation.topicGroups(), is(Collections.emptyMap()));
+        assertEquals(topicCreation.defaultTopicGroup(), topicCreation.findFirstGroup(FOO_TOPIC));
+        topicCreation.addTopic(FOO_TOPIC);
+        assertFalse(topicCreation.isTopicCreationRequired(FOO_TOPIC));
+
+        // verify new topic properties
+        NewTopic topicSpec = topicCreation.findFirstGroup(FOO_TOPIC).newTopic(FOO_TOPIC);
+        assertEquals(FOO_TOPIC, topicSpec.name());
+        assertEquals(replicas, topicSpec.replicationFactor());
+        assertEquals(partitions, topicSpec.numPartitions());
+        assertThat(topicSpec.configs(), is(topicProps));
+    }
+
+    @Test
+    public void topicCreationWithOneGroup() {
+        short fooReplicas = 3;
+        int partitions = 5;
+        sourceProps = defaultConnectorPropsWithTopicCreation();
+        sourceProps.put(TOPIC_CREATION_GROUPS_CONFIG, String.join(",", FOO_GROUP));
+        sourceProps.put(DEFAULT_TOPIC_CREATION_PREFIX + PARTITIONS_CONFIG, String.valueOf(partitions));
+        sourceProps.put(TOPIC_CREATION_PREFIX + FOO_GROUP + "." + INCLUDE_REGEX_CONFIG, FOO_REGEX);
+        sourceProps.put(TOPIC_CREATION_PREFIX + FOO_GROUP + "." + EXCLUDE_REGEX_CONFIG, BAR_REGEX);
+        sourceProps.put(TOPIC_CREATION_PREFIX + FOO_GROUP + "." + REPLICATION_FACTOR_CONFIG, String.valueOf(fooReplicas));
+
+        Map topicProps = new HashMap<>();
+        topicProps.put(CLEANUP_POLICY_CONFIG, CLEANUP_POLICY_COMPACT);
+        topicProps.forEach((k, v) -> sourceProps.put(TOPIC_CREATION_PREFIX + FOO_GROUP + "." + k, v));
+
+        // verify config creation
+        sourceConfig = new SourceConnectorConfig(MOCK_PLUGINS, sourceProps, true);
+        assertTrue(sourceConfig.usesTopicCreation());
+        assertEquals(DEFAULT_REPLICATION_FACTOR, (short) sourceConfig.topicCreationReplicationFactor(DEFAULT_TOPIC_CREATION_GROUP));
+        assertEquals(partitions, (int) sourceConfig.topicCreationPartitions(DEFAULT_TOPIC_CREATION_GROUP));
+        assertThat(sourceConfig.topicCreationInclude(DEFAULT_TOPIC_CREATION_GROUP), is(Collections.singletonList(".*")));
+        assertThat(sourceConfig.topicCreationExclude(DEFAULT_TOPIC_CREATION_GROUP), is(Collections.emptyList()));
+        assertThat(sourceConfig.topicCreationOtherConfigs(DEFAULT_TOPIC_CREATION_GROUP), is(Collections.emptyMap()));
+
+        // verify topic creation group is instantiated correctly
+        Map groups = TopicCreationGroup.configuredGroups(sourceConfig);
+        assertEquals(2, groups.size());
+        assertThat(groups.keySet(), hasItems(DEFAULT_TOPIC_CREATION_GROUP, FOO_GROUP));
+
+        // verify topic creation
+        TopicCreation topicCreation = TopicCreation.newTopicCreation(workerConfig, groups);
+        TopicCreationGroup defaultGroup = topicCreation.defaultTopicGroup();
+        // Default group will match all topics besides empty string
+        assertTrue(defaultGroup.matches(" "));
+        assertTrue(defaultGroup.matches(FOO_TOPIC));
+        assertTrue(defaultGroup.matches(BAR_TOPIC));
+        assertEquals(DEFAULT_TOPIC_CREATION_GROUP, defaultGroup.name());
+        TopicCreationGroup fooGroup = groups.get(FOO_GROUP);
+        assertFalse(fooGroup.matches(" "));
+        assertTrue(fooGroup.matches(FOO_TOPIC));
+        assertFalse(fooGroup.matches(BAR_TOPIC));
+        assertEquals(FOO_GROUP, fooGroup.name());
+
+        assertTrue(topicCreation.isTopicCreationEnabled());
+        assertTrue(topicCreation.isTopicCreationRequired(FOO_TOPIC));
+        assertEquals(1, topicCreation.topicGroups().size());
+        assertThat(topicCreation.topicGroups().keySet(), hasItems(FOO_GROUP));
+        assertEquals(fooGroup, topicCreation.findFirstGroup(FOO_TOPIC));
+        topicCreation.addTopic(FOO_TOPIC);
+        assertFalse(topicCreation.isTopicCreationRequired(FOO_TOPIC));
+
+        // verify new topic properties
+        NewTopic defaultTopicSpec = topicCreation.findFirstGroup(BAR_TOPIC).newTopic(BAR_TOPIC);
+        assertEquals(BAR_TOPIC, defaultTopicSpec.name());
+        assertEquals(DEFAULT_REPLICATION_FACTOR, defaultTopicSpec.replicationFactor());
+        assertEquals(partitions, defaultTopicSpec.numPartitions());
+        assertThat(defaultTopicSpec.configs(), is(Collections.emptyMap()));
+
+        NewTopic fooTopicSpec = topicCreation.findFirstGroup(FOO_TOPIC).newTopic(FOO_TOPIC);
+        assertEquals(FOO_TOPIC, fooTopicSpec.name());
+        assertEquals(fooReplicas, fooTopicSpec.replicationFactor());
+        assertEquals(partitions, fooTopicSpec.numPartitions());
+        assertThat(fooTopicSpec.configs(), is(topicProps));
+    }
+
+    @Test
+    public void topicCreationWithOneGroupAndCombinedRegex() {
+        short fooReplicas = 3;
+        int partitions = 5;
+        sourceProps = defaultConnectorPropsWithTopicCreation();
+        sourceProps.put(TOPIC_CREATION_GROUPS_CONFIG, String.join(",", FOO_GROUP));
+        sourceProps.put(DEFAULT_TOPIC_CREATION_PREFIX + PARTITIONS_CONFIG, String.valueOf(partitions));
+        // Setting here but they should be ignored for the default group
+        sourceProps.put(TOPIC_CREATION_PREFIX + FOO_GROUP + "." + INCLUDE_REGEX_CONFIG, String.join("|", FOO_REGEX, BAR_REGEX));
+        sourceProps.put(TOPIC_CREATION_PREFIX + FOO_GROUP + "." + REPLICATION_FACTOR_CONFIG, String.valueOf(fooReplicas));
+
+        Map topicProps = new HashMap<>();
+        topicProps.put(CLEANUP_POLICY_CONFIG, CLEANUP_POLICY_COMPACT);
+        topicProps.forEach((k, v) -> sourceProps.put(TOPIC_CREATION_PREFIX + FOO_GROUP + "." + k, v));
+
+        // verify config creation
+        sourceConfig = new SourceConnectorConfig(MOCK_PLUGINS, sourceProps, true);
+        assertTrue(sourceConfig.usesTopicCreation());
+        assertEquals(DEFAULT_REPLICATION_FACTOR, (short) sourceConfig.topicCreationReplicationFactor(DEFAULT_TOPIC_CREATION_GROUP));
+        assertEquals(partitions, (int) sourceConfig.topicCreationPartitions(DEFAULT_TOPIC_CREATION_GROUP));
+        assertThat(sourceConfig.topicCreationInclude(DEFAULT_TOPIC_CREATION_GROUP), is(Collections.singletonList(".*")));
+        assertThat(sourceConfig.topicCreationExclude(DEFAULT_TOPIC_CREATION_GROUP), is(Collections.emptyList()));
+        assertThat(sourceConfig.topicCreationOtherConfigs(DEFAULT_TOPIC_CREATION_GROUP), is(Collections.emptyMap()));
+
+        // verify topic creation group is instantiated correctly
+        Map groups = TopicCreationGroup.configuredGroups(sourceConfig);
+        assertEquals(2, groups.size());
+        assertThat(groups.keySet(), hasItems(DEFAULT_TOPIC_CREATION_GROUP, FOO_GROUP));
+
+        // verify topic creation
+        TopicCreation topicCreation = TopicCreation.newTopicCreation(workerConfig, groups);
+        TopicCreationGroup defaultGroup = topicCreation.defaultTopicGroup();
+        // Default group will match all topics besides empty string
+        assertTrue(defaultGroup.matches(" "));
+        assertTrue(defaultGroup.matches(FOO_TOPIC));
+        assertTrue(defaultGroup.matches(BAR_TOPIC));
+        assertEquals(DEFAULT_TOPIC_CREATION_GROUP, defaultGroup.name());
+        TopicCreationGroup fooGroup = groups.get(FOO_GROUP);
+        assertFalse(fooGroup.matches(" "));
+        assertTrue(fooGroup.matches(FOO_TOPIC));
+        assertTrue(fooGroup.matches(BAR_TOPIC));
+        assertEquals(FOO_GROUP, fooGroup.name());
+
+        assertTrue(topicCreation.isTopicCreationEnabled());
+        assertTrue(topicCreation.isTopicCreationRequired(FOO_TOPIC));
+        assertTrue(topicCreation.isTopicCreationRequired(BAR_TOPIC));
+        assertEquals(1, topicCreation.topicGroups().size());
+        assertThat(topicCreation.topicGroups().keySet(), hasItems(FOO_GROUP));
+        assertEquals(fooGroup, topicCreation.findFirstGroup(FOO_TOPIC));
+        assertEquals(fooGroup, topicCreation.findFirstGroup(BAR_TOPIC));
+        topicCreation.addTopic(FOO_TOPIC);
+        topicCreation.addTopic(BAR_TOPIC);
+        assertFalse(topicCreation.isTopicCreationRequired(FOO_TOPIC));
+        assertFalse(topicCreation.isTopicCreationRequired(BAR_TOPIC));
+
+        // verify new topic properties
+        NewTopic fooTopicSpec = topicCreation.findFirstGroup(FOO_TOPIC).newTopic(FOO_TOPIC);
+        assertEquals(FOO_TOPIC, fooTopicSpec.name());
+        assertEquals(fooReplicas, fooTopicSpec.replicationFactor());
+        assertEquals(partitions, fooTopicSpec.numPartitions());
+        assertThat(fooTopicSpec.configs(), is(topicProps));
+
+        NewTopic barTopicSpec = topicCreation.findFirstGroup(BAR_TOPIC).newTopic(BAR_TOPIC);
+        assertEquals(BAR_TOPIC, barTopicSpec.name());
+        assertEquals(fooReplicas, barTopicSpec.replicationFactor());
+        assertEquals(partitions, barTopicSpec.numPartitions());
+        assertThat(barTopicSpec.configs(), is(topicProps));
+    }
+
+    @Test
+    public void topicCreationWithTwoGroups() {
+        short fooReplicas = 3;
+        int partitions = 5;
+        int barPartitions = 1;
+
+        sourceProps = defaultConnectorPropsWithTopicCreation();
+        sourceProps.put(TOPIC_CREATION_GROUPS_CONFIG, String.join(",", FOO_GROUP, BAR_GROUP));
+        sourceProps.put(DEFAULT_TOPIC_CREATION_PREFIX + PARTITIONS_CONFIG, String.valueOf(partitions));
+        // Setting here but they should be ignored for the default group
+        sourceProps.put(TOPIC_CREATION_PREFIX + FOO_GROUP + "." + INCLUDE_REGEX_CONFIG, FOO_TOPIC);
+        sourceProps.put(TOPIC_CREATION_PREFIX + FOO_GROUP + "." + REPLICATION_FACTOR_CONFIG, String.valueOf(fooReplicas));
+        sourceProps.put(TOPIC_CREATION_PREFIX + BAR_GROUP + "." + INCLUDE_REGEX_CONFIG, BAR_REGEX);
+        sourceProps.put(TOPIC_CREATION_PREFIX + BAR_GROUP + "." + PARTITIONS_CONFIG, String.valueOf(barPartitions));
+
+        Map fooTopicProps = new HashMap<>();
+        fooTopicProps.put(RETENTION_MS_CONFIG, String.valueOf(TimeUnit.DAYS.toMillis(30)));
+        fooTopicProps.forEach((k, v) -> sourceProps.put(TOPIC_CREATION_PREFIX + FOO_GROUP + "." + k, v));
+
+        Map barTopicProps = new HashMap<>();
+        barTopicProps.put(CLEANUP_POLICY_CONFIG, CLEANUP_POLICY_COMPACT);
+        barTopicProps.forEach((k, v) -> sourceProps.put(TOPIC_CREATION_PREFIX + BAR_GROUP + "." + k, v));
+
+        // verify config creation
+        sourceConfig = new SourceConnectorConfig(MOCK_PLUGINS, sourceProps, true);
+        assertTrue(sourceConfig.usesTopicCreation());
+        assertEquals(DEFAULT_REPLICATION_FACTOR, (short) sourceConfig.topicCreationReplicationFactor(DEFAULT_TOPIC_CREATION_GROUP));
+        assertEquals(partitions, (int) sourceConfig.topicCreationPartitions(DEFAULT_TOPIC_CREATION_GROUP));
+        assertThat(sourceConfig.topicCreationInclude(DEFAULT_TOPIC_CREATION_GROUP), is(Collections.singletonList(".*")));
+        assertThat(sourceConfig.topicCreationExclude(DEFAULT_TOPIC_CREATION_GROUP), is(Collections.emptyList()));
+        assertThat(sourceConfig.topicCreationOtherConfigs(DEFAULT_TOPIC_CREATION_GROUP), is(Collections.emptyMap()));
+
+        // verify topic creation group is instantiated correctly
+        Map groups = TopicCreationGroup.configuredGroups(sourceConfig);
+        assertEquals(3, groups.size());
+        assertThat(groups.keySet(), hasItems(DEFAULT_TOPIC_CREATION_GROUP, FOO_GROUP, BAR_GROUP));
+
+        // verify topic creation
+        TopicCreation topicCreation = TopicCreation.newTopicCreation(workerConfig, groups);
+        TopicCreationGroup defaultGroup = topicCreation.defaultTopicGroup();
+        // Default group will match all topics besides empty string
+        assertTrue(defaultGroup.matches(" "));
+        assertTrue(defaultGroup.matches(FOO_TOPIC));
+        assertTrue(defaultGroup.matches(BAR_TOPIC));
+        assertEquals(DEFAULT_TOPIC_CREATION_GROUP, defaultGroup.name());
+        TopicCreationGroup fooGroup = groups.get(FOO_GROUP);
+        assertFalse(fooGroup.matches(" "));
+        assertTrue(fooGroup.matches(FOO_TOPIC));
+        assertFalse(fooGroup.matches(BAR_TOPIC));
+        assertEquals(FOO_GROUP, fooGroup.name());
+        TopicCreationGroup barGroup = groups.get(BAR_GROUP);
+        assertTrue(barGroup.matches(BAR_TOPIC));
+        assertFalse(barGroup.matches(FOO_TOPIC));
+        assertEquals(BAR_GROUP, barGroup.name());
+
+        assertTrue(topicCreation.isTopicCreationEnabled());
+        assertTrue(topicCreation.isTopicCreationRequired(FOO_TOPIC));
+        assertTrue(topicCreation.isTopicCreationRequired(BAR_TOPIC));
+        assertEquals(2, topicCreation.topicGroups().size());
+        assertThat(topicCreation.topicGroups().keySet(), hasItems(FOO_GROUP, BAR_GROUP));
+        assertEquals(fooGroup, topicCreation.findFirstGroup(FOO_TOPIC));
+        assertEquals(barGroup, topicCreation.findFirstGroup(BAR_TOPIC));
+        topicCreation.addTopic(FOO_TOPIC);
+        topicCreation.addTopic(BAR_TOPIC);
+        assertFalse(topicCreation.isTopicCreationRequired(FOO_TOPIC));
+        assertFalse(topicCreation.isTopicCreationRequired(BAR_TOPIC));
+
+        // verify new topic properties
+        String otherTopic = "any-other-topic";
+        NewTopic defaultTopicSpec = topicCreation.findFirstGroup(otherTopic).newTopic(otherTopic);
+        assertEquals(otherTopic, defaultTopicSpec.name());
+        assertEquals(DEFAULT_REPLICATION_FACTOR, defaultTopicSpec.replicationFactor());
+        assertEquals(partitions, defaultTopicSpec.numPartitions());
+        assertThat(defaultTopicSpec.configs(), is(Collections.emptyMap()));
+
+        NewTopic fooTopicSpec = topicCreation.findFirstGroup(FOO_TOPIC).newTopic(FOO_TOPIC);
+        assertEquals(FOO_TOPIC, fooTopicSpec.name());
+        assertEquals(fooReplicas, fooTopicSpec.replicationFactor());
+        assertEquals(partitions, fooTopicSpec.numPartitions());
+        assertThat(fooTopicSpec.configs(), is(fooTopicProps));
+
+        NewTopic barTopicSpec = topicCreation.findFirstGroup(BAR_TOPIC).newTopic(BAR_TOPIC);
+        assertEquals(BAR_TOPIC, barTopicSpec.name());
+        assertEquals(DEFAULT_REPLICATION_FACTOR, barTopicSpec.replicationFactor());
+        assertEquals(barPartitions, barTopicSpec.numPartitions());
+        assertThat(barTopicSpec.configs(), is(barTopicProps));
+    }
+
+    @Test
+    public void testTopicCreationWithSingleTransformation() {
+        sourceProps = defaultConnectorPropsWithTopicCreation();
+        sourceProps.put(TOPIC_CREATION_GROUPS_CONFIG, String.join(",", FOO_GROUP, BAR_GROUP));
+        String xformName = "example";
+        String castType = "int8";
+        sourceProps.put("transforms", xformName);
+        sourceProps.put("transforms." + xformName + ".type", Cast.Value.class.getName());
+        sourceProps.put("transforms." + xformName + ".spec", castType);
+
+        sourceConfig = new SourceConnectorConfig(MOCK_PLUGINS, sourceProps, true);
+
+        Map groups = TopicCreationGroup.configuredGroups(sourceConfig);
+        TopicCreation topicCreation = TopicCreation.newTopicCreation(workerConfig, groups);
+
+        assertTrue(topicCreation.isTopicCreationEnabled());
+        assertTrue(topicCreation.isTopicCreationRequired(FOO_TOPIC));
+        assertThat(topicCreation.defaultTopicGroup(), is(groups.get(DEFAULT_TOPIC_CREATION_GROUP)));
+        assertEquals(2, topicCreation.topicGroups().size());
+        assertThat(topicCreation.topicGroups().keySet(), hasItems(FOO_GROUP, BAR_GROUP));
+        assertEquals(topicCreation.defaultTopicGroup(), topicCreation.findFirstGroup(FOO_TOPIC));
+        topicCreation.addTopic(FOO_TOPIC);
+        assertFalse(topicCreation.isTopicCreationRequired(FOO_TOPIC));
+
+        List> transformations = sourceConfig.transformations();
+        assertEquals(1, transformations.size());
+        Cast xform = (Cast) transformations.get(0);
+        SourceRecord transformed = xform.apply(new SourceRecord(null, null, "topic", 0, null, null, Schema.INT8_SCHEMA, 42));
+        assertEquals(Schema.Type.INT8, transformed.valueSchema().type());
+        assertEquals((byte) 42, transformed.value());
+    }
+
+    @Test
+    public void topicCreationWithTwoGroupsAndTwoTransformations() {
+        short fooReplicas = 3;
+        int partitions = 5;
+        int barPartitions = 1;
+
+        sourceProps = defaultConnectorPropsWithTopicCreation();
+        sourceProps.put(TOPIC_CREATION_GROUPS_CONFIG, String.join(",", FOO_GROUP, BAR_GROUP));
+        sourceProps.put(DEFAULT_TOPIC_CREATION_PREFIX + PARTITIONS_CONFIG, String.valueOf(partitions));
+        // Setting here but they should be ignored for the default group
+        sourceProps.put(TOPIC_CREATION_PREFIX + FOO_GROUP + "." + INCLUDE_REGEX_CONFIG, FOO_TOPIC);
+        sourceProps.put(TOPIC_CREATION_PREFIX + FOO_GROUP + "." + REPLICATION_FACTOR_CONFIG, String.valueOf(fooReplicas));
+        sourceProps.put(TOPIC_CREATION_PREFIX + BAR_GROUP + "." + INCLUDE_REGEX_CONFIG, BAR_REGEX);
+        sourceProps.put(TOPIC_CREATION_PREFIX + BAR_GROUP + "." + PARTITIONS_CONFIG, String.valueOf(barPartitions));
+
+        String castName = "cast";
+        String castType = "int8";
+        sourceProps.put("transforms." + castName + ".type", Cast.Value.class.getName());
+        sourceProps.put("transforms." + castName + ".spec", castType);
+
+        String regexRouterName = "regex";
+        sourceProps.put("transforms." + regexRouterName + ".type", RegexRouter.class.getName());
+        sourceProps.put("transforms." + regexRouterName + ".regex", "(.*)");
+        sourceProps.put("transforms." + regexRouterName + ".replacement", "prefix-$1");
+
+        sourceProps.put("transforms", String.join(",", castName, regexRouterName));
+
+        Map fooTopicProps = new HashMap<>();
+        fooTopicProps.put(RETENTION_MS_CONFIG, String.valueOf(TimeUnit.DAYS.toMillis(30)));
+        fooTopicProps.forEach((k, v) -> sourceProps.put(TOPIC_CREATION_PREFIX + FOO_GROUP + "." + k, v));
+
+        Map barTopicProps = new HashMap<>();
+        barTopicProps.put(CLEANUP_POLICY_CONFIG, CLEANUP_POLICY_COMPACT);
+        barTopicProps.forEach((k, v) -> sourceProps.put(TOPIC_CREATION_PREFIX + BAR_GROUP + "." + k, v));
+
+        // verify config creation
+        sourceConfig = new SourceConnectorConfig(MOCK_PLUGINS, sourceProps, true);
+        assertTrue(sourceConfig.usesTopicCreation());
+        assertEquals(DEFAULT_REPLICATION_FACTOR, (short) sourceConfig.topicCreationReplicationFactor(DEFAULT_TOPIC_CREATION_GROUP));
+        assertEquals(partitions, (int) sourceConfig.topicCreationPartitions(DEFAULT_TOPIC_CREATION_GROUP));
+        assertThat(sourceConfig.topicCreationInclude(DEFAULT_TOPIC_CREATION_GROUP), is(Collections.singletonList(".*")));
+        assertThat(sourceConfig.topicCreationExclude(DEFAULT_TOPIC_CREATION_GROUP), is(Collections.emptyList()));
+        assertThat(sourceConfig.topicCreationOtherConfigs(DEFAULT_TOPIC_CREATION_GROUP), is(Collections.emptyMap()));
+
+        // verify topic creation group is instantiated correctly
+        Map groups = TopicCreationGroup.configuredGroups(sourceConfig);
+        assertEquals(3, groups.size());
+        assertThat(groups.keySet(), hasItems(DEFAULT_TOPIC_CREATION_GROUP, FOO_GROUP, BAR_GROUP));
+
+        // verify topic creation
+        TopicCreation topicCreation = TopicCreation.newTopicCreation(workerConfig, groups);
+        TopicCreationGroup defaultGroup = topicCreation.defaultTopicGroup();
+        // Default group will match all topics besides empty string
+        assertTrue(defaultGroup.matches(" "));
+        assertTrue(defaultGroup.matches(FOO_TOPIC));
+        assertTrue(defaultGroup.matches(BAR_TOPIC));
+        assertEquals(DEFAULT_TOPIC_CREATION_GROUP, defaultGroup.name());
+        TopicCreationGroup fooGroup = groups.get(FOO_GROUP);
+        assertFalse(fooGroup.matches(" "));
+        assertTrue(fooGroup.matches(FOO_TOPIC));
+        assertFalse(fooGroup.matches(BAR_TOPIC));
+        assertEquals(FOO_GROUP, fooGroup.name());
+        TopicCreationGroup barGroup = groups.get(BAR_GROUP);
+        assertTrue(barGroup.matches(BAR_TOPIC));
+        assertFalse(barGroup.matches(FOO_TOPIC));
+        assertEquals(BAR_GROUP, barGroup.name());
+
+        assertTrue(topicCreation.isTopicCreationEnabled());
+        assertTrue(topicCreation.isTopicCreationRequired(FOO_TOPIC));
+        assertTrue(topicCreation.isTopicCreationRequired(BAR_TOPIC));
+        assertEquals(2, topicCreation.topicGroups().size());
+        assertThat(topicCreation.topicGroups().keySet(), hasItems(FOO_GROUP, BAR_GROUP));
+        assertEquals(fooGroup, topicCreation.findFirstGroup(FOO_TOPIC));
+        assertEquals(barGroup, topicCreation.findFirstGroup(BAR_TOPIC));
+        topicCreation.addTopic(FOO_TOPIC);
+        topicCreation.addTopic(BAR_TOPIC);
+        assertFalse(topicCreation.isTopicCreationRequired(FOO_TOPIC));
+        assertFalse(topicCreation.isTopicCreationRequired(BAR_TOPIC));
+
+        // verify new topic properties
+        String otherTopic = "any-other-topic";
+        NewTopic defaultTopicSpec = topicCreation.findFirstGroup(otherTopic).newTopic(otherTopic);
+        assertEquals(otherTopic, defaultTopicSpec.name());
+        assertEquals(DEFAULT_REPLICATION_FACTOR, defaultTopicSpec.replicationFactor());
+        assertEquals(partitions, defaultTopicSpec.numPartitions());
+        assertThat(defaultTopicSpec.configs(), is(Collections.emptyMap()));
+
+        NewTopic fooTopicSpec = topicCreation.findFirstGroup(FOO_TOPIC).newTopic(FOO_TOPIC);
+        assertEquals(FOO_TOPIC, fooTopicSpec.name());
+        assertEquals(fooReplicas, fooTopicSpec.replicationFactor());
+        assertEquals(partitions, fooTopicSpec.numPartitions());
+        assertThat(fooTopicSpec.configs(), is(fooTopicProps));
+
+        NewTopic barTopicSpec = topicCreation.findFirstGroup(BAR_TOPIC).newTopic(BAR_TOPIC);
+        assertEquals(BAR_TOPIC, barTopicSpec.name());
+        assertEquals(DEFAULT_REPLICATION_FACTOR, barTopicSpec.replicationFactor());
+        assertEquals(barPartitions, barTopicSpec.numPartitions());
+        assertThat(barTopicSpec.configs(), is(barTopicProps));
+
+        List> transformations = sourceConfig.transformations();
+        assertEquals(2, transformations.size());
+
+        Cast castXForm = (Cast) transformations.get(0);
+        SourceRecord transformed = castXForm.apply(new SourceRecord(null, null, "topic", 0, null, null, Schema.INT8_SCHEMA, 42));
+        assertEquals(Schema.Type.INT8, transformed.valueSchema().type());
+        assertEquals((byte) 42, transformed.value());
+
+        RegexRouter regexRouterXForm = (RegexRouter) transformations.get(1);
+        transformed = regexRouterXForm.apply(new SourceRecord(null, null, "topic", 0, null, null, Schema.INT8_SCHEMA, 42));
+        assertEquals("prefix-topic", transformed.topic());
+    }
+}
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
new file mode 100644
index 0000000000000..4c270d25f16b3
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/util/clusters/EmbeddedConnectCluster.java
@@ -0,0 +1,781 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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 com.fasterxml.jackson.core.type.TypeReference;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import org.apache.kafka.common.utils.Exit;
+import org.apache.kafka.connect.errors.ConnectException;
+import org.apache.kafka.connect.runtime.rest.entities.ActiveTopicsInfo;
+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.ServerInfo;
+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.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.List;
+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;
+import static org.apache.kafka.connect.runtime.ConnectorConfig.VALUE_CONVERTER_CLASS_CONFIG;
+import static org.apache.kafka.connect.runtime.WorkerConfig.BOOTSTRAP_SERVERS_CONFIG;
+import static org.apache.kafka.connect.runtime.WorkerConfig.REST_HOST_NAME_CONFIG;
+import static org.apache.kafka.connect.runtime.WorkerConfig.REST_PORT_CONFIG;
+import static org.apache.kafka.connect.runtime.distributed.DistributedConfig.CONFIG_STORAGE_REPLICATION_FACTOR_CONFIG;
+import static org.apache.kafka.connect.runtime.distributed.DistributedConfig.CONFIG_TOPIC_CONFIG;
+import static org.apache.kafka.connect.runtime.distributed.DistributedConfig.OFFSET_STORAGE_REPLICATION_FACTOR_CONFIG;
+import static org.apache.kafka.connect.runtime.distributed.DistributedConfig.OFFSET_STORAGE_TOPIC_CONFIG;
+import static org.apache.kafka.connect.runtime.distributed.DistributedConfig.STATUS_STORAGE_REPLICATION_FACTOR_CONFIG;
+import static org.apache.kafka.connect.runtime.distributed.DistributedConfig.STATUS_STORAGE_TOPIC_CONFIG;
+
+/**
+ * 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. Methods on the same {@code EmbeddedConnectCluster} are
+ * not guaranteed to be thread-safe.
+ */
+public class EmbeddedConnectCluster {
+
+    private static final Logger log = LoggerFactory.getLogger(EmbeddedConnectCluster.class);
+
+    public static final int DEFAULT_NUM_BROKERS = 1;
+    public static final int DEFAULT_NUM_WORKERS = 1;
+    private static final Properties DEFAULT_BROKER_CONFIG = new Properties();
+    private static final String REST_HOST_NAME = "localhost";
+
+    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 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 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);
+        }
+    };
+
+    /**
+     * 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);
+        }
+    };
+
+    /**
+     * Start the connect cluster and the embedded Kafka and Zookeeper cluster.
+     */
+    public void start() {
+        if (maskExitProcedures) {
+            Exit.setExitProcedure(exitProcedure);
+            Exit.setHaltProcedure(haltProcedure);
+        }
+        kafkaCluster.before();
+        startConnect();
+    }
+
+    /**
+     * 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() {
+        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 {
+            if (maskExitProcedures) {
+                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()) {
+        }
+        if (toRemove != null) {
+            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);
+        }
+    }
+
+    /**
+     * Determine whether the Connect cluster has any workers running.
+     *
+     * @return true if any worker is running, or false otherwise
+     */
+    public boolean anyWorkersRunning() {
+        return workers().stream().anyMatch(WorkerHandle::isRunning);
+    }
+
+    /**
+     * Determine whether the Connect cluster has all workers running.
+     *
+     * @return true if all workers are running, or false otherwise
+     */
+    public boolean allWorkersRunning() {
+        return workers().stream().allMatch(WorkerHandle::isRunning);
+    }
+
+    @SuppressWarnings("deprecation")
+    public void startConnect() {
+        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);
+        workerProps.put(REST_PORT_CONFIG, "0"); // use a random available port
+
+        String internalTopicsReplFactor = String.valueOf(numBrokers);
+        putIfAbsent(workerProps, GROUP_ID_CONFIG, "connect-integration-test-" + connectClusterName);
+        putIfAbsent(workerProps, OFFSET_STORAGE_TOPIC_CONFIG, "connect-offset-topic-" + connectClusterName);
+        putIfAbsent(workerProps, OFFSET_STORAGE_REPLICATION_FACTOR_CONFIG, internalTopicsReplFactor);
+        putIfAbsent(workerProps, CONFIG_TOPIC_CONFIG, "connect-config-topic-" + connectClusterName);
+        putIfAbsent(workerProps, CONFIG_STORAGE_REPLICATION_FACTOR_CONFIG, internalTopicsReplFactor);
+        putIfAbsent(workerProps, STATUS_STORAGE_TOPIC_CONFIG, "connect-storage-topic-" + connectClusterName);
+        putIfAbsent(workerProps, STATUS_STORAGE_REPLICATION_FACTOR_CONFIG, internalTopicsReplFactor);
+        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 < 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 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 String configureConnector(String connName, Map connConfig) {
+        String url = endpointForResource(String.format("connectors/%s/config", connName));
+        return putConnectorConfig(url, connConfig);
+    }
+
+    /**
+     * Validate a given connector configuration. If the configuration validates or
+     * has a configuration error, an instance of {@link ConfigInfos} is returned. If the validation fails
+     * an exception is thrown.
+     *
+     * @param connClassName the name of the connector class
+     * @param connConfig    the intended configuration
+     * @throws ConnectRestException if the REST api returns error status
+     * @throws ConnectException if the configuration fails to serialize/deserialize or if the request failed to send
+     */
+    public ConfigInfos validateConnectorConfig(String connClassName, Map connConfig) {
+        String url = endpointForResource(String.format("connector-plugins/%s/config/validate", connClassName));
+        String response = putConnectorConfig(url, connConfig);
+        ConfigInfos configInfos;
+        try {
+            configInfos = new ObjectMapper().readValue(response, ConfigInfos.class);
+        } catch (IOException e) {
+            throw new ConnectException("Unable deserialize response into a ConfigInfos object");
+        }
+        return configInfos;
+    }
+
+    /**
+     * Execute a PUT request with the given connector configuration on the given URL endpoint.
+     *
+     * @param url        the full URL of the endpoint that corresponds to the given REST resource
+     * @param connConfig the intended configuration
+     * @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
+     */
+    protected String putConnectorConfig(String url, Map connConfig) {
+        ObjectMapper mapper = new ObjectMapper();
+        String content;
+        try {
+            content = mapper.writeValueAsString(connConfig);
+        } catch (IOException e) {
+            throw new ConnectException("Could not serialize connector configuration and 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 ConnectRestException if the REST API returns error status
+     * @throws ConnectException for any other error.
+     */
+    public void deleteConnector(String connName) {
+        String url = endpointForResource(String.format("connectors/%s", connName));
+        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));
+        }
+    }
+
+    /**
+     * Pause an existing connector.
+     *
+     * @param connName name of the connector to be paused
+     * @throws ConnectRestException if the REST API returns error status
+     * @throws ConnectException for any other error.
+     */
+    public void pauseConnector(String connName) {
+        String url = endpointForResource(String.format("connectors/%s/pause", connName));
+        Response response = requestPut(url, "");
+        if (response.getStatus() >= Response.Status.BAD_REQUEST.getStatusCode()) {
+            throw new ConnectRestException(response.getStatus(),
+                "Could not execute PUT request. Error response: " + responseToString(response));
+        }
+    }
+
+    /**
+     * Resume an existing connector.
+     *
+     * @param connName name of the connector to be resumed
+     * @throws ConnectRestException if the REST API returns error status
+     * @throws ConnectException for any other error.
+     */
+    public void resumeConnector(String connName) {
+        String url = endpointForResource(String.format("connectors/%s/resume", connName));
+        Response response = requestPut(url, "");
+        if (response.getStatus() >= Response.Status.BAD_REQUEST.getStatusCode()) {
+            throw new ConnectRestException(response.getStatus(),
+                "Could not execute PUT request. Error response: " + responseToString(response));
+        }
+    }
+
+    /**
+     * Restart an existing connector.
+     *
+     * @param connName name of the connector to be restarted
+     * @throws ConnectRestException if the REST API returns error status
+     * @throws ConnectException for any other error.
+     */
+    public void restartConnector(String connName) {
+        String url = endpointForResource(String.format("connectors/%s/restart", connName));
+        Response response = requestPost(url, "", Collections.emptyMap());
+        if (response.getStatus() >= Response.Status.BAD_REQUEST.getStatusCode()) {
+            throw new ConnectRestException(response.getStatus(),
+                "Could not execute POST 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 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 {
+            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 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 active topics of a connector running in this cluster.
+     *
+     * @param connectorName name of the connector
+     * @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 ActiveTopicsInfo connectorTopics(String connectorName) {
+        ObjectMapper mapper = new ObjectMapper();
+        String url = endpointForResource(String.format("connectors/%s/topics", connectorName));
+        Response response = requestGet(url);
+        try {
+            if (response.getStatus() < Response.Status.BAD_REQUEST.getStatusCode()) {
+                Map>> activeTopics = mapper
+                        .readerFor(new TypeReference>>>() { })
+                        .readValue(responseToString(response));
+                return new ActiveTopicsInfo(connectorName,
+                        activeTopics.get(connectorName).getOrDefault("topics", Collections.emptyList()));
+            }
+        } catch (IOException 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));
+    }
+
+    /**
+     * Reset the set of active topics of a connector running in this cluster.
+     *
+     * @param connectorName name of the connector
+     * @throws ConnectRestException if the HTTP request to the REST API failed with a valid status code.
+     * @throws ConnectException for any other error.
+     */
+    public void resetConnectorTopics(String connectorName) {
+        String url = endpointForResource(String.format("connectors/%s/topics/reset", connectorName));
+        Response response = requestPut(url, null);
+        if (response.getStatus() >= Response.Status.BAD_REQUEST.getStatusCode()) {
+            throw new ConnectRestException(response.getStatus(),
+                    "Resetting active topics for connector " + connectorName + " failed. "
+                    + "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;
+    }
+
+    /**
+     * 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;
+    }
+
+    private static void putIfAbsent(Map props, String propertyKey, String propertyValue) {
+        if (!props.containsKey(propertyKey)) {
+            props.put(propertyKey, propertyValue);
+        }
+    }
+
+    /**
+     * Return the handle to the Kafka cluster this Connect cluster connects to.
+     *
+     * @return the Kafka cluster handle
+     */
+    public EmbeddedKafkaCluster kafka() {
+        return kafkaCluster;
+    }
+
+    /**
+     * 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 the response to the GET request
+     * @throws ConnectException if execution of the GET request fails
+     */
+    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);
+                }
+            }
+            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();
+            }
+        } catch (IOException e) {
+            log.error("Could not execute " + httpMethod + " request to " + url, e);
+            throw new ConnectException(e);
+        }
+    }
+
+    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 {
+        private String name = UUID.randomUUID().toString();
+        private Map workerProps = new HashMap<>();
+        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;
+            return this;
+        }
+
+        public Builder workerProps(Map workerProps) {
+            this.workerProps = workerProps;
+            return this;
+        }
+
+        public Builder numWorkers(int numWorkers) {
+            this.numWorkers = numWorkers;
+            return this;
+        }
+
+        public Builder numBrokers(int numBrokers) {
+            this.numBrokers = numBrokers;
+            return this;
+        }
+
+        public Builder brokerProps(Properties brokerProps) {
+            this.brokerProps = 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, 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..6e1f87e307e18
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/util/clusters/EmbeddedConnectClusterAssertions.java
@@ -0,0 +1,464 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.clients.admin.TopicDescription;
+import org.apache.kafka.connect.runtime.AbstractStatus;
+import org.apache.kafka.connect.runtime.rest.entities.ActiveTopicsInfo;
+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.Arrays;
+import java.util.Collection;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.BiFunction;
+import java.util.stream.Collectors;
+
+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 VALIDATION_DURATION_MS = TimeUnit.SECONDS.toMillis(30);
+    public static final long CONNECTOR_SETUP_DURATION_MS = TimeUnit.SECONDS.toMillis(30);
+    private static final long CONNECT_INTERNAL_TOPIC_UPDATES_DURATION_MS = TimeUnit.SECONDS.toMillis(60);
+
+    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 at least the requested number of workers are up and running.
+     *
+     * @param numBrokers the number of online brokers
+     */
+    public void assertExactlyNumBrokersAreUp(int numBrokers, String detailMessage) throws InterruptedException {
+        try {
+            waitForCondition(
+                () -> checkBrokersUp(numBrokers, (actual, expected) -> actual == expected).orElse(false),
+                WORKER_SETUP_DURATION_MS,
+                "Didn't meet the exact requested number of online brokers: " + numBrokers);
+        } catch (AssertionError e) {
+            throw new AssertionError(detailMessage, e);
+        }
+    }
+
+    /**
+     * Confirm that the requested number of brokers are up and running.
+     *
+     * @param numBrokers the number of online brokers
+     * @return true if at least {@code numBrokers} are up; false otherwise
+     */
+    protected Optional checkBrokersUp(int numBrokers, BiFunction comp) {
+        try {
+            int numRunning = connect.kafka().runningBrokers().size();
+            return Optional.of(comp.apply(numRunning, numBrokers));
+        } catch (Exception e) {
+            log.error("Could not check running brokers.", e);
+            return Optional.empty();
+        }
+    }
+
+    /**
+     * Assert that the topics with the specified names do not exist.
+     *
+     * @param topicNames the names of the topics that are expected to not exist
+     */
+    public void assertTopicsDoNotExist(String... topicNames) throws InterruptedException {
+        Set topicNameSet = new HashSet<>(Arrays.asList(topicNames));
+        AtomicReference> existingTopics = new AtomicReference<>(topicNameSet);
+        waitForCondition(
+            () -> checkTopicsExist(topicNameSet, (actual, expected) -> {
+                existingTopics.set(actual);
+                return actual.isEmpty();
+            }).orElse(false),
+            CONNECTOR_SETUP_DURATION_MS,
+            "Unexpectedly found topics " + existingTopics.get());
+    }
+
+    /**
+     * Assert that the topics with the specified names do exist.
+     *
+     * @param topicNames the names of the topics that are expected to exist
+     */
+    public void assertTopicsExist(String... topicNames) throws InterruptedException {
+        Set topicNameSet = new HashSet<>(Arrays.asList(topicNames));
+        AtomicReference> missingTopics = new AtomicReference<>(topicNameSet);
+        waitForCondition(
+            () -> checkTopicsExist(topicNameSet, (actual, expected) -> {
+                Set missing = new HashSet<>(expected);
+                missing.removeAll(actual);
+                missingTopics.set(missing);
+                return missing.isEmpty();
+            }).orElse(false),
+            CONNECTOR_SETUP_DURATION_MS,
+            "Didn't find the topics " + missingTopics.get());
+    }
+
+    protected Optional checkTopicsExist(Set topicNames, BiFunction, Set, Boolean> comp) {
+        try {
+            Map> topics = connect.kafka().describeTopics(topicNames);
+            Set actualExistingTopics = topics.entrySet()
+                                                     .stream()
+                                                     .filter(e -> e.getValue().isPresent())
+                                                     .map(Map.Entry::getKey)
+                                                     .collect(Collectors.toSet());
+            return Optional.of(comp.apply(actualExistingTopics, topicNames));
+        } catch (Exception e) {
+            log.error("Failed to describe the topic(s): {}.", topicNames, e);
+            return Optional.empty();
+        }
+    }
+
+    /**
+     * Assert that the named topic is configured to have the specified replication factor and
+     * number of partitions.
+     *
+     * @param topicName  the name of the topic that is expected to exist
+     * @param replicas   the replication factor
+     * @param partitions the number of partitions
+     * @param detailMessage the assertion message
+     */
+    public void assertTopicSettings(String topicName, int replicas, int partitions, String detailMessage)
+            throws InterruptedException {
+        try {
+            waitForCondition(
+                () -> checkTopicSettings(
+                    topicName,
+                    replicas,
+                    partitions
+                ).orElse(false),
+                VALIDATION_DURATION_MS,
+                "Topic " + topicName + " does not exist or does not have exactly "
+                        + partitions + " partitions or at least "
+                        + replicas + " per partition");
+        } catch (AssertionError e) {
+            throw new AssertionError(detailMessage, e);
+        }
+    }
+
+    protected Optional checkTopicSettings(String topicName, int replicas, int partitions) {
+        try {
+            Map> topics = connect.kafka().describeTopics(topicName);
+            TopicDescription topicDesc = topics.get(topicName).orElse(null);
+            boolean result = topicDesc != null
+                    && topicDesc.name().equals(topicName)
+                    && topicDesc.partitions().size() == partitions
+                    && topicDesc.partitions().stream().allMatch(p -> p.replicas().size() >= replicas);
+            return Optional.of(result);
+        } catch (Exception e) {
+            log.error("Failed to describe the topic: {}.", topicName, e);
+            return Optional.empty();
+        }
+    }
+
+    /**
+     * Assert that the required number of errors are produced by a connector config validation.
+     *
+     * @param connectorClass the class of the connector to validate
+     * @param connConfig     the intended configuration
+     * @param numErrors      the number of errors expected
+     */
+    public void assertExactlyNumErrorsOnConnectorConfigValidation(String connectorClass, Map connConfig,
+        int numErrors, String detailMessage) throws InterruptedException {
+        try {
+            waitForCondition(
+                () -> checkValidationErrors(
+                    connectorClass,
+                    connConfig,
+                    numErrors,
+                    (actual, expected) -> actual == expected
+                ).orElse(false),
+                VALIDATION_DURATION_MS,
+                "Didn't meet the exact requested number of validation errors: " + numErrors);
+        } catch (AssertionError e) {
+            throw new AssertionError(detailMessage, e);
+        }
+    }
+
+    /**
+     * Confirm that the requested number of errors are produced by {@link EmbeddedConnectCluster#validateConnectorConfig}.
+     *
+     * @param connectorClass the class of the connector to validate
+     * @param connConfig     the intended configuration
+     * @param numErrors      the number of errors expected
+     * @return true if exactly {@code numErrors} are produced by the validation; false otherwise
+     */
+    protected Optional checkValidationErrors(String connectorClass, Map connConfig,
+        int numErrors, BiFunction comp) {
+        try {
+            int numErrorsProduced = connect.validateConnectorConfig(connectorClass, connConfig).errorCount();
+            return Optional.of(comp.apply(numErrorsProduced, numErrors));
+        } catch (Exception e) {
+            log.error("Could not check config validation error count.", 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();
+        }
+    }
+
+    /**
+     * Assert that a connector's set of active topics matches the given collection of topic names.
+     *
+     * @param connectorName the connector name
+     * @param topics a collection of topics to compare against
+     * @param detailMessage the assertion message
+     * @throws InterruptedException
+     */
+    public void assertConnectorActiveTopics(String connectorName, Collection topics, String detailMessage) throws InterruptedException {
+        try {
+            waitForCondition(
+                () -> checkConnectorActiveTopics(connectorName, topics).orElse(false),
+                CONNECT_INTERNAL_TOPIC_UPDATES_DURATION_MS,
+                "Connector active topics don't match the expected collection");
+        } catch (AssertionError e) {
+            throw new AssertionError(detailMessage, e);
+        }
+    }
+
+    /**
+     * Check whether a connector's set of active topics matches the given collection of topic names.
+     *
+     * @param connectorName the connector name
+     * @param topics a collection of topics to compare against
+     * @return true if the connector's active topics matches the given collection; false otherwise
+     */
+    protected Optional checkConnectorActiveTopics(String connectorName, Collection topics) {
+        try {
+            ActiveTopicsInfo info = connect.connectorTopics(connectorName);
+            boolean result = info != null
+                    && topics.size() == info.topics().size()
+                    && topics.containsAll(info.topics());
+            log.debug("Found connector {} using topics: {}", connectorName, info.topics());
+            return Optional.of(result);
+        } catch (Exception e) {
+            log.error("Could not check connector {} state info.", connectorName, 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
new file mode 100644
index 0000000000000..9af828e362ba8
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/util/clusters/EmbeddedKafkaCluster.java
@@ -0,0 +1,468 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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 kafka.server.BrokerState;
+import kafka.server.KafkaConfig;
+import kafka.server.KafkaConfig$;
+import kafka.server.KafkaServer;
+import kafka.server.RunningAsBroker;
+import kafka.utils.CoreUtils;
+import kafka.utils.TestUtils;
+import kafka.zk.EmbeddedZookeeper;
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.clients.admin.AdminClientConfig;
+import org.apache.kafka.clients.admin.DescribeTopicsResult;
+import org.apache.kafka.clients.admin.NewTopic;
+import org.apache.kafka.clients.admin.TopicDescription;
+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.KafkaProducer;
+import org.apache.kafka.clients.producer.ProducerConfig;
+import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.KafkaFuture;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.config.SslConfigs;
+import org.apache.kafka.common.config.types.Password;
+import org.apache.kafka.common.errors.InvalidReplicationFactorException;
+import org.apache.kafka.common.errors.UnknownTopicOrPartitionException;
+import org.apache.kafka.common.network.ListenerName;
+import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.connect.errors.ConnectException;
+import org.junit.rules.ExternalResource;
+import org.junit.rules.TemporaryFolder;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Properties;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Predicate;
+import java.util.stream.Collectors;
+
+import static org.apache.kafka.clients.consumer.ConsumerConfig.AUTO_OFFSET_RESET_CONFIG;
+import static org.apache.kafka.clients.consumer.ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG;
+import static org.apache.kafka.clients.consumer.ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG;
+import static org.apache.kafka.clients.consumer.ConsumerConfig.GROUP_ID_CONFIG;
+import static org.apache.kafka.clients.consumer.ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG;
+import static org.apache.kafka.clients.consumer.ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG;
+
+/**
+ * Setup an embedded Kafka cluster with specified number of brokers and specified broker properties. To be used for
+ * integration tests.
+ */
+public class EmbeddedKafkaCluster extends ExternalResource {
+
+    private static final Logger log = LoggerFactory.getLogger(EmbeddedKafkaCluster.class);
+
+    private static final long DEFAULT_PRODUCE_SEND_DURATION_MS = TimeUnit.SECONDS.toMillis(120); 
+
+    // Kafka Config
+    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");
+    private KafkaProducer producer;
+
+    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() {
+        start();
+    }
+
+    @Override
+    protected void after() {
+        stop();
+    }
+
+    /**
+     * 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());
+
+        putIfAbsent(brokerConfig, KafkaConfig$.MODULE$.HostNameProp(), "localhost");
+        putIfAbsent(brokerConfig, KafkaConfig$.MODULE$.DeleteTopicEnableProp(), true);
+        putIfAbsent(brokerConfig, KafkaConfig$.MODULE$.GroupInitialRebalanceDelayMsProp(), 0);
+        putIfAbsent(brokerConfig, KafkaConfig$.MODULE$.OffsetsTopicReplicationFactorProp(), (short) brokers.length);
+        putIfAbsent(brokerConfig, KafkaConfig$.MODULE$.AutoCreateTopicsEnableProp(), false);
+
+        Object listenerConfig = brokerConfig.get(KafkaConfig$.MODULE$.InterBrokerListenerNameProp());
+        if (listenerConfig != null) {
+            listenerName = new ListenerName(listenerConfig.toString());
+        }
+
+        for (int i = 0; i < brokers.length; i++) {
+            brokerConfig.put(KafkaConfig$.MODULE$.BrokerIdProp(), i);
+            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<>();
+        producerProps.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers());
+        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");
+        producer = new KafkaProducer<>(producerProps);
+    }
+
+    public void stopOnlyKafka() {
+        stop(false, false);
+    }
+
+    private void stop() {
+        stop(true, true);
+    }
+
+    private void stop(boolean deleteLogDirs, boolean stopZK) {
+        try {
+            if (producer != null) {
+                producer.close();
+            }
+        } catch (Exception e) {
+            log.error("Could not shutdown producer ", e);
+            throw new RuntimeException("Could not shutdown producer", e);
+        }
+
+        for (KafkaServer broker : brokers) {
+            try {
+                broker.shutdown();
+            } catch (Throwable t) {
+                String msg = String.format("Could not shutdown 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 {
+            if (stopZK) {
+                zookeeper.shutdown();
+            }
+        } catch (Throwable t) {
+            String msg = String.format("Could not shutdown zookeeper at %s", zKConnectString());
+            log.error(msg, t);
+            throw new RuntimeException(msg, t);
+        }
+    }
+
+    private static void putIfAbsent(final Properties props, final String propertyKey, final Object propertyValue) {
+        if (!props.containsKey(propertyKey)) {
+            props.put(propertyKey, propertyValue);
+        }
+    }
+
+    private String createLogDir() {
+        TemporaryFolder tmpFolder = new TemporaryFolder();
+        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() {
+        return Arrays.stream(brokers)
+                .map(this::address)
+                .collect(Collectors.joining(","));
+    }
+
+    public String address(KafkaServer server) {
+        return server.config().hostName() + ":" + server.boundPort(listenerName);
+    }
+
+    public String zKConnectString() {
+        return "127.0.0.1:" + zookeeper.port();
+    }
+
+    /**
+     * Get the brokers that have a {@link RunningAsBroker} state.
+     *
+     * @return the list of {@link KafkaServer} instances that are running;
+     *         never null but  possibly empty
+     */
+    public Set runningBrokers() {
+        return brokersInState(state -> state.currentState() == RunningAsBroker.state());
+    }
+
+    /**
+     * Get the brokers whose state match the given predicate.
+     *
+     * @return the list of {@link KafkaServer} instances with states that match the predicate;
+     *         never null but  possibly empty
+     */
+    public Set brokersInState(Predicate desiredState) {
+        return Arrays.stream(brokers)
+                     .filter(b -> hasState(b, desiredState))
+                     .collect(Collectors.toSet());
+    }
+
+    protected boolean hasState(KafkaServer server, Predicate desiredState) {
+        try {
+            return desiredState.test(server.brokerState());
+        } catch (Throwable e) {
+            // Broker failed to respond.
+            return false;
+        }
+    }
+
+    /**
+     * Get the topic descriptions of the named topics. The value of the map entry will be empty
+     * if the topic does not exist.
+     *
+     * @param topicNames the names of the topics to describe
+     * @return the map of optional {@link TopicDescription} keyed by the topic name
+     */
+    public Map> describeTopics(String... topicNames) {
+        return describeTopics(new HashSet<>(Arrays.asList(topicNames)));
+    }
+
+    /**
+     * Get the topic descriptions of the named topics. The value of the map entry will be empty
+     * if the topic does not exist.
+     *
+     * @param topicNames the names of the topics to describe
+     * @return the map of optional {@link TopicDescription} keyed by the topic name
+     */
+    public Map> describeTopics(Set topicNames) {
+        Map> results = new HashMap<>();
+        log.info("Describing topics {}", topicNames);
+        try (Admin admin = createAdminClient()) {
+            DescribeTopicsResult result = admin.describeTopics(topicNames);
+            Map> byName = result.values();
+            for (Map.Entry> entry : byName.entrySet()) {
+                String topicName = entry.getKey();
+                try {
+                    TopicDescription desc = entry.getValue().get();
+                    results.put(topicName, Optional.of(desc));
+                    log.info("Found topic {} : {}", topicName, desc);
+                } catch (ExecutionException e) {
+                    Throwable cause = e.getCause();
+                    if (cause instanceof UnknownTopicOrPartitionException) {
+                        results.put(topicName, Optional.empty());
+                        log.info("Found non-existant topic {}", topicName);
+                        continue;
+                    }
+                    throw new AssertionError("Could not describe topic(s)" + topicNames, e);
+                }
+            }
+        } catch (Exception e) {
+            throw new AssertionError("Could not describe topic(s) " + topicNames, e);
+        }
+        log.info("Found topics {}", results);
+        return results;
+    }
+
+    /**
+     * Create a Kafka topic with 1 partition and a replication factor of 1.
+     *
+     * @param topic The name of the topic.
+     */
+    public void createTopic(String topic) {
+        createTopic(topic, 1);
+    }
+
+    /**
+     * Create a Kafka topic with given partition and a replication factor of 1.
+     *
+     * @param topic The name of the topic.
+     */
+    public void createTopic(String topic, int partitions) {
+        createTopic(topic, partitions, 1, new HashMap<>());
+    }
+
+    /**
+     * Create a Kafka topic with the given parameters.
+     *
+     * @param topic       The name of the topic.
+     * @param partitions  The number of partitions for this topic.
+     * @param replication The replication factor for (partitions of) this topic.
+     * @param topicConfig Additional topic-level configuration settings.
+     */
+    public void createTopic(String topic, int partitions, int replication, Map topicConfig) {
+        if (replication > brokers.length) {
+            throw new InvalidReplicationFactorException("Insufficient brokers ("
+                    + brokers.length + ") for desired replication (" + replication + ")");
+        }
+
+        log.debug("Creating topic { name: {}, partitions: {}, replication: {}, config: {} }",
+                topic, partitions, replication, topicConfig);
+        final NewTopic newTopic = new NewTopic(topic, partitions, (short) replication);
+        newTopic.configs(topicConfig);
+
+        try (final Admin adminClient = createAdminClient()) {
+            adminClient.createTopics(Collections.singletonList(newTopic)).all().get();
+        } catch (final InterruptedException | ExecutionException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    public void produce(String topic, String value) {
+        produce(topic, null, null, value);
+    }
+
+    public void produce(String topic, String key, String value) {
+        produce(topic, null, key, value);
+    }
+
+    public void produce(String topic, Integer partition, String key, String value) {
+        ProducerRecord msg = new ProducerRecord<>(topic, partition, key == null ? null : key.getBytes(), value == null ? null : value.getBytes());
+        try {
+            producer.send(msg).get(DEFAULT_PRODUCE_SEND_DURATION_MS, TimeUnit.MILLISECONDS);
+        } catch (Exception e) {
+            throw new KafkaException("Could not produce message: " + msg, e);
+        }
+    }
+
+    public Admin createAdminClient() {
+        final Properties adminClientConfig = new Properties();
+        adminClientConfig.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers());
+        final Object listeners = brokerConfig.get(KafkaConfig$.MODULE$.ListenersProp());
+        if (listeners != null && listeners.toString().contains("SSL")) {
+            adminClientConfig.put(SslConfigs.SSL_TRUSTSTORE_LOCATION_CONFIG, brokerConfig.get(SslConfigs.SSL_TRUSTSTORE_LOCATION_CONFIG));
+            adminClientConfig.put(SslConfigs.SSL_TRUSTSTORE_PASSWORD_CONFIG, ((Password) brokerConfig.get(SslConfigs.SSL_TRUSTSTORE_PASSWORD_CONFIG)).value());
+            adminClientConfig.put(CommonClientConfigs.SECURITY_PROTOCOL_CONFIG, "SSL");
+        }
+        return Admin.create(adminClientConfig);
+    }
+
+    /**
+     * Consume at least n records in a given duration or throw an exception.
+     *
+     * @param n the number of expected records in this topic.
+     * @param maxDuration the max duration to wait for these records (in milliseconds).
+     * @param topics the topics to subscribe and consume records from.
+     * @return a {@link ConsumerRecords} collection containing at least n records.
+     */
+    public ConsumerRecords consume(int n, long maxDuration, String... topics) {
+        Map>> records = new HashMap<>();
+        int consumedRecords = 0;
+        try (KafkaConsumer consumer = createConsumerAndSubscribeTo(Collections.emptyMap(), topics)) {
+            final long startMillis = System.currentTimeMillis();
+            long allowedDuration = maxDuration;
+            while (allowedDuration > 0) {
+                log.debug("Consuming from {} for {} millis.", Arrays.toString(topics), allowedDuration);
+                ConsumerRecords rec = consumer.poll(Duration.ofMillis(allowedDuration));
+                if (rec.isEmpty()) {
+                    allowedDuration = maxDuration - (System.currentTimeMillis() - startMillis);
+                    continue;
+                }
+                for (TopicPartition partition: rec.partitions()) {
+                    final List> r = rec.records(partition);
+                    records.computeIfAbsent(partition, t -> new ArrayList<>()).addAll(r);
+                    consumedRecords += r.size();
+                }
+                if (consumedRecords >= n) {
+                    return new ConsumerRecords<>(records);
+                }
+                allowedDuration = maxDuration - (System.currentTimeMillis() - startMillis);
+            }
+        }
+
+        throw new RuntimeException("Could not find enough records. found " + consumedRecords + ", expected " + n);
+    }
+
+    public KafkaConsumer createConsumer(Map consumerProps) {
+        Map props = new HashMap<>(consumerProps);
+
+        putIfAbsent(props, GROUP_ID_CONFIG, UUID.randomUUID().toString());
+        putIfAbsent(props, BOOTSTRAP_SERVERS_CONFIG, bootstrapServers());
+        putIfAbsent(props, ENABLE_AUTO_COMMIT_CONFIG, "false");
+        putIfAbsent(props, AUTO_OFFSET_RESET_CONFIG, "earliest");
+        putIfAbsent(props, KEY_DESERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.ByteArrayDeserializer");
+        putIfAbsent(props, VALUE_DESERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.ByteArrayDeserializer");
+
+        KafkaConsumer consumer;
+        try {
+            consumer = new KafkaConsumer<>(props);
+        } catch (Throwable t) {
+            throw new ConnectException("Failed to create consumer", t);
+        }
+        return consumer;
+    }
+
+    public KafkaConsumer createConsumerAndSubscribeTo(Map consumerProps, String... topics) {
+        KafkaConsumer consumer = createConsumer(consumerProps);
+        consumer.subscribe(Arrays.asList(topics));
+        return consumer;
+    }
+
+    private static void putIfAbsent(final Map props, final String propertyKey, final Object propertyValue) {
+        if (!props.containsKey(propertyKey)) {
+            props.put(propertyKey, propertyValue);
+        }
+    }
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/util/clusters/UngracefulShutdownException.java b/connect/runtime/src/test/java/org/apache/kafka/connect/util/clusters/UngracefulShutdownException.java
new file mode 100644
index 0000000000000..2f2b030efa1b2
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/util/clusters/UngracefulShutdownException.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.util.clusters;
+
+import org.apache.kafka.common.KafkaException;
+
+/**
+ * An exception that can be used from within an {@code Exit.Procedure} to mask exit or halt calls
+ * and signify that the service terminated abruptly. It's intended to be used only from within
+ * integration tests.
+ */
+public class UngracefulShutdownException extends KafkaException {
+    public UngracefulShutdownException(String s) {
+        super(s);
+    }
+
+    public UngracefulShutdownException(String s, Throwable throwable) {
+        super(s, throwable);
+    }
+
+    public UngracefulShutdownException(Throwable throwable) {
+        super(throwable);
+    }
+}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/util/clusters/WorkerHandle.java b/connect/runtime/src/test/java/org/apache/kafka/connect/util/clusters/WorkerHandle.java
new file mode 100644
index 0000000000000..4d947940c585e
--- /dev/null
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/util/clusters/WorkerHandle.java
@@ -0,0 +1,121 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.cli.ConnectDistributed;
+import org.apache.kafka.connect.runtime.Connect;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.net.URI;
+import java.util.Map;
+import java.util.Objects;
+
+/**
+ * A handle to a worker executing in a Connect cluster.
+ */
+public class WorkerHandle {
+    private static final Logger log = LoggerFactory.getLogger(WorkerHandle.class);
+
+    private final String workerName;
+    private final Connect worker;
+
+    protected WorkerHandle(String workerName, Connect worker) {
+        this.workerName = workerName;
+        this.worker = worker;
+    }
+
+    /**
+     * Create and start a new worker with the given properties.
+     *
+     * @param name a name for this worker
+     * @param workerProperties the worker properties
+     * @return the worker's handle
+     */
+    public static WorkerHandle start(String name, Map workerProperties) {
+        return new WorkerHandle(name, new ConnectDistributed().startConnect(workerProperties));
+    }
+
+    /**
+     * Stop this worker.
+     */
+    public void stop() {
+        worker.stop();
+    }
+
+    /**
+     * Determine if this worker is running.
+     *
+     * @return true if the worker is running, or false otherwise
+     */
+    public boolean isRunning() {
+        return worker.isRunning();
+    }
+
+    /**
+     * 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/META-INF/services/org.apache.kafka.connect.rest.ConnectRestExtension b/connect/runtime/src/test/resources/META-INF/services/org.apache.kafka.connect.rest.ConnectRestExtension
new file mode 100644
index 0000000000000..0a1ef88924e1e
--- /dev/null
+++ b/connect/runtime/src/test/resources/META-INF/services/org.apache.kafka.connect.rest.ConnectRestExtension
@@ -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.
+
+org.apache.kafka.connect.runtime.isolation.PluginsTest$TestConnectRestExtension
\ No newline at end of file
diff --git a/connect/runtime/src/test/resources/log4j.properties b/connect/runtime/src/test/resources/log4j.properties
index d5e90fe788f76..176692deb7b2b 100644
--- a/connect/runtime/src/test/resources/log4j.properties
+++ b/connect/runtime/src/test/resources/log4j.properties
@@ -14,10 +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] %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.apache.kafka=ERROR
+log4j.logger.org.reflections=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/connect/runtime/src/test/resources/test-plugins/service-loader/test/plugins/ServiceLoadedClass.java b/connect/runtime/src/test/resources/test-plugins/service-loader/test/plugins/ServiceLoadedClass.java
new file mode 100644
index 0000000000000..98677ed43d65d
--- /dev/null
+++ b/connect/runtime/src/test/resources/test-plugins/service-loader/test/plugins/ServiceLoadedClass.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 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 17be48ce17443..e872b336e8573 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,17 +24,23 @@
 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;
 import org.apache.kafka.connect.transforms.util.SimpleConfig;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
-import java.util.Arrays;
+import java.util.EnumSet;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
@@ -44,6 +50,7 @@
 import static org.apache.kafka.connect.transforms.util.Requirements.requireStruct;
 
 public abstract class Cast> implements Transformation {
+    private static final Logger log = LoggerFactory.getLogger(Cast.class);
 
     // TODO: Currently we only support top-level field casting. Ideally we could use a dotted notation in the spec to
     // allow casting nested fields.
@@ -79,9 +86,16 @@ public String toString() {
 
     private static final String PURPOSE = "cast types";
 
-    private static final Set SUPPORTED_CAST_TYPES = new HashSet<>(
-            Arrays.asList(Schema.Type.INT8, Schema.Type.INT16, Schema.Type.INT32, Schema.Type.INT64,
-                    Schema.Type.FLOAT32, Schema.Type.FLOAT64, Schema.Type.BOOLEAN, Schema.Type.STRING)
+    private static final Set SUPPORTED_CAST_INPUT_TYPES = EnumSet.of(
+            Schema.Type.INT8, Schema.Type.INT16, Schema.Type.INT32, Schema.Type.INT64,
+                    Schema.Type.FLOAT32, Schema.Type.FLOAT64, Schema.Type.BOOLEAN,
+                            Schema.Type.STRING, Schema.Type.BYTES
+    );
+
+    private static final Set SUPPORTED_CAST_OUTPUT_TYPES = EnumSet.of(
+            Schema.Type.INT8, Schema.Type.INT16, Schema.Type.INT32, Schema.Type.INT64,
+                    Schema.Type.FLOAT32, Schema.Type.FLOAT64, Schema.Type.BOOLEAN,
+                            Schema.Type.STRING
     );
 
     // As a special case for casting the entire value (e.g. the incoming key is a int64 but you know it could be an
@@ -97,7 +111,7 @@ public void configure(Map props) {
         final SimpleConfig config = new SimpleConfig(CONFIG_DEF, props);
         casts = parseFieldTypes(config.getList(SPEC_CONFIG));
         wholeValueCastType = casts.get(WHOLE_VALUE_CAST);
-        schemaUpdateCache = new SynchronizedCache<>(new LRUCache(16));
+        schemaUpdateCache = new SynchronizedCache<>(new LRUCache<>(16));
     }
 
     @Override
@@ -121,14 +135,14 @@ public void close() {
 
     private R applySchemaless(R record) {
         if (wholeValueCastType != null) {
-            return newRecord(record, null, castValueToType(operatingValue(record), wholeValueCastType));
+            return newRecord(record, null, castValueToType(null, operatingValue(record), wholeValueCastType));
         }
 
         final Map value = requireMap(operatingValue(record), PURPOSE);
         final HashMap updatedValue = new HashMap<>(value);
         for (Map.Entry fieldSpec : casts.entrySet()) {
             String field = fieldSpec.getKey();
-            updatedValue.put(field, castValueToType(value.get(field), fieldSpec.getValue()));
+            updatedValue.put(field, castValueToType(null, value.get(field), fieldSpec.getValue()));
         }
         return newRecord(record, null, updatedValue);
     }
@@ -139,7 +153,7 @@ private R applyWithSchema(R record) {
 
         // Whole-record casting
         if (wholeValueCastType != null)
-            return newRecord(record, updatedSchema, castValueToType(operatingValue(record), wholeValueCastType));
+            return newRecord(record, updatedSchema, castValueToType(valueSchema, operatingValue(record), wholeValueCastType));
 
         // Casting within a struct
         final Struct value = requireStruct(operatingValue(record), PURPOSE);
@@ -148,7 +162,8 @@ private R applyWithSchema(R record) {
         for (Field field : value.schema().fields()) {
             final Object origFieldValue = value.get(field);
             final Schema.Type targetType = casts.get(field.name());
-            final Object newFieldValue = targetType != null ? castValueToType(origFieldValue, targetType) : origFieldValue;
+            final Object newFieldValue = targetType != null ? castValueToType(field.schema(), origFieldValue, targetType) : origFieldValue;
+            log.trace("Cast field '{}' from '{}' to '{}'", field.name(), origFieldValue, newFieldValue);
             updatedValue.put(updatedSchema.field(field.name()), newFieldValue);
         }
         return newRecord(record, updatedSchema, updatedValue);
@@ -165,20 +180,25 @@ private Schema getOrBuildSchema(Schema valueSchema) {
         } else {
             builder = SchemaUtil.copySchemaBasics(valueSchema, SchemaBuilder.struct());
             for (Field field : valueSchema.fields()) {
-                SchemaBuilder fieldBuilder =
-                        convertFieldType(casts.containsKey(field.name()) ? casts.get(field.name()) : field.schema().type());
-                if (field.schema().isOptional())
-                    fieldBuilder.optional();
-                if (field.schema().defaultValue() != null)
-                    fieldBuilder.defaultValue(castValueToType(field.schema().defaultValue(), fieldBuilder.type()));
-                builder.field(field.name(), fieldBuilder.build());
+                if (casts.containsKey(field.name())) {
+                    SchemaBuilder fieldBuilder = convertFieldType(casts.get(field.name()));
+                    if (field.schema().isOptional())
+                        fieldBuilder.optional();
+                    if (field.schema().defaultValue() != null) {
+                        Schema fieldSchema = field.schema();
+                        fieldBuilder.defaultValue(castValueToType(fieldSchema, fieldSchema.defaultValue(), fieldBuilder.type()));
+                    }
+                    builder.field(field.name(), fieldBuilder.build());
+                } else {
+                    builder.field(field.name(), field.schema());
+                }
             }
         }
 
         if (valueSchema.isOptional())
             builder.optional();
         if (valueSchema.defaultValue() != null)
-            builder.defaultValue(castValueToType(valueSchema.defaultValue(), builder.type()));
+            builder.defaultValue(castValueToType(valueSchema, valueSchema.defaultValue(), builder.type()));
 
         updatedSchema = builder.build();
         schemaUpdateCache.put(valueSchema, updatedSchema);
@@ -206,14 +226,26 @@ 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(Object value, Schema.Type targetType) {
+    private static Object castValueToType(Schema schema, Object value, Schema.Type targetType) {
         try {
             if (value == null) return null;
 
-            Schema.Type inferredType = ConnectSchema.schemaType(value.getClass());
+            Schema.Type inferredType = schema == null ? ConnectSchema.schemaType(value.getClass()) :
+                    schema.type();
             if (inferredType == null) {
                 throw new DataException("Cast transformation was passed a value of type " + value.getClass()
                         + " which is not supported by Connect's data API");
@@ -221,6 +253,11 @@ private static Object castValueToType(Object value, Schema.Type targetType) {
             // 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);
@@ -324,7 +361,12 @@ else if (value instanceof String)
     }
 
     private static String castToString(Object value) {
-        return value.toString();
+        if (value instanceof java.util.Date) {
+            java.util.Date dateValue = (java.util.Date) value;
+            return Values.dateFormatFor(dateValue).format(dateValue);
+        } else {
+            return value.toString();
+        }
     }
 
     protected abstract Schema operatingSchema(R record);
@@ -367,15 +409,19 @@ private enum FieldType {
     }
 
     private static Schema.Type validCastType(Schema.Type type, FieldType fieldType) {
-        if (!SUPPORTED_CAST_TYPES.contains(type)) {
-            String message = "Cast transformation does not support casting to/from " + type
-                    + "; supported types are " + SUPPORTED_CAST_TYPES;
-            switch (fieldType) {
-                case INPUT:
-                    throw new DataException(message);
-                case OUTPUT:
-                    throw new ConfigException(message);
-            }
+        switch (fieldType) {
+            case INPUT:
+                if (!SUPPORTED_CAST_INPUT_TYPES.contains(type)) {
+                    throw new DataException("Cast transformation does not support casting from " +
+                        type + "; supported types are " + SUPPORTED_CAST_INPUT_TYPES);
+                }
+                break;
+            case OUTPUT:
+                if (!SUPPORTED_CAST_OUTPUT_TYPES.contains(type)) {
+                    throw new ConfigException("Cast transformation does not support casting to " +
+                        type + "; supported types are " + SUPPORTED_CAST_OUTPUT_TYPES);
+                }
+                break;
         }
         return type;
     }
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/Filter.java b/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/Filter.java
new file mode 100644
index 0000000000000..d7fb54eaca9db
--- /dev/null
+++ b/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/Filter.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.connect.transforms;
+
+import java.util.Map;
+
+import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.connect.connector.ConnectRecord;
+
+/**
+ * Drops all records, filtering them from subsequent transformations in the chain.
+ * This is intended to be used conditionally to filter out records matching (or not matching)
+ * a particular {@link org.apache.kafka.connect.transforms.predicates.Predicate}.
+ * @param  The type of record.
+ */
+public class Filter> implements Transformation {
+
+    public static final String OVERVIEW_DOC = "Drops all records, filtering them from subsequent transformations in the chain. " +
+            "This is intended to be used conditionally to filter out records matching (or not matching) " +
+            "a particular Predicate.";
+    public static final ConfigDef CONFIG_DEF = new ConfigDef();
+
+    @Override
+    public R apply(R record) {
+        return null;
+    }
+
+    @Override
+    public ConfigDef config() {
+        return CONFIG_DEF;
+    }
+
+    @Override
+    public void close() {
+
+    }
+
+    @Override
+    public void configure(Map configs) {
+
+    }
+}
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..cad8d79e99f34 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 {
 
@@ -64,12 +64,14 @@ public abstract class Flatten> implements Transformat
     public void configure(Map props) {
         final SimpleConfig config = new SimpleConfig(CONFIG_DEF, props);
         delimiter = config.getString(DELIMITER_CONFIG);
-        schemaUpdateCache = new SynchronizedCache<>(new LRUCache(16));
+        schemaUpdateCache = new SynchronizedCache<>(new LRUCache<>(16));
     }
 
     @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/HoistField.java b/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/HoistField.java
index a852c797a3611..3614104e4588f 100644
--- a/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/HoistField.java
+++ b/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/HoistField.java
@@ -50,7 +50,7 @@ public abstract class HoistField> implements Transfor
     public void configure(Map props) {
         final SimpleConfig config = new SimpleConfig(CONFIG_DEF, props);
         fieldName = config.getString("field");
-        schemaUpdateCache = new SynchronizedCache<>(new LRUCache(16));
+        schemaUpdateCache = new SynchronizedCache<>(new LRUCache<>(16));
     }
 
     @Override
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..cbc820b65ccf0 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
@@ -122,12 +122,14 @@ public void configure(Map props) {
             throw new ConfigException(ConfigName.STATIC_VALUE, null, "No value specified for static field: " + staticField);
         }
 
-        schemaUpdateCache = new SynchronizedCache<>(new LRUCache(16));
+        schemaUpdateCache = new SynchronizedCache<>(new LRUCache<>(16));
     }
 
     @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/MaskField.java b/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/MaskField.java
index 050c549c97e16..0d61ac1a0596a 100644
--- a/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/MaskField.java
+++ b/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/MaskField.java
@@ -21,6 +21,7 @@
 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.data.Values;
 import org.apache.kafka.connect.errors.DataException;
 import org.apache.kafka.connect.transforms.util.NonEmptyListValidator;
 import org.apache.kafka.connect.transforms.util.SimpleConfig;
@@ -34,6 +35,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.Function;
 
 import static org.apache.kafka.connect.transforms.util.Requirements.requireMap;
 import static org.apache.kafka.connect.transforms.util.Requirements.requireStruct;
@@ -42,16 +44,23 @@ public abstract class MaskField> implements Transform
 
     public static final String OVERVIEW_DOC =
             "Mask specified fields with a valid null value for the field type (i.e. 0, false, empty string, and so on)."
-                    + "
            Use the concrete transformation type designed for the record key (" + Key.class.getName() + ") "
-                    + "or value (" + Value.class.getName() + ").";
+                    + "
            For numeric and string fields, an optional replacement value can be specified that is converted to the correct type."
+                    + "
            Use the concrete transformation type designed for the record key (" + Key.class.getName()
+                    + ") or value (" + Value.class.getName() + ").";
 
     private static final String FIELDS_CONFIG = "fields";
+    private static final String REPLACEMENT_CONFIG = "replacement";
 
     public static final ConfigDef CONFIG_DEF = new ConfigDef()
-            .define(FIELDS_CONFIG, ConfigDef.Type.LIST, ConfigDef.NO_DEFAULT_VALUE, new NonEmptyListValidator(), ConfigDef.Importance.HIGH, "Names of fields to mask.");
+            .define(FIELDS_CONFIG, ConfigDef.Type.LIST, ConfigDef.NO_DEFAULT_VALUE, new NonEmptyListValidator(),
+                    ConfigDef.Importance.HIGH, "Names of fields to mask.")
+            .define(REPLACEMENT_CONFIG, ConfigDef.Type.STRING, null, new ConfigDef.NonEmptyString(),
+                    ConfigDef.Importance.LOW, "Custom value replacement, that will be applied to all"
+                            + " 'fields' values (numeric or non-empty string values only).");
 
     private static final String PURPOSE = "mask fields";
 
+    private static final Map, Function> REPLACEMENT_MAPPING_FUNC = new HashMap<>();
     private static final Map, Object> PRIMITIVE_VALUE_MAPPING = new HashMap<>();
 
     static {
@@ -66,14 +75,26 @@ public abstract class MaskField> implements Transform
         PRIMITIVE_VALUE_MAPPING.put(BigDecimal.class, BigDecimal.ZERO);
         PRIMITIVE_VALUE_MAPPING.put(Date.class, new Date(0));
         PRIMITIVE_VALUE_MAPPING.put(String.class, "");
+
+        REPLACEMENT_MAPPING_FUNC.put(Byte.class, v -> Values.convertToByte(null, v));
+        REPLACEMENT_MAPPING_FUNC.put(Short.class, v -> Values.convertToShort(null, v));
+        REPLACEMENT_MAPPING_FUNC.put(Integer.class, v -> Values.convertToInteger(null, v));
+        REPLACEMENT_MAPPING_FUNC.put(Long.class, v -> Values.convertToLong(null, v));
+        REPLACEMENT_MAPPING_FUNC.put(Float.class, v -> Values.convertToFloat(null, v));
+        REPLACEMENT_MAPPING_FUNC.put(Double.class, v -> Values.convertToDouble(null, v));
+        REPLACEMENT_MAPPING_FUNC.put(String.class, Function.identity());
+        REPLACEMENT_MAPPING_FUNC.put(BigDecimal.class, BigDecimal::new);
+        REPLACEMENT_MAPPING_FUNC.put(BigInteger.class, BigInteger::new);
     }
 
     private Set maskedFields;
+    private String replacement;
 
     @Override
     public void configure(Map props) {
         final SimpleConfig config = new SimpleConfig(CONFIG_DEF, props);
         maskedFields = new HashSet<>(config.getList(FIELDS_CONFIG));
+        replacement = config.getString(REPLACEMENT_CONFIG);
     }
 
     @Override
@@ -104,9 +125,26 @@ private R applyWithSchema(R record) {
         return newRecord(record, updatedValue);
     }
 
-    private static Object masked(Object value) {
-        if (value == null)
+    private Object masked(Object value) {
+        if (value == null) {
             return null;
+        }
+        return replacement == null ? maskWithNullValue(value) : maskWithCustomReplacement(value, replacement);
+    }
+
+    private static Object maskWithCustomReplacement(Object value, String replacement) {
+        Function replacementMapper = REPLACEMENT_MAPPING_FUNC.get(value.getClass());
+        if (replacementMapper == null) {
+            throw new DataException("Cannot mask value of type " + value.getClass() + " with custom replacement.");
+        }
+        try {
+            return replacementMapper.apply(replacement);
+        } catch (NumberFormatException ex) {
+            throw new DataException("Unable to convert " + replacement + " (" + replacement.getClass() + ") to number", ex);
+        }
+    }
+
+    private static Object maskWithNullValue(Object value) {
         Object maskedValue = PRIMITIVE_VALUE_MAPPING.get(value.getClass());
         if (maskedValue == null) {
             if (value instanceof List)
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 ee08945416683..fb02577e99ac7 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
@@ -20,7 +20,9 @@
 import org.apache.kafka.common.cache.LRUCache;
 import org.apache.kafka.common.cache.SynchronizedCache;
 import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.common.config.ConfigDef.Importance;
 import org.apache.kafka.common.config.ConfigException;
+import org.apache.kafka.common.utils.ConfigUtils;
 import org.apache.kafka.connect.connector.ConnectRecord;
 import org.apache.kafka.connect.data.Field;
 import org.apache.kafka.connect.data.Schema;
@@ -44,17 +46,27 @@ public abstract class ReplaceField> implements Transf
             + "or value (" + Value.class.getName() + ").";
 
     interface ConfigName {
-        String BLACKLIST = "blacklist";
-        String WHITELIST = "whitelist";
+        String EXCLUDE = "exclude";
+        String INCLUDE = "include";
+
+        // for backwards compatibility
+        String INCLUDE_ALIAS = "whitelist";
+        String EXCLUDE_ALIAS = "blacklist";
+
         String RENAME = "renames";
     }
 
     public static final ConfigDef CONFIG_DEF = new ConfigDef()
-            .define(ConfigName.BLACKLIST, ConfigDef.Type.LIST, Collections.emptyList(), ConfigDef.Importance.MEDIUM,
-                    "Fields to exclude. This takes precedence over the whitelist.")
-            .define(ConfigName.WHITELIST, ConfigDef.Type.LIST, Collections.emptyList(), ConfigDef.Importance.MEDIUM,
+            .define(ConfigName.EXCLUDE, ConfigDef.Type.LIST, Collections.emptyList(), ConfigDef.Importance.MEDIUM,
+                    "Fields to exclude. This takes precedence over the fields to include.")
+            .define("blacklist", ConfigDef.Type.LIST, null, Importance.LOW,
+                    "Deprecated. Use " + ConfigName.EXCLUDE + " instead.")
+            .define(ConfigName.INCLUDE, ConfigDef.Type.LIST, Collections.emptyList(), ConfigDef.Importance.MEDIUM,
                     "Fields to include. If specified, only these fields will be used.")
+            .define("whitelist", ConfigDef.Type.LIST, null, Importance.LOW,
+                    "Deprecated. Use " + ConfigName.INCLUDE + " instead.")
             .define(ConfigName.RENAME, ConfigDef.Type.LIST, Collections.emptyList(), new ConfigDef.Validator() {
+                @SuppressWarnings("unchecked")
                 @Override
                 public void ensureValid(String name, Object value) {
                     parseRenameMappings((List) value);
@@ -68,8 +80,8 @@ public String toString() {
 
     private static final String PURPOSE = "field replacement";
 
-    private List blacklist;
-    private List whitelist;
+    private List exclude;
+    private List include;
     private Map renames;
     private Map reverseRenames;
 
@@ -77,13 +89,17 @@ public String toString() {
 
     @Override
     public void configure(Map configs) {
-        final SimpleConfig config = new SimpleConfig(CONFIG_DEF, configs);
-        blacklist = config.getList(ConfigName.BLACKLIST);
-        whitelist = config.getList(ConfigName.WHITELIST);
+        final SimpleConfig config = new SimpleConfig(CONFIG_DEF, ConfigUtils.translateDeprecatedConfigs(configs, new String[][]{
+            {ConfigName.INCLUDE, "whitelist"},
+            {ConfigName.EXCLUDE, "blacklist"},
+        }));
+
+        exclude = config.getList(ConfigName.EXCLUDE);
+        include = config.getList(ConfigName.INCLUDE);
         renames = parseRenameMappings(config.getList(ConfigName.RENAME));
         reverseRenames = invert(renames);
 
-        schemaUpdateCache = new SynchronizedCache<>(new LRUCache(16));
+        schemaUpdateCache = new SynchronizedCache<>(new LRUCache<>(16));
     }
 
     static Map parseRenameMappings(List mappings) {
@@ -107,7 +123,7 @@ static Map invert(Map source) {
     }
 
     boolean filter(String fieldName) {
-        return !blacklist.contains(fieldName) && (whitelist.isEmpty() || whitelist.contains(fieldName));
+        return !exclude.contains(fieldName) && (include.isEmpty() || include.contains(fieldName));
     }
 
     String renamed(String fieldName) {
@@ -122,7 +138,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/SetSchemaMetadata.java b/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/SetSchemaMetadata.java
index 901ac9f18ad88..fd3cbf314ff5a 100644
--- a/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/SetSchemaMetadata.java
+++ b/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/SetSchemaMetadata.java
@@ -24,12 +24,15 @@
 import org.apache.kafka.connect.data.Schema;
 import org.apache.kafka.connect.data.Struct;
 import org.apache.kafka.connect.transforms.util.SimpleConfig;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import java.util.Map;
 
 import static org.apache.kafka.connect.transforms.util.Requirements.requireSchema;
 
 public abstract class SetSchemaMetadata> implements Transformation {
+    private static final Logger log = LoggerFactory.getLogger(SetSchemaMetadata.class);
 
     public static final String OVERVIEW_DOC =
             "Set the schema name, version or both on the record's key (" + Key.class.getName() + ")"
@@ -76,6 +79,8 @@ public R apply(R record) {
                 isMap ? schema.keySchema() : null,
                 isMap || isArray ? schema.valueSchema() : null
         );
+        log.trace("Applying SetSchemaMetadata SMT. Original schema: {}, updated schema: {}",
+            schema, updatedSchema);
         return newRecord(record, updatedSchema);
     }
 
@@ -149,4 +154,4 @@ protected static Object updateSchemaIn(Object keyOrValue, Schema updatedSchema)
         }
         return keyOrValue;
     }
-}
\ No newline at end of file
+}
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..b92dc848fc0ca 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
@@ -245,7 +249,7 @@ public void configure(Map configs) {
         final String field = simpleConfig.getString(FIELD_CONFIG);
         final String type = simpleConfig.getString(TARGET_TYPE_CONFIG);
         String formatPattern = simpleConfig.getString(FORMAT_CONFIG);
-        schemaUpdateCache = new SynchronizedCache<>(new LRUCache(16));
+        schemaUpdateCache = new SynchronizedCache<>(new LRUCache<>(16));
 
         if (!VALID_TYPES.contains(type)) {
             throw new ConfigException("Unknown timestamp type in TimestampConverter: " + type + ". Valid values are "
@@ -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/TimestampRouter.java b/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/TimestampRouter.java
index 938ef0fc27729..388a7b941fa6d 100644
--- a/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/TimestampRouter.java
+++ b/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/TimestampRouter.java
@@ -25,8 +25,14 @@
 import java.util.Date;
 import java.util.Map;
 import java.util.TimeZone;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 
-public class TimestampRouter> implements Transformation {
+public class TimestampRouter> implements Transformation, AutoCloseable {
+
+    private static final Pattern TOPIC = Pattern.compile("${topic}", Pattern.LITERAL);
+
+    private static final Pattern TIMESTAMP = Pattern.compile("${timestamp}", Pattern.LITERAL);
 
     public static final String OVERVIEW_DOC =
             "Update the record's topic field as a function of the original topic value and the record timestamp."
@@ -72,7 +78,9 @@ public R apply(R record) {
             throw new DataException("Timestamp missing on record: " + record);
         }
         final String formattedTimestamp = timestampFormat.get().format(new Date(timestamp));
-        final String updatedTopic = topicFormat.replace("${topic}", record.topic()).replace("${timestamp}", formattedTimestamp);
+
+        final String replace1 = TOPIC.matcher(topicFormat).replaceAll(Matcher.quoteReplacement(record.topic()));
+        final String updatedTopic = TIMESTAMP.matcher(replace1).replaceAll(Matcher.quoteReplacement(formattedTimestamp));
         return record.newRecord(
                 updatedTopic, record.kafkaPartition(),
                 record.keySchema(), record.key(),
@@ -83,7 +91,7 @@ public R apply(R record) {
 
     @Override
     public void close() {
-        timestampFormat = null;
+        timestampFormat.remove();
     }
 
     @Override
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..8f843f4aaaef3 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;
 
@@ -54,7 +56,7 @@ public class ValueToKey> implements Transformation
     public void configure(Map configs) {
         final SimpleConfig config = new SimpleConfig(CONFIG_DEF, configs);
         fields = config.getList(FIELDS_CONFIG);
-        valueToKeySchemaCache = new SynchronizedCache<>(new LRUCache(16));
+        valueToKeySchemaCache = new SynchronizedCache<>(new LRUCache<>(16));
     }
 
     @Override
@@ -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/main/java/org/apache/kafka/connect/transforms/predicates/HasHeaderKey.java b/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/predicates/HasHeaderKey.java
new file mode 100644
index 0000000000000..f15d426c026a8
--- /dev/null
+++ b/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/predicates/HasHeaderKey.java
@@ -0,0 +1,68 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.transforms.predicates;
+
+import java.util.Iterator;
+import java.util.Map;
+
+import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.connect.connector.ConnectRecord;
+import org.apache.kafka.connect.header.Header;
+import org.apache.kafka.connect.transforms.util.SimpleConfig;
+
+/**
+ * A predicate which is true for records with at least one header with the configured name.
+ * @param  The type of connect record.
+ */
+public class HasHeaderKey> implements Predicate {
+
+    private static final String NAME_CONFIG = "name";
+    public static final String OVERVIEW_DOC = "A predicate which is true for records with at least one header with the configured name.";
+    public static final ConfigDef CONFIG_DEF = new ConfigDef()
+            .define(NAME_CONFIG, ConfigDef.Type.STRING, ConfigDef.NO_DEFAULT_VALUE,
+            new ConfigDef.NonEmptyString(), ConfigDef.Importance.MEDIUM,
+            "The header name.");
+    private String name;
+
+    @Override
+    public ConfigDef config() {
+        return CONFIG_DEF;
+    }
+
+    @Override
+    public boolean test(R record) {
+        Iterator
             headerIterator = record.headers().allWithName(name);
+        return headerIterator != null && headerIterator.hasNext();
+    }
+
+    @Override
+    public void close() {
+
+    }
+
+    @Override
+    public void configure(Map configs) {
+        this.name = new SimpleConfig(config(), configs).getString(NAME_CONFIG);
+    }
+
+    @Override
+    public String toString() {
+        return "HasHeaderKey{" +
+                "name='" + name + '\'' +
+                '}';
+    }
+}
diff --git a/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/predicates/RecordIsTombstone.java b/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/predicates/RecordIsTombstone.java
new file mode 100644
index 0000000000000..4a21eacc7ced3
--- /dev/null
+++ b/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/predicates/RecordIsTombstone.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.transforms.predicates;
+
+import java.util.Map;
+
+import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.connect.connector.ConnectRecord;
+
+/**
+ * A predicate which is true for records which are tombstones (i.e. have null value).
+ * @param  The type of connect record.
+ */
+public class RecordIsTombstone> implements Predicate {
+
+    public static final String OVERVIEW_DOC = "A predicate which is true for records which are tombstones (i.e. have null value).";
+    public static final ConfigDef CONFIG_DEF = new ConfigDef();
+
+    @Override
+    public ConfigDef config() {
+        return CONFIG_DEF;
+    }
+
+    @Override
+    public boolean test(R record) {
+        return record.value() == null;
+    }
+
+    @Override
+    public void close() {
+
+    }
+
+    @Override
+    public void configure(Map configs) {
+
+    }
+
+    @Override
+    public String toString() {
+        return "RecordIsTombstone{}";
+    }
+}
diff --git a/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/predicates/TopicNameMatches.java b/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/predicates/TopicNameMatches.java
new file mode 100644
index 0000000000000..3ea8f1ae95676
--- /dev/null
+++ b/connect/transforms/src/main/java/org/apache/kafka/connect/transforms/predicates/TopicNameMatches.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.transforms.predicates;
+
+import java.util.Map;
+import java.util.regex.Pattern;
+import java.util.regex.PatternSyntaxException;
+
+import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.common.config.ConfigException;
+import org.apache.kafka.connect.connector.ConnectRecord;
+import org.apache.kafka.connect.transforms.util.RegexValidator;
+import org.apache.kafka.connect.transforms.util.SimpleConfig;
+
+/**
+ * A predicate which is true for records with a topic name that matches the configured regular expression.
+ * @param  The type of connect record.
+ */
+public class TopicNameMatches> implements Predicate {
+
+    private static final String PATTERN_CONFIG = "pattern";
+
+    public static final String OVERVIEW_DOC = "A predicate which is true for records with a topic name that matches the configured regular expression.";
+
+    public static final ConfigDef CONFIG_DEF = new ConfigDef()
+            .define(PATTERN_CONFIG, ConfigDef.Type.STRING, ConfigDef.NO_DEFAULT_VALUE,
+            ConfigDef.CompositeValidator.of(new ConfigDef.NonEmptyString(), new RegexValidator()),
+            ConfigDef.Importance.MEDIUM,
+            "A Java regular expression for matching against the name of a record's topic.");
+    private Pattern pattern;
+
+    @Override
+    public ConfigDef config() {
+        return CONFIG_DEF;
+    }
+
+    @Override
+    public boolean test(R record) {
+        return record.topic() != null && pattern.matcher(record.topic()).matches();
+    }
+
+    @Override
+    public void close() {
+
+    }
+
+    @Override
+    public void configure(Map configs) {
+        SimpleConfig simpleConfig = new SimpleConfig(config(), configs);
+        Pattern result;
+        String value = simpleConfig.getString(PATTERN_CONFIG);
+        try {
+            result = Pattern.compile(value);
+        } catch (PatternSyntaxException e) {
+            throw new ConfigException(PATTERN_CONFIG, value, "entry must be a Java-compatible regular expression: " + e.getMessage());
+        }
+        this.pattern = result;
+    }
+
+    @Override
+    public String toString() {
+        return "TopicNameMatches{" +
+                "pattern=" + pattern +
+                '}';
+    }
+}
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 b190189b35d9a..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,16 +17,26 @@
 
 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;
 import org.apache.kafka.connect.source.SourceRecord;
 import org.junit.After;
 import org.junit.Test;
 
+import java.math.BigDecimal;
 import java.util.Collections;
+import java.util.Date;
 import java.util.HashMap;
 import java.util.Map;
 
@@ -37,6 +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_HOUR = TimeUnit.HOURS.toMillis(1);
+    private static final long MILLIS_PER_DAY = TimeUnit.DAYS.toMillis(1);
 
     @After
     public void teardown() {
@@ -59,6 +71,11 @@ public void testConfigInvalidTargetType() {
         xformKey.configure(Collections.singletonMap(Cast.SPEC_CONFIG, "foo:array"));
     }
 
+    @Test(expected = ConfigException.class)
+    public void testUnsupportedTargetType() {
+        xformKey.configure(Collections.singletonMap(Cast.SPEC_CONFIG, "foo:bytes"));
+    }
+
     @Test(expected = ConfigException.class)
     public void testConfigInvalidMap() {
         xformKey.configure(Collections.singletonMap(Cast.SPEC_CONFIG, "foo:int8:extra"));
@@ -169,6 +186,28 @@ public void castWholeRecordValueWithSchemaString() {
         assertEquals("42", transformed.value());
     }
 
+    @Test
+    public void castWholeBigDecimalRecordValueWithSchemaString() {
+        BigDecimal bigDecimal = new BigDecimal(42);
+        xformValue.configure(Collections.singletonMap(Cast.SPEC_CONFIG, "string"));
+        SourceRecord transformed = xformValue.apply(new SourceRecord(null, null, "topic", 0,
+                Decimal.schema(bigDecimal.scale()), bigDecimal));
+
+        assertEquals(Schema.Type.STRING, transformed.valueSchema().type());
+        assertEquals("42", transformed.value());
+    }
+
+    @Test
+    public void castWholeDateRecordValueWithSchemaString() {
+        Date timestamp = new Date(MILLIS_PER_DAY + 1); // day + 1msec to get a timestamp formatting.
+        xformValue.configure(Collections.singletonMap(Cast.SPEC_CONFIG, "string"));
+        SourceRecord transformed = xformValue.apply(new SourceRecord(null, null, "topic", 0,
+                Timestamp.SCHEMA, timestamp));
+
+        assertEquals(Schema.Type.STRING, transformed.valueSchema().type());
+        assertEquals(Values.dateFormatFor(timestamp).format(timestamp), transformed.value());
+    }
+
     @Test
     public void castWholeRecordDefaultValue() {
         // Validate default value in schema is correctly converted
@@ -287,10 +326,102 @@ 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() {
-        xformValue.configure(Collections.singletonMap(Cast.SPEC_CONFIG, "int8:int16,int16:int32,int32:int64,int64:boolean,float32:float64,float64:boolean,boolean:int8,string:int32,optional:int32"));
+        Date day = new Date(MILLIS_PER_DAY);
+        xformValue.configure(Collections.singletonMap(Cast.SPEC_CONFIG, "int8:int16,int16:int32,int32:int64,int64:boolean,float32:float64,float64:boolean,boolean:int8,string:int32,bigdecimal:string,date:string,optional:int32"));
 
         // Include an optional fields and fields with defaults to validate their values are passed through properly
         SchemaBuilder builder = SchemaBuilder.struct();
@@ -303,7 +434,10 @@ public void castFieldsWithSchema() {
         builder.field("float64", SchemaBuilder.float64().defaultValue(-1.125).build());
         builder.field("boolean", Schema.BOOLEAN_SCHEMA);
         builder.field("string", Schema.STRING_SCHEMA);
+        builder.field("bigdecimal", Decimal.schema(new BigDecimal(42).scale()));
+        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();
 
         Struct recordValue = new Struct(supportedTypesSchema);
@@ -314,7 +448,10 @@ public void castFieldsWithSchema() {
         recordValue.put("float32", 32.f);
         recordValue.put("float64", -64.);
         recordValue.put("boolean", true);
+        recordValue.put("bigdecimal", new BigDecimal(42));
+        recordValue.put("date", day);
         recordValue.put("string", "42");
+        recordValue.put("timestamp", new Date(0));
         // optional field intentionally omitted
 
         SourceRecord transformed = xformValue.apply(new SourceRecord(null, null, "topic", 0,
@@ -331,7 +468,25 @@ public void castFieldsWithSchema() {
         assertEquals(true, ((Struct) transformed.value()).schema().field("float64").schema().defaultValue());
         assertEquals((byte) 1, ((Struct) transformed.value()).get("boolean"));
         assertEquals(42, ((Struct) transformed.value()).get("string"));
+        assertEquals("42", ((Struct) transformed.value()).get("bigdecimal"));
+        assertEquals(Values.dateFormatFor(day).format(day), ((Struct) transformed.value()).get("date"));
+        assertEquals(new Date(0), ((Struct) transformed.value()).get("timestamp"));
         assertNull(((Struct) transformed.value()).get("optional"));
+
+        Schema transformedSchema = ((Struct) transformed.value()).schema();
+        assertEquals(Schema.INT16_SCHEMA.type(), transformedSchema.field("int8").schema().type());
+        assertEquals(Schema.OPTIONAL_INT32_SCHEMA.type(), transformedSchema.field("int16").schema().type());
+        assertEquals(Schema.INT64_SCHEMA.type(), transformedSchema.field("int32").schema().type());
+        assertEquals(Schema.BOOLEAN_SCHEMA.type(), transformedSchema.field("int64").schema().type());
+        assertEquals(Schema.FLOAT64_SCHEMA.type(), transformedSchema.field("float32").schema().type());
+        assertEquals(Schema.BOOLEAN_SCHEMA.type(), transformedSchema.field("float64").schema().type());
+        assertEquals(Schema.INT8_SCHEMA.type(), transformedSchema.field("boolean").schema().type());
+        assertEquals(Schema.INT32_SCHEMA.type(), transformedSchema.field("string").schema().type());
+        assertEquals(Schema.STRING_SCHEMA.type(), transformedSchema.field("bigdecimal").schema().type());
+        assertEquals(Schema.STRING_SCHEMA.type(), transformedSchema.field("date").schema().type());
+        assertEquals(Schema.OPTIONAL_INT32_SCHEMA.type(), transformedSchema.field("optional").schema().type());
+        // The following fields are not changed
+        assertEquals(Timestamp.SCHEMA.type(), transformedSchema.field("timestamp").schema().type());
     }
 
     @SuppressWarnings("unchecked")
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 d709054dbe043..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
@@ -141,6 +141,7 @@ public void testNestedMapWithDelimiter() {
 
         assertNull(transformed.valueSchema());
         assertTrue(transformed.value() instanceof Map);
+        @SuppressWarnings("unchecked")
         Map transformedMap = (Map) transformed.value();
         assertEquals(9, transformedMap.size());
         assertEquals((byte) 8, transformedMap.get("A#B#int8"));
@@ -181,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());
@@ -196,6 +237,7 @@ public void testOptionalFieldMap() {
 
         assertNull(transformed.valueSchema());
         assertTrue(transformed.value() instanceof Map);
+        @SuppressWarnings("unchecked")
         Map transformedMap = (Map) transformed.value();
 
         assertNull(transformedMap.get("B.opt_int32"));
@@ -211,6 +253,7 @@ public void testKey() {
 
         assertNull(transformed.keySchema());
         assertTrue(transformed.key() instanceof Map);
+        @SuppressWarnings("unchecked")
         Map transformedMap = (Map) transformed.key();
         assertEquals(12, transformedMap.get("A.B"));
     }
@@ -254,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..b572c092fc879 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
@@ -26,6 +26,7 @@
 import org.junit.Test;
 
 import java.util.Collections;
+import java.util.Date;
 import java.util.HashMap;
 import java.util.Map;
 
@@ -33,17 +34,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 +57,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 record = new SourceRecord(null, null, "test", 0, null, null, simpleStructSchema, simpleStruct, 789L);
+        final SourceRecord transformedRecord = xformValue.apply(record);
 
         assertEquals(simpleStructSchema.name(), transformedRecord.valueSchema().name());
         assertEquals(simpleStructSchema.version(), transformedRecord.valueSchema().version());
@@ -77,13 +79,13 @@ public void copySchemaAndInsertConfiguredFields() {
         assertEquals(0, ((Struct) transformedRecord.value()).getInt32("partition_field").intValue());
 
         assertEquals(Timestamp.builder().optional().build(), transformedRecord.valueSchema().field("timestamp_field").schema());
-        assertEquals(null, ((Struct) transformedRecord.value()).getInt64("timestamp_field"));
+        assertEquals(789L, ((Date) ((Struct) transformedRecord.value()).get("timestamp_field")).getTime());
 
         assertEquals(Schema.OPTIONAL_STRING_SCHEMA, transformedRecord.valueSchema().field("instance_id").schema());
         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 +99,102 @@ 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));
+                null, null, null, Collections.singletonMap("magic", 42L), 123L);
 
-        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(123L, ((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/MaskFieldTest.java b/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/MaskFieldTest.java
index 9211a462103fc..a354ca9b7a3be 100644
--- a/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/MaskFieldTest.java
+++ b/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/MaskFieldTest.java
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.connect.transforms;
 
+import org.apache.kafka.common.config.ConfigException;
 import org.apache.kafka.connect.data.Decimal;
 import org.apache.kafka.connect.data.Field;
 import org.apache.kafka.connect.data.Schema;
@@ -23,6 +24,7 @@
 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.errors.DataException;
 import org.apache.kafka.connect.sink.SinkRecord;
 import org.junit.Test;
 
@@ -36,13 +38,71 @@
 import java.util.List;
 import java.util.Map;
 
+import static java.util.Collections.singletonList;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
 
 public class MaskFieldTest {
 
-    private static MaskField transform(List fields) {
+    private static final Schema SCHEMA = SchemaBuilder.struct()
+            .field("magic", Schema.INT32_SCHEMA)
+            .field("bool", Schema.BOOLEAN_SCHEMA)
+            .field("byte", Schema.INT8_SCHEMA)
+            .field("short", Schema.INT16_SCHEMA)
+            .field("int", Schema.INT32_SCHEMA)
+            .field("long", Schema.INT64_SCHEMA)
+            .field("float", Schema.FLOAT32_SCHEMA)
+            .field("double", Schema.FLOAT64_SCHEMA)
+            .field("string", Schema.STRING_SCHEMA)
+            .field("date", org.apache.kafka.connect.data.Date.SCHEMA)
+            .field("time", Time.SCHEMA)
+            .field("timestamp", Timestamp.SCHEMA)
+            .field("decimal", Decimal.schema(0))
+            .field("array", SchemaBuilder.array(Schema.INT32_SCHEMA))
+            .field("map", SchemaBuilder.map(Schema.STRING_SCHEMA, Schema.STRING_SCHEMA))
+            .build();
+    private static final Map VALUES = new HashMap<>();
+    private static final Struct VALUES_WITH_SCHEMA = new Struct(SCHEMA);
+
+    static {
+        VALUES.put("magic", 42);
+        VALUES.put("bool", true);
+        VALUES.put("byte", (byte) 42);
+        VALUES.put("short", (short) 42);
+        VALUES.put("int", 42);
+        VALUES.put("long", 42L);
+        VALUES.put("float", 42f);
+        VALUES.put("double", 42d);
+        VALUES.put("string", "55.121.20.20");
+        VALUES.put("date", new Date());
+        VALUES.put("bigint", new BigInteger("42"));
+        VALUES.put("bigdec", new BigDecimal("42.0"));
+        VALUES.put("list", singletonList(42));
+        VALUES.put("map", Collections.singletonMap("key", "value"));
+
+        VALUES_WITH_SCHEMA.put("magic", 42);
+        VALUES_WITH_SCHEMA.put("bool", true);
+        VALUES_WITH_SCHEMA.put("byte", (byte) 42);
+        VALUES_WITH_SCHEMA.put("short", (short) 42);
+        VALUES_WITH_SCHEMA.put("int", 42);
+        VALUES_WITH_SCHEMA.put("long", 42L);
+        VALUES_WITH_SCHEMA.put("float", 42f);
+        VALUES_WITH_SCHEMA.put("double", 42d);
+        VALUES_WITH_SCHEMA.put("string", "hmm");
+        VALUES_WITH_SCHEMA.put("date", new Date());
+        VALUES_WITH_SCHEMA.put("time", new Date());
+        VALUES_WITH_SCHEMA.put("timestamp", new Date());
+        VALUES_WITH_SCHEMA.put("decimal", new BigDecimal(42));
+        VALUES_WITH_SCHEMA.put("array", Arrays.asList(1, 2, 3));
+        VALUES_WITH_SCHEMA.put("map", Collections.singletonMap("what", "what"));
+    }
+
+    private static MaskField transform(List fields, String replacement) {
         final MaskField xform = new MaskField.Value<>();
-        xform.configure(Collections.singletonMap("fields", fields));
+        Map props = new HashMap<>();
+        props.put("fields", fields);
+        props.put("replacement", replacement);
+        xform.configure(props);
         return xform;
     }
 
@@ -50,28 +110,32 @@ private static SinkRecord record(Schema schema, Object value) {
         return new SinkRecord("", 0, null, null, schema, value, 0);
     }
 
+    private static void checkReplacementWithSchema(String maskField, Object replacement) {
+        SinkRecord record = record(SCHEMA, VALUES_WITH_SCHEMA);
+        final Struct updatedValue = (Struct) transform(singletonList(maskField), String.valueOf(replacement)).apply(record).value();
+        assertEquals("Invalid replacement for " + maskField + " value", replacement, updatedValue.get(maskField));
+    }
+
+    private static void checkReplacementSchemaless(String maskField, Object replacement) {
+        checkReplacementSchemaless(singletonList(maskField), replacement);
+    }
+
+    @SuppressWarnings("unchecked")
+    private static void checkReplacementSchemaless(List maskFields, Object replacement) {
+        SinkRecord record = record(null, VALUES);
+        final Map updatedValue = (Map) transform(maskFields, String.valueOf(replacement))
+                .apply(record)
+                .value();
+        for (String maskField : maskFields) {
+            assertEquals("Invalid replacement for " + maskField + " value", replacement, updatedValue.get(maskField));
+        }
+    }
+
     @Test
-    public void schemaless() {
-        final Map value = new HashMap<>();
-        value.put("magic", 42);
-        value.put("bool", true);
-        value.put("byte", (byte) 42);
-        value.put("short", (short) 42);
-        value.put("int", 42);
-        value.put("long", 42L);
-        value.put("float", 42f);
-        value.put("double", 42d);
-        value.put("string", "blabla");
-        value.put("date", new Date());
-        value.put("bigint", new BigInteger("42"));
-        value.put("bigdec", new BigDecimal("42.0"));
-        value.put("list", Collections.singletonList(42));
-        value.put("map", Collections.singletonMap("key", "value"));
-
-        final List maskFields = new ArrayList<>(value.keySet());
+    public void testSchemaless() {
+        final List maskFields = new ArrayList<>(VALUES.keySet());
         maskFields.remove("magic");
-
-        final Map updatedValue = (Map) transform(maskFields).apply(record(null, value)).value();
+        @SuppressWarnings("unchecked") final Map updatedValue = (Map) transform(maskFields, null).apply(record(null, VALUES)).value();
 
         assertEquals(42, updatedValue.get("magic"));
         assertEquals(false, updatedValue.get("bool"));
@@ -90,50 +154,15 @@ public void schemaless() {
     }
 
     @Test
-    public void withSchema() {
-        Schema schema = SchemaBuilder.struct()
-                .field("magic", Schema.INT32_SCHEMA)
-                .field("bool", Schema.BOOLEAN_SCHEMA)
-                .field("byte", Schema.INT8_SCHEMA)
-                .field("short", Schema.INT16_SCHEMA)
-                .field("int", Schema.INT32_SCHEMA)
-                .field("long", Schema.INT64_SCHEMA)
-                .field("float", Schema.FLOAT32_SCHEMA)
-                .field("double", Schema.FLOAT64_SCHEMA)
-                .field("string", Schema.STRING_SCHEMA)
-                .field("date", org.apache.kafka.connect.data.Date.SCHEMA)
-                .field("time", Time.SCHEMA)
-                .field("timestamp", Timestamp.SCHEMA)
-                .field("decimal", Decimal.schema(0))
-                .field("array", SchemaBuilder.array(Schema.INT32_SCHEMA))
-                .field("map", SchemaBuilder.map(Schema.STRING_SCHEMA, Schema.STRING_SCHEMA))
-                .build();
-
-        final Struct value = new Struct(schema);
-        value.put("magic", 42);
-        value.put("bool", true);
-        value.put("byte", (byte) 42);
-        value.put("short", (short) 42);
-        value.put("int", 42);
-        value.put("long", 42L);
-        value.put("float", 42f);
-        value.put("double", 42d);
-        value.put("string", "hmm");
-        value.put("date", new Date());
-        value.put("time", new Date());
-        value.put("timestamp", new Date());
-        value.put("decimal", new BigDecimal(42));
-        value.put("array", Arrays.asList(1, 2, 3));
-        value.put("map", Collections.singletonMap("what", "what"));
-
-        final List maskFields = new ArrayList<>(schema.fields().size());
-        for (Field field: schema.fields()) {
+    public void testWithSchema() {
+        final List maskFields = new ArrayList<>(SCHEMA.fields().size());
+        for (Field field : SCHEMA.fields()) {
             if (!field.name().equals("magic")) {
                 maskFields.add(field.name());
             }
         }
 
-        final Struct updatedValue = (Struct) transform(maskFields).apply(record(schema, value)).value();
+        final Struct updatedValue = (Struct) transform(maskFields, null).apply(record(SCHEMA, VALUES_WITH_SCHEMA)).value();
 
         assertEquals(42, updatedValue.get("magic"));
         assertEquals(false, updatedValue.get("bool"));
@@ -152,4 +181,73 @@ public void withSchema() {
         assertEquals(Collections.emptyMap(), updatedValue.get("map"));
     }
 
+    @Test
+    public void testSchemalessWithReplacement() {
+        checkReplacementSchemaless("short", (short) 123);
+        checkReplacementSchemaless("byte", (byte) 123);
+        checkReplacementSchemaless("int", 123);
+        checkReplacementSchemaless("long", 123L);
+        checkReplacementSchemaless("float", 123.0f);
+        checkReplacementSchemaless("double", 123.0);
+        checkReplacementSchemaless("string", "123");
+        checkReplacementSchemaless("bigint", BigInteger.valueOf(123L));
+        checkReplacementSchemaless("bigdec", BigDecimal.valueOf(123.0));
+    }
+
+    @Test
+    public void testSchemalessUnsupportedReplacementType() {
+        String exMessage = "Cannot mask value of type";
+        Class exClass = DataException.class;
+
+        assertThrows(exMessage, exClass, () -> checkReplacementSchemaless("date", new Date()));
+        assertThrows(exMessage, exClass, () -> checkReplacementSchemaless(Arrays.asList("int", "date"), new Date()));
+        assertThrows(exMessage, exClass, () -> checkReplacementSchemaless("bool", false));
+        assertThrows(exMessage, exClass, () -> checkReplacementSchemaless("list", singletonList("123")));
+        assertThrows(exMessage, exClass, () -> checkReplacementSchemaless("map", Collections.singletonMap("123", "321")));
+    }
+
+    @Test
+    public void testWithSchemaAndReplacement() {
+        checkReplacementWithSchema("short", (short) 123);
+        checkReplacementWithSchema("byte", (byte) 123);
+        checkReplacementWithSchema("int", 123);
+        checkReplacementWithSchema("long", 123L);
+        checkReplacementWithSchema("float", 123.0f);
+        checkReplacementWithSchema("double", 123.0);
+        checkReplacementWithSchema("string", "123");
+        checkReplacementWithSchema("decimal", BigDecimal.valueOf(123.0));
+    }
+
+    @Test
+    public void testWithSchemaUnsupportedReplacementType() {
+        String exMessage = "Cannot mask value of type";
+        Class exClass = DataException.class;
+
+        assertThrows(exMessage, exClass, () -> checkReplacementWithSchema("time", new Date()));
+        assertThrows(exMessage, exClass, () -> checkReplacementWithSchema("timestamp", new Date()));
+        assertThrows(exMessage, exClass, () -> checkReplacementWithSchema("array", singletonList(123)));
+    }
+
+    @Test
+    public void testReplacementTypeMismatch() {
+        String exMessage = "Invalid value  for configuration replacement";
+        Class exClass = DataException.class;
+
+        assertThrows(exMessage, exClass, () -> checkReplacementSchemaless("byte", "foo"));
+        assertThrows(exMessage, exClass, () -> checkReplacementSchemaless("short", "foo"));
+        assertThrows(exMessage, exClass, () -> checkReplacementSchemaless("int", "foo"));
+        assertThrows(exMessage, exClass, () -> checkReplacementSchemaless("long", "foo"));
+        assertThrows(exMessage, exClass, () -> checkReplacementSchemaless("float", "foo"));
+        assertThrows(exMessage, exClass, () -> checkReplacementSchemaless("double", "foo"));
+        assertThrows(exMessage, exClass, () -> checkReplacementSchemaless("bigint", "foo"));
+        assertThrows(exMessage, exClass, () -> checkReplacementSchemaless("bigdec", "foo"));
+        assertThrows(exMessage, exClass, () -> checkReplacementSchemaless("int", new Date()));
+        assertThrows(exMessage, exClass, () -> checkReplacementSchemaless("int", new Object()));
+        assertThrows(exMessage, exClass, () -> checkReplacementSchemaless(Arrays.asList("string", "int"), "foo"));
+    }
+
+    @Test
+    public void testEmptyStringReplacementValue() {
+        assertThrows("String must be non-empty", ConfigException.class, () -> checkReplacementSchemaless("short", ""));
+    }
 }
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..c8aac95f1ac6f 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,10 +37,47 @@ public void teardown() {
         xform.close();
     }
 
+    @Test
+    public void tombstoneSchemaless() {
+        final Map props = new HashMap<>();
+        props.put("include", "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("include", "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<>();
-        props.put("blacklist", "dont");
+        props.put("exclude", "dont");
         props.put("renames", "abc:xyz,foo:bar");
 
         xform.configure(props);
@@ -63,7 +101,7 @@ public void schemaless() {
     @Test
     public void withSchema() {
         final Map props = new HashMap<>();
-        props.put("whitelist", "abc,foo");
+        props.put("include", "abc,foo");
         props.put("renames", "abc:xyz,foo:bar");
 
         xform.configure(props);
@@ -87,8 +125,47 @@ public void withSchema() {
         final Struct updatedValue = (Struct) transformedRecord.value();
 
         assertEquals(2, updatedValue.schema().fields().size());
-        assertEquals(new Integer(42), updatedValue.getInt32("xyz"));
+        assertEquals(Integer.valueOf(42), updatedValue.getInt32("xyz"));
         assertEquals(true, updatedValue.getBoolean("bar"));
     }
 
+    @Test
+    public void testIncludeBackwardsCompatibility() {
+        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 testExcludeBackwardsCompatibility() {
+        final Map props = new HashMap<>();
+        props.put("blacklist", "dont");
+        props.put("renames", "abc:xyz,foo:bar");
+
+        xform.configure(props);
+
+        final Map value = new HashMap<>();
+        value.put("dont", "whatever");
+        value.put("abc", 42);
+        value.put("foo", true);
+        value.put("etc", "etc");
+
+        final SinkRecord record = new SinkRecord("test", 0, null, null, null, value, 0);
+        final SinkRecord transformedRecord = xform.apply(record);
+
+        final Map updatedValue = (Map) transformedRecord.value();
+        assertEquals(3, updatedValue.size());
+        assertEquals(42, updatedValue.get("xyz"));
+        assertEquals(true, updatedValue.get("bar"));
+        assertEquals("etc", updatedValue.get("etc"));
+    }
 }
diff --git a/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/SetSchemaMetadataTest.java b/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/SetSchemaMetadataTest.java
index 257b382b9a07d..3a05f1f802a17 100644
--- a/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/SetSchemaMetadataTest.java
+++ b/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/SetSchemaMetadataTest.java
@@ -52,7 +52,7 @@ public void schemaVersionUpdate() {
         xform.configure(Collections.singletonMap("schema.version", 42));
         final SinkRecord record = new SinkRecord("", 0, null, null, SchemaBuilder.struct().build(), null, 0);
         final SinkRecord updatedRecord = xform.apply(record);
-        assertEquals(new Integer(42), updatedRecord.valueSchema().version());
+        assertEquals(Integer.valueOf(42), updatedRecord.valueSchema().version());
     }
 
     @Test
@@ -68,7 +68,7 @@ public void schemaNameAndVersionUpdate() {
         final SinkRecord updatedRecord = xform.apply(record);
 
         assertEquals("foo", updatedRecord.valueSchema().name());
-        assertEquals(new Integer(42), updatedRecord.valueSchema().version());
+        assertEquals(Integer.valueOf(42), updatedRecord.valueSchema().version());
     }
 
     @Test
@@ -94,7 +94,7 @@ public void schemaNameAndVersionUpdateWithStruct() {
         final SinkRecord updatedRecord = xform.apply(record);
 
         assertEquals("foo", updatedRecord.valueSchema().name());
-        assertEquals(new Integer(42), updatedRecord.valueSchema().version());
+        assertEquals(Integer.valueOf(42), updatedRecord.valueSchema().version());
 
         // Make sure the struct's schema and fields all point to the new schema
         assertMatchingSchema((Struct) updatedRecord.value(), updatedRecord.valueSchema());
@@ -125,7 +125,7 @@ public void updateSchemaOfStruct() {
 
     @Test
     public void updateSchemaOfNonStruct() {
-        Object value = new Integer(1);
+        Object value = Integer.valueOf(1);
         Object updatedValue = SetSchemaMetadata.updateSchemaIn(value, Schema.INT32_SCHEMA);
         assertSame(value, updatedValue);
     }
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/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/predicates/HasHeaderKeyTest.java b/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/predicates/HasHeaderKeyTest.java
new file mode 100644
index 0000000000000..f6d1d330c5320
--- /dev/null
+++ b/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/predicates/HasHeaderKeyTest.java
@@ -0,0 +1,123 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.transforms.predicates;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+import org.apache.kafka.common.config.ConfigException;
+import org.apache.kafka.common.config.ConfigValue;
+import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.header.Header;
+import org.apache.kafka.connect.source.SourceRecord;
+import org.apache.kafka.connect.transforms.util.SimpleConfig;
+import org.junit.Test;
+
+import static java.util.Collections.singletonList;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+
+public class HasHeaderKeyTest {
+
+    @Test
+    public void testNameRequiredInConfig() {
+        Map props = new HashMap<>();
+        ConfigException e = assertThrows(ConfigException.class, () -> config(props));
+        assertTrue(e.getMessage().contains("Missing required configuration \"name\""));
+    }
+
+    @Test
+    public void testNameMayNotBeEmptyInConfig() {
+        Map props = new HashMap<>();
+        props.put("name", "");
+        ConfigException e = assertThrows(ConfigException.class, () -> config(props));
+        assertTrue(e.getMessage().contains("String must be non-empty"));
+    }
+
+    @Test
+    public void testConfig() {
+        HasHeaderKey predicate = new HasHeaderKey<>();
+        predicate.config().validate(Collections.singletonMap("name", "foo"));
+
+        List configs = predicate.config().validate(Collections.singletonMap("name", ""));
+        assertEquals(singletonList("Invalid value  for configuration name: String must be non-empty"), configs.get(0).errorMessages());
+    }
+
+    @Test
+    public void testTest() {
+        HasHeaderKey predicate = new HasHeaderKey<>();
+        predicate.configure(Collections.singletonMap("name", "foo"));
+
+        assertTrue(predicate.test(recordWithHeaders("foo")));
+        assertTrue(predicate.test(recordWithHeaders("foo", "bar")));
+        assertTrue(predicate.test(recordWithHeaders("bar", "foo", "bar", "foo")));
+        assertFalse(predicate.test(recordWithHeaders("bar")));
+        assertFalse(predicate.test(recordWithHeaders("bar", "bar")));
+        assertFalse(predicate.test(recordWithHeaders()));
+        assertFalse(predicate.test(new SourceRecord(null, null, null, null, null)));
+
+    }
+
+    private SimpleConfig config(Map props) {
+        return new SimpleConfig(new HasHeaderKey().config(), props);
+    }
+
+    private SourceRecord recordWithHeaders(String... headers) {
+        return new SourceRecord(null, null, null, null, null, null, null, null, null,
+                Arrays.stream(headers).map(header -> new TestHeader(header)).collect(Collectors.toList()));
+    }
+
+    private static class TestHeader implements Header {
+
+        private final String key;
+
+        public TestHeader(String key) {
+            this.key = key;
+        }
+
+        @Override
+        public String key() {
+            return key;
+        }
+
+        @Override
+        public Schema schema() {
+            return null;
+        }
+
+        @Override
+        public Object value() {
+            return null;
+        }
+
+        @Override
+        public Header with(Schema schema, Object value) {
+            return null;
+        }
+
+        @Override
+        public Header rename(String key) {
+            return null;
+        }
+    }
+}
diff --git a/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/predicates/TopicNameMatchesTest.java b/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/predicates/TopicNameMatchesTest.java
new file mode 100644
index 0000000000000..b0cc34bc36213
--- /dev/null
+++ b/connect/transforms/src/test/java/org/apache/kafka/connect/transforms/predicates/TopicNameMatchesTest.java
@@ -0,0 +1,93 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF 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.transforms.predicates;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import org.apache.kafka.common.config.ConfigException;
+import org.apache.kafka.common.config.ConfigValue;
+import org.apache.kafka.connect.source.SourceRecord;
+import org.apache.kafka.connect.transforms.util.SimpleConfig;
+import org.junit.Test;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+
+public class TopicNameMatchesTest {
+
+    @Test
+    public void testPatternRequiredInConfig() {
+        Map props = new HashMap<>();
+        ConfigException e = assertThrows(ConfigException.class, () -> config(props));
+        assertTrue(e.getMessage().contains("Missing required configuration \"pattern\""));
+    }
+
+    @Test
+    public void testPatternMayNotBeEmptyInConfig() {
+        Map props = new HashMap<>();
+        props.put("pattern", "");
+        ConfigException e = assertThrows(ConfigException.class, () -> config(props));
+        assertTrue(e.getMessage().contains("String must be non-empty"));
+    }
+
+    @Test
+    public void testPatternIsValidRegexInConfig() {
+        Map props = new HashMap<>();
+        props.put("pattern", "[");
+        ConfigException e = assertThrows(ConfigException.class, () -> config(props));
+        assertTrue(e.getMessage().contains("Invalid regex"));
+    }
+
+    @Test
+    public void testConfig() {
+        TopicNameMatches predicate = new TopicNameMatches<>();
+        predicate.config().validate(Collections.singletonMap("pattern", "my-prefix-.*"));
+
+        List configs = predicate.config().validate(Collections.singletonMap("pattern", "*"));
+        List errorMsgs = configs.get(0).errorMessages();
+        assertEquals(1, errorMsgs.size());
+        assertTrue(errorMsgs.get(0).contains("Invalid regex"));
+    }
+
+    @Test
+    public void testTest() {
+        TopicNameMatches predicate = new TopicNameMatches<>();
+        predicate.configure(Collections.singletonMap("pattern", "my-prefix-.*"));
+
+        assertTrue(predicate.test(recordWithTopicName("my-prefix-")));
+        assertTrue(predicate.test(recordWithTopicName("my-prefix-foo")));
+        assertFalse(predicate.test(recordWithTopicName("x-my-prefix-")));
+        assertFalse(predicate.test(recordWithTopicName("x-my-prefix-foo")));
+        assertFalse(predicate.test(recordWithTopicName("your-prefix-")));
+        assertFalse(predicate.test(recordWithTopicName("your-prefix-foo")));
+        assertFalse(predicate.test(new SourceRecord(null, null, null, null, null)));
+
+    }
+
+    private SimpleConfig config(Map props) {
+        return new SimpleConfig(TopicNameMatches.CONFIG_DEF, props);
+    }
+
+    private SourceRecord recordWithTopicName(String topicName) {
+        return new SourceRecord(null, null, topicName, null, null);
+    }
+}
diff --git a/core/src/main/java/kafka/metrics/FilteringJmxReporter.java b/core/src/main/java/kafka/metrics/FilteringJmxReporter.java
new file mode 100644
index 0000000000000..be2ba14b3df24
--- /dev/null
+++ b/core/src/main/java/kafka/metrics/FilteringJmxReporter.java
@@ -0,0 +1,62 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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.metrics;
+
+import com.yammer.metrics.core.Metric;
+import com.yammer.metrics.core.MetricName;
+import com.yammer.metrics.core.MetricsRegistry;
+import com.yammer.metrics.reporting.JmxReporter;
+
+import java.util.function.Predicate;
+
+public class FilteringJmxReporter extends JmxReporter {
+
+    private volatile Predicate metricPredicate;
+
+    public FilteringJmxReporter(MetricsRegistry registry, Predicate metricPredicate) {
+        super(registry);
+        this.metricPredicate = metricPredicate;
+    }
+
+    @Override
+    public void onMetricAdded(MetricName name, Metric metric) {
+        if (metricPredicate.test(name)) {
+            super.onMetricAdded(name, metric);
+        }
+    }
+
+    @Override
+    public void onMetricRemoved(MetricName name) {
+        super.onMetricRemoved(name);
+    }
+
+    public void updatePredicate(Predicate predicate) {
+        this.metricPredicate = predicate;
+        // re-register metrics on update
+        getMetricsRegistry()
+            .allMetrics()
+            .forEach((name, metric) -> {
+                if (metricPredicate.test(name)) {
+                    super.onMetricAdded(name, metric);
+                } else {
+                    super.onMetricRemoved(name);
+                }
+            }
+            );
+    }
+}
diff --git a/core/src/main/java/kafka/metrics/KafkaYammerMetrics.java b/core/src/main/java/kafka/metrics/KafkaYammerMetrics.java
new file mode 100644
index 0000000000000..dd650fdd0f79e
--- /dev/null
+++ b/core/src/main/java/kafka/metrics/KafkaYammerMetrics.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 kafka.metrics;
+
+import com.yammer.metrics.core.MetricsRegistry;
+
+import org.apache.kafka.common.Reconfigurable;
+import org.apache.kafka.common.config.ConfigException;
+import org.apache.kafka.common.metrics.JmxReporter;
+
+import java.util.Map;
+import java.util.Set;
+import java.util.function.Predicate;
+
+/**
+ * This class encapsulates the default yammer metrics registry for Kafka server,
+ * and configures the set of exported JMX metrics for Yammer metrics.
+ *
+ * KafkaYammerMetrics.defaultRegistry() should always be used instead of Metrics.defaultRegistry()
+ */
+public class KafkaYammerMetrics implements Reconfigurable {
+
+    public static final KafkaYammerMetrics INSTANCE = new KafkaYammerMetrics();
+
+    /**
+     * convenience method to replace {@link com.yammer.metrics.Metrics#defaultRegistry()}
+     */
+    public static MetricsRegistry defaultRegistry() {
+        return INSTANCE.metricsRegistry;
+    }
+
+    private final MetricsRegistry metricsRegistry = new MetricsRegistry();
+    private final FilteringJmxReporter jmxReporter = new FilteringJmxReporter(metricsRegistry,
+        metricName -> true);
+
+    private KafkaYammerMetrics() {
+        jmxReporter.start();
+        Runtime.getRuntime().addShutdownHook(new Thread(jmxReporter::shutdown));
+    }
+
+    @Override
+    public void configure(Map configs) {
+        reconfigure(configs);
+    }
+
+    @Override
+    public Set reconfigurableConfigs() {
+        return JmxReporter.RECONFIGURABLE_CONFIGS;
+    }
+
+    @Override
+    public void validateReconfiguration(Map configs) throws ConfigException {
+        JmxReporter.compilePredicate(configs);
+    }
+
+    @Override
+    public void reconfigure(Map configs) {
+        Predicate mBeanPredicate = JmxReporter.compilePredicate(configs);
+        jmxReporter.updatePredicate(metricName -> mBeanPredicate.test(metricName.getMBeanName()));
+    }
+}
diff --git a/core/src/main/resources/common/message/GroupMetadataKey.json b/core/src/main/resources/common/message/GroupMetadataKey.json
new file mode 100644
index 0000000000000..2c39e68846652
--- /dev/null
+++ b/core/src/main/resources/common/message/GroupMetadataKey.json
@@ -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.
+
+{
+  "type": "data",
+  "name": "GroupMetadataKey",
+  "validVersions": "2",
+  "fields": [
+    { "name": "group", "type": "string", "versions": "2" }
+  ]
+}
diff --git a/core/src/main/resources/common/message/GroupMetadataValue.json b/core/src/main/resources/common/message/GroupMetadataValue.json
new file mode 100644
index 0000000000000..4d0f34d084c91
--- /dev/null
+++ b/core/src/main/resources/common/message/GroupMetadataValue.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.
+
+{
+  "type": "data",
+  "name": "GroupMetadataValue",
+  "validVersions": "0-3",
+  "fields": [
+    { "name": "protocolType", "versions": "0+", "type": "string"},
+    { "name": "generation", "versions": "0+", "type": "int32" },
+    { "name": "protocol", "versions": "0+", "type": "string", "nullableVersions": "0+" },
+    { "name": "leader", "versions": "0+", "type": "string", "nullableVersions": "0+" },
+    { "name": "currentStateTimestamp", "versions": "2+", "type": "int64", "default": -1, "ignorable": true},
+    { "name": "members", "versions": "0+", "type": "[]MemberMetadata" }
+  ],
+  "commonStructs": [
+    {
+      "name": "MemberMetadata",
+      "versions": "0-3",
+      "fields": [
+        { "name": "memberId", "versions": "0+", "type": "string" },
+        { "name": "groupInstanceId", "versions": "3+", "type": "string", "default": "null", "nullableVersions": "3+", "ignorable": true},
+        { "name": "clientId", "versions": "0+", "type": "string" },
+        { "name": "clientHost", "versions": "0+", "type": "string" },
+        { "name": "rebalanceTimeout", "versions": "1+", "type": "int32", "ignorable": true},
+        { "name": "sessionTimeout", "versions": "0+", "type": "int32" },
+        { "name": "subscription", "versions": "0+", "type": "bytes" },
+        { "name": "assignment", "versions": "0+", "type": "bytes" }
+      ]
+    }
+  ]
+}
diff --git a/core/src/main/resources/common/message/OffsetCommitKey.json b/core/src/main/resources/common/message/OffsetCommitKey.json
new file mode 100644
index 0000000000000..1ddd20a2bbda8
--- /dev/null
+++ b/core/src/main/resources/common/message/OffsetCommitKey.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.
+
+{
+  "type": "data",
+  "name": "OffsetCommitKey",
+  "validVersions": "0-1",
+  "fields": [
+    { "name": "group", "type": "string", "versions": "0-1" },
+    { "name": "topic", "type": "string", "versions": "0-1" },
+    { "name": "partition", "type": "int32", "versions": "0-1" }
+  ]
+}
diff --git a/core/src/main/resources/common/message/OffsetCommitValue.json b/core/src/main/resources/common/message/OffsetCommitValue.json
new file mode 100644
index 0000000000000..df524d014ec5e
--- /dev/null
+++ b/core/src/main/resources/common/message/OffsetCommitValue.json
@@ -0,0 +1,27 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements.  See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+{
+  "type": "data",
+  "name": "OffsetCommitValue",
+  "validVersions": "0-3",
+  "fields": [
+    { "name": "offset", "type": "int64", "versions": "0+" },
+    { "name": "leaderEpoch", "type": "int32", "versions": "3+", "default": -1, "ignorable": true},
+    { "name": "metadata", "type": "string", "versions": "0+" },
+    { "name": "commitTimestamp", "type": "int64", "versions": "0+" },
+    { "name": "expireTimestamp", "type": "int64", "versions": "1", "default": -1, "ignorable": true}
+  ]
+}
diff --git a/core/src/main/resources/common/message/TransactionLogKey.json b/core/src/main/resources/common/message/TransactionLogKey.json
new file mode 100644
index 0000000000000..f6fdbfbf9d625
--- /dev/null
+++ b/core/src/main/resources/common/message/TransactionLogKey.json
@@ -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.
+
+{
+  "type": "data",
+  "name": "TransactionLogKey",
+  "validVersions": "0",
+  "fields": [
+    { "name": "TransactionalId", "type": "string", "versions": "0"}
+  ]
+}
diff --git a/core/src/main/resources/common/message/TransactionLogValue.json b/core/src/main/resources/common/message/TransactionLogValue.json
new file mode 100644
index 0000000000000..153dc31253407
--- /dev/null
+++ b/core/src/main/resources/common/message/TransactionLogValue.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.
+
+{
+  "type": "data",
+  "name": "TransactionLogValue",
+  "validVersions": "0",
+  "fields": [
+    { "name": "ProducerId", "type": "int64", "versions": "0",
+      "about": "Producer id in use by the transactional id"},
+    { "name": "ProducerEpoch", "type": "int16", "versions": "0",
+      "about": "Epoch associated with the producer id"},
+    { "name": "TransactionTimeoutMs", "type": "int32", "versions": "0",
+      "about": "Transaction timeout in milliseconds"},
+    { "name": "TransactionStatus", "type": "int8", "versions": "0",
+      "about": "TransactionState the transaction is in"},
+    { "name": "TransactionPartitions", "type": "[]PartitionsSchema", "versions": "0", "nullableVersions": "0",
+      "about": "Set of partitions involved in the transaction", "fields": [
+      { "name": "Topic", "type": "string", "versions": "0"},
+      { "name": "PartitionIds", "type": "[]int32", "versions": "0"}]},
+    { "name": "TransactionLastUpdateTimestampMs", "type": "int64", "versions": "0",
+      "about": "Time the transaction was last updated"},
+    { "name": "TransactionStartTimestampMs", "type": "int64", "versions": "0",
+      "about": "Time the transaction was started"}
+  ]
+}
diff --git a/core/src/main/scala/kafka/Kafka.scala b/core/src/main/scala/kafka/Kafka.scala
index 25a72167d91d5..f01df66822194 100755
--- a/core/src/main/scala/kafka/Kafka.scala
+++ b/core/src/main/scala/kafka/Kafka.scala
@@ -18,16 +18,14 @@
 package kafka
 
 import java.util.Properties
-import java.util.concurrent.ConcurrentHashMap
 
-import sun.misc.{Signal, SignalHandler}
 import joptsimple.OptionParser
+import kafka.server.KafkaServerStartable
 import kafka.utils.Implicits._
-import kafka.server.{KafkaServer, KafkaServerStartable}
 import kafka.utils.{CommandLineUtils, Exit, Logging}
-import org.apache.kafka.common.utils.{OperatingSystem, Utils}
+import org.apache.kafka.common.utils.{Java, LoggingSignalHandler, OperatingSystem, Utils}
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 
 object Kafka extends Logging {
 
@@ -36,9 +34,19 @@ 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 || args.contains("--help")) {
+      CommandLineUtils.printUsageAndDie(optionParser,
+        "USAGE: java [options] %s server.properties [--override property=value]*".format(this.getClass.getCanonicalName.split('$').head))
+    }
 
-    if (args.length == 0) {
-      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))
@@ -55,39 +63,22 @@ object Kafka extends Logging {
     props
   }
 
-  private def registerLoggingSignalHandler(): Unit = {
-    val jvmSignalHandlers = new ConcurrentHashMap[String, SignalHandler]().asScala
-    val handler = new SignalHandler() {
-      override def handle(signal: Signal) {
-        info(s"Terminating process due to signal $signal")
-        jvmSignalHandlers.get(signal.getName).foreach(_.handle(signal))
-      }
-    }
-    def registerHandler(signalName: String) {
-      val oldHandler = Signal.handle(new Signal(signalName), handler)
-      if (oldHandler != null)
-        jvmSignalHandlers.put(signalName, oldHandler)
-    }
-
-    if (!OperatingSystem.IS_WINDOWS) {
-      registerHandler("TERM")
-      registerHandler("INT")
-      registerHandler("HUP")
-    }
-  }
-
   def main(args: Array[String]): Unit = {
     try {
       val serverProps = getPropsFromArgs(args)
       val kafkaServerStartable = KafkaServerStartable.fromProps(serverProps)
 
-      // register signal handler to log termination due to SIGTERM, SIGHUP and SIGINT (control-c)
-      registerLoggingSignalHandler()
+      try {
+        if (!OperatingSystem.IS_WINDOWS && !Java.isIbmJdk)
+          new LoggingSignalHandler().register()
+      } catch {
+        case e: ReflectiveOperationException =>
+          warn("Failed to register optional signal handler that logs a message when the process is terminated " +
+            s"by a signal. Reason for registration failure is: $e", e)
+      }
 
       // attach shutdown handler to catch terminating signals as well as normal termination
-      Runtime.getRuntime().addShutdownHook(new Thread("kafka-shutdown-hook") {
-        override def run(): Unit = kafkaServerStartable.shutdown()
-      })
+      Exit.addShutdownHook("kafka-shutdown-hook", kafkaServerStartable.shutdown())
 
       kafkaServerStartable.startup()
       kafkaServerStartable.awaitShutdown()
diff --git a/core/src/main/scala/kafka/admin/AclCommand.scala b/core/src/main/scala/kafka/admin/AclCommand.scala
index 2732f6c41a62e..5ca22ed297be8 100644
--- a/core/src/main/scala/kafka/admin/AclCommand.scala
+++ b/core/src/main/scala/kafka/admin/AclCommand.scala
@@ -17,42 +17,61 @@
 
 package kafka.admin
 
+import java.util.Properties
+
 import joptsimple._
-import kafka.security.auth._
+import joptsimple.util.EnumConverter
+import kafka.security.authorizer.{AclAuthorizer, AclEntry, AuthorizerUtils}
 import kafka.server.KafkaConfig
 import kafka.utils._
+import org.apache.kafka.clients.admin.{Admin, AdminClientConfig}
+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.Utils
+import org.apache.kafka.common.utils.{Utils, SecurityUtils => JSecurityUtils}
+import org.apache.kafka.server.authorizer.Authorizer
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
+import scala.collection.mutable
+import scala.io.StdIn
 
 object AclCommand extends Logging {
 
-  val Newline = scala.util.Properties.lineSeparator
-  val ResourceTypeToValidOperations = Map[ResourceType, Set[Operation]] (
-    Topic -> Set(Read, Write, Describe, Delete, DescribeConfigs, AlterConfigs, All),
-    Group -> Set(Read, Describe, All),
-    Cluster -> Set(Create, ClusterAction, DescribeConfigs, AlterConfigs, IdempotentWrite, Alter, Describe, All),
-    TransactionalId -> Set(Describe, Write, All)
-  )
+  val ClusterResourceFilter = new ResourcePatternFilter(JResourceType.CLUSTER, JResource.CLUSTER_NAME, PatternType.LITERAL)
+
+  private val Newline = scala.util.Properties.lineSeparator
 
-  def main(args: Array[String]) {
+  def main(args: Array[String]): Unit = {
 
     val opts = new AclCommandOptions(args)
 
-    if (opts.options.has(opts.helpOpt))
-      CommandLineUtils.printUsageAndDie(opts.parser, "Usage:")
+    CommandLineUtils.printHelpAndExitIfNeeded(opts, "This tool helps to manage acls on kafka.")
 
     opts.checkArgs()
 
+    val aclCommandService = {
+      if (opts.options.has(opts.bootstrapServerOpt)) {
+        new AdminClientService(opts)
+      } else {
+        val authorizerClassName = if (opts.options.has(opts.authorizerOpt))
+          opts.options.valueOf(opts.authorizerOpt)
+        else
+          classOf[AclAuthorizer].getName
+
+        new AuthorizerService(authorizerClassName, opts)
+      }
+    }
+
     try {
       if (opts.options.has(opts.addOpt))
-        addAcl(opts)
+        aclCommandService.addAcls()
       else if (opts.options.has(opts.removeOpt))
-        removeAcl(opts)
+        aclCommandService.removeAcls()
       else if (opts.options.has(opts.listOpt))
-        listAcl(opts)
+        aclCommandService.listAcls()
     } catch {
       case e: Throwable =>
         println(s"Error while executing ACL command: ${e.getMessage}")
@@ -61,127 +80,320 @@ object AclCommand extends Logging {
     }
   }
 
-  def withAuthorizer(opts: AclCommandOptions)(f: Authorizer => 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
-      }
+  sealed trait AclCommandService {
+    def addAcls(): Unit
+    def removeAcls(): Unit
+    def listAcls(): Unit
+  }
 
-    val authorizerClass = opts.options.valueOf(opts.authorizerOpt)
-    val authZ = CoreUtils.createObject[Authorizer](authorizerClass)
-    try {
-      authZ.configure(authorizerProperties.asJava)
-      f(authZ)
+  class AdminClientService(val opts: AclCommandOptions) extends AclCommandService with Logging {
+
+    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
+        new Properties()
+      props.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, opts.options.valueOf(opts.bootstrapServerOpt))
+      val adminClient = Admin.create(props)
+
+      try {
+        f(adminClient)
+      } finally {
+        adminClient.close()
+      }
     }
-    finally CoreUtils.swallow(authZ.close(), this)
-  }
 
-  private def addAcl(opts: AclCommandOptions) {
-    withAuthorizer(opts) { authorizer =>
+    def addAcls(): Unit = {
       val resourceToAcl = getResourceToAcls(opts)
+      withAdminClient(opts) { adminClient =>
+        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)).asJavaCollection
+          adminClient.createAcls(aclBindings).all().get()
+        }
+
+        listAcls(adminClient)
+      }
+    }
+
+    def removeAcls(): Unit = {
+      withAdminClient(opts) { adminClient =>
+        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(adminClient, 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(adminClient, acls, filter)
+          }
+        }
+
+        listAcls(adminClient)
+      }
+    }
+
+    def listAcls(): Unit = {
+      withAdminClient(opts) { adminClient =>
+        listAcls(adminClient)
+      }
+    }
 
-      if (resourceToAcl.values.exists(_.isEmpty))
-        CommandLineUtils.printUsageAndDie(opts.parser, "You must specify one of: --allow-principal, --deny-principal when trying to add ACLs.")
+    private def listAcls(adminClient: Admin): Unit = {
+      val filters = getResourceFilter(opts, dieIfNoResourceFound = false)
+      val listPrincipals = getPrincipals(opts, opts.listPrincipalsOpt)
+      val resourceToAcls = getAcls(adminClient, filters)
 
-      for ((resource, acls) <- resourceToAcl) {
-        println(s"Adding ACLs for resource `$resource`: $Newline ${acls.map("\t" + _).mkString(Newline)} $Newline")
-        authorizer.addAcls(acls, resource)
+      if (listPrincipals.isEmpty) {
+        printResourceAcls(resourceToAcls)
+      } else {
+        listPrincipals.foreach{principal =>
+          println(s"ACLs for principal `$principal`")
+          val filteredResourceToAcls = resourceToAcls.map { case (resource, acls) =>
+            resource -> acls.filter(acl => principal.toString.equals(acl.principal))
+          }.filter { case (_, acls) => acls.nonEmpty }
+          printResourceAcls(filteredResourceToAcls)
+        }
       }
+    }
 
-      listAcl(opts)
+    private def printResourceAcls(resourceToAcls: Map[ResourcePattern, Set[AccessControlEntry]]): Unit = {
+      for ((resource, acls) <- resourceToAcls)
+        println(s"Current ACLs for resource `$resource`: $Newline ${acls.map("\t" + _).mkString(Newline)} $Newline")
+    }
+
+    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, acl.toFilter)).toList.asJava
+        adminClient.deleteAcls(aclBindingFilters).all().get()
+      }
+    }
+
+    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 {
+          val results = for (filter <- filters) yield {
+            adminClient.describeAcls(new AclBindingFilter(filter, AccessControlEntryFilter.ANY)).values().get().asScala.toList
+          }
+          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 removeAcl(opts: AclCommandOptions) {
-    withAuthorizer(opts) { authorizer =>
+  class AuthorizerService(val authorizerClassName: String, val opts: AclCommandOptions) extends AclCommandService with Logging {
+
+    private def withAuthorizer()(f: Authorizer => Unit): Unit = {
+      // It is possible that zookeeper.set.acl could be true without SASL if mutual certificate authentication is configured.
+      // We will default the value of zookeeper.set.acl to true or false based on whether SASL is configured,
+      // but if SASL is not configured and zookeeper.set.acl is supposed to be true due to mutual certificate authentication
+      // then it will be up to the user to explicitly specify zookeeper.set.acl=true in the authorizer-properties.
+      val defaultProps = Map(KafkaConfig.ZkEnableSecureAclsProp -> JaasUtils.isZkSaslEnabled)
+      val authorizerPropertiesWithoutTls =
+        if (opts.options.has(opts.authorizerPropertiesOpt)) {
+          val authorizerProperties = opts.options.valuesOf(opts.authorizerPropertiesOpt).asScala
+          defaultProps ++ CommandLineUtils.parseKeyValueArgs(authorizerProperties, acceptMissingValue = false).asScala
+        } else {
+          defaultProps
+        }
+      val authorizerProperties =
+        if (opts.options.has(opts.zkTlsConfigFile)) {
+          // load in TLS configs both with and without the "authorizer." prefix
+          val validKeys = (KafkaConfig.ZkSslConfigToSystemPropertyMap.keys.toList ++ KafkaConfig.ZkSslConfigToSystemPropertyMap.keys.map("authorizer." + _).toList).asJava
+          authorizerPropertiesWithoutTls ++ Utils.loadProps(opts.options.valueOf(opts.zkTlsConfigFile), validKeys).asInstanceOf[java.util.Map[String, Any]].asScala
+        }
+        else
+          authorizerPropertiesWithoutTls
+
+      val authZ = AuthorizerUtils.createAuthorizer(authorizerClassName)
+      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.ifPresent { exception =>
+              println(s"Error while adding ACLs: ${exception.getMessage}")
+              println(Utils.stackTrace(exception))
+            }
+          }
+        }
 
-      for ((resource, acls) <- resourceToAcl) {
-        if (acls.isEmpty) {
-          if (confirmAction(opts, s"Are you sure you want to delete all ACLs for resource `$resource`? (y/n)"))
-            authorizer.removeAcls(resource)
+        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 {
-          if (confirmAction(opts, s"Are you sure you want to remove ACLs: $Newline ${acls.map("\t" + _).mkString(Newline)} $Newline from resource `$resource`? (y/n)"))
-            authorizer.removeAcls(acls, resource)
+          listPrincipals.foreach(principal => {
+            println(s"ACLs for principal `$principal`")
+            val filteredResourceToAcls =  resourceToAcls.map { case (resource, acls) =>
+              resource -> acls.filter(acl => principal.toString.equals(acl.principal))
+            }.filter { case (_, acls) => acls.nonEmpty }
+
+            for ((resource, acls) <- filteredResourceToAcls)
+              println(s"Current ACLs for resource `$resource`: $Newline ${acls.map("\t" + _).mkString(Newline)} $Newline")
+          })
         }
       }
+    }
 
-      listAcl(opts)
+    private def removeAcls(authorizer: Authorizer, 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.ifPresent { exception =>
+          println(s"Error while removing ACLs: ${exception.getMessage}")
+          println(Utils.stackTrace(exception))
+        }
+        result.aclBindingDeleteResults.forEach { deleteResult =>
+          deleteResult.exception.ifPresent { exception =>
+            println(s"Error while removing ACLs: ${exception.getMessage}")
+            println(Utils.stackTrace(exception))
+          }
+        }
+      }
     }
-  }
 
-  private def listAcl(opts: AclCommandOptions) {
-    withAuthorizer(opts) { authorizer =>
-      val resources = getResource(opts, dieIfNoResourceFound = false)
+    private def getAcls(authorizer: Authorizer, 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: Iterable[(Resource, Set[Acl])] =
-        if (resources.isEmpty) authorizer.getAcls()
-        else resources.map(resource => resource -> authorizer.getAcls(resource))
+      val resourceToAcls = mutable.Map[ResourcePattern, Set[AccessControlEntry]]().withDefaultValue(Set())
 
-      for ((resource, acls) <- resourceToAcls)
-        println(s"Current ACLs for resource `$resource`: $Newline ${acls.map("\t" + _).mkString(Newline)} $Newline")
+      aclBindings.foreach(aclBinding => resourceToAcls(aclBinding.pattern()) = resourceToAcls(aclBinding.pattern()) + aclBinding.entry())
+      resourceToAcls.toMap
+    }
+  }
+
+  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) =>
+        new ResourcePattern(filter.resourceType(), filter.name(), filter.patternType()) -> acls
     }
+
+    if (resourceToAcl.values.exists(_.isEmpty))
+      CommandLineUtils.printUsageAndDie(opts.parser, "You must specify one of: --allow-principal, --deny-principal when trying to add ACLs.")
+
+    resourceToAcl
   }
 
-  private def getResourceToAcls(opts: AclCommandOptions): Map[Resource, Set[Acl]] = {
-    var resourceToAcls = Map.empty[Resource, 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)) {
-      resourceToAcls ++= getCliResourceToAcls(opts)
+      resourceToAcls ++= getCliResourceFilterToAcls(opts)
     }
 
     //users are allowed to specify both --producer and --consumer options in a single command.
     if (opts.options.has(opts.producerOpt))
-      resourceToAcls ++= getProducerResourceToAcls(opts)
+      resourceToAcls ++= getProducerResourceFilterToAcls(opts)
 
     if (opts.options.has(opts.consumerOpt))
-      resourceToAcls ++= getConsumerResourceToAcls(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 getProducerResourceToAcls(opts: AclCommandOptions): Map[Resource, Set[Acl]] = {
-    val topics: Set[Resource] = getResource(opts).filter(_.resourceType == Topic)
-    val transactionalIds: Set[Resource] = getResource(opts).filter(_.resourceType == TransactionalId)
+  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 acls = getAcl(opts, Set(Write, Describe))
+    val topicAcls = getAcl(opts, Set(WRITE, DESCRIBE, CREATE))
+    val transactionalIdAcls = getAcl(opts, Set(WRITE, DESCRIBE))
 
-    //Write, Describe permission on topics, Create permission on cluster, Write, Describe on transactionalIds
-    topics.map(_ -> acls).toMap[Resource, Set[Acl]] ++
-      transactionalIds.map(_ -> acls).toMap[Resource, Set[Acl]] +
-      (Resource.ClusterResource -> (getAcl(opts, Set(Create)) ++
-        (if (enableIdempotence) getAcl(opts, Set(IdempotentWrite)) else Set.empty[Acl])))
+    //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(IDEMPOTENT_WRITE)))
+        else
+          Map.empty)
   }
 
-  private def getConsumerResourceToAcls(opts: AclCommandOptions): Map[Resource, Set[Acl]] = {
-    val resources = getResource(opts)
+  private def getConsumerResourceFilterToAcls(opts: AclCommandOptions): Map[ResourcePatternFilter, Set[AccessControlEntry]] = {
+    val filters = getResourceFilter(opts)
 
-    val topics: Set[Resource] = getResource(opts).filter(_.resourceType == Topic)
-    val groups: Set[Resource] = resources.filter(_.resourceType == Group)
+    val topics: Set[ResourcePatternFilter] = filters.filter(_.resourceType == JResourceType.TOPIC)
+    val groups: Set[ResourcePatternFilter] = filters.filter(_.resourceType == JResourceType.GROUP)
 
-    //Read,Describe on topic, Read on consumerGroup + Create on cluster
+    //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[Resource, Set[Acl]] ++
-      groups.map(_ -> getAcl(opts, Set(Read))).toMap[Resource, Set[Acl]]
+    topics.map(_ -> acls).toMap[ResourcePatternFilter, Set[AccessControlEntry]] ++
+      groups.map(_ -> getAcl(opts, Set(READ))).toMap[ResourcePatternFilter, Set[AccessControlEntry]]
   }
 
-  private def getCliResourceToAcls(opts: AclCommandOptions): Map[Resource, Set[Acl]] = {
+  private def getCliResourceFilterToAcls(opts: AclCommandOptions): Map[ResourcePatternFilter, Set[AccessControlEntry]] = {
     val acls = getAcl(opts)
-    val resources = getResource(opts)
-    resources.map(_ -> acls).toMap
+    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)
@@ -190,28 +402,29 @@ 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 => JSecurityUtils.operation(operation.trim)).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],
@@ -219,61 +432,77 @@ object AclCommand extends Logging {
     if (opts.options.has(hostOptionSpec))
       opts.options.valuesOf(hostOptionSpec).asScala.map(_.trim).toSet
     else if (opts.options.has(principalOptionSpec))
-      Set[String](Acl.WildCardHost)
+      Set[String](AclEntry.WildcardHost)
     else
       Set.empty[String]
   }
 
   private def getPrincipals(opts: AclCommandOptions, principalOptionSpec: ArgumentAcceptingOptionSpec[String]): Set[KafkaPrincipal] = {
     if (opts.options.has(principalOptionSpec))
-      opts.options.valuesOf(principalOptionSpec).asScala.map(s => KafkaPrincipal.fromString(s.trim)).toSet
+      opts.options.valuesOf(principalOptionSpec).asScala.map(s => JSecurityUtils.parseKafkaPrincipal(s.trim)).toSet
     else
       Set.empty[KafkaPrincipal]
   }
 
-  private def getResource(opts: AclCommandOptions, dieIfNoResourceFound: Boolean = true): Set[Resource] = {
-    var resources = Set.empty[Resource]
+  private def getResourceFilter(opts: AclCommandOptions, dieIfNoResourceFound: Boolean = true): Set[ResourcePatternFilter] = {
+    val patternType: PatternType = opts.options.valueOf(opts.resourcePatternType)
+
+    var resourceFilters = Set.empty[ResourcePatternFilter]
     if (opts.options.has(opts.topicOpt))
-      opts.options.valuesOf(opts.topicOpt).asScala.foreach(topic => resources += new Resource(Topic, topic.trim))
+      opts.options.valuesOf(opts.topicOpt).forEach(topic => resourceFilters += new ResourcePatternFilter(JResourceType.TOPIC, topic.trim, patternType))
 
-    if (opts.options.has(opts.clusterOpt) || opts.options.has(opts.idempotentOpt))
-      resources += Resource.ClusterResource
+    if (patternType == PatternType.LITERAL && (opts.options.has(opts.clusterOpt) || opts.options.has(opts.idempotentOpt)))
+      resourceFilters += ClusterResourceFilter
 
     if (opts.options.has(opts.groupOpt))
-      opts.options.valuesOf(opts.groupOpt).asScala.foreach(group => resources += new Resource(Group, group.trim))
+      opts.options.valuesOf(opts.groupOpt).forEach(group => resourceFilters += new ResourcePatternFilter(JResourceType.GROUP, group.trim, patternType))
 
     if (opts.options.has(opts.transactionalIdOpt))
-      opts.options.valuesOf(opts.transactionalIdOpt).asScala.foreach(transactionalId =>
-        resources += new Resource(TransactionalId, transactionalId))
+      opts.options.valuesOf(opts.transactionalIdOpt).forEach(transactionalId =>
+        resourceFilters += new ResourcePatternFilter(JResourceType.TRANSACTIONAL_ID, transactionalId, patternType))
+
+    if (opts.options.has(opts.delegationTokenOpt))
+      opts.options.valuesOf(opts.delegationTokenOpt).forEach(token => resourceFilters += new ResourcePatternFilter(JResourceType.DELEGATION_TOKEN, token.trim, patternType))
 
-    if (resources.isEmpty && dieIfNoResourceFound)
-      CommandLineUtils.printUsageAndDie(opts.parser, "You must provide at least one resource: --topic  or --cluster or --group ")
+    if (resourceFilters.isEmpty && dieIfNoResourceFound)
+      CommandLineUtils.printUsageAndDie(opts.parser, "You must provide at least one resource: --topic  or --cluster or --group  or --delegation-token ")
 
-    resources
+    resourceFilters
   }
 
   private def confirmAction(opts: AclCommandOptions, msg: String): Boolean = {
     if (opts.options.has(opts.forceOpt))
         return true
     println(msg)
-    Console.readLine().equalsIgnoreCase("y")
+    StdIn.readLine().equalsIgnoreCase("y")
   }
 
-  private def validateOperation(opts: AclCommandOptions, resourceToAcls: Map[Resource, Set[Acl]]) = {
+  private def validateOperation(opts: AclCommandOptions, resourceToAcls: Map[ResourcePatternFilter, Set[AccessControlEntry]]): Unit = {
     for ((resource, acls) <- resourceToAcls) {
-      val validOps = ResourceTypeToValidOperations(resource.resourceType)
+      val validOps = AclEntry.supportedOperations(resource.resourceType) + AclOperation.ALL
       if ((acls.map(_.operation) -- validOps).nonEmpty)
         CommandLineUtils.printUsageAndDie(opts.parser, s"ResourceType ${resource.resourceType} only supports operations ${validOps.mkString(",")}")
     }
   }
 
-  class AclCommandOptions(args: Array[String]) {
-    val parser = new OptionParser(false)
-    val authorizerOpt = parser.accepts("authorizer", "Fully qualified class name of the authorizer, defaults to kafka.security.auth.SimpleAclAuthorizer.")
+  class AclCommandOptions(args: Array[String]) extends CommandDefaultOptions(args) {
+    val CommandConfigDoc = "A property file containing configs to be passed to Admin Client."
+
+    val bootstrapServerOpt = parser.accepts("bootstrap-server", "A list of host/port pairs to use for establishing the connection to the Kafka cluster." +
+      " This list should be in the form host1:port1,host2:port2,... This config is required for acl management using admin client API.")
+      .withRequiredArg
+      .describedAs("server to connect to")
+      .ofType(classOf[String])
+
+    val commandConfigOpt = parser.accepts("command-config", CommandConfigDoc)
+      .withOptionalArg()
+      .describedAs("command-config")
+      .ofType(classOf[String])
+
+    val authorizerOpt = parser.accepts("authorizer", "Fully qualified class name of the authorizer, defaults to kafka.security.authorizer.AclAuthorizer.")
       .withRequiredArg
       .describedAs("authorizer")
       .ofType(classOf[String])
-      .defaultsTo(classOf[SimpleAclAuthorizer].getName)
 
     val authorizerPropertiesOpt = parser.accepts("authorizer-properties", "REQUIRED: properties required to configure an instance of Authorizer. " +
       "These are key=val pairs. For the default authorizer the example values are: zookeeper.connect=localhost:2181")
@@ -304,15 +533,32 @@ object AclCommand extends Logging {
       "used in combination with the --producer option. Note that idempotence is enabled automatically if " +
       "the producer is authorized to a particular transactional-id.")
 
+    val delegationTokenOpt = parser.accepts("delegation-token", "Delegation token to which ACLs should be added or removed. " +
+      "A value of * indicates ACL should apply to all tokens.")
+      .withRequiredArg
+      .describedAs("delegation-token")
+      .ofType(classOf[String])
+
+    val resourcePatternType = parser.accepts("resource-pattern-type", "The type of the resource pattern or pattern filter. " +
+      "When adding acls, this should be a specific pattern type, e.g. 'literal' or 'prefixed'. " +
+      "When listing or removing acls, a specific pattern type can be used to list or remove acls from specific resource patterns, " +
+      "or use the filter values of 'any' or 'match', where 'any' will match any pattern type, but will match the resource name exactly, " +
+      "where as 'match' will perform pattern matching to list or remove all acls that affect the supplied resource(s). " +
+      "WARNING: 'match', when used in combination with the '--remove' switch, should be used with care.")
+      .withRequiredArg()
+      .ofType(classOf[String])
+      .withValuesConvertedBy(new PatternTypeConverter())
+      .defaultsTo(PatternType.LITERAL)
+
     val addOpt = parser.accepts("add", "Indicates you are trying to add ACLs.")
     val removeOpt = parser.accepts("remove", "Indicates you are trying to remove ACLs.")
     val listOpt = parser.accepts("list", "List ACLs for the specified resource, use --topic  or --group  or --cluster to specify a resource.")
 
     val operationsOpt = parser.accepts("operation", "Operation that is being allowed or denied. Valid operation names are: " + Newline +
-      Operation.values.map("\t" + _).mkString(Newline) + Newline)
+      AclEntry.AclOperations.map("\t" + JSecurityUtils.operationName(_)).mkString(Newline) + Newline)
       .withRequiredArg
       .ofType(classOf[String])
-      .defaultsTo(All.name)
+      .defaultsTo(JSecurityUtils.operationName(AclOperation.ALL))
 
     val allowPrincipalsOpt = parser.accepts("allow-principal", "principal is in principalType:name format." +
       " Note that principalType must be supported by the Authorizer being used." +
@@ -332,6 +578,12 @@ object AclCommand extends Logging {
       .describedAs("deny-principal")
       .ofType(classOf[String])
 
+    val listPrincipalsOpt = parser.accepts("principal", "List ACLs for the specified principal. principal is in principalType:name format." +
+      " Note that principalType must be supported by the Authorizer being used. Multiple --principal option can be passed.")
+      .withOptionalArg()
+      .describedAs("principal")
+      .ofType(classOf[String])
+
     val allowHostsOpt = parser.accepts("allow-host", "Host from which principals listed in --allow-principal will have access. " +
       "If you have specified --allow-principal then the default for this option will be set to * which allows access from all hosts.")
       .withRequiredArg
@@ -345,19 +597,34 @@ object AclCommand extends Logging {
       .ofType(classOf[String])
 
     val producerOpt = parser.accepts("producer", "Convenience option to add/remove ACLs for producer role. " +
-      "This will generate ACLs that allows WRITE,DESCRIBE on topic and CREATE on cluster. ")
+      "This will generate ACLs that allows WRITE,DESCRIBE and CREATE on topic.")
 
     val consumerOpt = parser.accepts("consumer", "Convenience option to add/remove ACLs for consumer role. " +
       "This will generate ACLs that allows READ,DESCRIBE on topic and READ on group.")
 
-    val helpOpt = parser.accepts("help", "Print usage information.")
-
     val forceOpt = parser.accepts("force", "Assume Yes to all queries and do not prompt.")
 
-    val options = parser.parse(args: _*)
+    val zkTlsConfigFile = parser.accepts("zk-tls-config-file",
+      "Identifies the file where ZooKeeper client TLS connectivity properties for the authorizer are defined.  Any properties other than the following (with or without an \"authorizer.\" prefix) are ignored: " +
+        KafkaConfig.ZkSslConfigToSystemPropertyMap.keys.toList.sorted.mkString(", ") +
+        ". Note that if SASL is not configured and zookeeper.set.acl is supposed to be true due to mutual certificate authentication being used" +
+        " then it is necessary to explicitly specify --authorizer-properties zookeeper.set.acl=true")
+      .withRequiredArg().describedAs("Authorizer ZooKeeper TLS configuration").ofType(classOf[String])
+
+    options = parser.parse(args: _*)
+
+    def checkArgs(): Unit = {
+      if (options.has(bootstrapServerOpt) && options.has(authorizerOpt))
+        CommandLineUtils.printUsageAndDie(parser, "Only one of --bootstrap-server or --authorizer must be specified")
+
+      if (!options.has(bootstrapServerOpt))
+        CommandLineUtils.checkRequiredArgs(parser, options, authorizerPropertiesOpt)
+
+      if (options.has(commandConfigOpt) && !options.has(bootstrapServerOpt))
+        CommandLineUtils.printUsageAndDie(parser, "The --command-config option can only be used with --bootstrap-server option")
 
-    def checkArgs() {
-      CommandLineUtils.checkRequiredArgs(parser, options, authorizerPropertiesOpt)
+      if (options.has(authorizerPropertiesOpt) && options.has(bootstrapServerOpt))
+        CommandLineUtils.printUsageAndDie(parser, "The --authorizer-properties option can only be used with --authorizer option")
 
       val actions = Seq(addOpt, removeOpt, listOpt).count(options.has)
       if (actions != 1)
@@ -369,6 +636,9 @@ object AclCommand extends Logging {
       CommandLineUtils.checkInvalidArgs(parser, options, producerOpt, Set(operationsOpt, denyPrincipalsOpt, denyHostsOpt))
       CommandLineUtils.checkInvalidArgs(parser, options, consumerOpt, Set(operationsOpt, denyPrincipalsOpt, denyHostsOpt))
 
+      if (options.has(listPrincipalsOpt) && !options.has(listOpt))
+        CommandLineUtils.printUsageAndDie(parser, "The --principal option is only available if --list is set")
+
       if (options.has(producerOpt) && !options.has(topicOpt))
         CommandLineUtils.printUsageAndDie(parser, "With --producer you must specify a --topic")
 
@@ -379,5 +649,19 @@ object AclCommand extends Logging {
         CommandLineUtils.printUsageAndDie(parser, "With --consumer you must specify a --topic and a --group and no --cluster or --transactional-id option should be specified.")
     }
   }
+}
+
+class PatternTypeConverter extends EnumConverter[PatternType](classOf[PatternType]) {
+
+  override def convert(value: String): PatternType = {
+    val patternType = super.convert(value)
+    if (patternType.isUnknown)
+      throw new ValueConversionException("Unknown resource-pattern-type: " + value)
+
+    patternType
+  }
 
+  override def valuePattern: String = PatternType.values
+    .filter(_ != PatternType.UNKNOWN)
+    .mkString("|")
 }
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 1051993416ae4..0000000000000
--- a/core/src/main/scala/kafka/admin/AdminClient.scala
+++ /dev/null
@@ -1,490 +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.admin.AdminClient.DeleteRecordsResult
-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, RequestFutureAdapter}
-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.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.DescribeGroupsResponse.GroupMetadata
-import org.apache.kafka.common.requests.OffsetFetchResponse
-import org.apache.kafka.common.utils.{LogContext, KafkaThread, Time, Utils}
-import org.apache.kafka.common.{Cluster, 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 KafkaAdminClient.
-  * @see KafkaAdminClient
-  */
-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(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(_.protocolType == ConsumerProtocol.PROTOCOL_TYPE)
-    }
-  }
-
-  def listAllGroupsFlattened(): List[GroupOverview] = {
-    listAllGroups.values.flatten.toList
-  }
-
-  def listAllConsumerGroupsFlattened(): List[GroupOverview] = {
-    listAllGroupsFlattened.filter(_.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
-
-  /*
-   * Remove all the messages whose offset is smaller than the given offset of the corresponding partition
-   *
-   * DeleteRecordsResult contains either lowWatermark of the partition or exception. We list the possible exception
-   * and their interpretations below:
-   *
-   * - DisconnectException if leader node of the partition is not available. Need retry by user.
-   * - PolicyViolationException if the topic is configured as non-deletable.
-   * - TopicAuthorizationException if the topic doesn't exist and the user doesn't have the authority to create the topic
-   * - TimeoutException if response is not available within the timeout specified by either Future's timeout or AdminClient's request timeout
-   * - UnknownTopicOrPartitionException if the partition doesn't exist or if the user doesn't have the authority to describe the topic
-   * - NotLeaderForPartitionException if broker is not leader of the partition. Need retry by user.
-   * - OffsetOutOfRangeException if the offset is larger than high watermark of this partition
-   *
-   */
-
-  def deleteRecordsBefore(offsets: Map[TopicPartition, Long]): Future[Map[TopicPartition, DeleteRecordsResult]] = {
-    val metadataRequest = new MetadataRequest.Builder(offsets.keys.map(_.topic).toSet.toList.asJava, true)
-    val response = sendAnyNode(ApiKeys.METADATA, metadataRequest).asInstanceOf[MetadataResponse]
-    val errors = response.errors
-    if (!errors.isEmpty)
-      error(s"Metadata request contained errors: $errors")
-
-    val (partitionsWithoutError, partitionsWithError) = offsets.partition{ partitionAndOffset =>
-      !response.errors().containsKey(partitionAndOffset._1.topic())}
-
-    val (partitionsWithLeader, partitionsWithoutLeader) = partitionsWithoutError.partition{ partitionAndOffset =>
-      response.cluster().leaderFor(partitionAndOffset._1) != null}
-
-    val partitionsWithErrorResults = partitionsWithError.keys.map( partition =>
-      partition -> DeleteRecordsResult(DeleteRecordsResponse.INVALID_LOW_WATERMARK, response.errors().get(partition.topic()).exception())).toMap
-
-    val partitionsWithoutLeaderResults = partitionsWithoutLeader.mapValues( _ =>
-      DeleteRecordsResult(DeleteRecordsResponse.INVALID_LOW_WATERMARK, Errors.LEADER_NOT_AVAILABLE.exception()))
-
-    val partitionsGroupByLeader = partitionsWithLeader.groupBy(partitionAndOffset =>
-      response.cluster().leaderFor(partitionAndOffset._1))
-
-    // prepare requests and generate Future objects
-    val futures = partitionsGroupByLeader.map{ case (node, partitionAndOffsets) =>
-      val convertedMap: java.util.Map[TopicPartition, java.lang.Long] = partitionAndOffsets.mapValues(_.asInstanceOf[java.lang.Long]).asJava
-      val future = client.send(node, new DeleteRecordsRequest.Builder(requestTimeoutMs, convertedMap))
-      pendingFutures.add(future)
-      future.compose(new RequestFutureAdapter[ClientResponse, Map[TopicPartition, DeleteRecordsResult]]() {
-          override def onSuccess(response: ClientResponse, future: RequestFuture[Map[TopicPartition, DeleteRecordsResult]]) {
-            val deleteRecordsResponse = response.responseBody().asInstanceOf[DeleteRecordsResponse]
-            val result = deleteRecordsResponse.responses().asScala.mapValues(v => DeleteRecordsResult(v.lowWatermark, v.error.exception())).toMap
-            future.complete(result)
-            pendingFutures.remove(future)
-          }
-
-          override def onFailure(e: RuntimeException, future: RequestFuture[Map[TopicPartition, DeleteRecordsResult]]) {
-            val result = partitionAndOffsets.mapValues(_ => DeleteRecordsResult(DeleteRecordsResponse.INVALID_LOW_WATERMARK, e))
-            future.complete(result)
-            pendingFutures.remove(future)
-          }
-
-        })
-    }
-
-    // default output if not receiving DeleteRecordsResponse before timeout
-    val defaultResults = offsets.mapValues(_ =>
-      DeleteRecordsResult(DeleteRecordsResponse.INVALID_LOW_WATERMARK, Errors.REQUEST_TIMED_OUT.exception())) ++ partitionsWithErrorResults ++ partitionsWithoutLeaderResults
-
-    new CompositeFuture(time, defaultResults, futures.toList)
-  }
-
-  /**
-   * 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): GroupMetadata = {
-    val responseBody = send(coordinator, ApiKeys.DESCRIBE_GROUPS,
-        new DescribeGroupsRequest.Builder(Collections.singletonList(groupId)))
-    val response = responseBody.asInstanceOf[DescribeGroupsResponse]
-    val metadata = response.groups.get(groupId)
-    if (metadata == null)
-      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: DescribeGroupsResponse.GroupMetadata): Boolean =
-      metadata.error == Errors.NONE && (metadata.state == "Dead" || metadata.state == "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.state match {
-        case "Stable" =>
-          val assignment = ConsumerProtocol.deserializeAssignment(ByteBuffer.wrap(Utils.readBytes(consumer.memberAssignment)))
-          assignment.partitions.asScala.toList
-        case _ =>
-          List()
-      })
-    }.toList
-
-    ConsumerGroupSummary(metadata.state, metadata.protocol, Some(consumers), coordinator)
-  }
-
-  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)
-  }
-}
-
-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.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
-  }
-
-  case class DeleteRecordsResult(lowWatermark: Long, error: Exception)
-
-  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 time = Time.SYSTEM
-    val metrics = new Metrics(time)
-    val metadata = new Metadata(100L, 60 * 60 * 1000L, true)
-    val channelBuilder = ClientUtils.createChannelBuilder(config)
-    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 brokerAddresses = ClientUtils.parseAndValidateAddresses(brokerUrls)
-    val bootstrapCluster = Cluster.bootstrap(brokerAddresses)
-    metadata.update(bootstrapCluster, Collections.emptySet(), 0)
-
-    val selector = new Selector(
-      DefaultConnectionMaxIdleMs,
-      metrics,
-      time,
-      "admin",
-      channelBuilder,
-      new LogContext())
-
-    val networkClient = new NetworkClient(
-      selector,
-      metadata,
-      "admin-" + AdminClientIdSequence.getAndIncrement(),
-      DefaultMaxInFlightRequestsPerConnection,
-      DefaultReconnectBackoffMs,
-      DefaultReconnectBackoffMax,
-      DefaultSendBufferBytes,
-      DefaultReceiveBufferBytes,
-      requestTimeoutMs,
-      time,
-      true,
-      new ApiVersions,
-      new LogContext())
-
-    val highLevelClient = new ConsumerNetworkClient(
-      new LogContext(),
-      networkClient,
-      metadata,
-      time,
-      retryBackoffMs,
-      requestTimeoutMs.toLong,
-      Integer.MAX_VALUE)
-
-    new AdminClient(
-      time,
-      requestTimeoutMs,
-      retryBackoffMs,
-      highLevelClient,
-      bootstrapCluster.nodes.asScala.toList)
-  }
-}
diff --git a/core/src/main/scala/kafka/admin/AdminUtils.scala b/core/src/main/scala/kafka/admin/AdminUtils.scala
index 6665d2559a672..ce64be0db6e7f 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.{BrokerNotAvailableException, InvalidPartitionsException, InvalidReplicaAssignmentException, InvalidReplicationFactorException, InvalidTopicException, TopicExistsException, UnknownTopicOrPartitionException}
+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.
@@ -237,7 +211,7 @@ object AdminUtils extends Logging with AdminUtilities {
     */
   private[admin] def getRackAlternatedBrokerList(brokerRackMap: Map[Int, String]): IndexedSeq[Int] = {
     val brokersIteratorByRack = getInverseMap(brokerRackMap).map { case (rack, brokers) =>
-      (rack, brokers.toIterator)
+      (rack, brokers.iterator)
     }
     val racks = brokersIteratorByRack.keys.toArray.sorted
     val result = new mutable.ArrayBuffer[Int]
@@ -256,444 +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 has been deprecated and will be removed in a future release.", "0.11.0.0")
-  def isConsumerGroupActive(zkUtils: ZkUtils, group: String) = {
-    zkUtils.getConsumersInGroup(group).nonEmpty
-  }
-
-  /**
-   * Delete the whole directory of the given consumer group if the group is inactive.
-   *
-   * @param zkUtils Zookeeper utilities
-   * @param group Consumer group
-   * @return whether or not we deleted the consumer group information
-   */
-  @deprecated("This method has been deprecated and will be removed in a future release.", "0.11.0.0")
-  def deleteConsumerGroupInZK(zkUtils: ZkUtils, group: String) = {
-    if (!isConsumerGroupActive(zkUtils, group)) {
-      val dir = new ZKGroupDirs(group)
-      zkUtils.deletePathRecursive(dir.consumerGroupDir)
-      true
-    }
-    else false
-  }
-
-  /**
-   * Delete the given consumer group's information for the given topic in Zookeeper if the group is inactive.
-   * If the consumer group consumes no other topics, delete the whole consumer group directory.
-   *
-   * @param zkUtils Zookeeper utilities
-   * @param group Consumer group
-   * @param topic Topic of the consumer group information we wish to delete
-   * @return whether or not we deleted the consumer group information for the given topic
-   */
-  @deprecated("This method has been deprecated and will be removed in a future release.", "0.11.0.0")
-  def deleteConsumerGroupInfoForTopicInZK(zkUtils: ZkUtils, group: String, topic: String) = {
-    val topics = zkUtils.getTopicsByConsumerGroup(group)
-    if (topics == Seq(topic)) {
-      deleteConsumerGroupInZK(zkUtils, group)
-    }
-    else if (!isConsumerGroupActive(zkUtils, group)) {
-      val dir = new ZKGroupTopicDirs(group, topic)
-      zkUtils.deletePathRecursive(dir.consumerOwnerDir)
-      zkUtils.deletePathRecursive(dir.consumerOffsetDir)
-      true
-    }
-    else false
-  }
-
-  /**
-   * Delete every inactive consumer group's information about the given topic in Zookeeper.
-   *
-   * @param zkUtils Zookeeper utilities
-   * @param topic Topic of the consumer group information we wish to delete
-   */
-  @deprecated("This method has been deprecated and will be removed in a future release.", "0.11.0.0")
-  def deleteAllConsumerGroupInfoForTopicInZK(zkUtils: ZkUtils, topic: String) {
-    val groups = zkUtils.getAllConsumerGroupsForTopic(topic)
-    groups.foreach(group => deleteConsumerGroupInfoForTopicInZK(zkUtils, group, topic))
-  }
-
-  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)
-  }
-
-  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("Topic creation " + jsonPartitionData.toString)
-        zkUtils.createPersistentPath(zkPath, jsonPartitionData)
-      } else {
-        info("Topic update " + jsonPartitionData.toString)
-        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)
-  }
-
-  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)
-    }
-  }
-
-  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.encode(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
-   */
-  private def writeEntityConfig(zkUtils: ZkUtils, entityPath: String, config: Properties) {
-    val map = Map("version" -> 1, "config" -> config.asScala)
-    zkUtils.updatePersistentPath(entityPath, Json.encode(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 b25a8da20858f..c8a4b3d400c35 100644
--- a/core/src/main/scala/kafka/admin/BrokerApiVersionsCommand.scala
+++ b/core/src/main/scala/kafka/admin/BrokerApiVersionsCommand.scala
@@ -18,14 +18,33 @@
 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.CommandLineUtils
+import kafka.utils.{CommandDefaultOptions, CommandLineUtils}
+import kafka.utils.Implicits._
+import kafka.utils.Logging
 import org.apache.kafka.common.utils.Utils
-import org.apache.kafka.clients.CommonClientConfigs
-import joptsimple._
+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.jdk.CollectionConverters._
+import scala.util.{Failure, Success, Try}
 
 /**
  * A command for retrieving broker version information.
@@ -41,7 +60,7 @@ object BrokerApiVersionsCommand {
     val adminClient = createAdminClient(opts)
     adminClient.awaitBrokers()
     val brokerMap = adminClient.listAllBrokerVersionInfo()
-    brokerMap.foreach { case (broker, versionInfoOrError) =>
+    brokerMap.forKeyValue { (broker, versionInfoOrError) =>
       versionInfoOrError match {
         case Success(v) => out.print(s"${broker} -> ${v.toString(true)}\n")
         case Failure(v) => out.print(s"${broker} -> ERROR: ${v}\n")
@@ -59,11 +78,10 @@ object BrokerApiVersionsCommand {
     AdminClient.create(props)
   }
 
-  class BrokerVersionCommandOptions(args: Array[String]) {
+  class BrokerVersionCommandOptions(args: Array[String]) extends CommandDefaultOptions(args) {
     val BootstrapServerDoc = "REQUIRED: The server to connect to."
     val CommandConfigDoc = "A property file containing configs to be passed to Admin Client."
 
-    val parser = new OptionParser(false)
     val commandConfigOpt = parser.accepts("command-config", CommandConfigDoc)
                                  .withRequiredArg
                                  .describedAs("command config property file")
@@ -72,12 +90,253 @@ object BrokerApiVersionsCommand {
                                    .withRequiredArg
                                    .describedAs("server(s) to use for bootstrapping")
                                    .ofType(classOf[String])
-    val options = parser.parse(args : _*)
+    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", () => {
+      try {
+        while (running)
+          client.poll(time.timer(Long.MaxValue))
+      } catch {
+        case t: Throwable =>
+          error("admin-client-network-thread exited", t)
+      } finally {
+        pendingFutures.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 DefaultSocketConnectionSetupMs = CommonClientConfigs.SOCKET_CONNECTION_SETUP_TIMEOUT_MS_CONFIG
+    val DefaultSocketConnectionSetupMaxMs = CommonClientConfigs.SOCKET_CONNECTION_SETUP_TIMEOUT_MAX_MS_CONFIG
+    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.USE_ALL_DNS_IPS.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.SOCKET_CONNECTION_SETUP_TIMEOUT_MS_CONFIG,
+          ConfigDef.Type.LONG,
+          CommonClientConfigs.DEFAULT_SOCKET_CONNECTION_SETUP_TIMEOUT_MS,
+          ConfigDef.Importance.MEDIUM,
+          CommonClientConfigs.SOCKET_CONNECTION_SETUP_TIMEOUT_MS_DOC)
+        .define(
+          CommonClientConfigs.SOCKET_CONNECTION_SETUP_TIMEOUT_MAX_MS_CONFIG,
+          ConfigDef.Type.LONG,
+          CommonClientConfigs.DEFAULT_SOCKET_CONNECTION_SETUP_TIMEOUT_MAX_MS,
+          ConfigDef.Importance.MEDIUM,
+          CommonClientConfigs.SOCKET_CONNECTION_SETUP_TIMEOUT_MAX_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, logContext)
+      val requestTimeoutMs = config.getInt(CommonClientConfigs.REQUEST_TIMEOUT_MS_CONFIG)
+      val connectionSetupTimeoutMs = config.getLong(CommonClientConfigs.SOCKET_CONNECTION_SETUP_TIMEOUT_MS_CONFIG)
+      val connectionSetupTimeoutMaxMs = config.getLong(CommonClientConfigs.SOCKET_CONNECTION_SETUP_TIMEOUT_MAX_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,
+        connectionSetupTimeoutMs,
+        connectionSetupTimeoutMaxMs,
+        ClientDnsLookup.USE_ALL_DNS_IPS,
+        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 077ecce6df07e..7932c8dddb397 100644
--- a/core/src/main/scala/kafka/admin/ConfigCommand.scala
+++ b/core/src/main/scala/kafka/admin/ConfigCommand.scala
@@ -17,71 +17,116 @@
 
 package kafka.admin
 
-import java.util.Properties
+import java.nio.charset.StandardCharsets
+import java.util.concurrent.TimeUnit
+import java.util.{Collections, Properties}
 
 import joptsimple._
 import kafka.common.Config
-import kafka.common.InvalidConfigException
 import kafka.log.LogConfig
-import kafka.server.{ConfigEntityName, ConfigType, DynamicConfig}
-import kafka.utils.CommandLineUtils
+import kafka.server.DynamicConfig.QuotaConfigs
+import kafka.server.{ConfigEntityName, ConfigType, Defaults, DynamicBrokerConfig, DynamicConfig, KafkaConfig}
+import kafka.utils.{CommandDefaultOptions, CommandLineUtils, Exit, PasswordEncoder}
 import kafka.utils.Implicits._
 import kafka.zk.{AdminZkClient, KafkaZkClient}
-import kafka.zookeeper.ZooKeeperClient
+import org.apache.kafka.clients.admin.{Admin, AlterClientQuotasOptions, AlterConfigOp, AlterConfigsOptions, ConfigEntry, DescribeClusterOptions, DescribeConfigsOptions, ListTopicsOptions, ScramCredentialInfo, UserScramCredentialDeletion, UserScramCredentialUpsertion, Config => JConfig, ScramMechanism => PublicScramMechanism}
+import org.apache.kafka.clients.CommonClientConfigs
+import org.apache.kafka.common.config.ConfigResource
+import org.apache.kafka.common.config.types.Password
+import org.apache.kafka.common.errors.InvalidConfigurationException
+import org.apache.kafka.common.internals.Topic
+import org.apache.kafka.common.quota.{ClientQuotaAlteration, ClientQuotaEntity, ClientQuotaFilter, ClientQuotaFilterComponent}
 import org.apache.kafka.common.security.JaasUtils
-import org.apache.kafka.common.security.scram._
-import org.apache.kafka.common.utils.{Sanitizer, Utils}
+import org.apache.kafka.common.security.scram.internals.{ScramCredentialUtils, ScramFormatter, ScramMechanism}
+import org.apache.kafka.common.utils.{Sanitizer, Time, Utils}
+import org.apache.zookeeper.client.ZKClientConfig
 
+import scala.annotation.nowarn
+import scala.jdk.CollectionConverters._
 import scala.collection._
-import scala.collection.JavaConverters._
-
 
 /**
- * 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
+ * This script can be used to change configs for topics/clients/users/brokers/ips dynamically
  * An entity described or altered by the command may be one of:
  * 
              
- *     
            •  topic: --entity-type topics --entity-name 
- *     
            •  client: --entity-type clients --entity-name 
- *     
            •  user: --entity-type users --entity-name 
- *     
            •  : --entity-type users --entity-name  --entity-type clients --entity-name 
- *     
            •  broker: --entity-type brokers --entity-name 
+ *     
            •  topic: --topic  OR --entity-type topics --entity-name 
+ *     
            •  client: --client  OR --entity-type clients --entity-name 
+ *     
            •  user: --user  OR --entity-type users --entity-name 
+ *     
            •  : --user  --client  OR
+ *                          --entity-type users --entity-name  --entity-type clients --entity-name 
+ *     
            •  broker: --broker  OR --entity-type brokers --entity-name 
+ *     
            •  broker-logger: --broker-logger  OR --entity-type broker-loggers --entity-name 
+ *     
            •  ip: --ip  OR --entity-type ips --entity-name 
  * 
            
- * --entity-default may be used instead of --entity-name when describing or altering default configuration for users and clients.
- *
+ * --entity-type  --entity-default may be specified in place of --entity-type  --entity-name 
+ * when describing or altering default configuration for users, clients, brokers, or ips, respectively.
+ * Alternatively, --user-defaults, --client-defaults, --broker-defaults, or --ip-defaults may be specified in place of
+ * --entity-type  --entity-default, respectively.
  */
 object ConfigCommand extends Config {
 
+  val BrokerDefaultEntityName = ""
+  val BrokerLoggerConfigType = "broker-loggers"
+  val BrokerSupportedConfigTypes = ConfigType.all :+ BrokerLoggerConfigType
+  val ZkSupportedConfigTypes = ConfigType.all
   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
+  // ZooKeeper for bootstrapping. This allows all password configs to be stored encrypted in ZK,
+  // avoiding clear passwords in server.properties. For consistency with older versions, quota-related
+  // broker configs can still be updated using ZooKeeper at any time. ConfigCommand will be migrated
+  // to the new AdminClient later for these configs (KIP-248).
+  val BrokerConfigsUpdatableUsingZooKeeperWhileBrokerRunning = Set(
+    DynamicConfig.Broker.LeaderReplicationThrottledRateProp,
+    DynamicConfig.Broker.FollowerReplicationThrottledRateProp,
+    DynamicConfig.Broker.ReplicaAlterLogDirsIoMaxBytesPerSecondProp)
 
   def main(args: Array[String]): Unit = {
+    try {
+      val opts = new ConfigCommandOptions(args)
 
-    val opts = new ConfigCommandOptions(args)
+      CommandLineUtils.printHelpAndExitIfNeeded(opts, "This tool helps to manipulate and describe entity config for a topic, client, user, broker or ip")
 
-    if(args.length == 0)
-      CommandLineUtils.printUsageAndDie(opts.parser, "Add/Remove entity config for a topic, client, user or broker")
+      opts.checkArgs()
 
-    opts.checkArgs()
+      if (opts.options.has(opts.zkConnectOpt)) {
+        println(s"Warning: --zookeeper is deprecated and will be removed in a future version of Kafka.")
+        println(s"Use --bootstrap-server instead to specify a broker to connect to.")
+        processCommandWithZk(opts.options.valueOf(opts.zkConnectOpt), opts)
+      } else {
+        processCommand(opts)
+      }
+    } catch {
+      case e @ (_: IllegalArgumentException | _: InvalidConfigurationException | _: OptionException) =>
+        logger.debug(s"Failed config command with args '${args.mkString(" ")}'", e)
+        System.err.println(e.getMessage)
+        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)
+    }
+  }
 
-    val zooKeeperClient = new ZooKeeperClient(opts.options.valueOf(opts.zkConnectOpt), 30000, 30000, Int.MaxValue)
-    val zkClient = new KafkaZkClient(zooKeeperClient, JaasUtils.isZkSecurityEnabled())
+  private def processCommandWithZk(zkConnectString: String, opts: ConfigCommandOptions): Unit = {
+    val zkClientConfig = ZkSecurityMigrator.createZkClientConfigFromOption(opts.options, opts.zkTlsConfigFile)
+      .getOrElse(new ZKClientConfig())
+    val zkClient = KafkaZkClient(zkConnectString, JaasUtils.isZkSaslEnabled || KafkaConfig.zkTlsClientAuthEnabled(zkClientConfig), 30000, 30000,
+      Int.MaxValue, Time.SYSTEM, zkClientConfig = Some(zkClientConfig))
     val adminZkClient = new AdminZkClient(zkClient)
-
     try {
       if (opts.options.has(opts.alterOpt))
-        alterConfig(zkClient, opts, adminZkClient)
+        alterConfigWithZk(zkClient, opts, adminZkClient)
       else if (opts.options.has(opts.describeOpt))
-        describeConfig(zkClient, opts, adminZkClient)
-    } catch {
-      case e: Throwable =>
-        println("Error while executing config command " + e.getMessage)
-        println(Utils.stackTrace(e))
+        describeConfigWithZk(zkClient, opts, adminZkClient)
     } finally {
       zkClient.close()
     }
   }
 
-  private[admin] def alterConfig(zkClient: KafkaZkClient, opts: ConfigCommandOptions, adminZkClient: AdminZkClient) {
+  private[admin] def alterConfigWithZk(zkClient: KafkaZkClient, opts: ConfigCommandOptions, adminZkClient: AdminZkClient): Unit = {
     val configsToBeAdded = parseConfigsToBeAdded(opts)
     val configsToBeDeleted = parseConfigsToBeDeleted(opts)
     val entity = parseEntity(opts)
@@ -90,6 +135,26 @@ object ConfigCommand extends Config {
 
     if (entityType == ConfigType.User)
       preProcessScramCredentials(configsToBeAdded)
+    else if (entityType == ConfigType.Broker) {
+      // Replication quota configs may be updated using ZK at any time. Other dynamic broker configs
+      // may be updated using ZooKeeper only if the corresponding broker is not running. Dynamic broker
+      // configs at cluster-default level may be configured using ZK only if there are no brokers running.
+      val dynamicBrokerConfigs = configsToBeAdded.asScala.keySet.filterNot(BrokerConfigsUpdatableUsingZooKeeperWhileBrokerRunning.contains)
+      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 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 =>
+            require(zkClient.getBroker(brokerId).isEmpty, s"$errorMessage when broker $entityName is running. $info")
+          }
+        } else {
+          require(zkClient.getAllBrokersInCluster.isEmpty, s"$errorMessage for default cluster if any broker is running. $info")
+        }
+        preProcessBrokerConfigs(configsToBeAdded, perBrokerConfig)
+      }
+    }
 
     // compile the final set of configs
     val configs = adminZkClient.fetchEntityConfig(entityType, entityName)
@@ -97,17 +162,17 @@ object ConfigCommand extends Config {
     // fail the command if any of the configs to be deleted does not exist
     val invalidConfigs = configsToBeDeleted.filterNot(configs.containsKey(_))
     if (invalidConfigs.nonEmpty)
-      throw new InvalidConfigException(s"Invalid config(s): ${invalidConfigs.mkString(",")}")
+      throw new InvalidConfigurationException(s"Invalid config(s): ${invalidConfigs.mkString(",")}")
 
     configs ++= configsToBeAdded
     configsToBeDeleted.foreach(configs.remove(_))
 
     adminZkClient.changeConfigs(entityType, entityName, configs)
 
-    println(s"Completed Updating config for entity: $entity.")
+    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 {
@@ -128,7 +193,50 @@ object ConfigCommand extends Config {
     }
   }
 
-  private def describeConfig(zkClient: KafkaZkClient, opts: ConfigCommandOptions, adminZkClient: AdminZkClient) {
+  private[admin] def createPasswordEncoder(encoderConfigs: Map[String, String]): PasswordEncoder = {
+    encoderConfigs.get(KafkaConfig.PasswordEncoderSecretProp)
+    val encoderSecret = encoderConfigs.getOrElse(KafkaConfig.PasswordEncoderSecretProp,
+      throw new IllegalArgumentException("Password encoder secret not specified"))
+    new PasswordEncoder(new Password(encoderSecret),
+      None,
+      encoderConfigs.get(KafkaConfig.PasswordEncoderCipherAlgorithmProp).getOrElse(Defaults.PasswordEncoderCipherAlgorithm),
+      encoderConfigs.get(KafkaConfig.PasswordEncoderKeyLengthProp).map(_.toInt).getOrElse(Defaults.PasswordEncoderKeyLength),
+      encoderConfigs.get(KafkaConfig.PasswordEncoderIterationsProp).map(_.toInt).getOrElse(Defaults.PasswordEncoderIterations))
+  }
+
+  /**
+   * Pre-process broker configs provided to convert them to persistent format.
+   * 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): Unit = {
+    val passwordEncoderConfigs = new Properties
+    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.")
+      passwordEncoderConfigs.asScala.keySet.foreach(configsToBeAdded.remove)
+    }
+
+    DynamicBrokerConfig.validateConfigs(configsToBeAdded, perBrokerConfig)
+    val passwordConfigs = configsToBeAdded.asScala.keySet.filter(DynamicBrokerConfig.isPasswordConfig)
+    if (passwordConfigs.nonEmpty) {
+      require(passwordEncoderConfigs.containsKey(KafkaConfig.PasswordEncoderSecretProp),
+        s"${KafkaConfig.PasswordEncoderSecretProp} must be specified to update $passwordConfigs." +
+          " Other password encoder configs like cipher algorithm and iterations may also be specified" +
+          " to override the default encoding parameters. Password encoder configs will not be persisted" +
+          " in ZooKeeper."
+      )
+
+      val passwordEncoder = createPasswordEncoder(passwordEncoderConfigs.asScala)
+      passwordConfigs.foreach { configName =>
+        val encodedValue = passwordEncoder.encode(new Password(configsToBeAdded.getProperty(configName)))
+        configsToBeAdded.setProperty(configName, encodedValue)
+      }
+    }
+  }
+
+  private def describeConfigWithZk(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)
@@ -144,19 +252,25 @@ object ConfigCommand extends Config {
 
   private[admin] def parseConfigsToBeAdded(opts: ConfigCommandOptions): Properties = {
     val props = new Properties
+    if (opts.options.has(opts.addConfigFile)) {
+      val file = opts.options.valueOf(opts.addConfigFile)
+      props ++= Utils.loadProps(file)
+    }
     if (opts.options.has(opts.addConfig)) {
-      //split by commas, but avoid those in [], then into KV pairs
+      // Split list by commas, but avoid those in [], then into KV pairs
+      // Each KV pair is of format key=value, split them into key and value, using -1 as the limit for split() to
+      // include trailing empty strings. This is to support empty value (e.g. 'ssl.endpoint.identification.algorithm=')
       val pattern = "(?=[^\\]]*(?:\\[|$))"
       val configsToBeAdded = opts.options.valueOf(opts.addConfig)
         .split("," + pattern)
-        .map(_.split("""\s*=\s*""" + pattern))
+        .map(_.split("""\s*=\s*""" + pattern, -1))
       require(configsToBeAdded.forall(config => config.length == 2), "Invalid entity config: all configs to be added must be in the format \"key=val\".")
       //Create properties, parsing square brackets from values if necessary
       configsToBeAdded.foreach(pair => props.setProperty(pair(0).trim, pair(1).replaceAll("\\[?\\]?", "").trim))
-      if (props.containsKey(LogConfig.MessageFormatVersionProp)) {
-        println(s"WARNING: The configuration ${LogConfig.MessageFormatVersionProp}=${props.getProperty(LogConfig.MessageFormatVersionProp)} is specified. " +
-          s"This configuration will be ignored if the version is newer than the inter.broker.protocol.version specified in the broker.")
-      }
+    }
+    if (props.containsKey(LogConfig.MessageFormatVersionProp)) {
+      println(s"WARNING: The configuration ${LogConfig.MessageFormatVersionProp}=${props.getProperty(LogConfig.MessageFormatVersionProp)} is specified. " +
+        s"This configuration will be ignored if the version is newer than the inter.broker.protocol.version specified in the broker.")
     }
     props
   }
@@ -172,6 +286,344 @@ object ConfigCommand extends Config {
       Seq.empty
   }
 
+  private def processCommand(opts: ConfigCommandOptions): Unit = {
+    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))
+    val adminClient = Admin.create(props)
+
+    if (opts.options.has(opts.alterOpt) && opts.entityTypes.size != opts.entityNames.size)
+      throw new IllegalArgumentException(s"An entity name must be specified for every entity type")
+
+    try {
+      if (opts.options.has(opts.alterOpt))
+        alterConfig(adminClient, opts)
+      else if (opts.options.has(opts.describeOpt))
+        describeConfig(adminClient, opts)
+    } finally {
+      adminClient.close()
+    }
+  }
+
+  @nowarn("cat=deprecation")
+  private[admin] def alterConfig(adminClient: Admin, opts: ConfigCommandOptions): Unit = {
+    val entityTypes = opts.entityTypes
+    val entityNames = opts.entityNames
+    val entityTypeHead = entityTypes.head
+    val entityNameHead = entityNames.head
+    val configsToBeAddedMap = parseConfigsToBeAdded(opts).asScala.toMap // no need for mutability
+    val configsToBeAdded = configsToBeAddedMap.map { case (k, v) => (k, new ConfigEntry(k, v)) }
+    val configsToBeDeleted = parseConfigsToBeDeleted(opts)
+
+    entityTypeHead match {
+      case ConfigType.Topic =>
+        val oldConfig = getResourceConfig(adminClient, entityTypeHead, entityNameHead, includeSynonyms = false, describeAll = 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 configResource = new ConfigResource(ConfigResource.Type.TOPIC, entityNameHead)
+        val alterOptions = new AlterConfigsOptions().timeoutMs(30000).validateOnly(false)
+        val alterEntries = (configsToBeAdded.values.map(new AlterConfigOp(_, AlterConfigOp.OpType.SET))
+          ++ configsToBeDeleted.map { k => new AlterConfigOp(new ConfigEntry(k, ""), AlterConfigOp.OpType.DELETE) }
+        ).asJavaCollection
+        adminClient.incrementalAlterConfigs(Map(configResource -> alterEntries).asJava, alterOptions).all().get(60, TimeUnit.SECONDS)
+
+      case ConfigType.Broker =>
+        val oldConfig = getResourceConfig(adminClient, entityTypeHead, entityNameHead, includeSynonyms = false, describeAll = 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 configResource = new ConfigResource(ConfigResource.Type.BROKER, entityNameHead)
+        val alterOptions = new AlterConfigsOptions().timeoutMs(30000).validateOnly(false)
+        adminClient.alterConfigs(Map(configResource -> newConfig).asJava, alterOptions).all().get(60, TimeUnit.SECONDS)
+
+      case BrokerLoggerConfigType =>
+        val validLoggers = getResourceConfig(adminClient, entityTypeHead, entityNameHead, includeSynonyms = true, describeAll = false).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 configResource = new ConfigResource(ConfigResource.Type.BROKER_LOGGER, entityNameHead)
+        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)
+
+      case ConfigType.User | ConfigType.Client =>
+        val hasQuotaConfigsToAdd = configsToBeAdded.keys.exists(QuotaConfigs.isQuotaConfig)
+        val scramConfigsToAddMap = configsToBeAdded.filter(entry => ScramMechanism.isScram(entry._1))
+        val unknownConfigsToAdd = configsToBeAdded.keys.filterNot(key => ScramMechanism.isScram(key) || QuotaConfigs.isQuotaConfig(key))
+        val hasQuotaConfigsToDelete = configsToBeDeleted.exists(QuotaConfigs.isQuotaConfig)
+        val scramConfigsToDelete = configsToBeDeleted.filter(ScramMechanism.isScram)
+        val unknownConfigsToDelete = configsToBeDeleted.filterNot(key => ScramMechanism.isScram(key) || QuotaConfigs.isQuotaConfig(key))
+        if (entityTypeHead == ConfigType.Client || entityTypes.size == 2) { // size==2 for case where users is specified first on the command line, before clients
+          // either just a client or both a user and a client
+          if (unknownConfigsToAdd.nonEmpty || scramConfigsToAddMap.nonEmpty)
+            throw new IllegalArgumentException(s"Only quota configs can be added for '${ConfigType.Client}' using --bootstrap-server. Unexpected config names: ${unknownConfigsToAdd ++ scramConfigsToAddMap.keys}")
+          if (unknownConfigsToDelete.nonEmpty || scramConfigsToDelete.nonEmpty)
+            throw new IllegalArgumentException(s"Only quota configs can be deleted for '${ConfigType.Client}' using --bootstrap-server. Unexpected config names: ${unknownConfigsToDelete ++ scramConfigsToDelete}")
+        } else { // ConfigType.User
+          if (unknownConfigsToAdd.nonEmpty)
+            throw new IllegalArgumentException(s"Only quota and SCRAM credential configs can be added for '${ConfigType.User}' using --bootstrap-server. Unexpected config names: $unknownConfigsToAdd")
+          if (unknownConfigsToDelete.nonEmpty)
+            throw new IllegalArgumentException(s"Only quota and SCRAM credential configs can be deleted for '${ConfigType.User}' using --bootstrap-server. Unexpected config names: $unknownConfigsToDelete")
+          if (scramConfigsToAddMap.nonEmpty || scramConfigsToDelete.nonEmpty) {
+            if (entityNames.exists(_.isEmpty)) // either --entity-type users --entity-default or --user-defaults
+              throw new IllegalArgumentException("The use of --entity-default or --user-defaults is not allowed with User SCRAM Credentials using --bootstrap-server.")
+            if (hasQuotaConfigsToAdd || hasQuotaConfigsToDelete)
+              throw new IllegalArgumentException(s"Cannot alter both quota and SCRAM credential configs simultaneously for '${ConfigType.User}' using --bootstrap-server.")
+          }
+        }
+
+        if (hasQuotaConfigsToAdd || hasQuotaConfigsToDelete) {
+          alterQuotaConfigs(adminClient, entityTypes, entityNames, configsToBeAddedMap, configsToBeDeleted)
+        } else {
+          // handle altering user SCRAM credential configs
+          if (entityNames.size != 1)
+            // should never happen, if we get here then it is a bug
+            throw new IllegalStateException(s"Altering user SCRAM credentials should never occur for more zero or multiple users: $entityNames")
+          alterUserScramCredentialConfigs(adminClient, entityNames.head, scramConfigsToAddMap, scramConfigsToDelete)
+        }
+      case ConfigType.Ip =>
+        val unknownConfigs = (configsToBeAdded.keys ++ configsToBeDeleted).filterNot(key => DynamicConfig.Ip.names.contains(key))
+        if (unknownConfigs.nonEmpty)
+          throw new IllegalArgumentException(s"Only connection quota configs can be added for '${ConfigType.Ip}' using --bootstrap-server. Unexpected config names: ${unknownConfigs.mkString(",")}")
+        alterQuotaConfigs(adminClient, entityTypes, entityNames, configsToBeAddedMap, configsToBeDeleted)
+      case _ => throw new IllegalArgumentException(s"Unsupported entity type: $entityTypeHead")
+    }
+
+    if (entityNameHead.nonEmpty)
+      println(s"Completed updating config for ${entityTypeHead.dropRight(1)} $entityNameHead.")
+    else
+      println(s"Completed updating default config for $entityTypeHead in the cluster.")
+  }
+
+  private def alterUserScramCredentialConfigs(adminClient: Admin, user: String, scramConfigsToAddMap: Map[String, ConfigEntry], scramConfigsToDelete: Seq[String]) = {
+    val deletions = scramConfigsToDelete.map(mechanismName =>
+      new UserScramCredentialDeletion(user, PublicScramMechanism.fromMechanismName(mechanismName)))
+
+    def iterationsAndPasswordBytes(mechanism: ScramMechanism, credentialStr: String): (Integer, Array[Byte]) = {
+      val pattern = "(?:iterations=(\\-?[0-9]*),)?password=(.*)".r
+      val (iterations, password) = credentialStr match {
+        case pattern(iterations, password) => (if (iterations != null && iterations != "-1") iterations.toInt else DefaultScramIterations, password)
+        case _ => throw new IllegalArgumentException(s"Invalid credential property $mechanism=$credentialStr")
+      }
+      if (iterations < mechanism.minIterations)
+        throw new IllegalArgumentException(s"Iterations $iterations is less than the minimum ${mechanism.minIterations} required for ${mechanism.mechanismName}")
+      (iterations, password.getBytes(StandardCharsets.UTF_8))
+    }
+
+    val upsertions = scramConfigsToAddMap.map { case (mechanismName, configEntry) =>
+      val (iterations, passwordBytes) = iterationsAndPasswordBytes(ScramMechanism.forMechanismName(mechanismName), configEntry.value)
+      new UserScramCredentialUpsertion(user, new ScramCredentialInfo(PublicScramMechanism.fromMechanismName(mechanismName), iterations), passwordBytes)
+    }
+    // we are altering only a single user by definition, so we don't have to worry about one user succeeding and another
+    // failing; therefore just check the success of all the futures (since there will only be 1)
+    adminClient.alterUserScramCredentials((deletions ++ upsertions).toList.asJava).all.get(60, TimeUnit.SECONDS)
+  }
+
+  private def alterQuotaConfigs(adminClient: Admin, entityTypes: List[String], entityNames: List[String], configsToBeAddedMap: Map[String, String], configsToBeDeleted: Seq[String]) = {
+    // handle altering client/user quota configs
+    val oldConfig = getClientQuotasConfig(adminClient, entityTypes, entityNames)
+
+    val invalidConfigs = configsToBeDeleted.filterNot(oldConfig.contains)
+    if (invalidConfigs.nonEmpty)
+      throw new InvalidConfigurationException(s"Invalid config(s): ${invalidConfigs.mkString(",")}")
+
+    val alterEntityTypes = entityTypes.map {
+      case ConfigType.User => ClientQuotaEntity.USER
+      case ConfigType.Client => ClientQuotaEntity.CLIENT_ID
+      case ConfigType.Ip => ClientQuotaEntity.IP
+      case entType => throw new IllegalArgumentException(s"Unexpected entity type: $entType")
+    }
+    val alterEntityNames = entityNames.map(en => if (en.nonEmpty) en else null)
+
+    // Explicitly populate a HashMap to ensure nulls are recorded properly.
+    val alterEntityMap = new java.util.HashMap[String, String]
+    alterEntityTypes.zip(alterEntityNames).foreach { case (k, v) => alterEntityMap.put(k, v) }
+    val entity = new ClientQuotaEntity(alterEntityMap)
+
+    val alterOptions = new AlterClientQuotasOptions().validateOnly(false)
+    val alterOps = (configsToBeAddedMap.map { case (key, value) =>
+      val doubleValue = try value.toDouble catch {
+        case _: NumberFormatException =>
+          throw new IllegalArgumentException(s"Cannot parse quota configuration value for $key: $value")
+      }
+      new ClientQuotaAlteration.Op(key, doubleValue)
+    } ++ configsToBeDeleted.map(key => new ClientQuotaAlteration.Op(key, null))).asJavaCollection
+
+    adminClient.alterClientQuotas(Collections.singleton(new ClientQuotaAlteration(entity, alterOps)), alterOptions)
+      .all().get(60, TimeUnit.SECONDS)
+  }
+
+  private[admin] def describeConfig(adminClient: Admin, opts: ConfigCommandOptions): Unit = {
+    val entityTypes = opts.entityTypes
+    val entityNames = opts.entityNames
+    val describeAll = opts.options.has(opts.allOpt)
+
+    entityTypes.head match {
+      case ConfigType.Topic | ConfigType.Broker | BrokerLoggerConfigType =>
+        describeResourceConfig(adminClient, entityTypes.head, entityNames.headOption, describeAll)
+      case ConfigType.User | ConfigType.Client =>
+        describeClientQuotaAndUserScramCredentialConfigs(adminClient, entityTypes, entityNames)
+      case ConfigType.Ip =>
+        describeQuotaConfigs(adminClient, entityTypes, entityNames)
+      case entityType => throw new IllegalArgumentException(s"Invalid entity type: $entityType")
+    }
+  }
+
+  private def describeResourceConfig(adminClient: Admin, entityType: String, entityName: Option[String], describeAll: Boolean): Unit = {
+    val entities = entityName
+      .map(name => List(name))
+      .getOrElse(entityType match {
+        case ConfigType.Topic =>
+          adminClient.listTopics(new ListTopicsOptions().listInternal(true)).names().get().asScala.toSeq
+        case ConfigType.Broker | BrokerLoggerConfigType =>
+          adminClient.describeCluster(new DescribeClusterOptions()).nodes().get().asScala.map(_.idString).toSeq :+ BrokerDefaultEntityName
+        case entityType => throw new IllegalArgumentException(s"Invalid entity type: $entityType")
+      })
+
+    entities.foreach { entity =>
+      entity match {
+        case BrokerDefaultEntityName =>
+          println(s"Default configs for $entityType in the cluster are:")
+        case _ =>
+          val configSourceStr = if (describeAll) "All" else "Dynamic"
+          println(s"$configSourceStr configs for ${entityType.dropRight(1)} $entity are:")
+      }
+      getResourceConfig(adminClient, entityType, entity, includeSynonyms = true, describeAll).foreach { entry =>
+        val synonyms = entry.synonyms.asScala.map(synonym => s"${synonym.source}:${synonym.name}=${synonym.value}").mkString(", ")
+        println(s"  ${entry.name}=${entry.value} sensitive=${entry.isSensitive} synonyms={$synonyms}")
+      }
+    }
+  }
+
+  private def getResourceConfig(adminClient: Admin, entityType: String, entityName: String, includeSynonyms: Boolean, describeAll: Boolean) = {
+    def validateBrokerId(): Unit = try entityName.toInt catch {
+      case _: NumberFormatException =>
+        throw new IllegalArgumentException(s"The entity name for $entityType must be a valid integer broker id, found: $entityName")
+    }
+
+    val (configResourceType, dynamicConfigSource) = entityType match {
+      case ConfigType.Topic =>
+        if (!entityName.isEmpty)
+          Topic.validate(entityName)
+        (ConfigResource.Type.TOPIC, Some(ConfigEntry.ConfigSource.DYNAMIC_TOPIC_CONFIG))
+      case ConfigType.Broker => entityName match {
+        case BrokerDefaultEntityName =>
+          (ConfigResource.Type.BROKER, Some(ConfigEntry.ConfigSource.DYNAMIC_DEFAULT_BROKER_CONFIG))
+        case _ =>
+          validateBrokerId()
+          (ConfigResource.Type.BROKER, Some(ConfigEntry.ConfigSource.DYNAMIC_BROKER_CONFIG))
+      }
+      case BrokerLoggerConfigType =>
+        if (!entityName.isEmpty)
+          validateBrokerId()
+        (ConfigResource.Type.BROKER_LOGGER, None)
+      case entityType => throw new IllegalArgumentException(s"Invalid entity type: $entityType")
+    }
+
+    val configSourceFilter = if (describeAll)
+      None
+    else
+      dynamicConfigSource
+
+    val configResource = new ConfigResource(configResourceType, entityName)
+    val describeOptions = new DescribeConfigsOptions().includeSynonyms(includeSynonyms)
+    val configs = adminClient.describeConfigs(Collections.singleton(configResource), describeOptions)
+      .all.get(30, TimeUnit.SECONDS)
+    configs.get(configResource).entries.asScala
+      .filter(entry => configSourceFilter match {
+        case Some(configSource) => entry.source == configSource
+        case None => true
+      }).toSeq
+  }
+
+  private def describeQuotaConfigs(adminClient: Admin, entityTypes: List[String], entityNames: List[String]) = {
+    val quotaConfigs = getAllClientQuotasConfigs(adminClient, entityTypes, entityNames)
+    quotaConfigs.forKeyValue { (entity, entries) =>
+      val entityEntries = entity.entries.asScala
+
+      def entitySubstr(entityType: String): Option[String] =
+        entityEntries.get(entityType).map { name =>
+          val typeStr = entityType match {
+            case ClientQuotaEntity.USER => "user-principal"
+            case ClientQuotaEntity.CLIENT_ID => "client-id"
+            case ClientQuotaEntity.IP => "ip"
+          }
+          if (name != null) s"$typeStr '$name'"
+          else s"the default $typeStr"
+        }
+
+      val entityStr = (entitySubstr(ClientQuotaEntity.USER) ++
+                       entitySubstr(ClientQuotaEntity.CLIENT_ID) ++
+                       entitySubstr(ClientQuotaEntity.IP)).mkString(", ")
+      val entriesStr = entries.asScala.map(e => s"${e._1}=${e._2}").mkString(", ")
+      println(s"Quota configs for $entityStr are $entriesStr")
+    }
+  }
+
+  private def describeClientQuotaAndUserScramCredentialConfigs(adminClient: Admin, entityTypes: List[String], entityNames: List[String]) = {
+    describeQuotaConfigs(adminClient, entityTypes, entityNames)
+    // we describe user SCRAM credentials only when we are not describing client information
+    // and we are not given either --entity-default or --user-defaults
+    if (!entityTypes.contains(ConfigType.Client) && !entityNames.contains("")) {
+      val result = adminClient.describeUserScramCredentials(entityNames.asJava)
+      result.users.get(30, TimeUnit.SECONDS).asScala.foreach(user => {
+        try {
+          val description = result.description(user).get(30, TimeUnit.SECONDS)
+          val descriptionText = description.credentialInfos.asScala.map(info => s"${info.mechanism.mechanismName}=iterations=${info.iterations}").mkString(", ")
+          println(s"SCRAM credential configs for user-principal '$user' are $descriptionText")
+        } catch {
+          case e: Exception => println(s"Error retrieving SCRAM credential configs for user-principal '$user': ${e.getClass.getSimpleName}: ${e.getMessage}")
+        }
+      })
+    }
+  }
+
+  private def getClientQuotasConfig(adminClient: Admin, entityTypes: List[String], entityNames: List[String]): Map[String, java.lang.Double] = {
+    if (entityTypes.size != entityNames.size)
+      throw new IllegalArgumentException("Exactly one entity name must be specified for every entity type")
+    getAllClientQuotasConfigs(adminClient, entityTypes, entityNames).headOption.map(_._2.asScala).getOrElse(Map.empty)
+  }
+
+  private def getAllClientQuotasConfigs(adminClient: Admin, entityTypes: List[String], entityNames: List[String]) = {
+    val components = entityTypes.map(Some(_)).zipAll(entityNames.map(Some(_)), None, None).map { case (entityTypeOpt, entityNameOpt) =>
+      val entityType = entityTypeOpt match {
+        case Some(ConfigType.User) => ClientQuotaEntity.USER
+        case Some(ConfigType.Client) => ClientQuotaEntity.CLIENT_ID
+        case Some(ConfigType.Ip) => ClientQuotaEntity.IP
+        case Some(_) => throw new IllegalArgumentException(s"Unexpected entity type ${entityTypeOpt.get}")
+        case None => throw new IllegalArgumentException("More entity names specified than entity types")
+      }
+      entityNameOpt match {
+        case Some("") => ClientQuotaFilterComponent.ofDefaultEntity(entityType)
+        case Some(name) => ClientQuotaFilterComponent.ofEntity(entityType, name)
+        case None => ClientQuotaFilterComponent.ofEntityType(entityType)
+      }
+    }
+
+    adminClient.describeClientQuotas(ClientQuotaFilter.containsOnly(components.asJava)).entities.get(30, TimeUnit.SECONDS).asScala
+  }
+
   case class Entity(entityType: String, sanitizedName: Option[String]) {
     val entityPath = sanitizedName match {
       case Some(n) => entityType + "/" + n
@@ -239,27 +691,25 @@ object ConfigCommand extends Config {
   }
 
   private[admin] def parseEntity(opts: ConfigCommandOptions): ConfigEntity = {
-    val entityTypes = opts.options.valuesOf(opts.entityType).asScala
+    val entityTypes = opts.entityTypes
+    val entityNames = opts.entityNames
     if (entityTypes.head == ConfigType.User || entityTypes.head == ConfigType.Client)
-      parseQuotaEntity(opts)
+      parseClientQuotaEntity(opts, entityTypes, entityNames)
     else {
       // Exactly one entity type and at-most one entity name expected for other entities
-      val name = if (opts.options.has(opts.entityName)) Some(opts.options.valueOf(opts.entityName)) else None
+      val name = entityNames.headOption match {
+        case Some("") => Some(ConfigEntityName.Default)
+        case v => v
+      }
       ConfigEntity(Entity(entityTypes.head, name), None)
     }
   }
 
-  private def parseQuotaEntity(opts: ConfigCommandOptions): ConfigEntity = {
-    val types = opts.options.valuesOf(opts.entityType).asScala
-    val namesIterator = opts.options.valuesOf(opts.entityName).iterator
-    val names = opts.options.specs.asScala
-                    .filter(spec => spec.options.contains("entity-name") || spec.options.contains("entity-default"))
-                    .map(spec => if (spec.options.contains("entity-name")) namesIterator.next else "")
-
+  private def parseClientQuotaEntity(opts: ConfigCommandOptions, types: List[String], names: List[String]): ConfigEntity = {
     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
 
@@ -274,77 +724,197 @@ object ConfigCommand extends Config {
       }
     }
 
-    val entities = entityTypes.map(t => Entity(t, if (sortedNames.hasNext) Some(sanitizeName(t, sortedNames.next)) else None))
+    val entities = entityTypes.map(t => Entity(t, if (sortedNames.hasNext) Some(sanitizeName(t, sortedNames.next())) else None))
     ConfigEntity(entities.head, if (entities.size > 1) Some(entities(1)) else None)
   }
 
-  class ConfigCommandOptions(args: Array[String]) {
-    val parser = new OptionParser(false)
-    val zkConnectOpt = parser.accepts("zookeeper", "REQUIRED: The connection string for the zookeeper connection in the form host:port. " +
-            "Multiple URLS can be given to allow fail-over.")
+  class ConfigCommandOptions(args: Array[String]) extends CommandDefaultOptions(args) {
+
+    val zkConnectOpt = parser.accepts("zookeeper", "DEPRECATED. The connection string for the zookeeper connection in the form host:port. " +
+            "Multiple URLS can be given to allow fail-over. Replaced by --bootstrap-server, REQUIRED unless --bootstrap-server is given.")
             .withRequiredArg
             .describedAs("urls")
             .ofType(classOf[String])
+    val bootstrapServerOpt = parser.accepts("bootstrap-server", "The Kafka server to connect to. " +
+      "This is required for describing and altering broker configs.")
+      .withRequiredArg
+      .describedAs("server to connect to")
+      .ofType(classOf[String])
+    val commandConfigOpt = parser.accepts("command-config", "Property file containing configs to be passed to Admin Client. " +
+      "This is used only with --bootstrap-server option for describing and altering broker configs.")
+      .withRequiredArg
+      .describedAs("command config property file")
+      .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 allOpt = parser.accepts("all", "List all configs for the given topic, broker, or broker-logger entity (includes static configuration when the entity type is brokers)")
+
+    val entityType = parser.accepts("entity-type", "Type of entity (topics/clients/users/brokers/broker-loggers/ips)")
             .withRequiredArg
             .ofType(classOf[String])
-    val entityName = parser.accepts("entity-name", "Name of entity (topic name/client id/user principal name/broker id)")
+    val entityName = parser.accepts("entity-name", "Name of entity (topic name/client id/user principal name/broker id/ip)")
             .withRequiredArg
             .ofType(classOf[String])
-    val entityDefault = parser.accepts("entity-default", "Default entity name for clients/users (applies to corresponding entity type in command line)")
+    val entityDefault = parser.accepts("entity-default", "Default entity name for clients/users/brokers/ips (applies to corresponding entity type in command line)")
 
     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.Topic + "': " + LogConfig.configNames.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) +
+            "For entity-type '" + ConfigType.Ip + "': " + DynamicConfig.Ip.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])
+    val addConfigFile = parser.accepts("add-config-file", "Path to a properties file with configs to add. See add-config for a list of valid configurations.")
+            .withRequiredArg
+            .ofType(classOf[String])
     val deleteConfig = parser.accepts("delete-config", "config keys to remove 'k1,k2'")
             .withRequiredArg
             .ofType(classOf[String])
             .withValuesSeparatedBy(',')
-    val helpOpt = parser.accepts("help", "Print usage information.")
     val forceOpt = parser.accepts("force", "Suppress console prompts")
-    val options = parser.parse(args : _*)
+    val topic = parser.accepts("topic", "The topic's name.")
+      .withRequiredArg
+      .ofType(classOf[String])
+    val client = parser.accepts("client", "The client's ID.")
+      .withRequiredArg
+      .ofType(classOf[String])
+    val clientDefaults = parser.accepts("client-defaults", "The config defaults for all clients.")
+    val user = parser.accepts("user", "The user's principal name.")
+      .withRequiredArg
+      .ofType(classOf[String])
+    val userDefaults = parser.accepts("user-defaults", "The config defaults for all users.")
+    val broker = parser.accepts("broker", "The broker's ID.")
+      .withRequiredArg
+      .ofType(classOf[String])
+    val brokerDefaults = parser.accepts("broker-defaults", "The config defaults for all brokers.")
+    val brokerLogger = parser.accepts("broker-logger", "The broker's ID for its logger config.")
+      .withRequiredArg
+      .ofType(classOf[String])
+    val ipDefaults = parser.accepts("ip-defaults", "The config defaults for all IPs.")
+    val ip = parser.accepts("ip", "The IP address.")
+      .withRequiredArg
+      .ofType(classOf[String])
+    val zkTlsConfigFile = parser.accepts("zk-tls-config-file",
+      "Identifies the file where ZooKeeper client TLS connectivity properties are defined.  Any properties other than " +
+        KafkaConfig.ZkSslConfigToSystemPropertyMap.keys.toList.sorted.mkString(", ") + " are ignored.")
+      .withRequiredArg().describedAs("ZooKeeper TLS configuration").ofType(classOf[String])
+    options = parser.parse(args : _*)
+
+    private val entityFlags = List((topic, ConfigType.Topic),
+      (client, ConfigType.Client),
+      (user, ConfigType.User),
+      (broker, ConfigType.Broker),
+      (brokerLogger, BrokerLoggerConfigType),
+      (ip, ConfigType.Ip))
+
+    private val entityDefaultsFlags = List((clientDefaults, ConfigType.Client),
+      (userDefaults, ConfigType.User),
+      (brokerDefaults, ConfigType.Broker),
+      (ipDefaults, ConfigType.Ip))
+
+    private[admin] def entityTypes: List[String] = {
+      options.valuesOf(entityType).asScala.toList ++
+        (entityFlags ++ entityDefaultsFlags).filter(entity => options.has(entity._1)).map(_._2)
+    }
 
-    val allOpts: Set[OptionSpec[_]] = Set(alterOpt, describeOpt, entityType, entityName, addConfig, deleteConfig, helpOpt)
+    private[admin] def entityNames: List[String] = {
+      val namesIterator = options.valuesOf(entityName).iterator
+      options.specs.asScala
+        .filter(spec => spec.options.contains("entity-name") || spec.options.contains("entity-default"))
+        .map(spec => if (spec.options.contains("entity-name")) namesIterator.next else "").toList ++
+      entityFlags
+        .filter(entity => options.has(entity._1))
+        .map(entity => options.valueOf(entity._1)) ++
+      entityDefaultsFlags
+        .filter(entity => options.has(entity._1))
+        .map(_ => "")
+    }
 
-    def checkArgs() {
+    def checkArgs(): Unit = {
       // should have exactly one action
       val actions = Seq(alterOpt, describeOpt).count(options.has _)
-      if(actions != 1)
+      if (actions != 1)
         CommandLineUtils.printUsageAndDie(parser, "Command must include exactly one action: --describe, --alter")
-
       // check required args
-      CommandLineUtils.checkRequiredArgs(parser, options, zkConnectOpt, entityType)
       CommandLineUtils.checkInvalidArgs(parser, options, alterOpt, Set(describeOpt))
       CommandLineUtils.checkInvalidArgs(parser, options, describeOpt, Set(alterOpt, addConfig, deleteConfig))
-      val entityTypeVals = options.valuesOf(entityType).asScala
-      if(options.has(alterOpt)) {
-        if (entityTypeVals.contains(ConfigType.User) || entityTypeVals.contains(ConfigType.Client)) {
-          if (!options.has(entityName) && !options.has(entityDefault))
-            throw new IllegalArgumentException("--entity-name or --entity-default must be specified with --alter of users/clients")
-        } else if (!options.has(entityName))
-            throw new IllegalArgumentException(s"--entity-name must be specified with --alter of ${entityTypeVals}")
-
-        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")
-      }
+
+      val entityTypeVals = entityTypes
+      if (entityTypeVals.size != entityTypeVals.distinct.size)
+        throw new IllegalArgumentException(s"Duplicate entity type(s) specified: ${entityTypeVals.diff(entityTypeVals.distinct).mkString(",")}")
+
+      val (allowedEntityTypes, connectOptString) = if (options.has(bootstrapServerOpt))
+        (BrokerSupportedConfigTypes, "--bootstrap-server")
+      else
+        (ZkSupportedConfigTypes, "--zookeeper")
+
       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 (!allowedEntityTypes.contains(entityTypeVal))
+          throw new IllegalArgumentException(s"Invalid entity type $entityTypeVal, the 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")
+        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(entityName) || options.has(entityType) || options.has(entityDefault)) &&
+        (entityFlags ++ entityDefaultsFlags).exists(entity => options.has(entity._1)))
+        throw new IllegalArgumentException("--entity-{type,name,default} should not be used in conjunction with specific entity flags")
+
+      val hasEntityName = entityNames.exists(!_.isEmpty)
+      val hasEntityDefault = entityNames.exists(_.isEmpty)
+
+      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")
+
+      if (options.has(allOpt) && options.has(zkConnectOpt)) {
+        throw new IllegalArgumentException(s"--bootstrap-server must be specified for --all")
+      }
+
+      if (hasEntityName && (entityTypeVals.contains(ConfigType.Broker) || entityTypeVals.contains(BrokerLoggerConfigType))) {
+        Seq(entityName, broker, brokerLogger).filter(options.has(_)).map(options.valueOf(_)).foreach { brokerId =>
+          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 (hasEntityName && entityTypeVals.contains(ConfigType.Ip)) {
+        Seq(entityName, ip).filter(options.has(_)).map(options.valueOf(_)).foreach { ipEntity =>
+          if (!DynamicConfig.Ip.isValidIpEntity(ipEntity))
+            throw new IllegalArgumentException(s"The entity name for ${entityTypeVals.head} must be a valid IP or resolvable host, but it is: $ipEntity")
+        }
+      }
+
+      if (options.has(describeOpt) && entityTypeVals.contains(BrokerLoggerConfigType) && !hasEntityName)
+        throw new IllegalArgumentException(s"an 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) ||
+            entityTypeVals.contains(ConfigType.Ip)) {
+          if (!hasEntityName && !hasEntityDefault)
+            throw new IllegalArgumentException("an entity-name or default entity must be specified with --alter of users, clients, brokers or ips")
+        } else if (!hasEntityName)
+          throw new IllegalArgumentException(s"an entity name must be specified with --alter of ${entityTypeVals.mkString(",")}")
+
+        val isAddConfigPresent = options.has(addConfig)
+        val isAddConfigFilePresent = options.has(addConfigFile)
+        val isDeleteConfigPresent = options.has(deleteConfig)
+
+        if(isAddConfigPresent && isAddConfigFilePresent)
+          throw new IllegalArgumentException("Only one of --add-config or --add-config-file must be specified")
+
+        if(!isAddConfigPresent && !isAddConfigFilePresent && !isDeleteConfigPresent)
+          throw new IllegalArgumentException("At least one of --add-config, --add-config-file, or --delete-config must be specified with --alter")
+      }
     }
   }
 }
diff --git a/core/src/main/scala/kafka/admin/ConsumerGroupCommand.scala b/core/src/main/scala/kafka/admin/ConsumerGroupCommand.scala
index d71f062e3b407..d8a01abe1cbf2 100755
--- a/core/src/main/scala/kafka/admin/ConsumerGroupCommand.scala
+++ b/core/src/main/scala/kafka/admin/ConsumerGroupCommand.scala
@@ -17,104 +17,77 @@
 
 package kafka.admin
 
-import java.text.{ParseException, SimpleDateFormat}
-import java.util.{Date, Properties}
-import javax.xml.datatype.DatatypeFactory
-
-import joptsimple.{OptionParser, OptionSpec}
-import kafka.api.{OffsetFetchRequest, OffsetFetchResponse, OffsetRequest, PartitionOffsetRequestInfo}
-import kafka.client.ClientUtils
-import kafka.common.{OffsetMetadataAndError, TopicAndPartition}
-import kafka.utils.Implicits._
-import kafka.consumer.SimpleConsumer
+import java.time.{Duration, Instant}
+import java.util.Properties
+
+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.I0Itec.zkclient.exception.ZkNoNodeException
+import kafka.utils.Implicits._
+import org.apache.kafka.clients.admin._
+import org.apache.kafka.clients.consumer.OffsetAndMetadata
 import org.apache.kafka.clients.CommonClientConfigs
-import org.apache.kafka.clients.consumer.{ConsumerConfig, KafkaConsumer, OffsetAndMetadata}
-import org.apache.kafka.common.errors.BrokerNotAvailableException
+import org.apache.kafka.common.utils.Utils
 import org.apache.kafka.common.{KafkaException, Node, TopicPartition}
-import org.apache.kafka.common.internals.Topic
-import org.apache.kafka.common.network.ListenerName
+
+import scala.jdk.CollectionConverters._
+import scala.collection.mutable.ListBuffer
+import scala.collection.{Map, Seq, immutable, mutable}
+import scala.util.{Failure, Success, Try}
+import joptsimple.OptionSpec
 import org.apache.kafka.common.protocol.Errors
-import org.apache.kafka.common.security.JaasUtils
-import org.apache.kafka.common.security.auth.SecurityProtocol
-import org.apache.kafka.common.serialization.StringDeserializer
-import org.apache.kafka.common.utils.Utils
 
-import scala.collection.JavaConverters._
-import scala.collection.{Seq, Set, mutable}
+import scala.collection.immutable.TreeMap
+import scala.reflect.ClassTag
+import org.apache.kafka.common.requests.ListOffsetResponse
+import org.apache.kafka.common.ConsumerGroupState
+import joptsimple.OptionException
 
 object ConsumerGroupCommand extends Logging {
 
-  def main(args: Array[String]) {
-    val opts = new ConsumerGroupCommandOptions(args)
-
-    if (args.length == 0)
-      CommandLineUtils.printUsageAndDie(opts.parser, "List all consumer groups, describe a consumer group, delete consumer group info, or reset consumer group offsets.")
+  def main(args: Array[String]): Unit = {
 
-    // should have exactly one action
-    val actions = Seq(opts.listOpt, opts.describeOpt, opts.deleteOpt, opts.resetOffsetsOpt).count(opts.options.has _)
-    if (actions != 1)
-      CommandLineUtils.printUsageAndDie(opts.parser, "Command must include exactly one action: --list, --describe, --delete, --reset-offset")
+    val opts = new ConsumerGroupCommandOptions(args)
+    try {
+      opts.checkArgs()
+      CommandLineUtils.printHelpAndExitIfNeeded(opts, "This tool helps to list all consumer groups, describe a consumer group, delete consumer group info, or reset consumer group offsets.")
 
-    opts.checkArgs()
+      // should have exactly one action
+      val actions = Seq(opts.listOpt, opts.describeOpt, opts.deleteOpt, opts.resetOffsetsOpt, opts.deleteOffsetsOpt).count(opts.options.has)
+      if (actions != 1)
+        CommandLineUtils.printUsageAndDie(opts.parser, "Command must include exactly one action: --list, --describe, --delete, --reset-offsets, --delete-offsets")
 
-    val consumerGroupService = {
-      if (opts.useOldConsumer) {
-        System.err.println("Note: This will only show information about consumers that use ZooKeeper (not those using the Java consumer API).\n")
-        new ZkConsumerGroupService(opts)
-      } else {
-        System.err.println("Note: This will not show information about old Zookeeper-based consumers.\n")
-        new KafkaConsumerGroupService(opts)
-      }
+      run(opts)
+    } catch {
+      case e: OptionException =>
+        CommandLineUtils.printUsageAndDie(opts.parser, e.getMessage)
     }
+  }
 
+  def run(opts: ConsumerGroupCommandOptions): Unit = {
+    val consumerGroupService = new ConsumerGroupService(opts)
     try {
       if (opts.options.has(opts.listOpt))
-        consumerGroupService.listGroups().foreach(println(_))
-      else if (opts.options.has(opts.describeOpt)) {
-        val (state, assignments) = consumerGroupService.describeGroup()
-        val groupId = opts.options.valuesOf(opts.groupOpt).asScala.head
-        assignments match {
-          case None =>
-            // applies to both old and new consumer
-            printError(s"The consumer group '$groupId' does not exist.")
-          case Some(assignments) =>
-            if (opts.useOldConsumer)
-              printAssignment(assignments, false)
-            else
-              state match {
-                case Some("Dead") =>
-                  printError(s"Consumer group '$groupId' does not exist.")
-                case Some("Empty") =>
-                  System.err.println(s"Consumer group '$groupId' has no active members.")
-                  printAssignment(assignments, true)
-                case Some("PreparingRebalance") | Some("CompletingRebalance") =>
-                  System.err.println(s"Warning: Consumer group '$groupId' is rebalancing.")
-                  printAssignment(assignments, true)
-                case Some("Stable") =>
-                  printAssignment(assignments, true)
-                case other =>
-                  // the control should never reach here
-                  throw new KafkaException(s"Expected a valid consumer group state, but found '${other.getOrElse("NONE")}'.")
-              }
-        }
-      }
-      else if (opts.options.has(opts.deleteOpt)) {
-        consumerGroupService match {
-          case service: ZkConsumerGroupService => service.deleteGroups()
-          case _ => throw new IllegalStateException(s"delete is not supported for $consumerGroupService.")
-        }
-      }
+        consumerGroupService.listGroups()
+      else if (opts.options.has(opts.describeOpt))
+        consumerGroupService.describeGroups()
+      else if (opts.options.has(opts.deleteOpt))
+        consumerGroupService.deleteGroups()
       else if (opts.options.has(opts.resetOffsetsOpt)) {
         val offsetsToReset = consumerGroupService.resetOffsets()
         if (opts.options.has(opts.exportOpt)) {
-          val exported = consumerGroupService.exportOffsetsToReset(offsetsToReset)
+          val exported = consumerGroupService.exportOffsetsToCsv(offsetsToReset)
           println(exported)
         } else
           printOffsetsToReset(offsetsToReset)
       }
+      else if (opts.options.has(opts.deleteOffsetsOpt)) {
+        consumerGroupService.deleteOffsets()
+      }
     } catch {
+      case e: IllegalArgumentException =>
+        CommandLineUtils.printUsageAndDie(opts.parser, e.getMessage)
       case e: Throwable =>
         printError(s"Executing consumer group command failed due to ${e.getMessage}", Some(e))
     } finally {
@@ -122,439 +95,639 @@ object ConsumerGroupCommand extends Logging {
     }
   }
 
+  def consumerGroupStatesFromString(input: String): Set[ConsumerGroupState] = {
+    val parsedStates = input.split(',').map(s => ConsumerGroupState.parse(s.trim)).toSet
+    if (parsedStates.contains(ConsumerGroupState.UNKNOWN)) {
+      val validStates = ConsumerGroupState.values().filter(_ != ConsumerGroupState.UNKNOWN)
+      throw new IllegalArgumentException(s"Invalid state list '$input'. Valid states are: ${validStates.mkString(", ")}")
+    }
+    parsedStates
+  }
+
   val MISSING_COLUMN_VALUE = "-"
 
   def printError(msg: String, e: Option[Throwable] = None): Unit = {
-    println(s"Error: $msg")
-    e.foreach(debug("Exception in consumer group command", _))
+    println(s"\nError: $msg")
+    e.foreach(_.printStackTrace())
   }
 
-  def printAssignment(groupAssignment: Seq[PartitionAssignmentState], useNewConsumer: Boolean): Unit = {
-    print("\n%-30s %-10s %-15s %-15s %-10s %-50s".format("TOPIC", "PARTITION", "CURRENT-OFFSET", "LOG-END-OFFSET", "LAG", "CONSUMER-ID"))
-    if (useNewConsumer)
-      print("%-30s %s".format("HOST", "CLIENT-ID"))
-    println()
-
-    groupAssignment.foreach { consumerAssignment =>
-      print("%-30s %-10s %-15s %-15s %-10s %-50s".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)))
-      if (useNewConsumer)
-        print("%-30s %s".format(consumerAssignment.host.getOrElse(MISSING_COLUMN_VALUE), consumerAssignment.clientId.getOrElse(MISSING_COLUMN_VALUE)))
-      println()
+  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))
     }
   }
 
-  def printOffsetsToReset(groupAssignmentsToReset: Map[TopicPartition, OffsetAndMetadata]): Unit = {
-    print("\n%-30s %-10s %-15s".format("TOPIC", "PARTITION", "NEW-OFFSET"))
-    println()
-
-    groupAssignmentsToReset.foreach {
-      case (consumerAssignment, offsetAndMetadata) =>
-        print("%-30s %-10s %-15s".format(
-          consumerAssignment.topic(),
-          consumerAssignment.partition(),
-          offsetAndMetadata.offset()))
-        println()
-    }
-  }
-
-  protected case class PartitionAssignmentState(group: String, coordinator: Option[Node], topic: Option[String],
+  private[admin] case class PartitionAssignmentState(group: String, coordinator: Option[Node], topic: Option[String],
                                                 partition: Option[Int], offset: Option[Long], lag: Option[Long],
                                                 consumerId: Option[String], host: Option[String],
                                                 clientId: Option[String], logEndOffset: Option[Long])
 
-  sealed trait ConsumerGroupService {
+  private[admin] case class MemberAssignmentState(group: String, consumerId: String, host: String, clientId: String, groupInstanceId: String,
+                                             numPartitions: Int, assignment: List[TopicPartition])
 
-    def listGroups(): List[String]
+  private[admin] case class GroupState(group: String, coordinator: Node, assignmentStrategy: String, state: String, numMembers: Int)
 
-    def describeGroup(): (Option[String], Option[Seq[PartitionAssignmentState]]) = {
-      collectGroupAssignment(opts.options.valueOf(opts.groupOpt))
+  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 close(): Unit
-
-    protected def opts: ConsumerGroupCommandOptions
-
-    protected def getLogEndOffset(topicPartition: TopicPartition): LogOffsetResult
-
-    protected def collectGroupAssignment(group: String): (Option[String], Option[Seq[PartitionAssignmentState]])
-
-    protected def collectConsumerAssignment(group: String,
-                                            coordinator: Option[Node],
-                                            topicPartitions: Seq[TopicAndPartition],
-                                            getPartitionOffset: TopicAndPartition => Option[Long],
-                                            consumerIdOpt: Option[String],
-                                            hostOpt: Option[String],
-                                            clientIdOpt: Option[String]): Array[PartitionAssignmentState] = {
-      if (topicPartitions.isEmpty)
-        Array[PartitionAssignmentState](
-          PartitionAssignmentState(group, coordinator, None, None, None, getLag(None, None), consumerIdOpt, hostOpt, clientIdOpt, None)
-        )
-      else {
-        var assignmentRows: Array[PartitionAssignmentState] = Array()
-        topicPartitions
-          .sortBy(_.partition)
-          .foreach { topicPartition =>
-            assignmentRows = assignmentRows :+ describePartition(group, coordinator, topicPartition.topic, topicPartition.partition, getPartitionOffset(topicPartition),
-              consumerIdOpt, hostOpt, clientIdOpt)
-          }
-        assignmentRows
-      }
+    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
 
-    protected def getLag(offset: Option[Long], logEndOffset: Option[Long]): Option[Long] =
-      offset.filter(_ != -1).flatMap(offset => logEndOffset.map(_ - offset))
+      val fields =
+        if (classOf[CsvRecordWithGroup] == clazz) CsvRecordWithGroup.fields
+        else if (classOf[CsvRecordNoGroup] == clazz) CsvRecordNoGroup.fields
+        else throw new IllegalStateException(s"Unhandled class $clazz")
 
-    private def describePartition(group: String,
-                                  coordinator: Option[Node],
-                                  topic: String,
-                                  partition: Int,
-                                  offsetOpt: Option[Long],
-                                  consumerIdOpt: Option[String],
-                                  hostOpt: Option[String],
-                                  clientIdOpt: Option[String]): PartitionAssignmentState = {
-      def getDescribePartitionResult(logEndOffsetOpt: Option[Long]): PartitionAssignmentState =
-        PartitionAssignmentState(group, coordinator, Option(topic), Option(partition), offsetOpt,
-                                 getLag(offsetOpt, logEndOffsetOpt), consumerIdOpt, hostOpt,
-                                 clientIdOpt, logEndOffsetOpt)
-
-      getLogEndOffset(new TopicPartition(topic, partition)) match {
-        case LogOffsetResult.LogOffset(logEndOffset) => getDescribePartitionResult(Some(logEndOffset))
-        case LogOffsetResult.Unknown => getDescribePartitionResult(None)
-        case LogOffsetResult.Ignore => null
-      }
+      val schema = mapper.schemaFor(clazz).sortedBy(fields: _*)
+      schema
     }
+  }
 
+  class ConsumerGroupService(val opts: ConsumerGroupCommandOptions,
+                             private[admin] val configOverrides: Map[String, String] = Map.empty) {
 
-    def resetOffsets(): Map[TopicPartition, OffsetAndMetadata] = throw new UnsupportedOperationException
+    private val adminClient = createAdminClient(configOverrides)
 
-    def exportOffsetsToReset(assignmentsToReset: Map[TopicPartition, OffsetAndMetadata]): String = throw new UnsupportedOperationException
-  }
+    // 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
+    }
 
-  @deprecated("This class has been deprecated and will be removed in a future release.", "0.11.0.0")
-  class ZkConsumerGroupService(val opts: ConsumerGroupCommandOptions) extends ConsumerGroupService {
+    def listGroups(): Unit = {
+      if (opts.options.has(opts.stateOpt)) {
+        val stateValue = opts.options.valueOf(opts.stateOpt)
+        val states = if (stateValue == null || stateValue.isEmpty)
+          Set[ConsumerGroupState]()
+        else
+          consumerGroupStatesFromString(stateValue)
+        val listings = listConsumerGroupsWithState(states)
+        printGroupStates(listings.map(e => (e.groupId, e.state.get.toString)))
+      } else
+        listConsumerGroups().foreach(println(_))
+    }
 
-    private val zkUtils = {
-      val zkUrl = opts.options.valueOf(opts.zkConnectOpt)
-      ZkUtils(zkUrl, 30000, 30000, JaasUtils.isZkSecurityEnabled)
+    def listConsumerGroups(): List[String] = {
+      val result = adminClient.listConsumerGroups(withTimeoutMs(new ListConsumerGroupsOptions))
+      val listings = result.all.get.asScala
+      listings.map(_.groupId).toList
     }
 
-    def close() {
-      zkUtils.close()
+    def listConsumerGroupsWithState(states: Set[ConsumerGroupState]): List[ConsumerGroupListing] = {
+      val listConsumerGroupsOptions = withTimeoutMs(new ListConsumerGroupsOptions())
+      listConsumerGroupsOptions.inStates(states.asJava)
+      val result = adminClient.listConsumerGroups(listConsumerGroupsOptions)
+      result.all.get.asScala.toList
     }
 
-    def listGroups(): List[String] = {
-      zkUtils.getConsumerGroups().toList
+    private def printGroupStates(groupsAndStates: List[(String, String)]): Unit = {
+      // find proper columns width
+      var maxGroupLen = 15
+      for ((groupId, state) <- groupsAndStates) {
+        maxGroupLen = Math.max(maxGroupLen, groupId.length)
+      }
+      println(s"%${-maxGroupLen}s %s".format("GROUP", "STATE"))
+      for ((groupId, state) <- groupsAndStates) {
+        println(s"%${-maxGroupLen}s %s".format(groupId, state))
+      }
     }
 
-    def deleteGroups() {
-      if (opts.options.has(opts.groupOpt) && opts.options.has(opts.topicOpt))
-        deleteForTopic()
-      else if (opts.options.has(opts.groupOpt))
-        deleteForGroup()
-      else if (opts.options.has(opts.topicOpt))
-        deleteAllForTopic()
+    private def shouldPrintMemberState(group: String, state: Option[String], numRows: Option[Int]): Boolean = {
+      // numRows contains the number of data rows, if any, compiled from the API call in the caller method.
+      // if it's undefined or 0, there is no relevant group information to display.
+      numRows match {
+        case None =>
+          printError(s"The consumer group '$group' does not exist.")
+          false
+        case Some(num) => state match {
+          case Some("Dead") =>
+            printError(s"Consumer group '$group' does not exist.")
+          case Some("Empty") =>
+            Console.err.println(s"\nConsumer group '$group' has no active members.")
+          case Some("PreparingRebalance") | Some("CompletingRebalance") =>
+            Console.err.println(s"\nWarning: Consumer group '$group' is rebalancing.")
+          case Some("Stable") =>
+          case other =>
+            // the control should never reach here
+            throw new KafkaException(s"Expected a valid consumer group state, but found '${other.getOrElse("NONE")}'.")
+        }
+        !state.contains("Dead") && num > 0
+      }
     }
 
-    protected def collectGroupAssignment(group: String): (Option[String], Option[Seq[PartitionAssignmentState]]) = {
-      val props = if (opts.options.has(opts.commandConfigOpt)) Utils.loadProps(opts.options.valueOf(opts.commandConfigOpt)) else new Properties()
-      val channelSocketTimeoutMs = props.getProperty("channelSocketTimeoutMs", "600").toInt
-      val channelRetryBackoffMs = props.getProperty("channelRetryBackoffMsOpt", "300").toInt
-      if (!zkUtils.getConsumerGroups().contains(group))
-        return (None, None)
-
-      val topics = zkUtils.getTopicsByConsumerGroup(group)
-      val topicPartitions = getAllTopicPartitions(topics)
-      var groupConsumerIds = zkUtils.getConsumersInGroup(group)
-
-      // mapping of topic partition -> consumer id
-      val consumerIdByTopicPartition = topicPartitions.map { topicPartition =>
-        val owner = zkUtils.readDataMaybeNull(new ZKGroupTopicDirs(group, topicPartition.topic).consumerOwnerDir + "/" + topicPartition.partition)._1
-        topicPartition -> owner.map(o => o.substring(0, o.lastIndexOf('-'))).getOrElse(MISSING_COLUMN_VALUE)
-      }.toMap
+    private def size(colOpt: Option[Seq[Object]]): Option[Int] = colOpt.map(_.size)
+
+    private def printOffsets(offsets: Map[String, (Option[String], Option[Seq[PartitionAssignmentState]])]): Unit = {
+      for ((groupId, (state, assignments)) <- offsets) {
+        if (shouldPrintMemberState(groupId, state, size(assignments))) {
+          // find proper columns width
+          var (maxGroupLen, maxTopicLen, maxConsumerIdLen, maxHostLen) = (15, 15, 15, 15)
+          assignments match {
+            case None => // do nothing
+            case Some(consumerAssignments) =>
+              consumerAssignments.foreach { consumerAssignment =>
+                maxGroupLen = Math.max(maxGroupLen, consumerAssignment.group.length)
+                maxTopicLen = Math.max(maxTopicLen, consumerAssignment.topic.getOrElse(MISSING_COLUMN_VALUE).length)
+                maxConsumerIdLen = Math.max(maxConsumerIdLen, consumerAssignment.consumerId.getOrElse(MISSING_COLUMN_VALUE).length)
+                maxHostLen = Math.max(maxHostLen, consumerAssignment.host.getOrElse(MISSING_COLUMN_VALUE).length)
+              }
+          }
 
-      // mapping of consumer id -> list of topic partitions
-      val consumerTopicPartitions = consumerIdByTopicPartition groupBy{_._2} map {
-        case (key, value) => (key, value.unzip._1.toArray) }
+          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))
+                )
+              }
+          }
+        }
+      }
+    }
 
-      // mapping of consumer id -> list of subscribed topics
-      val topicsByConsumerId = zkUtils.getTopicsPerMemberId(group)
+    private def printMembers(members: Map[String, (Option[String], Option[Seq[MemberAssignmentState]])], verbose: Boolean): Unit = {
+      for ((groupId, (state, assignments)) <- members) {
+        if (shouldPrintMemberState(groupId, state, size(assignments))) {
+          // find proper columns width
+          var (maxGroupLen, maxConsumerIdLen, maxGroupInstanceIdLen, maxHostLen, maxClientIdLen, includeGroupInstanceId) = (15, 15, 17, 15, 15, false)
+          assignments match {
+            case None => // do nothing
+            case Some(memberAssignments) =>
+              memberAssignments.foreach { memberAssignment =>
+                maxGroupLen = Math.max(maxGroupLen, memberAssignment.group.length)
+                maxConsumerIdLen = Math.max(maxConsumerIdLen, memberAssignment.consumerId.length)
+                maxGroupInstanceIdLen =  Math.max(maxGroupInstanceIdLen, memberAssignment.groupInstanceId.length)
+                maxHostLen = Math.max(maxHostLen, memberAssignment.host.length)
+                maxClientIdLen = Math.max(maxClientIdLen, memberAssignment.clientId.length)
+                includeGroupInstanceId = includeGroupInstanceId || memberAssignment.groupInstanceId.length > 0
+              }
+          }
 
-      var assignmentRows = topicPartitions.flatMap { topicPartition =>
-        val partitionOffsets = getPartitionOffsets(group, List(topicPartition), channelSocketTimeoutMs, channelRetryBackoffMs)
-        val consumerId = consumerIdByTopicPartition.get(topicPartition)
-        // since consumer id is repeated in client id, leave host and client id empty
-        consumerId.foreach(id => groupConsumerIds = groupConsumerIds.filterNot(_ == id))
-        collectConsumerAssignment(group, None, List(topicPartition), partitionOffsets.get, consumerId, None, None)
+          if (includeGroupInstanceId) {
+            print(s"\n%${-maxGroupLen}s %${-maxConsumerIdLen}s %${-maxGroupInstanceIdLen}s %${-maxHostLen}s %${-maxClientIdLen}s %-15s "
+                .format("GROUP", "CONSUMER-ID", "GROUP-INSTANCE-ID", "HOST", "CLIENT-ID", "#PARTITIONS"))
+          } else {
+            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 =>
+                if (includeGroupInstanceId) {
+                  print(s"%${-maxGroupLen}s %${-maxConsumerIdLen}s %${-maxGroupInstanceIdLen}s %${-maxHostLen}s %${-maxClientIdLen}s %-15s ".format(
+                    memberAssignment.group, memberAssignment.consumerId, memberAssignment.groupInstanceId, memberAssignment.host,
+                    memberAssignment.clientId, memberAssignment.numPartitions))
+                } else {
+                  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))
+                }
+                println()
+              }
+          }
+        }
       }
+    }
 
-      assignmentRows ++= groupConsumerIds.sortBy(- consumerTopicPartitions.get(_).size).flatMap { consumerId =>
-        topicsByConsumerId(consumerId).flatMap { _ =>
-          // since consumers with no topic partitions are processed here, we pass empty for topic partitions and offsets
-          // since consumer id is repeated in client id, leave host and client id empty
-          collectConsumerAssignment(group, None, Array[TopicAndPartition](), Map[TopicAndPartition, Option[Long]](), Some(consumerId), None, None)
+    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()
         }
       }
+    }
 
-      (None, Some(assignmentRows))
+    def describeGroups(): Unit = {
+      val groupIds =
+        if (opts.options.has(opts.allGroupsOpt)) listConsumerGroups()
+        else opts.options.valuesOf(opts.groupOpt).asScala
+      val membersOptPresent = opts.options.has(opts.membersOpt)
+      val stateOptPresent = opts.options.has(opts.stateOpt)
+      val offsetsOptPresent = opts.options.has(opts.offsetsOpt)
+      val subActions = Seq(membersOptPresent, offsetsOptPresent, stateOptPresent).count(_ == true)
+
+      if (subActions == 0 || offsetsOptPresent) {
+        val offsets = collectGroupsOffsets(groupIds)
+        printOffsets(offsets)
+      } else if (membersOptPresent) {
+        val members = collectGroupsMembers(groupIds, opts.options.has(opts.verboseOpt))
+        printMembers(members, opts.options.has(opts.verboseOpt))
+      } else {
+        val states = collectGroupsState(groupIds)
+        printStates(states)
+      }
     }
 
-    private def getAllTopicPartitions(topics: Seq[String]): Seq[TopicAndPartition] = {
-      val topicPartitionMap = zkUtils.getPartitionsForTopics(topics)
-      topics.flatMap { topic =>
-        val partitions = topicPartitionMap.getOrElse(topic, Seq.empty)
-        partitions.map(TopicAndPartition(topic, _))
+    private def collectConsumerAssignment(group: String,
+                                          coordinator: Option[Node],
+                                          topicPartitions: Seq[TopicPartition],
+                                          getPartitionOffset: TopicPartition => Option[Long],
+                                          consumerIdOpt: Option[String],
+                                          hostOpt: Option[String],
+                                          clientIdOpt: Option[String]): Array[PartitionAssignmentState] = {
+      if (topicPartitions.isEmpty) {
+        Array[PartitionAssignmentState](
+          PartitionAssignmentState(group, coordinator, None, None, None, getLag(None, None), consumerIdOpt, hostOpt, clientIdOpt, None)
+        )
       }
+      else
+        describePartitions(group, coordinator, topicPartitions.sortBy(_.partition), getPartitionOffset, consumerIdOpt, hostOpt, clientIdOpt)
     }
 
-    protected def getLogEndOffset(topicPartition: TopicPartition): LogOffsetResult = {
-      zkUtils.getLeaderForPartition(topicPartition.topic, topicPartition.partition) match {
-        case Some(-1) => LogOffsetResult.Unknown
-        case Some(brokerId) =>
-          getZkConsumer(brokerId).map { consumer =>
-            val topicAndPartition = TopicAndPartition(topicPartition.topic, topicPartition.partition)
-            val request = OffsetRequest(Map(topicAndPartition -> PartitionOffsetRequestInfo(OffsetRequest.LatestTime, 1)))
-            val logEndOffset = consumer.getOffsetsBefore(request).partitionErrorAndOffsets(topicAndPartition).offsets.head
-            consumer.close()
-            LogOffsetResult.LogOffset(logEndOffset)
-          }.getOrElse(LogOffsetResult.Ignore)
-        case None =>
-          printError(s"No broker for partition '$topicPartition'")
-          LogOffsetResult.Ignore
-      }
-    }
-
-    private def getPartitionOffsets(group: String,
-                                    topicPartitions: Seq[TopicAndPartition],
-                                    channelSocketTimeoutMs: Int,
-                                    channelRetryBackoffMs: Int): Map[TopicAndPartition, Long] = {
-      val offsetMap = mutable.Map[TopicAndPartition, Long]()
-      val channel = ClientUtils.channelToOffsetManager(group, zkUtils, channelSocketTimeoutMs, channelRetryBackoffMs)
-      channel.send(OffsetFetchRequest(group, topicPartitions))
-      val offsetFetchResponse = OffsetFetchResponse.readFrom(channel.receive().payload())
-
-      offsetFetchResponse.requestInfo.foreach { case (topicAndPartition, offsetAndMetadata) =>
-        offsetAndMetadata match {
-          case OffsetMetadataAndError.NoOffset =>
-            val topicDirs = new ZKGroupTopicDirs(group, topicAndPartition.topic)
-            // this group may not have migrated off zookeeper for offsets storage (we don't expose the dual-commit option in this tool
-            // (meaning the lag may be off until all the consumers in the group have the same setting for offsets storage)
-            try {
-              val offset = zkUtils.readData(topicDirs.consumerOffsetDir + "/" + topicAndPartition.partition)._1.toLong
-              offsetMap.put(topicAndPartition, offset)
-            } catch {
-              case z: ZkNoNodeException =>
-                printError(s"Could not fetch offset from zookeeper for group '$group' partition '$topicAndPartition' due to missing offset data in zookeeper.", Some(z))
-            }
-          case offsetAndMetaData if offsetAndMetaData.error == Errors.NONE =>
-            offsetMap.put(topicAndPartition, offsetAndMetadata.offset)
-          case _ =>
-            printError(s"Could not fetch offset from kafka for group '$group' partition '$topicAndPartition' due to ${offsetAndMetadata.error.message}.")
-        }
+    private def getLag(offset: Option[Long], logEndOffset: Option[Long]): Option[Long] =
+      offset.filter(_ != -1).flatMap(offset => logEndOffset.map(_ - offset))
+
+    private def describePartitions(group: String,
+                                   coordinator: Option[Node],
+                                   topicPartitions: Seq[TopicPartition],
+                                   getPartitionOffset: TopicPartition => Option[Long],
+                                   consumerIdOpt: Option[String],
+                                   hostOpt: Option[String],
+                                   clientIdOpt: Option[String]): Array[PartitionAssignmentState] = {
+
+      def getDescribePartitionResult(topicPartition: TopicPartition, logEndOffsetOpt: Option[Long]): PartitionAssignmentState = {
+        val offset = getPartitionOffset(topicPartition)
+        PartitionAssignmentState(group, coordinator, Option(topicPartition.topic), Option(topicPartition.partition), offset,
+          getLag(offset, logEndOffsetOpt), consumerIdOpt, hostOpt, clientIdOpt, logEndOffsetOpt)
       }
-      channel.disconnect()
-      offsetMap.toMap
+
+      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
     }
 
-    private def deleteForGroup() {
-      val groups = opts.options.valuesOf(opts.groupOpt)
-      groups.asScala.foreach { group =>
-        try {
-          if (AdminUtils.deleteConsumerGroupInZK(zkUtils, group))
-            println(s"Deleted all consumer group information for group '$group' in zookeeper.")
-          else
-            printError(s"Delete for group '$group' failed because its consumers are still active.")
+    def resetOffsets(): Map[String, Map[TopicPartition, OffsetAndMetadata]] = {
+      val groupIds =
+        if (opts.options.has(opts.allGroupsOpt)) listConsumerGroups()
+        else opts.options.valuesOf(opts.groupOpt).asScala
+
+      val consumerGroups = adminClient.describeConsumerGroups(
+        groupIds.asJava,
+        withTimeoutMs(new DescribeConsumerGroupsOptions)
+      ).describedGroups()
+
+      val result =
+        consumerGroups.asScala.foldLeft(immutable.Map[String, Map[TopicPartition, OffsetAndMetadata]]()) {
+          case (acc, (groupId, groupDescription)) =>
+            groupDescription.get.state().toString match {
+              case "Empty" | "Dead" =>
+                val partitionsToReset = getPartitionsToReset(groupId)
+                val preparedOffsets = prepareOffsetsToReset(groupId, partitionsToReset)
+
+                // Dry-run is the default behavior if --execute is not specified
+                val dryRun = opts.options.has(opts.dryRunOpt) || !opts.options.has(opts.executeOpt)
+                if (!dryRun) {
+                  adminClient.alterConsumerGroupOffsets(
+                    groupId,
+                    preparedOffsets.asJava,
+                    withTimeoutMs(new AlterConsumerGroupOffsetsOptions)
+                  ).all.get
+                }
+                acc.updated(groupId, preparedOffsets)
+              case currentState =>
+                printError(s"Assignments can only be reset if the group '$groupId' is inactive, but the current state is $currentState.")
+                acc.updated(groupId, Map.empty)
+            }
         }
-        catch {
-          case e: ZkNoNodeException =>
-            printError(s"Delete for group '$group' failed because group does not exist.", Some(e))
+      result
+    }
+
+    def deleteOffsets(groupId: String, topics: List[String]): (Errors, Map[TopicPartition, Throwable]) = {
+      val partitionLevelResult = mutable.Map[TopicPartition, Throwable]()
+
+      val (topicWithPartitions, topicWithoutPartitions) = topics.partition(_.contains(":"))
+
+      val knownPartitions = topicWithPartitions.flatMap { topicArg =>
+        val split = topicArg.split(":")
+        split(1).split(",").map { partition =>
+          new TopicPartition(split(0), partition.toInt)
         }
       }
-    }
 
-    private def deleteForTopic() {
-      val groups = opts.options.valuesOf(opts.groupOpt)
-      val topic = opts.options.valueOf(opts.topicOpt)
-      Topic.validate(topic)
-      groups.asScala.foreach { group =>
-        try {
-          if (AdminUtils.deleteConsumerGroupInfoForTopicInZK(zkUtils, group, topic))
-            println(s"Deleted consumer group information for group '$group' topic '$topic' in zookeeper.")
-          else
-            printError(s"Delete for group '$group' topic '$topic' failed because its consumers are still active.")
+      // Get the partitions of topics that the user did not explicitly specify the partitions
+      val describeTopicsResult = adminClient.describeTopics(
+        topicWithoutPartitions.asJava,
+        withTimeoutMs(new DescribeTopicsOptions))
+
+      val unknownPartitions = describeTopicsResult.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
         }
-        catch {
-          case e: ZkNoNodeException =>
-            printError(s"Delete for group '$group' topic '$topic' failed because group does not exist.", Some(e))
+      }
+
+      val partitions = knownPartitions ++ unknownPartitions
+
+      val deleteResult = adminClient.deleteConsumerGroupOffsets(
+        groupId,
+        partitions.toSet.asJava,
+        withTimeoutMs(new DeleteConsumerGroupOffsetsOptions)
+      )
+
+      var topLevelException = Errors.NONE
+      Try(deleteResult.all.get) match {
+        case Success(_) =>
+        case Failure(e) => topLevelException = Errors.forException(e.getCause)
+      }
+
+      partitions.foreach { partition =>
+        Try(deleteResult.partitionResult(partition).get()) match {
+          case Success(_) => partitionLevelResult += partition -> null
+          case Failure(e) => partitionLevelResult += partition -> e
         }
       }
+
+      (topLevelException, partitionLevelResult)
     }
 
-    private def deleteAllForTopic() {
-      val topic = opts.options.valueOf(opts.topicOpt)
-      Topic.validate(topic)
-      AdminUtils.deleteAllConsumerGroupInfoForTopicInZK(zkUtils, topic)
-      println(s"Deleted consumer group information for all inactive consumer groups for topic '$topic' in zookeeper.")
+    def deleteOffsets(): Unit = {
+      val groupId = opts.options.valueOf(opts.groupOpt)
+      val topics = opts.options.valuesOf(opts.topicOpt).asScala.toList
+
+      val (topLevelResult, partitionLevelResult) = deleteOffsets(groupId, topics)
+
+      topLevelResult match {
+        case Errors.NONE =>
+          println(s"Request succeed for deleting offsets with topic ${topics.mkString(", ")} group $groupId")
+        case Errors.INVALID_GROUP_ID =>
+          printError(s"'$groupId' is not valid.")
+        case Errors.GROUP_ID_NOT_FOUND =>
+          printError(s"'$groupId' does not exist.")
+        case Errors.GROUP_AUTHORIZATION_FAILED =>
+          printError(s"Access to '$groupId' is not authorized.")
+        case Errors.NON_EMPTY_GROUP =>
+          printError(s"Deleting offsets of a consumer group '$groupId' is forbidden if the group is not empty.")
+        case Errors.GROUP_SUBSCRIBED_TO_TOPIC |
+             Errors.TOPIC_AUTHORIZATION_FAILED |
+             Errors.UNKNOWN_TOPIC_OR_PARTITION =>
+          printError(s"Encounter some partition level error, see the follow-up details:")
+        case _ =>
+          printError(s"Encounter some unknown error: $topLevelResult")
+      }
+
+      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 def getZkConsumer(brokerId: Int): Option[SimpleConsumer] = {
-      try {
-        zkUtils.getBrokerInfo(brokerId)
-          .map(_.getBrokerEndPoint(ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT)))
-          .map(endPoint => new SimpleConsumer(endPoint.host, endPoint.port, 10000, 100000, "ConsumerGroupCommand"))
-          .orElse(throw new BrokerNotAvailableException("Broker id %d does not exist".format(brokerId)))
-      } catch {
-        case t: Throwable =>
-          printError(s"Could not parse broker info due to ${t.getMessage}", Some(t))
-          None
+    private[admin] def describeConsumerGroups(groupIds: Seq[String]): mutable.Map[String, ConsumerGroupDescription] = {
+      adminClient.describeConsumerGroups(
+        groupIds.asJava,
+        withTimeoutMs(new DescribeConsumerGroupsOptions)
+      ).describedGroups().asScala.map {
+        case (groupId, groupDescriptionFuture) => (groupId, groupDescriptionFuture.get())
       }
     }
-  }
 
-  class KafkaConsumerGroupService(val opts: ConsumerGroupCommandOptions) extends ConsumerGroupService {
+    /**
+      * 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))
+    }
 
-    private val adminClient = createAdminClient()
+    /**
+      * 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.filter { case (tp, _) =>
+          !assignedTopicPartitions.contains(tp)
+        }.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
 
-    // `consumer` is only needed for `describe`, so we instantiate it lazily
-    private var consumer: KafkaConsumer[String, String] = null
+      groupOffsets
+    }
 
-    def listGroups(): List[String] = {
-      adminClient.listAllConsumerGroupsFlattened().map(_.groupId)
+    private[admin] def collectGroupMembers(groupId: String, verbose: Boolean): (Option[String], Option[Seq[MemberAssignmentState]]) = {
+      collectGroupsMembers(Seq(groupId), verbose)(groupId)
     }
 
-    protected def collectGroupAssignment(group: String): (Option[String], Option[Seq[PartitionAssignmentState]]) = {
-      val consumerGroupSummary = adminClient.describeConsumerGroup(group, opts.options.valueOf(opts.timeoutMsOpt))
-      (Some(consumerGroupSummary.state),
-        consumerGroupSummary.consumers match {
-          case None =>
-            None
-          case Some(consumers) =>
-            var assignedTopicPartitions = Array[TopicPartition]()
-            val offsets = adminClient.listGroupOffsets(group)
-            val rowsWithConsumer =
-              if (offsets.isEmpty)
-                List[PartitionAssignmentState]()
-              else {
-                consumers.sortWith(_.assignment.size > _.assignment.size).flatMap { consumerSummary =>
-                  val topicPartitions = consumerSummary.assignment.map(tp => TopicAndPartition(tp.topic, tp.partition))
-                  assignedTopicPartitions = assignedTopicPartitions ++ consumerSummary.assignment
-                  val partitionOffsets: Map[TopicAndPartition, Option[Long]] = consumerSummary.assignment.map { topicPartition =>
-                    new TopicAndPartition(topicPartition) -> offsets.get(topicPartition)
-                  }.toMap
-                  collectConsumerAssignment(group, Some(consumerGroupSummary.coordinator), topicPartitions,
-                    partitionOffsets, Some(s"${consumerSummary.consumerId}"), Some(s"${consumerSummary.host}"),
-                    Some(s"${consumerSummary.clientId}"))
-                }
-              }
+    private[admin] def collectGroupsMembers(groupIds: Seq[String], verbose: Boolean): TreeMap[String, (Option[String], Option[Seq[MemberAssignmentState]])] = {
+      val consumerGroups = describeConsumerGroups(groupIds)
+      TreeMap[String, (Option[String], Option[Seq[MemberAssignmentState]])]() ++ (for ((groupId, consumerGroup) <- consumerGroups) yield {
+        val state = consumerGroup.state.toString
+        val memberAssignmentStates = consumerGroup.members().asScala.map(consumer =>
+          MemberAssignmentState(
+            groupId,
+            consumer.consumerId,
+            consumer.host,
+            consumer.clientId,
+            consumer.groupInstanceId.orElse(""),
+            consumer.assignment.topicPartitions.size(),
+            if (verbose) consumer.assignment.topicPartitions.asScala.toList else List()
+          )).toList
+        groupId -> (Some(state), Option(memberAssignmentStates))
+      }).toMap
+    }
 
-            val rowsWithoutConsumer = offsets.filterNot {
-              case (topicPartition, offset) => assignedTopicPartitions.contains(topicPartition)
-              }.flatMap {
-                case (topicPartition, offset) =>
-                  val topicAndPartition = new TopicAndPartition(topicPartition)
-                  collectConsumerAssignment(group, Some(consumerGroupSummary.coordinator), Seq(topicAndPartition),
-                      Map(topicAndPartition -> Some(offset)), Some(MISSING_COLUMN_VALUE),
-                      Some(MISSING_COLUMN_VALUE), Some(MISSING_COLUMN_VALUE))
-                }
+    private[admin] def collectGroupState(groupId: String): GroupState = {
+      collectGroupsState(Seq(groupId))(groupId)
+    }
 
-            Some(rowsWithConsumer ++ rowsWithoutConsumer)
-      }
-      )
+    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
     }
 
-    protected def getLogEndOffset(topicPartition: TopicPartition): LogOffsetResult = {
-      val consumer = getConsumer()
-      val offsets = consumer.endOffsets(List(topicPartition).asJava)
-      val logStartOffset = offsets.get(topicPartition)
-      LogOffsetResult.LogOffset(logStartOffset)
+    private def getLogEndOffsets(groupId: String, topicPartitions: Seq[TopicPartition]): Map[TopicPartition, LogOffsetResult] = {
+      val endOffsets = topicPartitions.map { topicPartition =>
+        topicPartition -> OffsetSpec.latest
+      }.toMap
+      val offsets = adminClient.listOffsets(
+        endOffsets.asJava,
+        withTimeoutMs(new ListOffsetsOptions)
+      ).all.get
+      topicPartitions.map { topicPartition =>
+        Option(offsets.get(topicPartition)) match {
+          case Some(listOffsetsResultInfo) => topicPartition -> LogOffsetResult.LogOffset(listOffsetsResultInfo.offset)
+          case _ => topicPartition -> LogOffsetResult.Unknown
+        }
+      }.toMap
     }
 
-    protected def getLogStartOffset(topicPartition: TopicPartition): LogOffsetResult = {
-      val consumer = getConsumer()
-      val offsets = consumer.beginningOffsets(List(topicPartition).asJava)
-      val logStartOffset = offsets.get(topicPartition)
-      LogOffsetResult.LogOffset(logStartOffset)
+    private def getLogStartOffsets(groupId: String, topicPartitions: Seq[TopicPartition]): Map[TopicPartition, LogOffsetResult] = {
+      val startOffsets = topicPartitions.map { topicPartition =>
+        topicPartition -> OffsetSpec.earliest
+      }.toMap
+      val offsets = adminClient.listOffsets(
+        startOffsets.asJava,
+        withTimeoutMs(new ListOffsetsOptions)
+      ).all.get
+      topicPartitions.map { topicPartition =>
+        Option(offsets.get(topicPartition)) match {
+          case Some(listOffsetsResultInfo) => topicPartition -> LogOffsetResult.LogOffset(listOffsetsResultInfo.offset)
+          case _ => topicPartition -> LogOffsetResult.Unknown
+        }
+      }.toMap
     }
 
-    protected def getLogTimestampOffset(topicPartition: TopicPartition, timestamp: java.lang.Long): LogOffsetResult = {
-      val consumer = getConsumer()
-      consumer.assign(List(topicPartition).asJava)
-      val offsetsForTimes = consumer.offsetsForTimes(Map(topicPartition -> timestamp).asJava)
-      if (offsetsForTimes != null && !offsetsForTimes.isEmpty && offsetsForTimes.get(topicPartition) != null)
-        LogOffsetResult.LogOffset(offsetsForTimes.get(topicPartition).offset)
-      else {
-        getLogEndOffset(topicPartition)
-      }
+    private def getLogTimestampOffsets(groupId: String, topicPartitions: Seq[TopicPartition], timestamp: java.lang.Long): Map[TopicPartition, LogOffsetResult] = {
+      val timestampOffsets = topicPartitions.map { topicPartition =>
+        topicPartition -> OffsetSpec.forTimestamp(timestamp)
+      }.toMap
+      val offsets = adminClient.listOffsets(
+        timestampOffsets.asJava,
+        withTimeoutMs(new ListOffsetsOptions)
+      ).all.get
+      val (successfulOffsetsForTimes, unsuccessfulOffsetsForTimes) =
+        offsets.asScala.partition(_._2.offset != ListOffsetResponse.UNKNOWN_OFFSET)
+
+      val successfulLogTimestampOffsets = successfulOffsetsForTimes.map {
+        case (topicPartition, listOffsetsResultInfo) => topicPartition -> LogOffsetResult.LogOffset(listOffsetsResultInfo.offset)
+      }.toMap
+
+      successfulLogTimestampOffsets ++ getLogEndOffsets(groupId, unsuccessfulOffsetsForTimes.keySet.toSeq)
     }
 
-    def close() {
+    def close(): Unit = {
       adminClient.close()
-      if (consumer != null) consumer.close()
     }
 
-    private def createAdminClient(): 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))
-      AdminClient.create(props)
+      configOverrides.forKeyValue { (k, v) => props.put(k, v)}
+      Admin.create(props)
     }
 
-    private def getConsumer() = {
-      if (consumer == null)
-        consumer = createNewConsumer()
-      consumer
+    private def withTimeoutMs [T <: AbstractOptions[T]] (options : T) =  {
+      val t = opts.options.valueOf(opts.timeoutMsOpt).intValue()
+      options.timeoutMs(t)
     }
 
-    private def createNewConsumer(): 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.ENABLE_AUTO_COMMIT_CONFIG, "false")
-      properties.put(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG, "30000")
-      properties.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, deserializer)
-      properties.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, deserializer)
-      if (opts.options.has(opts.commandConfigOpt))
-        properties ++= Utils.loadProps(opts.options.valueOf(opts.commandConfigOpt))
-
-      new KafkaConsumer(properties)
-    }
-
-    override def resetOffsets(): Map[TopicPartition, OffsetAndMetadata] = {
-      val groupId = opts.options.valueOf(opts.groupOpt)
-      val consumerGroupSummary = adminClient.describeConsumerGroup(groupId, opts.options.valueOf(opts.timeoutMsOpt))
-      consumerGroupSummary.state match {
-        case "Empty" | "Dead" =>
-          val partitionsToReset = getPartitionsToReset(groupId)
-          val preparedOffsets = prepareOffsetsToReset(groupId, partitionsToReset)
-          val execute = opts.options.has(opts.executeOpt)
-          if (execute)
-            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
-      }
-    }
-
-    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 topicAndPartitions = topicArg.split(":")
-        val topic = topicAndPartitions(0)
-        topicAndPartitions(1).split(",").map(partition => new TopicPartition(topic, partition.toInt))
-      case topic => getConsumer().partitionsFor(topic).asScala
-        .map(partitionInfo => new TopicPartition(topic, partitionInfo.partition))
+        val topicPartitions = topicArg.split(":")
+        val topic = topicPartitions(0)
+        topicPartitions(1).split(",").map(partition => new TopicPartition(topic, partition.toInt))
+      case topic =>
+        val descriptionMap = adminClient.describeTopics(
+          Seq(topic).asJava,
+          withTimeoutMs(new DescribeTopicsOptions)
+        ).all().get.asScala
+        val r = descriptionMap.flatMap{ case(topic, description) =>
+          description.partitions().asScala.map{ tpInfo =>
+            new TopicPartition(topic, tpInfo.partition)
+          }
+        }
+        r
     }
 
     private def getPartitionsToReset(groupId: String): Seq[TopicPartition] = {
       if (opts.options.has(opts.allTopicsOpt)) {
-        val allTopicPartitions = adminClient.listGroupOffsets(groupId).keys.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
@@ -563,131 +736,213 @@ object ConsumerGroupCommand extends Logging {
       }
     }
 
-    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
+    private def getCommittedOffsets(groupId: String): Map[TopicPartition, OffsetAndMetadata] = {
+      adminClient.listConsumerGroupOffsets(
+        groupId,
+        withTimeoutMs(new ListConsumerGroupOffsetsOptions)
+      ).partitionsToOffsetAndMetadata.get.asScala
     }
 
-    private def prepareOffsetsToReset(groupId: String, partitionsToReset: Iterable[TopicPartition]): Map[TopicPartition, OffsetAndMetadata] = {
+    type GroupMetadata = immutable.Map[String, immutable.Map[TopicPartition, OffsetAndMetadata]]
+    private def parseResetPlan(resetPlanCsv: String): GroupMetadata = {
+      def updateGroupMetadata(group: String, topic: String, partition: Int, offset: Long, acc: GroupMetadata) = {
+        val topicPartition = new TopicPartition(topic, partition)
+        val offsetAndMetadata = new OffsetAndMetadata(offset)
+        val dataMap = acc.getOrElse(group, immutable.Map()).updated(topicPartition, offsetAndMetadata)
+        acc.updated(group, dataMap)
+      }
+      val csvReader = CsvUtils().readerFor[CsvRecordNoGroup]
+      val lines = resetPlanCsv.split("\n")
+      val isSingleGroupQuery = opts.options.valuesOf(opts.groupOpt).size() == 1
+      val isOldCsvFormat = lines.headOption.flatMap(line =>
+        Try(csvReader.readValue[CsvRecordNoGroup](line)).toOption).nonEmpty
+      // Single group CSV format: "topic,partition,offset"
+      val dataMap = if (isSingleGroupQuery && isOldCsvFormat) {
+        val group = opts.options.valueOf(opts.groupOpt)
+        lines.foldLeft(immutable.Map[String, immutable.Map[TopicPartition, OffsetAndMetadata]]()) { (acc, line) =>
+          val CsvRecordNoGroup(topic, partition, offset) = csvReader.readValue[CsvRecordNoGroup](line)
+          updateGroupMetadata(group, topic, partition, offset, acc)
+        }
+        // Multiple group CSV format: "group,topic,partition,offset"
+      } else {
+        val csvReader = CsvUtils().readerFor[CsvRecordWithGroup]
+        lines.foldLeft(immutable.Map[String, immutable.Map[TopicPartition, OffsetAndMetadata]]()) { (acc, line) =>
+          val CsvRecordWithGroup(group, topic, partition, offset) = csvReader.readValue[CsvRecordWithGroup](line)
+          updateGroupMetadata(group, topic, partition, offset, acc)
+        }
+      }
+      dataMap
+    }
+
+    private def prepareOffsetsToReset(groupId: String,
+                                      partitionsToReset: Seq[TopicPartition]): Map[TopicPartition, OffsetAndMetadata] = {
       if (opts.options.has(opts.resetToOffsetOpt)) {
         val offset = opts.options.valueOf(opts.resetToOffsetOpt)
-        partitionsToReset.map {
-          topicPartition =>
-            val newOffset: Long = checkOffsetRange(topicPartition, offset)
-            (topicPartition, new OffsetAndMetadata(newOffset))
-        }.toMap
+        checkOffsetsRange(groupId, partitionsToReset.map((_, offset)).toMap).map {
+          case (topicPartition, newOffset) => (topicPartition, new OffsetAndMetadata(newOffset))
+        }
       } else if (opts.options.has(opts.resetToEarliestOpt)) {
+        val logStartOffsets = getLogStartOffsets(groupId, partitionsToReset)
         partitionsToReset.map { topicPartition =>
-          getLogStartOffset(topicPartition) match {
-            case LogOffsetResult.LogOffset(offset) => (topicPartition, new OffsetAndMetadata(offset))
+          logStartOffsets.get(topicPartition) match {
+            case Some(LogOffsetResult.LogOffset(offset)) => (topicPartition, new OffsetAndMetadata(offset))
             case _ => CommandLineUtils.printUsageAndDie(opts.parser, s"Error getting starting offset of topic partition: $topicPartition")
           }
         }.toMap
       } else if (opts.options.has(opts.resetToLatestOpt)) {
+        val logEndOffsets = getLogEndOffsets(groupId, partitionsToReset)
         partitionsToReset.map { topicPartition =>
-          getLogEndOffset(topicPartition) match {
-            case LogOffsetResult.LogOffset(offset) => (topicPartition, new OffsetAndMetadata(offset))
+          logEndOffsets.get(topicPartition) match {
+            case Some(LogOffsetResult.LogOffset(offset)) => (topicPartition, new OffsetAndMetadata(offset))
             case _ => CommandLineUtils.printUsageAndDie(opts.parser, s"Error getting ending offset of topic partition: $topicPartition")
           }
         }.toMap
       } else if (opts.options.has(opts.resetShiftByOpt)) {
-        val currentCommittedOffsets = adminClient.listGroupOffsets(groupId)
-        partitionsToReset.map { topicPartition =>
+        val currentCommittedOffsets = getCommittedOffsets(groupId)
+        val requestedOffsets = partitionsToReset.map { topicPartition =>
           val shiftBy = opts.options.valueOf(opts.resetShiftByOpt)
           val currentOffset = currentCommittedOffsets.getOrElse(topicPartition,
-            throw new IllegalArgumentException(s"Cannot shift offset for partition $topicPartition since there is no current committed offset"))
-          val shiftedOffset = currentOffset + shiftBy
-          val newOffset: Long = checkOffsetRange(topicPartition, shiftedOffset)
-          (topicPartition, new OffsetAndMetadata(newOffset))
+            throw new IllegalArgumentException(s"Cannot shift offset for partition $topicPartition since there is no current committed offset")).offset
+          (topicPartition, currentOffset + shiftBy)
         }.toMap
+        checkOffsetsRange(groupId, requestedOffsets).map {
+          case (topicPartition, newOffset) => (topicPartition, new OffsetAndMetadata(newOffset))
+        }
       } else if (opts.options.has(opts.resetToDatetimeOpt)) {
+        val timestamp = Utils.getDateTime(opts.options.valueOf(opts.resetToDatetimeOpt))
+        val logTimestampOffsets = getLogTimestampOffsets(groupId, partitionsToReset, timestamp)
         partitionsToReset.map { topicPartition =>
-          val timestamp = getDateTime
-          val logTimestampOffset = getLogTimestampOffset(topicPartition, timestamp)
+          val logTimestampOffset = logTimestampOffsets.get(topicPartition)
           logTimestampOffset match {
-            case LogOffsetResult.LogOffset(offset) => (topicPartition, new OffsetAndMetadata(offset))
+            case Some(LogOffsetResult.LogOffset(offset)) => (topicPartition, new OffsetAndMetadata(offset))
             case _ => CommandLineUtils.printUsageAndDie(opts.parser, s"Error getting offset by timestamp of topic partition: $topicPartition")
           }
         }.toMap
       } else if (opts.options.has(opts.resetByDurationOpt)) {
+        val duration = opts.options.valueOf(opts.resetByDurationOpt)
+        val durationParsed = Duration.parse(duration)
+        val now = Instant.now()
+        durationParsed.negated().addTo(now)
+        val timestamp = now.minus(durationParsed).toEpochMilli
+        val logTimestampOffsets = getLogTimestampOffsets(groupId, partitionsToReset, timestamp)
         partitionsToReset.map { topicPartition =>
-          val duration = opts.options.valueOf(opts.resetByDurationOpt)
-          val now = new Date()
-          val durationParsed = DatatypeFactory.newInstance().newDuration(duration)
-          durationParsed.negate().addTo(now)
-          val timestamp = now.getTime
-          val logTimestampOffset = getLogTimestampOffset(topicPartition, timestamp)
+          val logTimestampOffset = logTimestampOffsets.get(topicPartition)
           logTimestampOffset match {
-            case LogOffsetResult.LogOffset(offset) => (topicPartition, new OffsetAndMetadata(offset))
+            case Some(LogOffsetResult.LogOffset(offset)) => (topicPartition, new OffsetAndMetadata(offset))
             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)
-        resetPlan.keySet.map { topicPartition =>
-          val newOffset: Long = checkOffsetRange(topicPartition, resetPlan(topicPartition).offset())
-          (topicPartition, new OffsetAndMetadata(newOffset))
-        }.toMap
+      } else if (resetPlanFromFile.isDefined) {
+        resetPlanFromFile.map(resetPlan => resetPlan.get(groupId).map { resetPlanForGroup =>
+          val requestedOffsets = resetPlanForGroup.keySet.map { topicPartition =>
+            topicPartition -> resetPlanForGroup(topicPartition).offset
+          }.toMap
+          checkOffsetsRange(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 = adminClient.listGroupOffsets(groupId)
-        partitionsToReset.map { topicPartition =>
-          currentCommittedOffsets.get(topicPartition).map { offset =>
-            (topicPartition, new OffsetAndMetadata(offset))
-          }.getOrElse(
-            getLogEndOffset(topicPartition) match {
-              case LogOffsetResult.LogOffset(offset) => (topicPartition, new OffsetAndMetadata(offset))
-              case _ => CommandLineUtils.printUsageAndDie(opts.parser, s"Error getting ending offset of topic partition: $topicPartition")
-            }
-          )
+        val currentCommittedOffsets = getCommittedOffsets(groupId)
+        val (partitionsToResetWithCommittedOffset, partitionsToResetWithoutCommittedOffset) =
+          partitionsToReset.partition(currentCommittedOffsets.keySet.contains(_))
+
+        val preparedOffsetsForPartitionsWithCommittedOffset = partitionsToResetWithCommittedOffset.map { topicPartition =>
+          (topicPartition, new OffsetAndMetadata(currentCommittedOffsets.get(topicPartition) match {
+            case Some(offset) => offset.offset
+            case None => throw new IllegalStateException(s"Expected a valid current offset for topic partition: $topicPartition")
+          }))
         }.toMap
+
+        val preparedOffsetsForPartitionsWithoutCommittedOffset = getLogEndOffsets(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")
+        }
+
+        preparedOffsetsForPartitionsWithCommittedOffset ++ preparedOffsetsForPartitionsWithoutCommittedOffset
       } else {
         CommandLineUtils.printUsageAndDie(opts.parser, "Option '%s' requires one of the following scenarios: %s".format(opts.resetOffsetsOpt, opts.allResetOffsetScenarioOpts) )
       }
     }
 
-    private def checkOffsetRange(topicPartition: TopicPartition, offset: Long) = {
-      getLogEndOffset(topicPartition) match {
-        case LogOffsetResult.LogOffset(endOffset) if offset > endOffset =>
-          warn(s"New offset ($offset) is higher than latest offset. Value will be set to $endOffset")
-          endOffset
+    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 =>
+            warn(s"New offset ($offset) is higher than latest offset for topic partition $topicPartition. Value will be set to $endOffset")
+            endOffset
+
+          case Some(_) => logStartOffsets.get(topicPartition) match {
+            case Some(LogOffsetResult.LogOffset(startOffset)) if offset < startOffset =>
+              warn(s"New offset ($offset) is lower than earliest offset for topic partition $topicPartition. Value will be set to $startOffset")
+              startOffset
+
+            case _ => offset
+          }
 
-        case _ => getLogStartOffset(topicPartition) match {
-          case LogOffsetResult.LogOffset(startOffset) if offset < startOffset =>
-            warn(s"New offset ($offset) is lower than earliest offset. Value will be set to $startOffset")
-            startOffset
+          case None => // the control should not reach here
+            throw new IllegalStateException(s"Unexpected non-existing offset value for topic partition $topicPartition")
+        })
+      }
+    }
 
-          case _ => offset
+    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("")
     }
 
-    private[admin] def getDateTime: java.lang.Long = {
-      val datetime: String = opts.options.valueOf(opts.resetToDatetimeOpt) match {
-        case ts if ts.split("T")(1).contains("+") || ts.split("T")(1).contains("-") || ts.split("T")(1).contains("Z") => ts.toString
-        case ts => s"${ts}Z"
+    def deleteGroups(): Map[String, Throwable] = {
+      val groupIds =
+        if (opts.options.has(opts.allGroupsOpt)) listConsumerGroups()
+        else opts.options.valuesOf(opts.groupOpt).asScala
+
+      val groupsToDelete = adminClient.deleteConsumerGroups(
+        groupIds.asJava,
+        withTimeoutMs(new DeleteConsumerGroupsOptions)
+      ).deletedGroups().asScala
+
+      val result = groupsToDelete.map { case (g, f) =>
+        Try(f.get) match {
+          case Success(_) => g -> null
+          case Failure(e) => g -> e
+        }
+      }
+
+      val (success, failed) = result.partition {
+        case (_, error) => error == null
       }
-      val date = {
-        try {
-          new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss.SSSXXX").parse(datetime)
-        } catch {
-          case e: ParseException => new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss.SSSX").parse(datetime)
+
+      if (failed.isEmpty) {
+        println(s"Deletion of requested consumer groups (${success.keySet.mkString("'", "', '", "'")}) was successful.")
+      }
+      else {
+        printError("Deletion of some consumer groups failed:")
+        failed.foreach {
+          case (group, error) => println(s"* Group '$group' could not be deleted due to: ${error.toString}")
         }
+        if (success.nonEmpty)
+          println(s"\nThese consumer groups were deleted successfully: ${success.keySet.mkString("'", "', '", "'")}")
       }
-      date.getTime
-    }
 
-    override 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" + _)
+      result.toMap
     }
-
   }
 
   sealed trait LogOffsetResult
@@ -698,10 +953,8 @@ object ConsumerGroupCommand extends Logging {
     case object Ignore extends LogOffsetResult
   }
 
-  class ConsumerGroupCommandOptions(args: Array[String]) {
-    val ZkConnectDoc = "REQUIRED (for consumer groups based on the old consumer): The connection string for the zookeeper connection in the form host:port. " +
-      "Multiple URLS can be given to allow fail-over."
-    val BootstrapServerDoc = "REQUIRED (for consumer groups based on the new consumer): The server to connect to."
+  class ConsumerGroupCommandOptions(args: Array[String]) extends CommandDefaultOptions(args) {
+    val BootstrapServerDoc = "REQUIRED: The server(s) to connect to."
     val GroupDoc = "The consumer group we wish to act on."
     val TopicDoc = "The topic whose consumer group information should be deleted or topic whose should be included in the reset offset process. " +
       "In `reset-offsets` case, partitions can be specified using this format: `topic1:0,1,2`, where 0,1,2 are the partition to be included in the process. " +
@@ -709,24 +962,21 @@ 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" + nl +
-      "Pass in groups with a single topic to just delete the given topic's partition offsets and ownership " +
-      "information for the given consumer groups. For instance --group g1 --group g2 --topic t1" + nl +
-      "Pass in just a topic to delete the given topic's partition offsets and ownership information " +
-      "for every consumer group. For instance --topic t1" + nl +
-      "WARNING: Group deletion only works for old ZK-based consumer groups, and one has to use it carefully to only delete groups that are not active."
-    val NewConsumerDoc = "Use the new consumer implementation. This is the default, so this option is deprecated and " +
-      "will be removed in a future release."
+      "over the entire consumer group. For instance --group g1 --group g2"
     val TimeoutMsDoc = "The timeout that can be set for some use cases. For example, it can be used when describing the group " +
       "to specify the maximum amount of time in milliseconds to wait before the group stabilizes (when the group is just created, " +
       "or is going through some changes)."
     val CommandConfigDoc = "Property file containing configs to be passed to Admin Client and Consumer."
     val ResetOffsetsDoc = "Reset offsets of consumer group. Supports one consumer group at the time, and instances should be inactive" + nl +
-      "Has 3 execution options: (default) to plan which offsets to reset, --execute to execute the reset-offsets process, and --export to export the results to a CSV format." + nl +
-      "Has the following scenarios to choose: --to-datetime, --by-period, --to-earliest, --to-latest, --shift-by, --from-file, --to-current. One scenario must be choose" + nl +
-      "To define the scope use: --all-topics or --topic. . One scope must be choose, unless you use '--from-file' scenario"
+      "Has 2 execution options: --dry-run (the default) to plan which offsets to reset, and --execute to update the offsets. " +
+      "Additionally, the --export option is used to export the results to a CSV format." + nl +
+      "You must choose one of the following reset specifications: --to-datetime, --by-period, --to-earliest, " +
+      "--to-latest, --shift-by, --from-file, --to-current." + nl +
+      "To define the scope use --all-topics or --topic. One scope must be specified unless you use '--from-file'."
+    val DryRunDoc = "Only show results without executing changes on Consumer Groups. Supported operations: reset-offsets."
     val ExecuteDoc = "Execute operation. Supported operations: reset-offsets."
     val ExportDoc = "Export operation execution to a CSV file. Supported operations: reset-offsets."
     val ResetToOffsetDoc = "Reset offsets to a specific offset."
@@ -736,13 +986,21 @@ object ConsumerGroupCommand extends Logging {
     val ResetToEarliestDoc = "Reset offsets to earliest offset."
     val ResetToLatestDoc = "Reset offsets to latest offset."
     val ResetToCurrentDoc = "Reset offsets to current offset."
-    val ResetShiftByDoc = "Reset offsets shifting current offset by 'n', where 'n' can be positive or negative"
+    val ResetShiftByDoc = "Reset offsets shifting current offset by 'n', where 'n' can be positive or negative."
+    val MembersDoc = "Describe members of the group. This option may be used with '--describe' and '--bootstrap-server' options only." + nl +
+      "Example: --bootstrap-server localhost:9092 --describe --group group1 --members"
+    val VerboseDoc = "Provide additional information, if any, when describing the group. This option may be used " +
+      "with '--offsets'/'--members'/'--state' and '--bootstrap-server' options only." + nl + "Example: --bootstrap-server localhost:9092 --describe --group group1 --members --verbose"
+    val OffsetsDoc = "Describe the group and list all topic partitions in the group along with their offset lag. " +
+      "This is the default sub-action of and may be used with '--describe' and '--bootstrap-server' options only." + nl +
+      "Example: --bootstrap-server localhost:9092 --describe --group group1 --offsets"
+    val StateDoc = "When specified with '--describe', includes the state of the group." + nl +
+      "Example: --bootstrap-server localhost:9092 --describe --group group1 --state" + nl +
+      "When specified with '--list', it displays the state of all groups. It can also be used to list groups with specific states." + nl +
+      "Example: --bootstrap-server localhost:9092 --list --state stable,empty" + nl +
+      "This option may be used with '--describe', '--list' and '--bootstrap-server' options only."
+    val DeleteOffsetsDoc = "Delete offsets of consumer group. Supports one consumer group at the time, and multiple topics."
 
-    val parser = new OptionParser(false)
-    val zkConnectOpt = parser.accepts("zookeeper", ZkConnectDoc)
-                             .withRequiredArg
-                             .describedAs("urls")
-                             .ofType(classOf[String])
     val bootstrapServerOpt = parser.accepts("bootstrap-server", BootstrapServerDoc)
                                    .withRequiredArg
                                    .describedAs("server to connect to")
@@ -758,8 +1016,8 @@ 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 newConsumerOpt = parser.accepts("new-consumer", NewConsumerDoc)
     val timeoutMsOpt = parser.accepts("timeout", TimeoutMsDoc)
                              .withRequiredArg
                              .describedAs("timeout (ms)")
@@ -770,6 +1028,8 @@ 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)
     val resetToOffsetOpt = parser.accepts("to-offset", ResetToOffsetDoc)
@@ -795,56 +1055,86 @@ object ConsumerGroupCommand extends Logging {
                              .withRequiredArg()
                              .describedAs("number-of-offsets")
                              .ofType(classOf[Long])
-    val options = parser.parse(args : _*)
+    val membersOpt = parser.accepts("members", MembersDoc)
+                           .availableIf(describeOpt)
+    val verboseOpt = parser.accepts("verbose", VerboseDoc)
+                           .availableIf(describeOpt)
+    val offsetsOpt = parser.accepts("offsets", OffsetsDoc)
+                           .availableIf(describeOpt)
+    val stateOpt = parser.accepts("state", StateDoc)
+                         .availableIf(describeOpt, listOpt)
+                         .withOptionalArg()
+                         .ofType(classOf[String])
 
-    val useOldConsumer = options.has(zkConnectOpt)
-    val describeOptPresent = options.has(describeOpt)
+    options = parser.parse(args : _*)
 
-    val allConsumerGroupLevelOpts: Set[OptionSpec[_]] = Set(listOpt, describeOpt, deleteOpt, resetOffsetsOpt)
-    val allResetOffsetScenarioOpts: Set[OptionSpec[_]] = Set(resetToOffsetOpt, resetShiftByOpt,
+    val allGroupSelectionScopeOpts = immutable.Set[OptionSpec[_]](groupOpt, allGroupsOpt)
+    val allConsumerGroupLevelOpts = immutable.Set[OptionSpec[_]](listOpt, describeOpt, deleteOpt, resetOffsetsOpt)
+    val allResetOffsetScenarioOpts = immutable.Set[OptionSpec[_]](resetToOffsetOpt, resetShiftByOpt,
       resetToDatetimeOpt, resetByDurationOpt, resetToEarliestOpt, resetToLatestOpt, resetToCurrentOpt, resetFromFileOpt)
+    val allDeleteOffsetsOpts = immutable.Set[OptionSpec[_]](groupOpt, topicOpt)
 
-    def checkArgs() {
-      // check required args
-      if (options.has(timeoutMsOpt) && (!describeOptPresent || useOldConsumer))
-        debug(s"Option '$timeoutMsOpt' is applicable only when both '$bootstrapServerOpt' and '$describeOpt' are used.")
+    def checkArgs(): Unit = {
 
-      if (useOldConsumer) {
-        if (options.has(bootstrapServerOpt))
-          CommandLineUtils.printUsageAndDie(parser, s"Option '$bootstrapServerOpt' is not valid with '$zkConnectOpt'.")
-        else if (options.has(newConsumerOpt))
-          CommandLineUtils.printUsageAndDie(parser, s"Option '$newConsumerOpt' is not valid with '$zkConnectOpt'.")
-      } else {
-        CommandLineUtils.checkRequiredArgs(parser, options, bootstrapServerOpt)
+      CommandLineUtils.checkRequiredArgs(parser, options, bootstrapServerOpt)
 
-        if (options.has(newConsumerOpt)) {
-          Console.err.println("The --new-consumer option is deprecated and will be removed in a future major release." +
-            "The new consumer is used by default if the --bootstrap-server option is provided.")
+      if (options.has(describeOpt)) {
+        if (!options.has(groupOpt) && !options.has(allGroupsOpt))
+          CommandLineUtils.printUsageAndDie(parser,
+            s"Option $describeOpt takes one of these options: ${allGroupSelectionScopeOpts.mkString(", ")}")
+        val mutuallyExclusiveOpts: Set[OptionSpec[_]] = Set(membersOpt, offsetsOpt, stateOpt)
+        if (mutuallyExclusiveOpts.toList.map(o => if (options.has(o)) 1 else 0).sum > 1) {
+          CommandLineUtils.printUsageAndDie(parser,
+            s"Option $describeOpt takes at most one of these options: ${mutuallyExclusiveOpts.mkString(", ")}")
         }
+        if (options.has(stateOpt) && options.valueOf(stateOpt) != null)
+          CommandLineUtils.printUsageAndDie(parser,
+            s"Option $describeOpt does not take a value for $stateOpt")
+      } else {
+        if (options.has(timeoutMsOpt))
+          debug(s"Option $timeoutMsOpt is applicable only when $describeOpt is used.")
+      }
+
+      if (options.has(deleteOpt)) {
+        if (!options.has(groupOpt) && !options.has(allGroupsOpt))
+          CommandLineUtils.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 (options.has(deleteOpt))
-          CommandLineUtils.printUsageAndDie(parser, s"Option '$deleteOpt' is only valid with '$zkConnectOpt'. Note that " +
-            "there's no need to delete group metadata for the new consumer as the group is deleted when the last " +
-            "committed offset for that group expires.")
+      if (options.has(deleteOffsetsOpt)) {
+        if (!options.has(groupOpt) || !options.has(topicOpt))
+          CommandLineUtils.printUsageAndDie(parser,
+            s"Option $deleteOffsetsOpt takes the following options: ${allDeleteOffsetsOpts.mkString(", ")}")
       }
 
-      if (describeOptPresent)
-        CommandLineUtils.checkRequiredArgs(parser, options, groupOpt)
-      if (options.has(deleteOpt) && !options.has(groupOpt) && !options.has(topicOpt))
-        CommandLineUtils.printUsageAndDie(parser, "Option %s either takes %s, %s, or both".format(deleteOpt, groupOpt, topicOpt))
-      if (options.has(resetOffsetsOpt))
-        CommandLineUtils.checkRequiredArgs(parser, options, groupOpt)
-        CommandLineUtils.checkInvalidArgs(parser, options, resetToOffsetOpt, allResetOffsetScenarioOpts - resetToOffsetOpt)
+      if (options.has(resetOffsetsOpt)) {
+        if (options.has(dryRunOpt) && options.has(executeOpt))
+          CommandLineUtils.printUsageAndDie(parser, s"Option $resetOffsetsOpt only accepts one of $executeOpt and $dryRunOpt")
+
+        if (!options.has(dryRunOpt) && !options.has(executeOpt)) {
+          Console.err.println("WARN: No action will be performed as the --execute option is missing." +
+            "In a future major release, the default behavior of this command will be to prompt the user before " +
+            "executing the reset rather than doing a dry run. You should add the --dry-run option explicitly " +
+            "if you are scripting this command and want to keep the current default behavior without prompting.")
+        }
+
+        if (!options.has(groupOpt) && !options.has(allGroupsOpt))
+          CommandLineUtils.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
new file mode 100644
index 0000000000000..6465b143e3f33
--- /dev/null
+++ b/core/src/main/scala/kafka/admin/DelegationTokenCommand.scala
@@ -0,0 +1,219 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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.text.SimpleDateFormat
+import java.util
+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.{Admin, CreateDelegationTokenOptions, DescribeDelegationTokenOptions, ExpireDelegationTokenOptions, RenewDelegationTokenOptions}
+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}
+
+import scala.jdk.CollectionConverters._
+import scala.collection.Set
+
+/**
+ * A command to manage delegation token.
+ */
+object DelegationTokenCommand extends Logging {
+
+  def main(args: Array[String]): Unit = {
+    val opts = new DelegationTokenCommandOptions(args)
+
+    CommandLineUtils.printHelpAndExitIfNeeded(opts, "This tool helps to create, renew, expire, or describe delegation tokens.")
+
+    // should have exactly one action
+    val actions = Seq(opts.createOpt, opts.renewOpt, opts.expiryOpt, opts.describeOpt).count(opts.options.has _)
+    if(actions != 1)
+      CommandLineUtils.printUsageAndDie(opts.parser, "Command must include exactly one action: --create, --renew, --expire or --describe")
+
+    opts.checkArgs()
+
+    val adminClient = createAdminClient(opts)
+
+    var exitCode = 0
+    try {
+      if(opts.options.has(opts.createOpt))
+        createToken(adminClient, opts)
+      else if(opts.options.has(opts.renewOpt))
+        renewToken(adminClient, opts)
+      else if(opts.options.has(opts.expiryOpt))
+        expireToken(adminClient, opts)
+      else if(opts.options.has(opts.describeOpt))
+        describeToken(adminClient, opts)
+    } catch {
+      case e: Throwable =>
+        println("Error while executing delegation token command : " + e.getMessage)
+        error(Utils.stackTrace(e))
+        exitCode = 1
+    } finally {
+      adminClient.close()
+      Exit.exit(exitCode)
+    }
+  }
+
+  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
+
+    println("Calling create token operation with renewers :" + renewerPrincipals +" , max-life-time-period :"+ maxLifeTimeMs)
+    val createDelegationTokenOptions = new CreateDelegationTokenOptions().maxlifeTimeMs(maxLifeTimeMs).renewers(renewerPrincipals)
+    val createResult = adminClient.createDelegationToken(createDelegationTokenOptions)
+    val token = createResult.delegationToken().get()
+    println("Created delegation token with tokenId : %s".format(token.tokenInfo.tokenId)); printToken(List(token))
+    token
+  }
+
+  def printToken(tokens: List[DelegationToken]): Unit = {
+    val dateFormat = new SimpleDateFormat("yyyy-MM-dd'T'HH:mm")
+    print("\n%-15s %-30s %-15s %-25s %-15s %-15s %-15s".format("TOKENID", "HMAC", "OWNER", "RENEWERS", "ISSUEDATE", "EXPIRYDATE", "MAXDATE"))
+    for (token <- tokens) {
+      val tokenInfo = token.tokenInfo
+      print("\n%-15s %-30s %-15s %-25s %-15s %-15s %-15s".format(
+        tokenInfo.tokenId,
+        token.hmacAsBase64String,
+        tokenInfo.owner,
+        tokenInfo.renewersAsString,
+        dateFormat.format(tokenInfo.issueTimestamp),
+        dateFormat.format(tokenInfo.expiryTimestamp),
+        dateFormat.format(tokenInfo.maxTimestamp)))
+      println()
+    }
+  }
+
+  private def getPrincipals(opts: DelegationTokenCommandOptions, principalOptionSpec: ArgumentAcceptingOptionSpec[String]): Option[util.List[KafkaPrincipal]] = {
+    if (opts.options.has(principalOptionSpec))
+      Some(opts.options.valuesOf(principalOptionSpec).asScala.map(s => SecurityUtils.parseKafkaPrincipal(s.trim)).toList.asJava)
+    else
+      None
+  }
+
+  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)
+    val renewResult = adminClient.renewDelegationToken(Base64.getDecoder.decode(hmac), new RenewDelegationTokenOptions().renewTimePeriodMs(renewTimePeriodMs))
+    val expiryTimeStamp = renewResult.expiryTimestamp().get()
+    val dateFormat = new SimpleDateFormat("yyyy-MM-dd'T'HH:mm")
+    println("Completed renew operation. New expiry date : %s".format(dateFormat.format(expiryTimeStamp)))
+    expiryTimeStamp
+  }
+
+  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)
+    val expireResult = adminClient.expireDelegationToken(Base64.getDecoder.decode(hmac), new ExpireDelegationTokenOptions().expiryTimePeriodMs(expiryTimePeriodMs))
+    val expiryTimeStamp = expireResult.expiryTimestamp().get()
+    val dateFormat = new SimpleDateFormat("yyyy-MM-dd'T'HH:mm")
+    println("Completed expire operation. New expiry date : %s".format(dateFormat.format(expiryTimeStamp)))
+    expiryTimeStamp
+  }
+
+  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.")
+    else
+      println("Calling describe token operation for owners :" + ownerPrincipals.get)
+
+    val describeResult = adminClient.describeDelegationToken(new DescribeDelegationTokenOptions().owners(ownerPrincipals.orNull))
+    val tokens = describeResult.delegationTokens().get().asScala.toList
+    println("Total number of tokens : %s".format(tokens.size)); printToken(tokens)
+    tokens
+  }
+
+  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))
+    Admin.create(props)
+  }
+
+  class DelegationTokenCommandOptions(args: Array[String]) extends CommandDefaultOptions(args) {
+    val BootstrapServerDoc = "REQUIRED: server(s) to use for bootstrapping."
+    val CommandConfigDoc = "REQUIRED: A property file containing configs to be passed to Admin Client. Token management" +
+      " operations are allowed in secure mode only. This config file is used to pass security related configs."
+
+    val bootstrapServerOpt = parser.accepts("bootstrap-server", BootstrapServerDoc)
+                                   .withRequiredArg
+                                   .ofType(classOf[String])
+    val commandConfigOpt = parser.accepts("command-config", CommandConfigDoc)
+      .withRequiredArg
+      .ofType(classOf[String])
+
+    val createOpt = parser.accepts("create", "Create a new delegation token. Use --renewer-principal option to pass renewers principals.")
+    val renewOpt = parser.accepts("renew",  "Renew delegation token. Use --renew-time-period option to set renew time period.")
+    val expiryOpt = parser.accepts("expire", "Expire delegation token. Use --expiry-time-period option to expire the token.")
+    val describeOpt = parser.accepts("describe", "Describe delegation tokens for the given principals. Use --owner-principal to pass owner/renewer principals." +
+      " If --owner-principal option is not supplied, all the user owned tokens and tokens where user have Describe permission will be returned.")
+
+    val ownerPrincipalsOpt = parser.accepts("owner-principal", "owner is a kafka principal. It is should be in principalType:name format.")
+      .withOptionalArg()
+      .ofType(classOf[String])
+
+    val renewPrincipalsOpt = parser.accepts("renewer-principal", "renewer is a kafka principal. It is should be in principalType:name format.")
+      .withOptionalArg()
+      .ofType(classOf[String])
+
+    val maxLifeTimeOpt = parser.accepts("max-life-time-period", "Max life period for the token in milliseconds. If the value is -1," +
+      " then token max life time will default to a server side config value (delegation.token.max.lifetime.ms).")
+      .withOptionalArg()
+      .ofType(classOf[Long])
+
+    val renewTimePeriodOpt = parser.accepts("renew-time-period", "Renew time period in milliseconds. If the value is -1, then the" +
+      " renew time period will default to a server side config value (delegation.token.expiry.time.ms).")
+      .withOptionalArg()
+      .ofType(classOf[Long])
+
+    val expiryTimePeriodOpt = parser.accepts("expiry-time-period", "Expiry time period in milliseconds. If the value is -1, then the" +
+      " token will get invalidated immediately." )
+      .withOptionalArg()
+      .ofType(classOf[Long])
+
+    val hmacOpt = parser.accepts("hmac", "HMAC of the delegation token")
+      .withOptionalArg
+      .ofType(classOf[String])
+
+    options = parser.parse(args : _*)
+
+    def checkArgs(): Unit = {
+      // check required args
+      CommandLineUtils.checkRequiredArgs(parser, options, bootstrapServerOpt, commandConfigOpt)
+
+      if (options.has(createOpt))
+        CommandLineUtils.checkRequiredArgs(parser, options, maxLifeTimeOpt)
+
+      if (options.has(renewOpt))
+        CommandLineUtils.checkRequiredArgs(parser, options, hmacOpt, renewTimePeriodOpt)
+
+      if (options.has(expiryOpt))
+        CommandLineUtils.checkRequiredArgs(parser, options, hmacOpt, expiryTimePeriodOpt)
+
+      // check invalid args
+      CommandLineUtils.checkInvalidArgs(parser, options, createOpt, Set(hmacOpt, renewTimePeriodOpt, expiryTimePeriodOpt, ownerPrincipalsOpt))
+      CommandLineUtils.checkInvalidArgs(parser, options, renewOpt, Set(renewPrincipalsOpt, maxLifeTimeOpt, expiryTimePeriodOpt, ownerPrincipalsOpt))
+      CommandLineUtils.checkInvalidArgs(parser, options, expiryOpt, Set(renewOpt, maxLifeTimeOpt, renewTimePeriodOpt, ownerPrincipalsOpt))
+      CommandLineUtils.checkInvalidArgs(parser, options, describeOpt, Set(renewTimePeriodOpt, maxLifeTimeOpt, hmacOpt, renewTimePeriodOpt, expiryTimePeriodOpt))
+    }
+  }
+}
diff --git a/core/src/main/scala/kafka/admin/DeleteRecordsCommand.scala b/core/src/main/scala/kafka/admin/DeleteRecordsCommand.scala
index 2715490ec2368..71ef6fd6f19b8 100644
--- a/core/src/main/scala/kafka/admin/DeleteRecordsCommand.scala
+++ b/core/src/main/scala/kafka/admin/DeleteRecordsCommand.scala
@@ -20,33 +20,54 @@ package kafka.admin
 import java.io.PrintStream
 import java.util.Properties
 
-import kafka.admin.AdminClient.DeleteRecordsResult
 import kafka.common.AdminCommandFailedException
-import kafka.utils.{CoreUtils, Json, CommandLineUtils}
+import kafka.utils.json.JsonValue
+import kafka.utils.{CommandDefaultOptions, CommandLineUtils, CoreUtils, Json}
+import org.apache.kafka.clients.admin.{Admin, RecordsToDelete}
+import org.apache.kafka.clients.CommonClientConfigs
 import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.utils.Utils
-import org.apache.kafka.clients.CommonClientConfigs
-import joptsimple._
+
+import scala.jdk.CollectionConverters._
+import scala.collection.Seq
 
 /**
  * A command for delete records of the given partitions down to the specified offset.
  */
 object DeleteRecordsCommand {
 
+  private[admin] val EarliestVersion = 1
+
   def main(args: Array[String]): Unit = {
     execute(args, System.out)
   }
 
   def parseOffsetJsonStringWithoutDedup(jsonData: String): Seq[(TopicPartition, Long)] = {
-    Json.parseFull(jsonData).toSeq.flatMap { js =>
-      js.asJsonObject.get("partitions").toSeq.flatMap { partitionsJs =>
-        partitionsJs.asJsonArray.iterator.map(_.asJsonObject).map { partitionJs =>
-          val topic = partitionJs("topic").to[String]
-          val partition = partitionJs("partition").to[Int]
-          val offset = partitionJs("offset").to[Long]
-          new TopicPartition(topic, partition) -> offset
-        }.toBuffer
-      }
+    Json.parseFull(jsonData) match {
+      case Some(js) =>
+        val version = js.asJsonObject.get("version") match {
+          case Some(jsonValue) => jsonValue.to[Int]
+          case None => EarliestVersion
+        }
+        parseJsonData(version, js)
+      case None => throw new AdminOperationException("The input string is not a valid JSON")
+    }
+  }
+
+  def parseJsonData(version: Int, js: JsonValue): Seq[(TopicPartition, Long)] = {
+    version match {
+      case 1 =>
+        js.asJsonObject.get("partitions") match {
+          case Some(partitions) =>
+            partitions.asJsonArray.iterator.map(_.asJsonObject).map { partitionJs =>
+              val topic = partitionJs("topic").to[String]
+              val partition = partitionJs("partition").to[Int]
+              val offset = partitionJs("offset").to[Long]
+              new TopicPartition(topic, partition) -> offset
+            }.toBuffer
+          case _ => throw new AdminOperationException("Missing partitions field");
+        }
+      case _ => throw new AdminOperationException(s"Not supported version field value $version")
     }
   }
 
@@ -61,35 +82,39 @@ object DeleteRecordsCommand {
     if (duplicatePartitions.nonEmpty)
       throw new AdminCommandFailedException("Offset json file contains duplicate topic partitions: %s".format(duplicatePartitions.mkString(",")))
 
+    val recordsToDelete = offsetSeq.map { case (topicPartition, offset) =>
+      (topicPartition, RecordsToDelete.beforeOffset(offset))
+    }.toMap.asJava
+
     out.println("Executing records delete operation")
-    val deleteRecordsResult: Map[TopicPartition, DeleteRecordsResult] = adminClient.deleteRecordsBefore(offsetSeq.toMap).get()
+    val deleteRecordsResult = adminClient.deleteRecords(recordsToDelete)
     out.println("Records delete operation completed:")
 
-    deleteRecordsResult.foreach{ case (tp, partitionResult) => {
-      if (partitionResult.error == null)
-        out.println(s"partition: $tp\tlow_watermark: ${partitionResult.lowWatermark}")
-      else
-        out.println(s"partition: $tp\terror: ${partitionResult.error.toString}")
-    }}
+    deleteRecordsResult.lowWatermarks.forEach { (tp, partitionResult) =>
+      try out.println(s"partition: $tp\tlow_watermark: ${partitionResult.get.lowWatermark}")
+      catch {
+        case e: Exception => out.println(s"partition: $tp\terror: ${e.getMessage}")
+      }
+    }
+
     adminClient.close()
   }
 
-  private def createAdminClient(opts: DeleteRecordsCommandOptions): AdminClient = {
+  private def createAdminClient(opts: DeleteRecordsCommandOptions): 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))
-    AdminClient.create(props)
+    Admin.create(props)
   }
 
-  class DeleteRecordsCommandOptions(args: Array[String]) {
+  class DeleteRecordsCommandOptions(args: Array[String]) extends CommandDefaultOptions(args) {
     val BootstrapServerDoc = "REQUIRED: The server to connect to."
     val offsetJsonFileDoc = "REQUIRED: The JSON file with offset per partition. The format to use is:\n" +
                                  "{\"partitions\":\n  [{\"topic\": \"foo\", \"partition\": 1, \"offset\": 1}],\n \"version\":1\n}"
     val CommandConfigDoc = "A property file containing configs to be passed to Admin Client."
 
-    val parser = new OptionParser(false)
     val bootstrapServerOpt = parser.accepts("bootstrap-server", BootstrapServerDoc)
                                    .withRequiredArg
                                    .describedAs("server(s) to use for bootstrapping")
@@ -103,7 +128,10 @@ object DeleteRecordsCommand {
                                    .describedAs("command config property file path")
                                    .ofType(classOf[String])
 
-    val options = parser.parse(args : _*)
+    options = parser.parse(args : _*)
+
+    CommandLineUtils.printHelpAndExitIfNeeded(this, "This tool helps to delete records of the given partitions down to the specified offset.")
+
     CommandLineUtils.checkRequiredArgs(parser, options, bootstrapServerOpt, offsetJsonFileOpt)
   }
 }
diff --git a/core/src/main/scala/kafka/admin/FeatureCommand.scala b/core/src/main/scala/kafka/admin/FeatureCommand.scala
new file mode 100644
index 0000000000000..aa2b93ef77018
--- /dev/null
+++ b/core/src/main/scala/kafka/admin/FeatureCommand.scala
@@ -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 kafka.admin
+
+import kafka.server.BrokerFeatures
+import kafka.utils.{CommandDefaultOptions, CommandLineUtils, Exit}
+import org.apache.kafka.clients.CommonClientConfigs
+import org.apache.kafka.clients.admin.{Admin, FeatureUpdate, UpdateFeaturesOptions}
+import org.apache.kafka.common.feature.{Features, SupportedVersionRange}
+import org.apache.kafka.common.utils.Utils
+import java.util.Properties
+
+import scala.collection.Seq
+import scala.collection.immutable.ListMap
+import scala.jdk.CollectionConverters._
+import joptsimple.OptionSpec
+
+import scala.concurrent.ExecutionException
+
+object FeatureCommand {
+
+  def main(args: Array[String]): Unit = {
+    val opts = new FeatureCommandOptions(args)
+    val featureApis = new FeatureApis(opts)
+    var exitCode = 0
+    try {
+      featureApis.execute()
+    } catch {
+      case e: IllegalArgumentException =>
+        printException(e)
+        opts.parser.printHelpOn(System.err)
+        exitCode = 1
+      case _: UpdateFeaturesException =>
+        exitCode = 1
+      case e: ExecutionException =>
+        val cause = if (e.getCause == null) e else e.getCause
+        printException(cause)
+        exitCode = 1
+      case e: Throwable =>
+        printException(e)
+        exitCode = 1
+    } finally {
+      featureApis.close()
+      Exit.exit(exitCode)
+    }
+  }
+
+  private def printException(exception: Throwable): Unit = {
+    System.err.println("\nError encountered when executing command: " + Utils.stackTrace(exception))
+  }
+}
+
+class UpdateFeaturesException(message: String) extends RuntimeException(message)
+
+/**
+ * A class that provides necessary APIs to bridge feature APIs provided by the the Admin client with
+ * the requirements of the CLI tool.
+ *
+ * @param opts the CLI options
+ */
+class FeatureApis(private var opts: FeatureCommandOptions) {
+  private var supportedFeatures = BrokerFeatures.createDefault().supportedFeatures
+  private var adminClient = FeatureApis.createAdminClient(opts)
+
+  private def pad(op: String): String = {
+    f"$op%11s"
+  }
+
+  private val addOp = pad("[Add]")
+  private val upgradeOp = pad("[Upgrade]")
+  private val deleteOp = pad("[Delete]")
+  private val downgradeOp = pad("[Downgrade]")
+
+  // For testing only.
+  private[admin] def setSupportedFeatures(newFeatures: Features[SupportedVersionRange]): Unit = {
+    supportedFeatures = newFeatures
+  }
+
+  // For testing only.
+  private[admin] def setOptions(newOpts: FeatureCommandOptions): Unit = {
+    adminClient.close()
+    adminClient = FeatureApis.createAdminClient(newOpts)
+    opts = newOpts
+  }
+
+  /**
+   * Describes the supported and finalized features. The request is issued to any of the provided
+   * bootstrap servers.
+   */
+  def describeFeatures(): Unit = {
+    val result = adminClient.describeFeatures.featureMetadata.get
+    val features = result.supportedFeatures.asScala.keys.toSet ++ result.finalizedFeatures.asScala.keys.toSet
+
+    features.toList.sorted.foreach {
+      feature =>
+        val output = new StringBuilder()
+        output.append(s"Feature: $feature")
+
+        val (supportedMinVersion, supportedMaxVersion) = {
+          val supportedVersionRange = result.supportedFeatures.get(feature)
+          if (supportedVersionRange == null) {
+            ("-", "-")
+          } else {
+            (supportedVersionRange.minVersion, supportedVersionRange.maxVersion)
+          }
+        }
+        output.append(s"\tSupportedMinVersion: $supportedMinVersion")
+        output.append(s"\tSupportedMaxVersion: $supportedMaxVersion")
+
+        val (finalizedMinVersionLevel, finalizedMaxVersionLevel) = {
+          val finalizedVersionRange = result.finalizedFeatures.get(feature)
+          if (finalizedVersionRange == null) {
+            ("-", "-")
+          } else {
+            (finalizedVersionRange.minVersionLevel, finalizedVersionRange.maxVersionLevel)
+          }
+        }
+        output.append(s"\tFinalizedMinVersionLevel: $finalizedMinVersionLevel")
+        output.append(s"\tFinalizedMaxVersionLevel: $finalizedMaxVersionLevel")
+
+        val epoch = {
+          if (result.finalizedFeaturesEpoch.isPresent) {
+            result.finalizedFeaturesEpoch.get.toString
+          } else {
+            "-"
+          }
+        }
+        output.append(s"\tEpoch: $epoch")
+
+        println(output)
+    }
+  }
+
+  /**
+   * Upgrades all features known to this tool to their highest max version levels. The method may
+   * add new finalized features if they were not finalized previously, but it does not delete
+   * any existing finalized feature. The results of the feature updates are written to STDOUT.
+   *
+   * NOTE: if the --dry-run CLI option is provided, this method only prints the expected feature
+   * updates to STDOUT, without applying them.
+   *
+   * @throws UpdateFeaturesException if at least one of the feature updates failed
+   */
+  def upgradeAllFeatures(): Unit = {
+    val metadata = adminClient.describeFeatures.featureMetadata.get
+    val existingFinalizedFeatures = metadata.finalizedFeatures
+    val updates = supportedFeatures.features.asScala.map {
+      case (feature, targetVersionRange) =>
+        val existingVersionRange = existingFinalizedFeatures.get(feature)
+        if (existingVersionRange == null) {
+          val updateStr =
+            addOp +
+            s"\tFeature: $feature" +
+            s"\tExistingFinalizedMaxVersion: -" +
+            s"\tNewFinalizedMaxVersion: ${targetVersionRange.max}"
+          (feature, Some((updateStr, new FeatureUpdate(targetVersionRange.max, false))))
+        } else {
+          if (targetVersionRange.max > existingVersionRange.maxVersionLevel) {
+            val updateStr =
+              upgradeOp +
+              s"\tFeature: $feature" +
+              s"\tExistingFinalizedMaxVersion: ${existingVersionRange.maxVersionLevel}" +
+              s"\tNewFinalizedMaxVersion: ${targetVersionRange.max}"
+            (feature, Some((updateStr, new FeatureUpdate(targetVersionRange.max, false))))
+          } else {
+            (feature, Option.empty)
+          }
+        }
+    }.filter {
+      case(_, updateInfo) => updateInfo.isDefined
+    }.map {
+      case(feature, updateInfo) => (feature, updateInfo.get)
+    }.toMap
+
+    if (updates.nonEmpty) {
+      maybeApplyFeatureUpdates(updates)
+    }
+  }
+
+  /**
+   * Downgrades existing finalized features to the highest max version levels known to this tool.
+   * The method may delete existing finalized features if they are no longer seen to be supported,
+   * but it does not add a feature that was not finalized previously. The results of the feature
+   * updates are written to STDOUT.
+   *
+   * NOTE: if the --dry-run CLI option is provided, this method only prints the expected feature
+   * updates to STDOUT, without applying them.
+   *
+   * @throws UpdateFeaturesException if at least one of the feature updates failed
+   */
+  def downgradeAllFeatures(): Unit = {
+    val metadata = adminClient.describeFeatures.featureMetadata.get
+    val existingFinalizedFeatures = metadata.finalizedFeatures
+    val supportedFeaturesMap = supportedFeatures.features
+    val updates = existingFinalizedFeatures.asScala.map {
+      case (feature, existingVersionRange) =>
+        val targetVersionRange = supportedFeaturesMap.get(feature)
+        if (targetVersionRange == null) {
+          val updateStr =
+            deleteOp +
+            s"\tFeature: $feature" +
+            s"\tExistingFinalizedMaxVersion: ${existingVersionRange.maxVersionLevel}" +
+            s"\tNewFinalizedMaxVersion: -"
+          (feature, Some(updateStr, new FeatureUpdate(0, true)))
+        } else {
+          if (targetVersionRange.max < existingVersionRange.maxVersionLevel) {
+            val updateStr =
+              downgradeOp +
+              s"\tFeature: $feature" +
+              s"\tExistingFinalizedMaxVersion: ${existingVersionRange.maxVersionLevel}" +
+              s"\tNewFinalizedMaxVersion: ${targetVersionRange.max}"
+            (feature, Some(updateStr, new FeatureUpdate(targetVersionRange.max, true)))
+          } else {
+            (feature, Option.empty)
+          }
+        }
+    }.filter {
+      case(_, updateInfo) => updateInfo.isDefined
+    }.map {
+      case(feature, updateInfo) => (feature, updateInfo.get)
+    }.toMap
+
+    if (updates.nonEmpty) {
+      maybeApplyFeatureUpdates(updates)
+    }
+  }
+
+  /**
+   * Applies the provided feature updates. If the --dry-run CLI option is provided, the method
+   * only prints the expected feature updates to STDOUT without applying them.
+   *
+   * @param updates the feature updates to be applied via the admin client
+   *
+   * @throws UpdateFeaturesException if at least one of the feature updates failed
+   */
+  private def maybeApplyFeatureUpdates(updates: Map[String, (String, FeatureUpdate)]): Unit = {
+    if (opts.hasDryRunOption) {
+      println("Expected feature updates:" + ListMap(
+        updates
+          .toSeq
+          .sortBy { case(feature, _) => feature} :_*)
+          .map { case(_, (updateStr, _)) => updateStr}
+          .mkString("\n"))
+    } else {
+      val result = adminClient.updateFeatures(
+        updates
+          .map { case(feature, (_, update)) => (feature, update)}
+          .asJava,
+        new UpdateFeaturesOptions())
+      val resultSortedByFeature = ListMap(
+        result
+          .values
+          .asScala
+          .toSeq
+          .sortBy { case(feature, _) => feature} :_*)
+      val failures = resultSortedByFeature.map {
+        case (feature, updateFuture) =>
+          val (updateStr, _) = updates(feature)
+          try {
+            updateFuture.get
+            println(updateStr + "\tResult: OK")
+            0
+          } catch {
+            case e: ExecutionException =>
+              val cause = if (e.getCause == null) e else e.getCause
+              println(updateStr + "\tResult: FAILED due to " + cause)
+              1
+            case e: Throwable =>
+              println(updateStr + "\tResult: FAILED due to " + e)
+              1
+          }
+      }.sum
+      if (failures > 0) {
+        throw new UpdateFeaturesException(s"$failures feature updates failed!")
+      }
+    }
+  }
+
+  def execute(): Unit = {
+    if (opts.hasDescribeOption) {
+      describeFeatures()
+    } else if (opts.hasUpgradeAllOption) {
+      upgradeAllFeatures()
+    } else if (opts.hasDowngradeAllOption) {
+      downgradeAllFeatures()
+    } else {
+      throw new IllegalStateException("Unexpected state: no CLI command could be executed.")
+    }
+  }
+
+  def close(): Unit = {
+    adminClient.close()
+  }
+}
+
+class FeatureCommandOptions(args: Array[String]) extends CommandDefaultOptions(args) {
+  private val bootstrapServerOpt = parser.accepts(
+      "bootstrap-server",
+      "REQUIRED: A comma-separated list of host:port pairs to use for establishing the connection" +
+      " to the Kafka cluster.")
+      .withRequiredArg
+      .describedAs("server to connect to")
+      .ofType(classOf[String])
+  private val commandConfigOpt = parser.accepts(
+    "command-config",
+    "Property file containing configs to be passed to Admin Client." +
+    " This is used with --bootstrap-server option when required.")
+    .withOptionalArg
+    .describedAs("command config property file")
+    .ofType(classOf[String])
+  private val describeOpt = parser.accepts(
+    "describe",
+    "Describe supported and finalized features from a random broker.")
+  private val upgradeAllOpt = parser.accepts(
+    "upgrade-all",
+    "Upgrades all finalized features to the maximum version levels known to the tool." +
+    " This command finalizes new features known to the tool that were never finalized" +
+    " previously in the cluster, but it is guaranteed to not delete any existing feature.")
+  private val downgradeAllOpt = parser.accepts(
+    "downgrade-all",
+    "Downgrades all finalized features to the maximum version levels known to the tool." +
+    " This command deletes unknown features from the list of finalized features in the" +
+    " cluster, but it is guaranteed to not add a new feature.")
+  private val dryRunOpt = parser.accepts(
+    "dry-run",
+    "Performs a dry-run of upgrade/downgrade mutations to finalized feature without applying them.")
+
+  options = parser.parse(args : _*)
+
+  checkArgs()
+
+  def has(builder: OptionSpec[_]): Boolean = options.has(builder)
+
+  def hasDescribeOption: Boolean = has(describeOpt)
+
+  def hasDryRunOption: Boolean = has(dryRunOpt)
+
+  def hasUpgradeAllOption: Boolean = has(upgradeAllOpt)
+
+  def hasDowngradeAllOption: Boolean = has(downgradeAllOpt)
+
+  def commandConfig: Properties = {
+    if (has(commandConfigOpt))
+      Utils.loadProps(options.valueOf(commandConfigOpt))
+    else
+      new Properties()
+  }
+
+  def bootstrapServers: String = options.valueOf(bootstrapServerOpt)
+
+  def checkArgs(): Unit = {
+    CommandLineUtils.printHelpAndExitIfNeeded(this, "This tool describes and updates finalized features.")
+    val numActions = Seq(describeOpt, upgradeAllOpt, downgradeAllOpt).count(has)
+    if (numActions != 1) {
+      CommandLineUtils.printUsageAndDie(
+        parser,
+        "Command must include exactly one action: --describe, --upgrade-all, --downgrade-all.")
+    }
+    CommandLineUtils.checkRequiredArgs(parser, options, bootstrapServerOpt)
+    if (hasDryRunOption && !hasUpgradeAllOption && !hasDowngradeAllOption) {
+      CommandLineUtils.printUsageAndDie(
+        parser,
+        "Command can contain --dry-run option only when either --upgrade-all or --downgrade-all actions are provided.")
+    }
+  }
+}
+
+object FeatureApis {
+  private def createAdminClient(opts: FeatureCommandOptions): Admin = {
+    val props = new Properties()
+    props.putAll(opts.commandConfig)
+    props.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, opts.bootstrapServers)
+    Admin.create(props)
+  }
+}
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..03737aa532294
--- /dev/null
+++ b/core/src/main/scala/kafka/admin/LeaderElectionCommand.scala
@@ -0,0 +1,289 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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.Implicits._
+import kafka.utils.Json
+import kafka.utils.Logging
+import org.apache.kafka.clients.admin.{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 org.apache.kafka.common.utils.Utils
+import scala.jdk.CollectionConverters._
+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.DEFAULT_API_TIMEOUT_MS_CONFIG, timeout.toMillis.toString)
+      props.setProperty(AdminClientConfig.REQUEST_TIMEOUT_MS_CONFIG, (timeout.toMillis / 2).toString)
+
+      Admin.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.forKeyValue { (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 6a167a23dcac4..ad0307f31b9cc 100644
--- a/core/src/main/scala/kafka/admin/LogDirsCommand.scala
+++ b/core/src/main/scala/kafka/admin/LogDirsCommand.scala
@@ -20,13 +20,12 @@ package kafka.admin
 import java.io.PrintStream
 import java.util.Properties
 
-import org.apache.kafka.clients.admin.{AdminClientConfig, DescribeLogDirsResult, AdminClient => JAdminClient}
-import org.apache.kafka.common.requests.DescribeLogDirsResponse.LogDirInfo
+import kafka.utils.{CommandDefaultOptions, CommandLineUtils, Json}
+import org.apache.kafka.clients.admin.{Admin, AdminClientConfig, DescribeLogDirsResult, LogDirDescription}
+import org.apache.kafka.common.utils.Utils
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 import scala.collection.Map
-import kafka.utils.{CommandLineUtils, Json}
-import joptsimple._
 
 /**
   * A command for querying log directory usage on the specified brokers
@@ -48,15 +47,15 @@ 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.allDescriptions.get().asScala.map { case (k, v) => k -> v.asScala }
 
         out.println(s"Received log directory information from brokers ${brokerList.mkString(",")}")
         out.println(formatAsJson(logDirInfosByBroker, topicList.toSet))
         adminClient.close()
     }
 
-    private def formatAsJson(logDirInfosByBroker: Map[Integer, Map[String, LogDirInfo]], topicSet: Set[String]): String = {
-        Json.encode(Map(
+    private def formatAsJson(logDirInfosByBroker: Map[Integer, Map[String, LogDirDescription]], topicSet: Set[String]): String = {
+        Json.encodeAsString(Map(
             "version" -> 1,
             "brokers" -> logDirInfosByBroker.map { case (broker, logDirInfos) =>
                 Map(
@@ -64,8 +63,8 @@ object LogDirsCommand {
                     "logDirs" -> logDirInfos.map { case (logDir, logDirInfo) =>
                         Map(
                             "logDir" -> logDir,
-                            "error" -> logDirInfo.error.exceptionName(),
-                            "partitions" -> logDirInfo.replicaInfos.asScala.filter { case (topicPartition, replicaInfo) =>
+                            "error" -> Option(logDirInfo.error).map(ex => ex.getClass.getName).orNull,
+                            "partitions" -> logDirInfo.replicaInfos.asScala.filter { case (topicPartition, _) =>
                                 topicSet.isEmpty || topicSet.contains(topicPartition.topic)
                             }.map { case (topicPartition, replicaInfo) =>
                                 Map(
@@ -73,28 +72,34 @@ object LogDirsCommand {
                                     "size" -> replicaInfo.size,
                                     "offsetLag" -> replicaInfo.offsetLag,
                                     "isFuture" -> replicaInfo.isFuture
-                                )
-                            }
-                        )
-                    }
-                )
-            }
-        ))
+                                ).asJava
+                            }.asJava
+                        ).asJava
+                    }.asJava
+                ).asJava
+            }.asJava
+        ).asJava)
     }
 
-    private def createAdminClient(opts: LogDirsCommandOptions): JAdminClient = {
-        val props = new Properties()
+    private def createAdminClient(opts: LogDirsCommandOptions): Admin = {
+        val props = if (opts.options.has(opts.commandConfigOpt))
+            Utils.loadProps(opts.options.valueOf(opts.commandConfigOpt))
+        else
+            new Properties()
         props.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, opts.options.valueOf(opts.bootstrapServerOpt))
-        props.put(AdminClientConfig.CLIENT_ID_CONFIG, "log-dirs-tool")
-        JAdminClient.create(props)
+        props.putIfAbsent(AdminClientConfig.CLIENT_ID_CONFIG, "log-dirs-tool")
+        Admin.create(props)
     }
 
-    class LogDirsCommandOptions(args: Array[String]) {
-        val parser = new OptionParser(false)
+    class LogDirsCommandOptions(args: Array[String]) extends CommandDefaultOptions(args){
         val bootstrapServerOpt = parser.accepts("bootstrap-server", "REQUIRED: the server(s) to use for bootstrapping")
           .withRequiredArg
           .describedAs("The server(s) to use for bootstrapping")
           .ofType(classOf[String])
+        val commandConfigOpt = parser.accepts("command-config", "Property file containing configs to be passed to Admin Client.")
+          .withRequiredArg
+          .describedAs("Admin client property file")
+          .ofType(classOf[String])
         val describeOpt = parser.accepts("describe", "Describe the specified log directories on the specified brokers.")
         val topicListOpt = parser.accepts("topic-list", "The list of topics to be queried in the form \"topic1,topic2,topic3\". " +
           "All topics will be queried if no topic list is specified")
@@ -108,7 +113,10 @@ object LogDirsCommand {
           .describedAs("Broker list")
           .ofType(classOf[String])
 
-        val options = parser.parse(args : _*)
+        options = parser.parse(args : _*)
+
+        CommandLineUtils.printHelpAndExitIfNeeded(this, "This tool helps to query log directory usage on the specified brokers.")
+
         CommandLineUtils.checkRequiredArgs(parser, options, bootstrapServerOpt, describeOpt)
     }
 }
diff --git a/core/src/main/scala/kafka/admin/PreferredReplicaLeaderElectionCommand.scala b/core/src/main/scala/kafka/admin/PreferredReplicaLeaderElectionCommand.scala
index c292fe61bea37..5d7620a9cfe15 100755
--- a/core/src/main/scala/kafka/admin/PreferredReplicaLeaderElectionCommand.scala
+++ b/core/src/main/scala/kafka/admin/PreferredReplicaLeaderElectionCommand.scala
@@ -16,68 +16,77 @@
  */
 package kafka.admin
 
-import joptsimple.OptionParser
-import kafka.utils._
-import org.I0Itec.zkclient.ZkClient
-import org.I0Itec.zkclient.exception.ZkNodeExistsException
-import kafka.common.{TopicAndPartition, AdminCommandFailedException}
+import scala.jdk.CollectionConverters._
 import collection._
-import org.apache.kafka.common.utils.Utils
+import java.util.Properties
+import java.util.concurrent.ExecutionException
+
+import joptsimple.OptionSpecBuilder
+import kafka.common.AdminCommandFailedException
+import kafka.utils._
+import kafka.utils.Implicits._
+import kafka.zk.KafkaZkClient
+import org.apache.kafka.clients.admin.{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 org.apache.kafka.common.security.JaasUtils
+import org.apache.kafka.common.utils.Time
+import org.apache.kafka.common.utils.Utils
+import org.apache.zookeeper.KeeperException.NodeExistsException
 
 object PreferredReplicaLeaderElectionCommand extends Logging {
 
   def main(args: Array[String]): Unit = {
-    val parser = new OptionParser(false)
-    val jsonFileOpt = parser.accepts("path-to-json-file", "The JSON file with the list of partitions " +
-      "for which preferred replica leader election should be done, in the following format - \n" +
-       "{\"partitions\":\n\t[{\"topic\": \"foo\", \"partition\": 1},\n\t {\"topic\": \"foobar\", \"partition\": 2}]\n}\n" +
-      "Defaults to all existing partitions")
-      .withRequiredArg
-      .describedAs("list of partitions for which preferred replica leader election needs to be triggered")
-      .ofType(classOf[String])
-    val zkConnectOpt = parser.accepts("zookeeper", "REQUIRED: The connection string for the zookeeper connection in the " +
-      "form host:port. Multiple URLS can be given to allow fail-over.")
-      .withRequiredArg
-      .describedAs("urls")
-      .ofType(classOf[String])
-      
-    if(args.length == 0)
-      CommandLineUtils.printUsageAndDie(parser, "This tool causes leadership for each partition to be transferred back to the 'preferred replica'," + 
-                                                " it can be used to balance leadership among the servers.")
+    val timeout = 30000
+    run(args, timeout)
+  }
 
-    val options = parser.parse(args : _*)
+  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.")
 
-    CommandLineUtils.checkRequiredArgs(parser, options, zkConnectOpt)
+    CommandLineUtils.checkRequiredArgs(commandOpts.parser, commandOpts.options)
 
-    val zkConnect = options.valueOf(zkConnectOpt)
-    var zkClient: ZkClient = null
-    var zkUtils: ZkUtils = null
-    try {
-      zkClient = ZkUtils.createZkClient(zkConnect, 30000, 30000)
-      zkUtils = ZkUtils(zkConnect, 
-                        30000,
-                        30000,
-                        JaasUtils.isZkSecurityEnabled())
-      val partitionsForPreferredReplicaElection =
-        if (!options.has(jsonFileOpt))
-          zkUtils.getAllPartitions()
+    if (commandOpts.options.has(commandOpts.bootstrapServerOpt) == commandOpts.options.has(commandOpts.zkConnectOpt)) {
+      CommandLineUtils.printUsageAndDie(commandOpts.parser, s"Exactly one of '${commandOpts.bootstrapServerOpt}' or '${commandOpts.zkConnectOpt}' must be provided")
+    }
+
+    val partitionsForPreferredReplicaElection =
+    if (commandOpts.options.has(commandOpts.jsonFileOpt))
+      Some(parsePreferredReplicaElectionData(Utils.readFileAsString(commandOpts.options.valueOf(commandOpts.jsonFileOpt))))
+    else
+      None
+
+    val preferredReplicaElectionCommand = if (commandOpts.options.has(commandOpts.zkConnectOpt)) {
+      println(s"Warning: --zookeeper is deprecated and will be removed in a future version of Kafka.")
+      println(s"Use --bootstrap-server instead to specify a broker to connect to.")
+      new ZkCommand(commandOpts.options.valueOf(commandOpts.zkConnectOpt),
+              JaasUtils.isZkSaslEnabled,
+              timeout)
+    } else {
+        val adminProps = if (commandOpts.options.has(commandOpts.adminClientConfigOpt))
+          Utils.loadProps(commandOpts.options.valueOf(commandOpts.adminClientConfigOpt))
         else
-          parsePreferredReplicaElectionData(Utils.readFileAsString(options.valueOf(jsonFileOpt)))
-      val preferredReplicaElectionCommand = new PreferredReplicaLeaderElectionCommand(zkUtils, partitionsForPreferredReplicaElection)
+          new Properties()
+        adminProps.setProperty(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, commandOpts.options.valueOf(commandOpts.bootstrapServerOpt))
+        adminProps.setProperty(AdminClientConfig.REQUEST_TIMEOUT_MS_CONFIG, timeout.toString)
+        adminProps.setProperty(AdminClientConfig.DEFAULT_API_TIMEOUT_MS_CONFIG, (timeout * 2).toString)
+        new AdminClientCommand(adminProps)
+    }
 
-      preferredReplicaElectionCommand.moveLeaderToPreferredReplica()
-    } catch {
-      case e: Throwable =>
-        println("Failed to start preferred replica election")
-        println(Utils.stackTrace(e))
+    try {
+      preferredReplicaElectionCommand.electPreferredLeaders(partitionsForPreferredReplicaElection)
     } finally {
-      if (zkClient != null)
-        zkClient.close()
+      preferredReplicaElectionCommand.close()
     }
   }
 
-  def parsePreferredReplicaElectionData(jsonString: String): immutable.Set[TopicAndPartition] = {
+  def parsePreferredReplicaElectionData(jsonString: String): collection.immutable.Set[TopicPartition] = {
     Json.parseFull(jsonString) match {
       case Some(js) =>
         js.asJsonObject.get("partitions") match {
@@ -86,7 +95,7 @@ object PreferredReplicaLeaderElectionCommand extends Logging {
             val partitions = partitionsRaw.map { p =>
               val topic = p("topic").to[String]
               val partition = p("partition").to[Int]
-              TopicAndPartition(topic, partition)
+              new TopicPartition(topic, partition)
             }.toBuffer
             val duplicatePartitions = CoreUtils.duplicates(partitions)
             if (duplicatePartitions.nonEmpty)
@@ -98,34 +107,194 @@ object PreferredReplicaLeaderElectionCommand extends Logging {
     }
   }
 
-  def writePreferredReplicaElectionData(zkUtils: ZkUtils,
-                                        partitionsUndergoingPreferredReplicaElection: scala.collection.Set[TopicAndPartition]) {
-    val zkPath = ZkUtils.PreferredReplicaLeaderElectionPath
-    val jsonData = ZkUtils.preferredReplicaLeaderElectionZkData(partitionsUndergoingPreferredReplicaElection)
+  def writePreferredReplicaElectionData(zkClient: KafkaZkClient,
+                                        partitionsUndergoingPreferredReplicaElection: Set[TopicPartition]): Unit = {
     try {
-      zkUtils.createPersistentPath(zkPath, jsonData)
-      println("Created preferred replica election path with %s".format(jsonData))
+      zkClient.createPreferredReplicaElection(partitionsUndergoingPreferredReplicaElection.toSet)
+      println("Created preferred replica election path with %s".format(partitionsUndergoingPreferredReplicaElection.mkString(",")))
     } catch {
-      case _: ZkNodeExistsException =>
-        val partitionsUndergoingPreferredReplicaElection =
-          PreferredReplicaLeaderElectionCommand.parsePreferredReplicaElectionData(zkUtils.readData(zkPath)._1)
+      case _: NodeExistsException =>
         throw new AdminOperationException("Preferred replica leader election currently in progress for " +
-          "%s. Aborting operation".format(partitionsUndergoingPreferredReplicaElection))
+          "%s. Aborting operation".format(zkClient.getPreferredReplicaElection.mkString(",")))
       case e2: Throwable => throw new AdminOperationException(e2.toString)
     }
   }
+
+  class PreferredReplicaLeaderElectionCommandOptions(args: Array[String]) extends CommandDefaultOptions(args) {
+    val jsonFileOpt = parser.accepts("path-to-json-file", "The JSON file with the list of partitions " +
+      "for which preferred replica leader election should be done, in the following format - \n" +
+      "{\"partitions\":\n\t[{\"topic\": \"foo\", \"partition\": 1},\n\t {\"topic\": \"foobar\", \"partition\": 2}]\n}\n" +
+      "Defaults to all existing partitions")
+      .withRequiredArg
+      .describedAs("list of partitions for which preferred replica leader election needs to be triggered")
+      .ofType(classOf[String])
+
+    private val zookeeperOptBuilder: OptionSpecBuilder = parser.accepts("zookeeper",
+      "DEPRECATED. The connection string for the zookeeper connection in the " +
+      "form host:port. Multiple URLS can be given to allow fail-over. " +
+      "Replaced by --bootstrap-server, REQUIRED unless --bootstrap-server is given.")
+    private val bootstrapOptBuilder: OptionSpecBuilder = 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 unless --zookeeper is given.")
+    parser.mutuallyExclusive(zookeeperOptBuilder, bootstrapOptBuilder)
+    val bootstrapServerOpt = bootstrapOptBuilder
+      .withRequiredArg
+      .describedAs("host:port")
+      .ofType(classOf[String])
+    val zkConnectOpt = zookeeperOptBuilder
+      .withRequiredArg
+      .describedAs("urls")
+      .ofType(classOf[String])
+
+    val adminClientConfigOpt = parser.accepts("admin.config",
+      "Admin client config properties file to pass to the admin client when --bootstrap-server is given.")
+      .availableIf(bootstrapServerOpt)
+      .withRequiredArg
+      .describedAs("config file")
+      .ofType(classOf[String])
+
+    options = parser.parse(args: _*)
+  }
+
+  /** Abstraction over different ways to perform a leader election */
+  trait Command {
+    /** 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
+  }
+
+  class ZkCommand(zkConnect: String, isSecure: Boolean, timeout: Int)
+    extends Command {
+    var zkClient: KafkaZkClient = null
+
+    val time = Time.SYSTEM
+    zkClient = KafkaZkClient(zkConnect, isSecure, timeout, timeout, Int.MaxValue, time)
+
+    override def electPreferredLeaders(partitionsFromUser: Option[Set[TopicPartition]]): Unit = {
+      try {
+        val topics =
+          partitionsFromUser match {
+            case Some(partitions) =>
+              partitions.map(_.topic).toSet
+            case None =>
+              zkClient.getAllPartitions.map(_.topic)
+          }
+
+        val partitionsFromZk = zkClient.getPartitionsForTopics(topics).flatMap{ case (topic, partitions) =>
+          partitions.map(new TopicPartition(topic, _))
+        }.toSet
+
+        val (validPartitions, invalidPartitions) =
+          partitionsFromUser match {
+            case Some(partitions) =>
+              partitions.partition(partitionsFromZk.contains)
+            case None =>
+              (zkClient.getAllPartitions, Set.empty)
+          }
+          PreferredReplicaLeaderElectionCommand.writePreferredReplicaElectionData(zkClient, validPartitions)
+
+        println("Successfully started preferred replica election for partitions %s".format(validPartitions))
+        invalidPartitions.foreach(p => println("Skipping preferred replica leader election for partition %s since it doesn't exist.".format(p)))
+      } catch {
+        case e: Throwable => throw new AdminCommandFailedException("Admin command failed", e)
+      }
+    }
+
+    override def close(): Unit = {
+      if (zkClient != null)
+        zkClient.close()
+    }
+  }
+
+  /** Election via AdminClient.electPreferredLeaders() */
+  class AdminClientCommand(adminClientProps: Properties)
+    extends Command with Logging {
+
+    val adminClient = Admin.create(adminClientProps)
+
+    override def electPreferredLeaders(partitionsFromUser: Option[Set[TopicPartition]]): Unit = {
+      val partitions = partitionsFromUser match {
+        case Some(partitionsFromUser) => partitionsFromUser.asJava
+        case None => null
+      }
+      debug(s"Calling AdminClient.electLeaders(ElectionType.PREFERRED, $partitions)")
+
+      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)
+          }
+
+          throw e
+        case e: Throwable =>
+          // We don't even know the attempted partitions
+          println("Error while making request")
+          e.printStackTrace()
+          return
+      }
+
+      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
+          }
+        } else {
+          succeeded += topicPartition
+        }
+      }
+
+      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.forKeyValue { (topicPartition, exception) =>
+          println(s"Error completing preferred leader election for partition: $topicPartition: $exception")
+          rootException.addSuppressed(exception)
+        }
+        throw rootException
+      }
+    }
+
+    override def close(): Unit = {
+      debug("Closing AdminClient")
+      adminClient.close()
+    }
+  }
 }
 
-class PreferredReplicaLeaderElectionCommand(zkUtils: ZkUtils, partitionsFromUser: scala.collection.Set[TopicAndPartition]) {
-  def moveLeaderToPreferredReplica() = {
+class PreferredReplicaLeaderElectionCommand(zkClient: KafkaZkClient, partitionsFromUser: scala.collection.Set[TopicPartition]) {
+  def moveLeaderToPreferredReplica(): Unit = {
     try {
       val topics = partitionsFromUser.map(_.topic).toSet
-      val partitionsFromZk = zkUtils.getPartitionsForTopics(topics.toSeq).flatMap{ case (topic, partitions) =>
-        partitions.map(TopicAndPartition(topic, _))
+      val partitionsFromZk = zkClient.getPartitionsForTopics(topics).flatMap { case (topic, partitions) =>
+        partitions.map(new TopicPartition(topic, _))
       }.toSet
 
       val (validPartitions, invalidPartitions) = partitionsFromUser.partition(partitionsFromZk.contains)
-      PreferredReplicaLeaderElectionCommand.writePreferredReplicaElectionData(zkUtils, validPartitions)
+      PreferredReplicaLeaderElectionCommand.writePreferredReplicaElectionData(zkClient, validPartitions)
 
       println("Successfully started preferred replica election for partitions %s".format(validPartitions))
       invalidPartitions.foreach(p => println("Skipping preferred replica leader election for partition %s since it doesn't exist.".format(p)))
diff --git a/core/src/main/scala/kafka/admin/ReassignPartitionsCommand.scala b/core/src/main/scala/kafka/admin/ReassignPartitionsCommand.scala
index e8aad7f167b4e..ad7b545417211 100755
--- a/core/src/main/scala/kafka/admin/ReassignPartitionsCommand.scala
+++ b/core/src/main/scala/kafka/admin/ReassignPartitionsCommand.scala
@@ -16,260 +16,1490 @@
  */
 package kafka.admin
 
-import java.util.Properties
+import java.util
+import java.util.Optional
 import java.util.concurrent.ExecutionException
 
-import scala.collection._
-import scala.collection.JavaConverters._
-import kafka.server.{ConfigType, DynamicConfig}
-import kafka.utils._
-import kafka.common.{AdminCommandFailedException, TopicAndPartition}
+import kafka.common.AdminCommandFailedException
 import kafka.log.LogConfig
-import org.I0Itec.zkclient.exception.ZkNodeExistsException
-import org.apache.kafka.common.utils.Utils
+import kafka.server.{ConfigType, DynamicConfig}
+import kafka.utils.{CommandDefaultOptions, CommandLineUtils, CoreUtils, Exit, Json, Logging}
+import kafka.utils.Implicits._
+import kafka.utils.json.JsonValue
+import kafka.zk.{AdminZkClient, KafkaZkClient}
+import org.apache.kafka.clients.admin.AlterConfigOp.OpType
+import org.apache.kafka.clients.admin.{Admin, AdminClientConfig, AlterConfigOp, ConfigEntry, NewPartitionReassignment, PartitionReassignment, TopicDescription}
+import org.apache.kafka.common.config.ConfigResource
+import org.apache.kafka.common.errors.{ReplicaNotAvailableException, UnknownTopicOrPartitionException}
 import org.apache.kafka.common.security.JaasUtils
-import org.apache.kafka.common.TopicPartitionReplica
-import org.apache.kafka.common.errors.ReplicaNotAvailableException
-import org.apache.kafka.clients.admin.{AdminClientConfig, AlterReplicaLogDirsOptions, AdminClient => JAdminClient}
-import LogConfig._
-import joptsimple.OptionParser
-import org.apache.kafka.clients.admin.DescribeReplicaLogDirsResult.ReplicaLogDirInfo
+import org.apache.kafka.common.utils.{Time, Utils}
+import org.apache.kafka.common.{KafkaException, KafkaFuture, TopicPartition, TopicPartitionReplica}
+
+import scala.jdk.CollectionConverters._
+import scala.collection.{Map, Seq, mutable}
+import scala.compat.java8.OptionConverters._
+import scala.math.Ordered.orderingToOrdered
 
 object ReassignPartitionsCommand extends Logging {
+  private[admin] val AnyLogDir = "any"
 
-  case class Throttle(interBrokerLimit: Long, replicaAlterLogDirsLimit: Long = -1, postUpdateAction: () => Unit = () => ())
+  val helpText = "This tool helps to move topic partitions between replicas."
 
-  private[admin] val NoThrottle = Throttle(-1, -1)
-  private[admin] val AnyLogDir = "any"
+  /**
+   * The earliest version of the partition reassignment JSON.  We will default to this
+   * version if no other version number is given.
+   */
+  private[admin] val EarliestVersion = 1
 
-  def main(args: Array[String]): Unit = {
+  /**
+   * The earliest version of the JSON for each partition reassignment topic.  We will
+   * default to this version if no other version number is given.
+   */
+  private[admin] val EarliestTopicsJsonVersion = 1
+
+  // Throttles that are set at the level of an individual broker.
+  private[admin] val brokerLevelLeaderThrottle =
+    DynamicConfig.Broker.LeaderReplicationThrottledRateProp
+  private[admin] val brokerLevelFollowerThrottle =
+    DynamicConfig.Broker.FollowerReplicationThrottledRateProp
+  private[admin] val brokerLevelLogDirThrottle =
+    DynamicConfig.Broker.ReplicaAlterLogDirsIoMaxBytesPerSecondProp
+  private[admin] val brokerLevelThrottles = Seq(
+    brokerLevelLeaderThrottle,
+    brokerLevelFollowerThrottle,
+    brokerLevelLogDirThrottle
+  )
+
+  // Throttles that are set at the level of an individual topic.
+  private[admin] val topicLevelLeaderThrottle =
+    LogConfig.LeaderReplicationThrottledReplicasProp
+  private[admin] val topicLevelFollowerThrottle =
+    LogConfig.FollowerReplicationThrottledReplicasProp
+  private[admin] val topicLevelThrottles = Seq(
+    topicLevelLeaderThrottle,
+    topicLevelFollowerThrottle
+  )
+
+  private[admin] val cannotExecuteBecauseOfExistingMessage = "Cannot execute because " +
+    "there is an existing partition assignment.  Use --additional to override this and " +
+    "create a new partition assignment in addition to the existing one. The --additional " +
+    "flag can also be used to change the throttle by resubmitting the current reassignment."
+
+  private[admin] val youMustRunVerifyPeriodicallyMessage = "Warning: You must run " +
+    "--verify periodically, until the reassignment completes, to ensure the throttle " +
+    "is removed."
+
+  /**
+   * A map from topic names to partition movements.
+   */
+  type MoveMap = mutable.Map[String, mutable.Map[Int, PartitionMove]]
+
+  /**
+   * A partition movement.  The source and destination brokers may overlap.
+   *
+   * @param sources         The source brokers.
+   * @param destinations    The destination brokers.
+   */
+  sealed case class PartitionMove(sources: mutable.Set[Int],
+                                  destinations: mutable.Set[Int]) { }
+
+  /**
+   * The state of a partition reassignment.  The current replicas and target replicas
+   * may overlap.
+   *
+   * @param currentReplicas The current replicas.
+   * @param targetReplicas  The target replicas.
+   * @param done            True if the reassignment is done.
+   */
+  sealed case class PartitionReassignmentState(currentReplicas: Seq[Int],
+                                               targetReplicas: Seq[Int],
+                                               done: Boolean) {}
+
+  /**
+   * The state of a replica log directory movement.
+   */
+  sealed trait LogDirMoveState {
+    /**
+     * True if the move is done without errors.
+     */
+    def done: Boolean
+  }
+
+  /**
+   * A replica log directory move state where the source log directory is missing.
+   *
+   * @param targetLogDir        The log directory that we wanted the replica to move to.
+   */
+  sealed case class MissingReplicaMoveState(targetLogDir: String)
+      extends LogDirMoveState {
+    override def done = false
+  }
+
+  /**
+   * A replica log directory move state where the source replica is missing.
+   *
+   * @param targetLogDir        The log directory that we wanted the replica to move to.
+   */
+  sealed case class MissingLogDirMoveState(targetLogDir: String)
+      extends LogDirMoveState {
+    override def done = false
+  }
 
+  /**
+   * A replica log directory move state where the move is in progress.
+   *
+   * @param currentLogDir       The current log directory.
+   * @param futureLogDir        The log directory that the replica is moving to.
+   * @param targetLogDir        The log directory that we wanted the replica to move to.
+   */
+  sealed case class ActiveMoveState(currentLogDir: String,
+                                    targetLogDir: String,
+                                    futureLogDir: String)
+      extends LogDirMoveState {
+    override def done = false
+  }
+
+  /**
+   * A replica log directory move state where there is no move in progress, but we did not
+   * reach the target log directory.
+   *
+   * @param currentLogDir       The current log directory.
+   * @param targetLogDir        The log directory that we wanted the replica to move to.
+   */
+  sealed case class CancelledMoveState(currentLogDir: String,
+                                       targetLogDir: String)
+      extends LogDirMoveState {
+    override def done = true
+  }
+
+  /**
+   * The completed replica log directory move state.
+   *
+   * @param targetLogDir        The log directory that we wanted the replica to move to.
+   */
+  sealed case class CompletedMoveState(targetLogDir: String)
+      extends LogDirMoveState {
+    override def done = true
+  }
+
+  /**
+   * An exception thrown to indicate that the command has failed, but we don't want to
+   * print a stack trace.
+   *
+   * @param message     The message to print out before exiting.  A stack trace will not
+   *                    be printed.
+   */
+  class TerseReassignmentFailureException(message: String) extends KafkaException(message) {
+  }
+
+  def main(args: Array[String]): Unit = {
     val opts = validateAndParseArgs(args)
-    val zkConnect = opts.options.valueOf(opts.zkConnectOpt)
-    val zkUtils = ZkUtils(zkConnect,
-                          30000,
-                          30000,
-                          JaasUtils.isZkSecurityEnabled())
-    val adminClientOpt = createAdminClient(opts)
+    var toClose: Option[AutoCloseable] = None
+    var failed = true
 
     try {
-      if(opts.options.has(opts.verifyOpt))
-        verifyAssignment(zkUtils, adminClientOpt, opts)
-      else if(opts.options.has(opts.generateOpt))
-        generateAssignment(zkUtils, opts)
-      else if (opts.options.has(opts.executeOpt))
-        executeAssignment(zkUtils, adminClientOpt, opts)
+      if (opts.options.has(opts.bootstrapServerOpt)) {
+        if (opts.options.has(opts.zkConnectOpt)) {
+          println("Warning: ignoring deprecated --zookeeper option because " +
+            "--bootstrap-server was specified.  The --zookeeper option will " +
+            "be removed in a future version of Kafka.")
+        }
+        val props = if (opts.options.has(opts.commandConfigOpt))
+          Utils.loadProps(opts.options.valueOf(opts.commandConfigOpt))
+        else
+          new util.Properties()
+        props.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, opts.options.valueOf(opts.bootstrapServerOpt))
+        props.putIfAbsent(AdminClientConfig.CLIENT_ID_CONFIG, "reassign-partitions-tool")
+        val adminClient = Admin.create(props)
+        toClose = Some(adminClient)
+        handleAction(adminClient, opts)
+      } else {
+        println("Warning: --zookeeper is deprecated, and will be removed in a future " +
+          "version of Kafka.")
+        val zkClient = KafkaZkClient(opts.options.valueOf(opts.zkConnectOpt),
+          JaasUtils.isZkSaslEnabled, 30000, 30000, Int.MaxValue, Time.SYSTEM)
+        toClose = Some(zkClient)
+        handleAction(zkClient, opts)
+      }
+      failed = false
     } catch {
+      case e: TerseReassignmentFailureException =>
+        println(e.getMessage)
       case e: Throwable =>
-        println("Partitions reassignment failed due to " + e.getMessage)
+        println("Error: " + e.getMessage)
         println(Utils.stackTrace(e))
-    } finally zkUtils.close()
+    } finally {
+      // Close the AdminClient or ZooKeeper client, as appropriate.
+      // It's good to do this after printing any error stack trace.
+      toClose.foreach(_.close())
+    }
+    // If the command failed, exit with a non-zero exit code.
+    if (failed) {
+      Exit.exit(1)
+    }
   }
 
-  private def createAdminClient(opts: ReassignPartitionsCommandOptions): Option[JAdminClient] = {
-    if (opts.options.has(opts.bootstrapServerOpt)) {
-      val props = new Properties()
-      props.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, opts.options.valueOf(opts.bootstrapServerOpt))
-      props.put(AdminClientConfig.CLIENT_ID_CONFIG, "reassign-partitions-tool")
-      Some(JAdminClient.create(props))
+  private def handleAction(adminClient: Admin,
+                           opts: ReassignPartitionsCommandOptions): Unit = {
+    if (opts.options.has(opts.verifyOpt)) {
+      verifyAssignment(adminClient,
+        Utils.readFileAsString(opts.options.valueOf(opts.reassignmentJsonFileOpt)),
+        opts.options.has(opts.preserveThrottlesOpt))
+    } else if (opts.options.has(opts.generateOpt)) {
+      generateAssignment(adminClient,
+        Utils.readFileAsString(opts.options.valueOf(opts.topicsToMoveJsonFileOpt)),
+        opts.options.valueOf(opts.brokerListOpt),
+        !opts.options.has(opts.disableRackAware))
+    } else if (opts.options.has(opts.executeOpt)) {
+      executeAssignment(adminClient,
+        opts.options.has(opts.additionalOpt),
+        Utils.readFileAsString(opts.options.valueOf(opts.reassignmentJsonFileOpt)),
+        opts.options.valueOf(opts.interBrokerThrottleOpt),
+        opts.options.valueOf(opts.replicaAlterLogDirsThrottleOpt),
+        opts.options.valueOf(opts.timeoutOpt))
+    } else if (opts.options.has(opts.cancelOpt)) {
+      cancelAssignment(adminClient,
+        Utils.readFileAsString(opts.options.valueOf(opts.reassignmentJsonFileOpt)),
+        opts.options.has(opts.preserveThrottlesOpt),
+        opts.options.valueOf(opts.timeoutOpt))
+    } else if (opts.options.has(opts.listOpt)) {
+      listReassignments(adminClient)
     } else {
-      None
+      throw new RuntimeException("Unsupported action.")
     }
   }
 
-  def verifyAssignment(zkUtils: ZkUtils, adminClientOpt: Option[JAdminClient], opts: ReassignPartitionsCommandOptions) {
-    val jsonFile = opts.options.valueOf(opts.reassignmentJsonFileOpt)
-    val jsonString = Utils.readFileAsString(jsonFile)
-    verifyAssignment(zkUtils, adminClientOpt, jsonString)
+  private def handleAction(zkClient: KafkaZkClient,
+                           opts: ReassignPartitionsCommandOptions): Unit = {
+    if (opts.options.has(opts.verifyOpt)) {
+      verifyAssignment(zkClient,
+        Utils.readFileAsString(opts.options.valueOf(opts.reassignmentJsonFileOpt)),
+        opts.options.has(opts.preserveThrottlesOpt))
+    } else if (opts.options.has(opts.generateOpt)) {
+      generateAssignment(zkClient,
+        Utils.readFileAsString(opts.options.valueOf(opts.topicsToMoveJsonFileOpt)),
+        opts.options.valueOf(opts.brokerListOpt),
+        !opts.options.has(opts.disableRackAware))
+    } else if (opts.options.has(opts.executeOpt)) {
+      executeAssignment(zkClient,
+        Utils.readFileAsString(opts.options.valueOf(opts.reassignmentJsonFileOpt)),
+        opts.options.valueOf(opts.interBrokerThrottleOpt))
+    } else {
+      throw new RuntimeException("Unsupported action.")
+    }
   }
 
-  def verifyAssignment(zkUtils: ZkUtils, adminClientOpt: Option[JAdminClient], jsonString: String): Unit = {
-    println("Status of partition reassignment: ")
-    val (partitionsToBeReassigned, replicaAssignment) = parsePartitionReassignmentData(jsonString)
-    val reassignedPartitionsStatus = checkIfPartitionReassignmentSucceeded(zkUtils, partitionsToBeReassigned.toMap)
-    val replicasReassignmentStatus = checkIfReplicaReassignmentSucceeded(adminClientOpt, replicaAssignment)
+  /**
+   * A result returned from verifyAssignment.
+   *
+   * @param partStates    A map from partitions to reassignment states.
+   * @param partsOngoing  True if there are any ongoing partition reassignments.
+   * @param moveStates    A map from log directories to movement states.
+   * @param movesOngoing  True if there are any ongoing moves that we know about.
+   */
+  case class VerifyAssignmentResult(partStates: Map[TopicPartition, PartitionReassignmentState],
+                                    partsOngoing: Boolean = false,
+                                    moveStates: Map[TopicPartitionReplica, LogDirMoveState] = Map.empty,
+                                    movesOngoing: Boolean = false)
 
-    reassignedPartitionsStatus.foreach { case (topicPartition, status) =>
-      status match {
-        case ReassignmentCompleted =>
-          println("Reassignment of partition %s completed successfully".format(topicPartition))
-        case ReassignmentFailed =>
-          println("Reassignment of partition %s failed".format(topicPartition))
-        case ReassignmentInProgress =>
-          println("Reassignment of partition %s is still in progress".format(topicPartition))
+  /**
+   * The entry point for the --verify command.
+   *
+   * @param adminClient           The AdminClient to use.
+   * @param jsonString            The JSON string to use for the topics and partitions to verify.
+   * @param preserveThrottles     True if we should avoid changing topic or broker throttles.
+   *
+   * @return                      A result that is useful for testing.
+   */
+  def verifyAssignment(adminClient: Admin, jsonString: String, preserveThrottles: Boolean)
+                      : VerifyAssignmentResult = {
+    val (targetParts, targetLogDirs) = parsePartitionReassignmentData(jsonString)
+    val (partStates, partsOngoing) = verifyPartitionAssignments(adminClient, targetParts)
+    val (moveStates, movesOngoing) = verifyReplicaMoves(adminClient, targetLogDirs)
+    if (!partsOngoing && !movesOngoing && !preserveThrottles) {
+      // If the partition assignments and replica assignments are done, clear any throttles
+      // that were set.  We have to clear all throttles, because we don't have enough
+      // information to know all of the source brokers that might have been involved in the
+      // previous reassignments.
+      clearAllThrottles(adminClient, targetParts)
+    }
+    VerifyAssignmentResult(partStates, partsOngoing, moveStates, movesOngoing)
+  }
+
+  /**
+   * Verify the partition reassignments specified by the user.
+   *
+   * @param adminClient           The AdminClient to use.
+   * @param targets               The partition reassignments specified by the user.
+   *
+   * @return                      A tuple of the partition reassignment states, and a
+   *                              boolean which is true if there are no ongoing
+   *                              reassignments (including reassignments not described
+   *                              in the JSON file.)
+   */
+  def verifyPartitionAssignments(adminClient: Admin,
+                                 targets: Seq[(TopicPartition, Seq[Int])])
+                                 : (Map[TopicPartition, PartitionReassignmentState], Boolean) = {
+    val (partStates, partsOngoing) = findPartitionReassignmentStates(adminClient, targets)
+    println(partitionReassignmentStatesToString(partStates))
+    (partStates, partsOngoing)
+  }
+
+  /**
+   * The deprecated entry point for the --verify command.
+   *
+   * @param zkClient              The ZooKeeper client to use.
+   * @param jsonString            The JSON string to use for the topics and partitions to verify.
+   * @param preserveThrottles     True if we should avoid changing topic or broker throttles.
+   *
+   * @return                      A result that is useful for testing.  Note that anything that
+   *                              would require AdminClient to see will be left out of this result.
+   */
+  def verifyAssignment(zkClient: KafkaZkClient, jsonString: String, preserveThrottles: Boolean)
+                       : VerifyAssignmentResult = {
+    val (targetParts, targetLogDirs) = parsePartitionReassignmentData(jsonString)
+    if (targetLogDirs.nonEmpty) {
+      throw new AdminCommandFailedException("bootstrap-server needs to be provided when " +
+        "replica reassignments are present.")
+    }
+    println("Warning: because you are using the deprecated --zookeeper option, the results " +
+      "may be incomplete.  Use --bootstrap-server instead for more accurate results.")
+    val (partStates, partsOngoing) = verifyPartitionAssignments(zkClient, targetParts.toMap)
+    if (!partsOngoing && !preserveThrottles) {
+      clearAllThrottles(zkClient, targetParts)
+    }
+    VerifyAssignmentResult(partStates, partsOngoing, Map.empty, false)
+  }
+
+  /**
+   * Verify the partition reassignments specified by the user.
+   *
+   * @param zkClient              The ZooKeeper client to use.
+   * @param targets               The partition reassignments specified by the user.
+   *
+   * @return                      A tuple of partition states and whether there are any
+   *                              ongoing reassignments found in the legacy reassign
+   *                              partitions ZNode.
+   */
+  def verifyPartitionAssignments(zkClient: KafkaZkClient,
+                                 targets: Map[TopicPartition, Seq[Int]])
+                                 : (Map[TopicPartition, PartitionReassignmentState], Boolean) = {
+    val (partStates, partsOngoing) = findPartitionReassignmentStates(zkClient, targets)
+    println(partitionReassignmentStatesToString(partStates))
+    (partStates, partsOngoing)
+  }
+
+  def compareTopicPartitions(a: TopicPartition, b: TopicPartition): Boolean = {
+    (a.topic(), a.partition()) < (b.topic(), b.partition())
+  }
+
+  def compareTopicPartitionReplicas(a: TopicPartitionReplica, b: TopicPartitionReplica): Boolean = {
+    (a.brokerId(), a.topic(), a.partition()) < (b.brokerId(), b.topic(), b.partition())
+  }
+
+  /**
+   * Convert partition reassignment states to a human-readable string.
+   *
+   * @param states      A map from topic partitions to states.
+   * @return            A string summarizing the partition reassignment states.
+   */
+  def partitionReassignmentStatesToString(states: Map[TopicPartition, PartitionReassignmentState])
+                                          : String = {
+    val bld = new mutable.ArrayBuffer[String]()
+    bld.append("Status of partition reassignment:")
+    states.keySet.toBuffer.sortWith(compareTopicPartitions).foreach {
+      topicPartition => {
+        val state = states(topicPartition)
+        if (state.done) {
+          if (state.currentReplicas.equals(state.targetReplicas)) {
+            bld.append("Reassignment of partition %s is complete.".
+              format(topicPartition.toString))
+          } else {
+            bld.append(s"There is no active reassignment of partition ${topicPartition}, " +
+              s"but replica set is ${state.currentReplicas.mkString(",")} rather than " +
+              s"${state.targetReplicas.mkString(",")}.")
+          }
+        } else {
+          bld.append("Reassignment of partition %s is still in progress.".format(topicPartition))
+        }
       }
     }
+    bld.mkString(System.lineSeparator())
+  }
+
+  /**
+   * Find the state of the specified partition reassignments.
+   *
+   * @param adminClient          The Admin client to use.
+   * @param targetReassignments  The reassignments we want to learn about.
+   *
+   * @return                     A tuple containing the reassignment states for each topic
+   *                             partition, plus whether there are any ongoing reassignments.
+   */
+  def findPartitionReassignmentStates(adminClient: Admin,
+                                      targetReassignments: Seq[(TopicPartition, Seq[Int])])
+                                      : (Map[TopicPartition, PartitionReassignmentState], Boolean) = {
+    val currentReassignments = adminClient.
+      listPartitionReassignments.reassignments.get().asScala
+    val (foundReassignments, notFoundReassignments) = targetReassignments.partition {
+      case (part, _) => currentReassignments.contains(part)
+    }
+    val foundResults: Seq[(TopicPartition, PartitionReassignmentState)] = foundReassignments.map {
+      case (part, targetReplicas) => (part,
+        PartitionReassignmentState(
+          currentReassignments(part).replicas.
+            asScala.map(i => i.asInstanceOf[Int]),
+          targetReplicas,
+          false))
+    }
+    val topicNamesToLookUp = new mutable.HashSet[String]()
+    notFoundReassignments.foreach {
+      case (part, _) =>
+        if (!currentReassignments.contains(part))
+          topicNamesToLookUp.add(part.topic)
+    }
+    val topicDescriptions = adminClient.
+      describeTopics(topicNamesToLookUp.asJava).values().asScala
+    val notFoundResults: Seq[(TopicPartition, PartitionReassignmentState)] = notFoundReassignments.map {
+      case (part, targetReplicas) =>
+        currentReassignments.get(part) match {
+          case Some(reassignment) => (part,
+            PartitionReassignmentState(
+              reassignment.replicas.asScala.map(_.asInstanceOf[Int]),
+              targetReplicas,
+              false))
+          case None =>
+            (part, topicDescriptionFutureToState(part.partition,
+              topicDescriptions(part.topic), targetReplicas))
+        }
+    }
+    val allResults = foundResults ++ notFoundResults
+    (allResults.toMap, currentReassignments.nonEmpty)
+  }
+
+  private def topicDescriptionFutureToState(partition: Int,
+                                            future: KafkaFuture[TopicDescription],
+                                            targetReplicas: Seq[Int]): PartitionReassignmentState = {
+    try {
+      val topicDescription = future.get()
+      if (topicDescription.partitions().size() < partition) {
+        throw new ExecutionException("Too few partitions found", new UnknownTopicOrPartitionException())
+      }
+      PartitionReassignmentState(
+        topicDescription.partitions.get(partition).replicas.asScala.map(_.id),
+        targetReplicas,
+        true)
+    } catch {
+      case t: ExecutionException =>
+        t.getCause match {
+          case _: UnknownTopicOrPartitionException =>
+            PartitionReassignmentState(Seq(), targetReplicas, true)
+          case e => throw e
+        }
+    }
+  }
+
+  /**
+   * Find the state of the specified partition reassignments.
+   *
+   * @param zkClient              The ZooKeeper client to use.
+   * @param targetReassignments   The reassignments we want to learn about.
+   *
+   * @return                      A tuple containing the reassignment states for each topic
+   *                              partition, plus whether there are any ongoing reassignments
+   *                              found in the legacy reassign partitions znode.
+   */
+  def findPartitionReassignmentStates(zkClient: KafkaZkClient,
+                                      targetReassignments: Map[TopicPartition, Seq[Int]])
+                                      : (Map[TopicPartition, PartitionReassignmentState], Boolean) = {
+    val partitionsBeingReassigned = zkClient.getPartitionReassignment
+    val results = new mutable.HashMap[TopicPartition, PartitionReassignmentState]()
+    targetReassignments.groupBy(_._1.topic).foreach {
+      case (topic, partitions) =>
+        val replicasForTopic = zkClient.getReplicaAssignmentForTopics(Set(topic))
+        partitions.foreach {
+          case (partition, targetReplicas) =>
+            val currentReplicas = replicasForTopic.getOrElse(partition, Seq())
+            results.put(partition, new PartitionReassignmentState(
+              currentReplicas, targetReplicas, !partitionsBeingReassigned.contains(partition)))
+        }
+    }
+    (results, partitionsBeingReassigned.nonEmpty)
+  }
+
+  /**
+   * Verify the replica reassignments specified by the user.
+   *
+   * @param adminClient           The AdminClient to use.
+   * @param targetReassignments   The replica reassignments specified by the user.
+   *
+   * @return                      A tuple of the replica states, and a boolean which is true
+   *                              if there are any ongoing replica moves.
+   *
+   *                              Note: Unlike in verifyPartitionAssignments, we will
+   *                              return false here even if there are unrelated ongoing
+   *                              reassignments. (We don't have an efficient API that
+   *                              returns all ongoing replica reassignments.)
+   */
+  def verifyReplicaMoves(adminClient: Admin,
+                         targetReassignments: Map[TopicPartitionReplica, String])
+                         : (Map[TopicPartitionReplica, LogDirMoveState], Boolean) = {
+    val moveStates = findLogDirMoveStates(adminClient, targetReassignments)
+    println(replicaMoveStatesToString(moveStates))
+    (moveStates, !moveStates.values.forall(_.done))
+  }
 
-    replicasReassignmentStatus.foreach { case (replica, status) =>
-      status match {
-        case ReassignmentCompleted =>
-          println("Reassignment of replica %s completed successfully".format(replica))
-        case ReassignmentFailed =>
-          println("Reassignment of replica %s failed".format(replica))
-        case ReassignmentInProgress =>
-          println("Reassignment of replica %s is still in progress".format(replica))
+  /**
+   * Find the state of the specified partition reassignments.
+   *
+   * @param adminClient           The AdminClient to use.
+   * @param targetMoves           The movements we want to learn about.  The map is keyed
+   *                              by TopicPartitionReplica, and its values are target log
+   *                              directories.
+   *
+   * @return                      The states for each replica movement.
+   */
+  def findLogDirMoveStates(adminClient: Admin,
+                           targetMoves: Map[TopicPartitionReplica, String])
+                           : Map[TopicPartitionReplica, LogDirMoveState] = {
+    val replicaLogDirInfos = adminClient.describeReplicaLogDirs(
+      targetMoves.keySet.asJava).all().get().asScala
+    targetMoves.map { case (replica, targetLogDir) =>
+      val moveState: LogDirMoveState = replicaLogDirInfos.get(replica) match {
+        case None => MissingReplicaMoveState(targetLogDir)
+        case Some(info) => if (info.getCurrentReplicaLogDir == null) {
+            MissingLogDirMoveState(targetLogDir)
+          } else if (info.getFutureReplicaLogDir == null) {
+            if (info.getCurrentReplicaLogDir.equals(targetLogDir)) {
+              CompletedMoveState(targetLogDir)
+            } else {
+              CancelledMoveState(info.getCurrentReplicaLogDir, targetLogDir)
+            }
+          } else {
+            ActiveMoveState(info.getCurrentReplicaLogDir(),
+              targetLogDir,
+              info.getFutureReplicaLogDir)
+          }
       }
+      (replica, moveState)
     }
+  }
 
-    removeThrottle(zkUtils, reassignedPartitionsStatus, replicasReassignmentStatus)
-  }
-
-  private[admin] def removeThrottle(zkUtils: ZkUtils,
-                                    reassignedPartitionsStatus: Map[TopicAndPartition, ReassignmentStatus],
-                                    replicasReassignmentStatus: Map[TopicPartitionReplica, ReassignmentStatus],
-                                    admin: AdminUtilities = AdminUtils): Unit = {
-
-    //If both partition assignment and replica reassignment have completed remove both the inter-broker and replica-alter-dir throttle
-    if (reassignedPartitionsStatus.forall { case (_, status) => status == ReassignmentCompleted } &&
-        replicasReassignmentStatus.forall { case (_, status) => status == ReassignmentCompleted }) {
-      var changed = false
-      //Remove the throttle limit from all brokers in the cluster
-      //(as we no longer know which specific brokers were involved in the move)
-      for (brokerId <- zkUtils.getAllBrokersInCluster().map(_.id)) {
-        val configs = admin.fetchEntityConfig(zkUtils, ConfigType.Broker, brokerId.toString)
-        // bitwise OR as we don't want to short-circuit
-        if (configs.remove(DynamicConfig.Broker.LeaderReplicationThrottledRateProp) != null
-          | configs.remove(DynamicConfig.Broker.FollowerReplicationThrottledRateProp) != null
-          | configs.remove(DynamicConfig.Broker.ReplicaAlterLogDirsIoMaxBytesPerSecondProp) != null){
-          admin.changeBrokerConfig(zkUtils, Seq(brokerId), configs)
-          changed = true
+  /**
+   * Convert replica move states to a human-readable string.
+   *
+   * @param states          A map from topic partition replicas to states.
+   * @return                A tuple of a summary string, and a boolean describing
+   *                        whether there are any active replica moves.
+   */
+  def replicaMoveStatesToString(states: Map[TopicPartitionReplica, LogDirMoveState])
+                                : String = {
+    val bld = new mutable.ArrayBuffer[String]
+    states.keySet.toBuffer.sortWith(compareTopicPartitionReplicas).foreach {
+      case replica =>
+        val state = states(replica)
+        state match {
+          case MissingLogDirMoveState(targetLogDir) =>
+            bld.append(s"Partition ${replica.topic()}-${replica.partition()} is not found " +
+              s"in any live log dir on broker ${replica.brokerId()}. There is likely an " +
+              s"offline log directory on the broker.")
+          case MissingReplicaMoveState(targetLogDir) =>
+            bld.append(s"Partition ${replica.topic()}-${replica.partition()} cannot be found " +
+              s"in any live log directory on broker ${replica.brokerId()}.")
+          case ActiveMoveState(currentLogDir, targetLogDir, futureLogDir) =>
+            if (targetLogDir.equals(futureLogDir)) {
+              bld.append(s"Reassignment of replica ${replica} is still in progress.")
+            } else {
+              bld.append(s"Partition ${replica.topic()}-${replica.partition()} on broker " +
+                s"${replica.brokerId()} is being moved to log dir ${futureLogDir} " +
+                s"instead of ${targetLogDir}.")
+            }
+          case CancelledMoveState(currentLogDir, targetLogDir) =>
+            bld.append(s"Partition ${replica.topic()}-${replica.partition()} on broker " +
+              s"${replica.brokerId()} is not being moved from log dir ${currentLogDir} to " +
+              s"${targetLogDir}.")
+          case CompletedMoveState(targetLogDir) =>
+            bld.append(s"Reassignment of replica ${replica} completed successfully.")
         }
+    }
+    bld.mkString(System.lineSeparator())
+  }
+
+  /**
+   * Clear all topic-level and broker-level throttles.
+   *
+   * @param adminClient     The AdminClient to use.
+   * @param targetParts     The target partitions loaded from the JSON file.
+   */
+  def clearAllThrottles(adminClient: Admin,
+                        targetParts: Seq[(TopicPartition, Seq[Int])]): Unit = {
+    val activeBrokers = adminClient.describeCluster().nodes().get().asScala.map(_.id()).toSet
+    val brokers = activeBrokers ++ targetParts.flatMap(_._2).toSet
+    println("Clearing broker-level throttles on broker%s %s".format(
+      if (brokers.size == 1) "" else "s", brokers.mkString(",")))
+    clearBrokerLevelThrottles(adminClient, brokers)
+
+    val topics = targetParts.map(_._1.topic()).toSet
+    println("Clearing topic-level throttles on topic%s %s".format(
+      if (topics.size == 1) "" else "s", topics.mkString(",")))
+    clearTopicLevelThrottles(adminClient, topics)
+  }
+
+  /**
+   * Clear all topic-level and broker-level throttles.
+   *
+   * @param zkClient        The ZooKeeper client to use.
+   * @param targetParts     The target partitions loaded from the JSON file.
+   */
+  def clearAllThrottles(zkClient: KafkaZkClient,
+                        targetParts: Seq[(TopicPartition, Seq[Int])]): Unit = {
+    val activeBrokers = zkClient.getAllBrokersInCluster.map(_.id).toSet
+    val brokers = activeBrokers ++ targetParts.flatMap(_._2).toSet
+    println("Clearing broker-level throttles on broker%s %s".format(
+      if (brokers.size == 1) "" else "s", brokers.mkString(",")))
+    clearBrokerLevelThrottles(zkClient, brokers)
+
+    val topics = targetParts.map(_._1.topic()).toSet
+    println("Clearing topic-level throttles on topic%s %s".format(
+      if (topics.size == 1) "" else "s", topics.mkString(",")))
+    clearTopicLevelThrottles(zkClient, topics)
+  }
+
+  /**
+   * Clear all throttles which have been set at the broker level.
+   *
+   * @param adminClient       The AdminClient to use.
+   * @param brokers           The brokers to clear the throttles for.
+   */
+  def clearBrokerLevelThrottles(adminClient: Admin, brokers: Set[Int]): Unit = {
+    val configOps = new util.HashMap[ConfigResource, util.Collection[AlterConfigOp]]()
+    brokers.foreach {
+      case brokerId => configOps.put(
+        new ConfigResource(ConfigResource.Type.BROKER, brokerId.toString),
+        brokerLevelThrottles.map(throttle => new AlterConfigOp(
+          new ConfigEntry(throttle, null), OpType.DELETE)).asJava)
+    }
+    adminClient.incrementalAlterConfigs(configOps).all().get()
+  }
+
+  /**
+   * Clear all throttles which have been set at the broker level.
+   *
+   * @param zkClient          The ZooKeeper client to use.
+   * @param brokers           The brokers to clear the throttles for.
+   */
+  def clearBrokerLevelThrottles(zkClient: KafkaZkClient, brokers: Set[Int]): Unit = {
+    val adminZkClient = new AdminZkClient(zkClient)
+    for (brokerId <- brokers) {
+      val configs = adminZkClient.fetchEntityConfig(ConfigType.Broker, brokerId.toString)
+      if (brokerLevelThrottles.flatMap(throttle => Option(configs.remove(throttle))).nonEmpty) {
+        adminZkClient.changeBrokerConfig(Seq(brokerId), configs)
+      }
+    }
+  }
+
+  /**
+   * Clear the reassignment throttles for the specified topics.
+   *
+   * @param adminClient           The AdminClient to use.
+   * @param topics                The topics to clear the throttles for.
+   */
+  def clearTopicLevelThrottles(adminClient: Admin, topics: Set[String]): Unit = {
+    val configOps = new util.HashMap[ConfigResource, util.Collection[AlterConfigOp]]()
+    topics.foreach {
+      topicName => configOps.put(
+        new ConfigResource(ConfigResource.Type.TOPIC, topicName),
+        topicLevelThrottles.map(throttle => new AlterConfigOp(new ConfigEntry(throttle, null),
+          OpType.DELETE)).asJava)
+    }
+    adminClient.incrementalAlterConfigs(configOps).all().get()
+  }
+
+  /**
+   * Clear the reassignment throttles for the specified topics.
+   *
+   * @param zkClient              The ZooKeeper client to use.
+   * @param topics                The topics to clear the throttles for.
+   */
+  def clearTopicLevelThrottles(zkClient: KafkaZkClient, topics: Set[String]): Unit = {
+    val adminZkClient = new AdminZkClient(zkClient)
+    for (topic <- topics) {
+      val configs = adminZkClient.fetchEntityConfig(ConfigType.Topic, topic)
+      if (topicLevelThrottles.flatMap(throttle => Option(configs.remove(throttle))).nonEmpty) {
+        adminZkClient.changeTopicConfig(topic, configs)
       }
+    }
+  }
+
+  /**
+   * The entry point for the --generate command.
+   *
+   * @param adminClient           The AdminClient to use.
+   * @param reassignmentJson      The JSON string to use for the topics to reassign.
+   * @param brokerListString      The comma-separated string of broker IDs to use.
+   * @param enableRackAwareness   True if rack-awareness should be enabled.
+   *
+   * @return                      A tuple containing the proposed assignment and the
+   *                              current assignment.
+   */
+  def generateAssignment(adminClient: Admin,
+                         reassignmentJson: String,
+                         brokerListString: String,
+                         enableRackAwareness: Boolean)
+                         : (Map[TopicPartition, Seq[Int]], Map[TopicPartition, Seq[Int]]) = {
+    val (brokersToReassign, topicsToReassign) =
+      parseGenerateAssignmentArgs(reassignmentJson, brokerListString)
+    val currentAssignments = getReplicaAssignmentForTopics(adminClient, topicsToReassign)
+    val brokerMetadatas = getBrokerMetadata(adminClient, brokersToReassign, enableRackAwareness)
+    val proposedAssignments = calculateAssignment(currentAssignments, brokerMetadatas)
+    println("Current partition replica assignment\n%s\n".
+      format(formatAsReassignmentJson(currentAssignments, Map.empty)))
+    println("Proposed partition reassignment configuration\n%s".
+      format(formatAsReassignmentJson(proposedAssignments, Map.empty)))
+    (proposedAssignments, currentAssignments)
+  }
 
-      //Remove the list of throttled replicas from all topics with partitions being moved
-      val topics = (reassignedPartitionsStatus.keySet.map(tp => tp.topic) ++ replicasReassignmentStatus.keySet.map(replica => replica.topic)).toSeq.distinct
-      for (topic <- topics) {
-        val configs = admin.fetchEntityConfig(zkUtils, ConfigType.Topic, topic)
-        // bitwise OR as we don't want to short-circuit
-        if (configs.remove(LogConfig.LeaderReplicationThrottledReplicasProp) != null
-          | configs.remove(LogConfig.FollowerReplicationThrottledReplicasProp) != null) {
-          admin.changeTopicConfig(zkUtils, topic, configs)
-          changed = true
+  /**
+   * The legacy entry point for the --generate command.
+   *
+   * @param zkClient              The ZooKeeper client to use.
+   * @param reassignmentJson      The JSON string to use for the topics to reassign.
+   * @param brokerListString      The comma-separated string of broker IDs to use.
+   * @param enableRackAwareness   True if rack-awareness should be enabled.
+   *
+   * @return                      A tuple containing the proposed assignment and the
+   *                              current assignment.
+   */
+  def generateAssignment(zkClient: KafkaZkClient,
+                         reassignmentJson: String,
+                         brokerListString: String,
+                         enableRackAwareness: Boolean)
+                         : (Map[TopicPartition, Seq[Int]], Map[TopicPartition, Seq[Int]]) = {
+    val (brokersToReassign, topicsToReassign) =
+      parseGenerateAssignmentArgs(reassignmentJson, brokerListString)
+    val currentAssignments = zkClient.getReplicaAssignmentForTopics(topicsToReassign.toSet)
+    val brokerMetadatas = getBrokerMetadata(zkClient, brokersToReassign, enableRackAwareness)
+    val proposedAssignments = calculateAssignment(currentAssignments, brokerMetadatas)
+    println("Current partition replica assignment\n%s\n".
+      format(formatAsReassignmentJson(currentAssignments, Map.empty)))
+    println("Proposed partition reassignment configuration\n%s".
+      format(formatAsReassignmentJson(proposedAssignments, Map.empty)))
+    (proposedAssignments, currentAssignments)
+  }
+
+  /**
+   * Calculate the new partition assignments to suggest in --generate.
+   *
+   * @param currentAssignment  The current partition assignments.
+   * @param brokerMetadatas    The rack information for each broker.
+   *
+   * @return                   A map from partitions to the proposed assignments for each.
+   */
+  def calculateAssignment(currentAssignment: Map[TopicPartition, Seq[Int]],
+                          brokerMetadatas: Seq[BrokerMetadata])
+                          : Map[TopicPartition, Seq[Int]] = {
+    val groupedByTopic = currentAssignment.groupBy { case (tp, _) => tp.topic }
+    val proposedAssignments = mutable.Map[TopicPartition, Seq[Int]]()
+    groupedByTopic.forKeyValue { (topic, assignment) =>
+      val (_, replicas) = assignment.head
+      val assignedReplicas = AdminUtils.
+        assignReplicasToBrokers(brokerMetadatas, assignment.size, replicas.size)
+      proposedAssignments ++= assignedReplicas.map { case (partition, replicas) =>
+        new TopicPartition(topic, partition) -> replicas
+      }
+    }
+    proposedAssignments
+  }
+
+  /**
+   * Get the current replica assignments for some topics.
+   *
+   * @param adminClient     The AdminClient to use.
+   * @param topics          The topics to get information about.
+   * @return                A map from partitions to broker assignments.
+   *                        If any topic can't be found, an exception will be thrown.
+   */
+  def getReplicaAssignmentForTopics(adminClient: Admin,
+                                    topics: Seq[String])
+                                    : Map[TopicPartition, Seq[Int]] = {
+    adminClient.describeTopics(topics.asJava).all().get().asScala.flatMap {
+      case (topicName, topicDescription) =>
+        topicDescription.partitions.asScala.map {
+          info => (new TopicPartition(topicName, info.partition),
+            info.replicas.asScala.map(_.id))
         }
+    }
+  }
+
+  /**
+   * Get the current replica assignments for some partitions.
+   *
+   * @param adminClient     The AdminClient to use.
+   * @param partitions      The partitions to get information about.
+   * @return                A map from partitions to broker assignments.
+   *                        If any topic can't be found, an exception will be thrown.
+   */
+  def getReplicaAssignmentForPartitions(adminClient: Admin,
+                                        partitions: Set[TopicPartition])
+                                        : Map[TopicPartition, Seq[Int]] = {
+    adminClient.describeTopics(partitions.map(_.topic).asJava).all().get().asScala.flatMap {
+      case (topicName, topicDescription) =>
+        topicDescription.partitions.asScala.flatMap {
+          info => if (partitions.contains(new TopicPartition(topicName, info.partition))) {
+            Some(new TopicPartition(topicName, info.partition()),
+                info.replicas.asScala.map(_.id))
+          } else {
+            None
+          }
       }
-      if (changed)
-        println("Throttle was removed.")
     }
   }
 
-  def generateAssignment(zkUtils: ZkUtils, opts: ReassignPartitionsCommandOptions) {
-    val topicsToMoveJsonFile = opts.options.valueOf(opts.topicsToMoveJsonFileOpt)
-    val brokerListToReassign = opts.options.valueOf(opts.brokerListOpt).split(',').map(_.toInt)
-    val duplicateReassignments = CoreUtils.duplicates(brokerListToReassign)
-    if (duplicateReassignments.nonEmpty)
-      throw new AdminCommandFailedException("Broker list contains duplicate entries: %s".format(duplicateReassignments.mkString(",")))
-    val topicsToMoveJsonString = Utils.readFileAsString(topicsToMoveJsonFile)
-    val disableRackAware = opts.options.has(opts.disableRackAware)
-    val (proposedAssignments, currentAssignments) = generateAssignment(zkUtils, brokerListToReassign, topicsToMoveJsonString, disableRackAware)
-    println("Current partition replica assignment\n%s\n".format(formatAsReassignmentJson(currentAssignments, Map.empty)))
-    println("Proposed partition reassignment configuration\n%s".format(formatAsReassignmentJson(proposedAssignments, Map.empty)))
+  /**
+   * Find the rack information for some brokers.
+   *
+   * @param adminClient         The AdminClient object.
+   * @param brokers             The brokers to gather metadata about.
+   * @param enableRackAwareness True if we should return rack information, and throw an
+   *                            exception if it is inconsistent.
+   *
+   * @return                    The metadata for each broker that was found.
+   *                            Brokers that were not found will be omitted.
+   */
+  def getBrokerMetadata(adminClient: Admin,
+                        brokers: Seq[Int],
+                        enableRackAwareness: Boolean): Seq[BrokerMetadata] = {
+    val brokerSet = brokers.toSet
+    val results = adminClient.describeCluster().nodes.get().asScala.
+      filter(node => brokerSet.contains(node.id)).
+      map {
+        node => if (enableRackAwareness && node.rack != null) {
+          BrokerMetadata(node.id, Some(node.rack))
+        } else {
+          BrokerMetadata(node.id, None)
+        }
+      }.toSeq
+    val numRackless = results.count(_.rack.isEmpty)
+    if (enableRackAwareness && numRackless != 0 && numRackless != results.size) {
+      throw new AdminOperationException("Not all brokers have rack information. Add " +
+        "--disable-rack-aware in command line to make replica assignment without rack " +
+        "information.")
+    }
+    results
   }
 
-  def generateAssignment(zkUtils: ZkUtils, brokerListToReassign: Seq[Int], topicsToMoveJsonString: String, disableRackAware: Boolean): (Map[TopicAndPartition, Seq[Int]], Map[TopicAndPartition, Seq[Int]]) = {
-    val topicsToReassign = ZkUtils.parseTopicsData(topicsToMoveJsonString)
+  /**
+   * Find the metadata for some brokers.
+   *
+   * @param zkClient              The ZooKeeper client to use.
+   * @param brokers               The brokers to gather metadata about.
+   * @param enableRackAwareness   True if we should return rack information, and throw an
+   *                              exception if it is inconsistent.
+   *
+   * @return                      The metadata for each broker that was found.
+   *                              Brokers that were not found will be omitted.
+   */
+  def getBrokerMetadata(zkClient: KafkaZkClient,
+                        brokers: Seq[Int],
+                        enableRackAwareness: Boolean): Seq[BrokerMetadata] = {
+    val adminZkClient = new AdminZkClient(zkClient)
+    adminZkClient.getBrokerMetadatas(if (enableRackAwareness)
+      RackAwareMode.Enforced else RackAwareMode.Disabled, Some(brokers))
+  }
+
+  /**
+   * Parse and validate data gathered from the command-line for --generate
+   * In particular, we parse the JSON and validate that duplicate brokers and
+   * topics don't appear.
+   *
+   * @param reassignmentJson       The JSON passed to --generate .
+   * @param brokerList             A list of brokers passed to --generate.
+   *
+   * @return                       A tuple of brokers to reassign, topics to reassign
+   */
+  def parseGenerateAssignmentArgs(reassignmentJson: String,
+                                  brokerList: String): (Seq[Int], Seq[String]) = {
+    val brokerListToReassign = brokerList.split(',').map(_.toInt)
+    val duplicateReassignments = CoreUtils.duplicates(brokerListToReassign)
+    if (duplicateReassignments.nonEmpty)
+      throw new AdminCommandFailedException("Broker list contains duplicate entries: %s".
+        format(duplicateReassignments.mkString(",")))
+    val topicsToReassign = parseTopicsData(reassignmentJson)
     val duplicateTopicsToReassign = CoreUtils.duplicates(topicsToReassign)
     if (duplicateTopicsToReassign.nonEmpty)
-      throw new AdminCommandFailedException("List of topics to reassign contains duplicate entries: %s".format(duplicateTopicsToReassign.mkString(",")))
-    val currentAssignment = zkUtils.getReplicaAssignmentForTopics(topicsToReassign)
+      throw new AdminCommandFailedException("List of topics to reassign contains duplicate entries: %s".
+        format(duplicateTopicsToReassign.mkString(",")))
+    (brokerListToReassign, topicsToReassign)
+  }
 
-    val groupedByTopic = currentAssignment.groupBy { case (tp, _) => tp.topic }
-    val rackAwareMode = if (disableRackAware) RackAwareMode.Disabled else RackAwareMode.Enforced
-    val brokerMetadatas = AdminUtils.getBrokerMetadatas(zkUtils, rackAwareMode, Some(brokerListToReassign))
+  /**
+   * The entry point for the --execute and --execute-additional commands.
+   *
+   * @param adminClient                 The AdminClient to use.
+   * @param additional                  Whether --additional was passed.
+   * @param reassignmentJson            The JSON string to use for the topics to reassign.
+   * @param interBrokerThrottle         The inter-broker throttle to use, or a negative
+   *                                    number to skip using a throttle.
+   * @param logDirThrottle              The replica log directory throttle to use, or a
+   *                                    negative number to skip using a throttle.
+   * @param timeoutMs                   The maximum time in ms to wait for log directory
+   *                                    replica assignment to begin.
+   * @param time                        The Time object to use.
+   */
+  def executeAssignment(adminClient: Admin,
+                        additional: Boolean,
+                        reassignmentJson: String,
+                        interBrokerThrottle: Long = -1L,
+                        logDirThrottle: Long = -1L,
+                        timeoutMs: Long = 10000L,
+                        time: Time = Time.SYSTEM): Unit = {
+    val (proposedParts, proposedReplicas) = parseExecuteAssignmentArgs(reassignmentJson)
+    val currentReassignments = adminClient.
+      listPartitionReassignments().reassignments().get().asScala
+    // If there is an existing assignment, check for --additional before proceeding.
+    // This helps avoid surprising users.
+    if (!additional && currentReassignments.nonEmpty) {
+      throw new TerseReassignmentFailureException(cannotExecuteBecauseOfExistingMessage)
+    }
+    verifyBrokerIds(adminClient, proposedParts.values.flatten.toSet)
+    val currentParts = getReplicaAssignmentForPartitions(adminClient, proposedParts.keySet.toSet)
+    println(currentPartitionReplicaAssignmentToString(proposedParts, currentParts))
 
-    val partitionsToBeReassigned = mutable.Map[TopicAndPartition, Seq[Int]]()
-    groupedByTopic.foreach { case (topic, assignment) =>
-      val (_, replicas) = assignment.head
-      val assignedReplicas = AdminUtils.assignReplicasToBrokers(brokerMetadatas, assignment.size, replicas.size)
-      partitionsToBeReassigned ++= assignedReplicas.map { case (partition, replicas) =>
-        TopicAndPartition(topic, partition) -> replicas
+    if (interBrokerThrottle >= 0 || logDirThrottle >= 0) {
+      println(youMustRunVerifyPeriodicallyMessage)
+
+      if (interBrokerThrottle >= 0) {
+        val moveMap = calculateProposedMoveMap(currentReassignments, proposedParts, currentParts)
+        modifyReassignmentThrottle(adminClient, moveMap, interBrokerThrottle)
       }
+
+      if (logDirThrottle >= 0) {
+        val movingBrokers = calculateMovingBrokers(proposedReplicas.keySet.toSet)
+        modifyLogDirThrottle(adminClient, movingBrokers, logDirThrottle)
+      }
+    }
+
+    // Execute the partition reassignments.
+    val errors = alterPartitionReassignments(adminClient, proposedParts)
+    if (errors.nonEmpty) {
+      throw new TerseReassignmentFailureException(
+        "Error reassigning partition(s):%n%s".format(
+          errors.keySet.toBuffer.sortWith(compareTopicPartitions).map { part =>
+            s"$part: ${errors(part).getMessage}"
+          }.mkString(System.lineSeparator())))
+    }
+    println("Successfully started partition reassignment%s for %s".format(
+      if (proposedParts.size == 1) "" else "s",
+      proposedParts.keySet.toBuffer.sortWith(compareTopicPartitions).mkString(",")))
+    if (proposedReplicas.nonEmpty) {
+      executeMoves(adminClient, proposedReplicas, timeoutMs, time)
     }
-    (partitionsToBeReassigned, currentAssignment)
   }
 
-  def executeAssignment(zkUtils: ZkUtils, adminClientOpt: Option[JAdminClient], opts: ReassignPartitionsCommandOptions) {
-    val reassignmentJsonFile =  opts.options.valueOf(opts.reassignmentJsonFileOpt)
-    val reassignmentJsonString = Utils.readFileAsString(reassignmentJsonFile)
-    val interBrokerThrottle = opts.options.valueOf(opts.interBrokerThrottleOpt)
-    val replicaAlterLogDirsThrottle = opts.options.valueOf(opts.replicaAlterLogDirsThrottleOpt)
-    val timeoutMs = opts.options.valueOf(opts.timeoutOpt)
-    executeAssignment(zkUtils, adminClientOpt, reassignmentJsonString, Throttle(interBrokerThrottle, replicaAlterLogDirsThrottle), timeoutMs)
+  /**
+   * Execute some partition log directory movements.
+   *
+   * @param adminClient                 The AdminClient to use.
+   * @param proposedReplicas            A map from TopicPartitionReplicas to the
+   *                                    directories to move them to.
+   * @param timeoutMs                   The maximum time in ms to wait for log directory
+   *                                    replica assignment to begin.
+   * @param time                        The Time object to use.
+   */
+  def executeMoves(adminClient: Admin,
+                   proposedReplicas: Map[TopicPartitionReplica, String],
+                   timeoutMs: Long,
+                   time: Time): Unit = {
+    val startTimeMs = time.milliseconds()
+    val pendingReplicas = new mutable.HashMap[TopicPartitionReplica, String]()
+    pendingReplicas ++= proposedReplicas
+    var done = false
+    do {
+      val completed = alterReplicaLogDirs(adminClient, pendingReplicas)
+      if (completed.nonEmpty) {
+        println("Successfully started log directory move%s for: %s".format(
+          if (completed.size == 1) "" else "s",
+          completed.toBuffer.sortWith(compareTopicPartitionReplicas).mkString(",")))
+      }
+      pendingReplicas --= completed
+      if (pendingReplicas.isEmpty) {
+        done = true
+      } else if (time.milliseconds() >= startTimeMs + timeoutMs) {
+        throw new TerseReassignmentFailureException(
+          "Timed out before log directory move%s could be started for: %s".format(
+            if (pendingReplicas.size == 1) "" else "s",
+            pendingReplicas.keySet.toBuffer.sortWith(compareTopicPartitionReplicas).
+              mkString(",")))
+      } else {
+        // If a replica has been moved to a new host and we also specified a particular
+        // log directory, we will have to keep retrying the alterReplicaLogDirs
+        // call.  It can't take effect until the replica is moved to that host.
+        time.sleep(100)
+      }
+    } while (!done)
   }
 
-  def executeAssignment(zkUtils: ZkUtils, adminClientOpt: Option[JAdminClient], reassignmentJsonString: String, throttle: Throttle, timeoutMs: Long = 10000L) {
-    val (partitionAssignment, replicaAssignment) = parseAndValidate(zkUtils, reassignmentJsonString)
-    val reassignPartitionsCommand = new ReassignPartitionsCommand(zkUtils, adminClientOpt, partitionAssignment.toMap, replicaAssignment)
+  /**
+   * Entry point for the --list command.
+   *
+   * @param adminClient   The AdminClient to use.
+   */
+  def listReassignments(adminClient: Admin): Unit = {
+    println(curReassignmentsToString(adminClient))
+  }
 
-    // If there is an existing rebalance running, attempt to change its throttle
-    if (zkUtils.pathExists(ZkUtils.ReassignPartitionsPath)) {
-      println("There is an existing assignment running.")
-      reassignPartitionsCommand.maybeLimit(throttle)
+  /**
+   * Convert the current partition reassignments to text.
+   *
+   * @param adminClient   The AdminClient to use.
+   * @return              A string describing the current partition reassignments.
+   */
+  def curReassignmentsToString(adminClient: Admin): String = {
+    val currentReassignments = adminClient.
+      listPartitionReassignments().reassignments().get().asScala
+    val text = currentReassignments.keySet.toBuffer.sortWith(compareTopicPartitions).map { part =>
+      val reassignment = currentReassignments(part)
+      val replicas = reassignment.replicas.asScala
+      val addingReplicas = reassignment.addingReplicas.asScala
+      val removingReplicas = reassignment.removingReplicas.asScala
+      "%s: replicas: %s.%s%s".format(part, replicas.mkString(","),
+        if (addingReplicas.isEmpty) "" else
+          " adding: %s.".format(addingReplicas.mkString(",")),
+        if (removingReplicas.isEmpty) "" else
+          " removing: %s.".format(removingReplicas.mkString(",")))
+    }.mkString(System.lineSeparator())
+    if (text.isEmpty) {
+      "No partition reassignments found."
     } else {
-      printCurrentAssignment(zkUtils, partitionAssignment.map(_._1.topic))
-      if (throttle.interBrokerLimit >= 0 || throttle.replicaAlterLogDirsLimit >= 0)
-        println(String.format("Warning: You must run Verify periodically, until the reassignment completes, to ensure the throttle is removed. You can also alter the throttle by rerunning the Execute command passing a new value."))
-      if (reassignPartitionsCommand.reassignPartitions(throttle, timeoutMs)) {
-        println("Successfully started reassignment of partitions.")
-      } else
-        println("Failed to reassign partitions %s".format(partitionAssignment))
+      "Current partition reassignments:%n%s".format(text)
     }
   }
 
-  def printCurrentAssignment(zkUtils: ZkUtils, topics: Seq[String]): Unit = {
-    // before starting assignment, output the current replica assignment to facilitate rollback
-    val currentPartitionReplicaAssignment = zkUtils.getReplicaAssignmentForTopics(topics)
-    println("Current partition replica assignment\n\n%s\n\nSave this to use as the --reassignment-json-file option during rollback"
-      .format(formatAsReassignmentJson(currentPartitionReplicaAssignment, Map.empty)))
+  /**
+   * Verify that all the brokers in an assignment exist.
+   *
+   * @param adminClient                 The AdminClient to use.
+   * @param brokers                     The broker IDs to verify.
+   */
+  def verifyBrokerIds(adminClient: Admin, brokers: Set[Int]): Unit = {
+    val allNodeIds = adminClient.describeCluster().nodes().get().asScala.map(_.id).toSet
+    brokers.find(!allNodeIds.contains(_)).map {
+      id => throw new AdminCommandFailedException(s"Unknown broker id ${id}")
+    }
   }
 
-  def formatAsReassignmentJson(partitionsToBeReassigned: Map[TopicAndPartition, Seq[Int]],
-                               replicaLogDirAssignment: Map[TopicPartitionReplica, String]): String = {
-    Json.encode(Map(
-      "version" -> 1,
-      "partitions" -> partitionsToBeReassigned.map { case (TopicAndPartition(topic, partition), replicas) =>
-        Map(
-          "topic" -> topic,
-          "partition" -> partition,
-          "replicas" -> replicas,
-          "log_dirs" -> replicas.map(r => replicaLogDirAssignment.getOrElse(new TopicPartitionReplica(topic, partition, r), AnyLogDir))
-        )
+  /**
+   * The entry point for the --execute command.
+   *
+   * @param zkClient                    The ZooKeeper client to use.
+   * @param reassignmentJson            The JSON string to use for the topics to reassign.
+   * @param interBrokerThrottle         The inter-broker throttle to use, or a negative number
+   *                                    to skip using a throttle.
+   */
+  def executeAssignment(zkClient: KafkaZkClient,
+                        reassignmentJson: String,
+                        interBrokerThrottle: Long): Unit = {
+    val (proposedParts, proposedReplicas) = parseExecuteAssignmentArgs(reassignmentJson)
+    if (proposedReplicas.nonEmpty) {
+      throw new AdminCommandFailedException("bootstrap-server needs to be provided when " +
+        "replica reassignments are present.")
+    }
+    verifyReplicasAndBrokersInAssignment(zkClient, proposedParts)
+
+    // Check for the presence of the legacy partition reassignment ZNode.  This actually
+    // won't detect all rebalances... only ones initiated by the legacy method.
+    // This is a limitation of the legacy ZK API.
+    val reassignPartitionsInProgress = zkClient.reassignPartitionsInProgress
+    if (reassignPartitionsInProgress) {
+      // Note: older versions of this tool would modify the broker quotas here (but not
+      // topic quotas, for some reason).  Since it might interfere with other ongoing
+      // reassignments, this behavior was dropped as part of the KIP-455 changes. The
+      // user can still alter existing throttles by resubmitting the current reassignment
+      // and providing the --additional flag.
+      throw new TerseReassignmentFailureException(cannotExecuteBecauseOfExistingMessage)
+    }
+    val currentParts = zkClient.getReplicaAssignmentForTopics(
+      proposedParts.map(_._1.topic()).toSet)
+    println(currentPartitionReplicaAssignmentToString(proposedParts, currentParts))
+
+    if (interBrokerThrottle >= 0) {
+      println(youMustRunVerifyPeriodicallyMessage)
+      val moveMap = calculateProposedMoveMap(Map.empty, proposedParts, currentParts)
+      val leaderThrottles = calculateLeaderThrottles(moveMap)
+      val followerThrottles = calculateFollowerThrottles(moveMap)
+      modifyTopicThrottles(zkClient, leaderThrottles, followerThrottles)
+      val reassigningBrokers = calculateReassigningBrokers(moveMap)
+      modifyBrokerThrottles(zkClient, reassigningBrokers, interBrokerThrottle)
+      println(s"The inter-broker throttle limit was set to ${interBrokerThrottle} B/s")
+    }
+    zkClient.createPartitionReassignment(proposedParts)
+    println("Successfully started partition reassignment%s for %s".format(
+      if (proposedParts.size == 1) "" else "s",
+      proposedParts.keySet.toBuffer.sortWith(compareTopicPartitions).mkString(",")))
+  }
+
+  /**
+   * Return the string which we want to print to describe the current partition assignment.
+   *
+   * @param proposedParts               The proposed partition assignment.
+   * @param currentParts                The current partition assignment.
+   *
+   * @return                            The string to print.  We will only print information about
+   *                                    partitions that appear in the proposed partition assignment.
+   */
+  def currentPartitionReplicaAssignmentToString(proposedParts: Map[TopicPartition, Seq[Int]],
+                                                currentParts: Map[TopicPartition, Seq[Int]]): String = {
+    "Current partition replica assignment%n%n%s%n%nSave this to use as the %s".
+        format(formatAsReassignmentJson(currentParts.filter { case (k, _) => proposedParts.contains(k) }.toMap, Map.empty),
+              "--reassignment-json-file option during rollback")
+  }
+
+  /**
+   * Verify that the replicas and brokers referenced in the given partition assignment actually
+   * exist.  This is necessary when using the deprecated ZK API, since ZooKeeper itself can't
+   * validate what we're applying.
+   *
+   * @param zkClient                    The ZooKeeper client to use.
+   * @param proposedParts               The partition assignment.
+   */
+  def verifyReplicasAndBrokersInAssignment(zkClient: KafkaZkClient,
+                                           proposedParts: Map[TopicPartition, Seq[Int]]): Unit = {
+    // check that all partitions in the proposed assignment exist in the cluster
+    val proposedTopics = proposedParts.map { case (tp, _) => tp.topic }
+    val existingAssignment = zkClient.getReplicaAssignmentForTopics(proposedTopics.toSet)
+    val nonExistentPartitions = proposedParts.map { case (tp, _) => tp }.filterNot(existingAssignment.contains)
+    if (nonExistentPartitions.nonEmpty)
+      throw new AdminCommandFailedException("The proposed assignment contains non-existent partitions: " +
+        nonExistentPartitions)
+
+    // check that all brokers in the proposed assignment exist in the cluster
+    val existingBrokerIDs = zkClient.getSortedBrokerList
+    val nonExistingBrokerIDs = proposedParts.toMap.values.flatten.filterNot(existingBrokerIDs.contains).toSet
+    if (nonExistingBrokerIDs.nonEmpty)
+      throw new AdminCommandFailedException("The proposed assignment contains non-existent brokerIDs: " + nonExistingBrokerIDs.mkString(","))
+  }
+
+  /**
+   * Execute the given partition reassignments.
+   *
+   * @param adminClient       The admin client object to use.
+   * @param reassignments     A map from topic names to target replica assignments.
+   * @return                  A map from partition objects to error strings.
+   */
+  def alterPartitionReassignments(adminClient: Admin,
+                                  reassignments: Map[TopicPartition, Seq[Int]]): Map[TopicPartition, Throwable] = {
+    val results: Map[TopicPartition, KafkaFuture[Void]] =
+      adminClient.alterPartitionReassignments(reassignments.map { case (part, replicas) =>
+        (part, Optional.of(new NewPartitionReassignment(replicas.map(Integer.valueOf(_)).asJava)))
+      }.asJava).values().asScala
+    results.flatMap {
+      case (part, future) => {
+        try {
+          future.get()
+          None
+        } catch {
+          case t: ExecutionException => Some(part, t.getCause())
+        }
       }
-    ))
+    }
   }
 
-  // Parses without deduplicating keys so the data can be checked before allowing reassignment to proceed
-  def parsePartitionReassignmentData(jsonData: String): (Seq[(TopicAndPartition, Seq[Int])], Map[TopicPartitionReplica, String]) = {
-    val partitionAssignment = mutable.ListBuffer.empty[(TopicAndPartition, Seq[Int])]
-    val replicaAssignment = mutable.Map.empty[TopicPartitionReplica, String]
-    for {
-      js <- Json.parseFull(jsonData).toSeq
-      partitionsSeq <- js.asJsonObject.get("partitions").toSeq
-      p <- partitionsSeq.asJsonArray.iterator
-    } {
-      val partitionFields = p.asJsonObject
-      val topic = partitionFields("topic").to[String]
-      val partition = partitionFields("partition").to[Int]
-      val newReplicas = partitionFields("replicas").to[Seq[Int]]
-      val newLogDirs = partitionFields.get("log_dirs") match {
-        case Some(jsonValue) => jsonValue.to[Seq[String]]
-        case None => newReplicas.map(r => AnyLogDir)
+  /**
+   * Cancel the given partition reassignments.
+   *
+   * @param adminClient       The admin client object to use.
+   * @param reassignments     The partition reassignments to cancel.
+   * @return                  A map from partition objects to error strings.
+   */
+  def cancelPartitionReassignments(adminClient: Admin,
+                                   reassignments: Set[TopicPartition])
+  : Map[TopicPartition, Throwable] = {
+    val results: Map[TopicPartition, KafkaFuture[Void]] =
+      adminClient.alterPartitionReassignments(reassignments.map {
+          (_, (None: Option[NewPartitionReassignment]).asJava)
+        }.toMap.asJava).values().asScala
+    results.flatMap { case (part, future) =>
+      try {
+        future.get()
+        None
+      } catch {
+        case t: ExecutionException => Some(part, t.getCause())
+      }
+    }
+  }
+
+  private def calculateCurrentMoveMap(currentReassignments: Map[TopicPartition, PartitionReassignment]): MoveMap = {
+    val moveMap = new mutable.HashMap[String, mutable.Map[Int, PartitionMove]]()
+    // Add the current reassignments to the move map.
+    currentReassignments.foreach { case (part, reassignment) =>
+      val move = PartitionMove(new mutable.HashSet[Int](), new mutable.HashSet[Int]())
+      reassignment.replicas.forEach { replica =>
+        move.sources += replica
+        move.destinations += replica
+      }
+      reassignment.addingReplicas.forEach(move.destinations += _)
+      reassignment.removingReplicas.forEach(move.destinations -= _)
+      val partMoves = moveMap.getOrElseUpdate(part.topic, new mutable.HashMap[Int, PartitionMove])
+      partMoves.put(part.partition, move)
+    }
+    moveMap
+  }
+
+  /**
+   * Calculate the global map of all partitions that are moving.
+   *
+   * @param currentReassignments    The currently active reassignments.
+   * @param proposedReassignments   The proposed reassignments (destinations replicas only).
+   * @param currentParts            The current location of the partitions that we are
+   *                                proposing to move.
+   * @return                        A map from topic name to partition map.
+   *                                The partition map is keyed on partition index and contains
+   *                                the movements for that partition.
+   */
+  def calculateProposedMoveMap(currentReassignments: Map[TopicPartition, PartitionReassignment],
+                               proposedReassignments: Map[TopicPartition, Seq[Int]],
+                               currentParts: Map[TopicPartition, Seq[Int]]): MoveMap = {
+    val moveMap = calculateCurrentMoveMap(currentReassignments)
+
+    // Add the proposed reassignments to the move map.  The proposals will overwrite
+    // the current reassignments.
+    proposedReassignments.foreach {
+      case (part, replicas) => {
+        val move = PartitionMove(new mutable.HashSet[Int](), new mutable.HashSet[Int]())
+        move.destinations ++= replicas
+        val partMoves = moveMap.getOrElseUpdate(part.topic(), new mutable.HashMap[Int, PartitionMove])
+        partMoves.put(part.partition(), move)
       }
-      if (newReplicas.size != newLogDirs.size)
-        throw new AdminCommandFailedException(s"Size of replicas list $newReplicas is different from " +
-          s"size of log dirs list $newLogDirs for partition ${TopicAndPartition(topic, partition)}")
-      partitionAssignment += (TopicAndPartition(topic, partition) -> newReplicas)
-      replicaAssignment ++= newReplicas.zip(newLogDirs).map { case (replica, logDir) =>
-        new TopicPartitionReplica(topic, partition, replica) -> logDir
-      }.filter(_._2 != AnyLogDir)
     }
-    (partitionAssignment, replicaAssignment)
+    // For partitions we are moving, add the current replica locations as sources.
+    // Ignore partitions that are not being moved.
+    moveMap.foreach {
+      case (topicName, partMap) =>
+        partMap.foreach {
+          case (partitionIndex, moves) =>
+            currentParts.get(new TopicPartition(topicName, partitionIndex)) match {
+              case None =>
+              case Some(replicas) => moves.sources ++= replicas
+            }
+        }
+    }
+    // Remove sources from destinations.  If something is a source, the data is already there,
+    // so it doesn't need to be treated as a destination (by having follower throttle applied, etc.)
+    moveMap.foreach {
+      case (_, partMap) =>
+        partMap.foreach {
+          case (_, moves) =>
+            moves.destinations --= moves.sources
+        }
+    }
+    moveMap
   }
 
-  def parseAndValidate(zkUtils: ZkUtils, reassignmentJsonString: String): (Seq[(TopicAndPartition, Seq[Int])], Map[TopicPartitionReplica, String]) = {
-    val (partitionsToBeReassigned, replicaAssignment) = parsePartitionReassignmentData(reassignmentJsonString)
+  /**
+   * Calculate the leader throttle configurations to use.
+   *
+   * @param moveMap   The movements.
+   * @return          A map from topic names to leader throttle configurations.
+   */
+  def calculateLeaderThrottles(moveMap: MoveMap): Map[String, String] = {
+    moveMap.map {
+      case (topicName, partMoveMap) => {
+        val components = new mutable.TreeSet[String]
+        partMoveMap.foreach {
+          case (partId, move) =>
+            move.sources.foreach(source => components.add("%d:%d".format(partId, source)))
+        }
+        (topicName, components.mkString(","))
+      }
+    }
+  }
+
+  /**
+   * Calculate the follower throttle configurations to use.
+   *
+   * @param moveMap   The movements.
+   * @return          A map from topic names to follower throttle configurations.
+   */
+  def calculateFollowerThrottles(moveMap: MoveMap): Map[String, String] = {
+    moveMap.map {
+      case (topicName, partMoveMap) => {
+        val components = new mutable.TreeSet[String]
+        partMoveMap.foreach {
+          case (partId, move) =>
+            move.destinations.foreach(destination =>
+              if (!move.sources.contains(destination)) {
+                components.add("%d:%d".format(partId, destination))
+              })
+        }
+        (topicName, components.mkString(","))
+      }
+    }
+  }
 
+  /**
+   * Calculate all the brokers which are involved in the given partition reassignments.
+   *
+   * @param moveMap       The partition movements.
+   * @return              A set of all the brokers involved.
+   */
+  def calculateReassigningBrokers(moveMap: MoveMap): Set[Int] = {
+    val reassigningBrokers = new mutable.TreeSet[Int]
+    moveMap.values.foreach {
+      _.values.foreach {
+        partMove =>
+          partMove.sources.foreach(reassigningBrokers.add)
+          partMove.destinations.foreach(reassigningBrokers.add)
+      }
+    }
+    reassigningBrokers.toSet
+  }
+
+  /**
+   * Calculate all the brokers which are involved in the given directory movements.
+   *
+   * @param replicaMoves  The replica movements.
+   * @return              A set of all the brokers involved.
+   */
+  def calculateMovingBrokers(replicaMoves: Set[TopicPartitionReplica]): Set[Int] = {
+    replicaMoves.map(_.brokerId())
+  }
+
+  /**
+   * Modify the topic configurations that control inter-broker throttling.
+   *
+   * @param adminClient         The adminClient object to use.
+   * @param leaderThrottles     A map from topic names to leader throttle configurations.
+   * @param followerThrottles   A map from topic names to follower throttle configurations.
+   */
+  def modifyTopicThrottles(adminClient: Admin,
+                           leaderThrottles: Map[String, String],
+                           followerThrottles: Map[String, String]): Unit = {
+    val configs = new util.HashMap[ConfigResource, util.Collection[AlterConfigOp]]()
+    val topicNames = leaderThrottles.keySet ++ followerThrottles.keySet
+    topicNames.foreach {
+      topicName =>
+        val ops = new util.ArrayList[AlterConfigOp]
+        leaderThrottles.get(topicName).foreach { value =>
+          ops.add(new AlterConfigOp(new ConfigEntry(topicLevelLeaderThrottle, value), OpType.SET))
+        }
+        followerThrottles.get(topicName).foreach { value =>
+          ops.add(new AlterConfigOp(new ConfigEntry(topicLevelFollowerThrottle, value), OpType.SET))
+        }
+        if (!ops.isEmpty) {
+          configs.put(new ConfigResource(ConfigResource.Type.TOPIC, topicName), ops)
+        }
+    }
+    adminClient.incrementalAlterConfigs(configs).all().get()
+  }
+
+  /**
+   * Modify the topic configurations that control inter-broker throttling.
+   *
+   * @param zkClient            The ZooKeeper client to use.
+   * @param leaderThrottles     A map from topic names to leader throttle configurations.
+   * @param followerThrottles   A map from topic names to follower throttle configurations.
+   */
+  def modifyTopicThrottles(zkClient: KafkaZkClient,
+                           leaderThrottles: Map[String, String],
+                           followerThrottles: Map[String, String]): Unit = {
+    val adminZkClient = new AdminZkClient(zkClient)
+    val topicNames = leaderThrottles.keySet ++ followerThrottles.keySet
+    topicNames.foreach {
+      topicName =>
+        val configs = adminZkClient.fetchEntityConfig(ConfigType.Topic, topicName)
+        leaderThrottles.get(topicName).map(configs.put(topicLevelLeaderThrottle, _))
+        followerThrottles.get(topicName).map(configs.put(topicLevelFollowerThrottle, _))
+        adminZkClient.changeTopicConfig(topicName, configs)
+    }
+  }
+
+  private def modifyReassignmentThrottle(admin: Admin, moveMap: MoveMap, interBrokerThrottle: Long): Unit = {
+    val leaderThrottles = calculateLeaderThrottles(moveMap)
+    val followerThrottles = calculateFollowerThrottles(moveMap)
+    modifyTopicThrottles(admin, leaderThrottles, followerThrottles)
+
+    val reassigningBrokers = calculateReassigningBrokers(moveMap)
+    modifyInterBrokerThrottle(admin, reassigningBrokers, interBrokerThrottle)
+  }
+
+  /**
+   * Modify the leader/follower replication throttles for a set of brokers.
+   *
+   * @param adminClient The Admin instance to use
+   * @param reassigningBrokers The set of brokers involved in the reassignment
+   * @param interBrokerThrottle The new throttle (ignored if less than 0)
+   */
+  def modifyInterBrokerThrottle(adminClient: Admin,
+                                reassigningBrokers: Set[Int],
+                                interBrokerThrottle: Long): Unit = {
+    if (interBrokerThrottle >= 0) {
+      val configs = new util.HashMap[ConfigResource, util.Collection[AlterConfigOp]]()
+      reassigningBrokers.foreach { brokerId =>
+        val ops = new util.ArrayList[AlterConfigOp]
+        ops.add(new AlterConfigOp(new ConfigEntry(brokerLevelLeaderThrottle,
+          interBrokerThrottle.toString), OpType.SET))
+        ops.add(new AlterConfigOp(new ConfigEntry(brokerLevelFollowerThrottle,
+          interBrokerThrottle.toString), OpType.SET))
+        configs.put(new ConfigResource(ConfigResource.Type.BROKER, brokerId.toString), ops)
+      }
+      adminClient.incrementalAlterConfigs(configs).all().get()
+      println(s"The inter-broker throttle limit was set to $interBrokerThrottle B/s")
+    }
+  }
+
+  /**
+   * Modify the log dir reassignment throttle for a set of brokers.
+   *
+   * @param admin The Admin instance to use
+   * @param movingBrokers The set of broker to alter the throttle of
+   * @param logDirThrottle The new throttle (ignored if less than 0)
+   */
+  def modifyLogDirThrottle(admin: Admin,
+                           movingBrokers: Set[Int],
+                           logDirThrottle: Long): Unit = {
+    if (logDirThrottle >= 0) {
+      val configs = new util.HashMap[ConfigResource, util.Collection[AlterConfigOp]]()
+      movingBrokers.foreach { brokerId =>
+        val ops = new util.ArrayList[AlterConfigOp]
+        ops.add(new AlterConfigOp(new ConfigEntry(brokerLevelLogDirThrottle, logDirThrottle.toString), OpType.SET))
+        configs.put(new ConfigResource(ConfigResource.Type.BROKER, brokerId.toString), ops)
+      }
+      admin.incrementalAlterConfigs(configs).all().get()
+      println(s"The replica-alter-dir throttle limit was set to $logDirThrottle B/s")
+    }
+  }
+
+  /**
+   * Modify the broker-level configurations for leader and follower throttling.
+   *
+   * @param zkClient            The ZooKeeper client to use.
+   * @param reassigningBrokers  The brokers to reconfigure.
+   * @param interBrokerThrottle The throttle value to set.
+   */
+  def modifyBrokerThrottles(zkClient: KafkaZkClient,
+                            reassigningBrokers: Set[Int],
+                            interBrokerThrottle: Long): Unit = {
+    val adminZkClient = new AdminZkClient(zkClient)
+    for (id <- reassigningBrokers) {
+      val configs = adminZkClient.fetchEntityConfig(ConfigType.Broker, id.toString)
+      configs.put(brokerLevelLeaderThrottle, interBrokerThrottle.toString)
+      configs.put(brokerLevelFollowerThrottle, interBrokerThrottle.toString)
+      adminZkClient.changeBrokerConfig(Seq(id), configs)
+    }
+  }
+
+  /**
+   * Parse the reassignment JSON string passed to the --execute command.
+   *
+   * @param reassignmentJson  The JSON string.
+   * @return                  A tuple of the partitions to be reassigned and the replicas
+   *                          to be reassigned.
+   */
+  def parseExecuteAssignmentArgs(reassignmentJson: String)
+      : (Map[TopicPartition, Seq[Int]], Map[TopicPartitionReplica, String]) = {
+    val (partitionsToBeReassigned, replicaAssignment) = parsePartitionReassignmentData(reassignmentJson)
     if (partitionsToBeReassigned.isEmpty)
-      throw new AdminCommandFailedException("Partition reassignment data file is empty")
+      throw new AdminCommandFailedException("Partition reassignment list cannot be empty")
     if (partitionsToBeReassigned.exists(_._2.isEmpty)) {
       throw new AdminCommandFailedException("Partition replica list cannot be empty")
     }
@@ -285,146 +1515,311 @@ object ReassignPartitionsCommand extends Logging {
         .mkString(". ")
       throw new AdminCommandFailedException("Partition replica lists may not contain duplicate entries: %s".format(duplicatesMsg))
     }
-    // check that all partitions in the proposed assignment exist in the cluster
-    val proposedTopics = partitionsToBeReassigned.map { case (tp, _) => tp.topic }.distinct
-    val existingAssignment = zkUtils.getReplicaAssignmentForTopics(proposedTopics)
-    val nonExistentPartitions = partitionsToBeReassigned.map { case (tp, _) => tp }.filterNot(existingAssignment.contains)
-    if (nonExistentPartitions.nonEmpty)
-      throw new AdminCommandFailedException("The proposed assignment contains non-existent partitions: " +
-        nonExistentPartitions)
+    (partitionsToBeReassigned.toMap, replicaAssignment)
+  }
 
-    // check that all brokers in the proposed assignment exist in the cluster
-    val existingBrokerIDs = zkUtils.getSortedBrokerList()
-    val nonExistingBrokerIDs = partitionsToBeReassigned.toMap.values.flatten.filterNot(existingBrokerIDs.contains).toSet
-    if (nonExistingBrokerIDs.nonEmpty)
-      throw new AdminCommandFailedException("The proposed assignment contains non-existent brokerIDs: " + nonExistingBrokerIDs.mkString(","))
+  /**
+   * The entry point for the --cancel command.
+   *
+   * @param adminClient           The AdminClient to use.
+   * @param jsonString            The JSON string to use for the topics and partitions to cancel.
+   * @param preserveThrottles     True if we should avoid changing topic or broker throttles.
+   * @param timeoutMs             The maximum time in ms to wait for log directory
+   *                              replica assignment to begin.
+   * @param time                  The Time object to use.
+   *
+   * @return                      A tuple of the partition reassignments that were cancelled,
+   *                              and the replica movements that were cancelled.
+   */
+  def cancelAssignment(adminClient: Admin,
+                       jsonString: String,
+                       preserveThrottles: Boolean,
+                       timeoutMs: Long = 10000L,
+                       time: Time = Time.SYSTEM)
+                       : (Set[TopicPartition], Set[TopicPartitionReplica]) = {
+    val (targetParts, targetReplicas) = parsePartitionReassignmentData(jsonString)
+    val targetPartsSet = targetParts.map(_._1).toSet
+    val curReassigningParts = adminClient.listPartitionReassignments(targetPartsSet.asJava).
+        reassignments().get().asScala.flatMap {
+      case (part, reassignment) => if (!reassignment.addingReplicas().isEmpty ||
+          !reassignment.removingReplicas().isEmpty) {
+        Some(part)
+      } else {
+        None
+      }
+    }.toSet
+    if (curReassigningParts.nonEmpty) {
+      val errors = cancelPartitionReassignments(adminClient, curReassigningParts)
+      if (errors.nonEmpty) {
+        throw new TerseReassignmentFailureException(
+          "Error cancelling partition reassignment%s for:%n%s".format(
+            if (errors.size == 1) "" else "s",
+            errors.keySet.toBuffer.sortWith(compareTopicPartitions).map {
+              part => s"${part}: ${errors(part).getMessage}"
+            }.mkString(System.lineSeparator())))
+      }
+      println("Successfully cancelled partition reassignment%s for: %s".format(
+        if (curReassigningParts.size == 1) "" else "s",
+        s"${curReassigningParts.toBuffer.sortWith(compareTopicPartitions).mkString(",")}"))
+    } else {
+      println("None of the specified partition reassignments are active.")
+    }
+    val curMovingParts = findLogDirMoveStates(adminClient, targetReplicas).flatMap {
+      case (part, moveState) => moveState match {
+        case state: ActiveMoveState => Some(part, state.currentLogDir)
+        case _ => None
+      }
+    }.toMap
+    if (curMovingParts.isEmpty) {
+      println("None of the specified partition moves are active.")
+    } else {
+      executeMoves(adminClient, curMovingParts, timeoutMs, time)
+    }
+    if (!preserveThrottles) {
+      clearAllThrottles(adminClient, targetParts)
+    }
+    (curReassigningParts, curMovingParts.keySet)
+  }
 
-    (partitionsToBeReassigned, replicaAssignment)
+  def formatAsReassignmentJson(partitionsToBeReassigned: Map[TopicPartition, Seq[Int]],
+                               replicaLogDirAssignment: Map[TopicPartitionReplica, String]): String = {
+    Json.encodeAsString(Map(
+      "version" -> 1,
+      "partitions" -> partitionsToBeReassigned.keySet.toBuffer.sortWith(compareTopicPartitions).map {
+        tp =>
+          val replicas = partitionsToBeReassigned(tp)
+          Map(
+            "topic" -> tp.topic,
+            "partition" -> tp.partition,
+            "replicas" -> replicas.asJava,
+            "log_dirs" -> replicas.map(r => replicaLogDirAssignment.getOrElse(new TopicPartitionReplica(tp.topic, tp.partition, r), AnyLogDir)).asJava
+          ).asJava
+      }.asJava
+    ).asJava)
   }
 
-  private def checkIfPartitionReassignmentSucceeded(zkUtils: ZkUtils, partitionsToBeReassigned: Map[TopicAndPartition, Seq[Int]])
-  :Map[TopicAndPartition, ReassignmentStatus] = {
-    val partitionsBeingReassigned = zkUtils.getPartitionsBeingReassigned().mapValues(_.newReplicas)
-    partitionsToBeReassigned.keys.map { topicAndPartition =>
-      (topicAndPartition, checkIfPartitionReassignmentSucceeded(zkUtils, topicAndPartition, partitionsToBeReassigned,
-        partitionsBeingReassigned))
-    }.toMap
+  def parseTopicsData(jsonData: String): Seq[String] = {
+    Json.parseFull(jsonData) match {
+      case Some(js) =>
+        val version = js.asJsonObject.get("version") match {
+          case Some(jsonValue) => jsonValue.to[Int]
+          case None => EarliestTopicsJsonVersion
+        }
+        parseTopicsData(version, js)
+      case None => throw new AdminOperationException("The input string is not a valid JSON")
+    }
   }
 
-  private def checkIfReplicaReassignmentSucceeded(adminClientOpt: Option[JAdminClient], replicaAssignment: Map[TopicPartitionReplica, String])
-  :Map[TopicPartitionReplica, ReassignmentStatus] = {
+  def parseTopicsData(version: Int, js: JsonValue): Seq[String] = {
+    version match {
+      case 1 =>
+        for {
+          partitionsSeq <- js.asJsonObject.get("topics").toSeq
+          p <- partitionsSeq.asJsonArray.iterator
+        } yield p.asJsonObject("topic").to[String]
+      case _ => throw new AdminOperationException(s"Not supported version field value $version")
+    }
+  }
 
-    val replicaLogDirInfos = {
-      if (replicaAssignment.nonEmpty) {
-        val adminClient = adminClientOpt.getOrElse(
-          throw new AdminCommandFailedException("bootstrap-server needs to be provided in order to reassign replica to the specified log directory"))
-        adminClient.describeReplicaLogDirs(replicaAssignment.keySet.asJava).all().get().asScala
-      } else {
-        Map.empty[TopicPartitionReplica, ReplicaLogDirInfo]
-      }
+  def parsePartitionReassignmentData(jsonData: String): (Seq[(TopicPartition, Seq[Int])], Map[TopicPartitionReplica, String]) = {
+    Json.parseFull(jsonData) match {
+      case Some(js) =>
+        val version = js.asJsonObject.get("version") match {
+          case Some(jsonValue) => jsonValue.to[Int]
+          case None => EarliestVersion
+        }
+        parsePartitionReassignmentData(version, js)
+      case None => throw new AdminOperationException("The input string is not a valid JSON")
     }
+  }
 
-    replicaAssignment.map { case (replica, newLogDir) =>
-      val status: ReassignmentStatus = replicaLogDirInfos.get(replica) match {
-        case Some(replicaLogDirInfo) =>
-          if (replicaLogDirInfo.getCurrentReplicaLogDir == null) {
-            println(s"Partition ${replica.topic()}-${replica.partition()} is not found in any live log dir on " +
-              s"broker ${replica.brokerId()}. There is likely offline log directory on the broker.")
-            ReassignmentFailed
-          } else if (replicaLogDirInfo.getFutureReplicaLogDir == newLogDir) {
-            ReassignmentInProgress
-          } else if (replicaLogDirInfo.getFutureReplicaLogDir != null) {
-            println(s"Partition ${replica.topic()}-${replica.partition()} on broker ${replica.brokerId()} " +
-              s"is being moved to log dir ${replicaLogDirInfo.getFutureReplicaLogDir} instead of $newLogDir")
-            ReassignmentFailed
-          } else if (replicaLogDirInfo.getCurrentReplicaLogDir == newLogDir) {
-            ReassignmentCompleted
-          } else {
-            println(s"Partition ${replica.topic()}-${replica.partition()} on broker ${replica.brokerId()} " +
-              s"is not being moved from log dir ${replicaLogDirInfo.getCurrentReplicaLogDir} to $newLogDir")
-            ReassignmentFailed
+  // Parses without deduplicating keys so the data can be checked before allowing reassignment to proceed
+  def parsePartitionReassignmentData(version:Int, jsonData: JsonValue): (Seq[(TopicPartition, Seq[Int])], Map[TopicPartitionReplica, String]) = {
+    version match {
+      case 1 =>
+        val partitionAssignment = mutable.ListBuffer.empty[(TopicPartition, Seq[Int])]
+        val replicaAssignment = mutable.Map.empty[TopicPartitionReplica, String]
+        for {
+          partitionsSeq <- jsonData.asJsonObject.get("partitions").toSeq
+          p <- partitionsSeq.asJsonArray.iterator
+        } {
+          val partitionFields = p.asJsonObject
+          val topic = partitionFields("topic").to[String]
+          val partition = partitionFields("partition").to[Int]
+          val newReplicas = partitionFields("replicas").to[Seq[Int]]
+          val newLogDirs = partitionFields.get("log_dirs") match {
+            case Some(jsonValue) => jsonValue.to[Seq[String]]
+            case None => newReplicas.map(_ => AnyLogDir)
           }
-        case None =>
-          println(s"Partition ${replica.topic()}-${replica.partition()} is not found in any live log dir on broker ${replica.brokerId()}.")
-          ReassignmentFailed
-      }
-      (replica, status)
-    }
-  }
-
-  def checkIfPartitionReassignmentSucceeded(zkUtils: ZkUtils, topicAndPartition: TopicAndPartition,
-                                            partitionsToBeReassigned: Map[TopicAndPartition, Seq[Int]],
-                                            partitionsBeingReassigned: Map[TopicAndPartition, Seq[Int]]): ReassignmentStatus = {
-    val newReplicas = partitionsToBeReassigned(topicAndPartition)
-    partitionsBeingReassigned.get(topicAndPartition) match {
-      case Some(_) => ReassignmentInProgress
-      case None =>
-        // check if the current replica assignment matches the expected one after reassignment
-        val assignedReplicas = zkUtils.getReplicasForPartition(topicAndPartition.topic, topicAndPartition.partition)
-        if(assignedReplicas == newReplicas)
-          ReassignmentCompleted
-        else {
-          println(("ERROR: Assigned replicas (%s) don't match the list of replicas for reassignment (%s)" +
-            " for partition %s").format(assignedReplicas.mkString(","), newReplicas.mkString(","), topicAndPartition))
-          ReassignmentFailed
+          if (newReplicas.size != newLogDirs.size)
+            throw new AdminCommandFailedException(s"Size of replicas list $newReplicas is different from " +
+              s"size of log dirs list $newLogDirs for partition ${new TopicPartition(topic, partition)}")
+          partitionAssignment += (new TopicPartition(topic, partition) -> newReplicas)
+          replicaAssignment ++= newReplicas.zip(newLogDirs).map { case (replica, logDir) =>
+            new TopicPartitionReplica(topic, partition, replica) -> logDir
+          }.filter(_._2 != AnyLogDir)
         }
+        (partitionAssignment, replicaAssignment)
+      case _ => throw new AdminOperationException(s"Not supported version field value $version")
     }
   }
 
   def validateAndParseArgs(args: Array[String]): ReassignPartitionsCommandOptions = {
     val opts = new ReassignPartitionsCommandOptions(args)
 
-    if(args.length == 0)
-      CommandLineUtils.printUsageAndDie(opts.parser, "This command moves topic partitions between replicas.")
+    CommandLineUtils.printHelpAndExitIfNeeded(opts, helpText)
 
-    // Should have exactly one action
-    val actions = Seq(opts.generateOpt, opts.executeOpt, opts.verifyOpt).count(opts.options.has _)
-    if(actions != 1)
-      CommandLineUtils.printUsageAndDie(opts.parser, "Command must include exactly one action: --generate, --execute or --verify")
-
-    CommandLineUtils.checkRequiredArgs(opts.parser, opts.options, opts.zkConnectOpt)
-
-    //Validate arguments for each action
-    if(opts.options.has(opts.verifyOpt)) {
-      if(!opts.options.has(opts.reassignmentJsonFileOpt))
-        CommandLineUtils.printUsageAndDie(opts.parser, "If --verify option is used, command must include --reassignment-json-file that was used during the --execute option")
-      CommandLineUtils.checkInvalidArgs(opts.parser, opts.options, opts.verifyOpt, Set(opts.interBrokerThrottleOpt, opts.replicaAlterLogDirsThrottleOpt, opts.topicsToMoveJsonFileOpt, opts.disableRackAware, opts.brokerListOpt))
-    }
-    else if(opts.options.has(opts.generateOpt)) {
-      if(!(opts.options.has(opts.topicsToMoveJsonFileOpt) && opts.options.has(opts.brokerListOpt)))
-        CommandLineUtils.printUsageAndDie(opts.parser, "If --generate option is used, command must include both --topics-to-move-json-file and --broker-list options")
-      CommandLineUtils.checkInvalidArgs(opts.parser, opts.options, opts.generateOpt, Set(opts.interBrokerThrottleOpt, opts.replicaAlterLogDirsThrottleOpt, opts.reassignmentJsonFileOpt))
+    // Determine which action we should perform.
+    val validActions = Seq(opts.generateOpt, opts.executeOpt, opts.verifyOpt,
+                           opts.cancelOpt, opts.listOpt)
+    val allActions = validActions.filter(opts.options.has _)
+    if (allActions.size != 1) {
+      CommandLineUtils.printUsageAndDie(opts.parser, "Command must include exactly one action: %s".format(
+        validActions.map("--" + _.options().get(0)).mkString(", ")))
     }
-    else if (opts.options.has(opts.executeOpt)){
-      if(!opts.options.has(opts.reassignmentJsonFileOpt))
-        CommandLineUtils.printUsageAndDie(opts.parser, "If --execute option is used, command must include --reassignment-json-file that was output " + "during the --generate option")
-      CommandLineUtils.checkInvalidArgs(opts.parser, opts.options, opts.executeOpt, Set(opts.topicsToMoveJsonFileOpt, opts.disableRackAware, opts.brokerListOpt))
+    val action = allActions(0)
+
+    // Check that we have either the --zookeeper option or the --bootstrap-server set.
+    // It would be nice to enforce that we can only have one of these options set at once.  Unfortunately,
+    // previous versions of this tool supported setting both options together.  To avoid breaking backwards
+    // compatibility, we will follow suit, for now.  This issue will eventually be resolved when we remove
+    // the --zookeeper option.
+    if (!opts.options.has(opts.zkConnectOpt) && !opts.options.has(opts.bootstrapServerOpt))
+      CommandLineUtils.printUsageAndDie(opts.parser, "Please specify --bootstrap-server")
+
+    // Make sure that we have all the required arguments for our action.
+    val requiredArgs = Map(
+      opts.verifyOpt -> collection.immutable.Seq(
+        opts.reassignmentJsonFileOpt
+      ),
+      opts.generateOpt -> collection.immutable.Seq(
+        opts.topicsToMoveJsonFileOpt,
+        opts.brokerListOpt
+      ),
+      opts.executeOpt -> collection.immutable.Seq(
+        opts.reassignmentJsonFileOpt
+      ),
+      opts.cancelOpt -> collection.immutable.Seq(
+        opts.reassignmentJsonFileOpt
+      ),
+      opts.listOpt -> collection.immutable.Seq.empty
+    )
+    CommandLineUtils.checkRequiredArgs(opts.parser, opts.options, requiredArgs(action): _*)
+
+    // Make sure that we didn't specify any arguments that are incompatible with our chosen action.
+    val permittedArgs = Map(
+      opts.verifyOpt -> Seq(
+        opts.bootstrapServerOpt,
+        opts.commandConfigOpt,
+        opts.preserveThrottlesOpt,
+        opts.zkConnectOpt
+      ),
+      opts.generateOpt -> Seq(
+        opts.bootstrapServerOpt,
+        opts.brokerListOpt,
+        opts.commandConfigOpt,
+        opts.disableRackAware,
+        opts.zkConnectOpt
+      ),
+      opts.executeOpt -> Seq(
+        opts.additionalOpt,
+        opts.bootstrapServerOpt,
+        opts.commandConfigOpt,
+        opts.interBrokerThrottleOpt,
+        opts.replicaAlterLogDirsThrottleOpt,
+        opts.timeoutOpt,
+        opts.zkConnectOpt
+      ),
+      opts.cancelOpt -> Seq(
+        opts.bootstrapServerOpt,
+        opts.commandConfigOpt,
+        opts.preserveThrottlesOpt,
+        opts.timeoutOpt
+      ),
+      opts.listOpt -> Seq(
+        opts.bootstrapServerOpt,
+        opts.commandConfigOpt
+      )
+    )
+    opts.options.specs.forEach(opt => {
+      if (!opt.equals(action) &&
+          !requiredArgs(action).contains(opt) &&
+          !permittedArgs(action).contains(opt)) {
+        CommandLineUtils.printUsageAndDie(opts.parser,
+          """Option "%s" can't be used with action "%s"""".format(opt, action))
+      }
+    })
+    if (!opts.options.has(opts.bootstrapServerOpt)) {
+      val bootstrapServerOnlyArgs = Seq(
+        opts.additionalOpt,
+        opts.cancelOpt,
+        opts.commandConfigOpt,
+        opts.replicaAlterLogDirsThrottleOpt,
+        opts.listOpt,
+        opts.timeoutOpt
+      )
+      bootstrapServerOnlyArgs.foreach {
+        opt => if (opts.options.has(opt)) {
+          throw new RuntimeException("You must specify --bootstrap-server " +
+            """when using "%s"""".format(opt))
+        }
+      }
     }
     opts
   }
 
-  class ReassignPartitionsCommandOptions(args: Array[String]) {
-    val parser = new OptionParser(false)
+  def alterReplicaLogDirs(adminClient: Admin,
+                          assignment: Map[TopicPartitionReplica, String])
+                          : Set[TopicPartitionReplica] = {
+    adminClient.alterReplicaLogDirs(assignment.asJava).values().asScala.flatMap {
+      case (replica, future) => {
+        try {
+          future.get()
+          Some(replica)
+        } catch {
+          case t: ExecutionException =>
+            t.getCause match {
+              // Ignore ReplicaNotAvailableException.  It is OK if the replica is not
+              // available at this moment.
+              case _: ReplicaNotAvailableException => None
+              case e: Throwable =>
+                throw new AdminCommandFailedException(s"Failed to alter dir for $replica", e)
+            }
+        }
+      }
+    }.toSet
+  }
+
+  sealed class ReassignPartitionsCommandOptions(args: Array[String]) extends CommandDefaultOptions(args)  {
+    // Actions
+    val verifyOpt = parser.accepts("verify", "Verify if the reassignment completed as specified by the --reassignment-json-file option. If there is a throttle engaged for the replicas specified, and the rebalance has completed, the throttle will be removed")
+    val generateOpt = parser.accepts("generate", "Generate a candidate partition reassignment configuration." +
+      " Note that this only generates a candidate assignment, it does not execute it.")
+    val executeOpt = parser.accepts("execute", "Kick off the reassignment as specified by the --reassignment-json-file option.")
+    val cancelOpt = parser.accepts("cancel", "Cancel an active reassignment.")
+    val listOpt = parser.accepts("list", "List all active partition reassignments.")
+
+    // Arguments
     val bootstrapServerOpt = parser.accepts("bootstrap-server", "the server(s) to use for bootstrapping. REQUIRED if " +
-                      "an absolution path of the log directory is specified for any replica in the reassignment json file")
+                      "an absolute path of the log directory is specified for any replica in the reassignment json file, " +
+                      "or if --zookeeper is not given.")
                       .withRequiredArg
                       .describedAs("Server(s) to use for bootstrapping")
                       .ofType(classOf[String])
-    val zkConnectOpt = parser.accepts("zookeeper", "REQUIRED: The connection string for the zookeeper connection in the " +
-                      "form host:port. Multiple URLS can be given to allow fail-over.")
+    val commandConfigOpt = parser.accepts("command-config", "Property file containing configs to be passed to Admin Client.")
+                      .withRequiredArg
+                      .describedAs("Admin client property file")
+                      .ofType(classOf[String])
+    val zkConnectOpt = parser.accepts("zookeeper", "DEPRECATED: The connection string for the zookeeper connection in the " +
+                      "form host:port. Multiple URLS can be given to allow fail-over.  Please use --bootstrap-server instead.")
                       .withRequiredArg
                       .describedAs("urls")
                       .ofType(classOf[String])
-    val generateOpt = parser.accepts("generate", "Generate a candidate partition reassignment configuration." +
-                      " Note that this only generates a candidate assignment, it does not execute it.")
-    val executeOpt = parser.accepts("execute", "Kick off the reassignment as specified by the --reassignment-json-file option.")
-    val verifyOpt = parser.accepts("verify", "Verify if the reassignment completed as specified by the --reassignment-json-file option. If there is a throttle engaged for the replicas specified, and the rebalance has completed, the throttle will be removed")
     val reassignmentJsonFileOpt = parser.accepts("reassignment-json-file", "The JSON file with the partition reassignment configuration" +
                       "The format to use is - \n" +
                       "{\"partitions\":\n\t[{\"topic\": \"foo\",\n\t  \"partition\": 1,\n\t  \"replicas\": [1,2,3],\n\t  \"log_dirs\": [\"dir1\",\"dir2\",\"dir3\"] }],\n\"version\":1\n}\n" +
                       "Note that \"log_dirs\" is optional. When it is specified, its length must equal the length of the replicas list. The value in this list " +
-                      "can be either \"any\" or the absolution path of the log directory on the broker. If absolute log directory path is specified, it is currently required that " +
-                      "the replica has not already been created on that broker. The replica will then be created in the specified log directory on the broker later.")
+                      "can be either \"any\" or the absolution path of the log directory on the broker. If absolute log directory path is specified, the replica will be moved to the specified log directory on the broker.")
                       .withRequiredArg
                       .describedAs("manual assignment json file path")
                       .ofType(classOf[String])
@@ -440,197 +1835,29 @@ object ReassignPartitionsCommand extends Logging {
                       .describedAs("brokerlist")
                       .ofType(classOf[String])
     val disableRackAware = parser.accepts("disable-rack-aware", "Disable rack aware replica assignment")
-    val interBrokerThrottleOpt = parser.accepts("throttle", "The movement of partitions between brokers will be throttled to this value (bytes/sec). Rerunning with this option, whilst a rebalance is in progress, will alter the throttle value. The throttle rate should be at least 1 KB/s.")
-                      .withRequiredArg()
-                      .describedAs("throttle")
-                      .ofType(classOf[Long])
-                      .defaultsTo(-1)
-    val replicaAlterLogDirsThrottleOpt = parser.accepts("replica-alter-log-dirs-throttle", "The movement of replicas between log directories on the same broker will be throttled to this value (bytes/sec). Rerunning with this option, whilst a rebalance is in progress, will alter the throttle value. The throttle rate should be at least 1 KB/s.")
-                      .withRequiredArg()
-                      .describedAs("replicaAlterLogDirsThrottle")
-                      .ofType(classOf[Long])
-                      .defaultsTo(-1)
-    val timeoutOpt = parser.accepts("timeout", "The maximum time in ms allowed to wait for partition reassignment execution to be successfully initiated")
+    val interBrokerThrottleOpt = parser.accepts("throttle", "The movement of partitions between brokers will be throttled to this value (bytes/sec). " +
+      "This option can be included with --execute when a reassignment is started, and it can be altered by resubmitting the current reassignment " +
+      "along with the --additional flag. The throttle rate should be at least 1 KB/s.")
+      .withRequiredArg()
+      .describedAs("throttle")
+      .ofType(classOf[Long])
+      .defaultsTo(-1)
+    val replicaAlterLogDirsThrottleOpt = parser.accepts("replica-alter-log-dirs-throttle",
+      "The movement of replicas between log directories on the same broker will be throttled to this value (bytes/sec). " +
+        "This option can be included with --execute when a reassignment is started, and it can be altered by resubmitting the current reassignment " +
+        "along with the --additional flag. The throttle rate should be at least 1 KB/s.")
+      .withRequiredArg()
+      .describedAs("replicaAlterLogDirsThrottle")
+      .ofType(classOf[Long])
+      .defaultsTo(-1)
+    val timeoutOpt = parser.accepts("timeout", "The maximum time in ms to wait for log directory replica assignment to begin.")
                       .withRequiredArg()
                       .describedAs("timeout")
                       .ofType(classOf[Long])
                       .defaultsTo(10000)
-    val options = parser.parse(args : _*)
+    val additionalOpt = parser.accepts("additional", "Execute this reassignment in addition to any " +
+      "other ongoing ones. This option can also be used to change the throttle of an ongoing reassignment.")
+    val preserveThrottlesOpt = parser.accepts("preserve-throttles", "Do not modify broker or topic throttles.")
+    options = parser.parse(args : _*)
   }
 }
-
-class ReassignPartitionsCommand(zkUtils: ZkUtils,
-                                adminClientOpt: Option[JAdminClient],
-                                proposedPartitionAssignment: Map[TopicAndPartition, Seq[Int]],
-                                proposedReplicaAssignment: Map[TopicPartitionReplica, String] = Map.empty,
-                                admin: AdminUtilities = AdminUtils)
-  extends Logging {
-
-  import ReassignPartitionsCommand._
-
-  def existingAssignment(): Map[TopicAndPartition, Seq[Int]] = {
-    val proposedTopics = proposedPartitionAssignment.keySet.map(_.topic).toSeq
-    zkUtils.getReplicaAssignmentForTopics(proposedTopics)
-  }
-
-  private def maybeThrottle(throttle: Throttle): Unit = {
-    if (throttle.interBrokerLimit >= 0)
-      assignThrottledReplicas(existingAssignment(), proposedPartitionAssignment)
-    maybeLimit(throttle)
-    if (throttle.interBrokerLimit >= 0 || throttle.replicaAlterLogDirsLimit >= 0)
-      throttle.postUpdateAction()
-    if (throttle.interBrokerLimit >= 0)
-      println(s"The inter-broker throttle limit was set to ${throttle.interBrokerLimit} B/s")
-    if (throttle.replicaAlterLogDirsLimit >= 0)
-      println(s"The replica-alter-dir throttle limit was set to ${throttle.replicaAlterLogDirsLimit} B/s")
-  }
-
-  /**
-    * 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) {
-    if (throttle.interBrokerLimit >= 0 || throttle.replicaAlterLogDirsLimit >= 0) {
-      val existingBrokers = existingAssignment().values.flatten.toSeq
-      val proposedBrokers = proposedPartitionAssignment.values.flatten.toSeq ++ proposedReplicaAssignment.keys.toSeq.map(_.brokerId())
-      val brokers = (existingBrokers ++ proposedBrokers).distinct
-
-      for (id <- brokers) {
-        val configs = admin.fetchEntityConfig(zkUtils, ConfigType.Broker, id.toString)
-        if (throttle.interBrokerLimit >= 0) {
-          configs.put(DynamicConfig.Broker.LeaderReplicationThrottledRateProp, throttle.interBrokerLimit.toString)
-          configs.put(DynamicConfig.Broker.FollowerReplicationThrottledRateProp, throttle.interBrokerLimit.toString)
-        }
-        if (throttle.replicaAlterLogDirsLimit >= 0)
-          configs.put(DynamicConfig.Broker.ReplicaAlterLogDirsIoMaxBytesPerSecondProp, throttle.replicaAlterLogDirsLimit.toString)
-
-        admin.changeBrokerConfig(zkUtils, Seq(id), configs)
-      }
-    }
-  }
-
-  /** Set throttles to replicas that are moving. Note: this method should only be used when the assignment is initiated. */
-  private[admin] def assignThrottledReplicas(existingPartitionAssignment: Map[TopicAndPartition, Seq[Int]],
-                                             proposedPartitionAssignment: Map[TopicAndPartition, Seq[Int]],
-                                             admin: AdminUtilities = AdminUtils): Unit = {
-    for (topic <- proposedPartitionAssignment.keySet.map(_.topic).toSeq) {
-      val existingPartitionAssignmentForTopic = existingPartitionAssignment.filter { case (tp, _) => tp.topic == topic }
-      val proposedPartitionAssignmentForTopic = proposedPartitionAssignment.filter { case (tp, _) => tp.topic == topic }
-
-      //Apply leader throttle to all replicas that exist before the re-balance.
-      val leader = format(preRebalanceReplicaForMovingPartitions(existingPartitionAssignmentForTopic, proposedPartitionAssignmentForTopic))
-
-      //Apply follower throttle to all "move destinations".
-      val follower = format(postRebalanceReplicasThatMoved(existingPartitionAssignmentForTopic, proposedPartitionAssignmentForTopic))
-
-      val configs = admin.fetchEntityConfig(zkUtils, ConfigType.Topic, topic)
-      configs.put(LeaderReplicationThrottledReplicasProp, leader)
-      configs.put(FollowerReplicationThrottledReplicasProp, follower)
-      admin.changeTopicConfig(zkUtils, topic, configs)
-
-      debug(s"Updated leader-throttled replicas for topic $topic with: $leader")
-      debug(s"Updated follower-throttled replicas for topic $topic with: $follower")
-    }
-  }
-
-  private def postRebalanceReplicasThatMoved(existing: Map[TopicAndPartition, Seq[Int]], proposed: Map[TopicAndPartition, Seq[Int]]): Map[TopicAndPartition, Seq[Int]] = {
-    //For each partition in the proposed list, filter out any replicas that exist now, and hence aren't being moved.
-    proposed.map { case (tp, proposedReplicas) =>
-      tp -> (proposedReplicas.toSet -- existing(tp)).toSeq
-    }
-  }
-
-  private def preRebalanceReplicaForMovingPartitions(existing: Map[TopicAndPartition, Seq[Int]], proposed: Map[TopicAndPartition, Seq[Int]]): Map[TopicAndPartition, Seq[Int]] = {
-    def moving(before: Seq[Int], after: Seq[Int]) = (after.toSet -- before.toSet).nonEmpty
-    //For any moving partition, throttle all the original (pre move) replicas (as any one might be a leader)
-    existing.filter { case (tp, preMoveReplicas) =>
-      proposed.contains(tp) && moving(preMoveReplicas, proposed(tp))
-    }
-  }
-
-  def format(moves: Map[TopicAndPartition, Seq[Int]]): String =
-    moves.flatMap { case (tp, moves) =>
-      moves.map(replicaId => s"${tp.partition}:${replicaId}")
-    }.mkString(",")
-
-  private def alterReplicaLogDirsIgnoreReplicaNotAvailable(replicaAssignment: Map[TopicPartitionReplica, String],
-                                                           adminClient: JAdminClient,
-                                                           timeoutMs: Long): Set[TopicPartitionReplica] = {
-    val alterReplicaLogDirsResult = adminClient.alterReplicaLogDirs(replicaAssignment.asJava, new AlterReplicaLogDirsOptions().timeoutMs(timeoutMs.toInt))
-    val replicasAssignedToFutureDir = alterReplicaLogDirsResult.values().asScala.flatMap { case (replica, future) => {
-      try {
-        future.get()
-        Some(replica)
-      } catch {
-        case t: ExecutionException =>
-          t.getCause match {
-            case e: ReplicaNotAvailableException => None // It is OK if the replica is not available at this moment
-            case e: Throwable => throw new AdminCommandFailedException(s"Failed to alter dir for $replica", e)
-          }
-      }
-    }}
-    replicasAssignedToFutureDir.toSet
-  }
-
-  def reassignPartitions(throttle: Throttle = NoThrottle, timeoutMs: Long = 10000L): Boolean = {
-    maybeThrottle(throttle)
-    try {
-      val validPartitions = proposedPartitionAssignment.filter { case (p, _) => validatePartition(zkUtils, p.topic, p.partition) }
-      if (validPartitions.isEmpty) false
-      else {
-        if (proposedReplicaAssignment.nonEmpty && adminClientOpt.isEmpty)
-          throw new AdminCommandFailedException("bootstrap-server needs to be provided in order to reassign replica to the specified log directory")
-        val startTimeMs = System.currentTimeMillis()
-
-        // Send AlterReplicaLogDirsRequest to allow broker to create replica in the right log dir later if the replica has not been created yet.
-        if (proposedReplicaAssignment.nonEmpty)
-          alterReplicaLogDirsIgnoreReplicaNotAvailable(proposedReplicaAssignment, adminClientOpt.get, timeoutMs)
-
-        // Create reassignment znode so that controller will send LeaderAndIsrRequest to create replica in the broker
-        val jsonReassignmentData = ZkUtils.formatAsReassignmentJson(validPartitions)
-        zkUtils.createPersistentPath(ZkUtils.ReassignPartitionsPath, jsonReassignmentData)
-
-        // Send AlterReplicaLogDirsRequest again to make sure broker will start to move replica to the specified log directory.
-        // It may take some time for controller to create replica in the broker. Retry if the replica has not been created.
-        var remainingTimeMs = startTimeMs + timeoutMs - System.currentTimeMillis()
-        val replicasAssignedToFutureDir = mutable.Set.empty[TopicPartitionReplica]
-        while (remainingTimeMs > 0 && replicasAssignedToFutureDir.size < proposedReplicaAssignment.size) {
-          replicasAssignedToFutureDir ++= alterReplicaLogDirsIgnoreReplicaNotAvailable(
-            proposedReplicaAssignment.filterKeys(replica => !replicasAssignedToFutureDir.contains(replica)), adminClientOpt.get, remainingTimeMs)
-          Thread.sleep(100)
-          remainingTimeMs = startTimeMs + timeoutMs - System.currentTimeMillis()
-        }
-        replicasAssignedToFutureDir.size == proposedReplicaAssignment.size
-      }
-    } catch {
-      case _: ZkNodeExistsException =>
-        val partitionsBeingReassigned = zkUtils.getPartitionsBeingReassigned()
-        throw new AdminCommandFailedException("Partition reassignment currently in " +
-          "progress for %s. Aborting operation".format(partitionsBeingReassigned))
-    }
-  }
-
-  def validatePartition(zkUtils: ZkUtils, topic: String, partition: Int): Boolean = {
-    // check if partition exists
-    val partitionsOpt = zkUtils.getPartitionsForTopics(List(topic)).get(topic)
-    partitionsOpt match {
-      case Some(partitions) =>
-        if(partitions.contains(partition)) {
-          true
-        } else {
-          error("Skipping reassignment of partition [%s,%d] ".format(topic, partition) +
-            "since it doesn't exist")
-          false
-        }
-      case None => error("Skipping reassignment of partition " +
-        "[%s,%d] since topic %s doesn't exist".format(topic, partition, topic))
-        false
-    }
-  }
-}
-
-sealed trait ReassignmentStatus { def status: Int }
-case object ReassignmentCompleted extends ReassignmentStatus { val status = 1 }
-case object ReassignmentInProgress extends ReassignmentStatus { val status = 0 }
-case object ReassignmentFailed extends ReassignmentStatus { val status = -1 }
-
diff --git a/core/src/main/scala/kafka/admin/TopicCommand.scala b/core/src/main/scala/kafka/admin/TopicCommand.scala
index bdd8aaf177304..95129e9b62032 100755
--- a/core/src/main/scala/kafka/admin/TopicCommand.scala
+++ b/core/src/main/scala/kafka/admin/TopicCommand.scala
@@ -17,247 +17,539 @@
 
 package kafka.admin
 
-import java.util.Properties
+import java.util
+import java.util.{Collections, Properties}
 
 import joptsimple._
 import kafka.common.AdminCommandFailedException
-import kafka.utils.Implicits._
-import kafka.consumer.Whitelist
 import kafka.log.LogConfig
 import kafka.server.ConfigType
+import kafka.utils.Implicits._
 import kafka.utils._
 import kafka.zk.{AdminZkClient, KafkaZkClient}
-import kafka.zookeeper.ZooKeeperClient
-import org.apache.kafka.common.errors.{InvalidTopicException, TopicExistsException}
+import org.apache.kafka.clients.CommonClientConfigs
+import org.apache.kafka.clients.admin.CreatePartitionsOptions
+import org.apache.kafka.clients.admin.CreateTopicsOptions
+import org.apache.kafka.clients.admin.DeleteTopicsOptions
+import org.apache.kafka.clients.admin.{Admin, ConfigEntry, ListTopicsOptions, NewPartitions, NewTopic, PartitionReassignment, 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.{ClusterAuthorizationException, InvalidTopicException, TopicExistsException, UnsupportedVersionException}
 import org.apache.kafka.common.internals.Topic
 import org.apache.kafka.common.security.JaasUtils
-import org.apache.kafka.common.utils.Utils
+import org.apache.kafka.common.utils.{Time, Utils}
 import org.apache.zookeeper.KeeperException.NodeExistsException
-import org.apache.kafka.common.TopicPartition
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 import scala.collection._
+import scala.compat.java8.OptionConverters._
+import scala.concurrent.ExecutionException
 
 object TopicCommand extends Logging {
 
   def main(args: Array[String]): Unit = {
-
     val opts = new TopicCommandOptions(args)
-
-    if(args.length == 0)
-      CommandLineUtils.printUsageAndDie(opts.parser, "Create, delete, describe, or change a topic.")
-
-    // should have exactly one action
-    val actions = Seq(opts.createOpt, opts.listOpt, opts.alterOpt, opts.describeOpt, opts.deleteOpt).count(opts.options.has _)
-    if(actions != 1)
-      CommandLineUtils.printUsageAndDie(opts.parser, "Command must include exactly one action: --list, --describe, --create, --alter or --delete")
-
     opts.checkArgs()
 
-    val zooKeeperClient = new ZooKeeperClient(opts.options.valueOf(opts.zkConnectOpt), 30000, 30000, Int.MaxValue)
-    val zkClient = new KafkaZkClient(zooKeeperClient, JaasUtils.isZkSecurityEnabled())
+    val topicService = if (opts.zkConnect.isDefined)
+      ZookeeperTopicService(opts.zkConnect)
+    else
+      AdminClientTopicService(opts.commandConfig, opts.bootstrapServer)
 
     var exitCode = 0
     try {
-      if(opts.options.has(opts.createOpt))
-        createTopic(zkClient, opts)
-      else if(opts.options.has(opts.alterOpt))
-        alterTopic(zkClient, opts)
-      else if(opts.options.has(opts.listOpt))
-        listTopics(zkClient, opts)
-      else if(opts.options.has(opts.describeOpt))
-        describeTopic(zkClient, opts)
-      else if(opts.options.has(opts.deleteOpt))
-        deleteTopic(zkClient, opts)
+      if (opts.hasCreateOption)
+        topicService.createTopic(opts)
+      else if (opts.hasAlterOption)
+        topicService.alterTopic(opts)
+      else if (opts.hasListOption)
+        topicService.listTopics(opts)
+      else if (opts.hasDescribeOption)
+        topicService.describeTopic(opts)
+      else if (opts.hasDeleteOption)
+        topicService.deleteTopic(opts)
     } catch {
+      case e: ExecutionException =>
+        if (e.getCause != null)
+          printException(e.getCause)
+        else
+          printException(e)
+        exitCode = 1
       case e: Throwable =>
-        println("Error while executing topic command : " + e.getMessage)
-        error(Utils.stackTrace(e))
+        printException(e)
         exitCode = 1
     } finally {
-      zkClient.close()
+      topicService.close()
       Exit.exit(exitCode)
     }
+  }
 
+  private def printException(e: Throwable): Unit = {
+    println("Error while executing topic command : " + e.getMessage)
+    error(Utils.stackTrace(e))
   }
 
-  private def getTopics(zkClient: KafkaZkClient, opts: TopicCommandOptions): Seq[String] = {
-    val allTopics = zkClient.getAllTopicsInCluster.sorted
-    if (opts.options.has(opts.topicOpt)) {
-      val topicsSpec = opts.options.valueOf(opts.topicOpt)
-      val topicsFilter = new Whitelist(topicsSpec)
-      allTopics.filter(topicsFilter.isTopicAllowed(_, excludeInternalTopics = false))
-    } else
-      allTopics
+  class CommandTopicPartition(opts: TopicCommandOptions) {
+    val name: String = opts.topic.get
+    val partitions: Option[Integer] = opts.partitions
+    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)
+    val rackAwareMode: RackAwareMode = opts.rackAwareMode
+
+    def hasReplicaAssignment: Boolean = replicaAssignment.isDefined
+    def hasPartitions: Boolean = partitions.isDefined
+    def ifTopicDoesntExist(): Boolean = opts.ifNotExists
   }
 
-  def createTopic(zkClient: KafkaZkClient, opts: TopicCommandOptions) {
-    val topic = opts.options.valueOf(opts.topicOpt)
-    val configs = parseTopicConfigsToBeAdded(opts)
-    val ifNotExists = opts.options.has(opts.ifNotExistsOpt)
-    if (Topic.hasCollisionChars(topic))
-      println("WARNING: Due to limitations in metric names, topics with a period ('.') or underscore ('_') could collide. To avoid issues it is best to use either, but not both.")
-    val adminZkClient = new AdminZkClient(zkClient)
-    try {
-      if (opts.options.has(opts.replicaAssignmentOpt)) {
-        val assignment = parseReplicaAssignment(opts.options.valueOf(opts.replicaAssignmentOpt))
-        adminZkClient.createOrUpdateTopicPartitionAssignmentPathInZK(topic, assignment, configs, update = false)
-      } else {
-        CommandLineUtils.checkRequiredArgs(opts.parser, opts.options, opts.partitionsOpt, opts.replicationFactorOpt)
-        val partitions = opts.options.valueOf(opts.partitionsOpt).intValue
-        val replicas = opts.options.valueOf(opts.replicationFactorOpt).intValue
-        val rackAwareMode = if (opts.options.has(opts.disableRackAware)) RackAwareMode.Disabled
-                            else RackAwareMode.Enforced
-        adminZkClient.createTopic(topic, partitions, replicas, configs, rackAwareMode)
+  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()
+    }
+  }
+
+  case class PartitionDescription(topic: String,
+                                  info: TopicPartitionInfo,
+                                  config: Option[JConfig],
+                                  markedForDeletion: Boolean,
+                                  reassignment: Option[PartitionReassignment]) {
+
+    private def minIsrCount: Option[Int] = {
+      config.map(_.get(TopicConfig.MIN_IN_SYNC_REPLICAS_CONFIG).value.toInt)
+    }
+
+    def isUnderReplicated: Boolean = {
+      getReplicationFactor(info, reassignment) - info.isr.size > 0
+    }
+
+    private def hasLeader: Boolean = {
+      info.leader != null
+    }
+
+    def isUnderMinIsr: Boolean = {
+      !hasLeader || minIsrCount.exists(info.isr.size < _)
+    }
+
+    def isAtMinIsrPartitions: Boolean =  {
+      minIsrCount.contains(info.isr.size)
+    }
+
+    def hasUnavailablePartitions(liveBrokers: Set[Int]): Boolean = {
+      !hasLeader || !liveBrokers.contains(info.leader.id)
+    }
+
+    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(","))
+      if (reassignment.nonEmpty) {
+        print("\tAdding Replicas: " + reassignment.get.addingReplicas().asScala.mkString(","))
+        print("\tRemoving Replicas: " + reassignment.get.removingReplicas().asScala.mkString(","))
       }
-      println("Created topic \"%s\".".format(topic))
-    } catch  {
-      case e: TopicExistsException => if (!ifNotExists) throw e
+      print(if (markedForDeletion) "\tMarkedForDeletion: true" else "")
+      println()
+    }
+
+  }
+
+  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.isUnderReplicated
+    }
+    private def shouldPrintUnavailablePartitions(partitionDescription: PartitionDescription): Boolean = {
+      opts.reportUnavailablePartitions && partitionDescription.hasUnavailablePartitions(liveBrokers)
     }
+    private def shouldPrintUnderMinIsrPartitions(partitionDescription: PartitionDescription): Boolean = {
+      opts.reportUnderMinIsrPartitions && partitionDescription.isUnderMinIsr
+    }
+    private def shouldPrintAtMinIsrPartitions(partitionDescription: PartitionDescription): Boolean = {
+      opts.reportAtMinIsrPartitions && partitionDescription.isAtMinIsrPartitions
+    }
+
+    private def shouldPrintTopicPartition(partitionDesc: PartitionDescription): Boolean = {
+      describeConfigs ||
+        shouldPrintUnderReplicatedPartitions(partitionDesc) ||
+        shouldPrintUnavailablePartitions(partitionDesc) ||
+        shouldPrintUnderMinIsrPartitions(partitionDesc) ||
+        shouldPrintAtMinIsrPartitions(partitionDesc)
+    }
+
+    def maybePrintPartitionDescription(desc: PartitionDescription): Unit = {
+      if (shouldPrintTopicPartition(desc))
+        desc.printDescription()
+    }
+  }
+
+  trait TopicService extends AutoCloseable {
+    def createTopic(opts: TopicCommandOptions): Unit = {
+      val topic = new CommandTopicPartition(opts)
+      if (Topic.hasCollisionChars(topic.name))
+        println("WARNING: Due to limitations in metric names, topics with a period ('.') or underscore ('_') could " +
+          "collide. To avoid issues it is best to use either, but not both.")
+      createTopic(topic)
+    }
+    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(topicIncludelist: Option[String], excludeInternalTopics: Boolean = false): Seq[String]
   }
 
-  def alterTopic(zkClient: KafkaZkClient, opts: TopicCommandOptions) {
-    val topics = getTopics(zkClient, opts)
-    val ifExists = opts.options.has(opts.ifExistsOpt)
-    if (topics.isEmpty && !ifExists) {
-      throw new IllegalArgumentException("Topic %s does not exist on ZK path %s".format(opts.options.valueOf(opts.topicOpt),
-          opts.options.valueOf(opts.zkConnectOpt)))
-    }
-    val adminZkClient = new AdminZkClient(zkClient)
-    topics.foreach { topic =>
-      val configs = adminZkClient.fetchEntityConfig(ConfigType.Topic, topic)
-      if(opts.options.has(opts.configOpt) || opts.options.has(opts.deleteConfigOpt)) {
-        println("WARNING: Altering topic configuration from this script has been deprecated and may be removed in future releases.")
-        println("         Going forward, please use kafka-configs.sh for this functionality")
-
-        val configsToBeAdded = parseTopicConfigsToBeAdded(opts)
-        val configsToBeDeleted = parseTopicConfigsToBeDeleted(opts)
-        // compile the final set of configs
-        configs ++= configsToBeAdded
-        configsToBeDeleted.foreach(config => configs.remove(config))
-        adminZkClient.changeTopicConfig(topic, configs)
-        println("Updated config for topic \"%s\".".format(topic))
+  object AdminClientTopicService {
+    def createAdminClient(commandConfig: Properties, bootstrapServer: Option[String]): Admin = {
+      bootstrapServer match {
+        case Some(serverList) => commandConfig.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, serverList)
+        case None =>
       }
+      Admin.create(commandConfig)
+    }
 
-      if(opts.options.has(opts.partitionsOpt)) {
-        if (topic == Topic.GROUP_METADATA_TOPIC_NAME) {
-          throw new IllegalArgumentException("The number of partitions for the offsets 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 nPartitions = opts.options.valueOf(opts.partitionsOpt).intValue
-        val existingAssignment = zkClient.getReplicaAssignmentForTopics(immutable.Set(topic)).map {
-          case (topicPartition, replicas) => topicPartition.partition -> replicas
+    def apply(commandConfig: Properties, bootstrapServer: Option[String]): AdminClientTopicService =
+      new AdminClientTopicService(createAdminClient(commandConfig, bootstrapServer))
+  }
+
+  case class AdminClientTopicService private (adminClient: Admin) extends TopicService {
+
+    override def createTopic(topic: CommandTopicPartition): Unit = {
+      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")
+
+      try {
+        val newTopic = if (topic.hasReplicaAssignment)
+          new NewTopic(topic.name, asJavaReplicaReassignment(topic.replicaAssignment.get))
+        else {
+          new NewTopic(
+            topic.name,
+            topic.partitions.asJava,
+            topic.replicationFactor.map(_.toShort).map(Short.box).asJava)
         }
-        if (existingAssignment.isEmpty)
-          throw new InvalidTopicException(s"The topic $topic does not exist")
-        val replicaAssignmentStr = opts.options.valueOf(opts.replicaAssignmentOpt)
-        val newAssignment = Option(replicaAssignmentStr).filter(_.nonEmpty).map { replicaAssignmentString =>
-          val startPartitionId = existingAssignment.size
-          val partitionList = replicaAssignmentString.split(",").drop(startPartitionId)
-          AdminUtils.parseReplicaAssignment(partitionList.mkString(","), startPartitionId)
+
+        val configsMap = topic.configsToAdd.stringPropertyNames()
+          .asScala
+          .map(name => name -> topic.configsToAdd.getProperty(name))
+          .toMap.asJava
+
+        newTopic.configs(configsMap)
+        val createResult = adminClient.createTopics(Collections.singleton(newTopic),
+          new CreateTopicsOptions().retryOnQuotaViolation(false))
+        createResult.all().get()
+        println(s"Created topic ${topic.name}.")
+      } catch {
+        case e : ExecutionException =>
+          if (e.getCause == null)
+            throw e
+          if (!(e.getCause.isInstanceOf[TopicExistsException] && topic.ifTopicDoesntExist()))
+            throw e.getCause
+      }
+    }
+
+    override def listTopics(opts: TopicCommandOptions): Unit = {
+      println(getTopics(opts.topic, opts.excludeInternalTopics).mkString("\n"))
+    }
+
+    override def alterTopic(opts: TopicCommandOptions): Unit = {
+      val topic = new CommandTopicPartition(opts)
+      val topics = getTopics(opts.topic, opts.excludeInternalTopics)
+      ensureTopicExists(topics, opts.topic, !opts.ifExists)
+
+      if (topics.nonEmpty) {
+        val topicsInfo = adminClient.describeTopics(topics.asJavaCollection).values()
+        val newPartitions = topics.map { topicName =>
+          if (topic.hasReplicaAssignment) {
+            val startPartitionId = topicsInfo.get(topicName).get().partitions().size()
+            val newAssignment = {
+              val replicaMap = topic.replicaAssignment.get.drop(startPartitionId)
+              new util.ArrayList(replicaMap.map(p => p._2.asJava).asJavaCollection).asInstanceOf[util.List[util.List[Integer]]]
+            }
+            topicName -> NewPartitions.increaseTo(topic.partitions.get, newAssignment)
+          } else {
+            topicName -> NewPartitions.increaseTo(topic.partitions.get)
+          }
+        }.toMap
+        adminClient.createPartitions(newPartitions.asJava,
+          new CreatePartitionsOptions().retryOnQuotaViolation(false)).all().get()
+      }
+    }
+
+    private def listAllReassignments(topicPartitions: util.Set[TopicPartition]): Map[TopicPartition, PartitionReassignment] = {
+      try {
+        adminClient.listPartitionReassignments(topicPartitions).reassignments().get().asScala
+      } catch {
+        case e: ExecutionException =>
+          e.getCause match {
+            case ex @ (_: UnsupportedVersionException | _: ClusterAuthorizationException) =>
+              logger.debug(s"Couldn't query reassignments through the AdminClient API: ${ex.getMessage}", ex)
+              Map()
+            case t => throw t
+          }
+      }
+    }
+
+    override def describeTopic(opts: TopicCommandOptions): Unit = {
+      val topics = getTopics(opts.topic, opts.excludeInternalTopics)
+      ensureTopicExists(topics, opts.topic, !opts.ifExists)
+
+      if (topics.nonEmpty) {
+        val allConfigs = adminClient.describeConfigs(topics.map(new ConfigResource(Type.TOPIC, _)).asJavaCollection).values()
+        val liveBrokers = adminClient.describeCluster().nodes().get().asScala.map(_.id())
+        val topicDescriptions = adminClient.describeTopics(topics.asJavaCollection).all().get().values().asScala
+        val describeOptions = new DescribeOptions(opts, liveBrokers.toSet)
+        val topicPartitions = topicDescriptions
+          .flatMap(td => td.partitions.iterator().asScala.map(p => new TopicPartition(td.name(), p.partition())))
+          .toSet.asJava
+        val reassignments = listAllReassignments(topicPartitions)
+
+        for (td <- topicDescriptions) {
+          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 hasNonDefault = config.entries().asScala.exists(!_.isDefault)
+            if (!opts.reportOverriddenConfigs || hasNonDefault) {
+              val numPartitions = td.partitions().size
+              val firstPartition = td.partitions.iterator.next()
+              val reassignment = reassignments.get(new TopicPartition(td.name, firstPartition.partition))
+              val topicDesc = TopicDescription(topicName, numPartitions, getReplicationFactor(firstPartition, reassignment), config, markedForDeletion = false)
+              topicDesc.printDescription()
+            }
+          }
+
+          if (describeOptions.describePartitions) {
+            for (partition <- sortedPartitions) {
+              val reassignment = reassignments.get(new TopicPartition(td.name, partition.partition))
+              val partitionDesc = PartitionDescription(topicName, partition, Some(config), markedForDeletion = false, reassignment)
+              describeOptions.maybePrintPartitionDescription(partitionDesc)
+            }
+          }
         }
-        val allBrokers = adminZkClient.getBrokerMetadatas()
-        adminZkClient.addPartitions(topic, existingAssignment, allBrokers, nPartitions, newAssignment)
-        println("Adding partitions succeeded!")
       }
     }
-  }
 
-  def listTopics(zkClient: KafkaZkClient, opts: TopicCommandOptions) {
-    val topics = getTopics(zkClient, opts)
-    for(topic <- topics) {
-      if (zkClient.isTopicMarkedForDeletion(topic)) {
-        println("%s - marked for deletion".format(topic))
+    override def deleteTopic(opts: TopicCommandOptions): Unit = {
+      val topics = getTopics(opts.topic, opts.excludeInternalTopics)
+      ensureTopicExists(topics, opts.topic, !opts.ifExists)
+      adminClient.deleteTopics(topics.asJavaCollection, new DeleteTopicsOptions().retryOnQuotaViolation(false))
+        .all().get()
+    }
+
+    override def getTopics(topicIncludelist: Option[String], excludeInternalTopics: Boolean = false): Seq[String] = {
+      val allTopics = if (excludeInternalTopics) {
+        adminClient.listTopics()
       } else {
-        println(topic)
+        adminClient.listTopics(new ListTopicsOptions().listInternal(true))
       }
+      doGetTopics(allTopics.names().get().asScala.toSeq.sorted, topicIncludelist, excludeInternalTopics)
     }
+
+    override def close(): Unit = adminClient.close()
   }
 
-  def deleteTopic(zkClient: KafkaZkClient, opts: TopicCommandOptions) {
-    val topics = getTopics(zkClient, opts)
-    val ifExists = opts.options.has(opts.ifExistsOpt)
-    if (topics.isEmpty && !ifExists) {
-      throw new IllegalArgumentException("Topic %s does not exist on ZK path %s".format(opts.options.valueOf(opts.topicOpt),
-          opts.options.valueOf(opts.zkConnectOpt)))
-    }
-    topics.foreach { topic =>
+  object ZookeeperTopicService {
+    def apply(zkConnect: Option[String]): ZookeeperTopicService =
+      new ZookeeperTopicService(KafkaZkClient(zkConnect.get, JaasUtils.isZkSaslEnabled, 30000, 30000,
+        Int.MaxValue, Time.SYSTEM))
+  }
+
+  case class ZookeeperTopicService(zkClient: KafkaZkClient) extends TopicService {
+
+    override def createTopic(topic: CommandTopicPartition): Unit = {
+      val adminZkClient = new AdminZkClient(zkClient)
       try {
-        if (Topic.isInternal(topic)) {
-          throw new AdminOperationException("Topic %s is a kafka internal topic and is not allowed to be marked for deletion.".format(topic))
-        } else {
-          zkClient.createDeleteTopicPath(topic)
-          println("Topic %s is marked for deletion.".format(topic))
-          println("Note: This will have no impact if delete.topic.enable is not set to true.")
+        if (topic.hasReplicaAssignment)
+          adminZkClient.createTopicWithAssignment(topic.name, topic.configsToAdd, topic.replicaAssignment.get)
+        else
+          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
+      }
+    }
+
+    override def listTopics(opts: TopicCommandOptions): Unit = {
+      val topics = getTopics(opts.topic, opts.excludeInternalTopics)
+      for(topic <- topics) {
+        if (zkClient.isTopicMarkedForDeletion(topic))
+          println(s"$topic - marked for deletion")
+        else
+          println(topic)
+      }
+    }
+
+    override def alterTopic(opts: TopicCommandOptions): Unit = {
+      val topics = getTopics(opts.topic, opts.excludeInternalTopics)
+      val tp = new CommandTopicPartition(opts)
+      ensureTopicExists(topics, opts.topic, !opts.ifExists)
+      val adminZkClient = new AdminZkClient(zkClient)
+      topics.foreach { topic =>
+        val configs = adminZkClient.fetchEntityConfig(ConfigType.Topic, topic)
+        if(opts.topicConfig.isDefined || opts.configsToDelete.isDefined) {
+          println("WARNING: Altering topic configuration from this script has been deprecated and may be removed in future releases.")
+          println("         Going forward, please use kafka-configs.sh for this functionality")
+
+          // compile the final set of configs
+          configs ++= tp.configsToAdd
+          tp.configsToDelete.foreach(config => configs.remove(config))
+          adminZkClient.changeTopicConfig(topic, configs)
+          println(s"Updated config for topic $topic.")
+        }
+
+        if(tp.hasPartitions) {
+          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.getFullReplicaAssignmentForTopics(immutable.Set(topic)).map {
+            case (topicPartition, assignment) => topicPartition.partition -> assignment
+          }
+          if (existingAssignment.isEmpty)
+            throw new InvalidTopicException(s"The topic $topic does not exist")
+          val newAssignment = tp.replicaAssignment.getOrElse(Map()).drop(existingAssignment.size)
+          val allBrokers = adminZkClient.getBrokerMetadatas()
+          val partitions: Integer = tp.partitions.getOrElse(1)
+          adminZkClient.addPartitions(topic, existingAssignment, allBrokers, partitions, Option(newAssignment).filter(_.nonEmpty))
+          println("Adding partitions succeeded!")
         }
-      } catch {
-        case _: NodeExistsException =>
-          println("Topic %s is already marked for deletion.".format(topic))
-        case e: AdminOperationException =>
-          throw e
-        case _: Throwable =>
-          throw new AdminOperationException("Error while deleting topic %s".format(topic))
       }
     }
-  }
 
-  def describeTopic(zkClient: KafkaZkClient, opts: TopicCommandOptions) {
-    val topics = getTopics(zkClient, opts)
-    val reportUnderReplicatedPartitions = opts.options.has(opts.reportUnderReplicatedPartitionsOpt)
-    val reportUnavailablePartitions = opts.options.has(opts.reportUnavailablePartitionsOpt)
-    val reportOverriddenConfigs = opts.options.has(opts.topicsWithOverridesOpt)
-    val liveBrokers = zkClient.getAllBrokersInCluster.map(_.id).toSet
-    val adminZkClient = new AdminZkClient(zkClient)
-
-    for (topic <- topics) {
-       zkClient.getPartitionAssignmentForTopics(immutable.Set(topic)).get(topic) match {
-        case Some(topicPartitionAssignment) =>
-          val describeConfigs: Boolean = !reportUnavailablePartitions && !reportUnderReplicatedPartitions
-          val describePartitions: Boolean = !reportOverriddenConfigs
-          val sortedPartitions = topicPartitionAssignment.toSeq.sortBy(_._1)
-          val markedForDeletion = zkClient.isTopicMarkedForDeletion(topic)
-          if (describeConfigs) {
-            val configs = adminZkClient.fetchEntityConfig(ConfigType.Topic, topic).asScala
-            if (!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("Topic:%s\tPartitionCount:%d\tReplicationFactor:%d\tConfigs:%s%s"
-                .format(topic, numPartitions, replicationFactor, configsAsString, markedForDeletionString))
+    override def describeTopic(opts: TopicCommandOptions): Unit = {
+      val topics = getTopics(opts.topic, opts.excludeInternalTopics)
+      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) {
+        zkClient.getPartitionAssignmentForTopics(immutable.Set(topic)).get(topic) match {
+          case Some(topicPartitionAssignment) =>
+            val markedForDeletion = zkClient.isTopicMarkedForDeletion(topic)
+            if (describeOptions.describeConfigs) {
+              val configs = adminZkClient.fetchEntityConfig(ConfigType.Topic, topic).asScala
+              if (!opts.reportOverriddenConfigs || configs.nonEmpty) {
+                val numPartitions = topicPartitionAssignment.size
+                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 (describePartitions) {
-            for ((partitionId, assignedReplicas) <- sortedPartitions) {
-              val leaderIsrEpoch = zkClient.getTopicPartitionState(new TopicPartition(topic, partitionId))
-              val inSyncReplicas = if (leaderIsrEpoch.isEmpty) Seq.empty[Int] else leaderIsrEpoch.get.leaderAndIsr.isr
-              val leader = if (leaderIsrEpoch.isEmpty) None else Option(leaderIsrEpoch.get.leaderAndIsr.leader)
-
-              if ((!reportUnderReplicatedPartitions && !reportUnavailablePartitions) ||
-                  (reportUnderReplicatedPartitions && inSyncReplicas.size < assignedReplicas.size) ||
-                  (reportUnavailablePartitions && (leader.isEmpty || !liveBrokers.contains(leader.get)))) {
-
-                val markedForDeletionString =
-                  if (markedForDeletion && !describeConfigs) "\tMarkedForDeletion: true" else ""
-                print("\tTopic: " + topic)
-                print("\tPartition: " + partitionId)
-                print("\tLeader: " + (if(leader.isDefined) leader.get else "none"))
-                print("\tReplicas: " + assignedReplicas.mkString(","))
-                print("\tIsr: " + inSyncReplicas.mkString(","))
-                print(markedForDeletionString)
-                println()
+            if (describeOptions.describePartitions) {
+              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, reassignment = None)
+                describeOptions.maybePrintPartitionDescription(partitionDesc)
               }
             }
+          case None =>
+            println("Topic " + topic + " doesn't exist!")
+        }
+      }
+    }
+
+    override def deleteTopic(opts: TopicCommandOptions): Unit = {
+      val topics = getTopics(opts.topic, opts.excludeInternalTopics)
+      ensureTopicExists(topics, opts.topic, !opts.ifExists)
+      topics.foreach { topic =>
+        try {
+          if (Topic.isInternal(topic)) {
+            throw new AdminOperationException(s"Topic $topic is a kafka internal topic and is not allowed to be marked for deletion.")
+          } else {
+            zkClient.createDeleteTopicPath(topic)
+            println(s"Topic $topic is marked for deletion.")
+            println("Note: This will have no impact if delete.topic.enable is not set to true.")
           }
-        case None =>
-          println("Topic " + topic + " doesn't exist!")
+        } catch {
+          case _: NodeExistsException =>
+            println(s"Topic $topic is already marked for deletion.")
+          case e: AdminOperationException =>
+            throw e
+          case e: Throwable =>
+            throw new AdminOperationException(s"Error while deleting topic $topic", e)
+        }
       }
     }
+
+    override def getTopics(topicIncludelist: Option[String], excludeInternalTopics: Boolean = false): Seq[String] = {
+      val allTopics = zkClient.getAllTopicsInCluster().toSeq.sorted
+      doGetTopics(allTopics, topicIncludelist, excludeInternalTopics)
+    }
+
+    override def close(): Unit = zkClient.close()
   }
 
+  /**
+    * ensures topic existence and throws exception if topic doesn't exist
+    *
+    * @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(foundTopics: Seq[String], requestedTopic: Option[String], requireTopicExists: Boolean): Unit = {
+    // 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"Topic '${requestedTopic.get}' does not exist as expected")
+    }
+  }
+
+  private def doGetTopics(allTopics: Seq[String], topicIncludeList: Option[String], excludeInternalTopics: Boolean): Seq[String] = {
+    if (topicIncludeList.isDefined) {
+      val topicsFilter = IncludeList(topicIncludeList.get)
+      allTopics.filter(topicsFilter.isTopicAllowed(_, excludeInternalTopics))
+    } else
+    allTopics.filterNot(Topic.isInternal(_) && excludeInternalTopics)
+  }
+
+
   def parseTopicConfigsToBeAdded(opts: TopicCommandOptions): Properties = {
-    val configsToBeAdded = opts.options.valuesOf(opts.configOpt).asScala.map(_.split("""\s*=\s*"""))
+    val configsToBeAdded = opts.topicConfig.getOrElse(Collections.emptyList()).asScala.map(_.split("""\s*=\s*"""))
     require(configsToBeAdded.forall(config => config.length == 2),
       "Invalid topic config: all configs to be added must be in the format \"key=val\".")
     val props = new Properties
@@ -265,132 +557,223 @@ object TopicCommand extends Logging {
     LogConfig.validate(props)
     if (props.containsKey(LogConfig.MessageFormatVersionProp)) {
       println(s"WARNING: The configuration ${LogConfig.MessageFormatVersionProp}=${props.getProperty(LogConfig.MessageFormatVersionProp)} is specified. " +
-      s"This configuration will be ignored if the version is newer than the inter.broker.protocol.version specified in the broker.")
+        s"This configuration will be ignored if the version is newer than the inter.broker.protocol.version specified in the broker.")
     }
     props
   }
 
   def parseTopicConfigsToBeDeleted(opts: TopicCommandOptions): Seq[String] = {
-    if (opts.options.has(opts.deleteConfigOpt)) {
-      val configsToBeDeleted = opts.options.valuesOf(opts.deleteConfigOpt).asScala.map(_.trim())
-      val propsToBeDeleted = new Properties
-      configsToBeDeleted.foreach(propsToBeDeleted.setProperty(_, ""))
-      LogConfig.validateNames(propsToBeDeleted)
-      configsToBeDeleted
-    }
-    else
-      Seq.empty
+    val configsToBeDeleted = opts.configsToDelete.getOrElse(Collections.emptyList()).asScala.map(_.trim())
+    val propsToBeDeleted = new Properties
+    configsToBeDeleted.foreach(propsToBeDeleted.setProperty(_, ""))
+    LogConfig.validateNames(propsToBeDeleted)
+    configsToBeDeleted
   }
 
   def parseReplicaAssignment(replicaAssignmentList: String): Map[Int, List[Int]] = {
     val partitionList = replicaAssignmentList.split(",")
-    val ret = new mutable.HashMap[Int, List[Int]]()
+    val ret = new mutable.LinkedHashMap[Int, List[Int]]()
     for (i <- 0 until partitionList.size) {
       val brokerList = partitionList(i).split(":").map(s => s.trim().toInt)
       val duplicateBrokers = CoreUtils.duplicates(brokerList)
       if (duplicateBrokers.nonEmpty)
-        throw new AdminCommandFailedException("Partition replica lists may not contain duplicate entries: %s".format(duplicateBrokers.mkString(",")))
+        throw new AdminCommandFailedException(s"Partition replica lists may not contain duplicate entries: ${duplicateBrokers.mkString(",")}")
       ret.put(i, brokerList.toList)
       if (ret(i).size != ret(0).size)
         throw new AdminOperationException("Partition " + i + " has different replication factor: " + brokerList)
     }
-    ret.toMap
+    ret
   }
 
-  class TopicCommandOptions(args: Array[String]) {
-    val parser = new OptionParser(false)
-    val zkConnectOpt = parser.accepts("zookeeper", "REQUIRED: The connection string for the zookeeper connection in the form host:port. " +
-                                      "Multiple URLS can be given to allow fail-over.")
-                           .withRequiredArg
-                           .describedAs("urls")
-                           .ofType(classOf[String])
-    val listOpt = parser.accepts("list", "List all available topics.")
-    val createOpt = parser.accepts("create", "Create a new topic.")
-    val deleteOpt = parser.accepts("delete", "Delete a topic")
-    val alterOpt = parser.accepts("alter", "Alter the number of partitions, replica assignment, and/or configuration for the topic.")
-    val describeOpt = parser.accepts("describe", "List details for the given topics.")
-    val helpOpt = parser.accepts("help", "Print usage information.")
-    val topicOpt = parser.accepts("topic", "The topic to be create, alter or describe. Can also accept a regular " +
-                                           "expression except for --create option")
+  def asJavaReplicaReassignment(original: Map[Int, List[Int]]): util.Map[Integer, util.List[Integer]] = {
+    original.map(f => Integer.valueOf(f._1) -> f._2.map(e => Integer.valueOf(e)).asJava).asJava
+  }
+
+  private def getReplicationFactor(tpi: TopicPartitionInfo, reassignment: Option[PartitionReassignment]): Int = {
+    // It is possible for a reassignment to complete between the time we have fetched its state and the time
+    // we fetch partition metadata. In ths case, we ignore the reassignment when determining replication factor.
+    def isReassignmentInProgress(ra: PartitionReassignment): Boolean = {
+      // Reassignment is still in progress as long as the removing and adding replicas are still present
+      val allReplicaIds = tpi.replicas.asScala.map(_.id).toSet
+      val changingReplicaIds = ra.removingReplicas.asScala.map(_.intValue).toSet ++ ra.addingReplicas.asScala.map(_.intValue).toSet
+      allReplicaIds.exists(changingReplicaIds.contains)
+    }
+
+    reassignment match {
+      case Some(ra) if isReassignmentInProgress(ra) => ra.replicas.asScala.diff(ra.addingReplicas.asScala).size
+      case _=> tpi.replicas.size
+    }
+  }
+
+  class TopicCommandOptions(args: Array[String]) extends CommandDefaultOptions(args) {
+    private val bootstrapServerOpt = parser.accepts("bootstrap-server", "REQUIRED: The Kafka server to connect to. In case of providing this, a direct Zookeeper connection won't be required.")
+      .withRequiredArg
+      .describedAs("server to connect to")
+      .ofType(classOf[String])
+    private val commandConfigOpt = parser.accepts("command-config", "Property file containing configs to be passed to Admin Client. " +
+      "This is used only with --bootstrap-server option for describing and altering broker configs.")
+      .withRequiredArg
+      .describedAs("command config property file")
+      .ofType(classOf[String])
+    private val zkConnectOpt = parser.accepts("zookeeper", "DEPRECATED, The connection string for the zookeeper connection in the form host:port. " +
+      "Multiple hosts can be given to allow fail-over.")
+      .withRequiredArg
+      .describedAs("hosts")
+      .ofType(classOf[String])
+    private val listOpt = parser.accepts("list", "List all available topics.")
+    private val createOpt = parser.accepts("create", "Create a new topic.")
+    private val deleteOpt = parser.accepts("delete", "Delete a topic")
+    private val alterOpt = parser.accepts("alter", "Alter the number of partitions, replica assignment, and/or configuration for the topic.")
+    private val describeOpt = parser.accepts("describe", "List details for the given topics.")
+    private val topicOpt = parser.accepts("topic", "The topic to create, alter, describe or delete. It also accepts a regular " +
+                                           "expression, except for --create option. Put topic name in double quotes and use the '\\' prefix " +
+                                           "to escape regular expression symbols; e.g. \"test\\.topic\".")
                          .withRequiredArg
                          .describedAs("topic")
                          .ofType(classOf[String])
-    val nl = System.getProperty("line.separator")
-    val configOpt = parser.accepts("config", "A topic configuration override for the topic being created or altered."  +
-                                             "The following is a list of valid configurations: " + nl + LogConfig.configNames.map("\t" + _).mkString(nl) + nl +
-                                             "See the Kafka documentation for full details on the topic configs.")
+    private val nl = System.getProperty("line.separator")
+    private val kafkaConfigsCanAlterTopicConfigsViaBootstrapServer =
+      " (the kafka-configs CLI supports altering topic configs with a --bootstrap-server option)"
+    private val configOpt = parser.accepts("config", "A topic configuration override for the topic being created or altered." +
+                                             " The following is a list of valid configurations: " + nl + LogConfig.configNames.map("\t" + _).mkString(nl) + nl +
+                                             "See the Kafka documentation for full details on the topic configs." +
+                                             " It is supported only in combination with --create if --bootstrap-server option is used" +
+                                             kafkaConfigsCanAlterTopicConfigsViaBootstrapServer + ".")
                            .withRequiredArg
                            .describedAs("name=value")
                            .ofType(classOf[String])
-    val deleteConfigOpt = parser.accepts("delete-config", "A topic configuration override to be removed for an existing topic (see the list of configurations under the --config option).")
+    private val deleteConfigOpt = parser.accepts("delete-config", "A topic configuration override to be removed for an existing topic (see the list of configurations under the --config option). " +
+      "Not supported with the --bootstrap-server option.")
                            .withRequiredArg
                            .describedAs("name")
                            .ofType(classOf[String])
-    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")
+    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). If not supplied for create, defaults to the cluster default.")
                            .withRequiredArg
                            .describedAs("# of partitions")
                            .ofType(classOf[java.lang.Integer])
-    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])
-    val replicaAssignmentOpt = parser.accepts("replica-assignment", "A list of manual partition-to-broker assignments for the topic being created or altered.")
+    private val replicaAssignmentOpt = parser.accepts("replica-assignment", "A list of manual partition-to-broker assignments for the topic being created or altered.")
                            .withRequiredArg
                            .describedAs("broker_id_for_part1_replica1 : broker_id_for_part1_replica2 , " +
                                         "broker_id_for_part2_replica1 : broker_id_for_part2_replica2 , ...")
                            .ofType(classOf[String])
-    val reportUnderReplicatedPartitionsOpt = parser.accepts("under-replicated-partitions",
-                                                            "if set when describing topics, only show under replicated partitions")
-    val reportUnavailablePartitionsOpt = parser.accepts("unavailable-partitions",
-                                                            "if set when describing topics, only show partitions whose leader is not available")
-    val topicsWithOverridesOpt = parser.accepts("topics-with-overrides",
-                                                "if set when describing topics, only show topics that have overridden configs")
-    val ifExistsOpt = parser.accepts("if-exists",
-                                     "if set when altering or deleting topics, the action will only execute if the topic exists")
-    val ifNotExistsOpt = parser.accepts("if-not-exists",
-                                        "if set when creating topics, the action will only execute if the topic does not already exist")
-
-    val disableRackAware = parser.accepts("disable-rack-aware", "Disable rack aware replica assignment")
-
-    val forceOpt = parser.accepts("force", "Suppress console prompts")
-
-    val options = parser.parse(args : _*)
+    private val reportUnderReplicatedPartitionsOpt = parser.accepts("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")
+    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.")
+    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")
+    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.")
+    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.")
+
+    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")
+
+    options = parser.parse(args : _*)
+
+    private val allTopicLevelOpts = immutable.Set[OptionSpec[_]](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
+
+    def hasCreateOption: Boolean = has(createOpt)
+    def hasAlterOption: Boolean = has(alterOpt)
+    def hasListOption: Boolean = has(listOpt)
+    def hasDescribeOption: Boolean = has(describeOpt)
+    def hasDeleteOption: Boolean = has(deleteOpt)
+
+    def zkConnect: Option[String] = valueAsOption(zkConnectOpt)
+    def bootstrapServer: Option[String] = valueAsOption(bootstrapServerOpt)
+    def commandConfig: Properties = if (has(commandConfigOpt)) Utils.loadProps(options.valueOf(commandConfigOpt)) else new Properties()
+    def topic: Option[String] = valueAsOption(topicOpt)
+    def partitions: Option[Integer] = valueAsOption(partitionsOpt)
+    def replicationFactor: Option[Integer] = valueAsOption(replicationFactorOpt)
+    def replicaAssignment: Option[Map[Int, List[Int]]] =
+      if (has(replicaAssignmentOpt) && !Option(options.valueOf(replicaAssignmentOpt)).getOrElse("").isEmpty)
+        Some(parseReplicaAssignment(options.valueOf(replicaAssignmentOpt)))
+      else
+        None
+    def rackAwareMode: RackAwareMode = if (has(disableRackAware)) RackAwareMode.Disabled else RackAwareMode.Enforced
+    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)
+    def excludeInternalTopics: Boolean = has(excludeInternalTopicOpt)
+    def topicConfig: Option[util.List[String]] = valuesAsOption(configOpt)
+    def configsToDelete: Option[util.List[String]] = valuesAsOption(deleteConfigOpt)
+
+    def checkArgs(): Unit = {
+      if (args.length == 0)
+        CommandLineUtils.printUsageAndDie(parser, "Create, delete, describe, or change a topic.")
+
+      CommandLineUtils.printHelpAndExitIfNeeded(this, "This tool helps to create, delete, describe, or change a topic.")
+
+      // should have exactly one action
+      val actions = Seq(createOpt, listOpt, alterOpt, describeOpt, deleteOpt).count(options.has)
+      if (actions != 1)
+        CommandLineUtils.printUsageAndDie(parser, "Command must include exactly one action: --list, --describe, --create, --alter or --delete")
 
-    val allTopicLevelOpts: Set[OptionSpec[_]] = Set(alterOpt, createOpt, describeOpt, listOpt, deleteOpt)
-
-    def checkArgs() {
       // check required args
-      CommandLineUtils.checkRequiredArgs(parser, options, zkConnectOpt)
-      if (!options.has(listOpt) && !options.has(describeOpt))
+      if (has(bootstrapServerOpt) == has(zkConnectOpt))
+        throw new IllegalArgumentException("Only one of --bootstrap-server or --zookeeper must be specified")
+
+      if (!has(bootstrapServerOpt))
+        CommandLineUtils.checkRequiredArgs(parser, options, zkConnectOpt)
+      if(has(describeOpt) && has(ifExistsOpt))
         CommandLineUtils.checkRequiredArgs(parser, options, topicOpt)
+      if (!has(listOpt) && !has(describeOpt))
+        CommandLineUtils.checkRequiredArgs(parser, options, topicOpt)
+      if (has(createOpt) && !has(replicaAssignmentOpt) && has(zkConnectOpt))
+        CommandLineUtils.checkRequiredArgs(parser, options, partitionsOpt, replicationFactorOpt)
+      if (has(bootstrapServerOpt) && has(alterOpt)) {
+        CommandLineUtils.checkInvalidArgsSet(parser, options, Set(bootstrapServerOpt, configOpt), Set(alterOpt),
+        Some(kafkaConfigsCanAlterTopicConfigsViaBootstrapServer))
+        CommandLineUtils.checkRequiredArgs(parser, options, partitionsOpt)
+      }
 
       // check invalid args
       CommandLineUtils.checkInvalidArgs(parser, options, configOpt, allTopicLevelOpts -- Set(alterOpt, createOpt))
-      CommandLineUtils.checkInvalidArgs(parser, options, deleteConfigOpt, allTopicLevelOpts -- 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))
       CommandLineUtils.checkInvalidArgs(parser, options, replicaAssignmentOpt, allTopicLevelOpts -- Set(createOpt,alterOpt))
       if(options.has(createOpt))
           CommandLineUtils.checkInvalidArgs(parser, options, replicaAssignmentOpt, Set(partitionsOpt, replicationFactorOpt))
       CommandLineUtils.checkInvalidArgs(parser, options, reportUnderReplicatedPartitionsOpt,
-        allTopicLevelOpts -- Set(describeOpt) + reportUnavailablePartitionsOpt + topicsWithOverridesOpt)
+        allTopicLevelOpts -- Set(describeOpt) ++ allReplicationReportOpts - reportUnderReplicatedPartitionsOpt + topicsWithOverridesOpt)
+      CommandLineUtils.checkInvalidArgs(parser, options, reportUnderMinIsrPartitionsOpt,
+        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 + topicsWithOverridesOpt)
+        allTopicLevelOpts -- Set(describeOpt) ++ allReplicationReportOpts - reportUnavailablePartitionsOpt + topicsWithOverridesOpt)
       CommandLineUtils.checkInvalidArgs(parser, options, topicsWithOverridesOpt,
-        allTopicLevelOpts -- Set(describeOpt) + reportUnderReplicatedPartitionsOpt + reportUnavailablePartitionsOpt)
-      CommandLineUtils.checkInvalidArgs(parser, options, ifExistsOpt, allTopicLevelOpts -- Set(alterOpt, deleteOpt))
+        allTopicLevelOpts -- Set(describeOpt) ++ allReplicationReportOpts)
+      CommandLineUtils.checkInvalidArgs(parser, options, ifExistsOpt, allTopicLevelOpts -- Set(alterOpt, deleteOpt, describeOpt))
       CommandLineUtils.checkInvalidArgs(parser, options, ifNotExistsOpt, allTopicLevelOpts -- Set(createOpt))
+      CommandLineUtils.checkInvalidArgs(parser, options, excludeInternalTopicOpt, allTopicLevelOpts -- Set(listOpt, describeOpt))
     }
   }
-
-  def askToProceed(): Unit = {
-    println("Are you sure you want to continue? [y/n]")
-    if (!Console.readLine().equalsIgnoreCase("y")) {
-      println("Ending your session")
-      Exit.exit(0)
-    }
-  }
-
 }
 
diff --git a/core/src/main/scala/kafka/admin/ZkSecurityMigrator.scala b/core/src/main/scala/kafka/admin/ZkSecurityMigrator.scala
index 82b7dac047e6a..db002fc7cdbeb 100644
--- a/core/src/main/scala/kafka/admin/ZkSecurityMigrator.scala
+++ b/core/src/main/scala/kafka/admin/ZkSecurityMigrator.scala
@@ -17,17 +17,21 @@
 
 package kafka.admin
 
-import joptsimple.OptionParser
-import org.I0Itec.zkclient.exception.ZkException
-import kafka.utils.{CommandLineUtils, Logging, ZkUtils}
+import joptsimple.{ArgumentAcceptingOptionSpec, OptionSet}
+import kafka.server.KafkaConfig
+import kafka.utils.{CommandDefaultOptions, CommandLineUtils, Exit, Logging}
+import kafka.utils.Implicits._
+import kafka.zk.{ControllerZNode, KafkaZkClient, ZkData, ZkSecurityMigratorUtils}
 import org.apache.kafka.common.security.JaasUtils
+import org.apache.kafka.common.utils.{Time, Utils}
 import org.apache.zookeeper.AsyncCallback.{ChildrenCallback, StatCallback}
-import org.apache.zookeeper.data.Stat
 import org.apache.zookeeper.KeeperException
 import org.apache.zookeeper.KeeperException.Code
+import org.apache.zookeeper.client.ZKClientConfig
+import org.apache.zookeeper.data.Stat
 
 import scala.annotation.tailrec
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 import scala.collection.mutable.Queue
 import scala.concurrent._
 import scala.concurrent.duration._
@@ -60,39 +64,34 @@ object ZkSecurityMigrator extends Logging {
   val usageMessage = ("ZooKeeper Migration Tool Help. This tool updates the ACLs of "
                       + "znodes as part of the process of setting up ZooKeeper "
                       + "authentication.")
+  val tlsConfigFileOption = "zk-tls-config-file"
 
-  def run(args: Array[String]) {
+  def run(args: Array[String]): Unit = {
     val jaasFile = System.getProperty(JaasUtils.JAVA_LOGIN_CONFIG_PARAM)
-    val parser = new OptionParser(false)
-    val zkAclOpt = parser.accepts("zookeeper.acl", "Indicates whether to make the Kafka znodes in ZooKeeper secure or unsecure."
-        + " The options are 'secure' and 'unsecure'").withRequiredArg().ofType(classOf[String])
-    val zkUrlOpt = parser.accepts("zookeeper.connect", "Sets the ZooKeeper connect string (ensemble). This parameter " +
-      "takes a comma-separated list of host:port pairs.").withRequiredArg().defaultsTo("localhost:2181").
-      ofType(classOf[String])
-    val zkSessionTimeoutOpt = parser.accepts("zookeeper.session.timeout", "Sets the ZooKeeper session timeout.").
-      withRequiredArg().ofType(classOf[java.lang.Integer]).defaultsTo(30000)
-    val zkConnectionTimeoutOpt = parser.accepts("zookeeper.connection.timeout", "Sets the ZooKeeper connection timeout.").
-      withRequiredArg().ofType(classOf[java.lang.Integer]).defaultsTo(30000)
-    val helpOpt = parser.accepts("help", "Print usage information.")
+    val opts = new ZkSecurityMigratorOptions(args)
 
-    val options = parser.parse(args : _*)
-    if (options.has(helpOpt))
-      CommandLineUtils.printUsageAndDie(parser, usageMessage)
+    CommandLineUtils.printHelpAndExitIfNeeded(opts, usageMessage)
 
-    if (jaasFile == null) {
-     val errorMsg = "No JAAS configuration file has been specified. Please make sure that you have set " +
-       "the system property %s".format(JaasUtils.JAVA_LOGIN_CONFIG_PARAM)
-     System.out.println("ERROR: %s".format(errorMsg))
-     throw new IllegalArgumentException("Incorrect configuration")
+    // Must have either SASL or TLS mutual authentication enabled to use this tool.
+    // Instantiate the client config we will use so that we take into account config provided via the CLI option
+    // and system properties passed via -D parameters if no CLI option is given.
+    val zkClientConfig = createZkClientConfigFromOption(opts.options, opts.zkTlsConfigFile).getOrElse(new ZKClientConfig())
+    val tlsClientAuthEnabled = KafkaConfig.zkTlsClientAuthEnabled(zkClientConfig)
+    if (jaasFile == null && !tlsClientAuthEnabled) {
+      val errorMsg = s"No JAAS configuration file has been specified and no TLS client certificate has been specified. Please make sure that you set " +
+        s"the system property ${JaasUtils.JAVA_LOGIN_CONFIG_PARAM} or provide a ZooKeeper client TLS configuration via --$tlsConfigFileOption  " +
+        s"identifying at least ${KafkaConfig.ZkSslClientEnableProp}, ${KafkaConfig.ZkClientCnxnSocketProp}, and ${KafkaConfig.ZkSslKeyStoreLocationProp}"
+      System.err.println("ERROR: %s".format(errorMsg))
+      throw new IllegalArgumentException("Incorrect configuration")
     }
 
-    if (!JaasUtils.isZkSecurityEnabled()) {
+    if (!tlsClientAuthEnabled && !JaasUtils.isZkSaslEnabled()) {
       val errorMsg = "Security isn't enabled, most likely the file isn't set properly: %s".format(jaasFile)
       System.out.println("ERROR: %s".format(errorMsg))
-      throw new IllegalArgumentException("Incorrect configuration") 
+      throw new IllegalArgumentException("Incorrect configuration")
     }
 
-    val zkAcl: Boolean = options.valueOf(zkAclOpt) match {
+    val zkAcl: Boolean = opts.options.valueOf(opts.zkAclOpt) match {
       case "secure" =>
         info("zookeeper.acl option is secure")
         true
@@ -100,50 +99,94 @@ object ZkSecurityMigrator extends Logging {
         info("zookeeper.acl option is unsecure")
         false
       case _ =>
-        CommandLineUtils.printUsageAndDie(parser, usageMessage)
+        CommandLineUtils.printUsageAndDie(opts.parser, usageMessage)
     }
-    val zkUrl = options.valueOf(zkUrlOpt)
-    val zkSessionTimeout = options.valueOf(zkSessionTimeoutOpt).intValue
-    val zkConnectionTimeout = options.valueOf(zkConnectionTimeoutOpt).intValue
-    val zkUtils = ZkUtils(zkUrl, zkSessionTimeout, zkConnectionTimeout, zkAcl)
-    val migrator = new ZkSecurityMigrator(zkUtils)
-    migrator.run()
+    val zkUrl = opts.options.valueOf(opts.zkUrlOpt)
+    val zkSessionTimeout = opts.options.valueOf(opts.zkSessionTimeoutOpt).intValue
+    val zkConnectionTimeout = opts.options.valueOf(opts.zkConnectionTimeoutOpt).intValue
+    val zkClient = KafkaZkClient(zkUrl, zkAcl, zkSessionTimeout, zkConnectionTimeout,
+      Int.MaxValue, Time.SYSTEM, zkClientConfig = Some(zkClientConfig))
+    val enablePathCheck = opts.options.has(opts.enablePathCheckOpt)
+    val migrator = new ZkSecurityMigrator(zkClient)
+    migrator.run(enablePathCheck)
   }
 
-  def main(args: Array[String]) {
+  def main(args: Array[String]): Unit = {
     try {
       run(args)
     } catch {
-        case e: Exception =>
+        case e: Exception => {
           e.printStackTrace()
+          // must exit with non-zero status so system tests will know we failed
+          Exit.exit(1)
+        }
     }
   }
+
+  def createZkClientConfigFromFile(filename: String) : ZKClientConfig = {
+    val zkTlsConfigFileProps = Utils.loadProps(filename, KafkaConfig.ZkSslConfigToSystemPropertyMap.keys.toList.asJava)
+    val zkClientConfig = new ZKClientConfig() // Initializes based on any system properties that have been set
+    // Now override any set system properties with explicitly-provided values from the config file
+    // Emit INFO logs due to camel-case property names encouraging mistakes -- help people see mistakes they make
+    info(s"Found ${zkTlsConfigFileProps.size()} ZooKeeper client configuration properties in file $filename")
+    zkTlsConfigFileProps.asScala.forKeyValue { (key, value) =>
+      info(s"Setting $key")
+      KafkaConfig.setZooKeeperClientProperty(zkClientConfig, key, value)
+    }
+    zkClientConfig
+  }
+
+  private[admin] def createZkClientConfigFromOption(options: OptionSet, option: ArgumentAcceptingOptionSpec[String]) : Option[ZKClientConfig] =
+    if (!options.has(option))
+      None
+    else
+      Some(createZkClientConfigFromFile(options.valueOf(option)))
+
+  class ZkSecurityMigratorOptions(args: Array[String]) extends CommandDefaultOptions(args) {
+    val zkAclOpt = parser.accepts("zookeeper.acl", "Indicates whether to make the Kafka znodes in ZooKeeper secure or unsecure."
+      + " The options are 'secure' and 'unsecure'").withRequiredArg().ofType(classOf[String])
+    val zkUrlOpt = parser.accepts("zookeeper.connect", "Sets the ZooKeeper connect string (ensemble). This parameter " +
+      "takes a comma-separated list of host:port pairs.").withRequiredArg().defaultsTo("localhost:2181").
+      ofType(classOf[String])
+    val zkSessionTimeoutOpt = parser.accepts("zookeeper.session.timeout", "Sets the ZooKeeper session timeout.").
+      withRequiredArg().ofType(classOf[java.lang.Integer]).defaultsTo(30000)
+    val zkConnectionTimeoutOpt = parser.accepts("zookeeper.connection.timeout", "Sets the ZooKeeper connection timeout.").
+      withRequiredArg().ofType(classOf[java.lang.Integer]).defaultsTo(30000)
+    val enablePathCheckOpt = parser.accepts("enable.path.check", "Checks if all the root paths exist in ZooKeeper " +
+      "before migration. If not, exit the command.")
+    val zkTlsConfigFile = parser.accepts(tlsConfigFileOption,
+      "Identifies the file where ZooKeeper client TLS connectivity properties are defined.  Any properties other than " +
+        KafkaConfig.ZkSslConfigToSystemPropertyMap.keys.mkString(", ") + " are ignored.")
+      .withRequiredArg().describedAs("ZooKeeper TLS configuration").ofType(classOf[String])
+    options = parser.parse(args : _*)
+  }
 }
 
-class ZkSecurityMigrator(zkUtils: ZkUtils) extends Logging {
+class ZkSecurityMigrator(zkClient: KafkaZkClient) extends Logging {
+  private val zkSecurityMigratorUtils = new ZkSecurityMigratorUtils(zkClient)
   private val futures = new Queue[Future[String]]
 
-  private def setAcl(path: String, setPromise: Promise[String]) = {
+  private def setAcl(path: String, setPromise: Promise[String]): Unit = {
     info("Setting ACL for path %s".format(path))
-    zkUtils.zkConnection.getZookeeper.setACL(path, zkUtils.defaultAcls(path), -1, SetACLCallback, setPromise)
+    zkSecurityMigratorUtils.currentZooKeeper.setACL(path, zkClient.defaultAcls(path).asJava, -1, SetACLCallback, setPromise)
   }
 
-  private def getChildren(path: String, childrenPromise: Promise[String]) = {
+  private def getChildren(path: String, childrenPromise: Promise[String]): Unit = {
     info("Getting children to set ACLs for path %s".format(path))
-    zkUtils.zkConnection.getZookeeper.getChildren(path, false, GetChildrenCallback, childrenPromise)
+    zkSecurityMigratorUtils.currentZooKeeper.getChildren(path, false, GetChildrenCallback, childrenPromise)
   }
 
-  private def setAclIndividually(path: String) = {
-    val setPromise = Promise[String]
+  private def setAclIndividually(path: String): Unit = {
+    val setPromise = Promise[String]()
     futures.synchronized {
       futures += setPromise.future
     }
     setAcl(path, setPromise)
   }
 
-  private def setAclsRecursively(path: String) = {
-    val setPromise = Promise[String]
-    val childrenPromise = Promise[String]
+  private def setAclsRecursively(path: String): Unit = {
+    val setPromise = Promise[String]()
+    val childrenPromise = Promise[String]()
     futures.synchronized {
       futures += setPromise.future
       futures += childrenPromise.future
@@ -156,8 +199,8 @@ class ZkSecurityMigrator(zkUtils: ZkUtils) extends Logging {
     def processResult(rc: Int,
                       path: String,
                       ctx: Object,
-                      children: java.util.List[String]) {
-      val zkHandle = zkUtils.zkConnection.getZookeeper
+                      children: java.util.List[String]): Unit = {
+      val zkHandle = zkSecurityMigratorUtils.currentZooKeeper
       val promise = ctx.asInstanceOf[Promise[String]]
       Code.get(rc) match {
         case Code.OK =>
@@ -178,10 +221,10 @@ class ZkSecurityMigrator(zkUtils: ZkUtils) 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))
       }
     }
   }
@@ -190,8 +233,8 @@ class ZkSecurityMigrator(zkUtils: ZkUtils) extends Logging {
     def processResult(rc: Int,
                       path: String,
                       ctx: Object,
-                      stat: Stat) {
-      val zkHandle = zkUtils.zkConnection.getZookeeper
+                      stat: Stat): Unit = {
+      val zkHandle = zkSecurityMigratorUtils.currentZooKeeper
       val promise = ctx.asInstanceOf[Promise[String]]
 
       Code.get(rc) match {
@@ -199,7 +242,7 @@ class ZkSecurityMigrator(zkUtils: ZkUtils) extends Logging {
           info("Successfully set ACLs for %s".format(path))
           promise success "done"
         case Code.CONNECTIONLOSS =>
-            zkHandle.setACL(path, zkUtils.defaultAcls(path), -1, SetACLCallback, ctx)
+            zkHandle.setACL(path, zkClient.defaultAcls(path).asJava, -1, SetACLCallback, ctx)
         case Code.NONODE =>
           warn("Znode is gone, it could be have been legitimately deleted: %s".format(path))
           promise success "done"
@@ -207,21 +250,26 @@ class ZkSecurityMigrator(zkUtils: ZkUtils) 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))
       }
     }
   }
 
-  private def run(): Unit = {
+  private def run(enablePathCheck: Boolean): Unit = {
     try {
       setAclIndividually("/")
-      for (path <- ZkUtils.SecureZkRootPaths) {
+      checkPathExistenceAndMaybeExit(enablePathCheck)
+      for (path <- ZkData.SecureRootPaths) {
         debug("Going to set ACL for %s".format(path))
-        zkUtils.makeSurePersistentPathExists(path)
-        setAclsRecursively(path)
+        if (path == ControllerZNode.path && !zkClient.pathExists(path)) {
+          debug("Ignoring to set ACL for %s, because it doesn't exist".format(path))
+        } else {
+          zkClient.makeSurePersistentPathExists(path)
+          setAclsRecursively(path)
+        }
       }
 
       @tailrec
@@ -232,15 +280,28 @@ class ZkSecurityMigrator(zkUtils: ZkUtils) extends Logging {
         future match {
           case Some(a) =>
             Await.result(a, 6000 millis)
-            futures.synchronized { futures.dequeue }
-            recurse
+            futures.synchronized { futures.dequeue() }
+            recurse()
           case None =>
         }
       }
       recurse()
 
     } finally {
-      zkUtils.close
+      zkClient.close()
+    }
+  }
+
+  private def checkPathExistenceAndMaybeExit(enablePathCheck: Boolean): Unit = {
+    val nonExistingSecureRootPaths = ZkData.SecureRootPaths.filterNot(zkClient.pathExists)
+    if (nonExistingSecureRootPaths.nonEmpty) {
+      println(s"Warning: The following secure root paths do not exist in ZooKeeper: ${nonExistingSecureRootPaths.mkString(",")}")
+      println("That might be due to an incorrect chroot is specified when executing the command.")
+      if (enablePathCheck) {
+        println("Exit the command.")
+        // must exit with non-zero status so system tests will know we failed
+        Exit.exit(1)
+      }
     }
   }
 }
diff --git a/core/src/main/scala/kafka/api/ApiUtils.scala b/core/src/main/scala/kafka/api/ApiUtils.scala
index 2145e8c5c672c..9be1e4bb87467 100644
--- a/core/src/main/scala/kafka/api/ApiUtils.scala
+++ b/core/src/main/scala/kafka/api/ApiUtils.scala
@@ -16,15 +16,15 @@
  */
 package kafka.api
 
-import java.nio._
-import kafka.common._
+import java.nio.ByteBuffer
+import java.nio.charset.StandardCharsets
+
+import org.apache.kafka.common.KafkaException
 
 /**
  * Helper functions specific to parsing or serializing requests and responses
  */
 object ApiUtils {
-  
-  val ProtocolEncoding = "UTF-8"
 
     /**
    * Read size prefixed string where the size is stored as a 2 byte short.
@@ -36,7 +36,7 @@ object ApiUtils {
       return null
     val bytes = new Array[Byte](size)
     buffer.get(bytes)
-    new String(bytes, ProtocolEncoding)
+    new String(bytes, StandardCharsets.UTF_8)
   }
   
   /**
@@ -44,11 +44,11 @@ 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 {
-      val encodedString = string.getBytes(ProtocolEncoding)
+      val encodedString = string.getBytes(StandardCharsets.UTF_8)
       if(encodedString.length > Short.MaxValue) {
         throw new KafkaException("String exceeds the maximum size of " + Short.MaxValue + ".")
       } else {
@@ -66,7 +66,7 @@ object ApiUtils {
     if(string == null) {
       2
     } else {
-      val encodedString = string.getBytes(ProtocolEncoding)
+      val encodedString = string.getBytes(StandardCharsets.UTF_8)
       if(encodedString.length > Short.MaxValue) {
         throw new KafkaException("String exceeds the maximum size of " + Short.MaxValue + ".")
       } else {
@@ -75,26 +75,4 @@ object ApiUtils {
     }
   }
   
-  /**
-   * Read an integer out of the bytebuffer from the current position and check that it falls within the given
-   * range. If not, throw KafkaException.
-   */
-  def readIntInRange(buffer: ByteBuffer, name: String, range: (Int, Int)): Int = {
-    val value = buffer.getInt
-    if(value < range._1 || value > range._2)
-      throw new KafkaException(name + " has value " + value + " which is not in the range " + range + ".")
-    else value
-  }
-
-  /**
-   * Read a short out of the bytebuffer from the current position and check that it falls within the given
-   * range. If not, throw KafkaException.
-   */
-  def readShortInRange(buffer: ByteBuffer, name: String, range: (Short, Short)): Short = {
-    val value = buffer.getShort
-    if(value < range._1 || value > range._2)
-      throw new KafkaException(name + " has value " + value + " which is not in the range " + range + ".")
-    else value
-  }
-  
 }
diff --git a/core/src/main/scala/kafka/api/ApiVersion.scala b/core/src/main/scala/kafka/api/ApiVersion.scala
index e509fc516ab13..44e987d327ceb 100644
--- a/core/src/main/scala/kafka/api/ApiVersion.scala
+++ b/core/src/main/scala/kafka/api/ApiVersion.scala
@@ -17,7 +17,13 @@
 
 package kafka.api
 
-import org.apache.kafka.common.record.RecordBatch
+import org.apache.kafka.common.config.ConfigDef.Validator
+import org.apache.kafka.common.config.ConfigException
+import org.apache.kafka.common.feature.{Features, FinalizedVersionRange, SupportedVersionRange}
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.record.{RecordBatch, RecordVersion}
+import org.apache.kafka.common.requests.{AbstractResponse, ApiVersionsResponse}
+import org.apache.kafka.common.requests.ApiVersionsResponse.DEFAULT_API_VERSIONS_RESPONSE
 
 /**
  * This class contains the different Kafka versions.
@@ -42,56 +48,151 @@ object ApiVersion {
   // This implicit is necessary due to: https://issues.scala-lang.org/browse/SI-8541
   implicit def orderingByVersion[A <: ApiVersion]: Ordering[A] = Ordering.by(_.id)
 
-  private val versionNameMap = Map(
-    "0.8.0" -> KAFKA_0_8_0,
-    "0.8.1" -> KAFKA_0_8_1,
-    "0.8.2" -> KAFKA_0_8_2,
-    "0.9.0" -> KAFKA_0_9_0,
+  val allVersions: Seq[ApiVersion] = Seq(
+    KAFKA_0_8_0,
+    KAFKA_0_8_1,
+    KAFKA_0_8_2,
+    KAFKA_0_9_0,
     // 0.10.0-IV0 is introduced for KIP-31/32 which changes the message format.
-    "0.10.0-IV0" -> KAFKA_0_10_0_IV0,
+    KAFKA_0_10_0_IV0,
     // 0.10.0-IV1 is introduced for KIP-36(rack awareness) and KIP-43(SASL handshake).
-    "0.10.0-IV1" -> KAFKA_0_10_0_IV1,
-    "0.10.0" -> KAFKA_0_10_0_IV1,
-
+    KAFKA_0_10_0_IV1,
     // introduced for JoinGroup protocol change in KIP-62
-    "0.10.1-IV0" -> KAFKA_0_10_1_IV0,
+    KAFKA_0_10_1_IV0,
     // 0.10.1-IV1 is introduced for KIP-74(fetch response size limit).
-    "0.10.1-IV1" -> KAFKA_0_10_1_IV1,
+    KAFKA_0_10_1_IV1,
     // introduced ListOffsetRequest v1 in KIP-79
-    "0.10.1-IV2" -> KAFKA_0_10_1_IV2,
-    "0.10.1" -> KAFKA_0_10_1_IV2,
+    KAFKA_0_10_1_IV2,
     // introduced UpdateMetadataRequest v3 in KIP-103
-    "0.10.2-IV0" -> KAFKA_0_10_2_IV0,
-    "0.10.2" -> KAFKA_0_10_2_IV0,
+    KAFKA_0_10_2_IV0,
     // KIP-98 (idempotent and transactional producer support)
-    "0.11.0-IV0" -> KAFKA_0_11_0_IV0,
+    KAFKA_0_11_0_IV0,
     // introduced DeleteRecordsRequest v0 and FetchRequest v4 in KIP-107
-    "0.11.0-IV1" -> KAFKA_0_11_0_IV1,
+    KAFKA_0_11_0_IV1,
     // Introduced leader epoch fetches to the replica fetcher via KIP-101
-    "0.11.0-IV2" -> KAFKA_0_11_0_IV2,
-    "0.11.0" -> KAFKA_0_11_0_IV2,
+    KAFKA_0_11_0_IV2,
     // Introduced LeaderAndIsrRequest V1, UpdateMetadataRequest V4 and FetchRequest V6 via KIP-112
-    "1.0-IV0" -> KAFKA_1_0_IV0,
-    "1.0" -> KAFKA_1_0_IV0
+    KAFKA_1_0_IV0,
+    // Introduced DeleteGroupsRequest V0 via KIP-229, plus KIP-227 incremental fetch requests,
+    // and KafkaStorageException for fetch requests.
+    KAFKA_1_1_IV0,
+    // Introduced OffsetsForLeaderEpochRequest V1 via KIP-279 (Fix log divergence between leader and follower after fast leader fail over)
+    KAFKA_2_0_IV0,
+    // Several request versions were bumped due to KIP-219 (Improve quota communication)
+    KAFKA_2_0_IV1,
+    // Introduced new schemas for group offset (v2) and group metadata (v2) (KIP-211)
+    KAFKA_2_1_IV0,
+    // New Fetch, OffsetsForLeaderEpoch, and ListOffsets schemas (KIP-320)
+    KAFKA_2_1_IV1,
+    // Support ZStandard Compression Codec (KIP-110)
+    KAFKA_2_1_IV2,
+    // Introduced broker generation (KIP-380), and
+    // 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,
+    // 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 adding_replicas and removing_replicas fields to LeaderAndIsrRequest
+    KAFKA_2_4_IV0,
+    // Flexible version support in inter-broker APIs
+    KAFKA_2_4_IV1,
+    // No new APIs, equivalent to 2.4-IV1
+    KAFKA_2_5_IV0,
+    // Introduced StopReplicaRequest V3 containing the leader epoch for each partition (KIP-570)
+    KAFKA_2_6_IV0,
+    // Introduced feature versioning support (KIP-584)
+    KAFKA_2_7_IV0,
+    // Bup Fetch protocol for Raft protocol (KIP-595)
+    KAFKA_2_7_IV1,
+    // Introduced AlterIsr (KIP-497)
+    KAFKA_2_7_IV2,
+    // Flexible versioning on ListOffsets, WriteTxnMarkers and OffsetsForLeaderEpoch.
+    KAFKA_2_8_IV0
   )
 
-  private val versionPattern = "\\.".r
+  // Map keys are the union of the short and full versions
+  private val versionMap: Map[String, ApiVersion] =
+    allVersions.map(v => v.version -> v).toMap ++ allVersions.groupBy(_.shortVersion).map { case (k, v) => k -> v.last }
+
+  /**
+   * Return an `ApiVersion` instance for `versionString`, which can be in a variety of formats (e.g. "0.8.0", "0.8.0.x",
+   * "0.10.0", "0.10.0-IV1"). `IllegalArgumentException` is thrown if `versionString` cannot be mapped to an `ApiVersion`.
+   */
+  def apply(versionString: String): ApiVersion = {
+    val versionSegments = versionString.split('.').toSeq
+    val numSegments = if (versionString.startsWith("0.")) 3 else 2
+    val key = versionSegments.take(numSegments).mkString(".")
+    versionMap.getOrElse(key, throw new IllegalArgumentException(s"Version `$versionString` is not a valid version"))
+  }
+
+  val latestVersion: ApiVersion = allVersions.last
 
-  def apply(version: String): ApiVersion = {
-    val versionsSeq = versionPattern.split(version)
-    val numSegments = if (version.startsWith("0.")) 3 else 2
-    val key = versionsSeq.take(numSegments).mkString(".")
-    versionNameMap.getOrElse(key, throw new IllegalArgumentException(s"Version `$version` is not a valid version"))
+  def isTruncationOnFetchSupported(version: ApiVersion): Boolean = version >= KAFKA_2_7_IV1
+
+  /**
+   * Return the minimum `ApiVersion` that supports `RecordVersion`.
+   */
+  def minSupportedFor(recordVersion: RecordVersion): ApiVersion = {
+    recordVersion match {
+      case RecordVersion.V0 => KAFKA_0_8_0
+      case RecordVersion.V1 => KAFKA_0_10_0_IV0
+      case RecordVersion.V2 => KAFKA_0_11_0_IV0
+      case _ => throw new IllegalArgumentException(s"Invalid message format version $recordVersion")
+    }
+  }
+
+  def apiVersionsResponse(throttleTimeMs: Int,
+                          maxMagic: Byte,
+                          latestSupportedFeatures: Features[SupportedVersionRange]): ApiVersionsResponse = {
+    apiVersionsResponse(
+      throttleTimeMs,
+      maxMagic,
+      latestSupportedFeatures,
+      Features.emptyFinalizedFeatures,
+      ApiVersionsResponse.UNKNOWN_FINALIZED_FEATURES_EPOCH
+    )
   }
 
-  def latestVersion = versionNameMap.values.max
+  def apiVersionsResponse(throttleTimeMs: Int,
+                          maxMagic: Byte,
+                          latestSupportedFeatures: Features[SupportedVersionRange],
+                          finalizedFeatures: Features[FinalizedVersionRange],
+                          finalizedFeaturesEpoch: Long): ApiVersionsResponse = {
+    val apiKeys = ApiVersionsResponse.defaultApiKeys(maxMagic)
+    if (maxMagic == RecordBatch.CURRENT_MAGIC_VALUE &&
+      throttleTimeMs == AbstractResponse.DEFAULT_THROTTLE_TIME)
+      return new ApiVersionsResponse(
+        ApiVersionsResponse.createApiVersionsResponseData(
+          DEFAULT_API_VERSIONS_RESPONSE.throttleTimeMs,
+          Errors.forCode(DEFAULT_API_VERSIONS_RESPONSE.data.errorCode),
+          apiKeys,
+          latestSupportedFeatures,
+          finalizedFeatures,
+          finalizedFeaturesEpoch)
+      )
 
+    new ApiVersionsResponse(
+      ApiVersionsResponse.createApiVersionsResponseData(
+        throttleTimeMs,
+        Errors.NONE,
+        apiKeys,
+        latestSupportedFeatures,
+        finalizedFeatures,
+        finalizedFeaturesEpoch)
+    )
+  }
 }
 
 sealed trait ApiVersion extends Ordered[ApiVersion] {
-  val version: String
-  val messageFormatVersion: Byte
-  val id: Int
+  def version: String
+  def shortVersion: String
+  def recordVersion: RecordVersion
+  def id: Int
+
+  def isAlterIsrSupported: Boolean = this >= KAFKA_2_7_IV2
 
   override def compare(that: ApiVersion): Int =
     ApiVersion.orderingByVersion.compare(this, that)
@@ -99,88 +200,252 @@ sealed trait ApiVersion extends Ordered[ApiVersion] {
   override def toString: String = version
 }
 
+/**
+ * For versions before 0.10.0, `version` and `shortVersion` were the same.
+ */
+sealed trait LegacyApiVersion extends ApiVersion {
+  def version = shortVersion
+}
+
+/**
+ * From 0.10.0 onwards, each version has a sub-version. For example, IV0 is the sub-version of 0.10.0-IV0.
+ */
+sealed trait DefaultApiVersion extends ApiVersion {
+  lazy val version = shortVersion + "-" + subVersion
+  protected def subVersion: String
+}
+
 // Keep the IDs in order of versions
-case object KAFKA_0_8_0 extends ApiVersion {
-  val version: String = "0.8.0.X"
-  val messageFormatVersion: Byte = RecordBatch.MAGIC_VALUE_V0
+case object KAFKA_0_8_0 extends LegacyApiVersion {
+  val shortVersion = "0.8.0"
+  val recordVersion = RecordVersion.V0
   val id: Int = 0
 }
 
-case object KAFKA_0_8_1 extends ApiVersion {
-  val version: String = "0.8.1.X"
-  val messageFormatVersion: Byte = RecordBatch.MAGIC_VALUE_V0
+case object KAFKA_0_8_1 extends LegacyApiVersion {
+  val shortVersion = "0.8.1"
+  val recordVersion = RecordVersion.V0
   val id: Int = 1
 }
 
-case object KAFKA_0_8_2 extends ApiVersion {
-  val version: String = "0.8.2.X"
-  val messageFormatVersion: Byte = RecordBatch.MAGIC_VALUE_V0
+case object KAFKA_0_8_2 extends LegacyApiVersion {
+  val shortVersion = "0.8.2"
+  val recordVersion = RecordVersion.V0
   val id: Int = 2
 }
 
-case object KAFKA_0_9_0 extends ApiVersion {
-  val version: String = "0.9.0.X"
-  val messageFormatVersion: Byte = RecordBatch.MAGIC_VALUE_V0
+case object KAFKA_0_9_0 extends LegacyApiVersion {
+  val shortVersion = "0.9.0"
+  val subVersion = ""
+  val recordVersion = RecordVersion.V0
   val id: Int = 3
 }
 
-case object KAFKA_0_10_0_IV0 extends ApiVersion {
-  val version: String = "0.10.0-IV0"
-  val messageFormatVersion: Byte = RecordBatch.MAGIC_VALUE_V1
+case object KAFKA_0_10_0_IV0 extends DefaultApiVersion {
+  val shortVersion = "0.10.0"
+  val subVersion = "IV0"
+  val recordVersion = RecordVersion.V1
   val id: Int = 4
 }
 
-case object KAFKA_0_10_0_IV1 extends ApiVersion {
-  val version: String = "0.10.0-IV1"
-  val messageFormatVersion: Byte = RecordBatch.MAGIC_VALUE_V1
+case object KAFKA_0_10_0_IV1 extends DefaultApiVersion {
+  val shortVersion = "0.10.0"
+  val subVersion = "IV1"
+  val recordVersion = RecordVersion.V1
   val id: Int = 5
 }
 
-case object KAFKA_0_10_1_IV0 extends ApiVersion {
-  val version: String = "0.10.1-IV0"
-  val messageFormatVersion: Byte = RecordBatch.MAGIC_VALUE_V1
+case object KAFKA_0_10_1_IV0 extends DefaultApiVersion {
+  val shortVersion = "0.10.1"
+  val subVersion = "IV0"
+  val recordVersion = RecordVersion.V1
   val id: Int = 6
 }
 
-case object KAFKA_0_10_1_IV1 extends ApiVersion {
-  val version: String = "0.10.1-IV1"
-  val messageFormatVersion: Byte = RecordBatch.MAGIC_VALUE_V1
+case object KAFKA_0_10_1_IV1 extends DefaultApiVersion {
+  val shortVersion = "0.10.1"
+  val subVersion = "IV1"
+  val recordVersion = RecordVersion.V1
   val id: Int = 7
 }
 
-case object KAFKA_0_10_1_IV2 extends ApiVersion {
-  val version: String = "0.10.1-IV2"
-  val messageFormatVersion: Byte = RecordBatch.MAGIC_VALUE_V1
+case object KAFKA_0_10_1_IV2 extends DefaultApiVersion {
+  val shortVersion = "0.10.1"
+  val subVersion = "IV2"
+  val recordVersion = RecordVersion.V1
   val id: Int = 8
 }
 
-case object KAFKA_0_10_2_IV0 extends ApiVersion {
-  val version: String = "0.10.2-IV0"
-  val messageFormatVersion: Byte = RecordBatch.MAGIC_VALUE_V1
+case object KAFKA_0_10_2_IV0 extends DefaultApiVersion {
+  val shortVersion = "0.10.2"
+  val subVersion = "IV0"
+  val recordVersion = RecordVersion.V1
   val id: Int = 9
 }
 
-case object KAFKA_0_11_0_IV0 extends ApiVersion {
-  val version: String = "0.11.0-IV0"
-  val messageFormatVersion: Byte = RecordBatch.MAGIC_VALUE_V2
+case object KAFKA_0_11_0_IV0 extends DefaultApiVersion {
+  val shortVersion = "0.11.0"
+  val subVersion = "IV0"
+  val recordVersion = RecordVersion.V2
   val id: Int = 10
 }
 
-case object KAFKA_0_11_0_IV1 extends ApiVersion {
-  val version: String = "0.11.0-IV1"
-  val messageFormatVersion: Byte = RecordBatch.MAGIC_VALUE_V2
+case object KAFKA_0_11_0_IV1 extends DefaultApiVersion {
+  val shortVersion = "0.11.0"
+  val subVersion = "IV1"
+  val recordVersion = RecordVersion.V2
   val id: Int = 11
 }
 
-case object KAFKA_0_11_0_IV2 extends ApiVersion {
-  val version: String = "0.11.0-IV2"
-  val messageFormatVersion: Byte = RecordBatch.MAGIC_VALUE_V2
+case object KAFKA_0_11_0_IV2 extends DefaultApiVersion {
+  val shortVersion = "0.11.0"
+  val subVersion = "IV2"
+  val recordVersion = RecordVersion.V2
   val id: Int = 12
 }
 
-case object KAFKA_1_0_IV0 extends ApiVersion {
-  val version: String = "1.0-IV0"
-  val messageFormatVersion: Byte = RecordBatch.MAGIC_VALUE_V2
+case object KAFKA_1_0_IV0 extends DefaultApiVersion {
+  val shortVersion = "1.0"
+  val subVersion = "IV0"
+  val recordVersion = RecordVersion.V2
   val id: Int = 13
 }
 
+case object KAFKA_1_1_IV0 extends DefaultApiVersion {
+  val shortVersion = "1.1"
+  val subVersion = "IV0"
+  val recordVersion = RecordVersion.V2
+  val id: Int = 14
+}
+
+case object KAFKA_2_0_IV0 extends DefaultApiVersion {
+  val shortVersion: String = "2.0"
+  val subVersion = "IV0"
+  val recordVersion = RecordVersion.V2
+  val id: Int = 15
+}
+
+case object KAFKA_2_0_IV1 extends DefaultApiVersion {
+  val shortVersion: String = "2.0"
+  val subVersion = "IV1"
+  val recordVersion = RecordVersion.V2
+  val id: Int = 16
+}
+
+case object KAFKA_2_1_IV0 extends DefaultApiVersion {
+  val shortVersion: String = "2.1"
+  val subVersion = "IV0"
+  val recordVersion = RecordVersion.V2
+  val id: Int = 17
+}
+
+case object KAFKA_2_1_IV1 extends DefaultApiVersion {
+  val shortVersion: String = "2.1"
+  val subVersion = "IV1"
+  val recordVersion = RecordVersion.V2
+  val id: Int = 18
+}
+
+case object KAFKA_2_1_IV2 extends DefaultApiVersion {
+  val shortVersion: String = "2.1"
+  val subVersion = "IV2"
+  val recordVersion = RecordVersion.V2
+  val id: Int = 19
+}
+
+case object KAFKA_2_2_IV0 extends DefaultApiVersion {
+  val shortVersion: String = "2.2"
+  val subVersion = "IV0"
+  val recordVersion = RecordVersion.V2
+  val id: Int = 20
+}
+
+case object KAFKA_2_2_IV1 extends DefaultApiVersion {
+  val shortVersion: String = "2.2"
+  val subVersion = "IV1"
+  val recordVersion = RecordVersion.V2
+  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_4_IV0 extends DefaultApiVersion {
+  val shortVersion: String = "2.4"
+  val subVersion = "IV0"
+  val recordVersion = RecordVersion.V2
+  val id: Int = 24
+}
+
+case object KAFKA_2_4_IV1 extends DefaultApiVersion {
+  val shortVersion: String = "2.4"
+  val subVersion = "IV1"
+  val recordVersion = RecordVersion.V2
+  val id: Int = 25
+}
+
+case object KAFKA_2_5_IV0 extends DefaultApiVersion {
+  val shortVersion: String = "2.5"
+  val subVersion = "IV0"
+  val recordVersion = RecordVersion.V2
+  val id: Int = 26
+}
+
+case object KAFKA_2_6_IV0 extends DefaultApiVersion {
+  val shortVersion: String = "2.6"
+  val subVersion = "IV0"
+  val recordVersion = RecordVersion.V2
+  val id: Int = 27
+}
+
+case object KAFKA_2_7_IV0 extends DefaultApiVersion {
+  val shortVersion: String = "2.7"
+  val subVersion = "IV0"
+  val recordVersion = RecordVersion.V2
+  val id: Int = 28
+}
+
+case object KAFKA_2_7_IV1 extends DefaultApiVersion {
+  val shortVersion: String = "2.7"
+  val subVersion = "IV1"
+  val recordVersion = RecordVersion.V2
+  val id: Int = 29
+}
+
+case object KAFKA_2_7_IV2 extends DefaultApiVersion {
+  val shortVersion: String = "2.7"
+  val subVersion = "IV2"
+  val recordVersion = RecordVersion.V2
+  val id: Int = 30
+}
+
+case object KAFKA_2_8_IV0 extends DefaultApiVersion {
+  val shortVersion: String = "2.8"
+  val subVersion = "IV0"
+  val recordVersion = RecordVersion.V2
+  val id: Int = 31
+}
+
+object ApiVersionValidator extends Validator {
+
+  override def ensureValid(name: String, value: Any): Unit = {
+    try {
+      ApiVersion(value.toString)
+    } catch {
+      case e: IllegalArgumentException => throw new ConfigException(name, value.toString, e.getMessage)
+    }
+  }
+
+  override def toString: String = "[" + ApiVersion.allVersions.map(_.version).distinct.mkString(", ") + "]"
+}
diff --git a/core/src/main/scala/kafka/api/FetchRequest.scala b/core/src/main/scala/kafka/api/FetchRequest.scala
deleted file mode 100644
index 0379559e2db2c..0000000000000
--- a/core/src/main/scala/kafka/api/FetchRequest.scala
+++ /dev/null
@@ -1,270 +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 kafka.utils.nonthreadsafe
-import kafka.api.ApiUtils._
-import kafka.common.TopicAndPartition
-import kafka.consumer.ConsumerConfig
-import java.util.concurrent.atomic.AtomicInteger
-import java.nio.ByteBuffer
-
-import org.apache.kafka.common.protocol.ApiKeys
-
-import scala.collection.mutable.ArrayBuffer
-import scala.util.Random
-
-case class PartitionFetchInfo(offset: Long, fetchSize: Int)
-
-@deprecated("This object has been deprecated and will be removed in a future release.", "0.11.0.0")
-object FetchRequest {
-
-  private val random = new Random
-
-  val CurrentVersion = 3.shortValue
-  val DefaultMaxWait = 0
-  val DefaultMinBytes = 0
-  val DefaultMaxBytes = Int.MaxValue
-  val DefaultCorrelationId = 0
-
-  def readFrom(buffer: ByteBuffer): FetchRequest = {
-    val versionId = buffer.getShort
-    val correlationId = buffer.getInt
-    val clientId = readShortString(buffer)
-    val replicaId = buffer.getInt
-    val maxWait = buffer.getInt
-    val minBytes = buffer.getInt
-    val maxBytes = if (versionId < 3) DefaultMaxBytes else buffer.getInt
-    val topicCount = buffer.getInt
-    val pairs = (1 to topicCount).flatMap(_ => {
-      val topic = readShortString(buffer)
-      val partitionCount = buffer.getInt
-      (1 to partitionCount).map(_ => {
-        val partitionId = buffer.getInt
-        val offset = buffer.getLong
-        val fetchSize = buffer.getInt
-        (TopicAndPartition(topic, partitionId), PartitionFetchInfo(offset, fetchSize))
-      })
-    })
-    FetchRequest(versionId, correlationId, clientId, replicaId, maxWait, minBytes, maxBytes, Vector(pairs:_*))
-  }
-
-  def shuffle(requestInfo: Seq[(TopicAndPartition, PartitionFetchInfo)]): Seq[(TopicAndPartition, PartitionFetchInfo)] = {
-    val groupedByTopic = requestInfo.groupBy { case (tp, _) => tp.topic }.map { case (topic, values) =>
-      topic -> random.shuffle(values)
-    }
-    random.shuffle(groupedByTopic.toSeq).flatMap { case (_, partitions) =>
-      partitions.map { case (tp, fetchInfo) => tp -> fetchInfo }
-    }
-  }
-
-  def batchByTopic[T](s: Seq[(TopicAndPartition, T)]): Seq[(String, Seq[(Int, T)])] = {
-    val result = new ArrayBuffer[(String, ArrayBuffer[(Int, T)])]
-    s.foreach { case (TopicAndPartition(t, p), value) =>
-      if (result.isEmpty || result.last._1 != t)
-        result += (t -> new ArrayBuffer)
-      result.last._2 += (p -> value)
-    }
-    result
-  }
-
-}
-
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.11.0.0")
-case class FetchRequest(versionId: Short = FetchRequest.CurrentVersion,
-                        correlationId: Int = FetchRequest.DefaultCorrelationId,
-                        clientId: String = ConsumerConfig.DefaultClientId,
-                        replicaId: Int = Request.OrdinaryConsumerId,
-                        maxWait: Int = FetchRequest.DefaultMaxWait,
-                        minBytes: Int = FetchRequest.DefaultMinBytes,
-                        maxBytes: Int = FetchRequest.DefaultMaxBytes,
-                        requestInfo: Seq[(TopicAndPartition, PartitionFetchInfo)])
-        extends RequestOrResponse(Some(ApiKeys.FETCH.id)) {
-
-  /**
-    * Partitions the request info into a list of lists (one for each topic) while preserving request info ordering
-    */
-  private type PartitionInfos = Seq[(Int, PartitionFetchInfo)]
-  private lazy val requestInfoGroupedByTopic: Seq[(String, PartitionInfos)] = FetchRequest.batchByTopic(requestInfo)
-
-  /** Public constructor for the clients */
-  @deprecated("The order of partitions in `requestInfo` is relevant, so this constructor is deprecated in favour of the " +
-    "one that takes a Seq", since = "0.10.1.0")
-  def this(correlationId: Int,
-           clientId: String,
-           maxWait: Int,
-           minBytes: Int,
-           maxBytes: Int,
-           requestInfo: Map[TopicAndPartition, PartitionFetchInfo]) {
-    this(versionId = FetchRequest.CurrentVersion,
-         correlationId = correlationId,
-         clientId = clientId,
-         replicaId = Request.OrdinaryConsumerId,
-         maxWait = maxWait,
-         minBytes = minBytes,
-         maxBytes = maxBytes,
-         requestInfo = FetchRequest.shuffle(requestInfo.toSeq))
-  }
-
-  /** Public constructor for the clients */
-  def this(correlationId: Int,
-           clientId: String,
-           maxWait: Int,
-           minBytes: Int,
-           maxBytes: Int,
-           requestInfo: Seq[(TopicAndPartition, PartitionFetchInfo)]) {
-    this(versionId = FetchRequest.CurrentVersion,
-      correlationId = correlationId,
-      clientId = clientId,
-      replicaId = Request.OrdinaryConsumerId,
-      maxWait = maxWait,
-      minBytes = minBytes,
-      maxBytes = maxBytes,
-      requestInfo = requestInfo)
-  }
-
-  def writeTo(buffer: ByteBuffer) {
-    buffer.putShort(versionId)
-    buffer.putInt(correlationId)
-    writeShortString(buffer, clientId)
-    buffer.putInt(replicaId)
-    buffer.putInt(maxWait)
-    buffer.putInt(minBytes)
-    if (versionId >= 3)
-      buffer.putInt(maxBytes)
-    buffer.putInt(requestInfoGroupedByTopic.size) // topic count
-    requestInfoGroupedByTopic.foreach {
-      case (topic, partitionFetchInfos) =>
-        writeShortString(buffer, topic)
-        buffer.putInt(partitionFetchInfos.size) // partition count
-        partitionFetchInfos.foreach {
-          case (partition, PartitionFetchInfo(offset, fetchSize)) =>
-            buffer.putInt(partition)
-            buffer.putLong(offset)
-            buffer.putInt(fetchSize)
-        }
-    }
-  }
-
-  def sizeInBytes: Int = {
-    2 + /* versionId */
-    4 + /* correlationId */
-    shortStringLength(clientId) +
-    4 + /* replicaId */
-    4 + /* maxWait */
-    4 + /* minBytes */
-    (if (versionId >= 3) 4 /* maxBytes */ else 0) +
-    4 + /* topic count */
-    requestInfoGroupedByTopic.foldLeft(0)((foldedTopics, currTopic) => {
-      val (topic, partitionFetchInfos) = currTopic
-      foldedTopics +
-      shortStringLength(topic) +
-      4 + /* partition count */
-      partitionFetchInfos.size * (
-        4 + /* partition id */
-        8 + /* offset */
-        4 /* fetch size */
-      )
-    })
-  }
-
-  def isFromFollower = Request.isValidBrokerId(replicaId)
-
-  def isFromOrdinaryConsumer = replicaId == Request.OrdinaryConsumerId
-
-  def isFromLowLevelConsumer = replicaId == Request.DebuggingConsumerId
-
-  def numPartitions = requestInfo.size
-
-  override def toString: String = {
-    describe(true)
-  }
-
-  override def describe(details: Boolean): String = {
-    val fetchRequest = new StringBuilder
-    fetchRequest.append("Name: " + this.getClass.getSimpleName)
-    fetchRequest.append("; Version: " + versionId)
-    fetchRequest.append("; CorrelationId: " + correlationId)
-    fetchRequest.append("; ClientId: " + clientId)
-    fetchRequest.append("; ReplicaId: " + replicaId)
-    fetchRequest.append("; MaxWait: " + maxWait + " ms")
-    fetchRequest.append("; MinBytes: " + minBytes + " bytes")
-    fetchRequest.append("; MaxBytes:" + maxBytes + " bytes")
-    if(details)
-      fetchRequest.append("; RequestInfo: " + requestInfo.mkString(","))
-    fetchRequest.toString()
-  }
-}
-
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.11.0.0")
-@nonthreadsafe
-class FetchRequestBuilder() {
-  private val correlationId = new AtomicInteger(0)
-  private var versionId = FetchRequest.CurrentVersion
-  private var clientId = ConsumerConfig.DefaultClientId
-  private var replicaId = Request.OrdinaryConsumerId
-  private var maxWait = FetchRequest.DefaultMaxWait
-  private var minBytes = FetchRequest.DefaultMinBytes
-  private var maxBytes = FetchRequest.DefaultMaxBytes
-  private val requestMap = new collection.mutable.ArrayBuffer[(TopicAndPartition, PartitionFetchInfo)]
-
-  def addFetch(topic: String, partition: Int, offset: Long, fetchSize: Int) = {
-    requestMap.append((TopicAndPartition(topic, partition), PartitionFetchInfo(offset, fetchSize)))
-    this
-  }
-
-  def clientId(clientId: String): FetchRequestBuilder = {
-    this.clientId = clientId
-    this
-  }
-
-  /**
-   * Only for internal use. Clients shouldn't set replicaId.
-   */
-  private[kafka] def replicaId(replicaId: Int): FetchRequestBuilder = {
-    this.replicaId = replicaId
-    this
-  }
-
-  def maxWait(maxWait: Int): FetchRequestBuilder = {
-    this.maxWait = maxWait
-    this
-  }
-
-  def minBytes(minBytes: Int): FetchRequestBuilder = {
-    this.minBytes = minBytes
-    this
-  }
-
-  def maxBytes(maxBytes: Int): FetchRequestBuilder = {
-    this.maxBytes = maxBytes
-    this
-  }
-
-  def requestVersion(versionId: Short): FetchRequestBuilder = {
-    this.versionId = versionId
-    this
-  }
-
-  def build() = {
-    val fetchRequest = FetchRequest(versionId, correlationId.getAndIncrement, clientId, replicaId, maxWait, minBytes,
-      maxBytes, new ArrayBuffer() ++ requestMap)
-    requestMap.clear()
-    fetchRequest
-  }
-}
diff --git a/core/src/main/scala/kafka/api/FetchResponse.scala b/core/src/main/scala/kafka/api/FetchResponse.scala
deleted file mode 100644
index ae2f19c4bcc1f..0000000000000
--- a/core/src/main/scala/kafka/api/FetchResponse.scala
+++ /dev/null
@@ -1,174 +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.nio.ByteBuffer
-
-import kafka.common.TopicAndPartition
-import kafka.message.{ByteBufferMessageSet, MessageSet}
-import kafka.api.ApiUtils._
-import org.apache.kafka.common.protocol.Errors
-
-import scala.collection._
-
-object FetchResponsePartitionData {
-  def readFrom(buffer: ByteBuffer): FetchResponsePartitionData = {
-    val error = Errors.forCode(buffer.getShort)
-    val hw = buffer.getLong
-    val messageSetSize = buffer.getInt
-    val messageSetBuffer = buffer.slice()
-    messageSetBuffer.limit(messageSetSize)
-    buffer.position(buffer.position() + messageSetSize)
-    new FetchResponsePartitionData(error, hw, new ByteBufferMessageSet(messageSetBuffer))
-  }
-
-  val headerSize =
-    2 + /* error code */
-    8 + /* high watermark */
-    4 /* messageSetSize */
-}
-
-case class FetchResponsePartitionData(error: Errors = Errors.NONE, hw: Long = -1L, messages: MessageSet) {
-  val sizeInBytes = FetchResponsePartitionData.headerSize + messages.sizeInBytes
-}
-
-object TopicData {
-  def readFrom(buffer: ByteBuffer): TopicData = {
-    val topic = readShortString(buffer)
-    val partitionCount = buffer.getInt
-    val topicPartitionDataPairs = (1 to partitionCount).map(_ => {
-      val partitionId = buffer.getInt
-      val partitionData = FetchResponsePartitionData.readFrom(buffer)
-      (partitionId, partitionData)
-    })
-    TopicData(topic, Seq(topicPartitionDataPairs:_*))
-  }
-
-  def headerSize(topic: String) =
-    shortStringLength(topic) +
-    4 /* partition count */
-}
-
-case class TopicData(topic: String, partitionData: Seq[(Int, FetchResponsePartitionData)]) {
-  val sizeInBytes =
-    TopicData.headerSize(topic) + partitionData.foldLeft(0)((folded, data) => {
-      folded + data._2.sizeInBytes + 4
-    }                                  /*_ + _.sizeInBytes + 4*/)
-
-  val headerSize = TopicData.headerSize(topic)
-}
-
-@deprecated("This object has been deprecated and will be removed in a future release.", "1.0.0")
-object FetchResponse {
-
-  // The request version is used to determine which fields we can expect in the response
-  def readFrom(buffer: ByteBuffer, requestVersion: Int): FetchResponse = {
-    val correlationId = buffer.getInt
-    val throttleTime = if (requestVersion > 0) buffer.getInt else 0
-    val topicCount = buffer.getInt
-    val pairs = (1 to topicCount).flatMap(_ => {
-      val topicData = TopicData.readFrom(buffer)
-      topicData.partitionData.map { case (partitionId, partitionData) =>
-        (TopicAndPartition(topicData.topic, partitionId), partitionData)
-      }
-    })
-    FetchResponse(correlationId, Vector(pairs:_*), requestVersion, throttleTime)
-  }
-
-  type FetchResponseEntry = (Int, FetchResponsePartitionData)
-
-  def batchByTopic(data: Seq[(TopicAndPartition, FetchResponsePartitionData)]): Seq[(String, Seq[FetchResponseEntry])] =
-    FetchRequest.batchByTopic(data)
-
-  // Returns the size of the response header
-  def headerSize(requestVersion: Int): Int = {
-    val throttleTimeSize = if (requestVersion > 0) 4 else 0
-    4 + /* correlationId */
-    4 + /* topic count */
-    throttleTimeSize
-  }
-
-  // Returns the size of entire fetch response in bytes (including the header size)
-  def responseSize(dataGroupedByTopic: Seq[(String, Seq[FetchResponseEntry])],
-                   requestVersion: Int): Int = {
-    headerSize(requestVersion) +
-    dataGroupedByTopic.foldLeft(0) { case (folded, (topic, partitionDataMap)) =>
-      val topicData = TopicData(topic, partitionDataMap.map {
-        case (partitionId, partitionData) => (partitionId, partitionData)
-      })
-      folded + topicData.sizeInBytes
-    }
-  }
-}
-
-@deprecated("This object has been deprecated and will be removed in a future release.", "1.0.0")
-case class FetchResponse(correlationId: Int,
-                         data: Seq[(TopicAndPartition, FetchResponsePartitionData)],
-                         requestVersion: Int = 0,
-                         throttleTimeMs: Int = 0)
-  extends RequestOrResponse() {
-
-  /**
-   * Partitions the data into a map of maps (one for each topic).
-   */
-  private lazy val dataByTopicAndPartition = data.toMap
-  lazy val dataGroupedByTopic = FetchResponse.batchByTopic(data)
-  val headerSizeInBytes = FetchResponse.headerSize(requestVersion)
-  lazy val sizeInBytes = FetchResponse.responseSize(dataGroupedByTopic, requestVersion)
-
-  /*
-   * Writes the header of the FetchResponse to the input buffer
-   */
-  def writeHeaderTo(buffer: ByteBuffer) = {
-    buffer.putInt(sizeInBytes)
-    buffer.putInt(correlationId)
-    // Include the throttleTime only if the client can read it
-    if (requestVersion > 0)
-      buffer.putInt(throttleTimeMs)
-
-    buffer.putInt(dataGroupedByTopic.size) // topic count
-  }
-  /*
-   * FetchResponse uses [sendfile](http://man7.org/linux/man-pages/man2/sendfile.2.html)
-   * api for data transfer through the FetchResponseSend, so `writeTo` aren't actually being used.
-   * It is implemented as an empty function to conform to `RequestOrResponse.writeTo`
-   * abstract method signature.
-   */
-  def writeTo(buffer: ByteBuffer): Unit = throw new UnsupportedOperationException
-
-  override def describe(details: Boolean): String = toString
-
-  private def partitionDataFor(topic: String, partition: Int): FetchResponsePartitionData = {
-    val topicAndPartition = TopicAndPartition(topic, partition)
-    dataByTopicAndPartition.get(topicAndPartition) match {
-      case Some(partitionData) => partitionData
-      case _ =>
-        throw new IllegalArgumentException(
-          "No partition %s in fetch response %s".format(topicAndPartition, this.toString))
-    }
-  }
-
-  def messageSet(topic: String, partition: Int): ByteBufferMessageSet =
-    partitionDataFor(topic, partition).messages.asInstanceOf[ByteBufferMessageSet]
-
-  def highWatermark(topic: String, partition: Int) = partitionDataFor(topic, partition).hw
-
-  def hasError = dataByTopicAndPartition.values.exists(_.error != Errors.NONE)
-
-  def error(topic: String, partition: Int) = partitionDataFor(topic, partition).error
-}
diff --git a/core/src/main/scala/kafka/api/GroupCoordinatorRequest.scala b/core/src/main/scala/kafka/api/GroupCoordinatorRequest.scala
deleted file mode 100644
index 6ee6ae73ffecf..0000000000000
--- a/core/src/main/scala/kafka/api/GroupCoordinatorRequest.scala
+++ /dev/null
@@ -1,74 +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.nio.ByteBuffer
-
-import org.apache.kafka.common.protocol.ApiKeys
-
-@deprecated("This object has been deprecated and will be removed in a future release.", "1.0.0")
-object GroupCoordinatorRequest {
-  val CurrentVersion = 0.shortValue
-  val DefaultClientId = ""
-
-  def readFrom(buffer: ByteBuffer) = {
-    // envelope
-    val versionId = buffer.getShort
-    val correlationId = buffer.getInt
-    val clientId = ApiUtils.readShortString(buffer)
-
-    // request
-    val group = ApiUtils.readShortString(buffer)
-    GroupCoordinatorRequest(group, versionId, correlationId, clientId)
-  }
-
-}
-
-@deprecated("This object has been deprecated and will be removed in a future release.", "1.0.0")
-case class GroupCoordinatorRequest(group: String,
-                                   versionId: Short = GroupCoordinatorRequest.CurrentVersion,
-                                   correlationId: Int = 0,
-                                   clientId: String = GroupCoordinatorRequest.DefaultClientId)
-  extends RequestOrResponse(Some(ApiKeys.FIND_COORDINATOR.id)) {
-
-  def sizeInBytes =
-    2 + /* versionId */
-    4 + /* correlationId */
-    ApiUtils.shortStringLength(clientId) +
-    ApiUtils.shortStringLength(group)
-
-  def writeTo(buffer: ByteBuffer) {
-    // envelope
-    buffer.putShort(versionId)
-    buffer.putInt(correlationId)
-    ApiUtils.writeShortString(buffer, clientId)
-
-    // consumer metadata request
-    ApiUtils.writeShortString(buffer, group)
-  }
-
-  def describe(details: Boolean) = {
-    val consumerMetadataRequest = new StringBuilder
-    consumerMetadataRequest.append("Name: " + this.getClass.getSimpleName)
-    consumerMetadataRequest.append("; Version: " + versionId)
-    consumerMetadataRequest.append("; CorrelationId: " + correlationId)
-    consumerMetadataRequest.append("; ClientId: " + clientId)
-    consumerMetadataRequest.append("; Group: " + group)
-    consumerMetadataRequest.toString()
-  }
-}
diff --git a/core/src/main/scala/kafka/api/GroupCoordinatorResponse.scala b/core/src/main/scala/kafka/api/GroupCoordinatorResponse.scala
deleted file mode 100644
index 7fd5fe3de6775..0000000000000
--- a/core/src/main/scala/kafka/api/GroupCoordinatorResponse.scala
+++ /dev/null
@@ -1,60 +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.nio.ByteBuffer
-import kafka.cluster.BrokerEndPoint
-import org.apache.kafka.common.protocol.Errors
-
-@deprecated("This object has been deprecated and will be removed in a future release.", "1.0.0")
-object GroupCoordinatorResponse {
-  val CurrentVersion = 0
-
-  private val NoBrokerEndpointOpt = Some(BrokerEndPoint(id = -1, host = "", port = -1))
-
-  def readFrom(buffer: ByteBuffer) = {
-    val correlationId = buffer.getInt
-    val error = Errors.forCode(buffer.getShort)
-    val broker = BrokerEndPoint.readFrom(buffer)
-    val coordinatorOpt = if (error == Errors.NONE)
-      Some(broker)
-    else
-      None
-
-    GroupCoordinatorResponse(coordinatorOpt, error, correlationId)
-  }
-
-}
-
-@deprecated("This object has been deprecated and will be removed in a future release.", "1.0.0")
-case class GroupCoordinatorResponse (coordinatorOpt: Option[BrokerEndPoint], error: Errors, correlationId: Int)
-  extends RequestOrResponse() {
-
-  def sizeInBytes =
-    4 + /* correlationId */
-    2 + /* error code */
-    coordinatorOpt.orElse(GroupCoordinatorResponse.NoBrokerEndpointOpt).get.sizeInBytes
-
-  def writeTo(buffer: ByteBuffer) {
-    buffer.putInt(correlationId)
-    buffer.putShort(error.code)
-    coordinatorOpt.orElse(GroupCoordinatorResponse.NoBrokerEndpointOpt).foreach(_.writeTo(buffer))
-  }
-
-  def describe(details: Boolean) = toString
-}
diff --git a/core/src/main/scala/kafka/api/LeaderAndIsr.scala b/core/src/main/scala/kafka/api/LeaderAndIsr.scala
index 7a83cf3321b92..05952aa58d3a9 100644
--- a/core/src/main/scala/kafka/api/LeaderAndIsr.scala
+++ b/core/src/main/scala/kafka/api/LeaderAndIsr.scala
@@ -17,13 +17,13 @@
 
 package kafka.api
 
-import kafka.utils._
-
 object LeaderAndIsr {
   val initialLeaderEpoch: Int = 0
   val initialZKVersion: Int = 0
   val NoLeader: Int = -1
+  val NoEpoch: Int = -1
   val LeaderDuringDelete: Int = -2
+  val EpochDuringDelete: Int = -2
 
   def apply(leader: Int, isr: List[Int]): LeaderAndIsr = LeaderAndIsr(leader, initialLeaderEpoch, isr, initialZKVersion)
 
@@ -42,7 +42,21 @@ case class LeaderAndIsr(leader: Int,
 
   def newEpochAndZkVersion = newLeaderAndIsr(leader, isr)
 
+  def leaderOpt: Option[Int] = {
+    if (leader == LeaderAndIsr.NoLeader) None else Some(leader)
+  }
+
+  def equalsIgnoreZk(other: LeaderAndIsr): Boolean = {
+    if (this == other) {
+      true
+    } else if (other == null) {
+      false
+    } else {
+      leader == other.leader && leaderEpoch == other.leaderEpoch && isr.equals(other.isr)
+    }
+  }
+
   override def toString: String = {
-    Json.encode(Map("leader" -> leader, "leader_epoch" -> leaderEpoch, "isr" -> isr))
+    s"LeaderAndIsr(leader=$leader, leaderEpoch=$leaderEpoch, isr=$isr, zkVersion=$zkVersion)"
   }
 }
diff --git a/core/src/main/scala/kafka/api/OffsetCommitRequest.scala b/core/src/main/scala/kafka/api/OffsetCommitRequest.scala
deleted file mode 100644
index bffcec3fad49f..0000000000000
--- a/core/src/main/scala/kafka/api/OffsetCommitRequest.scala
+++ /dev/null
@@ -1,181 +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.nio.ByteBuffer
-
-import kafka.api.ApiUtils._
-import kafka.common.{OffsetAndMetadata, TopicAndPartition}
-import kafka.utils.Logging
-import org.apache.kafka.common.protocol.ApiKeys
-
-import scala.collection._
-
-@deprecated("This object has been deprecated and will be removed in a future release.", "1.0.0")
-object OffsetCommitRequest extends Logging {
-  val CurrentVersion: Short = 2
-  val DefaultClientId = ""
-
-  def readFrom(buffer: ByteBuffer): OffsetCommitRequest = {
-    // Read values from the envelope
-    val versionId = buffer.getShort
-    assert(versionId == 0 || versionId == 1 || versionId == 2,
-           "Version " + versionId + " is invalid for OffsetCommitRequest. Valid versions are 0, 1 or 2.")
-
-    val correlationId = buffer.getInt
-    val clientId = readShortString(buffer)
-
-    // Read the OffsetRequest 
-    val groupId = readShortString(buffer)
-
-    // version 1 and 2 specific fields
-    val groupGenerationId: Int =
-      if (versionId >= 1)
-        buffer.getInt
-      else
-        org.apache.kafka.common.requests.OffsetCommitRequest.DEFAULT_GENERATION_ID
-
-    val memberId: String =
-      if (versionId >= 1)
-        readShortString(buffer)
-      else
-        org.apache.kafka.common.requests.OffsetCommitRequest.DEFAULT_MEMBER_ID
-
-    // version 2 specific fields
-    val retentionMs: Long =
-      if (versionId >= 2)
-        buffer.getLong
-      else
-        org.apache.kafka.common.requests.OffsetCommitRequest.DEFAULT_RETENTION_TIME
-
-    val topicCount = buffer.getInt
-    val pairs = (1 to topicCount).flatMap(_ => {
-      val topic = readShortString(buffer)
-      val partitionCount = buffer.getInt
-      (1 to partitionCount).map(_ => {
-        val partitionId = buffer.getInt
-        val offset = buffer.getLong
-        val timestamp = {
-          // version 1 specific field
-          if (versionId == 1)
-            buffer.getLong
-          else
-            org.apache.kafka.common.requests.OffsetCommitRequest.DEFAULT_TIMESTAMP
-        }
-        val metadata = readShortString(buffer)
-
-        (TopicAndPartition(topic, partitionId), OffsetAndMetadata(offset, metadata, timestamp))
-      })
-    })
-
-    OffsetCommitRequest(groupId, immutable.Map(pairs:_*), versionId, correlationId, clientId, groupGenerationId, memberId, retentionMs)
-  }
-}
-
-@deprecated("This object has been deprecated and will be removed in a future release.", "1.0.0")
-case class OffsetCommitRequest(groupId: String,
-                               requestInfo: immutable.Map[TopicAndPartition, OffsetAndMetadata],
-                               versionId: Short = OffsetCommitRequest.CurrentVersion,
-                               correlationId: Int = 0,
-                               clientId: String = OffsetCommitRequest.DefaultClientId,
-                               groupGenerationId: Int = org.apache.kafka.common.requests.OffsetCommitRequest.DEFAULT_GENERATION_ID,
-                               memberId: String =  org.apache.kafka.common.requests.OffsetCommitRequest.DEFAULT_MEMBER_ID,
-                               retentionMs: Long = org.apache.kafka.common.requests.OffsetCommitRequest.DEFAULT_RETENTION_TIME)
-    extends RequestOrResponse(Some(ApiKeys.OFFSET_COMMIT.id)) {
-
-  assert(versionId == 0 || versionId == 1 || versionId == 2,
-         "Version " + versionId + " is invalid for OffsetCommitRequest. Valid versions are 0, 1 or 2.")
-
-  lazy val requestInfoGroupedByTopic = requestInfo.groupBy(_._1.topic)
-
-  def writeTo(buffer: ByteBuffer) {
-    // Write envelope
-    buffer.putShort(versionId)
-    buffer.putInt(correlationId)
-    writeShortString(buffer, clientId)
-
-    // Write OffsetCommitRequest
-    writeShortString(buffer, groupId)             // consumer group
-
-    // version 1 and 2 specific data
-    if (versionId >= 1) {
-      buffer.putInt(groupGenerationId)
-      writeShortString(buffer, memberId)
-    }
-
-    // version 2 or above specific data
-    if (versionId >= 2) {
-      buffer.putLong(retentionMs)
-    }
-
-    buffer.putInt(requestInfoGroupedByTopic.size) // number of topics
-    requestInfoGroupedByTopic.foreach( t1 => { // topic -> Map[TopicAndPartition, OffsetMetadataAndError]
-      writeShortString(buffer, t1._1) // topic
-      buffer.putInt(t1._2.size)       // number of partitions for this topic
-      t1._2.foreach( t2 => {
-        buffer.putInt(t2._1.partition)
-        buffer.putLong(t2._2.offset)
-        // version 1 specific data
-        if (versionId == 1)
-          buffer.putLong(t2._2.commitTimestamp)
-        writeShortString(buffer, t2._2.metadata)
-      })
-    })
-  }
-
-  override def sizeInBytes =
-    2 + /* versionId */
-    4 + /* correlationId */
-    shortStringLength(clientId) +
-    shortStringLength(groupId) +
-    (if (versionId >= 1) 4 /* group generation id */ + shortStringLength(memberId) else 0) +
-    (if (versionId >= 2) 8 /* retention time */ else 0) +
-    4 + /* topic count */
-    requestInfoGroupedByTopic.foldLeft(0)((count, topicAndOffsets) => {
-      val (topic, offsets) = topicAndOffsets
-      count +
-      shortStringLength(topic) + /* topic */
-      4 + /* number of partitions */
-      offsets.foldLeft(0)((innerCount, offsetAndMetadata) => {
-        innerCount +
-        4 /* partition */ +
-        8 /* offset */ +
-        (if (versionId == 1) 8 else 0) /* timestamp */ +
-        shortStringLength(offsetAndMetadata._2.metadata)
-      })
-    })
-
-  override def describe(details: Boolean): String = {
-    val offsetCommitRequest = new StringBuilder
-    offsetCommitRequest.append("Name: " + this.getClass.getSimpleName)
-    offsetCommitRequest.append("; Version: " + versionId)
-    offsetCommitRequest.append("; CorrelationId: " + correlationId)
-    offsetCommitRequest.append("; ClientId: " + clientId)
-    offsetCommitRequest.append("; GroupId: " + groupId)
-    offsetCommitRequest.append("; GroupGenerationId: " + groupGenerationId)
-    offsetCommitRequest.append("; MemberId: " + memberId)
-    offsetCommitRequest.append("; RetentionMs: " + retentionMs)
-    if(details)
-      offsetCommitRequest.append("; RequestInfo: " + requestInfo.mkString(","))
-    offsetCommitRequest.toString()
-  }
-
-  override def toString = {
-    describe(details = true)
-  }
-}
diff --git a/core/src/main/scala/kafka/api/OffsetCommitResponse.scala b/core/src/main/scala/kafka/api/OffsetCommitResponse.scala
deleted file mode 100644
index 07adc76234f8e..0000000000000
--- a/core/src/main/scala/kafka/api/OffsetCommitResponse.scala
+++ /dev/null
@@ -1,82 +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.nio.ByteBuffer
-
-import kafka.utils.Logging
-import kafka.common.TopicAndPartition
-import org.apache.kafka.common.protocol.Errors
-
-@deprecated("This object has been deprecated and will be removed in a future release.", "1.0.0")
-object OffsetCommitResponse extends Logging {
-  val CurrentVersion: Short = 0
-
-  def readFrom(buffer: ByteBuffer): OffsetCommitResponse = {
-    val correlationId = buffer.getInt
-    val topicCount = buffer.getInt
-    val pairs = (1 to topicCount).flatMap(_ => {
-      val topic = ApiUtils.readShortString(buffer)
-      val partitionCount = buffer.getInt
-      (1 to partitionCount).map(_ => {
-        val partitionId = buffer.getInt
-        val error = Errors.forCode(buffer.getShort)
-        (TopicAndPartition(topic, partitionId), error)
-      })
-    })
-    OffsetCommitResponse(Map(pairs:_*), correlationId)
-  }
-}
-
-@deprecated("This object has been deprecated and will be removed in a future release.", "1.0.0")
-case class OffsetCommitResponse(commitStatus: Map[TopicAndPartition, Errors],
-                                correlationId: Int = 0)
-    extends RequestOrResponse() {
-
-  lazy val commitStatusGroupedByTopic = commitStatus.groupBy(_._1.topic)
-
-  def hasError = commitStatus.values.exists(_ != Errors.NONE)
-
-  def writeTo(buffer: ByteBuffer) {
-    buffer.putInt(correlationId)
-    buffer.putInt(commitStatusGroupedByTopic.size)
-    commitStatusGroupedByTopic.foreach { case(topic, statusMap) =>
-      ApiUtils.writeShortString(buffer, topic)
-      buffer.putInt(statusMap.size) // partition count
-      statusMap.foreach { case(topicAndPartition, error) =>
-        buffer.putInt(topicAndPartition.partition)
-        buffer.putShort(error.code)
-      }
-    }
-  }
-
-  override def sizeInBytes =
-    4 + /* correlationId */
-    4 + /* topic count */
-    commitStatusGroupedByTopic.foldLeft(0)((count, partitionStatusMap) => {
-      val (topic, partitionStatus) = partitionStatusMap
-      count +
-      ApiUtils.shortStringLength(topic) +
-      4 + /* partition count */
-      partitionStatus.size * ( 4 /* partition */  + 2 /* error code */)
-    })
-
-  override def describe(details: Boolean):String = { toString }
-
-}
-
diff --git a/core/src/main/scala/kafka/api/OffsetFetchRequest.scala b/core/src/main/scala/kafka/api/OffsetFetchRequest.scala
deleted file mode 100644
index c24078d52fe06..0000000000000
--- a/core/src/main/scala/kafka/api/OffsetFetchRequest.scala
+++ /dev/null
@@ -1,108 +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.nio.ByteBuffer
-
-import kafka.api.ApiUtils._
-import kafka.common.TopicAndPartition
-import kafka.utils.Logging
-import org.apache.kafka.common.protocol.ApiKeys
-
-@deprecated("This object has been deprecated and will be removed in a future release.", "1.0.0")
-object OffsetFetchRequest extends Logging {
-  val CurrentVersion: Short = 2
-  val DefaultClientId = ""
-
-  def readFrom(buffer: ByteBuffer): OffsetFetchRequest = {
-    // Read values from the envelope
-    val versionId = buffer.getShort
-    val correlationId = buffer.getInt
-    val clientId = readShortString(buffer)
-
-    // Read the OffsetFetchRequest
-    val consumerGroupId = readShortString(buffer)
-    val topicCount = buffer.getInt
-    val pairs = (1 to topicCount).flatMap(_ => {
-      val topic = readShortString(buffer)
-      val partitionCount = buffer.getInt
-      (1 to partitionCount).map(_ => {
-        val partitionId = buffer.getInt
-        TopicAndPartition(topic, partitionId)
-      })
-    })
-    OffsetFetchRequest(consumerGroupId, pairs, versionId, correlationId, clientId)
-  }
-}
-
-@deprecated("This object has been deprecated and will be removed in a future release.", "1.0.0")
-case class OffsetFetchRequest(groupId: String,
-                              requestInfo: Seq[TopicAndPartition],
-                              versionId: Short = OffsetFetchRequest.CurrentVersion,
-                              correlationId: Int = 0,
-                              clientId: String = OffsetFetchRequest.DefaultClientId)
-    extends RequestOrResponse(Some(ApiKeys.OFFSET_FETCH.id)) {
-
-  lazy val requestInfoGroupedByTopic = requestInfo.groupBy(_.topic)
-
-  def writeTo(buffer: ByteBuffer) {
-    // Write envelope
-    buffer.putShort(versionId)
-    buffer.putInt(correlationId)
-    writeShortString(buffer, clientId)
-
-    // Write OffsetFetchRequest
-    writeShortString(buffer, groupId)             // consumer group
-    buffer.putInt(requestInfoGroupedByTopic.size) // number of topics
-    requestInfoGroupedByTopic.foreach( t1 => { // (topic, Seq[TopicAndPartition])
-      writeShortString(buffer, t1._1) // topic
-      buffer.putInt(t1._2.size)       // number of partitions for this topic
-      t1._2.foreach( t2 => {
-        buffer.putInt(t2.partition)
-      })
-    })
-  }
-
-  override def sizeInBytes =
-    2 + /* versionId */
-    4 + /* correlationId */
-    shortStringLength(clientId) +
-    shortStringLength(groupId) +
-    4 + /* topic count */
-    requestInfoGroupedByTopic.foldLeft(0)((count, t) => {
-      count + shortStringLength(t._1) + /* topic */
-      4 + /* number of partitions */
-      t._2.size * 4 /* partition */
-    })
-
-  override def describe(details: Boolean): String = {
-    val offsetFetchRequest = new StringBuilder
-    offsetFetchRequest.append("Name: " + this.getClass.getSimpleName)
-    offsetFetchRequest.append("; Version: " + versionId)
-    offsetFetchRequest.append("; CorrelationId: " + correlationId)
-    offsetFetchRequest.append("; ClientId: " + clientId)
-    offsetFetchRequest.append("; GroupId: " + groupId)
-    if (details)
-      offsetFetchRequest.append("; RequestInfo: " + requestInfo.mkString(","))
-    offsetFetchRequest.toString()
-  }
-
-  override def toString: String = {
-    describe(details = true)
-  }
-}
diff --git a/core/src/main/scala/kafka/api/OffsetFetchResponse.scala b/core/src/main/scala/kafka/api/OffsetFetchResponse.scala
deleted file mode 100644
index b875dcc5a01e2..0000000000000
--- a/core/src/main/scala/kafka/api/OffsetFetchResponse.scala
+++ /dev/null
@@ -1,107 +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.nio.ByteBuffer
-
-import kafka.api.ApiUtils._
-import kafka.common.{TopicAndPartition, OffsetMetadataAndError}
-import kafka.utils.Logging
-
-import org.apache.kafka.common.protocol.Errors
-
-@deprecated("This object has been deprecated and will be removed in a future release.", "1.0.0")
-object OffsetFetchResponse extends Logging {
-
-  def readFrom(buffer: ByteBuffer): OffsetFetchResponse = {
-    readFrom(buffer, OffsetFetchRequest.CurrentVersion)
-  }
-
-  def readFrom(buffer: ByteBuffer, requestVersion: Int): OffsetFetchResponse = {
-    val correlationId = buffer.getInt
-    val topicCount = buffer.getInt
-    val pairs = (1 to topicCount).flatMap(_ => {
-      val topic = readShortString(buffer)
-      val partitionCount = buffer.getInt
-      (1 to partitionCount).map(_ => {
-        val partitionId = buffer.getInt
-        val offset = buffer.getLong
-        val metadata = readShortString(buffer)
-        val error = Errors.forCode(buffer.getShort)
-        (TopicAndPartition(topic, partitionId), OffsetMetadataAndError(offset, metadata, error))
-      })
-    })
-
-    val error = requestVersion match {
-      case 0 | 1 => Errors.NONE
-      case _ => Errors.forCode(buffer.getShort)
-    }
-
-    OffsetFetchResponse(Map(pairs:_*), requestVersion, correlationId, error)
-  }
-}
-
-@deprecated("This object has been deprecated and will be removed in a future release.", "1.0.0")
-case class OffsetFetchResponse(requestInfo: Map[TopicAndPartition, OffsetMetadataAndError],
-                               requestVersion: Int = OffsetFetchRequest.CurrentVersion,
-                               correlationId: Int = 0,
-                               error: Errors = Errors.NONE)
-    extends RequestOrResponse() {
-
-  lazy val requestInfoGroupedByTopic = requestInfo.groupBy(_._1.topic)
-
-  def writeTo(buffer: ByteBuffer) {
-    buffer.putInt(correlationId)
-    buffer.putInt(requestInfoGroupedByTopic.size) // number of topics
-    requestInfoGroupedByTopic.foreach( t1 => { // topic -> Map[TopicAndPartition, OffsetMetadataAndError]
-      writeShortString(buffer, t1._1) // topic
-      buffer.putInt(t1._2.size)       // number of partitions for this topic
-      t1._2.foreach( t2 => { // TopicAndPartition -> OffsetMetadataAndError
-        buffer.putInt(t2._1.partition)
-        buffer.putLong(t2._2.offset)
-        writeShortString(buffer, t2._2.metadata)
-        buffer.putShort(t2._2.error.code)
-      })
-    })
-
-    // the top level error_code was introduced in v2
-    if (requestVersion > 1)
-      buffer.putShort(error.code)
-  }
-
-  override def sizeInBytes =
-    4 + /* correlationId */
-    4 + /* topic count */
-    requestInfoGroupedByTopic.foldLeft(0)((count, topicAndOffsets) => {
-      val (topic, offsets) = topicAndOffsets
-      count +
-      shortStringLength(topic) + /* topic */
-      4 + /* number of partitions */
-      offsets.foldLeft(0)((innerCount, offsetsAndMetadata) => {
-        innerCount +
-        4 /* partition */ +
-        8 /* offset */ +
-        shortStringLength(offsetsAndMetadata._2.metadata) +
-        2 /* error */
-      })
-    }) +
-    (if (requestVersion > 1) 2 else 0) /* error */
-
-  override def describe(details: Boolean):String = { toString }
-}
-
diff --git a/core/src/main/scala/kafka/api/OffsetRequest.scala b/core/src/main/scala/kafka/api/OffsetRequest.scala
deleted file mode 100644
index f5483b1acf7e2..0000000000000
--- a/core/src/main/scala/kafka/api/OffsetRequest.scala
+++ /dev/null
@@ -1,126 +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.nio.ByteBuffer
-
-import kafka.api.ApiUtils._
-import kafka.common.TopicAndPartition
-import org.apache.kafka.common.protocol.ApiKeys
-
-@deprecated("This object has been deprecated and will be removed in a future release.", "1.0.0")
-object OffsetRequest {
-  val CurrentVersion = 0.shortValue
-  val DefaultClientId = ""
-
-  val SmallestTimeString = "smallest"
-  val LargestTimeString = "largest"
-  val LatestTime = -1L
-  val EarliestTime = -2L
-
-  def readFrom(buffer: ByteBuffer): OffsetRequest = {
-    val versionId = buffer.getShort
-    val correlationId = buffer.getInt
-    val clientId = readShortString(buffer)
-    val replicaId = buffer.getInt
-    val topicCount = buffer.getInt
-    val pairs = (1 to topicCount).flatMap(_ => {
-      val topic = readShortString(buffer)
-      val partitionCount = buffer.getInt
-      (1 to partitionCount).map(_ => {
-        val partitionId = buffer.getInt
-        val time = buffer.getLong
-        val maxNumOffsets = buffer.getInt
-        (TopicAndPartition(topic, partitionId), PartitionOffsetRequestInfo(time, maxNumOffsets))
-      })
-    })
-    OffsetRequest(Map(pairs:_*), versionId= versionId, clientId = clientId, correlationId = correlationId, replicaId = replicaId)
-  }
-}
-
-case class PartitionOffsetRequestInfo(time: Long, maxNumOffsets: Int)
-
-@deprecated("This object has been deprecated and will be removed in a future release.", "1.0.0")
-case class OffsetRequest(requestInfo: Map[TopicAndPartition, PartitionOffsetRequestInfo],
-                         versionId: Short = OffsetRequest.CurrentVersion,
-                         correlationId: Int = 0,
-                         clientId: String = OffsetRequest.DefaultClientId,
-                         replicaId: Int = Request.OrdinaryConsumerId)
-    extends RequestOrResponse(Some(ApiKeys.LIST_OFFSETS.id)) {
-
-  def this(requestInfo: Map[TopicAndPartition, PartitionOffsetRequestInfo], correlationId: Int, replicaId: Int) = this(requestInfo, OffsetRequest.CurrentVersion, correlationId, OffsetRequest.DefaultClientId, replicaId)
-
-  lazy val requestInfoGroupedByTopic = requestInfo.groupBy(_._1.topic)
-
-  def writeTo(buffer: ByteBuffer) {
-    buffer.putShort(versionId)
-    buffer.putInt(correlationId)
-    writeShortString(buffer, clientId)
-    buffer.putInt(replicaId)
-
-    buffer.putInt(requestInfoGroupedByTopic.size) // topic count
-    requestInfoGroupedByTopic.foreach {
-      case((topic, partitionInfos)) =>
-        writeShortString(buffer, topic)
-        buffer.putInt(partitionInfos.size) // partition count
-        partitionInfos.foreach {
-          case (TopicAndPartition(_, partition), partitionInfo) =>
-            buffer.putInt(partition)
-            buffer.putLong(partitionInfo.time)
-            buffer.putInt(partitionInfo.maxNumOffsets)
-        }
-    }
-  }
-
-  def sizeInBytes =
-    2 + /* versionId */
-    4 + /* correlationId */
-    shortStringLength(clientId) +
-    4 + /* replicaId */
-    4 + /* topic count */
-    requestInfoGroupedByTopic.foldLeft(0)((foldedTopics, currTopic) => {
-      val (topic, partitionInfos) = currTopic
-      foldedTopics +
-      shortStringLength(topic) +
-      4 + /* partition count */
-      partitionInfos.size * (
-        4 + /* partition */
-        8 + /* time */
-        4 /* maxNumOffsets */
-      )
-    })
-
-  def isFromOrdinaryClient = replicaId == Request.OrdinaryConsumerId
-  def isFromDebuggingClient = replicaId == Request.DebuggingConsumerId
-
-  override def toString: String = {
-    describe(true)
-  }
-
-  override def describe(details: Boolean): String = {
-    val offsetRequest = new StringBuilder
-    offsetRequest.append("Name: " + this.getClass.getSimpleName)
-    offsetRequest.append("; Version: " + versionId)
-    offsetRequest.append("; CorrelationId: " + correlationId)
-    offsetRequest.append("; ClientId: " + clientId)
-    offsetRequest.append("; ReplicaId: " + replicaId)
-    if(details)
-      offsetRequest.append("; RequestInfo: " + requestInfo.mkString(","))
-    offsetRequest.toString()
-  }
-}
diff --git a/core/src/main/scala/kafka/api/OffsetResponse.scala b/core/src/main/scala/kafka/api/OffsetResponse.scala
deleted file mode 100644
index 95a31d501b459..0000000000000
--- a/core/src/main/scala/kafka/api/OffsetResponse.scala
+++ /dev/null
@@ -1,102 +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.nio.ByteBuffer
-import kafka.common.TopicAndPartition
-import kafka.api.ApiUtils._
-import org.apache.kafka.common.protocol.Errors
-
-
-@deprecated("This object has been deprecated and will be removed in a future release.", "1.0.0")
-object OffsetResponse {
-
-  def readFrom(buffer: ByteBuffer): OffsetResponse = {
-    val correlationId = buffer.getInt
-    val numTopics = buffer.getInt
-    val pairs = (1 to numTopics).flatMap(_ => {
-      val topic = readShortString(buffer)
-      val numPartitions = buffer.getInt
-      (1 to numPartitions).map(_ => {
-        val partition = buffer.getInt
-        val error = Errors.forCode(buffer.getShort)
-        val numOffsets = buffer.getInt
-        val offsets = (1 to numOffsets).map(_ => buffer.getLong)
-        (TopicAndPartition(topic, partition), PartitionOffsetsResponse(error, offsets))
-      })
-    })
-    OffsetResponse(correlationId, Map(pairs:_*))
-  }
-
-}
-
-
-case class PartitionOffsetsResponse(error: Errors, offsets: Seq[Long]) {
-  override def toString: String = {
-    new String("error: " + error.exceptionName + " offsets: " + offsets.mkString)
-  }
-}
-
-
-@deprecated("This object has been deprecated and will be removed in a future release.", "1.0.0")
-case class OffsetResponse(correlationId: Int,
-                          partitionErrorAndOffsets: Map[TopicAndPartition, PartitionOffsetsResponse])
-    extends RequestOrResponse() {
-
-  lazy val offsetsGroupedByTopic = partitionErrorAndOffsets.groupBy(_._1.topic)
-
-  def hasError = partitionErrorAndOffsets.values.exists(_.error != Errors.NONE)
-
-  val sizeInBytes = {
-    4 + /* correlation id */
-    4 + /* topic count */
-    offsetsGroupedByTopic.foldLeft(0)((foldedTopics, currTopic) => {
-      val (topic, errorAndOffsetsMap) = currTopic
-      foldedTopics +
-      shortStringLength(topic) +
-      4 + /* partition count */
-      errorAndOffsetsMap.foldLeft(0)((foldedPartitions, currPartition) => {
-        foldedPartitions +
-        4 + /* partition id */
-        2 + /* partition error */
-        4 + /* offset array length */
-        currPartition._2.offsets.size * 8 /* offset */
-      })
-    })
-  }
-
-  def writeTo(buffer: ByteBuffer) {
-    buffer.putInt(correlationId)
-    buffer.putInt(offsetsGroupedByTopic.size) // topic count
-    offsetsGroupedByTopic.foreach {
-      case((topic, errorAndOffsetsMap)) =>
-        writeShortString(buffer, topic)
-        buffer.putInt(errorAndOffsetsMap.size) // partition count
-        errorAndOffsetsMap.foreach {
-          case((TopicAndPartition(_, partition), errorAndOffsets)) =>
-            buffer.putInt(partition)
-            buffer.putShort(errorAndOffsets.error.code)
-            buffer.putInt(errorAndOffsets.offsets.size) // offset array length
-            errorAndOffsets.offsets.foreach(buffer.putLong(_))
-        }
-    }
-  }
-
-  override def describe(details: Boolean):String = { toString }
-}
-
diff --git a/core/src/main/scala/kafka/api/ProducerRequest.scala b/core/src/main/scala/kafka/api/ProducerRequest.scala
deleted file mode 100644
index 9cdb14b6ba207..0000000000000
--- a/core/src/main/scala/kafka/api/ProducerRequest.scala
+++ /dev/null
@@ -1,148 +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.nio._
-
-import kafka.api.ApiUtils._
-import kafka.common._
-import kafka.message._
-import org.apache.kafka.common.protocol.ApiKeys
-
-@deprecated("This object has been deprecated and will be removed in a future release.", "1.0.0")
-object ProducerRequest {
-  val CurrentVersion = 2.shortValue
-
-  def readFrom(buffer: ByteBuffer): ProducerRequest = {
-    val versionId: Short = buffer.getShort
-    val correlationId: Int = buffer.getInt
-    val clientId: String = readShortString(buffer)
-    val requiredAcks: Short = buffer.getShort
-    val ackTimeoutMs: Int = buffer.getInt
-    //build the topic structure
-    val topicCount = buffer.getInt
-    val partitionDataPairs = (1 to topicCount).flatMap(_ => {
-      // process topic
-      val topic = readShortString(buffer)
-      val partitionCount = buffer.getInt
-      (1 to partitionCount).map(_ => {
-        val partition = buffer.getInt
-        val messageSetSize = buffer.getInt
-        val messageSetBuffer = new Array[Byte](messageSetSize)
-        buffer.get(messageSetBuffer,0,messageSetSize)
-        (TopicAndPartition(topic, partition), new ByteBufferMessageSet(ByteBuffer.wrap(messageSetBuffer)))
-      })
-    })
-
-    ProducerRequest(versionId, correlationId, clientId, requiredAcks, ackTimeoutMs, collection.mutable.Map(partitionDataPairs:_*))
-  }
-}
-
-@deprecated("This object has been deprecated and will be removed in a future release.", "1.0.0")
-case class ProducerRequest(versionId: Short = ProducerRequest.CurrentVersion,
-                           correlationId: Int,
-                           clientId: String,
-                           requiredAcks: Short,
-                           ackTimeoutMs: Int,
-                           data: collection.mutable.Map[TopicAndPartition, ByteBufferMessageSet])
-    extends RequestOrResponse(Some(ApiKeys.PRODUCE.id)) {
-
-  /**
-   * Partitions the data into a map of maps (one for each topic).
-   */
-  private lazy val dataGroupedByTopic = data.groupBy(_._1.topic)
-  val topicPartitionMessageSizeMap = data.map(r => r._1 -> r._2.sizeInBytes).toMap
-
-  def this(correlationId: Int,
-           clientId: String,
-           requiredAcks: Short,
-           ackTimeoutMs: Int,
-           data: collection.mutable.Map[TopicAndPartition, ByteBufferMessageSet]) =
-    this(ProducerRequest.CurrentVersion, correlationId, clientId, requiredAcks, ackTimeoutMs, data)
-
-  def writeTo(buffer: ByteBuffer) {
-    buffer.putShort(versionId)
-    buffer.putInt(correlationId)
-    writeShortString(buffer, clientId)
-    buffer.putShort(requiredAcks)
-    buffer.putInt(ackTimeoutMs)
-
-    //save the topic structure
-    buffer.putInt(dataGroupedByTopic.size) //the number of topics
-    dataGroupedByTopic.foreach {
-      case (topic, topicAndPartitionData) =>
-        writeShortString(buffer, topic) //write the topic
-        buffer.putInt(topicAndPartitionData.size) //the number of partitions
-        topicAndPartitionData.foreach(partitionAndData => {
-          val partition = partitionAndData._1.partition
-          val partitionMessageData = partitionAndData._2
-          val bytes = partitionMessageData.buffer
-          buffer.putInt(partition)
-          buffer.putInt(bytes.limit())
-          buffer.put(bytes)
-          bytes.rewind
-        })
-    }
-  }
-
-  def sizeInBytes: Int = {
-    2 + /* versionId */
-    4 + /* correlationId */
-    shortStringLength(clientId) + /* client id */
-    2 + /* requiredAcks */
-    4 + /* ackTimeoutMs */
-    4 + /* number of topics */
-    dataGroupedByTopic.foldLeft(0)((foldedTopics, currTopic) => {
-      foldedTopics +
-      shortStringLength(currTopic._1) +
-      4 + /* the number of partitions */
-      {
-        currTopic._2.foldLeft(0)((foldedPartitions, currPartition) => {
-          foldedPartitions +
-          4 + /* partition id */
-          4 + /* byte-length of serialized messages */
-          currPartition._2.sizeInBytes
-        })
-      }
-    })
-  }
-
-  def numPartitions = data.size
-
-  override def toString: String = {
-    describe(true)
-  }
-
-  override def describe(details: Boolean): String = {
-    val producerRequest = new StringBuilder
-    producerRequest.append("Name: " + this.getClass.getSimpleName)
-    producerRequest.append("; Version: " + versionId)
-    producerRequest.append("; CorrelationId: " + correlationId)
-    producerRequest.append("; ClientId: " + clientId)
-    producerRequest.append("; RequiredAcks: " + requiredAcks)
-    producerRequest.append("; AckTimeoutMs: " + ackTimeoutMs + " ms")
-    if(details)
-      producerRequest.append("; TopicAndPartition: " + topicPartitionMessageSizeMap.mkString(","))
-    producerRequest.toString()
-  }
-
-  def emptyData(){
-    data.clear()
-  }
-}
-
diff --git a/core/src/main/scala/kafka/api/ProducerResponse.scala b/core/src/main/scala/kafka/api/ProducerResponse.scala
deleted file mode 100644
index 2d3c9cc1396dd..0000000000000
--- a/core/src/main/scala/kafka/api/ProducerResponse.scala
+++ /dev/null
@@ -1,110 +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.nio.ByteBuffer
-import kafka.message.Message
-import org.apache.kafka.common.protocol.Errors
-
-import scala.collection.Map
-import kafka.common.TopicAndPartition
-import kafka.api.ApiUtils._
-
-@deprecated("This object has been deprecated and will be removed in a future release.", "1.0.0")
-object ProducerResponse {
-  // readFrom assumes that the response is written using V2 format
-  def readFrom(buffer: ByteBuffer): ProducerResponse = {
-    val correlationId = buffer.getInt
-    val topicCount = buffer.getInt
-    val statusPairs = (1 to topicCount).flatMap(_ => {
-      val topic = readShortString(buffer)
-      val partitionCount = buffer.getInt
-      (1 to partitionCount).map(_ => {
-        val partition = buffer.getInt
-        val error = Errors.forCode(buffer.getShort)
-        val offset = buffer.getLong
-        val timestamp = buffer.getLong
-        (TopicAndPartition(topic, partition), ProducerResponseStatus(error, offset, timestamp))
-      })
-    })
-
-    val throttleTime = buffer.getInt
-    ProducerResponse(correlationId, Map(statusPairs:_*), ProducerRequest.CurrentVersion, throttleTime)
-  }
-}
-
-case class ProducerResponseStatus(var error: Errors, offset: Long, timestamp: Long = Message.NoTimestamp)
-
-@deprecated("This object has been deprecated and will be removed in a future release.", "1.0.0")
-case class ProducerResponse(correlationId: Int,
-                            status: Map[TopicAndPartition, ProducerResponseStatus],
-                            requestVersion: Int = 0,
-                            throttleTime: Int = 0)
-    extends RequestOrResponse() {
-
-  /**
-   * Partitions the status map into a map of maps (one for each topic).
-   */
-  private lazy val statusGroupedByTopic = status.groupBy(_._1.topic)
-
-  def hasError = status.values.exists(_.error != Errors.NONE)
-
-  val sizeInBytes = {
-    val throttleTimeSize = if (requestVersion > 0) 4 else 0
-    val groupedStatus = statusGroupedByTopic
-    4 + /* correlation id */
-    4 + /* topic count */
-    groupedStatus.foldLeft (0) ((foldedTopics, currTopic) => {
-      foldedTopics +
-      shortStringLength(currTopic._1) +
-      4 + /* partition count for this topic */
-      currTopic._2.size * {
-        4 + /* partition id */
-        2 + /* error code */
-        8 + /* offset */
-        8 /* timestamp */
-      }
-    }) +
-    throttleTimeSize
-  }
-
-  def writeTo(buffer: ByteBuffer) {
-    val groupedStatus = statusGroupedByTopic
-    buffer.putInt(correlationId)
-    buffer.putInt(groupedStatus.size) // topic count
-
-    groupedStatus.foreach(topicStatus => {
-      val (topic, errorsAndOffsets) = topicStatus
-      writeShortString(buffer, topic)
-      buffer.putInt(errorsAndOffsets.size) // partition count
-      errorsAndOffsets.foreach {
-        case((TopicAndPartition(_, partition), ProducerResponseStatus(error, nextOffset, timestamp))) =>
-          buffer.putInt(partition)
-          buffer.putShort(error.code)
-          buffer.putLong(nextOffset)
-          buffer.putLong(timestamp)
-      }
-    })
-    // Throttle time is only supported on V1 style requests
-    if (requestVersion > 0)
-      buffer.putInt(throttleTime)
-  }
-
-  override def describe(details: Boolean):String = { toString }
-}
-
diff --git a/core/src/main/scala/kafka/api/Request.scala b/core/src/main/scala/kafka/api/Request.scala
new file mode 100644
index 0000000000000..653b5f653ac52
--- /dev/null
+++ b/core/src/main/scala/kafka/api/Request.scala
@@ -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 kafka.api
+
+object Request {
+  val OrdinaryConsumerId: Int = -1
+  val DebuggingConsumerId: Int = -2
+  val FutureLocalReplicaId: Int = -3
+
+  // Broker ids are non-negative int.
+  def isValidBrokerId(brokerId: Int): Boolean = brokerId >= 0
+
+  def describeReplicaId(replicaId: Int): String = {
+    replicaId match {
+      case OrdinaryConsumerId => "consumer"
+      case DebuggingConsumerId => "debug consumer"
+      case FutureLocalReplicaId => "future local replica"
+      case id if isValidBrokerId(id) => s"replica [$id]"
+      case id => s"invalid replica [$id]"
+    }
+  }
+}
diff --git a/core/src/main/scala/kafka/api/RequestOrResponse.scala b/core/src/main/scala/kafka/api/RequestOrResponse.scala
deleted file mode 100644
index a0fe129c40abe..0000000000000
--- a/core/src/main/scala/kafka/api/RequestOrResponse.scala
+++ /dev/null
@@ -1,48 +0,0 @@
-package kafka.api
-
-/**
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-import java.nio._
-import kafka.network.RequestChannel
-import kafka.utils.Logging
-
-object Request {
-  val OrdinaryConsumerId: Int = -1
-  val DebuggingConsumerId: Int = -2
-  val FutureLocalReplicaId: Int = -3
-
-  // Broker ids are non-negative int.
-  def isValidBrokerId(brokerId: Int): Boolean = brokerId >= 0
-}
-
-
-abstract class RequestOrResponse(val requestId: Option[Short] = None) extends Logging {
-
-  def sizeInBytes: Int
-
-  def writeTo(buffer: ByteBuffer): Unit
-
-  def handleError(e: Throwable, requestChannel: RequestChannel, request: RequestChannel.Request): Unit = {}
-
-  /* The purpose of this API is to return a string description of the Request mainly for the purpose of request logging.
-  *  This API has no meaning for a Response object.
-   * @param details If this is false, omit the parts of the request description that are proportional to the number of
-   *                topics or partitions. This is mainly to control the amount of request logging. */
-  def describe(details: Boolean): String
-}
-
diff --git a/core/src/main/scala/kafka/api/TopicMetadata.scala b/core/src/main/scala/kafka/api/TopicMetadata.scala
deleted file mode 100644
index e4d730c95af6d..0000000000000
--- a/core/src/main/scala/kafka/api/TopicMetadata.scala
+++ /dev/null
@@ -1,153 +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 kafka.cluster.BrokerEndPoint
-import java.nio.ByteBuffer
-import kafka.api.ApiUtils._
-import kafka.utils.Logging
-import org.apache.kafka.common.protocol.Errors
-
-object TopicMetadata {
-
-  val NoLeaderNodeId = -1
-
-  def readFrom(buffer: ByteBuffer, brokers: Map[Int, BrokerEndPoint]): TopicMetadata = {
-    val error = Errors.forCode(readShortInRange(buffer, "error code", (-1, Short.MaxValue)))
-    val topic = readShortString(buffer)
-    val numPartitions = readIntInRange(buffer, "number of partitions", (0, Int.MaxValue))
-    val partitionsMetadata: Array[PartitionMetadata] = new Array[PartitionMetadata](numPartitions)
-    for(i <- 0 until numPartitions) {
-      val partitionMetadata = PartitionMetadata.readFrom(buffer, brokers)
-      partitionsMetadata(i) = partitionMetadata
-    }
-    new TopicMetadata(topic, partitionsMetadata, error)
-  }
-}
-
-case class TopicMetadata(topic: String, partitionsMetadata: Seq[PartitionMetadata], error: Errors = Errors.NONE) extends Logging {
-  def sizeInBytes: Int = {
-    2 /* error code */ +
-    shortStringLength(topic) +
-    4 + partitionsMetadata.map(_.sizeInBytes).sum /* size and partition data array */
-  }
-
-  def writeTo(buffer: ByteBuffer) {
-    /* error code */
-    buffer.putShort(error.code)
-    /* topic */
-    writeShortString(buffer, topic)
-    /* number of partitions */
-    buffer.putInt(partitionsMetadata.size)
-    partitionsMetadata.foreach(m => m.writeTo(buffer))
-  }
-
-  override def toString: String = {
-    val topicMetadataInfo = new StringBuilder
-    topicMetadataInfo.append("{TopicMetadata for topic %s -> ".format(topic))
-    error match {
-      case Errors.NONE =>
-        partitionsMetadata.foreach { partitionMetadata =>
-          partitionMetadata.error match {
-            case Errors.NONE =>
-              topicMetadataInfo.append("\nMetadata for partition [%s,%d] is %s".format(topic,
-                partitionMetadata.partitionId, partitionMetadata.toString()))
-            case Errors.REPLICA_NOT_AVAILABLE =>
-              // this error message means some replica other than the leader is not available. The consumer
-              // doesn't care about non leader replicas, so ignore this
-              topicMetadataInfo.append("\nMetadata for partition [%s,%d] is %s".format(topic,
-                partitionMetadata.partitionId, partitionMetadata.toString()))
-            case error: Errors =>
-              topicMetadataInfo.append("\nMetadata for partition [%s,%d] is not available due to %s".format(topic,
-                partitionMetadata.partitionId, error.exceptionName))
-          }
-        }
-      case error: Errors =>
-        topicMetadataInfo.append("\nNo partition metadata for topic %s due to %s".format(topic,
-          error.exceptionName))
-    }
-    topicMetadataInfo.append("}")
-    topicMetadataInfo.toString()
-  }
-}
-
-object PartitionMetadata {
-
-  def readFrom(buffer: ByteBuffer, brokers: Map[Int, BrokerEndPoint]): PartitionMetadata = {
-    val error = Errors.forCode(readShortInRange(buffer, "error code", (-1, Short.MaxValue)))
-    val partitionId = readIntInRange(buffer, "partition id", (0, Int.MaxValue)) /* partition id */
-    val leaderId = buffer.getInt
-    val leader = brokers.get(leaderId)
-
-    /* list of all replicas */
-    val numReplicas = readIntInRange(buffer, "number of all replicas", (0, Int.MaxValue))
-    val replicaIds = (0 until numReplicas).map(_ => buffer.getInt)
-    val replicas = replicaIds.map(brokers)
-
-    /* list of in-sync replicas */
-    val numIsr = readIntInRange(buffer, "number of in-sync replicas", (0, Int.MaxValue))
-    val isrIds = (0 until numIsr).map(_ => buffer.getInt)
-    val isr = isrIds.map(brokers)
-
-    new PartitionMetadata(partitionId, leader, replicas, isr, error)
-  }
-}
-
-case class PartitionMetadata(partitionId: Int,
-                             leader: Option[BrokerEndPoint],
-                             replicas: Seq[BrokerEndPoint],
-                             isr: Seq[BrokerEndPoint] = Seq.empty,
-                             error: Errors = Errors.NONE) extends Logging {
-  def sizeInBytes: Int = {
-    2 /* error code */ +
-    4 /* partition id */ +
-    4 /* leader */ +
-    4 + 4 * replicas.size /* replica array */ +
-    4 + 4 * isr.size /* isr array */
-  }
-
-  def writeTo(buffer: ByteBuffer) {
-    buffer.putShort(error.code)
-    buffer.putInt(partitionId)
-
-    /* leader */
-    val leaderId = leader.fold(TopicMetadata.NoLeaderNodeId)(leader => leader.id)
-    buffer.putInt(leaderId)
-
-    /* number of replicas */
-    buffer.putInt(replicas.size)
-    replicas.foreach(r => buffer.putInt(r.id))
-
-    /* number of in-sync replicas */
-    buffer.putInt(isr.size)
-    isr.foreach(r => buffer.putInt(r.id))
-  }
-
-  override def toString: String = {
-    val partitionMetadataString = new StringBuilder
-    partitionMetadataString.append("\tpartition " + partitionId)
-    partitionMetadataString.append("\tleader: " + leader.getOrElse("none"))
-    partitionMetadataString.append("\treplicas: " + replicas.mkString(","))
-    partitionMetadataString.append("\tisr: " + isr.mkString(","))
-    partitionMetadataString.append("\tisUnderReplicated: " + (isr.size < replicas.size))
-    partitionMetadataString.toString()
-  }
-
-}
-
-
diff --git a/core/src/main/scala/kafka/api/TopicMetadataRequest.scala b/core/src/main/scala/kafka/api/TopicMetadataRequest.scala
deleted file mode 100644
index 217cedc495983..0000000000000
--- a/core/src/main/scala/kafka/api/TopicMetadataRequest.scala
+++ /dev/null
@@ -1,72 +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.nio.ByteBuffer
-
-import kafka.api.ApiUtils._
-import kafka.utils.Logging
-import org.apache.kafka.common.protocol.ApiKeys
-
-@deprecated("This object has been deprecated and will be removed in a future release.", "1.0.0")
-object TopicMetadataRequest extends Logging {
-  val CurrentVersion = 0.shortValue
-  val DefaultClientId = ""
-}
-
-@deprecated("This object has been deprecated and will be removed in a future release.", "1.0.0")
-case class TopicMetadataRequest(versionId: Short,
-                                correlationId: Int,
-                                clientId: String,
-                                topics: Seq[String])
- extends RequestOrResponse(Some(ApiKeys.METADATA.id)){
-
-  def this(topics: Seq[String], correlationId: Int) =
-    this(TopicMetadataRequest.CurrentVersion, correlationId, TopicMetadataRequest.DefaultClientId, topics)
-
-  def writeTo(buffer: ByteBuffer) {
-    buffer.putShort(versionId)
-    buffer.putInt(correlationId)
-    writeShortString(buffer, clientId)
-    buffer.putInt(topics.size)
-    topics.foreach(topic => writeShortString(buffer, topic))
-  }
-
-  def sizeInBytes: Int = {
-    2 +  /* version id */
-    4 + /* correlation id */
-    shortStringLength(clientId)  + /* client id */
-    4 + /* number of topics */
-    topics.foldLeft(0)(_ + shortStringLength(_)) /* topics */
-  }
-
-  override def toString: String = {
-    describe(true)
-  }
-
-  override def describe(details: Boolean): String = {
-    val topicMetadataRequest = new StringBuilder
-    topicMetadataRequest.append("Name: " + this.getClass.getSimpleName)
-    topicMetadataRequest.append("; Version: " + versionId)
-    topicMetadataRequest.append("; CorrelationId: " + correlationId)
-    topicMetadataRequest.append("; ClientId: " + clientId)
-    if(details)
-      topicMetadataRequest.append("; Topics: " + topics.mkString(","))
-    topicMetadataRequest.toString()
-  }
-}
diff --git a/core/src/main/scala/kafka/api/TopicMetadataResponse.scala b/core/src/main/scala/kafka/api/TopicMetadataResponse.scala
deleted file mode 100644
index 5447ea00daa22..0000000000000
--- a/core/src/main/scala/kafka/api/TopicMetadataResponse.scala
+++ /dev/null
@@ -1,57 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package kafka.api
-
-import kafka.cluster.BrokerEndPoint
-import java.nio.ByteBuffer
-
-@deprecated("This object has been deprecated and will be removed in a future release.", "1.0.0")
-object TopicMetadataResponse {
-
-  def readFrom(buffer: ByteBuffer): TopicMetadataResponse = {
-    val correlationId = buffer.getInt
-    val brokerCount = buffer.getInt
-    val brokers = (0 until brokerCount).map(_ => BrokerEndPoint.readFrom(buffer))
-    val brokerMap = brokers.map(b => (b.id, b)).toMap
-    val topicCount = buffer.getInt
-    val topicsMetadata = (0 until topicCount).map(_ => TopicMetadata.readFrom(buffer, brokerMap))
-    new TopicMetadataResponse(brokers, topicsMetadata, correlationId)
-  }
-}
-
-@deprecated("This object has been deprecated and will be removed in a future release.", "1.0.0")
-case class TopicMetadataResponse(brokers: Seq[BrokerEndPoint],
-                                 topicsMetadata: Seq[TopicMetadata],
-                                 correlationId: Int)
-    extends RequestOrResponse() {
-  val sizeInBytes: Int = {
-    4 + 4 + brokers.map(_.sizeInBytes).sum + 4 + topicsMetadata.map(_.sizeInBytes).sum
-  }
-
-  def writeTo(buffer: ByteBuffer) {
-    buffer.putInt(correlationId)
-    /* brokers */
-    buffer.putInt(brokers.size)
-    brokers.foreach(_.writeTo(buffer))
-    /* topic metadata */
-    buffer.putInt(topicsMetadata.length)
-    topicsMetadata.foreach(_.writeTo(buffer))
-  }
-
-  override def describe(details: Boolean):String = { toString }
-}
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..11a956d40b536
--- /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.jdk.CollectionConverters._
+
+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/client/ClientUtils.scala b/core/src/main/scala/kafka/client/ClientUtils.scala
deleted file mode 100755
index 557325627db05..0000000000000
--- a/core/src/main/scala/kafka/client/ClientUtils.scala
+++ /dev/null
@@ -1,217 +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.client
-
-import org.apache.kafka.common.protocol.Errors
-
-import scala.collection._
-import kafka.cluster._
-import kafka.api._
-import kafka.producer._
-import kafka.common.{BrokerEndPointNotAvailableException, KafkaException}
-import kafka.utils.{CoreUtils, Logging}
-import java.util.Properties
-
-import util.Random
-import kafka.network.BlockingChannel
-import kafka.utils.ZkUtils
-import java.io.IOException
-
-import org.apache.kafka.common.security.auth.SecurityProtocol
-
- /**
- * Helper functions common to clients (producer, consumer, or admin)
- */
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.11.0.0")
-object ClientUtils extends Logging {
-
-  /**
-   * Used by the producer to send a metadata request since it has access to the ProducerConfig
-   * @param topics The topics for which the metadata needs to be fetched
-   * @param brokers The brokers in the cluster as configured on the producer through metadata.broker.list
-   * @param producerConfig The producer's config
-   * @return topic metadata response
-   */
-  @deprecated("This method has been deprecated and will be removed in a future release.", "0.10.0.0")
-  def fetchTopicMetadata(topics: Set[String], brokers: Seq[BrokerEndPoint], producerConfig: ProducerConfig, correlationId: Int): TopicMetadataResponse = {
-    var fetchMetaDataSucceeded: Boolean = false
-    var i: Int = 0
-    val topicMetadataRequest = new TopicMetadataRequest(TopicMetadataRequest.CurrentVersion, correlationId, producerConfig.clientId, topics.toSeq)
-    var topicMetadataResponse: TopicMetadataResponse = null
-    var t: Throwable = null
-    // shuffle the list of brokers before sending metadata requests so that most requests don't get routed to the
-    // same broker
-    val shuffledBrokers = Random.shuffle(brokers)
-    while(i < shuffledBrokers.size && !fetchMetaDataSucceeded) {
-      val producer: SyncProducer = ProducerPool.createSyncProducer(producerConfig, shuffledBrokers(i))
-      info("Fetching metadata from broker %s with correlation id %d for %d topic(s) %s".format(shuffledBrokers(i), correlationId, topics.size, topics))
-      try {
-        topicMetadataResponse = producer.send(topicMetadataRequest)
-        fetchMetaDataSucceeded = true
-      }
-      catch {
-        case e: Throwable =>
-          warn("Fetching topic metadata with correlation id %d for topics [%s] from broker [%s] failed"
-            .format(correlationId, topics, shuffledBrokers(i).toString), e)
-          t = e
-      } finally {
-        i = i + 1
-        producer.close()
-      }
-    }
-    if(!fetchMetaDataSucceeded) {
-      throw new KafkaException("fetching topic metadata for topics [%s] from broker [%s] failed".format(topics, shuffledBrokers), t)
-    } else {
-      debug("Successfully fetched metadata for %d topic(s) %s".format(topics.size, topics))
-    }
-    topicMetadataResponse
-  }
-
-  /**
-   * Used by a non-producer client to send a metadata request
-   * @param topics The topics for which the metadata needs to be fetched
-   * @param brokers The brokers in the cluster as configured on the client
-   * @param clientId The client's identifier
-   * @return topic metadata response
-   */
-  def fetchTopicMetadata(topics: Set[String], brokers: Seq[BrokerEndPoint], clientId: String, timeoutMs: Int,
-                         correlationId: Int = 0): TopicMetadataResponse = {
-    val props = new Properties()
-    props.put("metadata.broker.list", brokers.map(_.connectionString).mkString(","))
-    props.put("client.id", clientId)
-    props.put("request.timeout.ms", timeoutMs.toString)
-    val producerConfig = new ProducerConfig(props)
-    fetchTopicMetadata(topics, brokers, producerConfig, correlationId)
-  }
-
-  /**
-   * Parse a list of broker urls in the form host1:port1, host2:port2, ...
-   */
-  def parseBrokerList(brokerListStr: String): Seq[BrokerEndPoint] = {
-    val brokersStr = CoreUtils.parseCsvList(brokerListStr)
-
-    brokersStr.zipWithIndex.map { case (address, brokerId) =>
-      BrokerEndPoint.createBrokerEndPoint(brokerId, address)
-    }
-  }
-
-  /**
-   * Creates a blocking channel to a random broker
-   */
-  def channelToAnyBroker(zkUtils: ZkUtils, socketTimeoutMs: Int = 3000) : BlockingChannel = {
-    var channel: BlockingChannel = null
-    var connected = false
-    while (!connected) {
-      val allBrokers = getPlaintextBrokerEndPoints(zkUtils)
-      Random.shuffle(allBrokers).find { broker =>
-        trace("Connecting to broker %s:%d.".format(broker.host, broker.port))
-        try {
-          channel = new BlockingChannel(broker.host, broker.port, BlockingChannel.UseDefaultBufferSize, BlockingChannel.UseDefaultBufferSize, socketTimeoutMs)
-          channel.connect()
-          debug("Created channel to broker %s:%d.".format(channel.host, channel.port))
-          true
-        } catch {
-          case _: Exception =>
-            if (channel != null) channel.disconnect()
-            channel = null
-            info("Error while creating channel to %s:%d.".format(broker.host, broker.port))
-            false
-        }
-      }
-      connected = channel != null
-    }
-
-    channel
-  }
-
-   /**
-    * Returns the first end point from each broker with the PLAINTEXT security protocol.
-    */
-  def getPlaintextBrokerEndPoints(zkUtils: ZkUtils): Seq[BrokerEndPoint] = {
-    zkUtils.getAllBrokersInCluster().map { broker =>
-      broker.endPoints.collectFirst {
-        case endPoint if endPoint.securityProtocol == SecurityProtocol.PLAINTEXT =>
-          new BrokerEndPoint(broker.id, endPoint.host, endPoint.port)
-      }.getOrElse(throw new BrokerEndPointNotAvailableException(s"End point with security protocol PLAINTEXT not found for broker ${broker.id}"))
-    }
-  }
-
-   /**
-    * Creates a blocking channel to the offset manager of the given group
-    */
-   def channelToOffsetManager(group: String, zkUtils: ZkUtils, socketTimeoutMs: Int = 3000, retryBackOffMs: Int = 1000) = {
-     var queryChannel = channelToAnyBroker(zkUtils)
-
-     var offsetManagerChannelOpt: Option[BlockingChannel] = None
-
-     while (offsetManagerChannelOpt.isEmpty) {
-
-       var coordinatorOpt: Option[BrokerEndPoint] = None
-
-       while (coordinatorOpt.isEmpty) {
-         try {
-           if (!queryChannel.isConnected)
-             queryChannel = channelToAnyBroker(zkUtils)
-           debug("Querying %s:%d to locate offset manager for %s.".format(queryChannel.host, queryChannel.port, group))
-           queryChannel.send(GroupCoordinatorRequest(group))
-           val response = queryChannel.receive()
-           val consumerMetadataResponse =  GroupCoordinatorResponse.readFrom(response.payload())
-           debug("Consumer metadata response: " + consumerMetadataResponse.toString)
-           if (consumerMetadataResponse.error == Errors.NONE)
-             coordinatorOpt = consumerMetadataResponse.coordinatorOpt
-           else {
-             debug("Query to %s:%d to locate offset manager for %s failed - will retry in %d milliseconds."
-                  .format(queryChannel.host, queryChannel.port, group, retryBackOffMs))
-             Thread.sleep(retryBackOffMs)
-           }
-         }
-         catch {
-           case _: IOException =>
-             info("Failed to fetch consumer metadata from %s:%d.".format(queryChannel.host, queryChannel.port))
-             queryChannel.disconnect()
-         }
-       }
-
-       val coordinator = coordinatorOpt.get
-       if (coordinator.host == queryChannel.host && coordinator.port == queryChannel.port) {
-         offsetManagerChannelOpt = Some(queryChannel)
-       } else {
-         val connectString = "%s:%d".format(coordinator.host, coordinator.port)
-         var offsetManagerChannel: BlockingChannel = null
-         try {
-           debug("Connecting to offset manager %s.".format(connectString))
-           offsetManagerChannel = new BlockingChannel(coordinator.host, coordinator.port,
-                                                      BlockingChannel.UseDefaultBufferSize,
-                                                      BlockingChannel.UseDefaultBufferSize,
-                                                      socketTimeoutMs)
-           offsetManagerChannel.connect()
-           offsetManagerChannelOpt = Some(offsetManagerChannel)
-           queryChannel.disconnect()
-         }
-         catch {
-           case _: IOException => // offsets manager may have moved
-             info("Error while connecting to %s.".format(connectString))
-             if (offsetManagerChannel != null) offsetManagerChannel.disconnect()
-             Thread.sleep(retryBackOffMs)
-             offsetManagerChannelOpt = None // just in case someone decides to change shutdownChannel to not swallow exceptions
-         }
-       }
-     }
-
-     offsetManagerChannelOpt.get
-   }
- }
diff --git a/core/src/main/scala/kafka/cluster/Broker.scala b/core/src/main/scala/kafka/cluster/Broker.scala
index a148dfd754030..46483d044d95d 100755
--- a/core/src/main/scala/kafka/cluster/Broker.scala
+++ b/core/src/main/scala/kafka/cluster/Broker.scala
@@ -17,138 +17,40 @@
 
 package kafka.cluster
 
-import kafka.common.{BrokerEndPointNotAvailableException, BrokerNotAvailableException, KafkaException}
-import kafka.utils.Json
-import org.apache.kafka.common.Node
+import java.util
+
+import kafka.common.BrokerEndPointNotAvailableException
+import kafka.server.KafkaConfig
+import org.apache.kafka.common.feature.{Features, SupportedVersionRange}
+import org.apache.kafka.common.feature.Features._
+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.common.utils.Time
+import org.apache.kafka.server.authorizer.AuthorizerServerInfo
 
-/**
- * A Kafka broker.
- * A broker has an id and a collection of end-points.
- * Each end-point is (host, port, protocolType).
- */
-object Broker {
+import scala.collection.Seq
+import scala.jdk.CollectionConverters._
 
-  private val HostKey = "host"
-  private val PortKey = "port"
-  private val VersionKey = "version"
-  private val EndpointsKey = "endpoints"
-  private val RackKey = "rack"
-  private val JmxPortKey = "jmx_port"
-  private val ListenerSecurityProtocolMapKey = "listener_security_protocol_map"
-  private val TimestampKey = "timestamp"
-
-  /**
-    * Create a broker object from id and JSON string.
-    *
-    * @param id
-    * @param brokerInfoString
-    *
-    * Version 1 JSON schema for a broker is:
-    * {
-    *   "version":1,
-    *   "host":"localhost",
-    *   "port":9092
-    *   "jmx_port":9999,
-    *   "timestamp":"2233345666"
-    * }
-    *
-    * Version 2 JSON schema for a broker is:
-    * {
-    *   "version":2,
-    *   "host":"localhost",
-    *   "port":9092,
-    *   "jmx_port":9999,
-    *   "timestamp":"2233345666",
-    *   "endpoints":["PLAINTEXT://host1:9092", "SSL://host1:9093"]
-    * }
-    *
-    * Version 3 JSON schema for a broker is:
-    * {
-    *   "version":3,
-    *   "host":"localhost",
-    *   "port":9092,
-    *   "jmx_port":9999,
-    *   "timestamp":"2233345666",
-    *   "endpoints":["PLAINTEXT://host1:9092", "SSL://host1:9093"],
-    *   "rack":"dc1"
-    * }
-    *
-    * Version 4 (current) JSON schema for a broker is:
-    * {
-    *   "version":4,
-    *   "host":"localhost",
-    *   "port":9092,
-    *   "jmx_port":9999,
-    *   "timestamp":"2233345666",
-    *   "endpoints":["CLIENT://host1:9092", "REPLICATION://host1:9093"],
-    *   "listener_security_protocol_map":{"CLIENT":"SSL", "REPLICATION":"PLAINTEXT"},
-    *   "rack":"dc1"
-    * }
-    */
-  def createBroker(id: Int, brokerInfoString: String): Broker = {
-    if (brokerInfoString == null)
-      throw new BrokerNotAvailableException(s"Broker id $id does not exist")
-    try {
-      Json.parseFull(brokerInfoString) match {
-        case Some(js) =>
-          val brokerInfo = js.asJsonObject
-          val version = brokerInfo(VersionKey).to[Int]
-
-          val endpoints =
-            if (version < 1)
-              throw new KafkaException(s"Unsupported version of broker registration: $brokerInfoString")
-            else if (version == 1) {
-              val host = brokerInfo(HostKey).to[String]
-              val port = brokerInfo(PortKey).to[Int]
-              val securityProtocol = SecurityProtocol.PLAINTEXT
-              val endPoint = new EndPoint(host, port, ListenerName.forSecurityProtocol(securityProtocol), securityProtocol)
-              Seq(endPoint)
-            }
-            else {
-              val securityProtocolMap = brokerInfo.get(ListenerSecurityProtocolMapKey).map(
-                _.to[Map[String, String]].map { case (listenerName, securityProtocol) =>
-                  new ListenerName(listenerName) -> SecurityProtocol.forName(securityProtocol)
-                })
-              val listeners = brokerInfo(EndpointsKey).to[Seq[String]]
-              listeners.map(EndPoint.createEndPoint(_, securityProtocolMap))
-            }
-
-          val rack = brokerInfo.get(RackKey).flatMap(_.to[Option[String]])
-          Broker(id, endpoints, rack)
-        case None =>
-          throw new BrokerNotAvailableException(s"Broker id $id does not exist")
-      }
-    } catch {
-      case t: Throwable =>
-        throw new KafkaException(s"Failed to parse the broker info from zookeeper: $brokerInfoString", t)
-    }
-  }
-
-  def toJson(version: Int, id: Int, host: String, port: Int, advertisedEndpoints: Seq[EndPoint], jmxPort: Int,
-             rack: Option[String]): String = {
-    val jsonMap = collection.mutable.Map(VersionKey -> version,
-      HostKey -> host,
-      PortKey -> port,
-      EndpointsKey -> advertisedEndpoints.map(_.connectionString).toArray,
-      JmxPortKey -> jmxPort,
-      TimestampKey -> Time.SYSTEM.milliseconds().toString
-    )
-    rack.foreach(rack => if (version >= 3) jsonMap += (RackKey -> rack))
-
-    if (version >= 4) {
-      jsonMap += (ListenerSecurityProtocolMapKey -> advertisedEndpoints.map { endPoint =>
-        endPoint.listenerName.value -> endPoint.securityProtocol.name
-      }.toMap)
-    }
+object Broker {
+  private[cluster] case class ServerInfo(clusterResource: ClusterResource,
+                                         brokerId: Int,
+                                         endpoints: util.List[Endpoint],
+                                         interBrokerEndpoint: Endpoint) extends AuthorizerServerInfo
 
-    Json.encode(jsonMap)
+  def apply(id: Int, endPoints: Seq[EndPoint], rack: Option[String]): Broker = {
+    new Broker(id, endPoints, rack, emptySupportedFeatures)
   }
 }
 
-case class Broker(id: Int, endPoints: Seq[EndPoint], rack: Option[String]) {
+/**
+ * A Kafka broker.
+ *
+ * @param id          a broker id
+ * @param endPoints   a collection of EndPoint. Each end-point is (host, port, listener name, security protocol).
+ * @param rack        an optional rack
+ * @param features    supported features
+ */
+case class Broker(id: Int, endPoints: Seq[EndPoint], rack: Option[String], features: Features[SupportedVersionRange]) {
 
   private val endPointsMap = endPoints.map { endPoint =>
     endPoint.listenerName -> endPoint
@@ -158,26 +60,39 @@ case class Broker(id: Int, endPoints: Seq[EndPoint], rack: Option[String]) {
     throw new IllegalArgumentException(s"There is more than one end point with the same listener name: ${endPoints.mkString(",")}")
 
   override def toString: String =
-    s"$id : ${endPointsMap.values.mkString("(",",",")")} : ${rack.orNull}"
+    s"$id : ${endPointsMap.values.mkString("(",",",")")} : ${rack.orNull} : $features"
 
   def this(id: Int, host: String, port: Int, listenerName: ListenerName, protocol: SecurityProtocol) = {
-    this(id, Seq(EndPoint(host, port, listenerName, protocol)), None)
+    this(id, Seq(EndPoint(host, port, listenerName, protocol)), None, emptySupportedFeatures)
   }
 
   def this(bep: BrokerEndPoint, listenerName: ListenerName, protocol: SecurityProtocol) = {
     this(bep.id, bep.host, bep.port, listenerName, protocol)
   }
 
-  def getNode(listenerName: ListenerName): Node = {
-    val endpoint = endPointsMap.getOrElse(listenerName,
-      throw new BrokerEndPointNotAvailableException(s"End point with listener name ${listenerName.value} not found for broker $id"))
-    new Node(id, endpoint.host, endpoint.port, rack.orNull)
+  def node(listenerName: ListenerName): Node =
+    getNode(listenerName).getOrElse {
+      throw new BrokerEndPointNotAvailableException(s"End point with listener name ${listenerName.value} not found " +
+        s"for broker $id")
+    }
+
+  def getNode(listenerName: ListenerName): Option[Node] =
+    endPointsMap.get(listenerName).map(endpoint => new Node(id, endpoint.host, endpoint.port, rack.orNull))
+
+  def brokerEndPoint(listenerName: ListenerName): BrokerEndPoint = {
+    val endpoint = endPoint(listenerName)
+    new BrokerEndPoint(id, endpoint.host, endpoint.port)
   }
 
-  def getBrokerEndPoint(listenerName: ListenerName): BrokerEndPoint = {
-    val endpoint = endPointsMap.getOrElse(listenerName,
+  def endPoint(listenerName: ListenerName): EndPoint = {
+    endPointsMap.getOrElse(listenerName,
       throw new BrokerEndPointNotAvailableException(s"End point with listener name ${listenerName.value} not found for broker $id"))
-    new BrokerEndPoint(id, endpoint.host, endpoint.port)
   }
 
+  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/BrokerEndPoint.scala b/core/src/main/scala/kafka/cluster/BrokerEndPoint.scala
index 847e959d76f71..b2b36af09da31 100644
--- a/core/src/main/scala/kafka/cluster/BrokerEndPoint.scala
+++ b/core/src/main/scala/kafka/cluster/BrokerEndPoint.scala
@@ -19,7 +19,7 @@ package kafka.cluster
 import java.nio.ByteBuffer
 
 import kafka.api.ApiUtils._
-import kafka.common.KafkaException
+import org.apache.kafka.common.KafkaException
 import org.apache.kafka.common.utils.Utils._
 
 object BrokerEndPoint {
@@ -76,4 +76,8 @@ case class BrokerEndPoint(id: Int, host: String, port: Int) {
     4 + /* broker Id */
     4 + /* port */
     shortStringLength(host)
+
+  override def toString: String = {
+    s"BrokerEndPoint(id=$id, host=$host:$port)"
+  }
 }
diff --git a/core/src/main/scala/kafka/cluster/Cluster.scala b/core/src/main/scala/kafka/cluster/Cluster.scala
deleted file mode 100644
index 75bbec054aed8..0000000000000
--- a/core/src/main/scala/kafka/cluster/Cluster.scala
+++ /dev/null
@@ -1,45 +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.cluster
-
-import scala.collection._
-
-/**
- * The set of active brokers in the cluster
- */
-private[kafka] class Cluster {
-  
-  private val brokers = new mutable.HashMap[Int, Broker]
-  
-  def this(brokerList: Iterable[Broker]) {
-    this()
-	  for(broker <- brokerList)
-      brokers.put(broker.id, broker)
-  }
-
-  def getBroker(id: Int): Option[Broker] = brokers.get(id)
-  
-  def add(broker: Broker) = brokers.put(broker.id, broker)
-  
-  def remove(id: Int) = brokers.remove(id)
-  
-  def size = brokers.size
-  
-  override def toString: String =
-    "Cluster(" + brokers.values.mkString(", ") + ")"  
-}
diff --git a/core/src/main/scala/kafka/cluster/EndPoint.scala b/core/src/main/scala/kafka/cluster/EndPoint.scala
index 57ef0da7047fa..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 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 91b86ee42bc4f..f47eac69b9cb9 100755
--- a/core/src/main/scala/kafka/cluster/Partition.scala
+++ b/core/src/main/scala/kafka/cluster/Partition.scala
@@ -16,280 +16,616 @@
  */
 package kafka.cluster
 
-
 import java.util.concurrent.locks.ReentrantReadWriteLock
+import java.util.{Optional, Properties}
 
-import com.yammer.metrics.core.Gauge
-import kafka.api.LeaderAndIsr
-import kafka.api.Request
-import kafka.controller.KafkaController
-import kafka.log.{LogAppendInfo, LogConfig}
+import kafka.api.{ApiVersion, LeaderAndIsr}
+import kafka.common.UnexpectedAppendOffsetException
+import kafka.controller.{KafkaController, StateChangeLogger}
+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
-import org.apache.kafka.common.TopicPartition
-import org.apache.kafka.common.errors.{NotEnoughReplicasException, NotLeaderForPartitionException, PolicyViolationException}
+import kafka.zk.{AdminZkClient, KafkaZkClient}
+import kafka.zookeeper.ZooKeeperClientException
+import org.apache.kafka.common.errors._
+import org.apache.kafka.common.message.FetchResponseData
+import org.apache.kafka.common.message.LeaderAndIsrRequestData.LeaderAndIsrPartitionState
+import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData.EpochEndOffset
 import org.apache.kafka.common.protocol.Errors
-import org.apache.kafka.common.protocol.Errors._
-import org.apache.kafka.common.record.MemoryRecords
-import org.apache.kafka.common.requests.EpochEndOffset._
-import org.apache.kafka.common.requests.{EpochEndOffset, LeaderAndIsrRequest}
+import org.apache.kafka.common.record.FileRecords.TimestampAndOffset
+import org.apache.kafka.common.record.{MemoryRecords, RecordBatch}
+import org.apache.kafka.common.requests._
+import org.apache.kafka.common.requests.OffsetsForLeaderEpochResponse.{UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET}
 import org.apache.kafka.common.utils.Time
+import org.apache.kafka.common.{IsolationLevel, TopicPartition}
+
+import scala.collection.{Map, Seq}
+import scala.jdk.CollectionConverters._
+
+trait IsrChangeListener {
+  def markExpand(): Unit
+  def markShrink(): Unit
+  def markFailed(): Unit
+}
+
+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) 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)
+    newVersionOpt
+  }
+
+  override def expandIsr(controllerEpoch: Int, leaderAndIsr: LeaderAndIsr): Option[Int] = {
+    val newVersionOpt = updateIsr(controllerEpoch, leaderAndIsr)
+    newVersionOpt
+  }
+
+  private def updateIsr(controllerEpoch: Int, leaderAndIsr: LeaderAndIsr): Option[Int] = {
+    val (updateSucceeded, newVersion) = ReplicationUtils.updateLeaderAndIsr(zkClient, topicPartition,
+      leaderAndIsr, controllerEpoch)
+
+    if (updateSucceeded) {
+      Some(newVersion)
+    } else {
+      None
+    }
+  }
+}
+
+class DelayedOperations(topicPartition: TopicPartition,
+                        produce: DelayedOperationPurgatory[DelayedProduce],
+                        fetch: DelayedOperationPurgatory[DelayedFetch],
+                        deleteRecords: DelayedOperationPurgatory[DelayedDeleteRecords]) {
+
+  def checkAndCompleteAll(): Unit = {
+    val requestKey = TopicPartitionOperationKey(topicPartition)
+    fetch.checkAndComplete(requestKey)
+    produce.checkAndComplete(requestKey)
+    deleteRecords.checkAndComplete(requestKey)
+  }
+
+  def numDelayedDelete: Int = deleteRecords.numDelayed
+}
+
+object Partition extends KafkaMetricsGroup {
+  def apply(topicPartition: TopicPartition,
+            time: Time,
+            replicaManager: ReplicaManager): Partition = {
+
+    val isrChangeListener = new IsrChangeListener {
+      override def markExpand(): Unit = {
+        replicaManager.recordIsrChange(topicPartition)
+        replicaManager.isrExpandRate.mark()
+      }
+
+      override def markShrink(): Unit = {
+        replicaManager.recordIsrChange(topicPartition)
+        replicaManager.isrShrinkRate.mark()
+      }
+
+      override def markFailed(): Unit = replicaManager.failedIsrUpdatesRate.mark()
+    }
+
+    val zkIsrBackingStore = new ZkPartitionStateStore(
+      topicPartition,
+      replicaManager.zkClient)
+
+    val delayedOperations = new DelayedOperations(
+      topicPartition,
+      replicaManager.delayedProducePurgatory,
+      replicaManager.delayedFetchPurgatory,
+      replicaManager.delayedDeleteRecordsPurgatory)
+
+    new Partition(topicPartition,
+      replicaLagTimeMaxMs = replicaManager.config.replicaLagTimeMaxMs,
+      interBrokerProtocolVersion = replicaManager.config.interBrokerProtocolVersion,
+      localBrokerId = replicaManager.config.brokerId,
+      time = time,
+      stateStore = zkIsrBackingStore,
+      isrChangeListener = isrChangeListener,
+      delayedOperations = delayedOperations,
+      metadataCache = replicaManager.metadataCache,
+      logManager = replicaManager.logManager,
+      alterIsrManager = replicaManager.alterIsrManager)
+  }
+
+  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)
+  }
+}
+
+
+sealed trait AssignmentState {
+  def replicas: Seq[Int]
+  def replicationFactor: Int = replicas.size
+  def isAddingReplica(brokerId: Int): Boolean = false
+}
+
+case class OngoingReassignmentState(addingReplicas: Seq[Int],
+                                    removingReplicas: Seq[Int],
+                                    replicas: Seq[Int]) extends AssignmentState {
+
+  override def replicationFactor: Int = replicas.diff(addingReplicas).size // keep the size of the original replicas
+  override def isAddingReplica(replicaId: Int): Boolean = addingReplicas.contains(replicaId)
+}
+
+case class SimpleAssignmentState(replicas: Seq[Int]) extends AssignmentState
+
+
+
+sealed trait IsrState {
+  /**
+   * Includes only the in-sync replicas which have been committed to ZK.
+   */
+  def isr: Set[Int]
+
+  /**
+   * This set may include un-committed ISR members following an expansion. This "effective" ISR is used for advancing
+   * the high watermark as well as determining which replicas are required for acks=all produce requests.
+   *
+   * Only applicable as of IBP 2.7-IV2, for older versions this will return the committed ISR
+   *
+   */
+  def maximalIsr: Set[Int]
+
+  /**
+   * Indicates if we have an AlterIsr request inflight.
+   */
+  def isInflight: Boolean
+}
+
+case class PendingExpandIsr(
+  isr: Set[Int],
+  newInSyncReplicaId: Int
+) extends IsrState {
+  val maximalIsr = isr + newInSyncReplicaId
+  val isInflight = true
+
+  override def toString: String = {
+    s"PendingExpandIsr(isr=$isr" +
+      s", newInSyncReplicaId=$newInSyncReplicaId" +
+      ")"
+  }
+}
+
+case class PendingShrinkIsr(
+  isr: Set[Int],
+  outOfSyncReplicaIds: Set[Int]
+) extends IsrState  {
+  val maximalIsr = isr
+  val isInflight = true
+
+  override def toString: String = {
+    s"PendingShrinkIsr(isr=$isr" +
+      s", outOfSyncReplicaIds=$outOfSyncReplicaIds" +
+      ")"
+  }
+}
+
+case class CommittedIsr(
+  isr: Set[Int]
+) extends IsrState {
+  val maximalIsr = isr
+  val isInflight = false
+
+  override def toString: String = {
+    s"CommittedIsr(isr=$isr" +
+      ")"
+  }
+}
 
-import scala.collection.JavaConverters._
-import scala.collection.Map
 
 /**
  * Data structure that represents a topic partition. The leader maintains the AR, ISR, CUR, RAR
+ *
+ * Concurrency notes:
+ * 1) Partition is thread-safe. Operations on partitions may be invoked concurrently from different
+ *    request handler threads
+ * 2) ISR updates are synchronized using a read-write lock. Read lock is used to check if an update
+ *    is required to avoid acquiring write lock in the common case of replica fetch when no update
+ *    is performed. ISR update condition is checked a second time under write lock before performing
+ *    the update
+ * 3) Various other operations like leader changes are processed while holding the ISR write lock.
+ *    This can introduce delays in produce and replica fetch requests, but these operations are typically
+ *    infrequent.
+ * 4) HW updates are synchronized using ISR read lock. @Log lock is acquired during the update with
+ *    locking order Partition lock -> Log lock.
+ * 5) lock is used to prevent the follower replica from being updated while ReplicaAlterDirThread is
+ *    executing maybeReplaceCurrentWithFutureReplica() to replace follower replica with the future replica.
  */
-class Partition(val topic: String,
-                val partitionId: Int,
+class Partition(val topicPartition: TopicPartition,
+                val replicaLagTimeMaxMs: Long,
+                interBrokerProtocolVersion: ApiVersion,
+                localBrokerId: Int,
                 time: Time,
-                replicaManager: ReplicaManager,
-                val isOffline: Boolean = false) extends Logging with KafkaMetricsGroup {
-
-  val topicPartition = new TopicPartition(topic, partitionId)
-
-  // Do not use replicaManager if this partition is ReplicaManager.OfflinePartition
-  private val localBrokerId = if (!isOffline) replicaManager.config.brokerId else -1
-  private val logManager = if (!isOffline) replicaManager.logManager else null
-  private val zkClient = if (!isOffline) replicaManager.zkClient else null
-  // allReplicasMap includes both assigned replicas and the future replica if there is ongoing replica movement
-  private val allReplicasMap = new Pool[Int, Replica]
+                stateStore: PartitionStateStore,
+                isrChangeListener: IsrChangeListener,
+                delayedOperations: DelayedOperations,
+                metadataCache: MetadataCache,
+                logManager: LogManager,
+                alterIsrManager: AlterIsrManager) extends Logging with KafkaMetricsGroup {
+
+  def topic: String = topicPartition.topic
+  def partitionId: Int = topicPartition.partition
+
+  private val stateChangeLogger = new StateChangeLogger(localBrokerId, inControllerContext = false, None)
+  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
+
+  // lock to prevent the follower replica log update while checking if the log dir could be replaced with future log.
+  private val futureLogLock = new Object()
   private var zkVersion: Int = LeaderAndIsr.initialZKVersion
   @volatile private var leaderEpoch: Int = LeaderAndIsr.initialLeaderEpoch - 1
+  // start offset for 'leaderEpoch' above (leader epoch of the current leader for this partition),
+  // 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 private[cluster] var isrState: IsrState = CommittedIsr(Set.empty)
+  @volatile var assignmentState: AssignmentState = SimpleAssignmentState(Seq.empty)
+
+  private val useAlterIsr: Boolean = interBrokerProtocolVersion.isAlterIsrSupported
+
+  // 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
    * the controller sends it a start replica command containing the leader for each partition that the broker hosts.
    * In addition to the leader, the controller can also send the epoch of the controller that elected the leader for
    * each partition. */
-  private var controllerEpoch: Int = KafkaController.InitialControllerEpoch - 1
+  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
-    )
-  }
+  newGauge("UnderReplicated", () => if (isUnderReplicated) 1 else 0, tags)
+  newGauge("InSyncReplicasCount", () => if (isLeader) isrState.isr.size else 0, tags)
+  newGauge("UnderMinIsr", () => if (isUnderMinIsr) 1 else 0, tags)
+  newGauge("AtMinIsr", () => if (isAtMinIsr) 1 else 0, tags)
+  newGauge("ReplicasCount", () => if (isLeader) assignmentState.replicationFactor else 0, tags)
+  newGauge("LastStableOffsetLag", () => log.map(_.lastStableOffsetLag).getOrElse(0), tags)
 
-  private def isLeaderReplicaLocal: Boolean = leaderReplicaIfLocal.isDefined
+  def isUnderReplicated: Boolean = isLeader && (assignmentState.replicationFactor - isrState.isr.size) > 0
 
-  def isUnderReplicated: Boolean =
-    isLeaderReplicaLocal && inSyncReplicas.size < assignedReplicas.size
+  def isUnderMinIsr: Boolean = leaderLogIfLocal.exists { isrState.isr.size < _.config.minInSyncReplicas }
 
-  def isUnderMinIsr: Boolean = {
-    leaderReplicaIfLocal match {
-      case Some(leaderReplica) =>
-        inSyncReplicas.size < leaderReplica.log.get.config.minInSyncReplicas
-      case None =>
-        false
-    }
-  }
+  def isAtMinIsr: Boolean = leaderLogIfLocal.exists { isrState.isr.size == _.config.minInSyncReplicas }
+
+  def isReassigning: Boolean = assignmentState.isInstanceOf[OngoingReassignmentState]
+
+  def isAddingLocalReplica: Boolean = assignmentState.isAddingReplica(localBrokerId)
+
+  def isAddingReplica(replicaId: Int): Boolean = assignmentState.isAddingReplica(replicaId)
+
+  def inSyncReplicaIds: Set[Int] = isrState.isr
 
   /**
     * Create the future replica if 1) the current replica is not in the given log directory and 2) the future replica
     * 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 = {
-    // The readLock is needed to make sure that while the caller checks the log directory of the
+  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.
-    inReadLock(leaderIsrUpdateLock) {
-      val currentReplica = getReplica().get
-      if (currentReplica.log.get.dir.getParent == logDir)
-        false
-      else if (getReplica(Request.FutureLocalReplicaId).isDefined) {
-        val futureReplicaLogDir = getReplica(Request.FutureLocalReplicaId).get.log.get.dir.getParent
-        if (futureReplicaLogDir != logDir)
-          throw new IllegalStateException(s"The future log dir $futureReplicaLogDir of $topicPartition is different from the requested log dir $logDir")
+    inWriteLock(leaderIsrUpdateLock) {
+      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 {
-        getOrCreateReplica(Request.FutureLocalReplicaId)
-        true
+        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 =>
+            createLogIfNotExists(isNew = false, isFutureReplica = true, highWatermarkCheckpoints)
+            true
+        }
       }
     }
   }
 
-  def getOrCreateReplica(replicaId: Int = localBrokerId, isNew: Boolean = false): Replica = {
-    allReplicasMap.getAndMaybePut(replicaId, {
-      if (isReplicaLocal(replicaId)) {
-        val adminZkClient = new AdminZkClient(zkClient)
-        val prop = adminZkClient.fetchEntityConfig(ConfigType.Topic, topic)
-        val config = LogConfig.fromProps(logManager.defaultConfig.originals,
-                                         prop)
-        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.")
+    }
+  }
+
+  // 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 NotLeaderOrFollowerException(s"Replica with id $replicaId is not available on broker $localBrokerId")
+  }
+
+  private def checkCurrentLeaderEpoch(remoteLeaderEpochOpt: Optional[Integer]): Errors = {
+    if (!remoteLeaderEpochOpt.isPresent) {
+      Errors.NONE
+    } else {
+      val remoteLeaderEpoch = remoteLeaderEpochOpt.get
+      val localLeaderEpoch = leaderEpoch
+      if (localLeaderEpoch > remoteLeaderEpoch)
+        Errors.FENCED_LEADER_EPOCH
+      else if (localLeaderEpoch < remoteLeaderEpoch)
+        Errors.UNKNOWN_LEADER_EPOCH
+      else
+        Errors.NONE
+    }
+  }
+
+  private def getLocalLog(currentLeaderEpoch: Optional[Integer],
+                          requireLeader: Boolean): Either[Log, Errors] = {
+    checkCurrentLeaderEpoch(currentLeaderEpoch) match {
+      case Errors.NONE =>
+        if (requireLeader && !isLeader) {
+          Right(Errors.NOT_LEADER_OR_FOLLOWER)
+        } else {
+          log match {
+            case Some(partitionLog) =>
+              Left(partitionLog)
+            case _ =>
+              Right(Errors.NOT_LEADER_OR_FOLLOWER)
+          }
+        }
+      case error =>
+        Right(error)
+    }
   }
 
-  def getReplica(replicaId: Int = localBrokerId): Option[Replica] = Option(allReplicasMap.get(replicaId))
+  def localLogOrException: Log = log.getOrElse {
+    throw new NotLeaderOrFollowerException(s"Log for partition $topicPartition is not available " +
+      s"on broker $localBrokerId")
+  }
 
-  def leaderReplicaIfLocal: Option[Replica] =
-    leaderReplicaIdOpt.filter(_ == localBrokerId).flatMap(getReplica)
+  def futureLocalLogOrException: Log = futureLog.getOrElse {
+    throw new NotLeaderOrFollowerException(s"Future log for partition $topicPartition is not available " +
+      s"on broker $localBrokerId")
+  }
 
-  def addReplicaIfNotExists(replica: Replica): Replica =
-    allReplicasMap.putIfNotExists(replica.brokerId, replica)
+  def leaderLogIfLocal: Option[Log] = {
+    log.filter(_ => isLeader)
+  }
 
-  def assignedReplicas: Set[Replica] =
-    allReplicasMap.values.filter(replica => Request.isValidBrokerId(replica.brokerId)).toSet
+  /**
+   * 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 ""} log for " +
+          s"partition $topicPartition with leader epoch $currentLeaderEpoch. The current leader " +
+          s"is $leaderReplicaIdOpt and the current epoch $leaderEpoch")
+    }
+  }
+
+  // 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)
+  }
 
-  def allReplicas: Set[Replica] =
-    allReplicasMap.values.toSet
+  // remoteReplicas will be called in the hot path, and must be inexpensive
+  def remoteReplicas: Iterable[Replica] =
+    remoteReplicasMap.values
 
-  private def removeReplica(replicaId: Int) {
-    allReplicasMap.remove(replicaId)
+  def futureReplicaDirChanged(newDestinationDir: String): Boolean = {
+    inReadLock(leaderIsrUpdateLock) {
+      futureLog.exists(_.parentDir != newDestinationDir)
+    }
   }
 
-  def removeFutureLocalReplica() {
+  def removeFutureLocalReplica(deleteFromLogDir: Boolean = true): Unit = {
     inWriteLock(leaderIsrUpdateLock) {
-      allReplicasMap.remove(Request.FutureLocalReplicaId)
+      futureLog = None
+      if (deleteFromLogDir)
+        logManager.asyncDelete(topicPartition, isFuture = true)
     }
   }
 
-  // Return true iff the future log has caught up with the current log for this partition
+  // Return true if the future replica exists and it has caught up with the current replica for this partition
   // 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 = getReplica().get
-    val futureReplica = getReplica(Request.FutureLocalReplicaId).get
-    if (replica.logEndOffset == futureReplica.logEndOffset) {
-      // 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) {
-        if (replica.logEndOffset == futureReplica.logEndOffset) {
-          logManager.replaceCurrentWithFutureLog(topicPartition)
-          replica.log = futureReplica.log
-          futureReplica.log = None
-          allReplicasMap.remove(Request.FutureLocalReplicaId)
-          true
-        } else false
-      }
-    } else false
+    // lock to prevent the log append by followers while checking if the log dir could be replaced with future log.
+    futureLogLock.synchronized {
+      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) {
+          futureLog match {
+            case Some(futurePartitionLog) =>
+              if (log.exists(_.logEndOffset == futurePartitionLog.logEndOffset)) {
+                logManager.replaceCurrentWithFutureLog(topicPartition)
+                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
+              false
+          }
+        }
+      } else false
+    }
   }
 
-  def delete() {
+  /**
+   * Delete the partition. Note that deleting the partition does not delete the underlying logs.
+   * The logs are deleted by the ReplicaManager after having deleted the partition.
+   */
+  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()
+      assignmentState = SimpleAssignmentState(Seq.empty)
+      log = None
+      futureLog = None
+      isrState = CommittedIsr(Set.empty)
       leaderReplicaIdOpt = None
-      removePartitionMetrics()
-      logManager.asyncDelete(topicPartition)
-      logManager.asyncDelete(topicPartition, isFuture = true)
+      leaderEpochStartOffsetOpt = None
+      Partition.removeMetrics(topicPartition)
     }
   }
 
   def getLeaderEpoch: Int = this.leaderEpoch
 
+  def getZkVersion: Int = this.zkVersion
+
   /**
    * Make the local replica the leader by resetting LogEndOffset for remote replicas (there could be old LogEndOffset
    * 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
+      controllerEpoch = partitionState.controllerEpoch
+
+      val isr = partitionState.isr.asScala.map(_.toInt).toSet
+      val addingReplicas = partitionState.addingReplicas.asScala.map(_.toInt)
+      val removingReplicas = partitionState.removingReplicas.asScala.map(_.toInt)
+
+      updateAssignmentAndIsr(
+        assignment = partitionState.replicas.asScala.map(_.toInt),
+        isr = isr,
+        addingReplicas = addingReplicas,
+        removingReplicas = removingReplicas
+      )
+      try {
+        createLogIfNotExists(partitionState.isNew, isFutureReplica = false, highWatermarkCheckpoints)
+      } catch {
+        case e: ZooKeeperClientException =>
+          stateChangeLogger.error(s"A ZooKeeper client exception has occurred and makeLeader will be skipping the " +
+            s"state change for the partition $topicPartition with leader epoch: $leaderEpoch ", e)
+
+          return false
+      }
 
-      info(s"$topicPartition starts at Leader Epoch ${partitionStateInfo.basePartitionState.leaderEpoch} from offset ${getReplica().get.logEndOffset.messageOffset}. Previous Leader Epoch was: $leaderEpoch")
+      val leaderLog = localLogOrException
+      val leaderEpochStartOffset = leaderLog.logEndOffset
+      stateChangeLogger.info(s"Leader $topicPartition starts at leader epoch ${partitionState.leaderEpoch} from " +
+        s"offset $leaderEpochStartOffset with high watermark ${leaderLog.highWatermark} " +
+        s"ISR ${isr.mkString("[", ",", "]")} addingReplicas ${addingReplicas.mkString("[", ",", "]")} " +
+        s"removingReplicas ${removingReplicas.mkString("[", ",", "]")}. 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
-      newAssignedReplicas.foreach(id => getOrCreateReplica(id, partitionStateInfo.isNew))
+      leaderEpoch = partitionState.leaderEpoch
+      leaderEpochStartOffsetOpt = Some(leaderEpochStartOffset)
+      zkVersion = partitionState.zkVersion
 
-      zkVersion = partitionStateInfo.basePartitionState.zkVersion
-      val isNewLeader = leaderReplicaIdOpt.map(_ != localBrokerId).getOrElse(true)
+      // Clear any pending AlterIsr requests and check replica state
+      alterIsrManager.clearPending(topicPartition)
 
-      val leaderReplica = getReplica().get
-      val curLeaderLogEndOffset = leaderReplica.logEndOffset.messageOffset
+      // 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.
+      leaderLog.maybeAssignEpochStartOffset(leaderEpoch, leaderEpochStartOffset)
+
+      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 (isrState.isr.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)
+        }
       }
       // 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)
@@ -299,27 +635,51 @@ class Partition(val topic: String,
 
   /**
    *  Make the local replica the follower by setting the new leader and ISR to empty
-   *  If the leader replica id does not change, return false to indicate the replica manager
+   *  If the leader replica id does not change and the new epoch is equal or one
+   *  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: Int = 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
-      zkVersion = partitionStateInfo.basePartitionState.zkVersion
-
-      if (leaderReplicaIdOpt.isDefined && leaderReplicaIdOpt.get == newLeaderBrokerId) {
-        false
+      controllerEpoch = partitionState.controllerEpoch
+
+      updateAssignmentAndIsr(
+        assignment = partitionState.replicas.asScala.iterator.map(_.toInt).toSeq,
+        isr = Set.empty[Int],
+        addingReplicas = partitionState.addingReplicas.asScala.map(_.toInt),
+        removingReplicas = partitionState.removingReplicas.asScala.map(_.toInt)
+      )
+      try {
+        createLogIfNotExists(partitionState.isNew, isFutureReplica = false, highWatermarkCheckpoints)
+      } catch {
+        case e: ZooKeeperClientException =>
+          stateChangeLogger.error(s"A ZooKeeper client exception has occurred. makeFollower will be skipping the " +
+            s"state change for the partition $topicPartition with leader epoch: $leaderEpoch.", e)
+
+          return false
       }
-      else {
+
+      val followerLog = localLogOrException
+      val leaderEpochEndOffset = followerLog.logEndOffset
+      stateChangeLogger.info(s"Follower $topicPartition starts at leader epoch ${partitionState.leaderEpoch} from " +
+        s"offset $leaderEpochEndOffset with high watermark ${followerLog.highWatermark}. " +
+        s"Previous leader epoch was $leaderEpoch.")
+
+      leaderEpoch = partitionState.leaderEpoch
+      leaderEpochStartOffsetOpt = None
+      zkVersion = partitionState.zkVersion
+
+      // Since we might have been a leader previously, still clear any pending AlterIsr requests
+      alterIsrManager.clearPending(topicPartition)
+
+      if (leaderReplicaIdOpt.contains(newLeaderBrokerId) && leaderEpoch == oldLeaderEpoch) {
+        false
+      } else {
         leaderReplicaIdOpt = Some(newLeaderBrokerId)
         true
       }
@@ -328,69 +688,137 @@ class Partition(val topic: String,
 
   /**
    * Update the follower's state in the leader based on the last fetch request. See
-   * [[kafka.cluster.Replica#updateLogReadResult]] for details.
+   * [[Replica.updateFetchState()]] 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): 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)
+
+        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 this in-sync replica needs to be added to the ISR.
+        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) {
+          // the leader log may be updated by ReplicaAlterLogDirsThread so the following method must be in lock of
+          // leaderIsrUpdateLock to prevent adding new hw to invalid log.
+          inReadLock(leaderIsrUpdateLock) {
+            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
+   * @param addingReplicas An ordered sequence of all broker ids that will be added to the
+    *                       assignment
+   * @param removingReplicas An ordered sequence of all broker ids that will be removed from
+    *                         the assignment
+   */
+  def updateAssignmentAndIsr(assignment: Seq[Int],
+                             isr: Set[Int],
+                             addingReplicas: Seq[Int],
+                             removingReplicas: Seq[Int]): Unit = {
+    val newRemoteReplicas = assignment.filter(_ != localBrokerId)
+    val removedReplicas = remoteReplicasMap.keys.filter(!newRemoteReplicas.contains(_))
+
+    // due to code paths accessing remoteReplicasMap without a lock,
+    // first add the new replicas and then remove the old ones
+    newRemoteReplicas.foreach(id => remoteReplicasMap.getAndMaybePut(id, new Replica(id, topicPartition)))
+    remoteReplicasMap.removeAll(removedReplicas)
+
+    if (addingReplicas.nonEmpty || removingReplicas.nonEmpty)
+      assignmentState = OngoingReassignmentState(addingReplicas, removingReplicas, assignment)
+    else
+      assignmentState = SimpleAssignmentState(assignment)
+    isrState = CommittedIsr(isr)
   }
 
   /**
    * Check and maybe expand the ISR of the partition.
-   * A replica will be added to ISR if its LEO >= current hw of the partition.
+   * A replica will be added to ISR if its LEO >= current hw of the partition and it is caught up to
+   * an offset within the current leader epoch. A replica must be caught up to the current leader
+   * epoch before it can join ISR, because otherwise, if there is committed data between current
+   * leader's HW and LEO, the replica may become the leader before it fetches the committed data
+   * and the data will be lost.
    *
    * Technically, a replica shouldn't be in ISR if it hasn't caught up for longer than replicaLagTimeMaxMs,
    * even if its log end offset is >= HW. However, to be consistent with how the follower determines
    * 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 = {
-    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
-          if (!inSyncReplicas.contains(replica) &&
-             assignedReplicas.map(_.brokerId).contains(replicaId) &&
-             replica.logEndOffset.offsetDiff(leaderHW) >= 0) {
-            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
+  private def maybeExpandIsr(followerReplica: Replica, followerFetchTimeMs: Long): Unit = {
+    val needsIsrUpdate = canAddReplicaToIsr(followerReplica.brokerId) && inReadLock(leaderIsrUpdateLock) {
+      needsExpandIsr(followerReplica)
+    }
+    if (needsIsrUpdate) {
+      inWriteLock(leaderIsrUpdateLock) {
+        // check if this replica needs to be added to the ISR
+        if (needsExpandIsr(followerReplica)) {
+          expandIsr(followerReplica.brokerId)
+        }
       }
     }
   }
 
+  private def needsExpandIsr(followerReplica: Replica): Boolean = {
+    canAddReplicaToIsr(followerReplica.brokerId) && isFollowerAtHighwatermark(followerReplica)
+  }
+
+  private def canAddReplicaToIsr(followerReplicaId: Int): Boolean = {
+    val current = isrState
+    !current.isInflight && !current.isr.contains(followerReplicaId)
+  }
+
+  private def isFollowerAtHighwatermark(followerReplica: Replica): Boolean = {
+    leaderLogIfLocal.exists { leaderLog =>
+      val followerEndOffset = followerReplica.logEndOffset
+      followerEndOffset >= leaderLog.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).
@@ -400,40 +828,40 @@ class Partition(val topic: String,
    * 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 curMaximalIsr = isrState.maximalIsr
 
-        def numAcks = curInSyncReplicas.count { r =>
-          if (!r.isLocal)
-            if (r.logEndOffset.messageOffset >= requiredOffset) {
-              trace(s"Replica ${r.brokerId} received offset $requiredOffset")
-              true
-            }
-            else
-              false
-          else
-            true /* also count the local (leader) replica */
-        }
+        if (isTraceEnabled) {
+          def logEndOffsetString: ((Int, Long)) => String = {
+            case (brokerId, logEndOffset) => s"broker $brokerId: $logEndOffset"
+          }
 
-        trace(s"$numAcks acks satisfied with acks = -1")
+          val curInSyncReplicaObjects = (curMaximalIsr - 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}
 
-        val minIsr = leaderReplica.log.get.config.minInSyncReplicas
+          trace(s"Progress awaiting ISR acks for offset $requiredOffset: " +
+            s"acked: ${ackedReplicas.map(logEndOffsetString)}, " +
+            s"awaiting ${awaitingReplicas.map(logEndOffsetString)}")
+        }
 
-        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 <= curMaximalIsr.size)
             (true, Errors.NONE)
           else
             (true, Errors.NOT_ENOUGH_REPLICAS_AFTER_APPEND)
         } else
           (false, Errors.NONE)
       case None =>
-        (false, Errors.NOT_LEADER_FOR_PARTITION)
+        (false, Errors.NOT_LEADER_OR_FOLLOWER)
     }
   }
 
@@ -451,24 +879,44 @@ class Partition(val topic: String,
    * follower's log end offset may keep falling behind the HW (determined by the leader's log end offset) and therefore
    * will never be added to ISR.
    *
-   * Returns true if the HW was incremented, and false otherwise.
-   * Note There is no need to acquire the leaderIsrUpdate lock here
-   * since all callers of this private API acquire that lock
+   * With the addition of AlterIsr, we also consider newly added replicas as part of the ISR when advancing
+   * the HW. These replicas have not yet been committed to the ISR by the controller, so we could revert to the previously
+   * committed ISR. However, adding additional replicas to the ISR makes it more restrictive and therefor safe. We call
+   * this set the "maximal" ISR. See KIP-497 for more details
+   *
+   * Note There is no need to acquire the leaderIsrUpdate lock here since all callers of this private API acquire that lock
+   *
+   * @return true if the HW was incremented, and false otherwise.
    */
-  private def maybeIncrementLeaderHW(leaderReplica: Replica, curTime: Long = time.milliseconds): Boolean = {
-    val allLogEndOffsets = assignedReplicas.filter { replica =>
-      curTime - replica.lastCaughtUpTimeMs <= replicaManager.config.replicaLagTimeMaxMs || inSyncReplicas.contains(replica)
-    }.map(_.logEndOffset)
-    val newHighWatermark = allLogEndOffsets.min(new LogOffsetMetadata.OffsetOrdering)
-    val oldHighWatermark = leaderReplica.highWatermark
-    if (oldHighWatermark.messageOffset < newHighWatermark.messageOffset || oldHighWatermark.onOlderSegment(newHighWatermark)) {
-      leaderReplica.highWatermark = newHighWatermark
-      debug(s"High watermark updated to $newHighWatermark")
-      true
-    } else  {
-      debug(s"Skipping update high watermark since new hw $newHighWatermark is not larger than old hw $oldHighWatermark." +
-        s"All LEOs are ${allLogEndOffsets.mkString(",")}")
-      false
+  private def maybeIncrementLeaderHW(leaderLog: Log, curTime: Long = time.milliseconds): Boolean = {
+    // maybeIncrementLeaderHW is in the hot path, the following code is written to
+    // avoid unnecessary collection generation
+    var newHighWatermark = leaderLog.logEndOffsetMetadata
+    remoteReplicasMap.values.foreach { replica =>
+      // Note here we are using the "maximal", see explanation above
+      if (replica.logEndOffsetMetadata.messageOffset < newHighWatermark.messageOffset &&
+        (curTime - replica.lastCaughtUpTimeMs <= replicaLagTimeMaxMs || isrState.maximalIsr.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
     }
   }
 
@@ -478,45 +926,54 @@ class Partition(val topic: String,
    * Low watermark will increase when the leader broker receives either FetchRequest or DeleteRecordsRequest.
    */
   def lowWatermarkIfLeader: Long = {
-    if (!isLeaderReplicaLocal)
-      throw new NotLeaderForPartitionException("Leader not local for partition %s on broker %d".format(topicPartition, localBrokerId))
-    val logStartOffsets = allReplicas.collect {
-      case replica if replicaManager.metadataCache.isBrokerAlive(replica.brokerId) || replica.brokerId == Request.FutureLocalReplicaId => replica.logStartOffset
+    if (!isLeader)
+      throw new NotLeaderOrFollowerException(s"Leader not local for partition $topicPartition on broker $localBrokerId")
+
+    // 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)
-  }
-
-  def maybeShrinkIsr(replicaMaxLagTimeMs: Long) {
-    val leaderHWIncremented = inWriteLock(leaderIsrUpdateLock) {
-      leaderReplicaIfLocal match {
-        case Some(leaderReplica) =>
-          val outOfSyncReplicas = getOutOfSyncReplicas(leaderReplica, replicaMaxLagTimeMs)
-          if(outOfSyncReplicas.nonEmpty) {
-            val newInSyncReplicas = inSyncReplicas -- outOfSyncReplicas
-            assert(newInSyncReplicas.nonEmpty)
-            info("Shrinking ISR from %s to %s".format(inSyncReplicas.map(_.brokerId).mkString(","),
-              newInSyncReplicas.map(_.brokerId).mkString(",")))
-            // update ISR in zk and in cache
-            updateIsr(newInSyncReplicas)
-            // we may need to increment high watermark since ISR could be down to 1
-
-            replicaManager.isrShrinkRate.mark()
-            maybeIncrementLeaderHW(leaderReplica)
-          } else {
-            false
-          }
+  private def tryCompleteDelayedRequests(): Unit = delayedOperations.checkAndCompleteAll()
+
+  def maybeShrinkIsr(): Unit = {
+    val needsIsrUpdate = !isrState.isInflight && inReadLock(leaderIsrUpdateLock) {
+      needsShrinkIsr()
+    }
+    val leaderHWIncremented = needsIsrUpdate && inWriteLock(leaderIsrUpdateLock) {
+      leaderLogIfLocal.exists { leaderLog =>
+        val outOfSyncReplicaIds = getOutOfSyncReplicas(replicaLagTimeMaxMs)
+        if (outOfSyncReplicaIds.nonEmpty) {
+          val outOfSyncReplicaLog = outOfSyncReplicaIds.map { replicaId =>
+            s"(brokerId: $replicaId, endOffset: ${getReplicaOrException(replicaId).logEndOffset})"
+          }.mkString(" ")
+          val newIsrLog = (isrState.isr -- outOfSyncReplicaIds).mkString(",")
+          info(s"Shrinking ISR from ${isrState.isr.mkString(",")} to $newIsrLog. " +
+               s"Leader: (highWatermark: ${leaderLog.highWatermark}, endOffset: ${leaderLog.logEndOffset}). " +
+               s"Out of sync replicas: $outOfSyncReplicaLog.")
+
+          shrinkIsr(outOfSyncReplicaIds)
 
-        case None => false // do nothing if no longer leader
+          // we may need to increment high watermark since ISR could be down to 1
+          maybeIncrementLeaderHW(leaderLog)
+        } else {
+          false
+        }
       }
     }
 
@@ -525,75 +982,251 @@ class Partition(val topic: String,
       tryCompleteDelayedRequests()
   }
 
-  def getOutOfSyncReplicas(leaderReplica: Replica, maxLagMs: Long): Set[Replica] = {
-    /**
-     * there are two cases that will be handled here -
-     * 1. Stuck followers: If the leo of the replica hasn't been updated for maxLagMs ms,
-     *                     the follower is stuck and should be removed from the ISR
-     * 2. Slow followers: If the replica has not read up to the leo within the last maxLagMs ms,
-     *                    then the follower is lagging and should be removed from the ISR
-     * Both these cases are handled by checking the lastCaughtUpTimeMs which represents
-     * the last time when the replica was fully caught up. If either of the above conditions
-     * is violated, that replica is considered to be out of sync
-     *
-     **/
-    val candidateReplicas = inSyncReplicas - leaderReplica
+  private def needsShrinkIsr(): Boolean = {
+    leaderLogIfLocal.exists { _ => getOutOfSyncReplicas(replicaLagTimeMaxMs).nonEmpty }
+  }
 
-    val laggingReplicas = candidateReplicas.filter(r => (time.milliseconds - r.lastCaughtUpTimeMs) > maxLagMs)
-    if (laggingReplicas.nonEmpty)
-      debug("Lagging replicas are %s".format(laggingReplicas.map(_.brokerId).mkString(",")))
+  private def isFollowerOutOfSync(replicaId: Int,
+                                  leaderEndOffset: Long,
+                                  currentTimeMs: Long,
+                                  maxLagMs: Long): Boolean = {
+    val followerReplica = getReplicaOrException(replicaId)
+    followerReplica.logEndOffset != leaderEndOffset &&
+      (currentTimeMs - followerReplica.lastCaughtUpTimeMs) > maxLagMs
+  }
 
-    laggingReplicas
+  /**
+   * 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 -
+   * 1. Stuck followers: If the leo of the replica hasn't been updated for maxLagMs ms,
+   *                     the follower is stuck and should be removed from the ISR
+   * 2. Slow followers: If the replica has not read up to the leo within the last maxLagMs ms,
+   *                    then the follower is lagging and should be removed from the ISR
+   * Both these cases are handled by checking the lastCaughtUpTimeMs which represents
+   * the last time when the replica was fully caught up. If either of the above conditions
+   * is violated, that replica is considered to be out of sync
+   *
+   * If an ISR update is in-flight, we will return an empty set here
+   **/
+  def getOutOfSyncReplicas(maxLagMs: Long): Set[Int] = {
+    val current = isrState
+    if (!current.isInflight) {
+      val candidateReplicaIds = current.isr - localBrokerId
+      val currentTimeMs = time.milliseconds()
+      val leaderEndOffset = localLogOrException.logEndOffset
+      candidateReplicaIds.filter(replicaId => isFollowerOutOfSync(replicaId, leaderEndOffset, currentTimeMs, maxLagMs))
+    } else {
+      Set.empty
+    }
   }
 
-  def appendRecordsToFutureReplica(records: MemoryRecords) {
-    getReplica(Request.FutureLocalReplicaId).get.log.get.appendAsFollower(records)
+  private def doAppendRecordsToFollowerOrFutureReplica(records: MemoryRecords, isFuture: Boolean): Option[LogAppendInfo] = {
+    if (isFuture) {
+      // The read lock is needed to handle race condition if request handler thread tries to
+      // remove future replica after receiving AlterReplicaLogDirsRequest.
+      inReadLock(leaderIsrUpdateLock) {
+        // Note the replica may be undefined if it is removed by a non-ReplicaAlterLogDirsThread before
+        // this method is called
+        futureLog.map { _.appendAsFollower(records) }
+      }
+    } else {
+      // The lock is needed to prevent the follower replica from being updated while ReplicaAlterDirThread
+      // is executing maybeReplaceCurrentWithFutureReplica() to replace follower replica with the future replica.
+      futureLogLock.synchronized {
+        Some(localLogOrException.appendAsFollower(records))
+      }
+    }
   }
 
-  def appendRecordsToFollower(records: MemoryRecords) {
-    // 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.
-    inReadLock(leaderIsrUpdateLock) {
-      getReplica().get.log.get.appendAsFollower(records)
+  def appendRecordsToFollowerOrFutureReplica(records: MemoryRecords, isFuture: Boolean): Option[LogAppendInfo] = {
+    try {
+      doAppendRecordsToFollowerOrFutureReplica(records, isFuture)
+    } catch {
+      case e: UnexpectedAppendOffsetException =>
+        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
+          // fetch its first offset from the leader.
+          // We handle this case here instead of Log#append() because we will need to remove the
+          // segment that start with log start offset and create a new one with earlier offset
+          // (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 ${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)
+        } else
+          throw e
     }
   }
 
-  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 = isrState.isr.size
 
           // Avoid writing to leader if there are not enough insync replicas to make it safe
           if (inSyncSize < minIsr && requiredAcks == -1) {
-            throw new NotEnoughReplicasException("Number of insync replicas for partition %s is [%d], below required minimum [%d]"
-              .format(topicPartition, inSyncSize, minIsr))
+            throw new NotEnoughReplicasException(s"The size of the current ISR ${isrState.isr} " +
+              s"is insufficient to satisfy the min.isr requirement of $minIsr for partition $topicPartition")
           }
 
-          val info = log.appendAsLeader(records, leaderEpoch = this.leaderEpoch, isFromClient)
-          // probably unblock some follower fetch requests since log end offset has been updated
-          replicaManager.tryCompleteDelayedFetch(TopicPartitionOperationKey(this.topic, this.partitionId))
+          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"
+          throw new NotLeaderOrFollowerException("Leader not local for partition %s on broker %d"
             .format(topicPartition, localBrokerId))
       }
     }
 
-    // some delayed operations may be unblocked after HW changed
-    if (leaderHWIncremented)
-      tryCompleteDelayedRequests()
+    info.copy(leaderHwChange = if (leaderHWIncremented) LeaderHwChange.Increased else LeaderHwChange.Same)
+  }
+
+  def readRecords(lastFetchedEpoch: Optional[Integer],
+                  fetchOffset: Long,
+                  currentLeaderEpoch: Optional[Integer],
+                  maxBytes: Int,
+                  fetchIsolation: FetchIsolation,
+                  fetchOnlyFromLeader: Boolean,
+                  minOneMessage: Boolean): LogReadInfo = inReadLock(leaderIsrUpdateLock) {
+    // decide whether to only fetch from leader
+    val localLog = localLogWithEpochOrException(currentLeaderEpoch, fetchOnlyFromLeader)
+
+    // 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
+
+    lastFetchedEpoch.ifPresent { fetchEpoch =>
+      val epochEndOffset = lastOffsetForLeaderEpoch(currentLeaderEpoch, fetchEpoch, fetchOnlyFromLeader = false)
+      val error = Errors.forCode(epochEndOffset.errorCode)
+      if (error != Errors.NONE) {
+        throw error.exception()
+      }
+
+      if (epochEndOffset.endOffset == UNDEFINED_EPOCH_OFFSET || epochEndOffset.leaderEpoch == UNDEFINED_EPOCH) {
+        throw new OffsetOutOfRangeException("Could not determine the end offset of the last fetched epoch " +
+          s"$lastFetchedEpoch from the request")
+      }
 
-    info
+      if (epochEndOffset.leaderEpoch < fetchEpoch || epochEndOffset.endOffset < fetchOffset) {
+        val emptyFetchData = FetchDataInfo(
+          fetchOffsetMetadata = LogOffsetMetadata(fetchOffset),
+          records = MemoryRecords.EMPTY,
+          firstEntryIncomplete = false,
+          abortedTransactions = None
+        )
+
+        val divergingEpoch = new FetchResponseData.EpochEndOffset()
+          .setEpoch(epochEndOffset.leaderEpoch)
+          .setEndOffset(epochEndOffset.endOffset)
+
+        return LogReadInfo(
+          fetchedData = emptyFetchData,
+          divergingEpoch = Some(divergingEpoch),
+          highWatermark = initialHighWatermark,
+          logStartOffset = initialLogStartOffset,
+          logEndOffset = initialLogEndOffset,
+          lastStableOffset = initialLastStableOffset)
+      }
+    }
+
+    val fetchedData = localLog.read(fetchOffset, maxBytes, fetchIsolation, minOneMessage)
+    LogReadInfo(
+      fetchedData = fetchedData,
+      divergingEpoch = None,
+      highWatermark = initialHighWatermark,
+      logStartOffset = initialLogStartOffset,
+      logEndOffset = initialLogEndOffset,
+      lastStableOffset = initialLastStableOffset)
+  }
+
+  def fetchOffsetForTimestamp(timestamp: Long,
+                              isolationLevel: Option[IsolationLevel],
+                              currentLeaderEpoch: Optional[Integer],
+                              fetchOnlyFromLeader: Boolean): Option[TimestampAndOffset] = inReadLock(leaderIsrUpdateLock) {
+    // decide whether to only fetch from leader
+    val localLog = localLogWithEpochOrException(currentLeaderEpoch, fetchOnlyFromLeader)
+
+    val lastFetchableOffset = isolationLevel match {
+      case Some(IsolationLevel.READ_COMMITTED) => localLog.lastStableOffset
+      case Some(IsolationLevel.READ_UNCOMMITTED) => localLog.highWatermark
+      case None => localLog.logEndOffset
+    }
+
+    val epochLogString = if (currentLeaderEpoch.isPresent) {
+      s"epoch ${currentLeaderEpoch.get}"
+    } else {
+      "unknown epoch"
+    }
+
+    // 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 > 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 (${localLog.highWatermark}) is lagging behind the " +
+        s"start offset from the beginning of this epoch ($epochStart)."))
+
+    def getOffsetByTimestamp: Option[TimestampAndOffset] = {
+      logManager.getLog(topicPartition).flatMap(log => log.fetchOffsetByTimestamp(timestamp))
+    }
+
+    // If we're in the lagging HW state after a leader election, throw OffsetNotAvailable for "latest" offset
+    // or for a timestamp lookup that is beyond the last fetchable offset.
+    timestamp match {
+      case ListOffsetRequest.LATEST_TIMESTAMP =>
+        maybeOffsetsError.map(e => throw e)
+          .orElse(Some(new TimestampAndOffset(RecordBatch.NO_TIMESTAMP, lastFetchableOffset, Optional.of(leaderEpoch))))
+      case ListOffsetRequest.EARLIEST_TIMESTAMP =>
+        getOffsetByTimestamp
+      case _ =>
+        getOffsetByTimestamp.filter(timestampAndOffset => timestampAndOffset.offset < lastFetchableOffset)
+          .orElse(maybeOffsetsError.map(e => throw e))
+    }
+  }
+
+  def fetchOffsetSnapshot(currentLeaderEpoch: Optional[Integer],
+                          fetchOnlyFromLeader: Boolean): LogOffsetSnapshot = inReadLock(leaderIsrUpdateLock) {
+    // decide whether to only fetch from leader
+    val localLog = localLogWithEpochOrException(currentLeaderEpoch, fetchOnlyFromLeader)
+    localLog.fetchOffsetSnapshot
+  }
+
+  def legacyFetchOffsetsForTimestamp(timestamp: Long,
+                                     maxNumOffsets: Int,
+                                     isFromConsumer: Boolean,
+                                     fetchOnlyFromLeader: Boolean): Seq[Long] = inReadLock(leaderIsrUpdateLock) {
+    val localLog = localLogWithEpochOrException(Optional.empty(), fetchOnlyFromLeader)
+    val allOffsets = localLog.legacyFetchOffsetsBefore(timestamp, maxNumOffsets)
+
+    if (!isFromConsumer) {
+      allOffsets
+    } else {
+      val hw = localLog.highWatermark
+      if (allOffsets.exists(_ > hw))
+        hw +: allOffsets.dropWhile(_ > hw)
+      else
+        allOffsets
+    }
   }
 
   def logStartOffset: Long = {
     inReadLock(leaderIsrUpdateLock) {
-      leaderReplicaIfLocal.map(_.log.get.logStartOffset).getOrElse(-1)
+      leaderLogIfLocal.map(_.logStartOffset).getOrElse(-1)
     }
   }
 
@@ -603,18 +1236,26 @@ class Partition(val topic: String,
    *
    * Return low watermark of the partition.
    */
-  def deleteRecordsOnLeader(offset: Long): Long = {
-    inReadLock(leaderIsrUpdateLock) {
-      leaderReplicaIfLocal match {
-        case Some(leaderReplica) =>
-          if (!leaderReplica.log.get.config.delete)
-            throw new PolicyViolationException("Records of partition %s can not be deleted due to the configured policy".format(topicPartition))
-          leaderReplica.maybeIncrementLogStartOffset(offset)
-          lowWatermarkIfLeader
-        case None =>
-          throw new NotLeaderForPartitionException("Leader not local for partition %s on broker %d"
-            .format(topicPartition, localBrokerId))
-      }
+  def deleteRecordsOnLeader(offset: Long): LogDeleteRecordsResult = inReadLock(leaderIsrUpdateLock) {
+    leaderLogIfLocal match {
+      case Some(leaderLog) =>
+        if (!leaderLog.config.delete)
+          throw new PolicyViolationException(s"Records of partition $topicPartition can not be deleted due to the configured policy")
+
+        val convertedOffset = if (offset == DeleteRecordsRequest.HIGH_WATERMARK)
+          leaderLog.highWatermark
+        else
+          offset
+
+        if (convertedOffset < 0)
+          throw new OffsetOutOfRangeException(s"The offset $convertedOffset for partition $topicPartition is not valid")
+
+        leaderLog.maybeIncrementLogStartOffset(convertedOffset, ClientRecordDeletion)
+        LogDeleteRecordsResult(
+          requestedOffset = convertedOffset,
+          lowWatermark = lowWatermarkIfLeader)
+      case None =>
+        throw new NotLeaderOrFollowerException(s"Leader not local for partition $topicPartition on broker $localBrokerId")
     }
   }
 
@@ -624,9 +1265,9 @@ class Partition(val topic: String,
     * @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.
+    // is executing maybeReplaceCurrentWithFutureReplica() to replace follower replica with the future replica.
     inReadLock(leaderIsrUpdateLock) {
       logManager.truncateTo(Map(topicPartition -> offset), isFuture = isFuture)
     }
@@ -638,71 +1279,217 @@ class Partition(val topic: String,
     * @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.
+    // is executing maybeReplaceCurrentWithFutureReplica() to replace follower replica with the future replica.
     inReadLock(leaderIsrUpdateLock) {
       logManager.truncateFullyAndStartAt(topicPartition, newOffset, isFuture = isFuture)
     }
   }
 
   /**
-    * @param leaderEpoch Requested leader epoch
-    * @return The last offset of messages published under this leader epoch.
-    */
-  def lastOffsetForLeaderEpoch(leaderEpoch: Int): EpochEndOffset = {
+   * Find the (exclusive) last offset of the largest epoch less than or equal to the requested epoch.
+   *
+   * @param currentLeaderEpoch The expected epoch of the current leader (if known)
+   * @param leaderEpoch Requested leader epoch
+   * @param fetchOnlyFromLeader Whether or not to require servicing only from the leader
+   *
+   * @return The requested leader epoch and the end offset of this leader epoch, or if the requested
+   *         leader epoch is unknown, the leader epoch less than the requested leader epoch and the end offset
+   *         of this leader epoch. The end offset of a leader epoch is defined as the start
+   *         offset of the first leader epoch larger than the leader epoch, or else the log end
+   *         offset if the leader epoch is the latest leader epoch.
+   */
+  def lastOffsetForLeaderEpoch(currentLeaderEpoch: Optional[Integer],
+                               leaderEpoch: Int,
+                               fetchOnlyFromLeader: Boolean): EpochEndOffset = {
     inReadLock(leaderIsrUpdateLock) {
-      leaderReplicaIfLocal match {
-        case Some(leaderReplica) =>
-          new EpochEndOffset(NONE, leaderReplica.epochs.get.endOffsetFor(leaderEpoch))
-        case None =>
-          new EpochEndOffset(NOT_LEADER_FOR_PARTITION, UNDEFINED_EPOCH_OFFSET)
+      val localLogOrError = getLocalLog(currentLeaderEpoch, fetchOnlyFromLeader)
+      localLogOrError match {
+        case Left(localLog) =>
+          localLog.endOffsetForEpoch(leaderEpoch) match {
+            case Some(epochAndOffset) => new EpochEndOffset()
+              .setPartition(partitionId)
+              .setErrorCode(Errors.NONE.code)
+              .setLeaderEpoch(epochAndOffset.leaderEpoch)
+              .setEndOffset(epochAndOffset.offset)
+            case None => new EpochEndOffset()
+              .setPartition(partitionId)
+              .setErrorCode(Errors.NONE.code)
+          }
+        case Right(error) => new EpochEndOffset()
+          .setPartition(partitionId)
+          .setErrorCode(error.code)
       }
     }
   }
 
-  private def updateIsr(newIsr: Set[Replica]) {
-    val newLeaderAndIsr = new LeaderAndIsr(localBrokerId, leaderEpoch, newIsr.map(_.brokerId).toList, zkVersion)
-    val (updateSucceeded,newVersion) = ReplicationUtils.updateLeaderAndIsr(zkClient, topic, partitionId,
-      newLeaderAndIsr, controllerEpoch, zkVersion)
+  private[cluster] def expandIsr(newInSyncReplica: Int): Unit = {
+    if (useAlterIsr) {
+      expandIsrWithAlterIsr(newInSyncReplica)
+    } else {
+      expandIsrWithZk(newInSyncReplica)
+    }
+  }
 
-    if(updateSucceeded) {
-      replicaManager.recordIsrChange(topicPartition)
-      inSyncReplicas = newIsr
-      zkVersion = newVersion
-      trace("ISR updated to [%s] and zkVersion updated to [%d]".format(newIsr.mkString(","), zkVersion))
+  private def expandIsrWithAlterIsr(newInSyncReplica: Int): Unit = {
+    // This is called from maybeExpandIsr which holds the ISR write lock
+    if (!isrState.isInflight) {
+      // When expanding the ISR, we can safely assume the new replica will make it into the ISR since this puts us in
+      // a more constrained state for advancing the HW.
+      sendAlterIsrRequest(PendingExpandIsr(isrState.isr, newInSyncReplica))
     } else {
-      replicaManager.failedIsrUpdatesRate.mark()
-      info("Cached zkVersion [%d] not equal to that in zookeeper, skip updating ISR".format(zkVersion))
+      trace(s"ISR update in-flight, not adding new in-sync replica $newInSyncReplica")
     }
   }
 
+  private def expandIsrWithZk(newInSyncReplica: Int): Unit = {
+    val newInSyncReplicaIds = isrState.isr + newInSyncReplica
+    info(s"Expanding ISR from ${isrState.isr.mkString(",")} to ${newInSyncReplicaIds.mkString(",")}")
+    val newLeaderAndIsr = new LeaderAndIsr(localBrokerId, leaderEpoch, newInSyncReplicaIds.toList, zkVersion)
+    val zkVersionOpt = stateStore.expandIsr(controllerEpoch, newLeaderAndIsr)
+    if (zkVersionOpt.isDefined) {
+      isrChangeListener.markExpand()
+    }
+    maybeUpdateIsrAndVersionWithZk(newInSyncReplicaIds, zkVersionOpt)
+  }
+
+  private[cluster] def shrinkIsr(outOfSyncReplicas: Set[Int]): Unit = {
+    if (useAlterIsr) {
+      shrinkIsrWithAlterIsr(outOfSyncReplicas)
+    } else {
+      shrinkIsrWithZk(isrState.isr -- outOfSyncReplicas)
+    }
+  }
+
+  private def shrinkIsrWithAlterIsr(outOfSyncReplicas: Set[Int]): Unit = {
+    // This is called from maybeShrinkIsr which holds the ISR write lock
+    if (!isrState.isInflight) {
+      // When shrinking the ISR, we cannot assume that the update will succeed as this could erroneously advance the HW
+      // We update pendingInSyncReplicaIds here simply to prevent any further ISR updates from occurring until we get
+      // the next LeaderAndIsr
+      sendAlterIsrRequest(PendingShrinkIsr(isrState.isr, outOfSyncReplicas))
+    } else {
+      trace(s"ISR update in-flight, not removing out-of-sync replicas $outOfSyncReplicas")
+    }
+  }
+
+  private def shrinkIsrWithZk(newIsr: Set[Int]): Unit = {
+    val newLeaderAndIsr = new LeaderAndIsr(localBrokerId, leaderEpoch, newIsr.toList, zkVersion)
+    val zkVersionOpt = stateStore.shrinkIsr(controllerEpoch, newLeaderAndIsr)
+    if (zkVersionOpt.isDefined) {
+      isrChangeListener.markShrink()
+    }
+    maybeUpdateIsrAndVersionWithZk(newIsr, zkVersionOpt)
+  }
+
+  private def maybeUpdateIsrAndVersionWithZk(isr: Set[Int], zkVersionOpt: Option[Int]): Unit = {
+    zkVersionOpt match {
+      case Some(newVersion) =>
+        isrState = CommittedIsr(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")
+        isrChangeListener.markFailed()
+    }
+  }
+
+  private def sendAlterIsrRequest(proposedIsrState: IsrState): Unit = {
+    val isrToSend: Set[Int] = proposedIsrState match {
+      case PendingExpandIsr(isr, newInSyncReplicaId) => isr + newInSyncReplicaId
+      case PendingShrinkIsr(isr, outOfSyncReplicaIds) => isr -- outOfSyncReplicaIds
+      case state =>
+        throw new IllegalStateException(s"Invalid state $state for `AlterIsr` request for partition $topicPartition")
+    }
+
+    val newLeaderAndIsr = new LeaderAndIsr(localBrokerId, leaderEpoch, isrToSend.toList, zkVersion)
+    val alterIsrItem = AlterIsrItem(topicPartition, newLeaderAndIsr, handleAlterIsrResponse(proposedIsrState))
+
+    if (!alterIsrManager.enqueue(alterIsrItem)) {
+      isrChangeListener.markFailed()
+      throw new IllegalStateException(s"Failed to enqueue `AlterIsr` request with state " +
+        s"$newLeaderAndIsr for partition $topicPartition")
+    }
+
+    isrState = proposedIsrState
+    debug(s"Sent `AlterIsr` request to change state to $newLeaderAndIsr after transition to $proposedIsrState")
+  }
+
   /**
-   * remove deleted log metrics
+   * This is called for each partition in the body of an AlterIsr response. For errors which are non-retryable we simply
+   * give up. This leaves [[Partition.isrState]] in an in-flight state (either pending shrink or pending expand).
+   * Since our error was non-retryable we are okay staying in this state until we see new metadata from UpdateMetadata
+   * or LeaderAndIsr
    */
-  def removePartitionMetrics() {
-    removeMetric("UnderReplicated", tags)
-    removeMetric("UnderMinIsr", tags)
-    removeMetric("InSyncReplicasCount", tags)
-    removeMetric("ReplicasCount", tags)
-    removeMetric("LastStableOffsetLag", tags)
+  private def handleAlterIsrResponse(proposedIsrState: IsrState)(result: Either[Errors, LeaderAndIsr]): Unit = {
+    inWriteLock(leaderIsrUpdateLock) {
+      if (isrState != proposedIsrState) {
+        // This means isrState was updated through leader election or some other mechanism before we got the AlterIsr
+        // response. We don't know what happened on the controller exactly, but we do know this response is out of date
+        // so we ignore it.
+        debug(s"Ignoring failed ISR update to $proposedIsrState since we have already updated state to $isrState")
+        return
+      }
+
+      result match {
+        case Left(error: Errors) =>
+          isrChangeListener.markFailed()
+          error match {
+            case Errors.UNKNOWN_TOPIC_OR_PARTITION =>
+              debug(s"Controller failed to update ISR to $proposedIsrState since it doesn't know about this topic or partition. Giving up.")
+            case Errors.FENCED_LEADER_EPOCH =>
+              debug(s"Controller failed to update ISR to $proposedIsrState since we sent an old leader epoch. Giving up.")
+            case Errors.INVALID_UPDATE_VERSION =>
+              debug(s"Controller failed to update ISR to $proposedIsrState due to invalid zk version. Giving up.")
+            case _ =>
+              warn(s"Controller failed to update ISR to $proposedIsrState due to unexpected $error. Retrying.")
+              sendAlterIsrRequest(proposedIsrState)
+          }
+        case Right(leaderAndIsr: LeaderAndIsr) =>
+          // Success from controller, still need to check a few things
+          if (leaderAndIsr.leaderEpoch != leaderEpoch) {
+            debug(s"Ignoring ISR from AlterIsr with ${leaderAndIsr} since we have a stale leader epoch $leaderEpoch.")
+            isrChangeListener.markFailed()
+          } else if (leaderAndIsr.zkVersion <= zkVersion) {
+            debug(s"Ignoring ISR from AlterIsr with ${leaderAndIsr} since we have a newer version $zkVersion.")
+            isrChangeListener.markFailed()
+          } else {
+            isrState = CommittedIsr(leaderAndIsr.isr.toSet)
+            zkVersion = leaderAndIsr.zkVersion
+            info(s"ISR updated from AlterIsr to ${isrState.isr.mkString(",")} and version updated to [$zkVersion]")
+            proposedIsrState match {
+              case PendingExpandIsr(_, _) => isrChangeListener.markExpand()
+              case PendingShrinkIsr(_, _) => isrChangeListener.markShrink()
+              case _ => // nothing to do, shouldn't get here
+            }
+          }
+      }
+    }
   }
 
   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("; Replicas: " + assignmentState.replicas.mkString(","))
+    partitionString.append("; ISR: " + isrState.isr.mkString(","))
+    assignmentState match {
+      case OngoingReassignmentState(adding, removing, _) =>
+        partitionString.append("; AddingReplicas: " + adding.mkString(","))
+        partitionString.append("; RemovingReplicas: " + removing.mkString(","))
+      case _ =>
+    }
     partitionString.toString
   }
 }
diff --git a/core/src/main/scala/kafka/cluster/Replica.scala b/core/src/main/scala/kafka/cluster/Replica.scala
index e41e389e22de2..ffc1c99147819 100644
--- a/core/src/main/scala/kafka/cluster/Replica.scala
+++ b/core/src/main/scala/kafka/cluster/Replica.scala
@@ -18,26 +18,17 @@
 package kafka.cluster
 
 import kafka.log.Log
+import kafka.server.LogOffsetMetadata
 import kafka.utils.Logging
-import kafka.server.{LogOffsetMetadata, LogReadResult}
-import kafka.common.KafkaException
 import org.apache.kafka.common.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)
+
+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
+  @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
@@ -51,14 +42,13 @@ 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
+  def logStartOffset: Long = _logStartOffset
 
-  def lastCaughtUpTimeMs = _lastCaughtUpTimeMs
+  def logEndOffsetMetadata: LogOffsetMetadata = _logEndOffsetMetadata
 
-  val epochs = log.map(_.leaderEpochCache)
+  def logEndOffset: Long = logEndOffsetMetadata.messageOffset
 
-  info(s"Replica loaded for partition $topicPartition with initial high watermark $initialHighWatermarkValue")
-  log.foreach(_.onHighWatermarkIncremented(initialHighWatermarkValue))
+  def lastCaughtUpTimeMs: Long = _lastCaughtUpTimeMs
 
   /*
    * If the FetchRequest reads up to the log end offset of the leader when the current fetch request is received,
@@ -72,105 +62,41 @@ 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): Unit = {
+    if (followerFetchOffsetMetadata.messageOffset >= leaderEndOffset)
+      _lastCaughtUpTimeMs = math.max(_lastCaughtUpTimeMs, followerFetchTimeMs)
+    else if (followerFetchOffsetMetadata.messageOffset >= lastFetchLeaderLogEndOffset)
       _lastCaughtUpTimeMs = math.max(_lastCaughtUpTimeMs, lastFetchTimeMs)
 
-    logStartOffset = logReadResult.followerLogStartOffset
-    logEndOffset = logReadResult.info.fetchOffsetMetadata
-    lastFetchLeaderLogEndOffset = logReadResult.leaderLogEndOffset
-    lastFetchTimeMs = logReadResult.fetchTimeMs
+    _logStartOffset = followerStartOffset
+    _logEndOffsetMetadata = followerFetchOffsetMetadata
+    lastFetchLeaderLogEndOffset = leaderEndOffset
+    lastFetchTimeMs = followerFetchTimeMs
   }
 
-  def resetLastCaughtUpTime(curLeaderLogEndOffset: Long, curTimeMs: Long, lastCaughtUpTimeMs: Long) {
+  def resetLastCaughtUpTime(curLeaderLogEndOffset: Long, curTimeMs: Long, lastCaughtUpTimeMs: Long): Unit = {
     lastFetchLeaderLogEndOffset = curLeaderLogEndOffset
     lastFetchTimeMs = curTimeMs
     _lastCaughtUpTimeMs = lastCaughtUpTimeMs
+    trace(s"Reset state of replica to $this")
   }
 
-  private def logEndOffset_=(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 logEndOffset: LogOffsetMetadata =
-    if (isLocal)
-      log.get.logEndOffsetMetadata
-    else
-      logEndOffsetMetadata
-
-  /**
-   * 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) {
-      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")
-    }
-  }
-
-  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"))
-  }
-
-  def convertHWToLocalOffsetMetadata() = {
-    if (isLocal) {
-      highWatermarkMetadata = log.get.convertToOffsetMetadata(highWatermarkMetadata.messageOffset)
-    } else {
-      throw new KafkaException(s"Should not construct complete high watermark on partition $topicPartition's non-local replica $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", lastCaughtUpTimeMs=$lastCaughtUpTimeMs")
+    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(")")
+    replicaString.toString
   }
 
   override def equals(that: Any): Boolean = that match {
@@ -179,18 +105,4 @@ class Replica(val brokerId: Int,
   }
 
   override def hashCode: Int = 31 + topicPartition.hashCode + 17 * brokerId
-
-  override def toString: String = {
-    val replicaString = new StringBuilder
-    replicaString.append("ReplicaId: " + brokerId)
-    replicaString.append("; Topic: " + topicPartition.topic)
-    replicaString.append("; Partition: " + topicPartition.partition)
-    replicaString.append("; isLocal: " + isLocal)
-    replicaString.append("; lastCaughtUpTimeMs: " + lastCaughtUpTimeMs)
-    if (isLocal) {
-      replicaString.append("; Highwatermark: " + highWatermark)
-      replicaString.append("; LastStableOffset: " + lastStableOffset)
-    }
-    replicaString.toString
-  }
 }
diff --git a/core/src/main/scala/kafka/common/AppInfo.scala b/core/src/main/scala/kafka/common/AppInfo.scala
deleted file mode 100644
index f77bdf55a69b0..0000000000000
--- a/core/src/main/scala/kafka/common/AppInfo.scala
+++ /dev/null
@@ -1,54 +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.common
-
-import com.yammer.metrics.core.Gauge
-import kafka.metrics.KafkaMetricsGroup
-import org.apache.kafka.common.utils.AppInfoParser
-
-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
-    }
-
-  }
-}
diff --git a/core/src/main/scala/kafka/common/AuthorizationException.scala b/core/src/main/scala/kafka/common/AuthorizationException.scala
deleted file mode 100644
index 55919a59896aa..0000000000000
--- a/core/src/main/scala/kafka/common/AuthorizationException.scala
+++ /dev/null
@@ -1,36 +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.common
-
-/**
- * Exception thrown when a principal is not authorized to perform an operation.
- * @param message
- */
-abstract class AuthorizationException(message: String) extends RuntimeException(message) {
-}
-
-class TopicAuthorizationException(message: String) extends AuthorizationException(message) {
-  def this() = this(null)
-}
-
-class GroupAuthorizationException(message: String) extends AuthorizationException(message) {
-  def this() = this(null)
-}
-
-class ClusterAuthorizationException(message: String) extends AuthorizationException(message) {
-  def this() = this(null)
-}
diff --git a/core/src/main/scala/kafka/common/BrokerNotAvailableException.scala b/core/src/main/scala/kafka/common/BrokerNotAvailableException.scala
deleted file mode 100644
index 611bed6f3181f..0000000000000
--- a/core/src/main/scala/kafka/common/BrokerNotAvailableException.scala
+++ /dev/null
@@ -1,22 +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.common
-
-class BrokerNotAvailableException(message: String) extends RuntimeException(message) {
-  def this() = this(null)
-}
\ No newline at end of file
diff --git a/core/src/main/scala/kafka/common/ClientIdAndTopic.scala b/core/src/main/scala/kafka/common/ClientIdAndTopic.scala
deleted file mode 100644
index 5825aad2c8d1a..0000000000000
--- a/core/src/main/scala/kafka/common/ClientIdAndTopic.scala
+++ /dev/null
@@ -1,35 +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.common
-
-/**
- * Convenience case class since (clientId, topic) pairs are used in the creation
- * of many Stats objects.
- */
-trait ClientIdTopic {
-}
-
-case class ClientIdAndTopic(clientId: String, topic: String) extends ClientIdTopic {
-  override def toString = "%s-%s".format(clientId, topic)
-}
-
-case class ClientIdAllTopics(clientId: String) extends ClientIdTopic {
-  override def toString = "%s-%s".format(clientId, "AllTopics")
-}
-
-
diff --git a/core/src/main/scala/kafka/common/Config.scala b/core/src/main/scala/kafka/common/Config.scala
index d24fb0df6352c..f56cca8bd0528 100644
--- a/core/src/main/scala/kafka/common/Config.scala
+++ b/core/src/main/scala/kafka/common/Config.scala
@@ -19,18 +19,19 @@ package kafka.common
 
 import util.matching.Regex
 import kafka.utils.Logging
+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 + "*")
 
     rgx.findFirstIn(value) match {
       case Some(t) =>
         if (!t.equals(value))
-          throw new InvalidConfigException(prop + " " + value + " is illegal, contains a character other than ASCII alphanumerics, '.', '_' and '-'")
-      case None => throw new InvalidConfigException(prop + " " + value + " is illegal, contains a character other than ASCII alphanumerics, '.', '_' and '-'")
+          throw new InvalidConfigurationException(prop + " " + value + " is illegal, contains a character other than ASCII alphanumerics, '.', '_' and '-'")
+      case None => throw new InvalidConfigurationException(prop + " " + value + " is illegal, contains a character other than ASCII alphanumerics, '.', '_' and '-'")
     }
   }
 }
diff --git a/core/src/main/scala/kafka/common/ConsumerCoordinatorNotAvailableException.scala b/core/src/main/scala/kafka/common/ConsumerCoordinatorNotAvailableException.scala
deleted file mode 100644
index 8e02d264e9447..0000000000000
--- a/core/src/main/scala/kafka/common/ConsumerCoordinatorNotAvailableException.scala
+++ /dev/null
@@ -1,22 +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.common
-
-class ConsumerCoordinatorNotAvailableException(message: String) extends RuntimeException(message) {
-  def this() = this(null)
-}
\ No newline at end of file
diff --git a/core/src/main/scala/kafka/common/ConsumerRebalanceFailedException.scala b/core/src/main/scala/kafka/common/ConsumerRebalanceFailedException.scala
deleted file mode 100644
index 2f4c2cf70220a..0000000000000
--- a/core/src/main/scala/kafka/common/ConsumerRebalanceFailedException.scala
+++ /dev/null
@@ -1,26 +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.common
-
-/**
- * Thrown when a request is made for broker but no brokers with that topic
- * exist.
- */
-class ConsumerRebalanceFailedException(message: String) extends RuntimeException(message) {
-  def this() = this(null)
-}
diff --git a/core/src/main/scala/kafka/common/ControllerMovedException.scala b/core/src/main/scala/kafka/common/ControllerMovedException.scala
deleted file mode 100644
index 39cf36d34059a..0000000000000
--- a/core/src/main/scala/kafka/common/ControllerMovedException.scala
+++ /dev/null
@@ -1,23 +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.common
-
-class ControllerMovedException(message: String, cause: Throwable) extends RuntimeException(message, cause) {
-  def this(message: String) = this(message, null)
-  def this() = this(null, null)
-}
\ No newline at end of file
diff --git a/core/src/main/scala/kafka/common/ErrorMapping.scala b/core/src/main/scala/kafka/common/ErrorMapping.scala
deleted file mode 100644
index 9f290731a7075..0000000000000
--- a/core/src/main/scala/kafka/common/ErrorMapping.scala
+++ /dev/null
@@ -1,117 +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.common
-
-import java.nio.ByteBuffer
-
-import kafka.message.InvalidMessageException
-import org.apache.kafka.common.errors.InvalidTopicException
-
-import scala.Predef._
-
-/**
- * A bi-directional mapping between error codes and exceptions
- */
-object ErrorMapping {
-  val EmptyByteBuffer = ByteBuffer.allocate(0)
-
-  val UnknownCode: Short = -1
-  val NoError: Short = 0
-  val OffsetOutOfRangeCode: Short = 1
-  val InvalidMessageCode: Short = 2
-  val UnknownTopicOrPartitionCode: Short = 3
-  val InvalidFetchSizeCode: Short = 4
-  val LeaderNotAvailableCode: Short = 5
-  val NotLeaderForPartitionCode: Short = 6
-  val RequestTimedOutCode: Short = 7
-  val BrokerNotAvailableCode: Short = 8
-  val ReplicaNotAvailableCode: Short = 9
-  val MessageSizeTooLargeCode: Short = 10
-  val StaleControllerEpochCode: Short = 11
-  val OffsetMetadataTooLargeCode: Short = 12
-  val StaleLeaderEpochCode: Short = 13
-  val OffsetsLoadInProgressCode: Short = 14
-  val ConsumerCoordinatorNotAvailableCode: Short = 15
-  val NotCoordinatorForConsumerCode: Short = 16
-  val InvalidTopicCode: Short = 17
-  val MessageSetSizeTooLargeCode: Short = 18
-  val NotEnoughReplicasCode: Short = 19
-  val NotEnoughReplicasAfterAppendCode: Short = 20
-  // 21: InvalidRequiredAcks
-  // 22: IllegalConsumerGeneration
-  // 23: INCONSISTENT_PARTITION_ASSIGNMENT_STRATEGY
-  // 24: UNKNOWN_PARTITION_ASSIGNMENT_STRATEGY
-  // 25: UNKNOWN_CONSUMER_ID
-  // 26: INVALID_SESSION_TIMEOUT
-  // 27: REBALANCE_IN_PROGRESS
-  // 28: INVALID_COMMIT_OFFSET_SIZE
-  val TopicAuthorizationCode: Short = 29
-  val GroupAuthorizationCode: Short = 30
-  val ClusterAuthorizationCode: Short = 31
-  // 32: INVALID_TIMESTAMP
-  // 33: UNSUPPORTED_SASL_MECHANISM
-  // 34: ILLEGAL_SASL_STATE
-  // 35: UNSUPPORTED_VERSION
-  // 36: TOPIC_ALREADY_EXISTS
-  // 37: INVALID_PARTITIONS
-  // 38: INVALID_REPLICATION_FACTOR
-  // 39: INVALID_REPLICA_ASSIGNMENT
-  // 40: INVALID_CONFIG
-  // 41: NOT_CONTROLLER
-  // 42: INVALID_REQUEST
-
-  private val exceptionToCode =
-    Map[Class[Throwable], Short](
-      classOf[OffsetOutOfRangeException].asInstanceOf[Class[Throwable]] -> OffsetOutOfRangeCode,
-      classOf[InvalidMessageException].asInstanceOf[Class[Throwable]] -> InvalidMessageCode,
-      classOf[UnknownTopicOrPartitionException].asInstanceOf[Class[Throwable]] -> UnknownTopicOrPartitionCode,
-      classOf[InvalidMessageSizeException].asInstanceOf[Class[Throwable]] -> InvalidFetchSizeCode,
-      classOf[LeaderNotAvailableException].asInstanceOf[Class[Throwable]] -> LeaderNotAvailableCode,
-      classOf[NotLeaderForPartitionException].asInstanceOf[Class[Throwable]] -> NotLeaderForPartitionCode,
-      classOf[RequestTimedOutException].asInstanceOf[Class[Throwable]] -> RequestTimedOutCode,
-      classOf[BrokerNotAvailableException].asInstanceOf[Class[Throwable]] -> BrokerNotAvailableCode,
-      classOf[ReplicaNotAvailableException].asInstanceOf[Class[Throwable]] -> ReplicaNotAvailableCode,
-      classOf[MessageSizeTooLargeException].asInstanceOf[Class[Throwable]] -> MessageSizeTooLargeCode,
-      classOf[ControllerMovedException].asInstanceOf[Class[Throwable]] -> StaleControllerEpochCode,
-      classOf[OffsetMetadataTooLargeException].asInstanceOf[Class[Throwable]] -> OffsetMetadataTooLargeCode,
-      classOf[OffsetsLoadInProgressException].asInstanceOf[Class[Throwable]] -> OffsetsLoadInProgressCode,
-      classOf[ConsumerCoordinatorNotAvailableException].asInstanceOf[Class[Throwable]] -> ConsumerCoordinatorNotAvailableCode,
-      classOf[NotCoordinatorForConsumerException].asInstanceOf[Class[Throwable]] -> NotCoordinatorForConsumerCode,
-      classOf[InvalidTopicException].asInstanceOf[Class[Throwable]] -> InvalidTopicCode,
-      classOf[MessageSetSizeTooLargeException].asInstanceOf[Class[Throwable]] -> MessageSetSizeTooLargeCode,
-      classOf[NotEnoughReplicasException].asInstanceOf[Class[Throwable]] -> NotEnoughReplicasCode,
-      classOf[NotEnoughReplicasAfterAppendException].asInstanceOf[Class[Throwable]] -> NotEnoughReplicasAfterAppendCode,
-      classOf[TopicAuthorizationException].asInstanceOf[Class[Throwable]] -> TopicAuthorizationCode,
-      classOf[GroupAuthorizationException].asInstanceOf[Class[Throwable]] -> GroupAuthorizationCode,
-      classOf[ClusterAuthorizationException].asInstanceOf[Class[Throwable]] -> ClusterAuthorizationCode
-    ).withDefaultValue(UnknownCode)
-
-  /* invert the mapping */
-  private val codeToException =
-    (Map[Short, Class[Throwable]]() ++ exceptionToCode.iterator.map(p => (p._2, p._1))).withDefaultValue(classOf[UnknownException])
-
-  def codeFor(exception: Class[Throwable]): Short = exceptionToCode(exception)
-
-  def maybeThrowException(code: Short) =
-    if(code != 0)
-      throw codeToException(code).newInstance()
-
-  def exceptionFor(code: Short) : Throwable = codeToException(code).newInstance()
-
-  def exceptionNameFor(code: Short) : String = codeToException(code).getName()
-}
diff --git a/core/src/main/scala/kafka/common/FailedToSendMessageException.scala b/core/src/main/scala/kafka/common/FailedToSendMessageException.scala
deleted file mode 100644
index 6aee4142d843a..0000000000000
--- a/core/src/main/scala/kafka/common/FailedToSendMessageException.scala
+++ /dev/null
@@ -1,23 +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.common
-
-/**
- * Indicates a producer pool initialization problem
-*/
-class FailedToSendMessageException(message: String, t: Throwable) extends RuntimeException(message, t) {
-}
diff --git a/core/src/main/scala/kafka/common/InconsistentBrokerMetadataException.scala b/core/src/main/scala/kafka/common/InconsistentBrokerMetadataException.scala
new file mode 100644
index 0000000000000..2b11512e44ce2
--- /dev/null
+++ b/core/src/main/scala/kafka/common/InconsistentBrokerMetadataException.scala
@@ -0,0 +1,27 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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.common
+
+/**
+ * Indicates the BrokerMetadata stored in logDirs is not consistent across logDirs.
+ */
+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/common/InconsistentClusterIdException.scala b/core/src/main/scala/kafka/common/InconsistentClusterIdException.scala
new file mode 100644
index 0000000000000..6868dd8780de9
--- /dev/null
+++ b/core/src/main/scala/kafka/common/InconsistentClusterIdException.scala
@@ -0,0 +1,27 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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.common
+
+/**
+ * Indicates the clusterId stored in logDirs is not consistent with the clusterIs stored in ZK.
+ */
+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/IndexOffsetOverflowException.scala b/core/src/main/scala/kafka/common/IndexOffsetOverflowException.scala
new file mode 100644
index 0000000000000..5dd9b43e9e843
--- /dev/null
+++ b/core/src/main/scala/kafka/common/IndexOffsetOverflowException.scala
@@ -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 kafka.common
+
+/**
+ * Indicates that an attempt was made to append a message whose offset could cause the index offset to overflow.
+ */
+class IndexOffsetOverflowException(message: String, cause: Throwable) extends org.apache.kafka.common.KafkaException(message, cause) {
+  def this(message: String) = this(message, null)
+}
diff --git a/core/src/main/scala/kafka/common/InterBrokerSendThread.scala b/core/src/main/scala/kafka/common/InterBrokerSendThread.scala
index 70dae354aa451..11e1aa8de3920 100644
--- a/core/src/main/scala/kafka/common/InterBrokerSendThread.scala
+++ b/core/src/main/scala/kafka/common/InterBrokerSendThread.scala
@@ -16,24 +16,33 @@
  */
 package kafka.common
 
+import java.util.{ArrayDeque, ArrayList, Collection, Collections, HashMap, Iterator}
+import java.util.Map.Entry
+
 import kafka.utils.ShutdownableThread
-import org.apache.kafka.clients.{ClientResponse, NetworkClient, RequestCompletionHandler}
+import org.apache.kafka.clients.{ClientRequest, ClientResponse, KafkaClient, RequestCompletionHandler}
 import org.apache.kafka.common.Node
+import org.apache.kafka.common.errors.AuthenticationException
 import org.apache.kafka.common.internals.FatalExitError
 import org.apache.kafka.common.requests.AbstractRequest
 import org.apache.kafka.common.utils.Time
 
+import scala.jdk.CollectionConverters._
 
 /**
  *  Class for inter-broker send thread that utilize a non-blocking network client.
  */
 abstract class InterBrokerSendThread(name: String,
-                                     networkClient: NetworkClient,
+                                     networkClient: KafkaClient,
                                      time: Time,
                                      isInterruptible: Boolean = true)
   extends ShutdownableThread(name, isInterruptible) {
 
   def generateRequests(): Iterable[RequestAndCompletionHandler]
+  def requestTimeoutMs: Int
+  private val unsentRequests = new UnsentRequests
+
+  def hasUnsentRequests = unsentRequests.iterator().hasNext
 
   override def shutdown(): Unit = {
     initiateShutdown()
@@ -42,36 +51,28 @@ abstract class InterBrokerSendThread(name: String,
     awaitShutdown()
   }
 
-  override def doWork() {
-    val now = time.milliseconds()
-    var pollTimeout = Long.MaxValue
+  override def doWork(): Unit = {
+    var now = time.milliseconds()
 
-    try {
-      for (request: RequestAndCompletionHandler <- generateRequests()) {
-        val destination = Integer.toString(request.destination.id())
-        val completionHandler = request.handler
-        val clientRequest = networkClient.newClientRequest(destination,
+    generateRequests().foreach { request =>
+      val completionHandler = request.handler
+      unsentRequests.put(request.destination,
+        networkClient.newClientRequest(
+          request.destination.idString,
           request.request,
           now,
           true,
-          completionHandler)
-
-        if (networkClient.ready(request.destination, now)) {
-          networkClient.send(clientRequest, now)
-        } else {
-          val header = clientRequest.makeHeader(request.request.latestAllowedVersion)
-          val disconnectResponse: ClientResponse = new ClientResponse(header, completionHandler, destination,
-            now /* createdTimeMs */ , now /* receivedTimeMs */ , true /* disconnected */ , null /* versionMismatch */ ,
-            null /* responseBody */)
-
-          // poll timeout would be the minimum of connection delay if there are any dest yet to be reached;
-          // otherwise it is infinity
-          pollTimeout = Math.min(pollTimeout, networkClient.connectionDelay(request.destination, now))
+          requestTimeoutMs,
+          completionHandler))
+    }
 
-          completionHandler.onComplete(disconnectResponse)
-        }
-      }
-      networkClient.poll(pollTimeout, now)
+    try {
+      val timeout = sendRequests(now)
+      networkClient.poll(timeout, now)
+      now = time.milliseconds()
+      checkDisconnects(now)
+      failExpiredRequests(now)
+      unsentRequests.clean()
     } catch {
       case e: FatalExitError => throw e
       case t: Throwable =>
@@ -84,9 +85,118 @@ abstract class InterBrokerSendThread(name: String,
     }
   }
 
-  def wakeup(): Unit = networkClient.wakeup()
+  private def sendRequests(now: Long): Long = {
+    var pollTimeout = Long.MaxValue
+    for (node <- unsentRequests.nodes.asScala) {
+      val requestIterator = unsentRequests.requestIterator(node)
+      while (requestIterator.hasNext) {
+        val request = requestIterator.next
+        if (networkClient.ready(node, now)) {
+          networkClient.send(request, now)
+          requestIterator.remove()
+        } else
+          pollTimeout = Math.min(pollTimeout, networkClient.connectionDelay(node, now))
+      }
+    }
+    pollTimeout
+  }
+
+  private def checkDisconnects(now: Long): Unit = {
+    // any disconnects affecting requests that have already been transmitted will be handled
+    // by NetworkClient, so we just need to check whether connections for any of the unsent
+    // requests have been disconnected; if they have, then we complete the corresponding future
+    // and set the disconnect flag in the ClientResponse
+    val iterator = unsentRequests.iterator()
+    while (iterator.hasNext) {
+      val entry = iterator.next
+      val (node, requests) = (entry.getKey, entry.getValue)
+      if (!requests.isEmpty && networkClient.connectionFailed(node)) {
+        iterator.remove()
+        for (request <- requests.asScala) {
+          val authenticationException = networkClient.authenticationException(node)
+          if (authenticationException != null)
+            error(s"Failed to send the following request due to authentication error: $request")
+          completeWithDisconnect(request, now, authenticationException)
+        }
+      }
+    }
+  }
+
+  private def failExpiredRequests(now: Long): Unit = {
+    // clear all expired unsent requests
+    val timedOutRequests = unsentRequests.removeAllTimedOut(now)
+    for (request <- timedOutRequests.asScala) {
+      debug(s"Failed to send the following request after ${request.requestTimeoutMs} ms: $request")
+      completeWithDisconnect(request, now, null)
+    }
+  }
+
+  def completeWithDisconnect(request: ClientRequest,
+                             now: Long,
+                             authenticationException: AuthenticationException): Unit = {
+    val handler = request.callback
+    handler.onComplete(new ClientResponse(request.makeHeader(request.requestBuilder().latestAllowedVersion()),
+      handler, request.destination, now /* createdTimeMs */ , now /* receivedTimeMs */ , true /* disconnected */ ,
+      null /* versionMismatch */ , authenticationException, null))
+  }
 
+  def wakeup(): Unit = networkClient.wakeup()
 }
 
-case class RequestAndCompletionHandler(destination: Node, request: AbstractRequest.Builder[_ <: AbstractRequest],
-                                       handler: RequestCompletionHandler)
\ No newline at end of file
+case class RequestAndCompletionHandler(destination: Node,
+                                       request: AbstractRequest.Builder[_ <: AbstractRequest],
+                                       handler: RequestCompletionHandler)
+
+private class UnsentRequests {
+  private val unsent = new HashMap[Node, ArrayDeque[ClientRequest]]
+
+  def put(node: Node, request: ClientRequest): Unit = {
+    var requests = unsent.get(node)
+    if (requests == null) {
+      requests = new ArrayDeque[ClientRequest]
+      unsent.put(node, requests)
+    }
+    requests.add(request)
+  }
+
+  def removeAllTimedOut(now: Long): Collection[ClientRequest] = {
+    val expiredRequests = new ArrayList[ClientRequest]
+    for (requests <- unsent.values.asScala) {
+      val requestIterator = requests.iterator
+      var foundExpiredRequest = false
+      while (requestIterator.hasNext && !foundExpiredRequest) {
+        val request = requestIterator.next
+        val elapsedMs = Math.max(0, now - request.createdTimeMs)
+        if (elapsedMs > request.requestTimeoutMs) {
+          expiredRequests.add(request)
+          requestIterator.remove()
+          foundExpiredRequest = true
+        }
+      }
+    }
+    expiredRequests
+  }
+
+  def clean(): Unit = {
+    val iterator = unsent.values.iterator
+    while (iterator.hasNext) {
+      val requests = iterator.next
+      if (requests.isEmpty)
+        iterator.remove()
+    }
+  }
+
+  def iterator(): Iterator[Entry[Node, ArrayDeque[ClientRequest]]] = {
+    unsent.entrySet().iterator()
+  }
+
+  def requestIterator(node: Node): Iterator[ClientRequest] = {
+    val requests = unsent.get(node)
+    if (requests == null)
+      Collections.emptyIterator[ClientRequest]
+    else
+      requests.iterator
+  }
+
+  def nodes = unsent.keySet
+}
diff --git a/core/src/main/scala/kafka/common/InvalidConfigException.scala b/core/src/main/scala/kafka/common/InvalidConfigException.scala
deleted file mode 100644
index 643784638deb1..0000000000000
--- a/core/src/main/scala/kafka/common/InvalidConfigException.scala
+++ /dev/null
@@ -1,25 +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.common
-
-/**
- * Indicates that the given config parameter has invalid value
- */
-class InvalidConfigException(message: String) extends RuntimeException(message) {
-  def this() = this(null)
-}
diff --git a/core/src/main/scala/kafka/common/InvalidMessageSizeException.scala b/core/src/main/scala/kafka/common/InvalidMessageSizeException.scala
deleted file mode 100644
index 6a7bb47cc3bf4..0000000000000
--- a/core/src/main/scala/kafka/common/InvalidMessageSizeException.scala
+++ /dev/null
@@ -1,26 +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.common
-
-/**
- * Indicates the client has requested a range no longer available on the server
- */
-class InvalidMessageSizeException(message: String) extends RuntimeException(message) {
-  def this() = this(null)
-}
-
diff --git a/core/src/main/scala/kafka/common/InvalidOffsetException.scala b/core/src/main/scala/kafka/common/InvalidOffsetException.scala
deleted file mode 100644
index c6811d7c0271e..0000000000000
--- a/core/src/main/scala/kafka/common/InvalidOffsetException.scala
+++ /dev/null
@@ -1,22 +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.common
-
-class InvalidOffsetException(message: String) extends RuntimeException(message) {
-  def this() = this(null)
-}
\ No newline at end of file
diff --git a/core/src/main/scala/kafka/common/KafkaException.scala b/core/src/main/scala/kafka/common/KafkaException.scala
index e72d151c09e97..9c34dd9bd78b4 100644
--- a/core/src/main/scala/kafka/common/KafkaException.scala
+++ b/core/src/main/scala/kafka/common/KafkaException.scala
@@ -17,8 +17,10 @@
 package kafka.common
 
 /**
- * Generic Kafka exception
-*/
+ * Usage of this class is discouraged. Use org.apache.kafka.common.KafkaException instead.
+ *
+ * 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/KafkaStorageException.scala b/core/src/main/scala/kafka/common/KafkaStorageException.scala
deleted file mode 100644
index e0ecff3ec45aa..0000000000000
--- a/core/src/main/scala/kafka/common/KafkaStorageException.scala
+++ /dev/null
@@ -1,27 +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.common
-
-/**
- * Kafka exception caused by disk-related IOException
- * This class is deprecated and will be replaced by org.apache.kafka.common.errors.KafkaStorageException
- */
-@Deprecated
-class KafkaStorageException(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/LeaderElectionNotNeededException.scala b/core/src/main/scala/kafka/common/LeaderElectionNotNeededException.scala
deleted file mode 100644
index ca89d2516bff7..0000000000000
--- a/core/src/main/scala/kafka/common/LeaderElectionNotNeededException.scala
+++ /dev/null
@@ -1,27 +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.common
-
-/**
- * This exception is thrown when new leader election is not necessary.
- */
-class LeaderElectionNotNeededException(message: String, cause: Throwable) extends RuntimeException(message, cause) {
-  def this(message: String) = this(message, null)
-  def this() = this(null, null)
-}
-
-
diff --git a/core/src/main/scala/kafka/common/LeaderNotAvailableException.scala b/core/src/main/scala/kafka/common/LeaderNotAvailableException.scala
deleted file mode 100644
index 972728edb0951..0000000000000
--- a/core/src/main/scala/kafka/common/LeaderNotAvailableException.scala
+++ /dev/null
@@ -1,26 +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.common
-
-/**
- * Thrown when a request is made for partition, but no leader exists for that partition
- */
-class LeaderNotAvailableException(message: String, cause: Throwable) extends RuntimeException(message, cause) {
-  def this(message: String) = this(message, null)
-  def this() = this(null, null)
-}
diff --git a/core/src/main/scala/kafka/common/LogSegmentOffsetOverflowException.scala b/core/src/main/scala/kafka/common/LogSegmentOffsetOverflowException.scala
new file mode 100644
index 0000000000000..2de5906109af9
--- /dev/null
+++ b/core/src/main/scala/kafka/common/LogSegmentOffsetOverflowException.scala
@@ -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 kafka.common
+
+import kafka.log.LogSegment
+
+/**
+ * Indicates that the log segment contains one or more messages that overflow the offset (and / or time) index. This is
+ * not a typical scenario, and could only happen when brokers have log segments that were created before the patch for
+ * KAFKA-5413. With KAFKA-6264, we have the ability to split such log segments into multiple log segments such that we
+ * do not have any segments with offset overflow.
+ */
+class LogSegmentOffsetOverflowException(val segment: LogSegment, val offset: Long)
+  extends org.apache.kafka.common.KafkaException(s"Detected offset overflow at offset $offset in segment $segment") {
+}
diff --git a/core/src/main/scala/kafka/common/MessageFormatter.scala b/core/src/main/scala/kafka/common/MessageFormatter.scala
index ef3c7238b7b11..7826eb7264a13 100644
--- a/core/src/main/scala/kafka/common/MessageFormatter.scala
+++ b/core/src/main/scala/kafka/common/MessageFormatter.scala
@@ -17,23 +17,12 @@
 
 package kafka.common
 
-import java.io.PrintStream
-import java.util.Properties
-
-import org.apache.kafka.clients.consumer.ConsumerRecord
-
 /**
   * Typical implementations of this interface convert a `ConsumerRecord` into a type that can then be passed to
   * a `PrintStream`.
   *
   * This is used by the `ConsoleConsumer`.
   */
-trait MessageFormatter {
-
-  def init(props: Properties) {}
-
-  def writeTo(consumerRecord: ConsumerRecord[Array[Byte], Array[Byte]], output: PrintStream): Unit
-
-  def close() {}
-
+@deprecated("This class is deprecated and will be replaced by org.apache.kafka.common.MessageFormatter.", "2.7.0")
+trait MessageFormatter extends org.apache.kafka.common.MessageFormatter {
 }
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/MessageSetSizeTooLargeException.scala b/core/src/main/scala/kafka/common/MessageSetSizeTooLargeException.scala
deleted file mode 100644
index 94a616ed3972a..0000000000000
--- a/core/src/main/scala/kafka/common/MessageSetSizeTooLargeException.scala
+++ /dev/null
@@ -1,22 +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.common
-
-class MessageSetSizeTooLargeException(message: String) extends RuntimeException(message) {
-  def this() = this(null)
-}
diff --git a/core/src/main/scala/kafka/common/MessageSizeTooLargeException.scala b/core/src/main/scala/kafka/common/MessageSizeTooLargeException.scala
deleted file mode 100644
index 2d18324fb7d2e..0000000000000
--- a/core/src/main/scala/kafka/common/MessageSizeTooLargeException.scala
+++ /dev/null
@@ -1,22 +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.common
-
-class MessageSizeTooLargeException(message: String) extends RuntimeException(message) {
-  def this() = this(null)
-}
\ No newline at end of file
diff --git a/core/src/main/scala/kafka/common/MessageStreamsExistException.scala b/core/src/main/scala/kafka/common/MessageStreamsExistException.scala
deleted file mode 100644
index b904ed04d7387..0000000000000
--- a/core/src/main/scala/kafka/common/MessageStreamsExistException.scala
+++ /dev/null
@@ -1,23 +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.common
-
-/**
- * Indicates a createMessageStreams can't be called more than once
-*/
-class MessageStreamsExistException(message: String, t: Throwable) extends RuntimeException(message, t) {
-}
diff --git a/core/src/main/scala/kafka/common/NoBrokersForPartitionException.scala b/core/src/main/scala/kafka/common/NoBrokersForPartitionException.scala
deleted file mode 100644
index 4577b298032ab..0000000000000
--- a/core/src/main/scala/kafka/common/NoBrokersForPartitionException.scala
+++ /dev/null
@@ -1,26 +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.common
-
-/**
- * Thrown when a request is made for broker but no brokers with that topic
- * exist.
- */
-class NoBrokersForPartitionException(message: String) extends RuntimeException(message) {
-  def this() = this(null)
-}
\ No newline at end of file
diff --git a/core/src/main/scala/kafka/common/NoEpochForPartitionException.scala b/core/src/main/scala/kafka/common/NoEpochForPartitionException.scala
deleted file mode 100644
index ef823f2fb58e7..0000000000000
--- a/core/src/main/scala/kafka/common/NoEpochForPartitionException.scala
+++ /dev/null
@@ -1,25 +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.common
-
-/**
- * Thrown when a get epoch request is made for partition, but no epoch exists for that partition
- */
-class NoEpochForPartitionException(message: String) extends RuntimeException(message) {
-  def this() = this(null)
-}
\ No newline at end of file
diff --git a/core/src/main/scala/kafka/common/NoReplicaOnlineException.scala b/core/src/main/scala/kafka/common/NoReplicaOnlineException.scala
deleted file mode 100644
index b66c8fc82c15b..0000000000000
--- a/core/src/main/scala/kafka/common/NoReplicaOnlineException.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.common
-
-
-/**
- * This exception is thrown by the leader elector in the controller when leader election fails for a partition since
- * all the leader candidate replicas for a partition are offline; the set of candidates may or may not be limited
- * to just the in sync replicas depending upon whether unclean leader election is allowed to occur.
- */
-class NoReplicaOnlineException(message: String, cause: Throwable) extends RuntimeException(message, cause) {
-  def this(message: String) = this(message, null)
-  def this() = this(null, null)
-}
\ No newline at end of file
diff --git a/core/src/main/scala/kafka/common/NotAssignedReplicaException.scala b/core/src/main/scala/kafka/common/NotAssignedReplicaException.scala
deleted file mode 100644
index 409d11275280c..0000000000000
--- a/core/src/main/scala/kafka/common/NotAssignedReplicaException.scala
+++ /dev/null
@@ -1,23 +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.common
-
-class NotAssignedReplicaException(message: String, cause: Throwable) extends RuntimeException(message, cause) {
-  def this(message: String) = this(message, null)
-  def this() = this(null, null)
-}
diff --git a/core/src/main/scala/kafka/common/NotCoordinatorForConsumerException.scala b/core/src/main/scala/kafka/common/NotCoordinatorForConsumerException.scala
deleted file mode 100644
index 1eb74be038eaa..0000000000000
--- a/core/src/main/scala/kafka/common/NotCoordinatorForConsumerException.scala
+++ /dev/null
@@ -1,22 +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.common
-
-class NotCoordinatorForConsumerException(message: String) extends RuntimeException(message) {
-  def this() = this(null)
-}
\ No newline at end of file
diff --git a/core/src/main/scala/kafka/common/NotEnoughReplicasAfterAppendException.scala b/core/src/main/scala/kafka/common/NotEnoughReplicasAfterAppendException.scala
deleted file mode 100644
index c4f9def6162e9..0000000000000
--- a/core/src/main/scala/kafka/common/NotEnoughReplicasAfterAppendException.scala
+++ /dev/null
@@ -1,27 +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.common
-
-/**
- * Number of insync replicas for the partition is lower than min.insync.replicas
- * This exception is raised when the low ISR size is discovered *after* the message
- * was already appended to the log. Producer retries will cause duplicates.
- */
-class NotEnoughReplicasAfterAppendException(message: String) extends RuntimeException(message) {
-  def this() = this(null)
-}
diff --git a/core/src/main/scala/kafka/common/NotEnoughReplicasException.scala b/core/src/main/scala/kafka/common/NotEnoughReplicasException.scala
deleted file mode 100644
index bfbe0ee4a5a15..0000000000000
--- a/core/src/main/scala/kafka/common/NotEnoughReplicasException.scala
+++ /dev/null
@@ -1,25 +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.common
-
-/**
- * Message was rejected because number of insync replicas for the partition is lower than min.insync.replicas
- */
-class NotEnoughReplicasException(message: String) extends RuntimeException(message) {
-    def this() = this(null)
-}
diff --git a/core/src/main/scala/kafka/common/NotLeaderForPartitionException.scala b/core/src/main/scala/kafka/common/NotLeaderForPartitionException.scala
deleted file mode 100644
index b4558f89f0a23..0000000000000
--- a/core/src/main/scala/kafka/common/NotLeaderForPartitionException.scala
+++ /dev/null
@@ -1,25 +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.common
-
-/**
- * Thrown when a request is made for partition on a broker that is NOT a leader for that partition
- */
-class NotLeaderForPartitionException(message: String) extends RuntimeException(message) {
-  def this() = this(null)
-}
\ No newline at end of file
diff --git a/core/src/main/scala/kafka/common/OffsetAndMetadata.scala b/core/src/main/scala/kafka/common/OffsetAndMetadata.scala
new file mode 100644
index 0000000000000..632c86376e4da
--- /dev/null
+++ b/core/src/main/scala/kafka/common/OffsetAndMetadata.scala
@@ -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 kafka.common
+
+import java.util.Optional
+
+case class OffsetAndMetadata(offset: Long,
+                             leaderEpoch: Optional[Integer],
+                             metadata: String,
+                             commitTimestamp: Long,
+                             expireTimestamp: Option[Long]) {
+
+
+  override def toString: String  = {
+    s"OffsetAndMetadata(offset=$offset" +
+      s", leaderEpoch=$leaderEpoch" +
+      s", metadata=$metadata" +
+      s", commitTimestamp=$commitTimestamp" +
+      s", expireTimestamp=$expireTimestamp)"
+  }
+}
+
+object OffsetAndMetadata {
+  val NoMetadata: String = ""
+
+  def apply(offset: Long, metadata: String, commitTimestamp: Long): OffsetAndMetadata = {
+    OffsetAndMetadata(offset, Optional.empty(), metadata, commitTimestamp, None)
+  }
+
+  def apply(offset: Long, metadata: String, commitTimestamp: Long, expireTimestamp: Long): OffsetAndMetadata = {
+    OffsetAndMetadata(offset, Optional.empty(), metadata, commitTimestamp, Some(expireTimestamp))
+  }
+}
diff --git a/core/src/main/scala/kafka/common/OffsetMetadataAndError.scala b/core/src/main/scala/kafka/common/OffsetMetadataAndError.scala
deleted file mode 100644
index 2cf9bb40eec09..0000000000000
--- a/core/src/main/scala/kafka/common/OffsetMetadataAndError.scala
+++ /dev/null
@@ -1,81 +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.common
-
-import org.apache.kafka.common.protocol.Errors
-
-case class OffsetMetadata(offset: Long, metadata: String = OffsetMetadata.NoMetadata) {
-  override def toString = "OffsetMetadata[%d,%s]"
-    .format(offset,
-    if (metadata != null && metadata.length > 0) metadata else "NO_METADATA")
-}
-
-object OffsetMetadata {
-  val InvalidOffset: Long = -1L
-  val NoMetadata: String = ""
-
-  val InvalidOffsetMetadata = OffsetMetadata(OffsetMetadata.InvalidOffset, OffsetMetadata.NoMetadata)
-}
-
-case class OffsetAndMetadata(offsetMetadata: OffsetMetadata,
-                             commitTimestamp: Long = org.apache.kafka.common.requests.OffsetCommitRequest.DEFAULT_TIMESTAMP,
-                             expireTimestamp: Long = org.apache.kafka.common.requests.OffsetCommitRequest.DEFAULT_TIMESTAMP) {
-
-  def offset = offsetMetadata.offset
-
-  def metadata = offsetMetadata.metadata
-
-  override def toString = "[%s,CommitTime %d,ExpirationTime %d]".format(offsetMetadata, commitTimestamp, expireTimestamp)
-}
-
-object OffsetAndMetadata {
-  def apply(offset: Long, metadata: String, commitTimestamp: Long, expireTimestamp: Long) = new OffsetAndMetadata(OffsetMetadata(offset, metadata), commitTimestamp, expireTimestamp)
-
-  def apply(offset: Long, metadata: String, timestamp: Long) = new OffsetAndMetadata(OffsetMetadata(offset, metadata), timestamp)
-
-  def apply(offset: Long, metadata: String) = new OffsetAndMetadata(OffsetMetadata(offset, metadata))
-
-  def apply(offset: Long) = new OffsetAndMetadata(OffsetMetadata(offset, OffsetMetadata.NoMetadata))
-}
-
-case class OffsetMetadataAndError(offsetMetadata: OffsetMetadata, error: Errors = Errors.NONE) {
-  def offset = offsetMetadata.offset
-
-  def metadata = offsetMetadata.metadata
-
-  override def toString = "[%s, Error=%s]".format(offsetMetadata, error)
-}
-
-object OffsetMetadataAndError {
-  val NoOffset = OffsetMetadataAndError(OffsetMetadata.InvalidOffsetMetadata, Errors.NONE)
-  val GroupLoading = OffsetMetadataAndError(OffsetMetadata.InvalidOffsetMetadata, Errors.COORDINATOR_LOAD_IN_PROGRESS)
-  val UnknownMember = OffsetMetadataAndError(OffsetMetadata.InvalidOffsetMetadata, Errors.UNKNOWN_MEMBER_ID)
-  val NotCoordinatorForGroup = OffsetMetadataAndError(OffsetMetadata.InvalidOffsetMetadata, Errors.NOT_COORDINATOR)
-  val GroupCoordinatorNotAvailable = OffsetMetadataAndError(OffsetMetadata.InvalidOffsetMetadata, Errors.COORDINATOR_NOT_AVAILABLE)
-  val UnknownTopicOrPartition = OffsetMetadataAndError(OffsetMetadata.InvalidOffsetMetadata, Errors.UNKNOWN_TOPIC_OR_PARTITION)
-  val IllegalGroupGenerationId = OffsetMetadataAndError(OffsetMetadata.InvalidOffsetMetadata, Errors.ILLEGAL_GENERATION)
-
-  def apply(offset: Long) = new OffsetMetadataAndError(OffsetMetadata(offset, OffsetMetadata.NoMetadata), Errors.NONE)
-
-  def apply(error: Errors) = new OffsetMetadataAndError(OffsetMetadata.InvalidOffsetMetadata, error)
-
-  def apply(offset: Long, metadata: String, error: Errors) = new OffsetMetadataAndError(OffsetMetadata(offset, metadata), error)
-}
-
-
-
diff --git a/core/src/main/scala/kafka/common/OffsetMetadataTooLargeException.scala b/core/src/main/scala/kafka/common/OffsetMetadataTooLargeException.scala
deleted file mode 100644
index 50edb273b3a79..0000000000000
--- a/core/src/main/scala/kafka/common/OffsetMetadataTooLargeException.scala
+++ /dev/null
@@ -1,27 +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.common
-
-/**
- * Indicates the client has specified offset metadata that exceeds the configured
- * maximum size in bytes
- */
-class OffsetMetadataTooLargeException(message: String) extends RuntimeException(message) {
-  def this() = this(null)
-}
-
diff --git a/core/src/main/scala/kafka/common/OffsetOutOfRangeException.scala b/core/src/main/scala/kafka/common/OffsetOutOfRangeException.scala
deleted file mode 100644
index 0a2514cc0d99f..0000000000000
--- a/core/src/main/scala/kafka/common/OffsetOutOfRangeException.scala
+++ /dev/null
@@ -1,26 +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.common
-
-/**
- * Indicates the client has requested a range no longer available on the server
- */
-class OffsetOutOfRangeException(message: String) extends RuntimeException(message) {
-  def this() = this(null)
-}
-
diff --git a/core/src/main/scala/kafka/common/OffsetsLoadInProgressException.scala b/core/src/main/scala/kafka/common/OffsetsLoadInProgressException.scala
deleted file mode 100644
index 1c8e96eefc7f0..0000000000000
--- a/core/src/main/scala/kafka/common/OffsetsLoadInProgressException.scala
+++ /dev/null
@@ -1,26 +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.common
-
-/**
- * Indicates that offsets are currently being loaded from disk into the cache so offset fetch requests cannot be satisfied.
- */
-class OffsetsLoadInProgressException(message: String) extends RuntimeException(message) {
-  def this() = this(null)
-}
-
diff --git a/core/src/main/scala/kafka/common/OffsetsOutOfOrderException.scala b/core/src/main/scala/kafka/common/OffsetsOutOfOrderException.scala
new file mode 100644
index 0000000000000..f8daaa4a181b2
--- /dev/null
+++ b/core/src/main/scala/kafka/common/OffsetsOutOfOrderException.scala
@@ -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 kafka.common
+
+/**
+ * Indicates the follower received records with non-monotonically increasing offsets
+ */
+class OffsetsOutOfOrderException(message: String) extends RuntimeException(message) {
+}
+
diff --git a/core/src/main/scala/kafka/common/QueueFullException.scala b/core/src/main/scala/kafka/common/QueueFullException.scala
deleted file mode 100644
index 27c04821ef54d..0000000000000
--- a/core/src/main/scala/kafka/common/QueueFullException.scala
+++ /dev/null
@@ -1,23 +0,0 @@
-package kafka.common
-
-/**
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-/* Indicates the queue for sending messages is full of unsent messages */
-class QueueFullException(message: String) extends RuntimeException(message) {
-  def this() = this(null)
-}
diff --git a/core/src/main/scala/kafka/common/RecordValidationException.scala b/core/src/main/scala/kafka/common/RecordValidationException.scala
new file mode 100644
index 0000000000000..baa7d72557604
--- /dev/null
+++ b/core/src/main/scala/kafka/common/RecordValidationException.scala
@@ -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 kafka.common
+
+import org.apache.kafka.common.errors.ApiException
+import org.apache.kafka.common.requests.ProduceResponse.RecordError
+
+import scala.collection.Seq
+
+class RecordValidationException(val invalidException: ApiException,
+                                val recordErrors: Seq[RecordError])
+  extends RuntimeException(invalidException) {
+}
diff --git a/core/src/main/scala/kafka/common/ReplicaNotAvailableException.scala b/core/src/main/scala/kafka/common/ReplicaNotAvailableException.scala
deleted file mode 100644
index f1d1eadc546ec..0000000000000
--- a/core/src/main/scala/kafka/common/ReplicaNotAvailableException.scala
+++ /dev/null
@@ -1,26 +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.common
-
-/**
- * Thrown when a request is made for partition, but no leader exists for that partition
- */
-class ReplicaNotAvailableException(cause: Throwable, message: String = "") extends RuntimeException(cause) {
-  def this() = this(null, "")
-  def this(message: String) = this(null, message)
-}
\ No newline at end of file
diff --git a/core/src/main/scala/kafka/common/RequestTimedOutException.scala b/core/src/main/scala/kafka/common/RequestTimedOutException.scala
deleted file mode 100644
index faedea8b40b5c..0000000000000
--- a/core/src/main/scala/kafka/common/RequestTimedOutException.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.common
-
-
-/**
- * Thrown when a produce request times out - i.e., if one or more partitions it
- * sends messages to receives fewer than the requiredAcks that is specified in
- * the produce request.
- */
-class RequestTimedOutException(message: String) extends RuntimeException(message) {
-  def this() = this(null)
-}
\ No newline at end of file
diff --git a/core/src/main/scala/kafka/common/StreamEndException.scala b/core/src/main/scala/kafka/common/StreamEndException.scala
deleted file mode 100644
index a9410bcb907b3..0000000000000
--- a/core/src/main/scala/kafka/common/StreamEndException.scala
+++ /dev/null
@@ -1,23 +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.common
-
-/**
- * An exception that indicates KafkaStream has ended.
- */
-class StreamEndException() extends RuntimeException {
-}
diff --git a/core/src/main/scala/kafka/common/TopicAndPartition.scala b/core/src/main/scala/kafka/common/TopicAndPartition.scala
deleted file mode 100644
index 6c276952f5aa7..0000000000000
--- a/core/src/main/scala/kafka/common/TopicAndPartition.scala
+++ /dev/null
@@ -1,30 +0,0 @@
-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
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-/**
- * Convenience case class since (topic, partition) pairs are ubiquitous.
- */
-case class TopicAndPartition(topic: String, partition: Int) {
-
-  def this(topicPartition: TopicPartition) = this(topicPartition.topic, topicPartition.partition)
-
-  override def toString: String = s"$topic-$partition"
-}
diff --git a/core/src/main/scala/kafka/common/UnavailableProducerException.scala b/core/src/main/scala/kafka/common/UnavailableProducerException.scala
deleted file mode 100644
index 885c98df63f3b..0000000000000
--- a/core/src/main/scala/kafka/common/UnavailableProducerException.scala
+++ /dev/null
@@ -1,24 +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.common
-
-/**
- * Indicates a producer pool initialization problem
-*/
-class UnavailableProducerException(message: String) extends RuntimeException(message) {
-  def this() = this(null)
-}
diff --git a/core/src/main/scala/kafka/common/UnexpectedAppendOffsetException.scala b/core/src/main/scala/kafka/common/UnexpectedAppendOffsetException.scala
new file mode 100644
index 0000000000000..e719a93006d31
--- /dev/null
+++ b/core/src/main/scala/kafka/common/UnexpectedAppendOffsetException.scala
@@ -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 kafka.common
+
+/**
+ * Indicates the follower or the future replica received records from the leader (or current
+ * replica) with first offset less than expected next offset. 
+ * @param firstOffset The first offset of the records to append
+ * @param lastOffset  The last offset of the records to append
+ */
+class UnexpectedAppendOffsetException(val message: String,
+                                      val firstOffset: Long,
+                                      val lastOffset: Long) extends RuntimeException(message) {
+}
diff --git a/core/src/main/scala/kafka/common/UnknownException.scala b/core/src/main/scala/kafka/common/UnknownException.scala
deleted file mode 100644
index 6cf0fc9db316b..0000000000000
--- a/core/src/main/scala/kafka/common/UnknownException.scala
+++ /dev/null
@@ -1,23 +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.common
-
-/**
- * If we don't know what else it is, call it this
- */
-class UnknownException extends RuntimeException
diff --git a/core/src/main/scala/kafka/common/UnknownMagicByteException.scala b/core/src/main/scala/kafka/common/UnknownMagicByteException.scala
deleted file mode 100644
index 544d42687ba8e..0000000000000
--- a/core/src/main/scala/kafka/common/UnknownMagicByteException.scala
+++ /dev/null
@@ -1,26 +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.common
-
-/**
- * Indicates the client has requested a range no longer available on the server
- */
-class UnknownMagicByteException(message: String) extends RuntimeException(message) {
-  def this() = this(null)
-}
-
diff --git a/core/src/main/scala/kafka/common/UnknownTopicOrPartitionException.scala b/core/src/main/scala/kafka/common/UnknownTopicOrPartitionException.scala
deleted file mode 100644
index f382d16de9c39..0000000000000
--- a/core/src/main/scala/kafka/common/UnknownTopicOrPartitionException.scala
+++ /dev/null
@@ -1,26 +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.common
-
-/**
- * Indicates one of the following situation:
- * 1. Producer does not have the partition metadata for this id upon sending messages
- * 2. Broker does not have the specified partition by id upon receiving messages
- */
-class UnknownTopicOrPartitionException(message: String) extends RuntimeException(message) {
-  def this() = this(null)
-}
diff --git a/core/src/main/scala/kafka/common/ZkNodeChangeNotificationListener.scala b/core/src/main/scala/kafka/common/ZkNodeChangeNotificationListener.scala
index 4cae80c5579e9..c49ec085cc7f0 100644
--- a/core/src/main/scala/kafka/common/ZkNodeChangeNotificationListener.scala
+++ b/core/src/main/scala/kafka/common/ZkNodeChangeNotificationListener.scala
@@ -16,6 +16,7 @@
  */
 package kafka.common
 
+import java.nio.charset.StandardCharsets.UTF_8
 import java.util.concurrent.LinkedBlockingQueue
 import java.util.concurrent.atomic.AtomicBoolean
 
@@ -24,11 +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: String)
+  def processNotification(notificationMessage: Array[Byte]): Unit
 }
 
 /**
@@ -56,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()
@@ -74,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) {
@@ -83,12 +87,7 @@ class ZkNodeChangeNotificationListener(private val zkClient: KafkaZkClient,
         for (notification <- notifications) {
           val changeId = changeNumber(notification)
           if (changeId > lastExecutedChange) {
-            val changeZnode = seqNodeRoot + "/" + notification
-            val (data, _) = zkClient.getDataAndStat(changeZnode)
-            data match {
-              case Some(d) => notificationHandler.processNotification(d)
-              case None => warn(s"read null data from $changeZnode when processing notification $notification")
-            }
+            processNotification(notification)
             lastExecutedChange = changeId
           }
         }
@@ -100,13 +99,25 @@ class ZkNodeChangeNotificationListener(private val zkClient: KafkaZkClient,
     }
   }
 
+  private def processNotification(notification: String): Unit = {
+    val changeZnode = seqNodeRoot + "/" + notification
+    val (data, _) = zkClient.getDataAndStat(changeZnode)
+    data match {
+      case Some(d) => Try(notificationHandler.processNotification(d)) match {
+        case Failure(e) => error(s"error processing change notification ${new String(d, UTF_8)} from $changeZnode", e)
+        case _ =>
+      }
+      case None => warn(s"read null data from $changeZnode")
+    }
+  }
+
   private def addChangeNotification(): Unit = {
     if (!isClosed.get && queue.peek() == null)
       queue.put(new ChangeNotification)
   }
 
   class ChangeNotification {
-    def process(): Unit = processNotifications
+    def process(): Unit = processNotifications()
   }
 
   /**
@@ -115,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)
@@ -132,18 +143,17 @@ class ZkNodeChangeNotificationListener(private val zkClient: KafkaZkClient,
   private def changeNumber(name: String): Long = name.substring(seqNodePrefix.length).toLong
 
   class ChangeEventProcessThread(name: String) extends ShutdownableThread(name = name) {
-    override def doWork(): Unit = queue.take().process
+    override def doWork(): Unit = queue.take().process()
   }
 
   object ChangeNotificationHandler extends ZNodeChildChangeHandler {
     override val path: String = seqNodeRoot
-    override def handleChildChange(): Unit = addChangeNotification
+    override def handleChildChange(): Unit = addChangeNotification()
   }
 
   object ZkStateChangeHandler extends  StateChangeHandler {
     override val name: String = StateChangeHandlers.zkNodeChangeListenerHandler(seqNodeRoot)
-    override def afterInitializingSession(): Unit = addChangeNotification
-    override def onReconnectionTimeout(): Unit = error("Reconnection timeout.")
+    override def afterInitializingSession(): Unit = addChangeNotification()
   }
 }
 
diff --git a/core/src/main/scala/kafka/consumer/BaseConsumer.scala b/core/src/main/scala/kafka/consumer/BaseConsumer.scala
deleted file mode 100644
index 04ac2d9f1a570..0000000000000
--- a/core/src/main/scala/kafka/consumer/BaseConsumer.scala
+++ /dev/null
@@ -1,161 +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.consumer
-
-import java.util.{Collections, Properties}
-import java.util.regex.Pattern
-
-import kafka.api.OffsetRequest
-import kafka.common.StreamEndException
-import kafka.message.Message
-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
-
-/**
- * A base consumer used to abstract both old and new consumer
- * this class should be removed (along with BaseProducer)
- * once we deprecate old consumer
- */
-@deprecated("This trait has been deprecated and will be removed in a future release. " +
-            "Please use org.apache.kafka.clients.consumer.KafkaConsumer instead.", "0.11.0.0")
-trait BaseConsumer {
-  def receive(): BaseConsumerRecord
-  def stop()
-  def cleanup()
-  def commit()
-}
-
-@deprecated("This class has been deprecated and will be removed in a future release. " +
-            "Please use org.apache.kafka.clients.consumer.ConsumerRecord instead.", "0.11.0.0")
-case class BaseConsumerRecord(topic: String,
-                              partition: Int,
-                              offset: Long,
-                              timestamp: Long = Message.NoTimestamp,
-                              timestampType: TimestampType = TimestampType.NO_TIMESTAMP_TYPE,
-                              key: Array[Byte],
-                              value: Array[Byte],
-                              headers: Headers = new RecordHeaders())
-
-@deprecated("This class has been deprecated and will be removed in a future release. " +
-            "Please use org.apache.kafka.clients.consumer.KafkaConsumer instead.", "0.11.0.0")
-class NewShinyConsumer(topic: Option[String], partitionId: Option[Int], offset: Option[Long], whitelist: Option[String], consumerProps: Properties, val timeoutMs: Long = Long.MaxValue) extends BaseConsumer {
-  import org.apache.kafka.clients.consumer.KafkaConsumer
-
-  val consumer = new KafkaConsumer[Array[Byte], Array[Byte]](consumerProps)
-  consumerInit()
-  var recordIter = consumer.poll(0).iterator
-
-  def consumerInit() {
-    (topic, partitionId, offset, whitelist) match {
-      case (Some(topic), Some(partitionId), Some(offset), None) =>
-        seek(topic, partitionId, offset)
-      case (Some(topic), Some(partitionId), None, None) =>
-        // default to latest if no offset is provided
-        seek(topic, partitionId, OffsetRequest.LatestTime)
-      case (Some(topic), None, None, None) =>
-        consumer.subscribe(Collections.singletonList(topic))
-      case (None, None, None, Some(whitelist)) =>
-        consumer.subscribe(Pattern.compile(whitelist))
-      case _ =>
-        throw new IllegalArgumentException("An invalid combination of arguments is provided. " +
-            "Exactly one of 'topic' or 'whitelist' must be provided. " +
-            "If 'topic' is provided, an optional 'partition' may also be provided. " +
-            "If 'partition' is provided, an optional 'offset' may also be provided, otherwise, consumption starts from the end of the partition.")
-    }
-  }
-
-  def seek(topic: String, partitionId: Int, offset: Long) {
-    val topicPartition = new TopicPartition(topic, partitionId)
-    consumer.assign(Collections.singletonList(topicPartition))
-    offset match {
-      case OffsetRequest.EarliestTime => consumer.seekToBeginning(Collections.singletonList(topicPartition))
-      case OffsetRequest.LatestTime => consumer.seekToEnd(Collections.singletonList(topicPartition))
-      case _ => consumer.seek(topicPartition, offset)
-    }
-  }
-
-  override def receive(): BaseConsumerRecord = {
-    if (!recordIter.hasNext) {
-      recordIter = consumer.poll(timeoutMs).iterator
-      if (!recordIter.hasNext)
-        throw new ConsumerTimeoutException
-    }
-
-    val record = recordIter.next
-    BaseConsumerRecord(record.topic,
-                       record.partition,
-                       record.offset,
-                       record.timestamp,
-                       record.timestampType,
-                       record.key,
-                       record.value,
-                       record.headers)
-  }
-
-  override def stop() {
-    this.consumer.wakeup()
-  }
-
-  override def cleanup() {
-    this.consumer.close()
-  }
-
-  override def commit() {
-    this.consumer.commitSync()
-  }
-}
-
-@deprecated("This class has been deprecated and will be removed in a future release. " +
-            "Please use org.apache.kafka.clients.consumer.KafkaConsumer instead.", "0.11.0.0")
-class OldConsumer(topicFilter: TopicFilter, consumerProps: Properties) extends BaseConsumer {
-  import kafka.serializer.DefaultDecoder
-
-  val consumerConnector = Consumer.create(new ConsumerConfig(consumerProps))
-  val stream: KafkaStream[Array[Byte], Array[Byte]] =
-    consumerConnector.createMessageStreamsByFilter(topicFilter, 1, new DefaultDecoder(), new DefaultDecoder()).head
-  val iter = stream.iterator
-
-  override def receive(): BaseConsumerRecord = {
-    if (!iter.hasNext())
-      throw new StreamEndException
-
-    val messageAndMetadata = iter.next
-    BaseConsumerRecord(messageAndMetadata.topic,
-                       messageAndMetadata.partition,
-                       messageAndMetadata.offset,
-                       messageAndMetadata.timestamp,
-                       messageAndMetadata.timestampType,
-                       messageAndMetadata.key,
-                       messageAndMetadata.message, 
-                       new RecordHeaders())
-  }
-
-  override def stop() {
-    this.consumerConnector.shutdown()
-  }
-
-  override def cleanup() {
-    this.consumerConnector.shutdown()
-  }
-
-  override def commit() {
-    this.consumerConnector.commitOffsets
-  }
-}
diff --git a/core/src/main/scala/kafka/consumer/BaseConsumerRecord.scala b/core/src/main/scala/kafka/consumer/BaseConsumerRecord.scala
new file mode 100644
index 0000000000000..7628b6b65168c
--- /dev/null
+++ b/core/src/main/scala/kafka/consumer/BaseConsumerRecord.scala
@@ -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 kafka.consumer
+
+import org.apache.kafka.common.header.Headers
+import org.apache.kafka.common.header.internals.RecordHeaders
+import org.apache.kafka.common.record.{RecordBatch, TimestampType}
+
+@deprecated("This class has been deprecated and will be removed in a future release. " +
+  "Please use org.apache.kafka.clients.consumer.ConsumerRecord instead.", "0.11.0.0")
+case class BaseConsumerRecord(topic: String,
+                              partition: Int,
+                              offset: Long,
+                              timestamp: Long = RecordBatch.NO_TIMESTAMP,
+                              timestampType: TimestampType = TimestampType.NO_TIMESTAMP_TYPE,
+                              key: Array[Byte],
+                              value: Array[Byte],
+                              headers: Headers = new RecordHeaders())
diff --git a/core/src/main/scala/kafka/consumer/ConsumerConfig.scala b/core/src/main/scala/kafka/consumer/ConsumerConfig.scala
deleted file mode 100644
index bea0307389a1d..0000000000000
--- a/core/src/main/scala/kafka/consumer/ConsumerConfig.scala
+++ /dev/null
@@ -1,205 +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.consumer
-
-import java.util.Properties
-import kafka.api.OffsetRequest
-import kafka.utils._
-import kafka.common.{InvalidConfigException, Config}
-import java.util.Locale
-
-@deprecated("This object has been deprecated and will be removed in a future release. " +
-            "Please use org.apache.kafka.clients.consumer.ConsumerConfig instead.", "0.11.0.0")
-object ConsumerConfig extends Config {
-  val RefreshMetadataBackoffMs = 200
-  val SocketTimeout = 30 * 1000
-  val SocketBufferSize = 64*1024
-  val FetchSize = 1024 * 1024
-  val MaxFetchSize = 10*FetchSize
-  val NumConsumerFetchers = 1
-  val DefaultFetcherBackoffMs = 1000
-  val AutoCommit = true
-  val AutoCommitInterval = 60 * 1000
-  val MaxQueuedChunks = 2
-  val MaxRebalanceRetries = 4
-  val AutoOffsetReset = OffsetRequest.LargestTimeString
-  val ConsumerTimeoutMs = -1
-  val MinFetchBytes = 1
-  val MaxFetchBytes = 50 * 1024 * 1024
-  val MaxFetchWaitMs = 100
-  val MirrorTopicsWhitelist = ""
-  val MirrorTopicsBlacklist = ""
-  val MirrorConsumerNumThreads = 1
-  val OffsetsChannelBackoffMs = 1000
-  val OffsetsChannelSocketTimeoutMs = 10000
-  val OffsetsCommitMaxRetries = 5
-  val OffsetsStorage = "zookeeper"
-
-  val MirrorTopicsWhitelistProp = "mirror.topics.whitelist"
-  val MirrorTopicsBlacklistProp = "mirror.topics.blacklist"
-  val ExcludeInternalTopics = true
-  val DefaultPartitionAssignmentStrategy = "range" /* select between "range", and "roundrobin" */
-  val MirrorConsumerNumThreadsProp = "mirror.consumer.numthreads"
-  val DefaultClientId = ""
-
-  def validate(config: ConsumerConfig) {
-    validateClientId(config.clientId)
-    validateGroupId(config.groupId)
-    validateAutoOffsetReset(config.autoOffsetReset)
-    validateOffsetsStorage(config.offsetsStorage)
-    validatePartitionAssignmentStrategy(config.partitionAssignmentStrategy)
-  }
-
-  def validateClientId(clientId: String) {
-    validateChars("client.id", clientId)
-  }
-
-  def validateGroupId(groupId: String) {
-    validateChars("group.id", groupId)
-  }
-
-  def validateAutoOffsetReset(autoOffsetReset: String) {
-    autoOffsetReset match {
-      case OffsetRequest.SmallestTimeString =>
-      case OffsetRequest.LargestTimeString =>
-      case _ => throw new InvalidConfigException("Wrong value " + autoOffsetReset + " of auto.offset.reset in ConsumerConfig; " +
-                                                 "Valid values are " + OffsetRequest.SmallestTimeString + " and " + OffsetRequest.LargestTimeString)
-    }
-  }
-
-  def validateOffsetsStorage(storage: String) {
-    storage match {
-      case "zookeeper" =>
-      case "kafka" =>
-      case _ => throw new InvalidConfigException("Wrong value " + storage + " of offsets.storage in consumer config; " +
-                                                 "Valid values are 'zookeeper' and 'kafka'")
-    }
-  }
-
-  def validatePartitionAssignmentStrategy(strategy: String) {
-    strategy match {
-      case "range" =>
-      case "roundrobin" =>
-      case _ => throw new InvalidConfigException("Wrong value " + strategy + " of partition.assignment.strategy in consumer config; " +
-        "Valid values are 'range' and 'roundrobin'")
-    }
-  }
-}
-
-@deprecated("This class has been deprecated and will be removed in a future release. " +
-            "Please use org.apache.kafka.clients.consumer.ConsumerConfig instead.", "0.11.0.0")
-class ConsumerConfig private (val props: VerifiableProperties) extends ZKConfig(props) {
-  import ConsumerConfig._
-
-  def this(originalProps: Properties) {
-    this(new VerifiableProperties(originalProps))
-    props.verify()
-  }
-
-  /** a string that uniquely identifies a set of consumers within the same consumer group */
-  val groupId = props.getString("group.id")
-
-  /** consumer id: generated automatically if not set.
-   *  Set this explicitly for only testing purpose. */
-  val consumerId: Option[String] = Option(props.getString("consumer.id", null))
-
-  /** the socket timeout for network requests. Its value should be at least fetch.wait.max.ms. */
-  val socketTimeoutMs = props.getInt("socket.timeout.ms", SocketTimeout)
-
-  /** the socket receive buffer for network requests */
-  val socketReceiveBufferBytes = props.getInt("socket.receive.buffer.bytes", SocketBufferSize)
-
-  /** the number of bytes of messages to attempt to fetch from each partition */
-  val fetchMessageMaxBytes = props.getInt("fetch.message.max.bytes", FetchSize)
-
-  /** the number threads used to fetch data */
-  val numConsumerFetchers = props.getInt("num.consumer.fetchers", NumConsumerFetchers)
-
-  /** if true, periodically commit to zookeeper the offset of messages already fetched by the consumer */
-  val autoCommitEnable = props.getBoolean("auto.commit.enable", AutoCommit)
-
-  /** the frequency in ms that the consumer offsets are committed to zookeeper */
-  val autoCommitIntervalMs = props.getInt("auto.commit.interval.ms", AutoCommitInterval)
-
-  /** max number of message chunks buffered for consumption, each chunk can be up to fetch.message.max.bytes*/
-  val queuedMaxMessages = props.getInt("queued.max.message.chunks", MaxQueuedChunks)
-
-  /** max number of retries during rebalance */
-  val rebalanceMaxRetries = props.getInt("rebalance.max.retries", MaxRebalanceRetries)
-
-  /** the minimum amount of data the server should return for a fetch request. If insufficient data is available the request will block */
-  val fetchMinBytes = props.getInt("fetch.min.bytes", MinFetchBytes)
-
-  /** the maximum amount of data the server should return for a fetch request */
-  val fetchMaxBytes = props.getInt("fetch.max.bytes", MaxFetchBytes)
-
-  /** the maximum amount of time the server will block before answering the fetch request if there isn't sufficient data to immediately satisfy fetch.min.bytes */
-  val fetchWaitMaxMs = props.getInt("fetch.wait.max.ms", MaxFetchWaitMs)
-  require(fetchWaitMaxMs <= socketTimeoutMs, "socket.timeout.ms should always be at least fetch.wait.max.ms" +
-    " to prevent unnecessary socket timeouts")
-
-  /** backoff time between retries during rebalance */
-  val rebalanceBackoffMs = props.getInt("rebalance.backoff.ms", zkSyncTimeMs)
-
-  /** backoff time to refresh the leader of a partition after it loses the current leader */
-  val refreshLeaderBackoffMs = props.getInt("refresh.leader.backoff.ms", RefreshMetadataBackoffMs)
-
-  /** backoff time to reconnect the offsets channel or to retry offset fetches/commits */
-  val offsetsChannelBackoffMs = props.getInt("offsets.channel.backoff.ms", OffsetsChannelBackoffMs)
-  /** socket timeout to use when reading responses for Offset Fetch/Commit requests. This timeout will also be used for
-   *  the ConsumerMetdata requests that are used to query for the offset coordinator. */
-  val offsetsChannelSocketTimeoutMs = props.getInt("offsets.channel.socket.timeout.ms", OffsetsChannelSocketTimeoutMs)
-
-  /** Retry the offset commit up to this many times on failure. This retry count only applies to offset commits during
-    * shut-down. It does not apply to commits from the auto-commit thread. It also does not apply to attempts to query
-    * for the offset coordinator before committing offsets. i.e., if a consumer metadata request fails for any reason,
-    * it is retried and that retry does not count toward this limit. */
-  val offsetsCommitMaxRetries = props.getInt("offsets.commit.max.retries", OffsetsCommitMaxRetries)
-
-  /** Specify whether offsets should be committed to "zookeeper" (default) or "kafka" */
-  val offsetsStorage = props.getString("offsets.storage", OffsetsStorage).toLowerCase(Locale.ROOT)
-
-  /** If you are using "kafka" as offsets.storage, you can dual commit offsets to ZooKeeper (in addition to Kafka). This
-    * is required during migration from zookeeper-based offset storage to kafka-based offset storage. With respect to any
-    * given consumer group, it is safe to turn this off after all instances within that group have been migrated to
-    * the new jar that commits offsets to the broker (instead of directly to ZooKeeper). */
-  val dualCommitEnabled = props.getBoolean("dual.commit.enabled", offsetsStorage == "kafka")
-
-  /* what to do if an offset is out of range.
-     smallest : automatically reset the offset to the smallest offset
-     largest : automatically reset the offset to the largest offset
-     anything else: throw exception to the consumer */
-  val autoOffsetReset = props.getString("auto.offset.reset", AutoOffsetReset)
-
-  /** throw a timeout exception to the consumer if no message is available for consumption after the specified interval */
-  val consumerTimeoutMs = props.getInt("consumer.timeout.ms", ConsumerTimeoutMs)
-
-  /**
-   * Client id is specified by the kafka consumer client, used to distinguish different clients
-   */
-  val clientId = props.getString("client.id", groupId)
-
-  /** Whether messages from internal topics (such as offsets) should be exposed to the consumer. */
-  val excludeInternalTopics = props.getBoolean("exclude.internal.topics", ExcludeInternalTopics)
-
-  /** Select a strategy for assigning partitions to consumer streams. Possible values: range, roundrobin */
-  val partitionAssignmentStrategy = props.getString("partition.assignment.strategy", DefaultPartitionAssignmentStrategy)
-
-  validate(this)
-}
-
diff --git a/core/src/main/scala/kafka/consumer/ConsumerConnector.scala b/core/src/main/scala/kafka/consumer/ConsumerConnector.scala
deleted file mode 100644
index f6d4a74a4d7bf..0000000000000
--- a/core/src/main/scala/kafka/consumer/ConsumerConnector.scala
+++ /dev/null
@@ -1,126 +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.consumer
-
-import kafka.common.{OffsetAndMetadata, TopicAndPartition}
-import kafka.javaapi.consumer.ConsumerRebalanceListener
-
-import scala.collection._
-import kafka.utils.Logging
-import kafka.serializer._
-
-/**
- *  Main interface for consumer
- */
-@deprecated("This trait has been deprecated and will be removed in a future release.", "0.11.0.0")
-trait ConsumerConnector {
-
-  /**
-   *  Create a list of MessageStreams for each topic.
-   *
-   *  @param topicCountMap  a map of (topic, #streams) pair
-   *  @return a map of (topic, list of  KafkaStream) pairs.
-   *          The number of items in the list is #streams. Each stream supports
-   *          an iterator over message/metadata pairs.
-   */
-  def createMessageStreams(topicCountMap: Map[String,Int]): Map[String, List[KafkaStream[Array[Byte],Array[Byte]]]]
-
-  /**
-   *  Create a list of MessageStreams for each topic.
-   *
-   *  @param topicCountMap  a map of (topic, #streams) pair
-   *  @param keyDecoder Decoder to decode the key portion of the message
-   *  @param valueDecoder Decoder to decode the value portion of the message
-   *  @return a map of (topic, list of  KafkaStream) pairs.
-   *          The number of items in the list is #streams. Each stream supports
-   *          an iterator over message/metadata pairs.
-   */
-  def createMessageStreams[K,V](topicCountMap: Map[String,Int],
-                                keyDecoder: Decoder[K],
-                                valueDecoder: Decoder[V])
-    : Map[String,List[KafkaStream[K,V]]]
-
-  /**
-   *  Create a list of message streams for all topics that match a given filter.
-   *
-   *  @param topicFilter Either a Whitelist or Blacklist TopicFilter object.
-   *  @param numStreams Number of streams to return
-   *  @param keyDecoder Decoder to decode the key portion of the message
-   *  @param valueDecoder Decoder to decode the value portion of the message
-   *  @return a list of KafkaStream each of which provides an
-   *          iterator over message/metadata pairs over allowed topics.
-   */
-  def createMessageStreamsByFilter[K,V](topicFilter: TopicFilter,
-                                        numStreams: Int = 1,
-                                        keyDecoder: Decoder[K] = new DefaultDecoder(),
-                                        valueDecoder: Decoder[V] = new DefaultDecoder())
-    : Seq[KafkaStream[K,V]]
-
-  /**
-   *  Commit the offsets of all broker partitions connected by this connector.
-   */
-  def commitOffsets(retryOnFailure: Boolean)
-
-  /**
-   * KAFKA-1743: This method added for backward compatibility.
-   */
-  def commitOffsets()
-
-  /**
-   * Commit offsets from an external offsets map.
-   * @param offsetsToCommit the offsets to be committed.
-   */
-  def commitOffsets(offsetsToCommit: immutable.Map[TopicAndPartition, OffsetAndMetadata], retryOnFailure: Boolean)
-
-  /**
-   * Wire in a consumer rebalance listener to be executed when consumer rebalance occurs.
-   * @param listener The consumer rebalance listener to wire in
-   */
-  def setConsumerRebalanceListener(listener: ConsumerRebalanceListener)
-
-  /**
-   *  Shut down the connector
-   */
-  def shutdown()
-}
-
-@deprecated("This object has been deprecated and will be removed in a future release. " +
-            "Please use org.apache.kafka.clients.consumer.Consumer instead.", "0.11.0.0")
-object Consumer extends Logging {
-  /**
-   *  Create a ConsumerConnector
-   *
-   *  @param config  at the minimum, need to specify the groupid of the consumer and the zookeeper
-   *                 connection string zookeeper.connect.
-   */
-  def create(config: ConsumerConfig): ConsumerConnector = {
-    val consumerConnect = new ZookeeperConsumerConnector(config)
-    consumerConnect
-  }
-
-  /**
-   *  Create a ConsumerConnector
-   *
-   *  @param config  at the minimum, need to specify the groupid of the consumer and the zookeeper
-   *                 connection string zookeeper.connect.
-   */
-  def createJavaConsumerConnector(config: ConsumerConfig): kafka.javaapi.consumer.ConsumerConnector = {
-    val consumerConnect = new kafka.javaapi.consumer.ZookeeperConsumerConnector(config)
-    consumerConnect
-  }
-}
diff --git a/core/src/main/scala/kafka/consumer/ConsumerFetcherManager.scala b/core/src/main/scala/kafka/consumer/ConsumerFetcherManager.scala
deleted file mode 100755
index 0a6b82e5fc104..0000000000000
--- a/core/src/main/scala/kafka/consumer/ConsumerFetcherManager.scala
+++ /dev/null
@@ -1,164 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package kafka.consumer
-
-import kafka.server.{AbstractFetcherManager, AbstractFetcherThread, BrokerAndInitialOffset}
-import kafka.cluster.{BrokerEndPoint, Cluster}
-import org.apache.kafka.common.TopicPartition
-import org.apache.kafka.common.utils.Time
-
-import scala.collection.immutable
-import collection.mutable.HashMap
-import scala.collection.mutable
-import java.util.concurrent.locks.ReentrantLock
-
-import kafka.utils.CoreUtils.inLock
-import kafka.utils.ZkUtils
-import kafka.utils.ShutdownableThread
-import kafka.client.ClientUtils
-import java.util.concurrent.atomic.AtomicInteger
-
-/**
- *  Usage:
- *  Once ConsumerFetcherManager is created, startConnections() and stopAllConnections() can be called repeatedly
- *  until shutdown() is called.
- */
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.11.0.0")
-class ConsumerFetcherManager(private val consumerIdString: String,
-                             private val config: ConsumerConfig,
-                             private val zkUtils : ZkUtils)
-        extends AbstractFetcherManager("ConsumerFetcherManager-%d".format(Time.SYSTEM.milliseconds),
-                                       config.clientId, config.numConsumerFetchers) {
-  private var partitionMap: immutable.Map[TopicPartition, PartitionTopicInfo] = null
-  private val noLeaderPartitionSet = new mutable.HashSet[TopicPartition]
-  private val lock = new ReentrantLock
-  private val cond = lock.newCondition()
-  private var leaderFinderThread: ShutdownableThread = null
-  private val correlationId = new AtomicInteger(0)
-
-  private class LeaderFinderThread(name: String) extends ShutdownableThread(name) {
-    // thread responsible for adding the fetcher to the right broker when leader is available
-    override def doWork() {
-      val leaderForPartitionsMap = new HashMap[TopicPartition, BrokerEndPoint]
-      lock.lock()
-      try {
-        while (noLeaderPartitionSet.isEmpty) {
-          trace("No partition for leader election.")
-          cond.await()
-        }
-
-        trace("Partitions without leader %s".format(noLeaderPartitionSet))
-        val brokers = ClientUtils.getPlaintextBrokerEndPoints(zkUtils)
-        val topicsMetadata = ClientUtils.fetchTopicMetadata(noLeaderPartitionSet.map(m => m.topic).toSet,
-                                                            brokers,
-                                                            config.clientId,
-                                                            config.socketTimeoutMs,
-                                                            correlationId.getAndIncrement).topicsMetadata
-        if(isDebugEnabled) topicsMetadata.foreach(topicMetadata => debug(topicMetadata.toString()))
-        topicsMetadata.foreach { tmd =>
-          val topic = tmd.topic
-          tmd.partitionsMetadata.foreach { pmd =>
-            val topicAndPartition = new TopicPartition(topic, pmd.partitionId)
-            if(pmd.leader.isDefined && noLeaderPartitionSet.contains(topicAndPartition)) {
-              val leaderBroker = pmd.leader.get
-              leaderForPartitionsMap.put(topicAndPartition, leaderBroker)
-              noLeaderPartitionSet -= topicAndPartition
-            }
-          }
-        }
-      } catch {
-        case t: Throwable => {
-            if (!isRunning.get())
-              throw t /* If this thread is stopped, propagate this exception to kill the thread. */
-            else
-              warn("Failed to find leader for %s".format(noLeaderPartitionSet), t)
-          }
-      } finally {
-        lock.unlock()
-      }
-
-      try {
-        addFetcherForPartitions(leaderForPartitionsMap.map { case (topicPartition, broker) =>
-          topicPartition -> BrokerAndInitialOffset(broker, partitionMap(topicPartition).getFetchOffset())}
-        )
-      } catch {
-        case t: Throwable =>
-          if (!isRunning.get())
-            throw t /* If this thread is stopped, propagate this exception to kill the thread. */
-          else {
-            warn("Failed to add leader for partitions %s; will retry".format(leaderForPartitionsMap.keySet.mkString(",")), t)
-            lock.lock()
-            noLeaderPartitionSet ++= leaderForPartitionsMap.keySet
-            lock.unlock()
-          }
-        }
-
-      shutdownIdleFetcherThreads()
-      Thread.sleep(config.refreshLeaderBackoffMs)
-    }
-  }
-
-  override def createFetcherThread(fetcherId: Int, sourceBroker: BrokerEndPoint): AbstractFetcherThread = {
-    new ConsumerFetcherThread(
-      "ConsumerFetcherThread-%s-%d-%d".format(consumerIdString, fetcherId, sourceBroker.id),
-      config, sourceBroker, partitionMap, this)
-  }
-
-  def startConnections(topicInfos: Iterable[PartitionTopicInfo], cluster: Cluster) {
-    leaderFinderThread = new LeaderFinderThread(consumerIdString + "-leader-finder-thread")
-    leaderFinderThread.start()
-
-    inLock(lock) {
-      partitionMap = topicInfos.map(tpi => (new TopicPartition(tpi.topic, tpi.partitionId), tpi)).toMap
-      noLeaderPartitionSet ++= topicInfos.map(tpi => new TopicPartition(tpi.topic, tpi.partitionId))
-      cond.signalAll()
-    }
-  }
-
-  def stopConnections() {
-    /*
-     * Stop the leader finder thread first before stopping fetchers. Otherwise, if there are more partitions without
-     * leader, then the leader finder thread will process these partitions (before shutting down) and add fetchers for
-     * these partitions.
-     */
-    info("Stopping leader finder thread")
-    if (leaderFinderThread != null) {
-      leaderFinderThread.shutdown()
-      leaderFinderThread = null
-    }
-
-    info("Stopping all fetchers")
-    closeAllFetchers()
-
-    // no need to hold the lock for the following since leaderFindThread and all fetchers have been stopped
-    partitionMap = null
-    noLeaderPartitionSet.clear()
-
-    info("All connections stopped")
-  }
-
-  def addPartitionsWithError(partitionList: Iterable[TopicPartition]) {
-    debug("adding partitions with error %s".format(partitionList))
-    inLock(lock) {
-      if (partitionMap != null) {
-        noLeaderPartitionSet ++= partitionList
-        cond.signalAll()
-      }
-    }
-  }
-}
diff --git a/core/src/main/scala/kafka/consumer/ConsumerFetcherThread.scala b/core/src/main/scala/kafka/consumer/ConsumerFetcherThread.scala
deleted file mode 100644
index 705dc249bf302..0000000000000
--- a/core/src/main/scala/kafka/consumer/ConsumerFetcherThread.scala
+++ /dev/null
@@ -1,154 +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.consumer
-
-import kafka.api.{FetchRequestBuilder, FetchResponsePartitionData, OffsetRequest, Request}
-import kafka.cluster.BrokerEndPoint
-import kafka.message.ByteBufferMessageSet
-import kafka.server.{AbstractFetcherThread, PartitionFetchState}
-import AbstractFetcherThread.ResultWithPartitions
-import kafka.common.{ErrorMapping, TopicAndPartition}
-
-import scala.collection.Map
-import ConsumerFetcherThread._
-import org.apache.kafka.common.TopicPartition
-import org.apache.kafka.common.protocol.Errors
-import org.apache.kafka.common.record.MemoryRecords
-import org.apache.kafka.common.requests.EpochEndOffset
-
-
-@deprecated("This class has been deprecated and will be removed in a future release. " +
-            "Please use org.apache.kafka.clients.consumer.internals.Fetcher instead.", "0.11.0.0")
-class ConsumerFetcherThread(name: String,
-                            val config: ConsumerConfig,
-                            sourceBroker: BrokerEndPoint,
-                            partitionMap: Map[TopicPartition, PartitionTopicInfo],
-                            val consumerFetcherManager: ConsumerFetcherManager)
-        extends AbstractFetcherThread(name = name,
-                                      clientId = config.clientId,
-                                      sourceBroker = sourceBroker,
-                                      fetchBackOffMs = config.refreshLeaderBackoffMs,
-                                      isInterruptible = true,
-                                      includeLogTruncation = false) {
-
-  type REQ = FetchRequest
-  type PD = PartitionData
-
-  private val clientId = config.clientId
-  private val fetchSize = config.fetchMessageMaxBytes
-
-  private val simpleConsumer = new SimpleConsumer(sourceBroker.host, sourceBroker.port, config.socketTimeoutMs,
-    config.socketReceiveBufferBytes, config.clientId)
-
-  private val fetchRequestBuilder = new FetchRequestBuilder().
-    clientId(clientId).
-    replicaId(Request.OrdinaryConsumerId).
-    maxWait(config.fetchWaitMaxMs).
-    minBytes(config.fetchMinBytes).
-    requestVersion(3) // for now, the old consumer is pinned to the old message format through the fetch request
-
-  override def initiateShutdown(): Boolean = {
-    val justShutdown = super.initiateShutdown()
-    if (justShutdown && isInterruptible)
-      simpleConsumer.disconnectToHandleJavaIOBug()
-    justShutdown
-  }
-
-  override def shutdown(): Unit = {
-    super.shutdown()
-    simpleConsumer.close()
-  }
-
-  // process fetched data
-  def processPartitionData(topicPartition: TopicPartition, fetchOffset: Long, partitionData: PartitionData) {
-    val pti = partitionMap(topicPartition)
-    if (pti.getFetchOffset != fetchOffset)
-      throw new RuntimeException("Offset doesn't match for partition [%s,%d] pti offset: %d fetch offset: %d"
-                                .format(topicPartition.topic, topicPartition.partition, pti.getFetchOffset, fetchOffset))
-    pti.enqueue(partitionData.underlying.messages.asInstanceOf[ByteBufferMessageSet])
-  }
-
-  // handle a partition whose offset is out of range and return a new fetch offset
-  def handleOffsetOutOfRange(topicPartition: TopicPartition): Long = {
-    val startTimestamp = config.autoOffsetReset match {
-      case OffsetRequest.SmallestTimeString => OffsetRequest.EarliestTime
-      case _ => OffsetRequest.LatestTime
-    }
-    val topicAndPartition = TopicAndPartition(topicPartition.topic, topicPartition.partition)
-    val newOffset = simpleConsumer.earliestOrLatestOffset(topicAndPartition, startTimestamp, Request.OrdinaryConsumerId)
-    val pti = partitionMap(topicPartition)
-    pti.resetFetchOffset(newOffset)
-    pti.resetConsumeOffset(newOffset)
-    newOffset
-  }
-
-  // any logic for partitions whose leader has changed
-  def handlePartitionsWithErrors(partitions: Iterable[TopicPartition]) {
-    if (partitions.nonEmpty) {
-      removePartitions(partitions.toSet)
-      consumerFetcherManager.addPartitionsWithError(partitions)
-    }
-  }
-
-  protected def buildFetchRequest(partitionMap: collection.Seq[(TopicPartition, PartitionFetchState)]): ResultWithPartitions[FetchRequest] = {
-    partitionMap.foreach { case ((topicPartition, partitionFetchState)) =>
-      if (partitionFetchState.isReadyForFetch)
-        fetchRequestBuilder.addFetch(topicPartition.topic, topicPartition.partition, partitionFetchState.fetchOffset, fetchSize)
-    }
-
-    ResultWithPartitions(new FetchRequest(fetchRequestBuilder.build()), Set())
-  }
-
-  protected def fetch(fetchRequest: FetchRequest): Seq[(TopicPartition, PartitionData)] =
-    simpleConsumer.fetch(fetchRequest.underlying).data.map { case (TopicAndPartition(t, p), value) =>
-      new TopicPartition(t, p) -> new PartitionData(value)
-    }
-
-  override def buildLeaderEpochRequest(allPartitions: Seq[(TopicPartition, PartitionFetchState)]): ResultWithPartitions[Map[TopicPartition, Int]] = {
-    ResultWithPartitions(Map(), Set())
-  }
-
-  override def fetchEpochsFromLeader(partitions: Map[TopicPartition, Int]): Map[TopicPartition, EpochEndOffset] = { Map() }
-
-  override def maybeTruncate(fetchedEpochs: Map[TopicPartition, EpochEndOffset]): ResultWithPartitions[Map[TopicPartition, Long]] = {
-    ResultWithPartitions(Map(), Set())
-  }
-}
-
-@deprecated("This object has been deprecated and will be removed in a future release. " +
-            "Please use org.apache.kafka.clients.consumer.internals.Fetcher instead.", "0.11.0.0")
-object ConsumerFetcherThread {
-
-  class FetchRequest(val underlying: kafka.api.FetchRequest) extends AbstractFetcherThread.FetchRequest {
-    private lazy val tpToOffset: Map[TopicPartition, Long] = underlying.requestInfo.map { case (tp, fetchInfo) =>
-      new TopicPartition(tp.topic, tp.partition) -> fetchInfo.offset
-    }.toMap
-    def isEmpty: Boolean = underlying.requestInfo.isEmpty
-    def offset(topicPartition: TopicPartition): Long = tpToOffset(topicPartition)
-    override def toString = underlying.toString
-  }
-
-  class PartitionData(val underlying: FetchResponsePartitionData) extends AbstractFetcherThread.PartitionData {
-    def error = underlying.error
-    def toRecords: MemoryRecords = underlying.messages.asInstanceOf[ByteBufferMessageSet].asRecords
-    def highWatermark: Long = underlying.hw
-    def exception: Option[Throwable] =
-      if (error == Errors.NONE) None else Some(ErrorMapping.exceptionFor(error.code))
-    override def toString = underlying.toString
-  }
-}
diff --git a/core/src/main/scala/kafka/consumer/ConsumerIterator.scala b/core/src/main/scala/kafka/consumer/ConsumerIterator.scala
deleted file mode 100755
index f096c55af3102..0000000000000
--- a/core/src/main/scala/kafka/consumer/ConsumerIterator.scala
+++ /dev/null
@@ -1,123 +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.consumer
-
-import kafka.utils.{IteratorTemplate, Logging}
-import java.util.concurrent.{TimeUnit, BlockingQueue}
-import kafka.serializer.Decoder
-import java.util.concurrent.atomic.AtomicReference
-import kafka.message.{MessageAndOffset, MessageAndMetadata}
-import kafka.common.{KafkaException, MessageSizeTooLargeException}
-
-
-/**
- * An iterator that blocks until a value can be read from the supplied queue.
- * The iterator takes a shutdownCommand object which can be added to the queue to trigger a shutdown
- *
- */
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.11.0.0")
-class ConsumerIterator[K, V](private val channel: BlockingQueue[FetchedDataChunk],
-                             consumerTimeoutMs: Int,
-                             private val keyDecoder: Decoder[K],
-                             private val valueDecoder: Decoder[V],
-                             val clientId: String)
-  extends IteratorTemplate[MessageAndMetadata[K, V]] with Logging {
-
-  private val current: AtomicReference[Iterator[MessageAndOffset]] = new AtomicReference(null)
-  private var currentTopicInfo: PartitionTopicInfo = null
-  private var consumedOffset: Long = -1L
-  private val consumerTopicStats = ConsumerTopicStatsRegistry.getConsumerTopicStat(clientId)
-
-  override def next(): MessageAndMetadata[K, V] = {
-    val item = super.next()
-    if(consumedOffset < 0)
-      throw new KafkaException("Offset returned by the message set is invalid %d".format(consumedOffset))
-    currentTopicInfo.resetConsumeOffset(consumedOffset)
-    val topic = currentTopicInfo.topic
-    trace("Setting %s consumed offset to %d".format(topic, consumedOffset))
-    consumerTopicStats.getConsumerTopicStats(topic).messageRate.mark()
-    consumerTopicStats.getConsumerAllTopicStats().messageRate.mark()
-    item
-  }
-
-  protected def makeNext(): MessageAndMetadata[K, V] = {
-    var currentDataChunk: FetchedDataChunk = null
-    // if we don't have an iterator, get one
-    var localCurrent = current.get()
-    if(localCurrent == null || !localCurrent.hasNext) {
-      if (consumerTimeoutMs < 0)
-        currentDataChunk = channel.take
-      else {
-        currentDataChunk = channel.poll(consumerTimeoutMs, TimeUnit.MILLISECONDS)
-        if (currentDataChunk == null) {
-          // reset state to make the iterator re-iterable
-          resetState()
-          throw new ConsumerTimeoutException
-        }
-      }
-      if(currentDataChunk eq ZookeeperConsumerConnector.shutdownCommand) {
-        debug("Received the shutdown command")
-        return allDone
-      } else {
-        currentTopicInfo = currentDataChunk.topicInfo
-        val cdcFetchOffset = currentDataChunk.fetchOffset
-        val ctiConsumeOffset = currentTopicInfo.getConsumeOffset
-        if (ctiConsumeOffset < cdcFetchOffset) {
-          error("consumed offset: %d doesn't match fetch offset: %d for %s;\n Consumer may lose data"
-            .format(ctiConsumeOffset, cdcFetchOffset, currentTopicInfo))
-          currentTopicInfo.resetConsumeOffset(cdcFetchOffset)
-        }
-        localCurrent = currentDataChunk.messages.iterator
-
-        current.set(localCurrent)
-      }
-      // if we just updated the current chunk and it is empty that means the fetch size is too small!
-      if(currentDataChunk.messages.validBytes == 0)
-        throw new MessageSizeTooLargeException("Found a message larger than the maximum fetch size of this consumer on topic " +
-                                               "%s partition %d at fetch offset %d. Increase the fetch size, or decrease the maximum message size the broker will allow."
-                                               .format(currentDataChunk.topicInfo.topic, currentDataChunk.topicInfo.partitionId, currentDataChunk.fetchOffset))
-    }
-    var item = localCurrent.next()
-    // reject the messages that have already been consumed
-    while (item.offset < currentTopicInfo.getConsumeOffset && localCurrent.hasNext) {
-      item = localCurrent.next()
-    }
-    consumedOffset = item.nextOffset
-
-    item.message.ensureValid() // validate checksum of message to ensure it is valid
-
-    new MessageAndMetadata(currentTopicInfo.topic,
-                           currentTopicInfo.partitionId,
-                           item.message,
-                           item.offset,
-                           keyDecoder,
-                           valueDecoder,
-                           item.message.timestamp,
-                           item.message.timestampType)
-  }
-
-  def clearCurrentChunk() {
-    debug("Clearing the current data chunk for this consumer iterator")
-    current.set(null)
-  }
-}
-
-@deprecated("This class has been deprecated and will be removed in a future release. " +
-            "Please use org.apache.kafka.common.errors.TimeoutException instead.", "0.11.0.0")
-class ConsumerTimeoutException() extends RuntimeException()
-
diff --git a/core/src/main/scala/kafka/consumer/ConsumerTopicStats.scala b/core/src/main/scala/kafka/consumer/ConsumerTopicStats.scala
deleted file mode 100644
index d13b327e9c15f..0000000000000
--- a/core/src/main/scala/kafka/consumer/ConsumerTopicStats.scala
+++ /dev/null
@@ -1,69 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package kafka.consumer
-
-import kafka.utils.{Pool, threadsafe, Logging}
-import java.util.concurrent.TimeUnit
-import kafka.metrics.KafkaMetricsGroup
-import kafka.common.{ClientIdTopic, ClientIdAllTopics, ClientIdAndTopic}
-
-@threadsafe
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.11.0.0")
-class ConsumerTopicMetrics(metricId: ClientIdTopic) extends KafkaMetricsGroup {
-  val tags = metricId match {
-    case ClientIdAndTopic(clientId, topic) => Map("clientId" -> clientId, "topic" -> topic)
-    case ClientIdAllTopics(clientId) => Map("clientId" -> clientId)
-  }
-
-  val messageRate = newMeter("MessagesPerSec", "messages", TimeUnit.SECONDS, tags)
-  val byteRate = newMeter("BytesPerSec", "bytes", TimeUnit.SECONDS, tags)
-}
-
-/**
- * Tracks metrics for each topic the given consumer client has consumed data from.
- * @param clientId The clientId of the given consumer client.
- */
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.11.0.0")
-class ConsumerTopicStats(clientId: String) extends Logging {
-  private val valueFactory = (k: ClientIdAndTopic) => new ConsumerTopicMetrics(k)
-  private val stats = new Pool[ClientIdAndTopic, ConsumerTopicMetrics](Some(valueFactory))
-  private val allTopicStats = new ConsumerTopicMetrics(new ClientIdAllTopics(clientId)) // to differentiate from a topic named AllTopics
-
-  def getConsumerAllTopicStats(): ConsumerTopicMetrics = allTopicStats
-
-  def getConsumerTopicStats(topic: String): ConsumerTopicMetrics = {
-    stats.getAndMaybePut(new ClientIdAndTopic(clientId, topic))
-  }
-}
-
-/**
- * Stores the topic stats information of each consumer client in a (clientId -> ConsumerTopicStats) map.
- */
-@deprecated("This object has been deprecated and will be removed in a future release.", "0.11.0.0")
-object ConsumerTopicStatsRegistry {
-  private val valueFactory = (k: String) => new ConsumerTopicStats(k)
-  private val globalStats = new Pool[String, ConsumerTopicStats](Some(valueFactory))
-
-  def getConsumerTopicStat(clientId: String) = {
-    globalStats.getAndMaybePut(clientId)
-  }
-
-  def removeConsumerTopicStat(clientId: String) {
-    globalStats.remove(clientId)
-  }
-}
\ No newline at end of file
diff --git a/core/src/main/scala/kafka/consumer/FetchRequestAndResponseStats.scala b/core/src/main/scala/kafka/consumer/FetchRequestAndResponseStats.scala
deleted file mode 100644
index 462a85ba063e6..0000000000000
--- a/core/src/main/scala/kafka/consumer/FetchRequestAndResponseStats.scala
+++ /dev/null
@@ -1,82 +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.consumer
-
-import java.util.concurrent.TimeUnit
-
-import kafka.common.{ClientIdAllBrokers, ClientIdBroker, ClientIdAndBroker}
-import kafka.metrics.{KafkaMetricsGroup, KafkaTimer}
-import kafka.utils.Pool
-
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.11.0.0")
-class FetchRequestAndResponseMetrics(metricId: ClientIdBroker) extends KafkaMetricsGroup {
-  val tags = metricId match {
-    case ClientIdAndBroker(clientId, brokerHost, brokerPort) =>
-      Map("clientId" -> clientId, "brokerHost" -> brokerHost,
-      "brokerPort" -> brokerPort.toString)
-    case ClientIdAllBrokers(clientId) =>
-      Map("clientId" -> clientId)
-  }
-
-  val requestTimer = new KafkaTimer(newTimer("FetchRequestRateAndTimeMs", TimeUnit.MILLISECONDS, TimeUnit.SECONDS, tags))
-  val requestSizeHist = newHistogram("FetchResponseSize", biased = true, tags)
-  val throttleTimeStats = newTimer("FetchRequestThrottleRateAndTimeMs", TimeUnit.MILLISECONDS, TimeUnit.SECONDS, tags)
-}
-
-/**
- * Tracks metrics of the requests made by a given consumer client to all brokers, and the responses obtained from the brokers.
- * @param clientId ClientId of the given consumer
- */
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.11.0.0")
-class FetchRequestAndResponseStats(clientId: String) {
-  private val valueFactory = (k: ClientIdBroker) => new FetchRequestAndResponseMetrics(k)
-  private val stats = new Pool[ClientIdBroker, FetchRequestAndResponseMetrics](Some(valueFactory))
-  private val allBrokersStats = new FetchRequestAndResponseMetrics(new ClientIdAllBrokers(clientId))
-
-  def getFetchRequestAndResponseAllBrokersStats(): FetchRequestAndResponseMetrics = allBrokersStats
-
-  def getFetchRequestAndResponseStats(brokerHost: String, brokerPort: Int): FetchRequestAndResponseMetrics = {
-    stats.getAndMaybePut(new ClientIdAndBroker(clientId, brokerHost, brokerPort))
-  }
-}
-
-/**
- * Stores the fetch request and response stats information of each consumer client in a (clientId -> FetchRequestAndResponseStats) map.
- */
-@deprecated("This object has been deprecated and will be removed in a future release.", "0.11.0.0")
-object FetchRequestAndResponseStatsRegistry {
-  private val valueFactory = (k: String) => new FetchRequestAndResponseStats(k)
-  private val globalStats = new Pool[String, FetchRequestAndResponseStats](Some(valueFactory))
-
-  def getFetchRequestAndResponseStats(clientId: String) = {
-    globalStats.getAndMaybePut(clientId)
-  }
-
-  def removeConsumerFetchRequestAndResponseStats(clientId: String) {
-    val pattern = (".*" + clientId + ".*").r
-    val keys = globalStats.keys
-    for (key <- keys) {
-      pattern.findFirstIn(key) match {
-        case Some(_) => globalStats.remove(key)
-        case _ =>
-      }
-    }
-  }
-}
-
-
diff --git a/core/src/main/scala/kafka/consumer/FetchedDataChunk.scala b/core/src/main/scala/kafka/consumer/FetchedDataChunk.scala
deleted file mode 100644
index 91eb874150048..0000000000000
--- a/core/src/main/scala/kafka/consumer/FetchedDataChunk.scala
+++ /dev/null
@@ -1,25 +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.consumer
-
-import kafka.message.ByteBufferMessageSet
-
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.11.0.0")
-case class FetchedDataChunk(messages: ByteBufferMessageSet,
-                            topicInfo: PartitionTopicInfo,
-                            fetchOffset: Long)
diff --git a/core/src/main/scala/kafka/consumer/KafkaStream.scala b/core/src/main/scala/kafka/consumer/KafkaStream.scala
deleted file mode 100644
index 914cedd07fd54..0000000000000
--- a/core/src/main/scala/kafka/consumer/KafkaStream.scala
+++ /dev/null
@@ -1,53 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package kafka.consumer
-
-
-import java.util.concurrent.BlockingQueue
-import kafka.serializer.Decoder
-import kafka.message.MessageAndMetadata
-
-@deprecated("This class has been deprecated and will be removed in a future release. " +
-            "Please use org.apache.kafka.streams.KafkaStreams instead.", "0.11.0.0")
-class KafkaStream[K,V](private val queue: BlockingQueue[FetchedDataChunk],
-                        consumerTimeoutMs: Int,
-                        private val keyDecoder: Decoder[K],
-                        private val valueDecoder: Decoder[V],
-                        val clientId: String)
-   extends Iterable[MessageAndMetadata[K,V]] with java.lang.Iterable[MessageAndMetadata[K,V]] {
-
-  private val iter: ConsumerIterator[K,V] =
-    new ConsumerIterator[K,V](queue, consumerTimeoutMs, keyDecoder, valueDecoder, clientId)
-
-  /**
-   *  Create an iterator over messages in the stream.
-   */
-  def iterator: ConsumerIterator[K,V] = iter
-
-  /**
-   * This method clears the queue being iterated during the consumer rebalancing. This is mainly
-   * to reduce the number of duplicates received by the consumer
-   */
-  def clear() {
-    iter.clearCurrentChunk()
-  }
-
-  override def toString: String = {
-     "%s kafka stream".format(clientId)
-  }
-}
diff --git a/core/src/main/scala/kafka/consumer/PartitionAssignor.scala b/core/src/main/scala/kafka/consumer/PartitionAssignor.scala
deleted file mode 100755
index 5d4fb8bdd0a05..0000000000000
--- a/core/src/main/scala/kafka/consumer/PartitionAssignor.scala
+++ /dev/null
@@ -1,172 +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.consumer
-
-import kafka.common.TopicAndPartition
-import kafka.utils.{Pool, CoreUtils, ZkUtils, Logging}
-
-import scala.collection.mutable
-
-@deprecated("This trait has been deprecated and will be removed in a future release. " +
-            "Please use org.apache.kafka.clients.consumer.internals.PartitionAssignor instead.", "0.11.0.0")
-trait PartitionAssignor {
-
-  /**
-   * Assigns partitions to consumer instances in a group.
-   * @return An assignment map of partition to this consumer group. This includes assignments for threads that belong
-   *         to the same consumer group.
-   */
-  def assign(ctx: AssignmentContext): Pool[String, mutable.Map[TopicAndPartition, ConsumerThreadId]]
-
-}
-
-@deprecated("This object has been deprecated and will be removed in a future release. " +
-            "Please use org.apache.kafka.clients.consumer.internals.PartitionAssignor instead.", "0.11.0.0")
-object PartitionAssignor {
-  def createInstance(assignmentStrategy: String) = assignmentStrategy match {
-    case "roundrobin" => new RoundRobinAssignor()
-    case _ => new RangeAssignor()
-  }
-}
-
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.11.0.0")
-class AssignmentContext(group: String, val consumerId: String, excludeInternalTopics: Boolean, zkUtils: ZkUtils) {
-  val myTopicThreadIds: collection.Map[String, collection.Set[ConsumerThreadId]] = {
-    val myTopicCount = TopicCount.constructTopicCount(group, consumerId, zkUtils, excludeInternalTopics)
-    myTopicCount.getConsumerThreadIdsPerTopic
-  }
-
-  val consumersForTopic: collection.Map[String, List[ConsumerThreadId]] =
-    zkUtils.getConsumersPerTopic(group, excludeInternalTopics)
-
-  // Some assignment strategies require knowledge of all topics consumed by any member of the group
-  val partitionsForTopic: collection.Map[String, Seq[Int]] =
-    zkUtils.getPartitionsForTopics(consumersForTopic.keySet.toSeq)
-
-  val consumers: Seq[String] = zkUtils.getConsumersInGroup(group).sorted
-}
-
-/**
- * The round-robin partition assignor lays out all the available partitions and all the available consumer threads. It
- * then proceeds to do a round-robin assignment from partition to consumer thread. If the subscriptions of all consumer
- * 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 consumer threads.)
- */
-@deprecated("This class has been deprecated and will be removed in a future release. " +
-            "Please use org.apache.kafka.clients.consumer.RoundRobinAssignor instead.", "0.11.0.0")
-class RoundRobinAssignor() extends PartitionAssignor with Logging {
-
-  def assign(ctx: AssignmentContext) = {
-
-    val valueFactory = (_: String) => new mutable.HashMap[TopicAndPartition, ConsumerThreadId]
-    val partitionAssignment =
-      new Pool[String, mutable.Map[TopicAndPartition, ConsumerThreadId]](Some(valueFactory))
-
-    if (ctx.consumersForTopic.nonEmpty) {
-      // Collect consumer thread ids across all topics, remove duplicates, and sort to ensure determinism
-      val allThreadIds = ctx.consumersForTopic.flatMap { case (topic, threadIds) =>
-         threadIds
-      }.toSet.toSeq.sorted
-
-      val threadAssignor = CoreUtils.circularIterator(allThreadIds)
-
-      info("Starting round-robin assignment with consumers " + ctx.consumers)
-      val allTopicPartitions = ctx.partitionsForTopic.flatMap { case (topic, partitions) =>
-        info("Consumer %s rebalancing the following partitions for topic %s: %s"
-          .format(ctx.consumerId, topic, partitions))
-        partitions.map(partition => {
-          TopicAndPartition(topic, partition)
-        })
-      }.toSeq.sortWith((topicPartition1, topicPartition2) => {
-        /*
-         * Randomize the order by taking the hashcode to reduce the likelihood of all partitions of a given topic ending
-         * up on one consumer (if it has a high enough stream count).
-         */
-        topicPartition1.toString.hashCode < topicPartition2.toString.hashCode
-      })
-
-      allTopicPartitions.foreach(topicPartition => {
-        val threadId = threadAssignor.dropWhile(threadId => !ctx.consumersForTopic(topicPartition.topic).contains(threadId)).next
-        // record the partition ownership decision
-        val assignmentForConsumer = partitionAssignment.getAndMaybePut(threadId.consumer)
-        assignmentForConsumer += (topicPartition -> threadId)
-      })
-    }
-
-    // assign Map.empty for the consumers which are not associated with topic partitions
-    ctx.consumers.foreach(consumerId => partitionAssignment.getAndMaybePut(consumerId))
-    partitionAssignment
-  }
-}
-
-/**
- * Range partitioning works on a per-topic basis. For each topic, we lay out the available partitions in numeric order
- * and the consumer threads in lexicographic order. We then divide the number of partitions by the total number of
- * consumer streams (threads) to determine the number of partitions to assign to each consumer. If it does not evenly
- * divide, then the first few consumers will have one extra partition. For example, suppose there are two consumers C1
- * and C2 with two streams each, and there are five available partitions (p0, p1, p2, p3, p4). So each consumer thread
- * will get at least one partition and the first consumer thread will get one extra partition. So the assignment will be:
- * p0 -> C1-0, p1 -> C1-0, p2 -> C1-1, p3 -> C2-0, p4 -> C2-1
- */
-@deprecated("This class has been deprecated and will be removed in a future release. " +
-            "Please use org.apache.kafka.clients.consumer.RangeAssignor instead.", "0.11.0.0")
-class RangeAssignor() extends PartitionAssignor with Logging {
-
-  def assign(ctx: AssignmentContext) = {
-    val valueFactory = (_: String) => new mutable.HashMap[TopicAndPartition, ConsumerThreadId]
-    val partitionAssignment =
-      new Pool[String, mutable.Map[TopicAndPartition, ConsumerThreadId]](Some(valueFactory))
-    for (topic <- ctx.myTopicThreadIds.keySet) {
-      val curConsumers = ctx.consumersForTopic(topic)
-      val curPartitions: Seq[Int] = ctx.partitionsForTopic(topic)
-
-      val nPartsPerConsumer = curPartitions.size / curConsumers.size
-      val nConsumersWithExtraPart = curPartitions.size % curConsumers.size
-
-      info("Consumer " + ctx.consumerId + " rebalancing the following partitions: " + curPartitions +
-        " for topic " + topic + " with consumers: " + curConsumers)
-
-      for (consumerThreadId <- curConsumers) {
-        val myConsumerPosition = curConsumers.indexOf(consumerThreadId)
-        assert(myConsumerPosition >= 0)
-        val startPart = nPartsPerConsumer * myConsumerPosition + myConsumerPosition.min(nConsumersWithExtraPart)
-        val nParts = nPartsPerConsumer + (if (myConsumerPosition + 1 > nConsumersWithExtraPart) 0 else 1)
-
-        /**
-         *   Range-partition the sorted partitions to consumers for better locality.
-         *  The first few consumers pick up an extra partition, if any.
-         */
-        if (nParts <= 0)
-          warn("No broker partitions consumed by consumer thread " + consumerThreadId + " for topic " + topic)
-        else {
-          for (i <- startPart until startPart + nParts) {
-            val partition = curPartitions(i)
-            info(consumerThreadId + " attempting to claim partition " + partition)
-            // record the partition ownership decision
-            val assignmentForConsumer = partitionAssignment.getAndMaybePut(consumerThreadId.consumer)
-            assignmentForConsumer += (TopicAndPartition(topic, partition) -> consumerThreadId)
-          }
-        }
-      }
-    }
-
-    // assign Map.empty for the consumers which are not associated with topic partitions
-    ctx.consumers.foreach(consumerId => partitionAssignment.getAndMaybePut(consumerId))
-    partitionAssignment
-  }
-}
diff --git a/core/src/main/scala/kafka/consumer/PartitionTopicInfo.scala b/core/src/main/scala/kafka/consumer/PartitionTopicInfo.scala
deleted file mode 100644
index 9a0879a198730..0000000000000
--- a/core/src/main/scala/kafka/consumer/PartitionTopicInfo.scala
+++ /dev/null
@@ -1,80 +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.consumer
-
-import java.util.concurrent._
-import java.util.concurrent.atomic._
-import kafka.message._
-import kafka.utils.Logging
-
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.11.0.0")
-class PartitionTopicInfo(val topic: String,
-                         val partitionId: Int,
-                         private val chunkQueue: BlockingQueue[FetchedDataChunk],
-                         private val consumedOffset: AtomicLong,
-                         private val fetchedOffset: AtomicLong,
-                         private val fetchSize: AtomicInteger,
-                         private val clientId: String) extends Logging {
-
-  debug("initial consumer offset of " + this + " is " + consumedOffset.get)
-  debug("initial fetch offset of " + this + " is " + fetchedOffset.get)
-
-  private val consumerTopicStats = ConsumerTopicStatsRegistry.getConsumerTopicStat(clientId)
-
-  def getConsumeOffset() = consumedOffset.get
-
-  def getFetchOffset() = fetchedOffset.get
-
-  def resetConsumeOffset(newConsumeOffset: Long) = {
-    consumedOffset.set(newConsumeOffset)
-    debug("reset consume offset of " + this + " to " + newConsumeOffset)
-  }
-
-  def resetFetchOffset(newFetchOffset: Long) = {
-    fetchedOffset.set(newFetchOffset)
-    debug("reset fetch offset of ( %s ) to %d".format(this, newFetchOffset))
-  }
-
-  /**
-   * Enqueue a message set for processing.
-   */
-  def enqueue(messages: ByteBufferMessageSet) {
-    val size = messages.validBytes
-    if(size > 0) {
-      val next = messages.shallowIterator.toSeq.last.nextOffset
-      trace("Updating fetch offset = " + fetchedOffset.get + " to " + next)
-      chunkQueue.put(new FetchedDataChunk(messages, this, fetchedOffset.get))
-      fetchedOffset.set(next)
-      debug("updated fetch offset of (%s) to %d".format(this, next))
-      consumerTopicStats.getConsumerTopicStats(topic).byteRate.mark(size)
-      consumerTopicStats.getConsumerAllTopicStats().byteRate.mark(size)
-    } else if(messages.sizeInBytes > 0) {
-      chunkQueue.put(new FetchedDataChunk(messages, this, fetchedOffset.get))
-    }
-  }
-
-  override def toString: String = topic + ":" + partitionId.toString + ": fetched offset = " + fetchedOffset.get +
-    ": consumed offset = " + consumedOffset.get
-}
-
-@deprecated("This object has been deprecated and will be removed in a future release.", "0.11.0.0")
-object PartitionTopicInfo {
-  val InvalidOffset = -1L
-
-  def isOffsetInvalid(offset: Long) = offset < 0L
-}
\ No newline at end of file
diff --git a/core/src/main/scala/kafka/consumer/SimpleConsumer.scala b/core/src/main/scala/kafka/consumer/SimpleConsumer.scala
deleted file mode 100644
index b30c9cefd96b0..0000000000000
--- a/core/src/main/scala/kafka/consumer/SimpleConsumer.scala
+++ /dev/null
@@ -1,199 +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.consumer
-
-
-import java.nio.channels.{AsynchronousCloseException, ClosedByInterruptException}
-import java.util.concurrent.TimeUnit
-
-import kafka.api._
-import kafka.network._
-import kafka.utils._
-import kafka.common.{ErrorMapping, TopicAndPartition}
-import org.apache.kafka.common.network.{NetworkReceive}
-import org.apache.kafka.common.protocol.Errors
-import org.apache.kafka.common.utils.Utils._
-
-/**
- * A consumer of kafka messages
- */
-@deprecated("This class has been deprecated and will be removed in a future release. " +
-            "Please use org.apache.kafka.clients.consumer.KafkaConsumer instead.", "0.11.0.0")
-@threadsafe
-class SimpleConsumer(val host: String,
-                     val port: Int,
-                     val soTimeout: Int,
-                     val bufferSize: Int,
-                     val clientId: String) extends Logging {
-
-  ConsumerConfig.validateClientId(clientId)
-  private val lock = new Object()
-  private val blockingChannel = new BlockingChannel(host, port, bufferSize, BlockingChannel.UseDefaultBufferSize, soTimeout)
-  private val fetchRequestAndResponseStats = FetchRequestAndResponseStatsRegistry.getFetchRequestAndResponseStats(clientId)
-  private var isClosed = false
-
-  private def connect(): BlockingChannel = {
-    close
-    blockingChannel.connect()
-    blockingChannel
-  }
-
-  private def disconnect() = {
-    debug("Disconnecting from " + formatAddress(host, port))
-    blockingChannel.disconnect()
-  }
-
-  private def reconnect() {
-    disconnect()
-    connect()
-  }
-
-  /**
-   * Unblock thread by closing channel and triggering AsynchronousCloseException if a read operation is in progress.
-   *
-   * This handles a bug found in Java 1.7 and below, where interrupting a thread can not correctly unblock
-   * the thread from waiting on ReadableByteChannel.read().
-   */
-  def disconnectToHandleJavaIOBug() = {
-    disconnect()
-  }
-
-  def close() {
-    lock synchronized {
-      disconnect()
-      isClosed = true
-    }
-  }
-
-  private def sendRequest(request: RequestOrResponse): NetworkReceive = {
-    lock synchronized {
-      var response: NetworkReceive = null
-      try {
-        getOrMakeConnection()
-        blockingChannel.send(request)
-        response = blockingChannel.receive()
-      } catch {
-        case e : ClosedByInterruptException =>
-          throw e
-        // Should not observe this exception when running Kafka with Java 1.8
-        case e: AsynchronousCloseException =>
-          throw e
-        case e : Throwable =>
-          info("Reconnect due to error:", e)
-          // retry once
-          try {
-            reconnect()
-            blockingChannel.send(request)
-            response = blockingChannel.receive()
-          } catch {
-            case e: Throwable =>
-              disconnect()
-              throw e
-          }
-      }
-      response
-    }
-  }
-
-  def send(request: TopicMetadataRequest): TopicMetadataResponse = {
-    val response = sendRequest(request)
-    TopicMetadataResponse.readFrom(response.payload())
-  }
-
-  def send(request: GroupCoordinatorRequest): GroupCoordinatorResponse = {
-    val response = sendRequest(request)
-    GroupCoordinatorResponse.readFrom(response.payload())
-  }
-
-  /**
-   *  Fetch a set of messages from a topic.
-   *
-   *  @param request  specifies the topic name, topic partition, starting byte offset, maximum bytes to be fetched.
-   *  @return a set of fetched messages
-   */
-  def fetch(request: FetchRequest): FetchResponse = {
-    var response: NetworkReceive = null
-    val specificTimer = fetchRequestAndResponseStats.getFetchRequestAndResponseStats(host, port).requestTimer
-    val aggregateTimer = fetchRequestAndResponseStats.getFetchRequestAndResponseAllBrokersStats.requestTimer
-    aggregateTimer.time {
-      specificTimer.time {
-        response = sendRequest(request)
-      }
-    }
-    val fetchResponse = FetchResponse.readFrom(response.payload(), request.versionId)
-    val fetchedSize = fetchResponse.sizeInBytes
-    fetchRequestAndResponseStats.getFetchRequestAndResponseStats(host, port).requestSizeHist.update(fetchedSize)
-    fetchRequestAndResponseStats.getFetchRequestAndResponseAllBrokersStats.requestSizeHist.update(fetchedSize)
-    fetchRequestAndResponseStats.getFetchRequestAndResponseStats(host, port).throttleTimeStats.update(fetchResponse.throttleTimeMs, TimeUnit.MILLISECONDS)
-    fetchRequestAndResponseStats.getFetchRequestAndResponseAllBrokersStats.throttleTimeStats.update(fetchResponse.throttleTimeMs, TimeUnit.MILLISECONDS)
-    fetchResponse
-  }
-
-  /**
-   *  Get a list of valid offsets (up to maxSize) before the given time.
-   *  @param request a [[kafka.api.OffsetRequest]] object.
-   *  @return a [[kafka.api.OffsetResponse]] object.
-   */
-  def getOffsetsBefore(request: OffsetRequest) = OffsetResponse.readFrom(sendRequest(request).payload())
-
-  /**
-   * Commit offsets for a topic
-   * Version 0 of the request will commit offsets to Zookeeper and version 1 and above will commit offsets to Kafka.
-   * @param request a [[kafka.api.OffsetCommitRequest]] object.
-   * @return a [[kafka.api.OffsetCommitResponse]] object.
-   */
-  def commitOffsets(request: OffsetCommitRequest) = {
-    // TODO: With KAFKA-1012, we have to first issue a ConsumerMetadataRequest and connect to the coordinator before
-    // we can commit offsets.
-    OffsetCommitResponse.readFrom(sendRequest(request).payload())
-  }
-
-  /**
-   * Fetch offsets for a topic
-   * Version 0 of the request will fetch offsets from Zookeeper and version 1 and above will fetch offsets from Kafka.
-   * @param request a [[kafka.api.OffsetFetchRequest]] object.
-   * @return a [[kafka.api.OffsetFetchResponse]] object.
-   */
-  def fetchOffsets(request: OffsetFetchRequest) = OffsetFetchResponse.readFrom(sendRequest(request).payload(), request.versionId)
-
-  private def getOrMakeConnection() {
-    if(!isClosed && !blockingChannel.isConnected) {
-      connect()
-    }
-  }
-
-  /**
-   * Get the earliest or latest offset of a given topic, partition.
-   * @param topicAndPartition Topic and partition of which the offset is needed.
-   * @param earliestOrLatest A value to indicate earliest or latest offset.
-   * @param consumerId Id of the consumer which could be a consumer client, SimpleConsumerShell or a follower broker.
-   * @return Requested offset.
-   */
-  def earliestOrLatestOffset(topicAndPartition: TopicAndPartition, earliestOrLatest: Long, consumerId: Int): Long = {
-    val request = OffsetRequest(requestInfo = Map(topicAndPartition -> PartitionOffsetRequestInfo(earliestOrLatest, 1)),
-                                clientId = clientId,
-                                replicaId = consumerId)
-    val partitionErrorAndOffset = getOffsetsBefore(request).partitionErrorAndOffsets(topicAndPartition)
-    val offset = partitionErrorAndOffset.error match {
-      case Errors.NONE => partitionErrorAndOffset.offsets.head
-      case _ => throw ErrorMapping.exceptionFor(partitionErrorAndOffset.error.code)
-    }
-    offset
-  }
-}
-
diff --git a/core/src/main/scala/kafka/consumer/TopicCount.scala b/core/src/main/scala/kafka/consumer/TopicCount.scala
deleted file mode 100755
index 2cabcaea0b746..0000000000000
--- a/core/src/main/scala/kafka/consumer/TopicCount.scala
+++ /dev/null
@@ -1,145 +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.consumer
-
-import scala.collection._
-import kafka.utils.{Json, ZKGroupDirs, ZkUtils, Logging, CoreUtils}
-import kafka.common.KafkaException
-
-@deprecated("This trait has been deprecated and will be removed in a future release.", "0.11.0.0")
-private[kafka] trait TopicCount {
-
-  def getConsumerThreadIdsPerTopic: Map[String, Set[ConsumerThreadId]]
-  def getTopicCountMap: Map[String, Int]
-  def pattern: String
-
-}
-
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.11.0.0")
-case class ConsumerThreadId(consumer: String, threadId: Int) extends Ordered[ConsumerThreadId] {
-  override def toString = "%s-%d".format(consumer, threadId)
-
-  def compare(that: ConsumerThreadId) = toString.compare(that.toString)
-}
-
-@deprecated("This object has been deprecated and will be removed in a future release.", "0.11.0.0")
-private[kafka] object TopicCount extends Logging {
-  val whiteListPattern = "white_list"
-  val blackListPattern = "black_list"
-  val staticPattern = "static"
-
-  def makeThreadId(consumerIdString: String, threadId: Int) = consumerIdString + "-" + threadId
-
-  def makeConsumerThreadIdsPerTopic(consumerIdString: String,
-                                    topicCountMap: Map[String,  Int]) = {
-    val consumerThreadIdsPerTopicMap = new mutable.HashMap[String, Set[ConsumerThreadId]]()
-    for ((topic, nConsumers) <- topicCountMap) {
-      val consumerSet = new mutable.HashSet[ConsumerThreadId]
-      assert(nConsumers >= 1)
-      for (i <- 0 until nConsumers)
-        consumerSet += ConsumerThreadId(consumerIdString, i)
-      consumerThreadIdsPerTopicMap.put(topic, consumerSet)
-    }
-    consumerThreadIdsPerTopicMap
-  }
-
-  def constructTopicCount(group: String, consumerId: String, zkUtils: ZkUtils, excludeInternalTopics: Boolean) : TopicCount = {
-    val dirs = new ZKGroupDirs(group)
-    val topicCountString = zkUtils.readData(dirs.consumerRegistryDir + "/" + consumerId)._1
-    var subscriptionPattern: String = null
-    var topMap: Map[String, Int] = null
-    try {
-      Json.parseFull(topicCountString) match {
-        case Some(js) =>
-          val consumerRegistrationMap = js.asJsonObject
-          consumerRegistrationMap.get("pattern") match {
-            case Some(pattern) => subscriptionPattern = pattern.to[String]
-            case None => throw new KafkaException("error constructing TopicCount : " + topicCountString)
-          }
-          consumerRegistrationMap.get("subscription") match {
-            case Some(sub) => topMap = sub.to[Map[String, Int]]
-            case None => throw new KafkaException("error constructing TopicCount : " + topicCountString)
-          }
-        case None => throw new KafkaException("error constructing TopicCount : " + topicCountString)
-      }
-    } catch {
-      case e: Throwable =>
-        error("error parsing consumer json string " + topicCountString, e)
-        throw e
-    }
-
-    val hasWhiteList = whiteListPattern.equals(subscriptionPattern)
-    val hasBlackList = blackListPattern.equals(subscriptionPattern)
-
-    if (topMap.isEmpty || !(hasWhiteList || hasBlackList)) {
-      new StaticTopicCount(consumerId, topMap)
-    } else {
-      val regex = topMap.head._1
-      val numStreams = topMap.head._2
-      val filter =
-        if (hasWhiteList)
-          new Whitelist(regex)
-        else
-          new Blacklist(regex)
-      new WildcardTopicCount(zkUtils, consumerId, filter, numStreams, excludeInternalTopics)
-    }
-  }
-
-  def constructTopicCount(consumerIdString: String, topicCount: Map[String, Int]) =
-    new StaticTopicCount(consumerIdString, topicCount)
-
-  def constructTopicCount(consumerIdString: String, filter: TopicFilter, numStreams: Int, zkUtils: ZkUtils, excludeInternalTopics: Boolean) =
-    new WildcardTopicCount(zkUtils, consumerIdString, filter, numStreams, excludeInternalTopics)
-
-}
-
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.11.0.0")
-private[kafka] class StaticTopicCount(val consumerIdString: String,
-                                val topicCountMap: Map[String, Int])
-                                extends TopicCount {
-
-  def getConsumerThreadIdsPerTopic = TopicCount.makeConsumerThreadIdsPerTopic(consumerIdString, topicCountMap)
-
-  def getTopicCountMap = topicCountMap
-
-  def pattern = TopicCount.staticPattern
-}
-
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.11.0.0")
-private[kafka] class WildcardTopicCount(zkUtils: ZkUtils,
-                                        consumerIdString: String,
-                                        topicFilter: TopicFilter,
-                                        numStreams: Int,
-                                        excludeInternalTopics: Boolean) extends TopicCount {
-  def getConsumerThreadIdsPerTopic = {
-    val wildcardTopics = zkUtils.getChildrenParentMayNotExist(ZkUtils.BrokerTopicsPath)
-                         .filter(topic => topicFilter.isTopicAllowed(topic, excludeInternalTopics))
-    TopicCount.makeConsumerThreadIdsPerTopic(consumerIdString, Map(wildcardTopics.map((_, numStreams)): _*))
-  }
-
-  def getTopicCountMap = Map(CoreUtils.JSONEscapeString(topicFilter.regex) -> numStreams)
-
-  def pattern: String = {
-    topicFilter match {
-      case _: Whitelist => TopicCount.whiteListPattern
-      case _: Blacklist => TopicCount.blackListPattern
-    }
-  }
-
-}
-
diff --git a/core/src/main/scala/kafka/consumer/TopicEventHandler.scala b/core/src/main/scala/kafka/consumer/TopicEventHandler.scala
deleted file mode 100644
index 2852e9bb92800..0000000000000
--- a/core/src/main/scala/kafka/consumer/TopicEventHandler.scala
+++ /dev/null
@@ -1,25 +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.consumer
-
-@deprecated("This trait has been deprecated and will be removed in a future release.", "0.11.0.0")
-trait TopicEventHandler[T] {
-
-  def handleTopicEvent(allTopics: Seq[T])
-
-}
diff --git a/core/src/main/scala/kafka/consumer/TopicFilter.scala b/core/src/main/scala/kafka/consumer/TopicFilter.scala
deleted file mode 100644
index b71b01acd5ed7..0000000000000
--- a/core/src/main/scala/kafka/consumer/TopicFilter.scala
+++ /dev/null
@@ -1,71 +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.consumer
-
-import kafka.utils.Logging
-import java.util.regex.{Pattern, PatternSyntaxException}
-
-import org.apache.kafka.common.internals.Topic
-
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.11.0.0")
-sealed abstract class TopicFilter(rawRegex: String) extends Logging {
-
-  val regex = rawRegex
-          .trim
-          .replace(',', '|')
-          .replace(" ", "")
-          .replaceAll("""^["']+""","")
-          .replaceAll("""["']+$""","") // property files may bring quotes
-
-  try {
-    Pattern.compile(regex)
-  }
-  catch {
-    case _: PatternSyntaxException =>
-      throw new RuntimeException(regex + " is an invalid regex.")
-  }
-
-  override def toString = regex
-
-  def isTopicAllowed(topic: String, excludeInternalTopics: Boolean): Boolean
-}
-
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.11.0.0")
-case class Whitelist(rawRegex: String) extends TopicFilter(rawRegex) {
-  override def isTopicAllowed(topic: String, excludeInternalTopics: Boolean) = {
-    val allowed = topic.matches(regex) && !(Topic.isInternal(topic) && excludeInternalTopics)
-
-    debug("%s %s".format(
-      topic, if (allowed) "allowed" else "filtered"))
-
-    allowed
-  }
-}
-
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.11.0.0")
-case class Blacklist(rawRegex: String) extends TopicFilter(rawRegex) {
-  override def isTopicAllowed(topic: String, excludeInternalTopics: Boolean) = {
-    val allowed = (!topic.matches(regex)) && !(Topic.isInternal(topic) && excludeInternalTopics)
-
-    debug("%s %s".format(
-      topic, if (allowed) "allowed" else "filtered"))
-
-    allowed
-  }
-}
-
diff --git a/core/src/main/scala/kafka/consumer/ZookeeperConsumerConnector.scala b/core/src/main/scala/kafka/consumer/ZookeeperConsumerConnector.scala
deleted file mode 100755
index bb5fc0fc2d87d..0000000000000
--- a/core/src/main/scala/kafka/consumer/ZookeeperConsumerConnector.scala
+++ /dev/null
@@ -1,1039 +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.consumer
-
-import java.net.InetAddress
-import java.util.UUID
-import java.util.concurrent._
-import java.util.concurrent.atomic._
-import java.util.concurrent.locks.ReentrantLock
-
-import com.yammer.metrics.core.Gauge
-import kafka.api._
-import kafka.client.ClientUtils
-import kafka.cluster._
-import kafka.common._
-import kafka.javaapi.consumer.ConsumerRebalanceListener
-import kafka.metrics._
-import kafka.network.BlockingChannel
-import kafka.serializer._
-import kafka.utils.CoreUtils.inLock
-import kafka.utils.ZkUtils._
-import kafka.utils._
-import org.I0Itec.zkclient.exception.ZkNodeExistsException
-import org.I0Itec.zkclient.{IZkChildListener, IZkDataListener, IZkStateListener}
-import org.apache.kafka.common.protocol.Errors
-import org.apache.kafka.common.security.JaasUtils
-import org.apache.kafka.common.utils.Time
-import org.apache.zookeeper.Watcher.Event.KeeperState
-
-import scala.collection._
-import scala.collection.JavaConverters._
-
-/**
- * This class handles the consumers interaction with zookeeper
- *
- * Directories:
- * 1. Consumer id registry:
- * /consumers/[group_id]/ids/[consumer_id] -> topic1,...topicN
- * A consumer has a unique consumer id within a consumer group. A consumer registers its id as an ephemeral znode
- * and puts all topics that it subscribes to as the value of the znode. The znode is deleted when the client is gone.
- * A consumer subscribes to event changes of the consumer id registry within its group.
- *
- * The consumer id is picked up from configuration, instead of the sequential id assigned by ZK. Generated sequential
- * ids are hard to recover during temporary connection loss to ZK, since it's difficult for the client to figure out
- * whether the creation of a sequential znode has succeeded or not. More details can be found at
- * (http://wiki.apache.org/hadoop/ZooKeeper/ErrorHandling)
- *
- * 2. Broker node registry:
- * /brokers/[0...N] --> { "host" : "host:port",
- *                        "topics" : {"topic1": ["partition1" ... "partitionN"], ...,
- *                                    "topicN": ["partition1" ... "partitionN"] } }
- * This is a list of all present broker brokers. A unique logical node id is configured on each broker node. A broker
- * node registers itself on start-up and creates a znode with the logical node id under /brokers. The value of the znode
- * is a JSON String that contains (1) the host name and the port the broker is listening to, (2) a list of topics that
- * the broker serves, (3) a list of logical partitions assigned to each topic on the broker.
- * A consumer subscribes to event changes of the broker node registry.
- *
- * 3. Partition owner registry:
- * /consumers/[group_id]/owner/[topic]/[broker_id-partition_id] --> consumer_node_id
- * This stores the mapping before broker partitions and consumers. Each partition is owned by a unique consumer
- * within a consumer group. The mapping is reestablished after each rebalancing.
- *
- * 4. Consumer offset tracking:
- * /consumers/[group_id]/offsets/[topic]/[broker_id-partition_id] --> offset_counter_value
- * Each consumer tracks the offset of the latest message consumed for each partition.
- *
- */
-@deprecated("This object has been deprecated and will be removed in a future release.", "0.11.0.0")
-private[kafka] object ZookeeperConsumerConnector {
-  val shutdownCommand: FetchedDataChunk = new FetchedDataChunk(null, null, -1L)
-}
-
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.11.0.0")
-private[kafka] class ZookeeperConsumerConnector(val config: ConsumerConfig,
-                                                val enableFetcher: Boolean) // for testing only
-        extends ConsumerConnector with Logging with KafkaMetricsGroup {
-
-  private val isShuttingDown = new AtomicBoolean(false)
-  private val rebalanceLock = new Object
-  private var fetcher: Option[ConsumerFetcherManager] = None
-  private var zkUtils: ZkUtils = null
-  private var topicRegistry = new Pool[String, Pool[Int, PartitionTopicInfo]]
-  private val checkpointedZkOffsets = new Pool[TopicAndPartition, Long]
-  private val topicThreadIdAndQueues = new Pool[(String, ConsumerThreadId), BlockingQueue[FetchedDataChunk]]
-  private val scheduler = new KafkaScheduler(threads = 1, threadNamePrefix = "kafka-consumer-scheduler-")
-  private val messageStreamCreated = new AtomicBoolean(false)
-
-  private var sessionExpirationListener: ZKSessionExpireListener = null
-  private var topicPartitionChangeListener: ZKTopicPartitionChangeListener = null
-  private var loadBalancerListener: ZKRebalancerListener = null
-
-  private var offsetsChannel: BlockingChannel = null
-  private val offsetsChannelLock = new Object
-
-  private var wildcardTopicWatcher: ZookeeperTopicEventWatcher = null
-  private var consumerRebalanceListener: ConsumerRebalanceListener = null
-
-  // useful for tracking migration of consumers to store offsets in kafka
-  private val kafkaCommitMeter = newMeter("KafkaCommitsPerSec", "commits", TimeUnit.SECONDS, Map("clientId" -> config.clientId))
-  private val zkCommitMeter = newMeter("ZooKeeperCommitsPerSec", "commits", TimeUnit.SECONDS, Map("clientId" -> config.clientId))
-  private val rebalanceTimer = new KafkaTimer(newTimer("RebalanceRateAndTime", TimeUnit.MILLISECONDS, TimeUnit.SECONDS, Map("clientId" -> config.clientId)))
-
-  newGauge(
-    "yammer-metrics-count",
-    new Gauge[Int] {
-      def value = {
-        com.yammer.metrics.Metrics.defaultRegistry().allMetrics().size()
-      }
-    }
-  )
-
-  val consumerIdString = {
-    var consumerUuid : String = null
-    config.consumerId match {
-      case Some(consumerId) // for testing only
-      => consumerUuid = consumerId
-      case None // generate unique consumerId automatically
-      => val uuid = UUID.randomUUID()
-      consumerUuid = "%s-%d-%s".format(
-        InetAddress.getLocalHost.getHostName, System.currentTimeMillis,
-        uuid.getMostSignificantBits().toHexString.substring(0,8))
-    }
-    config.groupId + "_" + consumerUuid
-  }
-  this.logIdent = "[" + consumerIdString + "], "
-
-  connectZk()
-  createFetcher()
-  ensureOffsetManagerConnected()
-
-  if (config.autoCommitEnable) {
-    scheduler.startup
-    info("starting auto committer every " + config.autoCommitIntervalMs + " ms")
-    scheduler.schedule("kafka-consumer-autocommit",
-                       autoCommit _,
-                       delay = config.autoCommitIntervalMs,
-                       period = config.autoCommitIntervalMs,
-                       unit = TimeUnit.MILLISECONDS)
-  }
-
-  KafkaMetricsReporter.startReporters(config.props)
-  AppInfo.registerInfo()
-
-  def this(config: ConsumerConfig) = this(config, true)
-
-  def createMessageStreams(topicCountMap: Map[String,Int]): Map[String, List[KafkaStream[Array[Byte],Array[Byte]]]] =
-    createMessageStreams(topicCountMap, new DefaultDecoder(), new DefaultDecoder())
-
-  def createMessageStreams[K,V](topicCountMap: Map[String,Int], keyDecoder: Decoder[K], valueDecoder: Decoder[V])
-      : Map[String, List[KafkaStream[K,V]]] = {
-    if (messageStreamCreated.getAndSet(true))
-      throw new MessageStreamsExistException(this.getClass.getSimpleName +
-                                   " can create message streams at most once",null)
-    consume(topicCountMap, keyDecoder, valueDecoder)
-  }
-
-  def createMessageStreamsByFilter[K,V](topicFilter: TopicFilter,
-                                        numStreams: Int,
-                                        keyDecoder: Decoder[K] = new DefaultDecoder(),
-                                        valueDecoder: Decoder[V] = new DefaultDecoder()) = {
-    val wildcardStreamsHandler = new WildcardStreamsHandler[K,V](topicFilter, numStreams, keyDecoder, valueDecoder)
-    wildcardStreamsHandler.streams
-  }
-
-  def setConsumerRebalanceListener(listener: ConsumerRebalanceListener) {
-    if (messageStreamCreated.get())
-      throw new MessageStreamsExistException(this.getClass.getSimpleName +
-        " can only set consumer rebalance listener before creating streams",null)
-    consumerRebalanceListener = listener
-  }
-
-  private def createFetcher() {
-    if (enableFetcher)
-      fetcher = Some(new ConsumerFetcherManager(consumerIdString, config, zkUtils))
-  }
-
-  private def connectZk() {
-    info("Connecting to zookeeper instance at " + config.zkConnect)
-    zkUtils = ZkUtils(config.zkConnect,
-                      config.zkSessionTimeoutMs,
-                      config.zkConnectionTimeoutMs,
-                      JaasUtils.isZkSecurityEnabled())
-  }
-
-  // Blocks until the offset manager is located and a channel is established to it.
-  private def ensureOffsetManagerConnected() {
-    if (config.offsetsStorage == "kafka") {
-      if (offsetsChannel == null || !offsetsChannel.isConnected)
-        offsetsChannel = ClientUtils.channelToOffsetManager(config.groupId, zkUtils,
-          config.offsetsChannelSocketTimeoutMs, config.offsetsChannelBackoffMs)
-
-      debug("Connected to offset manager %s:%d.".format(offsetsChannel.host, offsetsChannel.port))
-    }
-  }
-
-  def shutdown() {
-    val canShutdown = isShuttingDown.compareAndSet(false, true)
-    if (canShutdown) {
-      info("ZKConsumerConnector shutting down")
-      val startTime = System.nanoTime()
-      KafkaMetricsGroup.removeAllConsumerMetrics(config.clientId)
-      if (wildcardTopicWatcher != null)
-        wildcardTopicWatcher.shutdown()
-      rebalanceLock synchronized {
-        try {
-          if (config.autoCommitEnable)
-            scheduler.shutdown()
-          fetcher.foreach(_.stopConnections())
-          sendShutdownToAllQueues()
-          if (config.autoCommitEnable)
-            commitOffsets(true)
-          if (zkUtils != null) {
-            zkUtils.close()
-            zkUtils = null
-          }
-
-          if (offsetsChannel != null) offsetsChannel.disconnect()
-        } catch {
-          case e: Throwable =>
-            fatal("error during consumer connector shutdown", e)
-        }
-        info("ZKConsumerConnector shutdown completed in " + (System.nanoTime() - startTime) / 1000000 + " ms")
-      }
-    }
-  }
-
-  def consume[K, V](topicCountMap: scala.collection.Map[String,Int], keyDecoder: Decoder[K], valueDecoder: Decoder[V])
-      : Map[String,List[KafkaStream[K,V]]] = {
-    debug("entering consume ")
-    if (topicCountMap == null)
-      throw new RuntimeException("topicCountMap is null")
-
-    val topicCount = TopicCount.constructTopicCount(consumerIdString, topicCountMap)
-
-    val topicThreadIds = topicCount.getConsumerThreadIdsPerTopic
-
-    // make a list of (queue,stream) pairs, one pair for each threadId
-    val queuesAndStreams = topicThreadIds.values.map(threadIdSet =>
-      threadIdSet.map(_ => {
-        val queue =  new LinkedBlockingQueue[FetchedDataChunk](config.queuedMaxMessages)
-        val stream = new KafkaStream[K,V](
-          queue, config.consumerTimeoutMs, keyDecoder, valueDecoder, config.clientId)
-        (queue, stream)
-      })
-    ).flatten.toList
-
-    val dirs = new ZKGroupDirs(config.groupId)
-    registerConsumerInZK(dirs, consumerIdString, topicCount)
-    reinitializeConsumer(topicCount, queuesAndStreams)
-
-    loadBalancerListener.kafkaMessageAndMetadataStreams.asInstanceOf[Map[String, List[KafkaStream[K,V]]]]
-  }
-
-  // this API is used by unit tests only
-  def getTopicRegistry: Pool[String, Pool[Int, PartitionTopicInfo]] = topicRegistry
-
-  private def registerConsumerInZK(dirs: ZKGroupDirs, consumerIdString: String, topicCount: TopicCount) {
-    info("begin registering consumer " + consumerIdString + " in ZK")
-    val timestamp = Time.SYSTEM.milliseconds.toString
-    val consumerRegistrationInfo = Json.encode(Map("version" -> 1, "subscription" -> topicCount.getTopicCountMap, "pattern" -> topicCount.pattern,
-                                                  "timestamp" -> timestamp))
-    val zkWatchedEphemeral = new ZKCheckedEphemeral(dirs.
-                                                    consumerRegistryDir + "/" + consumerIdString,
-                                                    consumerRegistrationInfo,
-                                                    zkUtils.zkConnection.getZookeeper,
-                                                    false)
-    zkWatchedEphemeral.create()
-
-    info("end registering consumer " + consumerIdString + " in ZK")
-  }
-
-  private def sendShutdownToAllQueues() = {
-    for (queue <- topicThreadIdAndQueues.values.toSet[BlockingQueue[FetchedDataChunk]]) {
-      debug("Clearing up queue")
-      queue.clear()
-      queue.put(ZookeeperConsumerConnector.shutdownCommand)
-      debug("Cleared queue and sent shutdown command")
-    }
-  }
-
-  def autoCommit() {
-    trace("auto committing")
-    try {
-      commitOffsets(isAutoCommit = false)
-    }
-    catch {
-      case t: Throwable =>
-      // log it and let it go
-        error("exception during autoCommit: ", t)
-    }
-  }
-
-  def commitOffsetToZooKeeper(topicPartition: TopicAndPartition, offset: Long) {
-    if (checkpointedZkOffsets.get(topicPartition) != offset) {
-      val topicDirs = new ZKGroupTopicDirs(config.groupId, topicPartition.topic)
-      zkUtils.updatePersistentPath(topicDirs.consumerOffsetDir + "/" + topicPartition.partition, offset.toString)
-      checkpointedZkOffsets.put(topicPartition, offset)
-      zkCommitMeter.mark()
-    }
-  }
-
-  /**
-   * KAFKA-1743: This method added for backward compatibility.
-   */
-  def commitOffsets { commitOffsets(true) }
-
-  def commitOffsets(isAutoCommit: Boolean) {
-
-    val offsetsToCommit =
-      immutable.Map(topicRegistry.values.flatMap { partitionTopicInfos =>
-        partitionTopicInfos.values.map { info =>
-          TopicAndPartition(info.topic, info.partitionId) -> OffsetAndMetadata(info.getConsumeOffset())
-        }
-      }.toSeq: _*)
-
-    commitOffsets(offsetsToCommit, isAutoCommit)
-
-  }
-
-  def commitOffsets(offsetsToCommit: immutable.Map[TopicAndPartition, OffsetAndMetadata], isAutoCommit: Boolean) {
-    trace("OffsetMap: %s".format(offsetsToCommit))
-    var retriesRemaining = 1 + (if (isAutoCommit) 0 else config.offsetsCommitMaxRetries) // no retries for commits from auto-commit
-    var done = false
-    while (!done) {
-      val committed = offsetsChannelLock synchronized {
-        // committed when we receive either no error codes or only MetadataTooLarge errors
-        if (offsetsToCommit.size > 0) {
-          if (config.offsetsStorage == "zookeeper") {
-            offsetsToCommit.foreach { case (topicAndPartition, offsetAndMetadata) =>
-              commitOffsetToZooKeeper(topicAndPartition, offsetAndMetadata.offset)
-            }
-            true
-          } else {
-            val offsetCommitRequest = OffsetCommitRequest(config.groupId, offsetsToCommit, clientId = config.clientId)
-            ensureOffsetManagerConnected()
-            try {
-              kafkaCommitMeter.mark(offsetsToCommit.size)
-              offsetsChannel.send(offsetCommitRequest)
-              val offsetCommitResponse = OffsetCommitResponse.readFrom(offsetsChannel.receive().payload())
-              trace("Offset commit response: %s.".format(offsetCommitResponse))
-
-              val (commitFailed, retryableIfFailed, shouldRefreshCoordinator, errorCount) = {
-                offsetCommitResponse.commitStatus.foldLeft(false, false, false, 0) { case (folded, (topicPartition, error)) =>
-
-                  if (error == Errors.NONE && config.dualCommitEnabled) {
-                    val offset = offsetsToCommit(topicPartition).offset
-                    commitOffsetToZooKeeper(topicPartition, offset)
-                  }
-
-                  (folded._1 || // update commitFailed
-                    error != Errors.NONE,
-
-                    folded._2 || // update retryableIfFailed - (only metadata too large is not retryable)
-                      (error != Errors.NONE && error != Errors.OFFSET_METADATA_TOO_LARGE),
-
-                    folded._3 || // update shouldRefreshCoordinator
-                      error == Errors.NOT_COORDINATOR ||
-                      error == Errors.COORDINATOR_NOT_AVAILABLE,
-
-                    // update error count
-                    folded._4 + (if (error != Errors.NONE) 1 else 0))
-                }
-              }
-              debug(errorCount + " errors in offset commit response.")
-
-
-              if (shouldRefreshCoordinator) {
-                debug("Could not commit offsets (because offset coordinator has moved or is unavailable).")
-                offsetsChannel.disconnect()
-              }
-
-              if (commitFailed && retryableIfFailed)
-                false
-              else
-                true
-            }
-            catch {
-              case t: Throwable =>
-                error("Error while committing offsets.", t)
-                offsetsChannel.disconnect()
-                false
-            }
-          }
-        } else {
-          debug("No updates to offsets since last commit.")
-          true
-        }
-      }
-
-      done = {
-        retriesRemaining -= 1
-        retriesRemaining == 0 || committed
-      }
-
-      if (!done) {
-        debug("Retrying offset commit in %d ms".format(config.offsetsChannelBackoffMs))
-        Thread.sleep(config.offsetsChannelBackoffMs)
-      }
-    }
-  }
-
-  private def fetchOffsetFromZooKeeper(topicPartition: TopicAndPartition) = {
-    val dirs = new ZKGroupTopicDirs(config.groupId, topicPartition.topic)
-    val offsetString = zkUtils.readDataMaybeNull(dirs.consumerOffsetDir + "/" + topicPartition.partition)._1
-    offsetString match {
-      case Some(offsetStr) => (topicPartition, OffsetMetadataAndError(offsetStr.toLong))
-      case None => (topicPartition, OffsetMetadataAndError.NoOffset)
-    }
-  }
-
-  private def fetchOffsets(partitions: Seq[TopicAndPartition]) = {
-    if (partitions.isEmpty)
-      Some(OffsetFetchResponse(Map.empty))
-    else if (config.offsetsStorage == "zookeeper") {
-      val offsets = partitions.map(fetchOffsetFromZooKeeper)
-      Some(OffsetFetchResponse(immutable.Map(offsets:_*)))
-    } else {
-      val offsetFetchRequest = OffsetFetchRequest(groupId = config.groupId, requestInfo = partitions, clientId = config.clientId)
-
-      var offsetFetchResponseOpt: Option[OffsetFetchResponse] = None
-      while (!isShuttingDown.get && !offsetFetchResponseOpt.isDefined) {
-        offsetFetchResponseOpt = offsetsChannelLock synchronized {
-          ensureOffsetManagerConnected()
-          try {
-            offsetsChannel.send(offsetFetchRequest)
-            val offsetFetchResponse = OffsetFetchResponse.readFrom(offsetsChannel.receive().payload())
-            trace("Offset fetch response: %s.".format(offsetFetchResponse))
-
-            val (leaderChanged, loadInProgress) =
-              offsetFetchResponse.requestInfo.values.foldLeft(false, false) { case (folded, offsetMetadataAndError) =>
-                (folded._1 || (offsetMetadataAndError.error == Errors.NOT_COORDINATOR),
-                 folded._2 || (offsetMetadataAndError.error == Errors.COORDINATOR_LOAD_IN_PROGRESS))
-              }
-
-            if (leaderChanged) {
-              offsetsChannel.disconnect()
-              debug("Could not fetch offsets (because offset manager has moved).")
-              None // retry
-            }
-            else if (loadInProgress) {
-              debug("Could not fetch offsets (because offset cache is being loaded).")
-              None // retry
-            }
-            else {
-              if (config.dualCommitEnabled) {
-                // if dual-commit is enabled (i.e., if a consumer group is migrating offsets to kafka), then pick the
-                // maximum between offsets in zookeeper and kafka.
-                val kafkaOffsets = offsetFetchResponse.requestInfo
-                val mostRecentOffsets = kafkaOffsets.map { case (topicPartition, kafkaOffset) =>
-                  val zkOffset = fetchOffsetFromZooKeeper(topicPartition)._2.offset
-                  val mostRecentOffset = zkOffset.max(kafkaOffset.offset)
-                  (topicPartition, OffsetMetadataAndError(mostRecentOffset, kafkaOffset.metadata, Errors.NONE))
-                }
-                Some(OffsetFetchResponse(mostRecentOffsets))
-              }
-              else
-                Some(offsetFetchResponse)
-            }
-          }
-          catch {
-            case e: Exception =>
-              warn("Error while fetching offsets from %s:%d. Possible cause: %s".format(offsetsChannel.host, offsetsChannel.port, e.getMessage))
-              offsetsChannel.disconnect()
-              None // retry
-          }
-        }
-
-        if (offsetFetchResponseOpt.isEmpty) {
-          debug("Retrying offset fetch in %d ms".format(config.offsetsChannelBackoffMs))
-          Thread.sleep(config.offsetsChannelBackoffMs)
-        }
-      }
-
-      offsetFetchResponseOpt
-    }
-  }
-
-
-  class ZKSessionExpireListener(val dirs: ZKGroupDirs,
-                                 val consumerIdString: String,
-                                 val topicCount: TopicCount,
-                                 val loadBalancerListener: ZKRebalancerListener)
-    extends IZkStateListener {
-    @throws[Exception]
-    def handleStateChanged(state: KeeperState) {
-      // do nothing, since zkclient will do reconnect for us.
-    }
-
-    /**
-     * Called after the zookeeper session has expired and a new session has been created. You would have to re-create
-     * any ephemeral nodes here.
-     *
-     * @throws Exception
-     *             On any error.
-     */
-    @throws[Exception]
-    def handleNewSession() {
-      /**
-       *  When we get a SessionExpired event, we lost all ephemeral nodes and zkclient has reestablished a
-       *  connection for us. We need to release the ownership of the current consumer and re-register this
-       *  consumer in the consumer registry and trigger a rebalance.
-       */
-      info("ZK expired; release old broker parition ownership; re-register consumer " + consumerIdString)
-      loadBalancerListener.resetState()
-      registerConsumerInZK(dirs, consumerIdString, topicCount)
-      // explicitly trigger load balancing for this consumer
-      loadBalancerListener.syncedRebalance()
-      // There is no need to resubscribe to child and state changes.
-      // The child change watchers will be set inside rebalance when we read the children list.
-    }
-
-    override def handleSessionEstablishmentError(error: Throwable): Unit = {
-      fatal("Could not establish session with zookeeper", error)
-    }
-  }
-
-  class ZKTopicPartitionChangeListener(val loadBalancerListener: ZKRebalancerListener)
-    extends IZkDataListener {
-
-    def handleDataChange(dataPath : String, data: Object) {
-      try {
-        info("Topic info for path " + dataPath + " changed to " + data.toString + ", triggering rebalance")
-        // queue up the rebalance event
-        loadBalancerListener.rebalanceEventTriggered()
-        // There is no need to re-subscribe the watcher since it will be automatically
-        // re-registered upon firing of this event by zkClient
-      } catch {
-        case e: Throwable => error("Error while handling topic partition change for data path " + dataPath, e )
-      }
-    }
-
-    @throws[Exception]
-    def handleDataDeleted(dataPath : String) {
-      // TODO: This need to be implemented when we support delete topic
-      warn("Topic for path " + dataPath + " gets deleted, which should not happen at this time")
-    }
-  }
-
-  class ZKRebalancerListener(val group: String, val consumerIdString: String,
-                             val kafkaMessageAndMetadataStreams: mutable.Map[String,List[KafkaStream[_,_]]])
-    extends IZkChildListener {
-
-    private val partitionAssignor = PartitionAssignor.createInstance(config.partitionAssignmentStrategy)
-
-    private var isWatcherTriggered = false
-    private val lock = new ReentrantLock
-    private val cond = lock.newCondition()
-
-    @volatile private var allTopicsOwnedPartitionsCount = 0
-    newGauge("OwnedPartitionsCount",
-      new Gauge[Int] {
-        def value() = allTopicsOwnedPartitionsCount
-      },
-      Map("clientId" -> config.clientId, "groupId" -> config.groupId))
-
-    private def ownedPartitionsCountMetricTags(topic: String) = Map("clientId" -> config.clientId, "groupId" -> config.groupId, "topic" -> topic)
-
-    private val watcherExecutorThread = new Thread(consumerIdString + "_watcher_executor") {
-      override def run() {
-        info("starting watcher executor thread for consumer " + consumerIdString)
-        var doRebalance = false
-        while (!isShuttingDown.get) {
-          try {
-            lock.lock()
-            try {
-              if (!isWatcherTriggered)
-                cond.await(1000, TimeUnit.MILLISECONDS) // wake up periodically so that it can check the shutdown flag
-            } finally {
-              doRebalance = isWatcherTriggered
-              isWatcherTriggered = false
-              lock.unlock()
-            }
-            if (doRebalance)
-              syncedRebalance
-          } catch {
-            case t: Throwable => error("error during syncedRebalance", t)
-          }
-        }
-        info("stopping watcher executor thread for consumer " + consumerIdString)
-      }
-    }
-    watcherExecutorThread.start()
-
-    @throws[Exception]
-    def handleChildChange(parentPath : String, curChilds : java.util.List[String]) {
-      rebalanceEventTriggered()
-    }
-
-    def rebalanceEventTriggered() {
-      inLock(lock) {
-        isWatcherTriggered = true
-        cond.signalAll()
-      }
-    }
-
-    private def deletePartitionOwnershipFromZK(topic: String, partition: Int) {
-      val topicDirs = new ZKGroupTopicDirs(group, topic)
-      val znode = topicDirs.consumerOwnerDir + "/" + partition
-      zkUtils.deletePath(znode)
-      debug("Consumer " + consumerIdString + " releasing " + znode)
-    }
-
-    private def releasePartitionOwnership(localTopicRegistry: Pool[String, Pool[Int, PartitionTopicInfo]])= {
-      info("Releasing partition ownership")
-      for ((topic, infos) <- localTopicRegistry) {
-        for(partition <- infos.keys) {
-          deletePartitionOwnershipFromZK(topic, partition)
-        }
-        removeMetric("OwnedPartitionsCount", ownedPartitionsCountMetricTags(topic))
-        localTopicRegistry.remove(topic)
-      }
-      allTopicsOwnedPartitionsCount = 0
-    }
-
-    def resetState() {
-      topicRegistry.clear
-    }
-
-    def syncedRebalance() {
-      rebalanceLock synchronized {
-        rebalanceTimer.time {
-          for (i <- 0 until config.rebalanceMaxRetries) {
-            if(isShuttingDown.get())  {
-              return
-            }
-            info("begin rebalancing consumer " + consumerIdString + " try #" + i)
-            var done = false
-            var cluster: Cluster = null
-            try {
-              cluster = zkUtils.getCluster()
-              done = rebalance(cluster)
-            } catch {
-              case e: Throwable =>
-                /** occasionally, we may hit a ZK exception because the ZK state is changing while we are iterating.
-                  * For example, a ZK node can disappear between the time we get all children and the time we try to get
-                  * the value of a child. Just let this go since another rebalance will be triggered.
-                  **/
-                info("exception during rebalance ", e)
-            }
-            info("end rebalancing consumer " + consumerIdString + " try #" + i)
-            if (done) {
-              return
-            } else {
-              /* Here the cache is at a risk of being stale. To take future rebalancing decisions correctly, we should
-               * clear the cache */
-              info("Rebalancing attempt failed. Clearing the cache before the next rebalancing operation is triggered")
-            }
-            // stop all fetchers and clear all the queues to avoid data duplication
-            closeFetchersForQueues(cluster, kafkaMessageAndMetadataStreams, topicThreadIdAndQueues.map(q => q._2))
-            Thread.sleep(config.rebalanceBackoffMs)
-          }
-        }
-      }
-
-      throw new ConsumerRebalanceFailedException(consumerIdString + " can't rebalance after " + config.rebalanceMaxRetries +" retries")
-    }
-
-    private def rebalance(cluster: Cluster): Boolean = {
-      val myTopicThreadIdsMap = TopicCount.constructTopicCount(
-        group, consumerIdString, zkUtils, config.excludeInternalTopics).getConsumerThreadIdsPerTopic
-      val brokers = zkUtils.getAllBrokersInCluster()
-      if (brokers.size == 0) {
-        // This can happen in a rare case when there are no brokers available in the cluster when the consumer is started.
-        // We log a warning and register for child changes on brokers/id so that rebalance can be triggered when the brokers
-        // are up.
-        warn("no brokers found when trying to rebalance.")
-        zkUtils.subscribeChildChanges(BrokerIdsPath, loadBalancerListener)
-        true
-      }
-      else {
-        /**
-         * fetchers must be stopped to avoid data duplication, since if the current
-         * rebalancing attempt fails, the partitions that are released could be owned by another consumer.
-         * But if we don't stop the fetchers first, this consumer would continue returning data for released
-         * partitions in parallel. So, not stopping the fetchers leads to duplicate data.
-         */
-        closeFetchers(cluster, kafkaMessageAndMetadataStreams, myTopicThreadIdsMap)
-        if (consumerRebalanceListener != null) {
-          info("Invoking rebalance listener before relasing partition ownerships.")
-          consumerRebalanceListener.beforeReleasingPartitions(
-            if (topicRegistry.size == 0)
-              new java.util.HashMap[String, java.util.Set[java.lang.Integer]]
-            else
-              topicRegistry.map(topics =>
-                topics._1 -> topics._2.keys   // note this is incorrect, see KAFKA-2284
-              ).toMap.asJava.asInstanceOf[java.util.Map[String, java.util.Set[java.lang.Integer]]]
-          )
-        }
-        releasePartitionOwnership(topicRegistry)
-        val assignmentContext = new AssignmentContext(group, consumerIdString, config.excludeInternalTopics, zkUtils)
-        val globalPartitionAssignment = partitionAssignor.assign(assignmentContext)
-        val partitionAssignment = globalPartitionAssignment.get(assignmentContext.consumerId)
-        val currentTopicRegistry = new Pool[String, Pool[Int, PartitionTopicInfo]](
-          valueFactory = Some((_: String) => new Pool[Int, PartitionTopicInfo]))
-
-        // fetch current offsets for all topic-partitions
-        val topicPartitions = partitionAssignment.keySet.toSeq
-
-        val offsetFetchResponseOpt = fetchOffsets(topicPartitions)
-
-        if (isShuttingDown.get || !offsetFetchResponseOpt.isDefined)
-          false
-        else {
-          val offsetFetchResponse = offsetFetchResponseOpt.get
-          topicPartitions.foreach { case tp@ TopicAndPartition(topic, partition) =>
-            val offset = offsetFetchResponse.requestInfo(tp).offset
-            val threadId = partitionAssignment(tp)
-            addPartitionTopicInfo(currentTopicRegistry, partition, topic, offset, threadId)
-          }
-
-          /**
-           * move the partition ownership here, since that can be used to indicate a truly successful re-balancing attempt
-           * A rebalancing attempt is completed successfully only after the fetchers have been started correctly
-           */
-          if(reflectPartitionOwnershipDecision(partitionAssignment)) {
-            allTopicsOwnedPartitionsCount = partitionAssignment.size
-
-            partitionAssignment.view.groupBy { case (topicPartition, _) => topicPartition.topic }
-                                      .foreach { case (topic, partitionThreadPairs) =>
-              newGauge("OwnedPartitionsCount",
-                new Gauge[Int] {
-                  def value() = partitionThreadPairs.size
-                },
-                ownedPartitionsCountMetricTags(topic))
-            }
-
-            topicRegistry = currentTopicRegistry
-            // Invoke beforeStartingFetchers callback if the consumerRebalanceListener is set.
-            if (consumerRebalanceListener != null) {
-              info("Invoking rebalance listener before starting fetchers.")
-
-              // Partition assignor returns the global partition assignment organized as a map of [TopicPartition, ThreadId]
-              // per consumer, and we need to re-organize it to a map of [Partition, ThreadId] per topic before passing
-              // to the rebalance callback.
-              val partitionAssginmentGroupByTopic = globalPartitionAssignment.values.flatten.groupBy[String] {
-                case (topicPartition, _) => topicPartition.topic
-              }
-              val partitionAssigmentMapForCallback = partitionAssginmentGroupByTopic.map({
-                case (topic, partitionOwnerShips) =>
-                  val partitionOwnershipForTopicScalaMap = partitionOwnerShips.map({
-                    case (topicAndPartition, consumerThreadId) =>
-                      (topicAndPartition.partition: Integer) -> consumerThreadId
-                  }).toMap
-                  topic -> partitionOwnershipForTopicScalaMap.asJava
-              })
-              consumerRebalanceListener.beforeStartingFetchers(
-                consumerIdString,
-                partitionAssigmentMapForCallback.asJava
-              )
-            }
-            updateFetcher(cluster)
-            true
-          } else {
-            false
-          }
-        }
-      }
-    }
-
-    private def closeFetchersForQueues(cluster: Cluster,
-                                       messageStreams: Map[String,List[KafkaStream[_,_]]],
-                                       queuesToBeCleared: Iterable[BlockingQueue[FetchedDataChunk]]) {
-      val allPartitionInfos = topicRegistry.values.map(p => p.values).flatten
-      fetcher.foreach { f =>
-        f.stopConnections()
-        clearFetcherQueues(allPartitionInfos, cluster, queuesToBeCleared, messageStreams)
-        /**
-        * here, we need to commit offsets before stopping the consumer from returning any more messages
-        * from the current data chunk. Since partition ownership is not yet released, this commit offsets
-        * call will ensure that the offsets committed now will be used by the next consumer thread owning the partition
-        * for the current data chunk. Since the fetchers are already shutdown and this is the last chunk to be iterated
-        * by the consumer, there will be no more messages returned by this iterator until the rebalancing finishes
-        * successfully and the fetchers restart to fetch more data chunks
-        **/
-        if (config.autoCommitEnable) {
-          info("Committing all offsets after clearing the fetcher queues")
-          commitOffsets(true)
-        }
-      }
-    }
-
-    private def clearFetcherQueues(topicInfos: Iterable[PartitionTopicInfo], cluster: Cluster,
-                                   queuesTobeCleared: Iterable[BlockingQueue[FetchedDataChunk]],
-                                   messageStreams: Map[String,List[KafkaStream[_,_]]]) {
-
-      // Clear all but the currently iterated upon chunk in the consumer thread's queue
-      queuesTobeCleared.foreach(_.clear)
-      info("Cleared all relevant queues for this fetcher")
-
-      // Also clear the currently iterated upon chunk in the consumer threads
-      if(messageStreams != null)
-         messageStreams.foreach(_._2.foreach(s => s.clear()))
-
-      info("Cleared the data chunks in all the consumer message iterators")
-
-    }
-
-    private def closeFetchers(cluster: Cluster, messageStreams: Map[String,List[KafkaStream[_,_]]],
-                              relevantTopicThreadIdsMap: Map[String, Set[ConsumerThreadId]]) {
-      // only clear the fetcher queues for certain topic partitions that *might* no longer be served by this consumer
-      // after this rebalancing attempt
-      val queuesTobeCleared = topicThreadIdAndQueues.filter(q => relevantTopicThreadIdsMap.contains(q._1._1)).map(q => q._2)
-      closeFetchersForQueues(cluster, messageStreams, queuesTobeCleared)
-    }
-
-    private def updateFetcher(cluster: Cluster) {
-      // update partitions for fetcher
-      var allPartitionInfos : List[PartitionTopicInfo] = Nil
-      for (partitionInfos <- topicRegistry.values)
-        for (partition <- partitionInfos.values)
-          allPartitionInfos ::= partition
-      info("Consumer " + consumerIdString + " selected partitions : " +
-        allPartitionInfos.sortWith((s,t) => s.partitionId < t.partitionId).map(_.toString).mkString(","))
-
-      fetcher.foreach(_.startConnections(allPartitionInfos, cluster))
-    }
-
-    private def reflectPartitionOwnershipDecision(partitionAssignment: Map[TopicAndPartition, ConsumerThreadId]): Boolean = {
-      var successfullyOwnedPartitions : List[(String, Int)] = Nil
-      val partitionOwnershipSuccessful = partitionAssignment.map { partitionOwner =>
-        val topic = partitionOwner._1.topic
-        val partition = partitionOwner._1.partition
-        val consumerThreadId = partitionOwner._2
-        val partitionOwnerPath = zkUtils.getConsumerPartitionOwnerPath(group, topic, partition)
-        try {
-          zkUtils.createEphemeralPathExpectConflict(partitionOwnerPath, consumerThreadId.toString)
-          info(consumerThreadId + " successfully owned partition " + partition + " for topic " + topic)
-          successfullyOwnedPartitions ::= (topic, partition)
-          true
-        } catch {
-          case _: ZkNodeExistsException =>
-            // The node hasn't been deleted by the original owner. So wait a bit and retry.
-            info("waiting for the partition ownership to be deleted: " + partition + " for topic " + topic)
-            false
-        }
-      }
-      val hasPartitionOwnershipFailed = partitionOwnershipSuccessful.foldLeft(0)((sum, decision) => sum + (if(decision) 0 else 1))
-      /* even if one of the partition ownership attempt has failed, return false */
-      if(hasPartitionOwnershipFailed > 0) {
-        // remove all paths that we have owned in ZK
-        successfullyOwnedPartitions.foreach(topicAndPartition => deletePartitionOwnershipFromZK(topicAndPartition._1, topicAndPartition._2))
-        false
-      }
-      else true
-    }
-
-    private def addPartitionTopicInfo(currentTopicRegistry: Pool[String, Pool[Int, PartitionTopicInfo]],
-                                      partition: Int, topic: String,
-                                      offset: Long, consumerThreadId: ConsumerThreadId) {
-      val partTopicInfoMap = currentTopicRegistry.getAndMaybePut(topic)
-
-      val queue = topicThreadIdAndQueues.get((topic, consumerThreadId))
-      val consumedOffset = new AtomicLong(offset)
-      val fetchedOffset = new AtomicLong(offset)
-      val partTopicInfo = new PartitionTopicInfo(topic,
-                                                 partition,
-                                                 queue,
-                                                 consumedOffset,
-                                                 fetchedOffset,
-                                                 new AtomicInteger(config.fetchMessageMaxBytes),
-                                                 config.clientId)
-      partTopicInfoMap.put(partition, partTopicInfo)
-      debug(partTopicInfo + " selected new offset " + offset)
-      checkpointedZkOffsets.put(TopicAndPartition(topic, partition), offset)
-    }
-  }
-
-  private def reinitializeConsumer[K,V](
-      topicCount: TopicCount,
-      queuesAndStreams: List[(LinkedBlockingQueue[FetchedDataChunk],KafkaStream[K,V])]) {
-
-    val dirs = new ZKGroupDirs(config.groupId)
-
-    // listener to consumer and partition changes
-    if (loadBalancerListener == null) {
-      val topicStreamsMap = new mutable.HashMap[String,List[KafkaStream[K,V]]]
-      loadBalancerListener = new ZKRebalancerListener(
-        config.groupId, consumerIdString, topicStreamsMap.asInstanceOf[scala.collection.mutable.Map[String, List[KafkaStream[_,_]]]])
-    }
-
-    // create listener for session expired event if not exist yet
-    if (sessionExpirationListener == null)
-      sessionExpirationListener = new ZKSessionExpireListener(
-        dirs, consumerIdString, topicCount, loadBalancerListener)
-
-    // create listener for topic partition change event if not exist yet
-    if (topicPartitionChangeListener == null)
-      topicPartitionChangeListener = new ZKTopicPartitionChangeListener(loadBalancerListener)
-
-    val topicStreamsMap = loadBalancerListener.kafkaMessageAndMetadataStreams
-
-    // map of {topic -> Set(thread-1, thread-2, ...)}
-    val consumerThreadIdsPerTopic: Map[String, Set[ConsumerThreadId]] =
-      topicCount.getConsumerThreadIdsPerTopic
-
-    val allQueuesAndStreams = topicCount match {
-      case _: WildcardTopicCount =>
-        /*
-         * Wild-card consumption streams share the same queues, so we need to
-         * duplicate the list for the subsequent zip operation.
-         */
-        (1 to consumerThreadIdsPerTopic.keySet.size).flatMap(_ => queuesAndStreams).toList
-      case _: StaticTopicCount =>
-        queuesAndStreams
-    }
-
-    val topicThreadIds = consumerThreadIdsPerTopic.map { case (topic, threadIds) =>
-      threadIds.map((topic, _))
-    }.flatten
-
-    require(topicThreadIds.size == allQueuesAndStreams.size,
-      "Mismatch between thread ID count (%d) and queue count (%d)"
-      .format(topicThreadIds.size, allQueuesAndStreams.size))
-    val threadQueueStreamPairs = topicThreadIds.zip(allQueuesAndStreams)
-
-    threadQueueStreamPairs.foreach(e => {
-      val topicThreadId = e._1
-      val q = e._2._1
-      topicThreadIdAndQueues.put(topicThreadId, q)
-      debug("Adding topicThreadId %s and queue %s to topicThreadIdAndQueues data structure".format(topicThreadId, q.toString))
-      newGauge(
-        "FetchQueueSize",
-        new Gauge[Int] {
-          def value = q.size
-        },
-        Map("clientId" -> config.clientId,
-          "topic" -> topicThreadId._1,
-          "threadId" -> topicThreadId._2.threadId.toString)
-      )
-    })
-
-    val groupedByTopic = threadQueueStreamPairs.groupBy(_._1._1)
-    groupedByTopic.foreach(e => {
-      val topic = e._1
-      val streams = e._2.map(_._2._2).toList
-      topicStreamsMap += (topic -> streams)
-      debug("adding topic %s and %d streams to map.".format(topic, streams.size))
-    })
-
-    // listener to consumer and partition changes
-    zkUtils.subscribeStateChanges(sessionExpirationListener)
-
-    zkUtils.subscribeChildChanges(dirs.consumerRegistryDir, loadBalancerListener)
-
-    topicStreamsMap.foreach { topicAndStreams =>
-      // register on broker partition path changes
-      val topicPath = BrokerTopicsPath + "/" + topicAndStreams._1
-      zkUtils.subscribeDataChanges(topicPath, topicPartitionChangeListener)
-    }
-
-    // explicitly trigger load balancing for this consumer
-    loadBalancerListener.syncedRebalance()
-  }
-
-  class WildcardStreamsHandler[K,V](topicFilter: TopicFilter,
-                                  numStreams: Int,
-                                  keyDecoder: Decoder[K],
-                                  valueDecoder: Decoder[V])
-                                extends TopicEventHandler[String] {
-
-    if (messageStreamCreated.getAndSet(true))
-      throw new RuntimeException("Each consumer connector can create " +
-        "message streams by filter at most once.")
-
-    private val wildcardQueuesAndStreams = (1 to numStreams)
-      .map(_ => {
-        val queue = new LinkedBlockingQueue[FetchedDataChunk](config.queuedMaxMessages)
-        val stream = new KafkaStream[K,V](queue,
-                                          config.consumerTimeoutMs,
-                                          keyDecoder,
-                                          valueDecoder,
-                                          config.clientId)
-        (queue, stream)
-    }).toList
-
-     // bootstrap with existing topics
-    private var wildcardTopics =
-      zkUtils.getChildrenParentMayNotExist(BrokerTopicsPath)
-        .filter(topic => topicFilter.isTopicAllowed(topic, config.excludeInternalTopics))
-
-    private val wildcardTopicCount = TopicCount.constructTopicCount(
-      consumerIdString, topicFilter, numStreams, zkUtils, config.excludeInternalTopics)
-
-    val dirs = new ZKGroupDirs(config.groupId)
-    registerConsumerInZK(dirs, consumerIdString, wildcardTopicCount)
-    reinitializeConsumer(wildcardTopicCount, wildcardQueuesAndStreams)
-
-    /*
-     * Topic events will trigger subsequent synced rebalances.
-     */
-    info("Creating topic event watcher for topics " + topicFilter)
-    wildcardTopicWatcher = new ZookeeperTopicEventWatcher(zkUtils, this)
-
-    def handleTopicEvent(allTopics: Seq[String]) {
-      debug("Handling topic event")
-
-      val updatedTopics = allTopics.filter(topic => topicFilter.isTopicAllowed(topic, config.excludeInternalTopics))
-
-      val addedTopics = updatedTopics filterNot (wildcardTopics contains)
-      if (addedTopics.nonEmpty)
-        info("Topic event: added topics = %s"
-                             .format(addedTopics))
-
-      /*
-       * TODO: Deleted topics are interesting (and will not be a concern until
-       * 0.8 release). We may need to remove these topics from the rebalance
-       * listener's map in reinitializeConsumer.
-       */
-      val deletedTopics = wildcardTopics filterNot (updatedTopics contains)
-      if (deletedTopics.nonEmpty)
-        info("Topic event: deleted topics = %s"
-                             .format(deletedTopics))
-
-      wildcardTopics = updatedTopics
-      info("Topics to consume = %s".format(wildcardTopics))
-
-      if (addedTopics.nonEmpty || deletedTopics.nonEmpty)
-        reinitializeConsumer(wildcardTopicCount, wildcardQueuesAndStreams)
-    }
-
-    def streams: Seq[KafkaStream[K,V]] =
-      wildcardQueuesAndStreams.map(_._2)
-  }
-}
diff --git a/core/src/main/scala/kafka/consumer/ZookeeperTopicEventWatcher.scala b/core/src/main/scala/kafka/consumer/ZookeeperTopicEventWatcher.scala
deleted file mode 100644
index 8ce204e3b9114..0000000000000
--- a/core/src/main/scala/kafka/consumer/ZookeeperTopicEventWatcher.scala
+++ /dev/null
@@ -1,103 +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.consumer
-
-import scala.collection.JavaConverters._
-import kafka.utils.{ZkUtils, Logging}
-import org.I0Itec.zkclient.{IZkStateListener, IZkChildListener}
-import org.apache.zookeeper.Watcher.Event.KeeperState
-
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.11.0.0")
-class ZookeeperTopicEventWatcher(val zkUtils: ZkUtils,
-    val eventHandler: TopicEventHandler[String]) extends Logging {
-
-  val lock = new Object()
-
-  startWatchingTopicEvents()
-
-  private def startWatchingTopicEvents() {
-    val topicEventListener = new ZkTopicEventListener()
-    zkUtils.makeSurePersistentPathExists(ZkUtils.BrokerTopicsPath)
-
-    zkUtils.subscribeStateChanges(new ZkSessionExpireListener(topicEventListener))
-
-    val topics = zkUtils.subscribeChildChanges(ZkUtils.BrokerTopicsPath, topicEventListener).getOrElse {
-      throw new AssertionError(s"Expected ${ZkUtils.BrokerTopicsPath} to exist, but it does not. ")
-    }
-
-    // call to bootstrap topic list
-    topicEventListener.handleChildChange(ZkUtils.BrokerTopicsPath, topics.asJava)
-  }
-
-  private def stopWatchingTopicEvents() { zkUtils.unsubscribeAll() }
-
-  def shutdown() {
-    lock.synchronized {
-      info("Shutting down topic event watcher.")
-      if (zkUtils != null) {
-        stopWatchingTopicEvents()
-      }
-      else {
-        warn("Cannot shutdown since the embedded zookeeper client has already closed.")
-      }
-    }
-  }
-
-  class ZkTopicEventListener extends IZkChildListener {
-
-    @throws[Exception]
-    def handleChildChange(parent: String, children: java.util.List[String]) {
-      lock.synchronized {
-        try {
-          if (zkUtils != null) {
-            val latestTopics = zkUtils.getChildren(ZkUtils.BrokerTopicsPath)
-            debug("all topics: %s".format(latestTopics))
-            eventHandler.handleTopicEvent(latestTopics)
-          }
-        }
-        catch {
-          case e: Throwable =>
-            error("error in handling child changes", e)
-        }
-      }
-    }
-
-  }
-
-  class ZkSessionExpireListener(val topicEventListener: ZkTopicEventListener)
-    extends IZkStateListener {
-
-    @throws[Exception]
-    def handleStateChanged(state: KeeperState) { }
-
-    @throws[Exception]
-    def handleNewSession() {
-      lock.synchronized {
-        if (zkUtils != null) {
-          info("ZK expired: resubscribing topic event listener to topic registry")
-          zkUtils.subscribeChildChanges(ZkUtils.BrokerTopicsPath, topicEventListener)
-        }
-      }
-    }
-
-    override def handleSessionEstablishmentError(error: Throwable): Unit = {
-      //no-op ZookeeperConsumerConnector should log error.
-    }
-  }
-}
-
diff --git a/core/src/main/scala/kafka/consumer/package.html b/core/src/main/scala/kafka/consumer/package.html
deleted file mode 100644
index 9e06da8f0c0ef..0000000000000
--- a/core/src/main/scala/kafka/consumer/package.html
+++ /dev/null
@@ -1,19 +0,0 @@
-
-This is the consumer API for kafka.
\ No newline at end of file
diff --git a/core/src/main/scala/kafka/controller/ControllerChannelManager.scala b/core/src/main/scala/kafka/controller/ControllerChannelManager.scala
index a072978f8e3dc..f217776109cad 100755
--- a/core/src/main/scala/kafka/controller/ControllerChannelManager.scala
+++ b/core/src/main/scala/kafka/controller/ControllerChannelManager.scala
@@ -17,55 +17,58 @@
 package kafka.controller
 
 import java.net.SocketTimeoutException
-import java.util.concurrent.{BlockingQueue, LinkedBlockingQueue}
+import java.util.concurrent.{BlockingQueue, LinkedBlockingQueue, TimeUnit}
 
-import com.yammer.metrics.core.Gauge
+import com.yammer.metrics.core.{Gauge, Timer}
 import kafka.api._
 import kafka.cluster.Broker
-import kafka.common.KafkaException
 import kafka.metrics.KafkaMetricsGroup
 import kafka.server.KafkaConfig
 import kafka.utils._
+import kafka.utils.Implicits._
 import org.apache.kafka.clients._
+import org.apache.kafka.common.message.LeaderAndIsrRequestData.LeaderAndIsrPartitionState
+import org.apache.kafka.common.message.StopReplicaRequestData.{StopReplicaPartitionState, StopReplicaTopicState}
 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.{Node, TopicPartition}
-import org.slf4j.event.Level
+import org.apache.kafka.common.{KafkaException, Node, Reconfigurable, TopicPartition}
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 import scala.collection.mutable.HashMap
-import scala.collection.{Set, mutable}
-
+import scala.collection.{Seq, 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 + "]: "
 
-  newGauge(
-    "TotalQueueSize",
-    new Gauge[Int] {
-      def value: Int = brokerLock synchronized {
-        brokerStateInfo.values.iterator.map(_.messageQueue.size).sum
-      }
+  newGauge("TotalQueueSize",
+    () => brokerLock synchronized {
+      brokerStateInfo.values.iterator.map(_.messageQueue.size).sum
     }
   )
 
-  controllerContext.liveBrokers.foreach(addNewBroker)
-
   def startup() = {
+    controllerContext.liveOrShuttingDownBrokers.foreach(addNewBroker)
+
     brokerLock synchronized {
       brokerStateInfo.foreach(brokerState => startRequestSendThread(brokerState._1))
     }
@@ -73,24 +76,24 @@ class ControllerChannelManager(controllerContext: ControllerContext, config: Kaf
 
   def shutdown() = {
     brokerLock synchronized {
-      brokerStateInfo.values.foreach(removeExistingBroker)
+      brokerStateInfo.values.toList.foreach(removeExistingBroker)
     }
   }
 
-  def sendRequest(brokerId: Int, apiKey: ApiKeys, request: AbstractRequest.Builder[_ <: AbstractRequest],
-                  callback: AbstractResponse => Unit = null) {
+  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))
+          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)) {
@@ -100,26 +103,36 @@ 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 brokerNode = broker.getNode(config.interBrokerListenerName)
+    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(
-        config.interBrokerSecurityProtocol,
+        controllerToBrokerSecurityProtocol,
         JaasContext.Type.SERVER,
         config,
-        config.interBrokerListenerName,
+        controllerToBrokerListenerName,
         config.saslMechanismInterBrokerProtocol,
-        config.saslInterBrokerHandshakeRequestEnable
+        time,
+        config.saslInterBrokerHandshakeRequestEnable,
+        logContext
       )
+      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,
@@ -131,7 +144,7 @@ class ControllerChannelManager(controllerContext: ControllerContext, config: Kaf
         channelBuilder,
         logContext
       )
-      new NetworkClient(
+      val networkClient = new NetworkClient(
         selector,
         new ManualMetadataUpdater(Seq(brokerNode).asJava),
         config.brokerId.toString,
@@ -141,60 +154,64 @@ class ControllerChannelManager(controllerContext: ControllerContext, config: Kaf
         Selectable.USE_DEFAULT_BUFFER_SIZE,
         Selectable.USE_DEFAULT_BUFFER_SIZE,
         config.requestTimeoutMs,
+        config.connectionSetupTimeoutMs,
+        config.connectionSetupTimeoutMaxMs,
+        ClientDnsLookup.USE_ALL_DNS_IPS,
         time,
         false,
         new ApiVersions,
         logContext
       )
+      (networkClient, reconfigurableChannelBuilder)
     }
     val threadName = threadNamePrefix match {
       case None => s"Controller-${config.brokerId}-to-broker-${broker.id}-send-thread"
       case Some(name) => s"$name:Controller-${config.brokerId}-to-broker-${broker.id}-send-thread"
     }
 
+    val requestRateAndQueueTimeMetrics = newTimer(
+      RequestRateAndQueueTimeMetricName, TimeUnit.MILLISECONDS, TimeUnit.SECONDS, brokerMetricTags(broker.id)
+    )
+
     val requestThread = new RequestSendThread(config.brokerId, controllerContext, messageQueue, networkClient,
-      brokerNode, config, time, stateChangeLogger, threadName)
+      brokerNode, config, time, requestRateAndQueueTimeMetrics, stateChangeLogger, threadName)
     requestThread.setDaemon(false)
 
-    val queueSizeGauge = newGauge(
-      QueueSizeMetricName,
-      new Gauge[Int] {
-        def value: Int = messageQueue.size
-      },
-      queueSizeTags(broker.id)
-    )
+    val queueSizeGauge = newGauge(QueueSizeMetricName, () => messageQueue.size, brokerMetricTags(broker.id))
 
-    brokerStateInfo.put(broker.id, new ControllerBrokerStateInfo(networkClient, brokerNode, messageQueue,
-      requestThread, queueSizeGauge))
+    brokerStateInfo.put(broker.id, ControllerBrokerStateInfo(networkClient, brokerNode, messageQueue,
+      requestThread, queueSizeGauge, requestRateAndQueueTimeMetrics, reconfigurableChannelBuilder))
   }
 
-  private def queueSizeTags(brokerId: Int) = Map("broker-id" -> brokerId.toString)
+  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()
-      removeMetric(QueueSizeMetricName, queueSizeTags(brokerState.brokerNode.id))
+      removeMetric(QueueSizeMetricName, brokerMetricTags(brokerState.brokerNode.id))
+      removeMetric(RequestRateAndQueueTimeMetricName, brokerMetricTags(brokerState.brokerNode.id))
       brokerStateInfo.remove(brokerState.brokerNode.id)
     } catch {
       case e: Throwable => error("Error while removing broker by the controller", e)
     }
   }
 
-  protected def startRequestSendThread(brokerId: Int) {
+  protected def startRequestSendThread(brokerId: Int): Unit = {
     val requestThread = brokerStateInfo(brokerId).requestSendThread
     if (requestThread.getState == Thread.State.NEW)
       requestThread.start()
   }
 }
 
-case class QueueItem(apiKey: ApiKeys, request: AbstractRequest.Builder[_ <: AbstractRequest],
-                     callback: AbstractResponse => Unit)
+case class QueueItem(apiKey: ApiKeys, request: AbstractControlRequest.Builder[_ <: AbstractControlRequest],
+                     callback: AbstractResponse => Unit, enqueueTimeMs: Long)
 
 class RequestSendThread(val controllerId: Int,
                         val controllerContext: ControllerContext,
@@ -203,21 +220,26 @@ class RequestSendThread(val controllerId: Int,
                         val brokerNode: Node,
                         val config: KafkaConfig,
                         val time: Time,
+                        val requestRateAndQueueTimeMetrics: Timer,
                         val stateChangeLogger: StateChangeLogger,
                         name: String)
   extends ShutdownableThread(name = name) {
 
+  logIdent = s"[RequestSendThread controllerId=$controllerId] "
+
   private val socketTimeoutMs = config.controllerSocketTimeoutMs
 
   override def doWork(): Unit = {
 
-    def backoff(): Unit = CoreUtils.swallow(Thread.sleep(100), this, Level.TRACE)
+    def backoff(): Unit = pause(100, TimeUnit.MILLISECONDS)
+
+    val QueueItem(apiKey, requestBuilder, callback, enqueueTimeMs) = queue.take()
+    requestRateAndQueueTimeMetrics.update(time.milliseconds() - enqueueTimeMs, TimeUnit.MILLISECONDS)
 
-    val QueueItem(apiKey, requestBuilder, callback) = queue.take()
     var clientResponse: ClientResponse = null
     try {
       var isSendSuccessful = false
-      while (isRunning.get() && !isSendSuccessful) {
+      while (isRunning && !isSendSuccessful) {
         // if a broker goes down for a long time, then at some point the controller's zookeeper listener will trigger a
         // removeBroker which will invoke shutdown() on this thread. At that point, we will stop retrying.
         try {
@@ -248,8 +270,9 @@ class RequestSendThread(val controllerId: Int,
 
         val response = clientResponse.responseBody
 
-        stateChangeLogger.withControllerEpoch(controllerContext.epoch).trace("Received response " +
-          s"${response.toString(requestHeader.apiVersion)} for a request sent to broker $brokerNode")
+        stateChangeLogger.withControllerEpoch(controllerContext.epoch).trace(s"Received response " +
+          s"$response for request $api with correlation id " +
+          s"${requestHeader.correlationId} sent to broker $brokerNode")
 
         if (callback != null) {
           callback(response)
@@ -281,17 +304,50 @@ class RequestSendThread(val controllerId: Int,
     }
   }
 
+  override def initiateShutdown(): Boolean = {
+    if (super.initiateShutdown()) {
+      networkClient.initiateClose()
+      true
+    } else
+      false
+  }
+}
+
+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)
+  }
+
 }
 
-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]]
+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, mutable.Map[TopicPartition, StopReplicaPartitionState]]
   val updateMetadataRequestBrokerSet = mutable.Set.empty[Int]
-  val updateMetadataRequestPartitionInfoMap = mutable.Map.empty[TopicPartition, UpdateMetadataRequest.PartitionState]
+  val updateMetadataRequestPartitionInfoMap = mutable.Map.empty[TopicPartition, UpdateMetadataPartitionState]
 
-  def newBatch() {
+  def sendEvent(event: ControllerEvent): Unit
+
+  def sendRequest(brokerId: Int,
+                  request: AbstractControlRequest.Builder[_ <: AbstractControlRequest],
+                  callback: AbstractResponse => Unit = null): Unit
+
+  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 " +
@@ -302,69 +358,89 @@ class ControllerBrokerRequestBatch(controller: KafkaController, stateChangeLogge
     if (updateMetadataRequestBrokerSet.nonEmpty)
       throw new IllegalStateException("Controller to broker state change requests batch is not empty while creating a " +
         s"new one. Some UpdateMetadata state changes to brokers $updateMetadataRequestBrokerSet with partition info " +
-        s"updateMetadataRequestPartitionInfoMap might be lost ")
+        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) {
-    brokerIds.filter(b => b >= 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))
+  def addStopReplicaRequestForBrokers(brokerIds: Seq[Int],
+                                      topicPartition: TopicPartition,
+                                      deletePartition: Boolean): Unit = {
+    // A sentinel (-2) is used as an epoch if the topic is queued for deletion. It overrides
+    // any existing epoch.
+    val leaderEpoch = if (controllerContext.isTopicQueuedUpForDeletion(topicPartition.topic)) {
+      LeaderAndIsr.EpochDuringDelete
+    } else {
+      controllerContext.partitionLeadershipInfo(topicPartition)
+        .map(_.leaderAndIsr.leaderEpoch)
+        .getOrElse(LeaderAndIsr.NoEpoch)
+    }
+
+    brokerIds.filter(_ >= 0).foreach { brokerId =>
+      val result = stopReplicaRequestMap.getOrElseUpdate(brokerId, mutable.Map.empty)
+      val alreadyDelete = result.get(topicPartition).exists(_.deletePartition)
+      result.put(topicPartition, new StopReplicaPartitionState()
+          .setPartitionIndex(topicPartition.partition())
+          .setLeaderEpoch(leaderEpoch)
+          .setDeletePartition(alreadyDelete || 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(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 =>
@@ -372,124 +448,215 @@ class ControllerBrokerRequestBatch(controller: KafkaController, stateChangeLogge
       }
     }
 
-    val filteredPartitions = {
-      val givenPartitions = if (partitions.isEmpty)
-        controllerContext.partitionLeadershipInfo.keySet
-      else
-        partitions
-      if (controller.topicDeletionManager.partitionsToBeDeleted.isEmpty)
-        givenPartitions
-      else
-        givenPartitions -- controller.topicDeletionManager.partitionsToBeDeleted
-    }
-
     updateMetadataRequestBrokerSet ++= brokerIds.filter(_ >= 0)
-    filteredPartitions.foreach(partition => updateMetadataRequestPartitionInfo(partition, beingDeleted = false))
-    controller.topicDeletionManager.partitionsToBeDeleted.foreach(partition => updateMetadataRequestPartitionInfo(partition, beingDeleted = true))
+    partitions.foreach(partition => updateMetadataRequestPartitionInfo(partition,
+      beingDeleted = controllerContext.topicsToBeDeleted.contains(partition.topic)))
   }
 
-  def sendRequestsToBrokers(controllerEpoch: Int) {
-    try {
-      val stateChangeLog = stateChangeLogger.withControllerEpoch(controllerEpoch)
-
-      val leaderAndIsrRequestVersion: Short =
-        if (controller.config.interBrokerProtocolVersion >= KAFKA_1_0_IV0) 1
-        else 0
-
-      leaderAndIsrRequestMap.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")
+  private def sendLeaderAndIsrRequest(controllerEpoch: Int, stateChangeLog: StateChangeLogger): Unit = {
+    val leaderAndIsrRequestVersion: Short =
+      if (config.interBrokerProtocolVersion >= KAFKA_2_4_IV1) 4
+      else if (config.interBrokerProtocolVersion >= KAFKA_2_4_IV0) 3
+      else if (config.interBrokerProtocolVersion >= KAFKA_2_2_IV0) 2
+      else if (config.interBrokerProtocolVersion >= KAFKA_1_0_IV0) 1
+      else 0
+
+    leaderAndIsrRequestMap.forKeyValue { (broker, leaderAndIsrPartitionStates) =>
+      if (controllerContext.liveOrShuttingDownBrokerIds.contains(broker)) {
+        val leaderIds = mutable.Set.empty[Int]
+        var numBecomeLeaders = 0
+        leaderAndIsrPartitionStates.forKeyValue { (topicPartition, state) =>
+          leaderIds += state.leader
+          val typeOfRequest = if (broker == state.leader) {
+            numBecomeLeaders += 1
+            "become-leader"
+          } else {
+            "become-follower"
+          }
+          if (stateChangeLog.isTraceEnabled)
+            stateChangeLog.trace(s"Sending $typeOfRequest LeaderAndIsr request $state to broker $broker for partition $topicPartition")
         }
-        val leaderIds = leaderAndIsrPartitionStates.map(_._2.basePartitionState.leader).toSet
+        stateChangeLog.info(s"Sending LeaderAndIsr request to broker $broker with $numBecomeLeaders become-leader " +
+          s"and ${leaderAndIsrPartitionStates.size - numBecomeLeaders} become-follower partitions")
         val leaders = controllerContext.liveOrShuttingDownBrokers.filter(b => leaderIds.contains(b.id)).map {
-          _.getNode(controller.config.interBrokerListenerName)
+          _.node(config.interBrokerListenerName)
         }
+        val brokerEpoch = controllerContext.liveBrokerIdAndEpochs(broker)
         val leaderAndIsrRequestBuilder = new LeaderAndIsrRequest.Builder(leaderAndIsrRequestVersion, controllerId,
-          controllerEpoch, leaderAndIsrPartitionStates.asJava, leaders.asJava)
-        controller.sendRequest(broker, ApiKeys.LEADER_AND_ISR, leaderAndIsrRequestBuilder,
-          (r: AbstractResponse) => controller.eventManager.put(controller.LeaderAndIsrResponseReceived(r, broker)))
+          controllerEpoch, brokerEpoch, leaderAndIsrPartitionStates.values.toBuffer.asJava, leaders.asJava)
+        sendRequest(broker, leaderAndIsrRequestBuilder, (r: AbstractResponse) => {
+          val leaderAndIsrResponse = r.asInstanceOf[LeaderAndIsrResponse]
+          sendEvent(LeaderAndIsrResponseReceived(leaderAndIsrResponse, broker))
+        })
       }
-      leaderAndIsrRequestMap.clear()
+    }
+    leaderAndIsrRequestMap.clear()
+  }
 
-      updateMetadataRequestPartitionInfoMap.foreach { case (tp, partitionState) =>
-        stateChangeLog.trace(s"Sending UpdateMetadata request $partitionState to brokers $updateMetadataRequestBrokerSet " +
-          s"for partition $tp")
+  private def sendUpdateMetadataRequests(controllerEpoch: Int, stateChangeLog: StateChangeLogger): Unit = {
+    stateChangeLog.info(s"Sending UpdateMetadata request to brokers $updateMetadataRequestBrokerSet " +
+      s"for ${updateMetadataRequestPartitionInfoMap.size} partitions")
+
+    val partitionStates = updateMetadataRequestPartitionInfoMap.values.toBuffer
+    val updateMetadataRequestVersion: Short =
+      if (config.interBrokerProtocolVersion >= KAFKA_2_4_IV1) 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 {
+        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
+
+    updateMetadataRequestBrokerSet.intersect(controllerContext.liveOrShuttingDownBrokerIds).foreach { broker =>
+      val brokerEpoch = controllerContext.liveBrokerIdAndEpochs(broker)
+      val updateMetadataRequestBuilder = new UpdateMetadataRequest.Builder(updateMetadataRequestVersion,
+        controllerId, controllerEpoch, brokerEpoch, partitionStates.asJava, liveBrokers.asJava)
+      sendRequest(broker, updateMetadataRequestBuilder, (r: AbstractResponse) => {
+        val updateMetadataResponse = r.asInstanceOf[UpdateMetadataResponse]
+        sendEvent(UpdateMetadataResponseReceived(updateMetadataResponse, broker))
+      })
+
+    }
+    updateMetadataRequestBrokerSet.clear()
+    updateMetadataRequestPartitionInfoMap.clear()
+  }
+
+  private def sendStopReplicaRequests(controllerEpoch: Int, stateChangeLog: StateChangeLogger): Unit = {
+    val traceEnabled = stateChangeLog.isTraceEnabled
+    val stopReplicaRequestVersion: Short =
+      if (config.interBrokerProtocolVersion >= KAFKA_2_6_IV0) 3
+      else if (config.interBrokerProtocolVersion >= KAFKA_2_4_IV1) 2
+      else if (config.interBrokerProtocolVersion >= KAFKA_2_2_IV0) 1
+      else 0
+
+    def responseCallback(brokerId: Int, isPartitionDeleted: TopicPartition => Boolean)
+                        (response: AbstractResponse): Unit = {
+      val stopReplicaResponse = response.asInstanceOf[StopReplicaResponse]
+      val partitionErrorsForDeletingTopics = mutable.Map.empty[TopicPartition, Errors]
+      stopReplicaResponse.partitionErrors.forEach { pe =>
+        val tp = new TopicPartition(pe.topicName, pe.partitionIndex)
+        if (controllerContext.isTopicDeletionInProgress(pe.topicName) &&
+            isPartitionDeleted(tp)) {
+          partitionErrorsForDeletingTopics += tp -> Errors.forCode(pe.errorCode)
+        }
+      }
+      if (partitionErrorsForDeletingTopics.nonEmpty)
+        sendEvent(TopicDeletionStopReplicaResponseReceived(brokerId, stopReplicaResponse.error,
+          partitionErrorsForDeletingTopics))
+    }
 
-      val partitionStates = Map.empty ++ updateMetadataRequestPartitionInfoMap
-      val updateMetadataRequestVersion: Short =
-        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 updateMetadataRequest = {
-        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.getNode(listenerName)
-            val endPoints = Seq(new EndPoint(node.host, node.port, securityProtocol, listenerName))
-            new UpdateMetadataRequest.Broker(broker.id, endPoints.asJava, broker.rack.orNull)
+    stopReplicaRequestMap.forKeyValue { (brokerId, partitionStates) =>
+      if (controllerContext.liveOrShuttingDownBrokerIds.contains(brokerId)) {
+        if (traceEnabled)
+          partitionStates.forKeyValue { (topicPartition, partitionState) =>
+            stateChangeLog.trace(s"Sending StopReplica request $partitionState to " +
+              s"broker $brokerId for partition $topicPartition")
           }
+
+        val brokerEpoch = controllerContext.liveBrokerIdAndEpochs(brokerId)
+        if (stopReplicaRequestVersion >= 3) {
+          val stopReplicaTopicState = mutable.Map.empty[String, StopReplicaTopicState]
+          partitionStates.forKeyValue { (topicPartition, partitionState) =>
+            val topicState = stopReplicaTopicState.getOrElseUpdate(topicPartition.topic,
+              new StopReplicaTopicState().setTopicName(topicPartition.topic))
+            topicState.partitionStates().add(partitionState)
+          }
+
+          stateChangeLog.info(s"Sending StopReplica request for ${partitionStates.size} " +
+            s"replicas to broker $brokerId")
+          val stopReplicaRequestBuilder = new StopReplicaRequest.Builder(
+            stopReplicaRequestVersion, controllerId, controllerEpoch, brokerEpoch,
+            false, stopReplicaTopicState.values.toBuffer.asJava)
+          sendRequest(brokerId, stopReplicaRequestBuilder,
+            responseCallback(brokerId, tp => partitionStates.get(tp).exists(_.deletePartition)))
         } else {
-          controllerContext.liveOrShuttingDownBrokers.map { broker =>
-            val endPoints = broker.endPoints.map { endPoint =>
-              new UpdateMetadataRequest.EndPoint(endPoint.host, endPoint.port, endPoint.securityProtocol, endPoint.listenerName)
+          var numPartitionStateWithDelete = 0
+          var numPartitionStateWithoutDelete = 0
+          val topicStatesWithDelete = mutable.Map.empty[String, StopReplicaTopicState]
+          val topicStatesWithoutDelete = mutable.Map.empty[String, StopReplicaTopicState]
+
+          partitionStates.forKeyValue { (topicPartition, partitionState) =>
+            val topicStates = if (partitionState.deletePartition()) {
+              numPartitionStateWithDelete += 1
+              topicStatesWithDelete
+            } else {
+              numPartitionStateWithoutDelete += 1
+              topicStatesWithoutDelete
             }
-            new UpdateMetadataRequest.Broker(broker.id, endPoints.asJava, broker.rack.orNull)
+            val topicState = topicStates.getOrElseUpdate(topicPartition.topic,
+              new StopReplicaTopicState().setTopicName(topicPartition.topic))
+            topicState.partitionStates().add(partitionState)
           }
-        }
-        new UpdateMetadataRequest.Builder(updateMetadataRequestVersion, controllerId, controllerEpoch, partitionStates.asJava,
-          liveBrokers.asJava)
-      }
 
-      updateMetadataRequestBrokerSet.foreach { broker =>
-        controller.sendRequest(broker, ApiKeys.UPDATE_METADATA, updateMetadataRequest, null)
-      }
-      updateMetadataRequestBrokerSet.clear()
-      updateMetadataRequestPartitionInfoMap.clear()
-
-      stopReplicaRequestMap.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(",")}")
-
-        val (replicasToGroup, replicasToNotGroup) = replicaInfoList.partition(r => !r.deletePartition && r.callback == null)
-
-        // 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(controllerId, controllerEpoch, false,
-          replicasToGroup.map(_.replica.topicPartition).toSet.asJava)
-        controller.sendRequest(broker, ApiKeys.STOP_REPLICA, stopReplicaRequest)
-
-        replicasToNotGroup.foreach { r =>
-          val stopReplicaRequest = new StopReplicaRequest.Builder(
-              controllerId, controllerEpoch, r.deletePartition,
-              Set(r.replica.topicPartition).asJava)
-          controller.sendRequest(broker, ApiKeys.STOP_REPLICA, stopReplicaRequest, r.callback)
+          if (topicStatesWithDelete.nonEmpty) {
+            stateChangeLog.info(s"Sending StopReplica request (delete = true) for " +
+              s"$numPartitionStateWithDelete replicas to broker $brokerId")
+            val stopReplicaRequestBuilder = new StopReplicaRequest.Builder(
+              stopReplicaRequestVersion, controllerId, controllerEpoch, brokerEpoch,
+              true, topicStatesWithDelete.values.toBuffer.asJava)
+            sendRequest(brokerId, stopReplicaRequestBuilder, responseCallback(brokerId, _ => true))
+          }
+
+          if (topicStatesWithoutDelete.nonEmpty) {
+            stateChangeLog.info(s"Sending StopReplica request (delete = false) for " +
+              s"$numPartitionStateWithoutDelete replicas to broker $brokerId")
+            val stopReplicaRequestBuilder = new StopReplicaRequest.Builder(
+              stopReplicaRequestVersion, controllerId, controllerEpoch, brokerEpoch,
+              false, topicStatesWithoutDelete.values.toBuffer.asJava)
+            sendRequest(brokerId, stopReplicaRequestBuilder)
+          }
         }
       }
-      stopReplicaRequestMap.clear()
+    }
+
+    stopReplicaRequestMap.clear()
+  }
+
+  def sendRequestsToBrokers(controllerEpoch: Int): Unit = {
+    try {
+      val stateChangeLog = stateChangeLogger.withControllerEpoch(controllerEpoch)
+      sendLeaderAndIsrRequest(controllerEpoch, stateChangeLog)
+      sendUpdateMetadataRequests(controllerEpoch, stateChangeLog)
+      sendStopReplicaRequests(controllerEpoch, stateChangeLog)
     } catch {
       case e: Throwable =>
         if (leaderAndIsrRequestMap.nonEmpty) {
           error("Haven't been able to send leader and isr requests, current state of " +
-              s"the map is $leaderAndIsrRequestMap. Exception message: $e")
+            s"the map is $leaderAndIsrRequestMap. Exception message: $e")
         }
         if (updateMetadataRequestBrokerSet.nonEmpty) {
           error(s"Haven't been able to send metadata update requests to brokers $updateMetadataRequestBrokerSet, " +
-                s"current state of the partition info is $updateMetadataRequestPartitionInfoMap. Exception message: $e")
+            s"current state of the partition info is $updateMetadataRequestPartitionInfoMap. Exception message: $e")
         }
         if (stopReplicaRequestMap.nonEmpty) {
           error("Haven't been able to send stop replica requests, current state of " +
-              s"the map is $stopReplicaRequestMap. Exception message: $e")
+            s"the map is $stopReplicaRequestMap. Exception message: $e")
         }
         throw new IllegalStateException(e)
     }
@@ -500,8 +667,7 @@ case class ControllerBrokerStateInfo(networkClient: NetworkClient,
                                      brokerNode: Node,
                                      messageQueue: BlockingQueue[QueueItem],
                                      requestSendThread: RequestSendThread,
-                                     queueSizeGauge: Gauge[Int])
-
-case class StopReplicaRequestInfo(replica: PartitionAndReplica, deletePartition: Boolean, callback: AbstractResponse => Unit)
+                                     queueSizeGauge: Gauge[Int],
+                                     requestRateAndTimeMetrics: Timer,
+                                     reconfigurableChannelBuilder: Option[Reconfigurable])
 
-class Callbacks(val stopReplicaResponseCallback: (AbstractResponse, Int) => Unit = (_, _ ) => ())
diff --git a/core/src/main/scala/kafka/controller/ControllerContext.scala b/core/src/main/scala/kafka/controller/ControllerContext.scala
index 541bce82d4870..042830118223f 100644
--- a/core/src/main/scala/kafka/controller/ControllerContext.scala
+++ b/core/src/main/scala/kafka/controller/ControllerContext.scala
@@ -18,43 +18,215 @@
 package kafka.controller
 
 import kafka.cluster.Broker
-import org.apache.kafka.common.TopicPartition
+import kafka.utils.Implicits._
+import org.apache.kafka.common.{TopicPartition, Uuid}
 
-import scala.collection.{Seq, Set, mutable}
+import scala.collection.{Map, Seq, Set, mutable}
+
+object ReplicaAssignment {
+  def apply(replicas: Seq[Int]): ReplicaAssignment = {
+    apply(replicas, Seq.empty, Seq.empty)
+  }
+
+  val empty: ReplicaAssignment = apply(Seq.empty)
+}
+
+
+/**
+ * @param replicas the sequence of brokers assigned to the partition. It includes the set of brokers
+ *                 that were added (`addingReplicas`) and removed (`removingReplicas`).
+ * @param addingReplicas the replicas that are being added if there is a pending reassignment
+ * @param removingReplicas the replicas that are being removed if there is a pending reassignment
+ */
+case class ReplicaAssignment private (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(target: Seq[Int]): ReplicaAssignment = {
+    val fullReplicaSet = (target ++ originReplicas).distinct
+    ReplicaAssignment(
+      fullReplicaSet,
+      fullReplicaSet.diff(originReplicas),
+      fullReplicaSet.diff(target)
+    )
+  }
+
+  def removeReplica(replica: Int): ReplicaAssignment = {
+    ReplicaAssignment(
+      replicas.filterNot(_ == replica),
+      addingReplicas.filterNot(_ == replica),
+      removingReplicas.filterNot(_ == replica)
+    )
+  }
+
+  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 preferredReplicaImbalanceCount = 0
+  val shuttingDownBrokerIds = mutable.Set.empty[Int]
+  private val liveBrokers = mutable.Set.empty[Broker]
+  private val liveBrokerEpochs = mutable.Map.empty[Int, Long]
+  var epoch: Int = KafkaController.InitialControllerEpoch
+  var epochZkVersion: Int = KafkaController.InitialControllerEpochZkVersion
 
-  var controllerChannelManager: ControllerChannelManager = null
+  val allTopics = mutable.Set.empty[String]
+  var topicIds = mutable.Map.empty[String, Uuid]
+  var topicNames = mutable.Map.empty[Uuid, String]
+  val partitionAssignments = mutable.Map.empty[String, mutable.Map[Int, ReplicaAssignment]]
+  private 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.empty[Int, Set[TopicPartition]]
 
-  var shuttingDownBrokerIds: mutable.Set[Int] = mutable.Set.empty
-  var epoch: Int = KafkaController.InitialControllerEpoch - 1
-  var epochZkVersion: Int = KafkaController.InitialControllerEpochZkVersion - 1
-  var allTopics: Set[String] = Set.empty
-  var partitionReplicaAssignment: mutable.Map[TopicPartition, Seq[Int]] = mutable.Map.empty
-  var partitionLeadershipInfo: mutable.Map[TopicPartition, LeaderIsrAndControllerEpoch] = mutable.Map.empty
-  val partitionsBeingReassigned: mutable.Map[TopicPartition, ReassignedPartitionsContext] = mutable.Map.empty
-  val replicasOnOfflineDirs: mutable.Map[Int, Set[TopicPartition]] = mutable.Map.empty
+  val topicsToBeDeleted = mutable.Set.empty[String]
 
-  private var liveBrokersUnderlying: Set[Broker] = Set.empty
-  private var liveBrokerIdsUnderlying: Set[Int] = Set.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]
 
-  // setter
-  def liveBrokers_=(brokers: Set[Broker]) {
-    liveBrokersUnderlying = brokers
-    liveBrokerIdsUnderlying = liveBrokersUnderlying.map(_.id)
+  private def clearTopicsState(): Unit = {
+    allTopics.clear()
+    topicIds.clear()
+    topicNames.clear()
+    partitionAssignments.clear()
+    partitionLeadershipInfo.clear()
+    partitionsBeingReassigned.clear()
+    replicasOnOfflineDirs.clear()
+    partitionStates.clear()
+    offlinePartitionCount = 0
+    preferredReplicaImbalanceCount = 0
+    replicaStates.clear()
   }
 
-  // getter
-  def liveBrokers = liveBrokersUnderlying.filter(broker => !shuttingDownBrokerIds.contains(broker.id))
-  def liveBrokerIds = liveBrokerIdsUnderlying -- shuttingDownBrokerIds
+  def addTopicId(topic: String, id: Uuid): Unit = {
+    if (!allTopics.contains(topic))
+      throw new IllegalStateException(s"topic $topic is not contained in all topics.")
 
-  def liveOrShuttingDownBrokerIds = liveBrokerIdsUnderlying
-  def liveOrShuttingDownBrokers = liveBrokersUnderlying
+    topicIds.get(topic).foreach { existingId =>
+      if (!existingId.equals(id))
+        throw new IllegalStateException(s"topic ID map already contained ID for topic " +
+          s"$topic and new ID $id did not match existing ID $existingId")
+    }
+    topicNames.get(id).foreach { existingName =>
+      if (!existingName.equals(topic))
+        throw new IllegalStateException(s"topic name map already contained ID " +
+          s"$id and new name $topic did not match existing name $existingName")
+    }
+    topicIds.put(topic, id)
+    topicNames.put(id, topic)
+  }
+
+  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)
+      .getOrElse(topicPartition.partition, ReplicaAssignment.empty)
+  }
+
+  def updatePartitionFullReplicaAssignment(topicPartition: TopicPartition, newAssignment: ReplicaAssignment): Unit = {
+    val assignments = partitionAssignments.getOrElseUpdate(topicPartition.topic, mutable.Map.empty)
+    val previous = assignments.put(topicPartition.partition, newAssignment)
+    val leadershipInfo = partitionLeadershipInfo.get(topicPartition)
+    updatePreferredReplicaImbalanceMetric(topicPartition, previous, leadershipInfo,
+      Some(newAssignment), leadershipInfo)
+  }
+
+  def partitionReplicaAssignmentForTopic(topic : String): Map[TopicPartition, Seq[Int]] = {
+    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] = {
+    partitionAssignments.flatMap {
+      case (topic, topicReplicaAssignment) => topicReplicaAssignment.map {
+        case (partition, _) => new TopicPartition(topic, partition)
+      }
+    }.toSet
+  }
+
+  def setLiveBrokers(brokerAndEpochs: Map[Broker, Long]): Unit = {
+    clearLiveBrokers()
+    addLiveBrokers(brokerAndEpochs)
+  }
+
+  private def clearLiveBrokers(): Unit = {
+    liveBrokers.clear()
+    liveBrokerEpochs.clear()
+  }
+
+  def addLiveBrokers(brokerAndEpochs: Map[Broker, Long]): Unit = {
+    liveBrokers ++= brokerAndEpochs.keySet
+    liveBrokerEpochs ++= brokerAndEpochs.map { case (broker, brokerEpoch) => (broker.id, brokerEpoch) }
+  }
+
+  def removeLiveBrokers(brokerIds: Set[Int]): Unit = {
+    liveBrokers --= liveBrokers.filter(broker => brokerIds.contains(broker.id))
+    liveBrokerEpochs --= brokerIds
+  }
+
+  def updateBrokerMetadata(oldMetadata: Broker, newMetadata: Broker): Unit = {
+    liveBrokers -= oldMetadata
+    liveBrokers += newMetadata
+  }
+
+  // getter
+  def liveBrokerIds: Set[Int] = liveBrokerEpochs.keySet.diff(shuttingDownBrokerIds)
+  def liveOrShuttingDownBrokerIds: Set[Int] = liveBrokerEpochs.keySet
+  def liveOrShuttingDownBrokers: Set[Broker] = liveBrokers
+  def liveBrokerIdAndEpochs: Map[Int, Long] = liveBrokerEpochs
+  def liveOrShuttingDownBroker(brokerId: Int): Option[Broker] = liveOrShuttingDownBrokers.find(_.id == brokerId)
 
   def partitionsOnBroker(brokerId: Int): Set[TopicPartition] = {
-    partitionReplicaAssignment.collect {
-      case (topicPartition, replicas) if replicas.contains(brokerId) => topicPartition
+    partitionAssignments.flatMap {
+      case (topic, topicReplicaAssignment) => topicReplicaAssignment.filter {
+        case (_, partitionAssignment) => partitionAssignment.replicas.contains(brokerId)
+      }.map {
+        case (partition, _) => new TopicPartition(topic, partition)
+      }
     }.toSet
   }
 
@@ -68,27 +240,49 @@ class ControllerContext {
 
   def replicasOnBrokers(brokerIds: Set[Int]): Set[PartitionAndReplica] = {
     brokerIds.flatMap { brokerId =>
-      partitionReplicaAssignment.collect { case (topicPartition, replicas) if replicas.contains(brokerId) =>
-        PartitionAndReplica(topicPartition, brokerId)
+      partitionAssignments.flatMap {
+        case (topic, topicReplicaAssignment) => topicReplicaAssignment.collect {
+          case (partition, partitionAssignment) if partitionAssignment.replicas.contains(brokerId) =>
+            PartitionAndReplica(new TopicPartition(topic, partition), brokerId)
+        }
       }
-    }.toSet
+    }
   }
 
   def replicasForTopic(topic: String): Set[PartitionAndReplica] = {
-    partitionReplicaAssignment
-      .filter { case (topicPartition, _) => topicPartition.topic == topic }
-      .flatMap { case (topicPartition, replicas) =>
-        replicas.map(PartitionAndReplica(topicPartition, _))
-      }.toSet
+    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] =
-    partitionReplicaAssignment.keySet.filter(topicPartition => topicPartition.topic == topic)
+  def partitionsForTopic(topic: String): collection.Set[TopicPartition] = {
+    partitionAssignments.getOrElse(topic, mutable.Map.empty).map {
+      case (partition, _) => new TopicPartition(topic, partition)
+    }.toSet
+  }
 
-  def allLiveReplicas(): Set[PartitionAndReplica] = {
-    replicasOnBrokers(liveBrokerIds).filter { partitionAndReplica =>
-      isReplicaOnline(partitionAndReplica.replica, partitionAndReplica.topicPartition)
+  /**
+    * 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] = {
@@ -98,10 +292,226 @@ class ControllerContext {
     }
   }
 
-  def removeTopic(topic: String) = {
-    partitionLeadershipInfo = partitionLeadershipInfo.filter { case (topicPartition, _) => topicPartition.topic != topic }
-    partitionReplicaAssignment = partitionReplicaAssignment.filter { case (topicPartition, _) => topicPartition.topic != topic }
+  def resetContext(): Unit = {
+    topicsToBeDeleted.clear()
+    topicsWithDeletionStarted.clear()
+    topicsIneligibleForDeletion.clear()
+    shuttingDownBrokerIds.clear()
+    epoch = 0
+    epochZkVersion = 0
+    clearTopicsState()
+    clearLiveBrokers()
+  }
+
+  def setAllTopics(topics: Set[String]): Unit = {
+    allTopics.clear()
+    allTopics ++= topics
+  }
+
+  def removeTopic(topic: String): Unit = {
+    // Metric is cleaned when the topic is queued up for deletion so
+    // we don't clean it twice. We clean it only if it is deleted
+    // directly.
+    if (!topicsToBeDeleted.contains(topic))
+      cleanPreferredReplicaImbalanceMetric(topic)
+    topicsToBeDeleted -= topic
+    topicsWithDeletionStarted -= topic
     allTopics -= topic
+    topicIds.remove(topic).foreach { topicId =>
+      topicNames.remove(topicId)
+    }
+    partitionAssignments.remove(topic).foreach { assignments =>
+      assignments.keys.foreach { partition =>
+        partitionLeadershipInfo.remove(new TopicPartition(topic, partition))
+      }
+    }
+  }
+
+  def queueTopicDeletion(topics: Set[String]): Unit = {
+    topicsToBeDeleted ++= topics
+    topics.foreach(cleanPreferredReplicaImbalanceMetric)
+  }
+
+  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)
+  }
+
+  private def updatePartitionStateMetrics(partition: TopicPartition,
+                                          currentState: PartitionState,
+                                          targetState: PartitionState): Unit = {
+    if (!isTopicDeletionInProgress(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
+  }
+
+  def putPartitionLeadershipInfo(partition: TopicPartition,
+                                 leaderIsrAndControllerEpoch: LeaderIsrAndControllerEpoch): Unit = {
+    val previous = partitionLeadershipInfo.put(partition, leaderIsrAndControllerEpoch)
+    val replicaAssignment = partitionFullReplicaAssignment(partition)
+    updatePreferredReplicaImbalanceMetric(partition, Some(replicaAssignment), previous,
+      Some(replicaAssignment), Some(leaderIsrAndControllerEpoch))
+  }
+
+  def partitionLeadershipInfo(partition: TopicPartition): Option[LeaderIsrAndControllerEpoch] = {
+    partitionLeadershipInfo.get(partition)
+  }
+
+  def partitionsLeadershipInfo: Map[TopicPartition, LeaderIsrAndControllerEpoch] =
+    partitionLeadershipInfo
+
+  def partitionsWithLeaders: Set[TopicPartition] =
+    partitionLeadershipInfo.keySet.filter(tp => !isTopicQueuedUpForDeletion(tp.topic))
+
+  def partitionsWithOfflineLeader: Set[TopicPartition] = {
+    partitionLeadershipInfo.filter { case (topicPartition, leaderIsrAndControllerEpoch) =>
+      !isReplicaOnline(leaderIsrAndControllerEpoch.leaderAndIsr.leader, topicPartition) &&
+        !isTopicQueuedUpForDeletion(topicPartition.topic)
+    }.keySet
+  }
+
+  def partitionLeadersOnBroker(brokerId: Int): Set[TopicPartition] = {
+    partitionLeadershipInfo.filter { case (topicPartition, leaderIsrAndControllerEpoch) =>
+      !isTopicQueuedUpForDeletion(topicPartition.topic) &&
+        leaderIsrAndControllerEpoch.leaderAndIsr.leader == brokerId &&
+        partitionReplicaAssignment(topicPartition).size > 1
+    }.keySet
+  }
+
+  def clearPartitionLeadershipInfo(): Unit = partitionLeadershipInfo.clear()
+
+  def partitionWithLeadersCount: Int = partitionLeadershipInfo.size
+
+  private def updatePreferredReplicaImbalanceMetric(partition: TopicPartition,
+                                                    oldReplicaAssignment: Option[ReplicaAssignment],
+                                                    oldLeadershipInfo: Option[LeaderIsrAndControllerEpoch],
+                                                    newReplicaAssignment: Option[ReplicaAssignment],
+                                                    newLeadershipInfo: Option[LeaderIsrAndControllerEpoch]): Unit = {
+    if (!isTopicQueuedUpForDeletion(partition.topic)) {
+      oldReplicaAssignment.foreach { replicaAssignment =>
+        oldLeadershipInfo.foreach { leadershipInfo =>
+          if (!hasPreferredLeader(replicaAssignment, leadershipInfo))
+            preferredReplicaImbalanceCount -= 1
+        }
+      }
+
+      newReplicaAssignment.foreach { replicaAssignment =>
+        newLeadershipInfo.foreach { leadershipInfo =>
+          if (!hasPreferredLeader(replicaAssignment, leadershipInfo))
+            preferredReplicaImbalanceCount += 1
+        }
+      }
+    }
+  }
+
+  private def cleanPreferredReplicaImbalanceMetric(topic: String): Unit = {
+    partitionAssignments.getOrElse(topic, mutable.Map.empty).forKeyValue { (partition, replicaAssignment) =>
+      partitionLeadershipInfo.get(new TopicPartition(topic, partition)).foreach { leadershipInfo =>
+        if (!hasPreferredLeader(replicaAssignment, leadershipInfo))
+          preferredReplicaImbalanceCount -= 1
+      }
+    }
+  }
+
+  private def hasPreferredLeader(replicaAssignment: ReplicaAssignment,
+                                 leadershipInfo: LeaderIsrAndControllerEpoch): Boolean = {
+    val preferredReplica = replicaAssignment.replicas.head
+    if (replicaAssignment.isBeingReassigned && replicaAssignment.addingReplicas.contains(preferredReplica))
+      // reassigning partitions are not counted as imbalanced until the new replica joins the ISR (completes reassignment)
+      !leadershipInfo.leaderAndIsr.isr.contains(preferredReplica)
+    else
+      leadershipInfo.leaderAndIsr.leader == preferredReplica
+  }
+
+  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 396a39dd76bff..b5ae3ff086e9c 100644
--- a/core/src/main/scala/kafka/controller/ControllerEventManager.scala
+++ b/core/src/main/scala/kafka/controller/ControllerEventManager.scala
@@ -17,60 +17,143 @@
 
 package kafka.controller
 
-import java.util.concurrent.LinkedBlockingQueue
+import java.util.ArrayList
+import java.util.concurrent.atomic.AtomicBoolean
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeUnit}
 import java.util.concurrent.locks.ReentrantLock
 
-import kafka.metrics.KafkaTimer
+import kafka.metrics.{KafkaMetricsGroup, KafkaTimer}
 import kafka.utils.CoreUtils.inLock
 import kafka.utils.ShutdownableThread
+import org.apache.kafka.common.utils.Time
 
 import scala.collection._
 
 object ControllerEventManager {
   val ControllerEventThreadName = "controller-event-thread"
+  val EventQueueTimeMetricName = "EventQueueTimeMs"
+  val EventQueueSizeMetricName = "EventQueueSize"
 }
-class ControllerEventManager(rateAndTimeMetrics: Map[ControllerState, KafkaTimer],
-                             eventProcessedListener: ControllerEvent => Unit) {
+
+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],
+                             eventQueueTimeTimeoutMs: Long = 300000) extends KafkaMetricsGroup {
+  import ControllerEventManager._
 
   @volatile private var _state: ControllerState = ControllerState.Idle
   private val putLock = new ReentrantLock()
-  private val queue = new LinkedBlockingQueue[ControllerEvent]
-  private val thread = new ControllerEventThread(ControllerEventManager.ControllerEventThreadName)
+  private val queue = new LinkedBlockingQueue[QueuedEvent]
+  // Visible for test
+  private[controller] var thread = new ControllerEventThread(ControllerEventThreadName)
+
+  private val eventQueueTimeHist = newHistogram(EventQueueTimeMetricName)
+
+  newGauge(EventQueueSizeMetricName, () => queue.size)
 
   def state: ControllerState = _state
 
   def start(): Unit = thread.start()
 
-  def close(): Unit = thread.shutdown()
+  def close(): Unit = {
+    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.clear()
-    queue.put(event)
+  def clearAndPut(event: ControllerEvent): QueuedEvent = inLock(putLock){
+    val preemptedEvents = new ArrayList[QueuedEvent]()
+    queue.drainTo(preemptedEvents)
+    preemptedEvents.forEach(_.preempt(processor))
+    put(event)
   }
 
-  class ControllerEventThread(name: String) extends ShutdownableThread(name = name) {
+  def isEmpty: Boolean = queue.isEmpty
+
+  class ControllerEventThread(name: String) extends ShutdownableThread(name = name, isInterruptible = false) {
+    logIdent = s"[ControllerEventThread controllerId=$controllerId] "
+
     override def doWork(): Unit = {
-      val controllerEvent = queue.take()
-      _state = controllerEvent.state
-
-      try {
-        rateAndTimeMetrics(state).time {
-          controllerEvent.process()
-        }
-      } catch {
-        case e: Throwable => error(s"Error processing event $controllerEvent", e)
+      val dequeued = pollFromEventQueue()
+      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() - dequeued.enqueueTimeMs)
+
+          try {
+            def process(): Unit = dequeued.process(processor)
+
+            rateAndTimeMetrics.get(state) match {
+              case Some(timer) => timer.time { process() }
+              case None => process()
+            }
+          } catch {
+            case e: Throwable => error(s"Uncaught error processing event $controllerEvent", e)
+          }
+
+          _state = ControllerState.Idle
       }
+    }
+  }
 
-      try eventProcessedListener(controllerEvent)
-      catch {
-        case e: Throwable => error(s"Error while invoking listener for processed event $controllerEvent", e)
+  private def pollFromEventQueue(): QueuedEvent = {
+    val count = eventQueueTimeHist.count()
+    if (count != 0) {
+      val event  = queue.poll(eventQueueTimeTimeoutMs, TimeUnit.MILLISECONDS)
+      if (event == null) {
+        eventQueueTimeHist.clear()
+        queue.take()
+      } else {
+        event
       }
-
-      _state = ControllerState.Idle
+    } else {
+      queue.take()
     }
   }
 
diff --git a/core/src/main/scala/kafka/controller/ControllerState.scala b/core/src/main/scala/kafka/controller/ControllerState.scala
index 2bb63e8c0f197..f84240536dd85 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 {
@@ -86,7 +88,35 @@ object ControllerState {
     def value = 11
   }
 
+  case object ControllerShutdown extends ControllerState {
+    def value = 12
+  }
+
+  case object UncleanLeaderElectionEnable extends ControllerState {
+    def value = 13
+  }
+
+  case object TopicUncleanLeaderElectionEnable extends ControllerState {
+    def value = 14
+  }
+
+  case object ListPartitionReassignment extends ControllerState {
+    def value = 15
+  }
+
+  case object UpdateMetadataResponseReceived extends ControllerState {
+    def value = 16
+
+    override protected def hasRateAndTimeMetric: Boolean = false
+  }
+
+  case object UpdateFeatures extends ControllerState {
+    def value = 17
+  }
+
   val values: Seq[ControllerState] = Seq(Idle, ControllerChange, BrokerChange, TopicChange, TopicDeletion,
-    PartitionReassignment, AutoLeaderBalance, ManualLeaderBalance, ControlledShutdown, IsrChange, LeaderAndIsrResponseReceived,
-    LogDirChange)
+    AlterPartitionReassignment, AutoLeaderBalance, ManualLeaderBalance, ControlledShutdown, IsrChange,
+    LeaderAndIsrResponseReceived, LogDirChange, ControllerShutdown, UncleanLeaderElectionEnable,
+    TopicUncleanLeaderElectionEnable, ListPartitionReassignment, UpdateMetadataResponseReceived,
+    UpdateFeatures)
 }
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..dffa88841aac3
--- /dev/null
+++ b/core/src/main/scala/kafka/controller/Election.scala
@@ -0,0 +1,157 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package kafka.controller
+
+import kafka.api.LeaderAndIsr
+import org.apache.kafka.common.TopicPartition
+
+import scala.collection.Seq
+
+case class ElectionResult(topicPartition: TopicPartition, leaderAndIsr: Option[LeaderAndIsr], liveReplicas: Seq[Int])
+
+object Election {
+
+  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, controllerContext)
+        val newLeaderAndIsrOpt = leaderOpt.map { 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 3615b7ddb3965..fe14d427487c2 100644
--- a/core/src/main/scala/kafka/controller/KafkaController.scala
+++ b/core/src/main/scala/kafka/controller/KafkaController.scala
@@ -16,117 +16,140 @@
  */
 package kafka.controller
 
-import java.util.concurrent.{CountDownLatch, TimeUnit}
+import java.util
+import java.util.concurrent.TimeUnit
 
-import com.yammer.metrics.core.Gauge
 import kafka.admin.AdminOperationException
 import kafka.api._
 import kafka.common._
+import kafka.controller.KafkaController.AlterIsrCallback
+import kafka.cluster.Broker
+import kafka.controller.KafkaController.{AlterReassignmentsCallback, ElectLeadersCallback, ListReassignmentsCallback, UpdateFeaturesCallback}
 import kafka.metrics.{KafkaMetricsGroup, KafkaTimer}
 import kafka.server._
 import kafka.utils._
+import kafka.utils.Implicits._
 import kafka.zk.KafkaZkClient.UpdateLeaderAndIsrResult
-import kafka.zk._
+import kafka.zk.TopicZNode.TopicIdReplicaAssignment
+import kafka.zk.{FeatureZNodeStatus, _}
 import kafka.zookeeper.{StateChangeHandler, ZNodeChangeHandler, ZNodeChildChangeHandler}
+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}
+import org.apache.kafka.common.errors.{BrokerNotAvailableException, ControllerMovedException, StaleBrokerEpochException}
+import org.apache.kafka.common.message.{AlterIsrRequestData, AlterIsrResponseData}
+import org.apache.kafka.common.feature.{Features, FinalizedVersionRange}
+import org.apache.kafka.common.message.UpdateFeaturesRequestData
 import org.apache.kafka.common.metrics.Metrics
-import org.apache.kafka.common.protocol.{ApiKeys, Errors}
-import org.apache.kafka.common.requests.{AbstractRequest, AbstractResponse, LeaderAndIsrResponse, StopReplicaResponse}
-import org.apache.kafka.common.utils.Time
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.requests.{AbstractControlRequest, ApiError, LeaderAndIsrResponse, UpdateFeaturesRequest, UpdateMetadataResponse}
+import org.apache.kafka.common.utils.{Time, Utils}
 import org.apache.zookeeper.KeeperException
-import org.apache.zookeeper.KeeperException.{Code, NodeExistsException}
+import org.apache.zookeeper.KeeperException.Code
 
-import scala.collection._
-import scala.util.Try
+import scala.collection.{Map, Seq, Set, immutable, mutable}
+import scala.collection.mutable.ArrayBuffer
+import scala.jdk.CollectionConverters._
+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 = 1
-  val InitialControllerEpochZkVersion = 1
+  val InitialControllerEpoch = 0
+  val InitialControllerEpochZkVersion = 0
+
+  type ElectLeadersCallback = Map[TopicPartition, Either[ApiError, Int]] => Unit
+  type ListReassignmentsCallback = Either[Map[TopicPartition, ReplicaAssignment], ApiError] => Unit
+  type AlterReassignmentsCallback = Either[Map[TopicPartition, ApiError], ApiError] => Unit
+  type AlterIsrCallback = Either[Map[TopicPartition, Either[Errors, LeaderAndIsr]], Errors] => Unit
+  type UpdateFeaturesCallback = Either[ApiError, Map[String, ApiError]] => Unit
 }
 
-class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Time, metrics: Metrics, 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,
+                      brokerFeatures: BrokerFeatures,
+                      featureCache: FinalizedFeatureCache,
+                      threadNamePrefix: Option[String] = None)
+  extends ControllerEventProcessor with Logging with KafkaMetricsGroup {
+
   this.logIdent = s"[Controller id=${config.brokerId}] "
 
+  @volatile private var brokerInfo = initialBrokerInfo
+  @volatile private var _brokerEpoch = initialBrokerEpoch
+
+  private val isAlterIsrEnabled = config.interBrokerProtocolVersion.isAlterIsrSupported
   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(controllerContext.stats.rateAndTimeMetrics,
-    _ => updateMetrics())
-
-  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, topicDeletionManager, zkClient, mutable.Map.empty, new ControllerBrokerRequestBatch(this, stateChangeLogger))
-
-  private val controllerChangeHandler = new ControllerChangeHandler(this, eventManager)
-  private val brokerChangeHandler = new BrokerChangeHandler(this, eventManager)
-  private val topicChangeHandler = new TopicChangeHandler(this, eventManager)
-  private val topicDeletionHandler = new TopicDeletionHandler(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 brokerModificationsHandlers: mutable.Map[Int, BrokerModificationsHandler] = mutable.Map.empty
+  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)
 
   @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
-
-  newGauge(
-    "ActiveControllerCount",
-    new Gauge[Int] {
-      def value = if (isActive) 1 else 0
-    }
-  )
-
-  newGauge(
-    "OfflinePartitionsCount",
-    new Gauge[Int] {
-      def value: Int = offlinePartitionCount
-    }
-  )
-
-  newGauge(
-    "PreferredReplicaImbalanceCount",
-    new Gauge[Int] {
-      def value: Int = preferredReplicaImbalanceCount
-    }
-  )
-
-  newGauge(
-    "ControllerState",
-    new Gauge[Byte] {
-      def value: Byte = state.value
-    }
-  )
-
-  newGauge(
-    "GlobalTopicCount",
-    new Gauge[Int] {
-      def value: Int = globalTopicCount
-    }
-  )
-
-  newGauge(
-    "GlobalPartitionCount",
-    new Gauge[Int] {
-      def value: Int = globalPartitionCount
-    }
-  )
+  @volatile private var topicsToDeleteCount = 0
+  @volatile private var replicasToDeleteCount = 0
+  @volatile private var ineligibleTopicsToDeleteCount = 0
+  @volatile private var ineligibleReplicasToDeleteCount = 0
+
+  /* single-thread scheduler to clean expired tokens */
+  private val tokenCleanScheduler = new KafkaScheduler(threads = 1, threadNamePrefix = "delegation-token-cleaner")
+
+  newGauge("ActiveControllerCount", () => if (isActive) 1 else 0)
+  newGauge("OfflinePartitionsCount", () => offlinePartitionCount)
+  newGauge("PreferredReplicaImbalanceCount", () => preferredReplicaImbalanceCount)
+  newGauge("ControllerState", () => state.value)
+  newGauge("GlobalTopicCount", () => globalTopicCount)
+  newGauge("GlobalPartitionCount", () => globalPartitionCount)
+  newGauge("TopicsToDeleteCount", () => topicsToDeleteCount)
+  newGauge("ReplicasToDeleteCount", () => replicasToDeleteCount)
+  newGauge("TopicsIneligibleToDeleteCount", () => ineligibleTopicsToDeleteCount)
+  newGauge("ReplicasIneligibleToDeleteCount", () => ineligibleReplicasToDeleteCount)
 
   /**
    * Returns true if this broker is the current controller.
    */
   def isActive: Boolean = activeControllerId == config.brokerId
 
+  def brokerEpoch: Long = _brokerEpoch
+
   def epoch: Int = controllerContext.epoch
 
   /**
@@ -137,14 +160,15 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
   def startup() = {
     zkClient.registerStateChangeHandler(new StateChangeHandler {
       override val name: String = StateChangeHandlers.ControllerHandler
-      override def onReconnectionTimeout(): Unit = error("Reconnection timeout.")
       override def afterInitializingSession(): Unit = {
-        eventManager.put(Reelect)
+        eventManager.put(RegisterBrokerAndReelect)
       }
       override def beforeInitializingSession(): Unit = {
-        val expireEvent = new Expire
-        eventManager.clearAndPut(expireEvent)
-        expireEvent.waitUntilProcessed()
+        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
+        queuedEvent.awaitProcessing()
       }
     })
     eventManager.put(Startup)
@@ -156,57 +180,70 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
    * it shuts down the partition and replica state machines. If not, those are a no-op. In addition to that, it also
    * shuts down the controller channel manager, if one exists (i.e. if it was the current controller)
    */
-  def shutdown() = {
+  def shutdown(): Unit = {
     eventManager.close()
     onControllerResignation()
   }
 
   /**
-   * On controlled shutdown shutdown, the controller first determines the partitions that the
+   * On controlled shutdown, the controller first determines the partitions that the
    * shutting down broker leads, and moves leadership of those partitions to another broker
    * that is in that partition's ISR.
    *
    * @param id Id of the broker to shutdown.
+   * @param brokerEpoch The broker epoch in the controlled shutdown request
    * @return The number of partitions that the broker still leads.
    */
-  def controlledShutdown(id: Int, controlledShutdownCallback: Try[Set[TopicPartition]] => Unit): Unit = {
-    val controlledShutdownEvent = ControlledShutdown(id, controlledShutdownCallback)
+  def controlledShutdown(id: Int, brokerEpoch: Long, controlledShutdownCallback: Try[Set[TopicPartition]] => Unit): Unit = {
+    val controlledShutdownEvent = ControlledShutdown(id, brokerEpoch, controlledShutdownCallback)
     eventManager.put(controlledShutdownEvent)
   }
 
+  private[kafka] def updateBrokerInfo(newBrokerInfo: BrokerInfo): Unit = {
+    this.brokerInfo = newBrokerInfo
+    zkClient.updateBrokerInfo(newBrokerInfo)
+  }
+
+  private[kafka] def enableDefaultUncleanLeaderElection(): Unit = {
+    eventManager.put(UncleanLeaderElectionEnable)
+  }
+
+  private[kafka] def enableTopicUncleanLeaderElection(topic: String): Unit = {
+    if (isActive) {
+      eventManager.put(TopicUncleanLeaderElectionEnable(topic))
+    }
+  }
+
   private def state: ControllerState = eventManager.state
 
   /**
    * This callback is invoked by the zookeeper leader elector on electing the current broker as the new controller.
    * It does the following things on the become-controller state change -
-   * 1. Registers controller epoch changed listener
-   * 2. Increments the controller epoch
-   * 3. Initializes the controller's context object that holds cache objects for current topics, live brokers and
+   * 1. Initializes the controller's context object that holds cache objects for current topics, live brokers and
    *    leaders for all existing partitions.
-   * 4. Starts the controller's channel manager
-   * 5. Starts the replica state machine
-   * 6. Starts the partition state machine
+   * 2. Starts the controller's channel manager
+   * 3. Starts the replica state machine
+   * 4. Starts the partition state machine
    * 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() {
-    info("Reading controller epoch from ZooKeeper")
-    readControllerEpochFromZooKeeper()
-    info("Incrementing controller epoch in ZooKeeper")
-    incrementControllerEpoch()
+  private def onControllerFailover(): Unit = {
+    maybeSetupFeatureVersioning()
+
     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)
     nodeChangeHandlers.foreach(zkClient.registerZNodeChangeHandlerAndCheckExistence)
 
     info("Deleting log dir event notifications")
-    zkClient.deleteLogDirEventNotifications()
+    zkClient.deleteLogDirEventNotifications(controllerContext.epochZkVersion)
     info("Deleting isr change notifications")
-    zkClient.deleteIsrChangeNotifications()
+    zkClient.deleteIsrChangeNotifications(controllerContext.epochZkVersion)
     info("Initializing controller context")
     initializeControllerContext()
     info("Fetching topic deletions in progress")
@@ -219,21 +256,186 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
     // they can process the LeaderAndIsrRequests that are generated by replicaStateMachine.startup() and
     // partitionStateMachine.startup().
     info("Sending update metadata request")
-    sendUpdateMetadataRequest(controllerContext.liveOrShuttingDownBrokerIds.toSeq)
+    sendUpdateMetadataRequest(controllerContext.liveOrShuttingDownBrokerIds.toSeq, Set.empty)
 
     replicaStateMachine.startup()
     partitionStateMachine.startup()
 
     info(s"Ready to serve as the new controller with epoch $epoch")
-    maybeTriggerPartitionReassignment()
+
+    initializePartitionReassignments()
     topicDeletionManager.tryTopicDeletion()
     val pendingPreferredReplicaElections = fetchPendingPreferredReplicaElections()
-    onPreferredReplicaElection(pendingPreferredReplicaElections)
+    onReplicaElection(pendingPreferredReplicaElections, ElectionType.PREFERRED, ZkTriggered)
     info("Starting the controller scheduler")
     kafkaScheduler.startup()
     if (config.autoLeaderRebalanceEnable) {
       scheduleAutoLeaderRebalanceTask(delay = 5, unit = TimeUnit.SECONDS)
     }
+
+    if (config.tokenAuthEnabled) {
+      info("starting the token expiry check scheduler")
+      tokenCleanScheduler.startup()
+      tokenCleanScheduler.schedule(name = "delete-expired-tokens",
+        fun = () => tokenManager.expireTokens(),
+        period = config.delegationTokenExpiryCheckIntervalMs,
+        unit = TimeUnit.MILLISECONDS)
+    }
+  }
+
+  private def createFeatureZNode(newNode: FeatureZNode): Int = {
+    info(s"Creating FeatureZNode at path: ${FeatureZNode.path} with contents: $newNode")
+    zkClient.createFeatureZNode(newNode)
+    val (_, newVersion) = zkClient.getDataAndVersion(FeatureZNode.path)
+    newVersion
+  }
+
+  private def updateFeatureZNode(updatedNode: FeatureZNode): Int = {
+    info(s"Updating FeatureZNode at path: ${FeatureZNode.path} with contents: $updatedNode")
+    zkClient.updateFeatureZNode(updatedNode)
+  }
+
+  /**
+   * This method enables the feature versioning system (KIP-584).
+   *
+   * Development in Kafka (from a high level) is organized into features. Each feature is tracked by
+   * a name and a range of version numbers. A feature can be of two types:
+   *
+   * 1. Supported feature:
+   * A supported feature is represented by a name (string) and a range of versions (defined by a
+   * SupportedVersionRange). It refers to a feature that a particular broker advertises support for.
+   * Each broker advertises the version ranges of its own supported features in its own
+   * BrokerIdZNode. The contents of the advertisement are specific to the particular broker and
+   * do not represent any guarantee of a cluster-wide availability of the feature for any particular
+   * range of versions.
+   *
+   * 2. Finalized feature:
+   * A finalized feature is represented by a name (string) and a range of version levels (defined
+   * by a FinalizedVersionRange). Whenever the feature versioning system (KIP-584) is
+   * enabled, the finalized features are stored in the cluster-wide common FeatureZNode.
+   * In comparison to a supported feature, the key difference is that a finalized feature exists
+   * in ZK only when it is guaranteed to be supported by any random broker in the cluster for a
+   * specified range of version levels. Also, the controller is the only entity modifying the
+   * information about finalized features.
+   *
+   * This method sets up the FeatureZNode with enabled status, which means that the finalized
+   * features stored in the FeatureZNode are active. The enabled status should be written by the
+   * controller to the FeatureZNode only when the broker IBP config is greater than or equal to
+   * KAFKA_2_7_IV0.
+   *
+   * There are multiple cases handled here:
+   *
+   * 1. New cluster bootstrap:
+   *    A new Kafka cluster (i.e. it is deployed first time) is almost always started with IBP config
+   *    setting greater than or equal to KAFKA_2_7_IV0. We would like to start the cluster with all
+   *    the possible supported features finalized immediately. Assuming this is the case, the
+   *    controller will start up and notice that the FeatureZNode is absent in the new cluster,
+   *    it will then create a FeatureZNode (with enabled status) containing the entire list of
+   *    supported features as its finalized features.
+   *
+   * 2. Broker binary upgraded, but IBP config set to lower than KAFKA_2_7_IV0:
+   *    Imagine there was an existing Kafka cluster with IBP config less than KAFKA_2_7_IV0, and the
+   *    broker binary has now been upgraded to a newer version that supports the feature versioning
+   *    system (KIP-584). But the IBP config is still set to lower than KAFKA_2_7_IV0, and may be
+   *    set to a higher value later. In this case, we want to start with no finalized features and
+   *    allow the user to finalize them whenever they are ready i.e. in the future whenever the
+   *    user sets IBP config to be greater than or equal to KAFKA_2_7_IV0, then the user could start
+   *    finalizing the features. This process ensures we do not enable all the possible features
+   *    immediately after an upgrade, which could be harmful to Kafka.
+   *    This is how we handle such a case:
+   *      - Before the IBP config upgrade (i.e. IBP config set to less than KAFKA_2_7_IV0), the
+   *        controller will start up and check if the FeatureZNode is absent.
+   *        - If the node is absent, it will react by creating a FeatureZNode with disabled status
+   *          and empty finalized features.
+   *        - Otherwise, if a node already exists in enabled status then the controller will just
+   *          flip the status to disabled and clear the finalized features.
+   *      - After the IBP config upgrade (i.e. IBP config set to greater than or equal to
+   *        KAFKA_2_7_IV0), when the controller starts up it will check if the FeatureZNode exists
+   *        and whether it is disabled.
+   *         - If the node is in disabled status, the controller won’t upgrade all features immediately.
+   *           Instead it will just switch the FeatureZNode status to enabled status. This lets the
+   *           user finalize the features later.
+   *         - Otherwise, if a node already exists in enabled status then the controller will leave
+   *           the node umodified.
+   *
+   * 3. Broker binary upgraded, with existing cluster IBP config >= KAFKA_2_7_IV0:
+   *    Imagine there was an existing Kafka cluster with IBP config >= KAFKA_2_7_IV0, and the broker
+   *    binary has just been upgraded to a newer version (that supports IBP config KAFKA_2_7_IV0 and
+   *    higher). The controller will start up and find that a FeatureZNode is already present with
+   *    enabled status and existing finalized features. In such a case, the controller leaves the node
+   *    unmodified.
+   *
+   * 4. Broker downgrade:
+   *    Imagine that a Kafka cluster exists already and the IBP config is greater than or equal to
+   *    KAFKA_2_7_IV0. Then, the user decided to downgrade the cluster by setting IBP config to a
+   *    value less than KAFKA_2_7_IV0. This means the user is also disabling the feature versioning
+   *    system (KIP-584). In this case, when the controller starts up with the lower IBP config, it
+   *    will switch the FeatureZNode status to disabled with empty features.
+   */
+  private def enableFeatureVersioning(): Unit = {
+    val (mayBeFeatureZNodeBytes, version) = zkClient.getDataAndVersion(FeatureZNode.path)
+    if (version == ZkVersion.UnknownVersion) {
+      val newVersion = createFeatureZNode(new FeatureZNode(FeatureZNodeStatus.Enabled,
+                                          brokerFeatures.defaultFinalizedFeatures))
+      featureCache.waitUntilEpochOrThrow(newVersion, config.zkConnectionTimeoutMs)
+    } else {
+      val existingFeatureZNode = FeatureZNode.decode(mayBeFeatureZNodeBytes.get)
+      val newFeatures = existingFeatureZNode.status match {
+        case FeatureZNodeStatus.Enabled => existingFeatureZNode.features
+        case FeatureZNodeStatus.Disabled =>
+          if (!existingFeatureZNode.features.empty()) {
+            warn(s"FeatureZNode at path: ${FeatureZNode.path} with disabled status" +
+                 s" contains non-empty features: ${existingFeatureZNode.features}")
+          }
+          Features.emptyFinalizedFeatures
+      }
+      val newFeatureZNode = new FeatureZNode(FeatureZNodeStatus.Enabled, newFeatures)
+      if (!newFeatureZNode.equals(existingFeatureZNode)) {
+        val newVersion = updateFeatureZNode(newFeatureZNode)
+        featureCache.waitUntilEpochOrThrow(newVersion, config.zkConnectionTimeoutMs)
+      }
+    }
+  }
+
+  /**
+   * Disables the feature versioning system (KIP-584).
+   *
+   * Sets up the FeatureZNode with disabled status. This status means the feature versioning system
+   * (KIP-584) is disabled, and, the finalized features stored in the FeatureZNode are not relevant.
+   * This status should be written by the controller to the FeatureZNode only when the broker
+   * IBP config is less than KAFKA_2_7_IV0.
+   *
+   * NOTE:
+   * 1. When this method returns, existing finalized features (if any) will be cleared from the
+   *    FeatureZNode.
+   * 2. This method, unlike enableFeatureVersioning() need not wait for the FinalizedFeatureCache
+   *    to be updated, because, such updates to the cache (via FinalizedFeatureChangeListener)
+   *    are disabled when IBP config is < than KAFKA_2_7_IV0.
+   */
+  private def disableFeatureVersioning(): Unit = {
+    val newNode = FeatureZNode(FeatureZNodeStatus.Disabled, Features.emptyFinalizedFeatures())
+    val (mayBeFeatureZNodeBytes, version) = zkClient.getDataAndVersion(FeatureZNode.path)
+    if (version == ZkVersion.UnknownVersion) {
+      createFeatureZNode(newNode)
+    } else {
+      val existingFeatureZNode = FeatureZNode.decode(mayBeFeatureZNodeBytes.get)
+      if (existingFeatureZNode.status == FeatureZNodeStatus.Disabled &&
+          !existingFeatureZNode.features.empty()) {
+        warn(s"FeatureZNode at path: ${FeatureZNode.path} with disabled status" +
+             s" contains non-empty features: ${existingFeatureZNode.features}")
+      }
+      if (!newNode.equals(existingFeatureZNode)) {
+        updateFeatureZNode(newNode)
+      }
+    }
+  }
+
+  private def maybeSetupFeatureVersioning(): Unit = {
+    if (config.isFeatureVersioningSupported) {
+      enableFeatureVersioning()
+    } else {
+      disableFeatureVersioning()
+    }
   }
 
   private def scheduleAutoLeaderRebalanceTask(delay: Long, unit: TimeUnit): Unit = {
@@ -245,16 +447,14 @@ 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)
     zkClient.unregisterZNodeChangeHandler(partitionReassignmentHandler.path)
     zkClient.unregisterZNodeChangeHandler(preferredReplicaElectionHandler.path)
     zkClient.unregisterZNodeChildChangeHandler(logDirEventNotificationHandler.path)
-
-    // reset topic deletion manager
-    topicDeletionManager.reset()
+    unregisterBrokerModificationsHandler(brokerModificationsHandlers.keySet)
 
     // shutdown leader rebalance scheduler
     kafkaScheduler.shutdown()
@@ -262,6 +462,14 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
     preferredReplicaImbalanceCount = 0
     globalTopicCount = 0
     globalPartitionCount = 0
+    topicsToDeleteCount = 0
+    replicasToDeleteCount = 0
+    ineligibleTopicsToDeleteCount = 0
+    ineligibleReplicasToDeleteCount = 0
+
+    // stop token expiry check scheduler
+    if (tokenCleanScheduler.isStarted)
+      tokenCleanScheduler.shutdown()
 
     // de-register partition ISR listener for on-going partition reassignment task
     unregisterPartitionReassignmentIsrChangeHandlers()
@@ -274,7 +482,8 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
     replicaStateMachine.shutdown()
     zkClient.unregisterZNodeChildChangeHandler(brokerChangeHandler.path)
 
-    resetControllerContext()
+    controllerChannelManager.shutdown()
+    controllerContext.resetContext()
 
     info("Resigned")
   }
@@ -282,10 +491,12 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
   /*
    * This callback is invoked by the controller's LogDirEventNotificationListener with the list of broker ids who
    * have experienced new log directory failures. In response the controller should send LeaderAndIsrRequest
-   * to all these brokers to query the state of their replicas
+   * 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)
     replicaStateMachine.handleStateChanges(replicasOnBrokers.toSeq, OnlineReplica)
   }
@@ -304,15 +515,18 @@ 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
-    // send update metadata request to all live and shutting down brokers. Old brokers will get to know of the new
-    // broker via this update.
+    val existingBrokers = controllerContext.liveOrShuttingDownBrokerIds.diff(newBrokersSet)
+    // Send update metadata request to all the existing brokers in the cluster so that they know about the new brokers
+    // via this update. No need to include any partition states in the request since there are no partition state changes.
+    sendUpdateMetadataRequest(existingBrokers.toSeq, Set.empty)
+    // Send update metadata request to all the new brokers in the cluster with a full set of partition states for initialization.
     // In cases of controlled shutdown leaders will not be elected when a new broker comes up. So at least in the
-    // common controlled shutdown case, the metadata will reach the new brokers faster
-    sendUpdateMetadataRequest(controllerContext.liveOrShuttingDownBrokerIds.toSeq)
+    // common controlled shutdown case, the metadata will reach the new brokers faster.
+    sendUpdateMetadataRequest(newBrokers, controllerContext.partitionsWithLeaders)
     // the very first thing to do when a new broker comes up is send it the entire list of partitions that it is
     // supposed to host. Based on that the broker starts the high watermark threads for the input list of partitions
     val allReplicasOnNewBrokers = controllerContext.replicasOnBrokers(newBrokersSet)
@@ -321,33 +535,65 @@ 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(eventManager, brokerId)
+      zkClient.registerZNodeChangeHandlerAndCheckExistence(brokerModificationsHandler)
+      brokerModificationsHandlers.put(brokerId, brokerModificationsHandler)
+    }
+  }
+
+  private def unregisterBrokerModificationsHandler(brokerIds: Iterable[Int]): Unit = {
+    debug(s"Unregister BrokerModifications handler for $brokerIds")
+    brokerIds.foreach { brokerId =>
+      brokerModificationsHandlers.remove(brokerId).foreach(handler => zkClient.unregisterZNodeChangeHandler(handler.path))
+    }
   }
 
   /*
    * 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 =
       deadBrokers.filter(id => controllerContext.shuttingDownBrokerIds.remove(id))
-    info(s"Removed $deadBrokersThatWereShuttingDown from list of shutting down brokers.")
+    if (deadBrokersThatWereShuttingDown.nonEmpty)
+      info(s"Removed ${deadBrokersThatWereShuttingDown.mkString(",")} from list of shutting down brokers.")
     val allReplicasOnDeadBrokers = controllerContext.replicasOnBrokers(deadBrokers.toSet)
     onReplicasBecomeOffline(allReplicasOnDeadBrokers)
+
+    unregisterBrokerModificationsHandler(deadBrokers)
+  }
+
+  private def onBrokerUpdate(updatedBrokerId: Int): Unit = {
+    info(s"Broker info update callback for $updatedBrokerId")
+    sendUpdateMetadataRequest(controllerContext.liveOrShuttingDownBrokerIds.toSeq, Set.empty)
   }
 
   /**
@@ -365,12 +611,10 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
     val (newOfflineReplicasForDeletion, newOfflineReplicasNotForDeletion) =
       newOfflineReplicas.partition(p => topicDeletionManager.isTopicQueuedUpForDeletion(p.topic))
 
-    val partitionsWithoutLeader = controllerContext.partitionLeadershipInfo.filter(partitionAndLeader =>
-      !controllerContext.isReplicaOnline(partitionAndLeader._2.leaderAndIsr.leader, partitionAndLeader._1) &&
-        !topicDeletionManager.isTopicQueuedUpForDeletion(partitionAndLeader._1.topic)).keySet
+    val partitionsWithOfflineLeader = controllerContext.partitionsWithOfflineLeader
 
     // trigger OfflinePartition state for all partitions whose current leader is one amongst the newOfflineReplicas
-    partitionStateMachine.handleStateChanges(partitionsWithoutLeader.toSeq, OfflinePartition)
+    partitionStateMachine.handleStateChanges(partitionsWithOfflineLeader.toSeq, OfflinePartition)
     // trigger OnlinePartition state changes for offline or new partitions
     partitionStateMachine.triggerOnlinePartitionStateChange()
     // trigger OfflineReplica state change for those newly offline replicas
@@ -384,10 +628,10 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
       topicDeletionManager.failReplicaDeletion(newOfflineReplicasForDeletion)
     }
 
-    // If replica failure did not require leader re-election, inform brokers of the offline replica
+    // If replica failure did not require leader re-election, inform brokers of the offline brokers
     // Note that during leader re-election, brokers update their metadata
-    if (partitionsWithoutLeader.isEmpty) {
-      sendUpdateMetadataRequest(controllerContext.liveOrShuttingDownBrokerIds.toSeq)
+    if (partitionsWithOfflineLeader.isEmpty) {
+      sendUpdateMetadataRequest(controllerContext.liveOrShuttingDownBrokerIds.toSeq, Set.empty)
     }
   }
 
@@ -397,189 +641,288 @@ 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.
+   *
+   *   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
+   *
+   * 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.
+   *
+   * 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
    *
-   * 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.
+   *   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.
    *
-   * 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)
+   * 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
    *
-   * Note that we have to update AR in ZK with RAR last since it's the only place where we store OAR persistently.
+   * 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.
-      removePartitionFromReassignedPartitions(topicPartition)
-      info(s"Removed partition $topicPartition from the list of reassigned partitions in zookeeper")
-      controllerContext.partitionsBeingReassigned.remove(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))
     }
   }
 
-  private def watchIsrChangesForReassignedPartition(partition: TopicPartition,
-                                                    reassignedPartitionContext: ReassignedPartitionsContext) {
-    val reassignIsrChangeHandler = new PartitionReassignmentIsrChangeHandler(this, eventManager, partition)
-    reassignedPartitionContext.reassignIsrChangeHandler = reassignIsrChangeHandler
-    // register listener on the leader and isr path to wait until they catch up with the current leader
-    zkClient.registerZNodeChangeHandler(reassignIsrChangeHandler)
+  /**
+   * 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
+   */
+  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)
+    }
+
+    if (!isAlterIsrEnabled) {
+      val reassignIsrChangeHandler = new PartitionReassignmentIsrChangeHandler(eventManager, topicPartition)
+      zkClient.registerZNodeChangeHandler(reassignIsrChangeHandler)
+    }
+
+    controllerContext.partitionsBeingReassigned.add(topicPartition)
   }
 
-  private def initiateReassignReplicasForTopicPartition(topicPartition: TopicPartition,
-                                                reassignedPartitionContext: ReassignedPartitionsContext) {
-    val newReplicas = reassignedPartitionContext.newReplicas
-    val topic = topicPartition.topic
-    try {
-      val assignedReplicasOpt = controllerContext.partitionReplicaAssignment.get(topicPartition)
-      assignedReplicasOpt match {
-        case Some(assignedReplicas) =>
-          if (assignedReplicas == newReplicas) {
-            throw new KafkaException(s"Partition $topicPartition to be reassigned is already assigned to replicas " +
-              s"${newReplicas.mkString(",")}. Ignoring request for partition reassignment")
-          } else {
-            info(s"Handling reassignment of partition $topicPartition to new replicas ${newReplicas.mkString(",")}")
-            // first register ISR change listener
-            watchIsrChangesForReassignedPartition(topicPartition, reassignedPartitionContext)
-            controllerContext.partitionsBeingReassigned.put(topicPartition, reassignedPartitionContext)
-            // mark topic ineligible for deletion for the partitions being reassigned
-            topicDeletionManager.markTopicIneligibleForDeletion(Set(topic))
-            onPartitionReassignment(topicPartition, reassignedPartitionContext)
+  /**
+   * 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
+   *
+   * 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.
+   *
+   * @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(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")
+        new ApiError(Errors.UNKNOWN_TOPIC_OR_PARTITION, "The partition does not exist.")
+      } else {
+        val assignedReplicas = controllerContext.partitionReplicaAssignment(tp)
+        if (assignedReplicas.nonEmpty) {
+          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)
           }
-        case None => throw new KafkaException(s"Attempt to reassign partition $topicPartition that doesn't exist")
+        } else {
+          new ApiError(Errors.UNKNOWN_TOPIC_OR_PARTITION, "The partition does not exist.")
+        }
       }
-    } catch {
-      case e: Throwable =>
-        error(s"Error completing reassignment of partition $topicPartition", e)
-        // remove the partition from the admin path to unblock the admin client
-        removePartitionFromReassignedPartitions(topicPartition)
+
+      tp -> apiError
     }
   }
 
-  private def onPreferredReplicaElection(partitions: Set[TopicPartition], isTriggeredByAutoRebalance: Boolean = false) {
-    info(s"Starting preferred replica leader election for partitions ${partitions.mkString(",")}")
+  /**
+    * 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[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 {
-      partitionStateMachine.handleStateChanges(partitions.toSeq, OnlinePartition, Option(PreferredReplicaPartitionLeaderElectionStrategy))
-    } catch {
-      case e: Throwable => error(s"Error completing preferred replica leader election for partitions ${partitions.mkString(",")}", e)
-    } finally {
-      removePartitionsFromPreferredReplicaElection(partitions, isTriggeredByAutoRebalance)
-    }
-  }
+      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)
+      }
 
-  private def incrementControllerEpoch(): Unit = {
-    val newControllerEpoch = controllerContext.epoch + 1
-    val setDataResponse = zkClient.setControllerEpochRaw(newControllerEpoch, controllerContext.epochZkVersion)
-    setDataResponse.resultCode match {
-      case Code.OK =>
-        controllerContext.epochZkVersion = setDataResponse.stat.getVersion
-        controllerContext.epoch = newControllerEpoch
-      case Code.NONODE =>
-        // if path doesn't exist, this is the first controller whose epoch should be 1
-        // the following call can still fail if another controller gets elected between checking if the path exists and
-        // trying to create the controller epoch path
-        val createResponse = zkClient.createControllerEpochRaw(KafkaController.InitialControllerEpoch)
-        createResponse.resultCode match {
-          case Code.OK =>
-            controllerContext.epoch = KafkaController.InitialControllerEpoch
-            controllerContext.epochZkVersion = KafkaController.InitialControllerEpochZkVersion
-          case Code.NODEEXISTS =>
-            throw new ControllerMovedException("Controller moved to another broker. Aborting controller startup procedure")
-          case _ =>
-            val exception = createResponse.resultException.get
-            error("Error while incrementing controller epoch", exception)
-            throw exception
+      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
         }
-      case _ =>
-        throw new ControllerMovedException("Controller moved to another broker. Aborting controller startup procedure")
+      }
+
+      results
+    } finally {
+      if (electionTrigger != AdminClientTriggered) {
+        removePartitionsFromPreferredReplicaElection(partitions, electionTrigger == AutoTriggered)
+      }
     }
-    info(s"Epoch incremented to ${controllerContext.epoch}")
   }
 
-  private def initializeControllerContext() {
+  private def initializeControllerContext(): Unit = {
     // update controller cache with delete topic information
-    controllerContext.liveBrokers = zkClient.getAllBrokersInCluster.toSet
-    controllerContext.allTopics = zkClient.getAllTopicsInCluster.toSet
+    val curBrokerAndEpochs = zkClient.getAllBrokerAndEpochsInCluster
+    val (compatibleBrokerAndEpochs, incompatibleBrokerAndEpochs) = partitionOnFeatureCompatibility(curBrokerAndEpochs)
+    if (!incompatibleBrokerAndEpochs.isEmpty) {
+      warn("Ignoring registration of new brokers due to incompatibilities with finalized features: " +
+        incompatibleBrokerAndEpochs.map { case (broker, _) => broker.id }.toSeq.sorted.mkString(","))
+    }
+    controllerContext.setLiveBrokers(compatibleBrokerAndEpochs)
+    info(s"Initialized broker epochs cache: ${controllerContext.liveBrokerIdAndEpochs}")
+    controllerContext.setAllTopics(zkClient.getAllTopicsInCluster(true))
     registerPartitionModificationsHandlers(controllerContext.allTopics.toSeq)
-    controllerContext.partitionReplicaAssignment = mutable.Map.empty ++ zkClient.getReplicaAssignmentForTopics(controllerContext.allTopics.toSet)
-    controllerContext.partitionLeadershipInfo = new mutable.HashMap[TopicPartition, LeaderIsrAndControllerEpoch]
-    controllerContext.shuttingDownBrokerIds = mutable.Set.empty[Int]
+    val replicaAssignmentAndTopicIds = zkClient.getReplicaAssignmentAndTopicIdForTopics(controllerContext.allTopics.toSet)
+    processTopicIds(replicaAssignmentAndTopicIds)
+
+    replicaAssignmentAndTopicIds.foreach { case TopicIdReplicaAssignment(_, _, assignments) =>
+      assignments.foreach { case (topicPartition, replicaAssignment) =>
+        controllerContext.updatePartitionFullReplicaAssignment(topicPartition, replicaAssignment)
+        if (replicaAssignment.isBeingReassigned)
+          controllerContext.partitionsBeingReassigned.add(topicPartition)
+      }
+    }
+    controllerContext.clearPartitionLeadershipInfo()
+    controllerContext.shuttingDownBrokerIds.clear()
+    // register broker modifications handlers
+    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}")
@@ -589,10 +932,10 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
     val partitionsUndergoingPreferredReplicaElection = zkClient.getPreferredReplicaElection
     // check if they are already completed or topic was deleted
     val partitionsThatCompletedPreferredReplicaElection = partitionsUndergoingPreferredReplicaElection.filter { partition =>
-      val replicasOpt = controllerContext.partitionReplicaAssignment.get(partition)
-      val topicDeleted = replicasOpt.isEmpty
+      val replicas = controllerContext.partitionReplicaAssignment(partition)
+      val topicDeleted = replicas.isEmpty
       val successful =
-        if (!topicDeleted) controllerContext.partitionLeadershipInfo(partition).leaderAndIsr.leader == replicasOpt.get.head else false
+        if (!topicDeleted) controllerContext.partitionLeadershipInfo(partition).get.leaderAndIsr.leader == replicas.head else false
       successful || topicDeleted
     }
     val pendingPreferredReplicaElectionsIgnoringTopicDeletion = partitionsUndergoingPreferredReplicaElection -- partitionsThatCompletedPreferredReplicaElection
@@ -605,167 +948,145 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
     pendingPreferredReplicaElections
   }
 
-  private def resetControllerContext(): Unit = {
-    if (controllerContext.controllerChannelManager != null) {
-      controllerContext.controllerChannelManager.shutdown()
-      controllerContext.controllerChannelManager = null
+  /**
+   * 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)
     }
-    controllerContext.shuttingDownBrokerIds.clear()
-    controllerContext.epoch = 0
-    controllerContext.epochZkVersion = 0
-    controllerContext.allTopics = Set.empty
-    controllerContext.partitionReplicaAssignment.clear()
-    controllerContext.partitionLeadershipInfo.clear()
-    controllerContext.partitionsBeingReassigned.clear()
-    controllerContext.liveBrokers = Set.empty
-  }
-
-  private def initializePartitionReassignment() {
-    // read the partitions being reassigned from zookeeper path /admin/reassign_partitions
-    val partitionsBeingReassigned = zkClient.getPartitionReassignment
-    // check if they are already completed or topic was deleted
-    val reassignedPartitions = partitionsBeingReassigned.filter { case (tp, reassignmentReplicas) =>
-      controllerContext.partitionReplicaAssignment.get(tp) match {
-        case None => true // topic deleted
-        case Some(currentReplicas) => currentReplicas == reassignmentReplicas // reassignment completed
-      }
-    }.keys
-    reassignedPartitions.foreach(removePartitionFromReassignedPartitions)
-    val partitionsToReassign = partitionsBeingReassigned -- reassignedPartitions
-    controllerContext.partitionsBeingReassigned ++= partitionsToReassign.mapValues(new ReassignedPartitionsContext(_))
-    info(s"Partitions being reassigned: $partitionsBeingReassigned")
-    info(s"Partitions already reassigned: $reassignedPartitions")
-    info(s"Resuming reassignment of partitions: $partitionsToReassign")
   }
 
   private def fetchTopicDeletionsInProgress(): (Set[String], Set[String]) = {
     val topicsToBeDeleted = zkClient.getTopicDeletions.toSet
-    val topicsWithOfflineReplicas = controllerContext.partitionReplicaAssignment.filter { case (partition, replicas) =>
-      replicas.exists(r => !controllerContext.isReplicaOnline(r, partition))
-    }.keySet.map(_.topic)
-    val topicsForWhichPartitionReassignmentIsInProgress = controllerContext.partitionsBeingReassigned.keySet.map(_.topic)
+    val topicsWithOfflineReplicas = controllerContext.allTopics.filter { topic => {
+      val replicasForTopic = controllerContext.replicasForTopic(topic)
+      replicasForTopic.exists(r => !controllerContext.isReplicaOnline(r.replica, r.topicPartition))
+    }}
+    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 maybeTriggerPartitionReassignment() {
-    controllerContext.partitionsBeingReassigned.foreach { case (tp, reassignContext) =>
-      initiateReassignReplicasForTopicPartition(tp, reassignContext)
-    }
-  }
-
-  private def startChannelManager() {
-    controllerContext.controllerChannelManager = new ControllerChannelManager(controllerContext, config, time, metrics,
-      stateChangeLogger, threadNamePrefix)
-    controllerContext.controllerChannelManager.startup()
-  }
-
-  private def updateLeaderAndIsrCache(partitions: Seq[TopicPartition] = controllerContext.partitionReplicaAssignment.keys.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)
+    leaderIsrAndControllerEpochs.forKeyValue { (partition, leaderIsrAndControllerEpoch) =>
+      controllerContext.putPartitionLeadershipInfo(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)
+      }
     }
   }
 
   private def moveReassignedPartitionLeaderIfRequired(topicPartition: TopicPartition,
-                                                      reassignedPartitionContext: ReassignedPartitionsContext) {
-    val reassignedReplicas = reassignedPartitionContext.newReplicas
-    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.partitionReplicaAssignment.put(topicPartition, reassignedReplicas)
-    if (!reassignedPartitionContext.newReplicas.contains(currentLeader)) {
+                                                      newAssignment: ReplicaAssignment): Unit = {
+    val reassignedReplicas = newAssignment.replicas
+    val currentLeader = controllerContext.partitionLeadershipInfo(topicPartition).get.leaderAndIsr.leader
+
+    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]) {
-    val partitionsAndReplicasForThisTopic = controllerContext.partitionReplicaAssignment.filter(_._1.topic == partition.topic)
-    partitionsAndReplicasForThisTopic.put(partition, replicas)
-    val setDataResponse = zkClient.setTopicAssignmentRaw(partition.topic, partitionsAndReplicasForThisTopic.toMap)
+  private def updateReplicaAssignmentForPartition(topicPartition: TopicPartition, assignment: ReplicaAssignment): Unit = {
+    val topicAssignment = mutable.Map() ++=
+      controllerContext.partitionFullReplicaAssignmentForTopic(topicPartition.topic) +=
+      (topicPartition -> assignment)
+
+    val setDataResponse = zkClient.setTopicAssignmentRaw(topicPartition.topic,
+      controllerContext.topicIds(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.partitionReplicaAssignment.put(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)
+          // the isNew flag, when set to true, makes sure that when a replica possibly resided
+          // in a logDir that is offline, we refrain from just creating a new replica in a good
+          // logDir. This is exactly the behavior we want for the original replicas, but not
+          // for the replicas we add in this reassignment. For new replicas, want to be able
+          // to assign to one of the good logDirs.
+          brokerRequestBatch.addLeaderAndIsrRequestForBrokers(assignment.originReplicas, topicPartition,
+            updatedLeaderIsrAndControllerEpoch, assignment, isNew = false)
+          brokerRequestBatch.addLeaderAndIsrRequestForBrokers(assignment.addingReplicas, topicPartition,
+            updatedLeaderIsrAndControllerEpoch, assignment, isNew = true)
           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.info(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)
@@ -773,83 +1094,90 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
 
   private[controller] def unregisterPartitionModificationsHandlers(topics: Seq[String]) = {
     topics.foreach { topic =>
-      partitionModificationsHandlers.remove(topic)
-        .foreach(handler => zkClient.unregisterZNodeChangeHandler(handler.path))
+      partitionModificationsHandlers.remove(topic).foreach(handler => zkClient.unregisterZNodeChangeHandler(handler.path))
     }
   }
 
-  private def unregisterPartitionReassignmentIsrChangeHandlers() {
-    controllerContext.partitionsBeingReassigned.values.foreach { reassignedPartitionsContext =>
-      zkClient.unregisterZNodeChangeHandler(reassignedPartitionsContext.reassignIsrChangeHandler.path)
+  private def unregisterPartitionReassignmentIsrChangeHandlers(): Unit = {
+    if (!isAlterIsrEnabled) {
+      controllerContext.partitionsBeingReassigned.foreach { tp =>
+        val path = TopicPartitionStateZNode.path(tp)
+        zkClient.unregisterZNodeChangeHandler(path)
+      }
     }
   }
 
-  private def readControllerEpochFromZooKeeper() {
-    // initialize the controller epoch and zk version by reading from zookeeper
-    val epochAndStatOpt = zkClient.getControllerEpoch
-    epochAndStatOpt.foreach { case (epoch, stat) =>
-      controllerContext.epoch = epoch
-      controllerContext.epochZkVersion = stat.getVersion
-      info(s"Initialized controller epoch to ${controllerContext.epoch} and zk version ${controllerContext.epochZkVersion}")
+  private def removePartitionFromReassigningPartitions(topicPartition: TopicPartition,
+                                                       assignment: ReplicaAssignment): Unit = {
+    if (controllerContext.partitionsBeingReassigned.contains(topicPartition)) {
+      if (!isAlterIsrEnabled) {
+        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")
     }
   }
 
-  private def removePartitionFromReassignedPartitions(topicPartition: TopicPartition) {
-    controllerContext.partitionsBeingReassigned.get(topicPartition).foreach { reassignContext =>
-      // stop watching the ISR changes for this partition
-      zkClient.unregisterZNodeChangeHandler(reassignContext.reassignIsrChangeHandler.path)
-    }
+  /**
+   * Remove partitions from an active zk-based reassignment (if one exists).
+   *
+   * @param shouldRemoveReassignment Predicate indicating which partition reassignments should be removed
+   */
+  private def maybeRemoveFromZkReassignment(shouldRemoveReassignment: (TopicPartition, Seq[Int]) => Boolean): Unit = {
+    if (!zkClient.reassignPartitionsInProgress)
+      return
 
-    val updatedPartitionsBeingReassigned = controllerContext.partitionsBeingReassigned - topicPartition
+    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()
+      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)
-      catch {
+      try {
+        zkClient.setOrCreatePartitionReassignment(updatedPartitionsBeingReassigned, controllerContext.epochZkVersion)
+      } catch {
         case e: KeeperException => throw new AdminOperationException(e)
       }
     }
-
-    controllerContext.partitionsBeingReassigned.remove(topicPartition)
   }
 
   private def removePartitionsFromPreferredReplicaElection(partitionsToBeRemoved: Set[TopicPartition],
-                                                   isTriggeredByAutoRebalance : Boolean) {
+                                                           isTriggeredByAutoRebalance : Boolean): Unit = {
     for (partition <- partitionsToBeRemoved) {
       // check the status
-      val currentLeader = controllerContext.partitionLeadershipInfo(partition).leaderAndIsr.leader
+      val currentLeader = controllerContext.partitionLeadershipInfo(partition).get.leaderAndIsr.leader
       val preferredReplica = controllerContext.partitionReplicaAssignment(partition).head
       if (currentLeader == preferredReplica) {
         info(s"Partition $partition completed preferred replica leader election. New leader is $preferredReplica")
       } else {
-        warn(s"Partition $partition failed to complete preferred replica leader election. Leader is $currentLeader")
+        warn(s"Partition $partition failed to complete preferred replica leader election to $preferredReplica. " +
+          s"Leader is still $currentLeader")
       }
     }
     if (!isTriggeredByAutoRebalance) {
-      zkClient.deletePreferredReplicaElection()
+      zkClient.deletePreferredReplicaElection(controllerContext.epochZkVersion)
       // Ensure we detect future preferred replica leader elections
-      eventManager.put(PreferredReplicaLeaderElection)
+      eventManager.put(ReplicaLeaderElection(None, ElectionType.PREFERRED, ZkTriggered))
     }
   }
 
-  private[controller] def sendRequest(brokerId: Int, apiKey: ApiKeys, request: AbstractRequest.Builder[_ <: AbstractRequest],
-                                      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] = Set.empty[TopicPartition]) {
+  private[controller] def sendUpdateMetadataRequest(brokers: Seq[Int], partitions: Set[TopicPartition]): Unit = {
     try {
       brokerRequestBatch.newBatch()
       brokerRequestBatch.addUpdateMetadataRequestForBrokers(brokers, partitions)
@@ -884,16 +1212,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) =
-            zkClient.updateLeaderAndIsr(immutable.Map(partition -> newLeaderAndIsr), epoch)
-          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
+          val UpdateLeaderAndIsrResult(finishedUpdates, _) =
+            zkClient.updateLeaderAndIsr(immutable.Map(partition -> newLeaderAndIsr), epoch, controllerContext.epochZkVersion)
+
+          finishedUpdates.get(partition) match {
+            case Some(Right(leaderAndIsr)) =>
+              val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(leaderAndIsr, epoch)
+              controllerContext.putPartitionLeadershipInfo(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")
@@ -905,15 +1236,16 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
   private def checkAndTriggerAutoLeaderRebalance(): Unit = {
     trace("Checking need to trigger auto leader balancing")
     val preferredReplicasForTopicsByBrokers: Map[Int, Map[TopicPartition, Seq[Int]]] =
-      controllerContext.partitionReplicaAssignment.filterNot { case (tp, _) =>
-        topicDeletionManager.isTopicQueuedUpForDeletion(tp.topic)
-      }.groupBy { case (_, assignedReplicas) => assignedReplicas.head }
-    debug(s"Preferred replicas by broker $preferredReplicasForTopicsByBrokers")
+      controllerContext.allPartitions.filterNot {
+        tp => topicDeletionManager.isTopicQueuedUpForDeletion(tp.topic)
+      }.map { tp =>
+        (tp, controllerContext.partitionReplicaAssignment(tp) )
+      }.toMap.groupBy { case (_, assignedReplicas) => assignedReplicas.head }
 
     // for each broker, check if a preferred replica election needs to be triggered
-    preferredReplicasForTopicsByBrokers.foreach { case (leaderBroker, topicPartitionsForBroker) =>
+    preferredReplicasForTopicsByBrokers.forKeyValue { (leaderBroker, topicPartitionsForBroker) =>
       val topicsNotInPreferredReplica = topicPartitionsForBroker.filter { case (topicPartition, _) =>
-        val leadershipInfo = controllerContext.partitionLeadershipInfo.get(topicPartition)
+        val leadershipInfo = controllerContext.partitionLeadershipInfo(topicPartition)
         leadershipInfo.exists(_.leaderAndIsr.leader != leaderBroker)
       }
       debug(s"Topics not in preferred replica for broker $leaderBroker $topicsNotInPreferredReplica")
@@ -924,154 +1256,172 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
       // check ratio and if greater than desired ratio, trigger a rebalance for the topic partitions
       // that need to be on this broker
       if (imbalanceRatio > (config.leaderImbalancePerBrokerPercentage.toDouble / 100)) {
-        topicsNotInPreferredReplica.keys.foreach { topicPartition =>
-          // do this check only if the broker is live and there are no partitions being reassigned currently
-          // and preferred replica election is not in progress
-          if (controllerContext.isReplicaOnline(leaderBroker, topicPartition) &&
-            controllerContext.partitionsBeingReassigned.isEmpty &&
-            !topicDeletionManager.isTopicQueuedUpForDeletion(topicPartition.topic) &&
-            controllerContext.allTopics.contains(topicPartition.topic)) {
-            onPreferredReplicaElection(Set(topicPartition), isTriggeredByAutoRebalance = true)
-          }
-        }
+        // do this check only if the broker is live and there are no partitions being reassigned currently
+        // and preferred replica election is not in progress
+        val candidatePartitions = topicsNotInPreferredReplica.keys.filter(tp =>
+          controllerContext.partitionsBeingReassigned.isEmpty &&
+          !topicDeletionManager.isTopicQueuedUpForDeletion(tp.topic) &&
+          controllerContext.allTopics.contains(tp.topic) &&
+          canPreferredReplicaBeLeader(tp)
+       )
+        onReplicaElection(candidatePartitions.toSet, ElectionType.PREFERRED, AutoTriggered)
       }
     }
   }
 
-  case object AutoPreferredReplicaLeaderElection extends ControllerEvent {
-
-    def state = ControllerState.AutoLeaderBalance
+  private def canPreferredReplicaBeLeader(tp: TopicPartition): Boolean = {
+    val assignment = controllerContext.partitionReplicaAssignment(tp)
+    val liveReplicas = assignment.filter(replica => controllerContext.isReplicaOnline(replica, tp))
+    val isr = controllerContext.partitionLeadershipInfo(tp).get.leaderAndIsr.isr
+    PartitionLeaderElectionAlgorithms
+      .preferredReplicaPartitionLeaderElection(assignment, isr, liveReplicas.toSet)
+      .nonEmpty
+  }
 
-    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 class ControlledShutdown(id: Int, controlledShutdownCallback: Try[Set[TopicPartition]] => Unit) extends ControllerEvent {
+  private def processUncleanLeaderElectionEnable(): Unit = {
+    if (!isActive) return
+    info("Unclean leader election has been enabled by default")
+    partitionStateMachine.triggerOnlinePartitionStateChange()
+  }
 
-    def state = ControllerState.ControlledShutdown
+  private def processTopicUncleanLeaderElectionEnable(topic: String): Unit = {
+    if (!isActive) return
+    info(s"Unclean leader election has been enabled for topic $topic")
+    partitionStateMachine.triggerOnlinePartitionStateChange(topic)
+  }
+
+  private def processControlledShutdown(id: Int, brokerEpoch: Long, controlledShutdownCallback: Try[Set[TopicPartition]] => Unit): Unit = {
+    val controlledShutdownResult = Try { doControlledShutdown(id, brokerEpoch) }
+    controlledShutdownCallback(controlledShutdownResult)
+  }
 
-    override def process(): Unit = {
-      val controlledShutdownResult = Try { doControlledShutdown(id) }
-      controlledShutdownCallback(controlledShutdownResult)
+  private def doControlledShutdown(id: Int, brokerEpoch: Long): Set[TopicPartition] = {
+    if (!isActive) {
+      throw new ControllerMovedException("Controller moved to another broker. Aborting controlled shutdown")
     }
 
-    private def doControlledShutdown(id: Int): 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)
-      }
-      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) =>
-            leaderIsrAndControllerEpoch.leaderAndIsr.leader == id && controllerContext.partitionReplicaAssignment(topicPartition).size > 1
-        }.keys
+    val partitionsToActOn = controllerContext.partitionsOnBroker(id).filter { partition =>
+      controllerContext.partitionReplicaAssignment(partition).size > 1 &&
+        controllerContext.partitionLeadershipInfo(partition).isDefined &&
+        !topicDeletionManager.isTopicQueuedUpForDeletion(partition.topic)
+    }
+    val (partitionsLedByBroker, partitionsFollowedByBroker) = partitionsToActOn.partition { partition =>
+      controllerContext.partitionLeadershipInfo(partition).get.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)
+    trace(s"All leaders = ${controllerContext.partitionsLeadershipInfo.mkString(",")}")
+    controllerContext.partitionLeadersOnBroker(id)
   }
 
-  case class LeaderAndIsrResponseReceived(LeaderAndIsrResponseObj: AbstractResponse, brokerId: Int) extends ControllerEvent {
-
-    def state = ControllerState.LeaderAndIsrResponseReceived
+  private def processUpdateMetadataResponseReceived(updateMetadataResponse: UpdateMetadataResponse, brokerId: Int): Unit = {
+    if (!isActive) return
 
-    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
-      }
-
-      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)))
-      }
+    if (updateMetadataResponse.error != Errors.NONE) {
+      stateChangeLogger.error(s"Received error ${updateMetadataResponse.error} in UpdateMetadata " +
+        s"response $updateMetadataResponse from broker $brokerId")
     }
   }
 
-  case class TopicDeletionStopReplicaResponseReceived(stopReplicaResponseObj: AbstractResponse, replicaId: Int) extends ControllerEvent {
+  private def processLeaderAndIsrResponseReceived(leaderAndIsrResponse: LeaderAndIsrResponse, brokerId: Int): Unit = {
+    if (!isActive) return
 
-    def state = ControllerState.TopicDeletion
+    if (leaderAndIsrResponse.error != Errors.NONE) {
+      stateChangeLogger.error(s"Received error ${leaderAndIsrResponse.error} in LeaderAndIsr " +
+        s"response $leaderAndIsrResponse from broker $brokerId")
+      return
+    }
 
-    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)))
-      }
+    val offlineReplicas = new ArrayBuffer[TopicPartition]()
+    val onlineReplicas = new ArrayBuffer[TopicPartition]()
+
+    leaderAndIsrResponse.partitions.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 object Startup extends ControllerEvent {
+    val previousOfflineReplicas = controllerContext.replicasOnOfflineDirs.getOrElse(brokerId, Set.empty[TopicPartition])
+    val currentOfflineReplicas = mutable.Set() ++= previousOfflineReplicas --= onlineReplicas ++= offlineReplicas
+    controllerContext.replicasOnOfflineDirs.put(brokerId, currentOfflineReplicas)
+    val newOfflineReplicas = currentOfflineReplicas.diff(previousOfflineReplicas)
 
-    def state = ControllerState.ControllerChange
+    if (newOfflineReplicas.nonEmpty) {
+      stateChangeLogger.info(s"Mark replicas ${newOfflineReplicas.mkString(",")} on broker $brokerId as offline")
+      onReplicasBecomeOffline(newOfflineReplicas.map(PartitionAndReplica(_, brokerId)))
+    }
+  }
 
-    override def process(): Unit = {
-      zkClient.registerZNodeChangeHandlerAndCheckExistence(controllerChangeHandler)
-      elect()
+  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")
+
+    val partitionsInError = if (requestError != Errors.NONE)
+      partitionErrors.keySet
+    else
+      partitionErrors.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 != partitionErrors.size) {
+      // some replicas could have been successfully deleted
+      val deletedReplicas = partitionErrors.keySet.diff(partitionsInError)
+      topicDeletionManager.completeReplicaDeletion(deletedReplicas.map(PartitionAndReplica(_, replicaId)))
     }
+  }
 
+  private def processStartup(): Unit = {
+    zkClient.registerZNodeChangeHandlerAndCheckExistence(controllerChangeHandler)
+    elect()
   }
 
   private def updateMetrics(): Unit = {
@@ -1079,27 +1429,37 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
       if (!isActive) {
         0
       } else {
-        controllerContext.partitionLeadershipInfo.count { case (tp, leadershipInfo) =>
-          !controllerContext.liveOrShuttingDownBrokerIds.contains(leadershipInfo.leaderAndIsr.leader) &&
-            !topicDeletionManager.isTopicQueuedUpForDeletion(tp.topic)
-        }
+        controllerContext.offlinePartitionCount
       }
 
     preferredReplicaImbalanceCount =
       if (!isActive) {
         0
       } else {
-        controllerContext.partitionReplicaAssignment.count { case (topicPartition, replicas) =>
-          val preferredReplica = replicas.head
-          val leadershipInfo = controllerContext.partitionLeadershipInfo.get(topicPartition)
-          leadershipInfo.map(_.leaderAndIsr.leader != preferredReplica).getOrElse(false) &&
-            !topicDeletionManager.isTopicQueuedUpForDeletion(topicPartition.topic)
-        }
+        controllerContext.preferredReplicaImbalanceCount
       }
 
     globalTopicCount = if (!isActive) 0 else controllerContext.allTopics.size
 
-    globalPartitionCount = if (!isActive) 0 else controllerContext.partitionLeadershipInfo.size
+    globalPartitionCount = if (!isActive) 0 else controllerContext.partitionWithLeadersCount
+
+    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
@@ -1112,13 +1472,32 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
   }
 
   private def triggerControllerMove(): Unit = {
-    onControllerResignation()
-    activeControllerId = -1
-    zkClient.deleteController()
+    activeControllerId = zkClient.getControllerId.getOrElse(-1)
+    if (!isActive) {
+      warn("Controller has already moved when trying to trigger controller movement")
+      return
+    }
+    try {
+      val expectedControllerEpochZkVersion = controllerContext.epochZkVersion
+      activeControllerId = -1
+      onControllerResignation()
+      zkClient.deleteController(expectedControllerEpochZkVersion)
+    } catch {
+      case _: ControllerMovedException =>
+        warn("Controller has already moved when trying to trigger controller movement")
+    }
+  }
+
+  private def maybeResign(): Unit = {
+    val wasActiveBeforeChange = isActive
+    zkClient.registerZNodeChangeHandlerAndCheckExistence(controllerChangeHandler)
+    activeControllerId = zkClient.getControllerId.getOrElse(-1)
+    if (wasActiveBeforeChange && !isActive) {
+      onControllerResignation()
+    }
   }
 
   private def elect(): Unit = {
-    val timestamp = time.milliseconds
     activeControllerId = zkClient.getControllerId.getOrElse(-1)
     /*
      * We can get here during the initial startup and the handleDeleted ZK callback. Because of the potential race condition,
@@ -1131,365 +1510,1051 @@ class KafkaController(val config: KafkaConfig, zkClient: KafkaZkClient, time: Ti
     }
 
     try {
-      zkClient.checkedEphemeralCreate(ControllerZNode.path, ControllerZNode.encode(config.brokerId, timestamp))
-      info(s"${config.brokerId} successfully elected as the controller")
+      val (epoch, epochZkVersion) = zkClient.registerControllerAndIncrementControllerEpoch(config.brokerId)
+      controllerContext.epoch = epoch
+      controllerContext.epochZkVersion = epochZkVersion
       activeControllerId = config.brokerId
+
+      info(s"${config.brokerId} successfully elected as the controller. Epoch incremented to ${controllerContext.epoch} " +
+        s"and epoch zk version is now ${controllerContext.epochZkVersion}")
+
       onControllerFailover()
     } catch {
-      case _: NodeExistsException =>
-        // If someone else has written the path, then
-        activeControllerId = zkClient.getControllerId.getOrElse(-1)
+      case e: ControllerMovedException =>
+        maybeResign()
 
         if (activeControllerId != -1)
-          debug(s"Broker $activeControllerId was elected as controller instead of broker ${config.brokerId}")
+          debug(s"Broker $activeControllerId was elected as controller instead of broker ${config.brokerId}", e)
         else
-          warn("A controller has been elected but just resigned, this will result in another round of election")
-
-      case e2: Throwable =>
-        error(s"Error while electing or becoming controller on broker ${config.brokerId}", e2)
+          warn("A controller has been elected but just resigned, this will result in another round of election", e)
+      case t: Throwable =>
+        error(s"Error while electing or becoming controller on broker ${config.brokerId}. " +
+          s"Trigger controller movement immediately", t)
         triggerControllerMove()
     }
   }
 
-  case object BrokerChange extends ControllerEvent {
-    override def state: ControllerState = ControllerState.BrokerChange
+  /**
+   * Partitions the provided map of brokers and epochs into 2 new maps:
+   *  - The first map contains only those brokers whose features were found to be compatible with
+   *    the existing finalized features.
+   *  - The second map contains only those brokers whose features were found to be incompatible with
+   *    the existing finalized features.
+   *
+   * @param brokersAndEpochs   the map to be partitioned
+   * @return                   two maps: first contains compatible brokers and second contains
+   *                           incompatible brokers as explained above
+   */
+  private def partitionOnFeatureCompatibility(brokersAndEpochs: Map[Broker, Long]): (Map[Broker, Long], Map[Broker, Long]) = {
+    // There can not be any feature incompatibilities when the feature versioning system is disabled
+    // or when the finalized feature cache is empty. Otherwise, we check if the non-empty contents
+    //  of the cache are compatible with the supported features of each broker.
+    brokersAndEpochs.partition {
+      case (broker, _) =>
+        !config.isFeatureVersioningSupported ||
+        !featureCache.get.exists(
+          latestFinalizedFeatures =>
+            BrokerFeatures.hasIncompatibleFeatures(broker.features, latestFinalizedFeatures.features))
+    }
+  }
 
-    override def process(): Unit = {
-      if (!isActive) return
-      val curBrokers = zkClient.getAllBrokersInCluster.toSet
-      val curBrokerIds = curBrokers.map(_.id)
-      val liveOrShuttingDownBrokerIds = controllerContext.liveOrShuttingDownBrokerIds
-      val newBrokerIds = curBrokerIds -- liveOrShuttingDownBrokerIds
-      val deadBrokerIds = liveOrShuttingDownBrokerIds -- curBrokerIds
-      val newBrokers = curBrokers.filter(broker => newBrokerIds(broker.id))
-      controllerContext.liveBrokers = curBrokers
-      val newBrokerIdsSorted = newBrokerIds.toSeq.sorted
-      val deadBrokerIdsSorted = deadBrokerIds.toSeq.sorted
-      val liveBrokerIdsSorted = curBrokerIds.toSeq.sorted
-      info(s"Newly added brokers: ${newBrokerIdsSorted.mkString(",")}, " +
-        s"deleted brokers: ${deadBrokerIdsSorted.mkString(",")}, all live brokers: ${liveBrokerIdsSorted.mkString(",")}")
+  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.diff(liveOrShuttingDownBrokerIds)
+    val deadBrokerIds = liveOrShuttingDownBrokerIds.diff(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) {
+      val (newCompatibleBrokerAndEpochs, newIncompatibleBrokerAndEpochs) =
+        partitionOnFeatureCompatibility(newBrokerAndEpochs)
+      if (!newIncompatibleBrokerAndEpochs.isEmpty) {
+        warn("Ignoring registration of new brokers due to incompatibilities with finalized features: " +
+          newIncompatibleBrokerAndEpochs.map { case (broker, _) => broker.id }.toSeq.sorted.mkString(","))
+      }
+      controllerContext.addLiveBrokers(newCompatibleBrokerAndEpochs)
+      onBrokerStartup(newBrokerIdsSorted)
+    }
+    if (bouncedBrokerIds.nonEmpty) {
+      controllerContext.removeLiveBrokers(bouncedBrokerIds)
+      onBrokerFailure(bouncedBrokerIdsSorted)
+      val (bouncedCompatibleBrokerAndEpochs, bouncedIncompatibleBrokerAndEpochs) =
+        partitionOnFeatureCompatibility(bouncedBrokerAndEpochs)
+      if (!bouncedIncompatibleBrokerAndEpochs.isEmpty) {
+        warn("Ignoring registration of bounced brokers due to incompatibilities with finalized features: " +
+          bouncedIncompatibleBrokerAndEpochs.map { case (broker, _) => broker.id }.toSeq.sorted.mkString(","))
+      }
+      controllerContext.addLiveBrokers(bouncedCompatibleBrokerAndEpochs)
+      onBrokerStartup(bouncedBrokerIdsSorted)
+    }
+    if (deadBrokerIds.nonEmpty) {
+      controllerContext.removeLiveBrokers(deadBrokerIds)
+      onBrokerFailure(deadBrokerIdsSorted)
+    }
 
-      newBrokers.foreach(controllerContext.controllerChannelManager.addBroker)
-      deadBrokerIds.foreach(controllerContext.controllerChannelManager.removeBroker)
-      if (newBrokerIds.nonEmpty)
-        onBrokerStartup(newBrokerIdsSorted)
-      if (deadBrokerIds.nonEmpty)
-        onBrokerFailure(deadBrokerIdsSorted)
+    if (newBrokerIds.nonEmpty || deadBrokerIds.nonEmpty || bouncedBrokerIds.nonEmpty) {
+      info(s"Updated broker epochs cache: ${controllerContext.liveBrokerIdAndEpochs}")
     }
   }
 
-  case object TopicChange extends ControllerEvent {
-    override def state: ControllerState = ControllerState.TopicChange
+  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 || !oldMetadata.features.equals(newMetadata.features)) {
+        info(s"Updated broker metadata: $oldMetadata -> $newMetadata")
+        controllerContext.updateBrokerMetadata(oldMetadata, newMetadata)
+        onBrokerUpdate(brokerId)
+      }
+    }
+  }
 
-    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
+  private def processTopicChange(): Unit = {
+    if (!isActive) return
+    val topics = zkClient.getAllTopicsInCluster(true)
+    val newTopics = topics -- controllerContext.allTopics
+    val deletedTopics = controllerContext.allTopics.diff(topics)
+    controllerContext.setAllTopics(topics)
+
+    registerPartitionModificationsHandlers(newTopics.toSeq)
+    val addedPartitionReplicaAssignment = zkClient.getReplicaAssignmentAndTopicIdForTopics(newTopics)
+    deletedTopics.foreach(controllerContext.removeTopic)
+    processTopicIds(addedPartitionReplicaAssignment)
+
+    addedPartitionReplicaAssignment.foreach { case TopicIdReplicaAssignment(_, _, newAssignments) =>
+      newAssignments.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) {
+      val partitionAssignments = addedPartitionReplicaAssignment
+        .map { case TopicIdReplicaAssignment(_, _, partitionsReplicas) => partitionsReplicas.keySet }
+        .reduce((s1, s2) => s1.union(s2))
+      onNewPartitionCreation(partitionAssignments)
+    }
+  }
+
+  private def processTopicIds(topicIdAssignments: Set[TopicIdReplicaAssignment]): Unit = {
+    val updated = zkClient.setTopicIds(topicIdAssignments.filter(_.topicId.isEmpty), controllerContext.epochZkVersion)
+    val allTopicIdAssignments = updated ++ topicIdAssignments.filter(_.topicId.isDefined)
+    allTopicIdAssignments.foreach(topicIdAssignment => controllerContext.addTopicId(topicIdAssignment.topic, topicIdAssignment.topicId.get))
+  }
 
-      registerPartitionModificationsHandlers(newTopics.toSeq)
-      val addedPartitionReplicaAssignment = zkClient.getReplicaAssignmentForTopics(newTopics)
-      controllerContext.partitionReplicaAssignment = controllerContext.partitionReplicaAssignment.filter(p =>
-        !deletedTopics.contains(p._1.topic))
-      controllerContext.partitionReplicaAssignment ++= addedPartitionReplicaAssignment
-      info(s"New topics: [$newTopics], deleted topics: [$deletedTopics], new partition replica assignment " +
-        s"[$addedPartitionReplicaAssignment]")
-      if (addedPartitionReplicaAssignment.nonEmpty)
-        onNewPartitionCreation(addedPartitionReplicaAssignment.keySet)
+  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 object LogDirEventNotification extends ControllerEvent {
-    override def state: ControllerState = ControllerState.LogDirChange
+  private def processPartitionModifications(topic: String): Unit = {
+    def restorePartitionReplicaAssignment(
+      topic: String,
+      newPartitionReplicaAssignment: Map[TopicPartition, ReplicaAssignment]
+    ): 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))
+        .map { case (tp, _) =>
+          tp -> controllerContext.partitionFullReplicaAssignment(tp)
+      }.toMap
+
+      zkClient.setTopicAssignment(topic,
+        controllerContext.topicIds(topic),
+        existingPartitionReplicaAssignment,
+        controllerContext.epochZkVersion)
+    }
 
-    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)
+    if (!isActive) return
+    val partitionReplicaAssignment = zkClient.getFullReplicaAssignmentForTopics(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))
+
+        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.forKeyValue { (topicPartition, assignedReplicas) =>
+        controllerContext.updatePartitionFullReplicaAssignment(topicPartition, assignedReplicas)
       }
+      onNewPartitionCreation(partitionsToBeAdded.keySet)
     }
   }
 
-  case class PartitionModifications(topic: String) extends ControllerEvent {
-    override def state: ControllerState = ControllerState.TopicChange
-
-    override def process(): Unit = {
-      if (!isActive) return
-      val partitionReplicaAssignment = zkClient.getReplicaAssignmentForTopics(immutable.Set(topic))
-      val partitionsToBeAdded = partitionReplicaAssignment.filter(p =>
-        !controllerContext.partitionReplicaAssignment.contains(p._1))
-      if (topicDeletionManager.isTopicQueuedUpForDeletion(topic))
-        error(s"Skipping adding partitions ${partitionsToBeAdded.map(_._1.partition).mkString(",")} for topic $topic " +
-          "since it is currently being deleted")
-      else {
-        if (partitionsToBeAdded.nonEmpty) {
-          info(s"New partitions to be added $partitionsToBeAdded")
-          controllerContext.partitionReplicaAssignment ++= partitionsToBeAdded
-          onNewPartitionCreation(partitionsToBeAdded.keySet)
+  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 (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.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)
     }
   }
 
-  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]
+
+      zkClient.getPartitionReassignment.forKeyValue { (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))
+        }
+      }
 
-    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)
+      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))
       }
-      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)
+      partitionsReassigned.keySet
+    } else {
+      Set.empty
+    }
+  }
+
+  /**
+   * 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.forKeyValue { (tp, targetReplicas) =>
+        val maybeApiError = targetReplicas.flatMap(validateReplicas(tp, _))
+        maybeApiError match {
+          case None =>
+            maybeBuildReassignment(tp, targetReplicas) match {
+              case Some(context) => partitionsToReassign.put(tp, context)
+              case None => reassignmentResults.put(tp, new ApiError(Errors.NO_REASSIGNMENT_IN_PROGRESS))
+            }
+          case Some(err) =>
+            reassignmentResults.put(tp, err)
         }
-      } 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)
       }
+
+      // 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))
+
+      reassignmentResults ++= maybeTriggerPartitionReassignment(partitionsToReassign)
+      callback(Left(reassignmentResults))
     }
   }
 
-  case object PartitionReassignment extends ControllerEvent {
-    override def state: ControllerState = ControllerState.PartitionReassignment
+  private def validateReplicas(topicPartition: TopicPartition, replicas: Seq[Int]): Option[ApiError] = {
+    val replicaSet = replicas.toSet
+    if (replicas.isEmpty)
+      Some(new ApiError(Errors.INVALID_REPLICA_ASSIGNMENT,
+          s"Empty replica list specified in partition reassignment."))
+    else if (replicas.size != replicaSet.size) {
+      Some(new ApiError(Errors.INVALID_REPLICA_ASSIGNMENT,
+          s"Duplicate replica ids in partition reassignment replica list: $replicas"))
+    } else if (replicas.exists(_ < 0))
+      Some(new ApiError(Errors.INVALID_REPLICA_ASSIGNMENT,
+          s"Invalid broker id in replica list: $replicas"))
+    else {
+      // Ensure that any new replicas are among the live brokers
+      val currentAssignment = controllerContext.partitionFullReplicaAssignment(topicPartition)
+      val newAssignment = currentAssignment.reassignTo(replicas)
+      val areNewReplicasAlive = newAssignment.addingReplicas.toSet.subsetOf(controllerContext.liveBrokerIds)
+      if (!areNewReplicasAlive)
+        Some(new ApiError(Errors.INVALID_REPLICA_ASSIGNMENT,
+          s"Replica assignment has brokers that are not alive. Replica list: " +
+            s"${newAssignment.addingReplicas}, live broker list: ${controllerContext.liveBrokerIds}"))
+      else None
+    }
+  }
+
+  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)
+      }
+    }
+  }
 
-    override def process(): Unit = {
-      if (!isActive) return
+  private def processPartitionReassignmentIsrChange(topicPartition: TopicPartition): Unit = {
+    if (!isActive) return
 
-      // 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
-        val partitionsToBeReassigned = partitionReassignment -- controllerContext.partitionsBeingReassigned.keys
-        partitionsToBeReassigned.foreach { case (tp, context) =>
-          if (topicDeletionManager.isTopicQueuedUpForDeletion(tp.topic)) {
-            error(s"Skipping reassignment of $tp since the topic is currently being deleted")
-            removePartitionFromReassignedPartitions(tp)
-          } else {
-            initiateReassignReplicasForTopicPartition(tp, ReassignedPartitionsContext(context))
-          }
+    if (controllerContext.partitionsBeingReassigned.contains(topicPartition)) {
+      maybeCompleteReassignment(topicPartition)
+    }
+  }
+
+  private def maybeCompleteReassignment(topicPartition: TopicPartition): Unit = {
+    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)
+    }
+  }
+
+  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
+      }
+
+      partitionsToList.foreach { tp =>
+        val assignment = controllerContext.partitionFullReplicaAssignment(tp)
+        if (assignment.isBeingReassigned) {
+          results += tp -> assignment
         }
       }
+
+      callback(Left(results))
     }
   }
 
-  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")
+  /**
+   * Returns the new FinalizedVersionRange for the feature, if there are no feature
+   * incompatibilities seen with all known brokers for the provided feature update.
+   * Otherwise returns an ApiError object containing Errors.INVALID_REQUEST.
+   *
+   * @param update   the feature update to be processed (this can not be meant to delete the feature)
+   *
+   * @return         the new FinalizedVersionRange or error, as described above.
+   */
+  private def newFinalizedVersionRangeOrIncompatibilityError(update: UpdateFeaturesRequestData.FeatureUpdateKey): Either[FinalizedVersionRange, ApiError] = {
+    if (UpdateFeaturesRequest.isDeleteRequest(update)) {
+      throw new IllegalArgumentException(s"Provided feature update can not be meant to delete the feature: $update")
+    }
+
+    val supportedVersionRange = brokerFeatures.supportedFeatures.get(update.feature)
+    if (supportedVersionRange == null) {
+      Right(new ApiError(Errors.INVALID_REQUEST,
+                         "Could not apply finalized feature update because the provided feature" +
+                         " is not supported."))
+    } else {
+      var newVersionRange: FinalizedVersionRange = null
+      try {
+        newVersionRange = new FinalizedVersionRange(supportedVersionRange.min, update.maxVersionLevel)
+      } catch {
+        case _: IllegalArgumentException => {
+          // This exception means the provided maxVersionLevel is invalid. It is handled below
+          // outside of this catch clause.
+        }
+      }
+      if (newVersionRange == null) {
+        Right(new ApiError(Errors.INVALID_REQUEST,
+          "Could not apply finalized feature update because the provided" +
+          s" maxVersionLevel:${update.maxVersionLevel} is lower than the" +
+          s" supported minVersion:${supportedVersionRange.min}."))
+      } else {
+        val newFinalizedFeature =
+          Features.finalizedFeatures(Utils.mkMap(Utils.mkEntry(update.feature, newVersionRange)))
+        val numIncompatibleBrokers = controllerContext.liveOrShuttingDownBrokers.count(broker => {
+          BrokerFeatures.hasIncompatibleFeatures(broker.features, newFinalizedFeature)
+        })
+        if (numIncompatibleBrokers == 0) {
+          Left(newVersionRange)
+        } else {
+          Right(new ApiError(Errors.INVALID_REQUEST,
+                             "Could not apply finalized feature update because" +
+                             " brokers were found to have incompatible versions for the feature."))
         }
       }
     }
   }
 
-  case object IsrChangeNotification extends ControllerEvent {
-    override def state: ControllerState = ControllerState.IsrChange
+  /**
+   * Validates a feature update on an existing FinalizedVersionRange.
+   * If the validation succeeds, then, the return value contains:
+   * 1. the new FinalizedVersionRange for the feature, if the feature update was not meant to delete the feature.
+   * 2. Option.empty, if the feature update was meant to delete the feature.
+   *
+   * If the validation fails, then returned value contains a suitable ApiError.
+   *
+   * @param update                 the feature update to be processed.
+   * @param existingVersionRange   the existing FinalizedVersionRange which can be empty when no
+   *                               FinalizedVersionRange exists for the associated feature
+   *
+   * @return                       the new FinalizedVersionRange to be updated into ZK or error
+   *                               as described above.
+   */
+  private def validateFeatureUpdate(update: UpdateFeaturesRequestData.FeatureUpdateKey,
+                                    existingVersionRange: Option[FinalizedVersionRange]): Either[Option[FinalizedVersionRange], ApiError] = {
+    def newVersionRangeOrError(update: UpdateFeaturesRequestData.FeatureUpdateKey): Either[Option[FinalizedVersionRange], ApiError] = {
+      newFinalizedVersionRangeOrIncompatibilityError(update)
+        .fold(versionRange => Left(Some(versionRange)), error => Right(error))
+    }
 
-    override def process(): Unit = {
-      if (!isActive) return
-      val sequenceNumbers = zkClient.getAllIsrChangeNotifications
-      try {
-        val partitions = zkClient.getPartitionsFromIsrChangeNotifications(sequenceNumbers)
-        if (partitions.nonEmpty) {
-          updateLeaderAndIsrCache(partitions)
-          processUpdateNotifications(partitions)
+    if (update.feature.isEmpty) {
+      // Check that the feature name is not empty.
+      Right(new ApiError(Errors.INVALID_REQUEST, "Feature name can not be empty."))
+    } else {
+      // We handle deletion requests separately from non-deletion requests.
+      if (UpdateFeaturesRequest.isDeleteRequest(update)) {
+        if (existingVersionRange.isEmpty) {
+          // Disallow deletion of a non-existing finalized feature.
+          Right(new ApiError(Errors.INVALID_REQUEST,
+                             "Can not delete non-existing finalized feature."))
+        } else {
+          Left(Option.empty)
         }
-      } finally {
-        // delete the notifications
-        zkClient.deleteIsrChangeNotifications(sequenceNumbers)
+      } else if (update.maxVersionLevel() < 1) {
+        // Disallow deletion of a finalized feature without allowDowngrade flag set.
+        Right(new ApiError(Errors.INVALID_REQUEST,
+                           s"Can not provide maxVersionLevel: ${update.maxVersionLevel} less" +
+                           s" than 1 without setting the allowDowngrade flag to true in the request."))
+      } else {
+        existingVersionRange.map(existing =>
+          if (update.maxVersionLevel == existing.max) {
+            // Disallow a case where target maxVersionLevel matches existing maxVersionLevel.
+            Right(new ApiError(Errors.INVALID_REQUEST,
+                               s"Can not ${if (update.allowDowngrade) "downgrade" else "upgrade"}" +
+                               s" a finalized feature from existing maxVersionLevel:${existing.max}" +
+                               " to the same value."))
+          } else if (update.maxVersionLevel < existing.max && !update.allowDowngrade) {
+            // Disallow downgrade of a finalized feature without the allowDowngrade flag set.
+            Right(new ApiError(Errors.INVALID_REQUEST,
+                               s"Can not downgrade finalized feature from existing" +
+                               s" maxVersionLevel:${existing.max} to provided" +
+                               s" maxVersionLevel:${update.maxVersionLevel} without setting the" +
+                               " allowDowngrade flag in the request."))
+          } else if (update.allowDowngrade && update.maxVersionLevel > existing.max) {
+            // Disallow a request that sets allowDowngrade flag without specifying a
+            // maxVersionLevel that's lower than the existing maxVersionLevel.
+            Right(new ApiError(Errors.INVALID_REQUEST,
+                               s"When the allowDowngrade flag set in the request, the provided" +
+                               s" maxVersionLevel:${update.maxVersionLevel} can not be greater than" +
+                               s" existing maxVersionLevel:${existing.max}."))
+          } else if (update.maxVersionLevel < existing.min) {
+            // Disallow downgrade of a finalized feature below the existing finalized
+            // minVersionLevel.
+            Right(new ApiError(Errors.INVALID_REQUEST,
+                               s"Can not downgrade finalized feature to maxVersionLevel:${update.maxVersionLevel}" +
+                               s" because it's lower than the existing minVersionLevel:${existing.min}."))
+          } else {
+            newVersionRangeOrError(update)
+          }
+        ).getOrElse(newVersionRangeOrError(update))
       }
     }
+  }
 
-    private def processUpdateNotifications(partitions: Seq[TopicPartition]) {
+  private def processFeatureUpdates(request: UpdateFeaturesRequest,
+                                    callback: UpdateFeaturesCallback): Unit = {
+    if (isActive) {
+      processFeatureUpdatesWithActiveController(request, callback)
+    } else {
+      callback(Left(new ApiError(Errors.NOT_CONTROLLER)))
+    }
+  }
+
+  private def processFeatureUpdatesWithActiveController(request: UpdateFeaturesRequest,
+                                                        callback: UpdateFeaturesCallback): Unit = {
+    val updates = request.data.featureUpdates
+    val existingFeatures = featureCache.get
+      .map(featuresAndEpoch => featuresAndEpoch.features.features().asScala)
+      .getOrElse(Map[String, FinalizedVersionRange]())
+    // A map with key being feature name and value being FinalizedVersionRange.
+    // This contains the target features to be eventually written to FeatureZNode.
+    val targetFeatures = scala.collection.mutable.Map[String, FinalizedVersionRange]() ++ existingFeatures
+    // A map with key being feature name and value being error encountered when the FeatureUpdate
+    // was applied.
+    val errors = scala.collection.mutable.Map[String, ApiError]()
+
+    // Below we process each FeatureUpdate using the following logic:
+    //  - If a FeatureUpdate is found to be valid, then:
+    //    - The corresponding entry in errors map would be updated to contain Errors.NONE.
+    //    - If the FeatureUpdate is an add or update request, then the targetFeatures map is updated
+    //      to contain the new FinalizedVersionRange for the feature.
+    //    - Otherwise if the FeatureUpdate is a delete request, then the feature is removed from the
+    //      targetFeatures map.
+    //  - Otherwise if a FeatureUpdate is found to be invalid, then:
+    //    - The corresponding entry in errors map would be updated with the appropriate ApiError.
+    //    - The entry in targetFeatures map is left untouched.
+    updates.asScala.iterator.foreach { update =>
+      validateFeatureUpdate(update, existingFeatures.get(update.feature())) match {
+        case Left(newVersionRangeOrNone) =>
+          newVersionRangeOrNone match {
+            case Some(newVersionRange) => targetFeatures += (update.feature() -> newVersionRange)
+            case None => targetFeatures -= update.feature()
+          }
+          errors += (update.feature() -> new ApiError(Errors.NONE))
+        case Right(featureUpdateFailureReason) =>
+          errors += (update.feature() -> featureUpdateFailureReason)
+      }
+    }
+
+    // If the existing and target features are the same, then, we skip the update to the
+    // FeatureZNode as no changes to the node are required. Otherwise, we replace the contents
+    // of the FeatureZNode with the new features. This may result in partial or full modification
+    // of the existing finalized features in ZK.
+    try {
+      if (!existingFeatures.equals(targetFeatures)) {
+        val newNode = new FeatureZNode(FeatureZNodeStatus.Enabled, Features.finalizedFeatures(targetFeatures.asJava))
+        val newVersion = updateFeatureZNode(newNode)
+        featureCache.waitUntilEpochOrThrow(newVersion, request.data().timeoutMs())
+      }
+    } catch {
+      // For all features that correspond to valid FeatureUpdate (i.e. error is Errors.NONE),
+      // we set the error as Errors.FEATURE_UPDATE_FAILED since the FeatureZNode update has failed
+      // for these. For the rest, the existing error is left untouched.
+      case e: Exception =>
+        warn(s"Processing of feature updates: $request failed due to error: $e")
+        errors.foreach { case (feature, apiError) =>
+          if (apiError.error() == Errors.NONE) {
+            errors(feature) = new ApiError(Errors.FEATURE_UPDATE_FAILED)
+          }
+        }
+    } finally {
+      callback(Right(errors))
+    }
+  }
+
+  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)
+
+        // During a partial upgrade, the controller may be on an IBP which assumes
+        // ISR changes through the `AlterIsr` API while some brokers are on an older
+        // IBP which assumes notification through Zookeeper. In this case, since the
+        // controller will not have registered watches for reassigning partitions, we
+        // can still rely on the batch ISR change notification path in order to
+        // complete the reassignment.
+        partitions.filter(controllerContext.partitionsBeingReassigned.contains).foreach { topicPartition =>
+          maybeCompleteReassignment(topicPartition)
+        }
+      }
+    } finally {
+      // delete the notifications
+      zkClient.deleteIsrChangeNotifications(sequenceNumbers, controllerContext.epochZkVersion)
+    }
   }
 
-  case object PreferredReplicaLeaderElection extends ControllerEvent {
-    override def state: ControllerState = ControllerState.ManualLeaderBalance
+  def electLeaders(
+    partitions: Set[TopicPartition],
+    electionType: ElectionType,
+    callback: ElectLeadersCallback
+  ): Unit = {
+    eventManager.put(ReplicaLeaderElection(Some(partitions), electionType, AdminClientTriggered, callback))
+  }
 
-    override def process(): Unit = {
-      if (!isActive) return
+  def listPartitionReassignments(partitions: Option[Set[TopicPartition]],
+                                 callback: ListReassignmentsCallback): Unit = {
+    eventManager.put(ListPartitionReassignments(partitions, callback))
+  }
 
+  def updateFeatures(request: UpdateFeaturesRequest,
+                     callback: UpdateFeaturesCallback): Unit = {
+    eventManager.put(UpdateFeatures(request, callback))
+  }
+
+  def alterPartitionReassignments(partitions: Map[TopicPartition, Option[Seq[Int]]],
+                                  callback: AlterReassignmentsCallback): Unit = {
+    eventManager.put(ApiPartitionReassignment(partitions, callback))
+  }
+
+  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 (zkClient.registerZNodeChangeHandlerAndCheckExistence(preferredReplicaElectionHandler)) {
-        val partitions = zkClient.getPreferredReplicaElection
-        val partitionsForTopicsToBeDeleted = partitions.filter(p => topicDeletionManager.isTopicQueuedUpForDeletion(p.topic))
-        if (partitionsForTopicsToBeDeleted.nonEmpty) {
-          error(s"Skipping preferred replica election for partitions $partitionsForTopicsToBeDeleted since the " +
-            "respective topics are being deleted")
+      if (electionTrigger == AdminClientTriggered || zkClient.registerZNodeChangeHandlerAndCheckExistence(preferredReplicaElectionHandler)) {
+        val partitions = partitionsFromAdminClientOpt match {
+          case Some(partitions) => partitions
+          case None => zkClient.getPreferredReplicaElection
+        }
+
+        val allPartitions = controllerContext.allPartitions
+        val (knownPartitions, unknownPartitions) = partitions.partition(tp => allPartitions.contains(tp))
+        unknownPartitions.foreach { p =>
+          info(s"Skipping replica leader election ($electionType) for partition $p by $electionTrigger since it doesn't exist.")
         }
-        onPreferredReplicaElection(partitions -- partitionsForTopicsToBeDeleted)
+
+        val (partitionsBeingDeleted, livePartitions) = knownPartitions.partition(partition =>
+            topicDeletionManager.isTopicQueuedUpForDeletion(partition.topic))
+        if (partitionsBeingDeleted.nonEmpty) {
+          warn(s"Skipping replica leader election ($electionType) for partitions $partitionsBeingDeleted " +
+            s"by $electionTrigger since the respective topics are being deleted")
+        }
+
+        // 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).get.leaderAndIsr.leader
+              currentLeader != preferredReplica
+
+            case ElectionType.UNCLEAN =>
+              val currentLeader = controllerContext.partitionLeadershipInfo(partition).get.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
+  def alterIsrs(alterIsrRequest: AlterIsrRequestData, callback: AlterIsrResponseData => Unit): Unit = {
+    val isrsToAlter = mutable.Map[TopicPartition, LeaderAndIsr]()
+
+    alterIsrRequest.topics.forEach { topicReq =>
+      topicReq.partitions.forEach { partitionReq =>
+        val tp = new TopicPartition(topicReq.name, partitionReq.partitionIndex)
+        val newIsr = partitionReq.newIsr().asScala.toList.map(_.toInt)
+        isrsToAlter.put(tp, new LeaderAndIsr(alterIsrRequest.brokerId, partitionReq.leaderEpoch, newIsr, partitionReq.currentIsrVersion))
+      }
+    }
 
-    override def process(): Unit = {
-      val wasActiveBeforeChange = isActive
-      zkClient.registerZNodeChangeHandlerAndCheckExistence(controllerChangeHandler)
-      activeControllerId = zkClient.getControllerId.getOrElse(-1)
-      if (wasActiveBeforeChange && !isActive) {
-        onControllerResignation()
+    def responseCallback(results: Either[Map[TopicPartition, Either[Errors, LeaderAndIsr]], Errors]): Unit = {
+      val resp = new AlterIsrResponseData()
+      results match {
+        case Right(error) =>
+          resp.setErrorCode(error.code)
+        case Left(partitionResults) =>
+          resp.setTopics(new util.ArrayList())
+          partitionResults
+            .groupBy { case (tp, _) => tp.topic }   // Group by topic
+            .foreach { case (topic, partitions) =>
+              // Add each topic part to the response
+              val topicResp = new AlterIsrResponseData.TopicData()
+                .setName(topic)
+                .setPartitions(new util.ArrayList())
+              resp.topics.add(topicResp)
+              partitions.foreach { case (tp, errorOrIsr) =>
+                // Add each partition part to the response (new ISR or error)
+                errorOrIsr match {
+                  case Left(error) => topicResp.partitions.add(
+                    new AlterIsrResponseData.PartitionData()
+                      .setPartitionIndex(tp.partition)
+                      .setErrorCode(error.code))
+                  case Right(leaderAndIsr) => topicResp.partitions.add(
+                    new AlterIsrResponseData.PartitionData()
+                      .setPartitionIndex(tp.partition)
+                      .setLeaderId(leaderAndIsr.leader)
+                      .setLeaderEpoch(leaderAndIsr.leaderEpoch)
+                      .setIsr(leaderAndIsr.isr.map(Integer.valueOf).asJava)
+                      .setCurrentIsrVersion(leaderAndIsr.zkVersion))
+                }
+            }
+          }
       }
+      callback.apply(resp)
     }
+
+    eventManager.put(AlterIsrReceived(alterIsrRequest.brokerId, alterIsrRequest.brokerEpoch, isrsToAlter, responseCallback))
   }
 
-  case object Reelect extends ControllerEvent {
-    override def state = ControllerState.ControllerChange
+  private def processAlterIsr(brokerId: Int, brokerEpoch: Long,
+                              isrsToAlter: Map[TopicPartition, LeaderAndIsr],
+                              callback: AlterIsrCallback): Unit = {
+
+    // Handle a few short-circuits
+    if (!isActive) {
+      callback.apply(Right(Errors.NOT_CONTROLLER))
+      return
+    }
+
+    val brokerEpochOpt = controllerContext.liveBrokerIdAndEpochs.get(brokerId)
+    if (brokerEpochOpt.isEmpty) {
+      info(s"Ignoring AlterIsr due to unknown broker $brokerId")
+      callback.apply(Right(Errors.STALE_BROKER_EPOCH))
+      return
+    }
+
+    if (!brokerEpochOpt.contains(brokerEpoch)) {
+      info(s"Ignoring AlterIsr due to stale broker epoch $brokerEpoch for broker $brokerId")
+      callback.apply(Right(Errors.STALE_BROKER_EPOCH))
+      return
+    }
+
+    val response = try {
+      val partitionResponses = mutable.HashMap[TopicPartition, Either[Errors, LeaderAndIsr]]()
+
+      // Determine which partitions we will accept the new ISR for
+      val adjustedIsrs: Map[TopicPartition, LeaderAndIsr] = isrsToAlter.flatMap {
+        case (tp: TopicPartition, newLeaderAndIsr: LeaderAndIsr) =>
+          controllerContext.partitionLeadershipInfo(tp) match {
+            case Some(leaderIsrAndControllerEpoch) =>
+              val currentLeaderAndIsr = leaderIsrAndControllerEpoch.leaderAndIsr
+              if (newLeaderAndIsr.leaderEpoch < currentLeaderAndIsr.leaderEpoch) {
+                partitionResponses(tp) = Left(Errors.FENCED_LEADER_EPOCH)
+                None
+              } else if (newLeaderAndIsr.equalsIgnoreZk(currentLeaderAndIsr)) {
+                // If a partition is already in the desired state, just return it
+                partitionResponses(tp) = Right(currentLeaderAndIsr)
+                None
+              } else {
+                Some(tp -> newLeaderAndIsr)
+              }
+            case None =>
+              partitionResponses(tp) = Left(Errors.UNKNOWN_TOPIC_OR_PARTITION)
+              None
+          }
+      }
+
+      // Do the updates in ZK
+      debug(s"Updating ISRs for partitions: ${adjustedIsrs.keySet}.")
+      val UpdateLeaderAndIsrResult(finishedUpdates, badVersionUpdates) = zkClient.updateLeaderAndIsr(
+        adjustedIsrs, controllerContext.epoch, controllerContext.epochZkVersion)
+
+      val successfulUpdates: Map[TopicPartition, LeaderAndIsr] = finishedUpdates.flatMap {
+        case (partition: TopicPartition, isrOrError: Either[Throwable, LeaderAndIsr]) =>
+          isrOrError match {
+            case Right(updatedIsr) =>
+              debug(s"ISR for partition $partition updated to [${updatedIsr.isr.mkString(",")}] and zkVersion updated to [${updatedIsr.zkVersion}]")
+              partitionResponses(partition) = Right(updatedIsr)
+              Some(partition -> updatedIsr)
+            case Left(error) =>
+              warn(s"Failed to update ISR for partition $partition", error)
+              partitionResponses(partition) = Left(Errors.forException(error))
+              None
+          }
+      }
+
+      badVersionUpdates.foreach(partition => {
+        debug(s"Failed to update ISR for partition $partition, bad ZK version")
+        partitionResponses(partition) = Left(Errors.INVALID_UPDATE_VERSION)
+      })
+
+      def processUpdateNotifications(partitions: Seq[TopicPartition]): Unit = {
+        val liveBrokers: Seq[Int] = controllerContext.liveOrShuttingDownBrokerIds.toSeq
+        sendUpdateMetadataRequest(liveBrokers, partitions.toSet)
+      }
+
+      // Update our cache and send out metadata updates
+      updateLeaderAndIsrCache(successfulUpdates.keys.toSeq)
+      processUpdateNotifications(isrsToAlter.keys.toSeq)
+
+      Left(partitionResponses)
+    } catch {
+      case e: Throwable =>
+        error(s"Error when processing AlterIsr for partitions: ${isrsToAlter.keys.toSeq}", e)
+        Right(Errors.UNKNOWN_SERVER_ERROR)
+    }
+
+    callback.apply(response)
 
-    override def process(): Unit = {
-      val wasActiveBeforeChange = isActive
-      zkClient.registerZNodeChangeHandlerAndCheckExistence(controllerChangeHandler)
-      activeControllerId = zkClient.getControllerId.getOrElse(-1)
-      if (wasActiveBeforeChange && !isActive) {
-        onControllerResignation()
+    // After we have returned the result of the `AlterIsr` request, we should check whether
+    // there are any reassignments which can be completed by a successful ISR expansion.
+    response.left.foreach { alterIsrResponses =>
+      alterIsrResponses.forKeyValue { (topicPartition, partitionResponse) =>
+        if (controllerContext.partitionsBeingReassigned.contains(topicPartition)) {
+          val isSuccessfulUpdate = partitionResponse.isRight
+          if (isSuccessfulUpdate) {
+            maybeCompleteReassignment(topicPartition)
+          }
+        }
       }
-      elect()
     }
   }
 
-  // We can't make this a case object due to the countDownLatch field
-  class Expire extends ControllerEvent {
-    private val countDownLatch = new CountDownLatch(1)
-    override def state = ControllerState.ControllerChange
+  private def processControllerChange(): Unit = {
+    maybeResign()
+  }
 
-    override def process(): Unit = {
-      countDownLatch.countDown()
-      activeControllerId = -1
-      onControllerResignation()
-    }
+  private def processReelect(): Unit = {
+    maybeResign()
+    elect()
+  }
+
+  private def processRegisterBrokerAndReelect(): Unit = {
+    _brokerEpoch = zkClient.registerBroker(brokerInfo)
+    processReelect()
+  }
 
-    def waitUntilProcessed(): Unit = {
-      countDownLatch.await()
+  private def processExpire(): Unit = {
+    activeControllerId = -1
+    onControllerResignation()
+  }
+
+
+  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 UpdateMetadataResponseReceived(response, brokerId) =>
+          processUpdateMetadataResponseReceived(response, brokerId)
+        case TopicDeletionStopReplicaResponseReceived(replicaId, requestError, partitionErrors) =>
+          processTopicDeletionStopReplicaResponseReceived(replicaId, requestError, partitionErrors)
+        case BrokerChange =>
+          processBrokerChange()
+        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 UpdateFeatures(request, callback) =>
+          processFeatureUpdates(request, callback)
+        case PartitionReassignmentIsrChange(partition) =>
+          processPartitionReassignmentIsrChange(partition)
+        case IsrChangeNotification =>
+          processIsrChangeNotification()
+        case AlterIsrReceived(brokerId, brokerEpoch, isrsToAlter, callback) =>
+          processAlterIsr(brokerId, brokerEpoch, isrsToAlter, callback)
+        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.preempt()
+  }
 }
 
-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)
+  override def handleChildChange(): Unit = {
+    eventManager.put(BrokerChange)
+  }
+}
+
+class BrokerModificationsHandler(eventManager: ControllerEventManager, brokerId: Int) extends ZNodeChangeHandler {
+  override val path: String = BrokerIdZNode.path(brokerId)
+
+  override def handleDataChange(): Unit = {
+    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)
 }
 
-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)
+  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)
+  override def handleCreation(): Unit = eventManager.put(ControllerChange)
+  override def handleDeletion(): Unit = eventManager.put(Reelect)
+  override def handleDataChange(): Unit = eventManager.put(ControllerChange)
 }
 
-case class ReassignedPartitionsContext(var newReplicas: Seq[Int] = Seq.empty,
-                                       var reassignIsrChangeHandler: PartitionReassignmentIsrChangeHandler = null)
-
 case class PartitionAndReplica(topicPartition: TopicPartition, replica: Int) {
   def topic: String = topicPartition.topic
   def partition: Int = topicPartition.partition
@@ -1523,5 +2588,167 @@ private[controller] class ControllerStats extends KafkaMetricsGroup {
 
 sealed trait ControllerEvent {
   def state: ControllerState
+  // preempt() is not executed by `ControllerEventThread` but by the main thread.
+  def preempt(): Unit
+}
+
+case object ControllerChange extends ControllerEvent {
+  override def state: ControllerState = ControllerState.ControllerChange
+  override def preempt(): Unit = {}
+}
+
+case object Reelect extends ControllerEvent {
+  override def state: ControllerState = ControllerState.ControllerChange
+  override def preempt(): Unit = {}
+}
+
+case object RegisterBrokerAndReelect extends ControllerEvent {
+  override def state: ControllerState = ControllerState.ControllerChange
+  override def preempt(): Unit = {}
+}
+
+case object Expire extends ControllerEvent {
+  override def state: ControllerState = ControllerState.ControllerChange
+  override def preempt(): Unit = {}
+}
+
+case object ShutdownEventThread extends ControllerEvent {
+  override def state: ControllerState = ControllerState.ControllerShutdown
+  override def preempt(): Unit = {}
+}
+
+case object AutoPreferredReplicaLeaderElection extends ControllerEvent {
+  override def state: ControllerState = ControllerState.AutoLeaderBalance
+  override def preempt(): Unit = {}
+}
+
+case object UncleanLeaderElectionEnable extends ControllerEvent {
+  override def state: ControllerState = ControllerState.UncleanLeaderElectionEnable
+  override def preempt(): Unit = {}
+}
+
+case class TopicUncleanLeaderElectionEnable(topic: String) extends ControllerEvent {
+  override def state: ControllerState = ControllerState.TopicUncleanLeaderElectionEnable
+  override def preempt(): Unit = {}
+}
+
+case class ControlledShutdown(id: Int, brokerEpoch: Long, controlledShutdownCallback: Try[Set[TopicPartition]] => Unit) extends ControllerEvent {
+  override def state: ControllerState = ControllerState.ControlledShutdown
+  override def preempt(): Unit = controlledShutdownCallback(Failure(new ControllerMovedException("Controller moved to another broker")))
+}
+
+case class LeaderAndIsrResponseReceived(leaderAndIsrResponse: LeaderAndIsrResponse, brokerId: Int) extends ControllerEvent {
+  override def state: ControllerState = ControllerState.LeaderAndIsrResponseReceived
+  override def preempt(): Unit = {}
+}
+
+case class UpdateMetadataResponseReceived(updateMetadataResponse: UpdateMetadataResponse, brokerId: Int) extends ControllerEvent {
+  override def state: ControllerState = ControllerState.UpdateMetadataResponseReceived
+  override def preempt(): Unit = {}
+}
+
+case class TopicDeletionStopReplicaResponseReceived(replicaId: Int,
+                                                    requestError: Errors,
+                                                    partitionErrors: Map[TopicPartition, Errors]) extends ControllerEvent {
+  override def state: ControllerState = ControllerState.TopicDeletion
+  override def preempt(): Unit = {}
+}
+
+case object Startup extends ControllerEvent {
+  override def state: ControllerState = ControllerState.ControllerChange
+  override def preempt(): Unit = {}
+}
+
+case object BrokerChange extends ControllerEvent {
+  override def state: ControllerState = ControllerState.BrokerChange
+  override def preempt(): Unit = {}
+}
+
+case class BrokerModifications(brokerId: Int) extends ControllerEvent {
+  override def state: ControllerState = ControllerState.BrokerChange
+  override def preempt(): Unit = {}
+}
+
+case object TopicChange extends ControllerEvent {
+  override def state: ControllerState = ControllerState.TopicChange
+  override def preempt(): Unit = {}
+}
+
+case object LogDirEventNotification extends ControllerEvent {
+  override def state: ControllerState = ControllerState.LogDirChange
+  override def preempt(): Unit = {}
+}
+
+case class PartitionModifications(topic: String) extends ControllerEvent {
+  override def state: ControllerState = ControllerState.TopicChange
+  override def preempt(): Unit = {}
+}
+
+case object TopicDeletion extends ControllerEvent {
+  override def state: ControllerState = ControllerState.TopicDeletion
+  override def preempt(): Unit = {}
+}
+
+case object ZkPartitionReassignment extends ControllerEvent {
+  override def state: ControllerState = ControllerState.AlterPartitionReassignment
+  override def preempt(): Unit = {}
+}
+
+case class ApiPartitionReassignment(reassignments: Map[TopicPartition, Option[Seq[Int]]],
+                                    callback: AlterReassignmentsCallback) extends ControllerEvent {
+  override def state: ControllerState = ControllerState.AlterPartitionReassignment
+  override def preempt(): Unit = callback(Right(new ApiError(Errors.NOT_CONTROLLER)))
+}
+
+case class PartitionReassignmentIsrChange(partition: TopicPartition) extends ControllerEvent {
+  override def state: ControllerState = ControllerState.AlterPartitionReassignment
+  override def preempt(): Unit = {}
+}
+
+case object IsrChangeNotification extends ControllerEvent {
+  override def state: ControllerState = ControllerState.IsrChange
+  override def preempt(): Unit = {}
+}
+
+case class AlterIsrReceived(brokerId: Int, brokerEpoch: Long, isrsToAlter: Map[TopicPartition, LeaderAndIsr],
+                            callback: AlterIsrCallback) extends ControllerEvent {
+  override def state: ControllerState = ControllerState.IsrChange
+  override def preempt(): Unit = {}
+}
+
+case class ReplicaLeaderElection(
+  partitionsFromAdminClientOpt: Option[Set[TopicPartition]],
+  electionType: ElectionType,
+  electionTrigger: ElectionTrigger,
+  callback: ElectLeadersCallback = _ => {}
+) extends ControllerEvent {
+  override def state: ControllerState = ControllerState.ManualLeaderBalance
+
+  override def preempt(): Unit = callback(
+    partitionsFromAdminClientOpt.fold(Map.empty[TopicPartition, Either[ApiError, Int]]) { partitions =>
+      partitions.iterator.map(partition => partition -> Left(new ApiError(Errors.NOT_CONTROLLER, null))).toMap
+    }
+  )
+}
+
+/**
+  * @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
+  override def preempt(): Unit = callback(Right(new ApiError(Errors.NOT_CONTROLLER, null)))
+}
+
+case class UpdateFeatures(request: UpdateFeaturesRequest,
+                          callback: UpdateFeaturesCallback) extends ControllerEvent {
+  override def state: ControllerState = ControllerState.UpdateFeatures
+  override def preempt(): Unit = {}
+}
+
+
+// Used only in test cases
+abstract class MockEvent(val state: ControllerState) extends ControllerEvent {
   def process(): Unit
+  def preempt(): Unit
 }
diff --git a/core/src/main/scala/kafka/controller/PartitionStateMachine.scala b/core/src/main/scala/kafka/controller/PartitionStateMachine.scala
index 2e27272158f4c..17d6fde95ced1 100755
--- a/core/src/main/scala/kafka/controller/PartitionStateMachine.scala
+++ b/core/src/main/scala/kafka/controller/PartitionStateMachine.scala
@@ -18,111 +18,162 @@ package kafka.controller
 
 import kafka.api.LeaderAndIsr
 import kafka.common.StateChangeFailedException
+import kafka.controller.Election._
 import kafka.server.KafkaConfig
+import kafka.utils.Implicits._
 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,
-                            topicDeletionManager: TopicDeletionManager,
-                            zkClient: KafkaZkClient,
-                            partitionState: mutable.Map[TopicPartition, PartitionState],
-                            controllerBrokerRequestBatch: ControllerBrokerRequestBatch) extends Logging {
-  private val controllerId = config.brokerId
-
-  this.logIdent = s"[PartitionStateMachine controllerId=$controllerId] "
-
+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()
+  def shutdown(): Unit = {
     info("Stopped partition state machine")
   }
 
+  /**
+   * 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() {
-    for (topicPartition <- controllerContext.partitionReplicaAssignment.keys) {
+  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 {
+      controllerContext.partitionLeadershipInfo(topicPartition) match {
         case Some(currentLeaderIsrAndEpoch) =>
           // 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
-            partitionState.put(topicPartition, OnlinePartition)
+            controllerContext.putPartitionState(topicPartition, OnlinePartition)
           else
-            partitionState.put(topicPartition, OfflinePartition)
+            controllerContext.putPartitionState(topicPartition, OfflinePartition)
         case None =>
-          partitionState.put(topicPartition, 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() {
-    // 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.
+  def handleStateChanges(
+    partitions: Seq[TopicPartition],
+    targetState: PartitionState
+  ): Map[TopicPartition, Either[Throwable, LeaderAndIsr]] = {
+    handleStateChanges(partitions, targetState, None)
   }
 
-  def handleStateChanges(partitions: Seq[TopicPartition], targetState: PartitionState,
-                         partitionLeaderElectionStrategyOpt: Option[PartitionLeaderElectionStrategy] = None): Unit = {
+  def handleStateChanges(
+    partitions: Seq[TopicPartition],
+    targetState: PartitionState,
+    leaderElectionStrategy: Option[PartitionLeaderElectionStrategy]
+  ): Map[TopicPartition, Either[Throwable, LeaderAndIsr]]
+
+}
+
+/**
+ * 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 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()
-        doHandleStateChanges(partitions, targetState, partitionLeaderElectionStrategyOpt)
+        val result = doHandleStateChanges(
+          partitions,
+          targetState,
+          partitionLeaderElectionStrategyOpt
+        )
         controllerBrokerRequestBatch.sendRequestsToBrokers(controllerContext.epoch)
+        result
       } catch {
-        case e: Throwable => error(s"Error while moving some partitions to $targetState state", e)
+        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.iterator.map(_ -> Left(e)).toMap
       }
+    } else {
+      Map.empty
     }
   }
 
-  def partitionsInState(state: PartitionState): Set[TopicPartition] = {
-    partitionState.filter { case (_, s) => s == state }.keySet.toSet
+  private def partitionState(partition: TopicPartition): PartitionState = {
+    controllerContext.partitionState(partition)
   }
 
   /**
@@ -146,49 +197,68 @@ 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]): Unit = {
+  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))
+    val traceEnabled = stateChangeLog.isTraceEnabled
+    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 " +
+          stateChangeLog.info(s"Changed partition $partition state from ${partitionState(partition)} to $targetState with " +
             s"assigned replicas ${controllerContext.partitionReplicaAssignment(partition).mkString(",")}")
-          partitionState.put(partition, NewPartition)
+          controllerContext.putPartitionState(partition, NewPartition)
         }
+        Map.empty
       case OnlinePartition =>
         val uninitializedPartitions = validPartitions.filter(partition => partitionState(partition) == NewPartition)
         val partitionsToElectLeader = validPartitions.filter(partition => partitionState(partition) == OfflinePartition || partitionState(partition) == OnlinePartition)
         if (uninitializedPartitions.nonEmpty) {
           val successfulInitializations = initializeLeaderAndIsrForPartitions(uninitializedPartitions)
           successfulInitializations.foreach { partition =>
-            stateChangeLog.trace(s"Changed partition $partition from ${partitionState(partition)} to $targetState with state " +
-              s"${controllerContext.partitionLeadershipInfo(partition).leaderAndIsr}")
-            partitionState.put(partition, OnlinePartition)
+            stateChangeLog.info(s"Changed partition $partition from ${partitionState(partition)} to $targetState with state " +
+              s"${controllerContext.partitionLeadershipInfo(partition).get.leaderAndIsr}")
+            controllerContext.putPartitionState(partition, OnlinePartition)
           }
         }
         if (partitionsToElectLeader.nonEmpty) {
-          val successfulElections = 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}")
-            partitionState.put(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.info(
+                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
           }
+
+          electionResults
+        } else {
+          Map.empty
         }
-      case OfflinePartition =>
-        validPartitions.foreach { partition =>
-          stateChangeLog.trace(s"Changed partition $partition state from ${partitionState(partition)} to $targetState")
-          partitionState.put(partition, OfflinePartition)
-        }
-      case NonExistentPartition =>
+      case OfflinePartition | NonExistentPartition =>
         validPartitions.foreach { partition =>
-          stateChangeLog.trace(s"Changed partition $partition state from ${partitionState(partition)} to $targetState")
-          partitionState.put(partition, NonExistentPartition)
+          if (traceEnabled)
+            stateChangeLog.trace(s"Changed partition $partition state from ${partitionState(partition)} to $targetState")
+          controllerContext.putPartitionState(partition, targetState)
         }
+        Map.empty
     }
   }
 
@@ -219,10 +289,13 @@ class PartitionStateMachine(config: KafkaConfig,
       partition -> leaderIsrAndControllerEpoch
     }.toMap
     val createResponses = try {
-      zkClient.createTopicPartitionStatesRaw(leaderIsrAndControllerEpochs)
+      zkClient.createTopicPartitionStatesRaw(leaderIsrAndControllerEpochs, controllerContext.epochZkVersion)
     } catch {
+      case e: ControllerMovedException =>
+        error("Controller moved to another broker when trying to create the topic partition state znode", e)
+        throw e
       case e: Exception =>
-        partitionsWithLiveReplicas.foreach { case (partition,_) => logFailedStateChange(partition, partitionState(partition), NewPartition, e) }
+        partitionsWithLiveReplicas.foreach { case (partition, _) => logFailedStateChange(partition, partitionState(partition), NewPartition, e) }
         Seq.empty
     }
     createResponses.foreach { createResponse =>
@@ -230,9 +303,9 @@ class PartitionStateMachine(config: KafkaConfig,
       val partition = createResponse.ctx.get.asInstanceOf[TopicPartition]
       val leaderIsrAndControllerEpoch = leaderIsrAndControllerEpochs(partition)
       if (code == Code.OK) {
-        controllerContext.partitionLeadershipInfo.put(partition, leaderIsrAndControllerEpoch)
+        controllerContext.putPartitionLeadershipInfo(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)
@@ -245,20 +318,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 The partitions that successfully had a leader elected.
+   * @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] = {
-    val successfulElections = mutable.Buffer.empty[TopicPartition]
+  private def electLeaderForPartitions(
+    partitions: Seq[TopicPartition],
+    partitionLeaderElectionStrategy: PartitionLeaderElectionStrategy
+  ): Map[TopicPartition, Either[Throwable, LeaderAndIsr]] = {
     var remaining = partitions
+    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
       }
+
+      finishedElections ++= finished
+
+      if (remaining.nonEmpty)
+        logger.info(s"Retrying leader election with strategy $partitionLeaderElectionStrategy for partitions $remaining")
     }
-    successfulElections
+
+    finishedElections.toMap
   }
 
   /**
@@ -267,154 +353,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 (partition, 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, leaderAndIsrOpt, recipients) =>
+    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, leaderAndIsrOpt, recipients) => partition -> recipients }.toMap
-    val adjustedLeaderAndIsrs = partitionsWithLeaders.map { case (partition, leaderAndIsrOpt, recipients) => partition -> leaderAndIsrOpt.get }.toMap
-    val UpdateLeaderAndIsrResult(successfulUpdates, updatesToRetry, failedUpdates) = zkClient.updateLeaderAndIsr(
-      adjustedLeaderAndIsrs, controllerContext.epoch)
-    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)
+    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)
+    finishedUpdates.forKeyValue { (partition, result) =>
+      result.foreach { leaderAndIsr =>
+        val replicaAssignment = controllerContext.partitionFullReplicaAssignment(partition)
+        val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(leaderAndIsr, controllerContext.epoch)
+        controllerContext.putPartitionLeadershipInfo(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)
-        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)
+    electionForPartitionWithoutLiveReplicas ++
+    partitionsWithLiveInSyncReplicas.map { case (partition, leaderAndIsr) =>
+      (partition, Option(leaderAndIsr), false)
     }
   }
 
-  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 liveReplicas = assignment.filter(replica => controllerContext.isReplicaOnline(replica, partition))
-      val isr = leaderIsrAndControllerEpoch.leaderAndIsr.isr
-      val leaderOpt = PartitionLeaderElectionAlgorithms.controlledShutdownPartitionLeaderElection(assignment, isr, liveReplicas.toSet, shuttingDownBrokers)
-      val newIsr = isr.filter(replica => !controllerContext.shuttingDownBrokerIds.contains(replica))
-      val newLeaderAndIsrOpt = leaderOpt.map(leader => leaderIsrAndControllerEpoch.leaderAndIsr.newLeaderAndIsr(leader, newIsr))
-      (partition, newLeaderAndIsrOpt, liveReplicas)
-    }
-  }
-
-  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 " +
@@ -435,10 +513,13 @@ class PartitionStateMachine(config: KafkaConfig,
 }
 
 object PartitionLeaderElectionAlgorithms {
-  def offlinePartitionLeaderElection(assignment: Seq[Int], isr: Seq[Int], liveReplicas: Set[Int], uncleanLeaderElectionEnabled: Boolean): Option[Int] = {
+  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) {
-        assignment.find(liveReplicas.contains)
+        val leaderOpt = assignment.find(liveReplicas.contains)
+        if (leaderOpt.isDefined)
+          controllerContext.stats.uncleanLeaderElectionRate.mark()
+        leaderOpt
       } else {
         None
       }
@@ -459,10 +540,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 85af764b635e2..f56d28dade458 100644
--- a/core/src/main/scala/kafka/controller/ReplicaStateMachine.scala
+++ b/core/src/main/scala/kafka/controller/ReplicaStateMachine.scala
@@ -19,59 +19,35 @@ package kafka.controller
 import kafka.api.LeaderAndIsr
 import kafka.common.StateChangeFailedException
 import kafka.server.KafkaConfig
+import kafka.utils.Implicits._
 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")
   }
 
@@ -79,32 +55,67 @@ 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() {
-    controllerContext.partitionReplicaAssignment.foreach { case (partition, replicas) =>
+  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, OfflineReplica and ReplicaDeletionIneligible
+ * 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, OfflineReplica and ReplicaDeletionIneligible
+ * 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).forKeyValue { (replicaId, replicas) =>
+          doHandleStateChanges(replicaId, replicas, targetState)
         }
         controllerBrokerRequestBatch.sendRequestsToBrokers(controllerContext.epoch)
       } catch {
+        case e: ControllerMovedException =>
+          error(s"Controller moved to another broker when moving some replicas to $targetState state", e)
+          throw e
         case e: Throwable => error(s"Error while moving some replicas to $targetState state", e)
       }
     }
@@ -142,104 +153,134 @@ 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 = {
+    val stateLogger = stateChangeLogger.withControllerEpoch(controllerContext.epoch)
+    val traceEnabled = stateLogger.isTraceEnabled
+    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
-          controllerContext.partitionLeadershipInfo.get(partition) match {
+          val currentState = controllerContext.replicaState(replica)
+
+          controllerContext.partitionLeadershipInfo(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)
+                if (traceEnabled)
+                  logSuccessfulTransition(stateLogger, replicaId, partition, currentState, NewReplica)
+                controllerContext.putReplicaState(replica, NewReplica)
               }
             case None =>
-              logSuccessfulTransition(replicaId, partition, replicaState(replica), NewReplica)
-              replicaState.put(replica, NewReplica)
+              if (traceEnabled)
+                logSuccessfulTransition(stateLogger, 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)) {
-                controllerContext.partitionReplicaAssignment.put(partition, assignment :+ replicaId)
+              val assignment = controllerContext.partitionFullReplicaAssignment(partition)
+              if (!assignment.replicas.contains(replicaId)) {
+                error(s"Adding replica ($replicaId) that is not part of the assignment $assignment")
+                val newAssignment = assignment.copy(replicas = assignment.replicas :+ replicaId)
+                controllerContext.updatePartitionFullReplicaAssignment(partition, newAssignment)
               }
             case _ =>
-              controllerContext.partitionLeadershipInfo.get(partition) match {
+              controllerContext.partitionLeadershipInfo(partition) match {
                 case Some(leaderIsrAndControllerEpoch) =>
                   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)
+          if (traceEnabled)
+            logSuccessfulTransition(stateLogger, 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(replica.topicPartition).isDefined
         }
-        val replicasToRemoveFromIsr = validReplicas.filter(replica => controllerContext.partitionLeadershipInfo.contains(replica.topicPartition))
-        val updatedLeaderIsrAndControllerEpochs = removeReplicasFromIsr(replicaId, replicasToRemoveFromIsr.map(_.topicPartition))
-        updatedLeaderIsrAndControllerEpochs.foreach { case (partition, leaderIsrAndControllerEpoch) =>
-          if (!topicDeletionManager.isPartitionToBeDeleted(partition)) {
+        val updatedLeaderIsrAndControllerEpochs = removeReplicasFromIsr(replicaId, replicasWithLeadershipInfo.map(_.topicPartition))
+        updatedLeaderIsrAndControllerEpochs.forKeyValue { (partition, leaderIsrAndControllerEpoch) =>
+          stateLogger.info(s"Partition $partition state changed to $leaderIsrAndControllerEpoch after removing replica $replicaId from the ISR as part of transition to $OfflineReplica")
+          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)
+          if (traceEnabled)
+            logSuccessfulTransition(stateLogger, replicaId, partition, currentState, OfflineReplica)
+          controllerContext.putReplicaState(replica, OfflineReplica)
+        }
+
+        replicasWithoutLeadershipInfo.foreach { replica =>
+          val currentState = controllerContext.replicaState(replica)
+          if (traceEnabled)
+            logSuccessfulTransition(stateLogger, 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)
+          if (traceEnabled)
+            logSuccessfulTransition(stateLogger, 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)
+          if (traceEnabled)
+            logSuccessfulTransition(stateLogger, 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)
+          if (traceEnabled)
+            logSuccessfulTransition(stateLogger, replicaId, replica.topicPartition, currentState, ReplicaDeletionSuccessful)
+          controllerContext.putReplicaState(replica, ReplicaDeletionSuccessful)
         }
       case NonExistentReplica =>
         validReplicas.foreach { replica =>
-          val currentAssignedReplicas = controllerContext.partitionReplicaAssignment(replica.topicPartition)
-          controllerContext.partitionReplicaAssignment.put(replica.topicPartition, currentAssignedReplicas.filterNot(_ == replica.replica))
-          logSuccessfulTransition(replicaId, replica.topicPartition, replicaState(replica), NonExistentReplica)
-          replicaState.remove(replica)
+          val currentState = controllerContext.replicaState(replica)
+          val newAssignedReplicas = controllerContext
+            .partitionFullReplicaAssignment(replica.topicPartition)
+            .removeReplica(replica.replica)
+
+          controllerContext.updatePartitionFullReplicaAssignment(replica.topicPartition, newAssignedReplicas)
+          if (traceEnabled)
+            logSuccessfulTransition(stateLogger, replicaId, replica.topicPartition, currentState, NonExistentReplica)
+          controllerContext.removeReplicaState(replica)
         }
     }
   }
@@ -251,17 +292,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
@@ -273,114 +320,117 @@ 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 (partition, 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.map { leaderAndIsr =>
+        leaderAndIsr.isr.contains(replicaId)
+      }.getOrElse(false)
     }
-    val UpdateLeaderAndIsrResult(successfulUpdates, updatesToRetry, failedUpdates) = zkClient.updateLeaderAndIsr(
-      adjustedLeaderAndIsrs, controllerContext.epoch)
-    val exceptionsForPartitionsWithNoLeaderAndIsrInZk = partitionsWithNoLeaderAndIsrInZk.flatMap { partition =>
-      if (!topicDeletionManager.isPartitionToBeDeleted(partition)) {
-        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.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) =>
+        (partition, result.map { leaderAndIsr =>
+          val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(leaderAndIsr, controllerContext.epoch)
+          controllerContext.putPartitionLeadershipInfo(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)
+    (result.toMap, partitionsWithNoLeaderAndIsrInZk)
   }
 
-  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)
-  }
-
-  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 logSuccessfulTransition(logger: StateChangeLogger, replicaId: Int, partition: TopicPartition,
+                                      currState: ReplicaState, targetState: ReplicaState): Unit = {
+    logger.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)
@@ -425,7 +475,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 b1d83947b14ca..0fd7274d128ad 100755
--- a/core/src/main/scala/kafka/controller/TopicDeletionManager.scala
+++ b/core/src/main/scala/kafka/controller/TopicDeletionManager.scala
@@ -16,11 +16,40 @@
  */
 package kafka.controller
 
+import kafka.server.KafkaConfig
 import kafka.utils.Logging
 import kafka.zk.KafkaZkClient
 import org.apache.kafka.common.TopicPartition
 
-import scala.collection.{Set, mutable}
+import scala.collection.Set
+import scala.collection.mutable
+
+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
+}
+
+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)
+  }
+}
 
 /**
  * This manages the state machine for topic deletion.
@@ -55,25 +84,25 @@ 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]
-  val partitionsToBeDeleted = mutable.Set.empty[TopicPartition]
-  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}] "
+  val isDeleteTopicEnabled: Boolean = config.deleteTopicEnable
 
   def init(initialTopicsToBeDeleted: Set[String], initialTopicsIneligibleForDeletion: Set[String]): Unit = {
+    info(s"Initializing manager with initial deletions: $initialTopicsToBeDeleted, " +
+      s"initial ineligible deletions: $initialTopicsIneligibleForDeletion")
+
     if (isDeleteTopicEnabled) {
-      topicsToBeDeleted ++= initialTopicsToBeDeleted
-      partitionsToBeDeleted ++= topicsToBeDeleted.flatMap(controllerContext.partitionsForTopic)
-      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)
+      client.deleteTopicDeletions(initialTopicsToBeDeleted.toSeq, controllerContext.epochZkVersion)
     }
   }
 
@@ -83,27 +112,15 @@ 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()
-      partitionsToBeDeleted.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
-      partitionsToBeDeleted ++= topics.flatMap(controllerContext.partitionsForTopic)
+      controllerContext.queueTopicDeletion(topics)
       resumeDeletions()
     }
   }
@@ -114,11 +131,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()
       }
     }
@@ -131,14 +148,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()
       }
     }
@@ -150,39 +167,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 isPartitionToBeDeleted(topicAndPartition: TopicPartition) = {
-    if (isDeleteTopicEnabled) {
-      partitionsToBeDeleted.contains(topicAndPartition)
+      controllerContext.isAnyReplicaInState(topic, ReplicaDeletionStarted)
     } else
       false
   }
 
   def isTopicQueuedUpForDeletion(topic: String): Boolean = {
     if (isDeleteTopicEnabled) {
-      topicsToBeDeleted.contains(topic)
+      controllerContext.isTopicQueuedUpForDeletion(topic)
     } else
       false
   }
@@ -193,10 +203,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()
   }
 
@@ -209,37 +219,31 @@ 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)
   }
 
   /**
    * If the topic is queued for deletion but deletion is not currently under progress, then deletion is retried for that topic
    * To ensure a successful retry, reset states for respective replicas from ReplicaDeletionIneligible to OfflineReplica state
-   *@param topic Topic for which deletion should be retried
+   * @param topics Topics for which deletion should be retried
    */
-  private def markTopicForDeletionRetry(topic: String) {
+  private def retryDeletionForIneligibleReplicas(topics: Set[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 = topics.flatMap(controllerContext.replicasInState(_, ReplicaDeletionIneligible))
+    debug(s"Retrying deletion of topics ${topics.mkString(",")} since replicas ${failedReplicas.mkString(",")} were not successfully deleted")
+    replicaStateMachine.handleStateChanges(failedReplicas.toSeq, OfflineReplica)
   }
 
-  private def completeDeleteTopic(topic: String) {
+  private def completeDeleteTopic(topic: String): Unit = {
     // 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)
-    val partitionsForDeletedTopic = controllerContext.partitionsForTopic(topic)
-    // move respective partition to OfflinePartition and NonExistentPartition state
-    controller.partitionStateMachine.handleStateChanges(partitionsForDeletedTopic.toSeq, OfflinePartition)
-    controller.partitionStateMachine.handleStateChanges(partitionsForDeletedTopic.toSeq, NonExistentPartition)
-    topicsToBeDeleted -= topic
-    partitionsToBeDeleted.retain(_.topic != topic)
-    zkClient.deleteTopicZNode(topic)
-    zkClient.deleteTopicConfigs(Seq(topic))
-    zkClient.deleteTopicDeletions(Seq(topic))
+    replicaStateMachine.handleStateChanges(replicasForDeletedTopic.toSeq, NonExistentReplica)
+    client.deleteTopic(topic, controllerContext.epochZkVersion)
     controllerContext.removeTopic(topic)
   }
 
@@ -250,108 +254,105 @@ 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]) {
-    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)
-    controller.sendUpdateMetadataRequest(controllerContext.liveOrShuttingDownBrokerIds.toSeq, partitions)
-    val partitionReplicaAssignmentByTopic = controllerContext.partitionReplicaAssignment.groupBy(p => p._1.topic)
-    topics.foreach { topic =>
-      onPartitionDeletion(partitionReplicaAssignmentByTopic(topic).keySet)
+  private def onTopicDeletion(topics: Set[String]): Unit = {
+    val unseenTopicsForDeletion = topics.diff(controllerContext.topicsWithDeletionStarted)
+    if (unseenTopicsForDeletion.nonEmpty) {
+      val unseenPartitionsForDeletion = unseenTopicsForDeletion.flatMap(controllerContext.partitionsForTopic)
+      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)
     }
-  }
 
-  /**
-   * Invoked by onPartitionDeletion. It is the 2nd step of topic deletion, the first being sending
-   * UpdateMetadata requests to all brokers to start rejecting requests for deleted topics. As part of starting deletion,
-   * the topics are added to the in progress list. As long as a topic is in the in progress list, deletion for that topic
-   * is never retried. A topic is removed from the in progress list when
-   * 1. Either the topic is successfully deleted OR
-   * 2. No replica for the topic is in ReplicaDeletionStarted state and at least one replica is in ReplicaDeletionIneligible state
-   * If the topic is queued for deletion but deletion is not currently under progress, then deletion is retried for that topic
-   * As part of starting deletion, all replicas are moved to the ReplicaDeletionStarted state where the controller sends
-   * the replicas a StopReplicaRequest (delete=true)
-   * This method does the following things -
-   * 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
-   */
-  private def startReplicaDeletion(replicasForTopicsToBeDeleted: Set[PartitionAndReplica]) {
-    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 replicasForDeletionRetry = aliveReplicasForTopic -- successfullyDeletedReplicas
-      // move dead replicas directly to failed state
-      controller.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)
-      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))))
-      if (deadReplicasForTopic.nonEmpty) {
-        debug(s"Dead Replicas (${deadReplicasForTopic.mkString(",")}) found for topic $topic")
-        markTopicIneligibleForDeletion(Set(topic))
-      }
-    }
+    // send update metadata so that brokers stop serving data for topics to be deleted
+    client.sendMetadataUpdate(topics.flatMap(controllerContext.partitionsForTopic))
+
+    onPartitionDeletion(topics)
   }
 
   /**
    * Invoked by onTopicDeletion with the list of partitions for topics to be deleted
    * It does the following -
-   * 1. Send UpdateMetadataRequest to all live brokers (that are not shutting down) for partitions that are being
-   *    deleted. The brokers start rejecting all client requests with UnknownTopicOrPartitionException
+   * 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 replicas for the partitions to OfflineReplica state. This will send StopReplicaRequest to the replicas
    *    and LeaderAndIsrRequest to the leader with the shrunk ISR. When the leader replica itself is moved to OfflineReplica state,
    *    it will skip sending the LeaderAndIsrRequest since the leader will be updated to -1
    * 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]) {
-    info(s"Partition deletion callback for ${partitionsToBeDeleted.mkString(",")}")
-    val replicasPerPartition = controllerContext.replicasForPartition(partitionsToBeDeleted)
-    startReplicaDeletion(replicasPerPartition)
+  private def onPartitionDeletion(topicsToBeDeleted: Set[String]): Unit = {
+    val allDeadReplicas = mutable.ListBuffer.empty[PartitionAndReplica]
+    val allReplicasForDeletionRetry = mutable.ListBuffer.empty[PartitionAndReplica]
+    val allTopicsIneligibleForDeletion = mutable.Set.empty[String]
+
+    topicsToBeDeleted.foreach { topic =>
+      val (aliveReplicas, deadReplicas) = controllerContext.replicasForTopic(topic).partition { r =>
+        controllerContext.isReplicaOnline(r.replica, r.topicPartition)
+      }
+
+      val successfullyDeletedReplicas = controllerContext.replicasInState(topic, ReplicaDeletionSuccessful)
+      val replicasForDeletionRetry = aliveReplicas.diff(successfullyDeletedReplicas)
+
+      allDeadReplicas ++= deadReplicas
+      allReplicasForDeletionRetry ++= replicasForDeletionRetry
+
+      if (deadReplicas.nonEmpty) {
+        debug(s"Dead Replicas (${deadReplicas.mkString(",")}) found for topic $topic")
+        allTopicsIneligibleForDeletion += topic
+      }
+    }
+
+    // move dead replicas directly to failed state
+    replicaStateMachine.handleStateChanges(allDeadReplicas, ReplicaDeletionIneligible)
+    // send stop replica to all followers that are not in the OfflineReplica state so they stop sending fetch requests to the leader
+    replicaStateMachine.handleStateChanges(allReplicasForDeletionRetry, OfflineReplica)
+    replicaStateMachine.handleStateChanges(allReplicasForDeletionRetry, ReplicaDeletionStarted)
+
+    if (allTopicsIneligibleForDeletion.nonEmpty) {
+      markTopicIneligibleForDeletion(allTopicsIneligibleForDeletion, reason = "offline replicas")
+    }
   }
 
   private def resumeDeletions(): Unit = {
-    val topicsQueuedForDeletion = Set.empty[String] ++ topicsToBeDeleted
+    val topicsQueuedForDeletion = Set.empty[String] ++ controllerContext.topicsToBeDeleted
+    val topicsEligibleForRetry = mutable.Set.empty[String]
+    val topicsEligibleForDeletion = mutable.Set.empty[String]
 
     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)) {
+          topicsEligibleForRetry += topic
         }
       }
-      // Try delete topic if it is eligible for deletion.
+
+      // Add topic to the eligible set 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")
+        topicsEligibleForDeletion += topic
       }
     }
+
+    // topic deletion retry will be kicked off
+    if (topicsEligibleForRetry.nonEmpty) {
+      retryDeletionForIneligibleReplicas(topicsEligibleForRetry)
+    }
+
+    // topic deletion will be kicked off
+    if (topicsEligibleForDeletion.nonEmpty) {
+      onTopicDeletion(topicsEligibleForDeletion)
+    }
   }
 }
diff --git a/core/src/main/scala/kafka/coordinator/group/DelayedHeartbeat.scala b/core/src/main/scala/kafka/coordinator/group/DelayedHeartbeat.scala
index 5f16acb6a85dc..3f402d9b5f747 100644
--- a/core/src/main/scala/kafka/coordinator/group/DelayedHeartbeat.scala
+++ b/core/src/main/scala/kafka/coordinator/group/DelayedHeartbeat.scala
@@ -25,12 +25,12 @@ import kafka.server.DelayedOperation
  */
 private[group] class DelayedHeartbeat(coordinator: GroupCoordinator,
                                       group: GroupMetadata,
-                                      member: MemberMetadata,
-                                      heartbeatDeadline: Long,
-                                      sessionTimeout: Long)
-  extends DelayedOperation(sessionTimeout, Some(group.lock)) {
+                                      memberId: String,
+                                      isPending: Boolean,
+                                      timeoutMs: Long)
+  extends DelayedOperation(timeoutMs, Some(group.lock)) {
 
-  override def tryComplete(): Boolean = coordinator.tryCompleteHeartbeat(group, member, heartbeatDeadline, forceComplete _)
-  override def onExpiration() = coordinator.onExpireHeartbeat(group, member, heartbeatDeadline)
+  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/DelayedJoin.scala b/core/src/main/scala/kafka/coordinator/group/DelayedJoin.scala
index c75c0d4a39747..92e8835d103ec 100644
--- a/core/src/main/scala/kafka/coordinator/group/DelayedJoin.scala
+++ b/core/src/main/scala/kafka/coordinator/group/DelayedJoin.scala
@@ -36,8 +36,15 @@ private[group] class DelayedJoin(coordinator: GroupCoordinator,
                                  rebalanceTimeout: Long) extends DelayedOperation(rebalanceTimeout, Some(group.lock)) {
 
   override def tryComplete(): Boolean = coordinator.tryCompleteJoin(group, forceComplete _)
-  override def onExpiration() = coordinator.onExpireJoin()
-  override def onComplete() = coordinator.onCompleteJoin(group)
+  override def onExpiration(): Unit = {
+    coordinator.onExpireJoin()
+    // try to complete delayed actions introduced by coordinator.onCompleteJoin
+    tryToCompleteDelayedAction()
+  }
+  override def onComplete(): Unit = coordinator.onCompleteJoin(group)
+
+  // TODO: remove this ugly chain after we move the action queue to handler thread
+  private def tryToCompleteDelayedAction(): Unit = coordinator.groupManager.replicaManager.tryCompleteActions()
 }
 
 /**
@@ -58,7 +65,7 @@ private[group] class InitialDelayedJoin(coordinator: GroupCoordinator,
   override def tryComplete(): Boolean = false
 
   override def onComplete(): Unit = {
-    group.inLock  {
+    group.inLock {
       if (group.newMemberAdded && remainingMs != 0) {
         group.newMemberAdded = false
         val delay = min(configuredRebalanceDelay, remainingMs)
diff --git a/core/src/main/scala/kafka/coordinator/group/GroupCoordinator.scala b/core/src/main/scala/kafka/coordinator/group/GroupCoordinator.scala
index 129eae4c40b47..cd8a4149e664d 100644
--- a/core/src/main/scala/kafka/coordinator/group/GroupCoordinator.scala
+++ b/core/src/main/scala/kafka/coordinator/group/GroupCoordinator.scala
@@ -27,15 +27,17 @@ 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.protocol.Errors
-import org.apache.kafka.common.record.RecordBatch.{NO_PRODUCER_EPOCH, NO_PRODUCER_ID}
+import org.apache.kafka.common.message.JoinGroupResponseData.JoinGroupResponseMember
+import org.apache.kafka.common.message.LeaveGroupRequestData.MemberIdentity
+import org.apache.kafka.common.metrics.Metrics
+import org.apache.kafka.common.metrics.stats.Meter
+import org.apache.kafka.common.protocol.{ApiKeys, Errors}
 import org.apache.kafka.common.requests._
 import org.apache.kafka.common.utils.Time
 
-import scala.collection.{Map, Seq, immutable}
+import scala.collection.{Map, Seq, Set, immutable, mutable}
 import scala.math.max
 
-
 /**
  * GroupCoordinator handles general group membership and offset management.
  *
@@ -54,11 +56,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 + "]: "
 
@@ -80,10 +104,9 @@ 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.")
-    if (enableMetadataExpiration)
-      groupManager.enableMetadataExpiration()
+    groupManager.startup(enableMetadataExpiration)
     isActive.set(true)
     info("Startup complete.")
   }
@@ -92,7 +115,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,134 +124,222 @@ class GroupCoordinator(val brokerId: Int,
     info("Shutdown complete.")
   }
 
+  /**
+   * Verify if the group has space to accept the joining member. The various
+   * criteria are explained below.
+   */
+  private def acceptJoiningMember(group: GroupMetadata, member: String): Boolean = {
+    group.currentState match {
+      // Always accept the request when the group is empty or dead
+      case Empty | Dead =>
+        true
+
+      // An existing member is accepted if it is already awaiting. New members are accepted
+      // up to the max group size. Note that the number of awaiting members is used here
+      // for two reasons:
+      // 1) the group size is not reliable as it could already be above the max group size
+      //    if the max group size was reduced.
+      // 2) using the number of awaiting members allows to kick out the last rejoining
+      //    members of the group.
+      case PreparingRebalance =>
+        (group.has(member) && group.get(member).isAwaitingJoin) ||
+          group.numAwaiting < groupConfig.groupMaxSize
+
+      // An existing member is accepted. New members are accepted up to the max group size.
+      // Note that the group size is used here. When the group transitions to CompletingRebalance,
+      // members which haven't rejoined are removed.
+      case CompletingRebalance | Stable =>
+        group.has(member) || group.size < groupConfig.groupMaxSize
+    }
+  }
+
   def handleJoinGroup(groupId: String,
                       memberId: String,
+                      groupInstanceId: Option[String],
+                      requireKnownMemberId: Boolean,
                       clientId: String,
                       clientHost: String,
                       rebalanceTimeoutMs: Int,
                       sessionTimeoutMs: Int,
                       protocolType: String,
                       protocols: List[(String, Array[Byte])],
-                      responseCallback: JoinCallback) {
-    if (!isActive.get) {
-      responseCallback(joinError(memberId, Errors.COORDINATOR_NOT_AVAILABLE))
-    } else if (!validGroupId(groupId)) {
-      responseCallback(joinError(memberId, Errors.INVALID_GROUP_ID))
-    } else if (!isCoordinatorForGroup(groupId)) {
-      responseCallback(joinError(memberId, Errors.NOT_COORDINATOR))
-    } else if (isCoordinatorLoadInProgress(groupId)) {
-      responseCallback(joinError(memberId, Errors.COORDINATOR_LOAD_IN_PROGRESS))
-    } else if (sessionTimeoutMs < groupConfig.groupMinSessionTimeoutMs ||
-               sessionTimeoutMs > groupConfig.groupMaxSessionTimeoutMs) {
-      responseCallback(joinError(memberId, Errors.INVALID_SESSION_TIMEOUT))
+                      responseCallback: JoinCallback): Unit = {
+    validateGroupStatus(groupId, ApiKeys.JOIN_GROUP).foreach { error =>
+      responseCallback(JoinGroupResult(memberId, error))
+      return
+    }
+
+    if (sessionTimeoutMs < groupConfig.groupMinSessionTimeoutMs ||
+      sessionTimeoutMs > groupConfig.groupMaxSessionTimeoutMs) {
+      responseCallback(JoinGroupResult(memberId, Errors.INVALID_SESSION_TIMEOUT))
     } else {
-      // only try to create the group if the group is not unknown AND
-      // the member id is UNKNOWN, if member is specified but group does not
-      // exist we should reject the request
-      groupManager.getGroup(groupId) match {
+      val isUnknownMember = memberId == JoinGroupRequest.UNKNOWN_MEMBER_ID
+      // group is created if it does not exist and the member id is UNKNOWN. if member
+      // is specified but group does not exist, request is rejected with UNKNOWN_MEMBER_ID
+      groupManager.getOrMaybeCreateGroup(groupId, isUnknownMember) match {
         case None =>
-          if (memberId != JoinGroupRequest.UNKNOWN_MEMBER_ID) {
-            responseCallback(joinError(memberId, Errors.UNKNOWN_MEMBER_ID))
-          } else {
-            val group = groupManager.addGroup(new GroupMetadata(groupId))
-            doJoinGroup(group, memberId, clientId, clientHost, rebalanceTimeoutMs, sessionTimeoutMs, protocolType, protocols, responseCallback)
+          responseCallback(JoinGroupResult(memberId, Errors.UNKNOWN_MEMBER_ID))
+        case Some(group) =>
+          group.inLock {
+            if (!acceptJoiningMember(group, memberId)) {
+              group.remove(memberId)
+              group.removeStaticMember(groupInstanceId)
+              responseCallback(JoinGroupResult(JoinGroupRequest.UNKNOWN_MEMBER_ID, Errors.GROUP_MAX_SIZE_REACHED))
+            } else if (isUnknownMember) {
+              doUnknownJoinGroup(group, groupInstanceId, requireKnownMemberId, clientId, clientHost, rebalanceTimeoutMs, sessionTimeoutMs, protocolType, protocols, responseCallback)
+            } else {
+              doJoinGroup(group, memberId, groupInstanceId, clientId, clientHost, rebalanceTimeoutMs, sessionTimeoutMs, protocolType, protocols, responseCallback)
+            }
+
+            // attempt to complete JoinGroup
+            if (group.is(PreparingRebalance)) {
+              joinPurgatory.checkAndComplete(GroupKey(group.groupId))
+            }
           }
+      }
+    }
+  }
 
-        case Some(group) =>
-          doJoinGroup(group, memberId, clientId, clientHost, rebalanceTimeoutMs, sessionTimeoutMs, protocolType, protocols, responseCallback)
+  private def doUnknownJoinGroup(group: GroupMetadata,
+                                 groupInstanceId: Option[String],
+                                 requireKnownMemberId: Boolean,
+                                 clientId: String,
+                                 clientHost: String,
+                                 rebalanceTimeoutMs: Int,
+                                 sessionTimeoutMs: Int,
+                                 protocolType: String,
+                                 protocols: List[(String, Array[Byte])],
+                                 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; 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(JoinGroupResult(JoinGroupRequest.UNKNOWN_MEMBER_ID, Errors.COORDINATOR_NOT_AVAILABLE))
+      } else if (!group.supportsProtocols(protocolType, MemberMetadata.plainProtocolSet(protocols))) {
+        responseCallback(JoinGroupResult(JoinGroupRequest.UNKNOWN_MEMBER_ID, Errors.INCONSISTENT_GROUP_PROTOCOL))
+      } else {
+        val newMemberId = group.generateMemberId(clientId, groupInstanceId)
+
+        if (group.hasStaticMember(groupInstanceId)) {
+          updateStaticMemberAndRebalance(group, newMemberId, groupInstanceId, protocols, responseCallback)
+        } 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 joins 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(JoinGroupResult(newMemberId, Errors.MEMBER_ID_REQUIRED))
+        } else {
+          info(s"${if (groupInstanceId.isDefined) "Static" else "Dynamic"} Member with unknown member id joins 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)
+        }
       }
     }
   }
 
   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(Empty) && (!group.protocolType.contains(protocolType) || !group.supportsProtocols(protocols.map(_._1).toSet))) {
-        // if the new member does not support the group protocol, reject it
-        responseCallback(joinError(memberId, Errors.INCONSISTENT_GROUP_PROTOCOL))
-      } else if (group.is(Empty) && (protocols.isEmpty || protocolType.isEmpty)) {
-        //reject if first member with empty group protocol or protocolType is empty
-        responseCallback(joinError(memberId, Errors.INCONSISTENT_GROUP_PROTOCOL))
-      } else if (memberId != JoinGroupRequest.UNKNOWN_MEMBER_ID && !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))
+      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(JoinGroupResult(memberId, Errors.COORDINATOR_NOT_AVAILABLE))
+      } else if (!group.supportsProtocols(protocolType, MemberMetadata.plainProtocolSet(protocols))) {
+        responseCallback(JoinGroupResult(memberId, Errors.INCONSISTENT_GROUP_PROTOCOL))
+      } else if (group.isPendingMember(memberId)) {
+        // 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 {
+          debug(s"Dynamic Member with specific member id $memberId joins group ${group.groupId} in " +
+            s"${group.currentState} state. Adding to the group now.")
+          addMemberAndRebalance(rebalanceTimeoutMs, sessionTimeoutMs, memberId, groupInstanceId,
+            clientId, clientHost, protocolType, protocols, group, responseCallback)
+        }
       } else {
-        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; 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))
-          case PreparingRebalance =>
-            if (memberId == JoinGroupRequest.UNKNOWN_MEMBER_ID) {
-              addMemberAndRebalance(rebalanceTimeoutMs, sessionTimeoutMs, clientId, clientHost, protocolType, protocols, group, responseCallback)
-            } else {
-              val member = group.get(memberId)
-              updateMemberAndRebalance(group, member, protocols, responseCallback)
-            }
+        val groupInstanceIdNotFound = groupInstanceId.isDefined && !group.hasStaticMember(groupInstanceId)
+        if (group.isStaticMemberFenced(memberId, groupInstanceId, "join-group")) {
+          // given member id doesn't match with the groupInstanceId. Inform duplicate instance to shut down immediately.
+          responseCallback(JoinGroupResult(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(JoinGroupResult(memberId, Errors.UNKNOWN_MEMBER_ID))
+        } else {
+          val member = group.get(memberId)
 
-          case CompletingRebalance =>
-            if (memberId == JoinGroupRequest.UNKNOWN_MEMBER_ID) {
-              addMemberAndRebalance(rebalanceTimeoutMs, sessionTimeoutMs, clientId, clientHost, protocolType, protocols, group, responseCallback)
-            } else {
-              val member = group.get(memberId)
+          group.currentState match {
+            case PreparingRebalance =>
+              updateMemberAndRebalance(group, member, protocols, s"Member ${member.memberId} joining group during ${group.currentState}", responseCallback)
+
+            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 (memberId == group.leaderId) {
+                  members = if (group.isLeader(memberId)) {
                     group.currentMemberMetadata
                   } else {
-                    Map.empty
+                    List.empty
                   },
                   memberId = memberId,
                   generationId = group.generationId,
-                  subProtocol = group.protocol,
-                  leaderId = group.leaderId,
+                  protocolType = group.protocolType,
+                  protocolName = group.protocolName,
+                  leaderId = group.leaderOrNull,
                   error = Errors.NONE))
               } else {
                 // member has changed metadata, so force a rebalance
-                updateMemberAndRebalance(group, member, protocols, responseCallback)
+                updateMemberAndRebalance(group, member, protocols, s"Updating metadata for member ${member.memberId} during ${group.currentState}", responseCallback)
               }
-            }
 
-          case Empty | Stable =>
-            if (memberId == JoinGroupRequest.UNKNOWN_MEMBER_ID) {
-              // if the member id is unknown, register the member to the group
-              addMemberAndRebalance(rebalanceTimeoutMs, sessionTimeoutMs, clientId, clientHost, protocolType, protocols, group, responseCallback)
-            } else {
+            case Stable =>
               val member = group.get(memberId)
-              if (memberId == group.leaderId || !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
+              if (group.isLeader(memberId)) {
+                // force a rebalance if the leader sends JoinGroup;
+                // This 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)
+                updateMemberAndRebalance(group, member, protocols, s"leader ${member.memberId} re-joining group during ${group.currentState}", responseCallback)
+              } else if (!member.matches(protocols)) {
+                updateMemberAndRebalance(group, member, protocols, s"Updating metadata for member ${member.memberId} during ${group.currentState}", 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,
+                  members = List.empty,
                   memberId = memberId,
                   generationId = group.generationId,
-                  subProtocol = group.protocol,
-                  leaderId = group.leaderId,
+                  protocolType = group.protocolType,
+                  protocolName = group.protocolName,
+                  leaderId = group.leaderOrNull,
                   error = Errors.NONE))
               }
-            }
-        }
 
-        if (group.is(PreparingRebalance))
-          joinPurgatory.checkAndComplete(GroupKey(group.groupId))
+            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(JoinGroupResult(memberId, Errors.UNKNOWN_MEMBER_ID))
+          }
+        }
       }
     }
   }
@@ -236,49 +347,78 @@ class GroupCoordinator(val brokerId: Int,
   def handleSyncGroup(groupId: String,
                       generation: Int,
                       memberId: String,
+                      protocolType: Option[String],
+                      protocolName: Option[String],
+                      groupInstanceId: Option[String],
                       groupAssignment: Map[String, Array[Byte]],
-                      responseCallback: SyncCallback) {
-    if (!isActive.get) {
-      responseCallback(Array.empty, Errors.COORDINATOR_NOT_AVAILABLE)
-    } else if (!isCoordinatorForGroup(groupId)) {
-      responseCallback(Array.empty, Errors.NOT_COORDINATOR)
-    } else {
-      groupManager.getGroup(groupId) match {
-        case None => responseCallback(Array.empty, Errors.UNKNOWN_MEMBER_ID)
-        case Some(group) => doSyncGroup(group, generation, memberId, groupAssignment, responseCallback)
-      }
+                      responseCallback: SyncCallback): Unit = {
+    validateGroupStatus(groupId, ApiKeys.SYNC_GROUP) match {
+      case Some(error) if error == Errors.COORDINATOR_LOAD_IN_PROGRESS =>
+        // The coordinator is loading, which means we've lost the state of the active rebalance and the
+        // 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(SyncGroupResult(Errors.REBALANCE_IN_PROGRESS))
+
+      case Some(error) => responseCallback(SyncGroupResult(error))
+
+      case None =>
+        groupManager.getGroup(groupId) match {
+          case None => responseCallback(SyncGroupResult(Errors.UNKNOWN_MEMBER_ID))
+          case Some(group) => doSyncGroup(group, generation, memberId, protocolType, protocolName,
+            groupInstanceId, groupAssignment, responseCallback)
+        }
     }
   }
 
   private def doSyncGroup(group: GroupMetadata,
                           generationId: Int,
                           memberId: String,
+                          protocolType: Option[String],
+                          protocolName: Option[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(Errors.COORDINATOR_NOT_AVAILABLE))
+      } else if (group.isStaticMemberFenced(memberId, groupInstanceId, "sync-group")) {
+        responseCallback(SyncGroupResult(Errors.FENCED_INSTANCE_ID))
+      } else if (!group.has(memberId)) {
+        responseCallback(SyncGroupResult(Errors.UNKNOWN_MEMBER_ID))
       } else if (generationId != group.generationId) {
-        responseCallback(Array.empty, Errors.ILLEGAL_GENERATION)
+        responseCallback(SyncGroupResult(Errors.ILLEGAL_GENERATION))
+      } else if (protocolType.isDefined && !group.protocolType.contains(protocolType.get)) {
+        responseCallback(SyncGroupResult(Errors.INCONSISTENT_GROUP_PROTOCOL))
+      } else if (protocolName.isDefined && !group.protocolName.contains(protocolName.get)) {
+        responseCallback(SyncGroupResult(Errors.INCONSISTENT_GROUP_PROTOCOL))
       } else {
         group.currentState match {
-          case Empty | Dead =>
-            responseCallback(Array.empty, Errors.UNKNOWN_MEMBER_ID)
+          case Empty =>
+            responseCallback(SyncGroupResult(Errors.UNKNOWN_MEMBER_ID))
 
           case PreparingRebalance =>
-            responseCallback(Array.empty, Errors.REBALANCE_IN_PROGRESS)
+            responseCallback(SyncGroupResult(Errors.REBALANCE_IN_PROGRESS))
 
           case CompletingRebalance =>
             group.get(memberId).awaitingSyncCallback = responseCallback
 
             // if this is the leader, then we can attempt to persist state and transition to stable
-            if (memberId == group.leaderId) {
+            if (group.isLeader(memberId)) {
               info(s"Assignment received from leader for group ${group.groupId} for generation ${group.generationId}")
 
               // fill any missing members with an empty assignment
-              val missing = group.allMembers -- groupAssignment.keySet
+              val missing = group.allMembers.diff(groupAssignment.keySet)
               val assignment = groupAssignment ++ missing.map(_ -> Array.empty[Byte]).toMap
 
+              if (missing.nonEmpty) {
+                warn(s"Setting empty assignments for members $missing of ${group.groupId} for generation ${group.generationId}")
+              }
+
               groupManager.storeGroup(group, assignment, (error: Errors) => {
                 group.inLock {
                   // another member may have joined the group while we were awaiting this callback,
@@ -287,7 +427,7 @@ class GroupCoordinator(val brokerId: Int,
                   if (group.is(CompletingRebalance) && generationId == group.generationId) {
                     if (error != Errors.NONE) {
                       resetAndPropagateAssignmentError(group, error)
-                      maybePrepareRebalance(group)
+                      maybePrepareRebalance(group, s"error when storing group assignment during SyncGroup (member: $memberId)")
                     } else {
                       setAndPropagateAssignment(group, assignment)
                       group.transitionTo(Stable)
@@ -295,108 +435,232 @@ 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(group.protocolType, group.protocolName, 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) {
-    if (!isActive.get) {
-      responseCallback(Errors.COORDINATOR_NOT_AVAILABLE)
-    } else if (!isCoordinatorForGroup(groupId)) {
-      responseCallback(Errors.NOT_COORDINATOR)
-    } else if (isCoordinatorLoadInProgress(groupId)) {
-      responseCallback(Errors.COORDINATOR_LOAD_IN_PROGRESS)
-    } else {
-      groupManager.getGroup(groupId) match {
-        case None =>
-          // 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 consumer to retry
-          // joining without specified consumer id,
-          responseCallback(Errors.UNKNOWN_MEMBER_ID)
-
-        case Some(group) =>
-          group.inLock {
-            if (group.is(Dead) || !group.has(memberId)) {
-              responseCallback(Errors.UNKNOWN_MEMBER_ID)
-            } else {
-              val member = group.get(memberId)
-              removeHeartbeatForLeavingMember(group, member)
-              debug(s"Member ${member.memberId} in group ${group.groupId} has left, removing it from the group")
-              removeMemberAndUpdateGroup(group, member)
-              responseCallback(Errors.NONE)
+  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 =>
+        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, "leave-group")) {
+                    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))
+              }
             }
+        }
+    }
+  }
+
+  def handleDeleteGroups(groupIds: Set[String]): Map[String, Errors] = {
+    val groupErrors = mutable.Map.empty[String, Errors]
+    val groupsEligibleForDeletion = mutable.ArrayBuffer[GroupMetadata]()
+
+    groupIds.foreach { groupId =>
+      validateGroupStatus(groupId, ApiKeys.DELETE_GROUPS) match {
+        case Some(error) =>
+          groupErrors += groupId -> error
+
+        case None =>
+          groupManager.getGroup(groupId) match {
+            case None =>
+              groupErrors += groupId ->
+                (if (groupManager.groupNotExists(groupId)) Errors.GROUP_ID_NOT_FOUND else Errors.NOT_COORDINATOR)
+            case Some(group) =>
+              group.inLock {
+                group.currentState match {
+                  case Dead =>
+                    groupErrors += groupId ->
+                      (if (groupManager.groupNotExists(groupId)) Errors.GROUP_ID_NOT_FOUND else Errors.NOT_COORDINATOR)
+                  case Empty =>
+                    group.transitionTo(Dead)
+                    groupsEligibleForDeletion += group
+                  case Stable | PreparingRebalance | CompletingRebalance =>
+                    groupErrors(groupId) = Errors.NON_EMPTY_GROUP
+                }
+              }
           }
       }
     }
+
+    if (groupsEligibleForDeletion.nonEmpty) {
+      val offsetsRemoved = groupManager.cleanupGroupMetadata(groupsEligibleForDeletion, _.removeAllOffsets())
+      groupErrors ++= groupsEligibleForDeletion.map(_.groupId -> Errors.NONE).toMap
+      info(s"The following groups were deleted: ${groupsEligibleForDeletion.map(_.groupId).mkString(", ")}. " +
+        s"A total of $offsetsRemoved offsets were removed.")
+    }
+
+    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) {
-    if (!isActive.get) {
-      responseCallback(Errors.COORDINATOR_NOT_AVAILABLE)
-    } else if (!isCoordinatorForGroup(groupId)) {
-      responseCallback(Errors.NOT_COORDINATOR)
-    } else if (isCoordinatorLoadInProgress(groupId)) {
-      // the group is still loading, so respond just blindly
-      responseCallback(Errors.NONE)
-    } else {
-      groupManager.getGroup(groupId) match {
-        case None =>
+                      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
+        responseCallback(Errors.NONE)
+      else
+        responseCallback(error)
+      return
+    }
+
+    groupManager.getGroup(groupId) match {
+      case None =>
+        responseCallback(Errors.UNKNOWN_MEMBER_ID)
+
+      case Some(group) => 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
+          // finding the correct coordinator and rejoin.
+          responseCallback(Errors.COORDINATOR_NOT_AVAILABLE)
+        } else if (group.isStaticMemberFenced(memberId, groupInstanceId, "heartbeat")) {
+          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 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; 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(Errors.UNKNOWN_MEMBER_ID)
-
-              case Empty =>
-                responseCallback(Errors.UNKNOWN_MEMBER_ID)
-
-              case CompletingRebalance =>
-                if (!group.has(memberId))
-                  responseCallback(Errors.UNKNOWN_MEMBER_ID)
-                else
-                  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 CompletingRebalance =>
+              // consumers may start sending heartbeat after join-group response, in which case
+              // we should treat them as normal hb request and reset the timer
+              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 PreparingRebalance =>
+                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 Dead =>
+              throw new IllegalStateException(s"Reached unexpected condition for Dead group $groupId")
           }
+        }
       }
     }
   }
@@ -404,138 +668,205 @@ class GroupCoordinator(val brokerId: Int,
   def handleTxnCommitOffsets(groupId: String,
                              producerId: Long,
                              producerEpoch: Short,
+                             memberId: String,
+                             groupInstanceId: Option[String],
+                             generationId: Int,
                              offsetMetadata: immutable.Map[TopicPartition, OffsetAndMetadata],
                              responseCallback: immutable.Map[TopicPartition, Errors] => Unit): Unit = {
-    validateGroup(groupId) match {
-      case Some(error) => responseCallback(offsetMetadata.mapValues(_ => error))
+    validateGroupStatus(groupId, ApiKeys.TXN_OFFSET_COMMIT) match {
+      case Some(error) => responseCallback(offsetMetadata.map { case (k, _) => k -> error })
       case None =>
-        val group = groupManager.getGroup(groupId).getOrElse(groupManager.addGroup(new GroupMetadata(groupId)))
-        doCommitOffsets(group, NoMemberId, NoGeneration, producerId, producerEpoch, offsetMetadata, responseCallback)
+        val group = groupManager.getGroup(groupId).getOrElse {
+          groupManager.addGroup(new GroupMetadata(groupId, Empty, time))
+        }
+        doTxnCommitOffsets(group, memberId, groupInstanceId, generationId, 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) {
-    validateGroup(groupId) match {
-      case Some(error) => responseCallback(offsetMetadata.mapValues(_ => error))
+                          responseCallback: immutable.Map[TopicPartition, Errors] => Unit): Unit = {
+    validateGroupStatus(groupId, ApiKeys.OFFSET_COMMIT) match {
+      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))
-              doCommitOffsets(group, memberId, generationId, NO_PRODUCER_ID, NO_PRODUCER_EPOCH,
-                offsetMetadata, responseCallback)
+              val group = groupManager.addGroup(new GroupMetadata(groupId, Empty, time))
+              doCommitOffsets(group, memberId, groupInstanceId, generationId, 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,
-              offsetMetadata, responseCallback)
+            doCommitOffsets(group, memberId, groupInstanceId, generationId, offsetMetadata, responseCallback)
         }
     }
   }
 
-  def handleTxnCompletion(producerId: Long,
-                          offsetsPartitions: Iterable[TopicPartition],
-                          transactionResult: TransactionResult) {
+  def scheduleHandleTxnCompletion(producerId: Long,
+                                  offsetsPartitions: Iterable[TopicPartition],
+                                  transactionResult: TransactionResult): Unit = {
     require(offsetsPartitions.forall(_.topic == Topic.GROUP_METADATA_TOPIC_NAME))
     val isCommit = transactionResult == TransactionResult.COMMIT
-    groupManager.handleTxnCompletion(producerId, offsetsPartitions.map(_.partition).toSet, isCommit)
+    groupManager.scheduleHandleTxnCompletion(producerId, offsetsPartitions.map(_.partition).toSet, isCommit)
+  }
+
+  private def doTxnCommitOffsets(group: GroupMetadata,
+                                 memberId: String,
+                                 groupInstanceId: Option[String],
+                                 generationId: Int,
+                                 producerId: Long,
+                                 producerEpoch: Short,
+                                 offsetMetadata: immutable.Map[TopicPartition, OffsetAndMetadata],
+                                 responseCallback: immutable.Map[TopicPartition, Errors] => Unit): 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; 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, "txn-commit-offsets")) {
+        responseCallback(offsetMetadata.map { case (k, _) => k -> Errors.FENCED_INSTANCE_ID })
+      } else if (memberId != JoinGroupRequest.UNKNOWN_MEMBER_ID && !group.has(memberId)) {
+        // Enforce member id when it is set.
+        responseCallback(offsetMetadata.map { case (k, _) => k -> Errors.UNKNOWN_MEMBER_ID })
+      } else if (generationId >= 0 && generationId != group.generationId) {
+        // Enforce generation check when it is set.
+        responseCallback(offsetMetadata.map { case (k, _) => k -> Errors.ILLEGAL_GENERATION })
+      } else {
+        groupManager.storeOffsets(group, memberId, offsetMetadata, responseCallback, producerId, producerEpoch)
+      }
+    }
   }
 
   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))
-      } 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))
+        // 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, "commit-offsets")) {
+        responseCallback(offsetMetadata.map { case (k, _) => k -> Errors.FENCED_INSTANCE_ID })
+      } else if (generationId < 0 && group.is(Empty)) {
+        // The group is only using Kafka to store offsets.
+        groupManager.storeOffsets(group, memberId, offsetMetadata, responseCallback)
       } 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}")
+        }
       }
     }
   }
 
-  def handleFetchOffsets(groupId: String,
-                         partitions: Option[Seq[TopicPartition]] = None): (Errors, Map[TopicPartition, OffsetFetchResponse.PartitionData]) = {
-    if (!isActive.get)
-      (Errors.COORDINATOR_NOT_AVAILABLE, Map())
-    else if (!isCoordinatorForGroup(groupId)) {
-      debug("Could not fetch offsets for group %s (not group coordinator).".format(groupId))
-      (Errors.NOT_COORDINATOR, Map())
-    } else if (isCoordinatorLoadInProgress(groupId))
-      (Errors.COORDINATOR_LOAD_IN_PROGRESS, Map())
-    else {
-      // return offsets blindly regardless the current group state since the group may be using
-      // Kafka commit storage without automatic group management
-      (Errors.NONE, groupManager.getOffsets(groupId, partitions))
+  def handleFetchOffsets(groupId: String, requireStable: Boolean, partitions: Option[Seq[TopicPartition]] = None):
+  (Errors, Map[TopicPartition, OffsetFetchResponse.PartitionData]) = {
+
+    validateGroupStatus(groupId, ApiKeys.OFFSET_FETCH) match {
+      case Some(error) => error -> Map.empty
+      case None =>
+        // return offsets blindly regardless the current group state since the group may be using
+        // Kafka commit storage without automatic group management
+        (Errors.NONE, groupManager.getOffsets(groupId, requireStable, partitions))
     }
   }
 
-  def handleListGroups(): (Errors, List[GroupOverview]) = {
+  def handleListGroups(states: Set[String]): (Errors, List[GroupOverview]) = {
     if (!isActive.get) {
       (Errors.COORDINATOR_NOT_AVAILABLE, List[GroupOverview]())
     } else {
       val errorCode = if (groupManager.isLoading) Errors.COORDINATOR_LOAD_IN_PROGRESS else Errors.NONE
-      (errorCode, groupManager.currentGroups.map(_.overview).toList)
+      // if states is empty, return all groups
+      val groups = if (states.isEmpty)
+        groupManager.currentGroups
+      else
+        groupManager.currentGroups.filter(g => states.contains(g.summary.state))
+      (errorCode, groups.map(_.overview).toList)
     }
   }
 
   def handleDescribeGroup(groupId: String): (Errors, GroupSummary) = {
-    if (!isActive.get) {
-      (Errors.COORDINATOR_NOT_AVAILABLE, GroupCoordinator.EmptyGroup)
-    } else if (!isCoordinatorForGroup(groupId)) {
-      (Errors.NOT_COORDINATOR, GroupCoordinator.EmptyGroup)
-    } else if (isCoordinatorLoadInProgress(groupId)) {
-      (Errors.COORDINATOR_LOAD_IN_PROGRESS, GroupCoordinator.EmptyGroup)
-    } else {
-      groupManager.getGroup(groupId) match {
-        case None => (Errors.NONE, GroupCoordinator.DeadGroup)
-        case Some(group) =>
-          group.inLock {
-            (Errors.NONE, group.summary)
-          }
-      }
+    validateGroupStatus(groupId, ApiKeys.DESCRIBE_GROUPS) match {
+      case Some(error) => (error, GroupCoordinator.EmptyGroup)
+      case None =>
+        groupManager.getGroup(groupId) match {
+          case None => (Errors.NONE, GroupCoordinator.DeadGroup)
+          case Some(group) =>
+            group.inLock {
+              (Errors.NONE, group.summary)
+            }
+        }
     }
   }
 
-  def handleDeletedPartitions(topicPartitions: Seq[TopicPartition]) {
-    groupManager.cleanupGroupMetadata(Some(topicPartitions))
+  def handleDeletedPartitions(topicPartitions: Seq[TopicPartition]): Unit = {
+    val offsetsRemoved = groupManager.cleanupGroupMetadata(groupManager.currentGroups, group => {
+      group.removeOffsets(topicPartitions)
+    })
+    info(s"Removed $offsetsRemoved offsets associated with deleted partitions: ${topicPartitions.mkString(", ")}.")
   }
 
-  private def validateGroup(groupId: String): Option[Errors] = {
-    if (!isActive.get)
+  private def isValidGroupId(groupId: String, api: ApiKeys): Boolean = {
+    api match {
+      case ApiKeys.OFFSET_COMMIT | ApiKeys.OFFSET_FETCH | ApiKeys.DESCRIBE_GROUPS | ApiKeys.DELETE_GROUPS =>
+        // For backwards compatibility, we support the offset commit APIs for the empty groupId, and also
+        // in DescribeGroups and DeleteGroups so that users can view and delete state of all groups.
+        groupId != null
+      case _ =>
+        // The remaining APIs are groups using Kafka for group coordination and must have a non-empty groupId
+        groupId != null && !groupId.isEmpty
+    }
+  }
+
+  /**
+   * Check that the groupId is valid, assigned to this coordinator and that the group has been loaded.
+   */
+  private def validateGroupStatus(groupId: String, api: ApiKeys): Option[Errors] = {
+    if (!isValidGroupId(groupId, api))
+      Some(Errors.INVALID_GROUP_ID)
+    else if (!isActive.get)
       Some(Errors.COORDINATOR_NOT_AVAILABLE)
-    else if (!isCoordinatorForGroup(groupId))
-      Some(Errors.NOT_COORDINATOR)
     else if (isCoordinatorLoadInProgress(groupId))
       Some(Errors.COORDINATOR_LOAD_IN_PROGRESS)
+    else if (!isCoordinatorForGroup(groupId))
+      Some(Errors.NOT_COORDINATOR)
     else
       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
@@ -545,59 +876,75 @@ class GroupCoordinator(val brokerId: Int,
         case Empty | Dead =>
         case PreparingRebalance =>
           for (member <- group.allMemberMetadata) {
-            if (member.awaitingJoinCallback != null) {
-              member.awaitingJoinCallback(joinError(member.memberId, Errors.NOT_COORDINATOR))
-              member.awaitingJoinCallback = null
-            }
+            group.maybeInvokeJoinCallback(member, JoinGroupResult(member.memberId, Errors.NOT_COORDINATOR))
           }
+
           joinPurgatory.checkAndComplete(GroupKey(group.groupId))
 
         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(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))
+      if (groupIsOverCapacity(group)) {
+        prepareRebalance(group, s"Freshly-loaded group is over capacity ($groupConfig.groupMaxSize). Rebalacing in order to give a chance for consumers to commit offsets")
+      }
+
       group.allMemberMetadata.foreach(completeAndScheduleNextHeartbeatExpiration(group, _))
     }
   }
 
-  def handleGroupImmigration(offsetTopicPartitionId: Int) {
-    groupManager.loadGroupsForPartition(offsetTopicPartitionId, onGroupLoaded)
+  /**
+   * 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 = {
+    info(s"Elected as the group coordinator for partition $offsetTopicPartitionId")
+    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 = {
+    info(s"Resigned as the group coordinator for partition $offsetTopicPartitionId")
     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 = {
+    val (protocolType, protocolName) = if (error == Errors.NONE)
+      (group.protocolType, group.protocolName)
+    else
+      (None, None)
     for (member <- group.allMemberMetadata) {
-      if (member.awaitingSyncCallback != null) {
-        member.awaitingSyncCallback(member.assignment, error)
-        member.awaitingSyncCallback = null
+      if (member.assignment.isEmpty && error == Errors.NONE) {
+        warn(s"Sending empty assignment to member ${member.memberId} of ${group.groupId} for generation ${group.generationId} with no errors")
+      }
 
+      if (group.maybeInvokeSyncCallback(member, SyncGroupResult(protocolType, protocolName, 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
@@ -607,36 +954,36 @@ class GroupCoordinator(val brokerId: Int,
     }
   }
 
-  private def validGroupId(groupId: String): Boolean = {
-    groupId != null && !groupId.isEmpty
-  }
-
-  private def joinError(memberId: String, error: Errors): JoinGroupResult = {
-    JoinGroupResult(
-      members = Map.empty,
-      memberId = memberId,
-      generationId = 0,
-      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) {
-    // complete current heartbeat expectation
-    member.latestHeartbeat = time.milliseconds()
+  private def completeAndScheduleNextHeartbeatExpiration(group: GroupMetadata, member: MemberMetadata): Unit = {
+    completeAndScheduleNextExpiration(group, member, member.sessionTimeoutMs)
+  }
+
+  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 newHeartbeatDeadline = member.latestHeartbeat + member.sessionTimeoutMs
-    val delayedHeartbeat = new DelayedHeartbeat(this, group, member, newHeartbeatDeadline, member.sessionTimeoutMs)
+    member.heartbeatSatisfied = false
+    val delayedHeartbeat = new DelayedHeartbeat(this, group, member.memberId, isPending = false, timeoutMs)
     heartbeatPurgatory.tryCompleteElseWatch(delayedHeartbeat, Seq(memberKey))
   }
 
-  private def removeHeartbeatForLeavingMember(group: GroupMetadata, member: MemberMetadata) {
+  /**
+    * Add pending member expiration to heartbeat purgatory
+    */
+  private def addPendingMemberExpiration(group: GroupMetadata, pendingMemberId: String, timeoutMs: Long): Unit = {
+    val pendingMemberKey = MemberKey(group.groupId, pendingMemberId)
+    val delayedHeartbeat = new DelayedHeartbeat(this, group, pendingMemberId, isPending = true, timeoutMs)
+    heartbeatPurgatory.tryCompleteElseWatch(delayedHeartbeat, Seq(pendingMemberKey))
+  }
+
+  private def removeHeartbeatForLeavingMember(group: GroupMetadata, member: MemberMetadata): Unit = {
     member.isLeaving = true
     val memberKey = MemberKey(member.groupId, member.memberId)
     heartbeatPurgatory.checkAndComplete(memberKey)
@@ -644,42 +991,138 @@ 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 memberId = clientId + "-" + group.generateMemberIdSuffix
-    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.awaitingJoinCallback = callback
+
+    member.isNew = true
+
     // update the newMemberAdded flag to indicate that the join group can be further delayed
     if (group.is(PreparingRebalance) && group.generationId == 0)
       group.newMemberAdded = true
 
-    group.add(member)
-    maybePrepareRebalance(group)
-    member
+    group.add(member, callback)
+
+    // The session timeout does not affect new members since they do not have their memberId and
+    // cannot send heartbeats. Furthermore, we cannot detect disconnects because sockets are muted
+    // while the JoinGroup is in purgatory. If the client does disconnect (e.g. because of a request
+    // timeout during a long rebalance), they may simply retry which will lead to a lot of defunct
+    // members in the rebalance. To prevent this going on indefinitely, we timeout JoinGroup requests
+    // for new members. If the new member is still there, we expect it to retry.
+    completeAndScheduleNextExpiration(group, member, NewMemberJoinTimeoutMs)
+
+    if (member.isStaticMember) {
+      info(s"Adding new static member $groupInstanceId to group ${group.groupId} with member id $memberId.")
+      group.addStaticMember(groupInstanceId, memberId)
+    } else {
+      group.removePendingMember(memberId)
+    }
+    maybePrepareRebalance(group, s"Adding new member $memberId with group instance id $groupInstanceId")
+  }
+
+  private def updateStaticMemberAndRebalance(group: GroupMetadata,
+                                             newMemberId: String,
+                                             groupInstanceId: Option[String],
+                                             protocols: List[(String, Array[Byte])],
+                                             responseCallback: JoinCallback): Unit = {
+    val oldMemberId = group.getStaticMemberId(groupInstanceId)
+    info(s"Static member $groupInstanceId of group ${group.groupId} with unknown member id rejoins, assigning new member id $newMemberId, while " +
+      s"old member id $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)
+    val oldProtocols = knownStaticMember.supportedProtocols
+
+    group.currentState match {
+      case Stable =>
+        // check if group's selectedProtocol of next generation will change, if not, simply store group to persist the
+        // updated static member, if yes, rebalance should be triggered to let the group's assignment and selectProtocol consistent
+        val selectedProtocolOfNextGeneration = group.selectProtocol
+        if (group.protocolName.contains(selectedProtocolOfNextGeneration)) {
+          info(s"Static member which joins during Stable stage and doesn't affect selectProtocol will not trigger rebalance.")
+          val groupAssignment: Map[String, Array[Byte]] = group.allMemberMetadata.map(member => member.memberId -> member.assignment).toMap
+          groupManager.storeGroup(group, groupAssignment, error => {
+            if (error != Errors.NONE) {
+              warn(s"Failed to persist metadata for group ${group.groupId}: ${error.message}")
+
+              // Failed to persist member.id of the given static member, revert the update of the static member in the group.
+              group.updateMember(knownStaticMember, oldProtocols, null)
+              val oldMember = group.replaceGroupInstance(newMemberId, oldMemberId, groupInstanceId)
+              completeAndScheduleNextHeartbeatExpiration(group, oldMember)
+              responseCallback(JoinGroupResult(
+                List.empty,
+                memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID,
+                generationId = group.generationId,
+                protocolType = group.protocolType,
+                protocolName = group.protocolName,
+                leaderId = currentLeader,
+                error = error
+              ))
+            } else {
+              group.maybeInvokeJoinCallback(member, JoinGroupResult(
+                members = List.empty,
+                memberId = newMemberId,
+                generationId = group.generationId,
+                protocolType = group.protocolType,
+                protocolName = group.protocolName,
+                // We want to avoid current leader performing trivial assignment while the group
+                // is in stable stage, because the new assignment in leader's next sync call
+                // won't be broadcast by a stable 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))
+            }
+          })
+        } else {
+          maybePrepareRebalance(group, s"Group's selectedProtocol will change because static member ${member.memberId} with instance id $groupInstanceId joined with change of protocol")
+        }
+      case CompletingRebalance =>
+        // if the group is in after-sync stage, upon getting a new join-group of a known static member
+        // we should still trigger a new rebalance, since the old member may already be sent to the leader
+        // for assignment, and hence when the assignment gets back there would be a mismatch of the old member id
+        // with the new replaced member id. As a result the new member id would not get any assignment.
+        prepareRebalance(group, s"Updating metadata for static member ${member.memberId} with instance id $groupInstanceId")
+      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 =>
+    }
   }
 
   private def updateMemberAndRebalance(group: GroupMetadata,
                                        member: MemberMetadata,
                                        protocols: List[(String, Array[Byte])],
-                                       callback: JoinCallback) {
-    member.supportedProtocols = protocols
-    member.awaitingJoinCallback = callback
-    maybePrepareRebalance(group)
+                                       reason: String,
+                                       callback: JoinCallback): Unit = {
+    group.updateMember(member, protocols, callback)
+    maybePrepareRebalance(group, reason)
   }
 
-  private def maybePrepareRebalance(group: GroupMetadata) {
+  private def maybePrepareRebalance(group: GroupMetadata, reason: String): Unit = {
     group.inLock {
       if (group.canRebalance)
-        prepareRebalance(group)
+        prepareRebalance(group, reason)
     }
   }
 
-  private def prepareRebalance(group: GroupMetadata) {
+  // package private for testing
+  private[group] 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)
@@ -696,43 +1139,74 @@ class GroupCoordinator(val brokerId: Int,
 
     group.transitionTo(PreparingRebalance)
 
-    info(s"Preparing to rebalance group ${group.groupId} with old generation ${group.generationId} " +
-      s"(${Topic.GROUP_METADATA_TOPIC_NAME}-${partitionFor(group.groupId)})")
+    info(s"Preparing to rebalance group ${group.groupId} in state ${group.currentState} with old generation " +
+      s"${group.generationId} (${Topic.GROUP_METADATA_TOPIC_NAME}-${partitionFor(group.groupId)}) (reason: $reason)")
 
     val groupKey = GroupKey(group.groupId)
     joinPurgatory.tryCompleteElseWatch(delayedRebalance, Seq(groupKey))
   }
 
-  private def removeMemberAndUpdateGroup(group: GroupMetadata, member: MemberMetadata) {
+  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, JoinGroupResult(JoinGroupRequest.UNKNOWN_MEMBER_ID, Errors.UNKNOWN_MEMBER_ID))
+
     group.remove(member.memberId)
+    group.removeStaticMember(member.groupInstanceId)
+
     group.currentState match {
       case Dead | Empty =>
-      case Stable | CompletingRebalance => maybePrepareRebalance(group)
+      case Stable | CompletingRebalance => maybePrepareRebalance(group, reason)
       case PreparingRebalance => joinPurgatory.checkAndComplete(GroupKey(group.groupId))
     }
   }
 
+  private def removePendingMemberAndUpdateGroup(group: GroupMetadata, memberId: String): Unit = {
+    group.removePendingMember(memberId)
+
+    if (group.is(PreparingRebalance)) {
+      joinPurgatory.checkAndComplete(GroupKey(group.groupId))
+    }
+  }
+
   def tryCompleteJoin(group: GroupMetadata, forceComplete: () => Boolean) = {
     group.inLock {
-      if (group.notYetRejoinedMembers.isEmpty)
+      if (group.hasAllMembersJoined)
         forceComplete()
       else false
     }
   }
 
-  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 =>
-        group.remove(failedMember.memberId)
-        // TODO: cut the socket connection to the client
+      val notYetRejoinedDynamicMembers = group.notYetRejoinedMembers.filterNot(_._2.isStaticMember)
+      if (notYetRejoinedDynamicMembers.nonEmpty) {
+        info(s"Group ${group.groupId} remove dynamic members " +
+          s"who haven't joined: ${notYetRejoinedDynamicMembers.keySet}")
+
+        notYetRejoinedDynamicMembers.values foreach { failedMember =>
+          removeHeartbeatForLeavingMember(group, failedMember)
+          group.remove(failedMember.memberId)
+          // 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 " +
@@ -752,55 +1226,85 @@ 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 (member.memberId == group.leaderId) {
+              members = if (group.isLeader(member.memberId)) {
                 group.currentMemberMetadata
               } else {
-                Map.empty
+                List.empty
               },
               memberId = member.memberId,
               generationId = group.generationId,
-              subProtocol = group.protocol,
-              leaderId = group.leaderId,
+              protocolType = group.protocolType,
+              protocolName = group.protocolName,
+              leaderId = group.leaderOrNull,
               error = Errors.NONE)
 
-            member.awaitingJoinCallback(joinResult)
-            member.awaitingJoinCallback = null
+            group.maybeInvokeJoinCallback(member, joinResult)
             completeAndScheduleNextHeartbeatExpiration(group, member)
+            member.isNew = false
           }
         }
       }
     }
   }
 
-  def tryCompleteHeartbeat(group: GroupMetadata, member: MemberMetadata, heartbeatDeadline: Long, forceComplete: () => Boolean) = {
+  def tryCompleteHeartbeat(group: GroupMetadata,
+                           memberId: String,
+                           isPending: Boolean,
+                           forceComplete: () => Boolean): Boolean = {
     group.inLock {
-      if (shouldKeepMemberAlive(member, heartbeatDeadline) || member.isLeaving)
+      // The group has been unloaded and invalid, we should complete the heartbeat.
+      if (group.is(Dead)) {
         forceComplete()
-      else false
+      } else if (isPending) {
+        // complete the heartbeat if the member has joined the group
+        if (group.has(memberId)) {
+          forceComplete()
+        } else false
+      } else if (shouldCompleteNonPendingHeartbeat(group, memberId)) {
+        forceComplete()
+      } else false
     }
   }
 
-  def onExpireHeartbeat(group: GroupMetadata, member: MemberMetadata, 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 (!shouldKeepMemberAlive(member, heartbeatDeadline)) {
-        info(s"Member ${member.memberId} in group ${group.groupId} has failed, removing it from the group")
-        removeMemberAndUpdateGroup(group, member)
+      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.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")
+        }
       }
     }
   }
 
-  def onCompleteHeartbeat() {
+  def onCompleteHeartbeat(): Unit = {
     // TODO: add metrics for complete heartbeats
   }
 
   def partitionFor(group: String): Int = groupManager.partitionFor(group)
 
-  private def shouldKeepMemberAlive(member: MemberMetadata, heartbeatDeadline: Long) =
-    member.awaitingJoinCallback != null ||
-      member.awaitingSyncCallback != null ||
-      member.latestHeartbeat + member.sessionTimeoutMs > heartbeatDeadline
+  private def groupIsOverCapacity(group: GroupMetadata): Boolean = {
+    group.size > groupConfig.groupMaxSize
+  }
 
   private def isCoordinatorForGroup(groupId: String) = groupManager.isGroupLocal(groupId)
 
@@ -814,18 +1318,19 @@ object GroupCoordinator {
   val NoProtocol = ""
   val NoLeader = ""
   val NoGeneration = -1
-  val NoMemberId = ""
   val NoMembers = List[MemberSummary]()
   val EmptyGroup = GroupSummary(NoState, NoProtocolType, NoProtocol, NoMembers)
   val DeadGroup = GroupSummary(Dead.toString, NoProtocolType, NoProtocol, NoMembers)
+  val NewMemberJoinTimeoutMs: Int = 5 * 60 * 1000
 
   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(
@@ -846,26 +1351,75 @@ 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,
+      groupMaxSize = config.groupMaxSize,
       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)
+  }
+
+  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,
                        groupMaxSessionTimeoutMs: 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,
+                           protocolType: Option[String],
+                           protocolName: Option[String],
                            leaderId: String,
                            error: Errors)
+
+object JoinGroupResult {
+  def apply(memberId: String, error: Errors): JoinGroupResult = {
+    JoinGroupResult(
+      members = List.empty,
+      memberId = memberId,
+      generationId = GroupCoordinator.NoGeneration,
+      protocolType = None,
+      protocolName = None,
+      leaderId = GroupCoordinator.NoLeader,
+      error = error)
+  }
+}
+
+case class SyncGroupResult(protocolType: Option[String],
+                           protocolName: Option[String],
+                           memberAssignment: Array[Byte],
+                           error: Errors)
+
+object SyncGroupResult {
+  def apply(error: Errors): SyncGroupResult = {
+    SyncGroupResult(None, None, Array.empty, error)
+  }
+}
+
+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 42ca6ea0a3f57..179d671920ed8 100644
--- a/core/src/main/scala/kafka/coordinator/group/GroupMetadata.scala
+++ b/core/src/main/scala/kafka/coordinator/group/GroupMetadata.scala
@@ -16,16 +16,26 @@
  */
 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 kafka.utils.Implicits._
+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.jdk.CollectionConverters._
 
-private[group] sealed trait GroupState
+private[group] sealed trait GroupState {
+  val validPreviousStates: Set[GroupState]
+}
 
 /**
  * Group is preparing to rebalance
@@ -40,7 +50,9 @@ private[group] sealed trait GroupState
  *             all members have left the group => Empty
  *             group is removed by partition emigration => Dead
  */
-private[group] case object PreparingRebalance extends GroupState
+private[group] case object PreparingRebalance extends GroupState {
+  val validPreviousStates: Set[GroupState] = Set(Stable, CompletingRebalance, Empty)
+}
 
 /**
  * Group is awaiting state assignment from the leader
@@ -55,7 +67,9 @@ private[group] case object PreparingRebalance extends GroupState
  *             member failure detected => PreparingRebalance
  *             group is removed by partition emigration => Dead
  */
-private[group] case object CompletingRebalance extends GroupState
+private[group] case object CompletingRebalance extends GroupState {
+  val validPreviousStates: Set[GroupState] = Set(PreparingRebalance)
+}
 
 /**
  * Group is stable
@@ -71,7 +85,9 @@ private[group] case object CompletingRebalance extends GroupState
  *             follower join-group with new metadata => PreparingRebalance
  *             group is removed by partition emigration => Dead
  */
-private[group] case object Stable extends GroupState
+private[group] case object Stable extends GroupState {
+  val validPreviousStates: Set[GroupState] = Set(CompletingRebalance)
+}
 
 /**
  * Group has no more members and its metadata is being removed
@@ -84,7 +100,9 @@ private[group] case object Stable extends GroupState
  *         allow offset fetch requests
  * transition: Dead is a final state before group metadata is cleaned up, so there are no transitions
  */
-private[group] case object Dead extends GroupState
+private[group] case object Dead extends GroupState {
+  val validPreviousStates: Set[GroupState] = Set(Stable, PreparingRebalance, CompletingRebalance, Empty, Dead)
+}
 
 /**
   * Group has no more members, but lingers until all offsets have expired. This state
@@ -101,23 +119,49 @@ private[group] case object Dead extends GroupState
   *             group is removed by partition emigration => Dead
   *             group is removed by expiration => Dead
   */
-private[group] case object Empty extends GroupState
+private[group] case object Empty extends GroupState {
+  val validPreviousStates: Set[GroupState] = Set(PreparingRebalance)
+}
 
 
-private object GroupMetadata {
-  private val validPreviousStates: Map[GroupState, Set[GroupState]] =
-    Map(Dead -> Set(Stable, PreparingRebalance, CompletingRebalance, Empty, Dead),
-      CompletingRebalance -> Set(PreparingRebalance),
-      Stable -> Set(CompletingRebalance),
-      PreparingRebalance -> Set(Stable, CompletingRebalance, Empty),
-      Empty -> Set(PreparingRebalance))
+private object GroupMetadata extends Logging {
+
+  def loadGroup(groupId: String,
+                initialState: GroupState,
+                generationId: Int,
+                protocolType: String,
+                protocolName: String,
+                leaderId: String,
+                currentStateTimestamp: Option[Long],
+                members: Iterable[MemberMetadata],
+                time: Time): GroupMetadata = {
+    val group = new GroupMetadata(groupId, initialState, time)
+    group.generationId = generationId
+    group.protocolType = if (protocolType == null || protocolType.isEmpty) None else Some(protocolType)
+    group.protocolName = Option(protocolName)
+    group.leaderId = Option(leaderId)
+    group.currentStateTimestamp = currentStateTimestamp
+    members.foreach(member => {
+      group.add(member, null)
+      if (member.isStaticMember) {
+        info(s"Static member $member.groupInstanceId of group $groupId loaded " +
+          s"with member id ${member.memberId} at generation ${group.generationId}.")
+        group.addStaticMember(member.groupInstanceId, member.memberId)
+      }
+    })
+    group.subscribedTopics = group.computeSubscribedTopics()
+    group
+  }
+
+  private val MemberIdDelimiter = "-"
 }
 
 /**
  * Case class used to represent group metadata for the ListGroups API
  */
 case class GroupOverview(groupId: String,
-                         protocolType: String)
+                         protocolType: String,
+                         state: String)
 
 /**
  * Case class used to represent group metadata for the DescribeGroup API
@@ -130,11 +174,11 @@ case class GroupSummary(state: String,
 /**
   * We cache offset commits along with their commit record offset. This enables us to ensure that the latest offset
   * commit is always materialized when we have a mix of transactional and regular offset commits. Without preserving
-  * information of the commit record offset, compaction of the offsets topic it self may result in the wrong offset commit
+  * information of the commit record offset, compaction of the offsets topic itself may result in the wrong offset commit
   * 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
 }
 
 /**
@@ -151,28 +195,34 @@ case class CommitRecordMetadataAndOffset(appendedBatchOffset: Option[Long], offs
  *  3. leader id
  */
 @nonthreadsafe
-private[group] class GroupMetadata(val groupId: String, initialState: GroupState = Empty) extends Logging {
-
-  private var state: GroupState = initialState
+private[group] class GroupMetadata(val groupId: String, initialState: GroupState, time: Time) extends Logging {
+  type JoinCallback = JoinGroupResult => Unit
 
   private[group] val lock = new ReentrantLock
 
-  private val members = new mutable.HashMap[String, MemberMetadata]
+  private var state: GroupState = initialState
+  var currentStateTimestamp: Option[Long] = Some(time.milliseconds())
+  var protocolType: Option[String] = None
+  var protocolName: Option[String] = None
+  var generationId = 0
+  private var leaderId: 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)
   private val offsets = new mutable.HashMap[TopicPartition, CommitRecordMetadataAndOffset]
-
   private val pendingOffsetCommits = new mutable.HashMap[TopicPartition, OffsetAndMetadata]
-
   private val pendingTransactionalOffsetCommits = new mutable.HashMap[Long, mutable.Map[TopicPartition, CommitRecordMetadataAndOffset]]()
-
   private var receivedTransactionalOffsetCommits = false
-
   private var receivedConsumerOffsetCommits = false
 
-  var protocolType: Option[String] = None
-  var generationId = 0
-  var leaderId: String = null
-  var protocol: String = null
+  // 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)
@@ -181,51 +231,190 @@ private[group] class GroupMetadata(val groupId: String, initialState: GroupState
   def not(groupState: GroupState) = state != groupState
   def has(memberId: String) = members.contains(memberId)
   def get(memberId: String) = members(memberId)
+  def size = members.size
+
+  def isLeader(memberId: String): Boolean = leaderId.contains(memberId)
+  def leaderOrNull: String = leaderId.orNull
+  def currentStateTimestampOrDefault: Long = currentStateTimestamp.getOrElse(-1)
+
+  def isConsumerGroup: Boolean = protocolType.contains(ConsumerProtocol.PROTOCOL_TYPE)
 
-  def add(member: MemberMetadata) {
+  def add(member: MemberMetadata, callback: JoinCallback = null): Unit = {
     if (members.isEmpty)
       this.protocolType = Some(member.protocolType)
 
     assert(groupId == member.groupId)
     assert(this.protocolType.orNull == member.protocolType)
-    assert(supportsProtocols(member.protocols))
+    assert(supportsProtocols(member.protocolType, MemberMetadata.plainProtocolSet(member.supportedProtocols)))
 
-    if (leaderId == null)
-      leaderId = member.memberId
+    if (leaderId.isEmpty)
+      leaderId = Some(member.memberId)
     members.put(member.memberId, member)
+    member.supportedProtocols.foreach{ case (protocol, _) => supportedProtocols(protocol) += 1 }
+    member.awaitingJoinCallback = callback
+    if (member.isAwaitingJoin)
+      numMembersAwaitingJoin += 1
   }
 
-  def remove(memberId: String) {
-    members.remove(memberId)
-    if (memberId == leaderId) {
-      leaderId = if (members.isEmpty) {
-        null
+  def remove(memberId: String): Unit = {
+    members.remove(memberId).foreach { member =>
+      member.supportedProtocols.foreach{ case (protocol, _) => supportedProtocols(protocol) -= 1 }
+      if (member.isAwaitingJoin)
+        numMembersAwaitingJoin -= 1
+    }
+
+    if (isLeader(memberId))
+      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 {
-        members.keys.head
+        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(oldMemberId, Errors.FENCED_INSTANCE_ID))
+
+    maybeInvokeSyncCallback(oldMember, SyncGroupResult(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(_.awaitingJoinCallback == null).toList
+  def notYetRejoinedMembers = members.filter(!_._2.isAwaitingJoin).toMap
+
+  def hasAllMembersJoined = members.size == numMembersAwaitingJoin && pendingMembers.isEmpty
 
   def allMembers = members.keySet
 
+  def allStaticMembers = staticMembers.keySet
+
+  // For testing only.
+  def allDynamicMembers = {
+    val dynamicMemberSet = new mutable.HashSet[String]
+    allMembers.foreach(memberId => dynamicMemberSet.add(memberId))
+    staticMembers.values.foreach(memberId => dynamicMemberSet.remove(memberId))
+    dynamicMemberSet.toSet
+  }
+
+  def numPending = pendingMembers.size
+
+  def numAwaiting: Int = numMembersAwaitingJoin
+
   def allMemberMetadata = members.values.toList
 
   def rebalanceTimeoutMs = members.values.foldLeft(0) { (timeout, member) =>
     timeout.max(member.rebalanceTimeoutMs)
   }
 
-  // TODO: decide if ids should be predictable or random
-  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],
+                           operation: 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)} during $operation, will " +
+        s"respond with instance fenced error")
+      true
+    } else
+      false
+  }
 
-  def canRebalance = GroupMetadata.validPreviousStates(PreparingRebalance).contains(state)
+  def canRebalance = PreparingRebalance.validPreviousStates.contains(state)
 
-  def transitionTo(groupState: GroupState) {
+  def transitionTo(groupState: GroupState): Unit = {
     assertValidTransition(groupState)
     state = groupState
+    currentStateTimestamp = Some(time.milliseconds())
   }
 
   def selectProtocol: String = {
@@ -236,68 +425,163 @@ private[group] class GroupMetadata(val groupId: String, initialState: GroupState
     val candidates = candidateProtocols
 
     // let each member vote for one of the protocols and choose the one with the most votes
-    val votes: List[(String, Int)] = allMemberMetadata
+    val (protocol, _) = allMemberMetadata
       .map(_.vote(candidates))
       .groupBy(identity)
-      .mapValues(_.size)
-      .toList
+      .maxBy { case (_, votes) => votes.size }
 
-    votes.maxBy(_._2)._1
+    protocol
   }
 
-  private def candidateProtocols = {
+  private def candidateProtocols: Set[String] = {
     // get the set of protocols that are commonly supported by all members
-    allMemberMetadata
-      .map(_.protocols)
-      .reduceLeft((commonProtocols, protocols) => commonProtocols & protocols)
+    val numMembers = members.size
+    supportedProtocols.filter(_._2 == numMembers).map(_._1).toSet
+  }
+
+  def supportsProtocols(memberProtocolType: String, memberProtocols: Set[String]): Boolean = {
+    if (is(Empty))
+      !memberProtocolType.isEmpty && memberProtocols.nonEmpty
+    else
+      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 && protocolName.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(protocolName.get))
+              ConsumerProtocol.deserializeVersion(buffer)
+              ConsumerProtocol.deserializeSubscription(buffer, 0).topics.asScala.toSet
+            }.reduceLeft(_ ++ _)
+          )
+        } catch {
+          case e: SchemaException =>
+            warn(s"Failed to parse Consumer Protocol ${ConsumerProtocol.PROTOCOL_TYPE}:${protocolName.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 supportsProtocols(memberProtocols: Set[String]) = {
-    members.isEmpty || (memberProtocols & candidateProtocols).nonEmpty
+  def updateMember(member: MemberMetadata,
+                   protocols: List[(String, Array[Byte])],
+                   callback: JoinCallback): Unit = {
+    member.supportedProtocols.foreach{ case (protocol, _) => supportedProtocols(protocol) -= 1 }
+    protocols.foreach{ case (protocol, _) => supportedProtocols(protocol) += 1 }
+    member.supportedProtocols = protocols
+
+    if (callback != null && !member.isAwaitingJoin) {
+      numMembersAwaitingJoin += 1
+    } else if (callback == null && member.isAwaitingJoin) {
+      numMembersAwaitingJoin -= 1
+    }
+    member.awaitingJoinCallback = callback
+  }
+
+  def maybeInvokeJoinCallback(member: MemberMetadata,
+                              joinGroupResult: JoinGroupResult): Unit = {
+    if (member.isAwaitingJoin) {
+      member.awaitingJoinCallback(joinGroupResult)
+      member.awaitingJoinCallback = null
+      numMembersAwaitingJoin -= 1
+    }
+  }
+
+  /**
+    * @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])
+  def initNextGeneration(): Unit = {
     if (members.nonEmpty) {
       generationId += 1
-      protocol = selectProtocol
+      protocolName = Some(selectProtocol)
+      subscribedTopics = computeSubscribedTopics()
       transitionTo(CompletingRebalance)
     } else {
       generationId += 1
-      protocol = null
+      protocolName = 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))}.toMap
+    members.map{ case (memberId, memberMetadata) => new JoinGroupResponseMember()
+        .setMemberId(memberId)
+        .setGroupInstanceId(memberMetadata.groupInstanceId.orNull)
+        .setMetadata(memberMetadata.metadata(protocolName.get))
+    }.toList
   }
 
   def summary: GroupSummary = {
     if (is(Stable)) {
-      val members = this.members.values.map { member => member.summary(protocol) }.toList
-      GroupSummary(state.toString, protocolType.getOrElse(""), protocol, members)
+      val protocol = protocolName.orNull
+      if (protocol == null)
+        throw new IllegalStateException("Invalid null group protocol for stable group")
+
+      val members = this.members.values.map { member => member.summary(protocol) }
+      GroupSummary(state.toString, protocolType.getOrElse(""), protocol, members.toList)
     } else {
-      val members = this.members.values.map{ member => member.summaryNoMetadata() }.toList
-      GroupSummary(state.toString, protocolType.getOrElse(""), GroupCoordinator.NoProtocol, members)
+      val members = this.members.values.map{ member => member.summaryNoMetadata() }
+      GroupSummary(state.toString, protocolType.getOrElse(""), GroupCoordinator.NoProtocol, members.toList)
     }
   }
 
   def overview: GroupOverview = {
-    GroupOverview(groupId, protocolType.getOrElse(""))
+    GroupOverview(groupId, protocolType.getOrElse(""), state.toString)
   }
 
   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 " +
@@ -322,18 +606,18 @@ 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,
       mutable.Map.empty[TopicPartition, CommitRecordMetadataAndOffset])
 
-    offsets.foreach { case (topicPartition, offsetAndMetadata) =>
+    offsets.forKeyValue { (topicPartition, offsetAndMetadata) =>
       producerOffsets.put(topicPartition, CommitRecordMetadataAndOffset(None, offsetAndMetadata))
     }
   }
@@ -359,7 +643,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)
@@ -377,7 +661,7 @@ private[group] class GroupMetadata(val groupId: String, initialState: GroupState
     val pendingOffsetsOpt = pendingTransactionalOffsetCommits.remove(producerId)
     if (isCommit) {
       pendingOffsetsOpt.foreach { pendingOffsets =>
-        pendingOffsets.foreach { case (topicPartition, commitRecordMetadataAndOffset) =>
+        pendingOffsets.forKeyValue { (topicPartition, commitRecordMetadataAndOffset) =>
           if (commitRecordMetadataAndOffset.appendedBatchOffset.isEmpty)
             throw new IllegalStateException(s"Trying to complete a transactional offset commit for producerId $producerId " +
               s"and groupId $groupId even though the offset commit record itself hasn't been appended to the log.")
@@ -398,16 +682,24 @@ private[group] class GroupMetadata(val groupId: String, initialState: GroupState
     }
   }
 
-  def activeProducers = pendingTransactionalOffsetCommits.keySet
+  def activeProducers: collection.Set[Long] = pendingTransactionalOffsetCommits.keySet
 
-  def hasPendingOffsetCommitsFromProducer(producerId: Long) =
+  def hasPendingOffsetCommitsFromProducer(producerId: Long): Boolean =
     pendingTransactionalOffsetCommits.contains(producerId)
 
+  def hasPendingOffsetCommitsForTopicPartition(topicPartition: TopicPartition): Boolean = {
+    pendingOffsetCommits.contains(topicPartition) ||
+      pendingTransactionalOffsetCommits.exists(
+        _._2.contains(topicPartition)
+      )
+  }
+
+  def removeAllOffsets(): immutable.Map[TopicPartition, OffsetAndMetadata] = removeOffsets(offsets.keySet.toSeq)
+
   def removeOffsets(topicPartitions: Seq[TopicPartition]): immutable.Map[TopicPartition, OffsetAndMetadata] = {
     topicPartitions.flatMap { topicPartition =>
-
       pendingOffsetCommits.remove(topicPartition)
-      pendingTransactionalOffsetCommits.foreach { case (_, pendingOffsets) =>
+      pendingTransactionalOffsetCommits.forKeyValue { (_, pendingOffsets) =>
         pendingOffsets.remove(topicPartition)
       }
       val removedOffset = offsets.remove(topicPartition)
@@ -415,18 +707,65 @@ private[group] class GroupMetadata(val groupId: String, initialState: GroupState
     }.toMap
   }
 
-  def removeExpiredOffsets(startMs: Long) : Map[TopicPartition, OffsetAndMetadata] = {
-    val expiredOffsets = offsets
-      .filter {
+  def removeExpiredOffsets(currentTimestamp: Long, offsetRetentionMs: Long): Map[TopicPartition, OffsetAndMetadata] = {
+
+    def getExpiredOffsets(baseTimestamp: CommitRecordMetadataAndOffset => Long,
+                          subscribedTopics: Set[String] = Set.empty): Map[TopicPartition, OffsetAndMetadata] = {
+      offsets.filter {
         case (topicPartition, commitRecordMetadataAndOffset) =>
-          commitRecordMetadataAndOffset.offsetAndMetadata.expireTimestamp < startMs && !pendingOffsetCommits.contains(topicPartition)
-      }
-      .map {
+          !subscribedTopics.contains(topicPartition.topic()) &&
+          !pendingOffsetCommits.contains(topicPartition) && {
+            commitRecordMetadataAndOffset.offsetAndMetadata.expireTimestamp match {
+              case None =>
+                // current version with no per partition retention
+                currentTimestamp - baseTimestamp(commitRecordMetadataAndOffset) >= offsetRetentionMs
+              case Some(expireTimestamp) =>
+                // older versions with explicit expire_timestamp field => old expiration semantics is used
+                currentTimestamp >= expireTimestamp
+            }
+          }
+      }.map {
         case (topicPartition, commitRecordOffsetAndMetadata) =>
           (topicPartition, commitRecordOffsetAndMetadata.offsetAndMetadata)
-      }
+      }.toMap
+    }
+
+    val expiredOffsets: Map[TopicPartition, OffsetAndMetadata] = protocolType match {
+      case Some(_) if is(Empty) =>
+        // no consumer exists in the group =>
+        // - if current state timestamp exists and retention period has passed since group became Empty,
+        //   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)
+        )
+
+      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
+        // expire offsets with no pending offset commit that retention period has passed since their last commit
+        getExpiredOffsets(_.offsetAndMetadata.commitTimestamp)
+
+      case _ =>
+        Map()
+    }
+
+    if (expiredOffsets.nonEmpty)
+      debug(s"Expired offsets from group '$groupId': ${expiredOffsets.keySet}")
+
     offsets --= expiredOffsets.keySet
-    expiredOffsets.toMap
+    expiredOffsets
   }
 
   def allOffsets = offsets.map { case (topicPartition, commitRecordMetadataAndOffset) =>
@@ -442,10 +781,10 @@ private[group] class GroupMetadata(val groupId: String, initialState: GroupState
 
   def hasOffsets = offsets.nonEmpty || pendingOffsetCommits.nonEmpty || pendingTransactionalOffsetCommits.nonEmpty
 
-  private def assertValidTransition(targetState: GroupState) {
-    if (!GroupMetadata.validPreviousStates(targetState).contains(state))
+  private def assertValidTransition(targetState: GroupState): Unit = {
+    if (!targetState.validPreviousStates.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))
+        .format(groupId, targetState.validPreviousStates.mkString(","), targetState, state))
   }
 
   override def toString: String = {
diff --git a/core/src/main/scala/kafka/coordinator/group/GroupMetadataManager.scala b/core/src/main/scala/kafka/coordinator/group/GroupMetadataManager.scala
index 7e6a643ddb52f..5fefc15a33604 100644
--- a/core/src/main/scala/kafka/coordinator/group/GroupMetadataManager.scala
+++ b/core/src/main/scala/kafka/coordinator/group/GroupMetadataManager.scala
@@ -20,39 +20,46 @@ package kafka.coordinator.group
 import java.io.PrintStream
 import java.nio.ByteBuffer
 import java.nio.charset.StandardCharsets
+import java.util.Optional
 import java.util.concurrent.TimeUnit
 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}
-import kafka.common.{MessageFormatter, _}
+import kafka.api.{ApiVersion, KAFKA_0_10_1_IV0, KAFKA_2_1_IV0, KAFKA_2_1_IV1, KAFKA_2_3_IV0}
+import kafka.common.OffsetAndMetadata
+import kafka.internals.generated.{GroupMetadataValue, OffsetCommitKey, OffsetCommitValue, GroupMetadataKey => GroupMetadataKeyData}
+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.Implicits._
 import kafka.utils._
 import kafka.zk.KafkaZkClient
 import org.apache.kafka.clients.consumer.ConsumerRecord
-import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.clients.consumer.internals.ConsumerProtocol
 import org.apache.kafka.common.internals.Topic
-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.metrics.Metrics
+import org.apache.kafka.common.metrics.stats.{Avg, Max, Meter}
+import org.apache.kafka.common.protocol.{ByteBufferAccessor, Errors, MessageUtil}
 import org.apache.kafka.common.record._
-import org.apache.kafka.common.requests.{IsolationLevel, OffsetFetchResponse}
+import org.apache.kafka.common.requests.OffsetFetchResponse.PartitionData
 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, MessageFormatter, TopicPartition}
 
-import scala.collection.JavaConverters._
 import scala.collection._
-import scala.collection.mutable.ListBuffer
+import scala.collection.mutable.ArrayBuffer
+import scala.jdk.CollectionConverters._
 
 class GroupMetadataManager(brokerId: Int,
                            interBrokerProtocolVersion: ApiVersion,
                            config: OffsetConfig,
-                           replicaManager: ReplicaManager,
+                           val 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)
 
@@ -81,6 +88,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*/
+  private val partitionLoadSensor = metrics.sensor(GroupMetadataManager.LoadTimeSensor)
+
+  partitionLoadSensor.add(metrics.metricName("partition-load-time-max",
+    GroupMetadataManager.MetricsGroup,
+    "The max time it took to load the partitions in the last 30sec"), new Max())
+  partitionLoadSensor.add(metrics.metricName("partition-load-time-avg",
+    GroupMetadataManager.MetricsGroup,
+    "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] = {
@@ -89,59 +126,58 @@ class GroupMetadataManager(brokerId: Int,
   }
 
   recreateGauge("NumOffsets",
-    new Gauge[Int] {
-      def value = groupMetadataCache.values.map(group => {
-        group.inLock { group.numOffsets }
-      }).sum
-    })
+    () => groupMetadataCache.values.map { group =>
+      group.inLock { group.numOffsets }
+    }.sum
+  )
 
   recreateGauge("NumGroups",
-    new Gauge[Int] {
-      def value = groupMetadataCache.size
-    })
+    () => groupMetadataCache.size
+  )
 
   recreateGauge("NumGroupsPreparingRebalance",
-    new Gauge[Int] {
-      def value(): Int = groupMetadataCache.values.count(group => {
-        group synchronized { group.is(PreparingRebalance) }
-      })
+    () => groupMetadataCache.values.count { group =>
+      group synchronized {
+        group.is(PreparingRebalance)
+      }
     })
 
   recreateGauge("NumGroupsCompletingRebalance",
-    new Gauge[Int] {
-      def value(): Int = groupMetadataCache.values.count(group => {
-        group synchronized { group.is(CompletingRebalance) }
-      })
+    () => groupMetadataCache.values.count { group =>
+      group synchronized {
+        group.is(CompletingRebalance)
+      }
     })
 
   recreateGauge("NumGroupsStable",
-    new Gauge[Int] {
-      def value(): Int = groupMetadataCache.values.count(group => {
-        group synchronized { group.is(Stable) }
-      })
+    () => groupMetadataCache.values.count { group =>
+      group synchronized {
+        group.is(Stable)
+      }
     })
 
   recreateGauge("NumGroupsDead",
-    new Gauge[Int] {
-      def value(): Int = groupMetadataCache.values.count(group => {
-        group synchronized { group.is(Dead) }
-      })
+    () => groupMetadataCache.values.count { group =>
+      group synchronized {
+        group.is(Dead)
+      }
     })
 
   recreateGauge("NumGroupsEmpty",
-    new Gauge[Int] {
-      def value(): Int = groupMetadataCache.values.count(group => {
-        group synchronized { group.is(Empty) }
-      })
+    () => groupMetadataCache.values.count { group =>
+      group synchronized {
+        group.is(Empty)
+      }
     })
 
-  def enableMetadataExpiration() {
+  def startup(enableMetadataExpiration: Boolean): Unit = {
     scheduler.startup()
-
-    scheduler.schedule(name = "delete-expired-group-metadata",
-      fun = cleanupGroupMetadata,
-      period = config.offsetsRetentionCheckIntervalMs,
-      unit = TimeUnit.MILLISECONDS)
+    if (enableMetadataExpiration) {
+      scheduler.schedule(name = "delete-expired-group-metadata",
+        fun = () => cleanupGroupMetadata(),
+        period = config.offsetsRetentionCheckIntervalMs,
+        unit = TimeUnit.MILLISECONDS)
+    }
   }
 
   def currentGroups: Iterable[GroupMetadata] = groupMetadataCache.values
@@ -156,7 +192,14 @@ class GroupMetadataManager(brokerId: Int,
 
   def isGroupLoading(groupId: String): Boolean = isPartitionLoading(partitionFor(groupId))
 
-  def isLoading(): Boolean = inLock(partitionLock) { loadingPartitions.nonEmpty }
+  def isLoading: Boolean = inLock(partitionLock) { loadingPartitions.nonEmpty }
+
+  // return true iff group is owned and the group doesn't exist
+  def groupNotExists(groupId: String) = inLock(partitionLock) {
+    isGroupLocal(groupId) && getGroup(groupId).forall { group =>
+      group.inLock(group.is(Dead))
+    }
+  }
 
   // visible for testing
   private[group] def isGroupOpenForProducer(producerId: Long, groupId: String) = openGroupsForProducer.get(producerId) match {
@@ -165,13 +208,25 @@ class GroupMetadataManager(brokerId: Int,
     case None =>
       false
   }
+
   /**
-   * Get the group associated with the given groupId, or null if not found
+   * Get the group associated with the given groupId or null if not found
    */
   def getGroup(groupId: String): Option[GroupMetadata] = {
     Option(groupMetadataCache.get(groupId))
   }
 
+  /**
+   * Get the group associated with the given groupId - the group is created if createIfNotExist
+   * is true - or null if not found
+   */
+  def getOrMaybeCreateGroup(groupId: String, createIfNotExist: Boolean): Option[GroupMetadata] = {
+    if (createIfNotExist)
+      Option(groupMetadataCache.getAndMaybePut(groupId, new GroupMetadata(groupId, Empty, time)))
+    else
+      Option(groupMetadataCache.get(groupId))
+  }
+
   /**
    * Add a group or get the group associated with the given groupId if it already exists
    */
@@ -189,18 +244,11 @@ class GroupMetadataManager(brokerId: Int,
                  responseCallback: Errors => Unit): Unit = {
     getMagic(partitionFor(group.groupId)) match {
       case Some(magicValue) =>
-        val groupMetadataValueVersion = {
-          if (interBrokerProtocolVersion < KAFKA_0_10_1_IV0)
-            0.toShort
-          else
-            GroupMetadataManager.CURRENT_GROUP_VALUE_SCHEMA_VERSION
-        }
-
         // We always use CREATE_TIME, like the producer. The conversion to LOG_APPEND_TIME (if necessary) happens automatically.
         val timestampType = TimestampType.CREATE_TIME
         val timestamp = time.milliseconds()
         val key = GroupMetadataManager.groupMetadataKey(group.groupId)
-        val value = GroupMetadataManager.groupMetadataValue(group, groupAssignment, version = groupMetadataValueVersion)
+        val value = GroupMetadataManager.groupMetadataValue(group, groupAssignment, interBrokerProtocolVersion)
 
         val records = {
           val buffer = ByteBuffer.allocate(AbstractRecords.estimateSizeInBytes(magicValue, compressionType,
@@ -215,14 +263,13 @@ 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"
               .format(responseStatus, groupMetadataPartition))
 
-          // construct the error status in the propagated assignment response
-          // in the cache
+          // construct the error status in the propagated assignment response in the cache
           val status = responseStatus(groupMetadataPartition)
 
           val responseError = if (status.error == Errors.NONE) {
@@ -238,7 +285,7 @@ class GroupMetadataManager(brokerId: Int,
                    | Errors.NOT_ENOUGH_REPLICAS_AFTER_APPEND =>
                 Errors.COORDINATOR_NOT_AVAILABLE
 
-              case Errors.NOT_LEADER_FOR_PARTITION
+              case Errors.NOT_LEADER_OR_FOLLOWER
                    | Errors.KAFKA_STORAGE_ERROR =>
                 Errors.NOT_COORDINATOR
 
@@ -280,7 +327,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)
@@ -302,7 +349,7 @@ class GroupMetadataManager(brokerId: Int,
 
     group.inLock {
       if (!group.hasReceivedConsistentOffsetCommits)
-        warn(s"group: ${group.groupId} with leader: ${group.leaderId} has received offset commits from consumers as well " +
+        warn(s"group: ${group.groupId} with leader: ${group.leaderOrNull} has received offset commits from consumers as well " +
           s"as transactional producers. Mixing both types of offset commits will generally result in surprises and " +
           s"should be avoided.")
     }
@@ -311,9 +358,8 @@ 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 {
       getMagic(partitionFor(group.groupId)) match {
         case Some(magicValue) =>
@@ -323,7 +369,7 @@ class GroupMetadataManager(brokerId: Int,
 
           val records = filteredOffsetMetadata.map { case (topicPartition, offsetAndMetadata) =>
             val key = GroupMetadataManager.offsetCommitKey(group.groupId, topicPartition)
-            val value = GroupMetadataManager.offsetCommitValue(offsetAndMetadata)
+            val value = GroupMetadataManager.offsetCommitValue(offsetAndMetadata, interBrokerProtocolVersion)
             new SimpleRecord(timestamp, key, value)
           }
           val offsetTopicPartition = new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, partitionFor(group.groupId))
@@ -339,12 +385,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)
@@ -352,7 +401,7 @@ class GroupMetadataManager(brokerId: Int,
             val responseError = group.inLock {
               if (status.error == Errors.NONE) {
                 if (!group.is(Dead)) {
-                  filteredOffsetMetadata.foreach { case (topicPartition, offsetAndMetadata) =>
+                  filteredOffsetMetadata.forKeyValue { (topicPartition, offsetAndMetadata) =>
                     if (isTxnOffsetCommit)
                       group.onTxnOffsetCommitAppend(producerId, topicPartition, CommitRecordMetadataAndOffset(Some(status.baseOffset), offsetAndMetadata))
                     else
@@ -364,7 +413,7 @@ class GroupMetadataManager(brokerId: Int,
                 if (!group.is(Dead)) {
                   if (!group.hasPendingOffsetCommitsFromProducer(producerId))
                     removeProducerGroup(producerId, group.groupId)
-                  filteredOffsetMetadata.foreach { case (topicPartition, offsetAndMetadata) =>
+                  filteredOffsetMetadata.forKeyValue { (topicPartition, offsetAndMetadata) =>
                     if (isTxnOffsetCommit)
                       group.failPendingTxnOffsetCommit(producerId, topicPartition)
                     else
@@ -382,7 +431,7 @@ class GroupMetadataManager(brokerId: Int,
                        | Errors.NOT_ENOUGH_REPLICAS_AFTER_APPEND =>
                     Errors.COORDINATOR_NOT_AVAILABLE
 
-                  case Errors.NOT_LEADER_FOR_PARTITION
+                  case Errors.NOT_LEADER_OR_FOLLOWER
                        | Errors.KAFKA_STORAGE_ERROR =>
                     Errors.NOT_COORDINATOR
 
@@ -426,7 +475,6 @@ class GroupMetadataManager(brokerId: Int,
             (topicPartition, Errors.NOT_COORDINATOR)
           }
           responseCallback(commitStatus)
-          None
       }
     }
   }
@@ -435,39 +483,42 @@ class GroupMetadataManager(brokerId: Int,
    * The most important guarantee that this API provides is that it should never return a stale offset. i.e., it either
    * returns the current offset or it begins to sync the cache from the log (and returns an error code).
    */
-  def getOffsets(groupId: String, topicPartitionsOpt: Option[Seq[TopicPartition]]): Map[TopicPartition, OffsetFetchResponse.PartitionData] = {
+  def getOffsets(groupId: String, requireStable: Boolean, topicPartitionsOpt: Option[Seq[TopicPartition]]): Map[TopicPartition, PartitionData] = {
     trace("Getting offsets of %s for group %s.".format(topicPartitionsOpt.getOrElse("all partitions"), groupId))
     val group = groupMetadataCache.get(groupId)
     if (group == null) {
       topicPartitionsOpt.getOrElse(Seq.empty[TopicPartition]).map { topicPartition =>
-        (topicPartition, new OffsetFetchResponse.PartitionData(OffsetFetchResponse.INVALID_OFFSET, "", Errors.NONE))
+        val partitionData = new PartitionData(OffsetFetchResponse.INVALID_OFFSET,
+          Optional.empty(), "", Errors.NONE)
+        topicPartition -> partitionData
       }.toMap
     } else {
       group.inLock {
         if (group.is(Dead)) {
           topicPartitionsOpt.getOrElse(Seq.empty[TopicPartition]).map { topicPartition =>
-            (topicPartition, new OffsetFetchResponse.PartitionData(OffsetFetchResponse.INVALID_OFFSET, "", Errors.NONE))
+            val partitionData = new PartitionData(OffsetFetchResponse.INVALID_OFFSET,
+              Optional.empty(), "", Errors.NONE)
+            topicPartition -> partitionData
           }.toMap
         } else {
-          topicPartitionsOpt match {
-            case None =>
-              // Return offsets for all partitions owned by this consumer group. (this only applies to consumers
-              // that commit offsets to Kafka.)
-              group.allOffsets.map { case (topicPartition, offsetAndMetadata) =>
-                topicPartition -> new OffsetFetchResponse.PartitionData(offsetAndMetadata.offset, offsetAndMetadata.metadata, Errors.NONE)
-              }
+          val topicPartitions = topicPartitionsOpt.getOrElse(group.allOffsets.keySet)
 
-            case Some(topicPartitions) =>
-              topicPartitionsOpt.getOrElse(Seq.empty[TopicPartition]).map { topicPartition =>
-                val partitionData = group.offset(topicPartition) match {
-                  case None =>
-                    new OffsetFetchResponse.PartitionData(OffsetFetchResponse.INVALID_OFFSET, "", Errors.NONE)
-                  case Some(offsetAndMetadata) =>
-                    new OffsetFetchResponse.PartitionData(offsetAndMetadata.offset, offsetAndMetadata.metadata, Errors.NONE)
-                }
-                topicPartition -> partitionData
-              }.toMap
-          }
+          topicPartitions.map { topicPartition =>
+            if (requireStable && group.hasPendingOffsetCommitsForTopicPartition(topicPartition)) {
+              topicPartition -> new PartitionData(OffsetFetchResponse.INVALID_OFFSET,
+                Optional.empty(), "", Errors.UNSTABLE_OFFSET_COMMIT)
+            } else {
+              val partitionData = group.offset(topicPartition) match {
+                case None =>
+                  new PartitionData(OffsetFetchResponse.INVALID_OFFSET,
+                    Optional.empty(), "", Errors.NONE)
+                case Some(offsetAndMetadata) =>
+                  new PartitionData(offsetAndMetadata.offset,
+                    offsetAndMetadata.leaderEpoch, offsetAndMetadata.metadata, Errors.NONE)
+              }
+              topicPartition -> partitionData
+            }
+          }.toMap
         }
       }
     }
@@ -476,81 +527,112 @@ class GroupMetadataManager(brokerId: Int,
   /**
    * Asynchronously read the partition from the offsets topic and populate the cache
    */
-  def loadGroupsForPartition(offsetsPartition: Int, onGroupLoaded: GroupMetadata => Unit) {
+  def scheduleLoadGroupAndOffsets(offsetsPartition: Int, onGroupLoaded: GroupMetadata => Unit): Unit = {
     val topicPartition = new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, offsetsPartition)
-    info(s"Scheduling loading of offsets and group metadata from $topicPartition")
-    scheduler.schedule(topicPartition.toString, doLoadGroupsAndOffsets)
+    if (addLoadingPartition(offsetsPartition)) {
+      info(s"Scheduling loading of offsets and group metadata from $topicPartition")
+      val startTimeMs = time.milliseconds()
+      scheduler.schedule(topicPartition.toString, () => loadGroupsAndOffsets(topicPartition, onGroupLoaded, startTimeMs))
+    } else {
+      info(s"Already loading offsets and group metadata from $topicPartition")
+    }
+  }
 
-    def doLoadGroupsAndOffsets() {
+  private[group] def loadGroupsAndOffsets(topicPartition: TopicPartition, onGroupLoaded: GroupMetadata => Unit, startTimeMs: java.lang.Long): Unit = {
+    try {
+      val schedulerTimeMs = time.milliseconds() - startTimeMs
+      debug(s"Started loading offsets and group metadata from $topicPartition")
+      doLoadGroupsAndOffsets(topicPartition, onGroupLoaded)
+      val endTimeMs = time.milliseconds()
+      val totalLoadingTimeMs = endTimeMs - startTimeMs
+      partitionLoadSensor.record(totalLoadingTimeMs.toDouble, endTimeMs, false)
+      info(s"Finished loading offsets and group metadata from $topicPartition "
+        + s"in $totalLoadingTimeMs milliseconds, of which $schedulerTimeMs milliseconds"
+        + s" was spent in the scheduler.")
+    } catch {
+      case t: Throwable => error(s"Error loading offsets from $topicPartition", t)
+    } finally {
       inLock(partitionLock) {
-        if (loadingPartitions.contains(offsetsPartition)) {
-          info(s"Offset load from $topicPartition already in progress.")
-          return
-        } else {
-          loadingPartitions.add(offsetsPartition)
-        }
-      }
-
-      try {
-        val startMs = time.milliseconds()
-        loadGroupsAndOffsets(topicPartition, onGroupLoaded)
-        info(s"Finished loading offsets and group metadata from $topicPartition in ${time.milliseconds() - startMs} milliseconds.")
-      } catch {
-        case t: Throwable => error(s"Error loading offsets from $topicPartition", t)
-      } finally {
-        inLock(partitionLock) {
-          ownedPartitions.add(offsetsPartition)
-          loadingPartitions.remove(offsetsPartition)
-        }
+        ownedPartitions.add(topicPartition.partition)
+        loadingPartitions.remove(topicPartition.partition)
       }
     }
   }
 
-  private[group] def loadGroupsAndOffsets(topicPartition: TopicPartition, onGroupLoaded: GroupMetadata => Unit) {
-    def highWaterMark = replicaManager.getLogEndOffset(topicPartition).getOrElse(-1L)
+  private def doLoadGroupsAndOffsets(topicPartition: TopicPartition, onGroupLoaded: GroupMetadata => Unit): Unit = {
+    def logEndOffset: Long = replicaManager.getLogEndOffset(topicPartition).getOrElse(-1L)
 
     replicaManager.getLog(topicPartition) match {
       case None =>
         warn(s"Attempted to load offsets and group metadata from $topicPartition, but found no log")
 
       case Some(log) =>
-        var currOffset = log.logStartOffset
-        lazy val buffer = ByteBuffer.allocate(config.loadBufferSize)
-
-        // 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 < highWaterMark && !shuttingDown.get()) {
-          val fetchDataInfo = log.read(currOffset, config.loadBufferSize, maxOffset = None,
-            minOneMessage = true, isolationLevel = IsolationLevel.READ_UNCOMMITTED)
-          val memRecords = fetchDataInfo.records match {
+        // 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: @unchecked) match {
             case records: MemoryRecords => records
             case fileRecords: FileRecords =>
-              buffer.clear()
-              val bufferRead = fileRecords.readInto(buffer, 0)
-              MemoryRecords.readableRecords(bufferRead)
+              val sizeInBytes = fileRecords.sizeInBytes
+              val bytesNeeded = Math.max(config.loadBufferSize, sizeInBytes)
+
+              // 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.loadBufferSize < bytesNeeded)
+                  warn(s"Loaded offsets and group metadata from $topicPartition with buffer larger ($bytesNeeded bytes) than " +
+                    s"configured offsets.load.buffer.size (${config.loadBufferSize} bytes)")
+
+                buffer = ByteBuffer.allocate(bytesNeeded)
+              } else {
+                buffer.clear()
+              }
+
+              fileRecords.readInto(buffer, 0)
+              MemoryRecords.readableRecords(buffer)
           }
 
-          memRecords.batches.asScala.foreach { batch =>
+          memRecords.batches.forEach { batch =>
             val isTxnOffsetCommit = batch.isTransactional
             if (batch.isControlBatch) {
-              val record = batch.iterator.next()
-              val controlRecord = ControlRecordType.parse(record.key)
-              if (controlRecord == ControlRecordType.COMMIT) {
-                pendingOffsets.getOrElse(batch.producerId, mutable.Map[GroupTopicPartition, CommitRecordMetadataAndOffset]())
-                  .foreach {
-                    case (groupTopicPartition, commitRecordMetadataAndOffset) =>
-                      if (!loadedOffsets.contains(groupTopicPartition) || loadedOffsets(groupTopicPartition).olderThan(commitRecordMetadataAndOffset))
-                        loadedOffsets.put(groupTopicPartition, commitRecordMetadataAndOffset)
-                  }
+              val recordIterator = batch.iterator
+              if (recordIterator.hasNext) {
+                val record = recordIterator.next()
+                val controlRecord = ControlRecordType.parse(record.key)
+                if (controlRecord == ControlRecordType.COMMIT) {
+                  pendingOffsets.getOrElse(batch.producerId, mutable.Map[GroupTopicPartition, CommitRecordMetadataAndOffset]())
+                    .foreach {
+                      case (groupTopicPartition, commitRecordMetadataAndOffset) =>
+                        if (!loadedOffsets.contains(groupTopicPartition) || loadedOffsets(groupTopicPartition).olderThan(commitRecordMetadataAndOffset))
+                          loadedOffsets.put(groupTopicPartition, commitRecordMetadataAndOffset)
+                    }
+                }
+                pendingOffsets.remove(batch.producerId)
               }
-              pendingOffsets.remove(batch.producerId)
             } else {
+              var batchBaseOffset: Option[Long] = None
               for (record <- batch.asScala) {
                 require(record.hasKey, "Group metadata/offset entry key should not be null")
+                if (batchBaseOffset.isEmpty)
+                  batchBaseOffset = Some(record.offset)
                 GroupMetadataManager.readMessageKey(record.key) match {
 
                   case offsetKey: OffsetKey =>
@@ -567,15 +649,15 @@ class GroupMetadataManager(brokerId: Int,
                     } else {
                       val offsetAndMetadata = GroupMetadataManager.readOffsetMessageValue(record.value)
                       if (isTxnOffsetCommit)
-                        pendingOffsets(batch.producerId).put(groupTopicPartition, CommitRecordMetadataAndOffset(Some(batch.baseOffset), offsetAndMetadata))
+                        pendingOffsets(batch.producerId).put(groupTopicPartition, CommitRecordMetadataAndOffset(batchBaseOffset, offsetAndMetadata))
                       else
-                        loadedOffsets.put(groupTopicPartition, CommitRecordMetadataAndOffset(Some(batch.baseOffset), offsetAndMetadata))
+                        loadedOffsets.put(groupTopicPartition, CommitRecordMetadataAndOffset(batchBaseOffset, offsetAndMetadata))
                     }
 
                   case groupMetadataKey: GroupMetadataKey =>
                     // load group metadata
                     val groupId = groupMetadataKey.key
-                    val groupMetadata = GroupMetadataManager.readGroupMessageValue(groupId, record.value)
+                    val groupMetadata = GroupMetadataManager.readGroupMessageValue(groupId, record.value, time)
                     if (groupMetadata != null) {
                       removedGroups.remove(groupId)
                       loadedGroups.put(groupId, groupMetadata)
@@ -595,19 +677,21 @@ class GroupMetadataManager(brokerId: Int,
 
         val (groupOffsets, emptyGroupOffsets) = loadedOffsets
           .groupBy(_._1.group)
-          .mapValues(_.map { case (groupTopicPartition, offset) => (groupTopicPartition.topicPartition, offset) })
-          .partition { case (group, _) => loadedGroups.contains(group) }
+          .map { case (k, v) =>
+            k -> v.map { case (groupTopicPartition, offset) => (groupTopicPartition.topicPartition, offset) }
+          }.partition { case (group, _) => loadedGroups.contains(group) }
 
         val pendingOffsetsByGroup = mutable.Map[String, mutable.Map[Long, mutable.Map[TopicPartition, CommitRecordMetadataAndOffset]]]()
-        pendingOffsets.foreach { case (producerId, producerOffsets) =>
+        pendingOffsets.forKeyValue { (producerId, producerOffsets) =>
           producerOffsets.keySet.map(_.group).foreach(addProducerGroup(producerId, _))
           producerOffsets
             .groupBy(_._1.group)
-            .mapValues(_.map { case (groupTopicPartition, offset) => (groupTopicPartition.topicPartition, offset)})
-            .foreach { case (group, offsets) =>
+            .forKeyValue { (group, offsets) =>
               val groupPendingOffsets = pendingOffsetsByGroup.getOrElseUpdate(group, mutable.Map.empty[Long, mutable.Map[TopicPartition, CommitRecordMetadataAndOffset]])
               val groupProducerOffsets = groupPendingOffsets.getOrElseUpdate(producerId, mutable.Map.empty[TopicPartition, CommitRecordMetadataAndOffset])
-              groupProducerOffsets ++= offsets
+              groupProducerOffsets ++= offsets.map { case (groupTopicPartition, offset) =>
+                (groupTopicPartition.topicPartition, offset)
+              }
             }
         }
 
@@ -624,8 +708,8 @@ class GroupMetadataManager(brokerId: Int,
 
         // load groups which store offsets in kafka, but which have no active members and thus no group
         // metadata stored in the log
-        (emptyGroupOffsets.keySet ++ pendingEmptyGroupOffsets.keySet).foreach { case(groupId) =>
-          val group = new GroupMetadata(groupId)
+        (emptyGroupOffsets.keySet ++ pendingEmptyGroupOffsets.keySet).foreach { groupId =>
+          val group = new GroupMetadata(groupId, Empty, time)
           val offsets = emptyGroupOffsets.getOrElse(groupId, Map.empty[TopicPartition, CommitRecordMetadataAndOffset])
           val pendingOffsets = pendingEmptyGroupOffsets.getOrElse(groupId, Map.empty[Long, mutable.Map[TopicPartition, CommitRecordMetadataAndOffset]])
           debug(s"Loaded group metadata $group with offsets $offsets and pending offsets $pendingOffsets")
@@ -641,7 +725,6 @@ class GroupMetadataManager(brokerId: Int,
             throw new IllegalStateException(s"Unexpected unload of active group $groupId while " +
               s"loading partition $topicPartition")
         }
-
     }
   }
 
@@ -649,18 +732,8 @@ class GroupMetadataManager(brokerId: Int,
                         pendingTransactionalOffsets: Map[Long, mutable.Map[TopicPartition, CommitRecordMetadataAndOffset]]): Unit = {
     // offsets are initialized prior to loading the group into the cache to ensure that clients see a consistent
     // view of the group's offsets
-    val loadedOffsets = offsets.mapValues { case CommitRecordMetadataAndOffset(commitRecordOffset, offsetAndMetadata) =>
-      // special handling for version 0:
-      // set the expiration time stamp as commit time stamp + server default retention time
-      val updatedOffsetAndMetadata =
-        if (offsetAndMetadata.expireTimestamp == org.apache.kafka.common.requests.OffsetCommitRequest.DEFAULT_TIMESTAMP)
-        offsetAndMetadata.copy(expireTimestamp = offsetAndMetadata.commitTimestamp + config.offsetsRetentionMs)
-      else
-        offsetAndMetadata
-      CommitRecordMetadataAndOffset(commitRecordOffset, updatedOffsetAndMetadata)
-    }
-    trace(s"Initialized offsets $loadedOffsets for group ${group.groupId}")
-    group.initializeOffsets(loadedOffsets, pendingTransactionalOffsets.toMap)
+    trace(s"Initialized offsets $offsets for group ${group.groupId}")
+    group.initializeOffsets(offsets, pendingTransactionalOffsets.toMap)
 
     val currentGroup = addGroup(group)
     if (group != currentGroup)
@@ -675,15 +748,16 @@ 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)
+    scheduler.schedule(topicPartition.toString, () => removeGroupsAndOffsets())
 
-    def removeGroupsAndOffsets() {
+    def removeGroupsAndOffsets(): Unit = {
       var numOffsetsRemoved = 0
       var numGroupsRemoved = 0
 
+      debug(s"Started unloading offsets and group metadata for $topicPartition")
       inLock(partitionLock) {
         // we need to guard the group removal in cache in the loading partition lock
         // to prevent coordinator's check-and-get-group race condition
@@ -707,20 +781,29 @@ class GroupMetadataManager(brokerId: Int,
 
   // visible for testing
   private[group] def cleanupGroupMetadata(): Unit = {
-    cleanupGroupMetadata(None)
+    val currentTimestamp = time.milliseconds()
+    val numOffsetsRemoved = cleanupGroupMetadata(groupMetadataCache.values, group => {
+      group.removeExpiredOffsets(currentTimestamp, config.offsetsRetentionMs)
+    })
+    offsetExpiredSensor.record(numOffsetsRemoved)
+    if (numOffsetsRemoved > 0)
+      info(s"Removed $numOffsetsRemoved expired offsets in ${time.milliseconds() - currentTimestamp} milliseconds.")
   }
 
-  def cleanupGroupMetadata(deletedTopicPartitions: Option[Seq[TopicPartition]]) {
-    val startMs = time.milliseconds()
+  /**
+    * This function is used to clean up group offsets given the groups and also a function that performs the offset deletion.
+    * @param groups Groups whose metadata are to be cleaned up
+    * @param selector A function that implements deletion of (all or part of) group offsets. This function is called while
+    *                 a group lock is held, therefore there is no need for the caller to also obtain a group lock.
+    * @return The cumulative number of offsets removed
+    */
+  def cleanupGroupMetadata(groups: Iterable[GroupMetadata], selector: GroupMetadata => Map[TopicPartition, OffsetAndMetadata]): Int = {
     var offsetsRemoved = 0
 
-    groupMetadataCache.foreach { case (groupId, group) =>
+    groups.foreach { group =>
+      val groupId = group.groupId
       val (removedOffsets, groupIsDead, generation) = group.inLock {
-        val removedOffsets = deletedTopicPartitions match {
-          case Some(topicPartitions) => group.removeOffsets(topicPartitions)
-          case None => group.removeExpiredOffsets(startMs)
-        }
-
+        val removedOffsets = selector(group)
         if (group.is(Empty) && !group.hasOffsets) {
           info(s"Group $groupId transitioned to Dead in generation ${group.generationId}")
           group.transitionTo(Dead)
@@ -728,17 +811,17 @@ class GroupMetadataManager(brokerId: Int,
         (removedOffsets, group.is(Dead), group.generationId)
       }
 
-      val offsetsPartition = partitionFor(groupId)
-      val appendPartition = new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, offsetsPartition)
-      getMagic(offsetsPartition) match {
-        case Some(magicValue) =>
-          // We always use CREATE_TIME, like the producer. The conversion to LOG_APPEND_TIME (if necessary) happens automatically.
-          val timestampType = TimestampType.CREATE_TIME
-          val timestamp = time.milliseconds()
+    val offsetsPartition = partitionFor(groupId)
+    val appendPartition = new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, offsetsPartition)
+    getMagic(offsetsPartition) match {
+      case Some(magicValue) =>
+        // We always use CREATE_TIME, like the producer. The conversion to LOG_APPEND_TIME (if necessary) happens automatically.
+        val timestampType = TimestampType.CREATE_TIME
+        val timestamp = time.milliseconds()
 
           replicaManager.nonOfflinePartition(appendPartition).foreach { partition =>
-            val tombstones = ListBuffer.empty[SimpleRecord]
-            removedOffsets.foreach { case (topicPartition, offsetAndMetadata) =>
+            val tombstones = ArrayBuffer.empty[SimpleRecord]
+            removedOffsets.forKeyValue { (topicPartition, offsetAndMetadata) =>
               trace(s"Removing expired/deleted offset and metadata for $groupId, $topicPartition: $offsetAndMetadata")
               val commitKey = GroupMetadataManager.offsetCommitKey(groupId, topicPartition)
               tombstones += new SimpleRecord(timestamp, commitKey, null)
@@ -760,8 +843,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 " +
@@ -780,19 +863,32 @@ class GroupMetadataManager(brokerId: Int,
       }
     }
 
-    info(s"Removed $offsetsRemoved expired offsets in ${time.milliseconds() - startMs} milliseconds.")
+    offsetsRemoved
   }
 
-  def handleTxnCompletion(producerId: Long, completedPartitions: Set[Int], isCommit: Boolean) {
+  /**
+   * Complete pending transactional offset commits of the groups of `producerId` from the provided
+   * `completedPartitions`. This method is invoked when a commit or abort marker is fully written
+   * to the log. It may be invoked when a group lock is held by the caller, for instance when delayed
+   * operations are completed while appending offsets for a group. Since we need to acquire one or
+   * more group metadata locks to handle transaction completion, this operation is scheduled on
+   * the scheduler thread to avoid deadlocks.
+   */
+  def scheduleHandleTxnCompletion(producerId: Long, completedPartitions: Set[Int], isCommit: Boolean): Unit = {
+    scheduler.schedule(s"handleTxnCompletion-$producerId", () =>
+      handleTxnCompletion(producerId, completedPartitions, isCommit))
+  }
+
+  private[group] def handleTxnCompletion(producerId: Long, completedPartitions: Set[Int], isCommit: Boolean): Unit = {
     val pendingGroups = groupsBelongingToPartitions(producerId, completedPartitions)
-    pendingGroups.foreach { case (groupId) =>
+    pendingGroups.foreach { groupId =>
       getGroup(groupId) match {
         case Some(group) => group.inLock {
           if (!group.is(Dead)) {
             group.completePendingTxnOffsetCommit(producerId, isCommit)
             removeProducerGroup(producerId, groupId)
           }
-       }
+        }
         case _ =>
           info(s"Group $groupId has moved away from $brokerId after transaction marker was written but before the " +
             s"cache was updated. The cache on the new group owner will be updated instead.")
@@ -817,7 +913,7 @@ class GroupMetadataManager(brokerId: Int,
   }
 
   private def removeGroupFromAllProducers(groupId: String) = openGroupsForProducer synchronized {
-    openGroupsForProducer.foreach { case (_, groups) =>
+    openGroupsForProducer.forKeyValue { (_, groups) =>
       groups.remove(groupId)
     }
   }
@@ -830,7 +926,7 @@ class GroupMetadataManager(brokerId: Int,
   }
 
 
-  def shutdown() {
+  def shutdown(): Unit = {
     shuttingDown.set(true)
     if (scheduler.isStarted)
       scheduler.shutdown()
@@ -860,11 +956,24 @@ class GroupMetadataManager(brokerId: Int,
    *
    * NOTE: this is for test only
    */
-  def addPartitionOwnership(partition: Int) {
+  private[group] def addPartitionOwnership(partition: Int): Unit = {
     inLock(partitionLock) {
       ownedPartitions.add(partition)
     }
   }
+
+  /**
+   * Add a partition to the loading partitions set. Return true if the partition was not
+   * already loading.
+   *
+   * Visible for testing
+   */
+  private[group] def addLoadingPartition(partition: Int): Boolean = {
+    inLock(partitionLock) {
+      loadingPartitions.add(partition)
+    }
+  }
+
 }
 
 /**
@@ -880,183 +989,61 @@ class GroupMetadataManager(brokerId: Int,
  *    -> value version 1:       [offset, metadata, commit_timestamp, expire_timestamp]
  *
  * key version 2:       group metadata
- *     -> value version 0:       [protocol_type, generation, protocol, leader, members]
+ *    -> value version 0:       [protocol_type, generation, protocol, leader, members]
  */
 object GroupMetadataManager {
-
-  private val CURRENT_OFFSET_KEY_SCHEMA_VERSION = 1.toShort
-  private val CURRENT_GROUP_KEY_SCHEMA_VERSION = 2.toShort
-
-  private val OFFSET_COMMIT_KEY_SCHEMA = new Schema(new Field("group", STRING),
-    new Field("topic", STRING),
-    new Field("partition", INT32))
-  private val OFFSET_KEY_GROUP_FIELD = OFFSET_COMMIT_KEY_SCHEMA.get("group")
-  private val OFFSET_KEY_TOPIC_FIELD = OFFSET_COMMIT_KEY_SCHEMA.get("topic")
-  private val OFFSET_KEY_PARTITION_FIELD = OFFSET_COMMIT_KEY_SCHEMA.get("partition")
-
-  private val OFFSET_COMMIT_VALUE_SCHEMA_V0 = new Schema(new Field("offset", INT64),
-    new Field("metadata", STRING, "Associated metadata.", ""),
-    new Field("timestamp", INT64))
-  private val OFFSET_VALUE_OFFSET_FIELD_V0 = OFFSET_COMMIT_VALUE_SCHEMA_V0.get("offset")
-  private val OFFSET_VALUE_METADATA_FIELD_V0 = OFFSET_COMMIT_VALUE_SCHEMA_V0.get("metadata")
-  private val OFFSET_VALUE_TIMESTAMP_FIELD_V0 = OFFSET_COMMIT_VALUE_SCHEMA_V0.get("timestamp")
-
-  private val OFFSET_COMMIT_VALUE_SCHEMA_V1 = new Schema(new Field("offset", INT64),
-    new Field("metadata", STRING, "Associated metadata.", ""),
-    new Field("commit_timestamp", INT64),
-    new Field("expire_timestamp", INT64))
-  private val OFFSET_VALUE_OFFSET_FIELD_V1 = OFFSET_COMMIT_VALUE_SCHEMA_V1.get("offset")
-  private val OFFSET_VALUE_METADATA_FIELD_V1 = OFFSET_COMMIT_VALUE_SCHEMA_V1.get("metadata")
-  private val OFFSET_VALUE_COMMIT_TIMESTAMP_FIELD_V1 = OFFSET_COMMIT_VALUE_SCHEMA_V1.get("commit_timestamp")
-  private val OFFSET_VALUE_EXPIRE_TIMESTAMP_FIELD_V1 = OFFSET_COMMIT_VALUE_SCHEMA_V1.get("expire_timestamp")
-
-  private val GROUP_METADATA_KEY_SCHEMA = new Schema(new Field("group", STRING))
-  private val GROUP_KEY_GROUP_FIELD = GROUP_METADATA_KEY_SCHEMA.get("group")
-
-  private val MEMBER_ID_KEY = "member_id"
-  private val CLIENT_ID_KEY = "client_id"
-  private val CLIENT_HOST_KEY = "client_host"
-  private val REBALANCE_TIMEOUT_KEY = "rebalance_timeout"
-  private val SESSION_TIMEOUT_KEY = "session_timeout"
-  private val SUBSCRIPTION_KEY = "subscription"
-  private val ASSIGNMENT_KEY = "assignment"
-
-  private val MEMBER_METADATA_V0 = new Schema(
-    new Field(MEMBER_ID_KEY, STRING),
-    new Field(CLIENT_ID_KEY, STRING),
-    new Field(CLIENT_HOST_KEY, STRING),
-    new Field(SESSION_TIMEOUT_KEY, INT32),
-    new Field(SUBSCRIPTION_KEY, BYTES),
-    new Field(ASSIGNMENT_KEY, BYTES))
-
-  private val MEMBER_METADATA_V1 = new Schema(
-    new Field(MEMBER_ID_KEY, 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"
-  private val LEADER_KEY = "leader"
-  private val MEMBERS_KEY = "members"
-
-  private val GROUP_METADATA_VALUE_SCHEMA_V0 = 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(MEMBERS_KEY, new ArrayOf(MEMBER_METADATA_V0)))
-
-  private val GROUP_METADATA_VALUE_SCHEMA_V1 = 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(MEMBERS_KEY, new ArrayOf(MEMBER_METADATA_V1)))
-
-
-  // map of versions to key schemas as data types
-  private val MESSAGE_TYPE_SCHEMAS = Map(
-    0 -> OFFSET_COMMIT_KEY_SCHEMA,
-    1 -> OFFSET_COMMIT_KEY_SCHEMA,
-    2 -> GROUP_METADATA_KEY_SCHEMA)
-
-  // map of version of offset value schemas
-  private val OFFSET_VALUE_SCHEMAS = Map(
-    0 -> OFFSET_COMMIT_VALUE_SCHEMA_V0,
-    1 -> OFFSET_COMMIT_VALUE_SCHEMA_V1)
-  private val CURRENT_OFFSET_VALUE_SCHEMA_VERSION = 1.toShort
-
-  // map of version of group metadata value schemas
-  private val GROUP_VALUE_SCHEMAS = Map(
-    0 -> GROUP_METADATA_VALUE_SCHEMA_V0,
-    1 -> GROUP_METADATA_VALUE_SCHEMA_V1)
-  private val CURRENT_GROUP_VALUE_SCHEMA_VERSION = 1.toShort
-
-  private val CURRENT_OFFSET_KEY_SCHEMA = schemaForKey(CURRENT_OFFSET_KEY_SCHEMA_VERSION)
-  private val CURRENT_GROUP_KEY_SCHEMA = schemaForKey(CURRENT_GROUP_KEY_SCHEMA_VERSION)
-
-  private val CURRENT_OFFSET_VALUE_SCHEMA = schemaForOffset(CURRENT_OFFSET_VALUE_SCHEMA_VERSION)
-  private val CURRENT_GROUP_VALUE_SCHEMA = schemaForGroup(CURRENT_GROUP_VALUE_SCHEMA_VERSION)
-
-  private def schemaForKey(version: Int) = {
-    val schemaOpt = MESSAGE_TYPE_SCHEMAS.get(version)
-    schemaOpt match {
-      case Some(schema) => schema
-      case _ => throw new KafkaException("Unknown offset schema version " + version)
-    }
-  }
-
-  private def schemaForOffset(version: Int) = {
-    val schemaOpt = OFFSET_VALUE_SCHEMAS.get(version)
-    schemaOpt match {
-      case Some(schema) => schema
-      case _ => throw new KafkaException("Unknown offset schema version " + version)
-    }
-  }
-
-  private def schemaForGroup(version: Int) = {
-    val schemaOpt = GROUP_VALUE_SCHEMAS.get(version)
-    schemaOpt match {
-      case Some(schema) => schema
-      case _ => throw new KafkaException("Unknown group metadata version " + version)
-    }
-  }
+  // Metrics names
+  val MetricsGroup: String = "group-coordinator-metrics"
+  val LoadTimeSensor: String = "GroupPartitionLoadTime"
 
   /**
    * Generates the key for offset commit message for given (group, topic, partition)
    *
+   * @param groupId the ID of the group to generate the key
+   * @param topicPartition the TopicPartition to generate the key
    * @return key for offset commit message
    */
-  private[group] def offsetCommitKey(group: String, topicPartition: TopicPartition,
-                                           versionId: Short = 0): 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)
-    key.set(OFFSET_KEY_PARTITION_FIELD, topicPartition.partition)
-
-    val byteBuffer = ByteBuffer.allocate(2 /* version */ + key.sizeOf)
-    byteBuffer.putShort(CURRENT_OFFSET_KEY_SCHEMA_VERSION)
-    key.writeTo(byteBuffer)
-    byteBuffer.array()
+  def offsetCommitKey(groupId: String, topicPartition: TopicPartition): Array[Byte] = {
+    MessageUtil.toVersionPrefixedBytes(OffsetCommitKey.HIGHEST_SUPPORTED_VERSION,
+      new OffsetCommitKey()
+        .setGroup(groupId)
+        .setTopic(topicPartition.topic)
+        .setPartition(topicPartition.partition))
   }
 
   /**
    * Generates the key for group metadata message for given group
    *
+   * @param groupId the ID of the group to generate the key
    * @return key bytes for group metadata message
    */
-  private[group] def groupMetadataKey(group: String): Array[Byte] = {
-    val key = new Struct(CURRENT_GROUP_KEY_SCHEMA)
-    key.set(GROUP_KEY_GROUP_FIELD, group)
-
-    val byteBuffer = ByteBuffer.allocate(2 /* version */ + key.sizeOf)
-    byteBuffer.putShort(CURRENT_GROUP_KEY_SCHEMA_VERSION)
-    key.writeTo(byteBuffer)
-    byteBuffer.array()
+  def groupMetadataKey(groupId: String): Array[Byte] = {
+    MessageUtil.toVersionPrefixedBytes(GroupMetadataKeyData.HIGHEST_SUPPORTED_VERSION,
+      new GroupMetadataKeyData()
+        .setGroup(groupId))
   }
 
   /**
    * Generates the payload for offset commit message from given offset and metadata
    *
    * @param offsetAndMetadata consumer's current offset and metadata
+   * @param apiVersion the api version
    * @return payload for offset commit message
    */
-  private[group] def offsetCommitValue(offsetAndMetadata: OffsetAndMetadata): Array[Byte] = {
-    // generate commit value with schema version 1
-    val value = new Struct(CURRENT_OFFSET_VALUE_SCHEMA)
-    value.set(OFFSET_VALUE_OFFSET_FIELD_V1, offsetAndMetadata.offset)
-    value.set(OFFSET_VALUE_METADATA_FIELD_V1, offsetAndMetadata.metadata)
-    value.set(OFFSET_VALUE_COMMIT_TIMESTAMP_FIELD_V1, offsetAndMetadata.commitTimestamp)
-    value.set(OFFSET_VALUE_EXPIRE_TIMESTAMP_FIELD_V1, offsetAndMetadata.expireTimestamp)
-    val byteBuffer = ByteBuffer.allocate(2 /* version */ + value.sizeOf)
-    byteBuffer.putShort(CURRENT_OFFSET_VALUE_SCHEMA_VERSION)
-    value.writeTo(byteBuffer)
-    byteBuffer.array()
+  def offsetCommitValue(offsetAndMetadata: OffsetAndMetadata,
+                        apiVersion: ApiVersion): Array[Byte] = {
+    val version =
+      if (apiVersion < KAFKA_2_1_IV0 || offsetAndMetadata.expireTimestamp.nonEmpty) 1.toShort
+      else if (apiVersion < KAFKA_2_1_IV1) 2.toShort
+      else 3.toShort
+    MessageUtil.toVersionPrefixedBytes(version, new OffsetCommitValue()
+      .setOffset(offsetAndMetadata.offset)
+      .setMetadata(offsetAndMetadata.metadata)
+      .setCommitTimestamp(offsetAndMetadata.commitTimestamp)
+      .setLeaderEpoch(offsetAndMetadata.leaderEpoch.orElse(RecordBatch.NO_PARTITION_LEADER_EPOCH))
+      // version 1 has a non empty expireTimestamp field
+      .setExpireTimestamp(offsetAndMetadata.expireTimestamp.getOrElse(OffsetCommitRequest.DEFAULT_TIMESTAMP))
+    )
   }
 
   /**
@@ -1065,75 +1052,58 @@ object GroupMetadataManager {
    *
    * @param groupMetadata current group metadata
    * @param assignment the assignment for the rebalancing generation
-   * @param version the version of the value message to use
+   * @param apiVersion the api version
    * @return payload for offset commit message
    */
-  private[group] def groupMetadataValue(groupMetadata: GroupMetadata,
-                                        assignment: Map[String, Array[Byte]],
-                                        version: Short = 0): Array[Byte] = {
-    val value = if (version == 0) new Struct(GROUP_METADATA_VALUE_SCHEMA_V0) else new Struct(CURRENT_GROUP_VALUE_SCHEMA)
-
-    value.set(PROTOCOL_TYPE_KEY, groupMetadata.protocolType.getOrElse(""))
-    value.set(GENERATION_KEY, groupMetadata.generationId)
-    value.set(PROTOCOL_KEY, groupMetadata.protocol)
-    value.set(LEADER_KEY, groupMetadata.leaderId)
-
-    val memberArray = groupMetadata.allMemberMetadata.map { memberMetadata =>
-      val memberStruct = value.instance(MEMBERS_KEY)
-      memberStruct.set(MEMBER_ID_KEY, memberMetadata.memberId)
-      memberStruct.set(CLIENT_ID_KEY, memberMetadata.clientId)
-      memberStruct.set(CLIENT_HOST_KEY, memberMetadata.clientHost)
-      memberStruct.set(SESSION_TIMEOUT_KEY, memberMetadata.sessionTimeoutMs)
-
-      if (version > 0)
-        memberStruct.set(REBALANCE_TIMEOUT_KEY, memberMetadata.rebalanceTimeoutMs)
-
-      val metadata = memberMetadata.metadata(groupMetadata.protocol)
-      memberStruct.set(SUBSCRIPTION_KEY, ByteBuffer.wrap(metadata))
-
-      val memberAssignment = assignment(memberMetadata.memberId)
-      assert(memberAssignment != null)
-
-      memberStruct.set(ASSIGNMENT_KEY, ByteBuffer.wrap(memberAssignment))
-
-      memberStruct
-    }
-
-    value.set(MEMBERS_KEY, memberArray.toArray)
-
-    val byteBuffer = ByteBuffer.allocate(2 /* version */ + value.sizeOf)
-    byteBuffer.putShort(version)
-    value.writeTo(byteBuffer)
-    byteBuffer.array()
+  def groupMetadataValue(groupMetadata: GroupMetadata,
+                         assignment: Map[String, Array[Byte]],
+                         apiVersion: ApiVersion): Array[Byte] = {
+
+    val version =
+      if (apiVersion < KAFKA_0_10_1_IV0) 0.toShort
+      else if (apiVersion < KAFKA_2_1_IV0) 1.toShort
+      else if (apiVersion < KAFKA_2_3_IV0) 2.toShort
+      else 3.toShort
+
+    MessageUtil.toVersionPrefixedBytes(version, new GroupMetadataValue()
+      .setProtocolType(groupMetadata.protocolType.getOrElse(""))
+      .setGeneration(groupMetadata.generationId)
+      .setProtocol(groupMetadata.protocolName.orNull)
+      .setLeader(groupMetadata.leaderOrNull)
+      .setCurrentStateTimestamp(groupMetadata.currentStateTimestampOrDefault)
+      .setMembers(groupMetadata.allMemberMetadata.map { memberMetadata =>
+        new GroupMetadataValue.MemberMetadata()
+          .setMemberId(memberMetadata.memberId)
+          .setClientId(memberMetadata.clientId)
+          .setClientHost(memberMetadata.clientHost)
+          .setSessionTimeout(memberMetadata.sessionTimeoutMs)
+          .setRebalanceTimeout(memberMetadata.rebalanceTimeoutMs)
+          .setGroupInstanceId(memberMetadata.groupInstanceId.orNull)
+          // The group is non-empty, so the current protocol must be defined
+          .setSubscription(groupMetadata.protocolName.map(memberMetadata.metadata)
+            .getOrElse(throw new IllegalStateException("Attempted to write non-empty group metadata with no defined protocol.")))
+          .setAssignment(assignment.getOrElse(memberMetadata.memberId,
+            throw new IllegalStateException(s"Attempted to write member ${memberMetadata.memberId} of group ${groupMetadata.groupId} with no assignment.")))
+      }.asJava))
   }
 
   /**
    * Decodes the offset messages' key
    *
    * @param buffer input byte-buffer
-   * @return an GroupTopicPartition object
+   * @return an OffsetKey or GroupMetadataKey object from the message
    */
   def readMessageKey(buffer: ByteBuffer): BaseKey = {
     val version = buffer.getShort
-    val keySchema = schemaForKey(version)
-    val key = keySchema.read(buffer)
-
-    if (version <= CURRENT_OFFSET_KEY_SCHEMA_VERSION) {
+    if (version >= OffsetCommitKey.LOWEST_SUPPORTED_VERSION && version <= OffsetCommitKey.HIGHEST_SUPPORTED_VERSION) {
       // version 0 and 1 refer to offset
-      val group = key.get(OFFSET_KEY_GROUP_FIELD).asInstanceOf[String]
-      val topic = key.get(OFFSET_KEY_TOPIC_FIELD).asInstanceOf[String]
-      val partition = key.get(OFFSET_KEY_PARTITION_FIELD).asInstanceOf[Int]
-
-      OffsetKey(version, GroupTopicPartition(group, new TopicPartition(topic, partition)))
-
-    } else if (version == CURRENT_GROUP_KEY_SCHEMA_VERSION) {
-      // version 2 refers to offset
-      val group = key.get(GROUP_KEY_GROUP_FIELD).asInstanceOf[String]
-
-      GroupMetadataKey(version, group)
-    } else {
-      throw new IllegalStateException("Unknown version " + version + " for group metadata message")
-    }
+      val key = new OffsetCommitKey(new ByteBufferAccessor(buffer), version)
+      OffsetKey(version, GroupTopicPartition(key.group, new TopicPartition(key.topic, key.partition)))
+    } else if (version >= GroupMetadataKeyData.LOWEST_SUPPORTED_VERSION && version <= GroupMetadataKeyData.HIGHEST_SUPPORTED_VERSION) {
+      // version 2 refers to group metadata
+      val key = new GroupMetadataKeyData(new ByteBufferAccessor(buffer), version)
+      GroupMetadataKey(version, key.group)
+    } else throw new IllegalStateException(s"Unknown group metadata message version: $version")
   }
 
   /**
@@ -1143,87 +1113,68 @@ object GroupMetadataManager {
    * @return an offset-metadata object from the message
    */
   def readOffsetMessageValue(buffer: ByteBuffer): OffsetAndMetadata = {
-    if (buffer == null) { // tombstone
-      null
-    } else {
+    // tombstone
+    if (buffer == null) null
+    else {
       val version = buffer.getShort
-      val valueSchema = schemaForOffset(version)
-      val value = valueSchema.read(buffer)
-
-      if (version == 0) {
-        val offset = value.get(OFFSET_VALUE_OFFSET_FIELD_V0).asInstanceOf[Long]
-        val metadata = value.get(OFFSET_VALUE_METADATA_FIELD_V0).asInstanceOf[String]
-        val timestamp = value.get(OFFSET_VALUE_TIMESTAMP_FIELD_V0).asInstanceOf[Long]
-
-        OffsetAndMetadata(offset, metadata, timestamp)
-      } else if (version == 1) {
-        val offset = value.get(OFFSET_VALUE_OFFSET_FIELD_V1).asInstanceOf[Long]
-        val metadata = value.get(OFFSET_VALUE_METADATA_FIELD_V1).asInstanceOf[String]
-        val commitTimestamp = value.get(OFFSET_VALUE_COMMIT_TIMESTAMP_FIELD_V1).asInstanceOf[Long]
-        val expireTimestamp = value.get(OFFSET_VALUE_EXPIRE_TIMESTAMP_FIELD_V1).asInstanceOf[Long]
-
-        OffsetAndMetadata(offset, metadata, commitTimestamp, expireTimestamp)
-      } else {
-        throw new IllegalStateException("Unknown offset message version")
-      }
+      if (version >= OffsetCommitValue.LOWEST_SUPPORTED_VERSION && version <= OffsetCommitValue.HIGHEST_SUPPORTED_VERSION) {
+        val value = new OffsetCommitValue(new ByteBufferAccessor(buffer), version)
+        OffsetAndMetadata(
+          offset = value.offset,
+          leaderEpoch = if (value.leaderEpoch == RecordBatch.NO_PARTITION_LEADER_EPOCH) Optional.empty() else Optional.of(value.leaderEpoch),
+          metadata = value.metadata,
+          commitTimestamp = value.commitTimestamp,
+          expireTimestamp = if (value.expireTimestamp == OffsetCommitRequest.DEFAULT_TIMESTAMP) None else Some(value.expireTimestamp))
+      } else throw new IllegalStateException(s"Unknown offset message version: $version")
     }
   }
 
   /**
-   * Decodes the group metadata messages' payload and retrieves its member metadatafrom it
+   * Decodes the group metadata messages' payload and retrieves its member metadata from it
    *
+   * @param groupId The ID of the group to be read
    * @param buffer input byte-buffer
+   * @param time the time instance to use
    * @return a group metadata object from the message
    */
-  def readGroupMessageValue(groupId: String, buffer: ByteBuffer): GroupMetadata = {
-    if (buffer == null) { // tombstone
-      null
-    } else {
+  def readGroupMessageValue(groupId: String, buffer: ByteBuffer, time: Time): GroupMetadata = {
+    // tombstone
+    if (buffer == null) null
+    else {
       val version = buffer.getShort
-      val valueSchema = schemaForGroup(version)
-      val value = valueSchema.read(buffer)
-
-      if (version == 0 || version == 1) {
-        val protocolType = value.get(PROTOCOL_TYPE_KEY).asInstanceOf[String]
-
-        val memberMetadataArray = value.getArray(MEMBERS_KEY)
-        val initialState = if (memberMetadataArray.isEmpty) Empty else Stable
-
-        val group = new GroupMetadata(groupId, initialState)
-
-        group.generationId = value.get(GENERATION_KEY).asInstanceOf[Int]
-        group.leaderId = value.get(LEADER_KEY).asInstanceOf[String]
-        group.protocol = value.get(PROTOCOL_KEY).asInstanceOf[String]
-
-        memberMetadataArray.foreach { memberMetadataObj =>
-          val memberMetadata = memberMetadataObj.asInstanceOf[Struct]
-          val memberId = memberMetadata.get(MEMBER_ID_KEY).asInstanceOf[String]
-          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,
-            protocolType, List((group.protocol, subscription)))
-
-          member.assignment = Utils.toArray(memberMetadata.get(ASSIGNMENT_KEY).asInstanceOf[ByteBuffer])
-
-          group.add(member)
+      if (version >= GroupMetadataValue.LOWEST_SUPPORTED_VERSION && version <= GroupMetadataValue.HIGHEST_SUPPORTED_VERSION) {
+        val value = new GroupMetadataValue(new ByteBufferAccessor(buffer), version)
+        val members = value.members.asScala.map { memberMetadata =>
+          new MemberMetadata(
+            memberId = memberMetadata.memberId,
+            groupId = groupId,
+            groupInstanceId = Option(memberMetadata.groupInstanceId),
+            clientId = memberMetadata.clientId,
+            clientHost = memberMetadata.clientHost,
+            rebalanceTimeoutMs = if (version == 0) memberMetadata.sessionTimeout else memberMetadata.rebalanceTimeout,
+            sessionTimeoutMs = memberMetadata.sessionTimeout,
+            protocolType = value.protocolType,
+            supportedProtocols = List((value.protocol, memberMetadata.subscription)),
+            assignment = memberMetadata.assignment)
         }
-
-        group
-      } else {
-        throw new IllegalStateException("Unknown group metadata message version")
-      }
+        GroupMetadata.loadGroup(
+          groupId = groupId,
+          initialState = if (members.isEmpty) Empty else Stable,
+          generationId = value.generation,
+          protocolType = value.protocolType,
+          protocolName = value.protocol,
+          leaderId = value.leader,
+          currentStateTimestamp = if (value.currentStateTimestamp == -1) None else Some(value.currentStateTimestamp),
+          members = members,
+          time = time)
+      } else throw new IllegalStateException(s"Unknown group metadata message version: $version")
     }
   }
 
   // Formatter for use with tools such as console consumer: Consumer should also set exclude.internal.topics to false.
-  // (specify --formatter "kafka.coordinator.GroupMetadataManager\$OffsetsMessageFormatter" when consuming __consumer_offsets)
+  // (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.
@@ -1244,7 +1195,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.
@@ -1253,7 +1204,7 @@ object GroupMetadataManager {
           val value = consumerRecord.value
           val formattedValue =
             if (value == null) "NULL"
-            else GroupMetadataManager.readGroupMessageValue(groupId, ByteBuffer.wrap(value)).toString
+            else GroupMetadataManager.readGroupMessageValue(groupId, ByteBuffer.wrap(value), Time.SYSTEM).toString
           output.write(groupId.getBytes(StandardCharsets.UTF_8))
           output.write("::".getBytes(StandardCharsets.UTF_8))
           output.write(formattedValue.getBytes(StandardCharsets.UTF_8))
@@ -1263,6 +1214,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.protocolName.orNull,
+        "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 b082b9bad2efd..e6c98a5b83663 100644
--- a/core/src/main/scala/kafka/coordinator/group/MemberMetadata.scala
+++ b/core/src/main/scala/kafka/coordinator/group/MemberMetadata.scala
@@ -20,15 +20,18 @@ 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],
                          assignment: Array[Byte])
 
+private object MemberMetadata {
+  def plainProtocolSet(supportedProtocols: List[(String, Array[Byte])]) = supportedProtocols.map(_._1).toSet
+}
+
 /**
  * Member metadata contains the following metadata:
  *
@@ -50,22 +53,32 @@ case class MemberSummary(memberId: String,
  *                            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,
                                     val sessionTimeoutMs: Int,
                                     val protocolType: String,
-                                    var supportedProtocols: List[(String, Array[Byte])]) {
+                                    var supportedProtocols: List[(String, Array[Byte])],
+                                    var assignment: Array[Byte] = Array.empty[Byte]) {
 
-  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 protocols = supportedProtocols.map(_._1).toSet
+  def isAwaitingJoin = awaitingJoinCallback != null
+  def isAwaitingSync = awaitingSyncCallback != null
 
   /**
    * Get metadata corresponding to the provided protocol.
@@ -78,6 +91,19 @@ private[group] class MemberMetadata(val memberId: String,
     }
   }
 
+  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
+    }
+  }
+
   /**
    * Check if the provided protocol metadata matches the currently stored metadata.
    */
@@ -95,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])
   }
 
   /**
@@ -117,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, " +
@@ -124,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 2b9a60fdf3e6e..55ec590852cd0 100644
--- a/core/src/main/scala/kafka/coordinator/group/OffsetConfig.scala
+++ b/core/src/main/scala/kafka/coordinator/group/OffsetConfig.scala
@@ -23,7 +23,8 @@ import kafka.message.{CompressionCodec, NoCompressionCodec}
  * Configuration settings for in-built offset management
  * @param maxMetadataSize The maximum allowed metadata for any offset commit.
  * @param loadBufferSize Batch size for reading from the offsets segments when loading offsets into the cache.
- * @param offsetsRetentionMs Offsets older than this retention period will be discarded.
+ * @param offsetsRetentionMs After a consumer group loses all its consumers (i.e. becomes empty) its offsets will be kept for this retention period before getting discarded.
+ *                           For standalone consumers (using manual assignment), offsets will be expired after the time of last commit plus this retention period.
  * @param offsetsRetentionCheckIntervalMs Frequency at which to check for expired offsets.
  * @param offsetsTopicNumPartitions The number of partitions for the offset commit topic (should not change after deployment).
  * @param offsetsTopicSegmentBytes The offsets topic segment bytes should be kept relatively small to facilitate faster
diff --git a/core/src/main/scala/kafka/coordinator/transaction/DelayedTxnMarker.scala b/core/src/main/scala/kafka/coordinator/transaction/DelayedTxnMarker.scala
deleted file mode 100644
index cf18b8196053a..0000000000000
--- a/core/src/main/scala/kafka/coordinator/transaction/DelayedTxnMarker.scala
+++ /dev/null
@@ -1,49 +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.coordinator.transaction
-
-import java.util.concurrent.TimeUnit
-import java.util.concurrent.locks.Lock
-
-import kafka.server.DelayedOperation
-import org.apache.kafka.common.protocol.Errors
-
-/**
-  * Delayed transaction state change operations that are added to the purgatory without timeout (i.e. these operations should never time out)
-  */
-private[transaction] class DelayedTxnMarker(txnMetadata: TransactionMetadata,
-                                           completionCallback: Errors => Unit,
-                                           lock: Lock)
-  extends DelayedOperation(TimeUnit.DAYS.toMillis(100 * 365), Some(lock)) {
-
-  override def tryComplete(): Boolean = {
-    txnMetadata.inLock {
-      if (txnMetadata.topicPartitions.isEmpty)
-        forceComplete()
-      else false
-    }
-  }
-
-  override def onExpiration(): Unit = {
-    // this should never happen
-    throw new IllegalStateException(s"Delayed write txn marker operation for metadata $txnMetadata has timed out, this should never happen.")
-  }
-
-  // TODO: if we will always return NONE upon completion, we can remove the error code in the param
-  override def onComplete(): Unit = completionCallback(Errors.NONE)
-
-}
diff --git a/core/src/main/scala/kafka/coordinator/transaction/ProducerIdManager.scala b/core/src/main/scala/kafka/coordinator/transaction/ProducerIdManager.scala
index 6be3c6babff2e..4ea86bbae33dc 100644
--- a/core/src/main/scala/kafka/coordinator/transaction/ProducerIdManager.scala
+++ b/core/src/main/scala/kafka/coordinator/transaction/ProducerIdManager.scala
@@ -16,9 +16,13 @@
  */
 package kafka.coordinator.transaction
 
-import kafka.common.KafkaException
-import kafka.utils.{Json, Logging, ZkUtils}
-import kafka.zk.KafkaZkClient
+import java.nio.charset.StandardCharsets
+
+import kafka.utils.{Json, Logging}
+import kafka.zk.{KafkaZkClient, ProducerIdBlockZNode}
+import org.apache.kafka.common.KafkaException
+
+import scala.jdk.CollectionConverters._
 
 /**
  * ProducerIdManager is the part of the transaction coordinator that provides ProducerIds in a unique way
@@ -31,17 +35,17 @@ object ProducerIdManager extends Logging {
   val CurrentVersion: Long = 1L
   val PidBlockSize: Long = 1000L
 
-  def generateProducerIdBlockJson(producerIdBlock: ProducerIdBlock): String = {
-    Json.encode(Map("version" -> CurrentVersion,
+  def generateProducerIdBlockJson(producerIdBlock: ProducerIdBlock): Array[Byte] = {
+    Json.encodeAsBytes(Map("version" -> CurrentVersion,
       "broker" -> producerIdBlock.brokerId,
       "block_start" -> producerIdBlock.blockStartId.toString,
-      "block_end" -> producerIdBlock.blockEndId.toString)
+      "block_end" -> producerIdBlock.blockEndId.toString).asJava
     )
   }
 
-  def parseProducerIdBlockData(jsonData: String): ProducerIdBlock = {
+  def parseProducerIdBlockData(jsonData: Array[Byte]): ProducerIdBlock = {
     try {
-      Json.parseFull(jsonData).map(_.asJsonObject).flatMap { js =>
+      Json.parseBytes(jsonData).map(_.asJsonObject).flatMap { js =>
         val brokerId = js("broker").to[Int]
         val blockStart = js("block_start").to[String].toLong
         val blockEnd = js("block_end").to[String].toLong
@@ -83,7 +87,7 @@ class ProducerIdManager(val brokerId: Int, val zkClient: KafkaZkClient) extends
     var zkWriteComplete = false
     while (!zkWriteComplete) {
       // refresh current producerId block from zookeeper again
-      val (dataOpt, zkVersion) = zkClient.getDataAndVersion(ZkUtils.ProducerIdBlockPath)
+      val (dataOpt, zkVersion) = zkClient.getDataAndVersion(ProducerIdBlockZNode.path)
 
       // generate the new producerId block
       currentProducerIdBlock = dataOpt match {
@@ -106,7 +110,7 @@ class ProducerIdManager(val brokerId: Int, val zkClient: KafkaZkClient) extends
       val newProducerIdBlockData = ProducerIdManager.generateProducerIdBlockJson(currentProducerIdBlock)
 
       // try to write the new producerId block into zookeeper
-      val (succeeded, version) = zkClient.conditionalUpdatePath(ZkUtils.ProducerIdBlockPath,
+      val (succeeded, version) = zkClient.conditionalUpdatePath(ProducerIdBlockZNode.path,
         newProducerIdBlockData, zkVersion, Some(checkProducerIdBlockZkData))
       zkWriteComplete = succeeded
 
@@ -115,21 +119,19 @@ class ProducerIdManager(val brokerId: Int, val zkClient: KafkaZkClient) extends
     }
   }
 
-  private def checkProducerIdBlockZkData(zkClient: KafkaZkClient, path: String, expectedData: String): (Boolean, Int) = {
+  private def checkProducerIdBlockZkData(zkClient: KafkaZkClient, path: String, expectedData: Array[Byte]): (Boolean, Int) = {
     try {
       val expectedPidBlock = ProducerIdManager.parseProducerIdBlockData(expectedData)
-      val (dataOpt, zkVersion) = zkClient.getDataAndVersion(ZkUtils.ProducerIdBlockPath)
-      dataOpt match {
-        case Some(data) =>
+      zkClient.getDataAndVersion(ProducerIdBlockZNode.path) match {
+        case (Some(data), zkVersion) =>
           val currProducerIdBLock = ProducerIdManager.parseProducerIdBlockData(data)
           (currProducerIdBLock == expectedPidBlock, zkVersion)
-        case None =>
-          (false, -1)
+        case (None, _) => (false, -1)
       }
     } catch {
       case e: Exception =>
-        warn(s"Error while checking for producerId block Zk data on path $path: expected data $expectedData", e)
-
+        warn(s"Error while checking for producerId block Zk data on path $path: expected data " +
+          s"${new String(expectedData, StandardCharsets.UTF_8)}", e)
         (false, -1)
     }
   }
@@ -148,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 2cc719d2f589d..147ff73629b8c 100644
--- a/core/src/main/scala/kafka/coordinator/transaction/TransactionCoordinator.scala
+++ b/core/src/main/scala/kafka/coordinator/transaction/TransactionCoordinator.scala
@@ -19,7 +19,7 @@ package kafka.coordinator.transaction
 import java.util.Properties
 import java.util.concurrent.atomic.AtomicBoolean
 
-import kafka.server.{DelayedOperationPurgatory, KafkaConfig, MetadataCache, ReplicaManager}
+import kafka.server.{KafkaConfig, MetadataCache, ReplicaManager}
 import kafka.utils.{Logging, Scheduler}
 import kafka.zk.KafkaZkClient
 import org.apache.kafka.common.TopicPartition
@@ -28,7 +28,7 @@ import org.apache.kafka.common.metrics.Metrics
 import org.apache.kafka.common.protocol.Errors
 import org.apache.kafka.common.record.RecordBatch
 import org.apache.kafka.common.requests.TransactionResult
-import org.apache.kafka.common.utils.{LogContext, Time}
+import org.apache.kafka.common.utils.{LogContext, ProducerIdAndEpoch, Time}
 
 object TransactionCoordinator {
 
@@ -52,14 +52,12 @@ object TransactionCoordinator {
       config.requestTimeoutMs)
 
     val producerIdManager = new ProducerIdManager(config.brokerId, zkClient)
-    // 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,
-      txnMarkerPurgatory, time, logContext)
+      time, logContext)
 
     new TransactionCoordinator(config.brokerId, txnConfig, scheduler, producerIdManager, txnStateManager, txnMarkerChannelManager,
       time, logContext)
@@ -104,6 +102,7 @@ class TransactionCoordinator(brokerId: Int,
 
   def handleInitProducerId(transactionalId: String,
                            transactionTimeoutMs: Int,
+                           expectedProducerIdAndEpoch: Option[ProducerIdAndEpoch],
                            responseCallback: InitProducerIdCallback): Unit = {
 
     if (transactionalId == null) {
@@ -119,28 +118,31 @@ class TransactionCoordinator(brokerId: Int,
       // check transactionTimeoutMs is not larger than the broker configured maximum allowed value
       responseCallback(initTransactionError(Errors.INVALID_TRANSACTION_TIMEOUT))
     } else {
-      val coordinatorEpochAndMetadata = txnManager.getTransactionState(transactionalId).right.flatMap {
+      val coordinatorEpochAndMetadata = txnManager.getTransactionState(transactionalId).flatMap {
         case None =>
           val producerId = producerIdManager.generateProducerId()
           val createdMetadata = new TransactionMetadata(transactionalId = transactionalId,
             producerId = producerId,
+            lastProducerId = RecordBatch.NO_PRODUCER_ID,
             producerEpoch = RecordBatch.NO_PRODUCER_EPOCH,
+            lastProducerEpoch = RecordBatch.NO_PRODUCER_EPOCH,
             txnTimeoutMs = transactionTimeoutMs,
             state = Empty,
             topicPartitions = collection.mutable.Set.empty[TopicPartition],
             txnLastUpdateTimestamp = time.milliseconds())
-          txnManager.putTransactionStateIfNotExists(transactionalId, createdMetadata)
+          txnManager.putTransactionStateIfNotExists(createdMetadata)
 
         case Some(epochAndTxnMetadata) => Right(epochAndTxnMetadata)
       }
 
-      val result: ApiResult[(Int, TxnTransitMetadata)] = coordinatorEpochAndMetadata.right.flatMap {
+      val result: ApiResult[(Int, TxnTransitMetadata)] = coordinatorEpochAndMetadata.flatMap {
         existingEpochAndMetadata =>
           val coordinatorEpoch = existingEpochAndMetadata.coordinatorEpoch
           val txnMetadata = existingEpochAndMetadata.transactionMetadata
 
           txnMetadata.inLock {
-            prepareInitProduceIdTransit(transactionalId, transactionTimeoutMs, coordinatorEpoch, txnMetadata)
+            prepareInitProducerIdTransit(transactionalId, transactionTimeoutMs, coordinatorEpoch, txnMetadata,
+              expectedProducerIdAndEpoch)
           }
       }
 
@@ -159,10 +161,11 @@ class TransactionCoordinator(brokerId: Int,
               }
             }
 
-            handleEndTransaction(transactionalId,
+            endTransaction(transactionalId,
               newMetadata.producerId,
               newMetadata.producerEpoch,
               TransactionResult.ABORT,
+              isFromClient = false,
               sendRetriableErrorCallback)
           } else {
             def sendPidResponseCallback(error: Errors): Unit = {
@@ -183,13 +186,32 @@ class TransactionCoordinator(brokerId: Int,
     }
   }
 
-  private def prepareInitProduceIdTransit(transactionalId: String,
+  private def prepareInitProducerIdTransit(transactionalId: String,
                                           transactionTimeoutMs: Int,
                                           coordinatorEpoch: Int,
-                                          txnMetadata: TransactionMetadata): ApiResult[(Int, TxnTransitMetadata)] = {
+                                          txnMetadata: TransactionMetadata,
+                                          expectedProducerIdAndEpoch: Option[ProducerIdAndEpoch]): ApiResult[(Int, TxnTransitMetadata)] = {
+
+    def isValidProducerId(producerIdAndEpoch: ProducerIdAndEpoch): Boolean = {
+      // If a producer ID and epoch are provided by the request, fence the producer unless one of the following is true:
+      //   1. The producer epoch is equal to -1, which implies that the metadata was just created. This is the case of a
+      //      producer recovering from an UNKNOWN_PRODUCER_ID error, and it is safe to return the newly-generated
+      //      producer ID.
+      //   2. The expected producer ID matches the ID in current metadata (the epoch will be checked when we try to
+      //      increment it)
+      //   3. The expected producer ID matches the previous one and the expected epoch is exhausted, in which case this
+      //      could be a retry after a valid epoch bump that the producer never received the response for
+      txnMetadata.producerEpoch == RecordBatch.NO_PRODUCER_EPOCH ||
+        producerIdAndEpoch.producerId == txnMetadata.producerId ||
+        (producerIdAndEpoch.producerId == txnMetadata.lastProducerId && TransactionMetadata.isEpochExhausted(producerIdAndEpoch.epoch))
+    }
+
     if (txnMetadata.pendingTransitionInProgress) {
       // return a retriable exception to let the client backoff and retry
       Left(Errors.CONCURRENT_TRANSACTIONS)
+    }
+    else if (!expectedProducerIdAndEpoch.forall(isValidProducerId)) {
+      Left(Errors.PRODUCER_FENCED)
     } else {
       // caller should have synchronized on txnMetadata already
       txnMetadata.state match {
@@ -198,14 +220,22 @@ class TransactionCoordinator(brokerId: Int,
           Left(Errors.CONCURRENT_TRANSACTIONS)
 
         case CompleteAbort | CompleteCommit | Empty =>
-          val transitMetadata = if (txnMetadata.isProducerEpochExhausted) {
-            val newProducerId = producerIdManager.generateProducerId()
-            txnMetadata.prepareProducerIdRotation(newProducerId, transactionTimeoutMs, time.milliseconds())
-          } else {
-            txnMetadata.prepareIncrementProducerEpoch(transactionTimeoutMs, time.milliseconds())
-          }
+          val transitMetadataResult =
+            // If the epoch is exhausted and the expected epoch (if provided) matches it, generate a new producer ID
+            if (txnMetadata.isProducerEpochExhausted &&
+                expectedProducerIdAndEpoch.forall(_.epoch == txnMetadata.producerEpoch)) {
+              val newProducerId = producerIdManager.generateProducerId()
+              Right(txnMetadata.prepareProducerIdRotation(newProducerId, transactionTimeoutMs, time.milliseconds(),
+                expectedProducerIdAndEpoch.isDefined))
+            } else {
+              txnMetadata.prepareIncrementProducerEpoch(transactionTimeoutMs, expectedProducerIdAndEpoch.map(_.epoch),
+                time.milliseconds())
+            }
 
-          Right(coordinatorEpoch, transitMetadata)
+          transitMetadataResult match {
+            case Right(transitMetadata) => Right((coordinatorEpoch, transitMetadata))
+            case Left(err) => Left(err)
+          }
 
         case Ongoing =>
           // indicate to abort the current ongoing txn first. Note that this epoch is never returned to the
@@ -236,7 +266,7 @@ class TransactionCoordinator(brokerId: Int,
     } else {
       // try to update the transaction metadata and append the updated metadata to txn log;
       // if there is no such metadata treat it as invalid producerId mapping error.
-      val result: ApiResult[(Int, TxnTransitMetadata)] = txnManager.getTransactionState(transactionalId).right.flatMap {
+      val result: ApiResult[(Int, TxnTransitMetadata)] = txnManager.getTransactionState(transactionalId).flatMap {
         case None => Left(Errors.INVALID_PRODUCER_ID_MAPPING)
 
         case Some(epochAndMetadata) =>
@@ -248,7 +278,7 @@ class TransactionCoordinator(brokerId: Int,
             if (txnMetadata.producerId != producerId) {
               Left(Errors.INVALID_PRODUCER_ID_MAPPING)
             } else if (txnMetadata.producerEpoch != producerEpoch) {
-              Left(Errors.INVALID_PRODUCER_EPOCH)
+              Left(Errors.PRODUCER_FENCED)
             } else if (txnMetadata.pendingTransitionInProgress) {
               // return a retriable exception to let the client backoff and retry
               Left(Errors.CONCURRENT_TRANSACTIONS)
@@ -274,13 +304,40 @@ class TransactionCoordinator(brokerId: Int,
     }
   }
 
-  def handleTxnImmigration(txnTopicPartitionId: Int, coordinatorEpoch: Int) {
-    txnManager.loadTransactionsForTxnTopicPartition(txnTopicPartitionId, coordinatorEpoch, txnMarkerChannelManager.addTxnMarkersToSend)
+  /**
+   * 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 = {
+    info(s"Elected as the txn coordinator for partition $txnTopicPartitionId at epoch $coordinatorEpoch")
+    // 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 = {
+    info(s"Resigned as the txn coordinator for partition $txnTopicPartitionId at epoch $coordinatorEpoch")
+    coordinatorEpoch match {
+      case Some(epoch) =>
+        txnManager.removeTransactionsForTxnTopicPartition(txnTopicPartitionId, epoch)
+      case None =>
+        txnManager.removeTransactionsForTxnTopicPartition(txnTopicPartitionId)
+    }
+    txnMarkerChannelManager.removeMarkersForTxnTopicPartition(txnTopicPartitionId)
   }
 
   private def logInvalidStateTransitionAndReturnError(transactionalId: String,
@@ -296,10 +353,25 @@ class TransactionCoordinator(brokerId: Int,
                            producerEpoch: Short,
                            txnMarkerResult: TransactionResult,
                            responseCallback: EndTxnCallback): Unit = {
+    endTransaction(transactionalId,
+      producerId,
+      producerEpoch,
+      txnMarkerResult,
+      isFromClient = true,
+      responseCallback)
+  }
+
+  private def endTransaction(transactionalId: String,
+                             producerId: Long,
+                             producerEpoch: Short,
+                             txnMarkerResult: TransactionResult,
+                             isFromClient: Boolean,
+                             responseCallback: EndTxnCallback): Unit = {
+    var isEpochFence = false
     if (transactionalId == null || transactionalId.isEmpty)
       responseCallback(Errors.INVALID_REQUEST)
     else {
-      val preAppendResult: ApiResult[(Int, TxnTransitMetadata)] = txnManager.getTransactionState(transactionalId).right.flatMap {
+      val preAppendResult: ApiResult[(Int, TxnTransitMetadata)] = txnManager.getTransactionState(transactionalId).flatMap {
         case None =>
           Left(Errors.INVALID_PRODUCER_ID_MAPPING)
 
@@ -310,8 +382,9 @@ class TransactionCoordinator(brokerId: Int,
           txnMetadata.inLock {
             if (txnMetadata.producerId != producerId)
               Left(Errors.INVALID_PRODUCER_ID_MAPPING)
-            else if (producerEpoch < txnMetadata.producerEpoch)
-              Left(Errors.INVALID_PRODUCER_EPOCH)
+            // Strict equality is enforced on the client side requests, as they shouldn't bump the producer epoch.
+            else if ((isFromClient && producerEpoch != txnMetadata.producerEpoch) || producerEpoch < txnMetadata.producerEpoch)
+              Left(Errors.PRODUCER_FENCED)
             else if (txnMetadata.pendingTransitionInProgress && txnMetadata.pendingState.get != PrepareEpochFence)
               Left(Errors.CONCURRENT_TRANSACTIONS)
             else txnMetadata.state match {
@@ -324,8 +397,10 @@ class TransactionCoordinator(brokerId: Int,
                 if (nextState == PrepareAbort && txnMetadata.pendingState.contains(PrepareEpochFence)) {
                   // We should clear the pending state to make way for the transition to PrepareAbort and also bump
                   // the epoch in the transaction metadata we are about to append.
+                  isEpochFence = true
                   txnMetadata.pendingState = None
                   txnMetadata.producerEpoch = producerEpoch
+                  txnMetadata.lastProducerEpoch = RecordBatch.NO_PRODUCER_EPOCH
                 }
 
                 Right(coordinatorEpoch, txnMetadata.prepareAbortOrCommit(nextState, time.milliseconds()))
@@ -369,7 +444,7 @@ class TransactionCoordinator(brokerId: Int,
         case Right((coordinatorEpoch, newMetadata)) =>
           def sendTxnMarkersCallback(error: Errors): Unit = {
             if (error == Errors.NONE) {
-              val preSendResult: ApiResult[(TransactionMetadata, TxnTransitMetadata)] = txnManager.getTransactionState(transactionalId).right.flatMap {
+              val preSendResult: ApiResult[(TransactionMetadata, TxnTransitMetadata)] = txnManager.getTransactionState(transactionalId).flatMap {
                 case None =>
                   val errorMsg = s"The coordinator still owns the transaction partition for $transactionalId, but there is " +
                     s"no metadata in the cache; this is not expected"
@@ -383,7 +458,7 @@ class TransactionCoordinator(brokerId: Int,
                       if (txnMetadata.producerId != producerId)
                         Left(Errors.INVALID_PRODUCER_ID_MAPPING)
                       else if (txnMetadata.producerEpoch != producerEpoch)
-                        Left(Errors.INVALID_PRODUCER_EPOCH)
+                        Left(Errors.PRODUCER_FENCED)
                       else if (txnMetadata.pendingTransitionInProgress)
                         Left(Errors.CONCURRENT_TRANSACTIONS)
                       else txnMetadata.state match {
@@ -424,12 +499,28 @@ class TransactionCoordinator(brokerId: Int,
                   // the log append was successful
                   responseCallback(Errors.NONE)
 
-                  txnMarkerChannelManager.addTxnMarkersToSend(transactionalId, coordinatorEpoch, txnMarkerResult, txnMetadata, newPreSendMetadata)
+                  txnMarkerChannelManager.addTxnMarkersToSend(coordinatorEpoch, txnMarkerResult, txnMetadata, newPreSendMetadata)
               }
             } else {
               info(s"Aborting sending of transaction markers and returning $error error to client for $transactionalId's EndTransaction request of $txnMarkerResult, " +
                 s"since appending $newMetadata to transaction log with coordinator epoch $coordinatorEpoch failed")
 
+              if (isEpochFence) {
+                txnManager.getTransactionState(transactionalId).foreach {
+                  case None =>
+                    warn(s"The coordinator still owns the transaction partition for $transactionalId, but there is " +
+                      s"no metadata in the cache; this is not expected")
+
+                  case Some(epochAndMetadata) =>
+                    if (epochAndMetadata.coordinatorEpoch == coordinatorEpoch) {
+                      // This was attempted epoch fence that failed, so mark this state on the metadata
+                      epochAndMetadata.transactionMetadata.hasFailedEpochFence = true
+                      warn(s"The coordinator failed to write an epoch fence transition for producer $transactionalId to the transaction log " +
+                        s"with error $error. The epoch was increased to ${newMetadata.producerEpoch} but not returned to the client")
+                    }
+                }
+              }
+
               responseCallback(error)
             }
           }
@@ -443,48 +534,55 @@ class TransactionCoordinator(brokerId: Int,
 
   def partitionFor(transactionalId: String): Int = txnManager.partitionFor(transactionalId)
 
-  private def abortTimedOutTransactions(): Unit = {
+  private def onEndTransactionComplete(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.PRODUCER_FENCED |
+                  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")
+    }
+  }
+
+  private[transaction] def abortTimedOutTransactions(onComplete: TransactionalIdAndProducerIdEpoch => EndTxnCallback): Unit = {
+
     txnManager.timedOutTransactions().foreach { txnIdAndPidEpoch =>
-      txnManager.getTransactionState(txnIdAndPidEpoch.transactionalId).right.flatMap {
+      txnManager.getTransactionState(txnIdAndPidEpoch.transactionalId).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 =>
+            endTransaction(txnMetadata.transactionalId,
+              txnTransitMetadata.producerId,
+              txnTransitMetadata.producerEpoch,
+              TransactionResult.ABORT,
+              isFromClient = false,
+              onComplete(txnIdAndPidEpoch))
           }
       }
     }
@@ -493,11 +591,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(onEndTransactionComplete),
       txnConfig.abortTimedOutTransactionsIntervalMs,
       txnConfig.abortTimedOutTransactionsIntervalMs
     )
@@ -513,7 +611,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 2c7178e885ecf..68a2bedfc957d 100644
--- a/core/src/main/scala/kafka/coordinator/transaction/TransactionLog.scala
+++ b/core/src/main/scala/kafka/coordinator/transaction/TransactionLog.scala
@@ -16,26 +16,26 @@
  */
 package kafka.coordinator.transaction
 
-import kafka.common.{KafkaException, MessageFormatter}
-import org.apache.kafka.clients.consumer.ConsumerRecord
-import org.apache.kafka.common.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.internals.generated.{TransactionLogKey, TransactionLogValue}
+import org.apache.kafka.clients.consumer.ConsumerRecord
+import org.apache.kafka.common.protocol.{ByteBufferAccessor, MessageUtil}
+import org.apache.kafka.common.record.{CompressionType, Record, RecordBatch}
+import org.apache.kafka.common.{MessageFormatter, TopicPartition}
 
 import scala.collection.mutable
+import scala.jdk.CollectionConverters._
 
-/*
+/**
  * Messages stored for the transaction topic represent the producer id and transactional status of the corresponding
  * transactional id, which have versions for both the key and value fields. Key and value
  * versions are used to evolve the message formats:
  *
  * key version 0:               [transactionalId]
- *    -> value version 0:       [producer_id, producer_epoch, expire_timestamp, status, [topic [partition], timestamp]
+ *    -> value version 0:       [producer_id, producer_epoch, expire_timestamp, status, [topic, [partition] ], timestamp]
  */
 object TransactionLog {
 
@@ -54,91 +54,14 @@ object TransactionLog {
   val EnforcedCompressionType: CompressionType = CompressionType.NONE
   val EnforcedRequiredAcks: Short = (-1).toShort
 
-  // log message formats
-
-  private object KeySchema {
-    private val TXN_ID_KEY = "transactional_id"
-
-    private val V0 = new Schema(new Field(TXN_ID_KEY, STRING))
-    private val SCHEMAS = Map(0 -> V0)
-
-    val CURRENT_VERSION = 0.toShort
-    val CURRENT = schemaForKey(CURRENT_VERSION)
-
-    val TXN_ID_FIELD = V0.get(TXN_ID_KEY)
-
-    def ofVersion(version: Int): Option[Schema] = SCHEMAS.get(version)
-  }
-
-  private object ValueSchema {
-    private val ProducerIdKey = "producer_id"
-    private val ProducerEpochKey = "producer_epoch"
-    private val TxnTimeoutKey = "transaction_timeout"
-    private val TxnStatusKey = "transaction_status"
-    private val TxnPartitionsKey = "transaction_partitions"
-    private val TxnEntryTimestampKey = "transaction_entry_timestamp"
-    private val TxnStartTimestampKey = "transaction_start_timestamp"
-
-    private val PartitionIdsKey = "partition_ids"
-    private val TopicKey = "topic"
-    private val PartitionsSchema = new Schema(new Field(TopicKey, STRING),
-      new Field(PartitionIdsKey, new ArrayOf(INT32)))
-
-    private val V0 = new Schema(new Field(ProducerIdKey, INT64, "Producer id in use by the transactional id."),
-      new Field(ProducerEpochKey, INT16, "Epoch associated with the producer id"),
-      new Field(TxnTimeoutKey, INT32, "Transaction timeout in milliseconds"),
-      new Field(TxnStatusKey, INT8,
-        "TransactionState the transaction is in"),
-      new Field(TxnPartitionsKey, ArrayOf.nullable(PartitionsSchema),
-        "Set of partitions involved in the transaction"),
-      new Field(TxnEntryTimestampKey, INT64, "Time the transaction was last updated"),
-      new Field(TxnStartTimestampKey, INT64, "Time the transaction was started"))
-
-    private val Schemas = Map(0 -> V0)
-
-    val CurrentVersion = 0.toShort
-    val Current = schemaForValue(CurrentVersion)
-
-    val ProducerIdField = V0.get(ProducerIdKey)
-    val ProducerEpochField = V0.get(ProducerEpochKey)
-    val TxnTimeoutField = V0.get(TxnTimeoutKey)
-    val TxnStatusField = V0.get(TxnStatusKey)
-    val TxnPartitionsField = V0.get(TxnPartitionsKey)
-    val TxnEntryTimestampField = V0.get(TxnEntryTimestampKey)
-    val TxnStartTimestampField = V0.get(TxnStartTimestampKey)
-
-    val PartitionsTopicField = PartitionsSchema.get(TopicKey)
-    val PartitionIdsField = PartitionsSchema.get(PartitionIdsKey)
-
-    def ofVersion(version: Int): Option[Schema] = Schemas.get(version)
-  }
-
-  private def schemaForKey(version: Int) = {
-    KeySchema.ofVersion(version).getOrElse {
-      throw new KafkaException(s"Unknown transaction log message key schema version $version")
-    }
-  }
-
-  private def schemaForValue(version: Int) = {
-    ValueSchema.ofVersion(version).getOrElse {
-      throw new KafkaException(s"Unknown transaction log message value schema version $version")
-    }
-  }
-
   /**
     * Generates the bytes for transaction log message key
     *
     * @return key bytes
     */
-  private[coordinator] def keyToBytes(transactionalId: String): Array[Byte] = {
-    import KeySchema._
-    val key = new Struct(CURRENT)
-    key.set(TXN_ID_FIELD, transactionalId)
-
-    val byteBuffer = ByteBuffer.allocate(2 /* version */ + key.sizeOf)
-    byteBuffer.putShort(CURRENT_VERSION)
-    key.writeTo(byteBuffer)
-    byteBuffer.array()
+  private[transaction] def keyToBytes(transactionalId: String): Array[Byte] = {
+    MessageUtil.toVersionPrefixedBytes(TransactionLogKey.HIGHEST_SUPPORTED_VERSION,
+      new TransactionLogKey().setTransactionalId(transactionalId))
   }
 
   /**
@@ -146,39 +69,28 @@ object TransactionLog {
     *
     * @return value payload bytes
     */
-  private[coordinator] def valueToBytes(txnMetadata: TxnTransitMetadata): Array[Byte] = {
-    import ValueSchema._
-    val value = new Struct(Current)
-    value.set(ProducerIdField, txnMetadata.producerId)
-    value.set(ProducerEpochField, txnMetadata.producerEpoch)
-    value.set(TxnTimeoutField, txnMetadata.txnTimeoutMs)
-    value.set(TxnStatusField, txnMetadata.txnState.byte)
-    value.set(TxnEntryTimestampField, txnMetadata.txnLastUpdateTimestamp)
-    value.set(TxnStartTimestampField, txnMetadata.txnStartTimestamp)
-
-    if (txnMetadata.txnState == Empty) {
-      if (txnMetadata.topicPartitions.nonEmpty)
+  private[transaction] def valueToBytes(txnMetadata: TxnTransitMetadata): Array[Byte] = {
+    if (txnMetadata.txnState == Empty && txnMetadata.topicPartitions.nonEmpty)
         throw new IllegalStateException(s"Transaction is not expected to have any partitions since its state is ${txnMetadata.txnState}: $txnMetadata")
 
-      value.set(TxnPartitionsField, null)
-    } else {
-      // first group the topic partitions by their topic names
-      val topicAndPartitions = txnMetadata.topicPartitions.groupBy(_.topic())
-
-      val partitionArray = topicAndPartitions.map { case(topic, partitions) =>
-        val topicPartitionsStruct = value.instance(TxnPartitionsField)
-        val partitionIds: Array[Integer] = partitions.map(topicPartition => Integer.valueOf(topicPartition.partition())).toArray
-        topicPartitionsStruct.set(PartitionsTopicField, topic)
-        topicPartitionsStruct.set(PartitionIdsField, partitionIds)
-        topicPartitionsStruct
-      }
-      value.set(TxnPartitionsField, partitionArray.toArray)
-    }
-
-    val byteBuffer = ByteBuffer.allocate(2 /* version */ + value.sizeOf)
-    byteBuffer.putShort(CurrentVersion)
-    value.writeTo(byteBuffer)
-    byteBuffer.array()
+      val transactionPartitions = if (txnMetadata.txnState == Empty) null
+      else txnMetadata.topicPartitions
+        .groupBy(_.topic)
+        .map { case (topic, partitions) =>
+          new TransactionLogValue.PartitionsSchema()
+            .setTopic(topic)
+            .setPartitionIds(partitions.map(tp => Integer.valueOf(tp.partition)).toList.asJava)
+        }.toList.asJava
+
+    MessageUtil.toVersionPrefixedBytes(TransactionLogValue.HIGHEST_SUPPORTED_VERSION,
+      new TransactionLogValue()
+        .setProducerId(txnMetadata.producerId)
+        .setProducerEpoch(txnMetadata.producerEpoch)
+        .setTransactionTimeoutMs(txnMetadata.txnTimeoutMs)
+        .setTransactionStatus(txnMetadata.txnState.byte)
+        .setTransactionLastUpdateTimestampMs(txnMetadata.txnLastUpdateTimestamp)
+        .setTransactionStartTimestampMs(txnMetadata.txnStartTimestamp)
+        .setTransactionPartitions(transactionPartitions))
   }
 
   /**
@@ -188,15 +100,13 @@ object TransactionLog {
     */
   def readTxnRecordKey(buffer: ByteBuffer): TxnKey = {
     val version = buffer.getShort
-    val keySchema = schemaForKey(version)
-    val key = keySchema.read(buffer)
-
-    if (version == KeySchema.CURRENT_VERSION) {
-      val transactionalId = key.getString(KeySchema.TXN_ID_FIELD)
-      TxnKey(version, transactionalId)
-    } else {
-      throw new IllegalStateException(s"Unknown version $version from the transaction log message")
-    }
+    if (version >= TransactionLogKey.LOWEST_SUPPORTED_VERSION && version <= TransactionLogKey.HIGHEST_SUPPORTED_VERSION) {
+      val value = new TransactionLogKey(new ByteBufferAccessor(buffer), version)
+      TxnKey(
+        version = version,
+        transactionalId = value.transactionalId
+      )
+    } else throw new IllegalStateException(s"Unknown version $version from the transaction log message")
   }
 
   /**
@@ -204,70 +114,78 @@ object TransactionLog {
     *
     * @return a transaction metadata object from the message
     */
-  def readTxnRecordValue(transactionalId: String, buffer: ByteBuffer): TransactionMetadata = {
-    if (buffer == null) { // tombstone
-      null
-    } else {
-      import ValueSchema._
+  def readTxnRecordValue(transactionalId: String, buffer: ByteBuffer): Option[TransactionMetadata] = {
+    // tombstone
+    if (buffer == null) None
+    else {
       val version = buffer.getShort
-      val valueSchema = schemaForValue(version)
-      val value = valueSchema.read(buffer)
-
-      if (version == CurrentVersion) {
-        val producerId = value.getLong(ProducerIdField)
-        val epoch = value.getShort(ProducerEpochField)
-        val timeout = value.getInt(TxnTimeoutField)
-
-        val stateByte = value.getByte(TxnStatusField)
-        val state = TransactionMetadata.byteToState(stateByte)
-        val entryTimestamp = value.getLong(TxnEntryTimestampField)
-        val startTimestamp = value.getLong(TxnStartTimestampField)
-
-        val transactionMetadata = new TransactionMetadata(transactionalId, producerId, epoch, timeout, state,
-          mutable.Set.empty[TopicPartition],startTimestamp, entryTimestamp)
-
-        if (!state.equals(Empty)) {
-          val topicPartitionArray = value.getArray(TxnPartitionsField)
-
-          topicPartitionArray.foreach { memberMetadataObj =>
-            val memberMetadata = memberMetadataObj.asInstanceOf[Struct]
-            val topic = memberMetadata.getString(PartitionsTopicField)
-            val partitionIdArray = memberMetadata.getArray(PartitionIdsField)
-
-            val topicPartitions = partitionIdArray.map { partitionIdObj =>
-              val partitionId = partitionIdObj.asInstanceOf[Integer]
-              new TopicPartition(topic, partitionId)
-            }
-
-            transactionMetadata.addPartitions(topicPartitions.toSet)
-          }
-        }
-
-        transactionMetadata
-      } else {
-        throw new IllegalStateException(s"Unknown version $version from the transaction log message value")
-      }
+      if (version >= TransactionLogValue.LOWEST_SUPPORTED_VERSION && version <= TransactionLogValue.HIGHEST_SUPPORTED_VERSION) {
+        val value = new TransactionLogValue(new ByteBufferAccessor(buffer), version)
+        val transactionMetadata = new TransactionMetadata(
+          transactionalId = transactionalId,
+          producerId = value.producerId,
+          lastProducerId = RecordBatch.NO_PRODUCER_ID,
+          producerEpoch = value.producerEpoch,
+          lastProducerEpoch = RecordBatch.NO_PRODUCER_EPOCH,
+          txnTimeoutMs = value.transactionTimeoutMs,
+          state = TransactionMetadata.byteToState(value.transactionStatus),
+          topicPartitions = mutable.Set.empty[TopicPartition],
+          txnStartTimestamp = value.transactionStartTimestampMs,
+          txnLastUpdateTimestamp = value.transactionLastUpdateTimestampMs)
+
+        if (!transactionMetadata.state.equals(Empty))
+          value.transactionPartitions.forEach(partitionsSchema =>
+            transactionMetadata.addPartitions(partitionsSchema.partitionIds
+              .asScala
+              .map(partitionId => new TopicPartition(partitionsSchema.topic, partitionId))
+              .toSet)
+          )
+        Some(transactionMetadata)
+      } else throw new IllegalStateException(s"Unknown version $version from the transaction log message value")
     }
   }
 
   // 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) {
-  override def toString: String = transactionalId.toString
+  override def toString: String = transactionalId
 }
diff --git a/core/src/main/scala/kafka/coordinator/transaction/TransactionMarkerChannelManager.scala b/core/src/main/scala/kafka/coordinator/transaction/TransactionMarkerChannelManager.scala
index fa9d2c3fed57e..029ded837a87d 100644
--- a/core/src/main/scala/kafka/coordinator/transaction/TransactionMarkerChannelManager.scala
+++ b/core/src/main/scala/kafka/coordinator/transaction/TransactionMarkerChannelManager.scala
@@ -17,24 +17,26 @@
 package kafka.coordinator.transaction
 
 
+import java.util
+import java.util.concurrent.{BlockingQueue, ConcurrentHashMap, LinkedBlockingQueue}
+
+import kafka.api.KAFKA_2_8_IV0
 import kafka.common.{InterBrokerSendThread, RequestAndCompletionHandler}
 import kafka.metrics.KafkaMetricsGroup
-import kafka.server.{DelayedOperationPurgatory, KafkaConfig, MetadataCache}
+import kafka.server.{KafkaConfig, MetadataCache}
 import kafka.utils.{CoreUtils, Logging}
+import kafka.utils.Implicits._
 import org.apache.kafka.clients._
-import org.apache.kafka.common.{Node, TopicPartition}
 import org.apache.kafka.common.metrics.Metrics
 import org.apache.kafka.common.network._
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.requests.WriteTxnMarkersRequest.TxnMarkerEntry
 import org.apache.kafka.common.requests.{TransactionResult, WriteTxnMarkersRequest}
 import org.apache.kafka.common.security.JaasContext
 import org.apache.kafka.common.utils.{LogContext, Time}
-import org.apache.kafka.common.protocol.Errors
-import org.apache.kafka.common.requests.WriteTxnMarkersRequest.TxnMarkerEntry
-import com.yammer.metrics.core.Gauge
-import java.util
-import java.util.concurrent.{BlockingQueue, ConcurrentHashMap, LinkedBlockingQueue}
+import org.apache.kafka.common.{Node, Reconfigurable, TopicPartition}
 
-import collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 import scala.collection.{concurrent, immutable}
 
 object TransactionMarkerChannelManager {
@@ -42,7 +44,6 @@ object TransactionMarkerChannelManager {
             metrics: Metrics,
             metadataCache: MetadataCache,
             txnStateManager: TransactionStateManager,
-            txnMarkerPurgatory: DelayedOperationPurgatory[DelayedTxnMarker],
             time: Time,
             logContext: LogContext): TransactionMarkerChannelManager = {
     val channelBuilder = ChannelBuilders.clientChannelBuilder(
@@ -51,8 +52,14 @@ object TransactionMarkerChannelManager {
       config,
       config.interBrokerListenerName,
       config.saslMechanismInterBrokerProtocol,
-      config.saslInterBrokerHandshakeRequestEnable
+      time,
+      config.saslInterBrokerHandshakeRequestEnable,
+      logContext
     )
+    channelBuilder match {
+      case reconfigurable: Reconfigurable => config.addReconfigurable(reconfigurable)
+      case _ =>
+    }
     val selector = new Selector(
       NetworkReceive.UNLIMITED,
       config.connectionsMaxIdleMs,
@@ -74,6 +81,9 @@ object TransactionMarkerChannelManager {
       Selectable.USE_DEFAULT_BUFFER_SIZE,
       config.socketReceiveBufferBytes,
       config.requestTimeoutMs,
+      config.connectionSetupTimeoutMs,
+      config.connectionSetupTimeoutMaxMs,
+      ClientDnsLookup.USE_ALL_DNS_IPS,
       time,
       false,
       new ApiVersions,
@@ -84,7 +94,6 @@ object TransactionMarkerChannelManager {
       metadataCache,
       networkClient,
       txnStateManager,
-      txnMarkerPurgatory,
       time
     )
   }
@@ -108,7 +117,7 @@ class TxnMarkerQueue(@volatile var destination: Node) {
   }
 
   def forEachTxnTopicPartition[B](f:(Int, BlockingQueue[TxnIdAndMarkerEntry]) => B): Unit =
-    markersPerTxnTopicPartition.foreach { case (partition, queue) =>
+    markersPerTxnTopicPartition.forKeyValue { (partition, queue) =>
       if (!queue.isEmpty) f(partition, queue)
     }
 
@@ -122,7 +131,6 @@ class TransactionMarkerChannelManager(config: KafkaConfig,
                                       metadataCache: MetadataCache,
                                       networkClient: NetworkClient,
                                       txnStateManager: TransactionStateManager,
-                                      txnMarkerPurgatory: DelayedOperationPurgatory[DelayedTxnMarker],
                                       time: Time) extends InterBrokerSendThread("TxnMarkerSenderThread-" + config.brokerId, networkClient, time) with Logging with KafkaMetricsGroup {
 
   this.logIdent = "[Transaction Marker Channel Manager " + config.brokerId + "]: "
@@ -133,27 +141,23 @@ class TransactionMarkerChannelManager(config: KafkaConfig,
 
   private val markersQueueForUnknownBroker = new TxnMarkerQueue(Node.noNode)
 
-  private val txnLogAppendRetryQueue = new LinkedBlockingQueue[TxnLogAppend]()
+  private val txnLogAppendRetryQueue = new LinkedBlockingQueue[PendingCompleteTxn]()
 
-  newGauge(
-    "UnknownDestinationQueueSize",
-    new Gauge[Int] {
-      def value: Int = markersQueueForUnknownBroker.totalNumMarkers
-    }
-  )
+  private val transactionsWithPendingMarkers = new ConcurrentHashMap[String, PendingCompleteTxn]
 
-  newGauge(
-    "LogAppendRetryQueueSize",
-    new Gauge[Int] {
-      def value: Int = txnLogAppendRetryQueue.size
-    }
-  )
+  override val requestTimeoutMs: Int = config.requestTimeoutMs
+
+  val writeTxnMarkersRequestVersion: Short =
+    if (config.interBrokerProtocolVersion >= KAFKA_2_8_IV0) 1
+    else 0
+
+  newGauge("UnknownDestinationQueueSize", () => markersQueueForUnknownBroker.totalNumMarkers)
+  newGauge("LogAppendRetryQueueSize", () => txnLogAppendRetryQueue.size)
 
   override def generateRequests() = drainQueuedTransactionMarkers()
 
   override def shutdown(): Unit = {
     super.shutdown()
-    txnMarkerPurgatory.shutdown()
     markersQueuePerBroker.clear()
   }
 
@@ -165,7 +169,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,
@@ -179,9 +183,9 @@ class TransactionMarkerChannelManager(config: KafkaConfig,
   }
 
   def retryLogAppends(): Unit = {
-    val txnLogAppendRetries: java.util.List[TxnLogAppend] = new util.ArrayList[TxnLogAppend]()
+    val txnLogAppendRetries: java.util.List[PendingCompleteTxn] = new util.ArrayList[PendingCompleteTxn]()
     txnLogAppendRetryQueue.drainTo(txnLogAppendRetries)
-    txnLogAppendRetries.asScala.foreach { txnLogAppend =>
+    txnLogAppendRetries.forEach { txnLogAppend =>
       debug(s"Retry appending $txnLogAppend transaction log")
       tryAppendToLog(txnLogAppend)
     }
@@ -214,63 +218,86 @@ class TransactionMarkerChannelManager(config: KafkaConfig,
     }.filter { case (_, entries) => !entries.isEmpty }.map { case (node, entries) =>
       val markersToSend = entries.asScala.map(_.txnMarkerEntry).asJava
       val requestCompletionHandler = new TransactionMarkerRequestCompletionHandler(node.id, txnStateManager, this, entries)
-      RequestAndCompletionHandler(node, new WriteTxnMarkersRequest.Builder(markersToSend), requestCompletionHandler)
+      RequestAndCompletionHandler(node, new WriteTxnMarkersRequest.Builder(writeTxnMarkersRequestVersion, markersToSend), requestCompletionHandler)
     }
   }
 
-  def addTxnMarkersToSend(transactionalId: String,
-                          coordinatorEpoch: Int,
+  private def writeTxnCompletion(pendingCompleteTxn: PendingCompleteTxn): Unit = {
+    val transactionalId = pendingCompleteTxn.transactionalId
+    val txnMetadata = pendingCompleteTxn.txnMetadata
+    val newMetadata = pendingCompleteTxn.newMetadata
+    val coordinatorEpoch = pendingCompleteTxn.coordinatorEpoch
+
+    trace(s"Completed sending transaction markers for $transactionalId; begin transition " +
+      s"to ${newMetadata.txnState}")
+
+    txnStateManager.getTransactionState(transactionalId) match {
+      case Left(Errors.NOT_COORDINATOR) =>
+        info(s"No longer the coordinator for $transactionalId with coordinator epoch " +
+          s"$coordinatorEpoch; cancel appending $newMetadata to transaction log")
+
+      case Left(Errors.COORDINATOR_LOAD_IN_PROGRESS) =>
+        info(s"Loading the transaction partition that contains $transactionalId while my " +
+          s"current coordinator epoch is $coordinatorEpoch; so cancel appending $newMetadata to " +
+          s"transaction log since the loading process will continue the remaining work")
+
+      case Left(unexpectedError) =>
+        throw new IllegalStateException(s"Unhandled error $unexpectedError when fetching current transaction state")
+
+      case Right(Some(epochAndMetadata)) =>
+        if (epochAndMetadata.coordinatorEpoch == coordinatorEpoch) {
+          debug(s"Sending $transactionalId's transaction markers for $txnMetadata with " +
+            s"coordinator epoch $coordinatorEpoch succeeded, trying to append complete transaction log now")
+          tryAppendToLog(PendingCompleteTxn(transactionalId, coordinatorEpoch, txnMetadata, newMetadata))
+        } else {
+          info(s"The cached metadata $txnMetadata has changed to $epochAndMetadata after " +
+            s"completed sending the markers with coordinator epoch $coordinatorEpoch; abort " +
+            s"transiting the metadata to $newMetadata as it may have been updated by another process")
+        }
+
+      case Right(None) =>
+        val errorMsg = s"The coordinator still owns the transaction partition for $transactionalId, " +
+          s"but there is no metadata in the cache; this is not expected"
+        fatal(errorMsg)
+        throw new IllegalStateException(errorMsg)
+    }
+  }
+
+  def addTxnMarkersToSend(coordinatorEpoch: Int,
                           txnResult: TransactionResult,
                           txnMetadata: TransactionMetadata,
                           newMetadata: TxnTransitMetadata): Unit = {
+    val transactionalId = txnMetadata.transactionalId
+    val pendingCompleteTxn = PendingCompleteTxn(
+      transactionalId,
+      coordinatorEpoch,
+      txnMetadata,
+      newMetadata)
+
+    transactionsWithPendingMarkers.put(transactionalId, pendingCompleteTxn)
+    addTxnMarkersToBrokerQueue(transactionalId, txnMetadata.producerId,
+      txnMetadata.producerEpoch, txnResult, coordinatorEpoch, txnMetadata.topicPartitions.toSet)
+    maybeWriteTxnCompletion(transactionalId)
+  }
 
-    def appendToLogCallback(error: Errors): Unit = {
-      error match {
-        case Errors.NONE =>
-          trace(s"Completed sending transaction markers for $transactionalId as $txnResult")
-
-          txnStateManager.getTransactionState(transactionalId) match {
-            case Left(Errors.NOT_COORDINATOR) =>
-              info(s"No longer the coordinator for $transactionalId with coordinator epoch $coordinatorEpoch; cancel appending $newMetadata to transaction log")
-
-            case Left(Errors.COORDINATOR_LOAD_IN_PROGRESS) =>
-              info(s"Loading the transaction partition that contains $transactionalId while my current coordinator epoch is $coordinatorEpoch; " +
-                s"so cancel appending $newMetadata to transaction log since the loading process will continue the remaining work")
-
-            case Left(unexpectedError) =>
-              throw new IllegalStateException(s"Unhandled error $unexpectedError when fetching current transaction state")
-
-            case Right(Some(epochAndMetadata)) =>
-              if (epochAndMetadata.coordinatorEpoch == coordinatorEpoch) {
-                debug(s"Sending $transactionalId's transaction markers for $txnMetadata with coordinator epoch $coordinatorEpoch succeeded, trying to append complete transaction log now")
+  def numTxnsWithPendingMarkers: Int = transactionsWithPendingMarkers.size
 
-                tryAppendToLog(TxnLogAppend(transactionalId, coordinatorEpoch, txnMetadata, newMetadata))
-              } else {
-                info(s"The cached metadata $txnMetadata has changed to $epochAndMetadata after completed sending the markers with coordinator " +
-                  s"epoch $coordinatorEpoch; abort transiting the metadata to $newMetadata as it may have been updated by another process")
-              }
-
-            case Right(None) =>
-              val errorMsg = s"The coordinator still owns the transaction partition for $transactionalId, but there is " +
-                s"no metadata in the cache; this is not expected"
-              fatal(errorMsg)
-              throw new IllegalStateException(errorMsg)
-          }
+  private def hasPendingMarkersToWrite(txnMetadata: TransactionMetadata): Boolean = {
+    txnMetadata.inLock {
+      txnMetadata.topicPartitions.nonEmpty
+    }
+  }
 
-        case other =>
-          val errorMsg = s"Unexpected error ${other.exceptionName} before appending to txn log for $transactionalId"
-          fatal(errorMsg)
-          throw new IllegalStateException(errorMsg)
+  def maybeWriteTxnCompletion(transactionalId: String): Unit = {
+    Option(transactionsWithPendingMarkers.get(transactionalId)).foreach { pendingCompleteTxn =>
+      if (!hasPendingMarkersToWrite(pendingCompleteTxn.txnMetadata) &&
+          transactionsWithPendingMarkers.remove(transactionalId, pendingCompleteTxn)) {
+        writeTxnCompletion(pendingCompleteTxn)
       }
     }
-
-    val delayedTxnMarker = new DelayedTxnMarker(txnMetadata, appendToLogCallback, txnStateManager.stateReadLock)
-    txnMarkerPurgatory.tryCompleteElseWatch(delayedTxnMarker, Seq(transactionalId))
-
-    addTxnMarkersToBrokerQueue(transactionalId, txnMetadata.producerId, txnMetadata.producerEpoch, txnResult, coordinatorEpoch, txnMetadata.topicPartitions.toSet)
   }
 
-  private def tryAppendToLog(txnLogAppend: TxnLogAppend) = {
+  private def tryAppendToLog(txnLogAppend: PendingCompleteTxn): Unit = {
     // try to append to the transaction log
     def appendCallback(error: Errors): Unit =
       error match {
@@ -301,8 +328,11 @@ class TransactionMarkerChannelManager(config: KafkaConfig,
       _ == Errors.COORDINATOR_NOT_AVAILABLE)
   }
 
-  def addTxnMarkersToBrokerQueue(transactionalId: String, producerId: Long, producerEpoch: Short,
-                                 result: TransactionResult, coordinatorEpoch: Int,
+  def addTxnMarkersToBrokerQueue(transactionalId: String,
+                                 producerId: Long,
+                                 producerEpoch: Short,
+                                 result: TransactionResult,
+                                 coordinatorEpoch: Int,
                                  topicPartitions: immutable.Set[TopicPartition]): Unit = {
     val txnTopicPartition = txnStateManager.partitionFor(transactionalId)
     val partitionsByDestination: immutable.Map[Option[Node], immutable.Set[TopicPartition]] = topicPartitions.groupBy { topicPartition: TopicPartition =>
@@ -327,12 +357,12 @@ class TransactionMarkerChannelManager(config: KafkaConfig,
           txnStateManager.getTransactionState(transactionalId) match {
             case Left(error) =>
               info(s"Encountered $error trying to fetch transaction metadata for $transactionalId with coordinator epoch $coordinatorEpoch; cancel sending markers to its partition leaders")
-              txnMarkerPurgatory.cancelForKey(transactionalId)
+              transactionsWithPendingMarkers.remove(transactionalId)
 
             case Right(Some(epochAndMetadata)) =>
               if (epochAndMetadata.coordinatorEpoch != coordinatorEpoch) {
                 info(s"The cached metadata has changed to $epochAndMetadata (old coordinator epoch is $coordinatorEpoch) since preparing to send markers; cancel sending markers to its partition leaders")
-                txnMarkerPurgatory.cancelForKey(transactionalId)
+                transactionsWithPendingMarkers.remove(transactionalId)
               } else {
                 // if the leader of the partition is unknown, skip sending the txn marker since
                 // the partition is likely to be deleted already
@@ -345,7 +375,7 @@ class TransactionMarkerChannelManager(config: KafkaConfig,
                   topicPartitions.foreach(txnMetadata.removePartition)
                 }
 
-                txnMarkerPurgatory.checkAndComplete(transactionalId)
+                maybeWriteTxnCompletion(transactionalId)
               }
 
             case Right(None) =>
@@ -376,22 +406,19 @@ class TransactionMarkerChannelManager(config: KafkaConfig,
   }
 
   def removeMarkersForTxnId(transactionalId: String): Unit = {
-    // we do not need to clear the queue since it should have
-    // already been drained by the sender thread
-    txnMarkerPurgatory.cancelForKey(transactionalId)
-  }
-
-  def completeSendMarkersForTxnId(transactionalId: String): Unit = {
-    txnMarkerPurgatory.checkAndComplete(transactionalId)
+    transactionsWithPendingMarkers.remove(transactionalId)
   }
 }
 
 case class TxnIdAndMarkerEntry(txnId: String, txnMarkerEntry: TxnMarkerEntry)
 
-case class TxnLogAppend(transactionalId: String, coordinatorEpoch: Int, txnMetadata: TransactionMetadata, newMetadata: TxnTransitMetadata) {
+case class PendingCompleteTxn(transactionalId: String,
+                              coordinatorEpoch: Int,
+                              txnMetadata: TransactionMetadata,
+                              newMetadata: TxnTransitMetadata) {
 
   override def toString: String = {
-    "TxnLogAppend(" +
+    "PendingCompleteTxn(" +
       s"transactionalId=$transactionalId, " +
       s"coordinatorEpoch=$coordinatorEpoch, " +
       s"txnMetadata=$txnMetadata, " +
diff --git a/core/src/main/scala/kafka/coordinator/transaction/TransactionMarkerRequestCompletionHandler.scala b/core/src/main/scala/kafka/coordinator/transaction/TransactionMarkerRequestCompletionHandler.scala
index fefe767eb40b8..848e0fa65ceeb 100644
--- a/core/src/main/scala/kafka/coordinator/transaction/TransactionMarkerRequestCompletionHandler.scala
+++ b/core/src/main/scala/kafka/coordinator/transaction/TransactionMarkerRequestCompletionHandler.scala
@@ -24,7 +24,7 @@ import org.apache.kafka.common.protocol.Errors
 import org.apache.kafka.common.requests.WriteTxnMarkersResponse
 
 import scala.collection.mutable
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 
 class TransactionMarkerRequestCompletionHandler(brokerId: Int,
                                                 txnStateManager: TransactionStateManager,
@@ -89,10 +89,11 @@ class TransactionMarkerRequestCompletionHandler(brokerId: Int,
 
       val writeTxnMarkerResponse = response.responseBody.asInstanceOf[WriteTxnMarkersResponse]
 
+      val responseErrors = writeTxnMarkerResponse.errorsByProducerId;
       for (txnIdAndMarker <- txnIdAndMarkerEntries.asScala) {
         val transactionalId = txnIdAndMarker.txnId
         val txnMarker = txnIdAndMarker.txnMarkerEntry
-        val errors = writeTxnMarkerResponse.errors(txnMarker.producerId)
+        val errors = responseErrors.get(txnMarker.producerId)
 
         if (errors == null)
           throw new IllegalStateException(s"WriteTxnMarkerResponse does not contain expected error map for producer id ${txnMarker.producerId}")
@@ -143,10 +144,11 @@ class TransactionMarkerRequestCompletionHandler(brokerId: Int,
                       throw new IllegalStateException(s"Received fatal error ${error.exceptionName} while sending txn marker for $transactionalId")
 
                     case Errors.UNKNOWN_TOPIC_OR_PARTITION |
-                         Errors.NOT_LEADER_FOR_PARTITION |
+                         Errors.NOT_LEADER_OR_FOLLOWER |
                          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}")
@@ -192,7 +194,7 @@ class TransactionMarkerRequestCompletionHandler(brokerId: Int,
                   txnMarker.coordinatorEpoch,
                   retryPartitions.toSet)
               } else {
-                txnMarkerChannelManager.completeSendMarkersForTxnId(transactionalId)
+                txnMarkerChannelManager.maybeWriteTxnCompletion(transactionalId)
               }
             }
         }
diff --git a/core/src/main/scala/kafka/coordinator/transaction/TransactionMetadata.scala b/core/src/main/scala/kafka/coordinator/transaction/TransactionMetadata.scala
index ea82fb5f814af..e059b04a37374 100644
--- a/core/src/main/scala/kafka/coordinator/transaction/TransactionMetadata.scala
+++ b/core/src/main/scala/kafka/coordinator/transaction/TransactionMetadata.scala
@@ -20,6 +20,7 @@ import java.util.concurrent.locks.ReentrantLock
 
 import kafka.utils.{CoreUtils, Logging, nonthreadsafe}
 import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.protocol.Errors
 import org.apache.kafka.common.record.RecordBatch
 
 import scala.collection.{immutable, mutable}
@@ -85,13 +86,18 @@ private[transaction] case object PrepareEpochFence extends TransactionState { va
 
 private[transaction] object TransactionMetadata {
   def apply(transactionalId: String, producerId: Long, producerEpoch: Short, txnTimeoutMs: Int, timestamp: Long) =
-    new TransactionMetadata(transactionalId, producerId, producerEpoch, txnTimeoutMs, Empty,
-      collection.mutable.Set.empty[TopicPartition], timestamp, timestamp)
+    new TransactionMetadata(transactionalId, producerId, RecordBatch.NO_PRODUCER_ID, producerEpoch,
+      RecordBatch.NO_PRODUCER_EPOCH, txnTimeoutMs, Empty, collection.mutable.Set.empty[TopicPartition], timestamp, timestamp)
 
   def apply(transactionalId: String, producerId: Long, producerEpoch: Short, txnTimeoutMs: Int,
             state: TransactionState, timestamp: Long) =
-    new TransactionMetadata(transactionalId, producerId, producerEpoch, txnTimeoutMs, state,
-      collection.mutable.Set.empty[TopicPartition], timestamp, timestamp)
+    new TransactionMetadata(transactionalId, producerId, RecordBatch.NO_PRODUCER_ID, producerEpoch,
+      RecordBatch.NO_PRODUCER_EPOCH, txnTimeoutMs, state, collection.mutable.Set.empty[TopicPartition], timestamp, timestamp)
+
+  def apply(transactionalId: String, producerId: Long, lastProducerId: Long, producerEpoch: Short,
+            lastProducerEpoch: Short, txnTimeoutMs: Int, state: TransactionState, timestamp: Long) =
+    new TransactionMetadata(transactionalId, producerId, lastProducerId, producerEpoch, lastProducerEpoch,
+      txnTimeoutMs, state, collection.mutable.Set.empty[TopicPartition], timestamp, timestamp)
 
   def byteToState(byte: Byte): TransactionState = {
     byte match {
@@ -120,11 +126,15 @@ private[transaction] object TransactionMetadata {
       Dead -> Set(Empty, CompleteAbort, CompleteCommit),
       PrepareEpochFence -> Set(Ongoing)
     )
+
+  def isEpochExhausted(producerEpoch: Short): Boolean = producerEpoch >= Short.MaxValue - 1
 }
 
 // this is a immutable object representing the target transition of the transaction metadata
 private[transaction] case class TxnTransitMetadata(producerId: Long,
+                                                   lastProducerId: Long,
                                                    producerEpoch: Short,
+                                                   lastProducerEpoch: Short,
                                                    txnTimeoutMs: Int,
                                                    txnState: TransactionState,
                                                    topicPartitions: immutable.Set[TopicPartition],
@@ -133,7 +143,9 @@ private[transaction] case class TxnTransitMetadata(producerId: Long,
   override def toString: String = {
     "TxnTransitMetadata(" +
       s"producerId=$producerId, " +
+      s"lastProducerId=$lastProducerId, " +
       s"producerEpoch=$producerEpoch, " +
+      s"lastProducerEpoch=$lastProducerEpoch, " +
       s"txnTimeoutMs=$txnTimeoutMs, " +
       s"txnState=$txnState, " +
       s"topicPartitions=$topicPartitions, " +
@@ -145,7 +157,9 @@ private[transaction] case class TxnTransitMetadata(producerId: Long,
 /**
   *
   * @param producerId            producer id
+  * @param lastProducerId        last producer id assigned to the producer
   * @param producerEpoch         current epoch of the producer
+  * @param lastProducerEpoch     last epoch of the producer
   * @param txnTimeoutMs          timeout to be used to abort long running transactions
   * @param state                 current state of the transaction
   * @param topicPartitions       current set of partitions that are part of this transaction
@@ -155,7 +169,9 @@ private[transaction] case class TxnTransitMetadata(producerId: Long,
 @nonthreadsafe
 private[transaction] class TransactionMetadata(val transactionalId: String,
                                                var producerId: Long,
+                                               var lastProducerId: Long,
                                                var producerEpoch: Short,
+                                               var lastProducerEpoch: Short,
                                                var txnTimeoutMs: Int,
                                                var state: TransactionState,
                                                val topicPartitions: mutable.Set[TopicPartition],
@@ -167,6 +183,10 @@ private[transaction] class TransactionMetadata(val transactionalId: String,
   // initialized as the same as the current state
   var pendingState: Option[TransactionState] = None
 
+  // Indicates that during a previous attempt to fence a producer, the bumped epoch may not have been
+  // successfully written to the log. If this is true, we will not bump the epoch again when fencing
+  var hasFailedEpochFence: Boolean = false
+
   private[transaction] val lock = new ReentrantLock
 
   def inLock[T](fun: => T): T = CoreUtils.inLock(lock)(fun)
@@ -186,30 +206,73 @@ private[transaction] class TransactionMetadata(val transactionalId: String,
   // this is visible for test only
   def prepareNoTransit(): TxnTransitMetadata = {
     // do not call transitTo as it will set the pending state, a follow-up call to abort the transaction will set its pending state
-    TxnTransitMetadata(producerId, producerEpoch, txnTimeoutMs, state, topicPartitions.toSet, txnStartTimestamp, txnLastUpdateTimestamp)
+    TxnTransitMetadata(producerId, lastProducerId, producerEpoch, lastProducerEpoch, txnTimeoutMs, state, topicPartitions.toSet,
+      txnStartTimestamp, txnLastUpdateTimestamp)
   }
 
   def prepareFenceProducerEpoch(): TxnTransitMetadata = {
     if (producerEpoch == Short.MaxValue)
       throw new IllegalStateException(s"Cannot fence producer with epoch equal to Short.MaxValue since this would overflow")
 
-    prepareTransitionTo(PrepareEpochFence, producerId, (producerEpoch + 1).toShort, txnTimeoutMs, topicPartitions.toSet,
-      txnStartTimestamp, txnLastUpdateTimestamp)
+    // If we've already failed to fence an epoch (because the write to the log failed), we don't increase it again.
+    // This is safe because we never return the epoch to client if we fail to fence the epoch
+    val bumpedEpoch = if (hasFailedEpochFence) producerEpoch else (producerEpoch + 1).toShort
+
+    prepareTransitionTo(PrepareEpochFence, producerId, bumpedEpoch, RecordBatch.NO_PRODUCER_EPOCH, txnTimeoutMs,
+      topicPartitions.toSet, txnStartTimestamp, txnLastUpdateTimestamp)
   }
 
-  def prepareIncrementProducerEpoch(newTxnTimeoutMs: Int, updateTimestamp: Long): TxnTransitMetadata = {
+  def prepareIncrementProducerEpoch(newTxnTimeoutMs: Int,
+                                    expectedProducerEpoch: Option[Short],
+                                    updateTimestamp: Long): Either[Errors, TxnTransitMetadata] = {
     if (isProducerEpochExhausted)
       throw new IllegalStateException(s"Cannot allocate any more producer epochs for producerId $producerId")
 
-    val nextEpoch = if (producerEpoch == RecordBatch.NO_PRODUCER_EPOCH) 0 else producerEpoch + 1
-    prepareTransitionTo(Empty, producerId, nextEpoch.toShort, newTxnTimeoutMs, immutable.Set.empty[TopicPartition], -1,
-      updateTimestamp)
+    val bumpedEpoch = (producerEpoch + 1).toShort
+    val epochBumpResult: Either[Errors, (Short, Short)] = expectedProducerEpoch match {
+      case None =>
+        // If no expected epoch was provided by the producer, bump the current epoch and set the last epoch to -1
+        // In the case of a new producer, producerEpoch will be -1 and bumpedEpoch will be 0
+        Right(bumpedEpoch, RecordBatch.NO_PRODUCER_EPOCH)
+
+      case Some(expectedEpoch) =>
+        if (producerEpoch == RecordBatch.NO_PRODUCER_EPOCH || expectedEpoch == producerEpoch)
+          // If the expected epoch matches the current epoch, or if there is no current epoch, the producer is attempting
+          // to continue after an error and no other producer has been initialized. Bump the current and last epochs.
+          // The no current epoch case means this is a new producer; producerEpoch will be -1 and bumpedEpoch will be 0
+          Right(bumpedEpoch, producerEpoch)
+        else if (expectedEpoch == lastProducerEpoch)
+          // If the expected epoch matches the previous epoch, it is a retry of a successful call, so just return the
+          // current epoch without bumping. There is no danger of this producer being fenced, because a new producer
+          // calling InitProducerId would have caused the last epoch to be set to -1.
+          // Note that if the IBP is prior to 2.4.IV1, the lastProducerId and lastProducerEpoch will not be written to
+          // the transaction log, so a retry that spans a coordinator change will fail. We expect this to be a rare case.
+          Right(producerEpoch, lastProducerEpoch)
+        else {
+          // Otherwise, the producer has a fenced epoch and should receive an PRODUCER_FENCED error
+          info(s"Expected producer epoch $expectedEpoch does not match current " +
+            s"producer epoch $producerEpoch or previous producer epoch $lastProducerEpoch")
+          Left(Errors.PRODUCER_FENCED)
+        }
+    }
+
+    epochBumpResult match {
+      case Right((nextEpoch, lastEpoch)) => Right(prepareTransitionTo(Empty, producerId, nextEpoch, lastEpoch, newTxnTimeoutMs,
+        immutable.Set.empty[TopicPartition], -1, updateTimestamp))
+
+      case Left(err) => Left(err)
+    }
   }
 
-  def prepareProducerIdRotation(newProducerId: Long, newTxnTimeoutMs: Int, updateTimestamp: Long): TxnTransitMetadata = {
+  def prepareProducerIdRotation(newProducerId: Long,
+                                newTxnTimeoutMs: Int,
+                                updateTimestamp: Long,
+                                recordLastEpoch: Boolean): TxnTransitMetadata = {
     if (hasPendingTransaction)
       throw new IllegalStateException("Cannot rotate producer ids while a transaction is still pending")
-    prepareTransitionTo(Empty, newProducerId, 0, newTxnTimeoutMs, immutable.Set.empty[TopicPartition], -1, updateTimestamp)
+
+    prepareTransitionTo(Empty, newProducerId, 0, if (recordLastEpoch) producerEpoch else RecordBatch.NO_PRODUCER_EPOCH,
+      newTxnTimeoutMs, immutable.Set.empty[TopicPartition], -1, updateTimestamp)
   }
 
   def prepareAddPartitions(addedTopicPartitions: immutable.Set[TopicPartition], updateTimestamp: Long): TxnTransitMetadata = {
@@ -218,31 +281,34 @@ private[transaction] class TransactionMetadata(val transactionalId: String,
       case _ => txnStartTimestamp
     }
 
-    prepareTransitionTo(Ongoing, producerId, producerEpoch, txnTimeoutMs, (topicPartitions ++ addedTopicPartitions).toSet,
-      newTxnStartTimestamp, updateTimestamp)
+    prepareTransitionTo(Ongoing, producerId, producerEpoch, lastProducerEpoch, txnTimeoutMs,
+      (topicPartitions ++ addedTopicPartitions).toSet, newTxnStartTimestamp, updateTimestamp)
   }
 
   def prepareAbortOrCommit(newState: TransactionState, updateTimestamp: Long): TxnTransitMetadata = {
-    prepareTransitionTo(newState, producerId, producerEpoch, txnTimeoutMs, topicPartitions.toSet, txnStartTimestamp,
-      updateTimestamp)
+    prepareTransitionTo(newState, producerId, producerEpoch, lastProducerEpoch, txnTimeoutMs, topicPartitions.toSet,
+      txnStartTimestamp, updateTimestamp)
   }
 
   def prepareComplete(updateTimestamp: Long): TxnTransitMetadata = {
     val newState = if (state == PrepareCommit) CompleteCommit else CompleteAbort
-    prepareTransitionTo(newState, producerId, producerEpoch, txnTimeoutMs, Set.empty[TopicPartition], txnStartTimestamp,
-      updateTimestamp)
+
+    // Since the state change was successfully written to the log, unset the flag for a failed epoch fence
+    hasFailedEpochFence = false
+    prepareTransitionTo(newState, producerId, producerEpoch, lastProducerEpoch, txnTimeoutMs, Set.empty[TopicPartition],
+      txnStartTimestamp, updateTimestamp)
   }
 
   def prepareDead(): TxnTransitMetadata = {
-    prepareTransitionTo(Dead, producerId, producerEpoch, txnTimeoutMs, Set.empty[TopicPartition], txnStartTimestamp,
-      txnLastUpdateTimestamp)
+    prepareTransitionTo(Dead, producerId, producerEpoch, lastProducerEpoch, txnTimeoutMs, Set.empty[TopicPartition],
+      txnStartTimestamp, txnLastUpdateTimestamp)
   }
 
   /**
    * Check if the epochs have been exhausted for the current producerId. We do not allow the client to use an
    * epoch equal to Short.MaxValue to ensure that the coordinator will always be able to fence an existing producer.
    */
-  def isProducerEpochExhausted: Boolean = producerEpoch >= Short.MaxValue - 1
+  def isProducerEpochExhausted: Boolean = TransactionMetadata.isEpochExhausted(producerEpoch)
 
   private def hasPendingTransaction: Boolean = {
     state match {
@@ -254,6 +320,7 @@ private[transaction] class TransactionMetadata(val transactionalId: String,
   private def prepareTransitionTo(newState: TransactionState,
                                   newProducerId: Long,
                                   newEpoch: Short,
+                                  newLastEpoch: Short,
                                   newTxnTimeoutMs: Int,
                                   newTopicPartitions: immutable.Set[TopicPartition],
                                   newTxnStartTimestamp: Long,
@@ -270,7 +337,7 @@ private[transaction] class TransactionMetadata(val transactionalId: String,
 
     // check that the new state transition is valid and update the pending state if necessary
     if (TransactionMetadata.validPreviousStates(newState).contains(state)) {
-      val transitMetadata = TxnTransitMetadata(newProducerId, newEpoch, newTxnTimeoutMs, newState,
+      val transitMetadata = TxnTransitMetadata(newProducerId, producerId, newEpoch, newLastEpoch, newTxnTimeoutMs, newState,
         newTopicPartitions, newTxnStartTimestamp, updateTimestamp)
       debug(s"TransactionalId $transactionalId prepare transition from $state to $transitMetadata")
       pendingState = Some(newState)
@@ -314,14 +381,15 @@ private[transaction] class TransactionMetadata(val transactionalId: String,
           } else {
             txnTimeoutMs = transitMetadata.txnTimeoutMs
             producerEpoch = transitMetadata.producerEpoch
+            lastProducerEpoch = transitMetadata.lastProducerEpoch
             producerId = transitMetadata.producerId
+            lastProducerId = transitMetadata.lastProducerId
           }
 
         case Ongoing => // from addPartitions
           if (!validProducerEpoch(transitMetadata) ||
             !topicPartitions.subsetOf(transitMetadata.topicPartitions) ||
-            txnTimeoutMs != transitMetadata.txnTimeoutMs ||
-            txnStartTimestamp > transitMetadata.txnStartTimestamp) {
+            txnTimeoutMs != transitMetadata.txnTimeoutMs) {
 
             throwStateTransitionFailure(transitMetadata)
           } else {
@@ -411,6 +479,7 @@ private[transaction] class TransactionMetadata(val transactionalId: String,
       transactionalId == other.transactionalId &&
       producerId == other.producerId &&
       producerEpoch == other.producerEpoch &&
+      lastProducerEpoch == other.lastProducerEpoch &&
       txnTimeoutMs == other.txnTimeoutMs &&
       state.equals(other.state) &&
       topicPartitions.equals(other.topicPartitions) &&
diff --git a/core/src/main/scala/kafka/coordinator/transaction/TransactionStateManager.scala b/core/src/main/scala/kafka/coordinator/transaction/TransactionStateManager.scala
index da61077d3d201..2882d8630249c 100644
--- a/core/src/main/scala/kafka/coordinator/transaction/TransactionStateManager.scala
+++ b/core/src/main/scala/kafka/coordinator/transaction/TransactionStateManager.scala
@@ -22,33 +22,36 @@ import java.util.concurrent.TimeUnit
 import java.util.concurrent.atomic.AtomicBoolean
 import java.util.concurrent.locks.ReentrantReadWriteLock
 
-import kafka.common.KafkaException
-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.utils.Implicits._
 import kafka.zk.KafkaZkClient
-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.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.IsolationLevel
 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.jdk.CollectionConverters._
 import scala.collection.mutable
-import scala.collection.JavaConverters._
 
 
 object TransactionStateManager {
   // default transaction management config values
   val DefaultTransactionsMaxTimeoutMs: Int = TimeUnit.MINUTES.toMillis(15).toInt
   val DefaultTransactionalIdExpirationMs: Int = TimeUnit.DAYS.toMillis(7).toInt
-  val DefaultAbortTimedOutTransactionsIntervalMs: Int = TimeUnit.MINUTES.toMillis(1).toInt
+  val DefaultAbortTimedOutTransactionsIntervalMs: Int = TimeUnit.SECONDS.toMillis(10).toInt
   val DefaultRemoveExpiredTransactionalIdsIntervalMs: Int = TimeUnit.HOURS.toMillis(1).toInt
+
+  val MetricsGroup: String = "transaction-coordinator-metrics"
+  val LoadTimeSensor: String = "TransactionsPartitionLoadTime"
 }
 
 /**
@@ -59,9 +62,9 @@ object TransactionStateManager {
  * 3. the background expiration of the transaction as well as the transactional id.
  *
  * Delayed operation locking notes:
- * Delayed operations in TransactionStateManager use `stateLock.readLock` as the delayed operation
- * lock. Delayed operations are completed only if `stateLock.readLock` can be acquired.
- * Delayed callbacks may acquire `stateLock.readLock` or any of the `txnMetadata` locks.
+ * Delayed operations in TransactionStateManager use individual operation locks.
+ * Delayed callbacks may acquire `stateLock.readLock` or any of the `txnMetadata` locks,
+ * but we always require that `stateLock.readLock` be acquired first. In particular:
  * 
              
  * 
            • `stateLock.readLock` must never be acquired while holding `txnMetadata` lock.
              • 
  * 
            • `txnMetadata` lock must never be acquired while holding `stateLock.writeLock`.
              • 
@@ -73,11 +76,12 @@ 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 + "]: "
 
-  type SendTxnMarkersCallback = (String, Int, TransactionResult, TransactionMetadata, TxnTransitMetadata) => Unit
+  type SendTxnMarkersCallback = (Int, TransactionResult, TransactionMetadata, TxnTransitMetadata) => Unit
 
   /** shutting down flag */
   private val shuttingDown = new AtomicBoolean(false)
@@ -86,27 +90,31 @@ 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(TransactionStateManager.LoadTimeSensor)
+
+  partitionLoadSensor.add(metrics.metricName("partition-load-time-max",
+    TransactionStateManager.MetricsGroup,
+    "The max time it took to load the partitions in the last 30sec"), new Max())
+  partitionLoadSensor.add(metrics.metricName("partition-load-time-avg",
+    TransactionStateManager.MetricsGroup,
+    "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)
     }
   }
-  private[transaction] def stateReadLock = stateLock.readLock
 
   // this is best-effort expiration of an ongoing transaction which has been open for more than its
   // txn timeout value, we do not need to grab the lock on the metadata object upon checking its state
@@ -115,9 +123,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
@@ -135,12 +141,12 @@ class TransactionStateManager(brokerId: Int,
     }
   }
 
-  def enableTransactionalIdExpiration() {
+  def enableTransactionalIdExpiration(): Unit = {
     scheduler.schedule("transactionalId-expiration", () => {
       val now = time.milliseconds()
       inReadLock(stateLock) {
         val transactionalIdByPartition: Map[Int, mutable.Iterable[TransactionalIdCoordinatorEpochAndMetadata]] =
-          transactionMetadataCache.flatMap { case (partition, entry) =>
+          transactionMetadataCache.flatMap { case (_, entry) =>
             entry.metadataPerTransactionalId.filter { case (_, txnMetadata) => txnMetadata.state match {
               case Empty | CompleteCommit | CompleteAbort => true
               case _ => false
@@ -167,37 +173,32 @@ class TransactionStateManager(brokerId: Int,
             (topicPartition, records)
           }
 
-
         def removeFromCacheCallback(responses: collection.Map[TopicPartition, PartitionResponse]): Unit = {
-          responses.foreach { case (topicPartition, response) =>
-            response.error match {
-              case Errors.NONE =>
-                inReadLock(stateLock) {
-                  val toRemove = transactionalIdByPartition(topicPartition.partition())
-                  transactionMetadataCache.get(topicPartition.partition)
-                    .foreach { txnMetadataCacheEntry =>
-                      toRemove.foreach { idCoordinatorEpochAndMetadata =>
-                        val txnMetadata = txnMetadataCacheEntry.metadataPerTransactionalId.get(idCoordinatorEpochAndMetadata.transactionalId)
-                        txnMetadata.inLock {
-                          if (txnMetadataCacheEntry.coordinatorEpoch == idCoordinatorEpochAndMetadata.coordinatorEpoch
-                            && txnMetadata.pendingState.contains(Dead)
-                            && txnMetadata.producerEpoch == idCoordinatorEpochAndMetadata.transitMetadata.producerEpoch
-                          )
-                            txnMetadataCacheEntry.metadataPerTransactionalId.remove(idCoordinatorEpochAndMetadata.transactionalId)
-                          else {
-                             debug(s"failed to remove expired transactionalId: ${idCoordinatorEpochAndMetadata.transactionalId}" +
-                               s" from cache. pendingState: ${txnMetadata.pendingState} producerEpoch: ${txnMetadata.producerEpoch}" +
-                               s" expected producerEpoch: ${idCoordinatorEpochAndMetadata.transitMetadata.producerEpoch}" +
-                               s" coordinatorEpoch: ${txnMetadataCacheEntry.coordinatorEpoch} expected coordinatorEpoch: " +
-                               s"${idCoordinatorEpochAndMetadata.coordinatorEpoch}")
-                            txnMetadata.pendingState = None
-                          }
-                        }
-                      }
+          responses.forKeyValue { (topicPartition, response) =>
+            inReadLock(stateLock) {
+              val toRemove = transactionalIdByPartition(topicPartition.partition)
+              transactionMetadataCache.get(topicPartition.partition).foreach { txnMetadataCacheEntry =>
+                toRemove.foreach { idCoordinatorEpochAndMetadata =>
+                  val transactionalId = idCoordinatorEpochAndMetadata.transactionalId
+                  val txnMetadata = txnMetadataCacheEntry.metadataPerTransactionalId.get(transactionalId)
+                  txnMetadata.inLock {
+                    if (txnMetadataCacheEntry.coordinatorEpoch == idCoordinatorEpochAndMetadata.coordinatorEpoch
+                      && txnMetadata.pendingState.contains(Dead)
+                      && txnMetadata.producerEpoch == idCoordinatorEpochAndMetadata.transitMetadata.producerEpoch
+                      && response.error == Errors.NONE) {
+                      txnMetadataCacheEntry.metadataPerTransactionalId.remove(transactionalId)
+                    } else {
+                      warn(s"Failed to remove expired transactionalId: $transactionalId" +
+                        s" from cache. Tombstone append error code: ${response.error}," +
+                        s" pendingState: ${txnMetadata.pendingState}, producerEpoch: ${txnMetadata.producerEpoch}," +
+                        s" expected producerEpoch: ${idCoordinatorEpochAndMetadata.transitMetadata.producerEpoch}," +
+                        s" coordinatorEpoch: ${txnMetadataCacheEntry.coordinatorEpoch}, expected coordinatorEpoch: " +
+                        s"${idCoordinatorEpochAndMetadata.coordinatorEpoch}")
+                      txnMetadata.pendingState = None
                     }
+                  }
                 }
-              case _ =>
-                debug(s"writing transactionalId tombstones for partition: ${topicPartition.partition} failed with error: ${response.error.message()}")
+              }
             }
           }
         }
@@ -206,23 +207,22 @@ class TransactionStateManager(brokerId: Int,
           config.requestTimeoutMs,
           TransactionLog.EnforcedRequiredAcks,
           internalTopicsAllowed = true,
-          isFromClient = false,
+          origin = AppendOrigin.Coordinator,
           recordsPerPartition,
-          removeFromCacheCallback,
-          Some(stateLock.readLock)
-        )
+          removeFromCacheCallback)
       }
 
     }, delay = config.removeExpiredTransactionalIdsIntervalMs, period = config.removeExpiredTransactionalIdsIntervalMs)
   }
 
-  def getTransactionState(transactionalId: String): Either[Errors, Option[CoordinatorEpochAndTxnMetadata]] =
+  def getTransactionState(transactionalId: String): Either[Errors, Option[CoordinatorEpochAndTxnMetadata]] = {
     getAndMaybeAddTransactionState(transactionalId, None)
+  }
 
-  def putTransactionStateIfNotExists(transactionalId: String,
-                                     txnMetadata: TransactionMetadata): Either[Errors, CoordinatorEpochAndTxnMetadata] =
-    getAndMaybeAddTransactionState(transactionalId, Some(txnMetadata))
-      .right.map(_.getOrElse(throw new IllegalStateException(s"Unexpected empty transaction metadata returned while putting $txnMetadata")))
+  def putTransactionStateIfNotExists(txnMetadata: TransactionMetadata): Either[Errors, CoordinatorEpochAndTxnMetadata] = {
+    getAndMaybeAddTransactionState(txnMetadata.transactionalId, Some(txnMetadata)).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
@@ -237,8 +237,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) =>
@@ -289,51 +287,69 @@ class TransactionStateManager(brokerId: Int,
   private def loadTransactionMetadata(topicPartition: TopicPartition, coordinatorEpoch: Int): Pool[String, TransactionMetadata] =  {
     def logEndOffset = replicaManager.getLogEndOffset(topicPartition).getOrElse(-1L)
 
-    val startMs = time.milliseconds()
     val loadedTransactions = new Pool[String, TransactionMetadata]
 
     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) =>
-        lazy val buffer = ByteBuffer.allocate(config.transactionLogLoadBufferSize)
+        // 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, isolationLevel = IsolationLevel.READ_UNCOMMITTED)
-            val memRecords = fetchDataInfo.records match {
+            val fetchDataInfo = log.read(currOffset,
+              maxLength = config.transactionLogLoadBufferSize,
+              isolation = FetchLogEnd,
+              minOneMessage = true)
+
+            readAtLeastOneRecord = fetchDataInfo.records.sizeInBytes > 0
+
+            val memRecords = (fetchDataInfo.records: @unchecked) match {
               case records: MemoryRecords => records
               case fileRecords: FileRecords =>
+                val sizeInBytes = fileRecords.sizeInBytes
+                val bytesNeeded = Math.max(config.transactionLogLoadBufferSize, sizeInBytes)
+
+                // 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 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)
+                } else {
+                  buffer.clear()
+                }
                 buffer.clear()
-                val bufferRead = fileRecords.readInto(buffer, 0)
-                MemoryRecords.readableRecords(bufferRead)
+                fileRecords.readInto(buffer, 0)
+                MemoryRecords.readableRecords(buffer)
             }
 
-            memRecords.batches.asScala.foreach { batch =>
+            memRecords.batches.forEach { batch =>
               for (record <- batch.asScala) {
                 require(record.hasKey, "Transaction state log's key should not be null")
                 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")
           }
         } catch {
           case t: Throwable => error(s"Error loading transactions from transaction log $topicPartition", t)
@@ -345,42 +361,43 @@ 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(startTimeMs: java.lang.Long): Unit = {
+      val schedulerTimeMs = time.milliseconds() - startTimeMs
+      info(s"Loading transaction metadata from $topicPartition at epoch $coordinatorEpoch")
+      validateTransactionTopicPartitionCountIsStable()
+
       val loadedTransactions = loadTransactionMetadata(topicPartition, coordinatorEpoch)
+      val endTimeMs = time.milliseconds()
+      val totalLoadingTimeMs = endTimeMs - startTimeMs
+      partitionLoadSensor.record(totalLoadingTimeMs.toDouble, endTimeMs, false)
+      info(s"Finished loading ${loadedTransactions.size} transaction metadata from $topicPartition in " +
+        s"$totalLoadingTimeMs milliseconds, of which $schedulerTimeMs milliseconds was spent in the scheduler.")
 
       inWriteLock(stateLock) {
         if (loadingPartitions.contains(partitionAndLeaderEpoch)) {
@@ -399,7 +416,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
                 }
               }
           }
@@ -410,48 +427,48 @@ class TransactionStateManager(brokerId: Int,
           loadingPartitions.remove(partitionAndLeaderEpoch)
 
           transactionsPendingForCompletion.foreach { txnTransitMetadata =>
-            sendTxnMarkers(txnTransitMetadata.transactionalId, txnTransitMetadata.coordinatorEpoch, txnTransitMetadata.result, txnTransitMetadata.txnMetadata, txnTransitMetadata.transitMetadata)
+            sendTxnMarkers(txnTransitMetadata.coordinatorEpoch, txnTransitMetadata.result,
+              txnTransitMetadata.txnMetadata, txnTransitMetadata.transitMetadata)
           }
         }
       }
+
+      info(s"Completed loading transaction metadata from $topicPartition for coordinator epoch $coordinatorEpoch")
     }
 
-    scheduler.schedule(s"load-txns-for-partition-$topicPartition", loadTransactions)
+    val scheduleStartMs = time.milliseconds()
+    scheduler.schedule(s"load-txns-for-partition-$topicPartition", () => loadTransactions(scheduleStartMs))
+  }
+
+  def removeTransactionsForTxnTopicPartition(partitionId: Int): Unit = {
+    val topicPartition = new TopicPartition(Topic.TRANSACTION_STATE_TOPIC_NAME, partitionId)
+    inWriteLock(stateLock) {
+      loadingPartitions --= loadingPartitions.filter(_.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 = {
@@ -497,7 +514,7 @@ class TransactionStateManager(brokerId: Int,
                | Errors.REQUEST_TIMED_OUT => // note that for timed out request we return NOT_AVAILABLE error code to let client retry
             Errors.COORDINATOR_NOT_AVAILABLE
 
-          case Errors.NOT_LEADER_FOR_PARTITION
+          case Errors.NOT_LEADER_OR_FOLLOWER
                | Errors.KAFKA_STORAGE_ERROR =>
             Errors.NOT_COORDINATOR
 
@@ -615,10 +632,9 @@ class TransactionStateManager(brokerId: Int,
                 newMetadata.txnTimeoutMs.toLong,
                 TransactionLog.EnforcedRequiredAcks,
                 internalTopicsAllowed = true,
-                isFromClient = false,
+                origin = AppendOrigin.Coordinator,
                 recordsPerPartition,
-                updateCacheCallback,
-                delayedProduceLock = Some(stateLock.readLock))
+                updateCacheCallback)
 
               trace(s"Appending new metadata $newMetadata for transaction id $transactionalId with coordinator epoch $coordinatorEpoch to the local transaction log")
           }
@@ -626,7 +642,7 @@ class TransactionStateManager(brokerId: Int,
     }
   }
 
-  def shutdown() {
+  def shutdown(): Unit = {
     shuttingDown.set(true)
     loadingPartitions.clear()
     transactionMetadataCache.clear()
@@ -636,9 +652,15 @@ 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)
+private[transaction] case class CoordinatorEpochAndTxnMetadata(coordinatorEpoch: Int,
+                                                               transactionMetadata: TransactionMetadata)
 
 private[transaction] case class TransactionConfig(transactionalIdExpirationMs: Int = TransactionStateManager.DefaultTransactionalIdExpirationMs,
                                                   transactionMaxTimeoutMs: Int = TransactionStateManager.DefaultTransactionsMaxTimeoutMs,
@@ -651,7 +673,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/javaapi/FetchRequest.scala b/core/src/main/scala/kafka/javaapi/FetchRequest.scala
deleted file mode 100644
index fe8beaac4dc49..0000000000000
--- a/core/src/main/scala/kafka/javaapi/FetchRequest.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.javaapi
-
-import java.util
-
-import kafka.common.TopicAndPartition
-import kafka.api.{PartitionFetchInfo, Request}
-
-import scala.collection.JavaConverters._
-
-object FetchRequest {
-  private def seqToLinkedHashMap[K, V](s: Seq[(K, V)]): util.LinkedHashMap[K, V] = {
-    val map = new util.LinkedHashMap[K, V]
-    s.foreach { case (k, v) => map.put(k, v) }
-    map
-  }
-}
-
-class FetchRequest(correlationId: Int,
-                   clientId: String,
-                   maxWait: Int,
-                   minBytes: Int,
-                   requestInfo: util.LinkedHashMap[TopicAndPartition, PartitionFetchInfo]) {
-
-  @deprecated("The order of partitions in `requestInfo` is relevant, so this constructor is deprecated in favour of the " +
-    "one that takes a LinkedHashMap", since = "0.10.1.0")
-  def this(correlationId: Int, clientId: String, maxWait: Int, minBytes: Int,
-    requestInfo: java.util.Map[TopicAndPartition, PartitionFetchInfo]) {
-    this(correlationId, clientId, maxWait, minBytes,
-      FetchRequest.seqToLinkedHashMap(kafka.api.FetchRequest.shuffle(requestInfo.asScala.toSeq)))
-  }
-
-  val underlying = kafka.api.FetchRequest(
-    correlationId = correlationId,
-    clientId = clientId,
-    replicaId = Request.OrdinaryConsumerId,
-    maxWait = maxWait,
-    minBytes = minBytes,
-    requestInfo = requestInfo.asScala.toBuffer
-  )
-
-  override def toString = underlying.toString
-
-  override def equals(obj: Any): Boolean = {
-    obj match {
-      case null => false
-      case other: FetchRequest => this.underlying.equals(other.underlying)
-      case _ => false
-    }
-  }
-
-  override def hashCode = underlying.hashCode
-}
-
diff --git a/core/src/main/scala/kafka/javaapi/FetchResponse.scala b/core/src/main/scala/kafka/javaapi/FetchResponse.scala
deleted file mode 100644
index c9165558014f8..0000000000000
--- a/core/src/main/scala/kafka/javaapi/FetchResponse.scala
+++ /dev/null
@@ -1,44 +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.javaapi
-
-class FetchResponse(private val underlying: kafka.api.FetchResponse) {
-
-  def messageSet(topic: String, partition: Int): kafka.javaapi.message.ByteBufferMessageSet = {
-    import Implicits._
-    underlying.messageSet(topic, partition)
-  }
-
-  def highWatermark(topic: String, partition: Int) = underlying.highWatermark(topic, partition)
-
-  def hasError = underlying.hasError
-
-  def error(topic: String, partition: Int) = underlying.error(topic, partition)
-
-  def errorCode(topic: String, partition: Int) = error(topic, partition).code
-
-  override def equals(obj: Any): Boolean = {
-    obj match {
-      case null => false
-      case other: FetchResponse => this.underlying.equals(other.underlying)
-      case _ => false
-    }
-  }
-
-  override def hashCode = underlying.hashCode
-}
diff --git a/core/src/main/scala/kafka/javaapi/GroupCoordinatorResponse.scala b/core/src/main/scala/kafka/javaapi/GroupCoordinatorResponse.scala
deleted file mode 100644
index 096941c0c07a8..0000000000000
--- a/core/src/main/scala/kafka/javaapi/GroupCoordinatorResponse.scala
+++ /dev/null
@@ -1,50 +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.javaapi
-
-import java.nio.ByteBuffer
-import kafka.cluster.BrokerEndPoint
-
-class GroupCoordinatorResponse(private val underlying: kafka.api.GroupCoordinatorResponse) {
-
-  def error = underlying.error
-
-  def errorCode = error.code
-
-  def coordinator: BrokerEndPoint = {
-    import kafka.javaapi.Implicits._
-    underlying.coordinatorOpt
-  }
-
-  override def equals(obj: Any): Boolean = {
-    obj match {
-      case null => false
-      case other: GroupCoordinatorResponse => this.underlying.equals(other.underlying)
-      case _ => false
-    }
-  }
-
-  override def hashCode = underlying.hashCode
-
-  override def toString = underlying.toString
-
-}
-
-object GroupCoordinatorResponse {
-  def readFrom(buffer: ByteBuffer) = new GroupCoordinatorResponse(kafka.api.GroupCoordinatorResponse.readFrom(buffer))
-}
diff --git a/core/src/main/scala/kafka/javaapi/Implicits.scala b/core/src/main/scala/kafka/javaapi/Implicits.scala
deleted file mode 100644
index c69b0a3c3d87e..0000000000000
--- a/core/src/main/scala/kafka/javaapi/Implicits.scala
+++ /dev/null
@@ -1,50 +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.javaapi
-
-import kafka.utils.Logging
-
-private[javaapi] object Implicits extends Logging {
-
-  implicit def scalaMessageSetToJavaMessageSet(messageSet: kafka.message.ByteBufferMessageSet):
-     kafka.javaapi.message.ByteBufferMessageSet = {
-    new kafka.javaapi.message.ByteBufferMessageSet(messageSet.buffer)
-  }
-
-  implicit def toJavaFetchResponse(response: kafka.api.FetchResponse): kafka.javaapi.FetchResponse =
-    new kafka.javaapi.FetchResponse(response)
-
-  implicit def toJavaTopicMetadataResponse(response: kafka.api.TopicMetadataResponse): kafka.javaapi.TopicMetadataResponse =
-    new kafka.javaapi.TopicMetadataResponse(response)
-
-  implicit def toJavaOffsetResponse(response: kafka.api.OffsetResponse): kafka.javaapi.OffsetResponse =
-    new kafka.javaapi.OffsetResponse(response)
-
-  implicit def toJavaOffsetFetchResponse(response: kafka.api.OffsetFetchResponse): kafka.javaapi.OffsetFetchResponse =
-    new kafka.javaapi.OffsetFetchResponse(response)
-
-  implicit def toJavaOffsetCommitResponse(response: kafka.api.OffsetCommitResponse): kafka.javaapi.OffsetCommitResponse =
-    new kafka.javaapi.OffsetCommitResponse(response)
-
-  implicit def optionToJavaRef[T](opt: Option[T]): T = {
-    opt match {
-      case Some(obj) => obj
-      case None => null.asInstanceOf[T]
-    }
-  }
-
-}
diff --git a/core/src/main/scala/kafka/javaapi/OffsetCommitRequest.scala b/core/src/main/scala/kafka/javaapi/OffsetCommitRequest.scala
deleted file mode 100644
index 0c3c6516b07ae..0000000000000
--- a/core/src/main/scala/kafka/javaapi/OffsetCommitRequest.scala
+++ /dev/null
@@ -1,59 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package kafka.javaapi
-
-import kafka.common.{OffsetAndMetadata, TopicAndPartition}
-import scala.collection.JavaConverters._
-
-class OffsetCommitRequest(groupId: String,
-                          requestInfo: java.util.Map[TopicAndPartition, OffsetAndMetadata],
-                          correlationId: Int,
-                          clientId: String,
-                          versionId: Short) {
-  val underlying = {
-    val scalaMap: collection.immutable.Map[TopicAndPartition, OffsetAndMetadata] = requestInfo.asScala.toMap
-    kafka.api.OffsetCommitRequest(
-      groupId = groupId,
-      requestInfo = scalaMap,
-      versionId = versionId,
-      correlationId = correlationId,
-      clientId = clientId
-    )
-  }
-
-  def this(groupId: String,
-           requestInfo: java.util.Map[TopicAndPartition, OffsetAndMetadata],
-           correlationId: Int,
-           clientId: String) {
-
-    // by default bind to version 0 so that it commits to Zookeeper
-    this(groupId, requestInfo, correlationId, clientId, 0)
-  }
-
-  override def toString = underlying.toString
-
-  override def equals(obj: Any): Boolean = {
-    obj match {
-      case null => false
-      case other: OffsetCommitRequest => this.underlying.equals(other.underlying)
-      case _ => false
-    }
-  }
-
-  override def hashCode = underlying.hashCode
-}
diff --git a/core/src/main/scala/kafka/javaapi/OffsetCommitResponse.scala b/core/src/main/scala/kafka/javaapi/OffsetCommitResponse.scala
deleted file mode 100644
index c348eba22529e..0000000000000
--- a/core/src/main/scala/kafka/javaapi/OffsetCommitResponse.scala
+++ /dev/null
@@ -1,39 +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.javaapi
-
-import java.nio.ByteBuffer
-
-import kafka.common.TopicAndPartition
-import org.apache.kafka.common.protocol.Errors
-import scala.collection.JavaConverters._
-
-class OffsetCommitResponse(private val underlying: kafka.api.OffsetCommitResponse) {
-
-  def errors: java.util.Map[TopicAndPartition, Errors] = underlying.commitStatus.asJava
-
-  def hasError = underlying.hasError
-
-  def error(topicAndPartition: TopicAndPartition) = underlying.commitStatus(topicAndPartition)
-
-  def errorCode(topicAndPartition: TopicAndPartition) = error(topicAndPartition).code
-}
-
-object OffsetCommitResponse {
-  def readFrom(buffer: ByteBuffer) = new OffsetCommitResponse(kafka.api.OffsetCommitResponse.readFrom(buffer))
-}
diff --git a/core/src/main/scala/kafka/javaapi/OffsetFetchRequest.scala b/core/src/main/scala/kafka/javaapi/OffsetFetchRequest.scala
deleted file mode 100644
index 5f96439fba534..0000000000000
--- a/core/src/main/scala/kafka/javaapi/OffsetFetchRequest.scala
+++ /dev/null
@@ -1,60 +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.javaapi
-
-import kafka.common.TopicAndPartition
-import collection.JavaConverters._
-
-class OffsetFetchRequest(groupId: String,
-                         requestInfo: java.util.List[TopicAndPartition],
-                         versionId: Short,
-                         correlationId: Int,
-                         clientId: String) {
-
-  def this(groupId: String,
-           requestInfo: java.util.List[TopicAndPartition],
-           correlationId: Int,
-           clientId: String) {
-    // by default bind to version 0 so that it fetches from ZooKeeper
-    this(groupId, requestInfo, 0, correlationId, clientId)
-  }
-
-  val underlying = {
-    kafka.api.OffsetFetchRequest(
-      groupId = groupId,
-      requestInfo = requestInfo.asScala,
-      versionId = versionId,
-      correlationId = correlationId,
-      clientId = clientId
-    )
-  }
-
-  override def toString = underlying.toString
-
-  override def equals(obj: Any): Boolean = {
-    obj match {
-      case null => false
-      case other: OffsetFetchRequest => this.underlying.equals(other.underlying)
-      case _ => false
-    }
-  }
-
-  override def hashCode = underlying.hashCode
-}
-
-
diff --git a/core/src/main/scala/kafka/javaapi/OffsetFetchResponse.scala b/core/src/main/scala/kafka/javaapi/OffsetFetchResponse.scala
deleted file mode 100644
index 01aa8e8e88f14..0000000000000
--- a/core/src/main/scala/kafka/javaapi/OffsetFetchResponse.scala
+++ /dev/null
@@ -1,33 +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.javaapi
-
-import java.nio.ByteBuffer
-
-import kafka.common.{TopicAndPartition, OffsetMetadataAndError}
-import collection.JavaConverters._
-
-class OffsetFetchResponse(private val underlying: kafka.api.OffsetFetchResponse) {
-
-  def offsets: java.util.Map[TopicAndPartition, OffsetMetadataAndError] = underlying.requestInfo.asJava
-
-}
-
-object OffsetFetchResponse {
-  def readFrom(buffer: ByteBuffer) = new OffsetFetchResponse(kafka.api.OffsetFetchResponse.readFrom(buffer))
-}
diff --git a/core/src/main/scala/kafka/javaapi/OffsetRequest.scala b/core/src/main/scala/kafka/javaapi/OffsetRequest.scala
deleted file mode 100644
index 96b66ef6ee615..0000000000000
--- a/core/src/main/scala/kafka/javaapi/OffsetRequest.scala
+++ /dev/null
@@ -1,50 +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.javaapi
-
-import kafka.common.TopicAndPartition
-import kafka.api.{Request, PartitionOffsetRequestInfo}
-import scala.collection.JavaConverters._
-
-
-class OffsetRequest(requestInfo: java.util.Map[TopicAndPartition, PartitionOffsetRequestInfo],
-                    versionId: Short,
-                    clientId: String) {
-
-  val underlying = {
-    val scalaMap = requestInfo.asScala.toMap
-    kafka.api.OffsetRequest(
-      requestInfo = scalaMap,
-      versionId = versionId,
-      clientId = clientId,
-      replicaId = Request.OrdinaryConsumerId
-    )
-  }
-
-  override def toString = underlying.toString
-
-  override def equals(obj: Any): Boolean = {
-    obj match {
-      case null => false
-      case other: OffsetRequest => this.underlying.equals(other.underlying)
-      case _ => false
-    }
-  }
-
-  override def hashCode = underlying.hashCode
-}
diff --git a/core/src/main/scala/kafka/javaapi/OffsetResponse.scala b/core/src/main/scala/kafka/javaapi/OffsetResponse.scala
deleted file mode 100644
index cb2047f43e17e..0000000000000
--- a/core/src/main/scala/kafka/javaapi/OffsetResponse.scala
+++ /dev/null
@@ -1,45 +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.javaapi
-
-import kafka.common.TopicAndPartition
-
-class OffsetResponse(private val underlying: kafka.api.OffsetResponse) {
-
-  def hasError = underlying.hasError
-
-  def error(topic: String, partition: Int) =
-    underlying.partitionErrorAndOffsets(TopicAndPartition(topic, partition)).error
-
-  def errorCode(topic: String, partition: Int) = error(topic, partition).code
-
-  def offsets(topic: String, partition: Int) =
-    underlying.partitionErrorAndOffsets(TopicAndPartition(topic, partition)).offsets.toArray
-
-  override def equals(obj: Any): Boolean = {
-    obj match {
-      case null => false
-      case other: OffsetResponse => this.underlying.equals(other.underlying)
-      case _ => false
-    }
-  }
-
-  override def hashCode = underlying.hashCode
-
-  override def toString = underlying.toString
-}
diff --git a/core/src/main/scala/kafka/javaapi/TopicMetadata.scala b/core/src/main/scala/kafka/javaapi/TopicMetadata.scala
deleted file mode 100644
index 051445ceba547..0000000000000
--- a/core/src/main/scala/kafka/javaapi/TopicMetadata.scala
+++ /dev/null
@@ -1,67 +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.javaapi
-
-import kafka.cluster.BrokerEndPoint
-import scala.collection.JavaConverters._
-
-private[javaapi] object MetadataListImplicits {
-  implicit def toJavaTopicMetadataList(topicMetadataSeq: Seq[kafka.api.TopicMetadata]):
-  java.util.List[kafka.javaapi.TopicMetadata] = topicMetadataSeq.map(new kafka.javaapi.TopicMetadata(_)).asJava
-
-  implicit def toPartitionMetadataList(partitionMetadataSeq: Seq[kafka.api.PartitionMetadata]):
-  java.util.List[kafka.javaapi.PartitionMetadata] = partitionMetadataSeq.map(new kafka.javaapi.PartitionMetadata(_)).asJava
-}
-
-class TopicMetadata(private val underlying: kafka.api.TopicMetadata) {
-  def topic: String = underlying.topic
-
-  def partitionsMetadata: java.util.List[PartitionMetadata] = {
-    import kafka.javaapi.MetadataListImplicits._
-    underlying.partitionsMetadata
-  }
-
-  def error = underlying.error
-
-  def errorCode = error.code
-
-  def sizeInBytes: Int = underlying.sizeInBytes
-
-  override def toString = underlying.toString
-}
-
-
-class PartitionMetadata(private val underlying: kafka.api.PartitionMetadata) {
-  def partitionId: Int = underlying.partitionId
-
-  def leader: BrokerEndPoint = {
-    import kafka.javaapi.Implicits._
-    underlying.leader
-  }
-
-  def replicas: java.util.List[BrokerEndPoint] = underlying.replicas.asJava
-
-  def isr: java.util.List[BrokerEndPoint] = underlying.isr.asJava
-
-  def error = underlying.error
-
-  def errorCode = error.code
-
-  def sizeInBytes: Int = underlying.sizeInBytes
-
-  override def toString = underlying.toString
-}
diff --git a/core/src/main/scala/kafka/javaapi/TopicMetadataRequest.scala b/core/src/main/scala/kafka/javaapi/TopicMetadataRequest.scala
deleted file mode 100644
index fdb14cbaeaff0..0000000000000
--- a/core/src/main/scala/kafka/javaapi/TopicMetadataRequest.scala
+++ /dev/null
@@ -1,66 +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.javaapi
-
-import java.nio.ByteBuffer
-
-import kafka.api._
-import org.apache.kafka.common.protocol.ApiKeys
-
-import scala.collection.JavaConverters._
-
-class TopicMetadataRequest(val versionId: Short,
-                           val correlationId: Int,
-                           val clientId: String,
-                           val topics: java.util.List[String])
-    extends RequestOrResponse(Some(ApiKeys.METADATA.id)) {
-
-  val underlying: kafka.api.TopicMetadataRequest = new kafka.api.TopicMetadataRequest(versionId, correlationId, clientId, topics.asScala)
-
-  def this(topics: java.util.List[String]) =
-    this(kafka.api.TopicMetadataRequest.CurrentVersion, 0, kafka.api.TopicMetadataRequest.DefaultClientId, topics)
-
-  def this(topics: java.util.List[String], correlationId: Int) =
-    this(kafka.api.TopicMetadataRequest.CurrentVersion, correlationId, kafka.api.TopicMetadataRequest.DefaultClientId, topics)
-
-  def writeTo(buffer: ByteBuffer) = underlying.writeTo(buffer)
-
-  def sizeInBytes: Int = underlying.sizeInBytes
-
-  override def toString: String = {
-    describe(true)
-  }
-
-  override def describe(details: Boolean): String = {
-    val topicMetadataRequest = new StringBuilder
-    topicMetadataRequest.append("Name: " + this.getClass.getSimpleName)
-    topicMetadataRequest.append("; Version: " + versionId)
-    topicMetadataRequest.append("; CorrelationId: " + correlationId)
-    topicMetadataRequest.append("; ClientId: " + clientId)
-    if(details) {
-      topicMetadataRequest.append("; Topics: ")
-      val topicIterator = topics.iterator()
-      while (topicIterator.hasNext) {
-        val topic = topicIterator.next()
-        topicMetadataRequest.append("%s".format(topic))
-        if(topicIterator.hasNext)
-          topicMetadataRequest.append(",")
-      }
-    }
-    topicMetadataRequest.toString()
-  }
-}
diff --git a/core/src/main/scala/kafka/javaapi/TopicMetadataResponse.scala b/core/src/main/scala/kafka/javaapi/TopicMetadataResponse.scala
deleted file mode 100644
index 40f81d553611a..0000000000000
--- a/core/src/main/scala/kafka/javaapi/TopicMetadataResponse.scala
+++ /dev/null
@@ -1,38 +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.javaapi
-
-class TopicMetadataResponse(private val underlying: kafka.api.TopicMetadataResponse) {
-  def sizeInBytes: Int = underlying.sizeInBytes
-
-  def topicsMetadata: java.util.List[kafka.javaapi.TopicMetadata] = {
-    import kafka.javaapi.MetadataListImplicits._
-    underlying.topicsMetadata
-  }
-
-  override def equals(obj: Any): Boolean = {
-    obj match {
-      case null => false
-      case other: TopicMetadataResponse => this.underlying.equals(other.underlying)
-      case _ => false
-    }
-  }
-
-  override def hashCode = underlying.hashCode
-
-  override def toString = underlying.toString
-}
diff --git a/core/src/main/scala/kafka/javaapi/consumer/ConsumerConnector.java b/core/src/main/scala/kafka/javaapi/consumer/ConsumerConnector.java
deleted file mode 100644
index def977cfaf64a..0000000000000
--- a/core/src/main/scala/kafka/javaapi/consumer/ConsumerConnector.java
+++ /dev/null
@@ -1,91 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements. See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License. You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package kafka.javaapi.consumer;
-
-import kafka.common.OffsetAndMetadata;
-import kafka.common.TopicAndPartition;
-import kafka.consumer.KafkaStream;
-import kafka.consumer.TopicFilter;
-import kafka.serializer.Decoder;
-
-import java.util.List;
-import java.util.Map;
-
-/**
- * @deprecated since 0.11.0.0, this interface will be removed in a future release.
- */
-@Deprecated
-public interface ConsumerConnector {
-    /**
-     *  Create a list of MessageStreams of type T for each topic.
-     *
-     *  @param topicCountMap  a map of (topic, #streams) pair
-     *  @param keyDecoder a decoder that decodes the message key
-     *  @param valueDecoder a decoder that decodes the message itself
-     *  @return a map of (topic, list of  KafkaStream) pairs.
-     *          The number of items in the list is #streams. Each stream supports
-     *          an iterator over message/metadata pairs.
-     */
-    public  Map>>
-        createMessageStreams(Map topicCountMap, Decoder keyDecoder, Decoder valueDecoder);
-
-    public Map>> createMessageStreams(Map topicCountMap);
-
-    /**
-     *  Create a list of MessageAndTopicStreams containing messages of type T.
-     *
-     *  @param topicFilter a TopicFilter that specifies which topics to
-     *                    subscribe to (encapsulates a whitelist or a blacklist).
-     *  @param numStreams the number of message streams to return.
-     *  @param keyDecoder a decoder that decodes the message key
-     *  @param valueDecoder a decoder that decodes the message itself
-     *  @return a list of KafkaStream. Each stream supports an
-     *          iterator over its MessageAndMetadata elements.
-     */
-    public  List>
-        createMessageStreamsByFilter(TopicFilter topicFilter, int numStreams, Decoder keyDecoder, Decoder valueDecoder);
-
-    public List> createMessageStreamsByFilter(TopicFilter topicFilter, int numStreams);
-
-    public List> createMessageStreamsByFilter(TopicFilter topicFilter);
-
-    /**
-     *  Commit the offsets of all broker partitions connected by this connector.
-     */
-    public void commitOffsets();
-
-    public void commitOffsets(boolean retryOnFailure);
-
-    /**
-     *  Commit offsets using the provided offsets map
-     *
-     *  @param offsetsToCommit a map containing the offset to commit for each partition.
-     *  @param retryOnFailure enable retries on the offset commit if it fails.
-     */
-    public void commitOffsets(Map offsetsToCommit, boolean retryOnFailure);
-
-    /**
-     * Wire in a consumer rebalance listener to be executed when consumer rebalance occurs.
-     * @param listener The consumer rebalance listener to wire in
-     */
-    public void setConsumerRebalanceListener(ConsumerRebalanceListener listener);
-
-    /**
-     *  Shut down the connector
-     */
-    public void shutdown();
-}
diff --git a/core/src/main/scala/kafka/javaapi/consumer/ConsumerRebalanceListener.java b/core/src/main/scala/kafka/javaapi/consumer/ConsumerRebalanceListener.java
deleted file mode 100644
index ff23760a04792..0000000000000
--- a/core/src/main/scala/kafka/javaapi/consumer/ConsumerRebalanceListener.java
+++ /dev/null
@@ -1,54 +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.javaapi.consumer;
-
-
-import kafka.consumer.ConsumerThreadId;
-import java.util.Map;
-import java.util.Set;
-
-/**
- * This listener is used for execution of tasks defined by user when a consumer rebalance
- * occurs in {@link kafka.consumer.ZookeeperConsumerConnector}
- */
-/**
- * @deprecated since 0.11.0.0, this interface will be removed in a future release.
- */
-@Deprecated
-public interface ConsumerRebalanceListener {
-
-    /**
-     * This method is called after all the fetcher threads are stopped but before the
-     * ownership of partitions are released. Depending on whether auto offset commit is
-     * enabled or not, offsets may or may not have been committed.
-     * This listener is initially added to prevent duplicate messages on consumer rebalance
-     * in mirror maker, where offset auto commit is disabled to prevent data loss. It could
-     * also be used in more general cases.
-     * @param partitionOwnership The partition this consumer currently owns.
-     */
-    public void beforeReleasingPartitions(Map> partitionOwnership);
-
-    /**
-     * This method is called after the new partition assignment is finished but before fetcher
-     * threads start. A map of new global partition assignment is passed in as parameter.
-     * @param consumerId The consumer Id string of the consumer invoking this callback.
-     * @param globalPartitionAssignment A Map[topic, Map[Partition, ConsumerThreadId]]. It is the global partition
-     *                                  assignment of this consumer group.
-     */
-    public void beforeStartingFetchers(String consumerId, Map> globalPartitionAssignment);
-
-}
diff --git a/core/src/main/scala/kafka/javaapi/consumer/SimpleConsumer.scala b/core/src/main/scala/kafka/javaapi/consumer/SimpleConsumer.scala
deleted file mode 100644
index 188babbe8293b..0000000000000
--- a/core/src/main/scala/kafka/javaapi/consumer/SimpleConsumer.scala
+++ /dev/null
@@ -1,107 +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.javaapi.consumer
-
-import kafka.utils.threadsafe
-import kafka.javaapi.FetchResponse
-import kafka.javaapi.OffsetRequest
-
-/**
- * A consumer of kafka messages
- */
-@deprecated("This class has been deprecated and will be removed in a future release. " +
-            "Please use org.apache.kafka.clients.consumer.KafkaConsumer instead.", "0.11.0.0")
-@threadsafe
-class SimpleConsumer(val host: String,
-                     val port: Int,
-                     val soTimeout: Int,
-                     val bufferSize: Int,
-                     val clientId: String) {
-
-  private val underlying = new kafka.consumer.SimpleConsumer(host, port, soTimeout, bufferSize, clientId)
-
-  /**
-   *  Fetch a set of messages from a topic. This version of the fetch method
-   *  takes the Scala version of a fetch request (i.e.,
-   *  [[kafka.api.FetchRequest]] and is intended for use with the
-   *  [[kafka.api.FetchRequestBuilder]].
-   *
-   *  @param request  specifies the topic name, topic partition, starting byte offset, maximum bytes to be fetched.
-   *  @return a set of fetched messages
-   */
-  def fetch(request: kafka.api.FetchRequest): FetchResponse = {
-    import kafka.javaapi.Implicits._
-    underlying.fetch(request)
-  }
-
-  /**
-   *  Fetch a set of messages from a topic.
-   *
-   *  @param request specifies the topic name, topic partition, starting byte offset, maximum bytes to be fetched.
-   *  @return a set of fetched messages
-   */
-  def fetch(request: kafka.javaapi.FetchRequest): FetchResponse = {
-    fetch(request.underlying)
-  }
-
-  /**
-   *  Fetch metadata for a sequence of topics.
-   *
-   *  @param request specifies the versionId, clientId, sequence of topics.
-   *  @return metadata for each topic in the request.
-   */
-  def send(request: kafka.javaapi.TopicMetadataRequest): kafka.javaapi.TopicMetadataResponse = {
-    import kafka.javaapi.Implicits._
-    underlying.send(request.underlying)
-  }
-
-  /**
-   *  Get a list of valid offsets (up to maxSize) before the given time.
-   *
-   *  @param request a [[kafka.javaapi.OffsetRequest]] object.
-   *  @return a [[kafka.javaapi.OffsetResponse]] object.
-   */
-  def getOffsetsBefore(request: OffsetRequest): kafka.javaapi.OffsetResponse = {
-    import kafka.javaapi.Implicits._
-    underlying.getOffsetsBefore(request.underlying)
-  }
-
-  /**
-   * Commit offsets for a topic to Zookeeper
-   * @param request a [[kafka.javaapi.OffsetCommitRequest]] object.
-   * @return a [[kafka.javaapi.OffsetCommitResponse]] object.
-   */
-  def commitOffsets(request: kafka.javaapi.OffsetCommitRequest): kafka.javaapi.OffsetCommitResponse = {
-    import kafka.javaapi.Implicits._
-    underlying.commitOffsets(request.underlying)
-  }
-
-  /**
-   * Fetch offsets for a topic from Zookeeper
-   * @param request a [[kafka.javaapi.OffsetFetchRequest]] object.
-   * @return a [[kafka.javaapi.OffsetFetchResponse]] object.
-   */
-  def fetchOffsets(request: kafka.javaapi.OffsetFetchRequest): kafka.javaapi.OffsetFetchResponse = {
-    import kafka.javaapi.Implicits._
-    underlying.fetchOffsets(request.underlying)
-  }
-
-  def close() {
-    underlying.close
-  }
-}
diff --git a/core/src/main/scala/kafka/javaapi/consumer/ZookeeperConsumerConnector.scala b/core/src/main/scala/kafka/javaapi/consumer/ZookeeperConsumerConnector.scala
deleted file mode 100644
index d646938392d36..0000000000000
--- a/core/src/main/scala/kafka/javaapi/consumer/ZookeeperConsumerConnector.scala
+++ /dev/null
@@ -1,126 +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.javaapi.consumer
-
-import kafka.serializer._
-import kafka.consumer._
-import kafka.common.{OffsetAndMetadata, TopicAndPartition, MessageStreamsExistException}
-import java.util.concurrent.atomic.AtomicBoolean
-import scala.collection.JavaConverters._
-
-/**
- * This class handles the consumers interaction with zookeeper
- *
- * Directories:
- * 1. Consumer id registry:
- * /consumers/[group_id]/ids[consumer_id] -> topic1,...topicN
- * A consumer has a unique consumer id within a consumer group. A consumer registers its id as an ephemeral znode
- * and puts all topics that it subscribes to as the value of the znode. The znode is deleted when the client is gone.
- * A consumer subscribes to event changes of the consumer id registry within its group.
- *
- * The consumer id is picked up from configuration, instead of the sequential id assigned by ZK. Generated sequential
- * ids are hard to recover during temporary connection loss to ZK, since it's difficult for the client to figure out
- * whether the creation of a sequential znode has succeeded or not. More details can be found at
- * (http://wiki.apache.org/hadoop/ZooKeeper/ErrorHandling)
- *
- * 2. Broker node registry:
- * /brokers/[0...N] --> { "host" : "host:port",
- *                        "topics" : {"topic1": ["partition1" ... "partitionN"], ...,
- *                                    "topicN": ["partition1" ... "partitionN"] } }
- * This is a list of all present broker brokers. A unique logical node id is configured on each broker node. A broker
- * node registers itself on start-up and creates a znode with the logical node id under /brokers. The value of the znode
- * is a JSON String that contains (1) the host name and the port the broker is listening to, (2) a list of topics that
- * the broker serves, (3) a list of logical partitions assigned to each topic on the broker.
- * A consumer subscribes to event changes of the broker node registry.
- *
- * 3. Partition owner registry:
- * /consumers/[group_id]/owner/[topic]/[broker_id-partition_id] --> consumer_node_id
- * This stores the mapping before broker partitions and consumers. Each partition is owned by a unique consumer
- * within a consumer group. The mapping is reestablished after each rebalancing.
- *
- * 4. Consumer offset tracking:
- * /consumers/[group_id]/offsets/[topic]/[broker_id-partition_id] --> offset_counter_value
- * Each consumer tracks the offset of the latest message consumed for each partition.
- *
-*/
-
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.11.0.0")
-private[kafka] class ZookeeperConsumerConnector(val config: ConsumerConfig,
-                                                val enableFetcher: Boolean) // for testing only
-    extends ConsumerConnector {
-
-  private val underlying = new kafka.consumer.ZookeeperConsumerConnector(config, enableFetcher)
-  private val messageStreamCreated = new AtomicBoolean(false)
-
-  def this(config: ConsumerConfig) = this(config, true)
-
- // for java client
-  def createMessageStreams[K,V](
-        topicCountMap: java.util.Map[String,java.lang.Integer],
-        keyDecoder: Decoder[K],
-        valueDecoder: Decoder[V])
-      : java.util.Map[String,java.util.List[KafkaStream[K,V]]] = {
-
-    if (messageStreamCreated.getAndSet(true))
-      throw new MessageStreamsExistException(this.getClass.getSimpleName +
-                                   " can create message streams at most once",null)
-    val scalaTopicCountMap: Map[String, Int] = {
-      Map.empty[String, Int] ++ topicCountMap.asInstanceOf[java.util.Map[String, Int]].asScala
-    }
-    val scalaReturn = underlying.consume(scalaTopicCountMap, keyDecoder, valueDecoder)
-    val ret = new java.util.HashMap[String,java.util.List[KafkaStream[K,V]]]
-    for ((topic, streams) <- scalaReturn) {
-      val javaStreamList = new java.util.ArrayList[KafkaStream[K,V]]
-      for (stream <- streams)
-        javaStreamList.add(stream)
-      ret.put(topic, javaStreamList)
-    }
-    ret
-  }
-
-  def createMessageStreams(topicCountMap: java.util.Map[String,java.lang.Integer]): java.util.Map[String,java.util.List[KafkaStream[Array[Byte],Array[Byte]]]] =
-    createMessageStreams(topicCountMap, new DefaultDecoder(), new DefaultDecoder())
-
-  def createMessageStreamsByFilter[K,V](topicFilter: TopicFilter, numStreams: Int, keyDecoder: Decoder[K], valueDecoder: Decoder[V]) =
-    underlying.createMessageStreamsByFilter(topicFilter, numStreams, keyDecoder, valueDecoder).asJava
-
-  def createMessageStreamsByFilter(topicFilter: TopicFilter, numStreams: Int) =
-    createMessageStreamsByFilter(topicFilter, numStreams, new DefaultDecoder(), new DefaultDecoder())
-
-  def createMessageStreamsByFilter(topicFilter: TopicFilter) =
-    createMessageStreamsByFilter(topicFilter, 1, new DefaultDecoder(), new DefaultDecoder())
-
-  def commitOffsets() {
-    underlying.commitOffsets(true)
-  }
-
-  def commitOffsets(retryOnFailure: Boolean) {
-    underlying.commitOffsets(retryOnFailure)
-  }
-
-  def commitOffsets(offsetsToCommit: java.util.Map[TopicAndPartition, OffsetAndMetadata], retryOnFailure: Boolean) {
-    underlying.commitOffsets(offsetsToCommit.asScala.toMap, retryOnFailure)
-  }
-
-  def setConsumerRebalanceListener(consumerRebalanceListener: ConsumerRebalanceListener) {
-    underlying.setConsumerRebalanceListener(consumerRebalanceListener)
-  }
-
-  def shutdown() {
-    underlying.shutdown
-  }
-}
diff --git a/core/src/main/scala/kafka/javaapi/message/ByteBufferMessageSet.scala b/core/src/main/scala/kafka/javaapi/message/ByteBufferMessageSet.scala
deleted file mode 100644
index 590db83be42df..0000000000000
--- a/core/src/main/scala/kafka/javaapi/message/ByteBufferMessageSet.scala
+++ /dev/null
@@ -1,67 +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.javaapi.message
-
-import java.nio.ByteBuffer
-
-import kafka.common.LongRef
-import kafka.message._
-
-import scala.collection.JavaConverters._
-
-class ByteBufferMessageSet(val buffer: ByteBuffer) extends MessageSet {
-  private val underlying: kafka.message.ByteBufferMessageSet = new kafka.message.ByteBufferMessageSet(buffer)
-  
-  def this(compressionCodec: CompressionCodec, messages: java.util.List[Message]) {
-    this(new kafka.message.ByteBufferMessageSet(compressionCodec, new LongRef(0), messages.asScala: _*).buffer)
-  }
-
-  def this(messages: java.util.List[Message]) {
-    this(NoCompressionCodec, messages)
-  }
-
-  def validBytes: Int = underlying.validBytes
-
-  def getBuffer = buffer
-
-  override def iterator: java.util.Iterator[MessageAndOffset] = new java.util.Iterator[MessageAndOffset] {
-    val underlyingIterator = underlying.iterator
-    override def hasNext(): Boolean = {
-      underlyingIterator.hasNext
-    }
-
-    override def next(): MessageAndOffset = {
-      underlyingIterator.next
-    }
-
-    override def remove = throw new UnsupportedOperationException("remove API on MessageSet is not supported")
-  }
-
-  override def toString: String = underlying.toString
-
-  def sizeInBytes: Int = underlying.sizeInBytes
-
-  override def equals(other: Any): Boolean = {
-    other match {
-      case that: ByteBufferMessageSet => buffer.equals(that.buffer)
-      case _ => false
-    }
-  }
-
-
-  override def hashCode: Int = buffer.hashCode
-}
diff --git a/core/src/main/scala/kafka/javaapi/message/MessageSet.scala b/core/src/main/scala/kafka/javaapi/message/MessageSet.scala
deleted file mode 100644
index 80a67cd3546be..0000000000000
--- a/core/src/main/scala/kafka/javaapi/message/MessageSet.scala
+++ /dev/null
@@ -1,55 +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.javaapi.message
-
-
-import kafka.message.{MessageAndOffset, InvalidMessageException}
-
-
-/**
- * A set of messages. A message set has a fixed serialized form, though the container
- * for the bytes could be either in-memory or on disk. A The format of each message is
- * as follows:
- * 4 byte size containing an integer N
- * N message bytes as described in the message class
- */
-abstract class MessageSet extends java.lang.Iterable[MessageAndOffset] {
-
-  /**
-   * Provides an iterator over the messages in this set
-   */
-  def iterator: java.util.Iterator[MessageAndOffset]
-
-  /**
-   * Gives the total size of this message set in bytes
-   */
-  def sizeInBytes: Int
-
-  /**
-   * Validate the checksum of all the messages in the set. Throws an InvalidMessageException if the checksum doesn't
-   * match the payload for any message.
-   */
-  def validate(): Unit = {
-    val thisIterator = this.iterator
-    while(thisIterator.hasNext) {
-      val messageAndOffset = thisIterator.next
-      if(!messageAndOffset.message.isValid)
-        throw new InvalidMessageException
-    }
-  }
-}
diff --git a/core/src/main/scala/kafka/javaapi/producer/Producer.scala b/core/src/main/scala/kafka/javaapi/producer/Producer.scala
deleted file mode 100644
index b0b40b9fc0b1e..0000000000000
--- a/core/src/main/scala/kafka/javaapi/producer/Producer.scala
+++ /dev/null
@@ -1,52 +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.javaapi.producer
-
-import kafka.producer.ProducerConfig
-import kafka.producer.KeyedMessage
-import scala.collection.mutable
-
-@deprecated("This class has been deprecated and will be removed in a future release. " +
-            "Please use org.apache.kafka.clients.producer.KafkaProducer instead.", "0.10.0.0")
-class Producer[K,V](private val underlying: kafka.producer.Producer[K,V]) // for testing only
-{
-  def this(config: ProducerConfig) = this(new kafka.producer.Producer[K,V](config))
-  /**
-   * Sends the data to a single topic, partitioned by key, using either the
-   * synchronous or the asynchronous producer
-   * @param message the producer data object that encapsulates the topic, key and message data
-   */
-  def send(message: KeyedMessage[K,V]) {
-    underlying.send(message)
-  }
-
-  /**
-   * Use this API to send data to multiple topics
-   * @param messages list of producer data objects that encapsulate the topic, key and message data
-   */
-  def send(messages: java.util.List[KeyedMessage[K,V]]) {
-    import collection.JavaConversions._
-    underlying.send((messages: mutable.Buffer[KeyedMessage[K,V]]).toSeq: _*)
-  }
-
-  /**
-   * Close API to close the producer pool connections to all Kafka brokers. Also closes
-   * the zookeeper client connection if one exists
-   */
-  def close() = underlying.close()
-}
diff --git a/core/src/main/scala/kafka/log/AbstractIndex.scala b/core/src/main/scala/kafka/log/AbstractIndex.scala
index 9696d8d21152f..5d7de88b96a42 100644
--- a/core/src/main/scala/kafka/log/AbstractIndex.scala
+++ b/core/src/main/scala/kafka/log/AbstractIndex.scala
@@ -17,34 +17,92 @@
 
 package kafka.log
 
-import java.io.{File, RandomAccessFile}
-import java.nio.{ByteBuffer, MappedByteBuffer}
+import java.io.{Closeable, File, RandomAccessFile}
 import java.nio.channels.FileChannel
+import java.nio.file.Files
+import java.nio.{ByteBuffer, MappedByteBuffer}
 import java.util.concurrent.locks.{Lock, ReentrantLock}
 
-import kafka.log.IndexSearchType.IndexSearchEntity
+import kafka.common.IndexOffsetOverflowException
 import kafka.utils.CoreUtils.inLock
 import kafka.utils.{CoreUtils, Logging}
-import org.apache.kafka.common.utils.{MappedByteBuffers, OperatingSystem, Utils}
-
-import scala.math.ceil
+import org.apache.kafka.common.utils.{ByteBufferUnmapper, OperatingSystem, Utils}
 
 /**
  * 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 Logging {
+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
   @volatile
   private var _length: Long = _
-
   protected def entrySize: Int
 
+  /*
+   Kafka mmaps index files into memory, and all the read / write operations of the index is through OS page cache. This
+   avoids blocked disk I/O in most cases.
+
+   To the extent of our knowledge, all the modern operating systems use LRU policy or its variants to manage page
+   cache. Kafka always appends to the end of the index file, and almost all the index lookups (typically from in-sync
+   followers or consumers) are very close to the end of the index. So, the LRU cache replacement policy should work very
+   well with Kafka's index access pattern.
+
+   However, when looking up index, the standard binary search algorithm is not cache friendly, and can cause unnecessary
+   page faults (the thread is blocked to wait for reading some index entries from hard disk, as those entries are not
+   cached in the page cache).
+
+   For example, in an index with 13 pages, to lookup an entry in the last page (page #12), the standard binary search
+   algorithm will read index entries in page #0, 6, 9, 11, and 12.
+   page number: |0|1|2|3|4|5|6|7|8|9|10|11|12 |
+   steps:       |1| | | | | |3| | |4|  |5 |2/6|
+   In each page, there are hundreds log entries, corresponding to hundreds to thousands of kafka messages. When the
+   index gradually growing from the 1st entry in page #12 to the last entry in page #12, all the write (append)
+   operations are in page #12, and all the in-sync follower / consumer lookups read page #0,6,9,11,12. As these pages
+   are always used in each in-sync lookup, we can assume these pages are fairly recently used, and are very likely to be
+   in the page cache. When the index grows to page #13, the pages needed in a in-sync lookup change to #0, 7, 10, 12,
+   and 13:
+   page number: |0|1|2|3|4|5|6|7|8|9|10|11|12|13 |
+   steps:       |1| | | | | | |3| | | 4|5 | 6|2/7|
+   Page #7 and page #10 have not been used for a very long time. They are much less likely to be in the page cache, than
+   the other pages. The 1st lookup, after the 1st index entry in page #13 is appended, is likely to have to read page #7
+   and page #10 from disk (page fault), which can take up to more than a second. In our test, this can cause the
+   at-least-once produce latency to jump to about 1 second from a few ms.
+
+   Here, we use a more cache-friendly lookup algorithm:
+   if (target > indexEntry[end - N]) // if the target is in the last N entries of the index
+      binarySearch(end - N, end)
+   else
+      binarySearch(begin, end - N)
+
+   If possible, we only look up in the last N entries of the index. By choosing a proper constant N, all the in-sync
+   lookups should go to the 1st branch. We call the last N entries the "warm" section. As we frequently look up in this
+   relatively small section, the pages containing this section are more likely to be in the page cache.
+
+   We set N (_warmEntries) to 8192, because
+   1. This number is small enough to guarantee all the pages of the "warm" section is touched in every warm-section
+      lookup. So that, the entire warm section is really "warm".
+      When doing warm-section lookup, following 3 entries are always touched: indexEntry(end), indexEntry(end-N),
+      and indexEntry((end*2 -N)/2). If page size >= 4096, all the warm-section pages (3 or fewer) are touched, when we
+      touch those 3 entries. As of 2018, 4096 is the smallest page size for all the processors (x86-32, x86-64, MIPS,
+      SPARC, Power, ARM etc.).
+   2. This number is large enough to guarantee most of the in-sync lookups are in the warm-section. With default Kafka
+      settings, 8KB index corresponds to about 4MB (offset index) or 2.7MB (time index) log messages.
+
+   We can't set make N (_warmEntries) to be larger than 8192, as there is no simple way to guarantee all the "warm"
+   section pages are really warm (touched in every lookup) on a typical 4KB-page host.
+
+   In there future, we may use a backend thread to periodically touch the entire warm section. So that, we can
+   1) support larger warm section
+   2) make sure the warm section of low QPS topic-partitions are really warm.
+ */
+  protected def _warmEntries: Int = 8192 / entrySize
+
   protected val lock = new ReentrantLock
 
   @volatile
@@ -75,7 +133,7 @@ abstract class AbstractIndex[K, V](@volatile var file: File, val baseOffset: Lon
         idx.position(roundDownToExactMultiple(idx.limit(), entrySize))
       idx
     } finally {
-      CoreUtils.swallow(raf.close(), this)
+      CoreUtils.swallow(raf.close(), AbstractIndex)
     }
   }
 
@@ -83,23 +141,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
@@ -114,23 +176,26 @@ abstract class AbstractIndex[K, V](@volatile var file: File, val baseOffset: Lon
       val roundedNewSize = roundDownToExactMultiple(newSize, entrySize)
 
       if (_length == roundedNewSize) {
+        debug(s"Index ${file.getAbsolutePath} was not resized because it already has size $roundedNewSize")
         false
       } else {
         val raf = new RandomAccessFile(file, "rw")
         try {
           val position = mmap.position()
 
-          /* Windows won't let us modify the file length while the file is mmapped :-( */
-          if (OperatingSystem.IS_WINDOWS)
+          /* Windows or z/OS won't let us modify the file length while the file is mmapped :-( */
+          if (OperatingSystem.IS_WINDOWS || OperatingSystem.IS_ZOS)
             safeForceUnmap()
           raf.setLength(roundedNewSize)
           _length = roundedNewSize
           mmap = raf.getChannel().map(FileChannel.MapMode.READ_WRITE, 0, roundedNewSize)
           _maxEntries = mmap.limit() / entrySize
           mmap.position(position)
+          debug(s"Resized ${file.getAbsolutePath} to $roundedNewSize, position is ${mmap.position()} " +
+            s"and limit is ${mmap.limit()}")
           true
         } finally {
-          CoreUtils.swallow(raf.close(), this)
+          CoreUtils.swallow(raf.close(), AbstractIndex)
         }
       }
     }
@@ -141,40 +206,37 @@ 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()
     }
   }
 
   /**
-   * Delete this index file
+   * Delete this index file.
+   *
+   * @throws IOException if deletion fails due to an I/O error
+   * @return `true` if the file was deleted by this method; `false` if the file could not be deleted because it did
+   *         not exist
    */
-  def delete(): Boolean = {
-    info(s"Deleting index ${file.getAbsolutePath}")
-    inLock(lock) {
-      // On JVM, a memory mapping is typically unmapped by garbage collector.
-      // However, in some cases it can pause application threads(STW) for a long moment reading metadata from a physical disk.
-      // To prevent this, we forcefully cleanup memory mapping within proper execution which never affects API responsiveness.
-      // See https://issues.apache.org/jira/browse/KAFKA-4614 for the details.
-      safeForceUnmap()
-    }
-    file.delete()
+  def deleteIfExists(): Boolean = {
+    closeHandler()
+    Files.deleteIfExists(file.toPath)
   }
 
   /**
    * 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)
     }
@@ -183,14 +245,19 @@ 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()
   }
 
   def closeHandler(): Unit = {
+    // On JVM, a memory mapping is typically unmapped by garbage collector.
+    // However, in some cases it can pause application threads(STW) for a long moment reading metadata from a physical disk.
+    // To prevent this, we forcefully cleanup memory mapping within proper execution which never affects API responsiveness.
+    // See https://issues.apache.org/jira/browse/KAFKA-4614 for the details.
     inLock(lock) {
       safeForceUnmap()
     }
@@ -199,14 +266,14 @@ abstract class AbstractIndex[K, V](@volatile var file: File, val baseOffset: Lon
   /**
    * Do a basic sanity check on this index to detect obvious problems
    *
-   * @throws IllegalArgumentException if any problems are found
+   * @throws CorruptIndexException if any problems are found
    */
   def sanityCheck(): Unit
 
   /**
    * Remove all the entries from the index.
    */
-  def truncate(): Unit
+  protected def truncate(): Unit
 
   /**
    * Remove all entries from the index which have an offset greater than or equal to the given offset.
@@ -214,32 +281,62 @@ abstract class AbstractIndex[K, V](@volatile var file: File, val baseOffset: Lon
    */
   def truncateTo(offset: Long): Unit
 
+  /**
+   * Remove all the entries from the index and resize the index to the max index size.
+   */
+  def reset(): Unit = {
+    truncate()
+    resize(maxIndexSize)
+  }
+
+  /**
+   * Get offset relative to base offset of this index
+   * @throws IndexOffsetOverflowException
+   */
+  def relativeOffset(offset: Long): Int = {
+    val relativeOffset = toRelative(offset)
+    if (relativeOffset.isEmpty)
+      throw new IndexOffsetOverflowException(s"Integer overflow for offset: $offset (${file.getAbsoluteFile})")
+    relativeOffset.get
+  }
+
+  /**
+   * Check if a particular offset is valid to be appended to this index.
+   * @param offset The offset to check
+   * @return true if this offset is valid to be appended to this index; false otherwise
+   */
+  def canAppendOffset(offset: Long): Boolean = {
+    toRelative(offset).isDefined
+  }
+
   protected def safeForceUnmap(): Unit = {
-    try forceUnmap()
-    catch {
-      case t: Throwable => error(s"Error unmapping index $file", t)
+    if (mmap != null) {
+      try forceUnmap()
+      catch {
+        case t: Throwable => error(s"Error unmapping index $file", t)
+      }
     }
   }
 
   /**
    * Forcefully free the buffer's mmap.
    */
-  protected[log] def forceUnmap() {
-    try MappedByteBuffers.unmap(file.getAbsolutePath, mmap)
+  protected[log] def forceUnmap(): Unit = {
+    try ByteBufferUnmapper.unmap(file.getAbsolutePath, mmap)
     finally mmap = null // Accessing unmapped mmap crashes JVM by SEGV so we null it out to be safe
   }
 
   /**
-   * Execute the given function in a lock only if we are running on windows. We do this
-   * because Windows won't let us resize a file while it is mmapped. As a result we have to force unmap it
+   * Execute the given function in a lock only if we are running on windows or z/OS. We do this
+   * because Windows or z/OS won't let us resize a file while it is mmapped. As a result we have to force unmap it
    * and this requires synchronizing reads.
    */
   protected def maybeLock[T](lock: Lock)(fun: => T): T = {
-    if (OperatingSystem.IS_WINDOWS)
+    if (OperatingSystem.IS_WINDOWS || OperatingSystem.IS_ZOS)
       lock.lock()
     try fun
     finally {
-      if (OperatingSystem.IS_WINDOWS)
+      if (OperatingSystem.IS_WINDOWS || OperatingSystem.IS_ZOS)
         lock.unlock()
     }
   }
@@ -261,49 +358,58 @@ abstract class AbstractIndex[K, V](@volatile var file: File, val baseOffset: Lon
    * @param target The index key to look for
    * @return The slot found or -1 if the least entry in the index is larger than the target key or the index is empty
    */
-  protected def largestLowerBoundSlotFor(idx: ByteBuffer, target: Long, searchEntity: IndexSearchEntity): Int =
+  protected def largestLowerBoundSlotFor(idx: ByteBuffer, target: Long, searchEntity: IndexSearchType): Int =
     indexSlotRangeFor(idx, target, searchEntity)._1
 
   /**
    * Find the smallest entry greater than or equal the target key or value. If none can be found, -1 is returned.
    */
-  protected def smallestUpperBoundSlotFor(idx: ByteBuffer, target: Long, searchEntity: IndexSearchEntity): Int =
+  protected def smallestUpperBoundSlotFor(idx: ByteBuffer, target: Long, searchEntity: IndexSearchType): Int =
     indexSlotRangeFor(idx, target, searchEntity)._2
 
   /**
    * Lookup lower and upper bounds for the given target.
    */
-  private def indexSlotRangeFor(idx: ByteBuffer, target: Long, searchEntity: IndexSearchEntity): (Int, Int) = {
+  private def indexSlotRangeFor(idx: ByteBuffer, target: Long, searchEntity: IndexSearchType): (Int, Int) = {
     // check if the index is empty
     if(_entries == 0)
       return (-1, -1)
 
+    def binarySearch(begin: Int, end: Int) : (Int, Int) = {
+      // binary search for the entry
+      var lo = begin
+      var hi = end
+      while(lo < hi) {
+        val mid = (lo + hi + 1) >>> 1
+        val found = parseEntry(idx, mid)
+        val compareResult = compareIndexEntry(found, target, searchEntity)
+        if(compareResult > 0)
+          hi = mid - 1
+        else if(compareResult < 0)
+          lo = mid
+        else
+          return (mid, mid)
+      }
+      (lo, if (lo == _entries - 1) -1 else lo + 1)
+    }
+
+    val firstHotEntry = Math.max(0, _entries - 1 - _warmEntries)
+    // check if the target offset is in the warm section of the index
+    if(compareIndexEntry(parseEntry(idx, firstHotEntry), target, searchEntity) < 0) {
+      return binarySearch(firstHotEntry, _entries - 1)
+    }
+
     // check if the target offset is smaller than the least offset
     if(compareIndexEntry(parseEntry(idx, 0), target, searchEntity) > 0)
       return (-1, 0)
 
-    // binary search for the entry
-    var lo = 0
-    var hi = _entries - 1
-    while(lo < hi) {
-      val mid = ceil(hi/2.0 + lo/2.0).toInt
-      val found = parseEntry(idx, mid)
-      val compareResult = compareIndexEntry(found, target, searchEntity)
-      if(compareResult > 0)
-        hi = mid - 1
-      else if(compareResult < 0)
-        lo = mid
-      else
-        return (mid, mid)
-    }
-
-    (lo, if (lo == _entries - 1) -1 else lo + 1)
+    binarySearch(0, firstHotEntry)
   }
 
-  private def compareIndexEntry(indexEntry: IndexEntry, target: Long, searchEntity: IndexSearchEntity): Int = {
+  private def compareIndexEntry(indexEntry: IndexEntry, target: Long, searchEntity: IndexSearchType): Int = {
     searchEntity match {
-      case IndexSearchType.KEY => indexEntry.indexKey.compareTo(target)
-      case IndexSearchType.VALUE => indexEntry.indexValue.compareTo(target)
+      case IndexSearchType.KEY => java.lang.Long.compare(indexEntry.indexKey, target)
+      case IndexSearchType.VALUE => java.lang.Long.compare(indexEntry.indexValue, target)
     }
   }
 
@@ -313,9 +419,22 @@ abstract class AbstractIndex[K, V](@volatile var file: File, val baseOffset: Lon
    */
   private def roundDownToExactMultiple(number: Int, factor: Int) = factor * (number / factor)
 
+  private def toRelative(offset: Long): Option[Int] = {
+    val relativeOffset = offset - baseOffset
+    if (relativeOffset < 0 || relativeOffset > Int.MaxValue)
+      None
+    else
+      Some(relativeOffset.toInt)
+  }
+
+}
+
+object AbstractIndex extends Logging {
+  override val loggerName: String = classOf[AbstractIndex].getName
 }
 
-object IndexSearchType extends Enumeration {
-  type IndexSearchEntity = Value
-  val KEY, VALUE = Value
+sealed trait IndexSearchType
+object IndexSearchType {
+  case object KEY extends IndexSearchType
+  case object VALUE extends IndexSearchType
 }
diff --git a/core/src/main/scala/kafka/log/CorruptIndexException.scala b/core/src/main/scala/kafka/log/CorruptIndexException.scala
new file mode 100644
index 0000000000000..b39ee5b3bbbae
--- /dev/null
+++ b/core/src/main/scala/kafka/log/CorruptIndexException.scala
@@ -0,0 +1,20 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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
+
+class CorruptIndexException(message: String) extends RuntimeException(message)
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 c7ec3bd9682e7..d28f32c030aab 100644
--- a/core/src/main/scala/kafka/log/Log.scala
+++ b/core/src/main/scala/kafka/log/Log.scala
@@ -18,72 +18,143 @@
 package kafka.log
 
 import java.io.{File, IOException}
-import java.nio.file.Files
+import java.lang.{Long => JLong}
+import java.nio.file.{Files, NoSuchFileException}
 import java.text.NumberFormat
+import java.util.Map.{Entry => JEntry}
+import java.util.Optional
 import java.util.concurrent.atomic._
 import java.util.concurrent.{ConcurrentNavigableMap, ConcurrentSkipListMap, TimeUnit}
+import java.util.regex.Pattern
 
-import kafka.api.KAFKA_0_10_0_IV0
-import kafka.common.{InvalidOffsetException, KafkaException, LongRef}
+import kafka.api.{ApiVersion, KAFKA_0_10_0_IV0}
+import kafka.common.{LogSegmentOffsetOverflowException, LongRef, OffsetsOutOfOrderException, UnexpectedAppendOffsetException}
+import kafka.message.{BrokerCompressionCodec, CompressionCodec, NoCompressionCodec}
 import kafka.metrics.KafkaMetricsGroup
-import kafka.server.{BrokerTopicStats, FetchDataInfo, LogDirFailureChannel, LogOffsetMetadata}
+import kafka.server.checkpoints.LeaderEpochCheckpointFile
+import kafka.server.epoch.LeaderEpochFileCache
+import kafka.server.{BrokerTopicStats, FetchDataInfo, FetchHighWatermark, FetchIsolation, FetchLogEnd, FetchTxnCommitted, LogDirFailureChannel, LogOffsetMetadata, OffsetAndEpoch}
 import kafka.utils._
-import org.apache.kafka.common.errors.{CorruptRecordException, KafkaStorageException, OffsetOutOfRangeException, RecordBatchTooLargeException, RecordTooLargeException, UnsupportedForMessageFormatException}
+import org.apache.kafka.common.errors._
+import org.apache.kafka.common.message.FetchResponseData
+import org.apache.kafka.common.record.FileRecords.TimestampAndOffset
 import org.apache.kafka.common.record._
-import org.apache.kafka.common.requests.{IsolationLevel, ListOffsetRequest}
+import org.apache.kafka.common.requests.FetchResponse.AbortedTransaction
+import org.apache.kafka.common.requests.OffsetsForLeaderEpochResponse.UNDEFINED_EPOCH_OFFSET
+import org.apache.kafka.common.requests.ProduceResponse.RecordError
+import org.apache.kafka.common.requests.ListOffsetRequest
+import org.apache.kafka.common.utils.{Time, Utils}
+import org.apache.kafka.common.{InvalidRecordException, KafkaException, TopicPartition}
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 import scala.collection.mutable.{ArrayBuffer, ListBuffer}
-import scala.collection.{Seq, mutable}
-import com.yammer.metrics.core.Gauge
-import org.apache.kafka.common.utils.{Time, Utils}
-import kafka.message.{BrokerCompressionCodec, CompressionCodec, NoCompressionCodec}
-import kafka.server.checkpoints.{LeaderEpochCheckpointFile, LeaderEpochFile}
-import kafka.server.epoch.{LeaderEpochCache, LeaderEpochFileCache}
-import org.apache.kafka.common.TopicPartition
-import org.apache.kafka.common.requests.FetchResponse.AbortedTransaction
-import java.util.Map.{Entry => JEntry}
-import java.lang.{Long => JLong}
-import java.util.regex.Pattern
+import scala.collection.{Seq, Set, mutable}
 
 object LogAppendInfo {
-  val UnknownLogAppendInfo = LogAppendInfo(-1, -1, RecordBatch.NO_TIMESTAMP, -1L, RecordBatch.NO_TIMESTAMP, -1L,
-    RecordsProcessingStats.EMPTY, NoCompressionCodec, NoCompressionCodec, -1, -1, offsetsMonotonic = false)
+  val UnknownLogAppendInfo = LogAppendInfo(None, -1, None, RecordBatch.NO_TIMESTAMP, -1L, RecordBatch.NO_TIMESTAMP, -1L,
+    RecordConversionStats.EMPTY, NoCompressionCodec, NoCompressionCodec, -1, -1, offsetsMonotonic = false, -1L)
 
   def unknownLogAppendInfoWithLogStartOffset(logStartOffset: Long): LogAppendInfo =
-    LogAppendInfo(-1, -1, RecordBatch.NO_TIMESTAMP, -1L, RecordBatch.NO_TIMESTAMP, logStartOffset,
-      RecordsProcessingStats.EMPTY, NoCompressionCodec, NoCompressionCodec, -1, -1, offsetsMonotonic = false)
+    LogAppendInfo(None, -1, None, RecordBatch.NO_TIMESTAMP, -1L, RecordBatch.NO_TIMESTAMP, logStartOffset,
+      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: Seq[RecordError], errorMessage: String): LogAppendInfo =
+    LogAppendInfo(None, -1, None, RecordBatch.NO_TIMESTAMP, -1L, RecordBatch.NO_TIMESTAMP, logStartOffset,
+      RecordConversionStats.EMPTY, NoCompressionCodec, NoCompressionCodec, -1, -1,
+      offsetsMonotonic = false, -1L, recordErrors, errorMessage)
+}
+
+sealed trait LeaderHwChange
+object LeaderHwChange {
+  case object Increased extends LeaderHwChange
+  case object Same extends LeaderHwChange
+  case object None extends LeaderHwChange
 }
 
 /**
  * Struct to hold various quantities we compute about each message set before appending to the log
  *
  * @param firstOffset The first offset in the message set unless the message format is less than V2 and we are appending
- *                    to the follower. In that case, this will be the last offset for performance reasons.
+ *                    to the follower.
  * @param lastOffset The last offset in the message set
+ * @param lastLeaderEpoch The partition leader epoch corresponding to the last offset, if available.
  * @param maxTimestamp The maximum timestamp of the message set.
  * @param offsetOfMaxTimestamp The offset of the message with the maximum timestamp.
  * @param logAppendTime The log append time (if used) of the message set, otherwise Message.NoTimestamp
  * @param logStartOffset The start offset of the log at the time of this append.
- * @param recordsProcessingStats Statistics collected during record processing, `null` if `assignOffsets` is `false`
+ * @param recordConversionStats Statistics collected during record processing, `null` if `assignOffsets` is `false`
  * @param sourceCodec The source codec used in the message set (send by the producer)
  * @param targetCodec The target codec of the message set(after applying the broker compression configuration if any)
  * @param shallowCount The number of shallow messages
  * @param validBytes The number of valid bytes
  * @param offsetsMonotonic Are the offsets in this message set monotonically increasing
+ * @param lastOffsetOfFirstBatch The last offset of the first batch
+ * @param leaderHwChange Incremental if the high watermark needs to be increased after appending record.
+ *                       Same if high watermark is not changed. None is the default value and it means append failed
+ *
  */
-case class LogAppendInfo(var firstOffset: Long,
+case class LogAppendInfo(var firstOffset: Option[Long],
                          var lastOffset: Long,
+                         var lastLeaderEpoch: Option[Int],
                          var maxTimestamp: Long,
                          var offsetOfMaxTimestamp: Long,
                          var logAppendTime: Long,
                          var logStartOffset: Long,
-                         var recordsProcessingStats: RecordsProcessingStats,
+                         var recordConversionStats: RecordConversionStats,
                          sourceCodec: CompressionCodec,
                          targetCodec: CompressionCodec,
                          shallowCount: Int,
                          validBytes: Int,
-                         offsetsMonotonic: Boolean)
+                         offsetsMonotonic: Boolean,
+                         lastOffsetOfFirstBatch: Long,
+                         recordErrors: Seq[RecordError] = List(),
+                         errorMessage: String = null,
+                         leaderHwChange: LeaderHwChange = LeaderHwChange.None) {
+  /**
+   * 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
+   * older than 2, we use the last offset of the first batch as an approximation of the first
+   * offset to avoid decompressing the data.
+   */
+  def firstOrLastOffsetOfFirstBatch: Long = firstOffset.getOrElse(lastOffsetOfFirstBatch)
+
+  /**
+   * Get the (maximum) number of messages described by LogAppendInfo
+   * @return Maximum possible number of messages described by LogAppendInfo
+   */
+  def numMessages: Long = {
+    firstOffset match {
+      case Some(firstOffsetVal) if (firstOffsetVal >= 0 && lastOffset >= 0) => (lastOffset - firstOffsetVal + 1)
+      case _ => 0
+    }
+  }
+}
+
+/**
+ * Container class which represents a snapshot of the significant offsets for a partition. This allows fetching
+ * of these offsets atomically without the possibility of a leader change affecting their consistency relative
+ * to each other. See [[Log.fetchOffsetSnapshot()]].
+ */
+case class LogOffsetSnapshot(logStartOffset: Long,
+                             logEndOffset: LogOffsetMetadata,
+                             highWatermark: LogOffsetMetadata,
+                             lastStableOffset: LogOffsetMetadata)
+
+/**
+ * Another container which is used for lower level reads using  [[kafka.cluster.Partition.readRecords()]].
+ */
+case class LogReadInfo(fetchedData: FetchDataInfo,
+                       divergingEpoch: Option[FetchResponseData.EpochEndOffset],
+                       highWatermark: Long,
+                       logStartOffset: Long,
+                       logEndOffset: Long,
+                       lastStableOffset: Long)
 
 /**
  * A class used to hold useful metadata about a completed transaction. This is used to build
@@ -105,6 +176,37 @@ case class CompletedTxn(producerId: Long, firstOffset: Long, lastOffset: Long, i
   }
 }
 
+/**
+ * A class used to hold params required to decide to rotate a log segment or not.
+ */
+case class RollParams(maxSegmentMs: Long,
+                      maxSegmentBytes: Int,
+                      maxTimestampInMessages: Long,
+                      maxOffsetInMessages: Long,
+                      messagesSize: Int,
+                      now: Long)
+
+object RollParams {
+  def apply(config: LogConfig, appendInfo: LogAppendInfo, messagesSize: Int, now: Long): RollParams = {
+   new RollParams(config.maxSegmentMs,
+     config.segmentSize,
+     appendInfo.maxTimestamp,
+     appendInfo.lastOffset,
+     messagesSize, now)
+  }
+}
+
+sealed trait LogStartOffsetIncrementReason
+case object ClientRecordDeletion extends LogStartOffsetIncrementReason {
+  override def toString: String = "client delete records request"
+}
+case object LeaderOffsetIncremented extends LogStartOffsetIncrementReason {
+  override def toString: String = "leader offset increment"
+}
+case object SegmentDeletion extends LogStartOffsetIncrementReason {
+  override def toString: String = "segment deletion"
+}
+
 /**
  * An append-only log for storing messages.
  *
@@ -113,7 +215,7 @@ case class CompletedTxn(producerId: Long, firstOffset: Long, lastOffset: Long, i
  * 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 :
@@ -132,44 +234,41 @@ case class CompletedTxn(producerId: Long, firstOffset: Long, lastOffset: Long, i
  * @param time The time instance used for checking the clock
  * @param maxProducerIdExpirationMs The maximum amount of time to wait before a producer id is considered expired
  * @param producerIdExpirationCheckIntervalMs How often to check for producer ids which need to be expired
+ * @param hadCleanShutdown boolean flag to indicate if the Log had a clean/graceful shutdown last time. true means
+ *                         clean shutdown whereas false means a crash.
  */
 @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,
           scheduler: Scheduler,
           brokerTopicStats: BrokerTopicStats,
-          time: Time,
+          val time: Time,
           val maxProducerIdExpirationMs: Int,
           val producerIdExpirationCheckIntervalMs: Int,
           val topicPartition: TopicPartition,
           val producerStateManager: ProducerStateManager,
-          logDirFailureChannel: LogDirFailureChannel) extends Logging with KafkaMetricsGroup {
+          logDirFailureChannel: LogDirFailureChannel,
+          private val hadCleanShutdown: Boolean = true) extends Logging with KafkaMetricsGroup {
 
   import kafka.log.Log._
 
+  this.logIdent = s"[Log partition=$topicPartition, dir=${dir.getParent}] "
+
   /* 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 for index files of this log will be closed with either delete() or closeHandlers()
+
+  // 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
-  }
-
-  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
@@ -183,71 +282,251 @@ 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]
 
-  @volatile private var _leaderEpochCache: LeaderEpochCache = initializeLeaderEpochCache()
+  // Visible for testing
+  @volatile var leaderEpochCache: Option[LeaderEpochFileCache] = None
 
   locally {
-    val startMs = time.milliseconds
+    // create the log directory if it doesn't exist
+    Files.createDirectories(dir.toPath)
+
+    initializeLeaderEpochCache()
 
     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.clearAndFlushLatest(nextOffsetMetadata.messageOffset)
+    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.clearAndFlushEarliest(logStartOffset)
+    leaderEpochCache.foreach(_.truncateFromStart(logStartOffset))
 
-    loadProducerState(logEndOffset, reloadFromCleanShutdown = hasCleanShutdownFile)
+    // Any segment loading or recovery code must not use producerStateManager, so that we can build the full state here
+    // from scratch.
+    if (!producerStateManager.isEmpty)
+      throw new IllegalStateException("Producer state must be empty during log initialization")
 
-    info("Completed load of log %s with %d log segments, log start offset %d and log end offset %d in %d ms"
-      .format(name, segments.size(), logStartOffset, logEndOffset, time.milliseconds - startMs))
+    // Reload all snapshots into the ProducerStateManager cache, the intermediate ProducerStateManager used
+    // during log recovery may have deleted some files without the Log.producerStateManager instance witnessing the
+    // deletion.
+    producerStateManager.removeStraySnapshots(segments.values().asScala.map(_.baseOffset).toSeq)
+    loadProducerState(logEndOffset, reloadFromCleanShutdown = hadCleanShutdown)
   }
 
-  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
+  def dir: File = _dir
+
+  def parentDir: String = _parentDir
+
+  def parentDirFile: File = new File(_parentDir)
+
+  def initFileSize: Int = {
+    if (config.preallocate)
+      config.segmentSize
+    else
+      0
   }
 
-  newGauge("NumLogSegments",
-    new Gauge[Int] {
-      def value = numberOfSegments
-    },
-    tags)
+  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()
+  }
 
-  newGauge("LogStartOffset",
-    new Gauge[Long] {
-      def value = logStartOffset
-    },
-    tags)
+  private def checkIfMemoryMappedBufferClosed(): Unit = {
+    if (isMemoryMappedBufferClosed)
+      throw new KafkaStorageException(s"The memory mapped buffer for log of $topicPartition is already closed")
+  }
 
-  newGauge("LogEndOffset",
-    new Gauge[Long] {
-      def value = logEndOffset
-    },
-    tags)
+  def highWatermark: Long = highWatermarkMetadata.messageOffset
 
-  newGauge("Size",
-    new Gauge[Long] {
-      def value = size
-    },
-    tags)
+  /**
+   * 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
+  }
+
+  def updateHighWatermarkOffsetMetadata(hw: LogOffsetMetadata): Long = {
+    val newHighWatermark = if (hw.messageOffset < logStartOffset) {
+      updateHighWatermarkMetadata(LogOffsetMetadata(logStartOffset))
+      logStartOffset
+    } else if (hw.messageOffset > logEndOffset) {
+      updateHighWatermarkMetadata(logEndOffsetMetadata)
+      logEndOffset
+    } else {
+      updateHighWatermarkMetadata(hw)
+      hw.messageOffset
+    }
+    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")
 
-  scheduler.schedule(name = "PeriodicProducerExpirationCheck", fun = () => {
+    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
+  }
+
+  newGauge(LogMetricNames.NumLogSegments, () => numberOfSegments, tags)
+  newGauge(LogMetricNames.LogStartOffset, () => logStartOffset, tags)
+  newGauge(LogMetricNames.LogEndOffset, () => logEndOffset, tags)
+  newGauge(LogMetricNames.Size, () => size, tags)
+
+  val producerExpireCheck = scheduler.schedule(name = "PeriodicProducerExpirationCheck", fun = () => {
     lock synchronized {
       producerStateManager.removeExpiredProducers(time.milliseconds)
     }
@@ -256,46 +535,105 @@ class Log(@volatile var dir: File,
   /** The name of this log */
   def name  = dir.getName()
 
-  def leaderEpochCache = _leaderEpochCache
+  private def recordVersion: RecordVersion = config.messageFormatVersion.recordVersion
 
-  private def initializeLeaderEpochCache(): LeaderEpochCache = {
-    // create the log directory if it doesn't exist
-    Files.createDirectories(dir.toPath)
-    new LeaderEpochFileCache(topicPartition, () => logEndOffsetMetadata,
-      new LeaderEpochCheckpointFile(LeaderEpochFile.newFile(dir), logDirFailureChannel))
+  private def initializeLeaderEpochCache(): Unit = lock synchronized {
+    val leaderEpochFile = LeaderEpochCheckpointFile.newFile(dir)
+
+    def newLeaderEpochFileCache(): LeaderEpochFileCache = {
+      val checkpointFile = new LeaderEpochCheckpointFile(leaderEpochFile, logDirFailureChannel)
+      new LeaderEpochFileCache(topicPartition, () => logEndOffset, checkpointFile)
+    }
+
+    if (recordVersion.precedes(RecordVersion.V2)) {
+      val currentCache = if (leaderEpochFile.exists())
+        Some(newLeaderEpochFileCache())
+      else
+        None
+
+      if (currentCache.exists(_.nonEmpty))
+        warn(s"Deleting non-empty leader epoch cache due to incompatible message format $recordVersion")
+
+      Files.deleteIfExists(leaderEpochFile.toPath)
+      leaderEpochCache = None
+    } else {
+      leaderEpochCache = Some(newLeaderEpochFileCache())
+    }
   }
 
+  /**
+   * Removes any temporary files found in log directory, and creates a list of all .swap files which could be swapped
+   * in place of existing segment(s). For log splitting, we know that any .swap file whose base offset is higher than
+   * the smallest offset .clean file could be part of an incomplete split operation. Such .swap files are also deleted
+   * by this method.
+   * @return Set of .swap files that are valid to be swapped in as segment files
+   */
   private def removeTempFilesAndCollectSwapFiles(): Set[File] = {
-    var swapFiles = Set[File]()
+
+    def deleteIndicesIfExist(baseFile: File, suffix: String = ""): Unit = {
+      info(s"Deleting index files with suffix $suffix for baseFile $baseFile")
+      val offset = offsetFromFile(baseFile)
+      Files.deleteIfExists(Log.offsetIndexFile(dir, offset, suffix).toPath)
+      Files.deleteIfExists(Log.timeIndexFile(dir, offset, suffix).toPath)
+      Files.deleteIfExists(Log.transactionIndexFile(dir, offset, suffix).toPath)
+    }
+
+    val swapFiles = mutable.Set[File]()
+    val cleanFiles = mutable.Set[File]()
+    var minCleanedFileOffset = Long.MaxValue
 
     for (file <- dir.listFiles if file.isFile) {
       if (!file.canRead)
-        throw new IOException("Could not read file " + file)
+        throw new IOException(s"Could not read file $file")
       val filename = file.getName
-      if (filename.endsWith(DeletedFileSuffix) || filename.endsWith(CleanedFileSuffix)) {
-        // if the file ends in .deleted or .cleaned, delete it
+      if (filename.endsWith(DeletedFileSuffix)) {
+        debug(s"Deleting stray temporary file ${file.getAbsolutePath}")
         Files.deleteIfExists(file.toPath)
-      } else if(filename.endsWith(SwapFileSuffix)) {
+      } else if (filename.endsWith(CleanedFileSuffix)) {
+        minCleanedFileOffset = Math.min(offsetFromFileName(filename), minCleanedFileOffset)
+        cleanFiles += file
+      } else if (filename.endsWith(SwapFileSuffix)) {
         // we crashed in the middle of a swap operation, to recover:
-        // if a log, delete the .index file, complete the swap operation later
-        // if an index just delete it, it will be rebuilt
+        // if a log, delete the index files, complete the swap operation later
+        // if an index just delete the index files, they will be rebuilt
         val baseFile = new File(CoreUtils.replaceSuffix(file.getPath, SwapFileSuffix, ""))
+        info(s"Found file ${file.getAbsolutePath} from interrupted swap operation.")
         if (isIndexFile(baseFile)) {
-          Files.deleteIfExists(file.toPath)
+          deleteIndicesIfExist(baseFile)
         } else if (isLogFile(baseFile)) {
-          // delete the index files
-          val offset = offsetFromFile(baseFile)
-          Files.deleteIfExists(Log.offsetIndexFile(dir, offset).toPath)
-          Files.deleteIfExists(Log.timeIndexFile(dir, offset).toPath)
-          Files.deleteIfExists(Log.transactionIndexFile(dir, offset).toPath)
+          deleteIndicesIfExist(baseFile)
           swapFiles += file
         }
       }
     }
-    swapFiles
+
+    // KAFKA-6264: Delete all .swap files whose base offset is greater than the minimum .cleaned segment offset. Such .swap
+    // files could be part of an incomplete split operation that could not complete. See Log#splitOverflowedSegment
+    // for more details about the split operation.
+    val (invalidSwapFiles, validSwapFiles) = swapFiles.partition(file => offsetFromFile(file) >= minCleanedFileOffset)
+    invalidSwapFiles.foreach { file =>
+      debug(s"Deleting invalid swap file ${file.getAbsoluteFile} minCleanedFileOffset: $minCleanedFileOffset")
+      val baseFile = new File(CoreUtils.replaceSuffix(file.getPath, SwapFileSuffix, ""))
+      deleteIndicesIfExist(baseFile, SwapFileSuffix)
+      Files.deleteIfExists(file.toPath)
+    }
+
+    // Now that we have deleted all .swap files that constitute an incomplete split operation, let's delete all .clean files
+    cleanFiles.foreach { file =>
+      debug(s"Deleting stray .clean file ${file.getAbsolutePath}")
+      Files.deleteIfExists(file.toPath)
+    }
+
+    validSwapFiles
   }
 
-  // This method does not need to convert IOException to KafkaStorageException because it is only called before all logs are loaded
+  /**
+   * This method does not need to convert IOException to KafkaStorageException because it is only called before all logs are loaded
+   * It is possible that we encounter a segment with index offset overflow in which case the LogSegmentOffsetOverflowException
+   * will be thrown. Note that any segments that were opened before we encountered the exception will remain open and the
+   * caller is responsible for closing them appropriately, if needed.
+   * @throws LogSegmentOffsetOverflowException if the log directory contains a segment with messages that overflow the index offset
+   */
   private def loadSegmentFiles(): Unit = {
     // load segments in ascending order because transactional data from one segment may depend on the
     // segments that come before it
@@ -305,182 +643,242 @@ class Log(@volatile var dir: File,
         val offset = offsetFromFile(file)
         val logFile = Log.logFile(dir, offset)
         if (!logFile.exists) {
-          warn("Found an orphaned index file, %s, with no corresponding log file.".format(file.getAbsolutePath))
+          warn(s"Found an orphaned index file ${file.getAbsolutePath}, with no corresponding log file.")
           Files.deleteIfExists(file.toPath)
         }
       } else if (isLogFile(file)) {
         // if it's a log file, load the corresponding log segment
-        val startOffset = offsetFromFile(file)
-        val indexFile = Log.offsetIndexFile(dir, startOffset)
-        val timeIndexFile = Log.timeIndexFile(dir, startOffset)
-        val txnIndexFile = Log.transactionIndexFile(dir, startOffset)
-
-        val indexFileExists = indexFile.exists()
-        val timeIndexFileExists = timeIndexFile.exists()
-        val segment = new LogSegment(dir = dir,
-          startOffset = startOffset,
-          indexIntervalBytes = config.indexInterval,
-          maxIndexSize = config.maxIndexSize,
-          rollJitterMs = config.randomSegmentJitter,
+        val baseOffset = offsetFromFile(file)
+        val timeIndexFileNewlyCreated = !Log.timeIndexFile(dir, baseOffset).exists()
+        val segment = LogSegment.open(dir = dir,
+          baseOffset = baseOffset,
+          config,
           time = time,
           fileAlreadyExists = true)
 
-        if (indexFileExists) {
-          try {
-            segment.index.sanityCheck()
-            // Resize the time index file to 0 if it is newly created.
-            if (!timeIndexFileExists)
-              segment.timeIndex.resize(0)
-            segment.timeIndex.sanityCheck()
-            segment.txnIndex.sanityCheck()
-          } catch {
-            case e: java.lang.IllegalArgumentException =>
-              warn(s"Found a corrupted index file due to ${e.getMessage}}. deleting ${timeIndexFile.getAbsolutePath}, " +
-                s"${indexFile.getAbsolutePath}, and ${txnIndexFile.getAbsolutePath} and rebuilding index...")
-              Files.deleteIfExists(timeIndexFile.toPath)
-              Files.delete(indexFile.toPath)
-              segment.txnIndex.delete()
-              recoverSegment(segment)
-          }
-        } else {
-          error("Could not find offset index file corresponding to log file %s, rebuilding index...".format(segment.log.file.getAbsolutePath))
-          recoverSegment(segment)
+        try segment.sanityCheck(timeIndexFileNewlyCreated)
+        catch {
+          case _: NoSuchFileException =>
+            error(s"Could not find offset index file corresponding to log file ${segment.log.file.getAbsolutePath}, " +
+              "recovering segment and rebuilding index files...")
+            recoverSegment(segment)
+          case e: CorruptIndexException =>
+            warn(s"Found a corrupted index file corresponding to log file ${segment.log.file.getAbsolutePath} due " +
+              s"to ${e.getMessage}}, recovering segment and rebuilding index files...")
+            recoverSegment(segment)
         }
         addSegment(segment)
       }
     }
   }
 
-  private def recoverSegment(segment: LogSegment, leaderEpochCache: Option[LeaderEpochCache] = None): Int = lock synchronized {
-    val stateManager = new ProducerStateManager(topicPartition, dir, maxProducerIdExpirationMs)
-    stateManager.truncateAndReload(logStartOffset, segment.baseOffset, time.milliseconds)
-    logSegments(stateManager.mapEndOffset, segment.baseOffset).foreach { segment =>
-      val startOffset = math.max(segment.baseOffset, stateManager.mapEndOffset)
-      val fetchDataInfo = segment.read(startOffset, None, Int.MaxValue)
-      if (fetchDataInfo != null)
-        loadProducersFromLog(stateManager, fetchDataInfo.records)
-    }
-    stateManager.updateMapEndOffset(segment.baseOffset)
-
-    // take a snapshot for the first recovered segment to avoid reloading all the segments if we shutdown before we
-    // checkpoint the recovery point
-    stateManager.takeSnapshot()
-
-    val bytesTruncated = segment.recover(stateManager, leaderEpochCache)
-
+  /**
+   * Recover the given segment.
+   * @param segment Segment to recover
+   * @param leaderEpochCache Optional cache for updating the leader epoch during recovery
+   * @return The number of bytes truncated from the segment
+   * @throws LogSegmentOffsetOverflowException if the segment contains messages that cause index offset overflow
+   */
+  private def recoverSegment(segment: LogSegment,
+                             leaderEpochCache: Option[LeaderEpochFileCache] = None): Int = lock synchronized {
+    val producerStateManager = new ProducerStateManager(topicPartition, dir, maxProducerIdExpirationMs)
+    rebuildProducerState(segment.baseOffset, reloadFromCleanShutdown = false, producerStateManager)
+    val bytesTruncated = segment.recover(producerStateManager, leaderEpochCache)
     // once we have recovered the segment's data, take a snapshot to ensure that we won't
     // need to reload the same segment again while recovering another segment.
-    stateManager.takeSnapshot()
+    producerStateManager.takeSnapshot()
     bytesTruncated
   }
 
-  // This method does not need to convert IOException to KafkaStorageException because it is only called before all logs are loaded
+  /**
+   * This method does not need to convert IOException to KafkaStorageException because it is only called before all logs
+   * are loaded.
+   * @throws LogSegmentOffsetOverflowException if the swap file contains messages that cause the log segment offset to
+   *                                           overflow. Note that this is currently a fatal exception as we do not have
+   *                                           a way to deal with it. The exception is propagated all the way up to
+   *                                           KafkaServer#startup which will cause the broker to shut down if we are in
+   *                                           this situation. This is expected to be an extremely rare scenario in practice,
+   *                                           and manual intervention might be required to get out of it.
+   */
   private def completeSwapOperations(swapFiles: Set[File]): Unit = {
     for (swapFile <- swapFiles) {
       val logFile = new File(CoreUtils.replaceSuffix(swapFile.getPath, SwapFileSuffix, ""))
-      val startOffset = offsetFromFile(logFile)
-      val indexFile = new File(CoreUtils.replaceSuffix(logFile.getPath, LogFileSuffix, IndexFileSuffix) + SwapFileSuffix)
-      val index =  new OffsetIndex(indexFile, baseOffset = startOffset, maxIndexSize = config.maxIndexSize)
-      val timeIndexFile = new File(CoreUtils.replaceSuffix(logFile.getPath, LogFileSuffix, TimeIndexFileSuffix) + SwapFileSuffix)
-      val timeIndex = new TimeIndex(timeIndexFile, baseOffset = startOffset, maxIndexSize = config.maxIndexSize)
-      val txnIndexFile = new File(CoreUtils.replaceSuffix(logFile.getPath, LogFileSuffix, TxnIndexFileSuffix) + SwapFileSuffix)
-      val txnIndex = new TransactionIndex(startOffset, txnIndexFile)
-      val swapSegment = new LogSegment(FileRecords.open(swapFile),
-        index = index,
-        timeIndex = timeIndex,
-        txnIndex = txnIndex,
-        baseOffset = startOffset,
-        indexIntervalBytes = config.indexInterval,
-        rollJitterMs = config.randomSegmentJitter,
-        time = time)
-      info("Found log file %s from interrupted swap operation, repairing.".format(swapFile.getPath))
+      val baseOffset = offsetFromFile(logFile)
+      val swapSegment = LogSegment.open(swapFile.getParentFile,
+        baseOffset = baseOffset,
+        config,
+        time = time,
+        fileSuffix = SwapFileSuffix)
+      info(s"Found log file ${swapFile.getPath} from interrupted swap operation, repairing.")
       recoverSegment(swapSegment)
-      val oldSegments = logSegments(swapSegment.baseOffset, swapSegment.nextOffset())
-      replaceSegments(swapSegment, oldSegments.toSeq, isRecoveredSwapFile = true)
+
+      // We create swap files for two cases:
+      // (1) Log cleaning where multiple segments are merged into one, and
+      // (2) Log splitting where one segment is split into multiple.
+      //
+      // Both of these mean that the resultant swap segments be composed of the original set, i.e. the swap segment
+      // must fall within the range of existing segment(s). If we cannot find such a segment, it means the deletion
+      // of that segment was successful. In such an event, we should simply rename the .swap to .log without having to
+      // do a replace with an existing segment.
+      val oldSegments = logSegments(swapSegment.baseOffset, swapSegment.readNextOffset).filter { segment =>
+        segment.readNextOffset > swapSegment.baseOffset
+      }
+      replaceSegments(Seq(swapSegment), oldSegments.toSeq, isRecoveredSwapFile = true)
     }
   }
 
-  // Load the log segments from the log files on disk and return the next offset
-  // This method does not need to convert IOException to KafkaStorageException because it is only called before all logs are loaded
+  /**
+   * Load the log segments from the log files on disk and return the next offset.
+   * This method does not need to convert IOException to KafkaStorageException because it is only called before all logs
+   * are loaded.
+   * @throws LogSegmentOffsetOverflowException if we encounter a .swap file with messages that overflow index offset; or when
+   *                                           we find an unexpected number of .log files with overflow
+   */
   private def loadSegments(): Long = {
     // first do a pass through the files in the log directory and remove any temporary files
     // and find any interrupted swap operations
     val swapFiles = removeTempFilesAndCollectSwapFiles()
 
-    // now do a second pass and load all the log and index files
-    loadSegmentFiles()
+    // Now do a second pass and load all the log and index files.
+    // We might encounter legacy log segments with offset overflow (KAFKA-6264). We need to split such segments. When
+    // this happens, restart loading segment files from scratch.
+    retryOnOffsetOverflow {
+      // In case we encounter a segment with offset overflow, the retry logic will split it after which we need to retry
+      // loading of segments. In that case, we also need to close all segments that could have been left open in previous
+      // call to loadSegmentFiles().
+      logSegments.foreach(_.close())
+      segments.clear()
+      loadSegmentFiles()
+    }
 
     // Finally, complete any interrupted swap operations. To be crash-safe,
     // log files that are replaced by the swap segment should be renamed to .deleted
     // 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(new LogSegment(dir = dir,
-        startOffset = 0,
-        indexIntervalBytes = config.indexInterval,
-        maxIndexSize = config.maxIndexSize,
-        rollJitterMs = config.randomSegmentJitter,
-        time = time,
-        fileAlreadyExists = false,
-        initFileSize = this.initFileSize,
-        preallocate = config.preallocate))
-      0
-    } else if (!dir.getAbsolutePath.endsWith(Log.DeleteDirSuffix)) {
-      val nextOffset = recoverLog()
+    if (!dir.getAbsolutePath.endsWith(Log.DeleteDirSuffix)) {
+      val nextOffset = retryOnOffsetOverflow {
+        recoverLog()
+      }
+
       // reset the index size of the currently active log segment to allow more entries
-      activeSegment.index.resize(config.maxIndexSize)
-      activeSegment.timeIndex.resize(config.maxIndexSize)
+      activeSegment.resizeIndexes(config.maxIndexSize)
       nextOffset
-    } else 0
+    } else {
+       if (logSegments.isEmpty) {
+          addSegment(LogSegment.open(dir = dir,
+            baseOffset = 0,
+            config,
+            time = time,
+            initFileSize = this.initFileSize))
+       }
+      0
+    }
   }
 
-  private def updateLogEndOffset(messageOffset: Long) {
-    nextOffsetMetadata = new LogOffsetMetadata(messageOffset, activeSegment.baseOffset, activeSegment.size)
+  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 updateLogStartOffset(offset: Long): Unit = {
+    logStartOffset = offset
+
+    if (highWatermark < offset) {
+      updateHighWatermark(offset)
+    }
+
+    if (this.recoveryPoint < offset) {
+      this.recoveryPoint = offset
+    }
   }
 
   /**
    * Recover the log segments and return the next offset after recovery.
-   *
    * This method does not need to convert IOException to KafkaStorageException because it is only called before all
    * logs are loaded.
+   * @throws LogSegmentOffsetOverflowException if we encountered a legacy segment with offset overflow
    */
-  private def recoverLog(): Long = {
+  private[log] def recoverLog(): Long = {
     // if we have the clean shutdown marker, skip recovery
-    if (!hasCleanShutdownFile) {
-      // okay we need to actually recovery this log
+    if (!hadCleanShutdown) {
+      // okay we need to actually recover this log
       val unflushed = logSegments(this.recoveryPoint, Long.MaxValue).iterator
-      while (unflushed.hasNext) {
-        val segment = unflushed.next
-        info("Recovering unflushed segment %d in log %s.".format(segment.baseOffset, name))
+      var truncated = false
+
+      while (unflushed.hasNext && !truncated) {
+        val segment = unflushed.next()
+        info(s"Recovering unflushed segment ${segment.baseOffset}")
         val truncatedBytes =
           try {
-            recoverSegment(segment, Some(_leaderEpochCache))
+            recoverSegment(segment, leaderEpochCache)
           } catch {
             case _: InvalidOffsetException =>
               val startOffset = segment.baseOffset
-              warn("Found invalid offset during recovery for log " + dir.getName + ". Deleting the corrupt segment and " +
-                "creating an empty one with starting offset " + startOffset)
+              warn("Found invalid offset during recovery. Deleting the corrupt segment and " +
+                s"creating an empty one with starting offset $startOffset")
               segment.truncateTo(startOffset)
           }
         if (truncatedBytes > 0) {
           // we had an invalid message, delete all remaining log
-          warn("Corruption found in segment %d of log %s, truncating to offset %d.".format(segment.baseOffset, name,
-            segment.nextOffset()))
-          unflushed.foreach(deleteSegment)
+          warn(s"Corruption found in segment ${segment.baseOffset}, truncating to offset ${segment.readNextOffset}")
+          removeAndDeleteSegments(unflushed.toList,
+            asyncDelete = true,
+            reason = LogRecovery)
+          truncated = true
         }
       }
     }
-    recoveryPoint = activeSegment.nextOffset
+
+    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,
+          reason = LogRecovery)
+      }
+    }
+
+    if (logSegments.isEmpty) {
+      // no existing segments, create a new mutable segment beginning at logStartOffset
+      addSegment(LogSegment.open(dir = dir,
+        baseOffset = logStartOffset,
+        config,
+        time = time,
+        initFileSize = this.initFileSize,
+        preallocate = config.preallocate))
+    }
+
+    recoveryPoint = activeSegment.readNextOffset
     recoveryPoint
   }
 
-  private def loadProducerState(lastOffset: Long, reloadFromCleanShutdown: Boolean): Unit = lock synchronized {
+  // Rebuild producer state until lastOffset. This method may be called from the recovery code path, and thus must be
+  // free of all side-effects, i.e. it must not update any log-specific state.
+  private def rebuildProducerState(lastOffset: Long,
+                                   reloadFromCleanShutdown: Boolean,
+                                   producerStateManager: ProducerStateManager): Unit = lock synchronized {
     checkIfMemoryMappedBufferClosed()
-    val messageFormatVersion = config.messageFormatVersion.messageFormatVersion
-    info(s"Loading producer state from offset $lastOffset for partition $topicPartition with message " +
-      s"format version $messageFormatVersion")
+    val segments = logSegments
+    val offsetsToSnapshot =
+      if (segments.nonEmpty) {
+        val nextLatestSegmentBaseOffset = lowerSegment(segments.last.baseOffset).map(_.baseOffset)
+        Seq(nextLatestSegmentBaseOffset, Some(segments.last.baseOffset), Some(lastOffset))
+      } else {
+        Seq(Some(lastOffset))
+      }
+    info(s"Loading producer state till offset $lastOffset with message format version ${recordVersion.value}")
 
     // We want to avoid unnecessary scanning of the log to build the producer state when the broker is being
     // upgraded. The basic idea is to use the absence of producer snapshot files to detect the upgrade case,
@@ -494,13 +892,11 @@ class Log(@volatile var dir: File,
     // offset (see below). The next time the log is reloaded, we will load producer state using this snapshot
     // (or later snapshots). Otherwise, if there is no snapshot file, then we have to rebuild producer state
     // from the first segment.
-
-    if (producerStateManager.latestSnapshotOffset.isEmpty && (messageFormatVersion < RecordBatch.MAGIC_VALUE_V2 || reloadFromCleanShutdown)) {
+    if (recordVersion.value < RecordBatch.MAGIC_VALUE_V2 ||
+        (producerStateManager.latestSnapshotOffset.isEmpty && reloadFromCleanShutdown)) {
       // To avoid an expensive scan through all of the segments, we take empty snapshots from the start of the
       // last two segments and the last offset. This should avoid the full scan in the case that the log needs
       // truncation.
-      val nextLatestSegmentBaseOffset = lowerSegment(activeSegment.baseOffset).map(_.baseOffset)
-      val offsetsToSnapshot = Seq(nextLatestSegmentBaseOffset, Some(activeSegment.baseOffset), Some(lastOffset))
       offsetsToSnapshot.flatten.foreach { offset =>
         producerStateManager.updateMapEndOffset(offset)
         producerStateManager.takeSnapshot()
@@ -516,33 +912,39 @@ 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)
-          producerStateManager.takeSnapshot()
 
-          val fetchDataInfo = segment.read(startOffset, Some(lastOffset), Int.MaxValue)
+          if (offsetsToSnapshot.contains(Some(segment.baseOffset)))
+            producerStateManager.takeSnapshot()
+
+          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)
+            loadProducersFromRecords(producerStateManager, fetchDataInfo.records)
         }
       }
-
       producerStateManager.updateMapEndOffset(lastOffset)
-      updateFirstUnstableOffset()
+      producerStateManager.takeSnapshot()
     }
   }
 
-  private def loadProducersFromLog(producerStateManager: ProducerStateManager, records: Records): Unit = {
-    val loadedProducers = mutable.Map.empty[Long, ProducerAppendInfo]
-    val completedTxns = ListBuffer.empty[CompletedTxn]
-    records.batches.asScala.foreach { batch =>
-      if (batch.hasProducerId) {
-        val maybeCompletedTxn = updateProducers(batch, loadedProducers, isFromClient = false)
-        maybeCompletedTxn.foreach(completedTxns += _)
-      }
-    }
-    loadedProducers.values.foreach(producerStateManager.update)
-    completedTxns.foreach(producerStateManager.completeTxn)
+  private def loadProducerState(lastOffset: Long, reloadFromCleanShutdown: Boolean): Unit = lock synchronized {
+    rebuildProducerState(lastOffset, reloadFromCleanShutdown, producerStateManager)
+    maybeIncrementFirstUnstableOffset()
   }
 
   private[log] def activeProducersWithLastSequence: Map[Long, Int] = lock synchronized {
@@ -551,10 +953,13 @@ class Log(@volatile var dir: File,
     }
   }
 
-  /**
-   * Check if we have the "clean shutdown" file
-   */
-  private def hasCleanShutdownFile: Boolean = new File(dir.getParentFile, CleanShutdownFile).exists()
+  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
+    }
+  }
 
   /**
    * The number of segments in the log.
@@ -566,10 +971,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() {
-    debug(s"Closing log $name")
+  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
@@ -581,22 +987,23 @@ class Log(@volatile var dir: File,
   }
 
   /**
-    * Rename the directory of the log
-    *
-    * @throws KafkaStorageException if rename fails
-    */
-  def renameDir(name: String) {
+   * Rename the directory of the log
+   *
+   * @throws KafkaStorageException if rename fails
+   */
+  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))
-          producerStateManager.logDir = dir
+          _dir = renamedDir
+          _parentDir = renamedDir.getParent
+          logSegments.foreach(_.updateParentDir(renamedDir))
+          producerStateManager.updateParentDir(dir)
           // re-initialize leader epoch cache so that LeaderEpochCheckpointFile.checkpoint can correctly reference
           // the checkpoint file in renamed log directory
-          _leaderEpochCache = initializeLeaderEpochCache()
+          initializeLeaderEpochCache()
         }
       }
     }
@@ -605,8 +1012,8 @@ 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() {
-    debug(s"Closing handlers of log $name")
+  def closeHandlers(): Unit = {
+    debug("Closing handlers")
     lock synchronized {
       logSegments.foreach(_.closeHandlers())
       isMemoryMappedBufferClosed = true
@@ -617,12 +1024,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): LogAppendInfo = {
-    append(records, isFromClient, assignOffsets = true, leaderEpoch)
+  def appendAsLeader(records: MemoryRecords,
+                     leaderEpoch: Int,
+                     origin: AppendOrigin = AppendOrigin.Client,
+                     interBrokerProtocolVersion: ApiVersion = ApiVersion.latestVersion): LogAppendInfo = {
+    append(records, origin, interBrokerProtocolVersion, assignOffsets = true, leaderEpoch, ignoreRecordSize = false)
   }
 
   /**
@@ -633,7 +1044,13 @@ 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, assignOffsets = false, leaderEpoch = -1)
+    append(records,
+      origin = AppendOrigin.Replication,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      assignOffsets = false,
+      leaderEpoch = -1,
+      // disable to check the validation of record size since the record is already accepted by leader.
+      ignoreRecordSize = true)
   }
 
   /**
@@ -643,210 +1060,283 @@ 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
+   * @param ignoreRecordSize true to skip validation of record size.
    * @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, isFromClient: Boolean, assignOffsets: Boolean, leaderEpoch: Int): LogAppendInfo = {
+  private def append(records: MemoryRecords,
+                     origin: AppendOrigin,
+                     interBrokerProtocolVersion: ApiVersion,
+                     assignOffsets: Boolean,
+                     leaderEpoch: Int,
+                     ignoreRecordSize: Boolean): LogAppendInfo = {
     maybeHandleIOException(s"Error while appending records to $topicPartition in dir ${dir.getParent}") {
-      val appendInfo = analyzeAndValidateRecords(records, isFromClient = isFromClient)
+      val appendInfo = analyzeAndValidateRecords(records, origin, ignoreRecordSize)
 
       // return if we have no valid messages or if this is a duplicate of the last appended entry
-      if (appendInfo.shallowCount == 0)
-        return appendInfo
+      if (appendInfo.shallowCount == 0) appendInfo
+      else {
 
-      // trim any invalid bytes or partial messages before appending it to the on-disk log
-      var validRecords = trimInvalidBytes(records, appendInfo)
+        // trim any invalid bytes or partial messages before appending it to the on-disk log
+        var validRecords = trimInvalidBytes(records, appendInfo)
 
-      // they are valid, insert them in the log
-      lock synchronized {
-        checkIfMemoryMappedBufferClosed()
-        if (assignOffsets) {
-          // assign offsets to the message set
-          val offset = new LongRef(nextOffsetMetadata.messageOffset)
-          appendInfo.firstOffset = offset.value
-          val now = time.milliseconds
-          val validateAndOffsetAssignResult = try {
-            LogValidator.validateMessagesAndAssignOffsets(validRecords,
-              offset,
-              time,
-              now,
-              appendInfo.sourceCodec,
-              appendInfo.targetCodec,
-              config.compact,
-              config.messageFormatVersion.messageFormatVersion,
-              config.messageTimestampType,
-              config.messageTimestampDifferenceMaxMs,
-              leaderEpoch,
-              isFromClient)
-          } catch {
-            case e: IOException => throw new KafkaException("Error in validating messages while appending to log '%s'".format(name), e)
-          }
-          validRecords = validateAndOffsetAssignResult.validatedRecords
-          appendInfo.maxTimestamp = validateAndOffsetAssignResult.maxTimestamp
-          appendInfo.offsetOfMaxTimestamp = validateAndOffsetAssignResult.shallowOffsetOfMaxTimestamp
-          appendInfo.lastOffset = offset.value - 1
-          appendInfo.recordsProcessingStats = validateAndOffsetAssignResult.recordsProcessingStats
-          if (config.messageTimestampType == TimestampType.LOG_APPEND_TIME)
-            appendInfo.logAppendTime = now
-
-          // re-validate message sizes if there's a possibility that they have changed (due to re-compression or message
-          // format conversion)
-          if (validateAndOffsetAssignResult.messageSizeMaybeChanged) {
-            for (batch <- validRecords.batches.asScala) {
-              if (batch.sizeInBytes > config.maxMessageSize) {
-                // we record the original message set size instead of the trimmed size
-                // to be consistent with pre-compression bytesRejectedRate recording
-                brokerTopicStats.topicStats(topicPartition.topic).bytesRejectedRate.mark(records.sizeInBytes)
-                brokerTopicStats.allTopicsStats.bytesRejectedRate.mark(records.sizeInBytes)
-                throw new RecordTooLargeException("Message batch size is %d bytes which exceeds the maximum configured size of %d."
-                  .format(batch.sizeInBytes, config.maxMessageSize))
+        // they are valid, insert them in the log
+        lock synchronized {
+          checkIfMemoryMappedBufferClosed()
+          if (assignOffsets) {
+            // assign offsets to the message set
+            val offset = new LongRef(nextOffsetMetadata.messageOffset)
+            appendInfo.firstOffset = Some(offset.value)
+            val now = time.milliseconds
+            val validateAndOffsetAssignResult = try {
+              LogValidator.validateMessagesAndAssignOffsets(validRecords,
+                topicPartition,
+                offset,
+                time,
+                now,
+                appendInfo.sourceCodec,
+                appendInfo.targetCodec,
+                config.compact,
+                config.messageFormatVersion.recordVersion.value,
+                config.messageTimestampType,
+                config.messageTimestampDifferenceMaxMs,
+                leaderEpoch,
+                origin,
+                interBrokerProtocolVersion,
+                brokerTopicStats)
+            } catch {
+              case e: IOException =>
+                throw new KafkaException(s"Error validating messages while appending to log $name", e)
+            }
+            validRecords = validateAndOffsetAssignResult.validatedRecords
+            appendInfo.maxTimestamp = validateAndOffsetAssignResult.maxTimestamp
+            appendInfo.offsetOfMaxTimestamp = validateAndOffsetAssignResult.shallowOffsetOfMaxTimestamp
+            appendInfo.lastOffset = offset.value - 1
+            appendInfo.recordConversionStats = validateAndOffsetAssignResult.recordConversionStats
+            if (config.messageTimestampType == TimestampType.LOG_APPEND_TIME)
+              appendInfo.logAppendTime = now
+
+            // re-validate message sizes if there's a possibility that they have changed (due to re-compression or message
+            // format conversion)
+            if (!ignoreRecordSize && validateAndOffsetAssignResult.messageSizeMaybeChanged) {
+              validRecords.batches.forEach { batch =>
+                if (batch.sizeInBytes > config.maxMessageSize) {
+                  // we record the original message set size instead of the trimmed size
+                  // to be consistent with pre-compression bytesRejectedRate recording
+                  brokerTopicStats.topicStats(topicPartition.topic).bytesRejectedRate.mark(records.sizeInBytes)
+                  brokerTopicStats.allTopicsStats.bytesRejectedRate.mark(records.sizeInBytes)
+                  throw new RecordTooLargeException(s"Message batch size is ${batch.sizeInBytes} bytes in append to" +
+                    s"partition $topicPartition which exceeds the maximum configured size of ${config.maxMessageSize}.")
+                }
               }
             }
-          }
-        } else {
-          // we are taking the offsets we are given
-          if (!appendInfo.offsetsMonotonic || appendInfo.firstOffset < nextOffsetMetadata.messageOffset)
-            throw new IllegalArgumentException("Out of order offsets found in " + records.records.asScala.map(_.offset))
-        }
-
-        // 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)
-            _leaderEpochCache.assign(batch.partitionLeaderEpoch, batch.baseOffset)
-        }
-
-        // check messages set size may be exceed config.segmentSize
-        if (validRecords.sizeInBytes > config.segmentSize) {
-          throw new RecordBatchTooLargeException("Message batch size is %d bytes which exceeds the maximum configured segment size of %d."
-            .format(validRecords.sizeInBytes, config.segmentSize))
-        }
-
-        // 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)
-        maybeDuplicate.foreach { duplicate =>
-          appendInfo.firstOffset = duplicate.firstOffset
-          appendInfo.lastOffset = duplicate.lastOffset
-          appendInfo.logAppendTime = duplicate.timestamp
-          appendInfo.logStartOffset = logStartOffset
-          return appendInfo
-        }
+          } else {
+            // we are taking the offsets we are given
+            if (!appendInfo.offsetsMonotonic)
+              throw new OffsetsOutOfOrderException(s"Out of order offsets found in append to $topicPartition: " +
+                records.records.asScala.map(_.offset))
+
+            if (appendInfo.firstOrLastOffsetOfFirstBatch < nextOffsetMetadata.messageOffset) {
+              // we may still be able to recover if the log is empty
+              // one example: fetching from log start offset on the leader which is not batch aligned,
+              // which may happen as a result of AdminClient#deleteRecords()
+              val firstOffset = appendInfo.firstOffset match {
+                case Some(offset) => offset
+                case None => records.batches.asScala.head.baseOffset()
+              }
 
-        // maybe roll the log if this segment is full
-        val segment = maybeRoll(messagesSize = validRecords.sizeInBytes,
-          maxTimestampInMessages = appendInfo.maxTimestamp,
-          maxOffsetInMessages = appendInfo.lastOffset)
-
-        val logOffsetMetadata = LogOffsetMetadata(
-          messageOffset = appendInfo.firstOffset,
-          segmentBaseOffset = segment.baseOffset,
-          relativePositionInSegment = segment.size)
-
-        segment.append(firstOffset = appendInfo.firstOffset,
-          largestOffset = appendInfo.lastOffset,
-          largestTimestamp = appendInfo.maxTimestamp,
-          shallowOffsetOfMaxTimestamp = appendInfo.offsetOfMaxTimestamp,
-          records = validRecords)
-
-        // update the producer state
-        for ((producerId, producerAppendInfo) <- updatedProducers) {
-          producerAppendInfo.maybeCacheTxnFirstOffsetMetadata(logOffsetMetadata)
-          producerStateManager.update(producerAppendInfo)
-        }
+              val firstOrLast = if (appendInfo.firstOffset.isDefined) "First offset" else "Last offset of the first batch"
+              throw new UnexpectedAppendOffsetException(
+                s"Unexpected offset in append to $topicPartition. $firstOrLast " +
+                  s"${appendInfo.firstOrLastOffsetOfFirstBatch} is less than the next offset ${nextOffsetMetadata.messageOffset}. " +
+                  s"First 10 offsets in append: ${records.records.asScala.take(10).map(_.offset)}, last offset in" +
+                  s" append: ${appendInfo.lastOffset}. Log start offset = $logStartOffset",
+                firstOffset, appendInfo.lastOffset)
+            }
+          }
 
-        // 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)
-          segment.updateTxnIndex(completedTxn, lastStableOffset)
-        }
+          // update the epoch cache with the epoch stamped onto the message by the leader
+          validRecords.batches.forEach { batch =>
+            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()
+              }
+            }
+          }
 
-        // 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)
+          // check messages set size may be exceed config.segmentSize
+          if (validRecords.sizeInBytes > config.segmentSize) {
+            throw new RecordBatchTooLargeException(s"Message batch size is ${validRecords.sizeInBytes} bytes in append " +
+              s"to partition $topicPartition, which exceeds the maximum configured segment size of ${config.segmentSize}.")
+          }
 
-        // increment the log end offset
-        updateLogEndOffset(appendInfo.lastOffset + 1)
+          // 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(
+            logOffsetMetadata, validRecords, origin)
+
+          maybeDuplicate match {
+            case Some(duplicate) =>
+              appendInfo.firstOffset = Some(duplicate.firstOffset)
+              appendInfo.lastOffset = duplicate.lastOffset
+              appendInfo.logAppendTime = duplicate.timestamp
+              appendInfo.logStartOffset = logStartOffset
+            case None =>
+              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
+              updatedProducers.values.foreach(producerAppendInfo => 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.
+              completedTxns.foreach { completedTxn =>
+                val lastStableOffset = producerStateManager.lastStableOffset(completedTxn)
+                segment.updateTxnIndex(completedTxn, lastStableOffset)
+                producerStateManager.completeTxn(completedTxn)
+              }
 
-        // update the first unstable offset (which is used to compute LSO)
-        updateFirstUnstableOffset()
+              // 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)
 
-        trace("Appended message set to log %s with first offset: %d, next offset: %d, and messages: %s"
-          .format(this.name, appendInfo.firstOffset, nextOffsetMetadata.messageOffset, validRecords))
+              // update the first unstable offset (which is used to compute LSO)
+              maybeIncrementFirstUnstableOffset()
 
-        if (unflushedMessages >= config.flushInterval)
-          flush()
+              trace(s"Appended message set with last offset: ${appendInfo.lastOffset}, " +
+                s"first offset: ${appendInfo.firstOffset}, " +
+                s"next offset: ${nextOffsetMetadata.messageOffset}, " +
+                s"and messages: $validRecords")
 
-        appendInfo
+              if (unflushedMessages >= config.flushInterval) flush()
+          }
+          appendInfo
+        }
       }
     }
   }
 
-  def onHighWatermarkIncremented(highWatermark: Long): Unit = {
-    lock synchronized {
-      replicaHighWatermark = Some(highWatermark)
-      producerStateManager.onHighWatermarkUpdated(highWatermark)
-      updateFirstUnstableOffset()
+  def maybeAssignEpochStartOffset(leaderEpoch: Int, startOffset: Long): Unit = {
+    leaderEpochCache.foreach { cache =>
+      cache.assign(leaderEpoch, startOffset)
     }
   }
 
-  private def updateFirstUnstableOffset(): Unit = lock synchronized {
+  def latestEpoch: Option[Int] = leaderEpochCache.flatMap(_.latestEpoch)
+
+  def endOffsetForEpoch(leaderEpoch: Int): Option[OffsetAndEpoch] = {
+    leaderEpochCache.flatMap { cache =>
+      val (foundEpoch, foundOffset) = cache.endOffsetFor(leaderEpoch)
+      if (foundOffset == UNDEFINED_EPOCH_OFFSET)
+        None
+      else
+        Some(OffsetAndEpoch(foundOffset, foundEpoch))
+    }
+  }
+
+  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) {
-      debug(s"First unstable offset for ${this.name} updated to $updatedFirstStableOffset")
-      this.firstUnstableOffset = updatedFirstStableOffset
+    if (updatedFirstStableOffset != this.firstUnstableOffsetMetadata) {
+      debug(s"First unstable offset updated to $updatedFirstStableOffset")
+      this.firstUnstableOffsetMetadata = updatedFirstStableOffset
     }
   }
 
   /**
    * Increment the log start offset if the provided offset is larger.
    */
-  def maybeIncrementLogStartOffset(newLogStartOffset: Long) {
+  def maybeIncrementLogStartOffset(newLogStartOffset: Long, reason: LogStartOffsetIncrementReason): Unit = {
     // 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.
     maybeHandleIOException(s"Exception while increasing log start offset for $topicPartition to $newLogStartOffset in dir ${dir.getParent}") {
       lock synchronized {
+        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")
+
         checkIfMemoryMappedBufferClosed()
         if (newLogStartOffset > logStartOffset) {
-          info(s"Incrementing log start offset of partition $topicPartition to $newLogStartOffset in dir ${dir.getParent}")
-          logStartOffset = newLogStartOffset
-          _leaderEpochCache.clearAndFlushEarliest(logStartOffset)
-          producerStateManager.truncateHead(logStartOffset)
-          updateFirstUnstableOffset()
+          updateLogStartOffset(newLogStartOffset)
+          info(s"Incremented log start offset to $newLogStartOffset due to $reason")
+          leaderEpochCache.foreach(_.truncateFromStart(logStartOffset))
+          producerStateManager.onLogStartOffsetIncremented(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 upto 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(_.duplicateOf(batch)).foreach { duplicate =>
-          return (updatedProducers, completedTxns.toList, Some(duplicate))
+    var relativePositionInSegment = appendOffsetMetadata.relativePositionInSegment
+
+    records.batches.forEach { batch =>
+      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))
+          }
         }
 
-      val maybeCompletedTxn = updateProducers(batch, updatedProducers, isFromClient = isFromClient)
-      maybeCompletedTxn.foreach(completedTxns += _)
+        // 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(producerStateManager, batch, updatedProducers, firstOffsetMetadata, origin)
+        maybeCompletedTxn.foreach(completedTxns += _)
+      }
+
+      relativePositionInSegment += batch.sizeInBytes
     }
     (updatedProducers, completedTxns.toList, None)
   }
@@ -855,7 +1345,7 @@ class Log(@volatile var dir: File,
    * Validate the following:
    * 
                
    * 
              1.  each message matches its CRC
-   * 
              2.  each message size is valid
+   * 
              3.  each message size is valid (if ignoreRecordSize is false)
    * 
              4.  that the sequence numbers of the incoming record batches are consistent with the existing state and with each other.
    * 
              
    *
@@ -869,28 +1359,39 @@ 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,
+                                        ignoreRecordSize: Boolean): LogAppendInfo = {
     var shallowMessageCount = 0
     var validBytesCount = 0
-    var firstOffset = -1L
+    var firstOffset: Option[Long] = None
     var lastOffset = -1L
+    var lastLeaderEpoch = RecordBatch.NO_PARTITION_LEADER_EPOCH
     var sourceCodec: CompressionCodec = NoCompressionCodec
     var monotonic = true
     var maxTimestamp = RecordBatch.NO_TIMESTAMP
     var offsetOfMaxTimestamp = -1L
+    var readFirstMessage = false
+    var lastOffsetOfFirstBatch = -1L
 
-    for (batch <- records.batches.asScala) {
+    records.batches.forEach { batch =>
       // 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 should be 0, but it is ${batch.baseOffset}")
+      if (batch.magic >= RecordBatch.MAGIC_VALUE_V2 && origin == AppendOrigin.Client && 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.
       // When appending to the leader, we will update LogAppendInfo.baseOffset with the correct value. In the follower
       // case, validation will be more lenient.
-      if (firstOffset < 0)
-        firstOffset = if (batch.magic >= RecordBatch.MAGIC_VALUE_V2) batch.baseOffset else batch.lastOffset
+      // Also indicate whether we have the accurate first offset or not
+      if (!readFirstMessage) {
+        if (batch.magic >= RecordBatch.MAGIC_VALUE_V2)
+          firstOffset = Some(batch.baseOffset)
+        lastOffsetOfFirstBatch = batch.lastOffset
+        readFirstMessage = true
+      }
 
       // check that offsets are monotonically increasing
       if (lastOffset >= batch.lastOffset)
@@ -898,18 +1399,22 @@ class Log(@volatile var dir: File,
 
       // update the last offset seen
       lastOffset = batch.lastOffset
+      lastLeaderEpoch = batch.partitionLeaderEpoch
 
       // Check if the message sizes are valid.
       val batchSize = batch.sizeInBytes
-      if (batchSize > config.maxMessageSize) {
+      if (!ignoreRecordSize && batchSize > config.maxMessageSize) {
         brokerTopicStats.topicStats(topicPartition.topic).bytesRejectedRate.mark(records.sizeInBytes)
         brokerTopicStats.allTopicsStats.bytesRejectedRate.mark(records.sizeInBytes)
-        throw new RecordTooLargeException(s"The record batch size is $batchSize bytes which exceeds the maximum configured " +
-          s"value of ${config.maxMessageSize}.")
+        throw new RecordTooLargeException(s"The record batch size in the append to $topicPartition is $batchSize bytes " +
+          s"which exceeds the maximum configured value of ${config.maxMessageSize}.")
       }
 
       // 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
@@ -926,16 +1431,12 @@ class Log(@volatile var dir: File,
 
     // Apply broker-side compression if any
     val targetCodec = BrokerCompressionCodec.getTargetCompressionCodec(config.compressionType, sourceCodec)
-    LogAppendInfo(firstOffset, lastOffset, maxTimestamp, offsetOfMaxTimestamp, RecordBatch.NO_TIMESTAMP, logStartOffset,
-      RecordsProcessingStats.EMPTY, sourceCodec, targetCodec, shallowMessageCount, validBytesCount, monotonic)
-  }
-
-  private def updateProducers(batch: RecordBatch,
-                              producers: mutable.Map[Long, ProducerAppendInfo],
-                              isFromClient: Boolean): Option[CompletedTxn] = {
-    val producerId = batch.producerId
-    val appendInfo = producers.getOrElseUpdate(producerId, producerStateManager.prepareUpdate(producerId, isFromClient))
-    appendInfo.append(batch)
+    val lastLeaderEpochOpt: Option[Int] = if (lastLeaderEpoch != RecordBatch.NO_PARTITION_LEADER_EPOCH)
+      Some(lastLeaderEpoch)
+    else
+      None
+    LogAppendInfo(firstOffset, lastOffset, lastLeaderEpochOpt, maxTimestamp, offsetOfMaxTimestamp, RecordBatch.NO_TIMESTAMP, logStartOffset,
+      RecordConversionStats.EMPTY, sourceCodec, targetCodec, shallowMessageCount, validBytesCount, monotonic, lastOffsetOfFirstBatch)
   }
 
   /**
@@ -948,7 +1449,8 @@ class Log(@volatile var dir: File,
   private def trimInvalidBytes(records: MemoryRecords, info: LogAppendInfo): MemoryRecords = {
     val validBytes = info.validBytes
     if (validBytes < 0)
-      throw new CorruptRecordException("Illegal length of message set " + validBytes + " Message set cannot be appended to log. Possible causes are corrupted produce requests")
+      throw new CorruptRecordException(s"Cannot append record batch with illegal length $validBytes to " +
+        s"log for $topicPartition. A possible cause is a corrupted produce request.")
     if (validBytes == records.sizeInBytes) {
       records
     } else {
@@ -959,9 +1461,15 @@ class Log(@volatile var dir: File,
     }
   }
 
-  private[log] def readUncommitted(startOffset: Long, maxLength: Int, maxOffset: Option[Long] = None,
-                                   minOneMessage: Boolean = false): FetchDataInfo = {
-    read(startOffset, maxLength, maxOffset, minOneMessage, isolationLevel = IsolationLevel.READ_UNCOMMITTED)
+  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)
   }
 
   /**
@@ -969,79 +1477,73 @@ class Log(@volatile var dir: File,
    *
    * @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 isolationLevel The isolation level of the fetcher. The READ_UNCOMMITTED isolation level has the traditional
-   *                       read semantics (e.g. consumers are limited to fetching up to the high watermark). In
-   *                       READ_COMMITTED, consumers are limited to fetching up to the last stable offset. Additionally,
-   *                       in READ_COMMITTED, the transaction index is consulted after fetching to collect the list
-   *                       of aborted transactions in the fetch range which the consumer uses to filter the fetched
-   *                       records before they are returned to the user. Note that fetches from followers always use
-   *                       READ_UNCOMMITTED.
    * @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] = None, minOneMessage: Boolean = false,
-           isolationLevel: IsolationLevel): FetchDataInfo = {
+  def read(startOffset: Long,
+           maxLength: Int,
+           isolation: FetchIsolation,
+           minOneMessage: Boolean): FetchDataInfo = {
     maybeHandleIOException(s"Exception while reading from $topicPartition in dir ${dir.getParent}") {
-      trace("Reading %d bytes from offset %d in log %s of length %d bytes".format(maxLength, startOffset, name, size))
+      trace(s"Reading maximum $maxLength bytes at offset $startOffset from log with " +
+        s"total 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 (isolationLevel == IsolationLevel.READ_COMMITTED) 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)
-        throw new OffsetOutOfRangeException("Request for offset %d but we only have log segments in the range %d to %d.".format(startOffset, logStartOffset, next))
-
-      // 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
-      // continue to read from successive segments until we get some messages or we reach the end of the log
-      while (segmentEntry != null) {
-        val segment = segmentEntry.getValue
+      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 $endOffset.")
+
+      val maxOffsetMetadata = isolation match {
+        case FetchLogEnd => endOffsetMetadata
+        case FetchHighWatermark => fetchHighWatermarkMetadata
+        case FetchTxnCommitted => fetchLastStableOffsetMetadata
+      }
 
-        // 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
-          } else {
-            segment.size
-          }
+      if (startOffset == maxOffsetMetadata.messageOffset)
+        emptyFetchDataInfo(maxOffsetMetadata, includeAbortedTxns)
+      else if (startOffset > maxOffsetMetadata.messageOffset)
+        emptyFetchDataInfo(convertToOffsetMetadataOrThrow(startOffset), includeAbortedTxns)
+      else {
+        // 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
+        // continue to read from successive segments until we get some messages or we reach the end of the log
+        var done = segmentEntry == null
+        var fetchDataInfo: FetchDataInfo = null
+        while (!done) {
+          val segment = segmentEntry.getValue
+
+          val maxPosition =
+            // 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
+
+          fetchDataInfo = segment.read(startOffset, maxLength, maxPosition, minOneMessage)
+          if (fetchDataInfo != null) {
+            if (includeAbortedTxns)
+              fetchDataInfo = addAbortedTransactions(startOffset, segmentEntry, fetchDataInfo)
+          } else segmentEntry = segments.higherEntry(segmentEntry.getKey)
+
+          done = fetchDataInfo != null || segmentEntry == null
         }
-        val fetchInfo = segment.read(startOffset, maxOffset, maxLength, maxPosition, minOneMessage)
-        if (fetchInfo == null) {
-          segmentEntry = segments.higherEntry(segmentEntry.getKey)
-        } else {
-          return isolationLevel match {
-            case IsolationLevel.READ_UNCOMMITTED => fetchInfo
-            case IsolationLevel.READ_COMMITTED => addAbortedTransactions(startOffset, segmentEntry, fetchInfo)
-          }
+
+        if (fetchDataInfo != null) fetchDataInfo
+        else {
+          // okay we are beyond the end of the last segment with no data fetched although the start offset is in range,
+          // this can happen when all messages with offset larger than start offsets have been deleted.
+          // In this case, we will return the empty set with log end offset metadata
+          FetchDataInfo(nextOffsetMetadata, MemoryRecords.EMPTY)
         }
       }
-
-      // okay we are beyond the end of the last segment with no data fetched although the start offset is in range,
-      // this can happen when all messages with offset larger than start offsets have been deleted.
-      // In this case, we will return the empty set with log end offset metadata
-      FetchDataInfo(nextOffsetMetadata, MemoryRecords.EMPTY)
     }
   }
 
@@ -1103,7 +1605,7 @@ class Log(@volatile var dir: File,
    * @return The offset of the first message whose timestamp is greater than or equals to the given timestamp.
    *         None if no such message is found.
    */
-  def fetchOffsetsByTimestamp(targetTimestamp: Long): Option[TimestampOffset] = {
+  def fetchOffsetByTimestamp(targetTimestamp: Long): Option[TimestampAndOffset] = {
     maybeHandleIOException(s"Error while fetching offset by timestamp for $topicPartition in dir ${dir.getParent}") {
       debug(s"Searching offset for timestamp $targetTimestamp")
 
@@ -1118,36 +1620,83 @@ class Log(@volatile var dir: File,
       // constant time access while being safe to use with concurrent collections unlike `toArray`.
       val segmentsCopy = logSegments.toBuffer
       // For the earliest and latest, we do not need to return the timestamp.
-      if (targetTimestamp == ListOffsetRequest.EARLIEST_TIMESTAMP)
-        return Some(TimestampOffset(RecordBatch.NO_TIMESTAMP, logStartOffset))
-      else if (targetTimestamp == ListOffsetRequest.LATEST_TIMESTAMP)
-        return Some(TimestampOffset(RecordBatch.NO_TIMESTAMP, logEndOffset))
-
-      val targetSeg = {
-        // Get all the segments whose largest timestamp is smaller than target timestamp
-        val earlierSegs = segmentsCopy.takeWhile(_.largestTimestamp < targetTimestamp)
-        // We need to search the first segment whose largest timestamp is greater than the target timestamp if there is one.
-        if (earlierSegs.length < segmentsCopy.length)
-          Some(segmentsCopy(earlierSegs.length))
-        else
-          None
+      if (targetTimestamp == ListOffsetRequest.EARLIEST_TIMESTAMP) {
+        // The first cached epoch usually corresponds to the log start offset, but we have to verify this since
+        // it may not be true following a message format version bump as the epoch will not be available for
+        // log entries written in the older format.
+        val earliestEpochEntry = leaderEpochCache.flatMap(_.earliestEntry)
+        val epochOpt = earliestEpochEntry match {
+          case Some(entry) if entry.startOffset <= logStartOffset => Optional.of[Integer](entry.epoch)
+          case _ => Optional.empty[Integer]()
+        }
+        Some(new TimestampAndOffset(RecordBatch.NO_TIMESTAMP, logStartOffset, epochOpt))
+      } else if (targetTimestamp == ListOffsetRequest.LATEST_TIMESTAMP) {
+        val latestEpochOpt = leaderEpochCache.flatMap(_.latestEpoch).map(_.asInstanceOf[Integer])
+        val epochOptional = Optional.ofNullable(latestEpochOpt.orNull)
+        Some(new TimestampAndOffset(RecordBatch.NO_TIMESTAMP, logEndOffset, epochOptional))
+      } else {
+        // We need to search the first segment whose largest timestamp is >= the target timestamp if there is one.
+        val targetSeg = segmentsCopy.find(_.largestTimestamp >= targetTimestamp)
+        targetSeg.flatMap(_.findOffsetByTimestamp(targetTimestamp, logStartOffset))
       }
+    }
+  }
+
+  def legacyFetchOffsetsBefore(timestamp: Long, maxNumOffsets: Int): Seq[Long] = {
+    // Cache to avoid race conditions. `toBuffer` is faster than most alternatives and provides
+    // constant time access while being safe to use with concurrent collections unlike `toArray`.
+    val segments = logSegments.toBuffer
+    val lastSegmentHasSize = segments.last.size > 0
+
+    val offsetTimeArray =
+      if (lastSegmentHasSize)
+        new Array[(Long, Long)](segments.length + 1)
+      else
+        new Array[(Long, Long)](segments.length)
+
+    for (i <- segments.indices)
+      offsetTimeArray(i) = (math.max(segments(i).baseOffset, logStartOffset), segments(i).lastModified)
+    if (lastSegmentHasSize)
+      offsetTimeArray(segments.length) = (logEndOffset, time.milliseconds)
+
+    var startIndex = -1
+    timestamp match {
+      case ListOffsetRequest.LATEST_TIMESTAMP =>
+        startIndex = offsetTimeArray.length - 1
+      case ListOffsetRequest.EARLIEST_TIMESTAMP =>
+        startIndex = 0
+      case _ =>
+        var isFound = false
+        debug("Offset time array = " + offsetTimeArray.foreach(o => "%d, %d".format(o._1, o._2)))
+        startIndex = offsetTimeArray.length - 1
+        while (startIndex >= 0 && !isFound) {
+          if (offsetTimeArray(startIndex)._2 <= timestamp)
+            isFound = true
+          else
+            startIndex -= 1
+        }
+    }
 
-      targetSeg.flatMap(_.findOffsetByTimestamp(targetTimestamp, logStartOffset))
+    val retSize = maxNumOffsets.min(startIndex + 1)
+    val ret = new Array[Long](retSize)
+    for (j <- 0 until retSize) {
+      ret(j) = offsetTimeArray(startIndex)._1
+      startIndex -= 1
     }
+    // ensure that the returned seq is in descending order of offsets
+    ret.toSeq.sortBy(-_)
   }
 
   /**
-   * Given a message offset, find its corresponding offset metadata in the log.
-   * If the message offset is out of range, return unknown offset metadata
-   */
-  def convertToOffsetMetadata(offset: Long): LogOffsetMetadata = {
-    try {
-      val fetchDataInfo = readUncommitted(offset, 1)
-      fetchDataInfo.fetchOffsetMetadata
-    } catch {
-      case _: OffsetOutOfRangeException => LogOffsetMetadata.UnknownOffsetMetadata
-    }
+    * 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
   }
 
   /**
@@ -1158,16 +1707,18 @@ class Log(@volatile var dir: File,
    *                  (if there is one) and returns true iff it is deletable
    * @return The number of segments deleted
    */
-  private def deleteOldSegments(predicate: (LogSegment, Option[LogSegment]) => Boolean, reason: String): Int = {
+  private def deleteOldSegments(predicate: (LogSegment, Option[LogSegment]) => Boolean,
+                                reason: SegmentDeletionReason): Int = {
     lock synchronized {
       val deletable = deletableSegments(predicate)
       if (deletable.nonEmpty)
-        info(s"Found deletable segments with base offsets [${deletable.map(_.baseOffset).mkString(",")}] due to $reason")
-      deleteSegments(deletable)
+        deleteSegments(deletable, reason)
+      else
+        0
     }
   }
 
-  private def deleteSegments(deletable: Iterable[LogSegment]): Int = {
+  private def deleteSegments(deletable: Iterable[LogSegment], reason: SegmentDeletionReason): Int = {
     maybeHandleIOException(s"Error while deleting segments for $topicPartition in dir ${dir.getParent}") {
       val numToDelete = deletable.size
       if (numToDelete > 0) {
@@ -1177,8 +1728,8 @@ class Log(@volatile var dir: File,
         lock synchronized {
           checkIfMemoryMappedBufferClosed()
           // remove the segments for lookups
-          deletable.foreach(deleteSegment)
-          maybeIncrementLogStartOffset(segments.firstEntry.getValue.baseOffset)
+          removeAndDeleteSegments(deletable, asyncDelete = true, reason)
+          maybeIncrementLogStartOffset(segments.firstEntry.getValue.baseOffset, SegmentDeletion)
         }
       }
       numToDelete
@@ -1198,10 +1749,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) {
@@ -1224,25 +1774,34 @@ class Log(@volatile var dir: File,
   }
 
   /**
-   * Delete any log segments that have either expired due to time based retention
-   * or because the log size is > retentionSize
+   * If topic deletion is enabled, delete any log segments that have either expired due to time based retention
+   * or because the log size is > retentionSize.
+   *
+   * Whether or not deletion is enabled, delete any log segments that are before the log start offset
    */
   def deleteOldSegments(): Int = {
-    if (!config.delete) return 0
-    deleteRetentionMsBreachedSegments() + deleteRetentionSizeBreachedSegments() + deleteLogStartOffsetBreachedSegments()
+    if (config.delete) {
+      deleteRetentionMsBreachedSegments() + deleteRetentionSizeBreachedSegments() + deleteLogStartOffsetBreachedSegments()
+    } else {
+      deleteLogStartOffsetBreachedSegments()
+    }
   }
 
   private def deleteRetentionMsBreachedSegments(): Int = {
     if (config.retentionMs < 0) return 0
     val startMs = time.milliseconds
-    deleteOldSegments((segment, _) => startMs - segment.largestTimestamp > config.retentionMs,
-      reason = s"retention time ${config.retentionMs}ms breach")
+
+    def shouldDelete(segment: LogSegment, nextSegmentOpt: Option[LogSegment]): Boolean = {
+      startMs - segment.largestTimestamp > config.retentionMs
+    }
+
+    deleteOldSegments(shouldDelete, RetentionMsBreach)
   }
 
   private def deleteRetentionSizeBreachedSegments(): Int = {
     if (config.retentionSize < 0 || size < config.retentionSize) return 0
     var diff = size - config.retentionSize
-    def shouldDelete(segment: LogSegment, nextSegmentOpt: Option[LogSegment]) = {
+    def shouldDelete(segment: LogSegment, nextSegmentOpt: Option[LogSegment]): Boolean = {
       if (diff - segment.size >= 0) {
         diff -= segment.size
         true
@@ -1250,13 +1809,16 @@ class Log(@volatile var dir: File,
         false
       }
     }
-    deleteOldSegments(shouldDelete, reason = s"retention size in bytes ${config.retentionSize} breach")
+
+    deleteOldSegments(shouldDelete, RetentionSizeBreach)
   }
 
   private def deleteLogStartOffsetBreachedSegments(): Int = {
-    def shouldDelete(segment: LogSegment, nextSegmentOpt: Option[LogSegment]) =
+    def shouldDelete(segment: LogSegment, nextSegmentOpt: Option[LogSegment]): Boolean = {
       nextSegmentOpt.exists(_.baseOffset <= logStartOffset)
-    deleteOldSegments(shouldDelete, reason = s"log start offset $logStartOffset breach")
+    }
+
+    deleteOldSegments(shouldDelete, StartOffsetBreach)
   }
 
   def isFuture: Boolean = dir.getName.endsWith(Log.FutureDirSuffix)
@@ -1279,8 +1841,8 @@ class Log(@volatile var dir: File,
   /**
    * Roll the log over to a new empty log segment if necessary.
    *
-   * @param messagesSize The messages set size in bytes
-   * @param maxTimestampInMessages The maximum timestamp in the messages.
+   * @param messagesSize The messages set size in bytes.
+   * @param appendInfo log append information
    * logSegment will be rolled if one of the following conditions met
    * 
                
    * 
              1.  The logSegment is full
@@ -1290,17 +1852,19 @@ class Log(@volatile var dir: File,
    * 
              
    * @return The currently active segment after (perhaps) rolling to a new segment
    */
-  private def maybeRoll(messagesSize: Int, maxTimestampInMessages: Long, maxOffsetInMessages: Long): LogSegment = {
+  private def maybeRoll(messagesSize: Int, appendInfo: LogAppendInfo): LogSegment = {
     val segment = activeSegment
     val now = time.milliseconds
-    val reachedRollMs = segment.timeWaitedForRoll(now, maxTimestampInMessages) > config.segmentMs - segment.rollJitterMs
-    if (segment.size > config.segmentSize - messagesSize ||
-        (segment.size > 0 && reachedRollMs) ||
-        segment.index.isFull || segment.timeIndex.isFull || !segment.canConvertToRelativeOffset(maxOffsetInMessages)) {
-      debug(s"Rolling new log segment in $name (log_size = ${segment.size}/${config.segmentSize}}, " +
-          s"index_size = ${segment.index.entries}/${segment.index.maxEntries}, " +
-          s"time_index_size = ${segment.timeIndex.entries}/${segment.timeIndex.maxEntries}, " +
-          s"inactive_time_ms = ${segment.timeWaitedForRoll(now, maxTimestampInMessages)}/${config.segmentMs - segment.rollJitterMs}).")
+
+    val maxTimestampInMessages = appendInfo.maxTimestamp
+    val maxOffsetInMessages = appendInfo.lastOffset
+
+    if (segment.shouldRoll(RollParams(config, appendInfo, messagesSize, now))) {
+      debug(s"Rolling new log segment (log_size = ${segment.size}/${config.segmentSize}}, " +
+        s"offset_index_size = ${segment.offsetIndex.entries}/${segment.offsetIndex.maxEntries}, " +
+        s"time_index_size = ${segment.timeIndex.entries}/${segment.timeIndex.maxEntries}, " +
+        s"inactive_time_ms = ${segment.timeWaitedForRoll(now, maxTimestampInMessages)}/${config.segmentMs - segment.rollJitterMs}).")
+
       /*
         maxOffsetInMessages - Integer.MAX_VALUE is a heuristic value for the first offset in the set of messages.
         Since the offset in messages will not differ by more than Integer.MAX_VALUE, this is guaranteed <= the real
@@ -1310,8 +1874,13 @@ class Log(@volatile var dir: File,
         Integer.MAX_VALUE.toLong + 2 or more.  In this case, the prior behavior would roll a new log segment whose
         base offset was too low to contain the next message.  This edge case is possible when a replica is recovering a
         highly compacted topic from scratch.
-       */
-      roll(maxOffsetInMessages - Integer.MAX_VALUE)
+        Note that this is only required for pre-V2 message formats because these do not store the first message offset
+        in the header.
+      */
+      appendInfo.firstOffset match {
+        case Some(firstOffset) => roll(Some(firstOffset))
+        case None => roll(Some(maxOffsetInMessages - Integer.MAX_VALUE))
+      }
     } else {
       segment
     }
@@ -1323,27 +1892,44 @@ class Log(@volatile var dir: File,
    *
    * @return The newly rolled segment
    */
-  def roll(expectedNextOffset: Long = 0): LogSegment = {
+  def roll(expectedNextOffset: Option[Long] = None): LogSegment = {
     maybeHandleIOException(s"Error while rolling log segment for $topicPartition in dir ${dir.getParent}") {
-      val start = time.nanoseconds
+      val start = time.hiResClockMs()
       lock synchronized {
         checkIfMemoryMappedBufferClosed()
-        val newOffset = math.max(expectedNextOffset, logEndOffset)
+        val newOffset = math.max(expectedNextOffset.getOrElse(0L), logEndOffset)
         val logFile = Log.logFile(dir, newOffset)
-        val offsetIdxFile = offsetIndexFile(dir, newOffset)
-        val timeIdxFile = timeIndexFile(dir, newOffset)
-        val txnIdxFile = transactionIndexFile(dir, newOffset)
-        for (file <- List(logFile, offsetIdxFile, timeIdxFile, txnIdxFile) if file.exists) {
-          warn("Newly rolled segment file " + file.getName + " already exists; deleting it first")
-          file.delete()
-        }
 
-        Option(segments.lastEntry).foreach { entry =>
-          val seg = entry.getValue
-          seg.onBecomeInactiveSegment()
-          seg.index.trimToValidSize()
-          seg.timeIndex.trimToValidSize()
-          seg.log.trim()
+        if (segments.containsKey(newOffset)) {
+          // segment with the same base offset already exists and loaded
+          if (activeSegment.baseOffset == newOffset && activeSegment.size == 0) {
+            // We have seen this happen (see KAFKA-6388) after shouldRoll() returns true for an
+            // active segment of size zero because of one of the indexes is "full" (due to _maxEntries == 0).
+            warn(s"Trying to roll a new log segment with start offset $newOffset " +
+                 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}.")
+            removeAndDeleteSegments(Seq(activeSegment), asyncDelete = true, LogRoll)
+          } 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 " +
+                                     s"segment is ${segments.get(newOffset)}.")
+          }
+        } else if (!segments.isEmpty && newOffset < activeSegment.baseOffset) {
+          throw new KafkaException(
+            s"Trying to roll a new log segment for topic partition $topicPartition with " +
+            s"start offset $newOffset =max(provided offset = $expectedNextOffset, LEO = $logEndOffset) lower than start offset of the active segment $activeSegment")
+        } else {
+          val offsetIdxFile = offsetIndexFile(dir, newOffset)
+          val timeIdxFile = timeIndexFile(dir, newOffset)
+          val txnIdxFile = transactionIndexFile(dir, newOffset)
+
+          for (file <- List(logFile, offsetIdxFile, timeIdxFile, txnIdxFile) if file.exists) {
+            warn(s"Newly rolled segment file ${file.getAbsolutePath} already exists; deleting it first")
+            Files.delete(file.toPath)
+          }
+
+          Option(segments.lastEntry).foreach(_.getValue.onBecomeInactiveSegment())
         }
 
         // take a snapshot of the producer state to facilitate recovery. It is useful to have the snapshot
@@ -1354,25 +1940,22 @@ class Log(@volatile var dir: File,
         producerStateManager.updateMapEndOffset(newOffset)
         producerStateManager.takeSnapshot()
 
-        val segment = new LogSegment(dir,
-          startOffset = newOffset,
-          indexIntervalBytes = config.indexInterval,
-          maxIndexSize = config.maxIndexSize,
-          rollJitterMs = config.randomSegmentJitter,
+        val segment = LogSegment.open(dir,
+          baseOffset = newOffset,
+          config,
           time = time,
-          fileAlreadyExists = false,
           initFileSize = initFileSize,
           preallocate = config.preallocate)
-        val prev = addSegment(segment)
-        if (prev != null)
-          throw new KafkaException("Trying to roll a new log segment for topic partition %s with start offset %d while it already exists.".format(name, newOffset))
+        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)
 
-        info("Rolled new log segment for '" + name + "' in %.0f ms.".format((System.nanoTime - start) / (1000.0 * 1000.0)))
+        info(s"Rolled new log segment at offset $newOffset in ${time.hiResClockMs() - start} ms.")
 
         segment
       }
@@ -1382,7 +1965,7 @@ class Log(@volatile var dir: File,
   /**
    * The number of messages appended to the log since the last flush
    */
-  def unflushedMessages() = this.logEndOffset - this.recoveryPoint
+  def unflushedMessages: Long = this.logEndOffset - this.recoveryPoint
 
   /**
    * Flush all log segments
@@ -1394,72 +1977,37 @@ 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
-      debug("Flushing log '" + name + " up to offset " + offset + ", last flushed: " + lastFlushTime + " current time: " +
-        time.milliseconds + " unflushed = " + unflushedMessages)
-      for (segment <- logSegments(this.recoveryPoint, offset))
-        segment.flush()
+      if (offset > this.recoveryPoint) {
+        debug(s"Flushing log up to offset $offset, last flushed: $lastFlushTime,  current time: ${time.milliseconds()}, " +
+          s"unflushed: $unflushedMessages")
+        logSegments(this.recoveryPoint, offset).foreach(_.flush())
 
-      lock synchronized {
-        checkIfMemoryMappedBufferClosed()
-        if (offset > this.recoveryPoint) {
-          this.recoveryPoint = offset
-          lastFlushedTime.set(time.milliseconds)
+        lock synchronized {
+          checkIfMemoryMappedBufferClosed()
+          if (offset > this.recoveryPoint) {
+            this.recoveryPoint = offset
+            lastFlushedTime.set(time.milliseconds)
+          }
         }
       }
     }
   }
 
-  /**
-   * Cleanup old producer snapshots after the recovery point is checkpointed. It is useful to retain
-   * the snapshots from the recent segments in case we need to truncate and rebuild the producer state.
-   * Otherwise, we would always need to rebuild from the earliest segment.
-   *
-   * More specifically:
-   *
-   * 1. We always retain the producer snapshot from the last two segments. This solves the common case
-   * of truncating to an offset within the active segment, and the rarer case of truncating to the previous segment.
-   *
-   * 2. We only delete snapshots for offsets less than the recovery point. The recovery point is checkpointed
-   * periodically and it can be behind after a hard shutdown. Since recovery starts from the recovery point, the logic
-   * of rebuilding the producer snapshots in one pass and without loading older segments is simpler if we always
-   * have a producer snapshot for all segments being recovered.
-   *
-   * Return the minimum snapshots offset that was retained.
-   */
-  def deleteSnapshotsAfterRecoveryPointCheckpoint(): Long = {
-    val minOffsetToRetain = minSnapshotsOffsetToRetain
-    producerStateManager.deleteSnapshotsBefore(minOffsetToRetain)
-    minOffsetToRetain
-  }
-
-  // Visible for testing, see `deleteSnapshotsAfterRecoveryPointCheckpoint()` for details
-  private[log] def minSnapshotsOffsetToRetain: Long = {
-    lock synchronized {
-      val twoSegmentsMinOffset = lowerSegment(activeSegment.baseOffset).getOrElse(activeSegment).baseOffset
-      // Prefer segment base offset
-      val recoveryPointOffset = lowerSegment(recoveryPoint).map(_.baseOffset).getOrElse(recoveryPoint)
-      math.min(recoveryPointOffset, twoSegmentsMinOffset)
-    }
-  }
-
   private def lowerSegment(offset: Long): Option[LogSegment] =
     Option(segments.lowerEntry(offset)).map(_.getValue)
 
   /**
    * 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(_.delete())
-        segments.clear()
-        _leaderEpochCache.clear()
+        producerExpireCheck.cancel(true)
+        removeAndDeleteSegments(logSegments, asyncDelete = false, LogDeletion)
+        leaderEpochCache.foreach(_.clear())
         Utils.delete(dir)
         // File handlers will be closed if this log is deleted
         isMemoryMappedBufferClosed = true
@@ -1494,27 +2042,35 @@ class Log(@volatile var dir: File,
    * @param targetOffset The offset to truncate to, an upper bound on all offsets in the log after truncation is complete.
    * @return True iff targetOffset < logEndOffset
    */
-  private[log] def truncateTo(targetOffset: Long): Boolean = {
+  private[kafka] def truncateTo(targetOffset: Long): Boolean = {
     maybeHandleIOException(s"Error while truncating log to offset $targetOffset for $topicPartition in dir ${dir.getParent}") {
       if (targetOffset < 0)
-        throw new IllegalArgumentException("Cannot truncate to a negative offset (%d).".format(targetOffset))
+        throw new IllegalArgumentException(s"Cannot truncate partition $topicPartition to a negative offset (%d).".format(targetOffset))
       if (targetOffset >= logEndOffset) {
-        info("Truncating %s to %d has no effect as the largest offset in the log is %d.".format(name, targetOffset, logEndOffset - 1))
+        info(s"Truncating to $targetOffset has no effect as the largest offset in the log is ${logEndOffset - 1}")
+
+        // Always truncate epoch cache since we may have a conflicting epoch entry at the
+        // end of the log from the leader. This could happen if this broker was a leader
+        // and inserted the first start offset entry, but then failed to append any entries
+        // before another leader was elected.
+        lock synchronized {
+          leaderEpochCache.foreach(_.truncateFromEnd(logEndOffset))
+        }
+
         false
       } else {
-        info("Truncating log %s to offset %d.".format(name, targetOffset))
+        info(s"Truncating to offset $targetOffset")
         lock synchronized {
           checkIfMemoryMappedBufferClosed()
           if (segments.firstEntry.getValue.baseOffset > targetOffset) {
             truncateFullyAndStartAt(targetOffset)
           } else {
             val deletable = logSegments.filter(segment => segment.baseOffset > targetOffset)
-            deletable.foreach(deleteSegment)
+            removeAndDeleteSegments(deletable, asyncDelete = true, LogTruncation)
             activeSegment.truncateTo(targetOffset)
             updateLogEndOffset(targetOffset)
-            this.recoveryPoint = math.min(targetOffset, this.recoveryPoint)
-            this.logStartOffset = math.min(targetOffset, this.logStartOffset)
-            _leaderEpochCache.clearAndFlushLatest(targetOffset)
+            updateLogStartOffset(math.min(targetOffset, this.logStartOffset))
+            leaderEpochCache.foreach(_.truncateFromEnd(targetOffset))
             loadProducerState(targetOffset, reloadFromCleanShutdown = false)
           }
           true
@@ -1528,31 +2084,25 @@ 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 log '$name' at offset $newOffset")
+      debug(s"Truncate and start at offset $newOffset")
       lock synchronized {
         checkIfMemoryMappedBufferClosed()
-        val segmentsToDelete = logSegments.toList
-        segmentsToDelete.foreach(deleteSegment)
-        addSegment(new LogSegment(dir,
-          newOffset,
-          indexIntervalBytes = config.indexInterval,
-          maxIndexSize = config.maxIndexSize,
-          rollJitterMs = config.randomSegmentJitter,
+        removeAndDeleteSegments(logSegments, asyncDelete = true, LogTruncation)
+        addSegment(LogSegment.open(dir,
+          baseOffset = newOffset,
+          config = config,
           time = time,
-          fileAlreadyExists = false,
           initFileSize = initFileSize,
           preallocate = config.preallocate))
         updateLogEndOffset(newOffset)
-        _leaderEpochCache.clearAndFlush()
+        leaderEpochCache.foreach(_.clearAndFlush())
 
         producerStateManager.truncate()
         producerStateManager.updateMapEndOffset(newOffset)
-        updateFirstUnstableOffset()
-
-        this.recoveryPoint = math.min(newOffset, this.recoveryPoint)
-        this.logStartOffset = newOffset
+        maybeIncrementFirstUnstableOffset()
+        updateLogStartOffset(newOffset)
       }
     }
   }
@@ -1574,108 +2124,191 @@ 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 floor = segments.floorKey(from)
-      if (floor eq null)
-        segments.headMap(to).values.asScala
+      if (from > activeSegment.baseOffset)
+        Seq.empty
       else
-        segments.subMap(floor, true, to, false).values.asScala
+        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("Scheduling log segment %d for log %s for deletion.".format(segment.baseOffset, name))
-    lock synchronized {
-      segments.remove(segment.baseOffset)
-      asyncDeleteSegment(segment)
+  private def removeAndDeleteSegments(segments: Iterable[LogSegment],
+                                      asyncDelete: Boolean,
+                                      reason: SegmentDeletionReason): Unit = {
+    if (segments.nonEmpty) {
+      lock synchronized {
+        // 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
+        reason.logReason(this, toDelete)
+        toDelete.foreach { segment =>
+          this.segments.remove(segment.baseOffset)
+        }
+        deleteSegmentFiles(toDelete, asyncDelete)
+      }
     }
   }
 
   /**
-   * Perform an asynchronous delete on the given file if it exists (otherwise do nothing)
+   * 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.
    *
    * This method does not need to convert IOException (thrown from changeFileSuffixes) to KafkaStorageException because
    * it is either called before all logs are loaded or the caller will catch and handle IOException
    *
    * @throws IOException if the file can't be renamed and still exists
    */
-  private def asyncDeleteSegment(segment: LogSegment) {
-    segment.changeFileSuffixes("", Log.DeletedFileSuffix)
-    def deleteSeg() {
-      info("Deleting segment %d from log %s.".format(segment.baseOffset, name))
+  private def deleteSegmentFiles(segments: Iterable[LogSegment], asyncDelete: Boolean, deleteProducerStateSnapshots: Boolean = true): Unit = {
+    segments.foreach(_.changeFileSuffixes("", Log.DeletedFileSuffix))
+
+    def deleteSegments(): Unit = {
+      info(s"Deleting segment files ${segments.mkString(",")}")
       maybeHandleIOException(s"Error while deleting segments for $topicPartition in dir ${dir.getParent}") {
-        segment.delete()
+        segments.foreach { segment =>
+          segment.deleteIfExists()
+          if (deleteProducerStateSnapshots)
+            producerStateManager.removeAndDeleteSnapshot(segment.baseOffset)
+        }
       }
     }
-    scheduler.schedule("delete-file", deleteSeg _, delay = config.fileDeleteDelayMs)
+
+    if (asyncDelete)
+      scheduler.schedule("delete-file", () => deleteSegments(), delay = config.fileDeleteDelayMs)
+    else
+      deleteSegments()
   }
 
   /**
-   * Swap a new segment in place and delete one or more existing segments in a crash-safe manner. The old segments will
-   * be asynchronously deleted.
+   * Swap one or more new segment in place and delete one or more existing segments in a crash-safe manner. The old
+   * segments will be asynchronously deleted.
    *
    * This method does not need to convert IOException to KafkaStorageException because it is either called before all logs are loaded
    * or the caller will catch and handle IOException
    *
    * The sequence of operations is:
    * 
                
-   *   
              1.  Cleaner creates new segment with suffix .cleaned and invokes replaceSegments().
+   *   
              2.  Cleaner creates one or more new segments with suffix .cleaned and invokes replaceSegments().
    *        If broker crashes at this point, the clean-and-swap operation is aborted and
-   *        the .cleaned file is deleted on recovery in loadSegments().
-   *   
              3.  New segment is renamed .swap. If the broker crashes after this point before the whole
-   *        operation is completed, the swap operation is resumed on recovery as described in the next step.
+   *        the .cleaned files are deleted on recovery in loadSegments().
+   *   
              4.  New segments are renamed .swap. If the broker crashes before all segments were renamed to .swap, the
+   *        clean-and-swap operation is aborted - .cleaned as well as .swap files are deleted on recovery in
+   *        loadSegments(). We detect this situation by maintaining a specific order in which files are renamed from
+   *        .cleaned to .swap. Basically, files are renamed in descending order of offsets. On recovery, all .swap files
+   *        whose offset is greater than the minimum-offset .clean file are deleted.
+   *   
              5.  If the broker crashes after all new segments were renamed to .swap, the operation is completed, the swap
+   *        operation is resumed on recovery as described in the next step.
    *   
              6.  Old segment files are renamed to .deleted and asynchronous delete is scheduled.
    *        If the broker crashes, any .deleted files left behind are deleted on recovery in loadSegments().
    *        replaceSegments() is then invoked to complete the swap with newSegment recreated from
    *        the .swap file and oldSegments containing segments which were not renamed before the crash.
-   *   
              7.  Swap segment is renamed to replace the existing segment, completing this operation.
+   *   
              8.  Swap segment(s) are renamed to replace the existing segments, completing this operation.
    *        If the broker crashes, any .deleted files which may be left behind are deleted
    *        on recovery in loadSegments().
    * 
              
    *
-   * @param newSegment The new log segment to add to the log
+   * @param newSegments The new log segment to add to the log
    * @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(newSegment: 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
+      // but before this method is executed. We want to filter out those segments to avoid calling asyncDeleteSegment()
+      // multiple times for the same segment.
+      val sortedOldSegments = oldSegments.filter(seg => segments.containsKey(seg.baseOffset)).sortBy(_.baseOffset)
+
       checkIfMemoryMappedBufferClosed()
       // need to do this in two phases to be crash safe AND do the delete asynchronously
       // if we crash in the middle of this we complete the swap in loadSegments()
       if (!isRecoveredSwapFile)
-        newSegment.changeFileSuffixes(Log.CleanedFileSuffix, Log.SwapFileSuffix)
-      addSegment(newSegment)
+        sortedNewSegments.reverse.foreach(_.changeFileSuffixes(Log.CleanedFileSuffix, Log.SwapFileSuffix))
+      sortedNewSegments.reverse.foreach(addSegment(_))
+      val newSegmentBaseOffsets = sortedNewSegments.map(_.baseOffset).toSet
 
       // delete the old files
-      for (seg <- oldSegments) {
+      sortedOldSegments.foreach { seg =>
         // remove the index entry
-        if (seg.baseOffset != newSegment.baseOffset)
+        if (seg.baseOffset != sortedNewSegments.head.baseOffset)
           segments.remove(seg.baseOffset)
-        // delete segment
-        asyncDeleteSegment(seg)
+        // delete segment files, but do not delete producer state for segment objects which are being replaced.
+        deleteSegmentFiles(List(seg), asyncDelete = true, deleteProducerStateSnapshots = !newSegmentBaseOffsets.contains(seg.baseOffset))
       }
       // okay we are safe now, remove the swap suffix
-      newSegment.changeFileSuffixes(Log.SwapFileSuffix, "")
+      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()
     }
   }
 
@@ -1683,17 +2316,18 @@ class Log(@volatile var dir: File,
    * remove deleted log metrics
    */
   private[log] def removeLogMetrics(): Unit = {
-    removeMetric("NumLogSegments", tags)
-    removeMetric("LogStartOffset", tags)
-    removeMetric("LogEndOffset", tags)
-    removeMetric("Size", tags)
+    removeMetric(LogMetricNames.NumLogSegments, tags)
+    removeMetric(LogMetricNames.LogStartOffset, tags)
+    removeMetric(LogMetricNames.LogEndOffset, tags)
+    removeMetric(LogMetricNames.Size, tags)
   }
+
   /**
    * Add the given segment to the segments in this log. If this segment replaces an existing segment, delete it.
-   *
    * @param segment The segment to add
    */
-  def addSegment(segment: LogSegment) = this.segments.put(segment.baseOffset, segment)
+  @threadsafe
+  def addSegment(segment: LogSegment): LogSegment = this.segments.put(segment.baseOffset, segment)
 
   private def maybeHandleIOException[T](msg: => String)(fun: => T): T = {
     try {
@@ -1705,6 +2339,81 @@ class Log(@volatile var dir: File,
     }
   }
 
+  private[log] def retryOnOffsetOverflow[T](fn: => T): T = {
+    while (true) {
+      try {
+        return fn
+      } catch {
+        case e: LogSegmentOffsetOverflowException =>
+          info(s"Caught segment overflow error: ${e.getMessage}. Split segment and retry.")
+          splitOverflowedSegment(e.segment)
+      }
+    }
+    throw new IllegalStateException()
+  }
+
+  /**
+   * Split a segment into one or more segments such that there is no offset overflow in any of them. The
+   * resulting segments will contain the exact same messages that are present in the input segment. On successful
+   * completion of this method, the input segment will be deleted and will be replaced by the resulting new segments.
+   * See replaceSegments for recovery logic, in case the broker dies in the middle of this operation.
+   * 
              Note that this method assumes we have already determined that the segment passed in contains records that cause
+   * offset overflow.
              
+   * 
              The split logic overloads the use of .clean files that LogCleaner typically uses to make the process of replacing
+   * the input segment with multiple new segments atomic and recoverable in the event of a crash. See replaceSegments
+   * and completeSwapOperations for the implementation to make this operation recoverable on crashes.
              
+   * @param segment Segment to split
+   * @return List of new segments that replace the input segment
+   */
+  private[log] def splitOverflowedSegment(segment: LogSegment): List[LogSegment] = {
+    require(isLogFile(segment.log.file), s"Cannot split file ${segment.log.file.getAbsoluteFile}")
+    require(segment.hasOverflow, "Split operation is only permitted for segments with overflow")
+
+    info(s"Splitting overflowed segment $segment")
+
+    val newSegments = ListBuffer[LogSegment]()
+    try {
+      var position = 0
+      val sourceRecords = segment.log
+
+      while (position < sourceRecords.sizeInBytes) {
+        val firstBatch = sourceRecords.batchesFrom(position).asScala.head
+        val newSegment = LogCleaner.createNewCleanedSegment(this, firstBatch.baseOffset)
+        newSegments += newSegment
+
+        val bytesAppended = newSegment.appendFromFile(sourceRecords, position)
+        if (bytesAppended == 0)
+          throw new IllegalStateException(s"Failed to append records from position $position in $segment")
+
+        position += bytesAppended
+      }
+
+      // prepare new segments
+      var totalSizeOfNewSegments = 0
+      newSegments.foreach { splitSegment =>
+        splitSegment.onBecomeInactiveSegment()
+        splitSegment.flush()
+        splitSegment.lastModified = segment.lastModified
+        totalSizeOfNewSegments += splitSegment.log.sizeInBytes
+      }
+      // size of all the new segments combined must equal size of the original segment
+      if (totalSizeOfNewSegments != segment.log.sizeInBytes)
+        throw new IllegalStateException("Inconsistent segment sizes after split" +
+          s" before: ${segment.log.sizeInBytes} after: $totalSizeOfNewSegments")
+
+      // replace old segment with new ones
+      info(s"Replacing overflowed segment $segment with split segments $newSegments")
+      replaceSegments(newSegments.toList, List(segment))
+      newSegments.toList
+    } catch {
+      case e: Exception =>
+        newSegments.foreach { splitSegment =>
+          splitSegment.close()
+          splitSegment.deleteIfExists()
+        }
+        throw e
+    }
+  }
 }
 
 /**
@@ -1749,10 +2458,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,
@@ -1763,11 +2472,12 @@ object Log {
             time: Time = Time.SYSTEM,
             maxProducerIdExpirationMs: Int,
             producerIdExpirationCheckIntervalMs: Int,
-            logDirFailureChannel: LogDirFailureChannel): Log = {
+            logDirFailureChannel: LogDirFailureChannel,
+            lastShutdownClean: Boolean = true): Log = {
     val topicPartition = Log.parseTopicPartitionName(dir)
     val producerStateManager = new ProducerStateManager(topicPartition, dir, maxProducerIdExpirationMs)
     new Log(dir, config, logStartOffset, recoveryPoint, scheduler, brokerTopicStats, time, maxProducerIdExpirationMs,
-      producerIdExpirationCheckIntervalMs, topicPartition, producerStateManager, logDirFailureChannel)
+      producerIdExpirationCheckIntervalMs, topicPartition, producerStateManager, logDirFailureChannel, lastShutdownClean)
   }
 
   /**
@@ -1786,26 +2496,32 @@ object Log {
   }
 
   /**
-   * Construct a log file name in the given dir with the given base offset
+   * Construct a log file name in the given dir with the given base offset and the given suffix
    *
    * @param dir The directory in which the log will reside
    * @param offset The base offset of the log file
+   * @param suffix The suffix to be appended to the file name (e.g. "", ".deleted", ".cleaned", ".swap", etc.)
    */
-  def logFile(dir: File, offset: Long): File =
-    new File(dir, filenamePrefixFromOffset(offset) + LogFileSuffix)
+  def logFile(dir: File, offset: Long, suffix: String = ""): File =
+    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}"
   }
 
   /**
-    * Return a future directory name for the given topic partition. The name will be in the following
-    * format: topic-partition.uniqueId-future where topic, partition and uniqueId are variables.
-    */
+   * Return a future directory name for the given topic partition. The name will be in the following
+   * format: topic-partition.uniqueId-future where topic, partition and uniqueId are variables.
+   */
   def logFutureDirName(topicPartition: TopicPartition): String = {
     logDirNameWithSuffix(topicPartition, FutureDirSuffix)
   }
@@ -1816,30 +2532,35 @@ object Log {
   }
 
   /**
-    * Return a directory name for the given topic partition. The name will be in the following
-    * format: topic-partition where topic, partition are variables.
-    */
+   * Return a directory name for the given topic partition. The name will be in the following
+   * format: topic-partition where topic, partition are variables.
+   */
   def logDirName(topicPartition: TopicPartition): String = {
     s"${topicPartition.topic}-${topicPartition.partition}"
   }
 
   /**
-   * Construct an index file name in the given dir using the given base offset
+   * Construct an index file name in the given dir using the given base offset and the given suffix
    *
    * @param dir The directory in which the log will reside
    * @param offset The base offset of the log file
+   * @param suffix The suffix to be appended to the file name ("", ".deleted", ".cleaned", ".swap", etc.)
    */
-  def offsetIndexFile(dir: File, offset: Long): File =
-    new File(dir, filenamePrefixFromOffset(offset) + IndexFileSuffix)
+  def offsetIndexFile(dir: File, offset: Long, suffix: String = ""): File =
+    new File(dir, filenamePrefixFromOffset(offset) + IndexFileSuffix + suffix)
 
   /**
-   * Construct a time index file name in the given dir using the given base offset
+   * Construct a time index file name in the given dir using the given base offset and the given suffix
    *
    * @param dir The directory in which the log will reside
    * @param offset The base offset of the log file
+   * @param suffix The suffix to be appended to the file name ("", ".deleted", ".cleaned", ".swap", etc.)
    */
-  def timeIndexFile(dir: File, offset: Long): File =
-    new File(dir, filenamePrefixFromOffset(offset) + TimeIndexFileSuffix)
+  def timeIndexFile(dir: File, offset: Long, suffix: String = ""): File =
+    new File(dir, filenamePrefixFromOffset(offset) + TimeIndexFileSuffix + suffix)
+
+  def deleteFileIfExists(file: File, suffix: String = ""): Unit =
+    Files.deleteIfExists(new File(file.getPath + suffix).toPath)
 
   /**
    * Construct a producer id snapshot file using the given offset.
@@ -1850,20 +2571,30 @@ object Log {
   def producerSnapshotFile(dir: File, offset: Long): File =
     new File(dir, filenamePrefixFromOffset(offset) + ProducerSnapshotFileSuffix)
 
-  def transactionIndexFile(dir: File, offset: Long): File =
-    new File(dir, filenamePrefixFromOffset(offset) + TxnIndexFileSuffix)
+  /**
+   * Construct a transaction index file name in the given dir using the given base offset and the given suffix
+   *
+   * @param dir The directory in which the log will reside
+   * @param offset The base offset of the log file
+   * @param suffix The suffix to be appended to the file name ("", ".deleted", ".cleaned", ".swap", etc.)
+   */
+  def transactionIndexFile(dir: File, offset: Long, suffix: String = ""): File =
+    new File(dir, filenamePrefixFromOffset(offset) + TxnIndexFileSuffix + suffix)
 
-  def offsetFromFile(file: File): Long = {
-    val filename = file.getName
+  def offsetFromFileName(filename: String): Long = {
     filename.substring(0, filename.indexOf('.')).toLong
   }
 
+  def offsetFromFile(file: File): Long = {
+    offsetFromFileName(file.getName)
+  }
+
   /**
-    * Calculate a log's size (in bytes) based on its log segments
-    *
-    * @param segments The log segments to calculate the size of
-    * @return Sum of the log segments' sizes (in bytes)
-    */
+   * Calculate a log's size (in bytes) based on its log segments
+   *
+   * @param segments The log segments to calculate the size of
+   * @return Sum of the log segments' sizes (in bytes)
+   */
   def sizeInBytes(segments: Iterable[LogSegment]): Long =
     segments.map(_.size.toLong).sum
 
@@ -1912,4 +2643,104 @@ object Log {
   private def isLogFile(file: File): Boolean =
     file.getPath.endsWith(LogFileSuffix)
 
+  private def loadProducersFromRecords(producerStateManager: ProducerStateManager, records: Records): Unit = {
+    val loadedProducers = mutable.Map.empty[Long, ProducerAppendInfo]
+    val completedTxns = ListBuffer.empty[CompletedTxn]
+    records.batches.forEach { batch =>
+      if (batch.hasProducerId) {
+        val maybeCompletedTxn = updateProducers(
+          producerStateManager,
+          batch,
+          loadedProducers,
+          firstOffsetMetadata = None,
+          origin = AppendOrigin.Replication)
+        maybeCompletedTxn.foreach(completedTxns += _)
+      }
+    }
+    loadedProducers.values.foreach(producerStateManager.update)
+    completedTxns.foreach(producerStateManager.completeTxn)
+  }
+
+  private def updateProducers(producerStateManager: ProducerStateManager,
+                              batch: RecordBatch,
+                              producers: mutable.Map[Long, ProducerAppendInfo],
+                              firstOffsetMetadata: Option[LogOffsetMetadata],
+                              origin: AppendOrigin): Option[CompletedTxn] = {
+    val producerId = batch.producerId
+    val appendInfo = producers.getOrElseUpdate(producerId, producerStateManager.prepareUpdate(producerId, origin))
+    appendInfo.append(batch, firstOffsetMetadata)
+  }
+
+}
+
+object LogMetricNames {
+  val NumLogSegments: String = "NumLogSegments"
+  val LogStartOffset: String = "LogStartOffset"
+  val LogEndOffset: String = "LogEndOffset"
+  val Size: String = "Size"
+
+  def allMetricNames: List[String] = {
+    List(NumLogSegments, LogStartOffset, LogEndOffset, Size)
+  }
+}
+
+sealed trait SegmentDeletionReason {
+  def logReason(log: Log, toDelete: List[LogSegment]): Unit
+}
+
+case object RetentionMsBreach extends SegmentDeletionReason {
+  override def logReason(log: Log, toDelete: List[LogSegment]): Unit = {
+    val retentionMs = log.config.retentionMs
+    toDelete.foreach { segment =>
+      segment.largestRecordTimestamp match {
+        case Some(_) =>
+          log.info(s"Deleting segment $segment due to retention time ${retentionMs}ms breach based on the largest " +
+            s"record timestamp in the segment")
+        case None =>
+          log.info(s"Deleting segment $segment due to retention time ${retentionMs}ms breach based on the " +
+            s"last modified time of the segment")
+      }
+    }
+  }
+}
+
+case object RetentionSizeBreach extends SegmentDeletionReason {
+  override def logReason(log: Log, toDelete: List[LogSegment]): Unit = {
+    var size = log.size
+    toDelete.foreach { segment =>
+      size -= segment.size
+      log.info(s"Deleting segment $segment due to retention size ${log.config.retentionSize} breach. Log size " +
+        s"after deletion will be $size.")
+    }
+  }
+}
+
+case object StartOffsetBreach extends SegmentDeletionReason {
+  override def logReason(log: Log, toDelete: List[LogSegment]): Unit = {
+    log.info(s"Deleting segments due to log start offset ${log.logStartOffset} breach: ${toDelete.mkString(",")}")
+  }
+}
+
+case object LogRecovery extends SegmentDeletionReason {
+  override def logReason(log: Log, toDelete: List[LogSegment]): Unit = {
+    log.info(s"Deleting segments as part of log recovery: ${toDelete.mkString(",")}")
+  }
+}
+
+case object LogTruncation extends SegmentDeletionReason {
+  override def logReason(log: Log, toDelete: List[LogSegment]): Unit = {
+    log.info(s"Deleting segments as part of log truncation: ${toDelete.mkString(",")}")
+  }
+}
+
+case object LogRoll extends SegmentDeletionReason {
+  override def logReason(log: Log, toDelete: List[LogSegment]): Unit = {
+    log.info(s"Deleting segments as part of log roll: ${toDelete.mkString(",")}")
+  }
+}
+
+case object LogDeletion extends SegmentDeletionReason {
+  override def logReason(log: Log, toDelete: List[LogSegment]): Unit = {
+    log.info(s"Deleting segments as the log has been deleted: ${toDelete.mkString(",")}")
+  }
 }
diff --git a/core/src/main/scala/kafka/log/LogCleaner.scala b/core/src/main/scala/kafka/log/LogCleaner.scala
index 27229f3145f33..3225d9d1b314f 100644
--- a/core/src/main/scala/kafka/log/LogCleaner.scala
+++ b/core/src/main/scala/kafka/log/LogCleaner.scala
@@ -20,22 +20,24 @@ package kafka.log
 import java.io.{File, IOException}
 import java.nio._
 import java.util.Date
-import java.util.concurrent.{CountDownLatch, TimeUnit}
+import java.util.concurrent.TimeUnit
 
-import com.yammer.metrics.core.Gauge
 import kafka.common._
 import kafka.metrics.KafkaMetricsGroup
-import kafka.server.LogDirFailureChannel
+import kafka.server.{BrokerReconfigurable, KafkaConfig, LogDirFailureChannel}
 import kafka.utils._
-import org.apache.kafka.common.record._
-import org.apache.kafka.common.utils.Time
-import org.apache.kafka.common.TopicPartition
-import org.apache.kafka.common.errors.KafkaStorageException
+import org.apache.kafka.common.{KafkaException, TopicPartition}
+import org.apache.kafka.common.config.ConfigException
+import org.apache.kafka.common.errors.{CorruptRecordException, KafkaStorageException}
 import org.apache.kafka.common.record.MemoryRecords.RecordFilter
 import org.apache.kafka.common.record.MemoryRecords.RecordFilter.BatchRetention
+import org.apache.kafka.common.record._
+import org.apache.kafka.common.utils.Time
 
-import scala.collection.mutable
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
+import scala.collection.mutable.ListBuffer
+import scala.collection.{Iterable, Seq, Set, mutable}
+import scala.util.control.ControlThrowable
 
 /**
  * The cleaner is responsible for removing obsolete records from logs which have the "compact" retention strategy.
@@ -52,7 +54,7 @@ import scala.collection.JavaConverters._
  * To clean a log the cleaner first builds a mapping of key=>last_offset for the dirty section of the log. See kafka.log.OffsetMap for details of
  * the implementation of the mapping.
  *
- * Once the key=>offset map is built, the log is cleaned by recopying each log segment but omitting any key that appears in the offset map with a
+ * Once the key=>last_offset map is built, the log is cleaned by recopying each log segment but omitting any key that appears in the offset map with a
  * higher offset than what is found in the segment (i.e. messages with a key that appears in the dirty section of the log).
  *
  * To avoid segments shrinking to very small sizes with repeated cleanings we implement a rule by which if we will merge successive segments when
@@ -82,16 +84,20 @@ import scala.collection.JavaConverters._
  *    data from the transaction prior to reaching the offset of the marker. This follows the same logic used for
  *    tombstone deletion.
  *
- * @param config Configuration parameters for the cleaner
+ * @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 time A way to control the passage of time
  */
-class LogCleaner(val config: CleanerConfig,
+class LogCleaner(initialConfig: CleanerConfig,
                  val logDirs: Seq[File],
                  val logs: Pool[TopicPartition, Log],
                  val logDirFailureChannel: LogDirFailureChannel,
-                 time: Time = Time.SYSTEM) extends Logging with KafkaMetricsGroup {
+                 time: Time = Time.SYSTEM) extends Logging with KafkaMetricsGroup with BrokerReconfigurable
+{
+
+  /* Log cleaner configuration which may be dynamically updated */
+  @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)
@@ -104,72 +110,125 @@ class LogCleaner(val config: CleanerConfig,
                                         "bytes",
                                         time = time)
 
-  /* the threads */
-  private val cleaners = (0 until config.numThreads).map(new CleanerThread(_))
+  private[log] val cleaners = mutable.ArrayBuffer[CleanerThread]()
+
+  /**
+   * scala 2.12 does not support maxOption so we handle the empty manually.
+   * @param f to compute the result
+   * @return the max value (int value) or 0 if there is no cleaner
+   */
+  private def maxOverCleanerThreads(f: CleanerThread => Double): Int =
+    cleaners.foldLeft(0.0d)((max: Double, thread: CleanerThread) => math.max(max, f(thread))).toInt
+
 
   /* a metric to track the maximum utilization of any thread's buffer in the last cleaning */
   newGauge("max-buffer-utilization-percent",
-           new Gauge[Int] {
-             def value: Int = cleaners.map(_.lastStats).map(100 * _.bufferUtilization).max.toInt
-           })
+    () => maxOverCleanerThreads(_.lastStats.bufferUtilization) * 100)
+
   /* a metric to track the recopy rate of each thread's last cleaning */
-  newGauge("cleaner-recopy-percent",
-           new Gauge[Int] {
-             def value: Int = {
-               val stats = cleaners.map(_.lastStats)
-               val recopyRate = stats.map(_.bytesWritten).sum.toDouble / math.max(stats.map(_.bytesRead).sum, 1)
-               (100 * recopyRate).toInt
-             }
-           })
+  newGauge("cleaner-recopy-percent", () => {
+    val stats = cleaners.map(_.lastStats)
+    val recopyRate = stats.iterator.map(_.bytesWritten).sum.toDouble / math.max(stats.iterator.map(_.bytesRead).sum, 1)
+    (100 * recopyRate).toInt
+  })
+
   /* a metric to track the maximum cleaning time for the last cleaning from each thread */
   newGauge("max-clean-time-secs",
-           new Gauge[Int] {
-             def value: Int = cleaners.map(_.lastStats).map(_.elapsedSecs).max.toInt
-           })
+    () => maxOverCleanerThreads(_.lastStats.elapsedSecs))
+
+
+  // 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",
+    () => maxOverCleanerThreads(_.lastPreCleanStats.maxCompactionDelayMs.toDouble) / 1000)
+
+  newGauge("DeadThreadCount", () => deadThreadCount)
+
+  private[log] def deadThreadCount: Int = cleaners.count(_.isThreadFailed)
 
   /**
    * Start the background cleaning
    */
-  def startup() {
+  def startup(): Unit = {
     info("Starting the log cleaner")
-    cleaners.foreach(_.start())
+    (0 until config.numThreads).foreach { i =>
+      val cleaner = new CleanerThread(i)
+      cleaners += cleaner
+      cleaner.start()
+    }
   }
 
   /**
    * Stop the background cleaning
    */
-  def shutdown() {
+  def shutdown(): Unit = {
     info("Shutting down the log cleaner.")
     cleaners.foreach(_.shutdown())
+    cleaners.clear()
+  }
+
+  override def reconfigurableConfigs: Set[String] = {
+    LogCleaner.ReconfigurableConfigs
+  }
+
+  override def validateReconfiguration(newConfig: KafkaConfig): Unit = {
+    val newCleanerConfig = LogCleaner.cleanerConfig(newConfig)
+    val numThreads = newCleanerConfig.numThreads
+    val currentThreads = config.numThreads
+    if (numThreads < 1)
+      throw new ConfigException(s"Log cleaner threads should be at least 1")
+    if (numThreads < currentThreads / 2)
+      throw new ConfigException(s"Log cleaner threads cannot be reduced to less than half the current value $currentThreads")
+    if (numThreads > currentThreads * 2)
+      throw new ConfigException(s"Log cleaner threads cannot be increased to more than double the current value $currentThreads")
+
+  }
+
+  /**
+    * Reconfigure log clean config. This simply stops current log cleaners and creates new ones.
+    * That ensures that if any of the cleaners had failed, new cleaners are created to match the new config.
+    */
+  override def reconfigure(oldConfig: KafkaConfig, newConfig: KafkaConfig): Unit = {
+    config = LogCleaner.cleanerConfig(newConfig)
+    shutdown()
+    startup()
   }
 
   /**
    *  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
+   * Update checkpoint file to remove partitions if necessary.
    */
-  def updateCheckpoints(dataDir: File) {
-    cleanerManager.updateCheckpoints(dataDir, update=None)
+  def updateCheckpoints(dataDir: File, partitionToRemove: Option[TopicPartition] = None): Unit = {
+    cleanerManager.updateCheckpoints(dataDir, partitionToRemove = partitionToRemove)
   }
 
+  /**
+   * alter the checkpoint directory for the topicPartition, to remove the data in sourceLogDir, and add the data in destLogDir
+   */
   def alterCheckpointDir(topicPartition: TopicPartition, sourceLogDir: File, destLogDir: File): Unit = {
     cleanerManager.alterCheckpointDir(topicPartition, sourceLogDir, destLogDir)
   }
 
-  def handleLogDirFailure(dir: String) {
+  /**
+   * Stop cleaning logs in the provided directory
+   *
+   * @param dir     the absolute path of the log dir
+   */
+  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)
   }
 
@@ -177,15 +236,15 @@ class LogCleaner(val config: 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 a paused partition. This call blocks until the cleaning of a partition is resumed.
-   */
-  def resumeCleaning(topicPartition: TopicPartition) {
-    cleanerManager.resumeCleaning(topicPartition)
+    *  Resume the cleaning of paused partitions.
+    */
+  def resumeCleaning(topicPartitions: Iterable[TopicPartition]): Unit = {
+    cleanerManager.resumeCleaning(topicPartitions)
   }
 
   /**
@@ -209,16 +268,31 @@ class LogCleaner(val config: CleanerConfig,
     isCleaned
   }
 
+  /**
+    * To prevent race between retention and compaction,
+    * retention threads need to make this call to obtain:
+    * @return A list of log partitions that retention threads can safely work on
+    */
+  def pauseCleaningForNonCompactedPartitions(): Iterable[(TopicPartition, Log)] = {
+    cleanerManager.pauseCleaningForNonCompactedPartitions()
+  }
+
+  // Only for testing
+  private[kafka] def currentConfig: CleanerConfig = config
+
+  // Only for testing
+  private[log] def cleanerCount: Int = cleaners.size
+
   /**
    * 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)
-    extends ShutdownableThread(name = "kafka-log-cleaner-thread-" + threadId, isInterruptible = false) {
+  private[log] class CleanerThread(threadId: Int)
+    extends ShutdownableThread(name = s"kafka-log-cleaner-thread-$threadId", isInterruptible = false) {
 
-    override val loggerName = classOf[LogCleaner].getName
+    protected override def loggerName = classOf[LogCleaner].getName
 
-    if(config.dedupeBufferSize / config.numThreads > Int.MaxValue)
+    if (config.dedupeBufferSize / config.numThreads > Int.MaxValue)
       warn("Cannot use more than 2G of cleaner buffer space per cleaner thread, ignoring excess buffer space...")
 
     val cleaner = new Cleaner(id = threadId,
@@ -232,89 +306,120 @@ class LogCleaner(val config: CleanerConfig,
                               checkDone = checkDone)
 
     @volatile var lastStats: CleanerStats = new CleanerStats()
-    private val backOffWaitLatch = new CountDownLatch(1)
+    @volatile var lastPreCleanStats: PreCleanStats = new PreCleanStats()
 
-    private def checkDone(topicPartition: TopicPartition) {
-      if (!isRunning.get())
+    private def checkDone(topicPartition: TopicPartition): Unit = {
+      if (!isRunning)
         throw new ThreadShutdownException
       cleanerManager.checkCleaningAborted(topicPartition)
     }
 
     /**
      * 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() {
-      cleanOrSleep()
+    override def doWork(): Unit = {
+      val cleaned = tryCleanFilthiestLog()
+      if (!cleaned)
+        pause(config.backOffMs, TimeUnit.MILLISECONDS)
     }
 
-    override def shutdown() = {
-    	 initiateShutdown()
-    	 backOffWaitLatch.countDown()
-    	 awaitShutdown()
-     }
-
     /**
-     * Clean a log if there is a dirty log available, otherwise sleep for a bit
+     * Cleans a log if there is a dirty log available
+     * @return whether a log was cleaned
      */
-    private def cleanOrSleep() {
-      val cleaned = cleanerManager.grabFilthiestCompactedLog(time) match {
+    private def tryCleanFilthiestLog(): Boolean = {
+      try {
+        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
-          var endOffset = cleanable.firstDirtyOffset
+          this.lastPreCleanStats = preCleanStats
           try {
-            val (nextDirtyOffset, cleanerStats) = cleaner.clean(cleanable)
-            recordStats(cleaner.id, cleanable.log.name, cleanable.firstDirtyOffset, endOffset, cleanerStats)
-            endOffset = nextDirtyOffset
+            cleanLog(cleanable)
+            true
           } catch {
-            case _: LogCleaningAbortedException => // task can be aborted, let it go.
-            case _: KafkaStorageException => // partition is already offline. let it go.
-            case e: IOException =>
-              val msg = s"Failed to clean up log for ${cleanable.topicPartition} in dir ${cleanable.log.dir.getParent} due to IOException"
-              logDirFailureChannel.maybeAddOfflineLogDir(cleanable.log.dir.getParent, msg, e)
-          } finally {
-            cleanerManager.doneCleaning(cleanable.topicPartition, cleanable.log.dir.getParentFile, endOffset)
+            case e @ (_: ThreadShutdownException | _: ControlThrowable) => throw e
+            case e: Exception => throw new LogCleaningException(cleanable.log, e.getMessage, e)
           }
-          true
       }
       val deletable: Iterable[(TopicPartition, Log)] = cleanerManager.deletableLogs()
-      deletable.foreach{
-        case (topicPartition, log) =>
+      try {
+        deletable.foreach { case (_, log) =>
           try {
             log.deleteOldSegments()
-          } finally {
-            cleanerManager.doneDeleting(topicPartition)
+          } catch {
+            case e @ (_: ThreadShutdownException | _: ControlThrowable) => throw e
+            case e: Exception => throw new LogCleaningException(log, e.getMessage, e)
           }
+        }
+      } finally  {
+        cleanerManager.doneDeleting(deletable.map(_._1))
+      }
+
+      cleaned
+    }
+
+    private def cleanLog(cleanable: LogToClean): Unit = {
+      val startOffset = cleanable.firstDirtyOffset
+      var endOffset = startOffset
+      try {
+        val (nextDirtyOffset, cleanerStats) = cleaner.clean(cleanable)
+        endOffset = nextDirtyOffset
+        recordStats(cleaner.id, cleanable.log.name, startOffset, endOffset, cleanerStats)
+      } catch {
+        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.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.parentDirFile, endOffset)
       }
-      if (!cleaned)
-        backOffWaitLatch.await(config.backOffMs, TimeUnit.MILLISECONDS)
     }
 
     /**
      * 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 =
         "%n\tLog cleaner thread %d cleaned log %s (dirty section = [%d, %d])%n".format(id, name, from, to) +
-        "\t%,.1f MB of log processed in %,.1f seconds (%,.1f MB/sec).%n".format(mb(stats.bytesRead),
+        "\t%,.1f MB of log processed in %,.1f seconds (%,.1f MB/sec).%n".format(mb(stats.bytesRead.toDouble),
                                                                                 stats.elapsedSecs,
-                                                                                mb(stats.bytesRead/stats.elapsedSecs)) +
-        "\tIndexed %,.1f MB in %.1f seconds (%,.1f Mb/sec, %.1f%% of total time)%n".format(mb(stats.mapBytesRead),
+                                                                                mb(stats.bytesRead.toDouble / stats.elapsedSecs)) +
+        "\tIndexed %,.1f MB in %.1f seconds (%,.1f Mb/sec, %.1f%% of total time)%n".format(mb(stats.mapBytesRead.toDouble),
                                                                                            stats.elapsedIndexSecs,
-                                                                                           mb(stats.mapBytesRead)/stats.elapsedIndexSecs,
-                                                                                           100 * stats.elapsedIndexSecs/stats.elapsedSecs) +
+                                                                                           mb(stats.mapBytesRead.toDouble) / stats.elapsedIndexSecs,
+                                                                                           100 * stats.elapsedIndexSecs / stats.elapsedSecs) +
         "\tBuffer utilization: %.1f%%%n".format(100 * stats.bufferUtilization) +
-        "\tCleaned %,.1f MB in %.1f seconds (%,.1f Mb/sec, %.1f%% of total time)%n".format(mb(stats.bytesRead),
+        "\tCleaned %,.1f MB in %.1f seconds (%,.1f Mb/sec, %.1f%% of total time)%n".format(mb(stats.bytesRead.toDouble),
                                                                                            stats.elapsedSecs - stats.elapsedIndexSecs,
-                                                                                           mb(stats.bytesRead)/(stats.elapsedSecs - stats.elapsedIndexSecs), 100 * (stats.elapsedSecs - stats.elapsedIndexSecs).toDouble/stats.elapsedSecs) +
-        "\tStart size: %,.1f MB (%,d messages)%n".format(mb(stats.bytesRead), stats.messagesRead) +
-        "\tEnd size: %,.1f MB (%,d messages)%n".format(mb(stats.bytesWritten), stats.messagesWritten) +
+                                                                                           mb(stats.bytesRead.toDouble) / (stats.elapsedSecs - stats.elapsedIndexSecs), 100 * (stats.elapsedSecs - stats.elapsedIndexSecs) / stats.elapsedSecs) +
+        "\tStart size: %,.1f MB (%,d messages)%n".format(mb(stats.bytesRead.toDouble), stats.messagesRead) +
+        "\tEnd size: %,.1f MB (%,d messages)%n".format(mb(stats.bytesWritten.toDouble), stats.messagesWritten) +
         "\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))
       }
@@ -323,6 +428,37 @@ class LogCleaner(val config: CleanerConfig,
   }
 }
 
+object LogCleaner {
+  val ReconfigurableConfigs = Set(
+    KafkaConfig.LogCleanerThreadsProp,
+    KafkaConfig.LogCleanerDedupeBufferSizeProp,
+    KafkaConfig.LogCleanerDedupeBufferLoadFactorProp,
+    KafkaConfig.LogCleanerIoBufferSizeProp,
+    KafkaConfig.MessageMaxBytesProp,
+    KafkaConfig.LogCleanerIoMaxBytesPerSecondProp,
+    KafkaConfig.LogCleanerBackoffMsProp
+  )
+
+  def cleanerConfig(config: KafkaConfig): CleanerConfig = {
+    CleanerConfig(numThreads = config.logCleanerThreads,
+      dedupeBufferSize = config.logCleanerDedupeBufferSize,
+      dedupeBufferLoadFactor = config.logCleanerDedupeBufferLoadFactor,
+      ioBufferSize = config.logCleanerIoBufferSize,
+      maxMessageSize = config.messageMaxBytes,
+      maxIoBytesPerSecond = config.logCleanerIoMaxBytesPerSecond,
+      backOffMs = config.logCleanerBackoffMs,
+      enableCleaner = config.logCleanerEnable)
+
+  }
+
+  def createNewCleanedSegment(log: Log, baseOffset: Long): LogSegment = {
+    LogSegment.deleteIfExists(log.dir, baseOffset, fileSuffix = Log.CleanedFileSuffix)
+    LogSegment.open(log.dir, baseOffset, log.config, Time.SYSTEM,
+      fileSuffix = Log.CleanedFileSuffix, initFileSize = log.initFileSize, preallocate = log.config.preallocate)
+  }
+
+}
+
 /**
  * This class holds the actual logic for cleaning a log
  * @param id An identifier used for logging
@@ -341,11 +477,11 @@ private[log] class Cleaner(val id: Int,
                            dupBufferLoadFactor: Double,
                            throttler: Throttler,
                            time: Time,
-                           checkDone: (TopicPartition) => Unit) extends Logging {
+                           checkDone: TopicPartition => Unit) extends Logging {
 
-  override val loggerName = classOf[LogCleaner].getName
+  protected override def loggerName = classOf[LogCleaner].getName
 
-  this.logIdent = "Cleaner " + id + ": "
+  this.logIdent = s"Cleaner $id: "
 
   /* buffer used for read i/o */
   private var readBuffer = ByteBuffer.allocate(ioBufferSize)
@@ -393,11 +529,14 @@ private[log] class Cleaner(val id: Int,
     // this is the lower of the last active segment and the compaction lag
     val cleanableHorizonMs = log.logSegments(0, cleanable.firstUncleanableOffset).lastOption.map(_.lastModified).getOrElse(0L)
 
-
     // 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
@@ -415,38 +554,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) {
-
-    def deleteAndGetCleanedFile(file: File): File = {
-      val f = new File(file.getPath + Log.CleanedFileSuffix)
-      f.delete()
-      f
-    }
-
+                                 stats: CleanerStats,
+                                 transactionMetadata: CleanedTransactionMetadata): Unit = {
     // create a new segment with a suffix appended to the name of the log and indexes
-    val firstSegment = segments.head
-    val logFile = deleteAndGetCleanedFile(firstSegment.log.file)
-    val indexFile = deleteAndGetCleanedFile(firstSegment.index.file)
-    val timeIndexFile = deleteAndGetCleanedFile(firstSegment.timeIndex.file)
-    val txnIndexFile = deleteAndGetCleanedFile(firstSegment.txnIndex.file)
-
-    val startOffset = firstSegment.baseOffset
-    val records = FileRecords.open(logFile, false, log.initFileSize, log.config.preallocate)
-    val index = new OffsetIndex(indexFile, startOffset, firstSegment.index.maxIndexSize)
-    val timeIndex = new TimeIndex(timeIndexFile, startOffset, firstSegment.timeIndex.maxIndexSize)
-    val txnIndex = new TransactionIndex(startOffset, txnIndexFile)
-    val cleaned = new LogSegment(records, index, timeIndex, txnIndex, startOffset, firstSegment.indexIntervalBytes,
-      log.config.randomSegmentJitter, time)
+    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
@@ -454,29 +580,28 @@ 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(txnIndex))
-
-        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.")
-        cleanInto(log.topicPartition, currentSegment.log, cleaned, map, retainDeletes, log.config.maxMessageSize,
-          transactionMetadata, log.activeProducersWithLastSequence, stats)
-
+        transactionMetadata.addAbortedTransactions(abortedTransactions)
+
+        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, 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
+            // scratch once the split is complete.
+            info(s"Caught segment overflow error during cleaning: ${e.getMessage}")
+            log.splitOverflowedSegment(currentSegment)
+            throw new LogCleaningAbortedException()
+        }
         currentSegmentOpt = nextSegmentOpt
       }
 
-      // trim log segment
-      cleaned.log.trim()
-
-      // trim excess index
-      index.trimToValidSize()
-
-      // Append the last index entry
       cleaned.onBecomeInactiveSegment()
-
-      // trim time index
-      timeIndex.trimToValidSize()
-
       // flush new segment to disk before swap
       cleaned.flush()
 
@@ -485,12 +610,15 @@ private[log] class Cleaner(val id: Int,
       cleaned.lastModified = modified
 
       // swap in new segment
-      info("Swapping in cleaned segment %d for segment(s) %s in log %s.".format(cleaned.baseOffset, segments.map(_.baseOffset).mkString(","), log.name))
-      log.replaceSegments(cleaned, segments)
+      info(s"Swapping in cleaned segment $cleaned for segment(s) $segments in log $log")
+      log.replaceSegments(List(cleaned), segments)
     } catch {
       case e: LogCleaningAbortedException =>
-        cleaned.delete()
-        throw e
+        try cleaned.deleteIfExists()
+        catch {
+          case deleteException: Exception =>
+            e.addSuppressed(deleteException)
+        } finally throw e
     }
   }
 
@@ -502,7 +630,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
    */
@@ -510,22 +638,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
@@ -538,7 +679,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)
       }
     }
 
@@ -559,27 +700,60 @@ private[log] class Cleaner(val id: Int,
       position += result.bytesRead
 
       // if any messages are to be retained, write them out
-      val outputBuffer = result.output
+      val outputBuffer = result.outputBuffer
       if (outputBuffer.position() > 0) {
         outputBuffer.flip()
         val retained = MemoryRecords.readableRecords(outputBuffer)
         // it's OK not to hold the Log's lock in this case, because this segment is only accessed by other threads
         // after `Log.replaceSegments` (which acquires the lock) is called
-        dest.append(firstOffset = retained.batches.iterator.next().baseOffset,
-          largestOffset = result.maxOffset,
+        dest.append(largestOffset = result.maxOffset,
           largestTimestamp = result.maxTimestamp,
           shallowOffsetOfMaxTimestamp = result.shallowOffsetOfMaxTimestamp,
           records = retained)
         throttler.maybeThrottle(outputBuffer.limit())
       }
 
-      // if we read bytes but didn't get even one complete message, our I/O buffer is too small, grow it and try again
-      if (readBuffer.limit() > 0 && result.messagesRead == 0)
-        growBuffers(maxLogMessageSize)
+      // if we read bytes but didn't get even one complete batch, our I/O buffer is too small, grow it and try again
+      // `result.bytesRead` contains bytes from `messagesRead` and any discarded batches.
+      if (readBuffer.limit() > 0 && result.bytesRead == 0)
+        growBuffersOrFail(sourceRecords, position, maxLogMessageSize, records)
     }
     restoreBuffers()
   }
 
+
+  /**
+   * Grow buffers to process next batch of records from `sourceRecords.` Buffers are doubled in size
+   * up to a maximum of `maxLogMessageSize`. In some scenarios, a record could be bigger than the
+   * current maximum size configured for the log. For example:
+   *   1. A compacted topic using compression may contain a message set slightly larger than max.message.bytes
+   *   2. max.message.bytes of a topic could have been reduced after writing larger messages
+   * In these cases, grow the buffer to hold the next batch.
+   */
+  private def growBuffersOrFail(sourceRecords: FileRecords,
+                                position: Int,
+                                maxLogMessageSize: Int,
+                                memoryRecords: MemoryRecords): Unit = {
+
+    val maxSize = if (readBuffer.capacity >= maxLogMessageSize) {
+      val nextBatchSize = memoryRecords.firstBatchSize
+      val logDesc = s"log segment ${sourceRecords.file} at position $position"
+      if (nextBatchSize == null)
+        throw new IllegalStateException(s"Could not determine next batch size for $logDesc")
+      if (nextBatchSize <= 0)
+        throw new IllegalStateException(s"Invalid batch size $nextBatchSize for $logDesc")
+      if (nextBatchSize <= readBuffer.capacity)
+        throw new IllegalStateException(s"Batch size $nextBatchSize < buffer size ${readBuffer.capacity}, but not processed for $logDesc")
+      val bytesLeft = sourceRecords.channel.size - position
+      if (nextBatchSize > bytesLeft)
+        throw new CorruptRecordException(s"Log segment may be corrupt, batch size $nextBatchSize > $bytesLeft bytes left in segment for $logDesc")
+      nextBatchSize.intValue
+    } else
+      maxLogMessageSize
+
+    growBuffers(maxSize)
+  }
+
   private def shouldDiscardBatch(batch: RecordBatch,
                                  transactionMetadata: CleanedTransactionMetadata,
                                  retainTxnMarkers: Boolean): Boolean = {
@@ -604,13 +778,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
@@ -620,12 +795,12 @@ 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))
     val newSize = math.min(this.readBuffer.capacity * 2, maxBufferSize)
-    info("Growing cleaner I/O buffers from " + readBuffer.capacity + "bytes to " + newSize + " bytes.")
+    info(s"Growing cleaner I/O buffers from ${readBuffer.capacity} bytes to $newSize bytes.")
     this.readBuffer = ByteBuffer.allocate(newSize)
     this.writeBuffer = ByteBuffer.allocate(newSize)
   }
@@ -633,7 +808,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)
@@ -657,17 +832,17 @@ private[log] class Cleaner(val id: Int,
     while(segs.nonEmpty) {
       var group = List(segs.head)
       var logSize = segs.head.size.toLong
-      var indexSize = segs.head.index.sizeInBytes.toLong
+      var indexSize = segs.head.offsetIndex.sizeInBytes.toLong
       var timeIndexSize = segs.head.timeIndex.sizeInBytes.toLong
       segs = segs.tail
       while(segs.nonEmpty &&
             logSize + segs.head.size <= maxSize &&
-            indexSize + segs.head.index.sizeInBytes <= maxIndexSize &&
+            indexSize + segs.head.offsetIndex.sizeInBytes <= maxIndexSize &&
             timeIndexSize + segs.head.timeIndex.sizeInBytes <= maxIndexSize &&
             lastOffsetForFirstSegment(segs, firstUncleanableOffset) - group.last.baseOffset <= Int.MaxValue) {
         group = segs.head :: group
         logSize += segs.head.size
-        indexSize += segs.head.index.sizeInBytes
+        indexSize += segs.head.offsetIndex.sizeInBytes
         timeIndexSize += segs.head.timeIndex.sizeInBytes
         segs = segs.tail
       }
@@ -710,21 +885,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))
@@ -745,15 +926,22 @@ private[log] class Cleaner(val id: Int,
                                        segment: LogSegment,
                                        map: OffsetMap,
                                        startOffset: Long,
+                                       nextSegmentStartOffset: Long,
                                        maxLogMessageSize: Int,
                                        transactionMetadata: CleanedTransactionMetadata,
                                        stats: CleanerStats): Boolean = {
-    var position = segment.index.lookup(startOffset).position
+    var position = segment.offsetIndex.lookup(startOffset).position
     val maxDesiredMapSize = (map.slots * this.dupBufferLoadFactor).toInt
     while (position < segment.log.sizeInBytes) {
       checkDone(topicPartition)
       readBuffer.clear()
-      segment.log.readInto(readBuffer, position)
+      try {
+        segment.log.readInto(readBuffer, position)
+      } catch {
+        case e: Exception =>
+          throw new KafkaException(s"Failed to read from segment $segment of partition $topicPartition " +
+            "while loading offset map", e)
+      }
       val records = MemoryRecords.readableRecords(readBuffer)
       throttler.maybeThrottle(records.sizeInBytes)
 
@@ -769,15 +957,18 @@ private[log] class Cleaner(val id: Int,
             // Note that abort markers are supported in v2 and above, which means count is defined.
             stats.indexMessagesRead(batch.countOrNull)
           } else {
-            for (record <- batch.asScala) {
-              if (record.hasKey && record.offset >= startOffset) {
-                if (map.size < maxDesiredMapSize)
-                  map.put(record.key, record.offset)
-                else
-                  return true
+            val recordsIterator = batch.streamingIterator(decompressionBufferSupplier)
+            try {
+              for (record <- recordsIterator.asScala) {
+                if (record.hasKey && record.offset >= startOffset) {
+                  if (map.size < maxDesiredMapSize)
+                    map.put(record.key, record.offset)
+                  else
+                    return true
+                }
+                stats.indexMessagesRead(1)
               }
-              stats.indexMessagesRead(1)
-            }
+            } finally recordsIterator.close()
           }
         }
 
@@ -790,13 +981,36 @@ private[log] class Cleaner(val id: Int,
 
       // if we didn't read even one complete message, our read buffer may be too small
       if(position == startPosition)
-        growBuffers(maxLogMessageSize)
+        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
  */
@@ -813,78 +1027,78 @@ 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
   }
 
-  def elapsedSecs = (endTime - startTime)/1000.0
+  def elapsedSecs: Double = (endTime - startTime) / 1000.0
 
-  def elapsedIndexSecs = (mapCompleteTime - startTime)/1000.0
+  def elapsedIndexSecs: Double = (mapCompleteTime - startTime) / 1000.0
 
 }
 
 /**
- * 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
-  private[this] 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
+  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.
@@ -893,31 +1107,39 @@ private[log] class CleanedTransactionMetadata(val abortedTransactions: mutable.P
   def onControlBatchRead(controlBatch: RecordBatch): Boolean = {
     consumeAbortedTxnsUpTo(controlBatch.lastOffset)
 
-    val controlRecord = controlBatch.iterator.next()
-    val controlType = ControlRecordType.parse(controlRecord.key)
-    val producerId = controlBatch.producerId
-    controlType match {
-      case ControlRecordType.ABORT =>
-        ongoingAbortedTxns.remove(producerId) match {
-          // Retain the marker until all batches from the transaction have been removed
-          case Some(abortedTxnMetadata) if abortedTxnMetadata.lastObservedBatchOffset.isDefined =>
-            transactionIndex.foreach(_.append(abortedTxnMetadata.abortedTxn))
-            false
-          case _ => true
-        }
+    val controlRecordIterator = controlBatch.iterator
+    if (controlRecordIterator.hasNext) {
+      val controlRecord = controlRecordIterator.next()
+      val controlType = ControlRecordType.parse(controlRecord.key)
+      val producerId = controlBatch.producerId
+      controlType match {
+        case ControlRecordType.ABORT =>
+          ongoingAbortedTxns.remove(producerId) match {
+            // 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 =>
+              cleanedIndex.foreach(_.append(abortedTxnMetadata.abortedTxn))
+              false
+            case _ => true
+          }
 
-      case ControlRecordType.COMMIT =>
-        // This marker is eligible for deletion if we didn't traverse any batches from the transaction
-        !ongoingCommittedTxns.remove(producerId)
+        case ControlRecordType.COMMIT =>
+          // This marker is eligible for deletion if we didn't traverse any batches from the transaction
+          !ongoingCommittedTxns.remove(producerId)
 
-      case _ => false
+        case _ => false
+      }
+    } else {
+      // An empty control batch was already cleaned, so it's safe to discard
+      true
     }
   }
 
   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))
     }
   }
 
@@ -945,4 +1167,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 b48b757a3357d..eea889a7218e5 100755
--- a/core/src/main/scala/kafka/log/LogCleanerManager.scala
+++ b/core/src/main/scala/kafka/log/LogCleanerManager.scala
@@ -21,8 +21,7 @@ import java.io.File
 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,28 +31,40 @@ import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.utils.Time
 import org.apache.kafka.common.errors.KafkaStorageException
 
-import scala.collection.{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 object LogCleaningPaused 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)
 
 /**
- *  Manage the state of each partition being cleaned.
- *  If a partition is to be cleaned, it enters the LogCleaningInProgress state.
- *  While a partition is being cleaned, it can be requested to be aborted and paused. Then the partition first enters
- *  the LogCleaningAborted state. Once the cleaning task is aborted, the partition enters the LogCleaningPaused state.
- *  While a partition is in the LogCleaningPaused state, it won't be scheduled for cleaning again, until cleaning is
- *  requested to be resumed.
- */
+  * This class manages the state of each partition being cleaned.
+  * LogCleaningState defines the cleaning states that a TopicPartition can be in.
+  * 1. None                    : No cleaning state in a TopicPartition. In this state, it can become LogCleaningInProgress
+  *                              or LogCleaningPaused(1). Valid previous state are LogCleaningInProgress and LogCleaningPaused(1)
+  * 2. LogCleaningInProgress   : The cleaning is currently in progress. In this state, it can become None when log cleaning is finished
+  *                              or become LogCleaningAborted. Valid previous state is None.
+  * 3. LogCleaningAborted      : The cleaning abort is requested. In this state, it can become LogCleaningPaused(1).
+  *                              Valid previous state is LogCleaningInProgress.
+  * 4-a. LogCleaningPaused(1)  : The cleaning is paused once. No log cleaning can be done in this state.
+  *                              In this state, it can become None or LogCleaningPaused(2).
+  *                              Valid previous state is None, LogCleaningAborted or LogCleaningPaused(2).
+  * 4-b. LogCleaningPaused(i)  : The cleaning is paused i times where i>= 2. No log cleaning can be done in this state.
+  *                              In this state, it can become LogCleaningPaused(i-1) or LogCleaningPaused(i+1).
+  *                              Valid previous state is LogCleaningPaused(i-1) or LogCleaningPaused(i+1).
+  */
 private[log] class LogCleanerManager(val logDirs: Seq[File],
                                      val logs: Pool[TopicPartition, Log],
                                      val logDirFailureChannel: LogDirFailureChannel) extends Logging with KafkaMetricsGroup {
-
   import LogCleanerManager._
 
-  override val loggerName = classOf[LogCleaner].getName
+
+  protected override def loggerName = classOf[LogCleaner].getName
 
   // package-private for testing
   private[log] val offsetCheckpointFile = "cleaner-offset-checkpoint"
@@ -65,19 +76,53 @@ private[log] class LogCleanerManager(val logDirs: Seq[File],
   /* the set of logs currently being cleaned */
   private val inProgress = mutable.HashMap[TopicPartition, LogCleaningState]()
 
+  /* the set of uncleanable partitions (partitions that have raised an unexpected error during cleaning)
+   *   for each log directory */
+  private val uncleanablePartitions = mutable.HashMap[String, mutable.Set[TopicPartition]]()
+
   /* a global lock used to control all access to the in-progress set and the offset checkpoints */
   private val lock = new ReentrantLock
 
   /* for coordinating the pausing and the cleaning of a partition */
   private val pausedCleaningCond = lock.newCondition()
 
+  /* gauges for tracking the number of partitions marked as uncleanable for each log directory */
+  for (dir <- logDirs) {
+    newGauge("uncleanable-partitions-count",
+      () => inLock(lock) { uncleanablePartitions.get(dir.getAbsolutePath).map(_.size).getOrElse(0) },
+      Map("logDirectory" -> dir.getAbsolutePath)
+    )
+  }
+
+  /* gauges for tracking the number of uncleanable bytes from uncleanable partitions for each log directory */
+  for (dir <- logDirs) {
+    newGauge("uncleanable-bytes",
+      () => inLock(lock) {
+        uncleanablePartitions.get(dir.getAbsolutePath) match {
+          case Some(partitions) =>
+            val lastClean = allCleanerCheckpoints
+            val now = Time.SYSTEM.milliseconds
+            partitions.iterator.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 None => 0
+        }
+      },
+      Map("logDirectory" -> dir.getAbsolutePath)
+    )
+  }
+
   /* a gauge for tracking the cleanable ratio of the dirtiest log */
   @volatile private var dirtiestLogCleanableRatio = 0.0
-  newGauge("max-dirty-percent", new Gauge[Int] { def value = (100 * dirtiestLogCleanableRatio).toInt })
+  newGauge("max-dirty-percent", () => (100 * dirtiestLogCleanableRatio).toInt)
 
   /* a gauge for tracking the time since the last log cleaner run, in milli seconds */
-  @volatile private var timeOfLastRun : Long = Time.SYSTEM.milliseconds
-  newGauge("time-since-last-run-ms", new Gauge[Long] { def value = Time.SYSTEM.milliseconds - timeOfLastRun })
+  @volatile private var timeOfLastRun: Long = Time.SYSTEM.milliseconds
+  newGauge("time-since-last-run-ms", () => Time.SYSTEM.milliseconds - timeOfLastRun)
 
   /**
    * @return the position processed for all logs.
@@ -96,13 +141,30 @@ private[log] class LogCleanerManager(val logDirs: Seq[File],
     }
   }
 
+  /**
+    * Package private for unit test. Get the cleaning state of the partition.
+    */
+  private[log] def cleaningState(tp: TopicPartition): Option[LogCleaningState] = {
+    inLock(lock) {
+      inProgress.get(tp)
+    }
+  }
+
+  /**
+    * Package private for unit test. Set the cleaning state of the partition.
+    */
+  private[log] def setCleaningState(tp: TopicPartition, state: LogCleaningState): Unit = {
+    inLock(lock) {
+      inProgress.put(tp, state)
+    }
+  }
 
    /**
     * Choose the log to clean next and add it to the in-progress set. We recompute this
     * 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
@@ -110,20 +172,36 @@ private[log] class LogCleanerManager(val logDirs: Seq[File],
       val dirtyLogs = logs.filter {
         case (_, log) => log.config.compact  // match logs that are marked as compacted
       }.filterNot {
-        case (topicPartition, _) => inProgress.contains(topicPartition) // skip any logs already in-progress
+        case (topicPartition, log) =>
+          // skip any logs already in-progress and uncleanable partitions
+          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, partitionToUpdateOrAdd = 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)
@@ -132,12 +210,37 @@ private[log] class LogCleanerManager(val logDirs: Seq[File],
   }
 
   /**
-    * Find any logs that have compact and delete enabled
+    * Pause logs cleaning for logs that do not have compaction enabled
+    * and do not have other deletion or compaction in progress.
+    * This is to handle potential race between retention and cleaner threads when users
+    * switch topic configuration between compacted and non-compacted topic.
+    * @return retention logs that have log cleaning successfully paused
+    */
+  def pauseCleaningForNonCompactedPartitions(): Iterable[(TopicPartition, Log)] = {
+    inLock(lock) {
+      val deletableLogs = logs.filter {
+        case (_, log) => !log.config.compact // pick non-compacted logs
+      }.filterNot {
+        case (topicPartition, _) => inProgress.contains(topicPartition) // skip any logs already in-progress
+      }
+
+      deletableLogs.foreach {
+        case (topicPartition, _) => inProgress.put(topicPartition, LogCleaningPaused(1))
+      }
+      deletableLogs
+    }
+  }
+
+  /**
+    * Find any logs that have compaction enabled. Mark them as being cleaned
+    * Include logs without delete enabled, as they may have segments
+    * that precede the start offset.
     */
   def deletableLogs(): Iterable[(TopicPartition, Log)] = {
     inLock(lock) {
       val toClean = logs.filter { case (topicPartition, log) =>
-        !inProgress.contains(topicPartition) && isCompactAndDelete(log)
+        !inProgress.contains(topicPartition) && log.config.compact &&
+          !isUncleanablePartition(log, topicPartition)
       }
       toClean.foreach { case (tp, _) => inProgress.put(tp, LogCleaningInProgress) }
       toClean
@@ -150,12 +253,11 @@ 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(topicPartition)
+      resumeCleaning(Seq(topicPartition))
     }
-    info(s"The cleaning for partition $topicPartition is aborted")
   }
 
   /**
@@ -167,45 +269,49 @@ private[log] class LogCleanerManager(val logDirs: Seq[File],
    *     throws a LogCleaningAbortedException to stop the cleaning task.
    *  4. When the cleaning task is stopped, doneCleaning() is called, which sets the state of the partition as paused.
    *  5. abortAndPauseCleaning() waits until the state of the partition is changed to paused.
+   *  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 =>
-          inProgress.put(topicPartition, LogCleaningPaused)
-        case Some(state) =>
-          state match {
-            case LogCleaningInProgress =>
-              inProgress.put(topicPartition, LogCleaningAborted)
-            case LogCleaningPaused =>
-            case s =>
-              throw new IllegalStateException(s"Compaction for partition $topicPartition cannot be aborted and paused since it is in $s state.")
-          }
+          inProgress.put(topicPartition, LogCleaningPaused(1))
+        case Some(LogCleaningInProgress) =>
+          inProgress.put(topicPartition, LogCleaningAborted)
+        case Some(LogCleaningPaused(count)) =>
+          inProgress.put(topicPartition, LogCleaningPaused(count + 1))
+        case Some(s) =>
+          throw new IllegalStateException(s"Compaction for partition $topicPartition cannot be aborted and paused since it is in $s state.")
       }
-      while (!isCleaningInState(topicPartition, LogCleaningPaused))
+      while(!isCleaningInStatePaused(topicPartition))
         pausedCleaningCond.await(100, TimeUnit.MILLISECONDS)
     }
-    info(s"The cleaning for partition $topicPartition is aborted and paused")
   }
 
   /**
-   *  Resume the cleaning of a paused partition. This call blocks until the cleaning of a partition is resumed.
-   */
-  def resumeCleaning(topicPartition: TopicPartition) {
+    *  Resume the cleaning of paused partitions.
+    *  Each call of this function will undo one pause.
+    */
+  def resumeCleaning(topicPartitions: Iterable[TopicPartition]): Unit = {
     inLock(lock) {
-      inProgress.get(topicPartition) match {
-        case None =>
-          throw new IllegalStateException(s"Compaction for partition $topicPartition cannot be resumed since it is not paused.")
-        case Some(state) =>
-          state match {
-            case LogCleaningPaused =>
-              inProgress.remove(topicPartition)
-            case s =>
-              throw new IllegalStateException(s"Compaction for partition $topicPartition cannot be resumed since it is in $s state.")
+      topicPartitions.foreach {
+        topicPartition =>
+          inProgress.get(topicPartition) match {
+            case None =>
+              throw new IllegalStateException(s"Compaction for partition $topicPartition cannot be resumed since it is not paused.")
+            case Some(state) =>
+              state match {
+                case LogCleaningPaused(count) if count == 1 =>
+                  inProgress.remove(topicPartition)
+                case LogCleaningPaused(count) if count > 1 =>
+                  inProgress.put(topicPartition, LogCleaningPaused(count - 1))
+                case s =>
+                  throw new IllegalStateException(s"Compaction for partition $topicPartition cannot be resumed since it is in $s state.")
+              }
           }
       }
     }
-    info(s"Compaction for partition $topicPartition is resumed")
   }
 
   /**
@@ -222,23 +328,58 @@ private[log] class LogCleanerManager(val logDirs: Seq[File],
     }
   }
 
+  /**
+   *  Check if the cleaning for a partition is paused. The caller is expected to hold lock while making the call.
+   */
+  private def isCleaningInStatePaused(topicPartition: TopicPartition): Boolean = {
+    inProgress.get(topicPartition) match {
+      case None => false
+      case Some(state) =>
+        state match {
+          case _: LogCleaningPaused =>
+            true
+          case _ =>
+            false
+        }
+    }
+  }
+
   /**
    *  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)]) {
+  /**
+   * Update checkpoint file, adding or removing partitions if necessary.
+   *
+   * @param dataDir                       The File object to be updated
+   * @param partitionToUpdateOrAdd        The [TopicPartition, Long] map data to be updated. pass "none" if doing remove, not add
+   * @param topicPartitionToBeRemoved     The TopicPartition to be removed
+   */
+  def updateCheckpoints(dataDir: File, partitionToUpdateOrAdd: Option[(TopicPartition, Long)] = None,
+                        partitionToRemove: Option[TopicPartition] = None): Unit = {
     inLock(lock) {
       val checkpoint = checkpoints(dataDir)
       if (checkpoint != null) {
         try {
-          val existing = checkpoint.read().filterKeys(logs.keys) ++ update
-          checkpoint.write(existing)
+          val currentCheckpoint = checkpoint.read().filter { case (tp, _) => logs.keys.contains(tp) }.toMap
+          // remove the partition offset if any
+          var updatedCheckpoint = partitionToRemove match {
+            case Some(topicPartion) => currentCheckpoint - topicPartion
+            case None => currentCheckpoint
+          }
+          // update or add the partition offset if any
+          updatedCheckpoint = partitionToUpdateOrAdd match {
+            case Some(updatedOffset) => updatedCheckpoint + updatedOffset
+            case None => updatedCheckpoint
+          }
+
+          checkpoint.write(updatedCheckpoint)
         } catch {
           case e: KafkaStorageException =>
             error(s"Failed to access checkpoint file ${checkpoint.file.getName} in dir ${checkpoint.file.getParentFile.getAbsolutePath}", e)
@@ -247,39 +388,59 @@ private[log] class LogCleanerManager(val logDirs: Seq[File],
     }
   }
 
+  /**
+   * alter the checkpoint directory for the topicPartition, to remove the data in sourceLogDir, and add the data in destLogDir
+   */
   def alterCheckpointDir(topicPartition: TopicPartition, sourceLogDir: File, destLogDir: File): Unit = {
     inLock(lock) {
       try {
         checkpoints.get(sourceLogDir).flatMap(_.read().get(topicPartition)) match {
           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
-            updateCheckpoints(destLogDir, Option(topicPartition, offset))
+            debug(s"Removing the partition offset data in checkpoint file for '${topicPartition}' " +
+              s"from ${sourceLogDir.getAbsoluteFile} directory.")
+            updateCheckpoints(sourceLogDir, partitionToRemove = Option(topicPartition))
+
+            debug(s"Adding the partition offset data in checkpoint file for '${topicPartition}' " +
+              s"to ${destLogDir.getAbsoluteFile} directory.")
+            updateCheckpoints(destLogDir, partitionToUpdateOrAdd = Option(topicPartition, offset))
           case None =>
         }
       } catch {
         case e: KafkaStorageException =>
           error(s"Failed to access checkpoint file in dir ${sourceLogDir.getAbsolutePath}", e)
       }
+
+      val logUncleanablePartitions = uncleanablePartitions.getOrElse(sourceLogDir.toString, mutable.Set[TopicPartition]())
+      if (logUncleanablePartitions.contains(topicPartition)) {
+        logUncleanablePartitions.remove(topicPartition)
+        markPartitionUncleanable(destLogDir.toString, topicPartition)
+      }
     }
   }
 
-  def handleLogDirFailure(dir: String) {
-    info(s"Stopping cleaning logs in dir $dir")
+  /**
+   * Stop cleaning logs in the provided directory
+   *
+   * @param dir     the absolute path of the log dir
+   */
+  def handleLogDirFailure(dir: String): Unit = {
+    warn(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) {
+  /**
+   * Truncate the checkpointed offset for the given partition if its checkpointed offset is larger than the given offset
+   */
+  def maybeTruncateCheckpoint(dataDir: File, topicPartition: TopicPartition, offset: Long): Unit = {
     inLock(lock) {
       if (logs.get(topicPartition).config.compact) {
         val checkpoint = checkpoints(dataDir)
         if (checkpoint != null) {
           val existing = checkpoint.read()
           if (existing.getOrElse(topicPartition, 0L) > offset)
-            checkpoint.write(existing + (topicPartition -> offset))
+            checkpoint.write(mutable.Map() ++= existing += topicPartition -> offset)
         }
       }
     }
@@ -288,67 +449,146 @@ 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(topicPartition) match {
-        case LogCleaningInProgress =>
-          updateCheckpoints(dataDir, Option(topicPartition, endOffset))
+      inProgress.get(topicPartition) match {
+        case Some(LogCleaningInProgress) =>
+          updateCheckpoints(dataDir, partitionToUpdateOrAdd = Option(topicPartition, endOffset))
           inProgress.remove(topicPartition)
-        case LogCleaningAborted =>
-          inProgress.put(topicPartition, LogCleaningPaused)
+        case Some(LogCleaningAborted) =>
+          inProgress.put(topicPartition, LogCleaningPaused(1))
           pausedCleaningCond.signalAll()
+        case None =>
+          throw new IllegalStateException(s"State for partition $topicPartition should exist.")
         case s =>
           throw new IllegalStateException(s"In-progress partition $topicPartition cannot be in $s state.")
       }
     }
   }
 
-  def doneDeleting(topicPartition: TopicPartition): Unit = {
+  def doneDeleting(topicPartitions: Iterable[TopicPartition]): Unit = {
     inLock(lock) {
-      inProgress.remove(topicPartition)
+      topicPartitions.foreach {
+        topicPartition =>
+          inProgress.get(topicPartition) match {
+            case Some(LogCleaningInProgress) =>
+              inProgress.remove(topicPartition)
+            case Some(LogCleaningAborted) =>
+              inProgress.put(topicPartition, LogCleaningPaused(1))
+              pausedCleaningCond.signalAll()
+            case None =>
+              throw new IllegalStateException(s"State for partition $topicPartition should exist.")
+            case s =>
+              throw new IllegalStateException(s"In-progress partition $topicPartition cannot be in $s state.")
+          }
+      }
+    }
+  }
+
+  /**
+   * Returns an immutable set of the uncleanable partitions for a given log directory
+   * Only used for testing
+   */
+  private[log] def uncleanablePartitions(logDir: String): Set[TopicPartition] = {
+    var partitions: Set[TopicPartition] = Set()
+    inLock(lock) { partitions ++= uncleanablePartitions.getOrElse(logDir, partitions) }
+    partitions
+  }
+
+  def markPartitionUncleanable(logDir: String, partition: TopicPartition): Unit = {
+    inLock(lock) {
+      uncleanablePartitions.get(logDir) match {
+        case Some(partitions) =>
+          partitions.add(partition)
+        case None =>
+          uncleanablePartitions.put(logDir, mutable.Set(partition))
+      }
+    }
+  }
+
+  private def isUncleanablePartition(log: Log, topicPartition: TopicPartition): Boolean = {
+    inLock(lock) {
+      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.
+    val (firstDirtyOffset, forceUpdateCheckpoint) = {
+      val logStartOffset = log.logStartOffset
+      val checkpointDirtyOffset = lastCleanOffset.getOrElse(logStartOffset)
+
+      if (checkpointDirtyOffset < 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
+          warn(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)
       }
     }
 
-    // dirty log segments
-    val dirtyNonActiveSegments = log.logSegments(firstDirtyOffset, log.activeSegment.baseOffset)
-
-    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
@@ -356,23 +596,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.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}")
 
-    (firstDirtyOffset, firstUncleanableDirtyOffset)
+    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
+
+    (firstUncleanableOffset, cleanableBytes)
+  }
+
 }
diff --git a/core/src/main/scala/kafka/log/LogConfig.scala b/core/src/main/scala/kafka/log/LogConfig.scala
index 57b4112bcab0a..c0a6d9f0c2d76 100755
--- a/core/src/main/scala/kafka/log/LogConfig.scala
+++ b/core/src/main/scala/kafka/log/LogConfig.scala
@@ -19,17 +19,17 @@ package kafka.log
 
 import java.util.{Collections, Locale, Properties}
 
-import scala.collection.JavaConverters._
-import kafka.api.ApiVersion
-import kafka.message.{BrokerCompressionCodec, Message}
+import scala.jdk.CollectionConverters._
+import kafka.api.{ApiVersion, ApiVersionValidator}
+import kafka.message.BrokerCompressionCodec
 import kafka.server.{KafkaConfig, ThrottledReplicaListValidator}
 import kafka.utils.Implicits._
 import org.apache.kafka.common.errors.InvalidConfigurationException
 import org.apache.kafka.common.config.{AbstractConfig, ConfigDef, TopicConfig}
-import org.apache.kafka.common.record.TimestampType
+import org.apache.kafka.common.record.{LegacyRecord, TimestampType}
 import org.apache.kafka.common.utils.Utils
 
-import scala.collection.mutable
+import scala.collection.{Map, mutable}
 import org.apache.kafka.common.config.ConfigDef.{ConfigKey, ValidList, Validator}
 
 object Defaults {
@@ -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 " +
@@ -63,9 +64,11 @@ object Defaults {
   val LeaderReplicationThrottledReplicas = Collections.emptyList[String]()
   val FollowerReplicationThrottledReplicas = Collections.emptyList[String]()
   val MaxIdMapSnapshots = kafka.server.Defaults.MaxIdMapSnapshots
+  val MessageDownConversionEnable = kafka.server.Defaults.MessageDownConversionEnable
 }
 
-case class LogConfig(props: java.util.Map[_, _]) extends AbstractConfig(LogConfig.configDef, props, false) {
+case class LogConfig(props: java.util.Map[_, _], overriddenConfigs: Set[String] = Set.empty)
+  extends AbstractConfig(LogConfig.configDef, props, false) {
   /**
    * Important note: Any configuration parameter that is passed along from KafkaConfig to LogConfig
    * should also go in [[kafka.server.KafkaServer.copyKafkaConfigToLog]].
@@ -83,6 +86,7 @@ case class LogConfig(props: java.util.Map[_, _]) extends AbstractConfig(LogConfi
   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)
@@ -95,15 +99,21 @@ case class LogConfig(props: java.util.Map[_, _]) extends AbstractConfig(LogConfi
   val messageTimestampDifferenceMaxMs = getLong(LogConfig.MessageTimestampDifferenceMaxMsProp).longValue
   val LeaderReplicationThrottledReplicas = getList(LogConfig.LeaderReplicationThrottledReplicasProp)
   val FollowerReplicationThrottledReplicas = getList(LogConfig.FollowerReplicationThrottledReplicasProp)
+  val messageDownConversionEnable = getBoolean(LogConfig.MessageDownConversionEnableProp)
 
   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(4, (config: String) => "topicconfigs_" + config))
   }
 
   val SegmentBytesProp = TopicConfig.SEGMENT_BYTES_CONFIG
@@ -118,6 +128,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
@@ -130,6 +141,7 @@ object LogConfig {
   val MessageFormatVersionProp = TopicConfig.MESSAGE_FORMAT_VERSION_CONFIG
   val MessageTimestampTypeProp = TopicConfig.MESSAGE_TIMESTAMP_TYPE_CONFIG
   val MessageTimestampDifferenceMaxMsProp = TopicConfig.MESSAGE_TIMESTAMP_DIFFERENCE_MAX_MS_CONFIG
+  val MessageDownConversionEnableProp = TopicConfig.MESSAGE_DOWNCONVERSION_ENABLE_CONFIG
 
   // Leave these out of TopicConfig for now as they are replication quota configs
   val LeaderReplicationThrottledReplicasProp = "leader.replication.throttled.replicas"
@@ -148,6 +160,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
@@ -157,6 +170,7 @@ object LogConfig {
   val MessageFormatVersionDoc = TopicConfig.MESSAGE_FORMAT_VERSION_DOC
   val MessageTimestampTypeDoc = TopicConfig.MESSAGE_TIMESTAMP_TYPE_DOC
   val MessageTimestampDifferenceMaxMsDoc = TopicConfig.MESSAGE_TIMESTAMP_DIFFERENCE_MAX_MS_DOC
+  val MessageDownConversionEnableDoc = TopicConfig.MESSAGE_DOWNCONVERSION_ENABLE_DOC
 
   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 " +
@@ -167,9 +181,17 @@ 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 {
+  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 = {
@@ -192,11 +214,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)
       }
     }
@@ -204,6 +227,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._
@@ -211,9 +237,9 @@ object LogConfig {
     import org.apache.kafka.common.config.ConfigDef.ValidString._
 
     new LogConfigDef()
-      .define(SegmentBytesProp, INT, Defaults.SegmentSize, atLeast(Message.MinMessageOverhead), MEDIUM,
+      .define(SegmentBytesProp, INT, Defaults.SegmentSize, atLeast(LegacyRecord.RECORD_OVERHEAD_V0), MEDIUM,
         SegmentSizeDoc, KafkaConfig.LogSegmentBytesProp)
-      .define(SegmentMsProp, LONG, Defaults.SegmentMs, atLeast(0), MEDIUM, SegmentMsDoc,
+      .define(SegmentMsProp, LONG, Defaults.SegmentMs, atLeast(1), MEDIUM, SegmentMsDoc,
         KafkaConfig.LogRollTimeMillisProp)
       .define(SegmentJitterMsProp, LONG, Defaults.SegmentJitterMs, atLeast(0), MEDIUM, SegmentJitterMsDoc,
         KafkaConfig.LogRollTimeJitterMillisProp)
@@ -227,7 +253,7 @@ object LogConfig {
       .define(RetentionBytesProp, LONG, Defaults.RetentionSize, MEDIUM, RetentionSizeDoc,
         KafkaConfig.LogRetentionBytesProp)
       // can be negative. See kafka.log.LogManager.cleanupExpiredSegments
-      .define(RetentionMsProp, LONG, Defaults.RetentionMs, MEDIUM, RetentionMsDoc,
+      .define(RetentionMsProp, LONG, Defaults.RetentionMs, atLeast(-1), MEDIUM, RetentionMsDoc,
         KafkaConfig.LogRetentionTimeMillisProp)
       .define(MaxMessageBytesProp, INT, Defaults.MaxMessageSize, atLeast(0), MEDIUM, MaxMessageSizeDoc,
         KafkaConfig.MessageMaxBytesProp)
@@ -237,6 +263,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,
@@ -251,9 +279,9 @@ object LogConfig {
         MEDIUM, CompressionTypeDoc, KafkaConfig.CompressionTypeProp)
       .define(PreAllocateEnableProp, BOOLEAN, Defaults.PreAllocateEnable, MEDIUM, PreAllocateEnableDoc,
         KafkaConfig.LogPreAllocateProp)
-      .define(MessageFormatVersionProp, STRING, Defaults.MessageFormatVersion, MEDIUM, MessageFormatVersionDoc,
+      .define(MessageFormatVersionProp, STRING, Defaults.MessageFormatVersion, ApiVersionValidator, MEDIUM, MessageFormatVersionDoc,
         KafkaConfig.LogMessageFormatVersionProp)
-      .define(MessageTimestampTypeProp, STRING, Defaults.MessageTimestampType, MEDIUM, MessageTimestampTypeDoc,
+      .define(MessageTimestampTypeProp, STRING, Defaults.MessageTimestampType, in("CreateTime", "LogAppendTime"), MEDIUM, MessageTimestampTypeDoc,
         KafkaConfig.LogMessageTimestampTypeProp)
       .define(MessageTimestampDifferenceMaxMsProp, LONG, Defaults.MessageTimestampDifferenceMaxMs,
         atLeast(0), MEDIUM, MessageTimestampDifferenceMaxMsDoc, KafkaConfig.LogMessageTimestampDifferenceMaxMsProp)
@@ -261,6 +289,8 @@ object LogConfig {
         LeaderReplicationThrottledReplicasDoc, LeaderReplicationThrottledReplicasProp)
       .define(FollowerReplicationThrottledReplicasProp, LIST, Defaults.FollowerReplicationThrottledReplicas, ThrottledReplicaListValidator, MEDIUM,
         FollowerReplicationThrottledReplicasDoc, FollowerReplicationThrottledReplicasProp)
+      .define(MessageDownConversionEnableProp, BOOLEAN, Defaults.MessageDownConversionEnable, LOW,
+        MessageDownConversionEnableDoc, KafkaConfig.LogMessageDownConversionEnableProp)
   }
 
   def apply(): LogConfig = LogConfig(new Properties())
@@ -269,32 +299,80 @@ object LogConfig {
 
   def serverConfigName(configName: String): Option[String] = configDef.serverConfigName(configName)
 
+  def configType(configName: String): Option[ConfigDef.Type] = {
+    Option(configDef.configKeys.get(configName)).map(_.`type`)
+  }
+
   /**
    * Create a log config instance using the given properties and defaults
    */
   def fromProps(defaults: java.util.Map[_ <: Object, _ <: Object], overrides: Properties): LogConfig = {
     val props = new Properties()
-    defaults.asScala.foreach { case (k, v) => props.put(k, v) }
+    defaults.forEach { (k, v) => props.put(k, v) }
     props ++= overrides
-    LogConfig(props)
+    val overriddenKeys = overrides.keySet.asScala.map(_.asInstanceOf[String]).toSet
+    new LogConfig(props, overriddenKeys)
   }
 
   /**
    * 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)
   }
 
+  /**
+   * Map topic config to the broker config with highest priority. Some of these have additional synonyms
+   * that can be obtained using [[kafka.server.DynamicBrokerConfig#brokerConfigSynonyms]]
+   */
+  val TopicConfigSynonyms = Map(
+    SegmentBytesProp -> KafkaConfig.LogSegmentBytesProp,
+    SegmentMsProp -> KafkaConfig.LogRollTimeMillisProp,
+    SegmentJitterMsProp -> KafkaConfig.LogRollTimeJitterMillisProp,
+    SegmentIndexBytesProp -> KafkaConfig.LogIndexSizeMaxBytesProp,
+    FlushMessagesProp -> KafkaConfig.LogFlushIntervalMessagesProp,
+    FlushMsProp -> KafkaConfig.LogFlushIntervalMsProp,
+    RetentionBytesProp -> KafkaConfig.LogRetentionBytesProp,
+    RetentionMsProp -> KafkaConfig.LogRetentionTimeMillisProp,
+    MaxMessageBytesProp -> KafkaConfig.MessageMaxBytesProp,
+    IndexIntervalBytesProp -> KafkaConfig.LogIndexIntervalBytesProp,
+    DeleteRetentionMsProp -> KafkaConfig.LogCleanerDeleteRetentionMsProp,
+    MinCompactionLagMsProp -> KafkaConfig.LogCleanerMinCompactionLagMsProp,
+    MaxCompactionLagMsProp -> KafkaConfig.LogCleanerMaxCompactionLagMsProp,
+    FileDeleteDelayMsProp -> KafkaConfig.LogDeleteDelayMsProp,
+    MinCleanableDirtyRatioProp -> KafkaConfig.LogCleanerMinCleanRatioProp,
+    CleanupPolicyProp -> KafkaConfig.LogCleanupPolicyProp,
+    UncleanLeaderElectionEnableProp -> KafkaConfig.UncleanLeaderElectionEnableProp,
+    MinInSyncReplicasProp -> KafkaConfig.MinInSyncReplicasProp,
+    CompressionTypeProp -> KafkaConfig.CompressionTypeProp,
+    PreAllocateEnableProp -> KafkaConfig.LogPreAllocateProp,
+    MessageFormatVersionProp -> KafkaConfig.LogMessageFormatVersionProp,
+    MessageTimestampTypeProp -> KafkaConfig.LogMessageTimestampTypeProp,
+    MessageTimestampDifferenceMaxMsProp -> KafkaConfig.LogMessageTimestampDifferenceMaxMsProp,
+    MessageDownConversionEnableProp -> KafkaConfig.LogMessageDownConversionEnableProp
+  )
+
 }
diff --git a/core/src/main/scala/kafka/log/LogManager.scala b/core/src/main/scala/kafka/log/LogManager.scala
index 9a61be393c495..cd3024634e2b2 100755
--- a/core/src/main/scala/kafka/log/LogManager.scala
+++ b/core/src/main/scala/kafka/log/LogManager.scala
@@ -20,21 +20,21 @@ package kafka.log
 import java.io._
 import java.nio.file.Files
 import java.util.concurrent._
+import java.util.concurrent.atomic.AtomicInteger
 
-import com.yammer.metrics.core.Gauge
-import kafka.common.KafkaException
 import kafka.metrics.KafkaMetricsGroup
 import kafka.server.checkpoints.OffsetCheckpointFile
 import kafka.server.{BrokerState, RecoveringFromUncleanShutdown, _}
 import kafka.utils._
 import kafka.zk.KafkaZkClient
-import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.{KafkaException, TopicPartition}
 import org.apache.kafka.common.utils.Time
 import org.apache.kafka.common.errors.{KafkaStorageException, LogDirNotFoundException}
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 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.
@@ -50,9 +50,9 @@ import scala.collection.mutable.ArrayBuffer
 class LogManager(logDirs: Seq[File],
                  initialOfflineDirs: Seq[File],
                  val topicConfigs: Map[String, LogConfig], // note that this doesn't get updated after creation
-                 val defaultConfig: LogConfig,
+                 val initialDefaultConfig: LogConfig,
                  val cleanerConfig: CleanerConfig,
-                 ioThreads: Int,
+                 recoveryThreadsPerDataDir: Int,
                  val flushCheckMs: Long,
                  val flushRecoveryOffsetCheckpointMs: Long,
                  val flushStartOffsetCheckpointMs: Long,
@@ -75,9 +75,25 @@ class LogManager(logDirs: Seq[File],
   // from one log directory to another log directory on the same broker. The directory of the future log will be renamed
   // to replace the current log of the partition after the future log catches up with the current log
   private val futureLogs = new Pool[TopicPartition, Log]()
-  private val logsToBeDeleted = new LinkedBlockingQueue[Log]()
+  // Each element in the queue contains the log object to be deleted and the time it is scheduled for deletion.
+  private val logsToBeDeleted = new LinkedBlockingQueue[(Log, Long)]()
 
   private val _liveLogDirs: ConcurrentLinkedQueue[File] = createAndValidateLogDirs(logDirs, initialOfflineDirs)
+  @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
+  }
+
+  def currentDefaultConfig: LogConfig = _currentDefaultConfig
 
   def liveLogDirs: Seq[File] = {
     if (_liveLogDirs.size == logDirs.size)
@@ -95,35 +111,25 @@ class LogManager(logDirs: Seq[File],
   private val preferredLogDirs = new ConcurrentHashMap[TopicPartition, String]()
 
   private def offlineLogDirs: Iterable[File] = {
-    val logDirsSet = mutable.Set[File](logDirs: _*)
-    _liveLogDirs.asScala.foreach(logDirsSet -=)
+    val logDirsSet = mutable.Set[File]() ++= logDirs
+    _liveLogDirs.forEach(dir => logDirsSet -= dir)
     logDirsSet
   }
 
   loadLogs()
 
-  // public, so we can access this from kafka.admin.DeleteTopicTest
-  val cleaner: LogCleaner =
-    if(cleanerConfig.enableCleaner)
+  private[kafka] val cleaner: LogCleaner =
+    if (cleanerConfig.enableCleaner)
       new LogCleaner(cleanerConfig, liveLogDirs, currentLogs, logDirFailureChannel, time = time)
     else
       null
 
-  val offlineLogDirectoryCount = newGauge(
-    "OfflineLogDirectoryCount",
-    new Gauge[Int] {
-      def value = offlineLogDirs.size
-    }
-  )
+  newGauge("OfflineLogDirectoryCount", () => offlineLogDirs.size)
 
   for (dir <- logDirs) {
-    newGauge(
-      "LogDirectoryOffline",
-      new Gauge[Int] {
-        def value = if (_liveLogDirs.contains(dir)) 0 else 1
-      },
-      Map("logDirectory" -> dir.getAbsolutePath)
-    )
+    newGauge("LogDirectoryOffline",
+      () => if (_liveLogDirs.contains(dir)) 0 else 1,
+      Map("logDirectory" -> dir.getAbsolutePath))
   }
 
   /**
@@ -135,10 +141,8 @@ class LogManager(logDirs: Seq[File],
    * 
    */
   private def createAndValidateLogDirs(dirs: Seq[File], initialOfflineDirs: Seq[File]): ConcurrentLinkedQueue[File] = {
-    if(dirs.map(_.getCanonicalPath).toSet.size < dirs.size)
-      throw new KafkaException("Duplicate log directory found: " + dirs.mkString(", "))
-
     val liveLogDirs = new ConcurrentLinkedQueue[File]()
+    val canonicalPaths = mutable.HashSet.empty[String]
 
     for (dir <- dirs) {
       try {
@@ -146,13 +150,21 @@ class LogManager(logDirs: Seq[File],
           throw new IOException(s"Failed to load ${dir.getAbsolutePath} during broker startup")
 
         if (!dir.exists) {
-          info("Log directory '" + dir.getAbsolutePath + "' not found, creating it.")
+          info(s"Log directory ${dir.getAbsolutePath} not found, creating it.")
           val created = dir.mkdirs()
           if (!created)
-            throw new IOException("Failed to create data directory " + dir.getAbsolutePath)
+            throw new IOException(s"Failed to create data directory ${dir.getAbsolutePath}")
         }
         if (!dir.isDirectory || !dir.canRead)
-          throw new IOException(dir.getAbsolutePath + " is not a readable log directory.")
+          throw new IOException(s"${dir.getAbsolutePath} is not a readable log directory.")
+
+        // getCanonicalPath() throws IOException if a file system query fails or if the path is invalid (e.g. contains
+        // the Nul character). Since there's no easy way to distinguish between the two cases, we treat them the same
+        // and mark the log directory as offline.
+        if (!canonicalPaths.add(dir.getCanonicalPath))
+          throw new KafkaException(s"Duplicate log directory found: ${dirs.mkString(", ")}")
+
+
         liveLogDirs.add(dir)
       } catch {
         case e: IOException =>
@@ -160,16 +172,25 @@ class LogManager(logDirs: Seq[File],
       }
     }
     if (liveLogDirs.isEmpty) {
-      fatal(s"Shutdown broker because none of the specified log dirs from " + dirs.mkString(", ") + " can be created or validated")
+      fatal(s"Shutdown broker because none of the specified log dirs from ${dirs.mkString(", ")} can be created or validated")
       Exit.halt(1)
     }
 
     liveLogDirs
   }
 
-  // dir should be an absolute path
-  def handleLogDirFailure(dir: String) {
-    info(s"Stopping serving logs in dir $dir")
+  def resizeRecoveryThreadPool(newSize: Int): Unit = {
+    info(s"Resizing recovery thread pool size for each data dir from $numRecoveryThreadsPerDataDir to $newSize")
+    numRecoveryThreadsPerDataDir = newSize
+  }
+
+  /**
+   * The log directory failure handler. It will stop log cleaning in that directory.
+   *
+   * @param dir        the absolute path of the log directory
+   */
+  def handleLogDirFailure(dir: String): Unit = {
+    warn(s"Stopping serving logs in dir $dir")
     logCreationOrDeletionLock synchronized {
       _liveLogDirs.remove(new File(dir))
       if (_liveLogDirs.isEmpty) {
@@ -182,29 +203,24 @@ class LogManager(logDirs: Seq[File],
       if (cleaner != null)
         cleaner.handleLogDirFailure(dir)
 
-      val offlineCurrentTopicPartitions = currentLogs.collect {
-        case (tp, log) if log.dir.getParent == dir => tp
-      }
-      offlineCurrentTopicPartitions.foreach { topicPartition => {
-        val removedLog = currentLogs.remove(topicPartition)
-        if (removedLog != null) {
-          removedLog.closeHandlers()
-          removedLog.removeLogMetrics()
+      def removeOfflineLogs(logs: Pool[TopicPartition, Log]): Iterable[TopicPartition] = {
+        val offlineTopicPartitions: Iterable[TopicPartition] = logs.collect {
+          case (tp, log) if log.parentDir == dir => tp
         }
-      }}
+        offlineTopicPartitions.foreach { topicPartition => {
+          val removedLog = removeLogAndMetrics(logs, topicPartition)
+          removedLog.foreach {
+            log => log.closeHandlers()
+          }
+        }}
 
-      val offlineFutureTopicPartitions = futureLogs.collect {
-        case (tp, log) if log.dir.getParent == dir => tp
+        offlineTopicPartitions
       }
-      offlineFutureTopicPartitions.foreach { topicPartition => {
-        val removedLog = futureLogs.remove(topicPartition)
-        if (removedLog != null) {
-          removedLog.closeHandlers()
-          removedLog.removeLogMetrics()
-        }
-      }}
 
-      info(s"Logs for partitions ${offlineCurrentTopicPartitions.mkString(",")} are offline and " +
+      val offlineCurrentTopicPartitions = removeOfflineLogs(currentLogs)
+      val offlineFutureTopicPartitions = removeOfflineLogs(futureLogs)
+
+      warn(s"Logs for partitions ${offlineCurrentTopicPartitions.mkString(",")} are offline and " +
            s"logs for future partitions ${offlineFutureTopicPartitions.mkString(",")} are offline due to failure on log directory $dir")
       dirLocks.filter(_.file.getParent == dir).foreach(dir => CoreUtils.swallow(dir.destroy(), this))
     }
@@ -229,10 +245,19 @@ class LogManager(logDirs: Seq[File],
     }
   }
 
-  private def loadLog(logDir: File, recoveryPoints: Map[TopicPartition, Long], logStartOffsets: Map[TopicPartition, Long]): Unit = {
-    debug("Loading log '" + logDir.getName + "'")
+  private def addLogToBeDeleted(log: Log): Unit = {
+    this.logsToBeDeleted.add((log, time.milliseconds()))
+  }
+
+  // Only for testing
+  private[log] def hasLogsToBeDeleted: Boolean = !logsToBeDeleted.isEmpty
+
+  private[log] def loadLog(logDir: File,
+                      hadCleanShutdown: Boolean,
+                      recoveryPoints: Map[TopicPartition, Long],
+                      logStartOffsets: Map[TopicPartition, Long]): Log = {
     val topicPartition = Log.parseTopicPartitionName(logDir)
-    val config = topicConfigs.getOrElse(topicPartition.topic, defaultConfig)
+    val config = topicConfigs.getOrElse(topicPartition.topic, currentDefaultConfig)
     val logRecoveryPoint = recoveryPoints.getOrElse(topicPartition, 0L)
     val logStartOffset = logStartOffsets.getOrElse(topicPartition, 0L)
 
@@ -246,10 +271,11 @@ class LogManager(logDirs: Seq[File],
       scheduler = scheduler,
       time = time,
       brokerTopicStats = brokerTopicStats,
-      logDirFailureChannel = logDirFailureChannel)
+      logDirFailureChannel = logDirFailureChannel,
+      lastShutdownClean = hadCleanShutdown)
 
     if (logDir.getName.endsWith(Log.DeleteDirSuffix)) {
-      this.logsToBeDeleted.add(log)
+      addLogToBeDeleted(log)
     } else {
       val previous = {
         if (log.isFuture)
@@ -259,7 +285,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 " +
@@ -267,101 +293,115 @@ class LogManager(logDirs: Seq[File],
             s"for this partition. It is recommended to delete the partition in the log directory that is known to have failed recently.")
       }
     }
+
+    log
   }
 
   /**
    * Recover and load all logs in the given data directories
    */
-  private def loadLogs(): Unit = {
-    info("Loading logs.")
-    val startMs = time.milliseconds
+  private[log] def loadLogs(): Unit = {
+    info(s"Loading logs from log dirs $liveLogDirs")
+    val startMs = time.hiResClockMs()
     val threadPools = ArrayBuffer.empty[ExecutorService]
     val offlineDirs = mutable.Set.empty[(String, IOException)]
-    val jobs = mutable.Map.empty[File, Seq[Future[_]]]
+    val jobs = ArrayBuffer.empty[Seq[Future[_]]]
+    var numTotalLogs = 0
 
     for (dir <- liveLogDirs) {
+      val logDirAbsolutePath = dir.getAbsolutePath
+      var hadCleanShutdown: Boolean = false
       try {
-        val pool = Executors.newFixedThreadPool(ioThreads)
+        val pool = Executors.newFixedThreadPool(numRecoveryThreadsPerDataDir)
         threadPools.append(pool)
 
         val cleanShutdownFile = new File(dir, Log.CleanShutdownFile)
-
         if (cleanShutdownFile.exists) {
-          debug(s"Found clean shutdown file. Skipping recovery for all logs in data directory: ${dir.getAbsolutePath}")
+          info(s"Skipping recovery for all logs in $logDirAbsolutePath since clean shutdown file was found")
+          // Cache the clean shutdown status and use that for rest of log loading workflow. Delete the CleanShutdownFile
+          // so that if broker crashes while loading the log, it is considered hard shutdown during the next boot up. KAFKA-10471
+          cleanShutdownFile.delete()
+          hadCleanShutdown = true
         } else {
           // log recovery itself is being performed by `Log` class during initialization
+          info(s"Attempting recovery for all logs in $logDirAbsolutePath since no clean shutdown file was found")
           brokerState.newState(RecoveringFromUncleanShutdown)
         }
 
         var recoveryPoints = Map[TopicPartition, Long]()
         try {
-          recoveryPoints = this.recoveryPointCheckpoints(dir).read
+          recoveryPoints = this.recoveryPointCheckpoints(dir).read()
         } catch {
           case e: Exception =>
-            warn("Error occurred while reading recovery-point-offset-checkpoint file of directory " + dir, e)
-            warn("Resetting the recovery checkpoint to 0")
+            warn(s"Error occurred while reading recovery-point-offset-checkpoint file of directory " +
+              s"$logDirAbsolutePath, resetting the recovery checkpoint to 0", e)
         }
 
         var logStartOffsets = Map[TopicPartition, Long]()
         try {
-          logStartOffsets = this.logStartOffsetCheckpoints(dir).read
+          logStartOffsets = this.logStartOffsetCheckpoints(dir).read()
         } catch {
           case e: Exception =>
-            warn("Error occurred while reading log-start-offset-checkpoint file of directory " + dir, e)
+            warn(s"Error occurred while reading log-start-offset-checkpoint file of directory " +
+              s"$logDirAbsolutePath, resetting to the base offset of the first segment", e)
         }
 
-        val jobsForDir = for {
-          dirContent <- Option(dir.listFiles).toList
-          logDir <- dirContent if logDir.isDirectory
-        } yield {
-          CoreUtils.runnable {
+        val logsToLoad = Option(dir.listFiles).getOrElse(Array.empty).filter(_.isDirectory)
+        val numLogsLoaded = new AtomicInteger(0)
+        numTotalLogs += logsToLoad.length
+
+        val jobsForDir = logsToLoad.map { logDir =>
+          val runnable: Runnable = () => {
             try {
-              loadLog(logDir, recoveryPoints, logStartOffsets)
+              debug(s"Loading log $logDir")
+
+              val logLoadStartMs = time.hiResClockMs()
+              val log = loadLog(logDir, hadCleanShutdown, recoveryPoints, logStartOffsets)
+              val logLoadDurationMs = time.hiResClockMs() - logLoadStartMs
+              val currentNumLoaded = numLogsLoaded.incrementAndGet()
+
+              info(s"Completed load of $log with ${log.numberOfSegments} segments in ${logLoadDurationMs}ms " +
+                s"($currentNumLoaded/${logsToLoad.length} loaded in $logDirAbsolutePath)")
             } catch {
               case e: IOException =>
-                offlineDirs.add((dir.getAbsolutePath, e))
-                error("Error while loading log dir " + dir.getAbsolutePath, e)
+                offlineDirs.add((logDirAbsolutePath, e))
+                error(s"Error while loading log dir $logDirAbsolutePath", e)
             }
           }
+          runnable
         }
-        jobs(cleanShutdownFile) = jobsForDir.map(pool.submit)
+
+        jobs += jobsForDir.map(pool.submit)
       } catch {
         case e: IOException =>
-          offlineDirs.add((dir.getAbsolutePath, e))
-          error("Error while loading log dir " + dir.getAbsolutePath, e)
+          offlineDirs.add((logDirAbsolutePath, e))
+          error(s"Error while loading log dir $logDirAbsolutePath", e)
       }
     }
 
     try {
-      for ((cleanShutdownFile, dirJobs) <- jobs) {
+      for (dirJobs <- jobs) {
         dirJobs.foreach(_.get)
-        try {
-          cleanShutdownFile.delete()
-        } catch {
-          case e: IOException =>
-            offlineDirs.add((cleanShutdownFile.getParent, e))
-            error(s"Error while deleting the clean shutdown file $cleanShutdownFile", e)
-        }
       }
 
       offlineDirs.foreach { case (dir, e) =>
-        logDirFailureChannel.maybeAddOfflineLogDir(dir, s"Error while deleting the clean shutdown file in dir $dir", e)
+        logDirFailureChannel.maybeAddOfflineLogDir(dir, s"Error while loading log dir $dir", e)
       }
     } catch {
       case e: ExecutionException =>
-        error("There was an error in one of the threads during logs loading: " + e.getCause)
+        error(s"There was an error in one of the threads during logs loading: ${e.getCause}")
         throw e.getCause
     } finally {
       threadPools.foreach(_.shutdown())
     }
 
-    info(s"Logs loading complete in ${time.milliseconds - startMs} ms.")
+    info(s"Loaded $numTotalLogs logs in ${time.hiResClockMs() - startMs}ms.")
   }
 
   /**
    *  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))
@@ -386,11 +426,10 @@ class LogManager(logDirs: Seq[File],
                          delay = InitialTaskDelayMs,
                          period = flushStartOffsetCheckpointMs,
                          TimeUnit.MILLISECONDS)
-      scheduler.schedule("kafka-delete-logs",
+      scheduler.schedule("kafka-delete-logs", // will be rescheduled after each delete logs with a dynamic period
                          deleteLogs _,
                          delay = InitialTaskDelayMs,
-                         period = defaultConfig.fileDeleteDelayMs,
-                         TimeUnit.MILLISECONDS)
+                         unit = TimeUnit.MILLISECONDS)
     }
     if (cleanerConfig.enableCleaner)
       cleaner.startup()
@@ -399,7 +438,7 @@ class LogManager(logDirs: Seq[File],
   /**
    * Close all the logs
    */
-  def shutdown() {
+  def shutdown(): Unit = {
     info("Shutting down.")
 
     removeMetric("OfflineLogDirectoryCount")
@@ -415,21 +454,24 @@ class LogManager(logDirs: Seq[File],
       CoreUtils.swallow(cleaner.shutdown(), this)
     }
 
+    val localLogsByDir = logsByDir
+
     // close logs in each dir
     for (dir <- liveLogDirs) {
-      debug("Flushing and closing logs at " + dir)
+      debug(s"Flushing and closing logs at $dir")
 
-      val pool = Executors.newFixedThreadPool(ioThreads)
+      val pool = Executors.newFixedThreadPool(numRecoveryThreadsPerDataDir)
       threadPools.append(pool)
 
-      val logsInDir = logsByDir.getOrElse(dir.toString, Map()).values
+      val logs = logsInDir(localLogsByDir, dir).values
 
-      val jobsForDir = logsInDir map { log =>
-        CoreUtils.runnable {
+      val jobsForDir = logs.map { log =>
+        val runnable: Runnable = () => {
           // flush the log to ensure latest possible recovery point
           log.flush()
           log.close()
         }
+        runnable
       }
 
       jobs(dir) = jobsForDir.map(pool.submit).toSeq
@@ -437,23 +479,30 @@ class LogManager(logDirs: Seq[File],
 
     try {
       for ((dir, dirJobs) <- jobs) {
-        dirJobs.foreach(_.get)
+        val hasErrors = dirJobs.exists  { future =>
+          Try(future.get) match {
+            case Success(_) => false
+            case Failure(e) =>
+              warn(s"There was an error in one of the threads during LogManager shutdown: ${e.getCause}")
+              true
+          }
+        }
+
+        if (!hasErrors) {
+          val logs = logsInDir(localLogsByDir, dir)
 
-        // update the last flush point
-        debug("Updating recovery points at " + dir)
-        checkpointLogRecoveryOffsetsInDir(dir)
+          // update the last flush point
+          debug(s"Updating recovery points at $dir")
+          checkpointRecoveryOffsetsInDir(dir, logs)
 
-        debug("Updating log start offsets at " + dir)
-        checkpointLogStartOffsetsInDir(dir)
+          debug(s"Updating log start offsets at $dir")
+          checkpointLogStartOffsetsInDir(dir, logs)
 
-        // mark that the shutdown was clean by creating marker file
-        debug("Writing clean shutdown marker at " + dir)
-        CoreUtils.swallow(Files.createFile(new File(dir, Log.CleanShutdownFile).toPath), this)
+          // mark that the shutdown was clean by creating marker file
+          debug(s"Writing clean shutdown marker at $dir")
+          CoreUtils.swallow(Files.createFile(new File(dir, Log.CleanShutdownFile).toPath), this)
+        }
       }
-    } catch {
-      case e: ExecutionException =>
-        error("There was an error in one of the threads during LogManager shutdown: " + e.getCause)
-        throw e.getCause
     } finally {
       threadPools.foreach(_.shutdown())
       // regardless of whether the close succeeded, we need to unlock the data directories
@@ -469,8 +518,8 @@ 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) {
-    var truncated = false
+  def truncateTo(partitionOffsets: Map[TopicPartition, Long], isFuture: Boolean): Unit = {
+    val affectedLogs = ArrayBuffer.empty[Log]
     for ((topicPartition, truncateOffset) <- partitionOffsets) {
       val log = {
         if (isFuture)
@@ -480,24 +529,25 @@ class LogManager(logDirs: Seq[File],
       }
       // If the log does not exist, skip it
       if (log != null) {
-        //May need to abort and pause the cleaning of the log, and resume after truncation is done.
-        val needToStopCleaner = cleaner != null && truncateOffset < log.activeSegment.baseOffset
+        // May need to abort and pause the cleaning of the log, and resume after truncation is done.
+        val needToStopCleaner = truncateOffset < log.activeSegment.baseOffset
         if (needToStopCleaner && !isFuture)
-          cleaner.abortAndPauseCleaning(topicPartition)
+          abortAndPauseCleaning(topicPartition)
         try {
           if (log.truncateTo(truncateOffset))
-            truncated = true
+            affectedLogs += log
           if (needToStopCleaner && !isFuture)
-            cleaner.maybeTruncateCheckpoint(log.dir.getParentFile, topicPartition, log.activeSegment.baseOffset)
+            maybeTruncateCleanerCheckpointToActiveSegmentBaseOffset(log, topicPartition)
         } finally {
           if (needToStopCleaner && !isFuture)
-            cleaner.resumeCleaning(topicPartition)
+            resumeCleaning(topicPartition)
         }
       }
     }
 
-    if (truncated)
-      checkpointLogRecoveryOffsets()
+    for (dir <- affectedLogs.map(_.parentDirFile).distinct) {
+      checkpointRecoveryOffsetsInDir(dir)
+    }
   }
 
   /**
@@ -507,7 +557,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)
@@ -516,15 +566,18 @@ class LogManager(logDirs: Seq[File],
     }
     // If the log does not exist, skip it
     if (log != null) {
-        //Abort and pause the cleaning of the log, and resume after truncation is done.
-      if (cleaner != null && !isFuture)
-        cleaner.abortAndPauseCleaning(topicPartition)
-      log.truncateFullyAndStartAt(newOffset)
-      if (cleaner != null && !isFuture) {
-        cleaner.maybeTruncateCheckpoint(log.dir.getParentFile, topicPartition, log.activeSegment.baseOffset)
-        cleaner.resumeCleaning(topicPartition)
+      // Abort and pause the cleaning of the log, and resume after truncation is done.
+      if (!isFuture)
+        abortAndPauseCleaning(topicPartition)
+      try {
+        log.truncateFullyAndStartAt(newOffset)
+        if (!isFuture)
+          maybeTruncateCleanerCheckpointToActiveSegmentBaseOffset(log, topicPartition)
+      } finally {
+        if (!isFuture)
+          resumeCleaning(topicPartition)
       }
-      checkpointLogRecoveryOffsetsInDir(log.dir.getParentFile)
+      checkpointRecoveryOffsetsInDir(log.parentDirFile)
     }
   }
 
@@ -532,70 +585,113 @@ 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() {
-    liveLogDirs.foreach(checkpointLogRecoveryOffsetsInDir)
+  def checkpointLogRecoveryOffsets(): Unit = {
+    val logsByDirCached = logsByDir
+    liveLogDirs.foreach { logDir =>
+      val logsToCheckpoint = logsInDir(logsByDirCached, logDir)
+      checkpointRecoveryOffsetsInDir(logDir, logsToCheckpoint)
+    }
   }
 
   /**
    * 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() {
-    liveLogDirs.foreach(checkpointLogStartOffsetsInDir)
+  def checkpointLogStartOffsets(): Unit = {
+    val logsByDirCached = logsByDir
+    liveLogDirs.foreach { logDir =>
+      checkpointLogStartOffsetsInDir(logDir, logsInDir(logsByDirCached, logDir))
+    }
   }
 
   /**
-   * Make a checkpoint for all logs in provided directory.
+   * Checkpoint recovery offsets for all the logs in logDir.
+   *
+   * @param logDir the directory in which the logs to be checkpointed are
    */
-  private def checkpointLogRecoveryOffsetsInDir(dir: File): Unit = {
-    for {
-      partitionToLog <- logsByDir.get(dir.getAbsolutePath)
-      checkpoint <- recoveryPointCheckpoints.get(dir)
-    } {
-      try {
-        checkpoint.write(partitionToLog.mapValues(_.recoveryPoint))
-        allLogs.foreach(_.deleteSnapshotsAfterRecoveryPointCheckpoint())
-      } catch {
-        case e: IOException =>
-          logDirFailureChannel.maybeAddOfflineLogDir(dir.getAbsolutePath, s"Disk error while writing to recovery point " +
-            s"file in directory $dir", e)
+  // Only for testing
+  private[log] def checkpointRecoveryOffsetsInDir(logDir: File): Unit = {
+    checkpointRecoveryOffsetsInDir(logDir, logsInDir(logDir))
+  }
+
+  /**
+   * Checkpoint recovery offsets for all the provided logs.
+   *
+   * @param logDir the directory in which the logs are
+   * @param logsToCheckpoint the logs to be checkpointed
+   */
+  private def checkpointRecoveryOffsetsInDir(logDir: File, logsToCheckpoint: Map[TopicPartition, Log]): Unit = {
+    try {
+      recoveryPointCheckpoints.get(logDir).foreach { checkpoint =>
+        val recoveryOffsets = logsToCheckpoint.map { case (tp, log) => tp -> log.recoveryPoint }
+        checkpoint.write(recoveryOffsets)
       }
+    } catch {
+      case e: KafkaStorageException =>
+        error(s"Disk error while writing recovery offsets checkpoint in directory $logDir: ${e.getMessage}")
+      case e: IOException =>
+        logDirFailureChannel.maybeAddOfflineLogDir(logDir.getAbsolutePath,
+          s"Disk error while writing recovery offsets checkpoint in directory $logDir: ${e.getMessage}", e)
     }
   }
 
   /**
-   * Checkpoint log start offset for all logs in provided directory.
+   * Checkpoint log start offsets for all the provided logs in the provided directory.
+   *
+   * @param logDir the directory in which logs are checkpointed
+   * @param logsToCheckpoint the logs to be checkpointed
    */
-  private def checkpointLogStartOffsetsInDir(dir: File): Unit = {
-    for {
-      partitionToLog <- logsByDir.get(dir.getAbsolutePath)
-      checkpoint <- logStartOffsetCheckpoints.get(dir)
-    } {
-      try {
-        val logStartOffsets = partitionToLog.filter { case (_, log) =>
-          log.logStartOffset > log.logSegments.head.baseOffset
-        }.mapValues(_.logStartOffset)
+  private def checkpointLogStartOffsetsInDir(logDir: File, logsToCheckpoint: Map[TopicPartition, Log]): Unit = {
+    try {
+      logStartOffsetCheckpoints.get(logDir).foreach { checkpoint =>
+        val logStartOffsets = logsToCheckpoint.collect {
+          case (tp, log) if log.logStartOffset > log.logSegments.head.baseOffset => tp -> log.logStartOffset
+        }
         checkpoint.write(logStartOffsets)
-      } catch {
-        case e: IOException =>
-          logDirFailureChannel.maybeAddOfflineLogDir(dir.getAbsolutePath, s"Disk error while writing to logStartOffset file in directory $dir", e)
       }
+    } catch {
+      case e: KafkaStorageException =>
+        error(s"Disk error while writing log start offsets checkpoint in directory $logDir: ${e.getMessage}")
     }
   }
 
   // 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)
   }
 
+  /**
+   * Abort and pause cleaning of the provided partition and log a message about it.
+   */
   def abortAndPauseCleaning(topicPartition: TopicPartition): Unit = {
-    if (cleaner != null)
+    if (cleaner != null) {
       cleaner.abortAndPauseCleaning(topicPartition)
+      info(s"The cleaning for partition $topicPartition is aborted and paused")
+    }
   }
 
+  /**
+   * Resume cleaning of the provided partition and log a message about it.
+   */
+  private def resumeCleaning(topicPartition: TopicPartition): Unit = {
+    if (cleaner != null) {
+      cleaner.resumeCleaning(Seq(topicPartition))
+      info(s"Cleaning for partition $topicPartition is resumed")
+    }
+  }
+
+  /**
+   * Truncate the cleaner's checkpoint to the based offset of the active segment of
+   * the provided log.
+   */
+  private def maybeTruncateCleanerCheckpointToActiveSegmentBaseOffset(log: Log, topicPartition: TopicPartition): Unit = {
+    if (cleaner != null) {
+      cleaner.maybeTruncateCheckpoint(log.parentDirFile, topicPartition, log.activeSegment.baseOffset)
+    }
+  }
 
   /**
    * Get the log if it exists, otherwise return None
@@ -610,107 +706,189 @@ 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
+   * @param isFuture True if 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("Created log for partition [%s,%d] in %s with properties {%s}."
-            .format(topicPartition.topic,
-              topicPartition.partition,
-              logDir,
-              config.originals.asScala.mkString(", ")))
-          // Remove the preferred log dir since it has already been satisfied
-          preferredLogDirs.remove(topicPartition)
+        val logDir = logDirs
+          .iterator // 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"))
     }
   }
 
   /**
-   *  Delete logs marked for deletion.
+   *  Delete logs marked for deletion. Delete all logs for which `currentDefaultConfig.fileDeleteDelayMs`
+   *  has elapsed after the delete was scheduled. Logs for which this interval has not yet elapsed will be
+   *  considered for deletion in the next iteration of `deleteLogs`. The next iteration will be executed
+   *  after the remaining time for the first log that is not deleted. If there are no more `logsToBeDeleted`,
+   *  `deleteLogs` will be executed after `currentDefaultConfig.fileDeleteDelayMs`.
    */
   private def deleteLogs(): Unit = {
+    var nextDelayMs = 0L
     try {
-      while (!logsToBeDeleted.isEmpty) {
-        val removedLog = logsToBeDeleted.take()
+      def nextDeleteDelayMs: Long = {
+        if (!logsToBeDeleted.isEmpty) {
+          val (_, scheduleTimeMs) = logsToBeDeleted.peek()
+          scheduleTimeMs + currentDefaultConfig.fileDeleteDelayMs - time.milliseconds()
+        } else
+          currentDefaultConfig.fileDeleteDelayMs
+      }
+
+      while ({nextDelayMs = nextDeleteDelayMs; nextDelayMs <= 0}) {
+        val (removedLog, _) = logsToBeDeleted.take()
         if (removedLog != null) {
           try {
             removedLog.delete()
             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)
           }
         }
       }
     } catch {
       case e: Throwable =>
         error(s"Exception in kafka-delete-logs thread.", e)
+    } finally {
+      try {
+        scheduler.schedule("kafka-delete-logs",
+          deleteLogs _,
+          delay = nextDelayMs,
+          unit = TimeUnit.MILLISECONDS)
+      } catch {
+        case e: Throwable =>
+          if (scheduler.isStarted) {
+            // No errors should occur unless scheduler has been shutdown
+            error(s"Failed to schedule next delete in kafka-delete-logs thread", e)
+          }
+      }
     }
   }
 
@@ -725,19 +903,22 @@ class LogManager(logDirs: Seq[File],
       val sourceLog = currentLogs.get(topicPartition)
       val destLog = futureLogs.get(topicPartition)
 
+      info(s"Attempting to replace current log $sourceLog with $destLog for $topicPartition")
       if (sourceLog == null)
         throw new KafkaStorageException(s"The current replica for $topicPartition is offline")
       if (destLog == null)
         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.resumeCleaning(topicPartition)
+        cleaner.alterCheckpointDir(topicPartition, sourceLog.parentDirFile, destLog.parentDirFile)
+        resumeCleaning(topicPartition)
       }
 
       try {
@@ -745,9 +926,12 @@ 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()
-        checkpointLogRecoveryOffsetsInDir(sourceLog.dir.getParentFile)
-        checkpointLogStartOffsetsInDir(sourceLog.dir.getParentFile)
-        logsToBeDeleted.add(sourceLog)
+        val logDir = sourceLog.parentDirFile
+        val logsToCheckpoint = logsInDir(logDir)
+        checkpointRecoveryOffsetsInDir(logDir, logsToCheckpoint)
+        checkpointLogStartOffsetsInDir(logDir, logsToCheckpoint)
+        sourceLog.removeLogMetrics()
+        addLogToBeDeleted(sourceLog)
       } catch {
         case e: KafkaStorageException =>
           // If sourceLog's log directory is offline, we need close its handlers here.
@@ -767,49 +951,97 @@ class LogManager(logDirs: Seq[File],
     *
     * @param topicPartition TopicPartition that needs to be deleted
     * @param isFuture True iff the future log of the specified partition should be deleted
+    * @param checkpoint True if checkpoints must be written
     * @return the removed log
     */
-  def asyncDelete(topicPartition: TopicPartition, isFuture: Boolean = false): Log = {
-    val removedLog: Log = logCreationOrDeletionLock synchronized {
-      if (isFuture)
-        futureLogs.remove(topicPartition)
-      else
-        currentLogs.remove(topicPartition)
+  def asyncDelete(topicPartition: TopicPartition,
+                  isFuture: Boolean = false,
+                  checkpoint: Boolean = true): Option[Log] = {
+    val removedLog: Option[Log] = logCreationOrDeletionLock synchronized {
+      removeLogAndMetrics(if (isFuture) futureLogs else currentLogs, topicPartition)
     }
-    if (removedLog != null) {
-      //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)
-      }
-      removedLog.renameDir(Log.logDeleteDirName(topicPartition))
-      checkpointLogRecoveryOffsetsInDir(removedLog.dir.getParentFile)
-      checkpointLogStartOffsetsInDir(removedLog.dir.getParentFile)
-      logsToBeDeleted.add(removedLog)
-      info(s"Log for partition ${removedLog.topicPartition} is renamed to ${removedLog.dir.getAbsolutePath} and is scheduled for deletion")
-    } else if (offlineLogDirs.nonEmpty) {
-      throw new KafkaStorageException("Failed to delete log for " + topicPartition + " because it may be in one of the offline directories " + offlineLogDirs.mkString(","))
+    removedLog match {
+      case Some(removedLog) =>
+        // 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)
+          if (checkpoint) {
+            cleaner.updateCheckpoints(removedLog.parentDirFile, partitionToRemove = Option(topicPartition))
+          }
+        }
+        removedLog.renameDir(Log.logDeleteDirName(topicPartition))
+        if (checkpoint) {
+          val logDir = removedLog.parentDirFile
+          val logsToCheckpoint = logsInDir(logDir)
+          checkpointRecoveryOffsetsInDir(logDir, logsToCheckpoint)
+          checkpointLogStartOffsetsInDir(logDir, logsToCheckpoint)
+        }
+        addLogToBeDeleted(removedLog)
+        info(s"Log for partition ${removedLog.topicPartition} is renamed to ${removedLog.dir.getAbsolutePath} and is scheduled for deletion")
+
+      case None =>
+        if (offlineLogDirs.nonEmpty) {
+          throw new KafkaStorageException(s"Failed to delete log for ${if (isFuture) "future" else ""} $topicPartition because it may be in one of the offline directories ${offlineLogDirs.mkString(",")}")
+        }
     }
+
     removedLog
   }
 
   /**
-   * 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.
+   * Rename the directories of the given topic-partitions and add them in the queue for
+   * deletion. Checkpoints are updated once all the directories have been renamed.
+   *
+   * @param topicPartitions The set of topic-partitions to delete asynchronously
+   * @param errorHandler The error handler that will be called when a exception for a particular
+   *                     topic-partition is raised
    */
-  private def nextLogDir(): File = {
+  def asyncDelete(topicPartitions: Set[TopicPartition],
+                  errorHandler: (TopicPartition, Throwable) => Unit): Unit = {
+    val logDirs = mutable.Set.empty[File]
+
+    topicPartitions.foreach { topicPartition =>
+      try {
+        getLog(topicPartition).foreach { log =>
+          logDirs += log.parentDirFile
+          asyncDelete(topicPartition, checkpoint = false)
+        }
+        getLog(topicPartition, isFuture = true).foreach { log =>
+          logDirs += log.parentDirFile
+          asyncDelete(topicPartition, isFuture = true, checkpoint = false)
+        }
+      } catch {
+        case e: Throwable => errorHandler(topicPartition, e)
+      }
+    }
+
+    val logsByDirCached = logsByDir
+    logDirs.foreach { logDir =>
+      if (cleaner != null) cleaner.updateCheckpoints(logDir)
+      val logsToCheckpoint = logsInDir(logsByDirCached, logDir)
+      checkpointRecoveryOffsetsInDir(logDir, logsToCheckpoint)
+      checkpointLogStartOffsetsInDir(logDir, logsToCheckpoint)
+    }
+  }
+
+  /**
+   * 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 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).map { case (parent, logs) => parent -> logs.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
     }
   }
 
@@ -817,15 +1049,43 @@ 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
-    for(log <- allLogs; if !log.config.compact) {
-      debug("Garbage collecting '" + log.name + "'")
-      total += log.deleteOldSegments()
+
+    // clean current logs.
+    val deletableLogs = {
+      if (cleaner != null) {
+        // prevent cleaner from working on same partitions when changing cleanup policy
+        cleaner.pauseCleaningForNonCompactedPartitions()
+      } else {
+        currentLogs.filter {
+          case (_, log) => !log.config.compact
+        }
+      }
     }
-    debug("Log cleanup completed. " + total + " files deleted in " +
+
+    try {
+      deletableLogs.foreach {
+        case (topicPartition, log) =>
+          debug(s"Garbage collecting '${log.name}'")
+          total += log.deleteOldSegments()
+
+          val futureLog = futureLogs.get(topicPartition)
+          if (futureLog != null) {
+            // clean future logs
+            debug(s"Garbage collecting future log '${futureLog.name}'")
+            total += futureLog.deleteOldSegments()
+          }
+      }
+    } finally {
+      if (cleaner != null) {
+        cleaner.resumeCleaning(deletableLogs.map(_._1))
+      }
+    }
+
+    debug(s"Log cleanup completed. $total files deleted in " +
                   (time.milliseconds - startMs) / 1000 + " seconds")
   }
 
@@ -835,18 +1095,34 @@ 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, log) =>
-      topicPartition.topic() == topic
-    }.map { case (topicPartition, log) => log }
+    (currentLogs.toList ++ futureLogs.toList).collect {
+      case (topicPartition, log) if topicPartition.topic == topic => log
+    }
   }
 
   /**
    * Map of log dir to logs by topic and partitions in that dir
    */
   private def logsByDir: Map[String, Map[TopicPartition, Log]] = {
-    (this.currentLogs.toList ++ this.futureLogs.toList).groupBy {
-      case (_, log) => log.dir.getParent
-    }.mapValues(_.toMap)
+    // This code is called often by checkpoint processes and is written in a way that reduces
+    // allocations and CPU with many topic partitions.
+    // When changing this code please measure the changes with org.apache.kafka.jmh.server.CheckpointBench
+    val byDir = new mutable.AnyRefMap[String, mutable.AnyRefMap[TopicPartition, Log]]()
+    def addToDir(tp: TopicPartition, log: Log): Unit = {
+      byDir.getOrElseUpdate(log.parentDir, new mutable.AnyRefMap[TopicPartition, Log]()).put(tp, log)
+    }
+    currentLogs.foreachEntry(addToDir)
+    futureLogs.foreachEntry(addToDir)
+    byDir
+  }
+
+  private def logsInDir(dir: File): Map[TopicPartition, Log] = {
+    logsByDir.getOrElse(dir.getAbsolutePath, Map.empty)
+  }
+
+  private def logsInDir(cachedLogsByDir: Map[String, Map[TopicPartition, Log]],
+                        dir: File): Map[TopicPartition, Log] = {
+    cachedLogsByDir.getOrElse(dir.getAbsolutePath, Map.empty)
   }
 
   // logDir should be an absolute path
@@ -867,16 +1143,26 @@ class LogManager(logDirs: Seq[File],
     for ((topicPartition, log) <- currentLogs.toList ++ futureLogs.toList) {
       try {
         val timeSinceLastFlush = time.milliseconds - log.lastFlushTime
-        debug("Checking if flush is needed on " + topicPartition.topic + " flush interval  " + log.config.flushMs +
-              " last flushed " + log.lastFlushTime + " time since last flush: " + timeSinceLastFlush)
+        debug(s"Checking if flush is needed on ${topicPartition.topic} flush interval ${log.config.flushMs}" +
+              s" last flushed ${log.lastFlushTime} time since last flush: $timeSinceLastFlush")
         if(timeSinceLastFlush >= log.config.flushMs)
-          log.flush
+          log.flush()
       } catch {
         case e: Throwable =>
-          error("Error flushing topic " + topicPartition.topic, e)
+          error(s"Error flushing topic ${topicPartition.topic}", e)
       }
     }
   }
+
+  private def removeLogAndMetrics(logs: Pool[TopicPartition, Log], tp: TopicPartition): Option[Log] = {
+    val removedLog = logs.remove(tp)
+    if (removedLog != null) {
+      removedLog.removeLogMetrics()
+      Some(removedLog)
+    } else {
+      None
+    }
+  }
 }
 
 object LogManager {
@@ -894,32 +1180,30 @@ object LogManager {
             brokerTopicStats: BrokerTopicStats,
             logDirFailureChannel: LogDirFailureChannel): LogManager = {
     val defaultProps = KafkaServer.copyKafkaConfigToLog(config)
+
+    LogConfig.validateValues(defaultProps)
     val defaultLogConfig = LogConfig(defaultProps)
 
-    val (topicConfigs, failed) = zkClient.getLogConfigs(zkClient.getAllTopicsInCluster, defaultProps)
+    // read the log configurations from zookeeper
+    val (topicConfigs, failed) = zkClient.getLogConfigs(
+      zkClient.getAllTopicsInCluster(),
+      defaultProps
+    )
     if (!failed.isEmpty) throw failed.head._2
 
-    // read the log configurations from zookeeper
-    val cleanerConfig = CleanerConfig(numThreads = config.logCleanerThreads,
-      dedupeBufferSize = config.logCleanerDedupeBufferSize,
-      dedupeBufferLoadFactor = config.logCleanerDedupeBufferLoadFactor,
-      ioBufferSize = config.logCleanerIoBufferSize,
-      maxMessageSize = config.messageMaxBytes,
-      maxIoBytesPerSecond = config.logCleanerIoMaxBytesPerSecond,
-      backOffMs = config.logCleanerBackoffMs,
-      enableCleaner = config.logCleanerEnable)
+    val cleanerConfig = LogCleaner.cleanerConfig(config)
 
     new LogManager(logDirs = config.logDirs.map(new File(_).getAbsoluteFile),
       initialOfflineDirs = initialOfflineDirs.map(new File(_).getAbsoluteFile),
       topicConfigs = topicConfigs,
-      defaultConfig = defaultLogConfig,
+      initialDefaultConfig = defaultLogConfig,
       cleanerConfig = cleanerConfig,
-      ioThreads = config.numRecoveryThreadsPerDataDir,
+      recoveryThreadsPerDataDir = config.numRecoveryThreadsPerDataDir,
       flushCheckMs = config.logFlushSchedulerIntervalMs,
       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 6db2a508ac486..b43833d4a701d 100755
--- a/core/src/main/scala/kafka/log/LogSegment.scala
+++ b/core/src/main/scala/kafka/log/LogSegment.scala
@@ -17,68 +17,105 @@
 package kafka.log
 
 import java.io.{File, IOException}
-import java.nio.file.Files
+import java.nio.file.{Files, NoSuchFileException}
 import java.nio.file.attribute.FileTime
 import java.util.concurrent.TimeUnit
+
+import kafka.common.LogSegmentOffsetOverflowException
 import kafka.metrics.{KafkaMetricsGroup, KafkaTimer}
-import kafka.server.epoch.LeaderEpochCache
+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
+import org.apache.kafka.common.record.FileRecords.{LogOffsetPosition, TimestampAndOffset}
 import org.apache.kafka.common.record._
 import org.apache.kafka.common.utils.Time
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 import scala.math._
 
 /**
- * A segment of the log. Each segment has two components: a log and an index. The log is a FileMessageSet containing
+ * A segment of the log. Each segment has two components: a log and an index. The log is a FileRecords containing
  * the actual messages. The index is an OffsetIndex that maps from logical offsets to physical file positions. Each
  * segment has a base offset which is an offset <= the least offset of any message in this segment and > any offset in
  * any previous segment.
  *
  * 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 message set containing log entries
- * @param index The offset index
- * @param timeIndex The timestamp index
+ * @param log The file records containing log entries
+ * @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
+ * @param rollJitterMs The maximum random jitter subtracted from the scheduled segment roll time
  * @param time The time instance
  */
 @nonthreadsafe
-class LogSegment(val log: FileRecords,
-                 val index: OffsetIndex,
-                 val timeIndex: TimeIndex,
-                 val txnIndex: TransactionIndex,
-                 val baseOffset: Long,
-                 val indexIntervalBytes: Int,
-                 val rollJitterMs: Long,
-                 time: Time) extends Logging {
+class LogSegment private[log] (val log: FileRecords,
+                               val lazyOffsetIndex: LazyIndex[OffsetIndex],
+                               val lazyTimeIndex: LazyIndex[TimeIndex],
+                               val txnIndex: TransactionIndex,
+                               val baseOffset: Long,
+                               val indexIntervalBytes: Int,
+                               val rollJitterMs: Long,
+                               val time: Time) extends Logging {
+
+  def offsetIndex: OffsetIndex = lazyOffsetIndex.get
+
+  def timeIndex: TimeIndex = lazyTimeIndex.get
+
+  def shouldRoll(rollParams: RollParams): Boolean = {
+    val reachedRollMs = timeWaitedForRoll(rollParams.now, rollParams.maxTimestampInMessages) > rollParams.maxSegmentMs - rollJitterMs
+    size > rollParams.maxSegmentBytes - rollParams.messagesSize ||
+      (size > 0 && reachedRollMs) ||
+      offsetIndex.isFull || timeIndex.isFull || !canConvertToRelativeOffset(rollParams.maxOffsetInMessages)
+  }
+
+  def resizeIndexes(size: Int): Unit = {
+    offsetIndex.resize(size)
+    timeIndex.resize(size)
+  }
+
+  def sanityCheck(timeIndexFileNewlyCreated: Boolean): Unit = {
+    if (lazyOffsetIndex.file.exists) {
+      // Resize the time index file to 0 if it is newly created.
+      if (timeIndexFileNewlyCreated)
+        timeIndex.resize(0)
+      // Sanity checks for time index and offset index are skipped because
+      // we will recover the segments above the recovery point in recoverLog()
+      // in any case so sanity checking them here is redundant.
+      txnIndex.sanityCheck()
+    }
+    else throw new NoSuchFileException(s"Offset index file ${lazyOffsetIndex.file.getAbsolutePath} does not exist")
+  }
 
   private var created = time.milliseconds
 
   /* 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: Long = timeIndex.lastEntry.timestamp
-  @volatile private var offsetOfMaxTimestamp: Long = timeIndex.lastEntry.offset
-
-  def this(dir: File, startOffset: Long, indexIntervalBytes: Int, maxIndexSize: Int, rollJitterMs: Long, time: Time,
-           fileAlreadyExists: Boolean = false, initFileSize: Int = 0, preallocate: Boolean = false) =
-    this(FileRecords.open(Log.logFile(dir, startOffset), fileAlreadyExists, initFileSize, preallocate),
-         new OffsetIndex(Log.offsetIndexFile(dir, startOffset), baseOffset = startOffset, maxIndexSize = maxIndexSize),
-         new TimeIndex(Log.timeIndexFile(dir, startOffset), baseOffset = startOffset, maxIndexSize = maxIndexSize),
-         new TransactionIndex(startOffset, Log.transactionIndexFile(dir, startOffset)),
-         startOffset,
-         indexIntervalBytes,
-         rollJitterMs,
-         time)
+  @volatile private var _maxTimestampSoFar: Option[Long] = None
+  def maxTimestampSoFar_=(timestamp: Long): Unit = _maxTimestampSoFar = Some(timestamp)
+  def maxTimestampSoFar: Long = {
+    if (_maxTimestampSoFar.isEmpty)
+      _maxTimestampSoFar = Some(timeIndex.lastEntry.timestamp)
+    _maxTimestampSoFar.get
+  }
+
+  @volatile private var _offsetOfMaxTimestampSoFar: Option[Long] = None
+  def offsetOfMaxTimestampSoFar_=(offset: Long): Unit = _offsetOfMaxTimestampSoFar = Some(offset)
+  def offsetOfMaxTimestampSoFar: Long = {
+    if (_offsetOfMaxTimestampSoFar.isEmpty)
+      _offsetOfMaxTimestampSoFar = Some(timeIndex.lastEntry.offset)
+    _offsetOfMaxTimestampSoFar.get
+  }
 
   /* Return the size in bytes of this log segment */
   def size: Int = log.sizeInBytes()
@@ -87,7 +124,7 @@ class LogSegment(val log: FileRecords,
    * checks that the argument offset can be represented as an integer offset relative to the baseOffset.
    */
   def canConvertToRelativeOffset(offset: Long): Boolean = {
-    (offset - baseOffset) <= Integer.MAX_VALUE
+    offsetIndex.canAppendOffset(offset)
   }
 
   /**
@@ -96,46 +133,109 @@ class LogSegment(val log: FileRecords,
    *
    * It is assumed this method is being called from within a lock.
    *
-   * @param firstOffset The first offset in the message set.
    * @param largestOffset The last offset in the message set
    * @param largestTimestamp The largest timestamp in the message set.
    * @param shallowOffsetOfMaxTimestamp The offset of the message that has the largest timestamp in the messages to append.
    * @param records The log entries to append.
    * @return the physical position in the file of the appended records
+   * @throws LogSegmentOffsetOverflowException if the largest offset causes index offset overflow
    */
   @nonthreadsafe
-  def append(firstOffset: Long,
-             largestOffset: Long,
+  def append(largestOffset: Long,
              largestTimestamp: Long,
              shallowOffsetOfMaxTimestamp: Long,
              records: MemoryRecords): Unit = {
     if (records.sizeInBytes > 0) {
-      trace("Inserting %d bytes at offset %d at position %d with largest timestamp %d at shallow offset %d"
-          .format(records.sizeInBytes, firstOffset, log.sizeInBytes(), largestTimestamp, shallowOffsetOfMaxTimestamp))
+      trace(s"Inserting ${records.sizeInBytes} bytes at end offset $largestOffset at position ${log.sizeInBytes} " +
+            s"with largest timestamp $largestTimestamp at shallow offset $shallowOffsetOfMaxTimestamp")
       val physicalPosition = log.sizeInBytes()
       if (physicalPosition == 0)
         rollingBasedTimestamp = Some(largestTimestamp)
+
+      ensureOffsetInRange(largestOffset)
+
       // append the messages
-      require(canConvertToRelativeOffset(largestOffset), "largest offset in message set can not be safely converted to relative offset.")
       val appendedBytes = log.append(records)
-      trace(s"Appended $appendedBytes to ${log.file()} at offset $firstOffset")
+      trace(s"Appended $appendedBytes to ${log.file} at end offset $largestOffset")
       // Update the in memory max timestamp and corresponding offset.
       if (largestTimestamp > maxTimestampSoFar) {
         maxTimestampSoFar = largestTimestamp
-        offsetOfMaxTimestamp = shallowOffsetOfMaxTimestamp
+        offsetOfMaxTimestampSoFar = shallowOffsetOfMaxTimestamp
       }
       // append an entry to the index (if needed)
-      if(bytesSinceLastIndexEntry > indexIntervalBytes) {
-        index.append(firstOffset, physicalPosition)
-        timeIndex.maybeAppend(maxTimestampSoFar, offsetOfMaxTimestamp)
+      if (bytesSinceLastIndexEntry > indexIntervalBytes) {
+        offsetIndex.append(largestOffset, physicalPosition)
+        timeIndex.maybeAppend(maxTimestampSoFar, offsetOfMaxTimestampSoFar)
         bytesSinceLastIndexEntry = 0
       }
       bytesSinceLastIndexEntry += records.sizeInBytes
     }
   }
 
+  private def ensureOffsetInRange(offset: Long): Unit = {
+    if (!canConvertToRelativeOffset(offset))
+      throw new LogSegmentOffsetOverflowException(this, offset)
+  }
+
+  private def appendChunkFromFile(records: FileRecords, position: Int, bufferSupplier: BufferSupplier): Int = {
+    var bytesToAppend = 0
+    var maxTimestamp = Long.MinValue
+    var offsetOfMaxTimestamp = Long.MinValue
+    var maxOffset = Long.MinValue
+    var readBuffer = bufferSupplier.get(1024 * 1024)
+
+    def canAppend(batch: RecordBatch) =
+      canConvertToRelativeOffset(batch.lastOffset) &&
+        (bytesToAppend == 0 || bytesToAppend + batch.sizeInBytes < readBuffer.capacity)
+
+    // find all batches that are valid to be appended to the current log segment and
+    // determine the maximum offset and timestamp
+    val nextBatches = records.batchesFrom(position).asScala.iterator
+    for (batch <- nextBatches.takeWhile(canAppend)) {
+      if (batch.maxTimestamp > maxTimestamp) {
+        maxTimestamp = batch.maxTimestamp
+        offsetOfMaxTimestamp = batch.lastOffset
+      }
+      maxOffset = batch.lastOffset
+      bytesToAppend += batch.sizeInBytes
+    }
+
+    if (bytesToAppend > 0) {
+      // Grow buffer if needed to ensure we copy at least one batch
+      if (readBuffer.capacity < bytesToAppend)
+        readBuffer = bufferSupplier.get(bytesToAppend)
+
+      readBuffer.limit(bytesToAppend)
+      records.readInto(readBuffer, position)
+
+      append(maxOffset, maxTimestamp, offsetOfMaxTimestamp, MemoryRecords.readableRecords(readBuffer))
+    }
+
+    bufferSupplier.release(readBuffer)
+    bytesToAppend
+  }
+
+  /**
+   * Append records from a file beginning at the given position until either the end of the file
+   * is reached or an offset is found which is too large to convert to a relative offset for the indexes.
+   *
+   * @return the number of bytes appended to the log (may be less than the size of the input if an
+   *         offset is encountered which would overflow this segment)
+   */
+  def appendFromFile(records: FileRecords, start: Int): Int = {
+    var position = start
+    val bufferSupplier: BufferSupplier = new BufferSupplier.GrowableBufferSupplier
+    while (position < start + records.sizeInBytes) {
+      val bytesAppended = appendChunkFromFile(records, position, bufferSupplier)
+      if (bytesAppended == 0)
+        return position - start
+      position += bytesAppended
+    }
+    position - start
+  }
+
   @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))
@@ -145,12 +245,13 @@ class LogSegment(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)
@@ -170,7 +271,7 @@ class LogSegment(val log: FileRecords,
    */
   @threadsafe
   private[log] def translateOffset(offset: Long, startingFilePosition: Int = 0): LogOffsetPosition = {
-    val mapping = index.lookup(offset)
+    val mapping = offsetIndex.lookup(offset)
     log.searchForOffsetWithSize(offset, max(mapping.position, startingFilePosition))
   }
 
@@ -180,7 +281,6 @@ class LogSegment(val log: FileRecords,
    *
    * @param startOffset A lower bound on the first offset to include in the message set we read
    * @param maxSize The maximum number of bytes to include in the message set we read
-   * @param maxOffset An optional maximum offset for 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)
    *
@@ -188,12 +288,13 @@ class LogSegment(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("Invalid max size for log read (%d)".format(maxSize))
+      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
@@ -201,7 +302,7 @@ class LogSegment(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)
@@ -212,32 +313,14 @@ class LogSegment(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.read(startPosition, fetchSize),
+    FetchDataInfo(offsetMetadata, log.slice(startPosition, fetchSize),
       firstEntryIncomplete = adjustedMaxSize < startOffsetAndSize.size)
   }
 
    def fetchUpperBoundOffset(startOffsetPosition: OffsetPosition, fetchSize: Int): Option[Long] =
-     index.fetchUpperBoundOffset(startOffsetPosition, fetchSize).map(_.offset)
+     offsetIndex.fetchUpperBoundOffset(startOffsetPosition, fetchSize).map(_.offset)
 
   /**
    * Run recovery on the given segment. This will rebuild the index from the log file and lop off any invalid bytes
@@ -247,80 +330,91 @@ class LogSegment(val log: FileRecords,
    *                             the transaction index.
    * @param leaderEpochCache Optionally a cache for updating the leader epoch during recovery.
    * @return The number of bytes truncated from the log
+   * @throws LogSegmentOffsetOverflowException if the log segment contains an offset that causes the index offset to overflow
    */
   @nonthreadsafe
-  def recover(producerStateManager: ProducerStateManager, leaderEpochCache: Option[LeaderEpochCache] = None): Int = {
-    index.truncate()
-    index.resize(index.maxIndexSize)
-    timeIndex.truncate()
-    timeIndex.resize(timeIndex.maxIndexSize)
-    txnIndex.truncate()
+  def recover(producerStateManager: ProducerStateManager, leaderEpochCache: Option[LeaderEpochFileCache] = None): Int = {
+    offsetIndex.reset()
+    timeIndex.reset()
+    txnIndex.reset()
     var validBytes = 0
     var lastIndexEntry = 0
     maxTimestampSoFar = RecordBatch.NO_TIMESTAMP
     try {
       for (batch <- log.batches.asScala) {
         batch.ensureValid()
+        ensureOffsetInRange(batch.lastOffset)
 
         // The max timestamp is exposed at the batch level, so no need to iterate the records
         if (batch.maxTimestamp > maxTimestampSoFar) {
           maxTimestampSoFar = batch.maxTimestamp
-          offsetOfMaxTimestamp = batch.lastOffset
+          offsetOfMaxTimestampSoFar = batch.lastOffset
         }
 
         // Build offset index
-        if(validBytes - lastIndexEntry > indexIntervalBytes) {
-          val startOffset = batch.baseOffset
-          index.append(startOffset, validBytes)
-          timeIndex.maybeAppend(maxTimestampSoFar, offsetOfMaxTimestamp)
+        if (validBytes - lastIndexEntry > indexIntervalBytes) {
+          offsetIndex.append(batch.lastOffset, validBytes)
+          timeIndex.maybeAppend(maxTimestampSoFar, offsetOfMaxTimestampSoFar)
           lastIndexEntry = validBytes
         }
         validBytes += batch.sizeInBytes()
 
         if (batch.magic >= RecordBatch.MAGIC_VALUE_V2) {
           leaderEpochCache.foreach { cache =>
-            if (batch.partitionLeaderEpoch > cache.latestEpoch()) // this is to avoid unnecessary warning in cache.assign()
+            if (batch.partitionLeaderEpoch >= 0 && cache.latestEpoch.forall(batch.partitionLeaderEpoch > _))
               cache.assign(batch.partitionLeaderEpoch, batch.baseOffset)
           }
           updateProducerState(producerStateManager, batch)
         }
       }
     } 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)
       debug(s"Truncated $truncated invalid bytes at the end of segment ${log.file.getAbsoluteFile} during recovery")
 
     log.truncateTo(validBytes)
-    index.trimToValidSize()
+    offsetIndex.trimToValidSize()
     // A normally closed segment always appends the biggest timestamp ever seen into log segment, we do this as well.
-    timeIndex.maybeAppend(maxTimestampSoFar, offsetOfMaxTimestamp, skipFullCheck = true)
+    timeIndex.maybeAppend(maxTimestampSoFar, offsetOfMaxTimestampSoFar, skipFullCheck = true)
     timeIndex.trimToValidSize()
     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
-    offsetOfMaxTimestamp = lastTimeIndexEntry.offset
+    offsetOfMaxTimestampSoFar = lastTimeIndexEntry.offset
 
-    val offsetPosition = index.lookup(lastTimeIndexEntry.offset)
+    val offsetPosition = offsetIndex.lookup(lastTimeIndexEntry.offset)
     // Scan the rest of the messages to see if there is a larger timestamp after the last time index entry.
     val maxTimestampOffsetAfterLastEntry = log.largestTimestampAfter(offsetPosition.position)
     if (maxTimestampOffsetAfterLastEntry.timestamp > lastTimeIndexEntry.timestamp) {
       maxTimestampSoFar = maxTimestampOffsetAfterLastEntry.timestamp
-      offsetOfMaxTimestamp = maxTimestampOffsetAfterLastEntry.offset
+      offsetOfMaxTimestampSoFar = maxTimestampOffsetAfterLastEntry.offset
     }
   }
 
+  /**
+   * Check whether the last offset of the last batch in this segment overflows the indexes.
+   */
+  def hasOverflow: Boolean = {
+    val nextOffset = readNextOffset
+    nextOffset > baseOffset && !canConvertToRelativeOffset(nextOffset - 1)
+  }
+
   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 +
+    ", largestRecordTimestamp=" + largestRecordTimestamp +
+    ")"
 
   /**
    * Truncate off all index and log entries with offsets >= the given offset.
@@ -331,20 +425,23 @@ class LogSegment(val log: FileRecords,
    */
   @nonthreadsafe
   def truncateTo(offset: Long): Int = {
+    // Do offset translation before truncating the index to avoid needless scanning
+    // in case we truncate the full index
     val mapping = translateOffset(offset)
-    if (mapping == null)
-      return 0
-    index.truncateTo(offset)
+    offsetIndex.truncateTo(offset)
     timeIndex.truncateTo(offset)
     txnIndex.truncateTo(offset)
-    // after truncation, reset and allocate more space for the (new currently  active) index
-    index.resize(index.maxIndexSize)
+
+    // After truncation, reset and allocate more space for the (new currently active) index
+    offsetIndex.resize(offsetIndex.maxIndexSize)
     timeIndex.resize(timeIndex.maxIndexSize)
-    val bytesTruncated = log.truncateTo(mapping.position)
-    if(log.sizeInBytes == 0) {
+
+    val bytesTruncated = if (mapping == null) 0 else log.truncateTo(mapping.position)
+    if (log.sizeInBytes == 0) {
       created = time.milliseconds
       rollingBasedTimestamp = None
     }
+
     bytesSinceLastIndexEntry = 0
     if (maxTimestampSoFar >= 0)
       loadLargestTimestamp()
@@ -356,12 +453,12 @@ class LogSegment(val log: FileRecords,
    * Note that this is expensive.
    */
   @threadsafe
-  def nextOffset(): Long = {
-    val ms = read(index.lastOffset, None, log.sizeInBytes)
-    if (ms == null)
+  def readNextOffset: Long = {
+    val fetchData = read(offsetIndex.lastOffset, log.sizeInBytes)
+    if (fetchData == null)
       baseOffset
     else
-      ms.records.batches.asScala.lastOption
+      fetchData.records.batches.asScala.lastOption
         .map(_.nextOffset)
         .getOrElse(baseOffset)
   }
@@ -370,10 +467,10 @@ class LogSegment(val log: FileRecords,
    * Flush this log segment to disk
    */
   @threadsafe
-  def flush() {
+  def flush(): Unit = {
     LogFlushStats.logFlushTimer.time {
       log.flush()
-      index.flush()
+      offsetIndex.flush()
       timeIndex.flush()
       txnIndex.flush()
     }
@@ -383,30 +480,46 @@ class LogSegment(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))
-    index.file = new File(dir, index.file.getName)
-    timeIndex.file = new File(dir, timeIndex.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)))
-    index.renameTo(new File(CoreUtils.replaceSuffix(index.file.getPath, oldSuffix, newSuffix)))
-    timeIndex.renameTo(new File(CoreUtils.replaceSuffix(timeIndex.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)))
   }
 
   /**
-   * Append the largest time index entry to the time index when this log segment become inactive segment.
-   * This entry will be used to decide when to delete the segment.
+   * Append the largest time index entry to the time index and trim the log and indexes.
+   *
+   * The time index entry appended will be used to decide when to delete the segment.
    */
-  def onBecomeInactiveSegment() {
-    timeIndex.maybeAppend(maxTimestampSoFar, offsetOfMaxTimestamp, skipFullCheck = true)
+  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)
+    }
   }
 
   /**
@@ -420,17 +533,24 @@ class LogSegment(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.
    *
@@ -452,24 +572,24 @@ class LogSegment(val log: FileRecords,
    * @param startingOffset The starting offset to search.
    * @return the timestamp and offset of the first message that meets the requirements. None will be returned if there is no such message.
    */
-  def findOffsetByTimestamp(timestamp: Long, startingOffset: Long = baseOffset): Option[TimestampOffset] = {
+  def findOffsetByTimestamp(timestamp: Long, startingOffset: Long = baseOffset): Option[TimestampAndOffset] = {
     // Get the index entry with a timestamp less than or equal to the target timestamp
     val timestampOffset = timeIndex.lookup(timestamp)
-    val position = index.lookup(math.max(timestampOffset.offset, startingOffset)).position
+    val position = offsetIndex.lookup(math.max(timestampOffset.offset, startingOffset)).position
 
     // Search the timestamp
-    Option(log.searchForTimestamp(timestamp, position, startingOffset)).map { timestampAndOffset =>
-      TimestampOffset(timestampAndOffset.timestamp, timestampAndOffset.offset)
-    }
+    Option(log.searchForTimestamp(timestamp, position, startingOffset))
   }
 
   /**
    * Close this log segment
    */
-  def close() {
-    CoreUtils.swallow(timeIndex.maybeAppend(maxTimestampSoFar, offsetOfMaxTimestamp, skipFullCheck = true), this)
-    CoreUtils.swallow(index.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)
   }
@@ -477,9 +597,9 @@ class LogSegment(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(index.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)
   }
@@ -487,19 +607,25 @@ class LogSegment(val log: FileRecords,
   /**
    * Delete this log segment from the filesystem.
    */
-  def delete() {
-    val deletedLog = log.delete()
-    val deletedIndex = index.delete()
-    val deletedTimeIndex = timeIndex.delete()
-    val deletedTxnIndex = txnIndex.delete()
-    if (!deletedLog && log.file.exists)
-      throw new IOException("Delete of log " + log.file.getName + " failed.")
-    if (!deletedIndex && index.file.exists)
-      throw new IOException("Delete of index " + index.file.getName + " failed.")
-    if (!deletedTimeIndex && timeIndex.file.exists)
-      throw new IOException("Delete of time index " + timeIndex.file.getName + " failed.")
-    if (!deletedTxnIndex && txnIndex.file.exists)
-      throw new IOException("Delete of transaction index " + txnIndex.file.getName + " failed.")
+  def deleteIfExists(): Unit = {
+    def delete(delete: () => Boolean, fileType: String, file: File, logIfMissing: Boolean): Unit = {
+      try {
+        if (delete())
+          info(s"Deleted $fileType ${file.getAbsolutePath}.")
+        else if (logIfMissing)
+          info(s"Failed to delete $fileType ${file.getAbsolutePath} because it does not exist.")
+      }
+      catch {
+        case e: IOException => throw new IOException(s"Delete of $fileType ${file.getAbsolutePath} failed.", e)
+      }
+    }
+
+    CoreUtils.tryAll(Seq(
+      () => delete(log.deleteIfExists _, "log", log.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)
+    ))
   }
 
   /**
@@ -507,6 +633,11 @@ class LogSegment(val log: FileRecords,
    */
   def lastModified = log.file.lastModified
 
+  /**
+   * The largest timestamp this segment contains, if maxTimestampSoFar >= 0, otherwise None.
+   */
+  def largestRecordTimestamp: Option[Long] = if (maxTimestampSoFar >= 0) Some(maxTimestampSoFar) else None
+
   /**
    * The largest timestamp this segment contains.
    */
@@ -518,8 +649,33 @@ class LogSegment(val log: FileRecords,
   def lastModified_=(ms: Long) = {
     val fileTime = FileTime.fromMillis(ms)
     Files.setLastModifiedTime(log.file.toPath, fileTime)
-    Files.setLastModifiedTime(index.file.toPath, fileTime)
-    Files.setLastModifiedTime(timeIndex.file.toPath, fileTime)
+    Files.setLastModifiedTime(lazyOffsetIndex.file.toPath, fileTime)
+    Files.setLastModifiedTime(lazyTimeIndex.file.toPath, fileTime)
+  }
+
+}
+
+object LogSegment {
+
+  def open(dir: File, baseOffset: Long, config: LogConfig, time: Time, fileAlreadyExists: Boolean = false,
+           initFileSize: Int = 0, preallocate: Boolean = false, fileSuffix: String = ""): LogSegment = {
+    val maxIndexSize = config.maxIndexSize
+    new LogSegment(
+      FileRecords.open(Log.logFile(dir, baseOffset, fileSuffix), fileAlreadyExists, initFileSize, preallocate),
+      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,
+      rollJitterMs = config.randomSegmentJitter,
+      time)
+  }
+
+  def deleteIfExists(dir: File, baseOffset: Long, fileSuffix: String = ""): Unit = {
+    Log.deleteFileIfExists(Log.offsetIndexFile(dir, baseOffset, fileSuffix))
+    Log.deleteFileIfExists(Log.timeIndexFile(dir, baseOffset, fileSuffix))
+    Log.deleteFileIfExists(Log.transactionIndexFile(dir, baseOffset, fileSuffix))
+    Log.deleteFileIfExists(Log.logFile(dir, baseOffset, fileSuffix))
   }
 }
 
diff --git a/core/src/main/scala/kafka/log/LogValidator.scala b/core/src/main/scala/kafka/log/LogValidator.scala
index 15750e9cd065c..58706fa481d10 100644
--- a/core/src/main/scala/kafka/log/LogValidator.scala
+++ b/core/src/main/scala/kafka/log/LogValidator.scala
@@ -18,17 +18,49 @@ package kafka.log
 
 import java.nio.ByteBuffer
 
-import kafka.common.LongRef
-import kafka.message.{CompressionCodec, NoCompressionCodec}
+import kafka.api.{ApiVersion, KAFKA_2_1_IV0}
+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, UnsupportedForMessageFormatException}
-import org.apache.kafka.common.record._
+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.protocol.Errors
+import org.apache.kafka.common.requests.ProduceResponse.RecordError
 import org.apache.kafka.common.utils.Time
 
-import scala.collection.mutable
-import scala.collection.JavaConverters._
+import scala.collection.{Seq, mutable}
+import scala.jdk.CollectionConverters._
+import scala.collection.mutable.ArrayBuffer
 
-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:
@@ -45,44 +77,99 @@ 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): 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): 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)
+      validateMessagesAndAssignOffsetsCompressed(records, topicPartition, offsetCounter, time, now, sourceCodec, targetCodec, compactedTopic,
+        magic, timestampType, timestampDiffMaxMs, partitionLeaderEpoch, origin, interBrokerProtocolVersion, brokerTopicStats)
     }
   }
 
-  private def validateBatch(batch: RecordBatch, isFromClient: Boolean, toMagic: Byte): Unit = {
-    if (isFromClient) {
-      if (batch.hasProducerId && batch.baseSequence < 0)
-        throw new InvalidRecordException(s"Invalid sequence number ${batch.baseSequence} in record batch " +
-          s"with producerId ${batch.producerId}")
+  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 (batch.isControlBatch)
-        throw new InvalidRecordException("Clients are not allowed to write control records")
+    // 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")
+      }
+    }
 
-      if (Option(batch.countOrNull).contains(0))
-        throw new InvalidRecordException("Record batches must contain at least one record")
+    batch
+  }
+
+  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) {
+          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) {
+          brokerTopicStats.allTopicsStats.invalidOffsetOrSequenceRecordsPerSec.mark()
+          throw new InvalidRecordException(s"Invalid reported count for record batch: $count in topic partition $topicPartition.")
+        }
+
+        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 in topic partition $topicPartition.")
+        }
+      }
+
+      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.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)
@@ -92,23 +179,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): Option[ApiRecordError] = {
+    if (!record.hasMagic(batch.magic)) {
+      brokerTopicStats.allTopicsStats.invalidMagicNumberRecordsPerSec.mark()
+      return Some(ApiRecordError(Errors.INVALID_RECORD, new RecordError(batchIndex,
+        s"Record $record's magic does not match outer magic ${batch.magic} in topic partition $topicPartition.")))
+    }
 
     // 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).orElse {
+      validateTimestamp(batch, record, batchIndex, now, timestampType, timestampDiffMaxMs)
+    }
   }
 
   private def convertAndAssignOffsetsNonCompressed(records: MemoryRecords,
+                                                   topicPartition: TopicPartition,
                                                    offsetCounter: LongRef,
                                                    compactedTopic: Boolean,
                                                    time: Time,
@@ -117,7 +217,9 @@ 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)
 
@@ -130,56 +232,76 @@ private[kafka] object LogValidator extends Logging {
     val builder = MemoryRecords.builder(newBuffer, toMagicValue, CompressionType.NONE, timestampType,
       offsetCounter.value, now, producerId, producerEpoch, sequence, isTransactional, partitionLeaderEpoch)
 
-    for (batch <- records.batches.asScala) {
-      validateBatch(batch, isFromClient, toMagicValue)
+    val firstBatch = getFirstBatchAndMaybeValidateNoMoreBatches(records, NoCompressionCodec)
 
-      for (record <- batch.asScala) {
-        validateRecord(batch, record, now, timestampType, timestampDiffMaxMs, compactedTopic)
-        builder.appendWithOffset(offsetCounter.getAndIncrement(), record)
+    records.batches.forEach { batch =>
+      validateBatch(topicPartition, firstBatch, batch, origin, toMagicValue, brokerTopicStats)
+
+      val recordErrors = new ArrayBuffer[ApiRecordError](0)
+      for ((record, batchIndex) <- batch.asScala.view.zipWithIndex) {
+        validateRecord(batch, topicPartition, record, batchIndex, now, timestampType,
+          timestampDiffMaxMs, compactedTopic, brokerTopicStats).foreach(recordError => recordErrors += recordError)
+        // we fail the batch if any record fails, so we stop appending if any record fails
+        if (recordErrors.isEmpty)
+          builder.appendWithOffset(offsetCounter.getAndIncrement(), record)
       }
+
+      processRecordErrors(recordErrors)
     }
 
     val convertedRecords = builder.build()
 
     val info = builder.info
-    val recordsProcessingStats = new RecordsProcessingStats(builder.uncompressedBytesWritten,
-      builder.numRecords, time.nanoseconds - now)
+    val recordConversionStats = new RecordConversionStats(builder.uncompressedBytesWritten,
+      builder.numRecords, time.nanoseconds - startNanos)
     ValidationAndOffsetAssignResult(
       validatedRecords = convertedRecords,
       maxTimestamp = info.maxTimestamp,
       shallowOffsetOfMaxTimestamp = info.shallowOffsetOfMaxTimestamp,
       messageSizeMaybeChanged = true,
-      recordsProcessingStats = recordsProcessingStats)
+      recordConversionStats = recordConversionStats)
   }
 
-  private def assignOffsetsNonCompressed(records: MemoryRecords,
+  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
 
-    for (batch <- records.batches.asScala) {
-      validateBatch(batch, isFromClient, magic)
+    val firstBatch = getFirstBatchAndMaybeValidateNoMoreBatches(records, NoCompressionCodec)
+
+    records.batches.forEach { batch =>
+      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)
+      val recordErrors = new ArrayBuffer[ApiRecordError](0)
+      // this is a hot path and we want to avoid any unnecessary allocations.
+      var batchIndex = 0
+      batch.forEach { record =>
+        validateRecord(batch, topicPartition, record, batchIndex, now, timestampType,
+          timestampDiffMaxMs, compactedTopic, brokerTopicStats).foreach(recordError => recordErrors += recordError)
+
         val offset = offsetCounter.getAndIncrement()
         if (batch.magic > RecordBatch.MAGIC_VALUE_V0 && record.timestamp > maxBatchTimestamp) {
           maxBatchTimestamp = record.timestamp
           offsetOfMaxBatchTimestamp = offset
         }
+        batchIndex += 1
       }
 
+      processRecordErrors(recordErrors)
+
       if (batch.magic > RecordBatch.MAGIC_VALUE_V0 && maxBatchTimestamp > maxTimestamp) {
         maxTimestamp = maxBatchTimestamp
         offsetOfMaxTimestamp = offsetOfMaxBatchTimestamp
@@ -211,17 +333,17 @@ private[kafka] object LogValidator extends Logging {
       maxTimestamp = maxTimestamp,
       shallowOffsetOfMaxTimestamp = offsetOfMaxTimestamp,
       messageSizeMaybeChanged = false,
-      recordsProcessingStats = RecordsProcessingStats.EMPTY)
+      recordConversionStats = RecordConversionStats.EMPTY)
   }
 
   /**
    * 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,
@@ -232,83 +354,129 @@ private[kafka] object LogValidator extends Logging {
                                                  timestampType: TimestampType,
                                                  timestampDiffMaxMs: Long,
                                                  partitionLeaderEpoch: Int,
-                                                 isFromClient: Boolean): ValidationAndOffsetAssignResult = {
-
-      // No in place assignment situation 1 and 2
-      var inPlaceAssignment = sourceCodec == targetCodec && toMagic > RecordBatch.MAGIC_VALUE_V0
-
-      var maxTimestamp = RecordBatch.NO_TIMESTAMP
-      val expectedInnerOffset = new LongRef(0)
-      val validatedRecords = new mutable.ArrayBuffer[Record]
-
-      var uncompressedSizeInBytes = 0
-
-      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
-
-        for (record <- batch.asScala) {
-          validateRecord(batch, record, now, timestampType, timestampDiffMaxMs, compactedTopic)
-          if (sourceCodec != NoCompressionCodec && record.isCompressed)
-            throw new InvalidRecordException("Compressed outer record should not have an inner record with a " +
-              s"compression attribute set: $record")
-
-          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())
-              inPlaceAssignment = false
-            if (record.timestamp > maxTimestamp)
-              maxTimestamp = record.timestamp
-          }
+                                                 origin: AppendOrigin,
+                                                 interBrokerProtocolVersion: ApiVersion,
+                                                 brokerTopicStats: BrokerTopicStats): ValidationAndOffsetAssignResult = {
+
+    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")
+
+    def validateRecordCompression(batchIndex: Int, record: Record): Option[ApiRecordError] = {
+      if (sourceCodec != NoCompressionCodec && record.isCompressed)
+        Some(ApiRecordError(Errors.INVALID_RECORD, new RecordError(batchIndex,
+          s"Compressed outer record should not have an inner record with a compression attribute set: $record")))
+      else None
+    }
+
+    // No in place assignment situation 1
+    var inPlaceAssignment = sourceCodec == targetCodec
 
-          // No in place assignment situation 4
-          if (!record.hasMagic(toMagic))
-            inPlaceAssignment = false
+    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)
+
+    // 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
+
+    records.batches.forEach { batch =>
+      validateBatch(topicPartition, firstBatch, batch, origin, toMagic, brokerTopicStats)
+      uncompressedSizeInBytes += AbstractRecords.recordBatchHeaderSizeInBytes(toMagic, batch.compressionType())
+
+      // 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)
+
+      try {
+        val recordErrors = new ArrayBuffer[ApiRecordError](0)
+        // this is a hot path and we want to avoid any unnecessary allocations.
+        var batchIndex = 0
+        recordsIterator.forEachRemaining { record =>
+          val expectedOffset = expectedInnerOffset.getAndIncrement()
+          val recordError = validateRecordCompression(batchIndex, record).orElse {
+            validateRecord(batch, topicPartition, record, batchIndex, now,
+              timestampType, timestampDiffMaxMs, compactedTopic, brokerTopicStats).orElse {
+              if (batch.magic > RecordBatch.MAGIC_VALUE_V0 && toMagic > RecordBatch.MAGIC_VALUE_V0) {
+                if (record.timestamp > maxTimestamp)
+                  maxTimestamp = record.timestamp
+
+                // 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.
+                if (record.offset != expectedOffset)
+                  inPlaceAssignment = false
+              }
+              None
+            }
+          }
 
-          validatedRecords += record
+          recordError match {
+            case Some(e) => recordErrors += e
+            case None =>
+              uncompressedSizeInBytes += record.sizeInBytes()
+              validatedRecords += record
+          }
+         batchIndex += 1
         }
+        processRecordErrors(recordErrors)
+      } 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 = offsetCounter.addAndGet(validatedRecords.size) - 1
 
-        batch.setLastOffset(lastOffset)
+      batch.setLastOffset(lastOffset)
 
-        if (timestampType == TimestampType.LOG_APPEND_TIME)
-          maxTimestamp = now
+      if (timestampType == TimestampType.LOG_APPEND_TIME)
+        maxTimestamp = now
 
-        if (toMagic >= RecordBatch.MAGIC_VALUE_V1)
-          batch.setMaxTimestamp(timestampType, 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_V2)
+        batch.setPartitionLeaderEpoch(partitionLeaderEpoch)
 
-        val recordsProcessingStats = new RecordsProcessingStats(uncompressedSizeInBytes, 0, -1)
-        ValidationAndOffsetAssignResult(validatedRecords = records,
-          maxTimestamp = maxTimestamp,
-          shallowOffsetOfMaxTimestamp = lastOffset,
-          messageSizeMaybeChanged = false,
-          recordsProcessingStats = recordsProcessingStats)
-      }
+      val recordConversionStats = new RecordConversionStats(uncompressedSizeInBytes, 0, 0)
+      ValidationAndOffsetAssignResult(validatedRecords = records,
+        maxTimestamp = maxTimestamp,
+        shallowOffsetOfMaxTimestamp = lastOffset,
+        messageSizeMaybeChanged = false,
+        recordConversionStats = recordConversionStats)
+    }
   }
 
   private def buildRecordsAndAssignOffsets(magic: Byte,
@@ -323,8 +491,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)
@@ -345,7 +512,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 recordsProcessingStats = new RecordsProcessingStats(uncompresssedSizeInBytes + builder.uncompressedBytesWritten,
+    val recordConversionStats = new RecordConversionStats(uncompressedSizeInBytes + builder.uncompressedBytesWritten,
       conversionCount, time.nanoseconds - startNanos)
 
     ValidationAndOffsetAssignResult(
@@ -353,37 +520,59 @@ private[kafka] object LogValidator extends Logging {
       maxTimestamp = info.maxTimestamp,
       shallowOffsetOfMaxTimestamp = info.shallowOffsetOfMaxTimestamp,
       messageSizeMaybeChanged = true,
-      recordsProcessingStats = recordsProcessingStats)
+      recordConversionStats = recordConversionStats)
   }
 
-  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): Option[ApiRecordError] = {
+    if (compactedTopic && !record.hasKey) {
+      brokerTopicStats.allTopicsStats.noKeyCompactedTopicRecordsPerSec.mark()
+      Some(ApiRecordError(Errors.INVALID_RECORD, new RecordError(batchIndex,
+        s"Compacted topic cannot accept message without key in topic partition $topicPartition.")))
+    } else None
   }
 
-  /**
-   * This method validates the timestamps of a message.
-   * If the message is using create time, this method checks if it is within acceptable range.
-   */
   private def validateTimestamp(batch: RecordBatch,
                                 record: Record,
+                                batchIndex: Int,
                                 now: Long,
                                 timestampType: TimestampType,
-                                timestampDiffMaxMs: Long) {
+                                timestampDiffMaxMs: Long): Option[ApiRecordError] = {
     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}]")
-    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.")
+      Some(ApiRecordError(Errors.INVALID_TIMESTAMP, new RecordError(batchIndex,
+        s"Timestamp ${record.timestamp} of message with offset ${record.offset} is " +
+          s"out of range. The timestamp should be within [${now - timestampDiffMaxMs}, " +
+          s"${now + timestampDiffMaxMs}]")))
+    else if (batch.timestampType == TimestampType.LOG_APPEND_TIME)
+      Some(ApiRecordError(Errors.INVALID_TIMESTAMP, new RecordError(batchIndex,
+        s"Invalid timestamp type in message $record. Producer should not set timestamp " +
+          "type to LogAppendTime.")))
+    else None
+  }
+
+  private def processRecordErrors(recordErrors: Seq[ApiRecordError]): Unit = {
+    if (recordErrors.nonEmpty) {
+      val errors = recordErrors.map(_.recordError)
+      if (recordErrors.exists(_.apiError == Errors.INVALID_TIMESTAMP)) {
+        throw new RecordValidationException(new InvalidTimestampException(
+          "One or more records have been rejected due to invalid timestamp"), errors)
+      } else {
+        throw new RecordValidationException(new InvalidRecordException(
+          "One or more records have been rejected"), errors)
+      }
+    }
   }
 
   case class ValidationAndOffsetAssignResult(validatedRecords: MemoryRecords,
                                              maxTimestamp: Long,
                                              shallowOffsetOfMaxTimestamp: Long,
                                              messageSizeMaybeChanged: Boolean,
-                                             recordsProcessingStats: RecordsProcessingStats)
+                                             recordConversionStats: RecordConversionStats)
 
+  private case class ApiRecordError(apiError: Errors, recordError: RecordError)
 }
diff --git a/core/src/main/scala/kafka/log/OffsetIndex.scala b/core/src/main/scala/kafka/log/OffsetIndex.scala
index eb1584250fe9e..5719ee3814f0f 100755
--- a/core/src/main/scala/kafka/log/OffsetIndex.scala
+++ b/core/src/main/scala/kafka/log/OffsetIndex.scala
@@ -21,7 +21,8 @@ import java.io.File
 import java.nio.ByteBuffer
 
 import kafka.utils.CoreUtils.inLock
-import kafka.common.InvalidOffsetException
+import kafka.utils.Logging
+import org.apache.kafka.common.errors.InvalidOffsetException
 
 /**
  * An index that maps offsets to physical file locations for a particular log segment. This index may be sparse:
@@ -50,15 +51,16 @@ import kafka.common.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
 
   /* 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()))
+  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
@@ -67,7 +69,7 @@ class OffsetIndex(_file: File, baseOffset: Long, maxIndexSize: Int = -1, writabl
     inLock(lock) {
       _entries match {
         case 0 => OffsetPosition(baseOffset, 0)
-        case s => parseEntry(mmap, s - 1).asInstanceOf[OffsetPosition]
+        case s => parseEntry(mmap, s - 1)
       }
     }
   }
@@ -90,7 +92,7 @@ class OffsetIndex(_file: File, baseOffset: Long, maxIndexSize: Int = -1, writabl
       if(slot == -1)
         OffsetPosition(baseOffset, 0)
       else
-        parseEntry(idx, slot).asInstanceOf[OffsetPosition]
+        parseEntry(idx, slot)
     }
   }
 
@@ -106,7 +108,7 @@ class OffsetIndex(_file: File, baseOffset: Long, maxIndexSize: Int = -1, writabl
       if (slot == -1)
         None
       else
-        Some(parseEntry(idx, slot).asInstanceOf[OffsetPosition])
+        Some(parseEntry(idx, slot))
     }
   }
 
@@ -114,8 +116,8 @@ class OffsetIndex(_file: File, baseOffset: Long, maxIndexSize: Int = -1, writabl
 
   private def physical(buffer: ByteBuffer, n: Int): Int = buffer.getInt(n * entrySize + 4)
 
-  override def parseEntry(buffer: ByteBuffer, n: Int): IndexEntry = {
-      OffsetPosition(baseOffset + relativeOffset(buffer, n), physical(buffer, n))
+  override protected def parseEntry(buffer: ByteBuffer, n: Int): OffsetPosition = {
+    OffsetPosition(baseOffset + relativeOffset(buffer, n), physical(buffer, n))
   }
 
   /**
@@ -125,36 +127,37 @@ class OffsetIndex(_file: File, baseOffset: Long, maxIndexSize: Int = -1, writabl
    */
   def entry(n: Int): OffsetPosition = {
     maybeLock(lock) {
-      if(n >= _entries)
-        throw new IllegalArgumentException("Attempt to fetch the %dth entry from an index of size %d.".format(n, _entries))
-      val idx = mmap.duplicate
-      OffsetPosition(relativeOffset(idx, n), physical(idx, n))
+      if (n >= _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
    */
-  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))
-        mmap.putInt((offset - baseOffset).toInt)
+        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() + ".")
+        require(_entries * entrySize == mmap.position(), s"$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)
@@ -178,21 +181,27 @@ 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
+      debug(s"Truncated index ${file.getAbsolutePath} to $entries entries;" +
+        s" position is now ${mmap.position()} and last offset is now ${_lastOffset}")
     }
   }
 
-  override def sanityCheck() {
-    require(_entries == 0 || _lastOffset > baseOffset,
-            s"Corrupt index found, index file (${file.getAbsolutePath}) has non-zero size but the last offset " +
-                s"is ${_lastOffset} which is no larger than the base offset $baseOffset.")
-    require(length % entrySize == 0,
-            "Index file " + file.getAbsolutePath + " is corrupt, found " + length +
-            " bytes which is not positive or not a multiple of 8.")
+  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.")
+    if (length % entrySize != 0)
+      throw new CorruptIndexException(s"Index file ${file.getAbsolutePath} is corrupt, found $length bytes which is " +
+        s"neither positive nor a multiple of $entrySize.")
   }
 
 }
+
+object OffsetIndex extends Logging {
+  override val loggerName: String = classOf[OffsetIndex].getName
+}
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 da84b19abdfe7..b4fa26755f6b9 100644
--- a/core/src/main/scala/kafka/log/ProducerStateManager.scala
+++ b/core/src/main/scala/kafka/log/ProducerStateManager.scala
@@ -16,50 +16,39 @@
  */
 package kafka.log
 
-import java.io._
+import java.io.File
 import java.nio.ByteBuffer
-import java.nio.file.Files
+import java.nio.channels.FileChannel
+import java.nio.file.{Files, StandardOpenOption}
+import java.util.concurrent.ConcurrentSkipListMap
 
-import kafka.common.KafkaException
 import kafka.log.Log.offsetFromFile
 import kafka.server.LogOffsetMetadata
 import kafka.utils.{Logging, nonthreadsafe, threadsafe}
-import org.apache.kafka.common.TopicPartition
+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, EndTransactionMarker, RecordBatch}
+import org.apache.kafka.common.record.{ControlRecordType, DefaultRecordBatch, EndTransactionMarker, RecordBatch}
 import org.apache.kafka.common.utils.{ByteUtils, Crc32C}
 
+import scala.jdk.CollectionConverters._
 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) {
+private[log] case class TxnMetadata(
+  producerId: Long,
+  firstOffset: LogOffsetMetadata,
+  var lastOffset: Option[Long] = None
+) {
   def this(producerId: Long, firstOffset: Long) = this(producerId, LogOffsetMetadata(firstOffset))
 
   override def toString: String = {
@@ -70,14 +59,20 @@ private[log] case class TxnMetadata(producerId: Long, var firstOffset: LogOffset
   }
 }
 
-private[log] object ProducerIdEntry {
+private[log] object ProducerStateEntry {
   private[log] val NumBatchesToRetain = 5
-  def empty(producerId: Long) = new ProducerIdEntry(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 = lastSeq - offsetDelta
-  def firstOffset = lastOffset - offsetDelta
+  def firstSeq: Int =  DefaultRecordBatch.decrementSequence(lastSeq, offsetDelta)
+  def firstOffset: Long = lastOffset - offsetDelta
 
   override def toString: String = {
     "BatchMetadata(" +
@@ -90,30 +85,34 @@ private[log] case class BatchMetadata(lastSeq: Int, lastOffset: Long, offsetDelt
 }
 
 // the batchMetadata is ordered such that the batch with the lowest sequence is at the head of the queue while the
-// batch with the highest sequence is at the tail of the queue. We will retain at most ProducerIdEntry.NumBatchesToRetain
+// batch with the highest sequence is at the tail of the queue. We will retain at most ProducerStateEntry.NumBatchesToRetain
 // elements in the queue. When the queue is at capacity, we remove the first element to make space for the incoming batch.
-private[log] class ProducerIdEntry(val producerId: Long, val batchMetadata: mutable.Queue[BatchMetadata],
-                                   var producerEpoch: Short, var coordinatorEpoch: Int,
-                                   var currentTxnFirstOffset: Option[Long]) {
+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 (batchMetadata.isEmpty) RecordBatch.NO_SEQUENCE else batchMetadata.front.firstSeq
-  def firstOffset: Long = if (batchMetadata.isEmpty) -1L else batchMetadata.front.firstOffset
+  def firstSeq: Int = if (isEmpty) RecordBatch.NO_SEQUENCE else batchMetadata.front.firstSeq
 
-  def lastSeq: Int = if (batchMetadata.isEmpty) RecordBatch.NO_SEQUENCE else batchMetadata.last.lastSeq
-  def lastDataOffset: Long = if (batchMetadata.isEmpty) -1L else batchMetadata.last.lastOffset
-  def lastTimestamp = if (batchMetadata.isEmpty) RecordBatch.NO_TIMESTAMP else batchMetadata.last.timestamp
-  def lastOffsetDelta : Int = if (batchMetadata.isEmpty) 0 else batchMetadata.last.offsetDelta
+  def firstDataOffset: Long = if (isEmpty) -1L else batchMetadata.front.firstOffset
 
-  def addBatchMetadata(producerEpoch: Short, lastSeq: Int, lastOffset: Long, offsetDelta: Int, timestamp: Long) = {
-    maybeUpdateEpoch(producerEpoch)
+  def lastSeq: Int = if (isEmpty) RecordBatch.NO_SEQUENCE else batchMetadata.last.lastSeq
 
-    if (batchMetadata.size == ProducerIdEntry.NumBatchesToRetain)
-      batchMetadata.dequeue()
+  def lastDataOffset: Long = if (isEmpty) -1L else batchMetadata.last.lastOffset
+
+  def lastOffsetDelta : Int = if (isEmpty) 0 else batchMetadata.last.offsetDelta
 
-    batchMetadata.enqueue(BatchMetadata(lastSeq, lastOffset, offsetDelta, timestamp))
+  def isEmpty: Boolean = batchMetadata.isEmpty
+
+  def addBatch(producerEpoch: Short, lastSeq: Int, lastOffset: Long, offsetDelta: Int, timestamp: Long): Unit = {
+    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
@@ -123,29 +122,43 @@ private[log] class ProducerIdEntry(val producerId: Long, val batchMetadata: muta
     }
   }
 
-  def removeBatchesOlderThan(offset: Long) = batchMetadata.dropWhile(_.lastOffset < offset)
+  private def addBatchMetadata(batch: BatchMetadata): Unit = {
+    if (batchMetadata.size == ProducerStateEntry.NumBatchesToRetain)
+      batchMetadata.dequeue()
+    batchMetadata.enqueue(batch)
+  }
+
+  def update(nextEntry: ProducerStateEntry): Unit = {
+    maybeUpdateProducerEpoch(nextEntry.producerEpoch)
+    while (nextEntry.batchMetadata.nonEmpty)
+      addBatchMetadata(nextEntry.batchMetadata.dequeue())
+    this.coordinatorEpoch = nextEntry.coordinatorEpoch
+    this.currentTxnFirstOffset = nextEntry.currentTxnFirstOffset
+    this.lastTimestamp = nextEntry.lastTimestamp
+  }
 
-  def duplicateOf(batch: RecordBatch): Option[BatchMetadata] = {
-    if (batch.producerEpoch() != producerEpoch)
+  def findDuplicateBatch(batch: RecordBatch): Option[BatchMetadata] = {
+    if (batch.producerEpoch != producerEpoch)
        None
     else
-      batchWithSequenceRange(batch.baseSequence(), batch.lastSequence())
+      batchWithSequenceRange(batch.baseSequence, batch.lastSequence)
   }
 
   // Return the batch metadata of the cached batch having the exact sequence range, if any.
   def batchWithSequenceRange(firstSeq: Int, lastSeq: Int): Option[BatchMetadata] = {
-    val duplicate = batchMetadata.filter { case(metadata) =>
+    val duplicate = batchMetadata.filter { metadata =>
       firstSeq == metadata.firstSeq && lastSeq == metadata.lastSeq
     }
     duplicate.headOption
   }
 
   override def toString: String = {
-    "ProducerIdEntry(" +
+    "ProducerStateEntry(" +
       s"producerId=$producerId, " +
       s"producerEpoch=$producerEpoch, " +
       s"currentTxnFirstOffset=$currentTxnFirstOffset, " +
       s"coordinatorEpoch=$coordinatorEpoch, " +
+      s"lastTimestamp=$lastTimestamp, " +
       s"batchMetadata=$batchMetadata"
   }
 }
@@ -159,158 +172,175 @@ private[log] class ProducerIdEntry(val producerId: Long, val batchMetadata: muta
  * @param producerId The id of the producer appending to the log
  * @param currentEntry  The current entry associated with the producer id which contains metadata for a fixed number of
  *                      the most recent appends made by the producer. Validation of the first incoming append will
- *                      be made against the lastest append in the current entry. New appends will replace older appends
+ *                      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 replication 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,
-                                      currentEntry: ProducerIdEntry,
-                                      validationType: ValidationType) {
+private[log] class ProducerAppendInfo(val topicPartition: TopicPartition,
+                                      val producerId: Long,
+                                      val currentEntry: ProducerStateEntry,
+                                      val origin: AppendOrigin) extends Logging {
 
   private val transactions = ListBuffer.empty[TxnMetadata]
+  private val updatedEntry = ProducerStateEntry.empty(producerId)
 
-  private def maybeValidateAppend(producerEpoch: Short, firstSeq: Int, lastSeq: Int) = {
-    validationType match {
-      case ValidationType.None =>
+  updatedEntry.producerEpoch = currentEntry.producerEpoch
+  updatedEntry.coordinatorEpoch = currentEntry.coordinatorEpoch
+  updatedEntry.lastTimestamp = currentEntry.lastTimestamp
+  updatedEntry.currentTxnFirstOffset = currentEntry.currentTxnFirstOffset
 
-      case ValidationType.EpochOnly =>
-        checkEpoch(producerEpoch)
-
-      case ValidationType.Full =>
-        checkEpoch(producerEpoch)
-        checkSequence(producerEpoch, firstSeq, lastSeq)
+  private def maybeValidateDataBatch(producerEpoch: Short, firstSeq: Int, offset: Long): Unit = {
+    checkProducerEpoch(producerEpoch, offset)
+    if (origin == AppendOrigin.Client) {
+      checkSequence(producerEpoch, firstSeq, offset)
     }
   }
 
-  private def checkEpoch(producerEpoch: Short): Unit = {
-    if (isFenced(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), ${currentEntry.producerEpoch} (server epoch)")
+  private def checkProducerEpoch(producerEpoch: Short, offset: Long): Unit = {
+    if (producerEpoch < updatedEntry.producerEpoch) {
+      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 {
+        // Starting from 2.7, we replaced ProducerFenced error with InvalidProducerEpoch in the
+        // producer send response callback to differentiate from the former fatal exception,
+        // letting client abort the ongoing transaction and retry.
+        throw new InvalidProducerEpochException(message)
+      }
     }
   }
 
-  private def checkSequence(producerEpoch: Short, firstSeq: Int, lastSeq: Int): Unit = {
-    if (producerEpoch != currentEntry.producerEpoch) {
-      if (firstSeq != 0) {
-        if (currentEntry.producerEpoch != RecordBatch.NO_PRODUCER_EPOCH) {
-          throw new OutOfOrderSequenceException(s"Invalid sequence number for new epoch: $producerEpoch " +
-            s"(request epoch), $firstSeq (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.")
+  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 at offset $offset in " +
+            s"partition $topicPartition: $producerEpoch (request epoch), $appendFirstSeq (seq. number)")
         }
       }
-    } else if (currentEntry.lastSeq == RecordBatch.NO_SEQUENCE && firstSeq != 0) {
-      // the epoch was bumped by a control record, so we expect the sequence number to be reset
-      throw new OutOfOrderSequenceException(s"Out of order sequence number for producerId $producerId: found $firstSeq " +
-        s"(incoming seq. number), but expected 0")
-    } else if (isDuplicate(firstSeq, lastSeq)) {
-      throw new DuplicateSequenceException(s"Duplicate sequence number for producerId $producerId: (incomingBatch.firstSeq, " +
-        s"incomingBatch.lastSeq): ($firstSeq, $lastSeq).")
-    } else if (!inSequence(firstSeq, lastSeq)) {
-      throw new OutOfOrderSequenceException(s"Out of order sequence number for producerId $producerId: $firstSeq " +
-        s"(incoming seq. number), ${currentEntry.lastSeq} (current end sequence number)")
+    } else {
+      val currentLastSeq = if (!updatedEntry.isEmpty)
+        updatedEntry.lastSeq
+      else if (producerEpoch == currentEntry.producerEpoch)
+        currentEntry.lastSeq
+      else
+        RecordBatch.NO_SEQUENCE
+
+      // If there is no current producer epoch (possibly because all producer records have been deleted due to
+      // retention or the DeleteRecords API) accept writes with any sequence number
+      if (!(currentEntry.producerEpoch == RecordBatch.NO_PRODUCER_EPOCH || inSequence(currentLastSeq, appendFirstSeq))) {
+        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)")
+      }
     }
   }
 
-  private def isDuplicate(firstSeq: Int, lastSeq: Int): Boolean = {
-    ((lastSeq != 0 && currentEntry.firstSeq != Int.MaxValue && lastSeq < currentEntry.firstSeq)
-      || currentEntry.batchWithSequenceRange(firstSeq, lastSeq).isDefined)
+  private def inSequence(lastSeq: Int, nextSeq: Int): Boolean = {
+    nextSeq == lastSeq + 1L || (nextSeq == 0 && lastSeq == Int.MaxValue)
   }
 
-  private def inSequence(firstSeq: Int, lastSeq: Int): Boolean = {
-    firstSeq == currentEntry.lastSeq + 1L || (firstSeq == 0 && currentEntry.lastSeq == Int.MaxValue)
-  }
-
-  private def isFenced(producerEpoch: Short): Boolean = {
-    producerEpoch < currentEntry.producerEpoch
-  }
-
-  def append(batch: RecordBatch): Option[CompletedTxn] = {
+  def append(batch: RecordBatch, firstOffsetMetadataOpt: Option[LogOffsetMetadata]): Option[CompletedTxn] = {
     if (batch.isControlBatch) {
-      val record = batch.iterator.next()
-      val endTxnMarker = EndTransactionMarker.deserialize(record)
-      val completedTxn = appendEndTxnMarker(endTxnMarker, batch.producerEpoch, batch.baseOffset, record.timestamp)
-      Some(completedTxn)
+      val recordIterator = batch.iterator
+      if (recordIterator.hasNext) {
+        val record = recordIterator.next()
+        val endTxnMarker = EndTransactionMarker.deserialize(record)
+        appendEndTxnMarker(endTxnMarker, batch.producerEpoch, batch.baseOffset, record.timestamp)
+      } else {
+        // An empty control batch means the entire transaction has been cleaned from the log, so no need to append
+        None
+      }
     } else {
-      append(batch.producerEpoch, batch.baseSequence, batch.lastSequence, batch.maxTimestamp, 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,
-             lastOffset: Long,
-             isTransactional: Boolean): Unit = {
-    maybeValidateAppend(epoch, firstSeq, lastSeq)
-
-    currentEntry.addBatchMetadata(epoch, lastSeq, lastOffset, lastSeq - firstSeq, lastTimestamp)
-
-    if (currentEntry.currentTxnFirstOffset.isDefined && !isTransactional)
-      throw new InvalidTxnStateException(s"Expected transactional write from producer $producerId")
-
-    if (isTransactional && currentEntry.currentTxnFirstOffset.isEmpty) {
-      val firstOffset = lastOffset - (lastSeq - firstSeq)
-      currentEntry.currentTxnFirstOffset = Some(firstOffset)
-      transactions += new TxnMetadata(producerId, firstOffset)
+  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 at " +
+          s"offset $firstOffsetMetadata in partition $topicPartition")
+
+      case None if isTransactional =>
+        // Began a new transaction
+        updatedEntry.currentTxnFirstOffset = Some(firstOffset)
+        transactions += TxnMetadata(producerId, firstOffsetMetadata)
+
+      case _ => // nothing to do
     }
   }
 
-  def appendEndTxnMarker(endTxnMarker: EndTransactionMarker,
-                         producerEpoch: Short,
-                         offset: Long,
-                         timestamp: Long): CompletedTxn = {
-    if (isFenced(producerEpoch))
-      throw new ProducerFencedException(s"Invalid producer epoch: $producerEpoch (zombie): ${currentEntry.producerEpoch} (current)")
-
-    if (currentEntry.coordinatorEpoch > endTxnMarker.coordinatorEpoch)
-      throw new TransactionCoordinatorFencedException(s"Invalid coordinator epoch: ${endTxnMarker.coordinatorEpoch} " +
-        s"(zombie), ${currentEntry.coordinatorEpoch} (current)")
-
-    currentEntry.maybeUpdateEpoch(producerEpoch)
+  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)")
+      }
+    }
+  }
 
-    val firstOffset = currentEntry.currentTxnFirstOffset match {
-      case Some(txnFirstOffset) => txnFirstOffset
-      case None =>
-        transactions += new TxnMetadata(producerId, offset)
-        offset
+  def appendEndTxnMarker(
+    endTxnMarker: EndTransactionMarker,
+    producerEpoch: Short,
+    offset: Long,
+    timestamp: Long
+  ): Option[CompletedTxn] = {
+    checkProducerEpoch(producerEpoch, offset)
+    checkCoordinatorEpoch(endTxnMarker, offset)
+
+    // Only emit the `CompletedTxn` for non-empty transactions. A transaction marker
+    // without any associated data will not have any impact on the last stable offset
+    // and would not need to be reflected in the transaction index.
+    val completedTxn = updatedEntry.currentTxnFirstOffset.map { firstOffset =>
+      CompletedTxn(producerId, firstOffset, offset, endTxnMarker.controlType == ControlRecordType.ABORT)
     }
 
-    currentEntry.currentTxnFirstOffset = None
-    currentEntry.coordinatorEpoch = endTxnMarker.coordinatorEpoch
-    CompletedTxn(producerId, firstOffset, offset, endTxnMarker.controlType == ControlRecordType.ABORT)
+    updatedEntry.maybeUpdateProducerEpoch(producerEpoch)
+    updatedEntry.currentTxnFirstOffset = None
+    updatedEntry.coordinatorEpoch = endTxnMarker.coordinatorEpoch
+    updatedEntry.lastTimestamp = timestamp
+
+    completedTxn
   }
 
-  def latestEntry: ProducerIdEntry = currentEntry
+  def toEntry: ProducerStateEntry = updatedEntry
 
   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, " +
-      s"producerEpoch=${currentEntry.producerEpoch}, " +
-      s"firstSequence=${currentEntry.firstSeq}, " +
-      s"lastSequence=${currentEntry.lastSeq}, " +
-      s"currentTxnFirstOffset=${currentEntry.currentTxnFirstOffset}, " +
-      s"coordinatorEpoch=${currentEntry.coordinatorEpoch}, " +
+      s"producerEpoch=${updatedEntry.producerEpoch}, " +
+      s"firstSequence=${updatedEntry.firstSeq}, " +
+      s"lastSequence=${updatedEntry.lastSeq}, " +
+      s"currentTxnFirstOffset=${updatedEntry.currentTxnFirstOffset}, " +
+      s"coordinatorEpoch=${updatedEntry.coordinatorEpoch}, " +
+      s"lastTimestamp=${updatedEntry.lastTimestamp}, " +
       s"startedTransactions=$transactions)"
   }
 }
@@ -347,7 +377,7 @@ object ProducerStateManager {
     new Field(CrcField, Type.UNSIGNED_INT32, "CRC of the snapshot data"),
     new Field(ProducerEntriesField, new ArrayOf(ProducerSnapshotEntrySchema), "The entries in the producer table"))
 
-  def readSnapshot(file: File): Iterable[ProducerIdEntry] = {
+  def readSnapshot(file: File): Iterable[ProducerStateEntry] = {
     try {
       val buffer = Files.readAllBytes(file.toPath)
       val struct = PidSnapshotMapSchema.read(ByteBuffer.wrap(buffer))
@@ -364,7 +394,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)
@@ -372,8 +402,12 @@ object ProducerStateManager {
         val offsetDelta = producerEntryStruct.getInt(OffsetDeltaField)
         val coordinatorEpoch = producerEntryStruct.getInt(CoordinatorEpochField)
         val currentTxnFirstOffset = producerEntryStruct.getLong(CurrentTxnFirstOffsetField)
-        val newEntry = new ProducerIdEntry(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 {
@@ -382,7 +416,7 @@ object ProducerStateManager {
     }
   }
 
-  private def writeSnapshot(file: File, entries: mutable.Map[Long, ProducerIdEntry]) {
+  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
@@ -409,34 +443,25 @@ object ProducerStateManager {
     val crc = Crc32C.compute(buffer, ProducerEntriesOffset, buffer.limit() - ProducerEntriesOffset)
     ByteUtils.writeUnsignedInt(buffer, CrcOffset, crc)
 
-    val fos = new FileOutputStream(file)
+    val fileChannel = FileChannel.open(file.toPath, StandardOpenOption.CREATE, StandardOpenOption.WRITE)
     try {
-      fos.write(buffer.array, buffer.arrayOffset, buffer.limit())
+      fileChannel.write(buffer)
+      fileChannel.force(true)
     } finally {
-      fos.close()
+      fileChannel.close()
     }
   }
 
   private def isSnapshotFile(file: File): Boolean = file.getName.endsWith(Log.ProducerSnapshotFileSuffix)
 
   // visible for testing
-  private[log] def listSnapshotFiles(dir: File): Seq[File] = {
+  private[log] def listSnapshotFiles(dir: File): Seq[SnapshotFile] = {
     if (dir.exists && dir.isDirectory) {
       Option(dir.listFiles).map { files =>
-        files.filter(f => f.isFile && isSnapshotFile(f)).toSeq
+        files.filter(f => f.isFile && isSnapshotFile(f)).map(SnapshotFile(_)).toSeq
       }.getOrElse(Seq.empty)
     } else Seq.empty
   }
-
-  // visible for testing
-  private[log] def deleteSnapshotsBefore(dir: File, offset: Long): Unit = deleteSnapshotFiles(dir, _ < offset)
-
-  private def deleteSnapshotFiles(dir: File, predicate: Long => Boolean = _ => true) {
-    listSnapshotFiles(dir).filter(file => predicate(offsetFromFile(file))).foreach { file =>
-      Files.deleteIfExists(file.toPath)
-    }
-  }
-
 }
 
 /**
@@ -457,12 +482,18 @@ object ProducerStateManager {
  */
 @nonthreadsafe
 class ProducerStateManager(val topicPartition: TopicPartition,
-                           @volatile var logDir: File,
+                           @volatile var _logDir: File,
                            val maxProducerIdExpirationMs: Int = 60 * 60 * 1000) extends Logging {
   import ProducerStateManager._
   import java.util
 
-  private val producers = mutable.Map.empty[Long, ProducerIdEntry]
+  this.logIdent = s"[ProducerStateManager partition=$topicPartition] "
+
+  private var snapshots: ConcurrentSkipListMap[java.lang.Long, SnapshotFile] = locally {
+    loadSnapshots()
+  }
+
+  private val producers = mutable.Map.empty[Long, ProducerStateEntry]
   private var lastMapOffset = 0L
   private var lastSnapOffset = 0L
 
@@ -472,6 +503,59 @@ class ProducerStateManager(val topicPartition: TopicPartition,
   // completed transactions whose markers are at offsets above the high watermark
   private val unreplicatedTxns = new util.TreeMap[Long, TxnMetadata]
 
+  /**
+   * Load producer state snapshots by scanning the _logDir.
+   */
+  private def loadSnapshots(): ConcurrentSkipListMap[java.lang.Long, SnapshotFile] = {
+    val tm = new ConcurrentSkipListMap[java.lang.Long, SnapshotFile]()
+    for (f <- listSnapshotFiles(_logDir)) {
+      tm.put(f.offset, f)
+    }
+    tm
+  }
+
+  /**
+   * Scans the log directory, gathering all producer state snapshot files. Snapshot files which do not have an offset
+   * corresponding to one of the provided offsets in segmentBaseOffsets will be removed, except in the case that there
+   * is a snapshot file at a higher offset than any offset in segmentBaseOffsets.
+   *
+   * The goal here is to remove any snapshot files which do not have an associated segment file, but not to remove the
+   * largest stray snapshot file which was emitted during clean shutdown.
+   */
+  private[log] def removeStraySnapshots(segmentBaseOffsets: Seq[Long]): Unit = {
+    val maxSegmentBaseOffset = if (segmentBaseOffsets.isEmpty) None else Some(segmentBaseOffsets.max)
+    val baseOffsets = segmentBaseOffsets.toSet
+    var latestStraySnapshot: Option[SnapshotFile] = None
+
+    val ss = loadSnapshots()
+    for (snapshot <- ss.values().asScala) {
+      val key = snapshot.offset
+      latestStraySnapshot match {
+        case Some(prev) =>
+          if (!baseOffsets.contains(key)) {
+            // this snapshot is now the largest stray snapshot.
+            prev.deleteIfExists()
+            ss.remove(prev.offset)
+            latestStraySnapshot = Some(snapshot)
+          }
+        case None =>
+          if (!baseOffsets.contains(key)) {
+            latestStraySnapshot = Some(snapshot)
+          }
+      }
+    }
+
+    // Check to see if the latestStraySnapshot is larger than the largest segment base offset, if it is not,
+    // delete the largestStraySnapshot.
+    for (strayOffset <- latestStraySnapshot.map(_.offset); maxOffset <- maxSegmentBaseOffset) {
+      if (strayOffset < maxOffset) {
+        Option(ss.remove(strayOffset)).foreach(_.deleteIfExists())
+      }
+    }
+
+    this.snapshots = ss
+  }
+
   /**
    * An unstable offset is one which is either undecided (i.e. its ultimate outcome is not yet known),
    * or one that is decided, but may not have been replicated (i.e. any transaction which has a COMMIT/ABORT
@@ -500,39 +584,38 @@ class ProducerStateManager(val topicPartition: TopicPartition,
 
   /**
    * The first undecided offset is the earliest transactional message which has not yet been committed
-   * or aborted.
+   * or aborted. Unlike [[firstUnstableOffset]], this does not reflect the state of replication (i.e.
+   * whether a completed transaction marker is beyond the high watermark).
    */
-  def firstUndecidedOffset: Option[Long] = Option(ongoingTxns.firstEntry).map(_.getValue.firstOffset.messageOffset)
+  private[log] def firstUndecidedOffset: Option[Long] = Option(ongoingTxns.firstEntry).map(_.getValue.firstOffset.messageOffset)
 
   /**
    * Returns the last offset of this map
    */
-  def mapEndOffset = lastMapOffset
+  def mapEndOffset: Long = lastMapOffset
 
   /**
    * Get a copy of the active producers
    */
-  def activeProducers: immutable.Map[Long, ProducerIdEntry] = producers.toMap
+  def activeProducers: immutable.Map[Long, ProducerStateEntry] = producers.toMap
 
   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) =>
+        case Some(snapshot) =>
           try {
-            info(s"Loading producer state from snapshot file '$file' for partition $topicPartition")
-            val loadedProducers = readSnapshot(file).filter { producerEntry =>
-              isProducerRetained(producerEntry, logStartOffset) && !isProducerExpired(currentTime, producerEntry)
-            }
+            info(s"Loading producer state from snapshot file '$snapshot'")
+            val loadedProducers = readSnapshot(snapshot.file).filter { producerEntry => !isProducerExpired(currentTime, producerEntry) }
             loadedProducers.foreach(loadProducerEntry)
-            lastSnapOffset = offsetFromFile(file)
+            lastSnapOffset = snapshot.offset
             lastMapOffset = lastSnapOffset
             return
           } catch {
             case e: CorruptSnapshotException =>
-              warn(s"Failed to load producer snapshot from '$file': ${e.getMessage}")
-              Files.deleteIfExists(file.toPath)
+              warn(s"Failed to load producer snapshot from '${snapshot.file}': ${e.getMessage}")
+              removeAndDeleteSnapshot(snapshot.offset)
           }
         case None =>
           lastSnapOffset = logStartOffset
@@ -543,7 +626,7 @@ class ProducerStateManager(val topicPartition: TopicPartition,
   }
 
   // visible for testing
-  private[log] def loadProducerEntry(entry: ProducerIdEntry): Unit = {
+  private[log] def loadProducerEntry(entry: ProducerStateEntry): Unit = {
     val producerId = entry.producerId
     producers.put(producerId, entry)
     entry.currentTxnFirstOffset.foreach { offset =>
@@ -551,28 +634,31 @@ class ProducerStateManager(val topicPartition: TopicPartition,
     }
   }
 
-  private def isProducerExpired(currentTimeMs: Long, producerIdEntry: ProducerIdEntry): Boolean =
-    producerIdEntry.currentTxnFirstOffset.isEmpty && currentTimeMs - producerIdEntry.lastTimestamp >= maxProducerIdExpirationMs
+  private def isProducerExpired(currentTimeMs: Long, producerState: ProducerStateEntry): Boolean =
+    producerState.currentTxnFirstOffset.isEmpty && currentTimeMs - producerState.lastTimestamp >= maxProducerIdExpirationMs
 
   /**
    * Expire any producer ids which have been idle longer than the configured maximum expiration timeout.
    */
-  def removeExpiredProducers(currentTimeMs: Long) {
-    producers.retain { case (producerId, lastEntry) =>
-      !isProducerExpired(currentTimeMs, lastEntry)
-    }
+  def removeExpiredProducers(currentTimeMs: Long): Unit = {
+    producers --= producers.filter { case (_, lastEntry) => isProducerExpired(currentTimeMs, lastEntry) }.keySet
   }
 
   /**
    * 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
-    })
+    snapshots.values().asScala.foreach { snapshot =>
+      if (snapshot.offset > logEndOffset || snapshot.offset <= logStartOffset) {
+        removeAndDeleteSnapshot(snapshot.offset)
+      }
+    }
 
     if (logEndOffset != mapEndOffset) {
       producers.clear()
@@ -583,20 +669,13 @@ class ProducerStateManager(val topicPartition: TopicPartition,
       unreplicatedTxns.clear()
       loadFromSnapshot(logStartOffset, currentTimeMs)
     } else {
-      truncateHead(logStartOffset)
+      onLogStartOffsetIncremented(logStartOffset)
     }
   }
 
-  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
-
-    new ProducerAppendInfo(producerId, lastEntry(producerId).getOrElse(ProducerIdEntry.empty(producerId)), validationToPerform)
+  def prepareUpdate(producerId: Long, origin: AppendOrigin): ProducerAppendInfo = {
+    val currentEntry = lastEntry(producerId).getOrElse(ProducerStateEntry.empty(producerId))
+    new ProducerAppendInfo(topicPartition, producerId, currentEntry, origin)
   }
 
   /**
@@ -604,12 +683,19 @@ class ProducerStateManager(val topicPartition: TopicPartition,
    */
   def update(appendInfo: ProducerAppendInfo): Unit = {
     if (appendInfo.producerId == RecordBatch.NO_PRODUCER_ID)
-      throw new IllegalArgumentException(s"Invalid producer id ${appendInfo.producerId} passed to update")
+      throw new IllegalArgumentException(s"Invalid producer id ${appendInfo.producerId} passed to update " +
+        s"for partition $topicPartition")
 
     trace(s"Updated producer ${appendInfo.producerId} state to $appendInfo")
+    val updatedEntry = appendInfo.toEntry
+    producers.get(appendInfo.producerId) match {
+      case Some(currentEntry) =>
+        currentEntry.update(updatedEntry)
+
+      case None =>
+        producers.put(appendInfo.producerId, updatedEntry)
+    }
 
-    val entry = appendInfo.latestEntry
-    producers.put(appendInfo.producerId, entry)
     appendInfo.startedTransactions.foreach { txn =>
       ongoingTxns.put(txn.firstOffset.messageOffset, txn)
     }
@@ -622,7 +708,7 @@ class ProducerStateManager(val topicPartition: TopicPartition,
   /**
    * Get the last written entry for the given producer id.
    */
-  def lastEntry(producerId: Long): Option[ProducerIdEntry] = producers.get(producerId)
+  def lastEntry(producerId: Long): Option[ProducerStateEntry] = producers.get(producerId)
 
   /**
    * Take a snapshot at the current end offset if one does not already exist.
@@ -630,65 +716,47 @@ class ProducerStateManager(val topicPartition: TopicPartition,
   def takeSnapshot(): Unit = {
     // If not a new offset, then it is not worth taking another snapshot
     if (lastMapOffset > lastSnapOffset) {
-      val snapshotFile = Log.producerSnapshotFile(logDir, lastMapOffset)
-      debug(s"Writing producer snapshot for partition $topicPartition at offset $lastMapOffset")
-      writeSnapshot(snapshotFile, producers)
+      val snapshotFile = SnapshotFile(Log.producerSnapshotFile(_logDir, lastMapOffset))
+      info(s"Writing producer snapshot at offset $lastMapOffset")
+      writeSnapshot(snapshotFile.file, producers)
+      snapshots.put(snapshotFile.offset, snapshotFile)
 
       // Update the last snap offset according to the serialized map
       lastSnapOffset = lastMapOffset
     }
   }
 
+  /**
+   * Update the parentDir for this ProducerStateManager and all of the snapshot files which it manages.
+   */
+  def updateParentDir(parentDir: File): Unit ={
+    _logDir = parentDir
+    snapshots.forEach((_, s) => s.updateParentDir(parentDir))
+  }
+
   /**
    * Get the last offset (exclusive) of the latest snapshot file.
    */
-  def latestSnapshotOffset: Option[Long] = latestSnapshotFile.map(file => offsetFromFile(file))
+  def latestSnapshotOffset: Option[Long] = latestSnapshotFile.map(_.offset)
 
   /**
    * Get the last offset (exclusive) of the oldest snapshot file.
    */
-  def oldestSnapshotOffset: Option[Long] = oldestSnapshotFile.map(file => offsetFromFile(file))
-
-  private def isProducerRetained(producerIdEntry: ProducerIdEntry, logStartOffset: Long): Boolean = {
-    producerIdEntry.removeBatchesOlderThan(logStartOffset)
-    producerIdEntry.lastDataOffset >= logStartOffset
-  }
+  def oldestSnapshotOffset: Option[Long] = oldestSnapshotFile.map(_.offset)
 
   /**
-   * 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.
+   * Remove any unreplicated transactions lower than the provided logStartOffset and bring the lastMapOffset forward
+   * if necessary.
    */
-  def truncateHead(logStartOffset: Long) {
-    val evictedProducerEntries = producers.filter { case (_, producerIdEntry) =>
-      !isProducerRetained(producerIdEntry, logStartOffset)
-    }
-    val evictedProducerIds = evictedProducerEntries.keySet
-
-    producers --= evictedProducerIds
-    removeEvictedOngoingTransactions(evictedProducerIds)
+  def onLogStartOffsetIncremented(logStartOffset: Long): Unit = {
     removeUnreplicatedTransactions(logStartOffset)
 
     if (lastMapOffset < logStartOffset)
       lastMapOffset = logStartOffset
 
-    deleteSnapshotsBefore(logStartOffset)
     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) {
@@ -702,49 +770,83 @@ 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()
-    deleteSnapshotFiles(logDir)
+    snapshots.values().asScala.foreach { snapshot =>
+      removeAndDeleteSnapshot(snapshot.offset)
+    }
     lastSnapOffset = 0L
     lastMapOffset = 0L
   }
 
   /**
-   * 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("Attempted to complete a transaction which was not started")
+      throw new IllegalArgumentException(s"Attempted to complete transaction $completedTxn on partition $topicPartition " +
+        s"which was not started")
 
     txnMetadata.lastOffset = Some(completedTxn.lastOffset)
     unreplicatedTxns.put(completedTxn.firstOffset, txnMetadata)
-
-    val lastStableOffset = firstUndecidedOffset.getOrElse(completedTxn.lastOffset + 1)
-    lastStableOffset
   }
 
   @threadsafe
-  def deleteSnapshotsBefore(offset: Long): Unit = ProducerStateManager.deleteSnapshotsBefore(logDir, offset)
+  def deleteSnapshotsBefore(offset: Long): Unit = {
+    snapshots.subMap(0, offset).values().asScala.foreach { snapshot =>
+      removeAndDeleteSnapshot(snapshot.offset)
+    }
+  }
 
-  private def oldestSnapshotFile: Option[File] = {
-    val files = listSnapshotFiles
-    if (files.nonEmpty)
-      Some(files.minBy(offsetFromFile))
-    else
-      None
+  private def oldestSnapshotFile: Option[SnapshotFile] = {
+    Option(snapshots.firstEntry()).map(_.getValue)
   }
 
-  private def latestSnapshotFile: Option[File] = {
-    val files = listSnapshotFiles
-    if (files.nonEmpty)
-      Some(files.maxBy(offsetFromFile))
-    else
-      None
+  private def latestSnapshotFile: Option[SnapshotFile] = {
+    Option(snapshots.lastEntry()).map(_.getValue)
+  }
+
+  /**
+   * Removes the producer state snapshot file metadata corresponding to the provided offset if it exists from this
+   * ProducerStateManager, and deletes the backing snapshot file.
+   */
+  private[log] def removeAndDeleteSnapshot(snapshotOffset: Long): Unit = {
+    Option(snapshots.remove(snapshotOffset)).foreach(_.deleteIfExists())
+  }
+}
+
+case class SnapshotFile private[log] (private var _file: File,
+                                      offset: Long) {
+  def deleteIfExists(): Boolean = {
+    Files.deleteIfExists(file.toPath)
   }
 
-  private def listSnapshotFiles: Seq[File] = ProducerStateManager.listSnapshotFiles(logDir)
+  def updateParentDir(parentDir: File): Unit = {
+    _file = new File(parentDir, _file.getName)
+  }
 
+  def file: File = {
+    _file
+  }
+}
+
+object SnapshotFile {
+  def apply(file: File): SnapshotFile = {
+    val offset = offsetFromFile(file)
+    SnapshotFile(file, offset)
+  }
 }
diff --git a/core/src/main/scala/kafka/log/TimeIndex.scala b/core/src/main/scala/kafka/log/TimeIndex.scala
index 47ab2e527eb6f..6b2739cabcaf6 100644
--- a/core/src/main/scala/kafka/log/TimeIndex.scala
+++ b/core/src/main/scala/kafka/log/TimeIndex.scala
@@ -20,9 +20,9 @@ package kafka.log
 import java.io.File
 import java.nio.ByteBuffer
 
-import kafka.common.InvalidOffsetException
-import kafka.utils.CoreUtils._
+import kafka.utils.CoreUtils.inLock
 import kafka.utils.Logging
+import org.apache.kafka.common.errors.InvalidOffsetException
 import org.apache.kafka.common.record.RecordBatch
 
 /**
@@ -51,12 +51,16 @@ 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) with Logging {
+    extends AbstractIndex(_file, baseOffset, maxIndexSize, writable) {
+  import TimeIndex._
 
   @volatile private var _lastEntry = lastEntryFromIndexFile
 
   override def entrySize = 12
 
+  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
 
@@ -73,7 +77,7 @@ class TimeIndex(_file: File, baseOffset: Long, maxIndexSize: Int = -1, writable:
     inLock(lock) {
       _entries match {
         case 0 => TimestampOffset(RecordBatch.NO_TIMESTAMP, baseOffset)
-        case s => parseEntry(mmap, s - 1).asInstanceOf[TimestampOffset]
+        case s => parseEntry(mmap, s - 1)
       }
     }
   }
@@ -86,13 +90,13 @@ 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))
-      val idx = mmap.duplicate
-      TimestampOffset(timestamp(idx, n), relativeOffset(idx, n))
+        throw new IllegalArgumentException(s"Attempt to fetch the ${n}th entry from  time index ${file.getAbsolutePath} " +
+          s"which has size ${_entries}.")
+      parseEntry(mmap, n)
     }
   }
 
-  override def parseEntry(buffer: ByteBuffer, n: Int): IndexEntry = {
+  override def parseEntry(buffer: ByteBuffer, n: Int): TimestampOffset = {
     TimestampOffset(timestamp(buffer, n), baseOffset + relativeOffset(buffer, n))
   }
 
@@ -106,7 +110,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 + ").")
@@ -117,21 +121,21 @@ class TimeIndex(_file: File, baseOffset: Long, maxIndexSize: Int = -1, writable:
       // 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))
+        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 (_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))
+        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))
+        trace(s"Adding index entry $timestamp => $offset to ${file.getAbsolutePath}.")
         mmap.putLong(timestamp)
-        mmap.putInt((offset - baseOffset).toInt)
+        mmap.putInt(relativeOffset(offset))
         _entries += 1
         _lastEntry = TimestampOffset(timestamp, offset)
-        require(_entries * entrySize == mmap.position(), _entries + " entries but file position in index is " + mmap.position() + ".")
+        require(_entries * entrySize == mmap.position(), s"${_entries} entries but file position in index is ${mmap.position()}.")
       }
     }
   }
@@ -150,10 +154,8 @@ class TimeIndex(_file: File, baseOffset: Long, maxIndexSize: Int = -1, writable:
       val slot = largestLowerBoundSlotFor(idx, targetTimestamp, IndexSearchType.KEY)
       if (slot == -1)
         TimestampOffset(RecordBatch.NO_TIMESTAMP, baseOffset)
-      else {
-        val entry = parseEntry(idx, slot).asInstanceOf[TimestampOffset]
-        TimestampOffset(entry.timestamp, entry.offset)
-      }
+      else
+        parseEntry(idx, slot)
     }
   }
 
@@ -163,7 +165,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)
@@ -197,25 +199,31 @@ 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
+      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
-    require(_entries == 0 || (lastTimestamp >= timestamp(mmap, 0)),
-      s"Corrupt time index found, time index file (${file.getAbsolutePath}) has non-zero size but the last timestamp " +
-          s"is $lastTimestamp which is no larger than the first timestamp ${timestamp(mmap, 0)}")
-    require(_entries == 0 || lastOffset >= baseOffset,
-      s"Corrupt time index found, time index file (${file.getAbsolutePath}) has non-zero size but the last offset " +
-          s"is $lastOffset which is smaller than the first offset $baseOffset")
-    require(length % entrySize == 0,
-      "Time index file " + file.getAbsolutePath + " is corrupt, found " + length +
-          " bytes which is not positive or not a multiple of 12.")
+    if (_entries != 0 && lastTimestamp < timestamp(mmap, 0))
+      throw new CorruptIndexException(s"Corrupt time index found, time index file (${file.getAbsolutePath}) has " +
+        s"non-zero size but the last timestamp is $lastTimestamp which is less than the first timestamp " +
+        s"${timestamp(mmap, 0)}")
+    if (_entries != 0 && lastOffset < baseOffset)
+      throw new CorruptIndexException(s"Corrupt time index found, time index file (${file.getAbsolutePath}) has " +
+        s"non-zero size but the last offset is $lastOffset which is less than the first offset $baseOffset")
+    if (length % entrySize != 0)
+      throw new CorruptIndexException(s"Time index file ${file.getAbsolutePath} is corrupt, found $length bytes " +
+        s"which is neither positive nor a multiple of $entrySize.")
   }
 }
+
+object TimeIndex extends Logging {
+  override val loggerName: String = classOf[TimeIndex].getName
+}
diff --git a/core/src/main/scala/kafka/log/TransactionIndex.scala b/core/src/main/scala/kafka/log/TransactionIndex.scala
index 9afe009664729..9152bc41ab353 100644
--- a/core/src/main/scala/kafka/log/TransactionIndex.scala
+++ b/core/src/main/scala/kafka/log/TransactionIndex.scala
@@ -19,7 +19,7 @@ package kafka.log
 import java.io.{File, IOException}
 import java.nio.ByteBuffer
 import java.nio.channels.FileChannel
-import java.nio.file.StandardOpenOption
+import java.nio.file.{Files, StandardOpenOption}
 
 import kafka.utils.{Logging, nonthreadsafe}
 import org.apache.kafka.common.KafkaException
@@ -42,34 +42,44 @@ 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 delete(): Boolean = {
-    maybeChannel.forall { channel =>
-      channel.force(true)
-      close()
-      file.delete()
-    }
+  def file: File = _file
+
+  def updateParentDir(parentDir: File): Unit = _file = new File(parentDir, file.getName)
+
+  /**
+   * Delete this index.
+   *
+   * @throws IOException if deletion fails due to an I/O error
+   * @return `true` if the file was deleted by this method; `false` if the file could not be deleted because it did
+   *         not exist
+   */
+  def deleteIfExists(): Boolean = {
+    close()
+    Files.deleteIfExists(file.toPath)
   }
 
-  private def channel: FileChannel = {
+  private def channel(): FileChannel = {
     maybeChannel match {
       case Some(channel) => channel
       case None => openChannel()
@@ -77,14 +87,17 @@ class TransactionIndex(val startOffset: Long, @volatile var file: File) extends
   }
 
   private def openChannel(): FileChannel = {
-    val channel = FileChannel.open(file.toPath, StandardOpenOption.READ, StandardOpenOption.WRITE,
-      StandardOpenOption.CREATE)
+    val channel = FileChannel.open(file.toPath, StandardOpenOption.CREATE, StandardOpenOption.READ,
+      StandardOpenOption.WRITE)
     maybeChannel = Some(channel)
     channel.position(channel.size)
     channel
   }
 
-  def truncate() = {
+  /**
+   * Remove all the entries from the index. Unlike `AbstractIndex`, this index is not resized ahead of time.
+   */
+  def reset(): Unit = {
     maybeChannel.foreach(_.truncate(0))
     lastOffset = None
   }
@@ -98,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 = {
@@ -106,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
       }
@@ -131,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
@@ -140,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)
             }
           }
         }
@@ -171,11 +184,17 @@ class TransactionIndex(val startOffset: Long, @volatile var file: File) extends
     TxnIndexSearchResult(abortedTransactions.toList, isComplete = false)
   }
 
+  /**
+   * Do a basic sanity check on this index to detect obvious problems.
+   *
+   * @throws CorruptIndexException if any problems are found.
+   */
   def sanityCheck(): Unit = {
     val buffer = ByteBuffer.allocate(AbortedTxn.TotalSize)
     for ((abortedTxn, _) <- iterator(() => buffer)) {
-      require(abortedTxn.lastOffset >= startOffset, s"Last offset of aborted transaction $abortedTxn is less than " +
-        s"start offset $startOffset")
+      if (abortedTxn.lastOffset < 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/log/package.html b/core/src/main/scala/kafka/log/package.html
index c6ebf0c26bec3..ee2f72e0095f7 100644
--- a/core/src/main/scala/kafka/log/package.html
+++ b/core/src/main/scala/kafka/log/package.html
@@ -21,4 +21,4 @@
 The entry point for this system is LogManager. LogManager is responsible for holding all the logs, and handing them out by topic/partition. It also handles the enforcement of the
 flush policy and retention policies.
 
-The Log itself is made up of log segments. A log is a FileMessageSet that contains the data and an OffsetIndex that supports reads by offset on the log.
\ No newline at end of file
+The Log itself is made up of log segments. A log is a FileRecords that contains the data and an OffsetIndex that supports reads by offset on the log.
\ No newline at end of file
diff --git a/core/src/main/scala/kafka/message/ByteBufferMessageSet.scala b/core/src/main/scala/kafka/message/ByteBufferMessageSet.scala
deleted file mode 100644
index 62e2125fabdb4..0000000000000
--- a/core/src/main/scala/kafka/message/ByteBufferMessageSet.scala
+++ /dev/null
@@ -1,197 +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.message
-
-import java.nio.ByteBuffer
-
-import kafka.common.LongRef
-import kafka.utils.Logging
-import org.apache.kafka.common.record._
-
-import scala.collection.JavaConverters._
-
-object ByteBufferMessageSet {
-
-  private def create(offsetAssigner: OffsetAssigner,
-                     compressionCodec: CompressionCodec,
-                     timestampType: TimestampType,
-                     messages: Message*): ByteBuffer = {
-    if (messages.isEmpty)
-      MessageSet.Empty.buffer
-    else {
-      val buffer = ByteBuffer.allocate(math.min(math.max(MessageSet.messageSetSize(messages) / 2, 1024), 1 << 16))
-      val builder = MemoryRecords.builder(buffer, messages.head.magic, CompressionType.forId(compressionCodec.codec),
-        timestampType, offsetAssigner.baseOffset)
-
-      for (message <- messages)
-        builder.appendWithOffset(offsetAssigner.nextAbsoluteOffset(), message.asRecord)
-
-      builder.build().buffer
-    }
-  }
-
-}
-
-private object OffsetAssigner {
-
-  def apply(offsetCounter: LongRef, size: Int): OffsetAssigner =
-    new OffsetAssigner(offsetCounter.value to offsetCounter.addAndGet(size))
-
-}
-
-private class OffsetAssigner(offsets: Seq[Long]) {
-  private var index = 0
-
-  def nextAbsoluteOffset(): Long = {
-    val result = offsets(index)
-    index += 1
-    result
-  }
-
-  def baseOffset = offsets.head
-
-  def toInnerOffset(offset: Long): Long = offset - offsets.head
-
-}
-
-/**
- * A sequence of messages stored in a byte buffer
- *
- * There are two ways to create a ByteBufferMessageSet
- *
- * Option 1: From a ByteBuffer which already contains the serialized message set. Consumers will use this method.
- *
- * Option 2: Give it a list of messages along with instructions relating to serialization format. Producers will use this method.
- *
- *
- * Message format v1 has the following changes:
- * - For non-compressed messages, timestamp and timestamp type attributes have been added. The offsets of
- *   the messages remain absolute offsets.
- * - For compressed messages, timestamp and timestamp type attributes have been added and inner offsets (IO) are used
- *   for inner messages of compressed messages (see offset calculation details below). The timestamp type
- *   attribute is only set in wrapper messages. Inner messages always have CreateTime as the timestamp type in attributes.
- *
- * We set the timestamp in the following way:
- * For non-compressed messages: the timestamp and timestamp type message attributes are set and used.
- * For compressed messages:
- * 1. Wrapper messages' timestamp type attribute is set to the proper value
- * 2. Wrapper messages' timestamp is set to:
- *    - the max timestamp of inner messages if CreateTime is used
- *    - the current server time if wrapper message's timestamp = LogAppendTime.
- *      In this case the wrapper message timestamp is used and all the timestamps of inner messages are ignored.
- * 3. Inner messages' timestamp will be:
- *    - used when wrapper message's timestamp type is CreateTime
- *    - ignored when wrapper message's timestamp type is LogAppendTime
- * 4. Inner messages' timestamp type will always be ignored with one exception: producers must set the inner message
- *    timestamp type to CreateTime, otherwise the messages will be rejected by broker.
- *
- * Absolute offsets are calculated in the following way:
- * Ideally the conversion from relative offset(RO) to absolute offset(AO) should be:
- *
- * AO = AO_Of_Last_Inner_Message + RO
- *
- * However, note that the message sets sent by producers are compressed in a streaming way.
- * And the relative offset of an inner message compared with the last inner message is not known until
- * the last inner message is written.
- * Unfortunately we are not able to change the previously written messages after the last message is written to
- * the message set when stream compression is used.
- *
- * To solve this issue, we use the following solution:
- *
- * 1. When the producer creates a message set, it simply writes all the messages into a compressed message set with
- *    offset 0, 1, ... (inner offset).
- * 2. The broker will set the offset of the wrapper message to the absolute offset of the last message in the
- *    message set.
- * 3. When a consumer sees the message set, it first decompresses the entire message set to find out the inner
- *    offset (IO) of the last inner message. Then it computes RO and AO of previous messages:
- *
- *    RO = IO_of_a_message - IO_of_the_last_message
- *    AO = AO_Of_Last_Inner_Message + RO
- *
- * 4. This solution works for compacted message sets as well.
- *
- */
-class ByteBufferMessageSet(val buffer: ByteBuffer) extends MessageSet with Logging {
-
-  private[kafka] def this(compressionCodec: CompressionCodec,
-                          offsetCounter: LongRef,
-                          timestampType: TimestampType,
-                          messages: Message*) {
-    this(ByteBufferMessageSet.create(OffsetAssigner(offsetCounter, messages.size), compressionCodec,
-      timestampType, messages:_*))
-  }
-
-  def this(compressionCodec: CompressionCodec, offsetCounter: LongRef, messages: Message*) {
-    this(compressionCodec, offsetCounter, TimestampType.CREATE_TIME, messages:_*)
-  }
-
-  def this(compressionCodec: CompressionCodec, offsetSeq: Seq[Long], messages: Message*) {
-    this(ByteBufferMessageSet.create(new OffsetAssigner(offsetSeq), compressionCodec,
-      TimestampType.CREATE_TIME, messages:_*))
-  }
-
-  def this(compressionCodec: CompressionCodec, messages: Message*) {
-    this(compressionCodec, new LongRef(0L), messages: _*)
-  }
-
-  def this(messages: Message*) {
-    this(NoCompressionCodec, messages: _*)
-  }
-
-  def getBuffer = buffer
-
-  override def asRecords: MemoryRecords = MemoryRecords.readableRecords(buffer.duplicate())
-
-  /** default iterator that iterates over decompressed messages */
-  override def iterator: Iterator[MessageAndOffset] = internalIterator()
-
-  /** iterator over compressed messages without decompressing */
-  def shallowIterator: Iterator[MessageAndOffset] = internalIterator(isShallow = true)
-
-  /** When flag isShallow is set to be true, we do a shallow iteration: just traverse the first level of messages. **/
-  private def internalIterator(isShallow: Boolean = false): Iterator[MessageAndOffset] = {
-    if (isShallow)
-      asRecords.batches.asScala.iterator.map(MessageAndOffset.fromRecordBatch)
-    else
-      asRecords.records.asScala.iterator.map(MessageAndOffset.fromRecord)
-  }
-
-  /**
-   * The total number of bytes in this message set, including any partial trailing messages
-   */
-  def sizeInBytes: Int = buffer.limit()
-
-  /**
-   * The total number of bytes in this message set not including any partial, trailing messages
-   */
-  def validBytes: Int = asRecords.validBytes
-
-  /**
-   * Two message sets are equal if their respective byte buffers are equal
-   */
-  override def equals(other: Any): Boolean = {
-    other match {
-      case that: ByteBufferMessageSet =>
-        buffer.equals(that.buffer)
-      case _ => false
-    }
-  }
-
-  override def hashCode: Int = buffer.hashCode
-
-}
diff --git a/core/src/main/scala/kafka/message/CompressionCodec.scala b/core/src/main/scala/kafka/message/CompressionCodec.scala
index a485271bcb72c..b174feaea42dd 100644
--- a/core/src/main/scala/kafka/message/CompressionCodec.scala
+++ b/core/src/main/scala/kafka/message/CompressionCodec.scala
@@ -19,6 +19,8 @@ package kafka.message
 
 import java.util.Locale
 
+import kafka.common.UnknownCodecException
+
 object CompressionCodec {
   def getCompressionCodec(codec: Int): CompressionCodec = {
     codec match {
@@ -26,7 +28,8 @@ object CompressionCodec {
       case GZIPCompressionCodec.codec => GZIPCompressionCodec
       case SnappyCompressionCodec.codec => SnappyCompressionCodec
       case LZ4CompressionCodec.codec => LZ4CompressionCodec
-      case _ => throw new kafka.common.UnknownCodecException("%d is an unknown compression codec".format(codec))
+      case ZStdCompressionCodec.codec => ZStdCompressionCodec
+      case _ => throw new UnknownCodecException("%d is an unknown compression codec".format(codec))
     }
   }
   def getCompressionCodec(name: String): CompressionCodec = {
@@ -35,6 +38,7 @@ object CompressionCodec {
       case GZIPCompressionCodec.name => GZIPCompressionCodec
       case SnappyCompressionCodec.name => SnappyCompressionCodec
       case LZ4CompressionCodec.name => LZ4CompressionCodec
+      case ZStdCompressionCodec.name => ZStdCompressionCodec
       case _ => throw new kafka.common.UnknownCodecException("%s is an unknown compression codec".format(name))
     }
   }
@@ -42,8 +46,8 @@ object CompressionCodec {
 
 object BrokerCompressionCodec {
 
-  val brokerCompressionCodecs = List(UncompressedCodec, SnappyCompressionCodec, LZ4CompressionCodec, GZIPCompressionCodec, ProducerCompressionCodec)
-  val brokerCompressionOptions = brokerCompressionCodecs.map(codec => codec.name)
+  val brokerCompressionCodecs = List(UncompressedCodec, ZStdCompressionCodec, LZ4CompressionCodec, SnappyCompressionCodec, GZIPCompressionCodec, ProducerCompressionCodec)
+  val brokerCompressionOptions: List[String] = brokerCompressionCodecs.map(codec => codec.name)
 
   def isValid(compressionType: String): Boolean = brokerCompressionOptions.contains(compressionType.toLowerCase(Locale.ROOT))
 
@@ -66,8 +70,8 @@ sealed trait CompressionCodec { def codec: Int; def name: String }
 sealed trait BrokerCompressionCodec { def name: String }
 
 case object DefaultCompressionCodec extends CompressionCodec with BrokerCompressionCodec {
-  val codec = GZIPCompressionCodec.codec
-  val name = GZIPCompressionCodec.name
+  val codec: Int = GZIPCompressionCodec.codec
+  val name: String = GZIPCompressionCodec.name
 }
 
 case object GZIPCompressionCodec extends CompressionCodec with BrokerCompressionCodec {
@@ -85,6 +89,11 @@ case object LZ4CompressionCodec extends CompressionCodec with BrokerCompressionC
   val name = "lz4"
 }
 
+case object ZStdCompressionCodec extends CompressionCodec with BrokerCompressionCodec {
+  val codec = 4
+  val name = "zstd"
+}
+
 case object NoCompressionCodec extends CompressionCodec with BrokerCompressionCodec {
   val codec = 0
   val name = "none"
diff --git a/core/src/main/scala/kafka/message/InvalidMessageException.scala b/core/src/main/scala/kafka/message/InvalidMessageException.scala
deleted file mode 100644
index ef83500aea5ae..0000000000000
--- a/core/src/main/scala/kafka/message/InvalidMessageException.scala
+++ /dev/null
@@ -1,33 +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.message
-
-import org.apache.kafka.common.errors.CorruptRecordException
-
-/**
- * Indicates that a message failed its checksum and is corrupt
- *
- * InvalidMessageException extends CorruptRecordException for temporary compatibility with the old Scala clients.
- * We want to update the server side code to use and catch the new CorruptRecordException.
- * Because ByteBufferMessageSet.scala and Message.scala are used in both server and client code having
- * InvalidMessageException extend CorruptRecordException allows us to change server code without affecting the client.
- */
-class InvalidMessageException(message: String, throwable: Throwable) extends CorruptRecordException(message, throwable) {
-  def this(message: String) = this(null, null)
-  def this() = this(null)
-}
diff --git a/core/src/main/scala/kafka/message/Message.scala b/core/src/main/scala/kafka/message/Message.scala
deleted file mode 100755
index a46990160b6a1..0000000000000
--- a/core/src/main/scala/kafka/message/Message.scala
+++ /dev/null
@@ -1,379 +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.message
-
-import java.nio._
-
-import org.apache.kafka.common.record.{CompressionType, LegacyRecord, TimestampType}
-
-import scala.math._
-import org.apache.kafka.common.utils.{ByteUtils, Crc32}
-
-/**
- * Constants related to messages
- */
-object Message {
-
-  /**
-   * The current offset and size for all the fixed-length fields
-   */
-  val CrcOffset = 0
-  val CrcLength = 4
-  val MagicOffset = CrcOffset + CrcLength
-  val MagicLength = 1
-  val AttributesOffset = MagicOffset + MagicLength
-  val AttributesLength = 1
-  // Only message format version 1 has the timestamp field.
-  val TimestampOffset = AttributesOffset + AttributesLength
-  val TimestampLength = 8
-  val KeySizeOffset_V0 = AttributesOffset + AttributesLength
-  val KeySizeOffset_V1 = TimestampOffset + TimestampLength
-  val KeySizeLength = 4
-  val KeyOffset_V0 = KeySizeOffset_V0 + KeySizeLength
-  val KeyOffset_V1 = KeySizeOffset_V1 + KeySizeLength
-  val ValueSizeLength = 4
-
-  private val MessageHeaderSizeMap = Map (
-    (0: Byte) -> (CrcLength + MagicLength + AttributesLength + KeySizeLength + ValueSizeLength),
-    (1: Byte) -> (CrcLength + MagicLength + AttributesLength + TimestampLength + KeySizeLength + ValueSizeLength))
-
-  /**
-   * The amount of overhead bytes in a message
-   * This value is only used to check if the message size is valid or not. So the minimum possible message bytes is
-   * used here, which comes from a message in message format V0 with empty key and value.
-   */
-  val MinMessageOverhead = KeyOffset_V0 + ValueSizeLength
-  
-  /**
-   * The "magic" value
-   * When magic value is 0, the message uses absolute offset and does not have a timestamp field.
-   * When magic value is 1, the message uses relative offset and has a timestamp field.
-   */
-  val MagicValue_V0: Byte = 0
-  val MagicValue_V1: Byte = 1
-  val CurrentMagicValue: Byte = 1
-
-  /**
-   * Specifies the mask for the compression code. 3 bits to hold the compression codec.
-   * 0 is reserved to indicate no compression
-   */
-  val CompressionCodeMask: Int = 0x07
-  /**
-   * Specifies the mask for timestamp type. 1 bit at the 4th least significant bit.
-   * 0 for CreateTime, 1 for LogAppendTime
-   */
-  val TimestampTypeMask: Byte = 0x08
-  val TimestampTypeAttributeBitOffset: Int = 3
-
-  /**
-   * Compression code for uncompressed messages
-   */
-  val NoCompression: Int = 0
-
-  /**
-   * To indicate timestamp is not defined so "magic" value 0 will be used.
-   */
-  val NoTimestamp: Long = -1
-
-  /**
-   * Give the header size difference between different message versions.
-   */
-  def headerSizeDiff(fromMagicValue: Byte, toMagicValue: Byte) : Int =
-    MessageHeaderSizeMap(toMagicValue) - MessageHeaderSizeMap(fromMagicValue)
-
-
-  def fromRecord(record: LegacyRecord): Message = {
-    val wrapperTimestamp: Option[Long] = if (record.wrapperRecordTimestamp == null) None else Some(record.wrapperRecordTimestamp)
-    val wrapperTimestampType = Option(record.wrapperRecordTimestampType)
-    new Message(record.buffer, wrapperTimestamp, wrapperTimestampType)
-  }
-}
-
-/**
- * A message. The format of an N byte message is the following:
- *
- * 1. 4 byte CRC32 of the message
- * 2. 1 byte "magic" identifier to allow format changes, value is 0 or 1
- * 3. 1 byte "attributes" identifier to allow annotations on the message independent of the version
- *    bit 0 ~ 2 : Compression codec.
- *      0 : no compression
- *      1 : gzip
- *      2 : snappy
- *      3 : lz4
- *    bit 3 : Timestamp type
- *      0 : create time
- *      1 : log append time
- *    bit 4 ~ 7 : reserved
- * 4. (Optional) 8 byte timestamp only if "magic" identifier is greater than 0
- * 5. 4 byte key length, containing length K
- * 6. K byte key
- * 7. 4 byte payload length, containing length V
- * 8. V byte payload
- *
- * Default constructor wraps an existing ByteBuffer with the Message object with no change to the contents.
- * @param buffer the byte buffer of this message.
- * @param wrapperMessageTimestamp the wrapper message timestamp, which is only defined when the message is an inner
- *                                message of a compressed message.
- * @param wrapperMessageTimestampType the wrapper message timestamp type, which is only defined when the message is an
- *                                    inner message of a compressed message.
- */
-class Message(val buffer: ByteBuffer,
-              private val wrapperMessageTimestamp: Option[Long] = None,
-              private val wrapperMessageTimestampType: Option[TimestampType] = None) {
-  
-  import kafka.message.Message._
-
-  private[message] def asRecord: LegacyRecord = wrapperMessageTimestamp match {
-    case None => new LegacyRecord(buffer)
-    case Some(timestamp) => new LegacyRecord(buffer, timestamp, wrapperMessageTimestampType.orNull)
-  }
-
-  /**
-   * A constructor to create a Message
-   * @param bytes The payload of the message
-   * @param key The key of the message (null, if none)
-   * @param timestamp The timestamp of the message.
-   * @param timestampType The timestamp type of the message.
-   * @param codec The compression codec used on the contents of the message (if any)
-   * @param payloadOffset The offset into the payload array used to extract payload
-   * @param payloadSize The size of the payload to use
-   * @param magicValue the magic value to use
-   */
-  def this(bytes: Array[Byte], 
-           key: Array[Byte],
-           timestamp: Long,
-           timestampType: TimestampType,
-           codec: CompressionCodec, 
-           payloadOffset: Int, 
-           payloadSize: Int,
-           magicValue: Byte) = {
-    this(ByteBuffer.allocate(Message.CrcLength +
-                             Message.MagicLength +
-                             Message.AttributesLength +
-                             (if (magicValue == Message.MagicValue_V0) 0
-                              else Message.TimestampLength) +
-                             Message.KeySizeLength + 
-                             (if(key == null) 0 else key.length) + 
-                             Message.ValueSizeLength + 
-                             (if(bytes == null) 0 
-                              else if(payloadSize >= 0) payloadSize 
-                              else bytes.length - payloadOffset)))
-    validateTimestampAndMagicValue(timestamp, magicValue)
-    // skip crc, we will fill that in at the end
-    buffer.position(MagicOffset)
-    buffer.put(magicValue)
-    val attributes: Byte = LegacyRecord.computeAttributes(magicValue, CompressionType.forId(codec.codec), timestampType)
-    buffer.put(attributes)
-    // Only put timestamp when "magic" value is greater than 0
-    if (magic > MagicValue_V0)
-      buffer.putLong(timestamp)
-    if(key == null) {
-      buffer.putInt(-1)
-    } else {
-      buffer.putInt(key.length)
-      buffer.put(key, 0, key.length)
-    }
-    val size = if(bytes == null) -1
-               else if(payloadSize >= 0) payloadSize 
-               else bytes.length - payloadOffset
-    buffer.putInt(size)
-    if(bytes != null)
-      buffer.put(bytes, payloadOffset, size)
-    buffer.rewind()
-
-    // now compute the checksum and fill it in
-    ByteUtils.writeUnsignedInt(buffer, CrcOffset, computeChecksum)
-  }
-
-  def this(bytes: Array[Byte], key: Array[Byte], timestamp: Long, timestampType: TimestampType, codec: CompressionCodec, magicValue: Byte) =
-    this(bytes = bytes, key = key, timestamp = timestamp, timestampType = timestampType, codec = codec, payloadOffset = 0, payloadSize = -1, magicValue = magicValue)
-
-  def this(bytes: Array[Byte], key: Array[Byte], timestamp: Long, codec: CompressionCodec, magicValue: Byte) =
-    this(bytes = bytes, key = key, timestamp = timestamp, timestampType = TimestampType.CREATE_TIME, codec = codec, payloadOffset = 0, payloadSize = -1, magicValue = magicValue)
-  
-  def this(bytes: Array[Byte], timestamp: Long, codec: CompressionCodec, magicValue: Byte) =
-    this(bytes = bytes, key = null, timestamp = timestamp, codec = codec, magicValue = magicValue)
-  
-  def this(bytes: Array[Byte], key: Array[Byte], timestamp: Long, magicValue: Byte) =
-    this(bytes = bytes, key = key, timestamp = timestamp, codec = NoCompressionCodec, magicValue = magicValue)
-    
-  def this(bytes: Array[Byte], timestamp: Long, magicValue: Byte) =
-    this(bytes = bytes, key = null, timestamp = timestamp, codec = NoCompressionCodec, magicValue = magicValue)
-
-  def this(bytes: Array[Byte]) =
-    this(bytes = bytes, key = null, timestamp = Message.NoTimestamp, codec = NoCompressionCodec, magicValue = Message.CurrentMagicValue)
-    
-  /**
-   * Compute the checksum of the message from the message contents
-   */
-  def computeChecksum: Long =
-    Crc32.crc32(buffer, MagicOffset, buffer.limit() - MagicOffset)
-  
-  /**
-   * Retrieve the previously computed CRC for this message
-   */
-  def checksum: Long = ByteUtils.readUnsignedInt(buffer, CrcOffset)
-  
-    /**
-   * Returns true if the crc stored with the message matches the crc computed off the message contents
-   */
-  def isValid: Boolean = checksum == computeChecksum
-  
-  /**
-   * Throw an InvalidMessageException if isValid is false for this message
-   */
-  def ensureValid() {
-    if(!isValid)
-      throw new InvalidMessageException(s"Message is corrupt (stored crc = ${checksum}, computed crc = ${computeChecksum})")
-  }
-  
-  /**
-   * The complete serialized size of this message in bytes (including crc, header attributes, etc)
-   */
-  def size: Int = buffer.limit()
-
-  /**
-   * The position where the key size is stored.
-   */
-  private def keySizeOffset = {
-    if (magic == MagicValue_V0) KeySizeOffset_V0
-    else KeySizeOffset_V1
-  }
-
-  /**
-   * The length of the key in bytes
-   */
-  def keySize: Int = buffer.getInt(keySizeOffset)
-  
-  /**
-   * Does the message have a key?
-   */
-  def hasKey: Boolean = keySize >= 0
-  
-  /**
-   * The position where the payload size is stored
-   */
-  private def payloadSizeOffset = {
-    if (magic == MagicValue_V0) KeyOffset_V0 + max(0, keySize)
-    else KeyOffset_V1 + max(0, keySize)
-  }
-  
-  /**
-   * The length of the message value in bytes
-   */
-  def payloadSize: Int = buffer.getInt(payloadSizeOffset)
-  
-  /**
-   * Is the payload of this message null
-   */
-  def isNull: Boolean = payloadSize < 0
-  
-  /**
-   * The magic version of this message
-   */
-  def magic: Byte = buffer.get(MagicOffset)
-  
-  /**
-   * The attributes stored with this message
-   */
-  def attributes: Byte = buffer.get(AttributesOffset)
-
-  /**
-   * The timestamp of the message, only available when the "magic" value is greater than 0
-   * When magic > 0, The timestamp of a message is determined in the following way:
-   * 1. wrapperMessageTimestampType = None and wrapperMessageTimestamp is None - Uncompressed message, timestamp and timestamp type are in the message.
-   * 2. wrapperMessageTimestampType = LogAppendTime and wrapperMessageTimestamp is defined - Compressed message using LogAppendTime
-   * 3. wrapperMessageTimestampType = CreateTime and wrapperMessageTimestamp is defined - Compressed message using CreateTime
-   */
-  def timestamp: Long = {
-    if (magic == MagicValue_V0)
-      Message.NoTimestamp
-    // Case 2
-    else if (wrapperMessageTimestampType.exists(_ == TimestampType.LOG_APPEND_TIME) && wrapperMessageTimestamp.isDefined)
-      wrapperMessageTimestamp.get
-    else // case 1, 3
-      buffer.getLong(Message.TimestampOffset)
-  }
-
-  /**
-   * The timestamp type of the message
-   */
-  def timestampType = LegacyRecord.timestampType(magic, wrapperMessageTimestampType.orNull, attributes)
-
-  /**
-   * The compression codec used with this message
-   */
-  def compressionCodec: CompressionCodec = 
-    CompressionCodec.getCompressionCodec(buffer.get(AttributesOffset) & CompressionCodeMask)
-  
-  /**
-   * A ByteBuffer containing the content of the message
-   */
-  def payload: ByteBuffer = sliceDelimited(payloadSizeOffset)
-  
-  /**
-   * A ByteBuffer containing the message key
-   */
-  def key: ByteBuffer = sliceDelimited(keySizeOffset)
-
-  /**
-   * Read a size-delimited byte buffer starting at the given offset
-   */
-  private def sliceDelimited(start: Int): ByteBuffer = {
-    val size = buffer.getInt(start)
-    if(size < 0) {
-      null
-    } else {
-      var b = buffer.duplicate()
-      b.position(start + 4)
-      b = b.slice()
-      b.limit(size)
-      b.rewind
-      b
-    }
-  }
-
-  /**
-   * Validate the timestamp and "magic" value
-   */
-  private def validateTimestampAndMagicValue(timestamp: Long, magic: Byte) {
-    if (magic != MagicValue_V0 && magic != MagicValue_V1)
-      throw new IllegalArgumentException(s"Invalid magic value $magic")
-    if (timestamp < 0 && timestamp != NoTimestamp)
-      throw new IllegalArgumentException(s"Invalid message timestamp $timestamp")
-    if (magic == MagicValue_V0 && timestamp != NoTimestamp)
-      throw new IllegalArgumentException(s"Invalid timestamp $timestamp. Timestamp must be $NoTimestamp when magic = $MagicValue_V0")
-  }
-
-  override def toString: String = {
-    if (magic == MagicValue_V0)
-      s"Message(magic = $magic, attributes = $attributes, crc = $checksum, key = $key, payload = $payload)"
-    else
-      s"Message(magic = $magic, attributes = $attributes, $timestampType = $timestamp, crc = $checksum, key = $key, payload = $payload)"
-  }
-
-  override def equals(any: Any): Boolean = {
-    any match {
-      case that: Message => this.buffer.equals(that.buffer)
-      case _ => false
-    }
-  }
-  
-  override def hashCode(): Int = buffer.hashCode
-
-}
diff --git a/core/src/main/scala/kafka/message/MessageAndMetadata.scala b/core/src/main/scala/kafka/message/MessageAndMetadata.scala
deleted file mode 100755
index 5c09cafdbac34..0000000000000
--- a/core/src/main/scala/kafka/message/MessageAndMetadata.scala
+++ /dev/null
@@ -1,39 +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.message
-
-import kafka.serializer.Decoder
-import org.apache.kafka.common.record.TimestampType
-import org.apache.kafka.common.utils.Utils
-
-case class MessageAndMetadata[K, V](topic: String,
-                                    partition: Int,
-                                    private val rawMessage: Message,
-                                    offset: Long,
-                                    keyDecoder: Decoder[K], valueDecoder: Decoder[V],
-                                    timestamp: Long = Message.NoTimestamp,
-                                    timestampType: TimestampType = TimestampType.CREATE_TIME) {
-
-  /**
-   * Return the decoded message key and payload
-   */
-  def key(): K = if(rawMessage.key == null) null.asInstanceOf[K] else keyDecoder.fromBytes(Utils.readBytes(rawMessage.key))
-
-  def message(): V = if(rawMessage.isNull) null.asInstanceOf[V] else valueDecoder.fromBytes(Utils.readBytes(rawMessage.payload))
-}
-
diff --git a/core/src/main/scala/kafka/message/MessageAndOffset.scala b/core/src/main/scala/kafka/message/MessageAndOffset.scala
deleted file mode 100644
index 8de0f81bdfc61..0000000000000
--- a/core/src/main/scala/kafka/message/MessageAndOffset.scala
+++ /dev/null
@@ -1,54 +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.message
-
-import org.apache.kafka.common.record.{AbstractLegacyRecordBatch, Record, RecordBatch}
-
-object MessageAndOffset {
-  def fromRecordBatch(batch: RecordBatch): MessageAndOffset = {
-    batch match {
-      case legacyBatch: AbstractLegacyRecordBatch =>
-        MessageAndOffset(Message.fromRecord(legacyBatch.outerRecord), legacyBatch.lastOffset)
-
-      case _ =>
-        throw new IllegalArgumentException(s"Illegal batch type ${batch.getClass}. The older message format classes " +
-          s"only support conversion from ${classOf[AbstractLegacyRecordBatch]}, which is used for magic v0 and v1")
-    }
-  }
-
-  def fromRecord(record: Record): MessageAndOffset = {
-    record match {
-      case legacyBatch: AbstractLegacyRecordBatch =>
-        MessageAndOffset(Message.fromRecord(legacyBatch.outerRecord), legacyBatch.lastOffset)
-
-      case _ =>
-        throw new IllegalArgumentException(s"Illegal record type ${record.getClass}. The older message format classes " +
-          s"only support conversion from ${classOf[AbstractLegacyRecordBatch]}, which is used for magic v0 and v1")
-    }
-  }
-}
-
-case class MessageAndOffset(message: Message, offset: Long) {
-  
-  /**
-   * Compute the offset of the next message in the log
-   */
-  def nextOffset: Long = offset + 1
-
-}
-
diff --git a/core/src/main/scala/kafka/message/MessageLengthException.scala b/core/src/main/scala/kafka/message/MessageLengthException.scala
deleted file mode 100644
index 45d32a5d4aed4..0000000000000
--- a/core/src/main/scala/kafka/message/MessageLengthException.scala
+++ /dev/null
@@ -1,24 +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.message
-
-/**
- * Indicates the presence of a message that exceeds the maximum acceptable 
- * length (whatever that happens to be)
- */
-class MessageLengthException(message: String) extends RuntimeException(message)
diff --git a/core/src/main/scala/kafka/message/MessageSet.scala b/core/src/main/scala/kafka/message/MessageSet.scala
deleted file mode 100644
index 915def081a212..0000000000000
--- a/core/src/main/scala/kafka/message/MessageSet.scala
+++ /dev/null
@@ -1,111 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- * 
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package kafka.message
-
-import java.nio._
-
-import org.apache.kafka.common.record.Records
-
-/**
- * Message set helper functions
- */
-object MessageSet {
-
-  val MessageSizeLength = 4
-  val OffsetLength = 8
-  val LogOverhead = MessageSizeLength + OffsetLength
-  val Empty = new ByteBufferMessageSet(ByteBuffer.allocate(0))
-  
-  /**
-   * The size of a message set containing the given messages
-   */
-  def messageSetSize(messages: Iterable[Message]): Int =
-    messages.foldLeft(0)(_ + entrySize(_))
-
-  /**
-   * The size of a size-delimited entry in a message set
-   */
-  def entrySize(message: Message): Int = LogOverhead + message.size
-
-  /**
-   * Validate that all "magic" values in `messages` are the same and return their magic value and max timestamp
-   */
-  def magicAndLargestTimestamp(messages: Seq[Message]): MagicAndTimestamp = {
-    val firstMagicValue = messages.head.magic
-    var largestTimestamp = Message.NoTimestamp
-    for (message <- messages) {
-      if (message.magic != firstMagicValue)
-        throw new IllegalStateException("Messages in the same message set must have same magic value")
-      if (firstMagicValue > Message.MagicValue_V0)
-        largestTimestamp = math.max(largestTimestamp, message.timestamp)
-    }
-    MagicAndTimestamp(firstMagicValue, largestTimestamp)
-  }
-
-}
-
-case class MagicAndTimestamp(magic: Byte, timestamp: Long)
-
-/**
- * A set of messages with offsets. A message set has a fixed serialized form, though the container
- * for the bytes could be either in-memory or on disk. The format of each message is
- * as follows:
- * 8 byte message offset number
- * 4 byte size containing an integer N
- * N message bytes as described in the Message class
- */
-abstract class MessageSet extends Iterable[MessageAndOffset] {
-
-  /**
-   * Provides an iterator over the message/offset pairs in this set
-   */
-  def iterator: Iterator[MessageAndOffset]
-  
-  /**
-   * Gives the total size of this message set in bytes
-   */
-  def sizeInBytes: Int
-
-  /**
-   * Get the client representation of the message set
-   */
-  def asRecords: Records
-
-  /**
-   * Print this message set's contents. If the message set has more than 100 messages, just
-   * print the first 100.
-   */
-  override def toString: String = {
-    val builder = new StringBuilder()
-    builder.append(getClass.getSimpleName + "(")
-    val iter = this.asRecords.batches.iterator
-    var i = 0
-    while(iter.hasNext && i < 100) {
-      val message = iter.next
-      builder.append(message)
-      if(iter.hasNext)
-        builder.append(", ")
-      i += 1
-    }
-    if(iter.hasNext)
-      builder.append("...")
-    builder.append(")")
-    builder.toString
-  }
-
-}
diff --git a/core/src/main/scala/kafka/message/package.html b/core/src/main/scala/kafka/message/package.html
deleted file mode 100644
index c666d81d374bc..0000000000000
--- a/core/src/main/scala/kafka/message/package.html
+++ /dev/null
@@ -1,19 +0,0 @@
-
-Messages and everything related to them.
\ No newline at end of file
diff --git a/core/src/main/scala/kafka/metrics/KafkaCSVMetricsReporter.scala b/core/src/main/scala/kafka/metrics/KafkaCSVMetricsReporter.scala
index 81c20a7e117bc..0d8354728ac8d 100755
--- a/core/src/main/scala/kafka/metrics/KafkaCSVMetricsReporter.scala
+++ b/core/src/main/scala/kafka/metrics/KafkaCSVMetricsReporter.scala
@@ -20,7 +20,6 @@
 
 package kafka.metrics
 
-import com.yammer.metrics.Metrics
 import java.io.File
 import java.nio.file.Files
 
@@ -45,14 +44,14 @@ 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)
         csvDir = new File(props.getString("kafka.csv.metrics.dir", "kafka_metrics"))
         Utils.delete(csvDir)
         Files.createDirectories(csvDir.toPath())
-        underlying = new CsvReporter(Metrics.defaultRegistry(), csvDir)
+        underlying = new CsvReporter(KafkaYammerMetrics.defaultRegistry(), csvDir)
         if (props.getBoolean("kafka.csv.metrics.reporter.enabled", default = false)) {
           initialized = true
           startReporter(metricsConfig.pollingIntervalSecs)
@@ -62,7 +61,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,13 +72,13 @@ private class KafkaCSVMetricsReporter extends KafkaMetricsReporter
   }
 
 
-  override def stopReporter() {
+  override def stopReporter(): Unit = {
     synchronized {
       if (initialized && running) {
         underlying.shutdown()
         running = false
         info("Stopped Kafka CSV metrics reporter")
-        underlying = new CsvReporter(Metrics.defaultRegistry(), csvDir)
+        underlying = new CsvReporter(KafkaYammerMetrics.defaultRegistry(), csvDir)
       }
     }
   }
diff --git a/core/src/main/scala/kafka/metrics/KafkaMetricsConfig.scala b/core/src/main/scala/kafka/metrics/KafkaMetricsConfig.scala
index ad9eb20f3e72d..b13a1b9350fbb 100755
--- a/core/src/main/scala/kafka/metrics/KafkaMetricsConfig.scala
+++ b/core/src/main/scala/kafka/metrics/KafkaMetricsConfig.scala
@@ -20,7 +20,9 @@
 
 package kafka.metrics
 
-import kafka.utils.{VerifiableProperties, CoreUtils}
+import kafka.server.{Defaults, KafkaConfig}
+import kafka.utils.{CoreUtils, VerifiableProperties}
+import scala.collection.Seq
 
 class KafkaMetricsConfig(props: VerifiableProperties) {
 
@@ -28,10 +30,12 @@ class KafkaMetricsConfig(props: VerifiableProperties) {
    * Comma-separated list of reporter types. These classes should be on the
    * classpath and will be instantiated at run-time.
    */
-  val reporters = CoreUtils.parseCsvList(props.getString("kafka.metrics.reporters", ""))
+  val reporters: Seq[String] = CoreUtils.parseCsvList(props.getString(KafkaConfig.KafkaMetricsReporterClassesProp,
+    Defaults.KafkaMetricReporterClasses))
 
   /**
    * The metrics polling interval (in seconds).
    */
-  val pollingIntervalSecs = props.getInt("kafka.metrics.polling.interval.secs", 10)
+  val pollingIntervalSecs: Int = props.getInt(KafkaConfig.KafkaMetricsPollingIntervalSecondsProp,
+    Defaults.KafkaMetricsPollingIntervalSeconds)
 }
diff --git a/core/src/main/scala/kafka/metrics/KafkaMetricsGroup.scala b/core/src/main/scala/kafka/metrics/KafkaMetricsGroup.scala
index 75eaafd720a05..a63be1fe2693c 100644
--- a/core/src/main/scala/kafka/metrics/KafkaMetricsGroup.scala
+++ b/core/src/main/scala/kafka/metrics/KafkaMetricsGroup.scala
@@ -19,17 +19,10 @@ package kafka.metrics
 
 import java.util.concurrent.TimeUnit
 
-import com.yammer.metrics.Metrics
-import com.yammer.metrics.core.{Gauge, MetricName}
-import kafka.consumer.{ConsumerTopicStatsRegistry, FetchRequestAndResponseStatsRegistry}
-import kafka.producer.{ProducerRequestStatsRegistry, ProducerStatsRegistry, ProducerTopicStatsRegistry}
+import com.yammer.metrics.core.{Gauge, MetricName, Meter, Histogram, Timer}
 import kafka.utils.Logging
 import org.apache.kafka.common.utils.Sanitizer
 
-import scala.collection.immutable
-import scala.collection.JavaConverters._
-
-
 trait KafkaMetricsGroup extends Logging {
 
   /**
@@ -64,93 +57,27 @@ trait KafkaMetricsGroup extends Logging {
       nameBuilder.append(name)
     }
 
-    val scope: String = KafkaMetricsGroup.toScope(tags).getOrElse(null)
-    val tagsName = KafkaMetricsGroup.toMBeanName(tags)
+    val scope: String = toScope(tags).orNull
+    val tagsName = toMBeanName(tags)
     tagsName.foreach(nameBuilder.append(",").append(_))
 
     new MetricName(group, typeName, name, scope, nameBuilder.toString)
   }
 
-  def newGauge[T](name: String, metric: Gauge[T], tags: scala.collection.Map[String, String] = Map.empty) =
-    Metrics.defaultRegistry().newGauge(metricName(name, tags), metric)
-
-  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 newHistogram(name: String, biased: Boolean = true, tags: scala.collection.Map[String, String] = Map.empty) =
-    Metrics.defaultRegistry().newHistogram(metricName(name, tags), biased)
+  def newGauge[T](name: String, metric: Gauge[T], tags: scala.collection.Map[String, String] = Map.empty): Gauge[T] =
+    KafkaYammerMetrics.defaultRegistry().newGauge(metricName(name, tags), metric)
 
-  def newTimer(name: String, durationUnit: TimeUnit, rateUnit: TimeUnit, tags: scala.collection.Map[String, String] = Map.empty) =
-    Metrics.defaultRegistry().newTimer(metricName(name, tags), durationUnit, rateUnit)
+  def newMeter(name: String, eventType: String, timeUnit: TimeUnit, tags: scala.collection.Map[String, String] = Map.empty): Meter =
+    KafkaYammerMetrics.defaultRegistry().newMeter(metricName(name, tags), eventType, timeUnit)
 
-  def removeMetric(name: String, tags: scala.collection.Map[String, String] = Map.empty) =
-    Metrics.defaultRegistry().removeMetric(metricName(name, tags))
+  def newHistogram(name: String, biased: Boolean = true, tags: scala.collection.Map[String, String] = Map.empty): Histogram =
+    KafkaYammerMetrics.defaultRegistry().newHistogram(metricName(name, tags), biased)
 
+  def newTimer(name: String, durationUnit: TimeUnit, rateUnit: TimeUnit, tags: scala.collection.Map[String, String] = Map.empty): Timer =
+    KafkaYammerMetrics.defaultRegistry().newTimer(metricName(name, tags), durationUnit, rateUnit)
 
-}
-
-object KafkaMetricsGroup extends KafkaMetricsGroup with Logging {
-  /**
-   * To make sure all the metrics be de-registered after consumer/producer close, the metric names should be
-   * put into the metric name set.
-   */
-  private val consumerMetricNameList: immutable.List[MetricName] = immutable.List[MetricName](
-    // kafka.consumer.ZookeeperConsumerConnector
-    new MetricName("kafka.consumer", "ZookeeperConsumerConnector", "FetchQueueSize"),
-    new MetricName("kafka.consumer", "ZookeeperConsumerConnector", "KafkaCommitsPerSec"),
-    new MetricName("kafka.consumer", "ZookeeperConsumerConnector", "ZooKeeperCommitsPerSec"),
-    new MetricName("kafka.consumer", "ZookeeperConsumerConnector", "RebalanceRateAndTime"),
-    new MetricName("kafka.consumer", "ZookeeperConsumerConnector", "OwnedPartitionsCount"),
-
-    // kafka.consumer.ConsumerFetcherManager
-    new MetricName("kafka.consumer", "ConsumerFetcherManager", "MaxLag"),
-    new MetricName("kafka.consumer", "ConsumerFetcherManager", "MinFetchRate"),
-
-    // kafka.server.AbstractFetcherThread <-- kafka.consumer.ConsumerFetcherThread
-    new MetricName("kafka.server", "FetcherLagMetrics", "ConsumerLag"),
-
-    // kafka.consumer.ConsumerTopicStats <-- kafka.consumer.{ConsumerIterator, PartitionTopicInfo}
-    new MetricName("kafka.consumer", "ConsumerTopicMetrics", "MessagesPerSec"),
-
-    // kafka.consumer.ConsumerTopicStats
-    new MetricName("kafka.consumer", "ConsumerTopicMetrics", "BytesPerSec"),
-
-    // kafka.server.AbstractFetcherThread <-- kafka.consumer.ConsumerFetcherThread
-    new MetricName("kafka.server", "FetcherStats", "BytesPerSec"),
-    new MetricName("kafka.server", "FetcherStats", "RequestsPerSec"),
-
-    // kafka.consumer.FetchRequestAndResponseStats <-- kafka.consumer.SimpleConsumer
-    new MetricName("kafka.consumer", "FetchRequestAndResponseMetrics", "FetchResponseSize"),
-    new MetricName("kafka.consumer", "FetchRequestAndResponseMetrics", "FetchRequestRateAndTimeMs"),
-    new MetricName("kafka.consumer", "FetchRequestAndResponseMetrics", "FetchRequestThrottleRateAndTimeMs"),
-
-    /**
-     * ProducerRequestStats <-- SyncProducer
-     * metric for SyncProducer in fetchTopicMetaData() needs to be removed when consumer is closed.
-     */
-    new MetricName("kafka.producer", "ProducerRequestMetrics", "ProducerRequestRateAndTimeMs"),
-    new MetricName("kafka.producer", "ProducerRequestMetrics", "ProducerRequestSize")
-  )
-
-  private val producerMetricNameList: immutable.List[MetricName] = immutable.List[MetricName](
-    // kafka.producer.ProducerStats <-- DefaultEventHandler <-- Producer
-    new MetricName("kafka.producer", "ProducerStats", "SerializationErrorsPerSec"),
-    new MetricName("kafka.producer", "ProducerStats", "ResendsPerSec"),
-    new MetricName("kafka.producer", "ProducerStats", "FailedSendsPerSec"),
-
-    // kafka.producer.ProducerSendThread
-    new MetricName("kafka.producer.async", "ProducerSendThread", "ProducerQueueSize"),
-
-    // kafka.producer.ProducerTopicStats <-- kafka.producer.{Producer, async.DefaultEventHandler}
-    new MetricName("kafka.producer", "ProducerTopicMetrics", "MessagesPerSec"),
-    new MetricName("kafka.producer", "ProducerTopicMetrics", "DroppedMessagesPerSec"),
-    new MetricName("kafka.producer", "ProducerTopicMetrics", "BytesPerSec"),
-
-    // kafka.producer.ProducerRequestStats <-- SyncProducer
-    new MetricName("kafka.producer", "ProducerRequestMetrics", "ProducerRequestRateAndTimeMs"),
-    new MetricName("kafka.producer", "ProducerRequestMetrics", "ProducerRequestSize"),
-    new MetricName("kafka.producer", "ProducerRequestMetrics", "ProducerRequestThrottleRateAndTimeMs")
-  )
+  def removeMetric(name: String, tags: scala.collection.Map[String, String] = Map.empty): Unit =
+    KafkaYammerMetrics.defaultRegistry().removeMetric(metricName(name, tags))
 
   private def toMBeanName(tags: collection.Map[String, String]): Option[String] = {
     val filteredTags = tags.filter { case (_, tagValue) => tagValue != "" }
@@ -175,42 +102,6 @@ object KafkaMetricsGroup extends KafkaMetricsGroup with Logging {
     else None
   }
 
-  @deprecated("This method has been deprecated and will be removed in a future release.", "0.11.0.0")
-  def removeAllConsumerMetrics(clientId: String) {
-    FetchRequestAndResponseStatsRegistry.removeConsumerFetchRequestAndResponseStats(clientId)
-    ConsumerTopicStatsRegistry.removeConsumerTopicStat(clientId)
-    ProducerRequestStatsRegistry.removeProducerRequestStats(clientId)
-    removeAllMetricsInList(KafkaMetricsGroup.consumerMetricNameList, clientId)
-  }
-
-  @deprecated("This method has been deprecated and will be removed in a future release.", "0.10.0.0")
-  def removeAllProducerMetrics(clientId: String) {
-    ProducerRequestStatsRegistry.removeProducerRequestStats(clientId)
-    ProducerTopicStatsRegistry.removeProducerTopicStats(clientId)
-    ProducerStatsRegistry.removeProducerStats(clientId)
-    removeAllMetricsInList(KafkaMetricsGroup.producerMetricNameList, clientId)
-  }
-
-  private def removeAllMetricsInList(metricNameList: immutable.List[MetricName], clientId: String) {
-    metricNameList.foreach(metric => {
-      val pattern = (".*clientId=" + clientId + ".*").r
-      val registeredMetrics = Metrics.defaultRegistry().allMetrics().keySet().asScala
-      for (registeredMetric <- registeredMetrics) {
-        if (registeredMetric.getGroup == metric.getGroup &&
-          registeredMetric.getName == metric.getName &&
-          registeredMetric.getType == metric.getType) {
-          pattern.findFirstIn(registeredMetric.getMBeanName) match {
-            case Some(_) => {
-              val beforeRemovalSize = Metrics.defaultRegistry().allMetrics().keySet().size
-              Metrics.defaultRegistry().removeMetric(registeredMetric)
-              val afterRemovalSize = Metrics.defaultRegistry().allMetrics().keySet().size
-              trace("Removing metric %s. Metrics registry size reduced from %d to %d".format(
-                registeredMetric, beforeRemovalSize, afterRemovalSize))
-            }
-            case _ =>
-          }
-        }
-      }
-    })
-  }
 }
+
+object KafkaMetricsGroup extends KafkaMetricsGroup
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/metrics/LinuxIoMetricsCollector.scala b/core/src/main/scala/kafka/metrics/LinuxIoMetricsCollector.scala
new file mode 100644
index 0000000000000..17de008580aec
--- /dev/null
+++ b/core/src/main/scala/kafka/metrics/LinuxIoMetricsCollector.scala
@@ -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 kafka.metrics
+
+import java.nio.file.{Files, Paths}
+
+import org.apache.kafka.common.utils.Time
+import org.slf4j.Logger
+
+import scala.jdk.CollectionConverters._
+
+/**
+ * Retrieves Linux /proc/self/io metrics.
+ */
+class LinuxIoMetricsCollector(procRoot: String, val time: Time, val logger: Logger) {
+  import LinuxIoMetricsCollector._
+  var lastUpdateMs = -1L
+  var cachedReadBytes = 0L
+  var cachedWriteBytes = 0L
+  val path = Paths.get(procRoot, "self", "io")
+
+  def readBytes(): Long = this.synchronized {
+    val curMs = time.milliseconds()
+    if (curMs != lastUpdateMs) {
+      updateValues(curMs)
+    }
+    cachedReadBytes
+  }
+
+  def writeBytes(): Long = this.synchronized {
+    val curMs = time.milliseconds()
+    if (curMs != lastUpdateMs) {
+      updateValues(curMs)
+    }
+    cachedWriteBytes
+  }
+
+  /**
+   * Read /proc/self/io.
+   *
+   * Generally, each line in this file contains a prefix followed by a colon and a number.
+   *
+   * For example, it might contain this:
+   * rchar: 4052
+   * wchar: 0
+   * syscr: 13
+   * syscw: 0
+   * read_bytes: 0
+   * write_bytes: 0
+   * cancelled_write_bytes: 0
+   */
+  def updateValues(now: Long): Boolean = this.synchronized {
+    try {
+      cachedReadBytes = -1
+      cachedWriteBytes = -1
+      val lines = Files.readAllLines(path).asScala
+      lines.foreach(line => {
+        if (line.startsWith(READ_BYTES_PREFIX)) {
+          cachedReadBytes = line.substring(READ_BYTES_PREFIX.size).toLong
+        } else if (line.startsWith(WRITE_BYTES_PREFIX)) {
+          cachedWriteBytes = line.substring(WRITE_BYTES_PREFIX.size).toLong
+        }
+      })
+      lastUpdateMs = now
+      true
+    } catch {
+      case t: Throwable => {
+        logger.warn("Unable to update IO metrics", t)
+        false
+      }
+    }
+  }
+
+  def usable(): Boolean = {
+    if (path.toFile().exists()) {
+      updateValues(time.milliseconds())
+    } else {
+      logger.debug(s"disabling IO metrics collection because ${path} does not exist.")
+      false
+    }
+  }
+}
+
+object LinuxIoMetricsCollector {
+  val READ_BYTES_PREFIX = "read_bytes: "
+  val WRITE_BYTES_PREFIX = "write_bytes: "
+}
diff --git a/core/src/main/scala/kafka/network/BlockingChannel.scala b/core/src/main/scala/kafka/network/BlockingChannel.scala
deleted file mode 100644
index 3493ad34e1852..0000000000000
--- a/core/src/main/scala/kafka/network/BlockingChannel.scala
+++ /dev/null
@@ -1,135 +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.network
-
-import java.net.InetSocketAddress
-import java.nio.channels._
-
-import kafka.api.RequestOrResponse
-import kafka.utils.{CoreUtils, Logging, nonthreadsafe}
-import org.apache.kafka.common.network.NetworkReceive
-
-
-@deprecated("This object has been deprecated and will be removed in a future release.", "0.11.0.0")
-object BlockingChannel{
-  val UseDefaultBufferSize = -1
-}
-
-/**
- *  A simple blocking channel with timeouts correctly enabled.
- *
- */
-@nonthreadsafe
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.11.0.0")
-class BlockingChannel( val host: String, 
-                       val port: Int, 
-                       val readBufferSize: Int, 
-                       val writeBufferSize: Int, 
-                       val readTimeoutMs: Int ) extends Logging {
-  private var connected = false
-  private var channel: SocketChannel = null
-  private var readChannel: ReadableByteChannel = null
-  private var writeChannel: GatheringByteChannel = null
-  private val lock = new Object()
-  private val connectTimeoutMs = readTimeoutMs
-  private var connectionId: String = ""
-
-  def connect() = lock synchronized  {
-    if(!connected) {
-      try {
-        channel = SocketChannel.open()
-        if(readBufferSize > 0)
-          channel.socket.setReceiveBufferSize(readBufferSize)
-        if(writeBufferSize > 0)
-          channel.socket.setSendBufferSize(writeBufferSize)
-        channel.configureBlocking(true)
-        channel.socket.setSoTimeout(readTimeoutMs)
-        channel.socket.setKeepAlive(true)
-        channel.socket.setTcpNoDelay(true)
-        channel.socket.connect(new InetSocketAddress(host, port), connectTimeoutMs)
-
-        writeChannel = channel
-        // Need to create a new ReadableByteChannel from input stream because SocketChannel doesn't implement read with timeout
-        // See: http://stackoverflow.com/questions/2866557/timeout-for-socketchannel-doesnt-work
-        readChannel = Channels.newChannel(channel.socket().getInputStream)
-        connected = true
-        val localHost = channel.socket.getLocalAddress.getHostAddress
-        val localPort = channel.socket.getLocalPort
-        val remoteHost = channel.socket.getInetAddress.getHostAddress
-        val remotePort = channel.socket.getPort
-        connectionId = localHost + ":" + localPort + "-" + remoteHost + ":" + remotePort
-        // settings may not match what we requested above
-        val msg = "Created socket with SO_TIMEOUT = %d (requested %d), SO_RCVBUF = %d (requested %d), SO_SNDBUF = %d (requested %d), connectTimeoutMs = %d."
-        debug(msg.format(channel.socket.getSoTimeout,
-                         readTimeoutMs,
-                         channel.socket.getReceiveBufferSize, 
-                         readBufferSize,
-                         channel.socket.getSendBufferSize,
-                         writeBufferSize,
-                         connectTimeoutMs))
-
-      } catch {
-        case _: Throwable => disconnect()
-      }
-    }
-  }
-  
-  def disconnect() = lock synchronized {
-    if(channel != null) {
-      CoreUtils.swallow(channel.close(), this)
-      CoreUtils.swallow(channel.socket.close(), this)
-      channel = null
-      writeChannel = null
-    }
-    // closing the main socket channel *should* close the read channel
-    // but let's do it to be sure.
-    if(readChannel != null) {
-      CoreUtils.swallow(readChannel.close(), this)
-      readChannel = null
-    }
-    connected = false
-  }
-
-  def isConnected = connected
-
-  def send(request: RequestOrResponse): Long = {
-    if(!connected)
-      throw new ClosedChannelException()
-
-    val send = new RequestOrResponseSend(connectionId, request)
-    send.writeCompletely(writeChannel)
-  }
-  
-  def receive(): NetworkReceive = {
-    if(!connected)
-      throw new ClosedChannelException()
-
-    val response = readCompletely(readChannel)
-    response.payload().rewind()
-
-    response
-  }
-
-  private def readCompletely(channel: ReadableByteChannel): NetworkReceive = {
-    val response = new NetworkReceive
-    while (!response.complete())
-      response.readFromReadableChannel(channel)
-    response
-  }
-
-}
diff --git a/core/src/main/scala/kafka/network/RequestChannel.scala b/core/src/main/scala/kafka/network/RequestChannel.scala
index 7cc861957badd..0c03d70188165 100644
--- a/core/src/main/scala/kafka/network/RequestChannel.scala
+++ b/core/src/main/scala/kafka/network/RequestChannel.scala
@@ -21,38 +21,51 @@ import java.net.InetAddress
 import java.nio.ByteBuffer
 import java.util.concurrent._
 
+import com.fasterxml.jackson.databind.JsonNode
 import com.typesafe.scalalogging.Logger
-import com.yammer.metrics.core.{Gauge, Meter}
+import com.yammer.metrics.core.Meter
+import kafka.log.LogConfig
 import kafka.metrics.KafkaMetricsGroup
-import kafka.network.RequestChannel.{BaseRequest, SendAction, ShutdownRequest, NoOpAction, CloseConnectionAction}
-import kafka.utils.{Logging, NotNothing}
+import kafka.server.KafkaConfig
+import kafka.utils.{Logging, NotNothing, Pool}
+import kafka.utils.Implicits._
+import org.apache.kafka.common.config.types.Password
+import org.apache.kafka.common.config.ConfigResource
 import org.apache.kafka.common.memory.MemoryPool
+import org.apache.kafka.common.message.IncrementalAlterConfigsRequestData
+import org.apache.kafka.common.message.IncrementalAlterConfigsRequestData._
 import org.apache.kafka.common.network.Send
-import org.apache.kafka.common.protocol.{ApiKeys, Errors}
+import org.apache.kafka.common.protocol.{ApiKeys, Errors, ObjectSerializationCache}
 import org.apache.kafka.common.requests._
 import org.apache.kafka.common.security.auth.KafkaPrincipal
 import org.apache.kafka.common.utils.{Sanitizer, Time}
 
+import scala.annotation.nowarn
 import scala.collection.mutable
+import scala.jdk.CollectionConverters._
 import scala.reflect.ClassTag
 
 object RequestChannel extends Logging {
   private val requestLogger = Logger("kafka.request.logger")
 
+  val RequestQueueSizeMetric = "RequestQueueSize"
+  val ResponseQueueSizeMetric = "ResponseQueueSize"
+  val ProcessorMetricTag = "processor"
+
   def isRequestLoggingEnabled: Boolean = requestLogger.underlying.isDebugEnabled
 
   sealed trait BaseRequest
   case object ShutdownRequest extends BaseRequest
 
   case class Session(principal: KafkaPrincipal, clientAddress: InetAddress) {
-    val sanitizedUser = Sanitizer.sanitize(principal.getName)
+    val sanitizedUser: String = Sanitizer.sanitize(principal.getName)
   }
 
-  class Metrics {
+  class Metrics(allowDisabledApis: Boolean = false) {
 
     private val metricsMap = mutable.Map[String, RequestMetrics]()
 
-    (ApiKeys.values.toSeq.map(_.name) ++
+    (ApiKeys.values.toSeq.filter(_.isEnabled || allowDisabledApis).map(_.name) ++
         Seq(RequestMetrics.consumerFetchMetricName, RequestMetrics.followFetchMetricName)).foreach { name =>
       metricsMap.put(name, new RequestMetrics(name))
     }
@@ -68,25 +81,36 @@ object RequestChannel extends Logging {
                 val context: RequestContext,
                 val startTimeNanos: Long,
                 memoryPool: MemoryPool,
-                @volatile private var buffer: ByteBuffer,
-                metrics: RequestChannel.Metrics) extends BaseRequest {
+                @volatile var buffer: ByteBuffer,
+                metrics: RequestChannel.Metrics,
+                val envelope: Option[RequestChannel.Request] = None) extends BaseRequest {
     // These need to be volatile because the readers are in the network thread and the writers are in the request
     // handler threads or the purgatory threads
     @volatile var requestDequeueTimeNanos = -1L
     @volatile var apiLocalCompleteTimeNanos = -1L
     @volatile var responseCompleteTimeNanos = -1L
     @volatile var responseDequeueTimeNanos = -1L
-    @volatile var apiRemoteCompleteTimeNanos = -1L
     @volatile var messageConversionsTimeNanos = 0L
+    @volatile var apiThrottleTimeMs = 0L
     @volatile var temporaryMemoryBytes = 0L
     @volatile var recordNetworkThreadTimeCallback: Option[Long => Unit] = None
 
     val session = Session(context.principal, context.clientAddress)
+
     private val bodyAndSize: RequestAndSize = context.parseRequest(buffer)
 
+    // This is constructed on creation of a Request so that the JSON representation is computed before the request is
+    // processed by the api layer. Otherwise, a ProduceRequest can occur without its data (ie. it goes into purgatory).
+    val requestLog: Option[JsonNode] =
+      if (RequestChannel.isRequestLoggingEnabled) Some(RequestConvertToJson.request(loggableRequest))
+      else None
+
     def header: RequestHeader = context.header
+
     def sizeOfBodyInBytes: Int = bodyAndSize.size
 
+    def sizeInBytes: Int = header.size(new ObjectSerializationCache) + sizeOfBodyInBytes
+
     //most request types are parsed entirely into objects at this point. for those we can release the underlying buffer.
     //some (like produce, or any time the schema contains fields of types BYTES or NULLABLE_BYTES) retain a reference
     //to the buffer. for those requests we cannot release the buffer early, but only when request processing is done.
@@ -94,9 +118,43 @@ object RequestChannel extends Logging {
       releaseBuffer()
     }
 
-    def requestDesc(details: Boolean): String = s"$header -- ${body[AbstractRequest].toString(details)}"
+    def isForwarded: Boolean = envelope.isDefined
+
+    def buildResponseSend(abstractResponse: AbstractResponse): Send = {
+      envelope match {
+        case Some(request) =>
+          val responseBytes = context.buildResponseEnvelopePayload(abstractResponse)
+          val envelopeResponse = new EnvelopeResponse(responseBytes, Errors.NONE)
+          request.context.buildResponseSend(envelopeResponse)
+        case None =>
+          context.buildResponseSend(abstractResponse)
+      }
+    }
+
+    def responseNode(response: AbstractResponse): Option[JsonNode] = {
+      if (RequestChannel.isRequestLoggingEnabled)
+        Some(RequestConvertToJson.response(response, context.apiVersion))
+      else
+        None
+    }
 
-    def body[T <: AbstractRequest](implicit classTag: ClassTag[T], nn: NotNothing[T]): T = {
+    def headerForLoggingOrThrottling(): RequestHeader = {
+      envelope match {
+        case Some(request) =>
+          request.context.header
+        case None =>
+          context.header
+      }
+    }
+
+    def requestDesc(details: Boolean): String = {
+      val forwardDescription = envelope.map { request =>
+        s"Forwarded request: ${request.context} "
+      }.getOrElse("")
+      s"$forwardDescription$header -- ${loggableRequest.toString(details)}"
+    }
+
+    def body[T <: AbstractRequest](implicit classTag: ClassTag[T], @nowarn("cat=unused") nn: NotNothing[T]): T = {
       bodyAndSize.request match {
         case r: T => r
         case r =>
@@ -104,25 +162,61 @@ object RequestChannel extends Logging {
       }
     }
 
+    def loggableRequest: AbstractRequest = {
+
+      def loggableValue(resourceType: ConfigResource.Type, name: String, value: String): String = {
+        val maybeSensitive = resourceType match {
+          case ConfigResource.Type.BROKER => KafkaConfig.maybeSensitive(KafkaConfig.configType(name))
+          case ConfigResource.Type.TOPIC => KafkaConfig.maybeSensitive(LogConfig.configType(name))
+          case ConfigResource.Type.BROKER_LOGGER => false
+          case _ => true
+        }
+        if (maybeSensitive) Password.HIDDEN else value
+      }
+
+      bodyAndSize.request match {
+        case alterConfigs: AlterConfigsRequest =>
+          val loggableConfigs = alterConfigs.configs().asScala.map { case (resource, config) =>
+            val loggableEntries = new AlterConfigsRequest.Config(config.entries.asScala.map { entry =>
+                new AlterConfigsRequest.ConfigEntry(entry.name, loggableValue(resource.`type`, entry.name, entry.value))
+            }.asJavaCollection)
+            (resource, loggableEntries)
+          }.asJava
+          new AlterConfigsRequest.Builder(loggableConfigs, alterConfigs.validateOnly).build(alterConfigs.version())
+
+        case alterConfigs: IncrementalAlterConfigsRequest =>
+          val resources = new AlterConfigsResourceCollection(alterConfigs.data.resources.size)
+          alterConfigs.data.resources.forEach { resource =>
+            val newResource = new AlterConfigsResource()
+              .setResourceName(resource.resourceName)
+              .setResourceType(resource.resourceType)
+            resource.configs.forEach { config =>
+              newResource.configs.add(new AlterableConfig()
+                .setName(config.name)
+                .setValue(loggableValue(ConfigResource.Type.forId(resource.resourceType), config.name, config.value))
+                .setConfigOperation(config.configOperation))
+            }
+            resources.add(newResource)
+          }
+          val data = new IncrementalAlterConfigsRequestData()
+            .setValidateOnly(alterConfigs.data().validateOnly())
+            .setResources(resources)
+          new IncrementalAlterConfigsRequest.Builder(data).build(alterConfigs.version)
+
+        case _ =>
+          bodyAndSize.request
+      }
+    }
+
     trace(s"Processor $processor received request: ${requestDesc(true)}")
 
-    def requestThreadTimeNanos = {
+    def requestThreadTimeNanos: Long = {
       if (apiLocalCompleteTimeNanos == -1L) apiLocalCompleteTimeNanos = Time.SYSTEM.nanoseconds
       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.
-      // This may be read in a network thread before the actual update happens in KafkaApis which will cause us to
-      // see a negative value here. In that case, use responseCompleteTimeNanos as apiLocalCompleteTimeNanos.
-      if (apiLocalCompleteTimeNanos < 0)
-        apiLocalCompleteTimeNanos = responseCompleteTimeNanos
-      // If the apiRemoteCompleteTimeNanos is not set (i.e., for requests that do not go through a purgatory), then it is
-      // the same as responseCompleteTimeNanos.
-      if (apiRemoteCompleteTimeNanos < 0)
-        apiRemoteCompleteTimeNanos = responseCompleteTimeNanos
 
       /**
        * Converts nanos to millis with micros precision as additional decimal places in the request log have low
@@ -136,8 +230,7 @@ object RequestChannel extends Logging {
 
       val requestQueueTimeMs = nanosToMs(requestDequeueTimeNanos - startTimeNanos)
       val apiLocalTimeMs = nanosToMs(apiLocalCompleteTimeNanos - requestDequeueTimeNanos)
-      val apiRemoteTimeMs = nanosToMs(apiRemoteCompleteTimeNanos - apiLocalCompleteTimeNanos)
-      val apiThrottleTimeMs = nanosToMs(responseCompleteTimeNanos - apiRemoteCompleteTimeNanos)
+      val apiRemoteTimeMs = nanosToMs(responseCompleteTimeNanos - apiLocalCompleteTimeNanos)
       val responseQueueTimeMs = nanosToMs(responseDequeueTimeNanos - responseCompleteTimeNanos)
       val responseSendTimeMs = nanosToMs(endTimeNanos - responseDequeueTimeNanos)
       val messageConversionsTimeMs = nanosToMs(messageConversionsTimeNanos)
@@ -154,11 +247,11 @@ object RequestChannel extends Logging {
       val metricNames = fetchMetricNames :+ header.apiKey.name
       metricNames.foreach { metricName =>
         val m = metrics(metricName)
-        m.requestRate.mark()
+        m.requestRate(header.apiVersion).mark()
         m.requestQueueTimeHist.update(Math.round(requestQueueTimeMs))
         m.localTimeHist.update(Math.round(apiLocalTimeMs))
         m.remoteTimeHist.update(Math.round(apiRemoteTimeMs))
-        m.throttleTimeHist.update(Math.round(apiThrottleTimeMs))
+        m.throttleTimeHist.update(apiThrottleTimeMs)
         m.responseQueueTimeHist.update(Math.round(responseQueueTimeMs))
         m.responseSendTimeHist.update(Math.round(responseSendTimeMs))
         m.totalTimeHist.update(Math.round(totalTimeMs))
@@ -176,36 +269,25 @@ object RequestChannel extends Logging {
       recordNetworkThreadTimeCallback.foreach(record => record(networkThreadTimeNanos))
 
       if (isRequestLoggingEnabled) {
-        val detailsEnabled = requestLogger.underlying.isTraceEnabled
-        val responseString = response.responseAsString.getOrElse(
-          throw new IllegalStateException("responseAsString should always be defined if request logging is enabled"))
-
-        val builder = new StringBuilder(256)
-        builder.append("Completed request:").append(requestDesc(detailsEnabled))
-          .append(",response:").append(responseString)
-          .append(" from connection ").append(context.connectionId)
-          .append(";totalTime:").append(totalTimeMs)
-          .append(",requestQueueTime:").append(requestQueueTimeMs)
-          .append(",localTime:").append(apiLocalTimeMs)
-          .append(",remoteTime:").append(apiRemoteTimeMs)
-          .append(",throttleTime:").append(apiThrottleTimeMs)
-          .append(",responseQueueTime:").append(responseQueueTimeMs)
-          .append(",sendTime:").append(responseSendTimeMs)
-          .append(",securityProtocol:").append(context.securityProtocol)
-          .append(",principal:").append(session.principal)
-          .append(",listener:").append(context.listenerName.value)
-        if (temporaryMemoryBytes > 0)
-          builder.append(",temporaryMemoryBytes:").append(temporaryMemoryBytes)
-        if (messageConversionsTimeMs > 0)
-          builder.append(",messageConversionsTime:").append(messageConversionsTimeMs)
-        requestLogger.debug(builder.toString)
+        val desc = RequestConvertToJson.requestDescMetrics(header, requestLog, response.responseLog,
+          context, session, isForwarded,
+          totalTimeMs, requestQueueTimeMs, apiLocalTimeMs,
+          apiRemoteTimeMs, apiThrottleTimeMs, responseQueueTimeMs,
+          responseSendTimeMs, temporaryMemoryBytes,
+          messageConversionsTimeMs)
+        requestLogger.debug("Completed request:" + desc.toString)
       }
     }
 
     def releaseBuffer(): Unit = {
-      if (buffer != null) {
-        memoryPool.release(buffer)
-        buffer = null
+      envelope match {
+        case Some(request) =>
+          request.releaseBuffer()
+        case None =>
+          if (buffer != null) {
+            memoryPool.release(buffer)
+            buffer = null
+          }
       }
     }
 
@@ -214,79 +296,131 @@ object RequestChannel extends Logging {
       s"session=$session, " +
       s"listenerName=${context.listenerName}, " +
       s"securityProtocol=${context.securityProtocol}, " +
-      s"buffer=$buffer)"
+      s"buffer=$buffer, " +
+      s"envelope=$envelope)"
 
   }
 
-  /** responseAsString should only be defined if request logging is enabled */
-  class Response(val request: Request, val responseSend: Option[Send], val responseAction: ResponseAction,
-                 val responseAsString: Option[String]) {
-    request.responseCompleteTimeNanos = Time.SYSTEM.nanoseconds
-    if (request.apiLocalCompleteTimeNanos == -1L) request.apiLocalCompleteTimeNanos = Time.SYSTEM.nanoseconds
+  sealed abstract class Response(val request: Request) {
 
     def processor: Int = request.processor
 
-    override def toString =
-      s"Response(request=$request, responseSend=$responseSend, responseAction=$responseAction), responseAsString=$responseAsString"
+    def responseLog: Option[JsonNode] = None
+
+    def onComplete: Option[Send => Unit] = None
+
+    override def toString: String
+  }
+
+  /** responseLogValue should only be defined if request logging is enabled */
+  class SendResponse(request: Request,
+                     val responseSend: Send,
+                     val responseLogValue: Option[JsonNode],
+                     val onCompleteCallback: Option[Send => Unit]) extends Response(request) {
+    override def responseLog: Option[JsonNode] = responseLogValue
+
+    override def onComplete: Option[Send => Unit] = onCompleteCallback
+
+    override def toString: String =
+      s"Response(type=Send, request=$request, send=$responseSend, asString=$responseLogValue)"
+  }
+
+  class NoOpResponse(request: Request) extends Response(request) {
+    override def toString: String =
+      s"Response(type=NoOp, request=$request)"
+  }
+
+  class CloseConnectionResponse(request: Request) extends Response(request) {
+    override def toString: String =
+      s"Response(type=CloseConnection, request=$request)"
+  }
+
+  class StartThrottlingResponse(request: Request) extends Response(request) {
+    override def toString: String =
+      s"Response(type=StartThrottling, request=$request)"
   }
 
-  sealed trait ResponseAction
-  case object SendAction extends ResponseAction
-  case object NoOpAction extends ResponseAction
-  case object CloseConnectionAction extends ResponseAction
+  class EndThrottlingResponse(request: Request) extends Response(request) {
+    override def toString: String =
+      s"Response(type=EndThrottling, request=$request)"
+  }
 }
 
-class RequestChannel(val numProcessors: Int, val queueSize: Int) extends KafkaMetricsGroup {
-  val metrics = new RequestChannel.Metrics
-  private var responseListeners: List[(Int) => Unit] = Nil
+class RequestChannel(val queueSize: Int,
+                     val metricNamePrefix: String,
+                     time: Time,
+                     allowDisabledApis: Boolean = false) extends KafkaMetricsGroup {
+  import RequestChannel._
+  val metrics = new RequestChannel.Metrics(allowDisabledApis)
   private val requestQueue = new ArrayBlockingQueue[BaseRequest](queueSize)
-  private val responseQueues = new Array[BlockingQueue[RequestChannel.Response]](numProcessors)
-  for(i <- 0 until numProcessors)
-    responseQueues(i) = new LinkedBlockingQueue[RequestChannel.Response]()
-
-  newGauge(
-    "RequestQueueSize",
-    new Gauge[Int] {
-      def value = requestQueue.size
-    }
-  )
+  private val processors = new ConcurrentHashMap[Int, Processor]()
+  val requestQueueSizeMetricName = metricNamePrefix.concat(RequestQueueSizeMetric)
+  val responseQueueSizeMetricName = metricNamePrefix.concat(ResponseQueueSizeMetric)
+
+  newGauge(requestQueueSizeMetricName, () => requestQueue.size)
 
-  newGauge("ResponseQueueSize", new Gauge[Int]{
-    def value = responseQueues.foldLeft(0) {(total, q) => total + q.size()}
+  newGauge(responseQueueSizeMetricName, () => {
+    processors.values.asScala.foldLeft(0) {(total, processor) =>
+      total + processor.responseQueueSize
+    }
   })
 
-  for (i <- 0 until numProcessors) {
-    newGauge("ResponseQueueSize",
-      new Gauge[Int] {
-        def value = responseQueues(i).size()
-      },
-      Map("processor" -> i.toString)
-    )
+  def addProcessor(processor: Processor): Unit = {
+    if (processors.putIfAbsent(processor.id, processor) != null)
+      warn(s"Unexpected processor with processorId ${processor.id}")
+
+    newGauge(responseQueueSizeMetricName, () => processor.responseQueueSize,
+      Map(ProcessorMetricTag -> processor.id.toString))
+  }
+
+  def removeProcessor(processorId: Int): Unit = {
+    processors.remove(processorId)
+    removeMetric(responseQueueSizeMetricName, Map(ProcessorMetricTag -> processorId.toString))
   }
 
   /** 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.responseAction match {
-        case SendAction =>
-          s"Sending ${requestHeader.apiKey} response to client ${requestHeader.clientId} of ${response.responseSend.get.size} bytes."
-        case NoOpAction =>
+      val requestHeader = response.request.headerForLoggingOrThrottling()
+      val message = response match {
+        case sendResponse: SendResponse =>
+          s"Sending ${requestHeader.apiKey} response to client ${requestHeader.clientId} of ${sendResponse.responseSend.size} bytes."
+        case _: NoOpResponse =>
           s"Not sending ${requestHeader.apiKey} response to client ${requestHeader.clientId} as it's not required."
-        case CloseConnectionAction =>
+        case _: CloseConnectionResponse =>
           s"Closing connection for client ${requestHeader.clientId} due to error during ${requestHeader.apiKey}."
+        case _: StartThrottlingResponse =>
+          s"Notifying channel throttling has started for client ${requestHeader.clientId} for ${requestHeader.apiKey}"
+        case _: EndThrottlingResponse =>
+          s"Notifying channel throttling has ended for client ${requestHeader.clientId} for ${requestHeader.apiKey}"
       }
       trace(message)
     }
 
-    responseQueues(response.processor).put(response)
-    for(onResponse <- responseListeners)
-      onResponse(response.processor)
+    response match {
+      // We should only send one of the following per request
+      case _: SendResponse | _: NoOpResponse | _: CloseConnectionResponse =>
+        val request = response.request
+        val timeNanos = time.nanoseconds()
+        request.responseCompleteTimeNanos = timeNanos
+        if (request.apiLocalCompleteTimeNanos == -1L)
+          request.apiLocalCompleteTimeNanos = timeNanos
+      // For a given request, these may happen in addition to one in the previous section, skip updating the metrics
+      case _: StartThrottlingResponse | _: EndThrottlingResponse => ()
+    }
+
+    val processor = processors.get(response.processor)
+    // The processor may be null if it was shutdown. In this case, the connections
+    // are closed, so the response is dropped.
+    if (processor != null) {
+      processor.enqueueResponse(response)
+    }
   }
 
   /** Get the next request or block until specified time has elapsed */
@@ -297,29 +431,17 @@ class RequestChannel(val numProcessors: Int, val queueSize: Int) extends KafkaMe
   def receiveRequest(): RequestChannel.BaseRequest =
     requestQueue.take()
 
-  /** Get a response for the given processor if there is one */
-  def receiveResponse(processor: Int): RequestChannel.Response = {
-    val response = responseQueues(processor).poll()
-    if (response != null)
-      response.request.responseDequeueTimeNanos = Time.SYSTEM.nanoseconds
-    response
-  }
-
-  def addResponseListener(onResponse: Int => Unit) {
-    responseListeners ::= onResponse
-  }
-
-  def updateErrorMetrics(apiKey: ApiKeys, errors: collection.Map[Errors, Integer]) {
-    errors.foreach { case (error, count) =>
+  def updateErrorMetrics(apiKey: ApiKeys, errors: collection.Map[Errors, Integer]): Unit = {
+    errors.forKeyValue { (error, count) =>
       metrics(apiKey.name).markErrorMeter(error, count)
     }
   }
 
-  def clear() {
+  def clear(): Unit = {
     requestQueue.clear()
   }
 
-  def shutdown() {
+  def shutdown(): Unit = {
     clear()
     metrics.close()
   }
@@ -347,17 +469,19 @@ object RequestMetrics {
 }
 
 class RequestMetrics(name: String) extends KafkaMetricsGroup {
+
   import RequestMetrics._
 
   val tags = Map("request" -> name)
-  val requestRate = newMeter(RequestsPerSec, "requests", TimeUnit.SECONDS, tags)
+  val requestRateInternal = new Pool[Short, Meter]()
   // time a request spent in a request queue
   val requestQueueTimeHist = newHistogram(RequestQueueTimeMs, biased = true, tags)
   // time a request takes to be processed at the local broker
   val localTimeHist = newHistogram(LocalTimeMs, biased = true, tags)
   // time a request takes to wait on remote brokers (currently only relevant to fetch and produce requests)
   val remoteTimeHist = newHistogram(RemoteTimeMs, biased = true, tags)
-  // time a request is throttled
+  // time a request is throttled, not part of the request processing time (throttling is done at the client level
+  // for clients that support KIP-219 and by muting the channel for the rest)
   val throttleTimeHist = newHistogram(ThrottleTimeMs, biased = true, tags)
   // time a response spent in a response queue
   val responseQueueTimeHist = newHistogram(ResponseQueueTimeMs, biased = true, tags)
@@ -383,6 +507,10 @@ class RequestMetrics(name: String) extends KafkaMetricsGroup {
   private val errorMeters = mutable.Map[Errors, ErrorMeter]()
   Errors.values.foreach(error => errorMeters.put(error, new ErrorMeter(name, error)))
 
+  def requestRate(version: Short): Meter = {
+    requestRateInternal.getAndMaybePut(version, newMeter("RequestsPerSec", "requests", TimeUnit.SECONDS, tags + ("version" -> version.toString)))
+  }
+
   class ErrorMeter(name: String, error: Errors) {
     private val tags = Map("request" -> name, "error" -> error.name)
 
@@ -410,12 +538,12 @@ 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)
   }
 
   def removeMetrics(): Unit = {
-    removeMetric(RequestsPerSec, tags)
+    for (version <- requestRateInternal.keys) removeMetric(RequestsPerSec, tags + ("version" -> version.toString))
     removeMetric(RequestQueueTimeMs, tags)
     removeMetric(LocalTimeMs, tags)
     removeMetric(RemoteTimeMs, tags)
diff --git a/core/src/main/scala/kafka/network/RequestConvertToJson.scala b/core/src/main/scala/kafka/network/RequestConvertToJson.scala
new file mode 100644
index 0000000000000..243c2a4ca50ed
--- /dev/null
+++ b/core/src/main/scala/kafka/network/RequestConvertToJson.scala
@@ -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 kafka.network
+
+import com.fasterxml.jackson.databind.JsonNode
+import com.fasterxml.jackson.databind.node.{BooleanNode, DoubleNode, JsonNodeFactory, LongNode, ObjectNode, TextNode}
+import kafka.network.RequestChannel.Session
+import org.apache.kafka.common.message._
+import org.apache.kafka.common.network.ClientInformation
+import org.apache.kafka.common.requests._
+
+object RequestConvertToJson {
+  def request(request: AbstractRequest): JsonNode = {
+    request match {
+      case req: AddOffsetsToTxnRequest => AddOffsetsToTxnRequestDataJsonConverter.write(req.data, request.version)
+      case req: AddPartitionsToTxnRequest => AddPartitionsToTxnRequestDataJsonConverter.write(req.data, request.version)
+      case req: AlterClientQuotasRequest => AlterClientQuotasRequestDataJsonConverter.write(req.data, request.version)
+      case req: AlterConfigsRequest => AlterConfigsRequestDataJsonConverter.write(req.data, request.version)
+      case req: AlterIsrRequest => AlterIsrRequestDataJsonConverter.write(req.data, request.version)
+      case req: AlterPartitionReassignmentsRequest => AlterPartitionReassignmentsRequestDataJsonConverter.write(req.data, request.version)
+      case req: AlterReplicaLogDirsRequest => AlterReplicaLogDirsRequestDataJsonConverter.write(req.data, request.version)
+      case res: AlterUserScramCredentialsRequest => AlterUserScramCredentialsRequestDataJsonConverter.write(res.data, request.version)
+      case req: ApiVersionsRequest => ApiVersionsRequestDataJsonConverter.write(req.data, request.version)
+      case req: BeginQuorumEpochRequest => BeginQuorumEpochRequestDataJsonConverter.write(req.data, request.version)
+      case req: ControlledShutdownRequest => ControlledShutdownRequestDataJsonConverter.write(req.data, request.version)
+      case req: CreateAclsRequest => CreateAclsRequestDataJsonConverter.write(req.data, request.version)
+      case req: CreateDelegationTokenRequest => CreateDelegationTokenRequestDataJsonConverter.write(req.data, request.version)
+      case req: CreatePartitionsRequest => CreatePartitionsRequestDataJsonConverter.write(req.data, request.version)
+      case req: CreateTopicsRequest => CreateTopicsRequestDataJsonConverter.write(req.data, request.version)
+      case req: DeleteAclsRequest => DeleteAclsRequestDataJsonConverter.write(req.data, request.version)
+      case req: DeleteGroupsRequest => DeleteGroupsRequestDataJsonConverter.write(req.data, request.version)
+      case req: DeleteRecordsRequest => DeleteRecordsRequestDataJsonConverter.write(req.data, request.version)
+      case req: DeleteTopicsRequest => DeleteTopicsRequestDataJsonConverter.write(req.data, request.version)
+      case req: DescribeAclsRequest => DescribeAclsRequestDataJsonConverter.write(req.data, request.version)
+      case req: DescribeClientQuotasRequest => DescribeClientQuotasRequestDataJsonConverter.write(req.data, request.version)
+      case req: DescribeConfigsRequest => DescribeConfigsRequestDataJsonConverter.write(req.data, request.version)
+      case req: DescribeDelegationTokenRequest => DescribeDelegationTokenRequestDataJsonConverter.write(req.data, request.version)
+      case req: DescribeGroupsRequest => DescribeGroupsRequestDataJsonConverter.write(req.data, request.version)
+      case req: DescribeLogDirsRequest => DescribeLogDirsRequestDataJsonConverter.write(req.data, request.version)
+      case req: DescribeQuorumRequest => DescribeQuorumRequestDataJsonConverter.write(req.data, request.version)
+      case res: DescribeUserScramCredentialsRequest => DescribeUserScramCredentialsRequestDataJsonConverter.write(res.data, request.version)
+      case req: ElectLeadersRequest => ElectLeadersRequestDataJsonConverter.write(req.data, request.version)
+      case req: EndTxnRequest => EndTxnRequestDataJsonConverter.write(req.data, request.version)
+      case req: EndQuorumEpochRequest => EndQuorumEpochRequestDataJsonConverter.write(req.data, request.version)
+      case req: EnvelopeRequest => EnvelopeRequestDataJsonConverter.write(req.data, request.version)
+      case req: ExpireDelegationTokenRequest => ExpireDelegationTokenRequestDataJsonConverter.write(req.data, request.version)
+      case req: FetchRequest => FetchRequestDataJsonConverter.write(req.data, request.version)
+      case req: FindCoordinatorRequest => FindCoordinatorRequestDataJsonConverter.write(req.data, request.version)
+      case req: HeartbeatRequest => HeartbeatRequestDataJsonConverter.write(req.data, request.version)
+      case req: IncrementalAlterConfigsRequest => IncrementalAlterConfigsRequestDataJsonConverter.write(req.data, request.version)
+      case req: InitProducerIdRequest => InitProducerIdRequestDataJsonConverter.write(req.data, request.version)
+      case req: JoinGroupRequest => JoinGroupRequestDataJsonConverter.write(req.data, request.version)
+      case req: LeaderAndIsrRequest => LeaderAndIsrRequestDataJsonConverter.write(req.data, request.version)
+      case req: LeaveGroupRequest => LeaveGroupRequestDataJsonConverter.write(req.data, request.version)
+      case req: ListGroupsRequest => ListGroupsRequestDataJsonConverter.write(req.data, request.version)
+      case req: ListOffsetRequest => ListOffsetRequestDataJsonConverter.write(req.data, request.version)
+      case req: ListPartitionReassignmentsRequest => ListPartitionReassignmentsRequestDataJsonConverter.write(req.data, request.version)
+      case req: MetadataRequest => MetadataRequestDataJsonConverter.write(req.data, request.version)
+      case req: OffsetCommitRequest => OffsetCommitRequestDataJsonConverter.write(req.data, request.version)
+      case req: OffsetDeleteRequest => OffsetDeleteRequestDataJsonConverter.write(req.data, request.version)
+      case req: OffsetFetchRequest => OffsetFetchRequestDataJsonConverter.write(req.data, request.version)
+      case req: OffsetsForLeaderEpochRequest => OffsetForLeaderEpochRequestDataJsonConverter.write(req.data, request.version)
+      case req: ProduceRequest => ProduceRequestDataJsonConverter.write(req.data, request.version, false)
+      case req: RenewDelegationTokenRequest => RenewDelegationTokenRequestDataJsonConverter.write(req.data, request.version)
+      case req: SaslAuthenticateRequest => SaslAuthenticateRequestDataJsonConverter.write(req.data, request.version)
+      case req: SaslHandshakeRequest => SaslHandshakeRequestDataJsonConverter.write(req.data, request.version)
+      case req: StopReplicaRequest => StopReplicaRequestDataJsonConverter.write(req.data, request.version)
+      case req: SyncGroupRequest => SyncGroupRequestDataJsonConverter.write(req.data, request.version)
+      case req: TxnOffsetCommitRequest => TxnOffsetCommitRequestDataJsonConverter.write(req.data, request.version)
+      case req: UpdateFeaturesRequest => UpdateFeaturesRequestDataJsonConverter.write(req.data, request.version)
+      case req: UpdateMetadataRequest => UpdateMetadataRequestDataJsonConverter.write(req.data, request.version)
+      case req: VoteRequest => VoteRequestDataJsonConverter.write(req.data, request.version)
+      case req: WriteTxnMarkersRequest => WriteTxnMarkersRequestDataJsonConverter.write(req.data, request.version)
+      case _ => throw new IllegalStateException(s"ApiKey ${request.apiKey} is not currently handled in `request`, the " +
+        "code should be updated to do so.");
+    }
+  }
+
+  def response(response: AbstractResponse, version: Short): JsonNode = {
+    response match {
+      case res: AddOffsetsToTxnResponse => AddOffsetsToTxnResponseDataJsonConverter.write(res.data, version)
+      case res: AddPartitionsToTxnResponse => AddPartitionsToTxnResponseDataJsonConverter.write(res.data, version)
+      case res: AlterClientQuotasResponse => AlterClientQuotasResponseDataJsonConverter.write(res.data, version)
+      case res: AlterConfigsResponse => AlterConfigsResponseDataJsonConverter.write(res.data, version)
+      case res: AlterIsrResponse => AlterIsrResponseDataJsonConverter.write(res.data, version)
+      case res: AlterPartitionReassignmentsResponse => AlterPartitionReassignmentsResponseDataJsonConverter.write(res.data, version)
+      case res: AlterReplicaLogDirsResponse => AlterReplicaLogDirsResponseDataJsonConverter.write(res.data, version)
+      case res: AlterUserScramCredentialsResponse => AlterUserScramCredentialsResponseDataJsonConverter.write(res.data, version)
+      case res: ApiVersionsResponse => ApiVersionsResponseDataJsonConverter.write(res.data, version)
+      case res: BeginQuorumEpochResponse => BeginQuorumEpochResponseDataJsonConverter.write(res.data, version)
+      case res: ControlledShutdownResponse => ControlledShutdownResponseDataJsonConverter.write(res.data, version)
+      case res: CreateAclsResponse => CreateAclsResponseDataJsonConverter.write(res.data, version)
+      case res: CreateDelegationTokenResponse => CreateDelegationTokenResponseDataJsonConverter.write(res.data, version)
+      case res: CreatePartitionsResponse => CreatePartitionsResponseDataJsonConverter.write(res.data, version)
+      case res: CreateTopicsResponse => CreateTopicsResponseDataJsonConverter.write(res.data, version)
+      case res: DeleteAclsResponse => DeleteAclsResponseDataJsonConverter.write(res.data, version)
+      case res: DeleteGroupsResponse => DeleteGroupsResponseDataJsonConverter.write(res.data, version)
+      case res: DeleteRecordsResponse => DeleteRecordsResponseDataJsonConverter.write(res.data, version)
+      case res: DeleteTopicsResponse => DeleteTopicsResponseDataJsonConverter.write(res.data, version)
+      case res: DescribeAclsResponse => DescribeAclsResponseDataJsonConverter.write(res.data, version)
+      case res: DescribeClientQuotasResponse => DescribeClientQuotasResponseDataJsonConverter.write(res.data, version)
+      case res: DescribeConfigsResponse => DescribeConfigsResponseDataJsonConverter.write(res.data, version)
+      case res: DescribeDelegationTokenResponse => DescribeDelegationTokenResponseDataJsonConverter.write(res.data, version)
+      case res: DescribeGroupsResponse => DescribeGroupsResponseDataJsonConverter.write(res.data, version)
+      case res: DescribeLogDirsResponse => DescribeLogDirsResponseDataJsonConverter.write(res.data, version)
+      case res: DescribeQuorumResponse => DescribeQuorumResponseDataJsonConverter.write(res.data, version)
+      case res: DescribeUserScramCredentialsResponse => DescribeUserScramCredentialsResponseDataJsonConverter.write(res.data, version)
+      case res: ElectLeadersResponse => ElectLeadersResponseDataJsonConverter.write(res.data, version)
+      case res: EndTxnResponse => EndTxnResponseDataJsonConverter.write(res.data, version)
+      case res: EndQuorumEpochResponse => EndQuorumEpochResponseDataJsonConverter.write(res.data, version)
+      case res: EnvelopeResponse => EnvelopeResponseDataJsonConverter.write(res.data, version)
+      case res: ExpireDelegationTokenResponse => ExpireDelegationTokenResponseDataJsonConverter.write(res.data, version)
+      case res: FetchResponse[_] => FetchResponseDataJsonConverter.write(res.data, version, false)
+      case res: FindCoordinatorResponse => FindCoordinatorResponseDataJsonConverter.write(res.data, version)
+      case res: HeartbeatResponse => HeartbeatResponseDataJsonConverter.write(res.data, version)
+      case res: IncrementalAlterConfigsResponse => IncrementalAlterConfigsResponseDataJsonConverter.write(res.data, version)
+      case res: InitProducerIdResponse => InitProducerIdResponseDataJsonConverter.write(res.data, version)
+      case res: JoinGroupResponse => JoinGroupResponseDataJsonConverter.write(res.data, version)
+      case res: LeaderAndIsrResponse => LeaderAndIsrResponseDataJsonConverter.write(res.data, version)
+      case res: LeaveGroupResponse => LeaveGroupResponseDataJsonConverter.write(res.data, version)
+      case res: ListGroupsResponse => ListGroupsResponseDataJsonConverter.write(res.data, version)
+      case res: ListOffsetResponse => ListOffsetResponseDataJsonConverter.write(res.data, version)
+      case res: ListPartitionReassignmentsResponse => ListPartitionReassignmentsResponseDataJsonConverter.write(res.data, version)
+      case res: MetadataResponse => MetadataResponseDataJsonConverter.write(res.data, version)
+      case res: OffsetCommitResponse => OffsetCommitResponseDataJsonConverter.write(res.data, version)
+      case res: OffsetDeleteResponse => OffsetDeleteResponseDataJsonConverter.write(res.data, version)
+      case res: OffsetFetchResponse => OffsetFetchResponseDataJsonConverter.write(res.data, version)
+      case res: OffsetsForLeaderEpochResponse => OffsetForLeaderEpochResponseDataJsonConverter.write(res.data, version)
+      case res: ProduceResponse => ProduceResponseDataJsonConverter.write(res.data, version)
+      case res: RenewDelegationTokenResponse => RenewDelegationTokenResponseDataJsonConverter.write(res.data, version)
+      case res: SaslAuthenticateResponse => SaslAuthenticateResponseDataJsonConverter.write(res.data, version)
+      case res: SaslHandshakeResponse => SaslHandshakeResponseDataJsonConverter.write(res.data, version)
+      case res: StopReplicaResponse => StopReplicaResponseDataJsonConverter.write(res.data, version)
+      case res: SyncGroupResponse => SyncGroupResponseDataJsonConverter.write(res.data, version)
+      case res: TxnOffsetCommitResponse => TxnOffsetCommitResponseDataJsonConverter.write(res.data, version)
+      case res: UpdateFeaturesResponse => UpdateFeaturesResponseDataJsonConverter.write(res.data, version)
+      case res: UpdateMetadataResponse => UpdateMetadataResponseDataJsonConverter.write(res.data, version)
+      case res: WriteTxnMarkersResponse => WriteTxnMarkersResponseDataJsonConverter.write(res.data, version)
+      case res: VoteResponse => VoteResponseDataJsonConverter.write(res.data, version)
+      case _ => throw new IllegalStateException(s"ApiKey ${response.apiKey} is not currently handled in `response`, the " +
+        "code should be updated to do so.");
+    }
+  }
+
+  def requestHeaderNode(header: RequestHeader): JsonNode = {
+    val node = RequestHeaderDataJsonConverter.write(header.data, header.headerVersion, false).asInstanceOf[ObjectNode]
+    node.set("requestApiKeyName", new TextNode(header.apiKey.toString))
+    node
+  }
+
+  def requestDesc(header: RequestHeader, requestNode: Option[JsonNode], isForwarded: Boolean): JsonNode = {
+    val node = new ObjectNode(JsonNodeFactory.instance)
+    node.set("isForwarded", if (isForwarded) BooleanNode.TRUE else BooleanNode.FALSE)
+    node.set("requestHeader", requestHeaderNode(header))
+    node.set("request", requestNode.getOrElse(new TextNode("")))
+    node
+  }
+
+  def clientInfoNode(clientInfo: ClientInformation): JsonNode = {
+    val node = new ObjectNode(JsonNodeFactory.instance)
+    node.set("softwareName", new TextNode(clientInfo.softwareName))
+    node.set("softwareVersion", new TextNode(clientInfo.softwareVersion))
+    node
+  }
+
+  def requestDescMetrics(header: RequestHeader, requestNode: Option[JsonNode], responseNode: Option[JsonNode],
+                         context: RequestContext, session: Session, isForwarded: Boolean,
+                         totalTimeMs: Double, requestQueueTimeMs: Double, apiLocalTimeMs: Double,
+                         apiRemoteTimeMs: Double, apiThrottleTimeMs: Long, responseQueueTimeMs: Double,
+                         responseSendTimeMs: Double, temporaryMemoryBytes: Long,
+                         messageConversionsTimeMs: Double): JsonNode = {
+    val node = requestDesc(header, requestNode, isForwarded).asInstanceOf[ObjectNode]
+    node.set("response", responseNode.getOrElse(new TextNode("")))
+    node.set("connection", new TextNode(context.connectionId))
+    node.set("totalTimeMs", new DoubleNode(totalTimeMs))
+    node.set("requestQueueTimeMs", new DoubleNode(requestQueueTimeMs))
+    node.set("localTimeMs", new DoubleNode(apiLocalTimeMs))
+    node.set("remoteTimeMs", new DoubleNode(apiRemoteTimeMs))
+    node.set("throttleTimeMs", new LongNode(apiThrottleTimeMs))
+    node.set("responseQueueTimeMs", new DoubleNode(responseQueueTimeMs))
+    node.set("sendTimeMs", new DoubleNode(responseSendTimeMs))
+    node.set("securityProtocol", new TextNode(context.securityProtocol.toString))
+    node.set("principal", new TextNode(session.principal.toString))
+    node.set("listener", new TextNode(context.listenerName.value))
+    node.set("clientInformation", clientInfoNode(context.clientInformation))
+    if (temporaryMemoryBytes > 0)
+      node.set("temporaryMemoryBytes", new LongNode(temporaryMemoryBytes))
+    if (messageConversionsTimeMs > 0)
+      node.set("messageConversionsTime", new DoubleNode(messageConversionsTimeMs))
+    node
+  }
+}
diff --git a/core/src/main/scala/kafka/network/RequestOrResponseSend.scala b/core/src/main/scala/kafka/network/RequestOrResponseSend.scala
deleted file mode 100644
index 7a14e5e7cacec..0000000000000
--- a/core/src/main/scala/kafka/network/RequestOrResponseSend.scala
+++ /dev/null
@@ -1,53 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package kafka.network
-
-import java.nio.ByteBuffer
-import java.nio.channels.GatheringByteChannel
-
-import kafka.api.RequestOrResponse
-import kafka.utils.Logging
-import org.apache.kafka.common.network.NetworkSend
-
-object RequestOrResponseSend {
-  def serialize(request: RequestOrResponse): ByteBuffer = {
-    val buffer = ByteBuffer.allocate(request.sizeInBytes + request.requestId.fold(0)(_ => 2))
-    request.requestId.foreach(buffer.putShort)
-    request.writeTo(buffer)
-    buffer.rewind()
-    buffer
-  }
-}
-
-class RequestOrResponseSend(val dest: String, val buffer: ByteBuffer) extends NetworkSend(dest, buffer) with Logging {
-
-  def this(dest: String, request: RequestOrResponse) {
-    this(dest, RequestOrResponseSend.serialize(request))
-  }
-
-  def writeCompletely(channel: GatheringByteChannel): Long = {
-    var totalWritten = 0L
-    while(!completed()) {
-      val written = writeTo(channel)
-      trace(written + " bytes written.")
-      totalWritten += written
-    }
-    totalWritten
-  }
-
-}
diff --git a/core/src/main/scala/kafka/network/SocketServer.scala b/core/src/main/scala/kafka/network/SocketServer.scala
index 200bfe2bc5b3a..93d9b38a76088 100644
--- a/core/src/main/scala/kafka/network/SocketServer.scala
+++ b/core/src/main/scala/kafka/network/SocketServer.scala
@@ -19,182 +19,443 @@ package kafka.network
 
 import java.io.IOException
 import java.net._
-import java.nio.channels._
-import java.nio.channels.{Selector => NSelector}
+import java.nio.ByteBuffer
+import java.nio.channels.{Selector => NSelector, _}
+import java.util
+import java.util.Optional
 import java.util.concurrent._
 import java.util.concurrent.atomic._
-
-import com.yammer.metrics.core.Gauge
 import kafka.cluster.{BrokerEndPoint, EndPoint}
-import kafka.common.KafkaException
 import kafka.metrics.KafkaMetricsGroup
+import kafka.network.ConnectionQuotas._
+import kafka.network.Processor._
+import kafka.network.RequestChannel.{CloseConnectionResponse, EndThrottlingResponse, NoOpResponse, SendResponse, StartThrottlingResponse}
+import kafka.network.SocketServer._
 import kafka.security.CredentialProvider
-import kafka.server.KafkaConfig
+import kafka.server.{BrokerReconfigurable, DynamicConfig, KafkaConfig}
 import kafka.utils._
+import kafka.utils.Implicits._
+import org.apache.kafka.common.config.ConfigException
+import org.apache.kafka.common.errors.InvalidRequestException
 import org.apache.kafka.common.memory.{MemoryPool, SimpleMemoryPool}
 import org.apache.kafka.common.metrics._
-import org.apache.kafka.common.metrics.stats.Meter
-import org.apache.kafka.common.network.{ChannelBuilder, ChannelBuilders, KafkaChannel, ListenerName, Selectable, Send, Selector => KSelector}
-import org.apache.kafka.common.requests.{RequestContext, RequestHeader}
-import org.apache.kafka.common.security.auth.SecurityProtocol
+import org.apache.kafka.common.metrics.stats.{Avg, CumulativeSum, Meter, Rate}
+import org.apache.kafka.common.network.KafkaChannel.ChannelMuteEvent
+import org.apache.kafka.common.network.{ChannelBuilder, ChannelBuilders, ClientInformation, NetworkSend, KafkaChannel, ListenerName, ListenerReconfigurable, Selectable, Send, Selector => KSelector}
+import org.apache.kafka.common.protocol.{ApiKeys, Errors}
+import org.apache.kafka.common.requests.{ApiVersionsRequest, EnvelopeRequest, EnvelopeResponse, RequestContext, RequestHeader}
+import org.apache.kafka.common.security.auth.{KafkaPrincipal, KafkaPrincipalSerde, SecurityProtocol}
 import org.apache.kafka.common.utils.{KafkaThread, LogContext, Time}
+import org.apache.kafka.common.{Endpoint, KafkaException, MetricName, Reconfigurable}
 import org.slf4j.event.Level
 
 import scala.collection._
-import JavaConverters._
+import scala.collection.mutable.{ArrayBuffer, Buffer}
+import scala.compat.java8.OptionConverters._
+import scala.jdk.CollectionConverters._
 import scala.util.control.ControlThrowable
 
 /**
- * An NIO socket server. The threading model is
- *   1 Acceptor thread that handles new connections
- *   Acceptor has N Processor threads that each have their own selector and read requests from sockets
- *   M Handler threads that handle requests and produce responses back to the processor threads for writing.
+ * Handles new connections, requests and responses to and from broker.
+ * Kafka supports two types of request planes :
+ *  - data-plane :
+ *    - Handles requests from clients and other brokers in the cluster.
+ *    - The threading model is
+ *      1 Acceptor thread per listener, that handles new connections.
+ *      It is possible to configure multiple data-planes by specifying multiple "," separated endpoints for "listeners" in KafkaConfig.
+ *      Acceptor has N Processor threads that each have their own selector and read requests from sockets
+ *      M Handler threads that handle requests and produce responses back to the processor threads for writing.
+ *  - control-plane :
+ *    - Handles requests from controller. This is optional and can be configured by specifying "control.plane.listener.name".
+ *      If not configured, the controller requests are handled by the data-plane.
+ *    - The threading model is
+ *      1 Acceptor thread that handles new connections
+ *      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,
+                   val allowDisabledApis: Boolean = false)
+  extends Logging with KafkaMetricsGroup with BrokerReconfigurable {
 
-  private val endpoints = config.listeners.map(l => l.listenerName -> l).toMap
-  private val numProcessorThreads = config.numNetworkThreads
   private val maxQueuedRequests = config.queuedMaxRequests
-  private val totalProcessorThreads = numProcessorThreads * endpoints.size
-
-  private val maxConnectionsPerIp = config.maxConnectionsPerIp
-  private val maxConnectionsPerIpOverrides = config.maxConnectionsPerIpOverrides
 
   private val logContext = new LogContext(s"[SocketServer brokerId=${config.brokerId}] ")
   this.logIdent = logContext.logPrefix
 
   private val memoryPoolSensor = metrics.sensor("MemoryPoolUtilization")
-  private val memoryPoolDepletedPercentMetricName = metrics.metricName("MemoryPoolAvgDepletedPercent", "socket-server-metrics")
-  private val memoryPoolDepletedTimeMetricName = metrics.metricName("MemoryPoolDepletedTimeTotal", "socket-server-metrics")
+  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
-  val requestChannel = new RequestChannel(totalProcessorThreads, maxQueuedRequests)
-  private val processors = new Array[Processor](totalProcessorThreads)
-
-  private[network] val acceptors = mutable.Map[EndPoint, Acceptor]()
-  private var connectionQuotas: ConnectionQuotas = _
+  // data-plane
+  private val dataPlaneProcessors = new ConcurrentHashMap[Int, Processor]()
+  private[network] val dataPlaneAcceptors = new ConcurrentHashMap[EndPoint, Acceptor]()
+  val dataPlaneRequestChannel = new RequestChannel(maxQueuedRequests, DataPlaneMetricPrefix, time, allowDisabledApis)
+  // 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, time, allowDisabledApis))
+
+  private var nextProcessorId = 0
+  val connectionQuotas = new ConnectionQuotas(config, time, metrics)
+  private var startedProcessingRequests = false
   private var stoppedProcessingRequests = false
 
   /**
-   * Start the socket server
+   * Starts the socket server and creates all the Acceptors and the Processors. The Acceptors
+   * start listening at this stage so that the bound port is known when this method completes
+   * even when ephemeral ports are used. Acceptors and Processors are started if `startProcessingRequests`
+   * is true. If not, acceptors and processors are only started when [[kafka.network.SocketServer#startProcessingRequests()]]
+   * is invoked. Delayed starting of acceptors and 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. Incoming connections on this server are processed
+   * when processors start up and invoke [[org.apache.kafka.common.network.Selector#poll]].
+   *
+   * @param startProcessingRequests Flag indicating whether `Processor`s must be started.
    */
-  def startup() {
+  def startup(startProcessingRequests: Boolean = true): Unit = {
     this.synchronized {
+      createControlPlaneAcceptorAndProcessor(config.controlPlaneListener)
+      createDataPlaneAcceptorsAndProcessors(config.numNetworkThreads, config.dataPlaneListeners)
+      if (startProcessingRequests) {
+        this.startProcessingRequests()
+      }
+    }
+
+    newGauge(s"${DataPlaneMetricPrefix}NetworkProcessorAvgIdlePercent", () => SocketServer.this.synchronized {
+      val ioWaitRatioMetricNames = dataPlaneProcessors.values.asScala.iterator.map { p =>
+        metrics.metricName("io-wait-ratio", MetricsGroup, p.metricTags)
+      }
+      ioWaitRatioMetricNames.map { metricName =>
+        Option(metrics.metric(metricName)).fold(0.0)(m => Math.min(m.metricValue.asInstanceOf[Double], 1.0))
+      }.sum / dataPlaneProcessors.size
+    })
+    newGauge(s"${ControlPlaneMetricPrefix}NetworkProcessorAvgIdlePercent", () => SocketServer.this.synchronized {
+      val ioWaitRatioMetricName = controlPlaneProcessorOpt.map { p =>
+        metrics.metricName("io-wait-ratio", MetricsGroup, p.metricTags)
+      }
+      ioWaitRatioMetricName.map { metricName =>
+        Option(metrics.metric(metricName)).fold(0.0)(m => Math.min(m.metricValue.asInstanceOf[Double], 1.0))
+      }.getOrElse(Double.NaN)
+    })
+    newGauge("MemoryPoolAvailable", () => memoryPool.availableMemory)
+    newGauge("MemoryPoolUsed", () => memoryPool.size() - memoryPool.availableMemory)
+    newGauge(s"${DataPlaneMetricPrefix}ExpiredConnectionsKilledCount", () => SocketServer.this.synchronized {
+      val expiredConnectionsKilledCountMetricNames = dataPlaneProcessors.values.asScala.iterator.map { p =>
+        metrics.metricName("expired-connections-killed-count", MetricsGroup, p.metricTags)
+      }
+      expiredConnectionsKilledCountMetricNames.map { metricName =>
+        Option(metrics.metric(metricName)).fold(0.0)(m => m.metricValue.asInstanceOf[Double])
+      }.sum
+    })
+    newGauge(s"${ControlPlaneMetricPrefix}ExpiredConnectionsKilledCount", () => SocketServer.this.synchronized {
+      val expiredConnectionsKilledCountMetricNames = controlPlaneProcessorOpt.map { p =>
+        metrics.metricName("expired-connections-killed-count", MetricsGroup, p.metricTags)
+      }
+      expiredConnectionsKilledCountMetricNames.map { metricName =>
+        Option(metrics.metric(metricName)).fold(0.0)(m => m.metricValue.asInstanceOf[Double])
+      }.getOrElse(0.0)
+    })
+  }
 
-      connectionQuotas = new ConnectionQuotas(maxConnectionsPerIp, maxConnectionsPerIpOverrides)
+  /**
+   * Start processing requests and new connections. This method is used for delayed starting of
+   * all the acceptors and processors if [[kafka.network.SocketServer#startup]] was invoked with
+   * `startProcessingRequests=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.
+   *
+   * @param authorizerFutures Future per [[EndPoint]] used to wait before starting the processor
+   *                          corresponding to the [[EndPoint]]
+   */
+  def startProcessingRequests(authorizerFutures: Map[Endpoint, CompletableFuture[Void]] = Map.empty): Unit = {
+    info("Starting socket server acceptors and processors")
+    this.synchronized {
+      if (!startedProcessingRequests) {
+        startControlPlaneProcessorAndAcceptor(authorizerFutures)
+        startDataPlaneProcessorsAndAcceptors(authorizerFutures)
+        startedProcessingRequests = true
+      } else {
+        info("Socket server acceptors and processors already started")
+      }
+    }
+    info("Started socket server acceptors and processors")
+  }
 
-      val sendBufferSize = config.socketSendBufferBytes
-      val recvBufferSize = config.socketReceiveBufferBytes
-      val brokerId = config.brokerId
+  /**
+   * Starts processors of the provided acceptor and the acceptor itself.
+   *
+   * Before starting them, we ensure that authorizer has all the metadata to authorize
+   * requests on that endpoint by waiting on the provided future.
+   */
+  private def startAcceptorAndProcessors(threadPrefix: String,
+                                         endpoint: EndPoint,
+                                         acceptor: Acceptor,
+                                         authorizerFutures: Map[Endpoint, CompletableFuture[Void]] = Map.empty): Unit = {
+    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(threadPrefix)
+    debug(s"Start acceptor thread on listener ${endpoint.listenerName}")
+    if (!acceptor.isStarted()) {
+      KafkaThread.nonDaemon(
+        s"${threadPrefix}-kafka-socket-acceptor-${endpoint.listenerName}-${endpoint.securityProtocol}-${endpoint.port}",
+        acceptor
+      ).start()
+      acceptor.awaitStartup()
+    }
+    info(s"Started $threadPrefix acceptor and processor(s) for endpoint : ${endpoint.listenerName}")
+  }
 
-      var processorBeginIndex = 0
-      config.listeners.foreach { endpoint =>
-        val listenerName = endpoint.listenerName
-        val securityProtocol = endpoint.securityProtocol
-        val processorEndIndex = processorBeginIndex + numProcessorThreads
+  /**
+   * Starts processors of all the data-plane acceptors and all the acceptors of this server.
+   *
+   * We start inter-broker listener before other listeners. This allows authorization metadata for
+   * other listeners to be stored in Kafka topics in this cluster.
+   */
+  private def startDataPlaneProcessorsAndAcceptors(authorizerFutures: Map[Endpoint, CompletableFuture[Void]]): Unit = {
+    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.filter { case (k, _) => k != interBrokerListener }.values
+    orderedAcceptors.foreach { acceptor =>
+      val endpoint = acceptor.endPoint
+      startAcceptorAndProcessors(DataPlaneThreadPrefix, endpoint, acceptor, authorizerFutures)
+    }
+  }
 
-        for (i <- processorBeginIndex until processorEndIndex)
-          processors(i) = newProcessor(i, connectionQuotas, listenerName, securityProtocol, memoryPool)
+  /**
+   * Start the processor of control-plane acceptor and the acceptor of this server.
+   */
+  private def startControlPlaneProcessorAndAcceptor(authorizerFutures: Map[Endpoint, CompletableFuture[Void]]): Unit = {
+    controlPlaneAcceptorOpt.foreach { controlPlaneAcceptor =>
+      val endpoint = config.controlPlaneListener.get
+      startAcceptorAndProcessors(ControlPlaneThreadPrefix, endpoint, controlPlaneAcceptor, authorizerFutures)
+    }
+  }
 
-        val acceptor = new Acceptor(endpoint, sendBufferSize, recvBufferSize, brokerId,
-          processors.slice(processorBeginIndex, processorEndIndex), connectionQuotas)
-        acceptors.put(endpoint, acceptor)
-        KafkaThread.nonDaemon(s"kafka-socket-acceptor-$listenerName-$securityProtocol-${endpoint.port}", acceptor).start()
-        acceptor.awaitStartup()
+  private def endpoints = config.listeners.map(l => l.listenerName -> l).toMap
 
-        processorBeginIndex = processorEndIndex
-      }
+  private def createDataPlaneAcceptorsAndProcessors(dataProcessorsPerListener: Int,
+                                                    endpoints: Seq[EndPoint]): Unit = {
+    endpoints.foreach { endpoint =>
+      connectionQuotas.addListener(config, endpoint.listenerName)
+      val dataPlaneAcceptor = createAcceptor(endpoint, DataPlaneMetricPrefix)
+      addDataPlaneProcessors(dataPlaneAcceptor, endpoint, dataProcessorsPerListener)
+      dataPlaneAcceptors.put(endpoint, dataPlaneAcceptor)
+      info(s"Created data-plane acceptor and processors for endpoint : ${endpoint.listenerName}")
     }
+  }
 
-    newGauge("NetworkProcessorAvgIdlePercent",
-      new Gauge[Double] {
-        private val ioWaitRatioMetricNames = processors.map { p =>
-          metrics.metricName("io-wait-ratio", "socket-server-metrics", p.metricTags)
-        }
+  private def createControlPlaneAcceptorAndProcessor(endpointOpt: Option[EndPoint]): Unit = {
+    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, isPrivilegedListener = true)
+      controlPlaneAcceptorOpt = Some(controlPlaneAcceptor)
+      controlPlaneProcessorOpt = Some(controlPlaneProcessor)
+      val listenerProcessors = new ArrayBuffer[Processor]()
+      listenerProcessors += controlPlaneProcessor
+      controlPlaneRequestChannelOpt.foreach(_.addProcessor(controlPlaneProcessor))
+      nextProcessorId += 1
+      controlPlaneAcceptor.addProcessors(listenerProcessors, ControlPlaneThreadPrefix)
+      info(s"Created control-plane acceptor and processor for endpoint : ${endpoint.listenerName}")
+    }
+  }
 
-        def value = ioWaitRatioMetricNames.map { metricName =>
-          Option(metrics.metric(metricName)).fold(0.0)(_.value)
-        }.sum / totalProcessorThreads
-      }
-    )
-    newGauge("MemoryPoolAvailable",
-      new Gauge[Long] {
-        def value = memoryPool.availableMemory()
-      }
-    )
-    newGauge("MemoryPoolUsed",
-      new Gauge[Long] {
-        def value = memoryPool.size() - memoryPool.availableMemory()
-      }
-    )
-    info("Started " + acceptors.size + " acceptor threads")
+  private def createAcceptor(endPoint: EndPoint, metricPrefix: String) : Acceptor = {
+    val sendBufferSize = config.socketSendBufferBytes
+    val recvBufferSize = config.socketReceiveBufferBytes
+    val brokerId = config.brokerId
+    new Acceptor(endPoint, sendBufferSize, recvBufferSize, brokerId, connectionQuotas, metricPrefix, time)
   }
 
-  // register the processor threads for notification of responses
-  requestChannel.addResponseListener(id => processors(id).wakeup())
+  private def addDataPlaneProcessors(acceptor: Acceptor, endpoint: EndPoint, newProcessorsPerListener: Int): Unit = {
+    val listenerName = endpoint.listenerName
+    val securityProtocol = endpoint.securityProtocol
+    val listenerProcessors = new ArrayBuffer[Processor]()
+    val isPrivilegedListener = controlPlaneRequestChannelOpt.isEmpty && config.interBrokerListenerName == listenerName
+
+    for (_ <- 0 until newProcessorsPerListener) {
+      val processor = newProcessor(nextProcessorId, dataPlaneRequestChannel, connectionQuotas,
+        listenerName, securityProtocol, memoryPool, isPrivilegedListener)
+      listenerProcessors += processor
+      dataPlaneRequestChannel.addProcessor(processor)
+      nextProcessorId += 1
+    }
+    listenerProcessors.foreach(p => dataPlaneProcessors.put(p.id, p))
+    acceptor.addProcessors(listenerProcessors, DataPlaneThreadPrefix)
+  }
 
   /**
-    * Stop processing requests and new connections.
-    */
-  def stopProcessingRequests() = {
+   * Stop processing requests and new connections.
+   */
+  def stopProcessingRequests(): Unit = {
     info("Stopping socket server request processors")
     this.synchronized {
-      acceptors.values.foreach(_.shutdown)
-      processors.foreach(_.shutdown)
-      requestChannel.clear()
+      dataPlaneAcceptors.asScala.values.foreach(_.initiateShutdown())
+      dataPlaneAcceptors.asScala.values.foreach(_.awaitShutdown())
+      controlPlaneAcceptorOpt.foreach(_.initiateShutdown())
+      controlPlaneAcceptorOpt.foreach(_.awaitShutdown())
+      dataPlaneRequestChannel.clear()
+      controlPlaneRequestChannelOpt.foreach(_.clear())
       stoppedProcessingRequests = true
     }
     info("Stopped socket server request processors")
   }
 
+  def resizeThreadPool(oldNumNetworkThreads: Int, newNumNetworkThreads: Int): Unit = synchronized {
+    info(s"Resizing network thread pool size for each data-plane listener from $oldNumNetworkThreads to $newNumNetworkThreads")
+    if (newNumNetworkThreads > oldNumNetworkThreads) {
+      dataPlaneAcceptors.forEach { (endpoint, acceptor) =>
+        addDataPlaneProcessors(acceptor, endpoint, newNumNetworkThreads - oldNumNetworkThreads)
+      }
+    } else if (newNumNetworkThreads < oldNumNetworkThreads)
+      dataPlaneAcceptors.asScala.values.foreach(_.removeProcessors(oldNumNetworkThreads - newNumNetworkThreads, dataPlaneRequestChannel))
+  }
+
   /**
-    * Shutdown the socket server. If still processing requests, shutdown
-    * acceptors and processors first.
-    */
-  def shutdown() = {
+   * Shutdown the socket server. If still processing requests, shutdown
+   * acceptors and processors first.
+   */
+  def shutdown(): Unit = {
     info("Shutting down socket server")
     this.synchronized {
       if (!stoppedProcessingRequests)
         stopProcessingRequests()
-      requestChannel.shutdown()
+      dataPlaneRequestChannel.shutdown()
+      controlPlaneRequestChannelOpt.foreach(_.shutdown())
+      connectionQuotas.close()
     }
     info("Shutdown completed")
   }
 
   def boundPort(listenerName: ListenerName): Int = {
     try {
-      acceptors(endpoints(listenerName)).serverChannel.socket.getLocalPort
+      val acceptor = dataPlaneAcceptors.get(endpoints(listenerName))
+      if (acceptor != null) {
+        acceptor.serverChannel.socket.getLocalPort
+      } else {
+        controlPlaneAcceptorOpt.map (_.serverChannel.socket().getLocalPort).getOrElse(throw new KafkaException("Could not find listenerName : " + listenerName + " in data-plane or control-plane"))
+      }
     } catch {
-      case e: Exception => throw new KafkaException("Tried to check server's port before server was started or checked for port of non-existing protocol", e)
+      case e: Exception =>
+        throw new KafkaException("Tried to check server's port before server was started or checked for port of non-existing protocol", e)
+    }
+  }
+
+  def addListeners(listenersAdded: Seq[EndPoint]): Unit = synchronized {
+    info(s"Adding data-plane listeners for endpoints $listenersAdded")
+    createDataPlaneAcceptorsAndProcessors(config.numNetworkThreads, listenersAdded)
+    listenersAdded.foreach { endpoint =>
+      val acceptor = dataPlaneAcceptors.get(endpoint)
+      startAcceptorAndProcessors(DataPlaneThreadPrefix, endpoint, acceptor)
+    }
+  }
+
+  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 { acceptor =>
+        acceptor.initiateShutdown()
+        acceptor.awaitShutdown()
+      }
     }
   }
 
-  /* `protected` for test usage */
-  protected[network] def newProcessor(id: Int, connectionQuotas: ConnectionQuotas, listenerName: ListenerName,
-                                      securityProtocol: SecurityProtocol, memoryPool: MemoryPool): Processor = {
+  override def reconfigurableConfigs: Set[String] = SocketServer.ReconfigurableConfigs
+
+  override def validateReconfiguration(newConfig: KafkaConfig): Unit = {
+
+  }
+
+  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)
+    }
+    val maxConnectionRate = newConfig.maxConnectionCreationRate
+    if (maxConnectionRate != oldConfig.maxConnectionCreationRate) {
+      info(s"Updating broker-wide maxConnectionCreationRate: $maxConnectionRate")
+      connectionQuotas.updateBrokerMaxConnectionRate(maxConnectionRate)
+    }
+  }
+
+  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.forKeyValue { (endpoint, future) =>
+      if (endpoint.listenerName == Optional.of(acceptor.endPoint.listenerName.value))
+        future.join()
+    }
+  }
+
+  // `protected` for test usage
+  protected[network] def newProcessor(id: Int, requestChannel: RequestChannel, connectionQuotas: ConnectionQuotas, listenerName: ListenerName,
+                                      securityProtocol: SecurityProtocol, memoryPool: MemoryPool, isPrivilegedListener: Boolean): Processor = {
     new Processor(id,
       time,
       config.socketRequestMaxBytes,
       requestChannel,
       connectionQuotas,
       config.connectionsMaxIdleMs,
+      config.failedAuthenticationDelayMs,
       listenerName,
       securityProtocol,
       config,
       metrics,
       credentialProvider,
       memoryPool,
-      logContext
+      logContext,
+      isPrivilegedListener = isPrivilegedListener,
+      allowDisabledApis = allowDisabledApis
     )
   }
 
-  /* For test usage */
+  // For test usage
   private[network] def connectionCount(address: InetAddress): Int =
     Option(connectionQuotas).fold(0)(_.get(address))
 
-  /* For test usage */
-  private[network] def processor(index: Int): Processor = processors(index)
+  // For test usage
+  private[network] def dataPlaneProcessor(index: Int): Processor = dataPlaneProcessors.get(index)
+
+}
+
+object SocketServer {
+  val MetricsGroup = "socket-server-metrics"
+  val DataPlaneThreadPrefix = "data-plane"
+  val ControlPlaneThreadPrefix = "control-plane"
+  val DataPlaneMetricPrefix = ""
+  val ControlPlaneMetricPrefix = "ControlPlane"
+
+  val ReconfigurableConfigs = Set(
+    KafkaConfig.MaxConnectionsPerIpProp,
+    KafkaConfig.MaxConnectionsPerIpOverridesProp,
+    KafkaConfig.MaxConnectionsProp,
+    KafkaConfig.MaxConnectionCreationRateProp)
 
+  val ListenerReconfigurableConfigs = Set(KafkaConfig.MaxConnectionsProp, KafkaConfig.MaxConnectionCreationRateProp)
 }
 
 /**
@@ -214,14 +475,23 @@ private[kafka] abstract class AbstractServerThread(connectionQuotas: ConnectionQ
   def wakeup(): Unit
 
   /**
-   * Initiates a graceful shutdown by signaling to stop and waiting for the shutdown to complete
+   * Initiates a graceful shutdown by signaling to stop
    */
-  def shutdown(): Unit = {
-    alive.set(false)
-    wakeup()
-    shutdownLatch.await()
+  def initiateShutdown(): Unit = {
+    if (alive.getAndSet(false))
+      wakeup()
   }
 
+  /**
+   * Wait for the thread to completely shutdown
+   */
+  def awaitShutdown(): Unit = shutdownLatch.await
+
+  /**
+   * Returns true if the thread is completely started
+   */
+  def isStarted(): Boolean = startupLatch.getCount == 0
+
   /**
    * Wait for the thread to completely start up
    */
@@ -249,14 +519,18 @@ 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("Closing connection from " + channel.socket.getRemoteSocketAddress())
-      connectionQuotas.dec(channel.socket.getInetAddress)
-      CoreUtils.swallow(channel.socket().close(), this, Level.ERROR)
-      CoreUtils.swallow(channel.close(), this, Level.ERROR)
+      debug(s"Closing connection from ${channel.socket.getRemoteSocketAddress}")
+      connectionQuotas.dec(listenerName, channel.socket.getInetAddress)
+      closeSocket(channel)
     }
   }
+
+  protected def closeSocket(channel: SocketChannel): Unit = {
+    CoreUtils.swallow(channel.socket().close(), this, Level.ERROR)
+    CoreUtils.swallow(channel.close(), this, Level.ERROR)
+  }
 }
 
 /**
@@ -266,49 +540,80 @@ private[kafka] class Acceptor(val endPoint: EndPoint,
                               val sendBufferSize: Int,
                               val recvBufferSize: Int,
                               brokerId: Int,
-                              processors: Array[Processor],
-                              connectionQuotas: ConnectionQuotas) extends AbstractServerThread(connectionQuotas) with KafkaMetricsGroup {
+                              connectionQuotas: ConnectionQuotas,
+                              metricPrefix: String,
+                              time: Time) extends AbstractServerThread(connectionQuotas) with KafkaMetricsGroup {
 
   private val nioSelector = NSelector.open()
   val serverChannel = openServerSocket(endPoint.host, endPoint.port)
+  private val processors = new ArrayBuffer[Processor]()
+  private val processorsStarted = new AtomicBoolean
+  private val blockedPercentMeter = newMeter(s"${metricPrefix}AcceptorBlockedPercent",
+    "blocked time", TimeUnit.NANOSECONDS, Map(ListenerMetricTag -> endPoint.listenerName.value))
+  private var currentProcessorIndex = 0
+  private[network] val throttledSockets = new mutable.PriorityQueue[DelayedCloseSocket]()
+
+  private[network] case class DelayedCloseSocket(socket: SocketChannel, endThrottleTimeMs: Long) extends Ordered[DelayedCloseSocket] {
+    override def compare(that: DelayedCloseSocket): Int = endThrottleTimeMs compare that.endThrottleTimeMs
+  }
+
+  private[network] def addProcessors(newProcessors: Buffer[Processor], processorThreadPrefix: String): Unit = synchronized {
+    processors ++= newProcessors
+    if (processorsStarted.get)
+      startProcessors(newProcessors, processorThreadPrefix)
+  }
+
+  private[network] def startProcessors(processorThreadPrefix: String): Unit = synchronized {
+    if (!processorsStarted.getAndSet(true)) {
+      startProcessors(processors, processorThreadPrefix)
+    }
+  }
 
-  this.synchronized {
+  private def startProcessors(processors: Seq[Processor], processorThreadPrefix: String): Unit = synchronized {
     processors.foreach { processor =>
-      KafkaThread.nonDaemon(s"kafka-network-thread-$brokerId-${endPoint.listenerName}-${endPoint.securityProtocol}-${processor.id}",
-        processor).start()
+      KafkaThread.nonDaemon(
+        s"${processorThreadPrefix}-kafka-network-thread-$brokerId-${endPoint.listenerName}-${endPoint.securityProtocol}-${processor.id}",
+        processor
+      ).start()
+    }
+  }
+
+  private[network] def removeProcessors(removeCount: Int, requestChannel: RequestChannel): Unit = synchronized {
+    // Shutdown `removeCount` processors. Remove them from the processor list first so that no more
+    // connections are assigned. Shutdown the removed processors, closing the selector and its connections.
+    // The processors are then removed from `requestChannel` and any pending responses to these processors are dropped.
+    val toRemove = processors.takeRight(removeCount)
+    processors.remove(processors.size - removeCount, removeCount)
+    toRemove.foreach(_.initiateShutdown())
+    toRemove.foreach(_.awaitShutdown())
+    toRemove.foreach(processor => requestChannel.removeProcessor(processor.id))
+  }
+
+  override def initiateShutdown(): Unit = {
+    super.initiateShutdown()
+    synchronized {
+      processors.foreach(_.initiateShutdown())
+    }
+  }
+
+  override def awaitShutdown(): Unit = {
+    super.awaitShutdown()
+    synchronized {
+      processors.foreach(_.awaitShutdown())
     }
   }
 
   /**
    * Accept loop that checks for new connection attempts
    */
-  def run() {
+  def run(): Unit = {
     serverChannel.register(nioSelector, SelectionKey.OP_ACCEPT)
     startupComplete()
     try {
-      var currentProcessor = 0
       while (isRunning) {
         try {
-          val ready = nioSelector.select(500)
-          if (ready > 0) {
-            val keys = nioSelector.selectedKeys()
-            val iter = keys.iterator()
-            while (iter.hasNext && isRunning) {
-              try {
-                val key = iter.next
-                iter.remove()
-                if (key.isAcceptable)
-                  accept(key, processors(currentProcessor))
-                else
-                  throw new IllegalStateException("Unrecognized key state for acceptor thread.")
-
-                // round robin to the next processor thread
-                currentProcessor = (currentProcessor + 1) % processors.length
-              } catch {
-                case e: Throwable => error("Error while accepting connection", e)
-              }
-            }
-          }
+          acceptNewConnections()
+          closeThrottledConnections()
         }
         catch {
           // We catch all the throwables to prevent the acceptor thread from exiting on exceptions due
@@ -319,19 +624,21 @@ private[kafka] class Acceptor(val endPoint: EndPoint,
         }
       }
     } finally {
-      debug("Closing server socket and selector.")
+      debug("Closing server socket, selector, and any throttled sockets.")
       CoreUtils.swallow(serverChannel.close(), this, Level.ERROR)
       CoreUtils.swallow(nioSelector.close(), this, Level.ERROR)
+      throttledSockets.foreach(throttledSocket => closeSocket(throttledSocket.socket))
+      throttledSockets.clear()
       shutdownComplete()
     }
   }
 
-  /*
+  /**
    * Create a server socket to listen for connections on.
    */
   private def openServerSocket(host: String, port: Int): ServerSocketChannel = {
     val socketAddress =
-      if(host == null || host.trim.isEmpty)
+      if (host == null || host.trim.isEmpty)
         new InetSocketAddress(port)
       else
         new InetSocketAddress(host, port)
@@ -342,52 +649,130 @@ private[kafka] class Acceptor(val endPoint: EndPoint,
 
     try {
       serverChannel.socket.bind(socketAddress)
-      info("Awaiting socket connections on %s:%d.".format(socketAddress.getHostString, serverChannel.socket.getLocalPort))
+      info(s"Awaiting socket connections on ${socketAddress.getHostString}:${serverChannel.socket.getLocalPort}.")
     } catch {
       case e: SocketException =>
-        throw new KafkaException("Socket server failed to bind to %s:%d: %s.".format(socketAddress.getHostString, port, e.getMessage), e)
+        throw new KafkaException(s"Socket server failed to bind to ${socketAddress.getHostString}:$port: ${e.getMessage}.", e)
     }
     serverChannel
   }
 
-  /*
+  /**
+   * Listen for new connections and assign accepted connections to processors using round-robin.
+   */
+  private def acceptNewConnections(): Unit = {
+    val ready = nioSelector.select(500)
+    if (ready > 0) {
+      val keys = nioSelector.selectedKeys()
+      val iter = keys.iterator()
+      while (iter.hasNext && isRunning) {
+        try {
+          val key = iter.next
+          iter.remove()
+
+          if (key.isAcceptable) {
+            accept(key).foreach { socketChannel =>
+              // Assign the channel to the next processor (using round-robin) to which the
+              // channel can be added without blocking. If newConnections queue is full on
+              // all processors, block until the last one is able to accept a connection.
+              var retriesLeft = synchronized(processors.length)
+              var processor: Processor = null
+              do {
+                retriesLeft -= 1
+                processor = synchronized {
+                  // adjust the index (if necessary) and retrieve the processor atomically for
+                  // correct behaviour in case the number of processors is reduced dynamically
+                  currentProcessorIndex = currentProcessorIndex % processors.length
+                  processors(currentProcessorIndex)
+                }
+                currentProcessorIndex += 1
+              } while (!assignNewConnection(socketChannel, processor, retriesLeft == 0))
+            }
+          } else
+            throw new IllegalStateException("Unrecognized key state for acceptor thread.")
+        } catch {
+          case e: Throwable => error("Error while accepting connection", e)
+        }
+      }
+    }
+  }
+
+  /**
    * Accept a new connection
    */
-  def accept(key: SelectionKey, processor: Processor) {
+  private def accept(key: SelectionKey): Option[SocketChannel] = {
     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)
       if (sendBufferSize != Selectable.USE_DEFAULT_BUFFER_SIZE)
         socketChannel.socket().setSendBufferSize(sendBufferSize)
-
-      debug("Accepted connection from %s on %s and assigned it to processor %d, sendBufferSize [actual|requested]: [%d|%d] recvBufferSize [actual|requested]: [%d|%d]"
-            .format(socketChannel.socket.getRemoteSocketAddress, socketChannel.socket.getLocalSocketAddress, processor.id,
-                  socketChannel.socket.getSendBufferSize, sendBufferSize,
-                  socketChannel.socket.getReceiveBufferSize, recvBufferSize))
-
-      processor.accept(socketChannel)
+      Some(socketChannel)
     } catch {
       case e: TooManyConnectionsException =>
-        info("Rejected connection from %s, address already has the configured maximum of %d connections.".format(e.ip, e.count))
-        close(socketChannel)
+        info(s"Rejected connection from ${e.ip}, address already has the configured maximum of ${e.count} connections.")
+        close(endPoint.listenerName, socketChannel)
+        None
+      case e: ConnectionThrottledException =>
+        val ip = socketChannel.socket.getInetAddress
+        debug(s"Delaying closing of connection from $ip for ${e.throttleTimeMs} ms")
+        val endThrottleTimeMs = e.startThrottleTimeMs + e.throttleTimeMs
+        throttledSockets += DelayedCloseSocket(socketChannel, endThrottleTimeMs)
+        None
+    }
+  }
+
+  /**
+   * Close sockets for any connections that have been throttled.
+   */
+  private def closeThrottledConnections(): Unit = {
+    val timeMs = time.milliseconds
+    while (throttledSockets.headOption.exists(_.endThrottleTimeMs < timeMs)) {
+      val closingSocket = throttledSockets.dequeue()
+      debug(s"Closing socket from ip ${closingSocket.socket.getRemoteAddress}")
+      closeSocket(closingSocket.socket)
     }
   }
 
+  private def assignNewConnection(socketChannel: SocketChannel, processor: Processor, mayBlock: Boolean): Boolean = {
+    if (processor.accept(socketChannel, mayBlock, blockedPercentMeter)) {
+      debug(s"Accepted connection from ${socketChannel.socket.getRemoteSocketAddress} on" +
+        s" ${socketChannel.socket.getLocalSocketAddress} and assigned it to processor ${processor.id}," +
+        s" sendBufferSize [actual|requested]: [${socketChannel.socket.getSendBufferSize}|$sendBufferSize]" +
+        s" recvBufferSize [actual|requested]: [${socketChannel.socket.getReceiveBufferSize}|$recvBufferSize]")
+      true
+    } else
+      false
+  }
+
   /**
    * Wakeup the thread for selection.
    */
   @Override
-  def wakeup = nioSelector.wakeup()
+  def wakeup(): Unit = nioSelector.wakeup()
 
 }
 
+private[kafka] object Processor {
+  val IdlePercentMetricName = "IdlePercent"
+  val NetworkProcessorMetricTag = "networkProcessor"
+  val ListenerMetricTag = "listener"
+
+  val ConnectionQueueSize = 20
+}
+
 /**
  * Thread that processes all requests from a single connection. There are N of these running in parallel
  * each of which has its own selector
+ *
+ * @param isPrivilegedListener The privileged listener flag is used as one factor to determine whether
+ *                             a certain request is forwarded or not. When the control plane is defined,
+ *                             the control plane processor would be fellow broker's choice for sending
+ *                             forwarding requests; if the control plane is not defined, the processor
+ *                             relying on the inter broker listener would be acting as the privileged listener.
  */
 private[kafka] class Processor(val id: Int,
                                time: Time,
@@ -395,13 +780,17 @@ private[kafka] class Processor(val id: Int,
                                requestChannel: RequestChannel,
                                connectionQuotas: ConnectionQuotas,
                                connectionsMaxIdleMs: Long,
+                               failedAuthenticationDelayMs: Int,
                                listenerName: ListenerName,
                                securityProtocol: SecurityProtocol,
                                config: KafkaConfig,
                                metrics: Metrics,
                                credentialProvider: CredentialProvider,
                                memoryPool: MemoryPool,
-                               logContext: LogContext) extends AbstractServerThread(connectionQuotas) with KafkaMetricsGroup {
+                               logContext: LogContext,
+                               connectionQueueSize: Int = ConnectionQueueSize,
+                               isPrivilegedListener: Boolean = false,
+                               allowDisabledApis: Boolean = false) extends AbstractServerThread(connectionQuotas) with KafkaMetricsGroup {
 
   private object ConnectionId {
     def fromString(s: String): Option[ConnectionId] = s.split("-") match {
@@ -418,46 +807,64 @@ private[kafka] class Processor(val id: Int,
     override def toString: String = s"$localHost:$localPort-$remoteHost:$remotePort-$index"
   }
 
-  private val newConnections = new ConcurrentLinkedQueue[SocketChannel]()
+  private val newConnections = new ArrayBlockingQueue[SocketChannel](connectionQueueSize)
   private val inflightResponses = mutable.Map[String, RequestChannel.Response]()
+  private val responseQueue = new LinkedBlockingDeque[RequestChannel.Response]()
+
   private[kafka] val metricTags = mutable.LinkedHashMap(
-    "listener" -> listenerName.value,
-    "networkProcessor" -> id.toString
+    ListenerMetricTag -> listenerName.value,
+    NetworkProcessorMetricTag -> id.toString
   ).asJava
 
-  newGauge("IdlePercent",
-    new Gauge[Double] {
-      def value = {
-        Option(metrics.metric(metrics.metricName("io-wait-ratio", "socket-server-metrics", metricTags))).fold(0.0)(_.value)
-      }
-    },
+  newGauge(IdlePercentMetricName, () => {
+    Option(metrics.metric(metrics.metricName("io-wait-ratio", MetricsGroup, metricTags))).fold(0.0)(m =>
+      Math.min(m.metricValue.asInstanceOf[Double], 1.0))
+  },
     // for compatibility, only add a networkProcessor tag to the Yammer Metrics alias (the equivalent Selector metric
     // also includes the listener name)
-    Map("networkProcessor" -> id.toString)
+    Map(NetworkProcessorMetricTag -> id.toString)
   )
 
+  val expiredConnectionsKilledCount = new CumulativeSum()
+  private val expiredConnectionsKilledCountMetricName = metrics.metricName("expired-connections-killed-count", MetricsGroup, metricTags)
+  metrics.addMetric(expiredConnectionsKilledCountMetricName, expiredConnectionsKilledCount)
+
   private val selector = createSelector(
-      ChannelBuilders.serverChannelBuilder(listenerName, securityProtocol, config, credentialProvider.credentialCache))
+    ChannelBuilders.serverChannelBuilder(listenerName,
+      listenerName == config.interBrokerListenerName,
+      securityProtocol,
+      config,
+      credentialProvider.credentialCache,
+      credentialProvider.tokenCache,
+      time,
+      logContext))
   // Visible to override for testing
-  protected[network] def createSelector(channelBuilder: ChannelBuilder): KSelector = new KSelector(
-    maxRequestSize,
-    connectionsMaxIdleMs,
-    metrics,
-    time,
-    "socket-server",
-    metricTags,
-    false,
-    true,
-    channelBuilder,
-    memoryPool,
-    logContext)
+  protected[network] def createSelector(channelBuilder: ChannelBuilder): KSelector = {
+    channelBuilder match {
+      case reconfigurable: Reconfigurable => config.addReconfigurable(reconfigurable)
+      case _ =>
+    }
+    new KSelector(
+      maxRequestSize,
+      connectionsMaxIdleMs,
+      failedAuthenticationDelayMs,
+      metrics,
+      time,
+      "socket-server",
+      metricTags,
+      false,
+      true,
+      channelBuilder,
+      memoryPool,
+      logContext)
+  }
 
   // Connection ids have the format `localAddr:localPort-remoteAddr:remotePort-index`. The index is a
   // non-negative incrementing value that ensures that even if remotePort is reused after a connection is
   // 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) {
@@ -470,6 +877,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
@@ -481,47 +889,59 @@ private[kafka] class Processor(val id: Int,
         }
       }
     } finally {
-      debug("Closing selector - processor " + id)
+      debug(s"Closing selector - processor $id")
       CoreUtils.swallow(closeAll(), this, Level.ERROR)
       shutdownComplete()
     }
   }
 
-  private def processException(errorMessage: String, throwable: Throwable) {
+  private[network] 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)
+      error(s"Closing socket for $channelId because of error", throwable)
       close(channelId)
     }
     processException(errorMessage, throwable)
   }
 
-  private def processNewResponses() {
-    var curr: RequestChannel.Response = null
-    while ({curr = requestChannel.receiveResponse(id); curr != null}) {
-      val channelId = curr.request.context.connectionId
+  private def processNewResponses(): Unit = {
+    var currentResponse: RequestChannel.Response = null
+    while ({currentResponse = dequeueResponse(); currentResponse != null}) {
+      val channelId = currentResponse.request.context.connectionId
       try {
-        curr.responseAction match {
-          case RequestChannel.NoOpAction =>
+        currentResponse match {
+          case response: NoOpResponse =>
             // There is no response to send to the client, we need to read more pipelined requests
             // that are sitting in the server's socket buffer
-            updateRequestMetrics(curr)
-            trace("Socket server received empty response to send, registering for read: " + curr)
-            openOrClosingChannel(channelId).foreach(c => selector.unmute(c.id))
-          case RequestChannel.SendAction =>
-            val responseSend = curr.responseSend.getOrElse(
-              throw new IllegalStateException(s"responseSend must be defined for SendAction, response: $curr"))
-            sendResponse(curr, responseSend)
-          case RequestChannel.CloseConnectionAction =>
-            updateRequestMetrics(curr)
+            updateRequestMetrics(response)
+            trace(s"Socket server received empty response to send, registering for read: $response")
+            // Try unmuting the channel. If there was no quota violation and the channel has not been throttled,
+            // it will be unmuted immediately. If the channel has been throttled, it will be unmuted only if the
+            // throttling delay has already passed by now.
+            handleChannelMuteEvent(channelId, ChannelMuteEvent.RESPONSE_SENT)
+            tryUnmuteChannel(channelId)
+
+          case response: SendResponse =>
+            sendResponse(response, response.responseSend)
+          case response: CloseConnectionResponse =>
+            updateRequestMetrics(response)
             trace("Closing socket connection actively according to the response code.")
             close(channelId)
+          case _: StartThrottlingResponse =>
+            handleChannelMuteEvent(channelId, ChannelMuteEvent.THROTTLE_STARTED)
+          case _: EndThrottlingResponse =>
+            // Try unmuting the channel. The channel will be unmuted only if the response has already been sent out to
+            // the client.
+            handleChannelMuteEvent(channelId, ChannelMuteEvent.THROTTLE_ENDED)
+            tryUnmuteChannel(channelId)
+          case _ =>
+            throw new IllegalArgumentException(s"Unknown response type: ${currentResponse.getClass}")
         }
       } catch {
         case e: Throwable =>
@@ -530,8 +950,8 @@ private[kafka] class Processor(val id: Int,
     }
   }
 
-  /* `protected` for test usage */
-  protected[network] def sendResponse(response: RequestChannel.Response, responseSend: Send) {
+  // `protected` for test usage
+  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
@@ -543,33 +963,88 @@ private[kafka] class Processor(val id: Int,
     // removed from the Selector after discarding any pending staged receives.
     // `openOrClosingChannel` can be None if the selector closed the connection because it was idle for too long
     if (openOrClosingChannel(connectionId).isDefined) {
-      selector.send(responseSend)
+      selector.send(new NetworkSend(connectionId, responseSend))
       inflightResponses += (connectionId -> response)
     }
   }
 
-  private def poll() {
-    try selector.poll(300)
+  private def poll(): Unit = {
+    val pollTimeout = if (newConnections.isEmpty) 300 else 0
+    try selector.poll(pollTimeout)
     catch {
       case e @ (_: IllegalStateException | _: IOException) =>
         // The exception is not re-thrown and any completed sends/receives/connections/disconnections
         // from this poll will be processed.
-        error(s"Processor $id poll failed due to illegal state or IO exception")
+        error(s"Processor $id poll failed", e)
     }
   }
 
-  private def processCompletedReceives() {
-    selector.completedReceives.asScala.foreach { receive =>
+  protected def parseRequestHeader(buffer: ByteBuffer): RequestHeader = {
+    val header = RequestHeader.parse(buffer)
+    if (header.apiKey.isEnabled || allowDisabledApis) {
+      header
+    } else {
+      throw new InvalidRequestException("Received request for disabled api key " + header.apiKey)
+    }
+  }
+
+  private def processCompletedReceives(): Unit = {
+    selector.completedReceives.forEach { receive =>
       try {
         openOrClosingChannel(receive.source) match {
           case Some(channel) =>
-            val header = RequestHeader.parse(receive.payload)
-            val context = new RequestContext(header, receive.source, channel.socketAddress,
-              channel.principal, listenerName, securityProtocol)
-            val req = new RequestChannel.Request(processor = id, context = context,
-              startTimeNanos = time.nanoseconds, memoryPool, receive.payload, requestChannel.metrics)
-            requestChannel.sendRequest(req)
-            selector.mute(receive.source)
+            val header = parseRequestHeader(receive.payload)
+            if (header.apiKey == ApiKeys.SASL_HANDSHAKE && channel.maybeBeginServerReauthentication(receive,
+              () => time.nanoseconds()))
+              trace(s"Begin re-authentication: $channel")
+            else {
+              val nowNanos = time.nanoseconds()
+              if (channel.serverAuthenticationSessionExpired(nowNanos)) {
+                // 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
+                val context = new RequestContext(header, connectionId, channel.socketAddress,
+                  channel.principal, listenerName, securityProtocol,
+                  channel.channelMetadataRegistry.clientInformation, isPrivilegedListener, channel.principalSerde)
+
+                var req = new RequestChannel.Request(processor = id, context = context,
+                  startTimeNanos = nowNanos, memoryPool, receive.payload, requestChannel.metrics, None)
+
+                if (req.header.apiKey == ApiKeys.ENVELOPE) {
+                  // Override the request context with the forwarded request context.
+                  // The envelope's context will be preserved in the forwarded context
+
+                  req = parseForwardedPrincipal(req, channel.principalSerde.asScala) match {
+                    case Some(forwardedPrincipal) =>
+                      buildForwardedRequestContext(req, forwardedPrincipal)
+
+                    case None =>
+                      val envelopeResponse = new EnvelopeResponse(Errors.PRINCIPAL_DESERIALIZATION_FAILURE)
+                      sendEnvelopeResponse(req, envelopeResponse)
+                      null
+                  }
+                }
+
+                if (req != null) {
+                  // KIP-511: ApiVersionsRequest is intercepted here to catch the client software name
+                  // and version. It is done here to avoid wiring things up to the api layer.
+                  if (header.apiKey == ApiKeys.API_VERSIONS) {
+                    val apiVersionsRequest = req.body[ApiVersionsRequest]
+                    if (apiVersionsRequest.isValid) {
+                      channel.channelMetadataRegistry.registerClientInformation(new ClientInformation(
+                        apiVersionsRequest.data.clientSoftwareName,
+                        apiVersionsRequest.data.clientSoftwareVersion))
+                    }
+                  }
+                  requestChannel.sendRequest(req)
+                  selector.mute(connectionId)
+                  handleChannelMuteEvent(connectionId, ChannelMuteEvent.REQUEST_RECEIVED)
+                }
+              }
+            }
           case None =>
             // This should never happen since completed receives are processed immediately after `poll()`
             throw new IllegalStateException(s"Channel ${receive.source} removed from selector before processing completed receive")
@@ -581,44 +1056,122 @@ private[kafka] class Processor(val id: Int,
           processChannelException(receive.source, s"Exception while processing request from ${receive.source}", e)
       }
     }
+    selector.clearCompletedReceives()
+  }
+
+  private def sendEnvelopeResponse(
+    envelopeRequest: RequestChannel.Request,
+    envelopeResponse: EnvelopeResponse
+  ): Unit = {
+    val envelopResponseSend = envelopeRequest.context.buildResponseSend(envelopeResponse)
+    enqueueResponse(new RequestChannel.SendResponse(
+      envelopeRequest,
+      envelopResponseSend,
+      None,
+      None
+    ))
+  }
+
+  private def parseForwardedPrincipal(
+    envelopeRequest: RequestChannel.Request,
+    principalSerde: Option[KafkaPrincipalSerde]
+  ): Option[KafkaPrincipal] = {
+    val envelope = envelopeRequest.body[EnvelopeRequest]
+    try {
+      principalSerde.map { serde =>
+        serde.deserialize(envelope.requestPrincipal())
+      }
+    } catch {
+      case e: Exception =>
+        warn(s"Failed to deserialize principal from envelope request $envelope", e)
+        None
+    }
+  }
+
+  private def buildForwardedRequestContext(
+    envelopeRequest: RequestChannel.Request,
+    forwardedPrincipal: KafkaPrincipal
+  ): RequestChannel.Request = {
+    val envelope = envelopeRequest.body[EnvelopeRequest]
+    val forwardedRequestBuffer = envelope.requestData.duplicate()
+    val forwardedHeader = parseRequestHeader(forwardedRequestBuffer)
+    val forwardedClientAddress = InetAddress.getByAddress(envelope.clientAddress)
+
+    val forwardedContext = new RequestContext(
+      forwardedHeader,
+      envelopeRequest.context.connectionId,
+      forwardedClientAddress,
+      forwardedPrincipal,
+      listenerName,
+      securityProtocol,
+      ClientInformation.EMPTY,
+      isPrivilegedListener
+    )
+
+    new RequestChannel.Request(
+      processor = id,
+      context = forwardedContext,
+      startTimeNanos = envelopeRequest.startTimeNanos,
+      memoryPool,
+      forwardedRequestBuffer,
+      requestChannel.metrics,
+      Some(envelopeRequest)
+    )
   }
 
-  private def processCompletedSends() {
-    selector.completedSends.asScala.foreach { send =>
+  private def processCompletedSends(): Unit = {
+    selector.completedSends.forEach { send =>
       try {
-        val resp = inflightResponses.remove(send.destination).getOrElse {
-          throw new IllegalStateException(s"Send for ${send.destination} completed, but not in `inflightResponses`")
+        val response = inflightResponses.remove(send.destinationId).getOrElse {
+          throw new IllegalStateException(s"Send for ${send.destinationId} completed, but not in `inflightResponses`")
         }
-        updateRequestMetrics(resp)
-        selector.unmute(send.destination)
+        updateRequestMetrics(response)
+
+        // Invoke send completion callback
+        response.onComplete.foreach(onComplete => onComplete(send))
+
+        // Try unmuting the channel. If there was no quota violation and the channel has not been throttled,
+        // it will be unmuted immediately. If the channel has been throttled, it will unmuted only if the throttling
+        // delay has already passed by now.
+        handleChannelMuteEvent(send.destinationId, ChannelMuteEvent.RESPONSE_SENT)
+        tryUnmuteChannel(send.destinationId)
       } catch {
-        case e: Throwable => processChannelException(send.destination,
-            s"Exception while processing completed send to ${send.destination}", e)
+        case e: Throwable => processChannelException(send.destinationId,
+          s"Exception while processing completed send to ${send.destinationId}", e)
       }
     }
+    selector.clearCompletedSends()
   }
 
-  private def updateRequestMetrics(response: RequestChannel.Response) {
+  private def updateRequestMetrics(response: RequestChannel.Response): Unit = {
     val request = response.request
     val networkThreadTimeNanos = openOrClosingChannel(request.context.connectionId).fold(0L)(_.getAndResetNetworkThreadTimeNanos())
     request.updateRequestMetrics(networkThreadTimeNanos, response)
   }
 
-  private def processDisconnected() {
-    selector.disconnected.keySet.asScala.foreach { connectionId =>
+  private def processDisconnected(): Unit = {
+    selector.disconnected.keySet.forEach { connectionId =>
       try {
         val remoteHost = ConnectionId.fromString(connectionId).getOrElse {
           throw new IllegalStateException(s"connectionId has unexpected format: $connectionId")
         }.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`
@@ -631,7 +1184,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))
@@ -641,26 +1194,44 @@ private[kafka] class Processor(val id: Int,
   /**
    * Queue up a new connection for reading
    */
-  def accept(socketChannel: SocketChannel) {
-    newConnections.add(socketChannel)
-    wakeup()
+  def accept(socketChannel: SocketChannel,
+             mayBlock: Boolean,
+             acceptorIdlePercentMeter: com.yammer.metrics.core.Meter): Boolean = {
+    val accepted = {
+      if (newConnections.offer(socketChannel))
+        true
+      else if (mayBlock) {
+        val startNs = time.nanoseconds
+        newConnections.put(socketChannel)
+        acceptorIdlePercentMeter.mark(time.nanoseconds() - startNs)
+        true
+      } else
+        false
+    }
+    if (accepted)
+      wakeup()
+    accepted
   }
 
   /**
-   * Register any new connections that have been queued up
+   * Register any new connections that have been queued up. The number of connections processed
+   * in each iteration is limited to ensure that traffic and connection close notifications of
+   * existing channels are handled promptly.
    */
-  private def configureNewConnections() {
-    while (!newConnections.isEmpty) {
+  private def configureNewConnections(): Unit = {
+    var connectionsProcessed = 0
+    while (connectionsProcessed < connectionQueueSize && !newConnections.isEmpty) {
       val channel = newConnections.poll()
       try {
         debug(s"Processor $id listening to new connection from ${channel.socket.getRemoteSocketAddress}")
         selector.register(connectionId(channel.socket), channel)
+        connectionsProcessed += 1
       } catch {
         // We explicitly catch all exceptions and close the socket to avoid a socket leak.
         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)
       }
     }
@@ -669,11 +1240,15 @@ private[kafka] class Processor(val id: Int,
   /**
    * Close the selector and all open connections
    */
-  private def closeAll() {
-    selector.channels.asScala.foreach { channel =>
+  private def closeAll(): Unit = {
+    while (!newConnections.isEmpty) {
+      newConnections.poll().close()
+    }
+    selector.channels.forEach { channel =>
       close(channel.id)
     }
     selector.close()
+    removeMetric(IdlePercentMetricName, Map(NetworkProcessorMetricTag -> id.toString))
   }
 
   // 'protected` to allow override for testing
@@ -687,46 +1262,235 @@ private[kafka] class Processor(val id: Int,
     connId
   }
 
+  private[network] def enqueueResponse(response: RequestChannel.Response): Unit = {
+    responseQueue.put(response)
+    wakeup()
+  }
+
+  private def dequeueResponse(): RequestChannel.Response = {
+    val response = responseQueue.poll()
+    if (response != null)
+      response.request.responseDequeueTimeNanos = Time.SYSTEM.nanoseconds
+    response
+  }
+
+  private[network] def responseQueueSize = responseQueue.size
+
   // Only for testing
   private[network] def inflightResponseCount: Int = inflightResponses.size
 
   // Visible for testing
   // Only methods that are safe to call on a disconnected channel should be invoked on 'openOrClosingChannel'.
   private[network] def openOrClosingChannel(connectionId: String): Option[KafkaChannel] =
-     Option(selector.channel(connectionId)).orElse(Option(selector.closingChannel(connectionId)))
+    Option(selector.channel(connectionId)).orElse(Option(selector.closingChannel(connectionId)))
+
+  // Indicate the specified channel that the specified channel mute-related event has happened so that it can change its
+  // mute state.
+  private def handleChannelMuteEvent(connectionId: String, event: ChannelMuteEvent): Unit = {
+    openOrClosingChannel(connectionId).foreach(c => c.handleChannelMuteEvent(event))
+  }
+
+  private def tryUnmuteChannel(connectionId: String) = {
+    openOrClosingChannel(connectionId).foreach(c => selector.unmute(c.id))
+  }
 
   /* For test usage */
   private[network] def channel(connectionId: String): Option[KafkaChannel] =
     Option(selector.channel(connectionId))
 
-  // Visible for testing
-  private[network] def numStagedReceives(connectionId: String): Int =
-    openOrClosingChannel(connectionId).map(c => selector.numStagedReceives(c)).getOrElse(0)
-
   /**
    * Wakeup the thread for selection.
    */
-  @Override
-  def wakeup = selector.wakeup()
+  override def wakeup() = selector.wakeup()
+
+  override def initiateShutdown(): Unit = {
+    super.initiateShutdown()
+    removeMetric("IdlePercent", Map("networkProcessor" -> id.toString))
+    metrics.removeMetric(expiredConnectionsKilledCountMetricName)
+  }
+}
+
+/**
+ * Interface for connection quota configuration. Connection quotas can be configured at the
+ * broker, listener or IP level.
+ */
+sealed trait ConnectionQuotaEntity {
+  def sensorName: String
+  def metricName: String
+  def sensorExpiration: Long
+  def metricTags: Map[String, String]
+}
+
+object ConnectionQuotas {
+  private val InactiveSensorExpirationTimeSeconds = TimeUnit.HOURS.toSeconds(1)
+  private val ConnectionRateSensorName = "Connection-Accept-Rate"
+  private val ConnectionRateMetricName = "connection-accept-rate"
+  private val IpMetricTag = "ip"
+  private val ListenerThrottlePrefix = ""
+  private val IpThrottlePrefix = "ip-"
+
+  private case class ListenerQuotaEntity(listenerName: String) extends ConnectionQuotaEntity {
+    override def sensorName: String = s"$ConnectionRateSensorName-$listenerName"
+    override def sensorExpiration: Long = Long.MaxValue
+    override def metricName: String = ConnectionRateMetricName
+    override def metricTags: Map[String, String] = Map(ListenerMetricTag -> listenerName)
+  }
+
+  private case object BrokerQuotaEntity extends ConnectionQuotaEntity {
+    override def sensorName: String = ConnectionRateSensorName
+    override def sensorExpiration: Long = Long.MaxValue
+    override def metricName: String = s"broker-$ConnectionRateMetricName"
+    override def metricTags: Map[String, String] = Map.empty
+  }
 
+  private case class IpQuotaEntity(ip: InetAddress) extends ConnectionQuotaEntity {
+    override def sensorName: String = s"$ConnectionRateSensorName-${ip.getHostAddress}"
+    override def sensorExpiration: Long = InactiveSensorExpirationTimeSeconds
+    override def metricName: String = ConnectionRateMetricName
+    override def metricTags: Map[String, String] = Map(IpMetricTag -> ip.getHostAddress)
+  }
 }
 
-class ConnectionQuotas(val defaultMax: Int, overrideQuotas: Map[String, Int]) {
+class ConnectionQuotas(config: KafkaConfig, time: Time, metrics: Metrics) extends Logging with AutoCloseable {
 
-  private val overrides = 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 interBrokerListenerName = config.interBrokerListenerName
   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[network] val maxConnectionsPerListener = mutable.Map[ListenerName, ListenerConnectionQuota]()
+  @volatile private var totalCount = 0
+  // updates to defaultConnectionRatePerIp or connectionRatePerIp must be synchronized on `counts`
+  @volatile private var defaultConnectionRatePerIp = DynamicConfig.Ip.DefaultConnectionCreationRate
+  private val connectionRatePerIp = new ConcurrentHashMap[InetAddress, Int]()
+  // sensor that tracks broker-wide connection creation rate and limit (quota)
+  private val brokerConnectionRateSensor = getOrCreateConnectionRateQuotaSensor(config.maxConnectionCreationRate, BrokerQuotaEntity)
+  private val maxThrottleTimeMs = TimeUnit.SECONDS.toMillis(config.quotaWindowSizeSeconds.toLong)
+
+  def inc(listenerName: ListenerName, address: InetAddress, acceptorBlockedPercentMeter: com.yammer.metrics.core.Meter): Unit = {
     counts.synchronized {
+      waitForConnectionSlot(listenerName, acceptorBlockedPercentMeter)
+
+      recordIpConnectionMaybeThrottle(listenerName, address)
       val count = counts.getOrElseUpdate(address, 0)
       counts.put(address, count + 1)
-      val max = overrides.getOrElse(address, defaultMax)
+      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 dec(address: InetAddress) {
+  private[network] def updateMaxConnectionsPerIp(maxConnectionsPerIp: Int): Unit = {
+    defaultMaxConnectionsPerIp = maxConnectionsPerIp
+  }
+
+  private[network] def updateMaxConnectionsPerIpOverride(overrideQuotas: Map[String, Int]): Unit = {
+    maxConnectionsPerIpOverrides = overrideQuotas.map { case (host, count) => (InetAddress.getByName(host), count) }
+  }
+
+  private[network] def updateBrokerMaxConnections(maxConnections: Int): Unit = {
+    counts.synchronized {
+      brokerMaxConnections = maxConnections
+      counts.notifyAll()
+    }
+  }
+
+  private[network] def updateBrokerMaxConnectionRate(maxConnectionRate: Int): Unit = {
+    // if there is a connection waiting on the rate throttle delay, we will let it wait the original delay even if
+    // the rate limit increases, because it is just one connection per listener and the code is simpler that way
+    updateConnectionRateQuota(maxConnectionRate, BrokerQuotaEntity)
+  }
+
+  /**
+   * Update the connection rate quota for a given IP and updates quota configs for updated IPs.
+   * If an IP is given, metric config will be updated only for the given IP, otherwise
+   * all metric configs will be checked and updated if required.
+   *
+   * @param ip ip to update or default if None
+   * @param maxConnectionRate new connection rate, or resets entity to default if None
+   */
+  def updateIpConnectionRateQuota(ip: Option[InetAddress], maxConnectionRate: Option[Int]): Unit = synchronized {
+    def isIpConnectionRateMetric(metricName: MetricName) = {
+      metricName.name == ConnectionRateMetricName &&
+      metricName.group == MetricsGroup &&
+      metricName.tags.containsKey(IpMetricTag)
+    }
+
+    def shouldUpdateQuota(metric: KafkaMetric, quotaLimit: Int) = {
+      quotaLimit != metric.config.quota.bound
+    }
+
+    ip match {
+      case Some(address) =>
+        // synchronize on counts to ensure reading an IP connection rate quota and creating a quota config is atomic
+        counts.synchronized {
+          maxConnectionRate match {
+            case Some(rate) =>
+              info(s"Updating max connection rate override for $address to $rate")
+              connectionRatePerIp.put(address, rate)
+            case None =>
+              info(s"Removing max connection rate override for $address")
+              connectionRatePerIp.remove(address)
+          }
+        }
+        updateConnectionRateQuota(connectionRateForIp(address), IpQuotaEntity(address))
+      case None =>
+        // synchronize on counts to ensure reading an IP connection rate quota and creating a quota config is atomic
+        counts.synchronized {
+          defaultConnectionRatePerIp = maxConnectionRate.getOrElse(DynamicConfig.Ip.DefaultConnectionCreationRate)
+        }
+        info(s"Updated default max IP connection rate to $defaultConnectionRatePerIp")
+        metrics.metrics.forEach { (metricName, metric) =>
+          if (isIpConnectionRateMetric(metricName)) {
+            val quota = connectionRateForIp(InetAddress.getByName(metricName.tags.get(IpMetricTag)))
+            if (shouldUpdateQuota(metric, quota)) {
+              debug(s"Updating existing connection rate quota config for ${metricName.tags} to $quota")
+              metric.config(rateQuotaMetricConfig(quota))
+            }
+          }
+        }
+    }
+  }
+
+  // Visible for testing
+  def connectionRateForIp(ip: InetAddress): Int = {
+    connectionRatePerIp.getOrDefault(ip, defaultConnectionRatePerIp)
+  }
+
+  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)
+        newListenerQuota.configure(config.valuesWithPrefixOverride(listenerName.configPrefix))
+      }
+      counts.notifyAll()
+    }
+  }
+
+  private[network] def removeListener(config: KafkaConfig, listenerName: ListenerName): Unit = {
+    counts.synchronized {
+      maxConnectionsPerListener.remove(listenerName).foreach { listenerQuota =>
+        listenerCounts.remove(listenerName)
+        // once listener is removed from maxConnectionsPerListener, no metrics will be recorded into listener's sensor
+        // so it is safe to remove sensor here
+        listenerQuota.close()
+        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"))
@@ -734,6 +1498,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
     }
   }
 
@@ -741,6 +1518,261 @@ 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 {
+      val startThrottleTimeMs = time.milliseconds
+      val throttleTimeMs = math.max(recordConnectionAndGetThrottleTimeMs(listenerName, startThrottleTimeMs), 0)
+
+      if (throttleTimeMs > 0 || !connectionSlotAvailable(listenerName)) {
+        val startNs = time.nanoseconds
+        val endThrottleTimeMs = startThrottleTimeMs + throttleTimeMs
+        var remainingThrottleTimeMs = throttleTimeMs
+        do {
+          counts.wait(remainingThrottleTimeMs)
+          remainingThrottleTimeMs = math.max(endThrottleTimeMs - time.milliseconds, 0)
+        } while (remainingThrottleTimeMs > 0 || !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 =
+    interBrokerListenerName == listenerName && listenerCounts.size > 1
+
+  private def maxListenerConnections(listenerName: ListenerName): Int =
+    maxConnectionsPerListener.get(listenerName).map(_.maxConnections).getOrElse(Int.MaxValue)
+
+  /**
+   * Calculates the delay needed to bring the observed connection creation rate to listener-level limit or to broker-wide
+   * limit, whichever the longest. The delay is capped to the quota window size defined by QuotaWindowSizeSecondsProp
+   *
+   * @param listenerName listener for which calculate the delay
+   * @param timeMs current time in milliseconds
+   * @return delay in milliseconds
+   */
+  private def recordConnectionAndGetThrottleTimeMs(listenerName: ListenerName, timeMs: Long): Long = {
+    def recordAndGetListenerThrottleTime(minThrottleTimeMs: Int): Int = {
+      maxConnectionsPerListener
+        .get(listenerName)
+        .map { listenerQuota =>
+          val listenerThrottleTimeMs = recordAndGetThrottleTimeMs(listenerQuota.connectionRateSensor, timeMs)
+          val throttleTimeMs = math.max(minThrottleTimeMs, listenerThrottleTimeMs)
+          // record throttle time due to hitting connection rate quota
+          if (throttleTimeMs > 0) {
+            listenerQuota.listenerConnectionRateThrottleSensor.record(throttleTimeMs.toDouble, timeMs)
+          }
+          throttleTimeMs
+        }
+        .getOrElse(0)
+    }
+
+    if (protectedListener(listenerName)) {
+      recordAndGetListenerThrottleTime(0)
+    } else {
+      val brokerThrottleTimeMs = recordAndGetThrottleTimeMs(brokerConnectionRateSensor, timeMs)
+      recordAndGetListenerThrottleTime(brokerThrottleTimeMs)
+    }
+  }
+
+  /**
+   * Record IP throttle time on the corresponding listener. To avoid over-recording listener/broker connection rate, we
+   * also un-record the listener and broker connection if the IP gets throttled.
+   *
+   * @param listenerName listener to un-record connection
+   * @param throttleMs IP throttle time to record for listener
+   * @param timeMs current time in milliseconds
+   */
+  private def updateListenerMetrics(listenerName: ListenerName, throttleMs: Long, timeMs: Long): Unit = {
+    if (!protectedListener(listenerName)) {
+      brokerConnectionRateSensor.record(-1.0, timeMs, false)
+    }
+    maxConnectionsPerListener
+      .get(listenerName)
+      .foreach { listenerQuota =>
+        listenerQuota.ipConnectionRateThrottleSensor.record(throttleMs.toDouble, timeMs)
+        listenerQuota.connectionRateSensor.record(-1.0, timeMs, false)
+      }
+  }
+
+  /**
+   * Calculates the delay needed to bring the observed connection creation rate to the IP limit.
+   * If the connection would cause an IP quota violation, un-record the connection for both IP,
+   * listener, and broker connection rate and throw a ConnectionThrottledException. Calls to
+   * this function must be performed with the counts lock to ensure that reading the IP
+   * connection rate quota and creating the sensor's metric config is atomic.
+   *
+   * @param listenerName listener to unrecord connection if throttled
+   * @param address ip address to record connection
+   */
+  private def recordIpConnectionMaybeThrottle(listenerName: ListenerName, address: InetAddress): Unit = {
+    val connectionRateQuota = connectionRateForIp(address)
+    val quotaEnabled = connectionRateQuota != DynamicConfig.Ip.UnlimitedConnectionCreationRate
+    if (quotaEnabled) {
+      val sensor = getOrCreateConnectionRateQuotaSensor(connectionRateQuota, IpQuotaEntity(address))
+      val timeMs = time.milliseconds
+      val throttleMs = recordAndGetThrottleTimeMs(sensor, timeMs)
+      if (throttleMs > 0) {
+        trace(s"Throttling $address for $throttleMs ms")
+        // unrecord the connection since we won't accept the connection
+        sensor.record(-1.0, timeMs, false)
+        updateListenerMetrics(listenerName, throttleMs, timeMs)
+        throw new ConnectionThrottledException(address, timeMs, throttleMs)
+      }
+    }
+  }
+
+  /**
+   * Records a new connection into a given connection acceptance rate sensor 'sensor' and returns throttle time
+   * in milliseconds if quota got violated
+   * @param sensor sensor to record connection
+   * @param timeMs current time in milliseconds
+   * @return throttle time in milliseconds if quota got violated, otherwise 0
+   */
+  private def recordAndGetThrottleTimeMs(sensor: Sensor, timeMs: Long): Int = {
+    try {
+      sensor.record(1.0, timeMs)
+      0
+    } catch {
+      case e: QuotaViolationException =>
+        val throttleTimeMs = QuotaUtils.boundedThrottleTime(e, maxThrottleTimeMs, timeMs).toInt
+        debug(s"Quota violated for sensor (${sensor.name}). Delay time: $throttleTimeMs ms")
+        throttleTimeMs
+    }
+  }
+
+  /**
+   * Creates sensor for tracking the connection creation rate and corresponding connection rate quota for a given
+   * listener or broker-wide, if listener is not provided.
+   * @param quotaLimit connection creation rate quota
+   * @param connectionQuotaEntity entity to create the sensor for
+   */
+  private def getOrCreateConnectionRateQuotaSensor(quotaLimit: Int, connectionQuotaEntity: ConnectionQuotaEntity): Sensor = {
+    Option(metrics.getSensor(connectionQuotaEntity.sensorName)).getOrElse {
+      val sensor = metrics.sensor(
+        connectionQuotaEntity.sensorName,
+        rateQuotaMetricConfig(quotaLimit),
+        connectionQuotaEntity.sensorExpiration
+      )
+      sensor.add(connectionRateMetricName(connectionQuotaEntity), new Rate, null)
+      sensor
+    }
+  }
+
+  /**
+   * Updates quota configuration for a given connection quota entity
+   */
+  private def updateConnectionRateQuota(quotaLimit: Int, connectionQuotaEntity: ConnectionQuotaEntity): Unit = {
+    Option(metrics.metric(connectionRateMetricName(connectionQuotaEntity))).foreach { metric =>
+      metric.config(rateQuotaMetricConfig(quotaLimit))
+      info(s"Updated ${connectionQuotaEntity.metricName} max connection creation rate to $quotaLimit")
+    }
+  }
+
+  private def connectionRateMetricName(connectionQuotaEntity: ConnectionQuotaEntity): MetricName = {
+    metrics.metricName(
+      connectionQuotaEntity.metricName,
+      MetricsGroup,
+      s"Tracking rate of accepting new connections (per second)",
+      connectionQuotaEntity.metricTags.asJava)
+  }
+
+  private def rateQuotaMetricConfig(quotaLimit: Int): MetricConfig = {
+    new MetricConfig()
+      .timeWindow(config.quotaWindowSizeSeconds.toLong, TimeUnit.SECONDS)
+      .samples(config.numQuotaSamples)
+      .quota(new Quota(quotaLimit, true))
+  }
+
+  def close(): Unit = {
+    metrics.removeSensor(brokerConnectionRateSensor.name)
+    maxConnectionsPerListener.values.foreach(_.close())
+  }
+
+  class ListenerConnectionQuota(lock: Object, listener: ListenerName) extends ListenerReconfigurable with AutoCloseable {
+    @volatile private var _maxConnections = Int.MaxValue
+    private[network] val connectionRateSensor = getOrCreateConnectionRateQuotaSensor(Int.MaxValue, ListenerQuotaEntity(listener.value))
+    private[network] val listenerConnectionRateThrottleSensor = createConnectionRateThrottleSensor(ListenerThrottlePrefix)
+    private[network] val ipConnectionRateThrottleSensor = createConnectionRateThrottleSensor(IpThrottlePrefix)
+
+    def maxConnections: Int = _maxConnections
+
+    override def listenerName(): ListenerName = listener
+
+    override def configure(configs: util.Map[String, _]): Unit = {
+      _maxConnections = maxConnections(configs)
+      updateConnectionRateQuota(maxConnectionCreationRate(configs), ListenerQuotaEntity(listener.value))
+    }
+
+    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(s"Invalid ${KafkaConfig.MaxConnectionsProp} $value")
+
+      val rate = maxConnectionCreationRate(configs)
+      if (rate <= 0)
+        throw new ConfigException(s"Invalid ${KafkaConfig.MaxConnectionCreationRateProp} $rate")
+    }
+
+    override def reconfigure(configs: util.Map[String, _]): Unit = {
+      lock.synchronized {
+        _maxConnections = maxConnections(configs)
+        updateConnectionRateQuota(maxConnectionCreationRate(configs), ListenerQuotaEntity(listener.value))
+        lock.notifyAll()
+      }
+    }
+
+    def close(): Unit = {
+      metrics.removeSensor(connectionRateSensor.name)
+      metrics.removeSensor(listenerConnectionRateThrottleSensor.name)
+      metrics.removeSensor(ipConnectionRateThrottleSensor.name)
+    }
+
+    private def maxConnections(configs: util.Map[String, _]): Int = {
+      Option(configs.get(KafkaConfig.MaxConnectionsProp)).map(_.toString.toInt).getOrElse(Int.MaxValue)
+    }
+
+    private def maxConnectionCreationRate(configs: util.Map[String, _]): Int = {
+      Option(configs.get(KafkaConfig.MaxConnectionCreationRateProp)).map(_.toString.toInt).getOrElse(Int.MaxValue)
+    }
+
+    /**
+     * Creates sensor for tracking the average throttle time on this listener due to hitting broker/listener connection
+     * rate or IP connection rate quota. The average is out of all throttle times > 0, which is consistent with the
+     * bandwidth and request quota throttle time metrics.
+     */
+    private def createConnectionRateThrottleSensor(throttlePrefix: String): Sensor = {
+      val sensor = metrics.sensor(s"${throttlePrefix}ConnectionRateThrottleTime-${listener.value}")
+      val metricName = metrics.metricName(s"${throttlePrefix}connection-accept-throttle-time",
+        MetricsGroup,
+        "Tracking average throttle-time, out of non-zero throttle times, per listener",
+        Map(ListenerMetricTag -> listener.value).asJava)
+      sensor.add(metricName, new Avg)
+      sensor
+    }
+  }
 }
 
-class TooManyConnectionsException(val ip: InetAddress, val count: Int) extends KafkaException("Too many connections from %s (maximum = %d)".format(ip, count))
+class TooManyConnectionsException(val ip: InetAddress, val count: Int) extends KafkaException(s"Too many connections from $ip (maximum = $count)")
+
+class ConnectionThrottledException(val ip: InetAddress, val startThrottleTimeMs: Long, val throttleTimeMs: Long)
+  extends KafkaException(s"$ip throttled for $throttleTimeMs")
diff --git a/core/src/main/scala/kafka/producer/BaseProducer.scala b/core/src/main/scala/kafka/producer/BaseProducer.scala
deleted file mode 100644
index 83d9aa75133b2..0000000000000
--- a/core/src/main/scala/kafka/producer/BaseProducer.scala
+++ /dev/null
@@ -1,74 +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.producer
-
-import java.util.Properties
-
-// A base producer used whenever we need to have options for both old and new producers;
-// this class will be removed once we fully rolled out 0.9
-@deprecated("This trait has been deprecated and will be removed in a future release. " +
-            "Please use org.apache.kafka.clients.producer.KafkaProducer instead.", "0.10.0.0")
-trait BaseProducer {
-  def send(topic: String, key: Array[Byte], value: Array[Byte])
-  def close()
-}
-
-@deprecated("This class has been deprecated and will be removed in a future release. " +
-            "Please use org.apache.kafka.clients.producer.KafkaProducer instead.", "0.10.0.0")
-class NewShinyProducer(producerProps: Properties) extends BaseProducer {
-  import org.apache.kafka.clients.producer.{KafkaProducer, ProducerRecord}
-  import org.apache.kafka.clients.producer.internals.ErrorLoggingCallback
-
-  // decide whether to send synchronously based on producer properties
-  val sync = producerProps.getProperty("producer.type", "async").equals("sync")
-
-  val producer = new KafkaProducer[Array[Byte],Array[Byte]](producerProps)
-
-  override def send(topic: String, key: Array[Byte], value: Array[Byte]) {
-    val record = new ProducerRecord[Array[Byte],Array[Byte]](topic, key, value)
-    if(sync) {
-      this.producer.send(record).get()
-    } else {
-      this.producer.send(record,
-        new ErrorLoggingCallback(topic, key, value, false))
-    }
-  }
-
-  override def close() {
-    this.producer.close()
-  }
-}
-
-@deprecated("This class has been deprecated and will be removed in a future release. " +
-            "Please use org.apache.kafka.clients.producer.KafkaProducer instead.", "0.10.0.0")
-class OldProducer(producerProps: Properties) extends BaseProducer {
-
-  // default to byte array partitioner
-  if (producerProps.getProperty("partitioner.class") == null)
-    producerProps.setProperty("partitioner.class", classOf[kafka.producer.ByteArrayPartitioner].getName)
-  val producer = new kafka.producer.Producer[Array[Byte], Array[Byte]](new ProducerConfig(producerProps))
-
-  override def send(topic: String, key: Array[Byte], value: Array[Byte]) {
-    this.producer.send(new KeyedMessage[Array[Byte], Array[Byte]](topic, key, value))
-  }
-
-  override def close() {
-    this.producer.close()
-  }
-}
-
diff --git a/core/src/main/scala/kafka/producer/BrokerPartitionInfo.scala b/core/src/main/scala/kafka/producer/BrokerPartitionInfo.scala
deleted file mode 100644
index e77a50cec9270..0000000000000
--- a/core/src/main/scala/kafka/producer/BrokerPartitionInfo.scala
+++ /dev/null
@@ -1,105 +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.producer
-
-import org.apache.kafka.common.protocol.Errors
-
-import collection.mutable.HashMap
-import kafka.api.TopicMetadata
-import kafka.common.KafkaException
-import kafka.utils.Logging
-import kafka.client.ClientUtils
-
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.10.0.0")
-class BrokerPartitionInfo(producerConfig: ProducerConfig,
-                          producerPool: ProducerPool,
-                          topicPartitionInfo: HashMap[String, TopicMetadata])
-        extends Logging {
-  val brokerList = producerConfig.brokerList
-  val brokers = ClientUtils.parseBrokerList(brokerList)
-
-  /**
-   * Return a sequence of (brokerId, numPartitions).
-   * @param topic the topic for which this information is to be returned
-   * @return a sequence of (brokerId, numPartitions). Returns a zero-length
-   * sequence if no brokers are available.
-   */
-  def getBrokerPartitionInfo(topic: String, correlationId: Int): Seq[PartitionAndLeader] = {
-    debug("Getting broker partition info for topic %s".format(topic))
-    // check if the cache has metadata for this topic
-    val topicMetadata = topicPartitionInfo.get(topic)
-    val metadata: TopicMetadata =
-      topicMetadata match {
-        case Some(m) => m
-        case None =>
-          // refresh the topic metadata cache
-          updateInfo(Set(topic), correlationId)
-          val topicMetadata = topicPartitionInfo.get(topic)
-          topicMetadata match {
-            case Some(m) => m
-            case None => throw new KafkaException("Failed to fetch topic metadata for topic: " + topic)
-          }
-      }
-    val partitionMetadata = metadata.partitionsMetadata
-    if(partitionMetadata.isEmpty) {
-      if(metadata.error != Errors.NONE) {
-        throw new KafkaException(metadata.error.exception)
-      } else {
-        throw new KafkaException("Topic metadata %s has empty partition metadata and no error code".format(metadata))
-      }
-    }
-    partitionMetadata.map { m =>
-      m.leader match {
-        case Some(leader) =>
-          debug("Partition [%s,%d] has leader %d".format(topic, m.partitionId, leader.id))
-          new PartitionAndLeader(topic, m.partitionId, Some(leader.id))
-        case None =>
-          debug("Partition [%s,%d] does not have a leader yet".format(topic, m.partitionId))
-          new PartitionAndLeader(topic, m.partitionId, None)
-      }
-    }.sortWith((s, t) => s.partitionId < t.partitionId)
-  }
-
-  /**
-   * It updates the cache by issuing a get topic metadata request to a random broker.
-   * @param topics the topics for which the metadata is to be fetched
-   */
-  def updateInfo(topics: Set[String], correlationId: Int) {
-    var topicsMetadata: Seq[TopicMetadata] = Nil
-    val topicMetadataResponse = ClientUtils.fetchTopicMetadata(topics, brokers, producerConfig, correlationId)
-    topicsMetadata = topicMetadataResponse.topicsMetadata
-    // throw partition specific exception
-    topicsMetadata.foreach(tmd =>{
-      trace("Metadata for topic %s is %s".format(tmd.topic, tmd))
-      if(tmd.error == Errors.NONE) {
-        topicPartitionInfo.put(tmd.topic, tmd)
-      } else
-        warn("Error while fetching metadata [%s] for topic [%s]: %s ".format(tmd, tmd.topic, tmd.error.exception.getClass))
-      tmd.partitionsMetadata.foreach(pmd =>{
-        if (pmd.error != Errors.NONE && pmd.error == Errors.LEADER_NOT_AVAILABLE) {
-          warn("Error while fetching metadata %s for topic partition [%s,%d]: [%s]".format(pmd, tmd.topic, pmd.partitionId,
-            pmd.error.exception.getClass))
-        } // any other error code (e.g. ReplicaNotAvailable) can be ignored since the producer does not need to access the replica and isr metadata
-      })
-    })
-    producerPool.updateProducer(topicsMetadata)
-  }
-
-}
-
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.10.0.0")
-case class PartitionAndLeader(topic: String, partitionId: Int, leaderBrokerIdOpt: Option[Int])
diff --git a/core/src/main/scala/kafka/producer/ByteArrayPartitioner.scala b/core/src/main/scala/kafka/producer/ByteArrayPartitioner.scala
deleted file mode 100755
index 7848456288f56..0000000000000
--- a/core/src/main/scala/kafka/producer/ByteArrayPartitioner.scala
+++ /dev/null
@@ -1,30 +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.producer
-
-
-import kafka.utils._
-import org.apache.kafka.common.utils.Utils
-
-@deprecated("This class has been deprecated and will be removed in a future release. " +
-            "It has been replaced by org.apache.kafka.clients.producer.internals.DefaultPartitioner.", "0.10.0.0")
-class ByteArrayPartitioner(props: VerifiableProperties = null) extends Partitioner {
-  def partition(key: Any, numPartitions: Int): Int = {
-    Utils.abs(java.util.Arrays.hashCode(key.asInstanceOf[Array[Byte]])) % numPartitions
-  }
-}
diff --git a/core/src/main/scala/kafka/producer/DefaultPartitioner.scala b/core/src/main/scala/kafka/producer/DefaultPartitioner.scala
deleted file mode 100755
index f793811016b79..0000000000000
--- a/core/src/main/scala/kafka/producer/DefaultPartitioner.scala
+++ /dev/null
@@ -1,30 +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.producer
-
-
-import kafka.utils._
-import org.apache.kafka.common.utils.Utils
-
-@deprecated("This class has been deprecated and will be removed in a future release. " +
-            "It has been replaced by org.apache.kafka.clients.producer.internals.DefaultPartitioner.", "0.10.0.0")
-class DefaultPartitioner(props: VerifiableProperties = null) extends Partitioner {
-  def partition(key: Any, numPartitions: Int): Int = {
-    Utils.abs(key.hashCode) % numPartitions
-  }
-}
diff --git a/core/src/main/scala/kafka/producer/KeyedMessage.scala b/core/src/main/scala/kafka/producer/KeyedMessage.scala
deleted file mode 100644
index 84ea2328de553..0000000000000
--- a/core/src/main/scala/kafka/producer/KeyedMessage.scala
+++ /dev/null
@@ -1,44 +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.producer
-
-/**
- * A topic, key, and value.
- * If a partition key is provided it will override the key for the purpose of partitioning but will not be stored.
- */
-@deprecated("This class has been deprecated and will be removed in a future release. " +
-            "Please use org.apache.kafka.clients.producer.ProducerRecord instead.", "0.10.0.0")
-case class KeyedMessage[K, V](topic: String, key: K, partKey: Any, message: V) {
-  if(topic == null)
-    throw new IllegalArgumentException("Topic cannot be null.")
-  
-  def this(topic: String, message: V) = this(topic, null.asInstanceOf[K], null, message)
-  
-  def this(topic: String, key: K, message: V) = this(topic, key, key, message)
-  
-  def partitionKey = {
-    if(partKey != null)
-      partKey
-    else if(hasKey)
-      key
-    else
-      null  
-  }
-  
-  def hasKey = key != null
-}
diff --git a/core/src/main/scala/kafka/producer/Partitioner.scala b/core/src/main/scala/kafka/producer/Partitioner.scala
deleted file mode 100644
index 5d24692db5c56..0000000000000
--- a/core/src/main/scala/kafka/producer/Partitioner.scala
+++ /dev/null
@@ -1,35 +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.producer
-
-/**
- * A partitioner controls the mapping between user-provided keys and kafka partitions. Users can implement a custom
- * partitioner to change this mapping.
- * 
- * Implementations will be constructed via reflection and are required to have a constructor that takes a single 
- * VerifiableProperties instance--this allows passing configuration properties into the partitioner implementation.
- */
-@deprecated("This trait has been deprecated and will be removed in a future release. " +
-            "Please use org.apache.kafka.clients.producer.Partitioner instead.", "0.10.0.0")
-trait Partitioner {
-  /**
-   * Uses the key to calculate a partition bucket id for routing
-   * the data to the appropriate broker partition
-   * @return an integer between 0 and numPartitions-1
-   */
-  def partition(key: Any, numPartitions: Int): Int
-}
diff --git a/core/src/main/scala/kafka/producer/Producer.scala b/core/src/main/scala/kafka/producer/Producer.scala
deleted file mode 100755
index d6cf4c892a534..0000000000000
--- a/core/src/main/scala/kafka/producer/Producer.scala
+++ /dev/null
@@ -1,142 +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.producer
-
-import java.util.concurrent.atomic.AtomicBoolean
-import java.util.concurrent.{LinkedBlockingQueue, TimeUnit}
-
-import kafka.common.{AppInfo, QueueFullException}
-import kafka.metrics._
-import kafka.producer.async.{DefaultEventHandler, EventHandler, ProducerSendThread}
-import kafka.serializer.Encoder
-import kafka.utils._
-
-@deprecated("This class has been deprecated and will be removed in a future release. " +
-            "Please use org.apache.kafka.clients.producer.KafkaProducer instead.", "0.10.0.0")
-class Producer[K,V](val config: ProducerConfig,
-                    private val eventHandler: EventHandler[K,V])  // only for unit testing
-  extends Logging {
-
-  private val hasShutdown = new AtomicBoolean(false)
-  private val queue = new LinkedBlockingQueue[KeyedMessage[K,V]](config.queueBufferingMaxMessages)
-
-  private var sync: Boolean = true
-  private var producerSendThread: ProducerSendThread[K,V] = null
-  private val lock = new Object()
-
-  config.producerType match {
-    case "sync" =>
-    case "async" =>
-      sync = false
-      producerSendThread = new ProducerSendThread[K,V]("ProducerSendThread-" + config.clientId,
-                                                       queue,
-                                                       eventHandler,
-                                                       config.queueBufferingMaxMs,
-                                                       config.batchNumMessages,
-                                                       config.clientId)
-      producerSendThread.start()
-  }
-
-  private val producerTopicStats = ProducerTopicStatsRegistry.getProducerTopicStats(config.clientId)
-
-  KafkaMetricsReporter.startReporters(config.props)
-  AppInfo.registerInfo()
-
-  def this(config: ProducerConfig) =
-    this(config,
-         new DefaultEventHandler[K,V](config,
-                                      CoreUtils.createObject[Partitioner](config.partitionerClass, config.props),
-                                      CoreUtils.createObject[Encoder[V]](config.serializerClass, config.props),
-                                      CoreUtils.createObject[Encoder[K]](config.keySerializerClass, config.props),
-                                      new ProducerPool(config)))
-
-  /**
-   * Sends the data, partitioned by key to the topic using either the
-   * synchronous or the asynchronous producer
-   * @param messages the producer data object that encapsulates the topic, key and message data
-   */
-  def send(messages: KeyedMessage[K,V]*) {
-    lock synchronized {
-      if (hasShutdown.get)
-        throw new ProducerClosedException
-      recordStats(messages)
-      if (sync)
-        eventHandler.handle(messages)
-      else
-        asyncSend(messages)
-    }
-  }
-
-
-  private def recordStats(messages: Seq[KeyedMessage[K,V]]) {
-    for (message <- messages) {
-      producerTopicStats.getProducerTopicStats(message.topic).messageRate.mark()
-      producerTopicStats.getProducerAllTopicsStats.messageRate.mark()
-    }
-  }
-
-  private def asyncSend(messages: Seq[KeyedMessage[K,V]]) {
-    for (message <- messages) {
-      val added = config.queueEnqueueTimeoutMs match {
-        case 0  =>
-          queue.offer(message)
-        case _  =>
-          try {
-            if (config.queueEnqueueTimeoutMs < 0) {
-              queue.put(message)
-              true
-            } else {
-              queue.offer(message, config.queueEnqueueTimeoutMs, TimeUnit.MILLISECONDS)
-            }
-          }
-          catch {
-            case _: InterruptedException =>
-              false
-          }
-      }
-      if(!added) {
-        producerTopicStats.getProducerTopicStats(message.topic).droppedMessageRate.mark()
-        producerTopicStats.getProducerAllTopicsStats.droppedMessageRate.mark()
-        throw new QueueFullException("Event queue is full of unsent messages, could not send event: " + message.toString)
-      }else {
-        trace("Added to send queue an event: " + message.toString)
-        trace("Remaining queue size: " + queue.remainingCapacity)
-      }
-    }
-  }
-
-  /**
-   * Close API to close the producer pool connections to all Kafka brokers. Also closes
-   * the zookeeper client connection if one exists
-   */
-  def close() = {
-    lock synchronized {
-      val canShutdown = hasShutdown.compareAndSet(false, true)
-      if(canShutdown) {
-        info("Shutting down producer")
-        val startTime = System.nanoTime()
-        KafkaMetricsGroup.removeAllProducerMetrics(config.clientId)
-        if (producerSendThread != null)
-          producerSendThread.shutdown
-        eventHandler.close()
-        info("Producer shutdown completed in " + (System.nanoTime() - startTime) / 1000000 + " ms")
-      }
-    }
-  }
-}
-
-
diff --git a/core/src/main/scala/kafka/producer/ProducerClosedException.scala b/core/src/main/scala/kafka/producer/ProducerClosedException.scala
deleted file mode 100644
index 4f2f7316dabf6..0000000000000
--- a/core/src/main/scala/kafka/producer/ProducerClosedException.scala
+++ /dev/null
@@ -1,22 +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.producer
-
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.10.0.0")
-class ProducerClosedException() extends RuntimeException("producer already closed") {
-}
diff --git a/core/src/main/scala/kafka/producer/ProducerConfig.scala b/core/src/main/scala/kafka/producer/ProducerConfig.scala
deleted file mode 100755
index c2715d06fc380..0000000000000
--- a/core/src/main/scala/kafka/producer/ProducerConfig.scala
+++ /dev/null
@@ -1,121 +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.producer
-
-import async.AsyncProducerConfig
-import java.util.Properties
-import kafka.utils.{CoreUtils, VerifiableProperties}
-import kafka.message.NoCompressionCodec
-import kafka.common.{InvalidConfigException, Config}
-
-@deprecated("This object has been deprecated and will be removed in a future release. " +
-            "Please use org.apache.kafka.clients.producer.ProducerConfig instead.", "0.10.0.0")
-object ProducerConfig extends Config {
-  def validate(config: ProducerConfig) {
-    validateClientId(config.clientId)
-    validateBatchSize(config.batchNumMessages, config.queueBufferingMaxMessages)
-    validateProducerType(config.producerType)
-  }
-
-  def validateClientId(clientId: String) {
-    validateChars("client.id", clientId)
-  }
-
-  def validateBatchSize(batchSize: Int, queueSize: Int) {
-    if (batchSize > queueSize)
-      throw new InvalidConfigException("Batch size = " + batchSize + " can't be larger than queue size = " + queueSize)
-  }
-
-  def validateProducerType(producerType: String) {
-    producerType match {
-      case "sync" =>
-      case "async"=>
-      case _ => throw new InvalidConfigException("Invalid value " + producerType + " for producer.type, valid values are sync/async")
-    }
-  }
-}
-
-@deprecated("This class has been deprecated and will be removed in a future release. " +
-            "Please use org.apache.kafka.clients.producer.ProducerConfig instead.", "0.10.0.0")
-class ProducerConfig private (val props: VerifiableProperties)
-        extends AsyncProducerConfig with SyncProducerConfigShared {
-  import ProducerConfig._
-
-  def this(originalProps: Properties) {
-    this(new VerifiableProperties(originalProps))
-    props.verify()
-  }
-
-  /** This is for bootstrapping and the producer will only use it for getting metadata
-   * (topics, partitions and replicas). The socket connections for sending the actual data
-   * will be established based on the broker information returned in the metadata. The
-   * format is host1:port1,host2:port2, and the list can be a subset of brokers or
-   * a VIP pointing to a subset of brokers.
-   */
-  val brokerList = props.getString("metadata.broker.list")
-
-  /** the partitioner class for partitioning events amongst sub-topics */
-  val partitionerClass = props.getString("partitioner.class", "kafka.producer.DefaultPartitioner")
-
-  /** this parameter specifies whether the messages are sent asynchronously *
-   * or not. Valid values are - async for asynchronous send                 *
-   *                            sync for synchronous send                   */
-  val producerType = props.getString("producer.type", "sync")
-
-  /**
-   * This parameter allows you to specify the compression codec for all data generated *
-   * by this producer. The default is NoCompressionCodec
-   */
-  val compressionCodec = props.getCompressionCodec("compression.codec", NoCompressionCodec)
-
-  /** This parameter allows you to set whether compression should be turned *
-   *  on for particular topics
-   *
-   *  If the compression codec is anything other than NoCompressionCodec,
-   *
-   *    Enable compression only for specified topics if any
-   *
-   *    If the list of compressed topics is empty, then enable the specified compression codec for all topics
-   *
-   *  If the compression codec is NoCompressionCodec, compression is disabled for all topics
-   */
-  val compressedTopics = CoreUtils.parseCsvList(props.getString("compressed.topics", null))
-
-  /** The leader may be unavailable transiently, which can fail the sending of a message.
-    *  This property specifies the number of retries when such failures occur.
-    */
-  val messageSendMaxRetries = props.getInt("message.send.max.retries", 3)
-
-  /** Before each retry, the producer refreshes the metadata of relevant topics. Since leader
-    * election takes a bit of time, this property specifies the amount of time that the producer
-    * waits before refreshing the metadata.
-    */
-  val retryBackoffMs = props.getInt("retry.backoff.ms", 100)
-
-  /**
-   * The producer generally refreshes the topic metadata from brokers when there is a failure
-   * (partition missing, leader not available...). It will also poll regularly (default: every 10min
-   * so 600000ms). If you set this to a negative value, metadata will only get refreshed on failure.
-   * If you set this to zero, the metadata will get refreshed after each message sent (not recommended)
-   * Important note: the refresh happen only AFTER the message is sent, so if the producer never sends
-   * a message the metadata is never refreshed
-   */
-  val topicMetadataRefreshIntervalMs = props.getInt("topic.metadata.refresh.interval.ms", 600000)
-
-  validate(this)
-}
diff --git a/core/src/main/scala/kafka/producer/ProducerPool.scala b/core/src/main/scala/kafka/producer/ProducerPool.scala
deleted file mode 100644
index 6d4e4b7ccab12..0000000000000
--- a/core/src/main/scala/kafka/producer/ProducerPool.scala
+++ /dev/null
@@ -1,90 +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.producer
-
-import java.util.Properties
-
-import kafka.api.TopicMetadata
-import kafka.cluster.BrokerEndPoint
-import kafka.common.UnavailableProducerException
-import kafka.utils.Logging
-import kafka.utils.Implicits._
-
-import scala.collection.mutable.HashMap
-
-@deprecated("This object has been deprecated and will be removed in a future release.", "0.10.0.0")
-object ProducerPool {
-  /**
-   * Used in ProducerPool to initiate a SyncProducer connection with a broker.
-   */
-  def createSyncProducer(config: ProducerConfig, broker: BrokerEndPoint): SyncProducer = {
-    val props = new Properties()
-    props.put("host", broker.host)
-    props.put("port", broker.port.toString)
-    props ++= config.props.props
-    new SyncProducer(new SyncProducerConfig(props))
-  }
-}
-
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.10.0.0")
-class ProducerPool(val config: ProducerConfig) extends Logging {
-  private val syncProducers = new HashMap[Int, SyncProducer]
-  private val lock = new Object()
-
-  def updateProducer(topicMetadata: Seq[TopicMetadata]) {
-    val newBrokers = new collection.mutable.HashSet[BrokerEndPoint]
-    topicMetadata.foreach(tmd => {
-      tmd.partitionsMetadata.foreach(pmd => {
-        if(pmd.leader.isDefined) {
-          newBrokers += pmd.leader.get
-        }
-      })
-    })
-    lock synchronized {
-      newBrokers.foreach(b => {
-        if(syncProducers.contains(b.id)){
-          syncProducers(b.id).close()
-          syncProducers.put(b.id, ProducerPool.createSyncProducer(config, b))
-        } else
-          syncProducers.put(b.id, ProducerPool.createSyncProducer(config, b))
-      })
-    }
-  }
-
-  def getProducer(brokerId: Int) : SyncProducer = {
-    lock.synchronized {
-      val producer = syncProducers.get(brokerId)
-      producer match {
-        case Some(p) => p
-        case None => throw new UnavailableProducerException("Sync producer for broker id %d does not exist".format(brokerId))
-      }
-    }
-  }
-
-  /**
-   * Closes all the producers in the pool
-   */
-  def close() = {
-    lock.synchronized {
-      info("Closing all sync producers")
-      val iter = syncProducers.values.iterator
-      while(iter.hasNext)
-        iter.next.close
-    }
-  }
-}
diff --git a/core/src/main/scala/kafka/producer/ProducerRequestStats.scala b/core/src/main/scala/kafka/producer/ProducerRequestStats.scala
deleted file mode 100644
index 92bbbcfa807dc..0000000000000
--- a/core/src/main/scala/kafka/producer/ProducerRequestStats.scala
+++ /dev/null
@@ -1,69 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package kafka.producer
-
-import kafka.metrics.{KafkaTimer, KafkaMetricsGroup}
-import java.util.concurrent.TimeUnit
-import kafka.utils.Pool
-import kafka.common.{ClientIdAllBrokers, ClientIdBroker, ClientIdAndBroker}
-
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.10.0.0")
-class ProducerRequestMetrics(metricId: ClientIdBroker) extends KafkaMetricsGroup {
-  val tags = metricId match {
-    case ClientIdAndBroker(clientId, brokerHost, brokerPort) => Map("clientId" -> clientId, "brokerHost" -> brokerHost, "brokerPort" -> brokerPort.toString)
-    case ClientIdAllBrokers(clientId) => Map("clientId" -> clientId)
-  }
-
-  val requestTimer = new KafkaTimer(newTimer("ProducerRequestRateAndTimeMs", TimeUnit.MILLISECONDS, TimeUnit.SECONDS, tags))
-  val requestSizeHist = newHistogram("ProducerRequestSize", biased = true, tags)
-  val throttleTimeStats = newTimer("ProducerRequestThrottleRateAndTimeMs", TimeUnit.MILLISECONDS, TimeUnit.SECONDS, tags)
-}
-
-/**
- * Tracks metrics of requests made by a given producer client to all brokers.
- * @param clientId ClientId of the given producer
- */
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.10.0.0")
-class ProducerRequestStats(clientId: String) {
-  private val valueFactory = (k: ClientIdBroker) => new ProducerRequestMetrics(k)
-  private val stats = new Pool[ClientIdBroker, ProducerRequestMetrics](Some(valueFactory))
-  private val allBrokersStats = new ProducerRequestMetrics(new ClientIdAllBrokers(clientId))
-
-  def getProducerRequestAllBrokersStats(): ProducerRequestMetrics = allBrokersStats
-
-  def getProducerRequestStats(brokerHost: String, brokerPort: Int): ProducerRequestMetrics = {
-    stats.getAndMaybePut(new ClientIdAndBroker(clientId, brokerHost, brokerPort))
-  }
-}
-
-/**
- * Stores the request stats information of each producer client in a (clientId -> ProducerRequestStats) map.
- */
-@deprecated("This object has been deprecated and will be removed in a future release.", "0.10.0.0")
-object ProducerRequestStatsRegistry {
-  private val valueFactory = (k: String) => new ProducerRequestStats(k)
-  private val globalStats = new Pool[String, ProducerRequestStats](Some(valueFactory))
-
-  def getProducerRequestStats(clientId: String) = {
-    globalStats.getAndMaybePut(clientId)
-  }
-
-  def removeProducerRequestStats(clientId: String) {
-    globalStats.remove(clientId)
-  }
-}
-
diff --git a/core/src/main/scala/kafka/producer/ProducerStats.scala b/core/src/main/scala/kafka/producer/ProducerStats.scala
deleted file mode 100644
index 9466f26d13a2e..0000000000000
--- a/core/src/main/scala/kafka/producer/ProducerStats.scala
+++ /dev/null
@@ -1,46 +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.producer
-
-import kafka.metrics.KafkaMetricsGroup
-import java.util.concurrent.TimeUnit
-import kafka.utils.Pool
-
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.10.0.0")
-class ProducerStats(clientId: String) extends KafkaMetricsGroup {
-  val tags: Map[String, String] = Map("clientId" -> clientId)
-  val serializationErrorRate = newMeter("SerializationErrorsPerSec", "errors", TimeUnit.SECONDS, tags)
-  val resendRate = newMeter("ResendsPerSec", "resends", TimeUnit.SECONDS, tags)
-  val failedSendRate = newMeter("FailedSendsPerSec", "failed sends", TimeUnit.SECONDS, tags)
-}
-
-/**
- * Stores metrics of serialization and message sending activity of each producer client in a (clientId -> ProducerStats) map.
- */
-@deprecated("This object has been deprecated and will be removed in a future release.", "0.10.0.0")
-object ProducerStatsRegistry {
-  private val valueFactory = (k: String) => new ProducerStats(k)
-  private val statsRegistry = new Pool[String, ProducerStats](Some(valueFactory))
-
-  def getProducerStats(clientId: String) = {
-    statsRegistry.getAndMaybePut(clientId)
-  }
-
-  def removeProducerStats(clientId: String) {
-    statsRegistry.remove(clientId)
-  }
-}
diff --git a/core/src/main/scala/kafka/producer/ProducerTopicStats.scala b/core/src/main/scala/kafka/producer/ProducerTopicStats.scala
deleted file mode 100644
index 7bb9610c2c408..0000000000000
--- a/core/src/main/scala/kafka/producer/ProducerTopicStats.scala
+++ /dev/null
@@ -1,69 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package kafka.producer
-
-import kafka.metrics.KafkaMetricsGroup
-import kafka.common.{ClientIdTopic, ClientIdAllTopics, ClientIdAndTopic}
-import kafka.utils.{Pool, threadsafe}
-import java.util.concurrent.TimeUnit
-
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.10.0.0")
-@threadsafe
-class ProducerTopicMetrics(metricId: ClientIdTopic) extends KafkaMetricsGroup {
-  val tags = metricId match {
-    case ClientIdAndTopic(clientId, topic) => Map("clientId" -> clientId, "topic" -> topic)
-    case ClientIdAllTopics(clientId) => Map("clientId" -> clientId)
-  }
-
-  val messageRate = newMeter("MessagesPerSec", "messages", TimeUnit.SECONDS, tags)
-  val byteRate = newMeter("BytesPerSec", "bytes", TimeUnit.SECONDS, tags)
-  val droppedMessageRate = newMeter("DroppedMessagesPerSec", "drops", TimeUnit.SECONDS, tags)
-}
-
-/**
- * Tracks metrics for each topic the given producer client has produced data to.
- * @param clientId The clientId of the given producer client.
- */
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.10.0.0")
-class ProducerTopicStats(clientId: String) {
-  private val valueFactory = (k: ClientIdTopic) => new ProducerTopicMetrics(k)
-  private val stats = new Pool[ClientIdTopic, ProducerTopicMetrics](Some(valueFactory))
-  private val allTopicsStats = new ProducerTopicMetrics(new ClientIdAllTopics(clientId)) // to differentiate from a topic named AllTopics
-
-  def getProducerAllTopicsStats(): ProducerTopicMetrics = allTopicsStats
-
-  def getProducerTopicStats(topic: String): ProducerTopicMetrics = {
-    stats.getAndMaybePut(new ClientIdAndTopic(clientId, topic))
-  }
-}
-
-/**
- * Stores the topic stats information of each producer client in a (clientId -> ProducerTopicStats) map.
- */
-@deprecated("This object has been deprecated and will be removed in a future release.", "0.10.0.0")
-object ProducerTopicStatsRegistry {
-  private val valueFactory = (k: String) => new ProducerTopicStats(k)
-  private val globalStats = new Pool[String, ProducerTopicStats](Some(valueFactory))
-
-  def getProducerTopicStats(clientId: String) = {
-    globalStats.getAndMaybePut(clientId)
-  }
-
-  def removeProducerTopicStats(clientId: String) {
-    globalStats.remove(clientId)
-  }
-}
diff --git a/core/src/main/scala/kafka/producer/SyncProducer.scala b/core/src/main/scala/kafka/producer/SyncProducer.scala
deleted file mode 100644
index b132293349471..0000000000000
--- a/core/src/main/scala/kafka/producer/SyncProducer.scala
+++ /dev/null
@@ -1,169 +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.producer
-
-import java.util.Random
-import java.util.concurrent.TimeUnit
-
-import kafka.api._
-import kafka.network.{RequestOrResponseSend, BlockingChannel}
-import kafka.utils._
-import org.apache.kafka.common.network.NetworkReceive
-import org.apache.kafka.common.protocol.ApiKeys
-import org.apache.kafka.common.utils.Utils._
-
-@deprecated("This object has been deprecated and will be removed in a future release. " +
-            "Please use org.apache.kafka.clients.producer.KafkaProducer instead.", "0.10.0.0")
-object SyncProducer {
-  val RequestKey: Short = 0
-  val randomGenerator = new Random
-}
-
-/*
- * Send a message set.
- */
-@threadsafe
-@deprecated("This class has been deprecated and will be removed in a future release. " +
-            "Please use org.apache.kafka.clients.producer.KafkaProducer instead.", "0.10.0.0")
-class SyncProducer(val config: SyncProducerConfig) extends Logging {
-
-  private val lock = new Object()
-  @volatile private var shutdown: Boolean = false
-  private val blockingChannel = new BlockingChannel(config.host, config.port, BlockingChannel.UseDefaultBufferSize,
-    config.sendBufferBytes, config.requestTimeoutMs)
-  val producerRequestStats = ProducerRequestStatsRegistry.getProducerRequestStats(config.clientId)
-
-  trace("Instantiating Scala Sync Producer with properties: %s".format(config.props))
-
-  private def verifyRequest(request: RequestOrResponse) = {
-    /**
-     * This seems a little convoluted, but the idea is to turn on verification simply changing log4j settings
-     * Also, when verification is turned on, care should be taken to see that the logs don't fill up with unnecessary
-     * data. So, leaving the rest of the logging at TRACE, while errors should be logged at ERROR level
-     */
-    if (isDebugEnabled) {
-      val buffer = new RequestOrResponseSend("", request).buffer
-      trace("verifying sendbuffer of size " + buffer.limit())
-      val requestTypeId = buffer.getShort()
-      if(requestTypeId == ApiKeys.PRODUCE.id) {
-        val request = ProducerRequest.readFrom(buffer)
-        trace(request.toString)
-      }
-    }
-  }
-
-  /**
-   * Common functionality for the public send methods
-   */
-  private def doSend(request: RequestOrResponse, readResponse: Boolean = true): NetworkReceive = {
-    lock synchronized {
-      verifyRequest(request)
-      getOrMakeConnection()
-
-      var response: NetworkReceive = null
-      try {
-        blockingChannel.send(request)
-        if(readResponse)
-          response = blockingChannel.receive()
-        else
-          trace("Skipping reading response")
-      } catch {
-        case e: java.io.IOException =>
-          // no way to tell if write succeeded. Disconnect and re-throw exception to let client handle retry
-          disconnect()
-          throw e
-        case e: Throwable => throw e
-      }
-      response
-    }
-  }
-
-  /**
-   * Send a message. If the producerRequest had required.request.acks=0, then the
-   * returned response object is null
-   */
-  def send(producerRequest: ProducerRequest): ProducerResponse = {
-    val requestSize = producerRequest.sizeInBytes
-    producerRequestStats.getProducerRequestStats(config.host, config.port).requestSizeHist.update(requestSize)
-    producerRequestStats.getProducerRequestAllBrokersStats.requestSizeHist.update(requestSize)
-
-    var response: NetworkReceive = null
-    val specificTimer = producerRequestStats.getProducerRequestStats(config.host, config.port).requestTimer
-    val aggregateTimer = producerRequestStats.getProducerRequestAllBrokersStats.requestTimer
-    aggregateTimer.time {
-      specificTimer.time {
-        response = doSend(producerRequest, producerRequest.requiredAcks != 0)
-      }
-    }
-    if(producerRequest.requiredAcks != 0) {
-      val producerResponse = ProducerResponse.readFrom(response.payload)
-      producerRequestStats.getProducerRequestStats(config.host, config.port).throttleTimeStats.update(producerResponse.throttleTime, TimeUnit.MILLISECONDS)
-      producerRequestStats.getProducerRequestAllBrokersStats.throttleTimeStats.update(producerResponse.throttleTime, TimeUnit.MILLISECONDS)
-      producerResponse
-    }
-    else
-      null
-  }
-
-  def send(request: TopicMetadataRequest): TopicMetadataResponse = {
-    val response = doSend(request)
-    TopicMetadataResponse.readFrom(response.payload)
-  }
-
-  def close() = {
-    lock synchronized {
-      disconnect()
-      shutdown = true
-    }
-  }
-
-  /**
-   * Disconnect from current channel, closing connection.
-   * Side effect: channel field is set to null on successful disconnect
-   */
-  private def disconnect() {
-    try {
-      info("Disconnecting from " + formatAddress(config.host, config.port))
-      blockingChannel.disconnect()
-    } catch {
-      case e: Exception => error("Error on disconnect: ", e)
-    }
-  }
-
-  private def connect(): BlockingChannel = {
-    if (!blockingChannel.isConnected && !shutdown) {
-      try {
-        blockingChannel.connect()
-        info("Connected to " + formatAddress(config.host, config.port) + " for producing")
-      } catch {
-        case e: Exception => {
-          disconnect()
-          error("Producer connection to " + formatAddress(config.host, config.port) + " unsuccessful", e)
-          throw e
-        }
-      }
-    }
-    blockingChannel
-  }
-
-  private def getOrMakeConnection() {
-    if(!blockingChannel.isConnected) {
-      connect()
-    }
-  }
-}
diff --git a/core/src/main/scala/kafka/producer/SyncProducerConfig.scala b/core/src/main/scala/kafka/producer/SyncProducerConfig.scala
deleted file mode 100644
index 207779c558b3a..0000000000000
--- a/core/src/main/scala/kafka/producer/SyncProducerConfig.scala
+++ /dev/null
@@ -1,72 +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.producer
-
-import java.util.Properties
-import kafka.utils.VerifiableProperties
-
-@deprecated("This class has been deprecated and will be removed in a future release. " +
-            "Please use org.apache.kafka.clients.producer.ProducerConfig instead.", "0.10.0.0")
-class SyncProducerConfig private (val props: VerifiableProperties) extends SyncProducerConfigShared {
-  def this(originalProps: Properties) {
-    this(new VerifiableProperties(originalProps))
-    // no need to verify the property since SyncProducerConfig is supposed to be used internally
-  }
-
-  /** the broker to which the producer sends events */
-  val host = props.getString("host")
-
-  /** the port on which the broker is running */
-  val port = props.getInt("port")
-}
-
-@deprecated("This trait has been deprecated and will be removed in a future release. " +
-            "Please use org.apache.kafka.clients.producer.ProducerConfig instead.", "0.10.0.0")
-trait SyncProducerConfigShared {
-  val props: VerifiableProperties
-  
-  val sendBufferBytes = props.getInt("send.buffer.bytes", 100*1024)
-
-  /* the client application sending the producer requests */
-  val clientId = props.getString("client.id", SyncProducerConfig.DefaultClientId)
-
-  /*
-   * The number of acknowledgments the producer requires the leader to have received before considering a request complete.
-   * This controls the durability of the messages sent by the producer.
-   *
-   * request.required.acks = 0 - means the producer will not wait for any acknowledgement from the leader.
-   * request.required.acks = 1 - means the leader will write the message to its local log and immediately acknowledge
-   * request.required.acks = -1 - means the leader will wait for acknowledgement from all in-sync replicas before acknowledging the write
-   */
-
-  val requestRequiredAcks = props.getShortInRange("request.required.acks", SyncProducerConfig.DefaultRequiredAcks,(-1,1))
-
-  /*
-   * The ack timeout of the producer requests. Value must be non-negative and non-zero
-   */
-  val requestTimeoutMs = props.getIntInRange("request.timeout.ms", SyncProducerConfig.DefaultAckTimeoutMs,
-                                             (1, Integer.MAX_VALUE))
-}
-
-@deprecated("This object has been deprecated and will be removed in a future release. " +
-            "Please use org.apache.kafka.clients.producer.ProducerConfig instead.", "0.10.0.0")
-object SyncProducerConfig {
-  val DefaultClientId = ""
-  val DefaultRequiredAcks : Short = 0
-  val DefaultAckTimeoutMs = 10000
-}
diff --git a/core/src/main/scala/kafka/producer/async/AsyncProducerConfig.scala b/core/src/main/scala/kafka/producer/async/AsyncProducerConfig.scala
deleted file mode 100644
index cc3a79d44c1ee..0000000000000
--- a/core/src/main/scala/kafka/producer/async/AsyncProducerConfig.scala
+++ /dev/null
@@ -1,49 +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.producer.async
-
-import kafka.utils.VerifiableProperties
-
-@deprecated("This trait has been deprecated and will be removed in a future release. " +
-            "Please use org.apache.kafka.clients.producer.ProducerConfig instead.", "0.10.0.0")
-trait AsyncProducerConfig {
-  val props: VerifiableProperties
-
-  /* maximum time, in milliseconds, for buffering data on the producer queue */
-  val queueBufferingMaxMs = props.getInt("queue.buffering.max.ms", 5000)
-
-  /** the maximum size of the blocking queue for buffering on the producer */
-  val queueBufferingMaxMessages = props.getInt("queue.buffering.max.messages", 10000)
-
-  /**
-   * Timeout for event enqueue:
-   * 0: events will be enqueued immediately or dropped if the queue is full
-   * -ve: enqueue will block indefinitely if the queue is full
-   * +ve: enqueue will block up to this many milliseconds if the queue is full
-   */
-  val queueEnqueueTimeoutMs = props.getInt("queue.enqueue.timeout.ms", -1)
-
-  /** the number of messages batched at the producer */
-  val batchNumMessages = props.getInt("batch.num.messages", 200)
-
-  /** the serializer class for values */
-  val serializerClass = props.getString("serializer.class", "kafka.serializer.DefaultEncoder")
-  
-  /** the serializer class for keys (defaults to the same as for values) */
-  val keySerializerClass = props.getString("key.serializer.class", serializerClass)
-  
-}
diff --git a/core/src/main/scala/kafka/producer/async/DefaultEventHandler.scala b/core/src/main/scala/kafka/producer/async/DefaultEventHandler.scala
deleted file mode 100755
index 8c7465f9e4956..0000000000000
--- a/core/src/main/scala/kafka/producer/async/DefaultEventHandler.scala
+++ /dev/null
@@ -1,359 +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.producer.async
-
-import kafka.common._
-import kafka.message.{ByteBufferMessageSet, Message, NoCompressionCodec}
-import kafka.producer._
-import kafka.serializer.Encoder
-import kafka.utils._
-import org.apache.kafka.common.errors.{LeaderNotAvailableException, UnknownTopicOrPartitionException}
-import org.apache.kafka.common.protocol.Errors
-
-import scala.util.Random
-import scala.collection.{Map, Seq}
-import scala.collection.mutable.{ArrayBuffer, HashMap, Set}
-import java.util.concurrent.atomic._
-
-import kafka.api.{ProducerRequest, TopicMetadata}
-import org.apache.kafka.common.utils.{Time, Utils}
-import org.slf4j.event.Level
-
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.10.0.0")
-class DefaultEventHandler[K,V](config: ProducerConfig,
-                               private val partitioner: Partitioner,
-                               private val encoder: Encoder[V],
-                               private val keyEncoder: Encoder[K],
-                               private val producerPool: ProducerPool,
-                               private val topicPartitionInfos: HashMap[String, TopicMetadata] = new HashMap[String, TopicMetadata],
-                               private val time: Time = Time.SYSTEM)
-  extends EventHandler[K,V] with Logging {
-
-  val isSync = ("sync" == config.producerType)
-
-  val correlationId = new AtomicInteger(0)
-  val brokerPartitionInfo = new BrokerPartitionInfo(config, producerPool, topicPartitionInfos)
-
-  private val topicMetadataRefreshInterval = config.topicMetadataRefreshIntervalMs
-  private var lastTopicMetadataRefreshTime = 0L
-  private val topicMetadataToRefresh = Set.empty[String]
-  private val sendPartitionPerTopicCache = HashMap.empty[String, Int]
-
-  private val producerStats = ProducerStatsRegistry.getProducerStats(config.clientId)
-  private val producerTopicStats = ProducerTopicStatsRegistry.getProducerTopicStats(config.clientId)
-
-  def handle(events: Seq[KeyedMessage[K,V]]) {
-    val serializedData = serialize(events)
-    serializedData.foreach {
-      keyed =>
-        val dataSize = keyed.message.payloadSize
-        producerTopicStats.getProducerTopicStats(keyed.topic).byteRate.mark(dataSize)
-        producerTopicStats.getProducerAllTopicsStats.byteRate.mark(dataSize)
-    }
-    var outstandingProduceRequests = serializedData
-    var remainingRetries = config.messageSendMaxRetries + 1
-    val correlationIdStart = correlationId.get()
-    debug("Handling %d events".format(events.size))
-    while (remainingRetries > 0 && outstandingProduceRequests.nonEmpty) {
-      topicMetadataToRefresh ++= outstandingProduceRequests.map(_.topic)
-      if (topicMetadataRefreshInterval >= 0 &&
-          Time.SYSTEM.milliseconds - lastTopicMetadataRefreshTime > topicMetadataRefreshInterval) {
-        CoreUtils.swallow(brokerPartitionInfo.updateInfo(topicMetadataToRefresh.toSet, correlationId.getAndIncrement), this, Level.ERROR)
-        sendPartitionPerTopicCache.clear()
-        topicMetadataToRefresh.clear
-        lastTopicMetadataRefreshTime = Time.SYSTEM.milliseconds
-      }
-      outstandingProduceRequests = dispatchSerializedData(outstandingProduceRequests)
-      if (outstandingProduceRequests.nonEmpty) {
-        info("Back off for %d ms before retrying send. Remaining retries = %d".format(config.retryBackoffMs, remainingRetries-1))
-        // back off and update the topic metadata cache before attempting another send operation
-        Thread.sleep(config.retryBackoffMs)
-        // get topics of the outstanding produce requests and refresh metadata for those
-        CoreUtils.swallow(brokerPartitionInfo.updateInfo(outstandingProduceRequests.map(_.topic).toSet, correlationId.getAndIncrement), this, Level.ERROR)
-        sendPartitionPerTopicCache.clear()
-        remainingRetries -= 1
-        producerStats.resendRate.mark()
-      }
-    }
-    if(outstandingProduceRequests.nonEmpty) {
-      producerStats.failedSendRate.mark()
-      val correlationIdEnd = correlationId.get()
-      error("Failed to send requests for topics %s with correlation ids in [%d,%d]"
-        .format(outstandingProduceRequests.map(_.topic).toSet.mkString(","),
-        correlationIdStart, correlationIdEnd-1))
-      throw new FailedToSendMessageException("Failed to send messages after " + config.messageSendMaxRetries + " tries.", null)
-    }
-  }
-
-  private def dispatchSerializedData(messages: Seq[KeyedMessage[K,Message]]): Seq[KeyedMessage[K, Message]] = {
-    val partitionedDataOpt = partitionAndCollate(messages)
-    partitionedDataOpt match {
-      case Some(partitionedData) =>
-        val failedProduceRequests = new ArrayBuffer[KeyedMessage[K, Message]]
-        for ((brokerid, messagesPerBrokerMap) <- partitionedData) {
-          if (isTraceEnabled) {
-            messagesPerBrokerMap.foreach(partitionAndEvent =>
-              trace("Handling event for Topic: %s, Broker: %d, Partitions: %s".format(partitionAndEvent._1, brokerid, partitionAndEvent._2)))
-          }
-          val messageSetPerBrokerOpt = groupMessagesToSet(messagesPerBrokerMap)
-          messageSetPerBrokerOpt match {
-            case Some(messageSetPerBroker) =>
-              val failedTopicPartitions = send(brokerid, messageSetPerBroker)
-              failedTopicPartitions.foreach(topicPartition => {
-                messagesPerBrokerMap.get(topicPartition).foreach(failedProduceRequests.appendAll)
-              })
-            case None => // failed to group messages
-              messagesPerBrokerMap.values.foreach(m => failedProduceRequests.appendAll(m))
-          }
-        }
-        failedProduceRequests
-      case None => // failed to collate messages
-        messages
-    }
-  }
-
-  def serialize(events: Seq[KeyedMessage[K,V]]): Seq[KeyedMessage[K,Message]] = {
-    val serializedMessages = new ArrayBuffer[KeyedMessage[K,Message]](events.size)
-    events.foreach{e =>
-      try {
-        if(e.hasKey)
-          serializedMessages += new KeyedMessage[K,Message](
-            topic = e.topic,
-            key = e.key,
-            partKey = e.partKey,
-            message = new Message(key = keyEncoder.toBytes(e.key),
-                                  bytes = encoder.toBytes(e.message),
-                                  timestamp = time.milliseconds,
-                                  magicValue = Message.MagicValue_V1))
-        else
-          serializedMessages += new KeyedMessage[K,Message](
-            topic = e.topic,
-            key = e.key,
-            partKey = e.partKey,
-            message = new Message(bytes = encoder.toBytes(e.message),
-                                  timestamp = time.milliseconds,
-                                  magicValue = Message.MagicValue_V1))
-      } catch {
-        case t: Throwable =>
-          producerStats.serializationErrorRate.mark()
-          if (isSync) {
-            throw t
-          } else {
-            // currently, if in async mode, we just log the serialization error. We need to revisit
-            // this when doing kafka-496
-            error("Error serializing message for topic %s".format(e.topic), t)
-          }
-      }
-    }
-    serializedMessages
-  }
-
-  def partitionAndCollate(messages: Seq[KeyedMessage[K,Message]]): Option[Map[Int, collection.mutable.Map[TopicAndPartition, Seq[KeyedMessage[K,Message]]]]] = {
-    val ret = new HashMap[Int, collection.mutable.Map[TopicAndPartition, Seq[KeyedMessage[K,Message]]]]
-    try {
-      for (message <- messages) {
-        val topicPartitionsList = getPartitionListForTopic(message)
-        val partitionIndex = getPartition(message.topic, message.partitionKey, topicPartitionsList)
-        val brokerPartition = topicPartitionsList(partitionIndex)
-
-        // postpone the failure until the send operation, so that requests for other brokers are handled correctly
-        val leaderBrokerId = brokerPartition.leaderBrokerIdOpt.getOrElse(-1)
-
-        var dataPerBroker: HashMap[TopicAndPartition, Seq[KeyedMessage[K,Message]]] = null
-        ret.get(leaderBrokerId) match {
-          case Some(element) =>
-            dataPerBroker = element.asInstanceOf[HashMap[TopicAndPartition, Seq[KeyedMessage[K,Message]]]]
-          case None =>
-            dataPerBroker = new HashMap[TopicAndPartition, Seq[KeyedMessage[K,Message]]]
-            ret.put(leaderBrokerId, dataPerBroker)
-        }
-
-        val topicAndPartition = TopicAndPartition(message.topic, brokerPartition.partitionId)
-        var dataPerTopicPartition: ArrayBuffer[KeyedMessage[K,Message]] = null
-        dataPerBroker.get(topicAndPartition) match {
-          case Some(element) =>
-            dataPerTopicPartition = element.asInstanceOf[ArrayBuffer[KeyedMessage[K,Message]]]
-          case None =>
-            dataPerTopicPartition = new ArrayBuffer[KeyedMessage[K,Message]]
-            dataPerBroker.put(topicAndPartition, dataPerTopicPartition)
-        }
-        dataPerTopicPartition.append(message)
-      }
-      Some(ret)
-    }catch {    // Swallow recoverable exceptions and return None so that they can be retried.
-      case ute: UnknownTopicOrPartitionException => warn("Failed to collate messages by topic,partition due to: " + ute.getMessage); None
-      case lnae: LeaderNotAvailableException => warn("Failed to collate messages by topic,partition due to: " + lnae.getMessage); None
-      case oe: Throwable => error("Failed to collate messages by topic, partition due to: " + oe.getMessage); None
-    }
-  }
-
-  private def getPartitionListForTopic(m: KeyedMessage[K,Message]): Seq[PartitionAndLeader] = {
-    val topicPartitionsList = brokerPartitionInfo.getBrokerPartitionInfo(m.topic, correlationId.getAndIncrement)
-    debug("Broker partitions registered for topic: %s are %s"
-      .format(m.topic, topicPartitionsList.map(p => p.partitionId).mkString(",")))
-    val totalNumPartitions = topicPartitionsList.length
-    if(totalNumPartitions == 0)
-      throw new NoBrokersForPartitionException("Partition key = " + m.key)
-    topicPartitionsList
-  }
-
-  /**
-   * Retrieves the partition id and throws an UnknownTopicOrPartitionException if
-   * the value of partition is not between 0 and numPartitions-1
-   * @param topic The topic
-   * @param key the partition key
-   * @param topicPartitionList the list of available partitions
-   * @return the partition id
-   */
-  private def getPartition(topic: String, key: Any, topicPartitionList: Seq[PartitionAndLeader]): Int = {
-    val numPartitions = topicPartitionList.size
-    if(numPartitions <= 0)
-      throw new UnknownTopicOrPartitionException("Topic " + topic + " doesn't exist")
-    val partition =
-      if(key == null) {
-        // If the key is null, we don't really need a partitioner
-        // So we look up in the send partition cache for the topic to decide the target partition
-        val id = sendPartitionPerTopicCache.get(topic)
-        id match {
-          case Some(partitionId) =>
-            // directly return the partitionId without checking availability of the leader,
-            // since we want to postpone the failure until the send operation anyways
-            partitionId
-          case None =>
-            val availablePartitions = topicPartitionList.filter(_.leaderBrokerIdOpt.isDefined)
-            if (availablePartitions.isEmpty)
-              throw new LeaderNotAvailableException("No leader for any partition in topic " + topic)
-            val index = Utils.abs(Random.nextInt) % availablePartitions.size
-            val partitionId = availablePartitions(index).partitionId
-            sendPartitionPerTopicCache.put(topic, partitionId)
-            partitionId
-        }
-      } else
-        partitioner.partition(key, numPartitions)
-    if(partition < 0 || partition >= numPartitions)
-      throw new UnknownTopicOrPartitionException("Invalid partition id: " + partition + " for topic " + topic +
-        "; Valid values are in the inclusive range of [0, " + (numPartitions-1) + "]")
-    trace("Assigning message of topic %s and key %s to a selected partition %d".format(topic, if (key == null) "[none]" else key.toString, partition))
-    partition
-  }
-
-  /**
-   * Constructs and sends the produce request based on a map from (topic, partition) -> messages
-   *
-   * @param brokerId the broker that will receive the request
-   * @param messagesPerTopic the messages as a map from (topic, partition) -> messages
-   * @return the set (topic, partitions) messages which incurred an error sending or processing
-   */
-  private def send(brokerId: Int, messagesPerTopic: collection.mutable.Map[TopicAndPartition, ByteBufferMessageSet]) = {
-    if(brokerId < 0) {
-      warn("Failed to send data since partitions %s don't have a leader".format(messagesPerTopic.keys.mkString(",")))
-      messagesPerTopic.keys.toSeq
-    } else if(messagesPerTopic.nonEmpty) {
-      val currentCorrelationId = correlationId.getAndIncrement
-      val producerRequest = new ProducerRequest(currentCorrelationId, config.clientId, config.requestRequiredAcks,
-        config.requestTimeoutMs, messagesPerTopic)
-      var failedTopicPartitions = Seq.empty[TopicAndPartition]
-      try {
-        val syncProducer = producerPool.getProducer(brokerId)
-        debug("Producer sending messages with correlation id %d for topics %s to broker %d on %s:%d"
-          .format(currentCorrelationId, messagesPerTopic.keySet.mkString(","), brokerId, syncProducer.config.host, syncProducer.config.port))
-        val response = syncProducer.send(producerRequest)
-        debug("Producer sent messages with correlation id %d for topics %s to broker %d on %s:%d"
-          .format(currentCorrelationId, messagesPerTopic.keySet.mkString(","), brokerId, syncProducer.config.host, syncProducer.config.port))
-        if(response != null) {
-          if (response.status.size != producerRequest.data.size)
-            throw new KafkaException("Incomplete response (%s) for producer request (%s)".format(response, producerRequest))
-          if (isTraceEnabled) {
-            val successfullySentData = response.status.filter(_._2.error == Errors.NONE)
-            successfullySentData.foreach(m => messagesPerTopic(m._1).foreach(message =>
-              trace("Successfully sent message: %s".format(if(message.message.isNull) null else message.message.toString()))))
-          }
-          val failedPartitionsAndStatus = response.status.filter(_._2.error != Errors.NONE).toSeq
-          failedTopicPartitions = failedPartitionsAndStatus.map(partitionStatus => partitionStatus._1)
-          if(failedTopicPartitions.nonEmpty) {
-            val errorString = failedPartitionsAndStatus
-              .sortWith((p1, p2) => p1._1.topic.compareTo(p2._1.topic) < 0 ||
-                                    (p1._1.topic.compareTo(p2._1.topic) == 0 && p1._1.partition < p2._1.partition))
-              .map{
-                case(topicAndPartition, status) =>
-                  topicAndPartition.toString + ": " + status.error.exceptionName
-              }.mkString(",")
-            warn("Produce request with correlation id %d failed due to %s".format(currentCorrelationId, errorString))
-          }
-          failedTopicPartitions
-        } else {
-          Seq.empty[TopicAndPartition]
-        }
-      } catch {
-        case t: Throwable =>
-          warn("Failed to send producer request with correlation id %d to broker %d with data for partitions %s"
-            .format(currentCorrelationId, brokerId, messagesPerTopic.keys.mkString(",")), t)
-          messagesPerTopic.keys.toSeq
-      }
-    } else {
-      List.empty
-    }
-  }
-
-  private def groupMessagesToSet(messagesPerTopicAndPartition: collection.mutable.Map[TopicAndPartition, Seq[KeyedMessage[K, Message]]]) = {
-    /** enforce the compressed.topics config here.
-      * If the compression codec is anything other than NoCompressionCodec,
-      * Enable compression only for specified topics if any
-      * If the list of compressed topics is empty, then enable the specified compression codec for all topics
-      * If the compression codec is NoCompressionCodec, compression is disabled for all topics
-      */
-    try {
-      val messagesPerTopicPartition = messagesPerTopicAndPartition.map { case (topicAndPartition, messages) =>
-        val rawMessages = messages.map(_.message)
-        (topicAndPartition,
-          config.compressionCodec match {
-            case NoCompressionCodec =>
-              debug("Sending %d messages with no compression to %s".format(messages.size, topicAndPartition))
-              new ByteBufferMessageSet(NoCompressionCodec, rawMessages: _*)
-            case _ =>
-              config.compressedTopics.size match {
-                case 0 =>
-                  debug("Sending %d messages with compression codec %d to %s"
-                    .format(messages.size, config.compressionCodec.codec, topicAndPartition))
-                  new ByteBufferMessageSet(config.compressionCodec, rawMessages: _*)
-                case _ =>
-                  if (config.compressedTopics.contains(topicAndPartition.topic)) {
-                    debug("Sending %d messages with compression codec %d to %s"
-                      .format(messages.size, config.compressionCodec.codec, topicAndPartition))
-                    new ByteBufferMessageSet(config.compressionCodec, rawMessages: _*)
-                  }
-                  else {
-                    debug("Sending %d messages to %s with no compression as it is not in compressed.topics - %s"
-                      .format(messages.size, topicAndPartition, config.compressedTopics.toString))
-                    new ByteBufferMessageSet(NoCompressionCodec, rawMessages: _*)
-                  }
-              }
-          }
-          )
-      }
-      Some(messagesPerTopicPartition)
-    } catch {
-      case t: Throwable => error("Failed to group messages", t); None
-    }
-  }
-
-  def close() {
-    if (producerPool != null)
-      producerPool.close
-  }
-}
diff --git a/core/src/main/scala/kafka/producer/async/EventHandler.scala b/core/src/main/scala/kafka/producer/async/EventHandler.scala
deleted file mode 100644
index 44fb1eb822c55..0000000000000
--- a/core/src/main/scala/kafka/producer/async/EventHandler.scala
+++ /dev/null
@@ -1,37 +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.producer.async
-
-import kafka.producer.KeyedMessage
-
-/**
- * Handler that dispatches the batched data from the queue.
- */
-@deprecated("This trait has been deprecated and will be removed in a future release.", "0.10.0.0")
-trait EventHandler[K,V] {
-
-  /**
-   * Callback to dispatch the batched data and send it to a Kafka server
-   * @param events the data sent to the producer
-  */
-  def handle(events: Seq[KeyedMessage[K,V]])
-
-  /**
-   * Cleans up and shuts down the event handler
-  */
-  def close(): Unit
-}
diff --git a/core/src/main/scala/kafka/producer/async/IllegalQueueStateException.scala b/core/src/main/scala/kafka/producer/async/IllegalQueueStateException.scala
deleted file mode 100644
index 7779715a67472..0000000000000
--- a/core/src/main/scala/kafka/producer/async/IllegalQueueStateException.scala
+++ /dev/null
@@ -1,26 +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.producer.async
-
-/**
- * Indicates that the given config parameter has invalid value
- */
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.10.0.0")
-class IllegalQueueStateException(message: String) extends RuntimeException(message) {
-  def this() = this(null)
-}
diff --git a/core/src/main/scala/kafka/producer/async/MissingConfigException.scala b/core/src/main/scala/kafka/producer/async/MissingConfigException.scala
deleted file mode 100644
index a42678b2eb204..0000000000000
--- a/core/src/main/scala/kafka/producer/async/MissingConfigException.scala
+++ /dev/null
@@ -1,24 +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.producer.async
-
-/* Indicates any missing configuration parameter */
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.10.0.0")
-class MissingConfigException(message: String) extends RuntimeException(message) {
-  def this() = this(null)
-}
diff --git a/core/src/main/scala/kafka/producer/async/ProducerSendThread.scala b/core/src/main/scala/kafka/producer/async/ProducerSendThread.scala
deleted file mode 100644
index 03770930ddf5f..0000000000000
--- a/core/src/main/scala/kafka/producer/async/ProducerSendThread.scala
+++ /dev/null
@@ -1,114 +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.producer.async
-
-import kafka.utils.Logging
-import java.util.concurrent.{BlockingQueue, CountDownLatch, TimeUnit}
-
-import collection.mutable.ArrayBuffer
-import kafka.producer.KeyedMessage
-import kafka.metrics.KafkaMetricsGroup
-import com.yammer.metrics.core.Gauge
-import org.apache.kafka.common.utils.Time
-
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.10.0.0")
-class ProducerSendThread[K,V](val threadName: String,
-                              val queue: BlockingQueue[KeyedMessage[K,V]],
-                              val handler: EventHandler[K,V],
-                              val queueTime: Long,
-                              val batchSize: Int,
-                              val clientId: String) extends Thread(threadName) with Logging with KafkaMetricsGroup {
-
-  private val shutdownLatch = new CountDownLatch(1)
-  private val shutdownCommand = new KeyedMessage[K,V]("shutdown", null.asInstanceOf[K], null.asInstanceOf[V])
-
-  newGauge("ProducerQueueSize",
-          new Gauge[Int] {
-            def value = queue.size
-          },
-          Map("clientId" -> clientId))
-
-  override def run {
-    try {
-      processEvents
-    }catch {
-      case e: Throwable => error("Error in sending events: ", e)
-    }finally {
-      shutdownLatch.countDown
-    }
-  }
-
-  def shutdown(): Unit = {
-    info("Begin shutting down ProducerSendThread")
-    queue.put(shutdownCommand)
-    shutdownLatch.await
-    info("Shutdown ProducerSendThread complete")
-  }
-
-  private def processEvents() {
-    var lastSend = Time.SYSTEM.milliseconds
-    var events = new ArrayBuffer[KeyedMessage[K,V]]
-    var full: Boolean = false
-
-    // drain the queue until you get a shutdown command
-    Iterator.continually(queue.poll(scala.math.max(0, (lastSend + queueTime) - Time.SYSTEM.milliseconds), TimeUnit.MILLISECONDS))
-                      .takeWhile(item => if(item != null) item ne shutdownCommand else true).foreach {
-      currentQueueItem =>
-        val elapsed = Time.SYSTEM.milliseconds - lastSend
-        // check if the queue time is reached. This happens when the poll method above returns after a timeout and
-        // returns a null object
-        val expired = currentQueueItem == null
-        if(currentQueueItem != null) {
-          trace("Dequeued item for topic %s, partition key: %s, data: %s"
-              .format(currentQueueItem.topic, currentQueueItem.key, currentQueueItem.message))
-          events += currentQueueItem
-        }
-
-        // check if the batch size is reached
-        full = events.size >= batchSize
-
-        if(full || expired) {
-          if(expired)
-            debug(elapsed + " ms elapsed. Queue time reached. Sending..")
-          if(full)
-            debug("Batch full. Sending..")
-          // if either queue time has reached or batch size has reached, dispatch to event handler
-          tryToHandle(events)
-          lastSend = Time.SYSTEM.milliseconds
-          events = new ArrayBuffer[KeyedMessage[K,V]]
-        }
-    }
-    // send the last batch of events
-    tryToHandle(events)
-    if(queue.size > 0)
-      throw new IllegalQueueStateException("Invalid queue state! After queue shutdown, %d remaining items in the queue"
-        .format(queue.size))
-  }
-
-  def tryToHandle(events: Seq[KeyedMessage[K,V]]) {
-    val size = events.size
-    try {
-      debug("Handling " + size + " events")
-      if(size > 0)
-        handler.handle(events)
-    }catch {
-      case e: Throwable => error("Error in handling batch of " + size + " events", e)
-    }
-  }
-
-}
diff --git a/core/src/main/scala/kafka/raft/KafkaMetadataLog.scala b/core/src/main/scala/kafka/raft/KafkaMetadataLog.scala
new file mode 100644
index 0000000000000..039b299481dd4
--- /dev/null
+++ b/core/src/main/scala/kafka/raft/KafkaMetadataLog.scala
@@ -0,0 +1,164 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package kafka.raft
+
+import java.nio.file.NoSuchFileException
+import java.util.Optional
+
+import kafka.log.{AppendOrigin, Log}
+import kafka.server.{FetchHighWatermark, FetchLogEnd}
+import org.apache.kafka.common.record.{MemoryRecords, Records}
+import org.apache.kafka.common.{KafkaException, TopicPartition}
+import org.apache.kafka.raft
+import org.apache.kafka.raft.{LogAppendInfo, LogFetchInfo, LogOffsetMetadata, Isolation, ReplicatedLog}
+import org.apache.kafka.snapshot.FileRawSnapshotReader
+import org.apache.kafka.snapshot.FileRawSnapshotWriter
+import org.apache.kafka.snapshot.RawSnapshotReader
+import org.apache.kafka.snapshot.RawSnapshotWriter
+
+import scala.compat.java8.OptionConverters._
+
+class KafkaMetadataLog(
+  log: Log,
+  topicPartition: TopicPartition,
+  maxFetchSizeInBytes: Int = 1024 * 1024
+) extends ReplicatedLog {
+
+  override def read(startOffset: Long, readIsolation: Isolation): LogFetchInfo = {
+    val isolation = readIsolation match {
+      case Isolation.COMMITTED => FetchHighWatermark
+      case Isolation.UNCOMMITTED => FetchLogEnd
+      case _ => throw new IllegalArgumentException(s"Unhandled read isolation $readIsolation")
+    }
+
+    val fetchInfo = log.read(startOffset,
+      maxLength = maxFetchSizeInBytes,
+      isolation = isolation,
+      minOneMessage = true)
+
+    new LogFetchInfo(
+      fetchInfo.records,
+
+      new LogOffsetMetadata(
+        fetchInfo.fetchOffsetMetadata.messageOffset,
+        Optional.of(SegmentPosition(
+          fetchInfo.fetchOffsetMetadata.segmentBaseOffset,
+          fetchInfo.fetchOffsetMetadata.relativePositionInSegment))
+        )
+    )
+  }
+
+  override def appendAsLeader(records: Records, epoch: Int): LogAppendInfo = {
+    if (records.sizeInBytes == 0)
+      throw new IllegalArgumentException("Attempt to append an empty record set")
+
+    val appendInfo = log.appendAsLeader(records.asInstanceOf[MemoryRecords],
+      leaderEpoch = epoch,
+      origin = AppendOrigin.Coordinator)
+    new LogAppendInfo(appendInfo.firstOffset.getOrElse {
+      throw new KafkaException("Append failed unexpectedly")
+    }, appendInfo.lastOffset)
+  }
+
+  override def appendAsFollower(records: Records): LogAppendInfo = {
+    if (records.sizeInBytes == 0)
+      throw new IllegalArgumentException("Attempt to append an empty record set")
+
+    val appendInfo = log.appendAsFollower(records.asInstanceOf[MemoryRecords])
+    new LogAppendInfo(appendInfo.firstOffset.getOrElse {
+      throw new KafkaException("Append failed unexpectedly")
+    }, appendInfo.lastOffset)
+  }
+
+  override def lastFetchedEpoch: Int = {
+    log.latestEpoch.getOrElse(0)
+  }
+
+  override def endOffsetForEpoch(leaderEpoch: Int): Optional[raft.OffsetAndEpoch] = {
+    val endOffsetOpt = log.endOffsetForEpoch(leaderEpoch).map { offsetAndEpoch =>
+      new raft.OffsetAndEpoch(offsetAndEpoch.offset, offsetAndEpoch.leaderEpoch)
+    }
+    endOffsetOpt.asJava
+  }
+
+  override def endOffset: LogOffsetMetadata = {
+    val endOffsetMetadata = log.logEndOffsetMetadata
+    new LogOffsetMetadata(
+      endOffsetMetadata.messageOffset,
+      Optional.of(SegmentPosition(
+        endOffsetMetadata.segmentBaseOffset,
+        endOffsetMetadata.relativePositionInSegment))
+      )
+  }
+
+  override def startOffset: Long = {
+    log.logStartOffset
+  }
+
+  override def truncateTo(offset: Long): Unit = {
+    log.truncateTo(offset)
+  }
+
+  override def initializeLeaderEpoch(epoch: Int): Unit = {
+    log.maybeAssignEpochStartOffset(epoch, log.logEndOffset)
+  }
+
+  override def updateHighWatermark(offsetMetadata: LogOffsetMetadata): Unit = {
+    offsetMetadata.metadata.asScala match {
+      case Some(segmentPosition: SegmentPosition) => log.updateHighWatermarkOffsetMetadata(
+        new kafka.server.LogOffsetMetadata(
+          offsetMetadata.offset,
+          segmentPosition.baseOffset,
+          segmentPosition.relativePosition)
+      )
+      case _ =>
+        // FIXME: This API returns the new high watermark, which may be different from the passed offset
+        log.updateHighWatermark(offsetMetadata.offset)
+    }
+  }
+
+  override def flush(): Unit = {
+    log.flush()
+  }
+
+  override def lastFlushedOffset(): Long = {
+    log.recoveryPoint
+  }
+
+  /**
+   * Return the topic partition associated with the log.
+   */
+  override def topicPartition(): TopicPartition = {
+    topicPartition
+  }
+
+  override def createSnapshot(snapshotId: raft.OffsetAndEpoch): RawSnapshotWriter = {
+    FileRawSnapshotWriter.create(log.dir.toPath, snapshotId)
+  }
+
+  override def readSnapshot(snapshotId: raft.OffsetAndEpoch): Optional[RawSnapshotReader] = {
+    try {
+      Optional.of(FileRawSnapshotReader.open(log.dir.toPath, snapshotId))
+    } catch {
+      case e: NoSuchFileException => Optional.empty()
+    }
+  }
+
+  override def close(): Unit = {
+    log.close()
+  }
+}
diff --git a/core/src/main/scala/kafka/raft/KafkaNetworkChannel.scala b/core/src/main/scala/kafka/raft/KafkaNetworkChannel.scala
new file mode 100644
index 0000000000000..7f769c8463e09
--- /dev/null
+++ b/core/src/main/scala/kafka/raft/KafkaNetworkChannel.scala
@@ -0,0 +1,247 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package kafka.raft
+
+import java.net.InetSocketAddress
+import java.util
+import java.util.concurrent.ArrayBlockingQueue
+import java.util.concurrent.atomic.AtomicInteger
+
+import kafka.utils.Logging
+import org.apache.kafka.clients.{ClientRequest, ClientResponse, KafkaClient}
+import org.apache.kafka.common.message._
+import org.apache.kafka.common.protocol.{ApiKeys, ApiMessage, Errors}
+import org.apache.kafka.common.requests._
+import org.apache.kafka.common.utils.Time
+import org.apache.kafka.common.{KafkaException, Node}
+import org.apache.kafka.raft.{NetworkChannel, RaftMessage, RaftRequest, RaftResponse, RaftUtil}
+
+import scala.collection.mutable
+import scala.jdk.CollectionConverters._
+
+object KafkaNetworkChannel {
+
+  private[raft] def buildResponse(responseData: ApiMessage): AbstractResponse = {
+    responseData match {
+      case voteResponse: VoteResponseData =>
+        new VoteResponse(voteResponse)
+      case beginEpochResponse: BeginQuorumEpochResponseData =>
+        new BeginQuorumEpochResponse(beginEpochResponse)
+      case endEpochResponse: EndQuorumEpochResponseData =>
+        new EndQuorumEpochResponse(endEpochResponse)
+      case fetchResponse: FetchResponseData =>
+        new FetchResponse(fetchResponse)
+      case _ =>
+        throw new IllegalArgumentException(s"Unexpected type for responseData: $responseData")
+    }
+  }
+
+  private[raft] def buildRequest(requestData: ApiMessage): AbstractRequest.Builder[_ <: AbstractRequest] = {
+    requestData match {
+      case voteRequest: VoteRequestData =>
+        new VoteRequest.Builder(voteRequest)
+      case beginEpochRequest: BeginQuorumEpochRequestData =>
+        new BeginQuorumEpochRequest.Builder(beginEpochRequest)
+      case endEpochRequest: EndQuorumEpochRequestData =>
+        new EndQuorumEpochRequest.Builder(endEpochRequest)
+      case fetchRequest: FetchRequestData =>
+        // Since we already have the request, we go through a simplified builder
+        new AbstractRequest.Builder[FetchRequest](ApiKeys.FETCH) {
+          override def build(version: Short): FetchRequest = new FetchRequest(fetchRequest, version)
+        }
+      case _ =>
+        throw new IllegalArgumentException(s"Unexpected type for requestData: $requestData")
+    }
+  }
+
+  private[raft] def responseData(response: AbstractResponse): ApiMessage = {
+    response match {
+      case voteResponse: VoteResponse => voteResponse.data
+      case beginEpochResponse: BeginQuorumEpochResponse => beginEpochResponse.data
+      case endEpochResponse: EndQuorumEpochResponse => endEpochResponse.data
+      case fetchResponse: FetchResponse[_] => fetchResponse.data
+      case _ => throw new IllegalArgumentException(s"Unexpected type for response: $response")
+    }
+  }
+
+  private[raft] def requestData(request: AbstractRequest): ApiMessage = {
+    request match {
+      case voteRequest: VoteRequest => voteRequest.data
+      case beginEpochRequest: BeginQuorumEpochRequest => beginEpochRequest.data
+      case endEpochRequest: EndQuorumEpochRequest => endEpochRequest.data
+      case fetchRequest: FetchRequest => fetchRequest.data
+      case _ => throw new IllegalArgumentException(s"Unexpected type for request: $request")
+    }
+  }
+
+}
+
+class KafkaNetworkChannel(time: Time,
+                          client: KafkaClient,
+                          clientId: String,
+                          retryBackoffMs: Int,
+                          requestTimeoutMs: Int) extends NetworkChannel with Logging {
+  import KafkaNetworkChannel._
+
+  type ResponseHandler = AbstractResponse => Unit
+
+  private val correlationIdCounter = new AtomicInteger(0)
+  private val pendingInbound = mutable.Map.empty[Long, ResponseHandler]
+  private val undelivered = new ArrayBlockingQueue[RaftMessage](10)
+  private val pendingOutbound = new ArrayBlockingQueue[RaftRequest.Outbound](10)
+  private val endpoints = mutable.HashMap.empty[Int, Node]
+
+  override def newCorrelationId(): Int = correlationIdCounter.getAndIncrement()
+
+  private def buildClientRequest(req: RaftRequest.Outbound): ClientRequest = {
+    val destination = req.destinationId.toString
+    val request = buildRequest(req.data)
+    val correlationId = req.correlationId
+    val createdTimeMs = req.createdTimeMs
+    new ClientRequest(destination, request, correlationId, clientId, createdTimeMs, true,
+      requestTimeoutMs, null)
+  }
+
+  override def send(message: RaftMessage): Unit = {
+    message match {
+      case request: RaftRequest.Outbound =>
+        if (!pendingOutbound.offer(request))
+          throw new KafkaException("Pending outbound queue is full")
+
+      case response: RaftResponse.Outbound =>
+        pendingInbound.remove(response.correlationId).foreach { onResponseReceived: ResponseHandler =>
+          onResponseReceived(buildResponse(response.data))
+        }
+      case _ =>
+        throw new IllegalArgumentException("Unhandled message type " + message)
+    }
+  }
+
+  private def sendOutboundRequests(currentTimeMs: Long): Unit = {
+    while (!pendingOutbound.isEmpty) {
+      val request = pendingOutbound.peek()
+      endpoints.get(request.destinationId) match {
+        case Some(node) =>
+          if (client.connectionFailed(node)) {
+            pendingOutbound.poll()
+            val apiKey = ApiKeys.forId(request.data.apiKey)
+            val disconnectResponse = RaftUtil.errorResponse(apiKey, Errors.BROKER_NOT_AVAILABLE)
+            val success = undelivered.offer(new RaftResponse.Inbound(
+              request.correlationId, disconnectResponse, request.destinationId))
+            if (!success) {
+              throw new KafkaException("Undelivered queue is full")
+            }
+
+            // Make sure to reset the connection state
+            client.ready(node, currentTimeMs)
+          } else if (client.ready(node, currentTimeMs)) {
+            pendingOutbound.poll()
+            val clientRequest = buildClientRequest(request)
+            client.send(clientRequest, currentTimeMs)
+          } else {
+            // We will retry this request on the next poll
+            return
+          }
+
+        case None =>
+          pendingOutbound.poll()
+          val apiKey = ApiKeys.forId(request.data.apiKey)
+          val responseData = RaftUtil.errorResponse(apiKey, Errors.BROKER_NOT_AVAILABLE)
+          val response = new RaftResponse.Inbound(request.correlationId, responseData, request.destinationId)
+          if (!undelivered.offer(response))
+            throw new KafkaException("Undelivered queue is full")
+      }
+    }
+  }
+
+  def getConnectionInfo(nodeId: Int): Node = {
+    if (!endpoints.contains(nodeId))
+      null
+    else
+      endpoints(nodeId)
+  }
+
+  def allConnections(): Set[Node] = {
+    endpoints.values.toSet
+  }
+
+  private def buildInboundRaftResponse(response: ClientResponse): RaftResponse.Inbound = {
+    val header = response.requestHeader()
+    val data = if (response.authenticationException != null) {
+      RaftUtil.errorResponse(header.apiKey, Errors.CLUSTER_AUTHORIZATION_FAILED)
+    } else if (response.wasDisconnected) {
+      RaftUtil.errorResponse(header.apiKey, Errors.BROKER_NOT_AVAILABLE)
+    } else {
+      responseData(response.responseBody)
+    }
+    new RaftResponse.Inbound(header.correlationId, data, response.destination.toInt)
+  }
+
+  private def pollInboundResponses(timeoutMs: Long, inboundMessages: util.List[RaftMessage]): Unit = {
+    val responses = client.poll(timeoutMs, time.milliseconds())
+    for (response <- responses.asScala) {
+      inboundMessages.add(buildInboundRaftResponse(response))
+    }
+  }
+
+  private def drainInboundRequests(inboundMessages: util.List[RaftMessage]): Unit = {
+    undelivered.drainTo(inboundMessages)
+  }
+
+  private def pollInboundMessages(timeoutMs: Long): util.List[RaftMessage] = {
+    val pollTimeoutMs = if (!undelivered.isEmpty) {
+      0L
+    } else if (!pendingOutbound.isEmpty) {
+      retryBackoffMs
+    } else {
+      timeoutMs
+    }
+    val messages = new util.ArrayList[RaftMessage]
+    pollInboundResponses(pollTimeoutMs, messages)
+    drainInboundRequests(messages)
+    messages
+  }
+
+  override def receive(timeoutMs: Long): util.List[RaftMessage] = {
+    sendOutboundRequests(time.milliseconds())
+    pollInboundMessages(timeoutMs)
+  }
+
+  override def wakeup(): Unit = {
+    client.wakeup()
+  }
+
+  override def updateEndpoint(id: Int, address: InetSocketAddress): Unit = {
+    val node = new Node(id, address.getHostString, address.getPort)
+    endpoints.put(id, node)
+  }
+
+  def postInboundRequest(request: AbstractRequest, onResponseReceived: ResponseHandler): Unit = {
+    val data = requestData(request)
+    val correlationId = newCorrelationId()
+    val req = new RaftRequest.Inbound(correlationId, data, time.milliseconds())
+    pendingInbound.put(correlationId, onResponseReceived)
+    if (!undelivered.offer(req))
+      throw new KafkaException("Undelivered queue is full")
+    wakeup()
+  }
+
+  override def close(): Unit = {
+    client.close()
+  }
+
+}
diff --git a/core/src/main/scala/kafka/raft/SegmentPosition.scala b/core/src/main/scala/kafka/raft/SegmentPosition.scala
new file mode 100644
index 0000000000000..eb6a59f35d3bc
--- /dev/null
+++ b/core/src/main/scala/kafka/raft/SegmentPosition.scala
@@ -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 kafka.raft
+
+import org.apache.kafka.raft.OffsetMetadata
+
+case class SegmentPosition(baseOffset: Long, relativePosition: Int) extends OffsetMetadata {
+  override def toString: String = s"(segmentBaseOffset=$baseOffset,relativePositionInSegment=$relativePosition)"
+}
diff --git a/core/src/main/scala/kafka/raft/TimingWheelExpirationService.scala b/core/src/main/scala/kafka/raft/TimingWheelExpirationService.scala
new file mode 100644
index 0000000000000..c07661e5a5103
--- /dev/null
+++ b/core/src/main/scala/kafka/raft/TimingWheelExpirationService.scala
@@ -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 kafka.raft
+
+import java.util.concurrent.CompletableFuture
+
+import kafka.utils.ShutdownableThread
+import kafka.utils.timer.{Timer, TimerTask}
+import org.apache.kafka.common.errors.TimeoutException
+import org.apache.kafka.raft.ExpirationService
+
+object TimingWheelExpirationService {
+  private val WorkTimeoutMs: Long = 200L
+
+  class TimerTaskCompletableFuture[T](override val delayMs: Long) extends CompletableFuture[T] with TimerTask {
+    override def run(): Unit = {
+      completeExceptionally(new TimeoutException(
+        s"Future failed to be completed before timeout of $delayMs ms was reached"))
+    }
+  }
+}
+
+class TimingWheelExpirationService(timer: Timer) extends ExpirationService {
+  import TimingWheelExpirationService._
+
+  private val expirationReaper = new ExpiredOperationReaper()
+
+  expirationReaper.start()
+
+  override def failAfter[T](timeoutMs: Long): CompletableFuture[T] = {
+    val future = new TimerTaskCompletableFuture[T](timeoutMs)
+    future.whenComplete { (_, _) =>
+      future.cancel()
+    }
+    timer.add(future)
+    future
+  }
+
+  private class ExpiredOperationReaper extends ShutdownableThread(
+    name = "raft-expiration-reaper", isInterruptible = false) {
+
+    override def doWork(): Unit = {
+      timer.advanceClock(WorkTimeoutMs)
+    }
+  }
+
+  def shutdown(): Unit = {
+    expirationReaper.shutdown()
+  }
+}
diff --git a/core/src/main/scala/kafka/security/CredentialProvider.scala b/core/src/main/scala/kafka/security/CredentialProvider.scala
index 5d9d7ba283f70..9aa8bc915d4a7 100644
--- a/core/src/main/scala/kafka/security/CredentialProvider.scala
+++ b/core/src/main/scala/kafka/security/CredentialProvider.scala
@@ -17,19 +17,21 @@
 
 package kafka.security
 
-import java.util.{List, Properties}
+import java.util.{Collection, Properties}
 
 import org.apache.kafka.common.security.authenticator.CredentialCache
-import org.apache.kafka.common.security.scram.{ScramCredential, ScramCredentialUtils, ScramMechanism}
+import org.apache.kafka.common.security.scram.ScramCredential
 import org.apache.kafka.common.config.ConfigDef
 import org.apache.kafka.common.config.ConfigDef._
+import org.apache.kafka.common.security.scram.internals.{ScramCredentialUtils, ScramMechanism}
+import org.apache.kafka.common.security.token.delegation.internals.DelegationTokenCache
 
-class CredentialProvider(saslEnabledMechanisms: List[String]) {
+class CredentialProvider(scramMechanisms: Collection[String], val tokenCache: DelegationTokenCache) {
 
   val credentialCache = new CredentialCache
-  ScramCredentialUtils.createCache(credentialCache, saslEnabledMechanisms)
+  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/SecurityUtils.scala b/core/src/main/scala/kafka/security/SecurityUtils.scala
deleted file mode 100644
index 573a16b8c58f4..0000000000000
--- a/core/src/main/scala/kafka/security/SecurityUtils.scala
+++ /dev/null
@@ -1,53 +0,0 @@
-/**
-  * Licensed to the Apache Software Foundation (ASF) under one or more
-  * contributor license agreements.  See the NOTICE file distributed with
-  * this work for additional information regarding copyright ownership.
-  * The ASF licenses this file to You under the Apache License, Version 2.0
-  * (the "License"); you may not use this file except in compliance with
-  * the License.  You may obtain a copy of the License at
-  *
-  *    http://www.apache.org/licenses/LICENSE-2.0
-  *
-  * Unless required by applicable law or agreed to in writing, software
-  * distributed under the License is distributed on an "AS IS" BASIS,
-  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-  * See the License for the specific language governing permissions and
-  * limitations under the License.
-  */
-
-package kafka.security
-
-import kafka.security.auth.{Acl, Operation, PermissionType, Resource, ResourceType}
-import org.apache.kafka.common.acl.{AccessControlEntry, AclBinding, AclBindingFilter}
-import org.apache.kafka.common.protocol.Errors
-import org.apache.kafka.common.requests.ApiError
-import org.apache.kafka.common.resource.{Resource => AdminResource}
-import org.apache.kafka.common.security.auth.KafkaPrincipal
-
-import scala.util.{Failure, Success, Try}
-
-
-object SecurityUtils {
-
-  def convertToResourceAndAcl(filter: AclBindingFilter): Either[ApiError, (Resource, Acl)] = {
-    (for {
-      resourceType <- Try(ResourceType.fromJava(filter.resourceFilter.resourceType))
-      principal <- Try(KafkaPrincipal.fromString(filter.entryFilter.principal))
-      operation <- Try(Operation.fromJava(filter.entryFilter.operation))
-      permissionType <- Try(PermissionType.fromJava(filter.entryFilter.permissionType))
-      resource = Resource(resourceType, filter.resourceFilter.name)
-      acl = Acl(principal, permissionType, filter.entryFilter.host, operation)
-    } yield (resource, acl)) match {
-      case Failure(throwable) => Left(new ApiError(Errors.INVALID_REQUEST, throwable.getMessage))
-      case Success(s) => Right(s)
-    }
-  }
-
-  def convertToAclBinding(resource: Resource, acl: Acl): AclBinding = {
-    val adminResource = new AdminResource(resource.resourceType.toJava, resource.name)
-    val entry = new AccessControlEntry(acl.principal.toString, acl.host.toString,
-      acl.operation.toJava, acl.permissionType.toJava)
-    new AclBinding(adminResource, entry)
-  }
-
-}
diff --git a/core/src/main/scala/kafka/security/auth/Acl.scala b/core/src/main/scala/kafka/security/auth/Acl.scala
index 4e2cba4032f4e..befd9d27a61dd 100644
--- a/core/src/main/scala/kafka/security/auth/Acl.scala
+++ b/core/src/main/scala/kafka/security/auth/Acl.scala
@@ -17,61 +17,40 @@
 
 package kafka.security.auth
 
-import kafka.utils.Json
+import kafka.security.authorizer.AclEntry
+import org.apache.kafka.common.resource.ResourcePattern
 import org.apache.kafka.common.security.auth.KafkaPrincipal
-import org.apache.kafka.common.utils.SecurityUtils
 
+@deprecated("Use org.apache.kafka.common.acl.AclBinding", "Since 2.5")
 object Acl {
-  val WildCardPrincipal: KafkaPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "*")
-  val WildCardHost: String = "*"
+  val WildCardPrincipal: KafkaPrincipal = AclEntry.WildcardPrincipal
+  val WildCardHost: String = AclEntry.WildcardHost
+  val WildCardResource: String = ResourcePattern.WILDCARD_RESOURCE
   val AllowAllAcl = new Acl(WildCardPrincipal, Allow, WildCardHost, All)
-  val PrincipalKey = "principal"
-  val PermissionTypeKey = "permissionType"
-  val OperationKey = "operation"
-  val HostsKey = "host"
-  val VersionKey = "version"
-  val CurrentVersion = 1
-  val AclsKey = "acls"
+  val PrincipalKey = AclEntry.PrincipalKey
+  val PermissionTypeKey = AclEntry.PermissionTypeKey
+  val OperationKey = AclEntry.OperationKey
+  val HostsKey = AclEntry.HostsKey
+  val VersionKey = AclEntry.VersionKey
+  val CurrentVersion = AclEntry.CurrentVersion
+  val AclsKey = AclEntry.AclsKey
 
   /**
    *
-   * @param bytes of acls json string
-   *
-   * 
              
-      {
-        "version": 1,
-        "acls": [
-          {
-            "host":"host1",
-            "permissionType": "Deny",
-            "operation": "Read",
-            "principal": "User:alice"
-          }
-        ]
-      }
-   * 
              
-   *
-   * @return
+   * @see AclEntry
    */
   def fromBytes(bytes: Array[Byte]): Set[Acl] = {
-    if (bytes == null || bytes.isEmpty)
-      return collection.immutable.Set.empty[Acl]
-
-    Json.parseBytes(bytes).map(_.asJsonObject).map { js =>
-      //the acl json version.
-      require(js(VersionKey).to[Int] == CurrentVersion)
-      js(AclsKey).asJsonArray.iterator.map(_.asJsonObject).map { itemJs =>
-        val principal = SecurityUtils.parseKafkaPrincipal(itemJs(PrincipalKey).to[String])
-        val permissionType = PermissionType.fromString(itemJs(PermissionTypeKey).to[String])
-        val host = itemJs(HostsKey).to[String]
-        val operation = Operation.fromString(itemJs(OperationKey).to[String])
-        new Acl(principal, permissionType, host, operation)
-      }.toSet
-    }.getOrElse(Set.empty)
+    AclEntry.fromBytes(bytes)
+      .map(ace => Acl(ace.kafkaPrincipal,
+        PermissionType.fromJava(ace.permissionType()),
+        ace.host(),
+        Operation.fromJava(ace.operation())))
   }
 
   def toJsonCompatibleMap(acls: Set[Acl]): Map[String, Any] = {
-    Map(Acl.VersionKey -> Acl.CurrentVersion, Acl.AclsKey -> acls.map(acl => acl.toMap).toList)
+    AclEntry.toJsonCompatibleMap(acls.map(acl =>
+      AclEntry(acl.principal, acl.permissionType.toJava, acl.host, acl.operation.toJava)
+    ))
   }
 }
 
@@ -85,6 +64,7 @@ object Acl {
  * @param host A value of * indicates all hosts.
  * @param operation A value of ALL indicates all operations.
  */
+@deprecated("Use org.apache.kafka.common.acl.AclBinding", "Since 2.5")
 case class Acl(principal: KafkaPrincipal, permissionType: PermissionType, host: String, operation: Operation) {
 
   /**
diff --git a/core/src/main/scala/kafka/security/auth/Authorizer.scala b/core/src/main/scala/kafka/security/auth/Authorizer.scala
index 6f4ca0eb5225e..7509171313a53 100644
--- a/core/src/main/scala/kafka/security/auth/Authorizer.scala
+++ b/core/src/main/scala/kafka/security/auth/Authorizer.scala
@@ -32,12 +32,13 @@ 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 {
 
   /**
    * @param session The session being authenticated.
    * @param operation Type of operation client is trying to perform on resource.
-   * @param resource Resource the client is trying to access.
+   * @param resource Resource the client is trying to access. Resource pattern type is always literal in input resource.
    * @return true if the operation should be permitted, false otherwise
    */
   def authorize(session: Session, operation: Operation, resource: Resource): Boolean
@@ -45,42 +46,98 @@ trait Authorizer extends Configurable {
   /**
    * add the acls to resource, this is an additive operation so existing acls will not be overwritten, instead these new
    * acls will be added to existing acls.
+   *
+   * {code}
+   * // The following will add ACLs to the literal resource path 'foo', which will only affect the topic named 'foo':
+   * authorizer.addAcls(Set(acl1, acl2), Resource(Topic, "foo", LITERAL))
+   *
+   * // The following will add ACLs to the special literal topic resource path '*', which affects all topics:
+   * authorizer.addAcls(Set(acl1, acl2), Resource(Topic, "*", LITERAL))
+   *
+   * // The following will add ACLs to the prefixed resource path 'foo', which affects all topics whose name begins with 'foo':
+   * authorizer.addAcls(Set(acl1, acl2), Resource(Topic, "foo", PREFIXED))
+   * {code}
+   *
    * @param acls set of acls to add to existing acls
-   * @param resource the resource to which these acls should be attached.
+   * @param resource the resource path to which these acls should be attached.
+   *                the supplied resource will have a specific resource pattern type,
+   *                i.e. the resource pattern type will not be ``PatternType.ANY`` or ``PatternType.UNKNOWN``.
    */
   def addAcls(acls: Set[Acl], resource: Resource): Unit
 
   /**
    * remove these acls from the resource.
+   *
+   * {code}
+   * // The following will remove ACLs from the literal resource path 'foo', which will only affect the topic named 'foo':
+   * authorizer.removeAcls(Set(acl1, acl2), Resource(Topic, "foo", LITERAL))
+   *
+   * // The following will remove ACLs from the special literal topic resource path '*', which affects all topics:
+   * authorizer.removeAcls(Set(acl1, acl2), Resource(Topic, "*", LITERAL))
+   *
+   * // The following will remove ACLs from the prefixed resource path 'foo', which affects all topics whose name begins with 'foo':
+   * authorizer.removeAcls(Set(acl1, acl2), Resource(Topic, "foo", PREFIXED))
+   * {code}
+   *
    * @param acls set of acls to be removed.
-   * @param resource resource from which the acls should be removed.
+   * @param resource resource path from which the acls should be removed.
+   *                 the supplied resource will have a specific resource pattern type,
+   *                 i.e. the resource pattern type will not be ``PatternType.ANY`` or ``PatternType.UNKNOWN``.
    * @return true if some acl got removed, false if no acl was removed.
    */
   def removeAcls(acls: Set[Acl], resource: Resource): Boolean
 
   /**
    * remove a resource along with all of its acls from acl store.
-   * @param resource
+   *
+   * {code}
+   * // The following will remove all ACLs from the literal resource path 'foo', which will only affect the topic named 'foo':
+   * authorizer.removeAcls(Resource(Topic, "foo", LITERAL))
+   *
+   * // The following will remove all ACLs from the special literal topic resource path '*', which affects all topics:
+   * authorizer.removeAcls(Resource(Topic, "*", LITERAL))
+   *
+   * // The following will remove all ACLs from the prefixed resource path 'foo', which affects all topics whose name begins with 'foo':
+   * authorizer.removeAcls(Resource(Topic, "foo", PREFIXED))
+   * {code}
+   *
+   * @param resource the resource path from which these acls should be removed.
+   *                 the supplied resource will have a specific resource pattern type,
+   *                 i.e. the resource pattern type will not be ``PatternType.ANY`` or ``PatternType.UNKNOWN``.
    * @return
    */
   def removeAcls(resource: Resource): Boolean
 
   /**
-   * get set of acls for this resource
-   * @param resource
+   * get set of acls for the supplied resource
+   *
+   * {code}
+   * // The following will get all ACLs from the literal resource path 'foo', which will only affect the topic named 'foo':
+   * authorizer.removeAcls(Resource(Topic, "foo", LITERAL))
+   *
+   * // The following will get all ACLs from the special literal topic resource path '*', which affects all topics:
+   * authorizer.removeAcls(Resource(Topic, "*", LITERAL))
+   *
+   * // The following will get all ACLs from the prefixed resource path 'foo', which affects all topics whose name begins with 'foo':
+   * authorizer.removeAcls(Resource(Topic, "foo", PREFIXED))
+   * {code}
+   *
+   * @param resource the resource path to which the acls belong.
+   *                 the supplied resource will have a specific resource pattern type,
+   *                 i.e. the resource pattern type will not be ``PatternType.ANY`` or ``PatternType.UNKNOWN``.
    * @return empty set if no acls are found, otherwise the acls for the resource.
    */
   def getAcls(resource: Resource): Set[Acl]
 
   /**
    * get the acls for this principal.
-   * @param principal
+   * @param principal principal name.
    * @return empty Map if no acls exist for this principal, otherwise a map of resource -> acls for the principal.
    */
   def getAcls(principal: KafkaPrincipal): Map[Resource, Set[Acl]]
 
   /**
-   * gets the map of resource to acls for all resources.
+   * gets the map of resource paths to acls for all resources.
    */
   def getAcls(): Map[Resource, Set[Acl]]
 
@@ -90,4 +147,3 @@ trait Authorizer extends Configurable {
   def close(): Unit
 
 }
-
diff --git a/core/src/main/scala/kafka/security/auth/Operation.scala b/core/src/main/scala/kafka/security/auth/Operation.scala
index 0fa311a2ad9a9..3ae2384ea0156 100644
--- a/core/src/main/scala/kafka/security/auth/Operation.scala
+++ b/core/src/main/scala/kafka/security/auth/Operation.scala
@@ -22,56 +22,68 @@ import org.apache.kafka.common.acl.AclOperation
 /**
  * Different operations a client may perform on kafka resources.
  */
-
+@deprecated("Use org.apache.kafka.common.acl.AclOperation", "Since 2.5")
 sealed trait Operation extends BaseEnum {
   def toJava : AclOperation
 }
 
+@deprecated("Use org.apache.kafka.common.acl.AclOperation", "Since 2.5")
 case object Read extends Operation {
   val name = "Read"
   val toJava = AclOperation.READ
 }
+@deprecated("Use org.apache.kafka.common.acl.AclOperation", "Since 2.5")
 case object Write extends Operation {
   val name = "Write"
   val toJava = AclOperation.WRITE
 }
+@deprecated("Use org.apache.kafka.common.acl.AclOperation", "Since 2.5")
 case object Create extends Operation {
   val name = "Create"
   val toJava = AclOperation.CREATE
 }
+@deprecated("Use org.apache.kafka.common.acl.AclOperation", "Since 2.5")
 case object Delete extends Operation {
   val name = "Delete"
   val toJava = AclOperation.DELETE
 }
+@deprecated("Use org.apache.kafka.common.acl.AclOperation", "Since 2.5")
 case object Alter extends Operation {
   val name = "Alter"
   val toJava = AclOperation.ALTER
 }
+@deprecated("Use org.apache.kafka.common.acl.AclOperation", "Since 2.5")
 case object Describe extends Operation {
   val name = "Describe"
   val toJava = AclOperation.DESCRIBE
 }
+@deprecated("Use org.apache.kafka.common.acl.AclOperation", "Since 2.5")
 case object ClusterAction extends Operation {
   val name = "ClusterAction"
   val toJava = AclOperation.CLUSTER_ACTION
 }
+@deprecated("Use org.apache.kafka.common.acl.AclOperation", "Since 2.5")
 case object DescribeConfigs extends Operation {
   val name = "DescribeConfigs"
   val toJava = AclOperation.DESCRIBE_CONFIGS
 }
+@deprecated("Use org.apache.kafka.common.acl.AclOperation", "Since 2.5")
 case object AlterConfigs extends Operation {
   val name = "AlterConfigs"
   val toJava = AclOperation.ALTER_CONFIGS
 }
+@deprecated("Use org.apache.kafka.common.acl.AclOperation", "Since 2.5")
 case object IdempotentWrite extends Operation {
   val name = "IdempotentWrite"
   val toJava = AclOperation.IDEMPOTENT_WRITE
 }
+@deprecated("Use org.apache.kafka.common.acl.AclOperation", "Since 2.5")
 case object All extends Operation {
   val name = "All"
   val toJava = AclOperation.ALL
 }
 
+@deprecated("Use org.apache.kafka.common.acl.AclOperation", "Since 2.5")
 object Operation {
 
   def fromString(operation: String): Operation = {
@@ -79,7 +91,22 @@ object Operation {
     op.getOrElse(throw new KafkaException(operation + " not a valid operation name. The valid names are " + values.mkString(",")))
   }
 
-  def fromJava(operation: AclOperation): Operation = fromString(operation.toString.replaceAll("_", ""))
+  def fromJava(operation: AclOperation): Operation = {
+    operation match {
+      case AclOperation.READ => Read
+      case AclOperation.WRITE => Write
+      case AclOperation.CREATE => Create
+      case AclOperation.DELETE => Delete
+      case AclOperation.ALTER => Alter
+      case AclOperation.DESCRIBE => Describe
+      case AclOperation.CLUSTER_ACTION => ClusterAction
+      case AclOperation.ALTER_CONFIGS => AlterConfigs
+      case AclOperation.DESCRIBE_CONFIGS => DescribeConfigs
+      case AclOperation.IDEMPOTENT_WRITE => IdempotentWrite
+      case AclOperation.ALL => All
+      case _ => throw new KafkaException(operation + " is not a convertible operation name. The valid names are " + values.mkString(","))
+    }
+  }
 
   def values: Seq[Operation] = List(Read, Write, Create, Delete, Alter, Describe, ClusterAction, AlterConfigs,
      DescribeConfigs, IdempotentWrite, All)
diff --git a/core/src/main/scala/kafka/security/auth/PermissionType.scala b/core/src/main/scala/kafka/security/auth/PermissionType.scala
index c75e6f6f19d06..c5325b25342d1 100644
--- a/core/src/main/scala/kafka/security/auth/PermissionType.scala
+++ b/core/src/main/scala/kafka/security/auth/PermissionType.scala
@@ -19,20 +19,24 @@ package kafka.security.auth
 import kafka.common.{BaseEnum, KafkaException}
 import org.apache.kafka.common.acl.AclPermissionType
 
+@deprecated("Use org.apache.kafka.common.acl.AclPermissionType", "Since 2.5")
 sealed trait PermissionType extends BaseEnum {
   val toJava: AclPermissionType
 }
 
+@deprecated("Use org.apache.kafka.common.acl.AclPermissionType", "Since 2.5")
 case object Allow extends PermissionType {
   val name = "Allow"
   val toJava = AclPermissionType.ALLOW
 }
 
+@deprecated("Use org.apache.kafka.common.acl.AclPermissionType", "Since 2.5")
 case object Deny extends PermissionType {
   val name = "Deny"
   val toJava = AclPermissionType.DENY
 }
 
+@deprecated("Use org.apache.kafka.common.acl.AclPermissionType", "Since 2.5")
 object PermissionType {
   def fromString(permissionType: String): PermissionType = {
     val pType = values.find(pType => pType.name.equalsIgnoreCase(permissionType))
diff --git a/core/src/main/scala/kafka/security/auth/Resource.scala b/core/src/main/scala/kafka/security/auth/Resource.scala
index 311f5b5083a66..8045c68141687 100644
--- a/core/src/main/scala/kafka/security/auth/Resource.scala
+++ b/core/src/main/scala/kafka/security/auth/Resource.scala
@@ -16,31 +16,70 @@
  */
 package kafka.security.auth
 
+import kafka.common.KafkaException
+import kafka.security.authorizer.AclEntry
+import org.apache.kafka.common.resource.{PatternType, ResourcePattern}
+
+@deprecated("Use org.apache.kafka.common.resource.ResourcePattern", "Since 2.5")
 object Resource {
-  val Separator = ":"
+  val Separator = AclEntry.ResourceSeparator
   val ClusterResourceName = "kafka-cluster"
-  val ClusterResource = new Resource(Cluster, Resource.ClusterResourceName)
-  val ProducerIdResourceName = "producer-id"
-  val WildCardResource = "*"
+  val ClusterResource = Resource(Cluster, Resource.ClusterResourceName, PatternType.LITERAL)
+  val WildCardResource = AclEntry.WildcardResource
+
+  @deprecated("This resource name is not used by Kafka and will be removed in a future release", since = "2.1")
+  val ProducerIdResourceName = "producer-id" // This is not used since we don't have a producer id resource
 
   def fromString(str: String): Resource = {
-    str.split(Separator, 2) match {
-      case Array(resourceType, name, _*) => new Resource(ResourceType.fromString(resourceType), name)
-      case _ => throw new IllegalArgumentException("expected a string in format ResourceType:ResourceName but got " + str)
+    ResourceType.values.find(resourceType => str.startsWith(resourceType.name + Separator)) match {
+      case None => throw new KafkaException("Invalid resource string: '" + str + "'")
+      case Some(resourceType) =>
+        val remaining = str.substring(resourceType.name.length + 1)
+
+        PatternType.values.find(patternType => remaining.startsWith(patternType.name + Separator)) match {
+          case Some(patternType) =>
+            val name = remaining.substring(patternType.name.length + 1)
+            Resource(resourceType, name, patternType)
+
+          case None =>
+            Resource(resourceType, remaining, PatternType.LITERAL)
+        }
     }
   }
 }
 
 /**
  *
- * @param resourceType type of resource.
- * @param name name of the resource, for topic this will be topic name , for group it will be group name. For cluster type
+ * @param resourceType non-null type of resource.
+ * @param name non-null name of the resource, for topic this will be topic name , for group it will be group name. For cluster type
  *             it will be a constant string kafka-cluster.
+ * @param patternType non-null resource pattern type: literal, prefixed, etc.
  */
-case class Resource(resourceType: ResourceType, name: String) {
+@deprecated("Use org.apache.kafka.common.resource.ResourcePattern", "Since 2.5")
+case class Resource(resourceType: ResourceType, name: String, patternType: PatternType) {
+
+  if (!patternType.isSpecific)
+    throw new IllegalArgumentException(s"patternType must not be $patternType")
+
+  /**
+    * Create an instance of this class with the provided parameters.
+    * Resource pattern type would default to PatternType.LITERAL.
+    *
+    * @param resourceType non-null resource type
+    * @param name         non-null resource name
+    * @deprecated Since 2.0, use [[kafka.security.auth.Resource(ResourceType, String, PatternType)]]
+    */
+  @deprecated("Use Resource(ResourceType, String, PatternType", "Since 2.0")
+  def this(resourceType: ResourceType, name: String) = {
+    this(resourceType, name, PatternType.LITERAL)
+  }
+
+  def toPattern: ResourcePattern = {
+    new ResourcePattern(resourceType.toJava, name, patternType)
+  }
 
   override def toString: String = {
-    resourceType.name + Resource.Separator + name
+    resourceType.name + Resource.Separator + patternType + Resource.Separator + name
   }
 }
 
diff --git a/core/src/main/scala/kafka/security/auth/ResourceType.scala b/core/src/main/scala/kafka/security/auth/ResourceType.scala
index b046dddc00b9c..72b71c0b991dd 100644
--- a/core/src/main/scala/kafka/security/auth/ResourceType.scala
+++ b/core/src/main/scala/kafka/security/auth/ResourceType.scala
@@ -20,35 +20,57 @@ import kafka.common.{BaseEnum, KafkaException}
 import org.apache.kafka.common.protocol.Errors
 import org.apache.kafka.common.resource.{ResourceType => JResourceType}
 
-sealed trait ResourceType extends BaseEnum {
+@deprecated("Use org.apache.kafka.common.resource.ResourceType", "Since 2.5")
+sealed trait ResourceType extends BaseEnum with Ordered[ ResourceType ] {
   def error: Errors
   def toJava: JResourceType
+  // this method output will not include "All" Operation type
+  def supportedOperations: Set[Operation]
+
+  override def compare(that: ResourceType): Int = this.name compare that.name
 }
 
+@deprecated("Use org.apache.kafka.common.resource.ResourceType", "Since 2.5")
 case object Topic extends ResourceType {
   val name = "Topic"
   val error = Errors.TOPIC_AUTHORIZATION_FAILED
   val toJava = JResourceType.TOPIC
+  val supportedOperations = Set(Read, Write, Create, Describe, Delete, Alter, DescribeConfigs, AlterConfigs)
 }
 
+@deprecated("Use org.apache.kafka.common.resource.ResourceType", "Since 2.5")
 case object Group extends ResourceType {
   val name = "Group"
   val error = Errors.GROUP_AUTHORIZATION_FAILED
   val toJava = JResourceType.GROUP
+  val supportedOperations = Set(Read, Describe, Delete)
 }
 
+@deprecated("Use org.apache.kafka.common.resource.ResourceType", "Since 2.5")
 case object Cluster extends ResourceType {
   val name = "Cluster"
   val error = Errors.CLUSTER_AUTHORIZATION_FAILED
   val toJava = JResourceType.CLUSTER
+  val supportedOperations = Set(Create, ClusterAction, DescribeConfigs, AlterConfigs, IdempotentWrite, Alter, Describe)
 }
 
+@deprecated("Use org.apache.kafka.common.resource.ResourceType", "Since 2.5")
 case object TransactionalId extends ResourceType {
   val name = "TransactionalId"
   val error = Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED
   val toJava = JResourceType.TRANSACTIONAL_ID
+  val supportedOperations = Set(Describe, Write)
+}
+
+@deprecated("Use org.apache.kafka.common.resource.ResourceType", "Since 2.5")
+case object DelegationToken extends ResourceType {
+  val name = "DelegationToken"
+  val error = Errors.DELEGATION_TOKEN_AUTHORIZATION_FAILED
+  val toJava = JResourceType.DELEGATION_TOKEN
+  val supportedOperations : Set[Operation] = Set(Describe)
 }
 
+@deprecated("Use org.apache.kafka.common.resource.ResourceType", "Since 2.5")
 object ResourceType {
 
   def fromString(resourceType: String): ResourceType = {
@@ -56,7 +78,16 @@ object ResourceType {
     rType.getOrElse(throw new KafkaException(resourceType + " not a valid resourceType name. The valid names are " + values.mkString(",")))
   }
 
-  def values: Seq[ResourceType] = List(Topic, Group, Cluster, TransactionalId)
+  def fromJava(resourceType: JResourceType): ResourceType = {
+    resourceType match {
+      case JResourceType.TOPIC => Topic
+      case JResourceType.GROUP => Group
+      case JResourceType.CLUSTER => Cluster
+      case JResourceType.TRANSACTIONAL_ID => TransactionalId
+      case JResourceType.DELEGATION_TOKEN => DelegationToken
+      case _ => throw new KafkaException(resourceType + " is not a convertible resource type. The valid types are " + values.mkString(","))
+    }
+  }
 
-  def fromJava(operation: JResourceType): ResourceType = fromString(operation.toString.replaceAll("_", ""))
+  def values: Seq[ResourceType] = List(Topic, Group, Cluster, TransactionalId, DelegationToken)
 }
diff --git a/core/src/main/scala/kafka/security/auth/SimpleAclAuthorizer.scala b/core/src/main/scala/kafka/security/auth/SimpleAclAuthorizer.scala
index e1befc785f16a..653e154f7158b 100644
--- a/core/src/main/scala/kafka/security/auth/SimpleAclAuthorizer.scala
+++ b/core/src/main/scala/kafka/security/auth/SimpleAclAuthorizer.scala
@@ -17,309 +17,159 @@
 package kafka.security.auth
 
 import java.util
-import java.util.concurrent.locks.ReentrantReadWriteLock
 
-import com.typesafe.scalalogging.Logger
-import kafka.common.{NotificationHandler, ZkNodeChangeNotificationListener}
 import kafka.network.RequestChannel.Session
-import kafka.security.auth.SimpleAclAuthorizer.VersionedAcls
-import kafka.server.KafkaConfig
-import kafka.utils.CoreUtils.{inReadLock, inWriteLock}
+import kafka.security.auth.SimpleAclAuthorizer.BaseAuthorizer
+import kafka.security.authorizer.{AclAuthorizer, AuthorizerUtils, AuthorizerWrapper}
 import kafka.utils._
-import kafka.zk.{AclChangeNotificationSequenceZNode, AclChangeNotificationZNode, KafkaZkClient}
-import kafka.zookeeper.ZooKeeperClient
+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
+import org.apache.kafka.server.authorizer.{Action, AuthorizableRequestContext, AuthorizationResult}
 
-import scala.collection.JavaConverters._
-import scala.util.Random
+import scala.collection.mutable
+import scala.jdk.CollectionConverters._
 
+@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 = AuthorizerWrapper.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"
+      }
 
-  case class VersionedAcls(acls: Set[Acl], zkVersion: Int)
+      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 = null
-  private var aclChangeListener: ZkNodeChangeNotificationListener = null
 
-  private val aclCache = new scala.collection.mutable.HashMap[Resource, VersionedAcls]
-  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 zooKeeperClient = new ZooKeeperClient(zkUrl, zkSessionTimeOutMs, zkConnectionTimeoutMs, zkMaxInFlightRequests)
-
-    zkClient = new KafkaZkClient(zooKeeperClient, kafkaConfig.zkEnableSecureAcls)
-    zkClient.createAclPaths()
-
-    loadCache()
-
-    aclChangeListener = new ZkNodeChangeNotificationListener(zkClient, AclChangeNotificationZNode.path, AclChangeNotificationSequenceZNode.SequenceNumberPrefix, AclChangedNotificationHandler)
-    aclChangeListener.init()
+  override def configure(javaConfigs: util.Map[String, _]): Unit = {
+    aclAuthorizer.configure(javaConfigs)
   }
 
   override def authorize(session: Session, operation: Operation, resource: Resource): Boolean = {
-    val principal = session.principal
-    val host = session.clientAddress.getHostAddress
-    val acls = getAcls(resource) ++ getAcls(new Resource(resource.resourceType, Resource.WildCardResource))
-
-    // Check if there is any Deny acl match that would disallow this operation.
-    val denyMatch = aclMatch(operation, resource, principal, host, Deny, acls)
-
-    // 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)
-    }
-    val allowMatch = allowOps.exists(operation => aclMatch(operation, resource, principal, host, Allow, acls))
-
-    //we allow an operation if a user is a super user or 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 authorized = isSuperUser(operation, resource, principal, host) ||
-      isEmptyAclAndAuthorized(operation, resource, principal, host, acls) ||
-      (!denyMatch && allowMatch)
-
-    logAuditMessage(principal, authorized, operation, resource, host)
-    authorized
-  }
-
-  def isEmptyAclAndAuthorized(operation: Operation, resource: Resource, principal: KafkaPrincipal, host: String, acls: Set[Acl]): Boolean = {
-    if (acls.isEmpty) {
-      authorizerLogger.debug(s"No acl found for resource $resource, authorized = $shouldAllowEveryoneIfNoAclIsFound")
-      shouldAllowEveryoneIfNoAclIsFound
-    } else false
+    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(operations: 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) &&
-        (operations == acl.operation || acl.operation == All) &&
-        (acl.host == host || acl.host == Acl.WildCardHost)
-    }.exists { acl =>
-      authorizerLogger.debug(s"operation = $operations 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) {
-      inWriteLock(lock) {
-        updateResourceAcls(resource) { currentAcls =>
-          currentAcls ++ acls
-        }
-      }
+      val bindings = acls.map { acl => AuthorizerWrapper.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, AuthorizerWrapper.convertToAccessControlEntry(acl).toFilter)
     }
+    deleteAcls(filters)
   }
 
   override def removeAcls(resource: Resource): Boolean = {
-    inWriteLock(lock) {
-      val result = zkClient.deleteResource(resource)
-      updateCache(resource, VersionedAcls(Set(), 0))
-      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
-      }.toMap
-    }
+    val filter = new AclBindingFilter(ResourcePatternFilter.ANY,
+      new AccessControlEntryFilter(principal.toString, null, AclOperation.ANY, AclPermissionType.ANY))
+    acls(filter)
   }
 
-  override def getAcls(): Map[Resource, Set[Acl]] = {
-    inReadLock(lock) {
-      aclCache.mapValues(_.acls).toMap
-    }
+  def getMatchingAcls(resourceType: ResourceType, resourceName: String): Set[Acl] = {
+    val filter = new AclBindingFilter(new ResourcePatternFilter(resourceType.toJava, resourceName, PatternType.MATCH),
+      AccessControlEntryFilter.ANY)
+    acls(filter).flatMap(_._2).toSet
   }
 
-  def close() {
-    if (aclChangeListener != null) aclChangeListener.close()
-    if (zkClient != null) zkClient.close()
+  override def getAcls(): Map[Resource, Set[Acl]] = {
+    acls(AclBindingFilter.ANY)
   }
 
-  private def loadCache()  {
-    inWriteLock(lock) {
-      val resourceTypes = zkClient.getResourceTypes()
-      for (rType <- resourceTypes) {
-        val resourceType = ResourceType.fromString(rType)
-        val resourceNames = zkClient.getResourceNames(resourceType.name)
-        for (resourceName <- resourceNames) {
-          val versionedAcls = getAclsFromZk(Resource(resourceType, resourceName))
-          updateCache(new Resource(resourceType, resourceName), versionedAcls)
-        }
-      }
-    }
+  def close(): Unit = {
+    aclAuthorizer.close()
   }
 
-  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)
+  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.ifPresent(throwException) }
   }
 
-  /**
-    * 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) {
-          zkClient.conditionalSetOrCreateAclsForResource(resource, newAcls, currentVersionedAcls.zkVersion)
-        } 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 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.ifPresent(throwException) }
+    results.flatMap(_.aclBindingDeleteResults.asScala).foreach { result => result.exception.ifPresent(e => throw e) }
+    results.exists(r => r.aclBindingDeleteResults.asScala.exists(d => !d.exception.isPresent))
   }
 
-  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) {
-    if (versionedAcls.acls.nonEmpty) {
-      aclCache.put(resource, versionedAcls)
-    } else {
-      aclCache.remove(resource)
+  private def acls(filter: AclBindingFilter): Map[Resource, Set[Acl]] = {
+    val result = mutable.Map[Resource, mutable.Set[Acl]]()
+    aclAuthorizer.acls(filter).forEach { binding =>
+      val resource = AuthorizerWrapper.convertToResource(binding.pattern)
+      val acl = AuthorizerWrapper.convertToAcl(binding.entry)
+      result.getOrElseUpdate(resource, mutable.Set()).add(acl)
     }
+    result.mapValues(_.toSet).toMap
   }
 
-  private def updateAclChangedFlag(resource: Resource) {
-    zkClient.createAclChangeNotification(resource.toString)
-  }
-
-  private def backoffTime = {
-    retryBackoffMs + Random.nextInt(retryBackoffJitterMs)
+  // 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
   }
-
-  object AclChangedNotificationHandler extends NotificationHandler {
-    override def processNotification(notificationMessage: String) {
-      val resource: Resource = Resource.fromString(notificationMessage)
-      inWriteLock(lock) {
-        val versionedAcls = getAclsFromZk(resource)
-        updateCache(resource, versionedAcls)
-      }
-    }
-  }
-
-}
\ No newline at end of file
+}
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..0a60e51659c34
--- /dev/null
+++ b/core/src/main/scala/kafka/security/authorizer/AclAuthorizer.scala
@@ -0,0 +1,585 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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 com.typesafe.scalalogging.Logger
+import kafka.api.KAFKA_2_0_IV1
+import kafka.security.authorizer.AclAuthorizer.{AclSeqs, ResourceOrdering, VersionedAcls}
+import kafka.security.authorizer.AclEntry.ResourceSeparator
+import kafka.server.{KafkaConfig, KafkaServer}
+import kafka.utils._
+import kafka.utils.Implicits._
+import kafka.zk._
+import org.apache.kafka.common.Endpoint
+import org.apache.kafka.common.acl._
+import org.apache.kafka.common.acl.AclOperation._
+import org.apache.kafka.common.acl.AclPermissionType.{ALLOW, DENY}
+import org.apache.kafka.common.errors.{ApiException, InvalidRequestException, UnsupportedVersionException}
+import org.apache.kafka.common.protocol.ApiKeys
+import org.apache.kafka.common.resource._
+import org.apache.kafka.common.security.auth.KafkaPrincipal
+import org.apache.kafka.common.utils.{SecurityUtils, Time}
+import org.apache.kafka.server.authorizer.AclDeleteResult.AclBindingDeleteResult
+import org.apache.kafka.server.authorizer._
+import org.apache.zookeeper.client.ZKClientConfig
+
+import scala.annotation.nowarn
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.{Seq, mutable}
+import scala.jdk.CollectionConverters._
+import scala.util.{Failure, Random, Success, Try}
+
+object AclAuthorizer {
+  // Optional override zookeeper cluster configuration where acls will be stored. If not specified,
+  // acls will be stored in the same zookeeper where all other kafka broker metadata is stored.
+  val configPrefix = "authorizer."
+  val ZkUrlProp = s"${configPrefix}zookeeper.url"
+  val ZkConnectionTimeOutProp = s"${configPrefix}zookeeper.connection.timeout.ms"
+  val ZkSessionTimeOutProp = s"${configPrefix}zookeeper.session.timeout.ms"
+  val ZkMaxInFlightRequests = s"${configPrefix}zookeeper.max.in.flight.requests"
+
+  // 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[AclEntry], zkVersion: Int) {
+    def exists: Boolean = zkVersion != ZkVersion.UnknownVersion
+  }
+
+  class AclSeqs(seqs: Seq[AclEntry]*) {
+    def find(p: AclEntry => Boolean): Option[AclEntry] = {
+      // Lazily iterate through the inner `Seq` elements and stop as soon as we find a match
+      val it = seqs.iterator.flatMap(_.find(p))
+      if (it.hasNext) Some(it.next())
+      else None
+    }
+
+    def isEmpty: Boolean = !seqs.exists(_.nonEmpty)
+  }
+
+  val NoAcls = VersionedAcls(Set.empty, ZkVersion.UnknownVersion)
+  val WildcardHost = "*"
+
+  // Orders by resource type, then resource pattern type and finally reverse ordering by name.
+  class ResourceOrdering extends Ordering[ResourcePattern] {
+
+    def compare(a: ResourcePattern, b: ResourcePattern): Int = {
+      val rt = a.resourceType.compareTo(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
+      }
+    }
+  }
+
+  private[authorizer] def zkClientConfigFromKafkaConfigAndMap(kafkaConfig: KafkaConfig, configMap: mutable.Map[String, _<:Any]): Option[ZKClientConfig] = {
+    val zkSslClientEnable = configMap.get(AclAuthorizer.configPrefix + KafkaConfig.ZkSslClientEnableProp).
+      map(_.toString).getOrElse(kafkaConfig.zkSslClientEnable.toString).toBoolean
+    if (!zkSslClientEnable)
+      None
+    else {
+      // start with the base config from the Kafka configuration
+      // be sure to force creation since the zkSslClientEnable property in the kafkaConfig could be false
+      val zkClientConfig = KafkaServer.zkClientConfigFromKafkaConfig(kafkaConfig, true)
+      // add in any prefixed overlays
+      KafkaConfig.ZkSslConfigToSystemPropertyMap.foreach{ case (kafkaProp, sysProp) => {
+        val prefixedValue = configMap.get(AclAuthorizer.configPrefix + kafkaProp)
+        if (prefixedValue.isDefined)
+          zkClientConfig.get.setProperty(sysProp,
+            if (kafkaProp == KafkaConfig.ZkSslEndpointIdentificationAlgorithmProp)
+              (prefixedValue.get.toString.toUpperCase == "HTTPS").toString
+            else
+              prefixedValue.get.toString)
+      }}
+      zkClientConfig
+    }
+  }
+}
+
+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[ResourcePattern, VersionedAcls]()(new ResourceOrdering)
+  private val lock = new Object()
+
+  // 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.forKeyValue { (key, value) => props.put(key, value.toString) }
+
+    superUsers = configs.get(AclAuthorizer.SuperUsersProp).collect {
+      case str: String if str.nonEmpty => str.split(";").map(s => SecurityUtils.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 zkClientConfig = AclAuthorizer.zkClientConfigFromKafkaConfigAndMap(kafkaConfig, configs)
+    val time = Time.SYSTEM
+    zkClient = KafkaZkClient(zkUrl, kafkaConfig.zkEnableSecureAcls, zkSessionTimeOutMs, zkConnectionTimeoutMs,
+      zkMaxInFlightRequests, time, "kafka.security", "AclAuthorizer", name=Some("ACL authorizer"),
+      zkClientConfig = zkClientConfig)
+    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) {
+      lock synchronized {
+        aclsToCreate.forKeyValue { (resource, aclsWithIndex) =>
+          try {
+            updateResourceAcls(resource) { currentAcls =>
+              val newAcls = aclsWithIndex.map { case (acl, _) => new AclEntry(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
+  }
+
+  /**
+   *
+   * Concurrent updates:
+   * 
                
+   *   
              • If ACLs are created using [[kafka.security.authorizer.AclAuthorizer#createAcls]] while a delete is in
+   *   progress, these ACLs may or may not be considered for deletion depending on the order of updates.
+   *   The returned [[org.apache.kafka.server.authorizer.AclDeleteResult]] indicates which ACLs were deleted.
                • 
+   *   
              • If the provided filters use resource pattern type
+   *   [[org.apache.kafka.common.resource.PatternType#MATCH]] that needs to filter all resources to determine
+   *   matching ACLs, only ACLs that have already been propagated to the broker processing the ACL update will be
+   *   deleted. This may not include some ACLs that were persisted, but not yet propagated to all brokers. The
+   *   returned [[org.apache.kafka.server.authorizer.AclDeleteResult]] indicates which ACLs were deleted.
                • 
+   *   
              • If the provided filters use other resource pattern types that perform a direct match, all matching ACLs
+   *   from previously completed [[kafka.security.authorizer.AclAuthorizer#createAcls]] are guaranteed to be deleted.
                • 
+   * 
              
+   */
+  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
+    lock synchronized {
+      // Find all potentially matching resource patterns from the provided filters and ACL cache and apply the filters
+      val resources = aclCache.keys ++ filters.map(_._1.patternFilter).filter(_.matchesAtMostOne).flatMap(filterToResources)
+      val resourcesToUpdate = resources.map { resource =>
+        val matchingFilters = filters.filter { case (filter, _) =>
+          filter.patternFilter.matches(resource)
+        }
+        resource -> matchingFilters
+      }.toMap.filter(_._2.nonEmpty)
+
+      resourcesToUpdate.forKeyValue { (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(acl)
+                if (matches) {
+                  val binding = new AclBinding(resource, acl)
+                  deletedBindings.getOrElseUpdate(binding, index)
+                  resourceBindingsBeingDeleted.getOrElseUpdate(binding, index)
+                }
+                matches
+              }
+            }
+            currentAcls -- aclsToRemove
+          }
+        } catch {
+          case e: Exception =>
+            resourceBindingsBeingDeleted.keys.foreach { binding =>
+              deleteExceptions.getOrElseUpdate(binding, apiException(e))
+            }
+        }
+      }
+    }
+    val deletedResult = deletedBindings.groupBy(_._2).map { case (k, bindings) =>
+      k -> bindings.keys.map { binding => new AclBindingDeleteResult(binding, deleteExceptions.get(binding).orNull) }
+    }
+    (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] = {
+    val aclBindings = new util.ArrayList[AclBinding]()
+    aclCache.forKeyValue { case (resource, versionedAcls) =>
+      versionedAcls.acls.foreach { acl =>
+        val binding = new AclBinding(resource, acl.ace)
+        if (filter.matches(binding))
+          aclBindings.add(binding)
+      }
+    }
+    aclBindings
+  }
+
+  override def close(): Unit = {
+    aclChangeListeners.foreach(listener => listener.close())
+    if (zkClient != null) zkClient.close()
+  }
+
+  private def authorizeAction(requestContext: AuthorizableRequestContext, action: Action): AuthorizationResult = {
+    val resource = 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 = action.operation
+
+    def isEmptyAclAndAuthorized(acls: AclSeqs): 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: AclSeqs): Boolean = {
+      // Check if there are any Deny ACLs which would forbid this operation.
+      matchingAclExists(operation, resource, principal, host, DENY, acls)
+    }
+
+    def allowAclExists(acls: AclSeqs): 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[AclOperation](DESCRIBE, READ, WRITE, DELETE, ALTER)
+        case DESCRIBE_CONFIGS => Set[AclOperation](DESCRIBE_CONFIGS, ALTER_CONFIGS)
+        case _ => Set[AclOperation](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
+  }
+
+  @nowarn("cat=deprecation")
+  private def matchingAcls(resourceType: ResourceType, resourceName: String): AclSeqs = {
+    // this code is performance sensitive, make sure to run AclAuthorizerBenchmark after any changes
+
+    // save aclCache reference to a local val to get a consistent view of the cache during acl updates.
+    val aclCacheSnapshot = aclCache
+    val wildcard = aclCacheSnapshot.get(new ResourcePattern(resourceType, ResourcePattern.WILDCARD_RESOURCE, PatternType.LITERAL))
+      .map(_.acls.toBuffer)
+      .getOrElse(mutable.Buffer.empty)
+
+    val literal = aclCacheSnapshot.get(new ResourcePattern(resourceType, resourceName, PatternType.LITERAL))
+      .map(_.acls.toBuffer)
+      .getOrElse(mutable.Buffer.empty)
+
+    val prefixed = new ArrayBuffer[AclEntry]
+    aclCacheSnapshot
+      .from(new ResourcePattern(resourceType, resourceName, PatternType.PREFIXED))
+      .to(new ResourcePattern(resourceType, resourceName.take(1), PatternType.PREFIXED))
+      .forKeyValue { (resource, acls) =>
+        if (resourceName.startsWith(resource.name)) prefixed ++= acls.acls
+      }
+
+    new AclSeqs(prefixed, wildcard, literal)
+  }
+
+  private def matchingAclExists(operation: AclOperation,
+                                resource: ResourcePattern,
+                                principal: KafkaPrincipal,
+                                host: String,
+                                permissionType: AclPermissionType,
+                                acls: AclSeqs): Boolean = {
+    acls.find { acl =>
+      acl.permissionType == permissionType &&
+        (acl.kafkaPrincipal == principal || acl.kafkaPrincipal == AclEntry.WildcardPrincipal) &&
+        (operation == acl.operation || acl.operation == AclOperation.ALL) &&
+        (acl.host == host || acl.host == AclEntry.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 = {
+    lock synchronized  {
+      ZkAclStore.stores.foreach(store => {
+        val resourceTypes = zkClient.getResourceTypes(store.patternType)
+        for (rType <- resourceTypes) {
+          val resourceType = Try(SecurityUtils.resourceType(rType))
+          resourceType match {
+            case Success(resourceTypeObj) =>
+              val resourceNames = zkClient.getResourceNames(store.patternType, resourceTypeObj)
+              for (resourceName <- resourceNames) {
+                val resource = new ResourcePattern(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))
+  }
+
+  private def filterToResources(filter: ResourcePatternFilter): Set[ResourcePattern] = {
+    filter.patternType match {
+      case PatternType.LITERAL | PatternType.PREFIXED =>
+        Set(new ResourcePattern(filter.resourceType, filter.name, filter.patternType))
+      case PatternType.ANY =>
+        Set(new ResourcePattern(filter.resourceType, filter.name, PatternType.LITERAL),
+          new ResourcePattern(filter.resourceType, filter.name, PatternType.PREFIXED))
+      case _ => throw new IllegalArgumentException(s"Cannot determine matching resources for patternType $filter")
+    }
+  }
+
+  def logAuditMessage(requestContext: AuthorizableRequestContext, action: Action, authorized: Boolean): Unit = {
+    def logMessage: String = {
+      val principal = requestContext.principal
+      val operation = SecurityUtils.operationName(action.operation)
+      val host = requestContext.clientAddress.getHostAddress
+      val resourceType = SecurityUtils.resourceTypeName(action.resourcePattern.resourceType)
+      val resource = s"$resourceType$ResourceSeparator${action.resourcePattern.patternType}$ResourceSeparator${action.resourcePattern.name}"
+      val 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 " +
+        s"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: ResourcePattern)(getNewAcls: Set[AclEntry] => Set[AclEntry]): 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) {
+      info(s"Updated ACLs for $resource with new version ${newVersionedAcls.zkVersion}")
+      debug(s"Updated ACLs for $resource to $newVersionedAcls")
+      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: ResourcePattern): VersionedAcls = {
+    aclCache.getOrElse(resource, throw new IllegalArgumentException(s"ACLs do not exist in the cache for resource $resource"))
+  }
+
+  private def getAclsFromZk(resource: ResourcePattern): VersionedAcls = {
+    zkClient.getVersionedAclsForResource(resource)
+  }
+
+  private def updateCache(resource: ResourcePattern, versionedAcls: VersionedAcls): Unit = {
+    if (versionedAcls.acls.nonEmpty) {
+      aclCache = aclCache.updated(resource, versionedAcls)
+    } else {
+      aclCache -= resource
+    }
+  }
+
+  private def updateAclChangedFlag(resource: ResourcePattern): 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)
+    }
+  }
+
+  private[authorizer] def processAclChangeNotification(resource: ResourcePattern): Unit = {
+    lock synchronized {
+      val versionedAcls = getAclsFromZk(resource)
+      updateCache(resource, versionedAcls)
+    }
+  }
+
+  object AclChangedNotificationHandler extends AclChangeNotificationHandler {
+    override def processNotification(resource: ResourcePattern): Unit = {
+      processAclChangeNotification(resource)
+    }
+  }
+}
diff --git a/core/src/main/scala/kafka/security/authorizer/AclEntry.scala b/core/src/main/scala/kafka/security/authorizer/AclEntry.scala
new file mode 100644
index 0000000000000..2014916e7e4ba
--- /dev/null
+++ b/core/src/main/scala/kafka/security/authorizer/AclEntry.scala
@@ -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 kafka.security.authorizer
+
+import kafka.utils.Json
+import org.apache.kafka.common.acl.{AccessControlEntry, AclOperation, AclPermissionType}
+import org.apache.kafka.common.acl.AclOperation.{READ, WRITE, CREATE, DESCRIBE, DELETE, ALTER, DESCRIBE_CONFIGS, ALTER_CONFIGS, CLUSTER_ACTION, IDEMPOTENT_WRITE}
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.resource.{ResourcePattern, ResourceType}
+import org.apache.kafka.common.security.auth.KafkaPrincipal
+import org.apache.kafka.common.utils.SecurityUtils
+
+import scala.jdk.CollectionConverters._
+
+object AclEntry {
+  val WildcardPrincipal: KafkaPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "*")
+  val WildcardPrincipalString: String = WildcardPrincipal.toString
+  val WildcardHost: String = "*"
+  val WildcardResource: String = ResourcePattern.WILDCARD_RESOURCE
+
+  val ResourceSeparator = ":"
+  val ResourceTypes: Set[ResourceType] = ResourceType.values.toSet
+    .filterNot(t => t == ResourceType.UNKNOWN || t == ResourceType.ANY)
+  val AclOperations: Set[AclOperation] = AclOperation.values.toSet
+    .filterNot(t => t == AclOperation.UNKNOWN || t == AclOperation.ANY)
+
+  val PrincipalKey = "principal"
+  val PermissionTypeKey = "permissionType"
+  val OperationKey = "operation"
+  val HostsKey = "host"
+  val VersionKey = "version"
+  val CurrentVersion = 1
+  val AclsKey = "acls"
+
+  def apply(principal: KafkaPrincipal,
+            permissionType: AclPermissionType,
+            host: String,
+            operation: AclOperation): AclEntry = {
+    new AclEntry(new AccessControlEntry(if (principal == null) null else principal.toString,
+      host, operation, permissionType))
+  }
+
+  /**
+   * Parse JSON representation of ACLs
+   * @param bytes of acls json string
+   *
+   * 
              
+      {
+        "version": 1,
+        "acls": [
+          {
+            "host":"host1",
+            "permissionType": "Deny",
+            "operation": "Read",
+            "principal": "User:alice"
+          }
+        ]
+      }
+   * 
              
+   *
+   * @return set of AclEntry objects from the JSON string
+   */
+  def fromBytes(bytes: Array[Byte]): Set[AclEntry] = {
+    if (bytes == null || bytes.isEmpty)
+      return collection.immutable.Set.empty[AclEntry]
+
+    Json.parseBytes(bytes).map(_.asJsonObject).map { js =>
+      //the acl json version.
+      require(js(VersionKey).to[Int] == CurrentVersion)
+      js(AclsKey).asJsonArray.iterator.map(_.asJsonObject).map { itemJs =>
+        val principal = SecurityUtils.parseKafkaPrincipal(itemJs(PrincipalKey).to[String])
+        val permissionType = SecurityUtils.permissionType(itemJs(PermissionTypeKey).to[String])
+        val host = itemJs(HostsKey).to[String]
+        val operation = SecurityUtils.operation(itemJs(OperationKey).to[String])
+        AclEntry(principal, permissionType, host, operation)
+      }.toSet
+    }.getOrElse(Set.empty)
+  }
+
+  def toJsonCompatibleMap(acls: Set[AclEntry]): Map[String, Any] = {
+    Map(AclEntry.VersionKey -> AclEntry.CurrentVersion, AclEntry.AclsKey -> acls.map(acl => acl.toMap.asJava).toList.asJava)
+  }
+
+  def supportedOperations(resourceType: ResourceType): Set[AclOperation] = {
+    resourceType match {
+      case ResourceType.TOPIC => Set(READ, WRITE, CREATE, DESCRIBE, DELETE, ALTER, DESCRIBE_CONFIGS, ALTER_CONFIGS)
+      case ResourceType.GROUP => Set(READ, DESCRIBE, DELETE)
+      case ResourceType.CLUSTER => Set(CREATE, CLUSTER_ACTION, DESCRIBE_CONFIGS, ALTER_CONFIGS, IDEMPOTENT_WRITE, ALTER, DESCRIBE)
+      case ResourceType.TRANSACTIONAL_ID => Set(DESCRIBE, WRITE)
+      case ResourceType.DELEGATION_TOKEN => Set(DESCRIBE)
+      case _ => throw new IllegalArgumentException("Not a concrete resource type")
+    }
+  }
+
+  def authorizationError(resourceType: ResourceType): Errors = {
+    resourceType match {
+      case ResourceType.TOPIC => Errors.TOPIC_AUTHORIZATION_FAILED
+      case ResourceType.GROUP => Errors.GROUP_AUTHORIZATION_FAILED
+      case ResourceType.CLUSTER => Errors.CLUSTER_AUTHORIZATION_FAILED
+      case ResourceType.TRANSACTIONAL_ID => Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED
+      case ResourceType.DELEGATION_TOKEN => Errors.DELEGATION_TOKEN_AUTHORIZATION_FAILED
+      case _ => throw new IllegalArgumentException("Authorization error type not known")
+    }
+  }
+}
+
+class AclEntry(val ace: AccessControlEntry)
+  extends AccessControlEntry(ace.principal, ace.host, ace.operation, ace.permissionType) {
+
+  val kafkaPrincipal: KafkaPrincipal = if (principal == null)
+    null
+  else
+    SecurityUtils.parseKafkaPrincipal(principal)
+
+  def toMap: Map[String, Any] = {
+    Map(AclEntry.PrincipalKey -> principal,
+      AclEntry.PermissionTypeKey -> SecurityUtils.permissionTypeName(permissionType),
+      AclEntry.OperationKey -> SecurityUtils.operationName(operation),
+      AclEntry.HostsKey -> host)
+  }
+
+  override def hashCode(): Int = ace.hashCode()
+
+  override def equals(o: scala.Any): Boolean = super.equals(o) // to keep spotbugs happy
+
+  override def toString: String = {
+    "%s has %s permission for operations: %s from hosts: %s".format(principal, permissionType.name, operation, host)
+  }
+
+}
+
diff --git a/core/src/main/scala/kafka/security/authorizer/AuthorizerUtils.scala b/core/src/main/scala/kafka/security/authorizer/AuthorizerUtils.scala
new file mode 100644
index 0000000000000..0d670befbb7ff
--- /dev/null
+++ b/core/src/main/scala/kafka/security/authorizer/AuthorizerUtils.scala
@@ -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 kafka.security.authorizer
+
+import java.net.InetAddress
+
+import kafka.network.RequestChannel.Session
+import kafka.security.auth.{Authorizer => LegacyAuthorizer}
+import org.apache.kafka.common.acl._
+import org.apache.kafka.common.config.ConfigException
+import org.apache.kafka.common.resource.Resource
+import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
+import org.apache.kafka.common.utils.Utils
+import org.apache.kafka.server.authorizer.{AuthorizableRequestContext, Authorizer}
+
+import scala.annotation.nowarn
+
+
+object AuthorizerUtils {
+
+  @nowarn("cat=deprecation")
+  def createAuthorizer(className: String): Authorizer = {
+    Utils.newInstance(className, classOf[Object]) match {
+      case auth: Authorizer => auth
+      case auth: kafka.security.auth.Authorizer => new AuthorizerWrapper(auth)
+      case _ => throw new ConfigException(s"Authorizer does not implement ${classOf[Authorizer].getName} or ${classOf[LegacyAuthorizer].getName}.")
+    }
+  }
+
+  def validateAclBinding(aclBinding: AclBinding): Unit = {
+    if (aclBinding.isUnknown)
+      throw new IllegalArgumentException("ACL binding contains unknown elements")
+  }
+
+  def isClusterResource(name: String): Boolean = name.equals(Resource.CLUSTER_NAME)
+
+  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..9a8bf9da49fa7
--- /dev/null
+++ b/core/src/main/scala/kafka/security/authorizer/AuthorizerWrapper.scala
@@ -0,0 +1,178 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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, PermissionType, Resource, ResourceType}
+import kafka.security.authorizer.AuthorizerWrapper._
+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.protocol.Errors
+import org.apache.kafka.common.requests.ApiError
+import org.apache.kafka.common.resource.ResourcePattern
+import org.apache.kafka.common.utils.SecurityUtils.parseKafkaPrincipal
+import org.apache.kafka.server.authorizer.AclDeleteResult.AclBindingDeleteResult
+import org.apache.kafka.server.authorizer.{AuthorizableRequestContext, AuthorizerServerInfo, _}
+
+import scala.jdk.CollectionConverters._
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.{Seq, immutable, mutable}
+import scala.util.{Failure, Success, Try}
+
+@deprecated("Use kafka.security.authorizer.AclAuthorizer", "Since 2.5")
+object AuthorizerWrapper {
+
+  def convertToResourceAndAcl(filter: AclBindingFilter): Either[ApiError, (Resource, Acl)] = {
+    (for {
+      resourceType <- Try(ResourceType.fromJava(filter.patternFilter.resourceType))
+      principal <- Try(parseKafkaPrincipal(filter.entryFilter.principal))
+      operation <- Try(Operation.fromJava(filter.entryFilter.operation))
+      permissionType <- Try(PermissionType.fromJava(filter.entryFilter.permissionType))
+      resource = Resource(resourceType, filter.patternFilter.name, filter.patternFilter.patternType)
+      acl = Acl(principal, permissionType, filter.entryFilter.host, operation)
+    } yield (resource, acl)) match {
+      case Failure(throwable) => Left(new ApiError(Errors.INVALID_REQUEST, throwable.getMessage))
+      case Success(s) => Right(s)
+    }
+  }
+
+  def convertToAclBinding(resource: Resource, acl: Acl): AclBinding = {
+    val resourcePattern = new ResourcePattern(resource.resourceType.toJava, resource.name, resource.patternType)
+    new AclBinding(resourcePattern, convertToAccessControlEntry(acl))
+  }
+
+  def convertToAccessControlEntry(acl: Acl): AccessControlEntry = {
+    new AccessControlEntry(acl.principal.toString, acl.host.toString,
+      acl.operation.toJava, acl.permissionType.toJava)
+  }
+
+  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)
+  }
+}
+
+@deprecated("Use kafka.security.authorizer.AclAuthorizer", "Since 2.5")
+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, 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 =>
+        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) {
+        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 = 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 => convertToAclBinding(resource, acl)).filter(filter.matches)
+    }.asJava
+  }
+
+  override def close(): Unit = {
+    baseAuthorizer.close()
+  }
+}
diff --git a/core/src/main/scala/kafka/serializer/Encoder.scala b/core/src/main/scala/kafka/serializer/Encoder.scala
deleted file mode 100644
index b1277e1931595..0000000000000
--- a/core/src/main/scala/kafka/serializer/Encoder.scala
+++ /dev/null
@@ -1,79 +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.serializer
-
-import java.nio.ByteBuffer
-
-import kafka.utils.VerifiableProperties
-
-/**
- * An encoder is a method of turning objects into byte arrays.
- * An implementation is required to provide a constructor that
- * takes a VerifiableProperties instance.
- */
-trait Encoder[T] {
-  def toBytes(t: T): Array[Byte]
-}
-
-/**
- * The default implementation is a no-op, it just returns the same array it takes in
- */
-class DefaultEncoder(props: VerifiableProperties = null) extends Encoder[Array[Byte]] {
-  override def toBytes(value: Array[Byte]): Array[Byte] = value
-}
-
-class NullEncoder[T](props: VerifiableProperties = null) extends Encoder[T] {
-  override def toBytes(value: T): Array[Byte] = null
-}
-
-/**
- * The string encoder takes an optional parameter serializer.encoding which controls
- * the character set used in encoding the string into bytes.
- */
-class StringEncoder(props: VerifiableProperties = null) extends Encoder[String] {
-  val encoding = 
-    if(props == null) 
-      "UTF8" 
-    else 
-      props.getString("serializer.encoding", "UTF8")
-  
-  override def toBytes(s: String): Array[Byte] = 
-    if(s == null)
-      null
-    else
-      s.getBytes(encoding)
-}
-
-/**
-  * The long encoder translates longs into bytes.
-  */
-class LongEncoder(props: VerifiableProperties = null) extends Encoder[Long] {
-  override def toBytes(l: Long): Array[Byte] =
-    ByteBuffer.allocate(8).putLong(l).array()
-}
-
-/**
-  * The integer encoder translates integers into bytes.
-  */
-class IntegerEncoder(props: VerifiableProperties = null) extends Encoder[Integer] {
-  override def toBytes(i: Integer): Array[Byte] =
-    if(i == null)
-      null
-    else
-      ByteBuffer.allocate(4).putInt(i).array()
-}
diff --git a/core/src/main/scala/kafka/server/AbstractFetcherManager.scala b/core/src/main/scala/kafka/server/AbstractFetcherManager.scala
index c385d4fa8ae76..3d9189e1071dd 100755
--- a/core/src/main/scala/kafka/server/AbstractFetcherManager.scala
+++ b/core/src/main/scala/kafka/server/AbstractFetcherManager.scala
@@ -17,62 +17,98 @@
 
 package kafka.server
 
-import scala.collection.mutable
-import scala.collection.Set
-import scala.collection.Map
-import kafka.utils.Logging
 import kafka.cluster.BrokerEndPoint
 import kafka.metrics.KafkaMetricsGroup
-import com.yammer.metrics.core.Gauge
+import kafka.utils.Implicits._
+import kafka.utils.Logging
 import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.utils.Utils
 
-abstract class AbstractFetcherManager(protected val name: String, clientId: String, numFetchers: Int = 1)
+import scala.collection.{Map, Set, mutable}
+
+abstract class AbstractFetcherManager[T <: AbstractFetcherThread](val name: String, clientId: String, numFetchers: Int)
   extends Logging with KafkaMetricsGroup {
   // map of (source broker_id, fetcher_id per source broker) => fetcher.
   // package private for test
-  private[server] val fetcherThreadMap = new mutable.HashMap[BrokerIdAndFetcherId, AbstractFetcherThread]
-  private val mapLock = new Object
+  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) => {
-        fetcherThreadMapEntry._2.fetcherLagStats.stats.foldLeft(0L)((curMaxThread, fetcherLagStatsEntry) => {
-          curMaxThread.max(fetcherLagStatsEntry._2.lag)
-        }).max(curMaxAll)
-      })
-    },
-    Map("clientId" -> clientId)
-  )
-
-  newGauge(
-  "MinFetchRate", {
-    new Gauge[Double] {
-      // current min fetch rate across all fetchers/topics/partitions
-      def value = {
-        val headRate: Double =
-          fetcherThreadMap.headOption.map(_._2.fetcherStats.requestRate.oneMinuteRate).getOrElse(0)
-
-        fetcherThreadMap.foldLeft(headRate)((curMinAll, fetcherThreadMapEntry) => {
-          fetcherThreadMapEntry._2.fetcherStats.requestRate.oneMinuteRate.min(curMinAll)
-        })
+  private val tags = Map("clientId" -> clientId)
+
+  newGauge("MaxLag", () => {
+    // current max lag across all fetchers/topics/partitions
+    fetcherThreadMap.values.foldLeft(0L) { (curMaxLagAll, fetcherThread) =>
+      val maxLagThread = fetcherThread.fetcherLagStats.stats.values.foldLeft(0L)((curMaxLagThread, lagMetrics) =>
+        math.max(curMaxLagThread, lagMetrics.lag))
+      math.max(curMaxLagAll, maxLagThread)
+    }
+  }, tags)
+
+  newGauge("MinFetchRate", () => {
+    // current min fetch rate across all fetchers/topics/partitions
+    val headRate = fetcherThreadMap.values.headOption.map(_.fetcherStats.requestRate.oneMinuteRate).getOrElse(0.0)
+    fetcherThreadMap.values.foldLeft(headRate)((curMinAll, fetcherThread) =>
+      math.min(curMinAll, fetcherThread.fetcherStats.requestRate.oneMinuteRate))
+  }, tags)
+
+  newGauge("FailedPartitionsCount", () => failedPartitions.size, tags)
+
+  newGauge("DeadThreadCount", () => deadThreadCount, tags)
+
+  private[server] def deadThreadCount: Int = lock synchronized { fetcherThreadMap.values.count(_.isThreadFailed) }
+
+  def resizeThreadPool(newSize: Int): Unit = {
+    def migratePartitions(newSize: Int): Unit = {
+      fetcherThreadMap.forKeyValue { (id, thread) =>
+        val partitionStates = removeFetcherForPartitions(thread.partitions)
+        if (id.fetcherId >= newSize)
+          thread.shutdown()
+        val fetchStates = partitionStates.map { case (topicPartition, currentFetchState) =>
+          val initialFetchState = InitialFetchState(thread.sourceBroker,
+            currentLeaderEpoch = currentFetchState.currentLeaderEpoch,
+            initOffset = currentFetchState.fetchOffset)
+          topicPartition -> initialFetchState
+        }
+        addFetcherForPartitions(fetchStates)
       }
     }
-  },
-  Map("clientId" -> clientId)
-  )
+    lock synchronized {
+      val currentSize = numFetchersPerBroker
+      info(s"Resizing fetcher thread pool size from $currentSize to $newSize")
+      numFetchersPerBroker = newSize
+      if (newSize != currentSize) {
+        // We could just migrate some partitions explicitly to new threads. But this is currently
+        // reassigning all partitions using the new thread size so that hash-based allocation
+        // works with partition add/delete as it did before.
+        migratePartitions(newSize)
+      }
+      shutdownIdleFetcherThreads()
+    }
+  }
+
+  // Visible for testing
+  private[server] def getFetcher(topicPartition: TopicPartition): Option[T] = {
+    lock synchronized {
+      fetcherThreadMap.values.find { fetcherThread =>
+        fetcherThread.fetchState(topicPartition).isDefined
+      }
+    }
+  }
 
-  private def getFetcherId(topic: String, partitionId: Int) : Int = {
-    Utils.abs(31 * topic.hashCode() + partitionId) % numFetchers
+  // Visibility for testing
+  private[server] def getFetcherId(topicPartition: TopicPartition): Int = {
+    lock synchronized {
+      Utils.abs(31 * topicPartition.topic.hashCode() + topicPartition.partition) % numFetchersPerBroker
+    }
   }
 
   // This method is only needed by ReplicaAlterDirManager
-  def markPartitionsForTruncation(brokerId: Int, topicPartition: TopicPartition, truncationOffset: Long) {
-    mapLock synchronized {
-      val fetcherId = getFetcherId(topicPartition.topic, topicPartition.partition)
+  def markPartitionsForTruncation(brokerId: Int, topicPartition: TopicPartition, truncationOffset: Long): Unit = {
+    lock synchronized {
+      val fetcherId = getFetcherId(topicPartition)
       val brokerIdAndFetcherId = BrokerIdAndFetcherId(brokerId, fetcherId)
       fetcherThreadMap.get(brokerIdAndFetcherId).foreach { thread =>
         thread.markPartitionsForTruncation(topicPartition, truncationOffset)
@@ -81,24 +117,28 @@ abstract class AbstractFetcherManager(protected val name: String, clientId: Stri
   }
 
   // to be defined in subclass to create a specific fetcher
-  def createFetcherThread(fetcherId: Int, sourceBroker: BrokerEndPoint): AbstractFetcherThread
+  def createFetcherThread(fetcherId: Int, sourceBroker: BrokerEndPoint): T
 
-  def addFetcherForPartitions(partitionAndOffsets: Map[TopicPartition, BrokerAndInitialOffset]) {
-    mapLock synchronized {
-      val partitionsPerFetcher = partitionAndOffsets.groupBy { case(topicPartition, brokerAndInitialFetchOffset) =>
-        BrokerAndFetcherId(brokerAndInitialFetchOffset.broker, getFetcherId(topicPartition.topic, topicPartition.partition))}
+  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) {
+      def addAndStartFetcherThread(brokerAndFetcherId: BrokerAndFetcherId,
+                                   brokerIdAndFetcherId: BrokerIdAndFetcherId): T = {
         val fetcherThread = createFetcherThread(brokerAndFetcherId.fetcherId, brokerAndFetcherId.broker)
         fetcherThreadMap.put(brokerIdAndFetcherId, fetcherThread)
-        fetcherThread.start
+        fetcherThread.start()
+        fetcherThread
       }
 
       for ((brokerAndFetcherId, initialFetchOffsets) <- partitionsPerFetcher) {
         val brokerIdAndFetcherId = BrokerIdAndFetcherId(brokerAndFetcherId.broker.id, brokerAndFetcherId.fetcherId)
-        fetcherThreadMap.get(brokerIdAndFetcherId) match {
-          case Some(f) if f.sourceBroker.host == brokerAndFetcherId.broker.host && f.sourceBroker.port == brokerAndFetcherId.broker.port =>
+        val fetcherThread = fetcherThreadMap.get(brokerIdAndFetcherId) match {
+          case Some(currentFetcherThread) if currentFetcherThread.sourceBroker == brokerAndFetcherId.broker =>
             // reuse the fetcher thread
+            currentFetcherThread
           case Some(f) =>
             f.shutdown()
             addAndStartFetcherThread(brokerAndFetcherId, brokerIdAndFetcherId)
@@ -106,26 +146,31 @@ abstract class AbstractFetcherManager(protected val name: String, clientId: Stri
             addAndStartFetcherThread(brokerAndFetcherId, brokerIdAndFetcherId)
         }
 
-        fetcherThreadMap(brokerIdAndFetcherId).addPartitions(initialFetchOffsets.map { case (tp, brokerAndInitOffset) =>
-          tp -> brokerAndInitOffset.initOffset
-        })
+        addPartitionsToFetcherThread(fetcherThread, initialFetchOffsets)
       }
     }
+  }
 
-    info("Added fetcher for partitions %s".format(partitionAndOffsets.map { case (topicPartition, brokerAndInitialOffset) =>
-      "[" + topicPartition + ", initOffset " + brokerAndInitialOffset.initOffset + " to broker " + brokerAndInitialOffset.broker + "] "}))
+  protected def addPartitionsToFetcherThread(fetcherThread: T,
+                                             initialOffsetAndEpochs: collection.Map[TopicPartition, InitialFetchState]): Unit = {
+    fetcherThread.addPartitions(initialOffsetAndEpochs)
+    info(s"Added fetcher to broker ${fetcherThread.sourceBroker.id} for partitions $initialOffsetAndEpochs")
   }
 
-  def removeFetcherForPartitions(partitions: Set[TopicPartition]) {
-    mapLock synchronized {
+  def removeFetcherForPartitions(partitions: Set[TopicPartition]): Map[TopicPartition, PartitionFetchState] = {
+    val fetchStates = mutable.Map.empty[TopicPartition, PartitionFetchState]
+    lock synchronized {
       for (fetcher <- fetcherThreadMap.values)
-        fetcher.removePartitions(partitions)
+        fetchStates ++= fetcher.removePartitions(partitions)
+      failedPartitions.removeAll(partitions)
     }
-    info("Removed fetcher for partitions %s".format(partitions.mkString(",")))
+    if (partitions.nonEmpty)
+      info(s"Removed fetcher for partitions $partitions")
+    fetchStates
   }
 
-  def shutdownIdleFetcherThreads() {
-    mapLock synchronized {
+  def shutdownIdleFetcherThreads(): Unit = {
+    lock synchronized {
       val keysToBeRemoved = new mutable.HashSet[BrokerIdAndFetcherId]
       for ((key, fetcher) <- fetcherThreadMap) {
         if (fetcher.partitionCount <= 0) {
@@ -137,13 +182,13 @@ abstract class AbstractFetcherManager(protected val name: String, clientId: Stri
     }
   }
 
-  def closeAllFetchers() {
-    mapLock synchronized {
-      for ( (_, fetcher) <- fetcherThreadMap) {
+  def closeAllFetchers(): Unit = {
+    lock synchronized {
+      for ((_, fetcher) <- fetcherThreadMap) {
         fetcher.initiateShutdown()
       }
 
-      for ( (_, fetcher) <- fetcherThreadMap) {
+      for ((_, fetcher) <- fetcherThreadMap) {
         fetcher.shutdown()
       }
       fetcherThreadMap.clear()
@@ -151,8 +196,39 @@ abstract class AbstractFetcherManager(protected val name: String, clientId: 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 BrokerAndInitialOffset(broker: BrokerEndPoint, initOffset: Long)
+case class InitialFetchState(leader: BrokerEndPoint, currentLeaderEpoch: Int, initOffset: Long)
 
 case class BrokerIdAndFetcherId(brokerId: Int, fetcherId: Int)
diff --git a/core/src/main/scala/kafka/server/AbstractFetcherThread.scala b/core/src/main/scala/kafka/server/AbstractFetcherThread.scala
index b078073ba038e..9af23c6473be6 100755
--- a/core/src/main/scala/kafka/server/AbstractFetcherThread.scala
+++ b/core/src/main/scala/kafka/server/AbstractFetcherThread.scala
@@ -17,27 +17,36 @@
 
 package kafka.server
 
+import java.nio.ByteBuffer
+import java.util
+import java.util.Optional
 import java.util.concurrent.locks.ReentrantLock
 
 import kafka.cluster.BrokerEndPoint
 import kafka.utils.{DelayedItem, Pool, ShutdownableThread}
-import org.apache.kafka.common.errors.{CorruptRecordException, KafkaStorageException}
-import kafka.common.{ClientIdAndBroker, KafkaException}
+import kafka.utils.Implicits._
+import org.apache.kafka.common.errors._
+import kafka.common.ClientIdAndBroker
 import kafka.metrics.KafkaMetricsGroup
 import kafka.utils.CoreUtils.inLock
 import org.apache.kafka.common.protocol.Errors
-import AbstractFetcherThread._
 
 import scala.collection.{Map, Set, mutable}
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 import java.util.concurrent.TimeUnit
 import java.util.concurrent.atomic.AtomicLong
 
-import com.yammer.metrics.core.Gauge
-import org.apache.kafka.common.TopicPartition
-import org.apache.kafka.common.internals.{FatalExitError, PartitionStates}
-import org.apache.kafka.common.record.MemoryRecords
-import org.apache.kafka.common.requests.EpochEndOffset
+import kafka.log.LogAppendInfo
+import kafka.server.AbstractFetcherThread.ReplicaFetch
+import kafka.server.AbstractFetcherThread.ResultWithPartitions
+import org.apache.kafka.common.{InvalidRecordException, TopicPartition}
+import org.apache.kafka.common.internals.PartitionStates
+import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData.EpochEndOffset
+import org.apache.kafka.common.record.{FileRecords, MemoryRecords, Records}
+import org.apache.kafka.common.requests._
+import org.apache.kafka.common.requests.OffsetsForLeaderEpochResponse.{UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET}
+
+import scala.math._
 
 /**
  *  Abstract class for fetching data from multiple partitions from the same broker.
@@ -45,17 +54,17 @@ import org.apache.kafka.common.requests.EpochEndOffset
 abstract class AbstractFetcherThread(name: String,
                                      clientId: String,
                                      val sourceBroker: BrokerEndPoint,
+                                     failedPartitions: FailedPartitions,
                                      fetchBackOffMs: Int = 0,
                                      isInterruptible: Boolean = true,
-                                     includeLogTruncation: Boolean
-                                    )
+                                     val brokerTopicStats: BrokerTopicStats) //BrokerTopicStats's lifecycle managed by ReplicaManager
   extends ShutdownableThread(name, isInterruptible) {
 
-  type REQ <: FetchRequest
-  type PD <: PartitionData
+  type FetchData = FetchResponse.PartitionData[Records]
+  type EpochData = OffsetsForLeaderEpochRequest.PartitionData
 
-  private[server] val partitionStates = new PartitionStates[PartitionFetchState]
-  private val partitionMapLock = new ReentrantLock
+  private val partitionStates = new PartitionStates[PartitionFetchState]
+  protected val partitionMapLock = new ReentrantLock
   private val partitionMapCond = partitionMapLock.newCondition()
 
   private val metricId = ClientIdAndBroker(clientId, sourceBroker.host, sourceBroker.port)
@@ -65,25 +74,37 @@ abstract class AbstractFetcherThread(name: String,
   /* callbacks to be defined in subclass */
 
   // process fetched data
-  protected def processPartitionData(topicPartition: TopicPartition, fetchOffset: Long, partitionData: PD)
+  protected def processPartitionData(topicPartition: TopicPartition,
+                                     fetchOffset: Long,
+                                     partitionData: FetchData): Option[LogAppendInfo]
 
-  // handle a partition whose offset is out of range and return a new fetch offset
-  protected def handleOffsetOutOfRange(topicPartition: TopicPartition): Long
+  protected def truncate(topicPartition: TopicPartition, truncationState: OffsetTruncationState): Unit
 
-  // deal with partitions with errors, potentially due to leadership changes
-  protected def handlePartitionsWithErrors(partitions: Iterable[TopicPartition])
+  protected def truncateFullyAndStartAt(topicPartition: TopicPartition, offset: Long): Unit
+
+  protected def buildFetch(partitionMap: Map[TopicPartition, PartitionFetchState]): ResultWithPartitions[Option[ReplicaFetch]]
+
+  protected def latestEpoch(topicPartition: TopicPartition): Option[Int]
+
+  protected def logStartOffset(topicPartition: TopicPartition): Long
 
-  protected def buildLeaderEpochRequest(allPartitions: Seq[(TopicPartition, PartitionFetchState)]): ResultWithPartitions[Map[TopicPartition, Int]]
+  protected def logEndOffset(topicPartition: TopicPartition): Long
 
-  protected def fetchEpochsFromLeader(partitions: Map[TopicPartition, Int]): Map[TopicPartition, EpochEndOffset]
+  protected def endOffsetForEpoch(topicPartition: TopicPartition, epoch: Int): Option[OffsetAndEpoch]
 
-  protected def maybeTruncate(fetchedEpochs: Map[TopicPartition, EpochEndOffset]): ResultWithPartitions[Map[TopicPartition, Long]]
+  protected def fetchEpochEndOffsets(partitions: Map[TopicPartition, EpochData]): Map[TopicPartition, EpochEndOffset]
 
-  protected def buildFetchRequest(partitionMap: Seq[(TopicPartition, PartitionFetchState)]): ResultWithPartitions[REQ]
+  protected def fetchFromLeader(fetchRequest: FetchRequest.Builder): Map[TopicPartition, FetchData]
 
-  protected def fetch(fetchRequest: REQ): Seq[(TopicPartition, PD)]
+  protected def fetchEarliestOffsetFromLeader(topicPartition: TopicPartition, currentLeaderEpoch: Int): Long
 
-  override def shutdown() {
+  protected def fetchLatestOffsetFromLeader(topicPartition: TopicPartition, currentLeaderEpoch: Int): Long
+
+  protected val isOffsetForLeaderEpochSupported: Boolean
+
+  protected val isTruncationOnFetchSupported: Boolean
+
+  override def shutdown(): Unit = {
     initiateShutdown()
     inLock(partitionMapLock) {
       partitionMapCond.signalAll()
@@ -95,65 +116,209 @@ abstract class AbstractFetcherThread(name: String,
     fetcherLagStats.unregister()
   }
 
-  private def states() = partitionStates.partitionStates.asScala.map { state => state.topicPartition -> state.value }
-
-  override def doWork() {
+  override def doWork(): Unit = {
     maybeTruncate()
-    val fetchRequest = inLock(partitionMapLock) {
-      val ResultWithPartitions(fetchRequest, partitionsWithError) = buildFetchRequest(states)
-      if (fetchRequest.isEmpty) {
+    maybeFetch()
+  }
+
+  private def maybeFetch(): Unit = {
+    val fetchRequestOpt = inLock(partitionMapLock) {
+      val ResultWithPartitions(fetchRequestOpt, partitionsWithError) = buildFetch(partitionStates.partitionStateMap.asScala)
+
+      handlePartitionsWithErrors(partitionsWithError, "maybeFetch")
+
+      if (fetchRequestOpt.isEmpty) {
         trace(s"There are no active partitions. Back off for $fetchBackOffMs ms before sending a fetch request")
         partitionMapCond.await(fetchBackOffMs, TimeUnit.MILLISECONDS)
       }
-      handlePartitionsWithErrors(partitionsWithError)
-      fetchRequest
+
+      fetchRequestOpt
+    }
+
+    fetchRequestOpt.foreach { case ReplicaFetch(sessionPartitions, fetchRequest) =>
+      processFetchRequest(sessionPartitions, fetchRequest)
+    }
+  }
+
+  // deal with partitions with errors, potentially due to leadership changes
+  private def handlePartitionsWithErrors(partitions: Iterable[TopicPartition], methodName: String): Unit = {
+    if (partitions.nonEmpty) {
+      debug(s"Handling errors in $methodName for partitions $partitions")
+      delayPartitions(partitions, fetchBackOffMs)
+    }
+  }
+
+  /**
+   * Builds offset for leader epoch requests for partitions that are in the truncating phase based
+   * on latest epochs of the future replicas (the one that is fetching)
+   */
+  private def fetchTruncatingPartitions(): (Map[TopicPartition, EpochData], Set[TopicPartition]) = inLock(partitionMapLock) {
+    val partitionsWithEpochs = mutable.Map.empty[TopicPartition, EpochData]
+    val partitionsWithoutEpochs = mutable.Set.empty[TopicPartition]
+
+    partitionStates.partitionStateMap.forEach { (tp, state) =>
+      if (state.isTruncating) {
+        latestEpoch(tp) match {
+          case Some(epoch) if isOffsetForLeaderEpochSupported =>
+            partitionsWithEpochs += tp -> new EpochData(Optional.of(state.currentLeaderEpoch), epoch)
+          case _ =>
+            partitionsWithoutEpochs += tp
+        }
+      }
+    }
+
+    (partitionsWithEpochs, partitionsWithoutEpochs)
+  }
+
+  private def maybeTruncate(): Unit = {
+    val (partitionsWithEpochs, partitionsWithoutEpochs) = fetchTruncatingPartitions()
+    if (partitionsWithEpochs.nonEmpty) {
+      truncateToEpochEndOffsets(partitionsWithEpochs)
+    }
+    if (partitionsWithoutEpochs.nonEmpty) {
+      truncateToHighWatermark(partitionsWithoutEpochs)
+    }
+  }
+
+  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
     }
-    if (!fetchRequest.isEmpty)
-      processFetchRequest(fetchRequest)
   }
 
   /**
     * - Build a leader epoch fetch based on partitions that are in the Truncating phase
-    * - Issue LeaderEpochRequeust, retrieving the latest offset for each partition's
+    * - 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.
     */
-  def maybeTruncate(): Unit = {
-    val ResultWithPartitions(epochRequests, partitionsWithError) = inLock(partitionMapLock) { buildLeaderEpochRequest(states) }
-    handlePartitionsWithErrors(partitionsWithError)
+  private def truncateToEpochEndOffsets(latestEpochsForPartitions: Map[TopicPartition, EpochData]): Unit = {
+    val endOffsets = fetchEpochEndOffsets(latestEpochsForPartitions)
+    //Ensure we hold a lock during truncation.
+    inLock(partitionMapLock) {
+      //Check no leadership and no leader epoch changes happened whilst we were unlocked, fetching epochs
+      val epochEndOffsets = endOffsets.filter { case (tp, _) =>
+        val curPartitionState = partitionStates.stateValue(tp)
+        val partitionEpochRequest = latestEpochsForPartitions.getOrElse(tp, {
+          throw new IllegalStateException(
+            s"Leader replied with partition $tp not requested in OffsetsForLeaderEpoch request")
+        })
+        val leaderEpochInRequest = partitionEpochRequest.currentLeaderEpoch.get
+        curPartitionState != null && leaderEpochInRequest == curPartitionState.currentLeaderEpoch
+      }
 
-    if (epochRequests.nonEmpty) {
-      val fetchedEpochs = fetchEpochsFromLeader(epochRequests)
-      //Ensure we hold a lock during truncation.
-      inLock(partitionMapLock) {
-        //Check no leadership changes happened whilst we were unlocked, fetching epochs
-        val leaderEpochs = fetchedEpochs.filter { case (tp, _) => partitionStates.contains(tp) }
-        val ResultWithPartitions(fetchOffsets, partitionsWithError) = maybeTruncate(leaderEpochs)
-        handlePartitionsWithErrors(partitionsWithError)
-        markTruncationCompleteAndUpdateFetchOffset(fetchOffsets)
+      val ResultWithPartitions(fetchOffsets, partitionsWithError) = maybeTruncateToEpochEndOffsets(epochEndOffsets, latestEpochsForPartitions)
+      handlePartitionsWithErrors(partitionsWithError, "truncateToEpochEndOffsets")
+      updateFetchOffsetAndMaybeMarkTruncationComplete(fetchOffsets)
+    }
+  }
+
+  protected def truncateOnFetchResponse(epochEndOffsets: Map[TopicPartition, EpochEndOffset]): Unit = {
+    inLock(partitionMapLock) {
+      val ResultWithPartitions(fetchOffsets, partitionsWithError) = maybeTruncateToEpochEndOffsets(epochEndOffsets, Map.empty)
+      handlePartitionsWithErrors(partitionsWithError, "truncateOnFetchResponse")
+      updateFetchOffsetAndMaybeMarkTruncationComplete(fetchOffsets)
+    }
+  }
+
+  // Visible for testing
+  private[server] def truncateToHighWatermark(partitions: Set[TopicPartition]): Unit = inLock(partitionMapLock) {
+    val fetchOffsets = mutable.HashMap.empty[TopicPartition, OffsetTruncationState]
+
+    for (tp <- partitions) {
+      val partitionState = partitionStates.stateValue(tp)
+      if (partitionState != null) {
+        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)
       }
     }
+
+    updateFetchOffsetAndMaybeMarkTruncationComplete(fetchOffsets)
   }
 
-  private def processFetchRequest(fetchRequest: REQ) {
-    val partitionsWithError = mutable.Set[TopicPartition]()
+  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.forKeyValue { (tp, leaderEpochOffset) =>
+      Errors.forCode(leaderEpochOffset.errorCode) match {
+        case Errors.NONE =>
+          val offsetTruncationState = getOffsetTruncationState(tp, leaderEpochOffset)
+          if (doTruncate(tp, offsetTruncationState))
+            fetchOffsets.put(tp, offsetTruncationState)
+
+        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
+      }
+    }
 
-    var responseData: Seq[(TopicPartition, PD)] = Seq.empty
+    ResultWithPartitions(fetchOffsets, partitionsWithError)
+  }
+
+  /**
+   * 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
+      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
+      }
+    }
+  }
+
+  private def processFetchRequest(sessionPartitions: util.Map[TopicPartition, FetchRequest.PartitionData],
+                                  fetchRequest: FetchRequest.Builder): Unit = {
+    val partitionsWithError = mutable.Set[TopicPartition]()
+    val divergingEndOffsets = mutable.Map.empty[TopicPartition, EpochEndOffset]
+    var responseData: Map[TopicPartition, FetchData] = Map.empty
 
     try {
-      trace(s"Issuing fetch to broker ${sourceBroker.id}, request: $fetchRequest")
-      responseData = fetch(fetchRequest)
+      trace(s"Sending fetch request $fetchRequest")
+      responseData = fetchFromLeader(fetchRequest)
     } catch {
       case t: Throwable =>
-        if (isRunning.get) {
-          warn(s"Error in fetch to broker ${sourceBroker.id}, request $fetchRequest", t)
+        if (isRunning) {
+          warn(s"Error in response for fetch request $fetchRequest", t)
           inLock(partitionMapLock) {
             partitionsWithError ++= partitionStates.partitionSet.asScala
             // there is an error occurred while fetching partitions, sleep a while
-            // note that `ReplicaFetcherThread.handlePartitionsWithError` will also introduce the same delay for every
+            // note that `AbstractFetcherThread.handlePartitionsWithError` will also introduce the same delay for every
             // partition with error effectively doubling the delay. It would be good to improve this.
             partitionMapCond.await(fetchBackOffMs, TimeUnit.MILLISECONDS)
           }
@@ -164,172 +329,423 @@ abstract class AbstractFetcherThread(name: String,
     if (responseData.nonEmpty) {
       // process fetched data
       inLock(partitionMapLock) {
-
-        responseData.foreach { case (topicPartition, partitionData) =>
-          val topic = topicPartition.topic
-          val partitionId = topicPartition.partition
-          Option(partitionStates.stateValue(topicPartition)).foreach(currentPartitionFetchState =>
+        responseData.forKeyValue { (topicPartition, partitionData) =>
+          Option(partitionStates.stateValue(topicPartition)).foreach { currentFetchState =>
             // It's possible that a partition is removed and re-added or truncated when there is a pending fetch request.
-            // In this case, we only want to process the fetch response if the partition state is ready for fetch and the current offset is the same as the offset requested.
-            if (fetchRequest.offset(topicPartition) == currentPartitionFetchState.fetchOffset &&
-                currentPartitionFetchState.isReadyForFetch) {
+            // In this case, we only want to process the fetch response if the partition state is ready for fetch and
+            // the current offset is the same as the offset requested.
+            val fetchPartitionData = sessionPartitions.get(topicPartition)
+            if (fetchPartitionData != null && fetchPartitionData.fetchOffset == currentFetchState.fetchOffset && currentFetchState.isReadyForFetch) {
+              val requestEpoch = if (fetchPartitionData.currentLeaderEpoch.isPresent) Some(fetchPartitionData.currentLeaderEpoch.get().toInt) else None
               partitionData.error match {
                 case Errors.NONE =>
                   try {
-                    val records = partitionData.toRecords
-                    val newOffset = records.batches.asScala.lastOption.map(_.nextOffset).getOrElse(
-                      currentPartitionFetchState.fetchOffset)
-
-                    fetcherLagStats.getAndMaybePut(topic, partitionId).lag = Math.max(0L, partitionData.highWatermark - newOffset)
                     // Once we hand off the partition data to the subclass, we can't mess with it any more in this thread
-                    processPartitionData(topicPartition, currentPartitionFetchState.fetchOffset, partitionData)
-
-                    val validBytes = records.validBytes
-                    // ReplicaDirAlterThread may have removed topicPartition from the partitionStates after processing the partition data
-                    if (validBytes > 0 && partitionStates.contains(topicPartition)) {
-                      // Update partitionStates only if there is no exception during processPartitionData
-                      partitionStates.updateAndMoveToEnd(topicPartition, new PartitionFetchState(newOffset))
-                      fetcherStats.byteRate.mark(validBytes)
+                    val logAppendInfoOpt = processPartitionData(topicPartition, currentFetchState.fetchOffset,
+                      partitionData)
+
+                    logAppendInfoOpt.foreach { logAppendInfo =>
+                      val validBytes = logAppendInfo.validBytes
+                      val nextOffset = if (validBytes > 0) logAppendInfo.lastOffset + 1 else currentFetchState.fetchOffset
+                      val lag = Math.max(0L, partitionData.highWatermark - nextOffset)
+                      fetcherLagStats.getAndMaybePut(topicPartition).lag = lag
+
+                      // ReplicaDirAlterThread may have removed topicPartition from the partitionStates after processing the partition data
+                      if (validBytes > 0 && partitionStates.contains(topicPartition)) {
+                        // Update partitionStates only if there is no exception during processPartitionData
+                        val newFetchState = PartitionFetchState(nextOffset, Some(lag),
+                          currentFetchState.currentLeaderEpoch, state = Fetching,
+                          logAppendInfo.lastLeaderEpoch)
+                        partitionStates.updateAndMoveToEnd(topicPartition, newFetchState)
+                        fetcherStats.byteRate.mark(validBytes)
+                      }
+                    }
+                    if (isTruncationOnFetchSupported) {
+                      partitionData.divergingEpoch.ifPresent { divergingEpoch =>
+                        divergingEndOffsets += topicPartition -> new EpochEndOffset()
+                          .setPartition(topicPartition.partition)
+                          .setErrorCode(Errors.NONE.code)
+                          .setLeaderEpoch(divergingEpoch.epoch)
+                          .setEndOffset(divergingEpoch.endOffset)
+                      }
                     }
                   } 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
-                      // 2. If the message is corrupt due to a transient state in the log (truncation, partial writes can cause this), we simply continue and
-                      // should get fixed in the subsequent fetches
-                      error(s"Found invalid messages during fetch for partition $topicPartition offset ${currentPartitionFetchState.fetchOffset}", ime)
+                      // 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
+                      // 2. If the message is corrupt due to a transient state in the log (truncation, partial writes
+                      //    can cause this), we simply continue and should get fixed in the subsequent fetches
+                      error(s"Found invalid messages during fetch for partition $topicPartition " +
+                        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 ${currentPartitionFetchState.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 =>
-                  try {
-                    val newOffset = handleOffsetOutOfRange(topicPartition)
-                    partitionStates.updateAndMoveToEnd(topicPartition, new PartitionFetchState(newOffset))
-                    error(s"Current offset ${currentPartitionFetchState.fetchOffset} for partition $topicPartition out of range; reset offset to $newOffset")
-                  } catch {
-                    case e: FatalExitError => throw e
-                    case e: Throwable =>
-                      error(s"Error getting offset for partition $topicPartition from broker ${sourceBroker.id}", e)
-                      partitionsWithError += topicPartition
-                  }
-                case _ =>
-                  if (isRunning.get) {
-                    error(s"Error for partition $topicPartition from broker ${sourceBroker.id}", partitionData.exception.get)
+                  if (handleOutOfRangeError(topicPartition, currentFetchState, requestEpoch))
                     partitionsWithError += topicPartition
-                  }
+
+                case Errors.UNKNOWN_LEADER_EPOCH =>
+                  debug(s"Remote broker has a smaller leader epoch for partition $topicPartition than " +
+                    s"this replica's current leader epoch of ${currentFetchState.currentLeaderEpoch}.")
+                  partitionsWithError += topicPartition
+
+                case Errors.FENCED_LEADER_EPOCH =>
+                  if (onPartitionFenced(topicPartition, requestEpoch)) partitionsWithError += topicPartition
+
+                case Errors.NOT_LEADER_OR_FOLLOWER =>
+                  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"Received ${Errors.UNKNOWN_TOPIC_OR_PARTITION} from the leader for partition $topicPartition. " +
+                       "This error may be returned transiently when the partition is being created or deleted, but it is not " +
+                       "expected to persist.")
+                  partitionsWithError += topicPartition
+
+                case _ =>
+                  error(s"Error for partition $topicPartition at offset ${currentFetchState.fetchOffset}",
+                    partitionData.error.exception)
+                  partitionsWithError += topicPartition
               }
-            })
+            }
+          }
         }
       }
     }
 
-    if (partitionsWithError.nonEmpty)
-      debug(s"handling partitions with error for $partitionsWithError")
-    handlePartitionsWithErrors(partitionsWithError)
+    if (divergingEndOffsets.nonEmpty)
+      truncateOnFetchResponse(divergingEndOffsets)
+    if (partitionsWithError.nonEmpty) {
+      handlePartitionsWithErrors(partitionsWithError, "processFetchRequest")
+    }
   }
 
-  def markPartitionsForTruncation(topicPartition: TopicPartition, truncationOffset: Long) {
-    if (!includeLogTruncation)
-      throw new IllegalStateException("Truncation should not be requested if includeLogTruncation is disabled")
+  /**
+   * This is used to mark partitions for truncation in ReplicaAlterLogDirsThread after leader
+   * offsets are known.
+   */
+  def markPartitionsForTruncation(topicPartition: TopicPartition, truncationOffset: Long): Unit = {
     partitionMapLock.lockInterruptibly()
     try {
       Option(partitionStates.stateValue(topicPartition)).foreach { state =>
-        val newState = PartitionFetchState(math.min(truncationOffset, state.fetchOffset), state.delay, truncatingLog = true)
+        val newState = PartitionFetchState(math.min(truncationOffset, state.fetchOffset),
+          state.lag, state.currentLeaderEpoch, state.delay, state = Truncating,
+          lastFetchedEpoch = None)
         partitionStates.updateAndMoveToEnd(topicPartition, newState)
+        partitionMapCond.signalAll()
       }
-      partitionMapCond.signalAll()
     } finally partitionMapLock.unlock()
   }
 
-  def addPartitions(initialFetchOffsets: Map[TopicPartition, Long]) {
+  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")
+  }
+
+  /**
+   * Returns initial partition fetch state based on current state and the provided `initialFetchState`.
+   * From IBP 2.7 onwards, we can rely on truncation based on diverging data returned in fetch responses.
+   * For older versions, we can skip the truncation step iff the leader epoch matches the existing epoch.
+   */
+  private def partitionFetchState(tp: TopicPartition, initialFetchState: InitialFetchState, currentState: PartitionFetchState): PartitionFetchState = {
+    if (currentState != null && currentState.currentLeaderEpoch == initialFetchState.currentLeaderEpoch) {
+      currentState
+    } else if (initialFetchState.initOffset < 0) {
+      fetchOffsetAndTruncate(tp, initialFetchState.currentLeaderEpoch)
+    } else if (isTruncationOnFetchSupported) {
+      // With old message format, `latestEpoch` will be empty and we use Truncating state
+      // to truncate to high watermark.
+      val lastFetchedEpoch = latestEpoch(tp)
+      val state = if (lastFetchedEpoch.nonEmpty) Fetching else Truncating
+      PartitionFetchState(initialFetchState.initOffset, None, initialFetchState.currentLeaderEpoch,
+          state, lastFetchedEpoch)
+    } else {
+      PartitionFetchState(initialFetchState.initOffset, None, initialFetchState.currentLeaderEpoch,
+        state = Truncating, lastFetchedEpoch = None)
+    }
+  }
+
+  def addPartitions(initialFetchStates: Map[TopicPartition, InitialFetchState]): Set[TopicPartition] = {
     partitionMapLock.lockInterruptibly()
     try {
-      // If the partitionMap already has the topic/partition, then do not update the map with the old offset
-      val newPartitionToState = initialFetchOffsets.filter { case (tp, _) =>
-        !partitionStates.contains(tp)
-      }.map { case (tp, initialFetchOffset) =>
-        val fetchState =
-          if (initialFetchOffset < 0)
-            new PartitionFetchState(handleOffsetOutOfRange(tp), includeLogTruncation)
-          else
-            new PartitionFetchState(initialFetchOffset, includeLogTruncation)
-        tp -> fetchState
+      failedPartitions.removeAll(initialFetchStates.keySet)
+
+      initialFetchStates.forKeyValue { (tp, initialFetchState) =>
+        val currentState = partitionStates.stateValue(tp)
+        val updatedState = partitionFetchState(tp, initialFetchState, currentState)
+        partitionStates.updateAndMoveToEnd(tp, updatedState)
       }
-      val existingPartitionToState = states().toMap
-      partitionStates.set((existingPartitionToState ++ newPartitionToState).asJava)
+
       partitionMapCond.signalAll()
+      initialFetchStates.keySet
     } finally partitionMapLock.unlock()
   }
 
   /**
-    * Loop through all partitions, marking them as truncation complete and update the fetch offset
+    * Loop through all partitions, updating their fetch offset and maybe marking them as
+    * truncation completed if their offsetTruncationState indicates truncation completed
     *
-    * @param fetchOffsets the partitions to mark truncation complete
+    * @param fetchOffsets the partitions to update fetch offset and maybe mark truncation complete
     */
-  private def markTruncationCompleteAndUpdateFetchOffset(fetchOffsets: Map[TopicPartition, Long]) {
-    val newStates: Map[TopicPartition, PartitionFetchState] = partitionStates.partitionStates.asScala
-      .map { state =>
-        val maybeTruncationComplete = fetchOffsets.get(state.topicPartition()) match {
-          case Some(offset) => PartitionFetchState(offset, state.value.delay, truncatingLog = false)
-          case None => state.value()
+  private def updateFetchOffsetAndMaybeMarkTruncationComplete(fetchOffsets: Map[TopicPartition, OffsetTruncationState]): Unit = {
+    val newStates: Map[TopicPartition, PartitionFetchState] = partitionStates.partitionStateMap.asScala
+      .map { case (topicPartition, currentFetchState) =>
+        val maybeTruncationComplete = fetchOffsets.get(topicPartition) match {
+          case Some(offsetTruncationState) =>
+            val lastFetchedEpoch = latestEpoch(topicPartition)
+            val state = if (isTruncationOnFetchSupported || offsetTruncationState.truncationCompleted)
+              Fetching
+            else
+              Truncating
+            PartitionFetchState(offsetTruncationState.offset, currentFetchState.lag,
+              currentFetchState.currentLeaderEpoch, currentFetchState.delay, state, lastFetchedEpoch)
+          case None => currentFetchState
         }
-        (state.topicPartition(), maybeTruncationComplete)
-      }.toMap
+        (topicPartition, maybeTruncationComplete)
+      }
     partitionStates.set(newStates.asJava)
   }
 
-  def delayPartitions(partitions: Iterable[TopicPartition], delay: Long) {
+  /**
+   * Called from ReplicaFetcherThread and ReplicaAlterLogDirsThread maybeTruncate for each topic
+   * partition. Returns truncation offset and whether this is the final offset to truncate to
+   *
+   * For each topic partition, the offset to truncate to is calculated based on leader's returned
+   * epoch and offset:
+   *  -- If the leader replied with undefined epoch offset, we must use the high watermark. This can
+   *  happen if 1) the leader is still using message format older than KAFKA_0_11_0; 2) the follower
+   *  requested leader epoch < the first leader epoch known to the leader.
+   *  -- If the leader replied with the valid offset but undefined leader epoch, we truncate to
+   *  leader's offset if it is lower than follower's Log End Offset. This may happen if the
+   *  leader is on the inter-broker protocol version < KAFKA_2_0_IV0
+   *  -- If the leader replied with leader epoch not known to the follower, we truncate to the
+   *  end offset of the largest epoch that is smaller than the epoch the leader replied with, and
+   *  send OffsetsForLeaderEpochRequest with that leader epoch. In a more rare case, where the
+   *  follower was not tracking epochs smaller than the epoch the leader replied with, we
+   *  truncate the leader's offset (and do not send any more leader epoch requests).
+   *  -- Otherwise, truncate to min(leader's offset, end offset on the follower for epoch that
+   *  leader replied with, follower's Log End Offset).
+   *
+   * @param tp                    Topic partition
+   * @param leaderEpochOffset     Epoch end offset received from the leader for this topic partition
+   */
+  private def getOffsetTruncationState(tp: TopicPartition,
+                                       leaderEpochOffset: EpochEndOffset): OffsetTruncationState = inLock(partitionMapLock) {
+    if (leaderEpochOffset.endOffset == UNDEFINED_EPOCH_OFFSET) {
+      // truncate to initial offset which is the high watermark for follower replica. For
+      // future replica, it is either high watermark of the future replica or current
+      // replica's truncation offset (when the current replica truncates, it forces future
+      // replica's partition state to 'truncating' and sets initial offset to its truncation offset)
+      warn(s"Based on replica's leader epoch, leader replied with an unknown offset in $tp. " +
+           s"The initial fetch offset ${partitionStates.stateValue(tp).fetchOffset} will be used for truncation.")
+      OffsetTruncationState(partitionStates.stateValue(tp).fetchOffset, truncationCompleted = true)
+    } else if (leaderEpochOffset.leaderEpoch == UNDEFINED_EPOCH) {
+      // either leader or follower or both use inter-broker protocol version < KAFKA_2_0_IV0
+      // (version 0 of OffsetForLeaderEpoch request/response)
+      warn(s"Leader or replica is on protocol version where leader epoch is not considered in the OffsetsForLeaderEpoch response. " +
+           s"The leader's offset ${leaderEpochOffset.endOffset} will be used for truncation in $tp.")
+      OffsetTruncationState(min(leaderEpochOffset.endOffset, logEndOffset(tp)), truncationCompleted = true)
+    } else {
+      val replicaEndOffset = logEndOffset(tp)
+
+      // get (leader epoch, end offset) pair that corresponds to the largest leader epoch
+      // less than or equal to the requested epoch.
+      endOffsetForEpoch(tp, leaderEpochOffset.leaderEpoch) match {
+        case Some(OffsetAndEpoch(followerEndOffset, followerEpoch)) =>
+          if (followerEpoch != leaderEpochOffset.leaderEpoch) {
+            // the follower does not know about the epoch that leader replied with
+            // we truncate to the end offset of the largest epoch that is smaller than the
+            // epoch the leader replied with, and send another offset for leader epoch request
+            val intermediateOffsetToTruncateTo = min(followerEndOffset, replicaEndOffset)
+            info(s"Based on replica's leader epoch, leader replied with epoch ${leaderEpochOffset.leaderEpoch} " +
+              s"unknown to the replica for $tp. " +
+              s"Will truncate to $intermediateOffsetToTruncateTo and send another leader epoch request to the leader.")
+            OffsetTruncationState(intermediateOffsetToTruncateTo, truncationCompleted = false)
+          } else {
+            val offsetToTruncateTo = min(followerEndOffset, leaderEpochOffset.endOffset)
+            OffsetTruncationState(min(offsetToTruncateTo, replicaEndOffset), truncationCompleted = true)
+          }
+        case None =>
+          // This can happen if the follower was not tracking leader epochs at that point (before the
+          // upgrade, or if this broker is new). Since the leader replied with epoch <
+          // requested epoch from follower, so should be safe to truncate to leader's
+          // offset (this is the same behavior as post-KIP-101 and pre-KIP-279)
+          warn(s"Based on replica's leader epoch, leader replied with epoch ${leaderEpochOffset.leaderEpoch} " +
+            s"below any replica's tracked epochs for $tp. " +
+            s"The leader's offset only ${leaderEpochOffset.endOffset} will be used for truncation.")
+          OffsetTruncationState(min(leaderEpochOffset.endOffset, replicaEndOffset), truncationCompleted = true)
+      }
+    }
+  }
+
+  /**
+   * 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,
+                                    requestEpoch: Option[Int]): Boolean = {
+    try {
+      val newFetchState = fetchOffsetAndTruncate(topicPartition, fetchState.currentLeaderEpoch)
+      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 ${newFetchState.fetchOffset}")
+      false
+    } catch {
+      case _: FencedLeaderEpochException =>
+        onPartitionFenced(topicPartition, requestEpoch)
+
+      case e @ (_ : UnknownTopicOrPartitionException |
+                _ : UnknownLeaderEpochException |
+                _ : NotLeaderOrFollowerException) =>
+        info(s"Could not fetch offset for $topicPartition due to error: ${e.getMessage}")
+        true
+
+      case e: Throwable =>
+        error(s"Error getting offset for partition $topicPartition", e)
+        true
+    }
+  }
+
+  /**
+   * Handle a partition whose offset is out of range and return a new fetch offset.
+   */
+  protected def fetchOffsetAndTruncate(topicPartition: TopicPartition, currentLeaderEpoch: Int): PartitionFetchState = {
+    val replicaEndOffset = logEndOffset(topicPartition)
+
+    /**
+     * Unclean leader election: A follower goes down, in the meanwhile the leader keeps appending messages. The follower comes back up
+     * and before it has completely caught up with the leader's logs, all replicas in the ISR go down. The follower is now uncleanly
+     * elected as the new leader, and it starts appending messages from the client. The old leader comes back up, becomes a follower
+     * and it may discover that the current leader's end offset is behind its own end offset.
+     *
+     * In such a case, truncate the current follower's log to the current leader's end offset and continue fetching.
+     *
+     * There is a potential for a mismatch between the logs of the two replicas here. We don't fix this mismatch as of now.
+     */
+    val leaderEndOffset = fetchLatestOffsetFromLeader(topicPartition, currentLeaderEpoch)
+    if (leaderEndOffset < replicaEndOffset) {
+      warn(s"Reset fetch offset for partition $topicPartition from $replicaEndOffset to current " +
+        s"leader's latest offset $leaderEndOffset")
+      truncate(topicPartition, OffsetTruncationState(leaderEndOffset, truncationCompleted = true))
+
+      fetcherLagStats.getAndMaybePut(topicPartition).lag = 0
+      PartitionFetchState(leaderEndOffset, Some(0), currentLeaderEpoch,
+        state = Fetching, lastFetchedEpoch = latestEpoch(topicPartition))
+    } else {
+      /**
+       * If the leader's log end offset is greater than the follower's log end offset, there are two possibilities:
+       * 1. The follower could have been down for a long time and when it starts up, its end offset could be smaller than the leader's
+       * start offset because the leader has deleted old logs (log.logEndOffset < leaderStartOffset).
+       * 2. When unclean leader election occurs, it is possible that the old leader's high watermark is greater than
+       * the new leader's log end offset. So when the old leader truncates its offset to its high watermark and starts
+       * to fetch from the new leader, an OffsetOutOfRangeException will be thrown. After that some more messages are
+       * produced to the new leader. While the old leader is trying to handle the OffsetOutOfRangeException and query
+       * the log end offset of the new leader, the new leader's log end offset becomes higher than the follower's log end offset.
+       *
+       * In the first case, the follower's current log end offset is smaller than the leader's log start offset. So the
+       * follower should truncate all its logs, roll out a new segment and start to fetch from the current leader's log
+       * start offset.
+       * In the second case, the follower should just keep the current log segments and retry the fetch. In the second
+       * case, there will be some inconsistency of data between old and new leader. We are not solving it here.
+       * If users want to have strong consistency guarantees, appropriate configurations needs to be set for both
+       * brokers and producers.
+       *
+       * Putting the two cases together, the follower should fetch from the higher one of its replica log end offset
+       * and the current leader's log start offset.
+       */
+      val leaderStartOffset = fetchEarliestOffsetFromLeader(topicPartition, currentLeaderEpoch)
+      warn(s"Reset fetch offset for partition $topicPartition from $replicaEndOffset to current " +
+        s"leader's start offset $leaderStartOffset")
+      val offsetToFetch = Math.max(leaderStartOffset, replicaEndOffset)
+      // Only truncate log when current leader's log start offset is greater than follower's log end offset.
+      if (leaderStartOffset > replicaEndOffset)
+        truncateFullyAndStartAt(topicPartition, leaderStartOffset)
+
+      val initialLag = leaderEndOffset - offsetToFetch
+      fetcherLagStats.getAndMaybePut(topicPartition).lag = initialLag
+      PartitionFetchState(offsetToFetch, Some(initialLag), currentLeaderEpoch,
+        state = Fetching, lastFetchedEpoch = latestEpoch(topicPartition))
+    }
+  }
+
+  def delayPartitions(partitions: Iterable[TopicPartition], delay: Long): Unit = {
     partitionMapLock.lockInterruptibly()
     try {
       for (partition <- partitions) {
-        Option(partitionStates.stateValue(partition)).foreach (currentPartitionFetchState =>
-          if (!currentPartitionFetchState.isDelayed)
-            partitionStates.updateAndMoveToEnd(partition, PartitionFetchState(currentPartitionFetchState.fetchOffset, new DelayedItem(delay), currentPartitionFetchState.truncatingLog))
-        )
+        Option(partitionStates.stateValue(partition)).foreach { currentFetchState =>
+          if (!currentFetchState.isDelayed) {
+            partitionStates.updateAndMoveToEnd(partition, PartitionFetchState(currentFetchState.fetchOffset,
+              currentFetchState.lag, currentFetchState.currentLeaderEpoch, Some(new DelayedItem(delay)),
+              currentFetchState.state, currentFetchState.lastFetchedEpoch))
+          }
+        }
       }
       partitionMapCond.signalAll()
     } finally partitionMapLock.unlock()
   }
 
-  def removePartitions(topicPartitions: Set[TopicPartition]) {
+  def removePartitions(topicPartitions: Set[TopicPartition]): Map[TopicPartition, PartitionFetchState] = {
     partitionMapLock.lockInterruptibly()
     try {
-      topicPartitions.foreach { topicPartition =>
+      topicPartitions.map { topicPartition =>
+        val state = partitionStates.stateValue(topicPartition)
         partitionStates.remove(topicPartition)
-        fetcherLagStats.unregister(topicPartition.topic, topicPartition.partition)
-      }
+        fetcherLagStats.unregister(topicPartition)
+        topicPartition -> state
+      }.filter(_._2 != null).toMap
     } finally partitionMapLock.unlock()
   }
 
-  def partitionCount() = {
+  def partitionCount: Int = {
     partitionMapLock.lockInterruptibly()
     try partitionStates.size
     finally partitionMapLock.unlock()
   }
 
+  def partitions: Set[TopicPartition] = {
+    partitionMapLock.lockInterruptibly()
+    try partitionStates.partitionSet.asScala.toSet
+    finally partitionMapLock.unlock()
+  }
+
+  // Visible for testing
+  private[server] def fetchState(topicPartition: TopicPartition): Option[PartitionFetchState] = inLock(partitionMapLock) {
+    Option(partitionStates.stateValue(topicPartition))
+  }
+
+  protected def toMemoryRecords(records: Records): MemoryRecords = {
+    (records: @unchecked) match {
+      case r: MemoryRecords => r
+      case r: FileRecords =>
+        val buffer = ByteBuffer.allocate(r.sizeInBytes)
+        r.readInto(buffer, 0)
+        MemoryRecords.readableRecords(buffer)
+    }
+  }
 }
 
 object AbstractFetcherThread {
 
+  case class ReplicaFetch(partitionData: util.Map[TopicPartition, FetchRequest.PartitionData], fetchRequest: FetchRequest.Builder)
   case class ResultWithPartitions[R](result: R, partitionsWithError: Set[TopicPartition])
 
-  trait FetchRequest {
-    def isEmpty: Boolean
-    def offset(topicPartition: TopicPartition): Long
-  }
-
-  trait PartitionData {
-    def error: Errors
-    def exception: Option[Throwable]
-    def toRecords: MemoryRecords
-    def highWatermark: Long
-  }
-
 }
 
 object FetcherMetrics {
@@ -343,51 +759,38 @@ class FetcherLagMetrics(metricId: ClientIdTopicPartition) extends KafkaMetricsGr
   private[this] val lagVal = new AtomicLong(-1L)
   private[this] val tags = Map(
     "clientId" -> metricId.clientId,
-    "topic" -> metricId.topic,
-    "partition" -> metricId.partitionId.toString)
+    "topic" -> metricId.topicPartition.topic,
+    "partition" -> metricId.topicPartition.partition.toString)
 
-  newGauge(FetcherMetrics.ConsumerLag,
-    new Gauge[Long] {
-      def value = lagVal.get
-    },
-    tags
-  )
+  newGauge(FetcherMetrics.ConsumerLag, () => lagVal.get, 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)
   }
 }
 
 class FetcherLagStats(metricId: ClientIdAndBroker) {
-  private val valueFactory = (k: ClientIdTopicPartition) => new FetcherLagMetrics(k)
-  val stats = new Pool[ClientIdTopicPartition, FetcherLagMetrics](Some(valueFactory))
+  private val valueFactory = (k: TopicPartition) => new FetcherLagMetrics(ClientIdTopicPartition(metricId.clientId, k))
+  val stats = new Pool[TopicPartition, FetcherLagMetrics](Some(valueFactory))
 
-  def getAndMaybePut(topic: String, partitionId: Int): FetcherLagMetrics = {
-    stats.getAndMaybePut(ClientIdTopicPartition(metricId.clientId, topic, partitionId))
-  }
-
-  def isReplicaInSync(topic: String, partitionId: Int): Boolean = {
-    val fetcherLagMetrics = stats.get(ClientIdTopicPartition(metricId.clientId, topic, partitionId))
-    if (fetcherLagMetrics != null)
-      fetcherLagMetrics.lag <= 0
-    else
-      false
+  def getAndMaybePut(topicPartition: TopicPartition): FetcherLagMetrics = {
+    stats.getAndMaybePut(topicPartition)
   }
 
-  def unregister(topic: String, partitionId: Int) {
-    val lagMetrics = stats.remove(ClientIdTopicPartition(metricId.clientId, topic, partitionId))
+  def unregister(topicPartition: TopicPartition): Unit = {
+    val lagMetrics = stats.remove(topicPartition)
     if (lagMetrics != null) lagMetrics.unregister()
   }
 
-  def unregister() {
-    stats.keys.toBuffer.foreach { key: ClientIdTopicPartition =>
-      unregister(key.topic, key.partitionId)
+  def unregister(): Unit = {
+    stats.keys.toBuffer.foreach { key: TopicPartition =>
+      unregister(key)
     }
   }
 }
@@ -401,37 +804,71 @@ class FetcherStats(metricId: ClientIdAndBroker) extends KafkaMetricsGroup {
 
   val byteRate = newMeter(FetcherMetrics.BytesPerSec, "bytes", TimeUnit.SECONDS, tags)
 
-  def unregister() {
+  def unregister(): Unit = {
     removeMetric(FetcherMetrics.RequestsPerSec, tags)
     removeMetric(FetcherMetrics.BytesPerSec, tags)
   }
 
 }
 
-case class ClientIdTopicPartition(clientId: String, topic: String, partitionId: Int) {
-  override def toString = "%s-%s-%d".format(clientId, topic, partitionId)
+case class ClientIdTopicPartition(clientId: String, topicPartition: TopicPartition) {
+  override def toString: String = s"$clientId-$topicPartition"
+}
+
+sealed trait ReplicaState
+case object Truncating extends ReplicaState
+case object Fetching extends ReplicaState
+
+object PartitionFetchState {
+  def apply(offset: Long, lag: Option[Long], currentLeaderEpoch: Int, state: ReplicaState,
+            lastFetchedEpoch: Option[Int]): PartitionFetchState = {
+    PartitionFetchState(offset, lag, currentLeaderEpoch, None, state, lastFetchedEpoch)
+  }
 }
 
+
 /**
-  * case class to keep partition offset and its state(truncatingLog, delayed)
-  * This represents a partition as being either:
-  * (1) Truncating its log, for example having recently become a follower
-  * (2) Delayed, for example due to an error, where we subsequently back off a bit
-  * (3) ReadyForFetch, the is the active state where the thread is actively fetching data.
-  */
-case class PartitionFetchState(fetchOffset: Long, delay: DelayedItem, truncatingLog: Boolean = false) {
+ * case class to keep partition offset and its state(truncatingLog, delayed)
+ * This represents a partition as being either:
+ * (1) Truncating its log, for example having recently become a follower
+ * (2) Delayed, for example due to an error, where we subsequently back off a bit
+ * (3) ReadyForFetch, the is the active state where the thread is actively fetching data.
+ */
+case class PartitionFetchState(fetchOffset: Long,
+                               lag: Option[Long],
+                               currentLeaderEpoch: Int,
+                               delay: Option[DelayedItem],
+                               state: ReplicaState,
+                               lastFetchedEpoch: Option[Int]) {
+
+  def isReadyForFetch: Boolean = state == Fetching && !isDelayed
+
+  def isReplicaInSync: Boolean = lag.isDefined && lag.get <= 0
 
-  def this(offset: Long, truncatingLog: Boolean) = this(offset, new DelayedItem(0), truncatingLog)
+  def isTruncating: Boolean = state == Truncating && !isDelayed
 
-  def this(offset: Long, delay: DelayedItem) = this(offset, new DelayedItem(0), false)
+  def isDelayed: Boolean = delay.exists(_.getDelay(TimeUnit.MILLISECONDS) > 0)
 
-  def this(fetchOffset: Long) = this(fetchOffset, new DelayedItem(0))
+  override def toString: String = {
+    s"FetchState(fetchOffset=$fetchOffset" +
+      s", currentLeaderEpoch=$currentLeaderEpoch" +
+      s", lastFetchedEpoch=$lastFetchedEpoch" +
+      s", state=$state" +
+      s", lag=$lag" +
+      s", delay=${delay.map(_.delayMs).getOrElse(0)}ms" +
+      s")"
+  }
+}
 
-  def isReadyForFetch: Boolean = delay.getDelay(TimeUnit.MILLISECONDS) == 0 && !truncatingLog
+case class OffsetTruncationState(offset: Long, truncationCompleted: Boolean) {
 
-  def isTruncatingLog: Boolean = delay.getDelay(TimeUnit.MILLISECONDS) == 0 && truncatingLog
+  def this(offset: Long) = this(offset, true)
 
-  def isDelayed: Boolean = delay.getDelay(TimeUnit.MILLISECONDS) > 0
+  override def toString: String = "offset:%d-truncationCompleted:%b".format(offset, truncationCompleted)
+}
 
-  override def toString = "offset:%d-isReadyForFetch:%b-isTruncatingLog:%b".format(fetchOffset, isReadyForFetch, truncatingLog)
+case class OffsetAndEpoch(offset: Long, leaderEpoch: Int) {
+  override def toString: String = {
+    s"(offset=$offset, leaderEpoch=$leaderEpoch)"
+  }
 }
diff --git a/core/src/main/scala/kafka/server/ActionQueue.scala b/core/src/main/scala/kafka/server/ActionQueue.scala
new file mode 100644
index 0000000000000..1b6b8326fa4cf
--- /dev/null
+++ b/core/src/main/scala/kafka/server/ActionQueue.scala
@@ -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 kafka.server
+
+import java.util.concurrent.ConcurrentLinkedQueue
+
+import kafka.utils.Logging
+
+/**
+ * This queue is used to collect actions which need to be executed later. One use case is that ReplicaManager#appendRecords
+ * produces record changes so we need to check and complete delayed requests. In order to avoid conflicting locking,
+ * we add those actions to this queue and then complete them at the end of KafkaApis.handle() or DelayedJoin.onExpiration.
+ */
+class ActionQueue extends Logging {
+  private val queue = new ConcurrentLinkedQueue[() => Unit]()
+
+  /**
+   * add action to this queue.
+   * @param action action
+   */
+  def add(action: () => Unit): Unit = queue.add(action)
+
+  /**
+   * try to complete all delayed actions
+   */
+  def tryCompleteActions(): Unit = {
+    val maxToComplete = queue.size()
+    var count = 0
+    var done = false
+    while (!done && count < maxToComplete) {
+      try {
+        val action = queue.poll()
+        if (action == null) done = true
+        else action()
+      } catch {
+        case e: Throwable =>
+          error("failed to complete delayed actions", e)
+      } finally count += 1
+    }
+  }
+}
diff --git a/core/src/main/scala/kafka/server/AdminManager.scala b/core/src/main/scala/kafka/server/AdminManager.scala
index 8f69000125f3e..f1bd1e276c4f1 100644
--- a/core/src/main/scala/kafka/server/AdminManager.scala
+++ b/core/src/main/scala/kafka/server/AdminManager.scala
@@ -16,28 +16,46 @@
   */
 package kafka.server
 
+import java.util
 import java.util.{Collections, Properties}
 
 import kafka.admin.{AdminOperationException, AdminUtils}
 import kafka.common.TopicAlreadyMarkedForDeletionException
 import kafka.log.LogConfig
+import kafka.utils.Log4jController
 import kafka.metrics.KafkaMetricsGroup
+import kafka.server.DynamicConfig.QuotaConfigs
 import kafka.utils._
+import kafka.utils.Implicits._
 import kafka.zk.{AdminZkClient, KafkaZkClient}
-import org.apache.kafka.clients.admin.NewPartitions
-import org.apache.kafka.common.config.{AbstractConfig, ConfigDef, ConfigException, ConfigResource}
-import org.apache.kafka.common.errors.{ApiException, InvalidPartitionsException, InvalidReplicaAssignmentException, InvalidRequestException, PolicyViolationException, ReassignmentInProgressException, UnknownTopicOrPartitionException}
+import org.apache.kafka.clients.admin.{AlterConfigOp, ScramMechanism}
+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.ThrottlingQuotaExceededException
+import org.apache.kafka.common.errors.{ApiException, InvalidConfigurationException, InvalidPartitionsException, InvalidReplicaAssignmentException, InvalidRequestException, ReassignmentInProgressException, TopicExistsException, UnknownTopicOrPartitionException, UnsupportedVersionException}
 import org.apache.kafka.common.internals.Topic
+import org.apache.kafka.common.message.AlterUserScramCredentialsResponseData.AlterUserScramCredentialsResult
+import org.apache.kafka.common.message.CreatePartitionsRequestData.CreatePartitionsTopic
+import org.apache.kafka.common.message.CreateTopicsRequestData.CreatableTopic
+import org.apache.kafka.common.message.CreateTopicsResponseData.{CreatableTopicConfigs, CreatableTopicResult}
+import org.apache.kafka.common.message.{AlterUserScramCredentialsRequestData, AlterUserScramCredentialsResponseData, DescribeConfigsResponseData, DescribeUserScramCredentialsResponseData}
+import org.apache.kafka.common.message.DescribeConfigsRequestData.DescribeConfigsResource
+import org.apache.kafka.common.message.DescribeUserScramCredentialsResponseData.CredentialInfo
 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.requests.CreateTopicsRequest._
-import org.apache.kafka.common.requests.{AlterConfigsRequest, ApiError, DescribeConfigsResponse, Resource, ResourceType}
+import org.apache.kafka.common.security.scram.internals.{ScramMechanism => InternalScramMechanism}
 import org.apache.kafka.server.policy.{AlterConfigPolicy, CreateTopicPolicy}
 import org.apache.kafka.server.policy.CreateTopicPolicy.RequestMetadata
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.quota.{ClientQuotaAlteration, ClientQuotaEntity, ClientQuotaFilter, ClientQuotaFilterComponent}
+import org.apache.kafka.common.requests.CreateTopicsRequest._
+import org.apache.kafka.common.requests.DescribeConfigsResponse.ConfigSource
+import org.apache.kafka.common.requests.{AlterConfigsRequest, ApiError, DescribeConfigsResponse}
+import org.apache.kafka.common.security.scram.internals.{ScramCredentialUtils, ScramFormatter}
+import org.apache.kafka.common.utils.Sanitizer
 
-import scala.collection._
-import scala.collection.JavaConverters._
+import scala.collection.{Map, mutable, _}
+import scala.jdk.CollectionConverters._
 
 class AdminManager(val config: KafkaConfig,
                    val metrics: Metrics,
@@ -55,92 +73,154 @@ 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.")
   }
 
+  private def validateTopicCreatePolicy(topic: CreatableTopic,
+                                        resolvedNumPartitions: Int,
+                                        resolvedReplicationFactor: Short,
+                                        assignments: Map[Int, Seq[Int]]): Unit = {
+    createTopicPolicy.foreach { policy =>
+      // Use `null` for unset fields in the public API
+      val numPartitions: java.lang.Integer =
+        if (topic.assignments().isEmpty) resolvedNumPartitions else null
+      val replicationFactor: java.lang.Short =
+        if (topic.assignments().isEmpty) resolvedReplicationFactor else null
+      val javaAssignments = if (topic.assignments().isEmpty) {
+        null
+      } else {
+        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.forEach(config => javaConfigs.put(config.name, config.value))
+      policy.validate(new RequestMetadata(topic.name, numPartitions, replicationFactor,
+        javaAssignments, javaConfigs))
+    }
+  }
+
+  private def maybePopulateMetadataAndConfigs(metadataAndConfigs: Map[String, CreatableTopicResult],
+                                              topicName: String,
+                                              configs: Properties,
+                                              assignments: Map[Int, Seq[Int]]): Unit = {
+    metadataAndConfigs.get(topicName).foreach { result =>
+      val logConfig = LogConfig.fromProps(KafkaServer.copyKafkaConfigToLog(config), configs)
+      val createEntry = createTopicConfigEntry(logConfig, configs, includeSynonyms = false, includeDocumentation = false)(_, _)
+      val topicConfigs = logConfig.values.asScala.map { case (k, v) =>
+        val entry = createEntry(k, v)
+        new CreatableTopicConfigs()
+          .setName(k)
+          .setValue(entry.value)
+          .setIsSensitive(entry.isSensitive)
+          .setReadOnly(entry.readOnly)
+          .setConfigSource(entry.configSource)
+      }.toList.asJava
+      result.setConfigs(topicConfigs)
+      result.setNumPartitions(assignments.size)
+      result.setReplicationFactor(assignments(0).size.toShort)
+    }
+  }
+
   /**
     * Create topics and wait until the topics have been completely created.
     * The callback function will be triggered either when timeout, error or the topics are created.
     */
   def createTopics(timeout: Int,
                    validateOnly: Boolean,
-                   createInfo: Map[String, TopicDetails],
-                   responseCallback: Map[String, ApiError] => Unit) {
+                   toCreate: Map[String, CreatableTopic],
+                   includeConfigsAndMetadata: Map[String, CreatableTopicResult],
+                   controllerMutationQuota: ControllerMutationQuota,
+                   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 = createInfo.map { case (topic, arguments) =>
+    val metadata = toCreate.values.map(topic =>
       try {
-        val configs = new Properties()
-        arguments.configs.asScala.foreach { case (key, value) =>
-          configs.setProperty(key, value)
-        }
-        LogConfig.validate(configs)
-
-        val assignments = {
-          if ((arguments.numPartitions != NO_NUM_PARTITIONS || arguments.replicationFactor != NO_REPLICATION_FACTOR)
-            && !arguments.replicasAssignments.isEmpty)
-            throw new InvalidRequestException("Both numPartitions or replicationFactor and replicasAssignments were set. " +
-              "Both cannot be used at the same time.")
-          else if (!arguments.replicasAssignments.isEmpty) {
-            // Note: we don't check that replicaAssignment contains unknown brokers - unlike in add-partitions case,
-            // this follows the existing logic in TopicCommand
-            arguments.replicasAssignments.asScala.map { case (partitionId, replicas) =>
-              (partitionId.intValue, replicas.asScala.map(_.intValue))
-            }
-          } else
-            AdminUtils.assignReplicasToBrokers(brokers, arguments.numPartitions, arguments.replicationFactor)
+        if (metadataCache.contains(topic.name))
+          throw new TopicExistsException(s"Topic '${topic.name}' already exists.")
+
+        val nullConfigs = topic.configs.asScala.filter(_.value == null).map(_.name)
+        if (nullConfigs.nonEmpty)
+          throw new InvalidRequestException(s"Null value not supported for topic configs : ${nullConfigs.mkString(",")}")
+
+        if ((topic.numPartitions != NO_NUM_PARTITIONS || topic.replicationFactor != NO_REPLICATION_FACTOR)
+            && !topic.assignments().isEmpty) {
+          throw new InvalidRequestException("Both numPartitions or replicationFactor and replicasAssignments were set. " +
+            "Both cannot be used at the same time.")
         }
-        trace(s"Assignments for topic $topic are $assignments ")
 
-        createTopicPolicy match {
-          case Some(policy) =>
-            adminZkClient.validateCreateOrUpdateTopic(topic, assignments, configs, update = false)
+        val resolvedNumPartitions = if (topic.numPartitions == NO_NUM_PARTITIONS)
+          defaultNumPartitions else topic.numPartitions
+        val resolvedReplicationFactor = if (topic.replicationFactor == NO_REPLICATION_FACTOR)
+          defaultReplicationFactor else topic.replicationFactor
 
-            // Use `null` for unset fields in the public API
-            val numPartitions: java.lang.Integer =
-              if (arguments.numPartitions == NO_NUM_PARTITIONS) null else arguments.numPartitions
-            val replicationFactor: java.lang.Short =
-              if (arguments.replicationFactor == NO_REPLICATION_FACTOR) null else arguments.replicationFactor
-            val replicaAssignments = if (arguments.replicasAssignments.isEmpty) null else arguments.replicasAssignments
+        val assignments = if (topic.assignments.isEmpty) {
+          AdminUtils.assignReplicasToBrokers(
+            brokers, 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,
+          // this follows the existing logic in TopicCommand
+          topic.assignments.forEach { assignment =>
+            assignments(assignment.partitionIndex) = assignment.brokerIds.asScala.map(a => a: Int)
+          }
+          assignments
+        }
+        trace(s"Assignments for topic $topic are $assignments ")
 
-            policy.validate(new RequestMetadata(topic, numPartitions, replicationFactor, replicaAssignments,
-              arguments.configs))
+        val configs = new Properties()
+        topic.configs.forEach(entry => configs.setProperty(entry.name, entry.value))
+        adminZkClient.validateTopicCreate(topic.name, assignments, configs)
+        validateTopicCreatePolicy(topic, resolvedNumPartitions, resolvedReplicationFactor, assignments)
 
-            if (!validateOnly)
-              adminZkClient.createOrUpdateTopicPartitionAssignmentPathInZK(topic, assignments, configs, update = false)
+        // For responses with DescribeConfigs permission, populate metadata and configs. It is
+        // safe to populate it before creating the topic because the values are unset if the
+        // creation fails.
+        maybePopulateMetadataAndConfigs(includeConfigsAndMetadata, topic.name, configs, assignments)
 
-          case None =>
-            if (validateOnly)
-              adminZkClient.validateCreateOrUpdateTopic(topic, assignments, configs, update = false)
-            else
-              adminZkClient.createOrUpdateTopicPartitionAssignmentPathInZK(topic, assignments, configs, update = false)
+        if (validateOnly) {
+          CreatePartitionsMetadata(topic.name, assignments.keySet)
+        } else {
+          controllerMutationQuota.record(assignments.size)
+          adminZkClient.createTopicWithAssignment(topic.name, configs, assignments, validate = false)
+          CreatePartitionsMetadata(topic.name, assignments.keySet)
         }
-        CreatePartitionsMetadata(topic, assignments, ApiError.NONE)
       } catch {
         // Log client errors at a lower level than unexpected exceptions
-        case e@ (_: PolicyViolationException | _: ApiException) =>
-          info(s"Error processing create topic request for topic $topic with arguments $arguments", e)
-          CreatePartitionsMetadata(topic, Map(), ApiError.fromThrowable(e))
+        case e: TopicExistsException =>
+          debug(s"Topic creation failed since topic '${topic.name}' already exists.", e)
+          CreatePartitionsMetadata(topic.name, e)
+        case e: ThrottlingQuotaExceededException =>
+          debug(s"Topic creation not allowed because quota is violated. Delay time: ${e.throttleTimeMs}")
+          CreatePartitionsMetadata(topic.name, e)
+        case e: ApiException =>
+          info(s"Error processing create topic request $topic", e)
+          CreatePartitionsMetadata(topic.name, e)
+        case e: ConfigException =>
+          info(s"Error processing create topic request $topic", e)
+          CreatePartitionsMetadata(topic.name, new InvalidConfigurationException(e.getMessage, e.getCause))
         case e: Throwable =>
-          error(s"Error processing create topic request for topic $topic with arguments $arguments", e)
-          CreatePartitionsMetadata(topic, Map(), ApiError.fromThrowable(e))
-      }
-    }
+          error(s"Error processing create topic request $topic", e)
+          CreatePartitionsMetadata(topic.name, e)
+      }).toBuffer
 
     // 2. if timeout <= 0, validateOnly or no topics can proceed return immediately
     if (timeout <= 0 || validateOnly || !metadata.exists(_.error.is(Errors.NONE))) {
       val results = metadata.map { createTopicMetadata =>
         // ignore topics that already have errors
-        if (createTopicMetadata.error.isSuccess() && !validateOnly) {
+        if (createTopicMetadata.error.isSuccess && !validateOnly) {
           (createTopicMetadata.topic, new ApiError(Errors.REQUEST_TIMED_OUT, null))
         } else {
           (createTopicMetadata.topic, createTopicMetadata.error)
@@ -149,8 +229,9 @@ 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 = createInfo.keys.map(new TopicKey(_)).toSeq
+      val delayedCreate = new DelayedCreatePartitions(timeout, metadata, this,
+        responseCallback)
+      val delayedCreateKeys = toCreate.values.map(topic => TopicKey(topic.name)).toBuffer
       // try to complete the request immediately, otherwise put it into the purgatory
       topicPurgatory.tryCompleteElseWatch(delayedCreate, delayedCreateKeys)
     }
@@ -162,20 +243,24 @@ class AdminManager(val config: KafkaConfig,
     */
   def deleteTopics(timeout: Int,
                    topics: Set[String],
-                   responseCallback: Map[String, Errors] => Unit) {
-
+                   controllerMutationQuota: ControllerMutationQuota,
+                   responseCallback: Map[String, Errors] => Unit): Unit = {
     // 1. map over topics calling the asynchronous delete
     val metadata = topics.map { topic =>
         try {
+          controllerMutationQuota.record(metadataCache.numPartitions(topic).getOrElse(0).toDouble)
           adminZkClient.deleteTopic(topic)
           DeleteTopicMetadata(topic, Errors.NONE)
         } catch {
           case _: TopicAlreadyMarkedForDeletionException =>
             // swallow the exception, and still track deletion allowing multiple calls to wait for deletion
             DeleteTopicMetadata(topic, Errors.NONE)
+          case e: ThrottlingQuotaExceededException =>
+            debug(s"Topic deletion not allowed because quota is violated. Delay time: ${e.throttleTimeMs}")
+            DeleteTopicMetadata(topic, e)
           case e: Throwable =>
             error(s"Error processing delete topic request for topic $topic", e)
-            DeleteTopicMetadata(topic, Errors.forException(e))
+            DeleteTopicMetadata(topic, e)
         }
     }
 
@@ -193,38 +278,39 @@ class AdminManager(val config: KafkaConfig,
     } else {
       // 3. else pass the topics and errors to the delayed operation and set the keys
       val delayedDelete = new DelayedDeleteTopics(timeout, metadata.toSeq, this, responseCallback)
-      val delayedDeleteKeys = topics.map(new TopicKey(_)).toSeq
+      val delayedDeleteKeys = topics.map(TopicKey).toSeq
       // try to complete the request immediately, otherwise put it into the purgatory
       topicPurgatory.tryCompleteElseWatch(delayedDelete, delayedDeleteKeys)
     }
   }
 
   def createPartitions(timeout: Int,
-                       newPartitions: Map[String, NewPartitions],
+                       newPartitions: Seq[CreatePartitionsTopic],
                        validateOnly: Boolean,
-                       listenerName: ListenerName,
+                       controllerMutationQuota: ControllerMutationQuota,
                        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 metadata = newPartitions.map { newPartition =>
+      val topic = newPartition.name
 
-        val existingAssignment = zkClient.getReplicaAssignmentForTopics(immutable.Set(topic)).map {
-          case (topicPartition, replicas) => topicPartition.partition -> replicas
+      try {
+        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.")
 
         val oldNumPartitions = existingAssignment.size
-        val newNumPartitions = newPartition.totalCount
+        val newNumPartitions = newPartition.count
         val numPartitionsIncrement = newNumPartitions - oldNumPartitions
         if (numPartitionsIncrement < 0) {
           throw new InvalidPartitionsException(
@@ -233,7 +319,10 @@ class AdminManager(val config: KafkaConfig,
           throw new InvalidPartitionsException(s"Topic already has $oldNumPartitions partitions.")
         }
 
-        val reassignment = Option(newPartition.assignments).map(_.asScala.map(_.asScala.map(_.toInt))).map { assignments =>
+        val newPartitionsAssignment = Option(newPartition.assignments).map { assignmentMap =>
+          val assignments = assignmentMap.asScala.map {
+            createPartitionAssignment => createPartitionAssignment.brokerIds.asScala.map(_.toInt)
+          }
           val unknownBrokers = assignments.flatten.toSet -- allBrokerIds
           if (unknownBrokers.nonEmpty)
             throw new InvalidReplicaAssignmentException(
@@ -249,14 +338,25 @@ class AdminManager(val config: KafkaConfig,
           }.toMap
         }
 
-        val updatedReplicaAssignment = adminZkClient.addPartitions(topic, existingAssignment, allBrokers,
-          newPartition.totalCount, reassignment, validateOnly = validateOnly)
-        CreatePartitionsMetadata(topic, updatedReplicaAssignment, ApiError.NONE)
+        val assignmentForNewPartitions = adminZkClient.createNewPartitionsAssignment(
+          topic, existingAssignment, allBrokers, newPartition.count, newPartitionsAssignment)
+
+        if (validateOnly) {
+          CreatePartitionsMetadata(topic, (existingAssignment ++ assignmentForNewPartitions).keySet)
+        } else {
+          controllerMutationQuota.record(numPartitionsIncrement)
+          val updatedReplicaAssignment = adminZkClient.createPartitionsWithAssignment(
+            topic, existingAssignment, assignmentForNewPartitions)
+          CreatePartitionsMetadata(topic, updatedReplicaAssignment.keySet)
+        }
       } catch {
         case e: AdminOperationException =>
-          CreatePartitionsMetadata(topic, Map.empty, ApiError.fromThrowable(e))
+          CreatePartitionsMetadata(topic, e)
+        case e: ThrottlingQuotaExceededException =>
+          debug(s"Partition(s) creation not allowed because quota is violated. Delay time: ${e.throttleTimeMs}")
+          CreatePartitionsMetadata(topic, e)
         case e: ApiException =>
-          CreatePartitionsMetadata(topic, Map.empty, ApiError.fromThrowable(e))
+          CreatePartitionsMetadata(topic, e)
       }
     }
 
@@ -264,7 +364,7 @@ class AdminManager(val config: KafkaConfig,
     if (timeout <= 0 || validateOnly || !metadata.exists(_.error.is(Errors.NONE))) {
       val results = metadata.map { createPartitionMetadata =>
         // ignore topics that already have errors
-        if (createPartitionMetadata.error.isSuccess() && !validateOnly) {
+        if (createPartitionMetadata.error.isSuccess && !validateOnly) {
           (createPartitionMetadata.topic, new ApiError(Errors.REQUEST_TIMED_OUT, null))
         } else {
           (createPartitionMetadata.topic, createPartitionMetadata.error)
@@ -273,58 +373,73 @@ class AdminManager(val config: KafkaConfig,
       callback(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, callback)
-      val delayedCreateKeys = newPartitions.keySet.map(new TopicKey(_)).toSeq
+      val delayedCreate = new DelayedCreatePartitions(timeout, metadata, this, callback)
+      val delayedCreateKeys = newPartitions.map(createPartitionTopic => TopicKey(createPartitionTopic.name))
       // try to complete the request immediately, otherwise put it into the purgatory
       topicPurgatory.tryCompleteElseWatch(delayedCreate, delayedCreateKeys)
     }
   }
 
-  def describeConfigs(resourceToConfigNames: Map[Resource, Option[Set[String]]]): Map[Resource, DescribeConfigsResponse.Config] = {
-    resourceToConfigNames.map { case (resource, configNames) =>
-
-      def createResponseConfig(config: AbstractConfig, isReadOnly: Boolean, isDefault: String => Boolean): DescribeConfigsResponse.Config = {
-        val filteredConfigPairs = config.values.asScala.filter { case (configName, _) =>
-          /* Always returns true if configNames is None */
-          configNames.map(_.contains(configName)).getOrElse(true)
-        }.toIndexedSeq
+  def describeConfigs(resourceToConfigNames: List[DescribeConfigsResource],
+                      includeSynonyms: Boolean,
+                      includeDocumentation: Boolean): List[DescribeConfigsResponseData.DescribeConfigsResult] = {
+    resourceToConfigNames.map { case resource =>
 
-        val configEntries = filteredConfigPairs.map { case (name, value) =>
-          val configEntryType = config.typeOf(name)
-          val isSensitive = configEntryType == ConfigDef.Type.PASSWORD
-          val valueAsString =
-            if (isSensitive) null
-            else ConfigDef.convertToString(value, configEntryType)
-          new DescribeConfigsResponse.ConfigEntry(name, valueAsString, isSensitive, isDefault(name), isReadOnly)
-        }
+      def allConfigs(config: AbstractConfig) = {
+        config.originals.asScala.filter(_._2 != null) ++ config.nonInternalValues.asScala
+      }
+      def createResponseConfig(configs: Map[String, Any],
+                               createConfigEntry: (String, Any) => DescribeConfigsResponseData.DescribeConfigsResourceResult): DescribeConfigsResponseData.DescribeConfigsResult = {
+        val filteredConfigPairs = if (resource.configurationKeys == null)
+          configs.toBuffer
+        else
+          configs.filter { case (configName, _) =>
+            resource.configurationKeys.asScala.forall(_.contains(configName))
+          }.toBuffer
 
-        new DescribeConfigsResponse.Config(ApiError.NONE, configEntries.asJava)
+        val configEntries = filteredConfigPairs.map { case (name, value) => createConfigEntry(name, value) }
+        new DescribeConfigsResponseData.DescribeConfigsResult().setErrorCode(Errors.NONE.code)
+          .setConfigs(configEntries.asJava)
       }
 
       try {
-        val resourceConfig = resource.`type` match {
-
-          case ResourceType.TOPIC =>
-            val topic = resource.name
+        val configResult = ConfigResource.Type.forId(resource.resourceType) match {
+          case ConfigResource.Type.TOPIC =>
+            val topic = resource.resourceName
             Topic.validate(topic)
-            // Consider optimizing this by caching the configs or retrieving them from the `Log` when possible
-            val topicProps = adminZkClient.fetchEntityConfig(ConfigType.Topic, topic)
-            val logConfig = LogConfig.fromProps(KafkaServer.copyKafkaConfigToLog(config), topicProps)
-            createResponseConfig(logConfig, isReadOnly = false, name => !topicProps.containsKey(name))
-
-          case ResourceType.BROKER =>
-            val brokerId = try resource.name.toInt catch {
-              case _: NumberFormatException =>
-                throw new InvalidRequestException(s"Broker id must be an integer, but it is: ${resource.name}")
+            if (metadataCache.contains(topic)) {
+              // Consider optimizing this by caching the configs or retrieving them from the `Log` when possible
+              val topicProps = adminZkClient.fetchEntityConfig(ConfigType.Topic, topic)
+              val logConfig = LogConfig.fromProps(KafkaServer.copyKafkaConfigToLog(config), topicProps)
+              createResponseConfig(allConfigs(logConfig), createTopicConfigEntry(logConfig, topicProps, includeSynonyms, includeDocumentation))
+            } else {
+              new DescribeConfigsResponseData.DescribeConfigsResult().setErrorCode(Errors.UNKNOWN_TOPIC_OR_PARTITION.code)
+                .setConfigs(Collections.emptyList[DescribeConfigsResponseData.DescribeConfigsResourceResult])
             }
-            if (brokerId == config.brokerId)
-              createResponseConfig(config, isReadOnly = true, name => !config.originals.containsKey(name))
+
+          case ConfigResource.Type.BROKER =>
+            if (resource.resourceName == null || resource.resourceName.isEmpty)
+              createResponseConfig(config.dynamicConfig.currentDynamicDefaultConfigs,
+                  createBrokerConfigEntry(perBrokerConfig = false, includeSynonyms, includeDocumentation))
+            else if (resourceNameToBrokerId(resource.resourceName) == config.brokerId)
+              createResponseConfig(allConfigs(config),
+                  createBrokerConfigEntry(perBrokerConfig = true, includeSynonyms, includeDocumentation))
             else
-              throw new InvalidRequestException(s"Unexpected broker id, expected ${config.brokerId}, but received $brokerId")
+              throw new InvalidRequestException(s"Unexpected broker id, expected ${config.brokerId} or empty string, but received ${resource.resourceName}")
 
+          case ConfigResource.Type.BROKER_LOGGER =>
+            if (resource.resourceName == null || resource.resourceName.isEmpty)
+              throw new InvalidRequestException("Broker id must not be empty")
+            else if (resourceNameToBrokerId(resource.resourceName) != config.brokerId)
+              throw new InvalidRequestException(s"Unexpected broker id, expected ${config.brokerId} but received ${resource.resourceName}")
+            else
+              createResponseConfig(Log4jController.loggers,
+                (name, value) => new DescribeConfigsResponseData.DescribeConfigsResourceResult().setName(name)
+                  .setValue(value.toString).setConfigSource(ConfigSource.DYNAMIC_BROKER_LOGGER_CONFIG.id)
+                  .setIsSensitive(false).setReadOnly(false).setSynonyms(List.empty.asJava))
           case resourceType => throw new InvalidRequestException(s"Unsupported resource type: $resourceType")
         }
-        resource -> resourceConfig
+        configResult.setResourceName(resource.resourceName).setResourceType(resource.resourceType)
       } catch {
         case e: Throwable =>
           // Log client errors at a lower level than unexpected exceptions
@@ -333,52 +448,173 @@ class AdminManager(val config: KafkaConfig,
             info(message, e)
           else
             error(message, e)
-          resource -> new DescribeConfigsResponse.Config(ApiError.fromThrowable(e), Collections.emptyList[DescribeConfigsResponse.ConfigEntry])
+          val err = ApiError.fromThrowable(e)
+          new DescribeConfigsResponseData.DescribeConfigsResult()
+              .setResourceName(resource.resourceName)
+              .setResourceType(resource.resourceType)
+              .setErrorMessage(err.message)
+              .setErrorCode(err.error.code)
+              .setConfigs(Collections.emptyList[DescribeConfigsResponseData.DescribeConfigsResourceResult])
       }
-    }.toMap
+    }.toList
   }
 
-  def alterConfigs(configs: Map[Resource, AlterConfigsRequest.Config], validateOnly: Boolean): Map[Resource, ApiError] = {
+  def alterConfigs(configs: Map[ConfigResource, AlterConfigsRequest.Config], validateOnly: Boolean): Map[ConfigResource, ApiError] = {
     configs.map { case (resource, config) =>
+
+      try {
+        val nullUpdates = config.entries.asScala.filter(_.value == null).map(_.name)
+        if (nullUpdates.nonEmpty)
+          throw new InvalidRequestException(s"Null value not supported for : ${nullUpdates.mkString(",")}")
+
+        val configEntriesMap = config.entries.asScala.map(entry => (entry.name, entry.value)).toMap
+
+        val configProps = new Properties
+        config.entries.asScala.filter(_.value != null).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
+    if (!metadataCache.contains(topic))
+      throw new UnknownTopicOrPartitionException(s"The topic '$topic' does not exist.")
+
+    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: Seq[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)
+        case _ => throw new IllegalArgumentException(
+          s"Log level cannot be changed for OpType: ${alterConfigOp.opType()}")
+      }
+    }
+  }
+
+  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, Seq[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).filter { case (_, v) =>
+          v.size > 1
+        }.keySet
+        if (duplicateKeys.nonEmpty)
+          throw new InvalidRequestException(s"Error due to duplicate config keys : ${duplicateKeys.mkString(",")}")
+        val nullUpdates = alterConfigOps
+          .filter(entry => entry.configEntry.value == null && entry.opType() != OpType.DELETE)
+          .map(entry => s"${entry.opType}:${entry.configEntry.name}")
+        if (nullUpdates.nonEmpty)
+          throw new InvalidRequestException(s"Null value not supported for : ${nullUpdates.mkString(",")}")
+
+        val configEntriesMap = alterConfigOps.map(entry => (entry.configEntry.name, entry.configEntry.value)).toMap
+
         resource.`type` match {
-          case ResourceType.TOPIC =>
-            val topic = resource.name
+          case ConfigResource.Type.TOPIC =>
+            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
 
-            alterConfigPolicy match {
-              case Some(policy) =>
-                adminZkClient.validateTopicConfig(topic, properties)
+            val persistentProps = if (perBrokerConfig) adminZkClient.fetchEntityConfig(ConfigType.Broker, brokerId.get.toString)
+            else adminZkClient.fetchEntityConfig(ConfigType.Broker, ConfigEntityName.Default)
 
-                val configEntriesMap = config.entries.asScala.map(entry => (entry.name, entry.value)).toMap
-                policy.validate(new AlterConfigPolicy.RequestMetadata(
-                  new ConfigResource(ConfigResource.Type.TOPIC, resource.name), configEntriesMap.asJava))
+            val configProps = this.config.dynamicConfig.fromPersistentProps(persistentProps, perBrokerConfig)
+            prepareIncrementalConfigs(alterConfigOps, configProps, KafkaConfig.configKeys)
+            alterBrokerConfigs(resource, validateOnly, configProps, configEntriesMap)
 
-                if (!validateOnly)
-                  adminZkClient.changeTopicConfig(topic, properties)
-              case None =>
-                if (validateOnly)
-                  adminZkClient.validateTopicConfig(topic, properties)
-                else
-                  adminZkClient.changeTopicConfig(topic, properties)
-            }
+          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, but resource type is $resourceType")
+            throw new InvalidRequestException(s"AlterConfigs is only supported for topics and brokers, but resource type is $resourceType")
         }
       } catch {
-        case e: ConfigException =>
+        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] || e.isInstanceOf[PolicyViolationException])
+          val message = s"Error processing alter configs request for resource $resource, config $alterConfigOps"
+          if (e.isInstanceOf[ApiException])
             info(message, e)
           else
             error(message, e)
@@ -387,9 +623,670 @@ class AdminManager(val config: KafkaConfig,
     }.toMap
   }
 
-  def shutdown() {
+  private def validateLogLevelConfigs(alterConfigOps: Seq[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: Seq[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 =>
+      val configPropName = alterConfigOp.configEntry.name
+      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 = Option(configProps.getProperty(alterConfigOp.configEntry.name))
+            .orElse(Option(ConfigDef.convertToString(configKeys(configPropName).defaultValue, ConfigDef.Type.LIST)))
+            .getOrElse("")
+            .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 = Option(configProps.getProperty(alterConfigOp.configEntry.name))
+            .orElse(Option(ConfigDef.convertToString(configKeys(configPropName).defaultValue, ConfigDef.Type.LIST)))
+            .getOrElse("")
+            .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)
   }
+
+  private def resourceNameToBrokerId(resourceName: String): Int = {
+    try resourceName.toInt catch {
+      case _: NumberFormatException =>
+        throw new InvalidRequestException(s"Broker id must be an integer, but it is: $resourceName")
+    }
+  }
+
+  private def brokerSynonyms(name: String): List[String] = {
+    DynamicBrokerConfig.brokerConfigSynonyms(name, matchListenerOverride = true)
+  }
+
+  private def brokerDocumentation(name: String): String = {
+    config.documentationOf(name)
+  }
+
+  private def configResponseType(configType: Option[ConfigDef.Type]): DescribeConfigsResponse.ConfigType = {
+    if (configType.isEmpty)
+      DescribeConfigsResponse.ConfigType.UNKNOWN
+    else configType.get match {
+      case ConfigDef.Type.BOOLEAN => DescribeConfigsResponse.ConfigType.BOOLEAN
+      case ConfigDef.Type.STRING => DescribeConfigsResponse.ConfigType.STRING
+      case ConfigDef.Type.INT => DescribeConfigsResponse.ConfigType.INT
+      case ConfigDef.Type.SHORT => DescribeConfigsResponse.ConfigType.SHORT
+      case ConfigDef.Type.LONG => DescribeConfigsResponse.ConfigType.LONG
+      case ConfigDef.Type.DOUBLE => DescribeConfigsResponse.ConfigType.DOUBLE
+      case ConfigDef.Type.LIST => DescribeConfigsResponse.ConfigType.LIST
+      case ConfigDef.Type.CLASS => DescribeConfigsResponse.ConfigType.CLASS
+      case ConfigDef.Type.PASSWORD => DescribeConfigsResponse.ConfigType.PASSWORD
+      case _ => DescribeConfigsResponse.ConfigType.UNKNOWN
+    }
+  }
+
+  private def configSynonyms(name: String, synonyms: List[String], isSensitive: Boolean): List[DescribeConfigsResponseData.DescribeConfigsSynonym] = {
+    val dynamicConfig = config.dynamicConfig
+    val allSynonyms = mutable.Buffer[DescribeConfigsResponseData.DescribeConfigsSynonym]()
+
+    def maybeAddSynonym(map: Map[String, String], source: ConfigSource)(name: String): Unit = {
+      map.get(name).map { value =>
+        val configValue = if (isSensitive) null else value
+        allSynonyms += new DescribeConfigsResponseData.DescribeConfigsSynonym().setName(name).setValue(configValue).setSource(source.id)
+      }
+    }
+
+    synonyms.foreach(maybeAddSynonym(dynamicConfig.currentDynamicBrokerConfigs, ConfigSource.DYNAMIC_BROKER_CONFIG))
+    synonyms.foreach(maybeAddSynonym(dynamicConfig.currentDynamicDefaultConfigs, ConfigSource.DYNAMIC_DEFAULT_BROKER_CONFIG))
+    synonyms.foreach(maybeAddSynonym(dynamicConfig.staticBrokerConfigs, ConfigSource.STATIC_BROKER_CONFIG))
+    synonyms.foreach(maybeAddSynonym(dynamicConfig.staticDefaultConfigs, ConfigSource.DEFAULT_CONFIG))
+    allSynonyms.dropWhile(s => s.name != name).toList // e.g. drop listener overrides when describing base config
+  }
+
+  private def createTopicConfigEntry(logConfig: LogConfig, topicProps: Properties, includeSynonyms: Boolean, includeDocumentation: Boolean)
+                                    (name: String, value: Any): DescribeConfigsResponseData.DescribeConfigsResourceResult = {
+    val configEntryType = LogConfig.configType(name)
+    val isSensitive = KafkaConfig.maybeSensitive(configEntryType)
+    val valueAsString = if (isSensitive) null else ConfigDef.convertToString(value, configEntryType.orNull)
+    val allSynonyms = {
+      val list = LogConfig.TopicConfigSynonyms.get(name)
+        .map(s => configSynonyms(s, brokerSynonyms(s), isSensitive))
+        .getOrElse(List.empty)
+      if (!topicProps.containsKey(name))
+        list
+      else
+        new DescribeConfigsResponseData.DescribeConfigsSynonym().setName(name).setValue(valueAsString)
+          .setSource(ConfigSource.TOPIC_CONFIG.id) +: list
+    }
+    val source = if (allSynonyms.isEmpty) ConfigSource.DEFAULT_CONFIG.id else allSynonyms.head.source
+    val synonyms = if (!includeSynonyms) List.empty else allSynonyms
+    val dataType = configResponseType(configEntryType)
+    val configDocumentation = if (includeDocumentation) logConfig.documentationOf(name) else null
+    new DescribeConfigsResponseData.DescribeConfigsResourceResult()
+      .setName(name).setValue(valueAsString).setConfigSource(source)
+      .setIsSensitive(isSensitive).setReadOnly(false).setSynonyms(synonyms.asJava)
+      .setDocumentation(configDocumentation).setConfigType(dataType.id)
+  }
+
+  private def createBrokerConfigEntry(perBrokerConfig: Boolean, includeSynonyms: Boolean, includeDocumentation: Boolean)
+                                     (name: String, value: Any): DescribeConfigsResponseData.DescribeConfigsResourceResult = {
+    val allNames = brokerSynonyms(name)
+    val configEntryType = KafkaConfig.configType(name)
+    val isSensitive = KafkaConfig.maybeSensitive(configEntryType)
+    val valueAsString = if (isSensitive)
+      null
+    else value match {
+      case v: String => v
+      case _ => ConfigDef.convertToString(value, configEntryType.orNull)
+    }
+    val allSynonyms = configSynonyms(name, allNames, isSensitive)
+        .filter(perBrokerConfig || _.source == ConfigSource.DYNAMIC_DEFAULT_BROKER_CONFIG.id)
+    val synonyms = if (!includeSynonyms) List.empty else allSynonyms
+    val source = if (allSynonyms.isEmpty) ConfigSource.DEFAULT_CONFIG.id else allSynonyms.head.source
+    val readOnly = !DynamicBrokerConfig.AllDynamicConfigs.contains(name)
+
+    val dataType = configResponseType(configEntryType)
+    val configDocumentation = if (includeDocumentation) brokerDocumentation(name) else null
+    new DescribeConfigsResponseData.DescribeConfigsResourceResult().setName(name).setValue(valueAsString).setConfigSource(source)
+      .setIsSensitive(isSensitive).setReadOnly(readOnly).setSynonyms(synonyms.asJava)
+      .setDocumentation(configDocumentation).setConfigType(dataType.id)
+  }
+
+  private def sanitizeEntityName(entityName: String): String =
+    Option(entityName) match {
+      case None => ConfigEntityName.Default
+      case Some(name) => Sanitizer.sanitize(name)
+    }
+
+  private def desanitizeEntityName(sanitizedEntityName: String): String =
+    sanitizedEntityName match {
+      case ConfigEntityName.Default => null
+      case name => Sanitizer.desanitize(name)
+    }
+
+  private def parseAndSanitizeQuotaEntity(entity: ClientQuotaEntity): (Option[String], Option[String], Option[String]) = {
+    if (entity.entries.isEmpty)
+      throw new InvalidRequestException("Invalid empty client quota entity")
+
+    var user: Option[String] = None
+    var clientId: Option[String] = None
+    var ip: Option[String] = None
+    entity.entries.forEach { (entityType, entityName) =>
+      val sanitizedEntityName = Some(sanitizeEntityName(entityName))
+      entityType match {
+        case ClientQuotaEntity.USER => user = sanitizedEntityName
+        case ClientQuotaEntity.CLIENT_ID => clientId = sanitizedEntityName
+        case ClientQuotaEntity.IP => ip = sanitizedEntityName
+        case _ => throw new InvalidRequestException(s"Unhandled client quota entity type: ${entityType}")
+      }
+      if (entityName != null && entityName.isEmpty)
+        throw new InvalidRequestException(s"Empty ${entityType} not supported")
+    }
+    (user, clientId, ip)
+  }
+
+  private def userClientIdToEntity(user: Option[String], clientId: Option[String]): ClientQuotaEntity = {
+    new ClientQuotaEntity((user.map(u => ClientQuotaEntity.USER -> u) ++ clientId.map(c => ClientQuotaEntity.CLIENT_ID -> c)).toMap.asJava)
+  }
+
+  def describeClientQuotas(filter: ClientQuotaFilter): Map[ClientQuotaEntity, Map[String, Double]] = {
+    var userComponent: Option[ClientQuotaFilterComponent] = None
+    var clientIdComponent: Option[ClientQuotaFilterComponent] = None
+    var ipComponent: Option[ClientQuotaFilterComponent] = None
+    filter.components.forEach { component =>
+      component.entityType match {
+        case ClientQuotaEntity.USER =>
+          if (userComponent.isDefined)
+            throw new InvalidRequestException(s"Duplicate user filter component entity type")
+          userComponent = Some(component)
+        case ClientQuotaEntity.CLIENT_ID =>
+          if (clientIdComponent.isDefined)
+            throw new InvalidRequestException(s"Duplicate client filter component entity type")
+          clientIdComponent = Some(component)
+        case ClientQuotaEntity.IP =>
+          if (ipComponent.isDefined)
+            throw new InvalidRequestException(s"Duplicate ip filter component entity type")
+          ipComponent = Some(component)
+        case "" =>
+          throw new InvalidRequestException(s"Unexpected empty filter component entity type")
+        case et =>
+          // Supplying other entity types is not yet supported.
+          throw new UnsupportedVersionException(s"Custom entity type '${et}' not supported")
+      }
+    }
+    if ((userComponent.isDefined || clientIdComponent.isDefined) && ipComponent.isDefined)
+      throw new InvalidRequestException(s"Invalid entity filter component combination, IP filter component should not be used with " +
+        s"user or clientId filter component.")
+
+    val userClientQuotas = if (ipComponent.isEmpty)
+      handleDescribeClientQuotas(userComponent, clientIdComponent, filter.strict)
+    else
+      Map.empty
+
+    val ipQuotas = if (userComponent.isEmpty && clientIdComponent.isEmpty)
+      handleDescribeIpQuotas(ipComponent, filter.strict)
+    else
+      Map.empty
+
+    (userClientQuotas ++ ipQuotas).toMap
+  }
+
+  private def wantExact(component: Option[ClientQuotaFilterComponent]): Boolean = component.exists(_.`match` != null)
+
+  private def toOption(opt: java.util.Optional[String]): Option[String] = {
+    if (opt == null)
+      None
+    else if (opt.isPresent)
+      Some(opt.get)
+    else
+      Some(null)
+  }
+
+  private def sanitized(name: Option[String]): String = name.map(n => sanitizeEntityName(n)).getOrElse("")
+
+  private def fromProps(props: Map[String, String]): Map[String, Double] = {
+    props.map { case (key, value) =>
+      val doubleValue = try value.toDouble catch {
+        case _: NumberFormatException =>
+          throw new IllegalStateException(s"Unexpected client quota configuration value: $key -> $value")
+      }
+      key -> doubleValue
+    }
+  }
+
+  def handleDescribeClientQuotas(userComponent: Option[ClientQuotaFilterComponent],
+    clientIdComponent: Option[ClientQuotaFilterComponent], strict: Boolean): Map[ClientQuotaEntity, Map[String, Double]] = {
+
+    val user = userComponent.flatMap(c => toOption(c.`match`))
+    val clientId = clientIdComponent.flatMap(c => toOption(c.`match`))
+
+    val sanitizedUser = sanitized(user)
+    val sanitizedClientId = sanitized(clientId)
+
+    val exactUser = wantExact(userComponent)
+    val exactClientId = wantExact(clientIdComponent)
+
+    def wantExcluded(component: Option[ClientQuotaFilterComponent]): Boolean = strict && !component.isDefined
+    val excludeUser = wantExcluded(userComponent)
+    val excludeClientId = wantExcluded(clientIdComponent)
+
+    val userEntries = if (exactUser && excludeClientId)
+      Map((Some(user.get), None) -> adminZkClient.fetchEntityConfig(ConfigType.User, sanitizedUser))
+    else if (!excludeUser && !exactClientId)
+      adminZkClient.fetchAllEntityConfigs(ConfigType.User).map { case (name, props) =>
+        (Some(desanitizeEntityName(name)), None) -> props
+      }
+    else
+      Map.empty
+
+    val clientIdEntries = if (excludeUser && exactClientId)
+      Map((None, Some(clientId.get)) -> adminZkClient.fetchEntityConfig(ConfigType.Client, sanitizedClientId))
+    else if (!exactUser && !excludeClientId)
+      adminZkClient.fetchAllEntityConfigs(ConfigType.Client).map { case (name, props) =>
+        (None, Some(desanitizeEntityName(name))) -> props
+      }
+    else
+      Map.empty
+
+    val bothEntries = if (exactUser && exactClientId)
+      Map((Some(user.get), Some(clientId.get)) ->
+        adminZkClient.fetchEntityConfig(ConfigType.User, s"${sanitizedUser}/clients/${sanitizedClientId}"))
+    else if (!excludeUser && !excludeClientId)
+      adminZkClient.fetchAllChildEntityConfigs(ConfigType.User, ConfigType.Client).map { case (name, props) =>
+        val components = name.split("/")
+        if (components.size != 3 || components(1) != "clients")
+          throw new IllegalArgumentException(s"Unexpected config path: ${name}")
+        (Some(desanitizeEntityName(components(0))), Some(desanitizeEntityName(components(2)))) -> props
+      }
+    else
+      Map.empty
+
+    def matches(nameComponent: Option[ClientQuotaFilterComponent], name: Option[String]): Boolean = nameComponent match {
+      case Some(component) =>
+        toOption(component.`match`) match {
+          case Some(n) => name.exists(_ == n)
+          case None => name.isDefined
+        }
+      case None =>
+        !name.isDefined || !strict
+    }
+
+    (userEntries ++ clientIdEntries ++ bothEntries).flatMap { case ((u, c), p) =>
+      val quotaProps = p.asScala.filter { case (key, _) => QuotaConfigs.isQuotaConfig(key) }
+      if (quotaProps.nonEmpty && matches(userComponent, u) && matches(clientIdComponent, c))
+        Some(userClientIdToEntity(u, c) -> fromProps(quotaProps))
+      else
+        None
+    }.toMap
+  }
+
+  def handleDescribeIpQuotas(ipComponent: Option[ClientQuotaFilterComponent], strict: Boolean): Map[ClientQuotaEntity, Map[String, Double]] = {
+    val ip = ipComponent.flatMap(c => toOption(c.`match`))
+    val exactIp = wantExact(ipComponent)
+    val allIps = ipComponent.exists(_.`match` == null) || (ipComponent.isEmpty && !strict)
+    val ipEntries = if (exactIp)
+      Map(Some(ip.get) -> adminZkClient.fetchEntityConfig(ConfigType.Ip, sanitized(ip)))
+    else if (allIps)
+      adminZkClient.fetchAllEntityConfigs(ConfigType.Ip).map { case (name, props) =>
+        Some(desanitizeEntityName(name)) -> props
+      }
+    else
+      Map.empty
+
+    def ipToQuotaEntity(ip: Option[String]): ClientQuotaEntity = {
+      new ClientQuotaEntity(ip.map(ipName => ClientQuotaEntity.IP -> ipName).toMap.asJava)
+    }
+
+    ipEntries.flatMap { case (ip, props) =>
+      val ipQuotaProps = props.asScala.filter { case (key, _) => DynamicConfig.Ip.names.contains(key) }
+      if (ipQuotaProps.nonEmpty)
+        Some(ipToQuotaEntity(ip) -> fromProps(ipQuotaProps))
+      else
+        None
+    }
+  }
+
+  def alterClientQuotas(entries: Seq[ClientQuotaAlteration], validateOnly: Boolean): Map[ClientQuotaEntity, ApiError] = {
+    def alterEntityQuotas(entity: ClientQuotaEntity, ops: Iterable[ClientQuotaAlteration.Op]): Unit = {
+      val (path, configType, configKeys) = parseAndSanitizeQuotaEntity(entity) match {
+        case (Some(user), Some(clientId), None) => (user + "/clients/" + clientId, ConfigType.User, DynamicConfig.User.configKeys)
+        case (Some(user), None, None) => (user, ConfigType.User, DynamicConfig.User.configKeys)
+        case (None, Some(clientId), None) => (clientId, ConfigType.Client, DynamicConfig.Client.configKeys)
+        case (None, None, Some(ip)) =>
+          if (!DynamicConfig.Ip.isValidIpEntity(ip))
+            throw new InvalidRequestException(s"$ip is not a valid IP or resolvable host.")
+          (ip, ConfigType.Ip, DynamicConfig.Ip.configKeys)
+        case (_, _, Some(_)) => throw new InvalidRequestException(s"Invalid quota entity combination, " +
+          s"IP entity should not be used with user/client ID entity.")
+        case _ => throw new InvalidRequestException("Invalid client quota entity")
+      }
+
+      val props = adminZkClient.fetchEntityConfig(configType, path)
+      ops.foreach { op =>
+        op.value match {
+          case null =>
+            props.remove(op.key)
+          case value => configKeys.get(op.key) match {
+            case null =>
+              throw new InvalidRequestException(s"Invalid configuration key ${op.key}")
+            case key => key.`type` match {
+              case ConfigDef.Type.DOUBLE =>
+                props.setProperty(op.key, value.toString)
+              case ConfigDef.Type.LONG | ConfigDef.Type.INT =>
+                val epsilon = 1e-6
+                val intValue = if (key.`type` == ConfigDef.Type.LONG)
+                  (value + epsilon).toLong
+                else
+                  (value + epsilon).toInt
+                if ((intValue.toDouble - value).abs > epsilon)
+                  throw new InvalidRequestException(s"Configuration ${op.key} must be a ${key.`type`} value")
+                props.setProperty(op.key, intValue.toString)
+              case _ =>
+                throw new IllegalStateException(s"Unexpected config type ${key.`type`}")
+            }
+          }
+        }
+      }
+      if (!validateOnly)
+        adminZkClient.changeConfigs(configType, path, props)
+    }
+    entries.map { entry =>
+      val apiError = try {
+        alterEntityQuotas(entry.entity, entry.ops.asScala)
+        ApiError.NONE
+      } catch {
+        case e: Throwable =>
+          info(s"Error encountered while updating client quotas", e)
+          ApiError.fromThrowable(e)
+      }
+      entry.entity -> apiError
+    }.toMap
+  }
+
+  private val usernameMustNotBeEmptyMsg = "Username must not be empty"
+  private val errorProcessingDescribe = "Error processing describe user SCRAM credential configs request"
+  private val attemptToDescribeUserThatDoesNotExist = "Attempt to describe a user credential that does not exist"
+
+  def describeUserScramCredentials(users: Option[Seq[String]]): DescribeUserScramCredentialsResponseData = {
+    val describingAllUsers = !users.isDefined || users.get.isEmpty
+    val retval = new DescribeUserScramCredentialsResponseData()
+    val userResults = mutable.Map[String, DescribeUserScramCredentialsResponseData.DescribeUserScramCredentialsResult]()
+
+    def addToResultsIfHasScramCredential(user: String, userConfig: Properties, explicitUser: Boolean = false): Unit = {
+      val result = new DescribeUserScramCredentialsResponseData.DescribeUserScramCredentialsResult().setUser(user)
+      val configKeys = userConfig.stringPropertyNames
+      val hasScramCredential = ScramMechanism.values().toList.exists(key => key != ScramMechanism.UNKNOWN && configKeys.contains(key.mechanismName))
+      if (hasScramCredential) {
+        val credentialInfos = new util.ArrayList[CredentialInfo]
+        try {
+          ScramMechanism.values().filter(_ != ScramMechanism.UNKNOWN).foreach { mechanism =>
+            val propertyValue = userConfig.getProperty(mechanism.mechanismName)
+            if (propertyValue != null) {
+              val iterations = ScramCredentialUtils.credentialFromString(propertyValue).iterations
+              credentialInfos.add(new CredentialInfo().setMechanism(mechanism.`type`).setIterations(iterations))
+            }
+          }
+          result.setCredentialInfos(credentialInfos)
+        } catch {
+          case e: Exception => { // should generally never happen, but just in case bad data gets in...
+            val apiError = apiErrorFrom(e, errorProcessingDescribe)
+            result.setErrorCode(apiError.error.code).setErrorMessage(apiError.error.message)
+          }
+        }
+        userResults += (user -> result)
+      } else if (explicitUser) {
+        // it is an error to request credentials for a user that has no credentials
+        result.setErrorCode(Errors.RESOURCE_NOT_FOUND.code).setErrorMessage(s"$attemptToDescribeUserThatDoesNotExist: $user")
+        userResults += (user -> result)
+      }
+    }
+
+    def collectRetrievedResults(): Unit = {
+      if (describingAllUsers) {
+        val usersSorted = SortedSet.empty[String] ++ userResults.keys
+        usersSorted.foreach { user => retval.results.add(userResults(user)) }
+      } else {
+        // be sure to only include a single copy of a result for any user requested multiple times
+        users.get.distinct.foreach { user =>  retval.results.add(userResults(user)) }
+      }
+    }
+
+    try {
+      if (describingAllUsers)
+        adminZkClient.fetchAllEntityConfigs(ConfigType.User).foreach {
+          case (user, properties) => addToResultsIfHasScramCredential(user, properties) }
+      else {
+        // describing specific users
+        val illegalUsers = users.get.filter(_.isEmpty).toSet
+        illegalUsers.foreach { user =>
+          userResults += (user -> new DescribeUserScramCredentialsResponseData.DescribeUserScramCredentialsResult()
+            .setUser(user)
+            .setErrorCode(Errors.RESOURCE_NOT_FOUND.code)
+            .setErrorMessage(usernameMustNotBeEmptyMsg)) }
+        val duplicatedUsers = users.get.groupBy(identity).filter(
+          userAndOccurrencesTuple => userAndOccurrencesTuple._2.length > 1).keys
+        duplicatedUsers.filterNot(illegalUsers.contains).foreach { user =>
+          userResults += (user -> new DescribeUserScramCredentialsResponseData.DescribeUserScramCredentialsResult()
+            .setUser(user)
+            .setErrorCode(Errors.DUPLICATE_RESOURCE.code)
+            .setErrorMessage(s"Cannot describe SCRAM credentials for the same user twice in a single request: $user")) }
+        val usersToSkip = illegalUsers ++ duplicatedUsers
+        users.get.filterNot(usersToSkip.contains).foreach { user =>
+          try {
+            val userConfigs = adminZkClient.fetchEntityConfig(ConfigType.User, Sanitizer.sanitize(user))
+            addToResultsIfHasScramCredential(user, userConfigs, true)
+          } catch {
+            case e: Exception => {
+              val apiError = apiErrorFrom(e, errorProcessingDescribe)
+              userResults += (user -> new DescribeUserScramCredentialsResponseData.DescribeUserScramCredentialsResult()
+                .setUser(user)
+                .setErrorCode(apiError.error.code)
+                .setErrorMessage(apiError.error.message))
+            }
+          }
+        }
+      }
+      collectRetrievedResults()
+    } catch {
+      case e: Exception => {
+        // this should generally only happen when we get a failure trying to retrieve all user configs from ZooKeeper
+        val apiError = apiErrorFrom(e, errorProcessingDescribe)
+        retval.setErrorCode(apiError.error.code).setErrorMessage(apiError.messageWithFallback())
+      }
+    }
+    retval
+  }
+
+  def apiErrorFrom(e: Exception, message: String): ApiError = {
+    if (e.isInstanceOf[ApiException])
+      info(message, e)
+    else
+      error(message, e)
+    ApiError.fromThrowable(e)
+  }
+
+  case class requestStatus(user: String, mechanism: Option[ScramMechanism], legalRequest: Boolean, iterations: Int) {}
+
+  def alterUserScramCredentials(upsertions: Seq[AlterUserScramCredentialsRequestData.ScramCredentialUpsertion],
+                                deletions: Seq[AlterUserScramCredentialsRequestData.ScramCredentialDeletion]): AlterUserScramCredentialsResponseData = {
+
+    def scramMechanism(mechanism: Byte): ScramMechanism = {
+      ScramMechanism.fromType(mechanism)
+    }
+
+    def mechanismName(mechanism: Byte): String = {
+      scramMechanism(mechanism).mechanismName
+    }
+
+    val retval = new AlterUserScramCredentialsResponseData()
+
+    // fail any user that is invalid due to an empty user name, an unknown SCRAM mechanism, or unacceptable number of iterations
+    val maxIterations = 16384
+    val illegalUpsertions = upsertions.map(upsertion =>
+      if (upsertion.name.isEmpty)
+        requestStatus(upsertion.name, None, false, upsertion.iterations) // no determined mechanism -- empty user is the cause of failure
+      else {
+        val publicScramMechanism = scramMechanism(upsertion.mechanism)
+        if (publicScramMechanism == ScramMechanism.UNKNOWN) {
+          requestStatus(upsertion.name, Some(publicScramMechanism), false, upsertion.iterations) // unknown mechanism is the cause of failure
+        } else {
+          if (upsertion.iterations < InternalScramMechanism.forMechanismName(publicScramMechanism.mechanismName).minIterations
+            || upsertion.iterations > maxIterations) {
+            requestStatus(upsertion.name, Some(publicScramMechanism), false, upsertion.iterations) // known mechanism, bad iterations is the cause of failure
+          } else {
+            requestStatus(upsertion.name, Some(publicScramMechanism), true, upsertion.iterations) // legal
+          }
+        }
+      }).filter { !_.legalRequest }
+    val illegalDeletions = deletions.map(deletion =>
+      if (deletion.name.isEmpty) {
+        requestStatus(deletion.name, None, false, 0) // no determined mechanism -- empty user is the cause of failure
+      } else {
+        val publicScramMechanism = scramMechanism(deletion.mechanism)
+        requestStatus(deletion.name, Some(publicScramMechanism), publicScramMechanism != ScramMechanism.UNKNOWN, 0)
+      }).filter { !_.legalRequest }
+    // map user names to error messages
+    val unknownScramMechanismMsg = "Unknown SCRAM mechanism"
+    val tooFewIterationsMsg = "Too few iterations"
+    val tooManyIterationsMsg = "Too many iterations"
+    val illegalRequestsByUser =
+      illegalDeletions.map(requestStatus =>
+        if (requestStatus.user.isEmpty) {
+          (requestStatus.user, usernameMustNotBeEmptyMsg)
+        } else {
+          (requestStatus.user, unknownScramMechanismMsg)
+        }
+      ).toMap ++ illegalUpsertions.map(requestStatus =>
+        if (requestStatus.user.isEmpty) {
+          (requestStatus.user, usernameMustNotBeEmptyMsg)
+        } else if (requestStatus.mechanism == Some(ScramMechanism.UNKNOWN)) {
+          (requestStatus.user, unknownScramMechanismMsg)
+        } else {
+          (requestStatus.user, if (requestStatus.iterations > maxIterations) {tooManyIterationsMsg} else {tooFewIterationsMsg})
+        }
+      ).toMap
+
+    illegalRequestsByUser.forKeyValue { (user, errorMessage) =>
+      retval.results.add(new AlterUserScramCredentialsResult().setUser(user)
+        .setErrorCode(if (errorMessage == unknownScramMechanismMsg) {Errors.UNSUPPORTED_SASL_MECHANISM.code} else {Errors.UNACCEPTABLE_CREDENTIAL.code})
+        .setErrorMessage(errorMessage)) }
+
+    val invalidUsers = (illegalUpsertions ++ illegalDeletions).map(_.user).toSet
+    val initiallyValidUserMechanismPairs = (upsertions.filter(upsertion => !invalidUsers.contains(upsertion.name)).map(upsertion => (upsertion.name, upsertion.mechanism)) ++
+      deletions.filter(deletion => !invalidUsers.contains(deletion.name)).map(deletion => (deletion.name, deletion.mechanism)))
+
+    val usersWithDuplicateUserMechanismPairs = initiallyValidUserMechanismPairs.groupBy(identity).filter (
+      userMechanismPairAndOccurrencesTuple => userMechanismPairAndOccurrencesTuple._2.length > 1).keys.map(userMechanismPair => userMechanismPair._1).toSet
+    usersWithDuplicateUserMechanismPairs.foreach { user =>
+      retval.results.add(new AlterUserScramCredentialsResult()
+        .setUser(user)
+        .setErrorCode(Errors.DUPLICATE_RESOURCE.code).setErrorMessage("A user credential cannot be altered twice in the same request")) }
+
+    def potentiallyValidUserMechanismPairs = initiallyValidUserMechanismPairs.filter(pair => !usersWithDuplicateUserMechanismPairs.contains(pair._1))
+
+    val potentiallyValidUsers = potentiallyValidUserMechanismPairs.map(_._1).toSet
+    val configsByPotentiallyValidUser = potentiallyValidUsers.map(user => (user, adminZkClient.fetchEntityConfig(ConfigType.User, Sanitizer.sanitize(user)))).toMap
+
+    // check for deletion of a credential that does not exist
+    val invalidDeletions = deletions.filter(deletion => potentiallyValidUsers.contains(deletion.name)).filter(deletion =>
+      configsByPotentiallyValidUser(deletion.name).getProperty(mechanismName(deletion.mechanism)) == null)
+    val invalidUsersDueToInvalidDeletions = invalidDeletions.map(_.name).toSet
+    invalidUsersDueToInvalidDeletions.foreach { user =>
+      retval.results.add(new AlterUserScramCredentialsResult()
+        .setUser(user)
+        .setErrorCode(Errors.RESOURCE_NOT_FOUND.code).setErrorMessage("Attempt to delete a user credential that does not exist")) }
+
+    // now prepare the new set of property values for users that don't have any issues identified above,
+    // keeping track of ones that fail
+    val usersToTryToAlter = potentiallyValidUsers.diff(invalidUsersDueToInvalidDeletions)
+    val usersFailedToPrepareProperties = usersToTryToAlter.map(user => {
+      try {
+        // deletions: remove property keys
+        deletions.filter(deletion => usersToTryToAlter.contains(deletion.name)).foreach { deletion =>
+          configsByPotentiallyValidUser(deletion.name).remove(mechanismName(deletion.mechanism)) }
+        // upsertions: put property key/value
+        upsertions.filter(upsertion => usersToTryToAlter.contains(upsertion.name)).foreach { upsertion =>
+          val mechanism = InternalScramMechanism.forMechanismName(mechanismName(upsertion.mechanism))
+          val credential = new ScramFormatter(mechanism)
+            .generateCredential(upsertion.salt, upsertion.saltedPassword, upsertion.iterations)
+          configsByPotentiallyValidUser(upsertion.name).put(mechanismName(upsertion.mechanism), ScramCredentialUtils.credentialToString(credential)) }
+        (user) // success, 1 element, won't be matched
+      } catch {
+        case e: Exception =>
+          info(s"Error encountered while altering user SCRAM credentials", e)
+          (user, e) // fail, 2 elements, will be matched
+      }
+    }).collect { case (user: String, exception: Exception) => (user, exception) }.toMap
+
+    // now persist the properties we have prepared, again keeping track of whatever fails
+    val usersFailedToPersist = usersToTryToAlter.filterNot(usersFailedToPrepareProperties.contains).map(user => {
+      try {
+        adminZkClient.changeConfigs(ConfigType.User, Sanitizer.sanitize(user), configsByPotentiallyValidUser(user))
+        (user) // success, 1 element, won't be matched
+      } catch {
+        case e: Exception =>
+          info(s"Error encountered while altering user SCRAM credentials", e)
+          (user, e) // fail, 2 elements, will be matched
+      }
+    }).collect { case (user: String, exception: Exception) => (user, exception) }.toMap
+
+    // report failures
+    usersFailedToPrepareProperties.++(usersFailedToPersist).forKeyValue { (user, exception) =>
+      val error = Errors.forException(exception)
+      retval.results.add(new AlterUserScramCredentialsResult()
+        .setUser(user)
+        .setErrorCode(error.code)
+        .setErrorMessage(error.message)) }
+
+    // report successes
+    usersToTryToAlter.filterNot(usersFailedToPrepareProperties.contains).filterNot(usersFailedToPersist.contains).foreach { user =>
+      retval.results.add(new AlterUserScramCredentialsResult()
+        .setUser(user)
+        .setErrorCode(Errors.NONE.code)) }
+
+    retval
+  }
 }
diff --git a/core/src/main/scala/kafka/server/AlterIsrManager.scala b/core/src/main/scala/kafka/server/AlterIsrManager.scala
new file mode 100644
index 0000000000000..16fe620dee2d7
--- /dev/null
+++ b/core/src/main/scala/kafka/server/AlterIsrManager.scala
@@ -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 kafka.server
+
+import java.util
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicLong}
+import java.util.concurrent.{ConcurrentHashMap, TimeUnit}
+
+import kafka.api.LeaderAndIsr
+import kafka.metrics.KafkaMetricsGroup
+import kafka.utils.{Logging, Scheduler}
+import org.apache.kafka.clients.ClientResponse
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.message.{AlterIsrRequestData, AlterIsrResponseData}
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.requests.{AlterIsrRequest, AlterIsrResponse}
+import org.apache.kafka.common.utils.Time
+
+import scala.collection.mutable
+import scala.collection.mutable.ListBuffer
+import scala.jdk.CollectionConverters._
+
+/**
+ * Handles the sending of AlterIsr requests to the controller. Updating the ISR is an asynchronous operation,
+ * so partitions will learn about updates through LeaderAndIsr messages sent from the controller
+ */
+trait AlterIsrManager {
+  def start(): Unit
+
+  def enqueue(alterIsrItem: AlterIsrItem): Boolean
+
+  def clearPending(topicPartition: TopicPartition): Unit
+}
+
+case class AlterIsrItem(topicPartition: TopicPartition, leaderAndIsr: LeaderAndIsr, callback: Either[Errors, LeaderAndIsr] => Unit)
+
+class AlterIsrManagerImpl(val controllerChannelManager: BrokerToControllerChannelManager,
+                          val scheduler: Scheduler,
+                          val time: Time,
+                          val brokerId: Int,
+                          val brokerEpochSupplier: () => Long) extends AlterIsrManager with Logging with KafkaMetricsGroup {
+
+  // Used to allow only one pending ISR update per partition
+  private val unsentIsrUpdates: util.Map[TopicPartition, AlterIsrItem] = new ConcurrentHashMap[TopicPartition, AlterIsrItem]()
+
+  // Used to allow only one in-flight request at a time
+  private val inflightRequest: AtomicBoolean = new AtomicBoolean(false)
+
+  private val lastIsrPropagationMs = new AtomicLong(0)
+
+  override def start(): Unit = {
+    scheduler.schedule("send-alter-isr", propagateIsrChanges, 50, 50, TimeUnit.MILLISECONDS)
+  }
+
+  override def enqueue(alterIsrItem: AlterIsrItem): Boolean = {
+    unsentIsrUpdates.putIfAbsent(alterIsrItem.topicPartition, alterIsrItem) == null
+  }
+
+  override def clearPending(topicPartition: TopicPartition): Unit = {
+    unsentIsrUpdates.remove(topicPartition)
+  }
+
+  private def propagateIsrChanges(): Unit = {
+    if (!unsentIsrUpdates.isEmpty && inflightRequest.compareAndSet(false, true)) {
+      // Copy current unsent ISRs but don't remove from the map
+      val inflightAlterIsrItems = new ListBuffer[AlterIsrItem]()
+      unsentIsrUpdates.values().forEach(item => inflightAlterIsrItems.append(item))
+
+      val now = time.milliseconds()
+      lastIsrPropagationMs.set(now)
+      sendRequest(inflightAlterIsrItems.toSeq)
+    }
+  }
+
+  private def sendRequest(inflightAlterIsrItems: Seq[AlterIsrItem]): Unit = {
+    val message = buildRequest(inflightAlterIsrItems)
+
+    def clearInflightRequests(): Unit = {
+      // Be sure to clear the in-flight flag to allow future AlterIsr requests
+      if (!inflightRequest.compareAndSet(true, false)) {
+        throw new IllegalStateException("AlterIsr response callback called when no requests were in flight")
+      }
+    }
+
+    debug(s"Sending AlterIsr to controller $message")
+    // We will not timeout AlterISR request, instead letting it retry indefinitely
+    // until a response is received, or a new LeaderAndIsr overwrites the existing isrState
+    // which causes the inflight requests to be ignored.
+    controllerChannelManager.sendRequest(new AlterIsrRequest.Builder(message),
+      new ControllerRequestCompletionHandler {
+        override def onComplete(response: ClientResponse): Unit = {
+          try {
+            val body = response.responseBody().asInstanceOf[AlterIsrResponse]
+            handleAlterIsrResponse(body, message.brokerEpoch, inflightAlterIsrItems)
+          } finally {
+            clearInflightRequests()
+          }
+        }
+
+        override def onTimeout(): Unit = {
+          throw new IllegalStateException("Encountered unexpected timeout when sending AlterIsr to the controller")
+        }
+      }, Long.MaxValue)
+  }
+
+  private def buildRequest(inflightAlterIsrItems: Seq[AlterIsrItem]): AlterIsrRequestData = {
+    val message = new AlterIsrRequestData()
+      .setBrokerId(brokerId)
+      .setBrokerEpoch(brokerEpochSupplier.apply())
+      .setTopics(new util.ArrayList())
+
+    inflightAlterIsrItems.groupBy(_.topicPartition.topic).foreach(entry => {
+      val topicPart = new AlterIsrRequestData.TopicData()
+        .setName(entry._1)
+        .setPartitions(new util.ArrayList())
+      message.topics().add(topicPart)
+      entry._2.foreach(item => {
+        topicPart.partitions().add(new AlterIsrRequestData.PartitionData()
+          .setPartitionIndex(item.topicPartition.partition)
+          .setLeaderEpoch(item.leaderAndIsr.leaderEpoch)
+          .setNewIsr(item.leaderAndIsr.isr.map(Integer.valueOf).asJava)
+          .setCurrentIsrVersion(item.leaderAndIsr.zkVersion)
+        )
+      })
+    })
+    message
+  }
+
+  def handleAlterIsrResponse(alterIsrResponse: AlterIsrResponse,
+                             sentBrokerEpoch: Long,
+                             inflightAlterIsrItems: Seq[AlterIsrItem]): Unit = {
+    val data: AlterIsrResponseData = alterIsrResponse.data
+
+    Errors.forCode(data.errorCode) match {
+      case Errors.STALE_BROKER_EPOCH =>
+        warn(s"Broker had a stale broker epoch ($sentBrokerEpoch), retrying.")
+      case Errors.CLUSTER_AUTHORIZATION_FAILED =>
+        error(s"Broker is not authorized to send AlterIsr to controller",
+          Errors.CLUSTER_AUTHORIZATION_FAILED.exception("Broker is not authorized to send AlterIsr to controller"))
+      case Errors.NONE =>
+        // Collect partition-level responses to pass to the callbacks
+        val partitionResponses: mutable.Map[TopicPartition, Either[Errors, LeaderAndIsr]] =
+          new mutable.HashMap[TopicPartition, Either[Errors, LeaderAndIsr]]()
+        data.topics.forEach { topic =>
+          topic.partitions().forEach(partition => {
+            val tp = new TopicPartition(topic.name, partition.partitionIndex)
+            val error = Errors.forCode(partition.errorCode())
+            debug(s"Controller successfully handled AlterIsr request for $tp: $partition")
+            if (error == Errors.NONE) {
+              val newLeaderAndIsr = new LeaderAndIsr(partition.leaderId, partition.leaderEpoch,
+                partition.isr.asScala.toList.map(_.toInt), partition.currentIsrVersion)
+              partitionResponses(tp) = Right(newLeaderAndIsr)
+            } else {
+              partitionResponses(tp) = Left(error)
+            }
+          })
+        }
+
+        // Iterate across the items we sent rather than what we received to ensure we run the callback even if a
+        // partition was somehow erroneously excluded from the response. Note that these callbacks are run from
+        // the leaderIsrUpdateLock write lock in Partition#sendAlterIsrRequest
+        inflightAlterIsrItems.foreach(inflightAlterIsr =>
+          if (partitionResponses.contains(inflightAlterIsr.topicPartition)) {
+            try {
+              inflightAlterIsr.callback.apply(partitionResponses(inflightAlterIsr.topicPartition))
+            } finally {
+              // Regardless of callback outcome, we need to clear from the unsent updates map to unblock further updates
+              unsentIsrUpdates.remove(inflightAlterIsr.topicPartition)
+            }
+          } else {
+            // Don't remove this partition from the update map so it will get re-sent
+            warn(s"Partition ${inflightAlterIsr.topicPartition} was sent but not included in the response")
+          }
+        )
+      case e: Errors =>
+        warn(s"Controller returned an unexpected top-level error when handling AlterIsr request: $e")
+    }
+  }
+}
diff --git a/core/src/main/scala/kafka/server/BrokerFeatures.scala b/core/src/main/scala/kafka/server/BrokerFeatures.scala
new file mode 100644
index 0000000000000..dd84f9e73e70f
--- /dev/null
+++ b/core/src/main/scala/kafka/server/BrokerFeatures.scala
@@ -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 kafka.server
+
+import kafka.utils.Logging
+import org.apache.kafka.common.feature.{Features, FinalizedVersionRange, SupportedVersionRange}
+import org.apache.kafka.common.feature.Features._
+
+import scala.jdk.CollectionConverters._
+
+/**
+ * A class that encapsulates the latest features supported by the Broker and also provides APIs to
+ * check for incompatibilities between the features supported by the Broker and finalized features.
+ * This class is immutable in production. It provides few APIs to mutate state only for the purpose
+ * of testing.
+ */
+class BrokerFeatures private (@volatile var supportedFeatures: Features[SupportedVersionRange]) {
+  // For testing only.
+  def setSupportedFeatures(newFeatures: Features[SupportedVersionRange]): Unit = {
+    supportedFeatures = newFeatures
+  }
+
+  /**
+   * Returns the default finalized features that a new Kafka cluster with IBP config >= KAFKA_2_7_IV0
+   * needs to be bootstrapped with.
+   */
+  def defaultFinalizedFeatures: Features[FinalizedVersionRange] = {
+    Features.finalizedFeatures(
+      supportedFeatures.features.asScala.map {
+        case(name, versionRange) => (
+          name, new FinalizedVersionRange(versionRange.min, versionRange.max))
+      }.asJava)
+  }
+
+  /**
+   * Returns the set of feature names found to be incompatible.
+   * A feature incompatibility is a version mismatch between the latest feature supported by the
+   * Broker, and a provided finalized feature. This can happen because a provided finalized
+   * feature:
+   *  1) Does not exist in the Broker (i.e. it is unknown to the Broker).
+   *           [OR]
+   *  2) Exists but the FinalizedVersionRange does not match with the SupportedVersionRange
+   *     of the supported feature.
+   *
+   * @param finalized   The finalized features against which incompatibilities need to be checked for.
+   *
+   * @return            The subset of input features which are incompatible. If the returned object
+   *                    is empty, it means there were no feature incompatibilities found.
+   */
+  def incompatibleFeatures(finalized: Features[FinalizedVersionRange]): Features[FinalizedVersionRange] = {
+    BrokerFeatures.incompatibleFeatures(supportedFeatures, finalized, logIncompatibilities = true)
+  }
+}
+
+object BrokerFeatures extends Logging {
+
+  def createDefault(): BrokerFeatures = {
+    // The arguments are currently empty, but, in the future as we define features we should
+    // populate the required values here.
+    new BrokerFeatures(emptySupportedFeatures)
+  }
+
+  /**
+   * Returns true if any of the provided finalized features are incompatible with the provided
+   * supported features.
+   *
+   * @param supportedFeatures   The supported features to be compared
+   * @param finalizedFeatures   The finalized features to be compared
+   *
+   * @return                    - True if there are any feature incompatibilities found.
+   *                            - False otherwise.
+   */
+  def hasIncompatibleFeatures(supportedFeatures: Features[SupportedVersionRange],
+                              finalizedFeatures: Features[FinalizedVersionRange]): Boolean = {
+    !incompatibleFeatures(supportedFeatures, finalizedFeatures, logIncompatibilities = false).empty
+  }
+
+  private def incompatibleFeatures(supportedFeatures: Features[SupportedVersionRange],
+                                   finalizedFeatures: Features[FinalizedVersionRange],
+                                   logIncompatibilities: Boolean): Features[FinalizedVersionRange] = {
+    val incompatibleFeaturesInfo = finalizedFeatures.features.asScala.map {
+      case (feature, versionLevels) =>
+        val supportedVersions = supportedFeatures.get(feature)
+        if (supportedVersions == null) {
+          (feature, versionLevels, "{feature=%s, reason='Unsupported feature'}".format(feature))
+        } else if (versionLevels.isIncompatibleWith(supportedVersions)) {
+          (feature, versionLevels, "{feature=%s, reason='%s is incompatible with %s'}".format(
+            feature, versionLevels, supportedVersions))
+        } else {
+          (feature, versionLevels, null)
+        }
+    }.filter{ case(_, _, errorReason) => errorReason != null}.toList
+
+    if (logIncompatibilities && incompatibleFeaturesInfo.nonEmpty) {
+      warn("Feature incompatibilities seen: " +
+           incompatibleFeaturesInfo.map { case(_, _, errorReason) => errorReason }.mkString(", "))
+    }
+    Features.finalizedFeatures(
+      incompatibleFeaturesInfo.map { case(feature, versionLevels, _) => (feature, versionLevels) }.toMap.asJava)
+  }
+}
diff --git a/core/src/main/scala/kafka/server/BrokerMetadataCheckpoint.scala b/core/src/main/scala/kafka/server/BrokerMetadataCheckpoint.scala
index 8ac9864122e9c..167ce2717bda0 100755
--- a/core/src/main/scala/kafka/server/BrokerMetadataCheckpoint.scala
+++ b/core/src/main/scala/kafka/server/BrokerMetadataCheckpoint.scala
@@ -18,12 +18,19 @@
 package kafka.server
 
 import java.io._
-import java.nio.file.Files
+import java.nio.file.{Files, NoSuchFileException}
 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
@@ -37,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 {
@@ -56,7 +66,7 @@ class BrokerMetadataCheckpoint(val file: File) extends Logging {
   }
 
   def read(): Option[BrokerMetadata] = {
-    Files.deleteIfExists(new File(file + ".tmp").toPath()) // try to delete any existing temp files for cleanliness
+    Files.deleteIfExists(new File(file.getPath + ".tmp").toPath()) // try to delete any existing temp files for cleanliness
 
     lock synchronized {
       try {
@@ -65,12 +75,13 @@ 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)
         }
       } catch {
-        case _: FileNotFoundException =>
+        case _: NoSuchFileException =>
           warn("No meta.properties file under dir %s".format(file.getAbsolutePath()))
           None
         case e1: Exception =>
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/BrokerToControllerChannelManagerImpl.scala b/core/src/main/scala/kafka/server/BrokerToControllerChannelManagerImpl.scala
new file mode 100644
index 0000000000000..c0918ad7d03e8
--- /dev/null
+++ b/core/src/main/scala/kafka/server/BrokerToControllerChannelManagerImpl.scala
@@ -0,0 +1,228 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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.{LinkedBlockingDeque, TimeUnit}
+
+import kafka.common.{InterBrokerSendThread, RequestAndCompletionHandler}
+import kafka.utils.Logging
+import org.apache.kafka.clients._
+import org.apache.kafka.common.Node
+import org.apache.kafka.common.metrics.Metrics
+import org.apache.kafka.common.network._
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.requests.AbstractRequest
+import org.apache.kafka.common.security.JaasContext
+import org.apache.kafka.common.utils.{LogContext, Time}
+
+import scala.collection.mutable
+import scala.jdk.CollectionConverters._
+
+trait BrokerToControllerChannelManager {
+
+  /**
+   * Send request to the controller.
+   *
+   * @param request         The request to be sent.
+   * @param callback        Request completion callback.
+   * @param retryDeadlineMs The retry deadline which will only be checked after receiving a response.
+   *                        This means that in the worst case, the total timeout would be twice of
+   *                        the configured timeout.
+   */
+  def sendRequest(request: AbstractRequest.Builder[_ <: AbstractRequest],
+                  callback: ControllerRequestCompletionHandler,
+                  retryDeadlineMs: Long): Unit
+
+  def start(): Unit
+
+  def shutdown(): Unit
+}
+
+/**
+ * This class manages the connection between a broker and the controller. It runs a single
+ * {@link BrokerToControllerRequestThread} which uses the broker's metadata cache as its own metadata to find
+ * and connect to the controller. The channel is async and runs the network connection in the background.
+ * The maximum number of in-flight requests are set to one to ensure orderly response from the controller, therefore
+ * care must be taken to not block on outstanding requests for too long.
+ */
+class BrokerToControllerChannelManagerImpl(metadataCache: kafka.server.MetadataCache,
+                                           time: Time,
+                                           metrics: Metrics,
+                                           config: KafkaConfig,
+                                           channelName: String,
+                                           threadNamePrefix: Option[String] = None) extends BrokerToControllerChannelManager with Logging {
+  private val requestQueue = new LinkedBlockingDeque[BrokerToControllerQueueItem]
+  private val logContext = new LogContext(s"[broker-${config.brokerId}-to-controller] ")
+  private val manualMetadataUpdater = new ManualMetadataUpdater()
+  private val requestThread = newRequestThread
+
+  override def start(): Unit = {
+    requestThread.start()
+  }
+
+  override def shutdown(): Unit = {
+    requestThread.shutdown()
+    requestThread.awaitShutdown()
+    info(s"Broker to controller channel manager for $channelName shutdown")
+  }
+
+  private[server] def newRequestThread = {
+    val brokerToControllerListenerName = config.controlPlaneListenerName.getOrElse(config.interBrokerListenerName)
+    val brokerToControllerSecurityProtocol = config.controlPlaneSecurityProtocol.getOrElse(config.interBrokerSecurityProtocol)
+
+    val networkClient = {
+      val channelBuilder = ChannelBuilders.clientChannelBuilder(
+        brokerToControllerSecurityProtocol,
+        JaasContext.Type.SERVER,
+        config,
+        brokerToControllerListenerName,
+        config.saslMechanismInterBrokerProtocol,
+        time,
+        config.saslInterBrokerHandshakeRequestEnable,
+        logContext
+      )
+      val selector = new Selector(
+        NetworkReceive.UNLIMITED,
+        Selector.NO_IDLE_TIMEOUT_MS,
+        metrics,
+        time,
+        channelName,
+        Map("BrokerId" -> config.brokerId.toString).asJava,
+        false,
+        channelBuilder,
+        logContext
+      )
+      new NetworkClient(
+        selector,
+        manualMetadataUpdater,
+        config.brokerId.toString,
+        1,
+        50,
+        50,
+        Selectable.USE_DEFAULT_BUFFER_SIZE,
+        Selectable.USE_DEFAULT_BUFFER_SIZE,
+        config.requestTimeoutMs,
+        config.connectionSetupTimeoutMs,
+        config.connectionSetupTimeoutMaxMs,
+        ClientDnsLookup.USE_ALL_DNS_IPS,
+        time,
+        false,
+        new ApiVersions,
+        logContext
+      )
+    }
+    val threadName = threadNamePrefix match {
+      case None => s"broker-${config.brokerId}-to-controller-send-thread"
+      case Some(name) => s"$name:broker-${config.brokerId}-to-controller-send-thread"
+    }
+
+    new BrokerToControllerRequestThread(networkClient, manualMetadataUpdater, requestQueue, metadataCache, config,
+      brokerToControllerListenerName, time, threadName)
+  }
+
+  override def sendRequest(request: AbstractRequest.Builder[_ <: AbstractRequest],
+                           callback: ControllerRequestCompletionHandler,
+                           retryDeadlineMs: Long): Unit = {
+    requestQueue.put(BrokerToControllerQueueItem(request, callback, retryDeadlineMs))
+    requestThread.wakeup()
+  }
+}
+
+abstract class ControllerRequestCompletionHandler extends RequestCompletionHandler {
+
+  /**
+   * Fire when the request transmission time passes the caller defined deadline on the channel queue.
+   * It covers the total waiting time including retries which might be the result of individual request timeout.
+   */
+  def onTimeout(): Unit
+}
+
+case class BrokerToControllerQueueItem(request: AbstractRequest.Builder[_ <: AbstractRequest],
+                                       callback: ControllerRequestCompletionHandler,
+                                       deadlineMs: Long)
+
+class BrokerToControllerRequestThread(networkClient: KafkaClient,
+                                      metadataUpdater: ManualMetadataUpdater,
+                                      requestQueue: LinkedBlockingDeque[BrokerToControllerQueueItem],
+                                      metadataCache: kafka.server.MetadataCache,
+                                      config: KafkaConfig,
+                                      listenerName: ListenerName,
+                                      time: Time,
+                                      threadName: String)
+  extends InterBrokerSendThread(threadName, networkClient, time, isInterruptible = false) {
+
+  private var activeController: Option[Node] = None
+
+  override def requestTimeoutMs: Int = config.controllerSocketTimeoutMs
+
+  override def generateRequests(): Iterable[RequestAndCompletionHandler] = {
+    val requestsToSend = new mutable.Queue[RequestAndCompletionHandler]
+    val topRequest = requestQueue.poll()
+    if (topRequest != null) {
+      val request = RequestAndCompletionHandler(
+        activeController.get,
+        topRequest.request,
+        handleResponse(topRequest)
+      )
+
+      requestsToSend.enqueue(request)
+    }
+    requestsToSend
+  }
+
+  private[server] def handleResponse(request: BrokerToControllerQueueItem)(response: ClientResponse): Unit = {
+    if (hasTimedOut(request, response)) {
+      request.callback.onTimeout()
+    } else if (response.wasDisconnected()) {
+      activeController = None
+      requestQueue.putFirst(request)
+    } else if (response.responseBody().errorCounts().containsKey(Errors.NOT_CONTROLLER)) {
+      // just close the controller connection and wait for metadata cache update in doWork
+      networkClient.close(activeController.get.idString)
+      activeController = None
+      requestQueue.putFirst(request)
+    } else {
+      request.callback.onComplete(response)
+    }
+  }
+
+  private def hasTimedOut(request: BrokerToControllerQueueItem, response: ClientResponse): Boolean = {
+    response.receivedTimeMs() > request.deadlineMs
+  }
+
+  private[server] def backoff(): Unit = pause(100, TimeUnit.MILLISECONDS)
+
+  override def doWork(): Unit = {
+    if (activeController.isDefined) {
+      super.doWork()
+    } else {
+      debug("Controller isn't cached, looking for local metadata changes")
+      val controllerOpt = metadataCache.getControllerId.flatMap(metadataCache.getAliveBroker)
+      if (controllerOpt.isDefined) {
+        if (activeController.isEmpty || activeController.exists(_.id != controllerOpt.get.id))
+          info(s"Recorded new controller, from now on will use broker ${controllerOpt.get.id}")
+        activeController = Option(controllerOpt.get.node(listenerName))
+        metadataUpdater.setNodes(metadataCache.getAliveBrokers.map(_.node(listenerName)).asJava)
+      } else {
+        // need to backoff to avoid tight loops
+        debug("No controller defined in metadata cache, retrying after backoff")
+        backoff()
+      }
+    }
+  }
+}
diff --git a/core/src/main/scala/kafka/server/ClientQuotaManager.scala b/core/src/main/scala/kafka/server/ClientQuotaManager.scala
index 8ec27a3a4c03b..e32978cf1710e 100644
--- a/core/src/main/scala/kafka/server/ClientQuotaManager.scala
+++ b/core/src/main/scala/kafka/server/ClientQuotaManager.scala
@@ -16,54 +16,52 @@
  */
 package kafka.server
 
+import java.{lang, util}
 import java.util.concurrent.{ConcurrentHashMap, DelayQueue, TimeUnit}
 import java.util.concurrent.locks.ReentrantReadWriteLock
 
-import kafka.utils.{Logging, ShutdownableThread}
-import org.apache.kafka.common.MetricName
+import kafka.network.RequestChannel
+import kafka.network.RequestChannel._
+import kafka.server.ClientQuotaManager._
+import kafka.utils.{Logging, QuotaUtils, ShutdownableThread}
+import org.apache.kafka.common.{Cluster, MetricName}
 import org.apache.kafka.common.metrics._
-import org.apache.kafka.common.metrics.stats.{Avg, Rate, Total}
+import org.apache.kafka.common.metrics.Metrics
+import org.apache.kafka.common.metrics.stats.{Avg, CumulativeSum, Rate}
+import org.apache.kafka.common.security.auth.KafkaPrincipal
 import org.apache.kafka.common.utils.{Sanitizer, Time}
+import org.apache.kafka.server.quota.{ClientQuotaCallback, ClientQuotaEntity, ClientQuotaType}
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 
 /**
  * Represents the sensors aggregated per client
- * @param quotaEntity Quota entity representing ,  or 
+ * @param metricTags Quota metric tags for the client
  * @param quotaSensor @Sensor that tracks the quota
  * @param throttleTimeSensor @Sensor that tracks the throttle time
  */
-case class ClientSensors(quotaEntity: QuotaEntity, quotaSensor: Sensor, throttleTimeSensor: Sensor)
+case class ClientSensors(metricTags: Map[String, String], quotaSensor: Sensor, throttleTimeSensor: Sensor)
 
 /**
  * Configuration settings for quota management
- * @param quotaBytesPerSecondDefault The default bytes per second quota allocated to any client-id if
+ * @param quotaDefault The default allocated to any client-id if
  *        dynamic defaults or user quotas are not set
  * @param numQuotaSamples The number of samples to retain in memory
  * @param quotaWindowSizeSeconds The time span of each sample
  *
  */
-case class ClientQuotaManagerConfig(quotaBytesPerSecondDefault: Long =
-                                        ClientQuotaManagerConfig.QuotaBytesPerSecondDefault,
+case class ClientQuotaManagerConfig(quotaDefault: Long =
+                                        ClientQuotaManagerConfig.QuotaDefault,
                                     numQuotaSamples: Int =
                                         ClientQuotaManagerConfig.DefaultNumQuotaSamples,
                                     quotaWindowSizeSeconds: Int =
                                         ClientQuotaManagerConfig.DefaultQuotaWindowSizeSeconds)
 
 object ClientQuotaManagerConfig {
-  val QuotaBytesPerSecondDefault = Long.MaxValue
+  val QuotaDefault = Long.MaxValue
   // Always have 10 whole windows + 1 current window
   val DefaultNumQuotaSamples = 11
   val DefaultQuotaWindowSizeSeconds = 1
-  // Purge sensors after 1 hour of inactivity
-  val InactiveSensorExpirationTimeSeconds  = 3600
-  val QuotaRequestPercentDefault = Int.MaxValue.toDouble
-  val NanosToPercentagePerSecond = 100.0 / TimeUnit.SECONDS.toNanos(1)
-
-  val UnlimitedQuota = Quota.upperBound(Long.MaxValue)
-  val DefaultClientIdQuotaId = QuotaId(None, Some(ConfigEntityName.Default), Some(ConfigEntityName.Default))
-  val DefaultUserQuotaId = QuotaId(Some(ConfigEntityName.Default), None, None)
-  val DefaultUserClientIdQuotaId = QuotaId(Some(ConfigEntityName.Default), Some(ConfigEntityName.Default), Some(ConfigEntityName.Default))
 }
 
 object QuotaTypes {
@@ -71,11 +69,93 @@ object QuotaTypes {
   val ClientIdQuotaEnabled = 1
   val UserQuotaEnabled = 2
   val UserClientIdQuotaEnabled = 4
+  val CustomQuotas = 8 // No metric update optimizations are used with custom quotas
 }
 
-case class QuotaId(sanitizedUser: Option[String], clientId: Option[String], sanitizedClientId: Option[String])
+object ClientQuotaManager {
+  // Purge sensors after 1 hour of inactivity
+  val InactiveSensorExpirationTimeSeconds = 3600
+
+  val DefaultClientIdQuotaEntity = KafkaQuotaEntity(None, Some(DefaultClientIdEntity))
+  val DefaultUserQuotaEntity = KafkaQuotaEntity(Some(DefaultUserEntity), None)
+  val DefaultUserClientIdQuotaEntity = KafkaQuotaEntity(Some(DefaultUserEntity), Some(DefaultClientIdEntity))
+
+  sealed trait BaseUserEntity extends ClientQuotaEntity.ConfigEntity
+
+  case class UserEntity(sanitizedUser: String) extends BaseUserEntity {
+    override def entityType: ClientQuotaEntity.ConfigEntityType = ClientQuotaEntity.ConfigEntityType.USER
+    override def name: String = Sanitizer.desanitize(sanitizedUser)
+    override def toString: String = s"user $sanitizedUser"
+  }
+
+  case class ClientIdEntity(clientId: String) extends ClientQuotaEntity.ConfigEntity {
+    override def entityType: ClientQuotaEntity.ConfigEntityType = ClientQuotaEntity.ConfigEntityType.CLIENT_ID
+    override def name: String = clientId
+    override def toString: String = s"client-id $clientId"
+  }
+
+  case object DefaultUserEntity extends BaseUserEntity {
+    override def entityType: ClientQuotaEntity.ConfigEntityType = ClientQuotaEntity.ConfigEntityType.DEFAULT_USER
+    override def name: String = ConfigEntityName.Default
+    override def toString: String = "default user"
+  }
+
+  case object DefaultClientIdEntity extends ClientQuotaEntity.ConfigEntity {
+    override def entityType: ClientQuotaEntity.ConfigEntityType = ClientQuotaEntity.ConfigEntityType.DEFAULT_CLIENT_ID
+    override def name: String = ConfigEntityName.Default
+    override def toString: String = "default client-id"
+  }
+
+  case class KafkaQuotaEntity(userEntity: Option[BaseUserEntity],
+                              clientIdEntity: Option[ClientQuotaEntity.ConfigEntity]) extends ClientQuotaEntity {
+    override def configEntities: util.List[ClientQuotaEntity.ConfigEntity] =
+      (userEntity.toList ++ clientIdEntity.toList).asJava
+
+    def sanitizedUser: String = userEntity.map {
+      case entity: UserEntity => entity.sanitizedUser
+      case DefaultUserEntity => ConfigEntityName.Default
+    }.getOrElse("")
 
-case class QuotaEntity(quotaId: QuotaId, sanitizedUser: String, clientId: String, sanitizedClientId: String, quota: Quota)
+    def clientId: String = clientIdEntity.map(_.name).getOrElse("")
+
+    override def toString: String = {
+      val user = userEntity.map(_.toString).getOrElse("")
+      val clientId = clientIdEntity.map(_.toString).getOrElse("")
+      s"$user $clientId".trim
+    }
+  }
+
+  object DefaultTags {
+    val User = "user"
+    val ClientId = "client-id"
+  }
+
+  /**
+   * This calculates the amount of time needed to bring the metric within quota
+   * assuming that no new metrics are recorded.
+   *
+   * Basically, if O is the observed rate and T is the target rate over a window of W, to bring O down to T,
+   * we need to add a delay of X to W such that O * W / (W + X) = T.
+   * Solving for X, we get X = (O - T)/T * W.
+   */
+  def throttleTime(e: QuotaViolationException, timeMs: Long): Long = {
+    val difference = e.value - e.bound
+    // Use the precise window used by the rate calculation
+    val throttleTimeMs = difference / e.bound * windowSize(e.metric, timeMs)
+    Math.round(throttleTimeMs)
+  }
+
+  private def windowSize(metric: KafkaMetric, timeMs: Long): Long =
+    measurableAsRate(metric.metricName, metric.measurable).windowSize(metric.config, timeMs)
+
+  // Casting to Rate because we only use Rate in Quota computation
+  private def measurableAsRate(name: MetricName, measurable: Measurable): Rate = {
+    measurable match {
+      case r: Rate => r
+      case _ => throw new IllegalArgumentException(s"Metric $name is not a Rate metric, value $measurable")
+    }
+  }
+}
 
 /**
  * Helper class that records per-client metrics. It is also responsible for maintaining Quota usage statistics
@@ -102,45 +182,55 @@ case class QuotaEntity(quotaId: QuotaId, sanitizedUser: String, clientId: String
  * @param metrics @Metrics Metrics instance
  * @param quotaType Quota type of this quota manager
  * @param time @Time object to use
+ * @param threadNamePrefix The thread prefix to use
+ * @param clientQuotaCallback An optional @ClientQuotaCallback
  */
 class ClientQuotaManager(private val config: ClientQuotaManagerConfig,
                          private val metrics: Metrics,
                          private val quotaType: QuotaType,
                          private val time: Time,
-                         threadNamePrefix: String) extends Logging {
-  private val overriddenQuota = new ConcurrentHashMap[QuotaId, Quota]()
-  private val staticConfigClientIdQuota = Quota.upperBound(config.quotaBytesPerSecondDefault)
-  @volatile private var quotaTypesEnabled =
-    if (config.quotaBytesPerSecondDefault == Long.MaxValue) QuotaTypes.NoQuotas
-    else QuotaTypes.ClientIdQuotaEnabled
+                         private val threadNamePrefix: String,
+                         private val clientQuotaCallback: Option[ClientQuotaCallback] = None) extends Logging {
+
   private val lock = new ReentrantReadWriteLock()
-  private val delayQueue = new DelayQueue[ThrottledResponse]()
   private val sensorAccessor = new SensorAccess(lock, metrics)
-  private[server] val throttledRequestReaper = new ThrottledRequestReaper(delayQueue, threadNamePrefix)
+  private val quotaCallback = clientQuotaCallback.getOrElse(new DefaultQuotaCallback)
+  private val staticConfigClientIdQuota = Quota.upperBound(config.quotaDefault.toDouble)
+  private val clientQuotaType = QuotaType.toClientQuotaType(quotaType)
+
+  @volatile
+  private var quotaTypesEnabled = clientQuotaCallback match {
+    case Some(_) => QuotaTypes.CustomQuotas
+    case None =>
+      if (config.quotaDefault == Long.MaxValue) QuotaTypes.NoQuotas
+      else QuotaTypes.ClientIdQuotaEnabled
+  }
+
+  private val delayQueueSensor = metrics.sensor(quotaType.toString + "-delayQueue")
+  delayQueueSensor.add(metrics.metricName("queue-size", quotaType.toString,
+    "Tracks the size of the delay queue"), new CumulativeSum())
 
-  private val delayQueueSensor = metrics.sensor(quotaType + "-delayQueue")
-  delayQueueSensor.add(metrics.metricName("queue-size",
-                                      quotaType.toString,
-                                      "Tracks the size of the delay queue"), new Total())
-  start() // Use start method to keep findbugs happy
-  private def start() {
-    throttledRequestReaper.start()
+  private val delayQueue = new DelayQueue[ThrottledChannel]()
+  private[server] val throttledChannelReaper = new ThrottledChannelReaper(delayQueue, threadNamePrefix)
+  start() // Use start method to keep spotbugs happy
+  private def start(): Unit = {
+    throttledChannelReaper.start()
   }
 
   /**
-   * Reaper thread that triggers callbacks on all throttled requests
+   * Reaper thread that triggers channel unmute callbacks on all throttled channels
    * @param delayQueue DelayQueue to dequeue from
    */
-  class ThrottledRequestReaper(delayQueue: DelayQueue[ThrottledResponse], prefix: String) extends ShutdownableThread(
-    s"${prefix}ThrottledRequestReaper-${quotaType}", false) {
+  class ThrottledChannelReaper(delayQueue: DelayQueue[ThrottledChannel], prefix: String) extends ShutdownableThread(
+    s"${prefix}ThrottledChannelReaper-$quotaType", false) {
 
     override def doWork(): Unit = {
-      val response: ThrottledResponse = delayQueue.poll(1, TimeUnit.SECONDS)
-      if (response != null) {
+      val throttledChannel: ThrottledChannel = delayQueue.poll(1, TimeUnit.SECONDS)
+      if (throttledChannel != null) {
         // Decrement the size of the delay queue
         delayQueueSensor.record(-1)
-        trace("Response throttled for: " + response.throttleTimeMs + " ms")
-        response.execute()
+        // Notify the socket server that throttling is done for this channel, so that it can try to unmute the channel.
+        throttledChannel.notifyThrottlingDone()
       }
     }
   }
@@ -155,48 +245,46 @@ class ClientQuotaManager(private val config: ClientQuotaManagerConfig,
   def quotasEnabled: Boolean = quotaTypesEnabled != QuotaTypes.NoQuotas
 
   /**
-   * Records that a user/clientId changed some metric being throttled (produced/consumed bytes, request processing time etc.)
-   * If quota has been violated, callback is invoked after a delay, otherwise the callback is invoked immediately.
-   * Throttle time calculation may be overridden by sub-classes.
-   * @param sanitizedUser user principal of client
-   * @param clientId clientId that produced/fetched the data
-   * @param value amount of data in bytes or request processing time as a percentage
-   * @param callback Callback function. This will be triggered immediately if quota is not violated.
-   *                 If there is a quota violation, this callback will be triggered after a delay
-   * @return Number of milliseconds to delay the response in case of Quota violation.
-   *         Zero otherwise
+   * See {recordAndGetThrottleTimeMs}.
    */
-  def maybeRecordAndThrottle(sanitizedUser: String, clientId: String, value: Double, callback: Int => Unit): Int = {
+  def maybeRecordAndGetThrottleTimeMs(request: RequestChannel.Request, value: Double, timeMs: Long): Int = {
+    maybeRecordAndGetThrottleTimeMs(request.session, request.header.clientId, value, timeMs)
+  }
+
+  /**
+   * See {recordAndGetThrottleTimeMs}.
+   */
+  def maybeRecordAndGetThrottleTimeMs(session: Session, clientId: String, value: Double, timeMs: Long): Int = {
+    // Record metrics only if quotas are enabled.
     if (quotasEnabled) {
-      val clientSensors = getOrCreateQuotaSensors(sanitizedUser, clientId)
-      recordAndThrottleOnQuotaViolation(clientSensors, value, callback)
+      recordAndGetThrottleTimeMs(session, clientId, value, timeMs)
     } else {
-      // Don't record any metrics if quotas are not enabled at any level
-      val throttleTimeMs = 0
-      callback(throttleTimeMs)
-      throttleTimeMs
+      0
     }
   }
 
-  def recordAndThrottleOnQuotaViolation(clientSensors: ClientSensors, value: Double, callback: Int => Unit): Int = {
-    var throttleTimeMs = 0
+  /**
+   * Records that a user/clientId accumulated or would like to accumulate the provided amount at the
+   * the specified time, returns throttle time in milliseconds.
+   *
+   * @param session The session from which the user is extracted
+   * @param clientId The client id
+   * @param value The value to accumulate
+   * @param timeMs The time at which to accumulate the value
+   * @return The throttle time in milliseconds defines as the time to wait until the average
+   *         rate gets back to the defined quota
+   */
+  def recordAndGetThrottleTimeMs(session: Session, clientId: String, value: Double, timeMs: Long): Int = {
+    val clientSensors = getOrCreateQuotaSensors(session, clientId)
     try {
-      clientSensors.quotaSensor.record(value)
-      // trigger the callback immediately if quota is not violated
-      callback(0)
+      clientSensors.quotaSensor.record(value, timeMs, true)
+      0
     } catch {
-      case _: QuotaViolationException =>
-        // Compute the delay
-        val clientQuotaEntity = clientSensors.quotaEntity
-        val clientMetric = metrics.metrics().get(clientRateMetricName(clientQuotaEntity.sanitizedUser, clientQuotaEntity.clientId))
-        throttleTimeMs = throttleTime(clientMetric, getQuotaMetricConfig(clientQuotaEntity.quota)).toInt
-        clientSensors.throttleTimeSensor.record(throttleTimeMs)
-        // If delayed, add the element to the delayQueue
-        delayQueue.add(new ThrottledResponse(time, throttleTimeMs, callback))
-        delayQueueSensor.record()
-        debug("Quota violated for sensor (%s). Delay time: (%d)".format(clientSensors.quotaSensor.name(), throttleTimeMs))
+      case e: QuotaViolationException =>
+        val throttleTimeMs = throttleTime(e, timeMs).toInt
+        debug(s"Quota violated for sensor (${clientSensors.quotaSensor.name}). Delay time: ($throttleTimeMs)")
+        throttleTimeMs
     }
-    throttleTimeMs
   }
 
   /**
@@ -204,214 +292,169 @@ 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(session: Session, clientId: String, value: Double): Unit = {
+    val clientSensors = getOrCreateQuotaSensors(session, clientId)
     clientSensors.quotaSensor.record(value, time.milliseconds(), false)
   }
 
   /**
-   * Determines the quota-id for the client with the specified user principal
-   * and client-id and returns the quota entity that encapsulates the quota-id
-   * and the associated quota override or default quota.
+   * "Unrecord" the given value that has already been recorded for the given user/client by recording a negative value
+   * of the same quantity.
    *
+   * For a throttled fetch, the broker should return an empty response and thus should not record the value. Ideally,
+   * we would like to compute the throttle time before actually recording the value, but the current Sensor code
+   * couples value recording and quota checking very tightly. As a workaround, we will unrecord the value for the fetch
+   * in case of throttling. Rate keeps the sum of values that fall in each time window, so this should bring the
+   * overall sum back to the previous value.
    */
-  private def quotaEntity(sanitizedUser: String, clientId: String, sanitizedClientId: String) : QuotaEntity = {
-    quotaTypesEnabled match {
-      case QuotaTypes.NoQuotas | QuotaTypes.ClientIdQuotaEnabled =>
-        val quotaId = QuotaId(None, Some(clientId), Some(sanitizedClientId))
-        var quota = overriddenQuota.get(quotaId)
-        if (quota == null) {
-          quota = overriddenQuota.get(ClientQuotaManagerConfig.DefaultClientIdQuotaId)
-          if (quota == null)
-            quota = staticConfigClientIdQuota
-        }
-        QuotaEntity(quotaId, "", clientId, sanitizedClientId, quota)
-      case QuotaTypes.UserQuotaEnabled =>
-        val quotaId = QuotaId(Some(sanitizedUser), None, None)
-        var quota = overriddenQuota.get(quotaId)
-        if (quota == null) {
-          quota = overriddenQuota.get(ClientQuotaManagerConfig.DefaultUserQuotaId)
-          if (quota == null)
-            quota = ClientQuotaManagerConfig.UnlimitedQuota
-        }
-        QuotaEntity(quotaId, sanitizedUser, "", "", quota)
-      case QuotaTypes.UserClientIdQuotaEnabled =>
-        val quotaId = QuotaId(Some(sanitizedUser), Some(clientId), Some(sanitizedClientId))
-        var quota = overriddenQuota.get(quotaId)
-        if (quota == null) {
-          quota = overriddenQuota.get(QuotaId(Some(sanitizedUser), Some(ConfigEntityName.Default), Some(ConfigEntityName.Default)))
-          if (quota == null) {
-            quota = overriddenQuota.get(QuotaId(Some(ConfigEntityName.Default), Some(clientId), Some(sanitizedClientId)))
-            if (quota == null) {
-              quota = overriddenQuota.get(ClientQuotaManagerConfig.DefaultUserClientIdQuotaId)
-              if (quota == null)
-                quota = ClientQuotaManagerConfig.UnlimitedQuota
-            }
-          }
-        }
-        QuotaEntity(quotaId, sanitizedUser, clientId, sanitizedClientId, quota)
-      case _ =>
-        quotaEntityWithMultipleQuotaLevels(sanitizedUser, clientId, sanitizedClientId)
-    }
+  def unrecordQuotaSensor(request: RequestChannel.Request, value: Double, timeMs: Long): Unit = {
+    val clientSensors = getOrCreateQuotaSensors(request.session, request.header.clientId)
+    clientSensors.quotaSensor.record(value * (-1), timeMs, false)
   }
 
-  private def quotaEntityWithMultipleQuotaLevels(sanitizedUser: String, clientId: String, sanitizerClientId: String) : QuotaEntity = {
-    val userClientQuotaId = QuotaId(Some(sanitizedUser), Some(clientId), Some(sanitizerClientId))
-
-    val userQuotaId = QuotaId(Some(sanitizedUser), None, None)
-    val clientQuotaId = QuotaId(None, Some(clientId), Some(sanitizerClientId))
-    var quotaId = userClientQuotaId
-    var quotaConfigId = userClientQuotaId
-    // 1) /config/users//clients/
-    var quota = overriddenQuota.get(quotaConfigId)
-    if (quota == null) {
-      // 2) /config/users//clients/
-      quotaId = userClientQuotaId
-      quotaConfigId = QuotaId(Some(sanitizedUser), Some(ConfigEntityName.Default), Some(ConfigEntityName.Default))
-      quota = overriddenQuota.get(quotaConfigId)
-
-      if (quota == null) {
-        // 3) /config/users/
-        quotaId = userQuotaId
-        quotaConfigId = quotaId
-        quota = overriddenQuota.get(quotaConfigId)
-
-        if (quota == null) {
-          // 4) /config/users//clients/
-          quotaId = userClientQuotaId
-          quotaConfigId = QuotaId(Some(ConfigEntityName.Default), Some(clientId), Some(sanitizerClientId))
-          quota = overriddenQuota.get(quotaConfigId)
+  /**
+   * Returns maximum value 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
+    }
+  }
 
-          if (quota == null) {
-            // 5) /config/users//clients/
-            quotaId = userClientQuotaId
-            quotaConfigId = QuotaId(Some(ConfigEntityName.Default), Some(ConfigEntityName.Default), Some(ConfigEntityName.Default))
-            quota = overriddenQuota.get(quotaConfigId)
-
-            if (quota == null) {
-              // 6) /config/users/
-              quotaId = userQuotaId
-              quotaConfigId = QuotaId(Some(ConfigEntityName.Default), None, None)
-              quota = overriddenQuota.get(quotaConfigId)
-
-              if (quota == null) {
-                // 7) /config/clients/
-                quotaId = clientQuotaId
-                quotaConfigId = QuotaId(None, Some(clientId), Some(sanitizerClientId))
-                quota = overriddenQuota.get(quotaConfigId)
-
-                if (quota == null) {
-                  // 8) /config/clients/
-                  quotaId = clientQuotaId
-                  quotaConfigId = QuotaId(None, Some(ConfigEntityName.Default), Some(ConfigEntityName.Default))
-                  quota = overriddenQuota.get(quotaConfigId)
-
-                  if (quota == null) {
-                    quotaId = clientQuotaId
-                    quotaConfigId = null
-                    quota = staticConfigClientIdQuota
-                  }
-                }
-              }
-            }
-          }
-        }
-      }
+  /**
+   * Throttle a client by muting the associated channel for the given throttle time.
+   *
+   * @param request client request
+   * @param throttleTimeMs Duration in milliseconds for which the channel is to be muted.
+   * @param channelThrottlingCallback Callback for channel throttling
+   */
+  def throttle(request: RequestChannel.Request, throttleTimeMs: Int, channelThrottlingCallback: Response => Unit): Unit = {
+    if (throttleTimeMs > 0) {
+      val clientSensors = getOrCreateQuotaSensors(request.session, request.headerForLoggingOrThrottling().clientId)
+      clientSensors.throttleTimeSensor.record(throttleTimeMs)
+      val throttledChannel = new ThrottledChannel(request, time, throttleTimeMs, channelThrottlingCallback)
+      delayQueue.add(throttledChannel)
+      delayQueueSensor.record()
+      debug("Channel throttled for sensor (%s). Delay time: (%d)".format(clientSensors.quotaSensor.name(), throttleTimeMs))
     }
-    val quotaUser = if (quotaId == clientQuotaId) "" else sanitizedUser
-    val quotaClientId = if (quotaId == userQuotaId) "" else clientId
-    QuotaEntity(quotaId, quotaUser, quotaClientId, sanitizerClientId, quota)
   }
 
   /**
    * Returns the quota for the client with the specified (non-encoded) user principal and client-id.
-   * 
+   *
    * Note: this method is expensive, it is meant to be used by tests only
    */
-  def quota(user: String, clientId: String) = {
-    quotaEntity(Sanitizer.sanitize(user), clientId, Sanitizer.sanitize(clientId)).quota
+  def quota(user: String, clientId: String): Quota = {
+    val userPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, user)
+    quota(userPrincipal, clientId)
   }
 
-  /*
-   * This calculates the amount of time needed to bring the metric within quota
-   * assuming that no new metrics are recorded.
+  /**
+   * Returns the quota for the client with the specified user principal and client-id.
    *
-   * Basically, if O is the observed rate and T is the target rate over a window of W, to bring O down to T,
-   * we need to add a delay of X to W such that O * W / (W + X) = T.
-   * Solving for X, we get X = (O - T)/T * W.
+   * Note: this method is expensive, it is meant to be used by tests only
    */
-  protected def throttleTime(clientMetric: KafkaMetric, config: MetricConfig): Long = {
-    val rateMetric: Rate = measurableAsRate(clientMetric.metricName(), clientMetric.measurable())
-    val quota = config.quota()
-    val difference = clientMetric.value() - quota.bound
-    // Use the precise window used by the rate calculation
-    val throttleTimeMs = difference / quota.bound * rateMetric.windowSize(config, time.milliseconds())
-    throttleTimeMs.round
+  def quota(userPrincipal: KafkaPrincipal, clientId: String): Quota = {
+    val metricTags = quotaCallback.quotaMetricTags(clientQuotaType, userPrincipal, clientId)
+    Quota.upperBound(quotaLimit(metricTags))
   }
 
-  // Casting to Rate because we only use Rate in Quota computation
-  private def measurableAsRate(name: MetricName, measurable: Measurable): Rate = {
-    measurable match {
-      case r: Rate => r
-      case _ => throw new IllegalArgumentException(s"Metric $name is not a Rate metric, value $measurable")
-    }
+  private def quotaLimit(metricTags: util.Map[String, String]): Double = {
+    Option(quotaCallback.quotaLimit(clientQuotaType, metricTags)).map(_.toDouble).getOrElse(Long.MaxValue)
   }
 
-  /*
+  /**
+   * This calculates the amount of time needed to bring the metric within quota
+   * assuming that no new metrics are recorded.
+   *
+   * See {QuotaUtils.throttleTime} for the details.
+   */
+  protected def throttleTime(e: QuotaViolationException, timeMs: Long): Long = {
+    QuotaUtils.throttleTime(e, timeMs)
+  }
+
+  /**
    * This function either returns the sensors for a given client id or creates them if they don't exist
    * First sensor of the tuple is the quota enforcement sensor. Second one is the throttle time sensor
    */
-  def getOrCreateQuotaSensors(sanitizedUser: String, clientId: String): ClientSensors = {
-    val sanitizedClientId = Sanitizer.sanitize(clientId)
-    val clientQuotaEntity = quotaEntity(sanitizedUser, clientId, sanitizedClientId)
+  def getOrCreateQuotaSensors(session: Session, clientId: String): ClientSensors = {
+    // Use cached sanitized principal if using default callback
+    val metricTags = quotaCallback match {
+      case callback: DefaultQuotaCallback => callback.quotaMetricTags(session.sanitizedUser, clientId)
+      case _ => quotaCallback.quotaMetricTags(clientQuotaType, session.principal, clientId).asScala.toMap
+    }
     // Names of the sensors to access
-    ClientSensors(
-      clientQuotaEntity,
+    val sensors = ClientSensors(
+      metricTags,
       sensorAccessor.getOrCreate(
-        getQuotaSensorName(clientQuotaEntity.quotaId),
-        ClientQuotaManagerConfig.InactiveSensorExpirationTimeSeconds,
-        clientRateMetricName(clientQuotaEntity.sanitizedUser, clientQuotaEntity.clientId),
-        Some(getQuotaMetricConfig(clientQuotaEntity.quota)),
-        new Rate
+        getQuotaSensorName(metricTags),
+        ClientQuotaManager.InactiveSensorExpirationTimeSeconds,
+        registerQuotaMetrics(metricTags)
       ),
-      sensorAccessor.getOrCreate(getThrottleTimeSensorName(clientQuotaEntity.quotaId),
-        ClientQuotaManagerConfig.InactiveSensorExpirationTimeSeconds,
-        throttleMetricName(clientQuotaEntity),
-        None,
-        new Avg
+      sensorAccessor.getOrCreate(
+        getThrottleTimeSensorName(metricTags),
+        ClientQuotaManager.InactiveSensorExpirationTimeSeconds,
+        sensor => sensor.add(throttleMetricName(metricTags), new Avg)
       )
     )
+    if (quotaCallback.quotaResetRequired(clientQuotaType))
+      updateQuotaMetricConfigs()
+    sensors
   }
 
-  private def getThrottleTimeSensorName(quotaId: QuotaId): String = quotaType + "ThrottleTime-" + quotaId.sanitizedUser.getOrElse("") + ':' + quotaId.clientId.getOrElse("")
+  protected def registerQuotaMetrics(metricTags: Map[String, String])(sensor: Sensor): Unit = {
+    sensor.add(
+      clientQuotaMetricName(metricTags),
+      new Rate,
+      getQuotaMetricConfig(metricTags)
+    )
+  }
+
+  private def metricTagsToSensorSuffix(metricTags: Map[String, String]): String =
+    metricTags.values.mkString(":")
+
+  private def getThrottleTimeSensorName(metricTags: Map[String, String]): String =
+    s"${quotaType}ThrottleTime-${metricTagsToSensorSuffix(metricTags)}"
 
-  private def getQuotaSensorName(quotaId: QuotaId): String = quotaType + "-" + quotaId.sanitizedUser.getOrElse("") + ':' + quotaId.clientId.getOrElse("")
+  private def getQuotaSensorName(metricTags: Map[String, String]): String =
+    s"$quotaType-${metricTagsToSensorSuffix(metricTags)}"
 
-  protected def getQuotaMetricConfig(quota: Quota): MetricConfig = {
+  protected def getQuotaMetricConfig(metricTags: Map[String, String]): MetricConfig = {
+    getQuotaMetricConfig(quotaLimit(metricTags.asJava))
+  }
+
+  private def getQuotaMetricConfig(quotaLimit: Double): MetricConfig = {
     new MetricConfig()
-            .timeWindow(config.quotaWindowSizeSeconds, TimeUnit.SECONDS)
-            .samples(config.numQuotaSamples)
-            .quota(quota)
+      .timeWindow(config.quotaWindowSizeSeconds, TimeUnit.SECONDS)
+      .samples(config.numQuotaSamples)
+      .quota(new Quota(quotaLimit, true))
   }
 
   protected def getOrCreateSensor(sensorName: String, metricName: MetricName): Sensor = {
     sensorAccessor.getOrCreate(
-        sensorName,
-        ClientQuotaManagerConfig.InactiveSensorExpirationTimeSeconds,
-        metricName,
-        None,
-        new Rate
-      )
+      sensorName,
+      ClientQuotaManager.InactiveSensorExpirationTimeSeconds,
+      sensor => sensor.add(metricName, new Rate)
+    )
   }
 
   /**
    * Overrides quotas for ,  or  or the dynamic defaults
    * for any of these levels.
-   * @param sanitizedUser user to override if quota applies to  or 
-   * @param clientId client to override if quota applies to  or 
+   *
+   * @param sanitizedUser     user to override if quota applies to  or 
+   * @param clientId          client to override if quota applies to  or 
    * @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
+   * @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).
@@ -421,86 +464,229 @@ class ClientQuotaManager(private val config: ClientQuotaManagerConfig,
      */
     lock.writeLock().lock()
     try {
-      val quotaId = QuotaId(sanitizedUser, clientId, sanitizedClientId)
-      val userInfo = sanitizedUser match {
-        case Some(ConfigEntityName.Default) => "default user "
-        case Some(user) => "user " + user + " "
-        case None => ""
+      val userEntity = sanitizedUser.map {
+        case ConfigEntityName.Default => DefaultUserEntity
+        case user => UserEntity(user)
       }
-      val clientIdInfo = clientId match {
-        case Some(ConfigEntityName.Default) => "default client-id"
-        case Some(id) => "client-id " + id
-        case None => ""
+      val clientIdEntity = sanitizedClientId.map {
+        case ConfigEntityName.Default => DefaultClientIdEntity
+        case _ => ClientIdEntity(clientId.getOrElse(throw new IllegalStateException("Client-id not provided")))
       }
+      val quotaEntity = KafkaQuotaEntity(userEntity, clientIdEntity)
+
+      if (userEntity.nonEmpty) {
+        if (quotaEntity.clientIdEntity.nonEmpty)
+          quotaTypesEnabled |= QuotaTypes.UserClientIdQuotaEnabled
+        else
+          quotaTypesEnabled |= QuotaTypes.UserQuotaEnabled
+      } else if (clientIdEntity.nonEmpty)
+        quotaTypesEnabled |= QuotaTypes.ClientIdQuotaEnabled
+
       quota match {
-        case Some(newQuota) =>
-          info(s"Changing ${quotaType} quota for ${userInfo}${clientIdInfo} to $newQuota.bound}")
-          overriddenQuota.put(quotaId, newQuota)
-          (sanitizedUser, clientId) match {
-            case (Some(_), Some(_)) => quotaTypesEnabled |= QuotaTypes.UserClientIdQuotaEnabled
-            case (Some(_), None) => quotaTypesEnabled |= QuotaTypes.UserQuotaEnabled
-            case (None, Some(_)) => quotaTypesEnabled |= QuotaTypes.ClientIdQuotaEnabled
-            case (None, None) =>
-          }
-        case None =>
-          info(s"Removing ${quotaType} quota for ${userInfo}${clientIdInfo}")
-          overriddenQuota.remove(quotaId)
+        case Some(newQuota) => quotaCallback.updateQuota(clientQuotaType, quotaEntity, newQuota.bound)
+        case None => quotaCallback.removeQuota(clientQuotaType, quotaEntity)
       }
+      val updatedEntity = if (userEntity.contains(DefaultUserEntity) || clientIdEntity.contains(DefaultClientIdEntity))
+        None // more than one entity may need updating, so `updateQuotaMetricConfigs` will go through all metrics
+      else
+        Some(quotaEntity)
+      updateQuotaMetricConfigs(updatedEntity)
 
-      val quotaMetricName = clientRateMetricName(sanitizedUser.getOrElse(""), clientId.getOrElse(""))
-      val allMetrics = metrics.metrics()
+    } finally {
+      lock.writeLock().unlock()
+    }
+  }
 
-      // If multiple-levels of quotas are defined or if this is a default quota update, traverse metrics
-      // to find all affected values. Otherwise, update just the single matching one.
-      val singleUpdate = quotaTypesEnabled match {
-        case QuotaTypes.NoQuotas | QuotaTypes.ClientIdQuotaEnabled | QuotaTypes.UserQuotaEnabled | QuotaTypes.UserClientIdQuotaEnabled =>
-          !sanitizedUser.filter(_ == ConfigEntityName.Default).isDefined && !clientId.filter(_ == ConfigEntityName.Default).isDefined
-        case _ => false
+  /**
+   * Updates metrics configs. This is invoked when quota configs are updated in ZooKeeper
+   * or when partitions leaders change and custom callbacks that implement partition-based quotas
+   * have updated quotas.
+   *
+   * @param updatedQuotaEntity If set to one entity and quotas have only been enabled at one
+   *    level, then an optimized update is performed with a single metric update. If None is provided,
+   *    or if custom callbacks are used or if multi-level quotas have been enabled, all metric configs
+   *    are checked and updated if required.
+   */
+  def updateQuotaMetricConfigs(updatedQuotaEntity: Option[KafkaQuotaEntity] = None): Unit = {
+    val allMetrics = metrics.metrics()
+
+    // If using custom quota callbacks or if multiple-levels of quotas are defined or
+    // if this is a default quota update, traverse metrics to find all affected values.
+    // Otherwise, update just the single matching one.
+    val singleUpdate = quotaTypesEnabled match {
+      case QuotaTypes.NoQuotas | QuotaTypes.ClientIdQuotaEnabled | QuotaTypes.UserQuotaEnabled | QuotaTypes.UserClientIdQuotaEnabled =>
+        updatedQuotaEntity.nonEmpty
+      case _ => false
+    }
+    if (singleUpdate) {
+      val quotaEntity = updatedQuotaEntity.getOrElse(throw new IllegalStateException("Quota entity not specified"))
+      val user = quotaEntity.sanitizedUser
+      val clientId = quotaEntity.clientId
+      val metricTags = Map(DefaultTags.User -> user, DefaultTags.ClientId -> clientId)
+
+      val quotaMetricName = clientQuotaMetricName(metricTags)
+      // Change the underlying metric config if the sensor has been created
+      val metric = allMetrics.get(quotaMetricName)
+      if (metric != null) {
+        Option(quotaLimit(metricTags.asJava)).foreach { newQuota =>
+          info(s"Sensor for $quotaEntity already exists. Changing quota to $newQuota in MetricConfig")
+          metric.config(getQuotaMetricConfig(newQuota))
+        }
       }
-      if (singleUpdate) {
-          // Change the underlying metric config if the sensor has been created
-          val metric = allMetrics.get(quotaMetricName)
-          if (metric != null) {
-            val metricConfigEntity = quotaEntity(sanitizedUser.getOrElse(""), clientId.getOrElse(""), sanitizedClientId.getOrElse(""))
-            val newQuota = metricConfigEntity.quota
-            info(s"Sensor for ${userInfo}${clientIdInfo} already exists. Changing quota to ${newQuota.bound()} in MetricConfig")
-            metric.config(getQuotaMetricConfig(newQuota))
-          }
-      } else {
-          allMetrics.asScala.filterKeys(n => n.name == quotaMetricName.name && n.group == quotaMetricName.group).foreach {
-            case (metricName, metric) =>
-              val userTag = if (metricName.tags.containsKey("user")) metricName.tags.get("user") else ""
-              val clientIdTag = if (metricName.tags.containsKey("client-id")) metricName.tags.get("client-id") else ""
-              val metricConfigEntity = quotaEntity(userTag, clientIdTag, Sanitizer.sanitize(clientIdTag))
-              if (metricConfigEntity.quota != metric.config.quota) {
-                val newQuota = metricConfigEntity.quota
-                info(s"Sensor for quota-id ${metricConfigEntity.quotaId} already exists. Setting quota to ${newQuota.bound} in MetricConfig")
-                metric.config(getQuotaMetricConfig(newQuota))
-              }
+    } else {
+      val quotaMetricName = clientQuotaMetricName(Map.empty)
+      allMetrics.forEach { (metricName, metric) =>
+        if (metricName.name == quotaMetricName.name && metricName.group == quotaMetricName.group) {
+          val metricTags = metricName.tags
+          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))
+            }
           }
+        }
       }
-
-    } finally {
-      lock.writeLock().unlock()
     }
   }
 
-  protected def clientRateMetricName(sanitizedUser: String, clientId: String): MetricName = {
+  /**
+   * Returns the MetricName of the metric used for the quota. The name is used to create the
+   * metric but also to find the metric when the quota is changed.
+   */
+  protected def clientQuotaMetricName(quotaMetricTags: Map[String, String]): MetricName = {
     metrics.metricName("byte-rate", quotaType.toString,
-                   "Tracking byte-rate per user/client-id",
-                   "user", sanitizedUser,
-                   "client-id", clientId)
+      "Tracking byte-rate per user/client-id",
+      quotaMetricTags.asJava)
   }
 
-  private def throttleMetricName(quotaEntity: QuotaEntity): MetricName = {
+  private def throttleMetricName(quotaMetricTags: Map[String, String]): MetricName = {
     metrics.metricName("throttle-time",
-                       quotaType.toString,
-                       "Tracking average throttle-time per user/client-id",
-                       "user", quotaEntity.sanitizedUser,
-                       "client-id", quotaEntity.clientId)
+      quotaType.toString,
+      "Tracking average throttle-time per user/client-id",
+      quotaMetricTags.asJava)
   }
 
-  def shutdown() = {
-    throttledRequestReaper.shutdown()
+  def shutdown(): Unit = {
+    throttledChannelReaper.shutdown()
+  }
+
+  class DefaultQuotaCallback extends ClientQuotaCallback {
+    private val overriddenQuotas = new ConcurrentHashMap[ClientQuotaEntity, Quota]()
+
+    override def configure(configs: util.Map[String, _]): Unit = {}
+
+    override def quotaMetricTags(quotaType: ClientQuotaType, principal: KafkaPrincipal, clientId: String): util.Map[String, String] = {
+      quotaMetricTags(Sanitizer.sanitize(principal.getName), clientId).asJava
+    }
+
+    override def quotaLimit(quotaType: ClientQuotaType, metricTags: util.Map[String, String]): lang.Double = {
+      val sanitizedUser = metricTags.get(DefaultTags.User)
+      val clientId = metricTags.get(DefaultTags.ClientId)
+      var quota: Quota = null
+
+      if (sanitizedUser != null && clientId != null) {
+        val userEntity = Some(UserEntity(sanitizedUser))
+        val clientIdEntity = Some(ClientIdEntity(clientId))
+        if (!sanitizedUser.isEmpty && !clientId.isEmpty) {
+          // /config/users//clients/
+          quota = overriddenQuotas.get(KafkaQuotaEntity(userEntity, clientIdEntity))
+          if (quota == null) {
+            // /config/users//clients/
+            quota = overriddenQuotas.get(KafkaQuotaEntity(userEntity, Some(DefaultClientIdEntity)))
+          }
+          if (quota == null) {
+            // /config/users//clients/
+            quota = overriddenQuotas.get(KafkaQuotaEntity(Some(DefaultUserEntity), clientIdEntity))
+          }
+          if (quota == null) {
+            // /config/users//clients/
+            quota = overriddenQuotas.get(DefaultUserClientIdQuotaEntity)
+          }
+        } else if (!sanitizedUser.isEmpty) {
+          // /config/users/
+          quota = overriddenQuotas.get(KafkaQuotaEntity(userEntity, None))
+          if (quota == null) {
+            // /config/users/
+            quota = overriddenQuotas.get(DefaultUserQuotaEntity)
+          }
+        } else if (!clientId.isEmpty) {
+          // /config/clients/
+          quota = overriddenQuotas.get(KafkaQuotaEntity(None, clientIdEntity))
+          if (quota == null) {
+            // /config/clients/
+            quota = overriddenQuotas.get(DefaultClientIdQuotaEntity)
+          }
+          if (quota == null)
+            quota = staticConfigClientIdQuota
+        }
+      }
+      if (quota == null) null else quota.bound
+    }
+
+    override def updateClusterMetadata(cluster: Cluster): Boolean = {
+      // Default quota callback does not use any cluster metadata
+      false
+    }
+
+    override def updateQuota(quotaType: ClientQuotaType, entity: ClientQuotaEntity, newValue: Double): Unit = {
+      val quotaEntity = entity.asInstanceOf[KafkaQuotaEntity]
+      info(s"Changing $quotaType quota for $quotaEntity to $newValue")
+      overriddenQuotas.put(quotaEntity, new Quota(newValue, true))
+    }
+
+    override def removeQuota(quotaType: ClientQuotaType, entity: ClientQuotaEntity): Unit = {
+      val quotaEntity = entity.asInstanceOf[KafkaQuotaEntity]
+      info(s"Removing $quotaType quota for $quotaEntity")
+      overriddenQuotas.remove(quotaEntity)
+    }
+
+    override def quotaResetRequired(quotaType: ClientQuotaType): Boolean = false
+
+    def quotaMetricTags(sanitizedUser: String, clientId: String) : Map[String, String] = {
+      val (userTag, clientIdTag) = quotaTypesEnabled match {
+        case QuotaTypes.NoQuotas | QuotaTypes.ClientIdQuotaEnabled =>
+          ("", clientId)
+        case QuotaTypes.UserQuotaEnabled =>
+          (sanitizedUser, "")
+        case QuotaTypes.UserClientIdQuotaEnabled =>
+          (sanitizedUser, clientId)
+        case _ =>
+          val userEntity = Some(UserEntity(sanitizedUser))
+          val clientIdEntity = Some(ClientIdEntity(clientId))
+
+          var metricTags = (sanitizedUser, clientId)
+          // 1) /config/users//clients/
+          if (!overriddenQuotas.containsKey(KafkaQuotaEntity(userEntity, clientIdEntity))) {
+            // 2) /config/users//clients/
+            metricTags = (sanitizedUser, clientId)
+            if (!overriddenQuotas.containsKey(KafkaQuotaEntity(userEntity, Some(DefaultClientIdEntity)))) {
+              // 3) /config/users/
+              metricTags = (sanitizedUser, "")
+              if (!overriddenQuotas.containsKey(KafkaQuotaEntity(userEntity, None))) {
+                // 4) /config/users//clients/
+                metricTags = (sanitizedUser, clientId)
+                if (!overriddenQuotas.containsKey(KafkaQuotaEntity(Some(DefaultUserEntity), clientIdEntity))) {
+                  // 5) /config/users//clients/
+                  metricTags = (sanitizedUser, clientId)
+                  if (!overriddenQuotas.containsKey(DefaultUserClientIdQuotaEntity)) {
+                    // 6) /config/users/
+                    metricTags = (sanitizedUser, "")
+                    if (!overriddenQuotas.containsKey(DefaultUserQuotaEntity)) {
+                      // 7) /config/clients/
+                      // 8) /config/clients/
+                      // 9) static client-id quota
+                      metricTags = ("", clientId)
+                    }
+                  }
+                }
+              }
+            }
+          }
+          metricTags
+      }
+      Map(DefaultTags.User -> userTag, DefaultTags.ClientId -> clientIdTag)
+    }
+
+    override def close(): Unit = {}
   }
 }
diff --git a/core/src/main/scala/kafka/server/ClientRequestQuotaManager.scala b/core/src/main/scala/kafka/server/ClientRequestQuotaManager.scala
index 59fa4218acb43..2ceaab9c9afdf 100644
--- a/core/src/main/scala/kafka/server/ClientRequestQuotaManager.scala
+++ b/core/src/main/scala/kafka/server/ClientRequestQuotaManager.scala
@@ -19,39 +19,52 @@ package kafka.server
 import java.util.concurrent.TimeUnit
 
 import kafka.network.RequestChannel
+import kafka.utils.QuotaUtils
 import org.apache.kafka.common.MetricName
 import org.apache.kafka.common.metrics._
 import org.apache.kafka.common.utils.Time
+import org.apache.kafka.server.quota.ClientQuotaCallback
 
+import scala.jdk.CollectionConverters._
+
+object ClientRequestQuotaManager {
+  val QuotaRequestPercentDefault = Int.MaxValue.toDouble
+  val NanosToPercentagePerSecond = 100.0 / TimeUnit.SECONDS.toNanos(1)
+
+  private val ExemptSensorName = "exempt-" + QuotaType.Request
+}
 
 class ClientRequestQuotaManager(private val config: ClientQuotaManagerConfig,
                                 private val metrics: Metrics,
                                 private val time: Time,
-                                threadNamePrefix: String)
-                                extends ClientQuotaManager(config, metrics, QuotaType.Request, time, threadNamePrefix) {
-  val maxThrottleTimeMs = TimeUnit.SECONDS.toMillis(this.config.quotaWindowSizeSeconds)
-  def exemptSensor = getOrCreateSensor(exemptSensorName, exemptMetricName)
+                                private val threadNamePrefix: String,
+                                private val quotaCallback: Option[ClientQuotaCallback])
+    extends ClientQuotaManager(config, metrics, QuotaType.Request, time, threadNamePrefix, quotaCallback) {
 
-  def recordExempt(value: Double) {
+  private val maxThrottleTimeMs = TimeUnit.SECONDS.toMillis(this.config.quotaWindowSizeSeconds)
+  private val exemptMetricName = metrics.metricName("exempt-request-time",
+    QuotaType.Request.toString, "Tracking exempt-request-time utilization percentage")
+
+  lazy val exemptSensor: Sensor = getOrCreateSensor(ClientRequestQuotaManager.ExemptSensorName, exemptMetricName)
+
+  def recordExempt(value: Double): Unit = {
     exemptSensor.record(value)
   }
 
-  def maybeRecordAndThrottle(request: RequestChannel.Request, sendResponseCallback: Int => Unit): Unit = {
-    if (request.apiRemoteCompleteTimeNanos == -1) {
-      // When this callback is triggered, the remote API call has completed
-      request.apiRemoteCompleteTimeNanos = time.nanoseconds
-    }
-
+  /**
+    * Records that a user/clientId changed request processing time being throttled. If quota has been violated, return
+    * throttle time in milliseconds. Throttle time calculation may be overridden by sub-classes.
+    * @param request client request
+    * @return Number of milliseconds to throttle in case of quota violation. Zero otherwise
+    */
+  def maybeRecordAndGetThrottleTimeMs(request: RequestChannel.Request, timeMs: Long): Int = {
     if (quotasEnabled) {
-      val quotaSensors = getOrCreateQuotaSensors(request.session.sanitizedUser, request.header.clientId)
-      request.recordNetworkThreadTimeCallback = Some(timeNanos => recordNoThrottle(quotaSensors, nanosToPercentage(timeNanos)))
-
-      recordAndThrottleOnQuotaViolation(
-          quotaSensors,
-          nanosToPercentage(request.requestThreadTimeNanos),
-          sendResponseCallback)
+      request.recordNetworkThreadTimeCallback = Some(timeNanos => recordNoThrottle(
+        request.session, request.header.clientId, nanosToPercentage(timeNanos)))
+      recordAndGetThrottleTimeMs(request.session, request.header.clientId,
+        nanosToPercentage(request.requestThreadTimeNanos), timeMs)
     } else {
-      sendResponseCallback(0)
+      0
     }
   }
 
@@ -62,24 +75,16 @@ class ClientRequestQuotaManager(private val config: ClientQuotaManagerConfig,
     }
   }
 
-  override protected def throttleTime(clientMetric: KafkaMetric, config: MetricConfig): Long = {
-    math.min(super.throttleTime(clientMetric, config), maxThrottleTimeMs)
+  override protected def throttleTime(e: QuotaViolationException, timeMs: Long): Long = {
+    QuotaUtils.boundedThrottleTime(e, maxThrottleTimeMs, timeMs)
   }
 
-  override protected def clientRateMetricName(sanitizedUser: String, clientId: String): MetricName = {
+  override protected def clientQuotaMetricName(quotaMetricTags: Map[String, String]): MetricName = {
     metrics.metricName("request-time", QuotaType.Request.toString,
-                   "Tracking request-time per user/client-id",
-                   "user", sanitizedUser,
-                   "client-id", clientId)
-  }
-
-  private def exemptMetricName: MetricName = {
-    metrics.metricName("exempt-request-time", QuotaType.Request.toString,
-                   "Tracking exempt-request-time utilization percentage")
+      "Tracking request-time per user/client-id",
+      quotaMetricTags.asJava)
   }
 
-  private def exemptSensorName: String = "exempt-" + QuotaType.Request
-
-  private def nanosToPercentage(nanos: Long): Double = nanos * ClientQuotaManagerConfig.NanosToPercentagePerSecond
-
+  private def nanosToPercentage(nanos: Long): Double =
+    nanos * ClientRequestQuotaManager.NanosToPercentagePerSecond
 }
diff --git a/core/src/main/scala/kafka/server/ConfigHandler.scala b/core/src/main/scala/kafka/server/ConfigHandler.scala
index 390222d500ab3..c998ac2910d73 100644
--- a/core/src/main/scala/kafka/server/ConfigHandler.scala
+++ b/core/src/main/scala/kafka/server/ConfigHandler.scala
@@ -17,11 +17,14 @@
 
 package kafka.server
 
+import java.net.{InetAddress, UnknownHostException}
 import java.util.Properties
 
 import DynamicConfig.Broker._
 import kafka.api.ApiVersion
+import kafka.controller.KafkaController
 import kafka.log.{LogConfig, LogManager}
+import kafka.network.ConnectionQuotas
 import kafka.security.CredentialProvider
 import kafka.server.Constants._
 import kafka.server.QuotaFactory.QuotaManagers
@@ -33,42 +36,44 @@ import org.apache.kafka.common.metrics.Quota
 import org.apache.kafka.common.metrics.Quota._
 import org.apache.kafka.common.utils.Sanitizer
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
+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
 }
 
 /**
   * The TopicConfigHandler will process topic config changes in ZK.
   * The callback provides the topic name and the full properties set read from ZK
   */
-class TopicConfigHandler(private val logManager: LogManager, kafkaConfig: KafkaConfig, val quotas: QuotaManagers) extends ConfigHandler with Logging  {
+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()
-      props ++= logManager.defaultConfig.originals.asScala
-      topicConfig.asScala.foreach { case (key, value) =>
+      topicConfig.asScala.forKeyValue { (key, value) =>
         if (!configNamesToExclude.contains(key)) props.put(key, value)
       }
-      val logConfig = LogConfig(props)
-      if ((topicConfig.containsKey(LogConfig.RetentionMsProp)
-        || topicConfig.containsKey(LogConfig.MessageTimestampDifferenceMaxMsProp))
-        && logConfig.retentionMs < logConfig.messageTimestampDifferenceMaxMs)
-        warn(s"${LogConfig.RetentionMsProp} for topic $topic is set to ${logConfig.retentionMs}. It is smaller than " +
-          s"${LogConfig.MessageTimestampDifferenceMaxMsProp}'s value ${logConfig.messageTimestampDifferenceMaxMs}. " +
-          s"This may result in frequent log rolling.")
-      logs.foreach(_.config = logConfig)
+      val logConfig = LogConfig.fromProps(logManager.currentDefaultConfig.originals, props)
+      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) {
@@ -82,6 +87,10 @@ class TopicConfigHandler(private val logManager: LogManager, kafkaConfig: KafkaC
     }
     updateThrottledList(LogConfig.LeaderReplicationThrottledReplicasProp, quotas.leader)
     updateThrottledList(LogConfig.FollowerReplicationThrottledReplicasProp, quotas.follower)
+
+    if (Try(topicConfig.getProperty(KafkaConfig.UncleanLeaderElectionEnableProp).toBoolean).getOrElse(false)) {
+      kafkaController.enableTopicUncleanLeaderElection(topic)
+    }
   }
 
   def parseThrottledPartitions(topicConfig: Properties, brokerId: Int, prop: String): Seq[Int] = {
@@ -118,17 +127,17 @@ 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))
-        Some(new Quota(config.getProperty(DynamicConfig.Client.ProducerByteRateOverrideProp).toLong, true))
+        Some(new Quota(config.getProperty(DynamicConfig.Client.ProducerByteRateOverrideProp).toLong.toDouble, true))
       else
         None
     quotaManagers.produce.updateQuota(sanitizedUser, clientId, sanitizedClientId, producerQuota)
     val consumerQuota =
       if (config.containsKey(DynamicConfig.Client.ConsumerByteRateOverrideProp))
-        Some(new Quota(config.getProperty(DynamicConfig.Client.ConsumerByteRateOverrideProp).toLong, true))
+        Some(new Quota(config.getProperty(DynamicConfig.Client.ConsumerByteRateOverrideProp).toLong.toDouble, true))
       else
         None
     quotaManagers.fetch.updateQuota(sanitizedUser, clientId, sanitizedClientId, consumerQuota)
@@ -138,6 +147,12 @@ class QuotaConfigHandler(private val quotaManagers: QuotaManagers) {
       else
         None
     quotaManagers.request.updateQuota(sanitizedUser, clientId, sanitizedClientId, requestQuota)
+    val controllerMutationQuota =
+      if (config.containsKey(DynamicConfig.Client.ControllerMutationOverrideProp))
+        Some(new Quota(config.getProperty(DynamicConfig.Client.ControllerMutationOverrideProp).toDouble, true))
+      else
+        None
+    quotaManagers.controllerMutation.updateQuota(sanitizedUser, clientId, sanitizedClientId, controllerMutationQuota)
   }
 }
 
@@ -147,7 +162,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)
   }
 }
@@ -159,7 +174,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)
@@ -172,24 +187,46 @@ class UserConfigHandler(private val quotaManagers: QuotaManagers, val credential
   }
 }
 
+class IpConfigHandler(private val connectionQuotas: ConnectionQuotas) extends ConfigHandler with Logging {
+
+  def processConfigChanges(ip: String, config: Properties): Unit = {
+    val ipConnectionRateQuota = Option(config.getProperty(DynamicConfig.Ip.IpConnectionRateOverrideProp)).map(_.toInt)
+    val updatedIp = {
+      if (ip != ConfigEntityName.Default) {
+        try {
+          Some(InetAddress.getByName(ip))
+        } catch {
+          case _: UnknownHostException => throw new IllegalArgumentException(s"Unable to resolve address $ip")
+        }
+      } else
+        None
+    }
+    connectionQuotas.updateIpConnectionRateQuota(updatedIp, ipConnectionRateQuota)
+  }
+}
+
 /**
   * The BrokerConfigHandler will process individual broker config changes in ZK.
   * The callback provides the brokerId and the full properties set read from ZK.
   * This implementation reports the overrides to the respective ReplicationQuotaManager objects
   */
-class BrokerConfigHandler(private val brokerConfig: KafkaConfig, private val quotaManagers: QuotaManagers) extends ConfigHandler with Logging {
+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
       else
         DefaultReplicationThrottledRate
     }
-    if (brokerConfig.brokerId == brokerId.trim.toInt) {
-      quotaManagers.leader.updateQuota(upperBound(getOrDefault(LeaderReplicationThrottledRateProp)))
-      quotaManagers.follower.updateQuota(upperBound(getOrDefault(FollowerReplicationThrottledRateProp)))
-      quotaManagers.alterLogDirs.updateQuota(upperBound(getOrDefault(ReplicaAlterLogDirsIoMaxBytesPerSecondProp)))
+    if (brokerId == ConfigEntityName.Default)
+      brokerConfig.dynamicConfig.updateDefaultConfig(properties)
+    else if (brokerConfig.brokerId == brokerId.trim.toInt) {
+      brokerConfig.dynamicConfig.updateBrokerConfig(brokerConfig.brokerId, properties)
+      quotaManagers.leader.updateQuota(upperBound(getOrDefault(LeaderReplicationThrottledRateProp).toDouble))
+      quotaManagers.follower.updateQuota(upperBound(getOrDefault(FollowerReplicationThrottledRateProp).toDouble))
+      quotaManagers.alterLogDirs.updateQuota(upperBound(getOrDefault(ReplicaAlterLogDirsIoMaxBytesPerSecondProp).toDouble))
     }
   }
 }
@@ -201,9 +238,9 @@ object ThrottledReplicaListValidator extends Validator {
   override def ensureValid(name: String, value: Any): Unit = {
     def check(proposed: Seq[Any]): Unit = {
       if (!(proposed.forall(_.toString.trim.matches("([0-9]+:[0-9]+)?"))
-        || proposed.headOption.exists(_.toString.trim.equals("*"))))
+        || proposed.mkString.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)
@@ -212,6 +249,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/ControllerMutationQuotaManager.scala b/core/src/main/scala/kafka/server/ControllerMutationQuotaManager.scala
new file mode 100644
index 0000000000000..f011a6b36632d
--- /dev/null
+++ b/core/src/main/scala/kafka/server/ControllerMutationQuotaManager.scala
@@ -0,0 +1,282 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package kafka.server
+
+import kafka.network.RequestChannel
+import kafka.network.RequestChannel.Session
+import org.apache.kafka.common.MetricName
+import org.apache.kafka.common.errors.ThrottlingQuotaExceededException
+import org.apache.kafka.common.metrics.Metrics
+import org.apache.kafka.common.metrics.QuotaViolationException
+import org.apache.kafka.common.metrics.Sensor
+import org.apache.kafka.common.metrics.stats.Rate
+import org.apache.kafka.common.metrics.stats.TokenBucket
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.utils.Time
+import org.apache.kafka.server.quota.ClientQuotaCallback
+
+import scala.jdk.CollectionConverters._
+
+/**
+ * The ControllerMutationQuota trait defines a quota for a given user/clientId pair. Such
+ * quota is not meant to be cached forever but rather during the lifetime of processing
+ * a request.
+ */
+trait ControllerMutationQuota {
+  def isExceeded: Boolean
+  def record(permits: Double): Unit
+  def throttleTime: Int
+}
+
+/**
+ * Default quota used when quota is disabled.
+ */
+object UnboundedControllerMutationQuota extends ControllerMutationQuota {
+  override def isExceeded: Boolean = false
+  override def record(permits: Double): Unit = ()
+  override def throttleTime: Int = 0
+}
+
+/**
+ * The AbstractControllerMutationQuota is the base class of StrictControllerMutationQuota and
+ * PermissiveControllerMutationQuota.
+ *
+ * @param time @Time object to use
+ */
+abstract class AbstractControllerMutationQuota(private val time: Time) extends ControllerMutationQuota {
+  protected var lastThrottleTimeMs = 0L
+  protected var lastRecordedTimeMs = 0L
+
+  protected def updateThrottleTime(e: QuotaViolationException, timeMs: Long): Unit = {
+    lastThrottleTimeMs = ControllerMutationQuotaManager.throttleTimeMs(e, timeMs)
+    lastRecordedTimeMs = timeMs
+  }
+
+  override def throttleTime: Int = {
+    // If a throttle time has been recorded, we adjust it by deducting the time elapsed
+    // between the recording and now. We do this because `throttleTime` may be called
+    // long after having recorded it, especially when a request waits in the purgatory.
+    val deltaTimeMs = time.milliseconds - lastRecordedTimeMs
+    Math.max(0, lastThrottleTimeMs - deltaTimeMs).toInt
+  }
+}
+
+/**
+ * The StrictControllerMutationQuota defines a strict quota for a given user/clientId pair. The
+ * quota is strict meaning that 1) it does not accept any mutations once the quota is exhausted
+ * until it gets back to the defined rate; and 2) it does not throttle for any number of mutations
+ * if quota is not already exhausted.
+ *
+ * @param time @Time object to use
+ * @param quotaSensor @Sensor object with a defined quota for a given user/clientId pair
+ */
+class StrictControllerMutationQuota(private val time: Time,
+                                    private val quotaSensor: Sensor)
+    extends AbstractControllerMutationQuota(time) {
+
+  override def isExceeded: Boolean = lastThrottleTimeMs > 0
+
+  override def record(permits: Double): Unit = {
+    val timeMs = time.milliseconds
+    try {
+      quotaSensor synchronized {
+        quotaSensor.checkQuotas(timeMs)
+        quotaSensor.record(permits, timeMs, false)
+      }
+    } catch {
+      case e: QuotaViolationException =>
+        updateThrottleTime(e, timeMs)
+        throw new ThrottlingQuotaExceededException(lastThrottleTimeMs.toInt,
+          Errors.THROTTLING_QUOTA_EXCEEDED.message)
+    }
+  }
+}
+
+/**
+ * The PermissiveControllerMutationQuota defines a permissive quota for a given user/clientId pair.
+ * The quota is permissive meaning that 1) it does accept any mutations even if the quota is
+ * exhausted; and 2) it does throttle as soon as the quota is exhausted.
+ *
+ * @param time @Time object to use
+ * @param quotaSensor @Sensor object with a defined quota for a given user/clientId pair
+ */
+class PermissiveControllerMutationQuota(private val time: Time,
+                                        private val quotaSensor: Sensor)
+    extends AbstractControllerMutationQuota(time) {
+
+  override def isExceeded: Boolean = false
+
+  override def record(permits: Double): Unit = {
+    val timeMs = time.milliseconds
+    try {
+      quotaSensor.record(permits, timeMs, true)
+    } catch {
+      case e: QuotaViolationException =>
+        updateThrottleTime(e, timeMs)
+    }
+  }
+}
+
+object ControllerMutationQuotaManager {
+  val QuotaControllerMutationDefault = Int.MaxValue.toDouble
+
+  /**
+   * This calculates the amount of time needed to bring the TokenBucket within quota
+   * assuming that no new metrics are recorded.
+   *
+   * Basically, if a value < 0 is observed, the time required to bring it to zero is
+   * -value / refill rate (quota bound) * 1000.
+   */
+  def throttleTimeMs(e: QuotaViolationException, timeMs: Long): Long = {
+    e.metric().measurable() match {
+      case _: TokenBucket =>
+        Math.round(-e.value() / e.bound() * 1000)
+      case _ => throw new IllegalArgumentException(
+        s"Metric ${e.metric().metricName()} is not a TokenBucket metric, value ${e.metric().measurable()}")
+    }
+  }
+}
+
+/**
+ * The ControllerMutationQuotaManager is a specialized ClientQuotaManager used in the context
+ * of throttling controller's operations/mutations.
+ *
+ * @param config @ClientQuotaManagerConfig quota configs
+ * @param metrics @Metrics Metrics instance
+ * @param time @Time object to use
+ * @param threadNamePrefix The thread prefix to use
+ * @param quotaCallback @ClientQuotaCallback ClientQuotaCallback to use
+ */
+class ControllerMutationQuotaManager(private val config: ClientQuotaManagerConfig,
+                                     private val metrics: Metrics,
+                                     private val time: Time,
+                                     private val threadNamePrefix: String,
+                                     private val quotaCallback: Option[ClientQuotaCallback])
+    extends ClientQuotaManager(config, metrics, QuotaType.ControllerMutation, time, threadNamePrefix, quotaCallback) {
+
+  override protected def clientQuotaMetricName(quotaMetricTags: Map[String, String]): MetricName = {
+    metrics.metricName("tokens", QuotaType.ControllerMutation.toString,
+      "Tracking remaining tokens in the token bucket per user/client-id",
+      quotaMetricTags.asJava)
+  }
+
+  private def clientRateMetricName(quotaMetricTags: Map[String, String]): MetricName = {
+    metrics.metricName("mutation-rate", QuotaType.ControllerMutation.toString,
+      "Tracking mutation-rate per user/client-id",
+      quotaMetricTags.asJava)
+  }
+
+  override protected def registerQuotaMetrics(metricTags: Map[String, String])(sensor: Sensor): Unit = {
+    sensor.add(
+      clientRateMetricName(metricTags),
+      new Rate
+    )
+    sensor.add(
+      clientQuotaMetricName(metricTags),
+      new TokenBucket,
+      getQuotaMetricConfig(metricTags)
+    )
+  }
+
+  /**
+   * Records that a user/clientId accumulated or would like to accumulate the provided amount at the
+   * the specified time, returns throttle time in milliseconds. The quota is strict meaning that it
+   * does not accept any mutations once the quota is exhausted until it gets back to the defined rate.
+   *
+   * @param session The session from which the user is extracted
+   * @param clientId The client id
+   * @param value The value to accumulate
+   * @param timeMs The time at which to accumulate the value
+   * @return The throttle time in milliseconds defines as the time to wait until the average
+   *         rate gets back to the defined quota
+   */
+  override def recordAndGetThrottleTimeMs(session: Session, clientId: String, value: Double, timeMs: Long): Int = {
+    val clientSensors = getOrCreateQuotaSensors(session, clientId)
+    val quotaSensor = clientSensors.quotaSensor
+    try {
+      quotaSensor synchronized {
+        quotaSensor.checkQuotas(timeMs)
+        quotaSensor.record(value, timeMs, false)
+      }
+      0
+    } catch {
+      case e: QuotaViolationException =>
+        val throttleTimeMs = ControllerMutationQuotaManager.throttleTimeMs(e, timeMs).toInt
+        debug(s"Quota violated for sensor (${quotaSensor.name}). Delay time: ($throttleTimeMs)")
+        throttleTimeMs
+    }
+  }
+
+  /**
+   * Returns a StrictControllerMutationQuota for the given user/clientId pair or
+   * a UnboundedControllerMutationQuota$ if the quota is disabled.
+   *
+   * @param session The session from which the user is extracted
+   * @param clientId The client id
+   * @return ControllerMutationQuota
+   */
+  def newStrictQuotaFor(session: Session, clientId: String): ControllerMutationQuota = {
+    if (quotasEnabled) {
+      val clientSensors = getOrCreateQuotaSensors(session, clientId)
+      new StrictControllerMutationQuota(time, clientSensors.quotaSensor)
+    } else {
+      UnboundedControllerMutationQuota
+    }
+  }
+
+  def newStrictQuotaFor(request: RequestChannel.Request): ControllerMutationQuota =
+    newStrictQuotaFor(request.session, request.header.clientId)
+
+  /**
+   * Returns a PermissiveControllerMutationQuota for the given user/clientId pair or
+   * a UnboundedControllerMutationQuota$ if the quota is disabled.
+   *
+   * @param session The session from which the user is extracted
+   * @param clientId The client id
+   * @return ControllerMutationQuota
+   */
+  def newPermissiveQuotaFor(session: Session, clientId: String): ControllerMutationQuota = {
+    if (quotasEnabled) {
+      val clientSensors = getOrCreateQuotaSensors(session, clientId)
+      new PermissiveControllerMutationQuota(time, clientSensors.quotaSensor)
+    } else {
+      UnboundedControllerMutationQuota
+    }
+  }
+
+  def newPermissiveQuotaFor(request: RequestChannel.Request): ControllerMutationQuota =
+    newPermissiveQuotaFor(request.session, request.header.clientId)
+
+  /**
+   * Returns a ControllerMutationQuota based on `strictSinceVersion`. It returns a strict
+   * quota if the version is equal to or above of the `strictSinceVersion`, a permissive
+   * quota if the version is below, and a unbounded quota if the quota is disabled.
+   *
+   * When the quota is strictly enforced. Any operation above the quota is not allowed
+   * and rejected with a THROTTLING_QUOTA_EXCEEDED error.
+   *
+   * @param request The request to extract the user and the clientId from
+   * @param strictSinceVersion The version since quota is strict
+   * @return
+   */
+  def newQuotaFor(request: RequestChannel.Request, strictSinceVersion: Short): ControllerMutationQuota = {
+    if (request.header.apiVersion() >= strictSinceVersion)
+      newStrictQuotaFor(request)
+    else
+      newPermissiveQuotaFor(request)
+  }
+}
diff --git a/core/src/main/scala/kafka/server/DelayedCreatePartitions.scala b/core/src/main/scala/kafka/server/DelayedCreatePartitions.scala
index 0a9948311af0c..18fdcd26ebcf4 100644
--- a/core/src/main/scala/kafka/server/DelayedCreatePartitions.scala
+++ b/core/src/main/scala/kafka/server/DelayedCreatePartitions.scala
@@ -26,7 +26,21 @@ import scala.collection._
 /**
   * The create metadata maintained by the delayed create topic or create partitions operations.
   */
-case class CreatePartitionsMetadata(topic: String, replicaAssignments: Map[Int, Seq[Int]], error: ApiError)
+case class CreatePartitionsMetadata(topic: String, partitions: Set[Int], error: ApiError)
+
+object CreatePartitionsMetadata {
+  def apply(topic: String, partitions: Set[Int]): CreatePartitionsMetadata = {
+    CreatePartitionsMetadata(topic, partitions, ApiError.NONE)
+  }
+
+  def apply(topic: String, error: Errors): CreatePartitionsMetadata = {
+    CreatePartitionsMetadata(topic, Set.empty, new ApiError(error, null))
+  }
+
+  def apply(topic: String, throwable: Throwable): CreatePartitionsMetadata = {
+    CreatePartitionsMetadata(topic, Set.empty, ApiError.fromThrowable(throwable))
+  }
+}
 
 /**
   * A delayed create topic or create partitions operation that is stored in the topic purgatory.
@@ -46,7 +60,7 @@ class DelayedCreatePartitions(delayMs: Long,
     trace(s"Trying to complete operation for $createMetadata")
 
     val leaderlessPartitionCount = createMetadata.filter(_.error.isSuccess).foldLeft(0) { case (topicCounter, metadata) =>
-      topicCounter + missingLeaderCount(metadata.topic, metadata.replicaAssignments.keySet)
+      topicCounter + missingLeaderCount(metadata.topic, metadata.partitions)
     }
 
     if (leaderlessPartitionCount == 0) {
@@ -61,11 +75,11 @@ 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
-      if (metadata.error.isSuccess && missingLeaderCount(metadata.topic, metadata.replicaAssignments.keySet) > 0)
+      if (metadata.error.isSuccess && missingLeaderCount(metadata.topic, metadata.partitions) > 0)
         (metadata.topic, new ApiError(Errors.REQUEST_TIMED_OUT, null))
       else
         (metadata.topic, metadata.error)
@@ -83,6 +97,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 7a00bc1e2da90..317d0b89c3754 100644
--- a/core/src/main/scala/kafka/server/DelayedDeleteRecords.scala
+++ b/core/src/main/scala/kafka/server/DelayedDeleteRecords.scala
@@ -21,19 +21,21 @@ package kafka.server
 import java.util.concurrent.TimeUnit
 
 import kafka.metrics.KafkaMetricsGroup
-import org.apache.kafka.common.protocol.Errors
+import kafka.utils.Implicits._
 import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.message.DeleteRecordsResponseData
+import org.apache.kafka.common.protocol.Errors
 import org.apache.kafka.common.requests.DeleteRecordsResponse
 
 import scala.collection._
 
 
 case class DeleteRecordsPartitionStatus(requiredOffset: Long,
-                                        responseStatus: DeleteRecordsResponse.PartitionResponse) {
+                                        responseStatus: DeleteRecordsResponseData.DeleteRecordsPartitionResult) {
   @volatile var acksPending = false
 
   override def toString = "[acksPending: %b, error: %s, lowWatermark: %d, requiredOffset: %d]"
-    .format(acksPending, responseStatus.error.toString, responseStatus.lowWatermark, requiredOffset)
+    .format(acksPending, Errors.forCode(responseStatus.errorCode).toString, responseStatus.lowWatermark, requiredOffset)
 }
 
 /**
@@ -43,15 +45,15 @@ case class DeleteRecordsPartitionStatus(requiredOffset: Long,
 class DelayedDeleteRecords(delayMs: Long,
                            deleteRecordsStatus:  Map[TopicPartition, DeleteRecordsPartitionStatus],
                            replicaManager: ReplicaManager,
-                           responseCallback: Map[TopicPartition, DeleteRecordsResponse.PartitionResponse] => Unit)
+                           responseCallback: Map[TopicPartition, DeleteRecordsResponseData.DeleteRecordsPartitionResult] => Unit)
   extends DelayedOperation(delayMs) {
 
   // first update the acks pending variable according to the error code
-  deleteRecordsStatus.foreach { case (topicPartition, status) =>
-    if (status.responseStatus.error == Errors.NONE) {
+  deleteRecordsStatus.forKeyValue { (topicPartition, status) =>
+    if (status.responseStatus.errorCode == Errors.NONE.code) {
       // Timeout error state will be cleared when required acks are received
       status.acksPending = true
-      status.responseStatus.error = Errors.REQUEST_TIMED_OUT
+      status.responseStatus.setErrorCode(Errors.REQUEST_TIMED_OUT.code)
     } else {
       status.acksPending = false
     }
@@ -68,30 +70,30 @@ class DelayedDeleteRecords(delayMs: Long,
    */
   override def tryComplete(): Boolean = {
     // check for each partition if it still has pending acks
-    deleteRecordsStatus.foreach { case (topicPartition, status) =>
-      trace(s"Checking delete records satisfaction for ${topicPartition}, current status $status")
+    deleteRecordsStatus.forKeyValue { (topicPartition, status) =>
+      trace(s"Checking delete records satisfaction for $topicPartition, current status $status")
       // 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_OR_FOLLOWER, 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) {
           status.acksPending = false
-          status.responseStatus.error = error
-          status.responseStatus.lowWatermark = lw
+          status.responseStatus.setErrorCode(error.code)
+          status.responseStatus.setLowWatermark(lw)
         }
       }
     }
@@ -103,8 +105,8 @@ class DelayedDeleteRecords(delayMs: Long,
       false
   }
 
-  override def onExpiration() {
-    deleteRecordsStatus.foreach { case (topicPartition, status) =>
+  override def onExpiration(): Unit = {
+    deleteRecordsStatus.forKeyValue { (topicPartition, status) =>
       if (status.acksPending) {
         DelayedDeleteRecordsMetrics.recordExpiration(topicPartition)
       }
@@ -114,8 +116,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 +126,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..b9fb12057af44 100644
--- a/core/src/main/scala/kafka/server/DelayedDeleteTopics.scala
+++ b/core/src/main/scala/kafka/server/DelayedDeleteTopics.scala
@@ -26,6 +26,12 @@ import scala.collection._
   */
 case class DeleteTopicMetadata(topic: String, error: Errors)
 
+object DeleteTopicMetadata {
+  def apply(topic: String, throwable: Throwable): DeleteTopicMetadata = {
+    DeleteTopicMetadata(topic, Errors.forException(throwable))
+  }
+}
+
 /**
   * A delayed delete topics operation that can be created by the admin manager and watched
   * in the topic purgatory
@@ -57,7 +63,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/DelayedElectLeader.scala b/core/src/main/scala/kafka/server/DelayedElectLeader.scala
new file mode 100644
index 0000000000000..cd0a804058971
--- /dev/null
+++ b/core/src/main/scala/kafka/server/DelayedElectLeader.scala
@@ -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 kafka.server
+
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.requests.ApiError
+
+import scala.collection.{Map, mutable}
+
+/** A delayed elect leader operation that can be created by the replica manager and watched
+  * in the elect leader purgatory
+  */
+class DelayedElectLeader(
+  delayMs: Long,
+  expectedLeaders: Map[TopicPartition, Int],
+  results: Map[TopicPartition, ApiError],
+  replicaManager: ReplicaManager,
+  responseCallback: Map[TopicPartition, ApiError] => Unit
+) extends DelayedOperation(delayMs) {
+
+  private val waitingPartitions = mutable.Map() ++= expectedLeaders
+  private val fullResults = mutable.Map() ++= results
+
+
+  /**
+    * Call-back to execute when a delayed operation gets expired and hence forced to complete.
+    */
+  override def onExpiration(): Unit = {}
+
+  /**
+    * Process for completing an operation; This function needs to be defined
+    * in subclasses and will be called exactly once in forceComplete()
+    */
+  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, _) => tp -> new ApiError(Errors.REQUEST_TIMED_OUT, null)
+    }
+    responseCallback(timedOut ++ fullResults)
+  }
+
+  /**
+    * Try to complete the delayed operation by first checking if the operation
+    * can be completed by now. If yes execute the completion logic by calling
+    * forceComplete() and return true iff forceComplete returns true; otherwise return false
+    *
+    * This function needs to be defined in subclasses
+    */
+  override def tryComplete(): Boolean = {
+    updateWaiting()
+    debug(s"tryComplete() waitingPartitions: $waitingPartitions")
+    waitingPartitions.isEmpty && forceComplete()
+  }
+
+  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 e478053792deb..60e060817a3cf 100644
--- a/core/src/main/scala/kafka/server/DelayedFetch.scala
+++ b/core/src/main/scala/kafka/server/DelayedFetch.scala
@@ -21,16 +21,21 @@ import java.util.concurrent.TimeUnit
 
 import kafka.metrics.KafkaMetricsGroup
 import org.apache.kafka.common.TopicPartition
-import org.apache.kafka.common.errors.{NotLeaderForPartitionException, UnknownTopicOrPartitionException, KafkaStorageException}
+import org.apache.kafka.common.errors._
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.replica.ClientMetadata
 import org.apache.kafka.common.requests.FetchRequest.PartitionData
-import org.apache.kafka.common.requests.IsolationLevel
+import org.apache.kafka.common.requests.OffsetsForLeaderEpochResponse.{UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET}
 
 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 +
+      "]"
+  }
 }
 
 /**
@@ -40,17 +45,17 @@ case class FetchMetadata(fetchMinBytes: Int,
                          fetchMaxBytes: Int,
                          hardMaxBytesLimit: Boolean,
                          fetchOnlyLeader: Boolean,
-                         fetchOnlyCommitted: Boolean,
+                         fetchIsolation: FetchIsolation,
                          isFromFollower: Boolean,
                          replicaId: Int,
                          fetchPartitionStatus: Seq[(TopicPartition, FetchPartitionStatus)]) {
 
-  override def toString = "[minBytes: " + fetchMinBytes + ", " +
-    "maxBytes:" + fetchMaxBytes + ", " +
-    "onlyLeader:" + fetchOnlyLeader + ", " +
-    "onlyCommitted: " + fetchOnlyCommitted + ", " +
-    "replicaId: " + replicaId + ", " +
-    "partitionStatus: " + fetchPartitionStatus + "]"
+  override def toString = "FetchMetadata(minBytes=" + fetchMinBytes + ", " +
+    "maxBytes=" + fetchMaxBytes + ", " +
+    "onlyLeader=" + fetchOnlyLeader + ", " +
+    "fetchIsolation=" + fetchIsolation + ", " +
+    "replicaId=" + replicaId + ", " +
+    "partitionStatus=" + fetchPartitionStatus + ")"
 }
 /**
  * A delayed fetch operation that can be created by the replica manager and watched
@@ -60,7 +65,7 @@ class DelayedFetch(delayMs: Long,
                    fetchMetadata: FetchMetadata,
                    replicaManager: ReplicaManager,
                    quota: ReplicaQuota,
-                   isolationLevel: IsolationLevel,
+                   clientMetadata: Option[ClientMetadata],
                    responseCallback: Seq[(TopicPartition, FetchPartitionData)] => Unit)
   extends DelayedOperation(delayMs) {
 
@@ -68,77 +73,95 @@ 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 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: A diverging epoch was found, return response to trigger truncation
    * Upon completion, should return whatever data is available for each valid partition
    */
-  override def tryComplete() : Boolean = {
+  override def tryComplete(): Boolean = {
     var accumulatedSize = 0
-    var accumulatedThrottledSize = 0
     fetchMetadata.fetchPartitionStatus.foreach {
       case (topicPartition, fetchStatus) =>
         val fetchOffset = fetchStatus.startOffsetMetadata
+        val fetchLeaderEpoch = fetchStatus.fetchInfo.currentLeaderEpoch
         try {
           if (fetchOffset != LogOffsetMetadata.UnknownOffsetMetadata) {
-            val replica = replicaManager.getLeaderReplicaIfLocal(topicPartition)
-            val endOffset =
-              if (isolationLevel == IsolationLevel.READ_COMMITTED)
-                replica.lastStableOffset
-              else if (fetchMetadata.fetchOnlyCommitted)
-                replica.highWatermark
-              else
-                replica.logEndOffset
-
-            // Go directly to the check for Case D if the message offsets are the same. If the log segment
+            val partition = replicaManager.getPartitionOrException(topicPartition)
+            val offsetSnapshot = partition.fetchOffsetSnapshot(fetchLeaderEpoch, fetchMetadata.fetchOnlyLeader)
+
+            val endOffset = fetchMetadata.fetchIsolation match {
+              case FetchLogEnd => offsetSnapshot.logEndOffset
+              case FetchHighWatermark => offsetSnapshot.highWatermark
+              case FetchTxnCommitted => offsetSnapshot.lastStableOffset
+            }
+
+            // 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
-                debug("Satisfying fetch %s since it is fetching later segments of partition %s.".format(fetchMetadata, topicPartition))
+                // 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("Satisfying fetch %s immediately since it is fetching older segments.".format(fetchMetadata))
+                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.
-                if (!replicaManager.shouldLeaderThrottle(quota, topicPartition, fetchMetadata.replicaId))
+                if (!replicaManager.shouldLeaderThrottle(quota, partition, fetchMetadata.replicaId))
                   return forceComplete()
               } else if (fetchOffset.messageOffset < endOffset.messageOffset) {
                 // we take the partition fetch size as upper bound when accumulating the bytes (skip if a throttled partition)
                 val bytesAvailable = math.min(endOffset.positionDiff(fetchOffset), fetchStatus.fetchInfo.maxBytes)
-                if (quota.isThrottled(topicPartition))
-                  accumulatedThrottledSize += bytesAvailable
-                else
+                if (!replicaManager.shouldLeaderThrottle(quota, partition, fetchMetadata.replicaId))
                   accumulatedSize += bytesAvailable
               }
             }
+
+            // Case H: If truncation has caused diverging epoch while this request was in purgatory, return to trigger truncation
+            fetchStatus.fetchInfo.lastFetchedEpoch.ifPresent { fetchEpoch =>
+              val epochEndOffset = partition.lastOffsetForLeaderEpoch(fetchLeaderEpoch, fetchEpoch, fetchOnlyFromLeader = false)
+              if (epochEndOffset.errorCode != Errors.NONE.code()
+                  || epochEndOffset.endOffset == UNDEFINED_EPOCH_OFFSET
+                  || epochEndOffset.leaderEpoch == UNDEFINED_EPOCH) {
+                debug(s"Could not obtain last offset for leader epoch for partition $topicPartition, epochEndOffset=$epochEndOffset.")
+                return forceComplete()
+              } else if (epochEndOffset.leaderEpoch < fetchEpoch || epochEndOffset.endOffset < fetchStatus.fetchInfo.fetchOffset) {
+                debug(s"Satisfying fetch $fetchMetadata since it has diverging epoch requiring truncation for partition " +
+                  s"$topicPartition epochEndOffset=$epochEndOffset fetchEpoch=$fetchEpoch fetchOffset=${fetchStatus.fetchInfo.fetchOffset}.")
+                return forceComplete()
+              }
+            }
           }
         } catch {
-          case _: KafkaStorageException => // Case E
-            debug("Partition %s is in an offline log directory, satisfy %s immediately".format(topicPartition, fetchMetadata))
+          case _: NotLeaderOrFollowerException =>  // Case A or Case B
+            debug(s"Broker is no longer the leader or follower of $topicPartition, satisfy $fetchMetadata immediately")
+            return forceComplete()
+          case _: UnknownTopicOrPartitionException => // Case C
+            debug(s"Broker no longer knows of partition $topicPartition, satisfy $fetchMetadata immediately")
             return forceComplete()
-          case _: UnknownTopicOrPartitionException => // Case B
-            debug("Broker no longer know of %s, satisfy %s immediately".format(topicPartition, fetchMetadata))
+          case _: KafkaStorageException => // Case D
+            debug(s"Partition $topicPartition is in an offline log directory, satisfy $fetchMetadata immediately")
             return forceComplete()
-          case _: NotLeaderForPartitionException =>  // Case A
-            debug("Broker is no longer the leader of %s, satisfy %s immediately".format(topicPartition, fetchMetadata))
+          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 D
-    if (accumulatedSize >= fetchMetadata.fetchMinBytes
-      || ((accumulatedSize + accumulatedThrottledSize) >= fetchMetadata.fetchMinBytes && !quota.isQuotaExceeded()))
-      forceComplete()
+    // Case G
+    if (accumulatedSize >= fetchMetadata.fetchMinBytes)
+       forceComplete()
     else
       false
   }
 
-  override def onExpiration() {
+  override def onExpiration(): Unit = {
     if (fetchMetadata.isFromFollower)
       DelayedFetchMetrics.followerExpiredRequestMeter.mark()
     else
@@ -148,20 +171,31 @@ 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,
-      readOnlyCommitted = fetchMetadata.fetchOnlyCommitted,
+      fetchIsolation = fetchMetadata.fetchIsolation,
       fetchMaxBytes = fetchMetadata.fetchMaxBytes,
       hardMaxBytesLimit = fetchMetadata.hardMaxBytesLimit,
       readPartitionInfo = fetchMetadata.fetchPartitionStatus.map { case (tp, status) => tp -> status.fetchInfo },
-      quota = quota,
-      isolationLevel = isolationLevel)
+      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)
+      val isReassignmentFetch = fetchMetadata.isFromFollower &&
+        replicaManager.isAddingReplica(tp, fetchMetadata.replicaId)
+
+      tp -> FetchPartitionData(
+        result.error,
+        result.highWatermark,
+        result.leaderLogStartOffset,
+        result.info.records,
+        result.divergingEpoch,
+        result.lastStableOffset,
+        result.info.abortedTransactions,
+        result.preferredReadReplica,
+        isReassignmentFetch)
     }
 
     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..cf522ab8f6505
--- /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: Seq[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: Seq[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(): Unit = {
+    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 894d30e2d7c4b..09fd337aa9f4a 100644
--- a/core/src/main/scala/kafka/server/DelayedOperation.scala
+++ b/core/src/main/scala/kafka/server/DelayedOperation.scala
@@ -19,11 +19,10 @@ package kafka.server
 
 import java.util.concurrent._
 import java.util.concurrent.atomic._
-import java.util.concurrent.locks.{Lock, ReentrantLock, ReentrantReadWriteLock}
+import java.util.concurrent.locks.{Lock, ReentrantLock}
 
-import com.yammer.metrics.core.Gauge
 import kafka.metrics.KafkaMetricsGroup
-import kafka.utils.CoreUtils.{inReadLock, inWriteLock}
+import kafka.utils.CoreUtils.inLock
 import kafka.utils._
 import kafka.utils.timer._
 
@@ -42,9 +41,13 @@ import scala.collection.mutable.ListBuffer
  * forceComplete().
  *
  * A subclass of DelayedOperation needs to provide an implementation of both onComplete() and tryComplete().
+ *
+ * Noted that if you add a future delayed operation that calls ReplicaManager.appendRecords() in onComplete()
+ * like DelayedJoin, you must be aware that this operation's onExpiration() needs to call actionQueue.tryCompleteAction().
  */
 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)
   // Visible for testing
@@ -99,20 +102,24 @@ abstract class DelayedOperation(override val delayMs: Long,
   def tryComplete(): Boolean
 
   /**
-   * Thread-safe variant of tryComplete() that attempts completion only if the lock can be acquired
-   * without blocking.
+   * Thread-safe variant of tryComplete() and call extra function if first tryComplete returns false
+   * @param f else function to be executed after first tryComplete returns false
+   * @return result of tryComplete
    */
-  private[server] def maybeTryComplete(): Boolean = {
-    if (lock.tryLock()) {
-      try {
-        tryComplete()
-      } finally {
-        lock.unlock()
-      }
-    } else
-      false
+  private[server] def safeTryCompleteOrElse(f: => Unit): Boolean = inLock(lock) {
+    if (tryComplete()) true
+    else {
+      f
+      // last completion check
+      tryComplete()
+    }
   }
 
+  /**
+   * Thread-safe variant of tryComplete()
+   */
+  private[server] def safeTryComplete(): Boolean = inLock(lock)(tryComplete())
+
   /*
    * run() method defines a task that is executed on timeout
    */
@@ -124,6 +131,8 @@ abstract class DelayedOperation(override val delayMs: Long,
 
 object DelayedOperationPurgatory {
 
+  private val Shards = 512 // Shard the watcher list to reduce lock contention
+
   def apply[T <: DelayedOperation](purgatoryName: String,
                                    brokerId: Int = 0,
                                    purgeInterval: Int = 1000,
@@ -145,11 +154,25 @@ final class DelayedOperationPurgatory[T <: DelayedOperation](purgatoryName: Stri
                                                              reaperEnabled: Boolean = true,
                                                              timerEnabled: Boolean = true)
         extends Logging with KafkaMetricsGroup {
-
   /* a list of operation watching keys */
-  private val watchersForKey = new Pool[Any, Watchers](Some((key: Any) => new Watchers(key)))
+  private class WatcherList {
+    val watchersByKey = new Pool[Any, Watchers](Some((key: Any) => new Watchers(key)))
+
+    val watchersLock = new ReentrantLock()
+
+    /*
+     * Return all the current watcher lists,
+     * note that the returned watchers may be removed from the list by other threads
+     */
+    def allWatchers = {
+      watchersByKey.values
+    }
+  }
 
-  private val removeWatchersLock = new ReentrantReadWriteLock()
+  private val watcherLists = Array.fill[WatcherList](DelayedOperationPurgatory.Shards)(new WatcherList)
+  private def watcherList(key: Any): WatcherList = {
+    watcherLists(Math.abs(key.hashCode() % watcherLists.length))
+  }
 
   // the number of estimated total operations in the purgatory
   private[this] val estimatedTotalOperations = new AtomicInteger(0)
@@ -158,22 +181,8 @@ final class DelayedOperationPurgatory[T <: DelayedOperation](purgatoryName: Stri
   private val expirationReaper = new ExpiredOperationReaper()
 
   private val metricsTags = Map("delayedOperation" -> purgatoryName)
-
-  newGauge(
-    "PurgatorySize",
-    new Gauge[Int] {
-      def value: Int = watched
-    },
-    metricsTags
-  )
-
-  newGauge(
-    "NumDelayedOperations",
-    new Gauge[Int] {
-      def value: Int = delayed
-    },
-    metricsTags
-  )
+  newGauge("PurgatorySize", () => watched, metricsTags)
+  newGauge("NumDelayedOperations", () => numDelayed, metricsTags)
 
   if (reaperEnabled)
     expirationReaper.start()
@@ -194,38 +203,38 @@ final class DelayedOperationPurgatory[T <: DelayedOperation](purgatoryName: Stri
   def tryCompleteElseWatch(operation: T, watchKeys: Seq[Any]): Boolean = {
     assert(watchKeys.nonEmpty, "The watch key list can't be empty")
 
-    // The cost of tryComplete() is typically proportional to the number of keys. Calling
-    // tryComplete() for each key is going to be expensive if there are many keys. Instead,
-    // we do the check in the following way. Call tryComplete(). If the operation is not completed,
-    // we just add the operation to all keys. Then we call tryComplete() again. At this time, if
-    // the operation is still not completed, we are guaranteed that it won't miss any future triggering
-    // event since the operation is already on the watcher list for all keys. This does mean that
-    // if the operation is completed (by another thread) between the two tryComplete() calls, the
-    // operation is unnecessarily added for watch. However, this is a less severe issue since the
-    // expire reaper will clean it up periodically.
-
-    // At this point the only thread that can attempt this operation is this current thread
-    // Hence it is safe to tryComplete() without a lock
-    var isCompletedByMe = operation.tryComplete()
-    if (isCompletedByMe)
-      return true
-
-    var watchCreated = false
-    for(key <- watchKeys) {
-      // If the operation is already completed, stop adding it to the rest of the watcher list.
-      if (operation.isCompleted)
-        return false
-      watchForOperation(key, operation)
-
-      if (!watchCreated) {
-        watchCreated = true
-        estimatedTotalOperations.incrementAndGet()
-      }
-    }
-
-    isCompletedByMe = operation.maybeTryComplete()
-    if (isCompletedByMe)
-      return true
+    // The cost of tryComplete() is typically proportional to the number of keys. Calling tryComplete() for each key is
+    // going to be expensive if there are many keys. Instead, we do the check in the following way through safeTryCompleteOrElse().
+    // If the operation is not completed, we just add the operation to all keys. Then we call tryComplete() again. At
+    // this time, if the operation is still not completed, we are guaranteed that it won't miss any future triggering
+    // event since the operation is already on the watcher list for all keys.
+    //
+    // ==============[story about lock]==============
+    // Through safeTryCompleteOrElse(), we hold the operation's lock while adding the operation to watch list and doing
+    // the tryComplete() check. This is to avoid a potential deadlock between the callers to tryCompleteElseWatch() and
+    // checkAndComplete(). For example, the following deadlock can happen if the lock is only held for the final tryComplete()
+    // 1) thread_a holds readlock of stateLock from TransactionStateManager
+    // 2) thread_a is executing tryCompleteElseWatch()
+    // 3) thread_a adds op to watch list
+    // 4) thread_b requires writelock of stateLock from TransactionStateManager (blocked by thread_a)
+    // 5) thread_c calls checkAndComplete() and holds lock of op
+    // 6) thread_c is waiting readlock of stateLock to complete op (blocked by thread_b)
+    // 7) thread_a is waiting lock of op to call the final tryComplete() (blocked by thread_c)
+    //
+    // Note that even with the current approach, deadlocks could still be introduced. For example,
+    // 1) thread_a calls tryCompleteElseWatch() and gets lock of op
+    // 2) thread_a adds op to watch list
+    // 3) thread_a calls op#tryComplete and tries to require lock_b
+    // 4) thread_b holds lock_b and calls checkAndComplete()
+    // 5) thread_b sees op from watch list
+    // 6) thread_b needs lock of op
+    // To avoid the above scenario, we recommend DelayedOperationPurgatory.checkAndComplete() be called without holding
+    // any exclusive lock. Since DelayedOperationPurgatory.checkAndComplete() completes delayed operations asynchronously,
+    // holding a exclusive lock to make the call is often unnecessary.
+    if (operation.safeTryCompleteOrElse {
+      watchKeys.foreach(key => watchForOperation(key, operation))
+      if (watchKeys.nonEmpty) estimatedTotalOperations.incrementAndGet()
+    }) return true
 
     // if it cannot be completed by now and hence is watched, add to the expire queue also
     if (!operation.isCompleted) {
@@ -247,11 +256,14 @@ final class DelayedOperationPurgatory[T <: DelayedOperation](purgatoryName: Stri
    * @return the number of completed operations during this process
    */
   def checkAndComplete(key: Any): Int = {
-    val watchers = inReadLock(removeWatchersLock) { watchersForKey.get(key) }
-    if(watchers == null)
+    val wl = watcherList(key)
+    val watchers = inLock(wl.watchersLock) { wl.watchersByKey.get(key) }
+    val numCompleted = if (watchers == null)
       0
     else
       watchers.tryCompleteWatched()
+    debug(s"Request key $key unblocked $numCompleted $purgatoryName operations")
+    numCompleted
   }
 
   /**
@@ -259,38 +271,37 @@ final class DelayedOperationPurgatory[T <: DelayedOperation](purgatoryName: Stri
    * on multiple lists, and some of its watched entries may still be in the watch lists
    * even when it has been completed, this number may be larger than the number of real operations watched
    */
-  def watched: Int = allWatchers.map(_.countWatched).sum
+  def watched: Int = {
+    watcherLists.foldLeft(0) { case (sum, watcherList) => sum + watcherList.allWatchers.map(_.countWatched).sum }
+  }
 
   /**
    * 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
     */
   def cancelForKey(key: Any): List[T] = {
-    inWriteLock(removeWatchersLock) {
-      val watchers = watchersForKey.remove(key)
+    val wl = watcherList(key)
+    inLock(wl.watchersLock) {
+      val watchers = wl.watchersByKey.remove(key)
       if (watchers != null)
         watchers.cancel()
       else
         Nil
     }
   }
-  /*
-   * Return all the current watcher lists,
-   * note that the returned watchers may be removed from the list by other threads
-   */
-  private def allWatchers = inReadLock(removeWatchersLock) { watchersForKey.values }
 
   /*
    * 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) {
-    inReadLock(removeWatchersLock) {
-      val watcher = watchersForKey.getAndMaybePut(key)
+  private def watchForOperation(key: Any, operation: T): Unit = {
+    val wl = watcherList(key)
+    inLock(wl.watchersLock) {
+      val watcher = wl.watchersByKey.getAndMaybePut(key)
       watcher.watch(operation)
     }
   }
@@ -298,14 +309,15 @@ 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) {
-    inWriteLock(removeWatchersLock) {
+  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
-      if (watchersForKey.get(key) != watchers)
+      if (wl.watchersByKey.get(key) != watchers)
         return
 
       if (watchers != null && watchers.isEmpty) {
-        watchersForKey.remove(key)
+        wl.watchersByKey.remove(key)
       }
     }
   }
@@ -313,10 +325,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)
   }
 
   /**
@@ -331,7 +345,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)
     }
 
@@ -345,7 +359,7 @@ final class DelayedOperationPurgatory[T <: DelayedOperation](purgatoryName: Stri
         if (curr.isCompleted) {
           // another thread has completed this operation, just remove it
           iter.remove()
-        } else if (curr.maybeTryComplete()) {
+        } else if (curr.safeTryComplete()) {
           iter.remove()
           completed += 1
         }
@@ -389,19 +403,21 @@ 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 = allWatchers.map(_.purgeCompleted()).sum
+      val purged = watcherLists.foldLeft(0) {
+        case (sum, watcherList) => sum + watcherList.allWatchers.map(_.purgeCompleted()).sum
+      }
       debug("Purged %d elements from watch lists.".format(purged))
     }
   }
@@ -413,7 +429,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 718ed241b20bc..964da379778de 100644
--- a/core/src/main/scala/kafka/server/DelayedProduce.scala
+++ b/core/src/main/scala/kafka/server/DelayedProduce.scala
@@ -17,16 +17,15 @@
 
 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.Implicits._
 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._
@@ -34,8 +33,8 @@ import scala.collection._
 case class ProducePartitionStatus(requiredOffset: Long, responseStatus: PartitionResponse) {
   @volatile var acksPending = false
 
-  override def toString = "[acksPending: %b, error: %d, startOffset: %d, requiredOffset: %d]"
-    .format(acksPending, responseStatus.error.code, responseStatus.baseOffset, requiredOffset)
+  override def toString = s"[acksPending: $acksPending, error: ${responseStatus.error.code}, " +
+    s"startOffset: ${responseStatus.baseOffset}, requiredOffset: $requiredOffset]"
 }
 
 /**
@@ -44,8 +43,7 @@ case class ProducePartitionStatus(requiredOffset: Long, responseStatus: Partitio
 case class ProduceMetadata(produceRequiredAcks: Short,
                            produceStatus: Map[TopicPartition, ProducePartitionStatus]) {
 
-  override def toString = "[requiredAcks: %d, partitionStatus: %s]"
-    .format(produceRequiredAcks, produceStatus)
+  override def toString = s"[requiredAcks: $produceRequiredAcks, partitionStatus: $produceStatus]"
 }
 
 /**
@@ -60,7 +58,7 @@ class DelayedProduce(delayMs: Long,
   extends DelayedOperation(delayMs, lockOpt) {
 
   // first update the acks pending variable according to the error code
-  produceMetadata.produceStatus.foreach { case (topicPartition, status) =>
+  produceMetadata.produceStatus.forKeyValue { (topicPartition, status) =>
     if (status.responseStatus.error == Errors.NONE) {
       // Timeout error state will be cleared when required acks are received
       status.acksPending = true
@@ -69,7 +67,7 @@ class DelayedProduce(delayMs: Long,
       status.acksPending = false
     }
 
-    trace("Initial partition status for %s is %s".format(topicPartition, status))
+    trace(s"Initial partition status for $topicPartition is $status")
   }
 
   /**
@@ -84,20 +82,19 @@ class DelayedProduce(delayMs: Long,
    */
   override def tryComplete(): Boolean = {
     // check for each partition if it still has pending acks
-    produceMetadata.produceStatus.foreach { case (topicPartition, status) =>
+    produceMetadata.produceStatus.forKeyValue { (topicPartition, status) =>
       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) 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
@@ -113,9 +110,10 @@ class DelayedProduce(delayMs: Long,
       false
   }
 
-  override def onExpiration() {
-    produceMetadata.produceStatus.foreach { case (topicPartition, status) =>
+  override def onExpiration(): Unit = {
+    produceMetadata.produceStatus.forKeyValue { (topicPartition, status) =>
       if (status.acksPending) {
+        debug(s"Expiring produce request for partition $topicPartition with status $status")
         DelayedProduceMetrics.recordExpiration(topicPartition)
       }
     }
@@ -124,8 +122,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 +139,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
new file mode 100644
index 0000000000000..536a296383aac
--- /dev/null
+++ b/core/src/main/scala/kafka/server/DelegationTokenManager.scala
@@ -0,0 +1,512 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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.nio.ByteBuffer
+import java.nio.charset.StandardCharsets
+import java.security.InvalidKeyException
+import java.util.Base64
+
+import javax.crypto.spec.SecretKeySpec
+import javax.crypto.{Mac, SecretKey}
+import kafka.common.{NotificationHandler, ZkNodeChangeNotificationListener}
+import kafka.metrics.KafkaMetricsGroup
+import kafka.utils.{CoreUtils, Json, Logging}
+import kafka.zk.{DelegationTokenChangeNotificationSequenceZNode, DelegationTokenChangeNotificationZNode, DelegationTokensZNode, KafkaZkClient}
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.security.auth.KafkaPrincipal
+import org.apache.kafka.common.security.scram.internals.{ScramFormatter, ScramMechanism}
+import org.apache.kafka.common.security.scram.ScramCredential
+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.{Sanitizer, SecurityUtils, Time}
+
+import scala.jdk.CollectionConverters._
+import scala.collection.mutable
+
+object DelegationTokenManager {
+  val DefaultHmacAlgorithm = "HmacSHA512"
+  val OwnerKey ="owner"
+  val RenewersKey = "renewers"
+  val IssueTimestampKey = "issueTimestamp"
+  val MaxTimestampKey = "maxTimestamp"
+  val ExpiryTimestampKey = "expiryTimestamp"
+  val TokenIdKey = "tokenId"
+  val VersionKey = "version"
+  val CurrentVersion = 1
+  val ErrorTimestamp = -1
+
+  /**
+   *
+   * @param tokenId
+   * @param secretKey
+   * @return
+   */
+  def createHmac(tokenId: String, secretKey: String) : Array[Byte] = {
+    createHmac(tokenId, createSecretKey(secretKey.getBytes(StandardCharsets.UTF_8)))
+  }
+
+  /**
+   * Convert the byte[] to a secret key
+   * @param keybytes the byte[] to create the secret key from
+   * @return the secret key
+   */
+  def createSecretKey(keybytes: Array[Byte]) : SecretKey = {
+    new SecretKeySpec(keybytes, DefaultHmacAlgorithm)
+  }
+
+  /**
+   *
+   *
+   * @param tokenId
+   * @param secretKey
+   * @return
+   */
+  def createBase64HMAC(tokenId: String, secretKey: SecretKey) : String = {
+    val hmac = createHmac(tokenId, secretKey)
+    Base64.getEncoder.encodeToString(hmac)
+  }
+
+  /**
+   * Compute HMAC of the identifier using the secret key
+   * @param tokenId the bytes of the identifier
+   * @param secretKey  the secret key
+   * @return  String of the generated hmac
+   */
+  def createHmac(tokenId: String, secretKey: SecretKey) : Array[Byte] = {
+    val mac =  Mac.getInstance(DefaultHmacAlgorithm)
+    try
+      mac.init(secretKey)
+    catch {
+      case ike: InvalidKeyException => throw new IllegalArgumentException("Invalid key to HMAC computation", ike);
+    }
+    mac.doFinal(tokenId.getBytes(StandardCharsets.UTF_8))
+  }
+
+  def toJsonCompatibleMap(token: DelegationToken):  Map[String, Any] = {
+    val tokenInfo = token.tokenInfo
+    val tokenInfoMap = mutable.Map[String, Any]()
+    tokenInfoMap(VersionKey) = CurrentVersion
+    tokenInfoMap(OwnerKey) = Sanitizer.sanitize(tokenInfo.ownerAsString)
+    tokenInfoMap(RenewersKey) = tokenInfo.renewersAsString.asScala.map(e => Sanitizer.sanitize(e)).asJava
+    tokenInfoMap(IssueTimestampKey) = tokenInfo.issueTimestamp
+    tokenInfoMap(MaxTimestampKey) = tokenInfo.maxTimestamp
+    tokenInfoMap(ExpiryTimestampKey) = tokenInfo.expiryTimestamp
+    tokenInfoMap(TokenIdKey) = tokenInfo.tokenId()
+    tokenInfoMap.toMap
+  }
+
+  def fromBytes(bytes: Array[Byte]): Option[TokenInformation] = {
+    if (bytes == null || bytes.isEmpty)
+      return None
+
+    Json.parseBytes(bytes) match {
+      case Some(js) =>
+        val mainJs = js.asJsonObject
+        require(mainJs(VersionKey).to[Int] == CurrentVersion)
+        val owner = SecurityUtils.parseKafkaPrincipal(Sanitizer.desanitize(mainJs(OwnerKey).to[String]))
+        val renewerStr = mainJs(RenewersKey).to[Seq[String]]
+        val renewers = renewerStr.map(Sanitizer.desanitize(_)).map(SecurityUtils.parseKafkaPrincipal(_))
+        val issueTimestamp = mainJs(IssueTimestampKey).to[Long]
+        val expiryTimestamp = mainJs(ExpiryTimestampKey).to[Long]
+        val maxTimestamp = mainJs(MaxTimestampKey).to[Long]
+        val tokenId = mainJs(TokenIdKey).to[String]
+
+        val tokenInfo = new TokenInformation(tokenId, owner, renewers.asJava,
+          issueTimestamp, maxTimestamp, expiryTimestamp)
+
+        Some(tokenInfo)
+      case None =>
+        None
+    }
+  }
+
+  def filterToken(requestedPrincipal: KafkaPrincipal, owners : Option[List[KafkaPrincipal]], token: TokenInformation, authorizeToken: String => Boolean) : Boolean = {
+
+    val allow =
+    //exclude tokens which are not requested
+      if (!owners.isEmpty && !owners.get.exists(owner => token.ownerOrRenewer(owner))) {
+        false
+        //Owners and the renewers can describe their own tokens
+      } else if (token.ownerOrRenewer(requestedPrincipal)) {
+        true
+        // Check permission for non-owned tokens
+      } else if ((authorizeToken(token.tokenId))) {
+        true
+      }
+      else {
+        false
+      }
+
+    allow
+  }
+}
+
+class DelegationTokenManager(val config: KafkaConfig,
+                             val tokenCache: DelegationTokenCache,
+                             val time: Time,
+                             val zkClient: KafkaZkClient) extends Logging with KafkaMetricsGroup {
+  this.logIdent = s"[Token Manager on Broker ${config.brokerId}]: "
+
+  import DelegationTokenManager._
+
+  type CreateResponseCallback = CreateTokenResult => Unit
+  type RenewResponseCallback = (Errors, Long) => Unit
+  type ExpireResponseCallback = (Errors, Long) => Unit
+  type DescribeResponseCallback = (Errors, List[DelegationToken]) => Unit
+
+  val secretKey = {
+    val keyBytes =  if (config.tokenAuthEnabled) config.delegationTokenSecretKey.value.getBytes(StandardCharsets.UTF_8) else null
+    if (keyBytes == null || keyBytes.length == 0) null
+    else
+      createSecretKey(keyBytes)
+  }
+
+  val tokenMaxLifetime: Long = config.delegationTokenMaxLifeMs
+  val defaultTokenRenewTime: Long = config.delegationTokenExpiryTimeMs
+  val tokenRemoverScanInterval: Long = config.delegationTokenExpiryCheckIntervalMs
+  private val lock = new Object()
+  private var tokenChangeListener: ZkNodeChangeNotificationListener = null
+
+  def startup() = {
+    if (config.tokenAuthEnabled) {
+      zkClient.createDelegationTokenPaths()
+      loadCache()
+      tokenChangeListener = new ZkNodeChangeNotificationListener(zkClient, DelegationTokenChangeNotificationZNode.path, DelegationTokenChangeNotificationSequenceZNode.SequenceNumberPrefix, TokenChangedNotificationHandler)
+      tokenChangeListener.init()
+    }
+  }
+
+  def shutdown() = {
+    if (config.tokenAuthEnabled) {
+      if (tokenChangeListener != null) tokenChangeListener.close()
+    }
+  }
+
+  private def loadCache(): Unit = {
+    lock.synchronized {
+      val tokens = zkClient.getChildren(DelegationTokensZNode.path)
+      info(s"Loading the token cache. Total token count: ${tokens.size}")
+      for (tokenId <- tokens) {
+        try {
+          getTokenFromZk(tokenId) match {
+            case Some(token) => updateCache(token)
+            case None =>
+          }
+        } catch {
+          case ex: Throwable => error(s"Error while getting Token for tokenId: $tokenId", ex)
+        }
+      }
+    }
+  }
+
+  private def getTokenFromZk(tokenId: String): Option[DelegationToken] = {
+    zkClient.getDelegationTokenInfo(tokenId) match {
+      case Some(tokenInformation) => {
+        val hmac = createHmac(tokenId, secretKey)
+        Some(new DelegationToken(tokenInformation, hmac))
+      }
+      case None =>
+        None
+    }
+  }
+
+  /**
+   *
+   * @param token
+   */
+  private def updateCache(token: DelegationToken): Unit = {
+    val hmacString = token.hmacAsBase64String
+    val scramCredentialMap =  prepareScramCredentials(hmacString)
+    tokenCache.updateCache(token, scramCredentialMap.asJava)
+  }
+  /**
+   * @param hmacString
+   */
+  private def prepareScramCredentials(hmacString: String) : Map[String, ScramCredential] = {
+    val scramCredentialMap = mutable.Map[String, ScramCredential]()
+
+    def scramCredential(mechanism: ScramMechanism): ScramCredential = {
+      new ScramFormatter(mechanism).generateCredential(hmacString, mechanism.minIterations)
+    }
+
+    for (mechanism <- ScramMechanism.values)
+      scramCredentialMap(mechanism.mechanismName) = scramCredential(mechanism)
+
+    scramCredentialMap.toMap
+  }
+
+  /**
+   *
+   * @param owner
+   * @param renewers
+   * @param maxLifeTimeMs
+   * @param responseCallback
+   */
+  def createToken(owner: KafkaPrincipal,
+                  renewers: List[KafkaPrincipal],
+                  maxLifeTimeMs: Long,
+                  responseCallback: CreateResponseCallback): Unit = {
+
+    if (!config.tokenAuthEnabled) {
+      responseCallback(CreateTokenResult(-1, -1, -1, "", Array[Byte](), Errors.DELEGATION_TOKEN_AUTH_DISABLED))
+    } else {
+      lock.synchronized {
+        val tokenId = CoreUtils.generateUuidAsBase64()
+
+        val issueTimeStamp = time.milliseconds
+        val maxLifeTime = if (maxLifeTimeMs <= 0) tokenMaxLifetime else Math.min(maxLifeTimeMs, tokenMaxLifetime)
+        val maxLifeTimeStamp = issueTimeStamp + maxLifeTime
+        val expiryTimeStamp = Math.min(maxLifeTimeStamp, issueTimeStamp + defaultTokenRenewTime)
+
+        val tokenInfo = new TokenInformation(tokenId, owner, renewers.asJava, issueTimeStamp, maxLifeTimeStamp, expiryTimeStamp)
+
+        val hmac = createHmac(tokenId, secretKey)
+        val token = new DelegationToken(tokenInfo, hmac)
+        updateToken(token)
+        info(s"Created a delegation token: $tokenId for owner: $owner")
+        responseCallback(CreateTokenResult(issueTimeStamp, expiryTimeStamp, maxLifeTimeStamp, tokenId, hmac, Errors.NONE))
+      }
+    }
+  }
+
+  /**
+   *
+   * @param principal
+   * @param hmac
+   * @param renewLifeTimeMs
+   * @param renewCallback
+   */
+  def renewToken(principal: KafkaPrincipal,
+                 hmac: ByteBuffer,
+                 renewLifeTimeMs: Long,
+                 renewCallback: RenewResponseCallback): Unit = {
+
+    if (!config.tokenAuthEnabled) {
+      renewCallback(Errors.DELEGATION_TOKEN_AUTH_DISABLED, -1)
+    } else {
+      lock.synchronized  {
+        getToken(hmac) match {
+          case Some(token) => {
+            val now = time.milliseconds
+            val tokenInfo =  token.tokenInfo
+
+            if (!allowedToRenew(principal, tokenInfo)) {
+              renewCallback(Errors.DELEGATION_TOKEN_OWNER_MISMATCH, -1)
+            } else if (tokenInfo.maxTimestamp < now || tokenInfo.expiryTimestamp < now) {
+              renewCallback(Errors.DELEGATION_TOKEN_EXPIRED, -1)
+            } else {
+              val renewLifeTime = if (renewLifeTimeMs < 0) defaultTokenRenewTime else renewLifeTimeMs
+              val renewTimeStamp = now + renewLifeTime
+              val expiryTimeStamp = Math.min(tokenInfo.maxTimestamp, renewTimeStamp)
+              tokenInfo.setExpiryTimestamp(expiryTimeStamp)
+
+              updateToken(token)
+              info(s"Delegation token renewed for token: ${tokenInfo.tokenId} for owner: ${tokenInfo.owner}")
+              renewCallback(Errors.NONE, expiryTimeStamp)
+            }
+          }
+          case None => renewCallback(Errors.DELEGATION_TOKEN_NOT_FOUND, -1)
+        }
+      }
+    }
+  }
+
+  /**
+   * @param token
+   */
+  private def updateToken(token: DelegationToken): Unit = {
+    zkClient.setOrCreateDelegationToken(token)
+    updateCache(token)
+    zkClient.createTokenChangeNotification(token.tokenInfo.tokenId())
+  }
+
+  /**
+   *
+   * @param hmac
+   * @return
+   */
+  private def getToken(hmac: ByteBuffer): Option[DelegationToken] = {
+    try {
+      val byteArray = new Array[Byte](hmac.remaining)
+      hmac.get(byteArray)
+      val base64Pwd = Base64.getEncoder.encodeToString(byteArray)
+      val tokenInfo = tokenCache.tokenForHmac(base64Pwd)
+      if (tokenInfo == null) None else Some(new DelegationToken(tokenInfo, byteArray))
+    } catch {
+      case e: Exception =>
+        error("Exception while getting token for hmac", e)
+        None
+    }
+  }
+
+  /**
+   *
+   * @param principal
+   * @param tokenInfo
+   * @return
+   */
+  private def allowedToRenew(principal: KafkaPrincipal, tokenInfo: TokenInformation): Boolean = {
+    if (principal.equals(tokenInfo.owner) || tokenInfo.renewers.asScala.toList.contains(principal)) true else false
+  }
+
+  /**
+   *
+   * @param tokenId
+   * @return
+   */
+  def getToken(tokenId: String): Option[DelegationToken] = {
+    val tokenInfo = tokenCache.token(tokenId)
+    if (tokenInfo != null) Some(getToken(tokenInfo)) else None
+  }
+
+  /**
+   *
+   * @param tokenInfo
+   * @return
+   */
+  private def getToken(tokenInfo: TokenInformation): DelegationToken = {
+    val hmac = createHmac(tokenInfo.tokenId, secretKey)
+    new DelegationToken(tokenInfo, hmac)
+  }
+
+  /**
+   *
+   * @param principal
+   * @param hmac
+   * @param expireLifeTimeMs
+   * @param expireResponseCallback
+   */
+  def expireToken(principal: KafkaPrincipal,
+                  hmac: ByteBuffer,
+                  expireLifeTimeMs: Long,
+                  expireResponseCallback: ExpireResponseCallback): Unit = {
+
+    if (!config.tokenAuthEnabled) {
+      expireResponseCallback(Errors.DELEGATION_TOKEN_AUTH_DISABLED, -1)
+    } else {
+      lock.synchronized  {
+        getToken(hmac) match {
+          case Some(token) =>  {
+            val tokenInfo =  token.tokenInfo
+            val now = time.milliseconds
+
+            if (!allowedToRenew(principal, tokenInfo)) {
+              expireResponseCallback(Errors.DELEGATION_TOKEN_OWNER_MISMATCH, -1)
+            } else if (tokenInfo.maxTimestamp < now || tokenInfo.expiryTimestamp < now) {
+              expireResponseCallback(Errors.DELEGATION_TOKEN_EXPIRED, -1)
+            } else if (expireLifeTimeMs < 0) { //expire immediately
+              removeToken(tokenInfo.tokenId)
+              info(s"Token expired for token: ${tokenInfo.tokenId} for owner: ${tokenInfo.owner}")
+              expireResponseCallback(Errors.NONE, now)
+            } else {
+              //set expiry time stamp
+              val expiryTimeStamp = Math.min(tokenInfo.maxTimestamp, now + expireLifeTimeMs)
+              tokenInfo.setExpiryTimestamp(expiryTimeStamp)
+
+              updateToken(token)
+              info(s"Updated expiry time for token: ${tokenInfo.tokenId} for owner: ${tokenInfo.owner}")
+              expireResponseCallback(Errors.NONE, expiryTimeStamp)
+            }
+          }
+          case None => expireResponseCallback(Errors.DELEGATION_TOKEN_NOT_FOUND, -1)
+        }
+      }
+    }
+  }
+
+  /**
+   *
+   * @param tokenId
+   */
+  private def removeToken(tokenId: String): Unit = {
+    zkClient.deleteDelegationToken(tokenId)
+    removeCache(tokenId)
+    zkClient.createTokenChangeNotification(tokenId)
+  }
+
+  /**
+   *
+   * @param tokenId
+   */
+  private def removeCache(tokenId: String): Unit = {
+    tokenCache.removeCache(tokenId)
+  }
+
+  /**
+   *
+   * @return
+   */
+  def expireTokens(): Unit = {
+    lock.synchronized {
+      for (tokenInfo <- getAllTokenInformation) {
+        val now = time.milliseconds
+        if (tokenInfo.maxTimestamp < now || tokenInfo.expiryTimestamp < now) {
+          info(s"Delegation token expired for token: ${tokenInfo.tokenId} for owner: ${tokenInfo.owner}")
+          removeToken(tokenInfo.tokenId)
+        }
+      }
+    }
+  }
+
+  def getAllTokenInformation: List[TokenInformation] = tokenCache.tokens.asScala.toList
+
+  def getTokens(filterToken: TokenInformation => Boolean): List[DelegationToken] = {
+    getAllTokenInformation.filter(filterToken).map(token => getToken(token))
+  }
+
+  object TokenChangedNotificationHandler extends NotificationHandler {
+    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")
+        getTokenFromZk(tokenId) match {
+          case Some(token) => updateCache(token)
+          case None => removeCache(tokenId)
+        }
+      }
+    }
+  }
+
+}
+
+case class CreateTokenResult(issueTimestamp: Long,
+                             expiryTimestamp: Long,
+                             maxTimestamp: Long,
+                             tokenId: String,
+                             hmac: Array[Byte],
+                             error: Errors) {
+
+  override def equals(other: Any): Boolean = {
+    other match {
+      case that: CreateTokenResult =>
+        error.equals(that.error) &&
+          tokenId.equals(that.tokenId) &&
+          issueTimestamp.equals(that.issueTimestamp) &&
+          expiryTimestamp.equals(that.expiryTimestamp) &&
+          maxTimestamp.equals(that.maxTimestamp) &&
+          (hmac sameElements that.hmac)
+      case _ => false
+    }
+  }
+
+  override def hashCode(): Int = {
+    val fields = Seq(issueTimestamp, expiryTimestamp, maxTimestamp, tokenId, hmac, error)
+    fields.map(_.hashCode()).foldLeft(0)((a, b) => 31 * a + b)
+  }
+}
diff --git a/core/src/main/scala/kafka/server/DynamicBrokerConfig.scala b/core/src/main/scala/kafka/server/DynamicBrokerConfig.scala
new file mode 100755
index 0000000000000..7563b38a5d6be
--- /dev/null
+++ b/core/src/main/scala/kafka/server/DynamicBrokerConfig.scala
@@ -0,0 +1,948 @@
+/**
+  * Licensed to the Apache Software Foundation (ASF) under one or more
+  * contributor license agreements.  See the NOTICE file distributed with
+  * this work for additional information regarding copyright ownership.
+  * The ASF licenses this file to You under the Apache License, Version 2.0
+  * (the "License"); you may not use this file except in compliance with
+  * the License.  You may obtain a copy of the License at
+  *
+  * http://www.apache.org/licenses/LICENSE-2.0
+  *
+  * Unless required by applicable law or agreed to in writing, software
+  * distributed under the License is distributed on an "AS IS" BASIS,
+  * WITHOUT WARRANTIES 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
+import java.util.{Collections, Properties}
+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.utils.Implicits._
+import kafka.zk.{AdminZkClient, KafkaZkClient}
+import org.apache.kafka.common.Reconfigurable
+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}
+import org.apache.kafka.common.security.authenticator.LoginManager
+import org.apache.kafka.common.utils.Utils
+
+import scala.collection._
+import scala.jdk.CollectionConverters._
+
+/**
+  * Dynamic broker configurations are stored in ZooKeeper and may be defined at two levels:
+  * 
                
+  *   
              • Per-broker configs persisted at /configs/brokers/{brokerId}: These can be described/altered
+  *       using AdminClient using the resource name brokerId.
                • 
+  *   
              • Cluster-wide defaults persisted at /configs/brokers/<default>: These can be described/altered
+  *       using AdminClient using an empty resource name.
                • 
+  * 
              
+  * The order of precedence for broker configs is:
+  * 
                
+  *   
              1. DYNAMIC_BROKER_CONFIG: stored in ZK at /configs/brokers/{brokerId}
                2. 
+  *   
              2. DYNAMIC_DEFAULT_BROKER_CONFIG: stored in ZK at /configs/brokers/<default>
                3. 
+  *   
              3. STATIC_BROKER_CONFIG: properties that broker is started up with, typically from server.properties file
                4. 
+  *   
              4. DEFAULT_CONFIG: Default configs defined in KafkaConfig
                5. 
+  * 
              
+  * Log configs use topic config overrides if defined and fallback to broker defaults using the order of precedence above.
+  * Topic config overrides may use a different config name from the default broker config.
+  * See [[kafka.log.LogConfig#TopicConfigSynonyms]] for the mapping.
+  * 
              
+  * AdminClient returns all config synonyms in the order of precedence when configs are described with
+  * includeSynonyms. In addition to configs that may be defined with the same name at different levels,
+  * some configs have additional synonyms.
+  * 
              
+  * 
                
+  *   
              • Listener configs may be defined using the prefix listener.name.{listenerName}.{configName}. These may be
+  *       configured as dynamic or static broker configs. Listener configs have higher precedence than the base configs
+  *       that don't specify the listener name. Listeners without a listener config use the base config. Base configs
+  *       may be defined only as STATIC_BROKER_CONFIG or DEFAULT_CONFIG and cannot be updated dynamically.
              • 
+  *   
              • Some configs may be defined using multiple properties. For example, log.roll.ms and
+  *       log.roll.hours refer to the same config that may be defined in milliseconds or hours. The order of
+  *       precedence of these synonyms is described in the docs of these configs in [[kafka.server.KafkaConfig]].
                • 
+  * 
              
+  *
+  */
+object DynamicBrokerConfig {
+
+  private[server] val DynamicSecurityConfigs = SslConfigs.RECONFIGURABLE_CONFIGS.asScala
+
+  val AllDynamicConfigs = DynamicSecurityConfigs ++
+    LogCleaner.ReconfigurableConfigs ++
+    DynamicLogConfig.ReconfigurableConfigs ++
+    DynamicThreadPool.ReconfigurableConfigs ++
+    Set(KafkaConfig.MetricReporterClassesProp) ++
+    DynamicListenerConfig.ReconfigurableConfigs ++
+    SocketServer.ReconfigurableConfigs
+
+  private val ClusterLevelListenerConfigs = Set(KafkaConfig.MaxConnectionsProp, KafkaConfig.MaxConnectionCreationRateProp)
+  private val PerBrokerConfigs = (DynamicSecurityConfigs ++ DynamicListenerConfig.ReconfigurableConfigs).diff(
+    ClusterLevelListenerConfigs)
+  private val ListenerMechanismConfigs = Set(KafkaConfig.SaslJaasConfigProp,
+    KafkaConfig.SaslLoginCallbackHandlerClassProp,
+    KafkaConfig.SaslLoginClassProp,
+    KafkaConfig.SaslServerCallbackHandlerClassProp,
+    KafkaConfig.ConnectionsMaxReauthMsProp)
+
+  private val ReloadableFileConfigs = Set(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG, SslConfigs.SSL_TRUSTSTORE_LOCATION_CONFIG)
+
+  val ListenerConfigRegex = """listener\.name\.[^.]*\.(.*)""".r
+
+  private val DynamicPasswordConfigs = {
+    val passwordConfigs = KafkaConfig.configKeys.filter(_._2.`type` == ConfigDef.Type.PASSWORD).keySet
+    AllDynamicConfigs.intersect(passwordConfigs)
+  }
+
+  def isPasswordConfig(name: String): Boolean = DynamicBrokerConfig.DynamicPasswordConfigs.exists(name.endsWith)
+
+  def brokerConfigSynonyms(name: String, matchListenerOverride: Boolean): List[String] = {
+    name match {
+      case KafkaConfig.LogRollTimeMillisProp | KafkaConfig.LogRollTimeHoursProp =>
+        List(KafkaConfig.LogRollTimeMillisProp, KafkaConfig.LogRollTimeHoursProp)
+      case KafkaConfig.LogRollTimeJitterMillisProp | KafkaConfig.LogRollTimeJitterHoursProp =>
+        List(KafkaConfig.LogRollTimeJitterMillisProp, KafkaConfig.LogRollTimeJitterHoursProp)
+      case KafkaConfig.LogFlushIntervalMsProp => // LogFlushSchedulerIntervalMsProp is used as default
+        List(KafkaConfig.LogFlushIntervalMsProp, KafkaConfig.LogFlushSchedulerIntervalMsProp)
+      case KafkaConfig.LogRetentionTimeMillisProp | KafkaConfig.LogRetentionTimeMinutesProp | KafkaConfig.LogRetentionTimeHoursProp =>
+        List(KafkaConfig.LogRetentionTimeMillisProp, KafkaConfig.LogRetentionTimeMinutesProp, KafkaConfig.LogRetentionTimeHoursProp)
+      case ListenerConfigRegex(baseName) if matchListenerOverride =>
+        // `ListenerMechanismConfigs` are specified as listenerPrefix.mechanism.
+        // and other listener configs are specified as listenerPrefix.
+        // Add  as a synonym in both cases.
+        val mechanismConfig = ListenerMechanismConfigs.find(baseName.endsWith)
+        List(name, mechanismConfig.getOrElse(baseName))
+      case _ => List(name)
+    }
+  }
+
+  def validateConfigs(props: Properties, perBrokerConfig: Boolean): Unit = {
+    def checkInvalidProps(invalidPropNames: Set[String], errorMessage: String): Unit = {
+      if (invalidPropNames.nonEmpty)
+        throw new ConfigException(s"$errorMessage: $invalidPropNames")
+    }
+    checkInvalidProps(nonDynamicConfigs(props), "Cannot update these configs dynamically")
+    checkInvalidProps(securityConfigsWithoutListenerPrefix(props),
+      "These security configs can be dynamically updated only per-listener using the listener prefix")
+    validateConfigTypes(props)
+    if (!perBrokerConfig) {
+      checkInvalidProps(perBrokerConfigs(props),
+        "Cannot update these configs at default cluster level, broker id must be specified")
+    }
+  }
+
+  private def perBrokerConfigs(props: Properties): Set[String] = {
+    val configNames = props.asScala.keySet
+    def perBrokerListenerConfig(name: String): Boolean = {
+      name match {
+        case ListenerConfigRegex(baseName) => !ClusterLevelListenerConfigs.contains(baseName)
+        case _ => false
+      }
+    }
+    configNames.intersect(PerBrokerConfigs) ++ configNames.filter(perBrokerListenerConfig)
+  }
+
+  private def nonDynamicConfigs(props: Properties): Set[String] = {
+    props.asScala.keySet.intersect(DynamicConfig.Broker.nonDynamicProps)
+  }
+
+  private def securityConfigsWithoutListenerPrefix(props: Properties): Set[String] = {
+    DynamicSecurityConfigs.filter(props.containsKey)
+  }
+
+  private def validateConfigTypes(props: Properties): Unit = {
+    val baseProps = new Properties
+    props.asScala.foreach {
+      case (ListenerConfigRegex(baseName), v) => baseProps.put(baseName, v)
+      case (k, v) => baseProps.put(k, v)
+    }
+    DynamicConfig.Broker.validate(baseProps)
+  }
+
+  private[server] def addDynamicConfigs(configDef: ConfigDef): Unit = {
+    KafkaConfig.configKeys.forKeyValue { (configName, config) =>
+      if (AllDynamicConfigs.contains(configName)) {
+        configDef.define(config.name, config.`type`, config.defaultValue, config.validator,
+          config.importance, config.documentation, config.group, config.orderInGroup, config.width,
+          config.displayName, config.dependents, config.recommender)
+      }
+    }
+  }
+
+  private[server] def dynamicConfigUpdateModes: util.Map[String, String] = {
+    AllDynamicConfigs.map { name =>
+      val mode = if (PerBrokerConfigs.contains(name)) "per-broker" else "cluster-wide"
+      (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 {
+
+  private[server] val staticBrokerConfigs = ConfigDef.convertToStringMapWithPasswordValues(kafkaConfig.originalsFromThisConfig).asScala
+  private[server] val staticDefaultConfigs = ConfigDef.convertToStringMapWithPasswordValues(KafkaConfig.defaultValues.asJava).asScala
+  private val dynamicBrokerConfigs = mutable.Map[String, String]()
+  private val dynamicDefaultConfigs = mutable.Map[String, String]()
+  private val reconfigurables = mutable.Buffer[Reconfigurable]()
+  private val brokerReconfigurables = mutable.Buffer[BrokerReconfigurable]()
+  private val lock = new ReentrantReadWriteLock
+  private var currentConfig = kafkaConfig
+  private val dynamicConfigPasswordEncoder = maybeCreatePasswordEncoder(kafkaConfig.passwordEncoderSecret)
+
+  private[server] def initialize(zkClient: KafkaZkClient): Unit = {
+    currentConfig = new KafkaConfig(kafkaConfig.props, false, None)
+    val adminZkClient = new AdminZkClient(zkClient)
+    updateDefaultConfig(adminZkClient.fetchEntityConfig(ConfigType.Broker, ConfigEntityName.Default))
+    val props = adminZkClient.fetchEntityConfig(ConfigType.Broker, kafkaConfig.brokerId.toString)
+    val brokerConfig = maybeReEncodePasswords(props, adminZkClient)
+    updateBrokerConfig(kafkaConfig.brokerId, brokerConfig)
+  }
+
+  /**
+   * Clear all cached values. This is used to clear state on broker shutdown to avoid
+   * exceptions in tests when broker is restarted. These fields are re-initialized when
+   * broker starts up.
+   */
+  private[server] def clear(): Unit = {
+    dynamicBrokerConfigs.clear()
+    dynamicDefaultConfigs.clear()
+    reconfigurables.clear()
+    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(kafkaServer.kafkaYammerMetrics)
+    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)
+    addBrokerReconfigurable(new DynamicLogConfig(kafkaServer.logManager, kafkaServer))
+    addBrokerReconfigurable(new DynamicListenerConfig(kafkaServer))
+    addBrokerReconfigurable(kafkaServer.socketServer)
+  }
+
+  def addReconfigurable(reconfigurable: Reconfigurable): Unit = CoreUtils.inWriteLock(lock) {
+    verifyReconfigurableConfigs(reconfigurable.reconfigurableConfigs.asScala)
+    reconfigurables += reconfigurable
+  }
+
+  def addBrokerReconfigurable(reconfigurable: BrokerReconfigurable): Unit = CoreUtils.inWriteLock(lock) {
+    verifyReconfigurableConfigs(reconfigurable.reconfigurableConfigs)
+    brokerReconfigurables += reconfigurable
+  }
+
+  def removeReconfigurable(reconfigurable: Reconfigurable): Unit = CoreUtils.inWriteLock(lock) {
+    reconfigurables -= reconfigurable
+  }
+
+  private def verifyReconfigurableConfigs(configNames: Set[String]): Unit = CoreUtils.inWriteLock(lock) {
+    val nonDynamic = configNames.filter(DynamicConfig.Broker.nonDynamicProps.contains)
+    require(nonDynamic.isEmpty, s"Reconfigurable contains non-dynamic configs $nonDynamic")
+  }
+
+  // Visibility for testing
+  private[server] def currentKafkaConfig: KafkaConfig = CoreUtils.inReadLock(lock) {
+    currentConfig
+  }
+
+  private[server] def currentDynamicBrokerConfigs: Map[String, String] = CoreUtils.inReadLock(lock) {
+    dynamicBrokerConfigs.clone()
+  }
+
+  private[server] def currentDynamicDefaultConfigs: Map[String, String] = CoreUtils.inReadLock(lock) {
+    dynamicDefaultConfigs.clone()
+  }
+
+  private[server] def updateBrokerConfig(brokerId: Int, persistentProps: Properties): Unit = CoreUtils.inWriteLock(lock) {
+    try {
+      val props = fromPersistentProps(persistentProps, perBrokerConfig = true)
+      dynamicBrokerConfigs.clear()
+      dynamicBrokerConfigs ++= props.asScala
+      updateCurrentConfig()
+    } catch {
+      case e: Exception => error(s"Per-broker configs of $brokerId could not be applied: $persistentProps", e)
+    }
+  }
+
+  private[server] def updateDefaultConfig(persistentProps: Properties): Unit = CoreUtils.inWriteLock(lock) {
+    try {
+      val props = fromPersistentProps(persistentProps, perBrokerConfig = false)
+      dynamicDefaultConfigs.clear()
+      dynamicDefaultConfigs ++= props.asScala
+      updateCurrentConfig()
+    } catch {
+      case e: Exception => error(s"Cluster default configs could not be applied: $persistentProps", e)
+    }
+  }
+
+  /**
+   * All config updates through ZooKeeper are triggered through actual changes in values stored in ZooKeeper.
+   * For some configs like SSL keystores and truststores, we also want to reload the store if it was modified
+   * in-place, even though the actual value of the file path and password haven't changed. This scenario alone
+   * is handled here when a config update request using admin client is processed by AdminManager. If any of
+   * the SSL configs have changed, then the update will not be done here, but will be handled later when ZK
+   * changes are processed. At the moment, only listener configs are considered for reloading.
+   */
+  private[server] def reloadUpdatedFilesWithoutConfigChange(newProps: Properties): Unit = CoreUtils.inWriteLock(lock) {
+    reconfigurables
+      .filter(reconfigurable => ReloadableFileConfigs.exists(reconfigurable.reconfigurableConfigs.contains))
+      .foreach {
+        case reconfigurable: ListenerReconfigurable =>
+          val kafkaProps = validatedKafkaProps(newProps, perBrokerConfig = true)
+          val newConfig = new KafkaConfig(kafkaProps.asJava, false, None)
+          processListenerReconfigurable(reconfigurable, newConfig, Collections.emptyMap(), validateOnly = false, reloadOnly = true)
+        case reconfigurable =>
+          trace(s"Files will not be reloaded without config change for $reconfigurable")
+      }
+  }
+
+  private def maybeCreatePasswordEncoder(secret: Option[Password]): Option[PasswordEncoder] = {
+   secret.map { secret =>
+      new PasswordEncoder(secret,
+        kafkaConfig.passwordEncoderKeyFactoryAlgorithm,
+        kafkaConfig.passwordEncoderCipherAlgorithm,
+        kafkaConfig.passwordEncoderKeyLength,
+        kafkaConfig.passwordEncoderIterations)
+    }
+  }
+
+  private def passwordEncoder: PasswordEncoder = {
+    dynamicConfigPasswordEncoder.getOrElse(throw new ConfigException("Password encoder secret not configured"))
+  }
+
+  private[server] def toPersistentProps(configProps: Properties, perBrokerConfig: Boolean): Properties = {
+    val props = configProps.clone().asInstanceOf[Properties]
+
+    def encodePassword(configName: String, value: String): Unit = {
+      if (value != null) {
+        if (!perBrokerConfig)
+          throw new ConfigException("Password config can be defined only at broker level")
+        props.setProperty(configName, passwordEncoder.encode(new Password(value)))
+      }
+    }
+    configProps.asScala.forKeyValue { (name, value) =>
+      if (isPasswordConfig(name))
+        encodePassword(name, value)
+    }
+    props
+  }
+
+  private[server] def fromPersistentProps(persistentProps: Properties,
+                                          perBrokerConfig: Boolean): Properties = {
+    val props = persistentProps.clone().asInstanceOf[Properties]
+
+    // Remove all invalid configs from `props`
+    removeInvalidConfigs(props, perBrokerConfig)
+    def removeInvalidProps(invalidPropNames: Set[String], errorMessage: String): Unit = {
+      if (invalidPropNames.nonEmpty) {
+        invalidPropNames.foreach(props.remove)
+        error(s"$errorMessage: $invalidPropNames")
+      }
+    }
+    removeInvalidProps(nonDynamicConfigs(props), "Non-dynamic configs configured in ZooKeeper will be ignored")
+    removeInvalidProps(securityConfigsWithoutListenerPrefix(props),
+      "Security configs can be dynamically updated only using listener prefix, base configs will be ignored")
+    if (!perBrokerConfig)
+      removeInvalidProps(perBrokerConfigs(props), "Per-broker configs defined at default cluster level will be ignored")
+
+    def decodePassword(configName: String, value: String): Unit = {
+      if (value != null) {
+        try {
+          props.setProperty(configName, passwordEncoder.decode(value).value)
+        } catch {
+          case e: Exception =>
+            error(s"Dynamic password config $configName could not be decoded, ignoring.", e)
+            props.remove(configName)
+        }
+      }
+    }
+
+    props.asScala.forKeyValue { (name, value) =>
+      if (isPasswordConfig(name))
+        decodePassword(name, value)
+    }
+    props
+  }
+
+  // If the secret has changed, password.encoder.old.secret contains the old secret that was used
+  // to encode the configs in ZK. Decode passwords using the old secret and update ZK with values
+  // encoded using the current secret. Ignore any errors during decoding since old secret may not
+  // have been removed during broker restart.
+  private def maybeReEncodePasswords(persistentProps: Properties, adminZkClient: AdminZkClient): Properties = {
+    val props = persistentProps.clone().asInstanceOf[Properties]
+    if (props.asScala.keySet.exists(isPasswordConfig)) {
+      maybeCreatePasswordEncoder(kafkaConfig.passwordEncoderOldSecret).foreach { passwordDecoder =>
+        persistentProps.asScala.forKeyValue { (configName, value) =>
+          if (isPasswordConfig(configName) && value != null) {
+            val decoded = try {
+              Some(passwordDecoder.decode(value).value)
+            } catch {
+              case _: Exception =>
+                debug(s"Dynamic password config $configName could not be decoded using old secret, new secret will be used.")
+                None
+            }
+            decoded.foreach { value => props.put(configName, passwordEncoder.encode(new Password(value))) }
+          }
+        }
+        adminZkClient.changeBrokerConfig(Some(kafkaConfig.brokerId), props)
+      }
+    }
+    props
+  }
+
+  /**
+   * Validate the provided configs `propsOverride` and return the full Kafka configs with
+   * the configured defaults and these overrides.
+   *
+   * Note: The caller must acquire the read or write lock before invoking this method.
+   */
+  private def validatedKafkaProps(propsOverride: Properties, perBrokerConfig: Boolean): Map[String, String] = {
+    val propsResolved = DynamicBrokerConfig.resolveVariableConfigs(propsOverride)
+    validateConfigs(propsResolved, perBrokerConfig)
+    val newProps = mutable.Map[String, String]()
+    newProps ++= staticBrokerConfigs
+    if (perBrokerConfig) {
+      overrideProps(newProps, dynamicDefaultConfigs)
+      overrideProps(newProps, propsResolved.asScala)
+    } else {
+      overrideProps(newProps, propsResolved.asScala)
+      overrideProps(newProps, dynamicBrokerConfigs)
+    }
+    newProps
+  }
+
+  private[server] def validate(props: Properties, perBrokerConfig: Boolean): Unit = CoreUtils.inReadLock(lock) {
+    val newProps = validatedKafkaProps(props, perBrokerConfig)
+    processReconfiguration(newProps, validateOnly = true)
+  }
+
+  private def removeInvalidConfigs(props: Properties, perBrokerConfig: Boolean): Unit = {
+    try {
+      validateConfigTypes(props)
+      props.asScala
+    } catch {
+      case e: Exception =>
+        val invalidProps = props.asScala.filter { case (k, v) =>
+          val props1 = new Properties
+          props1.put(k, v)
+          try {
+            validateConfigTypes(props1)
+            false
+          } catch {
+            case _: Exception => true
+          }
+        }
+        invalidProps.keys.foreach(props.remove)
+        val configSource = if (perBrokerConfig) "broker" else "default cluster"
+        error(s"Dynamic $configSource config contains invalid values: $invalidProps, these configs will be ignored", e)
+    }
+  }
+
+  private[server] def maybeReconfigure(reconfigurable: Reconfigurable, oldConfig: KafkaConfig, newConfig: util.Map[String, _]): Unit = {
+    if (reconfigurable.reconfigurableConfigs.asScala.exists(key => oldConfig.originals.get(key) != newConfig.get(key)))
+      reconfigurable.reconfigure(newConfig)
+  }
+
+  /**
+   * Returns the change in configurations between the new props and current props by returning a
+   * map of the changed configs, as well as the set of deleted keys
+   */
+  private def updatedConfigs(newProps: java.util.Map[String, _], currentProps: java.util.Map[String, _]):
+  (mutable.Map[String, _], Set[String]) = {
+    val changeMap = newProps.asScala.filter {
+      case (k, v) => v != currentProps.get(k)
+    }
+    val deletedKeySet = currentProps.asScala.filter {
+      case (k, _) => !newProps.containsKey(k)
+    }.keySet
+    (changeMap, deletedKeySet)
+  }
+
+  /**
+    * Updates values in `props` with the new values from `propsOverride`. Synonyms of updated configs
+    * are removed from `props` to ensure that the config with the higher precedence is applied. For example,
+    * if `log.roll.ms` was defined in server.properties and `log.roll.hours` is configured dynamically,
+    * `log.roll.hours` from the dynamic configuration will be used and `log.roll.ms` will be removed from
+    * `props` (even though `log.roll.hours` is secondary to `log.roll.ms`).
+    */
+  private def overrideProps(props: mutable.Map[String, String], propsOverride: mutable.Map[String, String]): Unit = {
+    propsOverride.forKeyValue { (k, v) =>
+      // Remove synonyms of `k` to ensure the right precedence is applied. But disable `matchListenerOverride`
+      // so that base configs corresponding to listener configs are not removed. Base configs should not be removed
+      // since they may be used by other listeners. It is ok to retain them in `props` since base configs cannot be
+      // dynamically updated and listener-specific configs have the higher precedence.
+      brokerConfigSynonyms(k, matchListenerOverride = false).foreach(props.remove)
+      props.put(k, v)
+    }
+  }
+
+  private def updateCurrentConfig(): Unit = {
+    val newProps = mutable.Map[String, String]()
+    newProps ++= staticBrokerConfigs
+    overrideProps(newProps, dynamicDefaultConfigs)
+    overrideProps(newProps, dynamicBrokerConfigs)
+    val oldConfig = currentConfig
+    val (newConfig, brokerReconfigurablesToUpdate) = processReconfiguration(newProps, validateOnly = false)
+    if (newConfig ne currentConfig) {
+      currentConfig = newConfig
+      kafkaConfig.updateCurrentConfig(newConfig)
+
+      // Process BrokerReconfigurable updates after current config is updated
+      brokerReconfigurablesToUpdate.foreach(_.reconfigure(oldConfig, newConfig))
+    }
+  }
+
+  private def processReconfiguration(newProps: Map[String, String], validateOnly: Boolean): (KafkaConfig, List[BrokerReconfigurable]) = {
+    val newConfig = new KafkaConfig(newProps.asJava, !validateOnly, None)
+    val (changeMap, deletedKeySet) = updatedConfigs(newConfig.originalsFromThisConfig, currentConfig.originals)
+    if (changeMap.nonEmpty || deletedKeySet.nonEmpty) {
+      try {
+        val customConfigs = new util.HashMap[String, Object](newConfig.originalsFromThisConfig) // non-Kafka configs
+        newConfig.valuesFromThisConfig.keySet.forEach(customConfigs.remove(_))
+        reconfigurables.foreach {
+          case listenerReconfigurable: ListenerReconfigurable =>
+            processListenerReconfigurable(listenerReconfigurable, newConfig, customConfigs, validateOnly, reloadOnly = false)
+          case reconfigurable =>
+            if (needsReconfiguration(reconfigurable.reconfigurableConfigs, changeMap.keySet, deletedKeySet))
+              processReconfigurable(reconfigurable, changeMap.keySet, newConfig.valuesFromThisConfig, customConfigs, validateOnly)
+        }
+
+        // BrokerReconfigurable updates are processed after config is updated. Only do the validation here.
+        val brokerReconfigurablesToUpdate = mutable.Buffer[BrokerReconfigurable]()
+        brokerReconfigurables.foreach { reconfigurable =>
+          if (needsReconfiguration(reconfigurable.reconfigurableConfigs.asJava, changeMap.keySet, deletedKeySet)) {
+            reconfigurable.validateReconfiguration(newConfig)
+            if (!validateOnly)
+              brokerReconfigurablesToUpdate += reconfigurable
+          }
+        }
+        (newConfig, brokerReconfigurablesToUpdate.toList)
+      } catch {
+        case e: Exception =>
+          if (!validateOnly)
+            error(s"Failed to update broker configuration with configs : ${newConfig.originalsFromThisConfig}", e)
+          throw new ConfigException("Invalid dynamic configuration", e)
+      }
+    }
+    else
+      (currentConfig, List.empty)
+  }
+
+  private def needsReconfiguration(reconfigurableConfigs: util.Set[String], updatedKeys: Set[String], deletedKeys: Set[String]): Boolean = {
+    reconfigurableConfigs.asScala.intersect(updatedKeys).nonEmpty ||
+      reconfigurableConfigs.asScala.intersect(deletedKeys).nonEmpty
+  }
+
+  private def processListenerReconfigurable(listenerReconfigurable: ListenerReconfigurable,
+                                            newConfig: KafkaConfig,
+                                            customConfigs: util.Map[String, Object],
+                                            validateOnly: Boolean,
+                                            reloadOnly:  Boolean): Unit = {
+    val listenerName = listenerReconfigurable.listenerName
+    val oldValues = currentConfig.valuesWithPrefixOverride(listenerName.configPrefix)
+    val newValues = newConfig.valuesFromThisConfigWithPrefixOverride(listenerName.configPrefix)
+    val (changeMap, deletedKeys) = updatedConfigs(newValues, oldValues)
+    val updatedKeys = changeMap.keySet
+    val configsChanged = needsReconfiguration(listenerReconfigurable.reconfigurableConfigs, updatedKeys, deletedKeys)
+    // if `reloadOnly`, reconfigure if configs haven't changed. Otherwise reconfigure if configs have changed
+    if (reloadOnly != configsChanged)
+      processReconfigurable(listenerReconfigurable, updatedKeys, newValues, customConfigs, validateOnly)
+  }
+
+  private def processReconfigurable(reconfigurable: Reconfigurable,
+                                    updatedConfigNames: Set[String],
+                                    allNewConfigs: util.Map[String, _],
+                                    newCustomConfigs: util.Map[String, Object],
+                                    validateOnly: Boolean): Unit = {
+    val newConfigs = new util.HashMap[String, Object]
+    allNewConfigs.forEach { (k, v) => newConfigs.put(k, v.asInstanceOf[AnyRef]) }
+    newConfigs.putAll(newCustomConfigs)
+    try {
+      reconfigurable.validateReconfiguration(newConfigs)
+    } catch {
+      case e: ConfigException => throw e
+      case _: Exception =>
+        throw new ConfigException(s"Validation of dynamic config update of $updatedConfigNames failed with class ${reconfigurable.getClass}")
+    }
+
+    if (!validateOnly) {
+      info(s"Reconfiguring $reconfigurable, updated configs: $updatedConfigNames custom configs: $newCustomConfigs")
+      reconfigurable.reconfigure(newConfigs)
+    }
+  }
+}
+
+trait BrokerReconfigurable {
+
+  def reconfigurableConfigs: Set[String]
+
+  def validateReconfiguration(newConfig: KafkaConfig): Unit
+
+  def reconfigure(oldConfig: KafkaConfig, newConfig: KafkaConfig): Unit
+}
+
+object DynamicLogConfig {
+  // Exclude message.format.version for now since we need to check that the version
+  // is supported on all brokers in the cluster.
+  val ExcludedConfigs = Set(KafkaConfig.LogMessageFormatVersionProp)
+
+  val ReconfigurableConfigs = LogConfig.TopicConfigSynonyms.values.toSet -- ExcludedConfigs
+  val KafkaConfigToLogConfigName = LogConfig.TopicConfigSynonyms.map { case (k, v) => (v, k) }
+}
+
+class DynamicLogConfig(logManager: LogManager, server: KafkaServer) extends BrokerReconfigurable with Logging {
+
+  override def reconfigurableConfigs: Set[String] = {
+    DynamicLogConfig.ReconfigurableConfigs
+  }
+
+  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.
+  }
+
+  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.filter { case (k, _) => log.config.overriddenConfigs.contains(k) }
+
+      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)
+    newConfig.valuesFromThisConfig.forEach { (k, v) =>
+      if (DynamicLogConfig.ReconfigurableConfigs.contains(k)) {
+        DynamicLogConfig.KafkaConfigToLogConfigName.get(k).foreach { configName =>
+          if (v == null)
+             newBrokerDefaults.remove(configName)
+          else
+            newBrokerDefaults.put(configName, v.asInstanceOf[AnyRef])
+        }
+      }
+    }
+
+    logManager.reconfigureDefaultLogConfig(LogConfig(newBrokerDefaults))
+
+    updateLogsConfig(newBrokerDefaults.asScala)
+
+    if (logManager.currentDefaultConfig.uncleanLeaderElectionEnable && !origUncleanLeaderElectionEnable) {
+      server.kafkaController.enableDefaultUncleanLeaderElection()
+    }
+  }
+}
+
+object DynamicThreadPool {
+  val ReconfigurableConfigs = Set(
+    KafkaConfig.NumIoThreadsProp,
+    KafkaConfig.NumNetworkThreadsProp,
+    KafkaConfig.NumReplicaFetchersProp,
+    KafkaConfig.NumRecoveryThreadsPerDataDirProp,
+    KafkaConfig.BackgroundThreadsProp)
+}
+
+class DynamicThreadPool(server: KafkaServer) extends BrokerReconfigurable {
+
+  override def reconfigurableConfigs: Set[String] = {
+    DynamicThreadPool.ReconfigurableConfigs
+  }
+
+  override def validateReconfiguration(newConfig: KafkaConfig): Unit = {
+    newConfig.values.forEach { (k, v) =>
+      if (DynamicThreadPool.ReconfigurableConfigs.contains(k)) {
+        val newValue = v.asInstanceOf[Int]
+        val oldValue = currentValue(k)
+        if (newValue != oldValue) {
+          val errorMsg = s"Dynamic thread count update validation failed for $k=$v"
+          if (newValue <= 0)
+            throw new ConfigException(s"$errorMsg, value should be at least 1")
+          if (newValue < oldValue / 2)
+            throw new ConfigException(s"$errorMsg, value should be at least half the current value $oldValue")
+          if (newValue > oldValue * 2)
+            throw new ConfigException(s"$errorMsg, value should not be greater than double the current value $oldValue")
+        }
+      }
+    }
+  }
+
+  override def reconfigure(oldConfig: KafkaConfig, newConfig: KafkaConfig): Unit = {
+    if (newConfig.numIoThreads != oldConfig.numIoThreads)
+      server.dataPlaneRequestHandlerPool.resizeThreadPool(newConfig.numIoThreads)
+    if (newConfig.numNetworkThreads != oldConfig.numNetworkThreads)
+      server.socketServer.resizeThreadPool(oldConfig.numNetworkThreads, newConfig.numNetworkThreads)
+    if (newConfig.numReplicaFetchers != oldConfig.numReplicaFetchers)
+      server.replicaManager.replicaFetcherManager.resizeThreadPool(newConfig.numReplicaFetchers)
+    if (newConfig.numRecoveryThreadsPerDataDir != oldConfig.numRecoveryThreadsPerDataDir)
+      server.getLogManager.resizeRecoveryThreadPool(newConfig.numRecoveryThreadsPerDataDir)
+    if (newConfig.backgroundThreads != oldConfig.backgroundThreads)
+      server.kafkaScheduler.resizeThreadPool(newConfig.backgroundThreads)
+  }
+
+  private def currentValue(name: String): Int = {
+    name match {
+      case KafkaConfig.NumIoThreadsProp => server.config.numIoThreads
+      case KafkaConfig.NumNetworkThreadsProp => server.config.numNetworkThreads
+      case KafkaConfig.NumReplicaFetchersProp => server.config.numReplicaFetchers
+      case KafkaConfig.NumRecoveryThreadsPerDataDirProp => server.config.numRecoveryThreadsPerDataDir
+      case KafkaConfig.BackgroundThreadsProp => server.config.backgroundThreads
+      case n => throw new IllegalStateException(s"Unexpected config $n")
+    }
+  }
+}
+
+class DynamicMetricsReporters(brokerId: Int, server: KafkaServer) extends Reconfigurable {
+
+  private val dynamicConfig = server.config.dynamicConfig
+  private val metrics = server.metrics
+  private val propsOverride = Map[String, AnyRef](KafkaConfig.BrokerIdProp -> brokerId.toString)
+  private val currentReporters = mutable.Map[String, MetricsReporter]()
+
+  createReporters(dynamicConfig.currentKafkaConfig.getList(KafkaConfig.MetricReporterClassesProp),
+    Collections.emptyMap[String, Object])
+
+  private[server] def currentMetricsReporters: List[MetricsReporter] = currentReporters.values.toList
+
+  override def configure(configs: util.Map[String, _]): Unit = {}
+
+  override def reconfigurableConfigs(): util.Set[String] = {
+    val configs = new util.HashSet[String]()
+    configs.add(KafkaConfig.MetricReporterClassesProp)
+    currentReporters.values.foreach {
+      case reporter: Reconfigurable => configs.addAll(reporter.reconfigurableConfigs)
+      case _ =>
+    }
+    configs
+  }
+
+  override def validateReconfiguration(configs: util.Map[String, _]): Unit = {
+    val updatedMetricsReporters = metricsReporterClasses(configs)
+
+    // Ensure all the reporter classes can be loaded and have a default constructor
+    updatedMetricsReporters.foreach { className =>
+      val clazz = Utils.loadClass(className, classOf[MetricsReporter])
+      clazz.getConstructor()
+    }
+
+    // Validate the new configuration using every reconfigurable reporter instance that is not being deleted
+    currentReporters.values.foreach {
+      case reporter: Reconfigurable =>
+        if (updatedMetricsReporters.contains(reporter.getClass.getName))
+          reporter.validateReconfiguration(configs)
+      case _ =>
+    }
+  }
+
+  override def reconfigure(configs: util.Map[String, _]): Unit = {
+    val updatedMetricsReporters = metricsReporterClasses(configs)
+    val deleted = currentReporters.keySet.toSet -- updatedMetricsReporters
+    deleted.foreach(removeReporter)
+    currentReporters.values.foreach {
+      case reporter: Reconfigurable => dynamicConfig.maybeReconfigure(reporter, dynamicConfig.currentKafkaConfig, configs)
+      case _ =>
+    }
+    val added = updatedMetricsReporters.filterNot(currentReporters.keySet)
+    createReporters(added.asJava, configs)
+  }
+
+  private def createReporters(reporterClasses: util.List[String],
+                              updatedConfigs: util.Map[String, _]): Unit = {
+    val props = new util.HashMap[String, AnyRef]
+    updatedConfigs.forEach { (k, v) => props.put(k, v.asInstanceOf[AnyRef]) }
+    propsOverride.forKeyValue { (k, v) => props.put(k, v) }
+    val reporters = dynamicConfig.currentKafkaConfig.getConfiguredInstances(reporterClasses, classOf[MetricsReporter], props)
+    reporters.forEach { reporter =>
+      metrics.addReporter(reporter)
+      currentReporters += reporter.getClass.getName -> reporter
+    }
+    server.notifyClusterListeners(reporters.asScala)
+    server.notifyMetricsReporters(reporters.asScala)
+  }
+
+  private def removeReporter(className: String): Unit = {
+    currentReporters.remove(className).foreach(metrics.removeReporter)
+  }
+
+  private def metricsReporterClasses(configs: util.Map[String, _]): mutable.Buffer[String] = {
+    configs.get(KafkaConfig.MetricReporterClassesProp).asInstanceOf[util.List[String]].asScala
+  }
+}
+object DynamicListenerConfig {
+
+  val ReconfigurableConfigs = Set(
+    // Listener configs
+    KafkaConfig.AdvertisedListenersProp,
+    KafkaConfig.ListenersProp,
+    KafkaConfig.ListenerSecurityProtocolMapProp,
+
+    // SSL configs
+    KafkaConfig.PrincipalBuilderClassProp,
+    KafkaConfig.SslProtocolProp,
+    KafkaConfig.SslProviderProp,
+    KafkaConfig.SslCipherSuitesProp,
+    KafkaConfig.SslEnabledProtocolsProp,
+    KafkaConfig.SslKeystoreTypeProp,
+    KafkaConfig.SslKeystoreLocationProp,
+    KafkaConfig.SslKeystorePasswordProp,
+    KafkaConfig.SslKeyPasswordProp,
+    KafkaConfig.SslTruststoreTypeProp,
+    KafkaConfig.SslTruststoreLocationProp,
+    KafkaConfig.SslTruststorePasswordProp,
+    KafkaConfig.SslKeyManagerAlgorithmProp,
+    KafkaConfig.SslTrustManagerAlgorithmProp,
+    KafkaConfig.SslEndpointIdentificationAlgorithmProp,
+    KafkaConfig.SslSecureRandomImplementationProp,
+    KafkaConfig.SslClientAuthProp,
+    KafkaConfig.SslEngineFactoryClassProp,
+
+    // SASL configs
+    KafkaConfig.SaslMechanismInterBrokerProtocolProp,
+    KafkaConfig.SaslJaasConfigProp,
+    KafkaConfig.SaslEnabledMechanismsProp,
+    KafkaConfig.SaslKerberosServiceNameProp,
+    KafkaConfig.SaslKerberosKinitCmdProp,
+    KafkaConfig.SaslKerberosTicketRenewWindowFactorProp,
+    KafkaConfig.SaslKerberosTicketRenewJitterProp,
+    KafkaConfig.SaslKerberosMinTimeBeforeReloginProp,
+    KafkaConfig.SaslKerberosPrincipalToLocalRulesProp,
+    KafkaConfig.SaslLoginRefreshWindowFactorProp,
+    KafkaConfig.SaslLoginRefreshWindowJitterProp,
+    KafkaConfig.SaslLoginRefreshMinPeriodSecondsProp,
+    KafkaConfig.SaslLoginRefreshBufferSecondsProp,
+
+    // Connection limit configs
+    KafkaConfig.MaxConnectionsProp,
+    KafkaConfig.MaxConnectionCreationRateProp
+  )
+}
+
+class DynamicClientQuotaCallback(brokerId: Int, server: KafkaServer) extends Reconfigurable {
+
+  override def configure(configs: util.Map[String, _]): Unit = {}
+
+  override def reconfigurableConfigs(): util.Set[String] = {
+    val configs = new util.HashSet[String]()
+    server.quotaManagers.clientQuotaCallback.foreach {
+      case callback: Reconfigurable => configs.addAll(callback.reconfigurableConfigs)
+      case _ =>
+    }
+    configs
+  }
+
+  override def validateReconfiguration(configs: util.Map[String, _]): Unit = {
+    server.quotaManagers.clientQuotaCallback.foreach {
+      case callback: Reconfigurable => callback.validateReconfiguration(configs)
+      case _ =>
+    }
+  }
+
+  override def reconfigure(configs: util.Map[String, _]): Unit = {
+    val config = server.config
+    server.quotaManagers.clientQuotaCallback.foreach {
+      case callback: Reconfigurable =>
+        config.dynamicConfig.maybeReconfigure(callback, config.dynamicConfig.currentKafkaConfig, configs)
+        true
+      case _ => false
+    }
+  }
+}
+
+class DynamicListenerConfig(server: KafkaServer) extends BrokerReconfigurable with Logging {
+
+  override def reconfigurableConfigs: Set[String] = {
+    DynamicListenerConfig.ReconfigurableConfigs
+  }
+
+  def validateReconfiguration(newConfig: KafkaConfig): Unit = {
+    val oldConfig = server.config
+    val newListeners = listenersToMap(newConfig.listeners)
+    val newAdvertisedListeners = listenersToMap(newConfig.advertisedListeners)
+    val oldListeners = listenersToMap(oldConfig.listeners)
+    if (!newAdvertisedListeners.keySet.subsetOf(newListeners.keySet))
+      throw new ConfigException(s"Advertised listeners '$newAdvertisedListeners' must be a subset of listeners '$newListeners'")
+    if (!newListeners.keySet.subsetOf(newConfig.listenerSecurityProtocolMap.keySet))
+      throw new ConfigException(s"Listeners '$newListeners' must be subset of listener map '${newConfig.listenerSecurityProtocolMap}'")
+    newListeners.keySet.intersect(oldListeners.keySet).foreach { listenerName =>
+      def immutableListenerConfigs(kafkaConfig: KafkaConfig, prefix: String): Map[String, AnyRef] = {
+        kafkaConfig.originalsWithPrefix(prefix, true).asScala.filter { case (key, _) =>
+          // skip the reconfigurable configs
+          !DynamicSecurityConfigs.contains(key) && !SocketServer.ListenerReconfigurableConfigs.contains(key)
+        }
+      }
+      if (immutableListenerConfigs(newConfig, listenerName.configPrefix) != immutableListenerConfigs(oldConfig, listenerName.configPrefix))
+        throw new ConfigException(s"Configs cannot be updated dynamically for existing listener $listenerName, " +
+          "restart broker or create a new listener for update")
+      if (oldConfig.listenerSecurityProtocolMap(listenerName) != newConfig.listenerSecurityProtocolMap(listenerName))
+        throw new ConfigException(s"Security protocol cannot be updated for existing listener $listenerName")
+    }
+    if (!newAdvertisedListeners.contains(newConfig.interBrokerListenerName))
+      throw new ConfigException(s"Advertised listener must be specified for inter-broker listener ${newConfig.interBrokerListenerName}")
+  }
+
+  def reconfigure(oldConfig: KafkaConfig, newConfig: KafkaConfig): Unit = {
+    val newListeners = newConfig.listeners
+    val newListenerMap = listenersToMap(newListeners)
+    val oldListeners = oldConfig.listeners
+    val oldListenerMap = listenersToMap(oldListeners)
+    val listenersRemoved = oldListeners.filterNot(e => newListenerMap.contains(e.listenerName))
+    val listenersAdded = newListeners.filterNot(e => oldListenerMap.contains(e.listenerName))
+
+    // Clear SASL login cache to force re-login
+    if (listenersAdded.nonEmpty || listenersRemoved.nonEmpty)
+      LoginManager.closeAll()
+
+    server.socketServer.removeListeners(listenersRemoved)
+    if (listenersAdded.nonEmpty)
+      server.socketServer.addListeners(listenersAdded)
+
+    server.kafkaController.updateBrokerInfo(server.createBrokerInfo)
+  }
+
+  private def listenersToMap(listeners: Seq[EndPoint]): Map[ListenerName, EndPoint] =
+    listeners.map(e => (e.listenerName, e)).toMap
+
+}
+
+
+
diff --git a/core/src/main/scala/kafka/server/DynamicConfig.scala b/core/src/main/scala/kafka/server/DynamicConfig.scala
index ddfdff83dc877..f00c610776aa4 100644
--- a/core/src/main/scala/kafka/server/DynamicConfig.scala
+++ b/core/src/main/scala/kafka/server/DynamicConfig.scala
@@ -17,14 +17,17 @@
 
 package kafka.server
 
+import java.net.{InetAddress, UnknownHostException}
 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.Importance._
 import org.apache.kafka.common.config.ConfigDef.Range._
 import org.apache.kafka.common.config.ConfigDef.Type._
-import scala.collection.JavaConverters._
+
+import scala.jdk.CollectionConverters._
 
 /**
   * Class used to hold dynamic configs. These are configs which have no physical manifestation in the server.properties
@@ -33,15 +36,15 @@ import scala.collection.JavaConverters._
 object DynamicConfig {
 
   object Broker {
-    //Properties
+    // Properties
     val LeaderReplicationThrottledRateProp = "leader.replication.throttled.rate"
     val FollowerReplicationThrottledRateProp = "follower.replication.throttled.rate"
     val ReplicaAlterLogDirsIoMaxBytesPerSecondProp = "replica.alter.log.dirs.io.max.bytes.per.second"
 
-    //Defaults
+    // Defaults
     val DefaultReplicationThrottledRate = ReplicationQuotaManagerConfig.QuotaBytesPerSecondDefault
 
-    //Documentation
+    // Documentation
     val LeaderReplicationThrottledRateDoc = "A long representing the upper bound (bytes/sec) on replication traffic for leaders enumerated in the " +
       s"property ${LogConfig.LeaderReplicationThrottledReplicasProp} (for each topic). This property can be only set dynamically. It is suggested that the " +
       s"limit be kept above 1MB/s for accurate behaviour."
@@ -51,65 +54,117 @@ object DynamicConfig {
     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."
 
-    //Definitions
-    private val brokerConfigDef = new ConfigDef()
-      //round minimum value down, to make it easier for users.
+    // Definitions
+    val brokerConfigDef = new ConfigDef()
+      // Round minimum value down, to make it easier for users.
       .define(LeaderReplicationThrottledRateProp, LONG, DefaultReplicationThrottledRate, atLeast(0), MEDIUM, LeaderReplicationThrottledRateDoc)
       .define(FollowerReplicationThrottledRateProp, LONG, DefaultReplicationThrottledRate, atLeast(0), MEDIUM, FollowerReplicationThrottledRateDoc)
       .define(ReplicaAlterLogDirsIoMaxBytesPerSecondProp, LONG, DefaultReplicationThrottledRate, atLeast(0), MEDIUM, ReplicaAlterLogDirsIoMaxBytesPerSecondDoc)
+    DynamicBrokerConfig.addDynamicConfigs(brokerConfigDef)
+    val nonDynamicProps = KafkaConfig.configNames.toSet -- brokerConfigDef.names.asScala
 
     def names = brokerConfigDef.names
 
-    def validate(props: Properties) = DynamicConfig.validate(brokerConfigDef, props)
+    def validate(props: Properties) = DynamicConfig.validate(brokerConfigDef, props, customPropsAllowed = true)
   }
 
-  object Client {
-    //Properties
+  object QuotaConfigs {
     val ProducerByteRateOverrideProp = "producer_byte_rate"
     val ConsumerByteRateOverrideProp = "consumer_byte_rate"
     val RequestPercentageOverrideProp = "request_percentage"
+    val ControllerMutationOverrideProp = "controller_mutation_rate"
+    private val configNames = Set(ProducerByteRateOverrideProp, ConsumerByteRateOverrideProp,
+      RequestPercentageOverrideProp, ControllerMutationOverrideProp)
 
-    //Defaults
-    val DefaultProducerOverride = ClientQuotaManagerConfig.QuotaBytesPerSecondDefault
-    val DefaultConsumerOverride = ClientQuotaManagerConfig.QuotaBytesPerSecondDefault
-    val DefaultRequestOverride = ClientQuotaManagerConfig.QuotaRequestPercentDefault
+    def isQuotaConfig(name: String): Boolean = configNames.contains(name)
+  }
 
-    //Documentation
+  object Client {
+    // Properties
+    val ProducerByteRateOverrideProp = QuotaConfigs.ProducerByteRateOverrideProp
+    val ConsumerByteRateOverrideProp = QuotaConfigs.ConsumerByteRateOverrideProp
+    val RequestPercentageOverrideProp = QuotaConfigs.RequestPercentageOverrideProp
+    val ControllerMutationOverrideProp = QuotaConfigs.ControllerMutationOverrideProp
+
+    // Defaults
+    val DefaultProducerOverride = ClientQuotaManagerConfig.QuotaDefault
+    val DefaultConsumerOverride = ClientQuotaManagerConfig.QuotaDefault
+    val DefaultRequestOverride = ClientRequestQuotaManager.QuotaRequestPercentDefault
+    val DefaultControllerMutationOverride = ControllerMutationQuotaManager.QuotaControllerMutationDefault
+
+    // Documentation
     val ProducerOverrideDoc = "A rate representing the upper bound (bytes/sec) for producer traffic."
     val ConsumerOverrideDoc = "A rate representing the upper bound (bytes/sec) for consumer traffic."
     val RequestOverrideDoc = "A percentage representing the upper bound of time spent for processing requests."
+    val ControllerMutationOverrideDoc = "The rate at which mutations are accepted for the create topics request, " +
+      "the create partitions request and the delete topics request. The rate is accumulated by the number of partitions created or deleted."
 
-    //Definitions
+    // Definitions
     private val clientConfigs = new ConfigDef()
       .define(ProducerByteRateOverrideProp, LONG, DefaultProducerOverride, MEDIUM, ProducerOverrideDoc)
       .define(ConsumerByteRateOverrideProp, LONG, DefaultConsumerOverride, MEDIUM, ConsumerOverrideDoc)
       .define(RequestPercentageOverrideProp, DOUBLE, DefaultRequestOverride, MEDIUM, RequestOverrideDoc)
+      .define(ControllerMutationOverrideProp, DOUBLE, DefaultControllerMutationOverride, MEDIUM, ControllerMutationOverrideDoc)
+
+    def configKeys = clientConfigs.configKeys
 
     def names = clientConfigs.names
 
-    def validate(props: Properties) = DynamicConfig.validate(clientConfigs, props)
+    def validate(props: Properties) = DynamicConfig.validate(clientConfigs, props, customPropsAllowed = false)
   }
 
   object User {
-
-    //Definitions
+    // Definitions
     private val userConfigs = CredentialProvider.userCredentialConfigs
       .define(Client.ProducerByteRateOverrideProp, LONG, Client.DefaultProducerOverride, MEDIUM, Client.ProducerOverrideDoc)
       .define(Client.ConsumerByteRateOverrideProp, LONG, Client.DefaultConsumerOverride, MEDIUM, Client.ConsumerOverrideDoc)
       .define(Client.RequestPercentageOverrideProp, DOUBLE, Client.DefaultRequestOverride, MEDIUM, Client.RequestOverrideDoc)
+      .define(Client.ControllerMutationOverrideProp, DOUBLE, Client.DefaultControllerMutationOverride, MEDIUM, Client.ControllerMutationOverrideDoc)
+
+    def configKeys = userConfigs.configKeys
 
     def names = userConfigs.names
 
-    def validate(props: Properties) = DynamicConfig.validate(userConfigs, props)
+    def validate(props: Properties) = DynamicConfig.validate(userConfigs, props, customPropsAllowed = false)
+  }
+
+  object Ip {
+    val IpConnectionRateOverrideProp = "connection_creation_rate"
+    val UnlimitedConnectionCreationRate = Int.MaxValue
+    val DefaultConnectionCreationRate = UnlimitedConnectionCreationRate
+    val IpOverrideDoc = "An int representing the upper bound of connections accepted for the specified IP."
+
+    private val ipConfigs = new ConfigDef()
+      .define(IpConnectionRateOverrideProp, INT, DefaultConnectionCreationRate, atLeast(0), MEDIUM, IpOverrideDoc)
+
+    def configKeys = ipConfigs.configKeys
+
+    def names = ipConfigs.names
+
+    def validate(props: Properties) = DynamicConfig.validate(ipConfigs, props, customPropsAllowed = false)
+
+    def isValidIpEntity(ip: String): Boolean = {
+      if (ip != ConfigEntityName.Default) {
+        try {
+          InetAddress.getByName(ip)
+        } catch {
+          case _: UnknownHostException => return false
+        }
+      }
+      true
+    }
   }
 
-  private def validate(configDef: ConfigDef, props: Properties) = {
-    //Validate Names
+  private def validate(configDef: ConfigDef, props: Properties, customPropsAllowed: Boolean) = {
+    // Validate Names
     val names = configDef.names()
-    props.keys.asScala.foreach { name =>
-      require(names.contains(name), s"Unknown Dynamic Configuration '$name'.")
+    val propKeys = props.keySet.asScala.map(_.asInstanceOf[String])
+    if (!customPropsAllowed) {
+      val unknownKeys = propKeys.filter(!names.contains(_))
+      require(unknownKeys.isEmpty, s"Unknown Dynamic Configuration: $unknownKeys.")
     }
-    //ValidateValues
-    configDef.parse(props)
+    val propResolved = DynamicBrokerConfig.resolveVariableConfigs(props)
+    // ValidateValues
+    configDef.parse(propResolved)
   }
 }
diff --git a/core/src/main/scala/kafka/server/DynamicConfigManager.scala b/core/src/main/scala/kafka/server/DynamicConfigManager.scala
index 457742d55d068..3eed3827782d7 100644
--- a/core/src/main/scala/kafka/server/DynamicConfigManager.scala
+++ b/core/src/main/scala/kafka/server/DynamicConfigManager.scala
@@ -17,20 +17,19 @@
 
 package kafka.server
 
-import kafka.common.{NotificationHandler, ZkNodeChangeNotificationListener}
-import kafka.utils.Json
-import kafka.utils.Logging
-import kafka.utils.ZkUtils
+import java.nio.charset.StandardCharsets
 
-import scala.collection._
-import scala.collection.JavaConverters._
-import kafka.admin.AdminUtils
+import kafka.common.{NotificationHandler, ZkNodeChangeNotificationListener}
+import kafka.utils.{Json, Logging}
 import kafka.utils.json.JsonObject
-import kafka.zk.{AdminZkClient, KafkaZkClient}
+import kafka.zk.{AdminZkClient, ConfigEntityChangeNotificationSequenceZNode, ConfigEntityChangeNotificationZNode, KafkaZkClient}
 import org.apache.kafka.common.config.types.Password
-import org.apache.kafka.common.security.scram.ScramMechanism
+import org.apache.kafka.common.security.scram.internals.ScramMechanism
 import org.apache.kafka.common.utils.Time
 
+import scala.jdk.CollectionConverters._
+import scala.collection._
+
 /**
  * Represents all the entities that can be configured via ZK
  */
@@ -39,7 +38,8 @@ object ConfigType {
   val Client = "clients"
   val User = "users"
   val Broker = "brokers"
-  val all = Seq(Topic, Client, User, Broker)
+  val Ip = "ips"
+  val all = Seq(Topic, Client, User, Broker, Ip)
 }
 
 object ConfigEntityName {
@@ -91,33 +91,34 @@ class DynamicConfigManager(private val zkClient: KafkaZkClient,
   val adminZkClient = new AdminZkClient(zkClient)
 
   object ConfigChangedNotificationHandler extends NotificationHandler {
-    override def processNotification(json: String) = {
+    override def processNotification(jsonBytes: Array[Byte]) = {
       // Ignore non-json notifications because they can be from the deprecated TopicConfigManager
-      Json.parseFull(json).foreach { js =>
+      Json.parseBytes(jsonBytes).foreach { js =>
         val jsObject = js.asJsonObjectOption.getOrElse {
           throw new IllegalArgumentException("Config change notification has an unexpected value. The format is:" +
             """{"version" : 1, "entity_type":"topics/clients", "entity_name" : "topic_name/client_id"} or """ +
             """{"version" : 2, "entity_path":"entity_type/entity_name"}. """ +
-            s"Received: $json")
+            s"Received: ${new String(jsonBytes, StandardCharsets.UTF_8)}")
         }
         jsObject("version").to[Int] match {
-          case 1 => processEntityConfigChangeVersion1(json, jsObject)
-          case 2 => processEntityConfigChangeVersion2(json, jsObject)
+          case 1 => processEntityConfigChangeVersion1(jsonBytes, jsObject)
+          case 2 => processEntityConfigChangeVersion2(jsonBytes, jsObject)
           case version => throw new IllegalArgumentException("Config change notification has unsupported version " +
             s"'$version', supported versions are 1 and 2.")
         }
       }
     }
 
-    private def processEntityConfigChangeVersion1(json: String, 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 " +
-          s"'clients' or 'topics'. Received: $json")
+          s"'clients' or 'topics'. Received: ${new String(jsonBytes, StandardCharsets.UTF_8)}")
       }
 
       val entity = js.get("entity_name").flatMap(_.to[Option[String]]).getOrElse {
-        throw new IllegalArgumentException("Version 1 config change notification does not specify 'entity_name'. Received: " + json)
+        throw new IllegalArgumentException("Version 1 config change notification does not specify 'entity_name'. " +
+          s"Received: ${new String(jsonBytes, StandardCharsets.UTF_8)}")
       }
 
       val entityConfig = adminZkClient.fetchEntityConfig(entityType, entity)
@@ -126,10 +127,11 @@ class DynamicConfigManager(private val zkClient: KafkaZkClient,
 
     }
 
-    private def processEntityConfigChangeVersion2(json: String, 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'. Received: $json")
+        throw new IllegalArgumentException(s"Version 2 config change notification must specify 'entity_path'. " +
+          s"Received: ${new String(jsonBytes, StandardCharsets.UTF_8)}")
       }
 
       val index = entityPath.indexOf('/')
@@ -137,7 +139,7 @@ class DynamicConfigManager(private val zkClient: KafkaZkClient,
       if (index < 0 || !configHandlers.contains(rootEntityType)) {
         val entityTypes = configHandlers.keys.map(entityType => s"'$entityType'/").mkString(", ")
         throw new IllegalArgumentException("Version 2 config change notification must have 'entity_path' starting with " +
-          s"one of $entityTypes. Received: $json")
+          s"one of $entityTypes. Received: ${new String(jsonBytes, StandardCharsets.UTF_8)}")
       }
       val fullSanitizedEntityName = entityPath.substring(index + 1)
 
@@ -151,8 +153,8 @@ class DynamicConfigManager(private val zkClient: KafkaZkClient,
     }
   }
 
-  private val configChangeListener = new ZkNodeChangeNotificationListener(zkClient, ZkUtils.ConfigChangesPath,
-    AdminUtils.EntityConfigChangeZnodePrefix, ConfigChangedNotificationHandler)
+  private val configChangeListener = new ZkNodeChangeNotificationListener(zkClient, ConfigEntityChangeNotificationZNode.path,
+    ConfigEntityChangeNotificationSequenceZNode.SequenceNumberPrefix, ConfigChangedNotificationHandler)
 
   /**
    * Begin watching for config changes
diff --git a/core/src/main/scala/kafka/server/FetchDataInfo.scala b/core/src/main/scala/kafka/server/FetchDataInfo.scala
index cbd54c0e088d3..a55fe23e1eca6 100644
--- a/core/src/main/scala/kafka/server/FetchDataInfo.scala
+++ b/core/src/main/scala/kafka/server/FetchDataInfo.scala
@@ -20,6 +20,11 @@ package kafka.server
 import org.apache.kafka.common.record.Records
 import org.apache.kafka.common.requests.FetchResponse.AbortedTransaction
 
+sealed trait FetchIsolation
+case object FetchLogEnd extends FetchIsolation
+case object FetchHighWatermark extends FetchIsolation
+case object FetchTxnCommitted extends FetchIsolation
+
 case class FetchDataInfo(fetchOffsetMetadata: LogOffsetMetadata,
                          records: Records,
                          firstEntryIncomplete: Boolean = false,
diff --git a/core/src/main/scala/kafka/server/FetchSession.scala b/core/src/main/scala/kafka/server/FetchSession.scala
new file mode 100644
index 0000000000000..c6280f384dfd7
--- /dev/null
+++ b/core/src/main/scala/kafka/server/FetchSession.scala
@@ -0,0 +1,790 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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
+import java.util.Optional
+import java.util.concurrent.{ThreadLocalRandom, TimeUnit}
+
+import kafka.metrics.KafkaMetricsGroup
+import kafka.utils.Logging
+import org.apache.kafka.common.TopicPartition
+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.{ImplicitLinkedHashCollection, Time, Utils}
+
+import scala.math.Ordered.orderingToOrdered
+import scala.collection.{mutable, _}
+
+object FetchSession {
+  type REQ_MAP = util.Map[TopicPartition, FetchRequest.PartitionData]
+  type RESP_MAP = util.LinkedHashMap[TopicPartition, FetchResponse.PartitionData[Records]]
+  type CACHE_MAP = ImplicitLinkedHashCollection[CachedPartition]
+  type RESP_MAP_ITER = util.Iterator[util.Map.Entry[TopicPartition, FetchResponse.PartitionData[Records]]]
+
+  val NUM_INCREMENTAL_FETCH_SESSISONS = "NumIncrementalFetchSessions"
+  val NUM_INCREMENTAL_FETCH_PARTITIONS_CACHED = "NumIncrementalFetchPartitionsCached"
+  val INCREMENTAL_FETCH_SESSIONS_EVICTIONS_PER_SEC = "IncrementalFetchSessionEvictionsPerSec"
+  val EVICTIONS = "evictions"
+
+  def partitionsToLogString(partitions: util.Collection[TopicPartition], traceEnabled: Boolean): String = {
+    if (traceEnabled) {
+      "(" + Utils.join(partitions, ", ") + ")"
+    } else {
+      s"${partitions.size} partition(s)"
+    }
+  }
+}
+
+/**
+  * A cached partition.
+  *
+  * The broker maintains a set of these objects for each incremental fetch session.
+  * When an incremental fetch request is made, any partitions which are not explicitly
+  * enumerated in the fetch request are loaded from the cache.  Similarly, when an
+  * incremental fetch response is being prepared, any partitions that have not changed
+  * are left out of the response.
+  *
+  * We store many of these objects, so it is important for them to be memory-efficient.
+  * That is why we store topic and partition separately rather than storing a TopicPartition
+  * object.  The TP object takes up more memory because it is a separate JVM object, and
+  * because it stores the cached hash code in memory.
+  *
+  * Note that fetcherLogStartOffset is the LSO of the follower performing the fetch, whereas
+  * localLogStartOffset is the log start offset of the partition on this broker.
+  */
+class CachedPartition(val topic: String,
+                      val partition: Int,
+                      var maxBytes: Int,
+                      var fetchOffset: Long,
+                      var highWatermark: Long,
+                      var leaderEpoch: Optional[Integer],
+                      var fetcherLogStartOffset: Long,
+                      var localLogStartOffset: Long,
+                      var lastFetchedEpoch: Optional[Integer])
+    extends ImplicitLinkedHashCollection.Element {
+
+  var cachedNext: Int = ImplicitLinkedHashCollection.INVALID_INDEX
+  var cachedPrev: Int = ImplicitLinkedHashCollection.INVALID_INDEX
+
+  override def next: Int = cachedNext
+  override def setNext(next: Int): Unit = this.cachedNext = next
+  override def prev: Int = cachedPrev
+  override def setPrev(prev: Int): Unit = this.cachedPrev = prev
+
+  def this(topic: String, partition: Int) =
+    this(topic, partition, -1, -1, -1, Optional.empty(), -1, -1, Optional.empty[Integer])
+
+  def this(part: TopicPartition) =
+    this(part.topic, part.partition)
+
+  def this(part: TopicPartition, reqData: FetchRequest.PartitionData) =
+    this(part.topic, part.partition, reqData.maxBytes, reqData.fetchOffset, -1,
+      reqData.currentLeaderEpoch, reqData.logStartOffset, -1, reqData.lastFetchedEpoch)
+
+  def this(part: TopicPartition, reqData: FetchRequest.PartitionData,
+           respData: FetchResponse.PartitionData[Records]) =
+    this(part.topic, part.partition, reqData.maxBytes, reqData.fetchOffset, respData.highWatermark,
+      reqData.currentLeaderEpoch, reqData.logStartOffset, respData.logStartOffset, reqData.lastFetchedEpoch)
+
+  def reqData = new FetchRequest.PartitionData(fetchOffset, fetcherLogStartOffset, maxBytes, leaderEpoch, lastFetchedEpoch)
+
+  def updateRequestParams(reqData: FetchRequest.PartitionData): Unit = {
+    // Update our cached request parameters.
+    maxBytes = reqData.maxBytes
+    fetchOffset = reqData.fetchOffset
+    fetcherLogStartOffset = reqData.logStartOffset
+    leaderEpoch = reqData.currentLeaderEpoch
+    lastFetchedEpoch = reqData.lastFetchedEpoch
+  }
+
+  /**
+    * Determine whether or not the specified cached partition should be included in the FetchResponse we send back to
+    * the fetcher and update it if requested.
+    *
+    * This function should be called while holding the appropriate session lock.
+    *
+    * @param respData partition data
+    * @param updateResponseData if set to true, update this CachedPartition with new request and response data.
+    * @return True if this partition should be included in the response; false if it can be omitted.
+    */
+  def maybeUpdateResponseData(respData: FetchResponse.PartitionData[Records], updateResponseData: Boolean): Boolean = {
+    // Check the response data.
+    var mustRespond = false
+    if ((respData.records != null) && (respData.records.sizeInBytes > 0)) {
+      // Partitions with new data are always included in the response.
+      mustRespond = true
+    }
+    if (highWatermark != respData.highWatermark) {
+      mustRespond = true
+      if (updateResponseData)
+        highWatermark = respData.highWatermark
+    }
+    if (localLogStartOffset != respData.logStartOffset) {
+      mustRespond = true
+      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.
+      // This ensures that when the error goes away, we re-send the partition.
+      if (updateResponseData)
+        highWatermark = -1
+      mustRespond = true
+    }
+    if (respData.divergingEpoch.isPresent) {
+      // Partitions with diverging epoch are always included in response to trigger truncation.
+      mustRespond = true
+    }
+    mustRespond
+  }
+
+  override def hashCode: Int = (31 * partition) + topic.hashCode
+
+  def canEqual(that: Any) = that.isInstanceOf[CachedPartition]
+
+  override def equals(that: Any): Boolean =
+    that match {
+      case that: CachedPartition =>
+        this.eq(that) ||
+          (that.canEqual(this) &&
+            this.partition.equals(that.partition) &&
+            this.topic.equals(that.topic))
+      case _ => false
+    }
+
+  override def toString: String = synchronized {
+    "CachedPartition(topic=" + topic +
+      ", partition=" + partition +
+      ", maxBytes=" + maxBytes +
+      ", fetchOffset=" + fetchOffset +
+      ", highWatermark=" + highWatermark +
+      ", fetcherLogStartOffset=" + fetcherLogStartOffset +
+      ", localLogStartOffset=" + localLogStartOffset  +
+        ")"
+  }
+}
+
+/**
+  * The fetch session.
+  *
+  * Each fetch session is protected by its own lock, which must be taken before mutable
+  * fields are read or modified.  This includes modification of the session partition map.
+  *
+  * @param id           The unique fetch session ID.
+  * @param privileged   True if this session is privileged.  Sessions crated by followers
+  *                     are privileged; sesssion created by consumers are not.
+  * @param partitionMap The CachedPartitionMap.
+  * @param creationMs   The time in milliseconds when this session was created.
+  * @param lastUsedMs   The last used time in milliseconds.  This should only be updated by
+  *                     FetchSessionCache#touch.
+  * @param epoch        The fetch session sequence number.
+  */
+class FetchSession(val id: Int,
+                   val privileged: Boolean,
+                   val partitionMap: FetchSession.CACHE_MAP,
+                   val creationMs: Long,
+                   var lastUsedMs: Long,
+                   var epoch: Int) {
+  // This is used by the FetchSessionCache to store the last known size of this session.
+  // If this is -1, the Session is not in the cache.
+  var cachedSize = -1
+
+  def size: Int = synchronized {
+    partitionMap.size
+  }
+
+  def isEmpty: Boolean = synchronized {
+    partitionMap.isEmpty
+  }
+
+  def lastUsedKey: LastUsedKey = synchronized {
+    LastUsedKey(lastUsedMs, id)
+  }
+
+  def evictableKey: EvictableKey = synchronized {
+    EvictableKey(privileged, cachedSize, id)
+  }
+
+  def metadata: JFetchMetadata = synchronized { new JFetchMetadata(id, epoch) }
+
+  def getFetchOffset(topicPartition: TopicPartition): Option[Long] = synchronized {
+    Option(partitionMap.find(new CachedPartition(topicPartition))).map(_.fetchOffset)
+  }
+
+  type TL = util.ArrayList[TopicPartition]
+
+  // Update the cached partition data based on the request.
+  def update(fetchData: FetchSession.REQ_MAP,
+             toForget: util.List[TopicPartition],
+             reqMetadata: JFetchMetadata): (TL, TL, TL) = synchronized {
+    val added = new TL
+    val updated = new TL
+    val removed = new TL
+    fetchData.forEach { (topicPart, reqData) =>
+      val newCachedPart = new CachedPartition(topicPart, reqData)
+      val cachedPart = partitionMap.find(newCachedPart)
+      if (cachedPart == null) {
+        partitionMap.mustAdd(newCachedPart)
+        added.add(topicPart)
+      } else {
+        cachedPart.updateRequestParams(reqData)
+        updated.add(topicPart)
+      }
+    }
+    toForget.forEach { p =>
+      if (partitionMap.remove(new CachedPartition(p.topic, p.partition)))
+        removed.add(p)
+    }
+    (added, updated, removed)
+  }
+
+  override def toString: String = synchronized {
+    "FetchSession(id=" + id +
+      ", privileged=" + privileged +
+      ", partitionMap.size=" + partitionMap.size +
+      ", creationMs=" + creationMs +
+      ", lastUsedMs=" + lastUsedMs +
+      ", epoch=" + epoch + ")"
+  }
+}
+
+trait FetchContext extends Logging {
+  /**
+    * Get the fetch offset for a given partition.
+    */
+  def getFetchOffset(part: TopicPartition): Option[Long]
+
+  /**
+    * Apply a function to each partition in the fetch request.
+    */
+  def foreachPartition(fun: (TopicPartition, FetchRequest.PartitionData) => Unit): Unit
+
+  /**
+    * Get the response size to be used for quota computation. Since we are returning an empty response in case of
+    * throttling, we are not supposed to update the context until we know that we are not going to throttle.
+    */
+  def getResponseSize(updates: FetchSession.RESP_MAP, versionId: Short): Int
+
+  /**
+    * Updates the fetch context with new partition information.  Generates response data.
+    * The response data may require subsequent down-conversion.
+    */
+  def updateAndGenerateResponseData(updates: FetchSession.RESP_MAP): FetchResponse[Records]
+
+  def partitionsToLogString(partitions: util.Collection[TopicPartition]): String =
+    FetchSession.partitionsToLogString(partitions, isTraceEnabled)
+
+  /**
+    * Return an empty throttled response due to quota violation.
+    */
+  def getThrottledResponse(throttleTimeMs: Int): FetchResponse[Records] =
+    new FetchResponse(Errors.NONE, new FetchSession.RESP_MAP, throttleTimeMs, INVALID_SESSION_ID)
+}
+
+/**
+  * The fetch context for a fetch request that had a session error.
+  */
+class SessionErrorContext(val error: Errors,
+                          val reqMetadata: JFetchMetadata) extends FetchContext {
+  override def getFetchOffset(part: TopicPartition): Option[Long] = None
+
+  override def foreachPartition(fun: (TopicPartition, FetchRequest.PartitionData) => Unit): Unit = {}
+
+  override def getResponseSize(updates: FetchSession.RESP_MAP, versionId: Short): Int = {
+    FetchResponse.sizeOf(versionId, (new FetchSession.RESP_MAP).entrySet.iterator)
+  }
+
+  // Because of the fetch session error, we don't know what partitions were supposed to be in this request.
+  override def updateAndGenerateResponseData(updates: FetchSession.RESP_MAP): FetchResponse[Records] = {
+    debug(s"Session error fetch context returning $error")
+    new FetchResponse(error, new FetchSession.RESP_MAP, 0, INVALID_SESSION_ID)
+  }
+}
+
+/**
+  * The fetch context for a sessionless fetch request.
+  *
+  * @param fetchData          The partition data from the fetch request.
+  */
+class SessionlessFetchContext(val fetchData: util.Map[TopicPartition, FetchRequest.PartitionData]) extends FetchContext {
+  override def getFetchOffset(part: TopicPartition): Option[Long] =
+    Option(fetchData.get(part)).map(_.fetchOffset)
+
+  override def foreachPartition(fun: (TopicPartition, FetchRequest.PartitionData) => Unit): Unit = {
+    fetchData.forEach(fun(_, _))
+  }
+
+  override def getResponseSize(updates: FetchSession.RESP_MAP, versionId: Short): Int = {
+    FetchResponse.sizeOf(versionId, updates.entrySet.iterator)
+  }
+
+  override def updateAndGenerateResponseData(updates: FetchSession.RESP_MAP): FetchResponse[Records] = {
+    debug(s"Sessionless fetch context returning ${partitionsToLogString(updates.keySet)}")
+    new FetchResponse(Errors.NONE, updates, 0, INVALID_SESSION_ID)
+  }
+}
+
+/**
+  * The fetch context for a full fetch request.
+  *
+  * @param time               The clock to use.
+  * @param cache              The fetch session cache.
+  * @param reqMetadata        The request metadata.
+  * @param fetchData          The partition data from the fetch request.
+  * @param isFromFollower     True if this fetch request came from a follower.
+  */
+class FullFetchContext(private val time: Time,
+                       private val cache: FetchSessionCache,
+                       private val reqMetadata: JFetchMetadata,
+                       private val fetchData: util.Map[TopicPartition, FetchRequest.PartitionData],
+                       private val isFromFollower: Boolean) extends FetchContext {
+  override def getFetchOffset(part: TopicPartition): Option[Long] =
+    Option(fetchData.get(part)).map(_.fetchOffset)
+
+  override def foreachPartition(fun: (TopicPartition, FetchRequest.PartitionData) => Unit): Unit = {
+    fetchData.forEach(fun(_, _))
+  }
+
+  override def getResponseSize(updates: FetchSession.RESP_MAP, versionId: Short): Int = {
+    FetchResponse.sizeOf(versionId, updates.entrySet.iterator)
+  }
+
+  override def updateAndGenerateResponseData(updates: FetchSession.RESP_MAP): FetchResponse[Records] = {
+    def createNewSession: FetchSession.CACHE_MAP = {
+      val cachedPartitions = new FetchSession.CACHE_MAP(updates.size)
+      updates.forEach { (part, respData) =>
+        val reqData = fetchData.get(part)
+        cachedPartitions.mustAdd(new CachedPartition(part, reqData, respData))
+      }
+      cachedPartitions
+    }
+    val responseSessionId = cache.maybeCreateSession(time.milliseconds(), isFromFollower,
+        updates.size, () => createNewSession)
+    debug(s"Full fetch context with session id $responseSessionId returning " +
+      s"${partitionsToLogString(updates.keySet)}")
+    new FetchResponse(Errors.NONE, updates, 0, responseSessionId)
+  }
+}
+
+/**
+  * The fetch context for an incremental fetch request.
+  *
+  * @param time         The clock to use.
+  * @param reqMetadata  The request metadata.
+  * @param session      The incremental fetch request session.
+  */
+class IncrementalFetchContext(private val time: Time,
+                              private val reqMetadata: JFetchMetadata,
+                              private val session: FetchSession) extends FetchContext {
+
+  override def getFetchOffset(tp: TopicPartition): Option[Long] = session.getFetchOffset(tp)
+
+  override def foreachPartition(fun: (TopicPartition, FetchRequest.PartitionData) => Unit): Unit = {
+    // Take the session lock and iterate over all the cached partitions.
+    session.synchronized {
+      session.partitionMap.forEach { part =>
+        fun(new TopicPartition(part.topic, part.partition), part.reqData)
+      }
+    }
+  }
+
+  // Iterator that goes over the given partition map and selects partitions that need to be included in the response.
+  // If updateFetchContextAndRemoveUnselected is set to true, the fetch context will be updated for the selected
+  // partitions and also remove unselected ones as they are encountered.
+  private class PartitionIterator(val iter: FetchSession.RESP_MAP_ITER,
+                                  val updateFetchContextAndRemoveUnselected: Boolean)
+    extends FetchSession.RESP_MAP_ITER {
+    var nextElement: util.Map.Entry[TopicPartition, FetchResponse.PartitionData[Records]] = null
+
+    override def hasNext: Boolean = {
+      while ((nextElement == null) && iter.hasNext) {
+        val element = iter.next()
+        val topicPart = element.getKey
+        val respData = element.getValue
+        val cachedPart = session.partitionMap.find(new CachedPartition(topicPart))
+        val mustRespond = cachedPart.maybeUpdateResponseData(respData, updateFetchContextAndRemoveUnselected)
+        if (mustRespond) {
+          nextElement = element
+          if (updateFetchContextAndRemoveUnselected) {
+            session.partitionMap.remove(cachedPart)
+            session.partitionMap.mustAdd(cachedPart)
+          }
+        } else {
+          if (updateFetchContextAndRemoveUnselected) {
+            iter.remove()
+          }
+        }
+      }
+      nextElement != null
+    }
+
+    override def next(): util.Map.Entry[TopicPartition, FetchResponse.PartitionData[Records]] = {
+      if (!hasNext) throw new NoSuchElementException
+      val element = nextElement
+      nextElement = null
+      element
+    }
+
+    override def remove() = throw new UnsupportedOperationException
+  }
+
+  override def getResponseSize(updates: FetchSession.RESP_MAP, versionId: Short): Int = {
+    session.synchronized {
+      val expectedEpoch = JFetchMetadata.nextEpoch(reqMetadata.epoch)
+      if (session.epoch != expectedEpoch) {
+        FetchResponse.sizeOf(versionId, (new FetchSession.RESP_MAP).entrySet.iterator)
+      } else {
+        // Pass the partition iterator which updates neither the fetch context nor the partition map.
+        FetchResponse.sizeOf(versionId, new PartitionIterator(updates.entrySet.iterator, false))
+      }
+    }
+  }
+
+  override def updateAndGenerateResponseData(updates: FetchSession.RESP_MAP): FetchResponse[Records] = {
+    session.synchronized {
+      // Check to make sure that the session epoch didn't change in between
+      // creating this fetch context and generating this response.
+      val expectedEpoch = JFetchMetadata.nextEpoch(reqMetadata.epoch)
+      if (session.epoch != expectedEpoch) {
+        info(s"Incremental fetch session ${session.id} expected epoch $expectedEpoch, but " +
+          s"got ${session.epoch}.  Possible duplicate request.")
+        new FetchResponse(Errors.INVALID_FETCH_SESSION_EPOCH, new FetchSession.RESP_MAP, 0, session.id)
+      } else {
+        // Iterate over the update list using PartitionIterator. This will prune updates which don't need to be sent
+        val partitionIter = new PartitionIterator(updates.entrySet.iterator, true)
+        while (partitionIter.hasNext) {
+          partitionIter.next()
+        }
+        debug(s"Incremental fetch context with session id ${session.id} returning " +
+          s"${partitionsToLogString(updates.keySet)}")
+        new FetchResponse(Errors.NONE, updates, 0, session.id)
+      }
+    }
+  }
+
+  override def getThrottledResponse(throttleTimeMs: Int): FetchResponse[Records] = {
+    session.synchronized {
+      // Check to make sure that the session epoch didn't change in between
+      // creating this fetch context and generating this response.
+      val expectedEpoch = JFetchMetadata.nextEpoch(reqMetadata.epoch)
+      if (session.epoch != expectedEpoch) {
+        info(s"Incremental fetch session ${session.id} expected epoch $expectedEpoch, but " +
+          s"got ${session.epoch}.  Possible duplicate request.")
+        new FetchResponse(Errors.INVALID_FETCH_SESSION_EPOCH, new FetchSession.RESP_MAP, throttleTimeMs, session.id)
+      } else {
+        new FetchResponse(Errors.NONE, new FetchSession.RESP_MAP, throttleTimeMs, session.id)
+      }
+    }
+  }
+}
+
+case class LastUsedKey(lastUsedMs: Long, id: Int) extends Comparable[LastUsedKey] {
+  override def compareTo(other: LastUsedKey): Int =
+    (lastUsedMs, id) compare (other.lastUsedMs, other.id)
+}
+
+case class EvictableKey(privileged: Boolean, size: Int, id: Int) extends Comparable[EvictableKey] {
+  override def compareTo(other: EvictableKey): Int =
+    (privileged, size, id) compare (other.privileged, other.size, other.id)
+}
+
+/**
+  * Caches fetch sessions.
+  *
+  * See tryEvict for an explanation of the cache eviction strategy.
+  *
+  * The FetchSessionCache is thread-safe because all of its methods are synchronized.
+  * Note that individual fetch sessions have their own locks which are separate from the
+  * FetchSessionCache lock.  In order to avoid deadlock, the FetchSessionCache lock
+  * must never be acquired while an individual FetchSession lock is already held.
+  *
+  * @param maxEntries The maximum number of entries that can be in the cache.
+  * @param evictionMs The minimum time that an entry must be unused in order to be evictable.
+  */
+class FetchSessionCache(private val maxEntries: Int,
+                        private val evictionMs: Long) extends Logging with KafkaMetricsGroup {
+  private var numPartitions: Long = 0
+
+  // A map of session ID to FetchSession.
+  private val sessions = new mutable.HashMap[Int, FetchSession]
+
+  // Maps last used times to sessions.
+  private val lastUsed = new util.TreeMap[LastUsedKey, FetchSession]
+
+  // A map containing sessions which can be evicted by both privileged and
+  // unprivileged sessions.
+  private val evictableByAll = new util.TreeMap[EvictableKey, FetchSession]
+
+  // A map containing sessions which can be evicted by privileged sessions.
+  private val evictableByPrivileged = new util.TreeMap[EvictableKey, FetchSession]
+
+  // Set up metrics.
+  removeMetric(FetchSession.NUM_INCREMENTAL_FETCH_SESSISONS)
+  newGauge(FetchSession.NUM_INCREMENTAL_FETCH_SESSISONS, () => FetchSessionCache.this.size)
+  removeMetric(FetchSession.NUM_INCREMENTAL_FETCH_PARTITIONS_CACHED)
+  newGauge(FetchSession.NUM_INCREMENTAL_FETCH_PARTITIONS_CACHED, () => FetchSessionCache.this.totalPartitions)
+  removeMetric(FetchSession.INCREMENTAL_FETCH_SESSIONS_EVICTIONS_PER_SEC)
+  private[server] val evictionsMeter = newMeter(FetchSession.INCREMENTAL_FETCH_SESSIONS_EVICTIONS_PER_SEC,
+    FetchSession.EVICTIONS, TimeUnit.SECONDS, Map.empty)
+
+  /**
+    * Get a session by session ID.
+    *
+    * @param sessionId  The session ID.
+    * @return           The session, or None if no such session was found.
+    */
+  def get(sessionId: Int): Option[FetchSession] = synchronized {
+    sessions.get(sessionId)
+  }
+
+  /**
+    * Get the number of entries currently in the fetch session cache.
+    */
+  def size: Int = synchronized {
+    sessions.size
+  }
+
+  /**
+    * Get the total number of cached partitions.
+    */
+  def totalPartitions: Long = synchronized {
+    numPartitions
+  }
+
+  /**
+    * Creates a new random session ID.  The new session ID will be positive and unique on this broker.
+    *
+    * @return   The new session ID.
+    */
+  def newSessionId(): Int = synchronized {
+    var id = 0
+    do {
+      id = ThreadLocalRandom.current().nextInt(1, Int.MaxValue)
+    } while (sessions.contains(id) || id == INVALID_SESSION_ID)
+    id
+  }
+
+  /**
+    * Try to create a new session.
+    *
+    * @param now                The current time in milliseconds.
+    * @param privileged         True if the new entry we are trying to create is privileged.
+    * @param size               The number of cached partitions in the new entry we are trying to create.
+    * @param createPartitions   A callback function which creates the map of cached partitions.
+    * @return                   If we created a session, the ID; INVALID_SESSION_ID otherwise.
+    */
+  def maybeCreateSession(now: Long,
+                         privileged: Boolean,
+                         size: Int,
+                         createPartitions: () => FetchSession.CACHE_MAP): Int =
+  synchronized {
+    // If there is room, create a new session entry.
+    if ((sessions.size < maxEntries) ||
+        tryEvict(privileged, EvictableKey(privileged, size, 0), now)) {
+      val partitionMap = createPartitions()
+      val session = new FetchSession(newSessionId(), privileged, partitionMap,
+          now, now, JFetchMetadata.nextEpoch(INITIAL_EPOCH))
+      debug(s"Created fetch session ${session.toString}")
+      sessions.put(session.id, session)
+      touch(session, now)
+      session.id
+    } else {
+      debug(s"No fetch session created for privileged=$privileged, size=$size.")
+      INVALID_SESSION_ID
+    }
+  }
+
+  /**
+    * Try to evict an entry from the session cache.
+    *
+    * A proposed new element A may evict an existing element B if:
+    * 1. A is privileged and B is not, or
+    * 2. B is considered "stale" because it has been inactive for a long time, or
+    * 3. A contains more partitions than B, and B is not recently created.
+    *
+    * @param privileged True if the new entry we would like to add is privileged.
+    * @param key        The EvictableKey for the new entry we would like to add.
+    * @param now        The current time in milliseconds.
+    * @return           True if an entry was evicted; false otherwise.
+    */
+  def tryEvict(privileged: Boolean, key: EvictableKey, now: Long): Boolean = synchronized {
+    // Try to evict an entry which is stale.
+    val lastUsedEntry = lastUsed.firstEntry
+    if (lastUsedEntry == null) {
+      trace("There are no cache entries to evict.")
+      false
+    } else if (now - lastUsedEntry.getKey.lastUsedMs > evictionMs) {
+      val session = lastUsedEntry.getValue
+      trace(s"Evicting stale FetchSession ${session.id}.")
+      remove(session)
+      evictionsMeter.mark()
+      true
+    } else {
+      // If there are no stale entries, check the first evictable entry.
+      // If it is less valuable than our proposed entry, evict it.
+      val map = if (privileged) evictableByPrivileged else evictableByAll
+      val evictableEntry = map.firstEntry
+      if (evictableEntry == null) {
+        trace("No evictable entries found.")
+        false
+      } else if (key.compareTo(evictableEntry.getKey) < 0) {
+        trace(s"Can't evict ${evictableEntry.getKey} with ${key.toString}")
+        false
+      } else {
+        trace(s"Evicting ${evictableEntry.getKey} with ${key.toString}.")
+        remove(evictableEntry.getValue)
+        evictionsMeter.mark()
+        true
+      }
+    }
+  }
+
+  def remove(sessionId: Int): Option[FetchSession] = synchronized {
+    get(sessionId) match {
+      case None => None
+      case Some(session) => remove(session)
+    }
+  }
+
+  /**
+    * Remove an entry from the session cache.
+    *
+    * @param session  The session.
+    *
+    * @return         The removed session, or None if there was no such session.
+    */
+  def remove(session: FetchSession): Option[FetchSession] = synchronized {
+    val evictableKey = session.synchronized {
+      lastUsed.remove(session.lastUsedKey)
+      session.evictableKey
+    }
+    evictableByAll.remove(evictableKey)
+    evictableByPrivileged.remove(evictableKey)
+    val removeResult = sessions.remove(session.id)
+    if (removeResult.isDefined) {
+      numPartitions = numPartitions - session.cachedSize
+    }
+    removeResult
+  }
+
+  /**
+    * Update a session's position in the lastUsed and evictable trees.
+    *
+    * @param session  The session.
+    * @param now      The current time in milliseconds.
+    */
+  def touch(session: FetchSession, now: Long): Unit = synchronized {
+    session.synchronized {
+      // Update the lastUsed map.
+      lastUsed.remove(session.lastUsedKey)
+      session.lastUsedMs = now
+      lastUsed.put(session.lastUsedKey, session)
+
+      val oldSize = session.cachedSize
+      if (oldSize != -1) {
+        val oldEvictableKey = session.evictableKey
+        evictableByPrivileged.remove(oldEvictableKey)
+        evictableByAll.remove(oldEvictableKey)
+        numPartitions = numPartitions - oldSize
+      }
+      session.cachedSize = session.size
+      val newEvictableKey = session.evictableKey
+      if ((!session.privileged) || (now - session.creationMs > evictionMs)) {
+        evictableByPrivileged.put(newEvictableKey, session)
+      }
+      if (now - session.creationMs > evictionMs) {
+        evictableByAll.put(newEvictableKey, session)
+      }
+      numPartitions = numPartitions + session.cachedSize
+    }
+  }
+}
+
+class FetchManager(private val time: Time,
+                   private val cache: FetchSessionCache) extends Logging {
+  def newContext(reqMetadata: JFetchMetadata,
+                 fetchData: FetchSession.REQ_MAP,
+                 toForget: util.List[TopicPartition],
+                 isFollower: Boolean): FetchContext = {
+    val context = if (reqMetadata.isFull) {
+      var removedFetchSessionStr = ""
+      if (reqMetadata.sessionId != INVALID_SESSION_ID) {
+        // Any session specified in a FULL fetch request will be closed.
+        if (cache.remove(reqMetadata.sessionId).isDefined) {
+          removedFetchSessionStr = s" Removed fetch session ${reqMetadata.sessionId}."
+        }
+      }
+      var suffix = ""
+      val context = if (reqMetadata.epoch == FINAL_EPOCH) {
+        // If the epoch is FINAL_EPOCH, don't try to create a new session.
+        suffix = " Will not try to create a new session."
+        new SessionlessFetchContext(fetchData)
+      } else {
+        new FullFetchContext(time, cache, reqMetadata, fetchData, isFollower)
+      }
+      debug(s"Created a new full FetchContext with ${partitionsToLogString(fetchData.keySet)}."+
+        s"${removedFetchSessionStr}${suffix}")
+      context
+    } else {
+      cache.synchronized {
+        cache.get(reqMetadata.sessionId) match {
+          case None => {
+            debug(s"Session error for ${reqMetadata.sessionId}: no such session ID found.")
+            new SessionErrorContext(Errors.FETCH_SESSION_ID_NOT_FOUND, reqMetadata)
+          }
+          case Some(session) => session.synchronized {
+            if (session.epoch != reqMetadata.epoch) {
+              debug(s"Session error for ${reqMetadata.sessionId}: expected epoch " +
+                s"${session.epoch}, but got ${reqMetadata.epoch} instead.");
+              new SessionErrorContext(Errors.INVALID_FETCH_SESSION_EPOCH, reqMetadata)
+            } else {
+              val (added, updated, removed) = session.update(fetchData, toForget, reqMetadata)
+              if (session.isEmpty) {
+                debug(s"Created a new sessionless FetchContext and closing session id ${session.id}, " +
+                  s"epoch ${session.epoch}: after removing ${partitionsToLogString(removed)}, " +
+                  s"there are no more partitions left.")
+                cache.remove(session)
+                new SessionlessFetchContext(fetchData)
+              } else {
+                cache.touch(session, time.milliseconds())
+                session.epoch = JFetchMetadata.nextEpoch(session.epoch)
+                debug(s"Created a new incremental FetchContext for session id ${session.id}, " +
+                  s"epoch ${session.epoch}: added ${partitionsToLogString(added)}, " +
+                  s"updated ${partitionsToLogString(updated)}, " +
+                  s"removed ${partitionsToLogString(removed)}")
+                new IncrementalFetchContext(time, reqMetadata, session)
+              }
+            }
+          }
+        }
+      }
+    }
+    context
+  }
+
+  def partitionsToLogString(partitions: util.Collection[TopicPartition]): String =
+    FetchSession.partitionsToLogString(partitions, isTraceEnabled)
+}
diff --git a/core/src/main/scala/kafka/server/FinalizedFeatureCache.scala b/core/src/main/scala/kafka/server/FinalizedFeatureCache.scala
new file mode 100644
index 0000000000000..d7a90d0ba7bc1
--- /dev/null
+++ b/core/src/main/scala/kafka/server/FinalizedFeatureCache.scala
@@ -0,0 +1,162 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package kafka.server
+
+import kafka.utils.Logging
+import org.apache.kafka.common.feature.{Features, FinalizedVersionRange}
+
+import scala.concurrent.TimeoutException
+import scala.math.max
+
+// Raised whenever there was an error in updating the FinalizedFeatureCache with features.
+class FeatureCacheUpdateException(message: String) extends RuntimeException(message) {
+}
+
+// Helper class that represents finalized features along with an epoch value.
+case class FinalizedFeaturesAndEpoch(features: Features[FinalizedVersionRange], epoch: Long) {
+  override def toString(): String = {
+    s"FinalizedFeaturesAndEpoch(features=$features, epoch=$epoch)"
+  }
+}
+
+/**
+ * A common mutable cache containing the latest finalized features and epoch. By default the contents of
+ * the cache are empty. This cache needs to be populated at least once for its contents to become
+ * non-empty. Currently the main reader of this cache is the read path that serves an ApiVersionsRequest,
+ * returning the features information in the response. This cache is typically updated asynchronously
+ * whenever the finalized features and epoch values are modified in ZK by the KafkaController.
+ * This cache is thread-safe for reads and writes.
+ *
+ * @see FinalizedFeatureChangeListener
+ */
+class FinalizedFeatureCache(private val brokerFeatures: BrokerFeatures) extends Logging {
+  @volatile private var featuresAndEpoch: Option[FinalizedFeaturesAndEpoch] = Option.empty
+
+  /**
+   * @return   the latest known FinalizedFeaturesAndEpoch or empty if not defined in the cache.
+   */
+  def get: Option[FinalizedFeaturesAndEpoch] = {
+    featuresAndEpoch
+  }
+
+  def isEmpty: Boolean = {
+    featuresAndEpoch.isEmpty
+  }
+
+  /**
+   * Waits no more than timeoutMs for the cache's epoch to reach an epoch >= minExpectedEpoch.
+   *
+   * @param minExpectedEpoch   the minimum expected epoch to be reached by the cache
+   *                           (should be >= 0)
+   * @param timeoutMs          the timeout (in milli seconds)
+   *
+   * @throws                   TimeoutException if the cache's epoch has not reached at least
+   *                           minExpectedEpoch within timeoutMs.
+   */
+  def waitUntilEpochOrThrow(minExpectedEpoch: Long, timeoutMs: Long): Unit = {
+    if(minExpectedEpoch < 0L) {
+      throw new IllegalArgumentException(
+        s"Expected minExpectedEpoch >= 0, but $minExpectedEpoch was provided.")
+    }
+    waitUntilConditionOrThrow(
+      () => featuresAndEpoch.isDefined && featuresAndEpoch.get.epoch >= minExpectedEpoch,
+      timeoutMs)
+  }
+
+  /**
+   * Clears all existing finalized features and epoch from the cache.
+   */
+  def clear(): Unit = {
+    synchronized {
+      featuresAndEpoch = Option.empty
+      notifyAll()
+    }
+    info("Cleared cache")
+  }
+
+  /**
+   * Updates the cache to the latestFeatures, and updates the existing epoch to latestEpoch.
+   * Expects that the latestEpoch should be always greater than the existing epoch (when the
+   * existing epoch is defined).
+   *
+   * @param latestFeatures   the latest finalized features to be set in the cache
+   * @param latestEpoch      the latest epoch value to be set in the cache
+   *
+   * @throws                 FeatureCacheUpdateException if the cache update operation fails
+   *                         due to invalid parameters or incompatibilities with the broker's
+   *                         supported features. In such a case, the existing cache contents are
+   *                         not modified.
+   */
+  def updateOrThrow(latestFeatures: Features[FinalizedVersionRange], latestEpoch: Long): Unit = {
+    val latest = FinalizedFeaturesAndEpoch(latestFeatures, latestEpoch)
+    val existing = featuresAndEpoch.map(item => item.toString()).getOrElse("")
+    if (featuresAndEpoch.isDefined && featuresAndEpoch.get.epoch > latest.epoch) {
+      val errorMsg = s"FinalizedFeatureCache update failed due to invalid epoch in new $latest." +
+        s" The existing cache contents are $existing."
+      throw new FeatureCacheUpdateException(errorMsg)
+    } else {
+      val incompatibleFeatures = brokerFeatures.incompatibleFeatures(latest.features)
+      if (!incompatibleFeatures.empty) {
+        val errorMsg = "FinalizedFeatureCache update failed since feature compatibility" +
+          s" checks failed! Supported ${brokerFeatures.supportedFeatures} has incompatibilities" +
+          s" with the latest $latest."
+        throw new FeatureCacheUpdateException(errorMsg)
+      } else {
+        val logMsg = s"Updated cache from existing $existing to latest $latest."
+        synchronized {
+          featuresAndEpoch = Some(latest)
+          notifyAll()
+        }
+        info(logMsg)
+      }
+    }
+  }
+
+  /**
+   * Causes the current thread to wait no more than timeoutMs for the specified condition to be met.
+   * It is guaranteed that the provided condition will always be invoked only from within a
+   * synchronized block.
+   *
+   * @param waitCondition   the condition to be waited upon:
+   *                         - if the condition returns true, then, the wait will stop.
+   *                         - if the condition returns false, it means the wait must continue until
+   *                           timeout.
+   *
+   * @param timeoutMs       the timeout (in milli seconds)
+   *
+   * @throws                TimeoutException if the condition is not met within timeoutMs.
+   */
+  private def waitUntilConditionOrThrow(waitCondition: () => Boolean, timeoutMs: Long): Unit = {
+    if(timeoutMs < 0L) {
+      throw new IllegalArgumentException(s"Expected timeoutMs >= 0, but $timeoutMs was provided.")
+    }
+    val waitEndTimeNanos = System.nanoTime() + (timeoutMs * 1000000)
+    synchronized {
+      while (!waitCondition()) {
+        val nowNanos = System.nanoTime()
+        if (nowNanos > waitEndTimeNanos) {
+          throw new TimeoutException(
+            s"Timed out after waiting for ${timeoutMs}ms for required condition to be met." +
+              s" Current epoch: ${featuresAndEpoch.map(fe => fe.epoch).getOrElse("")}.")
+        }
+        val sleepTimeMs = max(1L, (waitEndTimeNanos - nowNanos) / 1000000)
+        wait(sleepTimeMs)
+      }
+    }
+  }
+}
diff --git a/core/src/main/scala/kafka/server/FinalizedFeatureChangeListener.scala b/core/src/main/scala/kafka/server/FinalizedFeatureChangeListener.scala
new file mode 100644
index 0000000000000..ac4600fa42a60
--- /dev/null
+++ b/core/src/main/scala/kafka/server/FinalizedFeatureChangeListener.scala
@@ -0,0 +1,257 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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.{CountDownLatch, LinkedBlockingQueue, TimeUnit}
+
+import kafka.utils.{Logging, ShutdownableThread}
+import kafka.zk.{FeatureZNode, FeatureZNodeStatus, KafkaZkClient, ZkVersion}
+import kafka.zookeeper.{StateChangeHandler, ZNodeChangeHandler}
+import org.apache.kafka.common.internals.FatalExitError
+
+import scala.concurrent.TimeoutException
+
+/**
+ * Listens to changes in the ZK feature node, via the ZK client. Whenever a change notification
+ * is received from ZK, the feature cache in FinalizedFeatureCache is asynchronously updated
+ * to the latest features read from ZK. The cache updates are serialized through a single
+ * notification processor thread.
+ *
+ * @param finalizedFeatureCache   the finalized feature cache
+ * @param zkClient                the Zookeeper client
+ */
+class FinalizedFeatureChangeListener(private val finalizedFeatureCache: FinalizedFeatureCache,
+                                     private val zkClient: KafkaZkClient) extends Logging {
+
+  /**
+   * Helper class used to update the FinalizedFeatureCache.
+   *
+   * @param featureZkNodePath   the path to the ZK feature node to be read
+   * @param maybeNotifyOnce     an optional latch that can be used to notify the caller when an
+   *                            updateOrThrow() operation is over
+   */
+  private class FeatureCacheUpdater(featureZkNodePath: String, maybeNotifyOnce: Option[CountDownLatch]) {
+
+    def this(featureZkNodePath: String) = this(featureZkNodePath, Option.empty)
+
+    /**
+     * Updates the feature cache in FinalizedFeatureCache with the latest features read from the
+     * ZK node in featureZkNodePath. If the cache update is not successful, then, a suitable
+     * exception is raised.
+     *
+     * NOTE: if a notifier was provided in the constructor, then, this method can be invoked exactly
+     * once successfully. A subsequent invocation will raise an exception.
+     *
+     * @throws   IllegalStateException, if a non-empty notifier was provided in the constructor, and
+     *           this method is called again after a successful previous invocation.
+     * @throws   FeatureCacheUpdateException, if there was an error in updating the
+     *           FinalizedFeatureCache.
+     */
+    def updateLatestOrThrow(): Unit = {
+      maybeNotifyOnce.foreach(notifier => {
+        if (notifier.getCount != 1) {
+          throw new IllegalStateException(
+            "Can not notify after updateLatestOrThrow was called more than once successfully.")
+        }
+      })
+
+      debug(s"Reading feature ZK node at path: $featureZkNodePath")
+      val (mayBeFeatureZNodeBytes, version) = zkClient.getDataAndVersion(featureZkNodePath)
+
+      // There are 4 cases:
+      //
+      // (empty dataBytes, valid version)       => The empty dataBytes will fail FeatureZNode deserialization.
+      //                                           FeatureZNode, when present in ZK, can not have empty contents.
+      // (non-empty dataBytes, valid version)   => This is a valid case, and should pass FeatureZNode deserialization
+      //                                           if dataBytes contains valid data.
+      // (empty dataBytes, unknown version)     => This is a valid case, and this can happen if the FeatureZNode
+      //                                           does not exist in ZK.
+      // (non-empty dataBytes, unknown version) => This case is impossible, since, KafkaZkClient.getDataAndVersion
+      //                                           API ensures that unknown version is returned only when the
+      //                                           ZK node is absent. Therefore dataBytes should be empty in such
+      //                                           a case.
+      if (version == ZkVersion.UnknownVersion) {
+        info(s"Feature ZK node at path: $featureZkNodePath does not exist")
+        finalizedFeatureCache.clear()
+      } else {
+        var maybeFeatureZNode: Option[FeatureZNode] = Option.empty
+        try {
+          maybeFeatureZNode = Some(FeatureZNode.decode(mayBeFeatureZNodeBytes.get))
+        } catch {
+          case e: IllegalArgumentException => {
+            error(s"Unable to deserialize feature ZK node at path: $featureZkNodePath", e)
+            finalizedFeatureCache.clear()
+          }
+        }
+        maybeFeatureZNode.foreach(featureZNode => {
+          featureZNode.status match {
+            case FeatureZNodeStatus.Disabled => {
+              info(s"Feature ZK node at path: $featureZkNodePath is in disabled status.")
+              finalizedFeatureCache.clear()
+            }
+            case FeatureZNodeStatus.Enabled => {
+              finalizedFeatureCache.updateOrThrow(featureZNode.features, version)
+            }
+            case _ => throw new IllegalStateException(s"Unexpected FeatureZNodeStatus found in $featureZNode")
+          }
+        })
+      }
+
+      maybeNotifyOnce.foreach(notifier => notifier.countDown())
+    }
+
+    /**
+     * Waits until at least a single updateLatestOrThrow completes successfully. This method returns
+     * immediately if an updateLatestOrThrow call had already completed successfully.
+     *
+     * @param waitTimeMs   the timeout for the wait operation
+     *
+     * @throws             TimeoutException if the wait can not be completed in waitTimeMs
+     *                     milli seconds
+     */
+    def awaitUpdateOrThrow(waitTimeMs: Long): Unit = {
+      maybeNotifyOnce.foreach(notifier => {
+        if (!notifier.await(waitTimeMs, TimeUnit.MILLISECONDS)) {
+          throw new TimeoutException(
+            s"Timed out after waiting for ${waitTimeMs}ms for FeatureCache to be updated.")
+        }
+      })
+    }
+  }
+
+  /**
+   * A shutdownable thread to process feature node change notifications that are populated into the
+   * queue. If any change notification can not be processed successfully (unless it is due to an
+   * interrupt), the thread treats it as a fatal event and triggers Broker exit.
+   *
+   * @param name   name of the thread
+   */
+  private class ChangeNotificationProcessorThread(name: String) extends ShutdownableThread(name = name) {
+    override def doWork(): Unit = {
+      try {
+        queue.take.updateLatestOrThrow()
+      } catch {
+        case ie: InterruptedException =>
+          // While the queue is empty and this thread is blocking on taking an item from the queue,
+          // a concurrent call to FinalizedFeatureChangeListener.close() could interrupt the thread
+          // and cause an InterruptedException to be raised from queue.take(). In such a case, it is
+          // safe to ignore the exception if the thread is being shutdown. We raise the exception
+          // here again, because, it is ignored by ShutdownableThread if it is shutting down.
+          throw ie
+        case e: Exception => {
+          error("Failed to process feature ZK node change event. The broker will eventually exit.", e)
+          throw new FatalExitError(1)
+        }
+      }
+    }
+  }
+
+  // Feature ZK node change handler.
+  object FeatureZNodeChangeHandler extends ZNodeChangeHandler {
+    override val path: String = FeatureZNode.path
+
+    override def handleCreation(): Unit = {
+      info(s"Feature ZK node created at path: $path")
+      queue.add(new FeatureCacheUpdater(path))
+    }
+
+    override def handleDataChange(): Unit = {
+      info(s"Feature ZK node updated at path: $path")
+      queue.add(new FeatureCacheUpdater(path))
+    }
+
+    override def handleDeletion(): Unit = {
+      warn(s"Feature ZK node deleted at path: $path")
+      // This event may happen, rarely (ex: ZK corruption or operational error).
+      // In such a case, we prefer to just log a warning and treat the case as if the node is absent,
+      // and populate the FinalizedFeatureCache with empty finalized features.
+      queue.add(new FeatureCacheUpdater(path))
+    }
+  }
+
+  object ZkStateChangeHandler extends StateChangeHandler {
+    val path: String = FeatureZNode.path
+
+    override val name: String = path
+
+    override def afterInitializingSession(): Unit = {
+      queue.add(new FeatureCacheUpdater(path))
+    }
+  }
+
+  private val queue = new LinkedBlockingQueue[FeatureCacheUpdater]
+
+  private val thread = new ChangeNotificationProcessorThread("feature-zk-node-event-process-thread")
+
+  /**
+   * This method initializes the feature ZK node change listener. Optionally, it also ensures to
+   * update the FinalizedFeatureCache once with the latest contents of the feature ZK node
+   * (if the node exists). This step helps ensure that feature incompatibilities (if any) in brokers
+   * are conveniently detected before the initOrThrow() method returns to the caller. If feature
+   * incompatibilities are detected, this method will throw an Exception to the caller, and the Broker
+   * will exit eventually.
+   *
+   * @param waitOnceForCacheUpdateMs   # of milli seconds to wait for feature cache to be updated once.
+   *                                   (should be > 0)
+   *
+   * @throws Exception if feature incompatibility check could not be finished in a timely manner
+   */
+  def initOrThrow(waitOnceForCacheUpdateMs: Long): Unit = {
+    if (waitOnceForCacheUpdateMs <= 0) {
+      throw new IllegalArgumentException(
+        s"Expected waitOnceForCacheUpdateMs > 0, but provided: $waitOnceForCacheUpdateMs")
+    }
+
+    thread.start()
+    zkClient.registerStateChangeHandler(ZkStateChangeHandler)
+    zkClient.registerZNodeChangeHandlerAndCheckExistence(FeatureZNodeChangeHandler)
+    val ensureCacheUpdateOnce = new FeatureCacheUpdater(
+      FeatureZNodeChangeHandler.path, Some(new CountDownLatch(1)))
+    queue.add(ensureCacheUpdateOnce)
+    try {
+      ensureCacheUpdateOnce.awaitUpdateOrThrow(waitOnceForCacheUpdateMs)
+    } catch {
+      case e: Exception => {
+        close()
+        throw e
+      }
+    }
+  }
+
+  /**
+   * Closes the feature ZK node change listener by unregistering the listener from ZK client,
+   * clearing the queue and shutting down the ChangeNotificationProcessorThread.
+   */
+  def close(): Unit = {
+    zkClient.unregisterStateChangeHandler(ZkStateChangeHandler.name)
+    zkClient.unregisterZNodeChangeHandler(FeatureZNodeChangeHandler.path)
+    queue.clear()
+    thread.shutdown()
+    thread.join()
+  }
+
+  // For testing only.
+  def isListenerInitiated: Boolean = {
+    thread.isRunning && thread.isAlive
+  }
+
+  // For testing only.
+  def isListenerDead: Boolean = {
+    !thread.isRunning && !thread.isAlive
+  }
+}
diff --git a/core/src/main/scala/kafka/server/ForwardingManager.scala b/core/src/main/scala/kafka/server/ForwardingManager.scala
new file mode 100644
index 0000000000000..54bef3eefc8bf
--- /dev/null
+++ b/core/src/main/scala/kafka/server/ForwardingManager.scala
@@ -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 kafka.server
+
+import java.nio.ByteBuffer
+
+import kafka.network.RequestChannel
+import kafka.utils.Logging
+import org.apache.kafka.clients.ClientResponse
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.requests.{AbstractRequest, AbstractResponse, EnvelopeRequest, EnvelopeResponse, RequestHeader}
+import org.apache.kafka.common.utils.Time
+
+import scala.compat.java8.OptionConverters._
+import scala.concurrent.TimeoutException
+
+class ForwardingManager(channelManager: BrokerToControllerChannelManager,
+                        time: Time,
+                        retryTimeoutMs: Long) extends Logging {
+
+  def forwardRequest(request: RequestChannel.Request,
+                     responseCallback: AbstractResponse => Unit): Unit = {
+    val principalSerde = request.context.principalSerde.asScala.getOrElse(
+      throw new IllegalArgumentException(s"Cannot deserialize principal from request $request " +
+        "since there is no serde defined")
+    )
+    val serializedPrincipal = principalSerde.serialize(request.context.principal)
+    val forwardRequestBuffer = request.buffer.duplicate()
+    forwardRequestBuffer.flip()
+    val envelopeRequest = new EnvelopeRequest.Builder(
+      forwardRequestBuffer,
+      serializedPrincipal,
+      request.context.clientAddress.getAddress
+    )
+
+    class ForwardingResponseHandler extends ControllerRequestCompletionHandler {
+      override def onComplete(clientResponse: ClientResponse): Unit = {
+        val envelopeResponse = clientResponse.responseBody.asInstanceOf[EnvelopeResponse]
+        val envelopeError = envelopeResponse.error()
+        val requestBody = request.body[AbstractRequest]
+
+        val response = if (envelopeError != Errors.NONE) {
+          // An envelope error indicates broker misconfiguration (e.g. the principal serde
+          // might not be defined on the receiving broker). In this case, we do not return
+          // the error directly to the client since it would not be expected. Instead we
+          // return `UNKNOWN_SERVER_ERROR` so that the user knows that there is a problem
+          // on the broker.
+          debug(s"Forwarded request $request failed with an error in the envelope response $envelopeError")
+          requestBody.getErrorResponse(Errors.UNKNOWN_SERVER_ERROR.exception)
+        } else {
+          parseResponse(envelopeResponse.responseData, requestBody, request.header)
+        }
+        responseCallback(response)
+      }
+
+      override def onTimeout(): Unit = {
+        debug(s"Forwarding of the request $request failed due to timeout exception")
+        val response = request.body[AbstractRequest].getErrorResponse(new TimeoutException)
+        responseCallback(response)
+      }
+    }
+
+    val currentTime = time.milliseconds()
+    val deadlineMs =
+      if (Long.MaxValue - currentTime < retryTimeoutMs)
+        Long.MaxValue
+      else
+        currentTime + retryTimeoutMs
+
+    channelManager.sendRequest(envelopeRequest, new ForwardingResponseHandler, deadlineMs)
+  }
+
+  private def parseResponse(
+    buffer: ByteBuffer,
+    request: AbstractRequest,
+    header: RequestHeader
+  ): AbstractResponse = {
+    try {
+      AbstractResponse.parseResponse(buffer, header)
+    } catch {
+      case e: Exception =>
+        error(s"Failed to parse response from envelope for request with header $header", e)
+        request.getErrorResponse(Errors.UNKNOWN_SERVER_ERROR.exception)
+    }
+  }
+
+}
diff --git a/core/src/main/scala/kafka/server/KafkaApis.scala b/core/src/main/scala/kafka/server/KafkaApis.scala
index a31b6c31beba3..3e01fb5f5581d 100644
--- a/core/src/main/scala/kafka/server/KafkaApis.scala
+++ b/core/src/main/scala/kafka/server/KafkaApis.scala
@@ -17,50 +17,82 @@
 
 package kafka.server
 
-import java.nio.ByteBuffer
+import java.lang.{Byte => JByte}
 import java.lang.{Long => JLong}
-import java.util.{Collections, Properties}
+import java.nio.ByteBuffer
 import java.util
+import java.util.{Collections, Optional}
 import java.util.concurrent.ConcurrentHashMap
 import java.util.concurrent.atomic.AtomicInteger
-
 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, OffsetMetadata}
-import kafka.server.QuotaFactory.{QuotaManagers, UnboundedQuota}
-import kafka.controller.KafkaController
-import kafka.coordinator.group.{GroupCoordinator, JoinGroupResult}
+import kafka.common.OffsetAndMetadata
+import kafka.controller.{KafkaController, ReplicaAssignment}
+import kafka.coordinator.group.{GroupCoordinator, JoinGroupResult, LeaveGroupResult, SyncGroupResult}
 import kafka.coordinator.transaction.{InitProducerIdResult, TransactionCoordinator}
-import kafka.log.{Log, LogManager, TimestampOffset}
+import kafka.log.AppendOrigin
+import kafka.message.ZStdCompressionCodec
 import kafka.network.RequestChannel
-import kafka.network.RequestChannel.{CloseConnectionAction, NoOpAction, SendAction}
-import kafka.security.SecurityUtils
-import kafka.security.auth.{Resource, _}
+import kafka.security.authorizer.{AclEntry, AuthorizerUtils}
+import kafka.server.QuotaFactory.{QuotaManagers, UnboundedQuota}
 import kafka.utils.{CoreUtils, Logging}
+import kafka.utils.Implicits._
 import kafka.zk.{AdminZkClient, KafkaZkClient}
+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.{FatalExitError, Topic}
 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.CreatePartitionsResponseData.CreatePartitionsTopicResult
+import org.apache.kafka.common.message.{AddOffsetsToTxnResponseData, AlterClientQuotasResponseData, AlterConfigsResponseData, AlterPartitionReassignmentsResponseData, AlterReplicaLogDirsResponseData, CreateAclsResponseData, CreatePartitionsResponseData, CreateTopicsResponseData, DeleteAclsResponseData, DeleteGroupsResponseData, DeleteRecordsResponseData, DeleteTopicsResponseData, DescribeAclsResponseData, DescribeClientQuotasResponseData, DescribeConfigsResponseData, DescribeGroupsResponseData, DescribeLogDirsResponseData, EndTxnResponseData, ExpireDelegationTokenResponseData, FindCoordinatorResponseData, HeartbeatResponseData, InitProducerIdResponseData, JoinGroupResponseData, LeaveGroupResponseData, ListGroupsResponseData, ListOffsetResponseData, ListPartitionReassignmentsResponseData, MetadataResponseData, 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.CreateAclsResponseData.AclCreationResult
+import org.apache.kafka.common.message.DeleteTopicsResponseData.{DeletableTopicResult, DeletableTopicResultCollection}
+import org.apache.kafka.common.message.DeleteRecordsResponseData.{DeleteRecordsPartitionResult, DeleteRecordsTopicResult}
+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.message.ListOffsetRequestData.ListOffsetPartition
+import org.apache.kafka.common.message.ListOffsetResponseData.{ListOffsetPartitionResponse, ListOffsetTopicResponse}
+import org.apache.kafka.common.message.OffsetForLeaderEpochRequestData.OffsetForLeaderTopic
+import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData
+import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData.{EpochEndOffset, OffsetForLeaderTopicResult, OffsetForLeaderTopicResultCollection}
 import org.apache.kafka.common.metrics.Metrics
-import org.apache.kafka.common.network.ListenerName
+import org.apache.kafka.common.network.{ListenerName, Send}
 import org.apache.kafka.common.protocol.{ApiKeys, Errors}
-import org.apache.kafka.common.record.{ControlRecordType, EndTransactionMarker, MemoryRecords, RecordBatch, RecordsProcessingStats}
-import org.apache.kafka.common.requests.CreateAclsResponse.AclCreationResponse
-import org.apache.kafka.common.requests.DeleteAclsResponse.{AclDeletionResult, AclFilterResponse}
-import org.apache.kafka.common.requests.{Resource => RResource, ResourceType => RResourceType, _}
+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.FindCoordinatorRequest.CoordinatorType
 import org.apache.kafka.common.requests.ProduceResponse.PartitionResponse
-import org.apache.kafka.common.utils.{Time, Utils}
+import org.apache.kafka.common.requests._
+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.{ProducerIdAndEpoch, Time, Utils}
 import org.apache.kafka.common.{Node, TopicPartition}
-import org.apache.kafka.common.requests.{SaslAuthenticateResponse, SaslHandshakeResponse}
-import org.apache.kafka.common.resource.{Resource => AdminResource}
-import org.apache.kafka.common.acl.{AccessControlEntry, AclBinding}
-import DescribeLogDirsResponse.LogDirInfo
+import org.apache.kafka.common.message.AlterConfigsResponseData.AlterConfigsResourceResponse
+import org.apache.kafka.common.message.MetadataResponseData.{MetadataResponsePartition, MetadataResponseTopic}
+import org.apache.kafka.server.authorizer._
 
-import scala.collection.{mutable, _}
-import scala.collection.JavaConverters._
+import scala.compat.java8.OptionConverters._
+import scala.jdk.CollectionConverters._
 import scala.collection.mutable.ArrayBuffer
+import scala.collection.{Map, Seq, Set, immutable, mutable}
 import scala.util.{Failure, Success, Try}
+import kafka.coordinator.group.GroupOverview
+
+import scala.annotation.nowarn
 
 /**
  * Logic to handle the various Kafka requests
@@ -71,6 +103,7 @@ class KafkaApis(val requestChannel: RequestChannel,
                 val groupCoordinator: GroupCoordinator,
                 val txnCoordinator: TransactionCoordinator,
                 val controller: KafkaController,
+                val forwardingManager: ForwardingManager,
                 val zkClient: KafkaZkClient,
                 val brokerId: Int,
                 val config: KafkaConfig,
@@ -78,24 +111,87 @@ class KafkaApis(val requestChannel: RequestChannel,
                 val metrics: Metrics,
                 val authorizer: Option[Authorizer],
                 val quotas: QuotaManagers,
+                val fetchManager: FetchManager,
                 brokerTopicStats: BrokerTopicStats,
                 val clusterId: String,
-                time: Time) extends Logging {
+                time: Time,
+                val tokenManager: DelegationTokenManager,
+                val brokerFeatures: BrokerFeatures,
+                val finalizedFeatureCache: FinalizedFeatureCache) extends ApiRequestHandler with Logging {
 
+  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.")
   }
 
+  private def maybeHandleInvalidEnvelope(
+    envelope: RequestChannel.Request,
+    forwardedApiKey: ApiKeys
+  ): Boolean = {
+    def sendEnvelopeError(error: Errors): Unit = {
+      sendErrorResponseMaybeThrottle(envelope, error.exception)
+    }
+
+    if (!config.metadataQuorumEnabled || !envelope.context.fromPrivilegedListener) {
+      // If the designated forwarding request is not coming from a privileged listener, or
+      // forwarding is not enabled yet, we would not handle the request.
+      closeConnection(envelope, Collections.emptyMap())
+      true
+    } else if (!authorize(envelope.context, CLUSTER_ACTION, CLUSTER, CLUSTER_NAME)) {
+      // Forwarding request must have CLUSTER_ACTION authorization to reduce the risk of impersonation.
+      sendEnvelopeError(Errors.CLUSTER_AUTHORIZATION_FAILED)
+      true
+    } else if (!forwardedApiKey.forwardable) {
+      sendEnvelopeError(Errors.INVALID_REQUEST)
+      true
+    } else if (!controller.isActive) {
+      sendEnvelopeError(Errors.NOT_CONTROLLER)
+      true
+    } else {
+      false
+    }
+  }
+
+  private def isForwardingEnabled(request: RequestChannel.Request): Boolean = {
+    config.metadataQuorumEnabled && request.context.principalSerde.isPresent
+  }
+
+  private def maybeForward(
+    request: RequestChannel.Request,
+    handler: RequestChannel.Request => Unit
+  ): Unit = {
+    def responseCallback(response: AbstractResponse): Unit = {
+      sendForwardedResponse(request, response)
+    }
+
+    if (!request.isForwarded && !controller.isActive && isForwardingEnabled(request)) {
+      forwardingManager.forwardRequest(request, responseCallback)
+    } else {
+      // When the KIP-500 mode is off or the principal serde is undefined, forwarding is not supported,
+      // therefore requests are handled directly.
+      handler(request)
+    }
+  }
+
   /**
    * Top-level method that handles all requests and multiplexes to the right api
    */
-  def handle(request: RequestChannel.Request) {
+  override 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}")
+
+      request.envelope.foreach { envelope =>
+        if (maybeHandleInvalidEnvelope(envelope, request.header.apiKey)) {
+          return
+        }
+      }
+
       request.header.apiKey match {
         case ApiKeys.PRODUCE => handleProduceRequest(request)
         case ApiKeys.FETCH => handleFetchRequest(request)
@@ -116,8 +212,8 @@ class KafkaApis(val requestChannel: RequestChannel,
         case ApiKeys.LIST_GROUPS => handleListGroupsRequest(request)
         case ApiKeys.SASL_HANDSHAKE => handleSaslHandshakeRequest(request)
         case ApiKeys.API_VERSIONS => handleApiVersionsRequest(request)
-        case ApiKeys.CREATE_TOPICS => handleCreateTopicsRequest(request)
-        case ApiKeys.DELETE_TOPICS => handleDeleteTopicsRequest(request)
+        case ApiKeys.CREATE_TOPICS => maybeForward(request, handleCreateTopicsRequest)
+        case ApiKeys.DELETE_TOPICS => maybeForward(request, handleDeleteTopicsRequest)
         case ApiKeys.DELETE_RECORDS => handleDeleteRecordsRequest(request)
         case ApiKeys.INIT_PRODUCER_ID => handleInitProducerIdRequest(request)
         case ApiKeys.OFFSET_FOR_LEADER_EPOCH => handleOffsetForLeaderEpochRequest(request)
@@ -127,172 +223,276 @@ class KafkaApis(val requestChannel: RequestChannel,
         case ApiKeys.WRITE_TXN_MARKERS => handleWriteTxnMarkersRequest(request)
         case ApiKeys.TXN_OFFSET_COMMIT => handleTxnOffsetCommitRequest(request)
         case ApiKeys.DESCRIBE_ACLS => handleDescribeAcls(request)
-        case ApiKeys.CREATE_ACLS => handleCreateAcls(request)
-        case ApiKeys.DELETE_ACLS => handleDeleteAcls(request)
-        case ApiKeys.ALTER_CONFIGS => handleAlterConfigsRequest(request)
+        case ApiKeys.CREATE_ACLS => maybeForward(request, handleCreateAcls)
+        case ApiKeys.DELETE_ACLS => maybeForward(request, handleDeleteAcls)
+        case ApiKeys.ALTER_CONFIGS => maybeForward(request, handleAlterConfigsRequest)
         case ApiKeys.DESCRIBE_CONFIGS => handleDescribeConfigsRequest(request)
         case ApiKeys.ALTER_REPLICA_LOG_DIRS => handleAlterReplicaLogDirsRequest(request)
         case ApiKeys.DESCRIBE_LOG_DIRS => handleDescribeLogDirsRequest(request)
         case ApiKeys.SASL_AUTHENTICATE => handleSaslAuthenticateRequest(request)
-        case ApiKeys.CREATE_PARTITIONS => handleCreatePartitionsRequest(request)
+        case ApiKeys.CREATE_PARTITIONS => maybeForward(request, handleCreatePartitionsRequest)
+        case ApiKeys.CREATE_DELEGATION_TOKEN => maybeForward(request, handleCreateTokenRequest)
+        case ApiKeys.RENEW_DELEGATION_TOKEN => maybeForward(request, handleRenewTokenRequest)
+        case ApiKeys.EXPIRE_DELEGATION_TOKEN => maybeForward(request, handleExpireTokenRequest)
+        case ApiKeys.DESCRIBE_DELEGATION_TOKEN => handleDescribeTokensRequest(request)
+        case ApiKeys.DELETE_GROUPS => handleDeleteGroupsRequest(request)
+        case ApiKeys.ELECT_LEADERS => handleElectReplicaLeader(request)
+        case ApiKeys.INCREMENTAL_ALTER_CONFIGS => maybeForward(request, handleIncrementalAlterConfigsRequest)
+        case ApiKeys.ALTER_PARTITION_REASSIGNMENTS => maybeForward(request, handleAlterPartitionReassignmentsRequest)
+        case ApiKeys.LIST_PARTITION_REASSIGNMENTS => handleListPartitionReassignmentsRequest(request)
+        case ApiKeys.OFFSET_DELETE => handleOffsetDeleteRequest(request)
+        case ApiKeys.DESCRIBE_CLIENT_QUOTAS => handleDescribeClientQuotasRequest(request)
+        case ApiKeys.ALTER_CLIENT_QUOTAS => maybeForward(request, handleAlterClientQuotasRequest)
+        case ApiKeys.DESCRIBE_USER_SCRAM_CREDENTIALS => handleDescribeUserScramCredentialsRequest(request)
+        case ApiKeys.ALTER_USER_SCRAM_CREDENTIALS => maybeForward(request, handleAlterUserScramCredentialsRequest)
+        case ApiKeys.ALTER_ISR => handleAlterIsrRequest(request)
+        case ApiKeys.UPDATE_FEATURES => maybeForward(request, handleUpdateFeatures)
+        case ApiKeys.ENVELOPE => throw new IllegalStateException(
+          "Envelope request should not be handled directly in top level API")
+
+        // Until we are ready to integrate the Raft layer, these APIs are treated as
+        // unexpected and we just close the connection.
+        case ApiKeys.VOTE => closeConnection(request, util.Collections.emptyMap())
+        case ApiKeys.BEGIN_QUORUM_EPOCH => closeConnection(request, util.Collections.emptyMap())
+        case ApiKeys.END_QUORUM_EPOCH => closeConnection(request, util.Collections.emptyMap())
+        case ApiKeys.DESCRIBE_QUORUM => closeConnection(request, util.Collections.emptyMap())
       }
     } catch {
       case e: FatalExitError => throw e
       case e: Throwable => handleError(request, e)
     } finally {
-      request.apiLocalCompleteTimeNanos = time.nanoseconds
+      // try to complete delayed action. In order to avoid conflicting locking, the actions to complete delayed requests
+      // are kept in a queue. We add the logic to check the ReplicaManager queue at the end of KafkaApis.handle() and the
+      // expiration thread for certain delayed operations (e.g. DelayedJoin)
+      replicaManager.tryCompleteActions()
+      // The local completion time may be set while processing the request. Only record it if it's unset.
+      if (request.apiLocalCompleteTimeNanos < 0)
+        request.apiLocalCompleteTimeNanos = time.nanoseconds
     }
   }
 
-  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))
       }
     }
 
-    if (authorize(request.session, ClusterAction, Resource.ClusterResource)) {
+    authorizeClusterOperation(request, CLUSTER_ACTION)
+    if (isBrokerEpochStale(leaderAndIsrRequest.brokerEpoch)) {
+      // 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}")
+      sendResponseExemptThrottle(request, leaderAndIsrRequest.getErrorResponse(0, Errors.STALE_BROKER_EPOCH.exception))
+    } else {
       val response = replicaManager.becomeLeaderOrFollower(correlationId, leaderAndIsrRequest, onLeadershipChange)
       sendResponseExemptThrottle(request, response)
-    } else {
-      sendResponseMaybeThrottle(request, throttleTimeMs => leaderAndIsrRequest.getErrorResponse(throttleTimeMs,
-        Errors.CLUSTER_AUTHORIZATION_FAILED.exception))
     }
   }
 
-  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]
-
-    if (authorize(request.session, ClusterAction, Resource.ClusterResource)) {
-      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.
-      result.foreach { case (topicPartition, error) =>
-        if (error == Errors.NONE && stopReplicaRequest.deletePartitions && topicPartition.topic == GROUP_METADATA_TOPIC_NAME) {
-          groupCoordinator.handleGroupEmigration(topicPartition.partition)
+    authorizeClusterOperation(request, CLUSTER_ACTION)
+    if (isBrokerEpochStale(stopReplicaRequest.brokerEpoch)) {
+      // 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 StopReplica request with broker epoch " +
+        s"${stopReplicaRequest.brokerEpoch} smaller than the current broker epoch ${controller.brokerEpoch}")
+      sendResponseExemptThrottle(request, new StopReplicaResponse(
+        new StopReplicaResponseData().setErrorCode(Errors.STALE_BROKER_EPOCH.code)))
+    } else {
+      val partitionStates = stopReplicaRequest.partitionStates().asScala
+      val (result, error) = replicaManager.stopReplicas(
+        request.context.correlationId,
+        stopReplicaRequest.controllerId,
+        stopReplicaRequest.controllerEpoch,
+        stopReplicaRequest.brokerEpoch,
+        partitionStates)
+      // 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.forKeyValue { (topicPartition, error) =>
+        if (error == Errors.NONE) {
+          if (topicPartition.topic == GROUP_METADATA_TOPIC_NAME
+              && partitionStates(topicPartition).deletePartition) {
+            groupCoordinator.onResignation(topicPartition.partition)
+          } else if (topicPartition.topic == TRANSACTION_STATE_TOPIC_NAME
+                     && partitionStates(topicPartition).deletePartition) {
+            val partitionState = partitionStates(topicPartition)
+            val leaderEpoch = if (partitionState.leaderEpoch >= 0)
+                Some(partitionState.leaderEpoch)
+            else
+              None
+            txnCoordinator.onResignation(topicPartition.partition, coordinatorEpoch = leaderEpoch)
+          }
         }
       }
-      sendResponseExemptThrottle(request, new StopReplicaResponse(error, result.asJava))
-    } else {
-      val result = stopReplicaRequest.partitions.asScala.map((_, Errors.CLUSTER_AUTHORIZATION_FAILED)).toMap
-      sendResponseMaybeThrottle(request, _ =>
-        new StopReplicaResponse(Errors.CLUSTER_AUTHORIZATION_FAILED, 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]
 
-    if (authorize(request.session, ClusterAction, Resource.ClusterResource)) {
+    authorizeClusterOperation(request, CLUSTER_ACTION)
+    if (isBrokerEpochStale(updateMetadataRequest.brokerEpoch)) {
+      // 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(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.forEach { partitionState =>
+          adminManager.tryCompleteDelayedTopicOperations(partitionState.topicName)
         }
       }
-      sendResponseExemptThrottle(request, new UpdateMetadataResponse(Errors.NONE))
-    } else {
-      sendResponseMaybeThrottle(request, _ => new UpdateMetadataResponse(Errors.CLUSTER_AUTHORIZATION_FAILED))
+      quotas.clientQuotaCallback.foreach { callback =>
+        if (callback.updateClusterMetadata(metadataCache.getClusterMetadata(clusterId, request.context.listenerName))) {
+          quotas.fetch.updateQuotaMetricConfigs()
+          quotas.produce.updateQuotaMetricConfigs()
+          quotas.request.updateQuotaMetricConfigs()
+          quotas.controllerMutation.updateQuotaMetricConfigs()
+        }
+      }
+      if (replicaManager.hasDelayedElectionOperations) {
+        updateMetadataRequest.partitionStates.forEach { partitionState =>
+          val tp = new TopicPartition(partitionState.topicName, partitionState.partitionIndex)
+          replicaManager.tryCompleteElection(TopicPartitionOperationKey(tp))
+        }
+      }
+      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, 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.forKeyValue { (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, new Resource(Group, offsetCommitRequest.groupId))) {
+    if (!authorize(request.context, 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, new Resource(Topic, topicPartition.topic)))
-          unauthorizedTopicErrors += (topicPartition -> Errors.TOPIC_AUTHORIZATION_FAILED)
-        else if (!metadataCache.contains(topicPartition.topic))
-          nonExistingTopicErrors += (topicPartition -> Errors.UNKNOWN_TOPIC_OR_PARTITION)
-        else
-          authorizedTopicRequestInfoBldr += (topicPartition -> partitionData)
+      val authorizedTopicRequestInfoBldr = immutable.Map.newBuilder[TopicPartition, OffsetCommitRequestData.OffsetCommitRequestPartition]
+
+      val topics = offsetCommitRequest.data.topics.asScala
+      val authorizedTopics = filterByAuthorized(request.context, READ, TOPIC, topics)(_.name)
+      for (topicData <- topics) {
+        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) {
@@ -300,10 +500,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 {
@@ -314,90 +518,114 @@ class KafkaApis(val requestChannel: RequestChannel,
       } else {
         // for version 1 and beyond store offsets in offset manager
 
-        // compute the retention time based on the request version:
-        // if it is v1 or not specified by user, we can use the default retention
-        val offsetRetention =
-          if (header.apiVersion <= 1 ||
-            offsetCommitRequest.retentionTime == OffsetCommitRequest.DEFAULT_RETENTION_TIME)
-            groupCoordinator.offsetConfig.offsetsRetentionMs
-          else
-            offsetCommitRequest.retentionTime
-
-        // commit timestamp is always set to now.
         // "default" expiration timestamp is now + retention (and retention may be overridden if v2)
         // expire timestamp is computed differently for v1 and v2.
-        //   - If v1 and no explicit commit timestamp is provided we use default expiration timestamp.
-        //   - If v1 and explicit commit timestamp is provided we calculate retention from that explicit commit timestamp
-        //   - If v2 we use the default expiration timestamp
+        //   - If v1 and no explicit commit timestamp is provided we treat it the same as v5.
+        //   - If v1 and explicit retention time is provided we calculate expiration timestamp based on that
+        //   - 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 defaultExpireTimestamp = offsetRetention + currentTimestamp
-        val partitionData = authorizedTopicRequestInfo.mapValues { partitionData =>
-          val metadata = if (partitionData.metadata == null) OffsetMetadata.NoMetadata else partitionData.metadata
-          new OffsetAndMetadata(
-            offsetMetadata = OffsetMetadata(partitionData.offset, metadata),
-            commitTimestamp = currentTimestamp,
-            expireTimestamp = {
-              if (partitionData.timestamp == OffsetCommitRequest.DEFAULT_TIMESTAMP)
-                defaultExpireTimestamp
-              else
-                offsetRetention + partitionData.timestamp
+        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.commitTimestamp match {
+              case OffsetCommitRequest.DEFAULT_TIMESTAMP => currentTimestamp
+              case customTimestamp => customTimestamp
+            },
+            expireTimestamp = offsetCommitRequest.data.retentionTimeMs match {
+              case OffsetCommitRequest.DEFAULT_RETENTION_TIME => None
+              case retentionTime => Some(currentTimestamp + retentionTime)
             }
           )
         }
 
         // 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
+    val requestSize = request.sizeInBytes
 
-    if (produceRequest.isTransactional) {
-      if (!authorize(request.session, Write, new Resource(TransactionalId, produceRequest.transactionalId))) {
+    val (hasIdempotentRecords, hasTransactionalRecords) = {
+      val flags = RequestUtils.flags(produceRequest)
+      (flags.getKey, flags.getValue)
+    }
+    if (hasTransactionalRecords) {
+      val isAuthorizedTransactional = produceRequest.transactionalId != null &&
+        authorize(request.context, 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.isIdempotent && !authorize(request.session, IdempotentWrite, Resource.ClusterResource)) {
+    } else if (hasIdempotentRecords && !authorize(request.context, IDEMPOTENT_WRITE, CLUSTER, CLUSTER_NAME)) {
       sendErrorResponseMaybeThrottle(request, Errors.CLUSTER_AUTHORIZATION_FAILED.exception)
       return
     }
 
     val unauthorizedTopicResponses = mutable.Map[TopicPartition, PartitionResponse]()
     val nonExistingTopicResponses = mutable.Map[TopicPartition, PartitionResponse]()
+    val invalidRequestResponses = mutable.Map[TopicPartition, PartitionResponse]()
     val authorizedRequestInfo = mutable.Map[TopicPartition, MemoryRecords]()
-
-    for ((topicPartition, memoryRecords) <- produceRequest.partitionRecordsOrFail.asScala) {
-      if (!authorize(request.session, Write, new Resource(Topic, topicPartition.topic)))
+    // cache the result to avoid redundant authorization calls
+    val authorizedTopics = filterByAuthorized(request.context, WRITE, TOPIC,
+      produceRequest.data().topicData().asScala)(_.name())
+
+    produceRequest.data.topicData.forEach(topic => topic.partitionData.forEach { partition =>
+      val topicPartition = new TopicPartition(topic.name, partition.index)
+      // This caller assumes the type is MemoryRecords and that is true on current serialization
+      // We cast the type to avoid causing big change to code base.
+      // https://issues.apache.org/jira/browse/KAFKA-10698
+      val memoryRecords = partition.records.asInstanceOf[MemoryRecords]
+      if (!authorizedTopics.contains(topicPartition.topic))
         unauthorizedTopicResponses += topicPartition -> new PartitionResponse(Errors.TOPIC_AUTHORIZATION_FAILED)
-      else if (!metadataCache.contains(topicPartition.topic))
+      else if (!metadataCache.contains(topicPartition))
         nonExistingTopicResponses += topicPartition -> new PartitionResponse(Errors.UNKNOWN_TOPIC_OR_PARTITION)
       else
-        authorizedRequestInfo += (topicPartition -> memoryRecords)
-    }
+        try {
+          ProduceRequest.validateRecords(request.header.apiVersion, memoryRecords)
+          authorizedRequestInfo += (topicPartition -> memoryRecords)
+        } catch {
+          case e: ApiException =>
+            invalidRequestResponses += topicPartition -> new PartitionResponse(Errors.forException(e))
+        }
+    })
 
     // the callback for sending a produce response
-    def sendResponseCallback(responseStatus: Map[TopicPartition, PartitionResponse]) {
-
-      val mergedResponseStatus = responseStatus ++ unauthorizedTopicResponses ++ nonExistingTopicResponses
+    // The construction of ProduceResponse is able to accept auto-generated protocol data so
+    // KafkaApis#handleProduceRequest should apply auto-generated protocol to avoid extra conversion.
+    // https://issues.apache.org/jira/browse/KAFKA-10730
+    @nowarn("cat=deprecation")
+    def sendResponseCallback(responseStatus: Map[TopicPartition, PartitionResponse]): Unit = {
+      val mergedResponseStatus = responseStatus ++ unauthorizedTopicResponses ++ nonExistingTopicResponses ++ invalidRequestResponses
       var errorInResponse = false
 
-      mergedResponseStatus.foreach { case (topicPartition, status) =>
+      mergedResponseStatus.forKeyValue { (topicPartition, status) =>
         if (status.error != Errors.NONE) {
           errorInResponse = true
           debug("Produce request with correlation id %d from client %s on partition %s failed due to %s".format(
@@ -408,43 +636,52 @@ class KafkaApis(val requestChannel: RequestChannel,
         }
       }
 
-      def produceResponseCallback(bandwidthThrottleTimeMs: Int) {
-        if (produceRequest.acks == 0) {
-          // no operation needed if producer request.required.acks = 0; however, if there is any error in handling
-          // the request, since no response is expected by the producer, the server will close socket server so that
-          // the producer client will know that some error has happened and will refresh its metadata
-          if (errorInResponse) {
-            val exceptionsSummary = mergedResponseStatus.map { case (topicPartition, status) =>
-              topicPartition -> status.error.exceptionName
-            }.mkString(", ")
-            info(
-              s"Closing connection due to error during produce request with correlation id ${request.header.correlationId} " +
-                s"from client id ${request.header.clientId} with ack=0\n" +
-                s"Topic and partition to exceptions: $exceptionsSummary"
-            )
-            closeConnection(request, new ProduceResponse(mergedResponseStatus.asJava).errorCounts)
-          } else {
-            sendNoOpResponseExemptThrottle(request)
-          }
+      // Record both bandwidth and request quota-specific values and throttle by muting the channel if any of the quotas
+      // have been violated. If both quotas have been violated, use the max throttle time between the two quotas. Note
+      // that the request quota is not enforced if acks == 0.
+      val timeMs = time.milliseconds()
+      val bandwidthThrottleTimeMs = quotas.produce.maybeRecordAndGetThrottleTimeMs(request, requestSize, timeMs)
+      val requestThrottleTimeMs =
+        if (produceRequest.acks == 0) 0
+        else quotas.request.maybeRecordAndGetThrottleTimeMs(request, timeMs)
+      val maxThrottleTimeMs = Math.max(bandwidthThrottleTimeMs, requestThrottleTimeMs)
+      if (maxThrottleTimeMs > 0) {
+        request.apiThrottleTimeMs = maxThrottleTimeMs
+        if (bandwidthThrottleTimeMs > requestThrottleTimeMs) {
+          quotas.produce.throttle(request, bandwidthThrottleTimeMs, requestChannel.sendResponse)
         } else {
-          sendResponseMaybeThrottle(request, requestThrottleMs =>
-            new ProduceResponse(mergedResponseStatus.asJava, bandwidthThrottleTimeMs + requestThrottleMs))
+          quotas.request.throttle(request, requestThrottleTimeMs, requestChannel.sendResponse)
         }
       }
 
-      // When this callback is triggered, the remote API call has completed
-      request.apiRemoteCompleteTimeNanos = time.nanoseconds
-
-      quotas.produce.maybeRecordAndThrottle(
-        request.session.sanitizedUser,
-        request.header.clientId,
-        numBytesAppended,
-        produceResponseCallback)
+      // Send the response immediately. In case of throttling, the channel has already been muted.
+      if (produceRequest.acks == 0) {
+        // no operation needed if producer request.required.acks = 0; however, if there is any error in handling
+        // the request, since no response is expected by the producer, the server will close socket server so that
+        // the producer client will know that some error has happened and will refresh its metadata
+        if (errorInResponse) {
+          val exceptionsSummary = mergedResponseStatus.map { case (topicPartition, status) =>
+            topicPartition -> status.error.exceptionName
+          }.mkString(", ")
+          info(
+            s"Closing connection due to error during produce request with correlation id ${request.header.correlationId} " +
+              s"from client id ${request.header.clientId} with ack=0\n" +
+              s"Topic and partition to exceptions: $exceptionsSummary"
+          )
+          closeConnection(request, new ProduceResponse(mergedResponseStatus.asJava).errorCounts)
+        } else {
+          // Note that although request throttling is exempt for acks == 0, the channel may be throttled due to
+          // bandwidth quota violation.
+          sendNoOpResponseExemptThrottle(request)
+        }
+      } else {
+        sendResponse(request, Some(new ProduceResponse(mergedResponseStatus.asJava, maxThrottleTimeMs)), None)
+      }
     }
 
-    def processingStatsCallback(processingStats: Map[TopicPartition, RecordsProcessingStats]): Unit = {
-      processingStats.foreach { case (tp, info) =>
-        updateRecordsProcessingStats(request, tp, info)
+    def processingStatsCallback(processingStats: FetchResponseStats): Unit = {
+      processingStats.forKeyValue { (tp, info) =>
+        updateRecordConversionStats(request, tp, info)
       }
     }
 
@@ -458,10 +695,10 @@ class KafkaApis(val requestChannel: RequestChannel,
         timeout = produceRequest.timeout.toLong,
         requiredAcks = produceRequest.acks,
         internalTopicsAllowed = internalTopicsAllowed,
-        isFromClient = true,
+        origin = AppendOrigin.Client,
         entriesPerPartition = authorizedRequestInfo,
         responseCallback = sendResponseCallback,
-        processingStatsCallback = processingStatsCallback)
+        recordConversionStatsCallback = processingStatsCallback)
 
       // if the request is put into the purgatory, it will have a held reference and hence cannot be garbage collected;
       // hence we clear its data here in order to let GC reclaim its memory since it is already appended to log
@@ -472,379 +709,496 @@ class KafkaApis(val requestChannel: RequestChannel,
   /**
    * Handle a fetch request
    */
-  def handleFetchRequest(request: RequestChannel.Request) {
-    val fetchRequest = request.body[FetchRequest]
+  def handleFetchRequest(request: RequestChannel.Request): Unit = {
     val versionId = request.header.apiVersion
     val clientId = request.header.clientId
+    val fetchRequest = request.body[FetchRequest]
+    val fetchContext = fetchManager.newContext(
+      fetchRequest.metadata,
+      fetchRequest.fetchData,
+      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
+    }
 
-    val unauthorizedTopicResponseData = mutable.ArrayBuffer[(TopicPartition, FetchResponse.PartitionData)]()
-    val nonExistingTopicResponseData = mutable.ArrayBuffer[(TopicPartition, FetchResponse.PartitionData)]()
-    val authorizedRequestInfo = mutable.ArrayBuffer[(TopicPartition, FetchRequest.PartitionData)]()
-
-    for ((topicPartition, partitionData) <- fetchRequest.fetchData.asScala) {
-      if (!authorize(request.session, Read, new Resource(Topic, topicPartition.topic)))
-        unauthorizedTopicResponseData += topicPartition -> new FetchResponse.PartitionData(Errors.TOPIC_AUTHORIZATION_FAILED,
-          FetchResponse.INVALID_HIGHWATERMARK, FetchResponse.INVALID_LAST_STABLE_OFFSET,
-          FetchResponse.INVALID_LOG_START_OFFSET, null, MemoryRecords.EMPTY)
-      else if (!metadataCache.contains(topicPartition.topic))
-        nonExistingTopicResponseData += topicPartition -> new FetchResponse.PartitionData(Errors.UNKNOWN_TOPIC_OR_PARTITION,
-          FetchResponse.INVALID_HIGHWATERMARK, FetchResponse.INVALID_LAST_STABLE_OFFSET,
-          FetchResponse.INVALID_LOG_START_OFFSET, null, MemoryRecords.EMPTY)
-      else
-        authorizedRequestInfo += (topicPartition -> partitionData)
-    }
-
-    def convertedPartitionData(tp: TopicPartition, data: FetchResponse.PartitionData) = {
-
-      // Down-conversion of the fetched records is needed when the stored magic version is
-      // greater than that supported by the client (as indicated by the fetch request version). If the
-      // configured magic version for the topic is less than or equal to that supported by the version of the
-      // fetch request, we skip the iteration through the records in order to check the magic version since we
-      // know it must be supported. However, if the magic version is changed from a higher version back to a
-      // lower version, this check will no longer be valid and we will fail to down-convert the messages
-      // which were written in the new format prior to the version downgrade.
-      replicaManager.getMagic(tp).flatMap { magic =>
-        val downConvertMagic = {
-          if (magic > RecordBatch.MAGIC_VALUE_V0 && versionId <= 1 && !data.records.hasCompatibleMagic(RecordBatch.MAGIC_VALUE_V0))
-            Some(RecordBatch.MAGIC_VALUE_V0)
-          else if (magic > RecordBatch.MAGIC_VALUE_V1 && versionId <= 3 && !data.records.hasCompatibleMagic(RecordBatch.MAGIC_VALUE_V1))
-            Some(RecordBatch.MAGIC_VALUE_V1)
+    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)
+    }
+
+    val erroneous = mutable.ArrayBuffer[(TopicPartition, FetchResponse.PartitionData[Records])]()
+    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.context, CLUSTER_ACTION, CLUSTER, CLUSTER_NAME)) {
+        fetchContext.foreachPartition { (topicPartition, data) =>
+          if (!metadataCache.contains(topicPartition))
+            erroneous += topicPartition -> errorResponse(Errors.UNKNOWN_TOPIC_OR_PARTITION)
           else
-            None
+            interesting += (topicPartition -> data)
         }
-
-        downConvertMagic.map { magic =>
-          trace(s"Down converting records from partition $tp to message format version $magic for fetch request from $clientId")
-          val converted = data.records.downConvert(magic, fetchRequest.fetchData.get(tp).fetchOffset, time)
-          updateRecordsProcessingStats(request, tp, converted.recordsProcessingStats)
-          new FetchResponse.PartitionData(data.error, data.highWatermark, FetchResponse.INVALID_LAST_STABLE_OFFSET,
-            data.logStartOffset, data.abortedTransactions, converted.records)
+      } else {
+        fetchContext.foreachPartition { (part, _) =>
+          erroneous += part -> errorResponse(Errors.TOPIC_AUTHORIZATION_FAILED)
         }
-
-      }.getOrElse(data)
+      }
+    } else {
+      // Regular Kafka consumers need READ permission on each partition they are fetching.
+      val partitionDatas = new mutable.ArrayBuffer[(TopicPartition, FetchRequest.PartitionData)]
+      fetchContext.foreachPartition { (topicPartition, partitionData) =>
+        partitionDatas += topicPartition -> partitionData
+      }
+      val authorizedTopics = filterByAuthorized(request.context, READ, TOPIC, partitionDatas)(_._1.topic)
+      partitionDatas.foreach { case (topicPartition, data) =>
+        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)
+        else
+          interesting += (topicPartition -> data)
+      }
     }
 
-    // the callback for process a fetch response, invoked before throttling
-    def processResponseCallback(responsePartitionData: Seq[(TopicPartition, FetchPartitionData)]) {
-      val partitionData = {
-        responsePartitionData.map { case (tp, data) =>
-          val abortedTransactions = data.abortedTransactions.map(_.asJava).orNull
-          val lastStableOffset = data.lastStableOffset.getOrElse(FetchResponse.INVALID_LAST_STABLE_OFFSET)
-          tp -> new FetchResponse.PartitionData(data.error, data.highWatermark, lastStableOffset,
-            data.logStartOffset, abortedTransactions, data.records)
-        }
+    def maybeDownConvertStorageError(error: Errors, version: Short): Errors = {
+      // If consumer sends FetchRequest V5 or earlier, the client library is not guaranteed to recognize the error code
+      // for KafkaStorageException. In this case the client library will translate KafkaStorageException to
+      // UnknownServerException which is not retriable. We can ensure that consumer will update metadata and retry
+      // by converting the KafkaStorageException to NotLeaderForPartitionException in the response if FetchRequest version <= 5
+      if (error == Errors.KAFKA_STORAGE_ERROR && versionId <= 5) {
+        Errors.NOT_LEADER_OR_FOLLOWER
+      } else {
+        error
       }
+    }
 
-      val mergedPartitionData = partitionData ++ unauthorizedTopicResponseData ++ nonExistingTopicResponseData
-      val fetchedPartitionData = new util.LinkedHashMap[TopicPartition, FetchResponse.PartitionData]()
+    def maybeConvertFetchedData(tp: TopicPartition,
+                                partitionData: FetchResponse.PartitionData[Records]): FetchResponse.PartitionData[BaseRecords] = {
+      val logConfig = replicaManager.getLogConfig(tp)
 
-      mergedPartitionData.foreach { case (topicPartition, data) =>
-        if (data.error != Errors.NONE)
-          debug(s"Fetch request with correlation id ${request.header.correlationId} from client $clientId " +
-            s"on partition $topicPartition failed due to ${data.error.exceptionName}")
+      if (logConfig.exists(_.compressionType == ZStdCompressionCodec.name) && versionId < 10) {
+        trace(s"Fetching messages is disabled for ZStandard compressed partition $tp. Sending unsupported version response to $clientId.")
+        errorResponse(Errors.UNSUPPORTED_COMPRESSION_TYPE)
+      } else {
+        // Down-conversion of the fetched records is needed when the stored magic version is
+        // greater than that supported by the client (as indicated by the fetch request version). If the
+        // configured magic version for the topic is less than or equal to that supported by the version of the
+        // fetch request, we skip the iteration through the records in order to check the magic version since we
+        // know it must be supported. However, if the magic version is changed from a higher version back to a
+        // lower version, this check will no longer be valid and we will fail to down-convert the messages
+        // which were written in the new format prior to the version downgrade.
+        val unconvertedRecords = partitionData.records
+        val downConvertMagic =
+          logConfig.map(_.messageFormatVersion.recordVersion.value).flatMap { magic =>
+            if (magic > RecordBatch.MAGIC_VALUE_V0 && versionId <= 1 && !unconvertedRecords.hasCompatibleMagic(RecordBatch.MAGIC_VALUE_V0))
+              Some(RecordBatch.MAGIC_VALUE_V0)
+            else if (magic > RecordBatch.MAGIC_VALUE_V1 && versionId <= 3 && !unconvertedRecords.hasCompatibleMagic(RecordBatch.MAGIC_VALUE_V1))
+              Some(RecordBatch.MAGIC_VALUE_V1)
+            else
+              None
+          }
 
-        fetchedPartitionData.put(topicPartition, data)
+        downConvertMagic match {
+          case Some(magic) =>
+            // For fetch requests from clients, check if down-conversion is disabled for the particular partition
+            if (!fetchRequest.isFromFollower && !logConfig.forall(_.messageDownConversionEnable)) {
+              trace(s"Conversion to message format ${downConvertMagic.get} is disabled for partition $tp. Sending unsupported version response to $clientId.")
+              errorResponse(Errors.UNSUPPORTED_VERSION)
+            } else {
+              try {
+                trace(s"Down converting records from partition $tp to message format version $magic for fetch request from $clientId")
+                // Because down-conversion is extremely memory intensive, we want to try and delay the down-conversion as much
+                // 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.
+                val error = maybeDownConvertStorageError(partitionData.error, versionId)
+                new FetchResponse.PartitionData[BaseRecords](error, partitionData.highWatermark,
+                  partitionData.lastStableOffset, partitionData.logStartOffset,
+                  partitionData.preferredReadReplica, partitionData.abortedTransactions,
+                  new LazyDownConversionRecords(tp, unconvertedRecords, magic, fetchContext.getFetchOffset(tp).get, time))
+              } catch {
+                case e: UnsupportedCompressionTypeException =>
+                  trace("Received unsupported compression type error during down-conversion", e)
+                  errorResponse(Errors.UNSUPPORTED_COMPRESSION_TYPE)
+              }
+            }
+          case None =>
+            val error = maybeDownConvertStorageError(partitionData.error, versionId)
+            new FetchResponse.PartitionData[BaseRecords](error,
+              partitionData.highWatermark,
+              partitionData.lastStableOffset,
+              partitionData.logStartOffset,
+              partitionData.preferredReadReplica,
+              partitionData.abortedTransactions,
+              partitionData.divergingEpoch,
+              unconvertedRecords)
+        }
       }
+    }
 
-      // fetch response callback invoked after any throttling
-      def fetchResponseCallback(bandwidthThrottleTimeMs: Int) {
-        def createResponse(requestThrottleTimeMs: Int): FetchResponse = {
-          val convertedData = new util.LinkedHashMap[TopicPartition, FetchResponse.PartitionData]
-          fetchedPartitionData.asScala.foreach { case (tp, partitionData) =>
-            convertedData.put(tp, convertedPartitionData(tp, partitionData))
-          }
-          val response = new FetchResponse(convertedData, bandwidthThrottleTimeMs + requestThrottleTimeMs)
-          response.responseData.asScala.foreach { case (topicPartition, data) =>
-            // record the bytes out metrics only when the response is being sent
-            brokerTopicStats.updateBytesOut(topicPartition.topic, fetchRequest.isFromFollower, data.records.sizeInBytes)
-          }
-          response
+    // the callback for process a fetch response, invoked before throttling
+    def processResponseCallback(responsePartitionData: Seq[(TopicPartition, FetchPartitionData)]): Unit = {
+      val partitions = new util.LinkedHashMap[TopicPartition, FetchResponse.PartitionData[Records]]
+      val reassigningPartitions = mutable.Set[TopicPartition]()
+      responsePartitionData.foreach { case (tp, data) =>
+        val abortedTransactions = data.abortedTransactions.map(_.asJava).orNull
+        val lastStableOffset = data.lastStableOffset.getOrElse(FetchResponse.INVALID_LAST_STABLE_OFFSET)
+        if (data.isReassignmentFetch)
+          reassigningPartitions.add(tp)
+        val error = maybeDownConvertStorageError(data.error, versionId)
+        partitions.put(tp, new FetchResponse.PartitionData(
+          error,
+          data.highWatermark,
+          lastStableOffset,
+          data.logStartOffset,
+          data.preferredReadReplica.map(int2Integer).asJava,
+          abortedTransactions,
+          data.divergingEpoch.asJava,
+          data.records))
+      }
+      erroneous.foreach { case (tp, data) => partitions.put(tp, data) }
+
+      var unconvertedFetchResponse: FetchResponse[Records] = null
+
+      def createResponse(throttleTimeMs: Int): FetchResponse[BaseRecords] = {
+        // Down-convert messages for each partition if required
+        val convertedData = new util.LinkedHashMap[TopicPartition, FetchResponse.PartitionData[BaseRecords]]
+        unconvertedFetchResponse.responseData.forEach { (tp, unconvertedPartitionData) =>
+          if (unconvertedPartitionData.error != Errors.NONE)
+            debug(s"Fetch request with correlation id ${request.header.correlationId} from client $clientId " +
+              s"on partition $tp failed due to ${unconvertedPartitionData.error.exceptionName}")
+          convertedData.put(tp, maybeConvertFetchedData(tp, unconvertedPartitionData))
         }
 
-        if (fetchRequest.isFromFollower)
-          sendResponseExemptThrottle(request, createResponse(0))
-        else
-          sendResponseMaybeThrottle(request, requestThrottleMs => createResponse(requestThrottleMs))
+        // Prepare fetch response from converted data
+        val response = new FetchResponse(unconvertedFetchResponse.error, convertedData, throttleTimeMs,
+          unconvertedFetchResponse.sessionId)
+        // record the bytes out metrics only when the response is being sent
+        response.responseData.forEach { (tp, data) =>
+          brokerTopicStats.updateBytesOut(tp.topic, fetchRequest.isFromFollower, reassigningPartitions.contains(tp), data.records.sizeInBytes)
+        }
+        response
       }
 
-      // When this callback is triggered, the remote API call has completed.
-      // Record time before any byte-rate throttling.
-      request.apiRemoteCompleteTimeNanos = time.nanoseconds
+      def updateConversionStats(send: Send): Unit = {
+        send match {
+          case send: MultiRecordsSend if send.recordConversionStats != null =>
+            send.recordConversionStats.asScala.toMap.foreach {
+              case (tp, stats) => updateRecordConversionStats(request, tp, stats)
+            }
+          case _ =>
+        }
+      }
 
       if (fetchRequest.isFromFollower) {
         // We've already evaluated against the quota and are good to go. Just need to record it now.
-        val responseSize = sizeOfThrottledPartitions(versionId, fetchRequest, mergedPartitionData, quotas.leader)
+        unconvertedFetchResponse = fetchContext.updateAndGenerateResponseData(partitions)
+        val responseSize = sizeOfThrottledPartitions(versionId, unconvertedFetchResponse, quotas.leader)
         quotas.leader.record(responseSize)
-        fetchResponseCallback(bandwidthThrottleTimeMs = 0)
+        trace(s"Sending Fetch response with partitions.size=${unconvertedFetchResponse.responseData.size}, " +
+          s"metadata=${unconvertedFetchResponse.sessionId}")
+        sendResponseExemptThrottle(request, createResponse(0), Some(updateConversionStats))
       } else {
         // Fetch size used to determine throttle time is calculated before any down conversions.
         // This may be slightly different from the actual response size. But since down conversions
-        // result in data being loaded into memory, it is better to do this after throttling to avoid OOM.
-        val response = new FetchResponse(fetchedPartitionData, 0)
-        val responseStruct = response.toStruct(versionId)
-        quotas.fetch.maybeRecordAndThrottle(request.session.sanitizedUser, clientId, responseStruct.sizeOf,
-          fetchResponseCallback)
+        // result in data being loaded into memory, we should do this only when we are not going to throttle.
+        //
+        // Record both bandwidth and request quota-specific values and throttle by muting the channel if any of the
+        // quotas have been violated. If both quotas have been violated, use the max throttle time between the two
+        // quotas. When throttled, we unrecord the recorded bandwidth quota value
+        val responseSize = fetchContext.getResponseSize(partitions, versionId)
+        val timeMs = time.milliseconds()
+        val requestThrottleTimeMs = quotas.request.maybeRecordAndGetThrottleTimeMs(request, timeMs)
+        val bandwidthThrottleTimeMs = quotas.fetch.maybeRecordAndGetThrottleTimeMs(request, responseSize, timeMs)
+
+        val maxThrottleTimeMs = math.max(bandwidthThrottleTimeMs, requestThrottleTimeMs)
+        if (maxThrottleTimeMs > 0) {
+          request.apiThrottleTimeMs = maxThrottleTimeMs
+          // Even if we need to throttle for request quota violation, we should "unrecord" the already recorded value
+          // from the fetch quota because we are going to return an empty response.
+          quotas.fetch.unrecordQuotaSensor(request, responseSize, timeMs)
+          if (bandwidthThrottleTimeMs > requestThrottleTimeMs) {
+            quotas.fetch.throttle(request, bandwidthThrottleTimeMs, requestChannel.sendResponse)
+          } else {
+            quotas.request.throttle(request, requestThrottleTimeMs, requestChannel.sendResponse)
+          }
+          // If throttling is required, return an empty response.
+          unconvertedFetchResponse = fetchContext.getThrottledResponse(maxThrottleTimeMs)
+        } else {
+          // Get the actual response. This will update the fetch context.
+          unconvertedFetchResponse = fetchContext.updateAndGenerateResponseData(partitions)
+          trace(s"Sending Fetch response with partitions.size=$responseSize, metadata=${unconvertedFetchResponse.sessionId}")
+        }
+
+        // Send the response immediately.
+        sendResponse(request, Some(createResponse(maxThrottleTimeMs)), Some(updateConversionStats))
       }
     }
 
-    if (authorizedRequestInfo.isEmpty)
+    // 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(Math.min(fetchRequest.maxBytes, config.fetchMaxBytes), maxQuotaWindowBytes)
+    val fetchMinBytes = Math.min(fetchRequest.minBytes, fetchMaxBytes)
+    if (interesting.isEmpty)
       processResponseCallback(Seq.empty)
     else {
       // call the replica manager to fetch messages from the local replica
       replicaManager.fetchMessages(
         fetchRequest.maxWait.toLong,
         fetchRequest.replicaId,
-        fetchRequest.minBytes,
-        fetchRequest.maxBytes,
+        fetchMinBytes,
+        fetchMaxBytes,
         versionId <= 2,
-        authorizedRequestInfo,
+        interesting,
         replicationQuota(fetchRequest),
         processResponseCallback,
-        fetchRequest.isolationLevel)
+        fetchRequest.isolationLevel,
+        clientMetadata)
+    }
+  }
+
+  class SelectingIterator(val partitions: util.LinkedHashMap[TopicPartition, FetchResponse.PartitionData[Records]],
+                          val quota: ReplicationQuotaManager)
+                          extends util.Iterator[util.Map.Entry[TopicPartition, FetchResponse.PartitionData[Records]]] {
+    val iter = partitions.entrySet().iterator()
+
+    var nextElement: util.Map.Entry[TopicPartition, FetchResponse.PartitionData[Records]] = null
+
+    override def hasNext: Boolean = {
+      while ((nextElement == null) && iter.hasNext()) {
+        val element = iter.next()
+        if (quota.isThrottled(element.getKey)) {
+          nextElement = element
+        }
+      }
+      nextElement != null
+    }
+
+    override def next(): util.Map.Entry[TopicPartition, FetchResponse.PartitionData[Records]] = {
+      if (!hasNext()) throw new NoSuchElementException()
+      val element = nextElement
+      nextElement = null
+      element
     }
+
+    override def remove(): Unit = throw new UnsupportedOperationException()
   }
 
+  // Traffic from both in-sync and out of sync replicas are accounted for in replication quota to ensure total replication
+  // traffic doesn't exceed quota.
   private def sizeOfThrottledPartitions(versionId: Short,
-                                        fetchRequest: FetchRequest,
-                                        mergedPartitionData: Seq[(TopicPartition, FetchResponse.PartitionData)],
+                                        unconvertedResponse: FetchResponse[Records],
                                         quota: ReplicationQuotaManager): Int = {
-    val partitionData = new util.LinkedHashMap[TopicPartition, FetchResponse.PartitionData]
-    mergedPartitionData.foreach { case (tp, data) =>
-      if (quota.isThrottled(tp))
-        partitionData.put(tp, data)
-    }
-    FetchResponse.sizeOf(versionId, partitionData)
+    val iter = new SelectingIterator(unconvertedResponse.responseData, quota)
+    FetchResponse.sizeOf(versionId, iter)
   }
 
   def replicationQuota(fetchRequest: FetchRequest): ReplicaQuota =
     if (fetchRequest.isFromFollower) quotas.leader else UnboundedQuota
 
-  def handleListOffsetRequest(request: RequestChannel.Request) {
-    val version = request.header.apiVersion()
+  def handleListOffsetRequest(request: RequestChannel.Request): Unit = {
+    val version = request.header.apiVersion
 
-    val mergedResponseMap = if (version == 0)
+    val topics = if (version == 0)
       handleListOffsetRequestV0(request)
     else
       handleListOffsetRequestV1AndAbove(request)
 
-    sendResponseMaybeThrottle(request, requestThrottleMs => new ListOffsetResponse(requestThrottleMs, mergedResponseMap.asJava))
+    sendResponseMaybeThrottle(request, requestThrottleMs => new ListOffsetResponse(new ListOffsetResponseData()
+      .setThrottleTimeMs(requestThrottleMs)
+      .setTopics(topics.asJava)))
   }
 
-  private def handleListOffsetRequestV0(request : RequestChannel.Request) : Map[TopicPartition, ListOffsetResponse.PartitionData] = {
+  private def handleListOffsetRequestV0(request : RequestChannel.Request) : List[ListOffsetTopicResponse] = {
     val correlationId = request.header.correlationId
     val clientId = request.header.clientId
     val offsetRequest = request.body[ListOffsetRequest]
 
-    val (authorizedRequestInfo, unauthorizedRequestInfo) = offsetRequest.offsetData.asScala.partition {
-      case (topicPartition, _) => authorize(request.session, Describe, new Resource(Topic, topicPartition.topic))
-    }
+    val (authorizedRequestInfo, unauthorizedRequestInfo) = partitionSeqByAuthorized(request.context,
+        DESCRIBE, TOPIC, offsetRequest.topics.asScala.toSeq)(_.name)
 
-    val unauthorizedResponseStatus = unauthorizedRequestInfo.mapValues(_ =>
-      new ListOffsetResponse.PartitionData(Errors.TOPIC_AUTHORIZATION_FAILED, List[JLong]().asJava)
+    val unauthorizedResponseStatus = unauthorizedRequestInfo.map(topic =>
+      new ListOffsetTopicResponse()
+        .setName(topic.name)
+        .setPartitions(topic.partitions.asScala.map(partition =>
+          new ListOffsetPartitionResponse()
+            .setPartitionIndex(partition.partitionIndex)
+            .setErrorCode(Errors.TOPIC_AUTHORIZATION_FAILED.code)).asJava)
     )
 
-    val responseMap = authorizedRequestInfo.map {case (topicPartition, partitionData) =>
-      try {
-        // ensure leader exists
-        val localReplica = if (offsetRequest.replicaId != ListOffsetRequest.DEBUGGING_REPLICA_ID)
-          replicaManager.getLeaderReplicaIfLocal(topicPartition)
-        else
-          replicaManager.getReplicaOrException(topicPartition)
-        val offsets = {
-          val allOffsets = fetchOffsets(replicaManager.logManager,
-            topicPartition,
-            partitionData.timestamp,
-            partitionData.maxNumOffsets)
-          if (offsetRequest.replicaId != ListOffsetRequest.CONSUMER_REPLICA_ID) {
-            allOffsets
-          } else {
-            val hw = localReplica.highWatermark.messageOffset
-            if (allOffsets.exists(_ > hw))
-              hw +: allOffsets.dropWhile(_ > hw)
-            else
-              allOffsets
-          }
+    val responseTopics = authorizedRequestInfo.map { topic =>
+      val responsePartitions = topic.partitions.asScala.map { partition =>
+        val topicPartition = new TopicPartition(topic.name, partition.partitionIndex)
+
+        try {
+          val offsets = replicaManager.legacyFetchOffsetsForTimestamp(
+            topicPartition = topicPartition,
+            timestamp = partition.timestamp,
+            maxNumOffsets = partition.maxNumOffsets,
+            isFromConsumer = offsetRequest.replicaId == ListOffsetRequest.CONSUMER_REPLICA_ID,
+            fetchOnlyFromLeader = offsetRequest.replicaId != ListOffsetRequest.DEBUGGING_REPLICA_ID)
+          new ListOffsetPartitionResponse()
+            .setPartitionIndex(partition.partitionIndex)
+            .setErrorCode(Errors.NONE.code)
+            .setOldStyleOffsets(offsets.map(JLong.valueOf).asJava)
+        } catch {
+          // NOTE: UnknownTopicOrPartitionException and NotLeaderOrFollowerException are special cases since these error messages
+          // are typically transient and there is no value in logging the entire stack trace for the same
+          case e @ (_ : UnknownTopicOrPartitionException |
+                    _ : NotLeaderOrFollowerException |
+                    _ : KafkaStorageException) =>
+            debug("Offset request with correlation id %d from client %s on partition %s failed due to %s".format(
+              correlationId, clientId, topicPartition, e.getMessage))
+            new ListOffsetPartitionResponse()
+              .setPartitionIndex(partition.partitionIndex)
+              .setErrorCode(Errors.forException(e).code)
+          case e: Throwable =>
+            error("Error while responding to offset request", e)
+            new ListOffsetPartitionResponse()
+              .setPartitionIndex(partition.partitionIndex)
+              .setErrorCode(Errors.forException(e).code)
         }
-        (topicPartition, new ListOffsetResponse.PartitionData(Errors.NONE, offsets.map(new JLong(_)).asJava))
-      } catch {
-        // NOTE: UnknownTopicOrPartitionException and NotLeaderForPartitionException are special cased since these error messages
-        // are typically transient and there is no value in logging the entire stack trace for the same
-        case e @ (_ : UnknownTopicOrPartitionException |
-                  _ : NotLeaderForPartitionException |
-                  _ : KafkaStorageException) =>
-          debug("Offset request with correlation id %d from client %s on partition %s failed due to %s".format(
-            correlationId, clientId, topicPartition, e.getMessage))
-          (topicPartition, new ListOffsetResponse.PartitionData(Errors.forException(e), List[JLong]().asJava))
-        case e: Throwable =>
-          error("Error while responding to offset request", e)
-          (topicPartition, new ListOffsetResponse.PartitionData(Errors.forException(e), List[JLong]().asJava))
       }
+      new ListOffsetTopicResponse().setName(topic.name).setPartitions(responsePartitions.asJava)
     }
-    responseMap ++ unauthorizedResponseStatus
+    (responseTopics ++ unauthorizedResponseStatus).toList
   }
 
-  private def handleListOffsetRequestV1AndAbove(request : RequestChannel.Request): Map[TopicPartition, ListOffsetResponse.PartitionData] = {
+  private def handleListOffsetRequestV1AndAbove(request : RequestChannel.Request): List[ListOffsetTopicResponse] = {
     val correlationId = request.header.correlationId
     val clientId = request.header.clientId
     val offsetRequest = request.body[ListOffsetRequest]
+    val version = request.header.apiVersion
 
-    val (authorizedRequestInfo, unauthorizedRequestInfo) = offsetRequest.partitionTimestamps.asScala.partition {
-      case (topicPartition, _) => authorize(request.session, Describe, new Resource(Topic, topicPartition.topic))
+    def buildErrorResponse(e: Errors, partition: ListOffsetPartition): ListOffsetPartitionResponse = {
+      new ListOffsetPartitionResponse()
+        .setPartitionIndex(partition.partitionIndex)
+        .setErrorCode(e.code)
+        .setTimestamp(ListOffsetResponse.UNKNOWN_TIMESTAMP)
+        .setOffset(ListOffsetResponse.UNKNOWN_OFFSET)
     }
 
-    val unauthorizedResponseStatus = unauthorizedRequestInfo.mapValues(_ => {
-      new ListOffsetResponse.PartitionData(Errors.TOPIC_AUTHORIZATION_FAILED,
-                                           ListOffsetResponse.UNKNOWN_TIMESTAMP,
-                                           ListOffsetResponse.UNKNOWN_OFFSET)
-    })
+    val (authorizedRequestInfo, unauthorizedRequestInfo) = partitionSeqByAuthorized(request.context,
+        DESCRIBE, TOPIC, offsetRequest.topics.asScala.toSeq)(_.name)
 
-    val responseMap = authorizedRequestInfo.map { case (topicPartition, timestamp) =>
-      if (offsetRequest.duplicatePartitions().contains(topicPartition)) {
-        debug(s"OffsetRequest with correlation id $correlationId from client $clientId on partition $topicPartition " +
-            s"failed because the partition is duplicated in the request.")
-        (topicPartition, new ListOffsetResponse.PartitionData(Errors.INVALID_REQUEST,
-                                                              ListOffsetResponse.UNKNOWN_TIMESTAMP,
-                                                              ListOffsetResponse.UNKNOWN_OFFSET))
-      } else {
-        try {
-          // ensure leader exists
-          val localReplica = if (offsetRequest.replicaId != ListOffsetRequest.DEBUGGING_REPLICA_ID)
-            replicaManager.getLeaderReplicaIfLocal(topicPartition)
-          else
-            replicaManager.getReplicaOrException(topicPartition)
+    val unauthorizedResponseStatus = unauthorizedRequestInfo.map(topic =>
+      new ListOffsetTopicResponse()
+        .setName(topic.name)
+        .setPartitions(topic.partitions.asScala.map(partition =>
+          buildErrorResponse(Errors.TOPIC_AUTHORIZATION_FAILED, partition)).asJava)
+    )
 
-          val fromConsumer = offsetRequest.replicaId == ListOffsetRequest.CONSUMER_REPLICA_ID
-          val found = if (fromConsumer) {
-            val lastFetchableOffset = offsetRequest.isolationLevel match {
-              case IsolationLevel.READ_COMMITTED => localReplica.lastStableOffset.messageOffset
-              case IsolationLevel.READ_UNCOMMITTED => localReplica.highWatermark.messageOffset
+    val responseTopics = authorizedRequestInfo.map { topic =>
+      val responsePartitions = topic.partitions.asScala.map { partition =>
+        val topicPartition = new TopicPartition(topic.name, partition.partitionIndex)
+        if (offsetRequest.duplicatePartitions.contains(topicPartition)) {
+          debug(s"OffsetRequest with correlation id $correlationId from client $clientId on partition $topicPartition " +
+              s"failed because the partition is duplicated in the request.")
+          buildErrorResponse(Errors.INVALID_REQUEST, partition)
+        } else {
+          try {
+            val fetchOnlyFromLeader = offsetRequest.replicaId != ListOffsetRequest.DEBUGGING_REPLICA_ID
+            val isClientRequest = offsetRequest.replicaId == ListOffsetRequest.CONSUMER_REPLICA_ID
+            val isolationLevelOpt = if (isClientRequest)
+              Some(offsetRequest.isolationLevel)
+            else
+              None
+
+            val foundOpt = replicaManager.fetchOffsetForTimestamp(topicPartition,
+              partition.timestamp,
+              isolationLevelOpt,
+              if (partition.currentLeaderEpoch == ListOffsetResponse.UNKNOWN_EPOCH) Optional.empty() else Optional.of(partition.currentLeaderEpoch),
+              fetchOnlyFromLeader)
+
+            val response = foundOpt match {
+              case Some(found) =>
+                val partitionResponse = new ListOffsetPartitionResponse()
+                  .setPartitionIndex(partition.partitionIndex)
+                  .setErrorCode(Errors.NONE.code)
+                  .setTimestamp(found.timestamp)
+                  .setOffset(found.offset)
+                if (found.leaderEpoch.isPresent && version >= 4)
+                  partitionResponse.setLeaderEpoch(found.leaderEpoch.get)
+                partitionResponse
+              case None =>
+                buildErrorResponse(Errors.NONE, partition)
             }
+            response
+          } catch {
+            // NOTE: These exceptions are special cases since these error messages are typically transient or the client
+            // would have received a clear exception and there is no value in logging the entire stack trace for the same
+            case e @ (_ : UnknownTopicOrPartitionException |
+                      _ : NotLeaderForPartitionException |
+                      _ : UnknownLeaderEpochException |
+                      _ : FencedLeaderEpochException |
+                      _ : KafkaStorageException |
+                      _ : UnsupportedForMessageFormatException) =>
+              e.printStackTrace()
+              debug(s"Offset request with correlation id $correlationId from client $clientId on " +
+                  s"partition $topicPartition failed due to ${e.getMessage}")
+              buildErrorResponse(Errors.forException(e), partition)
+
+            // Only V5 and newer ListOffset calls should get OFFSET_NOT_AVAILABLE
+            case e: OffsetNotAvailableException =>
+              if (request.header.apiVersion >= 5) {
+                buildErrorResponse(Errors.forException(e), partition)
+              } else {
+                buildErrorResponse(Errors.LEADER_NOT_AVAILABLE, partition)
+              }
 
-            if (timestamp == ListOffsetRequest.LATEST_TIMESTAMP)
-              TimestampOffset(RecordBatch.NO_TIMESTAMP, lastFetchableOffset)
-            else {
-              def allowed(timestampOffset: TimestampOffset): Boolean =
-                timestamp == ListOffsetRequest.EARLIEST_TIMESTAMP || timestampOffset.offset < lastFetchableOffset
-
-              fetchOffsetForTimestamp(topicPartition, timestamp)
-                .filter(allowed).getOrElse(TimestampOffset.Unknown)
-            }
-          } else {
-            fetchOffsetForTimestamp(topicPartition, timestamp)
-              .getOrElse(TimestampOffset.Unknown)
+            case e: Throwable =>
+              error("Error while responding to offset request", e)
+              buildErrorResponse(Errors.forException(e), partition)
           }
-
-          (topicPartition, new ListOffsetResponse.PartitionData(Errors.NONE, found.timestamp, found.offset))
-        } catch {
-          // NOTE: These exceptions are special cased since these error messages are typically transient or the client
-          // would have received a clear exception and there is no value in logging the entire stack trace for the same
-          case e @ (_ : UnknownTopicOrPartitionException |
-                    _ : NotLeaderForPartitionException |
-                    _ : KafkaStorageException |
-                    _ : UnsupportedForMessageFormatException) =>
-            debug(s"Offset request with correlation id $correlationId from client $clientId on " +
-                s"partition $topicPartition failed due to ${e.getMessage}")
-            (topicPartition, new ListOffsetResponse.PartitionData(Errors.forException(e),
-                                                                  ListOffsetResponse.UNKNOWN_TIMESTAMP,
-                                                                  ListOffsetResponse.UNKNOWN_OFFSET))
-          case e: Throwable =>
-            error("Error while responding to offset request", e)
-            (topicPartition, new ListOffsetResponse.PartitionData(Errors.forException(e),
-                                                                  ListOffsetResponse.UNKNOWN_TIMESTAMP,
-                                                                  ListOffsetResponse.UNKNOWN_OFFSET))
         }
       }
+      new ListOffsetTopicResponse().setName(topic.name).setPartitions(responsePartitions.asJava)
     }
-    responseMap ++ unauthorizedResponseStatus
-  }
-
-  def fetchOffsets(logManager: LogManager, topicPartition: TopicPartition, timestamp: Long, maxNumOffsets: Int): Seq[Long] = {
-    logManager.getLog(topicPartition) match {
-      case Some(log) =>
-        fetchOffsetsBefore(log, timestamp, maxNumOffsets)
-      case None =>
-        if (timestamp == ListOffsetRequest.LATEST_TIMESTAMP || timestamp == ListOffsetRequest.EARLIEST_TIMESTAMP)
-          Seq(0L)
-        else
-          Nil
-    }
-  }
-
-  private def fetchOffsetForTimestamp(topicPartition: TopicPartition, timestamp: Long): Option[TimestampOffset] = {
-    replicaManager.getLog(topicPartition) match {
-      case Some(log) =>
-        log.fetchOffsetsByTimestamp(timestamp)
-      case None =>
-        throw new UnknownTopicOrPartitionException(s"$topicPartition does not exist on the broker.")
-    }
-  }
-
-  private[server] def fetchOffsetsBefore(log: Log, timestamp: Long, maxNumOffsets: Int): Seq[Long] = {
-    // Cache to avoid race conditions. `toBuffer` is faster than most alternatives and provides
-    // constant time access while being safe to use with concurrent collections unlike `toArray`.
-    val segments = log.logSegments.toBuffer
-    val lastSegmentHasSize = segments.last.size > 0
-
-    val offsetTimeArray =
-      if (lastSegmentHasSize)
-        new Array[(Long, Long)](segments.length + 1)
-      else
-        new Array[(Long, Long)](segments.length)
-
-    for (i <- segments.indices)
-      offsetTimeArray(i) = (math.max(segments(i).baseOffset, log.logStartOffset), segments(i).lastModified)
-    if (lastSegmentHasSize)
-      offsetTimeArray(segments.length) = (log.logEndOffset, time.milliseconds)
-
-    var startIndex = -1
-    timestamp match {
-      case ListOffsetRequest.LATEST_TIMESTAMP =>
-        startIndex = offsetTimeArray.length - 1
-      case ListOffsetRequest.EARLIEST_TIMESTAMP =>
-        startIndex = 0
-      case _ =>
-        var isFound = false
-        debug("Offset time array = " + offsetTimeArray.foreach(o => "%d, %d".format(o._1, o._2)))
-        startIndex = offsetTimeArray.length - 1
-        while (startIndex >= 0 && !isFound) {
-          if (offsetTimeArray(startIndex)._2 <= timestamp)
-            isFound = true
-          else
-            startIndex -= 1
-        }
-    }
-
-    val retSize = maxNumOffsets.min(startIndex + 1)
-    val ret = new Array[Long](retSize)
-    for (j <- 0 until retSize) {
-      ret(j) = offsetTimeArray(startIndex)._1
-      startIndex -= 1
-    }
-    // ensure that the returned seq is in descending order of offsets
-    ret.toSeq.sortBy(-_)
+    (responseTopics ++ unauthorizedResponseStatus).toList
   }
 
   private def createTopic(topic: String,
                           numPartitions: Int,
                           replicationFactor: Int,
-                          properties: Properties = new Properties()): MetadataResponse.TopicMetadata = {
+                          properties: util.Properties = new util.Properties()): MetadataResponseTopic = {
     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())
+      metadataResponseTopic(Errors.LEADER_NOT_AVAILABLE, topic, isInternal(topic), 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())
+        metadataResponseTopic(Errors.LEADER_NOT_AVAILABLE, topic, isInternal(topic), 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())
+        metadataResponseTopic(Errors.forException(ex), topic, isInternal(topic), util.Collections.emptyList())
     }
   }
 
-  private def createInternalTopic(topic: String): MetadataResponse.TopicMetadata = {
+  private def metadataResponseTopic(error: Errors, topic: String, isInternal: Boolean,
+                                    partitionData: util.List[MetadataResponsePartition]): MetadataResponseTopic = {
+    new MetadataResponseTopic()
+      .setErrorCode(error.code)
+      .setName(topic)
+      .setIsInternal(isInternal)
+      .setPartitions(partitionData)
+  }
+
+  private def createInternalTopic(topic: String): MetadataResponseTopic = {
     if (topic == null)
       throw new IllegalArgumentException("topic must not be null")
 
@@ -857,7 +1211,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())
+          metadataResponseTopic(Errors.COORDINATOR_NOT_AVAILABLE, topic, true, util.Collections.emptyList())
         } else {
           createTopic(topic, config.offsetsTopicPartitions, config.offsetsTopicReplicationFactor.toInt,
             groupCoordinator.offsetsTopicConfigs)
@@ -868,7 +1222,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())
+          metadataResponseTopic(Errors.COORDINATOR_NOT_AVAILABLE, topic, true, util.Collections.emptyList())
         } else {
           createTopic(topic, config.transactionTopicPartitions, config.transactionTopicReplicationFactor.toInt,
             txnCoordinator.transactionTopicConfigs)
@@ -877,31 +1231,42 @@ class KafkaApis(val requestChannel: RequestChannel,
     }
   }
 
-  private def getOrCreateInternalTopic(topic: String, listenerName: ListenerName): MetadataResponse.TopicMetadata = {
+  private def getOrCreateInternalTopic(topic: String, listenerName: ListenerName): MetadataResponseData.MetadataResponseTopic = {
     val topicMetadata = metadataCache.getTopicMetadata(Set(topic), listenerName)
     topicMetadata.headOption.getOrElse(createInternalTopic(topic))
   }
 
-  private def getTopicMetadata(allowAutoTopicCreation: Boolean, topics: Set[String], listenerName: ListenerName,
-                               errorUnavailableEndpoints: Boolean): Seq[MetadataResponse.TopicMetadata] = {
-    val topicResponses = metadataCache.getTopicMetadata(topics, listenerName, errorUnavailableEndpoints)
+  private def getTopicMetadata(allowAutoTopicCreation: Boolean, isFetchAllMetadata: Boolean,
+                               topics: Set[String], listenerName: ListenerName,
+                               errorUnavailableEndpoints: Boolean,
+                               errorUnavailableListeners: Boolean): Seq[MetadataResponseTopic] = {
+    val topicResponses = metadataCache.getTopicMetadata(topics, listenerName,
+        errorUnavailableEndpoints, errorUnavailableListeners)
+
     if (topics.isEmpty || topicResponses.size == topics.size) {
       topicResponses
     } else {
-      val nonExistentTopics = topics -- topicResponses.map(_.topic).toSet
-      val responsesForNonExistentTopics = nonExistentTopics.map { topic =>
+      val nonExistentTopics = topics.diff(topicResponses.map(_.name).toSet)
+      val responsesForNonExistentTopics = nonExistentTopics.flatMap { topic =>
         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())
-          else
-            topicMetadata
+          Some(
+            if (topicMetadata.errorCode == Errors.COORDINATOR_NOT_AVAILABLE.code)
+              metadataResponseTopic(Errors.INVALID_REPLICATION_FACTOR, topic, true, util.Collections.emptyList())
+            else
+              topicMetadata
+          )
+        } else if (isFetchAllMetadata) {
+          // A metadata request for all topics should never result in topic auto creation, but a topic may be deleted
+          // in between the creation of the topics parameter and topicResponses, so make sure to return None for this case.
+          None
         } else if (allowAutoTopicCreation && config.autoCreateTopicsEnable) {
-          createTopic(topic, config.numPartitions, config.defaultReplicationFactor)
+          Some(createTopic(topic, config.numPartitions, config.defaultReplicationFactor))
         } else {
-          new MetadataResponse.TopicMetadata(Errors.UNKNOWN_TOPIC_OR_PARTITION, topic, false, java.util.Collections.emptyList())
+          Some(metadataResponseTopic(Errors.UNKNOWN_TOPIC_OR_PARTITION, topic, false, util.Collections.emptyList()))
         }
       }
+
       topicResponses ++ responsesForNonExistentTopics
     }
   }
@@ -909,61 +1274,82 @@ 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
 
-    val topics =
-      // Handle old metadata request logic. Version 0 has no way to specify "no topics".
-      if (requestVersion == 0) {
-        if (metadataRequest.topics() == null || metadataRequest.topics.isEmpty)
-          metadataCache.getAllTopics()
-        else
-          metadataRequest.topics.asScala.toSet
-      } else {
-        if (metadataRequest.isAllTopics)
-          metadataCache.getAllTopics()
-        else
-          metadataRequest.topics.asScala.toSet
-      }
-
-    var (authorizedTopics, unauthorizedForDescribeTopics) =
-      topics.partition(topic => authorize(request.session, Describe, new Resource(Topic, topic)))
+    val topics = if (metadataRequest.isAllTopics)
+      metadataCache.getAllTopics()
+    else
+      metadataRequest.topics.asScala.toSet
 
+    val authorizedForDescribeTopics = filterByAuthorized(request.context, DESCRIBE, TOPIC,
+      topics, logIfDenied = !metadataRequest.isAllTopics)(identity)
+    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)) {
-          authorizedTopics --= nonExistingTopics
-          unauthorizedForCreateTopics ++= nonExistingTopics
+        if (!authorize(request.context, CREATE, CLUSTER, CLUSTER_NAME, logIfDenied = false)) {
+          val authorizedForCreateTopics = filterByAuthorized(request.context, CREATE, TOPIC,
+            nonExistingTopics)(identity)
+          unauthorizedForCreateTopics = nonExistingTopics.diff(authorizedForCreateTopics)
+          authorizedTopics = authorizedTopics.diff(unauthorizedForCreateTopics)
         }
       }
     }
 
     val unauthorizedForCreateTopicMetadata = unauthorizedForCreateTopics.map(topic =>
-      new MetadataResponse.TopicMetadata(Errors.TOPIC_AUTHORIZATION_FAILED, topic, isInternal(topic),
-        java.util.Collections.emptyList()))
+      metadataResponseTopic(Errors.TOPIC_AUTHORIZATION_FAILED, topic, isInternal(topic), 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 =
       // In case of all topics, don't include topics unauthorized for Describe
       if ((requestVersion == 0 && (metadataRequest.topics == null || metadataRequest.topics.isEmpty)) || metadataRequest.isAllTopics)
-        Set.empty[MetadataResponse.TopicMetadata]
+        Set.empty[MetadataResponseTopic]
       else
         unauthorizedForDescribeTopics.map(topic =>
-          new MetadataResponse.TopicMetadata(Errors.TOPIC_AUTHORIZATION_FAILED, topic, false, java.util.Collections.emptyList()))
+          metadataResponseTopic(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
     val errorUnavailableEndpoints = requestVersion == 0
+    // In versions 5 and below, we returned LEADER_NOT_AVAILABLE if a matching listener was not found on the leader.
+    // From version 6 onwards, we return LISTENER_NOT_FOUND to enable diagnosis of configuration errors.
+    val errorUnavailableListeners = requestVersion >= 6
     val topicMetadata =
       if (authorizedTopics.isEmpty)
-        Seq.empty[MetadataResponse.TopicMetadata]
-      else
-        getTopicMetadata(metadataRequest.allowAutoTopicCreation, authorizedTopics, request.context.listenerName,
-          errorUnavailableEndpoints)
+        Seq.empty[MetadataResponseTopic]
+      else {
+        getTopicMetadata(
+          metadataRequest.allowAutoTopicCreation,
+          metadataRequest.isAllTopics,
+          authorizedTopics,
+          request.context.listenerName,
+          errorUnavailableEndpoints,
+          errorUnavailableListeners
+        )
+      }
+
+    var clusterAuthorizedOperations = Int.MinValue
+    if (request.header.apiVersion >= 8) {
+      // get cluster authorized operations
+      if (metadataRequest.data.includeClusterAuthorizedOperations) {
+        if (authorize(request.context, DESCRIBE, CLUSTER, CLUSTER_NAME))
+          clusterAuthorizedOperations = authorizedOperations(request, Resource.CLUSTER)
+        else
+          clusterAuthorizedOperations = 0
+      }
+
+      // get topic authorized operations
+      if (metadataRequest.data.includeTopicAuthorizedOperations) {
+        topicMetadata.foreach { topicData =>
+          topicData.setTopicAuthorizedOperations(authorizedOperations(request, new Resource(ResourceType.TOPIC, topicData.name)))
+        }
+      }
+    }
 
     val completeTopicMetadata = topicMetadata ++ unauthorizedForCreateTopicMetadata ++ unauthorizedForDescribeTopicMetadata
 
@@ -973,54 +1359,56 @@ class KafkaApis(val requestChannel: RequestChannel,
       brokers.mkString(","), request.header.correlationId, request.header.clientId))
 
     sendResponseMaybeThrottle(request, requestThrottleMs =>
-      new MetadataResponse(
-        requestThrottleMs,
-        brokers.map(_.getNode(request.context.listenerName)).asJava,
-        clusterId,
-        metadataCache.getControllerId.getOrElse(MetadataResponse.NO_CONTROLLER_ID),
-        completeTopicMetadata.asJava
+       MetadataResponse.prepareResponse(
+         requestVersion,
+         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, new Resource(Topic, partition.topic))
+    def partitionByAuthorized(seq: Seq[TopicPartition]): (Seq[TopicPartition], Seq[TopicPartition]) =
+      partitionSeqByAuthorized(request.context, DESCRIBE, TOPIC, seq)(_.topic)
 
     def createResponse(requestThrottleMs: Int): AbstractResponse = {
       val offsetFetchResponse =
         // reject the request if not authorized to the group
-        if (!authorize(request.session, Describe, new Resource(Group, offsetFetchRequest.groupId)))
+        if (!authorize(request.context, 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) = partitionByAuthorized(
+              offsetFetchRequest.partitions.asScala)
 
             // version 0 reads offsets from ZK
             val authorizedPartitionData = authorizedPartitions.map { topicPartition =>
               try {
-                if (!metadataCache.contains(topicPartition.topic))
+                if (!metadataCache.contains(topicPartition))
                   (topicPartition, OffsetFetchResponse.UNKNOWN_PARTITION)
                 else {
                   val payloadOpt = zkClient.getConsumerOffset(offsetFetchRequest.groupId, topicPartition)
                   payloadOpt match {
                     case Some(payload) =>
-                      (topicPartition, new OffsetFetchResponse.PartitionData(
-                          payload.toLong, OffsetFetchResponse.NO_METADATA, Errors.NONE))
+                      (topicPartition, new OffsetFetchResponse.PartitionData(payload.toLong,
+                        Optional.empty(), OffsetFetchResponse.NO_METADATA, Errors.NONE))
                     case None =>
                       (topicPartition, OffsetFetchResponse.UNKNOWN_PARTITION)
                   }
                 }
               } catch {
                 case e: Throwable =>
-                  (topicPartition, new OffsetFetchResponse.PartitionData(
-                      OffsetFetchResponse.INVALID_OFFSET, OffsetFetchResponse.NO_METADATA, Errors.forException(e)))
+                  (topicPartition, new OffsetFetchResponse.PartitionData(OffsetFetchResponse.INVALID_OFFSET,
+                    Optional.empty(), OffsetFetchResponse.NO_METADATA, Errors.forException(e)))
               }
             }.toMap
 
@@ -1029,19 +1417,20 @@ class KafkaApis(val requestChannel: RequestChannel,
           } else {
             // versions 1 and above read offsets from Kafka
             if (offsetFetchRequest.isAllPartitions) {
-              val (error, allPartitionData) = groupCoordinator.handleFetchOffsets(offsetFetchRequest.groupId)
+              val (error, allPartitionData) = groupCoordinator.handleFetchOffsets(offsetFetchRequest.groupId, offsetFetchRequest.requireStable)
               if (error != Errors.NONE)
                 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) }
+                val (authorizedPartitionData, _) = partitionMapByAuthorized(request.context,
+                  DESCRIBE, TOPIC, allPartitionData)(_.topic)
                 new OffsetFetchResponse(requestThrottleMs, Errors.NONE, authorizedPartitionData.asJava)
               }
             } else {
-              val (authorizedPartitions, unauthorizedPartitions) = offsetFetchRequest.partitions.asScala
-                .partition(authorizeTopicDescribe)
+              val (authorizedPartitions, unauthorizedPartitions) = partitionByAuthorized(
+                offsetFetchRequest.partitions.asScala)
               val (error, authorizedPartitionData) = groupCoordinator.handleFetchOffsets(offsetFetchRequest.groupId,
-                Some(authorizedPartitions))
+                offsetFetchRequest.requireStable, Some(authorizedPartitions))
               if (error != Errors.NONE)
                 offsetFetchRequest.getErrorResponse(requestThrottleMs, error)
               else {
@@ -1059,25 +1448,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, new Resource(Group, findCoordinatorRequest.coordinatorKey)))
+    if (findCoordinatorRequest.data.keyType == CoordinatorType.GROUP.id &&
+        !authorize(request.context, DESCRIBE, GROUP, findCoordinatorRequest.data.key))
       sendErrorResponseMaybeThrottle(request, Errors.GROUP_AUTHORIZATION_FAILED.exception)
-    else if (findCoordinatorRequest.coordinatorType == FindCoordinatorRequest.CoordinatorType.TRANSACTION &&
-        !authorize(request.session, Describe, new Resource(TransactionalId, findCoordinatorRequest.coordinatorKey)))
+    else if (findCoordinatorRequest.data.keyType == CoordinatorType.TRANSACTION.id &&
+        !authorize(request.context, 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)
 
@@ -1086,18 +1475,31 @@ class KafkaApis(val requestChannel: RequestChannel,
       }
 
       def createResponse(requestThrottleMs: Int): AbstractResponse = {
-        val responseBody = if (topicMetadata.error != Errors.NONE) {
-          new FindCoordinatorResponse(requestThrottleMs, Errors.COORDINATOR_NOT_AVAILABLE, Node.noNode)
+        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.errorCode != Errors.NONE.code) {
+          createFindCoordinatorResponse(Errors.COORDINATOR_NOT_AVAILABLE, Node.noNode)
         } else {
-          val coordinatorEndpoint = topicMetadata.partitionMetadata.asScala
-            .find(_.partition == partition)
-            .map(_.leader)
+          val coordinatorEndpoint = topicMetadata.partitions.asScala
+            .find(_.partitionIndex == partition)
+            .filter(_.leaderId != MetadataResponse.NO_LEADER_ID)
+            .flatMap(metadata => metadataCache.getAliveBroker(metadata.leaderId))
+            .flatMap(_.getNode(request.context.listenerName))
+            .filterNot(_.isEmpty)
 
           coordinatorEndpoint match {
-            case Some(endpoint) if !endpoint.isEmpty =>
-              new FindCoordinatorResponse(requestThrottleMs, Errors.NONE, endpoint)
+            case Some(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."
@@ -1108,48 +1510,110 @@ 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 = {
+        describeGroupsResponseData.setThrottleTimeMs(requestThrottleMs)
+        new DescribeGroupsResponse(describeGroupsResponseData)
+      }
+      sendResponseMaybeThrottle(request, createResponse)
+    }
+
     val describeRequest = request.body[DescribeGroupsRequest]
+    val describeGroupsResponseData = new DescribeGroupsResponseData()
 
-    val groups = describeRequest.groupIds.asScala.map { groupId =>
-      if (!authorize(request.session, Describe, new Resource(Group, groupId))) {
-        groupId -> DescribeGroupsResponse.GroupMetadata.forError(Errors.GROUP_AUTHORIZATION_FAILED)
+    describeRequest.data.groups.forEach { groupId =>
+      if (!authorize(request.context, 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 =>
-          val metadata = ByteBuffer.wrap(member.metadata)
-          val assignment = ByteBuffer.wrap(member.assignment)
-          new DescribeGroupsResponse.GroupMember(member.memberId, member.clientId, member.clientHost, metadata, assignment)
+          new DescribeGroupsResponseData.DescribedGroupMember()
+            .setMemberId(member.memberId)
+            .setGroupInstanceId(member.groupInstanceId.orNull)
+            .setClientId(member.clientId)
+            .setClientHost(member.clientHost)
+            .setMemberAssignment(member.assignment)
+            .setMemberMetadata(member.metadata)
         }
-        groupId -> new DescribeGroupsResponse.GroupMetadata(error, summary.state, summary.protocolType,
-          summary.protocol, members.asJava)
+
+        val describedGroup = new DescribeGroupsResponseData.DescribedGroup()
+          .setErrorCode(error.code)
+          .setGroupId(groupId)
+          .setGroupState(summary.state)
+          .setProtocolType(summary.protocolType)
+          .setProtocolData(summary.protocol)
+          .setMembers(members.asJava)
+
+        if (request.header.apiVersion >= 3) {
+          if (error == Errors.NONE && describeRequest.data.includeAuthorizedOperations) {
+            describedGroup.setAuthorizedOperations(authorizedOperations(request, new Resource(ResourceType.GROUP, groupId)))
+          }
+        }
+
+        describeGroupsResponseData.groups.add(describedGroup)
       }
-    }.toMap
+    }
 
-    sendResponseMaybeThrottle(request, requestThrottleMs => new DescribeGroupsResponse(requestThrottleMs, groups.asJava))
+    sendResponseCallback(describeGroupsResponseData)
   }
 
-  def handleListGroupsRequest(request: RequestChannel.Request) {
-    if (!authorize(request.session, Describe, Resource.ClusterResource)) {
+  def handleListGroupsRequest(request: RequestChannel.Request): Unit = {
+    val listGroupsRequest = request.body[ListGroupsRequest]
+    val states = if (listGroupsRequest.data.statesFilter == null)
+      // Handle a null array the same as empty
+      immutable.Set[String]()
+    else
+      listGroupsRequest.data.statesFilter.asScala.toSet
+
+    def createResponse(throttleMs: Int, groups: List[GroupOverview], error: Errors): AbstractResponse = {
+       new ListGroupsResponse(new ListGroupsResponseData()
+            .setErrorCode(error.code)
+            .setGroups(groups.map { group =>
+                val listedGroup = new ListGroupsResponseData.ListedGroup()
+                  .setGroupId(group.groupId)
+                  .setProtocolType(group.protocolType)
+                  .setGroupState(group.state.toString)
+                listedGroup
+            }.asJava)
+            .setThrottleTimeMs(throttleMs)
+        )
+    }
+    val (error, groups) = groupCoordinator.handleListGroups(states)
+    if (authorize(request.context, DESCRIBE, CLUSTER, CLUSTER_NAME))
+      // With describe cluster access all groups are returned. We keep this alternative for backward compatibility.
       sendResponseMaybeThrottle(request, requestThrottleMs =>
-        request.body[ListGroupsRequest].getErrorResponse(requestThrottleMs, Errors.CLUSTER_AUTHORIZATION_FAILED.exception))
-    } else {
-      val (error, groups) = groupCoordinator.handleListGroups()
-      val allGroups = groups.map { group => new ListGroupsResponse.Group(group.groupId, group.protocolType) }
+        createResponse(requestThrottleMs, groups, error))
+    else {
+      val filteredGroups = groups.filter(group => authorize(request.context, DESCRIBE, GROUP, group.groupId))
       sendResponseMaybeThrottle(request, requestThrottleMs =>
-        new ListGroupsResponse(requestThrottleMs, error, allGroups.asJava))
+        createResponse(requestThrottleMs, filteredGroups, error))
     }
   }
 
-  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 protocolName = if (request.context.apiVersion() >= 7)
+          joinResult.protocolName.orNull
+        else
+          joinResult.protocolName.getOrElse(GroupCoordinator.NoProtocol)
+
+        val responseBody = new JoinGroupResponse(
+          new JoinGroupResponseData()
+            .setThrottleTimeMs(requestThrottleMs)
+            .setErrorCode(joinResult.error.code)
+            .setGenerationId(joinResult.generationId)
+            .setProtocolType(joinResult.protocolType.orNull)
+            .setProtocolName(protocolName)
+            .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))
@@ -1158,121 +1622,202 @@ class KafkaApis(val requestChannel: RequestChannel,
       sendResponseMaybeThrottle(request, createResponse)
     }
 
-    if (!authorize(request.session, Read, new Resource(Group, joinGroupRequest.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())
-      )
+    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(JoinGroupRequest.UNKNOWN_MEMBER_ID, Errors.UNSUPPORTED_VERSION))
+    } else if (!authorize(request.context, READ, GROUP, joinGroupRequest.data.groupId)) {
+      sendResponseCallback(JoinGroupResult(JoinGroupRequest.UNKNOWN_MEMBER_ID, Errors.GROUP_AUTHORIZATION_FAILED))
     } else {
-      // let the coordinator to handle join-group
-      val protocols = joinGroupRequest.groupProtocols().asScala.map(protocol =>
-        (protocol.name, Utils.toArray(protocol.metadata))).toList
+      val groupInstanceId = Option(joinGroupRequest.data.groupInstanceId)
+
+      // Only return MEMBER_ID_REQUIRED error if joinGroupRequest version is >= 4
+      // and groupInstanceId is configured to unknown.
+      val requireKnownMemberId = joinGroupRequest.version >= 4 && groupInstanceId.isEmpty
+
+      // let the coordinator handle join-group
+      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, new Resource(Group, syncGroupRequest.groupId()))) {
-      sendResponseCallback(Array[Byte](), Errors.GROUP_AUTHORIZATION_FAILED)
+        new SyncGroupResponse(
+          new SyncGroupResponseData()
+            .setErrorCode(syncGroupResult.error.code)
+            .setProtocolType(syncGroupResult.protocolType.orNull)
+            .setProtocolName(syncGroupResult.protocolName.orNull)
+            .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(Errors.UNSUPPORTED_VERSION))
+    } else if (!syncGroupRequest.areMandatoryProtocolTypeAndNamePresent()) {
+      // Starting from version 5, ProtocolType and ProtocolName fields are mandatory.
+      sendResponseCallback(SyncGroupResult(Errors.INCONSISTENT_GROUP_PROTOCOL))
+    } else if (!authorize(request.context, READ, GROUP, syncGroupRequest.data.groupId)) {
+      sendResponseCallback(SyncGroupResult(Errors.GROUP_AUTHORIZATION_FAILED))
     } else {
+      val assignmentMap = immutable.Map.newBuilder[String, Array[Byte]]
+      syncGroupRequest.data.assignments.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.protocolType),
+        Option(syncGroupRequest.data.protocolName),
+        Option(syncGroupRequest.data.groupInstanceId),
+        assignmentMap.result(),
         sendResponseCallback
       )
     }
   }
 
-  def handleHeartbeatRequest(request: RequestChannel.Request) {
-    val heartbeatRequest = request.body[HeartbeatRequest]
+  def handleDeleteGroupsRequest(request: RequestChannel.Request): Unit = {
+    val deleteGroupsRequest = request.body[DeleteGroupsRequest]
+    val groups = deleteGroupsRequest.data.groupsNames.asScala.distinct
 
-    // the callback for sending a heartbeat response
-    def sendResponseCallback(error: Errors) {
-      def createResponse(requestThrottleMs: Int): AbstractResponse = {
-        val response = new HeartbeatResponse(requestThrottleMs, error)
-        trace("Sending heartbeat response %s for correlation id %d to client %s."
-          .format(response, request.header.correlationId, request.header.clientId))
-        response
-      }
-      sendResponseMaybeThrottle(request, createResponse)
-    }
+    val (authorizedGroups, unauthorizedGroups) = partitionSeqByAuthorized(request.context, DELETE, GROUP,
+      groups)(identity)
 
-    if (!authorize(request.session, Read, new Resource(Group, heartbeatRequest.groupId))) {
+    val groupDeletionResult = groupCoordinator.handleDeleteGroups(authorizedGroups.toSet) ++
+      unauthorizedGroups.map(_ -> Errors.GROUP_AUTHORIZATION_FAILED)
+
+    sendResponseMaybeThrottle(request, requestThrottleMs => {
+      val deletionCollections = new DeletableGroupResultCollection()
+      groupDeletionResult.forKeyValue { (groupId, error) =>
+        deletionCollections.add(new DeletableGroupResult()
+          .setGroupId(groupId)
+          .setErrorCode(error.code)
+        )
+      }
+
+      new DeleteGroupsResponse(new DeleteGroupsResponseData()
+        .setResults(deletionCollections)
+        .setThrottleTimeMs(requestThrottleMs)
+      )
+    })
+  }
+
+  def handleHeartbeatRequest(request: RequestChannel.Request): Unit = {
+    val heartbeatRequest = request.body[HeartbeatRequest]
+
+    // the callback for sending a heartbeat response
+    def sendResponseCallback(error: Errors): Unit = {
+      def createResponse(requestThrottleMs: Int): AbstractResponse = {
+        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
+      }
+      sendResponseMaybeThrottle(request, createResponse)
+    }
+
+    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.context, 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(requestThrottleMs, error)
-        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, new Resource(Group, leaveGroupRequest.groupId))) {
-      sendResponseMaybeThrottle(request, requestThrottleMs =>
-        new LeaveGroupResponse(requestThrottleMs, Errors.GROUP_AUTHORIZATION_FAILED))
+    if (!authorize(request.context, READ, GROUP, leaveGroupRequest.data.groupId)) {
+      sendResponseMaybeThrottle(request, requestThrottleMs => {
+        new LeaveGroupResponse(new LeaveGroupResponseData()
+          .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.groupId,
-        leaveGroupRequest.memberId,
+        leaveGroupRequest.data.groupId,
+        members,
         sendResponseCallback)
     }
   }
 
-  def handleSaslHandshakeRequest(request: RequestChannel.Request) {
-    sendResponseMaybeThrottle(request, _ => new SaslHandshakeResponse(Errors.ILLEGAL_SASL_STATE, config.saslEnabledMechanisms))
+  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) {
-    sendResponseMaybeThrottle(request, _ => new SaslAuthenticateResponse(Errors.ILLEGAL_SASL_STATE,
-        "SaslAuthenticate request received after successful authentication"))
+  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.
@@ -1280,182 +1825,287 @@ 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.messageFormatVersion)
+      else if (!apiVersionRequest.isValid)
+        apiVersionRequest.getErrorResponse(requestThrottleMs, Errors.INVALID_REQUEST.exception)
+      else {
+        val supportedFeatures = brokerFeatures.supportedFeatures
+        val finalizedFeaturesOpt = finalizedFeatureCache.get
+        finalizedFeaturesOpt match {
+          case Some(finalizedFeatures) => ApiVersion.apiVersionsResponse(
+            requestThrottleMs,
+            config.interBrokerProtocolVersion.recordVersion.value,
+            supportedFeatures,
+            finalizedFeatures.features,
+            finalizedFeatures.epoch)
+          case None => ApiVersion.apiVersionsResponse(
+            requestThrottleMs,
+            config.interBrokerProtocolVersion.recordVersion.value,
+            supportedFeatures)
+        }
+      }
     }
     sendResponseMaybeThrottle(request, createResponseCallback)
   }
 
-  def handleCreateTopicsRequest(request: RequestChannel.Request) {
-    val createTopicsRequest = request.body[CreateTopicsRequest]
+  def handleCreateTopicsRequest(request: RequestChannel.Request): Unit = {
+    val controllerMutationQuota = quotas.controllerMutation.newQuotaFor(request, strictSinceVersion = 6)
 
-    def sendResponseCallback(results: Map[String, ApiError]): Unit = {
+    def sendResponseCallback(results: CreatableTopicResultCollection): Unit = {
       def createResponse(requestThrottleMs: Int): AbstractResponse = {
-        val responseBody = new CreateTopicsResponse(requestThrottleMs, results.asJava)
-        trace(s"Sending create topics response $responseBody for correlation id ${request.header.correlationId} to client ${request.header.clientId}.")
+        val responseData = new CreateTopicsResponseData()
+          .setThrottleTimeMs(requestThrottleMs)
+          .setTopics(results)
+        val responseBody = new CreateTopicsResponse(responseData)
+        trace(s"Sending create topics response $responseData for correlation id " +
+          s"${request.header.correlationId} to client ${request.header.clientId}.")
         responseBody
       }
-      sendResponseMaybeThrottle(request, createResponse)
+      sendResponseMaybeThrottleWithControllerQuota(controllerMutationQuota, request, createResponse)
     }
 
+    val createTopicsRequest = request.body[CreateTopicsRequest]
+    val results = new CreatableTopicResultCollection(createTopicsRequest.data.topics.size)
     if (!controller.isActive) {
-      val results = createTopicsRequest.topics.asScala.map { case (topic, _) =>
-        (topic, new ApiError(Errors.NOT_CONTROLLER, null))
-      }
-      sendResponseCallback(results)
-    } else if (!authorize(request.session, Create, Resource.ClusterResource)) {
-      val results = createTopicsRequest.topics.asScala.map { case (topic, _) =>
-        (topic, new ApiError(Errors.CLUSTER_AUTHORIZATION_FAILED, null))
+      createTopicsRequest.data.topics.forEach { topic =>
+        results.add(new CreatableTopicResult().setName(topic.name)
+          .setErrorCode(Errors.NOT_CONTROLLER.code))
       }
       sendResponseCallback(results)
     } else {
-      val (validTopics, duplicateTopics) = createTopicsRequest.topics.asScala.partition { case (topic, _) =>
-        !createTopicsRequest.duplicateTopics.contains(topic)
+      createTopicsRequest.data.topics.forEach { topic =>
+        results.add(new CreatableTopicResult().setName(topic.name))
       }
-
-      // Special handling to add duplicate topics to the response
-      def sendResponseWithDuplicatesCallback(results: Map[String, ApiError]): Unit = {
-
-        val duplicatedTopicsResults =
-          if (duplicateTopics.nonEmpty) {
-            val errorMessage = s"Create topics request from client `${request.header.clientId}` contains multiple entries " +
-              s"for the following topics: ${duplicateTopics.keySet.mkString(",")}"
-            // We can send the error message in the response for version 1, so we don't have to log it any more
-            if (request.header.apiVersion == 0)
-              warn(errorMessage)
-            duplicateTopics.keySet.map((_, new ApiError(Errors.INVALID_REQUEST, errorMessage))).toMap
-          } else Map.empty
-
-        val completeResults = results ++ duplicatedTopicsResults
-        sendResponseCallback(completeResults)
+      val hasClusterAuthorization = authorize(request.context, CREATE, CLUSTER, CLUSTER_NAME,
+        logIfDenied = false)
+      val topics = createTopicsRequest.data.topics.asScala.map(_.name)
+      val authorizedTopics =
+        if (hasClusterAuthorization) topics.toSet
+        else filterByAuthorized(request.context, CREATE, TOPIC, topics)(identity)
+      val authorizedForDescribeConfigs = filterByAuthorized(request.context, DESCRIBE_CONFIGS, TOPIC,
+        topics, logIfDenied = false)(identity).map(name => name -> results.find(name)).toMap
+
+      results.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 (!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.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) =>
+          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(Errors.NONE.code)
+          }
+        }
+        sendResponseCallback(results)
       }
-
       adminManager.createTopics(
-        createTopicsRequest.timeout,
-        createTopicsRequest.validateOnly,
-        validTopics,
-        sendResponseWithDuplicatesCallback
-      )
+        createTopicsRequest.data.timeoutMs,
+        createTopicsRequest.data.validateOnly,
+        toCreate,
+        authorizedForDescribeConfigs,
+        controllerMutationQuota,
+        handleCreateTopicsResults)
     }
   }
 
   def handleCreatePartitionsRequest(request: RequestChannel.Request): Unit = {
     val createPartitionsRequest = request.body[CreatePartitionsRequest]
+    val controllerMutationQuota = quotas.controllerMutation.newQuotaFor(request, strictSinceVersion = 3)
 
     def sendResponseCallback(results: Map[String, ApiError]): Unit = {
       def createResponse(requestThrottleMs: Int): AbstractResponse = {
-        val responseBody = new CreatePartitionsResponse(requestThrottleMs, results.asJava)
+        val createPartitionsResults = results.map {
+          case (topic, error) => new CreatePartitionsTopicResult()
+            .setName(topic)
+            .setErrorCode(error.error.code)
+            .setErrorMessage(error.message)
+        }.toSeq
+        val responseBody = new CreatePartitionsResponse(new CreatePartitionsResponseData()
+          .setThrottleTimeMs(requestThrottleMs)
+          .setResults(createPartitionsResults.asJava))
         trace(s"Sending create partitions response $responseBody for correlation id ${request.header.correlationId} to " +
           s"client ${request.header.clientId}.")
         responseBody
       }
-      sendResponseMaybeThrottle(request, createResponse)
+      sendResponseMaybeThrottleWithControllerQuota(controllerMutationQuota, request, createResponse)
     }
 
     if (!controller.isActive) {
-      val result = createPartitionsRequest.newPartitions.asScala.map { case (topic, _) =>
-        (topic, new ApiError(Errors.NOT_CONTROLLER, null))
-      }
+      val result = createPartitionsRequest.data.topics.asScala.map { topic =>
+        (topic.name, new ApiError(Errors.NOT_CONTROLLER, null))
+      }.toMap
       sendResponseCallback(result)
     } else {
       // Special handling to add duplicate topics to the response
-      val dupes = createPartitionsRequest.duplicates.asScala
-      val notDuped = createPartitionsRequest.newPartitions.asScala -- dupes
-      val (authorized, unauthorized) = notDuped.partition { case (topic, _) =>
-        authorize(request.session, Alter, new Resource(Topic, topic))
-      }
-
-      val (queuedForDeletion, valid) = authorized.partition { case (topic, _) =>
-        controller.topicDeletionManager.isTopicQueuedUpForDeletion(topic)
+      val topics = createPartitionsRequest.data.topics.asScala.toSeq
+      val dupes = topics.groupBy(_.name)
+        .filter { _._2.size > 1 }
+        .keySet
+      val notDuped = topics.filterNot(topic => dupes.contains(topic.name))
+      val (authorized, unauthorized) = partitionSeqByAuthorized(request.context, ALTER, TOPIC,
+        notDuped)(_.name)
+
+      val (queuedForDeletion, valid) = authorized.partition { topic =>
+        controller.topicDeletionManager.isTopicQueuedUpForDeletion(topic.name)
       }
 
       val errors = dupes.map(_ -> new ApiError(Errors.INVALID_REQUEST, "Duplicate topic in request.")) ++
-        unauthorized.keySet.map(_ -> new ApiError(Errors.TOPIC_AUTHORIZATION_FAILED, "The topic authorization is failed.")) ++
-        queuedForDeletion.keySet.map(_ -> new ApiError(Errors.INVALID_TOPIC_EXCEPTION, "The topic is queued for deletion."))
+        unauthorized.map(_.name -> new ApiError(Errors.TOPIC_AUTHORIZATION_FAILED, "The topic authorization is failed.")) ++
+        queuedForDeletion.map(_.name -> new ApiError(Errors.INVALID_TOPIC_EXCEPTION, "The topic is queued for deletion."))
 
-      adminManager.createPartitions(createPartitionsRequest.timeout, valid, createPartitionsRequest.validateOnly,
-        request.context.listenerName, result => sendResponseCallback(result ++ errors))
+      adminManager.createPartitions(
+        createPartitionsRequest.data.timeoutMs,
+        valid,
+        createPartitionsRequest.data.validateOnly,
+        controllerMutationQuota,
+        result => sendResponseCallback(result ++ errors))
     }
   }
 
-  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]()
+  def handleDeleteTopicsRequest(request: RequestChannel.Request): Unit = {
+    val controllerMutationQuota = quotas.controllerMutation.newQuotaFor(request, strictSinceVersion = 5)
 
-    for (topic <- deleteTopicRequest.topics.asScala) {
-      if (!authorize(request.session, Delete, new Resource(Topic, topic)))
-        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 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)
+      sendResponseMaybeThrottleWithControllerQuota(controllerMutationQuota, 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.forEach { 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
+      deleteTopicRequest.data.topicNames.forEach { topic =>
+        results.add(new DeletableTopicResult()
+          .setName(topic)
+          .setErrorCode(error.code))
+      }
       sendResponseCallback(results)
     } else {
+      deleteTopicRequest.data.topicNames.forEach { topic =>
+        results.add(new DeletableTopicResult()
+          .setName(topic))
+      }
+      val authorizedTopics = filterByAuthorized(request.context, DELETE, TOPIC,
+        results.asScala)(_.name)
+      results.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,
+          toDelete,
+          controllerMutationQuota,
+          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 unauthorizedTopicResponses = mutable.Map[TopicPartition, DeleteRecordsPartitionResult]()
+    val nonExistingTopicResponses = mutable.Map[TopicPartition, DeleteRecordsPartitionResult]()
     val authorizedForDeleteTopicOffsets = mutable.Map[TopicPartition, Long]()
 
-    for ((topicPartition, offset) <- deleteRecordsRequest.partitionOffsets.asScala) {
-      if (!authorize(request.session, Delete, new Resource(Topic, topicPartition.topic)))
-        unauthorizedTopicResponses += topicPartition -> new DeleteRecordsResponse.PartitionResponse(
-          DeleteRecordsResponse.INVALID_LOW_WATERMARK, Errors.TOPIC_AUTHORIZATION_FAILED)
-      else if (!metadataCache.contains(topicPartition.topic))
-        nonExistingTopicResponses += topicPartition -> new DeleteRecordsResponse.PartitionResponse(
-          DeleteRecordsResponse.INVALID_LOW_WATERMARK, Errors.UNKNOWN_TOPIC_OR_PARTITION)
+    val topics = deleteRecordsRequest.data.topics.asScala
+    val authorizedTopics = filterByAuthorized(request.context, DELETE, TOPIC, topics)(_.name)
+    val deleteTopicPartitions = topics.flatMap { deleteTopic =>
+      deleteTopic.partitions.asScala.map { deletePartition =>
+        new TopicPartition(deleteTopic.name, deletePartition.partitionIndex) -> deletePartition.offset
+      }
+    }
+    for ((topicPartition, offset) <- deleteTopicPartitions) {
+      if (!authorizedTopics.contains(topicPartition.topic))
+        unauthorizedTopicResponses += topicPartition -> new DeleteRecordsPartitionResult()
+          .setLowWatermark(DeleteRecordsResponse.INVALID_LOW_WATERMARK)
+          .setErrorCode(Errors.TOPIC_AUTHORIZATION_FAILED.code)
+      else if (!metadataCache.contains(topicPartition))
+        nonExistingTopicResponses += topicPartition -> new DeleteRecordsPartitionResult()
+          .setLowWatermark(DeleteRecordsResponse.INVALID_LOW_WATERMARK)
+          .setErrorCode(Errors.UNKNOWN_TOPIC_OR_PARTITION.code)
       else
         authorizedForDeleteTopicOffsets += (topicPartition -> offset)
     }
 
     // the callback for sending a DeleteRecordsResponse
-    def sendResponseCallback(authorizedTopicResponses: Map[TopicPartition, DeleteRecordsResponse.PartitionResponse]) {
+    def sendResponseCallback(authorizedTopicResponses: Map[TopicPartition, DeleteRecordsPartitionResult]): Unit = {
       val mergedResponseStatus = authorizedTopicResponses ++ unauthorizedTopicResponses ++ nonExistingTopicResponses
-      mergedResponseStatus.foreach { case (topicPartition, status) =>
-        if (status.error != Errors.NONE) {
+      mergedResponseStatus.forKeyValue { (topicPartition, status) =>
+        if (status.errorCode != Errors.NONE.code) {
           debug("DeleteRecordsRequest with correlation id %d from client %s on partition %s failed due to %s".format(
             request.header.correlationId,
             request.header.clientId,
             topicPartition,
-            status.error.exceptionName))
+            Errors.forCode(status.errorCode).exceptionName))
         }
       }
 
       sendResponseMaybeThrottle(request, requestThrottleMs =>
-        new DeleteRecordsResponse(requestThrottleMs, mergedResponseStatus.asJava))
+        new DeleteRecordsResponse(new DeleteRecordsResponseData()
+          .setThrottleTimeMs(requestThrottleMs)
+          .setTopics(new DeleteRecordsResponseData.DeleteRecordsTopicResultCollection(mergedResponseStatus.groupBy(_._1.topic).map { case (topic, partitionMap) => {
+            new DeleteRecordsTopicResult()
+              .setName(topic)
+              .setPartitions(new DeleteRecordsResponseData.DeleteRecordsPartitionResultCollection(partitionMap.map { case (topicPartition, partitionResult) => {
+                new DeleteRecordsPartitionResult().setPartitionIndex(topicPartition.partition)
+                  .setLowWatermark(partitionResult.lowWatermark)
+                  .setErrorCode(partitionResult.errorCode)
+              }
+              }.toList.asJava.iterator()))
+          }
+          }.toList.asJava.iterator()))))
     }
 
     if (authorizedForDeleteTopicOffsets.isEmpty)
@@ -1463,7 +2113,7 @@ class KafkaApis(val requestChannel: RequestChannel,
     else {
       // call the replica manager to append messages to the replicas
       replicaManager.deleteRecords(
-        deleteRecordsRequest.timeout.toLong,
+        deleteRecordsRequest.data.timeoutMs.toLong,
         authorizedForDeleteTopicOffsets,
         sendResponseCallback)
     }
@@ -1471,57 +2121,96 @@ 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, new Resource(TransactionalId, transactionalId))) {
+      if (!authorize(request.context, 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.context, 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 finalError =
+          if (initProducerIdRequest.version < 4 && result.error == Errors.PRODUCER_FENCED) {
+            // For older clients, they could not understand the new PRODUCER_FENCED error code,
+            // so we need to return the INVALID_PRODUCER_EPOCH to have the same client handling logic.
+            Errors.INVALID_PRODUCER_EPOCH
+          } else {
+            result.error
+          }
+        val responseData = new InitProducerIdResponseData()
+          .setProducerId(result.producerId)
+          .setProducerEpoch(result.producerEpoch)
+          .setThrottleTimeMs(requestThrottleMs)
+          .setErrorCode(finalError.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)
+
+    val producerIdAndEpoch = (initProducerIdRequest.data.producerId, initProducerIdRequest.data.producerEpoch) match {
+      case (RecordBatch.NO_PRODUCER_ID, RecordBatch.NO_PRODUCER_EPOCH) => Right(None)
+      case (RecordBatch.NO_PRODUCER_ID, _) | (_, RecordBatch.NO_PRODUCER_EPOCH) => Left(Errors.INVALID_REQUEST)
+      case (_, _) => Right(Some(new ProducerIdAndEpoch(initProducerIdRequest.data.producerId, initProducerIdRequest.data.producerEpoch)))
+    }
+
+    producerIdAndEpoch match {
+      case Right(producerIdAndEpoch) => txnCoordinator.handleInitProducerId(transactionalId, initProducerIdRequest.data.transactionTimeoutMs,
+        producerIdAndEpoch, sendResponseCallback)
+      case Left(error) => sendErrorResponseMaybeThrottle(request, error.exception)
+    }
   }
 
   def handleEndTxnRequest(request: RequestChannel.Request): Unit = {
     ensureInterBrokerVersion(KAFKA_0_11_0_IV0)
     val endTxnRequest = request.body[EndTxnRequest]
-    val transactionalId = endTxnRequest.transactionalId
+    val transactionalId = endTxnRequest.data.transactionalId
 
-    if (authorize(request.session, Write, new Resource(TransactionalId, transactionalId))) {
-      def sendResponseCallback(error: Errors) {
+    if (authorize(request.context, 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}.")
+          val finalError =
+            if (endTxnRequest.version < 2 && error == Errors.PRODUCER_FENCED) {
+              // For older clients, they could not understand the new PRODUCER_FENCED error code,
+              // so we need to return the INVALID_PRODUCER_EPOCH to have the same client handling logic.
+              Errors.INVALID_PRODUCER_EPOCH
+            } else {
+              error
+            }
+          val responseBody = new EndTxnResponse(new EndTxnResponseData()
+            .setErrorCode(finalError.code)
+            .setThrottleTimeMs(requestThrottleMs))
+          trace(s"Completed ${endTxnRequest.data.transactionalId}'s EndTxnRequest " +
+            s"with committed: ${endTxnRequest.data.committed}, " +
+            s"errors: $error from client ${request.header.clientId}.")
           responseBody
         }
         sendResponseMaybeThrottle(request, createResponse)
       }
 
-      txnCoordinator.handleEndTransaction(endTxnRequest.transactionalId,
-        endTxnRequest.producerId,
-        endTxnRequest.producerEpoch,
-        endTxnRequest.command,
+      txnCoordinator.handleEndTransaction(endTxnRequest.data.transactionalId,
+        endTxnRequest.data.producerId,
+        endTxnRequest.data.producerEpoch,
+        endTxnRequest.result(),
         sendResponseCallback)
     } else
       sendResponseMaybeThrottle(request, requestThrottleMs =>
-        new EndTxnResponse(requestThrottleMs, Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED))
+        new EndTxnResponse(new EndTxnResponseData()
+            .setErrorCode(Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED.code)
+            .setThrottleTimeMs(requestThrottleMs))
+      )
   }
 
   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
@@ -1546,7 +2235,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
@@ -1556,7 +2245,7 @@ class KafkaApis(val requestChannel: RequestChannel,
         // as soon as the end transaction marker has been written for a transactional offset commit,
         // call to the group coordinator to materialize the offsets into the cache
         try {
-          groupCoordinator.handleTxnCompletion(producerId, successfulOffsetsPartitions, result)
+          groupCoordinator.scheduleHandleTxnCompletion(producerId, successfulOffsetsPartitions, result)
         } catch {
           case e: Exception =>
             error(s"Received an exception while trying to update the offsets cache on transaction marker append", e)
@@ -1580,7 +2269,7 @@ class KafkaApis(val requestChannel: RequestChannel,
       val partitionsWithCompatibleMessageFormat = new mutable.ArrayBuffer[TopicPartition]
 
       val currentErrors = new ConcurrentHashMap[TopicPartition, Errors]()
-      marker.partitions.asScala.foreach { partition =>
+      marker.partitions.forEach { partition =>
         replicaManager.getMagic(partition) match {
           case Some(magic) =>
             if (magic < RecordBatch.MAGIC_VALUE_V2)
@@ -1612,7 +2301,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))
       }
@@ -1620,7 +2309,7 @@ class KafkaApis(val requestChannel: RequestChannel,
 
     // No log appends were written as all partitions had incorrect log format
     // so we need to send the error response
-    if (skippedMarkers == markers.size())
+    if (skippedMarkers == markers.size)
       sendResponseExemptThrottle(request, new WriteTxnMarkersResponse(errors))
   }
 
@@ -1632,9 +2321,9 @@ class KafkaApis(val requestChannel: RequestChannel,
   def handleAddPartitionToTxnRequest(request: RequestChannel.Request): Unit = {
     ensureInterBrokerVersion(KAFKA_0_11_0_IV0)
     val addPartitionsToTxnRequest = request.body[AddPartitionsToTxnRequest]
-    val transactionalId = addPartitionsToTxnRequest.transactionalId
+    val transactionalId = addPartitionsToTxnRequest.data.transactionalId
     val partitionsToAdd = addPartitionsToTxnRequest.partitions.asScala
-    if (!authorize(request.session, Write, new Resource(TransactionalId, transactionalId)))
+    if (!authorize(request.context, WRITE, TRANSACTIONAL_ID, transactionalId))
       sendResponseMaybeThrottle(request, requestThrottleMs =>
         addPartitionsToTxnRequest.getErrorResponse(requestThrottleMs, Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED.exception))
     else {
@@ -1642,11 +2331,12 @@ class KafkaApis(val requestChannel: RequestChannel,
       val nonExistingTopicErrors = mutable.Map[TopicPartition, Errors]()
       val authorizedPartitions = mutable.Set[TopicPartition]()
 
+      val authorizedTopics = filterByAuthorized(request.context, WRITE, TOPIC,
+        partitionsToAdd.filterNot(tp => Topic.isInternal(tp.topic)))(_.topic)
       for (topicPartition <- partitionsToAdd) {
-        if (org.apache.kafka.common.internals.Topic.isInternal(topicPartition.topic) ||
-            !authorize(request.session, Write, new Resource(Topic, topicPartition.topic)))
+        if (!authorizedTopics.contains(topicPartition.topic))
           unauthorizedTopicErrors += topicPartition -> Errors.TOPIC_AUTHORIZATION_FAILED
-        else if (!metadataCache.contains(topicPartition.topic))
+        else if (!metadataCache.contains(topicPartition))
           nonExistingTopicErrors += topicPartition -> Errors.UNKNOWN_TOPIC_OR_PARTITION
         else
           authorizedPartitions.add(topicPartition)
@@ -1663,18 +2353,28 @@ class KafkaApis(val requestChannel: RequestChannel,
       } else {
         def sendResponseCallback(error: Errors): Unit = {
           def createResponse(requestThrottleMs: Int): AbstractResponse = {
+            val finalError =
+              if (addPartitionsToTxnRequest.version < 2 && error == Errors.PRODUCER_FENCED) {
+                // For older clients, they could not understand the new PRODUCER_FENCED error code,
+                // so we need to return the old INVALID_PRODUCER_EPOCH to have the same client handling logic.
+                Errors.INVALID_PRODUCER_EPOCH
+              } else {
+                error
+              }
+
             val responseBody: AddPartitionsToTxnResponse = new AddPartitionsToTxnResponse(requestThrottleMs,
-              partitionsToAdd.map{tp => (tp, error)}.toMap.asJava)
+              partitionsToAdd.map{tp => (tp, finalError)}.toMap.asJava)
             trace(s"Completed $transactionalId's AddPartitionsToTxnRequest with partitions $partitionsToAdd: errors: $error from client ${request.header.clientId}")
             responseBody
           }
 
+
           sendResponseMaybeThrottle(request, createResponse)
         }
 
         txnCoordinator.handleAddPartitionsToTransaction(transactionalId,
-          addPartitionsToTxnRequest.producerId,
-          addPartitionsToTxnRequest.producerEpoch,
+          addPartitionsToTxnRequest.data.producerId,
+          addPartitionsToTxnRequest.data.producerEpoch,
           authorizedPartitions,
           sendResponseCallback)
       }
@@ -1684,20 +2384,37 @@ class KafkaApis(val requestChannel: RequestChannel,
   def handleAddOffsetsToTxnRequest(request: RequestChannel.Request): Unit = {
     ensureInterBrokerVersion(KAFKA_0_11_0_IV0)
     val addOffsetsToTxnRequest = request.body[AddOffsetsToTxnRequest]
-    val transactionalId = addOffsetsToTxnRequest.transactionalId
-    val groupId = addOffsetsToTxnRequest.consumerGroupId
+    val transactionalId = addOffsetsToTxnRequest.data.transactionalId
+    val groupId = addOffsetsToTxnRequest.data.groupId
     val offsetTopicPartition = new TopicPartition(GROUP_METADATA_TOPIC_NAME, groupCoordinator.partitionFor(groupId))
 
-    if (!authorize(request.session, Write, new Resource(TransactionalId, transactionalId)))
+    if (!authorize(request.context, WRITE, TRANSACTIONAL_ID, transactionalId))
       sendResponseMaybeThrottle(request, requestThrottleMs =>
-        new AddOffsetsToTxnResponse(requestThrottleMs, Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED))
-    else if (!authorize(request.session, Read, new Resource(Group, groupId)))
+        new AddOffsetsToTxnResponse(new AddOffsetsToTxnResponseData()
+          .setErrorCode(Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED.code)
+          .setThrottleTimeMs(requestThrottleMs)))
+    else if (!authorize(request.context, READ, GROUP, groupId))
       sendResponseMaybeThrottle(request, requestThrottleMs =>
-        new AddOffsetsToTxnResponse(requestThrottleMs, Errors.GROUP_AUTHORIZATION_FAILED))
+        new AddOffsetsToTxnResponse(new AddOffsetsToTxnResponseData()
+          .setErrorCode(Errors.GROUP_AUTHORIZATION_FAILED.code)
+          .setThrottleTimeMs(requestThrottleMs))
+      )
     else {
       def sendResponseCallback(error: Errors): Unit = {
         def createResponse(requestThrottleMs: Int): AbstractResponse = {
-          val responseBody: AddOffsetsToTxnResponse = new AddOffsetsToTxnResponse(requestThrottleMs, error)
+          val finalError =
+            if (addOffsetsToTxnRequest.version < 2 && error == Errors.PRODUCER_FENCED) {
+              // For older clients, they could not understand the new PRODUCER_FENCED error code,
+              // so we need to return the old INVALID_PRODUCER_EPOCH to have the same client handling logic.
+              Errors.INVALID_PRODUCER_EPOCH
+            } else {
+              error
+            }
+
+          val responseBody: AddOffsetsToTxnResponse = new AddOffsetsToTxnResponse(
+            new AddOffsetsToTxnResponseData()
+              .setErrorCode(finalError.code)
+              .setThrottleTimeMs(requestThrottleMs))
           trace(s"Completed $transactionalId's AddOffsetsToTxnRequest for group $groupId on partition " +
             s"$offsetTopicPartition: errors: $error from client ${request.header.clientId}")
           responseBody
@@ -1706,8 +2423,8 @@ class KafkaApis(val requestChannel: RequestChannel,
       }
 
       txnCoordinator.handleAddPartitionsToTransaction(transactionalId,
-        addOffsetsToTxnRequest.producerId,
-        addOffsetsToTxnRequest.producerEpoch,
+        addOffsetsToTxnRequest.data.producerId,
+        addOffsetsToTxnRequest.data.producerEpoch,
         Set(offsetTopicPartition),
         sendResponseCallback)
     }
@@ -1720,34 +2437,48 @@ 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, new Resource(TransactionalId, txnOffsetCommitRequest.transactionalId)))
+    if (!authorize(request.context, WRITE, TRANSACTIONAL_ID, txnOffsetCommitRequest.data.transactionalId))
       sendErrorResponseMaybeThrottle(request, Errors.TRANSACTIONAL_ID_AUTHORIZATION_FAILED.exception)
-    else if (!authorize(request.session, Read, new Resource(Group, txnOffsetCommitRequest.consumerGroupId)))
+    else if (!authorize(request.context, 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 committedOffsets = txnOffsetCommitRequest.offsets.asScala
+      val authorizedTopics = filterByAuthorized(request.context, READ, TOPIC, committedOffsets)(_._1.topic)
 
-      for ((topicPartition, commitedOffset) <- txnOffsetCommitRequest.offsets.asScala) {
-        if (!authorize(request.session, Read, new Resource(Topic, topicPartition.topic)))
+      for ((topicPartition, commitedOffset) <- committedOffsets) {
+        if (!authorizedTopics.contains(topicPartition.topic))
           unauthorizedTopicErrors += topicPartition -> Errors.TOPIC_AUTHORIZATION_FAILED
-        else if (!metadataCache.contains(topicPartition.topic))
+        else if (!metadataCache.contains(topicPartition))
           nonExistingTopicErrors += topicPartition -> Errors.UNKNOWN_TOPIC_OR_PARTITION
         else
           authorizedTopicCommittedOffsets += (topicPartition -> commitedOffset)
       }
 
       // the callback for sending an offset commit response
-      def sendResponseCallback(authorizedTopicErrors: Map[TopicPartition, Errors]) {
-        val combinedCommitStatus = authorizedTopicErrors ++ unauthorizedTopicErrors ++ nonExistingTopicErrors
+      def sendResponseCallback(authorizedTopicErrors: Map[TopicPartition, Errors]): Unit = {
+        val combinedCommitStatus = mutable.Map() ++= authorizedTopicErrors ++= unauthorizedTopicErrors ++= nonExistingTopicErrors
         if (isDebugEnabled)
-          combinedCommitStatus.foreach { case (topicPartition, error) =>
+          combinedCommitStatus.forKeyValue { (topicPartition, error) =>
             if (error != Errors.NONE) {
               debug(s"TxnOffsetCommit with correlation id ${header.correlationId} from client ${header.clientId} " +
                 s"on partition $topicPartition failed due to ${error.exceptionName}")
             }
           }
+
+        // We need to replace COORDINATOR_LOAD_IN_PROGRESS with COORDINATOR_NOT_AVAILABLE
+        // for older producer client from 0.11 to prior 2.0, which could potentially crash due
+        // to unexpected loading error. This bug is fixed later by KAFKA-7296. Clients using
+        // txn commit protocol >= 2 (version 2.3 and onwards) are guaranteed to have
+        // the fix to check for the loading error.
+        if (txnOffsetCommitRequest.version < 2) {
+          combinedCommitStatus ++= combinedCommitStatus.collect {
+            case (tp, error) if error == Errors.COORDINATOR_LOAD_IN_PROGRESS => tp -> Errors.COORDINATOR_NOT_AVAILABLE
+          }
+        }
+
         sendResponseMaybeThrottle(request, requestThrottleMs =>
           new TxnOffsetCommitResponse(requestThrottleMs, combinedCommitStatus.asJava))
       }
@@ -1757,9 +2488,12 @@ 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,
+          txnOffsetCommitRequest.data.memberId,
+          Option(txnOffsetCommitRequest.data.groupInstanceId),
+          txnOffsetCommitRequest.data.generationId,
           offsetMetadata,
           sendResponseCallback)
       }
@@ -1767,79 +2501,94 @@ class KafkaApis(val requestChannel: RequestChannel,
   }
 
   private def convertTxnOffsets(offsetsMap: immutable.Map[TopicPartition, TxnOffsetCommitRequest.CommittedOffset]): immutable.Map[TopicPartition, OffsetAndMetadata] = {
-    val offsetRetention = groupCoordinator.offsetConfig.offsetsRetentionMs
     val currentTimestamp = time.milliseconds
-    val defaultExpireTimestamp = offsetRetention + currentTimestamp
     offsetsMap.map { case (topicPartition, partitionData) =>
-      val metadata = if (partitionData.metadata == null) OffsetMetadata.NoMetadata else partitionData.metadata
+      val metadata = if (partitionData.metadata == null) OffsetAndMetadata.NoMetadata else partitionData.metadata
       topicPartition -> new OffsetAndMetadata(
-        offsetMetadata = OffsetMetadata(partitionData.offset, metadata),
+        offset = partitionData.offset,
+        leaderEpoch = partitionData.leaderEpoch,
+        metadata = metadata,
         commitTimestamp = currentTimestamp,
-        expireTimestamp = defaultExpireTimestamp)
+        expireTimestamp = None)
     }
   }
 
   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 DescribeAclsResponse(new DescribeAclsResponseData()
+            .setErrorCode(Errors.SECURITY_DISABLED.code)
+            .setErrorMessage("No Authorizer is configured on the broker")
+            .setThrottleTimeMs(requestThrottleMs),
+          describeAclsRequest.version))
       case Some(auth) =>
-        val filter = describeAclsRequest.filter()
-        val returnedAcls = auth.getAcls.toSeq.flatMap { case (resource, acls) =>
-          acls.flatMap { acl =>
-            val fixture = new AclBinding(new AdminResource(resource.resourceType.toJava, resource.name),
-                new AccessControlEntry(acl.principal.toString, acl.host.toString, acl.operation.toJava, acl.permissionType.toJava))
-            if (filter.matches(fixture)) Some(fixture)
-            else None
-          }
-        }
+        val filter = describeAclsRequest.filter
+        val returnedAcls = new util.HashSet[AclBinding]()
+        auth.acls(filter).forEach(returnedAcls.add)
         sendResponseMaybeThrottle(request, requestThrottleMs =>
-          new DescribeAclsResponse(requestThrottleMs, ApiError.NONE, returnedAcls.asJava))
+          new DescribeAclsResponse(new DescribeAclsResponseData()
+            .setThrottleTimeMs(requestThrottleMs)
+            .setResources(DescribeAclsResponse.aclsResources(returnedAcls)),
+          describeAclsRequest.version))
     }
   }
 
   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 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) &&
-                    !resource.name.equals(Resource.ClusterResourceName))
-                  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 allBindings = createAclsRequest.aclCreations.asScala.map(CreateAclsRequest.aclBinding)
+        val errorResults = mutable.Map[AclBinding, AclCreateResult]()
+        val validBindings = new ArrayBuffer[AclBinding]
+        allBindings.foreach { 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)
+          } else
+            validBindings += acl
+        }
+
+        val createResults = auth.createAcls(request.context, validBindings.asJava).asScala.map(_.toCompletableFuture)
+
+        def sendResponseCallback(): Unit = {
+          val aclCreationResults = allBindings.map { acl =>
+            val result = errorResults.getOrElse(acl, createResults(validBindings.indexOf(acl)).get)
+            val creationResult = new AclCreationResult()
+            result.exception.asScala.foreach { throwable =>
+              val apiError = ApiError.fromThrowable(throwable)
+              creationResult
+                .setErrorCode(apiError.error.code)
+                .setErrorMessage(apiError.message)
+            }
+            creationResult
           }
+          sendResponseMaybeThrottle(request, requestThrottleMs =>
+            new CreateAclsResponse(new CreateAclsResponseData()
+              .setThrottleTimeMs(requestThrottleMs)
+              .setResults(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 =>
@@ -1847,162 +2596,784 @@ 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 AdminResource(resource.resourceType.toJava, resource.name),
-              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 sendResponseCallback(): Unit = {
+          val filterResults = deleteResults.map(_.get).map(DeleteAclsResponse.filterResult).asJava
+          sendResponseMaybeThrottle(request, requestThrottleMs =>
+            new DeleteAclsResponse(
+              new DeleteAclsResponseData()
+                .setThrottleTimeMs(requestThrottleMs)
+                .setFilterResults(filterResults),
+              deleteAclsRequest.version))
         }
-
-        val filterResponses = filters.indices.map { i =>
-          filterResponseMap.getOrElse(i, new AclFilterResponse(Seq.empty.asJava))
-        }.asJava
-        sendResponseMaybeThrottle(request, requestThrottleMs => new DeleteAclsResponse(requestThrottleMs, filterResponses))
+        alterAclsPurgatory.tryCompleteElseWatch(config.connectionsMaxIdleMs, deleteResults, sendResponseCallback)
     }
   }
 
   def handleOffsetForLeaderEpochRequest(request: RequestChannel.Request): Unit = {
     val offsetForLeaderEpoch = request.body[OffsetsForLeaderEpochRequest]
-    val requestInfo = offsetForLeaderEpoch.epochsByTopicPartition()
-    authorizeClusterAction(request)
+    val topics = offsetForLeaderEpoch.data.topics.asScala.toSeq
+
+    // 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 (authorizedTopics, unauthorizedTopics) =
+      if (authorize(request.context, CLUSTER_ACTION, CLUSTER, CLUSTER_NAME, logIfDenied = false))
+        (topics, Seq.empty[OffsetForLeaderTopic])
+      else partitionSeqByAuthorized(request.context, DESCRIBE, TOPIC, topics)(_.topic)
+
+    val endOffsetsForAuthorizedPartitions = replicaManager.lastOffsetForLeaderEpoch(authorizedTopics)
+    val endOffsetsForUnauthorizedPartitions = unauthorizedTopics.map { offsetForLeaderTopic =>
+      val partitions = offsetForLeaderTopic.partitions.asScala.map { offsetForLeaderPartition =>
+        new EpochEndOffset()
+          .setPartition(offsetForLeaderPartition.partition)
+          .setErrorCode(Errors.TOPIC_AUTHORIZATION_FAILED.code)
+      }
+
+      new OffsetForLeaderTopicResult()
+        .setTopic(offsetForLeaderTopic.topic)
+        .setPartitions(partitions.toList.asJava)
+    }
+
+    val endOffsetsForAllTopics = new OffsetForLeaderTopicResultCollection(
+      (endOffsetsForAuthorizedPartitions ++ endOffsetsForUnauthorizedPartitions).asJava.iterator
+    )
 
-    val lastOffsetForLeaderEpoch = replicaManager.lastOffsetForLeaderEpoch(requestInfo.asScala).asJava
-    sendResponseExemptThrottle(request, new OffsetsForLeaderEpochResponse(lastOffsetForLeaderEpoch))
+    sendResponseMaybeThrottle(request, requestThrottleMs =>
+      new OffsetsForLeaderEpochResponse(new OffsetForLeaderEpochResponseData()
+        .setThrottleTimeMs(requestThrottleMs)
+        .setTopics(endOffsetsForAllTopics)))
   }
 
   def handleAlterConfigsRequest(request: RequestChannel.Request): Unit = {
     val alterConfigsRequest = request.body[AlterConfigsRequest]
-    val (authorizedResources, unauthorizedResources) = alterConfigsRequest.configs.asScala.partition { case (resource, _) =>
+    val (authorizedResources, unauthorizedResources) = alterConfigsRequest.configs.asScala.toMap.partition { case (resource, _) =>
       resource.`type` match {
-        case RResourceType.BROKER =>
-          authorize(request.session, AlterConfigs, Resource.ClusterResource)
-        case RResourceType.TOPIC =>
-          authorize(request.session, AlterConfigs, new Resource(Topic, resource.name))
+        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.context, ALTER_CONFIGS, CLUSTER, CLUSTER_NAME)
+        case ConfigResource.Type.TOPIC =>
+          authorize(request.context, 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)
+    }
+    def responseCallback(requestThrottleMs: Int): AlterConfigsResponse = {
+      val data = new AlterConfigsResponseData()
+        .setThrottleTimeMs(requestThrottleMs)
+      (authorizedResult ++ unauthorizedResult).foreach{ case (resource, error) =>
+        data.responses().add(new AlterConfigsResourceResponse()
+          .setErrorCode(error.error.code)
+          .setErrorMessage(error.message)
+          .setResourceName(resource.name)
+          .setResourceType(resource.`type`.id))
+      }
+      new AlterConfigsResponse(data)
     }
-    sendResponseMaybeThrottle(request, requestThrottleMs =>
-      new AlterConfigsResponse(requestThrottleMs, (authorizedResult ++ unauthorizedResult).asJava))
+    sendResponseMaybeThrottle(request, responseCallback)
+  }
+
+  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(session: RequestChannel.Session, resource: RResource): ApiError = {
+  private def configsAuthorizationApiError(resource: ConfigResource): ApiError = {
     val error = resource.`type` match {
-      case RResourceType.BROKER => Errors.CLUSTER_AUTHORIZATION_FAILED
-      case RResourceType.TOPIC => Errors.TOPIC_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))
+      }.toBuffer
+    }.toMap
+
+    val (authorizedResources, unauthorizedResources) = configs.partition { case (resource, _) =>
+      resource.`type` match {
+        case ConfigResource.Type.BROKER | ConfigResource.Type.BROKER_LOGGER =>
+          authorize(request.context, ALTER_CONFIGS, CLUSTER, CLUSTER_NAME)
+        case ConfigResource.Type.TOPIC =>
+          authorize(request.context, 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(
+      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 RResourceType.BROKER => authorize(request.session, DescribeConfigs, Resource.ClusterResource)
-        case RResourceType.TOPIC =>
-          authorize(request.session, DescribeConfigs, new Resource(Topic, resource.name))
-        case rt => throw new InvalidRequestException(s"Unexpected resource type $rt for resource ${resource.name}")
+    val (authorizedResources, unauthorizedResources) = describeConfigsRequest.data.resources.asScala.partition { resource =>
+      ConfigResource.Type.forId(resource.resourceType) match {
+        case ConfigResource.Type.BROKER | ConfigResource.Type.BROKER_LOGGER =>
+          authorize(request.context, DESCRIBE_CONFIGS, CLUSTER, CLUSTER_NAME)
+        case ConfigResource.Type.TOPIC =>
+          authorize(request.context, DESCRIBE_CONFIGS, TOPIC, resource.resourceName)
+        case rt => throw new InvalidRequestException(s"Unexpected resource type $rt for resource ${resource.resourceName}")
       }
     }
-    val authorizedConfigs = adminManager.describeConfigs(authorizedResources.map { resource =>
-      resource -> Option(describeConfigsRequest.configNames(resource)).map(_.asScala.toSet)
-    }.toMap)
+    val authorizedConfigs = adminManager.describeConfigs(authorizedResources.toList, describeConfigsRequest.data.includeSynonyms, describeConfigsRequest.data.includeDocumentation)
     val unauthorizedConfigs = unauthorizedResources.map { resource =>
-      val error = configsAuthorizationApiError(request.session, resource)
-      resource -> new DescribeConfigsResponse.Config(error, Collections.emptyList[DescribeConfigsResponse.ConfigEntry])
+      val error = ConfigResource.Type.forId(resource.resourceType) match {
+        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.resourceName}")
+      }
+      new DescribeConfigsResponseData.DescribeConfigsResult().setErrorCode(error.code)
+        .setErrorMessage(error.message)
+        .setConfigs(Collections.emptyList[DescribeConfigsResponseData.DescribeConfigsResourceResult])
+        .setResourceName(resource.resourceName)
+        .setResourceType(resource.resourceType)
     }
 
     sendResponseMaybeThrottle(request, requestThrottleMs =>
-      new DescribeConfigsResponse(requestThrottleMs, (authorizedConfigs ++ unauthorizedConfigs).asJava))
+      new DescribeConfigsResponse(new DescribeConfigsResponseData().setThrottleTimeMs(requestThrottleMs)
+        .setResults((authorizedConfigs ++ unauthorizedConfigs).asJava)))
   }
 
   def handleAlterReplicaLogDirsRequest(request: RequestChannel.Request): Unit = {
     val alterReplicaDirsRequest = request.body[AlterReplicaLogDirsRequest]
-    val responseMap = {
-      if (authorize(request.session, Alter, Resource.ClusterResource))
-        replicaManager.alterReplicaLogDirs(alterReplicaDirsRequest.partitionDirs.asScala)
-      else
-        alterReplicaDirsRequest.partitionDirs.asScala.keys.map((_, Errors.CLUSTER_AUTHORIZATION_FAILED)).toMap
+    if (authorize(request.context, ALTER, CLUSTER, CLUSTER_NAME)) {
+      val result = replicaManager.alterReplicaLogDirs(alterReplicaDirsRequest.partitionDirs.asScala)
+      sendResponseMaybeThrottle(request, requestThrottleMs =>
+        new AlterReplicaLogDirsResponse(new AlterReplicaLogDirsResponseData()
+          .setResults(result.groupBy(_._1.topic).map {
+            case (topic, errors) => new AlterReplicaLogDirsResponseData.AlterReplicaLogDirTopicResult()
+              .setTopicName(topic)
+              .setPartitions(errors.map {
+                case (tp, error) => new AlterReplicaLogDirsResponseData.AlterReplicaLogDirPartitionResult()
+                  .setPartitionIndex(tp.partition)
+                  .setErrorCode(error.code)
+              }.toList.asJava)
+          }.toList.asJava)
+          .setThrottleTimeMs(requestThrottleMs)))
+    } else {
+      sendResponseMaybeThrottle(request, requestThrottleMs =>
+        alterReplicaDirsRequest.getErrorResponse(requestThrottleMs, Errors.CLUSTER_AUTHORIZATION_FAILED.exception))
     }
-    sendResponseMaybeThrottle(request, requestThrottleMs => new AlterReplicaLogDirsResponse(requestThrottleMs, responseMap.asJava))
   }
 
   def handleDescribeLogDirsRequest(request: RequestChannel.Request): Unit = {
     val describeLogDirsDirRequest = request.body[DescribeLogDirsRequest]
     val logDirInfos = {
-      if (authorize(request.session, Describe, Resource.ClusterResource)) {
+      if (authorize(request.context, DESCRIBE, CLUSTER, CLUSTER_NAME)) {
         val partitions =
           if (describeLogDirsDirRequest.isAllTopicPartitions)
             replicaManager.logManager.allLogs.map(_.topicPartition).toSet
           else
-            describeLogDirsDirRequest.topicPartitions().asScala
+            describeLogDirsDirRequest.data.topics.asScala.flatMap(
+              logDirTopic => logDirTopic.partitionIndex.asScala.map(partitionIndex =>
+                new TopicPartition(logDirTopic.topic, partitionIndex))).toSet
 
         replicaManager.describeLogDirs(partitions)
       } else {
-        Map.empty[String, LogDirInfo]
+        List.empty[DescribeLogDirsResponseData.DescribeLogDirsResult]
       }
     }
-    sendResponseMaybeThrottle(request, throttleTimeMs => new DescribeLogDirsResponse(throttleTimeMs, logDirInfos.asJava))
+    sendResponseMaybeThrottle(request, throttleTimeMs => new DescribeLogDirsResponse(new DescribeLogDirsResponseData()
+      .setThrottleTimeMs(throttleTimeMs)
+      .setResults(logDirInfos.asJava)))
   }
 
-  def authorizeClusterAction(request: RequestChannel.Request): Unit = {
-    if (!authorize(request.session, ClusterAction, Resource.ClusterResource))
-      throw new ClusterAuthorizationException(s"Request $request is not authorized.")
+  def handleCreateTokenRequest(request: RequestChannel.Request): Unit = {
+    val createTokenRequest = request.body[CreateDelegationTokenRequest]
+
+    // the callback for sending a create token response
+    def sendResponseCallback(createResult: CreateTokenResult): Unit = {
+      trace(s"Sending create token response for correlation id ${request.header.correlationId} " +
+        s"to client ${request.header.clientId}.")
+      sendResponseMaybeThrottle(request, requestThrottleMs =>
+        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 =>
+        CreateDelegationTokenResponse.prepareResponse(requestThrottleMs, Errors.DELEGATION_TOKEN_REQUEST_NOT_ALLOWED, request.context.principal))
+    else {
+      val renewerList = createTokenRequest.data.renewers.asScala.toList.map(entry =>
+        new KafkaPrincipal(entry.principalType, entry.principalName))
+
+      if (renewerList.exists(principal => principal.getPrincipalType != KafkaPrincipal.USER_TYPE)) {
+        sendResponseMaybeThrottle(request, requestThrottleMs =>
+          CreateDelegationTokenResponse.prepareResponse(requestThrottleMs, Errors.INVALID_PRINCIPAL_TYPE, request.context.principal))
+      }
+      else {
+        tokenManager.createToken(
+          request.context.principal,
+          renewerList,
+          createTokenRequest.data.maxLifetimeMs,
+          sendResponseCallback
+        )
+      }
+    }
   }
 
-  def authorizeClusterAlter(request: RequestChannel.Request): Unit = {
-    if (!authorize(request.session, Alter, Resource.ClusterResource))
-      throw new ClusterAuthorizationException(s"Request $request is not authorized.")
+  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): Unit = {
+      trace("Sending renew token response for correlation id %d to client %s."
+        .format(request.header.correlationId, request.header.clientId))
+      sendResponseMaybeThrottle(request, requestThrottleMs =>
+        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.context.principal,
+        ByteBuffer.wrap(renewTokenRequest.data.hmac),
+        renewTokenRequest.data.renewPeriodMs,
+        sendResponseCallback
+      )
+    }
+  }
+
+  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): 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(
+            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.context.principal,
+        expireTokenRequest.hmac(),
+        expireTokenRequest.expiryTimePeriod(),
+        sendResponseCallback
+      )
+    }
   }
 
-  def authorizeClusterDescribe(request: RequestChannel.Request): Unit = {
-    if (!authorize(request.session, Describe, Resource.ClusterResource))
+  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]): Unit = {
+      sendResponseMaybeThrottle(request, requestThrottleMs =>
+        new DescribeDelegationTokenResponse(requestThrottleMs, error, tokenDetails.asJava))
+      trace("Sending describe token response for correlation id %d to client %s."
+        .format(request.header.correlationId, request.header.clientId))
+    }
+
+    if (!allowTokenRequests(request))
+      sendResponseCallback(Errors.DELEGATION_TOKEN_REQUEST_NOT_ALLOWED, List.empty)
+    else if (!config.tokenAuthEnabled)
+      sendResponseCallback(Errors.DELEGATION_TOKEN_AUTH_DISABLED, List.empty)
+    else {
+      val requestPrincipal = request.context.principal
+
+      if (describeTokenRequest.ownersListEmpty()) {
+        sendResponseCallback(Errors.NONE, List())
+      }
+      else {
+        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.context, DESCRIBE, DELEGATION_TOKEN, tokenId)
+        def eligible(token: TokenInformation) = DelegationTokenManager.filterToken(requestPrincipal, owners, token, authorizeToken)
+        val tokens =  tokenManager.getTokens(eligible)
+        sendResponseCallback(Errors.NONE, tokens)
+      }
+    }
+  }
+
+  def allowTokenRequests(request: RequestChannel.Request): Boolean = {
+    val protocol = request.context.securityProtocol
+    if (request.context.principal.tokenAuthenticated ||
+      protocol == SecurityProtocol.PLAINTEXT ||
+      // disallow requests from 1-way SSL
+      (protocol == SecurityProtocol.SSL && request.context.principal == KafkaPrincipal.ANONYMOUS))
+      false
+    else
+      true
+  }
+
+  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 }
+          .forKeyValue { (topic, ps) =>
+            val electionResult = new ReplicaElectionResult()
+
+            electionResult.setTopic(topic)
+            ps.forKeyValue { (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.context, 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
+
+      sendResponseCallback(error)(partitionErrors)
+    } else {
+      val partitions = if (electionRequest.data.topicPartitions == null) {
+        metadataCache.getAllPartitions()
+      } else {
+        electionRequest.topicPartitions
+      }
+
+      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.context, DELETE, GROUP, groupId)) {
+      val topics = offsetDeleteRequest.data.topics.asScala
+      val authorizedTopics = filterByAuthorized(request.context, READ, TOPIC, topics)(_.name)
+
+      val topicPartitionErrors = mutable.Map[TopicPartition, Errors]()
+      val topicPartitions = mutable.ArrayBuffer[TopicPartition]()
+
+      for (topic <- topics) {
+        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
+        }
+      }
+
+      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).forKeyValue { (topic, topicPartitions) =>
+            val partitions = new OffsetDeleteResponseData.OffsetDeleteResponsePartitionCollection
+            topicPartitions.forKeyValue { (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 {
+      sendResponseMaybeThrottle(request, requestThrottleMs =>
+        offsetDeleteRequest.getErrorResponse(requestThrottleMs, Errors.GROUP_AUTHORIZATION_FAILED))
+    }
+  }
+
+  def handleDescribeClientQuotasRequest(request: RequestChannel.Request): Unit = {
+    val describeClientQuotasRequest = request.body[DescribeClientQuotasRequest]
+
+    if (authorize(request.context, DESCRIBE_CONFIGS, CLUSTER, CLUSTER_NAME)) {
+      val result = adminManager.describeClientQuotas(describeClientQuotasRequest.filter)
+
+      val entriesData = result.iterator.map { case (quotaEntity, quotaValues) =>
+        val entityData = quotaEntity.entries.asScala.iterator.map { case (entityType, entityName) =>
+          new DescribeClientQuotasResponseData.EntityData()
+            .setEntityType(entityType)
+            .setEntityName(entityName)
+        }.toBuffer
+
+        val valueData = quotaValues.iterator.map { case (key, value) =>
+          new DescribeClientQuotasResponseData.ValueData()
+            .setKey(key)
+            .setValue(value)
+        }.toBuffer
+
+        new DescribeClientQuotasResponseData.EntryData()
+          .setEntity(entityData.asJava)
+          .setValues(valueData.asJava)
+      }.toBuffer
+
+      sendResponseMaybeThrottle(request, requestThrottleMs =>
+        new DescribeClientQuotasResponse(new DescribeClientQuotasResponseData()
+          .setThrottleTimeMs(requestThrottleMs)
+          .setEntries(entriesData.asJava)))
+    } else {
+      sendResponseMaybeThrottle(request, requestThrottleMs =>
+        describeClientQuotasRequest.getErrorResponse(requestThrottleMs, Errors.CLUSTER_AUTHORIZATION_FAILED.exception))
+    }
+  }
+
+  def handleAlterClientQuotasRequest(request: RequestChannel.Request): Unit = {
+    val alterClientQuotasRequest = request.body[AlterClientQuotasRequest]
+
+    if (authorize(request.context, ALTER_CONFIGS, CLUSTER, CLUSTER_NAME)) {
+      val result = adminManager.alterClientQuotas(alterClientQuotasRequest.entries.asScala,
+        alterClientQuotasRequest.validateOnly)
+
+      val entriesData = result.iterator.map { case (quotaEntity, apiError) =>
+        val entityData = quotaEntity.entries.asScala.iterator.map { case (key, value) =>
+          new AlterClientQuotasResponseData.EntityData()
+            .setEntityType(key)
+            .setEntityName(value)
+        }.toBuffer
+
+        new AlterClientQuotasResponseData.EntryData()
+          .setErrorCode(apiError.error.code)
+          .setErrorMessage(apiError.message)
+          .setEntity(entityData.asJava)
+      }.toBuffer
+
+      sendResponseMaybeThrottle(request, requestThrottleMs =>
+        new AlterClientQuotasResponse(new AlterClientQuotasResponseData()
+          .setThrottleTimeMs(requestThrottleMs)
+          .setEntries(entriesData.asJava)))
+    } else {
+      sendResponseMaybeThrottle(request, requestThrottleMs =>
+        alterClientQuotasRequest.getErrorResponse(requestThrottleMs, Errors.CLUSTER_AUTHORIZATION_FAILED.exception))
+    }
+  }
+
+  def handleDescribeUserScramCredentialsRequest(request: RequestChannel.Request): Unit = {
+    val describeUserScramCredentialsRequest = request.body[DescribeUserScramCredentialsRequest]
+
+    if (authorize(request.context, DESCRIBE, CLUSTER, CLUSTER_NAME)) {
+      val result = adminManager.describeUserScramCredentials(
+        Option(describeUserScramCredentialsRequest.data.users).map(_.asScala.map(_.name).toList))
+      sendResponseMaybeThrottle(request, requestThrottleMs =>
+        new DescribeUserScramCredentialsResponse(result.setThrottleTimeMs(requestThrottleMs)))
+    } else {
+      sendResponseMaybeThrottle(request, requestThrottleMs =>
+        describeUserScramCredentialsRequest.getErrorResponse(requestThrottleMs, Errors.CLUSTER_AUTHORIZATION_FAILED.exception))
+    }
+  }
+
+  def handleAlterUserScramCredentialsRequest(request: RequestChannel.Request): Unit = {
+    val alterUserScramCredentialsRequest = request.body[AlterUserScramCredentialsRequest]
+
+    if (!controller.isActive) {
+      sendResponseMaybeThrottle(request, requestThrottleMs =>
+        alterUserScramCredentialsRequest.getErrorResponse(requestThrottleMs, Errors.NOT_CONTROLLER.exception))
+    } else if (authorize(request.context, ALTER, CLUSTER, CLUSTER_NAME)) {
+      val result = adminManager.alterUserScramCredentials(
+        alterUserScramCredentialsRequest.data.upsertions().asScala, alterUserScramCredentialsRequest.data.deletions().asScala)
+      sendResponseMaybeThrottle(request, requestThrottleMs =>
+        new AlterUserScramCredentialsResponse(result.setThrottleTimeMs(requestThrottleMs)))
+    } else {
+      sendResponseMaybeThrottle(request, requestThrottleMs =>
+        alterUserScramCredentialsRequest.getErrorResponse(requestThrottleMs, Errors.CLUSTER_AUTHORIZATION_FAILED.exception))
+    }
+  }
+
+  def handleAlterIsrRequest(request: RequestChannel.Request): Unit = {
+    val alterIsrRequest = request.body[AlterIsrRequest]
+    authorizeClusterOperation(request, CLUSTER_ACTION)
+
+    if (!controller.isActive)
+      sendResponseExemptThrottle(request, alterIsrRequest.getErrorResponse(
+        AbstractResponse.DEFAULT_THROTTLE_TIME, Errors.NOT_CONTROLLER.exception))
+    else
+      controller.alterIsrs(alterIsrRequest.data, alterIsrResp =>
+        sendResponseExemptThrottle(request, new AlterIsrResponse(alterIsrResp))
+      )
+  }
+
+  def handleUpdateFeatures(request: RequestChannel.Request): Unit = {
+    val updateFeaturesRequest = request.body[UpdateFeaturesRequest]
+
+    def sendResponseCallback(errors: Either[ApiError, Map[String, ApiError]]): Unit = {
+      def createResponse(throttleTimeMs: Int): UpdateFeaturesResponse = {
+        errors match {
+          case Left(topLevelError) =>
+            UpdateFeaturesResponse.createWithErrors(
+              topLevelError,
+              Collections.emptyMap(),
+              throttleTimeMs)
+          case Right(featureUpdateErrors) =>
+            UpdateFeaturesResponse.createWithErrors(
+              ApiError.NONE,
+              featureUpdateErrors.asJava,
+              throttleTimeMs)
+        }
+      }
+      sendResponseMaybeThrottle(request, requestThrottleMs => createResponse(requestThrottleMs))
+    }
+
+    if (!authorize(request.context, ALTER, CLUSTER, CLUSTER_NAME)) {
+      sendResponseCallback(Left(new ApiError(Errors.CLUSTER_AUTHORIZATION_FAILED)))
+    } else if (!controller.isActive) {
+      sendResponseCallback(Left(new ApiError(Errors.NOT_CONTROLLER)))
+    } else if (!config.isFeatureVersioningSupported) {
+      sendResponseCallback(Left(new ApiError(Errors.INVALID_REQUEST, "Feature versioning system is disabled.")))
+    } else {
+      controller.updateFeatures(updateFeaturesRequest, sendResponseCallback)
+    }
+  }
+
+  // private package for testing
+  private[server] def authorize(requestContext: RequestContext,
+                                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(requestContext, actions).get(0) == AuthorizationResult.ALLOWED
+    }
+  }
+
+  // private package for testing
+  private[server] def filterByAuthorized[T](requestContext: RequestContext,
+                                            operation: AclOperation,
+                                            resourceType: ResourceType,
+                                            resources: Iterable[T],
+                                            logIfAllowed: Boolean = true,
+                                            logIfDenied: Boolean = true)(resourceName: T => String): Set[String] = {
+    authorizer match {
+      case Some(authZ) =>
+        val resourceNameToCount = CoreUtils.groupMapReduce(resources)(resourceName)(_ => 1)(_ + _)
+        val actions = resourceNameToCount.iterator.map { case (resourceName, count) =>
+          val resource = new ResourcePattern(resourceType, resourceName, PatternType.LITERAL)
+          new Action(operation, resource, count, logIfAllowed, logIfDenied)
+        }.toBuffer
+        authZ.authorize(requestContext, actions.asJava).asScala
+          .zip(resourceNameToCount.keySet)
+          .collect { case (authzResult, resourceName) if authzResult == AuthorizationResult.ALLOWED =>
+            resourceName
+          }.toSet
+      case None => resources.iterator.map(resourceName).toSet
+    }
+  }
+
+  private def partitionSeqByAuthorized[T](requestContext: RequestContext,
+                                        operation: AclOperation,
+                                        resourceType: ResourceType,
+                                        resources: Seq[T],
+                                        logIfAllowed: Boolean = true,
+                                        logIfDenied: Boolean = true)(resourceName: T => String): (Seq[T], Seq[T]) = {
+    authorizer match {
+      case Some(_) =>
+        val authorizedResourceNames = filterByAuthorized(requestContext, operation, resourceType,
+          resources, logIfAllowed, logIfDenied)(resourceName)
+        resources.partition(resource => authorizedResourceNames.contains(resourceName(resource)))
+      case None => (resources, Seq.empty)
+    }
+  }
+
+  private def partitionMapByAuthorized[K, V](requestContext: RequestContext,
+                                             operation: AclOperation,
+                                             resourceType: ResourceType,
+                                             resources: Map[K, V],
+                                             logIfAllowed: Boolean = true,
+                                             logIfDenied: Boolean = true)(resourceName: K => String): (Map[K, V], Map[K, V]) = {
+    authorizer match {
+      case Some(_) =>
+        val authorizedResourceNames = filterByAuthorized(requestContext, operation, resourceType,
+          resources.keySet, logIfAllowed, logIfDenied)(resourceName)
+        resources.partition { case (k, _) => authorizedResourceNames.contains(resourceName(k)) }
+      case None => (resources, Map.empty)
+    }
+  }
+
+  private def authorizeClusterOperation(request: RequestChannel.Request, operation: AclOperation): Unit = {
+    if (!authorize(request.context, operation, CLUSTER, CLUSTER_NAME))
       throw new ClusterAuthorizationException(s"Request $request is not authorized.")
   }
 
-  private def updateRecordsProcessingStats(request: RequestChannel.Request, tp: TopicPartition,
-                                           processingStats: RecordsProcessingStats): Unit = {
-    val conversionCount = processingStats.numRecordsConverted
+  private def authorizedOperations(request: RequestChannel.Request, resource: Resource): Int = {
+    val supportedOps = AclEntry.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,
+                                          tp: TopicPartition,
+                                          conversionStats: RecordConversionStats): Unit = {
+    val conversionCount = conversionStats.numRecordsConverted
     if (conversionCount > 0) {
       request.header.apiKey match {
         case ApiKeys.PRODUCE =>
@@ -2014,73 +3385,147 @@ class KafkaApis(val requestChannel: RequestChannel,
         case _ =>
           throw new IllegalStateException("Message conversion info is recorded only for Produce/Fetch requests")
       }
-      request.messageConversionsTimeNanos = processingStats.conversionTimeNanos
+      request.messageConversionsTimeNanos = conversionStats.conversionTimeNanos
     }
-    request.temporaryMemoryBytes = processingStats.temporaryMemoryBytes
+    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".format(request.body[AbstractRequest]), e)
+    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)
     else
       sendErrorResponseExemptThrottle(request, e)
   }
 
-  private def sendResponseMaybeThrottle(request: RequestChannel.Request, createResponse: Int => AbstractResponse): Unit = {
-    quotas.request.maybeRecordAndThrottle(request,
-      throttleTimeMs => sendResponse(request, Some(createResponse(throttleTimeMs))))
+  private def sendForwardedResponse(
+    request: RequestChannel.Request,
+    response: AbstractResponse
+  ): Unit = {
+    // For forwarded requests, we take the throttle time from the broker that
+    // the request was forwarded to
+    val throttleTimeMs = response.throttleTimeMs()
+    quotas.request.throttle(request, throttleTimeMs, requestChannel.sendResponse)
+    sendResponse(request, Some(response), None)
+  }
+
+  // Throttle the channel if the request quota is enabled but has been violated. Regardless of throttling, send the
+  // response immediately.
+  private def sendResponseMaybeThrottle(request: RequestChannel.Request,
+                                        createResponse: Int => AbstractResponse): Unit = {
+    val throttleTimeMs = maybeRecordAndGetThrottleTimeMs(request)
+    // Only throttle non-forwarded requests
+    if (!request.isForwarded)
+      quotas.request.throttle(request, throttleTimeMs, requestChannel.sendResponse)
+    sendResponse(request, Some(createResponse(throttleTimeMs)), None)
   }
 
-  private def sendErrorResponseMaybeThrottle(request: RequestChannel.Request, error: Throwable) {
-    quotas.request.maybeRecordAndThrottle(request, sendErrorOrCloseConnection(request, error))
+  private def sendErrorResponseMaybeThrottle(request: RequestChannel.Request, error: Throwable): Unit = {
+    val throttleTimeMs = maybeRecordAndGetThrottleTimeMs(request)
+    // Only throttle non-forwarded requests or cluster authorization failures
+    if (error.isInstanceOf[ClusterAuthorizationException] || !request.isForwarded)
+      quotas.request.throttle(request, throttleTimeMs, requestChannel.sendResponse)
+    sendErrorOrCloseConnection(request, error, throttleTimeMs)
+  }
+
+  private def maybeRecordAndGetThrottleTimeMs(request: RequestChannel.Request): Int = {
+    val throttleTimeMs = quotas.request.maybeRecordAndGetThrottleTimeMs(request, time.milliseconds())
+    request.apiThrottleTimeMs = throttleTimeMs
+    throttleTimeMs
+  }
+
+  /**
+   * Throttle the channel if the controller mutations quota or the request quota have been violated.
+   * Regardless of throttling, send the response immediately.
+   */
+  private def sendResponseMaybeThrottleWithControllerQuota(controllerMutationQuota: ControllerMutationQuota,
+                                                           request: RequestChannel.Request,
+                                                           createResponse: Int => AbstractResponse): Unit = {
+    val timeMs = time.milliseconds
+    val controllerThrottleTimeMs = controllerMutationQuota.throttleTime
+    val requestThrottleTimeMs = quotas.request.maybeRecordAndGetThrottleTimeMs(request, timeMs)
+    val maxThrottleTimeMs = Math.max(controllerThrottleTimeMs, requestThrottleTimeMs)
+    // Only throttle non-forwarded requests
+    if (maxThrottleTimeMs > 0 && !request.isForwarded) {
+      request.apiThrottleTimeMs = maxThrottleTimeMs
+      if (controllerThrottleTimeMs > requestThrottleTimeMs) {
+        quotas.controllerMutation.throttle(request, controllerThrottleTimeMs, requestChannel.sendResponse)
+      } else {
+        quotas.request.throttle(request, requestThrottleTimeMs, requestChannel.sendResponse)
+      }
+    }
+
+    sendResponse(request, Some(createResponse(maxThrottleTimeMs)), None)
   }
 
-  private def sendResponseExemptThrottle(request: RequestChannel.Request, response: AbstractResponse): Unit = {
+  private def sendResponseExemptThrottle(request: RequestChannel.Request,
+                                         response: AbstractResponse,
+                                         onComplete: Option[Send => Unit] = None): Unit = {
     quotas.request.maybeRecordExempt(request)
-    sendResponse(request, Some(response))
+    sendResponse(request, Some(response), onComplete)
   }
 
   private def sendErrorResponseExemptThrottle(request: RequestChannel.Request, error: Throwable): Unit = {
     quotas.request.maybeRecordExempt(request)
-    sendErrorOrCloseConnection(request, error)(throttleMs = 0)
+    sendErrorOrCloseConnection(request, error, 0)
   }
 
-  private def sendErrorOrCloseConnection(request: RequestChannel.Request, error: Throwable)(throttleMs: Int): Unit = {
+  private def sendErrorOrCloseConnection(request: RequestChannel.Request, error: Throwable, throttleMs: Int): Unit = {
     val requestBody = request.body[AbstractRequest]
     val response = requestBody.getErrorResponse(throttleMs, error)
     if (response == null)
       closeConnection(request, requestBody.errorCounts(error))
     else
-      sendResponse(request, Some(response))
+      sendResponse(request, Some(response), None)
   }
 
   private def sendNoOpResponseExemptThrottle(request: RequestChannel.Request): Unit = {
     quotas.request.maybeRecordExempt(request)
-    sendResponse(request, None)
+    sendResponse(request, None, None)
   }
 
   private def closeConnection(request: RequestChannel.Request, errorCounts: java.util.Map[Errors, Integer]): Unit = {
     // This case is used when the request handler has encountered an error, but the client
     // does not expect a response (e.g. when produce request has acks set to 0)
     requestChannel.updateErrorMetrics(request.header.apiKey, errorCounts.asScala)
-    requestChannel.sendResponse(new RequestChannel.Response(request, None, CloseConnectionAction, None))
+    requestChannel.sendResponse(new RequestChannel.CloseConnectionResponse(request))
   }
 
-  private def sendResponse(request: RequestChannel.Request, responseOpt: Option[AbstractResponse]): Unit = {
+  private def sendResponse(request: RequestChannel.Request,
+                           responseOpt: Option[AbstractResponse],
+                           onComplete: Option[Send => Unit]): Unit = {
     // Update error metrics for each error code in the response including Errors.NONE
     responseOpt.foreach(response => requestChannel.updateErrorMetrics(request.header.apiKey, response.errorCounts.asScala))
 
-    responseOpt match {
+    val response = responseOpt match {
       case Some(response) =>
-        val responseSend = request.context.buildResponse(response)
-        val responseString =
-          if (RequestChannel.isRequestLoggingEnabled) Some(response.toString(request.context.header.apiVersion))
-          else None
-        requestChannel.sendResponse(new RequestChannel.Response(request, Some(responseSend), SendAction, responseString))
+        new RequestChannel.SendResponse(
+          request,
+          request.buildResponseSend(response),
+          request.responseNode(response),
+          onComplete
+        )
       case None =>
-        requestChannel.sendResponse(new RequestChannel.Response(request, None, NoOpAction, None))
+        new RequestChannel.NoOpResponse(request)
+    }
+
+    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 {
+      // brokerEpochInRequest > controller.brokerEpoch is possible in rare scenarios where the controller gets notified
+      // about the new broker epoch and sends a control request with this epoch before the broker learns about it
+      brokerEpochInRequest < controller.brokerEpoch
     }
   }
 
diff --git a/core/src/main/scala/kafka/server/KafkaConfig.scala b/core/src/main/scala/kafka/server/KafkaConfig.scala
index e24659c72a622..93b840576e09a 100755
--- a/core/src/main/scala/kafka/server/KafkaConfig.scala
+++ b/core/src/main/scala/kafka/server/KafkaConfig.scala
@@ -17,40 +17,52 @@
 
 package kafka.server
 
-import java.util.Properties
+import java.util
+import java.util.{Collections, Locale, Properties}
 
-import kafka.api.{ApiVersion, KAFKA_0_10_0_IV1}
+import kafka.api.{ApiVersion, ApiVersionValidator, KAFKA_0_10_0_IV1, KAFKA_2_1_IV0, KAFKA_2_7_IV0}
 import kafka.cluster.EndPoint
-import kafka.consumer.ConsumerConfig
 import kafka.coordinator.group.OffsetConfig
 import kafka.coordinator.transaction.{TransactionLog, TransactionStateManager}
-import kafka.message.{BrokerCompressionCodec, CompressionCodec, Message, MessageSet}
+import kafka.message.{BrokerCompressionCodec, CompressionCodec, ZStdCompressionCodec}
+import kafka.security.authorizer.AuthorizerUtils
 import kafka.utils.CoreUtils
 import kafka.utils.Implicits._
 import org.apache.kafka.clients.CommonClientConfigs
-import org.apache.kafka.common.config.ConfigDef.ValidList
+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.types.Password
+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.TimestampType
+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 org.apache.zookeeper.client.ZKClientConfig
 
-import scala.collection.JavaConverters._
-import scala.collection.Map
+import scala.jdk.CollectionConverters._
+import scala.collection.{Map, Seq}
 
 object Defaults {
   /** ********* Zookeeper Configuration ***********/
-  val ZkSessionTimeoutMs = 6000
+  val ZkSessionTimeoutMs = 18000
   val ZkSyncTimeMs = 2000
   val ZkEnableSecureAcls = false
   val ZkMaxInFlightRequests = 10
+  val ZkSslClientEnable = false
+  val ZkSslProtocol = "TLSv1.2"
+  val ZkSslEndpointIdentificationAlgorithm = "HTTPS"
+  val ZkSslCrlEnable = false
+  val ZkSslOcspEnable = false
 
   /** ********* General Configuration ***********/
   val BrokerIdGenerationEnable = true
   val MaxReservedBrokerId = 1000
   val BrokerId = -1
-  val MessageMaxBytes = 1000000 + MessageSet.LogOverhead
+  val MessageMaxBytes = 1024 * 1024 + Records.LOG_OVERHEAD
   val NumNetworkThreads = 3
   val NumIoThreads = 8
   val BackgroundThreads = 10
@@ -73,8 +85,13 @@ object Defaults {
   val SocketRequestMaxBytes: Int = 100 * 1024 * 1024
   val MaxConnectionsPerIp: Int = Int.MaxValue
   val MaxConnectionsPerIpOverrides: String = ""
+  val MaxConnections: Int = Int.MaxValue
+  val MaxConnectionCreationRate: Int = Int.MaxValue
   val ConnectionsMaxIdleMs = 10 * 60 * 1000L
   val RequestTimeoutMs = 30000
+  val ConnectionSetupTimeoutMs = CommonClientConfigs.DEFAULT_SOCKET_CONNECTION_SETUP_TIMEOUT_MS
+  val ConnectionSetupTimeoutMaxMs = CommonClientConfigs.DEFAULT_SOCKET_CONNECTION_SETUP_TIMEOUT_MAX_MS
+  val FailedAuthenticationDelayMs = 100
 
   /** ********* Log Configuration ***********/
   val NumPartitions = 1
@@ -99,6 +116,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
@@ -114,15 +132,16 @@ object Defaults {
   val NumRecoveryThreadsPerDataDir = 1
   val AutoCreateTopicsEnable = true
   val MinInSyncReplicas = 1
+  val MessageDownConversionEnable = true
 
   /** ********* Replication configuration ***********/
   val ControllerSocketTimeoutMs = RequestTimeoutMs
   val ControllerMessageQueueSize = Int.MaxValue
   val DefaultReplicationFactor = 1
-  val ReplicaLagTimeMaxMs = 10000L
-  val ReplicaSocketTimeoutMs = ConsumerConfig.SocketTimeout
-  val ReplicaSocketReceiveBufferBytes = ConsumerConfig.SocketBufferSize
-  val ReplicaFetchMaxBytes = ConsumerConfig.FetchSize
+  val ReplicaLagTimeMaxMs = 30000L
+  val ReplicaSocketTimeoutMs = 30 * 1000
+  val ReplicaSocketReceiveBufferBytes = 64 * 1024
+  val ReplicaFetchMaxBytes = 1024 * 1024
   val ReplicaFetchWaitMaxMs = 500
   val ReplicaFetchMinBytes = 1
   val ReplicaFetchResponseMaxBytes = 10 * 1024 * 1024
@@ -146,8 +165,9 @@ object Defaults {
 
   /** ********* Group coordinator configuration ***********/
   val GroupMinSessionTimeoutMs = 6000
-  val GroupMaxSessionTimeoutMs = 300000
+  val GroupMaxSessionTimeoutMs = 1800000
   val GroupInitialRebalanceDelayMs = 3000
+  val GroupMaxSize: Int = Int.MaxValue
 
   /** ********* Offset management configuration ***********/
   val OffsetMetadataMaxSize = OffsetConfig.DefaultMaxMetadataSize
@@ -156,7 +176,7 @@ object Defaults {
   val OffsetsTopicPartitions: Int = OffsetConfig.DefaultOffsetsTopicNumPartitions
   val OffsetsTopicSegmentBytes: Int = OffsetConfig.DefaultOffsetsTopicSegmentBytes
   val OffsetsTopicCompressionCodec: Int = OffsetConfig.DefaultOffsetsTopicCompressionCodec.codec
-  val OffsetsRetentionMinutes: Int = 24 * 60
+  val OffsetsRetentionMinutes: Int = 7 * 24 * 60
   val OffsetsRetentionCheckIntervalMs: Long = OffsetConfig.DefaultOffsetsRetentionCheckIntervalMs
   val OffsetCommitTimeoutMs = OffsetConfig.DefaultOffsetCommitTimeoutMs
   val OffsetCommitRequiredAcks = OffsetConfig.DefaultOffsetCommitRequiredAcks
@@ -172,15 +192,21 @@ object Defaults {
   val TransactionsAbortTimedOutTransactionsCleanupIntervalMS = TransactionStateManager.DefaultAbortTimedOutTransactionsIntervalMs
   val TransactionsRemoveExpiredTransactionsCleanupIntervalMS = TransactionStateManager.DefaultRemoveExpiredTransactionalIdsIntervalMs
 
+  /** ********* Fetch Configuration **************/
+  val MaxIncrementalFetchSessionCacheSlots = 1000
+  val FetchMaxBytes = 55 * 1024 * 1024
+
   /** ********* Quota Configuration ***********/
-  val ProducerQuotaBytesPerSecondDefault = ClientQuotaManagerConfig.QuotaBytesPerSecondDefault
-  val ConsumerQuotaBytesPerSecondDefault = ClientQuotaManagerConfig.QuotaBytesPerSecondDefault
+  val ProducerQuotaBytesPerSecondDefault = ClientQuotaManagerConfig.QuotaDefault
+  val ConsumerQuotaBytesPerSecondDefault = ClientQuotaManagerConfig.QuotaDefault
   val NumQuotaSamples: Int = ClientQuotaManagerConfig.DefaultNumQuotaSamples
   val QuotaWindowSizeSeconds: Int = ClientQuotaManagerConfig.DefaultQuotaWindowSizeSeconds
   val NumReplicationQuotaSamples: Int = ReplicationQuotaManagerConfig.DefaultNumQuotaSamples
   val ReplicationQuotaWindowSizeSeconds: Int = ReplicationQuotaManagerConfig.DefaultQuotaWindowSizeSeconds
   val NumAlterLogDirsReplicationQuotaSamples: Int = ReplicationQuotaManagerConfig.DefaultNumQuotaSamples
   val AlterLogDirsReplicationQuotaWindowSizeSeconds: Int = ReplicationQuotaManagerConfig.DefaultQuotaWindowSizeSeconds
+  val NumControllerQuotaSamples: Int = ClientQuotaManagerConfig.DefaultNumQuotaSamples
+  val ControllerQuotaWindowSizeSeconds: Int = ClientQuotaManagerConfig.DefaultQuotaWindowSizeSeconds
 
   /** ********* Transaction Configuration ***********/
   val TransactionalIdExpirationMsDefault = 604800000
@@ -196,6 +222,11 @@ object Defaults {
   val MetricReporterClasses = ""
   val MetricRecordingLevel = Sensor.RecordingLevel.INFO.toString()
 
+
+  /** ********* Kafka Yammer Metrics Reporter Configuration ***********/
+  val KafkaMetricReporterClasses = ""
+  val KafkaMetricsPollingIntervalSeconds = 10
+
   /** ********* SSL configuration ***********/
   val SslProtocol = SslConfigs.DEFAULT_SSL_PROTOCOL
   val SslEnabledProtocols = SslConfigs.DEFAULT_SSL_ENABLED_PROTOCOLS
@@ -203,27 +234,45 @@ object Defaults {
   val SslTruststoreType = SslConfigs.DEFAULT_SSL_TRUSTSTORE_TYPE
   val SslKeyManagerAlgorithm = SslConfigs.DEFAULT_SSL_KEYMANGER_ALGORITHM
   val SslTrustManagerAlgorithm = SslConfigs.DEFAULT_SSL_TRUSTMANAGER_ALGORITHM
-  val SslClientAuthRequired = "required"
-  val SslClientAuthRequested = "requested"
-  val SslClientAuthNone = "none"
-  val SslClientAuth = SslClientAuthNone
+  val SslEndpointIdentificationAlgorithm = SslConfigs.DEFAULT_SSL_ENDPOINT_IDENTIFICATION_ALGORITHM
+  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 ***********/
+  val ConnectionsMaxReauthMsDefault = 0L
 
   /** ********* Sasl configuration ***********/
   val SaslMechanismInterBrokerProtocol = SaslConfigs.DEFAULT_SASL_MECHANISM
-  val SaslEnabledMechanisms = SaslConfigs.DEFAULT_SASL_ENABLED_MECHANISMS
+  val SaslEnabledMechanisms = BrokerSecurityConfigs.DEFAULT_SASL_ENABLED_MECHANISMS
   val SaslKerberosKinitCmd = SaslConfigs.DEFAULT_KERBEROS_KINIT_CMD
   val SaslKerberosTicketRenewWindowFactor = SaslConfigs.DEFAULT_KERBEROS_TICKET_RENEW_WINDOW_FACTOR
   val SaslKerberosTicketRenewJitter = SaslConfigs.DEFAULT_KERBEROS_TICKET_RENEW_JITTER
   val SaslKerberosMinTimeBeforeRelogin = SaslConfigs.DEFAULT_KERBEROS_MIN_TIME_BEFORE_RELOGIN
   val SaslKerberosPrincipalToLocalRules = BrokerSecurityConfigs.DEFAULT_SASL_KERBEROS_PRINCIPAL_TO_LOCAL_RULES
+  val SaslLoginRefreshWindowFactor = SaslConfigs.DEFAULT_LOGIN_REFRESH_WINDOW_FACTOR
+  val SaslLoginRefreshWindowJitter = SaslConfigs.DEFAULT_LOGIN_REFRESH_WINDOW_JITTER
+  val SaslLoginRefreshMinPeriodSeconds = SaslConfigs.DEFAULT_LOGIN_REFRESH_MIN_PERIOD_SECONDS
+  val SaslLoginRefreshBufferSeconds = SaslConfigs.DEFAULT_LOGIN_REFRESH_BUFFER_SECONDS
+
+  /** ********* Delegation Token configuration ***********/
+  val DelegationTokenMaxLifeTimeMsDefault = 7 * 24 * 60 * 60 * 1000L
+  val DelegationTokenExpiryTimeMsDefault = 24 * 60 * 60 * 1000L
+  val DelegationTokenExpiryCheckIntervalMsDefault = 1 * 60 * 60 * 1000L
+
+  /** ********* Password encryption configuration for dynamic configs *********/
+  val PasswordEncoderCipherAlgorithm = "AES/CBC/PKCS5Padding"
+  val PasswordEncoderKeyLength = 128
+  val PasswordEncoderIterations = 4096
 }
 
 object KafkaConfig {
 
   private val LogConfigPrefix = "log."
 
-  def main(args: Array[String]) {
-    System.out.println(configDef.toHtmlTable)
+  def main(args: Array[String]): Unit = {
+    System.out.println(configDef.toHtml(4, (config: String) => "brokerconfigs_" + config,
+      DynamicBrokerConfig.dynamicConfigUpdateModes))
   }
 
   /** ********* Zookeeper Configuration ***********/
@@ -233,6 +282,62 @@ object KafkaConfig {
   val ZkSyncTimeMsProp = "zookeeper.sync.time.ms"
   val ZkEnableSecureAclsProp = "zookeeper.set.acl"
   val ZkMaxInFlightRequestsProp = "zookeeper.max.in.flight.requests"
+  val ZkSslClientEnableProp = "zookeeper.ssl.client.enable"
+  val ZkClientCnxnSocketProp = "zookeeper.clientCnxnSocket"
+  val ZkSslKeyStoreLocationProp = "zookeeper.ssl.keystore.location"
+  val ZkSslKeyStorePasswordProp = "zookeeper.ssl.keystore.password"
+  val ZkSslKeyStoreTypeProp = "zookeeper.ssl.keystore.type"
+  val ZkSslTrustStoreLocationProp = "zookeeper.ssl.truststore.location"
+  val ZkSslTrustStorePasswordProp = "zookeeper.ssl.truststore.password"
+  val ZkSslTrustStoreTypeProp = "zookeeper.ssl.truststore.type"
+  val ZkSslProtocolProp = "zookeeper.ssl.protocol"
+  val ZkSslEnabledProtocolsProp = "zookeeper.ssl.enabled.protocols"
+  val ZkSslCipherSuitesProp = "zookeeper.ssl.cipher.suites"
+  val ZkSslEndpointIdentificationAlgorithmProp = "zookeeper.ssl.endpoint.identification.algorithm"
+  val ZkSslCrlEnableProp = "zookeeper.ssl.crl.enable"
+  val ZkSslOcspEnableProp = "zookeeper.ssl.ocsp.enable"
+
+  // a map from the Kafka config to the corresponding ZooKeeper Java system property
+  private[kafka] val ZkSslConfigToSystemPropertyMap: Map[String, String] = Map(
+    ZkSslClientEnableProp -> ZKClientConfig.SECURE_CLIENT,
+    ZkClientCnxnSocketProp -> ZKClientConfig.ZOOKEEPER_CLIENT_CNXN_SOCKET,
+    ZkSslKeyStoreLocationProp -> "zookeeper.ssl.keyStore.location",
+    ZkSslKeyStorePasswordProp -> "zookeeper.ssl.keyStore.password",
+    ZkSslKeyStoreTypeProp -> "zookeeper.ssl.keyStore.type",
+    ZkSslTrustStoreLocationProp -> "zookeeper.ssl.trustStore.location",
+    ZkSslTrustStorePasswordProp -> "zookeeper.ssl.trustStore.password",
+    ZkSslTrustStoreTypeProp -> "zookeeper.ssl.trustStore.type",
+    ZkSslProtocolProp -> "zookeeper.ssl.protocol",
+    ZkSslEnabledProtocolsProp -> "zookeeper.ssl.enabledProtocols",
+    ZkSslCipherSuitesProp -> "zookeeper.ssl.ciphersuites",
+    ZkSslEndpointIdentificationAlgorithmProp -> "zookeeper.ssl.hostnameVerification",
+    ZkSslCrlEnableProp -> "zookeeper.ssl.crl",
+    ZkSslOcspEnableProp -> "zookeeper.ssl.ocsp")
+
+  private[kafka] def getZooKeeperClientProperty(clientConfig: ZKClientConfig, kafkaPropName: String): Option[String] = {
+    Option(clientConfig.getProperty(ZkSslConfigToSystemPropertyMap(kafkaPropName)))
+  }
+
+  private[kafka] def setZooKeeperClientProperty(clientConfig: ZKClientConfig, kafkaPropName: String, kafkaPropValue: Any): Unit = {
+    clientConfig.setProperty(ZkSslConfigToSystemPropertyMap(kafkaPropName),
+      kafkaPropName match {
+        case ZkSslEndpointIdentificationAlgorithmProp => (kafkaPropValue.toString.toUpperCase == "HTTPS").toString
+        case ZkSslEnabledProtocolsProp | ZkSslCipherSuitesProp => kafkaPropValue match {
+          case list: java.util.List[_] => list.asInstanceOf[java.util.List[_]].asScala.mkString(",")
+          case _ => kafkaPropValue.toString
+        }
+        case _ => kafkaPropValue.toString
+    })
+  }
+
+  // For ZooKeeper TLS client authentication to be enabled the client must (at a minimum) configure itself as using TLS
+  // with both a client connection socket and a key store location explicitly set.
+  private[kafka] def zkTlsClientAuthEnabled(zkClientConfig: ZKClientConfig) = {
+    getZooKeeperClientProperty(zkClientConfig, ZkSslClientEnableProp).getOrElse("false") == "true" &&
+      getZooKeeperClientProperty(zkClientConfig, ZkClientCnxnSocketProp).isDefined &&
+      getZooKeeperClientProperty(zkClientConfig, ZkSslKeyStoreLocationProp).isDefined
+  }
+
   /** ********* General Configuration ***********/
   val BrokerIdGenerationEnableProp = "broker.id.generation.enable"
   val MaxReservedBrokerIdProp = "reserved.broker.max.id"
@@ -245,6 +350,10 @@ object KafkaConfig {
   val QueuedMaxRequestsProp = "queued.max.requests"
   val QueuedMaxBytesProp = "queued.max.request.bytes"
   val RequestTimeoutMsProp = CommonClientConfigs.REQUEST_TIMEOUT_MS_CONFIG
+  val ConnectionSetupTimeoutMsProp = CommonClientConfigs.SOCKET_CONNECTION_SETUP_TIMEOUT_MS_CONFIG
+  val ConnectionSetupTimeoutMaxMsProp = CommonClientConfigs.SOCKET_CONNECTION_SETUP_TIMEOUT_MAX_MS_CONFIG
+  val EnableMetadataQuorumProp = "enable.metadata.quorum"
+
   /************* Authorizer Configuration ***********/
   val AuthorizerClassNameProp = "authorizer.class.name"
   /** ********* Socket Server Configuration ***********/
@@ -255,12 +364,16 @@ object KafkaConfig {
   val AdvertisedPortProp = "advertised.port"
   val AdvertisedListenersProp = "advertised.listeners"
   val ListenerSecurityProtocolMapProp = "listener.security.protocol.map"
+  val ControlPlaneListenerNameProp = "control.plane.listener.name"
   val SocketSendBufferBytesProp = "socket.send.buffer.bytes"
   val SocketReceiveBufferBytesProp = "socket.receive.buffer.bytes"
   val SocketRequestMaxBytesProp = "socket.request.max.bytes"
   val MaxConnectionsPerIpProp = "max.connections.per.ip"
   val MaxConnectionsPerIpOverridesProp = "max.connections.per.ip.overrides"
+  val MaxConnectionsProp = "max.connections"
+  val MaxConnectionCreationRateProp = "max.connection.creation.rate"
   val ConnectionsMaxIdleMsProp = "connections.max.idle.ms"
+  val FailedAuthenticationDelayMsProp = "connection.failed.authentication.delay.ms"
   /***************** rack configuration *************/
   val RackProp = "broker.rack"
   /** ********* Log Configuration ***********/
@@ -292,6 +405,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"
@@ -310,6 +424,7 @@ object KafkaConfig {
   val MinInSyncReplicasProp = "min.insync.replicas"
   val CreateTopicPolicyClassNameProp = "create.topic.policy.class.name"
   val AlterConfigPolicyClassNameProp = "alter.config.policy.class.name"
+  val LogMessageDownConversionEnableProp = LogConfigPrefix + "message.downconversion.enable"
   /** ********* Replication configuration ***********/
   val ControllerSocketTimeoutMsProp = "controller.socket.timeout.ms"
   val DefaultReplicationFactorProp = "default.replication.factor"
@@ -333,6 +448,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"
@@ -341,6 +457,7 @@ object KafkaConfig {
   val GroupMinSessionTimeoutMsProp = "group.min.session.timeout.ms"
   val GroupMaxSessionTimeoutMsProp = "group.max.session.timeout.ms"
   val GroupInitialRebalanceDelayMsProp = "group.initial.rebalance.delay.ms"
+  val GroupMaxSizeProp = "group.max.size"
   /** ********* Offset management configuration ***********/
   val OffsetMetadataMaxSizeProp = "offset.metadata.max.bytes"
   val OffsetsLoadBufferSizeProp = "offsets.load.buffer.size"
@@ -363,15 +480,22 @@ object KafkaConfig {
   val TransactionsAbortTimedOutTransactionCleanupIntervalMsProp = "transaction.abort.timed.out.transaction.cleanup.interval.ms"
   val TransactionsRemoveExpiredTransactionalIdCleanupIntervalMsProp = "transaction.remove.expired.transaction.cleanup.interval.ms"
 
+  /** ********* Fetch Configuration **************/
+  val MaxIncrementalFetchSessionCacheSlots = "max.incremental.fetch.session.cache.slots"
+  val FetchMaxBytes = "fetch.max.bytes"
+
   /** ********* Quota Configuration ***********/
   val ProducerQuotaBytesPerSecondDefaultProp = "quota.producer.default"
   val ConsumerQuotaBytesPerSecondDefaultProp = "quota.consumer.default"
   val NumQuotaSamplesProp = "quota.window.num"
   val NumReplicationQuotaSamplesProp = "replication.quota.window.num"
   val NumAlterLogDirsReplicationQuotaSamplesProp = "alter.log.dirs.replication.quota.window.num"
+  val NumControllerQuotaSamplesProp = "controller.quota.window.num"
   val QuotaWindowSizeSecondsProp = "quota.window.size.seconds"
   val ReplicationQuotaWindowSizeSecondsProp = "replication.quota.window.size.seconds"
   val AlterLogDirsReplicationQuotaWindowSizeSecondsProp = "alter.log.dirs.replication.quota.window.size.seconds"
+  val ControllerQuotaWindowSizeSecondsProp = "controller.quota.window.size.seconds"
+  val ClientQuotaCallbackClassProp = "client.quota.callback.class"
 
   val DeleteTopicEnableProp = "delete.topic.enable"
   val CompressionTypeProp = "compression.type"
@@ -382,8 +506,14 @@ object KafkaConfig {
   val MetricReporterClassesProp: String = CommonClientConfigs.METRIC_REPORTER_CLASSES_CONFIG
   val MetricRecordingLevelProp: String = CommonClientConfigs.METRICS_RECORDING_LEVEL_CONFIG
 
+  /** ********* Kafka Yammer Metrics Reporters Configuration ***********/
+  val KafkaMetricsReporterClassesProp = "kafka.metrics.reporters"
+  val KafkaMetricsPollingIntervalSecondsProp = "kafka.metrics.polling.interval.secs"
+
   /** ******** 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
@@ -394,33 +524,101 @@ object KafkaConfig {
   val SslKeystoreLocationProp = SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG
   val SslKeystorePasswordProp = SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG
   val SslKeyPasswordProp = SslConfigs.SSL_KEY_PASSWORD_CONFIG
+  val SslKeystoreKeyProp = SslConfigs.SSL_KEYSTORE_KEY_CONFIG
+  val SslKeystoreCertificateChainProp = SslConfigs.SSL_KEYSTORE_CERTIFICATE_CHAIN_CONFIG
   val SslTruststoreTypeProp = SslConfigs.SSL_TRUSTSTORE_TYPE_CONFIG
   val SslTruststoreLocationProp = SslConfigs.SSL_TRUSTSTORE_LOCATION_CONFIG
   val SslTruststorePasswordProp = SslConfigs.SSL_TRUSTSTORE_PASSWORD_CONFIG
+  val SslTruststoreCertificatesProp = SslConfigs.SSL_TRUSTSTORE_CERTIFICATES_CONFIG
   val SslKeyManagerAlgorithmProp = SslConfigs.SSL_KEYMANAGER_ALGORITHM_CONFIG
   val SslTrustManagerAlgorithmProp = SslConfigs.SSL_TRUSTMANAGER_ALGORITHM_CONFIG
   val SslEndpointIdentificationAlgorithmProp = SslConfigs.SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG
   val SslSecureRandomImplementationProp = SslConfigs.SSL_SECURE_RANDOM_IMPLEMENTATION_CONFIG
   val SslClientAuthProp = BrokerSecurityConfigs.SSL_CLIENT_AUTH_CONFIG
+  val SslPrincipalMappingRulesProp = BrokerSecurityConfigs.SSL_PRINCIPAL_MAPPING_RULES_CONFIG
+  var SslEngineFactoryClassProp = SslConfigs.SSL_ENGINE_FACTORY_CLASS_CONFIG
 
   /** ********* SASL Configuration ****************/
   val SaslMechanismInterBrokerProtocolProp = "sasl.mechanism.inter.broker.protocol"
+  val SaslJaasConfigProp = SaslConfigs.SASL_JAAS_CONFIG
   val SaslEnabledMechanismsProp = BrokerSecurityConfigs.SASL_ENABLED_MECHANISMS_CONFIG
+  val SaslServerCallbackHandlerClassProp = BrokerSecurityConfigs.SASL_SERVER_CALLBACK_HANDLER_CLASS
+  val SaslClientCallbackHandlerClassProp = SaslConfigs.SASL_CLIENT_CALLBACK_HANDLER_CLASS
+  val SaslLoginClassProp = SaslConfigs.SASL_LOGIN_CLASS
+  val SaslLoginCallbackHandlerClassProp = SaslConfigs.SASL_LOGIN_CALLBACK_HANDLER_CLASS
   val SaslKerberosServiceNameProp = SaslConfigs.SASL_KERBEROS_SERVICE_NAME
   val SaslKerberosKinitCmdProp = SaslConfigs.SASL_KERBEROS_KINIT_CMD
   val SaslKerberosTicketRenewWindowFactorProp = SaslConfigs.SASL_KERBEROS_TICKET_RENEW_WINDOW_FACTOR
   val SaslKerberosTicketRenewJitterProp = SaslConfigs.SASL_KERBEROS_TICKET_RENEW_JITTER
   val SaslKerberosMinTimeBeforeReloginProp = SaslConfigs.SASL_KERBEROS_MIN_TIME_BEFORE_RELOGIN
   val SaslKerberosPrincipalToLocalRulesProp = BrokerSecurityConfigs.SASL_KERBEROS_PRINCIPAL_TO_LOCAL_RULES_CONFIG
+  val SaslLoginRefreshWindowFactorProp = SaslConfigs.SASL_LOGIN_REFRESH_WINDOW_FACTOR
+  val SaslLoginRefreshWindowJitterProp = SaslConfigs.SASL_LOGIN_REFRESH_WINDOW_JITTER
+  val SaslLoginRefreshMinPeriodSecondsProp = SaslConfigs.SASL_LOGIN_REFRESH_MIN_PERIOD_SECONDS
+  val SaslLoginRefreshBufferSecondsProp = SaslConfigs.SASL_LOGIN_REFRESH_BUFFER_SECONDS
+
+  /** ********* Delegation Token Configuration ****************/
+  val DelegationTokenSecretKeyAliasProp = "delegation.token.master.key"
+  val DelegationTokenSecretKeyProp = "delegation.token.secret.key"
+  val DelegationTokenMaxLifeTimeProp = "delegation.token.max.lifetime.ms"
+  val DelegationTokenExpiryTimeMsProp = "delegation.token.expiry.time.ms"
+  val DelegationTokenExpiryCheckIntervalMsProp = "delegation.token.expiry.check.interval.ms"
+
+  /** ********* Password encryption configuration for dynamic configs *********/
+  val PasswordEncoderSecretProp = "password.encoder.secret"
+  val PasswordEncoderOldSecretProp = "password.encoder.old.secret"
+  val PasswordEncoderKeyFactoryAlgorithmProp = "password.encoder.keyfactory.algorithm"
+  val PasswordEncoderCipherAlgorithmProp = "password.encoder.cipher.algorithm"
+  val PasswordEncoderKeyLengthProp =  "password.encoder.key.length"
+  val PasswordEncoderIterationsProp =  "password.encoder.iterations"
 
   /* Documentation */
   /** ********* Zookeeper Configuration ***********/
-  val ZkConnectDoc = "Zookeeper host string"
+  val ZkConnectDoc = "Specifies the ZooKeeper connection string in the form hostname:port where host and port are the " +
+  "host and port of a ZooKeeper server. To allow connecting through other ZooKeeper nodes when that ZooKeeper machine is " +
+  "down you can also specify multiple hosts in the form hostname1:port1,hostname2:port2,hostname3:port3.\n" +
+  "The server can also have a ZooKeeper chroot path as part of its ZooKeeper connection string which puts its data under some path in the global ZooKeeper namespace. " +
+  "For example to give a chroot path of /chroot/path you would give the connection string as hostname1:port1,hostname2:port2,hostname3:port3/chroot/path."
   val ZkSessionTimeoutMsDoc = "Zookeeper session timeout"
   val ZkConnectionTimeoutMsDoc = "The max time that the client waits to establish a connection to zookeeper. If not set, the value in " + ZkSessionTimeoutMsProp + " is used"
   val ZkSyncTimeMsDoc = "How far a ZK follower can be behind a ZK leader"
   val ZkEnableSecureAclsDoc = "Set client to use secure ACLs"
   val ZkMaxInFlightRequestsDoc = "The maximum number of unacknowledged requests the client will send to Zookeeper before blocking."
+  val ZkSslClientEnableDoc = "Set client to use TLS when connecting to ZooKeeper." +
+    " An explicit value overrides any value set via the zookeeper.client.secure system property (note the different name)." +
+    s" Defaults to false if neither is set; when true, $ZkClientCnxnSocketProp must be set (typically to org.apache.zookeeper.ClientCnxnSocketNetty); other values to set may include " +
+    ZkSslConfigToSystemPropertyMap.keys.toList.sorted.filter(x => x != ZkSslClientEnableProp && x != ZkClientCnxnSocketProp).mkString("", ", ", "")
+  val ZkClientCnxnSocketDoc = "Typically set to org.apache.zookeeper.ClientCnxnSocketNetty when using TLS connectivity to ZooKeeper." +
+    s" Overrides any explicit value set via the same-named ${ZkSslConfigToSystemPropertyMap(ZkClientCnxnSocketProp)} system property."
+  val ZkSslKeyStoreLocationDoc = "Keystore location when using a client-side certificate with TLS connectivity to ZooKeeper." +
+    s" Overrides any explicit value set via the ${ZkSslConfigToSystemPropertyMap(ZkSslKeyStoreLocationProp)} system property (note the camelCase)."
+  val ZkSslKeyStorePasswordDoc = "Keystore password when using a client-side certificate with TLS connectivity to ZooKeeper." +
+    s" Overrides any explicit value set via the ${ZkSslConfigToSystemPropertyMap(ZkSslKeyStorePasswordProp)} system property (note the camelCase)." +
+    " Note that ZooKeeper does not support a key password different from the keystore password, so be sure to set the key password in the keystore to be identical to the keystore password; otherwise the connection attempt to Zookeeper will fail."
+  val ZkSslKeyStoreTypeDoc = "Keystore type when using a client-side certificate with TLS connectivity to ZooKeeper." +
+    s" Overrides any explicit value set via the ${ZkSslConfigToSystemPropertyMap(ZkSslKeyStoreTypeProp)} system property (note the camelCase)." +
+    " The default value of null means the type will be auto-detected based on the filename extension of the keystore."
+  val ZkSslTrustStoreLocationDoc = "Truststore location when using TLS connectivity to ZooKeeper." +
+    s" Overrides any explicit value set via the ${ZkSslConfigToSystemPropertyMap(ZkSslTrustStoreLocationProp)} system property (note the camelCase)."
+  val ZkSslTrustStorePasswordDoc = "Truststore password when using TLS connectivity to ZooKeeper." +
+    s" Overrides any explicit value set via the ${ZkSslConfigToSystemPropertyMap(ZkSslTrustStorePasswordProp)} system property (note the camelCase)."
+  val ZkSslTrustStoreTypeDoc = "Truststore type when using TLS connectivity to ZooKeeper." +
+    s" Overrides any explicit value set via the ${ZkSslConfigToSystemPropertyMap(ZkSslTrustStoreTypeProp)} system property (note the camelCase)." +
+    " The default value of null means the type will be auto-detected based on the filename extension of the truststore."
+  val ZkSslProtocolDoc = "Specifies the protocol to be used in ZooKeeper TLS negotiation." +
+    s" An explicit value overrides any value set via the same-named ${ZkSslConfigToSystemPropertyMap(ZkSslProtocolProp)} system property."
+  val ZkSslEnabledProtocolsDoc = "Specifies the enabled protocol(s) in ZooKeeper TLS negotiation (csv)." +
+    s" Overrides any explicit value set via the ${ZkSslConfigToSystemPropertyMap(ZkSslEnabledProtocolsProp)} system property (note the camelCase)." +
+    s" The default value of null means the enabled protocol will be the value of the ${KafkaConfig.ZkSslProtocolProp} configuration property."
+  val ZkSslCipherSuitesDoc = "Specifies the enabled cipher suites to be used in ZooKeeper TLS negotiation (csv)." +
+    s""" Overrides any explicit value set via the ${ZkSslConfigToSystemPropertyMap(ZkSslCipherSuitesProp)} system property (note the single word \"ciphersuites\").""" +
+    " The default value of null means the list of enabled cipher suites is determined by the Java runtime being used."
+  val ZkSslEndpointIdentificationAlgorithmDoc = "Specifies whether to enable hostname verification in the ZooKeeper TLS negotiation process, with (case-insensitively) \"https\" meaning ZooKeeper hostname verification is enabled and an explicit blank value meaning it is disabled (disabling it is only recommended for testing purposes)." +
+    s""" An explicit value overrides any \"true\" or \"false\" value set via the ${ZkSslConfigToSystemPropertyMap(ZkSslEndpointIdentificationAlgorithmProp)} system property (note the different name and values; true implies https and false implies blank)."""
+  val ZkSslCrlEnableDoc = "Specifies whether to enable Certificate Revocation List in the ZooKeeper TLS protocols." +
+    s" Overrides any explicit value set via the ${ZkSslConfigToSystemPropertyMap(ZkSslCrlEnableProp)} system property (note the shorter name)."
+  val ZkSslOcspEnableDoc = "Specifies whether to enable Online Certificate Status Protocol in the ZooKeeper TLS protocols." +
+    s" Overrides any explicit value set via the ${ZkSslConfigToSystemPropertyMap(ZkSslOcspEnableProp)} system property (note the shorter name)."
   /** ********* General Configuration ***********/
   val BrokerIdGenerationEnableDoc = s"Enable automatic broker id generation on the server. When enabled the value configured for $MaxReservedBrokerIdProp should be reviewed."
   val MaxReservedBrokerIdDoc = "Max number that can be used for a broker.id"
@@ -428,45 +626,53 @@ 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"
   val BackgroundThreadsDoc = "The number of threads to use for various background processing tasks"
-  val QueuedMaxRequestsDoc = "The number of queued requests allowed before blocking the network threads"
+  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 ConnectionSetupTimeoutMsDoc = CommonClientConfigs.SOCKET_CONNECTION_SETUP_TIMEOUT_MS_DOC
+  val ConnectionSetupTimeoutMaxMsDoc = CommonClientConfigs.SOCKET_CONNECTION_SETUP_TIMEOUT_MAX_MS_DOC
   /************* 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" +
+  val PortDoc = "DEPRECATED: only used when listeners is not set. " +
+  "Use listeners instead. \n" +
   "the port to listen and accept connections on"
-  val HostNameDoc = "DEPRECATED: only used when `listeners` is not set. " +
-  "Use `listeners` instead. \n" +
+  val HostNameDoc = "DEPRECATED: only used when listeners is not set. " +
+  "Use listeners instead. \n" +
   "hostname of broker. If this is set, it will only bind to this address. If this is not set, it will bind to all interfaces"
   val ListenersDoc = "Listener List - Comma-separated list of URIs we will listen on and the listener names." +
-  s" If the listener name is not a security protocol, $ListenerSecurityProtocolMapProp must also be set.\n" +
-  " Specify hostname as 0.0.0.0 to bind to all interfaces.\n" +
-  " Leave hostname empty to bind to default interface.\n" +
-  " Examples of legal listener lists:\n" +
-  " PLAINTEXT://myhost:9092,SSL://:9091\n" +
-  " CLIENT://0.0.0.0:9092,REPLICATION://localhost:9093\n"
-  val AdvertisedHostNameDoc = "DEPRECATED: only used when `advertised.listeners` or `listeners` are not set. " +
-  "Use `advertised.listeners` instead. \n" +
+    s" If the listener name is not a security protocol, $ListenerSecurityProtocolMapProp must also be set.\n" +
+    " Listener names and port numbers must be unique.\n" +
+    " Specify hostname as 0.0.0.0 to bind to all interfaces.\n" +
+    " Leave hostname empty to bind to default interface.\n" +
+    " Examples of legal listener lists:\n" +
+    " PLAINTEXT://myhost:9092,SSL://:9091\n" +
+    " CLIENT://0.0.0.0:9092,REPLICATION://localhost:9093\n"
+  val AdvertisedHostNameDoc = "DEPRECATED: only used when advertised.listeners or listeners are not set. " +
+  "Use advertised.listeners instead. \n" +
   "Hostname to publish to ZooKeeper for clients to use. In IaaS environments, this may " +
   "need to be different from the interface to which the broker binds. If this is not set, " +
-  "it will use the value for `host.name` if configured. Otherwise " +
+  "it will use the value for host.name if configured. Otherwise " +
   "it will use the value returned from java.net.InetAddress.getCanonicalHostName()."
-  val AdvertisedPortDoc = "DEPRECATED: only used when `advertised.listeners` or `listeners` are not set. " +
-  "Use `advertised.listeners` instead. \n" +
+  val AdvertisedPortDoc = "DEPRECATED: only used when advertised.listeners or listeners are not set. " +
+  "Use advertised.listeners instead. \n" +
   "The port to publish to ZooKeeper for clients to use. In IaaS environments, this may " +
   "need to be different from the port to which the broker binds. If this is not set, " +
   "it will publish the same port that the broker binds to."
-  val AdvertisedListenersDoc = "Listeners to publish to ZooKeeper for clients to use, if different than the `listeners` config property." +
-  " In IaaS environments, this may need to be different from the interface to which the broker binds." +
-  " If this is not set, the value for `listeners` will be used." +
-  " Unlike `listeners` it is not valid to advertise the 0.0.0.0 meta-address."
+  val AdvertisedListenersDoc = s"Listeners to publish to ZooKeeper for clients to use, if different than the $ListenersProp config property." +
+    " In IaaS environments, this may need to be different from the interface to which the broker binds." +
+    s" If this is not set, the value for $ListenersProp will be used." +
+    s" Unlike $ListenersProp, it is not valid to advertise the 0.0.0.0 meta-address.\n" +
+    s" Also unlike $ListenersProp, there can be duplicated ports in this property," +
+    " so that one listener can be configured to advertise another listener's address." +
+    " This can be useful in some cases where external load balancers are used."
   val ListenerSecurityProtocolMapDoc = "Map between listener names and security protocols. This must be defined for " +
     "the same security protocol to be usable in more than one port or IP. For example, internal and " +
     "external traffic can be separated even if SSL is required for both. Concretely, the user could define listeners " +
@@ -474,15 +680,46 @@ object KafkaConfig {
     "separated by a colon and map entries are separated by commas. Each listener name should only appear once in the map. " +
     "Different security (SSL and SASL) settings can be configured for each listener by adding a normalised " +
     "prefix (the listener name is lowercased) to the config name. For example, to set a different keystore for the " +
-    "INTERNAL listener, a config with name `listener.name.internal.ssl.keystore.location` would be set. " +
-    "If the config for the listener name is not set, the config will fallback to the generic config (i.e. `ssl.keystore.location`). "
-
-  val SocketSendBufferBytesDoc = "The SO_SNDBUF buffer of the socket sever sockets. If the value is -1, the OS default will be used."
-  val SocketReceiveBufferBytesDoc = "The SO_RCVBUF buffer of the socket sever sockets. If the value is -1, the OS default will be used."
+    "INTERNAL listener, a config with name listener.name.internal.ssl.keystore.location would be set. " +
+    "If the config for the listener name is not set, the config will fallback to the generic config (i.e. ssl.keystore.location). "
+  val controlPlaneListenerNameDoc = "Name of listener used for communication between controller and brokers. " +
+    s"Broker will use the $ControlPlaneListenerNameProp to locate the endpoint in $ListenersProp list, to listen for connections from the controller. " +
+    "For example, if a broker's config is :\n" +
+    "listeners = INTERNAL://192.1.1.8:9092, EXTERNAL://10.1.1.5:9093, CONTROLLER://192.1.1.8:9094\n" +
+    "listener.security.protocol.map = INTERNAL:PLAINTEXT, EXTERNAL:SSL, CONTROLLER:SSL\n" +
+    "control.plane.listener.name = CONTROLLER\n" +
+    "On startup, the broker will start listening on \"192.1.1.8:9094\" with security protocol \"SSL\".\n" +
+    s"On controller side, when it discovers a broker's published endpoints through zookeeper, it will use the $ControlPlaneListenerNameProp " +
+    "to find the endpoint, which it will use to establish connection to the broker.\n" +
+    "For example, if the broker's published endpoints on zookeeper are :\n" +
+    "\"endpoints\" : [\"INTERNAL://broker1.example.com:9092\",\"EXTERNAL://broker1.example.com:9093\",\"CONTROLLER://broker1.example.com:9094\"]\n" +
+    " and the controller's config is :\n" +
+    "listener.security.protocol.map = INTERNAL:PLAINTEXT, EXTERNAL:SSL, CONTROLLER:SSL\n" +
+    "control.plane.listener.name = CONTROLLER\n" +
+    "then controller will use \"broker1.example.com:9094\" with security protocol \"SSL\" to connect to the broker.\n" +
+    "If not explicitly configured, the default value will be null and there will be no dedicated endpoints for controller connections."
+
+  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 SocketRequestMaxBytesDoc = "The maximum number of bytes in a socket request"
-  val MaxConnectionsPerIpDoc = "The maximum number of connections we allow from each ip address"
-  val MaxConnectionsPerIpOverridesDoc = "Per-ip or hostname overrides to the default maximum number of connections"
+  val MaxConnectionsPerIpDoc = "The maximum number of connections we allow from each ip address. This can be set to 0 if there are overrides " +
+    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 MaxConnectionCreationRateDoc = "The maximum connection creation rate we allow in the broker at any time. Listener-level limits " +
+    s"may also be configured by prefixing the config name with the listener prefix, for example, listener.name.internal.$MaxConnectionCreationRateProp." +
+    "Broker-wide connection rate limit should be configured based on broker capacity while listener limits should be configured based on " +
+    "application requirements. New connections will be throttled if either the listener or the broker limit is reached, with the exception " +
+    "of inter-broker listener. Connections on the inter-broker listener will be throttled only when the listener-level rate limit is reached."
   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."
   /************* Rack Configuration **************/
   val RackDoc = "Rack of the broker. This will be used in rack aware replication assignment for fault tolerance. Examples: `RACK1`, `us-east-1d`"
   /** ********* Log Configuration ***********/
@@ -496,7 +733,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"
 
@@ -510,10 +747,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 "
@@ -551,6 +794,7 @@ object KafkaConfig {
     "implement the org.apache.kafka.server.policy.CreateTopicPolicy interface."
   val AlterConfigPolicyClassNameDoc = "The alter configs policy class that should be used for validation. The class should " +
     "implement the org.apache.kafka.server.policy.AlterConfigPolicy interface."
+  val LogMessageDownConversionEnableDoc = TopicConfig.MESSAGE_DOWNCONVERSION_ENABLE_DOC;
 
   /** ********* Replication configuration ***********/
   val ControllerSocketTimeoutMsDoc = "The socket timeout for controller-to-broker channels"
@@ -566,7 +810,7 @@ object KafkaConfig {
     "message.max.bytes (broker config) or max.message.bytes (topic config)."
   val ReplicaFetchWaitMaxMsDoc = "max wait time for each fetcher request issued by follower replicas. This value should always be less than the " +
   "replica.lag.time.max.ms at all times to prevent frequent shrinking of ISR for low throughput topics"
-  val ReplicaFetchMinBytesDoc = "Minimum bytes expected for each fetch response. If not enough bytes, wait up to replicaMaxWaitTimeMs"
+  val ReplicaFetchMinBytesDoc = "Minimum bytes expected for each fetch response. If not enough bytes, wait up to replica.fetch.wait.max.ms (broker config)."
   val ReplicaFetchResponseMaxBytesDoc = "Maximum bytes expected for the entire fetch response. Records are fetched in batches, " +
     "and if the first record batch in the first non-empty partition of the fetch is larger than this value, the record batch " +
     "will still be returned to ensure that progress can be made. As such, this is not an absolute maximum. The maximum " +
@@ -579,7 +823,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"
@@ -591,39 +835,48 @@ 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."
   val ControlledShutdownEnableDoc = "Enable controlled shutdown of the server"
-  /** ********* Consumer coordinator configuration ***********/
+  /** ********* Group coordinator configuration ***********/
   val GroupMinSessionTimeoutMsDoc = "The minimum allowed session timeout for registered consumers. Shorter timeouts result in quicker failure detection at the cost of more frequent consumer heartbeating, which can overwhelm broker resources."
   val GroupMaxSessionTimeoutMsDoc = "The maximum allowed session timeout for registered consumers. Longer timeouts give consumers more time to process messages in between heartbeats at the cost of a longer time to detect failures."
   val GroupInitialRebalanceDelayMsDoc = "The amount of time the group coordinator will wait for more consumers to join a new group before performing the first rebalance. A longer delay means potentially fewer rebalances, but increases the time until processing begins."
+  val GroupMaxSizeDoc = "The maximum number of consumers that a single consumer group can accommodate."
   /** ********* Offset management configuration ***********/
   val OffsetMetadataMaxSizeDoc = "The maximum size for a metadata entry associated with an offset commit"
-  val OffsetsLoadBufferSizeDoc = "Batch size for reading from the offsets segments when loading offsets into the cache."
+  val OffsetsLoadBufferSizeDoc = "Batch size for reading from the offsets segments when loading offsets into the cache (soft-limit, overridden if records are too large)."
   val OffsetsTopicReplicationFactorDoc = "The replication factor for the offsets topic (set higher to ensure availability). " +
   "Internal topic creation will fail until the cluster size meets this replication factor requirement."
   val OffsetsTopicPartitionsDoc = "The number of partitions for the offset commit topic (should not change after deployment)"
   val OffsetsTopicSegmentBytesDoc = "The offsets topic segment bytes should be kept relatively small in order to facilitate faster log compaction and cache loads"
   val OffsetsTopicCompressionCodecDoc = "Compression codec for the offsets topic - compression may be used to achieve \"atomic\" commits"
-  val OffsetsRetentionMinutesDoc = "Offsets older than this retention period will be discarded"
+  val OffsetsRetentionMinutesDoc = "After a consumer group loses all its consumers (i.e. becomes empty) its offsets will be kept for this retention period before getting discarded. " +
+    "For standalone consumers (using manual assignment), offsets will be expired after the time of last commit plus this retention period."
   val OffsetsRetentionCheckIntervalMsDoc = "Frequency at which to check for stale offsets"
   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"
   /** ********* 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."
-  val TransactionsLoadBufferSizeDoc = "Batch size for reading from the transaction log segments when loading producer ids and transactions into the cache."
+  val TransactionsLoadBufferSizeDoc = "Batch size for reading from the transaction log segments when loading producer ids and transactions into the cache (soft-limit, overridden if records are too large)."
   val TransactionsTopicReplicationFactorDoc = "The replication factor for the transaction topic (set higher to ensure availability). " +
     "Internal topic creation will fail until the cluster size meets this replication factor requirement."
   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 Configuration **************/
+  val MaxIncrementalFetchSessionCacheSlotsDoc = "The maximum number of incremental fetch sessions that we will maintain."
+  val FetchMaxBytesDoc = "The maximum number of bytes we will return for a fetch request. Must be at least 1024."
 
   /** ********* Quota Configuration ***********/
   val ProducerQuotaBytesPerSecondDefaultDoc = "DEPRECATED: Used only when dynamic default quotas are not configured for ,  or  in Zookeeper. " +
@@ -633,18 +886,20 @@ object KafkaConfig {
   val NumQuotaSamplesDoc = "The number of samples to retain in memory for client quotas"
   val NumReplicationQuotaSamplesDoc = "The number of samples to retain in memory for replication quotas"
   val NumAlterLogDirsReplicationQuotaSamplesDoc = "The number of samples to retain in memory for alter log dirs replication quotas"
+  val NumControllerQuotaSamplesDoc = "The number of samples to retain in memory for controller mutation quotas"
   val QuotaWindowSizeSecondsDoc = "The time span of each sample for client quotas"
   val ReplicationQuotaWindowSizeSecondsDoc = "The time span of each sample for replication quotas"
   val AlterLogDirsReplicationQuotaWindowSizeSecondsDoc = "The time span of each sample for alter log dirs replication quotas"
-  /** ********* 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 ControllerQuotaWindowSizeSecondsDoc = "The time span of each sample for controller mutations quotas"
+
+  val ClientQuotaCallbackClassDoc = "The fully qualified name of a class that implements the ClientQuotaCallback interface, " +
+    "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."
 
   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 " +
-  "('gzip', 'snappy', 'lz4'). It additionally accepts 'uncompressed' which is equivalent to no compression; and " +
+  "('gzip', 'snappy', 'lz4', 'zstd'). It additionally accepts 'uncompressed' which is equivalent to no compression; and " +
   "'producer' which means retain the original compression codec set by the producer."
 
   /** ********* Kafka Metrics Configuration ***********/
@@ -653,8 +908,21 @@ object KafkaConfig {
   val MetricReporterClassesDoc = CommonClientConfigs.METRIC_REPORTER_CLASSES_DOC
   val MetricRecordingLevelDoc = CommonClientConfigs.METRICS_RECORDING_LEVEL_DOC
 
+
+  /** ********* Kafka Yammer Metrics Reporter Configuration ***********/
+  val KafkaMetricsReporterClassesDoc = "A list of classes to use as Yammer metrics custom reporters." +
+    " The reporters should implement kafka.metrics.KafkaMetricsReporter trait. If a client wants" +
+    " to expose JMX operations on a custom reporter, the custom reporter needs to additionally implement an MBean" +
+    " trait that extends kafka.metrics.KafkaMetricsReporterMBean trait so that the registered MBean is compliant with" +
+    " the standard MBean convention."
+
+  val KafkaMetricsPollingIntervalSecondsDoc = s"The metrics polling interval (in seconds) which can be used" +
+    s" in $KafkaMetricsReporterClassesProp implementations."
+
   /** ******** 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
@@ -665,24 +933,57 @@ object KafkaConfig {
   val SslKeystoreLocationDoc = SslConfigs.SSL_KEYSTORE_LOCATION_DOC
   val SslKeystorePasswordDoc = SslConfigs.SSL_KEYSTORE_PASSWORD_DOC
   val SslKeyPasswordDoc = SslConfigs.SSL_KEY_PASSWORD_DOC
+  val SslKeystoreKeyDoc = SslConfigs.SSL_KEYSTORE_KEY_DOC
+  val SslKeystoreCertificateChainDoc = SslConfigs.SSL_KEYSTORE_CERTIFICATE_CHAIN_DOC
   val SslTruststoreTypeDoc = SslConfigs.SSL_TRUSTSTORE_TYPE_DOC
   val SslTruststorePasswordDoc = SslConfigs.SSL_TRUSTSTORE_PASSWORD_DOC
   val SslTruststoreLocationDoc = SslConfigs.SSL_TRUSTSTORE_LOCATION_DOC
+  val SslTruststoreCertificatesDoc = SslConfigs.SSL_TRUSTSTORE_CERTIFICATES_DOC
   val SslKeyManagerAlgorithmDoc = SslConfigs.SSL_KEYMANAGER_ALGORITHM_DOC
   val SslTrustManagerAlgorithmDoc = SslConfigs.SSL_TRUSTMANAGER_ALGORITHM_DOC
   val SslEndpointIdentificationAlgorithmDoc = SslConfigs.SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_DOC
   val SslSecureRandomImplementationDoc = SslConfigs.SSL_SECURE_RANDOM_IMPLEMENTATION_DOC
   val SslClientAuthDoc = BrokerSecurityConfigs.SSL_CLIENT_AUTH_DOC
+  val SslPrincipalMappingRulesDoc = BrokerSecurityConfigs.SSL_PRINCIPAL_MAPPING_RULES_DOC
+  val SslEngineFactoryClassDoc = SslConfigs.SSL_ENGINE_FACTORY_CLASS_DOC
 
   /** ********* Sasl Configuration ****************/
   val SaslMechanismInterBrokerProtocolDoc = "SASL mechanism used for inter-broker communication. Default is GSSAPI."
-  val SaslEnabledMechanismsDoc = SaslConfigs.SASL_ENABLED_MECHANISMS_DOC
+  val SaslJaasConfigDoc = SaslConfigs.SASL_JAAS_CONFIG_DOC
+  val SaslEnabledMechanismsDoc = BrokerSecurityConfigs.SASL_ENABLED_MECHANISMS_DOC
+  val SaslServerCallbackHandlerClassDoc = BrokerSecurityConfigs.SASL_SERVER_CALLBACK_HANDLER_CLASS_DOC
+  val SaslClientCallbackHandlerClassDoc = SaslConfigs.SASL_CLIENT_CALLBACK_HANDLER_CLASS_DOC
+  val SaslLoginClassDoc = SaslConfigs.SASL_LOGIN_CLASS_DOC
+  val SaslLoginCallbackHandlerClassDoc = SaslConfigs.SASL_LOGIN_CALLBACK_HANDLER_CLASS_DOC
   val SaslKerberosServiceNameDoc = SaslConfigs.SASL_KERBEROS_SERVICE_NAME_DOC
   val SaslKerberosKinitCmdDoc = SaslConfigs.SASL_KERBEROS_KINIT_CMD_DOC
   val SaslKerberosTicketRenewWindowFactorDoc = SaslConfigs.SASL_KERBEROS_TICKET_RENEW_WINDOW_FACTOR_DOC
   val SaslKerberosTicketRenewJitterDoc = SaslConfigs.SASL_KERBEROS_TICKET_RENEW_JITTER_DOC
   val SaslKerberosMinTimeBeforeReloginDoc = SaslConfigs.SASL_KERBEROS_MIN_TIME_BEFORE_RELOGIN_DOC
   val SaslKerberosPrincipalToLocalRulesDoc = BrokerSecurityConfigs.SASL_KERBEROS_PRINCIPAL_TO_LOCAL_RULES_DOC
+  val SaslLoginRefreshWindowFactorDoc = SaslConfigs.SASL_LOGIN_REFRESH_WINDOW_FACTOR_DOC
+  val SaslLoginRefreshWindowJitterDoc = SaslConfigs.SASL_LOGIN_REFRESH_WINDOW_JITTER_DOC
+  val SaslLoginRefreshMinPeriodSecondsDoc = SaslConfigs.SASL_LOGIN_REFRESH_MIN_PERIOD_SECONDS_DOC
+  val SaslLoginRefreshBufferSecondsDoc = SaslConfigs.SASL_LOGIN_REFRESH_BUFFER_SECONDS_DOC
+
+  /** ********* Delegation Token Configuration ****************/
+  val DelegationTokenSecretKeyAliasDoc = s"DEPRECATED: An alias for $DelegationTokenSecretKeyProp, which should be used instead of this config."
+  val DelegationTokenSecretKeyDoc = "Secret key to generate and verify delegation tokens. The same key must be configured across all the brokers. " +
+    " If the key is not set or set to empty string, brokers will disable the delegation token support."
+  val DelegationTokenMaxLifeTimeDoc = "The token has a maximum lifetime beyond which it cannot be renewed anymore. Default value 7 days."
+  val DelegationTokenExpiryTimeMsDoc = "The token validity time in miliseconds before the token needs to be renewed. Default value 1 day."
+  val DelegationTokenExpiryCheckIntervalDoc = "Scan interval to remove expired delegation tokens."
+
+  /** ********* Password encryption configuration for dynamic configs *********/
+  val PasswordEncoderSecretDoc = "The secret used for encoding dynamically configured passwords for this broker."
+  val PasswordEncoderOldSecretDoc = "The old secret that was used for encoding dynamically configured passwords. " +
+    "This is required only when the secret is updated. If specified, all dynamically encoded passwords are " +
+    s"decoded using this old secret and re-encoded using $PasswordEncoderSecretProp when broker starts up."
+  val PasswordEncoderKeyFactoryAlgorithmDoc = "The SecretKeyFactory algorithm used for encoding dynamically configured passwords. " +
+    "Default is PBKDF2WithHmacSHA512 if available and PBKDF2WithHmacSHA1 otherwise."
+  val PasswordEncoderCipherAlgorithmDoc = "The Cipher algorithm used for encoding dynamically configured passwords."
+  val PasswordEncoderKeyLengthDoc =  "The key length used for encoding dynamically configured passwords."
+  val PasswordEncoderIterationsDoc =  "The iteration count used for encoding dynamically configured passwords."
 
   private val configDef = {
     import ConfigDef.Importance._
@@ -699,6 +1000,20 @@ object KafkaConfig {
       .define(ZkSyncTimeMsProp, INT, Defaults.ZkSyncTimeMs, LOW, ZkSyncTimeMsDoc)
       .define(ZkEnableSecureAclsProp, BOOLEAN, Defaults.ZkEnableSecureAcls, HIGH, ZkEnableSecureAclsDoc)
       .define(ZkMaxInFlightRequestsProp, INT, Defaults.ZkMaxInFlightRequests, atLeast(1), HIGH, ZkMaxInFlightRequestsDoc)
+      .define(ZkSslClientEnableProp, BOOLEAN, Defaults.ZkSslClientEnable, MEDIUM, ZkSslClientEnableDoc)
+      .define(ZkClientCnxnSocketProp, STRING, null, MEDIUM, ZkClientCnxnSocketDoc)
+      .define(ZkSslKeyStoreLocationProp, STRING, null, MEDIUM, ZkSslKeyStoreLocationDoc)
+      .define(ZkSslKeyStorePasswordProp, PASSWORD, null, MEDIUM, ZkSslKeyStorePasswordDoc)
+      .define(ZkSslKeyStoreTypeProp, STRING, null, MEDIUM, ZkSslKeyStoreTypeDoc)
+      .define(ZkSslTrustStoreLocationProp, STRING, null, MEDIUM, ZkSslTrustStoreLocationDoc)
+      .define(ZkSslTrustStorePasswordProp, PASSWORD, null, MEDIUM, ZkSslTrustStorePasswordDoc)
+      .define(ZkSslTrustStoreTypeProp, STRING, null, MEDIUM, ZkSslTrustStoreTypeDoc)
+      .define(ZkSslProtocolProp, STRING, Defaults.ZkSslProtocol, LOW, ZkSslProtocolDoc)
+      .define(ZkSslEnabledProtocolsProp, LIST, null, LOW, ZkSslEnabledProtocolsDoc)
+      .define(ZkSslCipherSuitesProp, LIST, null, LOW, ZkSslCipherSuitesDoc)
+      .define(ZkSslEndpointIdentificationAlgorithmProp, STRING, Defaults.ZkSslEndpointIdentificationAlgorithm, LOW, ZkSslEndpointIdentificationAlgorithmDoc)
+      .define(ZkSslCrlEnableProp, BOOLEAN, Defaults.ZkSslCrlEnable, LOW, ZkSslCrlEnableDoc)
+      .define(ZkSslOcspEnableProp, BOOLEAN, Defaults.ZkSslOcspEnable, LOW, ZkSslOcspEnableDoc)
 
       /** ********* General Configuration ***********/
       .define(BrokerIdGenerationEnableProp, BOOLEAN, Defaults.BrokerIdGenerationEnable, MEDIUM, BrokerIdGenerationEnableDoc)
@@ -712,6 +1027,11 @@ 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(ConnectionSetupTimeoutMsProp, LONG, Defaults.ConnectionSetupTimeoutMs, MEDIUM, ConnectionSetupTimeoutMsDoc)
+      .define(ConnectionSetupTimeoutMaxMsProp, LONG, Defaults.ConnectionSetupTimeoutMaxMs, MEDIUM, ConnectionSetupTimeoutMaxMsDoc)
+
+      // Experimental flag to turn on APIs required for the internal metadata quorum (KIP-500)
+      .defineInternal(EnableMetadataQuorumProp, BOOLEAN, false, LOW)
 
       /************* Authorizer Configuration ***********/
       .define(AuthorizerClassNameProp, STRING, Defaults.AuthorizerClassName, LOW, AuthorizerClassNameDoc)
@@ -724,12 +1044,16 @@ object KafkaConfig {
       .define(AdvertisedPortProp, INT, null, HIGH, AdvertisedPortDoc)
       .define(AdvertisedListenersProp, STRING, null, HIGH, AdvertisedListenersDoc)
       .define(ListenerSecurityProtocolMapProp, STRING, Defaults.ListenerSecurityProtocolMap, LOW, ListenerSecurityProtocolMapDoc)
+      .define(ControlPlaneListenerNameProp, STRING, null, HIGH, controlPlaneListenerNameDoc)
       .define(SocketSendBufferBytesProp, INT, Defaults.SocketSendBufferBytes, HIGH, SocketSendBufferBytesDoc)
       .define(SocketReceiveBufferBytesProp, INT, Defaults.SocketReceiveBufferBytes, HIGH, SocketReceiveBufferBytesDoc)
       .define(SocketRequestMaxBytesProp, INT, Defaults.SocketRequestMaxBytes, atLeast(1), HIGH, SocketRequestMaxBytesDoc)
-      .define(MaxConnectionsPerIpProp, INT, Defaults.MaxConnectionsPerIp, atLeast(1), MEDIUM, MaxConnectionsPerIpDoc)
+      .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(MaxConnectionCreationRateProp, INT, Defaults.MaxConnectionCreationRate, atLeast(0), MEDIUM, MaxConnectionCreationRateDoc)
       .define(ConnectionsMaxIdleMsProp, LONG, Defaults.ConnectionsMaxIdleMs, MEDIUM, ConnectionsMaxIdleMsDoc)
+      .define(FailedAuthenticationDelayMsProp, INT, Defaults.FailedAuthenticationDelayMs, atLeast(0), LOW, FailedAuthenticationDelayMsDoc)
 
       /************ Rack Configuration ******************/
       .define(RackProp, STRING, null, MEDIUM, RackDoc)
@@ -738,7 +1062,7 @@ object KafkaConfig {
       .define(NumPartitionsProp, INT, Defaults.NumPartitions, atLeast(1), MEDIUM, NumPartitionsDoc)
       .define(LogDirProp, STRING, Defaults.LogDir, HIGH, LogDirDoc)
       .define(LogDirsProp, STRING, null, HIGH, LogDirsDoc)
-      .define(LogSegmentBytesProp, INT, Defaults.LogSegmentBytes, atLeast(Message.MinMessageOverhead), HIGH, LogSegmentBytesDoc)
+      .define(LogSegmentBytesProp, INT, Defaults.LogSegmentBytes, atLeast(LegacyRecord.RECORD_OVERHEAD_V0), HIGH, LogSegmentBytesDoc)
 
       .define(LogRollTimeMillisProp, LONG, null, HIGH, LogRollTimeMillisDoc)
       .define(LogRollTimeHoursProp, INT, Defaults.LogRollHours, atLeast(1), HIGH, LogRollTimeHoursDoc)
@@ -763,6 +1087,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)
@@ -775,11 +1100,12 @@ object KafkaConfig {
       .define(NumRecoveryThreadsPerDataDirProp, INT, Defaults.NumRecoveryThreadsPerDataDir, atLeast(1), HIGH, NumRecoveryThreadsPerDataDirDoc)
       .define(AutoCreateTopicsEnableProp, BOOLEAN, Defaults.AutoCreateTopicsEnable, HIGH, AutoCreateTopicsEnableDoc)
       .define(MinInSyncReplicasProp, INT, Defaults.MinInSyncReplicas, atLeast(1), HIGH, MinInSyncReplicasDoc)
-      .define(LogMessageFormatVersionProp, STRING, Defaults.LogMessageFormatVersion, MEDIUM, LogMessageFormatVersionDoc)
+      .define(LogMessageFormatVersionProp, STRING, Defaults.LogMessageFormatVersion, ApiVersionValidator, MEDIUM, LogMessageFormatVersionDoc)
       .define(LogMessageTimestampTypeProp, STRING, Defaults.LogMessageTimestampType, in("CreateTime", "LogAppendTime"), MEDIUM, LogMessageTimestampTypeDoc)
       .define(LogMessageTimestampDifferenceMaxMsProp, LONG, Defaults.LogMessageTimestampDifferenceMaxMs, MEDIUM, LogMessageTimestampDifferenceMaxMsDoc)
       .define(CreateTopicPolicyClassNameProp, CLASS, null, LOW, CreateTopicPolicyClassNameDoc)
       .define(AlterConfigPolicyClassNameProp, CLASS, null, LOW, AlterConfigPolicyClassNameDoc)
+      .define(LogMessageDownConversionEnableProp, BOOLEAN, Defaults.MessageDownConversionEnable, LOW, LogMessageDownConversionEnableDoc)
 
       /** ********* Replication configuration ***********/
       .define(ControllerSocketTimeoutMsProp, INT, Defaults.ControllerSocketTimeoutMs, MEDIUM, ControllerSocketTimeoutMsDoc)
@@ -802,8 +1128,9 @@ object KafkaConfig {
       .define(LeaderImbalanceCheckIntervalSecondsProp, LONG, Defaults.LeaderImbalanceCheckIntervalSeconds, HIGH, LeaderImbalanceCheckIntervalSecondsDoc)
       .define(UncleanLeaderElectionEnableProp, BOOLEAN, Defaults.UncleanLeaderElectionEnable, HIGH, UncleanLeaderElectionEnableDoc)
       .define(InterBrokerSecurityProtocolProp, STRING, Defaults.InterBrokerSecurityProtocol, MEDIUM, InterBrokerSecurityProtocolDoc)
-      .define(InterBrokerProtocolVersionProp, STRING, Defaults.InterBrokerProtocolVersion, MEDIUM, InterBrokerProtocolVersionDoc)
+      .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)
@@ -814,6 +1141,7 @@ object KafkaConfig {
       .define(GroupMinSessionTimeoutMsProp, INT, Defaults.GroupMinSessionTimeoutMs, MEDIUM, GroupMinSessionTimeoutMsDoc)
       .define(GroupMaxSessionTimeoutMsProp, INT, Defaults.GroupMaxSessionTimeoutMs, MEDIUM, GroupMaxSessionTimeoutMsDoc)
       .define(GroupInitialRebalanceDelayMsProp, INT, Defaults.GroupInitialRebalanceDelayMs, MEDIUM, GroupInitialRebalanceDelayMsDoc)
+      .define(GroupMaxSizeProp, INT, Defaults.GroupMaxSize, atLeast(1), MEDIUM, GroupMaxSizeDoc)
 
       /** ********* Offset management configuration ***********/
       .define(OffsetMetadataMaxSizeProp, INT, Defaults.OffsetMetadataMaxSize, HIGH, OffsetMetadataMaxSizeDoc)
@@ -840,21 +1168,36 @@ object KafkaConfig {
       .define(TransactionsAbortTimedOutTransactionCleanupIntervalMsProp, INT, Defaults.TransactionsAbortTimedOutTransactionsCleanupIntervalMS, atLeast(1), LOW, TransactionsAbortTimedOutTransactionsIntervalMsDoc)
       .define(TransactionsRemoveExpiredTransactionalIdCleanupIntervalMsProp, INT, Defaults.TransactionsRemoveExpiredTransactionsCleanupIntervalMS, atLeast(1), LOW, TransactionsRemoveExpiredTransactionsIntervalMsDoc)
 
+      /** ********* Fetch Configuration **************/
+      .define(MaxIncrementalFetchSessionCacheSlots, INT, Defaults.MaxIncrementalFetchSessionCacheSlots, atLeast(0), MEDIUM, MaxIncrementalFetchSessionCacheSlotsDoc)
+      .define(FetchMaxBytes, INT, Defaults.FetchMaxBytes, atLeast(1024), MEDIUM, FetchMaxBytesDoc)
+
       /** ********* Kafka Metrics Configuration ***********/
       .define(MetricNumSamplesProp, INT, Defaults.MetricNumSamples, atLeast(1), LOW, MetricNumSamplesDoc)
       .define(MetricSampleWindowMsProp, LONG, Defaults.MetricSampleWindowMs, atLeast(1), LOW, MetricSampleWindowMsDoc)
       .define(MetricReporterClassesProp, LIST, Defaults.MetricReporterClasses, LOW, MetricReporterClassesDoc)
       .define(MetricRecordingLevelProp, STRING, Defaults.MetricRecordingLevel, LOW, MetricRecordingLevelDoc)
 
+      /** ********* Kafka Yammer Metrics Reporter Configuration for docs ***********/
+      .define(KafkaMetricsReporterClassesProp, LIST, Defaults.KafkaMetricReporterClasses, LOW, KafkaMetricsReporterClassesDoc)
+      .define(KafkaMetricsPollingIntervalSecondsProp, INT, Defaults.KafkaMetricsPollingIntervalSeconds, atLeast(1), LOW, KafkaMetricsPollingIntervalSecondsDoc)
+
       /** ********* Quota configuration ***********/
       .define(ProducerQuotaBytesPerSecondDefaultProp, LONG, Defaults.ProducerQuotaBytesPerSecondDefault, atLeast(1), HIGH, ProducerQuotaBytesPerSecondDefaultDoc)
       .define(ConsumerQuotaBytesPerSecondDefaultProp, LONG, Defaults.ConsumerQuotaBytesPerSecondDefault, atLeast(1), HIGH, ConsumerQuotaBytesPerSecondDefaultDoc)
       .define(NumQuotaSamplesProp, INT, Defaults.NumQuotaSamples, atLeast(1), LOW, NumQuotaSamplesDoc)
       .define(NumReplicationQuotaSamplesProp, INT, Defaults.NumReplicationQuotaSamples, atLeast(1), LOW, NumReplicationQuotaSamplesDoc)
       .define(NumAlterLogDirsReplicationQuotaSamplesProp, INT, Defaults.NumAlterLogDirsReplicationQuotaSamples, atLeast(1), LOW, NumAlterLogDirsReplicationQuotaSamplesDoc)
+      .define(NumControllerQuotaSamplesProp, INT, Defaults.NumControllerQuotaSamples, atLeast(1), LOW, NumControllerQuotaSamplesDoc)
       .define(QuotaWindowSizeSecondsProp, INT, Defaults.QuotaWindowSizeSeconds, atLeast(1), LOW, QuotaWindowSizeSecondsDoc)
       .define(ReplicationQuotaWindowSizeSecondsProp, INT, Defaults.ReplicationQuotaWindowSizeSeconds, atLeast(1), LOW, ReplicationQuotaWindowSizeSecondsDoc)
       .define(AlterLogDirsReplicationQuotaWindowSizeSecondsProp, INT, Defaults.AlterLogDirsReplicationQuotaWindowSizeSeconds, atLeast(1), LOW, AlterLogDirsReplicationQuotaWindowSizeSecondsDoc)
+      .define(ControllerQuotaWindowSizeSecondsProp, INT, Defaults.ControllerQuotaWindowSizeSeconds, atLeast(1), LOW, ControllerQuotaWindowSizeSecondsDoc)
+      .define(ClientQuotaCallbackClassProp, CLASS, null, LOW, ClientQuotaCallbackClassDoc)
+
+      /** ********* General Security Configuration ****************/
+      .define(ConnectionsMaxReauthMsProp, LONG, Defaults.ConnectionsMaxReauthMsDefault, MEDIUM, ConnectionsMaxReauthMsDoc)
+      .define(securityProviderClassProp, STRING, null, LOW, securityProviderClassDoc)
 
       /** ********* SSL Configuration ****************/
       .define(PrincipalBuilderClassProp, CLASS, null, MEDIUM, PrincipalBuilderClassDoc)
@@ -865,29 +1208,58 @@ object KafkaConfig {
       .define(SslKeystoreLocationProp, STRING, null, MEDIUM, SslKeystoreLocationDoc)
       .define(SslKeystorePasswordProp, PASSWORD, null, MEDIUM, SslKeystorePasswordDoc)
       .define(SslKeyPasswordProp, PASSWORD, null, MEDIUM, SslKeyPasswordDoc)
+      .define(SslKeystoreKeyProp, PASSWORD, null, MEDIUM, SslKeystoreKeyDoc)
+      .define(SslKeystoreCertificateChainProp, PASSWORD, null, MEDIUM, SslKeystoreCertificateChainDoc)
       .define(SslTruststoreTypeProp, STRING, Defaults.SslTruststoreType, MEDIUM, SslTruststoreTypeDoc)
       .define(SslTruststoreLocationProp, STRING, null, MEDIUM, SslTruststoreLocationDoc)
       .define(SslTruststorePasswordProp, PASSWORD, null, MEDIUM, SslTruststorePasswordDoc)
+      .define(SslTruststoreCertificatesProp, PASSWORD, null, MEDIUM, SslTruststoreCertificatesDoc)
       .define(SslKeyManagerAlgorithmProp, STRING, Defaults.SslKeyManagerAlgorithm, MEDIUM, SslKeyManagerAlgorithmDoc)
       .define(SslTrustManagerAlgorithmProp, STRING, Defaults.SslTrustManagerAlgorithm, MEDIUM, SslTrustManagerAlgorithmDoc)
-      .define(SslEndpointIdentificationAlgorithmProp, STRING, null, LOW, SslEndpointIdentificationAlgorithmDoc)
+      .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(SslCipherSuitesProp, LIST, null, MEDIUM, SslCipherSuitesDoc)
+      .define(SslClientAuthProp, STRING, Defaults.SslClientAuthentication, in(Defaults.SslClientAuthenticationValidValues:_*), MEDIUM, SslClientAuthDoc)
+      .define(SslCipherSuitesProp, LIST, Collections.emptyList(), MEDIUM, SslCipherSuitesDoc)
+      .define(SslPrincipalMappingRulesProp, STRING, Defaults.SslPrincipalMappingRules, LOW, SslPrincipalMappingRulesDoc)
+      .define(SslEngineFactoryClassProp, CLASS, null, LOW, SslEngineFactoryClassDoc)
 
       /** ********* Sasl Configuration ****************/
       .define(SaslMechanismInterBrokerProtocolProp, STRING, Defaults.SaslMechanismInterBrokerProtocol, MEDIUM, SaslMechanismInterBrokerProtocolDoc)
+      .define(SaslJaasConfigProp, PASSWORD, null, MEDIUM, SaslJaasConfigDoc)
       .define(SaslEnabledMechanismsProp, LIST, Defaults.SaslEnabledMechanisms, MEDIUM, SaslEnabledMechanismsDoc)
+      .define(SaslServerCallbackHandlerClassProp, CLASS, null, MEDIUM, SaslServerCallbackHandlerClassDoc)
+      .define(SaslClientCallbackHandlerClassProp, CLASS, null, MEDIUM, SaslClientCallbackHandlerClassDoc)
+      .define(SaslLoginClassProp, CLASS, null, MEDIUM, SaslLoginClassDoc)
+      .define(SaslLoginCallbackHandlerClassProp, CLASS, null, MEDIUM, SaslLoginCallbackHandlerClassDoc)
       .define(SaslKerberosServiceNameProp, STRING, null, MEDIUM, SaslKerberosServiceNameDoc)
       .define(SaslKerberosKinitCmdProp, STRING, Defaults.SaslKerberosKinitCmd, MEDIUM, SaslKerberosKinitCmdDoc)
       .define(SaslKerberosTicketRenewWindowFactorProp, DOUBLE, Defaults.SaslKerberosTicketRenewWindowFactor, MEDIUM, SaslKerberosTicketRenewWindowFactorDoc)
       .define(SaslKerberosTicketRenewJitterProp, DOUBLE, Defaults.SaslKerberosTicketRenewJitter, MEDIUM, SaslKerberosTicketRenewJitterDoc)
       .define(SaslKerberosMinTimeBeforeReloginProp, LONG, Defaults.SaslKerberosMinTimeBeforeRelogin, MEDIUM, SaslKerberosMinTimeBeforeReloginDoc)
       .define(SaslKerberosPrincipalToLocalRulesProp, LIST, Defaults.SaslKerberosPrincipalToLocalRules, MEDIUM, SaslKerberosPrincipalToLocalRulesDoc)
-
+      .define(SaslLoginRefreshWindowFactorProp, DOUBLE, Defaults.SaslLoginRefreshWindowFactor, MEDIUM, SaslLoginRefreshWindowFactorDoc)
+      .define(SaslLoginRefreshWindowJitterProp, DOUBLE, Defaults.SaslLoginRefreshWindowJitter, MEDIUM, SaslLoginRefreshWindowJitterDoc)
+      .define(SaslLoginRefreshMinPeriodSecondsProp, SHORT, Defaults.SaslLoginRefreshMinPeriodSeconds, MEDIUM, SaslLoginRefreshMinPeriodSecondsDoc)
+      .define(SaslLoginRefreshBufferSecondsProp, SHORT, Defaults.SaslLoginRefreshBufferSeconds, MEDIUM, SaslLoginRefreshBufferSecondsDoc)
+      /** ********* Delegation Token Configuration ****************/
+      .define(DelegationTokenSecretKeyAliasProp, PASSWORD, null, MEDIUM, DelegationTokenSecretKeyAliasDoc)
+      .define(DelegationTokenSecretKeyProp, PASSWORD, null, MEDIUM, DelegationTokenSecretKeyDoc)
+      .define(DelegationTokenMaxLifeTimeProp, LONG, Defaults.DelegationTokenMaxLifeTimeMsDefault, atLeast(1), MEDIUM, DelegationTokenMaxLifeTimeDoc)
+      .define(DelegationTokenExpiryTimeMsProp, LONG, Defaults.DelegationTokenExpiryTimeMsDefault, atLeast(1), MEDIUM, DelegationTokenExpiryTimeMsDoc)
+      .define(DelegationTokenExpiryCheckIntervalMsProp, LONG, Defaults.DelegationTokenExpiryCheckIntervalMsDefault, atLeast(1), LOW, DelegationTokenExpiryCheckIntervalDoc)
+
+      /** ********* Password encryption configuration for dynamic configs *********/
+      .define(PasswordEncoderSecretProp, PASSWORD, null, MEDIUM, PasswordEncoderSecretDoc)
+      .define(PasswordEncoderOldSecretProp, PASSWORD, null, MEDIUM, PasswordEncoderOldSecretDoc)
+      .define(PasswordEncoderKeyFactoryAlgorithmProp, STRING, null, LOW, PasswordEncoderKeyFactoryAlgorithmDoc)
+      .define(PasswordEncoderCipherAlgorithmProp, STRING, Defaults.PasswordEncoderCipherAlgorithm, LOW, PasswordEncoderCipherAlgorithmDoc)
+      .define(PasswordEncoderKeyLengthProp, INT, Defaults.PasswordEncoderKeyLength, atLeast(8), LOW, PasswordEncoderKeyLengthDoc)
+      .define(PasswordEncoderIterationsProp, INT, Defaults.PasswordEncoderIterations, atLeast(1024), LOW, PasswordEncoderIterationsDoc)
   }
 
-  def configNames() = configDef.names().asScala.toList.sorted
+  def configNames: Seq[String] = configDef.names.asScala.toBuffer.sorted
+  private[server] def defaultValues: Map[String, _] = configDef.defaultValues.asScala
+  private[server] def configKeys: Map[String, ConfigKey] = configDef.configKeys.asScala
 
   def fromProps(props: Properties): KafkaConfig =
     fromProps(props, true)
@@ -907,11 +1279,78 @@ object KafkaConfig {
 
   def apply(props: java.util.Map[_, _]): KafkaConfig = new KafkaConfig(props, true)
 
+  private def typeOf(name: String): Option[ConfigDef.Type] = Option(configDef.configKeys.get(name)).map(_.`type`)
+
+  def configType(configName: String): Option[ConfigDef.Type] = {
+    val configType = configTypeExact(configName)
+    if (configType.isDefined) {
+      return configType
+    }
+    typeOf(configName) match {
+      case Some(t) => Some(t)
+      case None =>
+        DynamicBrokerConfig.brokerConfigSynonyms(configName, matchListenerOverride = true).flatMap(typeOf).headOption
+    }
+  }
+
+  private def configTypeExact(exactName: String): Option[ConfigDef.Type] = {
+    val configType = typeOf(exactName).orNull
+    if (configType != null) {
+      Some(configType)
+    } else {
+      val configKey = DynamicConfig.Broker.brokerConfigDef.configKeys().get(exactName)
+      if (configKey != null) {
+        Some(configKey.`type`)
+      } else {
+        None
+      }
+    }
+  }
+
+  def maybeSensitive(configType: Option[ConfigDef.Type]): Boolean = {
+    // If we can't determine the config entry type, treat it as a sensitive config to be safe
+    configType.isEmpty || configType.contains(ConfigDef.Type.PASSWORD)
+  }
 }
 
-class KafkaConfig(val props: java.util.Map[_, _], doLog: Boolean) extends AbstractConfig(KafkaConfig.configDef, props, doLog) {
+class KafkaConfig(val props: java.util.Map[_, _], doLog: Boolean, dynamicConfigOverride: Option[DynamicBrokerConfig])
+  extends AbstractConfig(KafkaConfig.configDef, props, doLog) {
+
+  def this(props: java.util.Map[_, _]) = this(props, true, None)
+  def this(props: java.util.Map[_, _], doLog: Boolean) = this(props, doLog, None)
+  // Cache the current config to avoid acquiring read lock to access from dynamicConfig
+  @volatile private var currentConfig = this
+  private[server] val dynamicConfig = dynamicConfigOverride.getOrElse(new DynamicBrokerConfig(this))
+
+  private[server] def updateCurrentConfig(newConfig: KafkaConfig): Unit = {
+    this.currentConfig = newConfig
+  }
 
-  def this(props: java.util.Map[_, _]) = this(props, true)
+  // The following captures any system properties impacting ZooKeeper TLS configuration
+  // and defines the default values this instance will use if no explicit config is given.
+  // We make it part of each instance rather than the object to facilitate testing.
+  private val zkClientConfigViaSystemProperties = new ZKClientConfig()
+
+  override def originals: util.Map[String, AnyRef] =
+    if (this eq currentConfig) super.originals else currentConfig.originals
+  override def values: util.Map[String, _] =
+    if (this eq currentConfig) super.values else currentConfig.values
+  override def nonInternalValues: util.Map[String, _] =
+    if (this eq currentConfig) super.nonInternalValues else currentConfig.values
+  override def originalsStrings: util.Map[String, String] =
+    if (this eq currentConfig) super.originalsStrings else currentConfig.originalsStrings
+  override def originalsWithPrefix(prefix: String): util.Map[String, AnyRef] =
+    if (this eq currentConfig) super.originalsWithPrefix(prefix) else currentConfig.originalsWithPrefix(prefix)
+  override def valuesWithPrefixOverride(prefix: String): util.Map[String, AnyRef] =
+    if (this eq currentConfig) super.valuesWithPrefixOverride(prefix) else currentConfig.valuesWithPrefixOverride(prefix)
+  override def get(key: String): AnyRef =
+    if (this eq currentConfig) super.get(key) else currentConfig.get(key)
+
+  //  During dynamic update, we use the values from this config, these are only used in DynamicBrokerConfig
+  private[server] def originalsFromThisConfig: util.Map[String, AnyRef] = super.originals
+  private[server] def valuesFromThisConfig: util.Map[String, _] = super.values
+  private[server] def valuesFromThisConfigWithPrefixOverride(prefix: String): util.Map[String, AnyRef] =
+    super.valuesWithPrefixOverride(prefix)
 
   /** ********* Zookeeper Configuration ***********/
   val zkConnect: String = getString(KafkaConfig.ZkConnectProp)
@@ -921,18 +1360,99 @@ class KafkaConfig(val props: java.util.Map[_, _], doLog: Boolean) extends Abstra
   val zkSyncTimeMs: Int = getInt(KafkaConfig.ZkSyncTimeMsProp)
   val zkEnableSecureAcls: Boolean = getBoolean(KafkaConfig.ZkEnableSecureAclsProp)
   val zkMaxInFlightRequests: Int = getInt(KafkaConfig.ZkMaxInFlightRequestsProp)
+
+  private def zkBooleanConfigOrSystemPropertyWithDefaultValue(propKey: String): Boolean = {
+    // Use the system property if it exists and the Kafka config value was defaulted rather than actually provided
+    // Need to translate any system property value from true/false (String) to true/false (Boolean)
+    val actuallyProvided = originals.containsKey(propKey)
+    if (actuallyProvided) getBoolean(propKey) else {
+      val sysPropValue = KafkaConfig.getZooKeeperClientProperty(zkClientConfigViaSystemProperties, propKey)
+      sysPropValue match {
+        case Some("true") => true
+        case Some(_) => false
+        case _ => getBoolean(propKey) // not specified so use the default value
+      }
+    }
+  }
+
+  private def zkStringConfigOrSystemPropertyWithDefaultValue(propKey: String): String = {
+    // Use the system property if it exists and the Kafka config value was defaulted rather than actually provided
+    val actuallyProvided = originals.containsKey(propKey)
+    if (actuallyProvided) getString(propKey) else {
+      val sysPropValue = KafkaConfig.getZooKeeperClientProperty(zkClientConfigViaSystemProperties, propKey)
+      sysPropValue match {
+        case Some(_) => sysPropValue.get
+        case _ => getString(propKey) // not specified so use the default value
+      }
+    }
+  }
+
+  private def zkOptionalStringConfigOrSystemProperty(propKey: String): Option[String] = {
+    Option(getString(propKey)) match {
+      case config: Some[String] => config
+      case _ => KafkaConfig.getZooKeeperClientProperty(zkClientConfigViaSystemProperties, propKey)
+    }
+  }
+  private def zkPasswordConfigOrSystemProperty(propKey: String): Option[Password] = {
+    Option(getPassword(propKey)) match {
+      case config: Some[Password] => config
+      case _ => {
+        val sysProp = KafkaConfig.getZooKeeperClientProperty (zkClientConfigViaSystemProperties, propKey)
+        if (sysProp.isDefined) Some (new Password (sysProp.get) ) else None
+      }
+    }
+  }
+  private def zkListConfigOrSystemProperty(propKey: String): Option[util.List[String]] = {
+    Option(getList(propKey)) match {
+      case config: Some[util.List[String]] => config
+      case _ => {
+        val sysProp = KafkaConfig.getZooKeeperClientProperty(zkClientConfigViaSystemProperties, propKey)
+        if (sysProp.isDefined) Some(sysProp.get.split("\\s*,\\s*").toList.asJava) else None
+      }
+    }
+  }
+
+  val zkSslClientEnable = zkBooleanConfigOrSystemPropertyWithDefaultValue(KafkaConfig.ZkSslClientEnableProp)
+  val zkClientCnxnSocketClassName = zkOptionalStringConfigOrSystemProperty(KafkaConfig.ZkClientCnxnSocketProp)
+  val zkSslKeyStoreLocation = zkOptionalStringConfigOrSystemProperty(KafkaConfig.ZkSslKeyStoreLocationProp)
+  val zkSslKeyStorePassword = zkPasswordConfigOrSystemProperty(KafkaConfig.ZkSslKeyStorePasswordProp)
+  val zkSslKeyStoreType = zkOptionalStringConfigOrSystemProperty(KafkaConfig.ZkSslKeyStoreTypeProp)
+  val zkSslTrustStoreLocation = zkOptionalStringConfigOrSystemProperty(KafkaConfig.ZkSslTrustStoreLocationProp)
+  val zkSslTrustStorePassword = zkPasswordConfigOrSystemProperty(KafkaConfig.ZkSslTrustStorePasswordProp)
+  val zkSslTrustStoreType = zkOptionalStringConfigOrSystemProperty(KafkaConfig.ZkSslTrustStoreTypeProp)
+  val ZkSslProtocol = zkStringConfigOrSystemPropertyWithDefaultValue(KafkaConfig.ZkSslProtocolProp)
+  val ZkSslEnabledProtocols = zkListConfigOrSystemProperty(KafkaConfig.ZkSslEnabledProtocolsProp)
+  val ZkSslCipherSuites = zkListConfigOrSystemProperty(KafkaConfig.ZkSslCipherSuitesProp)
+  val ZkSslEndpointIdentificationAlgorithm = {
+    // Use the system property if it exists and the Kafka config value was defaulted rather than actually provided
+    // Need to translate any system property value from true/false to HTTPS/
+    val kafkaProp = KafkaConfig.ZkSslEndpointIdentificationAlgorithmProp
+    val actuallyProvided = originals.containsKey(kafkaProp)
+    if (actuallyProvided) getString(kafkaProp) else {
+      val sysPropValue = KafkaConfig.getZooKeeperClientProperty(zkClientConfigViaSystemProperties, kafkaProp)
+      sysPropValue match {
+        case Some("true") => "HTTPS"
+        case Some(_) => ""
+        case _ => getString(kafkaProp) // not specified so use the default value
+      }
+    }
+  }
+  val ZkSslCrlEnable = zkBooleanConfigOrSystemPropertyWithDefaultValue(KafkaConfig.ZkSslCrlEnableProp)
+  val ZkSslOcspEnable = zkBooleanConfigOrSystemPropertyWithDefaultValue(KafkaConfig.ZkSslOcspEnableProp)
   /** ********* General Configuration ***********/
   val brokerIdGenerationEnable: Boolean = getBoolean(KafkaConfig.BrokerIdGenerationEnableProp)
   val maxReservedBrokerId: Int = getInt(KafkaConfig.MaxReservedBrokerIdProp)
   var brokerId: Int = getInt(KafkaConfig.BrokerIdProp)
 
-  val numNetworkThreads = getInt(KafkaConfig.NumNetworkThreadsProp)
-  val backgroundThreads = getInt(KafkaConfig.BackgroundThreadsProp)
+  def numNetworkThreads = getInt(KafkaConfig.NumNetworkThreadsProp)
+  def backgroundThreads = getInt(KafkaConfig.BackgroundThreadsProp)
   val queuedMaxRequests = getInt(KafkaConfig.QueuedMaxRequestsProp)
   val queuedMaxBytes = getLong(KafkaConfig.QueuedMaxBytesProp)
-  val numIoThreads = getInt(KafkaConfig.NumIoThreadsProp)
-  val messageMaxBytes = getInt(KafkaConfig.MessageMaxBytesProp)
+  def numIoThreads = getInt(KafkaConfig.NumIoThreadsProp)
+  def messageMaxBytes = getInt(KafkaConfig.MessageMaxBytesProp)
   val requestTimeoutMs = getInt(KafkaConfig.RequestTimeoutMsProp)
+  val connectionSetupTimeoutMs = getLong(KafkaConfig.ConnectionSetupTimeoutMsProp)
+  val connectionSetupTimeoutMaxMs = getLong(KafkaConfig.ConnectionSetupTimeoutMaxMsProp)
 
   def getNumReplicaAlterLogDirsThreads: Int = {
     val numThreads: Integer = Option(getInt(KafkaConfig.NumReplicaAlterLogDirsThreadsProp)).getOrElse(logDirs.size)
@@ -940,7 +1460,14 @@ class KafkaConfig(val props: java.util.Map[_, _], doLog: Boolean) extends Abstra
   }
 
   /************* Authorizer Configuration ***********/
-  val authorizerClassName: String = getString(KafkaConfig.AuthorizerClassNameProp)
+  val authorizer: Option[Authorizer] = {
+    val className = getString(KafkaConfig.AuthorizerClassNameProp)
+    if (className == null || className.isEmpty)
+      None
+    else {
+      Some(AuthorizerUtils.createAuthorizer(className))
+    }
+  }
 
   /** ********* Socket Server Configuration ***********/
   val hostName = getString(KafkaConfig.HostNameProp)
@@ -954,51 +1481,56 @@ class KafkaConfig(val props: java.util.Map[_, _], doLog: Boolean) extends Abstra
   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)
+  def maxConnectionCreationRate = getInt(KafkaConfig.MaxConnectionCreationRateProp)
   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)
   val numPartitions = getInt(KafkaConfig.NumPartitionsProp)
   val logDirs = CoreUtils.parseCsvList(Option(getString(KafkaConfig.LogDirsProp)).getOrElse(getString(KafkaConfig.LogDirProp)))
-  val logSegmentBytes = getInt(KafkaConfig.LogSegmentBytesProp)
-  val logFlushIntervalMessages = getLong(KafkaConfig.LogFlushIntervalMessagesProp)
+  def logSegmentBytes = getInt(KafkaConfig.LogSegmentBytesProp)
+  def logFlushIntervalMessages = getLong(KafkaConfig.LogFlushIntervalMessagesProp)
   val logCleanerThreads = getInt(KafkaConfig.LogCleanerThreadsProp)
-  val numRecoveryThreadsPerDataDir = getInt(KafkaConfig.NumRecoveryThreadsPerDataDirProp)
+  def numRecoveryThreadsPerDataDir = getInt(KafkaConfig.NumRecoveryThreadsPerDataDirProp)
   val logFlushSchedulerIntervalMs = getLong(KafkaConfig.LogFlushSchedulerIntervalMsProp)
   val logFlushOffsetCheckpointIntervalMs = getInt(KafkaConfig.LogFlushOffsetCheckpointIntervalMsProp).toLong
   val logFlushStartOffsetCheckpointIntervalMs = getInt(KafkaConfig.LogFlushStartOffsetCheckpointIntervalMsProp).toLong
   val logCleanupIntervalMs = getLong(KafkaConfig.LogCleanupIntervalMsProp)
-  val logCleanupPolicy = getList(KafkaConfig.LogCleanupPolicyProp)
+  def logCleanupPolicy = getList(KafkaConfig.LogCleanupPolicyProp)
   val offsetsRetentionMinutes = getInt(KafkaConfig.OffsetsRetentionMinutesProp)
   val offsetsRetentionCheckIntervalMs = getLong(KafkaConfig.OffsetsRetentionCheckIntervalMsProp)
-  val logRetentionBytes = getLong(KafkaConfig.LogRetentionBytesProp)
+  def logRetentionBytes = getLong(KafkaConfig.LogRetentionBytesProp)
   val logCleanerDedupeBufferSize = getLong(KafkaConfig.LogCleanerDedupeBufferSizeProp)
   val logCleanerDedupeBufferLoadFactor = getDouble(KafkaConfig.LogCleanerDedupeBufferLoadFactorProp)
   val logCleanerIoBufferSize = getInt(KafkaConfig.LogCleanerIoBufferSizeProp)
   val logCleanerIoMaxBytesPerSecond = getDouble(KafkaConfig.LogCleanerIoMaxBytesPerSecondProp)
-  val logCleanerDeleteRetentionMs = getLong(KafkaConfig.LogCleanerDeleteRetentionMsProp)
-  val logCleanerMinCompactionLagMs = getLong(KafkaConfig.LogCleanerMinCompactionLagMsProp)
+  def logCleanerDeleteRetentionMs = getLong(KafkaConfig.LogCleanerDeleteRetentionMsProp)
+  def logCleanerMinCompactionLagMs = getLong(KafkaConfig.LogCleanerMinCompactionLagMsProp)
+  def logCleanerMaxCompactionLagMs = getLong(KafkaConfig.LogCleanerMaxCompactionLagMsProp)
   val logCleanerBackoffMs = getLong(KafkaConfig.LogCleanerBackoffMsProp)
-  val logCleanerMinCleanRatio = getDouble(KafkaConfig.LogCleanerMinCleanRatioProp)
+  def logCleanerMinCleanRatio = getDouble(KafkaConfig.LogCleanerMinCleanRatioProp)
   val logCleanerEnable = getBoolean(KafkaConfig.LogCleanerEnableProp)
-  val logIndexSizeMaxBytes = getInt(KafkaConfig.LogIndexSizeMaxBytesProp)
-  val logIndexIntervalBytes = getInt(KafkaConfig.LogIndexIntervalBytesProp)
-  val logDeleteDelayMs = getLong(KafkaConfig.LogDeleteDelayMsProp)
-  val logRollTimeMillis: java.lang.Long = Option(getLong(KafkaConfig.LogRollTimeMillisProp)).getOrElse(60 * 60 * 1000L * getInt(KafkaConfig.LogRollTimeHoursProp))
-  val logRollTimeJitterMillis: java.lang.Long = Option(getLong(KafkaConfig.LogRollTimeJitterMillisProp)).getOrElse(60 * 60 * 1000L * getInt(KafkaConfig.LogRollTimeJitterHoursProp))
-  val logFlushIntervalMs: java.lang.Long = Option(getLong(KafkaConfig.LogFlushIntervalMsProp)).getOrElse(getLong(KafkaConfig.LogFlushSchedulerIntervalMsProp))
-  val logRetentionTimeMillis = getLogRetentionTimeMillis
-  val minInSyncReplicas = getInt(KafkaConfig.MinInSyncReplicasProp)
-  val logPreAllocateEnable: java.lang.Boolean = getBoolean(KafkaConfig.LogPreAllocateProp)
+  def logIndexSizeMaxBytes = getInt(KafkaConfig.LogIndexSizeMaxBytesProp)
+  def logIndexIntervalBytes = getInt(KafkaConfig.LogIndexIntervalBytesProp)
+  def logDeleteDelayMs = getLong(KafkaConfig.LogDeleteDelayMsProp)
+  def logRollTimeMillis: java.lang.Long = Option(getLong(KafkaConfig.LogRollTimeMillisProp)).getOrElse(60 * 60 * 1000L * getInt(KafkaConfig.LogRollTimeHoursProp))
+  def logRollTimeJitterMillis: java.lang.Long = Option(getLong(KafkaConfig.LogRollTimeJitterMillisProp)).getOrElse(60 * 60 * 1000L * getInt(KafkaConfig.LogRollTimeJitterHoursProp))
+  def logFlushIntervalMs: java.lang.Long = Option(getLong(KafkaConfig.LogFlushIntervalMsProp)).getOrElse(getLong(KafkaConfig.LogFlushSchedulerIntervalMsProp))
+  def minInSyncReplicas = getInt(KafkaConfig.MinInSyncReplicasProp)
+  def logPreAllocateEnable: java.lang.Boolean = getBoolean(KafkaConfig.LogPreAllocateProp)
   // We keep the user-provided String as `ApiVersion.apply` can choose a slightly different version (eg if `0.10.0`
   // is passed, `0.10.0-IV0` may be picked)
   val logMessageFormatVersionString = getString(KafkaConfig.LogMessageFormatVersionProp)
   val logMessageFormatVersion = ApiVersion(logMessageFormatVersionString)
-  val logMessageTimestampType = TimestampType.forName(getString(KafkaConfig.LogMessageTimestampTypeProp))
-  val logMessageTimestampDifferenceMaxMs: Long = getLong(KafkaConfig.LogMessageTimestampDifferenceMaxMsProp)
+  def logMessageTimestampType = TimestampType.forName(getString(KafkaConfig.LogMessageTimestampTypeProp))
+  def logMessageTimestampDifferenceMaxMs: Long = getLong(KafkaConfig.LogMessageTimestampDifferenceMaxMsProp)
+  def logMessageDownConversionEnable: Boolean = getBoolean(KafkaConfig.LogMessageDownConversionEnableProp)
 
   /** ********* Replication configuration ***********/
   val controllerSocketTimeoutMs: Int = getInt(KafkaConfig.ControllerSocketTimeoutMsProp)
@@ -1011,7 +1543,7 @@ class KafkaConfig(val props: java.util.Map[_, _], doLog: Boolean) extends Abstra
   val replicaFetchMinBytes = getInt(KafkaConfig.ReplicaFetchMinBytesProp)
   val replicaFetchResponseMaxBytes = getInt(KafkaConfig.ReplicaFetchResponseMaxBytesProp)
   val replicaFetchBackoffMs = getInt(KafkaConfig.ReplicaFetchBackoffMsProp)
-  val numReplicaFetchers = getInt(KafkaConfig.NumReplicaFetchersProp)
+  def numReplicaFetchers = getInt(KafkaConfig.NumReplicaFetchersProp)
   val replicaHighWatermarkCheckpointIntervalMs = getLong(KafkaConfig.ReplicaHighWatermarkCheckpointIntervalMsProp)
   val fetchPurgatoryPurgeIntervalRequests = getInt(KafkaConfig.FetchPurgatoryPurgeIntervalRequestsProp)
   val producerPurgatoryPurgeIntervalRequests = getInt(KafkaConfig.ProducerPurgatoryPurgeIntervalRequestsProp)
@@ -1019,9 +1551,7 @@ class KafkaConfig(val props: java.util.Map[_, _], doLog: Boolean) extends Abstra
   val autoLeaderRebalanceEnable = getBoolean(KafkaConfig.AutoLeaderRebalanceEnableProp)
   val leaderImbalancePerBrokerPercentage = getInt(KafkaConfig.LeaderImbalancePerBrokerPercentageProp)
   val leaderImbalanceCheckIntervalSeconds = getLong(KafkaConfig.LeaderImbalanceCheckIntervalSecondsProp)
-  val uncleanLeaderElectionEnable: java.lang.Boolean = getBoolean(KafkaConfig.UncleanLeaderElectionEnableProp)
-
-  val (interBrokerListenerName, interBrokerSecurityProtocol) = getInterBrokerListenerNameAndSecurityProtocol
+  def uncleanLeaderElectionEnable: java.lang.Boolean = getBoolean(KafkaConfig.UncleanLeaderElectionEnableProp)
 
   // We keep the user-provided String as `ApiVersion.apply` can choose a slightly different version (eg if `0.10.0`
   // is passed, `0.10.0-IV0` may be picked)
@@ -1033,10 +1563,17 @@ class KafkaConfig(val props: java.util.Map[_, _], doLog: Boolean) extends Abstra
   val controlledShutdownRetryBackoffMs = getLong(KafkaConfig.ControlledShutdownRetryBackoffMsProp)
   val controlledShutdownEnable = getBoolean(KafkaConfig.ControlledShutdownEnableProp)
 
+  /** ********* Feature configuration ***********/
+  def isFeatureVersioningSupported = interBrokerProtocolVersion >= KAFKA_2_7_IV0
+
+  /** ********* Experimental metadata quorum configuration ***********/
+  def metadataQuorumEnabled = getBoolean(KafkaConfig.EnableMetadataQuorumProp)
+
   /** ********* Group coordinator configuration ***********/
   val groupMinSessionTimeoutMs = getInt(KafkaConfig.GroupMinSessionTimeoutMsProp)
   val groupMaxSessionTimeoutMs = getInt(KafkaConfig.GroupMaxSessionTimeoutMsProp)
   val groupInitialRebalanceDelay = getInt(KafkaConfig.GroupInitialRebalanceDelayMsProp)
+  val groupMaxSize = getInt(KafkaConfig.GroupMaxSizeProp)
 
   /** ********* Offset management configuration ***********/
   val offsetMetadataMaxSize = getInt(KafkaConfig.OffsetMetadataMaxSizeProp)
@@ -1065,34 +1602,41 @@ class KafkaConfig(val props: java.util.Map[_, _], doLog: Boolean) extends Abstra
   val metricSampleWindowMs = getLong(KafkaConfig.MetricSampleWindowMsProp)
   val metricRecordingLevel = getString(KafkaConfig.MetricRecordingLevelProp)
 
-  /** ********* SSL Configuration **************/
-  val principalBuilderClass = getClass(KafkaConfig.PrincipalBuilderClassProp)
-  val sslProtocol = getString(KafkaConfig.SslProtocolProp)
-  val sslProvider = getString(KafkaConfig.SslProviderProp)
-  val sslEnabledProtocols = getList(KafkaConfig.SslEnabledProtocolsProp)
-  val sslKeystoreType = getString(KafkaConfig.SslKeystoreTypeProp)
-  val sslKeystoreLocation = getString(KafkaConfig.SslKeystoreLocationProp)
-  val sslKeystorePassword = getPassword(KafkaConfig.SslKeystorePasswordProp)
-  val sslKeyPassword = getPassword(KafkaConfig.SslKeyPasswordProp)
-  val sslTruststoreType = getString(KafkaConfig.SslTruststoreTypeProp)
-  val sslTruststoreLocation = getString(KafkaConfig.SslTruststoreLocationProp)
-  val sslTruststorePassword = getPassword(KafkaConfig.SslTruststorePasswordProp)
-  val sslKeyManagerAlgorithm = getString(KafkaConfig.SslKeyManagerAlgorithmProp)
-  val sslTrustManagerAlgorithm = getString(KafkaConfig.SslTrustManagerAlgorithmProp)
-  val sslClientAuth = getString(KafkaConfig.SslClientAuthProp)
-  val sslCipher = getList(KafkaConfig.SslCipherSuitesProp)
-
-  /** ********* Sasl Configuration **************/
-  val saslMechanismInterBrokerProtocol = getString(KafkaConfig.SaslMechanismInterBrokerProtocolProp)
-  val saslEnabledMechanisms = getList(KafkaConfig.SaslEnabledMechanismsProp)
-  val saslKerberosServiceName = getString(KafkaConfig.SaslKerberosServiceNameProp)
-  val saslKerberosKinitCmd = getString(KafkaConfig.SaslKerberosKinitCmdProp)
-  val saslKerberosTicketRenewWindowFactor = getDouble(KafkaConfig.SaslKerberosTicketRenewWindowFactorProp)
-  val saslKerberosTicketRenewJitter = getDouble(KafkaConfig.SaslKerberosTicketRenewJitterProp)
-  val saslKerberosMinTimeBeforeRelogin = getLong(KafkaConfig.SaslKerberosMinTimeBeforeReloginProp)
-  val saslKerberosPrincipalToLocalRules = getList(KafkaConfig.SaslKerberosPrincipalToLocalRulesProp)
+  /** ********* SSL/SASL Configuration **************/
+  // Security configs may be overridden for listeners, so it is not safe to use the base values
+  // Hence the base SSL/SASL configs are not fields of KafkaConfig, listener configs should be
+  // retrieved using KafkaConfig#valuesWithPrefixOverride
+  private def saslEnabledMechanisms(listenerName: ListenerName): Set[String] = {
+    val value = valuesWithPrefixOverride(listenerName.configPrefix).get(KafkaConfig.SaslEnabledMechanismsProp)
+    if (value != null)
+      value.asInstanceOf[util.List[String]].asScala.toSet
+    else
+      Set.empty[String]
+  }
+
+  def interBrokerListenerName = getInterBrokerListenerNameAndSecurityProtocol._1
+  def interBrokerSecurityProtocol = getInterBrokerListenerNameAndSecurityProtocol._2
+  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
 
+  /** ********* DelegationToken Configuration **************/
+  val delegationTokenSecretKey = Option(getPassword(KafkaConfig.DelegationTokenSecretKeyProp))
+    .getOrElse(getPassword(KafkaConfig.DelegationTokenSecretKeyAliasProp))
+  val tokenAuthEnabled = (delegationTokenSecretKey != null && !delegationTokenSecretKey.value.isEmpty)
+  val delegationTokenMaxLifeMs = getLong(KafkaConfig.DelegationTokenMaxLifeTimeProp)
+  val delegationTokenExpiryTimeMs = getLong(KafkaConfig.DelegationTokenExpiryTimeMsProp)
+  val delegationTokenExpiryCheckIntervalMs = getLong(KafkaConfig.DelegationTokenExpiryCheckIntervalMsProp)
+
+  /** ********* Password encryption configuration for dynamic configs *********/
+  def passwordEncoderSecret = Option(getPassword(KafkaConfig.PasswordEncoderSecretProp))
+  def passwordEncoderOldSecret = Option(getPassword(KafkaConfig.PasswordEncoderOldSecretProp))
+  def passwordEncoderCipherAlgorithm = getString(KafkaConfig.PasswordEncoderCipherAlgorithmProp)
+  def passwordEncoderKeyFactoryAlgorithm = Option(getString(KafkaConfig.PasswordEncoderKeyFactoryAlgorithmProp))
+  def passwordEncoderKeyLength = getInt(KafkaConfig.PasswordEncoderKeyLengthProp)
+  def passwordEncoderIterations = getInt(KafkaConfig.PasswordEncoderIterationsProp)
+
   /** ********* Quota Configuration **************/
   val producerQuotaBytesPerSecondDefault = getLong(KafkaConfig.ProducerQuotaBytesPerSecondDefaultProp)
   val consumerQuotaBytesPerSecondDefault = getLong(KafkaConfig.ConsumerQuotaBytesPerSecondDefaultProp)
@@ -1102,17 +1646,25 @@ class KafkaConfig(val props: java.util.Map[_, _], doLog: Boolean) extends Abstra
   val replicationQuotaWindowSizeSeconds = getInt(KafkaConfig.ReplicationQuotaWindowSizeSecondsProp)
   val numAlterLogDirsReplicationQuotaSamples = getInt(KafkaConfig.NumAlterLogDirsReplicationQuotaSamplesProp)
   val alterLogDirsReplicationQuotaWindowSizeSeconds = getInt(KafkaConfig.AlterLogDirsReplicationQuotaWindowSizeSecondsProp)
+  val numControllerQuotaSamples = getInt(KafkaConfig.NumControllerQuotaSamplesProp)
+  val controllerQuotaWindowSizeSeconds = getInt(KafkaConfig.ControllerQuotaWindowSizeSecondsProp)
 
-  /** ********* Transaction Configuration **************/
-  val transactionIdExpirationMs = getInt(KafkaConfig.TransactionalIdExpirationMsProp)
+  /** ********* Fetch Configuration **************/
+  val maxIncrementalFetchSessionCacheSlots = getInt(KafkaConfig.MaxIncrementalFetchSessionCacheSlots)
+  val fetchMaxBytes = getInt(KafkaConfig.FetchMaxBytes)
 
   val deleteTopicEnable = getBoolean(KafkaConfig.DeleteTopicEnableProp)
-  val compressionType = getString(KafkaConfig.CompressionTypeProp)
-  val listeners: Seq[EndPoint] = getListeners
-  val advertisedListeners: Seq[EndPoint] = getAdvertisedListeners
-  private[kafka] lazy val listenerSecurityProtocolMap = getListenerSecurityProtocolMap
+  def compressionType = getString(KafkaConfig.CompressionTypeProp)
+
+  def addReconfigurable(reconfigurable: Reconfigurable): Unit = {
+    dynamicConfig.addReconfigurable(reconfigurable)
+  }
 
-  private def getLogRetentionTimeMillis: Long = {
+  def removeReconfigurable(reconfigurable: Reconfigurable): Unit = {
+    dynamicConfig.removeReconfigurable(reconfigurable)
+  }
+
+  def logRetentionTimeMillis: Long = {
     val millisInMinute = 60L * 1000L
     val millisInHour = 60L * millisInMinute
 
@@ -1137,23 +1689,36 @@ class KafkaConfig(val props: java.util.Map[_, _], doLog: Boolean) extends Abstra
 
   // If the user did not define listeners but did define host or port, let's use them in backward compatible way
   // If none of those are defined, we default to PLAINTEXT://:9092
-  private def getListeners: Seq[EndPoint] = {
+  def listeners: Seq[EndPoint] = {
     Option(getString(KafkaConfig.ListenersProp)).map { listenerProp =>
       CoreUtils.listenerListToEndPoints(listenerProp, listenerSecurityProtocolMap)
     }.getOrElse(CoreUtils.listenerListToEndPoints("PLAINTEXT://" + hostName + ":" + port, listenerSecurityProtocolMap))
   }
 
+  def controlPlaneListener: Option[EndPoint] = {
+    controlPlaneListenerName.map { listenerName =>
+      listeners.filter(endpoint => endpoint.listenerName.value() == listenerName.value()).head
+    }
+  }
+
+  def dataPlaneListeners: Seq[EndPoint] = {
+    Option(getString(KafkaConfig.ControlPlaneListenerNameProp)) match {
+      case Some(controlPlaneListenerName) => listeners.filterNot(_.listenerName.value() == controlPlaneListenerName)
+      case None => listeners
+    }
+  }
+
   // If the user defined advertised listeners, we use those
-  // If he didn't but did define advertised host or port, we'll use those and fill in the missing value from regular host / port or defaults
+  // If they didn't but did define advertised host or port, we'll use those and fill in the missing value from regular host / port or defaults
   // If none of these are defined, we'll use the listeners
-  private def getAdvertisedListeners: Seq[EndPoint] = {
+  def advertisedListeners: Seq[EndPoint] = {
     val advertisedListenersProp = getString(KafkaConfig.AdvertisedListenersProp)
     if (advertisedListenersProp != null)
-      CoreUtils.listenerListToEndPoints(advertisedListenersProp, listenerSecurityProtocolMap)
+      CoreUtils.listenerListToEndPoints(advertisedListenersProp, listenerSecurityProtocolMap, requireDistinctPorts=false)
     else if (getString(KafkaConfig.AdvertisedHostNameProp) != null || getInt(KafkaConfig.AdvertisedPortProp) != null)
-      CoreUtils.listenerListToEndPoints("PLAINTEXT://" + advertisedHostName + ":" + advertisedPort, listenerSecurityProtocolMap)
+      CoreUtils.listenerListToEndPoints("PLAINTEXT://" + advertisedHostName + ":" + advertisedPort, listenerSecurityProtocolMap, requireDistinctPorts=false)
     else
-      getListeners
+      listeners
   }
 
   private def getInterBrokerListenerNameAndSecurityProtocol: (ListenerName, SecurityProtocol) = {
@@ -1174,6 +1739,19 @@ class KafkaConfig(val props: java.util.Map[_, _], doLog: Boolean) extends Abstra
     }
   }
 
+  private def getControlPlaneListenerNameAndSecurityProtocol: Option[(ListenerName, SecurityProtocol)] = {
+    Option(getString(KafkaConfig.ControlPlaneListenerNameProp)) match {
+      case Some(name) =>
+        val listenerName = ListenerName.normalised(name)
+        val securityProtocol = listenerSecurityProtocolMap.getOrElse(listenerName,
+          throw new ConfigException(s"Listener with ${listenerName.value} defined in " +
+            s"${KafkaConfig.ControlPlaneListenerNameProp} not found in ${KafkaConfig.ListenerSecurityProtocolMapProp}."))
+        Some(listenerName, securityProtocol)
+
+      case None => None
+   }
+  }
+
   private def getSecurityProtocol(protocolName: String, configName: String): SecurityProtocol = {
     try SecurityProtocol.forName(protocolName)
     catch {
@@ -1182,7 +1760,7 @@ class KafkaConfig(val props: java.util.Map[_, _], doLog: Boolean) extends Abstra
     }
   }
 
-  private def getListenerSecurityProtocolMap: Map[ListenerName, SecurityProtocol] = {
+  def listenerSecurityProtocolMap: Map[ListenerName, SecurityProtocol] = {
     getMap(KafkaConfig.ListenerSecurityProtocolMapProp, getString(KafkaConfig.ListenerSecurityProtocolMapProp))
       .map { case (listenerName, protocolName) =>
       ListenerName.normalised(listenerName) -> getSecurityProtocol(protocolName, KafkaConfig.ListenerSecurityProtocolMapProp)
@@ -1191,7 +1769,7 @@ class KafkaConfig(val props: java.util.Map[_, _], doLog: Boolean) extends Abstra
 
   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 {
@@ -1204,7 +1782,7 @@ class KafkaConfig(val props: java.util.Map[_, _], doLog: Boolean) extends Abstra
     require(logCleanerDedupeBufferSize / logCleanerThreads > 1024 * 1024, "log.cleaner.dedupe.buffer.size must be at least 1MB per cleaner thread.")
     require(replicaFetchWaitMaxMs <= replicaSocketTimeoutMs, "replica.socket.timeout.ms should always be at least replica.fetch.wait.max.ms" +
       " to prevent unnecessary socket timeouts")
-    require(replicaFetchWaitMaxMs <= replicaLagTimeMaxMs, "replica.fetch.wait.max.ms should always be at least replica.lag.time.max.ms" +
+    require(replicaFetchWaitMaxMs <= replicaLagTimeMaxMs, "replica.fetch.wait.max.ms should always be less than or equal to replica.lag.time.max.ms" +
       " to prevent frequent changes in ISR")
     require(offsetCommitRequiredAcks >= -1 && offsetCommitRequiredAcks <= offsetsTopicReplicationFactor,
       "offsets.commit.required.acks must be greater or equal -1 and less or equal to offsets.topic.replication.factor")
@@ -1224,14 +1802,48 @@ class KafkaConfig(val props: java.util.Map[_, _], doLog: Boolean) extends Abstra
     require(!advertisedListeners.exists(endpoint => endpoint.host=="0.0.0.0"),
       s"${KafkaConfig.AdvertisedListenersProp} cannot use the nonroutable meta-address 0.0.0.0. "+
       s"Use a routable IP address.")
-    require(interBrokerProtocolVersion >= logMessageFormatVersion,
-      s"log.message.format.version $logMessageFormatVersionString cannot be used when inter.broker.protocol.version is set to $interBrokerProtocolVersionString")
+
+    // validate controller.listener.name config
+    if (controlPlaneListenerName.isDefined) {
+      require(advertisedListenerNames.contains(controlPlaneListenerName.get),
+        s"${KafkaConfig.ControlPlaneListenerNameProp} must be a listener name defined in ${KafkaConfig.AdvertisedListenersProp}. " +
+        s"The valid options based on currently configured listeners are ${advertisedListenerNames.map(_.value).mkString(",")}")
+      // controlPlaneListenerName should be different from interBrokerListenerName
+      require(!controlPlaneListenerName.get.value().equals(interBrokerListenerName.value()),
+        s"${KafkaConfig.ControlPlaneListenerNameProp}, when defined, should have a different value from the inter broker listener name. " +
+        s"Currently they both have the value ${controlPlaneListenerName.get}")
+    }
+
+    val recordVersion = logMessageFormatVersion.recordVersion
+    require(interBrokerProtocolVersion.recordVersion.value >= recordVersion.value,
+      s"log.message.format.version $logMessageFormatVersionString can only be used when inter.broker.protocol.version " +
+      s"is set to version ${ApiVersion.minSupportedFor(recordVersion).shortVersion} or higher")
+
+    if (offsetsTopicCompressionCodec == ZStdCompressionCodec)
+      require(interBrokerProtocolVersion.recordVersion.value >= KAFKA_2_1_IV0.recordVersion.value,
+        "offsets.topic.compression.codec zstd can only be used when inter.broker.protocol.version " +
+        s"is set to version ${KAFKA_2_1_IV0.shortVersion} or higher")
+
     val interBrokerUsesSasl = interBrokerSecurityProtocol == SecurityProtocol.SASL_PLAINTEXT || interBrokerSecurityProtocol == SecurityProtocol.SASL_SSL
     require(!interBrokerUsesSasl || saslInterBrokerHandshakeRequestEnable || saslMechanismInterBrokerProtocol == SaslConfigs.GSSAPI_MECHANISM,
       s"Only GSSAPI mechanism is supported for inter-broker communication with SASL when inter.broker.protocol.version is set to $interBrokerProtocolVersionString")
-    require(!interBrokerUsesSasl || saslEnabledMechanisms.contains(saslMechanismInterBrokerProtocol),
+    require(!interBrokerUsesSasl || saslEnabledMechanisms(interBrokerListenerName).contains(saslMechanismInterBrokerProtocol),
       s"${KafkaConfig.SaslMechanismInterBrokerProtocolProp} must be included in ${KafkaConfig.SaslEnabledMechanismsProp} when SASL is used for inter-broker communication")
     require(queuedMaxBytes <= 0 || queuedMaxBytes >= socketRequestMaxBytes,
       s"${KafkaConfig.QueuedMaxBytesProp} must be larger or equal to ${KafkaConfig.SocketRequestMaxBytesProp}")
+
+    if (maxConnectionsPerIp == 0)
+      require(!maxConnectionsPerIpOverrides.isEmpty, s"${KafkaConfig.MaxConnectionsPerIpProp} can be set to zero only if" +
+        s" ${KafkaConfig.MaxConnectionsPerIpOverridesProp} property is set.")
+
+    val invalidAddresses = maxConnectionsPerIpOverrides.keys.filterNot(address => Utils.validHostPattern(address))
+    if (!invalidAddresses.isEmpty)
+      throw new IllegalArgumentException(s"${KafkaConfig.MaxConnectionsPerIpOverridesProp} contains invalid addresses : ${invalidAddresses.mkString(",")}")
+
+    if (connectionsMaxIdleMs >= 0)
+      require(failedAuthenticationDelayMs < connectionsMaxIdleMs,
+        s"${KafkaConfig.FailedAuthenticationDelayMsProp}=$failedAuthenticationDelayMs should always be less than" +
+        s" ${KafkaConfig.ConnectionsMaxIdleMsProp}=$connectionsMaxIdleMs to prevent failed" +
+        s" authentication responses from timing out")
   }
 }
diff --git a/core/src/main/scala/kafka/server/KafkaHealthcheck.scala b/core/src/main/scala/kafka/server/KafkaHealthcheck.scala
deleted file mode 100644
index 0edc07a5f6252..0000000000000
--- a/core/src/main/scala/kafka/server/KafkaHealthcheck.scala
+++ /dev/null
@@ -1,134 +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.net.InetAddress
-import java.util.Locale
-import java.util.concurrent.TimeUnit
-
-import kafka.api.ApiVersion
-import kafka.cluster.EndPoint
-import kafka.metrics.KafkaMetricsGroup
-import kafka.utils._
-import com.yammer.metrics.core.Gauge
-import org.I0Itec.zkclient.IZkStateListener
-import org.apache.kafka.common.security.auth.SecurityProtocol
-import org.apache.zookeeper.Watcher.Event.KeeperState
-
-import scala.collection.mutable.Set
-
-/**
- * This class registers the broker in zookeeper to allow 
- * other brokers and consumers to detect failures. It uses an ephemeral znode with the path:
- *   /brokers/ids/[0...N] --> advertisedHost:advertisedPort
- *   
- * Right now our definition of health is fairly naive. If we register in zk we are healthy, otherwise
- * we are dead.
- */
-class KafkaHealthcheck(brokerId: Int,
-                       advertisedEndpoints: Seq[EndPoint],
-                       zkUtils: ZkUtils,
-                       rack: Option[String],
-                       interBrokerProtocolVersion: ApiVersion) extends Logging {
-
-  private[server] val sessionExpireListener = new SessionExpireListener
-
-  def startup() {
-    zkUtils.subscribeStateChanges(sessionExpireListener)
-    register()
-  }
-
-  /**
-   * Register this broker as "alive" in zookeeper
-   */
-  def register() {
-    val jmxPort = System.getProperty("com.sun.management.jmxremote.port", "-1").toInt
-    val updatedEndpoints = advertisedEndpoints.map(endpoint =>
-      if (endpoint.host == null || endpoint.host.trim.isEmpty)
-        endpoint.copy(host = InetAddress.getLocalHost.getCanonicalHostName)
-      else
-        endpoint
-    )
-
-    // the default host and port are here for compatibility with older clients that only support PLAINTEXT
-    // we choose the first plaintext port, if there is one
-    // or we register an empty endpoint, which means that older clients will not be able to connect
-    val plaintextEndpoint = updatedEndpoints.find(_.securityProtocol == SecurityProtocol.PLAINTEXT).getOrElse(
-      new EndPoint(null, -1, null, null))
-    zkUtils.registerBrokerInZk(brokerId, plaintextEndpoint.host, plaintextEndpoint.port, updatedEndpoints, jmxPort, rack,
-      interBrokerProtocolVersion)
-  }
-
-  def shutdown(): Unit = sessionExpireListener.shutdown()
-
-  /**
-   *  When we get a SessionExpired event, it means that we have lost all ephemeral nodes and ZKClient has re-established
-   *  a connection for us. We need to re-register this broker in the broker registry. We rely on `handleStateChanged`
-   *  to record ZooKeeper connection state metrics.
-   */
-  class SessionExpireListener extends IZkStateListener with KafkaMetricsGroup {
-
-    private val metricNames = Set[String]()
-
-    private[server] val stateToMeterMap = {
-      import KeeperState._
-      val stateToEventTypeMap = Map(
-        Disconnected -> "Disconnects",
-        SyncConnected -> "SyncConnects",
-        AuthFailed -> "AuthFailures",
-        ConnectedReadOnly -> "ReadOnlyConnects",
-        SaslAuthenticated -> "SaslAuthentications",
-        Expired -> "Expires"
-      )
-      stateToEventTypeMap.map { case (state, eventType) =>
-        val name = s"ZooKeeper${eventType}PerSec"
-        metricNames += name
-        state -> newMeter(name, eventType.toLowerCase(Locale.ROOT), TimeUnit.SECONDS)
-      }
-    }
-
-    private[server] val sessionStateGauge =
-      newGauge("SessionState", new Gauge[String] {
-        override def value: String =
-          Option(zkUtils.zkConnection.getZookeeperState.toString).getOrElse("DISCONNECTED")
-      })
-
-    metricNames += "SessionState"
-
-    @throws[Exception]
-    override def handleStateChanged(state: KeeperState) {
-      stateToMeterMap.get(state).foreach(_.mark())
-    }
-
-    @throws[Exception]
-    override def handleNewSession() {
-      info("re-registering broker info in ZK for broker " + brokerId)
-      register()
-      info("done re-registering broker")
-      info("Subscribing to %s path to watch for new topics".format(ZkUtils.BrokerTopicsPath))
-    }
-
-    override def handleSessionEstablishmentError(error: Throwable) {
-      fatal("Could not establish session with zookeeper", error)
-    }
-
-    def shutdown(): Unit = metricNames.foreach(removeMetric(_))
-
-  }
-
-}
diff --git a/core/src/main/scala/kafka/server/KafkaRequestHandler.scala b/core/src/main/scala/kafka/server/KafkaRequestHandler.scala
index a4987817b8c15..f6d868ade7546 100755
--- a/core/src/main/scala/kafka/server/KafkaRequestHandler.scala
+++ b/core/src/main/scala/kafka/server/KafkaRequestHandler.scala
@@ -21,26 +21,35 @@ import kafka.network._
 import kafka.utils._
 import kafka.metrics.KafkaMetricsGroup
 import java.util.concurrent.{CountDownLatch, TimeUnit}
+import java.util.concurrent.atomic.AtomicInteger
 
 import com.yammer.metrics.core.Meter
 import org.apache.kafka.common.internals.FatalExitError
 import org.apache.kafka.common.utils.{KafkaThread, Time}
 
+import scala.collection.mutable
+import scala.jdk.CollectionConverters._
+
+trait ApiRequestHandler {
+  def handle(request: RequestChannel.Request): Unit
+}
+
 /**
  * A thread that answers kafka requests.
  */
 class KafkaRequestHandler(id: Int,
                           brokerId: Int,
                           val aggregateIdleMeter: Meter,
-                          val totalHandlerThreads: Int,
+                          val totalHandlerThreads: AtomicInteger,
                           val requestChannel: RequestChannel,
-                          apis: KafkaApis,
+                          apis: ApiRequestHandler,
                           time: Time) extends Runnable with Logging {
   this.logIdent = "[Kafka Request Handler " + id + " on Broker " + brokerId + "], "
-  private val latch = new CountDownLatch(1)
+  private val shutdownComplete = new CountDownLatch(1)
+  @volatile private var stopped = false
 
-  def run() {
-    while(true) {
+  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
       // time_window is independent of the number of threads, each recorded idle
@@ -50,12 +59,12 @@ class KafkaRequestHandler(id: Int,
       val req = requestChannel.receiveRequest(300)
       val endTime = time.nanoseconds
       val idleTime = endTime - startSelectTime
-      aggregateIdleMeter.mark(idleTime / totalHandlerThreads)
+      aggregateIdleMeter.mark(idleTime / totalHandlerThreads.get)
 
       req match {
         case RequestChannel.ShutdownRequest =>
           debug(s"Kafka request handler $id on broker $brokerId received shut down command")
-          latch.countDown()
+          shutdownComplete.countDown()
           return
 
         case request: RequestChannel.Request =>
@@ -65,41 +74,68 @@ class KafkaRequestHandler(id: Int,
             apis.handle(request)
           } catch {
             case e: FatalExitError =>
-              latch.countDown()
+              shutdownComplete.countDown()
               Exit.exit(e.statusCode)
             case e: Throwable => error("Exception when handling request", e)
           } finally {
-              request.releaseBuffer()
+            request.releaseBuffer()
           }
 
         case null => // continue
       }
     }
+    shutdownComplete.countDown()
+  }
+
+  def stop(): Unit = {
+    stopped = true
   }
 
   def initiateShutdown(): Unit = requestChannel.sendShutdownRequest()
 
-  def awaitShutdown(): Unit = latch.await()
+  def awaitShutdown(): Unit = shutdownComplete.await()
 
 }
 
 class KafkaRequestHandlerPool(val brokerId: Int,
                               val requestChannel: RequestChannel,
-                              val apis: KafkaApis,
+                              val apis: ApiRequestHandler,
                               time: Time,
-                              numThreads: Int) extends Logging with KafkaMetricsGroup {
+                              numThreads: Int,
+                              requestHandlerAvgIdleMetricName: String,
+                              logAndThreadNamePrefix : String) extends Logging with KafkaMetricsGroup {
 
+  private val threadPoolSize: AtomicInteger = new AtomicInteger(numThreads)
   /* a meter to track the average free capacity of the request handlers */
-  private val aggregateIdleMeter = newMeter("RequestHandlerAvgIdlePercent", "percent", TimeUnit.NANOSECONDS)
+  private val aggregateIdleMeter = newMeter(requestHandlerAvgIdleMetricName, "percent", TimeUnit.NANOSECONDS)
 
-  this.logIdent = "[Kafka Request Handler on Broker " + brokerId + "], "
-  val runnables = new Array[KafkaRequestHandler](numThreads)
-  for(i <- 0 until numThreads) {
-    runnables(i) = new KafkaRequestHandler(i, brokerId, aggregateIdleMeter, numThreads, requestChannel, apis, time)
-    KafkaThread.daemon("kafka-request-handler-" + i, runnables(i)).start()
+  this.logIdent = "[" + logAndThreadNamePrefix + " Kafka Request Handler on Broker " + brokerId + "], "
+  val runnables = new mutable.ArrayBuffer[KafkaRequestHandler](numThreads)
+  for (i <- 0 until numThreads) {
+    createHandler(i)
   }
 
-  def shutdown() {
+  def createHandler(id: Int): Unit = synchronized {
+    runnables += new KafkaRequestHandler(id, brokerId, aggregateIdleMeter, threadPoolSize, requestChannel, apis, time)
+    KafkaThread.daemon(logAndThreadNamePrefix + "-kafka-request-handler-" + id, runnables(id)).start()
+  }
+
+  def resizeThreadPool(newSize: Int): Unit = synchronized {
+    val currentSize = threadPoolSize.get
+    info(s"Resizing request handler thread pool size from $currentSize to $newSize")
+    if (newSize > currentSize) {
+      for (i <- currentSize until newSize) {
+        createHandler(i)
+      }
+    } else if (newSize < currentSize) {
+      for (i <- 1 to (currentSize - newSize)) {
+        runnables.remove(currentSize - i).stop()
+      }
+    }
+    threadPoolSize.set(newSize)
+  }
+
+  def shutdown(): Unit = synchronized {
     info("shutting down")
     for (handler <- runnables)
       handler.initiateShutdown()
@@ -115,39 +151,114 @@ 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()
+  }
+
+  // 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"))
+    metricTypeMap.put(BrokerTopicStats.ReassignmentBytesInPerSec, MeterWrapper(BrokerTopicStats.ReassignmentBytesInPerSec, "bytes"))
+    metricTypeMap.put(BrokerTopicStats.ReassignmentBytesOutPerSec, MeterWrapper(BrokerTopicStats.ReassignmentBytesOutPerSec, "bytes"))
+  }
+
+  // used for testing only
+  def metricMap: Map[String, MeterWrapper] = metricTypeMap.toMap
+
+  def messagesInRate: Meter = metricTypeMap.get(BrokerTopicStats.MessagesInPerSec).meter()
+
+  def bytesInRate: Meter = metricTypeMap.get(BrokerTopicStats.BytesInPerSec).meter()
+
+  def bytesOutRate: Meter = metricTypeMap.get(BrokerTopicStats.BytesOutPerSec).meter()
+
+  def bytesRejectedRate: Meter = metricTypeMap.get(BrokerTopicStats.BytesRejectedPerSec).meter()
+
+  private[server] def replicationBytesInRate: Option[Meter] =
+    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: Option[Meter] =
+    if (name.isEmpty) Some(metricTypeMap.get(BrokerTopicStats.ReplicationBytesOutPerSec).meter())
+    else None
+
+  private[server] def reassignmentBytesInPerSec: Option[Meter] =
+    if (name.isEmpty) Some(metricTypeMap.get(BrokerTopicStats.ReassignmentBytesInPerSec).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)
+
+  private[server] def reassignmentBytesOutPerSec: Option[Meter] =
+    if (name.isEmpty) Some(metricTypeMap.get(BrokerTopicStats.ReassignmentBytesOutPerSec).meter())
+    else None
+
+  def failedProduceRequestRate: Meter = metricTypeMap.get(BrokerTopicStats.FailedProduceRequestsPerSec).meter()
+
+  def failedFetchRequestRate: Meter = metricTypeMap.get(BrokerTopicStats.FailedFetchRequestsPerSec).meter()
+
+  def totalProduceRequestRate: Meter = metricTypeMap.get(BrokerTopicStats.TotalProduceRequestsPerSec).meter()
+
+  def totalFetchRequestRate: Meter = metricTypeMap.get(BrokerTopicStats.TotalFetchRequestsPerSec).meter()
+
+  def fetchMessageConversionsRate: Meter = metricTypeMap.get(BrokerTopicStats.FetchMessageConversionsPerSec).meter()
+
+  def produceMessageConversionsRate: Meter = metricTypeMap.get(BrokerTopicStats.ProduceMessageConversionsPerSec).meter()
+
+  def noKeyCompactedTopicRecordsPerSec: Meter = metricTypeMap.get(BrokerTopicStats.NoKeyCompactedTopicRecordsPerSec).meter()
+
+  def invalidMagicNumberRecordsPerSec: Meter = metricTypeMap.get(BrokerTopicStats.InvalidMagicNumberRecordsPerSec).meter()
+
+  def invalidMessageCrcRecordsPerSec: Meter = metricTypeMap.get(BrokerTopicStats.InvalidMessageCrcRecordsPerSec).meter()
+
+  def invalidOffsetOrSequenceRecordsPerSec: Meter = 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())
 }
 
 object BrokerTopicStats {
@@ -163,10 +274,19 @@ object BrokerTopicStats {
   val TotalFetchRequestsPerSec = "TotalFetchRequestsPerSec"
   val FetchMessageConversionsPerSec = "FetchMessageConversionsPerSec"
   val ProduceMessageConversionsPerSec = "ProduceMessageConversionsPerSec"
+  val ReassignmentBytesInPerSec = "ReassignmentBytesInPerSec"
+  val ReassignmentBytesOutPerSec = "ReassignmentBytesOutPerSec"
+
+  // 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))
 }
 
-class BrokerTopicStats {
+class BrokerTopicStats extends Logging {
   import BrokerTopicStats._
 
   private val stats = new Pool[String, BrokerTopicMetrics](Some(valueFactory))
@@ -175,26 +295,64 @@ 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) {
+  def updateReassignmentBytesIn(value: Long): Unit = {
+    allTopicsStats.reassignmentBytesInPerSec.foreach { metric =>
+      metric.mark(value)
+    }
+  }
+
+  def updateReassignmentBytesOut(value: Long): Unit = {
+    allTopicsStats.reassignmentBytesOutPerSec.foreach { metric =>
+      metric.mark(value)
+    }
+  }
+
+  // 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)
+      topicMetrics.closeMetric(BrokerTopicStats.ReassignmentBytesOutPerSec)
+    }
+  }
+
+  // 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)
+      topicMetrics.closeMetric(BrokerTopicStats.ReassignmentBytesInPerSec)
+    }
+  }
+
+  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, isReassignment: Boolean, value: Long): Unit = {
     if (isFollower) {
+      if (isReassignment)
+        updateReassignmentBytesOut(value)
       updateReplicationBytesOut(value)
     } else {
       topicStats(topic).bytesOutRate.mark(value)
@@ -202,10 +360,10 @@ class BrokerTopicStats {
     }
   }
 
-
   def close(): Unit = {
     allTopicsStats.close()
     stats.values.foreach(_.close())
-  }
 
+    info("Broker and topic stats closed")
+  }
 }
diff --git a/core/src/main/scala/kafka/server/KafkaServer.scala b/core/src/main/scala/kafka/server/KafkaServer.scala
index d026018f5cf89..3723eae7593e6 100755
--- a/core/src/main/scala/kafka/server/KafkaServer.scala
+++ b/core/src/main/scala/kafka/server/KafkaServer.scala
@@ -18,37 +18,39 @@
 package kafka.server
 
 import java.io.{File, IOException}
-import java.net.SocketTimeoutException
+import java.net.{InetAddress, SocketTimeoutException}
 import java.util
 import java.util.concurrent._
 import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger}
 
-import com.yammer.metrics.core.Gauge
-import kafka.api.KAFKA_0_9_0
+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
 import kafka.log.{LogConfig, LogManager}
-import kafka.metrics.{KafkaMetricsGroup, KafkaMetricsReporter}
+import kafka.metrics.{KafkaMetricsGroup, KafkaMetricsReporter, KafkaYammerMetrics, LinuxIoMetricsCollector}
 import kafka.network.SocketServer
 import kafka.security.CredentialProvider
-import kafka.security.auth.Authorizer
 import kafka.utils._
-import kafka.zk.KafkaZkClient
-import kafka.zookeeper.ZooKeeperClient
-import org.apache.kafka.clients.{ApiVersions, ManualMetadataUpdater, NetworkClient, NetworkClientUtils}
+import kafka.zk.{BrokerInfo, KafkaZkClient}
+import org.apache.kafka.clients.{ApiVersions, ClientDnsLookup, CommonClientConfigs, ManualMetadataUpdater, NetworkClient, NetworkClientUtils}
 import org.apache.kafka.common.internals.ClusterResourceListeners
-import org.apache.kafka.common.metrics.{JmxReporter, Metrics, _}
+import org.apache.kafka.common.message.ControlledShutdownRequestData
+import org.apache.kafka.common.metrics.{JmxReporter, Metrics, MetricsReporter, _}
 import org.apache.kafka.common.network._
 import org.apache.kafka.common.protocol.Errors
 import org.apache.kafka.common.requests.{ControlledShutdownRequest, ControlledShutdownResponse}
+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.{ClusterResource, Endpoint, Node}
+import org.apache.kafka.server.authorizer.Authorizer
+import org.apache.zookeeper.client.ZKClientConfig
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 import scala.collection.{Map, Seq, mutable}
 
 object KafkaServer {
@@ -68,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)
@@ -78,6 +81,7 @@ object KafkaServer {
     logProps.put(LogConfig.MessageFormatVersionProp, kafkaConfig.logMessageFormatVersion.version)
     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
   }
 
@@ -88,41 +92,75 @@ object KafkaServer {
       .timeWindow(kafkaConfig.metricSampleWindowMs, TimeUnit.MILLISECONDS)
   }
 
+  def zkClientConfigFromKafkaConfig(config: KafkaConfig, forceZkSslClientEnable: Boolean = false) =
+    if (!config.zkSslClientEnable && !forceZkSslClientEnable)
+      None
+    else {
+      val clientConfig = new ZKClientConfig()
+      KafkaConfig.setZooKeeperClientProperty(clientConfig, KafkaConfig.ZkSslClientEnableProp, "true")
+      config.zkClientCnxnSocketClassName.foreach(KafkaConfig.setZooKeeperClientProperty(clientConfig, KafkaConfig.ZkClientCnxnSocketProp, _))
+      config.zkSslKeyStoreLocation.foreach(KafkaConfig.setZooKeeperClientProperty(clientConfig, KafkaConfig.ZkSslKeyStoreLocationProp, _))
+      config.zkSslKeyStorePassword.foreach(x => KafkaConfig.setZooKeeperClientProperty(clientConfig, KafkaConfig.ZkSslKeyStorePasswordProp, x.value))
+      config.zkSslKeyStoreType.foreach(KafkaConfig.setZooKeeperClientProperty(clientConfig, KafkaConfig.ZkSslKeyStoreTypeProp, _))
+      config.zkSslTrustStoreLocation.foreach(KafkaConfig.setZooKeeperClientProperty(clientConfig, KafkaConfig.ZkSslTrustStoreLocationProp, _))
+      config.zkSslTrustStorePassword.foreach(x => KafkaConfig.setZooKeeperClientProperty(clientConfig, KafkaConfig.ZkSslTrustStorePasswordProp, x.value))
+      config.zkSslTrustStoreType.foreach(KafkaConfig.setZooKeeperClientProperty(clientConfig, KafkaConfig.ZkSslTrustStoreTypeProp, _))
+      KafkaConfig.setZooKeeperClientProperty(clientConfig, KafkaConfig.ZkSslProtocolProp, config.ZkSslProtocol)
+      config.ZkSslEnabledProtocols.foreach(KafkaConfig.setZooKeeperClientProperty(clientConfig, KafkaConfig.ZkSslEnabledProtocolsProp, _))
+      config.ZkSslCipherSuites.foreach(KafkaConfig.setZooKeeperClientProperty(clientConfig, KafkaConfig.ZkSslCipherSuitesProp, _))
+      KafkaConfig.setZooKeeperClientProperty(clientConfig, KafkaConfig.ZkSslEndpointIdentificationAlgorithmProp, config.ZkSslEndpointIdentificationAlgorithm)
+      KafkaConfig.setZooKeeperClientProperty(clientConfig, KafkaConfig.ZkSslCrlEnableProp, config.ZkSslCrlEnable.toString)
+      KafkaConfig.setZooKeeperClientProperty(clientConfig, KafkaConfig.ZkSslOcspEnableProp, config.ZkSslOcspEnable.toString)
+      Some(clientConfig)
+    }
+
+  val MIN_INCREMENTAL_FETCH_SESSION_EVICTION_MS: Long = 120000
 }
 
 /**
  * Represents the lifecycle of a single Kafka broker. Handles all functionality required
  * to start up and shutdown a single Kafka node.
  */
-class KafkaServer(val config: KafkaConfig, time: Time = Time.SYSTEM, threadNamePrefix: Option[String] = None, kafkaMetricsReporters: Seq[KafkaMetricsReporter] = List()) extends Logging with KafkaMetricsGroup {
+class KafkaServer(val config: KafkaConfig, time: Time = Time.SYSTEM, threadNamePrefix: Option[String] = None,
+                  kafkaMetricsReporters: Seq[KafkaMetricsReporter] = List()) extends Logging with KafkaMetricsGroup {
   private val startupComplete = new AtomicBoolean(false)
   private val isShuttingDown = new AtomicBoolean(false)
   private val isStartingUp = new AtomicBoolean(false)
 
   private var shutdownLatch = new CountDownLatch(1)
 
-  private val jmxPrefix: String = "kafka.server"
+  //properties for MetricsContext
+  private val metricsPrefix: String = "kafka.server"
+  private val KAFKA_CLUSTER_ID: String = "kafka.cluster.id"
+  private val KAFKA_BROKER_ID: String = "kafka.broker.id"
+
 
   private var logContext: LogContext = null
 
+  var kafkaYammerMetrics: KafkaYammerMetrics = null
   var metrics: Metrics = null
 
   val brokerState: BrokerState = new BrokerState
 
-  var apis: KafkaApis = null
+  var dataPlaneRequestProcessor: KafkaApis = null
+  var controlPlaneRequestProcessor: KafkaApis = null
+
   var authorizer: Option[Authorizer] = None
   var socketServer: SocketServer = null
-  var requestHandlerPool: KafkaRequestHandlerPool = null
+  var dataPlaneRequestHandlerPool: KafkaRequestHandlerPool = null
+  var controlPlaneRequestHandlerPool: KafkaRequestHandlerPool = null
 
   var logDirFailureChannel: LogDirFailureChannel = null
   var logManager: LogManager = null
 
   var replicaManager: ReplicaManager = null
   var adminManager: AdminManager = null
+  var tokenManager: DelegationTokenManager = null
 
   var dynamicConfigHandlers: Map[String, ConfigHandler] = null
   var dynamicConfigManager: DynamicConfigManager = null
   var credentialProvider: CredentialProvider = null
+  var tokenCache: DelegationTokenCache = null
 
   var groupCoordinator: GroupCoordinator = null
 
@@ -130,14 +168,18 @@ class KafkaServer(val config: KafkaConfig, time: Time = Time.SYSTEM, threadNameP
 
   var kafkaController: KafkaController = null
 
-  val kafkaScheduler = new KafkaScheduler(config.backgroundThreads)
+  var forwardingChannelManager: BrokerToControllerChannelManager = null
+
+  var alterIsrChannelManager: BrokerToControllerChannelManager = null
+
+  var kafkaScheduler: KafkaScheduler = null
 
-  var kafkaHealthcheck: KafkaHealthcheck = null
   var metadataCache: MetadataCache = null
   var quotaManagers: QuotaFactory.QuotaManagers = null
 
-  var zkUtils: ZkUtils = null
-  private var zkClient: KafkaZkClient = null
+  val zkClientConfig: ZKClientConfig = KafkaServer.zkClientConfigFromKafkaConfig(config).getOrElse(new ZKClientConfig())
+  private var _zkClient: KafkaZkClient = null
+
   val correlationId: AtomicInteger = new AtomicInteger(0)
   val brokerMetaPropsFile = "meta.properties"
   val brokerMetadataCheckpoints = config.logDirs.map(logDir => (logDir, new BrokerMetadataCheckpoint(new File(logDir + File.separator + brokerMetaPropsFile)))).toMap
@@ -145,108 +187,162 @@ class KafkaServer(val config: KafkaConfig, time: Time = Time.SYSTEM, threadNameP
   private var _clusterId: String = null
   private var _brokerTopicStats: BrokerTopicStats = null
 
+  private var _featureChangeListener: FinalizedFeatureChangeListener = null
+
+  val brokerFeatures: BrokerFeatures = BrokerFeatures.createDefault()
+
+  val featureCache: FinalizedFeatureCache = new FinalizedFeatureCache(brokerFeatures)
+
   def clusterId: String = _clusterId
 
+  // Visible for testing
+  private[kafka] def zkClient = _zkClient
+
   private[kafka] def brokerTopicStats = _brokerTopicStats
 
-  newGauge(
-    "BrokerState",
-    new Gauge[Int] {
-      def value = brokerState.currentState
-    }
-  )
+  private[kafka] def featureChangeListener = _featureChangeListener
 
-  newGauge(
-    "ClusterId",
-    new Gauge[String] {
-      def value = clusterId
-    }
-  )
+  newGauge("BrokerState", () => brokerState.currentState)
+  newGauge("ClusterId", () => clusterId)
+  newGauge("yammer-metrics-count", () =>  KafkaYammerMetrics.defaultRegistry.allMetrics.size)
 
-  newGauge(
-    "yammer-metrics-count",
-    new Gauge[Int] {
-      def value = {
-        com.yammer.metrics.Metrics.defaultRegistry.allMetrics.size
-      }
-    }
-  )
+  val linuxIoMetricsCollector = new LinuxIoMetricsCollector("/proc", time, logger.underlying)
+
+  if (linuxIoMetricsCollector.usable()) {
+    newGauge("linux-disk-read-bytes", () => linuxIoMetricsCollector.readBytes())
+    newGauge("linux-disk-write-bytes", () => linuxIoMetricsCollector.writeBytes())
+  }
 
   /**
    * 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")
 
-      if(isShuttingDown.get)
+      if (isShuttingDown.get)
         throw new IllegalStateException("Kafka server is still shutting down, cannot re-start!")
 
-      if(startupComplete.get)
+      if (startupComplete.get)
         return
 
       val canStartup = isStartingUp.compareAndSet(false, true)
       if (canStartup) {
         brokerState.newState(Starting)
 
-        /* start scheduler */
-        kafkaScheduler.startup()
-
         /* setup zookeeper */
-        zkUtils = initZk()
+        initZkClient(time)
+
+        /* initialize features */
+        _featureChangeListener = new FinalizedFeatureChangeListener(featureCache, _zkClient)
+        if (config.isFeatureVersioningSupported) {
+          _featureChangeListener.initOrThrow(config.zkConnectionTimeoutMs)
+        }
 
         /* Get or create cluster_id */
-        _clusterId = getOrGenerateClusterId(zkUtils)
+        _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
 
+        // initialize dynamic broker configs from ZooKeeper. Any updates made after this will be
+        // applied after DynamicConfigManager starts.
+        config.dynamicConfig.initialize(zkClient)
+
+        /* start scheduler */
+        kafkaScheduler = new KafkaScheduler(config.backgroundThreads)
+        kafkaScheduler.startup()
+
         /* create and configure metrics */
-        val reporters = config.getConfiguredInstances(KafkaConfig.MetricReporterClassesProp, classOf[MetricsReporter],
-            Map[String, AnyRef](KafkaConfig.BrokerIdProp -> (config.brokerId.toString)).asJava)
-        reporters.add(new JmxReporter(jmxPrefix))
+        kafkaYammerMetrics = KafkaYammerMetrics.INSTANCE
+        kafkaYammerMetrics.configure(config.originals)
+
+        val jmxReporter = new JmxReporter()
+        jmxReporter.configure(config.originals)
+
+        val reporters = new util.ArrayList[MetricsReporter]
+        reporters.add(jmxReporter)
+
         val metricConfig = KafkaServer.metricConfig(config)
-        metrics = new Metrics(metricConfig, reporters, time, true)
+        val metricsContext = createKafkaMetricsContext()
+        metrics = new Metrics(metricConfig, reporters, time, true, metricsContext)
 
         /* register broker metrics */
         _brokerTopicStats = new BrokerTopicStats
 
         quotaManagers = QuotaFactory.instantiate(config, metrics, time, threadNamePrefix.getOrElse(""))
-        notifyClusterListeners(kafkaMetricsReporters ++ reporters.asScala)
+        notifyClusterListeners(kafkaMetricsReporters ++ metrics.reporters.asScala)
 
         logDirFailureChannel = new LogDirFailureChannel(config.logDirs.size)
 
-        val zooKeeperClient = new ZooKeeperClient(config.zkConnect, config.zkSessionTimeoutMs, config.zkConnectionTimeoutMs, config.zkMaxInFlightRequests)
-        zkClient = new KafkaZkClient(zooKeeperClient, zkUtils.isSecure)
-
         /* start log manager */
         logManager = LogManager(config, initialOfflineDirs, zkClient, brokerState, kafkaScheduler, time, brokerTopicStats, logDirFailureChannel)
         logManager.startup()
 
         metadataCache = new MetadataCache(config.brokerId)
-        credentialProvider = new CredentialProvider(config.saslEnabledMechanisms)
-
-        socketServer = new SocketServer(config, metrics, time, credentialProvider)
-        socketServer.startup()
+        // Enable delegation token cache for all SCRAM mechanisms to simplify dynamic update.
+        // This keeps the cache up-to-date if new SCRAM mechanisms are enabled dynamically.
+        tokenCache = new DelegationTokenCache(ScramMechanism.mechanismNames)
+        credentialProvider = new CredentialProvider(ScramMechanism.mechanismNames, tokenCache)
+
+        // Create and start the socket server acceptor threads so that the bound port is known.
+        // Delay starting processors until the end of the initialization sequence to ensure
+        // that credentials have been loaded before processing authentications.
+        //
+        // Note that we allow the use of disabled APIs when experimental support for
+        // the internal metadata quorum has been enabled
+        socketServer = new SocketServer(config, metrics, time, credentialProvider,
+          allowDisabledApis = config.metadataQuorumEnabled)
+        socketServer.startup(startProcessingRequests = false)
 
         /* start replica manager */
+        alterIsrChannelManager = new BrokerToControllerChannelManagerImpl(
+          metadataCache, time, metrics, config, "alterIsrChannel", threadNamePrefix)
         replicaManager = createReplicaManager(isShuttingDown)
         replicaManager.startup()
+        alterIsrChannelManager.start()
+
+        val brokerInfo = createBrokerInfo
+        val brokerEpoch = zkClient.registerBroker(brokerInfo)
+
+        // 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)
+        tokenManager.startup()
 
         /* start kafka controller */
-        kafkaController = new KafkaController(config, zkClient, time, metrics, threadNamePrefix)
+        kafkaController = new KafkaController(config, zkClient, time, metrics, brokerInfo, brokerEpoch, tokenManager, brokerFeatures, featureCache, threadNamePrefix)
         kafkaController.startup()
 
+        var forwardingManager: ForwardingManager = null
+        if (config.metadataQuorumEnabled) {
+          /* start forwarding manager */
+          forwardingChannelManager = new BrokerToControllerChannelManagerImpl(metadataCache, time, metrics,
+            config, "forwardingChannel", threadNamePrefix)
+          forwardingChannelManager.start()
+          forwardingManager = new ForwardingManager(forwardingChannelManager, time, config.requestTimeoutMs.longValue())
+        }
+
         adminManager = new AdminManager(config, metrics, metadataCache, zkClient)
 
         /* 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 */
@@ -255,51 +351,63 @@ 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.map { case (ep, cs) =>
+              ep -> cs.toCompletableFuture
+            }
+          case None =>
+            brokerInfo.broker.endPoints.map { ep =>
+              ep.toJava -> CompletableFuture.completedFuture[Void](null)
+            }.toMap
         }
 
+        val fetchManager = new FetchManager(Time.SYSTEM,
+          new FetchSessionCache(config.maxIncrementalFetchSessionCacheSlots,
+            KafkaServer.MIN_INCREMENTAL_FETCH_SESSION_EVICTION_MS))
+
         /* start processing requests */
-        apis = new KafkaApis(socketServer.requestChannel, replicaManager, adminManager, groupCoordinator, transactionCoordinator,
-          kafkaController, zkClient, config.brokerId, config, metadataCache, metrics, authorizer, quotaManagers,
-          brokerTopicStats, clusterId, time)
+        dataPlaneRequestProcessor = new KafkaApis(socketServer.dataPlaneRequestChannel, replicaManager, adminManager, groupCoordinator, transactionCoordinator,
+          kafkaController, forwardingManager, zkClient, config.brokerId, config, metadataCache, metrics, authorizer, quotaManagers,
+          fetchManager, brokerTopicStats, clusterId, time, tokenManager, brokerFeatures, featureCache)
+
+        dataPlaneRequestHandlerPool = new KafkaRequestHandlerPool(config.brokerId, socketServer.dataPlaneRequestChannel, dataPlaneRequestProcessor, time,
+          config.numIoThreads, s"${SocketServer.DataPlaneMetricPrefix}RequestHandlerAvgIdlePercent", SocketServer.DataPlaneThreadPrefix)
 
-        requestHandlerPool = new KafkaRequestHandlerPool(config.brokerId, socketServer.requestChannel, apis, time,
-          config.numIoThreads)
+        socketServer.controlPlaneRequestChannelOpt.foreach { controlPlaneRequestChannel =>
+          controlPlaneRequestProcessor = new KafkaApis(controlPlaneRequestChannel, replicaManager, adminManager, groupCoordinator, transactionCoordinator,
+            kafkaController, forwardingManager, zkClient, config.brokerId, config, metadataCache, metrics, authorizer, quotaManagers,
+            fetchManager, brokerTopicStats, clusterId, time, tokenManager, brokerFeatures, featureCache)
+
+          controlPlaneRequestHandlerPool = new KafkaRequestHandlerPool(config.brokerId, socketServer.controlPlaneRequestChannelOpt.get, controlPlaneRequestProcessor, time,
+            1, s"${SocketServer.ControlPlaneMetricPrefix}RequestHandlerAvgIdlePercent", SocketServer.ControlPlaneThreadPrefix)
+        }
 
         Mx4jLoader.maybeLoad()
 
+        /* Add all reconfigurables for config change notification before starting config handlers */
+        config.dynamicConfig.addReconfigurables(this)
+
         /* start dynamic config manager */
-        dynamicConfigHandlers = Map[String, ConfigHandler](ConfigType.Topic -> new TopicConfigHandler(logManager, config, quotaManagers),
+        dynamicConfigHandlers = Map[String, ConfigHandler](ConfigType.Topic -> new TopicConfigHandler(logManager, config, quotaManagers, kafkaController),
                                                            ConfigType.Client -> new ClientIdConfigHandler(quotaManagers),
                                                            ConfigType.User -> new UserConfigHandler(quotaManagers, credentialProvider),
-                                                           ConfigType.Broker -> new BrokerConfigHandler(config, quotaManagers))
+                                                           ConfigType.Broker -> new BrokerConfigHandler(config, quotaManagers),
+                                                           ConfigType.Ip -> new IpConfigHandler(socketServer.connectionQuotas))
 
         // Create the config manager. start listening to notifications
         dynamicConfigManager = new DynamicConfigManager(zkClient, dynamicConfigHandlers)
         dynamicConfigManager.startup()
 
-        /* tell everyone we are alive */
-        val listeners = config.advertisedListeners.map { endpoint =>
-          if (endpoint.port == 0)
-            endpoint.copy(port = socketServer.boundPort(endpoint.listenerName))
-          else
-            endpoint
-        }
-        kafkaHealthcheck = new KafkaHealthcheck(config.brokerId, listeners, zkUtils, config.rack,
-          config.interBrokerProtocolVersion)
-        kafkaHealthcheck.startup()
-
-        // Now that the broker id is successfully registered via KafkaHealthcheck, checkpoint it
-        checkpointBrokerId(config.brokerId)
+        socketServer.startProcessingRequests(authorizerFutures)
 
         brokerState.newState(RunningAsBroker)
         shutdownLatch = new CountDownLatch(1)
         startupComplete.set(true)
         isStartingUp.set(false)
-        AppInfoParser.registerAppInfo(jmxPrefix, config.brokerId.toString, metrics)
+        AppInfoParser.registerAppInfo(metricsPrefix, config.brokerId.toString, metrics, time.milliseconds())
         info("started")
       }
     }
@@ -312,19 +420,44 @@ class KafkaServer(val config: KafkaConfig, time: Time = Time.SYSTEM, threadNameP
     }
   }
 
-  private def notifyClusterListeners(clusterListeners: Seq[AnyRef]): Unit = {
+  private[server] def notifyClusterListeners(clusterListeners: Seq[AnyRef]): Unit = {
     val clusterResourceListeners = new ClusterResourceListeners
     clusterResourceListeners.maybeAddAll(clusterListeners.asJava)
     clusterResourceListeners.onUpdate(new ClusterResource(clusterId))
   }
 
-  protected def createReplicaManager(isShuttingDown: AtomicBoolean): ReplicaManager =
+  private[server] def notifyMetricsReporters(metricsReporters: Seq[AnyRef]): Unit = {
+    val metricsContext = createKafkaMetricsContext()
+    metricsReporters.foreach {
+      case x: MetricsReporter => x.contextChange(metricsContext)
+      case _ => //do nothing
+    }
+  }
+
+  private[server] def createKafkaMetricsContext() : KafkaMetricsContext = {
+    val contextLabels = new util.HashMap[String, Object]
+    contextLabels.put(KAFKA_CLUSTER_ID, clusterId)
+    contextLabels.put(KAFKA_BROKER_ID, config.brokerId.toString)
+    contextLabels.putAll(config.originalsWithPrefix(CommonClientConfigs.METRICS_CONTEXT_PREFIX))
+    val metricsContext = new KafkaMetricsContext(metricsPrefix, contextLabels)
+    metricsContext
+  }
+
+  protected def createReplicaManager(isShuttingDown: AtomicBoolean): ReplicaManager = {
+    val alterIsrManager = new AlterIsrManagerImpl(alterIsrChannelManager, kafkaScheduler,
+      time, config.brokerId, () => kafkaController.brokerEpoch)
     new ReplicaManager(config, metrics, time, zkClient, kafkaScheduler, logManager, isShuttingDown, quotaManagers,
-      brokerTopicStats, metadataCache, logDirFailureChannel)
+      brokerTopicStats, metadataCache, logDirFailureChannel, alterIsrManager)
+  }
 
-  private def initZk(): ZkUtils = {
+  private def initZkClient(time: Time): Unit = {
     info(s"Connecting to zookeeper on ${config.zkConnect}")
 
+    def createZkClient(zkConnect: String, isSecure: Boolean) = {
+      KafkaZkClient(zkConnect, isSecure, config.zkSessionTimeoutMs, config.zkConnectionTimeoutMs,
+        config.zkMaxInFlightRequests, time, name = Some("Kafka server"), zkClientConfig = Some(zkClientConfig))
+    }
+
     val chrootIndex = config.zkConnect.indexOf("/")
     val chrootOption = {
       if (chrootIndex > 0) Some(config.zkConnect.substring(chrootIndex))
@@ -332,45 +465,64 @@ class KafkaServer(val config: KafkaConfig, time: Time = Time.SYSTEM, threadNameP
     }
 
     val secureAclsEnabled = config.zkEnableSecureAcls
-    val isZkSecurityEnabled = JaasUtils.isZkSecurityEnabled()
+    val isZkSecurityEnabled = JaasUtils.isZkSaslEnabled() || KafkaConfig.zkTlsClientAuthEnabled(zkClientConfig)
 
     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 ZooKeeper client TLS configuration identifying at least $KafkaConfig.ZkSslClientEnableProp, $KafkaConfig.ZkClientCnxnSocketProp, and $KafkaConfig.ZkSslKeyStoreLocationProp was not present and the " +
+        s"verification of the JAAS login file failed ${JaasUtils.zkSecuritySysConfigString}")
 
+    // make sure chroot path exists
     chrootOption.foreach { chroot =>
       val zkConnForChrootCreation = config.zkConnect.substring(0, chrootIndex)
-      val zkClientForChrootCreation = ZkUtils.withMetrics(zkConnForChrootCreation,
-                                              sessionTimeout = config.zkSessionTimeoutMs,
-                                              connectionTimeout = config.zkConnectionTimeoutMs,
-                                              secureAclsEnabled,
-                                              time)
-      zkClientForChrootCreation.makeSurePersistentPathExists(chroot)
+      val zkClient = createZkClient(zkConnForChrootCreation, secureAclsEnabled)
+      zkClient.makeSurePersistentPathExists(chroot)
       info(s"Created zookeeper path $chroot")
-      zkClientForChrootCreation.close()
+      zkClient.close()
     }
 
-    val zkUtils = ZkUtils.withMetrics(config.zkConnect,
-                          sessionTimeout = config.zkSessionTimeoutMs,
-                          connectionTimeout = config.zkConnectionTimeoutMs,
-                          secureAclsEnabled,
-                          time)
-    zkUtils.setupCommonPaths()
-    zkUtils
+    _zkClient = createZkClient(config.zkConnect, secureAclsEnabled)
+    _zkClient.createTopLevelPaths()
   }
 
-  def getOrGenerateClusterId(zkUtils: ZkUtils): String = {
-    zkUtils.getClusterId.getOrElse(zkUtils.createOrGetClusterId(CoreUtils.generateUuidAsBase64))
+  private def getOrGenerateClusterId(zkClient: KafkaZkClient): String = {
+    zkClient.getClusterId.getOrElse(zkClient.createOrGetClusterId(CoreUtils.generateUuidAsBase64()))
+  }
+
+  def createBrokerInfo: BrokerInfo = {
+    val endPoints = config.advertisedListeners.map(e => s"${e.host}:${e.port}")
+    zkClient.getAllBrokersInCluster.filter(_.id != config.brokerId).foreach { broker =>
+      val commonEndPoints = broker.endPoints.map(e => s"${e.host}:${e.port}").intersect(endPoints)
+      require(commonEndPoints.isEmpty, s"Configured end points ${commonEndPoints.mkString(",")} in" +
+        s" advertised listeners are already registered by broker ${broker.id}")
+    }
+
+    val listeners = config.advertisedListeners.map { endpoint =>
+      if (endpoint.port == 0)
+        endpoint.copy(port = socketServer.boundPort(endpoint.listenerName))
+      else
+        endpoint
+    }
+
+    val updatedEndpoints = listeners.map(endpoint =>
+      if (endpoint.host == null || endpoint.host.trim.isEmpty)
+        endpoint.copy(host = InetAddress.getLocalHost.getCanonicalHostName)
+      else
+        endpoint
+    )
+
+    val jmxPort = System.getProperty("com.sun.management.jmxremote.port", "-1").toInt
+    BrokerInfo(
+      Broker(config.brokerId, updatedEndpoints, config.rack, brokerFeatures.supportedFeatures),
+      config.interBrokerProtocolVersion,
+      jmxPort)
   }
 
   /**
-   *  Performs controlled shutdown
+   * Performs controlled shutdown
    */
-  private def controlledShutdown() {
+  private def controlledShutdown(): Unit = {
 
-    def node(broker: Broker): Node = {
-      val brokerEndPoint = broker.getBrokerEndPoint(config.interBrokerListenerName)
-      new Node(brokerEndPoint.id, brokerEndPoint.host, brokerEndPoint.port)
-    }
+    def node(broker: Broker): Node = broker.node(config.interBrokerListenerName)
 
     val socketTimeoutMs = config.controllerSocketTimeoutMs
 
@@ -383,7 +535,9 @@ class KafkaServer(val config: KafkaConfig, time: Time = Time.SYSTEM, threadNameP
           config,
           config.interBrokerListenerName,
           config.saslMechanismInterBrokerProtocol,
-          config.saslInterBrokerHandshakeRequestEnable)
+          time,
+          config.saslInterBrokerHandshakeRequestEnable,
+          logContext)
         val selector = new Selector(
           NetworkReceive.UNLIMITED,
           config.connectionsMaxIdleMs,
@@ -405,6 +559,9 @@ class KafkaServer(val config: KafkaConfig, time: Time = Time.SYSTEM, threadNameP
           Selectable.USE_DEFAULT_BUFFER_SIZE,
           Selectable.USE_DEFAULT_BUFFER_SIZE,
           config.requestTimeoutMs,
+          config.connectionSetupTimeoutMs,
+          config.connectionSetupTimeoutMaxMs,
+          ClientDnsLookup.USE_ALL_DNS_IPS,
           time,
           false,
           new ApiVersions,
@@ -425,22 +582,29 @@ class KafkaServer(val config: KafkaConfig, time: Time = Time.SYSTEM, threadNameP
           // 1. Find the controller and establish a connection to it.
 
           // Get the current controller info. This is to ensure we use the most recent info to issue the
-          // controlled shutdown request
-          val controllerId = zkUtils.getController()
-          //If this method returns None ignore and try again
-          zkUtils.getBrokerInfo(controllerId).foreach { broker =>
-            // if this is the first attempt, if the controller has changed or if an exception was thrown in a previous
-            // attempt, connect to the most recent controller
-            if (ioException || broker != prevController) {
-
-              ioException = false
-
-              if (prevController != null)
-                networkClient.close(node(prevController).idString)
-
-              prevController = broker
-              metadataUpdater.setNodes(Seq(node(prevController)).asJava)
-            }
+          // controlled shutdown request.
+          // If the controller id or the broker registration are missing, we sleep and retry (if there are remaining retries)
+          zkClient.getControllerId match {
+            case Some(controllerId) =>
+              zkClient.getBroker(controllerId) match {
+                case Some(broker) =>
+                  // if this is the first attempt, if the controller has changed or if an exception was thrown in a previous
+                  // attempt, connect to the most recent controller
+                  if (ioException || broker != prevController) {
+
+                    ioException = false
+
+                    if (prevController != null)
+                      networkClient.close(node(prevController).idString)
+
+                    prevController = broker
+                    metadataUpdater.setNodes(Seq(node(prevController)).asJava)
+                  }
+                case None =>
+                  info(s"Broker registration for controller $controllerId is not available (i.e. the Controller's ZK session expired)")
+              }
+            case None =>
+              info("No controller registered in ZooKeeper")
           }
 
           // 2. issue a controlled shutdown to the controller
@@ -451,27 +615,36 @@ class KafkaServer(val config: KafkaConfig, time: Time = Time.SYSTEM, threadNameP
                 throw new SocketTimeoutException(s"Failed to connect within $socketTimeoutMs ms")
 
               // send the controlled shutdown request
-              val controlledShutdownApiVersion: Short = if (config.interBrokerProtocolVersion < KAFKA_0_9_0) 0 else 1
-              val controlledShutdownRequest = new ControlledShutdownRequest.Builder(config.brokerId,
-                controlledShutdownApiVersion)
+              val controlledShutdownApiVersion: Short =
+                if (config.interBrokerProtocolVersion < KAFKA_0_9_0) 0
+                else if (config.interBrokerProtocolVersion < KAFKA_2_2_IV0) 1
+                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("Remaining partitions to move: %s".format(shutdownResponse.partitionsRemaining.asScala.mkString(",")))
-                info("Error code from controller: %d".format(shutdownResponse.error.code))
+                info(s"Remaining partitions to move: ${shutdownResponse.data.remainingPartitions}")
+                info(s"Error from controller: ${shutdownResponse.error}")
               }
             }
             catch {
               case ioe: IOException =>
                 ioException = true
-                warn("Error during controlled shutdown, possibly because leader movement took longer than the configured controller.socket.timeout.ms and/or request.timeout.ms: %s".format(ioe.getMessage))
+                warn("Error during controlled shutdown, possibly because leader movement took longer than the " +
+                  s"configured controller.socket.timeout.ms and/or request.timeout.ms: ${ioe.getMessage}")
                 // ignore and try again
             }
           }
@@ -506,7 +679,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")
 
@@ -520,9 +693,6 @@ class KafkaServer(val config: KafkaConfig, time: Time = Time.SYSTEM, threadNameP
         CoreUtils.swallow(controlledShutdown(), this)
         brokerState.newState(BrokerShuttingDown)
 
-        if (kafkaHealthcheck != null)
-          CoreUtils.swallow(kafkaHealthcheck.shutdown(), this)
-
         if (dynamicConfigManager != null)
           CoreUtils.swallow(dynamicConfigManager.shutdown(), this)
 
@@ -530,13 +700,17 @@ class KafkaServer(val config: KafkaConfig, time: Time = Time.SYSTEM, threadNameP
         // Socket server will be shutdown towards the end of the sequence.
         if (socketServer != null)
           CoreUtils.swallow(socketServer.stopProcessingRequests(), this)
-        if (requestHandlerPool != null)
-          CoreUtils.swallow(requestHandlerPool.shutdown(), this)
-
-        CoreUtils.swallow(kafkaScheduler.shutdown(), this)
-
-        if (apis != null)
-          CoreUtils.swallow(apis.close(), this)
+        if (dataPlaneRequestHandlerPool != null)
+          CoreUtils.swallow(dataPlaneRequestHandlerPool.shutdown(), this)
+        if (controlPlaneRequestHandlerPool != null)
+          CoreUtils.swallow(controlPlaneRequestHandlerPool.shutdown(), this)
+        if (kafkaScheduler != null)
+          CoreUtils.swallow(kafkaScheduler.shutdown(), this)
+
+        if (dataPlaneRequestProcessor != null)
+          CoreUtils.swallow(dataPlaneRequestProcessor.close(), this)
+        if (controlPlaneRequestProcessor != null)
+          CoreUtils.swallow(controlPlaneRequestProcessor.close(), this)
         CoreUtils.swallow(authorizer.foreach(_.close()), this)
         if (adminManager != null)
           CoreUtils.swallow(adminManager.shutdown(), this)
@@ -546,20 +720,33 @@ class KafkaServer(val config: KafkaConfig, time: Time = Time.SYSTEM, threadNameP
         if (groupCoordinator != null)
           CoreUtils.swallow(groupCoordinator.shutdown(), this)
 
+        if (tokenManager != null)
+          CoreUtils.swallow(tokenManager.shutdown(), this)
+
         if (replicaManager != null)
           CoreUtils.swallow(replicaManager.shutdown(), this)
+
+        if (alterIsrChannelManager != null)
+          CoreUtils.swallow(alterIsrChannelManager.shutdown(), this)
+
+        if (forwardingChannelManager != null)
+          CoreUtils.swallow(forwardingChannelManager.shutdown(), this)
+
         if (logManager != null)
           CoreUtils.swallow(logManager.shutdown(), this)
 
         if (kafkaController != null)
           CoreUtils.swallow(kafkaController.shutdown(), this)
-        if (zkUtils != null)
-          CoreUtils.swallow(zkUtils.close(), this)
+
+        if (featureChangeListener != null)
+          CoreUtils.swallow(featureChangeListener.close(), this)
+
         if (zkClient != null)
           CoreUtils.swallow(zkClient.close(), this)
 
         if (quotaManagers != null)
           CoreUtils.swallow(quotaManagers.shutdown(), this)
+
         // Even though socket server is stopped much earlier, controller can generate
         // response for controlled shutdown request. Shutdown server at the end to
         // avoid any failures (e.g. when metrics are recorded)
@@ -570,11 +757,14 @@ class KafkaServer(val config: KafkaConfig, time: Time = Time.SYSTEM, threadNameP
         if (brokerTopicStats != null)
           CoreUtils.swallow(brokerTopicStats.close(), this)
 
+        // Clear all reconfigurable instances stored in DynamicBrokerConfig
+        config.dynamicConfig.clear()
+
         brokerState.newState(NotRunning)
 
         startupComplete.set(false)
         isShuttingDown.set(false)
-        CoreUtils.swallow(AppInfoParser.unregisterAppInfo(jmxPrefix, config.brokerId.toString, metrics), this)
+        CoreUtils.swallow(AppInfoParser.unregisterAppInfo(metricsPrefix, config.brokerId.toString, metrics), this)
         shutdownLatch.countDown()
         info("shut down completed")
       }
@@ -592,77 +782,100 @@ class KafkaServer(val config: KafkaConfig, time: Time = Time.SYSTEM, threadNameP
    */
   def awaitShutdown(): Unit = shutdownLatch.await()
 
-  def getLogManager(): LogManager = logManager
+  def getLogManager: LogManager = logManager
 
   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 =>
+        case e: IOException =>
           offlineDirs += logDir
           error(s"Fail to read $brokerMetaPropsFile under log directory $logDir", e)
       }
     }
 
-    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 = new StringBuilder
 
+      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
+    * sequence id and configured brokerId, we increment the generated sequence id by KafkaConfig.MaxReservedBrokerId.
+    */
   private def generateBrokerId: Int = {
     try {
-      zkUtils.getBrokerSequenceId(config.maxReservedBrokerId)
+      zkClient.generateBrokerSequenceId() + config.maxReservedBrokerId
     } catch {
       case e: Exception =>
         error("Failed to generate broker.id due to ", e)
diff --git a/core/src/main/scala/kafka/server/KafkaServerStartable.scala b/core/src/main/scala/kafka/server/KafkaServerStartable.scala
index eb290d2b09623..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), reporters)
+    new KafkaServerStartable(KafkaConfig.fromProps(serverProps, false), reporters, threadNamePrefix)
   }
 }
 
-class KafkaServerStartable(val serverConfig: KafkaConfig, reporters: Seq[KafkaMetricsReporter]) extends Logging {
-  private val server = new KafkaServer(serverConfig, 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 serverConfig: KafkaConfig, reporters: Seq[KafkaMe
     }
   }
 
-  def shutdown() {
+  def shutdown(): Unit = {
     try server.shutdown()
     catch {
       case _: Throwable =>
@@ -58,7 +64,7 @@ class KafkaServerStartable(val serverConfig: KafkaConfig, reporters: Seq[KafkaMe
    * 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/LogDirFailureChannel.scala b/core/src/main/scala/kafka/server/LogDirFailureChannel.scala
index c78f04ecf824a..897d3fc36f8c1 100644
--- a/core/src/main/scala/kafka/server/LogDirFailureChannel.scala
+++ b/core/src/main/scala/kafka/server/LogDirFailureChannel.scala
@@ -45,9 +45,8 @@ class LogDirFailureChannel(logDirNum: Int) extends Logging {
    */
   def maybeAddOfflineLogDir(logDir: String, msg: => String, e: IOException): Unit = {
     error(msg, e)
-    if (offlineLogDirs.putIfAbsent(logDir, logDir) == null) {
+    if (offlineLogDirs.putIfAbsent(logDir, logDir) == null)
       offlineLogDirQueue.add(logDir)
-    }
   }
 
   /*
diff --git a/core/src/main/scala/kafka/server/LogOffsetMetadata.scala b/core/src/main/scala/kafka/server/LogOffsetMetadata.scala
index edc010e6a8573..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,9 +76,9 @@ 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 = messageOffset.toString + " [" + segmentBaseOffset + " : " + relativePositionInSegment + "]"
+  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 fd11bfa9a5b86..77d947740da6c 100755
--- a/core/src/main/scala/kafka/server/MetadataCache.scala
+++ b/core/src/main/scala/kafka/server/MetadataCache.scala
@@ -17,21 +17,27 @@
 
 package kafka.server
 
+import java.util
+import java.util.Collections
 import java.util.concurrent.locks.ReentrantReadWriteLock
 
-import scala.collection.{Seq, Set, mutable}
-import scala.collection.JavaConverters._
+import scala.collection.{mutable, Seq, Set}
+import scala.jdk.CollectionConverters._
 import kafka.cluster.{Broker, EndPoint}
 import kafka.api._
-import kafka.common.{BrokerEndPointNotAvailableException, TopicAndPartition}
 import kafka.controller.StateChangeLogger
 import kafka.utils.CoreUtils._
 import kafka.utils.Logging
+import kafka.utils.Implicits._
 import org.apache.kafka.common.internals.Topic
-import org.apache.kafka.common.{Node, TopicPartition}
+import org.apache.kafka.common.message.UpdateMetadataRequestData.UpdateMetadataPartitionState
+import org.apache.kafka.common.{Cluster, Node, PartitionInfo, TopicPartition}
+import org.apache.kafka.common.message.MetadataResponseData.MetadataResponseTopic
+import org.apache.kafka.common.message.MetadataResponseData.MetadataResponsePartition
 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
@@ -39,202 +45,335 @@ import org.apache.kafka.common.requests.{MetadataResponse, UpdateMetadataRequest
  */
 class MetadataCache(brokerId: Int) extends Logging {
 
-  private val cache = mutable.Map[String, mutable.Map[Int, UpdateMetadataRequest.PartitionState]]()
-  private var controllerId: Option[Int] = None
-  private val aliveBrokers = mutable.Map[Int, Broker]()
-  private val aliveNodes = mutable.Map[Int, collection.Map[ListenerName, Node]]()
   private val partitionMetadataLock = new ReentrantReadWriteLock()
+  //this is the cache state. every MetadataSnapshot instance is immutable, and updates (performed under a lock)
+  //replace the value with a completely new one. this means reads (which are not under any lock) need to grab
+  //the value of this var (into a val) ONCE and retain that read copy for the duration of their operation.
+  //multiple reads of this value risk getting different snapshots.
+  @volatile private var metadataSnapshot: MetadataSnapshot = MetadataSnapshot(partitionStates = mutable.AnyRefMap.empty,
+    controllerId = None, aliveBrokers = mutable.LongMap.empty, aliveNodes = mutable.LongMap.empty)
 
   this.logIdent = s"[MetadataCache brokerId=$brokerId] "
   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
+  // `List[Integer]` instead of `List[Int]` to avoid a collection copy.
   // filterUnavailableEndpoints exists to support v0 MetadataResponses
-  private def getEndpoints(brokers: Iterable[Int], listenerName: ListenerName, filterUnavailableEndpoints: Boolean): Seq[Node] = {
-    val result = new mutable.ArrayBuffer[Node](math.min(aliveBrokers.size, brokers.size))
-    brokers.foreach { brokerId =>
-      val endpoint = getAliveEndpoint(brokerId, listenerName) match {
-        case None => if (!filterUnavailableEndpoints) Some(new Node(brokerId, "", -1)) else None
-        case Some(node) => Some(node)
+  private def maybeFilterAliveReplicas(snapshot: MetadataSnapshot,
+                                       brokers: java.util.List[Integer],
+                                       listenerName: ListenerName,
+                                       filterUnavailableEndpoints: Boolean): java.util.List[Integer] = {
+    if (!filterUnavailableEndpoints) {
+      brokers
+    } else {
+      val res = new util.ArrayList[Integer](math.min(snapshot.aliveBrokers.size, brokers.size))
+      for (brokerId <- brokers.asScala) {
+        if (hasAliveEndpoint(snapshot, brokerId, listenerName))
+          res.add(brokerId)
       }
-      endpoint.foreach(result +=)
+      res
     }
-    result
   }
 
   // errorUnavailableEndpoints exists to support v0 MetadataResponses
-  private def getPartitionMetadata(topic: String, listenerName: ListenerName, errorUnavailableEndpoints: Boolean): Option[Iterable[MetadataResponse.PartitionMetadata]] = {
-    cache.get(topic).map { partitions =>
+  // If errorUnavailableListeners=true, return LISTENER_NOT_FOUND if listener is missing on the broker.
+  // Otherwise, return LEADER_NOT_AVAILABLE for broker unavailable and missing listener (Metadata response v5 and below).
+  private def getPartitionMetadata(snapshot: MetadataSnapshot, topic: String, listenerName: ListenerName, errorUnavailableEndpoints: Boolean,
+                                   errorUnavailableListeners: Boolean): Option[Iterable[MetadataResponsePartition]] = {
+    snapshot.partitionStates.get(topic).map { partitions =>
       partitions.map { case (partitionId, partitionState) =>
-        val topicPartition = TopicAndPartition(topic, partitionId)
-        val maybeLeader = getAliveEndpoint(partitionState.basePartitionState.leader, listenerName)
-        val replicas = partitionState.basePartitionState.replicas.asScala.map(_.toInt)
-        val replicaInfo = getEndpoints(replicas, listenerName, errorUnavailableEndpoints)
-        val offlineReplicaInfo = getEndpoints(partitionState.offlineReplicas.asScala.map(_.toInt), listenerName, errorUnavailableEndpoints)
+        val topicPartition = new TopicPartition(topic, partitionId.toInt)
+        val leaderBrokerId = partitionState.leader
+        val leaderEpoch = partitionState.leaderEpoch
+        val maybeLeader = getAliveEndpoint(snapshot, leaderBrokerId, listenerName)
+
+        val replicas = partitionState.replicas
+        val filteredReplicas = maybeFilterAliveReplicas(snapshot, replicas, listenerName, errorUnavailableEndpoints)
+
+        val isr = partitionState.isr
+        val filteredIsr = maybeFilterAliveReplicas(snapshot, isr, listenerName, errorUnavailableEndpoints)
+
+        val offlineReplicas = partitionState.offlineReplicas
 
         maybeLeader match {
           case None =>
-            debug(s"Error while fetching metadata for $topicPartition: leader not available")
-            new MetadataResponse.PartitionMetadata(Errors.LEADER_NOT_AVAILABLE, partitionId, Node.noNode(),
-              replicaInfo.asJava, java.util.Collections.emptyList(), offlineReplicaInfo.asJava)
+            val error = if (!snapshot.aliveBrokers.contains(leaderBrokerId)) { // we are already holding the read lock
+              debug(s"Error while fetching metadata for $topicPartition: leader not available")
+              Errors.LEADER_NOT_AVAILABLE
+            } else {
+              debug(s"Error while fetching metadata for $topicPartition: listener $listenerName " +
+                s"not found on leader $leaderBrokerId")
+              if (errorUnavailableListeners) Errors.LISTENER_NOT_FOUND else Errors.LEADER_NOT_AVAILABLE
+            }
 
-          case Some(leader) =>
-            val isr = partitionState.basePartitionState.isr.asScala.map(_.toInt)
-            val isrInfo = getEndpoints(isr, listenerName, errorUnavailableEndpoints)
+            new MetadataResponsePartition()
+              .setErrorCode(error.code)
+              .setPartitionIndex(partitionId.toInt)
+              .setLeaderId(MetadataResponse.NO_LEADER_ID)
+              .setLeaderEpoch(leaderEpoch)
+              .setReplicaNodes(filteredReplicas)
+              .setIsrNodes(filteredIsr)
+              .setOfflineReplicas(offlineReplicas)
 
-            if (replicaInfo.size < replicas.size) {
+          case Some(leader) =>
+            val error = if (filteredReplicas.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(",")}")
-
-              new MetadataResponse.PartitionMetadata(Errors.REPLICA_NOT_AVAILABLE, partitionId, leader,
-                replicaInfo.asJava, isrInfo.asJava, offlineReplicaInfo.asJava)
-            } else if (isrInfo.size < isr.size) {
+                s"following brokers ${replicas.asScala.filterNot(filteredReplicas.contains).mkString(",")}")
+              Errors.REPLICA_NOT_AVAILABLE
+            } else if (filteredIsr.size < isr.size) {
               debug(s"Error while fetching metadata for $topicPartition: in sync replica information not available for " +
-                s"following brokers ${isr.filterNot(isrInfo.map(_.id).contains).mkString(",")}")
-              new MetadataResponse.PartitionMetadata(Errors.REPLICA_NOT_AVAILABLE, partitionId, leader,
-                replicaInfo.asJava, isrInfo.asJava, offlineReplicaInfo.asJava)
+                s"following brokers ${isr.asScala.filterNot(filteredIsr.contains).mkString(",")}")
+              Errors.REPLICA_NOT_AVAILABLE
             } else {
-              new MetadataResponse.PartitionMetadata(Errors.NONE, partitionId, leader, replicaInfo.asJava,
-                isrInfo.asJava, offlineReplicaInfo.asJava)
+              Errors.NONE
             }
+
+            new MetadataResponsePartition()
+              .setErrorCode(error.code)
+              .setPartitionIndex(partitionId.toInt)
+              .setLeaderId(maybeLeader.map(_.id()).getOrElse(MetadataResponse.NO_LEADER_ID))
+              .setLeaderEpoch(leaderEpoch)
+              .setReplicaNodes(filteredReplicas)
+              .setIsrNodes(filteredIsr)
+              .setOfflineReplicas(offlineReplicas)
         }
       }
     }
   }
 
-  def getAliveEndpoint(brokerId: Int, listenerName: ListenerName): Option[Node] =
-    aliveNodes.get(brokerId).map { nodeMap =>
-      nodeMap.getOrElse(listenerName,
-        throw new BrokerEndPointNotAvailableException(s"Broker `$brokerId` does not have listener with name `$listenerName`"))
-    }
+  /**
+   * Check whether a broker is alive and has a registered listener matching the provided name.
+   * This method was added to avoid unnecessary allocations in [[maybeFilterAliveReplicas]], which is
+   * a hotspot in metadata handling.
+   */
+  private def hasAliveEndpoint(snapshot: MetadataSnapshot, brokerId: Int, listenerName: ListenerName): Boolean = {
+    snapshot.aliveNodes.get(brokerId).exists(_.contains(listenerName))
+  }
+
+  /**
+   * Get the endpoint matching the provided listener if the broker is alive. Note that listeners can
+   * be added dynamically, so a broker with a missing listener could be a transient error.
+   *
+   * @return None if broker is not alive or if the broker does not have a listener named `listenerName`.
+   */
+  private def getAliveEndpoint(snapshot: MetadataSnapshot, brokerId: Int, listenerName: ListenerName): Option[Node] = {
+    snapshot.aliveNodes.get(brokerId).flatMap(_.get(listenerName))
+  }
 
   // errorUnavailableEndpoints exists to support v0 MetadataResponses
-  def getTopicMetadata(topics: Set[String], listenerName: ListenerName, errorUnavailableEndpoints: Boolean = false): Seq[MetadataResponse.TopicMetadata] = {
-    inReadLock(partitionMetadataLock) {
-      topics.toSeq.flatMap { topic =>
-        getPartitionMetadata(topic, listenerName, errorUnavailableEndpoints).map { partitionMetadata =>
-          new MetadataResponse.TopicMetadata(Errors.NONE, topic, Topic.isInternal(topic), partitionMetadata.toBuffer.asJava)
-        }
+  def getTopicMetadata(topics: Set[String],
+                       listenerName: ListenerName,
+                       errorUnavailableEndpoints: Boolean = false,
+                       errorUnavailableListeners: Boolean = false): Seq[MetadataResponseTopic] = {
+    val snapshot = metadataSnapshot
+    topics.toSeq.flatMap { topic =>
+      getPartitionMetadata(snapshot, topic, listenerName, errorUnavailableEndpoints, errorUnavailableListeners).map { partitionMetadata =>
+        new MetadataResponseTopic()
+          .setErrorCode(Errors.NONE.code)
+          .setName(topic)
+          .setIsInternal(Topic.isInternal(topic))
+          .setPartitions(partitionMetadata.toBuffer.asJava)
       }
     }
   }
 
   def getAllTopics(): Set[String] = {
-    inReadLock(partitionMetadataLock) {
-      cache.keySet.toSet
-    }
+    getAllTopics(metadataSnapshot)
+  }
+
+  def getAllPartitions(): Set[TopicPartition] = {
+    metadataSnapshot.partitionStates.flatMap { case (topicName, partitionsAndStates) =>
+      partitionsAndStates.keys.map(partitionId => new TopicPartition(topicName, partitionId.toInt))
+    }.toSet
+  }
+
+  private def getAllTopics(snapshot: MetadataSnapshot): Set[String] = {
+    snapshot.partitionStates.keySet
+  }
+
+  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
   }
 
   def getNonExistingTopics(topics: Set[String]): Set[String] = {
-    inReadLock(partitionMetadataLock) {
-      topics -- cache.keySet
-    }
+    topics.diff(metadataSnapshot.partitionStates.keySet)
   }
 
-  def isBrokerAlive(brokerId: Int): Boolean = {
-    inReadLock(partitionMetadataLock) {
-      aliveBrokers.contains(brokerId)
-    }
+  def getAliveBroker(brokerId: Int): Option[Broker] = {
+    metadataSnapshot.aliveBrokers.get(brokerId)
   }
 
   def getAliveBrokers: Seq[Broker] = {
-    inReadLock(partitionMetadataLock) {
-      aliveBrokers.values.toBuffer
-    }
+    metadataSnapshot.aliveBrokers.values.toBuffer
   }
 
-  private def addOrUpdatePartitionInfo(topic: String,
+  private def addOrUpdatePartitionInfo(partitionStates: mutable.AnyRefMap[String, mutable.LongMap[UpdateMetadataPartitionState]],
+                                       topic: String,
                                        partitionId: Int,
-                                       stateInfo: UpdateMetadataRequest.PartitionState) {
-    inWriteLock(partitionMetadataLock) {
-      val infos = cache.getOrElseUpdate(topic, mutable.Map())
-      infos(partitionId) = stateInfo
-    }
+                                       stateInfo: UpdateMetadataPartitionState): Unit = {
+    val infos = partitionStates.getOrElseUpdate(topic, mutable.LongMap.empty)
+    infos(partitionId) = stateInfo
   }
 
-  def getPartitionInfo(topic: String, partitionId: Int): Option[UpdateMetadataRequest.PartitionState] = {
-    inReadLock(partitionMetadataLock) {
-      cache.get(topic).flatMap(_.get(partitionId))
-    }
+  def getPartitionInfo(topic: String, partitionId: Int): Option[UpdateMetadataPartitionState] = {
+    metadataSnapshot.partitionStates.get(topic).flatMap(_.get(partitionId))
+  }
+
+  def numPartitions(topic: String): Option[Int] = {
+    metadataSnapshot.partitionStates.get(topic).map(_.size)
   }
 
   // if the leader is not known, return None;
   // if the leader is known and corresponding node is available, return Some(node)
   // if the leader is known but corresponding node with the listener name is not available, return Some(NO_NODE)
   def getPartitionLeaderEndpoint(topic: String, partitionId: Int, listenerName: ListenerName): Option[Node] = {
-    inReadLock(partitionMetadataLock) {
-      cache.get(topic).flatMap(_.get(partitionId)) map { partitionInfo =>
-        val leaderId = partitionInfo.basePartitionState.leader
+    val snapshot = metadataSnapshot
+    snapshot.partitionStates.get(topic).flatMap(_.get(partitionId)) map { partitionInfo =>
+      val leaderId = partitionInfo.leader
 
-        aliveNodes.get(leaderId) match {
-          case Some(nodeMap) =>
-            nodeMap.getOrElse(listenerName, Node.noNode)
-          case None =>
-            Node.noNode
-        }
+      snapshot.aliveNodes.get(leaderId) match {
+        case Some(nodeMap) =>
+          nodeMap.getOrElse(listenerName, Node.noNode)
+        case None =>
+          Node.noNode
       }
     }
   }
 
-  def getControllerId: Option[Int] = controllerId
+  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 = {
+    val snapshot = metadataSnapshot
+    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.leader != LeaderAndIsr.LeaderDuringDelete }
+      .map { case (tp, state) =>
+        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).toBuffer.asJava,
+      partitions.toBuffer.asJava,
+      unauthorizedTopics, internalTopics,
+      snapshot.controllerId.map(id => node(id)).orNull)
+  }
 
   // This method returns the deleted TopicPartitions received from UpdateMetadataRequest
-  def updateCache(correlationId: Int, updateMetadataRequest: UpdateMetadataRequest): Seq[TopicPartition] = {
+  def updateMetadata(correlationId: Int, updateMetadataRequest: UpdateMetadataRequest): Seq[TopicPartition] = {
     inWriteLock(partitionMetadataLock) {
-      controllerId = updateMetadataRequest.controllerId match {
+
+      val aliveBrokers = new mutable.LongMap[Broker](metadataSnapshot.aliveBrokers.size)
+      val aliveNodes = new mutable.LongMap[collection.Map[ListenerName, Node]](metadataSnapshot.aliveNodes.size)
+      val controllerIdOpt = updateMetadataRequest.controllerId match {
           case id if id < 0 => None
           case id => Some(id)
         }
-      aliveNodes.clear()
-      aliveBrokers.clear()
-      updateMetadataRequest.liveBrokers.asScala.foreach { broker =>
+
+      updateMetadataRequest.liveBrokers.forEach { broker =>
         // `aliveNodes` is a hot path for metadata requests for large clusters, so we use java.util.HashMap which
         // is a bit faster than scala.collection.mutable.HashMap. When we drop support for Scala 2.10, we could
         // 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.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
       }
+      aliveNodes.get(brokerId).foreach { listenerMap =>
+        val listeners = listenerMap.keySet
+        if (!aliveNodes.values.forall(_.keySet == listeners))
+          error(s"Listeners are not identical across brokers: $aliveNodes")
+      }
 
       val deletedPartitions = new mutable.ArrayBuffer[TopicPartition]
-      updateMetadataRequest.partitionStates.asScala.foreach { case (tp, info) =>
+      if (!updateMetadataRequest.partitionStates.iterator.hasNext) {
+        metadataSnapshot = MetadataSnapshot(metadataSnapshot.partitionStates, controllerIdOpt, 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[UpdateMetadataPartitionState]](metadataSnapshot.partitionStates.size)
+        metadataSnapshot.partitionStates.forKeyValue { (topic, oldPartitionStates) =>
+          val copy = new mutable.LongMap[UpdateMetadataPartitionState](oldPartitionStates.size)
+          copy ++= oldPartitionStates
+          partitionStates(topic) = copy
+        }
+
+        val traceEnabled = stateChangeLogger.isTraceEnabled
         val controllerId = updateMetadataRequest.controllerId
         val controllerEpoch = updateMetadataRequest.controllerEpoch
-        if (info.basePartitionState.leader == LeaderAndIsr.LeaderDuringDelete) {
-          removePartitionInfo(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")
-          deletedPartitions += tp
-        } else {
-          addOrUpdatePartitionInfo(tp.topic, tp.partition, info)
-          stateChangeLogger.trace(s"Cached leader info $info for partition $tp in response to " +
-            s"UpdateMetadata request sent by controller $controllerId epoch $controllerEpoch with correlation id $correlationId")
+        val newStates = updateMetadataRequest.partitionStates.asScala
+        newStates.foreach { state =>
+          // per-partition logging here can be very expensive due going through all partitions in the cluster
+          val tp = new TopicPartition(state.topicName, state.partitionIndex)
+          if (state.leader == LeaderAndIsr.LeaderDuringDelete) {
+            removePartitionInfo(partitionStates, tp.topic, tp.partition)
+            if (traceEnabled)
+              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")
+            deletedPartitions += tp
+          } else {
+            addOrUpdatePartitionInfo(partitionStates, tp.topic, tp.partition, state)
+            if (traceEnabled)
+              stateChangeLogger.trace(s"Cached leader info $state for partition $tp in response to " +
+                s"UpdateMetadata request sent by controller $controllerId epoch $controllerEpoch with correlation id $correlationId")
+          }
         }
+        val cachedPartitionsCount = newStates.size - deletedPartitions.size
+        stateChangeLogger.info(s"Add $cachedPartitionsCount partitions and deleted ${deletedPartitions.size} partitions from metadata cache " +
+          s"in response to UpdateMetadata request sent by controller $controllerId epoch $controllerEpoch with correlation id $correlationId")
+
+        metadataSnapshot = MetadataSnapshot(partitionStates, controllerIdOpt, aliveBrokers, aliveNodes)
       }
       deletedPartitions
     }
   }
 
   def contains(topic: String): Boolean = {
-    inReadLock(partitionMetadataLock) {
-      cache.contains(topic)
-    }
+    metadataSnapshot.partitionStates.contains(topic)
   }
 
   def contains(tp: TopicPartition): Boolean = getPartitionInfo(tp.topic, tp.partition).isDefined
 
-  private def removePartitionInfo(topic: String, partitionId: Int): Boolean = {
-    cache.get(topic).exists { infos =>
+  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) cache.remove(topic)
+      if (infos.isEmpty) partitionStates.remove(topic)
       true
     }
   }
 
+  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/QuotaFactory.scala b/core/src/main/scala/kafka/server/QuotaFactory.scala
index 01441b573238c..9d4443b904ea8 100644
--- a/core/src/main/scala/kafka/server/QuotaFactory.scala
+++ b/core/src/main/scala/kafka/server/QuotaFactory.scala
@@ -20,46 +20,71 @@ import kafka.server.QuotaType._
 import kafka.utils.Logging
 import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.metrics.Metrics
+import org.apache.kafka.server.quota.ClientQuotaCallback
 import org.apache.kafka.common.utils.Time
+import org.apache.kafka.server.quota.ClientQuotaType
 
 object QuotaType  {
   case object Fetch extends QuotaType
   case object Produce extends QuotaType
   case object Request extends QuotaType
+  case object ControllerMutation extends QuotaType
   case object LeaderReplication extends QuotaType
   case object FollowerReplication extends QuotaType
   case object AlterLogDirsReplication extends QuotaType
+
+  def toClientQuotaType(quotaType: QuotaType): ClientQuotaType = {
+    quotaType match {
+      case QuotaType.Fetch => ClientQuotaType.FETCH
+      case QuotaType.Produce => ClientQuotaType.PRODUCE
+      case QuotaType.Request => ClientQuotaType.REQUEST
+      case QuotaType.ControllerMutation => ClientQuotaType.CONTROLLER_MUTATION
+      case _ => throw new IllegalArgumentException(s"Not a client quota type: $quotaType")
+    }
+  }
 }
+
 sealed trait QuotaType
 
 object QuotaFactory extends Logging {
 
   object UnboundedQuota extends ReplicaQuota {
     override def isThrottled(topicPartition: TopicPartition): Boolean = false
-    override def isQuotaExceeded(): Boolean = false
+    override def isQuotaExceeded: Boolean = false
+    def record(value: Long): Unit = ()
   }
 
   case class QuotaManagers(fetch: ClientQuotaManager,
                            produce: ClientQuotaManager,
                            request: ClientRequestQuotaManager,
+                           controllerMutation: ControllerMutationQuotaManager,
                            leader: ReplicationQuotaManager,
                            follower: ReplicationQuotaManager,
-                           alterLogDirs: ReplicationQuotaManager) {
-    def shutdown() {
-      fetch.shutdown
-      produce.shutdown
-      request.shutdown
+                           alterLogDirs: ReplicationQuotaManager,
+                           clientQuotaCallback: Option[ClientQuotaCallback]) {
+    def shutdown(): Unit = {
+      fetch.shutdown()
+      produce.shutdown()
+      request.shutdown()
+      controllerMutation.shutdown()
+      clientQuotaCallback.foreach(_.close())
     }
   }
 
   def instantiate(cfg: KafkaConfig, metrics: Metrics, time: Time, threadNamePrefix: String): QuotaManagers = {
+
+    val clientQuotaCallback = Option(cfg.getConfiguredInstance(KafkaConfig.ClientQuotaCallbackClassProp,
+      classOf[ClientQuotaCallback]))
     QuotaManagers(
-      new ClientQuotaManager(clientFetchConfig(cfg), metrics, Fetch, time, threadNamePrefix),
-      new ClientQuotaManager(clientProduceConfig(cfg), metrics, Produce, time, threadNamePrefix),
-      new ClientRequestQuotaManager(clientRequestConfig(cfg), metrics, time, threadNamePrefix),
+      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 ControllerMutationQuotaManager(clientControllerMutationConfig(cfg), metrics, time,
+        threadNamePrefix, clientQuotaCallback),
       new ReplicationQuotaManager(replicationConfig(cfg), metrics, LeaderReplication, time),
       new ReplicationQuotaManager(replicationConfig(cfg), metrics, FollowerReplication, time),
-      new ReplicationQuotaManager(alterLogDirsReplicationConfig(cfg), metrics, AlterLogDirsReplication, time)
+      new ReplicationQuotaManager(alterLogDirsReplicationConfig(cfg), metrics, AlterLogDirsReplication, time),
+      clientQuotaCallback
     )
   }
 
@@ -67,7 +92,7 @@ object QuotaFactory extends Logging {
     if (cfg.producerQuotaBytesPerSecondDefault != Long.MaxValue)
       warn(s"${KafkaConfig.ProducerQuotaBytesPerSecondDefaultProp} has been deprecated in 0.11.0.0 and will be removed in a future release. Use dynamic quota defaults instead.")
     ClientQuotaManagerConfig(
-      quotaBytesPerSecondDefault = cfg.producerQuotaBytesPerSecondDefault,
+      quotaDefault = cfg.producerQuotaBytesPerSecondDefault,
       numQuotaSamples = cfg.numQuotaSamples,
       quotaWindowSizeSeconds = cfg.quotaWindowSizeSeconds
     )
@@ -77,7 +102,7 @@ object QuotaFactory extends Logging {
     if (cfg.consumerQuotaBytesPerSecondDefault != Long.MaxValue)
       warn(s"${KafkaConfig.ConsumerQuotaBytesPerSecondDefaultProp} has been deprecated in 0.11.0.0 and will be removed in a future release. Use dynamic quota defaults instead.")
     ClientQuotaManagerConfig(
-      quotaBytesPerSecondDefault = cfg.consumerQuotaBytesPerSecondDefault,
+      quotaDefault = cfg.consumerQuotaBytesPerSecondDefault,
       numQuotaSamples = cfg.numQuotaSamples,
       quotaWindowSizeSeconds = cfg.quotaWindowSizeSeconds
     )
@@ -90,6 +115,13 @@ object QuotaFactory extends Logging {
     )
   }
 
+  def clientControllerMutationConfig(cfg: KafkaConfig): ClientQuotaManagerConfig = {
+    ClientQuotaManagerConfig(
+      numQuotaSamples = cfg.numControllerQuotaSamples,
+      quotaWindowSizeSeconds = cfg.controllerQuotaWindowSizeSeconds
+    )
+  }
+
   def replicationConfig(cfg: KafkaConfig): ReplicationQuotaManagerConfig = {
     ReplicationQuotaManagerConfig(
       numQuotaSamples = cfg.numReplicationQuotaSamples,
diff --git a/core/src/main/scala/kafka/server/ReplicaAlterLogDirsManager.scala b/core/src/main/scala/kafka/server/ReplicaAlterLogDirsManager.scala
index bded1256d2606..b45a76620c74c 100644
--- a/core/src/main/scala/kafka/server/ReplicaAlterLogDirsManager.scala
+++ b/core/src/main/scala/kafka/server/ReplicaAlterLogDirsManager.scala
@@ -17,22 +17,40 @@
 
 package kafka.server
 
-
 import kafka.cluster.BrokerEndPoint
+import org.apache.kafka.common.TopicPartition
 
 class ReplicaAlterLogDirsManager(brokerConfig: KafkaConfig,
                                  replicaManager: ReplicaManager,
                                  quotaManager: ReplicationQuotaManager,
                                  brokerTopicStats: BrokerTopicStats)
-  extends AbstractFetcherManager(s"ReplicaAlterLogDirsManager on broker ${brokerConfig.brokerId}",
-    "ReplicaAlterLogDirs", brokerConfig.getNumReplicaAlterLogDirsThreads) {
+  extends AbstractFetcherManager[ReplicaAlterLogDirsThread](
+    name = s"ReplicaAlterLogDirsManager on broker ${brokerConfig.brokerId}",
+    clientId = "ReplicaAlterLogDirs",
+    numFetchers = brokerConfig.getNumReplicaAlterLogDirsThreads) {
 
-  override def createFetcherThread(fetcherId: Int, sourceBroker: BrokerEndPoint): AbstractFetcherThread = {
+  override def createFetcherThread(fetcherId: Int, sourceBroker: BrokerEndPoint): ReplicaAlterLogDirsThread = {
     val threadName = s"ReplicaAlterLogDirsThread-$fetcherId"
-    new ReplicaAlterLogDirsThread(threadName, sourceBroker, brokerConfig, replicaManager, quotaManager, brokerTopicStats)
+    new ReplicaAlterLogDirsThread(threadName, sourceBroker, brokerConfig, failedPartitions, replicaManager,
+      quotaManager, brokerTopicStats)
+  }
+
+  override protected def addPartitionsToFetcherThread(fetcherThread: ReplicaAlterLogDirsThread,
+                                                      initialOffsetAndEpochs: collection.Map[TopicPartition, InitialFetchState]): 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() {
+  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 48c83d4383580..3315f30ee130c 100644
--- a/core/src/main/scala/kafka/server/ReplicaAlterLogDirsThread.scala
+++ b/core/src/main/scala/kafka/server/ReplicaAlterLogDirsThread.scala
@@ -17,54 +17,69 @@
 
 package kafka.server
 
-import java.nio.ByteBuffer
 import java.util
+import java.util.Optional
 
-import AbstractFetcherThread.ResultWithPartitions
-import kafka.cluster.BrokerEndPoint
-import org.apache.kafka.common.TopicPartition
-import org.apache.kafka.common.requests.EpochEndOffset._
-import org.apache.kafka.common.requests.{EpochEndOffset, FetchResponse, FetchRequest => JFetchRequest}
-import ReplicaAlterLogDirsThread.FetchRequest
-import ReplicaAlterLogDirsThread.PartitionData
 import kafka.api.Request
+import kafka.cluster.BrokerEndPoint
+import kafka.log.{LeaderOffsetIncremented, LogAppendInfo}
+import kafka.server.AbstractFetcherThread.ReplicaFetch
+import kafka.server.AbstractFetcherThread.ResultWithPartitions
 import kafka.server.QuotaFactory.UnboundedQuota
-import kafka.server.epoch.LeaderEpochCache
+import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.errors.KafkaStorageException
+import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData.EpochEndOffset
 import org.apache.kafka.common.protocol.{ApiKeys, Errors}
-import org.apache.kafka.common.record.{FileRecords, MemoryRecords}
+import org.apache.kafka.common.record.Records
+import org.apache.kafka.common.requests.FetchResponse.PartitionData
+import org.apache.kafka.common.requests.{FetchRequest, FetchResponse}
+import org.apache.kafka.common.requests.OffsetsForLeaderEpochResponse.UNDEFINED_EPOCH
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 import scala.collection.{Map, Seq, Set, mutable}
-
+import scala.compat.java8.OptionConverters._
 
 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,
-                                includeLogTruncation = true) {
-
-  type REQ = FetchRequest
-  type PD = PartitionData
+                                brokerTopicStats) {
 
   private val replicaId = brokerConfig.brokerId
   private val maxBytes = brokerConfig.replicaFetchResponseMaxBytes
   private val fetchSize = brokerConfig.replicaFetchMaxBytes
+  private var inProgressPartition: Option[TopicPartition] = None
 
-  private def epochCacheOpt(tp: TopicPartition): Option[LeaderEpochCache] =  replicaMgr.getReplica(tp).map(_.epochs.get)
+  override protected def latestEpoch(topicPartition: TopicPartition): Option[Int] = {
+    replicaMgr.futureLocalLogOrException(topicPartition).latestEpoch
+  }
+
+  override protected def logStartOffset(topicPartition: TopicPartition): Long = {
+    replicaMgr.futureLocalLogOrException(topicPartition).logStartOffset
+  }
+
+  override protected def logEndOffset(topicPartition: TopicPartition): Long = {
+    replicaMgr.futureLocalLogOrException(topicPartition).logEndOffset
+  }
+
+  override protected def endOffsetForEpoch(topicPartition: TopicPartition, epoch: Int): Option[OffsetAndEpoch] = {
+    replicaMgr.futureLocalLogOrException(topicPartition).endOffsetForEpoch(epoch)
+  }
 
-  def fetch(fetchRequest: FetchRequest): Seq[(TopicPartition, PartitionData)] = {
-    var partitionData: Seq[(TopicPartition, FetchResponse.PartitionData)] = null
-    val request = fetchRequest.underlying.build()
+  def fetchFromLeader(fetchRequest: FetchRequest.Builder): Map[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)
@@ -78,189 +93,205 @@ 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(",")}")
 
-    partitionData.map { case (key, value) =>
-      key -> new PartitionData(value)
-    }
+    partitionData.toMap
   }
 
   // process fetched data
-  def processPartitionData(topicPartition: TopicPartition, fetchOffset: Long, partitionData: PartitionData) {
-    val futureReplica = replicaMgr.getReplicaOrException(topicPartition, Request.FutureLocalReplicaId)
-    val partition = replicaMgr.getPartition(topicPartition).get
-    val records = partitionData.toRecords
-
-    if (fetchOffset != futureReplica.logEndOffset.messageOffset)
+  override def processPartitionData(topicPartition: TopicPartition,
+                                    fetchOffset: Long,
+                                    partitionData: PartitionData[Records]): Option[LogAppendInfo] = {
+    val partition = replicaMgr.nonOfflinePartition(topicPartition).get
+    val futureLog = partition.futureLocalLogOrException
+    val records = toMemoryRecords(partitionData.records)
+
+    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.messageOffset))
+        topicPartition, fetchOffset, futureLog.logEndOffset))
 
-    // Append the leader's messages to the log
-    partition.appendRecordsToFutureReplica(records)
-    futureReplica.highWatermark = new LogOffsetMetadata(partitionData.highWatermark)
-    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, LeaderOffsetIncremented)
 
     if (partition.maybeReplaceCurrentWithFutureReplica())
       removePartitions(Set(topicPartition))
 
     quota.record(records.sizeInBytes)
+    logAppendInfo
   }
 
-  def handleOffsetOutOfRange(topicPartition: TopicPartition): Long = {
-    val futureReplica = replicaMgr.getReplicaOrException(topicPartition, Request.FutureLocalReplicaId)
-    val currentReplica = replicaMgr.getReplicaOrException(topicPartition)
-    val partition = replicaMgr.getPartition(topicPartition).get
-    val logEndOffset: Long = currentReplica.logEndOffset.messageOffset
-
-    if (logEndOffset < futureReplica.logEndOffset.messageOffset) {
-      warn("Future replica for partition %s reset its fetch offset from %d to current replica's latest offset %d"
-        .format(topicPartition, futureReplica.logEndOffset.messageOffset, logEndOffset))
-      partition.truncateTo(logEndOffset, isFuture = true)
-      logEndOffset
-    } else {
-      val currentReplicaStartOffset: Long = currentReplica.logStartOffset
-      warn("Future replica for partition %s reset its fetch offset from %d to current replica's start offset %d"
-        .format(topicPartition, futureReplica.logEndOffset.messageOffset, currentReplicaStartOffset))
-      val offsetToFetch = Math.max(currentReplicaStartOffset, futureReplica.logEndOffset.messageOffset)
-      // Only truncate the log when current replica's log start offset is greater than future replica's log end offset.
-      if (currentReplicaStartOffset > futureReplica.logEndOffset.messageOffset)
-        partition.truncateFullyAndStartAt(currentReplicaStartOffset, isFuture = true)
-      offsetToFetch
+  override def addPartitions(initialFetchStates: Map[TopicPartition, InitialFetchState]): 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()
     }
   }
 
-  def handlePartitionsWithErrors(partitions: Iterable[TopicPartition]) {
-    if (partitions.nonEmpty)
-      delayPartitions(partitions, brokerConfig.replicaFetchBackoffMs.toLong)
+  override protected def fetchEarliestOffsetFromLeader(topicPartition: TopicPartition, leaderEpoch: Int): Long = {
+    val partition = replicaMgr.getPartitionOrException(topicPartition)
+    partition.localLogOrException.logStartOffset
   }
 
-  def buildLeaderEpochRequest(allPartitions: Seq[(TopicPartition, PartitionFetchState)]): ResultWithPartitions[Map[TopicPartition, Int]] = {
-    val partitionEpochOpts = allPartitions
-      .filter { case (_, state) => state.isTruncatingLog }
-      .map { case (tp, _) => tp -> epochCacheOpt(tp) }.toMap
-
-    val (partitionsWithEpoch, partitionsWithoutEpoch) = partitionEpochOpts.partition { case (tp, epochCacheOpt) => epochCacheOpt.nonEmpty }
-
-    val result = partitionsWithEpoch.map { case (tp, epochCacheOpt) => tp -> epochCacheOpt.get.latestEpoch() }
-    ResultWithPartitions(result, partitionsWithoutEpoch.keys.toSet)
+  override protected def fetchLatestOffsetFromLeader(topicPartition: TopicPartition, leaderEpoch: Int): Long = {
+    val partition = replicaMgr.getPartitionOrException(topicPartition)
+    partition.localLogOrException.logEndOffset
   }
 
-  def fetchEpochsFromLeader(partitions: Map[TopicPartition, Int]): Map[TopicPartition, EpochEndOffset] = {
-    partitions.map { case (tp, epoch) =>
+  /**
+   * Fetches offset for leader epoch from local replica for each given topic partitions
+   * @param partitions map of topic partition -> leader epoch of the future replica
+   * @return map of topic partition -> end offset for a requested leader epoch
+   */
+  override def fetchEpochEndOffsets(partitions: Map[TopicPartition, EpochData]): Map[TopicPartition, EpochEndOffset] = {
+    partitions.map { case (tp, epochData) =>
       try {
-        tp -> new EpochEndOffset(Errors.NONE, replicaMgr.getReplicaOrException(tp).epochs.get.endOffsetFor(epoch))
+        val endOffset = if (epochData.leaderEpoch == UNDEFINED_EPOCH) {
+          new EpochEndOffset()
+            .setPartition(tp.partition)
+            .setErrorCode(Errors.NONE.code)
+        } else {
+          val partition = replicaMgr.getPartitionOrException(tp)
+          partition.lastOffsetForLeaderEpoch(
+            currentLeaderEpoch = epochData.currentLeaderEpoch,
+            leaderEpoch = epochData.leaderEpoch,
+            fetchOnlyFromLeader = false)
+        }
+        tp -> endOffset
       } catch {
         case t: Throwable =>
           warn(s"Error when getting EpochEndOffset for $tp", t)
-          tp -> new EpochEndOffset(Errors.forException(t), UNDEFINED_EPOCH_OFFSET)
+          tp -> new EpochEndOffset()
+            .setPartition(tp.partition)
+            .setErrorCode(Errors.forException(t).code)
       }
     }
   }
 
-  def maybeTruncate(fetchedEpochs: Map[TopicPartition, EpochEndOffset]): ResultWithPartitions[Map[TopicPartition, Long]] = {
-    val fetchOffsets = scala.collection.mutable.HashMap.empty[TopicPartition, Long]
-    val partitionsWithError = mutable.Set[TopicPartition]()
-
-    fetchedEpochs.foreach { case (topicPartition, epochOffset) =>
-      try {
-        val futureReplica = replicaMgr.getReplicaOrException(topicPartition, Request.FutureLocalReplicaId)
-        val partition = replicaMgr.getPartition(topicPartition).get
+  override protected val isOffsetForLeaderEpochSupported: Boolean = true
+
+  override protected val isTruncationOnFetchSupported: Boolean = false
+
+  /**
+   * Truncate the log for each partition based on current replica's returned epoch and offset.
+   *
+   * The logic for finding the truncation offset is the same as in ReplicaFetcherThread
+   * and mainly implemented in AbstractFetcherThread.getOffsetTruncationState. One difference is
+   * that the initial fetch offset for topic partition could be set to the truncation offset of
+   * the current replica if that replica truncates. Otherwise, it is high watermark as in ReplicaFetcherThread.
+   *
+   * The reason we have to follow the leader epoch approach for truncating a future replica is to
+   * cover the case where a future replica is offline when the current replica truncates and
+   * re-replicates offsets that may have already been copied to the future replica. In that case,
+   * the future replica may miss "mark for truncation" event and must use the offset for leader epoch
+   * exchange with the current replica to truncate to the largest common log prefix for the topic partition
+   */
+  override def truncate(topicPartition: TopicPartition, truncationState: OffsetTruncationState): Unit = {
+    val partition = replicaMgr.getPartitionOrException(topicPartition)
+    partition.truncateTo(truncationState.offset, isFuture = true)
+  }
 
-        if (epochOffset.hasError) {
-          info(s"Retrying leaderEpoch request for partition $topicPartition as the current replica reported an error: ${epochOffset.error}")
-          partitionsWithError += topicPartition
-        } else {
-          val fetchOffset =
-            if (epochOffset.endOffset == UNDEFINED_EPOCH_OFFSET)
-              partitionStates.stateValue(topicPartition).fetchOffset
-            else if (epochOffset.endOffset >= futureReplica.logEndOffset.messageOffset)
-              futureReplica.logEndOffset.messageOffset
-            else
-              epochOffset.endOffset
-
-          partition.truncateTo(fetchOffset, isFuture = true)
-          fetchOffsets.put(topicPartition, fetchOffset)
-        }
-      } catch {
-        case e: KafkaStorageException =>
-          info(s"Failed to truncate $topicPartition", e)
-          partitionsWithError += topicPartition
-      }
-    }
-    ResultWithPartitions(fetchOffsets, partitionsWithError)
+  override protected def truncateFullyAndStartAt(topicPartition: TopicPartition, offset: Long): Unit = {
+    val partition = replicaMgr.getPartitionOrException(topicPartition)
+    partition.truncateFullyAndStartAt(offset, isFuture = true)
   }
 
-  def buildFetchRequest(partitionMap: Seq[(TopicPartition, PartitionFetchState)]): ResultWithPartitions[FetchRequest] = {
-    // Only include replica in the fetch request if it is not throttled.
-    val maxPartitionOpt = partitionMap.filter { case (topicPartition, partitionFetchState) =>
-      partitionFetchState.isReadyForFetch && !quota.isQuotaExceeded
+  private def nextReadyPartition(partitionMap: Map[TopicPartition, PartitionFetchState]): Option[(TopicPartition, PartitionFetchState)] = {
+    partitionMap.filter { case (_, partitionFetchState) =>
+      partitionFetchState.isReadyForFetch
     }.reduceLeftOption { (left, right) =>
-      if ((left._1.topic > right._1.topic()) || (left._1.topic == right._1.topic() && left._1.partition() >= right._1.partition()))
+      if ((left._1.topic < right._1.topic) || (left._1.topic == right._1.topic && left._1.partition < right._1.partition))
         left
       else
         right
     }
+  }
 
-    // Only move one replica at a time to increase its catch-up rate and thus reduce the time spent on moving any given replica
-    // Replicas are ordered by their TopicPartition
-    val requestMap = new util.LinkedHashMap[TopicPartition, JFetchRequest.PartitionData]
-    val partitionsWithError = mutable.Set[TopicPartition]()
+  private def selectPartitionToFetch(partitionMap: Map[TopicPartition, PartitionFetchState]): Option[(TopicPartition, PartitionFetchState)] = {
+    // Only move one partition at a time to increase its catch-up rate and thus reduce the time spent on
+    // moving any given replica. Replicas are selected in ascending order (lexicographically by topic) from the
+    // partitions that are ready to fetch. Once selected, we will continue fetching the same partition until it
+    // becomes unavailable or is removed.
 
-    if (maxPartitionOpt.nonEmpty) {
-      val (topicPartition, partitionFetchState) = maxPartitionOpt.get
-      try {
-        val logStartOffset = replicaMgr.getReplicaOrException(topicPartition, Request.FutureLocalReplicaId).logStartOffset
-        requestMap.put(topicPartition, new JFetchRequest.PartitionData(partitionFetchState.fetchOffset, logStartOffset, fetchSize))
-      } catch {
-        case e: KafkaStorageException =>
-          partitionsWithError += topicPartition
+    inProgressPartition.foreach { tp =>
+      val fetchStateOpt = partitionMap.get(tp)
+      fetchStateOpt.filter(_.isReadyForFetch).foreach { fetchState =>
+        return Some((tp, fetchState))
       }
     }
-    // Set maxWait and minBytes to 0 because the response should return immediately if
-    // the future log has caught up with the current log of the partition
-    val requestBuilder = JFetchRequest.Builder.forReplica(ApiKeys.FETCH.latestVersion, replicaId, 0, 0, requestMap).setMaxBytes(maxBytes)
-    ResultWithPartitions(new FetchRequest(requestBuilder), partitionsWithError)
-  }
-}
 
-object ReplicaAlterLogDirsThread {
+    inProgressPartition = None
 
-  private[server] class FetchRequest(val underlying: JFetchRequest.Builder) extends AbstractFetcherThread.FetchRequest {
-    def isEmpty: Boolean = underlying.fetchData.isEmpty
-    def offset(topicPartition: TopicPartition): Long = underlying.fetchData.asScala(topicPartition).fetchOffset
-    override def toString = underlying.toString
+    val nextPartitionOpt = nextReadyPartition(partitionMap)
+    nextPartitionOpt.foreach { case (tp, fetchState) =>
+      inProgressPartition = Some(tp)
+      info(s"Beginning/resuming copy of partition $tp from offset ${fetchState.fetchOffset}. " +
+        s"Including this partition, there are ${partitionMap.size} remaining partitions to copy by this thread.")
+    }
+    nextPartitionOpt
   }
 
-  private[server] class PartitionData(val underlying: FetchResponse.PartitionData) extends AbstractFetcherThread.PartitionData {
-
-    def error = underlying.error
+  private def buildFetchForPartition(tp: TopicPartition, fetchState: PartitionFetchState): ResultWithPartitions[Option[ReplicaFetch]] = {
+    val requestMap = new util.LinkedHashMap[TopicPartition, FetchRequest.PartitionData]
+    val partitionsWithError = mutable.Set[TopicPartition]()
 
-    def toRecords: MemoryRecords = {
-      if (underlying.records == MemoryRecords.EMPTY)
-        underlying.records.asInstanceOf[MemoryRecords]
-      else {
-        val buffer = ByteBuffer.allocate(underlying.records.sizeInBytes())
-        underlying.records.asInstanceOf[FileRecords].readInto(buffer, 0)
-        MemoryRecords.readableRecords(buffer)
-      }
+    try {
+      val logStartOffset = replicaMgr.futureLocalLogOrException(tp).logStartOffset
+      val lastFetchedEpoch = if (isTruncationOnFetchSupported)
+        fetchState.lastFetchedEpoch.map(_.asInstanceOf[Integer]).asJava
+      else
+        Optional.empty[Integer]
+      requestMap.put(tp, new FetchRequest.PartitionData(fetchState.fetchOffset, logStartOffset,
+        fetchSize, Optional.of(fetchState.currentLeaderEpoch), lastFetchedEpoch))
+    } catch {
+      case e: KafkaStorageException =>
+        debug(s"Failed to build fetch for $tp", e)
+        partitionsWithError += tp
     }
 
-    def highWatermark: Long = underlying.highWatermark
+    val fetchRequestOpt = if (requestMap.isEmpty) {
+      None
+    } else {
+      // Set maxWait and minBytes to 0 because the response should return immediately if
+      // the future log has caught up with the current log of the partition
+      val requestBuilder = FetchRequest.Builder.forReplica(ApiKeys.FETCH.latestVersion, replicaId, 0, 0, requestMap).setMaxBytes(maxBytes)
+      Some(ReplicaFetch(requestMap, requestBuilder))
+    }
 
-    def logStartOffset: Long = underlying.logStartOffset
+    ResultWithPartitions(fetchRequestOpt, partitionsWithError)
+  }
 
-    def exception: Option[Throwable] = error match {
-      case Errors.NONE => None
-      case e => Some(e.exception)
+  def buildFetch(partitionMap: Map[TopicPartition, PartitionFetchState]): ResultWithPartitions[Option[ReplicaFetch]] = {
+    // Only include replica in the fetch request if it is not throttled.
+    if (quota.isQuotaExceeded) {
+      ResultWithPartitions(None, Set.empty)
+    } else {
+      selectPartitionToFetch(partitionMap) match {
+        case Some((tp, fetchState)) =>
+          buildFetchForPartition(tp, fetchState)
+        case None =>
+          ResultWithPartitions(None, Set.empty)
+      }
     }
-
-    override def toString = underlying.toString
   }
-}
\ No newline at end of file
+
+}
diff --git a/core/src/main/scala/kafka/server/ReplicaFetcherBlockingSend.scala b/core/src/main/scala/kafka/server/ReplicaFetcherBlockingSend.scala
index 0bf2bd3c9fff7..212daa5ce1a71 100644
--- a/core/src/main/scala/kafka/server/ReplicaFetcherBlockingSend.scala
+++ b/core/src/main/scala/kafka/server/ReplicaFetcherBlockingSend.scala
@@ -26,16 +26,18 @@ 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._
+import scala.jdk.CollectionConverters._
 
 trait BlockingSend {
 
   def sendRequest(requestBuilder: AbstractRequest.Builder[_ <: AbstractRequest]): ClientResponse
 
-  def close()
+  def initiateClose(): Unit
+
+  def close(): Unit
 }
 
 class ReplicaFetcherBlockingSend(sourceBroker: BrokerEndPoint,
@@ -49,15 +51,23 @@ 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,
       brokerConfig,
       brokerConfig.interBrokerListenerName,
       brokerConfig.saslMechanismInterBrokerProtocol,
-      brokerConfig.saslInterBrokerHandshakeRequestEnable
+      time,
+      brokerConfig.saslInterBrokerHandshakeRequestEnable,
+      logContext
     )
+    val reconfigurableChannelBuilder = channelBuilder match {
+      case reconfigurable: Reconfigurable =>
+        brokerConfig.addReconfigurable(reconfigurable)
+        Some(reconfigurable)
+      case _ => None
+    }
     val selector = new Selector(
       NetworkReceive.UNLIMITED,
       brokerConfig.connectionsMaxIdleMs,
@@ -69,7 +79,7 @@ class ReplicaFetcherBlockingSend(sourceBroker: BrokerEndPoint,
       channelBuilder,
       logContext
     )
-    new NetworkClient(
+    val networkClient = new NetworkClient(
       selector,
       new ManualMetadataUpdater(),
       clientId,
@@ -79,14 +89,18 @@ class ReplicaFetcherBlockingSend(sourceBroker: BrokerEndPoint,
       Selectable.USE_DEFAULT_BUFFER_SIZE,
       brokerConfig.replicaSocketReceiveBufferBytes,
       brokerConfig.requestTimeoutMs,
+      brokerConfig.connectionSetupTimeoutMs,
+      brokerConfig.connectionSetupTimeoutMaxMs,
+      ClientDnsLookup.USE_ALL_DNS_IPS,
       time,
       false,
       new ApiVersions,
       logContext
     )
+    (networkClient, reconfigurableChannelBuilder)
   }
 
-  override def sendRequest(requestBuilder: Builder[_ <: AbstractRequest]): ClientResponse =  {
+  override def sendRequest(requestBuilder: Builder[_ <: AbstractRequest]): ClientResponse = {
     try {
       if (!NetworkClientUtils.awaitReady(networkClient, sourceNode, time, socketTimeout))
         throw new SocketTimeoutException(s"Failed to connect within $socketTimeout ms")
@@ -103,6 +117,11 @@ class ReplicaFetcherBlockingSend(sourceBroker: BrokerEndPoint,
     }
   }
 
+  override def initiateClose(): Unit = {
+    reconfigurableChannelBuilder.foreach(brokerConfig.removeReconfigurable)
+    networkClient.initiateClose()
+  }
+
   def close(): Unit = {
     networkClient.close()
   }
diff --git a/core/src/main/scala/kafka/server/ReplicaFetcherManager.scala b/core/src/main/scala/kafka/server/ReplicaFetcherManager.scala
index a50b0bbb82b5e..d547e1b5d769b 100644
--- a/core/src/main/scala/kafka/server/ReplicaFetcherManager.scala
+++ b/core/src/main/scala/kafka/server/ReplicaFetcherManager.scala
@@ -21,18 +21,25 @@ import kafka.cluster.BrokerEndPoint
 import org.apache.kafka.common.metrics.Metrics
 import org.apache.kafka.common.utils.Time
 
-class ReplicaFetcherManager(brokerConfig: KafkaConfig, protected val replicaManager: ReplicaManager, metrics: Metrics,
-                            time: Time, threadNamePrefix: Option[String] = None, quotaManager: ReplicationQuotaManager)
-      extends AbstractFetcherManager("ReplicaFetcherManager on broker " + brokerConfig.brokerId,
-        "Replica", brokerConfig.numReplicaFetchers) {
+class ReplicaFetcherManager(brokerConfig: KafkaConfig,
+                            protected val replicaManager: ReplicaManager,
+                            metrics: Metrics,
+                            time: Time,
+                            threadNamePrefix: Option[String] = None,
+                            quotaManager: ReplicationQuotaManager)
+      extends AbstractFetcherManager[ReplicaFetcherThread](
+        name = "ReplicaFetcherManager on broker " + brokerConfig.brokerId,
+        clientId = "Replica",
+        numFetchers = brokerConfig.numReplicaFetchers) {
 
-  override def createFetcherThread(fetcherId: Int, sourceBroker: BrokerEndPoint): AbstractFetcherThread = {
-    val prefix = threadNamePrefix.map(tp => s"${tp}:").getOrElse("")
+  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, metrics, time, quotaManager)
+    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 da94c4a8f2163..14c080339a7d3 100644
--- a/core/src/main/scala/kafka/server/ReplicaFetcherThread.scala
+++ b/core/src/main/scala/kafka/server/ReplicaFetcherThread.scala
@@ -17,110 +17,187 @@
 
 package kafka.server
 
-import java.util
-
-import AbstractFetcherThread.ResultWithPartitions
-import kafka.api.{FetchRequest => _, _}
-import kafka.cluster.{BrokerEndPoint, Replica}
-import kafka.log.LogConfig
-import kafka.server.ReplicaFetcherThread._
-import kafka.server.epoch.LeaderEpochCache
-import kafka.zk.AdminZkClient
-import org.apache.kafka.common.requests.EpochEndOffset._
+import java.util.Collections
+import java.util.Optional
+
+import kafka.api._
+import kafka.cluster.BrokerEndPoint
+import kafka.log.{LeaderOffsetIncremented, LogAppendInfo}
+import kafka.server.AbstractFetcherThread.ReplicaFetch
+import kafka.server.AbstractFetcherThread.ResultWithPartitions
+import kafka.utils.Implicits._
+import org.apache.kafka.clients.FetchSessionHandler
 import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.errors.KafkaStorageException
-import org.apache.kafka.common.internals.FatalExitError
+import org.apache.kafka.common.message.ListOffsetRequestData.{ListOffsetPartition, ListOffsetTopic}
+import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData.EpochEndOffset
 import org.apache.kafka.common.metrics.Metrics
 import org.apache.kafka.common.protocol.Errors
-import org.apache.kafka.common.record.MemoryRecords
-import org.apache.kafka.common.requests.{EpochEndOffset, FetchResponse, ListOffsetRequest, ListOffsetResponse, OffsetsForLeaderEpochRequest, OffsetsForLeaderEpochResponse, FetchRequest => JFetchRequest}
+import org.apache.kafka.common.record.{MemoryRecords, Records}
+import org.apache.kafka.common.requests._
 import org.apache.kafka.common.utils.{LogContext, Time}
-import scala.collection.JavaConverters._
+
+import scala.jdk.CollectionConverters._
 import scala.collection.{Map, mutable}
+import scala.compat.java8.OptionConverters._
 
 class ReplicaFetcherThread(name: String,
                            fetcherId: Int,
                            sourceBroker: BrokerEndPoint,
                            brokerConfig: KafkaConfig,
+                           failedPartitions: FailedPartitions,
                            replicaMgr: ReplicaManager,
                            metrics: Metrics,
                            time: Time,
-                           quota: ReplicationQuotaManager,
+                           quota: ReplicaQuota,
                            leaderEndpointBlockingSend: Option[BlockingSend] = None)
   extends AbstractFetcherThread(name = name,
                                 clientId = name,
                                 sourceBroker = sourceBroker,
+                                failedPartitions,
                                 fetchBackOffMs = brokerConfig.replicaFetchBackoffMs,
                                 isInterruptible = false,
-                                includeLogTruncation = true) {
-
-  type REQ = FetchRequest
-  type PD = PartitionData
+                                replicaMgr.brokerTopicStats) {
 
   private val replicaId = brokerConfig.brokerId
   private val logContext = new LogContext(s"[ReplicaFetcher replicaId=$replicaId, leaderId=${sourceBroker.id}, " +
     s"fetcherId=$fetcherId] ")
   this.logIdent = logContext.logPrefix
+
   private val leaderEndpoint = leaderEndpointBlockingSend.getOrElse(
     new ReplicaFetcherBlockingSend(sourceBroker, brokerConfig, metrics, time, fetcherId,
       s"broker-$replicaId-fetcher-$fetcherId", logContext))
-  private val fetchRequestVersion: Short =
-    if (brokerConfig.interBrokerProtocolVersion >= KAFKA_0_11_0_IV1) 5
+
+  // Visible for testing
+  private[server] val fetchRequestVersion: Short =
+    if (brokerConfig.interBrokerProtocolVersion >= KAFKA_2_7_IV1) 12
+    else 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
     else if (brokerConfig.interBrokerProtocolVersion >= KAFKA_0_11_0_IV0) 4
     else if (brokerConfig.interBrokerProtocolVersion >= KAFKA_0_10_1_IV1) 3
     else if (brokerConfig.interBrokerProtocolVersion >= KAFKA_0_10_0_IV0) 2
     else if (brokerConfig.interBrokerProtocolVersion >= KAFKA_0_9_0) 1
     else 0
+
+  // Visible for testing
+  private[server] val offsetForLeaderEpochRequestVersion: Short =
+    if (brokerConfig.interBrokerProtocolVersion >= KAFKA_2_8_IV0) 4
+    else 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
+
+  // Visible for testing
+  private[server] val listOffsetRequestVersion: Short =
+    if (brokerConfig.interBrokerProtocolVersion >= KAFKA_2_8_IV0) 6
+    else if (brokerConfig.interBrokerProtocolVersion >= KAFKA_2_2_IV1) 5
+    else if (brokerConfig.interBrokerProtocolVersion >= KAFKA_2_1_IV1) 4
+    else if (brokerConfig.interBrokerProtocolVersion >= KAFKA_2_0_IV1) 3
+    else if (brokerConfig.interBrokerProtocolVersion >= KAFKA_0_11_0_IV0) 2
+    else if (brokerConfig.interBrokerProtocolVersion >= KAFKA_0_10_1_IV2) 1
+    else 0
+
   private val maxWait = brokerConfig.replicaFetchWaitMaxMs
   private val minBytes = brokerConfig.replicaFetchMinBytes
   private val maxBytes = brokerConfig.replicaFetchResponseMaxBytes
   private val fetchSize = brokerConfig.replicaFetchMaxBytes
-  private val shouldSendLeaderEpochRequest: Boolean = brokerConfig.interBrokerProtocolVersion >= KAFKA_0_11_0_IV2
+  override protected val isOffsetForLeaderEpochSupported: Boolean = brokerConfig.interBrokerProtocolVersion >= KAFKA_0_11_0_IV2
+  override protected val isTruncationOnFetchSupported = ApiVersion.isTruncationOnFetchSupported(brokerConfig.interBrokerProtocolVersion)
+  val fetchSessionHandler = new FetchSessionHandler(logContext, sourceBroker.id)
 
-  private def epochCacheOpt(tp: TopicPartition): Option[LeaderEpochCache] =  replicaMgr.getReplica(tp).map(_.epochs.get)
+  override protected def latestEpoch(topicPartition: TopicPartition): Option[Int] = {
+    replicaMgr.localLogOrException(topicPartition).latestEpoch
+  }
+
+  override protected def logStartOffset(topicPartition: TopicPartition): Long = {
+    replicaMgr.localLogOrException(topicPartition).logStartOffset
+  }
+
+  override protected def logEndOffset(topicPartition: TopicPartition): Long = {
+    replicaMgr.localLogOrException(topicPartition).logEndOffset
+  }
+
+  override protected def endOffsetForEpoch(topicPartition: TopicPartition, epoch: Int): Option[OffsetAndEpoch] = {
+    replicaMgr.localLogOrException(topicPartition).endOffsetForEpoch(epoch)
+  }
 
   override def initiateShutdown(): Boolean = {
     val justShutdown = super.initiateShutdown()
     if (justShutdown) {
-      leaderEndpoint.close()
+      // This is thread-safe, so we don't expect any exceptions, but catch and log any errors
+      // to avoid failing the caller, especially during shutdown. We will attempt to close
+      // leaderEndpoint after the thread terminates.
+      try {
+        leaderEndpoint.initiateClose()
+      } catch {
+        case t: Throwable =>
+          error(s"Failed to initiate shutdown of leader endpoint $leaderEndpoint after initiating replica fetcher thread shutdown", t)
+      }
     }
     justShutdown
   }
 
+  override def awaitShutdown(): Unit = {
+    super.awaitShutdown()
+    // We don't expect any exceptions here, but catch and log any errors to avoid failing the caller,
+    // especially during shutdown. It is safe to catch the exception here without causing correctness
+    // issue because we are going to shutdown the thread and will not re-use the leaderEndpoint anyway.
+    try {
+      leaderEndpoint.close()
+    } catch {
+      case t: Throwable =>
+        error(s"Failed to close leader endpoint $leaderEndpoint after shutting down replica fetcher thread", t)
+    }
+  }
+
   // process fetched data
-  def processPartitionData(topicPartition: TopicPartition, fetchOffset: Long, partitionData: PartitionData) {
-    val replica = replicaMgr.getReplicaOrException(topicPartition)
-    val partition = replicaMgr.getPartition(topicPartition).get
-    val records = partitionData.toRecords
+  override def processPartitionData(topicPartition: TopicPartition,
+                                    fetchOffset: Long,
+                                    partitionData: FetchData): Option[LogAppendInfo] = {
+    val logTrace = isTraceEnabled
+    val partition = replicaMgr.nonOfflinePartition(topicPartition).get
+    val log = partition.localLogOrException
+    val records = toMemoryRecords(partitionData.records)
 
     maybeWarnIfOversizedRecords(records, topicPartition)
 
-    if (fetchOffset != replica.logEndOffset.messageOffset)
+    if (fetchOffset != log.logEndOffset)
       throw new IllegalStateException("Offset mismatch for partition %s: fetched offset = %d, log end offset = %d.".format(
-        topicPartition, fetchOffset, replica.logEndOffset.messageOffset))
+        topicPartition, fetchOffset, log.logEndOffset))
 
-    if (isTraceEnabled)
+    if (logTrace)
       trace("Follower has replica log end offset %d for partition %s. Received %d messages and leader hw %d"
-        .format(replica.logEndOffset.messageOffset, topicPartition, records.sizeInBytes, partitionData.highWatermark))
+        .format(log.logEndOffset, topicPartition, records.sizeInBytes, partitionData.highWatermark))
 
     // Append the leader's messages to the log
-    partition.appendRecordsToFollower(records)
+    val logAppendInfo = partition.appendRecordsToFollowerOrFutureReplica(records, isFuture = false)
 
-    if (isTraceEnabled)
+    if (logTrace)
       trace("Follower has replica log end offset %d after appending %d bytes of messages for partition %s"
-        .format(replica.logEndOffset.messageOffset, records.sizeInBytes, topicPartition))
-    val followerHighWatermark = replica.logEndOffset.messageOffset.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)
-    if (isTraceEnabled)
+
+    // 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, LeaderOffsetIncremented)
+    if (logTrace)
       trace(s"Follower set replica high watermark for partition $topicPartition to $followerHighWatermark")
+
+    // Traffic from both in-sync and out of sync replicas are accounted for in replication quota to ensure total replication
+    // traffic doesn't exceed quota.
     if (quota.isThrottled(topicPartition))
       quota.record(records.sizeInBytes)
-    replicaMgr.brokerTopicStats.updateReplicationBytesIn(records.sizeInBytes)
+
+    if (partition.isReassigning && partition.isAddingLocalReplica)
+      brokerTopicStats.updateReassignmentBytesIn(records.sizeInBytes)
+
+    brokerTopicStats.updateReplicationBytesIn(records.sizeInBytes)
+
+    logAppendInfo
   }
 
   def maybeWarnIfOversizedRecords(records: MemoryRecords, topicPartition: TopicPartition): Unit = {
@@ -132,123 +209,76 @@ class ReplicaFetcherThread(name: String,
         "equal or larger than your settings for max.message.bytes, both at a broker and topic level.")
   }
 
-  /**
-   * Handle a partition whose offset is out of range and return a new fetch offset.
-   */
-  def handleOffsetOutOfRange(topicPartition: TopicPartition): Long = {
-    val replica = replicaMgr.getReplicaOrException(topicPartition)
-    val partition = replicaMgr.getPartition(topicPartition).get
-
-    /**
-     * Unclean leader election: A follower goes down, in the meanwhile the leader keeps appending messages. The follower comes back up
-     * and before it has completely caught up with the leader's logs, all replicas in the ISR go down. The follower is now uncleanly
-     * elected as the new leader, and it starts appending messages from the client. The old leader comes back up, becomes a follower
-     * and it may discover that the current leader's end offset is behind its own end offset.
-     *
-     * In such a case, truncate the current follower's log to the current leader's end offset and continue fetching.
-     *
-     * There is a potential for a mismatch between the logs of the two replicas here. We don't fix this mismatch as of now.
-     */
-    val leaderEndOffset: Long = earliestOrLatestOffset(topicPartition, ListOffsetRequest.LATEST_TIMESTAMP)
-
-    if (leaderEndOffset < replica.logEndOffset.messageOffset) {
-      // Prior to truncating the follower's log, ensure that doing so is not disallowed by the configuration for unclean leader election.
-      // This situation could only happen if the unclean election configuration for a topic changes while a replica is down. Otherwise,
-      // we should never encounter this situation since a non-ISR leader cannot be elected if disallowed by the broker configuration.
-      val adminZkClient = new AdminZkClient(replicaMgr.zkClient)
-      if (!LogConfig.fromProps(brokerConfig.originals, adminZkClient.fetchEntityConfig(
-        ConfigType.Topic, topicPartition.topic)).uncleanLeaderElectionEnable) {
-        // Log a fatal error and shutdown the broker to ensure that data loss does not occur unexpectedly.
-        fatal(s"Exiting because log truncation is not allowed for partition $topicPartition, current leader's " +
-          s"latest offset $leaderEndOffset is less than replica's latest offset ${replica.logEndOffset.messageOffset}")
-        throw new FatalExitError
-      }
 
-      warn(s"Reset fetch offset for partition $topicPartition from ${replica.logEndOffset.messageOffset} to current " +
-        s"leader's latest offset $leaderEndOffset")
-      partition.truncateTo(leaderEndOffset, isFuture = false)
-      replicaMgr.replicaAlterLogDirsManager.markPartitionsForTruncation(brokerConfig.brokerId, topicPartition, leaderEndOffset)
-      leaderEndOffset
-    } else {
-      /**
-       * If the leader's log end offset is greater than the follower's log end offset, there are two possibilities:
-       * 1. The follower could have been down for a long time and when it starts up, its end offset could be smaller than the leader's
-       * start offset because the leader has deleted old logs (log.logEndOffset < leaderStartOffset).
-       * 2. When unclean leader election occurs, it is possible that the old leader's high watermark is greater than
-       * the new leader's log end offset. So when the old leader truncates its offset to its high watermark and starts
-       * to fetch from the new leader, an OffsetOutOfRangeException will be thrown. After that some more messages are
-       * produced to the new leader. While the old leader is trying to handle the OffsetOutOfRangeException and query
-       * the log end offset of the new leader, the new leader's log end offset becomes higher than the follower's log end offset.
-       *
-       * In the first case, the follower's current log end offset is smaller than the leader's log start offset. So the
-       * follower should truncate all its logs, roll out a new segment and start to fetch from the current leader's log
-       * start offset.
-       * In the second case, the follower should just keep the current log segments and retry the fetch. In the second
-       * case, there will be some inconsistency of data between old and new leader. We are not solving it here.
-       * If users want to have strong consistency guarantees, appropriate configurations needs to be set for both
-       * brokers and producers.
-       *
-       * Putting the two cases together, the follower should fetch from the higher one of its replica log end offset
-       * and the current leader's log start offset.
-       *
-       */
-      val leaderStartOffset: Long = earliestOrLatestOffset(topicPartition, ListOffsetRequest.EARLIEST_TIMESTAMP)
-      warn(s"Reset fetch offset for partition $topicPartition from ${replica.logEndOffset.messageOffset} to current " +
-        s"leader's start offset $leaderStartOffset")
-      val offsetToFetch = Math.max(leaderStartOffset, replica.logEndOffset.messageOffset)
-      // Only truncate log when current leader's log start offset is greater than follower's log end offset.
-      if (leaderStartOffset > replica.logEndOffset.messageOffset) {
-        partition.truncateFullyAndStartAt(leaderStartOffset, isFuture = false)
+  override protected def fetchFromLeader(fetchRequest: FetchRequest.Builder): Map[TopicPartition, FetchData] = {
+    try {
+      val clientResponse = leaderEndpoint.sendRequest(fetchRequest)
+      val fetchResponse = clientResponse.responseBody.asInstanceOf[FetchResponse[Records]]
+      if (!fetchSessionHandler.handleResponse(fetchResponse)) {
+        Map.empty
+      } else {
+        fetchResponse.responseData.asScala
       }
-      offsetToFetch
+    } catch {
+      case t: Throwable =>
+        fetchSessionHandler.handleError(t)
+        throw t
     }
   }
 
-  // any logic for partitions whose leader has changed
-  def handlePartitionsWithErrors(partitions: Iterable[TopicPartition]) {
-    if (partitions.nonEmpty)
-      delayPartitions(partitions, brokerConfig.replicaFetchBackoffMs.toLong)
+  override protected def fetchEarliestOffsetFromLeader(topicPartition: TopicPartition, currentLeaderEpoch: Int): Long = {
+    fetchOffsetFromLeader(topicPartition, currentLeaderEpoch, ListOffsetRequest.EARLIEST_TIMESTAMP)
   }
 
-  protected def fetch(fetchRequest: FetchRequest): Seq[(TopicPartition, PartitionData)] = {
-    val clientResponse = leaderEndpoint.sendRequest(fetchRequest.underlying)
-    val fetchResponse = clientResponse.responseBody.asInstanceOf[FetchResponse]
-    fetchResponse.responseData.asScala.toSeq.map { case (key, value) =>
-      key -> new PartitionData(value)
-    }
+  override protected def fetchLatestOffsetFromLeader(topicPartition: TopicPartition, currentLeaderEpoch: Int): Long = {
+    fetchOffsetFromLeader(topicPartition, currentLeaderEpoch, ListOffsetRequest.LATEST_TIMESTAMP)
   }
 
-  private def earliestOrLatestOffset(topicPartition: TopicPartition, earliestOrLatest: Long): Long = {
-    val requestBuilder = if (brokerConfig.interBrokerProtocolVersion >= KAFKA_0_10_1_IV2) {
-        val partitions = Map(topicPartition -> (earliestOrLatest: java.lang.Long))
-        ListOffsetRequest.Builder.forReplica(1, replicaId).setTargetTimes(partitions.asJava)
-      } else {
-        val partitions = Map(topicPartition -> new ListOffsetRequest.PartitionData(earliestOrLatest, 1))
-        ListOffsetRequest.Builder.forReplica(0, replicaId).setOffsetData(partitions.asJava)
-      }
+  private def fetchOffsetFromLeader(topicPartition: TopicPartition, currentLeaderEpoch: Int, earliestOrLatest: Long): Long = {
+    val topic = new ListOffsetTopic()
+      .setName(topicPartition.topic)
+      .setPartitions(Collections.singletonList(
+          new ListOffsetPartition()
+            .setPartitionIndex(topicPartition.partition)
+            .setCurrentLeaderEpoch(currentLeaderEpoch)
+            .setTimestamp(earliestOrLatest)))
+    val requestBuilder = ListOffsetRequest.Builder.forReplica(listOffsetRequestVersion, replicaId)
+      .setTargetTimes(Collections.singletonList(topic))
+
     val clientResponse = leaderEndpoint.sendRequest(requestBuilder)
     val response = clientResponse.responseBody.asInstanceOf[ListOffsetResponse]
-    val partitionData = response.responseData.get(topicPartition)
-    partitionData.error match {
+    val responsePartition = response.topics.asScala.find(_.name == topicPartition.topic).get
+      .partitions.asScala.find(_.partitionIndex == topicPartition.partition).get
+
+     Errors.forCode(responsePartition.errorCode) match {
       case Errors.NONE =>
         if (brokerConfig.interBrokerProtocolVersion >= KAFKA_0_10_1_IV2)
-          partitionData.offset
+          responsePartition.offset
         else
-          partitionData.offsets.get(0)
+          responsePartition.oldStyleOffsets.get(0)
       case error => throw error.exception
     }
   }
 
-  override def buildFetchRequest(partitionMap: Seq[(TopicPartition, PartitionFetchState)]): ResultWithPartitions[FetchRequest] = {
-    val requestMap = new util.LinkedHashMap[TopicPartition, JFetchRequest.PartitionData]
+  override def buildFetch(partitionMap: Map[TopicPartition, PartitionFetchState]): ResultWithPartitions[Option[ReplicaFetch]] = {
     val partitionsWithError = mutable.Set[TopicPartition]()
 
-    partitionMap.foreach { case (topicPartition, partitionFetchState) =>
+    val builder = fetchSessionHandler.newBuilder(partitionMap.size, false)
+    partitionMap.forKeyValue { (topicPartition, fetchState) =>
       // We will not include a replica in the fetch request if it should be throttled.
-      if (partitionFetchState.isReadyForFetch && !shouldFollowerThrottle(quota, topicPartition)) {
+      if (fetchState.isReadyForFetch && !shouldFollowerThrottle(quota, fetchState, topicPartition)) {
         try {
-          val logStartOffset = replicaMgr.getReplicaOrException(topicPartition).logStartOffset
-          requestMap.put(topicPartition, new JFetchRequest.PartitionData(partitionFetchState.fetchOffset, logStartOffset, fetchSize))
+          val logStartOffset = this.logStartOffset(topicPartition)
+          val lastFetchedEpoch = if (isTruncationOnFetchSupported)
+            fetchState.lastFetchedEpoch.map(_.asInstanceOf[Integer]).asJava
+          else
+            Optional.empty[Integer]
+          builder.add(topicPartition, new FetchRequest.PartitionData(
+            fetchState.fetchOffset,
+            logStartOffset,
+            fetchSize,
+            Optional.of(fetchState.currentLeaderEpoch),
+            lastFetchedEpoch))
         } catch {
           case _: KafkaStorageException =>
             // The replica has already been marked offline due to log directory failure and the original failure should have already been logged.
@@ -258,137 +288,86 @@ class ReplicaFetcherThread(name: String,
       }
     }
 
-    val requestBuilder = JFetchRequest.Builder.forReplica(fetchRequestVersion, replicaId, maxWait, minBytes, requestMap)
-      .setMaxBytes(maxBytes)
-    ResultWithPartitions(new FetchRequest(requestBuilder), partitionsWithError)
+    val fetchData = builder.build()
+    val fetchRequestOpt = if (fetchData.sessionPartitions.isEmpty && fetchData.toForget.isEmpty) {
+      None
+    } else {
+      val requestBuilder = FetchRequest.Builder
+        .forReplica(fetchRequestVersion, replicaId, maxWait, minBytes, fetchData.toSend)
+        .setMaxBytes(maxBytes)
+        .toForget(fetchData.toForget)
+        .metadata(fetchData.metadata)
+      Some(ReplicaFetch(fetchData.sessionPartitions(), requestBuilder))
+    }
+
+    ResultWithPartitions(fetchRequestOpt, partitionsWithError)
   }
 
   /**
-    * - Truncate the log to the leader's offset for each partition's epoch.
-    * - If the leader's offset is greater, we stick with the Log End Offset
-    *   otherwise we truncate to the leaders offset.
-    * - If the leader replied with undefined epoch offset we must use the high watermark
-    */
-  override def maybeTruncate(fetchedEpochs: Map[TopicPartition, EpochEndOffset]): ResultWithPartitions[Map[TopicPartition, Long]] = {
-    val fetchOffsets = scala.collection.mutable.HashMap.empty[TopicPartition, Long]
-    val partitionsWithError = mutable.Set[TopicPartition]()
-
-    fetchedEpochs.foreach { case (tp, epochOffset) =>
-      try {
-        val replica = replicaMgr.getReplicaOrException(tp)
-        val partition = replicaMgr.getPartition(tp).get
-
-        if (epochOffset.hasError) {
-          info(s"Retrying leaderEpoch request for partition ${replica.topicPartition} as the leader reported an error: ${epochOffset.error}")
-          partitionsWithError += tp
-        } else {
-          val fetchOffset =
-            if (epochOffset.endOffset == UNDEFINED_EPOCH_OFFSET) {
-              warn(s"Based on follower's leader epoch, leader replied with an unknown offset in ${replica.topicPartition}. " +
-                s"The initial fetch offset ${partitionStates.stateValue(tp).fetchOffset} will be used for truncation.")
-              partitionStates.stateValue(tp).fetchOffset
-            } else if (epochOffset.endOffset >= replica.logEndOffset.messageOffset)
-              logEndOffset(replica, epochOffset)
-            else
-              epochOffset.endOffset
-
-          partition.truncateTo(fetchOffset, isFuture = false)
-          replicaMgr.replicaAlterLogDirsManager.markPartitionsForTruncation(brokerConfig.brokerId, tp, fetchOffset)
-          fetchOffsets.put(tp, fetchOffset)
-        }
-      } catch {
-        case e: KafkaStorageException =>
-          info(s"Failed to truncate $tp", e)
-          partitionsWithError += tp
-      }
-    }
+   * Truncate the log for each partition's epoch based on leader's returned epoch and offset.
+   * The logic for finding the truncation offset is implemented in AbstractFetcherThread.getOffsetTruncationState
+   */
+  override def truncate(tp: TopicPartition, offsetTruncationState: OffsetTruncationState): Unit = {
+    val partition = replicaMgr.nonOfflinePartition(tp).get
+    val log = partition.localLogOrException
 
-    ResultWithPartitions(fetchOffsets, partitionsWithError)
-  }
+    partition.truncateTo(offsetTruncationState.offset, isFuture = false)
 
-  override def buildLeaderEpochRequest(allPartitions: Seq[(TopicPartition, PartitionFetchState)]): ResultWithPartitions[Map[TopicPartition, Int]] = {
-    val partitionEpochOpts = allPartitions
-      .filter { case (_, state) => state.isTruncatingLog }
-      .map { case (tp, _) => tp -> epochCacheOpt(tp) }.toMap
+    if (offsetTruncationState.offset < log.highWatermark)
+      warn(s"Truncating $tp to offset ${offsetTruncationState.offset} below high watermark " +
+        s"${log.highWatermark}")
 
-    val (partitionsWithEpoch, partitionsWithoutEpoch) = partitionEpochOpts.partition { case (tp, epochCacheOpt) => epochCacheOpt.nonEmpty }
+    // mark the future replica for truncation only when we do last truncation
+    if (offsetTruncationState.truncationCompleted)
+      replicaMgr.replicaAlterLogDirsManager.markPartitionsForTruncation(brokerConfig.brokerId, tp,
+        offsetTruncationState.offset)
+  }
 
-    debug(s"Build leaderEpoch request $partitionsWithEpoch")
-    val result = partitionsWithEpoch.map { case (tp, epochCacheOpt) => tp -> epochCacheOpt.get.latestEpoch() }
-    ResultWithPartitions(result, partitionsWithoutEpoch.keys.toSet)
+  override protected def truncateFullyAndStartAt(topicPartition: TopicPartition, offset: Long): Unit = {
+    val partition = replicaMgr.nonOfflinePartition(topicPartition).get
+    partition.truncateFullyAndStartAt(offset, isFuture = false)
   }
 
-  override def fetchEpochsFromLeader(partitions: Map[TopicPartition, Int]): Map[TopicPartition, EpochEndOffset] = {
-    var result: Map[TopicPartition, EpochEndOffset] = null
-    if (shouldSendLeaderEpochRequest) {
-      val partitionsAsJava = partitions.map { case (tp, epoch) => tp -> epoch.asInstanceOf[Integer] }.toMap.asJava
-      val epochRequest = new OffsetsForLeaderEpochRequest.Builder(partitionsAsJava)
-      try {
-        val response = leaderEndpoint.sendRequest(epochRequest)
-        result = response.responseBody.asInstanceOf[OffsetsForLeaderEpochResponse].responses.asScala
-        debug(s"Receive leaderEpoch response $result")
-      } catch {
-        case t: Throwable =>
-          warn(s"Error when sending leader epoch request for $partitions", t)
+  override def fetchEpochEndOffsets(partitions: Map[TopicPartition, EpochData]): Map[TopicPartition, EpochEndOffset] = {
 
-          // if we get any unexpected exception, mark all partitions with an error
-          result = partitions.map { case (tp, _) =>
-            tp -> new EpochEndOffset(Errors.forException(t), UNDEFINED_EPOCH_OFFSET)
-          }
-      }
-    } else {
-      // just generate a response with no error but UNDEFINED_OFFSET so that we can fall back to truncating using
-      // high watermark in maybeTruncate()
-      result = partitions.map { case (tp, _) =>
-        tp -> new EpochEndOffset(Errors.NONE, UNDEFINED_EPOCH_OFFSET)
-      }
+    if (partitions.isEmpty) {
+      debug("Skipping leaderEpoch request since all partitions do not have an epoch")
+      return Map.empty
     }
-    result
-  }
 
-  private def logEndOffset(replica: Replica, epochOffset: EpochEndOffset): Long = {
-    val logEndOffset = replica.logEndOffset.messageOffset
-    info(s"Based on follower's leader epoch, leader replied with an offset ${epochOffset.endOffset} >= the " +
-      s"follower's log end offset $logEndOffset in ${replica.topicPartition}. No truncation needed.")
-    logEndOffset
+    val epochRequest = OffsetsForLeaderEpochRequest.Builder.forFollower(offsetForLeaderEpochRequestVersion, partitions.asJava, brokerConfig.brokerId)
+    debug(s"Sending offset for leader epoch request $epochRequest")
+
+    try {
+      val response = leaderEndpoint.sendRequest(epochRequest)
+      val responseBody = response.responseBody.asInstanceOf[OffsetsForLeaderEpochResponse]
+      debug(s"Received leaderEpoch response $response")
+      responseBody.data.topics.asScala.flatMap { offsetForLeaderTopicResult =>
+        offsetForLeaderTopicResult.partitions().asScala.map { offsetForLeaderPartitionResult =>
+          val tp = new TopicPartition(offsetForLeaderTopicResult.topic, offsetForLeaderPartitionResult.partition)
+          tp -> offsetForLeaderPartitionResult
+        }
+      }.toMap
+    } catch {
+      case t: Throwable =>
+        warn(s"Error when sending leader epoch request for $partitions", t)
+
+        // if we get any unexpected exception, mark all partitions with an error
+        val error = Errors.forException(t)
+        partitions.map { case (tp, _) =>
+          tp -> new EpochEndOffset()
+            .setPartition(tp.partition)
+            .setErrorCode(error.code)
+        }
+    }
   }
 
   /**
    *  To avoid ISR thrashing, we only throttle a replica on the follower if it's in the throttled replica list,
    *  the quota is exceeded and the replica is not in sync.
    */
-  private def shouldFollowerThrottle(quota: ReplicaQuota, topicPartition: TopicPartition): Boolean = {
-    val isReplicaInSync = fetcherLagStats.isReplicaInSync(topicPartition.topic, topicPartition.partition)
-    quota.isThrottled(topicPartition) && quota.isQuotaExceeded && !isReplicaInSync
-  }
-}
-
-object ReplicaFetcherThread {
-
-  private[server] class FetchRequest(val underlying: JFetchRequest.Builder) extends AbstractFetcherThread.FetchRequest {
-    def isEmpty: Boolean = underlying.fetchData.isEmpty
-    def offset(topicPartition: TopicPartition): Long =
-      underlying.fetchData.asScala(topicPartition).fetchOffset
-    override def toString = underlying.toString
+  private def shouldFollowerThrottle(quota: ReplicaQuota, fetchState: PartitionFetchState, topicPartition: TopicPartition): Boolean = {
+    !fetchState.isReplicaInSync && quota.isThrottled(topicPartition) && quota.isQuotaExceeded
   }
 
-  private[server] class PartitionData(val underlying: FetchResponse.PartitionData) extends AbstractFetcherThread.PartitionData {
-
-    def error = underlying.error
-
-    def toRecords: MemoryRecords = {
-      underlying.records.asInstanceOf[MemoryRecords]
-    }
-
-    def highWatermark: Long = underlying.highWatermark
-
-    def logStartOffset: Long = underlying.logStartOffset
-
-    def exception: Option[Throwable] = error match {
-      case Errors.NONE => None
-      case e => Some(e.exception)
-    }
-
-    override def toString = underlying.toString
-  }
 }
diff --git a/core/src/main/scala/kafka/server/ReplicaManager.scala b/core/src/main/scala/kafka/server/ReplicaManager.scala
index 1344660a580f5..ea4402172f27b 100644
--- a/core/src/main/scala/kafka/server/ReplicaManager.scala
+++ b/core/src/main/scala/kafka/server/ReplicaManager.scala
@@ -17,38 +17,52 @@
 package kafka.server
 
 import java.io.File
+import java.util.Optional
 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.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.{Log, LogAppendInfo, LogManager}
+import kafka.log._
 import kafka.metrics.KafkaMetricsGroup
-import kafka.server.QuotaFactory.{QuotaManagers, UnboundedQuota}
-import kafka.server.checkpoints.OffsetCheckpointFile
+import kafka.server.{FetchMetadata => SFetchMetadata}
+import kafka.server.HostedPartition.Online
+import kafka.server.QuotaFactory.QuotaManagers
+import kafka.server.checkpoints.{LazyOffsetCheckpoints, OffsetCheckpointFile, OffsetCheckpoints}
 import kafka.utils._
+import kafka.utils.Implicits._
 import kafka.zk.KafkaZkClient
-import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.{ElectionType, IsolationLevel, 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.DeleteRecordsResponseData.DeleteRecordsPartitionResult
+import org.apache.kafka.common.message.{DescribeLogDirsResponseData, FetchResponseData, LeaderAndIsrResponseData}
+import org.apache.kafka.common.message.LeaderAndIsrResponseData.LeaderAndIsrPartitionError
+import org.apache.kafka.common.message.OffsetForLeaderEpochRequestData.{OffsetForLeaderTopic}
+import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData.{OffsetForLeaderTopicResult, EpochEndOffset}
+import org.apache.kafka.common.message.StopReplicaRequestData.StopReplicaPartitionState
 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.protocol.Errors.UNKNOWN_TOPIC_OR_PARTITION
-import org.apache.kafka.common.protocol.Errors.KAFKA_STORAGE_ERROR
+import org.apache.kafka.common.record.FileRecords.TimestampAndOffset
 import org.apache.kafka.common.record._
-import org.apache.kafka.common.requests.DescribeLogDirsResponse.{LogDirInfo, ReplicaInfo}
-import org.apache.kafka.common.requests.EpochEndOffset._
+import org.apache.kafka.common.replica.PartitionView.DefaultPartitionView
+import org.apache.kafka.common.replica.ReplicaView.DefaultReplicaView
+import org.apache.kafka.common.replica.{ClientMetadata, _}
 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.utils.Time
 
-import scala.collection.JavaConverters._
-import scala.collection._
+import scala.jdk.CollectionConverters._
+import scala.collection.{Map, Seq, Set, mutable}
+import scala.compat.java8.OptionConverters._
 
 /*
  * Result metadata of a log append operation on the log
@@ -67,23 +81,29 @@ case class LogDeleteRecordsResult(requestedOffset: Long, lowWatermark: Long, exc
   }
 }
 
-/*
+/**
  * Result metadata of a log read operation on the log
  * @param info @FetchDataInfo returned by the @Log read
- * @param hw high watermark of the local replica
- * @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 divergingEpoch Optional epoch and end offset which indicates the largest epoch such
+ *                       that subsequent records are known to diverge on the follower/consumer
+ * @param highWatermark high watermark of the local replica
+ * @param leaderLogStartOffset The log start offset of the leader at the time of the read
+ * @param leaderLogEndOffset The log end offset of the leader at the time of the read
+ * @param followerLogStartOffset The log start offset of the follower taken from the Fetch request
+ * @param fetchTimeMs The time the fetch was received
+ * @param lastStableOffset Current LSO or None if the result has an exception
+ * @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,
+                         divergingEpoch: Option[FetchResponseData.EpochEndOffset],
                          highWatermark: Long,
                          leaderLogStartOffset: Long,
                          leaderLogEndOffset: Long,
                          followerLogStartOffset: Long,
                          fetchTimeMs: Long,
-                         readSize: Int,
                          lastStableOffset: Option[Long],
+                         preferredReadReplica: Option[Int] = None,
                          exception: Option[Throwable] = None) {
 
   def error: Errors = exception match {
@@ -91,17 +111,23 @@ case class LogReadResult(info: FetchDataInfo,
     case Some(e) => Errors.forException(e)
   }
 
-  def updateLeaderReplicaInfo(leaderReplica: Replica): LogReadResult =
-    copy(highWatermark = leaderReplica.highWatermark.messageOffset,
-      leaderLogStartOffset = leaderReplica.logStartOffset,
-      leaderLogEndOffset = leaderReplica.logEndOffset.messageOffset)
-
   def withEmptyFetchInfo: LogReadResult =
     copy(info = FetchDataInfo(LogOffsetMetadata.UnknownOffsetMetadata, MemoryRecords.EMPTY))
 
-  override def toString =
-    s"Fetch Data: [$info], HW: [$highWatermark], leaderLogStartOffset: [$leaderLogStartOffset], leaderLogEndOffset: [$leaderLogEndOffset], " +
-    s"followerLogStartOffset: [$followerLogStartOffset], fetchTimeMs: [$fetchTimeMs], readSize: [$readSize], error: [$error]"
+  override def toString = {
+    "LogReadResult(" +
+      s"info=$info, " +
+      s"divergingEpoch=$divergingEpoch, " +
+      s"highWatermark=$highWatermark, " +
+      s"leaderLogStartOffset=$leaderLogStartOffset, " +
+      s"leaderLogEndOffset=$leaderLogEndOffset, " +
+      s"followerLogStartOffset=$followerLogStartOffset, " +
+      s"fetchTimeMs=$fetchTimeMs, " +
+      s"preferredReadReplica=$preferredReadReplica, " +
+      s"lastStableOffset=$lastStableOffset, " +
+      s"error=$error" +
+      ")"
+  }
 
 }
 
@@ -109,25 +135,56 @@ case class FetchPartitionData(error: Errors = Errors.NONE,
                               highWatermark: Long,
                               logStartOffset: Long,
                               records: Records,
+                              divergingEpoch: Option[FetchResponseData.EpochEndOffset],
                               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],
+                              isReassignmentFetch: Boolean)
+
+
+/**
+ * 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
 }
 
+case class IsrChangePropagationConfig(
+  // How often to check for ISR
+  checkIntervalMs: Long,
+
+  // Maximum time that an ISR change may be delayed before sending the notification
+  maxDelayMs: Long,
+
+  // Maximum time to await additional changes before sending the notification
+  lingerMs: Long
+)
+
 object ReplicaManager {
   val HighWatermarkFilename = "replication-offset-checkpoint"
-  val IsrChangePropagationBlackOut = 5000L
-  val IsrChangePropagationInterval = 60000L
-  val OfflinePartition = new Partition("", -1, null, null, isOffline = true)
+
+  // This field is mutable to allow overriding change notification behavior in test cases
+  @volatile var DefaultIsrPropagationConfig: IsrChangePropagationConfig = IsrChangePropagationConfig(
+    checkIntervalMs = 2500,
+    lingerMs = 5000,
+    maxDelayMs = 60000,
+  )
 }
 
 class ReplicaManager(val config: KafkaConfig,
@@ -144,7 +201,9 @@ class ReplicaManager(val config: KafkaConfig,
                      val delayedProducePurgatory: DelayedOperationPurgatory[DelayedProduce],
                      val delayedFetchPurgatory: DelayedOperationPurgatory[DelayedFetch],
                      val delayedDeleteRecordsPurgatory: DelayedOperationPurgatory[DelayedDeleteRecords],
-                     threadNamePrefix: Option[String]) extends Logging with KafkaMetricsGroup {
+                     val delayedElectLeaderPurgatory: DelayedOperationPurgatory[DelayedElectLeader],
+                     threadNamePrefix: Option[String],
+                     val alterIsrManager: AlterIsrManager) extends Logging with KafkaMetricsGroup {
 
   def this(config: KafkaConfig,
            metrics: Metrics,
@@ -157,7 +216,8 @@ class ReplicaManager(val config: KafkaConfig,
            brokerTopicStats: BrokerTopicStats,
            metadataCache: MetadataCache,
            logDirFailureChannel: LogDirFailureChannel,
-           threadNamePrefix: Option[String] = None) {
+           alterIsrManager: AlterIsrManager,
+           threadNamePrefix: Option[String] = None) = {
     this(config, metrics, time, zkClient, scheduler, logManager, isShuttingDown,
       quotaManagers, brokerTopicStats, metadataCache, logDirFailureChannel,
       DelayedOperationPurgatory[DelayedProduce](
@@ -169,33 +229,36 @@ class ReplicaManager(val config: KafkaConfig,
       DelayedOperationPurgatory[DelayedDeleteRecords](
         purgatoryName = "DeleteRecords", brokerId = config.brokerId,
         purgeInterval = config.deleteRecordsPurgatoryPurgeIntervalRequests),
-      threadNamePrefix)
+      DelayedOperationPurgatory[DelayedElectLeader](
+        purgatoryName = "ElectLeader", brokerId = config.brokerId),
+      threadNamePrefix, alterIsrManager)
   }
 
   /* epoch of the controller that last changed the leader */
-  @volatile var controllerEpoch: Int = KafkaController.InitialControllerEpoch - 1
+  @volatile var controllerEpoch: Int = KafkaController.InitialControllerEpoch
   private val localBrokerId = config.brokerId
-  private val allPartitions = new Pool[TopicPartition, Partition](valueFactory = Some(tp =>
-    new Partition(tp.topic, tp.partition, 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 isrChangeNotificationConfig = ReplicaManager.DefaultIsrPropagationConfig
   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 val lastIsrChangeMs = new AtomicLong(time.milliseconds())
+  private val lastIsrPropagationMs = new AtomicLong(time.milliseconds())
 
   private var logDirFailureHandler: LogDirFailureHandler = null
 
   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")
@@ -205,52 +268,37 @@ class ReplicaManager(val config: KafkaConfig,
     }
   }
 
-  val leaderCount = newGauge(
-    "LeaderCount",
-    new Gauge[Int] {
-      def value = leaderPartitionsIterator.size
-    }
-  )
-  val partitionCount = newGauge(
-    "PartitionCount",
-    new Gauge[Int] {
-      def value = allPartitions.size
-    }
-  )
-  val offlineReplicaCount = newGauge(
-    "OfflineReplicaCount",
-    new Gauge[Int] {
-      def value = offlinePartitionsIterator.size
-    }
-  )
-  val underReplicatedPartitions = newGauge(
-    "UnderReplicatedPartitions",
-    new Gauge[Int] {
-      def value = underReplicatedPartitionCount
-    }
-  )
-  val underMinIsrPartitionCount = newGauge(
-    "UnderMinIsrPartitionCount",
-    new Gauge[Int] {
-      def value = leaderPartitionsIterator.count(_.isUnderMinIsr)
-    }
-  )
+  // Visible for testing
+  private[server] val replicaSelectorOpt: Option[ReplicaSelector] = createReplicaSelector()
+
+  newGauge("LeaderCount", () => leaderPartitionsIterator.size)
+  // Visible for testing
+  private[kafka] val partitionCount = newGauge("PartitionCount", () => allPartitions.size)
+  newGauge("OfflineReplicaCount", () => offlinePartitionCount)
+  newGauge("UnderReplicatedPartitions", () => underReplicatedPartitionCount)
+  newGauge("UnderMinIsrPartitionCount", () => leaderPartitionsIterator.count(_.isUnderMinIsr))
+  newGauge("AtMinIsrPartitionCount", () => leaderPartitionsIterator.count(_.isAtMinIsr))
+  newGauge("ReassigningPartitions", () => reassigningPartitionsCount)
+
+  def reassigningPartitionsCount: Int = leaderPartitionsIterator.count(_.isReassigning)
 
-  val isrExpandRate = newMeter("IsrExpandsPerSec", "expands", TimeUnit.SECONDS)
-  val isrShrinkRate = newMeter("IsrShrinksPerSec", "shrinks", TimeUnit.SECONDS)
-  val failedIsrUpdatesRate = newMeter("FailedIsrUpdatesPerSec", "failedUpdates", TimeUnit.SECONDS)
+  val isrExpandRate: Meter = newMeter("IsrExpandsPerSec", "expands", TimeUnit.SECONDS)
+  val isrShrinkRate: Meter = newMeter("IsrShrinksPerSec", "shrinks", TimeUnit.SECONDS)
+  val failedIsrUpdatesRate: Meter = newMeter("FailedIsrUpdatesPerSec", "failedUpdates", TimeUnit.SECONDS)
 
   def underReplicatedPartitionCount: Int = leaderPartitionsIterator.count(_.isUnderReplicated)
 
-  def startHighWaterMarksCheckPointThread() = {
-    if(highWatermarkCheckPointThreadStarted.compareAndSet(false, true))
+  def startHighWatermarkCheckPointThread(): Unit = {
+    if (highWatermarkCheckPointThreadStarted.compareAndSet(false, true))
       scheduler.schedule("highwatermark-checkpoint", checkpointHighWatermarks _, period = config.replicaHighWatermarkCheckpointIntervalMs, unit = TimeUnit.MILLISECONDS)
   }
 
-  def recordIsrChange(topicPartition: TopicPartition) {
-    isrChangeSet synchronized {
-      isrChangeSet += topicPartition
-      lastIsrChangeMs.set(System.currentTimeMillis())
+  def recordIsrChange(topicPartition: TopicPartition): Unit = {
+    if (!config.interBrokerProtocolVersion.isAlterIsrSupported) {
+      isrChangeSet synchronized {
+        isrChangeSet += topicPartition
+        lastIsrChangeMs.set(time.milliseconds())
+      }
     }
   }
   /**
@@ -260,13 +308,13 @@ class ReplicaManager(val config: KafkaConfig,
    * 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.
    */
-  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)) {
-        ReplicationUtils.propagateIsrChanges(zkClient, isrChangeSet)
+        (lastIsrChangeMs.get() + isrChangeNotificationConfig.lingerMs < now ||
+          lastIsrPropagationMs.get() + isrChangeNotificationConfig.maxDelayMs < now)) {
+        zkClient.propagateIsrChanges(isrChangeSet)
         isrChangeSet.clear()
         lastIsrPropagationMs.set(now)
       }
@@ -283,44 +331,24 @@ 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))
-  }
+  def hasDelayedElectionOperations: Boolean = delayedElectLeaderPurgatory.numDelayed != 0
 
-  /**
-   * 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))
+  def tryCompleteElection(key: DelayedOperationKey): Unit = {
+    val completed = delayedElectLeaderPurgatory.checkAndComplete(key)
+    debug("Request key %s unblocked %d ElectLeader.".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 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)
-    scheduler.schedule("isr-change-propagation", maybePropagateIsrChanges _, period = 2500L, unit = TimeUnit.MILLISECONDS)
+    // If using AlterIsr, we don't need the znode ISR propagation
+    if (!config.interBrokerProtocolVersion.isAlterIsrSupported) {
+      scheduler.schedule("isr-change-propagation", maybePropagateIsrChanges _,
+        period = isrChangeNotificationConfig.checkIntervalMs, unit = TimeUnit.MILLISECONDS)
+    } else {
+      alterIsrManager.start()
+    }
     scheduler.schedule("shutdown-idle-replica-alter-log-dirs-thread", shutdownIdleReplicaAlterLogDirsThread _, period = 10000L, unit = TimeUnit.MILLISECONDS)
 
     // If inter-broker protocol (IBP) < 1.0, the controller will send LeaderAndIsrRequest V0 which does not include isNew field.
@@ -331,153 +359,298 @@ class ReplicaManager(val config: KafkaConfig,
     logDirFailureHandler.start()
   }
 
-  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)
-        throw new KafkaStorageException(s"Partition $topicPartition is on an offline disk")
-
-      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")
-      }
-
-      // Delete log and corresponding folders in case replica manager doesn't hold them anymore.
-      // This could happen when topic is being deleted while broker is down and recovers.
-      if (logManager.getLog(topicPartition).isDefined)
-        logManager.asyncDelete(topicPartition)
-      if (logManager.getLog(topicPartition, isFuture = true).isDefined)
-        logManager.asyncDelete(topicPartition, isFuture = true)
+  private def maybeRemoveTopicMetrics(topic: String): Unit = {
+    val topicHasOnlinePartition = allPartitions.values.exists {
+      case HostedPartition.Online(partition) => topic == partition.topic
+      case HostedPartition.None | HostedPartition.Offline => false
     }
-    stateChangeLogger.trace(s"Finished handling stop replica (delete=$deletePartition) for partition $topicPartition")
+    if (!topicHasOnlinePartition)
+      brokerTopicStats.removeMetrics(topic)
+  }
+
+  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) = {
+  def stopReplicas(correlationId: Int,
+                   controllerId: Int,
+                   controllerEpoch: Int,
+                   brokerEpoch: Long,
+                   partitionStates: Map[TopicPartition, StopReplicaPartitionState]
+                  ): (mutable.Map[TopicPartition, Errors], Errors) = {
     replicaStateChangeLock synchronized {
+      stateChangeLogger.info(s"Handling StopReplica request correlationId $correlationId from controller " +
+        s"$controllerId for ${partitionStates.size} partitions")
+      if (stateChangeLogger.isTraceEnabled)
+        partitionStates.forKeyValue { (topicPartition, partitionState) =>
+          stateChangeLogger.trace(s"Received StopReplica request $partitionState " +
+            s"correlation id $correlationId from controller $controllerId " +
+            s"epoch $controllerEpoch for partition $topicPartition")
+        }
+
       val responseMap = new collection.mutable.HashMap[TopicPartition, Errors]
-      if(stopReplicaRequest.controllerEpoch() < controllerEpoch) {
-        stateChangeLogger.warn("Received stop replica request from an old controller epoch " +
-          s"${stopReplicaRequest.controllerEpoch}. Latest known controller epoch is $controllerEpoch")
+      if (controllerEpoch < this.controllerEpoch) {
+        stateChangeLogger.warn(s"Ignoring StopReplica request from " +
+          s"controller $controllerId with correlation id $correlationId " +
+          s"since its controller epoch $controllerEpoch is old. " +
+          s"Latest known controller epoch is ${this.controllerEpoch}")
         (responseMap, Errors.STALE_CONTROLLER_EPOCH)
       } else {
-        val partitions = stopReplicaRequest.partitions.asScala
-        controllerEpoch = stopReplicaRequest.controllerEpoch
-        // First stop fetchers for all partitions, then stop the corresponding replicas
+        this.controllerEpoch = controllerEpoch
+
+        val stoppedPartitions = mutable.Map.empty[TopicPartition, StopReplicaPartitionState]
+        partitionStates.forKeyValue { (topicPartition, partitionState) =>
+          val deletePartition = partitionState.deletePartition
+
+          getPartition(topicPartition) match {
+            case HostedPartition.Offline =>
+              stateChangeLogger.warn(s"Ignoring StopReplica request (delete=$deletePartition) 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)
+
+            case HostedPartition.Online(partition) =>
+              val currentLeaderEpoch = partition.getLeaderEpoch
+              val requestLeaderEpoch = partitionState.leaderEpoch
+              // When a topic is deleted, the leader epoch is not incremented. To circumvent this,
+              // a sentinel value (EpochDuringDelete) overwriting any previous epoch is used.
+              // When an older version of the StopReplica request which does not contain the leader
+              // epoch, a sentinel value (NoEpoch) is used and bypass the epoch validation.
+              if (requestLeaderEpoch == LeaderAndIsr.EpochDuringDelete ||
+                  requestLeaderEpoch == LeaderAndIsr.NoEpoch ||
+                  requestLeaderEpoch > currentLeaderEpoch) {
+                stoppedPartitions += topicPartition -> partitionState
+                // Assume that everything will go right. It is overwritten in case of an error.
+                responseMap.put(topicPartition, Errors.NONE)
+              } else if (requestLeaderEpoch < currentLeaderEpoch) {
+                stateChangeLogger.warn(s"Ignoring StopReplica request (delete=$deletePartition) 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.FENCED_LEADER_EPOCH)
+              } else {
+                stateChangeLogger.info(s"Ignoring StopReplica request (delete=$deletePartition) 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.FENCED_LEADER_EPOCH)
+              }
+
+            case HostedPartition.None =>
+              // Delete log and corresponding folders in case replica manager doesn't hold them anymore.
+              // This could happen when topic is being deleted while broker is down and recovers.
+              stoppedPartitions += topicPartition -> partitionState
+              responseMap.put(topicPartition, Errors.NONE)
+          }
+        }
+
+        // First stop fetchers for all partitions.
+        val partitions = stoppedPartitions.keySet
         replicaFetcherManager.removeFetcherForPartitions(partitions)
         replicaAlterLogDirsManager.removeFetcherForPartitions(partitions)
-        for (topicPartition <- partitions){
-          try {
-            stopReplica(topicPartition, stopReplicaRequest.deletePartitions)
-            responseMap.put(topicPartition, Errors.NONE)
-          } catch {
+
+        // Second remove deleted partitions from the partition map. Fetchers rely on the
+        // ReplicaManager to get Partition's information so they must be stopped first.
+        val deletedPartitions = mutable.Set.empty[TopicPartition]
+        stoppedPartitions.forKeyValue { (topicPartition, partitionState) =>
+          if (partitionState.deletePartition) {
+            getPartition(topicPartition) match {
+              case hostedPartition@HostedPartition.Online(partition) =>
+                if (allPartitions.remove(topicPartition, hostedPartition)) {
+                  maybeRemoveTopicMetrics(topicPartition.topic)
+                  // Logs are not deleted here. They are deleted in a single batch later on.
+                  // This is done to avoid having to checkpoint for every deletions.
+                  partition.delete()
+                }
+
+              case _ =>
+            }
+
+            deletedPartitions += topicPartition
+          }
+
+          // 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)
+        }
+
+        // Third delete the logs and checkpoint.
+        logManager.asyncDelete(deletedPartitions, (topicPartition, exception) => {
+          exception match {
             case e: KafkaStorageException =>
-              stateChangeLogger.error(s"Ignoring stop replica (delete=${stopReplicaRequest.deletePartitions}) for " +
-                s"partition $topicPartition due to storage exception", e)
+              stateChangeLogger.error(s"Ignoring StopReplica request (delete=true) 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)
+
+            case e =>
+              stateChangeLogger.error(s"Ignoring StopReplica request (delete=true) from " +
+                s"controller $controllerId with correlation id $correlationId " +
+                s"epoch $controllerEpoch for partition $topicPartition due to an unexpected " +
+                s"${e.getClass.getName} exception: ${e.getMessage}")
+              responseMap.put(topicPartition, Errors.forException(e))
           }
-        }
+        })
+
         (responseMap, Errors.NONE)
       }
     }
   }
 
-  def getOrCreatePartition(topicPartition: TopicPartition): Partition =
-    allPartitions.getAndMaybePut(topicPartition)
-
-  def getPartition(topicPartition: TopicPartition): Option[Partition] =
-    Option(allPartitions.get(topicPartition))
-
-  def nonOfflinePartition(topicPartition: TopicPartition): Option[Partition] =
-    getPartition(topicPartition).filter(_ ne ReplicaManager.OfflinePartition)
+  def getPartition(topicPartition: TopicPartition): HostedPartition = {
+    Option(allPartitions.get(topicPartition)).getOrElse(HostedPartition.None)
+  }
 
-  private def nonOfflinePartitionsIterator: Iterator[Partition] =
-    allPartitions.values.iterator.filter(_ ne ReplicaManager.OfflinePartition)
+  def isAddingReplica(topicPartition: TopicPartition, replicaId: Int): Boolean = {
+    getPartition(topicPartition) match {
+      case Online(partition) => partition.isAddingReplica(replicaId)
+      case _ => false
+    }
+  }
 
-  private def offlinePartitionsIterator: Iterator[Partition] =
-    allPartitions.values.iterator.filter(_ eq ReplicaManager.OfflinePartition)
+  // Visible for testing
+  def createPartition(topicPartition: TopicPartition): Partition = {
+    val partition = Partition(topicPartition, time, this)
+    allPartitions.put(topicPartition, HostedPartition.Online(partition))
+    partition
+  }
 
-  def getReplicaOrException(topicPartition: TopicPartition, brokerId: Int): Replica = {
+  def nonOfflinePartition(topicPartition: TopicPartition): Option[Partition] = {
     getPartition(topicPartition) match {
-      case Some(partition) =>
-        if (partition eq ReplicaManager.OfflinePartition)
-          throw new KafkaStorageException(s"Replica $brokerId is in an offline log directory for partition $topicPartition")
-        else
-          partition.getReplica(brokerId).getOrElse(
-            throw new ReplicaNotAvailableException(s"Replica $brokerId is not available for partition $topicPartition"))
-      case None =>
-        throw new ReplicaNotAvailableException(s"Replica $brokerId is not available for partition $topicPartition")
+      case HostedPartition.Online(partition) => Some(partition)
+      case HostedPartition.None | HostedPartition.Offline => None
     }
   }
 
-  def getReplicaOrException(topicPartition: TopicPartition): Replica = getReplicaOrException(topicPartition, localBrokerId)
+  // 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.flatMap {
+      case HostedPartition.Online(partition) => Some(partition)
+      case HostedPartition.None | HostedPartition.Offline => None
+    }
+  }
 
-  def getLeaderReplicaIfLocal(topicPartition: TopicPartition): Replica =  {
-    val (_, replica) = getPartitionAndLeaderReplicaIfLocal(topicPartition)
-    replica
+  private def offlinePartitionCount: Int = {
+    allPartitions.values.iterator.count(_ == HostedPartition.Offline)
   }
 
-  def getPartitionAndLeaderReplicaIfLocal(topicPartition: TopicPartition): (Partition, Replica) =  {
-    val partitionOpt = getPartition(topicPartition)
-    partitionOpt match {
-      case None =>
-        throw new UnknownTopicOrPartitionException(s"Partition $topicPartition doesn't exist on $localBrokerId")
-      case Some(partition) =>
-        if (partition eq ReplicaManager.OfflinePartition)
-          throw new KafkaStorageException(s"Partition $topicPartition is in an offline log directory on broker $localBrokerId")
-        else partition.leaderReplicaIfLocal match {
-          case Some(leaderReplica) => (partition, leaderReplica)
-          case None =>
-            throw new NotLeaderForPartitionException(s"Leader not local for partition $topicPartition on broker $localBrokerId")
-        }
+  def getPartitionOrException(topicPartition: TopicPartition): Partition = {
+    getPartitionOrError(topicPartition) 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 getReplica(topicPartition: TopicPartition, replicaId: Int): Option[Replica] =
-    nonOfflinePartition(topicPartition).flatMap(_.getReplica(replicaId))
+  def getPartitionOrError(topicPartition: TopicPartition): Either[Errors, Partition] = {
+    getPartition(topicPartition) match {
+      case HostedPartition.Online(partition) =>
+        Right(partition)
+
+      case HostedPartition.Offline =>
+        Left(Errors.KAFKA_STORAGE_ERROR)
 
-  def getReplica(tp: TopicPartition): Option[Replica] = getReplica(tp, localBrokerId)
+      case HostedPartition.None if metadataCache.contains(topicPartition) =>
+        // The topic exists, but this broker is no longer a replica of it, so we return NOT_LEADER_OR_FOLLOWER 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.
+        Left(Errors.NOT_LEADER_OR_FOLLOWER)
 
-  def getLogDir(topicPartition: TopicPartition): Option[String] = {
-    getReplica(topicPartition).flatMap(_.log) match {
-      case Some(log) => Some(log.dir.getParent)
-      case None => None
+      case HostedPartition.None =>
+        Left(Errors.UNKNOWN_TOPIC_OR_PARTITION)
     }
   }
 
+  def localLogOrException(topicPartition: TopicPartition): Log = {
+    getPartitionOrException(topicPartition).localLogOrException
+  }
+
+  def futureLocalLogOrException(topicPartition: TopicPartition): Log = {
+    getPartitionOrException(topicPartition).futureLocalLogOrException
+  }
+
+  def futureLogExists(topicPartition: TopicPartition): Boolean = {
+    getPartitionOrException(topicPartition).futureLog.isDefined
+  }
+
+  def localLog(topicPartition: TopicPartition): Option[Log] = {
+    nonOfflinePartition(topicPartition).flatMap(_.log)
+  }
+
+  def getLogDir(topicPartition: TopicPartition): Option[String] = {
+    localLog(topicPartition).map(_.parentDir)
+  }
+
+  /**
+   * TODO: move this action queue to handle thread so we can simplify concurrency handling
+   */
+  private val actionQueue = new ActionQueue
+
+  def tryCompleteActions(): Unit = actionQueue.tryCompleteActions()
+
   /**
    * Append messages to leader replicas of the partition, and wait for them to be replicated to other replicas;
    * the callback function will be triggered either when timeout or the required acks are satisfied;
    * if the callback function itself is already synchronized on some object then pass this object to avoid deadlock.
+   *
+   * Noted that all pending delayed check operations are stored in a queue. All callers to ReplicaManager.appendRecords()
+   * are expected to call ActionQueue.tryCompleteActions for all affected partitions, without holding any conflicting
+   * locks.
    */
   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, RecordsProcessingStats] => 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, 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
+      }
+
+      actionQueue.add {
+        () =>
+          localProduceResults.foreach {
+            case (topicPartition, result) =>
+              val requestKey = TopicPartitionOperationKey(topicPartition)
+              result.info.leaderHwChange match {
+                case LeaderHwChange.Increased =>
+                  // some delayed operations may be unblocked after HW changed
+                  delayedProducePurgatory.checkAndComplete(requestKey)
+                  delayedFetchPurgatory.checkAndComplete(requestKey)
+                  delayedDeleteRecordsPurgatory.checkAndComplete(requestKey)
+                case LeaderHwChange.Same =>
+                  // probably unblock some follower fetch requests since log end offset has been updated
+                  delayedFetchPurgatory.checkAndComplete(requestKey)
+                case LeaderHwChange.None =>
+                  // nothing
+              }
+          }
       }
 
-      processingStatsCallback(localProduceResults.mapValues(_.info.recordsProcessingStats))
+      recordConversionStatsCallback(localProduceResults.map { case (k, v) => k -> v.info.recordConversionStats })
 
       if (delayedProduceRequestRequired(requiredAcks, entriesPerPartition, localProduceResults)) {
         // create delayed produce operation
@@ -485,7 +658,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
@@ -494,7 +667,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 {
@@ -502,7 +675,7 @@ class ReplicaManager(val config: KafkaConfig,
       // Just return an error and don't handle the request at all
       val responseStatus = entriesPerPartition.map { case (topicPartition, _) =>
         topicPartition -> new PartitionResponse(Errors.INVALID_REQUIRED_ACKS,
-          LogAppendInfo.UnknownLogAppendInfo.firstOffset, RecordBatch.NO_TIMESTAMP, LogAppendInfo.UnknownLogAppendInfo.logStartOffset)
+          LogAppendInfo.UnknownLogAppendInfo.firstOffset.getOrElse(-1), RecordBatch.NO_TIMESTAMP, LogAppendInfo.UnknownLogAppendInfo.logStartOffset)
       }
       responseCallback(responseStatus)
     }
@@ -520,20 +693,12 @@ class ReplicaManager(val config: KafkaConfig,
         (topicPartition, LogDeleteRecordsResult(-1L, -1L, Some(new InvalidTopicException(s"Cannot delete records of internal topic ${topicPartition.topic}"))))
       } else {
         try {
-          val (partition, replica) = getPartitionAndLeaderReplicaIfLocal(topicPartition)
-          val convertedOffset =
-            if (requestedOffset == DeleteRecordsRequest.HIGH_WATERMARK) {
-              replica.highWatermark.messageOffset
-            } else
-              requestedOffset
-          if (convertedOffset < 0)
-            throw new OffsetOutOfRangeException(s"The offset $convertedOffset for partition $topicPartition is not valid")
-
-          val lowWatermark = partition.deleteRecordsOnLeader(convertedOffset)
-          (topicPartition, LogDeleteRecordsResult(convertedOffset, lowWatermark))
+          val partition = getPartitionOrException(topicPartition)
+          val logDeleteResult = partition.deleteRecordsOnLeader(requestedOffset)
+          (topicPartition, logDeleteResult)
         } catch {
           case e@ (_: UnknownTopicOrPartitionException |
-                   _: NotLeaderForPartitionException |
+                   _: NotLeaderOrFollowerException |
                    _: OffsetOutOfRangeException |
                    _: PolicyViolationException |
                    _: KafkaStorageException) =>
@@ -552,61 +717,78 @@ class ReplicaManager(val config: KafkaConfig,
   // 1. the delete records operation on this partition is successful
   // 2. low watermark of this partition is smaller than the specified offset
   private def delayedDeleteRecordsRequired(localDeleteRecordsResults: Map[TopicPartition, LogDeleteRecordsResult]): Boolean = {
-    localDeleteRecordsResults.exists{ case (tp, deleteRecordsResult) =>
+    localDeleteRecordsResults.exists{ case (_, deleteRecordsResult) =>
       deleteRecordsResult.exception.isEmpty && deleteRecordsResult.lowWatermark < deleteRecordsResult.requestedOffset
     }
   }
 
-  /*
+  /**
    * For each pair of partition and log directory specified in the map, if the partition has already been created on
    * this broker, move its log files to the specified log directory. Otherwise, record the pair in the memory so that
    * the partition will be created in the specified log directory when broker receives LeaderAndIsrRequest for the partition later.
-   *
    */
   def alterReplicaLogDirs(partitionDirs: Map[TopicPartition, String]): Map[TopicPartition, Errors] = {
     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")
 
-          // Stop current replica movement if the destinationDir is different from the existing destination log directory
-          getReplica(topicPartition, Request.FutureLocalReplicaId) match {
-            case Some(futureReplica) =>
-              if (futureReplica.log.get.dir.getParent != destinationDir) {
+          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))
-                getPartition(topicPartition).get.removeFutureLocalReplica()
-                logManager.asyncDelete(topicPartition, isFuture = true)
+                partition.removeFutureLocalReplica()
               }
-            case None =>
+            case HostedPartition.Offline =>
+              throw new KafkaStorageException(s"Partition $topicPartition is offline")
+
+            case HostedPartition.None => // Do nothing
           }
 
           // If the log for this partition has not been created yet:
           // 1) Record the destination log directory in the memory so that the partition will be created in this log directory
           //    when broker receives LeaderAndIsrRequest for this partition later.
-          // 2) Respond with ReplicaNotAvailableException for this partition in the AlterReplicaLogDirsResponse
+          // 2) Respond with NotLeaderOrFollowerException for this partition in the AlterReplicaLogDirsResponse
           logManager.maybeUpdatePreferredLogDir(topicPartition, destinationDir)
-          // throw ReplicaNotAvailableException if replica does not exit for the given partition
-          getReplicaOrException(topicPartition)
+
+          // throw NotLeaderOrFollowerException if replica does not exist for the given partition
+          val partition = getPartitionOrException(topicPartition)
+          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 (getPartition(topicPartition).get.maybeCreateFutureReplica(destinationDir)) {
-            val futureReplica = getReplicaOrException(topicPartition, Request.FutureLocalReplicaId)
+          val highWatermarkCheckpoints = new LazyOffsetCheckpoints(this.highWatermarkCheckpoints)
+          if (partition.maybeCreateFutureReplica(destinationDir, highWatermarkCheckpoints)) {
+            val futureLog = futureLocalLogOrException(topicPartition)
             logManager.abortAndPauseCleaning(topicPartition)
-            replicaAlterLogDirsManager.addFetcherForPartitions(Map(topicPartition -> BrokerAndInitialOffset(
-              BrokerEndPoint(config.brokerId, "localhost", -1), futureReplica.highWatermark.messageOffset)))
+
+            val initialFetchState = InitialFetchState(BrokerEndPoint(config.brokerId, "localhost", -1),
+              partition.getLeaderEpoch, futureLog.highWatermark)
+            replicaAlterLogDirsManager.addFetcherForPartitions(Map(topicPartition -> initialFetchState))
           }
 
           (topicPartition, Errors.NONE)
         } catch {
-          case e@(_: LogDirNotFoundException |
+          case e@(_: InvalidTopicException |
+                  _: LogDirNotFoundException |
                   _: ReplicaNotAvailableException |
                   _: KafkaStorageException) =>
+            warn(s"Unable to alter log dirs for $topicPartition", e)
             (topicPartition, Errors.forException(e))
+          case e: NotLeaderOrFollowerException =>
+            // Retaining REPLICA_NOT_AVAILABLE exception for ALTER_REPLICA_LOG_DIRS for compatibility
+            warn(s"Unable to alter log dirs for $topicPartition", e)
+            (topicPartition, Errors.REPLICA_NOT_AVAILABLE)
           case t: Throwable =>
             error("Error while changing replica dir for partition %s".format(topicPartition), t)
             (topicPartition, Errors.forException(t))
@@ -623,8 +805,8 @@ class ReplicaManager(val config: KafkaConfig,
    * 2) size and lag of current and future logs for each partition in the given log directory. Only logs of the queried partitions
    *    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)
+  def describeLogDirs(partitions: Set[TopicPartition]): List[DescribeLogDirsResponseData.DescribeLogDirsResult] = {
+    val logsByDir = logManager.allLogs.groupBy(log => log.parentDir)
 
     config.logDirs.toSet.map { logDir: String =>
       val absolutePath = new File(logDir).getAbsolutePath
@@ -634,34 +816,48 @@ class ReplicaManager(val config: KafkaConfig,
 
         logsByDir.get(absolutePath) match {
           case Some(logs) =>
-            val replicaInfos = logs.filter { log =>
-              partitions.contains(log.topicPartition)
-            }.map { log =>
-              log.topicPartition -> new ReplicaInfo(log.size, getLogEndOffsetLag(log.topicPartition, log.logEndOffset, log.isFuture), log.isFuture)
-            }.toMap
-
-            (absolutePath, new LogDirInfo(Errors.NONE, replicaInfos.asJava))
+            val topicInfos = logs.groupBy(_.topicPartition.topic).map{case (topic, logs) =>
+              new DescribeLogDirsResponseData.DescribeLogDirsTopic().setName(topic).setPartitions(
+                logs.filter { log =>
+                  partitions.contains(log.topicPartition)
+                }.map { log =>
+                  new DescribeLogDirsResponseData.DescribeLogDirsPartition()
+                    .setPartitionSize(log.size)
+                    .setPartitionIndex(log.topicPartition.partition)
+                    .setOffsetLag(getLogEndOffsetLag(log.topicPartition, log.logEndOffset, log.isFuture))
+                    .setIsFutureKey(log.isFuture)
+                }.toList.asJava)
+            }.toList.asJava
+
+            new DescribeLogDirsResponseData.DescribeLogDirsResult().setLogDir(absolutePath)
+              .setErrorCode(Errors.NONE.code).setTopics(topicInfos)
           case None =>
-            (absolutePath, new LogDirInfo(Errors.NONE, Map.empty[TopicPartition, ReplicaInfo].asJava))
+            new DescribeLogDirsResponseData.DescribeLogDirsResult().setLogDir(absolutePath)
+              .setErrorCode(Errors.NONE.code)
         }
 
       } catch {
         case e: KafkaStorageException =>
-          (absolutePath, new LogDirInfo(Errors.KAFKA_STORAGE_ERROR, Map.empty[TopicPartition, ReplicaInfo].asJava))
+          warn("Unable to describe replica dirs for %s".format(absolutePath), e)
+          new DescribeLogDirsResponseData.DescribeLogDirsResult()
+            .setLogDir(absolutePath)
+            .setErrorCode(Errors.KAFKA_STORAGE_ERROR.code)
         case t: Throwable =>
           error(s"Error while describing replica in dir $absolutePath", t)
-          (absolutePath, new LogDirInfo(Errors.forException(t), Map.empty[TopicPartition, ReplicaInfo].asJava))
+          new DescribeLogDirsResponseData.DescribeLogDirsResult()
+            .setLogDir(absolutePath)
+            .setErrorCode(Errors.forException(t).code)
       }
-    }.toMap
+    }.toList
   }
 
   def getLogEndOffsetLag(topicPartition: TopicPartition, logEndOffset: Long, isFuture: Boolean): Long = {
-    getReplica(topicPartition) match {
-      case Some(replica) =>
+    localLog(topicPartition) match {
+      case Some(log) =>
         if (isFuture)
-          replica.logEndOffset.messageOffset - 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
@@ -670,7 +866,7 @@ class ReplicaManager(val config: KafkaConfig,
 
   def deleteRecords(timeout: Long,
                     offsetPerPartition: Map[TopicPartition, Long],
-                    responseCallback: Map[TopicPartition, DeleteRecordsResponse.PartitionResponse] => Unit) {
+                    responseCallback: Map[TopicPartition, DeleteRecordsPartitionResult] => Unit): Unit = {
     val timeBeforeLocalDeleteRecords = time.milliseconds
     val localDeleteRecordsResults = deleteRecordsOnLocalLog(offsetPerPartition)
     debug("Delete records on local log in %d ms".format(time.milliseconds - timeBeforeLocalDeleteRecords))
@@ -679,7 +875,10 @@ class ReplicaManager(val config: KafkaConfig,
       topicPartition ->
         DeleteRecordsPartitionStatus(
           result.requestedOffset, // requested offset
-          new DeleteRecordsResponse.PartitionResponse(result.lowWatermark, result.error)) // response status
+          new DeleteRecordsPartitionResult()
+            .setLowWatermark(result.lowWatermark)
+            .setErrorCode(result.error.code)
+            .setPartitionIndex(topicPartition.partition)) // response status
     }
 
     if (delayedDeleteRecordsRequired(localDeleteRecordsResults)) {
@@ -687,7 +886,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
@@ -695,7 +894,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)
     }
   }
@@ -721,10 +920,25 @@ 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] = {
-    trace("Append [%s] to local log ".format(entriesPerPartition))
+    val traceEnabled = isTraceEnabled
+    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
+    }
+
+    if (traceEnabled)
+      trace(s"Append [$entriesPerPartition] to local log")
+
     entriesPerPartition.map { case (topicPartition, records) =>
       brokerTopicStats.topicStats(topicPartition.topic).totalProduceRequestRate.mark()
       brokerTopicStats.allTopicsStats.totalProduceRequestRate.mark()
@@ -736,22 +950,9 @@ class ReplicaManager(val config: KafkaConfig,
           Some(new InvalidTopicException(s"Cannot append to internal topic ${topicPartition.topic}"))))
       } else {
         try {
-          val partitionOpt = getPartition(topicPartition)
-          val info = partitionOpt match {
-            case Some(partition) =>
-              if (partition eq ReplicaManager.OfflinePartition)
-                throw new KafkaStorageException(s"Partition $topicPartition is in an offline log directory on broker $localBrokerId")
-              partition.appendRecordsToLeader(records, isFromClient, requiredAcks)
-
-            case None => throw new UnknownTopicOrPartitionException("Partition %s doesn't exist on %d"
-              .format(topicPartition, localBrokerId))
-          }
-
-          val numAppendedMessages =
-            if (info.firstOffset == -1L || info.lastOffset == -1L)
-              0
-            else
-              info.lastOffset - info.firstOffset + 1
+          val partition = getPartitionOrException(topicPartition)
+          val info = partition.appendRecordsToLeader(records, origin, requiredAcks)
+          val numAppendedMessages = info.numMessages
 
           // update stats for successfully appended bytes and messages as bytesInRate and messageInRate
           brokerTopicStats.topicStats(topicPartition.topic).bytesInRate.mark(records.sizeInBytes)
@@ -759,39 +960,56 @@ class ReplicaManager(val config: KafkaConfig,
           brokerTopicStats.topicStats(topicPartition.topic).messagesInRate.mark(numAppendedMessages)
           brokerTopicStats.allTopicsStats.messagesInRate.mark(numAppendedMessages)
 
-          trace("%d bytes written to log %s-%d beginning at offset %d and ending at offset %d"
-            .format(records.sizeInBytes, topicPartition.topic, topicPartition.partition, info.firstOffset, info.lastOffset))
+          if (traceEnabled)
+            trace(s"${records.sizeInBytes} written to log $topicPartition beginning at offset " +
+              s"${info.firstOffset.getOrElse(-1)} and ending at offset ${info.lastOffset}")
+
           (topicPartition, LogAppendResult(info))
         } catch {
           // NOTE: Failed produce requests metric is not incremented for known exceptions
           // it is supposed to indicate un-expected failures of a broker in handling a produce request
           case e@ (_: UnknownTopicOrPartitionException |
-                   _: NotLeaderForPartitionException |
+                   _: NotLeaderOrFollowerException |
                    _: 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("Error processing append operation on partition %s".format(topicPartition), t)
+            val logStartOffset = processFailedRecord(topicPartition, t)
             (topicPartition, LogAppendResult(LogAppendInfo.unknownLogAppendInfoWithLogStartOffset(logStartOffset), Some(t)))
         }
       }
     }
   }
 
+  def fetchOffsetForTimestamp(topicPartition: TopicPartition,
+                              timestamp: Long,
+                              isolationLevel: Option[IsolationLevel],
+                              currentLeaderEpoch: Optional[Integer],
+                              fetchOnlyFromLeader: Boolean): Option[TimestampAndOffset] = {
+    val partition = getPartitionOrException(topicPartition)
+    partition.fetchOffsetForTimestamp(timestamp, isolationLevel, currentLeaderEpoch, fetchOnlyFromLeader)
+  }
+
+  def legacyFetchOffsetsForTimestamp(topicPartition: TopicPartition,
+                                     timestamp: Long,
+                                     maxNumOffsets: Int,
+                                     isFromConsumer: Boolean,
+                                     fetchOnlyFromLeader: Boolean): Seq[Long] = {
+    val partition = getPartitionOrException(topicPartition)
+    partition.legacyFetchOffsetsForTimestamp(timestamp, maxNumOffsets, isFromConsumer, fetchOnlyFromLeader)
+  }
+
   /**
-   * 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,
@@ -799,59 +1017,90 @@ 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 fetchOnlyCommitted = !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,
         fetchOnlyFromLeader = fetchOnlyFromLeader,
-        readOnlyCommitted = fetchOnlyCommitted,
+        fetchIsolation = fetchIsolation,
         fetchMaxBytes = fetchMaxBytes,
         hardMaxBytesLimit = hardMaxBytesLimit,
         readPartitionInfo = fetchInfos,
         quota = quota,
-        isolationLevel = isolationLevel)
-      if (isFromFollower) updateFollowerLogReadResults(replicaId, result)
+        clientMetadata = clientMetadata)
+      if (isFromFollower) updateFollowerFetchState(replicaId, result)
       else result
     }
 
     val logReadResults = readFromLog()
 
     // check if this fetch request can be satisfied right away
-    val logReadResultValues = logReadResults.map { case (_, v) => v }
-    val bytesReadable = logReadResultValues.map(_.info.records.sizeInBytes).sum
-    val errorReadingData = logReadResultValues.foldLeft(false) ((errorIncurred, readResult) =>
-      errorIncurred || (readResult.error != Errors.NONE))
+    var bytesReadable: Long = 0
+    var errorReadingData = false
+    var hasDivergingEpoch = false
+    val logReadResultMap = new mutable.HashMap[TopicPartition, LogReadResult]
+    logReadResults.foreach { case (topicPartition, logReadResult) =>
+      brokerTopicStats.topicStats(topicPartition.topic).totalFetchRequestRate.mark()
+      brokerTopicStats.allTopicsStats.totalFetchRequestRate.mark()
+
+      if (logReadResult.error != Errors.NONE)
+        errorReadingData = true
+      if (logReadResult.divergingEpoch.nonEmpty)
+        hasDivergingEpoch = true
+      bytesReadable = bytesReadable + logReadResult.info.records.sizeInBytes
+      logReadResultMap.put(topicPartition, logReadResult)
+    }
 
     // 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) we found a diverging epoch
+    if (timeout <= 0 || fetchInfos.isEmpty || bytesReadable >= fetchMinBytes || errorReadingData || hasDivergingEpoch) {
       val fetchPartitionData = logReadResults.map { case (tp, result) =>
-        tp -> FetchPartitionData(result.error, result.highWatermark, result.leaderLogStartOffset, result.info.records,
-          result.lastStableOffset, result.info.abortedTransactions)
+        val isReassignmentFetch = isFromFollower && isAddingReplica(tp, replicaId)
+        tp -> FetchPartitionData(
+          result.error,
+          result.highWatermark,
+          result.leaderLogStartOffset,
+          result.info.records,
+          result.divergingEpoch,
+          result.lastStableOffset,
+          result.info.abortedTransactions,
+          result.preferredReadReplica,
+          isReassignmentFetch)
       }
       responseCallback(fetchPartitionData)
     } else {
       // construct the fetch results from the read results
-      val fetchPartitionStatus = logReadResults.map { case (topicPartition, result) =>
-        val fetchInfo = fetchInfos.collectFirst {
-          case (tp, v) if tp == topicPartition => v
-        }.getOrElse(sys.error(s"Partition $topicPartition not found in fetchInfos"))
-        (topicPartition, FetchPartitionStatus(result.info.fetchOffsetMetadata, fetchInfo))
+      val fetchPartitionStatus = new mutable.ArrayBuffer[(TopicPartition, FetchPartitionStatus)]
+      fetchInfos.foreach { case (topicPartition, partitionData) =>
+        logReadResultMap.get(topicPartition).foreach(logReadResult => {
+          val logOffsetMetadata = logReadResult.info.fetchOffsetMetadata
+          fetchPartitionStatus += (topicPartition -> FetchPartitionStatus(logOffsetMetadata, partitionData))
+        })
       }
-      val fetchMetadata = FetchMetadata(fetchMinBytes, fetchMaxBytes, hardMaxBytesLimit, fetchOnlyFromLeader,
-        fetchOnlyCommitted, isFromFollower, replicaId, fetchPartitionStatus)
-      val delayedFetch = new DelayedFetch(timeout, fetchMetadata, this, quota, isolationLevel, responseCallback)
+      val fetchMetadata: SFetchMetadata = SFetchMetadata(fetchMinBytes, fetchMaxBytes, hardMaxBytesLimit,
+        fetchOnlyFromLeader, fetchIsolation, isFromFollower, replicaId, fetchPartitionStatus)
+      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
@@ -865,113 +1114,119 @@ class ReplicaManager(val config: KafkaConfig,
    */
   def readFromLocalLog(replicaId: Int,
                        fetchOnlyFromLeader: Boolean,
-                       readOnlyCommitted: Boolean,
+                       fetchIsolation: FetchIsolation,
                        fetchMaxBytes: Int,
                        hardMaxBytesLimit: Boolean,
                        readPartitionInfo: Seq[(TopicPartition, PartitionData)],
                        quota: ReplicaQuota,
-                       isolationLevel: IsolationLevel): Seq[(TopicPartition, LogReadResult)] = {
+                       clientMetadata: Option[ClientMetadata]): Seq[(TopicPartition, LogReadResult)] = {
+    val traceEnabled = isTraceEnabled
 
     def read(tp: TopicPartition, fetchInfo: PartitionData, limitBytes: Int, minOneMessage: Boolean): LogReadResult = {
       val offset = fetchInfo.fetchOffset
       val partitionFetchSize = fetchInfo.maxBytes
       val followerLogStartOffset = fetchInfo.logStartOffset
 
-      brokerTopicStats.topicStats(tp.topic).totalFetchRequestRate.mark()
-      brokerTopicStats.allTopicsStats.totalFetchRequestRate.mark()
-
+      val adjustedMaxBytes = math.min(fetchInfo.maxBytes, limitBytes)
       try {
-        trace(s"Fetching log segment for partition $tp, offset $offset, partition fetch size $partitionFetchSize, " +
-          s"remaining response limit $limitBytes" +
-          (if (minOneMessage) s", ignoring response/partition size limits" else ""))
-
-        // decide whether to only fetch from leader
-        val localReplica = if (fetchOnlyFromLeader)
-          getLeaderReplicaIfLocal(tp)
-        else
-          getReplicaOrException(tp)
-
-        val initialHighWatermark = localReplica.highWatermark.messageOffset
-        val lastStableOffset = if (isolationLevel == IsolationLevel.READ_COMMITTED)
-          Some(localReplica.lastStableOffset.messageOffset)
-        else
-          None
-
-        // decide whether to only fetch committed data (i.e. messages below high watermark)
-        val maxOffsetOpt = if (readOnlyCommitted)
-          Some(lastStableOffset.getOrElse(initialHighWatermark))
-        else
-          None
+        if (traceEnabled)
+          trace(s"Fetching log segment for partition $tp, offset $offset, partition fetch size $partitionFetchSize, " +
+            s"remaining response limit $limitBytes" +
+            (if (minOneMessage) s", ignoring response/partition size limits" else ""))
 
-        /* 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 initialLogEndOffset = localReplica.logEndOffset.messageOffset
-        val initialLogStartOffset = localReplica.logStartOffset
+        val partition = getPartitionOrException(tp)
         val fetchTimeMs = time.milliseconds
-        val logReadInfo = localReplica.log match {
-          case Some(log) =>
-            val adjustedFetchSize = math.min(partitionFetchSize, limitBytes)
 
-            // Try the read first, this tells us whether we need all of adjustedFetchSize for this partition
-            val fetch = log.read(offset, adjustedFetchSize, maxOffsetOpt, minOneMessage, isolationLevel)
+        // If we are the leader, determine the preferred read-replica
+        val preferredReadReplica = clientMetadata.flatMap(
+          metadata => findPreferredReadReplica(partition, 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),
+            divergingEpoch = None,
+            highWatermark = offsetSnapshot.highWatermark.messageOffset,
+            leaderLogStartOffset = offsetSnapshot.logStartOffset,
+            leaderLogEndOffset = offsetSnapshot.logEndOffset.messageOffset,
+            followerLogStartOffset = followerLogStartOffset,
+            fetchTimeMs = -1L,
+            lastStableOffset = Some(offsetSnapshot.lastStableOffset.messageOffset),
+            preferredReadReplica = preferredReadReplica,
+            exception = None)
+        } else {
+          // Try the read first, this tells us whether we need all of adjustedFetchSize for this partition
+          val readInfo: LogReadInfo = partition.readRecords(
+            lastFetchedEpoch = fetchInfo.lastFetchedEpoch,
+            fetchOffset = fetchInfo.fetchOffset,
+            currentLeaderEpoch = fetchInfo.currentLeaderEpoch,
+            maxBytes = adjustedMaxBytes,
+            fetchIsolation = fetchIsolation,
+            fetchOnlyFromLeader = fetchOnlyFromLeader,
+            minOneMessage = minOneMessage)
+
+          val fetchDataInfo = if (shouldLeaderThrottle(quota, partition, replicaId)) {
             // If the partition is being throttled, simply return an empty set.
-            if (shouldLeaderThrottle(quota, tp, replicaId))
-              FetchDataInfo(fetch.fetchOffsetMetadata, MemoryRecords.EMPTY)
+            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`
-            else if (!hardMaxBytesLimit && fetch.firstEntryIncomplete)
-              FetchDataInfo(fetch.fetchOffsetMetadata, MemoryRecords.EMPTY)
-            else fetch
+            FetchDataInfo(readInfo.fetchedData.fetchOffsetMetadata, MemoryRecords.EMPTY)
+          } else {
+            readInfo.fetchedData
+          }
 
-          case None =>
-            error(s"Leader for partition $tp does not have a local log")
-            FetchDataInfo(LogOffsetMetadata.UnknownOffsetMetadata, MemoryRecords.EMPTY)
+          LogReadResult(info = fetchDataInfo,
+            divergingEpoch = readInfo.divergingEpoch,
+            highWatermark = readInfo.highWatermark,
+            leaderLogStartOffset = readInfo.logStartOffset,
+            leaderLogEndOffset = readInfo.logEndOffset,
+            followerLogStartOffset = followerLogStartOffset,
+            fetchTimeMs = fetchTimeMs,
+            lastStableOffset = Some(readInfo.lastStableOffset),
+            preferredReadReplica = preferredReadReplica,
+            exception = None)
         }
-
-        LogReadResult(info = logReadInfo,
-                      highWatermark = initialHighWatermark,
-                      leaderLogStartOffset = initialLogStartOffset,
-                      leaderLogEndOffset = initialLogEndOffset,
-                      followerLogStartOffset = followerLogStartOffset,
-                      fetchTimeMs = fetchTimeMs,
-                      readSize = partitionFetchSize,
-                      lastStableOffset = lastStableOffset,
-                      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
         case e@ (_: UnknownTopicOrPartitionException |
-                 _: NotLeaderForPartitionException |
+                 _: NotLeaderOrFollowerException |
+                 _: UnknownLeaderEpochException |
+                 _: FencedLeaderEpochException |
                  _: ReplicaNotAvailableException |
                  _: KafkaStorageException |
                  _: OffsetOutOfRangeException) =>
           LogReadResult(info = FetchDataInfo(LogOffsetMetadata.UnknownOffsetMetadata, MemoryRecords.EMPTY),
-                        highWatermark = -1L,
-                        leaderLogStartOffset = -1L,
-                        leaderLogEndOffset = -1L,
-                        followerLogStartOffset = -1L,
-                        fetchTimeMs = -1L,
-                        readSize = partitionFetchSize,
-                        lastStableOffset = None,
-                        exception = Some(e))
+            divergingEpoch = None,
+            highWatermark = Log.UnknownOffset,
+            leaderLogStartOffset = Log.UnknownOffset,
+            leaderLogEndOffset = Log.UnknownOffset,
+            followerLogStartOffset = Log.UnknownOffset,
+            fetchTimeMs = -1L,
+            lastStableOffset = None,
+            exception = Some(e))
         case e: Throwable =>
           brokerTopicStats.topicStats(tp.topic).failedFetchRequestRate.mark()
           brokerTopicStats.allTopicsStats.failedFetchRequestRate.mark()
-          error(s"Error processing fetch operation on partition $tp, offset $offset", e)
+
+          val fetchSource = Request.describeReplicaId(replicaId)
+          error(s"Error processing fetch with max size $adjustedMaxBytes from $fetchSource " +
+            s"on partition $tp: $fetchInfo", e)
+
           LogReadResult(info = FetchDataInfo(LogOffsetMetadata.UnknownOffsetMetadata, MemoryRecords.EMPTY),
-                        highWatermark = -1L,
-                        leaderLogStartOffset = -1L,
-                        leaderLogEndOffset = -1L,
-                        followerLogStartOffset = -1L,
-                        fetchTimeMs = -1L,
-                        readSize = partitionFetchSize,
-                        lastStableOffset = None,
-                        exception = Some(e))
+            divergingEpoch = None,
+            highWatermark = Log.UnknownOffset,
+            leaderLogStartOffset = Log.UnknownOffset,
+            leaderLogEndOffset = Log.UnknownOffset,
+            followerLogStartOffset = Log.UnknownOffset,
+            fetchTimeMs = -1L,
+            lastStableOffset = None,
+            exception = Some(e))
       }
     }
 
@@ -990,19 +1245,60 @@ 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(partition: Partition,
+                               clientMetadata: ClientMetadata,
+                               replicaId: Int,
+                               fetchOffset: Long,
+                               currentTimeMs: Long): Option[Int] = {
+    partition.leaderReplicaIdOpt.flatMap { leaderReplicaId =>
+      // Don't look up preferred for follower fetches via normal replication
+      if (Request.isValidBrokerId(replicaId))
+        None
+      else {
+        replicaSelectorOpt.flatMap { replicaSelector =>
+          val replicaEndpoints = metadataCache.getPartitionReplicaEndpoints(partition.topicPartition,
+            new ListenerName(clientMetadata.listenerName))
+          val replicaInfos = 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 && replica.logStartOffset <= fetchOffset)
+            .map(replica => new DefaultReplicaView(
+              replicaEndpoints.getOrElse(replica.brokerId, Node.noNode()),
+              replica.logEndOffset,
+              currentTimeMs - replica.lastCaughtUpTimeMs))
+
+          val leaderReplica = new DefaultReplicaView(
+            replicaEndpoints.getOrElse(leaderReplicaId, Node.noNode()),
+            partition.localLogOrException.logEndOffset, 0L)
+          val replicaInfoSet = mutable.Set[ReplicaView]() ++= replicaInfos += leaderReplica
+
+          val partitionInfo = new DefaultPartitionView(replicaInfoSet.asJava, leaderReplica)
+          replicaSelector.select(partition.topicPartition, clientMetadata, partitionInfo).asScala.collect {
+            // 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
+            case selected if !selected.endpoint.isEmpty && selected != leaderReplica => selected.endpoint.id
+          }
+        }
+      }
+    }
+  }
+
   /**
    *  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)
-    }
-    quota.isThrottled(topicPartition) && quota.isQuotaExceeded && !isReplicaInSync
+  def shouldLeaderThrottle(quota: ReplicaQuota, partition: Partition, replicaId: Int): Boolean = {
+    val isReplicaInSync = partition.inSyncReplicaIds.contains(replicaId)
+    !isReplicaInSync && quota.isThrottled(partition.topicPartition) && quota.isQuotaExceeded
   }
 
-  def getMagic(topicPartition: TopicPartition): Option[Byte] =
-    getReplica(topicPartition).flatMap(_.log.map(_.config.messageFormatVersion.messageFormatVersion))
+  def getLogConfig(topicPartition: TopicPartition): Option[LogConfig] = localLog(topicPartition).map(_.config)
+
+  def getMagic(topicPartition: TopicPartition): Option[Byte] = getLogConfig(topicPartition).map(_.messageFormatVersion.recordVersion.value)
 
   def maybeUpdateMetadataCache(correlationId: Int, updateMetadataRequest: UpdateMetadataRequest) : Seq[TopicPartition] =  {
     replicaStateChangeLock synchronized {
@@ -1013,7 +1309,7 @@ class ReplicaManager(val config: KafkaConfig,
         stateChangeLogger.warn(stateControllerEpochErrorMessage)
         throw new ControllerMovedException(stateChangeLogger.messageWithPrefix(stateControllerEpochErrorMessage))
       } else {
-        val deletedPartitions = metadataCache.updateCache(correlationId, updateMetadataRequest)
+        val deletedPartitions = metadataCache.updateMetadata(correlationId, updateMetadataRequest)
         controllerEpoch = updateMetadataRequest.controllerEpoch
         deletedPartitions
       }
@@ -1023,108 +1319,181 @@ 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")
-    }
+    val startMs = time.milliseconds()
     replicaStateChangeLock synchronized {
-      if (leaderAndIsrRequest.controllerEpoch < controllerEpoch) {
-        stateChangeLogger.warn(s"Ignoring LeaderAndIsr request from controller ${leaderAndIsrRequest.controllerId} with " +
-          s"correlation id $correlationId since its controller epoch ${leaderAndIsrRequest.controllerEpoch} is old. " +
-          s"Latest known controller epoch is $controllerEpoch")
-        leaderAndIsrRequest.getErrorResponse(0, Errors.STALE_CONTROLLER_EPOCH.exception)
-      } else {
-        val responseMap = new mutable.HashMap[TopicPartition, Errors]
-        val controllerId = leaderAndIsrRequest.controllerId
-        controllerEpoch = leaderAndIsrRequest.controllerEpoch
-
-        // First check partition's leader epoch
-        val partitionState = new mutable.HashMap[Partition, LeaderAndIsrRequest.PartitionState]()
-        val newPartitions = leaderAndIsrRequest.partitionStates.asScala.keys.filter(topicPartition => getPartition(topicPartition).isEmpty)
-
-        leaderAndIsrRequest.partitionStates.asScala.foreach { case (topicPartition, stateInfo) =>
-          val partition = getOrCreatePartition(topicPartition)
-          val partitionLeaderEpoch = partition.getLeaderEpoch
-          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 (partitionLeaderEpoch < stateInfo.basePartitionState.leaderEpoch) {
-            // 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)
+      val controllerId = leaderAndIsrRequest.controllerId
+      val requestPartitionStates = leaderAndIsrRequest.partitionStates.asScala
+      stateChangeLogger.info(s"Handling LeaderAndIsr request correlationId $correlationId from controller " +
+        s"$controllerId for ${requestPartitionStates.size} partitions")
+      if (stateChangeLogger.isTraceEnabled)
+        requestPartitionStates.foreach { partitionState =>
+          stateChangeLogger.trace(s"Received LeaderAndIsr request $partitionState " +
+            s"correlation id $correlationId from controller $controllerId " +
+            s"epoch ${leaderAndIsrRequest.controllerEpoch}")
+        }
+
+      val response = {
+        if (leaderAndIsrRequest.controllerEpoch < controllerEpoch) {
+          stateChangeLogger.warn(s"Ignoring LeaderAndIsr request from controller $controllerId with " +
+            s"correlation id $correlationId since its controller epoch ${leaderAndIsrRequest.controllerEpoch} is old. " +
+            s"Latest known controller epoch is $controllerEpoch")
+          leaderAndIsrRequest.getErrorResponse(0, Errors.STALE_CONTROLLER_EPOCH.exception)
+        } else {
+          val responseMap = new mutable.HashMap[TopicPartition, Errors]
+          controllerEpoch = leaderAndIsrRequest.controllerEpoch
+
+          val partitionStates = new mutable.HashMap[Partition, LeaderAndIsrPartitionState]()
+
+          // First create the partition if it doesn't exist already
+          requestPartitionStates.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) =>
+                Some(partition)
+
+              case HostedPartition.None =>
+                val partition = Partition(topicPartition, time, this)
+                allPartitions.putIfNotExists(topicPartition, HostedPartition.Online(partition))
+                Some(partition)
+            }
+
+            // Next check partition's leader epoch
+            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.info(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 ${stateInfo.basePartitionState.leaderEpoch} is not higher than the current " +
-              s"leader epoch $partitionLeaderEpoch")
-            responseMap.put(topicPartition, Errors.STALE_CONTROLLER_EPOCH)
           }
-        }
 
-        val partitionsTobeLeader = partitionState.filter { case (_, stateInfo) =>
-          stateInfo.basePartitionState.leader == localBrokerId
-        }
-        val partitionsToBeFollower = partitionState -- partitionsTobeLeader.keys
+          val partitionsToBeLeader = partitionStates.filter { case (_, partitionState) =>
+            partitionState.leader == localBrokerId
+          }
+          val partitionsToBeFollower = partitionStates.filter { case (k, _) => !partitionsToBeLeader.contains(k) }
 
-        val partitionsBecomeLeader = if (partitionsTobeLeader.nonEmpty)
-          makeLeaders(controllerId, controllerEpoch, partitionsTobeLeader, correlationId, responseMap)
-        else
-          Set.empty[Partition]
-        val partitionsBecomeFollower = if (partitionsToBeFollower.nonEmpty)
-          makeFollowers(controllerId, controllerEpoch, partitionsToBeFollower, correlationId, responseMap)
-        else
-          Set.empty[Partition]
+          val highWatermarkCheckpoints = new LazyOffsetCheckpoints(this.highWatermarkCheckpoints)
+          val partitionsBecomeLeader = if (partitionsToBeLeader.nonEmpty)
+            makeLeaders(controllerId, controllerEpoch, partitionsToBeLeader, correlationId, responseMap,
+              highWatermarkCheckpoints)
+          else
+            Set.empty[Partition]
+          val partitionsBecomeFollower = if (partitionsToBeFollower.nonEmpty)
+            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.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 (getReplica(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
+          // 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
+          startHighWatermarkCheckPointThread()
+
+          maybeAddLogDirFetchers(partitionStates.keySet, highWatermarkCheckpoints)
+
+          replicaFetcherManager.shutdownIdleFetcherThreads()
+          replicaAlterLogDirsManager.shutdownIdleFetcherThreads()
+          onLeadershipChange(partitionsBecomeLeader, partitionsBecomeFollower)
+          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))
         }
+      }
+      val endMs = time.milliseconds()
+      val elapsedMs = endMs - startMs
+      stateChangeLogger.info(s"Finished LeaderAndIsr request in ${elapsedMs}ms correlationId $correlationId from controller " +
+        s"$controllerId for ${requestPartitionStates.size} partitions")
+      response
+    }
+  }
 
-        val newOnlineReplicas = newPartitions.flatMap(topicPartition => getReplica(topicPartition))
-        // Add future replica to partition's map
-        val futureReplicasAndInitialOffset = newOnlineReplicas.filter { replica =>
-          logManager.getLog(replica.topicPartition, isFuture = true).isDefined
-        }.map { replica =>
-          replica.topicPartition -> BrokerAndInitialOffset(BrokerEndPoint(config.brokerId, "localhost", -1), replica.highWatermark.messageOffset)
-        }.toMap
-        futureReplicasAndInitialOffset.keys.foreach(tp => getPartition(tp).get.getOrCreateReplica(Request.FutureLocalReplicaId))
-
-        // pause cleaning for partitions that are being moved and start ReplicaAlterDirThread to move replica from source dir to destination dir
-        futureReplicasAndInitialOffset.keys.foreach(logManager.abortAndPauseCleaning)
-        replicaAlterLogDirsManager.addFetcherForPartitions(futureReplicasAndInitialOffset)
-
-        replicaFetcherManager.shutdownIdleFetcherThreads()
-        replicaAlterLogDirsManager.shutdownIdleFetcherThreads()
-        onLeadershipChange(partitionsBecomeLeader, partitionsBecomeFollower)
-        new LeaderAndIsrResponse(Errors.NONE, responseMap.asJava)
+  private def maybeAddLogDirFetchers(partitions: Set[Partition],
+                                     offsetCheckpoints: OffsetCheckpoints): Unit = {
+    val futureReplicasAndInitialOffset = new mutable.HashMap[TopicPartition, InitialFetchState]
+    for (partition <- partitions) {
+      val topicPartition = partition.topicPartition
+      if (logManager.getLog(topicPartition, isFuture = true).isDefined) {
+        partition.log.foreach { log =>
+          val leader = BrokerEndPoint(config.brokerId, "localhost", -1)
+
+          // Add future replica log to partition's map
+          partition.createLogIfNotExists(
+            isNew = false,
+            isFutureReplica = true,
+            offsetCheckpoints)
+
+          // 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, log.highWatermark))
+        }
       }
     }
+
+    if (futureReplicasAndInitialOffset.nonEmpty)
+      replicaAlterLogDirsManager.addFetcherForPartitions(futureReplicasAndInitialOffset)
   }
 
   /*
@@ -1141,42 +1510,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 =>
-      stateChangeLogger.trace(s"Handling LeaderAndIsr request correlationId $correlationId from " +
-        s"controller $controllerId epoch $epoch starting the become-leader transition for " +
-        s"partition ${partition.topicPartition}")
-    }
-
-    for (partition <- partitionState.keys)
+                          responseMap: mutable.Map[TopicPartition, Errors],
+                          highWatermarkCheckpoints: OffsetCheckpoints): Set[Partition] = {
+    val traceEnabled = stateChangeLogger.isTraceEnabled
+    partitionStates.keys.foreach { partition =>
+      if (traceEnabled)
+        stateChangeLogger.trace(s"Handling LeaderAndIsr request correlationId $correlationId from " +
+          s"controller $controllerId epoch $controllerEpoch starting the become-leader transition for " +
+          s"partition ${partition.topicPartition}")
       responseMap.put(partition.topicPartition, Errors.NONE)
+    }
 
-    val partitionsToMakeLeaders: mutable.Set[Partition] = mutable.Set()
+    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))
+      stateChangeLogger.info(s"Stopped fetchers as part of LeaderAndIsr request correlationId $correlationId from " +
+        s"controller $controllerId epoch $controllerEpoch as part of the become-leader transition for " +
+        s"${partitionStates.size} partitions")
       // Update the partition information to be the leader
-      partitionState.foreach{ case (partition, partitionStateInfo) =>
+      partitionStates.forKeyValue { (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})")
-          } else
+          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)
@@ -1186,18 +1556,19 @@ 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 =>
-      stateChangeLogger.trace(s"Completed LeaderAndIsr request correlationId $correlationId from controller $controllerId " +
-        s"epoch $epoch for the become-leader transition for partition ${partition.topicPartition}")
-    }
+    if (traceEnabled)
+      partitionStates.keys.foreach { partition =>
+        stateChangeLogger.trace(s"Completed LeaderAndIsr request correlationId $correlationId from controller $controllerId " +
+          s"epoch $controllerEpoch for the become-leader transition for partition ${partition.topicPartition}")
+      }
 
     partitionsToMakeLeaders
   }
@@ -1221,233 +1592,260 @@ 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,
-                            partitionState: Map[Partition, LeaderAndIsrRequest.PartitionState],
+                            controllerEpoch: Int,
+                            partitionStates: Map[Partition, LeaderAndIsrPartitionState],
                             correlationId: Int,
-                            responseMap: mutable.Map[TopicPartition, Errors]) : Set[Partition] = {
-    partitionState.keys.foreach { partition =>
-      stateChangeLogger.trace(s"Handling LeaderAndIsr request correlationId $correlationId from controller $controllerId " +
-        s"epoch $epoch starting the become-follower transition for partition ${partition.topicPartition}")
-    }
-
-    for (partition <- partitionState.keys)
+                            responseMap: mutable.Map[TopicPartition, Errors],
+                            highWatermarkCheckpoints: OffsetCheckpoints) : Set[Partition] = {
+    val traceLoggingEnabled = stateChangeLogger.isTraceEnabled
+    partitionStates.forKeyValue { (partition, partitionState) =>
+      if (traceLoggingEnabled)
+        stateChangeLogger.trace(s"Handling LeaderAndIsr request correlationId $correlationId from controller $controllerId " +
+          s"epoch $controllerEpoch starting the become-follower transition for partition ${partition.topicPartition} with leader " +
+          s"${partitionState.leader}")
       responseMap.put(partition.topicPartition, Errors.NONE)
+    }
 
     val partitionsToMakeFollower: mutable.Set[Partition] = mutable.Set()
-
     try {
-
       // TODO: Delete leaders from LeaderAndIsrRequest
-      partitionState.foreach{ case (partition, partitionStateInfo) =>
+      partitionStates.forKeyValue { (partition, partitionState) =>
+        val newLeaderBrokerId = partitionState.leader
         try {
-          val newLeaderBrokerId = partitionStateInfo.basePartitionState.leader
           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(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}) since the " +
-              s"replica for the partition is offline due to disk error $e")
+              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 in dir $dirOpt", e)
+            error(s"Error while making broker the follower for partition $partition with leader " +
+              s"$newLeaderBrokerId in dir $dirOpt", e)
             responseMap.put(partition.topicPartition, Errors.KAFKA_STORAGE_ERROR)
         }
       }
 
       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}")
-      }
-
-      partitionsToMakeFollower.foreach { partition =>
-        val topicPartitionOperationKey = new TopicPartitionOperationKey(partition.topicPartition)
-        tryCompleteDelayedProduce(topicPartitionOperationKey)
-        tryCompleteDelayedFetch(topicPartitionOperationKey)
-      }
+      stateChangeLogger.info(s"Stopped fetchers as part of become-follower request from controller $controllerId " +
+        s"epoch $controllerEpoch with correlation id $correlationId for ${partitionsToMakeFollower.size} partitions")
 
       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")
+        completeDelayedFetchOrProduceRequests(partition.topicPartition)
       }
 
       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} since it is shutting down")
+        if (traceLoggingEnabled) {
+          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 $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 =>
-          partition.topicPartition -> BrokerAndInitialOffset(
-            metadataCache.getAliveBrokers.find(_.id == partition.leaderReplicaIdOpt.get).get.getBrokerEndPoint(config.interBrokerListenerName),
-            partition.getReplica().get.highWatermark.messageOffset)).toMap
-        replicaFetcherManager.addFetcherForPartitions(partitionsToMakeFollowerWithLeaderAndOffset)
+        val partitionsToMakeFollowerWithLeaderAndOffset = partitionsToMakeFollower.map { partition =>
+          val leader = metadataCache.getAliveBrokers.find(_.id == partition.leaderReplicaIdOpt.get).get
+            .brokerEndPoint(config.interBrokerListenerName)
+          val log = partition.localLogOrException
+          val fetchOffset = initialFetchOffset(log)
+          partition.topicPartition -> InitialFetchState(leader, partition.getLeaderEpoch, fetchOffset)
+        }.toMap
 
-        partitionsToMakeFollower.foreach { partition =>
-          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}")
-        }
+        replicaFetcherManager.addFetcherForPartitions(partitionsToMakeFollowerWithLeaderAndOffset)
       }
     } 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
     }
 
-    partitionState.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}")
-    }
+    if (traceLoggingEnabled)
+      partitionStates.keys.foreach { partition =>
+        stateChangeLogger.trace(s"Completed LeaderAndIsr request correlationId $correlationId from controller $controllerId " +
+          s"epoch $controllerEpoch for the become-follower transition for partition ${partition.topicPartition} with leader " +
+          s"${partitionStates(partition).leader}")
+      }
 
     partitionsToMakeFollower
   }
 
+  /**
+   * From IBP 2.7 onwards, we send latest fetch epoch in the request and truncate if a
+   * diverging epoch is returned in the response, avoiding the need for a separate
+   * OffsetForLeaderEpoch request.
+   */
+  private def initialFetchOffset(log: Log): Long = {
+    if (ApiVersion.isTruncationOnFetchSupported(config.interBrokerProtocolVersion) && log.latestEpoch.nonEmpty)
+      log.logEndOffset
+    else
+      log.highWatermark
+  }
+
   private def maybeShrinkIsr(): Unit = {
     trace("Evaluating ISR list of partitions to see which replicas can be removed from the ISR")
-    nonOfflinePartitionsIterator.foreach(_.maybeShrinkIsr(config.replicaLagTimeMaxMs))
+
+    // Shrink ISRs for non offline partitions
+    allPartitions.keys.foreach { topicPartition =>
+      nonOfflinePartition(topicPartition).foreach(_.maybeShrinkIsr())
+    }
   }
 
   /**
-   * Update the follower's fetch state in the leader based on the last fetch request and update `readResult`,
-   * if necessary.
+   * 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) =>
-              if (partition.updateReplicaLogReadResult(replica, readResult))
-                partition.leaderReplicaIfLocal.foreach { leaderReplica =>
-                  updatedReadResult = readResult.updateLeaderReplicaInfo(leaderReplica)
-                }
-            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)) {
+              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.assignmentState.replicas.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.messageOffset))
+    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.getReplica(localBrokerId).foreach(replicasList.add)
-      partition.getReplica(Request.FutureLocalReplicaId).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
-      try {
-        highWatermarkCheckpoints.get(dir).foreach(_.write(hwms))
-      } catch {
+  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(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
-  def handleLogDirFailure(dir: String) {
+  /**
+   * The log directory failure handler for the replica
+   *
+   * @param dir                     the absolute path of the log directory
+   * @param sendZkNotification      check if we need to send notification to zookeeper node (needed for unit test)
+   */
+  def handleLogDirFailure(dir: String, sendZkNotification: Boolean = true): Unit = {
     if (!logManager.isLogDirOnline(dir))
       return
-    info(s"Stopping serving replicas in dir $dir")
+    warn(s"Stopping serving replicas in dir $dir")
     replicaStateChangeLock synchronized {
       val newOfflinePartitions = nonOfflinePartitionsIterator.filter { partition =>
-        partition.getReplica(config.brokerId).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.getReplica(Request.FutureLocalReplicaId).exists { replica =>
-          replica.log.isDefined && replica.log.get.dir.getParent == dir
-        }
+        partition.futureLog.exists { _.parentDir == dir }
       }.toSet
 
       replicaFetcherManager.removeFetcherForPartitions(newOfflinePartitions)
       replicaAlterLogDirsManager.removeFetcherForPartitions(newOfflinePartitions ++ partitionsWithOfflineFutureReplica.map(_.topicPartition))
 
-      partitionsWithOfflineFutureReplica.foreach(partition => partition.removeFutureLocalReplica())
+      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 " +
+      warn(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.")
     }
     logManager.handleLogDirFailure(dir)
-    zkClient.propagateLogDirEvent(localBrokerId)
-    info(s"Stopped serving replicas in dir $dir")
+
+    if (sendZkNotification)
+      zkClient.propagateLogDirEvent(localBrokerId)
+    warn(s"Stopped serving replicas in dir $dir")
   }
 
-  def removeMetrics() {
+  def removeMetrics(): Unit = {
     removeMetric("LeaderCount")
     removeMetric("PartitionCount")
     removeMetric("OfflineReplicaCount")
     removeMetric("UnderReplicatedPartitions")
     removeMetric("UnderMinIsrPartitionCount")
+    removeMetric("AtMinIsrPartitionCount")
   }
 
   // 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)
@@ -1457,8 +1855,10 @@ class ReplicaManager(val config: KafkaConfig,
     delayedFetchPurgatory.shutdown()
     delayedProducePurgatory.shutdown()
     delayedDeleteRecordsPurgatory.shutdown()
+    delayedElectLeaderPurgatory.shutdown()
     if (checkpointHW)
       checkpointHighWatermarks()
+    replicaSelectorOpt.foreach(_.close)
     info("Shut down completely")
   }
 
@@ -1470,19 +1870,95 @@ class ReplicaManager(val config: KafkaConfig,
     new ReplicaAlterLogDirsManager(config, this, quotaManager, brokerTopicStats)
   }
 
-  def lastOffsetForLeaderEpoch(requestedEpochInfo: Map[TopicPartition, Integer]): Map[TopicPartition, EpochEndOffset] = {
-    requestedEpochInfo.map { case (tp, leaderEpoch) =>
-      val epochEndOffset = getPartition(tp) match {
-        case Some(partition) =>
-          if (partition eq ReplicaManager.OfflinePartition)
-            new EpochEndOffset(KAFKA_STORAGE_ERROR, UNDEFINED_EPOCH_OFFSET)
-          else
-            partition.lastOffsetForLeaderEpoch(leaderEpoch)
-        case None =>
-          new EpochEndOffset(UNKNOWN_TOPIC_OR_PARTITION, UNDEFINED_EPOCH_OFFSET)
+  protected def createReplicaSelector(): Option[ReplicaSelector] = {
+    config.replicaSelectorClassName.map { className =>
+      val tmpReplicaSelector: ReplicaSelector = CoreUtils.createObject[ReplicaSelector](className)
+      tmpReplicaSelector.configure(config.originals())
+      tmpReplicaSelector
+    }
+  }
+
+  def lastOffsetForLeaderEpoch(
+    requestedEpochInfo: Seq[OffsetForLeaderTopic]
+  ): Seq[OffsetForLeaderTopicResult] = {
+    requestedEpochInfo.map { offsetForLeaderTopic =>
+      val partitions = offsetForLeaderTopic.partitions.asScala.map { offsetForLeaderPartition =>
+        val tp = new TopicPartition(offsetForLeaderTopic.topic, offsetForLeaderPartition.partition)
+        getPartition(tp) match {
+          case HostedPartition.Online(partition) =>
+            val currentLeaderEpochOpt =
+              if (offsetForLeaderPartition.currentLeaderEpoch == RecordBatch.NO_PARTITION_LEADER_EPOCH)
+                Optional.empty[Integer]
+              else
+                Optional.of[Integer](offsetForLeaderPartition.currentLeaderEpoch)
+
+            partition.lastOffsetForLeaderEpoch(
+              currentLeaderEpochOpt,
+              offsetForLeaderPartition.leaderEpoch,
+              fetchOnlyFromLeader = true)
+
+          case HostedPartition.Offline =>
+            new EpochEndOffset()
+              .setPartition(offsetForLeaderPartition.partition)
+              .setErrorCode(Errors.KAFKA_STORAGE_ERROR.code)
+
+          case HostedPartition.None if metadataCache.contains(tp) =>
+            new EpochEndOffset()
+              .setPartition(offsetForLeaderPartition.partition)
+              .setErrorCode(Errors.NOT_LEADER_OR_FOLLOWER.code)
+
+          case HostedPartition.None =>
+            new EpochEndOffset()
+              .setPartition(offsetForLeaderPartition.partition)
+              .setErrorCode(Errors.UNKNOWN_TOPIC_OR_PARTITION.code)
+        }
       }
-      tp -> epochEndOffset
+
+      new OffsetForLeaderTopicResult()
+        .setTopic(offsetForLeaderTopic.topic)
+        .setPartitions(partitions.toList.asJava)
     }
   }
-}
 
+  def electLeaders(
+    controller: KafkaController,
+    partitions: Set[TopicPartition],
+    electionType: ElectionType,
+    responseCallback: Map[TopicPartition, ApiError] => Unit,
+    requestTimeout: Int
+  ): Unit = {
+
+    val deadline = time.milliseconds() + requestTimeout
+
+    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.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(failures)
+      }
+    }
+
+    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 84004e3f8e66a..3035cb1371858 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
@@ -51,8 +53,9 @@ object ReplicationQuotaManagerConfig {
 }
 
 trait ReplicaQuota {
+  def record(value: Long): Unit
   def isThrottled(topicPartition: TopicPartition): Boolean
-  def isQuotaExceeded(): Boolean
+  def isQuotaExceeded: Boolean
 }
 
 object Constants {
@@ -83,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.
@@ -99,12 +102,13 @@ class ReplicationQuotaManager(val config: ReplicationQuotaManagerConfig,
     *
     * @return
     */
-  override def isQuotaExceeded(): Boolean = {
+  override def isQuotaExceeded: Boolean = {
     try {
       sensor().checkQuotas()
     } catch {
       case qve: QuotaViolationException =>
-        trace("%s: Quota violated for sensor (%s), metric: (%s), metric-value: (%f), bound: (%f)".format(replicationType, sensor().name(), qve.metricName, qve.value, qve.bound))
+        trace(s"$replicationType: Quota violated for sensor (${sensor().name}), metric: (${qve.metric.metricName}), " +
+          s"metric-value: (${qve.value}), bound: (${qve.bound})")
         return true
     }
     false
@@ -129,13 +133,8 @@ class ReplicationQuotaManager(val config: ReplicationQuotaManagerConfig,
     *
     * @param value
     */
-  def record(value: Long) {
-    try {
-      sensor().record(value)
-    } catch {
-      case qve: QuotaViolationException =>
-        trace(s"Record: Quota violated, but ignored, for sensor (${sensor.name}), metric: (${qve.metricName}), value : (${qve.value}), bound: (${qve.bound}), recordedValue ($value)")
-    }
+  def record(value: Long): Unit = {
+    sensor().record(value.toDouble, time.milliseconds(), false)
   }
 
   /**
@@ -146,7 +145,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)
   }
 
@@ -156,7 +155,7 @@ class ReplicationQuotaManager(val config: ReplicationQuotaManagerConfig,
     * @param topic
     * @return
     */
-  def markThrottled(topic: String) {
+  def markThrottled(topic: String): Unit = {
     markThrottled(topic, AllReplicas)
   }
 
@@ -166,7 +165,7 @@ class ReplicationQuotaManager(val config: ReplicationQuotaManagerConfig,
     * @param topic
     * @return
     */
-  def removeThrottle(topic: String) {
+  def removeThrottle(topic: String): Unit = {
     throttledPartitions.remove(topic)
   }
 
@@ -175,10 +174,10 @@ class ReplicationQuotaManager(val config: ReplicationQuotaManagerConfig,
     *
     * @return
     */
-  def upperBound(): Long = {
+  def upperBound: Long = {
     inReadLock(lock) {
       if (quota != null)
-        quota.bound().toLong
+        quota.bound.toLong
       else
         Long.MaxValue
     }
@@ -195,9 +194,7 @@ class ReplicationQuotaManager(val config: ReplicationQuotaManagerConfig,
     sensorAccess.getOrCreate(
       replicationType.toString,
       InactiveSensorExpirationTimeSeconds,
-      rateMetricName,
-      Some(getQuotaMetricConfig(quota)),
-      new SimpleRate
+      sensor => sensor.add(rateMetricName, new SimpleRate, getQuotaMetricConfig(quota))
     )
   }
 }
diff --git a/core/src/main/scala/kafka/server/SensorAccess.scala b/core/src/main/scala/kafka/server/SensorAccess.scala
index 49e1fc00b8fa2..3a063f2aba714 100644
--- a/core/src/main/scala/kafka/server/SensorAccess.scala
+++ b/core/src/main/scala/kafka/server/SensorAccess.scala
@@ -18,8 +18,7 @@ package kafka.server
 
 import java.util.concurrent.locks.ReadWriteLock
 
-import org.apache.kafka.common.MetricName
-import org.apache.kafka.common.metrics.{MeasurableStat, MetricConfig, Metrics, Sensor}
+import org.apache.kafka.common.metrics.{Metrics, Sensor}
 
 /**
   * Class which centralises the logic for creating/accessing sensors.
@@ -29,8 +28,7 @@ import org.apache.kafka.common.metrics.{MeasurableStat, MetricConfig, Metrics, S
   */
 class SensorAccess(lock: ReadWriteLock, metrics: Metrics) {
 
-  def getOrCreate(sensorName: String, expirationTime: Long,
-                  metricName: => MetricName, config: => Option[MetricConfig], measure: => MeasurableStat): Sensor = {
+  def getOrCreate(sensorName: String, expirationTime: Long, registerMetrics: Sensor => Unit): Sensor = {
     var sensor: Sensor = null
 
     /* Acquire the read lock to fetch the sensor. It is safe to call getSensor from multiple threads.
@@ -61,8 +59,8 @@ class SensorAccess(lock: ReadWriteLock, metrics: Metrics) {
         // ensure that we initialise `ClientSensors` with non-null parameters.
         sensor = metrics.getSensor(sensorName)
         if (sensor == null) {
-          sensor = metrics.sensor(sensorName, config.orNull, expirationTime)
-          sensor.add(metricName, measure)
+          sensor = metrics.sensor(sensorName, null, expirationTime)
+          registerMetrics(sensor)
         }
       } finally {
         lock.writeLock().unlock()
diff --git a/core/src/main/scala/kafka/server/ThrottledChannel.scala b/core/src/main/scala/kafka/server/ThrottledChannel.scala
new file mode 100644
index 0000000000000..531ef5da43199
--- /dev/null
+++ b/core/src/main/scala/kafka/server/ThrottledChannel.scala
@@ -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 kafka.server
+
+import java.util.concurrent.{Delayed, TimeUnit}
+
+import kafka.network
+import kafka.network.RequestChannel
+import kafka.network.RequestChannel.Response
+import kafka.utils.Logging
+import org.apache.kafka.common.utils.Time
+
+
+/**
+  * Represents a request whose response has been delayed.
+  * @param request The request that has been delayed
+  * @param time Time instance to use
+  * @param throttleTimeMs Delay associated with this request
+  * @param channelThrottlingCallback Callback for channel throttling
+  */
+class ThrottledChannel(val request: RequestChannel.Request, val time: Time, val throttleTimeMs: Int,
+                       channelThrottlingCallback: Response => Unit)
+  extends Delayed with Logging {
+
+  private val endTimeNanos = time.nanoseconds() + TimeUnit.MILLISECONDS.toNanos(throttleTimeMs)
+
+  // Notify the socket server that throttling has started for this channel.
+  channelThrottlingCallback(new RequestChannel.StartThrottlingResponse(request))
+
+  // Notify the socket server that throttling has been done for this channel.
+  def notifyThrottlingDone(): Unit = {
+    trace(s"Channel throttled for: $throttleTimeMs ms")
+    channelThrottlingCallback(new network.RequestChannel.EndThrottlingResponse(request))
+  }
+
+  override def getDelay(unit: TimeUnit): Long = {
+    unit.convert(endTimeNanos - time.nanoseconds(), TimeUnit.NANOSECONDS)
+  }
+
+  override def compareTo(d: Delayed): Int = {
+    val other = d.asInstanceOf[ThrottledChannel]
+    java.lang.Long.compare(this.endTimeNanos, other.endTimeNanos)
+  }
+}
diff --git a/core/src/main/scala/kafka/server/ThrottledResponse.scala b/core/src/main/scala/kafka/server/ThrottledResponse.scala
deleted file mode 100644
index 214fa1f5c837a..0000000000000
--- a/core/src/main/scala/kafka/server/ThrottledResponse.scala
+++ /dev/null
@@ -1,46 +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.util.concurrent.{TimeUnit, Delayed}
-
-import org.apache.kafka.common.utils.Time
-
-
-/**
- * Represents a request whose response has been delayed.
- * @param time @Time instance to use
- * @param throttleTimeMs delay associated with this request
- * @param callback Callback to trigger after delayTimeMs milliseconds
- */
-private[server] class ThrottledResponse(val time: Time, val throttleTimeMs: Int, callback: Int => Unit) extends Delayed {
-  val endTime = time.milliseconds + throttleTimeMs
-
-  def execute() = callback(throttleTimeMs)
-
-  override def getDelay(unit: TimeUnit): Long = {
-    unit.convert(endTime - time.milliseconds, TimeUnit.MILLISECONDS)
-  }
-
-  override def compareTo(d: Delayed): Int = {
-    val other = d.asInstanceOf[ThrottledResponse]
-    if (this.endTime < other.endTime) -1
-    else if (this.endTime > other.endTime) 1
-    else 0
-  }
-}
diff --git a/core/src/main/scala/kafka/server/checkpoints/CheckpointFile.scala b/core/src/main/scala/kafka/server/checkpoints/CheckpointFile.scala
index 4c1011f9429d2..2ab9ab9b39318 100644
--- a/core/src/main/scala/kafka/server/checkpoints/CheckpointFile.scala
+++ b/core/src/main/scala/kafka/server/checkpoints/CheckpointFile.scala
@@ -33,6 +33,48 @@ trait CheckpointFileFormatter[T]{
   def fromLine(line: String): Option[T]
 }
 
+class CheckpointReadBuffer[T](location: String,
+                              reader: BufferedReader,
+                              version: Int,
+                              formatter: CheckpointFileFormatter[T]) extends Logging {
+  def read(): Seq[T] = {
+    def malformedLineException(line: String) =
+      new IOException(s"Malformed line in checkpoint file ($location): '$line'")
+
+    var line: String = null
+    try {
+      line = reader.readLine()
+      if (line == null)
+        return Seq.empty
+      line.toInt match {
+        case fileVersion if fileVersion == version =>
+          line = reader.readLine()
+          if (line == null)
+            return Seq.empty
+          val expectedSize = line.toInt
+          val entries = mutable.Buffer[T]()
+          line = reader.readLine()
+          while (line != null) {
+            val entry = formatter.fromLine(line)
+            entry match {
+              case Some(e) =>
+                entries += e
+                line = reader.readLine()
+              case _ => throw malformedLineException(line)
+            }
+          }
+          if (entries.size != expectedSize)
+            throw new IOException(s"Expected $expectedSize entries in checkpoint file ($location), but found only ${entries.size}")
+          entries
+        case _ =>
+          throw new IOException(s"Unrecognized version of the checkpoint file ($location): " + version)
+      }
+    } catch {
+      case _: NumberFormatException => throw malformedLineException(line)
+    }
+  }
+}
+
 class CheckpointFile[T](val file: File,
                         version: Int,
                         formatter: CheckpointFileFormatter[T],
@@ -45,7 +87,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
@@ -80,41 +122,12 @@ class CheckpointFile[T](val file: File,
   }
 
   def read(): Seq[T] = {
-    def malformedLineException(line: String) =
-      new IOException(s"Malformed line in checkpoint file (${file.getAbsolutePath}): $line'")
     lock synchronized {
       try {
-        val reader = new BufferedReader(new InputStreamReader(new FileInputStream(file), StandardCharsets.UTF_8))
-        var line: String = null
+        val reader = Files.newBufferedReader(path)
         try {
-          line = reader.readLine()
-          if (line == null)
-            return Seq.empty
-          line.toInt match {
-            case fileVersion if fileVersion == version =>
-              line = reader.readLine()
-              if (line == null)
-                return Seq.empty
-              val expectedSize = line.toInt
-              val entries = mutable.Buffer[T]()
-              line = reader.readLine()
-              while (line != null) {
-                val entry = formatter.fromLine(line)
-                entry match {
-                  case Some(e) =>
-                    entries += e
-                    line = reader.readLine()
-                  case _ => throw malformedLineException(line)
-                }
-              }
-              if (entries.size != expectedSize)
-                throw new IOException(s"Expected $expectedSize entries in checkpoint file (${file.getAbsolutePath}), but found only ${entries.size}")
-              entries
-            case _ =>
-              throw new IOException(s"Unrecognized version of the checkpoint file (${file.getAbsolutePath}): " + version)
-          }
-        } catch {
-          case _: NumberFormatException => throw malformedLineException(line)
+          val checkpointBuffer = new CheckpointReadBuffer[T](file.getAbsolutePath, reader, version, formatter)
+          checkpointBuffer.read()
         } finally {
           reader.close()
         }
diff --git a/core/src/main/scala/kafka/server/checkpoints/LeaderEpochCheckpointFile.scala b/core/src/main/scala/kafka/server/checkpoints/LeaderEpochCheckpointFile.scala
index a8db688ac135d..c9a0f62b297e8 100644
--- a/core/src/main/scala/kafka/server/checkpoints/LeaderEpochCheckpointFile.scala
+++ b/core/src/main/scala/kafka/server/checkpoints/LeaderEpochCheckpointFile.scala
@@ -25,19 +25,17 @@ import kafka.server.epoch.EpochEntry
 import scala.collection._
 
 trait LeaderEpochCheckpoint {
-  def write(epochs: Seq[EpochEntry])
+  def write(epochs: Iterable[EpochEntry]): Unit
   def read(): Seq[EpochEntry]
 }
 
-object LeaderEpochFile {
-  private val LeaderEpochCheckpointFilename = "leader-epoch-checkpoint"
-  def newFile(dir: File): File = new File(dir, LeaderEpochCheckpointFilename)
-}
-
 object LeaderEpochCheckpointFile {
+  private val LeaderEpochCheckpointFilename = "leader-epoch-checkpoint"
   private val WhiteSpacesPattern = Pattern.compile("\\s+")
   private val CurrentVersion = 0
 
+  def newFile(dir: File): File = new File(dir, LeaderEpochCheckpointFilename)
+
   object Formatter extends CheckpointFileFormatter[EpochEntry] {
 
     override def toLine(entry: EpochEntry): String = s"${entry.epoch} ${entry.startOffset}"
@@ -54,14 +52,22 @@ object LeaderEpochCheckpointFile {
 }
 
 /**
-  * This class persists a map of (LeaderEpoch => Offsets) to a file (for a certain replica)
-  */
+ * This class persists a map of (LeaderEpoch => Offsets) to a file (for a certain replica)
+ *
+ * The format in the LeaderEpoch checkpoint file is like this:
+ * -----checkpoint file begin------
+ * 0                <- LeaderEpochCheckpointFile.currentVersion
+ * 2                <- following entries size
+ * 0  1     <- the format is: leader_epoch(int32) start_offset(int64)
+ * 1  2
+ * -----checkpoint file end----------
+ */
 class LeaderEpochCheckpointFile(val file: File, logDirFailureChannel: LogDirFailureChannel = null) extends LeaderEpochCheckpoint {
   import LeaderEpochCheckpointFile._
 
   val checkpoint = new CheckpointFile[EpochEntry](file, CurrentVersion, Formatter, logDirFailureChannel, file.getParentFile.getParent)
 
-  def write(epochs: Seq[EpochEntry]): Unit = checkpoint.write(epochs)
+  def write(epochs: Iterable[EpochEntry]): Unit = checkpoint.write(epochs)
 
   def read(): Seq[EpochEntry] = checkpoint.read()
 }
diff --git a/core/src/main/scala/kafka/server/checkpoints/OffsetCheckpointFile.scala b/core/src/main/scala/kafka/server/checkpoints/OffsetCheckpointFile.scala
index 2769cb4240a52..722f60900f55a 100644
--- a/core/src/main/scala/kafka/server/checkpoints/OffsetCheckpointFile.scala
+++ b/core/src/main/scala/kafka/server/checkpoints/OffsetCheckpointFile.scala
@@ -45,19 +45,55 @@ object OffsetCheckpointFile {
 }
 
 trait OffsetCheckpoint {
-  def write(epochs: Seq[EpochEntry])
+  def write(epochs: Seq[EpochEntry]): Unit
   def read(): Seq[EpochEntry]
 }
 
 /**
-  * This class persists a map of (Partition => Offsets) to a file (for a certain replica)
-  */
+ * This class persists a map of (Partition => Offsets) to a file (for a certain replica)
+ *
+ * The format in the offset checkpoint file is like this:
+ *  -----checkpoint file begin------
+ *  0                <- OffsetCheckpointFile.currentVersion
+ *  2                <- following entries size
+ *  tp1  par1  1     <- the format is: TOPIC  PARTITION  OFFSET
+ *  tp1  par2  2
+ *  -----checkpoint file end----------
+ */
 class OffsetCheckpointFile(val file: File, logDirFailureChannel: LogDirFailureChannel = null) {
   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 29464663ea77a..03775439f00a6 100644
--- a/core/src/main/scala/kafka/server/epoch/LeaderEpochFileCache.scala
+++ b/core/src/main/scala/kafka/server/epoch/LeaderEpochFileCache.scala
@@ -16,113 +16,229 @@
   */
 package kafka.server.epoch
 
+import java.util
 import java.util.concurrent.locks.ReentrantReadWriteLock
 
-import kafka.server.LogOffsetMetadata
 import kafka.server.checkpoints.LeaderEpochCheckpoint
-import org.apache.kafka.common.requests.EpochEndOffset.{UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET}
 import kafka.utils.CoreUtils._
 import kafka.utils.Logging
 import org.apache.kafka.common.TopicPartition
-import scala.collection.mutable.ListBuffer
-
-trait LeaderEpochCache {
-  def assign(leaderEpoch: Int, offset: Long)
-  def latestEpoch(): Int
-  def endOffsetFor(epoch: Int): Long
-  def clearAndFlushLatest(offset: Long)
-  def clearAndFlushEarliest(offset: Long)
-  def clearAndFlush()
-  def clear()
-}
+import org.apache.kafka.common.requests.OffsetsForLeaderEpochResponse.{UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET}
+
+import scala.collection.{Seq, mutable}
+import scala.jdk.CollectionConverters._
 
 /**
-  * Represents a cache of (LeaderEpoch => Offset) mappings for a particular replica.
-  *
-  * Leader Epoch = epoch assigned to each leader by the controller.
-  * Offset = offset of the first message in each epoch.
-  *
-  * @param leo a function that determines the log end offset
-  * @param checkpoint the checkpoint file
-  */
-class LeaderEpochFileCache(topicPartition: TopicPartition, leo: () => LogOffsetMetadata, checkpoint: LeaderEpochCheckpoint) extends LeaderEpochCache with Logging {
+ * Represents a cache of (LeaderEpoch => Offset) mappings for a particular replica.
+ *
+ * Leader Epoch = epoch assigned to each leader by the controller.
+ * Offset = offset of the first message in each epoch.
+ *
+ * @param topicPartition the associated topic partition
+ * @param checkpoint the checkpoint file
+ * @param logEndOffset function to fetch the current log end offset
+ */
+class LeaderEpochFileCache(topicPartition: TopicPartition,
+                           logEndOffset: () => Long,
+                           checkpoint: LeaderEpochCheckpoint) extends Logging {
+  this.logIdent = s"[LeaderEpochCache $topicPartition] "
+
   private val lock = new ReentrantReadWriteLock()
-  private var epochs: ListBuffer[EpochEntry] = inWriteLock(lock) { ListBuffer(checkpoint.read(): _*) }
+  private val epochs = new util.TreeMap[Int, EpochEntry]()
+
+  inWriteLock(lock) {
+    checkpoint.read().foreach(assign)
+  }
 
   /**
     * Assigns the supplied Leader Epoch to the supplied Offset
     * Once the epoch is assigned it cannot be reassigned
-    *
-    * @param epoch
-    * @param offset
     */
-  override def assign(epoch: Int, offset: Long): Unit = {
+  def assign(epoch: Int, startOffset: Long): Unit = {
+    val entry = EpochEntry(epoch, startOffset)
+    if (assign(entry)) {
+      debug(s"Appended new epoch entry $entry. Cache now contains ${epochs.size} entries.")
+      flush()
+    }
+  }
+
+  private def assign(entry: EpochEntry): Boolean = {
+    if (entry.epoch < 0 || entry.startOffset < 0) {
+      throw new IllegalArgumentException(s"Received invalid partition leader epoch entry $entry")
+    }
+
+    def isUpdateNeeded: Boolean = {
+      latestEntry match {
+        case Some(lastEntry) =>
+          entry.epoch != lastEntry.epoch || entry.startOffset < lastEntry.startOffset
+        case None =>
+          true
+      }
+    }
+
+    // Check whether the append is needed before acquiring the write lock
+    // in order to avoid contention with readers in the common case
+    if (!isUpdateNeeded)
+      return false
+
     inWriteLock(lock) {
-      if (epoch >= 0 && epoch > latestEpoch && offset >= latestOffset) {
-        info(s"Updated PartitionLeaderEpoch. ${epochChangeMsg(epoch, offset)}. Cache now contains ${epochs.size} entries.")
-        epochs += EpochEntry(epoch, offset)
-        flush()
+      if (isUpdateNeeded) {
+        maybeTruncateNonMonotonicEntries(entry)
+        epochs.put(entry.epoch, entry)
+        true
       } else {
-        validateAndMaybeWarn(epoch, offset)
+        false
       }
     }
   }
 
   /**
-    * Returns the current Leader Epoch. This is the latest epoch
-    * which has messages assigned to it.
-    *
-    * @return
-    */
-  override def latestEpoch(): Int = {
+   * Remove any entries which violate monotonicity prior to appending a new entry
+   */
+  private def maybeTruncateNonMonotonicEntries(newEntry: EpochEntry): Unit = {
+    val removedEpochs = removeFromEnd { entry =>
+      entry.epoch >= newEntry.epoch || entry.startOffset >= newEntry.startOffset
+    }
+
+    if (removedEpochs.size > 1
+      || (removedEpochs.nonEmpty && removedEpochs.head.startOffset != newEntry.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 $newEntry caused truncation of conflicting entries $removedEpochs. " +
+        s"Cache now contains ${epochs.size} entries.")
+    }
+  }
+
+  private def removeFromEnd(predicate: EpochEntry => Boolean): Seq[EpochEntry] = {
+    removeWhileMatching(epochs.descendingMap.entrySet().iterator(), predicate)
+  }
+
+  private def removeFromStart(predicate: EpochEntry => Boolean): Seq[EpochEntry] = {
+    removeWhileMatching(epochs.entrySet().iterator(), predicate)
+  }
+
+  private def removeWhileMatching(
+    iterator: util.Iterator[util.Map.Entry[Int, EpochEntry]],
+    predicate: EpochEntry => Boolean
+  ): Seq[EpochEntry] = {
+    val removedEpochs = mutable.ListBuffer.empty[EpochEntry]
+
+    while (iterator.hasNext) {
+      val entry = iterator.next().getValue
+      if (predicate.apply(entry)) {
+        removedEpochs += entry
+        iterator.remove()
+      } else {
+        return removedEpochs
+      }
+    }
+
+    removedEpochs
+  }
+
+  def nonEmpty: Boolean = inReadLock(lock) {
+    !epochs.isEmpty
+  }
+
+  def latestEntry: Option[EpochEntry] = {
+    inReadLock(lock) {
+      Option(epochs.lastEntry).map(_.getValue)
+    }
+  }
+
+  /**
+   * Returns the current Leader Epoch if one exists. This is the latest epoch
+   * which has messages assigned to it.
+   */
+  def latestEpoch: Option[Int] = {
+    latestEntry.map(_.epoch)
+  }
+
+  def previousEpoch: Option[Int] = {
     inReadLock(lock) {
-      if (epochs.isEmpty) UNDEFINED_EPOCH else epochs.last.epoch
+      latestEntry.flatMap(entry => Option(epochs.lowerEntry(entry.epoch))).map(_.getKey)
     }
   }
 
   /**
-    * Returns the End Offset for a requested Leader Epoch.
+   * Get the earliest cached entry if one exists.
+   */
+  def earliestEntry: Option[EpochEntry] = {
+    inReadLock(lock) {
+      Option(epochs.firstEntry).map(_.getValue)
+    }
+  }
+
+  /**
+    * Returns the Leader Epoch and the End Offset for a requested Leader Epoch.
     *
-    * This is defined as the start offset of the first Leader Epoch larger than the
-    * Leader Epoch requested, or else the Log End Offset if the latest epoch was requested.
+    * The Leader Epoch returned is the largest epoch less than or equal to the requested Leader
+    * Epoch. The End Offset is the end offset of this epoch, which is defined as the start offset
+    * of the first Leader Epoch larger than the Leader Epoch requested, or else the Log End
+    * 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
-    * @return offset
+    * @param requestedEpoch requested leader epoch
+    * @return found leader epoch and end offset
     */
-  override def endOffsetFor(requestedEpoch: Int): Long = {
+  def endOffsetFor(requestedEpoch: Int): (Int, Long) = {
     inReadLock(lock) {
-      val offset =
-        if (requestedEpoch == latestEpoch) {
-          leo().messageOffset
-        }
-        else {
-          val subsequentEpochs = epochs.filter(e => e.epoch > requestedEpoch)
-          if (subsequentEpochs.isEmpty || requestedEpoch < epochs.head.epoch)
-            UNDEFINED_EPOCH_OFFSET
-          else
-            subsequentEpochs.head.startOffset
+      val epochAndOffset =
+        if (requestedEpoch == UNDEFINED_EPOCH) {
+          // This may happen if a bootstrapping follower sends a request with undefined epoch or
+          // a follower is on the older message format where leader epochs are not recorded
+          (UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET)
+        } else if (latestEpoch.contains(requestedEpoch)) {
+          // For the leader, the latest epoch is always the current leader epoch that is still being written to.
+          // Followers should not have any reason to query for the end offset of the current epoch, but a consumer
+          // might if it is verifying its committed offset following a group rebalance. In this case, we return
+          // the current log end offset which makes the truncation check work as expected.
+          (requestedEpoch, logEndOffset())
+        } else {
+          val higherEntry = epochs.higherEntry(requestedEpoch)
+          if (higherEntry == null) {
+            // The requested epoch is larger than any known epoch. This case should never be hit because
+            // the latest cached epoch is always the largest.
+            (UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET)
+          } else {
+            val floorEntry = epochs.floorEntry(requestedEpoch)
+            if (floorEntry == null) {
+              // The requested epoch is smaller than any known epoch, so we return the start offset of the first
+              // known epoch which is larger than it. This may be inaccurate as there could have been
+              // epochs in between, but the point is that the data has already been removed from the log
+              // and we want to ensure that the follower can replicate correctly beginning from the leader's
+              // start offset.
+              (requestedEpoch, higherEntry.getValue.startOffset)
+            } else {
+              // We have at least one previous epoch and one subsequent epoch. The result is the first
+              // prior epoch and the starting offset of the first subsequent epoch.
+              (floorEntry.getValue.epoch, higherEntry.getValue.startOffset)
+            }
+          }
         }
-      debug(s"Processed offset for epoch request for partition ${topicPartition} epoch:$requestedEpoch and returning offset $offset from epoch list of size ${epochs.size}")
-      offset
+      debug(s"Processed end offset request for epoch $requestedEpoch and returning epoch ${epochAndOffset._1} " +
+        s"with end offset ${epochAndOffset._2} from epoch cache of size ${epochs.size}")
+      epochAndOffset
     }
   }
 
   /**
     * Removes all epoch entries from the store with start offsets greater than or equal to the passed offset.
-    *
-    * @param offset
     */
-  override def clearAndFlushLatest(offset: Long): Unit = {
+  def truncateFromEnd(endOffset: Long): Unit = {
     inWriteLock(lock) {
-      val before = epochs
-      if (offset >= 0 && offset <= latestOffset()) {
-        epochs = epochs.filter(entry => entry.startOffset < offset)
+      if (endOffset >= 0 && latestEntry.exists(_.startOffset >= endOffset)) {
+        val removedEntries = removeFromEnd(_.startOffset >= endOffset)
+
         flush()
-        info(s"Cleared latest ${before.toSet.filterNot(epochs.toSet)} entries from epoch cache based on passed offset $offset leaving ${epochs.size} in EpochFile for partition $topicPartition")
+
+        debug(s"Cleared entries $removedEntries from epoch cache after " +
+          s"truncating to end offset $endOffset, leaving ${epochs.size} entries in the cache.")
       }
     }
   }
@@ -131,23 +247,24 @@ class LeaderEpochFileCache(topicPartition: TopicPartition, leo: () => LogOffsetM
     * Clears old epoch entries. This method searches for the oldest epoch < offset, updates the saved epoch offset to
     * be offset, then clears any previous epoch entries.
     *
-    * This method is exclusive: so clearEarliest(6) will retain an entry at offset 6.
+    * This method is exclusive: so truncateFromStart(6) will retain an entry at offset 6.
     *
-    * @param offset the offset to clear up to
+    * @param startOffset the offset to clear up to
     */
-  override def clearAndFlushEarliest(offset: Long): Unit = {
+  def truncateFromStart(startOffset: Long): Unit = {
     inWriteLock(lock) {
-      val before = epochs
-      if (offset >= 0 && earliestOffset() < offset) {
-        val earliest = epochs.filter(entry => entry.startOffset < offset)
-        if (earliest.nonEmpty) {
-          epochs = epochs --= earliest
-          //If the offset is less than the earliest offset remaining, add previous epoch back, but with an updated offset
-          if (offset < earliestOffset() || epochs.isEmpty)
-            new EpochEntry(earliest.last.epoch, offset) +=: epochs
-          flush()
-          info(s"Cleared earliest ${before.toSet.filterNot(epochs.toSet).size} entries from epoch cache based on passed offset $offset leaving ${epochs.size} in EpochFile for partition $topicPartition")
-        }
+      val removedEntries = removeFromStart { entry =>
+        entry.startOffset <= startOffset
+      }
+
+      removedEntries.lastOption.foreach { firstBeforeStartOffset =>
+        val updatedFirstEntry = EpochEntry(firstBeforeStartOffset.epoch, startOffset)
+        epochs.put(updatedFirstEntry.epoch, updatedFirstEntry)
+
+        flush()
+
+        debug(s"Cleared entries $removedEntries and rewrote first entry $updatedFirstEntry after " +
+          s"truncating to start offset $startOffset, leaving ${epochs.size} in the cache.")
       }
     }
   }
@@ -155,47 +272,31 @@ class LeaderEpochFileCache(topicPartition: TopicPartition, leo: () => LogOffsetM
   /**
     * Delete all entries.
     */
-  override def clearAndFlush() = {
+  def clearAndFlush(): Unit = {
     inWriteLock(lock) {
       epochs.clear()
       flush()
     }
   }
 
-  override def clear() = {
+  def clear(): Unit = {
     inWriteLock(lock) {
       epochs.clear()
     }
   }
 
-  def epochEntries(): ListBuffer[EpochEntry] = {
-    epochs
-  }
-
-  private def earliestOffset(): Long = {
-    if (epochs.isEmpty) -1 else epochs.head.startOffset
-  }
-
-  private def latestOffset(): Long = {
-    if (epochs.isEmpty) -1 else epochs.last.startOffset
-  }
+  // Visible for testing
+  def epochEntries: Seq[EpochEntry] = epochs.values.asScala.toSeq
 
   private def flush(): Unit = {
-    checkpoint.write(epochs)
+    checkpoint.write(epochs.values.asScala)
   }
 
-  def epochChangeMsg(epoch: Int, offset: Long) = s"New: {epoch:$epoch, offset:$offset}, Current: {epoch:$latestEpoch, offset$latestOffset} for Partition: $topicPartition"
-
-  def validateAndMaybeWarn(epoch: Int, offset: Long) = {
-    assert(epoch >= 0, s"Received a PartitionLeaderEpoch assignment for an epoch < 0. This should not happen. ${epochChangeMsg(epoch, offset)}")
-    if (epoch < latestEpoch())
-      warn(s"Received a PartitionLeaderEpoch assignment for an epoch < latestEpoch. " +
-        s"This implies messages have arrived out of order. ${epochChangeMsg(epoch, offset)}")
-    else if (offset < latestOffset())
-      warn(s"Received a PartitionLeaderEpoch assignment for an offset < latest offset for the most recent, stored PartitionLeaderEpoch. " +
-        s"This implies messages have arrived out of order. ${epochChangeMsg(epoch, offset)}")
-  }
 }
 
 // Mapping of epoch to the first offset of the subsequent epoch
-case class EpochEntry(epoch: Int, startOffset: Long)
+case class EpochEntry(epoch: Int, startOffset: Long) {
+  override def toString: String = {
+    s"EpochEntry(epoch=$epoch, startOffset=$startOffset)"
+  }
+}
diff --git a/core/src/main/scala/kafka/tools/ConsoleConsumer.scala b/core/src/main/scala/kafka/tools/ConsoleConsumer.scala
index 7d2d371015783..4de15b9f373f5 100755
--- a/core/src/main/scala/kafka/tools/ConsoleConsumer.scala
+++ b/core/src/main/scala/kafka/tools/ConsoleConsumer.scala
@@ -19,25 +19,25 @@ package kafka.tools
 
 import java.io.PrintStream
 import java.nio.charset.StandardCharsets
+import java.time.Duration
 import java.util.concurrent.CountDownLatch
-import java.util.{Locale, Properties, Random}
+import java.util.regex.Pattern
+import java.util.{Collections, Locale, Map, Properties, Random}
 
 import com.typesafe.scalalogging.LazyLogging
 import joptsimple._
-import kafka.api.OffsetRequest
-import kafka.common.{MessageFormatter, StreamEndException}
-import kafka.consumer._
-import kafka.message._
-import kafka.metrics.KafkaMetricsReporter
-import kafka.utils._
 import kafka.utils.Implicits._
-import org.apache.kafka.clients.consumer.{ConsumerConfig, ConsumerRecord}
-import org.apache.kafka.common.errors.{AuthenticationException, WakeupException}
+import kafka.utils.{Exit, _}
+import org.apache.kafka.clients.consumer.{Consumer, ConsumerConfig, ConsumerRecord, KafkaConsumer}
+import org.apache.kafka.common.{MessageFormatter, TopicPartition}
+import org.apache.kafka.common.errors.{AuthenticationException, TimeoutException, WakeupException}
 import org.apache.kafka.common.record.TimestampType
-import org.apache.kafka.common.serialization.Deserializer
+import org.apache.kafka.common.requests.ListOffsetRequest
+import org.apache.kafka.common.serialization.{ByteArrayDeserializer, Deserializer}
 import org.apache.kafka.common.utils.Utils
 
-import scala.collection.JavaConverters._
+import scala.annotation.nowarn
+import scala.jdk.CollectionConverters._
 
 /**
  * Consumer that dumps messages to standard out.
@@ -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,76 +62,46 @@ 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)
 
-    val consumer =
-      if (conf.useOldConsumer) {
-        checkZk(conf)
-        val props = getOldConsumerProps(conf)
-        checkAndMaybeDeleteOldPath(conf, props)
-        new OldConsumer(conf.filterSpec, props)
-      } else {
-        val timeoutMs = if (conf.timeoutMs >= 0) conf.timeoutMs else Long.MaxValue
-        if (conf.partitionArg.isDefined)
-          new NewShinyConsumer(Option(conf.topicArg), conf.partitionArg, Option(conf.offsetArg), None, getNewConsumerProps(conf), timeoutMs)
-        else
-          new NewShinyConsumer(Option(conf.topicArg), None, None, Option(conf.whitelistArg), getNewConsumerProps(conf), timeoutMs)
-      }
+    val consumerWrapper =
+      if (conf.partitionArg.isDefined)
+        new ConsumerWrapper(Option(conf.topicArg), conf.partitionArg, Option(conf.offsetArg), None, consumer, timeoutMs)
+      else
+        new ConsumerWrapper(Option(conf.topicArg), None, None, Option(conf.whitelistArg), consumer, timeoutMs)
 
-    addShutdownHook(consumer, conf)
+    addShutdownHook(consumerWrapper, conf)
 
-    try {
-      process(conf.maxMessages, conf.formatter, consumer, System.out, conf.skipMessageOnError)
-    } finally {
-      consumer.cleanup()
+    try process(conf.maxMessages, conf.formatter, consumerWrapper, System.out, conf.skipMessageOnError)
+    finally {
+      consumerWrapper.cleanup()
       conf.formatter.close()
       reportRecordCount()
 
-      // if we generated a random group id (as none specified explicitly) then avoid polluting zookeeper with persistent group data, this is a hack
-      if (conf.useOldConsumer && !conf.groupIdPassed)
-        ZkUtils.maybeDeletePath(conf.options.valueOf(conf.zkConnectOpt), "/consumers/" + conf.consumerProps.get("group.id"))
-
       shutdownLatch.countDown()
     }
   }
 
-  def checkZk(config: ConsumerConfig) {
-    if (!checkZkPathExists(config.options.valueOf(config.zkConnectOpt), "/brokers/ids")) {
-      System.err.println("No brokers found in ZK.")
-      Exit.exit(1)
-    }
-
-    if (!config.options.has(config.deleteConsumerOffsetsOpt) && config.options.has(config.resetBeginningOpt) &&
-      checkZkPathExists(config.options.valueOf(config.zkConnectOpt), "/consumers/" + config.consumerProps.getProperty("group.id") + "/offsets")) {
-      System.err.println("Found previous offset information for this group " + config.consumerProps.getProperty("group.id")
-        + ". Please use --delete-consumer-offsets to delete previous offsets metadata")
-      Exit.exit(1)
-    }
-  }
-
-  def addShutdownHook(consumer: BaseConsumer, conf: ConsumerConfig) {
-    Runtime.getRuntime.addShutdownHook(new Thread() {
-      override def run() {
-        consumer.stop()
+  def addShutdownHook(consumer: ConsumerWrapper, conf: ConsumerConfig): Unit = {
+    Exit.addShutdownHook("consumer-shutdown-hook", {
+        consumer.wakeup()
 
         shutdownLatch.await()
 
         if (conf.enableSystestEventsLogging) {
           System.out.println("shutdown_complete")
         }
-      }
     })
   }
 
-  def process(maxMessages: Integer, formatter: MessageFormatter, consumer: BaseConsumer, output: PrintStream, skipMessageOnError: Boolean) {
+  def process(maxMessages: Integer, formatter: MessageFormatter, consumer: ConsumerWrapper, output: PrintStream,
+              skipMessageOnError: Boolean): Unit = {
     while (messageCount < maxMessages || maxMessages == -1) {
-      val msg: BaseConsumerRecord = try {
+      val msg: ConsumerRecord[Array[Byte], Array[Byte]] = try {
         consumer.receive()
       } catch {
-        case _: StreamEndException =>
-          trace("Caught StreamEndException because consumer is shutdown, ignore and terminate.")
-          // Consumer is already closed
-          return
         case _: WakeupException =>
           trace("Caught WakeupException because consumer is shutdown, ignore and terminate.")
           // Consumer will be closed
@@ -161,63 +131,33 @@ object ConsoleConsumer extends Logging {
     }
   }
 
-  def reportRecordCount() {
+  def reportRecordCount(): Unit = {
     System.err.println(s"Processed a total of $messageCount messages")
   }
 
   def checkErr(output: PrintStream, formatter: MessageFormatter): Boolean = {
     val gotError = output.checkError()
     if (gotError) {
-      // This means no one is listening to our output stream any more, time to shutdown
+      // This means no one is listening to our output stream anymore, time to shutdown
       System.err.println("Unable to write to standard out, closing consumer.")
     }
     gotError
   }
 
-  def getOldConsumerProps(config: ConsumerConfig): Properties = {
-    val props = new Properties
-
-    props ++= config.consumerProps
-    props ++= config.extraConsumerProps
-    setAutoOffsetResetValue(config, props)
-    props.put("zookeeper.connect", config.zkConnectionStr)
-
-    if (config.timeoutMs >= 0)
-      props.put("consumer.timeout.ms", config.timeoutMs.toString)
-
-    props
-  }
-
-  def checkAndMaybeDeleteOldPath(config: ConsumerConfig, props: Properties) = {
-    val consumerGroupBasePath = "/consumers/" + props.getProperty("group.id")
-    if (config.options.has(config.deleteConsumerOffsetsOpt)) {
-      ZkUtils.maybeDeletePath(config.options.valueOf(config.zkConnectOpt), consumerGroupBasePath)
-    } else {
-      val resetToBeginning = OffsetRequest.SmallestTimeString == props.getProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG)
-      if (resetToBeginning && checkZkPathExists(config.options.valueOf(config.zkConnectOpt), consumerGroupBasePath + "/offsets")) {
-        System.err.println("Found previous offset information for this group " + props.getProperty("group.id")
-          + ". Please use --delete-consumer-offsets to delete previous offsets metadata")
-        Exit.exit(1)
-      }
-    }
-  }
-
-  def getNewConsumerProps(config: ConsumerConfig): Properties = {
+  private[tools] def consumerProps(config: ConsumerConfig): Properties = {
     val props = new Properties
-
     props ++= config.consumerProps
     props ++= config.extraConsumerProps
     setAutoOffsetResetValue(config, props)
     props.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, config.bootstrapServer)
-    props.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.ByteArrayDeserializer")
-    props.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.ByteArrayDeserializer")
-    props.put(ConsumerConfig.ISOLATION_LEVEL_CONFIG, config.isolationLevel)
+    CommandLineUtils.maybeMergeOptions(
+      props, ConsumerConfig.ISOLATION_LEVEL_CONFIG, config.options, config.isolationLevelOpt)
     props
   }
 
   /**
-    * Used by both getNewConsumerProps and getOldConsumerProps to retrieve the correct value for the
-    * consumer parameter 'auto.offset.reset'.
+    * Used by consumerProps to retrieve the correct value for the consumer parameter 'auto.offset.reset'.
+    *
     * Order of priority is:
     *   1. Explicitly set parameter via --consumer.property command line parameter
     *   2. Explicit --from-beginning given -> 'earliest'
@@ -226,11 +166,8 @@ 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) {
-    val (earliestConfigValue, latestConfigValue) = if (config.useOldConsumer)
-      (OffsetRequest.SmallestTimeString, OffsetRequest.LargestTimeString)
-    else
-      ("earliest", "latest")
+  def setAutoOffsetResetValue(config: ConsumerConfig, props: Properties): Unit = {
+    val (earliestConfigValue, latestConfigValue) = ("earliest", "latest")
 
     if (props.containsKey(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG)) {
       // auto.offset.reset parameter was specified on the command line
@@ -251,20 +188,15 @@ object ConsoleConsumer extends Logging {
     }
   }
 
-  class ConsumerConfig(args: Array[String]) {
-    val parser = new OptionParser(false)
+  class ConsumerConfig(args: Array[String]) extends CommandDefaultOptions(args) {
     val topicIdOpt = parser.accepts("topic", "The topic id to consume on.")
       .withRequiredArg
       .describedAs("topic")
       .ofType(classOf[String])
-    val whitelistOpt = parser.accepts("whitelist", "Whitelist of topics to include for consumption.")
+    val whitelistOpt = parser.accepts("whitelist", "Regular expression specifying whitelist of topics to include for consumption.")
       .withRequiredArg
       .describedAs("whitelist")
       .ofType(classOf[String])
-    val blacklistOpt = parser.accepts("blacklist", "Blacklist of topics to exclude from consumption.")
-      .withRequiredArg
-      .describedAs("blacklist")
-      .ofType(classOf[String])
     val partitionIdOpt = parser.accepts("partition", "The partition to consume from. Consumption " +
       "starts from the end of the partition unless '--offset' is specified.")
       .withRequiredArg
@@ -275,11 +207,6 @@ object ConsoleConsumer extends Logging {
       .describedAs("consume offset")
       .ofType(classOf[String])
       .defaultsTo("latest")
-    val zkConnectOpt = parser.accepts("zookeeper", "REQUIRED (only when using old consumer): The connection string for the zookeeper connection in the form host:port. " +
-      "Multiple URLS can be given to allow fail-over.")
-      .withRequiredArg
-      .describedAs("urls")
-      .ofType(classOf[String])
     val consumerPropertyOpt = parser.accepts("consumer-property", "A mechanism to pass user-defined properties in the form key=value to the consumer.")
       .withRequiredArg
       .describedAs("consumer_prop")
@@ -293,11 +220,27 @@ object ConsoleConsumer extends Logging {
       .describedAs("class")
       .ofType(classOf[String])
       .defaultsTo(classOf[DefaultMessageFormatter].getName)
-    val messageFormatterArgOpt = parser.accepts("property", "The properties to initialize the message formatter.")
+    val messageFormatterArgOpt = parser.accepts("property",
+    """The properties to initialize the message formatter. Default properties include:
+      | print.timestamp=true|false
+      | print.key=true|false
+      | print.offset=true|false
+      | print.partition=true|false
+      | print.headers=true|false
+      | print.value=true|false
+      | key.separator=
+      | line.separator=
+      | headers.separator=
+      | null.literal=
+      | key.deserializer=
+      | value.deserializer=
+      | header.deserializer=
+      |
+      |Users can also pass in customized properties for their formatter; more specifically, users can pass in properties keyed with 'key.deserializer.', 'value.deserializer.' and 'headers.deserializer.' prefixes to configure their deserializers."""
+      .stripMargin)
       .withRequiredArg
       .describedAs("prop")
       .ofType(classOf[String])
-    val deleteConsumerOffsetsOpt = parser.accepts("delete-consumer-offsets", "If specified, the consumer path in zookeeper is deleted when starting up")
     val resetBeginningOpt = parser.accepts("from-beginning", "If the consumer does not already have an established offset to consume from, " +
       "start with the earliest message present in the log rather than the latest message.")
     val maxMessagesOpt = parser.accepts("max-messages", "The maximum number of messages to consume before exiting. If not set, consumption is continual.")
@@ -310,15 +253,7 @@ object ConsoleConsumer extends Logging {
       .ofType(classOf[java.lang.Integer])
     val skipMessageOnErrorOpt = parser.accepts("skip-message-on-error", "If there is an error when processing a message, " +
       "skip it instead of halt.")
-    val csvMetricsReporterEnabledOpt = parser.accepts("csv-reporter-enabled", "If set, the CSV metrics reporter will be enabled")
-    val metricsDirectoryOpt = parser.accepts("metrics-dir", "If csv-reporter-enable is set, and this parameter is" +
-      "set, the csv metrics will be output here")
-      .withRequiredArg
-      .describedAs("metrics directory")
-      .ofType(classOf[java.lang.String])
-    val newConsumerOpt = parser.accepts("new-consumer", "Use the new consumer implementation. This is the default, so " +
-      "this option is deprecated and will be removed in a future release.")
-    val bootstrapServerOpt = parser.accepts("bootstrap-server", "REQUIRED (unless old consumer is used): The server to connect to.")
+    val bootstrapServerOpt = parser.accepts("bootstrap-server", "REQUIRED: The server(s) to connect to.")
       .withRequiredArg
       .describedAs("server to connect to")
       .ofType(classOf[String])
@@ -331,10 +266,10 @@ object ConsoleConsumer extends Logging {
       .describedAs("deserializer for values")
       .ofType(classOf[String])
     val enableSystestEventsLoggingOpt = parser.accepts("enable-systest-events",
-                                                       "Log lifecycle events of the consumer in addition to logging consumed " +
-                                                       "messages. (This is specific for system tests.)")
+      "Log lifecycle events of the consumer in addition to logging consumed " +
+        "messages. (This is specific for system tests.)")
     val isolationLevelOpt = parser.accepts("isolation-level",
-        "Set to read_committed in order to filter out transactional messages which are not committed. Set to read_uncommitted" +
+      "Set to read_committed in order to filter out transactional messages which are not committed. Set to read_uncommitted " +
         "to read all messages.")
       .withRequiredArg()
       .ofType(classOf[String])
@@ -345,16 +280,14 @@ object ConsoleConsumer extends Logging {
       .describedAs("consumer group id")
       .ofType(classOf[String])
 
-    if (args.length == 0)
-      CommandLineUtils.printUsageAndDie(parser, "The console consumer is a tool that reads data from Kafka and outputs it to standard output.")
+    options = tryParse(parser, args)
+
+    CommandLineUtils.printHelpAndExitIfNeeded(this, "This tool helps to read data from Kafka topics and outputs it to standard output.")
 
     var groupIdPassed = true
-    val options: OptionSet = tryParse(parser, args)
-    val useOldConsumer = options.has(zkConnectOpt)
     val enableSystestEventsLogging = options.has(enableSystestEventsLoggingOpt)
 
-    // If using old consumer, exactly one of whitelist/blacklist/topic is required.
-    // If using new consumer, topic must be specified.
+    // topic must be specified.
     var topicArg: String = null
     var whitelistArg: String = null
     var filterSpec: TopicFilter = null
@@ -363,7 +296,6 @@ object ConsoleConsumer extends Logging {
       Utils.loadProps(options.valueOf(consumerConfigOpt))
     else
       new Properties()
-    val zkConnectionStr = options.valueOf(zkConnectOpt)
     val fromBeginning = options.has(resetBeginningOpt)
     val partitionArg = if (options.has(partitionIdOpt)) Some(options.valueOf(partitionIdOpt).intValue) else None
     val skipMessageOnError = options.has(skipMessageOnErrorOpt)
@@ -374,8 +306,7 @@ object ConsoleConsumer extends Logging {
     val bootstrapServer = options.valueOf(bootstrapServerOpt)
     val keyDeserializer = options.valueOf(keyDeserializerOpt)
     val valueDeserializer = options.valueOf(valueDeserializerOpt)
-    val isolationLevel = options.valueOf(isolationLevelOpt).toString
-    val formatter: MessageFormatter = messageFormatterClass.newInstance().asInstanceOf[MessageFormatter]
+    val formatter: MessageFormatter = messageFormatterClass.getDeclaredConstructor().newInstance().asInstanceOf[MessageFormatter]
 
     if (keyDeserializer != null && !keyDeserializer.isEmpty) {
       formatterArgs.setProperty(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, keyDeserializer)
@@ -383,30 +314,14 @@ object ConsoleConsumer extends Logging {
     if (valueDeserializer != null && !valueDeserializer.isEmpty) {
       formatterArgs.setProperty(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, valueDeserializer)
     }
-    formatter.init(formatterArgs)
-
-    if (useOldConsumer) {
-      if (options.has(bootstrapServerOpt))
-        CommandLineUtils.printUsageAndDie(parser, s"Option $bootstrapServerOpt is not valid with $zkConnectOpt.")
-      else if (options.has(newConsumerOpt))
-        CommandLineUtils.printUsageAndDie(parser, s"Option $newConsumerOpt is not valid with $zkConnectOpt.")
-      val topicOrFilterOpt = List(topicIdOpt, whitelistOpt, blacklistOpt).filter(options.has)
-      if (topicOrFilterOpt.size != 1)
-        CommandLineUtils.printUsageAndDie(parser, "Exactly one of whitelist/blacklist/topic is required.")
-      topicArg = options.valueOf(topicOrFilterOpt.head)
-      filterSpec = if (options.has(blacklistOpt)) new Blacklist(topicArg) else new Whitelist(topicArg)
-      Console.err.println("Using the ConsoleConsumer with old consumer is deprecated and will be removed " +
-        s"in a future major release. Consider using the new consumer by passing $bootstrapServerOpt instead of ${zkConnectOpt}.")
-    } else {
-      val topicOrFilterOpt = List(topicIdOpt, whitelistOpt).filter(options.has)
-      if (topicOrFilterOpt.size != 1)
-        CommandLineUtils.printUsageAndDie(parser, "Exactly one of whitelist/topic is required.")
-      topicArg = options.valueOf(topicIdOpt)
-      whitelistArg = options.valueOf(whitelistOpt)
-    }
 
-    if (useOldConsumer && (partitionArg.isDefined || options.has(offsetOpt)))
-      CommandLineUtils.printUsageAndDie(parser, "Partition-offset based consumption is supported in the new consumer only.")
+    formatter.configure(formatterArgs.asScala.asJava)
+
+    val topicOrFilterOpt = List(topicIdOpt, whitelistOpt).filter(options.has)
+    if (topicOrFilterOpt.size != 1)
+      CommandLineUtils.printUsageAndDie(parser, "Exactly one of whitelist/topic is required.")
+    topicArg = options.valueOf(topicIdOpt)
+    whitelistArg = options.valueOf(whitelistOpt)
 
     if (partitionArg.isDefined) {
       if (!options.has(topicIdOpt))
@@ -423,54 +338,35 @@ object ConsoleConsumer extends Logging {
     val offsetArg =
       if (options.has(offsetOpt)) {
         options.valueOf(offsetOpt).toLowerCase(Locale.ROOT) match {
-          case "earliest" => OffsetRequest.EarliestTime
-          case "latest" => OffsetRequest.LatestTime
+          case "earliest" => ListOffsetRequest.EARLIEST_TIMESTAMP
+          case "latest" => ListOffsetRequest.LATEST_TIMESTAMP
           case offsetString =>
-            val offset =
-              try offsetString.toLong
-              catch {
-                case _: NumberFormatException => invalidOffset(offsetString)
-              }
-            if (offset < 0) invalidOffset(offsetString)
-            offset
+            try {
+              val offset = offsetString.toLong
+              if (offset < 0)
+                invalidOffset(offsetString)
+              offset
+            } catch {
+              case _: NumberFormatException => invalidOffset(offsetString)
+            }
         }
       }
-      else if (fromBeginning) OffsetRequest.EarliestTime
-      else OffsetRequest.LatestTime
-
-    if (!useOldConsumer) {
-      CommandLineUtils.checkRequiredArgs(parser, options, bootstrapServerOpt)
-
-      if (options.has(newConsumerOpt)) {
-        Console.err.println("The --new-consumer option is deprecated and will be removed in a future major release." +
-          "The new consumer is used by default if the --bootstrap-server option is provided.")
-      }
-    }
+      else if (fromBeginning) ListOffsetRequest.EARLIEST_TIMESTAMP
+      else ListOffsetRequest.LATEST_TIMESTAMP
 
-    if (options.has(csvMetricsReporterEnabledOpt)) {
-      val csvReporterProps = new Properties()
-      csvReporterProps.put("kafka.metrics.polling.interval.secs", "5")
-      csvReporterProps.put("kafka.metrics.reporters", "kafka.metrics.KafkaCSVMetricsReporter")
-      if (options.has(metricsDirectoryOpt))
-        csvReporterProps.put("kafka.csv.metrics.dir", options.valueOf(metricsDirectoryOpt))
-      else
-        csvReporterProps.put("kafka.csv.metrics.dir", "kafka_metrics")
-      csvReporterProps.put("kafka.csv.metrics.reporter.enabled", "true")
-      val verifiableProps = new VerifiableProperties(csvReporterProps)
-      KafkaMetricsReporter.startReporters(verifiableProps)
-    }
+    CommandLineUtils.checkRequiredArgs(parser, options, bootstrapServerOpt)
 
     // if the group id is provided in more than place (through different means) all values must be the same
     val groupIdsProvided = Set(
-        Option(options.valueOf(groupIdOpt)),                           // via --group
-        Option(consumerProps.get(ConsumerConfig.GROUP_ID_CONFIG)),     // via --consumer-property
-        Option(extraConsumerProps.get(ConsumerConfig.GROUP_ID_CONFIG)) // via --cosumer.config
-      ).flatten
+      Option(options.valueOf(groupIdOpt)), // via --group
+      Option(consumerProps.get(ConsumerConfig.GROUP_ID_CONFIG)), // via --consumer-property
+      Option(extraConsumerProps.get(ConsumerConfig.GROUP_ID_CONFIG)) // via --consumer.config
+    ).flatten
 
     if (groupIdsProvided.size > 1) {
       CommandLineUtils.printUsageAndDie(parser, "The group ids provided in different places (directly using '--group', "
-                                              + "via '--consumer-property', or via '--consumer.config') do not match. "
-                                              + s"Detected group ids: ${groupIdsProvided.mkString("'", "', '", "'")}")
+        + "via '--consumer-property', or via '--consumer.config') do not match. "
+        + s"Detected group ids: ${groupIdsProvided.mkString("'", "', '", "'")}")
     }
 
     groupIdsProvided.headOption match {
@@ -478,9 +374,16 @@ object ConsoleConsumer extends Logging {
         consumerProps.put(ConsumerConfig.GROUP_ID_CONFIG, group)
       case None =>
         consumerProps.put(ConsumerConfig.GROUP_ID_CONFIG, s"console-consumer-${new Random().nextInt(100000)}")
+        // By default, avoid unnecessary expansion of the coordinator cache since
+        // the auto-generated group and its offsets is not intended to be used again
+        if (!consumerProps.containsKey(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG))
+          consumerProps.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false")
         groupIdPassed = false
     }
 
+    if (groupIdPassed && partitionArg.isDefined)
+      CommandLineUtils.printUsageAndDie(parser, "Options group and partition cannot be specified together.")
+
     def tryParse(parser: OptionParser, args: Array[String]): OptionSet = {
       try
         parser.parse(args: _*)
@@ -491,46 +394,107 @@ object ConsoleConsumer extends Logging {
     }
   }
 
-  def checkZkPathExists(zkUrl: String, path: String): Boolean = {
-    try {
-      val zk = ZkUtils.createZkClient(zkUrl, 30 * 1000, 30 * 1000)
-      zk.exists(path)
-    } catch {
-      case _: Throwable => false
+  private[tools] class ConsumerWrapper(topic: Option[String], partitionId: Option[Int], offset: Option[Long], whitelist: Option[String],
+                                       consumer: Consumer[Array[Byte], Array[Byte]], val timeoutMs: Long = Long.MaxValue) {
+    consumerInit()
+    var recordIter = Collections.emptyList[ConsumerRecord[Array[Byte], Array[Byte]]]().iterator()
+
+    def consumerInit(): Unit = {
+      (topic, partitionId, offset, whitelist) match {
+        case (Some(topic), Some(partitionId), Some(offset), None) =>
+          seek(topic, partitionId, offset)
+        case (Some(topic), Some(partitionId), None, None) =>
+          // default to latest if no offset is provided
+          seek(topic, partitionId, ListOffsetRequest.LATEST_TIMESTAMP)
+        case (Some(topic), None, None, None) =>
+          consumer.subscribe(Collections.singletonList(topic))
+        case (None, None, None, Some(whitelist)) =>
+          consumer.subscribe(Pattern.compile(whitelist))
+        case _ =>
+          throw new IllegalArgumentException("An invalid combination of arguments is provided. " +
+            "Exactly one of 'topic' or 'whitelist' must be provided. " +
+            "If 'topic' is provided, an optional 'partition' may also be provided. " +
+            "If 'partition' is provided, an optional 'offset' may also be provided, otherwise, consumption starts from the end of the partition.")
+      }
+    }
+
+    def seek(topic: String, partitionId: Int, offset: Long): Unit = {
+      val topicPartition = new TopicPartition(topic, partitionId)
+      consumer.assign(Collections.singletonList(topicPartition))
+      offset match {
+        case ListOffsetRequest.EARLIEST_TIMESTAMP => consumer.seekToBeginning(Collections.singletonList(topicPartition))
+        case ListOffsetRequest.LATEST_TIMESTAMP => consumer.seekToEnd(Collections.singletonList(topicPartition))
+        case _ => consumer.seek(topicPartition, offset)
+      }
+    }
+
+    def resetUnconsumedOffsets(): Unit = {
+      val smallestUnconsumedOffsets = collection.mutable.Map[TopicPartition, Long]()
+      while (recordIter.hasNext) {
+        val record = recordIter.next()
+        val tp = new TopicPartition(record.topic, record.partition)
+        // avoid auto-committing offsets which haven't been consumed
+        smallestUnconsumedOffsets.getOrElseUpdate(tp, record.offset)
+      }
+      smallestUnconsumedOffsets.forKeyValue { (tp, offset) => consumer.seek(tp, offset) }
+    }
+
+    def receive(): ConsumerRecord[Array[Byte], Array[Byte]] = {
+      if (!recordIter.hasNext) {
+        recordIter = consumer.poll(Duration.ofMillis(timeoutMs)).iterator
+        if (!recordIter.hasNext)
+          throw new TimeoutException()
+      }
+
+      recordIter.next
     }
+
+    def wakeup(): Unit = {
+      this.consumer.wakeup()
+    }
+
+    def cleanup(): Unit = {
+      resetUnconsumedOffsets()
+      this.consumer.close()
+    }
+
   }
 }
 
 class DefaultMessageFormatter extends MessageFormatter {
+  var printTimestamp = false
   var printKey = false
   var printValue = true
-  var printTimestamp = false
-  var keySeparator = "\t".getBytes(StandardCharsets.UTF_8)
-  var lineSeparator = "\n".getBytes(StandardCharsets.UTF_8)
+  var printPartition = false
+  var printOffset = false
+  var printHeaders = false
+  var keySeparator = utfBytes("\t")
+  var lineSeparator = utfBytes("\n")
+  var headersSeparator = utfBytes(",")
+  var nullLiteral = utfBytes("null")
 
   var keyDeserializer: Option[Deserializer[_]] = None
   var valueDeserializer: Option[Deserializer[_]] = None
-
-  override def init(props: Properties) {
-    if (props.containsKey("print.timestamp"))
-      printTimestamp = props.getProperty("print.timestamp").trim.equalsIgnoreCase("true")
-    if (props.containsKey("print.key"))
-      printKey = props.getProperty("print.key").trim.equalsIgnoreCase("true")
-    if (props.containsKey("print.value"))
-      printValue = props.getProperty("print.value").trim.equalsIgnoreCase("true")
-    if (props.containsKey("key.separator"))
-      keySeparator = props.getProperty("key.separator").getBytes(StandardCharsets.UTF_8)
-    if (props.containsKey("line.separator"))
-      lineSeparator = props.getProperty("line.separator").getBytes(StandardCharsets.UTF_8)
-    // Note that `toString` will be called on the instance returned by `Deserializer.deserialize`
-    if (props.containsKey("key.deserializer"))
-      keyDeserializer = Some(Class.forName(props.getProperty("key.deserializer")).newInstance().asInstanceOf[Deserializer[_]])
-    // Note that `toString` will be called on the instance returned by `Deserializer.deserialize`
-    if (props.containsKey("value.deserializer"))
-      valueDeserializer = Some(Class.forName(props.getProperty("value.deserializer")).newInstance().asInstanceOf[Deserializer[_]])
+  var headersDeserializer: Option[Deserializer[_]] = None
+
+  override def configure(configs: Map[String, _]): Unit = {
+    getPropertyIfExists(configs, "print.timestamp", getBoolProperty).foreach(printTimestamp = _)
+    getPropertyIfExists(configs, "print.key", getBoolProperty).foreach(printKey = _)
+    getPropertyIfExists(configs, "print.offset", getBoolProperty).foreach(printOffset = _)
+    getPropertyIfExists(configs, "print.partition", getBoolProperty).foreach(printPartition = _)
+    getPropertyIfExists(configs, "print.headers", getBoolProperty).foreach(printHeaders = _)
+    getPropertyIfExists(configs, "print.value", getBoolProperty).foreach(printValue = _)
+    getPropertyIfExists(configs, "key.separator", getByteProperty).foreach(keySeparator = _)
+    getPropertyIfExists(configs, "line.separator", getByteProperty).foreach(lineSeparator = _)
+    getPropertyIfExists(configs, "headers.separator", getByteProperty).foreach(headersSeparator = _)
+    getPropertyIfExists(configs, "null.literal", getByteProperty).foreach(nullLiteral = _)
+
+    keyDeserializer = getPropertyIfExists(configs, "key.deserializer", getDeserializerProperty(true))
+    valueDeserializer = getPropertyIfExists(configs, "value.deserializer", getDeserializerProperty(false))
+    headersDeserializer = getPropertyIfExists(configs, "headers.deserializer", getDeserializerProperty(false))
   }
 
-  def writeTo(consumerRecord: ConsumerRecord[Array[Byte], Array[Byte]], output: PrintStream) {
+  def writeTo(consumerRecord: ConsumerRecord[Array[Byte], Array[Byte]], output: PrintStream): Unit = {
 
     def writeSeparator(columnSeparator: Boolean): Unit = {
       if (columnSeparator)
@@ -539,39 +503,103 @@ class DefaultMessageFormatter extends MessageFormatter {
         output.write(lineSeparator)
     }
 
-    def write(deserializer: Option[Deserializer[_]], sourceBytes: Array[Byte]) {
-      val nonNullBytes = Option(sourceBytes).getOrElse("null".getBytes(StandardCharsets.UTF_8))
-      val convertedBytes = deserializer.map(_.deserialize(null, nonNullBytes).toString.
-        getBytes(StandardCharsets.UTF_8)).getOrElse(nonNullBytes)
-      output.write(convertedBytes)
+    def deserialize(deserializer: Option[Deserializer[_]], sourceBytes: Array[Byte], topic: String) = {
+      val nonNullBytes = Option(sourceBytes).getOrElse(nullLiteral)
+      val convertedBytes = deserializer
+        .map(d => utfBytes(d.deserialize(topic, consumerRecord.headers, nonNullBytes).toString))
+        .getOrElse(nonNullBytes)
+      convertedBytes
     }
 
     import consumerRecord._
 
     if (printTimestamp) {
       if (timestampType != TimestampType.NO_TIMESTAMP_TYPE)
-        output.write(s"$timestampType:$timestamp".getBytes(StandardCharsets.UTF_8))
+        output.write(utfBytes(s"$timestampType:$timestamp"))
       else
-        output.write(s"NO_TIMESTAMP".getBytes(StandardCharsets.UTF_8))
-      writeSeparator(printKey || printValue)
+        output.write(utfBytes("NO_TIMESTAMP"))
+      writeSeparator(columnSeparator =  printOffset || printPartition || printHeaders || printKey || printValue)
+    }
+
+    if (printPartition) {
+      output.write(utfBytes("Partition:"))
+      output.write(utfBytes(partition().toString))
+      writeSeparator(columnSeparator = printOffset || printHeaders || printKey || printValue)
+    }
+
+    if (printOffset) {
+      output.write(utfBytes("Offset:"))
+      output.write(utfBytes(offset().toString))
+      writeSeparator(columnSeparator = printHeaders || printKey || printValue)
+    }
+
+    if (printHeaders) {
+      val headersIt = headers().iterator.asScala
+      if (headersIt.hasNext) {
+        headersIt.foreach { header =>
+          output.write(utfBytes(header.key() + ":"))
+          output.write(deserialize(headersDeserializer, header.value(), topic))
+          if (headersIt.hasNext) {
+            output.write(headersSeparator)
+          }
+        }
+      } else {
+        output.write(utfBytes("NO_HEADERS"))
+      }
+      writeSeparator(columnSeparator = printKey || printValue)
     }
 
     if (printKey) {
-      write(keyDeserializer, key)
-      writeSeparator(printValue)
+      output.write(deserialize(keyDeserializer, key, topic))
+      writeSeparator(columnSeparator = printValue)
     }
 
     if (printValue) {
-      write(valueDeserializer, value)
+      output.write(deserialize(valueDeserializer, value, topic))
       output.write(lineSeparator)
     }
   }
+
+  private def propertiesWithKeyPrefixStripped(prefix: String, configs: Map[String, _]): Map[String, _] = {
+    val newConfigs = collection.mutable.Map[String, Any]()
+    configs.asScala.foreach { case (key, value) =>
+      if (key.startsWith(prefix) && key.length > prefix.length)
+        newConfigs.put(key.substring(prefix.length), value)
+    }
+    newConfigs.asJava
+  }
+
+  private def utfBytes(str: String) = str.getBytes(StandardCharsets.UTF_8)
+
+  private def getByteProperty(configs: Map[String, _], key: String): Array[Byte] = {
+    utfBytes(configs.get(key).asInstanceOf[String])
+  }
+
+  private def getBoolProperty(configs: Map[String, _], key: String): Boolean = {
+    configs.get(key).asInstanceOf[String].trim.equalsIgnoreCase("true")
+  }
+
+  private def getDeserializerProperty(isKey: Boolean)(configs: Map[String, _], propertyName: String): Deserializer[_] = {
+    val deserializer = Class.forName(configs.get(propertyName).asInstanceOf[String]).newInstance().asInstanceOf[Deserializer[_]]
+    val deserializerConfig = propertiesWithKeyPrefixStripped(propertyName + ".", configs)
+      .asScala
+      .asJava
+    deserializer.configure(deserializerConfig, isKey)
+    deserializer
+  }
+
+  private def getPropertyIfExists[T](configs: Map[String, _], key: String, getter: (Map[String, _], String) => T): Option[T] = {
+    if (configs.containsKey(key))
+      Some(getter(configs, key))
+    else
+      None
+  }
 }
 
 class LoggingMessageFormatter extends MessageFormatter with LazyLogging {
   private val defaultWriter: DefaultMessageFormatter = new DefaultMessageFormatter
 
-  override def init(props: Properties): Unit = defaultWriter.init(props)
+  override def configure(configs: Map[String, _]): Unit = defaultWriter.configure(configs)
 
   def writeTo(consumerRecord: ConsumerRecord[Array[Byte], Array[Byte]], output: PrintStream): Unit = {
     import consumerRecord._
@@ -583,29 +611,22 @@ class LoggingMessageFormatter extends MessageFormatter with LazyLogging {
 }
 
 class NoOpMessageFormatter extends MessageFormatter {
-  override def init(props: Properties) {}
 
-  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) {
-    topicStr = props.getProperty("topic")
-    if (topicStr != null)
-      topicStr = topicStr + ":"
+  override def configure(configs: Map[String, _]): Unit = {
+    topicStr = if (configs.containsKey("topic"))
+      configs.get("topic").toString + ":"
     else
-      topicStr = ""
+      ""
   }
 
-  def writeTo(consumerRecord: ConsumerRecord[Array[Byte], Array[Byte]], output: PrintStream) {
-    import consumerRecord._
-    val chksum =
-      if (timestampType != TimestampType.NO_TIMESTAMP_TYPE)
-        new Message(value, key, timestamp, timestampType, NoCompressionCodec, 0, -1, Message.MagicValue_V1).checksum
-      else
-        new Message(value, key, Message.NoTimestamp, Message.MagicValue_V0).checksum
-    output.println(topicStr + "checksum:" + chksum)
+  @nowarn("cat=deprecation")
+  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 39bb0ff967105..7c221baf69a2b 100644
--- a/core/src/main/scala/kafka/tools/ConsoleProducer.scala
+++ b/core/src/main/scala/kafka/tools/ConsoleProducer.scala
@@ -17,50 +17,41 @@
 
 package kafka.tools
 
-import kafka.common._
-import kafka.message._
-import kafka.serializer._
-import kafka.utils.{CommandLineUtils, Exit, ToolsUtils}
-import kafka.utils.Implicits._
-import kafka.producer.{NewShinyProducer, OldProducer}
-import java.util.Properties
 import java.io._
 import java.nio.charset.StandardCharsets
+import java.util.Properties
 
-import joptsimple._
-import org.apache.kafka.clients.producer.{ProducerConfig, ProducerRecord}
+import joptsimple.{OptionException, OptionParser, OptionSet}
+import kafka.common._
+import kafka.message._
+import kafka.utils.Implicits._
+import kafka.utils.{CommandDefaultOptions, CommandLineUtils, Exit, ToolsUtils}
+import org.apache.kafka.clients.producer.internals.ErrorLoggingCallback
+import org.apache.kafka.clients.producer.{KafkaProducer, ProducerConfig, ProducerRecord}
+import org.apache.kafka.common.KafkaException
 import org.apache.kafka.common.utils.Utils
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 
 object ConsoleProducer {
 
-  def main(args: Array[String]) {
+  def main(args: Array[String]): Unit = {
 
     try {
         val config = new ProducerConfig(args)
-        val reader = Class.forName(config.readerClass).newInstance().asInstanceOf[MessageReader]
+        val reader = Class.forName(config.readerClass).getDeclaredConstructor().newInstance().asInstanceOf[MessageReader]
         reader.init(System.in, getReaderProps(config))
 
-        val producer =
-          if(config.useOldProducer) {
-            new OldProducer(getOldProducerProps(config))
-          } else {
-            new NewShinyProducer(getNewProducerProps(config))
-          }
+        val producer = new KafkaProducer[Array[Byte], Array[Byte]](producerProps(config))
 
-        Runtime.getRuntime.addShutdownHook(new Thread() {
-          override def run() {
-            producer.close()
-          }
-        })
+    Exit.addShutdownHook("producer-shutdown-hook", producer.close)
 
-        var message: ProducerRecord[Array[Byte], Array[Byte]] = null
+        var record: ProducerRecord[Array[Byte], Array[Byte]] = null
         do {
-          message = reader.readMessage()
-          if (message != null)
-            producer.send(message.topic, message.key, message.value)
-        } while (message != null)
+          record = reader.readMessage()
+          if (record != null)
+            send(producer, record, config.sync)
+        } while (record != null)
     } catch {
       case e: joptsimple.OptionException =>
         System.err.println(e.getMessage)
@@ -72,79 +63,80 @@ object ConsoleProducer {
     Exit.exit(0)
   }
 
+  private def send(producer: KafkaProducer[Array[Byte], Array[Byte]],
+                         record: ProducerRecord[Array[Byte], Array[Byte]], sync: Boolean): Unit = {
+    if (sync)
+      producer.send(record).get()
+    else
+      producer.send(record, new ErrorLoggingCallback(record.topic, record.key, record.value, false))
+  }
+
   def getReaderProps(config: ProducerConfig): Properties = {
     val props = new Properties
-    props.put("topic",config.topic)
+    props.put("topic", config.topic)
     props ++= config.cmdLineProps
     props
   }
 
-  def getOldProducerProps(config: ProducerConfig): Properties = {
-    val props = producerProps(config)
-
-    props.put("metadata.broker.list", config.brokerList)
-    props.put("compression.codec", config.compressionCodec)
-    props.put("producer.type", if(config.sync) "sync" else "async")
-    props.put("batch.num.messages", config.batchSize.toString)
-    props.put("message.send.max.retries", config.messageSendMaxRetries.toString)
-    props.put("retry.backoff.ms", config.retryBackoffMs.toString)
-    props.put("queue.buffering.max.ms", config.sendTimeout.toString)
-    props.put("queue.buffering.max.messages", config.queueSize.toString)
-    props.put("queue.enqueue.timeout.ms", config.queueEnqueueTimeoutMs.toString)
-    props.put("request.required.acks", config.requestRequiredAcks)
-    props.put("request.timeout.ms", config.requestTimeoutMs.toString)
-    props.put("key.serializer.class", config.keyEncoderClass)
-    props.put("serializer.class", config.valueEncoderClass)
-    props.put("send.buffer.bytes", config.socketBuffer.toString)
-    props.put("topic.metadata.refresh.interval.ms", config.metadataExpiryMs.toString)
-    props.put("client.id", "console-producer")
-
-    props
-  }
-
-  private def producerProps(config: ProducerConfig): Properties = {
+  def producerProps(config: ProducerConfig): Properties = {
     val props =
       if (config.options.has(config.producerConfigOpt))
         Utils.loadProps(config.options.valueOf(config.producerConfigOpt))
       else new Properties
+
     props ++= config.extraProducerProps
-    props
-  }
 
-  def getNewProducerProps(config: ProducerConfig): Properties = {
-    val props = producerProps(config)
+    if (config.bootstrapServer != null)
+      props.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, config.bootstrapServer)
+    else
+      props.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, config.brokerList)
 
-    props.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, config.brokerList)
     props.put(ProducerConfig.COMPRESSION_TYPE_CONFIG, config.compressionCodec)
-    props.put(ProducerConfig.SEND_BUFFER_CONFIG, config.socketBuffer.toString)
-    props.put(ProducerConfig.RETRY_BACKOFF_MS_CONFIG, config.retryBackoffMs.toString)
-    props.put(ProducerConfig.METADATA_MAX_AGE_CONFIG, config.metadataExpiryMs.toString)
-    props.put(ProducerConfig.MAX_BLOCK_MS_CONFIG, config.maxBlockMs.toString)
-    props.put(ProducerConfig.ACKS_CONFIG, config.requestRequiredAcks)
-    props.put(ProducerConfig.REQUEST_TIMEOUT_MS_CONFIG, config.requestTimeoutMs.toString)
-    props.put(ProducerConfig.RETRIES_CONFIG, config.messageSendMaxRetries.toString)
-    props.put(ProducerConfig.LINGER_MS_CONFIG, config.sendTimeout.toString)
-    props.put(ProducerConfig.BUFFER_MEMORY_CONFIG, config.maxMemoryBytes.toString)
-    props.put(ProducerConfig.BATCH_SIZE_CONFIG, config.maxPartitionMemoryBytes.toString)
-    props.put(ProducerConfig.CLIENT_ID_CONFIG, "console-producer")
+    if (props.getProperty(ProducerConfig.CLIENT_ID_CONFIG) == null)
+      props.put(ProducerConfig.CLIENT_ID_CONFIG, "console-producer")
     props.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.ByteArraySerializer")
     props.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.ByteArraySerializer")
 
+    CommandLineUtils.maybeMergeOptions(
+      props, ProducerConfig.LINGER_MS_CONFIG, config.options, config.sendTimeoutOpt)
+    CommandLineUtils.maybeMergeOptions(
+      props, ProducerConfig.ACKS_CONFIG, config.options, config.requestRequiredAcksOpt)
+    CommandLineUtils.maybeMergeOptions(
+      props, ProducerConfig.REQUEST_TIMEOUT_MS_CONFIG, config.options, config.requestTimeoutMsOpt)
+    CommandLineUtils.maybeMergeOptions(
+      props, ProducerConfig.RETRIES_CONFIG, config.options, config.messageSendMaxRetriesOpt)
+    CommandLineUtils.maybeMergeOptions(
+      props, ProducerConfig.RETRY_BACKOFF_MS_CONFIG, config.options, config.retryBackoffMsOpt)
+    CommandLineUtils.maybeMergeOptions(
+      props, ProducerConfig.SEND_BUFFER_CONFIG, config.options, config.socketBufferSizeOpt)
+    CommandLineUtils.maybeMergeOptions(
+      props, ProducerConfig.BUFFER_MEMORY_CONFIG, config.options, config.maxMemoryBytesOpt)
+    CommandLineUtils.maybeMergeOptions(
+      props, ProducerConfig.BATCH_SIZE_CONFIG, config.options, config.maxPartitionMemoryBytesOpt)
+    CommandLineUtils.maybeMergeOptions(
+      props, ProducerConfig.METADATA_MAX_AGE_CONFIG, config.options, config.metadataExpiryMsOpt)
+    CommandLineUtils.maybeMergeOptions(
+      props, ProducerConfig.MAX_BLOCK_MS_CONFIG, config.options, config.maxBlockMsOpt)
+
     props
   }
 
-  class ProducerConfig(args: Array[String]) {
-    val parser = new OptionParser(false)
+  class ProducerConfig(args: Array[String]) extends CommandDefaultOptions(args) {
     val topicOpt = parser.accepts("topic", "REQUIRED: The topic id to produce messages to.")
       .withRequiredArg
       .describedAs("topic")
       .ofType(classOf[String])
-    val brokerListOpt = parser.accepts("broker-list", "REQUIRED: The broker list string in the form HOST1:PORT1,HOST2:PORT2.")
+    val brokerListOpt = parser.accepts("broker-list", "DEPRECATED, use --bootstrap-server instead; ignored if --bootstrap-server is specified.  The broker list string in the form HOST1:PORT1,HOST2:PORT2.")
       .withRequiredArg
       .describedAs("broker-list")
       .ofType(classOf[String])
+    val bootstrapServerOpt = parser.accepts("bootstrap-server", "REQUIRED unless --broker-list(deprecated) is specified. The server(s) to connect to. The broker list string in the form HOST1:PORT1,HOST2:PORT2.")
+      .requiredUnless("broker-list")
+      .withRequiredArg
+      .describedAs("server to connect to")
+      .ofType(classOf[String])
     val syncOpt = parser.accepts("sync", "If set message send requests to the brokers are synchronously, one at a time as they arrive.")
-    val compressionCodecOpt = parser.accepts("compression-codec", "The compression codec: either 'none', 'gzip', 'snappy', or 'lz4'." +
+    val compressionCodecOpt = parser.accepts("compression-codec", "The compression codec: either 'none', 'gzip', 'snappy', 'lz4', or 'zstd'." +
                                                                   "If specified without value, then it defaults to 'gzip'")
                                     .withOptionalArg()
                                     .describedAs("compression-codec")
@@ -154,7 +146,7 @@ object ConsoleProducer {
       .describedAs("size")
       .ofType(classOf[java.lang.Integer])
       .defaultsTo(200)
-    val messageSendMaxRetriesOpt = parser.accepts("message-send-max-retries", "Brokers can fail receiving the message for multiple reasons, and being unavailable transiently is just one of them. This property specifies the number of retires before the producer give up and drop this message.")
+    val messageSendMaxRetriesOpt = parser.accepts("message-send-max-retries", "Brokers can fail receiving the message for multiple reasons, and being unavailable transiently is just one of them. This property specifies the number of retries before the producer give up and drop this message.")
       .withRequiredArg
       .ofType(classOf[java.lang.Integer])
       .defaultsTo(3)
@@ -168,17 +160,6 @@ object ConsoleProducer {
       .describedAs("timeout_ms")
       .ofType(classOf[java.lang.Integer])
       .defaultsTo(1000)
-    val queueSizeOpt = parser.accepts("queue-size", "If set and the producer is running in asynchronous mode, this gives the maximum amount of " +
-      " messages will queue awaiting sufficient batch size.")
-      .withRequiredArg
-      .describedAs("queue_size")
-      .ofType(classOf[java.lang.Integer])
-      .defaultsTo(10000)
-    val queueEnqueueTimeoutMsOpt = parser.accepts("queue-enqueuetimeout-ms", "Timeout for event enqueue")
-      .withRequiredArg
-      .describedAs("queue enqueuetimeout ms")
-      .ofType(classOf[java.lang.Integer])
-      .defaultsTo(Int.MaxValue)
     val requestRequiredAcksOpt = parser.accepts("request-required-acks", "The required acks of the producer requests")
       .withRequiredArg
       .describedAs("request required acks")
@@ -214,16 +195,6 @@ object ConsoleProducer {
       .describedAs("memory in bytes per partition")
       .ofType(classOf[java.lang.Long])
       .defaultsTo(16 * 1024L)
-    val valueEncoderOpt = parser.accepts("value-serializer", "The class name of the message encoder implementation to use for serializing values.")
-      .withRequiredArg
-      .describedAs("encoder_class")
-      .ofType(classOf[java.lang.String])
-      .defaultsTo(classOf[DefaultEncoder].getName)
-    val keyEncoderOpt = parser.accepts("key-serializer", "The class name of the message encoder implementation to use for serializing keys.")
-      .withRequiredArg
-      .describedAs("encoder_class")
-      .ofType(classOf[java.lang.String])
-      .defaultsTo(classOf[DefaultEncoder].getName)
     val messageReaderOpt = parser.accepts("line-reader", "The class name of the class to use for reading lines from standard in. " +
       "By default each line is read as a separate message.")
       .withRequiredArg
@@ -236,7 +207,10 @@ object ConsoleProducer {
       .ofType(classOf[java.lang.Integer])
       .defaultsTo(1024*100)
     val propertyOpt = parser.accepts("property", "A mechanism to pass user-defined properties in the form key=value to the message reader. " +
-      "This allows custom configuration for a user-defined message reader.")
+      "This allows custom configuration for a user-defined message reader. Default properties include:\n" +
+      "\tparse.key=true|false\n" +
+      "\tkey.separator=\n" +
+      "\tignore.error=true|false")
       .withRequiredArg
       .describedAs("prop")
       .ofType(classOf[String])
@@ -248,17 +222,21 @@ object ConsoleProducer {
       .withRequiredArg
       .describedAs("config file")
       .ofType(classOf[String])
-    val useOldProducerOpt = parser.accepts("old-producer", "Use the old producer implementation.")
 
-    val options = parser.parse(args : _*)
-    if(args.length == 0)
-      CommandLineUtils.printUsageAndDie(parser, "Read data from standard input and publish it to Kafka.")
-    CommandLineUtils.checkRequiredArgs(parser, options, topicOpt, brokerListOpt)
+    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)
 
-    val useOldProducer = options.has(useOldProducerOpt)
     val topic = options.valueOf(topicOpt)
+
+    val bootstrapServer = options.valueOf(bootstrapServerOpt)
     val brokerList = options.valueOf(brokerListOpt)
-    ToolsUtils.validatePortOrDie(parser,brokerList)
+
+    val brokerHostsAndPorts = options.valueOf(if (options.has(bootstrapServerOpt)) bootstrapServerOpt else brokerListOpt)
+    ToolsUtils.validatePortOrDie(parser, brokerHostsAndPorts)
+
     val sync = options.has(syncOpt)
     val compressionCodecOptionValue = options.valueOf(compressionCodecOpt)
     val compressionCodec = if (options.has(compressionCodecOpt))
@@ -267,24 +245,18 @@ object ConsoleProducer {
                              else compressionCodecOptionValue
                            else NoCompressionCodec.name
     val batchSize = options.valueOf(batchSizeOpt)
-    val sendTimeout = options.valueOf(sendTimeoutOpt)
-    val queueSize = options.valueOf(queueSizeOpt)
-    val queueEnqueueTimeoutMs = options.valueOf(queueEnqueueTimeoutMsOpt)
-    val requestRequiredAcks = options.valueOf(requestRequiredAcksOpt)
-    val requestTimeoutMs = options.valueOf(requestTimeoutMsOpt)
-    val messageSendMaxRetries = options.valueOf(messageSendMaxRetriesOpt)
-    val retryBackoffMs = options.valueOf(retryBackoffMsOpt)
-    val keyEncoderClass = options.valueOf(keyEncoderOpt)
-    val valueEncoderClass = options.valueOf(valueEncoderOpt)
     val readerClass = options.valueOf(messageReaderOpt)
-    val socketBuffer = options.valueOf(socketBufferSizeOpt)
     val cmdLineProps = CommandLineUtils.parseKeyValueArgs(options.valuesOf(propertyOpt).asScala)
     val extraProducerProps = CommandLineUtils.parseKeyValueArgs(options.valuesOf(producerPropertyOpt).asScala)
-    /* new producer related configs */
-    val maxMemoryBytes = options.valueOf(maxMemoryBytesOpt)
-    val maxPartitionMemoryBytes = options.valueOf(maxPartitionMemoryBytesOpt)
-    val metadataExpiryMs = options.valueOf(metadataExpiryMsOpt)
-    val maxBlockMs = options.valueOf(maxBlockMsOpt)
+
+    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 {
@@ -294,8 +266,9 @@ object ConsoleProducer {
     var keySeparator = "\t"
     var ignoreError = false
     var lineNumber = 0
+    var printPrompt = System.console != null
 
-    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")
@@ -308,7 +281,8 @@ object ConsoleProducer {
 
     override def readMessage() = {
       lineNumber += 1
-      print(">")
+      if (printPrompt)
+        print(">")
       (reader.readLine(), parseKey) match {
         case (null, _) => null
         case (line, true) =>
diff --git a/core/src/main/scala/kafka/tools/ConsumerPerformance.scala b/core/src/main/scala/kafka/tools/ConsumerPerformance.scala
index a3e60e652c4ab..44324e898c316 100644
--- a/core/src/main/scala/kafka/tools/ConsumerPerformance.scala
+++ b/core/src/main/scala/kafka/tools/ConsumerPerformance.scala
@@ -17,28 +17,21 @@
 
 package kafka.tools
 
+import java.text.SimpleDateFormat
+import java.time.Duration
 import java.util
-
-import scala.collection.JavaConverters._
 import java.util.concurrent.atomic.AtomicLong
-import java.nio.channels.ClosedByInterruptException
+import java.util.{Properties, Random}
 
+import com.typesafe.scalalogging.LazyLogging
+import joptsimple.OptionException
+import kafka.utils.{CommandLineUtils, ToolsUtils}
 import org.apache.kafka.clients.consumer.{ConsumerRebalanceListener, KafkaConsumer}
 import org.apache.kafka.common.serialization.ByteArrayDeserializer
 import org.apache.kafka.common.utils.Utils
 import org.apache.kafka.common.{Metric, MetricName, TopicPartition}
-import kafka.utils.{CommandLineUtils, ToolsUtils}
-import java.util.{Collections, Properties, Random}
-
-import kafka.consumer.Consumer
-import kafka.consumer.ConsumerConnector
-import kafka.consumer.KafkaStream
-import kafka.consumer.ConsumerTimeoutException
-import java.text.SimpleDateFormat
-import java.util.concurrent.atomic.AtomicBoolean
-
-import com.typesafe.scalalogging.LazyLogging
 
+import scala.jdk.CollectionConverters._
 import scala.collection.mutable
 
 /**
@@ -52,70 +45,38 @@ object ConsumerPerformance extends LazyLogging {
     logger.info("Starting consumer...")
     val totalMessagesRead = new AtomicLong(0)
     val totalBytesRead = new AtomicLong(0)
-    val consumerTimeout = new AtomicBoolean(false)
     var metrics: mutable.Map[MetricName, _ <: Metric] = null
     val joinGroupTimeInMs = new AtomicLong(0)
 
-    if (!config.hideHeader) {
-      printHeader(config.showDetailedStats, config.useOldConsumer)
-    }
+    if (!config.hideHeader)
+      printHeader(config.showDetailedStats)
 
     var startMs, endMs = 0L
-    if (!config.useOldConsumer) {
-      val consumer = new KafkaConsumer[Array[Byte], Array[Byte]](config.props)
-      consumer.subscribe(Collections.singletonList(config.topic))
-      startMs = System.currentTimeMillis
-      consume(consumer, List(config.topic), config.numMessages, 1000, config, totalMessagesRead, totalBytesRead, joinGroupTimeInMs, startMs)
-      endMs = System.currentTimeMillis
-
-      if (config.printMetrics) {
-        metrics = consumer.metrics().asScala
-      }
-      consumer.close()
-    } else {
-      import kafka.consumer.ConsumerConfig
-      val consumerConfig = new ConsumerConfig(config.props)
-      val consumerConnector: ConsumerConnector = Consumer.create(consumerConfig)
-      val topicMessageStreams = consumerConnector.createMessageStreams(Map(config.topic -> config.numThreads))
-      var threadList = List[ConsumerPerfThread]()
-      for (streamList <- topicMessageStreams.values)
-        for (i <- 0 until streamList.length)
-          threadList ::= new ConsumerPerfThread(i, "kafka-zk-consumer-" + i, streamList(i), config, totalMessagesRead, totalBytesRead, consumerTimeout)
-
-      logger.info("Sleeping for 1 second.")
-      Thread.sleep(1000)
-      logger.info("starting threads")
-      startMs = System.currentTimeMillis
-      for (thread <- threadList)
-        thread.start()
-      for (thread <- threadList)
-        thread.join()
-      endMs =
-        if (consumerTimeout.get()) System.currentTimeMillis - consumerConfig.consumerTimeoutMs
-        else System.currentTimeMillis
-      consumerConnector.shutdown()
+    val consumer = new KafkaConsumer[Array[Byte], Array[Byte]](config.props)
+    startMs = System.currentTimeMillis
+    consume(consumer, List(config.topic), config.numMessages, config.recordFetchTimeoutMs, config, totalMessagesRead, totalBytesRead, joinGroupTimeInMs, startMs)
+    endMs = System.currentTimeMillis
+
+    if (config.printMetrics) {
+      metrics = consumer.metrics.asScala
     }
+    consumer.close()
     val elapsedSecs = (endMs - startMs) / 1000.0
     val fetchTimeInMs = (endMs - startMs) - joinGroupTimeInMs.get
     if (!config.showDetailedStats) {
       val totalMBRead = (totalBytesRead.get * 1.0) / (1024 * 1024)
-      print("%s, %s, %.4f, %.4f, %d, %.4f".format(
+      println("%s, %s, %.4f, %.4f, %d, %.4f, %d, %d, %.4f, %.4f".format(
         config.dateFormat.format(startMs),
         config.dateFormat.format(endMs),
         totalMBRead,
         totalMBRead / elapsedSecs,
         totalMessagesRead.get,
-        totalMessagesRead.get / elapsedSecs
+        totalMessagesRead.get / elapsedSecs,
+        joinGroupTimeInMs.get,
+        fetchTimeInMs,
+        totalMBRead / (fetchTimeInMs / 1000.0),
+        totalMessagesRead.get / (fetchTimeInMs / 1000.0)
       ))
-      if (!config.useOldConsumer) {
-        print(", %d, %d, %.4f, %.4f".format(
-          joinGroupTimeInMs.get,
-          fetchTimeInMs,
-          totalMBRead / (fetchTimeInMs / 1000.0),
-          totalMessagesRead.get / (fetchTimeInMs / 1000.0)
-        ))
-      }
-      println()
     }
 
     if (metrics != null) {
@@ -124,13 +85,12 @@ object ConsumerPerformance extends LazyLogging {
 
   }
 
-  private[tools] def printHeader(showDetailedStats: Boolean, useOldConsumer: Boolean): Unit = {
-    val newFieldsInHeader = if (!useOldConsumer) ", rebalance.time.ms, fetch.time.ms, fetch.MB.sec, fetch.nMsg.sec" else ""
-    if (!showDetailedStats) {
-        println("start.time, end.time, data.consumed.in.MB, MB.sec, data.consumed.in.nMsg, nMsg.sec" + newFieldsInHeader)
-      } else {
-        println("time, threadId, data.consumed.in.MB, MB.sec, data.consumed.in.nMsg, nMsg.sec" + newFieldsInHeader)
-    }
+  private[tools] def printHeader(showDetailedStats: Boolean): Unit = {
+    val newFieldsInHeader = ", rebalance.time.ms, fetch.time.ms, fetch.MB.sec, fetch.nMsg.sec"
+    if (!showDetailedStats)
+      println("start.time, end.time, data.consumed.in.MB, MB.sec, data.consumed.in.nMsg, nMsg.sec" + newFieldsInHeader)
+    else
+      println("time, threadId, data.consumed.in.MB, MB.sec, data.consumed.in.nMsg, nMsg.sec" + newFieldsInHeader)
   }
 
   def consume(consumer: KafkaConsumer[Array[Byte], Array[Byte]],
@@ -141,33 +101,30 @@ object ConsumerPerformance extends LazyLogging {
               totalMessagesRead: AtomicLong,
               totalBytesRead: AtomicLong,
               joinTime: AtomicLong,
-              testStartTime: Long) {
+              testStartTime: Long): Unit = {
     var bytesRead = 0L
     var messagesRead = 0L
     var lastBytesRead = 0L
     var lastMessagesRead = 0L
-    var joinStart = 0L
+    var joinStart = System.currentTimeMillis
     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
       }})
-    consumer.poll(0)
-    consumer.seekToBeginning(Collections.emptyList())
 
     // Now start the benchmark
-    val startMs = System.currentTimeMillis
-    var lastReportTime: Long = startMs
-    var lastConsumedTime = System.currentTimeMillis
-    var currentTimeMillis = lastConsumedTime
+    var currentTimeMillis = System.currentTimeMillis
+    var lastReportTime: Long = currentTimeMillis
+    var lastConsumedTime = currentTimeMillis
 
     while (messagesRead < count && currentTimeMillis - lastConsumedTime <= timeout) {
-      val records = consumer.poll(100).asScala
+      val records = consumer.poll(Duration.ofMillis(100)).asScala
       currentTimeMillis = System.currentTimeMillis
       if (records.nonEmpty)
         lastConsumedTime = currentTimeMillis
@@ -180,7 +137,7 @@ object ConsumerPerformance extends LazyLogging {
 
         if (currentTimeMillis - lastReportTime >= config.reportingInterval) {
           if (config.showDetailedStats)
-            printNewConsumerProgress(0, bytesRead, lastBytesRead, messagesRead, lastMessagesRead,
+            printConsumerProgress(0, bytesRead, lastBytesRead, messagesRead, lastMessagesRead,
               lastReportTime, currentTimeMillis, config.dateFormat, joinTimeMsInSingleRound)
           joinTimeMsInSingleRound = 0L
           lastReportTime = currentTimeMillis
@@ -190,23 +147,14 @@ object ConsumerPerformance extends LazyLogging {
       }
     }
 
+    if (messagesRead < count)
+      println(s"WARNING: Exiting before consuming the expected number of messages: timeout ($timeout ms) exceeded. " +
+        "You can use the --timeout option to increase the timeout.")
     totalMessagesRead.set(messagesRead)
     totalBytesRead.set(bytesRead)
   }
 
-  def printOldConsumerProgress(id: Int,
-                               bytesRead: Long,
-                               lastBytesRead: Long,
-                               messagesRead: Long,
-                               lastMessagesRead: Long,
-                               startMs: Long,
-                               endMs: Long,
-                               dateFormat: SimpleDateFormat): Unit = {
-    printBasicProgress(id, bytesRead, lastBytesRead, messagesRead, lastMessagesRead, startMs, endMs, dateFormat)
-    println()
-  }
-
-  def printNewConsumerProgress(id: Int,
+  def printConsumerProgress(id: Int,
                                bytesRead: Long,
                                lastBytesRead: Long,
                                messagesRead: Long,
@@ -228,7 +176,7 @@ object ConsumerPerformance extends LazyLogging {
                                  startMs: Long,
                                  endMs: Long,
                                  dateFormat: SimpleDateFormat): Unit = {
-    val elapsedMs: Double = endMs - startMs
+    val elapsedMs: Double = (endMs - startMs).toDouble
     val totalMbRead = (bytesRead * 1.0) / (1024 * 1024)
     val intervalMbRead = ((bytesRead - lastBytesRead) * 1.0) / (1024 * 1024)
     val intervalMbPerSec = 1000.0 * intervalMbRead / elapsedMs
@@ -255,14 +203,14 @@ object ConsumerPerformance extends LazyLogging {
   }
 
   class ConsumerPerfConfig(args: Array[String]) extends PerfConfig(args) {
-    val zkConnectOpt = parser.accepts("zookeeper", "REQUIRED (only when using old consumer): The connection string for the zookeeper connection in the form host:port. " +
-      "Multiple URLS can be given to allow fail-over. This option is only used with the old consumer.")
+    val brokerListOpt = parser.accepts("broker-list", "DEPRECATED, use --bootstrap-server instead; ignored if --bootstrap-server is specified.  The broker list string in the form HOST1:PORT1,HOST2:PORT2.")
       .withRequiredArg
-      .describedAs("urls")
+      .describedAs("broker-list")
       .ofType(classOf[String])
-    val bootstrapServersOpt = parser.accepts("broker-list", "REQUIRED (unless old consumer is used): A broker list to use for connecting if using the new consumer.")
-      .withRequiredArg()
-      .describedAs("host")
+    val bootstrapServerOpt = parser.accepts("bootstrap-server", "REQUIRED unless --broker-list(deprecated) is specified. The server(s) to connect to.")
+      .requiredUnless("broker-list")
+      .withRequiredArg
+      .describedAs("server to connect to")
       .ofType(classOf[String])
     val topicOpt = parser.accepts("topic", "REQUIRED: The topic to consume from.")
       .withRequiredArg
@@ -285,68 +233,62 @@ object ConsumerPerformance extends LazyLogging {
       .describedAs("size")
       .ofType(classOf[java.lang.Integer])
       .defaultsTo(2 * 1024 * 1024)
-    val numThreadsOpt = parser.accepts("threads", "Number of processing threads.")
+    val numThreadsOpt = parser.accepts("threads", "DEPRECATED AND IGNORED: Number of processing threads.")
       .withRequiredArg
       .describedAs("count")
       .ofType(classOf[java.lang.Integer])
       .defaultsTo(10)
-    val numFetchersOpt = parser.accepts("num-fetch-threads", "Number of fetcher threads.")
+    val numFetchersOpt = parser.accepts("num-fetch-threads", "DEPRECATED AND IGNORED: Number of fetcher threads.")
       .withRequiredArg
       .describedAs("count")
       .ofType(classOf[java.lang.Integer])
       .defaultsTo(1)
-    val newConsumerOpt = parser.accepts("new-consumer", "Use the new consumer implementation. This is the default, so " +
-      "this option is deprecated and will be removed in a future release.")
     val consumerConfigOpt = parser.accepts("consumer.config", "Consumer config properties file.")
       .withRequiredArg
       .describedAs("config file")
       .ofType(classOf[String])
-    val printMetricsOpt = parser.accepts("print-metrics", "Print out the metrics. This only applies to new consumer.")
+    val printMetricsOpt = parser.accepts("print-metrics", "Print out the metrics.")
     val showDetailedStatsOpt = parser.accepts("show-detailed-stats", "If set, stats are reported for each reporting " +
       "interval as configured by reporting-interval")
+    val recordFetchTimeoutOpt = parser.accepts("timeout", "The maximum allowed time in milliseconds between returned records.")
+      .withOptionalArg()
+      .describedAs("milliseconds")
+      .ofType(classOf[Long])
+      .defaultsTo(10000)
+
+    try
+      options = parser.parse(args: _*)
+    catch {
+      case e: OptionException =>
+        CommandLineUtils.printUsageAndDie(parser, e.getMessage)
+    }
+
+    if(options.has(numThreadsOpt) || options.has(numFetchersOpt))
+      println("WARNING: option [threads] and [num-fetch-threads] have been deprecated and will be ignored by the test")
 
-    val options = parser.parse(args: _*)
+    CommandLineUtils.printHelpAndExitIfNeeded(this, "This tool helps in performance test for the full zookeeper consumer")
 
     CommandLineUtils.checkRequiredArgs(parser, options, topicOpt, numMessagesOpt)
 
-    val useOldConsumer = options.has(zkConnectOpt)
     val printMetrics = options.has(printMetricsOpt)
 
     val props = if (options.has(consumerConfigOpt))
       Utils.loadProps(options.valueOf(consumerConfigOpt))
     else
       new Properties
-    if (!useOldConsumer) {
-      CommandLineUtils.checkRequiredArgs(parser, options, bootstrapServersOpt)
 
-      if (options.has(newConsumerOpt)) {
-        Console.err.println("The --new-consumer option is deprecated and will be removed in a future major release." +
-          "The new consumer is used by default if the --bootstrap-server option is provided.")
-      }
+    import org.apache.kafka.clients.consumer.ConsumerConfig
+
+    val brokerHostsAndPorts = options.valueOf(if (options.has(bootstrapServerOpt)) bootstrapServerOpt else brokerListOpt)
+    props.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, brokerHostsAndPorts)
+    props.put(ConsumerConfig.GROUP_ID_CONFIG, options.valueOf(groupIdOpt))
+    props.put(ConsumerConfig.RECEIVE_BUFFER_CONFIG, options.valueOf(socketBufferSizeOpt).toString)
+    props.put(ConsumerConfig.MAX_PARTITION_FETCH_BYTES_CONFIG, options.valueOf(fetchSizeOpt).toString)
+    props.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, if (options.has(resetBeginningOffsetOpt)) "latest" else "earliest")
+    props.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, classOf[ByteArrayDeserializer])
+    props.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, classOf[ByteArrayDeserializer])
+    props.put(ConsumerConfig.CHECK_CRCS_CONFIG, "false")
 
-      import org.apache.kafka.clients.consumer.ConsumerConfig
-      props.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, options.valueOf(bootstrapServersOpt))
-      props.put(ConsumerConfig.GROUP_ID_CONFIG, options.valueOf(groupIdOpt))
-      props.put(ConsumerConfig.RECEIVE_BUFFER_CONFIG, options.valueOf(socketBufferSizeOpt).toString)
-      props.put(ConsumerConfig.MAX_PARTITION_FETCH_BYTES_CONFIG, options.valueOf(fetchSizeOpt).toString)
-      props.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, if (options.has(resetBeginningOffsetOpt)) "latest" else "earliest")
-      props.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, classOf[ByteArrayDeserializer])
-      props.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, classOf[ByteArrayDeserializer])
-      props.put(ConsumerConfig.CHECK_CRCS_CONFIG, "false")
-    } else {
-      if (options.has(bootstrapServersOpt))
-        CommandLineUtils.printUsageAndDie(parser, s"Option $bootstrapServersOpt is not valid with $zkConnectOpt.")
-      else if (options.has(newConsumerOpt))
-        CommandLineUtils.printUsageAndDie(parser, s"Option $newConsumerOpt is not valid with $zkConnectOpt.")
-      CommandLineUtils.checkRequiredArgs(parser, options, zkConnectOpt, numMessagesOpt)
-      props.put("group.id", options.valueOf(groupIdOpt))
-      props.put("socket.receive.buffer.bytes", options.valueOf(socketBufferSizeOpt).toString)
-      props.put("fetch.message.max.bytes", options.valueOf(fetchSizeOpt).toString)
-      props.put("auto.offset.reset", if (options.has(resetBeginningOffsetOpt)) "largest" else "smallest")
-      props.put("zookeeper.connect", options.valueOf(zkConnectOpt))
-      props.put("consumer.timeout.ms", "1000")
-      props.put("num.consumer.fetchers", options.valueOf(numFetchersOpt).toString)
-    }
     val numThreads = options.valueOf(numThreadsOpt).intValue
     val topic = options.valueOf(topicOpt)
     val numMessages = options.valueOf(numMessagesOpt).longValue
@@ -356,54 +298,7 @@ object ConsumerPerformance extends LazyLogging {
     val showDetailedStats = options.has(showDetailedStatsOpt)
     val dateFormat = new SimpleDateFormat(options.valueOf(dateFormatOpt))
     val hideHeader = options.has(hideHeaderOpt)
+    val recordFetchTimeoutMs = options.valueOf(recordFetchTimeoutOpt).longValue()
   }
 
-  class ConsumerPerfThread(threadId: Int,
-                           name: String,
-                           stream: KafkaStream[Array[Byte], Array[Byte]],
-                           config: ConsumerPerfConfig,
-                           totalMessagesRead: AtomicLong,
-                           totalBytesRead: AtomicLong,
-                           consumerTimeout: AtomicBoolean)
-    extends Thread(name) {
-
-    override def run() {
-      var bytesRead = 0L
-      var messagesRead = 0L
-      val startMs = System.currentTimeMillis
-      var lastReportTime: Long = startMs
-      var lastBytesRead = 0L
-      var lastMessagesRead = 0L
-
-      try {
-        val iter = stream.iterator
-        while (iter.hasNext && messagesRead < config.numMessages) {
-          val messageAndMetadata = iter.next()
-          messagesRead += 1
-          bytesRead += messageAndMetadata.message.length
-          val currentTimeMillis = System.currentTimeMillis
-
-          if (currentTimeMillis - lastReportTime >= config.reportingInterval) {
-            if (config.showDetailedStats)
-              printOldConsumerProgress(threadId, bytesRead, lastBytesRead, messagesRead, lastMessagesRead, lastReportTime, currentTimeMillis, config.dateFormat)
-            lastReportTime = currentTimeMillis
-            lastMessagesRead = messagesRead
-            lastBytesRead = bytesRead
-          }
-        }
-      } catch {
-        case _: InterruptedException =>
-        case _: ClosedByInterruptException =>
-        case _: ConsumerTimeoutException =>
-          consumerTimeout.set(true)
-        case e: Throwable => e.printStackTrace()
-      }
-      totalMessagesRead.addAndGet(messagesRead)
-      totalBytesRead.addAndGet(bytesRead)
-      if (config.showDetailedStats)
-        printOldConsumerProgress(threadId, bytesRead, lastBytesRead, messagesRead, lastMessagesRead, startMs, System.currentTimeMillis, config.dateFormat)
-
-    }
-
-  }
 }
diff --git a/core/src/main/scala/kafka/tools/DumpLogSegments.scala b/core/src/main/scala/kafka/tools/DumpLogSegments.scala
index fe82dc2d72f36..b07bb5f4d0102 100755
--- a/core/src/main/scala/kafka/tools/DumpLogSegments.scala
+++ b/core/src/main/scala/kafka/tools/DumpLogSegments.scala
@@ -18,90 +18,35 @@
 package kafka.tools
 
 import java.io._
-import java.nio.ByteBuffer
 
-import joptsimple.OptionParser
-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 kafka.utils.Implicits._
 import org.apache.kafka.common.record._
 import org.apache.kafka.common.utils.Utils
 
+import scala.jdk.CollectionConverters._
 import scala.collection.mutable
 import scala.collection.mutable.ArrayBuffer
-import scala.collection.JavaConverters._
 
 object DumpLogSegments {
 
-  def main(args: Array[String]) {
-    val parser = new OptionParser(false)
-    val printOpt = parser.accepts("print-data-log", "if set, printing the messages content when dumping data logs. Automatically set if any decoder option is specified.")
-    val verifyOpt = parser.accepts("verify-index-only", "if set, just verify the index log without printing its content.")
-    val indexSanityOpt = parser.accepts("index-sanity-check", "if set, just checks the index sanity without printing its content. " +
-      "This is the same check that is executed on broker startup to determine if an index needs rebuilding or not.")
-    val filesOpt = parser.accepts("files", "REQUIRED: The comma separated list of data and index log files to be dumped.")
-                           .withRequiredArg
-                           .describedAs("file1, file2, ...")
-                           .ofType(classOf[String])
-    val maxMessageSizeOpt = parser.accepts("max-message-size", "Size of largest message.")
-                                  .withRequiredArg
-                                  .describedAs("size")
-                                  .ofType(classOf[java.lang.Integer])
-                                  .defaultsTo(5 * 1024 * 1024)
-    val deepIterationOpt = parser.accepts("deep-iteration", "if set, uses deep instead of shallow iteration.")
-    val valueDecoderOpt = parser.accepts("value-decoder-class", "if set, used to deserialize the messages. This class should implement kafka.serializer.Decoder trait. Custom jar should be available in kafka/libs directory.")
-                               .withOptionalArg()
-                               .ofType(classOf[java.lang.String])
-                               .defaultsTo("kafka.serializer.StringDecoder")
-    val keyDecoderOpt = parser.accepts("key-decoder-class", "if set, used to deserialize the keys. This class should implement kafka.serializer.Decoder trait. Custom jar should be available in kafka/libs directory.")
-                               .withOptionalArg()
-                               .ofType(classOf[java.lang.String])
-                               .defaultsTo("kafka.serializer.StringDecoder")
-    val offsetsOpt = parser.accepts("offsets-decoder", "if set, log data will be parsed as offset data from the " +
-      "__consumer_offsets topic.")
-    val transactionLogOpt = parser.accepts("transaction-log-decoder", "if set, log data will be parsed as " +
-      "transaction metadata from the __transaction_state topic.")
-
-    val helpOpt = parser.accepts("help", "Print usage information.")
-
-    val options = parser.parse(args : _*)
-
-    if(args.length == 0 || options.has(helpOpt))
-      CommandLineUtils.printUsageAndDie(parser, "Parse a log file and dump its contents to the console, useful for debugging a seemingly corrupt log segment.")
+  // visible for testing
+  private[tools] val RecordIndent = "|"
 
-    CommandLineUtils.checkRequiredArgs(parser, options, filesOpt)
+  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()
 
-    val printDataLog = options.has(printOpt) ||
-      options.has(offsetsOpt) ||
-      options.has(transactionLogOpt) ||
-      options.has(valueDecoderOpt) ||
-      options.has(keyDecoderOpt)
-    val verifyOnly = options.has(verifyOpt)
-    val indexSanityOnly = options.has(indexSanityOpt)
-
-    val files = options.valueOf(filesOpt).split(",")
-    val maxMessageSize = options.valueOf(maxMessageSizeOpt).intValue()
-    val isDeepIteration = options.has(deepIterationOpt)
-
-    val messageParser = if (options.has(offsetsOpt)) {
-      new OffsetsMessageParser
-    } else if (options.has(transactionLogOpt)) {
-      new TransactionLogMessageParser
-    } else {
-      val valueDecoder: Decoder[_] = CoreUtils.createObject[Decoder[_]](options.valueOf(valueDecoderOpt), new VerifiableProperties)
-      val keyDecoder: Decoder[_] = CoreUtils.createObject[Decoder[_]](options.valueOf(keyDecoderOpt), new VerifiableProperties)
-      new DecoderMessageParser(keyDecoder, valueDecoder)
-    }
-
-    val misMatchesForIndexFilesMap = new mutable.HashMap[String, List[(Long, Long)]]
+    val misMatchesForIndexFilesMap = mutable.Map[String, List[(Long, Long)]]()
     val timeIndexDumpErrors = new TimeIndexDumpErrors
-    val nonConsecutivePairsForLogFilesMap = new mutable.HashMap[String, List[(Long, Long)]]
+    val nonConsecutivePairsForLogFilesMap = mutable.Map[String, List[(Long, Long)]]()
 
-    for(arg <- files) {
+    for (arg <- opts.files) {
       val file = new File(arg)
       println(s"Dumping $file")
 
@@ -109,11 +54,12 @@ object DumpLogSegments {
       val suffix = filename.substring(filename.lastIndexOf("."))
       suffix match {
         case Log.LogFileSuffix =>
-          dumpLog(file, printDataLog, nonConsecutivePairsForLogFilesMap, isDeepIteration, maxMessageSize , messageParser)
+          dumpLog(file, opts.shouldPrintDataLog, nonConsecutivePairsForLogFilesMap, opts.isDeepIteration,
+            opts.maxMessageSize, opts.messageParser)
         case Log.IndexFileSuffix =>
-          dumpIndex(file, indexSanityOnly, verifyOnly, misMatchesForIndexFilesMap, maxMessageSize)
+          dumpIndex(file, opts.indexSanityOnly, opts.verifyOnly, misMatchesForIndexFilesMap, opts.maxMessageSize)
         case Log.TimeIndexFileSuffix =>
-          dumpTimeIndex(file, indexSanityOnly, verifyOnly, timeIndexDumpErrors, maxMessageSize)
+          dumpTimeIndex(file, opts.indexSanityOnly, opts.verifyOnly, timeIndexDumpErrors, opts.maxMessageSize)
         case Log.ProducerSnapshotFileSuffix =>
           dumpProducerIdSnapshot(file)
         case Log.TxnIndexFileSuffix =>
@@ -123,23 +69,19 @@ object DumpLogSegments {
       }
     }
 
-    misMatchesForIndexFilesMap.foreach {
-      case (fileName, listOfMismatches) => {
-        System.err.println("Mismatches in :" + fileName)
-        listOfMismatches.foreach(m => {
-          System.err.println("  Index offset: %d, log offset: %d".format(m._1, m._2))
-        })
+    misMatchesForIndexFilesMap.forKeyValue { (fileName, listOfMismatches) =>
+      System.err.println(s"Mismatches in :$fileName")
+      listOfMismatches.foreach { case (indexOffset, logOffset) =>
+        System.err.println(s"  Index offset: $indexOffset, log offset: $logOffset")
       }
     }
 
     timeIndexDumpErrors.printErrors()
 
-    nonConsecutivePairsForLogFilesMap.foreach {
-      case (fileName, listOfNonConsecutivePairs) => {
-        System.err.println("Non-consecutive offsets in :" + fileName)
-        listOfNonConsecutivePairs.foreach(m => {
-          System.err.println("  %d is followed by %d".format(m._1, m._2))
-        })
+    nonConsecutivePairsForLogFilesMap.forKeyValue { (fileName, listOfNonConsecutivePairs) =>
+      System.err.println(s"Non-consecutive offsets in $fileName")
+      listOfNonConsecutivePairs.foreach { case (first, second) =>
+        System.err.println(s"  $first is followed by $second")
       }
     }
   }
@@ -155,9 +97,13 @@ object DumpLogSegments {
   private def dumpProducerIdSnapshot(file: File): Unit = {
     try {
       ProducerStateManager.readSnapshot(file).foreach { entry =>
-        println(s"producerId: ${entry.producerId} producerEpoch: ${entry.producerEpoch} " +
-          s"coordinatorEpoch: ${entry.coordinatorEpoch} currentTxnFirstOffset: ${entry.currentTxnFirstOffset} " +
-          s"cachedMetadata: ${entry.batchMetadata}")
+        print(s"producerId: ${entry.producerId} producerEpoch: ${entry.producerEpoch} " +
+          s"coordinatorEpoch: ${entry.coordinatorEpoch} currentTxnFirstOffset: ${entry.currentTxnFirstOffset} ")
+        entry.batchMetadata.headOption.foreach { metadata =>
+          print(s"firstSequence: ${metadata.firstSeq} lastSequence: ${metadata.lastSeq} " +
+            s"lastOffset: ${metadata.lastOffset} offsetDelta: ${metadata.offsetDelta} timestamp: ${metadata.timestamp}")
+        }
+        println()
       }
     } catch {
       case e: CorruptSnapshotException =>
@@ -166,11 +112,12 @@ object DumpLogSegments {
   }
 
   /* print out the contents of the index */
-  private def dumpIndex(file: File,
-                        indexSanityOnly: Boolean,
-                        verifyOnly: Boolean,
-                        misMatchesForIndexFilesMap: mutable.HashMap[String, List[(Long, Long)]],
-                        maxMessageSize: Int) {
+  // Visible for testing
+  private[tools] def dumpIndex(file: File,
+                               indexSanityOnly: Boolean,
+                               verifyOnly: Boolean,
+                               misMatchesForIndexFilesMap: mutable.Map[String, List[(Long, Long)]],
+                               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)
@@ -178,33 +125,36 @@ object DumpLogSegments {
 
     //Check that index passes sanityCheck, this is the check that determines if indexes will be rebuilt on startup or not.
     if (indexSanityOnly) {
-      index.sanityCheck
+      index.sanityCheck()
       println(s"$file passed sanity check.")
       return
     }
 
-    for(i <- 0 until index.entries) {
+    for (i <- 0 until index.entries) {
       val entry = index.entry(i)
-      val slice = fileRecords.read(entry.position, maxMessageSize)
-      val firstRecord = slice.records.iterator.next()
-      if (firstRecord.offset != entry.offset + index.baseOffset) {
+
+      // since it is a sparse file, in the event of a crash there may be many zero entries, stop if we see one
+      if (entry.offset == index.baseOffset && i > 0)
+        return
+
+      val slice = fileRecords.slice(entry.position, maxMessageSize)
+      val firstBatchLastOffset = slice.batches.iterator.next().lastOffset
+      if (firstBatchLastOffset != entry.offset) {
         var misMatchesSeq = misMatchesForIndexFilesMap.getOrElse(file.getAbsolutePath, List[(Long, Long)]())
-        misMatchesSeq ::=(entry.offset + index.baseOffset, firstRecord.offset)
+        misMatchesSeq ::= (entry.offset, firstBatchLastOffset)
         misMatchesForIndexFilesMap.put(file.getAbsolutePath, misMatchesSeq)
       }
-      // since it is a sparse file, in the event of a crash there may be many zero entries, stop if we see one
-      if(entry.offset == 0 && i > 0)
-        return
       if (!verifyOnly)
-        println("offset: %d position: %d".format(entry.offset + index.baseOffset, entry.position))
+        println(s"offset: ${entry.offset} position: ${entry.position}")
     }
   }
 
-  private def dumpTimeIndex(file: File,
-                            indexSanityOnly: Boolean,
-                            verifyOnly: Boolean,
-                            timeIndexDumpErrors: TimeIndexDumpErrors,
-                            maxMessageSize: Int) {
+  // Visible for testing
+  private[tools] def dumpTimeIndex(file: File,
+                                   indexSanityOnly: Boolean,
+                                   verifyOnly: Boolean,
+                                   timeIndexDumpErrors: TimeIndexDumpErrors,
+                                   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,61 +162,68 @@ object DumpLogSegments {
     val index = new OffsetIndex(indexFile, baseOffset = startOffset, writable = false)
     val timeIndex = new TimeIndex(file, baseOffset = startOffset, writable = false)
 
-    //Check that index passes sanityCheck, this is the check that determines if indexes will be rebuilt on startup or not.
-    if (indexSanityOnly) {
-      timeIndex.sanityCheck
-      println(s"$file passed sanity check.")
-      return
-    }
+    try {
+      //Check that index passes sanityCheck, this is the check that determines if indexes will be rebuilt on startup or not.
+      if (indexSanityOnly) {
+        timeIndex.sanityCheck()
+        println(s"$file passed sanity check.")
+        return
+      }
 
-    var prevTimestamp = RecordBatch.NO_TIMESTAMP
-    for(i <- 0 until timeIndex.entries) {
-      val entry = timeIndex.entry(i)
-      val position = index.lookup(entry.offset + timeIndex.baseOffset).position
-      val partialFileRecords = fileRecords.read(position, Int.MaxValue)
-      val batches = partialFileRecords.batches.asScala
-      var maxTimestamp = RecordBatch.NO_TIMESTAMP
-      // We first find the message by offset then check if the timestamp is correct.
-      batches.find(_.lastOffset >= entry.offset + timeIndex.baseOffset) match {
-        case None =>
-          timeIndexDumpErrors.recordShallowOffsetNotFound(file, entry.offset + timeIndex.baseOffset,
-            -1.toLong)
-        case Some(batch) if batch.lastOffset != entry.offset + timeIndex.baseOffset =>
-          timeIndexDumpErrors.recordShallowOffsetNotFound(file, entry.offset + timeIndex.baseOffset, batch.lastOffset)
-        case Some(batch) =>
-          for (record <- batch.asScala)
-            maxTimestamp = math.max(maxTimestamp, record.timestamp)
-
-          if (maxTimestamp != entry.timestamp)
-            timeIndexDumpErrors.recordMismatchTimeIndex(file, entry.timestamp, maxTimestamp)
-
-          if (prevTimestamp >= entry.timestamp)
-            timeIndexDumpErrors.recordOutOfOrderIndexTimestamp(file, entry.timestamp, prevTimestamp)
-
-          // since it is a sparse file, in the event of a crash there may be many zero entries, stop if we see one
-          if (entry.offset == 0 && i > 0)
-            return
+      var prevTimestamp = RecordBatch.NO_TIMESTAMP
+      for (i <- 0 until timeIndex.entries) {
+        val entry = timeIndex.entry(i)
+
+        // since it is a sparse file, in the event of a crash there may be many zero entries, stop if we see one
+        if (entry.offset == timeIndex.baseOffset && i > 0)
+          return
+
+        val position = index.lookup(entry.offset).position
+        val partialFileRecords = fileRecords.slice(position, Int.MaxValue)
+        val batches = partialFileRecords.batches.asScala
+        var maxTimestamp = RecordBatch.NO_TIMESTAMP
+        // We first find the message by offset then check if the timestamp is correct.
+        batches.find(_.lastOffset >= entry.offset) match {
+          case None =>
+            timeIndexDumpErrors.recordShallowOffsetNotFound(file, entry.offset,
+              -1.toLong)
+          case Some(batch) if batch.lastOffset != entry.offset =>
+            timeIndexDumpErrors.recordShallowOffsetNotFound(file, entry.offset, batch.lastOffset)
+          case Some(batch) =>
+            for (record <- batch.asScala)
+              maxTimestamp = math.max(maxTimestamp, record.timestamp)
+
+            if (maxTimestamp != entry.timestamp)
+              timeIndexDumpErrors.recordMismatchTimeIndex(file, entry.timestamp, maxTimestamp)
+
+            if (prevTimestamp >= entry.timestamp)
+              timeIndexDumpErrors.recordOutOfOrderIndexTimestamp(file, entry.timestamp, prevTimestamp)
+        }
+        if (!verifyOnly)
+          println(s"timestamp: ${entry.timestamp} offset: ${entry.offset}")
+        prevTimestamp = entry.timestamp
       }
-      if (!verifyOnly)
-        println("timestamp: %s offset: %s".format(entry.timestamp, timeIndex.baseOffset + entry.offset))
-      prevTimestamp = entry.timestamp
+    } finally {
+      fileRecords.closeHandlers()
+      index.closeHandler()
+      timeIndex.closeHandler()
     }
   }
 
-  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)
@@ -274,193 +231,113 @@ 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)
-      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.encode(Map("metadata" -> groupId))
-      val valueString = Json.encode(Map(
-          "protocolType" -> protocolType,
-          "protocol" -> group.protocol,
-          "generationId" -> group.generationId,
-          "assignment" -> assignment))
-
-      (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.HashMap[String, List[(Long, Long)]],
+                      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 messageSet = FileRecords.open(file, false)
-    var validBytes = 0L
-    var lastOffset = -1L
-
-    for (batch <- messageSet.batches.asScala) {
-      if (isDeepIteration) {
-        for (record <- batch.asScala) {
-          if (lastOffset == -1)
+    val fileRecords = FileRecords.open(file, false)
+    try {
+      var validBytes = 0L
+      var lastOffset = -1L
+
+      for (batch <- fileRecords.batches.asScala) {
+        printBatchLevel(batch, validBytes)
+        if (isDeepIteration) {
+          for (record <- batch.asScala) {
+            if (lastOffset == -1)
+              lastOffset = record.offset
+            else if (record.offset != lastOffset + 1) {
+              var nonConsecutivePairsSeq = nonConsecutivePairsForLogFilesMap.getOrElse(file.getAbsolutePath, List[(Long, Long)]())
+              nonConsecutivePairsSeq ::= (lastOffset, record.offset)
+              nonConsecutivePairsForLogFilesMap.put(file.getAbsolutePath, nonConsecutivePairsSeq)
+            }
             lastOffset = record.offset
-          else if (record.offset != lastOffset + 1) {
-            var nonConsecutivePairsSeq = nonConsecutivePairsForLogFilesMap.getOrElse(file.getAbsolutePath, List[(Long, Long)]())
-            nonConsecutivePairsSeq ::= (lastOffset, record.offset)
-            nonConsecutivePairsForLogFilesMap.put(file.getAbsolutePath, nonConsecutivePairsSeq)
-          }
-          lastOffset = record.offset
-
-          print("offset: " + record.offset + " position: " + validBytes +
-            " " + batch.timestampType + ": " + record.timestamp + " isvalid: " + record.isValid +
-            " keysize: " + record.keySize + " valuesize: " + record.valueSize + " magic: " + batch.magic +
-            " compresscodec: " + batch.compressionType)
-
-          if (batch.magic >= RecordBatch.MAGIC_VALUE_V2) {
-            print(" producerId: " + batch.producerId + " producerEpoch: " + batch.producerEpoch + " sequence: " + record.sequence +
-              " isTransactional: " + batch.isTransactional +
-              " headerKeys: " + record.headers.map(_.key).mkString("[", ",", "]"))
-          } else {
-            print(" crc: " + record.checksumOrNull)
-          }
 
-          if (batch.isControlBatch) {
-            val controlTypeId = ControlRecordType.parseTypeId(record.key)
-            ControlRecordType.fromTypeId(controlTypeId) match {
-              case ControlRecordType.ABORT | ControlRecordType.COMMIT =>
-                val endTxnMarker = EndTransactionMarker.deserialize(record)
-                print(s" endTxnMarker: ${endTxnMarker.controlType} coordinatorEpoch: ${endTxnMarker.coordinatorEpoch}")
-              case controlType =>
-                print(s" controlType: $controlType($controlTypeId)")
+            print(s"$RecordIndent offset: ${record.offset} isValid: ${record.isValid} crc: ${record.checksumOrNull}" +
+                s" keySize: ${record.keySize} valueSize: ${record.valueSize} ${batch.timestampType}: ${record.timestamp}" +
+                s" baseOffset: ${batch.baseOffset} lastOffset: ${batch.lastOffset} baseSequence: ${batch.baseSequence}" +
+                s" lastSequence: ${batch.lastSequence} producerEpoch: ${batch.producerEpoch} partitionLeaderEpoch: ${batch.partitionLeaderEpoch}" +
+                s" batchSize: ${batch.sizeInBytes} magic: ${batch.magic} compressType: ${batch.compressionType} position: ${validBytes}")
+
+
+            if (batch.magic >= RecordBatch.MAGIC_VALUE_V2) {
+              print(" sequence: " + record.sequence + " headerKeys: " + record.headers.map(_.key).mkString("[", ",", "]"))
+            } else {
+              print(s" crc: ${record.checksumOrNull} isvalid: ${record.isValid}")
+            }
+
+            if (batch.isControlBatch) {
+              val controlTypeId = ControlRecordType.parseTypeId(record.key)
+              ControlRecordType.fromTypeId(controlTypeId) match {
+                case ControlRecordType.ABORT | ControlRecordType.COMMIT =>
+                  val endTxnMarker = EndTransactionMarker.deserialize(record)
+                  print(s" endTxnMarker: ${endTxnMarker.controlType} coordinatorEpoch: ${endTxnMarker.coordinatorEpoch}")
+                case controlType =>
+                  print(s" controlType: $controlType($controlTypeId)")
+              }
+            } else if (printContents) {
+              val (key, payload) = parser.parse(record)
+              key.foreach(key => print(s" key: $key"))
+              payload.foreach(payload => print(s" payload: $payload"))
             }
-          } else if (printContents) {
-            val (key, payload) = parser.parse(record)
-            key.foreach(key => print(s" key: $key"))
-            payload.foreach(payload => print(s" payload: $payload"))
+            println()
           }
-          println()
         }
-      } else {
-        if (batch.magic >= RecordBatch.MAGIC_VALUE_V2)
-          print("baseOffset: " + batch.baseOffset + " lastOffset: " + batch.lastOffset +
-            " baseSequence: " + batch.baseSequence + " lastSequence: " + batch.lastSequence +
-            " producerId: " + batch.producerId + " producerEpoch: " + batch.producerEpoch +
-            " partitionLeaderEpoch: " + batch.partitionLeaderEpoch + " isTransactional: " + batch.isTransactional)
-        else
-          print("offset: " + batch.lastOffset)
-
-        println(" position: " + validBytes + " " + batch.timestampType + ": " + batch.maxTimestamp +
-          " isvalid: " + batch.isValid +
-          " size: " + batch.sizeInBytes + " magic: " + batch.magic +
-          " compresscodec: " + batch.compressionType + " crc: " + batch.checksum)
+        validBytes += batch.sizeInBytes
       }
-      validBytes += batch.sizeInBytes
-    }
-    val trailingBytes = messageSet.sizeInBytes - validBytes
-    if(trailingBytes > 0)
-      println("Found %d invalid bytes at the end of %s".format(trailingBytes, file.getName))
+      val trailingBytes = fileRecords.sizeInBytes - validBytes
+      if (trailingBytes > 0)
+        println(s"Found $trailingBytes invalid bytes at the end of ${file.getName}")
+    } finally fileRecords.closeHandlers()
+  }
+
+  private def printBatchLevel(batch: FileLogInputStream.FileChannelRecordBatch, accumulativeBytes: Long): Unit = {
+    if (batch.magic >= RecordBatch.MAGIC_VALUE_V2)
+      print("baseOffset: " + batch.baseOffset + " lastOffset: " + batch.lastOffset + " count: " + batch.countOrNull +
+        " baseSequence: " + batch.baseSequence + " lastSequence: " + batch.lastSequence +
+        " producerId: " + batch.producerId + " producerEpoch: " + batch.producerEpoch +
+        " partitionLeaderEpoch: " + batch.partitionLeaderEpoch + " isTransactional: " + batch.isTransactional +
+        " isControl: " + batch.isControlBatch)
+    else
+      print("offset: " + batch.lastOffset)
+
+    println(" position: " + accumulativeBytes + " " + batch.timestampType + ": " + batch.maxTimestamp +
+      " size: " + batch.sizeInBytes + " magic: " + batch.magic +
+      " compresscodec: " + batch.compressionType + " crc: " + batch.checksum + " isvalid: " + batch.isValid)
   }
 
   class TimeIndexDumpErrors {
-    val misMatchesForTimeIndexFilesMap = new mutable.HashMap[String, ArrayBuffer[(Long, Long)]]
-    val outOfOrderTimestamp = new mutable.HashMap[String, ArrayBuffer[(Long, Long)]]
-    val shallowOffsetNotFound = new mutable.HashMap[String, ArrayBuffer[(Long, Long)]]
+    val misMatchesForTimeIndexFilesMap = mutable.Map[String, ArrayBuffer[(Long, Long)]]()
+    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)
@@ -488,4 +365,71 @@ 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.")
+    val indexSanityOpt = parser.accepts("index-sanity-check", "if set, just checks the index sanity without printing its content. " +
+      "This is the same check that is executed on broker startup to determine if an index needs rebuilding or not.")
+    val filesOpt = parser.accepts("files", "REQUIRED: The comma separated list of data and index log files to be dumped.")
+      .withRequiredArg
+      .describedAs("file1, file2, ...")
+      .ofType(classOf[String])
+    val maxMessageSizeOpt = parser.accepts("max-message-size", "Size of largest message.")
+      .withRequiredArg
+      .describedAs("size")
+      .ofType(classOf[java.lang.Integer])
+      .defaultsTo(5 * 1024 * 1024)
+    val deepIterationOpt = parser.accepts("deep-iteration", "if set, uses deep instead of shallow iteration. Automatically set if print-data-log is enabled.")
+    val valueDecoderOpt = parser.accepts("value-decoder-class", "if set, used to deserialize the messages. This class should implement kafka.serializer.Decoder trait. Custom jar should be available in kafka/libs directory.")
+      .withOptionalArg()
+      .ofType(classOf[java.lang.String])
+      .defaultsTo("kafka.serializer.StringDecoder")
+    val keyDecoderOpt = parser.accepts("key-decoder-class", "if set, used to deserialize the keys. This class should implement kafka.serializer.Decoder trait. Custom jar should be available in kafka/libs directory.")
+      .withOptionalArg()
+      .ofType(classOf[java.lang.String])
+      .defaultsTo("kafka.serializer.StringDecoder")
+    val offsetsOpt = parser.accepts("offsets-decoder", "if set, log data will be parsed as offset data from the " +
+      "__consumer_offsets topic.")
+    val transactionLogOpt = parser.accepts("transaction-log-decoder", "if set, log data will be parsed as " +
+      "transaction metadata from the __transaction_state topic.")
+    options = parser.parse(args : _*)
+
+    def messageParser: MessageParser[_, _] =
+      if (options.has(offsetsOpt)) {
+        new OffsetsMessageParser
+      } else if (options.has(transactionLogOpt)) {
+        new TransactionLogMessageParser
+      } else {
+        val valueDecoder: Decoder[_] = CoreUtils.createObject[Decoder[_]](options.valueOf(valueDecoderOpt), new VerifiableProperties)
+        val keyDecoder: Decoder[_] = CoreUtils.createObject[Decoder[_]](options.valueOf(keyDecoderOpt), new VerifiableProperties)
+        new DecoderMessageParser(keyDecoder, valueDecoder)
+      }
+
+    lazy val shouldPrintDataLog: Boolean = options.has(printOpt) ||
+      options.has(offsetsOpt) ||
+      options.has(transactionLogOpt) ||
+      options.has(valueDecoderOpt) ||
+      options.has(keyDecoderOpt)
+
+    lazy val isDeepIteration: Boolean = options.has(deepIterationOpt) || shouldPrintDataLog
+    lazy val verifyOnly: Boolean = options.has(verifyOpt)
+    lazy val indexSanityOnly: Boolean = options.has(indexSanityOpt)
+    lazy val files = options.valueOf(filesOpt).split(",")
+    lazy val maxMessageSize = options.valueOf(maxMessageSizeOpt).intValue()
+
+    def checkArgs(): Unit = CommandLineUtils.checkRequiredArgs(parser, options, filesOpt)
+
+  }
 }
diff --git a/core/src/main/scala/kafka/tools/EndToEndLatency.scala b/core/src/main/scala/kafka/tools/EndToEndLatency.scala
index 3beaf827f5959..8a0e670570b1b 100755
--- a/core/src/main/scala/kafka/tools/EndToEndLatency.scala
+++ b/core/src/main/scala/kafka/tools/EndToEndLatency.scala
@@ -18,14 +18,18 @@
 package kafka.tools
 
 import java.nio.charset.StandardCharsets
+import java.time.Duration
 import java.util.{Arrays, Collections, Properties}
 
 import kafka.utils.Exit
+import org.apache.kafka.clients.admin.{Admin, NewTopic}
+import org.apache.kafka.clients.CommonClientConfigs
 import org.apache.kafka.clients.consumer.{ConsumerConfig, KafkaConsumer}
 import org.apache.kafka.clients.producer._
+import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.utils.Utils
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 import scala.util.Random
 
 
@@ -42,8 +46,10 @@ import scala.util.Random
 
 object EndToEndLatency {
   private val timeout: Long = 60000
+  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)
@@ -59,22 +65,22 @@ object EndToEndLatency {
     if (!List("1", "all").contains(producerAcks))
       throw new IllegalArgumentException("Latency testing requires synchronous acknowledgement. Please use 1 or all")
 
-    def loadProps: Properties = propsFile.map(Utils.loadProps).getOrElse(new Properties())
+    def loadPropsWithBootstrapServers: Properties = {
+      val props = propsFile.map(Utils.loadProps).getOrElse(new Properties())
+      props.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, brokerList)
+      props
+    }
 
-    val consumerProps = loadProps
-    consumerProps.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList)
+    val consumerProps = loadPropsWithBootstrapServers
     consumerProps.put(ConsumerConfig.GROUP_ID_CONFIG, "test-group-" + System.currentTimeMillis())
     consumerProps.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false")
     consumerProps.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "latest")
     consumerProps.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.ByteArrayDeserializer")
     consumerProps.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.ByteArrayDeserializer")
     consumerProps.put(ConsumerConfig.FETCH_MAX_WAIT_MS_CONFIG, "0") //ensure we have no temporal batching
-
     val consumer = new KafkaConsumer[Array[Byte], Array[Byte]](consumerProps)
-    consumer.subscribe(Collections.singletonList(topic))
 
-    val producerProps = loadProps
-    producerProps.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList)
+    val producerProps = loadPropsWithBootstrapServers
     producerProps.put(ProducerConfig.LINGER_MS_CONFIG, "0") //ensure writes are synchronous
     producerProps.put(ProducerConfig.MAX_BLOCK_MS_CONFIG, Long.MaxValue.toString)
     producerProps.put(ProducerConfig.ACKS_CONFIG, producerAcks.toString)
@@ -82,16 +88,28 @@ 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()
     }
 
-    //Ensure we are at latest offset. seekToEnd evaluates lazily, that is to say actually performs the seek only when
-    //a poll() or position() request is issued. Hence we need to poll after we seek to ensure we see our first write.
-    consumer.seekToEnd(Collections.emptyList())
-    consumer.poll(0)
+    // create topic if it does not exist
+    if (!consumer.listTopics().containsKey(topic)) {
+      try {
+        createTopic(topic, loadPropsWithBootstrapServers)
+      } catch {
+        case t: Throwable =>
+          finalise()
+          throw new RuntimeException(s"Failed to create topic $topic", t)
+      }
+    }
+
+    val topicPartitions = consumer.partitionsFor(topic).asScala
+      .map(p => new TopicPartition(p.topic(), p.partition())).asJava
+    consumer.assign(topicPartitions)
+    consumer.seekToEnd(topicPartitions)
+    consumer.assignment.forEach(consumer.position(_))
 
     var totalTime = 0.0
     val latencies = new Array[Long](numMessages)
@@ -103,7 +121,7 @@ object EndToEndLatency {
 
       //Send message (of random bytes) synchronously then immediately poll for it
       producer.send(new ProducerRecord[Array[Byte], Array[Byte]](topic, message)).get()
-      val recordIter = consumer.poll(timeout).iterator
+      val recordIter = consumer.poll(Duration.ofMillis(timeout)).iterator
 
       val elapsed = System.nanoTime - begin
 
@@ -129,7 +147,7 @@ object EndToEndLatency {
 
       //Report progress
       if (i % 1000 == 0)
-        println(i + "\t" + elapsed / 1000.0 / 1000.0)
+        println(i.toString + "\t" + elapsed / 1000.0 / 1000.0)
       totalTime += elapsed
       latencies(i) = elapsed / 1000 / 1000
     }
@@ -148,4 +166,14 @@ object EndToEndLatency {
   def randomBytesOfLen(random: Random, len: Int): Array[Byte] = {
     Array.fill(len)((random.nextInt(26) + 65).toByte)
   }
+
+  def createTopic(topic: String, props: Properties): Unit = {
+    println("Topic \"%s\" does not exist. Will create topic with %d partition(s) and replication factor = %d"
+              .format(topic, defaultNumPartitions, defaultReplicationFactor))
+
+    val adminClient = Admin.create(props)
+    val newTopic = new NewTopic(topic, defaultNumPartitions, defaultReplicationFactor)
+    try adminClient.createTopics(Collections.singleton(newTopic)).all().get()
+    finally Utils.closeQuietly(adminClient, "AdminClient")
+  }
 }
diff --git a/core/src/main/scala/kafka/tools/ExportZkOffsets.scala b/core/src/main/scala/kafka/tools/ExportZkOffsets.scala
deleted file mode 100644
index d8ce9b068acb7..0000000000000
--- a/core/src/main/scala/kafka/tools/ExportZkOffsets.scala
+++ /dev/null
@@ -1,132 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package kafka.tools
-
-import java.io.{FileOutputStream, OutputStreamWriter}
-import java.nio.charset.StandardCharsets
-
-import joptsimple._
-import kafka.utils.{CommandLineUtils, Exit, Logging, ZKGroupTopicDirs, ZkUtils}
-import org.apache.kafka.common.security.JaasUtils
-
-import scala.collection.JavaConverters._
-
-
-/**
- *  A utility that retrieves the offset of broker partitions in ZK and
- *  prints to an output file in the following format:
- *
- *  /consumers/group1/offsets/topic1/1-0:286894308
- *  /consumers/group1/offsets/topic1/2-0:284803985
- *
- *  This utility expects 3 arguments:
- *  1. Zk host:port string
- *  2. group name (all groups implied if omitted)
- *  3. output filename
- *
- *  To print debug message, add the following line to log4j.properties:
- *  log4j.logger.kafka.tools.ExportZkOffsets$=DEBUG
- *  (for eclipse debugging, copy log4j.properties to the binary directory in "core" such as core/bin)
- */
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.11.0.0")
-object ExportZkOffsets extends Logging {
-
-  def main(args: Array[String]) {
-    val parser = new OptionParser(false)
-    warn("WARNING: ExportZkOffsets is deprecated and will be dropped in a future release following 0.11.0.0.")
-
-    val zkConnectOpt = parser.accepts("zkconnect", "ZooKeeper connect string.")
-                            .withRequiredArg()
-                            .defaultsTo("localhost:2181")
-                            .ofType(classOf[String])
-    val groupOpt = parser.accepts("group", "Consumer group.")
-                            .withRequiredArg()
-                            .ofType(classOf[String])
-    val outFileOpt = parser.accepts("output-file", "Output file")
-                            .withRequiredArg()
-                            .ofType(classOf[String])
-    parser.accepts("help", "Print this message.")
-
-    if(args.length == 0)
-      CommandLineUtils.printUsageAndDie(parser, "Export consumer offsets to an output file.")
-
-    val options = parser.parse(args : _*)
-
-    if (options.has("help")) {
-       parser.printHelpOn(System.out)
-       Exit.exit(0)
-    }
-
-    CommandLineUtils.checkRequiredArgs(parser, options, zkConnectOpt, outFileOpt)
-
-    val zkConnect  = options.valueOf(zkConnectOpt)
-    val groups     = options.valuesOf(groupOpt)
-    val outfile    = options.valueOf(outFileOpt)
-
-    var zkUtils   : ZkUtils    = null
-    val fileWriter : OutputStreamWriter =
-        new OutputStreamWriter(new FileOutputStream(outfile), StandardCharsets.UTF_8)
-
-    try {
-      zkUtils = ZkUtils(zkConnect,
-                        30000,
-                        30000,
-                        JaasUtils.isZkSecurityEnabled())
-
-      var consumerGroups: Seq[String] = null
-
-      if (groups.size == 0) {
-        consumerGroups = zkUtils.getChildren(ZkUtils.ConsumersPath).toList
-      }
-      else {
-        consumerGroups = groups.asScala
-      }
-
-      for (consumerGrp <- consumerGroups) {
-        val topicsList = getTopicsList(zkUtils, consumerGrp)
-
-        for (topic <- topicsList) {
-          val bidPidList = getBrokeridPartition(zkUtils, consumerGrp, topic)
-
-          for (bidPid <- bidPidList) {
-            val zkGrpTpDir = new ZKGroupTopicDirs(consumerGrp,topic)
-            val offsetPath = zkGrpTpDir.consumerOffsetDir + "/" + bidPid
-            zkUtils.readDataMaybeNull(offsetPath)._1 match {
-              case Some(offsetVal) =>
-                fileWriter.write(offsetPath + ":" + offsetVal + "\n")
-                debug(offsetPath + " => " + offsetVal)
-              case None =>
-                error("Could not retrieve offset value from " + offsetPath)
-            }
-          }
-        }
-      }
-    }
-    finally {
-      fileWriter.flush()
-      fileWriter.close()
-    }
-  }
-
-  private def getBrokeridPartition(zkUtils: ZkUtils, consumerGroup: String, topic: String): List[String] =
-    zkUtils.getChildrenParentMayNotExist("/consumers/%s/offsets/%s".format(consumerGroup, topic)).toList
-
-  private def getTopicsList(zkUtils: ZkUtils, consumerGroup: String): List[String] =
-    zkUtils.getChildren("/consumers/%s/offsets".format(consumerGroup)).toList
-
-}
diff --git a/core/src/main/scala/kafka/tools/GetOffsetShell.scala b/core/src/main/scala/kafka/tools/GetOffsetShell.scala
index 4104dedb9e09a..76b16f57ec8be 100644
--- a/core/src/main/scala/kafka/tools/GetOffsetShell.scala
+++ b/core/src/main/scala/kafka/tools/GetOffsetShell.scala
@@ -18,13 +18,17 @@
  */
 package kafka.tools
 
-import kafka.consumer._
+import java.util.Properties
+
 import joptsimple._
-import kafka.api.{OffsetRequest, PartitionOffsetRequestInfo}
-import kafka.common.TopicAndPartition
-import kafka.client.ClientUtils
 import kafka.utils.{CommandLineUtils, Exit, ToolsUtils}
+import org.apache.kafka.clients.consumer.{ConsumerConfig, KafkaConsumer}
+import org.apache.kafka.common.{PartitionInfo, TopicPartition}
+import org.apache.kafka.common.requests.ListOffsetRequest
+import org.apache.kafka.common.serialization.ByteArrayDeserializer
 
+import scala.jdk.CollectionConverters._
+import scala.collection.Seq
 
 object GetOffsetShell {
 
@@ -43,24 +47,24 @@ object GetOffsetShell {
                            .describedAs("partition ids")
                            .ofType(classOf[String])
                            .defaultsTo("")
-    val timeOpt = parser.accepts("time", "timestamp of the offsets before that")
+    val timeOpt = parser.accepts("time", "timestamp of the offsets before that. [Note: No offset is returned, if the timestamp greater than recently commited record timestamp is given.]")
                            .withRequiredArg
                            .describedAs("timestamp/-1(latest)/-2(earliest)")
                            .ofType(classOf[java.lang.Long])
-                           .defaultsTo(-1)
-    val nOffsetsOpt = parser.accepts("offsets", "number of offsets returned")
+                           .defaultsTo(-1L)
+    parser.accepts("offsets", "DEPRECATED AND IGNORED: number of offsets returned")
                            .withRequiredArg
                            .describedAs("count")
                            .ofType(classOf[java.lang.Integer])
                            .defaultsTo(1)
-    val maxWaitMsOpt = parser.accepts("max-wait-ms", "The max amount of time each fetch request waits.")
+    parser.accepts("max-wait-ms", "DEPRECATED AND IGNORED: The max amount of time each fetch request waits.")
                            .withRequiredArg
                            .describedAs("ms")
                            .ofType(classOf[java.lang.Integer])
                            .defaultsTo(1000)
-                           
-   if(args.length == 0)
-      CommandLineUtils.printUsageAndDie(parser, "An interactive shell for getting consumer offsets.")
+
+   if (args.length == 0)
+      CommandLineUtils.printUsageAndDie(parser, "An interactive shell for getting topic offsets.")
 
     val options = parser.parse(args : _*)
 
@@ -69,41 +73,83 @@ object GetOffsetShell {
     val clientId = "GetOffsetShell"
     val brokerList = options.valueOf(brokerListOpt)
     ToolsUtils.validatePortOrDie(parser, brokerList)
-    val metadataTargetBrokers = ClientUtils.parseBrokerList(brokerList)
     val topic = options.valueOf(topicOpt)
-    val partitionList = options.valueOf(partitionOpt)
-    val time = options.valueOf(timeOpt).longValue
-    val nOffsets = options.valueOf(nOffsetsOpt).intValue
-    val maxWaitMs = options.valueOf(maxWaitMsOpt).intValue()
-
-    val topicsMetadata = ClientUtils.fetchTopicMetadata(Set(topic), metadataTargetBrokers, clientId, maxWaitMs).topicsMetadata
-    if(topicsMetadata.size != 1 || !topicsMetadata.head.topic.equals(topic)) {
-      System.err.println(("Error: no valid topic metadata for topic: %s, " + " probably the topic does not exist, run ").format(topic) +
-        "kafka-list-topic.sh to verify")
-      Exit.exit(1)
-    }
-    val partitions =
-      if(partitionList == "") {
-        topicsMetadata.head.partitionsMetadata.map(_.partitionId)
-      } else {
-        partitionList.split(",").map(_.toInt).toSeq
-      }
-    partitions.foreach { partitionId =>
-      val partitionMetadataOpt = topicsMetadata.head.partitionsMetadata.find(_.partitionId == partitionId)
-      partitionMetadataOpt match {
-        case Some(metadata) =>
-          metadata.leader match {
-            case Some(leader) =>
-              val consumer = new SimpleConsumer(leader.host, leader.port, 10000, 100000, clientId)
-              val topicAndPartition = TopicAndPartition(topic, partitionId)
-              val request = OffsetRequest(Map(topicAndPartition -> PartitionOffsetRequestInfo(time, nOffsets)))
-              val offsets = consumer.getOffsetsBefore(request).partitionErrorAndOffsets(topicAndPartition).offsets
-
-              println("%s:%d:%s".format(topic, partitionId, offsets.mkString(",")))
-            case None => System.err.println("Error: partition %d does not have a leader. Skip getting offsets".format(partitionId))
+    val partitionIdsRequested: Set[Int] = {
+      val partitionsString = options.valueOf(partitionOpt)
+      if (partitionsString.isEmpty)
+        Set.empty
+      else
+        partitionsString.split(",").map { partitionString =>
+          try partitionString.toInt
+          catch {
+            case _: NumberFormatException =>
+              System.err.println(s"--partitions expects a comma separated list of numeric partition ids, but received: $partitionsString")
+              Exit.exit(1)
           }
-        case None => System.err.println("Error: partition %d does not exist".format(partitionId))
+        }.toSet
+    }
+    val listOffsetsTimestamp = options.valueOf(timeOpt).longValue
+
+    val config = new Properties
+    config.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList)
+    config.setProperty(ConsumerConfig.CLIENT_ID_CONFIG, clientId)
+    val consumer = new KafkaConsumer(config, new ByteArrayDeserializer, new ByteArrayDeserializer)
+
+    val partitionInfos = listPartitionInfos(consumer, topic, partitionIdsRequested) match {
+      case None =>
+        System.err.println(s"Topic $topic does not exist")
+        Exit.exit(1)
+      case Some(p) if p.isEmpty =>
+        if (partitionIdsRequested.isEmpty)
+          System.err.println(s"Topic $topic has 0 partitions")
+        else
+          System.err.println(s"Topic $topic does not have any of the requested partitions ${partitionIdsRequested.mkString(",")}")
+        Exit.exit(1)
+      case Some(p) => p
+    }
+
+    if (partitionIdsRequested.nonEmpty) {
+      (partitionIdsRequested -- partitionInfos.map(_.partition)).foreach { partitionId =>
+        System.err.println(s"Error: partition $partitionId does not exist")
       }
     }
+
+    val topicPartitions = partitionInfos.sortBy(_.partition).flatMap { p =>
+      if (p.leader == null) {
+        System.err.println(s"Error: partition ${p.partition} does not have a leader. Skip getting offsets")
+        None
+      } else
+        Some(new TopicPartition(p.topic, p.partition))
+    }
+
+    /* Note that the value of the map can be null */
+    val partitionOffsets: collection.Map[TopicPartition, java.lang.Long] = listOffsetsTimestamp match {
+      case ListOffsetRequest.EARLIEST_TIMESTAMP => consumer.beginningOffsets(topicPartitions.asJava).asScala
+      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.map { case (k, x) =>
+          if (x == null) (k, null) else (k, x.offset: java.lang.Long)
+        }
+    }
+
+    partitionOffsets.toArray.sortBy { case (tp, _) => tp.partition }.foreach { case (tp, offset) =>
+      println(s"$topic:${tp.partition}:${Option(offset).getOrElse("")}")
+    }
+
+  }
+
+  /**
+   * Return the partition infos for `topic`. If the topic does not exist, `None` is returned.
+   */
+  private def listPartitionInfos(consumer: KafkaConsumer[_, _], topic: String, partitionIds: Set[Int]): Option[Seq[PartitionInfo]] = {
+    val partitionInfos = consumer.listTopics.asScala.filter { case (k, _) => k == topic }.values.flatMap(_.asScala).toBuffer
+    if (partitionInfos.isEmpty)
+      None
+    else if (partitionIds.isEmpty)
+      Some(partitionInfos)
+    else
+      Some(partitionInfos.filter(p => partitionIds.contains(p.partition)))
   }
+
 }
diff --git a/core/src/main/scala/kafka/tools/ImportZkOffsets.scala b/core/src/main/scala/kafka/tools/ImportZkOffsets.scala
deleted file mode 100644
index c345f94dff4ba..0000000000000
--- a/core/src/main/scala/kafka/tools/ImportZkOffsets.scala
+++ /dev/null
@@ -1,110 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- * 
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package kafka.tools
-
-import java.io.{BufferedReader, FileInputStream, InputStreamReader}
-import java.nio.charset.StandardCharsets
-
-import joptsimple._
-import kafka.utils.{CommandLineUtils, Exit, Logging, ZkUtils}
-import org.apache.kafka.common.security.JaasUtils
-
-
-/**
- *  A utility that updates the offset of broker partitions in ZK.
- *  
- *  This utility expects 2 input files as arguments:
- *  1. consumer properties file
- *  2. a file contains partition offsets data such as:
- *     (This output data file can be obtained by running kafka.tools.ExportZkOffsets)
- *  
- *     /consumers/group1/offsets/topic1/3-0:285038193
- *     /consumers/group1/offsets/topic1/1-0:286894308
- *     
- *  To print debug message, add the following line to log4j.properties:
- *  log4j.logger.kafka.tools.ImportZkOffsets$=DEBUG
- *  (for eclipse debugging, copy log4j.properties to the binary directory in "core" such as core/bin)
- */
-object ImportZkOffsets extends Logging {
-
-  def main(args: Array[String]) {
-    val parser = new OptionParser(false)
-    
-    val zkConnectOpt = parser.accepts("zkconnect", "ZooKeeper connect string.")
-                            .withRequiredArg()
-                            .defaultsTo("localhost:2181")
-                            .ofType(classOf[String])
-    val inFileOpt = parser.accepts("input-file", "Input file")
-                            .withRequiredArg()
-                            .ofType(classOf[String])
-    parser.accepts("help", "Print this message.")
-    
-    if(args.length == 0)
-      CommandLineUtils.printUsageAndDie(parser, "Import offsets to zookeeper from files.")
-            
-    val options = parser.parse(args : _*)
-    
-    if (options.has("help")) {
-       parser.printHelpOn(System.out)
-       Exit.exit(0)
-    }
-    
-    CommandLineUtils.checkRequiredArgs(parser, options, inFileOpt)
-    
-    val zkConnect           = options.valueOf(zkConnectOpt)
-    val partitionOffsetFile = options.valueOf(inFileOpt)
-
-    val zkUtils = ZkUtils(zkConnect, 30000, 30000, JaasUtils.isZkSecurityEnabled())
-    val partitionOffsets: Map[String,String] = getPartitionOffsetsFromFile(partitionOffsetFile)
-
-    updateZkOffsets(zkUtils, partitionOffsets)
-  }
-
-  private def getPartitionOffsetsFromFile(filename: String):Map[String,String] = {
-    val br = new BufferedReader(new InputStreamReader(new FileInputStream(filename), StandardCharsets.UTF_8))
-    try {
-      var partOffsetsMap: Map[String,String] = Map()
-
-      var s: String = br.readLine()
-      while ( s != null && s.length() >= 1) {
-        val tokens = s.split(":")
-
-        partOffsetsMap += tokens(0) -> tokens(1)
-        debug("adding node path [" + s + "]")
-
-        s = br.readLine()
-      }
-
-      partOffsetsMap
-    } finally {
-      br.close()
-    }
-  }
-  
-  private def updateZkOffsets(zkUtils: ZkUtils, partitionOffsets: Map[String,String]): Unit = {
-    for ((partition, offset) <- partitionOffsets) {
-      debug("updating [" + partition + "] with offset [" + offset + "]")
-      
-      try {
-        zkUtils.updatePersistentPath(partition, offset.toString)
-      } catch {
-        case e: Throwable => e.printStackTrace()
-      }
-    }
-  }
-}
diff --git a/core/src/main/scala/kafka/tools/JmxTool.scala b/core/src/main/scala/kafka/tools/JmxTool.scala
index 4a6a348d9a66a..f7ace833728d7 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.jdk.CollectionConverters._
 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 =
@@ -50,17 +51,22 @@ object JmxTool extends Logging {
         .describedAs("name")
         .ofType(classOf[String])
     val attributesOpt =
-      parser.accepts("attributes", "The whitelist of attributes to query. This is a comma-separated list. If no " +
+      parser.accepts("attributes", "The list of attributes to include in the query. This is a comma-separated list. If no " +
         "attributes are specified all objects will be queried.")
         .withRequiredArg
         .describedAs("name")
         .ofType(classOf[String])
-    val reportingIntervalOpt = parser.accepts("reporting-interval", "Interval in MS with which to poll jmx stats.")
+    val reportingIntervalOpt = parser.accepts("reporting-interval", "Interval in MS with which to poll jmx stats; default value is 2 seconds. " +
+      "Value of -1 equivalent to setting one-time to true")
       .withRequiredArg
       .describedAs("ms")
       .ofType(classOf[java.lang.Integer])
       .defaultsTo(2000)
-    val helpOpt = parser.accepts("help", "Print usage information.")
+    val oneTimeOpt = parser.accepts("one-time", "Flag to indicate run once only.")
+      .withRequiredArg
+      .describedAs("one-time")
+      .ofType(classOf[java.lang.Boolean])
+      .defaultsTo(false)
     val dateFormatOpt = parser.accepts("date-format", "The date format to use for formatting the time field. " +
       "See java.text.SimpleDateFormat for options.")
       .withRequiredArg
@@ -72,8 +78,25 @@ object JmxTool extends Logging {
         .describedAs("service-url")
         .ofType(classOf[String])
         .defaultsTo("service:jmx:rmi:///jndi/rmi://:9999/jmxrmi")
+    val reportFormatOpt = parser.accepts("report-format", "output format name: either 'original', 'properties', 'csv', 'tsv' ")
+      .withRequiredArg
+      .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.")
+
 
     if(args.length == 0)
       CommandLineUtils.printUsageAndDie(parser, "Dump JMX values to standard output.")
@@ -87,12 +110,19 @@ object JmxTool extends Logging {
 
     val url = new JMXServiceURL(options.valueOf(jmxServiceUrlOpt))
     val interval = options.valueOf(reportingIntervalOpt).intValue
-    val attributesWhitelistExists = options.has(attributesOpt)
-    val attributesWhitelist = if(attributesWhitelistExists) Some(options.valueOf(attributesOpt).split(",")) else None
+    val oneTime = interval < 0 || options.has(oneTimeOpt)
+    val attributesIncludeExists = options.has(attributesOpt)
+    val attributesInclude = if(attributesIncludeExists) Some(options.valueOf(attributesOpt).split(",").filterNot(_.equals(""))) else None
     val dateFormatExists = options.has(dateFormatOpt)
     val dateFormat = if(dateFormatExists) Some(new SimpleDateFormat(options.valueOf(dateFormatOpt))) else None
     val wait = options.has(waitOpt)
 
+    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
@@ -101,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 {
@@ -124,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
@@ -149,48 +190,86 @@ object JmxTool extends Logging {
     }
 
     val numExpectedAttributes: Map[ObjectName, Int] =
-      if (attributesWhitelistExists)
-        queries.map((_, attributesWhitelist.get.size)).toMap
-      else {
-        names.map{(name: ObjectName) =>
+      if (!attributesIncludeExists)
+        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 => attributesInclude.get.contains(attr.getName))
+            (name, expectedAttributes.size)}.toMap.filter(_._2 > 0)
+        else
+          queries.map((_, attributesInclude.get.length)).toMap
       }
 
+    if(numExpectedAttributes.isEmpty) {
+      CommandLineUtils.printUsageAndDie(parser, s"No matched attributes for the queried objects $queries.")
+    }
+
     // print csv header
-    val keys = List("time") ++ queryAttributes(mbsc, names, attributesWhitelist).keys.toArray.sorted
-    if(keys.size == numExpectedAttributes.values.sum + 1)
+    val keys = List("time") ++ queryAttributes(mbsc, names, attributesInclude).keys.toArray.sorted
+    if(reportFormatOriginal && keys.size == numExpectedAttributes.values.sum + 1) {
       println(keys.map("\"" + _ + "\"").mkString(","))
+    }
 
-    while(true) {
+    var keepGoing = true
+    while (keepGoing) {
       val start = System.currentTimeMillis
-      val attributes = queryAttributes(mbsc, names, attributesWhitelist)
+      val attributes = queryAttributes(mbsc, names, attributesInclude)
       attributes("time") = dateFormat match {
         case Some(dFormat) => dFormat.format(new Date)
         case None => System.currentTimeMillis().toString
       }
-      if(attributes.keySet.size == numExpectedAttributes.values.sum + 1)
-        println(keys.map(attributes(_)).mkString(","))
-      val sleep = max(0, interval - (System.currentTimeMillis - start))
-      Thread.sleep(sleep)
+      if(attributes.keySet.size == numExpectedAttributes.values.sum + 1) {
+        if(reportFormatOriginal) {
+          println(keys.map(attributes(_)).mkString(","))
+        }
+        else if(reportFormat.equals("properties")) {
+          keys.foreach( k => { println(k + "=" + attributes(k) ) } )
+        }
+        else if(reportFormat.equals("csv")) {
+          keys.foreach( k => { println(k + ",\"" + attributes(k) + "\"" ) } )
+        }
+        else { // tsv
+          keys.foreach( k => { println(k + "\t" + attributes(k) ) } )
+        }
+      }
+
+      if (oneTime) {
+        keepGoing = false
+      }
+      else {
+        val sleep = max(0, interval - (System.currentTimeMillis - start))
+        Thread.sleep(sleep)
+      }
     }
   }
 
-  def queryAttributes(mbsc: MBeanServerConnection, names: Iterable[ObjectName], attributesWhitelist: Option[Array[String]]) = {
+  def queryAttributes(mbsc: MBeanServerConnection, names: Iterable[ObjectName], attributesInclude: Option[Array[String]]): mutable.Map[String, Any] = {
     val attributes = new mutable.HashMap[String, Any]()
     for (name <- names) {
       val mbean = mbsc.getMBeanInfo(name)
       for (attrObj <- mbsc.getAttributes(name, mbean.getAttributes.map(_.getName)).asScala) {
         val attr = attrObj.asInstanceOf[Attribute]
-        attributesWhitelist match {
+        attributesInclude match {
           case Some(allowedAttributes) =>
             if (allowedAttributes.contains(attr.getName))
-              attributes(name + ":" + attr.getName) = attr.getValue
-          case None => attributes(name + ":" + attr.getName) = attr.getValue
+              attributes(name.toString + ":" + attr.getName) = attr.getValue
+          case None => attributes(name.toString + ":" + attr.getName) = attr.getValue
         }
       }
     }
     attributes
   }
 
+  def parseFormat(reportFormatOpt : String): String = reportFormatOpt match {
+    case "properties" => "properties"
+    case "csv" => "csv"
+    case "tsv" => "tsv"
+    case _ => "original"
+  }
 }
diff --git a/core/src/main/scala/kafka/tools/MirrorMaker.scala b/core/src/main/scala/kafka/tools/MirrorMaker.scala
index 618fd2a95b96c..c1146f2026503 100755
--- a/core/src/main/scala/kafka/tools/MirrorMaker.scala
+++ b/core/src/main/scala/kafka/tools/MirrorMaker.scala
@@ -17,45 +17,41 @@
 
 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}
 
-import com.yammer.metrics.core.Gauge
-import joptsimple.OptionParser
-import kafka.consumer.{BaseConsumer, BaseConsumerRecord, Blacklist, ConsumerIterator, ConsumerThreadId, ConsumerTimeoutException, TopicFilter, Whitelist, ZookeeperConsumerConnector, ConsumerConfig => OldConsumerConfig}
-import kafka.javaapi.consumer.ConsumerRebalanceListener
+import kafka.consumer.BaseConsumerRecord
 import kafka.metrics.KafkaMetricsGroup
-import kafka.serializer.DefaultDecoder
-import kafka.utils.{CommandLineUtils, CoreUtils, Logging, ZKConfig}
-import org.apache.kafka.clients.consumer
-import org.apache.kafka.clients.consumer.{CommitFailedException, Consumer, ConsumerRecord, KafkaConsumer, OffsetAndMetadata}
+import kafka.utils._
+import org.apache.kafka.clients.consumer._
 import org.apache.kafka.clients.producer.internals.ErrorLoggingCallback
 import org.apache.kafka.clients.producer.{KafkaProducer, ProducerConfig, ProducerRecord, RecordMetadata}
-import org.apache.kafka.common.TopicPartition
-import org.apache.kafka.common.serialization.ByteArrayDeserializer
-import org.apache.kafka.common.utils.Utils
-import org.apache.kafka.common.errors.WakeupException
+import org.apache.kafka.common.errors.{TimeoutException, WakeupException}
+import org.apache.kafka.common.record.RecordBatch
+import org.apache.kafka.common.serialization.{ByteArrayDeserializer, ByteArraySerializer}
+import org.apache.kafka.common.utils.{Time, Utils}
+import org.apache.kafka.common.{KafkaException, TopicPartition}
 
-import scala.collection.JavaConverters._
+import scala.annotation.nowarn
+import scala.jdk.CollectionConverters._
 import scala.collection.mutable.HashMap
 import scala.util.control.ControlThrowable
-import org.apache.kafka.clients.consumer.{ConsumerConfig => NewConsumerConfig}
-import org.apache.kafka.common.header.internals.RecordHeaders
-import org.apache.kafka.common.record.RecordBatch
+import scala.util.{Failure, Success, Try}
 
 /**
  * The mirror maker has the following architecture:
- * - There are N mirror maker thread shares one ZookeeperConsumerConnector and each owns a Kafka stream.
+ * - There are N mirror maker thread, each of which is equipped with a separate KafkaConsumer instance.
  * - All the mirror maker threads share one producer.
  * - Each mirror maker thread periodically flushes the producer and then commits all offsets.
  *
  * @note For mirror maker, the following settings are set by default to make sure there is no data loss:
- *       1. use new producer with following settings
+ *       1. use producer with following settings
  *            acks=all
- *            retries=max integer
+ *            delivery.timeout.ms=max integer
  *            max.block.ms=max long
  *            max.in.flight.requests.per.connection=1
  *       2. Consumer Settings
@@ -67,239 +63,31 @@ object MirrorMaker extends Logging with KafkaMetricsGroup {
 
   private[tools] var producer: MirrorMakerProducer = null
   private var mirrorMakerThreads: Seq[MirrorMakerThread] = null
-  private val isShuttingdown: AtomicBoolean = new AtomicBoolean(false)
+  private val isShuttingDown: AtomicBoolean = new AtomicBoolean(false)
   // Track the messages not successfully sent by mirror maker.
   private val numDroppedMessages: AtomicInteger = new AtomicInteger(0)
   private var messageHandler: MirrorMakerMessageHandler = null
   private var offsetCommitIntervalMs = 0
   private var abortOnSendFailure: Boolean = true
   @volatile private var exitingOnSendFailure: Boolean = false
+  private var lastSuccessfulCommitTime = -1L
+  private val time = Time.SYSTEM
 
   // If a message send failed after retries are exhausted. The offset of the messages will also be removed from
   // the unacked offset list to avoid offset commit being stuck on that offset. In this case, the offset of that
   // message was not really acked, but was skipped. This metric records the number of skipped offsets.
-  newGauge("MirrorMaker-numDroppedMessages",
-    new Gauge[Int] {
-      def value = numDroppedMessages.get()
-    })
+  newGauge("MirrorMaker-numDroppedMessages", () => numDroppedMessages.get())
 
-  def main(args: Array[String]) {
+  def main(args: Array[String]): Unit = {
 
     info("Starting mirror maker")
     try {
-      val parser = new OptionParser(false)
-
-      val consumerConfigOpt = parser.accepts("consumer.config",
-        "Embedded consumer config for consuming from the source cluster.")
-        .withRequiredArg()
-        .describedAs("config file")
-        .ofType(classOf[String])
-
-      val useNewConsumerOpt = parser.accepts("new.consumer",
-        "Use new consumer in mirror maker (this is the default).")
-
-      val producerConfigOpt = parser.accepts("producer.config",
-        "Embedded producer config.")
-        .withRequiredArg()
-        .describedAs("config file")
-        .ofType(classOf[String])
-
-      val numStreamsOpt = parser.accepts("num.streams",
-        "Number of consumption streams.")
-        .withRequiredArg()
-        .describedAs("Number of threads")
-        .ofType(classOf[java.lang.Integer])
-        .defaultsTo(1)
-
-      val whitelistOpt = parser.accepts("whitelist",
-        "Whitelist of topics to mirror.")
-        .withRequiredArg()
-        .describedAs("Java regex (String)")
-        .ofType(classOf[String])
-
-      val blacklistOpt = parser.accepts("blacklist",
-        "Blacklist of topics to mirror. Only old consumer supports blacklist.")
-        .withRequiredArg()
-        .describedAs("Java regex (String)")
-        .ofType(classOf[String])
-
-      val offsetCommitIntervalMsOpt = parser.accepts("offset.commit.interval.ms",
-        "Offset commit interval in ms.")
-        .withRequiredArg()
-        .describedAs("offset commit interval in millisecond")
-        .ofType(classOf[java.lang.Integer])
-        .defaultsTo(60000)
-
-      val consumerRebalanceListenerOpt = parser.accepts("consumer.rebalance.listener",
-        "The consumer rebalance listener to use for mirror maker consumer.")
-        .withRequiredArg()
-        .describedAs("A custom rebalance listener of type ConsumerRebalanceListener")
-        .ofType(classOf[String])
-
-      val rebalanceListenerArgsOpt = parser.accepts("rebalance.listener.args",
-        "Arguments used by custom rebalance listener for mirror maker consumer.")
-        .withRequiredArg()
-        .describedAs("Arguments passed to custom rebalance listener constructor as a string.")
-        .ofType(classOf[String])
-
-      val messageHandlerOpt = parser.accepts("message.handler",
-        "Message handler which will process every record in-between consumer and producer.")
-        .withRequiredArg()
-        .describedAs("A custom message handler of type MirrorMakerMessageHandler")
-        .ofType(classOf[String])
-
-      val messageHandlerArgsOpt = parser.accepts("message.handler.args",
-        "Arguments used by custom message handler for mirror maker.")
-        .withRequiredArg()
-        .describedAs("Arguments passed to message handler constructor.")
-        .ofType(classOf[String])
-
-      val abortOnSendFailureOpt = parser.accepts("abort.on.send.failure",
-        "Configure the mirror maker to exit on a failed send.")
-        .withRequiredArg()
-        .describedAs("Stop the entire mirror maker when a send failure occurs")
-        .ofType(classOf[String])
-        .defaultsTo("true")
-
-      val helpOpt = parser.accepts("help", "Print this message.")
-
-      if (args.length == 0)
-        CommandLineUtils.printUsageAndDie(parser, "Continuously copy data between two Kafka clusters.")
-
-
-      val options = parser.parse(args: _*)
-
-      if (options.has(helpOpt)) {
-        parser.printHelpOn(System.out)
-        sys.exit(0)
-      }
-
-      CommandLineUtils.checkRequiredArgs(parser, options, consumerConfigOpt, producerConfigOpt)
-
-      val consumerProps = Utils.loadProps(options.valueOf(consumerConfigOpt))
-      val useOldConsumer = consumerProps.containsKey(ZKConfig.ZkConnectProp)
-
-      if (useOldConsumer) {
-        if (options.has(useNewConsumerOpt)) {
-          error(s"The consumer configuration parameter `${ZKConfig.ZkConnectProp}` is not valid when using --new.consumer")
-          sys.exit(1)
-        }
-
-        if (consumerProps.containsKey(NewConsumerConfig.BOOTSTRAP_SERVERS_CONFIG)) {
-          error(s"The configuration parameters `${ZKConfig.ZkConnectProp}` (old consumer) and " +
-            s"`${NewConsumerConfig.BOOTSTRAP_SERVERS_CONFIG}` (new consumer) cannot be used together.")
-          sys.exit(1)
-        }
-
-        if (List(whitelistOpt, blacklistOpt).count(options.has) != 1) {
-          error("Exactly one of whitelist or blacklist is required.")
-          sys.exit(1)
-        }
-      } else {
-        if (options.has(blacklistOpt)) {
-          error("blacklist can not be used when using new consumer in mirror maker. Use whitelist instead.")
-          sys.exit(1)
-        }
-
-        if (!options.has(whitelistOpt)) {
-          error("whitelist must be specified when using new consumer in mirror maker.")
-          sys.exit(1)
-        }
-
-        if (!consumerProps.containsKey(NewConsumerConfig.PARTITION_ASSIGNMENT_STRATEGY_CONFIG))
-          System.err.println("WARNING: The default partition assignment strategy of the new-consumer-based mirror maker will " +
-            "change from 'range' to 'roundrobin' in an upcoming release (so that better load balancing can be achieved). If " +
-            "you prefer to make this switch in advance of that release add the following to the corresponding new-consumer " +
-            "config: 'partition.assignment.strategy=org.apache.kafka.clients.consumer.RoundRobinAssignor'")
-
-      }
-
-      abortOnSendFailure = options.valueOf(abortOnSendFailureOpt).toBoolean
-      offsetCommitIntervalMs = options.valueOf(offsetCommitIntervalMsOpt).intValue()
-      val numStreams = options.valueOf(numStreamsOpt).intValue()
-
-      Runtime.getRuntime.addShutdownHook(new Thread("MirrorMakerShutdownHook") {
-        override def run() {
-          cleanShutdown()
-        }
-      })
-
-      // create producer
-      val producerProps = Utils.loadProps(options.valueOf(producerConfigOpt))
-      val sync = producerProps.getProperty("producer.type", "async").equals("sync")
-      producerProps.remove("producer.type")
-      // Defaults to no data loss settings.
-      maybeSetDefaultProperty(producerProps, ProducerConfig.RETRIES_CONFIG, Int.MaxValue.toString)
-      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")
-      // Always set producer key and value serializer to ByteArraySerializer.
-      producerProps.setProperty(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.ByteArraySerializer")
-      producerProps.setProperty(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.ByteArraySerializer")
-      producer = new MirrorMakerProducer(sync, producerProps)
-
-      // Create consumers
-      val mirrorMakerConsumers = if (useOldConsumer) {
-        val customRebalanceListener: Option[ConsumerRebalanceListener] = {
-          val customRebalanceListenerClass = options.valueOf(consumerRebalanceListenerOpt)
-          if (customRebalanceListenerClass != null) {
-            val rebalanceListenerArgs = options.valueOf(rebalanceListenerArgsOpt)
-            if (rebalanceListenerArgs != null) {
-              Some(CoreUtils.createObject[ConsumerRebalanceListener](customRebalanceListenerClass, rebalanceListenerArgs))
-            } else {
-              Some(CoreUtils.createObject[ConsumerRebalanceListener](customRebalanceListenerClass))
-            }
-          } else {
-            None
-          }
-        }
-        createOldConsumers(
-          numStreams,
-          consumerProps,
-          customRebalanceListener,
-          Option(options.valueOf(whitelistOpt)),
-          Option(options.valueOf(blacklistOpt)))
-      } else {
-        val customRebalanceListener: Option[org.apache.kafka.clients.consumer.ConsumerRebalanceListener] = {
-          val customRebalanceListenerClass = options.valueOf(consumerRebalanceListenerOpt)
-          if (customRebalanceListenerClass != null) {
-            val rebalanceListenerArgs = options.valueOf(rebalanceListenerArgsOpt)
-            if (rebalanceListenerArgs != null) {
-                Some(CoreUtils.createObject[org.apache.kafka.clients.consumer.ConsumerRebalanceListener](customRebalanceListenerClass, rebalanceListenerArgs))
-            } else {
-                Some(CoreUtils.createObject[org.apache.kafka.clients.consumer.ConsumerRebalanceListener](customRebalanceListenerClass))
-            }
-          } else {
-            None
-          }
-        }
-        createNewConsumers(
-          numStreams,
-          consumerProps,
-          customRebalanceListener,
-          Option(options.valueOf(whitelistOpt)))
-      }
-
-      // Create mirror maker threads.
-      mirrorMakerThreads = (0 until numStreams) map (i =>
-        new MirrorMakerThread(mirrorMakerConsumers(i), i))
-
-      // Create and initialize message handler
-      val customMessageHandlerClass = options.valueOf(messageHandlerOpt)
-      val messageHandlerArgs = options.valueOf(messageHandlerArgsOpt)
-      messageHandler = {
-        if (customMessageHandlerClass != null) {
-          if (messageHandlerArgs != null)
-            CoreUtils.createObject[MirrorMakerMessageHandler](customMessageHandlerClass, messageHandlerArgs)
-          else
-            CoreUtils.createObject[MirrorMakerMessageHandler](customMessageHandlerClass)
-        } else {
-          defaultMirrorMakerMessageHandler
-        }
-      }
+      val opts = new MirrorMakerOptions(args)
+      CommandLineUtils.printHelpAndExitIfNeeded(opts, "This tool helps to continuously copy data between two Kafka clusters.")
+      opts.checkArgs()
     } catch {
-      case ct : ControlThrowable => throw ct
-      case t : Throwable =>
+      case ct: ControlThrowable => throw ct
+      case t: Throwable =>
         error("Exception when starting mirror maker.", t)
     }
 
@@ -307,43 +95,10 @@ object MirrorMaker extends Logging with KafkaMetricsGroup {
     mirrorMakerThreads.foreach(_.awaitShutdown())
   }
 
-  private def createOldConsumers(numStreams: Int,
-                                 consumerConfigProps: Properties,
-                                 customRebalanceListener: Option[ConsumerRebalanceListener],
-                                 whitelist: Option[String],
-                                 blacklist: Option[String]) : Seq[MirrorMakerBaseConsumer] = {
-    // Disable consumer auto offsets commit to prevent data loss.
-    maybeSetDefaultProperty(consumerConfigProps, "auto.commit.enable", "false")
-    // Set the consumer timeout so we will not block for low volume pipeline. The timeout is necessary to make sure
-    // Offsets are still committed for those low volume pipelines.
-    maybeSetDefaultProperty(consumerConfigProps, "consumer.timeout.ms", "10000")
-    // The default client id is group id, we manually set client id to groupId-index to avoid metric collision
-    val groupIdString = consumerConfigProps.getProperty("group.id")
-    val connectors = (0 until numStreams) map { i =>
-      consumerConfigProps.setProperty("client.id", groupIdString + "-" + i.toString)
-      val consumerConfig = new OldConsumerConfig(consumerConfigProps)
-      new ZookeeperConsumerConnector(consumerConfig)
-    }
-
-    // create filters
-    val filterSpec = if (whitelist.isDefined)
-      new Whitelist(whitelist.get)
-    else if (blacklist.isDefined)
-      new Blacklist(blacklist.get)
-    else
-      throw new IllegalArgumentException("Either whitelist or blacklist should be defined!")
-    (0 until numStreams) map { i =>
-      val consumer = new MirrorMakerOldConsumer(connectors(i), filterSpec)
-      val consumerRebalanceListener = new InternalRebalanceListenerForOldConsumer(consumer, customRebalanceListener)
-      connectors(i).setConsumerRebalanceListener(consumerRebalanceListener)
-      consumer
-    }
-  }
-
-  def createNewConsumers(numStreams: Int,
-                         consumerConfigProps: Properties,
-                         customRebalanceListener: Option[org.apache.kafka.clients.consumer.ConsumerRebalanceListener],
-                         whitelist: Option[String]) : Seq[MirrorMakerBaseConsumer] = {
+  def createConsumers(numStreams: Int,
+                      consumerConfigProps: Properties,
+                      customRebalanceListener: Option[ConsumerRebalanceListener],
+                      whitelist: Option[String]): Seq[ConsumerWrapper] = {
     // Disable consumer auto offsets commit to prevent data loss.
     maybeSetDefaultProperty(consumerConfigProps, "enable.auto.commit", "false")
     // Hardcode the deserializer to ByteArrayDeserializer
@@ -355,36 +110,57 @@ object MirrorMaker extends Logging with KafkaMetricsGroup {
       consumerConfigProps.setProperty("client.id", groupIdString + "-" + i.toString)
       new KafkaConsumer[Array[Byte], Array[Byte]](consumerConfigProps)
     }
-    whitelist.getOrElse(throw new IllegalArgumentException("White list cannot be empty for new consumer"))
-    consumers.map(consumer => new MirrorMakerNewConsumer(consumer, customRebalanceListener, whitelist))
+    whitelist.getOrElse(throw new IllegalArgumentException("White list cannot be empty"))
+    consumers.map(consumer => new ConsumerWrapper(consumer, customRebalanceListener, whitelist))
   }
 
-  def commitOffsets(mirrorMakerConsumer: MirrorMakerBaseConsumer) {
+  def commitOffsets(consumerWrapper: ConsumerWrapper): Unit = {
     if (!exitingOnSendFailure) {
-      trace("Committing offsets.")
-      try {
-        mirrorMakerConsumer.commit()
-      } catch {
-        case e: WakeupException =>
-          // we only call wakeup() once to close the consumer,
-          // so if we catch it in commit we can safely retry
-          // and re-throw to break the loop
-          mirrorMakerConsumer.commit()
-          throw e
-
-        case _: CommitFailedException =>
-          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 ${consumer.ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG} or reduce the number of records " +
-            s"handled on each iteration with ${consumer.ConsumerConfig.MAX_POLL_RECORDS_CONFIG}")
+      var retry = 0
+      var retryNeeded = true
+      while (retryNeeded) {
+        trace("Committing offsets.")
+        try {
+          consumerWrapper.commit()
+          lastSuccessfulCommitTime = time.milliseconds
+          retryNeeded = false
+        } catch {
+          case e: WakeupException =>
+            // we only call wakeup() once to close the consumer,
+            // so if we catch it in commit we can safely retry
+            // and re-throw to break the loop
+            commitOffsets(consumerWrapper)
+            throw e
+
+          case _: TimeoutException =>
+            Try(consumerWrapper.consumer.listTopics) match {
+              case Success(visibleTopics) =>
+                consumerWrapper.offsets --= consumerWrapper.offsets.keySet.filter(tp => !visibleTopics.containsKey(tp.topic))
+              case Failure(e) =>
+                warn("Failed to list all authorized topics after committing offsets timed out: ", e)
+            }
+
+            retry += 1
+            warn("Failed to commit offsets because the offset commit request processing can not be completed in time. " +
+              s"If you see this regularly, it could indicate that you need to increase the consumer's ${ConsumerConfig.DEFAULT_API_TIMEOUT_MS_CONFIG} " +
+              s"Last successful offset commit timestamp=$lastSuccessfulCommitTime, retry count=$retry")
+            Thread.sleep(100)
+
+          case _: CommitFailedException =>
+            retryNeeded = false
+            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}")
+        }
       }
     } else {
       info("Exiting on send failure, skip committing offsets.")
     }
   }
 
-  def cleanShutdown() {
-    if (isShuttingdown.compareAndSet(false, true)) {
+  def cleanShutdown(): Unit = {
+    if (isShuttingDown.compareAndSet(false, true)) {
       info("Start clean shutdown.")
       // Shutdown consumer threads.
       info("Shutting down consumer threads.")
@@ -398,14 +174,14 @@ 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)
       info("Property %s is overridden to %s - data loss or message reordering is possible.".format(propertyName, propertyValue))
   }
 
-  class MirrorMakerThread(mirrorMakerConsumer: MirrorMakerBaseConsumer,
+  class MirrorMakerThread(consumerWrapper: ConsumerWrapper,
                           val threadId: Int) extends Thread with Logging with KafkaMetricsGroup {
     private val threadName = "mirrormaker-thread-" + threadId
     private val shutdownLatch: CountDownLatch = new CountDownLatch(1)
@@ -415,27 +191,43 @@ object MirrorMaker extends Logging with KafkaMetricsGroup {
 
     setName(threadName)
 
-    override def run() {
-      info("Starting mirror maker thread " + threadName)
+    @nowarn("cat=deprecation")
+    private def toBaseConsumerRecord(record: ConsumerRecord[Array[Byte], Array[Byte]]): BaseConsumerRecord =
+      BaseConsumerRecord(record.topic,
+        record.partition,
+        record.offset,
+        record.timestamp,
+        record.timestampType,
+        record.key,
+        record.value,
+        record.headers)
+
+    override def run(): Unit = {
+      info(s"Starting mirror maker thread $threadName")
       try {
-        mirrorMakerConsumer.init()
+        consumerWrapper.init()
 
-        // We need the two while loop to make sure when old consumer is used, even there is no message we
-        // still commit offset. When new consumer is used, this is handled by poll(timeout).
+        // We needed two while loops due to the old consumer semantics, this can now be simplified
         while (!exitingOnSendFailure && !shuttingDown) {
           try {
-            while (!exitingOnSendFailure && !shuttingDown && mirrorMakerConsumer.hasData) {
-              val data = mirrorMakerConsumer.receive()
-              trace("Sending message with value size %d and offset %d".format(data.value.length, data.offset))
-              val records = messageHandler.handle(data)
-              records.asScala.foreach(producer.send)
+            while (!exitingOnSendFailure && !shuttingDown) {
+              val data = consumerWrapper.receive()
+              if (data.value != null) {
+                trace("Sending message with value size %d and offset %d.".format(data.value.length, data.offset))
+              } else {
+                trace("Sending message with null value and offset %d.".format(data.offset))
+              }
+              val records = messageHandler.handle(toBaseConsumerRecord(data))
+              records.forEach(producer.send)
               maybeFlushAndCommitOffsets()
             }
           } catch {
-            case _: ConsumerTimeoutException =>
-              trace("Caught ConsumerTimeoutException, continue iteration.")
+            case _: NoRecordsException =>
+              trace("Caught NoRecordsException, continue iteration.")
             case _: WakeupException =>
-              trace("Caught ConsumerWakeupException, continue iteration.")
+              trace("Caught WakeupException, continue iteration.")
+            case e: KafkaException if (shuttingDown || exitingOnSendFailure) =>
+              trace(s"Ignoring caught KafkaException during shutdown. sendFailure: $exitingOnSendFailure.", e)
           }
           maybeFlushAndCommitOffsets()
         }
@@ -450,39 +242,36 @@ object MirrorMaker extends Logging with KafkaMetricsGroup {
 
           // note that this commit is skipped if flush() fails which ensures that we don't lose messages
           info("Committing consumer offsets.")
-          commitOffsets(mirrorMakerConsumer)
+          commitOffsets(consumerWrapper)
         }, this)
 
         info("Shutting down consumer connectors.")
-        CoreUtils.swallow(mirrorMakerConsumer.stop(), this)
-        CoreUtils.swallow(mirrorMakerConsumer.cleanup(), this)
+        CoreUtils.swallow(consumerWrapper.wakeup(), this)
+        CoreUtils.swallow(consumerWrapper.close(), this)
         shutdownLatch.countDown()
         info("Mirror maker thread stopped")
         // if it exits accidentally, stop the entire mirror maker
-        if (!isShuttingdown.get()) {
+        if (!isShuttingDown.get()) {
           fatal("Mirror maker thread exited abnormally, stopping the whole mirror maker.")
           sys.exit(-1)
         }
       }
     }
 
-    def maybeFlushAndCommitOffsets() {
-      val commitRequested = mirrorMakerConsumer.commitRequested()
-      if (commitRequested || System.currentTimeMillis() - lastOffsetCommitMs > offsetCommitIntervalMs) {
+    def maybeFlushAndCommitOffsets(): Unit = {
+      if (System.currentTimeMillis() - lastOffsetCommitMs > offsetCommitIntervalMs) {
         debug("Committing MirrorMaker state.")
         producer.flush()
-        commitOffsets(mirrorMakerConsumer)
+        commitOffsets(consumerWrapper)
         lastOffsetCommitMs = System.currentTimeMillis()
-        if (commitRequested)
-          mirrorMakerConsumer.notifyCommit()
       }
     }
 
-    def shutdown() {
+    def shutdown(): Unit = {
       try {
-        info(threadName + " shutting down")
+        info(s"$threadName shutting down")
         shuttingDown = true
-        mirrorMakerConsumer.stop()
+        consumerWrapper.wakeup()
       }
       catch {
         case _: InterruptedException =>
@@ -490,7 +279,7 @@ object MirrorMaker extends Logging with KafkaMetricsGroup {
       }
     }
 
-    def awaitShutdown() {
+    def awaitShutdown(): Unit = {
       try {
         shutdownLatch.await()
         info("Mirror maker thread shutdown complete")
@@ -501,110 +290,23 @@ object MirrorMaker extends Logging with KafkaMetricsGroup {
     }
   }
 
-  private[kafka] trait MirrorMakerBaseConsumer extends BaseConsumer {
-    def init()
-    def commitRequested(): Boolean
-    def notifyCommit()
-    def requestAndWaitForCommit()
-    def hasData : Boolean
-  }
-
-  private class MirrorMakerOldConsumer(connector: ZookeeperConsumerConnector,
-                                       filterSpec: TopicFilter) extends MirrorMakerBaseConsumer {
-    private var iter: ConsumerIterator[Array[Byte], Array[Byte]] = null
-    private var immediateCommitRequested: Boolean = false
-    private var numCommitsNotified: Long = 0
-
-    override def init() {
-      // Creating one stream per each connector instance
-      val streams = connector.createMessageStreamsByFilter(filterSpec, 1, new DefaultDecoder(), new DefaultDecoder())
-      require(streams.size == 1)
-      val stream = streams.head
-      iter = stream.iterator()
-    }
-
-    override def requestAndWaitForCommit() {
-      this.synchronized {
-        // only wait() if mirrorMakerConsumer has been initialized and it has not been cleaned up.
-        if (iter != null) {
-          immediateCommitRequested = true
-          val nextNumCommitsNotified = numCommitsNotified + 1
-          do {
-            this.wait()
-          } while (numCommitsNotified < nextNumCommitsNotified)
-        }
-      }
-    }
-
-    override def notifyCommit() {
-      this.synchronized {
-        immediateCommitRequested = false
-        numCommitsNotified = numCommitsNotified + 1
-        this.notifyAll()
-      }
-    }
-
-    override def commitRequested(): Boolean = {
-      this.synchronized {
-        immediateCommitRequested
-      }
-    }
-
-    override def hasData = iter.hasNext()
-
-    override def receive() : BaseConsumerRecord = {
-      val messageAndMetadata = iter.next()
-      BaseConsumerRecord(messageAndMetadata.topic,
-                         messageAndMetadata.partition,
-                         messageAndMetadata.offset,
-                         messageAndMetadata.timestamp,
-                         messageAndMetadata.timestampType,
-                         messageAndMetadata.key,
-                         messageAndMetadata.message,
-                         new RecordHeaders())
-    }
-
-    override def stop() {
-      // Do nothing
-    }
-
-    override def cleanup() {
-      // We need to set the iterator to null and notify the rebalance listener thread.
-      // This is to handle the case that the consumer rebalance is triggered when the
-      // mirror maker thread is shutting down and the rebalance listener is waiting for the offset commit.
-      this.synchronized {
-        iter = null
-        if (immediateCommitRequested) {
-          notifyCommit()
-        }
-      }
-      connector.shutdown()
-    }
-
-    override def commit() {
-      connector.commitOffsets
-    }
-  }
-
-  // Only for testing
-  private[tools] class MirrorMakerNewConsumer(consumer: Consumer[Array[Byte], Array[Byte]],
-                                       customRebalanceListener: Option[org.apache.kafka.clients.consumer.ConsumerRebalanceListener],
-                                       whitelistOpt: Option[String])
-    extends MirrorMakerBaseConsumer {
+  // Visible for testing
+  private[tools] class ConsumerWrapper(private[tools] val consumer: Consumer[Array[Byte], Array[Byte]],
+                                       customRebalanceListener: Option[ConsumerRebalanceListener],
+                                       whitelistOpt: Option[String]) {
     val regex = whitelistOpt.getOrElse(throw new IllegalArgumentException("New consumer only supports whitelist."))
     var recordIter: java.util.Iterator[ConsumerRecord[Array[Byte], Array[Byte]]] = null
 
-    // TODO: we need to manually maintain the consumed offsets for new consumer
-    // since its internal consumed position is updated in batch rather than one
-    // record at a time, this can be resolved when we break the unification of both consumers
-    private val offsets = new HashMap[TopicPartition, Long]()
+    // We manually maintain the consumed offsets for historical reasons and it could be simplified
+    // Visible for testing
+    private[tools] val offsets = new HashMap[TopicPartition, Long]()
 
-    override def init() {
-      debug("Initiating new consumer")
-      val consumerRebalanceListener = new InternalRebalanceListenerForNewConsumer(this, customRebalanceListener)
+    def init(): Unit = {
+      debug("Initiating consumer")
+      val consumerRebalanceListener = new InternalRebalanceListener(this, customRebalanceListener)
       whitelistOpt.foreach { whitelist =>
         try {
-          consumer.subscribe(Pattern.compile(Whitelist(whitelist).regex), consumerRebalanceListener)
+          consumer.subscribe(Pattern.compile(IncludeList(whitelist).regex), consumerRebalanceListener)
         } catch {
           case pse: RuntimeException =>
             error(s"Invalid expression syntax: $whitelist")
@@ -613,90 +315,53 @@ object MirrorMaker extends Logging with KafkaMetricsGroup {
       }
     }
 
-    override def requestAndWaitForCommit() {
-      // Do nothing
-    }
-
-    override def notifyCommit() {
-      // Do nothing
-    }
-    
-    override def commitRequested(): Boolean = {
-      false
-    }
-
-    override def hasData = true
-
-    override def receive() : BaseConsumerRecord = {
+    def receive(): ConsumerRecord[Array[Byte], Array[Byte]] = {
       if (recordIter == null || !recordIter.hasNext) {
         // In scenarios where data does not arrive within offsetCommitIntervalMs and
         // offsetCommitIntervalMs is less than poll's timeout, offset commit will be delayed for any
         // 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(1000).iterator
+        recordIter = consumer.poll(Duration.ofSeconds(1L)).iterator
         if (!recordIter.hasNext)
-          throw new ConsumerTimeoutException
+          throw new NoRecordsException
       }
 
       val record = recordIter.next()
       val tp = new TopicPartition(record.topic, record.partition)
 
       offsets.put(tp, record.offset + 1)
-
-      BaseConsumerRecord(record.topic,
-                         record.partition,
-                         record.offset,
-                         record.timestamp,
-                         record.timestampType,
-                         record.key,
-                         record.value,
-                         record.headers)
+      record
     }
 
-    override def stop() {
+    def wakeup(): Unit = {
       consumer.wakeup()
     }
 
-    override def cleanup() {
+    def close(): Unit = {
       consumer.close()
     }
 
-    override def commit() {
-      consumer.commitSync(offsets.map { case (tp, offset) =>  (tp, new OffsetAndMetadata(offset, ""))}.asJava)
+    def commit(): Unit = {
+      consumer.commitSync(offsets.map { case (tp, offset) => (tp, new OffsetAndMetadata(offset)) }.asJava)
       offsets.clear()
     }
   }
 
-  private class InternalRebalanceListenerForNewConsumer(mirrorMakerConsumer: MirrorMakerBaseConsumer,
-                                                        customRebalanceListenerForNewConsumer: Option[org.apache.kafka.clients.consumer.ConsumerRebalanceListener])
-    extends org.apache.kafka.clients.consumer.ConsumerRebalanceListener {
-
-    override def onPartitionsRevoked(partitions: util.Collection[TopicPartition]) {
-      producer.flush()
-      commitOffsets(mirrorMakerConsumer)
-      customRebalanceListenerForNewConsumer.foreach(_.onPartitionsRevoked(partitions))
-    }
-
-    override def onPartitionsAssigned(partitions: util.Collection[TopicPartition]) {
-      customRebalanceListenerForNewConsumer.foreach(_.onPartitionsAssigned(partitions))
-    }
-  }
-
-  private class InternalRebalanceListenerForOldConsumer(mirrorMakerConsumer: MirrorMakerBaseConsumer,
-                                                        customRebalanceListenerForOldConsumer: Option[ConsumerRebalanceListener])
+  private class InternalRebalanceListener(consumerWrapper: ConsumerWrapper,
+                                          customRebalanceListener: Option[ConsumerRebalanceListener])
     extends ConsumerRebalanceListener {
 
-    override def beforeReleasingPartitions(partitionOwnership: java.util.Map[String, java.util.Set[java.lang.Integer]]) {
-      // The zookeeper listener thread, which executes this method, needs to wait for MirrorMakerThread to flush data and commit offset
-      mirrorMakerConsumer.requestAndWaitForCommit()
-      // invoke custom consumer rebalance listener
-      customRebalanceListenerForOldConsumer.foreach(_.beforeReleasingPartitions(partitionOwnership))
+    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 beforeStartingFetchers(consumerId: String,
-                                        partitionAssignment: java.util.Map[String, java.util.Map[java.lang.Integer, ConsumerThreadId]]) {
-      customRebalanceListenerForOldConsumer.foreach(_.beforeStartingFetchers(consumerId, partitionAssignment))
+    override def onPartitionsAssigned(partitions: util.Collection[TopicPartition]): Unit = {
+      customRebalanceListener.foreach(_.onPartitionsAssigned(partitions))
     }
   }
 
@@ -704,7 +369,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 {
@@ -713,23 +378,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.
@@ -749,14 +414,163 @@ object MirrorMaker extends Logging with KafkaMetricsGroup {
    * If message.handler.args is specified. A constructor that takes in a String as argument must exist.
    */
   trait MirrorMakerMessageHandler {
+    @nowarn("cat=deprecation")
     def handle(record: BaseConsumerRecord): util.List[ProducerRecord[Array[Byte], Array[Byte]]]
   }
 
   private[tools] object defaultMirrorMakerMessageHandler extends MirrorMakerMessageHandler {
+    @nowarn("cat=deprecation")
     override def handle(record: BaseConsumerRecord): util.List[ProducerRecord[Array[Byte], Array[Byte]]] = {
       val timestamp: java.lang.Long = if (record.timestamp == RecordBatch.NO_TIMESTAMP) null else record.timestamp
       Collections.singletonList(new ProducerRecord(record.topic, null, timestamp, record.key, record.value, record.headers))
     }
   }
 
+  // package-private for tests
+  private[tools] class NoRecordsException extends RuntimeException
+
+  class MirrorMakerOptions(args: Array[String]) extends CommandDefaultOptions(args) {
+
+    val consumerConfigOpt = parser.accepts("consumer.config",
+      "Embedded consumer config for consuming from the source cluster.")
+      .withRequiredArg()
+      .describedAs("config file")
+      .ofType(classOf[String])
+
+    parser.accepts("new.consumer",
+      "DEPRECATED Use new consumer in mirror maker (this is the default so this option will be removed in " +
+        "a future version).")
+
+    val producerConfigOpt = parser.accepts("producer.config",
+      "Embedded producer config.")
+      .withRequiredArg()
+      .describedAs("config file")
+      .ofType(classOf[String])
+
+    val numStreamsOpt = parser.accepts("num.streams",
+      "Number of consumption streams.")
+      .withRequiredArg()
+      .describedAs("Number of threads")
+      .ofType(classOf[java.lang.Integer])
+      .defaultsTo(1)
+
+    val whitelistOpt = parser.accepts("whitelist",
+      "Whitelist of topics to mirror.")
+      .withRequiredArg()
+      .describedAs("Java regex (String)")
+      .ofType(classOf[String])
+
+    val offsetCommitIntervalMsOpt = parser.accepts("offset.commit.interval.ms",
+      "Offset commit interval in ms.")
+      .withRequiredArg()
+      .describedAs("offset commit interval in millisecond")
+      .ofType(classOf[java.lang.Integer])
+      .defaultsTo(60000)
+
+    val consumerRebalanceListenerOpt = parser.accepts("consumer.rebalance.listener",
+      "The consumer rebalance listener to use for mirror maker consumer.")
+      .withRequiredArg()
+      .describedAs("A custom rebalance listener of type ConsumerRebalanceListener")
+      .ofType(classOf[String])
+
+    val rebalanceListenerArgsOpt = parser.accepts("rebalance.listener.args",
+      "Arguments used by custom rebalance listener for mirror maker consumer.")
+      .withRequiredArg()
+      .describedAs("Arguments passed to custom rebalance listener constructor as a string.")
+      .ofType(classOf[String])
+
+    val messageHandlerOpt = parser.accepts("message.handler",
+      "Message handler which will process every record in-between consumer and producer.")
+      .withRequiredArg()
+      .describedAs("A custom message handler of type MirrorMakerMessageHandler")
+      .ofType(classOf[String])
+
+    val messageHandlerArgsOpt = parser.accepts("message.handler.args",
+      "Arguments used by custom message handler for mirror maker.")
+      .withRequiredArg()
+      .describedAs("Arguments passed to message handler constructor.")
+      .ofType(classOf[String])
+
+    val abortOnSendFailureOpt = parser.accepts("abort.on.send.failure",
+      "Configure the mirror maker to exit on a failed send.")
+      .withRequiredArg()
+      .describedAs("Stop the entire mirror maker when a send failure occurs")
+      .ofType(classOf[String])
+      .defaultsTo("true")
+
+    options = parser.parse(args: _*)
+
+    def checkArgs() = {
+      CommandLineUtils.checkRequiredArgs(parser, options, consumerConfigOpt, producerConfigOpt)
+      val consumerProps = Utils.loadProps(options.valueOf(consumerConfigOpt))
+
+      if (!options.has(whitelistOpt)) {
+        error("whitelist must be specified")
+        sys.exit(1)
+      }
+
+      if (!consumerProps.containsKey(ConsumerConfig.PARTITION_ASSIGNMENT_STRATEGY_CONFIG))
+        System.err.println("WARNING: The default partition assignment strategy of the mirror maker will " +
+          "change from 'range' to 'roundrobin' in an upcoming release (so that better load balancing can be achieved). If " +
+          "you prefer to make this switch in advance of that release add the following to the corresponding " +
+          "config: 'partition.assignment.strategy=org.apache.kafka.clients.consumer.RoundRobinAssignor'")
+
+      abortOnSendFailure = options.valueOf(abortOnSendFailureOpt).toBoolean
+      offsetCommitIntervalMs = options.valueOf(offsetCommitIntervalMsOpt).intValue()
+      val numStreams = options.valueOf(numStreamsOpt).intValue()
+
+      Exit.addShutdownHook("MirrorMakerShutdownHook", cleanShutdown())
+
+      // create producer
+      val producerProps = Utils.loadProps(options.valueOf(producerConfigOpt))
+      val sync = producerProps.getProperty("producer.type", "async").equals("sync")
+      producerProps.remove("producer.type")
+      // Defaults to no data loss settings.
+      maybeSetDefaultProperty(producerProps, ProducerConfig.DELIVERY_TIMEOUT_MS_CONFIG, Int.MaxValue.toString)
+      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")
+      // 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)
+      producer = new MirrorMakerProducer(sync, producerProps)
+
+      // Create consumers
+      val customRebalanceListener: Option[ConsumerRebalanceListener] = {
+        val customRebalanceListenerClass = options.valueOf(consumerRebalanceListenerOpt)
+        if (customRebalanceListenerClass != null) {
+          val rebalanceListenerArgs = options.valueOf(rebalanceListenerArgsOpt)
+          if (rebalanceListenerArgs != null)
+            Some(CoreUtils.createObject[ConsumerRebalanceListener](customRebalanceListenerClass, rebalanceListenerArgs))
+          else
+            Some(CoreUtils.createObject[ConsumerRebalanceListener](customRebalanceListenerClass))
+        } else {
+          None
+        }
+      }
+      val mirrorMakerConsumers = createConsumers(
+        numStreams,
+        consumerProps,
+        customRebalanceListener,
+        Option(options.valueOf(whitelistOpt)))
+
+      // Create mirror maker threads.
+      mirrorMakerThreads = (0 until numStreams) map (i =>
+        new MirrorMakerThread(mirrorMakerConsumers(i), i))
+
+      // Create and initialize message handler
+      val customMessageHandlerClass = options.valueOf(messageHandlerOpt)
+      val messageHandlerArgs = options.valueOf(messageHandlerArgsOpt)
+      messageHandler = {
+        if (customMessageHandlerClass != null) {
+          if (messageHandlerArgs != null)
+            CoreUtils.createObject[MirrorMakerMessageHandler](customMessageHandlerClass, messageHandlerArgs)
+          else
+            CoreUtils.createObject[MirrorMakerMessageHandler](customMessageHandlerClass)
+        } else {
+          defaultMirrorMakerMessageHandler
+        }
+      }
+    }
+  }
 }
diff --git a/core/src/main/scala/kafka/tools/PerfConfig.scala b/core/src/main/scala/kafka/tools/PerfConfig.scala
index 264ae6ae373f6..836163c85fa19 100644
--- a/core/src/main/scala/kafka/tools/PerfConfig.scala
+++ b/core/src/main/scala/kafka/tools/PerfConfig.scala
@@ -17,11 +17,10 @@
 
 package kafka.tools
 
-import joptsimple.OptionParser
+import kafka.utils.CommandDefaultOptions
 
 
-class PerfConfig(args: Array[String]) {
-  val parser = new OptionParser(false)
+class PerfConfig(args: Array[String]) extends CommandDefaultOptions(args) {
   val numMessagesOpt = parser.accepts("messages", "REQUIRED: The number of messages to send or consume")
     .withRequiredArg
     .describedAs("count")
@@ -38,5 +37,4 @@ class PerfConfig(args: Array[String]) {
     .ofType(classOf[String])
     .defaultsTo("yyyy-MM-dd HH:mm:ss:SSS")
   val hideHeaderOpt = parser.accepts("hide-header", "If set, skips printing the header for the stats ")
-  val helpOpt = parser.accepts("help", "Print usage.")
 }
diff --git a/core/src/main/scala/kafka/tools/ProducerPerformance.scala b/core/src/main/scala/kafka/tools/ProducerPerformance.scala
deleted file mode 100644
index 365c9afc58703..0000000000000
--- a/core/src/main/scala/kafka/tools/ProducerPerformance.scala
+++ /dev/null
@@ -1,309 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package kafka.tools
-
-import kafka.metrics.KafkaMetricsReporter
-import kafka.producer.{NewShinyProducer, OldProducer}
-import kafka.utils.{CommandLineUtils, Exit, Logging, ToolsUtils, VerifiableProperties}
-import kafka.utils.Implicits._
-import kafka.message.CompressionCodec
-import kafka.serializer._
-import java.util.concurrent.{CountDownLatch, Executors}
-import java.util.concurrent.atomic.AtomicLong
-import java.util._
-import java.text.SimpleDateFormat
-import java.math.BigInteger
-import java.nio.charset.StandardCharsets
-
-import org.apache.kafka.common.utils.Utils
-import org.slf4j.LoggerFactory
-
-/**
- * Load test for the producer
- */
-@deprecated("This class will be replaced by org.apache.kafka.tools.ProducerPerformance after the old producer client is removed", "0.9.0.0")
-object ProducerPerformance extends Logging {
-
-  def main(args: Array[String]) {
-    val config = new ProducerPerfConfig(args)
-    if (!config.isFixedSize)
-      logger.info("WARN: Throughput will be slower due to changing message size per request")
-
-    val totalBytesSent = new AtomicLong(0)
-    val totalMessagesSent = new AtomicLong(0)
-    val executor = Executors.newFixedThreadPool(config.numThreads)
-    val allDone = new CountDownLatch(config.numThreads)
-    val startMs = System.currentTimeMillis
-    val rand = new java.util.Random
-
-    if (!config.hideHeader)
-      println("start.time, end.time, compression, message.size, batch.size, total.data.sent.in.MB, MB.sec, " +
-        "total.data.sent.in.nMsg, nMsg.sec")
-
-    for (i <- 0 until config.numThreads) {
-      executor.execute(new ProducerThread(i, config, totalBytesSent, totalMessagesSent, allDone, rand))
-    }
-
-    allDone.await()
-    val endMs = System.currentTimeMillis
-    val elapsedSecs = (endMs - startMs) / 1000.0
-    val totalMBSent = (totalBytesSent.get * 1.0) / (1024 * 1024)
-    println("%s, %s, %d, %d, %d, %.2f, %.4f, %d, %.4f".format(
-      config.dateFormat.format(startMs), config.dateFormat.format(endMs),
-      config.compressionCodec.codec, config.messageSize, config.batchSize, totalMBSent,
-      totalMBSent / elapsedSecs, totalMessagesSent.get, totalMessagesSent.get / elapsedSecs))
-    Exit.exit(0)
-  }
-
-  class ProducerPerfConfig(args: Array[String]) extends PerfConfig(args) {
-    val brokerListOpt = parser.accepts("broker-list", "REQUIRED: broker info the list of broker host and port for bootstrap.")
-      .withRequiredArg
-      .describedAs("hostname:port,..,hostname:port")
-      .ofType(classOf[String])
-    val producerConfigOpt = parser.accepts("producer.config", "Producer config properties file.")
-      .withRequiredArg
-      .describedAs("config file")
-      .ofType(classOf[String])
-    val topicsOpt = parser.accepts("topics", "REQUIRED: The comma separated list of topics to produce to")
-      .withRequiredArg
-      .describedAs("topic1,topic2..")
-      .ofType(classOf[String])
-    val producerRequestTimeoutMsOpt = parser.accepts("request-timeout-ms", "The producer request timeout in ms")
-      .withRequiredArg()
-      .ofType(classOf[java.lang.Integer])
-      .defaultsTo(3000)
-    val producerNumRetriesOpt = parser.accepts("producer-num-retries", "The producer retries number")
-      .withRequiredArg()
-      .ofType(classOf[java.lang.Integer])
-      .defaultsTo(3)
-    val producerRetryBackOffMsOpt = parser.accepts("producer-retry-backoff-ms", "The producer retry backoff time in milliseconds")
-      .withRequiredArg()
-      .ofType(classOf[java.lang.Integer])
-      .defaultsTo(100)
-    val producerRequestRequiredAcksOpt = parser.accepts("request-num-acks", "Number of acks required for producer request " +
-      "to complete")
-      .withRequiredArg()
-      .ofType(classOf[java.lang.Integer])
-      .defaultsTo(-1)
-    val varyMessageSizeOpt = parser.accepts("vary-message-size", "If set, message size will vary up to the given maximum.")
-    val syncOpt = parser.accepts("sync", "If set, messages are sent synchronously.")
-    val numThreadsOpt = parser.accepts("threads", "Number of sending threads.")
-      .withRequiredArg
-      .describedAs("number of threads")
-      .ofType(classOf[java.lang.Integer])
-      .defaultsTo(1)
-    val initialMessageIdOpt = parser.accepts("initial-message-id", "The is used for generating test data, If set, messages will be tagged with an " +
-      "ID and sent by producer starting from this ID sequentially. Message content will be String type and " +
-      "in the form of 'Message:000...1:xxx...'")
-      .withRequiredArg()
-      .describedAs("initial message id")
-      .ofType(classOf[java.lang.Integer])
-    val messageSendGapMsOpt = parser.accepts("message-send-gap-ms", "If set, the send thread will wait for specified time between two sends")
-      .withRequiredArg()
-      .describedAs("message send time gap")
-      .ofType(classOf[java.lang.Integer])
-      .defaultsTo(0)
-    val csvMetricsReporterEnabledOpt = parser.accepts("csv-reporter-enabled", "If set, the CSV metrics reporter will be enabled")
-    val metricsDirectoryOpt = parser.accepts("metrics-dir", "If csv-reporter-enable is set, and this parameter is" +
-      "set, the csv metrics will be outputted here")
-      .withRequiredArg
-      .describedAs("metrics directory")
-      .ofType(classOf[java.lang.String])
-    val useNewProducerOpt = parser.accepts("new-producer", "Use the new producer implementation.")
-    val messageSizeOpt = parser.accepts("message-size", "The size of each message.")
-      .withRequiredArg
-      .describedAs("size")
-      .ofType(classOf[java.lang.Integer])
-      .defaultsTo(100)
-    val batchSizeOpt = parser.accepts("batch-size", "Number of messages to write in a single batch.")
-      .withRequiredArg
-      .describedAs("size")
-      .ofType(classOf[java.lang.Integer])
-      .defaultsTo(200)
-    val compressionCodecOpt = parser.accepts("compression-codec", "If set, messages are sent compressed")
-      .withRequiredArg
-      .describedAs("supported codec: NoCompressionCodec as 0, GZIPCompressionCodec as 1, SnappyCompressionCodec as 2, LZ4CompressionCodec as 3")
-      .ofType(classOf[java.lang.Integer])
-      .defaultsTo(0)
-
-    val options = parser.parse(args: _*)
-    CommandLineUtils.checkRequiredArgs(parser, options, topicsOpt, brokerListOpt, numMessagesOpt)
-
-    val topicsStr = options.valueOf(topicsOpt)
-    val topics = topicsStr.split(",")
-    val numMessages = options.valueOf(numMessagesOpt).longValue
-    val reportingInterval = options.valueOf(reportingIntervalOpt).intValue
-    val dateFormat = new SimpleDateFormat(options.valueOf(dateFormatOpt))
-    val hideHeader = options.has(hideHeaderOpt)
-    val brokerList = options.valueOf(brokerListOpt)
-    ToolsUtils.validatePortOrDie(parser,brokerList)
-    val messageSize = options.valueOf(messageSizeOpt).intValue
-    var isFixedSize = !options.has(varyMessageSizeOpt)
-    var isSync = options.has(syncOpt)
-    var batchSize = options.valueOf(batchSizeOpt).intValue
-    var numThreads = options.valueOf(numThreadsOpt).intValue
-    val compressionCodec = CompressionCodec.getCompressionCodec(options.valueOf(compressionCodecOpt).intValue)
-    val seqIdMode = options.has(initialMessageIdOpt)
-    var initialMessageId: Int = 0
-    if (seqIdMode)
-      initialMessageId = options.valueOf(initialMessageIdOpt).intValue()
-    val producerRequestTimeoutMs = options.valueOf(producerRequestTimeoutMsOpt).intValue()
-    val producerRequestRequiredAcks = options.valueOf(producerRequestRequiredAcksOpt).intValue()
-    val producerNumRetries = options.valueOf(producerNumRetriesOpt).intValue()
-    val producerRetryBackoffMs = options.valueOf(producerRetryBackOffMsOpt).intValue()
-    val useNewProducer = options.has(useNewProducerOpt)
-
-    val csvMetricsReporterEnabled = options.has(csvMetricsReporterEnabledOpt)
-
-    val producerProps = if (options.has(producerConfigOpt))
-      Utils.loadProps(options.valueOf(producerConfigOpt))
-    else
-      new Properties()
-
-    if (csvMetricsReporterEnabled) {
-      val props = new Properties()
-      props.put("kafka.metrics.polling.interval.secs", "1")
-      props.put("kafka.metrics.reporters", "kafka.metrics.KafkaCSVMetricsReporter")
-      if (options.has(metricsDirectoryOpt))
-        props.put("kafka.csv.metrics.dir", options.valueOf(metricsDirectoryOpt))
-      else
-        props.put("kafka.csv.metrics.dir", "kafka_metrics")
-      props.put("kafka.csv.metrics.reporter.enabled", "true")
-      val verifiableProps = new VerifiableProperties(props)
-      KafkaMetricsReporter.startReporters(verifiableProps)
-    }
-
-    val messageSendGapMs = options.valueOf(messageSendGapMsOpt).intValue()
-  }
-
-  class ProducerThread(val threadId: Int,
-    val config: ProducerPerfConfig,
-    val totalBytesSent: AtomicLong,
-    val totalMessagesSent: AtomicLong,
-    val allDone: CountDownLatch,
-    val rand: Random) extends Runnable {
-    val seqIdNumDigit = 10 // no. of digits for max int value
-
-    val messagesPerThread = config.numMessages / config.numThreads
-    debug("Messages per thread = " + messagesPerThread)
-    val props = new Properties()
-    val producer =
-      if (config.useNewProducer) {
-        import org.apache.kafka.clients.producer.ProducerConfig
-        props ++= config.producerProps
-        props.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, config.brokerList)
-        props.put(ProducerConfig.SEND_BUFFER_CONFIG, (64 * 1024).toString)
-        props.put(ProducerConfig.CLIENT_ID_CONFIG, "producer-performance")
-        props.put(ProducerConfig.ACKS_CONFIG, config.producerRequestRequiredAcks.toString)
-        props.put(ProducerConfig.RETRIES_CONFIG, config.producerNumRetries.toString)
-        props.put(ProducerConfig.RETRY_BACKOFF_MS_CONFIG, config.producerRetryBackoffMs.toString)
-        props.put(ProducerConfig.COMPRESSION_TYPE_CONFIG, config.compressionCodec.name)
-        props.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.ByteArraySerializer")
-        props.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.ByteArraySerializer")
-        new NewShinyProducer(props)
-      } else {
-        props ++= config.producerProps
-        props.put("metadata.broker.list", config.brokerList)
-        props.put("compression.codec", config.compressionCodec.codec.toString)
-        props.put("send.buffer.bytes", (64 * 1024).toString)
-        if (!config.isSync) {
-          props.put("producer.type", "async")
-          props.put("batch.num.messages", config.batchSize.toString)
-          props.put("queue.enqueue.timeout.ms", "-1")
-        }
-        props.put("client.id", "producer-performance")
-        props.put("request.required.acks", config.producerRequestRequiredAcks.toString)
-        props.put("request.timeout.ms", config.producerRequestTimeoutMs.toString)
-        props.put("message.send.max.retries", config.producerNumRetries.toString)
-        props.put("retry.backoff.ms", config.producerRetryBackoffMs.toString)
-        props.put("serializer.class", classOf[DefaultEncoder].getName)
-        props.put("key.serializer.class", classOf[NullEncoder[Long]].getName)
-        new OldProducer(props)
-      }
-
-    // generate the sequential message ID
-    private val SEP = ":" // message field separator
-    private val messageIdLabel = "MessageID"
-    private val threadIdLabel = "ThreadID"
-    private val topicLabel = "Topic"
-    private var leftPaddedSeqId: String = ""
-
-    private def generateMessageWithSeqId(topic: String, msgId: Long, msgSize: Int): Array[Byte] = {
-      // Each thread gets a unique range of sequential no. for its ids.
-      // Eg. 1000 msg in 10 threads => 100 msg per thread
-      // thread 0 IDs :   0 ~  99
-      // thread 1 IDs : 100 ~ 199
-      // thread 2 IDs : 200 ~ 299
-      // . . .
-      leftPaddedSeqId = String.format("%0" + seqIdNumDigit + "d", long2Long(msgId))
-
-      val msgHeader = topicLabel + SEP +
-        topic + SEP +
-        threadIdLabel + SEP +
-        threadId + SEP +
-        messageIdLabel + SEP +
-        leftPaddedSeqId + SEP
-
-      val seqMsgString = String.format("%1$-" + msgSize + "s", msgHeader).replace(' ', 'x')
-      debug(seqMsgString)
-      seqMsgString.getBytes(StandardCharsets.UTF_8)
-    }
-
-    private def generateProducerData(topic: String, messageId: Long): Array[Byte] = {
-      val msgSize = if (config.isFixedSize) config.messageSize else 1 + rand.nextInt(config.messageSize)
-      if (config.seqIdMode) {
-        val seqId = config.initialMessageId + (messagesPerThread * threadId) + messageId
-        generateMessageWithSeqId(topic, seqId, msgSize)
-      } else {
-        new Array[Byte](msgSize)
-      }
-    }
-
-    override def run {
-      var bytesSent = 0L
-      var nSends = 0
-      var i: Long = 0L
-      var message: Array[Byte] = null
-
-      while (i < messagesPerThread) {
-        try {
-          config.topics.foreach(
-            topic => {
-              message = generateProducerData(topic, i)
-              producer.send(topic, BigInteger.valueOf(i).toByteArray, message)
-              bytesSent += message.size
-              nSends += 1
-              if (config.messageSendGapMs > 0)
-                Thread.sleep(config.messageSendGapMs)
-            })
-        } catch {
-          case e: Throwable => error("Error when sending message " + new String(message, StandardCharsets.UTF_8), e)
-        }
-        i += 1
-      }
-      try {
-        producer.close()
-      } catch {
-        case e: Throwable => error("Error when closing producer", e)
-      }
-      totalBytesSent.addAndGet(bytesSent)
-      totalMessagesSent.addAndGet(nSends)
-      allDone.countDown()
-    }
-  }
-}
diff --git a/core/src/main/scala/kafka/tools/ReplayLogProducer.scala b/core/src/main/scala/kafka/tools/ReplayLogProducer.scala
deleted file mode 100644
index ca9c111163a31..0000000000000
--- a/core/src/main/scala/kafka/tools/ReplayLogProducer.scala
+++ /dev/null
@@ -1,163 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package kafka.tools
-
-import joptsimple.OptionParser
-import java.util.concurrent.CountDownLatch
-import java.util.Properties
-import kafka.consumer._
-import kafka.utils.{ToolsUtils, CommandLineUtils, Logging, ZkUtils}
-import kafka.api.OffsetRequest
-import org.apache.kafka.clients.producer.{ProducerRecord, KafkaProducer, ProducerConfig}
-import scala.collection.JavaConverters._
-
-object ReplayLogProducer extends Logging {
-
-  private val GroupId: String = "replay-log-producer"
-
-  def main(args: Array[String]) {
-    val config = new Config(args)
-
-    // if there is no group specified then avoid polluting zookeeper with persistent group data, this is a hack
-    ZkUtils.maybeDeletePath(config.zkConnect, "/consumers/" + GroupId)
-    Thread.sleep(500)
-
-    // consumer properties
-    val consumerProps = new Properties
-    consumerProps.put("group.id", GroupId)
-    consumerProps.put("zookeeper.connect", config.zkConnect)
-    consumerProps.put("consumer.timeout.ms", "10000")
-    consumerProps.put("auto.offset.reset", OffsetRequest.SmallestTimeString)
-    consumerProps.put("fetch.message.max.bytes", (1024*1024).toString)
-    consumerProps.put("socket.receive.buffer.bytes", (2 * 1024 * 1024).toString)
-    val consumerConfig = new ConsumerConfig(consumerProps)
-    val consumerConnector: ConsumerConnector = Consumer.create(consumerConfig)
-    val topicMessageStreams = consumerConnector.createMessageStreams(Predef.Map(config.inputTopic -> config.numThreads))
-    var threadList = List[ZKConsumerThread]()
-    for (streamList <- topicMessageStreams.values)
-      for (stream <- streamList)
-        threadList ::= new ZKConsumerThread(config, stream)
-
-    for (thread <- threadList)
-      thread.start
-
-    threadList.foreach(_.shutdown)
-    consumerConnector.shutdown
-  }
-
-  class Config(args: Array[String]) {
-    val parser = new OptionParser(false)
-    val zkConnectOpt = parser.accepts("zookeeper", "REQUIRED: The connection string for the zookeeper connection in the form host:port. " +
-      "Multiple URLS can be given to allow fail-over.")
-      .withRequiredArg
-      .describedAs("zookeeper url")
-      .ofType(classOf[String])
-      .defaultsTo("127.0.0.1:2181")
-    val brokerListOpt = parser.accepts("broker-list", "REQUIRED: the broker list must be specified.")
-      .withRequiredArg
-      .describedAs("hostname:port")
-      .ofType(classOf[String])
-    val inputTopicOpt = parser.accepts("inputtopic", "REQUIRED: The topic to consume from.")
-      .withRequiredArg
-      .describedAs("input-topic")
-      .ofType(classOf[String])
-    val outputTopicOpt = parser.accepts("outputtopic", "REQUIRED: The topic to produce to")
-      .withRequiredArg
-      .describedAs("output-topic")
-      .ofType(classOf[String])
-    val numMessagesOpt = parser.accepts("messages", "The number of messages to send.")
-      .withRequiredArg
-      .describedAs("count")
-      .ofType(classOf[java.lang.Integer])
-      .defaultsTo(-1)
-    val numThreadsOpt = parser.accepts("threads", "Number of sending threads.")
-      .withRequiredArg
-      .describedAs("threads")
-      .ofType(classOf[java.lang.Integer])
-      .defaultsTo(1)
-    val reportingIntervalOpt = parser.accepts("reporting-interval", "Interval at which to print progress info.")
-      .withRequiredArg
-      .describedAs("size")
-      .ofType(classOf[java.lang.Integer])
-      .defaultsTo(5000)
-    val propertyOpt = parser.accepts("property", "A mechanism to pass properties in the form key=value to the producer. " +
-      "This allows the user to override producer properties that are not exposed by the existing command line arguments")
-      .withRequiredArg
-      .describedAs("producer properties")
-      .ofType(classOf[String])
-    val syncOpt = parser.accepts("sync", "If set message send requests to the brokers are synchronously, one at a time as they arrive.")
-
-    val options = parser.parse(args : _*)
-    
-    CommandLineUtils.checkRequiredArgs(parser, options, brokerListOpt, inputTopicOpt)
-
-    val zkConnect = options.valueOf(zkConnectOpt)
-    val brokerList = options.valueOf(brokerListOpt)
-    ToolsUtils.validatePortOrDie(parser,brokerList)
-    val numMessages = options.valueOf(numMessagesOpt).intValue
-    val numThreads = options.valueOf(numThreadsOpt).intValue
-    val inputTopic = options.valueOf(inputTopicOpt)
-    val outputTopic = options.valueOf(outputTopicOpt)
-    val reportingInterval = options.valueOf(reportingIntervalOpt).intValue
-    val isSync = options.has(syncOpt)
-    val producerProps = CommandLineUtils.parseKeyValueArgs(options.valuesOf(propertyOpt).asScala)
-    producerProps.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList)
-    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")
-  }
-
-  class ZKConsumerThread(config: Config, stream: KafkaStream[Array[Byte], Array[Byte]]) extends Thread with Logging {
-    val shutdownLatch = new CountDownLatch(1)
-    val producer = new KafkaProducer[Array[Byte],Array[Byte]](config.producerProps)
-
-    override def run() {
-      info("Starting consumer thread..")
-      var messageCount: Int = 0
-      try {
-        val iter =
-          if(config.numMessages >= 0)
-            stream.slice(0, config.numMessages)
-          else
-            stream
-        for (messageAndMetadata <- iter) {
-          try {
-            val response = producer.send(new ProducerRecord[Array[Byte],Array[Byte]](config.outputTopic, null,
-                                            messageAndMetadata.timestamp, messageAndMetadata.key(), messageAndMetadata.message()))
-            if(config.isSync) {
-              response.get()
-            }
-            messageCount += 1
-          }catch {
-            case ie: Exception => error("Skipping this message", ie)
-          }
-        }
-      }catch {
-        case e: ConsumerTimeoutException => error("consumer thread timing out", e)
-      }
-      info("Sent " + messageCount + " messages")
-      shutdownLatch.countDown
-      info("thread finished execution !" )
-    }
-
-    def shutdown() {
-      shutdownLatch.await
-      producer.close
-    }
-
-  }
-}
diff --git a/core/src/main/scala/kafka/tools/ReplicaVerificationTool.scala b/core/src/main/scala/kafka/tools/ReplicaVerificationTool.scala
index 71f33683dbe22..0060f3c69f909 100644
--- a/core/src/main/scala/kafka/tools/ReplicaVerificationTool.scala
+++ b/core/src/main/scala/kafka/tools/ReplicaVerificationTool.scala
@@ -17,25 +17,36 @@
 
 package kafka.tools
 
+import java.net.SocketTimeoutException
 import java.text.SimpleDateFormat
-import java.util.Date
+import java.util
 import java.util.concurrent.CountDownLatch
-import java.util.concurrent.atomic.AtomicReference
+import java.util.concurrent.atomic.{AtomicInteger, AtomicReference}
 import java.util.regex.{Pattern, PatternSyntaxException}
+import java.util.{Date, Optional, Properties}
 
 import joptsimple.OptionParser
 import kafka.api._
-import kafka.client.ClientUtils
-import kafka.cluster.BrokerEndPoint
-import kafka.common.TopicAndPartition
-import kafka.consumer.{ConsumerConfig, SimpleConsumer, Whitelist}
-import kafka.message.{ByteBufferMessageSet, MessageSet}
+import kafka.utils.IncludeList
 import kafka.utils._
-import org.apache.kafka.common.protocol.Errors
-import org.apache.kafka.common.utils.Time
+import org.apache.kafka.clients._
+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}
+import org.apache.kafka.common.protocol.{ApiKeys, Errors}
+import org.apache.kafka.common.record.MemoryRecords
+import org.apache.kafka.common.requests.AbstractRequest.Builder
+import org.apache.kafka.common.requests.{AbstractRequest, FetchResponse, ListOffsetRequest, FetchRequest => JFetchRequest}
+import org.apache.kafka.common.serialization.StringDeserializer
+import org.apache.kafka.common.utils.{LogContext, Time}
+import org.apache.kafka.common.{Node, TopicPartition}
+
+import scala.jdk.CollectionConverters._
+import scala.collection.Seq
 
 /**
- *  For verifying the consistency among replicas.
+ * For verifying the consistency among replicas.
  *
  *  1. start a fetcher on every broker.
  *  2. each fetcher does the following
@@ -44,11 +55,11 @@ import org.apache.kafka.common.utils.Time
  *    2.3 waits for all other fetchers to finish step 2.2
  *    2.4 one of the fetchers verifies the consistency of fetched results among replicas
  *
- *  The consistency verification is up to the high watermark. The tool reports the
- *  max lag between the verified offset and the high watermark among all partitions.
+ * The consistency verification is up to the high watermark. The tool reports the
+ * max lag between the verified offset and the high watermark among all partitions.
  *
- *  If a broker goes down, the verification of the partitions on that broker is delayed
- *  until the broker is up again.
+ * If a broker goes down, the verification of the partitions on that broker is delayed
+ * until the broker is up again.
  *
  * Caveats:
  * 1. The tools needs all brokers to be up at startup time.
@@ -56,7 +67,7 @@ import org.apache.kafka.common.utils.Time
  */
 
 object ReplicaVerificationTool extends Logging {
-  val clientId= "replicaVerificationTool"
+  val clientId = "replicaVerificationTool"
   val dateFormatString = "yyyy-MM-dd HH:mm:ss,SSS"
   val dateFormat = new SimpleDateFormat(dateFormatString)
 
@@ -74,7 +85,7 @@ object ReplicaVerificationTool extends Logging {
                          .withRequiredArg
                          .describedAs("bytes")
                          .ofType(classOf[java.lang.Integer])
-                         .defaultsTo(ConsumerConfig.FetchSize)
+                         .defaultsTo(ConsumerConfig.DEFAULT_MAX_PARTITION_FETCH_BYTES)
     val maxWaitMsOpt = parser.accepts("max-wait-ms", "The max amount of time each fetch request waits.")
                          .withRequiredArg
                          .describedAs("ms")
@@ -95,22 +106,27 @@ 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)
-    val topicWhiteListFiler = new Whitelist(regex)
+    val topicWhiteListFiler = new IncludeList(regex)
 
-    try {
-      Pattern.compile(regex)
-    }
+    try Pattern.compile(regex)
     catch {
       case _: PatternSyntaxException =>
-        throw new RuntimeException(regex + " is an invalid regex.")
+        throw new RuntimeException(s"$regex is an invalid regex.")
     }
 
     val fetchSize = options.valueOf(fetchSizeOpt).intValue
@@ -120,161 +136,184 @@ object ReplicaVerificationTool extends Logging {
     // getting topic metadata
     info("Getting topic metadata...")
     val brokerList = options.valueOf(brokerListOpt)
-    ToolsUtils.validatePortOrDie(parser,brokerList)
-    val metadataTargetBrokers = ClientUtils.parseBrokerList(brokerList)
-    val topicsMetadataResponse = ClientUtils.fetchTopicMetadata(Set[String](), metadataTargetBrokers, clientId, maxWaitMs)
-    val brokerMap = topicsMetadataResponse.brokers.map(b => (b.id, b)).toMap
-    val filteredTopicMetadata = topicsMetadataResponse.topicsMetadata.filter(
-        topicMetadata => if (topicWhiteListFiler.isTopicAllowed(topicMetadata.topic, excludeInternalTopics = false))
-          true
-        else
-          false
-    )
+    ToolsUtils.validatePortOrDie(parser, brokerList)
+
+    val (topicsMetadata, brokerInfo) = {
+      val adminClient = createAdminClient(brokerList)
+      try ((listTopicsMetadata(adminClient), brokerDetails(adminClient)))
+      finally CoreUtils.swallow(adminClient.close(), this)
+    }
+
+    val filteredTopicMetadata = topicsMetadata.filter { topicMetaData =>
+      topicWhiteListFiler.isTopicAllowed(topicMetaData.name, excludeInternalTopics = false)
+    }
 
     if (filteredTopicMetadata.isEmpty) {
-      error("No topics found. " + topicWhiteListOpt + ", if specified, is either filtering out all topics or there is no topic.")
+      error(s"No topics found. $topicWhiteListOpt if specified, is either filtering out all topics or there is no topic.")
       Exit.exit(1)
     }
 
-    val topicPartitionReplicaList: Seq[TopicPartitionReplica] = filteredTopicMetadata.flatMap(
-      topicMetadataResponse =>
-        topicMetadataResponse.partitionsMetadata.flatMap(
-          partitionMetadata =>
-            partitionMetadata.replicas.map(broker =>
-              TopicPartitionReplica(topic = topicMetadataResponse.topic, partitionId = partitionMetadata.partitionId, replicaId = broker.id))
-        )
-    )
-    debug("Selected topic partitions: " + topicPartitionReplicaList)
-    val topicAndPartitionsPerBroker: Map[Int, Seq[TopicAndPartition]] = topicPartitionReplicaList.groupBy(_.replicaId)
-      .map { case (brokerId, partitions) =>
-               brokerId -> partitions.map { partition => TopicAndPartition(partition.topic, partition.partitionId) } }
-    debug("Topic partitions per broker: " + topicAndPartitionsPerBroker)
-    val expectedReplicasPerTopicAndPartition: Map[TopicAndPartition, Int] =
-          topicPartitionReplicaList.groupBy(replica => TopicAndPartition(replica.topic, replica.partitionId))
-          .map { case (topicAndPartition, replicaSet) => topicAndPartition -> replicaSet.size }
-    debug("Expected replicas per topic partition: " + expectedReplicasPerTopicAndPartition)
-    val leadersPerBroker: Map[Int, Seq[TopicAndPartition]] = filteredTopicMetadata.flatMap { topicMetadataResponse =>
-      topicMetadataResponse.partitionsMetadata.map { partitionMetadata =>
-        (TopicAndPartition(topicMetadataResponse.topic, partitionMetadata.partitionId), partitionMetadata.leader.get.id)
+    val topicPartitionReplicas = filteredTopicMetadata.flatMap { topicMetadata =>
+      topicMetadata.partitions.asScala.flatMap { partitionMetadata =>
+        partitionMetadata.replicas.asScala.map { node =>
+          TopicPartitionReplica(topic = topicMetadata.name, partitionId = partitionMetadata.partition, replicaId = node.id)
+        }
+      }
+    }
+    debug(s"Selected topic partitions: $topicPartitionReplicas")
+    val brokerToTopicPartitions = topicPartitionReplicas.groupBy(_.replicaId).map { case (brokerId, partitions) =>
+      brokerId -> partitions.map { partition => new TopicPartition(partition.topic, partition.partitionId) }
+    }
+    debug(s"Topic partitions per broker: $brokerToTopicPartitions")
+    val expectedReplicasPerTopicPartition = topicPartitionReplicas.groupBy { replica =>
+      new TopicPartition(replica.topic, replica.partitionId)
+    }.map { case (topicAndPartition, replicaSet) => topicAndPartition -> replicaSet.size }
+    debug(s"Expected replicas per topic partition: $expectedReplicasPerTopicPartition")
+
+    val topicPartitions = filteredTopicMetadata.flatMap { topicMetaData =>
+      topicMetaData.partitions.asScala.map { partitionMetadata =>
+        new TopicPartition(topicMetaData.name, partitionMetadata.partition)
       }
-    }.groupBy(_._2).mapValues(topicAndPartitionAndLeaderIds => topicAndPartitionAndLeaderIds.map { case (topicAndPartition, _) =>
-       topicAndPartition
-     })
-    debug("Leaders per broker: " + leadersPerBroker)
-
-    val replicaBuffer = new ReplicaBuffer(expectedReplicasPerTopicAndPartition,
-                                          leadersPerBroker,
-                                          topicAndPartitionsPerBroker.size,
-                                          brokerMap,
-                                          initialOffsetTime,
-                                          reportInterval)
+    }
+
+    val consumerProps = consumerConfig(brokerList)
+
+    val replicaBuffer = new ReplicaBuffer(expectedReplicasPerTopicPartition,
+      initialOffsets(topicPartitions, consumerProps, initialOffsetTime),
+      brokerToTopicPartitions.size,
+      reportInterval)
     // create all replica fetcher threads
-    val verificationBrokerId = topicAndPartitionsPerBroker.head._1
-    val fetcherThreads: Iterable[ReplicaFetcher] = topicAndPartitionsPerBroker.map {
-      case (brokerId, topicAndPartitions) =>
-        new ReplicaFetcher(name = "ReplicaFetcher-" + brokerId,
-                           sourceBroker = brokerMap(brokerId),
-                           topicAndPartitions = topicAndPartitions,
-                           replicaBuffer = replicaBuffer,
-                           socketTimeout = 30000,
-                           socketBufferSize = 256000,
-                           fetchSize = fetchSize,
-                           maxWait = maxWaitMs,
-                           minBytes = 1,
-                           doVerification = brokerId == verificationBrokerId)
+    val verificationBrokerId = brokerToTopicPartitions.head._1
+    val counter = new AtomicInteger(0)
+    val fetcherThreads = brokerToTopicPartitions.map { case (brokerId, topicPartitions) =>
+      new ReplicaFetcher(name = s"ReplicaFetcher-$brokerId",
+        sourceBroker = brokerInfo(brokerId),
+        topicPartitions = topicPartitions,
+        replicaBuffer = replicaBuffer,
+        socketTimeout = 30000,
+        socketBufferSize = 256000,
+        fetchSize = fetchSize,
+        maxWait = maxWaitMs,
+        minBytes = 1,
+        doVerification = brokerId == verificationBrokerId,
+        consumerProps,
+        fetcherId = counter.incrementAndGet())
     }
 
-    Runtime.getRuntime.addShutdownHook(new Thread() {
-      override def run() {
+    Exit.addShutdownHook("ReplicaVerificationToolShutdownHook", {
         info("Stopping all fetchers")
         fetcherThreads.foreach(_.shutdown())
-      }
     })
     fetcherThreads.foreach(_.start())
-    println(ReplicaVerificationTool.getCurrentTimeString() + ": verification process is started.")
+    println(s"${ReplicaVerificationTool.getCurrentTimeString()}: verification process is started.")
 
   }
+
+  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): Map[Int, Node] = {
+    adminClient.describeCluster.nodes.get.asScala.map(n => (n.id, n)).toMap
+  }
+
+  private def createAdminClient(brokerUrl: String): Admin = {
+    val props = new Properties()
+    props.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, brokerUrl)
+    Admin.create(props)
+  }
+
+  private def initialOffsets(topicPartitions: Seq[TopicPartition], consumerConfig: Properties,
+                             initialOffsetTime: Long): collection.Map[TopicPartition, Long] = {
+    val consumer = createConsumer(consumerConfig)
+    try {
+      if (ListOffsetRequest.LATEST_TIMESTAMP == initialOffsetTime)
+        consumer.endOffsets(topicPartitions.asJava).asScala.map { case (k, v) => k -> v.longValue }
+      else if (ListOffsetRequest.EARLIEST_TIMESTAMP == initialOffsetTime)
+        consumer.beginningOffsets(topicPartitions.asJava).asScala.map { case (k, v) => k -> v.longValue }
+      else {
+        val timestampsToSearch = topicPartitions.map(tp => tp -> (initialOffsetTime: java.lang.Long)).toMap
+        consumer.offsetsForTimes(timestampsToSearch.asJava).asScala.map { case (k, v) => k -> v.offset }
+      }
+    } finally consumer.close()
+  }
+
+  private def consumerConfig(brokerUrl: String): Properties = {
+    val properties = new Properties()
+    properties.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, brokerUrl)
+    properties.put(ConsumerConfig.GROUP_ID_CONFIG, "ReplicaVerification")
+    properties.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, classOf[StringDeserializer])
+    properties.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, classOf[StringDeserializer])
+    properties
+  }
+
+  private def createConsumer(consumerConfig: Properties): KafkaConsumer[String, String] =
+    new KafkaConsumer(consumerConfig)
 }
 
-private case class TopicPartitionReplica(topic: String,  partitionId: Int,  replicaId: Int)
+private case class TopicPartitionReplica(topic: String, partitionId: Int, replicaId: Int)
 
 private case class MessageInfo(replicaId: Int, offset: Long, nextOffset: Long, checksum: Long)
 
-private class ReplicaBuffer(expectedReplicasPerTopicAndPartition: Map[TopicAndPartition, Int],
-                            leadersPerBroker: Map[Int, Seq[TopicAndPartition]],
+private class ReplicaBuffer(expectedReplicasPerTopicPartition: collection.Map[TopicPartition, Int],
+                            initialOffsets: collection.Map[TopicPartition, Long],
                             expectedNumFetchers: Int,
-                            brokerMap: Map[Int, BrokerEndPoint],
-                            initialOffsetTime: Long,
                             reportInterval: Long) extends Logging {
-  private val fetchOffsetMap = new Pool[TopicAndPartition, Long]
-  private val messageSetCache = new Pool[TopicAndPartition, Pool[Int, FetchResponsePartitionData]]
+  private val fetchOffsetMap = new Pool[TopicPartition, Long]
+  private val recordsCache = new Pool[TopicPartition, Pool[Int, FetchResponse.PartitionData[MemoryRecords]]]
   private val fetcherBarrier = new AtomicReference(new CountDownLatch(expectedNumFetchers))
   private val verificationBarrier = new AtomicReference(new CountDownLatch(1))
   @volatile private var lastReportTime = Time.SYSTEM.milliseconds
   private var maxLag: Long = -1L
   private var offsetWithMaxLag: Long = -1L
-  private var maxLagTopicAndPartition: TopicAndPartition = null
+  private var maxLagTopicAndPartition: TopicPartition = null
   initialize()
 
-  def createNewFetcherBarrier() {
+  def createNewFetcherBarrier(): Unit = {
     fetcherBarrier.set(new CountDownLatch(expectedNumFetchers))
   }
 
-  def getFetcherBarrier() = fetcherBarrier.get()
+  def getFetcherBarrier() = fetcherBarrier.get
 
-  def createNewVerificationBarrier() {
+  def createNewVerificationBarrier(): Unit = {
     verificationBarrier.set(new CountDownLatch(1))
   }
 
-  def getVerificationBarrier() = verificationBarrier.get()
+  def getVerificationBarrier() = verificationBarrier.get
 
-  private def initialize() {
-    for (topicAndPartition <- expectedReplicasPerTopicAndPartition.keySet)
-      messageSetCache.put(topicAndPartition, new Pool[Int, FetchResponsePartitionData])
+  private def initialize(): Unit = {
+    for (topicPartition <- expectedReplicasPerTopicPartition.keySet)
+      recordsCache.put(topicPartition, new Pool[Int, FetchResponse.PartitionData[MemoryRecords]])
     setInitialOffsets()
   }
 
-  private def offsetResponseStringWithError(offsetResponse: OffsetResponse): String = {
-    offsetResponse.partitionErrorAndOffsets.filter { case (_, partitionOffsetsResponse) =>
-      partitionOffsetsResponse.error != Errors.NONE
-    }.mkString
-  }
 
-  private def setInitialOffsets() {
-    for ((brokerId, topicAndPartitions) <- leadersPerBroker) {
-      val broker = brokerMap(brokerId)
-      val consumer = new SimpleConsumer(broker.host, broker.port, 10000, 100000, ReplicaVerificationTool.clientId)
-      val initialOffsetMap: Map[TopicAndPartition, PartitionOffsetRequestInfo] =
-        topicAndPartitions.map(topicAndPartition => topicAndPartition -> PartitionOffsetRequestInfo(initialOffsetTime, 1)).toMap
-      val offsetRequest = OffsetRequest(initialOffsetMap)
-      val offsetResponse = consumer.getOffsetsBefore(offsetRequest)
-      assert(!offsetResponse.hasError, offsetResponseStringWithError(offsetResponse))
-      offsetResponse.partitionErrorAndOffsets.foreach { case (topicAndPartition, partitionOffsetResponse) =>
-        fetchOffsetMap.put(topicAndPartition, partitionOffsetResponse.offsets.head)
-      }
-    }
+  private def setInitialOffsets(): Unit = {
+    for ((tp, offset) <- initialOffsets)
+      fetchOffsetMap.put(tp, offset)
   }
 
-  def addFetchedData(topicAndPartition: TopicAndPartition, replicaId: Int, partitionData: FetchResponsePartitionData) {
-    messageSetCache.get(topicAndPartition).put(replicaId, partitionData)
+  def addFetchedData(topicAndPartition: TopicPartition, replicaId: Int, partitionData: FetchResponse.PartitionData[MemoryRecords]): Unit = {
+    recordsCache.get(topicAndPartition).put(replicaId, partitionData)
   }
 
-  def getOffset(topicAndPartition: TopicAndPartition) = {
+  def getOffset(topicAndPartition: TopicPartition) = {
     fetchOffsetMap.get(topicAndPartition)
   }
 
-  def verifyCheckSum(println: String => Unit) {
+  def verifyCheckSum(println: String => Unit): Unit = {
     debug("Begin verification")
     maxLag = -1L
-    for ((topicAndPartition, fetchResponsePerReplica) <- messageSetCache) {
-      debug("Verifying " + topicAndPartition)
-      assert(fetchResponsePerReplica.size == expectedReplicasPerTopicAndPartition(topicAndPartition),
-            "fetched " + fetchResponsePerReplica.size + " replicas for " + topicAndPartition + ", but expected "
-            + expectedReplicasPerTopicAndPartition(topicAndPartition) + " replicas")
+    for ((topicPartition, fetchResponsePerReplica) <- recordsCache) {
+      debug(s"Verifying $topicPartition")
+      assert(fetchResponsePerReplica.size == expectedReplicasPerTopicPartition(topicPartition),
+        "fetched " + fetchResponsePerReplica.size + " replicas for " + topicPartition + ", but expected "
+          + expectedReplicasPerTopicPartition(topicPartition) + " replicas")
       val recordBatchIteratorMap = fetchResponsePerReplica.map { case (replicaId, fetchResponse) =>
-        replicaId -> fetchResponse.messages.asInstanceOf[ByteBufferMessageSet].asRecords.batches.iterator
+        replicaId -> fetchResponse.records.batches.iterator
       }
-      val maxHw = fetchResponsePerReplica.values.map(_.hw).max
+      val maxHw = fetchResponsePerReplica.values.map(_.highWatermark).max
 
       // Iterate one message at a time from every replica, until high watermark is reached.
       var isMessageInAllReplicas = true
@@ -286,7 +325,7 @@ private class ReplicaBuffer(expectedReplicasPerTopicAndPartition: Map[TopicAndPa
               val batch = recordBatchIterator.next()
 
               // only verify up to the high watermark
-              if (batch.lastOffset >= fetchResponsePerReplica.get(replicaId).hw)
+              if (batch.lastOffset >= fetchResponsePerReplica.get(replicaId).highWatermark)
                 isMessageInAllReplicas = false
               else {
                 messageInfoFromFirstReplicaOpt match {
@@ -295,15 +334,15 @@ private class ReplicaBuffer(expectedReplicasPerTopicAndPartition: Map[TopicAndPa
                       MessageInfo(replicaId, batch.lastOffset, batch.nextOffset, batch.checksum))
                   case Some(messageInfoFromFirstReplica) =>
                     if (messageInfoFromFirstReplica.offset != batch.lastOffset) {
-                      println(ReplicaVerificationTool.getCurrentTimeString + ": partition " + topicAndPartition
+                      println(ReplicaVerificationTool.getCurrentTimeString() + ": partition " + topicPartition
                         + ": replica " + messageInfoFromFirstReplica.replicaId + "'s offset "
                         + messageInfoFromFirstReplica.offset + " doesn't match replica "
                         + replicaId + "'s offset " + batch.lastOffset)
                       Exit.exit(1)
                     }
                     if (messageInfoFromFirstReplica.checksum != batch.checksum)
-                      println(ReplicaVerificationTool.getCurrentTimeString + ": partition "
-                        + topicAndPartition + " has unmatched checksum at offset " + batch.lastOffset + "; replica "
+                      println(ReplicaVerificationTool.getCurrentTimeString() + ": partition "
+                        + topicPartition + " has unmatched checksum at offset " + batch.lastOffset + "; replica "
                         + messageInfoFromFirstReplica.replicaId + "'s checksum " + messageInfoFromFirstReplica.checksum
                         + "; replica " + replicaId + "'s checksum " + batch.checksum)
                 }
@@ -313,20 +352,20 @@ private class ReplicaBuffer(expectedReplicasPerTopicAndPartition: Map[TopicAndPa
           } catch {
             case t: Throwable =>
               throw new RuntimeException("Error in processing replica %d in partition %s at offset %d."
-              .format(replicaId, topicAndPartition, fetchOffsetMap.get(topicAndPartition)), t)
+                .format(replicaId, topicPartition, fetchOffsetMap.get(topicPartition)), t)
           }
         }
         if (isMessageInAllReplicas) {
           val nextOffset = messageInfoFromFirstReplicaOpt.get.nextOffset
-          fetchOffsetMap.put(topicAndPartition, nextOffset)
-          debug(expectedReplicasPerTopicAndPartition(topicAndPartition) + " replicas match at offset " +
-                nextOffset + " for " + topicAndPartition)
+          fetchOffsetMap.put(topicPartition, nextOffset)
+          debug(s"${expectedReplicasPerTopicPartition(topicPartition)} replicas match at offset " +
+            s"$nextOffset for $topicPartition")
         }
       }
-      if (maxHw - fetchOffsetMap.get(topicAndPartition) > maxLag) {
-        offsetWithMaxLag = fetchOffsetMap.get(topicAndPartition)
+      if (maxHw - fetchOffsetMap.get(topicPartition) > maxLag) {
+        offsetWithMaxLag = fetchOffsetMap.get(topicPartition)
         maxLag = maxHw - offsetWithMaxLag
-        maxLagTopicAndPartition = topicAndPartition
+        maxLagTopicAndPartition = topicPartition
       }
       fetchResponsePerReplica.clear()
     }
@@ -334,51 +373,55 @@ private class ReplicaBuffer(expectedReplicasPerTopicAndPartition: Map[TopicAndPa
     if (currentTimeMs - lastReportTime > reportInterval) {
       println(ReplicaVerificationTool.dateFormat.format(new Date(currentTimeMs)) + ": max lag is "
         + maxLag + " for partition " + maxLagTopicAndPartition + " at offset " + offsetWithMaxLag
-        + " among " + messageSetCache.size + " partitions")
+        + " among " + recordsCache.size + " partitions")
       lastReportTime = currentTimeMs
     }
   }
 }
 
-private class ReplicaFetcher(name: String, sourceBroker: BrokerEndPoint, topicAndPartitions: Iterable[TopicAndPartition],
+private class ReplicaFetcher(name: String, sourceBroker: Node, topicPartitions: Iterable[TopicPartition],
                              replicaBuffer: ReplicaBuffer, socketTimeout: Int, socketBufferSize: Int,
-                             fetchSize: Int, maxWait: Int, minBytes: Int, doVerification: Boolean)
+                             fetchSize: Int, maxWait: Int, minBytes: Int, doVerification: Boolean, consumerConfig: Properties,
+                             fetcherId: Int)
   extends ShutdownableThread(name) {
-  val simpleConsumer = new SimpleConsumer(sourceBroker.host, sourceBroker.port, socketTimeout, socketBufferSize, ReplicaVerificationTool.clientId)
-  val fetchRequestBuilder = new FetchRequestBuilder().
-          clientId(ReplicaVerificationTool.clientId).
-          replicaId(Request.DebuggingConsumerId).
-          maxWait(maxWait).
-          minBytes(minBytes)
 
-  override def doWork() {
+  private val fetchEndpoint = new ReplicaFetcherBlockingSend(sourceBroker, new ConsumerConfig(consumerConfig), new Metrics(), Time.SYSTEM, fetcherId,
+    s"broker-${Request.DebuggingConsumerId}-fetcher-$fetcherId")
+
+  override def doWork(): Unit = {
 
     val fetcherBarrier = replicaBuffer.getFetcherBarrier()
     val verificationBarrier = replicaBuffer.getVerificationBarrier()
 
-    for (topicAndPartition <- topicAndPartitions)
-      fetchRequestBuilder.addFetch(topicAndPartition.topic, topicAndPartition.partition,
-        replicaBuffer.getOffset(topicAndPartition), fetchSize)
+    val requestMap = new util.LinkedHashMap[TopicPartition, JFetchRequest.PartitionData]
+    for (topicPartition <- topicPartitions)
+      requestMap.put(topicPartition, new JFetchRequest.PartitionData(replicaBuffer.getOffset(topicPartition),
+        0L, fetchSize, Optional.empty()))
 
-    val fetchRequest = fetchRequestBuilder.build()
-    debug("Issuing fetch request " + fetchRequest)
+    val fetchRequestBuilder = JFetchRequest.Builder.
+      forReplica(ApiKeys.FETCH.latestVersion, Request.DebuggingConsumerId, maxWait, minBytes, requestMap)
 
-    var response: FetchResponse = null
+    debug("Issuing fetch request ")
+
+    var fetchResponse: FetchResponse[MemoryRecords] = null
     try {
-      response = simpleConsumer.fetch(fetchRequest)
+      val clientResponse = fetchEndpoint.sendRequest(fetchRequestBuilder)
+      fetchResponse = clientResponse.responseBody.asInstanceOf[FetchResponse[MemoryRecords]]
     } catch {
       case t: Throwable =>
-        if (!isRunning.get)
+        if (!isRunning)
           throw t
     }
 
-    if (response != null) {
-      response.data.foreach { case (topicAndPartition, partitionData) =>
-        replicaBuffer.addFetchedData(topicAndPartition, sourceBroker.id, partitionData)
+    if (fetchResponse != null) {
+      fetchResponse.responseData.forEach { (tp, partitionData) =>
+        replicaBuffer.addFetchedData(tp, sourceBroker.id, partitionData)
       }
     } else {
-      for (topicAndPartition <- topicAndPartitions)
-        replicaBuffer.addFetchedData(topicAndPartition, sourceBroker.id, new FetchResponsePartitionData(messages = MessageSet.Empty))
+      val emptyResponse = new FetchResponse.PartitionData(Errors.NONE, FetchResponse.INVALID_HIGHWATERMARK,
+        FetchResponse.INVALID_LAST_STABLE_OFFSET, FetchResponse.INVALID_LOG_START_OFFSET, null, MemoryRecords.EMPTY)
+      for (topicAndPartition <- topicPartitions)
+        replicaBuffer.addFetchedData(topicAndPartition, sourceBroker.id, emptyResponse)
     }
 
     fetcherBarrier.countDown()
@@ -402,3 +445,68 @@ private class ReplicaFetcher(name: String, sourceBroker: BrokerEndPoint, topicAn
     debug("Done verification")
   }
 }
+
+private class ReplicaFetcherBlockingSend(sourceNode: Node,
+                                         consumerConfig: ConsumerConfig,
+                                         metrics: Metrics,
+                                         time: Time,
+                                         fetcherId: Int,
+                                         clientId: String) {
+
+  private val socketTimeout: Int = consumerConfig.getInt(ConsumerConfig.REQUEST_TIMEOUT_MS_CONFIG)
+
+  private val networkClient = {
+    val logContext = new LogContext()
+    val channelBuilder = org.apache.kafka.clients.ClientUtils.createChannelBuilder(consumerConfig, time, logContext)
+    val selector = new Selector(
+      NetworkReceive.UNLIMITED,
+      consumerConfig.getLong(ConsumerConfig.CONNECTIONS_MAX_IDLE_MS_CONFIG),
+      metrics,
+      time,
+      "replica-fetcher",
+      Map("broker-id" -> sourceNode.id.toString, "fetcher-id" -> fetcherId.toString).asJava,
+      false,
+      channelBuilder,
+      logContext
+    )
+    new NetworkClient(
+      selector,
+      new ManualMetadataUpdater(),
+      clientId,
+      1,
+      0,
+      0,
+      Selectable.USE_DEFAULT_BUFFER_SIZE,
+      consumerConfig.getInt(ConsumerConfig.RECEIVE_BUFFER_CONFIG),
+      consumerConfig.getInt(ConsumerConfig.REQUEST_TIMEOUT_MS_CONFIG),
+      consumerConfig.getLong(ConsumerConfig.SOCKET_CONNECTION_SETUP_TIMEOUT_MS_CONFIG),
+      consumerConfig.getLong(ConsumerConfig.SOCKET_CONNECTION_SETUP_TIMEOUT_MAX_MS_CONFIG),
+      ClientDnsLookup.USE_ALL_DNS_IPS,
+      time,
+      false,
+      new ApiVersions,
+      logContext
+    )
+  }
+
+  def sendRequest(requestBuilder: Builder[_ <: AbstractRequest]): ClientResponse = {
+    try {
+      if (!NetworkClientUtils.awaitReady(networkClient, sourceNode, time, socketTimeout))
+        throw new SocketTimeoutException(s"Failed to connect within $socketTimeout ms")
+      else {
+        val clientRequest = networkClient.newClientRequest(sourceNode.id.toString, requestBuilder,
+          time.milliseconds(), true)
+        NetworkClientUtils.sendAndReceive(networkClient, clientRequest, time)
+      }
+    }
+    catch {
+      case e: Throwable =>
+        networkClient.close(sourceNode.id.toString)
+        throw e
+    }
+  }
+
+  def close(): Unit = {
+    networkClient.close()
+  }
+}
diff --git a/core/src/main/scala/kafka/tools/SimpleConsumerPerformance.scala b/core/src/main/scala/kafka/tools/SimpleConsumerPerformance.scala
deleted file mode 100644
index 888d462f6354f..0000000000000
--- a/core/src/main/scala/kafka/tools/SimpleConsumerPerformance.scala
+++ /dev/null
@@ -1,165 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package kafka.tools
-
-import java.net.URI
-import java.text.SimpleDateFormat
-
-import com.typesafe.scalalogging.LazyLogging
-import kafka.api.{FetchRequestBuilder, OffsetRequest, PartitionOffsetRequestInfo}
-import kafka.consumer.SimpleConsumer
-import kafka.utils._
-import kafka.common.TopicAndPartition
-import org.apache.kafka.common.utils.Time
-
-
-/**
- * Performance test for the simple consumer
- */
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.11.0.0")
-object SimpleConsumerPerformance extends LazyLogging {
-
-  def main(args: Array[String]) {
-    logger.warn("WARNING: SimpleConsumerPerformance is deprecated and will be dropped in a future release following 0.11.0.0.")
-
-    val config = new ConsumerPerfConfig(args)
-    logger.info("Starting SimpleConsumer...")
-
-    if(!config.hideHeader) {
-      if(!config.showDetailedStats)
-        println("start.time, end.time, fetch.size, data.consumed.in.MB, MB.sec, data.consumed.in.nMsg, nMsg.sec")
-      else
-        println("time, fetch.size, data.consumed.in.MB, MB.sec, data.consumed.in.nMsg, nMsg.sec")
-    }
-
-    val consumer = new SimpleConsumer(config.url.getHost, config.url.getPort, 30*1000, 2*config.fetchSize, config.clientId)
-
-    // reset to latest or smallest offset
-    val topicAndPartition = TopicAndPartition(config.topic, config.partition)
-    val request = OffsetRequest(Map(
-      topicAndPartition -> PartitionOffsetRequestInfo(if (config.fromLatest) OffsetRequest.LatestTime else OffsetRequest.EarliestTime, 1)
-      ))
-    var offset: Long = consumer.getOffsetsBefore(request).partitionErrorAndOffsets(topicAndPartition).offsets.head
-
-    val startMs = System.currentTimeMillis
-    var done = false
-    var totalBytesRead = 0L
-    var totalMessagesRead = 0L
-    var consumedInterval = 0
-    var lastReportTime: Long = startMs
-    var lastBytesRead = 0L
-    var lastMessagesRead = 0L
-    while(!done) {
-      // TODO: add in the maxWait and minBytes for performance
-      val request = new FetchRequestBuilder()
-        .clientId(config.clientId)
-        .addFetch(config.topic, config.partition, offset, config.fetchSize)
-        .build()
-      val fetchResponse = consumer.fetch(request)
-
-      var messagesRead = 0
-      var bytesRead = 0
-      val messageSet = fetchResponse.messageSet(config.topic, config.partition)
-      for (message <- messageSet) {
-        messagesRead += 1
-        bytesRead += message.message.payloadSize
-      }
-
-      if(messagesRead == 0 || totalMessagesRead > config.numMessages)
-        done = true
-      else
-        // we only did one fetch so we find the offset for the first (head) messageset
-        offset = messageSet.last.nextOffset
-
-      totalBytesRead += bytesRead
-      totalMessagesRead += messagesRead
-      consumedInterval += messagesRead
-
-      if(consumedInterval > config.reportingInterval) {
-        if(config.showDetailedStats) {
-          val reportTime = System.currentTimeMillis
-          val elapsed = (reportTime - lastReportTime)/1000.0
-          val totalMBRead = ((totalBytesRead-lastBytesRead)*1.0)/(1024*1024)
-          println("%s, %d, %.4f, %.4f, %d, %.4f".format(config.dateFormat.format(reportTime), config.fetchSize,
-            (totalBytesRead*1.0)/(1024*1024), totalMBRead/elapsed,
-            totalMessagesRead, (totalMessagesRead-lastMessagesRead)/elapsed))
-        }
-        lastReportTime = Time.SYSTEM.milliseconds
-        lastBytesRead = totalBytesRead
-        lastMessagesRead = totalMessagesRead
-        consumedInterval = 0
-      }
-    }
-    val reportTime = System.currentTimeMillis
-    val elapsed = (reportTime - startMs) / 1000.0
-
-    if(!config.showDetailedStats) {
-      val totalMBRead = (totalBytesRead*1.0)/(1024*1024)
-      println("%s, %s, %d, %.4f, %.4f, %d, %.4f".format(config.dateFormat.format(startMs),
-        config.dateFormat.format(reportTime), config.fetchSize, totalMBRead, totalMBRead/elapsed,
-        totalMessagesRead, totalMessagesRead/elapsed))
-    }
-    Exit.exit(0)
-  }
-
-  class ConsumerPerfConfig(args: Array[String]) extends PerfConfig(args) {
-    val urlOpt = parser.accepts("server", "REQUIRED: The hostname of the server to connect to.")
-                           .withRequiredArg
-                           .describedAs("kafka://hostname:port")
-                           .ofType(classOf[String])
-    val topicOpt = parser.accepts("topic", "REQUIRED: The topic to consume from.")
-      .withRequiredArg
-      .describedAs("topic")
-      .ofType(classOf[String])
-    val resetBeginningOffsetOpt = parser.accepts("from-latest", "If the consumer does not already have an established " +
-      "offset to consume from, start with the latest message present in the log rather than the earliest message.")
-    val partitionOpt = parser.accepts("partition", "The topic partition to consume from.")
-                           .withRequiredArg
-                           .describedAs("partition")
-                           .ofType(classOf[java.lang.Integer])
-                           .defaultsTo(0)
-    val fetchSizeOpt = parser.accepts("fetch-size", "REQUIRED: The fetch size to use for consumption.")
-                           .withRequiredArg
-                           .describedAs("bytes")
-                           .ofType(classOf[java.lang.Integer])
-                           .defaultsTo(1024*1024)
-    val clientIdOpt = parser.accepts("clientId", "The ID of this client.")
-                           .withRequiredArg
-                           .describedAs("clientId")
-                           .ofType(classOf[String])
-                           .defaultsTo("SimpleConsumerPerformanceClient")
-    val showDetailedStatsOpt = parser.accepts("show-detailed-stats", "If set, stats are reported for each reporting " +
-      "interval as configured by reporting-interval")
-
-    val options = parser.parse(args : _*)
-
-    CommandLineUtils.checkRequiredArgs(parser, options, topicOpt, urlOpt, numMessagesOpt)
-
-    val url = new URI(options.valueOf(urlOpt))
-    val fetchSize = options.valueOf(fetchSizeOpt).intValue
-    val fromLatest = options.has(resetBeginningOffsetOpt)
-    val partition = options.valueOf(partitionOpt).intValue
-    val topic = options.valueOf(topicOpt)
-    val numMessages = options.valueOf(numMessagesOpt).longValue
-    val reportingInterval = options.valueOf(reportingIntervalOpt).intValue
-    val showDetailedStats = options.has(showDetailedStatsOpt)
-    val dateFormat = new SimpleDateFormat(options.valueOf(dateFormatOpt))
-    val hideHeader = options.has(hideHeaderOpt)
-    val clientId = options.valueOf(clientIdOpt).toString
-  }
-}
diff --git a/core/src/main/scala/kafka/tools/SimpleConsumerShell.scala b/core/src/main/scala/kafka/tools/SimpleConsumerShell.scala
deleted file mode 100755
index da8b698f13895..0000000000000
--- a/core/src/main/scala/kafka/tools/SimpleConsumerShell.scala
+++ /dev/null
@@ -1,263 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package kafka.tools
-
-import joptsimple._
-import kafka.utils._
-import kafka.consumer._
-import kafka.client.ClientUtils
-import kafka.api.{FetchRequestBuilder, OffsetRequest, Request}
-import kafka.cluster.BrokerEndPoint
-
-import scala.collection.JavaConverters._
-import kafka.common.{MessageFormatter, TopicAndPartition}
-import org.apache.kafka.clients.consumer.ConsumerRecord
-import org.apache.kafka.common.utils.{KafkaThread, Utils}
-
-/**
- * Command line program to dump out messages to standard out using the simple consumer
- */
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.11.0.0")
-object SimpleConsumerShell extends Logging {
-
-  def UseLeaderReplica = -1
-
-  def main(args: Array[String]): Unit = {
-    warn("WARNING: SimpleConsumerShell is deprecated and will be dropped in a future release following 0.11.0.0.")
-
-    val parser = new OptionParser(false)
-    val brokerListOpt = parser.accepts("broker-list", "REQUIRED: The list of hostname and port of the server to connect to.")
-                           .withRequiredArg
-                           .describedAs("hostname:port,...,hostname:port")
-                           .ofType(classOf[String])
-    val topicOpt = parser.accepts("topic", "REQUIRED: The topic to consume from.")
-                           .withRequiredArg
-                           .describedAs("topic")
-                           .ofType(classOf[String])
-    val partitionIdOpt = parser.accepts("partition", "The partition to consume from.")
-                           .withRequiredArg
-                           .describedAs("partition")
-                           .ofType(classOf[java.lang.Integer])
-                           .defaultsTo(0)
-    val replicaIdOpt = parser.accepts("replica", "The replica id to consume from, default -1 means leader broker.")
-                           .withRequiredArg
-                           .describedAs("replica id")
-                           .ofType(classOf[java.lang.Integer])
-                           .defaultsTo(UseLeaderReplica)
-    val offsetOpt = parser.accepts("offset", "The offset id to consume from, default to -2 which means from beginning; while value -1 means from end")
-                           .withRequiredArg
-                           .describedAs("consume offset")
-                           .ofType(classOf[java.lang.Long])
-                           .defaultsTo(OffsetRequest.EarliestTime)
-    val clientIdOpt = parser.accepts("clientId", "The ID of this client.")
-                           .withRequiredArg
-                           .describedAs("clientId")
-                           .ofType(classOf[String])
-                           .defaultsTo("SimpleConsumerShell")
-    val fetchSizeOpt = parser.accepts("fetchsize", "The fetch size of each request.")
-                           .withRequiredArg
-                           .describedAs("fetchsize")
-                           .ofType(classOf[java.lang.Integer])
-                           .defaultsTo(1024 * 1024)
-    val messageFormatterOpt = parser.accepts("formatter", "The name of a class to use for formatting kafka messages for display.")
-                           .withRequiredArg
-                           .describedAs("class")
-                           .ofType(classOf[String])
-                           .defaultsTo(classOf[DefaultMessageFormatter].getName)
-    val messageFormatterArgOpt = parser.accepts("property")
-                           .withRequiredArg
-                           .describedAs("prop")
-                           .ofType(classOf[String])
-    val printOffsetOpt = parser.accepts("print-offsets", "Print the offsets returned by the iterator")
-    val maxWaitMsOpt = parser.accepts("max-wait-ms", "The max amount of time each fetch request waits.")
-                           .withRequiredArg
-                           .describedAs("ms")
-                           .ofType(classOf[java.lang.Integer])
-                           .defaultsTo(1000)
-    val maxMessagesOpt = parser.accepts("max-messages", "The number of messages to consume")
-                           .withRequiredArg
-                           .describedAs("max-messages")
-                           .ofType(classOf[java.lang.Integer])
-                           .defaultsTo(Integer.MAX_VALUE)
-    val skipMessageOnErrorOpt = parser.accepts("skip-message-on-error", "If there is an error when processing a message, " +
-        "skip it instead of halt.")
-    val noWaitAtEndOfLogOpt = parser.accepts("no-wait-at-logend",
-        "If set, when the simple consumer reaches the end of the Log, it will stop, not waiting for new produced messages")
-
-    if(args.length == 0)
-      CommandLineUtils.printUsageAndDie(parser, "A low-level tool for fetching data directly from a particular replica.")
-
-    val options = parser.parse(args : _*)
-    CommandLineUtils.checkRequiredArgs(parser, options, brokerListOpt, topicOpt)
-
-    val topic = options.valueOf(topicOpt)
-    val partitionId = options.valueOf(partitionIdOpt).intValue()
-    val replicaId = options.valueOf(replicaIdOpt).intValue()
-    var startingOffset = options.valueOf(offsetOpt).longValue
-    val fetchSize = options.valueOf(fetchSizeOpt).intValue
-    val clientId = options.valueOf(clientIdOpt).toString
-    val maxWaitMs = options.valueOf(maxWaitMsOpt).intValue()
-    val maxMessages = options.valueOf(maxMessagesOpt).intValue
-
-    val skipMessageOnError = options.has(skipMessageOnErrorOpt)
-    val printOffsets = options.has(printOffsetOpt)
-    val noWaitAtEndOfLog = options.has(noWaitAtEndOfLogOpt)
-
-    val messageFormatterClass = Class.forName(options.valueOf(messageFormatterOpt))
-    val formatterArgs = CommandLineUtils.parseKeyValueArgs(options.valuesOf(messageFormatterArgOpt).asScala)
-
-    val fetchRequestBuilder = new FetchRequestBuilder()
-                       .clientId(clientId)
-                       .replicaId(Request.DebuggingConsumerId)
-                       .maxWait(maxWaitMs)
-                       .minBytes(ConsumerConfig.MinFetchBytes)
-
-    // getting topic metadata
-    info("Getting topic metadata...")
-    val brokerList = options.valueOf(brokerListOpt)
-    ToolsUtils.validatePortOrDie(parser,brokerList)
-    val metadataTargetBrokers = ClientUtils.parseBrokerList(brokerList)
-    val topicsMetadata = ClientUtils.fetchTopicMetadata(Set(topic), metadataTargetBrokers, clientId, maxWaitMs).topicsMetadata
-    if(topicsMetadata.size != 1 || !topicsMetadata.head.topic.equals(topic)) {
-      System.err.println(("Error: no valid topic metadata for topic: %s, " + "what we get from server is only: %s").format(topic, topicsMetadata))
-      Exit.exit(1)
-    }
-
-    // validating partition id
-    val partitionsMetadata = topicsMetadata.head.partitionsMetadata
-    val partitionMetadataOpt = partitionsMetadata.find(p => p.partitionId == partitionId)
-    if (partitionMetadataOpt.isEmpty) {
-      System.err.println("Error: partition %d does not exist for topic %s".format(partitionId, topic))
-      Exit.exit(1)
-    }
-
-    // validating replica id and initializing target broker
-    var fetchTargetBroker: BrokerEndPoint = null
-    var replicaOpt: Option[BrokerEndPoint] = null
-    if (replicaId == UseLeaderReplica) {
-      replicaOpt = partitionMetadataOpt.get.leader
-      if (replicaOpt.isEmpty) {
-        System.err.println("Error: user specifies to fetch from leader for partition (%s, %d) which has not been elected yet".format(topic, partitionId))
-        Exit.exit(1)
-      }
-    }
-    else {
-      val replicasForPartition = partitionMetadataOpt.get.replicas
-      replicaOpt = replicasForPartition.find(r => r.id == replicaId)
-      if(replicaOpt.isEmpty) {
-        System.err.println("Error: replica %d does not exist for partition (%s, %d)".format(replicaId, topic, partitionId))
-        Exit.exit(1)
-      }
-    }
-    fetchTargetBroker = replicaOpt.get
-
-    // initializing starting offset
-    if(startingOffset < OffsetRequest.EarliestTime) {
-      System.err.println("Invalid starting offset: %d".format(startingOffset))
-      Exit.exit(1)
-    }
-    if (startingOffset < 0) {
-      val simpleConsumer = new SimpleConsumer(fetchTargetBroker.host,
-                                              fetchTargetBroker.port,
-                                              ConsumerConfig.SocketTimeout,
-                                              ConsumerConfig.SocketBufferSize, clientId)
-      try {
-        startingOffset = simpleConsumer.earliestOrLatestOffset(TopicAndPartition(topic, partitionId), startingOffset,
-                                                               Request.DebuggingConsumerId)
-      } catch {
-        case t: Throwable =>
-          System.err.println("Error in getting earliest or latest offset due to: " + Utils.stackTrace(t))
-          Exit.exit(1)
-      } finally {
-        if (simpleConsumer != null)
-          simpleConsumer.close()
-      }
-    }
-
-    // initializing formatter
-    val formatter = messageFormatterClass.newInstance().asInstanceOf[MessageFormatter]
-    formatter.init(formatterArgs)
-
-    val replicaString = if(replicaId > 0) "leader" else "replica"
-    info("Starting simple consumer shell to partition [%s, %d], %s [%d], host and port: [%s, %d], from offset [%d]"
-                 .format(topic, partitionId, replicaString, replicaId,
-                         fetchTargetBroker.host,
-                         fetchTargetBroker.port, startingOffset))
-    val simpleConsumer = new SimpleConsumer(fetchTargetBroker.host,
-                                            fetchTargetBroker.port,
-                                            10000, 64*1024, clientId)
-    val thread = KafkaThread.nonDaemon("kafka-simpleconsumer-shell", new Runnable() {
-      def run() {
-        var offset = startingOffset
-        var numMessagesConsumed = 0
-        try {
-          while (numMessagesConsumed < maxMessages) {
-            val fetchRequest = fetchRequestBuilder
-                    .addFetch(topic, partitionId, offset, fetchSize)
-                    .build()
-            val fetchResponse = simpleConsumer.fetch(fetchRequest)
-            val messageSet = fetchResponse.messageSet(topic, partitionId)
-            if (messageSet.validBytes <= 0 && noWaitAtEndOfLog) {
-              println("Terminating. Reached the end of partition (%s, %d) at offset %d".format(topic, partitionId, offset))
-              return
-            }
-            debug("multi fetched " + messageSet.sizeInBytes + " bytes from offset " + offset)
-            for (messageAndOffset <- messageSet if numMessagesConsumed < maxMessages) {
-              try {
-                offset = messageAndOffset.nextOffset
-                if (printOffsets)
-                  System.out.println("next offset = " + offset)
-                val message = messageAndOffset.message
-                val key = if (message.hasKey) Utils.readBytes(message.key) else null
-                val value = if (message.isNull) null else Utils.readBytes(message.payload)
-                val serializedKeySize = if (message.hasKey) key.size else -1
-                val serializedValueSize = if (message.isNull) -1 else value.size
-                formatter.writeTo(new ConsumerRecord(topic, partitionId, offset, message.timestamp,
-                  message.timestampType, message.checksum, serializedKeySize, serializedValueSize, key, value), System.out)
-                numMessagesConsumed += 1
-              } catch {
-                case e: Throwable =>
-                  if (skipMessageOnError)
-                    error("Error processing message, skipping this message: ", e)
-                  else
-                    throw e
-              }
-              if (System.out.checkError()) {
-                // This means no one is listening to our output stream any more, time to shutdown
-                System.err.println("Unable to write to standard out, closing consumer.")
-                formatter.close()
-                simpleConsumer.close()
-                Exit.exit(1)
-              }
-            }
-          }
-        } catch {
-          case e: Throwable =>
-            error("Error consuming topic, partition, replica (%s, %d, %d) with offset [%d]".format(topic, partitionId, replicaId, offset), e)
-        } finally {
-          info(s"Consumed $numMessagesConsumed messages")
-        }
-      }
-    })
-    thread.start()
-    thread.join()
-    System.out.flush()
-    formatter.close()
-    simpleConsumer.close()
-  }
-}
diff --git a/core/src/main/scala/kafka/tools/StateChangeLogMerger.scala b/core/src/main/scala/kafka/tools/StateChangeLogMerger.scala
index a3c80d17bb0a4..de711e5d02f4b 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) {
@@ -166,7 +166,7 @@ object StateChangeLogMerger extends Logging {
    */
   def getNextLine(itr: Iterator[String]): LineIterator = {
     while (itr != null && itr.hasNext) {
-      val nextLine = itr.next
+      val nextLine = itr.next()
       dateRegex.findFirstIn(nextLine).foreach { d =>
         val date = dateFormat.parse(d)
         if ((date.equals(startDate) || date.after(startDate)) && (date.equals(endDate) || date.before(endDate))) {
diff --git a/core/src/main/scala/kafka/tools/StreamsResetter.java b/core/src/main/scala/kafka/tools/StreamsResetter.java
index 55392588bb9c8..f72a3b6905400 100644
--- a/core/src/main/scala/kafka/tools/StreamsResetter.java
+++ b/core/src/main/scala/kafka/tools/StreamsResetter.java
@@ -21,20 +21,32 @@
 import joptsimple.OptionSet;
 import joptsimple.OptionSpec;
 import joptsimple.OptionSpecBuilder;
-import org.apache.kafka.clients.admin.AdminClient;
+import kafka.utils.CommandLineUtils;
+import org.apache.kafka.clients.admin.Admin;
 import org.apache.kafka.clients.admin.DeleteTopicsResult;
-import org.apache.kafka.clients.admin.KafkaAdminClient;
+import org.apache.kafka.clients.admin.DescribeConsumerGroupsOptions;
+import org.apache.kafka.clients.admin.DescribeConsumerGroupsResult;
+import org.apache.kafka.clients.admin.MemberDescription;
+import org.apache.kafka.clients.admin.RemoveMembersFromConsumerGroupOptions;
+import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerConfig;
 import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.clients.consumer.OffsetAndTimestamp;
 import org.apache.kafka.common.KafkaFuture;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.annotation.InterfaceStability;
 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;
+import java.time.Duration;
+import java.time.Instant;
 import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.LinkedList;
@@ -42,12 +54,16 @@
 import java.util.Map;
 import java.util.Properties;
 import java.util.Set;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
+import java.util.stream.Collectors;
 
 /**
- * {@link StreamsResetter} resets the processing state of a Kafka Streams application so that, for example, you can reprocess its input from scratch.
+ * {@link StreamsResetter} resets the processing state of a Kafka Streams application so that, for example,
+ * you can reprocess its input from scratch.
  * 
                      
- * This class is not part of public API. For backward compatibility, use the provided script in "bin/" instead of calling this class directly from your code.
+ * This class is not part of public API. For backward compatibility,
+ * use the provided script in "bin/" instead of calling this class directly from your code.
  * 
                      
  * Resetting the processing state of an application includes the following actions:
  * 
                        
@@ -56,14 +72,17 @@
  * 
                      1. deleting any topics created internally by Kafka Streams for this application
                        2. 
  * 
                      
  * 
                      
- * Do only use this tool if no application instance is running. Otherwise, the application will get into an invalid state and crash or produce wrong results.
+ * Do only use this tool if no application instance is running.
+ * Otherwise, the application will get into an invalid state and crash or produce wrong results.
  * 
                      
  * If you run multiple application instances, running this tool once is sufficient.
- * However, you need to call {@code KafkaStreams#cleanUp()} before re-starting any instance (to clean local state store directory).
+ * However, you need to call {@code KafkaStreams#cleanUp()} before re-starting any instance
+ * (to clean local state store directory).
  * Otherwise, your application is in an invalid state.
  * 
                      
  * User output topics will not be deleted or modified by this tool.
- * If downstream applications consume intermediate or output topics, it is the user's responsibility to adjust those applications manually if required.
+ * If downstream applications consume intermediate or output topics,
+ * it is the user's responsibility to adjust those applications manually if required.
  */
 @InterfaceStability.Unstable
 public class StreamsResetter {
@@ -74,12 +93,43 @@ public class StreamsResetter {
     private static OptionSpec applicationIdOption;
     private static OptionSpec inputTopicsOption;
     private static OptionSpec intermediateTopicsOption;
+    private static OptionSpec toOffsetOption;
+    private static OptionSpec toDatetimeOption;
+    private static OptionSpec byDurationOption;
+    private static OptionSpecBuilder toEarliestOption;
+    private static OptionSpecBuilder toLatestOption;
+    private static OptionSpec fromFileOption;
+    private static OptionSpec shiftByOption;
     private static OptionSpecBuilder dryRunOption;
+    private static OptionSpec helpOption;
+    private static OptionSpec versionOption;
+    private static OptionSpecBuilder executeOption;
     private static OptionSpec commandConfigOption;
+    private static OptionSpecBuilder forceOption;
+
+    private final static String USAGE = "This tool helps to quickly reset an application in order to reprocess "
+            + "its data from scratch.\n"
+            + "* This tool resets offsets of input topics to the earliest available offset and it skips to the end of "
+            + "intermediate topics (topics that are input and output topics, e.g., used by deprecated through() method).\n"
+            + "* This tool deletes the internal topics that were created by Kafka Streams (topics starting with "
+            + "\"-\").\n"
+            + "You do not need to specify internal topics because the tool finds them automatically.\n"
+            + "* This tool will not delete output topics (if you want to delete them, you need to do it yourself "
+            + "with the bin/kafka-topics.sh command).\n"
+            + "* This tool will not clean up the local state on the stream application instances (the persisted "
+            + "stores used to cache aggregation results).\n"
+            + "You need to call KafkaStreams#cleanUp() in your application or manually delete them from the "
+            + "directory specified by \"state.dir\" configuration (/tmp/kafka-streams/ by default).\n"
+            + "* When long session timeout has been configured, active members could take longer to get expired on the "
+            + "broker thus blocking the reset job to complete. Use the \"--force\" option could remove those left-over "
+            + "members immediately. Make sure to stop all stream applications when this option is specified "
+            + "to avoid unexpected disruptions.\n\n"
+            + "*** Important! You will get wrong output if you don't clean up the local stores after running the "
+            + "reset tool!\n\n";
 
     private OptionSet options = null;
     private final List allTopics = new LinkedList<>();
-    private boolean dryRun = false;
+
 
     public int run(final String[] args) {
         return run(args, new Properties());
@@ -87,13 +137,13 @@ public int run(final String[] args) {
 
     public int run(final String[] args,
                    final Properties config) {
-        int exitCode = EXIT_CODE_SUCCESS;
-
-        KafkaAdminClient kafkaAdminClient = null;
+        int exitCode;
 
+        Admin adminClient = null;
         try {
             parseArguments(args);
-            dryRun = options.has(dryRunOption);
+
+            final boolean dryRun = options.has(dryRunOption);
 
             final String groupId = options.valueOf(applicationIdOption);
             final Properties properties = new Properties();
@@ -102,11 +152,11 @@ public int run(final String[] args,
             }
             properties.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, options.valueOf(bootstrapServerOption));
 
-            validateNoActiveConsumers(groupId, properties);
-            kafkaAdminClient = (KafkaAdminClient) AdminClient.create(properties);
+            adminClient = Admin.create(properties);
+            maybeDeleteActiveConsumers(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----");
@@ -114,40 +164,44 @@ public int run(final String[] args,
 
             final HashMap consumerConfig = new HashMap<>(config);
             consumerConfig.putAll(properties);
-            maybeResetInputAndSeekToEndIntermediateTopicOffsets(consumerConfig);
-            maybeDeleteInternalTopics(kafkaAdminClient);
-
+            exitCode = maybeResetInputAndSeekToEndIntermediateTopicOffsets(consumerConfig, 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(60, TimeUnit.SECONDS);
+            if (adminClient != null) {
+                adminClient.close(Duration.ofSeconds(60));
             }
         }
 
         return exitCode;
     }
 
-    private void validateNoActiveConsumers(final String groupId,
-                                           final Properties properties) {
-        kafka.admin.AdminClient olderAdminClient = null;
-        try {
-            olderAdminClient = kafka.admin.AdminClient.create(properties);
-            if (!olderAdminClient.describeConsumerGroup(groupId, 0).consumers().get().isEmpty()) {
-                throw new IllegalStateException("Consumer group '" + groupId + "' is still active. "
-                                                + "Make sure to stop all running application instances before running the reset tool.");
-            }
-        } finally {
-            if (olderAdminClient != null) {
-                olderAdminClient.close();
+    private void maybeDeleteActiveConsumers(final String groupId,
+                                            final Admin adminClient)
+        throws ExecutionException, InterruptedException {
+
+        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()) {
+            if (options.has(forceOption)) {
+                System.out.println("Force deleting all active members in the group: " + groupId);
+                adminClient.removeMembersFromConsumerGroup(groupId, new RemoveMembersFromConsumerGroupOptions()).all().get();
+            } else {
+                throw new IllegalStateException("Consumer group '" + groupId + "' is still active "
+                        + "and has following members: " + members + ". "
+                        + "Make sure to stop all running application instances before running the reset tool."
+                        + " You can use option '--force' to remove active members from the group.");
             }
         }
     }
 
-    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()
@@ -164,45 +218,110 @@ private void parseArguments(final String[] args) throws IOException {
             .ofType(String.class)
             .withValuesSeparatedBy(',')
             .describedAs("list");
-        intermediateTopicsOption = optionParser.accepts("intermediate-topics", "Comma-separated list of intermediate user topics (topics used in the through() method). For these topics, the tool will skip to the end.")
+        intermediateTopicsOption = optionParser.accepts("intermediate-topics", "Comma-separated list of intermediate user topics (topics that are input and output topics, e.g., used in the deprecated through() method). For these topics, the tool will skip to the end.")
             .withRequiredArg()
             .ofType(String.class)
             .withValuesSeparatedBy(',')
             .describedAs("list");
+        toOffsetOption = optionParser.accepts("to-offset", "Reset offsets to a specific offset.")
+            .withRequiredArg()
+            .ofType(Long.class);
+        toDatetimeOption = optionParser.accepts("to-datetime", "Reset offsets to offset from datetime. Format: 'YYYY-MM-DDTHH:mm:SS.sss'")
+            .withRequiredArg()
+            .ofType(String.class);
+        byDurationOption = optionParser.accepts("by-duration", "Reset offsets to offset by duration from current timestamp. Format: 'PnDTnHnMnS'")
+            .withRequiredArg()
+            .ofType(String.class);
+        toEarliestOption = optionParser.accepts("to-earliest", "Reset offsets to earliest offset.");
+        toLatestOption = optionParser.accepts("to-latest", "Reset offsets to latest offset.");
+        fromFileOption = optionParser.accepts("from-file", "Reset offsets to values defined in CSV file.")
+            .withRequiredArg()
+            .ofType(String.class);
+        shiftByOption = optionParser.accepts("shift-by", "Reset offsets shifting current offset by 'n', where 'n' can be positive or negative")
+            .withRequiredArg()
+            .describedAs("number-of-offsets")
+            .ofType(Long.class);
         commandConfigOption = optionParser.accepts("config-file", "Property file containing configs to be passed to admin clients and embedded consumer.")
             .withRequiredArg()
             .ofType(String.class)
             .describedAs("file name");
+        forceOption = optionParser.accepts("force", "Force the removal of members of the consumer group (intended to remove stopped members if a long session timeout was used). " +
+                "Make sure to shut down all stream applications when this option is specified to avoid unexpected rebalances.");
+        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.").forHelp();
+        versionOption = optionParser.accepts("version", "Print version information and exit.").forHelp();
 
-        // TODO: deprecated in 1.0; can be removed eventually
+        // 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.");
 
         try {
             options = optionParser.parse(args);
+            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 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 maybeResetInputAndSeekToEndIntermediateTopicOffsets(final Map consumerConfig) {
+    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,
+                                                                    final boolean dryRun)
+        throws IOException, ParseException {
+
         final List inputTopics = options.valuesOf(inputTopicsOption);
         final List intermediateTopics = options.valuesOf(intermediateTopicsOption);
-
+        int topicNotFound = EXIT_CODE_SUCCESS;
 
         final List notFoundInputTopics = new ArrayList<>();
         final List notFoundIntermediateTopics = new ArrayList<>();
 
-        String groupId = options.valueOf(applicationIdOption);
+        final String groupId = options.valueOf(applicationIdOption);
 
         if (inputTopics.size() == 0 && intermediateTopics.size() == 0) {
             System.out.println("No input or intermediate topics specified. Skipping seek.");
-            return;
+            return EXIT_CODE_SUCCESS;
         }
 
         if (inputTopics.size() != 0) {
-            System.out.println("Seek-to-beginning for input topics " + inputTopics);
+            System.out.println("Reset-offsets for input topics " + inputTopics);
         }
         if (intermediateTopics.size() != 0) {
             System.out.println("Seek-to-end for intermediate topics " + intermediateTopics);
@@ -225,16 +344,42 @@ private void maybeResetInputAndSeekToEndIntermediateTopicOffsets(final Map consu
             }
         }
 
+        if (!notFoundInputTopics.isEmpty()) {
+            System.out.println("Following input topics are not found, skipping them");
+            for (final String topic : notFoundInputTopics) {
+                System.out.println("Topic: " + topic);
+            }
+            topicNotFound = EXIT_CODE_ERROR;
+        }
+
+        if (!notFoundIntermediateTopics.isEmpty()) {
+            System.out.println("Following intermediate topics are not found, skipping them");
+            for (final String topic : notFoundIntermediateTopics) {
+                System.out.println("Topic:" + topic);
+            }
+            topicNotFound = EXIT_CODE_ERROR;
+        }
+
+        // Return early if there are no topics to reset (the consumer will raise an error if we
+        // try to poll with an empty subscription)
+        if (topicsToSubscribe.isEmpty()) {
+            return topicNotFound;
+        }
+
         final Properties config = new Properties();
         config.putAll(consumerConfig);
         config.setProperty(ConsumerConfig.GROUP_ID_CONFIG, groupId);
         config.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false");
 
-        try (final KafkaConsumer client = new KafkaConsumer<>(config, new ByteArrayDeserializer(), new ByteArrayDeserializer())) {
-            client.subscribe(topicsToSubscribe);
-            client.poll(1);
+        try (final KafkaConsumer client =
+                 new KafkaConsumer<>(config, new ByteArrayDeserializer(), new ByteArrayDeserializer())) {
+
+            final Collection partitions = topicsToSubscribe.stream().map(client::partitionsFor)
+                    .flatMap(Collection::stream)
+                    .map(info -> new TopicPartition(info.topic(), info.partition()))
+                    .collect(Collectors.toList());
+            client.assign(partitions);
 
-            final Set partitions = client.assignment();
             final Set inputTopicPartitions = new HashSet<>();
             final Set intermediateTopicPartitions = new HashSet<>();
 
@@ -249,9 +394,9 @@ private void maybeResetInputAndSeekToEndIntermediateTopicOffsets(final Map consu
                 }
             }
 
-            maybeSeekToBeginning(client, inputTopicPartitions);
+            maybeReset(groupId, client, inputTopicPartitions);
 
-            maybeSeekToEnd(client, intermediateTopicPartitions);
+            maybeSeekToEnd(groupId, client, intermediateTopicPartitions);
 
             if (!dryRun) {
                 for (final TopicPartition p : partitions) {
@@ -259,64 +404,196 @@ private void maybeResetInputAndSeekToEndIntermediateTopicOffsets(final Map consu
                 }
                 client.commitSync();
             }
-
-            if (notFoundInputTopics.size() > 0) {
-                System.out.println("Following input topics are not found, skipping them");
-                for (final String topic : notFoundInputTopics) {
-                    System.out.println("Topic: " + topic);
-                }
-            }
-
-            if (notFoundIntermediateTopics.size() > 0) {
-                System.out.println("Following intermediate topics are not found, skipping them");
-                for (final String topic : notFoundIntermediateTopics) {
-                    System.out.println("Topic:" + topic);
-                }
-            }
-
-        } catch (final RuntimeException e) {
+        } catch (final IOException | ParseException e) {
             System.err.println("ERROR: Resetting offsets failed.");
             throw e;
         }
         System.out.println("Done.");
+        return topicNotFound;
     }
 
-    private void maybeSeekToEnd(final KafkaConsumer client,
-                                final Set intermediateTopicPartitions) {
-
-        final String groupId = options.valueOf(applicationIdOption);
-        final List intermediateTopics = options.valuesOf(intermediateTopicsOption);
-
+    // visible for testing
+    public void maybeSeekToEnd(final String groupId,
+                               final Consumer client,
+                               final Set intermediateTopicPartitions) {
         if (intermediateTopicPartitions.size() > 0) {
             System.out.println("Following intermediate topics offsets will be reset to end (for consumer group " + groupId + ")");
-            for (final String topic : intermediateTopics) {
-                if (allTopics.contains(topic)) {
-                    System.out.println("Topic: " + topic);
+            for (final TopicPartition topicPartition : intermediateTopicPartitions) {
+                if (allTopics.contains(topicPartition.topic())) {
+                    System.out.println("Topic: " + topicPartition.topic());
                 }
             }
-            if (!dryRun) {
-                client.seekToEnd(intermediateTopicPartitions);
+            client.seekToEnd(intermediateTopicPartitions);
+        }
+    }
+
+    private void maybeReset(final String groupId,
+                            final Consumer client,
+                            final Set inputTopicPartitions)
+        throws IOException, ParseException {
+
+        if (inputTopicPartitions.size() > 0) {
+            System.out.println("Following input topics offsets will be reset to (for consumer group " + groupId + ")");
+            if (options.has(toOffsetOption)) {
+                resetOffsetsTo(client, inputTopicPartitions, options.valueOf(toOffsetOption));
+            } else if (options.has(toEarliestOption)) {
+                client.seekToBeginning(inputTopicPartitions);
+            } else if (options.has(toLatestOption)) {
+                client.seekToEnd(inputTopicPartitions);
+            } else if (options.has(shiftByOption)) {
+                shiftOffsetsBy(client, inputTopicPartitions, options.valueOf(shiftByOption));
+            } else if (options.has(toDatetimeOption)) {
+                final String ts = options.valueOf(toDatetimeOption);
+                final long timestamp = Utils.getDateTime(ts);
+                resetToDatetime(client, inputTopicPartitions, timestamp);
+            } else if (options.has(byDurationOption)) {
+                final String duration = options.valueOf(byDurationOption);
+                resetByDuration(client, inputTopicPartitions, Duration.parse(duration));
+            } else if (options.has(fromFileOption)) {
+                final String resetPlanPath = options.valueOf(fromFileOption);
+                final Map topicPartitionsAndOffset =
+                    getTopicPartitionOffsetFromResetPlan(resetPlanPath);
+                resetOffsetsFromResetPlan(client, inputTopicPartitions, topicPartitionsAndOffset);
+            } else {
+                client.seekToBeginning(inputTopicPartitions);
+            }
+
+            for (final TopicPartition p : inputTopicPartitions) {
+                System.out.println("Topic: " + p.topic() + " Partition: " + p.partition() + " Offset: " + client.position(p));
             }
         }
     }
 
-    private void maybeSeekToBeginning(final KafkaConsumer client,
-                                      final Set inputTopicPartitions) {
+    // visible for testing
+    public void resetOffsetsFromResetPlan(final Consumer client,
+                                          final Set inputTopicPartitions,
+                                          final Map topicPartitionsAndOffset) {
+        final Map endOffsets = client.endOffsets(inputTopicPartitions);
+        final Map beginningOffsets = client.beginningOffsets(inputTopicPartitions);
 
-        final List inputTopics = options.valuesOf(inputTopicsOption);
-        final String groupId = options.valueOf(applicationIdOption);
+        final Map validatedTopicPartitionsAndOffset =
+            checkOffsetRange(topicPartitionsAndOffset, beginningOffsets, endOffsets);
 
-        if (inputTopicPartitions.size() > 0) {
-            System.out.println("Following input topics offsets will be reset to beginning (for consumer group " + groupId + ")");
-            for (final String topic : inputTopics) {
-                if (allTopics.contains(topic)) {
-                    System.out.println("Topic: " + topic);
-                }
+        for (final TopicPartition topicPartition : inputTopicPartitions) {
+            client.seek(topicPartition, validatedTopicPartitionsAndOffset.get(topicPartition));
+        }
+    }
+
+    private Map getTopicPartitionOffsetFromResetPlan(final String resetPlanPath)
+        throws IOException, ParseException {
+
+        final String resetPlanCsv = Utils.readFileAsString(resetPlanPath);
+        return parseResetPlan(resetPlanCsv);
+    }
+
+    private void resetByDuration(final Consumer client,
+                                 final Set inputTopicPartitions,
+                                 final Duration duration) {
+        resetToDatetime(client, inputTopicPartitions, Instant.now().minus(duration).toEpochMilli());
+    }
+
+    private void resetToDatetime(final Consumer client,
+                                 final Set inputTopicPartitions,
+                                 final Long timestamp) {
+        final Map topicPartitionsAndTimes = new HashMap<>(inputTopicPartitions.size());
+        for (final TopicPartition topicPartition : inputTopicPartitions) {
+            topicPartitionsAndTimes.put(topicPartition, timestamp);
+        }
+
+        final Map topicPartitionsAndOffset = client.offsetsForTimes(topicPartitionsAndTimes);
+
+        for (final TopicPartition topicPartition : inputTopicPartitions) {
+            client.seek(topicPartition, topicPartitionsAndOffset.get(topicPartition).offset());
+        }
+    }
+
+    // visible for testing
+    public void shiftOffsetsBy(final Consumer client,
+                               final Set inputTopicPartitions,
+                               final long shiftBy) {
+        final Map endOffsets = client.endOffsets(inputTopicPartitions);
+        final Map beginningOffsets = client.beginningOffsets(inputTopicPartitions);
+
+        final Map topicPartitionsAndOffset = new HashMap<>(inputTopicPartitions.size());
+        for (final TopicPartition topicPartition : inputTopicPartitions) {
+            final long position = client.position(topicPartition);
+            final long offset = position + shiftBy;
+            topicPartitionsAndOffset.put(topicPartition, offset);
+        }
+
+        final Map validatedTopicPartitionsAndOffset =
+            checkOffsetRange(topicPartitionsAndOffset, beginningOffsets, endOffsets);
+
+        for (final TopicPartition topicPartition : inputTopicPartitions) {
+            client.seek(topicPartition, validatedTopicPartitionsAndOffset.get(topicPartition));
+        }
+    }
+
+    // visible for testing
+    public void resetOffsetsTo(final Consumer client,
+                               final Set inputTopicPartitions,
+                               final Long offset) {
+        final Map endOffsets = client.endOffsets(inputTopicPartitions);
+        final Map beginningOffsets = client.beginningOffsets(inputTopicPartitions);
+
+        final Map topicPartitionsAndOffset = new HashMap<>(inputTopicPartitions.size());
+        for (final TopicPartition topicPartition : inputTopicPartitions) {
+            topicPartitionsAndOffset.put(topicPartition, offset);
+        }
+
+        final Map validatedTopicPartitionsAndOffset =
+            checkOffsetRange(topicPartitionsAndOffset, beginningOffsets, endOffsets);
+
+        for (final TopicPartition topicPartition : inputTopicPartitions) {
+            client.seek(topicPartition, validatedTopicPartitionsAndOffset.get(topicPartition));
+        }
+    }
+
+
+    private Map parseResetPlan(final String resetPlanCsv) throws ParseException {
+        final Map topicPartitionAndOffset = new HashMap<>();
+        if (resetPlanCsv == null || resetPlanCsv.isEmpty()) {
+            throw new ParseException("Error parsing reset plan CSV file. It is empty,", 0);
+        }
+
+        final String[] resetPlanCsvParts = resetPlanCsv.split("\n");
+
+        for (final String line : resetPlanCsvParts) {
+            final String[] lineParts = line.split(",");
+            if (lineParts.length != 3) {
+                throw new ParseException("Reset plan CSV file is not following the format `TOPIC,PARTITION,OFFSET`.", 0);
             }
-            if (!dryRun) {
-                client.seekToBeginning(inputTopicPartitions);
+            final String topic = lineParts[0];
+            final int partition = Integer.parseInt(lineParts[1]);
+            final long offset = Long.parseLong(lineParts[2]);
+            final TopicPartition topicPartition = new TopicPartition(topic, partition);
+            topicPartitionAndOffset.put(topicPartition, offset);
+        }
+
+        return topicPartitionAndOffset;
+    }
+
+    private Map checkOffsetRange(final Map inputTopicPartitionsAndOffset,
+                                                       final Map beginningOffsets,
+                                                       final Map endOffsets) {
+        final Map validatedTopicPartitionsOffsets = new HashMap<>();
+        for (final Map.Entry topicPartitionAndOffset : inputTopicPartitionsAndOffset.entrySet()) {
+            final long endOffset = endOffsets.get(topicPartitionAndOffset.getKey());
+            final long offset = topicPartitionAndOffset.getValue();
+            if (offset < endOffset) {
+                final long beginningOffset = beginningOffsets.get(topicPartitionAndOffset.getKey());
+                if (offset > beginningOffset) {
+                    validatedTopicPartitionsOffsets.put(topicPartitionAndOffset.getKey(), offset);
+                } else {
+                    System.out.println("New offset (" + offset + ") is lower than earliest offset. Value will be set to " + beginningOffset);
+                    validatedTopicPartitionsOffsets.put(topicPartitionAndOffset.getKey(), beginningOffset);
+                }
+            } else {
+                System.out.println("New offset (" + offset + ") is higher than latest offset. Value will be set to " + endOffset);
+                validatedTopicPartitionsOffsets.put(topicPartitionAndOffset.getKey(), endOffset);
             }
         }
+        return validatedTopicPartitionsOffsets;
     }
 
     private boolean isInputTopic(final String topic) {
@@ -327,10 +604,9 @@ private boolean isIntermediateTopic(final String topic) {
         return options.valuesOf(intermediateTopicsOption).contains(topic);
     }
 
-    private void maybeDeleteInternalTopics(final KafkaAdminClient adminClient) {
-
+    private void maybeDeleteInternalTopics(final Admin adminClient, final boolean dryRun) {
         System.out.println("Deleting all internal/auto-created topics for application " + options.valueOf(applicationIdOption));
-        List topicsToDelete = new ArrayList<>();
+        final List topicsToDelete = new ArrayList<>();
         for (final String listing : allTopics) {
             if (isInternalTopic(listing)) {
                 if (!dryRun) {
@@ -346,8 +622,9 @@ private void maybeDeleteInternalTopics(final KafkaAdminClient adminClient) {
         System.out.println("Done.");
     }
 
-    private void doDelete(final List topicsToDelete,
-                          final KafkaAdminClient adminClient) {
+    // visible for testing
+    public void doDelete(final List topicsToDelete,
+                         final Admin adminClient) {
         boolean hasDeleteErrors = false;
         final DeleteTopicsResult deleteTopicsResult = adminClient.deleteTopics(topicsToDelete);
         final Map> results = deleteTopicsResult.values();
@@ -355,7 +632,7 @@ private void doDelete(final List topicsToDelete,
         for (final Map.Entry> entry : results.entrySet()) {
             try {
                 entry.getValue().get(30, TimeUnit.SECONDS);
-            } catch (Exception e) {
+            } catch (final Exception e) {
                 System.err.println("ERROR: deleting topic " + entry.getKey());
                 e.printStackTrace(System.err);
                 hasDeleteErrors = true;
@@ -366,30 +643,22 @@ private void doDelete(final List topicsToDelete,
         }
     }
 
-
     private boolean isInternalTopic(final String topicName) {
-        return topicName.startsWith(options.valueOf(applicationIdOption) + "-")
-            && (topicName.endsWith("-changelog") || topicName.endsWith("-repartition"));
+        // 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
+        // and don't consider them as internal topics even if they follow the same naming schema.
+        // Cf. https://issues.apache.org/jira/browse/KAFKA-7930
+        return !isInputTopic(topicName) && !isIntermediateTopic(topicName) && topicName.startsWith(options.valueOf(applicationIdOption) + "-")
+               && matchesInternalTopicFormat(topicName);
     }
 
-    private void printHelp(OptionParser parser) throws IOException {
-        System.err.println("The Streams Reset Tool allows you to quickly reset an application in order to reprocess "
-                + "its data from scratch.\n"
-                + "* This tool resets offsets of input topics to the earliest available offset and it skips to the end of "
-                + "intermediate topics (topics used in the through() method).\n"
-                + "* This tool deletes the internal topics that were created by Kafka Streams (topics starting with "
-                + "\"-\").\n"
-                + "You do not need to specify internal topics because the tool finds them automatically.\n"
-                + "* This tool will not delete output topics (if you want to delete them, you need to do it yourself "
-                + "with the bin/kafka-topics.sh command).\n"
-                + "* This tool will not clean up the local state on the stream application instances (the persisted "
-                + "stores used to cache aggregation results).\n"
-                + "You need to call KafkaStreams#cleanUp() in your application or manually delete them from the "
-                + "directory specified by \"state.dir\" configuration (/tmp/kafka-streams/ by default).\n\n"
-                + "*** Important! You will get wrong output if you don't clean up the local stores after running the "
-                + "reset tool!\n\n"
-        );
-        parser.printHelpOn(System.err);
+    // visible for testing
+    public boolean matchesInternalTopicFormat(final String topicName) {
+        return topicName.endsWith("-changelog") || topicName.endsWith("-repartition")
+               || topicName.endsWith("-subscription-registration-topic")
+               || topicName.endsWith("-subscription-response-topic")
+               || topicName.matches(".+-KTABLE-FK-JOIN-SUBSCRIPTION-REGISTRATION-\\d+-topic")
+               || topicName.matches(".+-KTABLE-FK-JOIN-SUBSCRIPTION-RESPONSE-\\d+-topic");
     }
 
     public static void main(final String[] args) {
diff --git a/core/src/main/scala/kafka/tools/TestRaftRequestHandler.scala b/core/src/main/scala/kafka/tools/TestRaftRequestHandler.scala
new file mode 100644
index 0000000000000..4fad0db4f1e83
--- /dev/null
+++ b/core/src/main/scala/kafka/tools/TestRaftRequestHandler.scala
@@ -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 kafka.tools
+
+import kafka.network.RequestChannel
+import kafka.network.RequestConvertToJson
+import kafka.raft.KafkaNetworkChannel
+import kafka.server.ApiRequestHandler
+import kafka.utils.Logging
+import org.apache.kafka.common.internals.FatalExitError
+import org.apache.kafka.common.protocol.{ApiKeys, Errors}
+import org.apache.kafka.common.requests.{AbstractRequest, AbstractResponse}
+import org.apache.kafka.common.utils.Time
+
+import scala.jdk.CollectionConverters._
+
+/**
+ * Simple request handler implementation for use by [[TestRaftServer]].
+ */
+class TestRaftRequestHandler(
+  networkChannel: KafkaNetworkChannel,
+  requestChannel: RequestChannel,
+  time: Time,
+) extends ApiRequestHandler with Logging {
+
+  override 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}")
+      request.header.apiKey match {
+        case ApiKeys.VOTE
+             | ApiKeys.BEGIN_QUORUM_EPOCH
+             | ApiKeys.END_QUORUM_EPOCH
+             | ApiKeys.FETCH =>
+          val requestBody = request.body[AbstractRequest]
+          networkChannel.postInboundRequest(requestBody, response =>
+            sendResponse(request, Some(response)))
+
+        case _ => throw new IllegalArgumentException(s"Unsupported api key: ${request.header.apiKey}")
+      }
+    } catch {
+      case e: FatalExitError => throw e
+      case e: Throwable => handleError(request, e)
+    } finally {
+      // The local completion time may be set while processing the request. Only record it if it's unset.
+      if (request.apiLocalCompleteTimeNanos < 0)
+        request.apiLocalCompleteTimeNanos = time.nanoseconds
+    }
+  }
+
+  private def handleError(request: RequestChannel.Request, err: Throwable): Unit = {
+    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]}", err)
+
+    val requestBody = request.body[AbstractRequest]
+    val response = requestBody.getErrorResponse(0, err)
+    if (response == null)
+      closeConnection(request, requestBody.errorCounts(err))
+    else
+      sendResponse(request, Some(response))
+  }
+
+  private def closeConnection(request: RequestChannel.Request, errorCounts: java.util.Map[Errors, Integer]): Unit = {
+    // This case is used when the request handler has encountered an error, but the client
+    // does not expect a response (e.g. when produce request has acks set to 0)
+    requestChannel.updateErrorMetrics(request.header.apiKey, errorCounts.asScala)
+    requestChannel.sendResponse(new RequestChannel.CloseConnectionResponse(request))
+  }
+
+  private def sendResponse(request: RequestChannel.Request,
+                           responseOpt: Option[AbstractResponse]): Unit = {
+    // Update error metrics for each error code in the response including Errors.NONE
+    responseOpt.foreach(response => requestChannel.updateErrorMetrics(request.header.apiKey, response.errorCounts.asScala))
+
+    val response = responseOpt match {
+      case Some(response) =>
+        val responseSend = request.context.buildResponseSend(response)
+        val responseString =
+          if (RequestChannel.isRequestLoggingEnabled) Some(RequestConvertToJson.response(response, request.context.apiVersion))
+          else None
+        new RequestChannel.SendResponse(request, responseSend, responseString, None)
+      case None =>
+        new RequestChannel.NoOpResponse(request)
+    }
+    sendResponse(response)
+  }
+
+  private def sendResponse(response: RequestChannel.Response): Unit = {
+    requestChannel.sendResponse(response)
+  }
+
+}
diff --git a/core/src/main/scala/kafka/tools/TestRaftServer.scala b/core/src/main/scala/kafka/tools/TestRaftServer.scala
new file mode 100644
index 0000000000000..8d179dccacd02
--- /dev/null
+++ b/core/src/main/scala/kafka/tools/TestRaftServer.scala
@@ -0,0 +1,626 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package kafka.tools
+
+import java.io.File
+import java.nio.file.Files
+import java.util.concurrent.atomic.{AtomicInteger, AtomicLong}
+import java.util.concurrent.{CountDownLatch, LinkedBlockingDeque, TimeUnit}
+import java.util.{Collections, Random}
+
+import joptsimple.OptionException
+import kafka.log.{Log, LogConfig, LogManager}
+import kafka.network.SocketServer
+import kafka.raft.{KafkaMetadataLog, KafkaNetworkChannel, TimingWheelExpirationService}
+import kafka.security.CredentialProvider
+import kafka.server.{BrokerTopicStats, KafkaConfig, KafkaRequestHandlerPool, KafkaServer, LogDirFailureChannel}
+import kafka.utils.timer.SystemTimer
+import kafka.utils.{CommandDefaultOptions, CommandLineUtils, CoreUtils, Exit, KafkaScheduler, Logging, ShutdownableThread}
+import org.apache.kafka.clients.{ApiVersions, ClientDnsLookup, ManualMetadataUpdater, NetworkClient}
+import org.apache.kafka.common.config.ConfigException
+import org.apache.kafka.common.metrics.Metrics
+import org.apache.kafka.common.metrics.stats.Percentiles.BucketSizing
+import org.apache.kafka.common.metrics.stats.{Meter, Percentile, Percentiles}
+import org.apache.kafka.common.network.{ChannelBuilders, NetworkReceive, Selectable, Selector}
+import org.apache.kafka.common.protocol.Writable
+import org.apache.kafka.common.security.JaasContext
+import org.apache.kafka.common.security.scram.internals.ScramMechanism
+import org.apache.kafka.common.security.token.delegation.internals.DelegationTokenCache
+import org.apache.kafka.common.utils.{LogContext, Time, Utils}
+import org.apache.kafka.common.{TopicPartition, protocol}
+import org.apache.kafka.raft.BatchReader.Batch
+import org.apache.kafka.raft.{BatchReader, FileBasedStateStore, KafkaRaftClient, QuorumState, RaftClient, RaftConfig, RecordSerde}
+
+import scala.jdk.CollectionConverters._
+
+/**
+ * This is an experimental server which is intended for testing the performance
+ * of the Raft implementation. It uses a hard-coded `__cluster_metadata` topic.
+ */
+class TestRaftServer(
+  val config: KafkaConfig,
+  val throughput: Int,
+  val recordSize: Int
+) extends Logging {
+  import kafka.tools.TestRaftServer._
+
+  private val partition = new TopicPartition("__cluster_metadata", 0)
+  private val time = Time.SYSTEM
+  private val shutdownLatch = new CountDownLatch(1)
+
+  var socketServer: SocketServer = _
+  var credentialProvider: CredentialProvider = _
+  var tokenCache: DelegationTokenCache = _
+  var dataPlaneRequestHandlerPool: KafkaRequestHandlerPool = _
+  var scheduler: KafkaScheduler = _
+  var metrics: Metrics = _
+  var ioThread: RaftIoThread = _
+  var networkChannel: KafkaNetworkChannel = _
+  var metadataLog: KafkaMetadataLog = _
+  var workloadGenerator: RaftWorkloadGenerator = _
+
+  def startup(): Unit = {
+    val logContext = new LogContext(s"[Raft id=${config.brokerId}] ")
+
+    metrics = new Metrics()
+    scheduler = new KafkaScheduler(threads = 1)
+
+    scheduler.startup()
+
+    tokenCache = new DelegationTokenCache(ScramMechanism.mechanismNames)
+    credentialProvider = new CredentialProvider(ScramMechanism.mechanismNames, tokenCache)
+
+    socketServer = new SocketServer(config, metrics, time, credentialProvider, allowDisabledApis = true)
+    socketServer.startup(startProcessingRequests = false)
+
+    val logDirName = Log.logDirName(partition)
+    val logDir = createLogDirectory(new File(config.logDirs.head), logDirName)
+
+    val raftConfig = new RaftConfig(config.originals)
+    val metadataLog = buildMetadataLog(logDir)
+    val networkChannel = buildNetworkChannel(raftConfig, logContext)
+
+    val raftClient = buildRaftClient(
+      raftConfig,
+      metadataLog,
+      networkChannel,
+      logContext,
+      logDir
+    )
+
+    workloadGenerator = new RaftWorkloadGenerator(
+      raftClient,
+      time,
+      recordsPerSec = 20000,
+      recordSize = 256
+    )
+
+    raftClient.register(workloadGenerator)
+    raftClient.initialize()
+
+    val requestHandler = new TestRaftRequestHandler(
+      networkChannel,
+      socketServer.dataPlaneRequestChannel,
+      time
+    )
+
+    dataPlaneRequestHandlerPool = new KafkaRequestHandlerPool(
+      config.brokerId,
+      socketServer.dataPlaneRequestChannel,
+      requestHandler,
+      time,
+      config.numIoThreads,
+      s"${SocketServer.DataPlaneMetricPrefix}RequestHandlerAvgIdlePercent",
+      SocketServer.DataPlaneThreadPrefix
+    )
+
+    socketServer.startProcessingRequests(Map.empty)
+    ioThread = new RaftIoThread(raftClient)
+    ioThread.start()
+    workloadGenerator.start()
+  }
+
+  def shutdown(): Unit = {
+    if (workloadGenerator != null)
+      CoreUtils.swallow(workloadGenerator.shutdown(), this)
+    if (ioThread != null)
+      CoreUtils.swallow(ioThread.shutdown(), this)
+    if (dataPlaneRequestHandlerPool != null)
+      CoreUtils.swallow(dataPlaneRequestHandlerPool.shutdown(), this)
+    if (socketServer != null)
+      CoreUtils.swallow(socketServer.shutdown(), this)
+    if (networkChannel != null)
+      CoreUtils.swallow(networkChannel.close(), this)
+    if (scheduler != null)
+      CoreUtils.swallow(scheduler.shutdown(), this)
+    if (metrics != null)
+      CoreUtils.swallow(metrics.close(), this)
+    if (metadataLog != null)
+      CoreUtils.swallow(metadataLog.close(), this)
+
+    shutdownLatch.countDown()
+  }
+
+  def awaitShutdown(): Unit = {
+    shutdownLatch.await()
+  }
+
+  private def buildNetworkChannel(raftConfig: RaftConfig,
+                                  logContext: LogContext): KafkaNetworkChannel = {
+    val netClient = buildNetworkClient(raftConfig, logContext)
+    val clientId = s"Raft-${config.brokerId}"
+    new KafkaNetworkChannel(time, netClient, clientId,
+      raftConfig.retryBackoffMs, raftConfig.requestTimeoutMs)
+  }
+
+  private def buildMetadataLog(logDir: File): KafkaMetadataLog = {
+    if (config.logDirs.size != 1) {
+      throw new ConfigException("There must be exactly one configured log dir")
+    }
+
+    val defaultProps = KafkaServer.copyKafkaConfigToLog(config)
+    LogConfig.validateValues(defaultProps)
+    val defaultLogConfig = LogConfig(defaultProps)
+
+    val log = Log(
+      dir = logDir,
+      config = defaultLogConfig,
+      logStartOffset = 0L,
+      recoveryPoint = 0L,
+      scheduler = scheduler,
+      brokerTopicStats = new BrokerTopicStats,
+      time = time,
+      maxProducerIdExpirationMs = config.transactionalIdExpirationMs,
+      producerIdExpirationCheckIntervalMs = LogManager.ProducerIdExpirationCheckIntervalMs,
+      logDirFailureChannel = new LogDirFailureChannel(5)
+    )
+    new KafkaMetadataLog(log, partition)
+  }
+
+  private def createLogDirectory(logDir: File, logDirName: String): File = {
+    val logDirPath = logDir.getAbsolutePath
+    val dir = new File(logDirPath, logDirName)
+    Files.createDirectories(dir.toPath)
+    dir
+  }
+
+  private def buildRaftClient(raftConfig: RaftConfig,
+                              metadataLog: KafkaMetadataLog,
+                              networkChannel: KafkaNetworkChannel,
+                              logContext: LogContext,
+                              logDir: File): KafkaRaftClient[Array[Byte]] = {
+    val quorumState = new QuorumState(
+      config.brokerId,
+      raftConfig.quorumVoterIds,
+      raftConfig.electionTimeoutMs,
+      raftConfig.fetchTimeoutMs,
+      new FileBasedStateStore(new File(logDir, "quorum-state")),
+      time,
+      logContext,
+      new Random()
+    )
+
+    val expirationTimer = new SystemTimer("raft-expiration-executor")
+    val expirationService = new TimingWheelExpirationService(expirationTimer)
+    val serde = new ByteArraySerde
+
+    new KafkaRaftClient(
+      raftConfig,
+      serde,
+      networkChannel,
+      metadataLog,
+      quorumState,
+      time,
+      expirationService,
+      logContext
+    )
+  }
+
+  private def buildNetworkClient(raftConfig: RaftConfig,
+                                 logContext: LogContext): NetworkClient = {
+    val channelBuilder = ChannelBuilders.clientChannelBuilder(
+      config.interBrokerSecurityProtocol,
+      JaasContext.Type.SERVER,
+      config,
+      config.interBrokerListenerName,
+      config.saslMechanismInterBrokerProtocol,
+      time,
+      config.saslInterBrokerHandshakeRequestEnable,
+      logContext
+    )
+
+    val metricGroupPrefix = "raft-channel"
+    val collectPerConnectionMetrics = false
+
+    val selector = new Selector(
+      NetworkReceive.UNLIMITED,
+      config.connectionsMaxIdleMs,
+      metrics,
+      time,
+      metricGroupPrefix,
+      Map.empty[String, String].asJava,
+      collectPerConnectionMetrics,
+      channelBuilder,
+      logContext
+    )
+
+    val clientId = s"broker-${config.brokerId}-raft-client"
+    val maxInflightRequestsPerConnection = 1
+    val reconnectBackoffMs = 50
+    val reconnectBackoffMsMs = 500
+    val discoverBrokerVersions = false
+
+    new NetworkClient(
+      selector,
+      new ManualMetadataUpdater(),
+      clientId,
+      maxInflightRequestsPerConnection,
+      reconnectBackoffMs,
+      reconnectBackoffMsMs,
+      Selectable.USE_DEFAULT_BUFFER_SIZE,
+      config.socketReceiveBufferBytes,
+      raftConfig.requestTimeoutMs,
+      config.connectionSetupTimeoutMs,
+      config.connectionSetupTimeoutMaxMs,
+      ClientDnsLookup.USE_ALL_DNS_IPS,
+      time,
+      discoverBrokerVersions,
+      new ApiVersions,
+      logContext
+    )
+  }
+
+  class RaftWorkloadGenerator(
+    client: KafkaRaftClient[Array[Byte]],
+    time: Time,
+    recordsPerSec: Int,
+    recordSize: Int
+  ) extends ShutdownableThread(name = "raft-workload-generator")
+    with RaftClient.Listener[Array[Byte]] {
+
+    sealed trait RaftEvent
+    case class HandleClaim(epoch: Int) extends RaftEvent
+    case object HandleResign extends RaftEvent
+    case class HandleCommit(reader: BatchReader[Array[Byte]]) extends RaftEvent
+    case object Shutdown extends RaftEvent
+
+    private val eventQueue = new LinkedBlockingDeque[RaftEvent]()
+    private val stats = new WriteStats(metrics, time, printIntervalMs = 5000)
+    private val payload = new Array[Byte](recordSize)
+    private val pendingAppends = new LinkedBlockingDeque[PendingAppend]()
+    private val recordCount = new AtomicInteger(0)
+    private val throttler = new ThroughputThrottler(time, recordsPerSec)
+
+    private var claimedEpoch: Option[Int] = None
+
+    override def handleClaim(epoch: Int): Unit = {
+      eventQueue.offer(HandleClaim(epoch))
+    }
+
+    override def handleResign(): Unit = {
+      eventQueue.offer(HandleResign)
+    }
+
+    override def handleCommit(reader: BatchReader[Array[Byte]]): Unit = {
+      eventQueue.offer(HandleCommit(reader))
+    }
+
+    override def initiateShutdown(): Boolean = {
+      val initiated = super.initiateShutdown()
+      eventQueue.offer(Shutdown)
+      initiated
+    }
+
+    private def sendNow(
+      leaderEpoch: Int,
+      currentTimeMs: Long
+    ): Unit = {
+      recordCount.incrementAndGet()
+
+      val offset = client.scheduleAppend(leaderEpoch, Collections.singletonList(payload))
+      if (offset == null) {
+        time.sleep(10)
+      } else {
+        pendingAppends.offer(PendingAppend(offset, currentTimeMs))
+      }
+    }
+
+    override def doWork(): Unit = {
+      val startTimeMs = time.milliseconds()
+      val eventTimeoutMs = claimedEpoch.map { leaderEpoch =>
+        val throttleTimeMs = throttler.maybeThrottle(recordCount.get() + 1, startTimeMs)
+        if (throttleTimeMs == 0) {
+          sendNow(leaderEpoch, startTimeMs)
+        }
+        throttleTimeMs
+      }.getOrElse(Long.MaxValue)
+
+      eventQueue.poll(eventTimeoutMs, TimeUnit.MILLISECONDS) match {
+        case HandleClaim(epoch) =>
+          claimedEpoch = Some(epoch)
+          throttler.reset()
+          pendingAppends.clear()
+          recordCount.set(0)
+
+        case HandleResign =>
+          claimedEpoch = None
+          pendingAppends.clear()
+
+        case HandleCommit(reader) =>
+          try {
+            while (reader.hasNext) {
+              val batch = reader.next()
+              claimedEpoch.foreach { leaderEpoch =>
+                handleLeaderCommit(leaderEpoch, batch)
+              }
+            }
+          } finally {
+            reader.close()
+          }
+
+        case _ =>
+      }
+    }
+
+    private def handleLeaderCommit(
+      leaderEpoch: Int,
+      batch: Batch[Array[Byte]]
+    ): Unit = {
+      val batchEpoch = batch.epoch()
+      var offset = batch.baseOffset
+      val currentTimeMs = time.milliseconds()
+
+      // We are only interested in batches written during the current leader's
+      // epoch since this allows us to rely on the local clock
+      if (batchEpoch != leaderEpoch) {
+        return
+      }
+
+      for (record <- batch.records.asScala) {
+        val pendingAppend = pendingAppends.peek()
+
+        if (pendingAppend == null || pendingAppend.offset != offset) {
+          throw new IllegalStateException(s"Unexpected append at offset $offset. The " +
+            s"next offset we expected was ${pendingAppend.offset}")
+        }
+
+        pendingAppends.poll()
+        val latencyMs = math.max(0, currentTimeMs - pendingAppend.appendTimeMs).toInt
+        stats.record(latencyMs, record.length, currentTimeMs)
+        offset += 1
+      }
+    }
+
+  }
+
+  class RaftIoThread(
+    client: KafkaRaftClient[Array[Byte]]
+  ) extends ShutdownableThread(
+    name = "raft-io-thread",
+    isInterruptible = false
+  ) {
+    override def doWork(): Unit = {
+      client.poll()
+    }
+
+    override def initiateShutdown(): Boolean = {
+      if (super.initiateShutdown()) {
+        client.shutdown(5000).whenComplete { (_, exception) =>
+          if (exception != null) {
+            error("Graceful shutdown of RaftClient failed", exception)
+          } else {
+            info("Completed graceful shutdown of RaftClient")
+          }
+        }
+        true
+      } else {
+        false
+      }
+    }
+
+    override def isRunning: Boolean = {
+      client.isRunning
+    }
+  }
+
+}
+
+object TestRaftServer extends Logging {
+
+  case class PendingAppend(
+    offset: Long,
+    appendTimeMs: Long
+  ) {
+    override def toString: String = {
+      s"PendingAppend(offset=$offset, appendTimeMs=$appendTimeMs)"
+    }
+  }
+
+  private class ByteArraySerde extends RecordSerde[Array[Byte]] {
+    override def recordSize(data: Array[Byte], context: Any): Int = {
+      data.length
+    }
+
+    override def write(data: Array[Byte], context: Any, out: Writable): Unit = {
+      out.writeByteArray(data)
+    }
+
+    override def read(input: protocol.Readable, size: Int): Array[Byte] = {
+      val data = new Array[Byte](size)
+      input.readArray(data)
+      data
+    }
+  }
+
+  private class LatencyHistogram(
+    metrics: Metrics,
+    name: String,
+    group: String
+  ) {
+    private val sensor = metrics.sensor(name)
+    private val latencyP75Name = metrics.metricName(s"$name.p75", group)
+    private val latencyP99Name = metrics.metricName(s"$name.p99", group)
+    private val latencyP999Name = metrics.metricName(s"$name.p999", group)
+
+    sensor.add(new Percentiles(
+      1000,
+      250.0,
+      BucketSizing.CONSTANT,
+      new Percentile(latencyP75Name, 75),
+      new Percentile(latencyP99Name, 99),
+      new Percentile(latencyP999Name, 99.9)
+    ))
+
+    private val p75 = metrics.metric(latencyP75Name)
+    private val p99 = metrics.metric(latencyP99Name)
+    private val p999 = metrics.metric(latencyP999Name)
+
+    def record(latencyMs: Int): Unit = sensor.record(latencyMs)
+    def currentP75: Double = p75.metricValue.asInstanceOf[Double]
+    def currentP99: Double = p99.metricValue.asInstanceOf[Double]
+    def currentP999: Double = p999.metricValue.asInstanceOf[Double]
+  }
+
+  private class ThroughputMeter(
+    metrics: Metrics,
+    name: String,
+    group: String
+  ) {
+    private val sensor = metrics.sensor(name)
+    private val throughputRateName = metrics.metricName(s"$name.rate", group)
+    private val throughputTotalName = metrics.metricName(s"$name.total", group)
+
+    sensor.add(new Meter(throughputRateName, throughputTotalName))
+
+    private val rate = metrics.metric(throughputRateName)
+
+    def record(bytes: Int): Unit = sensor.record(bytes)
+    def currentRate: Double = rate.metricValue.asInstanceOf[Double]
+  }
+
+  private class ThroughputThrottler(
+    time: Time,
+    targetRecordsPerSec: Int
+  ) {
+    private val startTimeMs = new AtomicLong(time.milliseconds())
+
+    require(targetRecordsPerSec > 0)
+
+    def reset(): Unit = {
+      this.startTimeMs.set(time.milliseconds())
+    }
+
+    def maybeThrottle(
+      currentCount: Int,
+      currentTimeMs: Long
+    ): Long = {
+      val targetDurationMs = math.round(currentCount / targetRecordsPerSec.toDouble * 1000)
+      if (targetDurationMs > 0) {
+        val targetDeadlineMs = startTimeMs.get() + targetDurationMs
+        if (targetDeadlineMs > currentTimeMs) {
+          val throttleDurationMs = targetDeadlineMs - currentTimeMs
+          return throttleDurationMs
+        }
+      }
+      0
+    }
+  }
+
+  private class WriteStats(
+    metrics: Metrics,
+    time: Time,
+    printIntervalMs: Long
+  ) {
+    private var lastReportTimeMs = time.milliseconds()
+    private val latency = new LatencyHistogram(metrics, name = "commit.latency", group = "kafka.raft")
+    private val throughput = new ThroughputMeter(metrics, name = "bytes.committed", group = "kafka.raft")
+
+    def record(
+      latencyMs: Int,
+      bytes: Int,
+      currentTimeMs: Long
+    ): Unit = {
+      throughput.record(bytes)
+      latency.record(latencyMs)
+
+      if (currentTimeMs - lastReportTimeMs >= printIntervalMs) {
+        printSummary()
+        this.lastReportTimeMs = currentTimeMs
+      }
+    }
+
+    private def printSummary(): Unit = {
+      println("Throughput (bytes/second): %.2f, Latency (ms): %.1f p75 %.1f p99 %.1f p999".format(
+        throughput.currentRate,
+        latency.currentP75,
+        latency.currentP99,
+        latency.currentP999
+      ))
+    }
+  }
+
+  class TestRaftServerOptions(args: Array[String]) extends CommandDefaultOptions(args) {
+    val configOpt = parser.accepts("config", "Required configured file")
+      .withRequiredArg
+      .describedAs("filename")
+      .ofType(classOf[String])
+
+    val throughputOpt = parser.accepts("throughput",
+      "The number of records per second the leader will write to the metadata topic")
+      .withRequiredArg
+      .describedAs("records/sec")
+      .ofType(classOf[Int])
+      .defaultsTo(5000)
+
+    val recordSizeOpt = parser.accepts("record-size", "The size of each record")
+      .withRequiredArg
+      .describedAs("size in bytes")
+      .ofType(classOf[Int])
+      .defaultsTo(256)
+
+    options = parser.parse(args : _*)
+  }
+
+  def main(args: Array[String]): Unit = {
+    val opts = new TestRaftServerOptions(args)
+    try {
+      CommandLineUtils.printHelpAndExitIfNeeded(opts,
+        "Standalone raft server for performance testing")
+
+      val configFile = opts.options.valueOf(opts.configOpt)
+      val serverProps = Utils.loadProps(configFile)
+      val config = KafkaConfig.fromProps(serverProps, doLog = false)
+      val throughput = opts.options.valueOf(opts.throughputOpt)
+      val recordSize = opts.options.valueOf(opts.recordSizeOpt)
+      val server = new TestRaftServer(config, throughput, recordSize)
+
+      Exit.addShutdownHook("raft-shutdown-hook", server.shutdown())
+
+      server.startup()
+      server.awaitShutdown()
+      Exit.exit(0)
+    } catch {
+      case e: OptionException =>
+        CommandLineUtils.printUsageAndDie(opts.parser, e.getMessage)
+      case e: Throwable =>
+        fatal("Exiting Kafka due to fatal exception", e)
+        Exit.exit(1)
+    }
+  }
+
+}
diff --git a/core/src/main/scala/kafka/tools/UpdateOffsetsInZK.scala b/core/src/main/scala/kafka/tools/UpdateOffsetsInZK.scala
deleted file mode 100755
index 20f1db20efb33..0000000000000
--- a/core/src/main/scala/kafka/tools/UpdateOffsetsInZK.scala
+++ /dev/null
@@ -1,93 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package kafka.tools
-
-import kafka.consumer.{ConsumerConfig, SimpleConsumer}
-import kafka.api.{OffsetRequest, PartitionOffsetRequestInfo}
-import kafka.common.{KafkaException, TopicAndPartition}
-import kafka.utils.{Exit, Logging, ZKGroupTopicDirs, ZkUtils}
-import org.apache.kafka.common.network.ListenerName
-import org.apache.kafka.common.security.JaasUtils
-import org.apache.kafka.common.security.auth.SecurityProtocol
-import org.apache.kafka.common.utils.Utils
-
-/**
- *  A utility that updates the offset of every broker partition to the offset of earliest or latest log segment file, in ZK.
- */
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.11.0.0")
-object UpdateOffsetsInZK extends Logging {
-  val Earliest = "earliest"
-  val Latest = "latest"
-
-  def main(args: Array[String]) {
-    warn("WARNING: UpdateOffsetsInZK is deprecated and will be dropped in releases following 0.11.0.0.")
-
-    if(args.length < 3)
-      usage
-    val config = new ConsumerConfig(Utils.loadProps(args(1)))
-    val zkUtils = ZkUtils(config.zkConnect, config.zkSessionTimeoutMs,
-        config.zkConnectionTimeoutMs, JaasUtils.isZkSecurityEnabled())
-    args(0) match {
-      case Earliest => getAndSetOffsets(zkUtils, OffsetRequest.EarliestTime, config, args(2))
-      case Latest => getAndSetOffsets(zkUtils, OffsetRequest.LatestTime, config, args(2))
-      case _ => usage
-    }
-  }
-
-  private def getAndSetOffsets(zkUtils: ZkUtils, offsetOption: Long, config: ConsumerConfig, topic: String): Unit = {
-    val partitionsPerTopicMap = zkUtils.getPartitionsForTopics(List(topic))
-    var partitions: Seq[Int] = Nil
-
-    partitionsPerTopicMap.get(topic) match {
-      case Some(l) =>  partitions = l.sortWith((s,t) => s < t)
-      case _ => throw new RuntimeException("Can't find topic " + topic)
-    }
-
-    var numParts = 0
-    for (partition <- partitions) {
-      val brokerHostingPartition = zkUtils.getLeaderForPartition(topic, partition)
-
-      val broker = brokerHostingPartition match {
-        case Some(b) => b
-        case None => throw new KafkaException("Broker " + brokerHostingPartition + " is unavailable. Cannot issue " +
-          "getOffsetsBefore request")
-      }
-
-      zkUtils.getBrokerInfo(broker) match {
-        case Some(brokerInfo) =>
-          val brokerEndPoint = brokerInfo.getBrokerEndPoint(ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT))
-          val consumer = new SimpleConsumer(brokerEndPoint.host, brokerEndPoint.port, 10000, 100 * 1024, "UpdateOffsetsInZk")
-          val topicAndPartition = TopicAndPartition(topic, partition)
-          val request = OffsetRequest(Map(topicAndPartition -> PartitionOffsetRequestInfo(offsetOption, 1)))
-          val offset = consumer.getOffsetsBefore(request).partitionErrorAndOffsets(topicAndPartition).offsets.head
-          val topicDirs = new ZKGroupTopicDirs(config.groupId, topic)
-
-          println("updating partition " + partition + " with new offset: " + offset)
-          zkUtils.updatePersistentPath(topicDirs.consumerOffsetDir + "/" + partition, offset.toString)
-          numParts += 1
-        case None => throw new KafkaException("Broker information for broker id %d does not exist in ZK".format(broker))
-      }
-    }
-    println("updated the offset for " + numParts + " partitions")
-  }
-
-  private def usage() = {
-    println("USAGE: " + UpdateOffsetsInZK.getClass.getName + " [earliest | latest] consumer.properties topic")
-    Exit.exit(1)
-  }
-}
diff --git a/core/src/main/scala/kafka/tools/VerifyConsumerRebalance.scala b/core/src/main/scala/kafka/tools/VerifyConsumerRebalance.scala
deleted file mode 100644
index 164595714d20e..0000000000000
--- a/core/src/main/scala/kafka/tools/VerifyConsumerRebalance.scala
+++ /dev/null
@@ -1,139 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package kafka.tools
-
-import joptsimple.OptionParser
-import org.apache.kafka.common.security._
-import kafka.utils.{CommandLineUtils, Exit, Logging, ZKGroupTopicDirs, ZkUtils}
-
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.11.0.0")
-object VerifyConsumerRebalance extends Logging {
-  def main(args: Array[String]) {
-    val parser = new OptionParser(false)
-    warn("WARNING: VerifyConsumerRebalance is deprecated and will be dropped in a future release following 0.11.0.0.")
-
-    val zkConnectOpt = parser.accepts("zookeeper.connect", "ZooKeeper connect string.").
-      withRequiredArg().defaultsTo("localhost:2181").ofType(classOf[String])
-    val groupOpt = parser.accepts("group", "Consumer group.").
-      withRequiredArg().ofType(classOf[String])
-    parser.accepts("help", "Print this message.")
-    
-    if(args.length == 0)
-      CommandLineUtils.printUsageAndDie(parser, "Validate that all partitions have a consumer for a given consumer group.")
-
-    val options = parser.parse(args : _*)
-
-    if (options.has("help")) {
-      parser.printHelpOn(System.out)
-      Exit.exit(0)
-    }
-
-    CommandLineUtils.checkRequiredArgs(parser, options, groupOpt)
-
-    val zkConnect = options.valueOf(zkConnectOpt)
-    val group = options.valueOf(groupOpt)
-
-    var zkUtils: ZkUtils = null
-    try {
-      zkUtils = ZkUtils(zkConnect,
-                        30000,
-                        30000, 
-                        JaasUtils.isZkSecurityEnabled())
-
-      debug("zkConnect = %s; group = %s".format(zkConnect, group))
-
-      // check if the rebalancing operation succeeded.
-      try {
-        if(validateRebalancingOperation(zkUtils, group))
-          println("Rebalance operation successful !")
-        else
-          println("Rebalance operation failed !")
-      } catch {
-        case e2: Throwable => error("Error while verifying current rebalancing operation", e2)
-      }
-    }
-    finally {
-      if (zkUtils != null)
-        zkUtils.close()
-    }
-  }
-
-  private def validateRebalancingOperation(zkUtils: ZkUtils, group: String): Boolean = {
-    info("Verifying rebalancing operation for consumer group " + group)
-    var rebalanceSucceeded: Boolean = true
-    /**
-     * A successful rebalancing operation would select an owner for each available partition
-     * This means that for each partition registered under /brokers/topics/[topic]/[broker-id], an owner exists
-     * under /consumers/[consumer_group]/owners/[topic]/[broker_id-partition_id]
-     */
-    val consumersPerTopicMap = zkUtils.getConsumersPerTopic(group, excludeInternalTopics = false)
-    val partitionsPerTopicMap = zkUtils.getPartitionsForTopics(consumersPerTopicMap.keySet.toSeq)
-
-    partitionsPerTopicMap.foreach { case (topic, partitions) =>
-      val topicDirs = new ZKGroupTopicDirs(group, topic)
-      info("Alive partitions for topic %s are %s ".format(topic, partitions.toString))
-      info("Alive consumers for topic %s => %s ".format(topic, consumersPerTopicMap.get(topic)))
-      val partitionsWithOwners = zkUtils.getChildrenParentMayNotExist(topicDirs.consumerOwnerDir)
-      if(partitionsWithOwners.isEmpty) {
-        error("No owners for any partitions for topic " + topic)
-        rebalanceSucceeded = false
-      }
-      debug("Children of " + topicDirs.consumerOwnerDir + " = " + partitionsWithOwners.toString)
-      val consumerIdsForTopic = consumersPerTopicMap.get(topic)
-
-      // for each available partition for topic, check if an owner exists
-      partitions.foreach { partition =>
-        // check if there is a node for [partition]
-        if(!partitionsWithOwners.contains(partition.toString)) {
-          error("No owner for partition [%s,%d]".format(topic, partition))
-          rebalanceSucceeded = false
-        }
-        // try reading the partition owner path for see if a valid consumer id exists there
-        val partitionOwnerPath = topicDirs.consumerOwnerDir + "/" + partition
-        val partitionOwner = zkUtils.readDataMaybeNull(partitionOwnerPath)._1 match {
-          case Some(m) => m
-          case None => null
-        }
-        if(partitionOwner == null) {
-          error("No owner for partition [%s,%d]".format(topic, partition))
-          rebalanceSucceeded = false
-        }
-        else {
-          // check if the owner is a valid consumer id
-          consumerIdsForTopic match {
-            case Some(consumerIds) =>
-              if(!consumerIds.map(c => c.toString).contains(partitionOwner)) {
-                error(("Owner %s for partition [%s,%d] is not a valid member of consumer " +
-                  "group %s").format(partitionOwner, topic, partition, group))
-                rebalanceSucceeded = false
-              }
-              else
-                info("Owner of partition [%s,%d] is %s".format(topic, partition, partitionOwner))
-            case None => {
-              error("No consumer ids registered for topic " + topic)
-              rebalanceSucceeded = false
-            }
-          }
-        }
-      }
-
-    }
-
-    rebalanceSucceeded
-  }
-}
diff --git a/core/src/main/scala/kafka/tools/ZooKeeperMainWrapper.scala b/core/src/main/scala/kafka/tools/ZooKeeperMainWrapper.scala
deleted file mode 100644
index 14d2ceb8a647e..0000000000000
--- a/core/src/main/scala/kafka/tools/ZooKeeperMainWrapper.scala
+++ /dev/null
@@ -1,40 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package kafka.tools
-
-import 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)
-  }
-}
-
-/**
- * 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.
- */
-object ZooKeeperMainWrapper {
-
-  def main(args: Array[String]): Unit = {
-    val main: ZooKeeperMainWrapper = new ZooKeeperMainWrapper(args)
-    main.runCmd()
-  }
-}
diff --git a/core/src/main/scala/kafka/utils/CommandDefaultOptions.scala b/core/src/main/scala/kafka/utils/CommandDefaultOptions.scala
new file mode 100644
index 0000000000000..2cdb408b4bb8c
--- /dev/null
+++ b/core/src/main/scala/kafka/utils/CommandDefaultOptions.scala
@@ -0,0 +1,27 @@
+/**
+  * Licensed to the Apache Software Foundation (ASF) under one or more
+  * contributor license agreements.  See the NOTICE file distributed with
+  * this work for additional information regarding copyright ownership.
+  * The ASF licenses this file to You under the Apache License, Version 2.0
+  * (the "License"); you may not use this file except in compliance with
+  * the License.  You may obtain a copy of the License at
+  *
+  * http://www.apache.org/licenses/LICENSE-2.0
+  *
+  * Unless required by applicable law or agreed to in writing, software
+  * distributed under the License is distributed on an "AS IS" BASIS,
+  * WITHOUT WARRANTIES 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 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.").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 edf473e3a708f..80726ce06b599 100644
--- a/core/src/main/scala/kafka/utils/CommandLineUtils.scala
+++ b/core/src/main/scala/kafka/utils/CommandLineUtils.scala
@@ -5,7 +5,7 @@
  * The ASF licenses this file to You under the Apache License, Version 2.0
  * (the "License"); you may not use this file except in compliance with
  * the License.  You may obtain a copy of the License at
- * 
+ *
  *    http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing, software
@@ -16,21 +16,55 @@
  */
  package kafka.utils
 
-import joptsimple.{OptionSpec, OptionSet, OptionParser}
-import scala.collection.Set
 import java.util.Properties
 
- /**
- * Helper functions for dealing with command line utilities
- */
+import joptsimple.{OptionParser, OptionSet, OptionSpec}
+
+import scala.collection.Set
+
+/**
+  * Helper functions for dealing with command line utilities
+  */
 object CommandLineUtils extends Logging {
+  /**
+    * Check if there are no options or `--help` option from command line
+    *
+    * @param commandOpts Acceptable options for a command
+    * @return true on matching the help check condition
+    */
+  def isPrintHelpNeeded(commandOpts: CommandDefaultOptions): Boolean = {
+    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, 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
+    */
+  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))
+      if (!options.has(arg))
         printUsageAndDie(parser, "Missing required argument \"" + arg + "\"")
     }
   }
@@ -38,11 +72,24 @@ object CommandLineUtils extends Logging {
   /**
    * Check that none of the listed options are present
    */
-  def checkInvalidArgs(parser: OptionParser, options: OptionSet, usedOption: OptionSpec[_], invalidOptions: Set[OptionSpec[_]]) {
-    if(options.has(usedOption)) {
-      for(arg <- invalidOptions) {
-        if(options.has(arg))
-          printUsageAndDie(parser, "Option \"" + usedOption + "\" can't be used with option\"" + arg + "\"")
+  def checkInvalidArgs(parser: OptionParser, options: OptionSet, usedOption: OptionSpec[_], invalidOptions: Set[OptionSpec[_]]): Unit = {
+    if (options.has(usedOption)) {
+      for (arg <- invalidOptions) {
+        if (options.has(arg))
+          printUsageAndDie(parser, "Option \"" + usedOption + "\" can't be used with option \"" + arg + "\"")
+      }
+    }
+  }
+
+  /**
+    * 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[_]],
+                         trailingAdditionalMessage: Option[String] = None): Unit = {
+    if (usedOptions.count(options.has) == usedOptions.size) {
+      for (arg <- invalidOptions) {
+        if (options.has(arg))
+          printUsageAndDie(parser, "Option combination \"" + usedOptions.mkString(",") + "\" can't be used with option \"" + arg + "\"" + trailingAdditionalMessage.getOrElse(""))
       }
     }
   }
@@ -56,24 +103,43 @@ 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
    */
   def parseKeyValueArgs(args: Iterable[String], acceptMissingValue: Boolean = true): Properties = {
-    val splits = args.map(_ split "=").filterNot(_.length == 0)
+    val splits = args.map(_.split("=", 2)).filterNot(_.length == 0)
 
     val props = new Properties
     for (a <- splits) {
-      if (a.length == 1) {
+      if (a.length == 1 || (a.length == 2 && a(1).isEmpty())) {
         if (acceptMissingValue) props.put(a(0), "")
         else throw new IllegalArgumentException(s"Missing value for key ${a(0)}")
       }
-      else if (a.length == 2) props.put(a(0), a(1))
-      else {
-        System.err.println("Invalid command line properties: " + args.mkString(" "))
-        Exit.exit(1)
-      }
+      else props.put(a(0), a(1))
     }
     props
   }
+
+  /**
+    * Merge the options into {@code props} for key {@code key}, with the following precedence, from high to low:
+    * 1) if {@code spec} is specified on {@code options} explicitly, use the value;
+    * 2) if {@code props} already has {@code key} set, keep it;
+    * 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]): Unit = {
+    if (options.has(spec) || !props.containsKey(key)) {
+      val value = options.valueOf(spec)
+      if (value == null)
+        props.remove(key)
+      else
+        props.put(key, value.toString)
+    }
+  }
 }
diff --git a/core/src/main/scala/kafka/utils/CoreUtils.scala b/core/src/main/scala/kafka/utils/CoreUtils.scala
index efd4d1ebd9f73..ff58a871881df 100755
--- a/core/src/main/scala/kafka/utils/CoreUtils.scala
+++ b/core/src/main/scala/kafka/utils/CoreUtils.scala
@@ -22,17 +22,21 @@ import java.nio._
 import java.nio.channels._
 import java.util.concurrent.locks.{Lock, ReadWriteLock}
 import java.lang.management._
-import java.util.{Properties, UUID}
+import java.util.{Base64, Properties, UUID}
+
+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
-import org.apache.kafka.common.utils.{Base64, KafkaThread, Utils}
+import org.apache.kafka.common.utils.Utils
 import org.slf4j.event.Level
 
+import scala.annotation.nowarn
+
 /**
  * General helper functions!
  *
@@ -44,34 +48,14 @@ import org.slf4j.event.Level
  * 2. It is the most general possible utility, not just the thing you needed in one particular place
  * 3. You have tests for it if it is nontrivial in any way
  */
-object CoreUtils extends Logging {
-
-  /**
-   * Return the smallest element in `traversable` 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)
+object CoreUtils {
+  private val logger = Logger(getClass)
 
   /**
-   * Wrap the given function in a java.lang.Runnable
-   * @param fun A function
-   * @return A Runnable that just executes the function
+   * Return the smallest element in `iterable` if it is not empty. Otherwise return `ifEmpty`.
    */
-  def runnable(fun: => Unit): Runnable =
-    new Runnable {
-      def run() = fun
-    }
-
-  /**
-    * Create a thread
-    *
-    * @param name The name of the thread
-    * @param daemon Whether the thread should block JVM shutdown
-    * @param fun The function to execute in the thread
-    * @return The unstarted thread
-    */
-  def newThread(name: String, daemon: Boolean)(fun: => Unit): Thread =
-    new KafkaThread(name, runnable(fun), daemon)
+  def min[A, B >: A](iterable: Iterable[A], ifEmpty: A)(implicit cmp: Ordering[B]): A =
+    if (iterable.isEmpty) ifEmpty else iterable.min(cmp)
 
   /**
     * Do the given action and log any exceptions thrown without rethrowing them.
@@ -80,7 +64,7 @@ object CoreUtils extends Logging {
     * @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 {
@@ -100,6 +84,30 @@ object CoreUtils extends Logging {
    */
   def delete(files: Seq[String]): Unit = files.foreach(f => Utils.delete(new File(f)))
 
+  /**
+   * Invokes every function in `all` even if one or more functions throws an exception.
+   *
+   * If any of the functions throws an exception, the first one will be rethrown at the end with subsequent exceptions
+   * added as suppressed exceptions.
+   */
+  // Note that this is a generalised version of `Utils.closeAll`. We could potentially make it more general by
+  // changing the signature to `def tryAll[R](all: Seq[() => R]): Seq[R]`
+  def tryAll(all: Seq[() => Unit]): Unit = {
+    var exception: Throwable = null
+    all.foreach { element =>
+      try element.apply()
+      catch {
+        case e: Throwable =>
+          if (exception != null)
+            exception.addSuppressed(e)
+          else
+            exception = e
+      }
+    }
+    if (exception != null)
+      throw exception
+  }
+
   /**
    * Register the given mbean with the platform mbean server,
    * unregistering any mbean that was there before. Note,
@@ -115,16 +123,15 @@ object CoreUtils extends Logging {
       val mbs = ManagementFactory.getPlatformMBeanServer()
       mbs synchronized {
         val objName = new ObjectName(name)
-        if(mbs.isRegistered(objName))
+        if (mbs.isRegistered(objName))
           mbs.unregisterMBean(objName)
         mbs.registerMBean(mbean, objName)
         true
       }
     } catch {
-      case e: Exception => {
-        error("Failed to register Mbean " + name, e)
+      case e: Exception =>
+        logger.error(s"Failed to register Mbean $name", e)
         false
-      }
     }
   }
 
@@ -132,11 +139,11 @@ object CoreUtils extends Logging {
    * 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)
-      if(mbs.isRegistered(objName))
+      if (mbs.isRegistered(objName))
         mbs.unregisterMBean(objName)
     }
   }
@@ -148,7 +155,7 @@ object CoreUtils extends Logging {
   def read(channel: ReadableByteChannel, buffer: ByteBuffer): Int = {
     channel.read(buffer) match {
       case -1 => throw new EOFException("Received -1 when reading from channel, socket has likely been closed.")
-      case n: Int => n
+      case n => n
     }
   }
 
@@ -174,11 +181,10 @@ object CoreUtils extends Logging {
    * Whitespace surrounding the comma will be removed.
    */
   def parseCsvList(csvList: String): Seq[String] = {
-    if(csvList == null || csvList.isEmpty)
+    if (csvList == null || csvList.isEmpty)
       Seq.empty[String]
-    else {
+    else
       csvList.split("\\s*,\\s*").filter(v => !v.equals(""))
-    }
   }
 
   /**
@@ -233,35 +239,10 @@ object CoreUtils extends Logging {
 
   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
    */
-  def duplicates[T](s: Traversable[T]): Iterable[T] = {
+  def duplicates[T](s: Iterable[T]): Iterable[T] = {
     s.groupBy(identity)
       .map { case (k, l) => (k, l.size)}
       .filter { case (_, l) => l > 1 }
@@ -269,14 +250,20 @@ object CoreUtils extends Logging {
   }
 
   def listenerListToEndPoints(listeners: String, securityProtocolMap: Map[ListenerName, SecurityProtocol]): Seq[EndPoint] = {
+    listenerListToEndPoints(listeners, securityProtocolMap, true)
+  }
+
+  def listenerListToEndPoints(listeners: String, securityProtocolMap: Map[ListenerName, SecurityProtocol], requireDistinctPorts: Boolean): Seq[EndPoint] = {
     def validate(endPoints: Seq[EndPoint]): Unit = {
       // filter port 0 for unit tests
       val portsExcludingZero = endPoints.map(_.port).filter(_ != 0)
-      val distinctPorts = portsExcludingZero.distinct
       val distinctListenerNames = endPoints.map(_.listenerName).distinct
 
-      require(distinctPorts.size == portsExcludingZero.size, s"Each listener must have a different port, listeners: $listeners")
       require(distinctListenerNames.size == endPoints.size, s"Each listener must have a different name, listeners: $listeners")
+      if (requireDistinctPorts) {
+        val distinctPorts = portsExcludingZero.distinct
+        require(distinctPorts.size == portsExcludingZero.size, s"Each listener must have a different port, listeners: $listeners")
+      }
     }
 
     val endPoints = try {
@@ -292,7 +279,7 @@ object CoreUtils extends Logging {
 
   def generateUuidAsBase64(): String = {
     val uuid = UUID.randomUUID()
-    Base64.urlEncoderNoPadding.encodeToString(getBytesFromUuid(uuid))
+    Base64.getUrlEncoder.withoutPadding.encodeToString(getBytesFromUuid(uuid))
   }
 
   def getBytesFromUuid(uuid: UUID): Array[Byte] = {
@@ -319,9 +306,8 @@ object CoreUtils extends Logging {
    * may be invoked more than once if multiple threads attempt to insert a key at the same
    * time, but the same inserted value will be returned to all threads.
    *
-   * In Scala 2.12, `ConcurrentMap.getOrElse` has the same behaviour as this method, but that
-   * is not the case in Scala 2.11. We can remove this method once we drop support for Scala
-   * 2.11.
+   * In Scala 2.12, `ConcurrentMap.getOrElse` has the same behaviour as this method, but JConcurrentMapWrapper that
+   * wraps Java maps does not.
    */
   def atomicGetOrUpdate[K, V](map: concurrent.Map[K, V], key: K, createValue: => V): V = {
     map.get(key) match {
@@ -331,4 +317,12 @@ object CoreUtils extends Logging {
         map.putIfAbsent(key, value).getOrElse(value)
     }
   }
+
+  @nowarn("cat=unused") // see below for explanation
+  def groupMapReduce[T, K, B](elements: Iterable[T])(key: T => K)(f: T => B)(reduce: (B, B) => B): Map[K, B] = {
+    // required for Scala 2.12 compatibility, unused in Scala 2.13 and hence we need to suppress the unused warning
+    import scala.collection.compat._
+    elements.groupMapReduce(key)(f)(reduce)
+  }
+
 }
diff --git a/core/src/main/scala/kafka/utils/DelayedItem.scala b/core/src/main/scala/kafka/utils/DelayedItem.scala
index 9d92b97e5807b..cfb87719a0058 100644
--- a/core/src/main/scala/kafka/utils/DelayedItem.scala
+++ b/core/src/main/scala/kafka/utils/DelayedItem.scala
@@ -23,7 +23,7 @@ import org.apache.kafka.common.utils.Time
 
 import scala.math._
 
-class DelayedItem(delayMs: Long) extends Delayed with Logging {
+class DelayedItem(val delayMs: Long) extends Delayed with Logging {
 
   private val dueMs = Time.SYSTEM.milliseconds + delayMs
 
diff --git a/core/src/main/scala/kafka/utils/Exit.scala b/core/src/main/scala/kafka/utils/Exit.scala
index 5819e97076f0a..ad17237571e56 100644
--- a/core/src/main/scala/kafka/utils/Exit.scala
+++ b/core/src/main/scala/kafka/utils/Exit.scala
@@ -34,20 +34,30 @@ object Exit {
     throw new AssertionError("halt should not return, but it did.")
   }
 
+  def addShutdownHook(name: String, shutdownHook: => Unit): Unit = {
+    JExit.addShutdownHook(name, () => shutdownHook)
+  }
+
   def setExitProcedure(exitProcedure: (Int, Option[String]) => Nothing): Unit =
     JExit.setExitProcedure(functionToProcedure(exitProcedure))
 
   def setHaltProcedure(haltProcedure: (Int, Option[String]) => Nothing): Unit =
     JExit.setHaltProcedure(functionToProcedure(haltProcedure))
 
+  def setShutdownHookAdder(shutdownHookAdder: (String, => Unit) => Unit): Unit = {
+    JExit.setShutdownHookAdder((name, runnable) => shutdownHookAdder(name, runnable.run))
+  }
+
   def resetExitProcedure(): Unit =
     JExit.resetExitProcedure()
 
   def resetHaltProcedure(): Unit =
     JExit.resetHaltProcedure()
 
+  def resetShutdownHookAdder(): Unit =
+    JExit.resetShutdownHookAdder()
+
   private def functionToProcedure(procedure: (Int, Option[String]) => Nothing) = new JExit.Procedure {
     def execute(statusCode: Int, message: String): Unit = procedure(statusCode, Option(message))
   }
-
 }
diff --git a/core/src/main/scala/kafka/utils/FileLock.scala b/core/src/main/scala/kafka/utils/FileLock.scala
index d1edc83bf8dcd..c635f76dff3a2 100644
--- a/core/src/main/scala/kafka/utils/FileLock.scala
+++ b/core/src/main/scala/kafka/utils/FileLock.scala
@@ -18,7 +18,7 @@
 
 import java.io._
 import java.nio.channels._
-import java.nio.file.{FileAlreadyExistsException, Files}
+import java.nio.file.StandardOpenOption
 
 /**
  * A file lock a la flock/funlock
@@ -27,18 +27,16 @@ import java.nio.file.{FileAlreadyExistsException, Files}
  */
 class FileLock(val file: File) extends Logging {
 
-  try Files.createFile(file.toPath) // create the file if it doesn't exist
-  catch { case _: FileAlreadyExistsException => }
-
-  private val channel = new RandomAccessFile(file, "rw").getChannel()
+  private val channel = FileChannel.open(file.toPath, StandardOpenOption.CREATE, StandardOpenOption.READ,
+    StandardOpenOption.WRITE)
   private var flock: java.nio.channels.FileLock = null
 
   /**
    * Lock the file or throw an exception if the lock is already held
    */
-  def lock() {
+  def lock(): Unit = {
     this synchronized {
-      trace("Acquiring lock on " + file.getAbsolutePath)
+      trace(s"Acquiring lock on ${file.getAbsolutePath}")
       flock = channel.lock()
     }
   }
@@ -48,7 +46,7 @@ class FileLock(val file: File) extends Logging {
    */
   def tryLock(): Boolean = {
     this synchronized {
-      trace("Acquiring lock on " + file.getAbsolutePath)
+      trace(s"Acquiring lock on ${file.getAbsolutePath}")
       try {
         // weirdly this method will return null if the lock is held by another
         // process, but will throw an exception if the lock is held by this process
@@ -64,9 +62,9 @@ class FileLock(val file: File) extends Logging {
   /**
    * Unlock the lock if it is held
    */
-  def unlock() {
+  def unlock(): Unit = {
     this synchronized {
-      trace("Releasing lock on " + file.getAbsolutePath)
+      trace(s"Releasing lock on ${file.getAbsolutePath}")
       if(flock != null)
         flock.release()
     }
diff --git a/core/src/main/scala/kafka/utils/Implicits.scala b/core/src/main/scala/kafka/utils/Implicits.scala
index 5196d45ea9154..fbd22ec6b8cfc 100644
--- a/core/src/main/scala/kafka/utils/Implicits.scala
+++ b/core/src/main/scala/kafka/utils/Implicits.scala
@@ -20,7 +20,8 @@ package kafka.utils
 import java.util
 import java.util.Properties
 
-import scala.collection.JavaConverters._
+import scala.annotation.nowarn
+import scala.jdk.CollectionConverters._
 
 /**
   * In order to have these implicits in scope, add the following import:
@@ -46,4 +47,21 @@ object Implicits {
 
   }
 
+  /**
+   * Exposes `forKeyValue` which maps to `foreachEntry` in Scala 2.13 and `foreach` in Scala 2.12
+   * (with the help of scala.collection.compat). `foreachEntry` avoids the tuple allocation and
+   * is more efficient.
+   *
+   * This was not named `foreachEntry` to avoid `unused import` warnings in Scala 2.13 (the implicit
+   * would not be triggered in Scala 2.13 since `Map.foreachEntry` would have precedence).
+   */
+  @nowarn("cat=unused-imports")
+  implicit class MapExtensionMethods[K, V](private val self: scala.collection.Map[K, V]) extends AnyVal {
+    import scala.collection.compat._
+    def forKeyValue[U](f: (K, V) => U): Unit = {
+      self.foreachEntry { (k, v) => f(k, v) }
+    }
+
+  }
+
 }
diff --git a/core/src/main/scala/kafka/utils/IteratorTemplate.scala b/core/src/main/scala/kafka/utils/IteratorTemplate.scala
deleted file mode 100644
index 7cd161e0e475e..0000000000000
--- a/core/src/main/scala/kafka/utils/IteratorTemplate.scala
+++ /dev/null
@@ -1,85 +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
-
-class State
-object DONE extends State
-object READY extends State
-object NOT_READY extends State
-object FAILED extends State
-
-/**
- * Transliteration of the iterator template in google collections. To implement an iterator
- * override makeNext and call allDone() when there is no more items
- */
-abstract class IteratorTemplate[T] extends Iterator[T] with java.util.Iterator[T] {
-  
-  private var state: State = NOT_READY
-  private var nextItem = null.asInstanceOf[T]
-
-  def next(): T = {
-    if(!hasNext())
-      throw new NoSuchElementException()
-    state = NOT_READY
-    if(nextItem == null)
-      throw new IllegalStateException("Expected item but none found.")
-    nextItem
-  }
-  
-  def peek(): T = {
-    if(!hasNext)
-      throw new NoSuchElementException()
-    nextItem
-  }
-  
-  def hasNext: Boolean = {
-    if(state == FAILED)
-      throw new IllegalStateException("Iterator is in failed state")
-    state match {
-      case DONE => false
-      case READY => true
-      case _ => maybeComputeNext()
-    }
-  }
-  
-  protected def makeNext(): T
-  
-  def maybeComputeNext(): Boolean = {
-    state = FAILED
-    nextItem = makeNext()
-    if(state == DONE) {
-      false
-    } else {
-      state = READY
-      true
-    }
-  }
-  
-  protected def allDone(): T = {
-    state = DONE
-    null.asInstanceOf[T]
-  }
-  
-  override def remove =
-    throw new UnsupportedOperationException("Removal not supported")
-
-  protected def resetState() {
-    state = NOT_READY
-  }
-}
-
diff --git a/core/src/main/scala/kafka/utils/Json.scala b/core/src/main/scala/kafka/utils/Json.scala
index ad40c495e5497..48bdc85dbb9ba 100644
--- a/core/src/main/scala/kafka/utils/Json.scala
+++ b/core/src/main/scala/kafka/utils/Json.scala
@@ -16,13 +16,11 @@
  */
 package kafka.utils
 
-import java.nio.charset.StandardCharsets
-
 import com.fasterxml.jackson.core.JsonProcessingException
 import com.fasterxml.jackson.databind.ObjectMapper
 import kafka.utils.json.JsonValue
 
-import scala.collection._
+import scala.reflect.ClassTag
 
 /**
  * Provides methods for parsing JSON with Jackson and encoding to JSON with a simple and naive custom implementation.
@@ -38,6 +36,15 @@ object Json {
     try Option(mapper.readTree(input)).map(JsonValue(_))
     catch { case _: JsonProcessingException => None }
 
+  /**
+   * Parse a JSON string into either a generic type T, or a JsonProcessingException in the case of
+   * exception.
+   */
+  def parseStringAs[T](input: String)(implicit tag: ClassTag[T]): Either[JsonProcessingException, T] = {
+    try Right(mapper.readValue(input, tag.runtimeClass).asInstanceOf[T])
+    catch { case e: JsonProcessingException => Left(e) }
+  }
+
   /**
    * Parse a JSON byte array into a JsonValue if possible. `None` is returned if `input` is not valid JSON.
    */
@@ -45,37 +52,29 @@ object Json {
     try Option(mapper.readTree(input)).map(JsonValue(_))
     catch { case _: JsonProcessingException => None }
 
+  def tryParseBytes(input: Array[Byte]): Either[JsonProcessingException, JsonValue] =
+    try Right(mapper.readTree(input)).map(JsonValue(_))
+    catch { case e: JsonProcessingException => Left(e) }
+
   /**
-   * Encode an object into a JSON string. This method accepts any type T where
-   *   T => null | Boolean | String | Number | Map[String, T] | Array[T] | Iterable[T]
-   * Any other type will result in an exception.
-   * 
-   * This method does not properly handle non-ascii characters. 
+   * Parse a JSON byte array into either a generic type T, or a JsonProcessingException in the case of exception.
    */
-  def encode(obj: Any): String = {
-    obj match {
-      case null => "null"
-      case b: Boolean => b.toString
-      case s: String => "\"" + s + "\""
-      case n: Number => n.toString
-      case m: Map[_, _] => "{" +
-        m.map {
-          case (k, v) => encode(k) + ":" + encode(v)
-          case elem => throw new IllegalArgumentException(s"Invalid map element '$elem' in $obj")
-        }.mkString(",") + "}"
-      case a: Array[_] => encode(a.toSeq)
-      case i: Iterable[_] => "[" + i.map(encode).mkString(",") + "]"
-      case other: AnyRef => throw new IllegalArgumentException(s"Unknown argument of type ${other.getClass}: $other")
-    }
+  def parseBytesAs[T](input: Array[Byte])(implicit tag: ClassTag[T]): Either[JsonProcessingException, T] = {
+    try Right(mapper.readValue(input, tag.runtimeClass).asInstanceOf[T])
+    catch { case e: JsonProcessingException => Left(e) }
   }
 
   /**
-   * Encode an object into a JSON value in bytes. This method accepts any type T where
-   *   T => null | Boolean | String | Number | Map[String, T] | Array[T] | Iterable[T]
-   * Any other type will result in an exception.
-   *
-   * This method does not properly handle non-ascii characters.
+   * Encode an object into a JSON string. This method accepts any type supported by Jackson's ObjectMapper in
+   * the default configuration. That is, Java collections are supported, but Scala collections are not (to avoid
+   * a jackson-scala dependency).
    */
-  def encodeAsBytes(obj: Any): Array[Byte] = encode(obj).getBytes(StandardCharsets.UTF_8)
+  def encodeAsString(obj: Any): String = mapper.writeValueAsString(obj)
 
+  /**
+   * Encode an object into a JSON value in bytes. This method accepts any type supported by Jackson's ObjectMapper in
+   * the default configuration. That is, Java collections are supported, but Scala collections are not (to avoid
+   * a jackson-scala dependency).
+   */
+  def encodeAsBytes(obj: Any): Array[Byte] = mapper.writeValueAsBytes(obj)
 }
diff --git a/core/src/main/scala/kafka/utils/KafkaScheduler.scala b/core/src/main/scala/kafka/utils/KafkaScheduler.scala
index d20fdd7fa55b9..f1585afa1ae2b 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,14 +80,13 @@ class KafkaScheduler(val threads: Int,
       executor = new ScheduledThreadPoolExecutor(threads)
       executor.setContinueExistingPeriodicTasksAfterShutdownPolicy(false)
       executor.setExecuteExistingDelayedTasksAfterShutdownPolicy(false)
-      executor.setThreadFactory(new ThreadFactory() {
-                                  def newThread(runnable: Runnable): Thread = 
-                                    new KafkaThread(threadNamePrefix + schedulerThreadId.getAndIncrement(), runnable, daemon)
-                                })
+      executor.setRemoveOnCancelPolicy(true)
+      executor.setThreadFactory(runnable =>
+        new KafkaThread(threadNamePrefix + schedulerThreadId.getAndIncrement(), runnable, daemon))
     }
   }
   
-  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
@@ -99,27 +99,42 @@ class KafkaScheduler(val threads: Int,
     }
   }
 
-  def schedule(name: String, fun: ()=>Unit, delay: Long, period: Long, unit: TimeUnit) {
+  def scheduleOnce(name: String, fun: () => Unit): Unit = {
+    schedule(name, fun, delay = 0L, period = -1L, unit = TimeUnit.MILLISECONDS)
+  }
+
+  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 {
-      ensureRunning
-      val runnable = CoreUtils.runnable {
+      ensureRunning()
+      val runnable: Runnable = () => {
         try {
           trace("Beginning execution of scheduled task '%s'.".format(name))
           fun()
         } catch {
-          case t: Throwable => error("Uncaught exception in scheduled task '" + name +"'", t)
+          case t: Throwable => error(s"Uncaught exception in scheduled task '$name'", t)
         } finally {
           trace("Completed execution of scheduled task '%s'.".format(name))
         }
       }
-      if(period >= 0)
+      if (period >= 0)
         executor.scheduleAtFixedRate(runnable, delay, period, unit)
       else
         executor.schedule(runnable, delay, unit)
     }
   }
+
+  /**
+   * Package private for testing.
+   */
+  private[kafka] def taskRunning(task: ScheduledFuture[_]): Boolean = {
+    executor.getQueue().contains(task)
+  }
+
+  def resizeThreadPool(newSize: Int): Unit = {
+    executor.setCorePoolSize(newSize)
+  }
   
   def isStarted: Boolean = {
     this synchronized {
diff --git a/core/src/main/scala/kafka/utils/Log4jController.scala b/core/src/main/scala/kafka/utils/Log4jController.scala
index 95d07334a61d5..a02fdb69e14d1 100755
--- a/core/src/main/scala/kafka/utils/Log4jController.scala
+++ b/core/src/main/scala/kafka/utils/Log4jController.scala
@@ -22,69 +22,113 @@ import java.util.Locale
 
 import org.apache.log4j.{Level, LogManager, Logger}
 
+import scala.collection.mutable
+import scala.jdk.CollectionConverters._
+
+
+object Log4jController {
+  val ROOT_LOGGER = "root"
+
+  private def resolveLevel(logger: Logger): String = {
+    var name = logger.getName
+    var level = logger.getLevel
+    while (level == null) {
+      val index = name.lastIndexOf(".")
+      if (index > 0) {
+        name = name.substring(0, index)
+        val ancestor = existingLogger(name)
+        if (ancestor != null) {
+          level = ancestor.getLevel
+        }
+      } else {
+        level = existingLogger(ROOT_LOGGER).getLevel
+      }
+    }
+    level.toString
+  }
 
-/**
- * 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 {
+  /**
+    * 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)
 
-  def getLoggers = {
-    val lst = new util.ArrayList[String]()
-    lst.add("root=" + existingLogger("root").getLevel.toString)
     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"))
+        logs.put(logger.getName, resolveLevel(logger))
       }
     }
-    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] = {
+    // we replace scala collection by java collection so mbean client is able to deserialize it without scala library.
+    new util.ArrayList[String](Log4jController.loggers.map {
+      case (logger, level) => s"$logger=$level"
+    }.toSeq.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.resolveLevel(log)
     }
     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/Logging.scala b/core/src/main/scala/kafka/utils/Logging.scala
index e409bba3cca7f..0221821e7473a 100755
--- a/core/src/main/scala/kafka/utils/Logging.scala
+++ b/core/src/main/scala/kafka/utils/Logging.scala
@@ -17,8 +17,8 @@
 
 package kafka.utils
 
-import com.typesafe.scalalogging.{LazyLogging, Logger}
-import org.slf4j.{Marker, MarkerFactory}
+import com.typesafe.scalalogging.Logger
+import org.slf4j.{LoggerFactory, Marker, MarkerFactory}
 
 
 object Log4jControllerRegistration {
@@ -38,13 +38,16 @@ private object Logging {
   private val FatalMarker: Marker = MarkerFactory.getMarker("FATAL")
 }
 
-trait Logging extends LazyLogging {
-  def loggerName: String = logger.underlying.getName
+trait Logging {
+
+  protected lazy val logger = Logger(LoggerFactory.getLogger(loggerName))
 
   protected var logIdent: String = _
 
   Log4jControllerRegistration
 
+  protected def loggerName: String = getClass.getName
+
   protected def msgWithLogIdent(msg: String): String =
     if (logIdent == null) msg else logIdent + msg
 
diff --git a/core/src/main/scala/kafka/utils/Mx4jLoader.scala b/core/src/main/scala/kafka/utils/Mx4jLoader.scala
index d9d1cb46a7628..e49f3a57f1ed1 100644
--- a/core/src/main/scala/kafka/utils/Mx4jLoader.scala
+++ b/core/src/main/scala/kafka/utils/Mx4jLoader.scala
@@ -45,7 +45,7 @@ object Mx4jLoader extends Logging {
       val processorName = new ObjectName("Server:name=XSLTProcessor")
 
       val httpAdaptorClass = Class.forName("mx4j.tools.adaptor.http.HttpAdaptor")
-      val httpAdaptor = httpAdaptorClass.newInstance()
+      val httpAdaptor = httpAdaptorClass.getDeclaredConstructor().newInstance()
       httpAdaptorClass.getMethod("setHost", classOf[String]).invoke(httpAdaptor, address.asInstanceOf[AnyRef])
       httpAdaptorClass.getMethod("setPort", Integer.TYPE).invoke(httpAdaptor, port.asInstanceOf[AnyRef])
 
@@ -53,15 +53,15 @@ object Mx4jLoader extends Logging {
       mbs.registerMBean(httpAdaptor, httpName)
 
       val xsltProcessorClass = Class.forName("mx4j.tools.adaptor.http.XSLTProcessor")
-      val xsltProcessor = xsltProcessorClass.newInstance()
+      val xsltProcessor = xsltProcessorClass.getDeclaredConstructor().newInstance()
       httpAdaptorClass.getMethod("setProcessor", Class.forName("mx4j.tools.adaptor.http.ProcessorMBean")).invoke(httpAdaptor, xsltProcessor.asInstanceOf[AnyRef])
       mbs.registerMBean(xsltProcessor, processorName)
       httpAdaptorClass.getMethod("start").invoke(httpAdaptor)
-      info("mx4j successfuly loaded")
+      info("mx4j successfully loaded")
       return true
     }
     catch {
-	  case _: ClassNotFoundException =>
+      case _: ClassNotFoundException =>
         info("Will not load MX4J, mx4j-tools.jar is not in the classpath")
       case e: Throwable =>
         warn("Could not start register mbean in JMX", e)
diff --git a/core/src/main/scala/kafka/utils/PasswordEncoder.scala b/core/src/main/scala/kafka/utils/PasswordEncoder.scala
new file mode 100644
index 0000000000000..f748a455c62bf
--- /dev/null
+++ b/core/src/main/scala/kafka/utils/PasswordEncoder.scala
@@ -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 kafka.utils
+
+import java.nio.charset.StandardCharsets
+import java.security.{AlgorithmParameters, NoSuchAlgorithmException, SecureRandom}
+import java.security.spec.AlgorithmParameterSpec
+import java.util.Base64
+
+import javax.crypto.{Cipher, SecretKeyFactory}
+import javax.crypto.spec._
+import kafka.utils.PasswordEncoder._
+import org.apache.kafka.common.config.ConfigException
+import org.apache.kafka.common.config.types.Password
+
+import scala.collection.Map
+
+object PasswordEncoder {
+  val KeyFactoryAlgorithmProp = "keyFactoryAlgorithm"
+  val CipherAlgorithmProp = "cipherAlgorithm"
+  val InitializationVectorProp = "initializationVector"
+  val KeyLengthProp = "keyLength"
+  val SaltProp = "salt"
+  val IterationsProp = "iterations"
+  val EncyrptedPasswordProp = "encryptedPassword"
+  val PasswordLengthProp = "passwordLength"
+}
+
+/**
+  * Password encoder and decoder implementation. Encoded passwords are persisted as a CSV map
+  * containing the encoded password in base64 and along with the properties used for encryption.
+  *
+  * @param secret The secret used for encoding and decoding
+  * @param keyFactoryAlgorithm  Key factory algorithm if configured. By default, PBKDF2WithHmacSHA512 is
+  *                             used if available, PBKDF2WithHmacSHA1 otherwise.
+  * @param cipherAlgorithm Cipher algorithm used for encoding.
+  * @param keyLength Key length used for encoding. This should be valid for the specified algorithms.
+  * @param iterations Iteration count used for encoding.
+  *
+  * The provided `keyFactoryAlgorithm`, 'cipherAlgorithm`, `keyLength` and `iterations` are used for encoding passwords.
+  * The values used for encoding are stored along with the encoded password and the stored values are used for decoding.
+  *
+  */
+class PasswordEncoder(secret: Password,
+                      keyFactoryAlgorithm: Option[String],
+                      cipherAlgorithm: String,
+                      keyLength: Int,
+                      iterations: Int) extends Logging {
+
+  private val secureRandom = new SecureRandom
+  private val cipherParamsEncoder = cipherParamsInstance(cipherAlgorithm)
+
+  def encode(password: Password): String = {
+    val salt = new Array[Byte](256)
+    secureRandom.nextBytes(salt)
+    val cipher = Cipher.getInstance(cipherAlgorithm)
+    val keyFactory = secretKeyFactory(keyFactoryAlgorithm)
+    val keySpec = secretKeySpec(keyFactory, cipherAlgorithm, keyLength, salt, iterations)
+    cipher.init(Cipher.ENCRYPT_MODE, keySpec)
+    val encryptedPassword = cipher.doFinal(password.value.getBytes(StandardCharsets.UTF_8))
+    val encryptedMap = Map(
+      KeyFactoryAlgorithmProp -> keyFactory.getAlgorithm,
+      CipherAlgorithmProp -> cipherAlgorithm,
+      KeyLengthProp -> keyLength,
+      SaltProp -> base64Encode(salt),
+      IterationsProp -> iterations.toString,
+      EncyrptedPasswordProp -> base64Encode(encryptedPassword),
+      PasswordLengthProp -> password.value.length
+    ) ++ cipherParamsEncoder.toMap(cipher.getParameters)
+    encryptedMap.map { case (k, v) => s"$k:$v" }.mkString(",")
+  }
+
+  def decode(encodedPassword: String): Password = {
+    val params = CoreUtils.parseCsvMap(encodedPassword)
+    val keyFactoryAlg = params(KeyFactoryAlgorithmProp)
+    val cipherAlg = params(CipherAlgorithmProp)
+    val keyLength = params(KeyLengthProp).toInt
+    val salt = base64Decode(params(SaltProp))
+    val iterations = params(IterationsProp).toInt
+    val encryptedPassword = base64Decode(params(EncyrptedPasswordProp))
+    val passwordLengthProp = params(PasswordLengthProp).toInt
+    val cipher = Cipher.getInstance(cipherAlg)
+    val keyFactory = secretKeyFactory(Some(keyFactoryAlg))
+    val keySpec = secretKeySpec(keyFactory, cipherAlg, keyLength, salt, iterations)
+    cipher.init(Cipher.DECRYPT_MODE, keySpec, cipherParamsEncoder.toParameterSpec(params))
+    val password = try {
+      val decrypted = cipher.doFinal(encryptedPassword)
+      new String(decrypted, StandardCharsets.UTF_8)
+    } catch {
+      case e: Exception => throw new ConfigException("Password could not be decoded", e)
+    }
+    if (password.length != passwordLengthProp) // Sanity check
+      throw new ConfigException("Password could not be decoded, sanity check of length failed")
+    new Password(password)
+  }
+
+  private def secretKeyFactory(keyFactoryAlg: Option[String]): SecretKeyFactory = {
+    keyFactoryAlg match {
+      case Some(algorithm) => SecretKeyFactory.getInstance(algorithm)
+      case None =>
+        try {
+          SecretKeyFactory.getInstance("PBKDF2WithHmacSHA512")
+        } catch {
+          case _: NoSuchAlgorithmException => SecretKeyFactory.getInstance("PBKDF2WithHmacSHA1")
+        }
+    }
+  }
+
+  private def secretKeySpec(keyFactory: SecretKeyFactory,
+                            cipherAlg: String,
+                            keyLength: Int,
+                            salt: Array[Byte], iterations: Int): SecretKeySpec = {
+    val keySpec = new PBEKeySpec(secret.value.toCharArray, salt, iterations, keyLength)
+    val algorithm = if (cipherAlg.indexOf('/') > 0) cipherAlg.substring(0, cipherAlg.indexOf('/')) else cipherAlg
+    new SecretKeySpec(keyFactory.generateSecret(keySpec).getEncoded, algorithm)
+  }
+
+  private def base64Encode(bytes: Array[Byte]): String = Base64.getEncoder.encodeToString(bytes)
+
+  private[utils] def base64Decode(encoded: String): Array[Byte] = Base64.getDecoder.decode(encoded)
+
+  private def cipherParamsInstance(cipherAlgorithm: String): CipherParamsEncoder = {
+    val aesPattern = "AES/(.*)/.*".r
+    cipherAlgorithm match {
+      case aesPattern("GCM") => new GcmParamsEncoder
+      case _ => new IvParamsEncoder
+    }
+  }
+
+  private trait CipherParamsEncoder {
+    def toMap(cipher: AlgorithmParameters): Map[String, String]
+    def toParameterSpec(paramMap: Map[String, String]): AlgorithmParameterSpec
+  }
+
+  private class IvParamsEncoder extends CipherParamsEncoder {
+    def toMap(cipherParams: AlgorithmParameters): Map[String, String] = {
+      if (cipherParams != null) {
+        val ivSpec = cipherParams.getParameterSpec(classOf[IvParameterSpec])
+        Map(InitializationVectorProp -> base64Encode(ivSpec.getIV))
+      } else
+        throw new IllegalStateException("Could not determine initialization vector for cipher")
+    }
+    def toParameterSpec(paramMap: Map[String, String]): AlgorithmParameterSpec = {
+      new IvParameterSpec(base64Decode(paramMap(InitializationVectorProp)))
+    }
+  }
+
+  private class GcmParamsEncoder extends CipherParamsEncoder {
+    def toMap(cipherParams: AlgorithmParameters): Map[String, String] = {
+      if (cipherParams != null) {
+        val spec = cipherParams.getParameterSpec(classOf[GCMParameterSpec])
+        Map(InitializationVectorProp -> base64Encode(spec.getIV),
+            "authenticationTagLength" -> spec.getTLen.toString)
+      } else
+        throw new IllegalStateException("Could not determine initialization vector for cipher")
+    }
+    def toParameterSpec(paramMap: Map[String, String]): AlgorithmParameterSpec = {
+      new GCMParameterSpec(paramMap("authenticationTagLength").toInt, base64Decode(paramMap(InitializationVectorProp)))
+    }
+  }
+}
diff --git a/core/src/main/scala/kafka/utils/Pool.scala b/core/src/main/scala/kafka/utils/Pool.scala
index 4ddf55736dde8..d64ff5d6538ff 100644
--- a/core/src/main/scala/kafka/utils/Pool.scala
+++ b/core/src/main/scala/kafka/utils/Pool.scala
@@ -19,17 +19,19 @@ package kafka.utils
 
 import java.util.concurrent._
 
-import collection.mutable
-import collection.JavaConverters._
-import kafka.common.KafkaException
+import org.apache.kafka.common.KafkaException
+
+import collection.Set
+import scala.jdk.CollectionConverters._
 
 class Pool[K,V](valueFactory: Option[K => V] = None) extends Iterable[(K, V)] {
 
   private val pool: ConcurrentMap[K, V] = new ConcurrentHashMap[K, V]
-  private val createLock = new Object
-  
+
   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)
 
   /**
@@ -56,21 +58,8 @@ class Pool[K,V](valueFactory: Option[K => V] = None) extends Iterable[(K, V)] {
     * @param createValue Factory function.
     * @return The final value associated with the key.
     */
-  def getAndMaybePut(key: K, createValue: => V): V = {
-    val current = pool.get(key)
-    if (current == null) {
-      createLock synchronized {
-        val current = pool.get(key)
-        if (current == null) {
-          val value = createValue
-          pool.put(key, value)
-          value
-        }
-        else current
-      }
-    }
-    else current
-  }
+  def getAndMaybePut(key: K, createValue: => V): V =
+    pool.computeIfAbsent(key, _ => createValue)
 
   def contains(id: K): Boolean = pool.containsKey(id)
   
@@ -80,11 +69,17 @@ class Pool[K,V](valueFactory: Option[K => V] = None) extends Iterable[(K, V)] {
 
   def remove(key: K, value: V): Boolean = pool.remove(key, value)
 
-  def keys: mutable.Set[K] = pool.keySet.asScala
+  def removeAll(keys: Iterable[K]): Unit = pool.keySet.removeAll(keys.asJavaCollection)
+
+  def keys: Set[K] = pool.keySet.asScala
 
   def values: Iterable[V] = pool.values.asScala
 
-  def clear() { pool.clear() }
+  def clear(): Unit = { pool.clear() }
+
+  def foreachEntry(f: (K, V) => Unit): Unit = {
+    pool.forEach((k, v) => f(k, v))
+  }
   
   override def size: Int = pool.size
   
@@ -94,7 +89,7 @@ class Pool[K,V](valueFactory: Option[K => V] = None) extends Iterable[(K, V)] {
     
     def hasNext: Boolean = iter.hasNext
     
-    def next: (K, V) = {
+    def next(): (K, V) = {
       val n = iter.next
       (n.getKey, n.getValue)
     }
diff --git a/core/src/main/scala/kafka/utils/QuotaUtils.scala b/core/src/main/scala/kafka/utils/QuotaUtils.scala
new file mode 100755
index 0000000000000..93d5cdebcc344
--- /dev/null
+++ b/core/src/main/scala/kafka/utils/QuotaUtils.scala
@@ -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 kafka.utils
+
+import org.apache.kafka.common.MetricName
+import org.apache.kafka.common.metrics.{KafkaMetric, Measurable, QuotaViolationException}
+import org.apache.kafka.common.metrics.stats.Rate
+
+/**
+ * Helper functions related to quotas
+ */
+object QuotaUtils {
+
+  /**
+   * This calculates the amount of time needed to bring the observed rate within quota
+   * assuming that no new metrics are recorded.
+   *
+   * If O is the observed rate and T is the target rate over a window of W, to bring O down to T,
+   * we need to add a delay of X to W such that O * W / (W + X) = T.
+   * Solving for X, we get X = (O - T)/T * W.
+   *
+   * @param timeMs current time in milliseconds
+   * @return Delay in milliseconds
+   */
+  def throttleTime(e: QuotaViolationException, timeMs: Long): Long = {
+    val difference = e.value - e.bound
+    // Use the precise window used by the rate calculation
+    val throttleTimeMs = difference / e.bound * windowSize(e.metric, timeMs)
+    Math.round(throttleTimeMs)
+  }
+
+  /**
+   * Calculates the amount of time needed to bring the observed rate within quota using the same algorithm as
+   * throttleTime() utility method but the returned value is capped to given maxThrottleTime
+   */
+  def boundedThrottleTime(e: QuotaViolationException, maxThrottleTime: Long, timeMs: Long): Long = {
+    math.min(throttleTime(e, timeMs), maxThrottleTime)
+  }
+
+  /**
+   * Returns window size of the given metric
+   *
+   * @param metric metric with measurable of type Rate
+   * @param timeMs current time in milliseconds
+   * @throws IllegalArgumentException if given measurable is not Rate
+   */
+  private def windowSize(metric: KafkaMetric, timeMs: Long): Long =
+    measurableAsRate(metric.metricName, metric.measurable).windowSize(metric.config, timeMs)
+
+  /**
+   * Casts provided Measurable to Rate
+   * @throws IllegalArgumentException if given measurable is not Rate
+   */
+  private def measurableAsRate(name: MetricName, measurable: Measurable): Rate = {
+    measurable match {
+      case r: Rate => r
+      case _ => throw new IllegalArgumentException(s"Metric $name is not a Rate metric, value $measurable")
+    }
+  }
+}
diff --git a/core/src/main/scala/kafka/utils/ReplicationUtils.scala b/core/src/main/scala/kafka/utils/ReplicationUtils.scala
index f5751d20b584b..e2733b8936fd9 100644
--- a/core/src/main/scala/kafka/utils/ReplicationUtils.scala
+++ b/core/src/main/scala/kafka/utils/ReplicationUtils.scala
@@ -18,76 +18,39 @@
 package kafka.utils
 
 import kafka.api.LeaderAndIsr
-import kafka.controller.{IsrChangeNotificationHandler, LeaderIsrAndControllerEpoch}
-import kafka.utils.ZkUtils._
+import kafka.controller.LeaderIsrAndControllerEpoch
 import kafka.zk._
 import org.apache.kafka.common.TopicPartition
-import org.apache.zookeeper.data.Stat
-
-import scala.collection._
 
 object ReplicationUtils extends Logging {
 
-  private val IsrChangeNotificationPrefix = "isr_change_"
-
-  def updateLeaderAndIsr(zkClient: KafkaZkClient, topic: String, partitionId: Int, newLeaderAndIsr: LeaderAndIsr, controllerEpoch: Int,
-    zkVersion: Int): (Boolean,Int) = {
-    debug(s"Updated ISR for $topic-$partitionId to ${newLeaderAndIsr.isr.mkString(",")}")
-    val path = getTopicPartitionLeaderAndIsrPath(topic, partitionId)
-    val newLeaderData = LeaderAndIsrZNode.encode(newLeaderAndIsr, controllerEpoch)
+  def updateLeaderAndIsr(zkClient: KafkaZkClient, partition: TopicPartition, newLeaderAndIsr: LeaderAndIsr,
+                         controllerEpoch: Int): (Boolean, Int) = {
+    debug(s"Updated ISR for $partition to ${newLeaderAndIsr.isr.mkString(",")}")
+    val path = TopicPartitionStateZNode.path(partition)
+    val newLeaderData = TopicPartitionStateZNode.encode(LeaderIsrAndControllerEpoch(newLeaderAndIsr, controllerEpoch))
     // use the epoch of the controller that made the leadership decision, instead of the current controller epoch
-    val updatePersistentPath: (Boolean, Int) = zkClient.conditionalUpdatePath(path, newLeaderData, zkVersion, Some(checkLeaderAndIsrZkData))
+    val updatePersistentPath: (Boolean, Int) = zkClient.conditionalUpdatePath(path, newLeaderData,
+      newLeaderAndIsr.zkVersion, Some(checkLeaderAndIsrZkData))
     updatePersistentPath
   }
 
-  def propagateIsrChanges(zkClient: KafkaZkClient, isrChangeSet: Set[TopicPartition]): Unit = {
-    val isrChangeNotificationPath: String = zkClient.createSequentialPersistentPath(
-      ZkUtils.IsrChangeNotificationPath + "/" + IsrChangeNotificationPrefix,
-      generateIsrChangeJson(isrChangeSet))
-    debug(s"Added $isrChangeNotificationPath for $isrChangeSet")
-  }
-
-  private def checkLeaderAndIsrZkData(zkClient: KafkaZkClient, path: String, expectedLeaderAndIsrInfo: String): (Boolean, Int) = {
+  private def checkLeaderAndIsrZkData(zkClient: KafkaZkClient, path: String, expectedLeaderAndIsrInfo: Array[Byte]): (Boolean, Int) = {
     try {
       val (writtenLeaderOpt, writtenStat) = zkClient.getDataAndStat(path)
-      val expectedLeader = parseLeaderAndIsr(expectedLeaderAndIsrInfo, path, writtenStat)
-      writtenLeaderOpt.foreach { writtenData =>
-        val writtenLeader = parseLeaderAndIsr(writtenData, path, writtenStat)
-        (expectedLeader,writtenLeader) match {
-          case (Some(expectedLeader),Some(writtenLeader)) =>
-            if(expectedLeader == writtenLeader)
-              return (true, writtenStat.getVersion())
-          case _ =>
+      val expectedLeaderOpt = TopicPartitionStateZNode.decode(expectedLeaderAndIsrInfo, writtenStat)
+      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
         }
       }
+      if (succeeded) (true, writtenStat.getVersion)
+      else (false, -1)
     } catch {
-      case _: Exception =>
+      case _: Exception => (false, -1)
     }
-    (false, -1)
-  }
-
-  def getLeaderIsrAndEpochForPartition(zkUtils: ZkUtils, topic: String, partition: Int): Option[LeaderIsrAndControllerEpoch] = {
-    val leaderAndIsrPath = getTopicPartitionLeaderAndIsrPath(topic, partition)
-    val (leaderAndIsrOpt, stat) = zkUtils.readDataMaybeNull(leaderAndIsrPath)
-    debug(s"Read leaderISR $leaderAndIsrOpt for $topic-$partition")
-    leaderAndIsrOpt.flatMap(leaderAndIsrStr => parseLeaderAndIsr(leaderAndIsrStr, leaderAndIsrPath, stat))
-  }
-
-  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))}
-  }
-
-  private def generateIsrChangeJson(isrChanges: Set[TopicPartition]): String = {
-    val partitions = isrChanges.map(tp => Map("topic" -> tp.topic, "partition" -> tp.partition)).toArray
-    Json.encode(Map("version" -> IsrChangeNotificationHandler.Version, "partitions" -> partitions))
   }
 
 }
diff --git a/core/src/main/scala/kafka/utils/ShutdownableThread.scala b/core/src/main/scala/kafka/utils/ShutdownableThread.scala
index 0922d15e93a03..0ca21c4ab166d 100644
--- a/core/src/main/scala/kafka/utils/ShutdownableThread.scala
+++ b/core/src/main/scala/kafka/utils/ShutdownableThread.scala
@@ -17,8 +17,7 @@
 
 package kafka.utils
 
-import java.util.concurrent.atomic.AtomicBoolean
-import java.util.concurrent.CountDownLatch
+import java.util.concurrent.{CountDownLatch, TimeUnit}
 
 import org.apache.kafka.common.internals.FatalExitError
 
@@ -26,34 +25,62 @@ abstract class ShutdownableThread(val name: String, val isInterruptible: Boolean
         extends Thread(name) with Logging {
   this.setDaemon(false)
   this.logIdent = "[" + name + "]: "
-  val isRunning: AtomicBoolean = new AtomicBoolean(true)
-  private val shutdownLatch = new CountDownLatch(1)
-
+  private val shutdownInitiated = new CountDownLatch(1)
+  private val shutdownComplete = new CountDownLatch(1)
+  @volatile private var isStarted: Boolean = false
+  
   def shutdown(): Unit = {
     initiateShutdown()
     awaitShutdown()
   }
 
-  def isShutdownComplete: Boolean = {
-    shutdownLatch.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 = {
-    if (isRunning.compareAndSet(true, false)) {
-      info("Shutting down")
-      if (isInterruptible)
-        interrupt()
-      true
-    } else
-      false
+    this.synchronized {
+      if (isRunning) {
+        info("Shutting down")
+        shutdownInitiated.countDown()
+        if (isInterruptible)
+          interrupt()
+        true
+      } else
+        false
+    }
   }
 
-    /**
+  /**
    * After calling initiateShutdown(), use this API to wait until the shutdown is complete
    */
   def awaitShutdown(): Unit = {
-    shutdownLatch.await()
-    info("Shutdown completed")
+    if (!isShutdownInitiated)
+      throw new IllegalStateException("initiateShutdown() was not called before awaitShutdown()")
+    else {
+      if (isStarted)
+        shutdownComplete.await()
+      info("Shutdown completed")
+    }
+  }
+
+  /**
+   *  Causes the current thread to wait until the shutdown is initiated,
+   *  or the specified waiting time elapses.
+   *
+   * @param timeout
+   * @param unit
+   */
+  def pause(timeout: Long, unit: TimeUnit): Unit = {
+    if (shutdownInitiated.await(timeout, unit))
+      trace("shutdownInitiated latch count reached zero. Shutdown called.")
   }
 
   /**
@@ -62,21 +89,25 @@ abstract class ShutdownableThread(val name: String, val isInterruptible: Boolean
   def doWork(): Unit
 
   override def run(): Unit = {
+    isStarted = true
     info("Starting")
     try {
-      while (isRunning.get)
+      while (isRunning)
         doWork()
     } catch {
       case e: FatalExitError =>
-        isRunning.set(false)
-        shutdownLatch.countDown()
+        shutdownInitiated.countDown()
+        shutdownComplete.countDown()
         info("Stopped")
         Exit.exit(e.statusCode())
       case e: Throwable =>
-        if (isRunning.get())
+        if (isRunning)
           error("Error due to", e)
+    } finally {
+       shutdownComplete.countDown()
     }
-    shutdownLatch.countDown()
     info("Stopped")
   }
+
+  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 e781cd6f7674b..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)
 
@@ -73,7 +73,7 @@ class Throttler(desiredRatePerSec: Double,
             time.sleep(sleepTime)
           }
         }
-        periodStartNs = now
+        periodStartNs = time.nanoseconds()
         observedSoFar = 0
       }
     }
@@ -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/ToolsUtils.scala b/core/src/main/scala/kafka/utils/ToolsUtils.scala
index 50e04f5719627..0f3de767fd806 100644
--- a/core/src/main/scala/kafka/utils/ToolsUtils.scala
+++ b/core/src/main/scala/kafka/utils/ToolsUtils.scala
@@ -52,12 +52,16 @@ object ToolsUtils {
         if (maxLengthOfDisplayName < mergedKeyName.length) {
           maxLengthOfDisplayName = mergedKeyName.length
         }
-        (mergedKeyName, value.value())
+        (mergedKeyName, value.metricValue)
     }
     println(s"\n%-${maxLengthOfDisplayName}s   %s".format("Metric Name", "Value"))
     sortedMap.foreach {
       case (metricName, value) =>
-        println(s"%-${maxLengthOfDisplayName}s : %.3f".format(metricName, value))
+        val specifier = value match {
+          case _ @ (_: java.lang.Float | _: java.lang.Double) => "%.3f"
+          case _ => "%s"
+        }
+        println(s"%-${maxLengthOfDisplayName}s : $specifier".format(metricName, value))
     }
   }
 }
diff --git a/core/src/main/scala/kafka/utils/TopicFilter.scala b/core/src/main/scala/kafka/utils/TopicFilter.scala
new file mode 100644
index 0000000000000..075c14722ee96
--- /dev/null
+++ b/core/src/main/scala/kafka/utils/TopicFilter.scala
@@ -0,0 +1,55 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package kafka.utils
+
+import java.util.regex.{Pattern, PatternSyntaxException}
+
+import org.apache.kafka.common.internals.Topic
+
+sealed abstract class TopicFilter(rawRegex: String) extends Logging {
+
+  val regex = rawRegex
+          .trim
+          .replace(',', '|')
+          .replace(" ", "")
+          .replaceAll("""^["']+""","")
+          .replaceAll("""["']+$""","") // property files may bring quotes
+
+  try {
+    Pattern.compile(regex)
+  }
+  catch {
+    case _: PatternSyntaxException =>
+      throw new RuntimeException(regex + " is an invalid regex.")
+  }
+
+  override def toString = regex
+
+  def isTopicAllowed(topic: String, excludeInternalTopics: Boolean): Boolean
+}
+
+case class IncludeList(rawRegex: String) extends TopicFilter(rawRegex) {
+  override def isTopicAllowed(topic: String, excludeInternalTopics: Boolean) = {
+    val allowed = topic.matches(regex) && !(Topic.isInternal(topic) && excludeInternalTopics)
+
+    debug("%s %s".format(
+      topic, if (allowed) "allowed" else "filtered"))
+
+    allowed
+  }
+}
diff --git a/core/src/main/scala/kafka/utils/VerifiableProperties.scala b/core/src/main/scala/kafka/utils/VerifiableProperties.scala
index de4f654c1ee59..bbacd95dc3b37 100755
--- a/core/src/main/scala/kafka/utils/VerifiableProperties.scala
+++ b/core/src/main/scala/kafka/utils/VerifiableProperties.scala
@@ -21,7 +21,7 @@ import java.util.Properties
 import java.util.Collections
 import scala.collection._
 import kafka.message.{CompressionCodec, NoCompressionCodec}
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 
 
 class VerifiableProperties(val props: Properties) extends Logging {
@@ -122,7 +122,7 @@ class VerifiableProperties(val props: Properties) extends Logging {
     require(v >= range._1 && v <= range._2, name + " has value " + v + " which is not in the range " + range + ".")
     v
   }
-  
+
   /**
    * Get a required argument as a double
    * @param name The property name
@@ -130,7 +130,7 @@ class VerifiableProperties(val props: Properties) extends Logging {
    * @throws IllegalArgumentException If the given property is not present
    */
   def getDouble(name: String): Double = getString(name).toDouble
-  
+
   /**
    * Get an optional argument as a double
    * @param name The property name
@@ -141,7 +141,7 @@ class VerifiableProperties(val props: Properties) extends Logging {
       getDouble(name)
     else
       default
-  } 
+  }
 
   /**
    * Read a boolean value from the properties instance
@@ -158,7 +158,7 @@ class VerifiableProperties(val props: Properties) extends Logging {
       v.toBoolean
     }
   }
-  
+
   def getBoolean(name: String) = getString(name).toBoolean
 
   /**
@@ -178,7 +178,7 @@ class VerifiableProperties(val props: Properties) extends Logging {
     require(containsKey(name), "Missing required property '" + name + "'")
     getProperty(name)
   }
-  
+
   /**
    * Get a Map[String, String] from a property list in the form k1:v2, k2:v2, ...
    */
@@ -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
new file mode 100644
index 0000000000000..9d3130e6685d3
--- /dev/null
+++ b/core/src/main/scala/kafka/utils/VersionInfo.scala
@@ -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 kafka.utils
+
+import org.apache.kafka.common.utils.AppInfoParser
+
+object VersionInfo {
+
+  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 60c2adfba6111..0000000000000
--- a/core/src/main/scala/kafka/utils/ZkUtils.scala
+++ /dev/null
@@ -1,1245 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package kafka.utils
-
-import java.util.concurrent.{CountDownLatch, TimeUnit}
-
-import kafka.admin._
-import kafka.api.{ApiVersion, KAFKA_0_10_0_IV1, LeaderAndIsr}
-import kafka.cluster._
-import kafka.common.{KafkaException, NoEpochForPartitionException, TopicAndPartition}
-import kafka.consumer.{ConsumerThreadId, TopicCount}
-import kafka.controller.{LeaderIsrAndControllerEpoch, ReassignedPartitionsContext}
-import kafka.metrics.KafkaMetricsGroup
-import kafka.server.ConfigType
-import kafka.utils.ZkUtils._
-
-import com.yammer.metrics.core.MetricName
-import org.I0Itec.zkclient.exception.{ZkBadVersionException, ZkException, 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.kafka.common.utils.Time
-import org.apache.zookeeper.AsyncCallback.{DataCallback, StringCallback}
-import org.apache.zookeeper.KeeperException.Code
-import org.apache.zookeeper.data.{ACL, Stat}
-import org.apache.zookeeper.{CreateMode, KeeperException, ZooDefs, ZooKeeper}
-
-import scala.collection._
-import scala.collection.JavaConverters._
-
-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 ProducerIdBlockPath = "/latest_producer_id_block"
-  // Important: it is necessary to add any new top level Zookeeper path to the Seq
-  val SecureZkRootPaths = Seq(AdminPath,
-                              BrokersPath,
-                              ClusterPath,
-                              ConfigPath,
-                              ControllerPath,
-                              ControllerEpochPath,
-                              IsrChangeNotificationPath,
-                              KafkaAclPath,
-                              KafkaAclChangesPath,
-                              ProducerIdBlockPath,
-                              LogDirEventNotificationPath)
-
-  // Important: it is necessary to add any new top level Zookeeper path that contains
-  //            sensitive information that should not be world readable to the Seq
-  val SensitiveZkRootPaths = Seq(ConfigUsersPath)
-
-  def withMetrics(zkUrl: String, sessionTimeout: Int, connectionTimeout: Int, isZkSecurityEnabled: Boolean,
-                  time: Time): ZkUtils = {
-    val (zkClient, zkConnection) = createZkClientAndConnection(zkUrl, sessionTimeout, connectionTimeout)
-    new ZkUtils(new ZooKeeperClientMetrics(zkClient, time), zkConnection, isZkSecurityEnabled)
-  }
-
-  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 = {
-    path != null && SensitiveZkRootPaths.exists(path.startsWith(_))
-  }
-
-  @deprecated("This is deprecated, use defaultAcls(isSecure, path) which doesn't make sensitive data world readable", since = "0.10.2.1")
-  def DefaultAcls(isSecure: Boolean): java.util.List[ACL] = defaultAcls(isSecure, "")
-
-  def defaultAcls(isSecure: Boolean, path: String): java.util.List[ACL] = {
-    if (isSecure) {
-      val list = new java.util.ArrayList[ACL]
-      list.addAll(ZooDefs.Ids.CREATOR_ALL_ACL)
-      if (!sensitivePath(path)) {
-        list.addAll(ZooDefs.Ids.READ_ACL_UNSAFE)
-      }
-      list
-    } else
-      ZooDefs.Ids.OPEN_ACL_UNSAFE
-  }
-
-  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
-
-  // Parses without deduplicating keys so the data can be checked before allowing reassignment to proceed
-  def parsePartitionReassignmentData(jsonData: String): Map[TopicAndPartition, Seq[Int]] = {
-    val seq = for {
-      js <- Json.parseFull(jsonData).toSeq
-      partitionsSeq <- js.asJsonObject.get("partitions").toSeq
-      p <- partitionsSeq.asJsonArray.iterator
-    } yield {
-      val partitionFields = p.asJsonObject
-      val topic = partitionFields("topic").to[String]
-      val partition = partitionFields("partition").to[Int]
-      val newReplicas = partitionFields("replicas").to[Seq[Int]]
-      TopicAndPartition(topic, partition) -> newReplicas
-    }
-    seq.toMap
-  }
-
-  def parseTopicsData(jsonData: String): Seq[String] = {
-    for {
-      js <- Json.parseFull(jsonData).toSeq
-      partitionsSeq <- js.asJsonObject.get("topics").toSeq
-      p <- partitionsSeq.asJsonArray.iterator
-    } yield p.asJsonObject("topic").to[String]
-  }
-
-  def controllerZkData(brokerId: Int, timestamp: Long): String = {
-    Json.encode(Map("version" -> 1, "brokerid" -> brokerId, "timestamp" -> timestamp.toString))
-  }
-
-  def preferredReplicaLeaderElectionZkData(partitions: scala.collection.Set[TopicAndPartition]): String = {
-    Json.encode(Map("version" -> 1, "partitions" -> partitions.map(tp => Map("topic" -> tp.topic, "partition" -> tp.partition))))
-  }
-
-  def formatAsReassignmentJson(partitionsToBeReassigned: Map[TopicAndPartition, Seq[Int]]): String = {
-    Json.encode(Map(
-      "version" -> 1,
-      "partitions" -> partitionsToBeReassigned.map { case (TopicAndPartition(topic, partition), replicas) =>
-        Map(
-          "topic" -> topic,
-          "partition" -> partition,
-          "replicas" -> replicas
-        )
-      }
-    ))
-  }
-
-}
-
-class ZooKeeperClientWrapper(val zkClient: ZkClient) {
-  def apply[T](method: ZkClient => T): T = method(zkClient)
-  def close(): Unit = {
-    if(zkClient != null)
-      zkClient.close()
-  }
-}
-
-class ZooKeeperClientMetrics(zkClient: ZkClient, val time: Time)
-    extends ZooKeeperClientWrapper(zkClient) with KafkaMetricsGroup {
-  val latencyMetric = newHistogram("ZooKeeperRequestLatencyMs")
-
-  override def metricName(name: String, metricTags: scala.collection.Map[String, String]): MetricName = {
-    explicitMetricName("kafka.server", "ZooKeeperClientMetrics", name, metricTags)
-  }
-
-  override def apply[T](method: ZkClient => T): T = {
-    val startNs = time.nanoseconds
-    val ret =
-      try method(zkClient)
-      finally latencyMetric.update(TimeUnit.NANOSECONDS.toMillis(time.nanoseconds - startNs))
-    ret
-  }
-
-  override def close(): Unit = {
-    if (latencyMetric != null)
-      removeMetric("ZooKeeperRequestLatencyMs")
-    super.close()
-  }
-}
-
-/**
- * Legacy class for interacting with ZooKeeper. Whenever possible, ``KafkaZkClient`` should be used instead.
- */
-class ZkUtils(zkClientWrap: ZooKeeperClientWrapper,
-              val zkConnection: ZkConnection,
-              val isSecure: Boolean) extends Logging {
-  // These are persistent ZK paths that should exist on kafka broker startup.
-  val persistentZkPaths = Seq(ConsumersPath,
-                              BrokerIdsPath,
-                              BrokerTopicsPath,
-                              ConfigChangesPath,
-                              getEntityConfigRootPath(ConfigType.Topic),
-                              getEntityConfigRootPath(ConfigType.Client),
-                              DeleteTopicsPath,
-                              BrokerSequenceIdPath,
-                              IsrChangeNotificationPath,
-                              ProducerIdBlockPath,
-                              LogDirEventNotificationPath)
-
-  /** Present for compatibility */
-  def this(zkClient: ZkClient, zkConnection: ZkConnection, isSecure: Boolean) =
-    this(new ZooKeeperClientWrapper(zkClient), zkConnection, isSecure)
-
-  // Visible for testing
-  val zkPath = new ZkPath(zkClientWrap)
-
-  import ZkUtils._
-
-  @deprecated("This is deprecated, use defaultAcls(path) which doesn't make sensitive data world readable", since = "0.10.2.1")
-  val DefaultAcls: java.util.List[ACL] = ZkUtils.defaultAcls(isSecure, "")
-
-  def defaultAcls(path: String): java.util.List[ACL] = ZkUtils.defaultAcls(isSecure, path)
-
-  def zkClient: ZkClient = zkClientWrap.zkClient
-
-  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("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) = {
-      val jsonMap = Map("version" -> "1", "id" -> id)
-      Json.encode(jsonMap)
-    }
-
-    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] = {
-    ReplicationUtils.getLeaderIsrAndEpochForPartition(this, topic, partition).map(_.leaderAndIsr)
-  }
-
-  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)
-  }
-
-  /**
-   * Register brokers with v4 json format (which includes multiple endpoints and rack) if
-   * the apiVersion is 0.10.0.X or above. Register the broker with v2 json format otherwise.
-   * Due to KAFKA-3100, 0.9.0.0 broker and old clients will break if JSON version is above 2.
-   * We include v2 to make it possible for the broker to migrate from 0.9.0.0 to 0.10.0.X or above without having to
-   * upgrade to 0.9.0.1 first (clients have to be upgraded to 0.9.0.1 in any case).
-   *
-   * This format also includes default endpoints for compatibility with older clients.
-   *
-   * @param id broker ID
-   * @param host broker host name
-   * @param port broker port
-   * @param advertisedEndpoints broker end points
-   * @param jmxPort jmx port
-   * @param rack broker rack
-   * @param apiVersion Kafka version the broker is running as
-   */
-  def registerBrokerInZk(id: Int,
-                         host: String,
-                         port: Int,
-                         advertisedEndpoints: Seq[EndPoint],
-                         jmxPort: Int,
-                         rack: Option[String],
-                         apiVersion: ApiVersion) {
-    val brokerIdPath = BrokerIdsPath + "/" + id
-    // see method documentation for reason why we do this
-    val version = if (apiVersion >= KAFKA_0_10_0_IV1) 4 else 2
-    val json = Broker.toJson(version, id, host, port, advertisedEndpoints, jmxPort, rack)
-    registerBrokerInZk(brokerIdPath, json)
-
-    info("Registered broker %d at path %s with addresses: %s".format(id, brokerIdPath, advertisedEndpoints.mkString(",")))
-  }
-
-  private def registerBrokerInZk(brokerIdPath: String, brokerInfo: String) {
-    try {
-      val zkCheckedEphemeral = new ZKCheckedEphemeral(brokerIdPath,
-                                                      brokerInfo,
-                                                      zkConnection.getZookeeper,
-                                                      isSecure)
-      zkClientWrap(_ => zkCheckedEphemeral.create())
-    } catch {
-      case _: ZkNodeExistsException =>
-        throw new RuntimeException("A broker is already registered on the path " + brokerIdPath
-                + ". This probably " + "indicates that you either have configured a brokerid that is already in use, or "
-                + "else you have shutdown this broker and restarted it faster than the zookeeper "
-                + "timeout so it appears to be re-registering.")
-    }
-  }
-
-  @deprecated("This method has been deprecated and will be removed in a future release.", "0.11.0.0")
-  def getConsumerPartitionOwnerPath(group: String, topic: String, partition: Int): String = {
-    val topicDirs = new ZKGroupTopicDirs(group, topic)
-    topicDirs.consumerOwnerDir + "/" + partition
-  }
-
-  def leaderAndIsrZkData(leaderAndIsr: LeaderAndIsr, controllerEpoch: Int): String = {
-    Json.encode(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.encode(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 (!zkClientWrap(zkClient => 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("conflict in " + path + " data: " + data + " stored data: " + storedData)
-          throw e
-        } else {
-          // otherwise, the creation succeeded, return normally
-          info(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 {
-      zkClientWrap(_.writeData(path, data))
-    } catch {
-      case _: ZkNoNodeException =>
-        createParentPath(path)
-        try {
-          zkPath.createPersistent(path, data, acl)
-        } catch {
-          case _: ZkNodeExistsException =>
-            zkClientWrap(_.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 = zkClientWrap(_.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 = zkClientWrap(_.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 persistent 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 {
-      zkClientWrap(_.writeData(path, data))
-    } catch {
-      case _: ZkNoNodeException =>
-        createParentPath(path)
-        zkPath.createEphemeral(path, data, acl)
-    }
-  }
-
-  def deletePath(path: String): Boolean = {
-    zkClientWrap(_.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 {
-      zkClientWrap(_.delete(path, expectedVersion))
-      true
-    } catch {
-      case _: ZkBadVersionException => false
-    }
-  }
-
-  def deletePathRecursive(path: String) {
-    zkClientWrap(_.deleteRecursive(path))
-  }
-
-  def subscribeDataChanges(path: String, listener: IZkDataListener): Unit =
-    zkClientWrap(_.subscribeDataChanges(path, listener))
-
-  def unsubscribeDataChanges(path: String, dataListener: IZkDataListener): Unit =
-    zkClientWrap(_.unsubscribeDataChanges(path, dataListener))
-
-  def subscribeStateChanges(listener: IZkStateListener): Unit =
-    zkClientWrap(_.subscribeStateChanges(listener))
-
-  def subscribeChildChanges(path: String, listener: IZkChildListener): Option[Seq[String]] =
-    Option(zkClientWrap(_.subscribeChildChanges(path, listener))).map(_.asScala)
-
-  def unsubscribeChildChanges(path: String, childListener: IZkChildListener): Unit =
-    zkClientWrap(_.unsubscribeChildChanges(path, childListener))
-
-  def unsubscribeAll(): Unit =
-    zkClientWrap(_.unsubscribeAll())
-
-  def readData(path: String): (String, Stat) = {
-    val stat: Stat = new Stat()
-    val dataStr: String = zkClientWrap(_.readData[String](path, stat))
-    (dataStr, stat)
-  }
-
-  def readDataMaybeNull(path: String): (Option[String], Stat) = {
-    val stat = new Stat()
-    val dataAndStat = try {
-                        val dataStr = zkClientWrap(_.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 = zkClientWrap(_.readData[String](path, stat))
-      (Option(data), stat.getVersion)
-    } catch {
-      case _: ZkNoNodeException => (None, stat.getVersion)
-    }
-  }
-
-  def getChildren(path: String): Seq[String] = zkClientWrap(_.getChildren(path)).asScala
-
-  def getChildrenParentMayNotExist(path: String): Seq[String] = {
-    try {
-      zkClientWrap(_.getChildren(path)).asScala
-    } catch {
-      case _: ZkNoNodeException => Nil
-    }
-  }
-
-  /**
-   * Check if the given path exists
-   */
-  def pathExists(path: String): Boolean = {
-    zkClientWrap(_.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(Broker.createBroker(node.toInt, brokerZKString))
-    }
-    cluster
-  }
-
-  def getPartitionLeaderAndIsrForTopics(topicAndPartitions: Set[TopicAndPartition]): mutable.Map[TopicAndPartition, LeaderIsrAndControllerEpoch] = {
-    val ret = new mutable.HashMap[TopicAndPartition, LeaderIsrAndControllerEpoch]
-    for(topicAndPartition <- topicAndPartitions) {
-      ReplicationUtils.getLeaderIsrAndEpochForPartition(this, topicAndPartition.topic, topicAndPartition.partition).foreach { leaderIsrAndControllerEpoch =>
-        ret.put(topicAndPartition, leaderIsrAndControllerEpoch)
-      }
-    }
-    ret
-  }
-
-  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(p => p._1 -> new ReassignedPartitionsContext(p._2))
-      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)
-      case None => Set.empty[TopicAndPartition]
-    }
-  }
-
-  def deletePartition(brokerId: Int, topic: String) {
-    val brokerIdPath = BrokerIdsPath + "/" + brokerId
-    zkClientWrap(_.delete(brokerIdPath))
-    val brokerPartTopicPath = ZkUtils.BrokerTopicsPath + "/" + topic + "/" + brokerId
-    zkClientWrap(_.delete(brokerPartTopicPath))
-  }
-
-  @deprecated("This method has been deprecated and will be removed in a future release.", "0.11.0.0")
-  def getConsumersInGroup(group: String): Seq[String] = {
-    val dirs = new ZKGroupDirs(group)
-    getChildren(dirs.consumerRegistryDir)
-  }
-
-  @deprecated("This method has been deprecated and will be removed in a future release.", "0.11.0.0")
-  def getConsumersPerTopic(group: String, excludeInternalTopics: Boolean): mutable.Map[String, List[ConsumerThreadId]] = {
-    val dirs = new ZKGroupDirs(group)
-    val consumers = getChildrenParentMayNotExist(dirs.consumerRegistryDir)
-    val consumersPerTopicMap = new mutable.HashMap[String, List[ConsumerThreadId]]
-    for (consumer <- consumers) {
-      val topicCount = TopicCount.constructTopicCount(group, consumer, this, excludeInternalTopics)
-      for ((topic, consumerThreadIdSet) <- topicCount.getConsumerThreadIdsPerTopic) {
-        for (consumerThreadId <- consumerThreadIdSet)
-          consumersPerTopicMap.get(topic) match {
-            case Some(curConsumers) => consumersPerTopicMap.put(topic, consumerThreadId :: curConsumers)
-            case _ => consumersPerTopicMap.put(topic, List(consumerThreadId))
-          }
-      }
-    }
-    for ((topic, consumerList) <- consumersPerTopicMap)
-      consumersPerTopicMap.put(topic, consumerList.sortWith((s, t) => s < t))
-    consumersPerTopicMap
-  }
-
-  @deprecated("This method has been deprecated and will be removed in a future release.", "0.11.0.0")
-  def getTopicsPerMemberId(group: String, excludeInternalTopics: Boolean = true): Map[String, List[String]] = {
-    val dirs = new ZKGroupDirs(group)
-    val memberIds = getChildrenParentMayNotExist(dirs.consumerRegistryDir)
-    memberIds.map { memberId =>
-      val topicCount = TopicCount.constructTopicCount(group, memberId, this, excludeInternalTopics)
-      memberId -> topicCount.getTopicCountMap.keys.toList
-    }.toMap
-  }
-
-  /**
-   * 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(Broker.createBroker(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 = zkClientWrap(_.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
-    }
-  }
-
-  @deprecated("This method has been deprecated and will be removed in a future release.", "0.11.0.0")
-  def getConsumerGroups() = {
-    getChildren(ConsumersPath)
-  }
-
-  @deprecated("This method has been deprecated and will be removed in a future release.", "0.11.0.0")
-  def getTopicsByConsumerGroup(consumerGroup:String) = {
-    getChildrenParentMayNotExist(new ZKGroupDirs(consumerGroup).consumerGroupOwnersDir)
-  }
-
-  @deprecated("This method has been deprecated and will be removed in a future release.", "0.11.0.0")
-  def getAllConsumerGroupsForTopic(topic: String): Set[String] = {
-    val groups = getChildrenParentMayNotExist(ConsumersPath)
-    if (groups == null) Set.empty
-    else {
-      groups.foldLeft(Set.empty[String]) {(consumerGroupsForTopic, group) =>
-        val topics = getChildren(new ZKGroupDirs(group).consumerGroupOffsetsDir)
-        if (topics.contains(topic)) consumerGroupsForTopic + group
-        else consumerGroupsForTopic
-      }
-    }
-  }
-
-  def close() {
-    zkClientWrap.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")
-  }
-}
-
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.11.0.0")
-class ZKGroupDirs(val group: String) {
-  def consumerDir = ConsumersPath
-  def consumerGroupDir = consumerDir + "/" + group
-  def consumerRegistryDir = consumerGroupDir + "/ids"
-  def consumerGroupOffsetsDir = consumerGroupDir + "/offsets"
-  def consumerGroupOwnersDir = consumerGroupDir + "/owners"
-}
-
-@deprecated("This class has been deprecated and will be removed in a future release.", "0.11.0.0")
-class ZKGroupTopicDirs(group: String, topic: String) extends ZKGroupDirs(group) {
-  def consumerOffsetDir = consumerGroupOffsetsDir + "/" + topic
-  def consumerOwnerDir = consumerGroupOwnersDir + "/" + topic
-}
-
-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(clientWrap: ZooKeeperClientWrapper) {
-
-  @volatile private var isNamespacePresent: Boolean = false
-
-  def checkNamespace() {
-    if (isNamespacePresent)
-      return
-
-    if (!clientWrap(_.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()
-    clientWrap(_.createPersistent(path, data, acls))
-  }
-
-  def createPersistent(path: String, createParents: Boolean, acls: java.util.List[ACL]) {
-    checkNamespace()
-    clientWrap(_.createPersistent(path, createParents, acls))
-  }
-
-  def createEphemeral(path: String, data: Object, acls: java.util.List[ACL]) {
-    checkNamespace()
-    clientWrap(_.createEphemeral(path, data, acls))
-  }
-
-  def createPersistentSequential(path: String, data: Object, acls: java.util.List[ACL]): String = {
-    checkNamespace()
-    clientWrap(_.createPersistentSequential(path, data, acls))
-  }
-}
-
-/**
- * Creates an ephemeral znode checking the session owner
- * in the case of conflict. In the regular case, the
- * znode is created and the create call returns OK. If
- * the call receives a node exists event, then it checks
- * if the session matches. If it does, then it returns OK,
- * and otherwise it fails the operation.
- */
-
-class ZKCheckedEphemeral(path: String,
-                         data: String,
-                         zkHandle: ZooKeeper,
-                         isSecure: Boolean) extends Logging {
-  private val createCallback = new CreateCallback
-  private val getDataCallback = new GetDataCallback
-  val latch: CountDownLatch = new CountDownLatch(1)
-  var result: Code = Code.OK
-  val defaultAcls = ZkUtils.defaultAcls(isSecure, path)
-
-  private class CreateCallback extends StringCallback {
-    def processResult(rc: Int,
-                      path: String,
-                      ctx: Object,
-                      name: String) {
-      Code.get(rc) match {
-        case Code.OK =>
-           setResult(Code.OK)
-        case Code.CONNECTIONLOSS =>
-          // try again
-          createEphemeral
-        case Code.NONODE =>
-          error("No node for path %s (could be the parent missing)".format(path))
-          setResult(Code.NONODE)
-        case Code.NODEEXISTS =>
-          zkHandle.getData(path, false, getDataCallback, null)
-        case Code.SESSIONEXPIRED =>
-          error("Session has expired while creating %s".format(path))
-          setResult(Code.SESSIONEXPIRED)
-        case Code.INVALIDACL =>
-          error("Invalid ACL")
-          setResult(Code.INVALIDACL)
-        case _ =>
-          warn("ZooKeeper event while creating registration node: %s %s".format(path, Code.get(rc)))
-          setResult(Code.get(rc))
-      }
-    }
-  }
-
-  private class GetDataCallback extends DataCallback {
-      def processResult(rc: Int,
-                        path: String,
-                        ctx: Object,
-                        readData: Array[Byte],
-                        stat: Stat) {
-        Code.get(rc) match {
-          case Code.OK =>
-                if (stat.getEphemeralOwner != zkHandle.getSessionId)
-                  setResult(Code.NODEEXISTS)
-                else
-                  setResult(Code.OK)
-          case Code.NONODE =>
-            info("The ephemeral node [%s] at %s has gone away while reading it, ".format(data, path))
-            createEphemeral
-          case Code.SESSIONEXPIRED =>
-            error("Session has expired while reading znode %s".format(path))
-            setResult(Code.SESSIONEXPIRED)
-          case Code.INVALIDACL =>
-            error("Invalid ACL")
-            setResult(Code.INVALIDACL)
-          case _ =>
-            warn("ZooKeeper event while getting znode data: %s %s".format(path, Code.get(rc)))
-            setResult(Code.get(rc))
-        }
-      }
-  }
-
-  private def createEphemeral() {
-    zkHandle.create(path,
-                    ZKStringSerializer.serialize(data),
-                    defaultAcls,
-                    CreateMode.EPHEMERAL,
-                    createCallback,
-                    null)
-  }
-
-  private def createRecursive(prefix: String, suffix: String) {
-    debug("Path: %s, Prefix: %s, Suffix: %s".format(path, prefix, suffix))
-    if(suffix.isEmpty()) {
-      createEphemeral
-    } else {
-      zkHandle.create(prefix,
-                      new Array[Byte](0),
-                      defaultAcls,
-                      CreateMode.PERSISTENT,
-                      new StringCallback() {
-                        def processResult(rc : Int,
-                                          path : String,
-                                          ctx : Object,
-                                          name : String) {
-                          Code.get(rc) match {
-                            case Code.OK | Code.NODEEXISTS =>
-                              // Nothing to do
-                            case Code.CONNECTIONLOSS =>
-                              // try again
-                              val suffix = ctx.asInstanceOf[String]
-                              createRecursive(path, suffix)
-                            case Code.NONODE =>
-                              error("No node for path %s (could be the parent missing)".format(path))
-                              setResult(Code.get(rc))
-                            case Code.SESSIONEXPIRED =>
-                              error("Session has expired while creating %s".format(path))
-                              setResult(Code.get(rc))
-                            case Code.INVALIDACL =>
-                              error("Invalid ACL")
-                              setResult(Code.INVALIDACL)
-                            case _ =>
-                              warn("ZooKeeper event while creating registration node: %s %s".format(path, Code.get(rc)))
-                              setResult(Code.get(rc))
-                          }
-                        }
-                  },
-                  suffix)
-      // Update prefix and suffix
-      val index = suffix.indexOf('/', 1) match {
-        case -1 => suffix.length
-        case x : Int => x
-      }
-      // Get new prefix
-      val newPrefix = prefix + suffix.substring(0, index)
-      // Get new suffix
-      val newSuffix = suffix.substring(index, suffix.length)
-      createRecursive(newPrefix, newSuffix)
-    }
-  }
-
-  private def setResult(code: Code) {
-    result = code
-    latch.countDown()
-  }
-
-  private def waitUntilResolved(): Code = {
-    latch.await()
-    result
-  }
-
-  def create() {
-    val index = path.indexOf('/', 1) match {
-        case -1 => path.length
-        case x : Int => x
-    }
-    val prefix = path.substring(0, index)
-    val suffix = path.substring(index, path.length)
-    debug(s"Path: $path, Prefix: $prefix, Suffix: $suffix")
-    info(s"Creating $path (is it secure? $isSecure)")
-    createRecursive(prefix, suffix)
-    val result = waitUntilResolved()
-    info("Result of znode creation is: %s".format(result))
-    result match {
-      case Code.OK =>
-        // Nothing to do
-      case _ =>
-        throw ZkException.create(KeeperException.create(result))
-    }
-  }
-}
diff --git a/core/src/main/scala/kafka/utils/json/DecodeJson.scala b/core/src/main/scala/kafka/utils/json/DecodeJson.scala
index eab1f2a27891a..71bfd61cee11e 100644
--- a/core/src/main/scala/kafka/utils/json/DecodeJson.scala
+++ b/core/src/main/scala/kafka/utils/json/DecodeJson.scala
@@ -17,10 +17,9 @@
 
 package kafka.utils.json
 
-import scala.collection._
-import scala.language.higherKinds
-import JavaConverters._
-import generic.CanBuildFrom
+import scala.collection.{Map, Seq}
+import scala.collection.compat._
+import scala.jdk.CollectionConverters._
 
 import com.fasterxml.jackson.databind.{JsonMappingException, JsonNode}
 
@@ -43,11 +42,7 @@ trait DecodeJson[T] {
   def decode(node: JsonNode): T =
     decodeEither(node) match {
       case Right(x) => x
-      case Left(x) =>
-        // Non-deprecated constructors were only introduced in Jackson 2.7, so stick with the deprecated one in case
-        // people have older versions of Jackson in their classpath. Once the Scala clients are removed, we can loosen
-        // this restriction.
-        throw new JsonMappingException(x)
+      case Left(x) => throw new JsonMappingException(null, x)
     }
 
 }
@@ -85,33 +80,27 @@ object DecodeJson {
       if (node.isTextual) Right(node.textValue) else Left(s"Expected `String` value, received $node")
   }
 
-  implicit def decodeOption[E](implicit decodeJson: DecodeJson[E]): DecodeJson[Option[E]] = new DecodeJson[Option[E]] {
-    def decodeEither(node: JsonNode): Either[String, Option[E]] = {
-      if (node.isNull) Right(None)
-      else decodeJson.decodeEither(node).right.map(Some(_))
-    }
+  implicit def decodeOption[E](implicit decodeJson: DecodeJson[E]): DecodeJson[Option[E]] = (node: JsonNode) => {
+    if (node.isNull) Right(None)
+    else decodeJson.decodeEither(node).map(Some(_))
   }
 
-  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]] {
-    def decodeEither(node: JsonNode): Either[String, S[E]] = {
-      if (node.isArray)
-        decodeIterator(node.elements.asScala)(decodeJson.decodeEither)
-      else Left(s"Expected JSON array, received $node")
-    }
+  implicit def decodeSeq[E, S[+T] <: Seq[E]](implicit decodeJson: DecodeJson[E], factory: Factory[E, S[E]]): DecodeJson[S[E]] = (node: JsonNode) => {
+    if (node.isArray)
+      decodeIterator(node.elements.asScala)(decodeJson.decodeEither)
+    else Left(s"Expected JSON array, received $node")
   }
 
-  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]] {
-    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)
-      else Left(s"Expected JSON object, received $node")
-    }
+  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]] = (node: JsonNode) => {
+    if (node.isObject)
+      decodeIterator(node.fields.asScala)(e => decodeJson.decodeEither(e.getValue).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 {
+      f(it.next()) match {
         case Right(x) => result += x
         case Left(x) => return Left(x)
       }
diff --git a/core/src/main/scala/kafka/utils/json/JsonArray.scala b/core/src/main/scala/kafka/utils/json/JsonArray.scala
index 0b81a4ae13853..c22eda8651d75 100644
--- a/core/src/main/scala/kafka/utils/json/JsonArray.scala
+++ b/core/src/main/scala/kafka/utils/json/JsonArray.scala
@@ -17,8 +17,8 @@
 
 package kafka.utils.json
 
-import scala.collection.{Iterator, JavaConverters}
-import JavaConverters._
+import scala.collection.Iterator
+import scala.jdk.CollectionConverters._
 
 import com.fasterxml.jackson.databind.node.ArrayNode
 
diff --git a/core/src/main/scala/kafka/utils/json/JsonObject.scala b/core/src/main/scala/kafka/utils/json/JsonObject.scala
index 8feb08b9a9e62..9bf91ae1a6b0a 100644
--- a/core/src/main/scala/kafka/utils/json/JsonObject.scala
+++ b/core/src/main/scala/kafka/utils/json/JsonObject.scala
@@ -19,7 +19,7 @@ package kafka.utils.json
 
 import com.fasterxml.jackson.databind.JsonMappingException
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 
 import com.fasterxml.jackson.databind.node.ObjectNode
 
@@ -31,7 +31,7 @@ import scala.collection.Iterator
 class JsonObject private[json] (protected val node: ObjectNode) extends JsonValue {
 
   def apply(name: String): JsonValue =
-    get(name).getOrElse(throw new JsonMappingException(s"No such field exists: `$name`"))
+    get(name).getOrElse(throw new JsonMappingException(null, s"No such field exists: `$name`"))
 
   def get(name: String): Option[JsonValue] = Option(node.get(name)).map(JsonValue(_))
 
diff --git a/core/src/main/scala/kafka/utils/json/JsonValue.scala b/core/src/main/scala/kafka/utils/json/JsonValue.scala
index cbc82c03d1e14..ff62c6c12d138 100644
--- a/core/src/main/scala/kafka/utils/json/JsonValue.scala
+++ b/core/src/main/scala/kafka/utils/json/JsonValue.scala
@@ -59,7 +59,7 @@ trait JsonValue {
    * If this is a JSON object, return an instance of JsonObject. Otherwise, throw a JsonMappingException.
    */
   def asJsonObject: JsonObject =
-    asJsonObjectOption.getOrElse(throw new JsonMappingException(s"Expected JSON object, received $node"))
+    asJsonObjectOption.getOrElse(throw new JsonMappingException(null, s"Expected JSON object, received $node"))
 
   /**
    * If this is a JSON object, return a JsonObject wrapped by a `Some`. Otherwise, return None.
@@ -76,7 +76,7 @@ trait JsonValue {
    * If this is a JSON array, return an instance of JsonArray. Otherwise, throw a JsonMappingException.
    */
   def asJsonArray: JsonArray =
-    asJsonArrayOption.getOrElse(throw new JsonMappingException(s"Expected JSON array, received $node"))
+    asJsonArrayOption.getOrElse(throw new JsonMappingException(null, s"Expected JSON array, received $node"))
 
   /**
    * If this is a JSON array, return a JsonArray wrapped by a `Some`. Otherwise, return None.
diff --git a/core/src/main/scala/kafka/utils/timer/Timer.scala b/core/src/main/scala/kafka/utils/timer/Timer.scala
index ae8caff863658..6973dcb2a7986 100644
--- a/core/src/main/scala/kafka/utils/timer/Timer.scala
+++ b/core/src/main/scala/kafka/utils/timer/Timer.scala
@@ -16,9 +16,9 @@
  */
 package kafka.utils.timer
 
-import java.util.concurrent.{DelayQueue, Executors, ThreadFactory, TimeUnit}
 import java.util.concurrent.atomic.AtomicInteger
 import java.util.concurrent.locks.ReentrantReadWriteLock
+import java.util.concurrent.{DelayQueue, Executors, TimeUnit}
 
 import kafka.utils.threadsafe
 import org.apache.kafka.common.utils.{KafkaThread, Time}
@@ -58,10 +58,8 @@ class SystemTimer(executorName: String,
                   startMs: Long = Time.SYSTEM.hiResClockMs) extends Timer {
 
   // timeout timer
-  private[this] val taskExecutor = Executors.newFixedThreadPool(1, new ThreadFactory() {
-    def newThread(runnable: Runnable): Thread =
-      KafkaThread.nonDaemon("executor-"+executorName, runnable)
-  })
+  private[this] val taskExecutor = Executors.newFixedThreadPool(1,
+    (runnable: Runnable) => KafkaThread.nonDaemon("executor-" + executorName, runnable))
 
   private[this] val delayQueue = new DelayQueue[TimerTaskList]()
   private[this] val taskCounter = new AtomicInteger(0)
@@ -95,8 +93,6 @@ class SystemTimer(executorName: String,
     }
   }
 
-  private[this] val reinsert = (timerTaskEntry: TimerTaskEntry) => addTimerTaskEntry(timerTaskEntry)
-
   /*
    * Advances the clock if there is an expired bucket. If there isn't any expired bucket when called,
    * waits up to timeoutMs before giving up.
@@ -107,8 +103,8 @@ class SystemTimer(executorName: String,
       writeLock.lock()
       try {
         while (bucket != null) {
-          timingWheel.advanceClock(bucket.getExpiration())
-          bucket.flush(reinsert)
+          timingWheel.advanceClock(bucket.getExpiration)
+          bucket.flush(addTimerTaskEntry)
           bucket = delayQueue.poll()
         }
       } finally {
@@ -122,9 +118,8 @@ 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/utils/timer/TimerTask.scala b/core/src/main/scala/kafka/utils/timer/TimerTask.scala
index 66238540d34d6..a1995c1df0a68 100644
--- a/core/src/main/scala/kafka/utils/timer/TimerTask.scala
+++ b/core/src/main/scala/kafka/utils/timer/TimerTask.scala
@@ -40,8 +40,6 @@ trait TimerTask extends Runnable {
     }
   }
 
-  private[timer] def getTimerTaskEntry(): TimerTaskEntry = {
-    timerTaskEntry
-  }
+  private[timer] def getTimerTaskEntry: TimerTaskEntry = timerTaskEntry
 
 }
diff --git a/core/src/main/scala/kafka/utils/timer/TimerTaskList.scala b/core/src/main/scala/kafka/utils/timer/TimerTaskList.scala
index 3dbfa8f3882c2..efd0480618090 100644
--- a/core/src/main/scala/kafka/utils/timer/TimerTaskList.scala
+++ b/core/src/main/scala/kafka/utils/timer/TimerTaskList.scala
@@ -43,9 +43,7 @@ private[timer] class TimerTaskList(taskCounter: AtomicInteger) extends Delayed {
   }
 
   // Get the bucket's expiration time
-  def getExpiration(): Long = {
-    expiration.get()
-  }
+  def getExpiration: Long = expiration.get
 
   // Apply the supplied function to each of tasks in this list
   def foreach(f: (TimerTask)=>Unit): Unit = {
@@ -105,7 +103,7 @@ private[timer] class TimerTaskList(taskCounter: AtomicInteger) extends Delayed {
   }
 
   // Remove all task entries and apply the supplied function to each of them
-  def flush(f: (TimerTaskEntry)=>Unit): Unit = {
+  def flush(f: TimerTaskEntry => Unit): Unit = {
     synchronized {
       var head = root.next
       while (head ne root) {
@@ -122,12 +120,8 @@ private[timer] class TimerTaskList(taskCounter: AtomicInteger) extends Delayed {
   }
 
   def compareTo(d: Delayed): Int = {
-
     val other = d.asInstanceOf[TimerTaskList]
-
-    if(getExpiration < other.getExpiration) -1
-    else if(getExpiration > other.getExpiration) 1
-    else 0
+    java.lang.Long.compare(getExpiration, other.getExpiration)
   }
 
 }
@@ -159,7 +153,7 @@ private[timer] class TimerTaskEntry(val timerTask: TimerTask, val expirationMs:
   }
 
   override def compare(that: TimerTaskEntry): Int = {
-    this.expirationMs compare that.expirationMs
+    java.lang.Long.compare(expirationMs, that.expirationMs)
   }
 }
 
diff --git a/core/src/main/scala/kafka/zk/AdminZkClient.scala b/core/src/main/scala/kafka/zk/AdminZkClient.scala
index e00b8e6a55a09..1e994b4ecac85 100644
--- a/core/src/main/scala/kafka/zk/AdminZkClient.scala
+++ b/core/src/main/scala/kafka/zk/AdminZkClient.scala
@@ -20,16 +20,24 @@ 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._
-import org.apache.kafka.common.TopicPartition
+import kafka.utils.Implicits._
+import org.apache.kafka.common.{TopicPartition, Uuid}
 import org.apache.kafka.common.errors._
 import org.apache.kafka.common.internals.Topic
 import org.apache.zookeeper.KeeperException.NodeExistsException
 
 import scala.collection.{Map, Seq}
 
+/**
+ * Provides admin related methods for interacting with ZooKeeper.
+ *
+ * This is an internal class and no compatibility guarantees are provided,
+ * see org.apache.kafka.clients.admin.AdminClient for publicly supported APIs.
+ */
 class AdminZkClient(zkClient: KafkaZkClient) extends Logging {
 
   /**
@@ -44,10 +52,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)
-    createOrUpdateTopicPartitionAssignmentPathInZK(topic, replicaAssignment, topicConfig)
+    createTopicWithAssignment(topic, topicConfig, replicaAssignment)
   }
 
   /**
@@ -75,55 +83,56 @@ class AdminZkClient(zkClient: KafkaZkClient) extends Logging {
   }
 
   /**
-   * Creates or Updates the partition assignment for a given topic
-   * @param topic
-   * @param partitionReplicaAssignment
-   * @param config
-   * @param update
+   * Create topic and optionally validate its parameters. Note that this method is used by the
+   * TopicCommand as well.
+   *
+   * @param topic The name of the topic
+   * @param config The config of the topic
+   * @param partitionReplicaAssignment The assignments of the topic
+   * @param validate Boolean indicating if parameters must be validated or not (true by default)
    */
-  def createOrUpdateTopicPartitionAssignmentPathInZK(topic: String,
-                                                     partitionReplicaAssignment: Map[Int, Seq[Int]],
-                                                     config: Properties = new Properties,
-                                                     update: Boolean = false) {
-    validateCreateOrUpdateTopic(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
-      zkClient.setOrCreateEntityConfigs(ConfigType.Topic, topic, config)
-    }
+  def createTopicWithAssignment(topic: String,
+                                config: Properties,
+                                partitionReplicaAssignment: Map[Int, Seq[Int]],
+                                validate: Boolean = true): Unit = {
+    if (validate)
+      validateTopicCreate(topic, partitionReplicaAssignment, config)
+
+    info(s"Creating topic $topic with configuration $config and initial partition " +
+      s"assignment $partitionReplicaAssignment")
+
+    // write out the config if there is any, this isn't transactional with the partition assignments
+    zkClient.setOrCreateEntityConfigs(ConfigType.Topic, topic, config)
 
     // create the partition assignment
-    writeTopicPartitionAssignment(topic, partitionReplicaAssignment, update)
+    writeTopicPartitionAssignment(topic, partitionReplicaAssignment.map { case (k, v) => k -> ReplicaAssignment(v) },
+      isUpdate = false)
   }
 
   /**
-   * Validate method to use before the topic creation or update
-   * @param topic
-   * @param partitionReplicaAssignment
-   * @param config
-   * @param update
+   * Validate topic creation parameters. Note that this method is indirectly used by the
+   * TopicCommand via the `createTopicWithAssignment` method.
+   *
+   * @param topic The name of the topic
+   * @param partitionReplicaAssignment The assignments of the topic
+   * @param config The config of the topic
    */
-  def validateCreateOrUpdateTopic(topic: String,
-                                  partitionReplicaAssignment: Map[Int, Seq[Int]],
-                                  config: Properties,
-                                  update: Boolean): Unit = {
-    // validate arguments
+  def validateTopicCreate(topic: String,
+                          partitionReplicaAssignment: Map[Int, Seq[Int]],
+                          config: Properties): Unit = {
     Topic.validate(topic)
 
-    if (!update) {
-      if (zkClient.topicExists(topic))
+    if (zkClient.topicExists(topic))
+      throw new TopicExistsException(s"Topic '$topic' already exists.")
+    else if (Topic.hasCollisionChars(topic)) {
+      val allTopics = zkClient.getAllTopicsInCluster()
+      // 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.")
-      else if (Topic.hasCollisionChars(topic)) {
-        val allTopics = zkClient.getAllTopicsInCluster
-        // 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(", ")}")
-        }
+      val collidingTopics = allTopics.filter(Topic.hasCollision(topic, _))
+      if (collidingTopics.nonEmpty) {
+        throw new InvalidTopicException(s"Topic '$topic' collides with existing topics: ${collidingTopics.mkString(", ")}")
       }
     }
 
@@ -135,21 +144,25 @@ class AdminZkClient(zkClient: KafkaZkClient) extends Logging {
         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)
+    val partitionSize = partitionReplicaAssignment.size
+    val sequenceSum = partitionSize * (partitionSize - 1) / 2
+    if (partitionReplicaAssignment.size != partitionReplicaAssignment.toSet.size ||
+        partitionReplicaAssignment.keys.filter(_ >= 0).sum != sequenceSum)
+        throw new InvalidReplicaAssignmentException("partitions should be a consecutive 0-based integer sequence")
+
+    LogConfig.validate(config)
   }
 
-  private def writeTopicPartitionAssignment(topic: String, replicaAssignment: Map[Int, Seq[Int]], update: Boolean) {
+  private 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 (!update) {
-        info("Topic creation " + assignment)
-        zkClient.createTopicAssignment(topic, assignment)
+      if (!isUpdate) {
+        val topicId = Uuid.randomUuid()
+        zkClient.createTopicAssignment(topic, topicId, assignment.map { case (k, v) => k -> v.replicas })
       } else {
-        info("Topic update " + assignment)
-        zkClient.setTopicAssignment(topic, assignment)
+        val topicIds = zkClient.getTopicIdsForTopics(Set(topic))
+        zkClient.setTopicAssignment(topic, topicIds(topic), assignment)
       }
       debug("Updated path %s with %s for replica assignment".format(TopicZNode.path(topic), assignment))
     } catch {
@@ -158,12 +171,11 @@ 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)
@@ -178,26 +190,61 @@ class AdminZkClient(zkClient: KafkaZkClient) extends Logging {
   }
 
   /**
-  * 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 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
-  */
+   * Add partitions to existing topic with optional replica assignment. Note that this
+   * method is used by the TopicCommand.
+   *
+   * @param topic Topic for adding partitions to
+   * @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
+   * @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]] = {
+
+    val proposedAssignmentForNewPartitions = createNewPartitionsAssignment(
+      topic,
+      existingAssignment,
+      allBrokers,
+      numPartitions,
+      replicaAssignment
+    )
+
+    if (validateOnly) {
+      (existingAssignment ++ proposedAssignmentForNewPartitions)
+        .map { case (k, v) => k -> v.replicas }
+    } else {
+      createPartitionsWithAssignment(topic, existingAssignment, proposedAssignmentForNewPartitions)
+        .map { case (k, v) => k -> v.replicas }
+    }
+  }
+
+  /**
+   * Create assignment to add the given number of partitions while validating the
+   * provided arguments.
+   *
+   * @param topic Topic for adding partitions to
+   * @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
+   * @return the assignment for the new partitions
+   */
+  def createNewPartitionsAssignment(topic: String,
+                                    existingAssignment: Map[Int, ReplicaAssignment],
+                                    allBrokers: Seq[BrokerMetadata],
+                                    numPartitions: Int = 1,
+                                    replicaAssignment: Option[Map[Int, Seq[Int]]] = None): Map[Int, ReplicaAssignment] = {
     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)
@@ -216,22 +263,40 @@ class AdminZkClient(zkClient: KafkaZkClient) extends Logging {
       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
-      createOrUpdateTopicPartitionAssignmentPathInZK(topic, proposedAssignment, update = true)
+
+    proposedAssignmentForNewPartitions.map { case (tp, replicas) =>
+      tp -> ReplicaAssignment(replicas, List(), List())
     }
-    proposedAssignment
+  }
+
+  /**
+   * Add partitions to the existing topic with the provided assignment. This method does
+   * not validate the provided assignments. Validation must be done beforehand.
+   *
+   * @param topic Topic for adding partitions to
+   * @param existingAssignment A map from partition id to its assignment
+   * @param newPartitionAssignment The assignments to add
+   * @return the updated replica assignment
+   */
+  def createPartitionsWithAssignment(topic: String,
+                                     existingAssignment: Map[Int, ReplicaAssignment],
+                                     newPartitionAssignment: Map[Int, ReplicaAssignment]): Map[Int, ReplicaAssignment] = {
+
+    info(s"Creating ${newPartitionAssignment.size} partitions for '$topic' with the following replica assignment: " +
+      s"$newPartitionAssignment.")
 
+    val combinedAssignment = existingAssignment ++ newPartitionAssignment
+
+    writeTopicPartitionAssignment(topic, combinedAssignment, isUpdate = true)
+
+    combinedAssignment
   }
 
   private def validateReplicaAssignment(replicaAssignment: Map[Int, Seq[Int]],
                                         expectedReplicationFactor: Int,
                                         availableBrokerIds: Set[Int]): Unit = {
 
-    replicaAssignment.foreach { case (partitionId, replicas) =>
+    replicaAssignment.forKeyValue { (partitionId, replicas) =>
       if (replicas.isEmpty)
         throw new InvalidReplicaAssignmentException(
           s"Cannot have replication factor of 0 for partition id $partitionId.")
@@ -254,11 +319,23 @@ class AdminZkClient(zkClient: KafkaZkClient) extends Logging {
       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 ${expectedReplicationFactor} while partitions [${partitions.mkString(", ")}] have " +
+        s"partition 0 has $expectedReplicationFactor while partitions [${partitions.mkString(", ")}] have " +
         s"replication factors [${repFactors.mkString(", ")}], respectively.")
     }
   }
 
+  def parseBroker(broker: String): Option[Int] = {
+    broker match {
+      case ConfigEntityName.Default => None
+      case _ =>
+        try Some(broker.toInt)
+        catch {
+          case _: NumberFormatException =>
+            throw new IllegalArgumentException(s"Error parsing broker $broker. The broker's Entity Name must be a single integer value")
+        }
+    }
+  }
+
   /**
    * Change the configs for a given entityType and entityName
    * @param entityType
@@ -267,20 +344,13 @@ class AdminZkClient(zkClient: KafkaZkClient) extends Logging {
    */
   def changeConfigs(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(entityName, configs)
       case ConfigType.Client => changeClientIdConfig(entityName, configs)
       case ConfigType.User => changeUserOrUserClientIdConfig(entityName, configs)
-      case ConfigType.Broker => changeBrokerConfig(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}")
+      case ConfigType.Broker => changeBrokerConfig(parseBroker(entityName), configs)
+      case ConfigType.Ip => changeIpConfig(entityName, configs)
+      case _ => throw new IllegalArgumentException(s"$entityType is not a known entityType. Should be one of ${ConfigType.all}")
     }
   }
 
@@ -294,7 +364,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)
   }
@@ -309,7 +379,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
@@ -317,6 +387,28 @@ class AdminZkClient(zkClient: KafkaZkClient) extends Logging {
     changeEntityConfig(ConfigType.User, sanitizedEntityName, configs)
   }
 
+  /**
+   * Validates the IP configs.
+   * @param ip ip for which configs are being validated
+   * @param configs properties to validate for the IP
+   */
+  def validateIpConfig(ip: String, configs: Properties): Unit = {
+    if (!DynamicConfig.Ip.isValidIpEntity(ip))
+      throw new AdminOperationException(s"$ip is not a valid IP or resolvable host.")
+    DynamicConfig.Ip.validate(configs)
+  }
+
+  /**
+   * Update the config for an IP. These overrides will be persisted between sessions, and will override any default
+   * IP properties.
+   * @param ip ip for which configs are being updated
+   * @param configs properties to update for the IP
+   */
+  def changeIpConfig(ip: String, configs: Properties): Unit = {
+    validateIpConfig(ip, configs)
+    changeEntityConfig(ConfigType.Ip, ip, configs)
+  }
+
   /**
    * validates the topic configs
    * @param topic
@@ -325,7 +417,7 @@ class AdminZkClient(zkClient: KafkaZkClient) extends Logging {
   def validateTopicConfig(topic: String, configs: Properties): Unit = {
     Topic.validate(topic)
     if (!zkClient.topicExists(topic))
-      throw new AdminOperationException("Topic \"%s\" does not exist.".format(topic))
+      throw new UnknownTopicOrPartitionException(s"Topic '$topic' does not exist.")
     // remove the topic overrides
     LogConfig.validate(configs)
   }
@@ -351,12 +443,34 @@ class AdminZkClient(zkClient: KafkaZkClient) extends Logging {
     * @param configs: The config to change, as properties
     */
   def changeBrokerConfig(brokers: Seq[Int], configs: Properties): Unit = {
-    DynamicConfig.Broker.validate(configs)
-    brokers.foreach { broker => changeEntityConfig(ConfigType.Broker, broker.toString, configs)
+    validateBrokerConfig(configs)
+    brokers.foreach {
+      broker => changeEntityConfig(ConfigType.Broker, broker.toString, configs)
     }
   }
 
-  private def changeEntityConfig(rootEntityType: String, fullSanitizedEntityName: String, configs: Properties) {
+  /**
+    * Override a broker override or broker default config. These overrides will be persisted between sessions, and will
+    * override any defaults entered in the broker's config files
+    *
+    * @param broker: The broker to apply config changes to or None to update dynamic default configs
+    * @param configs: The config to change, as properties
+    */
+  def changeBrokerConfig(broker: Option[Int], configs: Properties): Unit = {
+    validateBrokerConfig(configs)
+    changeEntityConfig(ConfigType.Broker, broker.map(_.toString).getOrElse(ConfigEntityName.Default), configs)
+  }
+
+  /**
+    * Validate dynamic broker configs. Since broker configs may contain custom configs, the validation
+    * only verifies that the provided config does not contain any static configs.
+    * @param configs configs to validate
+    */
+  def validateBrokerConfig(configs: Properties): Unit = {
+    DynamicConfig.Broker.validate(configs)
+  }
+
+  private def changeEntityConfig(rootEntityType: String, fullSanitizedEntityName: String, configs: Properties): Unit = {
     val sanitizedEntityPath = rootEntityType + '/' + fullSanitizedEntityName
     zkClient.setOrCreateEntityConfigs(rootEntityType, fullSanitizedEntityName, configs)
 
@@ -365,8 +479,8 @@ class AdminZkClient(zkClient: KafkaZkClient) extends Logging {
   }
 
   /**
-   * Read the entity (topic, broker, client, user or ) config (if any) from zk
-   * sanitizedEntityName is , , ,  or /clients/.
+   * Read the entity (topic, broker, client, user,  or ) config (if any) from zk
+   * sanitizedEntityName is , , , , /clients/ or .
    * @param rootEntityType
    * @param sanitizedEntityName
    * @return
@@ -380,7 +494,7 @@ class AdminZkClient(zkClient: KafkaZkClient) extends Logging {
    * @return
    */
   def getAllTopicConfigs(): Map[String, Properties] =
-    zkClient.getAllTopicsInCluster.map(topic => (topic, fetchEntityConfig(ConfigType.Topic, topic))).toMap
+    zkClient.getAllTopicsInCluster().map(topic => (topic, fetchEntityConfig(ConfigType.Topic, topic))).toMap
 
   /**
    * Gets all the entity configs for a given entityType
diff --git a/core/src/main/scala/kafka/zk/KafkaZkClient.scala b/core/src/main/scala/kafka/zk/KafkaZkClient.scala
index 110d676cc0c18..c8122d7f95615 100644
--- a/core/src/main/scala/kafka/zk/KafkaZkClient.scala
+++ b/core/src/main/scala/kafka/zk/KafkaZkClient.scala
@@ -16,45 +16,165 @@
 */
 package kafka.zk
 
-import java.nio.charset.StandardCharsets.UTF_8
 import java.util.Properties
 
+import com.yammer.metrics.core.MetricName
 import kafka.api.LeaderAndIsr
 import kafka.cluster.Broker
-import kafka.controller.{LeaderIsrAndControllerEpoch, ReassignedPartitionsContext}
+import kafka.controller.{KafkaController, LeaderIsrAndControllerEpoch, ReplicaAssignment}
 import kafka.log.LogConfig
-import kafka.security.auth.SimpleAclAuthorizer.VersionedAcls
-import kafka.security.auth.{Acl, Resource, ResourceType}
-
+import kafka.metrics.KafkaMetricsGroup
+import kafka.security.authorizer.AclAuthorizer.{NoAcls, VersionedAcls}
+import kafka.security.authorizer.AclEntry
 import kafka.server.ConfigType
-import kafka.utils._
+import kafka.utils.Logging
+import kafka.zk.TopicZNode.TopicIdReplicaAssignment
 import kafka.zookeeper._
-import org.apache.kafka.common.TopicPartition
-import org.apache.zookeeper.KeeperException.Code
+import org.apache.kafka.common.errors.ControllerMovedException
+import org.apache.kafka.common.resource.{PatternType, ResourcePattern, ResourceType}
+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, Uuid}
+import org.apache.zookeeper.KeeperException.{Code, NodeExistsException}
+import org.apache.zookeeper.OpResult.{CreateResult, ErrorResult, SetDataResult}
+import org.apache.zookeeper.client.ZKClientConfig
 import org.apache.zookeeper.data.{ACL, Stat}
-import org.apache.zookeeper.{CreateMode, KeeperException}
+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(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends Logging {
+class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boolean, time: Time) extends AutoCloseable with
+  Logging with KafkaMetricsGroup {
+
+  override def metricName(name: String, metricTags: scala.collection.Map[String, String]): MetricName = {
+    explicitMetricName("kafka.server", "ZooKeeperClientMetrics", name, metricTags)
+  }
+
+  private val latencyMetric = newHistogram("ZooKeeperRequestLatencyMs")
+
   import KafkaZkClient._
 
-  def createSequentialPersistentPath(path: String, data: String = ""): String = {
-    val createRequest = CreateRequest(path, data.getBytes("UTF-8"), acls(path), CreateMode.PERSISTENT_SEQUENTIAL)
+  // Only for testing
+  private[kafka] def currentZooKeeper: ZooKeeper = zooKeeperClient.currentZooKeeper
+
+  // This variable holds the Zookeeper session id at the moment a Broker gets registered in Zookeeper and the subsequent
+  // updates of the session id. It is possible that the session id changes over the time for 'Session expired'.
+  // This code is part of the work around done in the KAFKA-7165, once ZOOKEEPER-2985 is complete, this code must
+  // be deleted.
+  private var currentZooKeeperSessionId: Long = -1
+
+  /**
+   * Create a sequential persistent path. That is, the znode will not be automatically deleted upon client's disconnect
+   * and a monotonically increasing number will be appended to its name.
+   *
+   * @param path the path to create (with the monotonically increasing number appended)
+   * @param data the znode data
+   * @return the created path (including the appended monotonically increasing number)
+   */
+  private[kafka] def createSequentialPersistentPath(path: String, data: Array[Byte]): String = {
+    val createRequest = CreateRequest(path, data, defaultAcls(path), CreateMode.PERSISTENT_SEQUENTIAL)
     val createResponse = retryRequestUntilConnected(createRequest)
-    createResponse.resultException.foreach(e => throw e)
-    createResponse.path
+    createResponse.maybeThrow()
+    createResponse.name
+  }
+
+  /**
+    * Registers the broker in zookeeper and return the broker epoch.
+    * @param brokerInfo payload of the broker znode
+    * @return broker epoch (znode create transaction id)
+    */
+  def registerBroker(brokerInfo: BrokerInfo): Long = {
+    val path = brokerInfo.path
+    val stat = checkedEphemeralCreate(path, brokerInfo.toJsonBytes)
+    info(s"Registered broker ${brokerInfo.broker.id} at path $path with addresses: " +
+      s"${brokerInfo.broker.endPoints.map(_.connectionString).mkString(",")}, czxid (broker epoch): ${stat.getCzxid}")
+    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.
+   * @return the (updated controller epoch, epoch zkVersion) tuple
+   * @throws ControllerMovedException if fail to create /controller or fail to increment controller epoch.
+   */
+  def registerControllerAndIncrementControllerEpoch(controllerId: Int): (Int, Int) = {
+    val timestamp = time.milliseconds()
+
+    // Read /controller_epoch to get the current controller epoch and zkVersion,
+    // create /controller_epoch with initial value if not exists
+    val (curEpoch, curEpochZkVersion) = getControllerEpoch
+      .map(e => (e._1, e._2.getVersion))
+      .getOrElse(maybeCreateControllerEpochZNode())
+
+    // Create /controller and update /controller_epoch atomically
+    val newControllerEpoch = curEpoch + 1
+    val expectedControllerEpochZkVersion = curEpochZkVersion
+
+    debug(s"Try to create ${ControllerZNode.path} and increment controller epoch to $newControllerEpoch with expected controller epoch zkVersion $expectedControllerEpochZkVersion")
+
+    def checkControllerAndEpoch(): (Int, Int) = {
+      val curControllerId = getControllerId.getOrElse(throw new ControllerMovedException(
+        s"The ephemeral node at ${ControllerZNode.path} went away while checking whether the controller election succeeds. " +
+          s"Aborting controller startup procedure"))
+      if (controllerId == curControllerId) {
+        val (epoch, stat) = getControllerEpoch.getOrElse(
+          throw new IllegalStateException(s"${ControllerEpochZNode.path} existed before but goes away while trying to read it"))
+
+        // If the epoch is the same as newControllerEpoch, it is safe to infer that the returned epoch zkVersion
+        // is associated with the current broker during controller election because we already knew that the zk
+        // transaction succeeds based on the controller znode verification. Other rounds of controller
+        // election will result in larger epoch number written in zk.
+        if (epoch == newControllerEpoch)
+          return (newControllerEpoch, stat.getVersion)
+      }
+      throw new ControllerMovedException("Controller moved to another broker. Aborting controller startup procedure")
+    }
+
+    def tryCreateControllerZNodeAndIncrementEpoch(): (Int, Int) = {
+      val response = retryRequestUntilConnected(
+        MultiRequest(Seq(
+          CreateOp(ControllerZNode.path, ControllerZNode.encode(controllerId, timestamp), defaultAcls(ControllerZNode.path), CreateMode.EPHEMERAL),
+          SetDataOp(ControllerEpochZNode.path, ControllerEpochZNode.encode(newControllerEpoch), expectedControllerEpochZkVersion)))
+      )
+      response.resultCode match {
+        case Code.NODEEXISTS | Code.BADVERSION => checkControllerAndEpoch()
+        case Code.OK =>
+          val setDataResult = response.zkOpResults(1).rawOpResult.asInstanceOf[SetDataResult]
+          (newControllerEpoch, setDataResult.getStat.getVersion)
+        case code => throw KeeperException.create(code)
+      }
+    }
+
+    tryCreateControllerZNodeAndIncrementEpoch()
+  }
+
+  private def maybeCreateControllerEpochZNode(): (Int, Int) = {
+    createControllerEpochRaw(KafkaController.InitialControllerEpoch).resultCode match {
+      case Code.OK =>
+        info(s"Successfully created ${ControllerEpochZNode.path} with initial epoch ${KafkaController.InitialControllerEpoch}")
+        (KafkaController.InitialControllerEpoch, KafkaController.InitialControllerEpochZkVersion)
+      case Code.NODEEXISTS =>
+        val (epoch, stat) = getControllerEpoch.getOrElse(throw new IllegalStateException(s"${ControllerEpochZNode.path} existed before but goes away while trying to read it"))
+        (epoch, stat.getVersion)
+      case code =>
+        throw KeeperException.create(code)
+    }
+  }
+
+  def updateBrokerInfo(brokerInfo: BrokerInfo): Unit = {
+    val brokerIdPath = brokerInfo.path
+    val setDataRequest = SetDataRequest(brokerIdPath, brokerInfo.toJsonBytes, ZkVersion.MatchAnyVersion)
+    val response = retryRequestUntilConnected(setDataRequest)
+    response.maybeThrow()
+    info("Updated broker %d at path %s with addresses: %s".format(brokerInfo.broker.id, brokerIdPath, brokerInfo.broker.endPoints))
   }
 
   /**
@@ -72,31 +192,33 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
   /**
    * Sets topic partition states for the given partitions.
    * @param leaderIsrAndControllerEpochs the partition states of each partition whose state we wish to set.
+   * @param expectedControllerEpochZkVersion expected controller epoch zkVersion.
    * @return sequence of SetDataResponse whose contexts are the partitions they are associated with.
    */
-  def setTopicPartitionStatesRaw(leaderIsrAndControllerEpochs: Map[TopicPartition, LeaderIsrAndControllerEpoch]): Seq[SetDataResponse] = {
+  def setTopicPartitionStatesRaw(leaderIsrAndControllerEpochs: Map[TopicPartition, LeaderIsrAndControllerEpoch], expectedControllerEpochZkVersion: Int): Seq[SetDataResponse] = {
     val setDataRequests = leaderIsrAndControllerEpochs.map { case (partition, leaderIsrAndControllerEpoch) =>
       val path = TopicPartitionStateZNode.path(partition)
       val data = TopicPartitionStateZNode.encode(leaderIsrAndControllerEpoch)
       SetDataRequest(path, data, leaderIsrAndControllerEpoch.leaderAndIsr.zkVersion, Some(partition))
     }
-    retryRequestsUntilConnected(setDataRequests.toSeq)
+    retryRequestsUntilConnected(setDataRequests.toSeq, expectedControllerEpochZkVersion)
   }
 
   /**
    * Creates topic partition state znodes for the given partitions.
    * @param leaderIsrAndControllerEpochs the partition states of each partition whose state we wish to set.
+   * @param expectedControllerEpochZkVersion expected controller epoch zkVersion.
    * @return sequence of CreateResponse whose contexts are the partitions they are associated with.
    */
-  def createTopicPartitionStatesRaw(leaderIsrAndControllerEpochs: Map[TopicPartition, LeaderIsrAndControllerEpoch]): Seq[CreateResponse] = {
-    createTopicPartitions(leaderIsrAndControllerEpochs.keys.map(_.topic).toSet.toSeq)
-    createTopicPartition(leaderIsrAndControllerEpochs.keys.toSeq)
+  def createTopicPartitionStatesRaw(leaderIsrAndControllerEpochs: Map[TopicPartition, LeaderIsrAndControllerEpoch], expectedControllerEpochZkVersion: Int): Seq[CreateResponse] = {
+    createTopicPartitions(leaderIsrAndControllerEpochs.keys.map(_.topic).toSet.toSeq, expectedControllerEpochZkVersion)
+    createTopicPartition(leaderIsrAndControllerEpochs.keys.toSeq, expectedControllerEpochZkVersion)
     val createRequests = leaderIsrAndControllerEpochs.map { case (partition, leaderIsrAndControllerEpoch) =>
       val path = TopicPartitionStateZNode.path(partition)
       val data = TopicPartitionStateZNode.encode(leaderIsrAndControllerEpoch)
-      CreateRequest(path, data, acls(path), CreateMode.PERSISTENT, Some(partition))
+      CreateRequest(path, data, defaultAcls(path), CreateMode.PERSISTENT, Some(partition))
     }
-    retryRequestsUntilConnected(createRequests.toSeq)
+    retryRequestsUntilConnected(createRequests.toSeq, expectedControllerEpochZkVersion)
   }
 
   /**
@@ -117,39 +239,50 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
    */
   def createControllerEpochRaw(epoch: Int): CreateResponse = {
     val createRequest = CreateRequest(ControllerEpochZNode.path, ControllerEpochZNode.encode(epoch),
-      acls(ControllerEpochZNode.path), CreateMode.PERSISTENT)
+      defaultAcls(ControllerEpochZNode.path), CreateMode.PERSISTENT)
     retryRequestUntilConnected(createRequest)
   }
 
   /**
-   * Try to update the partition states of multiple partitions in zookeeper.
+   * Update the partition states of multiple partitions in zookeeper.
    * @param leaderAndIsrs The partition states to update.
    * @param controllerEpoch The current controller epoch.
+   * @param expectedControllerEpochZkVersion expected controller epoch zkVersion.
    * @return UpdateLeaderAndIsrResult instance containing per partition results.
    */
-  def updateLeaderAndIsr(leaderAndIsrs: Map[TopicPartition, LeaderAndIsr], controllerEpoch: Int): UpdateLeaderAndIsrResult = {
-    val successfulUpdates = mutable.Map.empty[TopicPartition, LeaderAndIsr]
-    val updatesToRetry = mutable.Buffer.empty[TopicPartition]
-    val failed = mutable.Map.empty[TopicPartition, Exception]
-    val leaderIsrAndControllerEpochs = leaderAndIsrs.map { case (partition, leaderAndIsr) => partition -> LeaderIsrAndControllerEpoch(leaderAndIsr, controllerEpoch) }
+  def updateLeaderAndIsr(
+    leaderAndIsrs: Map[TopicPartition, LeaderAndIsr],
+    controllerEpoch: Int,
+    expectedControllerEpochZkVersion: Int
+  ): UpdateLeaderAndIsrResult = {
+    val leaderIsrAndControllerEpochs = leaderAndIsrs.map { case (partition, leaderAndIsr) =>
+      partition -> LeaderIsrAndControllerEpoch(leaderAndIsr, controllerEpoch)
+    }
     val setDataResponses = try {
-      setTopicPartitionStatesRaw(leaderIsrAndControllerEpochs)
+      setTopicPartitionStatesRaw(leaderIsrAndControllerEpochs, expectedControllerEpochZkVersion)
     } 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)
   }
 
   /**
@@ -160,8 +293,10 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
    *         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 {
@@ -208,6 +343,12 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
   /**
    * Sets or creates the entity znode path with the given configs depending
    * on whether it already exists or not.
+   *
+   * If this is method is called concurrently, the last writer wins. In cases where we update configs and then
+   * partition assignment (i.e. create topic), it's possible for one thread to set this and the other to set the
+   * partition assignment. As such, the recommendation is to never call create topic for the same topic with different
+   * configs/partition assignment concurrently.
+   *
    * @param rootEntityType entity type
    * @param sanitizedEntityName entity name
    * @throws KeeperException if there is an error while setting or creating the znode
@@ -215,21 +356,25 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
   def setOrCreateEntityConfigs(rootEntityType: String, sanitizedEntityName: String, config: Properties) = {
 
     def set(configData: Array[Byte]): SetDataResponse = {
-      val setDataRequest = SetDataRequest(ConfigEntityZNode.path(rootEntityType, sanitizedEntityName), ConfigEntityZNode.encode(config), ZkVersion.NoVersion)
+      val setDataRequest = SetDataRequest(ConfigEntityZNode.path(rootEntityType, sanitizedEntityName),
+        configData, ZkVersion.MatchAnyVersion)
       retryRequestUntilConnected(setDataRequest)
     }
 
-    def create(configData: Array[Byte]) = {
+    def createOrSet(configData: Array[Byte]): Unit = {
       val path = ConfigEntityZNode.path(rootEntityType, sanitizedEntityName)
-      createRecursive(path, ConfigEntityZNode.encode(config))
+      try createRecursive(path, configData)
+      catch {
+        case _: NodeExistsException => set(configData).maybeThrow()
+      }
     }
 
     val configData = ConfigEntityZNode.encode(config)
 
     val setDataResponse = set(configData)
     setDataResponse.resultCode match {
-      case Code.NONODE => create(configData)
-      case _ => setDataResponse.resultException.foreach(e => throw e)
+      case Code.NONODE => createOrSet(configData)
+      case _ => setDataResponse.maybeThrow()
     }
   }
 
@@ -248,12 +393,11 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
    * @throws KeeperException if there is an error while setting or creating the znode
    */
   def createConfigChangeNotification(sanitizedEntityPath: String): Unit = {
+    makeSurePersistentPathExists(ConfigEntityChangeNotificationZNode.path)
     val path = ConfigEntityChangeNotificationSequenceZNode.createPath
-    val createRequest = CreateRequest(path, ConfigEntityChangeNotificationSequenceZNode.encode(sanitizedEntityPath), acls(path), CreateMode.PERSISTENT_SEQUENTIAL)
+    val createRequest = CreateRequest(path, ConfigEntityChangeNotificationSequenceZNode.encode(sanitizedEntityPath), defaultAcls(path), CreateMode.PERSISTENT_SEQUENTIAL)
     val createResponse = retryRequestUntilConnected(createRequest)
-    if (createResponse.resultCode != Code.OK) {
-      createResponse.resultException.foreach(e => throw e)
-    }
+    createResponse.maybeThrow()
   }
 
   /**
@@ -268,32 +412,65 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
       val brokerId = getDataResponse.ctx.get.asInstanceOf[Int]
       getDataResponse.resultCode match {
         case Code.OK =>
-          Option(BrokerIdZNode.decode(brokerId, getDataResponse.data))
+          Option(BrokerIdZNode.decode(brokerId, getDataResponse.data).broker)
         case Code.NONODE => None
         case _ => throw getDataResponse.resultException.get
       }
     }
   }
 
+  /**
+    * Gets all brokers with broker epoch in the cluster.
+    * @return map of broker to epoch in the cluster.
+    */
+  def getAllBrokerAndEpochsInCluster: Map[Broker, Long] = {
+    val brokerIds = getSortedBrokerList
+    val getDataRequests = brokerIds.map(brokerId => GetDataRequest(BrokerIdZNode.path(brokerId), ctx = Some(brokerId)))
+    val getDataResponses = retryRequestsUntilConnected(getDataRequests)
+    getDataResponses.flatMap { getDataResponse =>
+      val brokerId = getDataResponse.ctx.get.asInstanceOf[Int]
+      getDataResponse.resultCode match {
+        case Code.OK =>
+          Some((BrokerIdZNode.decode(brokerId, getDataResponse.data).broker, getDataResponse.stat.getCzxid))
+        case Code.NONODE => None
+        case _ => throw getDataResponse.resultException.get
+      }
+    }.toMap
+  }
+
+  /**
+    * Get a broker from ZK
+    * @return an optional Broker
+    */
+  def getBroker(brokerId: Int): Option[Broker] = {
+    val getDataRequest = GetDataRequest(BrokerIdZNode.path(brokerId))
+    val getDataResponse = retryRequestUntilConnected(getDataRequest)
+    getDataResponse.resultCode match {
+      case Code.OK =>
+        Option(BrokerIdZNode.decode(brokerId, getDataResponse.data).broker)
+      case Code.NONODE => None
+      case _ => throw getDataResponse.resultException.get
+    }
+  }
 
   /**
    * Gets the list of sorted broker Ids
    */
-  def getSortedBrokerList(): Seq[Int] =
-    getChildren(BrokerIdsZNode.path).map(_.toInt).sorted
+  def getSortedBrokerList: Seq[Int] = getChildren(BrokerIdsZNode.path).map(_.toInt).sorted
 
   /**
    * Gets all topics in the cluster.
+   * @param registerWatch indicates if a watch must be registered or not
    * @return sequence of topics in the cluster.
    */
-  def getAllTopicsInCluster: Seq[String] = {
-    val getChildrenResponse = retryRequestUntilConnected(GetChildrenRequest(TopicsZNode.path))
+  def getAllTopicsInCluster(registerWatch: Boolean = false): Set[String] = {
+    val getChildrenResponse = retryRequestUntilConnected(
+      GetChildrenRequest(TopicsZNode.path, registerWatch))
     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
     }
-
   }
 
   /**
@@ -305,38 +482,70 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
     pathExists(TopicZNode.path(topicName))
   }
 
+  /**
+   * Adds a topic ID to existing topic and replica assignments
+   * @param topicIdReplicaAssignments the TopicIDReplicaAssignments to add a topic ID to
+   * @return the updated TopicIdReplicaAssigments including the newly created topic IDs
+   */
+  def setTopicIds(topicIdReplicaAssignments: collection.Set[TopicIdReplicaAssignment],
+                  expectedControllerEpochZkVersion: Int): Set[TopicIdReplicaAssignment] = {
+    val updatedAssignments = topicIdReplicaAssignments.map {
+      case TopicIdReplicaAssignment(topic, None, assignments) =>
+        TopicIdReplicaAssignment(topic, Some(Uuid.randomUuid()), assignments)
+      case TopicIdReplicaAssignment(topic, Some(_), _) =>
+        throw new IllegalArgumentException("TopicIdReplicaAssignment for " + topic + " already contains a topic ID.")
+    }.toSet
+
+    val setDataRequests = updatedAssignments.map { case TopicIdReplicaAssignment(topic, topicIdOpt, assignments) =>
+      SetDataRequest(TopicZNode.path(topic), TopicZNode.encode(topicIdOpt.get, assignments), ZkVersion.MatchAnyVersion)
+    }.toSeq
+
+    retryRequestsUntilConnected(setDataRequests, expectedControllerEpochZkVersion)
+    updatedAssignments
+  }
+
   /**
    * Sets the topic znode with the given assignment.
    * @param topic the topic whose assignment is being set.
+   * @param topicId optional topic ID if the topic has one
    * @param assignment the partition to replica mapping to set for the given topic
+   * @param expectedControllerEpochZkVersion expected controller epoch zkVersion.
    * @return SetDataResponse
    */
-  def setTopicAssignmentRaw(topic: String, assignment: collection.Map[TopicPartition, Seq[Int]]): SetDataResponse = {
-    val setDataRequest = SetDataRequest(TopicZNode.path(topic), TopicZNode.encode(assignment), ZkVersion.NoVersion)
-    retryRequestUntilConnected(setDataRequest)
+  def setTopicAssignmentRaw(topic: String,
+                            topicId: Uuid,
+                            assignment: collection.Map[TopicPartition, ReplicaAssignment],
+                            expectedControllerEpochZkVersion: Int): SetDataResponse = {
+    val setDataRequest = SetDataRequest(TopicZNode.path(topic), TopicZNode.encode(topicId, assignment), ZkVersion.MatchAnyVersion)
+    retryRequestUntilConnected(setDataRequest, expectedControllerEpochZkVersion)
   }
 
   /**
    * Sets the topic znode with the given assignment.
    * @param topic the topic whose assignment is being set.
+   * @param topicId unique topic ID for the topic
    * @param assignment the partition to replica mapping to set for the given topic
+   * @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]]) = {
-    val setDataResponse = setTopicAssignmentRaw(topic, assignment)
-    if (setDataResponse.resultCode != Code.OK) {
-      setDataResponse.resultException.foreach(e => throw e)
-    }
+  def setTopicAssignment(topic: String,
+                         topicId: Uuid,
+                         assignment: Map[TopicPartition, ReplicaAssignment],
+                         expectedControllerEpochZkVersion: Int = ZkVersion.MatchAnyVersion) = {
+    val setDataResponse = setTopicAssignmentRaw(topic, topicId, assignment, expectedControllerEpochZkVersion)
+    setDataResponse.maybeThrow()
   }
 
   /**
    * Create the topic znode with the given assignment.
    * @param topic the topic whose assignment is being set.
+   * @param topicId unique topic ID for the topic
    * @param assignment the partition to replica mapping to set for the given topic
    * @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))
+  def createTopicAssignment(topic: String, topicId: Uuid, assignment: Map[TopicPartition, Seq[Int]]): Unit = {
+    val persistedAssignments = assignment.map { case (k, v) => k -> ReplicaAssignment(v) }
+    createRecursive(TopicZNode.path(topic), TopicZNode.encode(topicId, persistedAssignments))
   }
 
   /**
@@ -344,7 +553,7 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
    * @return sequence of znode names and not the absolute znode path.
    */
   def getAllLogDirEventNotifications: Seq[String] = {
-    val getChildrenResponse = retryRequestUntilConnected(GetChildrenRequest(LogDirEventNotificationZNode.path))
+    val getChildrenResponse = retryRequestUntilConnected(GetChildrenRequest(LogDirEventNotificationZNode.path, registerWatch = true))
     getChildrenResponse.resultCode match {
       case Code.OK => getChildrenResponse.children.map(LogDirEventNotificationSequenceZNode.sequenceNumber)
       case Code.NONODE => Seq.empty
@@ -373,40 +582,93 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
 
   /**
    * Deletes all log dir event notifications.
+   * @param expectedControllerEpochZkVersion expected controller epoch zkVersion.
    */
-  def deleteLogDirEventNotifications(): Unit = {
-    val getChildrenResponse = retryRequestUntilConnected(GetChildrenRequest(LogDirEventNotificationZNode.path))
+  def deleteLogDirEventNotifications(expectedControllerEpochZkVersion: Int): Unit = {
+    val getChildrenResponse = retryRequestUntilConnected(GetChildrenRequest(LogDirEventNotificationZNode.path, registerWatch = true))
     if (getChildrenResponse.resultCode == Code.OK) {
-      deleteLogDirEventNotifications(getChildrenResponse.children)
+      deleteLogDirEventNotifications(getChildrenResponse.children.map(LogDirEventNotificationSequenceZNode.sequenceNumber), expectedControllerEpochZkVersion)
     } else if (getChildrenResponse.resultCode != Code.NONODE) {
-      getChildrenResponse.resultException.foreach(e => throw e)
+      getChildrenResponse.maybeThrow()
     }
   }
 
   /**
    * Deletes the log dir event notifications associated with the given sequence numbers.
    * @param sequenceNumbers the sequence numbers associated with the log dir event notifications to be deleted.
+   * @param expectedControllerEpochZkVersion expected controller epoch zkVersion.
    */
-  def deleteLogDirEventNotifications(sequenceNumbers: Seq[String]): Unit = {
+  def deleteLogDirEventNotifications(sequenceNumbers: Seq[String], expectedControllerEpochZkVersion: Int): Unit = {
     val deleteRequests = sequenceNumbers.map { sequenceNumber =>
-      DeleteRequest(LogDirEventNotificationSequenceZNode.path(sequenceNumber), ZkVersion.NoVersion)
+      DeleteRequest(LogDirEventNotificationSequenceZNode.path(sequenceNumber), ZkVersion.MatchAnyVersion)
     }
-    retryRequestsUntilConnected(deleteRequests)
+    retryRequestsUntilConnected(deleteRequests, expectedControllerEpochZkVersion)
   }
 
   /**
-   * Gets the assignments for the given topics.
+   * Gets the topic IDs for the given topics.
+   * @param topics the topics we wish to retrieve the Topic IDs for
+   * @return the Topic IDs
+   * @throws IllegalStateException if a topic in topics does not have a Topic ID
+   */
+  def getTopicIdsForTopics(topics: Set[String]): Map[String, Uuid] = {
+    val getDataRequests = topics.map(topic => GetDataRequest(TopicZNode.path(topic), ctx = Some(topic)))
+    val getDataResponses = retryRequestsUntilConnected(getDataRequests.toSeq)
+    getDataResponses.map { getDataResponse =>
+      val topic = getDataResponse.ctx.get.asInstanceOf[String]
+      getDataResponse.resultCode match {
+        case Code.OK => Some(TopicZNode.decode(topic, getDataResponse.data))
+        case Code.NONODE => None
+        case _ => throw getDataResponse.resultException.get
+      }
+    }.map(_.get)
+      .map(topicIdAssignment => (topicIdAssignment.topic,
+        topicIdAssignment.topicId.getOrElse(
+          throw new IllegalStateException("Topic " + topicIdAssignment.topic + " does not have a topic ID."))))
+      .toMap
+  }
+
+  /**
+   * 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).map { case (k, v) => k -> v.replicas }
+  }
+
+  /**
+   * Gets the TopicID and replica assignments for the given topics.
+   * @param topics the topics whose partitions we wish to get the assignments for.
+   * @return the TopicIdReplicaAssignment for each partition for the given topics.
+   */
+  def getReplicaAssignmentAndTopicIdForTopics(topics: Set[String]): Set[TopicIdReplicaAssignment] = {
     val getDataRequests = topics.map(topic => GetDataRequest(TopicZNode.path(topic), ctx = Some(topic)))
     val getDataResponses = retryRequestsUntilConnected(getDataRequests.toSeq)
-    getDataResponses.flatMap { getDataResponse =>
+    getDataResponses.map { 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 => TopicIdReplicaAssignment(topic, None, Map.empty[TopicPartition, ReplicaAssignment])
+        case _ => throw getDataResponse.resultException.get
+      }
+    }.toSet
+  }
+
+  /**
+    * 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).assignment
+        case Code.NONODE => Map.empty[TopicPartition, ReplicaAssignment]
         case _ => throw getDataResponse.resultException.get
       }
     }.toMap
@@ -417,16 +679,16 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
    * @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 =>
       val topic = getDataResponse.ctx.get.asInstanceOf[String]
-       if (getDataResponse.resultCode == Code.OK) {
-        val partitionMap = TopicZNode.decode(topic, getDataResponse.data).map { case (k, v) => (k.partition, v) }
+      if (getDataResponse.resultCode == Code.OK) {
+        val partitionMap = TopicZNode.decode(topic, getDataResponse.data).assignment.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
       }
@@ -469,17 +731,32 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
     topicData.getOrElse(topicPartition, Seq.empty)
   }
 
+  /**
+   * Gets all partitions in the cluster
+   * @return all partitions in the cluster
+   */
+  def getAllPartitions: Set[TopicPartition] = {
+    val topics = getChildren(TopicsZNode.path)
+    if (topics == null) Set.empty
+    else {
+      topics.flatMap { topic =>
+        // The partitions path may not exist if the topic is in the process of being deleted
+        getChildren(TopicPartitionsZNode.path(topic)).map(_.toInt).map(new TopicPartition(topic, _))
+      }.toSet
+    }
+  }
+
   /**
    * Gets the data and version at the given zk path
    * @param path zk node path
-   * @return A tuple of 2 elements, where first element is zk node data as string
+   * @return A tuple of 2 elements, where first element is zk node data as an array of bytes
    *         and second element is zk node version.
-   *         returns (None, ZkVersion.NoVersion) if node doesn't exists and throws exception for any error
+   *         returns (None, ZkVersion.UnknownVersion) if node doesn't exist and throws exception for any error
    */
-  def getDataAndVersion(path: String): (Option[String], Int) = {
+  def getDataAndVersion(path: String): (Option[Array[Byte]], Int) = {
     val (data, stat) = getDataAndStat(path)
     stat match {
-      case ZkStat.NoStat => (data, ZkVersion.NoVersion)
+      case ZkStat.NoStat => (data, ZkVersion.UnknownVersion)
       case _ => (data, stat.getVersion)
     }
   }
@@ -487,22 +764,16 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
   /**
    * Gets the data and Stat at the given zk path
    * @param path zk node path
-   * @return A tuple of 2 elements, where first element is zk node data as string
+   * @return A tuple of 2 elements, where first element is zk node data as an array of bytes
    *         and second element is zk node stats.
    *         returns (None, ZkStat.NoStat) if node doesn't exists and throws exception for any error
    */
-  def getDataAndStat(path: String): (Option[String], Stat) = {
+  def getDataAndStat(path: String): (Option[Array[Byte]], Stat) = {
     val getDataRequest = GetDataRequest(path)
     val getDataResponse = retryRequestUntilConnected(getDataRequest)
 
     getDataResponse.resultCode match {
-      case Code.OK =>
-        if (getDataResponse.data == null)
-          (None, getDataResponse.stat)
-        else {
-          val data = Option(getDataResponse.data).map(new String(_, UTF_8))
-          (data, getDataResponse.stat)
-        }
+      case Code.OK => (Option(getDataResponse.data), getDataResponse.stat)
       case Code.NONODE => (None, ZkStat.NoStat)
       case _ => throw getDataResponse.resultException.get
     }
@@ -514,7 +785,7 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
    * @return list of child node names
    */
   def getChildren(path : String): Seq[String] = {
-    val getChildrenResponse = retryRequestUntilConnected(GetChildrenRequest(path))
+    val getChildrenResponse = retryRequestUntilConnected(GetChildrenRequest(path, registerWatch = true))
     getChildrenResponse.resultCode match {
       case Code.OK => getChildrenResponse.children
       case Code.NONODE => Seq.empty
@@ -524,22 +795,22 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
 
   /**
    * 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, ZkVersion.NoVersion)
+   * exist, the current version is not the expected version, etc.) return (false, ZkVersion.UnknownVersion)
    *
    * When there is a ConnectionLossException during the conditional update, ZookeeperClient 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 conditionalUpdatePath(path: String, data: String, expectVersion: Int,
-                            optionalChecker:Option[(KafkaZkClient, String, String) => (Boolean,Int)] = None): (Boolean, Int) = {
+  def conditionalUpdatePath(path: String, data: Array[Byte], expectVersion: Int,
+                            optionalChecker: Option[(KafkaZkClient, String, Array[Byte]) => (Boolean,Int)] = None): (Boolean, Int) = {
 
-    val setDataRequest = SetDataRequest(path, data.getBytes(UTF_8), expectVersion)
+    val setDataRequest = SetDataRequest(path, data, expectVersion)
     val setDataResponse = retryRequestUntilConnected(setDataRequest)
 
     setDataResponse.resultCode match {
       case Code.OK =>
         debug("Conditional update of path %s with value %s and expected version %d succeeded, returning the new version: %d"
-          .format(path, data, expectVersion, setDataResponse.stat.getVersion))
+          .format(path, Utils.utf8(data), expectVersion, setDataResponse.stat.getVersion))
         (true, setDataResponse.stat.getVersion)
 
       case Code.BADVERSION =>
@@ -548,18 +819,18 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
           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, setDataResponse.resultException.get.getMessage))
-            (false, ZkVersion.NoVersion)
+              .format(path, Utils.utf8(data), expectVersion, setDataResponse.resultException.get.getMessage))
+            (false, ZkVersion.UnknownVersion)
         }
 
       case Code.NONODE =>
-        debug("Conditional update of path %s with data %s and expected version %d failed due to %s".format(path, data,
-          expectVersion, setDataResponse.resultException.get.getMessage))
-        (false, ZkVersion.NoVersion)
+        debug("Conditional update of path %s with data %s and expected version %d failed due to %s".format(path,
+          Utils.utf8(data), expectVersion, setDataResponse.resultException.get.getMessage))
+        (false, ZkVersion.UnknownVersion)
 
       case _ =>
-        debug("Conditional update of path %s with data %s and expected version %d failed due to %s".format(path, data,
-          expectVersion, setDataResponse.resultException.get.getMessage))
+        debug("Conditional update of path %s with data %s and expected version %d failed due to %s".format(path,
+          Utils.utf8(data), expectVersion, setDataResponse.resultException.get.getMessage))
         throw setDataResponse.resultException.get
     }
   }
@@ -573,7 +844,6 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
     createRecursive(DeleteTopicsTopicZNode.path(topicName))
   }
 
-
   /**
    * Checks if topic is marked for deletion
    * @param topic
@@ -588,7 +858,7 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
    * @return sequence of topics marked for deletion.
    */
   def getTopicDeletions: Seq[String] = {
-    val getChildrenResponse = retryRequestUntilConnected(GetChildrenRequest(DeleteTopicsZNode.path))
+    val getChildrenResponse = retryRequestUntilConnected(GetChildrenRequest(DeleteTopicsZNode.path, registerWatch = true))
     getChildrenResponse.resultCode match {
       case Code.OK => getChildrenResponse.children
       case Code.NONODE => Seq.empty
@@ -599,21 +869,28 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
   /**
    * Remove the given topics from the topics marked for deletion.
    * @param topics the topics to remove.
+   * @param expectedControllerEpochZkVersion expected controller epoch zkVersion.
    */
-  def deleteTopicDeletions(topics: Seq[String]): Unit = {
-    val deleteRequests = topics.map(topic => DeleteRequest(DeleteTopicsTopicZNode.path(topic), ZkVersion.NoVersion))
-    retryRequestsUntilConnected(deleteRequests)
+  def deleteTopicDeletions(topics: Seq[String], expectedControllerEpochZkVersion: Int): Unit = {
+    val deleteRequests = topics.map(topic => DeleteRequest(DeleteTopicsTopicZNode.path(topic), ZkVersion.MatchAnyVersion))
+    retryRequestsUntilConnected(deleteRequests, expectedControllerEpochZkVersion)
   }
 
   /**
    * Returns all reassignments.
    * @return the reassignments for each partition.
    */
-  def getPartitionReassignment: Map[TopicPartition, Seq[Int]] = {
+  def getPartitionReassignment: collection.Map[TopicPartition, Seq[Int]] = {
     val getDataRequest = GetDataRequest(ReassignPartitionsZNode.path)
     val getDataResponse = retryRequestUntilConnected(getDataRequest)
     getDataResponse.resultCode match {
-      case  Code.OK => ReassignPartitionsZNode.decode(getDataResponse.data)
+      case Code.OK =>
+        ReassignPartitionsZNode.decode(getDataResponse.data) match {
+          case Left(e) =>
+            logger.warn(s"Ignoring partition reassignment due to invalid json: ${e.getMessage}", e)
+            Map.empty[TopicPartition, Seq[Int]]
+          case Right(assignments) => assignments
+        }
       case Code.NONODE => Map.empty
       case _ => throw getDataResponse.resultException.get
     }
@@ -624,19 +901,22 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
    * exists or not.
    *
    * @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
    */
-  def setOrCreatePartitionReassignment(reassignment: collection.Map[TopicPartition, Seq[Int]]): Unit = {
+  @Deprecated
+  def setOrCreatePartitionReassignment(reassignment: collection.Map[TopicPartition, Seq[Int]], expectedControllerEpochZkVersion: Int): Unit = {
 
     def set(reassignmentData: Array[Byte]): SetDataResponse = {
-      val setDataRequest = SetDataRequest(ReassignPartitionsZNode.path, reassignmentData, ZkVersion.NoVersion)
-      retryRequestUntilConnected(setDataRequest)
+      val setDataRequest = SetDataRequest(ReassignPartitionsZNode.path, reassignmentData, ZkVersion.MatchAnyVersion)
+      retryRequestUntilConnected(setDataRequest, expectedControllerEpochZkVersion)
     }
 
     def create(reassignmentData: Array[Byte]): CreateResponse = {
-      val createRequest = CreateRequest(ReassignPartitionsZNode.path, reassignmentData, acls(ReassignPartitionsZNode.path),
+      val createRequest = CreateRequest(ReassignPartitionsZNode.path, reassignmentData, defaultAcls(ReassignPartitionsZNode.path),
         CreateMode.PERSISTENT)
-      retryRequestUntilConnected(createRequest)
+      retryRequestUntilConnected(createRequest, expectedControllerEpochZkVersion)
     }
 
     val reassignmentData = ReassignPartitionsZNode.encode(reassignment)
@@ -644,50 +924,36 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
     setDataResponse.resultCode match {
       case Code.NONODE =>
         val createDataResponse = create(reassignmentData)
-        createDataResponse.resultException.foreach(e => throw e)
-      case _ => setDataResponse.resultException.foreach(e => throw e)
+        createDataResponse.maybeThrow()
+      case _ => setDataResponse.maybeThrow()
     }
   }
 
   /**
    * 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 setting or creating the znode
+   * @throws KeeperException if there is an error while creating the znode.
    */
   def createPartitionReassignment(reassignment: Map[TopicPartition, Seq[Int]])  = {
-    val createRequest = CreateRequest(ReassignPartitionsZNode.path, ReassignPartitionsZNode.encode(reassignment),
-      acls(ReassignPartitionsZNode.path), CreateMode.PERSISTENT)
-    val createResponse = retryRequestUntilConnected(createRequest)
-
-    if (createResponse.resultCode != Code.OK) {
-      throw createResponse.resultException.get
-    }
+    createRecursive(ReassignPartitionsZNode.path, ReassignPartitionsZNode.encode(reassignment))
   }
 
   /**
    * Deletes the partition reassignment znode.
+   * @param expectedControllerEpochZkVersion expected controller epoch zkVersion.
    */
-  def deletePartitionReassignment(): Unit = {
-    val deleteRequest = DeleteRequest(ReassignPartitionsZNode.path, ZkVersion.NoVersion)
-    retryRequestUntilConnected(deleteRequest)
+  def deletePartitionReassignment(expectedControllerEpochZkVersion: Int): Unit = {
+    deletePath(ReassignPartitionsZNode.path, expectedControllerEpochZkVersion)
   }
 
   /**
-   * Checks if reassign partitions is in progress
-   * @return true if reassign partitions is in progress, else false
+   * Checks if reassign partitions is in progress.
+   * @return true if reassign partitions is in progress, else false.
    */
-  def reassignPartitionsInProgress(): Boolean = {
+  def reassignPartitionsInProgress: Boolean = {
     pathExists(ReassignPartitionsZNode.path)
   }
 
-  /**
-   * Gets the partitions being reassigned for given topics
-   * @return ReassignedPartitionsContexts for each topic which are being reassigned.
-   */
-  def getPartitionsBeingReassigned(): Map[TopicPartition, ReassignedPartitionsContext] = {
-    getPartitionReassignment.mapValues(replicas => ReassignedPartitionsContext(replicas))
-  }
-
   /**
    * Gets topic partition states for the given partitions.
    * @param partitions the partitions for which we want to get states.
@@ -723,15 +989,28 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
 
   /**
    * Gets the leader for a given partition
-   * @param partition
+   * @param partition The partition for which we want to get leader.
    * @return optional integer if the leader exists and None otherwise.
    */
-  def getLeaderForPartition(partition: TopicPartition): Option[Int] = {
-    val leaderIsrEpoch = getTopicPartitionState(partition)
-    if (leaderIsrEpoch.isDefined)
-      Option(leaderIsrEpoch.get.leaderAndIsr.leader)
-    else
-      None
+  def getLeaderForPartition(partition: TopicPartition): Option[Int] =
+    getTopicPartitionState(partition).map(_.leaderAndIsr.leader)
+
+  /**
+   * Gets the in-sync replicas (ISR) for a specific topicPartition
+   * @param partition The partition for which we want to get ISR.
+   * @return optional ISR if exists and None otherwise
+   */
+  def getInSyncReplicasForPartition(partition: TopicPartition): Option[Seq[Int]] =
+    getTopicPartitionState(partition).map(_.leaderAndIsr.isr)
+
+
+  /**
+   * Gets the leader epoch for a specific topicPartition
+   * @param partition The partition for which we want to get the leader epoch
+   * @return optional integer if the leader exists and None otherwise
+   */
+  def getEpochForPartition(partition: TopicPartition): Option[Int] = {
+    getTopicPartitionState(partition).map(_.leaderAndIsr.leaderEpoch)
   }
 
   /**
@@ -739,7 +1018,7 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
    * @return sequence of znode names and not the absolute znode path.
    */
   def getAllIsrChangeNotifications: Seq[String] = {
-    val getChildrenResponse = retryRequestUntilConnected(GetChildrenRequest(IsrChangeNotificationZNode.path))
+    val getChildrenResponse = retryRequestUntilConnected(GetChildrenRequest(IsrChangeNotificationZNode.path, registerWatch = true))
     getChildrenResponse.resultCode match {
       case Code.OK => getChildrenResponse.children.map(IsrChangeNotificationSequenceZNode.sequenceNumber)
       case Code.NONODE => Seq.empty
@@ -768,25 +1047,36 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
 
   /**
    * Deletes all isr change notifications.
+   * @param expectedControllerEpochZkVersion expected controller epoch zkVersion.
    */
-  def deleteIsrChangeNotifications(): Unit = {
-    val getChildrenResponse = retryRequestUntilConnected(GetChildrenRequest(IsrChangeNotificationZNode.path))
+  def deleteIsrChangeNotifications(expectedControllerEpochZkVersion: Int): Unit = {
+    val getChildrenResponse = retryRequestUntilConnected(GetChildrenRequest(IsrChangeNotificationZNode.path, registerWatch = true))
     if (getChildrenResponse.resultCode == Code.OK) {
-      deleteIsrChangeNotifications(getChildrenResponse.children.map(IsrChangeNotificationSequenceZNode.sequenceNumber))
+      deleteIsrChangeNotifications(getChildrenResponse.children.map(IsrChangeNotificationSequenceZNode.sequenceNumber), expectedControllerEpochZkVersion)
     } else if (getChildrenResponse.resultCode != Code.NONODE) {
-      getChildrenResponse.resultException.foreach(e => throw e)
+      getChildrenResponse.maybeThrow()
     }
   }
 
   /**
    * Deletes the isr change notifications associated with the given sequence numbers.
    * @param sequenceNumbers the sequence numbers associated with the isr change notifications to be deleted.
+   * @param expectedControllerEpochZkVersion expected controller epoch zkVersion.
    */
-  def deleteIsrChangeNotifications(sequenceNumbers: Seq[String]): Unit = {
+  def deleteIsrChangeNotifications(sequenceNumbers: Seq[String], expectedControllerEpochZkVersion: Int): Unit = {
     val deleteRequests = sequenceNumbers.map { sequenceNumber =>
-      DeleteRequest(IsrChangeNotificationSequenceZNode.path(sequenceNumber), ZkVersion.NoVersion)
+      DeleteRequest(IsrChangeNotificationSequenceZNode.path(sequenceNumber), ZkVersion.MatchAnyVersion)
     }
-    retryRequestsUntilConnected(deleteRequests)
+    retryRequestsUntilConnected(deleteRequests, expectedControllerEpochZkVersion)
+  }
+
+  /**
+   * Creates preferred replica election znode with partitions undergoing election
+   * @param partitions
+   * @throws KeeperException if there is an error while creating the znode
+   */
+  def createPreferredReplicaElection(partitions: Set[TopicPartition]): Unit = {
+    createRecursive(PreferredReplicaElectionZNode.path, PreferredReplicaElectionZNode.encode(partitions))
   }
 
   /**
@@ -805,10 +1095,11 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
 
   /**
    * Deletes the preferred replica election znode.
+   * @param expectedControllerEpochZkVersion expected controller epoch zkVersion.
    */
-  def deletePreferredReplicaElection(): Unit = {
-    val deleteRequest = DeleteRequest(PreferredReplicaElectionZNode.path, ZkVersion.NoVersion)
-    retryRequestUntilConnected(deleteRequest)
+  def deletePreferredReplicaElection(expectedControllerEpochZkVersion: Int): Unit = {
+    val deleteRequest = DeleteRequest(PreferredReplicaElectionZNode.path, ZkVersion.MatchAnyVersion)
+    retryRequestUntilConnected(deleteRequest, expectedControllerEpochZkVersion)
   }
 
   /**
@@ -827,10 +1118,11 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
 
   /**
    * Deletes the controller znode.
+   * @param expectedControllerEpochZkVersion expected controller epoch zkVersion.
    */
-  def deleteController(): Unit = {
-    val deleteRequest = DeleteRequest(ControllerZNode.path, ZkVersion.NoVersion)
-    retryRequestUntilConnected(deleteRequest)
+  def deleteController(expectedControllerEpochZkVersion: Int): Unit = {
+    val deleteRequest = DeleteRequest(ControllerZNode.path, ZkVersion.MatchAnyVersion)
+    retryRequestUntilConnected(deleteRequest, expectedControllerEpochZkVersion)
   }
 
   /**
@@ -852,29 +1144,35 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
   /**
    * Recursively deletes the topic znode.
    * @param topic the topic whose topic znode we wish to delete.
+   * @param expectedControllerEpochZkVersion expected controller epoch zkVersion.
    */
-  def deleteTopicZNode(topic: String): Unit = {
-    deleteRecursive(TopicZNode.path(topic))
+  def deleteTopicZNode(topic: String, expectedControllerEpochZkVersion: Int): Unit = {
+    deleteRecursive(TopicZNode.path(topic), expectedControllerEpochZkVersion)
   }
 
   /**
    * Deletes the topic configs for the given topics.
    * @param topics the topics whose configs we wish to delete.
+   * @param expectedControllerEpochZkVersion expected controller epoch zkVersion.
    */
-  def deleteTopicConfigs(topics: Seq[String]): Unit = {
-    val deleteRequests = topics.map(topic => DeleteRequest(ConfigEntityZNode.path(ConfigType.Topic, topic), ZkVersion.NoVersion))
-    retryRequestsUntilConnected(deleteRequests)
+  def deleteTopicConfigs(topics: Seq[String], expectedControllerEpochZkVersion: Int): Unit = {
+    val deleteRequests = topics.map(topic => DeleteRequest(ConfigEntityZNode.path(ConfigType.Topic, topic),
+      ZkVersion.MatchAnyVersion))
+    retryRequestsUntilConnected(deleteRequests, expectedControllerEpochZkVersion)
   }
 
   //Acl management methods
 
   /**
-   * Creates the required zk nodes for Acl storage
+   * Creates the required zk nodes for Acl storage and Acl change storage.
    */
   def createAclPaths(): Unit = {
-    createRecursive(AclZNode.path, throwIfPathExists = false)
-    createRecursive(AclChangeNotificationZNode.path, throwIfPathExists = false)
-    ResourceType.values.foreach(resource => createRecursive(ResourceTypeZNode.path(resource.name), throwIfPathExists = false))
+    ZkAclStore.stores.foreach(store => {
+      createRecursive(store.aclPath, throwIfPathExists = false)
+      AclEntry.ResourceTypes.foreach(resourceType => createRecursive(store.path(resourceType), throwIfPathExists = false))
+    })
+
+    ZkAclChangeStore.stores.foreach(store => createRecursive(store.aclChangePath, throwIfPathExists = false))
   }
 
   /**
@@ -882,12 +1180,12 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
    * @param resource Resource to get VersionedAcls for
    * @return  VersionedAcls
    */
-  def getVersionedAclsForResource(resource: Resource): VersionedAcls = {
+  def getVersionedAclsForResource(resource: ResourcePattern): VersionedAcls = {
     val getDataRequest = GetDataRequest(ResourceZNode.path(resource))
     val getDataResponse = retryRequestUntilConnected(getDataRequest)
     getDataResponse.resultCode match {
       case Code.OK => ResourceZNode.decode(getDataResponse.data, getDataResponse.stat)
-      case Code.NONODE => VersionedAcls(Set(), -1)
+      case Code.NONODE => NoAcls
       case _ => throw getDataResponse.resultException.get
     }
   }
@@ -900,99 +1198,119 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
    * @param expectedVersion
    * @return true if the update was successful and the new version
    */
-  def conditionalSetOrCreateAclsForResource(resource: Resource, aclsSet: Set[Acl], expectedVersion: Int): (Boolean, Int) = {
+  def conditionalSetAclsForResource(resource: ResourcePattern,
+                                    aclsSet: Set[AclEntry],
+                                    expectedVersion: Int): (Boolean, Int) = {
     def set(aclData: Array[Byte],  expectedVersion: Int): SetDataResponse = {
       val setDataRequest = SetDataRequest(ResourceZNode.path(resource), aclData, expectedVersion)
       retryRequestUntilConnected(setDataRequest)
     }
 
-    def create(aclData: Array[Byte]): CreateResponse = {
-      val path = ResourceZNode.path(resource)
-      val createRequest = CreateRequest(path, aclData, acls(path), CreateMode.PERSISTENT)
-      retryRequestUntilConnected(createRequest)
-    }
+    if (expectedVersion < 0)
+      throw new IllegalArgumentException(s"Invalid version $expectedVersion provided for conditional update")
 
     val aclData = ResourceZNode.encode(aclsSet)
 
     val setDataResponse = set(aclData, expectedVersion)
     setDataResponse.resultCode match {
       case Code.OK => (true, setDataResponse.stat.getVersion)
-      case Code.NONODE => {
-        val createResponse = create(aclData)
-        createResponse.resultCode match {
-          case Code.OK => (true, 0)
-          case Code.NODEEXISTS => (false, 0)
-          case _ => throw createResponse.resultException.get
-        }
-      }
-      case Code.BADVERSION => (false, 0)
+      case Code.NONODE | Code.BADVERSION  => (false, ZkVersion.UnknownVersion)
       case _ => throw setDataResponse.resultException.get
     }
   }
 
+  def createAclsForResourceIfNotExists(resource: ResourcePattern, aclsSet: Set[AclEntry]): (Boolean, Int) = {
+    def create(aclData: Array[Byte]): CreateResponse = {
+      val path = ResourceZNode.path(resource)
+      val createRequest = CreateRequest(path, aclData, defaultAcls(path), CreateMode.PERSISTENT)
+      retryRequestUntilConnected(createRequest)
+    }
+
+    val aclData = ResourceZNode.encode(aclsSet)
+
+    val createResponse = create(aclData)
+    createResponse.resultCode match {
+      case Code.OK => (true, 0)
+      case Code.NODEEXISTS => (false, ZkVersion.UnknownVersion)
+      case _ => throw createResponse.resultException.get
+    }
+  }
+
   /**
-   * Creates Acl change notification message
-   * @param resourceName resource name
+   * Creates an Acl change notification message.
+   * @param resource resource pattern that has changed
    */
-  def createAclChangeNotification(resourceName: String): Unit = {
-    val path = AclChangeNotificationSequenceZNode.createPath
-    val createRequest = CreateRequest(path, AclChangeNotificationSequenceZNode.encode(resourceName), acls(path), CreateMode.PERSISTENT_SEQUENTIAL)
+  def createAclChangeNotification(resource: ResourcePattern): Unit = {
+    val aclChange = ZkAclStore(resource.patternType).changeStore.createChangeNode(resource)
+    val createRequest = CreateRequest(aclChange.path, aclChange.bytes, defaultAcls(aclChange.path), CreateMode.PERSISTENT_SEQUENTIAL)
     val createResponse = retryRequestUntilConnected(createRequest)
-    createResponse.resultException.foreach(e => throw e)
+    createResponse.maybeThrow()
   }
 
-  def propagateLogDirEvent(brokerId: Int) {
+  def propagateLogDirEvent(brokerId: Int): Unit = {
     val logDirEventNotificationPath: String = createSequentialPersistentPath(
       LogDirEventNotificationZNode.path + "/" + LogDirEventNotificationSequenceZNode.SequenceNumberPrefix,
-      new String(LogDirEventNotificationSequenceZNode.encode(brokerId), UTF_8))
+      LogDirEventNotificationSequenceZNode.encode(brokerId))
     debug(s"Added $logDirEventNotificationPath for broker $brokerId")
   }
 
+  def propagateIsrChanges(isrChangeSet: collection.Set[TopicPartition]): Unit = {
+    val isrChangeNotificationPath: String = createSequentialPersistentPath(IsrChangeNotificationSequenceZNode.path(),
+      IsrChangeNotificationSequenceZNode.encode(isrChangeSet))
+    debug(s"Added $isrChangeNotificationPath for $isrChangeSet")
+  }
+
   /**
    * Deletes all Acl change notifications.
    * @throws KeeperException if there is an error while deleting Acl change notifications
    */
   def deleteAclChangeNotifications(): Unit = {
-    val getChildrenResponse = retryRequestUntilConnected(GetChildrenRequest(AclChangeNotificationZNode.path))
-    if (getChildrenResponse.resultCode == Code.OK) {
-      deleteAclChangeNotifications(getChildrenResponse.children)
-    } else if (getChildrenResponse.resultCode != Code.NONODE) {
-      getChildrenResponse.resultException.foreach(e => throw e)
-    }
+    ZkAclChangeStore.stores.foreach(store => {
+      val getChildrenResponse = retryRequestUntilConnected(GetChildrenRequest(store.aclChangePath, registerWatch = true))
+      if (getChildrenResponse.resultCode == Code.OK) {
+        deleteAclChangeNotifications(store.aclChangePath, getChildrenResponse.children)
+      } else if (getChildrenResponse.resultCode != Code.NONODE) {
+        getChildrenResponse.maybeThrow()
+      }
+    })
   }
 
   /**
-   * Deletes the Acl change notifications associated with the given sequence nodes
-   * @param sequenceNodes
-   */
-  private def deleteAclChangeNotifications(sequenceNodes: Seq[String]): Unit = {
+    * Deletes the Acl change notifications associated with the given sequence nodes
+    *
+    * @param aclChangePath the root path
+    * @param sequenceNodes the name of the node to delete.
+    */
+  private def deleteAclChangeNotifications(aclChangePath: String, sequenceNodes: Seq[String]): Unit = {
     val deleteRequests = sequenceNodes.map { sequenceNode =>
-      DeleteRequest(AclChangeNotificationSequenceZNode.deletePath(sequenceNode), ZkVersion.NoVersion)
+      DeleteRequest(s"$aclChangePath/$sequenceNode", ZkVersion.MatchAnyVersion)
     }
 
     val deleteResponses = retryRequestsUntilConnected(deleteRequests)
     deleteResponses.foreach { deleteResponse =>
-      if (deleteResponse.resultCode != Code.OK && deleteResponse.resultCode != Code.NONODE) {
-        deleteResponse.resultException.foreach(e => throw e)
+      if (deleteResponse.resultCode != Code.NONODE) {
+        deleteResponse.maybeThrow()
       }
     }
   }
 
   /**
-   * Gets the resource types
+   * Gets the resource types, for which ACLs are stored, for the supplied resource pattern type.
+   * @param patternType The resource pattern type to retrieve the names for.
    * @return list of resource type names
    */
-  def getResourceTypes(): Seq[String] = {
-    getChildren(AclZNode.path)
+  def getResourceTypes(patternType: PatternType): Seq[String] = {
+    getChildren(ZkAclStore(patternType).aclPath)
   }
 
   /**
-   * Gets the resource names for a give resource type
-   * @param resourceType
+   * Gets the resource names, for which ACLs are stored, for a given resource type and pattern type
+   * @param patternType The resource pattern type to retrieve the names for.
+   * @param resourceType Resource type to retrieve the names for.
    * @return list of resource names
    */
-  def getResourceNames(resourceType: String): Seq[String] = {
-    getChildren(ResourceTypeZNode.path(resourceType))
+  def getResourceNames(patternType: PatternType, resourceType: ResourceType): Seq[String] = {
+    getChildren(ZkAclStore(patternType).path(resourceType))
   }
 
   /**
@@ -1000,7 +1318,7 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
    * @param resource
    * @return delete status
    */
-  def deleteResource(resource: Resource): Boolean = {
+  def deleteResource(resource: ResourcePattern): Boolean = {
     deleteRecursive(ResourceZNode.path(resource))
   }
 
@@ -1009,7 +1327,7 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
    * @param resource
    * @return existence status
    */
-  def resourceExists(resource: Resource): Boolean = {
+  def resourceExists(resource: ResourcePattern): Boolean = {
     pathExists(ResourceZNode.path(resource))
   }
 
@@ -1019,7 +1337,7 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
    * @param expectedVersion
    * @return return true if it succeeds, false otherwise (the current version is not the expected version)
    */
-  def conditionalDelete(resource: Resource, expectedVersion: Int): Boolean = {
+  def conditionalDelete(resource: ResourcePattern, expectedVersion: Int): Boolean = {
     val deleteRequest = DeleteRequest(ResourceZNode.path(resource), expectedVersion)
     val deleteResponse = retryRequestUntilConnected(deleteRequest)
     deleteResponse.resultCode match {
@@ -1031,11 +1349,93 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
 
   /**
    * Deletes the zk node recursively
-   * @param path
-   * @return  return true if it succeeds, false otherwise
+   * @param path path to delete
+   * @param expectedControllerEpochZkVersion expected controller epoch zkVersion.
+   * @param recursiveDelete enable recursive delete
+   * @return KeeperException if there is an error while deleting the path
+   */
+  def deletePath(path: String, expectedControllerEpochZkVersion: Int = ZkVersion.MatchAnyVersion, recursiveDelete: Boolean = true): Unit = {
+    if (recursiveDelete)
+      deleteRecursive(path, expectedControllerEpochZkVersion)
+    else {
+      val deleteRequest = DeleteRequest(path, ZkVersion.MatchAnyVersion)
+      val deleteResponse = retryRequestUntilConnected(deleteRequest, expectedControllerEpochZkVersion)
+      if (deleteResponse.resultCode != Code.OK && deleteResponse.resultCode != Code.NONODE) {
+          throw deleteResponse.resultException.get
+      }
+    }
+  }
+
+  /**
+   * Creates the required zk nodes for Delegation Token storage
+   */
+  def createDelegationTokenPaths(): Unit = {
+    createRecursive(DelegationTokenChangeNotificationZNode.path, throwIfPathExists = false)
+    createRecursive(DelegationTokensZNode.path, throwIfPathExists = false)
+  }
+
+  /**
+   * Creates Delegation Token change notification message
+   * @param tokenId token Id
    */
-  def deletePath(path: String): Boolean = {
-    deleteRecursive(path)
+  def createTokenChangeNotification(tokenId: String): Unit = {
+    val path = DelegationTokenChangeNotificationSequenceZNode.createPath
+    val createRequest = CreateRequest(path, DelegationTokenChangeNotificationSequenceZNode.encode(tokenId), defaultAcls(path), CreateMode.PERSISTENT_SEQUENTIAL)
+    val createResponse = retryRequestUntilConnected(createRequest)
+    createResponse.resultException.foreach(e => throw e)
+  }
+
+  /**
+   * Sets or creates token info znode with the given token details depending on whether it already
+   * exists or not.
+   *
+   * @param token the token to set on the token znode
+   * @throws KeeperException if there is an error while setting or creating the znode
+   */
+  def setOrCreateDelegationToken(token: DelegationToken): Unit = {
+
+    def set(tokenData: Array[Byte]): SetDataResponse = {
+      val setDataRequest = SetDataRequest(DelegationTokenInfoZNode.path(token.tokenInfo().tokenId()), tokenData, ZkVersion.MatchAnyVersion)
+      retryRequestUntilConnected(setDataRequest)
+    }
+
+    def create(tokenData: Array[Byte]): CreateResponse = {
+      val path = DelegationTokenInfoZNode.path(token.tokenInfo().tokenId())
+      val createRequest = CreateRequest(path, tokenData, defaultAcls(path), CreateMode.PERSISTENT)
+      retryRequestUntilConnected(createRequest)
+    }
+
+    val tokenInfo = DelegationTokenInfoZNode.encode(token)
+    val setDataResponse = set(tokenInfo)
+    setDataResponse.resultCode match {
+      case Code.NONODE =>
+        val createDataResponse = create(tokenInfo)
+        createDataResponse.maybeThrow()
+      case _ => setDataResponse.maybeThrow()
+    }
+  }
+
+  /**
+   * Gets the Delegation Token Info
+   * @return optional TokenInfo that is Some if the token znode exists and can be parsed and None otherwise.
+   */
+  def getDelegationTokenInfo(delegationTokenId: String): Option[TokenInformation] = {
+    val getDataRequest = GetDataRequest(DelegationTokenInfoZNode.path(delegationTokenId))
+    val getDataResponse = retryRequestUntilConnected(getDataRequest)
+    getDataResponse.resultCode match {
+      case Code.OK => DelegationTokenInfoZNode.decode(getDataResponse.data)
+      case Code.NONODE => None
+      case _ => throw getDataResponse.resultException.get
+    }
+  }
+
+  /**
+   * Deletes the given Delegation token node
+   * @param delegationTokenId
+   * @return delete status
+   */
+  def deleteDelegationToken(delegationTokenId: String): Boolean = {
+    deleteRecursive(DelegationTokenInfoZNode.path(delegationTokenId))
   }
 
   /**
@@ -1108,6 +1508,7 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
    * Close the underlying ZooKeeperClient.
    */
   def close(): Unit = {
+    removeMetric("ZooKeeperRequestLatencyMs")
     zooKeeperClient.close()
   }
 
@@ -1138,13 +1539,123 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
     if (setDataResponse.resultCode == Code.NONODE) {
       createConsumerOffset(group, topicPartition, offset)
     } else {
-      setDataResponse.resultException.foreach(e => throw e)
+      setDataResponse.maybeThrow()
     }
   }
 
+  /**
+    * Get the cluster id.
+    * @return optional cluster id in String.
+    */
+  def getClusterId: Option[String] = {
+    val getDataRequest = GetDataRequest(ClusterIdZNode.path)
+    val getDataResponse = retryRequestUntilConnected(getDataRequest)
+    getDataResponse.resultCode match {
+      case Code.OK => Some(ClusterIdZNode.fromJson(getDataResponse.data))
+      case Code.NONODE => None
+      case _ => throw getDataResponse.resultException.get
+    }
+  }
+
+  /**
+    * Return the ACLs of the node of the given path
+    * @param path the given path for the node
+    * @return the ACL array of the given node.
+    */
+  def getAcl(path: String): Seq[ACL] = {
+    val getAclRequest = GetAclRequest(path)
+    val getAclResponse = retryRequestUntilConnected(getAclRequest)
+    getAclResponse.resultCode match {
+      case Code.OK => getAclResponse.acl
+      case _ => throw getAclResponse.resultException.get
+    }
+  }
+
+  /**
+    * sets the ACLs to the node of the given path
+    * @param path the given path for the node
+    * @param acl the given acl for the node
+    */
+  def setAcl(path: String, acl: Seq[ACL]): Unit = {
+    val setAclRequest = SetAclRequest(path, acl, ZkVersion.MatchAnyVersion)
+    val setAclResponse = retryRequestUntilConnected(setAclRequest)
+    setAclResponse.maybeThrow()
+  }
+
+  /**
+    * Create the cluster Id. If the cluster id already exists, return the current cluster id.
+    * @return  cluster id
+    */
+  def createOrGetClusterId(proposedClusterId: String): String = {
+    try {
+      createRecursive(ClusterIdZNode.path, ClusterIdZNode.toJson(proposedClusterId))
+      proposedClusterId
+    } catch {
+      case _: NodeExistsException => getClusterId.getOrElse(
+        throw new KafkaException("Failed to get cluster id from Zookeeper. This can happen if /cluster/id is deleted from Zookeeper."))
+    }
+  }
+
+  /**
+    * Generate a broker id by updating the broker sequence id path in ZK and return the version of the path.
+    * The version is incremented by one on every update starting from 1.
+    * @return sequence number as the broker id
+    */
+  def generateBrokerSequenceId(): Int = {
+    val setDataRequest = SetDataRequest(BrokerSequenceIdZNode.path, Array.empty[Byte], ZkVersion.MatchAnyVersion)
+    val setDataResponse = retryRequestUntilConnected(setDataRequest)
+    setDataResponse.resultCode match {
+      case Code.OK => setDataResponse.stat.getVersion
+      case Code.NONODE =>
+        // maker sure the path exists
+        createRecursive(BrokerSequenceIdZNode.path, Array.empty[Byte], throwIfPathExists = false)
+        generateBrokerSequenceId()
+      case _ => throw setDataResponse.resultException.get
+    }
+  }
+
+  /**
+    * Pre-create top level paths in ZK if needed.
+    */
+  def createTopLevelPaths(): Unit = {
+    ZkData.PersistentZkPaths.foreach(makeSurePersistentPathExists(_))
+  }
+
+  /**
+    * Make sure a persistent path exists in ZK.
+    * @param path
+    */
+  def makeSurePersistentPathExists(path: String): Unit = {
+    createRecursive(path, data = null, throwIfPathExists = false)
+  }
+
+  def createFeatureZNode(nodeContents: FeatureZNode): Unit = {
+    val createRequest = CreateRequest(
+      FeatureZNode.path,
+      FeatureZNode.encode(nodeContents),
+      defaultAcls(FeatureZNode.path),
+      CreateMode.PERSISTENT)
+    val response = retryRequestUntilConnected(createRequest)
+    response.maybeThrow()
+  }
+
+  def updateFeatureZNode(nodeContents: FeatureZNode): Int = {
+    val setRequest = SetDataRequest(
+      FeatureZNode.path,
+      FeatureZNode.encode(nodeContents),
+      ZkVersion.MatchAnyVersion)
+    val response = retryRequestUntilConnected(setRequest)
+    response.maybeThrow()
+    response.stat.getVersion
+  }
+
+  def deleteFeatureZNode(): Unit = {
+    deletePath(FeatureZNode.path, ZkVersion.MatchAnyVersion, false)
+  }
+
   private def setConsumerOffset(group: String, topicPartition: TopicPartition, offset: Long): SetDataResponse = {
     val setDataRequest = SetDataRequest(ConsumerOffset.path(group, topicPartition.topic, topicPartition.partition),
-      ConsumerOffset.encode(offset), ZkVersion.NoVersion)
+      ConsumerOffset.encode(offset), ZkVersion.MatchAnyVersion)
     retryRequestUntilConnected(setDataRequest)
   }
 
@@ -1156,25 +1667,25 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
   /**
    * Deletes the given zk path recursively
    * @param path
+   * @param expectedControllerEpochZkVersion expected controller epoch zkVersion.
    * @return true if path gets deleted successfully, false if root path doesn't exist
    * @throws KeeperException if there is an error while deleting the znodes
    */
-  private[zk] def deleteRecursive(path: String): Boolean = {
-    val getChildrenResponse = retryRequestUntilConnected(GetChildrenRequest(path))
+  def deleteRecursive(path: String, expectedControllerEpochZkVersion: Int = ZkVersion.MatchAnyVersion): Boolean = {
+    val getChildrenResponse = retryRequestUntilConnected(GetChildrenRequest(path, registerWatch = true))
     getChildrenResponse.resultCode match {
       case Code.OK =>
-        getChildrenResponse.children.foreach(child => deleteRecursive(s"$path/$child"))
-        val deleteResponse = retryRequestUntilConnected(DeleteRequest(path, ZkVersion.NoVersion))
-        if (deleteResponse.resultCode != Code.OK && deleteResponse.resultCode != Code.NONODE) {
+        getChildrenResponse.children.foreach(child => deleteRecursive(s"$path/$child", expectedControllerEpochZkVersion))
+        val deleteResponse = retryRequestUntilConnected(DeleteRequest(path, ZkVersion.MatchAnyVersion), expectedControllerEpochZkVersion)
+        if (deleteResponse.resultCode != Code.OK && deleteResponse.resultCode != Code.NONODE)
           throw deleteResponse.resultException.get
-        }
         true
       case Code.NONODE => false
       case _ => throw getChildrenResponse.resultException.get
     }
   }
 
-  private[zk] def pathExists(path: String): Boolean = {
+  def pathExists(path: String): Boolean = {
     val existsRequest = ExistsRequest(path)
     val existsResponse = retryRequestUntilConnected(existsRequest)
     existsResponse.resultCode match {
@@ -1184,7 +1695,7 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
     }
   }
 
-  private[zk] def createRecursive(path: String, data: Array[Byte] = null, throwIfPathExists: Boolean = true) = {
+  private[kafka] def createRecursive(path: String, data: Array[Byte] = null, throwIfPathExists: Boolean = true) = {
 
     def parentPath(path: String): String = {
       val indexOfLastSlash = path.lastIndexOf("/")
@@ -1193,7 +1704,7 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
     }
 
     def createRecursive0(path: String): Unit = {
-      val createRequest = CreateRequest(path, null, acls(path), CreateMode.PERSISTENT)
+      val createRequest = CreateRequest(path, null, defaultAcls(path), CreateMode.PERSISTENT)
       var createResponse = retryRequestUntilConnected(createRequest)
       if (createResponse.resultCode == Code.NONODE) {
         createRecursive0(parentPath(path))
@@ -1206,58 +1717,72 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
       }
     }
 
-    val createRequest = CreateRequest(path, data, acls(path), CreateMode.PERSISTENT)
+    val createRequest = CreateRequest(path, data, defaultAcls(path), CreateMode.PERSISTENT)
     var createResponse = retryRequestUntilConnected(createRequest)
 
     if (throwIfPathExists && createResponse.resultCode == Code.NODEEXISTS) {
-      createResponse.resultException.foreach(e => throw e)
+      createResponse.maybeThrow()
     } else if (createResponse.resultCode == Code.NONODE) {
       createRecursive0(parentPath(path))
       createResponse = retryRequestUntilConnected(createRequest)
-      createResponse.resultException.foreach(e => throw e)
+      if (throwIfPathExists || createResponse.resultCode != Code.NODEEXISTS)
+        createResponse.maybeThrow()
     } else if (createResponse.resultCode != Code.NODEEXISTS)
-      createResponse.resultException.foreach(e => throw e)
+      createResponse.maybeThrow()
 
   }
 
-  private def createTopicPartition(partitions: Seq[TopicPartition]): Seq[CreateResponse] = {
+  private def createTopicPartition(partitions: Seq[TopicPartition], expectedControllerEpochZkVersion: Int): Seq[CreateResponse] = {
     val createRequests = partitions.map { partition =>
       val path = TopicPartitionZNode.path(partition)
-      CreateRequest(path, null, acls(path), CreateMode.PERSISTENT, Some(partition))
+      CreateRequest(path, null, defaultAcls(path), CreateMode.PERSISTENT, Some(partition))
     }
-    retryRequestsUntilConnected(createRequests)
+    retryRequestsUntilConnected(createRequests, expectedControllerEpochZkVersion)
   }
 
-  private def createTopicPartitions(topics: Seq[String]): 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, acls(path), CreateMode.PERSISTENT, Some(topic))
+      CreateRequest(path, null, defaultAcls(path), CreateMode.PERSISTENT, Some(topic))
     }
-    retryRequestsUntilConnected(createRequests)
+    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)
   }
 
-  private def acls(path: String): Seq[ACL] = {
-    import scala.collection.JavaConverters._
-    ZkUtils.defaultAcls(isSecure, path).asScala
+  def defaultAcls(path: String): Seq[ACL] = ZkData.defaultAcls(isSecure, path)
+
+  def secure: Boolean = isSecure
+
+  private[zk] def retryRequestUntilConnected[Req <: AsyncRequest](request: Req, expectedControllerZkVersion: Int = ZkVersion.MatchAnyVersion): Req#Response = {
+    retryRequestsUntilConnected(Seq(request), expectedControllerZkVersion).head
   }
 
-  private def retryRequestUntilConnected[Req <: AsyncRequest](request: Req): Req#Response = {
-    retryRequestsUntilConnected(Seq(request)).head
+  private def retryRequestsUntilConnected[Req <: AsyncRequest](requests: Seq[Req], expectedControllerZkVersion: Int): Seq[Req#Response] = {
+    expectedControllerZkVersion match {
+      case ZkVersion.MatchAnyVersion => retryRequestsUntilConnected(requests)
+      case version if version >= 0 =>
+        retryRequestsUntilConnected(requests.map(wrapRequestWithControllerEpochCheck(_, version)))
+          .map(unwrapResponseWithControllerEpochCheck(_).asInstanceOf[Req#Response])
+      case invalidVersion =>
+        throw new IllegalArgumentException(s"Expected controller epoch zkVersion $invalidVersion should be non-negative or equal to ${ZkVersion.MatchAnyVersion}")
+    }
   }
 
   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)
 
+      batchResponses.foreach(response => latencyMetric.update(response.metadata.responseTimeMs))
+
       // Only execute slow path if we find a response with CONNECTIONLOSS
       if (batchResponses.exists(_.resultCode == Code.CONNECTIONLOSS)) {
         val requestResponsePairs = remainingRequests.zip(batchResponses)
@@ -1280,49 +1805,113 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
     responses
   }
 
-  def checkedEphemeralCreate(path: String, data: Array[Byte]): Unit = {
+  private def checkedEphemeralCreate(path: String, data: Array[Byte]): Stat = {
     val checkedEphemeral = new CheckedEphemeral(path, data)
     info(s"Creating $path (is it secure? $isSecure)")
-    val code = checkedEphemeral.create()
-    info(s"Result of znode creation at $path is: $code")
-    code match {
-      case Code.OK =>
-      case _ => throw KeeperException.create(code)
-    }
+    val stat = checkedEphemeral.create()
+    info(s"Stat of the created znode at $path is: $stat")
+    stat
+  }
+
+  private def isZKSessionIdDiffFromCurrentZKSessionId(): Boolean = {
+    zooKeeperClient.sessionId != currentZooKeeperSessionId
+  }
+
+  private def isZKSessionTheEphemeralOwner(ephemeralOwnerId: Long): Boolean = {
+    ephemeralOwnerId == currentZooKeeperSessionId
+  }
+
+  private[zk] def shouldReCreateEphemeralZNode(ephemeralOwnerId: Long): Boolean = {
+    isZKSessionTheEphemeralOwner(ephemeralOwnerId) && isZKSessionIdDiffFromCurrentZKSessionId()
+  }
+
+  private def updateCurrentZKSessionId(newSessionId: Long): Unit = {
+    currentZooKeeperSessionId = newSessionId
   }
 
   private class CheckedEphemeral(path: String, data: Array[Byte]) extends Logging {
-    def create(): Code = {
-      val createRequest = CreateRequest(path, data, acls(path), CreateMode.EPHEMERAL)
-      val createResponse = retryRequestUntilConnected(createRequest)
-      val code = createResponse.resultCode
-      code match {
-        case Code.OK => code
-        case Code.NODEEXISTS => get()
-        case _ =>
+    def create(): Stat = {
+      val response = retryRequestUntilConnected(
+        MultiRequest(Seq(
+          CreateOp(path, null, defaultAcls(path), CreateMode.EPHEMERAL),
+          SetDataOp(path, data, 0)))
+      )
+      val stat = response.resultCode match {
+        case Code.OK =>
+          val setDataResult = response.zkOpResults(1).rawOpResult.asInstanceOf[SetDataResult]
+          setDataResult.getStat
+        case Code.NODEEXISTS =>
+          getAfterNodeExists()
+        case code =>
           error(s"Error while creating ephemeral at $path with return code: $code")
+          throw KeeperException.create(code)
+      }
+
+      // At this point, we need to save a reference to the zookeeper session id.
+      // This is done here since the Zookeeper session id may not be available at the Object creation time.
+      // This is assuming the 'retryRequestUntilConnected' method got connected and a valid session id is present.
+      // This code is part of the workaround done in the KAFKA-7165, once ZOOKEEPER-2985 is complete, this code
+      // must be deleted.
+      updateCurrentZKSessionId(zooKeeperClient.sessionId)
+
+      stat
+    }
+
+    // This method is part of the work around done in the KAFKA-7165, once ZOOKEEPER-2985 is complete, this code must
+    // be deleted.
+    private def delete(): Code = {
+      val deleteRequest = DeleteRequest(path, ZkVersion.MatchAnyVersion)
+      val deleteResponse = retryRequestUntilConnected(deleteRequest)
+      deleteResponse.resultCode match {
+        case code@ Code.OK => code
+        case code@ Code.NONODE => code
+        case code =>
+          error(s"Error while deleting ephemeral node at $path with return code: $code")
           code
       }
     }
 
-    private def get(): Code = {
+    private def reCreate(): Stat = {
+      val codeAfterDelete = delete()
+      val codeAfterReCreate = codeAfterDelete
+      debug(s"Result of znode ephemeral deletion at $path is: $codeAfterDelete")
+      if (codeAfterDelete == Code.OK || codeAfterDelete == Code.NONODE) {
+        create()
+      } else {
+        throw KeeperException.create(codeAfterReCreate)
+      }
+    }
+
+    private def getAfterNodeExists(): Stat = {
       val getDataRequest = GetDataRequest(path)
       val getDataResponse = retryRequestUntilConnected(getDataRequest)
-      val code = getDataResponse.resultCode
-      code match {
+      val ephemeralOwnerId = getDataResponse.stat.getEphemeralOwner
+      getDataResponse.resultCode match {
+        // At this point, the Zookeeper session could be different (due a 'Session expired') from the one that initially
+        // registered the Broker into the Zookeeper ephemeral node, but the znode is still present in ZooKeeper.
+        // The expected behaviour is that Zookeeper server removes the ephemeral node associated with the expired session
+        // but due an already reported bug in Zookeeper (ZOOKEEPER-2985) this is not happening, so, the following check
+        // will validate if this Broker got registered with the previous (expired) session and try to register again,
+        // deleting the ephemeral node and creating it again.
+        // This code is part of the work around done in the KAFKA-7165, once ZOOKEEPER-2985 is complete, this code must
+        // be deleted.
+        case Code.OK if shouldReCreateEphemeralZNode(ephemeralOwnerId) =>
+          info(s"Was not possible to create the ephemeral at $path, node already exists and owner " +
+            s"'$ephemeralOwnerId' does not match current session '${zooKeeperClient.sessionId}'" +
+            s", trying to delete and re-create it with the newest Zookeeper session")
+          reCreate()
+        case Code.OK if ephemeralOwnerId != zooKeeperClient.sessionId =>
+          error(s"Error while creating ephemeral at $path, node already exists and owner " +
+            s"'$ephemeralOwnerId' does not match current session '${zooKeeperClient.sessionId}'")
+          throw KeeperException.create(Code.NODEEXISTS)
         case Code.OK =>
-          if (getDataResponse.stat.getEphemeralOwner != zooKeeperClient.sessionId) {
-            error(s"Error while creating ephemeral at $path with return code: $code")
-            Code.NODEEXISTS
-          } else {
-            code
-          }
+          getDataResponse.stat
         case Code.NONODE =>
-          info(s"The ephemeral node at $path went away while reading it")
+          info(s"The ephemeral node at $path went away while reading it, attempting create() again")
           create()
-        case _ =>
-          error(s"Error while creating ephemeral at $path with return code: $code")
-          code
+        case code =>
+          error(s"Error while creating ephemeral at $path as it already exists and error getting the node data due to $code")
+          throw KeeperException.create(code)
       }
     }
   }
@@ -1331,13 +1920,97 @@ class KafkaZkClient(zooKeeperClient: ZooKeeperClient, isSecure: Boolean) extends
 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.
+   *
+   * The metric group and type are preserved by default for compatibility with previous versions.
+   */
+  def apply(connectString: String,
+            isSecure: Boolean,
+            sessionTimeoutMs: Int,
+            connectionTimeoutMs: Int,
+            maxInFlightRequests: Int,
+            time: Time,
+            metricGroup: String = "kafka.server",
+            metricType: String = "SessionExpireListener",
+            name: Option[String] = None,
+            zkClientConfig: Option[ZKClientConfig] = None) = {
+    val zooKeeperClient = new ZooKeeperClient(connectString, sessionTimeoutMs, connectionTimeoutMs, maxInFlightRequests,
+      time, metricGroup, metricType, name, zkClientConfig)
+    new KafkaZkClient(zooKeeperClient, isSecure, time)
+  }
+
+  // A helper function to transform a regular request into a MultiRequest
+  // with the check on controller epoch znode zkVersion.
+  // This is used for fencing zookeeper updates in controller.
+  private def wrapRequestWithControllerEpochCheck(request: AsyncRequest, expectedControllerZkVersion: Int): MultiRequest = {
+      val checkOp = CheckOp(ControllerEpochZNode.path, expectedControllerZkVersion)
+      request match {
+        case CreateRequest(path, data, acl, createMode, ctx) =>
+          MultiRequest(Seq(checkOp, CreateOp(path, data, acl, createMode)), ctx)
+        case DeleteRequest(path, version, ctx) =>
+          MultiRequest(Seq(checkOp, DeleteOp(path, version)), ctx)
+        case SetDataRequest(path, data, version, ctx) =>
+          MultiRequest(Seq(checkOp, SetDataOp(path, data, version)), ctx)
+        case _ => throw new IllegalStateException(s"$request does not need controller epoch check")
+      }
+  }
+
+  // A helper function to transform a MultiResponse with the check on
+  // controller epoch znode zkVersion back into a regular response.
+  // ControllerMovedException will be thrown if the controller epoch
+  // znode zkVersion check fails. This is used for fencing zookeeper
+  // updates in controller.
+  private def unwrapResponseWithControllerEpochCheck(response: AsyncResponse): AsyncResponse = {
+    response match {
+      case MultiResponse(resultCode, _, ctx, zkOpResults, responseMetadata) =>
+        zkOpResults match {
+          case Seq(ZkOpResult(checkOp: CheckOp, checkOpResult), zkOpResult) =>
+            checkOpResult match {
+              case errorResult: ErrorResult =>
+                if (checkOp.path.equals(ControllerEpochZNode.path)) {
+                  val errorCode = Code.get(errorResult.getErr)
+                  if (errorCode == Code.BADVERSION)
+                  // Throw ControllerMovedException when the zkVersionCheck is performed on the controller epoch znode and the check fails
+                    throw new ControllerMovedException(s"Controller epoch zkVersion check fails. Expected zkVersion = ${checkOp.version}")
+                  else if (errorCode != Code.OK)
+                    throw KeeperException.create(errorCode, checkOp.path)
+                }
+              case _ =>
+            }
+            val rawOpResult = zkOpResult.rawOpResult
+            zkOpResult.zkOp match {
+              case createOp: CreateOp =>
+                val name = rawOpResult match {
+                  case c: CreateResult => c.getPath
+                  case _ => null
+                }
+                CreateResponse(resultCode, createOp.path, ctx, name, responseMetadata)
+              case deleteOp: DeleteOp =>
+                DeleteResponse(resultCode, deleteOp.path, ctx, responseMetadata)
+              case setDataOp: SetDataOp =>
+                val stat = rawOpResult match {
+                  case s: SetDataResult => s.getStat
+                  case _ => null
+                }
+                SetDataResponse(resultCode, setDataOp.path, ctx, stat, responseMetadata)
+              case zkOp => throw new IllegalStateException(s"Unexpected zkOp: $zkOp")
+            }
+          case null => throw KeeperException.create(resultCode)
+          case _ => throw new IllegalStateException(s"Cannot unwrap $response because the first zookeeper op is not check op in original MultiRequest")
+        }
+      case _ => throw new IllegalStateException(s"Cannot unwrap $response because it is not a MultiResponse")
+    }
+  }
 }
diff --git a/core/src/main/scala/kafka/zk/ZkData.scala b/core/src/main/scala/kafka/zk/ZkData.scala
index be6efca246cc6..868db61856b83 100644
--- a/core/src/main/scala/kafka/zk/ZkData.scala
+++ b/core/src/main/scala/kafka/zk/ZkData.scala
@@ -17,30 +17,45 @@
 package kafka.zk
 
 import java.nio.charset.StandardCharsets.UTF_8
+import java.util
 import java.util.Properties
 
-import kafka.api.{ApiVersion, KAFKA_0_10_0_IV1, LeaderAndIsr}
+import com.fasterxml.jackson.annotation.JsonProperty
+import com.fasterxml.jackson.core.JsonProcessingException
+import kafka.api.{ApiVersion, KAFKA_0_10_0_IV1, KAFKA_2_7_IV0, LeaderAndIsr}
 import kafka.cluster.{Broker, EndPoint}
-import kafka.controller.{IsrChangeNotificationHandler, LeaderIsrAndControllerEpoch}
-import kafka.security.auth.{Acl, Resource}
-import kafka.security.auth.SimpleAclAuthorizer.VersionedAcls
+import kafka.common.{NotificationHandler, ZkNodeChangeNotificationListener}
+import kafka.controller.{IsrChangeNotificationHandler, LeaderIsrAndControllerEpoch, ReplicaAssignment}
+import kafka.security.authorizer.AclAuthorizer.VersionedAcls
+import kafka.security.authorizer.AclEntry
+import kafka.server.{ConfigType, DelegationTokenManager}
 import kafka.utils.Json
-import org.apache.kafka.common.TopicPartition
-import org.apache.zookeeper.data.Stat
+import kafka.utils.json.JsonObject
+import org.apache.kafka.common.{KafkaException, TopicPartition, Uuid}
+import org.apache.kafka.common.errors.UnsupportedVersionException
+import org.apache.kafka.common.feature.{Features, FinalizedVersionRange, SupportedVersionRange}
+import org.apache.kafka.common.feature.Features._
+import org.apache.kafka.common.network.ListenerName
+import org.apache.kafka.common.resource.{PatternType, ResourcePattern, ResourceType}
+import org.apache.kafka.common.security.auth.SecurityProtocol
+import org.apache.kafka.common.security.token.delegation.{DelegationToken, TokenInformation}
+import org.apache.kafka.common.utils.{SecurityUtils, Time}
+import org.apache.zookeeper.ZooDefs
+import org.apache.zookeeper.data.{ACL, Stat}
+
+import scala.beans.BeanProperty
+import scala.jdk.CollectionConverters._
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.{Map, Seq, immutable, mutable}
+import scala.util.{Failure, Success, Try}
 
 // This file contains objects for encoding/decoding data stored in ZooKeeper nodes (znodes).
 
-object LeaderAndIsrZNode {
-  def encode(leaderAndIsr: LeaderAndIsr, controllerEpoch: Int): String = {
-    Json.encode(Map("version" -> 1, "leader" -> leaderAndIsr.leader, "leader_epoch" -> leaderAndIsr.leaderEpoch,
-                "controller_epoch" -> controllerEpoch, "isr" -> leaderAndIsr.isr))
-  }
-}
-
 object ControllerZNode {
   def path = "/controller"
-  def encode(brokerId: Int, timestamp: Long): Array[Byte] =
-    Json.encodeAsBytes(Map("version" -> 1, "brokerid" -> brokerId, "timestamp" -> timestamp.toString))
+  def encode(brokerId: Int, timestamp: Long): Array[Byte] = {
+    Json.encodeAsBytes(Map("version" -> 1, "brokerid" -> brokerId, "timestamp" -> timestamp.toString).asJava)
+  }
   def decode(bytes: Array[Byte]): Option[Int] = Json.parseBytes(bytes).map { js =>
     js.asJsonObject("brokerid").to[Int]
   }
@@ -65,21 +80,197 @@ object BrokerIdsZNode {
   def encode: Array[Byte] = null
 }
 
+object BrokerInfo {
+
+  /**
+   * - Create a broker info with v5 json format if the apiVersion is 2.7.x or above.
+   * - Create a broker info with v4 json format (which includes multiple endpoints and rack) if
+   *   the apiVersion is 0.10.0.X or above but lesser than 2.7.x.
+   * - Register the broker with v2 json format otherwise.
+   *
+   * Due to KAFKA-3100, 0.9.0.0 broker and old clients will break if JSON version is above 2.
+   *
+   * We include v2 to make it possible for the broker to migrate from 0.9.0.0 to 0.10.0.X or above
+   * without having to upgrade to 0.9.0.1 first (clients have to be upgraded to 0.9.0.1 in
+   * any case).
+   */
+  def apply(broker: Broker, apiVersion: ApiVersion, jmxPort: Int): BrokerInfo = {
+    val version = {
+      if (apiVersion >= KAFKA_2_7_IV0)
+        5
+      else if (apiVersion >= KAFKA_0_10_0_IV1)
+        4
+      else
+        2
+    }
+    BrokerInfo(broker, version, jmxPort)
+  }
+
+}
+
+case class BrokerInfo(broker: Broker, version: Int, jmxPort: Int) {
+  val path: String = BrokerIdZNode.path(broker.id)
+  def toJsonBytes: Array[Byte] = BrokerIdZNode.encode(this)
+}
+
 object BrokerIdZNode {
+  private val HostKey = "host"
+  private val PortKey = "port"
+  private val VersionKey = "version"
+  private val EndpointsKey = "endpoints"
+  private val RackKey = "rack"
+  private val JmxPortKey = "jmx_port"
+  private val ListenerSecurityProtocolMapKey = "listener_security_protocol_map"
+  private val TimestampKey = "timestamp"
+  private val FeaturesKey = "features"
+
   def path(id: Int) = s"${BrokerIdsZNode.path}/$id"
-  def encode(id: Int,
-             host: String,
-             port: Int,
-             advertisedEndpoints: Seq[EndPoint],
-             jmxPort: Int,
-             rack: Option[String],
-             apiVersion: ApiVersion): Array[Byte] = {
-    val version = if (apiVersion >= KAFKA_0_10_0_IV1) 4 else 2
-    Broker.toJson(version, id, host, port, advertisedEndpoints, jmxPort, rack).getBytes(UTF_8)
+
+  /**
+   * Encode to JSON bytes.
+   *
+   * The JSON format includes a top level host and port for compatibility with older clients.
+   */
+  def encode(version: Int, host: String, port: Int, advertisedEndpoints: Seq[EndPoint], jmxPort: Int,
+             rack: Option[String], features: Features[SupportedVersionRange]): Array[Byte] = {
+    val jsonMap = collection.mutable.Map(VersionKey -> version,
+      HostKey -> host,
+      PortKey -> port,
+      EndpointsKey -> advertisedEndpoints.map(_.connectionString).toBuffer.asJava,
+      JmxPortKey -> jmxPort,
+      TimestampKey -> Time.SYSTEM.milliseconds().toString
+    )
+    rack.foreach(rack => if (version >= 3) jsonMap += (RackKey -> rack))
+
+    if (version >= 4) {
+      jsonMap += (ListenerSecurityProtocolMapKey -> advertisedEndpoints.map { endPoint =>
+        endPoint.listenerName.value -> endPoint.securityProtocol.name
+      }.toMap.asJava)
+    }
+
+    if (version >= 5) {
+      jsonMap += (FeaturesKey -> features.toMap)
+    }
+    Json.encodeAsBytes(jsonMap.asJava)
   }
 
-  def decode(id: Int, bytes: Array[Byte]): Broker = {
-    Broker.createBroker(id, new String(bytes, UTF_8))
+  def encode(brokerInfo: BrokerInfo): Array[Byte] = {
+    val broker = brokerInfo.broker
+    // the default host and port are here for compatibility with older clients that only support PLAINTEXT
+    // we choose the first plaintext port, if there is one
+    // or we register an empty endpoint, which means that older clients will not be able to connect
+    val plaintextEndpoint = broker.endPoints.find(_.securityProtocol == SecurityProtocol.PLAINTEXT).getOrElse(
+      new EndPoint(null, -1, null, null))
+    encode(brokerInfo.version, plaintextEndpoint.host, plaintextEndpoint.port, broker.endPoints, brokerInfo.jmxPort,
+      broker.rack, broker.features)
+  }
+
+  def featuresAsJavaMap(brokerInfo: JsonObject): util.Map[String, util.Map[String, java.lang.Short]] = {
+    FeatureZNode.asJavaMap(brokerInfo
+      .get(FeaturesKey)
+      .flatMap(_.to[Option[Map[String, Map[String, Int]]]])
+      .map(theMap => theMap.map {
+         case(featureName, versionsInfo) => featureName -> versionsInfo.map {
+           case(label, version) => label -> version.asInstanceOf[Short]
+         }.toMap
+      }.toMap)
+      .getOrElse(Map[String, Map[String, Short]]()))
+  }
+
+  /**
+    * Create a BrokerInfo object from id and JSON bytes.
+    *
+    * @param id
+    * @param jsonBytes
+    *
+    * Version 1 JSON schema for a broker is:
+    * {
+    *   "version":1,
+    *   "host":"localhost",
+    *   "port":9092
+    *   "jmx_port":9999,
+    *   "timestamp":"2233345666"
+    * }
+    *
+    * Version 2 JSON schema for a broker is:
+    * {
+    *   "version":2,
+    *   "host":"localhost",
+    *   "port":9092,
+    *   "jmx_port":9999,
+    *   "timestamp":"2233345666",
+    *   "endpoints":["PLAINTEXT://host1:9092", "SSL://host1:9093"]
+    * }
+    *
+    * Version 3 JSON schema for a broker is:
+    * {
+    *   "version":3,
+    *   "host":"localhost",
+    *   "port":9092,
+    *   "jmx_port":9999,
+    *   "timestamp":"2233345666",
+    *   "endpoints":["PLAINTEXT://host1:9092", "SSL://host1:9093"],
+    *   "rack":"dc1"
+    * }
+    *
+    * Version 4 JSON schema for a broker is:
+    * {
+    *   "version":4,
+    *   "host":"localhost",
+    *   "port":9092,
+    *   "jmx_port":9999,
+    *   "timestamp":"2233345666",
+    *   "endpoints":["CLIENT://host1:9092", "REPLICATION://host1:9093"],
+    *   "rack":"dc1"
+    * }
+    *
+    * Version 5 (current) JSON schema for a broker is:
+    * {
+    *   "version":5,
+    *   "host":"localhost",
+    *   "port":9092,
+    *   "jmx_port":9999,
+    *   "timestamp":"2233345666",
+    *   "endpoints":["CLIENT://host1:9092", "REPLICATION://host1:9093"],
+    *   "rack":"dc1",
+    *   "features": {"feature": {"min_version":1, "first_active_version":2, "max_version":3}}
+    * }
+    */
+  def decode(id: Int, jsonBytes: Array[Byte]): BrokerInfo = {
+    Json.tryParseBytes(jsonBytes) match {
+      case Right(js) =>
+        val brokerInfo = js.asJsonObject
+        val version = brokerInfo(VersionKey).to[Int]
+        val jmxPort = brokerInfo(JmxPortKey).to[Int]
+
+        val endpoints =
+          if (version < 1)
+            throw new KafkaException("Unsupported version of broker registration: " +
+              s"${new String(jsonBytes, UTF_8)}")
+          else if (version == 1) {
+            val host = brokerInfo(HostKey).to[String]
+            val port = brokerInfo(PortKey).to[Int]
+            val securityProtocol = SecurityProtocol.PLAINTEXT
+            val endPoint = new EndPoint(host, port, ListenerName.forSecurityProtocol(securityProtocol), securityProtocol)
+            Seq(endPoint)
+          }
+          else {
+            val securityProtocolMap = brokerInfo.get(ListenerSecurityProtocolMapKey).map(
+              _.to[Map[String, String]].map { case (listenerName, securityProtocol) =>
+                new ListenerName(listenerName) -> SecurityProtocol.forName(securityProtocol)
+              })
+            val listeners = brokerInfo(EndpointsKey).to[Seq[String]]
+            listeners.map(EndPoint.createEndPoint(_, securityProtocolMap))
+          }
+
+        val rack = brokerInfo.get(RackKey).flatMap(_.to[Option[String]])
+        val features = featuresAsJavaMap(brokerInfo)
+        BrokerInfo(
+          Broker(id, endpoints, rack, fromSupportedFeaturesMap(features)), version, jmxPort)
+      case Left(e) =>
+        throw new KafkaException(s"Failed to parse ZooKeeper registration for broker $id: " +
+          s"${new String(jsonBytes, UTF_8)}", e)
+    }
   }
 }
 
@@ -88,21 +279,63 @@ object TopicsZNode {
 }
 
 object TopicZNode {
+  case class TopicIdReplicaAssignment(topic: String,
+                                      topicId: Option[Uuid],
+                                      assignment: Map[TopicPartition, ReplicaAssignment])
   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 }
-    Json.encodeAsBytes(Map("version" -> 1, "partitions" -> assignmentJson))
+  def encode(topicId: Uuid,
+             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)
+    }
+
+    val topicAssignment = mutable.Map(
+      "version" -> 3,
+      "topic_id" -> topicId.toString,
+      "partitions" -> replicaAssignmentJson.asJava,
+      "adding_replicas" -> addingReplicasAssignmentJson.asJava,
+      "removing_replicas" -> removingReplicasAssignmentJson.asJava
+    )
+
+    Json.encodeAsBytes(topicAssignment.asJava)
   }
-  def decode(topic: String, bytes: Array[Byte]): Map[TopicPartition, Seq[Int]] = {
-    Json.parseBytes(bytes).flatMap { js =>
+  def decode(topic: String, bytes: Array[Byte]): TopicIdReplicaAssignment = {
+    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).map { js =>
       val assignmentJson = js.asJsonObject
+      val topicId = assignmentJson.get("topic_id").map(_.to[String]).map(Uuid.fromString)
+      val addingReplicasJsonOpt = assignmentJson.get("adding_replicas").map(_.asJsonObject)
+      val removingReplicasJsonOpt = assignmentJson.get("removing_replicas").map(_.asJsonObject)
       val partitionsJsonOpt = assignmentJson.get("partitions").map(_.asJsonObject)
-      partitionsJsonOpt.map { partitionsJson =>
+      val partitions = partitionsJsonOpt.map { partitionsJson =>
         partitionsJson.iterator.map { case (partition, replicas) =>
-          new TopicPartition(topic, partition.toInt) -> replicas.to[Seq[Int]]
-        }
-      }
-    }.map(_.toMap).getOrElse(Map.empty)
+          new TopicPartition(topic, partition.toInt) -> ReplicaAssignment(
+            replicas.to[Seq[Int]],
+            getReplicas(addingReplicasJsonOpt, partition),
+            getReplicas(removingReplicasJsonOpt, partition)
+          )
+        }.toMap
+      }.getOrElse(immutable.Map.empty[TopicPartition, ReplicaAssignment])
+
+      TopicIdReplicaAssignment(topic, topicId, partitions)
+    }.getOrElse(TopicIdReplicaAssignment(topic, None, Map.empty[TopicPartition, ReplicaAssignment]))
   }
 }
 
@@ -120,7 +353,7 @@ object TopicPartitionStateZNode {
     val leaderAndIsr = leaderIsrAndControllerEpoch.leaderAndIsr
     val controllerEpoch = leaderIsrAndControllerEpoch.controllerEpoch
     Json.encodeAsBytes(Map("version" -> 1, "leader" -> leaderAndIsr.leader, "leader_epoch" -> leaderAndIsr.leaderEpoch,
-      "controller_epoch" -> controllerEpoch, "isr" -> leaderAndIsr.isr))
+      "controller_epoch" -> controllerEpoch, "isr" -> leaderAndIsr.isr.asJava).asJava)
   }
   def decode(bytes: Array[Byte], stat: Stat): Option[LeaderIsrAndControllerEpoch] = {
     Json.parseBytes(bytes).map { js =>
@@ -142,13 +375,12 @@ object ConfigEntityTypeZNode {
 object ConfigEntityZNode {
   def path(entityType: String, entityName: String) = s"${ConfigEntityTypeZNode.path(entityType)}/$entityName"
   def encode(config: Properties): Array[Byte] = {
-    import scala.collection.JavaConverters._
-    Json.encodeAsBytes(Map("version" -> 1, "config" -> config.asScala))
+    Json.encodeAsBytes(Map("version" -> 1, "config" -> config).asJava)
   }
   def decode(bytes: Array[Byte]): Properties = {
     val props = new Properties()
     if (bytes != null) {
-      Json.parseBytes(bytes).map { js =>
+      Json.parseBytes(bytes).foreach { js =>
         val configOpt = js.asJsonObjectOption.flatMap(_.get("config").flatMap(_.asJsonObjectOption))
         configOpt.foreach(config => config.iterator.foreach { case (k, v) => props.setProperty(k, v.to[String]) })
       }
@@ -164,8 +396,8 @@ object ConfigEntityChangeNotificationZNode {
 object ConfigEntityChangeNotificationSequenceZNode {
   val SequenceNumberPrefix = "config_change_"
   def createPath = s"${ConfigEntityChangeNotificationZNode.path}/$SequenceNumberPrefix"
-  def encode(sanitizedEntityPath : String): Array[Byte] = Json.encodeAsBytes(Map("version" -> 2, "entity_path" -> sanitizedEntityPath))
-  def decode(bytes: Array[Byte]): String = new String(bytes, UTF_8)
+  def encode(sanitizedEntityPath: String): Array[Byte] = Json.encodeAsBytes(
+    Map("version" -> 2, "entity_path" -> sanitizedEntityPath).asJava)
 }
 
 object IsrChangeNotificationZNode {
@@ -174,10 +406,10 @@ object IsrChangeNotificationZNode {
 
 object IsrChangeNotificationSequenceZNode {
   val SequenceNumberPrefix = "isr_change_"
-  def path(sequenceNumber: String) = s"${IsrChangeNotificationZNode.path}/$SequenceNumberPrefix$sequenceNumber"
-  def encode(partitions: Set[TopicPartition]): Array[Byte] = {
-    val partitionsJson = partitions.map(partition => Map("topic" -> partition.topic, "partition" -> partition.partition))
-    Json.encodeAsBytes(Map("version" -> IsrChangeNotificationHandler.Version, "partitions" -> partitionsJson))
+  def path(sequenceNumber: String = "") = s"${IsrChangeNotificationZNode.path}/$SequenceNumberPrefix$sequenceNumber"
+  def encode(partitions: collection.Set[TopicPartition]): Array[Byte] = {
+    val partitionsJson = partitions.map(partition => Map("topic" -> partition.topic, "partition" -> partition.partition).asJava)
+    Json.encodeAsBytes(Map("version" -> IsrChangeNotificationHandler.Version, "partitions" -> partitionsJson.asJava).asJava)
   }
 
   def decode(bytes: Array[Byte]): Set[TopicPartition] = {
@@ -202,8 +434,9 @@ object LogDirEventNotificationSequenceZNode {
   val SequenceNumberPrefix = "log_dir_event_"
   val LogDirFailureEvent = 1
   def path(sequenceNumber: String) = s"${LogDirEventNotificationZNode.path}/$SequenceNumberPrefix$sequenceNumber"
-  def encode(brokerId: Int) =
-    Json.encodeAsBytes(Map("version" -> 1, "broker" -> brokerId, "event" -> LogDirFailureEvent))
+  def encode(brokerId: Int) = {
+    Json.encodeAsBytes(Map("version" -> 1, "broker" -> brokerId, "event" -> LogDirFailureEvent).asJava)
+  }
   def decode(bytes: Array[Byte]): Option[Int] = Json.parseBytes(bytes).map { js =>
     js.asJsonObject("broker").to[Int]
   }
@@ -222,35 +455,56 @@ object DeleteTopicsTopicZNode {
   def path(topic: String) = s"${DeleteTopicsZNode.path}/$topic"
 }
 
+/**
+ * The znode for initiating a partition reassignment.
+ * @deprecated Since 2.4, use the PartitionReassignment Kafka API instead.
+ */
 object ReassignPartitionsZNode {
+
+  /**
+    * The assignment of brokers for a `TopicPartition`.
+    *
+    * A replica assignment consists of a `topic`, `partition` and a list of `replicas`, which
+    * represent the broker ids that the `TopicPartition` is assigned to.
+    */
+  case class ReplicaAssignment(@BeanProperty @JsonProperty("topic") topic: String,
+                               @BeanProperty @JsonProperty("partition") partition: Int,
+                               @BeanProperty @JsonProperty("replicas") replicas: java.util.List[Int])
+
+  /**
+    * 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
+    */
+  @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(reassignment: collection.Map[TopicPartition, Seq[Int]]): Array[Byte] = {
-    val reassignmentJson = reassignment.map { case (tp, replicas) =>
-      Map("topic" -> tp.topic, "partition" -> tp.partition, "replicas" -> replicas)
-    }
-    Json.encodeAsBytes(Map("version" -> 1, "partitions" -> reassignmentJson))
-  }
-  def decode(bytes: Array[Byte]): Map[TopicPartition, Seq[Int]] = Json.parseBytes(bytes).flatMap { js =>
-    val reassignmentJson = js.asJsonObject
-    val partitionsJsonOpt = reassignmentJson.get("partitions")
-    partitionsJsonOpt.map { partitionsJson =>
-      partitionsJson.asJsonArray.iterator.map { partitionFieldsJs =>
-        val partitionFields = partitionFieldsJs.asJsonObject
-        val topic = partitionFields("topic").to[String]
-        val partition = partitionFields("partition").to[Int]
-        val replicas = partitionFields("replicas").to[Seq[Int]]
-        new TopicPartition(topic, partition) -> replicas
-      }
+
+  def encode(reassignmentMap: collection.Map[TopicPartition, Seq[Int]]): Array[Byte] = {
+    val reassignment = LegacyPartitionAssignment(1,
+      reassignmentMap.toSeq.map { case (tp, replicas) =>
+        ReplicaAssignment(tp.topic, tp.partition, replicas.asJava)
+      }.asJava
+    )
+    Json.encodeAsBytes(reassignment)
+  }
+
+  def decode(bytes: Array[Byte]): Either[JsonProcessingException, collection.Map[TopicPartition, Seq[Int]]] =
+    Json.parseBytesAs[LegacyPartitionAssignment](bytes).map { partitionAssignment =>
+      partitionAssignment.partitions.asScala.iterator.map { replicaAssignment =>
+        new TopicPartition(replicaAssignment.topic, replicaAssignment.partition) -> replicaAssignment.replicas.asScala
+      }.toMap
     }
-  }.map(_.toMap).getOrElse(Map.empty)
 }
 
 object PreferredReplicaElectionZNode {
   def path = s"${AdminZNode.path}/preferred_replica_election"
   def encode(partitions: Set[TopicPartition]): Array[Byte] = {
     val jsonMap = Map("version" -> 1,
-      "partitions" -> partitions.map(tp => Map("topic" -> tp.topic, "partition" -> tp.partition)))
-    Json.encodeAsBytes(jsonMap)
+      "partitions" -> partitions.map(tp => Map("topic" -> tp.topic, "partition" -> tp.partition).asJava).asJava)
+    Json.encodeAsBytes(jsonMap.asJava)
   }
   def decode(bytes: Array[Byte]): Set[TopicPartition] = Json.parseBytes(bytes).map { js =>
     val partitionsJson = js.asJsonObject("partitions").asJsonArray
@@ -263,14 +517,20 @@ object PreferredReplicaElectionZNode {
   }.map(_.toSet).getOrElse(Set.empty)
 }
 
+//old consumer path znode
+object ConsumerPathZNode {
+  def path = "/consumers"
+}
+
 object ConsumerOffset {
-  def path(group: String, topic: String, partition: Integer) = s"/consumers/${group}/offset/${topic}/${partition}"
+  def path(group: String, topic: String, partition: Integer) = s"${ConsumerPathZNode.path}/${group}/offsets/${topic}/${partition}"
   def encode(offset: Long): Array[Byte] = offset.toString.getBytes(UTF_8)
   def decode(bytes: Array[Byte]): Option[Long] = Option(bytes).map(new String(_, UTF_8).toLong)
 }
 
 object ZkVersion {
-  val NoVersion = -1
+  val MatchAnyVersion = -1 // if used in a conditional set, matches any version (the value should match ZooKeeper codebase)
+  val UnknownVersion = -2  // Version returned from get if node does not exist (internal constant for Kafka codebase, unused value in ZK)
 }
 
 object ZkStat {
@@ -283,38 +543,446 @@ object StateChangeHandlers {
 }
 
 /**
- * The root acl storage node. Under this node there will be one child node per resource type (Topic, Cluster, Group).
- * under each resourceType there will be a unique child for each resource instance and the data for that child will contain
- * list of its acls as a json object. Following gives an example:
- *
- * 
                      - * /kafka-acl/Topic/topic-1 => {"version": 1, "acls": [ { "host":"host1", "permissionType": "Allow","operation": "Read","principal": "User:alice"}]}
- * /kafka-acl/Cluster/kafka-cluster => {"version": 1, "acls": [ { "host":"host1", "permissionType": "Allow","operation": "Read","principal": "User:alice"}]}
- * /kafka-acl/Group/group-1 => {"version": 1, "acls": [ { "host":"host1", "permissionType": "Allow","operation": "Read","principal": "User:alice"}]}
- * 
                      
- */
-object AclZNode {
-  def path = "/kafka-acl"
+  * Acls for resources are stored in ZK under two root paths:
+  * 
                        
+  *   
                      • [[org.apache.kafka.common.resource.PatternType#LITERAL Literal]] patterns are stored under '/kafka-acl'.
+  *   The format is JSON. See [[kafka.zk.ResourceZNode]] for details.
                        • 
+  *   
                      • All other patterns are stored under '/kafka-acl-extended/pattern-type'.
+  *   The format is JSON. See [[kafka.zk.ResourceZNode]] for details.
                        • 
+  * 
                      
+  *
+  * Under each root node there will be one child node per resource type (Topic, Cluster, Group, etc).
+  * Under each resourceType there will be a unique child for each resource pattern and the data for that child will contain
+  * list of its acls as a json object. Following gives an example:
+  *
+  * 
                      +  * // Literal patterns:
+  * /kafka-acl/Topic/topic-1 => {"version": 1, "acls": [ { "host":"host1", "permissionType": "Allow","operation": "Read","principal": "User:alice"}]}
+  * /kafka-acl/Cluster/kafka-cluster => {"version": 1, "acls": [ { "host":"host1", "permissionType": "Allow","operation": "Read","principal": "User:alice"}]}
+  *
+  * // Prefixed patterns:
+  * /kafka-acl-extended/PREFIXED/Group/group-1 => {"version": 1, "acls": [ { "host":"host1", "permissionType": "Allow","operation": "Read","principal": "User:alice"}]}
+  * 
                      
+  *
+  * Acl change events are also stored under two paths:
+  * 
                        
+  *   
                      • [[org.apache.kafka.common.resource.PatternType#LITERAL Literal]] patterns are stored under '/kafka-acl-changes'.
+  *   The format is a UTF8 string in the form: <resource-type>:<resource-name>
                        • 
+  *   
                      • All other patterns are stored under '/kafka-acl-extended-changes'
+  *   The format is JSON, as defined by [[kafka.zk.ExtendedAclChangeEvent]]
                        • 
+  * 
                      
+  */
+sealed trait ZkAclStore {
+  val patternType: PatternType
+  val aclPath: String
+
+  def path(resourceType: ResourceType): String = s"$aclPath/${SecurityUtils.resourceTypeName(resourceType)}"
+
+  def path(resourceType: ResourceType, resourceName: String): String = s"$aclPath/${SecurityUtils.resourceTypeName(resourceType)}/$resourceName"
+
+  def changeStore: ZkAclChangeStore
+}
+
+object ZkAclStore {
+  private val storesByType: Map[PatternType, ZkAclStore] = PatternType.values
+    .filter(_.isSpecific)
+    .map(patternType => (patternType, create(patternType)))
+    .toMap
+
+  val stores: Iterable[ZkAclStore] = storesByType.values
+
+  val securePaths: Iterable[String] = stores
+    .flatMap(store => Set(store.aclPath, store.changeStore.aclChangePath))
+
+  def apply(patternType: PatternType): ZkAclStore = {
+    storesByType.get(patternType) match {
+      case Some(store) => store
+      case None => throw new KafkaException(s"Invalid pattern type: $patternType")
+    }
+  }
+
+  private def create(patternType: PatternType) = {
+    patternType match {
+      case PatternType.LITERAL => LiteralAclStore
+      case _ => new ExtendedAclStore(patternType)
+    }
+  }
+}
+
+object LiteralAclStore extends ZkAclStore {
+  val patternType: PatternType = PatternType.LITERAL
+  val aclPath: String = "/kafka-acl"
+
+  def changeStore: ZkAclChangeStore = LiteralAclChangeStore
+}
+
+class ExtendedAclStore(val patternType: PatternType) extends ZkAclStore {
+  if (patternType == PatternType.LITERAL)
+    throw new IllegalArgumentException("Literal pattern types are not supported")
+
+  val aclPath: String = s"${ExtendedAclZNode.path}/${patternType.name.toLowerCase}"
+
+  def changeStore: ZkAclChangeStore = ExtendedAclChangeStore
+}
+
+object ExtendedAclZNode {
+  def path = "/kafka-acl-extended"
+}
+
+trait AclChangeNotificationHandler {
+  def processNotification(resource: ResourcePattern): Unit
 }
 
-object ResourceTypeZNode {
-  def path(resourceType: String) = s"${AclZNode.path}/$resourceType"
+trait AclChangeSubscription extends AutoCloseable {
+  def close(): Unit
+}
+
+case class AclChangeNode(path: String, bytes: Array[Byte])
+
+sealed trait ZkAclChangeStore {
+  val aclChangePath: String
+  def createPath: String = s"$aclChangePath/${ZkAclChangeStore.SequenceNumberPrefix}"
+
+  def decode(bytes: Array[Byte]): ResourcePattern
+
+  protected def encode(resource: ResourcePattern): Array[Byte]
+
+  def createChangeNode(resource: ResourcePattern): AclChangeNode = AclChangeNode(createPath, encode(resource))
+
+  def createListener(handler: AclChangeNotificationHandler, zkClient: KafkaZkClient): AclChangeSubscription = {
+    val rawHandler: NotificationHandler = (bytes: Array[Byte]) => handler.processNotification(decode(bytes))
+
+    val aclChangeListener = new ZkNodeChangeNotificationListener(
+      zkClient, aclChangePath, ZkAclChangeStore.SequenceNumberPrefix, rawHandler)
+
+    aclChangeListener.init()
+
+    () => aclChangeListener.close()
+  }
+}
+
+object ZkAclChangeStore {
+  val stores: Iterable[ZkAclChangeStore] = List(LiteralAclChangeStore, ExtendedAclChangeStore)
+
+  def SequenceNumberPrefix = "acl_changes_"
+}
+
+case object LiteralAclChangeStore extends ZkAclChangeStore {
+  val name = "LiteralAclChangeStore"
+  val aclChangePath: String = "/kafka-acl-changes"
+
+  def encode(resource: ResourcePattern): Array[Byte] = {
+    if (resource.patternType != PatternType.LITERAL)
+      throw new IllegalArgumentException("Only literal resource patterns can be encoded")
+
+    val legacyName = resource.resourceType.toString + AclEntry.ResourceSeparator + resource.name
+    legacyName.getBytes(UTF_8)
+  }
+
+  def decode(bytes: Array[Byte]): ResourcePattern = {
+    val string = new String(bytes, UTF_8)
+    string.split(AclEntry.ResourceSeparator, 2) match {
+        case Array(resourceType, resourceName, _*) => new ResourcePattern(ResourceType.fromString(resourceType), resourceName, PatternType.LITERAL)
+        case _ => throw new IllegalArgumentException("expected a string in format ResourceType:ResourceName but got " + string)
+      }
+  }
+}
+
+case object ExtendedAclChangeStore extends ZkAclChangeStore {
+  val name = "ExtendedAclChangeStore"
+  val aclChangePath: String = "/kafka-acl-extended-changes"
+
+  def encode(resource: ResourcePattern): Array[Byte] = {
+    if (resource.patternType == PatternType.LITERAL)
+      throw new IllegalArgumentException("Literal pattern types are not supported")
+
+    Json.encodeAsBytes(ExtendedAclChangeEvent(
+      ExtendedAclChangeEvent.currentVersion,
+      resource.resourceType.name,
+      resource.name,
+      resource.patternType.name))
+  }
+
+  def decode(bytes: Array[Byte]): ResourcePattern = {
+    val changeEvent = Json.parseBytesAs[ExtendedAclChangeEvent](bytes) match {
+      case Right(event) => event
+      case Left(e) => throw new IllegalArgumentException("Failed to parse ACL change event", e)
+    }
+
+    changeEvent.toResource match {
+      case Success(r) => r
+      case Failure(e) => throw new IllegalArgumentException("Failed to convert ACL change event to resource", e)
+    }
+  }
 }
 
 object ResourceZNode {
-  def path(resource: Resource) = s"${AclZNode.path}/${resource.resourceType}/${resource.name}"
-  def encode(acls: Set[Acl]): Array[Byte] = Json.encodeAsBytes(Acl.toJsonCompatibleMap(acls))
-  def decode(bytes: Array[Byte], stat: Stat): VersionedAcls = VersionedAcls(Acl.fromBytes(bytes), stat.getVersion)
+  def path(resource: ResourcePattern): String = ZkAclStore(resource.patternType).path(resource.resourceType, resource.name)
+
+  def encode(acls: Set[AclEntry]): Array[Byte] = Json.encodeAsBytes(AclEntry.toJsonCompatibleMap(acls).asJava)
+  def decode(bytes: Array[Byte], stat: Stat): VersionedAcls = VersionedAcls(AclEntry.fromBytes(bytes), stat.getVersion)
+}
+
+object ExtendedAclChangeEvent {
+  val currentVersion: Int = 1
+}
+
+case class ExtendedAclChangeEvent(@BeanProperty @JsonProperty("version") version: Int,
+                                  @BeanProperty @JsonProperty("resourceType") resourceType: String,
+                                  @BeanProperty @JsonProperty("name") name: String,
+                                  @BeanProperty @JsonProperty("patternType") patternType: String) {
+  if (version > ExtendedAclChangeEvent.currentVersion)
+    throw new UnsupportedVersionException(s"Acl change event received for unsupported version: $version")
+
+  def toResource: Try[ResourcePattern] = {
+    for {
+      resType <- Try(ResourceType.fromString(resourceType))
+      patType <- Try(PatternType.fromString(patternType))
+      resource = new ResourcePattern(resType, name, patType)
+    } yield resource
+  }
 }
 
-object AclChangeNotificationZNode {
-  def path = "/kafka-acl-changes"
+object ClusterZNode {
+  def path = "/cluster"
 }
 
-object AclChangeNotificationSequenceZNode {
-  val SequenceNumberPrefix = "acl_changes_"
-  def createPath = s"${AclChangeNotificationZNode.path}/$SequenceNumberPrefix"
-  def deletePath(sequenceNode: String) = s"${AclChangeNotificationZNode.path}/${sequenceNode}"
-  def encode(resourceName : String): Array[Byte] = resourceName.getBytes(UTF_8)
+object ClusterIdZNode {
+  def path = s"${ClusterZNode.path}/id"
+
+  def toJson(id: String): Array[Byte] = {
+    Json.encodeAsBytes(Map("version" -> "1", "id" -> id).asJava)
+  }
+
+  def fromJson(clusterIdJson:  Array[Byte]): String = {
+    Json.parseBytes(clusterIdJson).map(_.asJsonObject("id").to[String]).getOrElse {
+      throw new KafkaException(s"Failed to parse the cluster id json $clusterIdJson")
+    }
+  }
+}
+
+object BrokerSequenceIdZNode {
+  def path = s"${BrokersZNode.path}/seqid"
+}
+
+object ProducerIdBlockZNode {
+  def path = "/latest_producer_id_block"
+}
+
+object DelegationTokenAuthZNode {
+  def path = "/delegation_token"
+}
+
+object DelegationTokenChangeNotificationZNode {
+  def path =  s"${DelegationTokenAuthZNode.path}/token_changes"
+}
+
+object DelegationTokenChangeNotificationSequenceZNode {
+  val SequenceNumberPrefix = "token_change_"
+  def createPath = s"${DelegationTokenChangeNotificationZNode.path}/$SequenceNumberPrefix"
+  def deletePath(sequenceNode: String) = s"${DelegationTokenChangeNotificationZNode.path}/${sequenceNode}"
+  def encode(tokenId : String): Array[Byte] = tokenId.getBytes(UTF_8)
   def decode(bytes: Array[Byte]): String = new String(bytes, UTF_8)
-}
\ No newline at end of file
+}
+
+object DelegationTokensZNode {
+  def path = s"${DelegationTokenAuthZNode.path}/tokens"
+}
+
+object DelegationTokenInfoZNode {
+  def path(tokenId: String) =  s"${DelegationTokensZNode.path}/$tokenId"
+  def encode(token: DelegationToken): Array[Byte] =  Json.encodeAsBytes(DelegationTokenManager.toJsonCompatibleMap(token).asJava)
+  def decode(bytes: Array[Byte]): Option[TokenInformation] = DelegationTokenManager.fromBytes(bytes)
+}
+
+/**
+ * Represents the status of the FeatureZNode.
+ *
+ * Enabled  -> This status means the feature versioning system (KIP-584) is enabled, and, the
+ *             finalized features stored in the FeatureZNode are active. This status is written by
+ *             the controller to the FeatureZNode only when the broker IBP config is greater than
+ *             or equal to KAFKA_2_7_IV0.
+ *
+ * Disabled -> This status means the feature versioning system (KIP-584) is disabled, and, the
+ *             the finalized features stored in the FeatureZNode is not relevant. This status is
+ *             written by the controller to the FeatureZNode only when the broker IBP config
+ *             is less than KAFKA_2_7_IV0.
+ */
+sealed trait FeatureZNodeStatus {
+  def id: Int
+}
+
+object FeatureZNodeStatus {
+  case object Disabled extends FeatureZNodeStatus {
+    val id: Int = 0
+  }
+
+  case object Enabled extends FeatureZNodeStatus {
+    val id: Int = 1
+  }
+
+  def withNameOpt(id: Int): Option[FeatureZNodeStatus] = {
+    id match {
+      case Disabled.id => Some(Disabled)
+      case Enabled.id => Some(Enabled)
+      case _ => Option.empty
+    }
+  }
+}
+
+/**
+ * Represents the contents of the ZK node containing finalized feature information.
+ *
+ * @param status     the status of the ZK node
+ * @param features   the cluster-wide finalized features
+ */
+case class FeatureZNode(status: FeatureZNodeStatus, features: Features[FinalizedVersionRange]) {
+}
+
+object FeatureZNode {
+  private val VersionKey = "version"
+  private val StatusKey = "status"
+  private val FeaturesKey = "features"
+
+  // V1 contains 'version', 'status' and 'features' keys.
+  val V1 = 1
+  val CurrentVersion = V1
+
+  def path = "/feature"
+
+  def asJavaMap(scalaMap: Map[String, Map[String, Short]]): util.Map[String, util.Map[String, java.lang.Short]] = {
+    scalaMap
+      .map {
+        case(featureName, versionInfo) => featureName -> versionInfo.map {
+          case(label, version) => label -> java.lang.Short.valueOf(version)
+        }.asJava
+      }.asJava
+  }
+
+  /**
+   * Encodes a FeatureZNode to JSON.
+   *
+   * @param featureZNode   FeatureZNode to be encoded
+   *
+   * @return               JSON representation of the FeatureZNode, as an Array[Byte]
+   */
+  def encode(featureZNode: FeatureZNode): Array[Byte] = {
+    val jsonMap = collection.mutable.Map(
+      VersionKey -> CurrentVersion,
+      StatusKey -> featureZNode.status.id,
+      FeaturesKey -> featureZNode.features.toMap)
+    Json.encodeAsBytes(jsonMap.asJava)
+  }
+
+  /**
+   * Decodes the contents of the feature ZK node from Array[Byte] to a FeatureZNode.
+   *
+   * @param jsonBytes   the contents of the feature ZK node
+   *
+   * @return            the FeatureZNode created from jsonBytes
+   *
+   * @throws IllegalArgumentException   if the Array[Byte] can not be decoded.
+   */
+  def decode(jsonBytes: Array[Byte]): FeatureZNode = {
+    Json.tryParseBytes(jsonBytes) match {
+      case Right(js) =>
+        val featureInfo = js.asJsonObject
+        val version = featureInfo(VersionKey).to[Int]
+        if (version < V1) {
+          throw new IllegalArgumentException(s"Unsupported version: $version of feature information: " +
+            s"${new String(jsonBytes, UTF_8)}")
+        }
+
+        val featuresMap = featureInfo
+          .get(FeaturesKey)
+          .flatMap(_.to[Option[Map[String, Map[String, Int]]]])
+
+        if (featuresMap.isEmpty) {
+          throw new IllegalArgumentException("Features map can not be absent in: " +
+            s"${new String(jsonBytes, UTF_8)}")
+        }
+        val features = asJavaMap(
+          featuresMap
+            .map(theMap => theMap.map {
+              case (featureName, versionInfo) => featureName -> versionInfo.map {
+                case (label, version) => label -> version.asInstanceOf[Short]
+              }
+            }).getOrElse(Map[String, Map[String, Short]]()))
+
+        val statusInt = featureInfo
+          .get(StatusKey)
+          .flatMap(_.to[Option[Int]])
+        if (statusInt.isEmpty) {
+          throw new IllegalArgumentException("Status can not be absent in feature information: " +
+            s"${new String(jsonBytes, UTF_8)}")
+        }
+        val status = FeatureZNodeStatus.withNameOpt(statusInt.get)
+        if (status.isEmpty) {
+          throw new IllegalArgumentException(
+            s"Malformed status: $statusInt found in feature information: ${new String(jsonBytes, UTF_8)}")
+        }
+
+        var finalizedFeatures: Features[FinalizedVersionRange] = null
+        try {
+          finalizedFeatures = fromFinalizedFeaturesMap(features)
+        } catch {
+          case e: Exception => throw new IllegalArgumentException(
+            "Unable to convert to finalized features from map: " + features, e)
+        }
+        FeatureZNode(status.get, finalizedFeatures)
+      case Left(e) =>
+        throw new IllegalArgumentException(s"Failed to parse feature information: " +
+          s"${new String(jsonBytes, UTF_8)}", e)
+    }
+  }
+}
+
+object ZkData {
+
+  // Important: it is necessary to add any new top level Zookeeper path to the Seq
+  val SecureRootPaths = Seq(AdminZNode.path,
+    BrokersZNode.path,
+    ClusterZNode.path,
+    ConfigZNode.path,
+    ControllerZNode.path,
+    ControllerEpochZNode.path,
+    IsrChangeNotificationZNode.path,
+    ProducerIdBlockZNode.path,
+    LogDirEventNotificationZNode.path,
+    DelegationTokenAuthZNode.path,
+    ExtendedAclZNode.path) ++ ZkAclStore.securePaths
+
+  // These are persistent ZK paths that should exist on kafka broker startup.
+  val PersistentZkPaths = Seq(
+    ConsumerPathZNode.path, // old consumer path
+    BrokerIdsZNode.path,
+    TopicsZNode.path,
+    ConfigEntityChangeNotificationZNode.path,
+    DeleteTopicsZNode.path,
+    BrokerSequenceIdZNode.path,
+    IsrChangeNotificationZNode.path,
+    ProducerIdBlockZNode.path,
+    LogDirEventNotificationZNode.path
+  ) ++ ConfigType.all.map(ConfigEntityTypeZNode.path)
+
+  val SensitiveRootPaths = Seq(
+    ConfigEntityTypeZNode.path(ConfigType.User),
+    ConfigEntityTypeZNode.path(ConfigType.Broker),
+    DelegationTokensZNode.path
+  )
+
+  def sensitivePath(path: String): Boolean = {
+    path != null && SensitiveRootPaths.exists(path.startsWith)
+  }
+
+  def defaultAcls(isSecure: Boolean, path: String): Seq[ACL] = {
+    //Old Consumer path is kept open as different consumers will write under this node.
+    if (!ConsumerPathZNode.path.equals(path) && isSecure) {
+      val acls = new ArrayBuffer[ACL]
+      acls ++= ZooDefs.Ids.CREATOR_ALL_ACL.asScala
+      if (!sensitivePath(path))
+        acls ++= ZooDefs.Ids.READ_ACL_UNSAFE.asScala
+      acls
+    } else ZooDefs.Ids.OPEN_ACL_UNSAFE.asScala
+  }
+}
diff --git a/core/src/main/scala/kafka/zk/ZkSecurityMigratorUtils.scala b/core/src/main/scala/kafka/zk/ZkSecurityMigratorUtils.scala
new file mode 100644
index 0000000000000..31a7ba2907379
--- /dev/null
+++ b/core/src/main/scala/kafka/zk/ZkSecurityMigratorUtils.scala
@@ -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 kafka.zk
+
+import org.apache.zookeeper.ZooKeeper
+
+/**
+ * This class should only be used in ZkSecurityMigrator tool.
+ * This class will be removed after we migrate ZkSecurityMigrator away from ZK's asynchronous API.
+ * @param kafkaZkClient
+ */
+class ZkSecurityMigratorUtils(val kafkaZkClient: KafkaZkClient) {
+
+  def currentZooKeeper: ZooKeeper = kafkaZkClient.currentZooKeeper
+
+}
diff --git a/core/src/main/scala/kafka/zookeeper/ZooKeeperClient.scala b/core/src/main/scala/kafka/zookeeper/ZooKeeperClient.scala
old mode 100644
new mode 100755
index ef170597f94e3..96ef6d5179368
--- a/core/src/main/scala/kafka/zookeeper/ZooKeeperClient.scala
+++ b/core/src/main/scala/kafka/zookeeper/ZooKeeperClient.scala
@@ -17,19 +17,32 @@
 
 package kafka.zookeeper
 
+import java.util.Locale
 import java.util.concurrent.locks.{ReentrantLock, ReentrantReadWriteLock}
-import java.util.concurrent.{ArrayBlockingQueue, ConcurrentHashMap, CountDownLatch, Semaphore, TimeUnit}
+import java.util.concurrent._
+import java.util.{List => JList}
 
+import com.yammer.metrics.core.MetricName
+import kafka.metrics.KafkaMetricsGroup
 import kafka.utils.CoreUtils.{inLock, inReadLock, inWriteLock}
-import kafka.utils.Logging
-import org.apache.zookeeper.AsyncCallback.{ACLCallback, Children2Callback, DataCallback, StatCallback, StringCallback, VoidCallback}
+import kafka.utils.{KafkaScheduler, Logging}
+import kafka.zookeeper.ZooKeeperClient._
+import org.apache.kafka.common.utils.Time
+import org.apache.zookeeper.AsyncCallback.{Children2Callback, DataCallback, StatCallback}
 import org.apache.zookeeper.KeeperException.Code
 import org.apache.zookeeper.Watcher.Event.{EventType, KeeperState}
 import org.apache.zookeeper.ZooKeeper.States
 import org.apache.zookeeper.data.{ACL, Stat}
-import org.apache.zookeeper.{CreateMode, KeeperException, WatchedEvent, Watcher, ZooKeeper}
+import org.apache.zookeeper._
+import org.apache.zookeeper.client.ZKClientConfig
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
+import scala.collection.Seq
+import scala.collection.mutable.Set
+
+object ZooKeeperClient {
+  val RetryBackoffMs = 1000
+}
 
 /**
  * A ZooKeeper client that encourages pipelined requests.
@@ -38,12 +51,34 @@ import scala.collection.JavaConverters._
  * @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
+ * @param zkClientConfig ZooKeeper client configuration, for TLS configs if desired
  */
 class ZooKeeperClient(connectString: String,
                       sessionTimeoutMs: Int,
                       connectionTimeoutMs: Int,
-                      maxInFlightRequests: Int) extends Logging {
-  this.logIdent = "[ZooKeeperClient] "
+                      maxInFlightRequests: Int,
+                      time: Time,
+                      metricGroup: String,
+                      metricType: String,
+                      name: Option[String],
+                      zkClientConfig: Option[ZKClientConfig]) 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,
+      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()
@@ -51,10 +86,57 @@ class ZooKeeperClient(connectString: String,
   private val zNodeChildChangeHandlers = new ConcurrentHashMap[String, ZNodeChildChangeHandler]().asScala
   private val inFlightRequests = new Semaphore(maxInFlightRequests)
   private val stateChangeHandlers = new ConcurrentHashMap[String, StateChangeHandler]().asScala
+  private[zookeeper] val reinitializeScheduler = new KafkaScheduler(threads = 1, s"zk-client-${threadPrefix}reinit-")
+  private var isFirstConnectionEstablished = false
+
+  private val metricNames = Set[String]()
+
+  // The state map has to be created before creating ZooKeeper since it's needed in the ZooKeeper callback.
+  private val stateToMeterMap = {
+    import KeeperState._
+    val stateToEventTypeMap = Map(
+      Disconnected -> "Disconnects",
+      SyncConnected -> "SyncConnects",
+      AuthFailed -> "AuthFailures",
+      ConnectedReadOnly -> "ReadOnlyConnects",
+      SaslAuthenticated -> "SaslAuthentications",
+      Expired -> "Expires"
+    )
+    stateToEventTypeMap.map { case (state, eventType) =>
+      val name = s"ZooKeeper${eventType}PerSec"
+      metricNames += name
+      state -> newMeter(name, eventType.toLowerCase(Locale.ROOT), TimeUnit.SECONDS)
+    }
+  }
+
+  private val clientConfig = zkClientConfig getOrElse new ZKClientConfig()
 
   info(s"Initializing a new session to $connectString.")
-  @volatile private var zooKeeper = new ZooKeeper(connectString, sessionTimeoutMs, ZooKeeperClientWatcher)
-  waitUntilConnected(connectionTimeoutMs, TimeUnit.MILLISECONDS)
+  // Fail-fast if there's an error during construction (so don't call initialize, which retries forever)
+  @volatile private var zooKeeper = new ZooKeeper(connectString, sessionTimeoutMs, ZooKeeperClientWatcher,
+    clientConfig)
+  private[zookeeper] def getClientConfig = clientConfig
+
+  newGauge("SessionState", () => connectionState.toString)
+
+  metricNames += "SessionState"
+
+  reinitializeScheduler.startup()
+  try waitUntilConnected(connectionTimeoutMs, TimeUnit.MILLISECONDS)
+  catch {
+    case e: Throwable =>
+      close()
+      throw e
+  }
+
+  override def metricName(name: String, metricTags: scala.collection.Map[String, String]): MetricName = {
+    explicitMetricName(metricGroup, metricType, name, metricTags)
+  }
+
+  /**
+   * Return the state of the ZooKeeper connection.
+   */
+  def connectionState: States = zooKeeper.getState
 
   /**
    * Send a request and wait for its response. See handle(Seq[AsyncRequest]) for details.
@@ -77,7 +159,7 @@ class ZooKeeperClient(connectString: String,
    * response type (e.g. Seq[CreateRequest] -> Seq[CreateResponse]). Otherwise, the most specific common supertype
    * will be used (e.g. Seq[AsyncRequest] -> Seq[AsyncResponse]).
    */
-  def handleRequests[Req <: AsyncRequest](requests: Seq[Req]): Seq[Req#Response] = inReadLock(initializationLock) {
+  def handleRequests[Req <: AsyncRequest](requests: Seq[Req]): Seq[Req#Response] = {
     if (requests.isEmpty)
       Seq.empty
     else {
@@ -87,10 +169,12 @@ class ZooKeeperClient(connectString: String,
       requests.foreach { request =>
         inFlightRequests.acquire()
         try {
-          send(request) { response =>
-            responseQueue.add(response)
-            inFlightRequests.release()
-            countDownLatch.countDown()
+          inReadLock(initializationLock) {
+            send(request) { response =>
+              responseQueue.add(response)
+              inFlightRequests.release()
+              countDownLatch.countDown()
+            }
           }
         } catch {
           case e: Throwable =>
@@ -103,53 +187,68 @@ class ZooKeeperClient(connectString: String,
     }
   }
 
-  private def send[Req <: AsyncRequest](request: Req)(processResponse: Req#Response => Unit): Unit = {
+  // Visibility to override for testing
+  private[zookeeper] def send[Req <: AsyncRequest](request: Req)(processResponse: Req#Response => Unit): Unit = {
     // Safe to cast as we always create a response of the right type
     def callback(response: AsyncResponse): Unit = processResponse(response.asInstanceOf[Req#Response])
 
-    request match {
+    def responseMetadata(sendTimeMs: Long) = new ResponseMetadata(sendTimeMs, receivedTimeMs = time.hiResClockMs())
+
+    val sendTimeMs = time.hiResClockMs()
+
+    // Cast to AsyncRequest to workaround a scalac bug that results in an false exhaustiveness warning
+    // with -Xlint:strict-unsealed-patmat
+    (request: AsyncRequest) match {
       case ExistsRequest(path, ctx) =>
         zooKeeper.exists(path, shouldWatch(request), new StatCallback {
-          override def processResult(rc: Int, path: String, ctx: Any, stat: Stat): Unit =
-            callback(ExistsResponse(Code.get(rc), path, Option(ctx), stat))
+          def processResult(rc: Int, path: String, ctx: Any, stat: Stat): Unit =
+            callback(ExistsResponse(Code.get(rc), path, Option(ctx), stat, responseMetadata(sendTimeMs)))
         }, ctx.orNull)
       case GetDataRequest(path, ctx) =>
         zooKeeper.getData(path, shouldWatch(request), new DataCallback {
-          override def processResult(rc: Int, path: String, ctx: Any, data: Array[Byte], stat: Stat): Unit =
-            callback(GetDataResponse(Code.get(rc), path, Option(ctx), data, stat))
+          def processResult(rc: Int, path: String, ctx: Any, data: Array[Byte], stat: Stat): Unit =
+            callback(GetDataResponse(Code.get(rc), path, Option(ctx), data, stat, responseMetadata(sendTimeMs)))
         }, ctx.orNull)
-      case GetChildrenRequest(path, ctx) =>
+      case GetChildrenRequest(path, _, ctx) =>
         zooKeeper.getChildren(path, shouldWatch(request), new Children2Callback {
-          override def processResult(rc: Int, path: String, ctx: Any, children: java.util.List[String], stat: Stat): Unit =
-            callback(GetChildrenResponse(Code.get(rc), path, Option(ctx),
-              Option(children).map(_.asScala).getOrElse(Seq.empty), stat))
+          def processResult(rc: Int, path: String, ctx: Any, children: JList[String], stat: Stat): Unit =
+            callback(GetChildrenResponse(Code.get(rc), path, Option(ctx), Option(children).map(_.asScala).getOrElse(Seq.empty),
+              stat, responseMetadata(sendTimeMs)))
         }, ctx.orNull)
       case CreateRequest(path, data, acl, createMode, ctx) =>
-        zooKeeper.create(path, data, acl.asJava, createMode, new StringCallback {
-          override def processResult(rc: Int, path: String, ctx: Any, name: String): Unit =
-            callback(CreateResponse(Code.get(rc), path, Option(ctx), name))
-        }, ctx.orNull)
+        zooKeeper.create(path, data, acl.asJava, createMode,
+          (rc, path, ctx, name) =>
+            callback(CreateResponse(Code.get(rc), path, Option(ctx), name, responseMetadata(sendTimeMs))),
+          ctx.orNull)
       case SetDataRequest(path, data, version, ctx) =>
-        zooKeeper.setData(path, data, version, new StatCallback {
-          override def processResult(rc: Int, path: String, ctx: Any, stat: Stat): Unit =
-            callback(SetDataResponse(Code.get(rc), path, Option(ctx), stat))
-        }, ctx.orNull)
+        zooKeeper.setData(path, data, version,
+          (rc, path, ctx, stat) =>
+            callback(SetDataResponse(Code.get(rc), path, Option(ctx), stat, responseMetadata(sendTimeMs))),
+          ctx.orNull)
       case DeleteRequest(path, version, ctx) =>
-        zooKeeper.delete(path, version, new VoidCallback {
-          override def processResult(rc: Int, path: String, ctx: Any): Unit =
-            callback(DeleteResponse(Code.get(rc), path, Option(ctx)))
-        }, ctx.orNull)
+        zooKeeper.delete(path, version,
+          (rc, path, ctx) => callback(DeleteResponse(Code.get(rc), path, Option(ctx), responseMetadata(sendTimeMs))),
+          ctx.orNull)
       case GetAclRequest(path, ctx) =>
-        zooKeeper.getACL(path, null, new ACLCallback {
-          override def processResult(rc: Int, path: String, ctx: Any, acl: java.util.List[ACL], stat: Stat): Unit = {
+        zooKeeper.getACL(path, null,
+          (rc, path, ctx, acl, stat) =>
             callback(GetAclResponse(Code.get(rc), path, Option(ctx), Option(acl).map(_.asScala).getOrElse(Seq.empty),
-              stat))
-        }}, ctx.orNull)
+              stat, responseMetadata(sendTimeMs))),
+          ctx.orNull)
       case SetAclRequest(path, acl, version, ctx) =>
-        zooKeeper.setACL(path, acl.asJava, version, new StatCallback {
-          override def processResult(rc: Int, path: String, ctx: Any, stat: Stat): Unit =
-            callback(SetAclResponse(Code.get(rc), path, Option(ctx), stat))
-        }, ctx.orNull)
+        zooKeeper.setACL(path, acl.asJava, version,
+          (rc, path, ctx, stat) =>
+            callback(SetAclResponse(Code.get(rc), path, Option(ctx), stat, responseMetadata(sendTimeMs))),
+          ctx.orNull)
+      case MultiRequest(zkOps, ctx) =>
+        def toZkOpResult(opResults: JList[OpResult]): Seq[ZkOpResult] =
+          Option(opResults).map(results => zkOps.zip(results.asScala).map { case (zkOp, result) =>
+            ZkOpResult(zkOp, result)
+          }).orNull
+        zooKeeper.multi(zkOps.map(_.toZookeeperOp).asJava,
+          (rc, path, ctx, opResults) =>
+            callback(MultiResponse(Code.get(rc), path, Option(ctx), toZkOpResult(opResults), responseMetadata(sendTimeMs))),
+          ctx.orNull)
     }
   }
 
@@ -166,19 +265,20 @@ class ZooKeeperClient(connectString: String,
     info("Waiting until connected.")
     var nanos = timeUnit.toNanos(timeout)
     inLock(isConnectedOrExpiredLock) {
-      var state = zooKeeper.getState
+      var state = connectionState
       while (!state.isConnected && state.isAlive) {
         if (nanos <= 0) {
           throw new ZooKeeperClientTimeoutException(s"Timed out waiting for connection while in state: $state")
         }
         nanos = isConnectedOrExpiredCondition.awaitNanos(nanos)
-        state = zooKeeper.getState
+        state = connectionState
       }
       if (state == States.AUTH_FAILED) {
         throw new ZooKeeperClientAuthFailedException("Auth failed either before or while waiting for connection")
       } else if (state == States.CLOSED) {
         throw new ZooKeeperClientExpiredException("Session expired either before or while waiting for connection")
       }
+      isFirstConnectionEstablished = true
     }
     info("Connected.")
   }
@@ -186,7 +286,7 @@ class ZooKeeperClient(connectString: String,
   // If this method is changed, the documentation for registerZNodeChangeHandler and/or registerZNodeChildChangeHandler
   // may need to be updated.
   private def shouldWatch(request: AsyncRequest): Boolean = request match {
-    case _: GetChildrenRequest => zNodeChildChangeHandlers.contains(request.path)
+    case GetChildrenRequest(_, registerWatch, _) => registerWatch && zNodeChildChangeHandlers.contains(request.path)
     case _: ExistsRequest | _: GetDataRequest => zNodeChangeHandlers.contains(request.path)
     case _ => throw new IllegalArgumentException(s"Request $request is not watchable")
   }
@@ -248,12 +348,21 @@ class ZooKeeperClient(connectString: String,
     stateChangeHandlers.remove(name)
   }
 
-  def close(): Unit = inWriteLock(initializationLock) {
+  def close(): Unit = {
     info("Closing.")
-    zNodeChangeHandlers.clear()
-    zNodeChildChangeHandlers.clear()
-    stateChangeHandlers.clear()
-    zooKeeper.close()
+
+    // 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().
+    reinitializeScheduler.shutdown()
+
+    inWriteLock(initializationLock) {
+      zNodeChangeHandlers.clear()
+      zNodeChildChangeHandlers.clear()
+      stateChangeHandlers.clear()
+      zooKeeper.close()
+      metricNames.foreach(removeMetric(_))
+    }
     info("Closed.")
   }
 
@@ -261,57 +370,97 @@ class ZooKeeperClient(connectString: String,
     zooKeeper.getSessionId
   }
 
-  private def initialize(): Unit = {
-    if (!zooKeeper.getState.isAlive) {
-      info(s"Initializing a new session to $connectString.")
-      var now = System.currentTimeMillis()
-      val threshold = now + connectionTimeoutMs
-      while (now < threshold) {
-        try {
-          zooKeeper.close()
-          zooKeeper = new ZooKeeper(connectString, sessionTimeoutMs, ZooKeeperClientWatcher)
-          waitUntilConnected(threshold - now, TimeUnit.MILLISECONDS)
-          return
-        } catch {
-          case _: Exception =>
-            now = System.currentTimeMillis()
-            if (now < threshold) {
-              Thread.sleep(1000)
-              now = System.currentTimeMillis()
-            }
+  // Only for testing
+  private[kafka] def currentZooKeeper: ZooKeeper = inReadLock(initializationLock) {
+    zooKeeper
+  }
+
+  private def reinitialize(): Unit = {
+    // Initialization callbacks are invoked outside of the lock to avoid deadlock potential since their completion
+    // may require additional Zookeeper requests, which will block to acquire the initialization lock
+    stateChangeHandlers.values.foreach(callBeforeInitializingSession _)
+
+    inWriteLock(initializationLock) {
+      if (!connectionState.isAlive) {
+        zooKeeper.close()
+        info(s"Initializing a new session to $connectString.")
+        // retry forever until ZooKeeper can be instantiated
+        var connected = false
+        while (!connected) {
+          try {
+            zooKeeper = new ZooKeeper(connectString, sessionTimeoutMs, ZooKeeperClientWatcher, clientConfig)
+            connected = true
+          } catch {
+            case e: Exception =>
+              info("Error when recreating ZooKeeper, retrying after a short sleep", e)
+              Thread.sleep(RetryBackoffMs)
+          }
         }
       }
-      info(s"Timed out waiting for connection during session initialization while in state: ${zooKeeper.getState}")
-      stateChangeHandlers.foreach {case (name, handler) => handler.onReconnectionTimeout()}
     }
+
+    stateChangeHandlers.values.foreach(callAfterInitializingSession _)
   }
 
   /**
-   * reinitialize method to use in unit tests
+   * Close the zookeeper client to force session reinitialization. This is visible for testing only.
    */
-  private[zookeeper] def reinitialize(): Unit = {
+  private[zookeeper] def forceReinitialize(): Unit = {
     zooKeeper.close()
-    initialize()
+    reinitialize()
+  }
+
+  private def callBeforeInitializingSession(handler: StateChangeHandler): Unit = {
+    try {
+      handler.beforeInitializingSession()
+    } catch {
+      case t: Throwable =>
+        error(s"Uncaught error in handler ${handler.name}", t)
+    }
+  }
+
+  private def callAfterInitializingSession(handler: StateChangeHandler): Unit = {
+    try {
+      handler.afterInitializingSession()
+    } catch {
+      case t: Throwable =>
+        error(s"Uncaught error in handler ${handler.name}", t)
+    }
+  }
+
+  // Visibility for testing
+  private[zookeeper] def scheduleReinitialize(name: String, message: String, delayMs: Long): Unit = {
+    reinitializeScheduler.schedule(name, () => {
+      info(message)
+      reinitialize()
+    }, delayMs, period = -1L, unit = TimeUnit.MILLISECONDS)
   }
 
-  private object ZooKeeperClientWatcher extends Watcher {
+  private def threadPrefix: String = name.map(n => n.replaceAll("\\s", "") + "-").getOrElse("")
+
+  // package level visibility for testing only
+  private[zookeeper] object ZooKeeperClientWatcher extends Watcher {
     override def process(event: WatchedEvent): Unit = {
-      debug("Received event: " + event)
+      debug(s"Received event: $event")
       Option(event.getPath) match {
         case None =>
+          val state = event.getState
+          stateToMeterMap.get(state).foreach(_.mark())
           inLock(isConnectedOrExpiredLock) {
             isConnectedOrExpiredCondition.signalAll()
           }
-          if (event.getState == KeeperState.AuthFailed) {
+          if (state == KeeperState.AuthFailed) {
             error("Auth failed.")
-            stateChangeHandlers.foreach {case (name, handler) => handler.onAuthFailure()}
-          } else if (event.getState == KeeperState.Expired) {
-            inWriteLock(initializationLock) {
-              info("Session expired.")
-              stateChangeHandlers.foreach {case (name, handler) => handler.beforeInitializingSession()}
-              initialize()
-              stateChangeHandlers.foreach {case (name, handler) => handler.afterInitializingSession()}
+            stateChangeHandlers.values.foreach(_.onAuthFailure())
+
+            // If this is during initial startup, we fail fast. Otherwise, schedule retry.
+            val initialized = inLock(isConnectedOrExpiredLock) {
+              isFirstConnectionEstablished
             }
+            if (initialized)
+              scheduleReinitialize("auth-failed", "Reinitializing due to auth failure.", RetryBackoffMs)
+          } else if (state == KeeperState.Expired) {
+            scheduleReinitialize("session-expired", "Session expired.", delayMs = 0L)
           }
         case Some(path) =>
           (event.getType: @unchecked) match {
@@ -330,7 +479,6 @@ trait StateChangeHandler {
   def beforeInitializingSession(): Unit = {}
   def afterInitializingSession(): Unit = {}
   def onAuthFailure(): Unit = {}
-  def onReconnectionTimeout(): Unit = {}
 }
 
 trait ZNodeChangeHandler {
@@ -345,6 +493,29 @@ trait ZNodeChildChangeHandler {
   def handleChildChange(): Unit = {}
 }
 
+// Thin wrapper for zookeeper.Op
+sealed trait ZkOp {
+  def toZookeeperOp: Op
+}
+
+case class CreateOp(path: String, data: Array[Byte], acl: Seq[ACL], createMode: CreateMode) extends ZkOp {
+  override def toZookeeperOp: Op = Op.create(path, data, acl.asJava, createMode)
+}
+
+case class DeleteOp(path: String, version: Int) extends ZkOp {
+  override def toZookeeperOp: Op = Op.delete(path, version)
+}
+
+case class SetDataOp(path: String, data: Array[Byte], version: Int) extends ZkOp {
+  override def toZookeeperOp: Op = Op.setData(path, data, version)
+}
+
+case class CheckOp(path: String, version: Int) extends ZkOp {
+  override def toZookeeperOp: Op = Op.check(path, version)
+}
+
+case class ZkOpResult(zkOp: ZkOp, rawOpResult: OpResult)
+
 sealed trait AsyncRequest {
   /**
    * This type member allows us to define methods that take requests and return responses with the correct types.
@@ -384,11 +555,18 @@ case class SetAclRequest(path: String, acl: Seq[ACL], version: Int, ctx: Option[
   type Response = SetAclResponse
 }
 
-case class GetChildrenRequest(path: String, ctx: Option[Any] = None) extends AsyncRequest {
+case class GetChildrenRequest(path: String, registerWatch: Boolean, ctx: Option[Any] = None) extends AsyncRequest {
   type Response = GetChildrenResponse
 }
 
-sealed trait AsyncResponse {
+case class MultiRequest(zkOps: Seq[ZkOp], ctx: Option[Any] = None) extends AsyncRequest {
+  type Response = MultiResponse
+
+  override def path: String = null
+}
+
+
+sealed abstract class AsyncResponse {
   def resultCode: Code
   def path: String
   def ctx: Option[Any]
@@ -396,15 +574,40 @@ sealed trait AsyncResponse {
   /** Return None if the result code is OK and KeeperException otherwise. */
   def resultException: Option[KeeperException] =
     if (resultCode == Code.OK) None else Some(KeeperException.create(resultCode, path))
+
+  /**
+   * Throw KeeperException if the result code is not OK.
+   */
+  def maybeThrow(): Unit = {
+    if (resultCode != Code.OK)
+      throw KeeperException.create(resultCode, path)
+  }
+
+  def metadata: ResponseMetadata
+}
+
+case class ResponseMetadata(sendTimeMs: Long, receivedTimeMs: Long) {
+  def responseTimeMs: Long = receivedTimeMs - sendTimeMs
 }
-case class CreateResponse(resultCode: Code, path: String, ctx: Option[Any], name: String) extends AsyncResponse
-case class DeleteResponse(resultCode: Code, path: String, ctx: Option[Any]) extends AsyncResponse
-case class ExistsResponse(resultCode: Code, path: String, ctx: Option[Any], stat: Stat) extends AsyncResponse
-case class GetDataResponse(resultCode: Code, path: String, ctx: Option[Any], data: Array[Byte], stat: Stat) extends AsyncResponse
-case class SetDataResponse(resultCode: Code, path: String, ctx: Option[Any], stat: Stat) extends AsyncResponse
-case class GetAclResponse(resultCode: Code, path: String, ctx: Option[Any], acl: Seq[ACL], stat: Stat) extends AsyncResponse
-case class SetAclResponse(resultCode: Code, path: String, ctx: Option[Any], stat: Stat) extends AsyncResponse
-case class GetChildrenResponse(resultCode: Code, path: String, ctx: Option[Any], children: Seq[String], stat: Stat) extends AsyncResponse
+
+case class CreateResponse(resultCode: Code, path: String, ctx: Option[Any], name: String,
+                          metadata: ResponseMetadata) extends AsyncResponse
+case class DeleteResponse(resultCode: Code, path: String, ctx: Option[Any],
+                          metadata: ResponseMetadata) extends AsyncResponse
+case class ExistsResponse(resultCode: Code, path: String, ctx: Option[Any], stat: Stat,
+                          metadata: ResponseMetadata) extends AsyncResponse
+case class GetDataResponse(resultCode: Code, path: String, ctx: Option[Any], data: Array[Byte], stat: Stat,
+                           metadata: ResponseMetadata) extends AsyncResponse
+case class SetDataResponse(resultCode: Code, path: String, ctx: Option[Any], stat: Stat,
+                           metadata: ResponseMetadata) extends AsyncResponse
+case class GetAclResponse(resultCode: Code, path: String, ctx: Option[Any], acl: Seq[ACL], stat: Stat,
+                          metadata: ResponseMetadata) extends AsyncResponse
+case class SetAclResponse(resultCode: Code, path: String, ctx: Option[Any], stat: Stat,
+                          metadata: ResponseMetadata) extends AsyncResponse
+case class GetChildrenResponse(resultCode: Code, path: String, ctx: Option[Any], children: Seq[String], stat: Stat,
+                               metadata: ResponseMetadata) extends AsyncResponse
+case class MultiResponse(resultCode: Code, path: String, ctx: Option[Any], zkOpResults: Seq[ZkOpResult],
+                         metadata: ResponseMetadata) extends AsyncResponse
 
 class ZooKeeperClientException(message: String) extends RuntimeException(message)
 class ZooKeeperClientExpiredException(message: String) extends ZooKeeperClientException(message)
diff --git a/core/src/main/scala/org/apache/zookeeper/ZooKeeperMainWithTlsSupportForKafka.scala b/core/src/main/scala/org/apache/zookeeper/ZooKeeperMainWithTlsSupportForKafka.scala
new file mode 100644
index 0000000000000..93674ad39e240
--- /dev/null
+++ b/core/src/main/scala/org/apache/zookeeper/ZooKeeperMainWithTlsSupportForKafka.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 org.apache.zookeeper
+
+import kafka.admin.ZkSecurityMigrator
+import org.apache.zookeeper.admin.ZooKeeperAdmin
+import org.apache.zookeeper.cli.CommandNotFoundException
+import org.apache.zookeeper.cli.MalformedCommandException
+import org.apache.zookeeper.client.ZKClientConfig
+
+import scala.jdk.CollectionConverters._
+
+object ZooKeeperMainWithTlsSupportForKafka {
+  val zkTlsConfigFileOption = "-zk-tls-config-file"
+  def main(args: Array[String]): Unit = {
+    val zkTlsConfigFileIndex = args.indexOf(zkTlsConfigFileOption)
+    val zooKeeperMain: ZooKeeperMain =
+      if (zkTlsConfigFileIndex < 0)
+        // no TLS config, so just pass args directly
+        new ZooKeeperMainWithTlsSupportForKafka(args, None)
+      else if (zkTlsConfigFileIndex == args.length - 1)
+        throw new IllegalArgumentException(s"Error: no filename provided with option $zkTlsConfigFileOption")
+      else
+        // found TLS config, so instantiate it and pass args without the two TLS config-related arguments
+        new ZooKeeperMainWithTlsSupportForKafka(
+          args.slice(0, zkTlsConfigFileIndex) ++ args.slice(zkTlsConfigFileIndex + 2, args.length),
+          Some(ZkSecurityMigrator.createZkClientConfigFromFile(args(zkTlsConfigFileIndex + 1))))
+    // The run method of ZooKeeperMain is package-private,
+    // therefore this code unfortunately must reside in the same org.apache.zookeeper package.
+    zooKeeperMain.run
+  }
+}
+
+class ZooKeeperMainWithTlsSupportForKafka(args: Array[String], val zkClientConfig: Option[ZKClientConfig])
+  extends ZooKeeperMain(args) with Watcher {
+
+  override def processZKCmd (co: ZooKeeperMain.MyCommandOptions): Boolean = {
+    // Unfortunately the usage() method is static, so it can't be overridden.
+    // This method is where usage() gets called.  We don't cover all possible calls
+    // to usage() -- we would have to implement the entire method to do that -- but
+    // the short implementation below covers most cases.
+    val args = co.getArgArray
+    val cmd = co.getCommand
+    if (args.length < 1) {
+      kafkaTlsUsage()
+      throw new MalformedCommandException("No command entered")
+    }
+
+    if (!ZooKeeperMain.commandMap.containsKey(cmd)) {
+      kafkaTlsUsage()
+      throw new CommandNotFoundException(s"Command not found $cmd")
+    }
+    super.processZKCmd(co)
+  }
+
+  def kafkaTlsUsage(): Unit = {
+    System.err.println("ZooKeeper -server host:port [-zk-tls-config-file ] cmd args")
+    ZooKeeperMain.commandMap.keySet.asScala.toList.sorted.foreach(cmd =>
+      System.err.println(s"\t$cmd ${ZooKeeperMain.commandMap.get(cmd)}"))
+  }
+
+  override def connectToZK(newHost: String) = {
+    // ZooKeeperAdmin has no constructor that supports passing in both readOnly and ZkClientConfig,
+    // and readOnly ends up being set to false when passing in a ZkClientConfig instance;
+    // therefore it is currently not possible for us to construct a ZooKeeperAdmin instance with
+    // both an explicit ZkClientConfig instance and a readOnly value of true.
+    val readOnlyRequested = cl.getOption("readonly") != null
+    if (readOnlyRequested && zkClientConfig.isDefined)
+      throw new IllegalArgumentException(
+        s"read-only mode (-r) is not supported with an explicit TLS config (${ZooKeeperMainWithTlsSupportForKafka.zkTlsConfigFileOption})")
+    if (zk != null && zk.getState.isAlive) zk.close()
+    host = newHost
+    zk = if (zkClientConfig.isDefined)
+      new ZooKeeperAdmin(host, cl.getOption("timeout").toInt, this, zkClientConfig.get)
+    else
+      new ZooKeeperAdmin(host, cl.getOption("timeout").toInt, this, readOnlyRequested)
+  }
+
+  override def process(event: WatchedEvent): Unit = {
+    if (getPrintWatches) {
+      ZooKeeperMain.printMessage("WATCHER::")
+      ZooKeeperMain.printMessage(event.toString)
+    }
+  }
+}
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/resources/minikdc-krb5.conf b/core/src/test/resources/minikdc-krb5.conf
index 0603404875525..20f1be5b1d39f 100644
--- a/core/src/test/resources/minikdc-krb5.conf
+++ b/core/src/test/resources/minikdc-krb5.conf
@@ -18,6 +18,8 @@
 [libdefaults]
 default_realm = {0}
 udp_preference_limit = 1
+default_tkt_enctypes=aes128-cts-hmac-sha1-96
+default_tgs_enctypes=aes128-cts-hmac-sha1-96
 
 [realms]
 {0} = '{'
diff --git a/core/src/test/scala/integration/kafka/admin/BrokerApiVersionsCommandTest.scala b/core/src/test/scala/integration/kafka/admin/BrokerApiVersionsCommandTest.scala
index 1f4b5e291d38e..5792c9c262df2 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
@@ -28,34 +30,43 @@ import org.apache.kafka.common.protocol.ApiKeys
 import org.junit.Assert.{assertEquals, assertFalse, assertNotNull, assertTrue}
 import org.junit.Test
 
+import scala.jdk.CollectionConverters._
+
 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)
     val content = new String(byteArrayOutputStream.toByteArray, StandardCharsets.UTF_8)
     val lineIter = content.split("\n").iterator
     assertTrue(lineIter.hasNext)
-    assertEquals(s"$brokerList (id: 0 rack: null) -> (", lineIter.next)
+    assertEquals(s"$brokerList (id: 0 rack: null) -> (", lineIter.next())
     val nodeApiVersions = NodeApiVersions.create
-    for (apiKey <- ApiKeys.values) {
+    val enabledApis = ApiKeys.enabledApis.asScala
+    for (apiKey <- enabledApis) {
       val apiVersion = nodeApiVersions.apiVersion(apiKey)
       assertNotNull(apiVersion)
+
       val versionRangeStr =
         if (apiVersion.minVersion == apiVersion.maxVersion) apiVersion.minVersion.toString
         else s"${apiVersion.minVersion} to ${apiVersion.maxVersion}"
-      val terminator = if (apiKey == ApiKeys.values.last) "" else ","
       val usableVersion = nodeApiVersions.latestUsableVersion(apiKey)
-      val line = s"\t${apiKey.name}(${apiKey.id}): $versionRangeStr [usable: $usableVersion]$terminator"
+      // Admin client should not see ENVELOPE supported versions as its a broker-internal API.
+      val usableVersionInfo = if (apiKey == ApiKeys.ENVELOPE) "UNSUPPORTED" else
+        s"$versionRangeStr [usable: $usableVersion]"
+
+      val terminator = if (apiKey == enabledApis.last) "" else ","
+
+      val line = s"\t${apiKey.name}(${apiKey.id}): $usableVersionInfo$terminator"
       assertTrue(lineIter.hasNext)
-      assertEquals(line, lineIter.next)
+      assertEquals(line, lineIter.next())
     }
     assertTrue(lineIter.hasNext)
-    assertEquals(")", lineIter.next)
+    assertEquals(")", lineIter.next())
     assertFalse(lineIter.hasNext)
   }
 }
diff --git a/core/src/test/scala/integration/kafka/admin/ReassignPartitionsIntegrationTest.scala b/core/src/test/scala/integration/kafka/admin/ReassignPartitionsIntegrationTest.scala
index 3f515288b482b..6d5d02d3caca5 100644
--- a/core/src/test/scala/integration/kafka/admin/ReassignPartitionsIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/admin/ReassignPartitionsIntegrationTest.scala
@@ -1,46 +1,733 @@
-/**
-  * Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE
-  * file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file
-  * to You under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the
-  * License. You may obtain a copy of the License at
-  *
-  * http://www.apache.org/licenses/LICENSE-2.0
-  *
-  * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
-  * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
-  * specific language governing permissions and limitations under the License.
-  */
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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.Closeable
+import java.util.{Collections, HashMap, List}
+
+import kafka.admin.ReassignPartitionsCommand._
+import kafka.api.KAFKA_2_7_IV1
+import kafka.server.{IsrChangePropagationConfig, KafkaConfig, KafkaServer, ReplicaManager}
+import kafka.utils.Implicits._
 import kafka.utils.TestUtils
-import kafka.zk.ZooKeeperTestHarness
-import org.junit.Test
+import kafka.zk.{KafkaZkClient, ZooKeeperTestHarness}
+import org.apache.kafka.clients.admin.{Admin, AdminClientConfig, AlterConfigOp, ConfigEntry, DescribeLogDirsResult, NewTopic}
+import org.apache.kafka.clients.producer.ProducerRecord
+import org.apache.kafka.common.config.ConfigResource
+import org.apache.kafka.common.network.ListenerName
+import org.apache.kafka.common.security.auth.SecurityProtocol
+import org.apache.kafka.common.utils.Utils
+import org.apache.kafka.common.{TopicPartition, TopicPartitionReplica}
+import org.junit.Assert.{assertEquals, assertFalse, assertTrue}
+import org.junit.rules.Timeout
+import org.junit.{After, Rule, Test}
+
+import scala.collection.{Map, Seq, mutable}
+import scala.jdk.CollectionConverters._
+
+class ReassignPartitionsIntegrationTest extends ZooKeeperTestHarness {
+  @Rule
+  def globalTimeout: Timeout = Timeout.millis(300000)
+
+  var cluster: ReassignPartitionsTestCluster = null
+
+  @After
+  override def tearDown(): Unit = {
+    Utils.closeQuietly(cluster, "ReassignPartitionsTestCluster")
+    super.tearDown()
+  }
+
+  val unthrottledBrokerConfigs =
+    0.to(4).map { brokerId =>
+      brokerId -> brokerLevelThrottles.map(throttle => (throttle, -1L)).toMap
+    }.toMap
+
+
+  @Test
+  def testReassignment(): Unit = {
+    cluster = new ReassignPartitionsTestCluster(zkConnect)
+    cluster.setup()
+    executeAndVerifyReassignment()
+  }
+
+  @Test
+  def testReassignmentWithAlterIsrDisabled(): Unit = {
+    // Test reassignment when the IBP is on an older version which does not use
+    // the `AlterIsr` API. In this case, the controller will register individual
+    // watches for each reassigning partition so that the reassignment can be
+    // completed as soon as the ISR is expanded.
+    val configOverrides = Map(KafkaConfig.InterBrokerProtocolVersionProp -> KAFKA_2_7_IV1.version)
+    cluster = new ReassignPartitionsTestCluster(zkConnect, configOverrides = configOverrides)
+    cluster.setup()
+    executeAndVerifyReassignment()
+  }
+
+  @Test
+  def testReassignmentCompletionDuringPartialUpgrade(): Unit = {
+    // Test reassignment during a partial upgrade when some brokers are relying on
+    // `AlterIsr` and some rely on the old notification logic through Zookeeper.
+    // In this test case, broker 0 starts up first on the latest IBP and is typically
+    // elected as controller. The three remaining brokers start up on the older IBP.
+    // We want to ensure that reassignment can still complete through the ISR change
+    // notification path even though the controller expects `AlterIsr`.
+
+    // Override change notification settings so that test is not delayed by ISR
+    // change notification delay
+    ReplicaManager.DefaultIsrPropagationConfig = IsrChangePropagationConfig(
+      checkIntervalMs = 500,
+      lingerMs = 100,
+      maxDelayMs = 500
+    )
+
+    val oldIbpConfig = Map(KafkaConfig.InterBrokerProtocolVersionProp -> KAFKA_2_7_IV1.version)
+    val brokerConfigOverrides = Map(1 -> oldIbpConfig, 2 -> oldIbpConfig, 3 -> oldIbpConfig)
+
+    cluster = new ReassignPartitionsTestCluster(zkConnect, brokerConfigOverrides = brokerConfigOverrides)
+    cluster.setup()
+
+    executeAndVerifyReassignment()
+  }
+
+  def executeAndVerifyReassignment(): Unit = {
+    val assignment = """{"version":1,"partitions":""" +
+      """[{"topic":"foo","partition":0,"replicas":[0,1,3],"log_dirs":["any","any","any"]},""" +
+      """{"topic":"bar","partition":0,"replicas":[3,2,0],"log_dirs":["any","any","any"]}""" +
+      """]}"""
+
+    // Check that the assignment has not yet been started yet.
+    val initialAssignment = Map(
+      new TopicPartition("foo", 0) ->
+        PartitionReassignmentState(Seq(0, 1, 2), Seq(0, 1, 3), true),
+      new TopicPartition("bar", 0) ->
+        PartitionReassignmentState(Seq(3, 2, 1), Seq(3, 2, 0), true)
+    )
+    waitForVerifyAssignment(cluster.adminClient, assignment, false,
+      VerifyAssignmentResult(initialAssignment))
+    waitForVerifyAssignment(zkClient, assignment, false,
+      VerifyAssignmentResult(initialAssignment))
+
+    // Execute the assignment
+    runExecuteAssignment(cluster.adminClient, false, assignment, -1L, -1L)
+    assertEquals(unthrottledBrokerConfigs,
+      describeBrokerLevelThrottles(unthrottledBrokerConfigs.keySet.toSeq))
+    val finalAssignment = Map(
+      new TopicPartition("foo", 0) ->
+        PartitionReassignmentState(Seq(0, 1, 3), Seq(0, 1, 3), true),
+      new TopicPartition("bar", 0) ->
+        PartitionReassignmentState(Seq(3, 2, 0), Seq(3, 2, 0), true)
+    )
+
+    // When using --zookeeper, we aren't able to see the new-style assignment
+    assertFalse(runVerifyAssignment(zkClient, assignment, false).movesOngoing)
+
+    // Wait for the assignment to complete
+    waitForVerifyAssignment(zkClient, assignment, false,
+      VerifyAssignmentResult(finalAssignment))
+
+    assertEquals(unthrottledBrokerConfigs,
+      describeBrokerLevelThrottles(unthrottledBrokerConfigs.keySet.toSeq))
+  }
+
+  /**
+   * Test running a quick reassignment with the --zookeeper option.
+   */
+  @Test
+  def testLegacyReassignment(): Unit = {
+    cluster = new ReassignPartitionsTestCluster(zkConnect)
+    cluster.setup()
+    val assignment = """{"version":1,"partitions":""" +
+      """[{"topic":"foo","partition":0,"replicas":[0,1,3],"log_dirs":["any","any","any"]},""" +
+      """{"topic":"bar","partition":0,"replicas":[3,2,0],"log_dirs":["any","any","any"]}""" +
+      """]}"""
+    // Execute the assignment
+    runExecuteAssignment(zkClient, assignment, -1L)
+    val finalAssignment = Map(
+      new TopicPartition("foo", 0) ->
+        PartitionReassignmentState(Seq(0, 1, 3), Seq(0, 1, 3), true),
+      new TopicPartition("bar", 0) ->
+        PartitionReassignmentState(Seq(3, 2, 0), Seq(3, 2, 0), true)
+    )
+    // Wait for the assignment to complete
+    waitForVerifyAssignment(cluster.adminClient, assignment, false,
+      VerifyAssignmentResult(finalAssignment))
+    waitForVerifyAssignment(zkClient, assignment, false,
+      VerifyAssignmentResult(finalAssignment))
+  }
 
-class ReassignPartitionsIntegrationTest extends ZooKeeperTestHarness with RackAwareTest {
+  @Test
+  def testHighWaterMarkAfterPartitionReassignment(): Unit = {
+    cluster = new ReassignPartitionsTestCluster(zkConnect)
+    cluster.setup()
+    val assignment = """{"version":1,"partitions":""" +
+      """[{"topic":"foo","partition":0,"replicas":[3,1,2],"log_dirs":["any","any","any"]}""" +
+      """]}"""
+
+    // Set the high water mark of foo-0 to 123 on its leader.
+    val part = new TopicPartition("foo", 0)
+    cluster.servers(0).replicaManager.logManager.truncateFullyAndStartAt(part, 123L, false)
+
+    // Execute the assignment
+    runExecuteAssignment(cluster.adminClient, false, assignment, -1L, -1L)
+    val finalAssignment = Map(part ->
+        PartitionReassignmentState(Seq(3, 1, 2), Seq(3, 1, 2), true))
+
+    // Wait for the assignment to complete
+    waitForVerifyAssignment(cluster.adminClient, assignment, false,
+      VerifyAssignmentResult(finalAssignment))
+
+    TestUtils.waitUntilTrue(() => {
+      cluster.servers(3).replicaManager.nonOfflinePartition(part).
+        flatMap(_.leaderLogIfLocal).isDefined
+      }, "broker 3 should be the new leader", pause = 10L)
+    assertEquals(s"Expected broker 3 to have the correct high water mark for the " +
+      "partition.", 123L, cluster.servers(3).replicaManager.
+      localLogOrException(part).highWatermark)
+  }
+
+  @Test
+  def testAlterReassignmentThrottle(): Unit = {
+    cluster = new ReassignPartitionsTestCluster(zkConnect)
+    cluster.setup()
+    cluster.produceMessages("foo", 0, 50)
+    cluster.produceMessages("baz", 2, 60)
+    val assignment = """{"version":1,"partitions":
+      [{"topic":"foo","partition":0,"replicas":[0,3,2],"log_dirs":["any","any","any"]},
+      {"topic":"baz","partition":2,"replicas":[3,2,1],"log_dirs":["any","any","any"]}
+      ]}"""
+
+    // Execute the assignment with a low throttle
+    val initialThrottle = 1L
+    runExecuteAssignment(cluster.adminClient, false, assignment, initialThrottle, -1L)
+    waitForInterBrokerThrottle(Set(0, 1, 2, 3), initialThrottle)
+
+    // Now update the throttle and verify the reassignment completes
+    val updatedThrottle = 300000L
+    runExecuteAssignment(cluster.adminClient, additional = true, assignment, updatedThrottle, -1L)
+    waitForInterBrokerThrottle(Set(0, 1, 2, 3), updatedThrottle)
+
+    val finalAssignment = Map(
+      new TopicPartition("foo", 0) ->
+        PartitionReassignmentState(Seq(0, 3, 2), Seq(0, 3, 2), true),
+      new TopicPartition("baz", 2) ->
+        PartitionReassignmentState(Seq(3, 2, 1), Seq(3, 2, 1), true))
+
+    // Now remove the throttles.
+    waitForVerifyAssignment(cluster.adminClient, assignment, false,
+      VerifyAssignmentResult(finalAssignment))
+    waitForBrokerLevelThrottles(unthrottledBrokerConfigs)
+  }
+
+  /**
+   * Test running a reassignment with the interBrokerThrottle set.
+   */
+  @Test
+  def testThrottledReassignment(): Unit = {
+    cluster = new ReassignPartitionsTestCluster(zkConnect)
+    cluster.setup()
+    cluster.produceMessages("foo", 0, 50)
+    cluster.produceMessages("baz", 2, 60)
+    val assignment = """{"version":1,"partitions":""" +
+      """[{"topic":"foo","partition":0,"replicas":[0,3,2],"log_dirs":["any","any","any"]},""" +
+      """{"topic":"baz","partition":2,"replicas":[3,2,1],"log_dirs":["any","any","any"]}""" +
+      """]}"""
+
+    // Check that the assignment has not yet been started yet.
+    val initialAssignment = Map(
+      new TopicPartition("foo", 0) ->
+        PartitionReassignmentState(Seq(0, 1, 2), Seq(0, 3, 2), true),
+      new TopicPartition("baz", 2) ->
+        PartitionReassignmentState(Seq(0, 2, 1), Seq(3, 2, 1), true))
+    assertEquals(VerifyAssignmentResult(initialAssignment),
+      runVerifyAssignment(cluster.adminClient, assignment, false))
+    assertEquals(VerifyAssignmentResult(initialAssignment),
+      runVerifyAssignment(zkClient, assignment, false))
+    assertEquals(unthrottledBrokerConfigs,
+      describeBrokerLevelThrottles(unthrottledBrokerConfigs.keySet.toSeq))
 
+    // Execute the assignment
+    val interBrokerThrottle = 300000L
+    runExecuteAssignment(cluster.adminClient, false, assignment, interBrokerThrottle, -1L)
+    waitForInterBrokerThrottle(Set(0, 1, 2, 3), interBrokerThrottle)
+
+    val finalAssignment = Map(
+      new TopicPartition("foo", 0) ->
+        PartitionReassignmentState(Seq(0, 3, 2), Seq(0, 3, 2), true),
+      new TopicPartition("baz", 2) ->
+        PartitionReassignmentState(Seq(3, 2, 1), Seq(3, 2, 1), true))
+
+    // Wait for the assignment to complete
+    TestUtils.waitUntilTrue(
+      () => {
+        // Check the reassignment status.
+        val result = runVerifyAssignment(cluster.adminClient, assignment, true)
+        if (!result.partsOngoing) {
+          true
+        } else {
+          assertTrue("Expected at least one partition reassignment to be ongoing when " +
+            s"result = ${result}", !result.partStates.forall(_._2.done))
+          assertEquals(Seq(0, 3, 2),
+            result.partStates(new TopicPartition("foo", 0)).targetReplicas)
+          assertEquals(Seq(3, 2, 1),
+            result.partStates(new TopicPartition("baz", 2)).targetReplicas)
+          logger.info(s"Current result: ${result}")
+          waitForInterBrokerThrottle(Set(0, 1, 2, 3), interBrokerThrottle)
+          false
+        }
+      }, "Expected reassignment to complete.")
+    waitForVerifyAssignment(cluster.adminClient, assignment, true,
+      VerifyAssignmentResult(finalAssignment))
+    waitForVerifyAssignment(zkClient, assignment, true,
+      VerifyAssignmentResult(finalAssignment))
+    // The throttles should still have been preserved, since we ran with --preserve-throttles
+    waitForInterBrokerThrottle(Set(0, 1, 2, 3), interBrokerThrottle)
+    // Now remove the throttles.
+    waitForVerifyAssignment(cluster.adminClient, assignment, false,
+      VerifyAssignmentResult(finalAssignment))
+    waitForBrokerLevelThrottles(unthrottledBrokerConfigs)
+  }
+
+  @Test
+  def testProduceAndConsumeWithReassignmentInProgress(): Unit = {
+    cluster = new ReassignPartitionsTestCluster(zkConnect)
+    cluster.setup()
+    cluster.produceMessages("baz", 2, 60)
+    val assignment = """{"version":1,"partitions":""" +
+      """[{"topic":"baz","partition":2,"replicas":[3,2,1],"log_dirs":["any","any","any"]}""" +
+      """]}"""
+    runExecuteAssignment(cluster.adminClient, false, assignment, 300L, -1L)
+    cluster.produceMessages("baz", 2, 100)
+    val consumer = TestUtils.createConsumer(cluster.brokerList)
+    val part = new TopicPartition("baz", 2)
+    try {
+      consumer.assign(Seq(part).asJava)
+      TestUtils.pollUntilAtLeastNumRecords(consumer, numRecords = 100)
+    } finally {
+      consumer.close()
+    }
+    TestUtils.removeReplicationThrottleForPartitions(cluster.adminClient, Seq(0,1,2,3), Set(part))
+    val finalAssignment = Map(part ->
+      PartitionReassignmentState(Seq(3, 2, 1), Seq(3, 2, 1), true))
+    waitForVerifyAssignment(cluster.adminClient, assignment, false,
+      VerifyAssignmentResult(finalAssignment))
+  }
+
+  /**
+   * Test running a reassignment and then cancelling it.
+   */
   @Test
-  def testRackAwareReassign() {
-    val rackInfo = Map(0 -> "rack1", 1 -> "rack2", 2 -> "rack2", 3 -> "rack1", 4 -> "rack3", 5 -> "rack3")
-    TestUtils.createBrokersInZk(toBrokerMetadata(rackInfo), zkUtils)
+  def testCancellation(): Unit = {
+    cluster = new ReassignPartitionsTestCluster(zkConnect)
+    cluster.setup()
+    cluster.produceMessages("foo", 0, 200)
+    cluster.produceMessages("baz", 1, 200)
+    val assignment = """{"version":1,"partitions":""" +
+      """[{"topic":"foo","partition":0,"replicas":[0,1,3],"log_dirs":["any","any","any"]},""" +
+      """{"topic":"baz","partition":1,"replicas":[0,2,3],"log_dirs":["any","any","any"]}""" +
+      """]}"""
+    assertEquals(unthrottledBrokerConfigs,
+      describeBrokerLevelThrottles(unthrottledBrokerConfigs.keySet.toSeq))
+    val interBrokerThrottle = 1L
+    runExecuteAssignment(cluster.adminClient, false, assignment, interBrokerThrottle, -1L)
+    waitForInterBrokerThrottle(Set(0, 1, 2, 3), interBrokerThrottle)
 
-    val numPartitions = 18
-    val replicationFactor = 3
+    // Verify that the reassignment is running.  The very low throttle should keep it
+    // from completing before this runs.
+    waitForVerifyAssignment(cluster.adminClient, assignment, true,
+      VerifyAssignmentResult(Map(
+        new TopicPartition("foo", 0) -> PartitionReassignmentState(Seq(0, 1, 3, 2), Seq(0, 1, 3), false),
+        new TopicPartition("baz", 1) -> PartitionReassignmentState(Seq(0, 2, 3, 1), Seq(0, 2, 3), false)),
+      true, Map(), false))
+    // Cancel the reassignment.
+    assertEquals((Set(
+        new TopicPartition("foo", 0),
+        new TopicPartition("baz", 1)
+      ), Set()), runCancelAssignment(cluster.adminClient, assignment, true))
+    // Broker throttles are still active because we passed --preserve-throttles
+    waitForInterBrokerThrottle(Set(0, 1, 2, 3), interBrokerThrottle)
+    // Cancelling the reassignment again should reveal nothing to cancel.
+    assertEquals((Set(), Set()), runCancelAssignment(cluster.adminClient, assignment, false))
+    // This time, the broker throttles were removed.
+    waitForBrokerLevelThrottles(unthrottledBrokerConfigs)
+    // Verify that there are no ongoing reassignments.
+    assertFalse(runVerifyAssignment(cluster.adminClient, assignment, false).partsOngoing)
+  }
+
+  private def waitForLogDirThrottle(throttledBrokers: Set[Int], logDirThrottle: Long): Unit = {
+    val throttledConfigMap = Map[String, Long](
+      brokerLevelLeaderThrottle -> -1,
+      brokerLevelFollowerThrottle -> -1,
+      brokerLevelLogDirThrottle -> logDirThrottle)
+    waitForBrokerThrottles(throttledBrokers, throttledConfigMap)
+  }
+
+  private def waitForInterBrokerThrottle(throttledBrokers: Set[Int], interBrokerThrottle: Long): Unit = {
+    val throttledConfigMap = Map[String, Long](
+      brokerLevelLeaderThrottle -> interBrokerThrottle,
+      brokerLevelFollowerThrottle -> interBrokerThrottle,
+      brokerLevelLogDirThrottle -> -1L)
+    waitForBrokerThrottles(throttledBrokers, throttledConfigMap)
+  }
 
-    // create a non rack aware assignment topic first
-    val createOpts = new kafka.admin.TopicCommand.TopicCommandOptions(Array(
-      "--partitions", numPartitions.toString,
-      "--replication-factor", replicationFactor.toString,
-      "--disable-rack-aware",
-      "--topic", "foo"))
-    kafka.admin.TopicCommand.createTopic(zkClient, createOpts)
+  private def waitForBrokerThrottles(throttledBrokers: Set[Int], throttleConfig: Map[String, Long]): Unit = {
+    val throttledBrokerConfigs = unthrottledBrokerConfigs.map { case (brokerId, unthrottledConfig) =>
+      val expectedThrottleConfig = if (throttledBrokers.contains(brokerId)) {
+        throttleConfig
+      } else {
+        unthrottledConfig
+      }
+      brokerId -> expectedThrottleConfig
+    }
+    waitForBrokerLevelThrottles(throttledBrokerConfigs)
+  }
 
-    val topicJson = """{"topics": [{"topic": "foo"}], "version":1}"""
-    val (proposedAssignment, currentAssignment) = ReassignPartitionsCommand.generateAssignment(zkUtils,
-      rackInfo.keys.toSeq.sorted, topicJson, disableRackAware = false)
+  private def waitForBrokerLevelThrottles(targetThrottles: Map[Int, Map[String, Long]]): Unit = {
+    var curThrottles: Map[Int, Map[String, Long]] = Map.empty
+    TestUtils.waitUntilTrue(() => {
+      curThrottles = describeBrokerLevelThrottles(targetThrottles.keySet.toSeq)
+      targetThrottles.equals(curThrottles)
+    }, s"timed out waiting for broker throttle to become ${targetThrottles}.  " +
+      s"Latest throttles were ${curThrottles}", pause = 25)
+  }
+
+  /**
+   * Describe the broker-level throttles in the cluster.
+   *
+   * @return                A map whose keys are broker IDs and whose values are throttle
+   *                        information.  The nested maps are keyed on throttle name.
+   */
+  private def describeBrokerLevelThrottles(brokerIds: Seq[Int]): Map[Int, Map[String, Long]] = {
+    brokerIds.map { brokerId =>
+      val props = zkClient.getEntityConfigs("brokers", brokerId.toString)
+      val throttles = brokerLevelThrottles.map { throttleName =>
+        (throttleName, props.getOrDefault(throttleName, "-1").asInstanceOf[String].toLong)
+      }.toMap
+      brokerId -> throttles
+    }.toMap
+  }
+
+  /**
+   * Test moving partitions between directories.
+   */
+  @Test
+  def testLogDirReassignment(): Unit = {
+    val topicPartition = new TopicPartition("foo", 0)
+
+    cluster = new ReassignPartitionsTestCluster(zkConnect)
+    cluster.setup()
+    cluster.produceMessages(topicPartition.topic, topicPartition.partition, 700)
+
+    val targetBrokerId = 0
+    val replicas = Seq(0, 1, 2)
+    val reassignment = buildLogDirReassignment(topicPartition, targetBrokerId, replicas)
+
+    // Start the replica move, but throttle it to be very slow so that it can't complete
+    // before our next checks happen.
+    val logDirThrottle = 1L
+    runExecuteAssignment(cluster.adminClient, additional = false, reassignment.json,
+      interBrokerThrottle = -1L, logDirThrottle)
+
+    // Check the output of --verify
+    waitForVerifyAssignment(cluster.adminClient, reassignment.json, true,
+      VerifyAssignmentResult(Map(
+          topicPartition -> PartitionReassignmentState(Seq(0, 1, 2), Seq(0, 1, 2), true)
+        ), false, Map(
+          new TopicPartitionReplica(topicPartition.topic, topicPartition.partition, 0) ->
+            ActiveMoveState(reassignment.currentDir, reassignment.targetDir, reassignment.targetDir)
+        ), true))
+    waitForLogDirThrottle(Set(0), logDirThrottle)
+
+    // Remove the throttle
+    cluster.adminClient.incrementalAlterConfigs(Collections.singletonMap(
+      new ConfigResource(ConfigResource.Type.BROKER, "0"),
+      Collections.singletonList(new AlterConfigOp(
+        new ConfigEntry(brokerLevelLogDirThrottle, ""), AlterConfigOp.OpType.DELETE)))).
+          all().get()
+    waitForBrokerLevelThrottles(unthrottledBrokerConfigs)
+
+    // Wait for the directory movement to complete.
+    waitForVerifyAssignment(cluster.adminClient, reassignment.json, true,
+        VerifyAssignmentResult(Map(
+          topicPartition -> PartitionReassignmentState(Seq(0, 1, 2), Seq(0, 1, 2), true)
+        ), false, Map(
+          new TopicPartitionReplica(topicPartition.topic, topicPartition.partition, 0) ->
+            CompletedMoveState(reassignment.targetDir)
+        ), false))
+
+    val info1 = new BrokerDirs(cluster.adminClient.describeLogDirs(0.to(4).
+        map(_.asInstanceOf[Integer]).asJavaCollection), 0)
+    assertEquals(reassignment.targetDir,
+      info1.curLogDirs.getOrElse(topicPartition, ""))
+  }
+
+  @Test
+  def testAlterLogDirReassignmentThrottle(): Unit = {
+    val topicPartition = new TopicPartition("foo", 0)
+
+    cluster = new ReassignPartitionsTestCluster(zkConnect)
+    cluster.setup()
+    cluster.produceMessages(topicPartition.topic, topicPartition.partition, 700)
+
+    val targetBrokerId = 0
+    val replicas = Seq(0, 1, 2)
+    val reassignment = buildLogDirReassignment(topicPartition, targetBrokerId, replicas)
+
+    // Start the replica move with a low throttle so it does not complete
+    val initialLogDirThrottle = 1L
+    runExecuteAssignment(cluster.adminClient, false, reassignment.json,
+      interBrokerThrottle = -1L, initialLogDirThrottle)
+    waitForLogDirThrottle(Set(0), initialLogDirThrottle)
+
+    // Now increase the throttle and verify that the log dir movement completes
+    val updatedLogDirThrottle = 3000000L
+    runExecuteAssignment(cluster.adminClient, additional = true, reassignment.json,
+      interBrokerThrottle = -1L, replicaAlterLogDirsThrottle = updatedLogDirThrottle)
+    waitForLogDirThrottle(Set(0), updatedLogDirThrottle)
+
+    waitForVerifyAssignment(cluster.adminClient, reassignment.json, true,
+      VerifyAssignmentResult(Map(
+        topicPartition -> PartitionReassignmentState(Seq(0, 1, 2), Seq(0, 1, 2), true)
+      ), false, Map(
+        new TopicPartitionReplica(topicPartition.topic, topicPartition.partition, targetBrokerId) ->
+          CompletedMoveState(reassignment.targetDir)
+      ), false))
+  }
+
+  case class LogDirReassignment(json: String, currentDir: String, targetDir: String)
+
+  private def buildLogDirReassignment(topicPartition: TopicPartition,
+                                      brokerId: Int,
+                                      replicas: Seq[Int]): LogDirReassignment = {
+
+    val describeLogDirsResult = cluster.adminClient.describeLogDirs(
+      0.to(4).map(_.asInstanceOf[Integer]).asJavaCollection)
+
+    val logDirInfo = new BrokerDirs(describeLogDirsResult, brokerId)
+    assertTrue(logDirInfo.futureLogDirs.isEmpty)
+
+    val currentDir = logDirInfo.curLogDirs(topicPartition)
+    val newDir = logDirInfo.logDirs.find(!_.equals(currentDir)).get
+
+    val logDirs = replicas.map { replicaId =>
+      if (replicaId == brokerId)
+        s""""$newDir""""
+      else
+        "\"any\""
+    }
+
+    val reassignmentJson =
+      s"""
+         | { "version": 1,
+         |  "partitions": [
+         |    {
+         |     "topic": "${topicPartition.topic}",
+         |     "partition": ${topicPartition.partition},
+         |     "replicas": [${replicas.mkString(",")}],
+         |     "log_dirs": [${logDirs.mkString(",")}]
+         |    }
+         |   ]
+         |  }
+         |""".stripMargin
+
+    LogDirReassignment(reassignmentJson, currentDir = currentDir, targetDir = newDir)
+  }
+
+  private def runVerifyAssignment(adminClient: Admin, jsonString: String,
+                                  preserveThrottles: Boolean) = {
+    println(s"==> verifyAssignment(adminClient, jsonString=${jsonString})")
+    verifyAssignment(adminClient, jsonString, preserveThrottles)
+  }
+
+  private def waitForVerifyAssignment(adminClient: Admin, jsonString: String,
+                                      preserveThrottles: Boolean,
+                                      expectedResult: VerifyAssignmentResult): Unit = {
+    var latestResult: VerifyAssignmentResult = null
+    TestUtils.waitUntilTrue(
+      () => {
+        latestResult = runVerifyAssignment(adminClient, jsonString, preserveThrottles)
+        expectedResult.equals(latestResult)
+      }, s"Timed out waiting for verifyAssignment result ${expectedResult}.  " +
+        s"The latest result was ${latestResult}", pause = 10L)
+  }
+
+  private def runVerifyAssignment(zkClient: KafkaZkClient, jsonString: String,
+                                  preserveThrottles: Boolean) = {
+    println(s"==> verifyAssignment(zkClient, jsonString=${jsonString})")
+    verifyAssignment(zkClient, jsonString, preserveThrottles)
+  }
+
+  private def waitForVerifyAssignment(zkClient: KafkaZkClient, jsonString: String,
+                                      preserveThrottles: Boolean,
+                                      expectedResult: VerifyAssignmentResult): Unit = {
+    var latestResult: VerifyAssignmentResult = null
+    TestUtils.waitUntilTrue(
+      () => {
+        println(s"==> verifyAssignment(zkClient, jsonString=${jsonString}, " +
+          s"preserveThrottles=${preserveThrottles})")
+        latestResult = verifyAssignment(zkClient, jsonString, preserveThrottles)
+        expectedResult.equals(latestResult)
+      }, s"Timed out waiting for verifyAssignment result ${expectedResult}.  " +
+        s"The latest result was ${latestResult}", pause = 10L)
+  }
+
+  private def runExecuteAssignment(adminClient: Admin,
+                                   additional: Boolean,
+                                   reassignmentJson: String,
+                                   interBrokerThrottle: Long,
+                                   replicaAlterLogDirsThrottle: Long) = {
+    println(s"==> executeAssignment(adminClient, additional=${additional}, " +
+      s"reassignmentJson=${reassignmentJson}, " +
+      s"interBrokerThrottle=${interBrokerThrottle}, " +
+      s"replicaAlterLogDirsThrottle=${replicaAlterLogDirsThrottle}))")
+    executeAssignment(adminClient, additional, reassignmentJson,
+      interBrokerThrottle, replicaAlterLogDirsThrottle)
+  }
+
+  private def runExecuteAssignment(zkClient: KafkaZkClient,
+                                   reassignmentJson: String,
+                                   interBrokerThrottle: Long) = {
+    println(s"==> executeAssignment(adminClient, " +
+      s"reassignmentJson=${reassignmentJson}, " +
+      s"interBrokerThrottle=${interBrokerThrottle})")
+    executeAssignment(zkClient, reassignmentJson, interBrokerThrottle)
+  }
+
+  private def runCancelAssignment(adminClient: Admin, jsonString: String,
+                                  preserveThrottles: Boolean) = {
+    println(s"==> cancelAssignment(adminClient, jsonString=${jsonString})")
+    cancelAssignment(adminClient, jsonString, preserveThrottles)
+  }
+
+  class BrokerDirs(result: DescribeLogDirsResult, val brokerId: Int) {
+    val logDirs = new mutable.HashSet[String]
+    val curLogDirs = new mutable.HashMap[TopicPartition, String]
+    val futureLogDirs = new mutable.HashMap[TopicPartition, String]
+    result.descriptions.get(brokerId).get().forEach {
+      case (logDirName, logDirInfo) => {
+        logDirs.add(logDirName)
+        logDirInfo.replicaInfos.forEach {
+          case (part, info) =>
+            if (info.isFuture) {
+              futureLogDirs.put(part, logDirName)
+            } else {
+              curLogDirs.put(part, logDirName)
+            }
+        }
+      }
+    }
+  }
+
+  class ReassignPartitionsTestCluster(
+    val zkConnect: String,
+    configOverrides: Map[String, String] = Map.empty,
+    brokerConfigOverrides: Map[Int, Map[String, String]] = Map.empty
+  ) extends Closeable {
+    val brokers = Map(
+      0 -> "rack0",
+      1 -> "rack0",
+      2 -> "rack1",
+      3 -> "rack1",
+      4 -> "rack1"
+    )
+
+    val topics = Map(
+      "foo" -> Seq(Seq(0, 1, 2), Seq(1, 2, 3)),
+      "bar" -> Seq(Seq(3, 2, 1)),
+      "baz" -> Seq(Seq(1, 0, 2), Seq(2, 0, 1), Seq(0, 2, 1))
+    )
+
+    val brokerConfigs = brokers.map {
+      case (brokerId, rack) =>
+        val config = TestUtils.createBrokerConfig(
+          nodeId = brokerId,
+          zkConnect = zkConnect,
+          rack = Some(rack),
+          enableControlledShutdown = false, // shorten test time
+          logDirCount = 3)
+        // shorter backoff to reduce test durations when no active partitions are eligible for fetching due to throttling
+        config.setProperty(KafkaConfig.ReplicaFetchBackoffMsProp, "100")
+        // Don't move partition leaders automatically.
+        config.setProperty(KafkaConfig.AutoLeaderRebalanceEnableProp, "false")
+        config.setProperty(KafkaConfig.ReplicaLagTimeMaxMsProp, "1000")
+        configOverrides.forKeyValue(config.setProperty)
+
+        brokerConfigOverrides.get(brokerId).foreach { overrides =>
+          overrides.forKeyValue(config.setProperty)
+        }
+
+        config
+    }.toBuffer
+
+    var servers = new mutable.ArrayBuffer[KafkaServer]
+
+    var brokerList: String = null
+
+    var adminClient: Admin = null
+
+    def setup(): Unit = {
+      createServers()
+      createTopics()
+    }
+
+    def createServers(): Unit = {
+      brokers.keySet.foreach { brokerId =>
+        servers += TestUtils.createServer(KafkaConfig(brokerConfigs(brokerId)))
+      }
+    }
+
+    def createTopics(): Unit = {
+      TestUtils.waitUntilBrokerMetadataIsPropagated(servers)
+      brokerList = TestUtils.bootstrapServers(servers,
+        ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT))
+      adminClient = Admin.create(Map[String, Object](
+        AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG -> brokerList
+      ).asJava)
+      adminClient.createTopics(topics.map {
+        case (topicName, parts) =>
+          val partMap = new HashMap[Integer, List[Integer]]()
+          parts.zipWithIndex.foreach {
+            case (part, index) => partMap.put(index, part.map(Integer.valueOf(_)).asJava)
+          }
+          new NewTopic(topicName, partMap)
+      }.toList.asJava).all().get()
+      topics.foreach {
+        case (topicName, parts) =>
+          parts.indices.foreach {
+            index => TestUtils.waitUntilMetadataIsPropagated(servers, topicName, index)
+          }
+      }
+    }
+
+    def produceMessages(topic: String, partition: Int, numMessages: Int): Unit = {
+      val records = (0 until numMessages).map(_ =>
+        new ProducerRecord[Array[Byte], Array[Byte]](topic, partition,
+          null, new Array[Byte](10000)))
+      TestUtils.produceMessages(servers, records, -1)
+    }
 
-    val assignment = proposedAssignment map { case (topicPartition, replicas) =>
-      (topicPartition.partition, replicas)
+    override def close(): Unit = {
+      brokerList = null
+      Utils.closeQuietly(adminClient, "adminClient")
+      adminClient = null
+      try {
+        TestUtils.shutdownServers(servers)
+      } finally {
+        servers.clear()
+      }
     }
-    checkReplicaDistribution(assignment, rackInfo, rackInfo.size, numPartitions, replicationFactor)
   }
 }
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..aa31173dfdfde
--- /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.jdk.CollectionConverters._
+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
deleted file mode 100644
index 0676a1a32637b..0000000000000
--- a/core/src/test/scala/integration/kafka/api/AdminClientIntegrationTest.scala
+++ /dev/null
@@ -1,1077 +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
-import java.util.{Collections, Properties}
-import java.util.Arrays.asList
-import java.util.concurrent.{ExecutionException, TimeUnit}
-import java.io.File
-import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger}
-
-import org.apache.kafka.clients.admin.KafkaAdminClientTest
-import org.apache.kafka.common.utils.{Time, Utils}
-import kafka.log.LogConfig
-import kafka.server.{Defaults, KafkaConfig, KafkaServer}
-import org.apache.kafka.clients.admin._
-import kafka.utils.{Logging, TestUtils, ZkUtils}
-import kafka.utils.Implicits._
-import org.apache.kafka.clients.admin.NewTopic
-import org.apache.kafka.clients.consumer.KafkaConsumer
-import org.apache.kafka.clients.producer.KafkaProducer
-import org.apache.kafka.clients.producer.ProducerRecord
-import org.apache.kafka.common.{KafkaFuture, TopicPartition, TopicPartitionReplica}
-import org.apache.kafka.common.acl._
-import org.apache.kafka.common.config.ConfigResource
-import org.apache.kafka.common.errors._
-import org.junit.{After, Before, Ignore, Rule, Test}
-import org.apache.kafka.common.requests.{DeleteRecordsRequest, MetadataResponse}
-import org.apache.kafka.common.resource.{Resource, ResourceType}
-import org.junit.rules.Timeout
-import org.junit.Assert._
-
-import scala.util.Random
-import scala.collection.JavaConverters._
-import java.lang.{Long => JLong}
-
-import scala.concurrent.duration.Duration
-import scala.concurrent.{Await, Future}
-
-/**
- * An integration test of the KafkaAdminClient.
- *
- * Also see {@link org.apache.kafka.clients.admin.KafkaAdminClientTest} for a unit test of the admin client.
- */
-class AdminClientIntegrationTest extends IntegrationTestHarness with Logging {
-
-  import AdminClientIntegrationTest._
-
-  @Rule
-  def globalTimeout = Timeout.millis(120000)
-
-  var client: AdminClient = null
-
-  val topic = "topic"
-  val partition = 0
-  val topicPartition = new TopicPartition(topic, partition)
-
-  @Before
-  override def setUp(): Unit = {
-    super.setUp
-    TestUtils.waitUntilBrokerMetadataIsPropagated(servers)
-  }
-
-  @After
-  override def tearDown(): Unit = {
-    if (client != null)
-      Utils.closeQuietly(client, "AdminClient")
-    super.tearDown()
-  }
-
-  val serverCount = 3
-  val consumerCount = 1
-  val producerCount = 1
-
-  override def generateConfigs = {
-    val cfgs = TestUtils.createBrokerConfigs(serverCount, zkConnect, interBrokerSecurityProtocol = Some(securityProtocol),
-      trustStoreFile = trustStoreFile, saslProperties = serverSaslProperties, logDirCount = 2)
-    cfgs.foreach { config =>
-      config.setProperty(KafkaConfig.ListenersProp, s"${listenerName.value}://localhost:${TestUtils.RandomPort}")
-      config.remove(KafkaConfig.InterBrokerSecurityProtocolProp)
-      config.setProperty(KafkaConfig.InterBrokerListenerNameProp, listenerName.value)
-      config.setProperty(KafkaConfig.ListenerSecurityProtocolMapProp, s"${listenerName.value}:${securityProtocol.name}")
-      config.setProperty(KafkaConfig.DeleteTopicEnableProp, "true")
-      // 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))
-        config.setProperty(KafkaConfig.SslTruststorePasswordProp, "some.invalid.pass")
-    }
-    cfgs.foreach(_ ++= serverConfig)
-    cfgs.map(KafkaConfig.fromProps)
-  }
-
-  def createConfig(): util.Map[String, Object] = {
-    val config = new util.HashMap[String, Object]
-    config.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList)
-    config.put(AdminClientConfig.REQUEST_TIMEOUT_MS_CONFIG, "20000")
-    val securityProps: util.Map[Object, Object] =
-      TestUtils.adminClientSecurityConfigs(securityProtocol, trustStoreFile, clientSaslProperties)
-    securityProps.asScala.foreach { case (key, value) => config.put(key.asInstanceOf[String], value) }
-    config
-  }
-
-  def waitForTopics(client: AdminClient, expectedPresent: Seq[String], expectedMissing: Seq[String]): Unit = {
-    TestUtils.waitUntilTrue(() => {
-        val topics = client.listTopics.names.get()
-        expectedPresent.forall(topicName => topics.contains(topicName)) &&
-          expectedMissing.forall(topicName => !topics.contains(topicName))
-      }, "timed out waiting for topics")
-  }
-
-  def assertFutureExceptionTypeEquals(future: KafkaFuture[_], clazz: Class[_ <: Throwable]): Unit = {
-    try {
-      future.get()
-      fail("Expected CompletableFuture.get to return an exception")
-    } catch {
-      case e: ExecutionException =>
-        val cause = e.getCause()
-        assertTrue("Expected an exception of type " + clazz.getName + "; got type " +
-            cause.getClass().getName, clazz.isInstance(cause))
-    }
-  }
-
-  @Test
-  def testClose(): Unit = {
-    val client = AdminClient.create(createConfig())
-    client.close()
-    client.close() // double close has no effect
-  }
-
-  @Test
-  def testListNodes(): Unit = {
-    client = AdminClient.create(createConfig())
-    val brokerStrs = brokerList.split(",").toList.sorted
-    var nodeStrs: List[String] = null
-    do {
-      val nodes = client.describeCluster().nodes().get().asScala
-      nodeStrs = nodes.map ( node => s"${node.host}:${node.port}" ).toList.sorted
-    } while (nodeStrs.size < brokerStrs.size)
-    assertEquals(brokerStrs.mkString(","), nodeStrs.mkString(","))
-  }
-
-  @Test
-  def testCreateDeleteTopics(): Unit = {
-    client = AdminClient.create(createConfig())
-    val topics = Seq("mytopic", "mytopic2")
-    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)
-    )
-    client.createTopics(newTopics.asJava, new CreateTopicsOptions().validateOnly(true)).all.get()
-    waitForTopics(client, List(), topics)
-
-    client.createTopics(newTopics.asJava).all.get()
-    waitForTopics(client, topics, List())
-
-    val results = client.createTopics(newTopics.asJava).values()
-    assertTrue(results.containsKey("mytopic"))
-    assertFutureExceptionTypeEquals(results.get("mytopic"), classOf[TopicExistsException])
-    assertTrue(results.containsKey("mytopic2"))
-    assertFutureExceptionTypeEquals(results.get("mytopic2"), classOf[TopicExistsException])
-
-    val topicToDescription = client.describeTopics(topics.asJava).all.get()
-    assertEquals(topics.toSet, topicToDescription.keySet.asScala)
-
-    val topic0 = topicToDescription.get("mytopic")
-    assertEquals(false, topic0.isInternal)
-    assertEquals("mytopic", topic0.name)
-    assertEquals(2, topic0.partitions.size)
-    val topic0Partition0 = topic0.partitions.get(0)
-    assertEquals(1, topic0Partition0.leader.id)
-    assertEquals(0, topic0Partition0.partition)
-    assertEquals(Seq(1, 2), topic0Partition0.isr.asScala.map(_.id))
-    assertEquals(Seq(1, 2), topic0Partition0.replicas.asScala.map(_.id))
-    val topic0Partition1 = topic0.partitions.get(1)
-    assertEquals(2, topic0Partition1.leader.id)
-    assertEquals(1, topic0Partition1.partition)
-    assertEquals(Seq(2, 0), topic0Partition1.isr.asScala.map(_.id))
-    assertEquals(Seq(2, 0), topic0Partition1.replicas.asScala.map(_.id))
-
-    val topic1 = topicToDescription.get("mytopic2")
-    assertEquals(false, topic1.isInternal)
-    assertEquals("mytopic2", topic1.name)
-    assertEquals(3, topic1.partitions.size)
-    for (partitionId <- 0 until 3) {
-      val partition = topic1.partitions.get(partitionId)
-      assertEquals(partitionId, partition.partition)
-      assertEquals(3, partition.replicas.size)
-      partition.replicas.asScala.foreach { replica =>
-        assertTrue(replica.id >= 0)
-        assertTrue(replica.id < serverCount)
-      }
-      assertEquals("No duplicate replica ids", partition.replicas.size, partition.replicas.asScala.map(_.id).distinct.size)
-
-      assertEquals(3, partition.isr.size)
-      assertEquals(partition.replicas, partition.isr)
-      assertTrue(partition.replicas.contains(partition.leader))
-    }
-
-    client.deleteTopics(topics.asJava).all.get()
-    waitForTopics(client, List(), topics)
-  }
-
-  /**
-    * describe should not auto create topics
-    */
-  @Test
-  def testDescribeNonExistingTopic(): Unit = {
-    client = AdminClient.create(createConfig())
-
-    val existingTopic = "existing-topic"
-    client.createTopics(Seq(existingTopic).map(new NewTopic(_, 1, 1)).asJava).all.get()
-    waitForTopics(client, Seq(existingTopic), List())
-
-    val nonExistingTopic = "non-existing"
-    val results = client.describeTopics(Seq(nonExistingTopic, existingTopic).asJava).values
-    assertEquals(existingTopic, results.get(existingTopic).get.name)
-    intercept[ExecutionException](results.get(nonExistingTopic).get).getCause.isInstanceOf[UnknownTopicOrPartitionException]
-    assertEquals(None, zkUtils.getTopicPartitionCount(nonExistingTopic))
-  }
-
-  @Test
-  def testDescribeCluster(): Unit = {
-    client = AdminClient.create(createConfig())
-    val nodes = client.describeCluster.nodes.get()
-    val clusterId = client.describeCluster().clusterId().get()
-    assertEquals(servers.head.apis.clusterId, clusterId)
-    val controller = client.describeCluster().controller().get()
-    assertEquals(servers.head.apis.metadataCache.getControllerId.
-      getOrElse(MetadataResponse.NO_CONTROLLER_ID), controller.id())
-    val brokers = brokerList.split(",")
-    assertEquals(brokers.size, nodes.size)
-    for (node <- nodes.asScala) {
-      val hostStr = s"${node.host}:${node.port}"
-      assertTrue(s"Unknown host:port pair $hostStr in brokerVersionInfos", brokers.contains(hostStr))
-    }
-  }
-
-  @Test
-  def testDescribeLogDirs(): Unit = {
-    client = AdminClient.create(createConfig())
-    val topic = "topic"
-    val leaderByPartition = TestUtils.createTopic(zkUtils, topic, 10, 1, servers, new Properties())
-    val partitionsByBroker = leaderByPartition.groupBy { case (partitionId, leaderId) => leaderId }.mapValues(_.keys.toSeq)
-    val brokers = (0 until serverCount).map(Integer.valueOf)
-    val logDirInfosByBroker = client.describeLogDirs(brokers.asJava).all.get
-
-    (0 until serverCount).foreach { brokerId =>
-      val server = servers.find(_.config.brokerId == brokerId).get
-      val expectedPartitions = partitionsByBroker(brokerId)
-      val logDirInfos = logDirInfosByBroker.get(brokerId)
-      val replicaInfos = logDirInfos.asScala.flatMap { case (logDir, logDirInfo) => logDirInfo.replicaInfos.asScala }.filterKeys(_.topic == topic)
-
-      assertEquals(expectedPartitions.toSet, replicaInfos.keys.map(_.partition).toSet)
-      logDirInfos.asScala.foreach { case (logDir, logDirInfo) =>
-        logDirInfo.replicaInfos.asScala.keys.foreach(tp =>
-          assertEquals(server.logManager.getLog(tp).get.dir.getParent, logDir)
-        )
-      }
-    }
-
-    client.close()
-  }
-
-  @Test
-  def testDescribeReplicaLogDirs(): Unit = {
-    client = AdminClient.create(createConfig())
-    val topic = "topic"
-    val leaderByPartition = TestUtils.createTopic(zkUtils, topic, 10, 1, servers, new Properties())
-    val replicas = leaderByPartition.map { case (partition, brokerId) => new TopicPartitionReplica(topic, partition, brokerId) }.toSeq
-
-    val replicaDirInfos = client.describeReplicaLogDirs(replicas.asJavaCollection).all.get
-    replicaDirInfos.asScala.foreach { case (topicPartitionReplica, replicaDirInfo) =>
-      val server = servers.find(_.config.brokerId == topicPartitionReplica.brokerId()).get
-      val tp = new TopicPartition(topicPartitionReplica.topic(), topicPartitionReplica.partition())
-      assertEquals(server.logManager.getLog(tp).get.dir.getParent, replicaDirInfo.getCurrentReplicaLogDir)
-    }
-  }
-
-  @Test
-  def testAlterReplicaLogDirs(): Unit = {
-    client = AdminClient.create(createConfig())
-    val topic = "topic"
-    val tp = new TopicPartition(topic, 0)
-    val randomNums = servers.map(server => server -> Random.nextInt(2)).toMap
-
-    // Generate two mutually exclusive replicaAssignment
-    val firstReplicaAssignment = servers.map { server =>
-      val logDir = new File(server.config.logDirs(randomNums(server))).getAbsolutePath
-      new TopicPartitionReplica(topic, 0, server.config.brokerId) -> logDir
-    }.toMap
-    val secondReplicaAssignment = servers.map { server =>
-      val logDir = new File(server.config.logDirs(1 - randomNums(server))).getAbsolutePath
-      new TopicPartitionReplica(topic, 0, server.config.brokerId) -> logDir
-    }.toMap
-
-    // Verify that replica can be created in the specified log directory
-    val futures = client.alterReplicaLogDirs(firstReplicaAssignment.asJava,
-      new AlterReplicaLogDirsOptions).values.asScala.values
-    futures.foreach { future =>
-      val exception = intercept[ExecutionException](future.get)
-      assertTrue(exception.getCause.isInstanceOf[ReplicaNotAvailableException])
-    }
-
-    TestUtils.createTopic(zkUtils, topic, 1, serverCount, servers, new Properties)
-    servers.foreach { server =>
-      val logDir = server.logManager.getLog(tp).get.dir.getParent
-      assertEquals(firstReplicaAssignment(new TopicPartitionReplica(topic, 0, server.config.brokerId)), logDir)
-    }
-
-    // Verify that replica can be moved to the specified log directory after the topic has been created
-    client.alterReplicaLogDirs(secondReplicaAssignment.asJava, new AlterReplicaLogDirsOptions).all.get
-    servers.foreach { server =>
-      TestUtils.waitUntilTrue(() => {
-        val logDir = server.logManager.getLog(tp).get.dir.getParent
-        secondReplicaAssignment(new TopicPartitionReplica(topic, 0, server.config.brokerId)) == logDir
-      }, "timed out waiting for replica movement", 6000L)
-    }
-
-    // Verify that replica can be moved to the specified log directory while the producer is sending messages
-    val running = new AtomicBoolean(true)
-    val numMessages = new AtomicInteger
-    import scala.concurrent.ExecutionContext.Implicits._
-    val producerFuture = Future {
-      val producer = TestUtils.createNewProducer(
-        TestUtils.getBrokerListStrFromServers(servers, protocol = securityProtocol),
-        securityProtocol = securityProtocol,
-        trustStoreFile = trustStoreFile,
-        retries = 0, // Producer should not have to retry when broker is moving replica between log directories.
-        requestTimeoutMs = 10000,
-        acks = -1
-      )
-      try {
-        while (running.get) {
-          val future = producer.send(new ProducerRecord(topic, s"xxxxxxxxxxxxxxxxxxxx-$numMessages".getBytes))
-          numMessages.incrementAndGet()
-          future.get(10, TimeUnit.SECONDS)
-        }
-        numMessages.get
-      } finally producer.close()
-    }
-
-    try {
-      TestUtils.waitUntilTrue(() => numMessages.get > 100, "timed out waiting for message produce", 6000L)
-      client.alterReplicaLogDirs(firstReplicaAssignment.asJava, new AlterReplicaLogDirsOptions).all.get
-      servers.foreach { server =>
-        TestUtils.waitUntilTrue(() => {
-          val logDir = server.logManager.getLog(tp).get.dir.getParent
-          firstReplicaAssignment(new TopicPartitionReplica(topic, 0, server.config.brokerId)) == logDir
-        }, "timed out waiting for replica movement", 6000L)
-      }
-    } finally running.set(false)
-
-    val finalNumMessages = Await.result(producerFuture, Duration(20, TimeUnit.SECONDS))
-
-    // Verify that all messages that are produced can be consumed
-    val consumerRecords = TestUtils.consumeTopicRecords(servers, topic, finalNumMessages, securityProtocol, trustStoreFile)
-    consumerRecords.zipWithIndex.foreach { case (consumerRecord, index) =>
-      assertEquals(s"xxxxxxxxxxxxxxxxxxxx-$index", new String(consumerRecord.value))
-    }
-  }
-
-  @Test
-  def testDescribeAndAlterConfigs(): Unit = {
-    client = AdminClient.create(createConfig)
-
-    // Create topics
-    val topic1 = "describe-alter-configs-topic-1"
-    val topicResource1 = new ConfigResource(ConfigResource.Type.TOPIC, topic1)
-    val topicConfig1 = new Properties
-    topicConfig1.setProperty(LogConfig.MaxMessageBytesProp, "500000")
-    topicConfig1.setProperty(LogConfig.RetentionMsProp, "60000000")
-    TestUtils.createTopic(zkUtils, topic1, 1, 1, servers, topicConfig1)
-
-    val topic2 = "describe-alter-configs-topic-2"
-    val topicResource2 = new ConfigResource(ConfigResource.Type.TOPIC, topic2)
-    TestUtils.createTopic(zkUtils, topic2, 1, 1, servers, new Properties)
-
-    // Describe topics and broker
-    val brokerResource1 = new ConfigResource(ConfigResource.Type.BROKER, servers(1).config.brokerId.toString)
-    val brokerResource2 = new ConfigResource(ConfigResource.Type.BROKER, servers(2).config.brokerId.toString)
-    val configResources = Seq(topicResource1, topicResource2, brokerResource1, brokerResource2)
-    val describeResult = client.describeConfigs(configResources.asJava)
-    val configs = describeResult.all.get
-
-    assertEquals(4, configs.size)
-
-    val maxMessageBytes1 = configs.get(topicResource1).get(LogConfig.MaxMessageBytesProp)
-    assertEquals(LogConfig.MaxMessageBytesProp, maxMessageBytes1.name)
-    assertEquals(topicConfig1.get(LogConfig.MaxMessageBytesProp), maxMessageBytes1.value)
-    assertFalse(maxMessageBytes1.isDefault)
-    assertFalse(maxMessageBytes1.isSensitive)
-    assertFalse(maxMessageBytes1.isReadOnly)
-
-    assertEquals(topicConfig1.get(LogConfig.RetentionMsProp),
-      configs.get(topicResource1).get(LogConfig.RetentionMsProp).value)
-
-    val maxMessageBytes2 = configs.get(topicResource2).get(LogConfig.MaxMessageBytesProp)
-    assertEquals(Defaults.MessageMaxBytes.toString, maxMessageBytes2.value)
-    assertEquals(LogConfig.MaxMessageBytesProp, maxMessageBytes2.name)
-    assertTrue(maxMessageBytes2.isDefault)
-    assertFalse(maxMessageBytes2.isSensitive)
-    assertFalse(maxMessageBytes2.isReadOnly)
-
-    assertEquals(servers(1).config.values.size, configs.get(brokerResource1).entries.size)
-    assertEquals(servers(1).config.brokerId.toString, configs.get(brokerResource1).get(KafkaConfig.BrokerIdProp).value)
-    val listenerSecurityProtocolMap = configs.get(brokerResource1).get(KafkaConfig.ListenerSecurityProtocolMapProp)
-    assertEquals(servers(1).config.getString(KafkaConfig.ListenerSecurityProtocolMapProp), listenerSecurityProtocolMap.value)
-    assertEquals(KafkaConfig.ListenerSecurityProtocolMapProp, listenerSecurityProtocolMap.name)
-    assertFalse(listenerSecurityProtocolMap.isDefault)
-    assertFalse(listenerSecurityProtocolMap.isSensitive)
-    assertTrue(listenerSecurityProtocolMap.isReadOnly)
-    val truststorePassword = configs.get(brokerResource1).get(KafkaConfig.SslTruststorePasswordProp)
-    assertEquals(KafkaConfig.SslTruststorePasswordProp, truststorePassword.name)
-    assertNull(truststorePassword.value)
-    assertFalse(truststorePassword.isDefault)
-    assertTrue(truststorePassword.isSensitive)
-    assertTrue(truststorePassword.isReadOnly)
-    val compressionType = configs.get(brokerResource1).get(KafkaConfig.CompressionTypeProp)
-    assertEquals(servers(1).config.compressionType.toString, compressionType.value)
-    assertEquals(KafkaConfig.CompressionTypeProp, compressionType.name)
-    assertTrue(compressionType.isDefault)
-    assertFalse(compressionType.isSensitive)
-    assertTrue(compressionType.isReadOnly)
-
-    assertEquals(servers(2).config.values.size, configs.get(brokerResource2).entries.size)
-    assertEquals(servers(2).config.brokerId.toString, configs.get(brokerResource2).get(KafkaConfig.BrokerIdProp).value)
-    assertEquals(servers(2).config.logCleanerThreads.toString,
-      configs.get(brokerResource2).get(KafkaConfig.LogCleanerThreadsProp).value)
-
-    checkValidAlterConfigs(zkUtils, servers, client, topicResource1, topicResource2)
-  }
-
-  @Test
-  def testCreatePartitions(): Unit = {
-    client = AdminClient.create(createConfig)
-
-    // Create topics
-    val topic1 = "create-partitions-topic-1"
-    TestUtils.createTopic(zkUtils, topic1, 1, 1, servers, new Properties)
-
-    val topic2 = "create-partitions-topic-2"
-    TestUtils.createTopic(zkUtils, topic2, 1, 2, servers, new Properties)
-
-    // 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 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 numPartitions(topic: String) =
-      partitions(topic).size
-
-    // validateOnly: try creating a new partition (no assignments), to bring the total to 3 partitions
-    var alterResult = client.createPartitions(Map(topic1 ->
-      NewPartitions.increaseTo(3)).asJava, validateOnly)
-    var altered = alterResult.values.get(topic1).get
-    assertEquals(1, numPartitions(topic1))
-
-    // try creating a new partition (no assignments), to bring the total to 3 partitions
-    alterResult = client.createPartitions(Map(topic1 ->
-      NewPartitions.increaseTo(3)).asJava, actuallyDoIt)
-    altered = alterResult.values.get(topic1).get
-    assertEquals(3, numPartitions(topic1))
-
-    // 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))
-    alterResult = client.createPartitions(Map(topic2 ->
-      NewPartitions.increaseTo(3, newPartition2Assignments)).asJava, validateOnly)
-    altered = alterResult.values.get(topic2).get
-    assertEquals(1, numPartitions(topic2))
-
-    // now try creating a new partition (with assignments), to bring the total to 3 partitions
-    alterResult = client.createPartitions(Map(topic2 ->
-      NewPartitions.increaseTo(3, newPartition2Assignments)).asJava, actuallyDoIt)
-    altered = alterResult.values.get(topic2).get
-    val actualPartitions2 = partitions(topic2)
-    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)
-
-    // loop over error cases calling with+without validate-only
-    for (option <- Seq(validateOnly, actuallyDoIt)) {
-      val desc = if (option.validateOnly()) "validateOnly" else "validateOnly=false"
-
-      // try a newCount which would be a decrease
-      alterResult = client.createPartitions(Map(topic1 ->
-        NewPartitions.increaseTo(1)).asJava, option)
-      try {
-        alterResult.values.get(topic1).get
-        fail(s"$desc: Expect InvalidPartitionsException when newCount is a decrease")
-      } catch {
-        case e: ExecutionException =>
-          assertTrue(desc, e.getCause.isInstanceOf[InvalidPartitionsException])
-          assertEquals(desc, "Topic currently has 3 partitions, which is higher than the requested 1.", e.getCause.getMessage)
-          assertEquals(desc, 3, numPartitions(topic1))
-      }
-
-      // try a newCount which would be a noop (without assignment)
-      alterResult = client.createPartitions(Map(topic2 ->
-        NewPartitions.increaseTo(3)).asJava, option)
-      try {
-        alterResult.values.get(topic2).get
-        fail(s"$desc: Expect InvalidPartitionsException when requesting a noop")
-      } catch {
-        case e: ExecutionException =>
-          assertTrue(desc, e.getCause.isInstanceOf[InvalidPartitionsException])
-          assertEquals(desc, "Topic already has 3 partitions.", e.getCause.getMessage)
-          assertEquals(desc, 3, numPartitions(topic2))
-      }
-
-      // try a newCount which would be a noop (where the assignment matches current state)
-      alterResult = client.createPartitions(Map(topic2 ->
-        NewPartitions.increaseTo(3, newPartition2Assignments)).asJava, option)
-      try {
-        alterResult.values.get(topic2).get
-      } catch {
-        case e: ExecutionException =>
-          assertTrue(desc, e.getCause.isInstanceOf[InvalidPartitionsException])
-          assertEquals(desc, "Topic already has 3 partitions.", e.getCause.getMessage)
-          assertEquals(desc, 3, numPartitions(topic2))
-      }
-
-      // try a newCount which would be a noop (where the assignment doesn't match current state)
-      alterResult = client.createPartitions(Map(topic2 ->
-        NewPartitions.increaseTo(3, newPartition2Assignments.asScala.reverse.toList.asJava)).asJava, option)
-      try {
-        alterResult.values.get(topic2).get
-      } catch {
-        case e: ExecutionException =>
-          assertTrue(desc, e.getCause.isInstanceOf[InvalidPartitionsException])
-          assertEquals(desc, "Topic already has 3 partitions.", e.getCause.getMessage)
-          assertEquals(desc, 3, numPartitions(topic2))
-      }
-
-      // try a bad topic name
-      val unknownTopic = "an-unknown-topic"
-      alterResult = client.createPartitions(Map(unknownTopic ->
-        NewPartitions.increaseTo(2)).asJava, option)
-      try {
-        alterResult.values.get(unknownTopic).get
-        fail(s"$desc: Expect InvalidTopicException when using an unknown topic")
-      } catch {
-        case e: ExecutionException =>
-          assertTrue(desc, e.getCause.isInstanceOf[UnknownTopicOrPartitionException])
-          assertEquals(desc, "The topic 'an-unknown-topic' does not exist.", e.getCause.getMessage)
-      }
-
-      // try an invalid newCount
-      alterResult = client.createPartitions(Map(topic1 ->
-        NewPartitions.increaseTo(-22)).asJava, option)
-      try {
-        altered = alterResult.values.get(topic1).get
-        fail(s"$desc: Expect InvalidPartitionsException when newCount is invalid")
-      } catch {
-        case e: ExecutionException =>
-          assertTrue(desc, e.getCause.isInstanceOf[InvalidPartitionsException])
-          assertEquals(desc, "Topic currently has 3 partitions, which is higher than the requested -22.",
-            e.getCause.getMessage)
-          assertEquals(desc, 3, numPartitions(topic1))
-      }
-
-      // try assignments where the number of brokers != replication factor
-      alterResult = client.createPartitions(Map(topic1 ->
-        NewPartitions.increaseTo(4, asList(asList(1, 2)))).asJava, option)
-      try {
-        altered = alterResult.values.get(topic1).get
-        fail(s"$desc: Expect InvalidPartitionsException when #brokers != replication factor")
-      } catch {
-        case e: ExecutionException =>
-          assertTrue(desc, e.getCause.isInstanceOf[InvalidReplicaAssignmentException])
-          assertEquals(desc, "Inconsistent replication factor between partitions, partition 0 has 1 " +
-            "while partitions [3] have replication factors [2], respectively.",
-            e.getCause.getMessage)
-          assertEquals(desc, 3, numPartitions(topic1))
-      }
-
-      // try #assignments < with the increase
-      alterResult = client.createPartitions(Map(topic1 ->
-        NewPartitions.increaseTo(6, asList(asList(1)))).asJava, option)
-      try {
-        altered = alterResult.values.get(topic1).get
-        fail(s"$desc: Expect InvalidReplicaAssignmentException when #assignments != newCount - oldCount")
-      } catch {
-        case e: ExecutionException =>
-          assertTrue(desc, e.getCause.isInstanceOf[InvalidReplicaAssignmentException])
-          assertEquals(desc, "Increasing the number of partitions by 3 but 1 assignments provided.", e.getCause.getMessage)
-          assertEquals(desc, 3, numPartitions(topic1))
-      }
-
-      // try #assignments > with the increase
-      alterResult = client.createPartitions(Map(topic1 ->
-        NewPartitions.increaseTo(4, asList(asList(1), asList(2)))).asJava, option)
-      try {
-        altered = alterResult.values.get(topic1).get
-        fail(s"$desc: Expect InvalidReplicaAssignmentException when #assignments != newCount - oldCount")
-      } catch {
-        case e: ExecutionException =>
-          assertTrue(desc, e.getCause.isInstanceOf[InvalidReplicaAssignmentException])
-          assertEquals(desc, "Increasing the number of partitions by 1 but 2 assignments provided.", e.getCause.getMessage)
-          assertEquals(desc, 3, numPartitions(topic1))
-      }
-
-      // try with duplicate brokers in assignments
-      alterResult = client.createPartitions(Map(topic1 ->
-        NewPartitions.increaseTo(4, asList(asList(1, 1)))).asJava, option)
-      try {
-        altered = alterResult.values.get(topic1).get
-        fail(s"$desc: Expect InvalidReplicaAssignmentException when assignments has duplicate brokers")
-      } catch {
-        case e: ExecutionException =>
-          assertTrue(desc, e.getCause.isInstanceOf[InvalidReplicaAssignmentException])
-          assertEquals(desc, "Duplicate brokers not allowed in replica assignment: 1, 1 for partition id 3.",
-            e.getCause.getMessage)
-          assertEquals(desc, 3, numPartitions(topic1))
-      }
-
-      // try assignments with differently sized inner lists
-      alterResult = client.createPartitions(Map(topic1 ->
-        NewPartitions.increaseTo(5, asList(asList(1), asList(1, 0)))).asJava, option)
-      try {
-        altered = alterResult.values.get(topic1).get
-        fail(s"$desc: Expect InvalidReplicaAssignmentException when assignments have differently sized inner lists")
-      } catch {
-        case e: ExecutionException =>
-          assertTrue(desc, e.getCause.isInstanceOf[InvalidReplicaAssignmentException])
-          assertEquals(desc, "Inconsistent replication factor between partitions, partition 0 has 1 " +
-            "while partitions [4] have replication factors [2], respectively.", e.getCause.getMessage)
-          assertEquals(desc, 3, numPartitions(topic1))
-      }
-
-      // try assignments with unknown brokers
-      alterResult = client.createPartitions(Map(topic1 ->
-        NewPartitions.increaseTo(4, asList(asList(12)))).asJava, option)
-      try {
-        altered = alterResult.values.get(topic1).get
-        fail(s"$desc: Expect InvalidReplicaAssignmentException when assignments contains an unknown broker")
-      } catch {
-        case e: ExecutionException =>
-          assertTrue(desc, e.getCause.isInstanceOf[InvalidReplicaAssignmentException])
-          assertEquals(desc, "Unknown broker(s) in replica assignment: 12.", e.getCause.getMessage)
-          assertEquals(desc, 3, numPartitions(topic1))
-      }
-
-      // try with empty assignments
-      alterResult = client.createPartitions(Map(topic1 ->
-        NewPartitions.increaseTo(4, Collections.emptyList())).asJava, option)
-      try {
-        altered = alterResult.values.get(topic1).get
-        fail(s"$desc: Expect InvalidReplicaAssignmentException when assignments is empty")
-      } catch {
-        case e: ExecutionException =>
-          assertTrue(desc, e.getCause.isInstanceOf[InvalidReplicaAssignmentException])
-          assertEquals(desc, "Increasing the number of partitions by 1 but 0 assignments provided.", e.getCause.getMessage)
-          assertEquals(desc, 3, numPartitions(topic1))
-      }
-    }
-
-    // a mixed success, failure response
-    alterResult = client.createPartitions(Map(
-      topic1 -> NewPartitions.increaseTo(4),
-      topic2 -> NewPartitions.increaseTo(2)).asJava, actuallyDoIt)
-    // assert that the topic1 now has 4 partitions
-    altered = alterResult.values.get(topic1).get
-    assertEquals(4, numPartitions(topic1))
-    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)
-        // assert that the topic2 still has 3 partitions
-        assertEquals(3, numPartitions(topic2))
-    }
-
-    // finally, try to add partitions to a topic queued for deletion
-    val deleteResult = client.deleteTopics(asList(topic1))
-    deleteResult.values.get(topic1).get
-    alterResult = client.createPartitions(Map(topic1 ->
-      NewPartitions.increaseTo(4)).asJava, validateOnly)
-    try {
-      altered = alterResult.values.get(topic1).get
-      fail("Expect InvalidTopicException when the topic is queued for deletion")
-    } catch {
-      case e: ExecutionException =>
-        assertTrue(e.getCause.isInstanceOf[InvalidTopicException])
-        assertEquals("The topic is queued for deletion.", e.getCause.getMessage)
-    }
-  }
-
-  @Test
-  def testSeekAfterDeleteRecords(): Unit = {
-    TestUtils.createTopic(zkUtils, topic, 2, serverCount, servers)
-
-    client = AdminClient.create(createConfig)
-
-    val consumer = consumers.head
-    subscribeAndWaitForAssignment(topic, consumer)
-
-    sendRecords(producers.head, 10, topicPartition)
-    consumer.seekToBeginning(Collections.singleton(topicPartition))
-    assertEquals(0L, consumer.position(topicPartition))
-
-    val result = client.deleteRecords(Map(topicPartition -> RecordsToDelete.beforeOffset(5L)).asJava)
-    val lowWatermark = result.lowWatermarks().get(topicPartition).get().lowWatermark()
-    assertEquals(5L, lowWatermark)
-
-    consumer.seekToBeginning(Collections.singletonList(topicPartition))
-    assertEquals(5L, consumer.position(topicPartition))
-
-    consumer.seek(topicPartition, 7L)
-    assertEquals(7L, consumer.position(topicPartition))
-
-    client.close()
-  }
-
-  @Test
-  @Ignore // Disabled temporarily until flakiness is resolved
-  def testLogStartOffsetCheckpoint(): Unit = {
-    TestUtils.createTopic(zkUtils, topic, 2, serverCount, servers)
-
-    client = AdminClient.create(createConfig)
-
-    subscribeAndWaitForAssignment(topic, consumers.head)
-
-    sendRecords(producers.head, 10, topicPartition)
-    var result = client.deleteRecords(Map(topicPartition -> RecordsToDelete.beforeOffset(5L)).asJava)
-    var lowWatermark = result.lowWatermarks().get(topicPartition).get().lowWatermark()
-    assertEquals(5L, lowWatermark)
-
-    for (i <- 0 until serverCount) {
-      killBroker(i)
-    }
-    restartDeadBrokers()
-
-    client.close()
-    brokerList = TestUtils.bootstrapServers(servers, listenerName)
-    client = AdminClient.create(createConfig)
-
-    TestUtils.waitUntilTrue(() => {
-      // Need to retry if leader is not available for the partition
-      result = client.deleteRecords(Map(topicPartition -> RecordsToDelete.beforeOffset(0L)).asJava)
-
-      val future = result.lowWatermarks().get(topicPartition)
-      try {
-        lowWatermark = future.get(1000L, TimeUnit.MILLISECONDS).lowWatermark()
-        lowWatermark == 5L
-      } catch {
-        case e: LeaderNotAvailableException => false
-      }
-
-    }, "Expected low watermark of the partition to be 5L")
-
-    client.close()
-  }
-
-  @Test
-  def testLogStartOffsetAfterDeleteRecords(): Unit = {
-    TestUtils.createTopic(zkUtils, topic, 2, serverCount, servers)
-
-    client = AdminClient.create(createConfig)
-
-    subscribeAndWaitForAssignment(topic, consumers.head)
-
-    sendRecords(producers.head, 10, topicPartition)
-    val result = client.deleteRecords(Map(topicPartition -> RecordsToDelete.beforeOffset(3L)).asJava)
-    val lowWatermark = result.lowWatermarks().get(topicPartition).get().lowWatermark()
-    assertEquals(3L, lowWatermark)
-
-    for (i <- 0 until serverCount)
-      assertEquals(3, servers(i).replicaManager.getReplica(topicPartition).get.logStartOffset)
-
-    client.close()
-  }
-
-  @Test
-  def testOffsetsForTimesAfterDeleteRecords(): Unit = {
-    TestUtils.createTopic(zkUtils, topic, 2, serverCount, servers)
-
-    client = AdminClient.create(createConfig)
-
-    val consumer = consumers.head
-    subscribeAndWaitForAssignment(topic, consumer)
-
-    sendRecords(producers.head, 10, topicPartition)
-    assertEquals(0L, consumer.offsetsForTimes(Map(topicPartition -> new JLong(0L)).asJava).get(topicPartition).offset())
-
-    var result = client.deleteRecords(Map(topicPartition -> RecordsToDelete.beforeOffset(5L)).asJava)
-    result.all().get()
-    assertEquals(5L, consumer.offsetsForTimes(Map(topicPartition -> new JLong(0L)).asJava).get(topicPartition).offset())
-
-    result = client.deleteRecords(Map(topicPartition -> RecordsToDelete.beforeOffset(DeleteRecordsRequest.HIGH_WATERMARK)).asJava)
-    result.all().get()
-    assertNull(consumer.offsetsForTimes(Map(topicPartition -> new JLong(0L)).asJava).get(topicPartition))
-
-    client.close()
-  }
-
-  private def subscribeAndWaitForAssignment(topic: String, consumer: KafkaConsumer[Array[Byte], Array[Byte]]): Unit = {
-    consumer.subscribe(Collections.singletonList(topic))
-    TestUtils.waitUntilTrue(() => {
-      consumer.poll(0)
-      !consumer.assignment.isEmpty
-    }, "Expected non-empty assignment")
-  }
-
-  private def sendRecords(producer: KafkaProducer[Array[Byte], Array[Byte]],
-                          numRecords: Int,
-                          topicPartition: TopicPartition): Unit = {
-    val futures = (0 until numRecords).map( i => {
-      val record = new ProducerRecord(topicPartition.topic, topicPartition.partition, s"$i".getBytes, s"$i".getBytes)
-      debug(s"Sending this record: $record")
-      producer.send(record)
-    })
-
-    futures.foreach(_.get)
-  }
-
-  @Test
-  def testInvalidAlterConfigs(): Unit = {
-    client = AdminClient.create(createConfig)
-    checkInvalidAlterConfigs(zkUtils, servers, client)
-  }
-
-  val ACL1 = new AclBinding(new Resource(ResourceType.TOPIC, "mytopic3"),
-      new AccessControlEntry("User:ANONYMOUS", "*", AclOperation.DESCRIBE, AclPermissionType.ALLOW))
-
-  /**
-   * Test that ACL operations are not possible when the authorizer is disabled.
-   * Also see {@link kafka.api.SaslSslAdminClientIntegrationTest} for tests of ACL operations
-   * when the authorizer is enabled.
-   */
-  @Test
-  def testAclOperations(): Unit = {
-    client = AdminClient.create(createConfig())
-    assertFutureExceptionTypeEquals(client.describeAcls(AclBindingFilter.ANY).values(), classOf[SecurityDisabledException])
-    assertFutureExceptionTypeEquals(client.createAcls(Collections.singleton(ACL1)).all(),
-        classOf[SecurityDisabledException])
-    assertFutureExceptionTypeEquals(client.deleteAcls(Collections.singleton(ACL1.toFilter())).all(),
-      classOf[SecurityDisabledException])
-    client.close()
-  }
-
-  /**
-    * Test closing the AdminClient with a generous timeout.  Calls in progress should be completed,
-    * since they can be done within the timeout.  New calls should receive timeouts.
-    */
-  @Test
-  def testDelayedClose(): Unit = {
-    client = AdminClient.create(createConfig())
-    val topics = Seq("mytopic", "mytopic2")
-    val newTopics = topics.map(new NewTopic(_, 1, 1))
-    val future = client.createTopics(newTopics.asJava, new CreateTopicsOptions().validateOnly(true)).all()
-    client.close(2, TimeUnit.HOURS)
-    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
-  }
-
-  /**
-    * Test closing the AdminClient with a timeout of 0, when there are calls with extremely long
-    * timeouts in progress.  The calls should be aborted after the hard shutdown timeout elapses.
-    */
-  @Test
-  def testForceClose(): Unit = {
-    val config = createConfig()
-    config.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:22")
-    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,
-      new CreateTopicsOptions().timeoutMs(900000)).all()
-    client.close(0, TimeUnit.MILLISECONDS)
-    assertFutureExceptionTypeEquals(future, classOf[TimeoutException])
-  }
-
-  /**
-    * Check that a call with a timeout does not complete before the minimum timeout has elapsed,
-    * even when the default request timeout is shorter.
-    */
-  @Test
-  def testMinimumRequestTimeouts(): Unit = {
-    val config = createConfig()
-    config.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:22")
-    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,
-      new CreateTopicsOptions().timeoutMs(2)).all()
-    assertFutureExceptionTypeEquals(future, classOf[TimeoutException])
-    val endTimeMs = Time.SYSTEM.milliseconds()
-    assertTrue("Expected the timeout to take at least one millisecond.", endTimeMs > startTimeMs);
-    client.close()
-  }
-
-  /**
-    * Test injecting timeouts for calls that are in flight.
-    */
-  @Test
-  def testCallInFlightTimeouts(): Unit = {
-    val config = createConfig()
-    config.put(AdminClientConfig.REQUEST_TIMEOUT_MS_CONFIG, "100000000")
-    val factory = new KafkaAdminClientTest.FailureInjectingTimeoutProcessorFactory()
-    val client = KafkaAdminClientTest.createInternal(new AdminClientConfig(config), factory)
-    val future = client.createTopics(Seq("mytopic", "mytopic2").map(new NewTopic(_, 1, 1)).asJava,
-        new CreateTopicsOptions().validateOnly(true)).all()
-    assertFutureExceptionTypeEquals(future, classOf[TimeoutException])
-    val future2 = client.createTopics(Seq("mytopic3", "mytopic4").map(new NewTopic(_, 1, 1)).asJava,
-      new CreateTopicsOptions().validateOnly(true)).all()
-    future2.get
-    client.close()
-    assertEquals(1, factory.failuresInjected)
-  }
-}
-
-object AdminClientIntegrationTest {
-
-  import org.scalatest.Assertions._
-
-  def checkValidAlterConfigs(zkUtils: ZkUtils, servers: Seq[KafkaServer], client: AdminClient,
-                             topicResource1: ConfigResource, topicResource2: ConfigResource): Unit = {
-    // Alter topics
-    var topicConfigEntries1 = Seq(
-      new ConfigEntry(LogConfig.FlushMsProp, "1000")
-    ).asJava
-
-    var topicConfigEntries2 = Seq(
-      new ConfigEntry(LogConfig.MinCleanableDirtyRatioProp, "0.9"),
-      new ConfigEntry(LogConfig.CompressionTypeProp, "lz4")
-    ).asJava
-
-    var alterResult = client.alterConfigs(Map(
-      topicResource1 -> new Config(topicConfigEntries1),
-      topicResource2 -> new Config(topicConfigEntries2)
-    ).asJava)
-
-    assertEquals(Set(topicResource1, topicResource2).asJava, alterResult.values.keySet)
-    alterResult.all.get
-
-    // Verify that topics were updated correctly
-    var describeResult = client.describeConfigs(Seq(topicResource1, topicResource2).asJava)
-    var configs = describeResult.all.get
-
-    assertEquals(2, configs.size)
-
-    assertEquals("1000", configs.get(topicResource1).get(LogConfig.FlushMsProp).value)
-    assertEquals(Defaults.MessageMaxBytes.toString,
-      configs.get(topicResource1).get(LogConfig.MaxMessageBytesProp).value)
-    assertEquals((Defaults.LogRetentionHours * 60 * 60 * 1000).toString,
-      configs.get(topicResource1).get(LogConfig.RetentionMsProp).value)
-
-    assertEquals("0.9", configs.get(topicResource2).get(LogConfig.MinCleanableDirtyRatioProp).value)
-    assertEquals("lz4", configs.get(topicResource2).get(LogConfig.CompressionTypeProp).value)
-
-    // Alter topics with validateOnly=true
-    topicConfigEntries1 = Seq(
-      new ConfigEntry(LogConfig.MaxMessageBytesProp, "10")
-    ).asJava
-
-    topicConfigEntries2 = Seq(
-      new ConfigEntry(LogConfig.MinCleanableDirtyRatioProp, "0.3")
-    ).asJava
-
-    alterResult = client.alterConfigs(Map(
-      topicResource1 -> new Config(topicConfigEntries1),
-      topicResource2 -> new Config(topicConfigEntries2)
-    ).asJava, new AlterConfigsOptions().validateOnly(true))
-
-    assertEquals(Set(topicResource1, topicResource2).asJava, alterResult.values.keySet)
-    alterResult.all.get
-
-    // Verify that topics were not updated due to validateOnly = true
-    describeResult = client.describeConfigs(Seq(topicResource1, topicResource2).asJava)
-    configs = describeResult.all.get
-
-    assertEquals(2, configs.size)
-
-    assertEquals(Defaults.MessageMaxBytes.toString,
-      configs.get(topicResource1).get(LogConfig.MaxMessageBytesProp).value)
-    assertEquals("0.9", configs.get(topicResource2).get(LogConfig.MinCleanableDirtyRatioProp).value)
-  }
-
-  def checkInvalidAlterConfigs(zkUtils: ZkUtils, servers: Seq[KafkaServer], client: AdminClient): Unit = {
-    // Create topics
-    val topic1 = "invalid-alter-configs-topic-1"
-    val topicResource1 = new ConfigResource(ConfigResource.Type.TOPIC, topic1)
-    TestUtils.createTopic(zkUtils, topic1, 1, 1, servers, new Properties())
-
-    val topic2 = "invalid-alter-configs-topic-2"
-    val topicResource2 = new ConfigResource(ConfigResource.Type.TOPIC, topic2)
-    TestUtils.createTopic(zkUtils, topic2, 1, 1, servers, new Properties)
-
-    val topicConfigEntries1 = Seq(
-      new ConfigEntry(LogConfig.MinCleanableDirtyRatioProp, "1.1"), // this value is invalid as it's above 1.0
-      new ConfigEntry(LogConfig.CompressionTypeProp, "lz4")
-    ).asJava
-
-    var topicConfigEntries2 = Seq(new ConfigEntry(LogConfig.CompressionTypeProp, "snappy")).asJava
-
-    val brokerResource = new ConfigResource(ConfigResource.Type.BROKER, servers.head.config.brokerId.toString)
-    val brokerConfigEntries = Seq(new ConfigEntry(KafkaConfig.CompressionTypeProp, "gzip")).asJava
-
-    // Alter configs: first and third are invalid, second is valid
-    var alterResult = client.alterConfigs(Map(
-      topicResource1 -> new Config(topicConfigEntries1),
-      topicResource2 -> new Config(topicConfigEntries2),
-      brokerResource -> new Config(brokerConfigEntries)
-    ).asJava)
-
-    assertEquals(Set(topicResource1, topicResource2, brokerResource).asJava, alterResult.values.keySet)
-    assertTrue(intercept[ExecutionException](alterResult.values.get(topicResource1).get).getCause.isInstanceOf[InvalidRequestException])
-    alterResult.values.get(topicResource2).get
-    assertTrue(intercept[ExecutionException](alterResult.values.get(brokerResource).get).getCause.isInstanceOf[InvalidRequestException])
-
-    // Verify that first and third resources were not updated and second was updated
-    var describeResult = client.describeConfigs(Seq(topicResource1, topicResource2, brokerResource).asJava)
-    var configs = describeResult.all.get
-    assertEquals(3, configs.size)
-
-    assertEquals(Defaults.LogCleanerMinCleanRatio.toString,
-      configs.get(topicResource1).get(LogConfig.MinCleanableDirtyRatioProp).value)
-    assertEquals(Defaults.CompressionType.toString,
-      configs.get(topicResource1).get(LogConfig.CompressionTypeProp).value)
-
-    assertEquals("snappy", configs.get(topicResource2).get(LogConfig.CompressionTypeProp).value)
-
-    assertEquals(Defaults.CompressionType.toString, configs.get(brokerResource).get(KafkaConfig.CompressionTypeProp).value)
-
-    // Alter configs with validateOnly = true: first and third are invalid, second is valid
-    topicConfigEntries2 = Seq(new ConfigEntry(LogConfig.CompressionTypeProp, "gzip")).asJava
-
-    alterResult = client.alterConfigs(Map(
-      topicResource1 -> new Config(topicConfigEntries1),
-      topicResource2 -> new Config(topicConfigEntries2),
-      brokerResource -> new Config(brokerConfigEntries)
-    ).asJava, new AlterConfigsOptions().validateOnly(true))
-
-    assertEquals(Set(topicResource1, topicResource2, brokerResource).asJava, alterResult.values.keySet)
-    assertTrue(intercept[ExecutionException](alterResult.values.get(topicResource1).get).getCause.isInstanceOf[InvalidRequestException])
-    alterResult.values.get(topicResource2).get
-    assertTrue(intercept[ExecutionException](alterResult.values.get(brokerResource).get).getCause.isInstanceOf[InvalidRequestException])
-
-    // Verify that no resources are updated since validate_only = true
-    describeResult = client.describeConfigs(Seq(topicResource1, topicResource2, brokerResource).asJava)
-    configs = describeResult.all.get
-    assertEquals(3, configs.size)
-
-    assertEquals(Defaults.LogCleanerMinCleanRatio.toString,
-      configs.get(topicResource1).get(LogConfig.MinCleanableDirtyRatioProp).value)
-    assertEquals(Defaults.CompressionType.toString,
-      configs.get(topicResource1).get(LogConfig.CompressionTypeProp).value)
-
-    assertEquals("snappy", configs.get(topicResource2).get(LogConfig.CompressionTypeProp).value)
-
-    assertEquals(Defaults.CompressionType.toString, configs.get(brokerResource).get(KafkaConfig.CompressionTypeProp).value)
-  }
-
-}
diff --git a/core/src/test/scala/integration/kafka/api/AdminClientWithPoliciesIntegrationTest.scala b/core/src/test/scala/integration/kafka/api/AdminClientWithPoliciesIntegrationTest.scala
index 8e8e8258cf23c..ea98f0b9c74d0 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, 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,8 +29,10 @@ 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._
+import scala.annotation.nowarn
+import scala.jdk.CollectionConverters._
 
 /**
   * Tests AdminClient calls when the broker is configured with policies like AlterConfigPolicy, CreateTopicPolicy, etc.
@@ -39,7 +41,7 @@ class AdminClientWithPoliciesIntegrationTest extends KafkaServerTestHarness with
 
   import AdminClientWithPoliciesIntegrationTest._
 
-  var client: AdminClient = null
+  var client: Admin = null
   val brokerCount = 3
 
   @Rule
@@ -47,7 +49,7 @@ class AdminClientWithPoliciesIntegrationTest extends KafkaServerTestHarness with
 
   @Before
   override def setUp(): Unit = {
-    super.setUp
+    super.setUp()
     TestUtils.waitUntilBrokerMetadataIsPropagated(servers)
   }
 
@@ -69,44 +71,45 @@ class AdminClientWithPoliciesIntegrationTest extends KafkaServerTestHarness with
 
   @Test
   def testValidAlterConfigs(): Unit = {
-    client = AdminClient.create(createConfig)
+    client = Admin.create(createConfig)
     // Create topics
     val topic1 = "describe-alter-configs-topic-1"
     val topicResource1 = new ConfigResource(ConfigResource.Type.TOPIC, topic1)
     val topicConfig1 = new Properties
     topicConfig1.setProperty(LogConfig.MaxMessageBytesProp, "500000")
     topicConfig1.setProperty(LogConfig.RetentionMsProp, "60000000")
-    TestUtils.createTopic(zkUtils, topic1, 1, 1, servers, topicConfig1)
+    createTopic(topic1, 1, 1, topicConfig1)
 
     val topic2 = "describe-alter-configs-topic-2"
     val topicResource2 = new ConfigResource(ConfigResource.Type.TOPIC, topic2)
-    TestUtils.createTopic(zkUtils, topic2, 1, 1, servers, new Properties)
+    createTopic(topic2, 1, 1)
 
-    AdminClientIntegrationTest.checkValidAlterConfigs(zkUtils, servers, client, topicResource1, topicResource2)
+    PlaintextAdminIntegrationTest.checkValidAlterConfigs(client, topicResource1, topicResource2)
   }
 
   @Test
   def testInvalidAlterConfigs(): Unit = {
-    client = AdminClient.create(createConfig)
-    AdminClientIntegrationTest.checkInvalidAlterConfigs(zkUtils, servers, client)
+    client = Admin.create(createConfig)
+    PlaintextAdminIntegrationTest.checkInvalidAlterConfigs(zkClient, servers, client)
   }
 
+  @nowarn("cat=deprecation")
   @Test
   def testInvalidAlterConfigsDueToPolicy(): Unit = {
-    client = AdminClient.create(createConfig)
+    client = Admin.create(createConfig)
 
     // Create topics
     val topic1 = "invalid-alter-configs-due-to-policy-topic-1"
     val topicResource1 = new ConfigResource(ConfigResource.Type.TOPIC, topic1)
-    TestUtils.createTopic(zkUtils, topic1, 1, 1, servers, new Properties())
+    createTopic(topic1, 1, 1)
 
     val topic2 = "invalid-alter-configs-due-to-policy-topic-2"
     val topicResource2 = new ConfigResource(ConfigResource.Type.TOPIC, topic2)
-    TestUtils.createTopic(zkUtils, topic2, 1, 1, servers, new Properties)
+    createTopic(topic2, 1, 1)
 
     val topic3 = "invalid-alter-configs-due-to-policy-topic-3"
     val topicResource3 = new ConfigResource(ConfigResource.Type.TOPIC, topic3)
-    TestUtils.createTopic(zkUtils, topic3, 1, 1, servers, new Properties)
+    createTopic(topic3, 1, 1)
 
     val topicConfigEntries1 = Seq(
       new ConfigEntry(LogConfig.MinCleanableDirtyRatioProp, "0.9"),
diff --git a/core/src/test/scala/integration/kafka/api/AuthorizerIntegrationTest.scala b/core/src/test/scala/integration/kafka/api/AuthorizerIntegrationTest.scala
index b701f13008e82..e8cce7abe8118 100644
--- a/core/src/test/scala/integration/kafka/api/AuthorizerIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/AuthorizerIntegrationTest.scala
@@ -12,55 +12,94 @@
  */
 package kafka.api
 
-import java.nio.ByteBuffer
+import java.time.Duration
 import java.util
 import java.util.concurrent.ExecutionException
 import java.util.regex.Pattern
-import java.util.{ArrayList, Collections, Properties}
+import java.util.{Collections, Optional, Properties}
 
-import kafka.admin.AdminClient
-import kafka.admin.ConsumerGroupCommand.ConsumerGroupCommandOptions
-import kafka.admin.ConsumerGroupCommand.KafkaConsumerGroupService
-import kafka.common.TopicAndPartition
+import kafka.admin.ConsumerGroupCommand.{ConsumerGroupCommandOptions, ConsumerGroupService}
 import kafka.log.LogConfig
-import kafka.network.SocketServer
-import kafka.security.auth._
+import kafka.security.authorizer.AclEntry
+import kafka.security.authorizer.AclEntry.WildcardHost
 import kafka.server.{BaseRequestTest, KafkaConfig}
 import kafka.utils.TestUtils
-import org.apache.kafka.clients.admin.NewPartitions
-import org.apache.kafka.clients.consumer.OffsetAndMetadata
-import org.apache.kafka.clients.consumer.internals.NoOpConsumerRebalanceListener
+import org.apache.kafka.clients.admin.{Admin, 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.AclOperation._
+import org.apache.kafka.common.acl.AclPermissionType.ALLOW
+import org.apache.kafka.common.acl.{AccessControlEntry, AccessControlEntryFilter, AclBindingFilter, AclOperation, AclPermissionType}
+import org.apache.kafka.common.config.internals.BrokerSecurityConfigs
+import org.apache.kafka.common.config.{ConfigResource, LogLevelConfig}
 import org.apache.kafka.common.errors._
 import org.apache.kafka.common.internals.Topic.GROUP_METADATA_TOPIC_NAME
-import org.apache.kafka.common.KafkaException
-import org.apache.kafka.common.protocol.{ApiKeys, Errors}
-import org.apache.kafka.common.requests.{Resource => RResource, ResourceType => RResourceType, _}
-import org.apache.kafka.common.acl.{AccessControlEntry, AccessControlEntryFilter, AclBinding, AclBindingFilter, AclOperation, AclPermissionType}
+import org.apache.kafka.common.message.{AddOffsetsToTxnRequestData, AlterPartitionReassignmentsRequestData, AlterReplicaLogDirsRequestData, ControlledShutdownRequestData, CreateAclsRequestData, CreatePartitionsRequestData, CreateTopicsRequestData, DeleteAclsRequestData, DeleteGroupsRequestData, DeleteRecordsRequestData, DeleteTopicsRequestData, DescribeConfigsRequestData, DescribeGroupsRequestData, DescribeLogDirsRequestData, FindCoordinatorRequestData, HeartbeatRequestData, IncrementalAlterConfigsRequestData, JoinGroupRequestData, ListPartitionReassignmentsRequestData, OffsetCommitRequestData, ProduceRequestData, SyncGroupRequestData}
+import org.apache.kafka.common.message.CreatePartitionsRequestData.CreatePartitionsTopic
+import org.apache.kafka.common.message.CreateTopicsRequestData.{CreatableTopic, CreatableTopicCollection}
+import org.apache.kafka.common.message.IncrementalAlterConfigsRequestData.{AlterConfigsResource, AlterableConfig, AlterableConfigCollection}
+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.ListOffsetRequestData.{ListOffsetPartition, ListOffsetTopic}
+import org.apache.kafka.common.message.OffsetForLeaderEpochRequestData.OffsetForLeaderPartition
+import org.apache.kafka.common.message.OffsetForLeaderEpochRequestData.OffsetForLeaderTopic
+import org.apache.kafka.common.message.OffsetForLeaderEpochRequestData.OffsetForLeaderTopicCollection
+import org.apache.kafka.common.message.StopReplicaRequestData.{StopReplicaPartitionState, StopReplicaTopicState}
+import org.apache.kafka.common.message.UpdateMetadataRequestData.{UpdateMetadataBroker, UpdateMetadataEndpoint, UpdateMetadataPartitionState}
 import org.apache.kafka.common.network.ListenerName
-import org.apache.kafka.common.record.{CompressionType, MemoryRecords, SimpleRecord}
-import org.apache.kafka.common.requests.CreateAclsRequest.AclCreation
-import org.apache.kafka.common.requests.CreateTopicsRequest.TopicDetails
-import org.apache.kafka.common.resource.{ResourceFilter, Resource => AdminResource, ResourceType => AdminResourceType}
-import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
-import org.apache.kafka.common.{Node, TopicPartition, requests}
+import org.apache.kafka.common.protocol.{ApiKeys, Errors}
+import org.apache.kafka.common.record.{CompressionType, MemoryRecords, RecordBatch, Records, SimpleRecord}
+import org.apache.kafka.common.requests._
+import org.apache.kafka.common.resource.PatternType.LITERAL
+import org.apache.kafka.common.resource.ResourceType._
+import org.apache.kafka.common.resource.{PatternType, Resource, ResourcePattern, ResourcePatternFilter, ResourceType}
+import org.apache.kafka.common.security.auth.{AuthenticationContext, KafkaPrincipal, KafkaPrincipalBuilder, SecurityProtocol}
+import org.apache.kafka.common.{ElectionType, IsolationLevel, Node, TopicPartition, requests}
+import org.apache.kafka.test.{TestUtils => JTestUtils}
 import org.junit.Assert._
-import org.junit.{After, Assert, Before, Test}
+import org.junit.{After, Before, Test}
+import org.scalatest.Assertions.intercept
 
-import scala.collection.JavaConverters._
+import scala.annotation.nowarn
+import scala.jdk.CollectionConverters._
 import scala.collection.mutable
 import scala.collection.mutable.Buffer
 
+object AuthorizerIntegrationTest {
+  val BrokerPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "broker")
+  val ClientPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "client")
+
+  val BrokerListenerName = "BROKER"
+  val ClientListenerName = "CLIENT"
+
+  class PrincipalBuilder extends KafkaPrincipalBuilder {
+    override def build(context: AuthenticationContext): KafkaPrincipal = {
+      context.listenerName match {
+        case BrokerListenerName => BrokerPrincipal
+        case ClientListenerName => ClientPrincipal
+        case listenerName => throw new IllegalArgumentException(s"No principal mapped to listener $listenerName")
+      }
+    }
+  }
+}
+
 class AuthorizerIntegrationTest extends BaseRequestTest {
+  import AuthorizerIntegrationTest._
 
-  override def numBrokers: Int = 1
-  val brokerId: Integer = 0
+  override def interBrokerListenerName: ListenerName = new ListenerName(BrokerListenerName)
+  override def listenerName: ListenerName = new ListenerName(ClientListenerName)
+  override def brokerCount: Int = 1
+
+  def clientPrincipal: KafkaPrincipal = ClientPrincipal
+  def brokerPrincipal: KafkaPrincipal = BrokerPrincipal
 
+  val clientPrincipalString: String = clientPrincipal.toString
+
+  val brokerId: Integer = 0
   val topic = "topic"
   val topicPattern = "topic.*"
-  val createTopic = "topic-new"
-  val deleteTopic = "topic-delete"
   val transactionalId = "transactional.id"
   val producerId = 83392L
   val part = 0
@@ -68,127 +107,130 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
   val clientId = "client-Id"
   val tp = new TopicPartition(topic, part)
   val logDir = "logDir"
-  val deleteRecordsPartition = new TopicPartition(deleteTopic, part)
-  val topicAndPartition = TopicAndPartition(topic, part)
   val group = "my-group"
-  val topicResource = new Resource(Topic, topic)
-  val groupResource = new Resource(Group, group)
-  val deleteTopicResource = new Resource(Topic, deleteTopic)
-  val transactionalIdResource = new Resource(TransactionalId, transactionalId)
-
-  val groupReadAcl = Map(groupResource -> Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Read)))
-  val groupDescribeAcl = Map(groupResource -> Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Describe)))
-  val clusterAcl = Map(Resource.ClusterResource -> Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, ClusterAction)))
-  val clusterCreateAcl = Map(Resource.ClusterResource -> Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Create)))
-  val clusterAlterAcl = Map(Resource.ClusterResource -> Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Alter)))
-  val clusterDescribeAcl = Map(Resource.ClusterResource -> Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Describe)))
-  val clusterIdempotentWriteAcl = Map(Resource.ClusterResource -> Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, IdempotentWrite)))
-  val topicReadAcl = Map(topicResource -> Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Read)))
-  val topicWriteAcl = Map(topicResource -> Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)))
-  val topicDescribeAcl = Map(topicResource -> Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Describe)))
-  val topicAlterAcl = Map(topicResource -> Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Alter)))
-  val topicDeleteAcl = Map(deleteTopicResource -> Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Delete)))
-  val topicDescribeConfigsAcl = Map(topicResource -> Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, DescribeConfigs)))
-  val topicAlterConfigsAcl = Map(topicResource -> Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, AlterConfigs)))
-  val transactionIdWriteAcl = Map(transactionalIdResource -> Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)))
-  val transactionalIdDescribeAcl = Map(transactionalIdResource -> Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Describe)))
-
-
-  val consumers = Buffer[KafkaConsumer[Array[Byte], Array[Byte]]]()
-  val producers = Buffer[KafkaProducer[Array[Byte], Array[Byte]]]()
-
-  val producerCount = 1
-  val consumerCount = 2
-  val producerConfig = new Properties
+  val protocolType = "consumer"
+  val protocolName = "consumer-range"
+  val clusterResource = new ResourcePattern(CLUSTER, Resource.CLUSTER_NAME, LITERAL)
+  val topicResource = new ResourcePattern(TOPIC, topic, LITERAL)
+  val groupResource = new ResourcePattern(GROUP, group, LITERAL)
+  val transactionalIdResource = new ResourcePattern(TRANSACTIONAL_ID, transactionalId, LITERAL)
+
+  val groupReadAcl = Map(groupResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)))
+  val groupDescribeAcl = Map(groupResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)))
+  val groupDeleteAcl = Map(groupResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DELETE, ALLOW)))
+  val clusterAcl = Map(clusterResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, CLUSTER_ACTION, ALLOW)))
+  val clusterCreateAcl = Map(clusterResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, CREATE, ALLOW)))
+  val clusterAlterAcl = Map(clusterResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, ALTER, ALLOW)))
+  val clusterDescribeAcl = Map(clusterResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)))
+  val clusterAlterConfigsAcl = Map(clusterResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, ALTER_CONFIGS, ALLOW)))
+  val clusterIdempotentWriteAcl = Map(clusterResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, IDEMPOTENT_WRITE, ALLOW)))
+  val topicCreateAcl = Map(topicResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, CREATE, ALLOW)))
+  val topicReadAcl = Map(topicResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)))
+  val topicWriteAcl = Map(topicResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)))
+  val topicDescribeAcl = Map(topicResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)))
+  val topicAlterAcl = Map(topicResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, ALTER, ALLOW)))
+  val topicDeleteAcl = Map(topicResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DELETE, ALLOW)))
+  val topicDescribeConfigsAcl = Map(topicResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE_CONFIGS, ALLOW)))
+  val topicAlterConfigsAcl = Map(topicResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, ALTER_CONFIGS, ALLOW)))
+  val transactionIdWriteAcl = Map(transactionalIdResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)))
+  val transactionalIdDescribeAcl = Map(transactionalIdResource -> Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)))
+
   val numRecords = 1
+  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 = {
-    properties.put(KafkaConfig.AuthorizerClassNameProp, classOf[SimpleAclAuthorizer].getName)
+  override def brokerPropertyOverrides(properties: Properties): Unit = {
+    properties.put(KafkaConfig.AuthorizerClassNameProp, "kafka.security.auth.SimpleAclAuthorizer")
     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")
+    properties.put(BrokerSecurityConfigs.PRINCIPAL_BUILDER_CLASS_CONFIG,
+      classOf[PrincipalBuilder].getName)
   }
 
-  val requestKeyToResponseDeserializer: Map[ApiKeys, Class[_ <: Any]] =
-    Map(ApiKeys.METADATA -> classOf[requests.MetadataResponse],
-      ApiKeys.PRODUCE -> classOf[requests.ProduceResponse],
-      ApiKeys.FETCH -> classOf[requests.FetchResponse],
-      ApiKeys.LIST_OFFSETS -> classOf[requests.ListOffsetResponse],
-      ApiKeys.OFFSET_COMMIT -> classOf[requests.OffsetCommitResponse],
-      ApiKeys.OFFSET_FETCH -> classOf[requests.OffsetFetchResponse],
-      ApiKeys.FIND_COORDINATOR -> classOf[FindCoordinatorResponse],
-      ApiKeys.UPDATE_METADATA -> classOf[requests.UpdateMetadataResponse],
-      ApiKeys.JOIN_GROUP -> classOf[JoinGroupResponse],
-      ApiKeys.SYNC_GROUP -> classOf[SyncGroupResponse],
-      ApiKeys.DESCRIBE_GROUPS -> classOf[DescribeGroupsResponse],
-      ApiKeys.HEARTBEAT -> classOf[HeartbeatResponse],
-      ApiKeys.LEAVE_GROUP -> classOf[LeaveGroupResponse],
-      ApiKeys.LEADER_AND_ISR -> classOf[requests.LeaderAndIsrResponse],
-      ApiKeys.STOP_REPLICA -> classOf[requests.StopReplicaResponse],
-      ApiKeys.CONTROLLED_SHUTDOWN -> classOf[requests.ControlledShutdownResponse],
-      ApiKeys.CREATE_TOPICS -> classOf[CreateTopicsResponse],
-      ApiKeys.DELETE_TOPICS -> classOf[requests.DeleteTopicsResponse],
-      ApiKeys.DELETE_RECORDS -> classOf[requests.DeleteRecordsResponse],
-      ApiKeys.OFFSET_FOR_LEADER_EPOCH -> classOf[OffsetsForLeaderEpochResponse],
-      ApiKeys.DESCRIBE_CONFIGS -> classOf[DescribeConfigsResponse],
-      ApiKeys.ALTER_CONFIGS -> classOf[AlterConfigsResponse],
-      ApiKeys.INIT_PRODUCER_ID -> classOf[InitProducerIdResponse],
-      ApiKeys.WRITE_TXN_MARKERS -> classOf[WriteTxnMarkersResponse],
-      ApiKeys.ADD_PARTITIONS_TO_TXN -> classOf[AddPartitionsToTxnResponse],
-      ApiKeys.ADD_OFFSETS_TO_TXN -> classOf[AddOffsetsToTxnResponse],
-      ApiKeys.END_TXN -> classOf[EndTxnResponse],
-      ApiKeys.TXN_OFFSET_COMMIT -> classOf[TxnOffsetCommitResponse],
-      ApiKeys.CREATE_ACLS -> classOf[CreateAclsResponse],
-      ApiKeys.DELETE_ACLS -> classOf[DeleteAclsResponse],
-      ApiKeys.DESCRIBE_ACLS -> classOf[DescribeAclsResponse],
-      ApiKeys.ALTER_REPLICA_LOG_DIRS -> classOf[AlterReplicaLogDirsResponse],
-      ApiKeys.DESCRIBE_LOG_DIRS -> classOf[DescribeLogDirsResponse],
-      ApiKeys.CREATE_PARTITIONS -> classOf[CreatePartitionsResponse]
-  )
-
+  @nowarn("cat=deprecation")
   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) => 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.FETCH -> ((resp: requests.FetchResponse[Records]) => resp.responseData.asScala.find(_._1 == tp).get._2.error),
+    ApiKeys.LIST_OFFSETS -> ((resp: requests.ListOffsetResponse) => {
+      Errors.forCode(
+        resp.data
+          .topics.asScala.find(_.name == topic).get
+          .partitions.asScala.find(_.partitionIndex == part).get
+          .errorCode
+      )
+    }),
+    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.DESCRIBE_GROUPS -> ((resp: DescribeGroupsResponse) => resp.groups.get(group).error),
+    ApiKeys.SYNC_GROUP -> ((resp: SyncGroupResponse) => Errors.forCode(resp.data.errorCode())),
+    ApiKeys.DESCRIBE_GROUPS -> ((resp: DescribeGroupsResponse) => {
+      Errors.forCode(resp.data.groups.asScala.find(g => group == g.groupId).head.errorCode)
+    }),
     ApiKeys.HEARTBEAT -> ((resp: HeartbeatResponse) => resp.error),
     ApiKeys.LEAVE_GROUP -> ((resp: LeaveGroupResponse) => resp.error),
-    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.DELETE_GROUPS -> ((resp: DeleteGroupsResponse) => resp.get(group)),
+    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) => resp.errors.asScala.find(_._1 == createTopic).get._2.error),
-    ApiKeys.DELETE_TOPICS -> ((resp: requests.DeleteTopicsResponse) => resp.errors.asScala.find(_._1 == deleteTopic).get._2),
-    ApiKeys.DELETE_RECORDS -> ((resp: requests.DeleteRecordsResponse) => resp.responses.get(deleteRecordsPartition).error),
-    ApiKeys.OFFSET_FOR_LEADER_EPOCH -> ((resp: OffsetsForLeaderEpochResponse) => resp.responses.get(tp).error),
+    ApiKeys.CREATE_TOPICS -> ((resp: CreateTopicsResponse) => Errors.forCode(resp.data.topics.find(topic).errorCode())),
+    ApiKeys.DELETE_TOPICS -> ((resp: requests.DeleteTopicsResponse) => Errors.forCode(resp.data.responses.find(topic).errorCode())),
+    ApiKeys.DELETE_RECORDS -> ((resp: requests.DeleteRecordsResponse) => Errors.forCode(
+      resp.data.topics.find(tp.topic).partitions.find(tp.partition).errorCode)),
+    ApiKeys.OFFSET_FOR_LEADER_EPOCH -> ((resp: OffsetsForLeaderEpochResponse) => Errors.forCode(
+      resp.data.topics.find(tp.topic).partitions.asScala.find(_.partition == tp.partition).get.errorCode)),
     ApiKeys.DESCRIBE_CONFIGS -> ((resp: DescribeConfigsResponse) =>
-      resp.configs.get(new RResource(RResourceType.TOPIC, tp.topic)).error.error),
+      Errors.forCode(resp.resultMap.get(new ConfigResource(ConfigResource.Type.TOPIC, tp.topic)).errorCode)),
     ApiKeys.ALTER_CONFIGS -> ((resp: AlterConfigsResponse) =>
-      resp.errors.get(new RResource(RResourceType.TOPIC, tp.topic)).error),
+      resp.errors.get(new ConfigResource(ConfigResource.Type.TOPIC, tp.topic)).error),
     ApiKeys.INIT_PRODUCER_ID -> ((resp: InitProducerIdResponse) => resp.error),
-    ApiKeys.WRITE_TXN_MARKERS -> ((resp: WriteTxnMarkersResponse) => resp.errors(producerId).get(tp)),
+    ApiKeys.WRITE_TXN_MARKERS -> ((resp: WriteTxnMarkersResponse) => resp.errorsByProducerId.get(producerId).get(tp)),
     ApiKeys.ADD_PARTITIONS_TO_TXN -> ((resp: AddPartitionsToTxnResponse) => resp.errors.get(tp)),
-    ApiKeys.ADD_OFFSETS_TO_TXN -> ((resp: AddOffsetsToTxnResponse) => resp.error),
+    ApiKeys.ADD_OFFSETS_TO_TXN -> ((resp: AddOffsetsToTxnResponse) => Errors.forCode(resp.data.errorCode)),
     ApiKeys.END_TXN -> ((resp: EndTxnResponse) => resp.error),
     ApiKeys.TXN_OFFSET_COMMIT -> ((resp: TxnOffsetCommitResponse) => resp.errors.get(tp)),
-    ApiKeys.CREATE_ACLS -> ((resp: CreateAclsResponse) => resp.aclCreationResponses.asScala.head.error.error),
+    ApiKeys.CREATE_ACLS -> ((resp: CreateAclsResponse) => Errors.forCode(resp.results.asScala.head.errorCode)),
     ApiKeys.DESCRIBE_ACLS -> ((resp: DescribeAclsResponse) => resp.error.error),
-    ApiKeys.DELETE_ACLS -> ((resp: DeleteAclsResponse) => resp.responses.asScala.head.error.error),
-    ApiKeys.ALTER_REPLICA_LOG_DIRS -> ((resp: AlterReplicaLogDirsResponse) => resp.responses.get(tp)),
+    ApiKeys.DELETE_ACLS -> ((resp: DeleteAclsResponse) => Errors.forCode(resp.filterResults.asScala.head.errorCode)),
+    ApiKeys.ALTER_REPLICA_LOG_DIRS -> ((resp: AlterReplicaLogDirsResponse) => Errors.forCode(resp.data.results.asScala
+      .find(x => x.topicName == tp.topic).get.partitions.asScala
+      .find(p => p.partitionIndex == tp.partition).get.errorCode)),
     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)
+      if (resp.data.results.size() > 0) Errors.forCode(resp.data.results.get(0).errorCode) else Errors.CLUSTER_AUTHORIZATION_FAILED),
+    ApiKeys.CREATE_PARTITIONS -> ((resp: CreatePartitionsResponse) => Errors.forCode(resp.data.results.asScala.head.errorCode())),
+    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,
@@ -202,13 +244,14 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
     ApiKeys.DESCRIBE_GROUPS -> groupDescribeAcl,
     ApiKeys.HEARTBEAT -> groupReadAcl,
     ApiKeys.LEAVE_GROUP -> groupReadAcl,
+    ApiKeys.DELETE_GROUPS -> groupDeleteAcl,
     ApiKeys.LEADER_AND_ISR -> clusterAcl,
     ApiKeys.STOP_REPLICA -> clusterAcl,
     ApiKeys.CONTROLLED_SHUTDOWN -> clusterAcl,
-    ApiKeys.CREATE_TOPICS -> clusterCreateAcl,
+    ApiKeys.CREATE_TOPICS -> topicCreateAcl,
     ApiKeys.DELETE_TOPICS -> topicDeleteAcl,
     ApiKeys.DELETE_RECORDS -> topicDeleteAcl,
-    ApiKeys.OFFSET_FOR_LEADER_EPOCH -> clusterAcl,
+    ApiKeys.OFFSET_FOR_LEADER_EPOCH -> topicDescribeAcl,
     ApiKeys.DESCRIBE_CONFIGS -> topicDescribeConfigsAcl,
     ApiKeys.ALTER_CONFIGS -> topicAlterConfigsAcl,
     ApiKeys.INIT_PRODUCER_ID -> (transactionIdWriteAcl ++ clusterIdempotentWriteAcl),
@@ -222,41 +265,28 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
     ApiKeys.DELETE_ACLS -> clusterAlterAcl,
     ApiKeys.ALTER_REPLICA_LOG_DIRS -> clusterAlterAcl,
     ApiKeys.DESCRIBE_LOG_DIRS -> clusterDescribeAcl,
-    ApiKeys.CREATE_PARTITIONS -> topicAlterAcl
-
+    ApiKeys.CREATE_PARTITIONS -> topicAlterAcl,
+    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() {
-    super.setUp()
-
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, ClusterAction)), Resource.ClusterResource)
+  override def setUp(): Unit = {
+    doSetup(createOffsetsTopic = false)
 
-    for (_ <- 0 until producerCount)
-      producers += TestUtils.createNewProducer(TestUtils.getBrokerListStrFromServers(servers),
-        maxBlockMs = 3000,
-        acks = 1)
+    // Allow inter-broker communication
+    addAndVerifyAcls(Set(new AccessControlEntry(brokerPrincipal.toString, WildcardHost, CLUSTER_ACTION, ALLOW)), clusterResource)
 
-    for (_ <- 0 until consumerCount)
-      consumers += TestUtils.createNewConsumer(TestUtils.getBrokerListStrFromServers(servers), groupId = group, securityProtocol = SecurityProtocol.PLAINTEXT)
-
-    // create the consumer offset topic
-    TestUtils.createTopic(zkUtils, GROUP_METADATA_TOPIC_NAME,
-      1,
-      1,
-      servers,
-      servers.head.groupCoordinator.offsetsTopicConfigs)
-    // create the test topic with all the brokers as replicas
-    TestUtils.createTopic(zkUtils, topic, 1, 1, this.servers)
-    TestUtils.createTopic(zkUtils, deleteTopic, 1, 1, this.servers)
+    TestUtils.createOffsetsTopic(zkClient, servers)
   }
 
   @After
-  override def tearDown() = {
-    producers.foreach(_.close())
-    consumers.foreach(_.wakeup())
-    consumers.foreach(_.close())
-    removeAllAcls()
+  override def tearDown(): Unit = {
+    adminClients.foreach(_.close())
+    removeAllClientAcls()
     super.tearDown()
   }
 
@@ -264,149 +294,351 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
     new requests.MetadataRequest.Builder(List(topic).asJava, allowAutoTopicCreation).build()
   }
 
-  private def createProduceRequest = {
-    requests.ProduceRequest.Builder.forCurrentMagic(1, 5000,
-      collection.mutable.Map(tp -> MemoryRecords.withRecords(CompressionType.NONE, new SimpleRecord("test".getBytes))).asJava).
-      build()
-  }
+  private def createProduceRequest =
+    requests.ProduceRequest.forCurrentMagic(new ProduceRequestData()
+      .setTopicData(new ProduceRequestData.TopicProduceDataCollection(
+        Collections.singletonList(new ProduceRequestData.TopicProduceData()
+          .setName(tp.topic()).setPartitionData(Collections.singletonList(
+          new ProduceRequestData.PartitionProduceData()
+            .setIndex(tp.partition())
+            .setRecords(MemoryRecords.withRecords(CompressionType.NONE, new SimpleRecord("test".getBytes))))))
+        .iterator))
+      .setAcks(1.toShort)
+      .setTimeoutMs(5000))
+      .build()
 
   private def createFetchRequest = {
     val partitionMap = new util.LinkedHashMap[TopicPartition, requests.FetchRequest.PartitionData]
-    partitionMap.put(tp, new requests.FetchRequest.PartitionData(0, 0, 100))
+    partitionMap.put(tp, new requests.FetchRequest.PartitionData(0, 0, 100, Optional.of(27)))
+    requests.FetchRequest.Builder.forConsumer(100, Int.MaxValue, partitionMap).build()
+  }
+
+  private def createFetchFollowerRequest = {
+    val partitionMap = new util.LinkedHashMap[TopicPartition, requests.FetchRequest.PartitionData]
+    partitionMap.put(tp, new requests.FetchRequest.PartitionData(0, 0, 100, Optional.of(27)))
     val version = ApiKeys.FETCH.latestVersion
     requests.FetchRequest.Builder.forReplica(version, 5000, 100, Int.MaxValue, partitionMap).build()
   }
 
   private def createListOffsetsRequest = {
     requests.ListOffsetRequest.Builder.forConsumer(false, IsolationLevel.READ_UNCOMMITTED).setTargetTimes(
-      Map(tp -> (0L: java.lang.Long)).asJava).
+      List(new ListOffsetTopic()
+        .setName(tp.topic)
+        .setPartitions(List(new ListOffsetPartition()
+          .setPartitionIndex(tp.partition)
+          .setTimestamp(0L)
+          .setCurrentLeaderEpoch(27)).asJava)).asJava
+      ).
       build()
   }
 
-  private def offsetsForLeaderEpochRequest = {
-    new OffsetsForLeaderEpochRequest.Builder().add(tp, 7).build()
+  private def offsetsForLeaderEpochRequest: OffsetsForLeaderEpochRequest = {
+    val epochs = new OffsetForLeaderTopicCollection()
+    epochs.add(new OffsetForLeaderTopic()
+      .setTopic(tp.topic())
+      .setPartitions(List(new OffsetForLeaderPartition()
+          .setPartition(tp.partition())
+          .setLeaderEpoch(7)
+          .setCurrentLeaderEpoch(27)).asJava))
+    OffsetsForLeaderEpochRequest.Builder.forConsumer(epochs).build()
   }
 
   private def createOffsetFetchRequest = {
-    new requests.OffsetFetchRequest.Builder(group, List(tp).asJava).build()
+    new requests.OffsetFetchRequest.Builder(group, false, List(tp).asJava, false).build()
   }
 
   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, partitionState, brokers).build()
+    new requests.UpdateMetadataRequest.Builder(version, brokerId, Int.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(protocolName)
+        .setMetadata("test".getBytes())
+    ).iterator())
+
+    new JoinGroupRequest.Builder(
+      new JoinGroupRequestData()
+        .setGroupId(group)
+        .setSessionTimeoutMs(10000)
+        .setMemberId(JoinGroupRequest.UNKNOWN_MEMBER_ID)
+        .setGroupInstanceId(null)
+        .setProtocolType(protocolType)
+        .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)
+        .setProtocolType(protocolType)
+        .setProtocolName(protocolName)
+        .setAssignments(Collections.emptyList())
+    ).build()
   }
 
   private def createDescribeGroupsRequest = {
-    new DescribeGroupsRequest.Builder(List(group).asJava).build()
+    new DescribeGroupsRequest.Builder(new DescribeGroupsRequestData().setGroups(List(group).asJava)).build()
   }
 
   private def createOffsetCommitRequest = {
     new requests.OffsetCommitRequest.Builder(
-      group, Map(tp -> new requests.OffsetCommitRequest.PartitionData(0, "metadata")).asJava).
-      setMemberId("").setGenerationId(1).setRetentionTime(1000).
-      build()
-  }
-
-  private def createPartitionsRequest = {
-    new CreatePartitionsRequest.Builder(
-      Map(topic -> NewPartitions.increaseTo(10)).asJava, 10000, true
+        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 heartbeatRequest = new HeartbeatRequest.Builder(group, 1, "").build()
-
-  private def leaveGroupRequest = new LeaveGroupRequest.Builder(group, "").build()
-
-  private def leaderAndIsrRequest = {
-    new requests.LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, brokerId, Int.MaxValue,
-      Map(tp -> new LeaderAndIsrRequest.PartitionState(Int.MaxValue, brokerId, Int.MaxValue, List(brokerId).asJava, 2, Seq(brokerId).asJava, false)).asJava,
+  private def createPartitionsRequest = {
+    val partitionTopic = new CreatePartitionsTopic()
+      .setName(topic)
+      .setCount(10)
+      .setAssignments(null)
+    val data = new CreatePartitionsRequestData()
+      .setTimeoutMs(10000)
+      .setValidateOnly(true)
+    data.topics().add(partitionTopic)
+    new CreatePartitionsRequest.Builder(data).build(0.toShort)
+  }
+
+  private def heartbeatRequest = new HeartbeatRequest.Builder(
+    new HeartbeatRequestData()
+      .setGroupId(group)
+      .setGenerationId(1)
+      .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(
+    new DeleteGroupsRequestData()
+      .setGroupsNames(Collections.singletonList(group))
+  ).build()
+
+  private def leaderAndIsrRequest: LeaderAndIsrRequest = {
+    new requests.LeaderAndIsrRequest.Builder(ApiKeys.LEADER_AND_ISR.latestVersion, brokerId, Int.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(brokerId, Int.MaxValue, true, Set(tp).asJava).build()
+  private def stopReplicaRequest: StopReplicaRequest = {
+    val topicStates = Seq(
+      new StopReplicaTopicState()
+        .setTopicName(tp.topic())
+        .setPartitionStates(Seq(new StopReplicaPartitionState()
+          .setPartitionIndex(tp.partition())
+          .setLeaderEpoch(LeaderAndIsr.initialLeaderEpoch + 2)
+          .setDeletePartition(true)).asJava)
+    ).asJava
+    new StopReplicaRequest.Builder(ApiKeys.STOP_REPLICA.latestVersion, brokerId, Int.MaxValue,
+      Long.MaxValue, false, topicStates).build()
+  }
 
-  private def controlledShutdownRequest = new requests.ControlledShutdownRequest.Builder(brokerId,
-    ApiKeys.CONTROLLED_SHUTDOWN.latestVersion).build()
+  private def controlledShutdownRequest: ControlledShutdownRequest = {
+    new ControlledShutdownRequest.Builder(
+      new ControlledShutdownRequestData()
+        .setBrokerId(brokerId)
+        .setBrokerEpoch(Long.MaxValue),
+      ApiKeys.CONTROLLED_SHUTDOWN.latestVersion).build()
+  }
 
-  private def createTopicsRequest =
-    new CreateTopicsRequest.Builder(Map(createTopic -> new TopicDetails(1, 1.toShort)).asJava, 0).build()
+  private def createTopicsRequest: CreateTopicsRequest = {
+    new CreateTopicsRequest.Builder(new CreateTopicsRequestData().setTopics(
+      new CreatableTopicCollection(Collections.singleton(new CreatableTopic().
+        setName(topic).setNumPartitions(1).
+        setReplicationFactor(1.toShort)).iterator))).build()
+  }
 
-  private def deleteTopicsRequest = new DeleteTopicsRequest.Builder(Set(deleteTopic).asJava, 5000).build()
+  private def deleteTopicsRequest: DeleteTopicsRequest = {
+    new DeleteTopicsRequest.Builder(
+      new DeleteTopicsRequestData()
+        .setTopicNames(Collections.singletonList(topic))
+        .setTimeoutMs(5000)).build()
+  }
 
-  private def deleteRecordsRequest = new DeleteRecordsRequest.Builder(5000, Collections.singletonMap(deleteRecordsPartition, 0L)).build()
+  private def deleteRecordsRequest = new DeleteRecordsRequest.Builder(
+    new DeleteRecordsRequestData()
+      .setTimeoutMs(5000)
+      .setTopics(Collections.singletonList(new DeleteRecordsRequestData.DeleteRecordsTopic()
+        .setName(tp.topic)
+        .setPartitions(Collections.singletonList(new DeleteRecordsRequestData.DeleteRecordsPartition()
+          .setPartitionIndex(tp.partition)
+          .setOffset(0L)))))).build()
 
   private def describeConfigsRequest =
-    new DescribeConfigsRequest.Builder(Collections.singleton(new RResource(RResourceType.TOPIC, tp.topic))).build()
+    new DescribeConfigsRequest.Builder(new DescribeConfigsRequestData().setResources(Collections.singletonList(
+      new DescribeConfigsRequestData.DescribeConfigsResource().setResourceType(ConfigResource.Type.TOPIC.id)
+        .setResourceName(tp.topic)))).build()
 
   private def alterConfigsRequest =
     new AlterConfigsRequest.Builder(
-      Collections.singletonMap(new RResource(RResourceType.TOPIC, tp.topic),
+      Collections.singletonMap(new ConfigResource(ConfigResource.Type.TOPIC, tp.topic),
         new AlterConfigsRequest.Config(Collections.singleton(
           new AlterConfigsRequest.ConfigEntry(LogConfig.MaxMessageBytesProp, "1000000")
         ))), true).build()
 
-  private def describeAclsRequest = new DescribeAclsRequest.Builder(AclBindingFilter.ANY).build()
-
-  private def createAclsRequest = new CreateAclsRequest.Builder(
-    Collections.singletonList(new AclCreation(new AclBinding(
-      new AdminResource(AdminResourceType.TOPIC, "mytopic"),
-      new AccessControlEntry("User:ANONYMOUS", "*", AclOperation.WRITE, AclPermissionType.DENY))))).build()
-
-  private def deleteAclsRequest = new DeleteAclsRequest.Builder(
-    Collections.singletonList(new AclBindingFilter(
-      new ResourceFilter(AdminResourceType.TOPIC, null),
-      new AccessControlEntryFilter("User:ANONYMOUS", "*", AclOperation.ANY, AclPermissionType.DENY)))).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 alterReplicaLogDirsRequest = new AlterReplicaLogDirsRequest.Builder(Collections.singletonMap(tp, logDir)).build()
+  private def describeAclsRequest = new DescribeAclsRequest.Builder(AclBindingFilter.ANY).build()
 
-  private def describeLogDirsRequest = new DescribeLogDirsRequest.Builder(Collections.singleton(tp)).build()
+  private def createAclsRequest: CreateAclsRequest = new CreateAclsRequest.Builder(
+    new CreateAclsRequestData().setCreations(Collections.singletonList(
+      new CreateAclsRequestData.AclCreation()
+        .setResourceType(ResourceType.TOPIC.code)
+        .setResourceName("mytopic")
+        .setResourcePatternType(PatternType.LITERAL.code)
+        .setPrincipal(clientPrincipalString)
+        .setHost("*")
+        .setOperation(AclOperation.WRITE.code)
+        .setPermissionType(AclPermissionType.DENY.code)))
+  ).build()
+
+  private def deleteAclsRequest: DeleteAclsRequest = new DeleteAclsRequest.Builder(
+    new DeleteAclsRequestData().setFilters(Collections.singletonList(
+      new DeleteAclsRequestData.DeleteAclsFilter()
+        .setResourceTypeFilter(ResourceType.TOPIC.code)
+        .setResourceNameFilter(null)
+        .setPatternTypeFilter(PatternType.LITERAL.code)
+        .setPrincipalFilter(clientPrincipalString)
+        .setHostFilter("*")
+        .setOperation(AclOperation.ANY.code)
+        .setPermissionType(AclPermissionType.DENY.code)))
+  ).build()
+
+  private def alterReplicaLogDirsRequest = {
+    val dir = new AlterReplicaLogDirsRequestData.AlterReplicaLogDir()
+      .setPath(logDir)
+    dir.topics.add(new AlterReplicaLogDirsRequestData.AlterReplicaLogDirTopic()
+      .setName(tp.topic)
+      .setPartitions(Collections.singletonList(tp.partition)))
+    val data = new AlterReplicaLogDirsRequestData();
+    data.dirs.add(dir)
+    new AlterReplicaLogDirsRequest.Builder(data).build()
+  }
+
+  private def describeLogDirsRequest = new DescribeLogDirsRequest.Builder(new DescribeLogDirsRequestData().setTopics(new DescribeLogDirsRequestData.DescribableLogDirTopicCollection(Collections.singleton(
+    new DescribeLogDirsRequestData.DescribableLogDirTopic().setTopic(tp.topic).setPartitionIndex(Collections.singletonList(tp.partition))).iterator()))).build()
 
   private def addPartitionsToTxnRequest = new AddPartitionsToTxnRequest.Builder(transactionalId, 1, 1, Collections.singletonList(tp)).build()
 
-  private def addOffsetsToTxnRequest = new AddOffsetsToTxnRequest.Builder(transactionalId, 1, 1, group).build()
-
+  private def addOffsetsToTxnRequest = new AddOffsetsToTxnRequest.Builder(
+    new AddOffsetsToTxnRequestData()
+      .setTransactionalId(transactionalId)
+      .setProducerId(1)
+      .setProducerEpoch(1)
+      .setGroupId(group)
+  ).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](
+      // First create the topic
+      ApiKeys.CREATE_TOPICS -> createTopicsRequest,
+
       ApiKeys.METADATA -> createMetadataRequest(allowAutoTopicCreation = true),
       ApiKeys.PRODUCE -> createProduceRequest,
       ApiKeys.FETCH -> createFetchRequest,
       ApiKeys.LIST_OFFSETS -> createListOffsetsRequest,
       ApiKeys.OFFSET_FETCH -> createOffsetFetchRequest,
       ApiKeys.FIND_COORDINATOR -> createFindCoordinatorRequest,
-      ApiKeys.UPDATE_METADATA -> createUpdateMetadataRequest,
       ApiKeys.JOIN_GROUP -> createJoinGroupRequest,
       ApiKeys.SYNC_GROUP -> createSyncGroupRequest,
       ApiKeys.DESCRIBE_GROUPS -> createDescribeGroupsRequest,
       ApiKeys.OFFSET_COMMIT -> createOffsetCommitRequest,
       ApiKeys.HEARTBEAT -> heartbeatRequest,
       ApiKeys.LEAVE_GROUP -> leaveGroupRequest,
-      ApiKeys.LEADER_AND_ISR -> leaderAndIsrRequest,
-      ApiKeys.STOP_REPLICA -> stopReplicaRequest,
-      ApiKeys.CONTROLLED_SHUTDOWN -> controlledShutdownRequest,
-      ApiKeys.CREATE_TOPICS -> createTopicsRequest,
-      ApiKeys.DELETE_TOPICS -> deleteTopicsRequest,
       ApiKeys.DELETE_RECORDS -> deleteRecordsRequest,
       ApiKeys.OFFSET_FOR_LEADER_EPOCH -> offsetsForLeaderEpochRequest,
       ApiKeys.DESCRIBE_CONFIGS -> describeConfigsRequest,
@@ -418,26 +650,39 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
       ApiKeys.DESCRIBE_LOG_DIRS -> describeLogDirsRequest,
       ApiKeys.CREATE_PARTITIONS -> createPartitionsRequest,
       ApiKeys.ADD_PARTITIONS_TO_TXN -> addPartitionsToTxnRequest,
-      ApiKeys.ADD_OFFSETS_TO_TXN -> addOffsetsToTxnRequest
+      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,
+
+      // Inter-broker APIs use an invalid broker epoch, so does not affect the test case
+      ApiKeys.UPDATE_METADATA -> createUpdateMetadataRequest,
+      ApiKeys.LEADER_AND_ISR -> leaderAndIsrRequest,
+      ApiKeys.STOP_REPLICA -> stopReplicaRequest,
+      ApiKeys.CONTROLLED_SHUTDOWN -> controlledShutdownRequest,
+
+      // Delete the topic last
+      ApiKeys.DELETE_TOPICS -> deleteTopicsRequest
     )
 
     for ((key, request) <- requestKeyToRequest) {
-      removeAllAcls()
+      removeAllClientAcls()
       val resources = requestKeysToAcls(key).map(_._1.resourceType).toSet
-      sendRequestAndVerifyResponseError(key, request, resources, isAuthorized = false)
+      sendRequestAndVerifyResponseError(request, resources, isAuthorized = false)
 
       val resourceToAcls = requestKeysToAcls(key)
       resourceToAcls.get(topicResource).foreach { acls =>
         val describeAcls = topicDescribeAcl(topicResource)
-        val isAuthorized =  describeAcls == acls
+        val isAuthorized = describeAcls == acls
         addAndVerifyAcls(describeAcls, topicResource)
-        sendRequestAndVerifyResponseError(key, request, resources, isAuthorized = isAuthorized)
-        removeAllAcls()
+        sendRequestAndVerifyResponseError(request, resources, isAuthorized = isAuthorized)
+        removeAllClientAcls()
       }
 
       for ((resource, acls) <- resourceToAcls)
         addAndVerifyAcls(acls, resource)
-      sendRequestAndVerifyResponseError(key, request, resources, isAuthorized = true)
+      sendRequestAndVerifyResponseError(request, resources, isAuthorized = true)
     }
   }
 
@@ -445,12 +690,7 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
    * even if the topic doesn't exist, request APIs should not leak the topic name
    */
   @Test
-  def testAuthorizationWithTopicNotExisting() {
-    adminZkClient.deleteTopic(topic)
-    TestUtils.verifyTopicDeletion(zkUtils, topic, 1, servers)
-    adminZkClient.deleteTopic(deleteTopic)
-    TestUtils.verifyTopicDeletion(zkUtils, deleteTopic, 1, servers)
-
+  def testAuthorizationWithTopicNotExisting(): Unit = {
     val requestKeyToRequest = mutable.LinkedHashMap[ApiKeys, AbstractRequest](
       ApiKeys.METADATA -> createMetadataRequest(allowAutoTopicCreation = false),
       ApiKeys.PRODUCE -> createProduceRequest,
@@ -462,555 +702,872 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
       ApiKeys.DELETE_RECORDS -> deleteRecordsRequest,
       ApiKeys.ADD_PARTITIONS_TO_TXN -> addPartitionsToTxnRequest,
       ApiKeys.ADD_OFFSETS_TO_TXN -> addOffsetsToTxnRequest,
-      ApiKeys.CREATE_PARTITIONS -> createPartitionsRequest
+      ApiKeys.CREATE_PARTITIONS -> createPartitionsRequest,
+      ApiKeys.DELETE_GROUPS -> deleteGroupsRequest,
+      ApiKeys.OFFSET_FOR_LEADER_EPOCH -> offsetsForLeaderEpochRequest,
+      ApiKeys.ELECT_LEADERS -> electLeadersRequest
     )
 
     for ((key, request) <- requestKeyToRequest) {
-      removeAllAcls()
+      removeAllClientAcls()
       val resources = requestKeysToAcls(key).map(_._1.resourceType).toSet
-      sendRequestAndVerifyResponseError(key, request, resources, isAuthorized = false, topicExists = false)
+      sendRequestAndVerifyResponseError(request, resources, isAuthorized = false, topicExists = false)
 
       val resourceToAcls = requestKeysToAcls(key)
       resourceToAcls.get(topicResource).foreach { acls =>
         val describeAcls = topicDescribeAcl(topicResource)
         val isAuthorized = describeAcls == acls
         addAndVerifyAcls(describeAcls, topicResource)
-        sendRequestAndVerifyResponseError(key, request, resources, isAuthorized = isAuthorized, topicExists = false)
-        removeAllAcls()
+        sendRequestAndVerifyResponseError(request, resources, isAuthorized = isAuthorized, topicExists = false)
+        removeAllClientAcls()
       }
 
       for ((resource, acls) <- resourceToAcls)
         addAndVerifyAcls(acls, resource)
-      sendRequestAndVerifyResponseError(key, request, resources, isAuthorized = true, topicExists = false)
+      sendRequestAndVerifyResponseError(request, resources, isAuthorized = true, topicExists = false)
     }
   }
 
   @Test
-  def testProduceWithNoTopicAccess() {
-    try {
-      sendRecords(numRecords, tp)
-      fail("should have thrown exception")
-    } catch {
-      case _: TopicAuthorizationException => //expected
+  def testCreateTopicAuthorizationWithClusterCreate(): Unit = {
+    removeAllClientAcls()
+    val resources = Set[ResourceType](TOPIC)
+
+    sendRequestAndVerifyResponseError(createTopicsRequest, resources, isAuthorized = false)
+
+    for ((resource, acls) <- clusterCreateAcl)
+      addAndVerifyAcls(acls, resource)
+    sendRequestAndVerifyResponseError(createTopicsRequest, resources, isAuthorized = true)
+  }
+
+  @Test
+  def testFetchFollowerRequest(): Unit = {
+    createTopic(topic)
+
+    val request = createFetchFollowerRequest
+
+    removeAllClientAcls()
+    val resources = Set(topicResource.resourceType, clusterResource.resourceType)
+    sendRequestAndVerifyResponseError(request, resources, isAuthorized = false)
+
+    val readAcls = topicReadAcl(topicResource)
+    addAndVerifyAcls(readAcls, topicResource)
+    sendRequestAndVerifyResponseError(request, resources, isAuthorized = false)
+
+    val clusterAcls = clusterAcl(clusterResource)
+    addAndVerifyAcls(clusterAcls, clusterResource)
+    sendRequestAndVerifyResponseError(request, resources, isAuthorized = true)
+  }
+
+  @Test
+  def testIncrementalAlterConfigsRequestRequiresClusterPermissionForBrokerLogger(): Unit = {
+    createTopic(topic)
+
+    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 request = new IncrementalAlterConfigsRequest.Builder(data).build()
+
+    removeAllClientAcls()
+    val resources = Set(topicResource.resourceType, clusterResource.resourceType)
+    sendRequestAndVerifyResponseError(request, resources, isAuthorized = false)
+
+    val clusterAcls = clusterAlterConfigsAcl(clusterResource)
+    addAndVerifyAcls(clusterAcls, clusterResource)
+    sendRequestAndVerifyResponseError(request, resources, isAuthorized = true)
+  }
+
+  @Test
+  def testOffsetsForLeaderEpochClusterPermission(): Unit = {
+    createTopic(topic)
+
+    val request = offsetsForLeaderEpochRequest
+
+    removeAllClientAcls()
+
+    val resources = Set(topicResource.resourceType, clusterResource.resourceType)
+    sendRequestAndVerifyResponseError(request, resources, isAuthorized = false)
+
+    // Although the OffsetsForLeaderEpoch API now accepts topic describe, we should continue
+    // allowing cluster action for backwards compatibility
+    val clusterAcls = clusterAcl(clusterResource)
+    addAndVerifyAcls(clusterAcls, clusterResource)
+    sendRequestAndVerifyResponseError(request, resources, isAuthorized = true)
+  }
+
+  @Test
+  def testProduceWithNoTopicAccess(): Unit = {
+    createTopic(topic)
+    val producer = createProducer()
+    intercept[TopicAuthorizationException] {
+      sendRecords(producer, numRecords, tp)
     }
   }
 
   @Test
-  def testProduceWithTopicDescribe() {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Describe)), topicResource)
-    try {
-      sendRecords(numRecords, tp)
-      fail("should have thrown exception")
-    } catch {
-      case e: TopicAuthorizationException =>
-        assertEquals(Collections.singleton(topic), e.unauthorizedTopics())
+  def testProduceWithTopicDescribe(): Unit = {
+    createTopic(topic)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), topicResource)
+    val producer = createProducer()
+    intercept[TopicAuthorizationException] {
+      sendRecords(producer, numRecords, tp)
     }
   }
 
   @Test
-  def testProduceWithTopicRead() {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Read)), topicResource)
-    try {
-      sendRecords(numRecords, tp)
-      fail("should have thrown exception")
-    } catch {
-      case e: TopicAuthorizationException =>
-        assertEquals(Collections.singleton(topic), e.unauthorizedTopics())
+  def testProduceWithTopicRead(): Unit = {
+    createTopic(topic)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), topicResource)
+    val producer = createProducer()
+    intercept[TopicAuthorizationException] {
+      sendRecords(producer, numRecords, tp)
     }
   }
 
   @Test
-  def testProduceWithTopicWrite() {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)), topicResource)
-    sendRecords(numRecords, tp)
+  def testProduceWithTopicWrite(): Unit = {
+    createTopic(topic)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    val producer = createProducer()
+    sendRecords(producer, numRecords, tp)
   }
 
   @Test
-  def testCreatePermissionNeededForWritingToNonExistentTopic() {
-    val newTopic = "newTopic"
-    val topicPartition = new TopicPartition(newTopic, 0)
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)), new Resource(Topic, newTopic))
-    try {
-      sendRecords(numRecords, topicPartition)
-      Assert.fail("should have thrown exception")
-    } catch {
-      case e: TopicAuthorizationException => assertEquals(Collections.singleton(newTopic), e.unauthorizedTopics())
+  def testCreatePermissionOnTopicToWriteToNonExistentTopic(): Unit = {
+    testCreatePermissionNeededToWriteToNonExistentTopic(TOPIC)
+  }
+
+  @Test
+  def testCreatePermissionOnClusterToWriteToNonExistentTopic(): Unit = {
+    testCreatePermissionNeededToWriteToNonExistentTopic(CLUSTER)
+  }
+
+  private def testCreatePermissionNeededToWriteToNonExistentTopic(resType: ResourceType): Unit = {
+    val newTopicResource = new ResourcePattern(TOPIC, topic, LITERAL)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), newTopicResource)
+    val producer = createProducer()
+    val e = intercept[TopicAuthorizationException] {
+      sendRecords(producer, numRecords, tp)
     }
+    assertEquals(Collections.singleton(tp.topic), e.unauthorizedTopics())
+
+    val resource = if (resType == ResourceType.TOPIC) newTopicResource else clusterResource
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, CREATE, ALLOW)), resource)
 
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Create)), Resource.ClusterResource)
-    sendRecords(numRecords, topicPartition)
+    sendRecords(producer, numRecords, tp)
   }
 
-  @Test(expected = classOf[TopicAuthorizationException])
+  @Test
   def testConsumeUsingAssignWithNoAccess(): Unit = {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)), topicResource)
-    sendRecords(1, tp)
-    removeAllAcls()
-    this.consumers.head.assign(List(tp).asJava)
-    consumeRecords(this.consumers.head)
+    createTopic(topic)
+
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    val producer = createProducer()
+    sendRecords(producer, 1, tp)
+    removeAllClientAcls()
+
+    val consumer = createConsumer()
+    consumer.assign(List(tp).asJava)
+    intercept[TopicAuthorizationException] {
+      consumeRecords(consumer)
+    }
   }
 
   @Test
   def testSimpleConsumeWithOffsetLookupAndNoGroupAccess(): Unit = {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)), topicResource)
-    sendRecords(1, tp)
-    removeAllAcls()
+    createTopic(topic)
 
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Read)), 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)
-      this.consumers.head.assign(List(tp).asJava)
-      consumeRecords(this.consumers.head)
-      Assert.fail("should have thrown exception")
-    } catch {
-      case e: GroupAuthorizationException => assertEquals(group, e.groupId())
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    val producer = createProducer()
+    sendRecords(producer, 1, tp)
+    removeAllClientAcls()
+
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), topicResource)
+
+    // note this still depends on group access because we haven't set offsets explicitly, which means
+    // they will first be fetched from the consumer coordinator (which requires group access)
+    val consumer = createConsumer()
+    consumer.assign(List(tp).asJava)
+    val e = intercept[GroupAuthorizationException] {
+      consumeRecords(consumer)
     }
+    assertEquals(group, e.groupId())
   }
 
   @Test
   def testSimpleConsumeWithExplicitSeekAndNoGroupAccess(): Unit = {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)), topicResource)
-    sendRecords(1, tp)
-    removeAllAcls()
+    createTopic(topic)
+
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    val producer = createProducer()
+    sendRecords(producer, 1, tp)
+    removeAllClientAcls()
 
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Read)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), topicResource)
 
     // in this case, we do an explicit seek, so there should be no need to query the coordinator at all
-    this.consumers.head.assign(List(tp).asJava)
-    this.consumers.head.seekToBeginning(List(tp).asJava)
-    consumeRecords(this.consumers.head)
+    val consumer = createConsumer()
+    consumer.assign(List(tp).asJava)
+    consumer.seekToBeginning(List(tp).asJava)
+    consumeRecords(consumer)
   }
 
-  @Test(expected = classOf[KafkaException])
-  def testConsumeWithoutTopicDescribeAccess() {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)), topicResource)
-    sendRecords(1, tp)
-    removeAllAcls()
+  @Test
+  def testConsumeWithoutTopicDescribeAccess(): Unit = {
+    createTopic(topic)
+
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    val producer = createProducer()
+    sendRecords(producer, 1, tp)
+    removeAllClientAcls()
 
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Read)), groupResource)
-    this.consumers.head.assign(List(tp).asJava)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
+    val consumer = createConsumer()
+    consumer.assign(List(tp).asJava)
 
-    // the consumer should raise an exception if it receives UNKNOWN_TOPIC_OR_PARTITION
-    // from the ListOffsets response when looking up the initial position.
-    consumeRecords(this.consumers.head)
+    val e = intercept[TopicAuthorizationException] {
+      consumeRecords(consumer)
+    }
+    assertEquals(Collections.singleton(topic), e.unauthorizedTopics())
   }
 
   @Test
-  def testConsumeWithTopicDescribe() {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)), topicResource)
-    sendRecords(1, tp)
-    removeAllAcls()
+  def testConsumeWithTopicDescribe(): Unit = {
+    createTopic(topic)
 
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Describe)), topicResource)
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Read)), groupResource)
-    try {
-      this.consumers.head.assign(List(tp).asJava)
-      consumeRecords(this.consumers.head)
-      Assert.fail("should have thrown exception")
-    } catch {
-      case e: TopicAuthorizationException => assertEquals(Collections.singleton(topic), e.unauthorizedTopics())
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    val producer = createProducer()
+    sendRecords(producer, 1, tp)
+    removeAllClientAcls()
+
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
+
+    val consumer = createConsumer()
+    consumer.assign(List(tp).asJava)
+    val e = intercept[TopicAuthorizationException] {
+      consumeRecords(consumer)
     }
+    assertEquals(Collections.singleton(topic), e.unauthorizedTopics())
   }
 
   @Test
-  def testConsumeWithTopicWrite() {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)), topicResource)
-    sendRecords(1, tp)
-    removeAllAcls()
+  def testConsumeWithTopicWrite(): Unit = {
+    createTopic(topic)
 
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)), topicResource)
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Read)), groupResource)
-    try {
-      this.consumers.head.assign(List(tp).asJava)
-      consumeRecords(this.consumers.head)
-      Assert.fail("should have thrown exception")
-    } catch {
-      case e: TopicAuthorizationException =>
-        assertEquals(Collections.singleton(topic), e.unauthorizedTopics())
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    val producer = createProducer()
+    sendRecords(producer, 1, tp)
+    removeAllClientAcls()
+
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
+
+    val consumer = createConsumer()
+    consumer.assign(List(tp).asJava)
+    val e = intercept[TopicAuthorizationException] {
+      consumeRecords(consumer)
     }
+    assertEquals(Collections.singleton(topic), e.unauthorizedTopics())
   }
 
   @Test
-  def testConsumeWithTopicAndGroupRead() {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)), topicResource)
-    sendRecords(1, tp)
-    removeAllAcls()
+  def testConsumeWithTopicAndGroupRead(): Unit = {
+    createTopic(topic)
+
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    val producer = createProducer()
+    sendRecords(producer, 1, tp)
+    removeAllClientAcls()
+
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
 
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Read)), topicResource)
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Read)), groupResource)
-    this.consumers.head.assign(List(tp).asJava)
-    consumeRecords(this.consumers.head)
+    val consumer = createConsumer()
+    consumer.assign(List(tp).asJava)
+    consumeRecords(consumer)
   }
 
+  @nowarn("cat=deprecation")
   @Test
-  def testPatternSubscriptionWithNoTopicAccess() {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)), topicResource)
-    sendRecords(1, tp)
-    removeAllAcls()
+  def testPatternSubscriptionWithNoTopicAccess(): Unit = {
+    createTopic(topic)
+
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    val producer = createProducer()
+    sendRecords(producer, 1, tp)
+    removeAllClientAcls()
 
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Read)), groupResource)
-    this.consumers.head.subscribe(Pattern.compile(topicPattern), new NoOpConsumerRebalanceListener)
-    this.consumers.head.poll(50)
-    assertTrue(this.consumers.head.subscription.isEmpty)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
+
+    val consumer = createConsumer()
+    consumer.subscribe(Pattern.compile(topicPattern), new NoOpConsumerRebalanceListener)
+    consumer.poll(0)
+    assertTrue(consumer.subscription.isEmpty)
   }
 
   @Test
-  def testPatternSubscriptionWithTopicDescribeOnlyAndGroupRead() {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)), topicResource)
-    sendRecords(1, tp)
-    removeAllAcls()
+  def testPatternSubscriptionWithTopicDescribeOnlyAndGroupRead(): Unit = {
+    createTopic(topic)
 
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Describe)), topicResource)
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Read)), groupResource)
-    val consumer = consumers.head
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    val producer = createProducer()
+    sendRecords(producer, 1, tp)
+    removeAllClientAcls()
+
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
+    val consumer = createConsumer()
     consumer.subscribe(Pattern.compile(topicPattern))
-    try {
+    val e = intercept[TopicAuthorizationException] {
       consumeRecords(consumer)
-      Assert.fail("Expected TopicAuthorizationException")
-    } catch {
-      case _: TopicAuthorizationException => //expected
     }
+    assertEquals(Collections.singleton(topic), e.unauthorizedTopics())
   }
 
+  @nowarn("cat=deprecation")
   @Test
-  def testPatternSubscriptionWithTopicAndGroupRead() {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)), topicResource)
-    sendRecords(1, tp)
+  def testPatternSubscriptionWithTopicAndGroupRead(): Unit = {
+    createTopic(topic)
+
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    val producer = createProducer()
+    sendRecords(producer, 1, tp)
 
     // create an unmatched topic
     val unmatchedTopic = "unmatched"
-    TestUtils.createTopic(zkUtils, unmatchedTopic, 1, 1, this.servers)
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)),  new Resource(Topic, unmatchedTopic))
-    sendRecords(1, new TopicPartition(unmatchedTopic, part))
-    removeAllAcls()
-
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Read)), topicResource)
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Read)), groupResource)
-    val consumer = consumers.head
+    createTopic(unmatchedTopic)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)),  new ResourcePattern(TOPIC, unmatchedTopic, LITERAL))
+    sendRecords(producer, 1, new TopicPartition(unmatchedTopic, part))
+    removeAllClientAcls()
+
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
+    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(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Read)),  new Resource(Topic,
-      GROUP_METADATA_TOPIC_NAME))
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)),  new ResourcePattern(TOPIC,
+      GROUP_METADATA_TOPIC_NAME, LITERAL))
     consumer.subscribe(Pattern.compile(GROUP_METADATA_TOPIC_NAME))
     consumer.poll(0)
     assertTrue(consumer.subscription().isEmpty)
     assertTrue(consumer.assignment().isEmpty)
   }
 
+  @nowarn("cat=deprecation")
   @Test
-  def testPatternSubscriptionMatchingInternalTopic() {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)), topicResource)
-    sendRecords(1, tp)
-    removeAllAcls()
+  def testPatternSubscriptionMatchingInternalTopic(): Unit = {
+    createTopic(topic)
 
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Read)), topicResource)
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Read)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    val producer = createProducer()
+    sendRecords(producer, 1, tp)
+    removeAllClientAcls()
+
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
 
-    val consumerConfig = new Properties
     consumerConfig.put(ConsumerConfig.EXCLUDE_INTERNAL_TOPICS_CONFIG, "false")
-    val consumer = TestUtils.createNewConsumer(TestUtils.getBrokerListStrFromServers(servers), groupId = group,
-      securityProtocol = SecurityProtocol.PLAINTEXT, props = Some(consumerConfig))
-    try {
-      // ensure that internal topics are not included if no permission
-      consumer.subscribe(Pattern.compile(".*"))
-      consumeRecords(consumer)
-      assertEquals(Set(topic).asJava, consumer.subscription)
+    val consumer = createConsumer()
+    // ensure that internal topics are not included if no permission
+    consumer.subscribe(Pattern.compile(".*"))
+    consumeRecords(consumer)
+    assertEquals(Set(topic).asJava, consumer.subscription)
 
-      // now authorize the user for the internal topic and verify that we can subscribe
-      addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Read)), Resource(Topic,
-        GROUP_METADATA_TOPIC_NAME))
-      consumer.subscribe(Pattern.compile(GROUP_METADATA_TOPIC_NAME))
-      consumer.poll(0)
-      assertEquals(Set(GROUP_METADATA_TOPIC_NAME), consumer.subscription.asScala)
-    } finally consumer.close()
+    // now authorize the user for the internal topic and verify that we can subscribe
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), new ResourcePattern(TOPIC,
+      GROUP_METADATA_TOPIC_NAME, LITERAL))
+    consumer.subscribe(Pattern.compile(GROUP_METADATA_TOPIC_NAME))
+    consumer.poll(0)
+    assertEquals(Set(GROUP_METADATA_TOPIC_NAME), consumer.subscription.asScala)
   }
 
   @Test
-  def testPatternSubscriptionMatchingInternalTopicWithDescribeOnlyPermission() {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)), topicResource)
-    sendRecords(1, tp)
-    removeAllAcls()
+  def testPatternSubscriptionMatchingInternalTopicWithDescribeOnlyPermission(): Unit = {
+    createTopic(topic)
 
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Read)), topicResource)
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Read)), groupResource)
-    val internalTopicResource = new Resource(Topic, GROUP_METADATA_TOPIC_NAME)
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Describe)), internalTopicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    val producer = createProducer()
+    sendRecords(producer, 1, tp)
+    removeAllClientAcls()
+
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
+    val internalTopicResource = new ResourcePattern(TOPIC, GROUP_METADATA_TOPIC_NAME, LITERAL)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), internalTopicResource)
 
-    val consumerConfig = new Properties
     consumerConfig.put(ConsumerConfig.EXCLUDE_INTERNAL_TOPICS_CONFIG, "false")
-    val consumer = TestUtils.createNewConsumer(TestUtils.getBrokerListStrFromServers(servers), groupId = group,
-      securityProtocol = SecurityProtocol.PLAINTEXT, props = Some(consumerConfig))
-    try {
-      consumer.subscribe(Pattern.compile(".*"))
+    val consumer = createConsumer()
+    consumer.subscribe(Pattern.compile(".*"))
+    val e = intercept[TopicAuthorizationException] {
       // It is possible that the first call returns records of "topic" and the second call throws TopicAuthorizationException
       consumeRecords(consumer)
       consumeRecords(consumer)
-      Assert.fail("Expected TopicAuthorizationException")
-    } catch {
-      case _: TopicAuthorizationException => //expected
-    } finally consumer.close()
+    }
+    assertEquals(Collections.singleton(GROUP_METADATA_TOPIC_NAME), e.unauthorizedTopics())
   }
 
   @Test
-  def testPatternSubscriptionNotMatchingInternalTopic() {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)), topicResource)
-    sendRecords(1, tp)
-    removeAllAcls()
+  def testPatternSubscriptionNotMatchingInternalTopic(): Unit = {
+    createTopic(topic)
+
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    val producer = createProducer()
+    sendRecords(producer, 1, tp)
+    removeAllClientAcls()
 
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Read)), topicResource)
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Read)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
 
-    val consumerConfig = new Properties
     consumerConfig.put(ConsumerConfig.EXCLUDE_INTERNAL_TOPICS_CONFIG, "false")
-    val consumer = TestUtils.createNewConsumer(TestUtils.getBrokerListStrFromServers(servers), groupId = group,
-      securityProtocol = SecurityProtocol.PLAINTEXT, props = Some(consumerConfig))
-    try {
-      consumer.subscribe(Pattern.compile(topicPattern))
-      consumeRecords(consumer)
-    } finally consumer.close()
-}
+    val consumer = createConsumer()
+    consumer.subscribe(Pattern.compile(topicPattern))
+    consumeRecords(consumer)
+  }
+
+  @Test
+  def testCreatePermissionOnTopicToReadFromNonExistentTopic(): Unit = {
+    testCreatePermissionNeededToReadFromNonExistentTopic("newTopic",
+      Set(new AccessControlEntry(clientPrincipalString, WildcardHost, CREATE, ALLOW)),
+      TOPIC)
+  }
 
   @Test
-  def testCreatePermissionNeededToReadFromNonExistentTopic() {
-    val newTopic = "newTopic"
+  def testCreatePermissionOnClusterToReadFromNonExistentTopic(): Unit = {
+    testCreatePermissionNeededToReadFromNonExistentTopic("newTopic",
+      Set(new AccessControlEntry(clientPrincipalString, WildcardHost, CREATE, ALLOW)),
+      CLUSTER)
+  }
+
+  private def testCreatePermissionNeededToReadFromNonExistentTopic(newTopic: String, acls: Set[AccessControlEntry], resType: ResourceType): Unit = {
     val topicPartition = new TopicPartition(newTopic, 0)
-    val newTopicResource = new Resource(Topic, newTopic)
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Read)), newTopicResource)
+    val newTopicResource = new ResourcePattern(TOPIC, newTopic, LITERAL)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), newTopicResource)
     addAndVerifyAcls(groupReadAcl(groupResource), groupResource)
-    addAndVerifyAcls(clusterAcl(Resource.ClusterResource), Resource.ClusterResource)
-    try {
-      this.consumers.head.assign(List(topicPartition).asJava)
-      consumeRecords(this.consumers.head)
-      Assert.fail("should have thrown exception")
-    } catch {
-      case e: TopicAuthorizationException =>
-        assertEquals(Collections.singleton(newTopic), e.unauthorizedTopics())
-    }
+    val consumer = createConsumer()
+    consumer.assign(List(topicPartition).asJava)
+    val unauthorizedTopics = intercept[TopicAuthorizationException] {
+      (0 until 10).foreach(_ => consumer.poll(Duration.ofMillis(50L)))
+    }.unauthorizedTopics
+    assertEquals(Collections.singleton(newTopic), unauthorizedTopics)
 
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)), newTopicResource)
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Create)), Resource.ClusterResource)
+    val resource = if (resType == TOPIC) newTopicResource else clusterResource
+    addAndVerifyAcls(acls, resource)
 
-    sendRecords(numRecords, topicPartition)
-    consumeRecords(this.consumers.head, topic = newTopic, part = 0)
+    TestUtils.waitUntilTrue(() => {
+      consumer.poll(Duration.ofMillis(50L))
+      this.zkClient.topicExists(newTopic)
+    }, "Expected topic was not created")
   }
 
-  @Test(expected = classOf[AuthorizationException])
-  def testCommitWithNoAccess() {
-    this.consumers.head.commitSync(Map(tp -> new OffsetAndMetadata(5)).asJava)
+  @Test
+  def testCreatePermissionMetadataRequestAutoCreate(): Unit = {
+    val readAcls = topicReadAcl(topicResource)
+    addAndVerifyAcls(readAcls, topicResource)
+    assertFalse(zkClient.topicExists(topic))
+
+    val metadataRequest = new MetadataRequest.Builder(List(topic).asJava, true).build()
+    val metadataResponse = connectAndReceive[MetadataResponse](metadataRequest)
+
+    assertEquals(Set().asJava, metadataResponse.topicsByError(Errors.NONE))
+
+    val createAcls = topicCreateAcl(topicResource)
+    addAndVerifyAcls(createAcls, topicResource)
+
+    // retry as topic being created can have MetadataResponse with Errors.LEADER_NOT_AVAILABLE
+    TestUtils.retry(JTestUtils.DEFAULT_MAX_WAIT_MS) {
+      val metadataResponse = connectAndReceive[MetadataResponse](metadataRequest)
+      assertEquals(Set(topic).asJava, metadataResponse.topicsByError(Errors.NONE))
+    }
   }
 
-  @Test(expected = classOf[KafkaException])
-  def testCommitWithNoTopicAccess() {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Read)), groupResource)
-    this.consumers.head.commitSync(Map(tp -> new OffsetAndMetadata(5)).asJava)
+  @Test(expected = classOf[GroupAuthorizationException])
+  def testCommitWithNoAccess(): Unit = {
+    val consumer = createConsumer()
+    consumer.commitSync(Map(tp -> new OffsetAndMetadata(5)).asJava)
+  }
+
+  @Test(expected = classOf[TopicAuthorizationException])
+  def testCommitWithNoTopicAccess(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, 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(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Read)), groupResource)
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)), topicResource)
-    this.consumers.head.commitSync(Map(tp -> new OffsetAndMetadata(5)).asJava)
+  def testCommitWithTopicWrite(): Unit = {
+    createTopic(topic)
+
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, 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(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Read)), groupResource)
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Describe)), topicResource)
-    this.consumers.head.commitSync(Map(tp -> new OffsetAndMetadata(5)).asJava)
+  def testCommitWithTopicDescribe(): Unit = {
+    createTopic(topic)
+
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, 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(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Read)), topicResource)
-    this.consumers.head.commitSync(Map(tp -> new OffsetAndMetadata(5)).asJava)
+  def testCommitWithNoGroupAccess(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), topicResource)
+    val consumer = createConsumer()
+    consumer.commitSync(Map(tp -> new OffsetAndMetadata(5)).asJava)
   }
 
   @Test
-  def testCommitWithTopicAndGroupRead() {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Read)), groupResource)
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Read)), topicResource)
-    this.consumers.head.commitSync(Map(tp -> new OffsetAndMetadata(5)).asJava)
+  def testCommitWithTopicAndGroupRead(): Unit = {
+    createTopic(topic)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), topicResource)
+    val consumer = createConsumer()
+    consumer.commitSync(Map(tp -> new OffsetAndMetadata(5)).asJava)
   }
 
-  @Test(expected = classOf[AuthorizationException])
-  def testOffsetFetchWithNoAccess() {
-    this.consumers.head.assign(List(tp).asJava)
-    this.consumers.head.position(tp)
+  @Test(expected = classOf[TopicAuthorizationException])
+  def testOffsetFetchWithNoAccess(): Unit = {
+    val consumer = createConsumer()
+    consumer.assign(List(tp).asJava)
+    consumer.position(tp)
   }
 
   @Test(expected = classOf[GroupAuthorizationException])
-  def testOffsetFetchWithNoGroupAccess() {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Read)), topicResource)
-    this.consumers.head.assign(List(tp).asJava)
-    this.consumers.head.position(tp)
+  def testOffsetFetchWithNoGroupAccess(): Unit = {
+    createTopic(topic)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, 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(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Read)), groupResource)
-    this.consumers.head.assign(List(tp).asJava)
-    this.consumers.head.position(tp)
+  @Test(expected = classOf[TopicAuthorizationException])
+  def testOffsetFetchWithNoTopicAccess(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
+    val consumer = createConsumer()
+    consumer.assign(List(tp).asJava)
+    consumer.position(tp)
   }
 
   @Test
-  def testFetchAllOffsetsTopicAuthorization() {
+  def testFetchAllOffsetsTopicAuthorization(): Unit = {
+    createTopic(topic)
+
     val offset = 15L
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Read)), groupResource)
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Read)), topicResource)
-    this.consumers.head.assign(List(tp).asJava)
-    this.consumers.head.commitSync(Map(tp -> new OffsetAndMetadata(offset)).asJava)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, 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(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Read)), groupResource)
+    removeAllClientAcls()
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, 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)
-    var offsetFetchResponse = sendOffsetFetchRequest(offsetFetchRequest, anySocketServer)
+    val offsetFetchRequest = new requests.OffsetFetchRequest.Builder(group, false, null, false).build()
+    var offsetFetchResponse = connectAndReceive[OffsetFetchResponse](offsetFetchRequest)
     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(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Describe)), topicResource)
-    offsetFetchResponse = sendOffsetFetchRequest(offsetFetchRequest, anySocketServer)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), topicResource)
+    offsetFetchResponse = connectAndReceive[OffsetFetchResponse](offsetFetchRequest)
     assertEquals(Errors.NONE, offsetFetchResponse.error)
     assertTrue(offsetFetchResponse.responseData.containsKey(tp))
     assertEquals(offset, offsetFetchResponse.responseData.get(tp).offset)
   }
 
   @Test
-  def testOffsetFetchTopicDescribe() {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Describe)), groupResource)
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Describe)), topicResource)
-    this.consumers.head.assign(List(tp).asJava)
-    this.consumers.head.position(tp)
+  def testOffsetFetchTopicDescribe(): Unit = {
+    createTopic(topic)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), topicResource)
+    val consumer = createConsumer()
+    consumer.assign(List(tp).asJava)
+    consumer.position(tp)
   }
 
   @Test
-  def testOffsetFetchWithTopicAndGroupRead() {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Read)), groupResource)
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Read)), topicResource)
-    this.consumers.head.assign(List(tp).asJava)
-    this.consumers.head.position(tp)
+  def testOffsetFetchWithTopicAndGroupRead(): Unit = {
+    createTopic(topic)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), topicResource)
+    val consumer = createConsumer()
+    consumer.assign(List(tp).asJava)
+    consumer.position(tp)
   }
 
   @Test(expected = classOf[TopicAuthorizationException])
-  def testListOffsetsWithNoTopicAccess() {
-    this.consumers.head.partitionsFor(topic)
+  def testMetadataWithNoTopicAccess(): Unit = {
+    val consumer = createConsumer()
+    consumer.partitionsFor(topic)
   }
 
   @Test
-  def testListOffsetsWithTopicDescribe() {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Describe)), topicResource)
-    this.consumers.head.partitionsFor(topic)
+  def testMetadataWithTopicDescribe(): Unit = {
+    createTopic(topic)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), topicResource)
+    val consumer = createConsumer()
+    consumer.partitionsFor(topic)
   }
 
-  @Test(expected = classOf[GroupAuthorizationException])
-  def testDescribeGroupApiWithNoGroupAcl() {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Describe)), topicResource)
-    AdminClient.createSimplePlaintext(brokerList).describeConsumerGroup(group)
+  @Test(expected = classOf[TopicAuthorizationException])
+  def testListOffsetsWithNoTopicAccess(): Unit = {
+    val consumer = createConsumer()
+    consumer.endOffsets(Set(tp).asJava)
+  }
+
+  @Test
+  def testListOffsetsWithTopicDescribe(): Unit = {
+    createTopic(topic)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), topicResource)
+    val consumer = createConsumer()
+    consumer.endOffsets(Set(tp).asJava)
+  }
+
+  @Test
+  def testDescribeGroupApiWithNoGroupAcl(): Unit = {
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, 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(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Describe)), groupResource)
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Describe)), topicResource)
-    AdminClient.createSimplePlaintext(brokerList).describeConsumerGroup(group)
+  def testDescribeGroupApiWithGroupDescribe(): Unit = {
+    createTopic(topic)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), topicResource)
+    createAdminClient().describeConsumerGroups(Seq(group).asJava).describedGroups().get(group).get()
   }
 
   @Test
-  def testDescribeGroupCliWithGroupDescribe() {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Describe)), groupResource)
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Describe)), topicResource)
+  def testDescribeGroupCliWithGroupDescribe(): Unit = {
+    createTopic(topic)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), topicResource)
 
     val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--group", group)
     val opts = new ConsumerGroupCommandOptions(cgcArgs)
-    val consumerGroupService = new KafkaConsumerGroupService(opts)
-    consumerGroupService.describeGroup()
+    val consumerGroupService = new ConsumerGroupService(opts)
+    consumerGroupService.describeGroups()
     consumerGroupService.close()
   }
 
   @Test
-  def testUnauthorizedDeleteTopicsWithoutDescribe() {
-    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)
+  def testListGroupApiWithAndWithoutListGroupAcls(): Unit = {
+    createTopic(topic)
+
+    // write some record to the topic
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, 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 AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), new ResourcePattern(GROUP, group2, LITERAL))
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), topicResource)
+    val consumer = createConsumer()
+    consumer.subscribe(Collections.singleton(topic))
+    consumeRecords(consumer)
+
+    val otherConsumerProps = new Properties
+    otherConsumerProps.put(ConsumerConfig.GROUP_ID_CONFIG, group2)
+    val otherConsumer = createConsumer(configOverrides = otherConsumerProps)
+    otherConsumer.subscribe(Collections.singleton(topic))
+    consumeRecords(otherConsumer)
+
+    val adminClient = createAdminClient()
+
+    // first use cluster describe permission
+    removeAllClientAcls()
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, 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
+    removeAllClientAcls()
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
+    // it should list only one group now
+    val groupList = adminClient.listConsumerGroups().all().get().asScala.toList
+    assertEquals(1, groupList.length)
+    assertEquals(group, groupList.head.groupId)
+
+    // now remove all acls and verify describe group access is required to list any group
+    removeAllClientAcls()
+    val listGroupResult = adminClient.listConsumerGroups()
+    assertEquals(List(), listGroupResult.errors().get().asScala.toList)
+    assertEquals(List(), listGroupResult.all().get().asScala.toList)
+    otherConsumer.close()
   }
 
   @Test
-  def testUnauthorizedDeleteTopicsWithDescribe() {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Describe)), deleteTopicResource)
-    val response = connectAndSend(deleteTopicsRequest, ApiKeys.DELETE_TOPICS)
-    val version = ApiKeys.DELETE_TOPICS.latestVersion
-    val deleteResponse = DeleteTopicsResponse.parse(response, version)
+  def testDeleteGroupApiWithDeleteGroupAcl(): Unit = {
+    createTopic(topic)
 
-    assertEquals(Errors.TOPIC_AUTHORIZATION_FAILED, deleteResponse.errors.asScala.head._2)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DELETE, ALLOW)), groupResource)
+    val consumer = createConsumer()
+    consumer.assign(List(tp).asJava)
+    consumer.commitSync(Map(tp -> new OffsetAndMetadata(5, "")).asJava)
+    createAdminClient().deleteConsumerGroups(Seq(group).asJava).deletedGroups().get(group).get()
   }
 
   @Test
-  def testDeleteTopicsWithWildCardAuth() {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Delete)), new Resource(Topic, "*"))
-    val response = connectAndSend(deleteTopicsRequest, ApiKeys.DELETE_TOPICS)
-    val version = ApiKeys.DELETE_TOPICS.latestVersion
-    val deleteResponse = DeleteTopicsResponse.parse(response, version)
+  def testDeleteGroupApiWithNoDeleteGroupAcl(): Unit = {
+    createTopic(topic)
 
-    assertEquals(Errors.NONE, deleteResponse.errors.asScala.head._2)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), topicResource)
+    val consumer = createConsumer()
+    consumer.assign(List(tp).asJava)
+    consumer.commitSync(Map(tp -> new OffsetAndMetadata(5, "")).asJava)
+    val result = createAdminClient().deleteConsumerGroups(Seq(group).asJava)
+    TestUtils.assertFutureExceptionTypeEquals(result.deletedGroups().get(group), classOf[GroupAuthorizationException])
   }
 
   @Test
-  def testUnauthorizedDeleteRecordsWithoutDescribe() {
-    val response = connectAndSend(deleteRecordsRequest, ApiKeys.DELETE_RECORDS)
-    val version = ApiKeys.DELETE_RECORDS.latestVersion
-    val deleteRecordsResponse = DeleteRecordsResponse.parse(response, version)
-    assertEquals(Errors.TOPIC_AUTHORIZATION_FAILED, deleteRecordsResponse.responses.asScala.head._2.error)
+  def testDeleteGroupApiWithNoDeleteGroupAcl2(): Unit = {
+    val result = createAdminClient().deleteConsumerGroups(Seq(group).asJava)
+    TestUtils.assertFutureExceptionTypeEquals(result.deletedGroups().get(group), classOf[GroupAuthorizationException])
   }
 
   @Test
-  def testUnauthorizedDeleteRecordsWithDescribe() {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Describe)), deleteTopicResource)
-    val response = connectAndSend(deleteRecordsRequest, ApiKeys.DELETE_RECORDS)
-    val version = ApiKeys.DELETE_RECORDS.latestVersion
-    val deleteRecordsResponse = DeleteRecordsResponse.parse(response, version)
-    assertEquals(Errors.TOPIC_AUTHORIZATION_FAILED, deleteRecordsResponse.responses.asScala.head._2.error)
+  def testDeleteGroupOffsetsWithAcl(): Unit = {
+    createTopic(topic)
+
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DELETE, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), topicResource)
+    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 testDeleteRecordsWithWildCardAuth() {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Delete)), new Resource(Topic, "*"))
-    val response = connectAndSend(deleteRecordsRequest, ApiKeys.DELETE_RECORDS)
-    val version = ApiKeys.DELETE_RECORDS.latestVersion
-    val deleteRecordsResponse = DeleteRecordsResponse.parse(response, version)
+  def testDeleteGroupOffsetsWithoutDeleteAcl(): Unit = {
+    createTopic(topic)
 
-    assertEquals(Errors.NONE, deleteRecordsResponse.responses.asScala.head._2.error)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, 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 testUnauthorizedCreatePartitions() {
-    val response = connectAndSend(createPartitionsRequest, ApiKeys.CREATE_PARTITIONS)
-    val version = ApiKeys.CREATE_PARTITIONS.latestVersion
-    val createPartitionsResponse = CreatePartitionsResponse.parse(response, version)
-    assertEquals(Errors.TOPIC_AUTHORIZATION_FAILED, createPartitionsResponse.errors.asScala.head._2.error)
+  def testDeleteGroupOffsetsWithDeleteAclWithoutTopicAcl(): Unit = {
+    createTopic(topic)
+    // Create the consumer group
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, READ, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, 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
+    removeAllClientAcls()
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DELETE, ALLOW)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, 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 testCreatePartitionsWithWildCardAuth() {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Alter)), new Resource(Topic, "*"))
-    val response = connectAndSend(createPartitionsRequest, ApiKeys.CREATE_PARTITIONS)
-    val version = ApiKeys.CREATE_PARTITIONS.latestVersion
-    val createPartitionsResponse = CreatePartitionsResponse.parse(response, version)
-    assertEquals(Errors.NONE, createPartitionsResponse.errors.asScala.head._2.error)
+  def testDeleteGroupOffsetsWithNoAcl(): Unit = {
+    val result = createAdminClient().deleteConsumerGroupOffsets(group, Set(tp).asJava)
+    TestUtils.assertFutureExceptionTypeEquals(result.all(), classOf[GroupAuthorizationException])
+  }
+
+  @Test
+  def testUnauthorizedDeleteTopicsWithoutDescribe(): Unit = {
+    val deleteResponse = connectAndReceive[DeleteTopicsResponse](deleteTopicsRequest)
+    assertEquals(Errors.TOPIC_AUTHORIZATION_FAILED.code, deleteResponse.data.responses.find(topic).errorCode)
+  }
+
+  @Test
+  def testUnauthorizedDeleteTopicsWithDescribe(): Unit = {
+    createTopic(topic)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), topicResource)
+    val deleteResponse = connectAndReceive[DeleteTopicsResponse](deleteTopicsRequest)
+    assertEquals(Errors.TOPIC_AUTHORIZATION_FAILED.code, deleteResponse.data.responses.find(topic).errorCode)
+  }
+
+  @Test
+  def testDeleteTopicsWithWildCardAuth(): Unit = {
+    createTopic(topic)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DELETE, ALLOW)), new ResourcePattern(TOPIC, "*", LITERAL))
+    val deleteResponse = connectAndReceive[DeleteTopicsResponse](deleteTopicsRequest)
+    assertEquals(Errors.NONE.code, deleteResponse.data.responses.find(topic).errorCode)
+  }
+
+  @Test
+  def testUnauthorizedDeleteRecordsWithoutDescribe(): Unit = {
+    val deleteRecordsResponse = connectAndReceive[DeleteRecordsResponse](deleteRecordsRequest)
+    assertEquals(Errors.TOPIC_AUTHORIZATION_FAILED.code, deleteRecordsResponse.data.topics.asScala.head.
+      partitions.asScala.head.errorCode)
+  }
+
+  @Test
+  def testUnauthorizedDeleteRecordsWithDescribe(): Unit = {
+    createTopic(topic)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), topicResource)
+    val deleteRecordsResponse = connectAndReceive[DeleteRecordsResponse](deleteRecordsRequest)
+    assertEquals(Errors.TOPIC_AUTHORIZATION_FAILED.code, deleteRecordsResponse.data.topics.asScala.head.
+      partitions.asScala.head.errorCode)
+  }
+
+  @Test
+  def testDeleteRecordsWithWildCardAuth(): Unit = {
+    createTopic(topic)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DELETE, ALLOW)), new ResourcePattern(TOPIC, "*", LITERAL))
+    val deleteRecordsResponse = connectAndReceive[DeleteRecordsResponse](deleteRecordsRequest)
+    assertEquals(Errors.NONE.code, deleteRecordsResponse.data.topics.asScala.head.
+      partitions.asScala.head.errorCode)
+  }
+
+  @Test
+  def testUnauthorizedCreatePartitions(): Unit = {
+    val createPartitionsResponse = connectAndReceive[CreatePartitionsResponse](createPartitionsRequest)
+    assertEquals(Errors.TOPIC_AUTHORIZATION_FAILED.code(), createPartitionsResponse.data.results.asScala.head.errorCode())
+  }
+
+  @Test
+  def testCreatePartitionsWithWildCardAuth(): Unit = {
+    createTopic(topic)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, ALTER, ALLOW)), new ResourcePattern(TOPIC, "*", LITERAL))
+    val createPartitionsResponse = connectAndReceive[CreatePartitionsResponse](createPartitionsRequest)
+    assertEquals(Errors.NONE.code(), createPartitionsResponse.data.results.asScala.head.errorCode())
   }
 
   @Test(expected = classOf[TransactionalIdAuthorizationException])
   def testTransactionalProducerInitTransactionsNoWriteTransactionalIdAcl(): Unit = {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Describe)), transactionalIdResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), transactionalIdResource)
     val producer = buildTransactionalProducer()
     producer.initTransactions()
   }
@@ -1023,38 +1580,40 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
 
   @Test
   def testSendOffsetsWithNoConsumerGroupDescribeAccess(): Unit = {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, ClusterAction)), Resource.ClusterResource)
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)), topicResource)
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)), transactionalIdResource)
+    createTopic(topic)
+
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, CLUSTER_ACTION, ALLOW)), clusterResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), transactionalIdResource)
     val producer = buildTransactionalProducer()
     producer.initTransactions()
     producer.beginTransaction()
-    try {
-      producer.sendOffsetsToTransaction(Map(new TopicPartition(topic, 0) -> new OffsetAndMetadata(0L)).asJava, group)
-      fail("Should have raised GroupAuthorizationException")
-    } catch {
-      case e: GroupAuthorizationException =>
+
+    intercept[GroupAuthorizationException] {
+      producer.sendOffsetsToTransaction(Map(tp -> new OffsetAndMetadata(0L)).asJava, group)
     }
   }
 
   @Test
   def testSendOffsetsWithNoConsumerGroupWriteAccess(): Unit = {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)), transactionalIdResource)
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Describe)), groupResource)
+    createTopic(topic)
+
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), transactionalIdResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), groupResource)
     val producer = buildTransactionalProducer()
     producer.initTransactions()
     producer.beginTransaction()
-    try {
-      producer.sendOffsetsToTransaction(Map(new TopicPartition(topic, 0) -> new OffsetAndMetadata(0L)).asJava, group)
-      fail("Should have raised GroupAuthorizationException")
-    } catch {
-      case e: GroupAuthorizationException =>
+
+    intercept[GroupAuthorizationException] {
+      producer.sendOffsetsToTransaction(Map(tp -> new OffsetAndMetadata(0L)).asJava, group)
     }
   }
 
   @Test
   def testIdempotentProducerNoIdempotentWriteAclInInitProducerId(): Unit = {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)), topicResource)
+    createTopic(topic)
+
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
     val producer = buildIdempotentProducer()
     try {
       // the InitProducerId is sent asynchronously, so we expect the error either in the callback
@@ -1078,8 +1637,10 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
 
   @Test
   def testIdempotentProducerNoIdempotentWriteAclInProduce(): Unit = {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)), topicResource)
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, IdempotentWrite)), Resource.ClusterResource)
+    createTopic(topic)
+
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, IDEMPOTENT_WRITE, ALLOW)), clusterResource)
 
     val producer = buildIdempotentProducer()
 
@@ -1087,8 +1648,8 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
     producer.send(new ProducerRecord[Array[Byte], Array[Byte]](topic, "hi".getBytes)).get()
 
     // revoke the IdempotentWrite permission
-    removeAllAcls()
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)), topicResource)
+    removeAllClientAcls()
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
 
     try {
       // the send should now fail with a cluster auth error
@@ -1111,155 +1672,154 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
 
   @Test
   def shouldInitTransactionsWhenAclSet(): Unit = {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)), transactionalIdResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), transactionalIdResource)
     val producer = buildTransactionalProducer()
     producer.initTransactions()
   }
 
   @Test
   def testTransactionalProducerTopicAuthorizationExceptionInSendCallback(): Unit = {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)), transactionalIdResource)
+    createTopic(topic)
+
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), transactionalIdResource)
     // add describe access so that we can fetch metadata
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Describe)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), topicResource)
     val producer = buildTransactionalProducer()
     producer.initTransactions()
     producer.beginTransaction()
-    try {
-      producer.send(new ProducerRecord(tp.topic, tp.partition, "1".getBytes, "1".getBytes)).get
-      Assert.fail("expected TopicAuthorizationException")
-    } catch {
-      case e: ExecutionException =>
-        e.getCause match {
-          case cause: TopicAuthorizationException =>
-            assertEquals(Set(topic), cause.unauthorizedTopics().asScala)
-          case other =>
-            fail("Unexpected failure cause in send callback")
-        }
-    }
+
+    val future = producer.send(new ProducerRecord(tp.topic, tp.partition, "1".getBytes, "1".getBytes))
+    val e = JTestUtils.assertFutureThrows(future, classOf[TopicAuthorizationException])
+    assertEquals(Set(topic), e.unauthorizedTopics.asScala)
   }
 
   @Test
   def testTransactionalProducerTopicAuthorizationExceptionInCommit(): Unit = {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)), transactionalIdResource)
+    createTopic(topic)
+
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), transactionalIdResource)
     // add describe access so that we can fetch metadata
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Describe)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)), topicResource)
     val producer = buildTransactionalProducer()
     producer.initTransactions()
     producer.beginTransaction()
-    try {
+
+    intercept[TopicAuthorizationException] {
       producer.send(new ProducerRecord(tp.topic, tp.partition, "1".getBytes, "1".getBytes))
       producer.commitTransaction()
-      Assert.fail("expected TopicAuthorizationException")
-    } catch {
-      case e: TopicAuthorizationException =>
-        assertEquals(Set(topic), e.unauthorizedTopics().asScala)
     }
   }
 
   @Test
   def shouldThrowTransactionalIdAuthorizationExceptionWhenNoTransactionAccessDuringSend(): Unit = {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)), transactionalIdResource)
+    createTopic(topic)
+
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), transactionalIdResource)
     val producer = buildTransactionalProducer()
     producer.initTransactions()
-    removeAllAcls()
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)), topicResource)
-    try {
-      producer.beginTransaction()
-      producer.send(new ProducerRecord(tp.topic, tp.partition, "1".getBytes, "1".getBytes)).get
-      Assert.fail("expected TransactionalIdAuthorizationException")
-    } catch {
-      case e: ExecutionException => assertTrue(s"expected TransactionalIdAuthorizationException, but got ${e.getCause}",
-        e.getCause.isInstanceOf[TransactionalIdAuthorizationException])
-    }
+    removeAllClientAcls()
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    producer.beginTransaction()
+    val future = producer.send(new ProducerRecord(tp.topic, tp.partition, "1".getBytes, "1".getBytes))
+    JTestUtils.assertFutureThrows(future, classOf[TransactionalIdAuthorizationException])
   }
 
   @Test
   def shouldThrowTransactionalIdAuthorizationExceptionWhenNoTransactionAccessOnEndTransaction(): Unit = {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)), transactionalIdResource)
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)), topicResource)
+    createTopic(topic)
+
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), transactionalIdResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
     val producer = buildTransactionalProducer()
     producer.initTransactions()
     producer.beginTransaction()
     producer.send(new ProducerRecord(tp.topic, tp.partition, "1".getBytes, "1".getBytes)).get
-    removeAllAcls()
-    try {
+    removeAllClientAcls()
+    intercept[TransactionalIdAuthorizationException] {
       producer.commitTransaction()
-      Assert.fail("expected TransactionalIdAuthorizationException")
-    } catch {
-      case _: TransactionalIdAuthorizationException => // ok
     }
   }
 
   @Test
   def shouldSuccessfullyAbortTransactionAfterTopicAuthorizationException(): Unit = {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)), transactionalIdResource)
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)), topicResource)
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Describe)), new Resource(Topic, deleteTopic))
+    createTopic(topic)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), transactionalIdResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, DESCRIBE, ALLOW)),
+      new ResourcePattern(TOPIC, topic, LITERAL))
     val producer = buildTransactionalProducer()
     producer.initTransactions()
     producer.beginTransaction()
     producer.send(new ProducerRecord(tp.topic, tp.partition, "1".getBytes, "1".getBytes)).get
     // try and add a partition resulting in TopicAuthorizationException
-    try {
-      producer.send(new ProducerRecord(deleteTopic, 0, "1".getBytes, "1".getBytes)).get
-    } catch {
-      case e: ExecutionException =>
-        assertTrue(e.getCause.isInstanceOf[TopicAuthorizationException])
-    }
+    val future = producer.send(new ProducerRecord("otherTopic", 0, "1".getBytes, "1".getBytes))
+    val e = JTestUtils.assertFutureThrows(future, classOf[TopicAuthorizationException])
+    assertEquals(Set("otherTopic"), e.unauthorizedTopics.asScala)
     // now rollback
     producer.abortTransaction()
   }
 
   @Test
   def shouldThrowTransactionalIdAuthorizationExceptionWhenNoTransactionAccessOnSendOffsetsToTxn(): Unit = {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)), transactionalIdResource)
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)), groupResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), transactionalIdResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), groupResource)
     val producer = buildTransactionalProducer()
     producer.initTransactions()
     producer.beginTransaction()
-    removeAllAcls()
-    try {
-      val offsets: util.Map[TopicPartition, OffsetAndMetadata] = Map(new TopicPartition(tp.topic, tp.partition) -> new OffsetAndMetadata(1L)).asJava
+    removeAllClientAcls()
+    intercept[TransactionalIdAuthorizationException] {
+      val offsets = Map(tp -> new OffsetAndMetadata(1L)).asJava
       producer.sendOffsetsToTransaction(offsets, group)
-      Assert.fail("expected TransactionalIdAuthorizationException")
-    } catch {
-      case _: TransactionalIdAuthorizationException => // ok
+      producer.commitTransaction()
     }
   }
 
   @Test
   def shouldSendSuccessfullyWhenIdempotentAndHasCorrectACL(): Unit = {
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, IdempotentWrite)), Resource.ClusterResource)
-    addAndVerifyAcls(Set(new Acl(KafkaPrincipal.ANONYMOUS, Allow, Acl.WildCardHost, Write)), topicResource)
+    createTopic(topic)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, IDEMPOTENT_WRITE, ALLOW)), clusterResource)
+    addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WildcardHost, WRITE, ALLOW)), topicResource)
     val producer = buildIdempotentProducer()
     producer.send(new ProducerRecord(tp.topic, tp.partition, "1".getBytes, "1".getBytes)).get
   }
 
-  def removeAllAcls() = {
-    servers.head.apis.authorizer.get.getAcls().keys.foreach { resource =>
-      servers.head.apis.authorizer.get.removeAcls(resource)
-      TestUtils.waitAndVerifyAcls(Set.empty[Acl], servers.head.apis.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 response = connectAndReceive[MetadataResponse](request)
+    assertEquals(Collections.emptyMap, response.errorCounts)
+    assertFalse("Cluster id not returned", response.clusterId.isEmpty)
+  }
+
+  def removeAllClientAcls(): Unit = {
+    val authorizer = servers.head.dataPlaneRequestProcessor.authorizer.get
+    val aclEntryFilter = new AccessControlEntryFilter(clientPrincipalString, null, AclOperation.ANY, AclPermissionType.ANY)
+    val aclFilter = new AclBindingFilter(ResourcePatternFilter.ANY, aclEntryFilter)
+
+    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, aclEntryFilter)
     }
   }
 
-  def sendRequestAndVerifyResponseError(apiKey: ApiKeys,
-                                        request: AbstractRequest,
-                                        resources: Set[ResourceType],
-                                        isAuthorized: Boolean,
-                                        topicExists: Boolean = true): AbstractResponse = {
-    val resp = connectAndSend(request, apiKey)
-    val response = requestKeyToResponseDeserializer(apiKey).getMethod("parse", classOf[ByteBuffer], classOf[Short]).invoke(
-      null, resp, request.version: java.lang.Short).asInstanceOf[AbstractResponse]
-    val error = requestKeyToError(apiKey).asInstanceOf[(AbstractResponse) => Errors](response)
+  private def sendRequestAndVerifyResponseError(request: AbstractRequest,
+                                                resources: Set[ResourceType],
+                                                isAuthorized: Boolean,
+                                                topicExists: Boolean = true): AbstractResponse = {
+    val apiKey = request.apiKey
+    val response = connectAndReceive[AbstractResponse](request)
+    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)
+          Set(Errors.UNKNOWN_TOPIC_OR_PARTITION, AclEntry.authorizationError(ResourceType.TOPIC))
         else
-          Set(Topic.error)
+          Set(AclEntry.authorizationError(ResourceType.TOPIC))
       } else {
-        Set(resourceType.error)
+        Set(AclEntry.authorizationError(resourceType))
       }
     }
 
@@ -1268,7 +1828,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
@@ -1277,9 +1837,11 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
     response
   }
 
-  private def sendRecords(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 =>
-      this.producers.head.send(new ProducerRecord(tp.topic(), tp.partition(), i.toString.getBytes, i.toString.getBytes))
+      producer.send(new ProducerRecord(tp.topic(), tp.partition(), i.toString.getBytes, i.toString.getBytes))
     }
     try {
       futures.foreach(_.get)
@@ -1288,57 +1850,44 @@ class AuthorizerIntegrationTest extends BaseRequestTest {
     }
   }
 
-  private def addAndVerifyAcls(acls: Set[Acl], resource: Resource) = {
-    servers.head.apis.authorizer.get.addAcls(acls, resource)
-    TestUtils.waitAndVerifyAcls(servers.head.apis.authorizer.get.getAcls(resource) ++ acls, servers.head.apis.authorizer.get, resource)
+  private def addAndVerifyAcls(acls: Set[AccessControlEntry], resource: ResourcePattern): Unit = {
+    TestUtils.addAndVerifyAcls(servers.head, acls, resource)
   }
 
   private def consumeRecords(consumer: Consumer[Array[Byte], Array[Byte]],
                              numRecords: Int = 1,
                              startingOffset: Int = 0,
                              topic: String = topic,
-                             part: Int = part) {
-    val records = new ArrayList[ConsumerRecord[Array[Byte], Array[Byte]]]()
-
-    TestUtils.waitUntilTrue(() => {
-      for (record <- consumer.poll(50).asScala)
-        records.add(record)
-      records.size == numRecords
-    }, "Failed to receive all expected records from the consumer")
+                             part: Int = part): Unit = {
+    val records = TestUtils.consumeRecords(consumer, numRecords)
 
     for (i <- 0 until numRecords) {
-      val record = records.get(i)
+      val record = records(i)
       val offset = startingOffset + i
-      assertEquals(topic, record.topic())
-      assertEquals(part, record.partition())
-      assertEquals(offset.toLong, record.offset())
+      assertEquals(topic, record.topic)
+      assertEquals(part, record.partition)
+      assertEquals(offset.toLong, record.offset)
     }
   }
 
-  private def sendOffsetFetchRequest(request: requests.OffsetFetchRequest,
-                                     socketServer: SocketServer): requests.OffsetFetchResponse = {
-    val response = connectAndSend(request, ApiKeys.OFFSET_FETCH, socketServer)
-    requests.OffsetFetchResponse.parse(response, request.version)
-  }
-
   private def buildTransactionalProducer(): KafkaProducer[Array[Byte], Array[Byte]] = {
-    val transactionalProperties = new Properties()
-    transactionalProperties.setProperty(ProducerConfig.TRANSACTIONAL_ID_CONFIG, transactionalId)
-    val producer = TestUtils.createNewProducer(TestUtils.getBrokerListStrFromServers(servers),
-      retries = 3,
-      props = Some(transactionalProperties))
-    producers += producer
-    producer
+    producerConfig.setProperty(ProducerConfig.TRANSACTIONAL_ID_CONFIG, transactionalId)
+    producerConfig.setProperty(ProducerConfig.ACKS_CONFIG, "all")
+    createProducer()
   }
 
   private def buildIdempotentProducer(): KafkaProducer[Array[Byte], Array[Byte]] = {
-    val idempotentProperties = new Properties()
-    idempotentProperties.setProperty(ProducerConfig.ENABLE_IDEMPOTENCE_CONFIG, "true")
-    val producer = TestUtils.createNewProducer(TestUtils.getBrokerListStrFromServers(servers),
-      retries = 3,
-      props = Some(idempotentProperties))
-    producers += producer
-    producer
+    producerConfig.setProperty(ProducerConfig.ENABLE_IDEMPOTENCE_CONFIG, "true")
+    producerConfig.setProperty(ProducerConfig.ACKS_CONFIG, "all")
+    createProducer()
+  }
+
+  private def createAdminClient(): Admin = {
+    val props = new Properties()
+    props.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList)
+    val adminClient = Admin.create(props)
+    adminClients += adminClient
+    adminClient
   }
 
 }
diff --git a/core/src/test/scala/integration/kafka/api/BaseAdminIntegrationTest.scala b/core/src/test/scala/integration/kafka/api/BaseAdminIntegrationTest.scala
new file mode 100644
index 0000000000000..a1fc75bebf8b1
--- /dev/null
+++ b/core/src/test/scala/integration/kafka/api/BaseAdminIntegrationTest.scala
@@ -0,0 +1,236 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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
+import java.util.Properties
+import java.util.concurrent.ExecutionException
+
+import kafka.security.authorizer.AclEntry
+import kafka.server.KafkaConfig
+import kafka.utils.Logging
+import kafka.utils.TestUtils._
+import org.apache.kafka.clients.admin.{Admin, AdminClientConfig, CreateTopicsOptions, CreateTopicsResult, DescribeClusterOptions, DescribeTopicsOptions, NewTopic, TopicDescription}
+import org.apache.kafka.common.acl.AclOperation
+import org.apache.kafka.common.errors.{TopicExistsException, UnknownTopicOrPartitionException}
+import org.apache.kafka.common.resource.ResourceType
+import org.apache.kafka.common.utils.Utils
+import org.junit.Assert._
+import org.junit.rules.Timeout
+import org.junit.{After, Before, Rule, Test}
+
+import scala.jdk.CollectionConverters._
+import scala.collection.Seq
+import scala.compat.java8.OptionConverters._
+
+/**
+ * Base integration test cases for [[Admin]]. Each test case added here will be executed
+ * in extending classes. Typically we prefer to write basic Admin functionality test cases in
+ * [[kafka.api.PlaintextAdminIntegrationTest]] rather than here to avoid unnecessary execution
+ * time to the build. However, if an admin API involves differing interactions with
+ * authentication/authorization layers, we may add the test case here.
+ */
+abstract class BaseAdminIntegrationTest extends IntegrationTestHarness with Logging {
+  def brokerCount = 3
+  override def logDirCount = 2
+
+  var client: Admin = _
+
+  @Rule
+  def globalTimeout: Timeout = Timeout.millis(120000)
+
+  @Before
+  override def setUp(): Unit = {
+    super.setUp()
+    waitUntilBrokerMetadataIsPropagated(servers)
+  }
+
+  @After
+  override def tearDown(): Unit = {
+    if (client != null)
+      Utils.closeQuietly(client, "AdminClient")
+    super.tearDown()
+  }
+
+  @Test
+  def testCreateDeleteTopics(): Unit = {
+    client = Admin.create(createConfig)
+    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.toShort),
+      new NewTopic("mytopic3", Option.empty[Integer].asJava, Option.empty[java.lang.Short].asJava)
+    )
+    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("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 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)
+
+    val topic0 = topicToDescription.get("mytopic")
+    assertEquals(false, topic0.isInternal)
+    assertEquals("mytopic", topic0.name)
+    assertEquals(2, topic0.partitions.size)
+    val topic0Partition0 = topic0.partitions.get(0)
+    assertEquals(1, topic0Partition0.leader.id)
+    assertEquals(0, topic0Partition0.partition)
+    assertEquals(Seq(1, 2), topic0Partition0.isr.asScala.map(_.id))
+    assertEquals(Seq(1, 2), topic0Partition0.replicas.asScala.map(_.id))
+    val topic0Partition1 = topic0.partitions.get(1)
+    assertEquals(2, topic0Partition1.leader.id)
+    assertEquals(1, topic0Partition1.partition)
+    assertEquals(Seq(2, 0), topic0Partition1.isr.asScala.map(_.id))
+    assertEquals(Seq(2, 0), topic0Partition1.replicas.asScala.map(_.id))
+
+    val topic1 = topicToDescription.get("mytopic2")
+    assertEquals(false, topic1.isInternal)
+    assertEquals("mytopic2", topic1.name)
+    assertEquals(3, topic1.partitions.size)
+    for (partitionId <- 0 until 3) {
+      val partition = topic1.partitions.get(partitionId)
+      assertEquals(partitionId, partition.partition)
+      assertEquals(3, partition.replicas.size)
+      partition.replicas.forEach { replica =>
+        assertTrue(replica.id >= 0)
+        assertTrue(replica.id < brokerCount)
+      }
+      assertEquals("No duplicate replica ids", partition.replicas.size, partition.replicas.asScala.map(_.id).distinct.size)
+
+      assertEquals(3, partition.isr.size)
+      assertEquals(partition.replicas, partition.isr)
+      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 testAuthorizedOperations(): Unit = {
+    client = Admin.create(createConfig)
+
+    // without includeAuthorizedOperations flag
+    var result = client.describeCluster
+    assertNull(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)
+    assertNull(topicResult.authorizedOperations)
+
+    //with includeAuthorizedOperations flag
+    topicResult = getTopicMetadata(client, topic, new DescribeTopicsOptions().includeAuthorizedOperations(true))
+    expectedOperations = AclEntry.supportedOperations(ResourceType.TOPIC).asJava
+    assertEquals(expectedOperations, topicResult.authorizedOperations)
+  }
+
+  def configuredClusterPermissions: Set[AclOperation] =
+    AclEntry.supportedOperations(ResourceType.CLUSTER)
+
+  override def modifyConfigs(configs: Seq[Properties]): Unit = {
+    super.modifyConfigs(configs)
+    configs.foreach { config =>
+      config.setProperty(KafkaConfig.DeleteTopicEnableProp, "true")
+      config.setProperty(KafkaConfig.GroupInitialRebalanceDelayMsProp, "0")
+      config.setProperty(KafkaConfig.AutoLeaderRebalanceEnableProp, "false")
+      config.setProperty(KafkaConfig.ControlledShutdownEnableProp, "false")
+      // 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))
+        config.setProperty(KafkaConfig.SslTruststorePasswordProp, "some.invalid.pass")
+    }
+  }
+
+  def createConfig: util.Map[String, Object] = {
+    val config = new util.HashMap[String, Object]
+    config.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList)
+    config.put(AdminClientConfig.REQUEST_TIMEOUT_MS_CONFIG, "20000")
+    val securityProps: util.Map[Object, Object] =
+      adminClientSecurityConfigs(securityProtocol, trustStoreFile, clientSaslProperties)
+    securityProps.forEach { (key, value) => config.put(key.asInstanceOf[String], value) }
+    config
+  }
+
+  def waitForTopics(client: Admin, expectedPresent: Seq[String], expectedMissing: Seq[String]): Unit = {
+    waitUntilTrue(() => {
+      val topics = client.listTopics.names.get()
+      expectedPresent.forall(topicName => topics.contains(topicName)) &&
+        expectedMissing.forall(topicName => !topics.contains(topicName))
+    }, "timed out waiting for topics")
+  }
+
+  def getTopicMetadata(client: Admin,
+                       topic: String,
+                       describeOptions: DescribeTopicsOptions = new DescribeTopicsOptions,
+                       expectedNumPartitionsOpt: Option[Int] = None): TopicDescription = {
+    var result: TopicDescription = null
+    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/BaseConsumerTest.scala b/core/src/test/scala/integration/kafka/api/BaseConsumerTest.scala
index e124468b923de..c66f3d11b9d9e 100644
--- a/core/src/test/scala/integration/kafka/api/BaseConsumerTest.scala
+++ b/core/src/test/scala/integration/kafka/api/BaseConsumerTest.scala
@@ -1,113 +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
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this 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.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.serialization.ByteArrayDeserializer
-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.jdk.CollectionConverters._
+import scala.collection.Seq
 
 /**
- * Integration tests for the new 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 producerCount = 1
-  val consumerCount = 2
-  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, "0")
-  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
-    TestUtils.createTopic(this.zkUtils, topic, 2, serverCount, this.servers)
-  }
+abstract class BaseConsumerTest extends AbstractConsumerTest {
 
   @Test
-  def testSimpleConsumption() {
+  def testSimpleConsumption(): Unit = {
     val numRecords = 10000
-    sendRecords(numRecords)
+    val producer = createProducer()
+    sendRecords(producer, numRecords, tp)
 
-    assertEquals(0, this.consumers.head.assignment.size)
-    this.consumers.head.assign(List(tp).asJava)
-    assertEquals(1, this.consumers.head.assignment.size)
+    val consumer = createConsumer()
+    assertEquals(0, consumer.assignment.size)
+    consumer.assign(List(tp).asJava)
+    assertEquals(1, consumer.assignment.size)
 
-    this.consumers.head.seek(tp, 0)
-    consumeAndVerifyRecords(consumer = this.consumers.head, numRecords = numRecords, startingOffset = 0)
+    consumer.seek(tp, 0)
+    consumeAndVerifyRecords(consumer = consumer, numRecords = numRecords, startingOffset = 0)
 
     // check async commit callbacks
-    val commitCallback = new CountConsumerCommitCallback()
-    this.consumers.head.commitAsync(commitCallback)
-    awaitCommitCallback(this.consumers.head, commitCallback)
+    sendAndAwaitAsyncCommit(consumer)
   }
 
   @Test
-  def testCoordinatorFailover() {
+  def testCoordinatorFailover(): Unit = {
     val listener = new TestConsumerReassignmentListener()
-    this.consumerConfig.setProperty(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG, "5000")
-    this.consumerConfig.setProperty(ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG, "2000")
-    val consumer0 = new KafkaConsumer(this.consumerConfig, new ByteArrayDeserializer(), new ByteArrayDeserializer())
-    consumers += consumer0
+    this.consumerConfig.setProperty(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG, "5001")
+    this.consumerConfig.setProperty(ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG, "1000")
+    // Use higher poll timeout to avoid consumer leaving the group due to timeout
+    this.consumerConfig.setProperty(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG, "15000")
+    val consumer = createConsumer()
 
-    consumer0.subscribe(List(topic).asJava, listener)
+    consumer.subscribe(List(topic).asJava, listener)
 
     // the initial subscription should cause a callback execution
-    consumer0.poll(2000)
-
+    awaitRebalance(consumer, listener)
     assertEquals(1, listener.callsToAssigned)
 
     // get metadata for the topic
     var parts: Seq[PartitionInfo] = null
     while (parts == null)
-      parts = consumer0.partitionsFor(Topic.GROUP_METADATA_TOPIC_NAME).asScala
+      parts = consumer.partitionsFor(Topic.GROUP_METADATA_TOPIC_NAME).asScala
     assertEquals(1, parts.size)
     assertNotNull(parts.head.leader())
 
@@ -115,212 +74,7 @@ abstract class BaseConsumerTest extends IntegrationTestHarness {
     val coordinator = parts.head.leader().id()
     this.servers(coordinator).shutdown()
 
-    consumer0.poll(5000)
-
     // the failover should not cause a rebalance
-    assertEquals(1, listener.callsToAssigned)
-    assertEquals(1, listener.callsToRevoked)
-  }
-
-  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
-    }
+    ensureNoRebalance(consumer, listener)
   }
-
-  protected def sendRecords(numRecords: Int): Seq[ProducerRecord[Array[Byte], Array[Byte]]] =
-    sendRecords(numRecords, tp)
-
-  protected def sendRecords(numRecords: Int, tp: TopicPartition): Seq[ProducerRecord[Array[Byte], Array[Byte]]] =
-    sendRecords(this.producers.head, numRecords, tp)
-
-  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]]
-    val maxIters = numRecords * 300
-    var iters = 0
-    while (records.size < numRecords) {
-      val polledRecords = consumer.poll(50).asScala
-      assertTrue(polledRecords.size <= maxPollRecords)
-      for (record <- polledRecords)
-        records += record
-      if (iters > maxIters)
-        throw new IllegalStateException("Failed to consume the expected records after " + iters + " iterations.")
-      iters += 1
-    }
-    records
-  }
-
-  protected def awaitCommitCallback[K, V](consumer: Consumer[K, V],
-                                          commitCallback: CountConsumerCommitCallback,
-                                          count: Int = 1): Unit = {
-    val started = System.currentTimeMillis()
-    while (commitCallback.successCount < count && System.currentTimeMillis() - started < 10000)
-      consumer.poll(50)
-    assertEquals(count, commitCallback.successCount)
-  }
-
-  protected class CountConsumerCommitCallback extends OffsetCommitCallback {
-    var successCount = 0
-    var failCount = 0
-
-    override def onComplete(offsets: util.Map[TopicPartition, OffsetAndMetadata], exception: Exception): Unit = {
-      if (exception == null)
-        successCount += 1
-      else
-        failCount += 1
-    }
-  }
-
-  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(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 0badda9157638..cec3e8d7b54ca 100644
--- a/core/src/test/scala/integration/kafka/api/BaseProducerSendTest.scala
+++ b/core/src/test/scala/integration/kafka/api/BaseProducerSendTest.scala
@@ -17,24 +17,29 @@
 
 package kafka.api
 
+import java.time.Duration
 import java.nio.charset.StandardCharsets
 import java.util.Properties
 import java.util.concurrent.TimeUnit
 
-import collection.JavaConverters._
 import kafka.integration.KafkaServerTestHarness
 import kafka.log.LogConfig
 import kafka.server.KafkaConfig
 import kafka.utils.TestUtils
-import org.apache.kafka.clients.consumer.{ConsumerRecord, KafkaConsumer}
+import org.apache.kafka.clients.consumer.KafkaConsumer
 import org.apache.kafka.clients.producer._
-import org.apache.kafka.common.{KafkaException, TopicPartition}
+import org.apache.kafka.common.errors.TimeoutException
 import org.apache.kafka.common.record.TimestampType
 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.mutable.{ArrayBuffer, Buffer}
+import scala.annotation.nowarn
+import scala.jdk.CollectionConverters._
+import scala.collection.mutable.Buffer
+import scala.concurrent.ExecutionException
 
 abstract class BaseProducerSendTest extends KafkaServerTestHarness {
 
@@ -53,13 +58,13 @@ abstract class BaseProducerSendTest extends KafkaServerTestHarness {
   private val numRecords = 100
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
-    consumer = TestUtils.createNewConsumer(TestUtils.getBrokerListStrFromServers(servers), securityProtocol = SecurityProtocol.PLAINTEXT)
+    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())
@@ -67,9 +72,23 @@ abstract class BaseProducerSendTest extends KafkaServerTestHarness {
     super.tearDown()
   }
 
-  protected def createProducer(brokerList: String, retries: Int = 0, lingerMs: Long = 0, props: Option[Properties] = None): KafkaProducer[Array[Byte],Array[Byte]] = {
-    val producer = TestUtils.createNewProducer(brokerList, securityProtocol = securityProtocol, trustStoreFile = trustStoreFile,
-      saslProperties = clientSaslProperties, retries = retries, lingerMs = lingerMs, props = props)
+  protected def createProducer(brokerList: String,
+                               lingerMs: Int = 0,
+                               deliveryTimeoutMs: Int = 2 * 60 * 1000,
+                               batchSize: Int = 16384,
+                               compressionType: String = "none",
+                               maxBlockMs: Long = 60 * 1000L,
+                               bufferSize: Long = 1024L * 1024L): KafkaProducer[Array[Byte],Array[Byte]] = {
+    val producer = TestUtils.createProducer(brokerList,
+      compressionType = compressionType,
+      securityProtocol = securityProtocol,
+      trustStoreFile = trustStoreFile,
+      saslProperties = clientSaslProperties,
+      lingerMs = lingerMs,
+      deliveryTimeoutMs = deliveryTimeoutMs,
+      maxBlockMs = maxBlockMs,
+      batchSize = batchSize,
+      bufferSize = bufferSize)
     registerProducer(producer)
   }
 
@@ -78,30 +97,22 @@ abstract class BaseProducerSendTest extends KafkaServerTestHarness {
     producer
   }
 
-  private def pollUntilNumRecords(numRecords: Int) : Seq[ConsumerRecord[Array[Byte], Array[Byte]]] = {
-    val records = new ArrayBuffer[ConsumerRecord[Array[Byte], Array[Byte]]]()
-    TestUtils.waitUntilTrue(() => {
-      records ++= consumer.poll(50).asScala
-      records.size == numRecords
-    }, s"Consumed ${records.size} records until timeout, but expected $numRecords records.")
-    records
-  }
-
   /**
    * testSendOffset checks the basic send API behavior
    *
    * 1. Send with null key/value/partition-id should be accepted; send with null topic should be rejected.
    * 2. Last message of the non-blocking send should return the correct offset metadata
    */
+  @nowarn("cat=deprecation")
   @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())
@@ -123,7 +134,7 @@ abstract class BaseProducerSendTest extends KafkaServerTestHarness {
 
     try {
       // create topic
-      TestUtils.createTopic(zkUtils, topic, 1, 2, servers)
+      createTopic(topic, 1, 2)
 
       // send a normal record
       val record0 = new ProducerRecord[Array[Byte], Array[Byte]](topic, partition, "key".getBytes(StandardCharsets.UTF_8),
@@ -166,32 +177,33 @@ abstract class BaseProducerSendTest extends KafkaServerTestHarness {
   }
 
   @Test
-  def testSendCompressedMessageWithCreateTime() {
-    val producerProps = new Properties()
-    producerProps.setProperty(ProducerConfig.COMPRESSION_TYPE_CONFIG, "gzip")
-    val producer = createProducer(brokerList = brokerList, lingerMs = Long.MaxValue, props = Some(producerProps))
+  def testSendCompressedMessageWithCreateTime(): Unit = {
+    val producer = createProducer(brokerList = brokerList,
+      compressionType = "gzip",
+      lingerMs = Int.MaxValue,
+      deliveryTimeoutMs = Int.MaxValue)
     sendAndVerifyTimestamp(producer, TimestampType.CREATE_TIME)
   }
 
   @Test
-  def testSendNonCompressedMessageWithCreateTime() {
-    val producer = createProducer(brokerList = brokerList, lingerMs = Long.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 {
-      TestUtils.createTopic(zkUtils, topic, 1, 2, servers)
+      createTopic(topic, 1, 2)
 
       val futures = for (i <- 1 to numRecords) yield {
         val record = new ProducerRecord(topic, partition, s"key$i".getBytes(StandardCharsets.UTF_8),
           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)
@@ -205,7 +217,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
@@ -215,7 +227,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)
@@ -239,14 +251,14 @@ abstract class BaseProducerSendTest extends KafkaServerTestHarness {
         topicProps.setProperty(LogConfig.MessageTimestampTypeProp, "LogAppendTime")
       else
         topicProps.setProperty(LogConfig.MessageTimestampTypeProp, "CreateTime")
-      TestUtils.createTopic(zkUtils, topic, 1, 2, servers, topicProps)
+      createTopic(topic, 1, 2, topicProps)
 
       val recordAndFutures = for (i <- 1 to numRecords) yield {
         val record = new ProducerRecord(topic, partition, baseTimestamp + i, s"key$i".getBytes(StandardCharsets.UTF_8),
           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)
@@ -266,12 +278,12 @@ 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 {
       // create topic
-      TestUtils.createTopic(zkUtils, topic, 1, 2, servers)
+      createTopic(topic, 1, 2)
 
       // non-blocking send a list of records
       val record0 = new ProducerRecord[Array[Byte], Array[Byte]](topic, null, "key".getBytes(StandardCharsets.UTF_8),
@@ -299,11 +311,11 @@ abstract class BaseProducerSendTest extends KafkaServerTestHarness {
    * The specified partition-id should be respected
    */
   @Test
-  def testSendToPartition() {
+  def testSendToPartition(): Unit = {
     val producer = createProducer(brokerList)
 
     try {
-      TestUtils.createTopic(zkUtils, topic, 2, 2, servers)
+      createTopic(topic, 2, 2)
       val partition = 1
 
       val now = System.currentTimeMillis()
@@ -321,7 +333,7 @@ abstract class BaseProducerSendTest extends KafkaServerTestHarness {
       consumer.assign(List(new TopicPartition(topic, partition)).asJava)
 
       // make sure the fetched messages also respect the partitioning and ordering
-      val records = pollUntilNumRecords(numRecords)
+      val records = TestUtils.consumeRecords(consumer, numRecords)
 
       records.zipWithIndex.foreach { case (record, i) =>
         assertEquals(topic, record.topic)
@@ -344,11 +356,11 @@ abstract class BaseProducerSendTest extends KafkaServerTestHarness {
     * succeed as long as the partition is included in the metadata.
     */
   @Test
-  def testSendBeforeAndAfterPartitionExpansion() {
-    val producer = createProducer(brokerList)
+  def testSendBeforeAndAfterPartitionExpansion(): Unit = {
+    val producer = createProducer(brokerList, maxBlockMs = 5 * 1000L)
 
     // create topic
-    TestUtils.createTopic(zkUtils, topic, 1, 2, servers)
+    createTopic(topic, 1, 2)
     val partition0 = 0
 
     var futures0 = (1 to numRecords).map { i =>
@@ -365,14 +377,17 @@ abstract class BaseProducerSendTest extends KafkaServerTestHarness {
     // Trying to send a record to a partition beyond topic's partition range before adding the partition should fail.
     val partition1 = 1
     try {
-      producer.send(new ProducerRecord(topic, partition1, null, "value".getBytes(StandardCharsets.UTF_8)))
+      producer.send(new ProducerRecord(topic, partition1, null, "value".getBytes(StandardCharsets.UTF_8))).get()
       fail("Should not allow sending a record to a partition not present in the metadata")
     } catch {
-      case _: KafkaException => // this is ok
+      case e: ExecutionException => e.getCause match {
+        case _: TimeoutException => // this is ok
+        case ex => throw new Exception("Sending to a partition not present in the metadata should result in a TimeoutException", ex)
+      }
     }
 
-    val existingAssignment = zkUtils.getReplicaAssignmentForTopics(List(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
@@ -407,10 +422,10 @@ abstract class BaseProducerSendTest extends KafkaServerTestHarness {
    * Test that flush immediately sends all accumulated requests.
    */
   @Test
-  def testFlush() {
-    val producer = createProducer(brokerList, lingerMs = Long.MaxValue)
+  def testFlush(): Unit = {
+    val producer = createProducer(brokerList, lingerMs = Int.MaxValue, deliveryTimeoutMs = Int.MaxValue)
     try {
-      TestUtils.createTopic(zkUtils, topic, 2, 2, servers)
+      createTopic(topic, 2, 2)
       val record = new ProducerRecord[Array[Byte], Array[Byte]](topic,
         "value".getBytes(StandardCharsets.UTF_8))
       for (_ <- 0 until 50) {
@@ -428,8 +443,8 @@ abstract class BaseProducerSendTest extends KafkaServerTestHarness {
    * Test close with zero timeout from caller thread
    */
   @Test
-  def testCloseWithZeroTimeoutFromCallerThread() {
-    TestUtils.createTopic(zkUtils, topic, 2, 2, servers)
+  def testCloseWithZeroTimeoutFromCallerThread(): Unit = {
+    createTopic(topic, 2, 2)
     val partition = 0
     consumer.assign(List(new TopicPartition(topic, partition)).asJava)
     val record0 = new ProducerRecord[Array[Byte], Array[Byte]](topic, partition, null,
@@ -437,20 +452,19 @@ abstract class BaseProducerSendTest extends KafkaServerTestHarness {
 
     // Test closing from caller thread.
     for (_ <- 0 until 50) {
-      val producer = createProducer(brokerList, lingerMs = Long.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()
           fail("No message should be sent successfully.")
         } catch {
-          case e: Exception =>
-            assertEquals("java.lang.IllegalStateException: Producer is closed forcefully.", e.getMessage)
+          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)
     }
   }
 
@@ -458,27 +472,27 @@ abstract class BaseProducerSendTest extends KafkaServerTestHarness {
    * Test close with zero and non-zero timeout from sender thread
    */
   @Test
-  def testCloseWithZeroTimeoutFromSenderThread() {
-    TestUtils.createTopic(zkUtils, topic, 1, 2, servers)
+  def testCloseWithZeroTimeoutFromSenderThread(): Unit = {
+    createTopic(topic, 1, 2)
     val partition = 0
     consumer.assign(List(new TopicPartition(topic, partition)).asJava)
     val record = new ProducerRecord[Array[Byte], Array[Byte]](topic, partition, null, "value".getBytes(StandardCharsets.UTF_8))
 
     // 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 = Long.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
@@ -489,7 +503,7 @@ abstract class BaseProducerSendTest extends KafkaServerTestHarness {
         producer.flush()
         assertTrue("All requests are complete.", responses.forall(_.isDone()))
         // Check the messages received by broker.
-        pollUntilNumRecords(numRecords)
+        TestUtils.pollUntilAtLeastNumRecords(consumer, numRecords)
       } finally {
         producer.close()
       }
diff --git a/core/src/test/scala/integration/kafka/api/BaseQuotaTest.scala b/core/src/test/scala/integration/kafka/api/BaseQuotaTest.scala
index d821f528637f7..23e649777658d 100644
--- a/core/src/test/scala/integration/kafka/api/BaseQuotaTest.scala
+++ b/core/src/test/scala/integration/kafka/api/BaseQuotaTest.scala
@@ -14,33 +14,39 @@
 
 package kafka.api
 
+import java.time.Duration
+import java.util.concurrent.TimeUnit
 import java.util.{Collections, HashMap, Properties}
 
-import kafka.server.{ClientQuotaManagerConfig, DynamicConfig, KafkaConfig, KafkaServer, QuotaId, QuotaType}
+import com.yammer.metrics.core.{Histogram, Meter}
+import kafka.api.QuotaTestClients._
+import kafka.metrics.KafkaYammerMetrics
+import kafka.server.{ClientQuotaManager, ClientQuotaManagerConfig, DynamicConfig, KafkaConfig, KafkaServer, QuotaType}
 import kafka.utils.TestUtils
+import org.apache.kafka.clients.admin.Admin
 import org.apache.kafka.clients.consumer.{ConsumerConfig, KafkaConsumer}
 import org.apache.kafka.clients.producer._
 import org.apache.kafka.clients.producer.internals.ErrorLoggingCallback
-import org.apache.kafka.common.{MetricName, TopicPartition}
+import org.apache.kafka.common.{Metric, MetricName, TopicPartition}
 import org.apache.kafka.common.metrics.{KafkaMetric, Quota}
+import org.apache.kafka.common.protocol.ApiKeys
+import org.apache.kafka.common.quota.ClientQuotaAlteration
+import org.apache.kafka.common.quota.ClientQuotaEntity
+import org.apache.kafka.common.security.auth.KafkaPrincipal
 import org.junit.Assert._
 import org.junit.{Before, Test}
+import org.scalatest.Assertions.fail
 
-abstract class BaseQuotaTest extends IntegrationTestHarness {
+import scala.collection.Map
+import scala.jdk.CollectionConverters._
 
-  def userPrincipal : String
-  def producerQuotaId : QuotaId
-  def consumerQuotaId : QuotaId
-  def overrideQuotas(producerQuota: Long, consumerQuota: Long, requestQuota: Double)
-  def removeQuotaOverrides()
+abstract class BaseQuotaTest extends IntegrationTestHarness {
 
-  override val serverCount = 2
-  val producerCount = 1
-  val consumerCount = 1
+  override val brokerCount = 2
 
-  private val producerBufferSize = 300000
   protected def producerClientId = "QuotasTestProducer-1"
   protected def consumerClientId = "QuotasTestConsumer-1"
+  protected def createQuotaTestClients(topic: String, leaderNode: KafkaServer): QuotaTestClients
 
   this.serverConfig.setProperty(KafkaConfig.ControlledShutdownEnableProp, "false")
   this.serverConfig.setProperty(KafkaConfig.OffsetsTopicReplicationFactorProp, "2")
@@ -48,8 +54,8 @@ abstract class BaseQuotaTest extends IntegrationTestHarness {
   this.serverConfig.setProperty(KafkaConfig.GroupMinSessionTimeoutMsProp, "100")
   this.serverConfig.setProperty(KafkaConfig.GroupMaxSessionTimeoutMsProp, "30000")
   this.serverConfig.setProperty(KafkaConfig.GroupInitialRebalanceDelayMsProp, "0")
-  this.producerConfig.setProperty(ProducerConfig.ACKS_CONFIG, "0")
-  this.producerConfig.setProperty(ProducerConfig.BUFFER_MEMORY_CONFIG, producerBufferSize.toString)
+  this.producerConfig.setProperty(ProducerConfig.ACKS_CONFIG, "-1")
+  this.producerConfig.setProperty(ProducerConfig.BUFFER_MEMORY_CONFIG, "300000")
   this.producerConfig.setProperty(ProducerConfig.CLIENT_ID_CONFIG, producerClientId)
   this.consumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "QuotasTest")
   this.consumerConfig.setProperty(ConsumerConfig.MAX_PARTITION_FETCH_BYTES_CONFIG, 4096.toString)
@@ -59,202 +65,320 @@ abstract class BaseQuotaTest extends IntegrationTestHarness {
   this.consumerConfig.setProperty(ConsumerConfig.FETCH_MAX_WAIT_MS_CONFIG, "0")
 
   // Low enough quota that a producer sending a small payload in a tight loop should get throttled
-  val defaultProducerQuota = 8000
-  val defaultConsumerQuota = 2500
-  val defaultRequestQuota = Int.MaxValue
+  val defaultProducerQuota: Long = 8000
+  val defaultConsumerQuota: Long = 2500
+  val defaultRequestQuota: Double = Long.MaxValue.toDouble
 
-  var leaderNode: KafkaServer = null
-  var followerNode: KafkaServer = null
-  private val topic1 = "topic-1"
+  val topic1 = "topic-1"
+  var leaderNode: KafkaServer = _
+  var followerNode: KafkaServer = _
+  var quotaTestClients: QuotaTestClients = _
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
 
     val numPartitions = 1
-    val leaders = TestUtils.createTopic(zkUtils, topic1, numPartitions, serverCount, servers)
+    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 producer = producers.head
-    val produced = produceUntilThrottled(producer, numRecords)
-    assertTrue("Should have been throttled", producerThrottleMetric.value > 0)
-    verifyProducerThrottleTimeMetric(producer)
+    val produced = quotaTestClients.produceUntilThrottled(numRecords)
+    quotaTestClients.verifyProduceThrottle(expectThrottle = true)
 
     // Consumer should read in a bursty manner and get throttled immediately
-    val consumer = consumers.head
-    consumeUntilThrottled(consumer, produced)
-    assertTrue("Should have been throttled", consumerThrottleMetric.value > 0)
-    verifyConsumerThrottleTimeMetric(consumer)
+    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)
     props.put(DynamicConfig.Client.ConsumerByteRateOverrideProp, Long.MaxValue.toString)
 
-    overrideQuotas(Long.MaxValue, Long.MaxValue, Int.MaxValue)
-    waitForQuotaUpdate(Long.MaxValue, Long.MaxValue, Int.MaxValue)
+    quotaTestClients.overrideQuotas(Long.MaxValue, Long.MaxValue, Long.MaxValue.toDouble)
+    quotaTestClients.waitForQuotaUpdate(Long.MaxValue, Long.MaxValue, Long.MaxValue.toDouble)
 
     val numRecords = 1000
-    assertEquals(numRecords, produceUntilThrottled(producers.head, numRecords))
-    assertEquals("Should not have been throttled", 0.0, producerThrottleMetric.value, 0.0)
+    assertEquals(numRecords, quotaTestClients.produceUntilThrottled(numRecords))
+    quotaTestClients.verifyProduceThrottle(expectThrottle = false)
 
     // The "client" consumer does not get throttled.
-    assertEquals(numRecords, consumeUntilThrottled(consumers.head, numRecords))
-    assertEquals("Should not have been throttled", 0.0, consumerThrottleMetric.value, 0.0)
+    assertEquals(numRecords, quotaTestClients.consumeUntilThrottled(numRecords))
+    quotaTestClients.verifyConsumeThrottle(expectThrottle = false)
   }
 
   @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, Long.MaxValue.toDouble)
+    quotaTestClients.waitForQuotaUpdate(2000, 250, Long.MaxValue.toDouble)
+
+    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
-    overrideQuotas(Long.MaxValue, Long.MaxValue, Int.MaxValue)
-    waitForQuotaUpdate(Long.MaxValue, Long.MaxValue, Int.MaxValue)
+    quotaTestClients.overrideQuotas(Long.MaxValue, Long.MaxValue, Long.MaxValue.toDouble)
+    quotaTestClients.waitForQuotaUpdate(Long.MaxValue, Long.MaxValue, Long.MaxValue.toDouble)
 
     val numRecords = 1000
-    assertEquals(numRecords, produceUntilThrottled(producers.head, numRecords))
-    assertEquals("Should not have been throttled", 0.0, producerThrottleMetric.value, 0.0)
-    assertEquals(numRecords, consumeUntilThrottled(consumers.head, numRecords))
-    assertEquals("Should not have been throttled", 0.0, consumerThrottleMetric.value, 0.0)
+    assertEquals(numRecords, quotaTestClients.produceUntilThrottled(numRecords))
+    quotaTestClients.verifyProduceThrottle(expectThrottle = false)
+    assertEquals(numRecords, quotaTestClients.consumeUntilThrottled(numRecords))
+    quotaTestClients.verifyConsumeThrottle(expectThrottle = false)
 
     // Delete producer and consumer quota overrides. Consumer and producer should now be
     // throttled since broker defaults are very small
-    removeQuotaOverrides()
-    val produced = produceUntilThrottled(producers.head, numRecords)
-    assertTrue("Should have been throttled", producerThrottleMetric.value > 0)
+    quotaTestClients.removeQuotaOverrides()
+    quotaTestClients.waitForQuotaUpdate(defaultProducerQuota, defaultConsumerQuota, defaultRequestQuota)
+    val produced = quotaTestClients.produceUntilThrottled(numRecords)
+    quotaTestClients.verifyProduceThrottle(expectThrottle = true)
 
     // Since producer may have been throttled after producing a couple of records,
     // consume from beginning till throttled
-    consumers.head.seekToBeginning(Collections.singleton(new TopicPartition(topic1, 0)))
-    consumeUntilThrottled(consumers.head, numRecords + produced)
-    assertTrue("Should have been throttled", consumerThrottleMetric.value > 0)
+    quotaTestClients.consumer.seekToBeginning(Collections.singleton(new TopicPartition(topic1, 0)))
+    quotaTestClients.consumeUntilThrottled(numRecords + produced)
+    quotaTestClients.verifyConsumeThrottle(expectThrottle = true)
   }
 
   @Test
-  def testThrottledRequest() {
-
-    overrideQuotas(Long.MaxValue, Long.MaxValue, 0.1)
-    waitForQuotaUpdate(Long.MaxValue, Long.MaxValue, 0.1)
+  def testThrottledRequest(): Unit = {
+    quotaTestClients.overrideQuotas(Long.MaxValue, Long.MaxValue, 0.1)
+    quotaTestClients.waitForQuotaUpdate(Long.MaxValue, Long.MaxValue, 0.1)
 
-    val consumer = consumers.head
+    val consumer = quotaTestClients.consumer
     consumer.subscribe(Collections.singleton(topic1))
     val endTimeMs = System.currentTimeMillis + 10000
     var throttled = false
-    while ((!throttled || exemptRequestMetric == null) && System.currentTimeMillis < endTimeMs) {
-      consumer.poll(100)
-      val throttleMetric = consumerRequestThrottleMetric
-      throttled = throttleMetric != null && throttleMetric.value > 0
+    while ((!throttled || quotaTestClients.exemptRequestMetric == null) && System.currentTimeMillis < endTimeMs) {
+      consumer.poll(Duration.ofMillis(100L))
+      val throttleMetric = quotaTestClients.throttleMetric(QuotaType.Request, consumerClientId)
+      throttled = throttleMetric != null && metricValue(throttleMetric) > 0
     }
 
     assertTrue("Should have been throttled", throttled)
-    verifyConsumerThrottleTimeMetric(consumer, Some(ClientQuotaManagerConfig.DefaultQuotaWindowSizeSeconds * 1000.0))
+    quotaTestClients.verifyConsumerClientThrottleTimeMetric(expectThrottle = true,
+      Some(ClientQuotaManagerConfig.DefaultQuotaWindowSizeSeconds * 1000.0))
 
-    assertNotNull("Exempt requests not recorded", exemptRequestMetric)
-    assertTrue("Exempt requests not recorded", exemptRequestMetric.value > 0)
+    val exemptMetric = quotaTestClients.exemptRequestMetric
+    assertNotNull("Exempt requests not recorded", exemptMetric)
+    assertTrue("Exempt requests not recorded", metricValue(exemptMetric) > 0)
   }
+}
+
+object QuotaTestClients {
+  val DefaultEntity: String = null
 
-  def produceUntilThrottled(p: KafkaProducer[Array[Byte], Array[Byte]], maxRecords: Int): Int = {
+  def metricValue(metric: Metric): Double = metric.metricValue().asInstanceOf[Double]
+}
+
+abstract class QuotaTestClients(topic: String,
+                                leaderNode: KafkaServer,
+                                producerClientId: String,
+                                consumerClientId: String,
+                                val producer: KafkaProducer[Array[Byte], Array[Byte]],
+                                val consumer: KafkaConsumer[Array[Byte], Array[Byte]],
+                                val adminClient: Admin) {
+
+  def overrideQuotas(producerQuota: Long, consumerQuota: Long, requestQuota: Double): Unit
+  def removeQuotaOverrides(): Unit
+
+  protected def userPrincipal: KafkaPrincipal
+  protected def quotaMetricTags(clientId: String): Map[String, String]
+
+  def produceUntilThrottled(maxRecords: Int, waitForRequestCompletion: Boolean = true): Int = {
     var numProduced = 0
     var throttled = false
     do {
       val payload = numProduced.toString.getBytes
-      p.send(new ProducerRecord[Array[Byte], Array[Byte]](topic1, null, null, payload),
-             new ErrorLoggingCallback(topic1, null, null, true)).get()
+      val future = producer.send(new ProducerRecord[Array[Byte], Array[Byte]](topic, null, null, payload),
+        new ErrorLoggingCallback(topic, null, null, true))
       numProduced += 1
-      val throttleMetric = producerThrottleMetric
-      throttled = throttleMetric != null && throttleMetric.value > 0
+      do {
+        val metric = throttleMetric(QuotaType.Produce, producerClientId)
+        throttled = metric != null && metricValue(metric) > 0
+      } while (!future.isDone && (!throttled || waitForRequestCompletion))
     } while (numProduced < maxRecords && !throttled)
     numProduced
   }
 
-  def consumeUntilThrottled(consumer: KafkaConsumer[Array[Byte], Array[Byte]], maxRecords: Int): Int = {
-    consumer.subscribe(Collections.singleton(topic1))
+  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
-      val throttleMetric = consumerThrottleMetric
-      throttled = throttleMetric != null && throttleMetric.value > 0
-    }  while (numConsumed < maxRecords && !throttled)
+      numConsumed += consumer.poll(Duration.ofMillis(100L)).count
+      val metric = throttleMetric(QuotaType.Fetch, consumerClientId)
+      throttled = metric != null && metricValue(metric) > 0
+    } 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) {
+    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
   }
 
-  def waitForQuotaUpdate(producerQuota: Long, consumerQuota: Long, requestQuota: Double) {
-    TestUtils.retry(10000) {
-      val quotaManagers = leaderNode.apis.quotas
-      val overrideProducerQuota = quotaManagers.produce.quota(userPrincipal, producerClientId)
-      val overrideConsumerQuota = quotaManagers.fetch.quota(userPrincipal, consumerClientId)
-      val overrideProducerRequestQuota = quotaManagers.request.quota(userPrincipal, producerClientId)
-      val overrideConsumerRequestQuota = quotaManagers.request.quota(userPrincipal, consumerClientId)
-
-      assertEquals(s"ClientId $producerClientId of user $userPrincipal must have producer quota", Quota.upperBound(producerQuota), overrideProducerQuota)
-      assertEquals(s"ClientId $consumerClientId of user $userPrincipal must have consumer quota", Quota.upperBound(consumerQuota), overrideConsumerQuota)
-      assertEquals(s"ClientId $producerClientId of user $userPrincipal must have request quota", Quota.upperBound(requestQuota), overrideProducerRequestQuota)
-      assertEquals(s"ClientId $consumerClientId of user $userPrincipal must have request quota", Quota.upperBound(requestQuota), overrideConsumerRequestQuota)
+  private def quota(quotaManager: ClientQuotaManager, userPrincipal: KafkaPrincipal, clientId: String): Quota = {
+    quotaManager.quota(userPrincipal, clientId)
+  }
+
+  private def verifyThrottleTimeRequestChannelMetric(apiKey: ApiKeys, metricNameSuffix: String,
+                                                     clientId: String, expectThrottle: Boolean): Unit = {
+    val throttleTimeMs = brokerRequestMetricsThrottleTimeMs(apiKey, metricNameSuffix)
+    if (expectThrottle)
+      assertTrue(s"Client with id=$clientId should have been throttled, $throttleTimeMs", throttleTimeMs > 0)
+    else
+      assertEquals(s"Client with id=$clientId should not have been throttled", 0.0, throttleTimeMs, 0.0)
+  }
+
+  def verifyProduceThrottle(expectThrottle: Boolean, verifyClientMetric: Boolean = true,
+                            verifyRequestChannelMetric: Boolean = true): Unit = {
+    verifyThrottleTimeMetric(QuotaType.Produce, producerClientId, expectThrottle)
+    if (verifyRequestChannelMetric)
+      verifyThrottleTimeRequestChannelMetric(ApiKeys.PRODUCE, "", producerClientId, expectThrottle)
+    if (verifyClientMetric)
+      verifyProducerClientThrottleTimeMetric(expectThrottle)
+  }
+
+  def verifyConsumeThrottle(expectThrottle: Boolean, verifyClientMetric: Boolean = true,
+                            verifyRequestChannelMetric: Boolean = true): Unit = {
+    verifyThrottleTimeMetric(QuotaType.Fetch, consumerClientId, expectThrottle)
+    if (verifyRequestChannelMetric)
+      verifyThrottleTimeRequestChannelMetric(ApiKeys.FETCH, "Consumer", consumerClientId, expectThrottle)
+    if (verifyClientMetric)
+      verifyConsumerClientThrottleTimeMetric(expectThrottle)
+  }
+
+  private def verifyThrottleTimeMetric(quotaType: QuotaType, clientId: String, expectThrottle: Boolean): Unit = {
+    val throttleMetricValue = metricValue(throttleMetric(quotaType, clientId))
+    if (expectThrottle) {
+      assertTrue(s"Client with id=$clientId should have been throttled", throttleMetricValue > 0)
+    } else {
+      assertTrue(s"Client with id=$clientId should not have been throttled", throttleMetricValue.isNaN)
+    }
+  }
+
+  private def throttleMetricName(quotaType: QuotaType, clientId: String): MetricName = {
+    leaderNode.metrics.metricName("throttle-time",
+      quotaType.toString,
+      quotaMetricTags(clientId).asJava)
+  }
+
+  def throttleMetric(quotaType: QuotaType, clientId: String): KafkaMetric = {
+    leaderNode.metrics.metrics.get(throttleMetricName(quotaType, clientId))
+  }
+
+  private def brokerRequestMetricsThrottleTimeMs(apiKey: ApiKeys, metricNameSuffix: String): Double = {
+    def yammerMetricValue(name: String): Double = {
+      val allMetrics = KafkaYammerMetrics.defaultRegistry.allMetrics.asScala
+      val (_, metric) = allMetrics.find { case (metricName, _) =>
+        metricName.getMBeanName.startsWith(name)
+      }.getOrElse(fail(s"Unable to find broker metric $name: allMetrics: ${allMetrics.keySet.map(_.getMBeanName)}"))
+      metric match {
+        case m: Meter => m.count.toDouble
+        case m: Histogram => m.max
+        case m => fail(s"Unexpected broker metric of class ${m.getClass}")
+      }
     }
+
+    yammerMetricValue(s"kafka.network:type=RequestMetrics,name=ThrottleTimeMs,request=${apiKey.name}$metricNameSuffix")
   }
 
-  private def verifyProducerThrottleTimeMetric(producer: KafkaProducer[_, _]) {
+  def exemptRequestMetric: KafkaMetric = {
+    val metricName = leaderNode.metrics.metricName("exempt-request-time", QuotaType.Request.toString, "")
+    leaderNode.metrics.metrics.get(metricName)
+  }
+
+  private 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))
     val maxMetric = producer.metrics.get(new MetricName("produce-throttle-time-max", "producer-metrics", "", tags))
 
-    TestUtils.waitUntilTrue(() => avgMetric.value > 0.0 && maxMetric.value > 0.0,
-        s"Producer throttle metric not updated: avg=${avgMetric.value} max=${maxMetric.value}")
+    if (expectThrottle) {
+      TestUtils.waitUntilTrue(() => metricValue(avgMetric) > 0.0 && metricValue(maxMetric) > 0.0,
+        s"Producer throttle metric not updated: avg=${metricValue(avgMetric)} max=${metricValue(maxMetric)}")
+    } else
+      assertEquals("Should not have been throttled", 0.0, metricValue(maxMetric), 0.0)
   }
 
-  private def verifyConsumerThrottleTimeMetric(consumer: KafkaConsumer[_, _], 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))
     val maxMetric = consumer.metrics.get(new MetricName("fetch-throttle-time-max", "consumer-fetch-manager-metrics", "", tags))
 
-    TestUtils.waitUntilTrue(() => avgMetric.value > 0.0 && maxMetric.value > 0.0,
-        s"Consumer throttle metric not updated: avg=${avgMetric.value} max=${maxMetric.value}")
-    maxThrottleTime.foreach(max => assertTrue(s"Maximum consumer throttle too high: ${maxMetric.value}", maxMetric.value <= max))
+    if (expectThrottle) {
+      TestUtils.waitUntilTrue(() => metricValue(avgMetric) > 0.0 && metricValue(maxMetric) > 0.0,
+        s"Consumer throttle metric not updated: avg=${metricValue(avgMetric)} max=${metricValue(maxMetric)}")
+      maxThrottleTime.foreach(max => assertTrue(s"Maximum consumer throttle too high: ${metricValue(maxMetric)}",
+        metricValue(maxMetric) <= max))
+    } else
+      assertEquals("Should not have been throttled", 0.0, metricValue(maxMetric), 0.0)
   }
 
-  private def throttleMetricName(quotaType: QuotaType, quotaId: QuotaId): MetricName = {
-    leaderNode.metrics.metricName("throttle-time",
-                                  quotaType.toString,
-                                  "Tracking throttle-time per user/client-id",
-                                  "user", quotaId.sanitizedUser.getOrElse(""),
-                                  "client-id", quotaId.clientId.getOrElse(""))
+  def clientQuotaEntity(user: Option[String], clientId: Option[String]): ClientQuotaEntity = {
+    var entries = Map.empty[String, String]
+    user.foreach(user => entries = entries ++ Map(ClientQuotaEntity.USER -> user))
+    clientId.foreach(clientId => entries = entries ++ Map(ClientQuotaEntity.CLIENT_ID -> clientId))
+    new ClientQuotaEntity(entries.asJava)
   }
 
-  def throttleMetric(quotaType: QuotaType, quotaId: QuotaId): KafkaMetric = {
-    leaderNode.metrics.metrics.get(throttleMetricName(quotaType, quotaId))
+  // None is translated to `null` which remove the quota
+  def clientQuotaAlteration(quotaEntity: ClientQuotaEntity,
+                            producerQuota: Option[Long],
+                            consumerQuota: Option[Long],
+                            requestQuota: Option[Double]): ClientQuotaAlteration = {
+    var ops = Seq.empty[ClientQuotaAlteration.Op]
+    def addOp(key: String, value: Option[Double]): Unit = {
+      ops = ops ++ Seq(new ClientQuotaAlteration.Op(key, value.map(Double.box).orNull))
+    }
+    addOp(DynamicConfig.Client.ProducerByteRateOverrideProp, producerQuota.map(_.toDouble))
+    addOp(DynamicConfig.Client.ConsumerByteRateOverrideProp, consumerQuota.map(_.toDouble))
+    addOp(DynamicConfig.Client.RequestPercentageOverrideProp, requestQuota)
+    new ClientQuotaAlteration(quotaEntity, ops.asJava)
   }
 
-  private def producerThrottleMetric = throttleMetric(QuotaType.Produce, producerQuotaId)
-  private def consumerThrottleMetric = throttleMetric(QuotaType.Fetch, consumerQuotaId)
-  private def consumerRequestThrottleMetric = throttleMetric(QuotaType.Request, consumerQuotaId)
-
-  private def exemptRequestMetric: KafkaMetric = {
-    val metricName = leaderNode.metrics.metricName("exempt-request-time", QuotaType.Request.toString, "")
-    leaderNode.metrics.metrics.get(metricName)
+  def alterClientQuotas(quotaAlterations: ClientQuotaAlteration *): Unit = {
+    adminClient.alterClientQuotas(quotaAlterations.asJava).all().get()
   }
 
-  def quotaProperties(producerQuota: Long, consumerQuota: Long, requestQuota: Double): Properties = {
-    val props = new Properties()
-    props.put(DynamicConfig.Client.ProducerByteRateOverrideProp, producerQuota.toString)
-    props.put(DynamicConfig.Client.ConsumerByteRateOverrideProp, consumerQuota.toString)
-    props.put(DynamicConfig.Client.RequestPercentageOverrideProp, requestQuota.toString)
-    props
+  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)
+      val overrideConsumerQuota = quota(quotaManagers.fetch, userPrincipal, consumerClientId)
+      val overrideProducerRequestQuota = quota(quotaManagers.request, userPrincipal, producerClientId)
+      val overrideConsumerRequestQuota = quota(quotaManagers.request, userPrincipal, consumerClientId)
+
+      assertEquals(s"ClientId $producerClientId of user $userPrincipal must have producer quota", Quota.upperBound(producerQuota.toDouble), overrideProducerQuota)
+      assertEquals(s"ClientId $consumerClientId of user $userPrincipal must have consumer quota", Quota.upperBound(consumerQuota.toDouble), overrideConsumerQuota)
+      assertEquals(s"ClientId $producerClientId of user $userPrincipal must have request quota", Quota.upperBound(requestQuota.toDouble), overrideProducerRequestQuota)
+      assertEquals(s"ClientId $consumerClientId of user $userPrincipal must have request quota", Quota.upperBound(requestQuota.toDouble), overrideConsumerRequestQuota)
+    }
   }
 }
diff --git a/core/src/test/scala/integration/kafka/api/ClientIdQuotaTest.scala b/core/src/test/scala/integration/kafka/api/ClientIdQuotaTest.scala
index 3e0832790b5de..19a58f9b5c2aa 100644
--- a/core/src/test/scala/integration/kafka/api/ClientIdQuotaTest.scala
+++ b/core/src/test/scala/integration/kafka/api/ClientIdQuotaTest.scala
@@ -14,45 +14,59 @@
 
 package kafka.api
 
-import java.util.Properties
-
-import kafka.server.{DynamicConfig, KafkaConfig, QuotaId}
+import kafka.server.{KafkaConfig, KafkaServer}
 import org.apache.kafka.common.security.auth.KafkaPrincipal
-import org.apache.kafka.common.utils.Sanitizer
 import org.junit.Before
 
 class ClientIdQuotaTest extends BaseQuotaTest {
 
-  override val userPrincipal = KafkaPrincipal.ANONYMOUS.getName
   override def producerClientId = "QuotasTestProducer-!@#$%^&*()"
   override def consumerClientId = "QuotasTestConsumer-!@#$%^&*()"
-  override val producerQuotaId = QuotaId(None, Some(producerClientId), Some(Sanitizer.sanitize(producerClientId)))
-  override val consumerQuotaId = QuotaId(None, Some(consumerClientId), Some(Sanitizer.sanitize(consumerClientId)))
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     this.serverConfig.setProperty(KafkaConfig.ProducerQuotaBytesPerSecondDefaultProp, defaultProducerQuota.toString)
     this.serverConfig.setProperty(KafkaConfig.ConsumerQuotaBytesPerSecondDefaultProp, defaultConsumerQuota.toString)
     super.setUp()
   }
-  override def overrideQuotas(producerQuota: Long, consumerQuota: Long, requestQuota: Double) {
-    val producerProps = new Properties()
-    producerProps.put(DynamicConfig.Client.ProducerByteRateOverrideProp, producerQuota.toString)
-    producerProps.put(DynamicConfig.Client.RequestPercentageOverrideProp, requestQuota.toString)
-    updateQuotaOverride(producerClientId, producerProps)
-
-    val consumerProps = new Properties()
-    consumerProps.put(DynamicConfig.Client.ConsumerByteRateOverrideProp, consumerQuota.toString)
-    consumerProps.put(DynamicConfig.Client.RequestPercentageOverrideProp, requestQuota.toString)
-    updateQuotaOverride(consumerClientId, consumerProps)
-  }
-  override def removeQuotaOverrides() {
-    val emptyProps = new Properties
-    updateQuotaOverride(producerClientId, emptyProps)
-    updateQuotaOverride(consumerClientId, emptyProps)
-  }
 
-  private def updateQuotaOverride(clientId: String, properties: Properties) {
-    adminZkClient.changeClientIdConfig(Sanitizer.sanitize(clientId), properties)
+  override def createQuotaTestClients(topic: String, leaderNode: KafkaServer): QuotaTestClients = {
+    val producer = createProducer()
+    val consumer = createConsumer()
+    val adminClient = createAdminClient()
+
+    new QuotaTestClients(topic, leaderNode, producerClientId, consumerClientId, producer, consumer, adminClient) {
+      override def userPrincipal: KafkaPrincipal = KafkaPrincipal.ANONYMOUS
+
+      override def quotaMetricTags(clientId: String): Map[String, String] = {
+        Map("user" -> "", "client-id" -> clientId)
+      }
+
+      override def overrideQuotas(producerQuota: Long, consumerQuota: Long, requestQuota: Double): Unit = {
+        alterClientQuotas(
+          clientQuotaAlteration(
+            clientQuotaEntity(None, Some(producerClientId)),
+            Some(producerQuota), None, Some(requestQuota)
+          ),
+          clientQuotaAlteration(
+            clientQuotaEntity(None, Some(consumerClientId)),
+            None, Some(consumerQuota), Some(requestQuota)
+          )
+        )
+      }
+
+      override def removeQuotaOverrides(): Unit = {
+        alterClientQuotas(
+          clientQuotaAlteration(
+            clientQuotaEntity(None, Some(producerClientId)),
+            None, None, None
+          ),
+          clientQuotaAlteration(
+            clientQuotaEntity(None, Some(consumerClientId)),
+            None, None, None
+          )
+        )
+      }
+    }
   }
 }
diff --git a/core/src/test/scala/integration/kafka/api/ConsumerBounceTest.scala b/core/src/test/scala/integration/kafka/api/ConsumerBounceTest.scala
index 0b421187625cc..51c308df4d01d 100644
--- a/core/src/test/scala/integration/kafka/api/ConsumerBounceTest.scala
+++ b/core/src/test/scala/integration/kafka/api/ConsumerBounceTest.scala
@@ -13,68 +13,61 @@
 
 package kafka.api
 
+import java.time
 import java.util.concurrent._
-import java.util.{Collection, Collections}
+import java.util.{Collection, Collections, Properties}
 
-import kafka.admin.AdminClient
 import kafka.server.KafkaConfig
-import kafka.utils.{CoreUtils, Logging, ShutdownableThread, TestUtils}
+import kafka.utils.{Logging, ShutdownableThread, TestUtils}
 import org.apache.kafka.clients.consumer._
-import org.apache.kafka.clients.producer.{ProducerConfig, ProducerRecord}
+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.Errors
+import org.apache.kafka.common.requests.{FindCoordinatorRequest, FindCoordinatorResponse}
 import org.junit.Assert._
-import org.junit.{After, Before, Ignore, Test}
-
-import scala.collection.JavaConverters._
+import org.junit.{After, Ignore, Test}
 
+import scala.annotation.nowarn
+import scala.jdk.CollectionConverters._
+import scala.collection.{Seq, mutable}
 
 /**
- * Integration tests for the new consumer that cover basic usage as well as server failures
+ * Integration tests for the consumer that cover basic usage as well as server failures
  */
-class ConsumerBounceTest extends IntegrationTestHarness with Logging {
-
-  val producerCount = 1
-  val consumerCount = 2
-  val serverCount = 3
-
-  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)
-
-  // configure the servers and clients
-  this.serverConfig.setProperty(KafkaConfig.OffsetsTopicReplicationFactorProp, "3") // don't want to lose offset
-  this.serverConfig.setProperty(KafkaConfig.OffsetsTopicPartitionsProp, "1")
-  this.serverConfig.setProperty(KafkaConfig.GroupMinSessionTimeoutMsProp, "10") // set small enough session timeout
-  this.serverConfig.setProperty(KafkaConfig.GroupInitialRebalanceDelayMsProp, "0")
-  this.serverConfig.setProperty(KafkaConfig.UncleanLeaderElectionEnableProp, "true")
-  this.serverConfig.setProperty(KafkaConfig.AutoCreateTopicsEnableProp, "false")
-  this.producerConfig.setProperty(ProducerConfig.ACKS_CONFIG, "all")
-  this.consumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "my-test")
-  this.consumerConfig.setProperty(ConsumerConfig.MAX_PARTITION_FETCH_BYTES_CONFIG, 4096.toString)
-  this.consumerConfig.setProperty(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG, "10000")
-  this.consumerConfig.setProperty(ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG, "3000")
-  this.consumerConfig.setProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest")
+  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 = {
-    FixedPortTestUtils.createBrokerConfigs(serverCount, zkConnect, enableControlledShutdown = false)
-      .map(KafkaConfig.fromProps(_, serverConfig))
+    generateKafkaConfigs()
   }
 
-  @Before
-  override def setUp() {
-    super.setUp()
-
-    // create the test topic with all the brokers as replicas
-    TestUtils.createTopic(this.zkUtils, topic, 1, serverCount, this.servers)
+  private def generateKafkaConfigs(maxGroupSize: String = maxGroupSize.toString): Seq[KafkaConfig] = {
+    val properties = new Properties
+    properties.put(KafkaConfig.OffsetsTopicReplicationFactorProp, "3") // don't want to lose offset
+    properties.put(KafkaConfig.OffsetsTopicPartitionsProp, "1")
+    properties.put(KafkaConfig.GroupMinSessionTimeoutMsProp, "10") // set small enough session timeout
+    properties.put(KafkaConfig.GroupInitialRebalanceDelayMsProp, "0")
+    properties.put(KafkaConfig.GroupMaxSizeProp, maxGroupSize)
+    properties.put(KafkaConfig.UncleanLeaderElectionEnableProp, "true")
+    properties.put(KafkaConfig.AutoCreateTopicsEnableProp, "false")
+
+    FixedPortTestUtils.createBrokerConfigs(brokerCount, zkConnect, enableControlledShutdown = false)
+      .map(KafkaConfig.fromProps(_, properties))
   }
 
   @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))
@@ -91,20 +84,21 @@ class ConsumerBounceTest extends IntegrationTestHarness with Logging {
    * 1. Produce a bunch of messages
    * 2. Then consume the messages while killing and restarting brokers at random
    */
-  def consumeWithBrokerFailures(numIters: Int) {
+  @nowarn("cat=deprecation")
+  def consumeWithBrokerFailures(numIters: Int): Unit = {
     val numRecords = 1000
-    sendRecords(numRecords)
-    this.producers.foreach(_.close)
+    val producer = createProducer()
+    sendRecords(producer, numRecords)
 
     var consumed = 0L
-    val consumer = this.consumers.head
+    val consumer = createConsumer()
 
     consumer.subscribe(Collections.singletonList(topic))
 
     val scheduler = new BounceBrokerScheduler(numIters)
     scheduler.start()
 
-    while (scheduler.isRunning.get()) {
+    while (scheduler.isRunning) {
       val records = consumer.poll(100).asScala
       assertEquals(Set(tp), consumer.assignment.asScala)
 
@@ -115,7 +109,7 @@ class ConsumerBounceTest extends IntegrationTestHarness 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())
@@ -129,24 +123,24 @@ class ConsumerBounceTest extends IntegrationTestHarness with Logging {
   @Test
   def testSeekAndCommitWithBrokerFailures() = seekAndCommitWithBrokerFailures(5)
 
-  def seekAndCommitWithBrokerFailures(numIters: Int) {
+  def seekAndCommitWithBrokerFailures(numIters: Int): Unit = {
     val numRecords = 1000
-    sendRecords(numRecords)
-    this.producers.foreach(_.close)
+    val producer = createProducer()
+    sendRecords(producer, numRecords)
 
-    val consumer = this.consumers.head
+    val consumer = createConsumer()
     consumer.assign(Collections.singletonList(tp))
     consumer.seek(tp, 0)
 
     // wait until all the followers have synced the last HW with leader
     TestUtils.waitUntilTrue(() => servers.forall(server =>
-      server.replicaManager.getReplica(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)
     scheduler.start()
 
-    while(scheduler.isRunning.get()) {
+    while(scheduler.isRunning) {
       val coin = TestUtils.random.nextInt(3)
       if (coin == 0) {
         info("Seeking to end of log")
@@ -160,29 +154,31 @@ class ConsumerBounceTest extends IntegrationTestHarness 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 = this.consumers.head
+    val consumer = createConsumer()
     consumer.subscribe(Collections.singleton(newtopic))
     executor.schedule(new Runnable {
-        def run() = TestUtils.createTopic(zkUtils, newtopic, serverCount, serverCount, servers)
+        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) {
         val futures = (0 until remainingRecords).map { i =>
-          this.producers.head.send(new ProducerRecord(topic, part, i.toString.getBytes, i.toString.getBytes))
+          producer.send(new ProducerRecord(topic, part, i.toString.getBytes, i.toString.getBytes))
         }
         futures.map { future =>
           try {
@@ -196,25 +192,29 @@ class ConsumerBounceTest extends IntegrationTestHarness with Logging {
       assertEquals(0, remainingRecords)
     }
 
+    val poller = new ConsumerAssignmentPoller(consumer, List(newtopic))
+    consumerPollers += poller
+    poller.start()
     sendRecords(numRecords, newtopic)
-    receiveRecords(consumer, numRecords, newtopic, 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, newtopic, 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
-    sendRecords(numRecords)
+    val producer = createProducer()
+    sendRecords(producer, numRecords)
 
     checkCloseGoodPath(numRecords, "group1")
     checkCloseWithCoordinatorFailure(numRecords, "group2", "group3")
@@ -226,9 +226,9 @@ class ConsumerBounceTest extends IntegrationTestHarness 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)
   }
@@ -239,23 +239,36 @@ class ConsumerBounceTest extends IntegrationTestHarness 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)
 
-    val adminClient = AdminClient.createSimplePlaintext(this.brokerList)
-    killBroker(adminClient.findCoordinator(dynamicGroup).id)
-    killBroker(adminClient.findCoordinator(manualGroup).id)
+    killBroker(findCoordinator(dynamicGroup))
+    killBroker(findCoordinator(manualGroup))
+
+    val future1 = submitCloseAndValidate(consumer1, Long.MaxValue, None, gracefulCloseTimeMs)
+
+    val future2 = submitCloseAndValidate(consumer2, Long.MaxValue, None, gracefulCloseTimeMs)
 
-    val future1 = submitCloseAndValidate(consumer1, Long.MaxValue, None, Some(gracefulCloseTimeMs))
-    val future2 = submitCloseAndValidate(consumer2, Long.MaxValue, None, Some(gracefulCloseTimeMs))
     future1.get
     future2.get
 
     restartDeadBrokers()
     checkClosedState(dynamicGroup, 0)
     checkClosedState(manualGroup, numRecords)
-    adminClient.close()
+  }
+
+  private def findCoordinator(group: String): Int = {
+    val request = new FindCoordinatorRequest.Builder(new FindCoordinatorRequestData()
+      .setKeyType(FindCoordinatorRequest.CoordinatorType.GROUP.id)
+      .setKey(group)).build()
+    var nodeId = -1
+    TestUtils.waitUntilTrue(() => {
+      val response = connectAndReceive[FindCoordinatorResponse](request)
+      nodeId = response.node.id
+      response.error == Errors.NONE
+    }, s"Failed to find coordinator for group $group")
+    nodeId
   }
 
   /**
@@ -263,7 +276,7 @@ class ConsumerBounceTest extends IntegrationTestHarness 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
@@ -274,51 +287,121 @@ class ConsumerBounceTest extends IntegrationTestHarness 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
   }
 
+  /**
+    * If we have a running consumer group of size N, configure consumer.group.max.size = N-1 and restart all brokers,
+    * the group should be forced to rebalance when it becomes hosted on a Coordinator with the new config.
+    * Then, 1 consumer should be left out of the group.
+    */
+  @Test
+  def testRollingBrokerRestartsWithSmallerMaxGroupSizeConfigDisruptsBigGroup(): Unit = {
+    val group = "group-max-size-test"
+    val topic = "group-max-size-test"
+    val maxGroupSize = 2
+    val consumerCount = maxGroupSize + 1
+    val partitionCount = 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 partitions = createTopicPartitions(topic, numPartitions = partitionCount, replicationFactor = brokerCount)
+
+    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)
+    for (serverIdx <- servers.indices) {
+      killBroker(serverIdx)
+      val config = newConfigs(serverIdx)
+      servers(serverIdx) = TestUtils.createServer(config, time = brokerTime(config.brokerId))
+      restartDeadBrokers()
+    }
+
+    def raisedExceptions: Seq[Throwable] = {
+      consumerPollers.flatten(_.thrownException)
+    }
+
+    // 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)
+
+    assertEquals(1, raisedExceptions.size)
+    assertTrue(raisedExceptions.head.isInstanceOf[GroupMaxSizeReachedException])
+  }
+
+  /**
+    * When we have the consumer group max size configured to X, the X+1th consumer trying to join should receive a fatal exception
+    */
+  @Test
+  def testConsumerReceivesFatalExceptionWhenGroupPassesMaxSize(): Unit = {
+    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
+    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
+    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)
+  }
+
   /**
    * Consumer is closed during rebalance. Close should leave group and close
    * immediately if rebalance is in progress. If brokers are not available,
    * close should terminate immediately without sending leave group.
    */
   @Test
-  def testCloseDuringRebalance() {
+  def testCloseDuringRebalance(): Unit = {
     val topic = "closetest"
-    TestUtils.createTopic(this.zkUtils, topic, 10, serverCount, this.servers)
+    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) {
+  @nowarn("cat=deprecation")
+  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)
+      executor.submit(() => {
+        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)
     }
 
     def createConsumerToRebalance(): Future[Any] = {
-      val consumer = createConsumer(groupId)
+      val consumer = createConsumerWithGroupId(groupId)
       val rebalanceSemaphore = new Semaphore(0)
       val future = subscribeAndPoll(consumer, Some(rebalanceSemaphore))
       // Wait for consumer to poll and trigger rebalance
@@ -327,15 +410,14 @@ class ConsumerBounceTest extends IntegrationTestHarness with Logging {
       assertFalse("Rebalance completed too early", future.isDone)
       future
     }
-
-    val consumer1 = createConsumer(groupId)
+    val consumer1 = createConsumerWithGroupId(groupId)
     waitForRebalance(2000, subscribeAndPoll(consumer1))
-    val consumer2 = createConsumer(groupId)
+    val consumer2 = createConsumerWithGroupId(groupId)
     waitForRebalance(2000, subscribeAndPoll(consumer2), consumer1)
     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)
@@ -353,39 +435,32 @@ class ConsumerBounceTest extends IntegrationTestHarness with Logging {
     closeFuture2.get(2000, TimeUnit.MILLISECONDS)
   }
 
-  private def createConsumer(groupId: String) : KafkaConsumer[Array[Byte], Array[Byte]] = {
-    this.consumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, groupId)
-    val consumer = createNewConsumer
-    consumers += consumer
-    consumer
-  }
+  private def createConsumerAndReceive(groupId: String, manualAssign: Boolean, numRecords: Int): KafkaConsumer[Array[Byte], Array[Byte]] = {
+    val consumer = createConsumerWithGroupId(groupId)
+    val consumerPoller = if (manualAssign)
+        subscribeConsumerAndStartPolling(consumer, List(), Set(tp))
+      else
+        subscribeConsumerAndStartPolling(consumer, List(topic))
 
-  private def createConsumerAndReceive(groupId: String, manualAssign: Boolean, numRecords: Int) : KafkaConsumer[Array[Byte], Array[Byte]] = {
-    val consumer = createConsumer(groupId)
-    if (manualAssign)
-      consumer.assign(Collections.singleton(tp))
-    else
-      consumer.subscribe(Collections.singleton(topic))
-    receiveRecords(consumer, numRecords)
+    receiveExactRecords(consumerPoller, numRecords)
+    consumerPoller.shutdown()
     consumer
   }
 
-  private def receiveRecords(consumer: KafkaConsumer[Array[Byte], Array[Byte]], numRecords: Int, topic: String = this.topic, timeoutMs: Long = 60000) {
-    var received = 0L
-    val endTimeMs = System.currentTimeMillis + timeoutMs
-    while (received < numRecords && System.currentTimeMillis < endTimeMs)
-      received += consumer.poll(1000).count()
-    assertEquals(numRecords, received)
+  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 {
+    executor.submit(() => {
       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)
       }
@@ -396,21 +471,21 @@ class ConsumerBounceTest extends IntegrationTestHarness 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 = createConsumer(groupId)
+    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)
-    assertTrue("Assigment did not complete on time", assignSemaphore.tryAcquire(1, TimeUnit.SECONDS))
+    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()
   }
 
@@ -431,12 +506,31 @@ class ConsumerBounceTest extends IntegrationTestHarness with Logging {
     }
   }
 
-  private def sendRecords(numRecords: Int, topic: String = this.topic) {
+  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): Unit = {
+    var partitionIndex = 0
+    def getPartition: Int = {
+      numPartitions match {
+        case Some(partitions) =>
+          val nextPart = partitionIndex % partitions
+          partitionIndex += 1
+          nextPart
+        case None => part
+      }
+    }
+
     val futures = (0 until numRecords).map { i =>
-      this.producers.head.send(new ProducerRecord(topic, part, i.toString.getBytes, i.toString.getBytes))
+      producer.send(new ProducerRecord(topic, getPartition, i.toString.getBytes, i.toString.getBytes))
     }
     futures.map(_.get)
   }
 
-
 }
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..c7b77c262a298
--- /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
new file mode 100644
index 0000000000000..4cb9e42cbaa86
--- /dev/null
+++ b/core/src/test/scala/integration/kafka/api/CustomQuotaCallbackTest.scala
@@ -0,0 +1,471 @@
+/**
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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.{lang, util}
+import java.util.concurrent.ConcurrentHashMap
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger}
+import java.util.Properties
+
+import kafka.api.GroupedUserPrincipalBuilder._
+import kafka.api.GroupedUserQuotaCallback._
+import kafka.server._
+import kafka.utils.JaasTestUtils.ScramLoginModule
+import kafka.utils.{JaasTestUtils, Logging, TestUtils}
+import kafka.zk.ConfigEntityChangeNotificationZNode
+import org.apache.kafka.clients.admin.{Admin, 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}
+import org.apache.kafka.common.config.SaslConfigs
+import org.apache.kafka.common.network.ListenerName
+import org.apache.kafka.common.security.auth._
+import org.apache.kafka.server.quota._
+import org.junit.Assert._
+import org.junit.{After, Before, Test}
+
+import scala.collection.mutable.ArrayBuffer
+import scala.jdk.CollectionConverters._
+
+class CustomQuotaCallbackTest extends IntegrationTestHarness with SaslSetup {
+
+  override protected def securityProtocol = SecurityProtocol.SASL_SSL
+  override protected def listenerName = new ListenerName("CLIENT")
+  override protected def interBrokerListenerName: ListenerName = new ListenerName("BROKER")
+
+  override protected lazy val trustStoreFile = Some(File.createTempFile("truststore", ".jks"))
+  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[Admin]()
+  private var producerWithoutQuota: KafkaProducer[Array[Byte], Array[Byte]] = _
+
+  val defaultRequestQuota = 1000
+  val defaultProduceQuota = 2000 * 1000 * 1000
+  val defaultConsumeQuota = 1000 * 1000 * 1000
+
+  @Before
+  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)
+    this.serverConfig.setProperty(KafkaConfig.ClientQuotaCallbackClassProp, classOf[GroupedUserQuotaCallback].getName)
+    this.serverConfig.setProperty(s"${listenerName.configPrefix}${KafkaConfig.PrincipalBuilderClassProp}",
+      classOf[GroupedUserPrincipalBuilder].getName)
+    this.serverConfig.setProperty(KafkaConfig.DeleteTopicEnableProp, "true")
+    super.setUp()
+    brokerList = TestUtils.bootstrapServers(servers, listenerName)
+
+    producerConfig.put(SaslConfigs.SASL_JAAS_CONFIG,
+      ScramLoginModule(JaasTestUtils.KafkaScramAdmin, JaasTestUtils.KafkaScramAdminPassword).toString)
+    producerWithoutQuota = createProducer()
+  }
+
+  @After
+  override def tearDown(): Unit = {
+    adminClients.foreach(_.close())
+    GroupedUserQuotaCallback.tearDown()
+    super.tearDown()
+  }
+
+  override def configureSecurityBeforeServersStart(): Unit = {
+    super.configureSecurityBeforeServersStart()
+    zkClient.makeSurePersistentPathExists(ConfigEntityChangeNotificationZNode.path)
+    createScramCredentials(zkConnect, JaasTestUtils.KafkaScramAdmin, JaasTestUtils.KafkaScramAdminPassword)
+  }
+
+  @Test
+  def testCustomQuotaCallback(): Unit = {
+    // Large quota override, should not throttle
+    var brokerId = 0
+    var user = createGroupWithOneUser("group0_user1", brokerId)
+    user.configureAndWaitForQuota(1000000, 2000000)
+    quotaLimitCalls.values.foreach(_.set(0))
+    user.produceConsume(expectProduceThrottle = false, expectConsumeThrottle = false)
+
+    // ClientQuotaCallback#quotaLimit is invoked by each quota manager once per throttled produce request for each client
+    assertEquals(1, quotaLimitCalls(ClientQuotaType.PRODUCE).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)
+    user.produceConsume(expectProduceThrottle = true, expectConsumeThrottle = true)
+
+    // Quota override deletion - verify default quota applied (large quota, no throttling)
+    user = addUser("group0_user2", brokerId)
+    user.removeQuotaOverrides()
+    user.waitForQuotaUpdate(defaultProduceQuota, defaultConsumeQuota, defaultRequestQuota)
+    user.removeThrottleMetrics() // since group was throttled before
+    user.produceConsume(expectProduceThrottle = false, expectConsumeThrottle = false)
+
+    // Make default quota smaller, should throttle
+    user.configureAndWaitForQuota(8000, 2500, divisor = 1, group = None)
+    user.produceConsume(expectProduceThrottle = true, expectConsumeThrottle = true)
+
+    // Configure large quota override, should not throttle
+    user = addUser("group0_user3", brokerId)
+    user.configureAndWaitForQuota(2000000, 2000000)
+    user.removeThrottleMetrics() // since group was throttled before
+    user.produceConsume(expectProduceThrottle = false, expectConsumeThrottle = false)
+
+    // Quota large enough for one partition, should not throttle
+    brokerId = 1
+    user = createGroupWithOneUser("group1_user1", brokerId)
+    user.configureAndWaitForQuota(8000 * 100, 2500 * 100)
+    user.produceConsume(expectProduceThrottle = false, expectConsumeThrottle = false)
+
+    // Create large number of partitions on another broker, should result in throttling on first partition
+    val largeTopic = "group1_largeTopic"
+    createTopic(largeTopic, numPartitions = 99, leader = 0)
+    user.waitForQuotaUpdate(8000, 2500, defaultRequestQuota)
+    user.produceConsume(expectProduceThrottle = true, expectConsumeThrottle = true)
+
+    // Remove quota override and test default quota applied with scaling based on partitions
+    user = addUser("group1_user2", brokerId)
+    user.waitForQuotaUpdate(defaultProduceQuota / 100, defaultConsumeQuota / 100, defaultRequestQuota)
+    user.removeThrottleMetrics() // since group was throttled before
+    user.produceConsume(expectProduceThrottle = false, expectConsumeThrottle = false)
+    user.configureAndWaitForQuota(8000 * 100, 2500 * 100, divisor=100, group = None)
+    user.produceConsume(expectProduceThrottle = true, expectConsumeThrottle = true)
+
+    // Remove the second topic with large number of partitions, verify no longer throttled
+    adminZkClient.deleteTopic(largeTopic)
+    user = addUser("group1_user3", brokerId)
+    user.waitForQuotaUpdate(8000 * 100, 2500 * 100, defaultRequestQuota)
+    user.removeThrottleMetrics() // since group was throttled before
+    user.produceConsume(expectProduceThrottle = false, expectConsumeThrottle = false)
+
+    // Alter configs of custom callback dynamically
+    val adminClient = createAdminClient()
+    val newProps = new Properties
+    newProps.put(GroupedUserQuotaCallback.DefaultProduceQuotaProp, "8000")
+    newProps.put(GroupedUserQuotaCallback.DefaultFetchQuotaProp, "2500")
+    TestUtils.incrementalAlterConfigs(servers, adminClient, newProps, perBrokerConfig = false)
+    user.waitForQuotaUpdate(8000, 2500, defaultRequestQuota)
+    user.produceConsume(expectProduceThrottle = true, expectConsumeThrottle = true)
+
+    assertEquals(brokerCount, callbackInstances.get)
+  }
+
+  /**
+   * Creates a group with one user and one topic with one partition.
+   * @param firstUser First user to create in the group
+   * @param brokerId The broker id to use as leader of the partition
+   */
+  private def createGroupWithOneUser(firstUser: String, brokerId: Int): GroupedUser = {
+    val user = addUser(firstUser, brokerId)
+    createTopic(user.topic, numPartitions = 1, brokerId)
+    user.configureAndWaitForQuota(defaultProduceQuota, defaultConsumeQuota, divisor = 1, group = None)
+    user
+  }
+
+  private def createTopic(topic: String, numPartitions: Int, leader: Int): Unit = {
+    val assignment = (0 until numPartitions).map { i => i -> Seq(leader) }.toMap
+    TestUtils.createTopic(zkClient, topic, assignment, servers)
+  }
+
+  private def createAdminClient(): Admin = {
+    val config = new util.HashMap[String, Object]
+    config.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG,
+      TestUtils.bootstrapServers(servers, new ListenerName("BROKER")))
+    clientSecurityProps("admin-client").asInstanceOf[util.Map[Object, Object]].forEach { (key, value) =>
+      config.put(key.toString, value)
+    }
+    config.put(SaslConfigs.SASL_JAAS_CONFIG,
+      ScramLoginModule(JaasTestUtils.KafkaScramAdmin, JaasTestUtils.KafkaScramAdminPassword).toString)
+    val adminClient = Admin.create(config)
+    adminClients += adminClient
+    adminClient
+  }
+
+  private def produceWithoutThrottle(topic: String, numRecords: Int): Unit = {
+    (0 until numRecords).foreach { i =>
+      val payload = i.toString.getBytes
+      producerWithoutQuota.send(new ProducerRecord[Array[Byte], Array[Byte]](topic, null, null, payload))
+    }
+  }
+
+  private def passwordForUser(user: String) = {
+    s"$user:secret"
+  }
+
+  private def addUser(user: String, leader: Int): GroupedUser = {
+    val adminClient = createAdminClient()
+    createScramCredentials(adminClient, user, passwordForUser(user))
+    waitForUserScramCredentialToAppearOnAllBrokers(user, kafkaClientSaslMechanism)
+    groupedUser(adminClient, user, leader)
+  }
+
+  private def groupedUser(adminClient: Admin, user: String, leader: Int): GroupedUser = {
+    val password = passwordForUser(user)
+    val userGroup = group(user)
+    val topic = s"${userGroup}_topic"
+    val producerClientId = s"$user:producer-client-id"
+    val consumerClientId = s"$user:consumer-client-id"
+
+    producerConfig.put(ProducerConfig.CLIENT_ID_CONFIG, producerClientId)
+    producerConfig.put(SaslConfigs.SASL_JAAS_CONFIG, ScramLoginModule(user, password).toString)
+
+    consumerConfig.put(ConsumerConfig.CLIENT_ID_CONFIG, consumerClientId)
+    consumerConfig.put(ConsumerConfig.MAX_PARTITION_FETCH_BYTES_CONFIG, 4096.toString)
+    consumerConfig.put(ConsumerConfig.GROUP_ID_CONFIG, s"$user-group")
+    consumerConfig.put(SaslConfigs.SASL_JAAS_CONFIG, ScramLoginModule(user, password).toString)
+
+    GroupedUser(user, userGroup, topic, servers(leader), producerClientId, consumerClientId,
+      createProducer(), createConsumer(), adminClient)
+  }
+
+  case class GroupedUser(user: String, userGroup: String, topic: String, leaderNode: KafkaServer,
+                         producerClientId: String, consumerClientId: String,
+                         override val producer: KafkaProducer[Array[Byte], Array[Byte]],
+                         override val consumer: KafkaConsumer[Array[Byte], Array[Byte]],
+                         override val adminClient: Admin) extends
+    QuotaTestClients(topic, leaderNode, producerClientId, consumerClientId, producer, consumer, adminClient) {
+
+    override def userPrincipal: KafkaPrincipal = GroupedUserPrincipal(user, userGroup)
+
+    override def quotaMetricTags(clientId: String): Map[String, String] = {
+      Map(GroupedUserQuotaCallback.QuotaGroupTag -> userGroup)
+    }
+
+    override def overrideQuotas(producerQuota: Long, consumerQuota: Long, requestQuota: Double): Unit = {
+      configureQuota(userGroup, producerQuota, consumerQuota, requestQuota)
+    }
+
+    override def removeQuotaOverrides(): Unit = {
+      alterClientQuotas(
+        clientQuotaAlteration(
+          clientQuotaEntity(Some(quotaEntityName(userGroup)), None),
+          None, None, None
+        )
+      )
+    }
+
+    def configureQuota(userGroup: String, producerQuota: Long, consumerQuota: Long, requestQuota: Double): Unit = {
+      alterClientQuotas(
+        clientQuotaAlteration(
+          clientQuotaEntity(Some(quotaEntityName(userGroup)), None),
+          Some(producerQuota), Some(consumerQuota), Some(requestQuota)
+        )
+      )
+    }
+
+    def configureAndWaitForQuota(produceQuota: Long, fetchQuota: Long, divisor: Int = 1,
+                                 group: Option[String] = Some(userGroup)): Unit = {
+      configureQuota(group.getOrElse(""), produceQuota, fetchQuota, defaultRequestQuota)
+      waitForQuotaUpdate(produceQuota / divisor, fetchQuota / divisor, defaultRequestQuota)
+    }
+
+    def produceConsume(expectProduceThrottle: Boolean, expectConsumeThrottle: Boolean): Unit = {
+      val numRecords = 1000
+      val produced = produceUntilThrottled(numRecords, waitForRequestCompletion = false)
+      // don't verify request channel metrics as it's difficult to write non flaky assertions
+      // given the specifics of this test (throttle metric removal followed by produce/consume
+      // until throttled)
+      verifyProduceThrottle(expectProduceThrottle, verifyClientMetric = false,
+        verifyRequestChannelMetric = false)
+      // make sure there are enough records on the topic to test consumer throttling
+      produceWithoutThrottle(topic, numRecords - produced)
+      consumeUntilThrottled(numRecords, waitForRequestCompletion = false)
+      verifyConsumeThrottle(expectConsumeThrottle, verifyClientMetric = false,
+        verifyRequestChannelMetric = false)
+    }
+
+    def removeThrottleMetrics(): Unit = {
+      def removeSensors(quotaType: QuotaType, clientId: String): Unit = {
+        val sensorSuffix = quotaMetricTags(clientId).values.mkString(":")
+        leaderNode.metrics.removeSensor(s"${quotaType}ThrottleTime-$sensorSuffix")
+        leaderNode.metrics.removeSensor(s"$quotaType-$sensorSuffix")
+      }
+      removeSensors(QuotaType.Produce, producerClientId)
+      removeSensors(QuotaType.Fetch, consumerClientId)
+      removeSensors(QuotaType.Request, producerClientId)
+      removeSensors(QuotaType.Request, consumerClientId)
+    }
+
+    private def quotaEntityName(userGroup: String): String = s"${userGroup}_"
+  }
+}
+
+object GroupedUserPrincipalBuilder {
+  def group(str: String): String = {
+    if (str.indexOf("_") <= 0)
+      ""
+    else
+      str.substring(0, str.indexOf("_"))
+  }
+}
+
+class GroupedUserPrincipalBuilder extends KafkaPrincipalBuilder {
+  override def build(context: AuthenticationContext): KafkaPrincipal = {
+    val securityProtocol = context.securityProtocol
+    if (securityProtocol == SecurityProtocol.SASL_PLAINTEXT || securityProtocol == SecurityProtocol.SASL_SSL) {
+      val user = context.asInstanceOf[SaslAuthenticationContext].server().getAuthorizationID
+      val userGroup = group(user)
+      if (userGroup.isEmpty)
+        new KafkaPrincipal(KafkaPrincipal.USER_TYPE, user)
+      else
+        GroupedUserPrincipal(user, userGroup)
+    } else
+      throw new IllegalStateException(s"Unexpected security protocol $securityProtocol")
+  }
+}
+
+case class GroupedUserPrincipal(user: String, userGroup: String) extends KafkaPrincipal(KafkaPrincipal.USER_TYPE, user)
+
+object GroupedUserQuotaCallback {
+  val QuotaGroupTag = "group"
+  val DefaultProduceQuotaProp = "default.produce.quota"
+  val DefaultFetchQuotaProp = "default.fetch.quota"
+  val UnlimitedQuotaMetricTags = new util.HashMap[String, String]
+  val quotaLimitCalls = Map(
+    ClientQuotaType.PRODUCE -> new AtomicInteger,
+    ClientQuotaType.FETCH -> new AtomicInteger,
+    ClientQuotaType.REQUEST -> new AtomicInteger
+  )
+  val callbackInstances = new AtomicInteger
+
+  def tearDown(): Unit = {
+    callbackInstances.set(0)
+    quotaLimitCalls.values.foreach(_.set(0))
+    UnlimitedQuotaMetricTags.clear()
+  }
+}
+
+/**
+ * Quota callback for a grouped user. Both user principals and topics of each group
+ * are prefixed with the group name followed by '_'. This callback defines quotas of different
+ * types at the group level. Group quotas are configured in ZooKeeper as user quotas with
+ * the entity name "${group}_". Default group quotas are configured in ZooKeeper as user quotas
+ * with the entity name "_".
+ *
+ * Default group quotas may also be configured using the configuration options
+ * "default.produce.quota" and "default.fetch.quota" which can be reconfigured dynamically
+ * without restarting the broker. This tests custom reconfigurable options for quota callbacks,
+ */
+class GroupedUserQuotaCallback extends ClientQuotaCallback with Reconfigurable with Logging {
+
+  var brokerId: Int = -1
+  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]()
+
+  override def configure(configs: util.Map[String, _]): Unit = {
+    brokerId = configs.get(KafkaConfig.BrokerIdProp).toString.toInt
+    callbackInstances.incrementAndGet
+  }
+
+  override def reconfigurableConfigs: util.Set[String] = {
+    Set(DefaultProduceQuotaProp, DefaultFetchQuotaProp).asJava
+  }
+
+  override def validateReconfiguration(configs: util.Map[String, _]): Unit = {
+    reconfigurableConfigs.forEach(configValue(configs, _))
+  }
+
+  override def reconfigure(configs: util.Map[String, _]): Unit = {
+    configValue(configs, DefaultProduceQuotaProp).foreach(value => quotas(ClientQuotaType.PRODUCE).put("", value.toDouble))
+    configValue(configs, DefaultFetchQuotaProp).foreach(value => quotas(ClientQuotaType.FETCH).put("", value.toDouble))
+    customQuotasUpdated.values.foreach(_.set(true))
+  }
+
+  private def configValue(configs: util.Map[String, _], key: String): Option[Long] = {
+    val value = configs.get(key)
+    if (value != null) Some(value.toString.toLong) else None
+  }
+
+  override def quotaMetricTags(quotaType: ClientQuotaType, principal: KafkaPrincipal, clientId: String): util.Map[String, String] = {
+    principal match {
+      case groupPrincipal: GroupedUserPrincipal =>
+        val userGroup = groupPrincipal.userGroup
+        val quotaLimit = quotaOrDefault(userGroup, quotaType)
+        if (quotaLimit != null)
+          Map(QuotaGroupTag -> userGroup).asJava
+        else
+          UnlimitedQuotaMetricTags
+      case _ =>
+        UnlimitedQuotaMetricTags
+    }
+  }
+
+  override def quotaLimit(quotaType: ClientQuotaType, metricTags: util.Map[String, String]): lang.Double = {
+    quotaLimitCalls(quotaType).incrementAndGet
+    val group = metricTags.get(QuotaGroupTag)
+    if (group != null) quotaOrDefault(group, quotaType) else null
+  }
+
+  override def updateClusterMetadata(cluster: Cluster): Boolean = {
+    val topicsByGroup = cluster.topics.asScala.groupBy(group)
+
+    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 }
+      val multiplier = if (totalPartitions == 0)
+        1
+      else if (partitionsOnThisBroker == 0)
+        1.0 / totalPartitions
+      else
+        partitionsOnThisBroker.toDouble / totalPartitions
+      partitionRatio.put(group, multiplier) != multiplier
+    }.exists(identity)
+  }
+
+  override def updateQuota(quotaType: ClientQuotaType, quotaEntity: ClientQuotaEntity, newValue: Double): Unit = {
+    quotas(quotaType).put(userGroup(quotaEntity), newValue)
+  }
+
+  override def removeQuota(quotaType: ClientQuotaType, quotaEntity: ClientQuotaEntity): Unit = {
+    quotas(quotaType).remove(userGroup(quotaEntity))
+  }
+
+  override def quotaResetRequired(quotaType: ClientQuotaType): Boolean = customQuotasUpdated(quotaType).getAndSet(false)
+
+  def close(): Unit = {}
+
+  private def userGroup(quotaEntity: ClientQuotaEntity): String = {
+    val configEntity = quotaEntity.configEntities.get(0)
+    if (configEntity.entityType == ClientQuotaEntity.ConfigEntityType.USER)
+      group(configEntity.name)
+    else
+      throw new IllegalArgumentException(s"Config entity type ${configEntity.entityType} is not supported")
+  }
+
+  private def quotaOrDefault(group: String, quotaType: ClientQuotaType): lang.Double = {
+    val quotaMap = quotas(quotaType)
+    var quotaLimit: Any = quotaMap.get(group)
+    if (quotaLimit == null)
+      quotaLimit = quotaMap.get("")
+    if (quotaLimit != null) scaledQuota(quotaType, group, quotaLimit.asInstanceOf[Double]) else null
+  }
+
+  private def scaledQuota(quotaType: ClientQuotaType, group: String, configuredQuota: Double): Double = {
+    if (quotaType == ClientQuotaType.REQUEST)
+      configuredQuota
+    else {
+      val multiplier = partitionRatio.get(group)
+      if (multiplier <= 0.0) configuredQuota else configuredQuota * multiplier
+    }
+  }
+}
+
+
diff --git a/core/src/test/scala/integration/kafka/api/DelegationTokenEndToEndAuthorizationTest.scala b/core/src/test/scala/integration/kafka/api/DelegationTokenEndToEndAuthorizationTest.scala
new file mode 100644
index 0000000000000..ad2aab7787519
--- /dev/null
+++ b/core/src/test/scala/integration/kafka/api/DelegationTokenEndToEndAuthorizationTest.scala
@@ -0,0 +1,126 @@
+/**
+  * Licensed to the Apache Software Foundation (ASF) under one or more
+  * contributor license agreements.  See the NOTICE file distributed with
+  * this work for additional information regarding copyright ownership.
+  * The ASF licenses this file to You under the Apache License, Version 2.0
+  * (the "License"); you may not use this file except in compliance with
+  * the License.  You may obtain a copy of the License at
+  *
+  * http://www.apache.org/licenses/LICENSE-2.0
+  *
+  * Unless required by applicable law or agreed to in writing, software
+  * distributed under the License is distributed on an "AS IS" BASIS,
+  * WITHOUT WARRANTIES 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.Properties
+
+import kafka.server.KafkaConfig
+import kafka.utils.{JaasTestUtils, TestUtils}
+import kafka.zk.ConfigEntityChangeNotificationZNode
+import org.apache.kafka.clients.admin.{Admin, AdminClientConfig, ScramCredentialInfo, UserScramCredentialAlteration, UserScramCredentialUpsertion, ScramMechanism => PublicScramMechanism}
+import org.apache.kafka.common.config.SaslConfigs
+import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
+import org.apache.kafka.common.security.scram.internals.ScramMechanism
+import org.apache.kafka.common.security.token.delegation.DelegationToken
+import org.junit.Assert._
+import org.junit.{Before, Test}
+
+import scala.jdk.CollectionConverters._
+
+class DelegationTokenEndToEndAuthorizationTest extends EndToEndAuthorizationTest {
+
+  val kafkaClientSaslMechanism = "SCRAM-SHA-256"
+  val kafkaServerSaslMechanisms = ScramMechanism.mechanismNames.asScala.toList
+
+  override protected def securityProtocol = SecurityProtocol.SASL_SSL
+
+  override protected val serverSaslProperties = Some(kafkaServerSaslProperties(kafkaServerSaslMechanisms, kafkaClientSaslMechanism))
+  override protected val clientSaslProperties = Some(kafkaClientSaslProperties(kafkaClientSaslMechanism))
+
+  override val clientPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, JaasTestUtils.KafkaScramUser)
+  private val clientPassword = JaasTestUtils.KafkaScramPassword
+
+  override val kafkaPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, JaasTestUtils.KafkaScramAdmin)
+  private val kafkaPassword = JaasTestUtils.KafkaScramAdminPassword
+
+  private val privilegedAdminClientConfig = new Properties()
+
+  this.serverConfig.setProperty(KafkaConfig.DelegationTokenSecretKeyProp, "testKey")
+
+  override def configureSecurityBeforeServersStart(): Unit = {
+    super.configureSecurityBeforeServersStart()
+    zkClient.makeSurePersistentPathExists(ConfigEntityChangeNotificationZNode.path)
+    // Create broker admin credentials before starting brokers
+    createScramCredentials(zkConnect, kafkaPrincipal.getName, kafkaPassword)
+  }
+
+  override def createPrivilegedAdminClient() = createScramAdminClient(kafkaClientSaslMechanism, kafkaPrincipal.getName, kafkaPassword)
+
+  override def configureSecurityAfterServersStart(): Unit = {
+    super.configureSecurityAfterServersStart()
+
+    // create scram credential for user "scram-user"
+    createScramCredentialsViaPrivilegedAdminClient(clientPrincipal.getName, clientPassword)
+    waitForUserScramCredentialToAppearOnAllBrokers(clientPrincipal.getName, kafkaClientSaslMechanism)
+
+    //create a token with "scram-user" credentials and a privileged token with scram-admin credentials
+    val tokens = createDelegationTokens()
+    val token = tokens._1
+    val privilegedToken = tokens._2
+
+    privilegedAdminClientConfig.putAll(adminClientConfig)
+
+    // pass token to client jaas config
+    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)
+    val privilegedClientLoginContext = JaasTestUtils.tokenClientLoginModule(privilegedToken.tokenInfo().tokenId(), privilegedToken.hmacAsBase64String())
+    privilegedAdminClientConfig.put(SaslConfigs.SASL_JAAS_CONFIG, privilegedClientLoginContext)
+  }
+
+  @Test
+  def testCreateUserWithDelegationToken(): Unit = {
+    val privilegedAdminClient = Admin.create(privilegedAdminClientConfig)
+    try {
+      val user = "user"
+      val results = privilegedAdminClient.alterUserScramCredentials(List[UserScramCredentialAlteration](
+        new UserScramCredentialUpsertion(user, new ScramCredentialInfo(PublicScramMechanism.SCRAM_SHA_256, 4096), "password")).asJava)
+      assertEquals(1, results.values.size)
+      val future = results.values.get(user)
+      future.get // make sure we haven't completed exceptionally
+    } finally {
+      privilegedAdminClient.close()
+    }
+  }
+
+  @Before
+  override def setUp(): Unit = {
+    startSasl(jaasSections(kafkaServerSaslMechanisms, Option(kafkaClientSaslMechanism), Both))
+    super.setUp()
+    privilegedAdminClientConfig.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList)
+  }
+
+  private def createDelegationTokens(): (DelegationToken, DelegationToken) = {
+    val adminClient = createScramAdminClient(kafkaClientSaslMechanism, clientPrincipal.getName, clientPassword)
+    try {
+      val privilegedAdminClient = createScramAdminClient(kafkaClientSaslMechanism, kafkaPrincipal.getName, kafkaPassword)
+      try {
+        val token = adminClient.createDelegationToken().delegationToken().get()
+        val privilegedToken = privilegedAdminClient.createDelegationToken().delegationToken().get()
+        //wait for tokens to reach all the brokers
+        TestUtils.waitUntilTrue(() => servers.forall(server => server.tokenCache.tokens().size() == 2),
+          "Timed out waiting for token to propagate to all servers")
+        (token, privilegedToken)
+      } finally {
+        privilegedAdminClient.close()
+      }
+    } finally {
+      adminClient.close()
+    }
+  }
+}
diff --git a/core/src/test/scala/integration/kafka/api/DescribeAuthorizedOperationsTest.scala b/core/src/test/scala/integration/kafka/api/DescribeAuthorizedOperationsTest.scala
new file mode 100644
index 0000000000000..98d9034757ea8
--- /dev/null
+++ b/core/src/test/scala/integration/kafka/api/DescribeAuthorizedOperationsTest.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.api
+
+import java.io.File
+import java.util
+import java.util.Properties
+
+import kafka.security.authorizer.{AclAuthorizer, AclEntry}
+import kafka.server.KafkaConfig
+import kafka.utils.{CoreUtils, JaasTestUtils, TestUtils}
+import org.apache.kafka.clients.admin._
+import org.apache.kafka.common.acl.AclOperation.{ALTER, CLUSTER_ACTION, DESCRIBE}
+import org.apache.kafka.common.acl.AclPermissionType.ALLOW
+import org.apache.kafka.common.acl._
+import org.apache.kafka.common.resource.{PatternType, Resource, ResourcePattern, ResourceType}
+import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
+import org.apache.kafka.common.utils.Utils
+import org.apache.kafka.server.authorizer.Authorizer
+import org.junit.Assert.{assertEquals, assertFalse, assertNull}
+import org.junit.{After, Before, Test}
+
+import scala.jdk.CollectionConverters._
+
+class DescribeAuthorizedOperationsTest extends IntegrationTestHarness with SaslSetup {
+  override val brokerCount = 1
+  this.serverConfig.setProperty(KafkaConfig.ZkEnableSecureAclsProp, "true")
+  this.serverConfig.setProperty(KafkaConfig.AuthorizerClassNameProp, classOf[AclAuthorizer].getName)
+
+  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(): Unit = {
+    val authorizer = CoreUtils.createObject[Authorizer](classOf[AclAuthorizer].getName)
+    val clusterResource = new ResourcePattern(ResourceType.CLUSTER, Resource.CLUSTER_NAME, PatternType.LITERAL)
+    val topicResource = new ResourcePattern(ResourceType.TOPIC, AclEntry.WildcardResource, PatternType.LITERAL)
+
+    try {
+      authorizer.configure(this.configs.head.originals())
+      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()
+    }
+  }
+
+  @Before
+  override def setUp(): Unit = {
+    startSasl(jaasSections(Seq("GSSAPI"), Some("GSSAPI"), Both, JaasTestUtils.KafkaServerContextName))
+    super.setUp()
+    TestUtils.waitUntilBrokerMetadataIsPropagated(servers)
+  }
+
+  private def accessControlEntry(userName: String, permissionType: AclPermissionType, operation: AclOperation): AccessControlEntry = {
+    new AccessControlEntry(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, userName).toString,
+      AclEntry.WildcardHost, operation, permissionType)
+  }
+
+  @After
+  override def tearDown(): Unit = {
+    if (client != null)
+      Utils.closeQuietly(client, "AdminClient")
+    super.tearDown()
+    closeSasl()
+  }
+
+  val group1Acl = new AclBinding(new ResourcePattern(ResourceType.GROUP, group1, PatternType.LITERAL),
+    new AccessControlEntry("User:" + JaasTestUtils.KafkaClientPrincipalUnqualifiedName2, "*", AclOperation.ALL, AclPermissionType.ALLOW))
+
+  val group2Acl = new AclBinding(new ResourcePattern(ResourceType.GROUP, group2, PatternType.LITERAL),
+    new AccessControlEntry("User:" + JaasTestUtils.KafkaClientPrincipalUnqualifiedName2, "*", AclOperation.DESCRIBE, AclPermissionType.ALLOW))
+
+  val group3Acl = new AclBinding(new ResourcePattern(ResourceType.GROUP, group3, PatternType.LITERAL),
+    new AccessControlEntry("User:" + JaasTestUtils.KafkaClientPrincipalUnqualifiedName2, "*", AclOperation.DELETE, AclPermissionType.ALLOW))
+
+  val clusterAllAcl = new AclBinding(new ResourcePattern(ResourceType.CLUSTER, Resource.CLUSTER_NAME, PatternType.LITERAL),
+    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)
+    adminClientConfig.put(AdminClientConfig.REQUEST_TIMEOUT_MS_CONFIG, "20000")
+    val securityProps: util.Map[Object, Object] =
+      TestUtils.adminClientSecurityConfigs(securityProtocol, trustStoreFile, clientSaslProperties)
+    securityProps.forEach { (key, value) => adminClientConfig.put(key.asInstanceOf[String], value) }
+    adminClientConfig
+  }
+
+  @Test
+  def testConsumerGroupAuthorizedOperations(): Unit = {
+    client = Admin.create(createConfig())
+
+    val results = client.createAcls(List(group1Acl, group2Acl, group3Acl).asJava)
+    assertEquals(Set(group1Acl, group2Acl, group3Acl), results.values.keySet.asScala)
+    results.all.get
+
+    val describeConsumerGroupsResult = client.describeConsumerGroups(Seq(group1, group2, group3).asJava,
+      new DescribeConsumerGroupsOptions().includeAuthorizedOperations(true))
+    assertEquals(3, describeConsumerGroupsResult.describedGroups().size())
+    val expectedOperations = AclEntry.supportedOperations(ResourceType.GROUP).asJava
+
+    val group1Description = describeConsumerGroupsResult.describedGroups().get(group1).get
+    assertEquals(expectedOperations, group1Description.authorizedOperations())
+
+    val group2Description = describeConsumerGroupsResult.describedGroups().get(group2).get
+    assertEquals(Set(AclOperation.DESCRIBE), group2Description.authorizedOperations().asScala.toSet)
+
+    val group3Description = describeConsumerGroupsResult.describedGroups().get(group3).get
+    assertEquals(Set(AclOperation.DESCRIBE, AclOperation.DELETE), group3Description.authorizedOperations().asScala.toSet)
+  }
+
+  @Test
+  def testClusterAuthorizedOperations(): Unit = {
+    client = Admin.create(createConfig())
+
+    // test without includeAuthorizedOperations flag
+    var clusterDescribeResult = client.describeCluster()
+    assertNull(clusterDescribeResult.authorizedOperations.get())
+
+    //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(clusterAllAcl).asJava)
+    assertEquals(Set(clusterAllAcl), results.values.keySet.asScala)
+    results.all.get
+
+    val expectedOperations = AclEntry.supportedOperations(ResourceType.CLUSTER).asJava
+
+    clusterDescribeResult = client.describeCluster(new DescribeClusterOptions().
+      includeAuthorizedOperations(true))
+    assertEquals(expectedOperations, clusterDescribeResult.authorizedOperations().get())
+  }
+
+  @Test
+  def testTopicAuthorizedOperations(): Unit = {
+    client = Admin.create(createConfig())
+    createTopic(topic1)
+    createTopic(topic2)
+
+    // test without includeAuthorizedOperations flag
+    var describeTopicsResult = client.describeTopics(Set(topic1, topic2).asJava).all.get()
+    assertNull(describeTopicsResult.get(topic1).authorizedOperations)
+    assertNull(describeTopicsResult.get(topic2).authorizedOperations)
+
+    //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 = AclEntry.supportedOperations(ResourceType.TOPIC).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 720d8b68733fc..f780f47b32f03 100644
--- a/core/src/test/scala/integration/kafka/api/EndToEndAuthorizationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/EndToEndAuthorizationTest.scala
@@ -5,7 +5,7 @@
  * The ASF licenses this file to You under the Apache License, Version 2.0
  * (the "License"); you may not use this file except in compliance with
  * the License.  You may obtain a copy of the License at
- * 
+ *
  *    http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing, software
@@ -17,24 +17,34 @@
 
 package kafka.api
 
+import com.yammer.metrics.core.Gauge
 import java.io.File
-import java.util.ArrayList
+import java.util.Collections
 import java.util.concurrent.ExecutionException
 
 import kafka.admin.AclCommand
-import kafka.common.TopicAndPartition
-import kafka.security.auth._
+import kafka.metrics.KafkaYammerMetrics
+import kafka.security.authorizer.AclAuthorizer
+import kafka.security.authorizer.AclEntry.WildcardHost
 import kafka.server._
 import kafka.utils._
-import org.apache.kafka.clients.consumer.{Consumer, ConsumerConfig, ConsumerRecord}
+import org.apache.kafka.clients.admin.Admin
+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, TimeoutException, TopicAuthorizationException}
+import org.apache.kafka.common.errors.{GroupAuthorizationException, TopicAuthorizationException}
+import org.apache.kafka.common.resource._
+import org.apache.kafka.common.resource.ResourceType._
+import org.apache.kafka.common.resource.PatternType.{LITERAL, PREFIXED}
+import org.apache.kafka.common.security.auth.KafkaPrincipal
 import org.junit.Assert._
 import org.junit.{After, Before, Test}
+import org.scalatest.Assertions.{assertThrows, fail, intercept}
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 
 /**
   * The test cases here verify that a producer authorized to publish to a topic
@@ -55,130 +65,155 @@ import scala.collection.JavaConverters._
   * would end up with ZooKeeperTestHarness twice.
   */
 abstract class EndToEndAuthorizationTest extends IntegrationTestHarness with SaslSetup {
-  override val producerCount = 1
-  override val consumerCount = 2
-  override val serverCount = 3
+  override val brokerCount = 3
 
-  override def configureSecurityBeforeServersStart() {
-    AclCommand.main(clusterAclArgs)
+  override def configureSecurityBeforeServersStart(): Unit = {
+    AclCommand.main(clusterActionArgs)
+    AclCommand.main(clusterAlterArgs)
+    AclCommand.main(topicBrokerReadAclArgs)
   }
 
   val numRecords = 1
-  val group = "group"
-  val topic = "e2etopic"
-  val topicWildcard = "*"
+  val groupPrefix = "gr"
+  val group = s"${groupPrefix}oup"
+  val topicPrefix = "e2e"
+  val topic = s"${topicPrefix}topic"
+  val wildcard = "*"
   val part = 0
   val tp = new TopicPartition(topic, part)
-  val topicAndPartition = TopicAndPartition(topic, part)
-  val clientPrincipal: String
-  val kafkaPrincipal: String
 
   override protected lazy val trustStoreFile = Some(File.createTempFile("truststore", ".jks"))
+  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 clientPrincipal: KafkaPrincipal
+  def kafkaPrincipal: KafkaPrincipal
 
-  val topicResource = new Resource(Topic, topic)
-  val groupResource = new Resource(Group, group)
-  val clusterResource = Resource.ClusterResource
-
-  // Arguments to AclCommand to set ACLs. There are three definitions here:
-  // 1- Provides read and write access to topic
-  // 2- Provides only write access to topic
-  // 3- Provides read access to consumer group
-  def clusterAclArgs: Array[String] = Array("--authorizer-properties",
-                                            s"zookeeper.connect=$zkConnect",
-                                            s"--add",
-                                            s"--cluster",
-                                            s"--operation=ClusterAction",
-                                            s"--allow-principal=$kafkaPrincipalType:$kafkaPrincipal")
+  // Arguments to AclCommand to set ACLs.
+  def clusterActionArgs: Array[String] = Array("--authorizer-properties",
+                                          s"zookeeper.connect=$zkConnect",
+                                          s"--add",
+                                          s"--cluster",
+                                          s"--operation=ClusterAction",
+                                          s"--allow-principal=$kafkaPrincipal")
+  // necessary to create SCRAM credentials via the admin client using the broker's credentials
+  // without this we would need to create the SCRAM credentials via ZooKeeper
+  def clusterAlterArgs: Array[String] = Array("--authorizer-properties",
+                                          s"zookeeper.connect=$zkConnect",
+                                          s"--add",
+                                          s"--cluster",
+                                          s"--operation=Alter",
+                                          s"--allow-principal=$kafkaPrincipal")
   def topicBrokerReadAclArgs: Array[String] = Array("--authorizer-properties",
-                                                    s"zookeeper.connect=$zkConnect",
-                                                    s"--add",
-                                                    s"--topic=$topicWildcard",
-                                                    s"--operation=Read",
-                                                    s"--allow-principal=$kafkaPrincipalType:$kafkaPrincipal")
-  def produceAclArgs: Array[String] = Array("--authorizer-properties",
+                                          s"zookeeper.connect=$zkConnect",
+                                          s"--add",
+                                          s"--topic=$wildcard",
+                                          s"--operation=Read",
+                                          s"--allow-principal=$kafkaPrincipal")
+  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=$clientPrincipal")
   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=$clientPrincipal")
   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=$clientPrincipal")
   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")
-  def consumeAclArgs: 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=$clientPrincipal")
+  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=$clientPrincipal")
   def groupAclArgs: Array[String] = Array("--authorizer-properties",
                                           s"zookeeper.connect=$zkConnect",
                                           s"--add",
                                           s"--group=$group",
                                           s"--operation=Read",
-                                          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))
+                                          s"--allow-principal=$clientPrincipal")
+  def produceConsumeWildcardAclArgs: Array[String] = Array("--authorizer-properties",
+                                          s"zookeeper.connect=$zkConnect",
+                                          s"--add",
+                                          s"--topic=$wildcard",
+                                          s"--group=$wildcard",
+                                          s"--consumer",
+                                          s"--producer",
+                                          s"--allow-principal=$clientPrincipal")
+  def produceConsumePrefixedAclsArgs: Array[String] = Array("--authorizer-properties",
+                                          s"zookeeper.connect=$zkConnect",
+                                          s"--add",
+                                          s"--topic=$topicPrefix",
+                                          s"--group=$groupPrefix",
+                                          s"--resource-pattern-type=prefixed",
+                                          s"--consumer",
+                                          s"--producer",
+                                          s"--allow-principal=$clientPrincipal")
+
+  def ClusterActionAndClusterAlterAcls = Set(new AccessControlEntry(kafkaPrincipal.toString, WildcardHost, CLUSTER_ACTION, ALLOW),
+    new AccessControlEntry(kafkaPrincipal.toString, WildcardHost, ALTER, ALLOW))
+  def TopicBrokerReadAcl = Set(new AccessControlEntry(kafkaPrincipal.toString, WildcardHost, READ, ALLOW))
+  def GroupReadAcl = Set(new AccessControlEntry(clientPrincipal.toString, WildcardHost, READ, ALLOW))
+  def TopicReadAcl = Set(new AccessControlEntry(clientPrincipal.toString, WildcardHost, READ, ALLOW))
+  def TopicWriteAcl = Set(new AccessControlEntry(clientPrincipal.toString, WildcardHost, WRITE, ALLOW))
+  def TopicDescribeAcl = Set(new AccessControlEntry(clientPrincipal.toString, WildcardHost, DESCRIBE, ALLOW))
+  def TopicCreateAcl = Set(new AccessControlEntry(clientPrincipal.toString, WildcardHost, CREATE, ALLOW))
   // 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")
   this.serverConfig.setProperty(KafkaConfig.MinInSyncReplicasProp, "3")
+  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()
-    AclCommand.main(topicBrokerReadAclArgs)
     servers.foreach { s =>
-      TestUtils.waitAndVerifyAcls(TopicBrokerReadAcl, s.apis.authorizer.get, new Resource(Topic, "*"))
+      TestUtils.waitAndVerifyAcls(ClusterActionAndClusterAlterAcls, 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
-    TestUtils.createTopic(zkUtils, topic, 1, 3, this.servers)
+    createTopic(topic, 1, 3)
   }
 
-  override def createNewProducer: KafkaProducer[Array[Byte], Array[Byte]] = {
-    TestUtils.createNewProducer(brokerList,
-                                maxBlockMs = 3000L,
-                                securityProtocol = this.securityProtocol,
-                                trustStoreFile = this.trustStoreFile,
-                                saslProperties = this.clientSaslProperties,
-                                props = Some(producerConfig))
-  }
-  
   /**
     * Closes MiniKDC last when tearing down.
     */
   @After
-  override def tearDown() {
-    consumers.foreach(_.wakeup())
+  override def tearDown(): Unit = {
     super.tearDown()
     closeSasl()
   }
@@ -188,52 +223,190 @@ abstract class EndToEndAuthorizationTest extends IntegrationTestHarness with Sas
     */
   @Test
   def testProduceConsumeViaAssign(): Unit = {
-    setAclsAndProduce()
-    consumers.head.assign(List(tp).asJava)
-    consumeRecords(this.consumers.head, numRecords)
+    setAclsAndProduce(tp)
+    val consumer = createConsumer()
+    consumer.assign(List(tp).asJava)
+    consumeRecords(consumer, numRecords)
+    confirmReauthenticationMetrics()
+  }
+
+  protected def confirmReauthenticationMetrics(): Unit = {
+    val expiredConnectionsKilledCountTotal = getGauge("ExpiredConnectionsKilledCount").value()
+    servers.foreach { s =>
+        val numExpiredKilled = TestUtils.totalMetricValue(s, "expired-connections-killed-count")
+        assertTrue("Should have been zero expired connections killed: " + numExpiredKilled + "(total=" + expiredConnectionsKilledCountTotal + ")", numExpiredKilled == 0)
+    }
+    assertEquals("Should have been zero expired connections killed total", 0, expiredConnectionsKilledCountTotal, 0.0)
+    servers.foreach { s =>
+      assertTrue("failed re-authentications not 0", TestUtils.totalMetricValue(s, "failed-reauthentication-total") == 0)
+    }
+  }
+
+  private def getGauge(metricName: String) = {
+    KafkaYammerMetrics.defaultRegistry.allMetrics.asScala
+           .filter { case (k, _) => k.getName == metricName }
+           .headOption
+           .getOrElse { fail( "Unable to find metric " + metricName ) }
+           ._2.asInstanceOf[Gauge[Double]]
   }
 
   @Test
   def testProduceConsumeViaSubscribe(): Unit = {
-    setAclsAndProduce()
-    consumers.head.subscribe(List(topic).asJava)
-    consumeRecords(this.consumers.head, numRecords)
+    setAclsAndProduce(tp)
+    val consumer = createConsumer()
+    consumer.subscribe(List(topic).asJava)
+    consumeRecords(consumer, numRecords)
+    confirmReauthenticationMetrics()
+  }
+
+  @Test
+  def testProduceConsumeWithWildcardAcls(): Unit = {
+    setWildcardResourceAcls()
+    val producer = createProducer()
+    sendRecords(producer, numRecords, tp)
+    val consumer = createConsumer()
+    consumer.subscribe(List(topic).asJava)
+    consumeRecords(consumer, numRecords)
+    confirmReauthenticationMetrics()
+  }
+
+  @Test
+  def testProduceConsumeWithPrefixedAcls(): Unit = {
+    setPrefixedResourceAcls()
+    val producer = createProducer()
+    sendRecords(producer, numRecords, tp)
+    val consumer = createConsumer()
+    consumer.subscribe(List(topic).asJava)
+    consumeRecords(consumer, numRecords)
+    confirmReauthenticationMetrics()
+  }
+
+  @Test
+  def testProduceConsumeTopicAutoCreateTopicCreateAcl(): Unit = {
+    // topic2 is not created on setup()
+    val tp2 = new TopicPartition("topic2", 0)
+    setAclsAndProduce(tp2)
+    val consumer = createConsumer()
+    consumer.assign(List(tp2).asJava)
+    consumeRecords(consumer, numRecords, topic = tp2.topic)
+    confirmReauthenticationMetrics()
+  }
+
+  private def setWildcardResourceAcls(): Unit = {
+    AclCommand.main(produceConsumeWildcardAclArgs)
+    servers.foreach { s =>
+      TestUtils.waitAndVerifyAcls(TopicReadAcl ++ TopicWriteAcl ++ TopicDescribeAcl ++ TopicCreateAcl ++ TopicBrokerReadAcl, s.dataPlaneRequestProcessor.authorizer.get, wildcardTopicResource)
+      TestUtils.waitAndVerifyAcls(GroupReadAcl, s.dataPlaneRequestProcessor.authorizer.get, wildcardGroupResource)
+    }
   }
 
-  protected def setAclsAndProduce() {
-    AclCommand.main(produceAclArgs)
-    AclCommand.main(consumeAclArgs)
+  private def setPrefixedResourceAcls(): Unit = {
+    AclCommand.main(produceConsumePrefixedAclsArgs)
     servers.foreach { s =>
-      TestUtils.waitAndVerifyAcls(TopicReadAcl ++ TopicWriteAcl ++ TopicDescribeAcl, s.apis.authorizer.get, topicResource)
-      TestUtils.waitAndVerifyAcls(GroupReadAcl, s.apis.authorizer.get, groupResource)
+      TestUtils.waitAndVerifyAcls(TopicReadAcl ++ TopicWriteAcl ++ TopicDescribeAcl ++ TopicCreateAcl, s.dataPlaneRequestProcessor.authorizer.get, prefixedTopicResource)
+      TestUtils.waitAndVerifyAcls(GroupReadAcl, s.dataPlaneRequestProcessor.authorizer.get, prefixedGroupResource)
+    }
+  }
+
+  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 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)
     }
-    sendRecords(numRecords, tp)
   }
 
   /**
-    * 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 = {
-    sendRecords(numRecords, tp)
+  @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 the unauthorized topic and throw; since we can now return data during the time we are updating
+    // metadata / fetching positions, it is possible that the authorized topic record is returned during this time.
+    consumer.assign(List(tp, tp2).asJava)
+    sendRecords(producer, numRecords, tp2)
+    var topic2RecordConsumed = false
+    def verifyNoRecords(records: ConsumerRecords[Array[Byte], Array[Byte]]): Boolean = {
+      assertEquals("Consumed records with unexpected partitions: " + records, Collections.singleton(tp2), records.partitions())
+      topic2RecordConsumed = true
+      false
+    }
+    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)
+    if (!topic2RecordConsumed) {
+      consumeRecordsIgnoreOneAuthorizationException(consumer, numRecords, startingOffset = 1, topic2)
+    }
+    sendRecords(producer, numRecords, tp)
+    consumeRecordsIgnoreOneAuthorizationException(consumer, numRecords, startingOffset = 0, 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
   def testNoProduceWithDescribeAcl(): Unit = {
     AclCommand.main(describeAclArgs)
     servers.foreach { s =>
-      TestUtils.waitAndVerifyAcls(TopicDescribeAcl, s.apis.authorizer.get, topicResource)
+      TestUtils.waitAndVerifyAcls(TopicDescribeAcl, s.dataPlaneRequestProcessor.authorizer.get, topicResource)
     }
     try{
-      sendRecords(numRecords, tp)
+      val producer = createProducer()
+      sendRecords(producer, numRecords, tp)
       fail("exception expected")
     } catch {
       case e: TopicAuthorizationException =>
         assertEquals(Set(topic).asJava, e.unauthorizedTopics())
     }
+    confirmReauthenticationMetrics()
   }
-  
+
    /**
     * Tests that a consumer fails to consume messages without the appropriate
     * ACL set.
@@ -241,72 +414,90 @@ abstract class EndToEndAuthorizationTest extends IntegrationTestHarness with Sas
   @Test(expected = classOf[KafkaException])
   def testNoConsumeWithoutDescribeAclViaAssign(): Unit = {
     noConsumeWithoutDescribeAclSetup()
-    consumers.head.assign(List(tp).asJava)
+    val consumer = createConsumer()
+    consumer.assign(List(tp).asJava)
     // the exception is expected when the consumer attempts to lookup offsets
-    consumeRecords(this.consumers.head)
+    consumeRecords(consumer)
+    confirmReauthenticationMetrics()
   }
-  
-  @Test(expected = classOf[TopicAuthorizationException])
+
+  @Test
   def testNoConsumeWithoutDescribeAclViaSubscribe(): Unit = {
     noConsumeWithoutDescribeAclSetup()
-    consumers.head.subscribe(List(topic).asJava)
+    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(this.consumers.head, 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 = {
-    AclCommand.main(produceAclArgs)
+    AclCommand.main(produceAclArgs(tp.topic))
     AclCommand.main(groupAclArgs)
     servers.foreach { s =>
-      TestUtils.waitAndVerifyAcls(TopicWriteAcl ++ TopicDescribeAcl, s.apis.authorizer.get, topicResource)
-      TestUtils.waitAndVerifyAcls(GroupReadAcl, s.apis.authorizer.get, groupResource)
+      TestUtils.waitAndVerifyAcls(TopicWriteAcl ++ TopicDescribeAcl ++ TopicCreateAcl, s.dataPlaneRequestProcessor.authorizer.get, topicResource)
+      TestUtils.waitAndVerifyAcls(GroupReadAcl, s.dataPlaneRequestProcessor.authorizer.get, groupResource)
     }
 
-    sendRecords(numRecords, tp)
+    val producer = createProducer()
+    sendRecords(producer, numRecords, tp)
 
     AclCommand.main(deleteDescribeAclArgs)
     AclCommand.main(deleteWriteAclArgs)
     servers.foreach { s =>
-      TestUtils.waitAndVerifyAcls(GroupReadAcl, s.apis.authorizer.get, groupResource)
+      TestUtils.waitAndVerifyAcls(GroupReadAcl, s.dataPlaneRequestProcessor.authorizer.get, groupResource)
     }
   }
-  
+
   @Test
   def testNoConsumeWithDescribeAclViaAssign(): Unit = {
     noConsumeWithDescribeAclSetup()
-    consumers.head.assign(List(tp).asJava)
+    val consumer = createConsumer()
+    consumer.assign(List(tp).asJava)
 
     try {
-      consumeRecords(this.consumers.head)
+      consumeRecords(consumer)
       fail("Topic authorization exception expected")
     } catch {
       case e: TopicAuthorizationException =>
         assertEquals(Set(topic).asJava, e.unauthorizedTopics())
     }
+    confirmReauthenticationMetrics()
   }
-  
+
   @Test
   def testNoConsumeWithDescribeAclViaSubscribe(): Unit = {
     noConsumeWithDescribeAclSetup()
-    consumers.head.subscribe(List(topic).asJava)
+    val consumer = createConsumer()
+    consumer.subscribe(List(topic).asJava)
 
     try {
-      consumeRecords(this.consumers.head)
+      consumeRecords(consumer)
       fail("Topic authorization exception expected")
     } catch {
       case e: TopicAuthorizationException =>
         assertEquals(Set(topic).asJava, e.unauthorizedTopics())
     }
+    confirmReauthenticationMetrics()
   }
-  
+
   private def noConsumeWithDescribeAclSetup(): Unit = {
-    AclCommand.main(produceAclArgs)
+    AclCommand.main(produceAclArgs(tp.topic))
     AclCommand.main(groupAclArgs)
     servers.foreach { s =>
-      TestUtils.waitAndVerifyAcls(TopicWriteAcl ++ TopicDescribeAcl, s.apis.authorizer.get, topicResource)
-      TestUtils.waitAndVerifyAcls(GroupReadAcl, s.apis.authorizer.get, groupResource)
+      TestUtils.waitAndVerifyAcls(TopicWriteAcl ++ TopicDescribeAcl ++ TopicCreateAcl, s.dataPlaneRequestProcessor.authorizer.get, topicResource)
+      TestUtils.waitAndVerifyAcls(GroupReadAcl, s.dataPlaneRequestProcessor.authorizer.get, groupResource)
     }
-    sendRecords(numRecords, tp)
+    val producer = createProducer()
+    sendRecords(producer, numRecords, tp)
   }
 
   /**
@@ -315,26 +506,31 @@ abstract class EndToEndAuthorizationTest extends IntegrationTestHarness with Sas
     */
   @Test
   def testNoGroupAcl(): Unit = {
-    AclCommand.main(produceAclArgs)
+    AclCommand.main(produceAclArgs(tp.topic))
     servers.foreach { s =>
-      TestUtils.waitAndVerifyAcls(TopicWriteAcl ++ TopicDescribeAcl, s.apis.authorizer.get, topicResource)
+      TestUtils.waitAndVerifyAcls(TopicWriteAcl ++ TopicDescribeAcl ++ TopicCreateAcl, s.dataPlaneRequestProcessor.authorizer.get, topicResource)
     }
-    sendRecords(numRecords, tp)
-    consumers.head.assign(List(tp).asJava)
+    val producer = createProducer()
+    sendRecords(producer, numRecords, tp)
+
+    val consumer = createConsumer()
+    consumer.assign(List(tp).asJava)
     try {
-      consumeRecords(this.consumers.head)
+      consumeRecords(consumer)
       fail("Topic authorization exception expected")
     } catch {
       case e: GroupAuthorizationException =>
         assertEquals(group, e.groupId())
     }
+    confirmReauthenticationMetrics()
   }
 
-  private def sendRecords(numRecords: Int, tp: TopicPartition) {
+  protected final 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")
-      this.producers.head.send(record)
+      producer.send(record)
     }
     try {
       futures.foreach(_.get)
@@ -343,28 +539,37 @@ abstract class EndToEndAuthorizationTest extends IntegrationTestHarness with Sas
     }
   }
 
-  protected def consumeRecords(consumer: Consumer[Array[Byte], Array[Byte]],
-                             numRecords: Int = 1,
-                             startingOffset: Int = 0,
-                             topic: String = topic,
-                             part: Int = part,
-                             timeout: Long = 10000) {
-    val records = new ArrayList[ConsumerRecord[Array[Byte], Array[Byte]]]()
-
-    val deadlineMs = System.currentTimeMillis() + timeout
-    while (records.size < numRecords && System.currentTimeMillis() < deadlineMs) {
-      for (record <- consumer.poll(50).asScala)
-        records.add(record)
-    }
-    if (records.size < numRecords)
-      throw new TimeoutException
+  protected final def consumeRecords(consumer: Consumer[Array[Byte], Array[Byte]],
+                                     numRecords: Int = 1,
+                                     startingOffset: Int = 0,
+                                     topic: String = topic,
+                                     part: Int = part,
+                                     timeout: Long = 10000): Unit = {
+    val records = TestUtils.consumeRecords(consumer, numRecords, timeout)
 
     for (i <- 0 until numRecords) {
-      val record = records.get(i)
+      val record = records(i)
       val offset = startingOffset + i
-      assertEquals(topic, record.topic())
-      assertEquals(part, record.partition())
-      assertEquals(offset.toLong, record.offset())
+      assertEquals(topic, record.topic)
+      assertEquals(part, record.partition)
+      assertEquals(offset.toLong, record.offset)
+    }
+  }
+
+  protected def createScramAdminClient(scramMechanism: String, user: String, password: String): Admin = {
+    createAdminClient(brokerList, securityProtocol, trustStoreFile, clientSaslProperties,
+      scramMechanism, user, password)
+  }
+
+  // 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 fce75b0033d0b..89cbf338a36ca 100644
--- a/core/src/test/scala/integration/kafka/api/EndToEndClusterIdTest.scala
+++ b/core/src/test/scala/integration/kafka/api/EndToEndClusterIdTest.scala
@@ -21,7 +21,6 @@ import java.util.concurrent.ExecutionException
 import java.util.concurrent.atomic.AtomicReference
 import java.util.Properties
 
-import kafka.common.TopicAndPartition
 import kafka.integration.KafkaServerTestHarness
 import kafka.server._
 import kafka.utils._
@@ -33,11 +32,9 @@ import org.apache.kafka.test.{TestUtils => _, _}
 import org.junit.Assert._
 import org.junit.{Before, Test}
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 import org.apache.kafka.test.TestUtils.isValidClusterId
 
-import scala.collection.mutable.ArrayBuffer
-
 /** The test cases here verify the following conditions.
   * 1. The ProducerInterceptor receives the cluster id after the onSend() method is called and before onAcknowledgement() method is called.
   * 2. The Serializer receives the cluster id before the serialize() method is called.
@@ -58,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)
     }
   }
@@ -69,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)
     }
   }
@@ -80,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)
     }
   }
@@ -100,7 +97,6 @@ class EndToEndClusterIdTest extends KafkaServerTestHarness {
   val topic = "e2etopic"
   val part = 0
   val tp = new TopicPartition(topic, part)
-  val topicAndPartition = TopicAndPartition(topic, part)
   this.serverConfig.setProperty(KafkaConfig.MetricReporterClassesProp, classOf[MockBrokerMetricsReporter].getName)
 
   override def generateConfigs = {
@@ -111,15 +107,15 @@ class EndToEndClusterIdTest extends KafkaServerTestHarness {
   }
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
     MockDeserializer.resetStaticVariables
     // create the consumer offset topic
-    TestUtils.createTopic(this.zkUtils, topic, 2, serverCount, this.servers)
+    createTopic(topic, 2, serverCount)
   }
 
   @Test
-  def testEndToEnd() {
+  def testEndToEnd(): Unit = {
     val appendStr = "mock"
     MockConsumerInterceptor.resetCounters()
     MockProducerInterceptor.resetCounters()
@@ -187,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")
@@ -204,18 +200,9 @@ class EndToEndClusterIdTest extends KafkaServerTestHarness {
                              numRecords: Int,
                              startingOffset: Int = 0,
                              topic: String = topic,
-                             part: Int = part) {
-    val records = new ArrayBuffer[ConsumerRecord[Array[Byte], Array[Byte]]]()
-    val maxIters = numRecords * 50
-    var iters = 0
-    while (records.size < numRecords) {
-      for (record <- consumer.poll(50).asScala) {
-        records += record
-      }
-      if (iters > maxIters)
-        throw new IllegalStateException("Failed to consume the expected records after " + iters + " iterations.")
-      iters += 1
-    }
+                             part: Int = part): Unit = {
+    val records = TestUtils.consumeRecords(consumer, numRecords)
+
     for (i <- 0 until numRecords) {
       val record = records(i)
       val offset = startingOffset + i
diff --git a/core/src/test/scala/integration/kafka/api/GroupAuthorizerIntegrationTest.scala b/core/src/test/scala/integration/kafka/api/GroupAuthorizerIntegrationTest.scala
new file mode 100644
index 0000000000000..8f52102cce67b
--- /dev/null
+++ b/core/src/test/scala/integration/kafka/api/GroupAuthorizerIntegrationTest.scala
@@ -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 kafka.api
+
+import java.util.Properties
+import java.util.concurrent.ExecutionException
+
+import kafka.api.GroupAuthorizerIntegrationTest._
+import kafka.security.auth.SimpleAclAuthorizer
+import kafka.security.authorizer.AclEntry.WildcardHost
+import kafka.server.{BaseRequestTest, KafkaConfig}
+import kafka.utils.TestUtils
+import org.apache.kafka.clients.consumer.ConsumerConfig
+import org.apache.kafka.clients.producer.ProducerRecord
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.acl.{AccessControlEntry, AclOperation, AclPermissionType}
+import org.apache.kafka.common.config.internals.BrokerSecurityConfigs
+import org.apache.kafka.common.errors.TopicAuthorizationException
+import org.apache.kafka.common.network.ListenerName
+import org.apache.kafka.common.resource.{PatternType, Resource, ResourcePattern, ResourceType}
+import org.apache.kafka.common.security.auth.{AuthenticationContext, KafkaPrincipal, KafkaPrincipalBuilder}
+import org.junit.Assert._
+import org.junit.{Before, Test}
+import org.scalatest.Assertions.intercept
+
+import scala.jdk.CollectionConverters._
+
+object GroupAuthorizerIntegrationTest {
+  val BrokerPrincipal = new KafkaPrincipal("Group", "broker")
+  val ClientPrincipal = new KafkaPrincipal("Group", "client")
+
+  val BrokerListenerName = "BROKER"
+  val ClientListenerName = "CLIENT"
+
+  class GroupPrincipalBuilder extends KafkaPrincipalBuilder {
+    override def build(context: AuthenticationContext): KafkaPrincipal = {
+      context.listenerName match {
+        case BrokerListenerName => BrokerPrincipal
+        case ClientListenerName => ClientPrincipal
+        case listenerName => throw new IllegalArgumentException(s"No principal mapped to listener $listenerName")
+      }
+    }
+  }
+}
+
+class GroupAuthorizerIntegrationTest extends BaseRequestTest {
+
+  val brokerId: Integer = 0
+
+  override def brokerCount: Int = 1
+  override def interBrokerListenerName: ListenerName = new ListenerName(BrokerListenerName)
+  override def listenerName: ListenerName = new ListenerName(ClientListenerName)
+
+  def brokerPrincipal: KafkaPrincipal = BrokerPrincipal
+  def clientPrincipal: KafkaPrincipal = ClientPrincipal
+
+  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")
+    properties.put(BrokerSecurityConfigs.PRINCIPAL_BUILDER_CLASS_CONFIG, classOf[GroupPrincipalBuilder].getName)
+  }
+
+  @Before
+  override def setUp(): Unit = {
+    doSetup(createOffsetsTopic = false)
+
+    // Allow inter-broker communication
+    TestUtils.addAndVerifyAcls(servers.head,
+      Set(createAcl(AclOperation.CLUSTER_ACTION, AclPermissionType.ALLOW, principal = BrokerPrincipal)),
+      new ResourcePattern(ResourceType.CLUSTER, Resource.CLUSTER_NAME, PatternType.LITERAL))
+
+    TestUtils.createOffsetsTopic(zkClient, servers)
+  }
+
+  private def createAcl(aclOperation: AclOperation,
+                        aclPermissionType: AclPermissionType,
+                        principal: KafkaPrincipal = ClientPrincipal): AccessControlEntry = {
+    new AccessControlEntry(principal.toString, WildcardHost, aclOperation, aclPermissionType)
+  }
+
+  @Test
+  def testUnauthorizedProduceAndConsume(): Unit = {
+    val topic = "topic"
+    val topicPartition = new TopicPartition("topic", 0)
+
+    createTopic(topic)
+
+    val producer = createProducer()
+    val produceException = intercept[ExecutionException] {
+      producer.send(new ProducerRecord[Array[Byte], Array[Byte]](topic, "message".getBytes)).get()
+    }.getCause
+    assertTrue(produceException.isInstanceOf[TopicAuthorizationException])
+    assertEquals(Set(topic), produceException.asInstanceOf[TopicAuthorizationException].unauthorizedTopics.asScala)
+
+    val consumer = createConsumer(configsToRemove = List(ConsumerConfig.GROUP_ID_CONFIG))
+    consumer.assign(List(topicPartition).asJava)
+    val consumeException = intercept[TopicAuthorizationException] {
+      TestUtils.pollUntilAtLeastNumRecords(consumer, numRecords = 1)
+    }
+    assertEquals(Set(topic), consumeException.unauthorizedTopics.asScala)
+  }
+
+  @Test
+  def testAuthorizedProduceAndConsume(): Unit = {
+    val topic = "topic"
+    val topicPartition = new TopicPartition("topic", 0)
+
+    createTopic(topic)
+
+    TestUtils.addAndVerifyAcls(servers.head,
+      Set(createAcl(AclOperation.WRITE, AclPermissionType.ALLOW)),
+      new ResourcePattern(ResourceType.TOPIC, topic, PatternType.LITERAL))
+    val producer = createProducer()
+    producer.send(new ProducerRecord[Array[Byte], Array[Byte]](topic, "message".getBytes)).get()
+
+    TestUtils.addAndVerifyAcls(servers.head,
+      Set(createAcl(AclOperation.READ, AclPermissionType.ALLOW)),
+      new ResourcePattern(ResourceType.TOPIC, topic, PatternType.LITERAL))
+    val consumer = createConsumer(configsToRemove = List(ConsumerConfig.GROUP_ID_CONFIG))
+    consumer.assign(List(topicPartition).asJava)
+    TestUtils.pollUntilAtLeastNumRecords(consumer, numRecords = 1)
+  }
+
+}
diff --git a/core/src/test/scala/integration/kafka/api/GroupCoordinatorIntegrationTest.scala b/core/src/test/scala/integration/kafka/api/GroupCoordinatorIntegrationTest.scala
index 2049e0a6a8ff2..af846283d8158 100644
--- a/core/src/test/scala/integration/kafka/api/GroupCoordinatorIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/GroupCoordinatorIntegrationTest.scala
@@ -21,12 +21,11 @@ import org.apache.kafka.common.TopicPartition
 import org.junit.Test
 import org.junit.Assert._
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 import java.util.Properties
 
 import org.apache.kafka.common.internals.Topic
 import org.apache.kafka.common.record.CompressionType
-import org.apache.kafka.common.security.auth.SecurityProtocol
 
 class GroupCoordinatorIntegrationTest extends KafkaServerTestHarness {
   val offsetsTopicCompressionCodec = CompressionType.GZIP
@@ -39,15 +38,13 @@ class GroupCoordinatorIntegrationTest extends KafkaServerTestHarness {
   }
 
   @Test
-  def testGroupCoordinatorPropagatesOfffsetsTopicCompressionCodec() {
-    val consumer = TestUtils.createNewConsumer(TestUtils.getBrokerListStrFromServers(servers),
-                                               securityProtocol = SecurityProtocol.PLAINTEXT)
+  def testGroupCoordinatorPropagatesOffsetsTopicCompressionCodec(): Unit = {
+    val consumer = TestUtils.createConsumer(TestUtils.getBrokerListStrFromServers(servers))
     val offsetMap = Map(
       new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, 0) -> new OffsetAndMetadata(10, "")
     ).asJava
     consumer.commitSync(offsetMap)
     val logManager = servers.head.getLogManager
-
     def getGroupMetadataLogOpt: Option[Log] =
       logManager.getLog(new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, 0))
 
diff --git a/core/src/test/scala/integration/kafka/api/GroupEndToEndAuthorizationTest.scala b/core/src/test/scala/integration/kafka/api/GroupEndToEndAuthorizationTest.scala
new file mode 100644
index 0000000000000..05243f0a7ab72
--- /dev/null
+++ b/core/src/test/scala/integration/kafka/api/GroupEndToEndAuthorizationTest.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.api
+
+import kafka.api.GroupEndToEndAuthorizationTest._
+import kafka.utils.JaasTestUtils
+import org.apache.kafka.common.config.internals.BrokerSecurityConfigs
+import org.apache.kafka.common.security.auth.{AuthenticationContext, KafkaPrincipal, KafkaPrincipalBuilder, SaslAuthenticationContext}
+
+object GroupEndToEndAuthorizationTest {
+  val GroupPrincipalType = "Group"
+  val ClientGroup = "testGroup"
+  class GroupPrincipalBuilder extends KafkaPrincipalBuilder {
+    override def build(context: AuthenticationContext): KafkaPrincipal = {
+      context match {
+        case ctx: SaslAuthenticationContext =>
+          if (ctx.server.getAuthorizationID == JaasTestUtils.KafkaScramUser)
+            new KafkaPrincipal(GroupPrincipalType, ClientGroup)
+          else
+            new KafkaPrincipal(GroupPrincipalType, ctx.server.getAuthorizationID)
+        case _ =>
+          KafkaPrincipal.ANONYMOUS
+      }
+    }
+  }
+}
+
+class GroupEndToEndAuthorizationTest extends SaslScramSslEndToEndAuthorizationTest {
+  override val clientPrincipal = new KafkaPrincipal(GroupPrincipalType, ClientGroup)
+  override val kafkaPrincipal = new KafkaPrincipal(GroupPrincipalType, JaasTestUtils.KafkaScramAdmin)
+  this.serverConfig.setProperty(BrokerSecurityConfigs.PRINCIPAL_BUILDER_CLASS_CONFIG, classOf[GroupPrincipalBuilder].getName)
+}
diff --git a/core/src/test/scala/integration/kafka/api/IntegrationTestHarness.scala b/core/src/test/scala/integration/kafka/api/IntegrationTestHarness.scala
index 4b272398ec4d8..49f2897c6dbea 100644
--- a/core/src/test/scala/integration/kafka/api/IntegrationTestHarness.scala
+++ b/core/src/test/scala/integration/kafka/api/IntegrationTestHarness.scala
@@ -17,42 +17,56 @@
 
 package kafka.api
 
-import org.apache.kafka.clients.producer.ProducerConfig
+import java.time.Duration
+
 import org.apache.kafka.clients.consumer.{ConsumerConfig, KafkaConsumer}
 import kafka.utils.TestUtils
 import kafka.utils.Implicits._
 import java.util.Properties
 
-import org.apache.kafka.clients.producer.KafkaProducer
+import org.apache.kafka.clients.producer.{KafkaProducer, ProducerConfig}
 import kafka.server.KafkaConfig
 import kafka.integration.KafkaServerTestHarness
-import org.apache.kafka.common.network.ListenerName
+import org.apache.kafka.clients.admin.{Admin, 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.Buffer
+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 brokerCount: Int
+  protected def logDirCount: Int = 1
 
-  val producerCount: Int
-  val consumerCount: Int
-  val serverCount: Int
-  var logDirCount: Int = 1
-  lazy val producerConfig = new Properties
-  lazy val consumerConfig = new Properties
-  lazy val serverConfig = new Properties
+  val producerConfig = new Properties
+  val consumerConfig = new Properties
+  val adminClientConfig = new Properties
+  val serverConfig = new Properties
 
-  val consumers = Buffer[KafkaConsumer[Array[Byte], Array[Byte]]]()
-  val producers = Buffer[KafkaProducer[Array[Byte], Array[Byte]]]()
+  private val consumers = mutable.Buffer[KafkaConsumer[_, _]]()
+  private val producers = mutable.Buffer[KafkaProducer[_, _]]()
+  private val adminClients = mutable.Buffer[Admin]()
 
   protected def interBrokerListenerName: ListenerName = listenerName
 
-  override def generateConfigs = {
-    val cfgs = TestUtils.createBrokerConfigs(serverCount, zkConnect, interBrokerSecurityProtocol = Some(securityProtocol),
+  protected def modifyConfigs(props: Seq[Properties]): Unit = {
+    props.foreach(_ ++= serverConfig)
+  }
+
+  override def generateConfigs: Seq[KafkaConfig] = {
+    val cfgs = TestUtils.createBrokerConfigs(brokerCount, zkConnect, interBrokerSecurityProtocol = Some(securityProtocol),
       trustStoreFile = trustStoreFile, saslProperties = serverSaslProperties, logDirCount = logDirCount)
-    cfgs.foreach { config =>
+    configureListeners(cfgs)
+    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)
 
@@ -63,50 +77,87 @@ 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() {
-    val producerSecurityProps = TestUtils.producerSecurityConfigs(securityProtocol, trustStoreFile, clientSaslProperties)
-    val consumerSecurityProps = TestUtils.consumerSecurityConfigs(securityProtocol, trustStoreFile, clientSaslProperties)
+  override def setUp(): Unit = {
+    doSetup(createOffsetsTopic = true)
+  }
+
+  def doSetup(createOffsetsTopic: Boolean): Unit = {
+    // Generate client security properties before starting the brokers in case certs are needed
+    producerConfig ++= clientSecurityProps("producer")
+    consumerConfig ++= clientSecurityProps("consumer")
+    adminClientConfig ++= clientSecurityProps("adminClient")
+
     super.setUp()
-    producerConfig.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, classOf[org.apache.kafka.common.serialization.ByteArraySerializer])
-    producerConfig.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, classOf[org.apache.kafka.common.serialization.ByteArraySerializer])
-    producerConfig ++= producerSecurityProps
-    consumerConfig.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, classOf[org.apache.kafka.common.serialization.ByteArrayDeserializer])
-    consumerConfig.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, classOf[org.apache.kafka.common.serialization.ByteArrayDeserializer])
-    consumerConfig ++= consumerSecurityProps
-    for (_ <- 0 until producerCount)
-      producers += createNewProducer
-    for (_ <- 0 until consumerCount) {
-      consumers += createNewConsumer
-    }
 
-    TestUtils.createOffsetsTopic(zkUtils, servers)
+    producerConfig.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList)
+    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)
+
+    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)
+
+    adminClientConfig.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList)
+
+    if (createOffsetsTopic)
+      TestUtils.createOffsetsTopic(zkClient, servers)
+  }
+
+  def clientSecurityProps(certAlias: String): Properties = {
+    TestUtils.securityConfigs(Mode.CLIENT, securityProtocol, trustStoreFile, certAlias, TestUtils.SslCertificateCn,
+      clientSaslProperties)
+  }
+
+  def createProducer[K, V](keySerializer: Serializer[K] = new ByteArraySerializer,
+                           valueSerializer: Serializer[V] = new ByteArraySerializer,
+                           configOverrides: Properties = new Properties): KafkaProducer[K, V] = {
+    val props = new Properties
+    props ++= producerConfig
+    props ++= configOverrides
+    val producer = new KafkaProducer[K, V](props, keySerializer, valueSerializer)
+    producers += producer
+    producer
   }
 
-  def createNewProducer: KafkaProducer[Array[Byte], Array[Byte]] = {
-      TestUtils.createNewProducer(brokerList,
-                                  securityProtocol = this.securityProtocol,
-                                  trustStoreFile = this.trustStoreFile,
-                                  saslProperties = this.clientSaslProperties,
-                                  props = Some(producerConfig))
+  def createConsumer[K, V](keyDeserializer: Deserializer[K] = new ByteArrayDeserializer,
+                           valueDeserializer: Deserializer[V] = new ByteArrayDeserializer,
+                           configOverrides: Properties = new Properties,
+                           configsToRemove: List[String] = List()): KafkaConsumer[K, V] = {
+    val props = new Properties
+    props ++= consumerConfig
+    props ++= configOverrides
+    configsToRemove.foreach(props.remove(_))
+    val consumer = new KafkaConsumer[K, V](props, keyDeserializer, valueDeserializer)
+    consumers += consumer
+    consumer
   }
 
-  def createNewConsumer: KafkaConsumer[Array[Byte], Array[Byte]] = {
-      TestUtils.createNewConsumer(brokerList,
-                                  securityProtocol = this.securityProtocol,
-                                  trustStoreFile = this.trustStoreFile,
-                                  saslProperties = this.clientSaslProperties,
-                                  props = Some(consumerConfig))
+  def createAdminClient(configOverrides: Properties = new Properties): Admin = {
+    val props = new Properties
+    props ++= adminClientConfig
+    props ++= configOverrides
+    val adminClient = Admin.create(props)
+    adminClients += adminClient
+    adminClient
   }
 
   @After
-  override def tearDown() {
-    producers.foreach(_.close())
-    consumers.foreach(_.close())
+  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 2839137eb112e..0000000000000
--- a/core/src/test/scala/integration/kafka/api/LegacyAdminClientTest.scala
+++ /dev/null
@@ -1,273 +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 java.util.concurrent.TimeUnit
-
-import kafka.admin.AdminClient
-import kafka.admin.AdminClient.DeleteRecordsResult
-import kafka.server.KafkaConfig
-import java.lang.{Long => JLong}
-import kafka.utils.{Logging, TestUtils}
-import org.apache.kafka.clients.consumer.{KafkaConsumer, ConsumerConfig}
-import org.apache.kafka.clients.producer.{KafkaProducer, ProducerRecord, ProducerConfig}
-import org.apache.kafka.common.TopicPartition
-import org.apache.kafka.common.protocol.{Errors, ApiKeys}
-import org.apache.kafka.common.requests.DeleteRecordsRequest
-import org.junit.{After, Before, Test}
-import org.junit.Assert._
-import scala.collection.JavaConverters._
-
-/**
-  * Tests for the deprecated Scala AdminClient.
-  */
-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)
-    TestUtils.createTopic(this.zkUtils, topic, 2, serverCount, this.servers)
-  }
-
-  @After
-  override def tearDown() {
-    client.close()
-    super.tearDown()
-  }
-
-  @Test
-  def testSeekToBeginningAfterDeleteRecords() {
-    val consumer = consumers.head
-    subscribeAndWaitForAssignment(topic, consumer)
-
-    sendRecords(producers.head, 10, tp)
-    consumer.seekToBeginning(Collections.singletonList(tp))
-    assertEquals(0L, consumer.position(tp))
-
-    client.deleteRecordsBefore(Map((tp, 5L))).get()
-    consumer.seekToBeginning(Collections.singletonList(tp))
-    assertEquals(5L, consumer.position(tp))
-
-    client.deleteRecordsBefore(Map((tp, DeleteRecordsRequest.HIGH_WATERMARK))).get()
-    consumer.seekToBeginning(Collections.singletonList(tp))
-    assertEquals(10L, consumer.position(tp))
-  }
-
-  @Test
-  def testConsumeAfterDeleteRecords() {
-    val consumer = consumers.head
-    subscribeAndWaitForAssignment(topic, consumer)
-
-    sendRecords(producers.head, 10, tp)
-    var messageCount = 0
-    TestUtils.waitUntilTrue(() => {
-      messageCount += consumer.poll(0).count()
-      messageCount == 10
-    }, "Expected 10 messages", 3000L)
-
-    client.deleteRecordsBefore(Map((tp, 3L))).get()
-    consumer.seek(tp, 1)
-    messageCount = 0
-    TestUtils.waitUntilTrue(() => {
-      messageCount += consumer.poll(0).count()
-      messageCount == 7
-    }, "Expected 7 messages", 3000L)
-
-    client.deleteRecordsBefore(Map((tp, 8L))).get()
-    consumer.seek(tp, 1)
-    messageCount = 0
-    TestUtils.waitUntilTrue(() => {
-      messageCount += consumer.poll(0).count()
-      messageCount == 2
-    }, "Expected 2 messages", 3000L)
-  }
-
-  @Test
-  def testLogStartOffsetCheckpoint() {
-    subscribeAndWaitForAssignment(topic, consumers.head)
-
-    sendRecords(producers.head, 10, tp)
-    assertEquals(DeleteRecordsResult(5L, null), client.deleteRecordsBefore(Map((tp, 5L))).get()(tp))
-
-    for (i <- 0 until serverCount)
-      killBroker(i)
-    restartDeadBrokers()
-
-    client.close()
-    brokerList = TestUtils.bootstrapServers(servers, listenerName)
-    client = AdminClient.createSimplePlaintext(brokerList)
-
-    TestUtils.waitUntilTrue(() => {
-      // Need to retry if leader is not available for the partition
-      client.deleteRecordsBefore(Map((tp, 0L))).get(1000L, TimeUnit.MILLISECONDS)(tp).equals(DeleteRecordsResult(5L, null))
-    }, "Expected low watermark of the partition to be 5L")
-  }
-
-  @Test
-  def testLogStartOffsetAfterDeleteRecords() {
-    subscribeAndWaitForAssignment(topic, consumers.head)
-
-    sendRecords(producers.head, 10, tp)
-    client.deleteRecordsBefore(Map((tp, 3L))).get()
-
-    for (i <- 0 until serverCount)
-      assertEquals(3, servers(i).replicaManager.getReplica(tp).get.logStartOffset)
-  }
-
-  @Test
-  def testOffsetsForTimesWhenOffsetNotFound() {
-    val consumer = consumers.head
-    assertNull(consumer.offsetsForTimes(Map(tp -> new JLong(0L)).asJava).get(tp))
-  }
-
-  @Test
-  def testOffsetsForTimesAfterDeleteRecords() {
-    val consumer = consumers.head
-    subscribeAndWaitForAssignment(topic, consumer)
-
-    sendRecords(producers.head, 10, tp)
-    assertEquals(0L, consumer.offsetsForTimes(Map(tp -> new JLong(0L)).asJava).get(tp).offset())
-
-    client.deleteRecordsBefore(Map((tp, 5L))).get()
-    assertEquals(5L, consumer.offsetsForTimes(Map(tp -> new JLong(0L)).asJava).get(tp).offset())
-
-    client.deleteRecordsBefore(Map((tp, DeleteRecordsRequest.HIGH_WATERMARK))).get()
-    assertNull(consumer.offsetsForTimes(Map(tp -> new JLong(0L)).asJava).get(tp))
-  }
-
-  @Test
-  def testDeleteRecordsWithException() {
-    subscribeAndWaitForAssignment(topic, consumers.head)
-
-    sendRecords(producers.head, 10, tp)
-    // Should get success result
-    assertEquals(DeleteRecordsResult(5L, null), client.deleteRecordsBefore(Map((tp, 5L))).get()(tp))
-    // OffsetOutOfRangeException if offset > high_watermark
-    assertEquals(DeleteRecordsResult(-1L, Errors.OFFSET_OUT_OF_RANGE.exception()), client.deleteRecordsBefore(Map((tp, 20))).get()(tp))
-
-    val nonExistPartition = new TopicPartition(topic, 3)
-    // UnknownTopicOrPartitionException if user tries to delete records of a non-existent partition
-    assertEquals(DeleteRecordsResult(-1L, Errors.LEADER_NOT_AVAILABLE.exception()),
-                 client.deleteRecordsBefore(Map((nonExistPartition, 20))).get()(nonExistPartition))
-  }
-
-  @Test
-  def testListGroups() {
-    subscribeAndWaitForAssignment(topic, consumers.head)
-
-    val groups = client.listAllGroupsFlattened
-    assertFalse(groups.isEmpty)
-    val group = groups.head
-    assertEquals(groupId, group.groupId)
-    assertEquals("consumer", group.protocolType)
-  }
-
-  @Test
-  def testListAllBrokerVersionInfo() {
-    subscribeAndWaitForAssignment(topic, consumers.head)
-
-    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(1, brokerVersionInfo.latestUsableVersion(ApiKeys.API_VERSIONS))
-    }
-  }
-
-  @Test
-  def testGetConsumerGroupSummary() {
-    subscribeAndWaitForAssignment(topic, consumers.head)
-
-    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() {
-    subscribeAndWaitForAssignment(topic, consumers.head)
-
-    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")
-  }
-
-  private def sendRecords(producer: KafkaProducer[Array[Byte], Array[Byte]],
-                          numRecords: Int,
-                          tp: TopicPartition) {
-    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")
-      producer.send(record)
-    }
-
-    futures.foreach(_.get)
-  }
-
-}
diff --git a/core/src/test/scala/integration/kafka/api/LogAppendTimeTest.scala b/core/src/test/scala/integration/kafka/api/LogAppendTimeTest.scala
index 4a97bea680b26..74a774c682f67 100644
--- a/core/src/test/scala/integration/kafka/api/LogAppendTimeTest.scala
+++ b/core/src/test/scala/integration/kafka/api/LogAppendTimeTest.scala
@@ -21,15 +21,11 @@ import java.util.concurrent.TimeUnit
 
 import kafka.server.KafkaConfig
 import kafka.utils.TestUtils
-import org.apache.kafka.clients.consumer.ConsumerRecord
 import org.apache.kafka.clients.producer.ProducerRecord
 import org.apache.kafka.common.record.TimestampType
 import org.junit.{Before, Test}
 import org.junit.Assert.{assertEquals, assertNotEquals, assertTrue}
 
-import scala.collection.JavaConverters._
-import scala.collection.mutable.ArrayBuffer
-
 /**
   * Tests where the broker is configured to use LogAppendTime. For tests where LogAppendTime is configured via topic
   * level configs, see the *ProducerSendTest classes.
@@ -37,7 +33,7 @@ import scala.collection.mutable.ArrayBuffer
 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)
@@ -46,14 +42,14 @@ class LogAppendTimeTest extends IntegrationTestHarness {
   private val topic = "topic"
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
-    TestUtils.createTopic(zkUtils, topic, servers = servers)
+    createTopic(topic)
   }
 
   @Test
-  def testProduceConsume() {
-    val producer = producers.head
+  def testProduceConsume(): Unit = {
+    val producer = createProducer()
     val now = System.currentTimeMillis()
     val createTime = now - TimeUnit.DAYS.toMillis(1)
     val producerRecords = (1 to 10).map(i => new ProducerRecord(topic, null, createTime, s"key$i".getBytes,
@@ -64,13 +60,9 @@ class LogAppendTimeTest extends IntegrationTestHarness {
       assertTrue(recordMetadata.timestamp < now + TimeUnit.SECONDS.toMillis(60))
     }
 
-    val consumer = consumers.head
+    val consumer = createConsumer()
     consumer.subscribe(Collections.singleton(topic))
-    val consumerRecords = new ArrayBuffer[ConsumerRecord[Array[Byte], Array[Byte]]]
-    TestUtils.waitUntilTrue(() => {
-      consumerRecords ++= consumer.poll(50).asScala
-      consumerRecords.size == producerRecords.size
-    }, s"Consumed ${consumerRecords.size} records until timeout instead of the expected ${producerRecords.size} records")
+    val consumerRecords = TestUtils.consumeRecords(consumer, producerRecords.size)
 
     consumerRecords.zipWithIndex.foreach { case (consumerRecord, index) =>
       val producerRecord = producerRecords(index)
diff --git a/core/src/test/scala/integration/kafka/api/MetricsTest.scala b/core/src/test/scala/integration/kafka/api/MetricsTest.scala
index 26022be7266e3..6c90362049c5c 100644
--- a/core/src/test/scala/integration/kafka/api/MetricsTest.scala
+++ b/core/src/test/scala/integration/kafka/api/MetricsTest.scala
@@ -17,24 +17,25 @@ import java.util.{Locale, Properties}
 import kafka.log.LogConfig
 import kafka.server.{KafkaConfig, KafkaServer}
 import kafka.utils.{JaasTestUtils, TestUtils}
-import com.yammer.metrics.Metrics
 import com.yammer.metrics.core.{Gauge, Histogram, Meter}
+import kafka.metrics.KafkaYammerMetrics
+import org.apache.kafka.clients.consumer.KafkaConsumer
 import org.apache.kafka.clients.producer.{KafkaProducer, ProducerConfig, ProducerRecord}
 import org.apache.kafka.common.{Metric, MetricName, TopicPartition}
 import org.apache.kafka.common.config.SaslConfigs
 import org.apache.kafka.common.errors.InvalidTopicException
 import org.apache.kafka.common.network.ListenerName
 import org.apache.kafka.common.security.auth.SecurityProtocol
+import org.apache.kafka.common.security.authenticator.TestJaasConfig
 import org.junit.{After, Before, Test}
 import org.junit.Assert._
+import org.scalatest.Assertions.fail
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 
 class MetricsTest extends IntegrationTestHarness with SaslSetup {
 
-  override val producerCount = 1
-  override val consumerCount = 1
-  override val serverCount = 1
+  override val brokerCount = 1
 
   override protected def listenerName = new ListenerName("CLIENT")
   private val kafkaClientSaslMechanism = "PLAIN"
@@ -43,6 +44,9 @@ class MetricsTest extends IntegrationTestHarness with SaslSetup {
     s"${listenerName.value.toLowerCase(Locale.ROOT)}.${JaasTestUtils.KafkaServerContextName}"
   this.serverConfig.setProperty(KafkaConfig.ZkEnableSecureAclsProp, "false")
   this.serverConfig.setProperty(KafkaConfig.AutoCreateTopicsEnableDoc, "false")
+  this.producerConfig.setProperty(ProducerConfig.LINGER_MS_CONFIG, "10")
+  // intentionally slow message down conversion via gzip compression to ensure we can measure the time it takes
+  this.producerConfig.setProperty(ProducerConfig.COMPRESSION_TYPE_CONFIG, "gzip")
   override protected def securityProtocol = SecurityProtocol.SASL_PLAINTEXT
   override protected val serverSaslProperties =
     Some(kafkaServerSaslProperties(kafkaServerSaslMechanisms, kafkaClientSaslMechanism))
@@ -71,26 +75,26 @@ class MetricsTest extends IntegrationTestHarness with SaslSetup {
     val topic = "topicWithOldMessageFormat"
     val props = new Properties
     props.setProperty(LogConfig.MessageFormatVersionProp, "0.9.0")
-    TestUtils.createTopic(this.zkUtils, topic, numPartitions = 1, replicationFactor = 1, this.servers, props)
+    createTopic(topic, numPartitions = 1, replicationFactor = 1, props)
     val tp = new TopicPartition(topic, 0)
 
     // Produce and consume some records
     val numRecords = 10
-    val recordSize = 1000
-    val producer = producers.head
+    val recordSize = 100000
+    val producer = createProducer()
     sendRecords(producer, numRecords, recordSize, tp)
 
-    val consumer = this.consumers.head
+    val consumer = createConsumer()
     consumer.assign(List(tp).asJava)
     consumer.seek(tp, 0)
     TestUtils.consumeRecords(consumer, numRecords)
 
-    verifyKafkaRateMetricsHaveCumulativeCount()
+    verifyKafkaRateMetricsHaveCumulativeCount(producer, consumer)
     verifyClientVersionMetrics(consumer.metrics, "Consumer")
-    verifyClientVersionMetrics(this.producers.head.metrics, "Producer")
+    verifyClientVersionMetrics(producer.metrics, "Producer")
 
     val server = servers.head
-    verifyBrokerMessageConversionMetrics(server, recordSize)
+    verifyBrokerMessageConversionMetrics(server, recordSize, tp)
     verifyBrokerErrorMetrics(servers.head)
     verifyBrokerZkMetrics(server, topic)
 
@@ -109,16 +113,17 @@ class MetricsTest extends IntegrationTestHarness with SaslSetup {
 
   // Create a producer that fails authentication to verify authentication failure metrics
   private def generateAuthenticationFailure(tp: TopicPartition): Unit = {
-    val producerProps = new Properties()
     val saslProps = new Properties()
-     // Temporary limit to reduce blocking before KIP-152 client-side changes are merged
-    saslProps.put(ProducerConfig.REQUEST_TIMEOUT_MS_CONFIG, "1000")
-    saslProps.put(ProducerConfig.MAX_BLOCK_MS_CONFIG, "1000")
-    saslProps.put(SaslConfigs.SASL_MECHANISM, "SCRAM-SHA-256")
+    saslProps.put(SaslConfigs.SASL_MECHANISM, kafkaClientSaslMechanism)
+    saslProps.put(SaslConfigs.SASL_JAAS_CONFIG, TestJaasConfig.jaasConfigProperty(kafkaClientSaslMechanism, "badUser", "badPass"))
     // Use acks=0 to verify error metric when connection is closed without a response
-    saslProps.put(ProducerConfig.ACKS_CONFIG, "0")
-    val producer = TestUtils.createNewProducer(brokerList, securityProtocol = securityProtocol,
-        trustStoreFile = trustStoreFile, saslProperties = Some(saslProps), props = Some(producerProps))
+    val producer = TestUtils.createProducer(brokerList,
+      acks = 0,
+      requestTimeoutMs = 1000,
+      maxBlockMs = 1000,
+      securityProtocol = securityProtocol,
+      trustStoreFile = trustStoreFile,
+      saslProperties = Some(saslProps))
 
     try {
       producer.send(new ProducerRecord(tp.topic, tp.partition, "key".getBytes, "value".getBytes)).get
@@ -129,7 +134,8 @@ class MetricsTest extends IntegrationTestHarness with SaslSetup {
     }
   }
 
-  private def verifyKafkaRateMetricsHaveCumulativeCount(): Unit =  {
+  private def verifyKafkaRateMetricsHaveCumulativeCount(producer: KafkaProducer[Array[Byte], Array[Byte]],
+                                                        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))
@@ -143,12 +149,10 @@ class MetricsTest extends IntegrationTestHarness with SaslSetup {
           totalExists || totalTimeExists)
     }
 
-    val consumer = this.consumers.head
     val consumerMetricNames = consumer.metrics.keySet.asScala.toSet
     consumerMetricNames.filter(_.name.endsWith("-rate"))
         .foreach(verify(_, consumerMetricNames))
 
-    val producer = this.producers.head
     val producerMetricNames = producer.metrics.keySet.asScala.toSet
     val producerExclusions = Set("compression-rate") // compression-rate is an Average metric, not Rate
     producerMetricNames.filter(_.name.endsWith("-rate"))
@@ -187,7 +191,7 @@ class MetricsTest extends IntegrationTestHarness with SaslSetup {
     verifyKafkaMetricRecorded("failed-authentication-total", metrics, "Broker", Some("socket-server-metrics"))
   }
 
-  private def verifyBrokerMessageConversionMetrics(server: KafkaServer, recordSize: Int): Unit = {
+  private def verifyBrokerMessageConversionMetrics(server: KafkaServer, recordSize: Int, tp: TopicPartition): Unit = {
     val requestMetricsPrefix = "kafka.network:type=RequestMetrics"
     val requestBytes = verifyYammerMetricRecorded(s"$requestMetricsPrefix,name=RequestBytes,request=Produce")
     val tempBytes = verifyYammerMetricRecorded(s"$requestMetricsPrefix,name=TemporaryMemoryBytes,request=Produce")
@@ -196,7 +200,6 @@ class MetricsTest extends IntegrationTestHarness with SaslSetup {
 
     verifyYammerMetricRecorded(s"kafka.server:type=BrokerTopicMetrics,name=ProduceMessageConversionsPerSec")
     verifyYammerMetricRecorded(s"$requestMetricsPrefix,name=MessageConversionsTimeMs,request=Produce", value => value > 0.0)
-
     verifyYammerMetricRecorded(s"$requestMetricsPrefix,name=RequestBytes,request=Fetch")
     verifyYammerMetricRecorded(s"$requestMetricsPrefix,name=TemporaryMemoryBytes,request=Fetch", value => value == 0.0)
 
@@ -205,27 +208,30 @@ class MetricsTest extends IntegrationTestHarness with SaslSetup {
   }
 
   private def verifyBrokerZkMetrics(server: KafkaServer, topic: String): Unit = {
-    // Latency is rounded to milliseconds, so we may need to retry some operations to get latency > 0.
-    val (_, recorded) = TestUtils.computeUntilTrue({
-      servers.head.zkUtils.getLeaderAndIsrForPartition(topic, 0)
-      yammerMetricValue("kafka.server:type=ZooKeeperClientMetrics,name=ZooKeeperRequestLatencyMs").asInstanceOf[Double]
-    })(latency => latency > 0.0)
-    assertTrue("ZooKeeper latency not recorded", recorded)
-
-    assertEquals(s"Unexpected ZK state ${server.zkUtils.zkConnection.getZookeeperState}",
-        "CONNECTED", yammerMetricValue("SessionState"))
+    val histogram = yammerHistogram("kafka.server:type=ZooKeeperClientMetrics,name=ZooKeeperRequestLatencyMs")
+    // Latency is rounded to milliseconds, so check the count instead
+    val initialCount = histogram.count
+    servers.head.zkClient.getLeaderForPartition(new TopicPartition(topic, 0))
+    val newCount = histogram.count
+    assertTrue("ZooKeeper latency not recorded",  newCount > initialCount)
+
+    val min = histogram.min
+    assertTrue(s"Min latency should not be negative: $min", min >= 0)
+
+    assertEquals(s"Unexpected ZK state", "CONNECTED", yammerMetricValue("SessionState"))
   }
 
   private def verifyBrokerErrorMetrics(server: KafkaServer): Unit = {
 
-    def errorMetricCount = Metrics.defaultRegistry.allMetrics.keySet.asScala.filter(_.getName == "ErrorsPerSec").size
+    def errorMetricCount = KafkaYammerMetrics.defaultRegistry.allMetrics.keySet.asScala.filter(_.getName == "ErrorsPerSec").size
 
     val startErrorMetricCount = errorMetricCount
     val errorMetricPrefix = "kafka.network:type=RequestMetrics,name=ErrorsPerSec"
     verifyYammerMetricRecorded(s"$errorMetricPrefix,request=Metadata,error=NONE")
 
     try {
-      consumers.head.partitionsFor("12{}!")
+      val consumer = createConsumer()
+      consumer.partitionsFor("12{}!")
     } catch {
       case _: InvalidTopicException => // expected
     }
@@ -237,7 +243,8 @@ class MetricsTest extends IntegrationTestHarness with SaslSetup {
     assertTrue(s"Too many error metrics $currentErrorMetricCount" , currentErrorMetricCount < 10)
 
     // Verify that error metric is updated with producer acks=0 when no response is sent
-    sendRecords(producers.head, 1, 100, new TopicPartition("non-existent", 0))
+    val producer = createProducer()
+    sendRecords(producer, numRecords = 1, recordSize = 100, new TopicPartition("non-existent", 0))
     verifyYammerMetricRecorded(s"$errorMetricPrefix,request=Metadata,error=LEADER_NOT_AVAILABLE")
   }
 
@@ -246,7 +253,7 @@ class MetricsTest extends IntegrationTestHarness with SaslSetup {
     val matchingMetrics = metrics.asScala.filter {
       case (metricName, _) => metricName.name == name && group.forall(_ == metricName.group)
     }
-    assertTrue(s"Metric not found $name", matchingMetrics.size > 0)
+    assertTrue(s"Metric not found $name", matchingMetrics.nonEmpty)
     verify(matchingMetrics.values)
   }
 
@@ -254,7 +261,7 @@ class MetricsTest extends IntegrationTestHarness with SaslSetup {
       group: Option[String]): Double = {
     // Use max value of all matching metrics since Selector metrics are recorded for each Processor
     verifyKafkaMetric(name, metrics, entity, group) { matchingMetrics =>
-      matchingMetrics.foldLeft(0.0)((max, metric) => Math.max(max, metric.value))
+      matchingMetrics.foldLeft(0.0)((max, metric) => Math.max(max, metric.metricValue.asInstanceOf[Double]))
     }
   }
 
@@ -265,7 +272,7 @@ class MetricsTest extends IntegrationTestHarness with SaslSetup {
   }
 
   private def yammerMetricValue(name: String): Any = {
-    val allMetrics = Metrics.defaultRegistry.allMetrics.asScala
+    val allMetrics = KafkaYammerMetrics.defaultRegistry.allMetrics.asScala
     val (_, metric) = allMetrics.find { case (n, _) => n.getMBeanName.endsWith(name) }
       .getOrElse(fail(s"Unable to find broker metric $name: allMetrics: ${allMetrics.keySet.map(_.getMBeanName)}"))
     metric match {
@@ -276,6 +283,16 @@ class MetricsTest extends IntegrationTestHarness with SaslSetup {
     }
   }
 
+  private def yammerHistogram(name: String): Histogram = {
+    val allMetrics = KafkaYammerMetrics.defaultRegistry.allMetrics.asScala
+    val (_, metric) = allMetrics.find { case (n, _) => n.getMBeanName.endsWith(name) }
+      .getOrElse(fail(s"Unable to find broker metric $name: allMetrics: ${allMetrics.keySet.map(_.getMBeanName)}"))
+    metric match {
+      case m: Histogram => m
+      case m => fail(s"Unexpected broker metric of class ${m.getClass}")
+    }
+  }
+
   private def verifyYammerMetricRecorded(name: String, verify: Double => Boolean = d => d > 0): Double = {
     val metricValue = yammerMetricValue(name).asInstanceOf[Double]
     assertTrue(s"Broker metric not recorded correctly for $name value $metricValue", verify(metricValue))
@@ -283,7 +300,7 @@ class MetricsTest extends IntegrationTestHarness with SaslSetup {
   }
 
   private def verifyNoRequestMetrics(errorMessage: String): Unit = {
-    val metrics = Metrics.defaultRegistry.allMetrics.asScala.filter { case (n, _) =>
+    val metrics = KafkaYammerMetrics.defaultRegistry.allMetrics.asScala.filter { case (n, _) =>
       n.getMBeanName.startsWith("kafka.network:type=RequestMetrics")
     }
     assertTrue(s"$errorMessage: ${metrics.keys}", metrics.isEmpty)
diff --git a/core/src/test/scala/integration/kafka/api/PlaintextAdminIntegrationTest.scala b/core/src/test/scala/integration/kafka/api/PlaintextAdminIntegrationTest.scala
new file mode 100644
index 0000000000000..c0e0de9fe8a4b
--- /dev/null
+++ b/core/src/test/scala/integration/kafka/api/PlaintextAdminIntegrationTest.scala
@@ -0,0 +1,2431 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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.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 kafka.log.LogConfig
+import kafka.security.authorizer.AclEntry
+import kafka.server.{Defaults, DynamicConfig, KafkaConfig, KafkaServer}
+import kafka.utils.TestUtils._
+import kafka.utils.{Log4jController, 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, ProducerRecord}
+import org.apache.kafka.common.acl.{AccessControlEntry, AclBinding, AclBindingFilter, AclOperation, AclPermissionType}
+import org.apache.kafka.common.config.{ConfigResource, LogLevelConfig}
+import org.apache.kafka.common.errors._
+import org.apache.kafka.common.requests.{DeleteRecordsRequest, MetadataResponse}
+import org.apache.kafka.common.resource.{PatternType, ResourcePattern, ResourceType}
+import org.apache.kafka.common.utils.{Time, Utils}
+import org.apache.kafka.common.{ConsumerGroupState, ElectionType, TopicPartition, TopicPartitionInfo, TopicPartitionReplica}
+import org.junit.Assert._
+import org.junit.{After, Before, Ignore, Test}
+import org.scalatest.Assertions.intercept
+import org.slf4j.LoggerFactory
+
+import scala.annotation.nowarn
+import scala.jdk.CollectionConverters._
+import scala.collection.Seq
+import scala.concurrent.duration.Duration
+import scala.concurrent.{Await, Future}
+import scala.util.Random
+
+/**
+ * An integration test of the KafkaAdminClient.
+ *
+ * Also see [[org.apache.kafka.clients.admin.KafkaAdminClientTest]] for unit tests of the admin client.
+ */
+class PlaintextAdminIntegrationTest extends BaseAdminIntegrationTest {
+  import PlaintextAdminIntegrationTest._
+
+  val topic = "topic"
+  val partition = 0
+  val topicPartition = new TopicPartition(topic, partition)
+
+  private var brokerLoggerConfigResource: ConfigResource = _
+  private val changedBrokerLoggers = scala.collection.mutable.Set[String]()
+
+  @Before
+  override def setUp(): Unit = {
+    super.setUp()
+    brokerLoggerConfigResource = new ConfigResource(
+      ConfigResource.Type.BROKER_LOGGER, servers.head.config.brokerId.toString)
+  }
+
+  @After
+  override def tearDown(): Unit = {
+    teardownBrokerLoggers()
+    super.tearDown()
+  }
+
+  @Test
+  def testClose(): Unit = {
+    val client = Admin.create(createConfig)
+    client.close()
+    client.close() // double close has no effect
+  }
+
+  @Test
+  def testListNodes(): Unit = {
+    client = Admin.create(createConfig)
+    val brokerStrs = brokerList.split(",").toList.sorted
+    var nodeStrs: List[String] = null
+    do {
+      val nodes = client.describeCluster().nodes().get().asScala
+      nodeStrs = nodes.map ( node => s"${node.host}:${node.port}" ).toList.sorted
+    } while (nodeStrs.size < brokerStrs.size)
+    assertEquals(brokerStrs.mkString(","), nodeStrs.mkString(","))
+  }
+
+  @Test
+  def testCreateExistingTopicsThrowTopicExistsException(): Unit = {
+    client = Admin.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 = Admin.create(createConfig)
+    val topics = Seq("mytopic")
+    val newTopics = Seq(new NewTopic("mytopic", 3, 3.toShort))
+    client.createTopics(newTopics.asJava).all.get()
+    waitForTopics(client, expectedPresent = topics, expectedMissing = List())
+
+    val controller = servers.find(_.config.brokerId == TestUtils.waitUntilControllerElected(zkClient)).get
+    controller.shutdown()
+    controller.awaitShutdown()
+    val topicDesc = client.describeTopics(topics.asJava).all.get()
+    assertEquals(topics.toSet, topicDesc.keySet.asScala)
+  }
+
+  /**
+    * describe should not auto create topics
+    */
+  @Test
+  def testDescribeNonExistingTopic(): Unit = {
+    client = Admin.create(createConfig)
+
+    val existingTopic = "existing-topic"
+    client.createTopics(Seq(existingTopic).map(new NewTopic(_, 1, 1.toShort)).asJava).all.get()
+    waitForTopics(client, Seq(existingTopic), List())
+
+    val nonExistingTopic = "non-existing"
+    val results = client.describeTopics(Seq(nonExistingTopic, existingTopic).asJava).values
+    assertEquals(existingTopic, results.get(existingTopic).get.name)
+    intercept[ExecutionException](results.get(nonExistingTopic).get).getCause.isInstanceOf[UnknownTopicOrPartitionException]
+    assertEquals(None, zkClient.getTopicPartitionCount(nonExistingTopic))
+  }
+
+  @Test
+  def testDescribeCluster(): Unit = {
+    client = Admin.create(createConfig)
+    val result = client.describeCluster
+    val nodes = result.nodes.get()
+    val clusterId = result.clusterId().get()
+    assertEquals(servers.head.dataPlaneRequestProcessor.clusterId, clusterId)
+    val controller = result.controller().get()
+    assertEquals(servers.head.dataPlaneRequestProcessor.metadataCache.getControllerId.
+      getOrElse(MetadataResponse.NO_CONTROLLER_ID), controller.id())
+    val brokers = brokerList.split(",")
+    assertEquals(brokers.size, nodes.size)
+    for (node <- nodes.asScala) {
+      val hostStr = s"${node.host}:${node.port}"
+      assertTrue(s"Unknown host:port pair $hostStr in brokerVersionInfos", brokers.contains(hostStr))
+    }
+  }
+
+  @Test
+  def testDescribeLogDirs(): Unit = {
+    client = Admin.create(createConfig)
+    val topic = "topic"
+    val leaderByPartition = createTopic(topic, numPartitions = 10, replicationFactor = 1)
+    val partitionsByBroker = leaderByPartition.groupBy { case (_, leaderId) => leaderId }.map { case (k, v) =>
+      k -> v.keys.toSeq
+    }
+    val brokers = (0 until brokerCount).map(Integer.valueOf)
+    val logDirInfosByBroker = client.describeLogDirs(brokers.asJava).allDescriptions.get
+
+    (0 until brokerCount).foreach { brokerId =>
+      val server = servers.find(_.config.brokerId == brokerId).get
+      val expectedPartitions = partitionsByBroker(brokerId)
+      val logDirInfos = logDirInfosByBroker.get(brokerId)
+      val replicaInfos = logDirInfos.asScala.flatMap { case (_, logDirInfo) =>
+        logDirInfo.replicaInfos.asScala
+      }.filter { case (k, _) => k.topic == topic }
+
+      assertEquals(expectedPartitions.toSet, replicaInfos.keys.map(_.partition).toSet)
+      logDirInfos.forEach { (logDir, logDirInfo) =>
+        logDirInfo.replicaInfos.asScala.keys.foreach(tp =>
+          assertEquals(server.logManager.getLog(tp).get.dir.getParent, logDir)
+        )
+      }
+    }
+  }
+
+  @Test
+  def testDescribeReplicaLogDirs(): Unit = {
+    client = Admin.create(createConfig)
+    val topic = "topic"
+    val leaderByPartition = createTopic(topic, numPartitions = 10, replicationFactor = 1)
+    val replicas = leaderByPartition.map { case (partition, brokerId) =>
+      new TopicPartitionReplica(topic, partition, brokerId)
+    }.toSeq
+
+    val replicaDirInfos = client.describeReplicaLogDirs(replicas.asJavaCollection).all.get
+    replicaDirInfos.forEach { (topicPartitionReplica, replicaDirInfo) =>
+      val server = servers.find(_.config.brokerId == topicPartitionReplica.brokerId()).get
+      val tp = new TopicPartition(topicPartitionReplica.topic(), topicPartitionReplica.partition())
+      assertEquals(server.logManager.getLog(tp).get.dir.getParent, replicaDirInfo.getCurrentReplicaLogDir)
+    }
+  }
+
+  @Test
+  def testAlterReplicaLogDirs(): Unit = {
+    client = Admin.create(createConfig)
+    val topic = "topic"
+    val tp = new TopicPartition(topic, 0)
+    val randomNums = servers.map(server => server -> Random.nextInt(2)).toMap
+
+    // Generate two mutually exclusive replicaAssignment
+    val firstReplicaAssignment = servers.map { server =>
+      val logDir = new File(server.config.logDirs(randomNums(server))).getAbsolutePath
+      new TopicPartitionReplica(topic, 0, server.config.brokerId) -> logDir
+    }.toMap
+    val secondReplicaAssignment = servers.map { server =>
+      val logDir = new File(server.config.logDirs(1 - randomNums(server))).getAbsolutePath
+      new TopicPartitionReplica(topic, 0, server.config.brokerId) -> logDir
+    }.toMap
+
+    // Verify that replica can be created in the specified log directory
+    val futures = client.alterReplicaLogDirs(firstReplicaAssignment.asJava,
+      new AlterReplicaLogDirsOptions).values.asScala.values
+    futures.foreach { future =>
+      val exception = intercept[ExecutionException](future.get)
+      assertTrue(exception.getCause.isInstanceOf[UnknownTopicOrPartitionException])
+    }
+
+    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)
+    }
+
+    // Verify that replica can be moved to the specified log directory after the topic has been created
+    client.alterReplicaLogDirs(secondReplicaAssignment.asJava, new AlterReplicaLogDirsOptions).all.get
+    servers.foreach { server =>
+      TestUtils.waitUntilTrue(() => {
+        val logDir = server.logManager.getLog(tp).get.dir.getParent
+        secondReplicaAssignment(new TopicPartitionReplica(topic, 0, server.config.brokerId)) == logDir
+      }, "timed out waiting for replica movement")
+    }
+
+    // Verify that replica can be moved to the specified log directory while the producer is sending messages
+    val running = new AtomicBoolean(true)
+    val numMessages = new AtomicInteger
+    import scala.concurrent.ExecutionContext.Implicits._
+    val producerFuture = Future {
+      val producer = TestUtils.createProducer(
+        TestUtils.getBrokerListStrFromServers(servers, protocol = securityProtocol),
+        securityProtocol = securityProtocol,
+        trustStoreFile = trustStoreFile,
+        retries = 0, // Producer should not have to retry when broker is moving replica between log directories.
+        requestTimeoutMs = 10000,
+        acks = -1
+      )
+      try {
+        while (running.get) {
+          val future = producer.send(new ProducerRecord(topic, s"xxxxxxxxxxxxxxxxxxxx-$numMessages".getBytes))
+          numMessages.incrementAndGet()
+          future.get(10, TimeUnit.SECONDS)
+        }
+        numMessages.get
+      } finally producer.close()
+    }
+
+    try {
+      TestUtils.waitUntilTrue(() => numMessages.get > 10, s"only $numMessages messages are produced before timeout. Producer future ${producerFuture.value}")
+      client.alterReplicaLogDirs(firstReplicaAssignment.asJava, new AlterReplicaLogDirsOptions).all.get
+      servers.foreach { server =>
+        TestUtils.waitUntilTrue(() => {
+          val logDir = server.logManager.getLog(tp).get.dir.getParent
+          firstReplicaAssignment(new TopicPartitionReplica(topic, 0, server.config.brokerId)) == logDir
+        }, s"timed out waiting for replica movement. Producer future ${producerFuture.value}")
+      }
+
+      val currentMessagesNum = numMessages.get
+      TestUtils.waitUntilTrue(() => numMessages.get - currentMessagesNum > 10,
+        s"only ${numMessages.get - currentMessagesNum} messages are produced within timeout after replica movement. Producer future ${producerFuture.value}")
+    } finally running.set(false)
+
+    val finalNumMessages = Await.result(producerFuture, Duration(20, TimeUnit.SECONDS))
+
+    // Verify that all messages that are produced can be consumed
+    val consumerRecords = TestUtils.consumeTopicRecords(servers, topic, finalNumMessages,
+      securityProtocol = securityProtocol, trustStoreFile = trustStoreFile)
+    consumerRecords.zipWithIndex.foreach { case (consumerRecord, index) =>
+      assertEquals(s"xxxxxxxxxxxxxxxxxxxx-$index", new String(consumerRecord.value))
+    }
+  }
+
+  @Test
+  def testDescribeAndAlterConfigs(): Unit = {
+    client = Admin.create(createConfig)
+
+    // Create topics
+    val topic1 = "describe-alter-configs-topic-1"
+    val topicResource1 = new ConfigResource(ConfigResource.Type.TOPIC, topic1)
+    val topicConfig1 = new Properties
+    topicConfig1.setProperty(LogConfig.MaxMessageBytesProp, "500000")
+    topicConfig1.setProperty(LogConfig.RetentionMsProp, "60000000")
+    createTopic(topic1, numPartitions = 1, replicationFactor = 1, topicConfig1)
+
+    val topic2 = "describe-alter-configs-topic-2"
+    val topicResource2 = new ConfigResource(ConfigResource.Type.TOPIC, topic2)
+    createTopic(topic2, numPartitions = 1, replicationFactor = 1)
+
+    // Describe topics and broker
+    val brokerResource1 = new ConfigResource(ConfigResource.Type.BROKER, servers(1).config.brokerId.toString)
+    val brokerResource2 = new ConfigResource(ConfigResource.Type.BROKER, servers(2).config.brokerId.toString)
+    val configResources = Seq(topicResource1, topicResource2, brokerResource1, brokerResource2)
+    val describeResult = client.describeConfigs(configResources.asJava)
+    val configs = describeResult.all.get
+
+    assertEquals(4, configs.size)
+
+    val maxMessageBytes1 = configs.get(topicResource1).get(LogConfig.MaxMessageBytesProp)
+    assertEquals(LogConfig.MaxMessageBytesProp, maxMessageBytes1.name)
+    assertEquals(topicConfig1.get(LogConfig.MaxMessageBytesProp), maxMessageBytes1.value)
+    assertFalse(maxMessageBytes1.isDefault)
+    assertFalse(maxMessageBytes1.isSensitive)
+    assertFalse(maxMessageBytes1.isReadOnly)
+
+    assertEquals(topicConfig1.get(LogConfig.RetentionMsProp),
+      configs.get(topicResource1).get(LogConfig.RetentionMsProp).value)
+
+    val maxMessageBytes2 = configs.get(topicResource2).get(LogConfig.MaxMessageBytesProp)
+    assertEquals(Defaults.MessageMaxBytes.toString, maxMessageBytes2.value)
+    assertEquals(LogConfig.MaxMessageBytesProp, maxMessageBytes2.name)
+    assertTrue(maxMessageBytes2.isDefault)
+    assertFalse(maxMessageBytes2.isSensitive)
+    assertFalse(maxMessageBytes2.isReadOnly)
+
+    assertEquals(servers(1).config.nonInternalValues.size, configs.get(brokerResource1).entries.size)
+    assertEquals(servers(1).config.brokerId.toString, configs.get(brokerResource1).get(KafkaConfig.BrokerIdProp).value)
+    val listenerSecurityProtocolMap = configs.get(brokerResource1).get(KafkaConfig.ListenerSecurityProtocolMapProp)
+    assertEquals(servers(1).config.getString(KafkaConfig.ListenerSecurityProtocolMapProp), listenerSecurityProtocolMap.value)
+    assertEquals(KafkaConfig.ListenerSecurityProtocolMapProp, listenerSecurityProtocolMap.name)
+    assertFalse(listenerSecurityProtocolMap.isDefault)
+    assertFalse(listenerSecurityProtocolMap.isSensitive)
+    assertFalse(listenerSecurityProtocolMap.isReadOnly)
+    val truststorePassword = configs.get(brokerResource1).get(KafkaConfig.SslTruststorePasswordProp)
+    assertEquals(KafkaConfig.SslTruststorePasswordProp, truststorePassword.name)
+    assertNull(truststorePassword.value)
+    assertFalse(truststorePassword.isDefault)
+    assertTrue(truststorePassword.isSensitive)
+    assertFalse(truststorePassword.isReadOnly)
+    val compressionType = configs.get(brokerResource1).get(KafkaConfig.CompressionTypeProp)
+    assertEquals(servers(1).config.compressionType, compressionType.value)
+    assertEquals(KafkaConfig.CompressionTypeProp, compressionType.name)
+    assertTrue(compressionType.isDefault)
+    assertFalse(compressionType.isSensitive)
+    assertFalse(compressionType.isReadOnly)
+
+    assertEquals(servers(2).config.nonInternalValues.size, configs.get(brokerResource2).entries.size)
+    assertEquals(servers(2).config.brokerId.toString, configs.get(brokerResource2).get(KafkaConfig.BrokerIdProp).value)
+    assertEquals(servers(2).config.logCleanerThreads.toString,
+      configs.get(brokerResource2).get(KafkaConfig.LogCleanerThreadsProp).value)
+
+    checkValidAlterConfigs(client, topicResource1, topicResource2)
+  }
+
+  @Test
+  def testCreatePartitions(): Unit = {
+    client = Admin.create(createConfig)
+
+    // Create topics
+    val topic1 = "create-partitions-topic-1"
+    createTopic(topic1, numPartitions = 1, replicationFactor = 1)
+
+    val topic2 = "create-partitions-topic-2"
+    createTopic(topic2, numPartitions = 1, replicationFactor = 2)
+
+    // assert that both the topics have 1 partition
+    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, expectedNumPartitionsOpt: Option[Int]): util.List[TopicPartitionInfo] = {
+      getTopicMetadata(client, topic, expectedNumPartitionsOpt = expectedNumPartitionsOpt).partitions
+    }
+
+    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 ->
+      NewPartitions.increaseTo(3)).asJava, validateOnly)
+    var altered = alterResult.values.get(topic1).get
+    assertEquals(1, numPartitions(topic1))
+
+    // try creating a new partition (no assignments), to bring the total to 3 partitions
+    alterResult = client.createPartitions(Map(topic1 ->
+      NewPartitions.increaseTo(3)).asJava, actuallyDoIt)
+    altered = alterResult.values.get(topic1).get
+    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))
+    alterResult = client.createPartitions(Map(topic2 ->
+      NewPartitions.increaseTo(3, newPartition2Assignments)).asJava, validateOnly)
+    altered = alterResult.values.get(topic2).get
+    assertEquals(1, numPartitions(topic2))
+
+    // now try creating a new partition (with assignments), to bring the total to 3 partitions
+    alterResult = client.createPartitions(Map(topic2 ->
+      NewPartitions.increaseTo(3, newPartition2Assignments)).asJava, actuallyDoIt)
+    altered = alterResult.values.get(topic2).get
+    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)
+
+    // loop over error cases calling with+without validate-only
+    for (option <- Seq(validateOnly, actuallyDoIt)) {
+      val desc = if (option.validateOnly()) "validateOnly" else "validateOnly=false"
+
+      // try a newCount which would be a decrease
+      alterResult = client.createPartitions(Map(topic1 ->
+        NewPartitions.increaseTo(1)).asJava, option)
+      try {
+        alterResult.values.get(topic1).get
+        fail(s"$desc: Expect InvalidPartitionsException when newCount is a decrease")
+      } catch {
+        case e: ExecutionException =>
+          assertTrue(desc, e.getCause.isInstanceOf[InvalidPartitionsException])
+          assertEquals(desc, "Topic currently has 3 partitions, which is higher than the requested 1.", e.getCause.getMessage)
+          assertEquals(desc, 3, numPartitions(topic1))
+      }
+
+      // try a newCount which would be a noop (without assignment)
+      alterResult = client.createPartitions(Map(topic2 ->
+        NewPartitions.increaseTo(3)).asJava, option)
+      try {
+        alterResult.values.get(topic2).get
+        fail(s"$desc: Expect InvalidPartitionsException when requesting a noop")
+      } catch {
+        case e: ExecutionException =>
+          assertTrue(desc, e.getCause.isInstanceOf[InvalidPartitionsException])
+          assertEquals(desc, "Topic already has 3 partitions.", e.getCause.getMessage)
+          assertEquals(desc, 3, numPartitions(topic2, Some(3)))
+      }
+
+      // try a newCount which would be a noop (where the assignment matches current state)
+      alterResult = client.createPartitions(Map(topic2 ->
+        NewPartitions.increaseTo(3, newPartition2Assignments)).asJava, option)
+      try {
+        alterResult.values.get(topic2).get
+      } catch {
+        case e: ExecutionException =>
+          assertTrue(desc, e.getCause.isInstanceOf[InvalidPartitionsException])
+          assertEquals(desc, "Topic already has 3 partitions.", e.getCause.getMessage)
+          assertEquals(desc, 3, numPartitions(topic2, Some(3)))
+      }
+
+      // try a newCount which would be a noop (where the assignment doesn't match current state)
+      alterResult = client.createPartitions(Map(topic2 ->
+        NewPartitions.increaseTo(3, newPartition2Assignments.asScala.reverse.toList.asJava)).asJava, option)
+      try {
+        alterResult.values.get(topic2).get
+      } catch {
+        case e: ExecutionException =>
+          assertTrue(desc, e.getCause.isInstanceOf[InvalidPartitionsException])
+          assertEquals(desc, "Topic already has 3 partitions.", e.getCause.getMessage)
+          assertEquals(desc, 3, numPartitions(topic2, Some(3)))
+      }
+
+      // try a bad topic name
+      val unknownTopic = "an-unknown-topic"
+      alterResult = client.createPartitions(Map(unknownTopic ->
+        NewPartitions.increaseTo(2)).asJava, option)
+      try {
+        alterResult.values.get(unknownTopic).get
+        fail(s"$desc: Expect InvalidTopicException when using an unknown topic")
+      } catch {
+        case e: ExecutionException =>
+          assertTrue(desc, e.getCause.isInstanceOf[UnknownTopicOrPartitionException])
+          assertEquals(desc, "The topic 'an-unknown-topic' does not exist.", e.getCause.getMessage)
+      }
+
+      // try an invalid newCount
+      alterResult = client.createPartitions(Map(topic1 ->
+        NewPartitions.increaseTo(-22)).asJava, option)
+      try {
+        altered = alterResult.values.get(topic1).get
+        fail(s"$desc: Expect InvalidPartitionsException when newCount is invalid")
+      } catch {
+        case e: ExecutionException =>
+          assertTrue(desc, e.getCause.isInstanceOf[InvalidPartitionsException])
+          assertEquals(desc, "Topic currently has 3 partitions, which is higher than the requested -22.",
+            e.getCause.getMessage)
+          assertEquals(desc, 3, numPartitions(topic1))
+      }
+
+      // try assignments where the number of brokers != replication factor
+      alterResult = client.createPartitions(Map(topic1 ->
+        NewPartitions.increaseTo(4, asList(asList(1, 2)))).asJava, option)
+      try {
+        altered = alterResult.values.get(topic1).get
+        fail(s"$desc: Expect InvalidPartitionsException when #brokers != replication factor")
+      } catch {
+        case e: ExecutionException =>
+          assertTrue(desc, e.getCause.isInstanceOf[InvalidReplicaAssignmentException])
+          assertEquals(desc, "Inconsistent replication factor between partitions, partition 0 has 1 " +
+            "while partitions [3] have replication factors [2], respectively.",
+            e.getCause.getMessage)
+          assertEquals(desc, 3, numPartitions(topic1))
+      }
+
+      // try #assignments < with the increase
+      alterResult = client.createPartitions(Map(topic1 ->
+        NewPartitions.increaseTo(6, asList(asList(1)))).asJava, option)
+      try {
+        altered = alterResult.values.get(topic1).get
+        fail(s"$desc: Expect InvalidReplicaAssignmentException when #assignments != newCount - oldCount")
+      } catch {
+        case e: ExecutionException =>
+          assertTrue(desc, e.getCause.isInstanceOf[InvalidReplicaAssignmentException])
+          assertEquals(desc, "Increasing the number of partitions by 3 but 1 assignments provided.", e.getCause.getMessage)
+          assertEquals(desc, 3, numPartitions(topic1))
+      }
+
+      // try #assignments > with the increase
+      alterResult = client.createPartitions(Map(topic1 ->
+        NewPartitions.increaseTo(4, asList(asList(1), asList(2)))).asJava, option)
+      try {
+        altered = alterResult.values.get(topic1).get
+        fail(s"$desc: Expect InvalidReplicaAssignmentException when #assignments != newCount - oldCount")
+      } catch {
+        case e: ExecutionException =>
+          assertTrue(desc, e.getCause.isInstanceOf[InvalidReplicaAssignmentException])
+          assertEquals(desc, "Increasing the number of partitions by 1 but 2 assignments provided.", e.getCause.getMessage)
+          assertEquals(desc, 3, numPartitions(topic1))
+      }
+
+      // try with duplicate brokers in assignments
+      alterResult = client.createPartitions(Map(topic1 ->
+        NewPartitions.increaseTo(4, asList(asList(1, 1)))).asJava, option)
+      try {
+        altered = alterResult.values.get(topic1).get
+        fail(s"$desc: Expect InvalidReplicaAssignmentException when assignments has duplicate brokers")
+      } catch {
+        case e: ExecutionException =>
+          assertTrue(desc, e.getCause.isInstanceOf[InvalidReplicaAssignmentException])
+          assertEquals(desc, "Duplicate brokers not allowed in replica assignment: 1, 1 for partition id 3.",
+            e.getCause.getMessage)
+          assertEquals(desc, 3, numPartitions(topic1))
+      }
+
+      // try assignments with differently sized inner lists
+      alterResult = client.createPartitions(Map(topic1 ->
+        NewPartitions.increaseTo(5, asList(asList(1), asList(1, 0)))).asJava, option)
+      try {
+        altered = alterResult.values.get(topic1).get
+        fail(s"$desc: Expect InvalidReplicaAssignmentException when assignments have differently sized inner lists")
+      } catch {
+        case e: ExecutionException =>
+          assertTrue(desc, e.getCause.isInstanceOf[InvalidReplicaAssignmentException])
+          assertEquals(desc, "Inconsistent replication factor between partitions, partition 0 has 1 " +
+            "while partitions [4] have replication factors [2], respectively.", e.getCause.getMessage)
+          assertEquals(desc, 3, numPartitions(topic1))
+      }
+
+      // try assignments with unknown brokers
+      alterResult = client.createPartitions(Map(topic1 ->
+        NewPartitions.increaseTo(4, asList(asList(12)))).asJava, option)
+      try {
+        altered = alterResult.values.get(topic1).get
+        fail(s"$desc: Expect InvalidReplicaAssignmentException when assignments contains an unknown broker")
+      } catch {
+        case e: ExecutionException =>
+          assertTrue(desc, e.getCause.isInstanceOf[InvalidReplicaAssignmentException])
+          assertEquals(desc, "Unknown broker(s) in replica assignment: 12.", e.getCause.getMessage)
+          assertEquals(desc, 3, numPartitions(topic1))
+      }
+
+      // try with empty assignments
+      alterResult = client.createPartitions(Map(topic1 ->
+        NewPartitions.increaseTo(4, Collections.emptyList())).asJava, option)
+      try {
+        altered = alterResult.values.get(topic1).get
+        fail(s"$desc: Expect InvalidReplicaAssignmentException when assignments is empty")
+      } catch {
+        case e: ExecutionException =>
+          assertTrue(desc, e.getCause.isInstanceOf[InvalidReplicaAssignmentException])
+          assertEquals(desc, "Increasing the number of partitions by 1 but 0 assignments provided.", e.getCause.getMessage)
+          assertEquals(desc, 3, numPartitions(topic1))
+      }
+    }
+
+    // a mixed success, failure response
+    alterResult = client.createPartitions(Map(
+      topic1 -> NewPartitions.increaseTo(4),
+      topic2 -> NewPartitions.increaseTo(2)).asJava, actuallyDoIt)
+    // assert that the topic1 now has 4 partitions
+    altered = alterResult.values.get(topic1).get
+    TestUtils.waitUntilTrue(() => numPartitions(topic1) == 4, "Timed out waiting for new partitions to appear")
+    try {
+      altered = alterResult.values.get(topic2).get
+    } catch {
+      case e: ExecutionException =>
+        assertTrue(e.getCause.isInstanceOf[InvalidPartitionsException])
+        assertEquals("Topic currently has 3 partitions, which is higher than the requested 2.", e.getCause.getMessage)
+        // assert that the topic2 still has 3 partitions
+        assertEquals(3, numPartitions(topic2))
+    }
+
+    // finally, try to add partitions to a topic queued for deletion
+    val deleteResult = client.deleteTopics(asList(topic1))
+    deleteResult.values.get(topic1).get
+    alterResult = client.createPartitions(Map(topic1 ->
+      NewPartitions.increaseTo(4)).asJava, validateOnly)
+    try {
+      altered = alterResult.values.get(topic1).get
+      fail("Expect InvalidTopicException when the topic is queued for deletion")
+    } catch {
+      case e: ExecutionException =>
+        assertTrue(e.getCause.isInstanceOf[InvalidTopicException])
+        assertEquals("The topic is queued for deletion.", e.getCause.getMessage)
+    }
+  }
+
+  @Test
+  def testSeekAfterDeleteRecords(): Unit = {
+    createTopic(topic, numPartitions = 2, replicationFactor = brokerCount)
+
+    client = Admin.create(createConfig)
+
+    val consumer = createConsumer()
+    subscribeAndWaitForAssignment(topic, consumer)
+
+    val producer = createProducer()
+    sendRecords(producer, 10, topicPartition)
+    consumer.seekToBeginning(Collections.singleton(topicPartition))
+    assertEquals(0L, consumer.position(topicPartition))
+
+    val result = client.deleteRecords(Map(topicPartition -> RecordsToDelete.beforeOffset(5L)).asJava)
+    val lowWatermark = result.lowWatermarks().get(topicPartition).get.lowWatermark
+    assertEquals(5L, lowWatermark)
+
+    consumer.seekToBeginning(Collections.singletonList(topicPartition))
+    assertEquals(5L, consumer.position(topicPartition))
+
+    consumer.seek(topicPartition, 7L)
+    assertEquals(7L, consumer.position(topicPartition))
+
+    client.deleteRecords(Map(topicPartition -> RecordsToDelete.beforeOffset(DeleteRecordsRequest.HIGH_WATERMARK)).asJava).all.get
+    consumer.seekToBeginning(Collections.singletonList(topicPartition))
+    assertEquals(10L, consumer.position(topicPartition))
+  }
+
+  @Test
+  def testLogStartOffsetCheckpoint(): Unit = {
+    createTopic(topic, numPartitions = 2, replicationFactor = brokerCount)
+
+    client = Admin.create(createConfig)
+
+    val consumer = createConsumer()
+    subscribeAndWaitForAssignment(topic, consumer)
+
+    val producer = createProducer()
+    sendRecords(producer, 10, topicPartition)
+    var result = client.deleteRecords(Map(topicPartition -> RecordsToDelete.beforeOffset(5L)).asJava)
+    var lowWatermark: Option[Long] = Some(result.lowWatermarks.get(topicPartition).get.lowWatermark)
+    assertEquals(Some(5), lowWatermark)
+
+    for (i <- 0 until brokerCount) {
+      killBroker(i)
+    }
+    restartDeadBrokers()
+
+    client.close()
+    brokerList = TestUtils.bootstrapServers(servers, listenerName)
+    client = Admin.create(createConfig)
+
+    TestUtils.waitUntilTrue(() => {
+      // Need to retry if leader is not available for the partition
+      result = client.deleteRecords(Map(topicPartition -> RecordsToDelete.beforeOffset(0L)).asJava)
+
+      lowWatermark = None
+      val future = result.lowWatermarks().get(topicPartition)
+      try {
+        lowWatermark = Some(future.get.lowWatermark)
+        lowWatermark.contains(5L)
+      } catch {
+        case e: ExecutionException if e.getCause.isInstanceOf[LeaderNotAvailableException] ||
+          e.getCause.isInstanceOf[NotLeaderOrFollowerException] => false
+        }
+    }, s"Expected low watermark of the partition to be 5 but got ${lowWatermark.getOrElse("no response within the timeout")}")
+  }
+
+  @Test
+  def testLogStartOffsetAfterDeleteRecords(): Unit = {
+    createTopic(topic, numPartitions = 2, replicationFactor = brokerCount)
+
+    client = Admin.create(createConfig)
+
+    val consumer = createConsumer()
+    subscribeAndWaitForAssignment(topic, consumer)
+
+    val producer = createProducer()
+    sendRecords(producer, 10, topicPartition)
+
+    val result = client.deleteRecords(Map(topicPartition -> RecordsToDelete.beforeOffset(3L)).asJava)
+    val lowWatermark = result.lowWatermarks.get(topicPartition).get.lowWatermark
+    assertEquals(3L, lowWatermark)
+
+    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 = 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.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.localLog(topicPartition).get.logStartOffset == expectedStartOffset
+      }, s"Expected follower to discover new log start offset $expectedStartOffset")
+
+      TestUtils.waitUntilTrue(() => {
+        servers(followerIndex).replicaManager.localLog(topicPartition).get.logEndOffset == expectedEndOffset
+      }, s"Expected follower to catch up to log end offset $expectedEndOffset")
+    }
+
+    // we will produce to topic and delete records while one follower is down
+    killBroker(followerIndex)
+
+    client = Admin.create(createConfig)
+    val producer = createProducer()
+    sendRecords(producer, 100, topicPartition)
+
+    val result = client.deleteRecords(Map(topicPartition -> RecordsToDelete.beforeOffset(3L)).asJava)
+    result.all().get()
+
+    // start the stopped broker to verify that it will be able to fetch from new log start offset
+    restartDeadBrokers()
+
+    waitForFollowerLog(expectedStartOffset=3L, expectedEndOffset=100L)
+
+    // after the new replica caught up, all replicas should have same log start offset
+    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)
+    sendRecords(producer, 100, topicPartition)
+    val result1 = client.deleteRecords(Map(topicPartition -> RecordsToDelete.beforeOffset(117L)).asJava)
+    result1.all().get()
+    restartDeadBrokers()
+    waitForFollowerLog(expectedStartOffset=117L, expectedEndOffset=200L)
+  }
+
+  @Test
+  def testAlterLogDirsAfterDeleteRecords(): Unit = {
+    client = Admin.create(createConfig)
+    createTopic(topic, numPartitions = 1, replicationFactor = brokerCount)
+    val expectedLEO = 100
+    val producer = createProducer()
+    sendRecords(producer, expectedLEO, topicPartition)
+
+    // delete records to move log start offset
+    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 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
+    val futureLogDir = servers(0).config.logDirs(1)
+    val futureReplica = new TopicPartitionReplica(topic, 0, servers(0).config.brokerId)
+
+    // Verify that replica can be moved to the specified log directory
+    client.alterReplicaLogDirs(Map(futureReplica -> futureLogDir).asJava).all.get
+    TestUtils.waitUntilTrue(() => {
+      futureLogDir == servers(0).logManager.getLog(topicPartition).get.dir.getParent
+    }, "timed out waiting for replica movement")
+
+    // once replica moved, its LSO and LEO should match other replicas
+    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 = brokerCount)
+
+    client = Admin.create(createConfig)
+
+    val consumer = createConsumer()
+    subscribeAndWaitForAssignment(topic, consumer)
+
+    val producer = createProducer()
+    sendRecords(producer, 10, topicPartition)
+    assertEquals(0L, consumer.offsetsForTimes(Map(topicPartition -> JLong.valueOf(0L)).asJava).get(topicPartition).offset())
+
+    var result = client.deleteRecords(Map(topicPartition -> RecordsToDelete.beforeOffset(5L)).asJava)
+    result.all.get
+    assertEquals(5L, consumer.offsetsForTimes(Map(topicPartition -> JLong.valueOf(0L)).asJava).get(topicPartition).offset())
+
+    result = client.deleteRecords(Map(topicPartition -> RecordsToDelete.beforeOffset(DeleteRecordsRequest.HIGH_WATERMARK)).asJava)
+    result.all.get
+    assertNull(consumer.offsetsForTimes(Map(topicPartition -> JLong.valueOf(0L)).asJava).get(topicPartition))
+  }
+
+  @Test
+  def testConsumeAfterDeleteRecords(): Unit = {
+    val consumer = createConsumer()
+    subscribeAndWaitForAssignment(topic, consumer)
+
+    client = Admin.create(createConfig)
+
+    val producer = createProducer()
+    sendRecords(producer, 10, topicPartition)
+    var messageCount = 0
+    TestUtils.consumeRecords(consumer, 10)
+
+    client.deleteRecords(Map(topicPartition -> RecordsToDelete.beforeOffset(3L)).asJava).all.get
+    consumer.seek(topicPartition, 1)
+    messageCount = 0
+    TestUtils.consumeRecords(consumer, 7)
+
+    client.deleteRecords(Map(topicPartition -> RecordsToDelete.beforeOffset(8L)).asJava).all.get
+    consumer.seek(topicPartition, 1)
+    messageCount = 0
+    TestUtils.consumeRecords(consumer, 2)
+  }
+
+  @Test
+  def testDeleteRecordsWithException(): Unit = {
+    val consumer = createConsumer()
+    subscribeAndWaitForAssignment(topic, consumer)
+
+    client = Admin.create(createConfig)
+
+    val producer = createProducer()
+    sendRecords(producer, 10, topicPartition)
+
+    assertEquals(5L, client.deleteRecords(Map(topicPartition -> RecordsToDelete.beforeOffset(5L)).asJava)
+      .lowWatermarks.get(topicPartition).get.lowWatermark)
+
+    // OffsetOutOfRangeException if offset > high_watermark
+    var cause = intercept[ExecutionException] {
+      client.deleteRecords(Map(topicPartition -> RecordsToDelete.beforeOffset(20L)).asJava).lowWatermarks.get(topicPartition).get
+    }.getCause
+    assertEquals(classOf[OffsetOutOfRangeException], cause.getClass)
+
+    val nonExistPartition = new TopicPartition(topic, 3)
+    // LeaderNotAvailableException if non existent partition
+    cause = intercept[ExecutionException] {
+      client.deleteRecords(Map(nonExistPartition -> RecordsToDelete.beforeOffset(20L)).asJava).lowWatermarks.get(nonExistPartition).get
+    }.getCause
+    assertEquals(classOf[LeaderNotAvailableException], cause.getClass)
+  }
+
+  @Test
+  def testDescribeConfigsForTopic(): Unit = {
+    createTopic(topic, numPartitions = 2, replicationFactor = brokerCount)
+    client = Admin.create(createConfig)
+
+    val existingTopic = new ConfigResource(ConfigResource.Type.TOPIC, topic)
+    client.describeConfigs(Collections.singletonList(existingTopic)).values.get(existingTopic).get()
+
+    val nonExistentTopic = new ConfigResource(ConfigResource.Type.TOPIC, "unknown")
+    val describeResult1 = client.describeConfigs(Collections.singletonList(nonExistentTopic))
+
+    assertTrue(intercept[ExecutionException](describeResult1.values.get(nonExistentTopic).get).getCause.isInstanceOf[UnknownTopicOrPartitionException])
+
+    val invalidTopic = new ConfigResource(ConfigResource.Type.TOPIC, "(invalid topic)")
+    val describeResult2 = client.describeConfigs(Collections.singletonList(invalidTopic))
+
+    assertTrue(intercept[ExecutionException](describeResult2.values.get(invalidTopic).get).getCause.isInstanceOf[InvalidTopicException])
+  }
+
+  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")
+  }
+
+  private def sendRecords(producer: KafkaProducer[Array[Byte], Array[Byte]],
+                          numRecords: Int,
+                          topicPartition: TopicPartition): Unit = {
+    val futures = (0 until numRecords).map( i => {
+      val record = new ProducerRecord(topicPartition.topic, topicPartition.partition, s"$i".getBytes, s"$i".getBytes)
+      debug(s"Sending this record: $record")
+      producer.send(record)
+    })
+
+    futures.foreach(_.get)
+  }
+
+  @Test
+  def testInvalidAlterConfigs(): Unit = {
+    client = Admin.create(createConfig)
+    checkInvalidAlterConfigs(zkClient, servers, client)
+  }
+
+  /**
+   * Test that ACL operations are not possible when the authorizer is disabled.
+   * Also see [[kafka.api.SaslSslAdminIntegrationTest.testAclOperations()]] for tests of ACL operations
+   * when the authorizer is enabled.
+   */
+  @Test
+  def testAclOperations(): Unit = {
+    val acl = new AclBinding(new ResourcePattern(ResourceType.TOPIC, "mytopic3", PatternType.LITERAL),
+      new AccessControlEntry("User:ANONYMOUS", "*", AclOperation.DESCRIBE, AclPermissionType.ALLOW))
+    client = Admin.create(createConfig)
+    assertFutureExceptionTypeEquals(client.describeAcls(AclBindingFilter.ANY).values(), classOf[SecurityDisabledException])
+    assertFutureExceptionTypeEquals(client.createAcls(Collections.singleton(acl)).all(),
+      classOf[SecurityDisabledException])
+    assertFutureExceptionTypeEquals(client.deleteAcls(Collections.singleton(acl.toFilter())).all(),
+      classOf[SecurityDisabledException])
+  }
+
+  /**
+    * Test closing the AdminClient with a generous timeout.  Calls in progress should be completed,
+    * since they can be done within the timeout.  New calls should receive timeouts.
+    */
+  @Test
+  def testDelayedClose(): Unit = {
+    client = Admin.create(createConfig)
+    val topics = Seq("mytopic", "mytopic2")
+    val newTopics = topics.map(new NewTopic(_, 1, 1.toShort))
+    val future = client.createTopics(newTopics.asJava, new CreateTopicsOptions().validateOnly(true)).all()
+    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(time.Duration.ofMinutes(30)) // multiple close-with-timeout should have no effect
+  }
+
+  /**
+    * Test closing the AdminClient with a timeout of 0, when there are calls with extremely long
+    * timeouts in progress.  The calls should be aborted after the hard shutdown timeout elapses.
+    */
+  @Test
+  def testForceClose(): Unit = {
+    val config = createConfig
+    config.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, s"localhost:${TestUtils.IncorrectBrokerPort}")
+    client = Admin.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.toShort)).asJava,
+      new CreateTopicsOptions().timeoutMs(900000)).all()
+    client.close(time.Duration.ZERO)
+    assertFutureExceptionTypeEquals(future, classOf[TimeoutException])
+  }
+
+  /**
+    * Check that a call with a timeout does not complete before the minimum timeout has elapsed,
+    * even when the default request timeout is shorter.
+    */
+  @Test
+  def testMinimumRequestTimeouts(): Unit = {
+    val config = createConfig
+    config.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, s"localhost:${TestUtils.IncorrectBrokerPort}")
+    config.put(AdminClientConfig.REQUEST_TIMEOUT_MS_CONFIG, "0")
+    client = Admin.create(config)
+    val startTimeMs = Time.SYSTEM.milliseconds()
+    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()
+    assertTrue("Expected the timeout to take at least one millisecond.", endTimeMs > startTimeMs)
+  }
+
+  /**
+    * Test injecting timeouts for calls that are in flight.
+    */
+  @Test
+  def testCallInFlightTimeouts(): Unit = {
+    val config = createConfig
+    config.put(AdminClientConfig.DEFAULT_API_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.toShort)).asJava,
+        new CreateTopicsOptions().validateOnly(true)).all()
+    assertFutureExceptionTypeEquals(future, classOf[TimeoutException])
+    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)
+  }
+
+  /**
+   * Test the consumer group APIs.
+   */
+  @Test
+  def testConsumerGroups(): Unit = {
+    val config = createConfig
+    client = Admin.create(config)
+    try {
+      // Verify that initially there are no consumer groups to list.
+      val list1 = client.listConsumerGroups()
+      assertTrue(0 == list1.all().get().size())
+      assertTrue(0 == list1.errors().get().size())
+      assertTrue(0 == list1.valid().get().size())
+      val testTopicName = "test_topic"
+      val testTopicName1 = testTopicName + "1"
+      val testTopicName2 = testTopicName + "2"
+      val testNumPartitions = 2
+
+      client.createTopics(util.Arrays.asList(
+        new NewTopic(testTopicName, testNumPartitions, 1.toShort),
+        new NewTopic(testTopicName1, testNumPartitions, 1.toShort),
+        new NewTopic(testTopicName2, testNumPartitions, 1.toShort)
+      )).all().get()
+      waitForTopics(client, List(testTopicName, testTopicName1, testTopicName2), List())
+
+      val producer = createProducer()
+      try {
+        producer.send(new ProducerRecord(testTopicName, 0, null, null)).get()
+      } finally {
+        Utils.closeQuietly(producer, "producer")
+      }
+
+      val EMPTY_GROUP_INSTANCE_ID = ""
+      val testGroupId = "test_group_id"
+      val testClientId = "test_client_id"
+      val testInstanceId1 = "test_instance_id_1"
+      val testInstanceId2 = "test_instance_id_2"
+      val fakeGroupId = "fake_group_id"
+
+      def createProperties(groupInstanceId: String): Properties = {
+        val newConsumerConfig = new Properties(consumerConfig)
+        // We need to disable the auto commit because after the members got removed from group, the offset commit
+        // will cause the member rejoining and the test will be flaky (check ConsumerCoordinator#OffsetCommitResponseHandler)
+        newConsumerConfig.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false")
+        newConsumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, testGroupId)
+        newConsumerConfig.setProperty(ConsumerConfig.CLIENT_ID_CONFIG, testClientId)
+        if (groupInstanceId != EMPTY_GROUP_INSTANCE_ID) {
+          newConsumerConfig.setProperty(ConsumerConfig.GROUP_INSTANCE_ID_CONFIG, groupInstanceId)
+        }
+        newConsumerConfig
+      }
+
+      // contains two static members and one dynamic member
+      val groupInstanceSet = Set(testInstanceId1, testInstanceId2, EMPTY_GROUP_INSTANCE_ID)
+      val consumerSet = groupInstanceSet.map { groupInstanceId => createConsumer(configOverrides = createProperties(groupInstanceId))}
+      val topicSet = Set(testTopicName, testTopicName1, testTopicName2)
+
+      val latch = new CountDownLatch(consumerSet.size)
+      try {
+        def createConsumerThread[K,V](consumer: KafkaConsumer[K,V], topic: String): Thread = {
+          new Thread {
+            override def run : Unit = {
+              consumer.subscribe(Collections.singleton(topic))
+              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
+              }
+            }
+          }
+        }
+
+        // Start consumers in a thread that will subscribe to a new group.
+        val consumerThreads = consumerSet.zip(topicSet).map(zipped => createConsumerThread(zipped._1, zipped._2))
+
+        try {
+          consumerThreads.foreach(_.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(group =>
+                group.groupId == testGroupId &&
+                group.state.get == ConsumerGroupState.STABLE)
+            matching.size == 1
+          }, s"Expected to be able to list $testGroupId")
+
+          TestUtils.waitUntilTrue(() => {
+            val options = new ListConsumerGroupsOptions().inStates(Set(ConsumerGroupState.STABLE).asJava)
+            val matching = client.listConsumerGroups(options).all.get.asScala.filter(group =>
+                group.groupId == testGroupId &&
+                group.state.get == ConsumerGroupState.STABLE)
+            matching.size == 1
+          }, s"Expected to be able to list $testGroupId in state Stable")
+
+          TestUtils.waitUntilTrue(() => {
+            val options = new ListConsumerGroupsOptions().inStates(Set(ConsumerGroupState.EMPTY).asJava)
+            val matching = client.listConsumerGroups(options).all.get.asScala.filter(
+                _.groupId == testGroupId)
+            matching.isEmpty
+          }, s"Expected to find zero groups")
+
+          val describeWithFakeGroupResult = client.describeConsumerGroups(Seq(testGroupId, fakeGroupId).asJava,
+            new DescribeConsumerGroupsOptions().includeAuthorizedOperations(true))
+          assertEquals(2, describeWithFakeGroupResult.describedGroups().size())
+
+          // Test that we can get information about the test consumer group.
+          assertTrue(describeWithFakeGroupResult.describedGroups().containsKey(testGroupId))
+          var testGroupDescription = describeWithFakeGroupResult.describedGroups().get(testGroupId).get()
+
+          assertEquals(testGroupId, testGroupDescription.groupId())
+          assertFalse(testGroupDescription.isSimpleConsumerGroup)
+          assertEquals(groupInstanceSet.size, testGroupDescription.members().size())
+          val members = testGroupDescription.members()
+          members.asScala.foreach(member => assertEquals(testClientId, member.clientId()))
+          val topicPartitionsByTopic = members.asScala.flatMap(_.assignment().topicPartitions().asScala).groupBy(_.topic())
+          topicSet.foreach { topic =>
+            val topicPartitions = topicPartitionsByTopic.getOrElse(topic, List.empty)
+            assertEquals(testNumPartitions, topicPartitions.size)
+          }
+
+          val expectedOperations = AclEntry.supportedOperations(ResourceType.GROUP).asJava
+          assertEquals(expectedOperations, testGroupDescription.authorizedOperations())
+
+          // Test that the fake group is listed as dead.
+          assertTrue(describeWithFakeGroupResult.describedGroups().containsKey(fakeGroupId))
+          val fakeGroupDescription = describeWithFakeGroupResult.describedGroups().get(fakeGroupId).get()
+
+          assertEquals(fakeGroupId, fakeGroupDescription.groupId())
+          assertEquals(0, fakeGroupDescription.members().size())
+          assertEquals("", fakeGroupDescription.partitionAssignor())
+          assertEquals(ConsumerGroupState.DEAD, fakeGroupDescription.state())
+          assertEquals(expectedOperations, fakeGroupDescription.authorizedOperations())
+
+          // Test that all() returns 2 results
+          assertEquals(2, describeWithFakeGroupResult.all().get().size())
+
+          // Test listConsumerGroupOffsets
+          TestUtils.waitUntilTrue(() => {
+            val parts = client.listConsumerGroupOffsets(testGroupId).partitionsToOffsetAndMetadata().get()
+            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
+          var deleteResult = client.deleteConsumerGroups(Seq(testGroupId, fakeGroupId).asJava)
+          assertEquals(2, deleteResult.deletedGroups().size())
+
+          // Deleting the fake group ID should get GroupIdNotFoundException.
+          assertTrue(deleteResult.deletedGroups().containsKey(fakeGroupId))
+          assertFutureExceptionTypeEquals(deleteResult.deletedGroups().get(fakeGroupId),
+            classOf[GroupIdNotFoundException])
+
+          // Deleting the real group ID should get GroupNotEmptyException
+          assertTrue(deleteResult.deletedGroups().containsKey(testGroupId))
+          assertFutureExceptionTypeEquals(deleteResult.deletedGroups().get(testGroupId),
+            classOf[GroupNotEmptyException])
+
+          // Test delete one correct static member
+          removeMembersResult = client.removeMembersFromConsumerGroup(testGroupId, new RemoveMembersFromConsumerGroupOptions(
+            Collections.singleton(new MemberToRemove(testInstanceId1))
+          ))
+
+          assertNull(removeMembersResult.all().get())
+          val validMemberFuture = removeMembersResult.memberResult(new MemberToRemove(testInstanceId1))
+          assertNull(validMemberFuture.get())
+
+          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)
+          assertEquals(consumerSet.size - 1, testGroupDescription.members().size())
+
+          // Delete all active members remaining (a static member + a dynamic member)
+          removeMembersResult = client.removeMembersFromConsumerGroup(testGroupId, new RemoveMembersFromConsumerGroupOptions())
+          assertNull(removeMembersResult.all().get())
+
+          // The group should contain no members now.
+          testGroupDescription = client.describeConsumerGroups(Seq(testGroupId).asJava,
+            new DescribeConsumerGroupsOptions().includeAuthorizedOperations(true))
+              .describedGroups().get(testGroupId).get()
+          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 {
+        consumerThreads.foreach {
+          case consumerThread =>
+            consumerThread.interrupt()
+            consumerThread.join()
+        }
+      }
+      } finally {
+        consumerSet.zip(groupInstanceSet).foreach(zipped => Utils.closeQuietly(zipped._1, zipped._2))
+      }
+    } finally {
+      Utils.closeQuietly(client, "adminClient")
+    }
+  }
+
+  @Test
+  def testDeleteConsumerGroupOffsets(): Unit = {
+    val config = createConfig
+    client = Admin.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 = Admin.create(createConfig)
+
+    val prefer0 = Seq(0, 1, 2)
+    val prefer1 = Seq(1, 2, 0)
+    val prefer2 = Seq(2, 0, 1)
+
+    val partition1 = new TopicPartition("elect-preferred-leaders-topic-1", 0)
+    TestUtils.createTopic(zkClient, partition1.topic, Map[Int, Seq[Int]](partition1.partition -> prefer0), servers)
+
+    val partition2 = new TopicPartition("elect-preferred-leaders-topic-2", 0)
+    TestUtils.createTopic(zkClient, partition2.topic, Map[Int, Seq[Int]](partition2.partition -> prefer0), servers)
+
+    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 = zkClient.getLeaderForPartition(partition1).get
+      val prior2 = zkClient.getLeaderForPartition(partition2).get
+
+      var m = Map.empty[TopicPartition, Seq[Int]]
+
+      if (prior1 != preferred)
+        m += partition1 -> newAssignment
+      if (prior2 != preferred)
+        m += partition2 -> newAssignment
+
+      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)
+      // Check the leader hasn't moved
+      TestUtils.assertLeader(client, partition1, prior1)
+      TestUtils.assertLeader(client, partition2, prior2)
+    }
+
+    // Check current leaders are 0
+    TestUtils.assertLeader(client, partition1, 0)
+    TestUtils.assertLeader(client, partition2, 0)
+
+    // Noop election
+    var electResult = client.electLeaders(ElectionType.PREFERRED, Set(partition1).asJava)
+    var exception = electResult.partitions.get.get(partition1).get
+    assertEquals(classOf[ElectionNotNeededException], exception.getClass)
+    TestUtils.assertLeader(client, partition1, 0)
+
+    // Noop election with null partitions
+    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.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
+    assertFalse(electResult.partitions.get.containsKey(partition2))
+    TestUtils.assertLeader(client, partition2, 0)
+
+    // meaningful election with null partitions
+    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.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.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()
+    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 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"))
+    TestUtils.assertLeader(client, partition1, 2)
+
+    // preferred leader unavailable with null argument
+    electResult = client.electLeaders(ElectionType.PREFERRED, null, shortTimeout)
+
+    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"))
+
+    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"))
+
+    TestUtils.assertLeader(client, partition1, 2)
+    TestUtils.assertLeader(client, partition2, 2)
+  }
+
+  @Test
+  def testElectUncleanLeadersForOnePartition(): Unit = {
+    // Case: unclean leader election with one topic partition
+    client = Admin.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 = Admin.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 = Admin.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 = Admin.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 = Admin.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 = Admin.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 = Admin.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 = Admin.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 = Admin.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 = Admin.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
+
+    // Test SET and APPEND on previously unset properties
+    var topic2AlterConfigs = Seq(
+      new AlterConfigOp(new ConfigEntry(LogConfig.MinCleanableDirtyRatioProp, "0.9"), AlterConfigOp.OpType.SET),
+      new AlterConfigOp(new ConfigEntry(LogConfig.CompressionTypeProp, "lz4"), AlterConfigOp.OpType.SET),
+      new AlterConfigOp(new ConfigEntry(LogConfig.CleanupPolicyProp, LogConfig.Compact), AlterConfigOp.OpType.APPEND)
+    ).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)
+    assertEquals("delete,compact", configs.get(topic2Resource).get(LogConfig.CleanupPolicyProp).value)
+
+    //verify subtract operation, including from an empty property
+    topic1AlterConfigs = Seq(
+      new AlterConfigOp(new ConfigEntry(LogConfig.CleanupPolicyProp, LogConfig.Compact), AlterConfigOp.OpType.SUBTRACT),
+      new AlterConfigOp(new ConfigEntry(LogConfig.LeaderReplicationThrottledReplicasProp, "0"), AlterConfigOp.OpType.SUBTRACT)
+    ).asJava
+
+    // subtract all from this list property
+    topic2AlterConfigs = Seq(
+      new AlterConfigOp(new ConfigEntry(LogConfig.CleanupPolicyProp, LogConfig.Compact + "," + LogConfig.Delete), AlterConfigOp.OpType.SUBTRACT)
+    ).asJavaCollection
+
+    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
+    describeResult = client.describeConfigs(Seq(topic1Resource, topic2Resource).asJava)
+    configs = describeResult.all.get
+
+    assertEquals(2, configs.size)
+
+    assertEquals("delete", configs.get(topic1Resource).get(LogConfig.CleanupPolicyProp).value)
+    assertEquals("1000", configs.get(topic1Resource).get(LogConfig.FlushMsProp).value) // verify previous change is still intact
+    assertEquals("", configs.get(topic1Resource).get(LogConfig.LeaderReplicationThrottledReplicasProp).value)
+    assertEquals("", configs.get(topic2Resource).get(LogConfig.CleanupPolicyProp).value )
+
+    // 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 testIncrementalAlterConfigsDeleteAndSetBrokerConfigs(): Unit = {
+    client = Admin.create(createConfig)
+    val broker0Resource = new ConfigResource(ConfigResource.Type.BROKER, "0")
+    client.incrementalAlterConfigs(Map(broker0Resource ->
+      Seq(new AlterConfigOp(new ConfigEntry(DynamicConfig.Broker.LeaderReplicationThrottledRateProp, "123"),
+          AlterConfigOp.OpType.SET),
+        new AlterConfigOp(new ConfigEntry(DynamicConfig.Broker.FollowerReplicationThrottledRateProp, "456"),
+          AlterConfigOp.OpType.SET)
+      ).asJavaCollection).asJava).all().get()
+    TestUtils.waitUntilTrue(() => {
+      val broker0Configs = client.describeConfigs(Seq(broker0Resource).asJava).
+        all().get().get(broker0Resource).entries().asScala.map {
+        case entry => (entry.name, entry.value)
+      }.toMap
+      ("123".equals(broker0Configs.getOrElse(DynamicConfig.Broker.LeaderReplicationThrottledRateProp, "")) &&
+        "456".equals(broker0Configs.getOrElse(DynamicConfig.Broker.FollowerReplicationThrottledRateProp, "")))
+    }, "Expected to see the broker properties we just set", pause=25)
+    client.incrementalAlterConfigs(Map(broker0Resource ->
+      Seq(new AlterConfigOp(new ConfigEntry(DynamicConfig.Broker.LeaderReplicationThrottledRateProp, ""),
+        AlterConfigOp.OpType.DELETE),
+        new AlterConfigOp(new ConfigEntry(DynamicConfig.Broker.FollowerReplicationThrottledRateProp, "654"),
+          AlterConfigOp.OpType.SET),
+        new AlterConfigOp(new ConfigEntry(DynamicConfig.Broker.ReplicaAlterLogDirsIoMaxBytesPerSecondProp, "987"),
+          AlterConfigOp.OpType.SET)
+      ).asJavaCollection).asJava).all().get()
+    TestUtils.waitUntilTrue(() => {
+      val broker0Configs = client.describeConfigs(Seq(broker0Resource).asJava).
+        all().get().get(broker0Resource).entries().asScala.map {
+        case entry => (entry.name, entry.value)
+      }.toMap
+      ("".equals(broker0Configs.getOrElse(DynamicConfig.Broker.LeaderReplicationThrottledRateProp, "")) &&
+        "654".equals(broker0Configs.getOrElse(DynamicConfig.Broker.FollowerReplicationThrottledRateProp, "")) &&
+        "987".equals(broker0Configs.getOrElse(DynamicConfig.Broker.ReplicaAlterLogDirsIoMaxBytesPerSecondProp, "")))
+    }, "Expected to see the broker properties we just modified", pause=25)
+  }
+
+  @Test
+  def testIncrementalAlterConfigsDeleteBrokerConfigs(): Unit = {
+    client = Admin.create(createConfig)
+    val broker0Resource = new ConfigResource(ConfigResource.Type.BROKER, "0")
+    client.incrementalAlterConfigs(Map(broker0Resource ->
+      Seq(new AlterConfigOp(new ConfigEntry(DynamicConfig.Broker.LeaderReplicationThrottledRateProp, "123"),
+        AlterConfigOp.OpType.SET),
+        new AlterConfigOp(new ConfigEntry(DynamicConfig.Broker.FollowerReplicationThrottledRateProp, "456"),
+          AlterConfigOp.OpType.SET),
+        new AlterConfigOp(new ConfigEntry(DynamicConfig.Broker.ReplicaAlterLogDirsIoMaxBytesPerSecondProp, "789"),
+          AlterConfigOp.OpType.SET)
+      ).asJavaCollection).asJava).all().get()
+    TestUtils.waitUntilTrue(() => {
+      val broker0Configs = client.describeConfigs(Seq(broker0Resource).asJava).
+        all().get().get(broker0Resource).entries().asScala.map {
+        case entry => (entry.name, entry.value)
+      }.toMap
+      ("123".equals(broker0Configs.getOrElse(DynamicConfig.Broker.LeaderReplicationThrottledRateProp, "")) &&
+        "456".equals(broker0Configs.getOrElse(DynamicConfig.Broker.FollowerReplicationThrottledRateProp, "")) &&
+        "789".equals(broker0Configs.getOrElse(DynamicConfig.Broker.ReplicaAlterLogDirsIoMaxBytesPerSecondProp, "")))
+    }, "Expected to see the broker properties we just set", pause=25)
+    client.incrementalAlterConfigs(Map(broker0Resource ->
+      Seq(new AlterConfigOp(new ConfigEntry(DynamicConfig.Broker.LeaderReplicationThrottledRateProp, ""),
+        AlterConfigOp.OpType.DELETE),
+        new AlterConfigOp(new ConfigEntry(DynamicConfig.Broker.FollowerReplicationThrottledRateProp, ""),
+          AlterConfigOp.OpType.DELETE),
+        new AlterConfigOp(new ConfigEntry(DynamicConfig.Broker.ReplicaAlterLogDirsIoMaxBytesPerSecondProp, ""),
+          AlterConfigOp.OpType.DELETE)
+      ).asJavaCollection).asJava).all().get()
+    TestUtils.waitUntilTrue(() => {
+      val broker0Configs = client.describeConfigs(Seq(broker0Resource).asJava).
+        all().get().get(broker0Resource).entries().asScala.map {
+        case entry => (entry.name, entry.value)
+      }.toMap
+      ("".equals(broker0Configs.getOrElse(DynamicConfig.Broker.LeaderReplicationThrottledRateProp, "")) &&
+        "".equals(broker0Configs.getOrElse(DynamicConfig.Broker.FollowerReplicationThrottledRateProp, "")) &&
+        "".equals(broker0Configs.getOrElse(DynamicConfig.Broker.ReplicaAlterLogDirsIoMaxBytesPerSecondProp, "")))
+    }, "Expected to see the broker properties we just removed to be deleted", pause=25)
+  }
+
+  @Test
+  def testInvalidIncrementalAlterConfigs(): Unit = {
+    client = Admin.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 = Admin.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 = Admin.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()
+  }
+
+  // Verify that createTopics and alterConfigs fail with null values
+  @Test
+  def testNullConfigs(): Unit = {
+
+    def validateLogConfig(compressionType: String): Unit = {
+      val logConfig = zkClient.getLogConfigs(Set(topic), Collections.emptyMap[String, AnyRef])._1(topic)
+
+      assertEquals(compressionType, logConfig.originals.get(LogConfig.CompressionTypeProp))
+      assertNull(logConfig.originals.get(LogConfig.MessageFormatVersionProp))
+      assertEquals(ApiVersion.latestVersion, logConfig.messageFormatVersion)
+    }
+
+    client = Admin.create(createConfig)
+    val invalidConfigs = Map[String, String](LogConfig.MessageFormatVersionProp -> null,
+      LogConfig.CompressionTypeProp -> "producer").asJava
+    val newTopic = new NewTopic(topic, 2, brokerCount.toShort)
+    val e1 = intercept[ExecutionException] {
+      client.createTopics(Collections.singletonList(newTopic.configs(invalidConfigs))).all.get()
+    }
+    assertTrue(s"Unexpected exception ${e1.getCause.getClass}", e1.getCause.isInstanceOf[InvalidRequestException])
+
+    val validConfigs = Map[String, String](LogConfig.CompressionTypeProp -> "producer").asJava
+    client.createTopics(Collections.singletonList(newTopic.configs(validConfigs))).all.get()
+    waitForTopics(client, expectedPresent = Seq(topic), expectedMissing = List())
+    validateLogConfig(compressionType = "producer")
+
+    val topicResource = new ConfigResource(ConfigResource.Type.TOPIC, topic)
+    val alterOps = Seq(
+      new AlterConfigOp(new ConfigEntry(LogConfig.MessageFormatVersionProp, null), AlterConfigOp.OpType.SET),
+      new AlterConfigOp(new ConfigEntry(LogConfig.CompressionTypeProp, "lz4"), AlterConfigOp.OpType.SET)
+    )
+    val e2 = intercept[ExecutionException] {
+      client.incrementalAlterConfigs(Map(topicResource -> alterOps.asJavaCollection).asJava).all.get
+    }
+    assertTrue(s"Unexpected exception ${e2.getCause.getClass}", e2.getCause.isInstanceOf[InvalidRequestException])
+    validateLogConfig(compressionType = "producer")
+  }
+
+  @Test
+  def testDescribeConfigsForLog4jLogLevels(): Unit = {
+    client = Admin.create(createConfig)
+    LoggerFactory.getLogger("kafka.cluster.Replica").trace("Message to create the logger")
+    val loggerConfig = describeBrokerLoggers()
+    val kafkaLogLevel = loggerConfig.get("kafka").value()
+    val logCleanerLogLevelConfig = loggerConfig.get("kafka.cluster.Replica")
+    // we expect the log level to be inherited from the first ancestor with a level configured
+    assertEquals(kafkaLogLevel, logCleanerLogLevelConfig.value())
+    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 = Admin.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 = Admin.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 = Admin.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 = Admin.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
+    */
+  @nowarn("cat=deprecation")
+  @Test
+  @Ignore // To be re-enabled once KAFKA-8779 is resolved
+  def testAlterConfigsForLog4jLogLevelsDoesNotWork(): Unit = {
+    client = Admin.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.
+   */
+  private 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 PlaintextAdminIntegrationTest {
+
+  @nowarn("cat=deprecation")
+  def checkValidAlterConfigs(client: Admin, topicResource1: ConfigResource, topicResource2: ConfigResource): Unit = {
+    // Alter topics
+    var topicConfigEntries1 = Seq(
+      new ConfigEntry(LogConfig.FlushMsProp, "1000")
+    ).asJava
+
+    var topicConfigEntries2 = Seq(
+      new ConfigEntry(LogConfig.MinCleanableDirtyRatioProp, "0.9"),
+      new ConfigEntry(LogConfig.CompressionTypeProp, "lz4")
+    ).asJava
+
+    var alterResult = client.alterConfigs(Map(
+      topicResource1 -> new Config(topicConfigEntries1),
+      topicResource2 -> new Config(topicConfigEntries2)
+    ).asJava)
+
+    assertEquals(Set(topicResource1, topicResource2).asJava, alterResult.values.keySet)
+    alterResult.all.get
+
+    // Verify that topics were updated correctly
+    var describeResult = client.describeConfigs(Seq(topicResource1, topicResource2).asJava)
+    var configs = describeResult.all.get
+
+    assertEquals(2, configs.size)
+
+    assertEquals("1000", configs.get(topicResource1).get(LogConfig.FlushMsProp).value)
+    assertEquals(Defaults.MessageMaxBytes.toString,
+      configs.get(topicResource1).get(LogConfig.MaxMessageBytesProp).value)
+    assertEquals((Defaults.LogRetentionHours * 60 * 60 * 1000).toString,
+      configs.get(topicResource1).get(LogConfig.RetentionMsProp).value)
+
+    assertEquals("0.9", configs.get(topicResource2).get(LogConfig.MinCleanableDirtyRatioProp).value)
+    assertEquals("lz4", configs.get(topicResource2).get(LogConfig.CompressionTypeProp).value)
+
+    // Alter topics with validateOnly=true
+    topicConfigEntries1 = Seq(
+      new ConfigEntry(LogConfig.MaxMessageBytesProp, "10")
+    ).asJava
+
+    topicConfigEntries2 = Seq(
+      new ConfigEntry(LogConfig.MinCleanableDirtyRatioProp, "0.3")
+    ).asJava
+
+    alterResult = client.alterConfigs(Map(
+      topicResource1 -> new Config(topicConfigEntries1),
+      topicResource2 -> new Config(topicConfigEntries2)
+    ).asJava, new AlterConfigsOptions().validateOnly(true))
+
+    assertEquals(Set(topicResource1, topicResource2).asJava, alterResult.values.keySet)
+    alterResult.all.get
+
+    // Verify that topics were not updated due to validateOnly = true
+    describeResult = client.describeConfigs(Seq(topicResource1, topicResource2).asJava)
+    configs = describeResult.all.get
+
+    assertEquals(2, configs.size)
+
+    assertEquals(Defaults.MessageMaxBytes.toString,
+      configs.get(topicResource1).get(LogConfig.MaxMessageBytesProp).value)
+    assertEquals("0.9", configs.get(topicResource2).get(LogConfig.MinCleanableDirtyRatioProp).value)
+  }
+
+  @nowarn("cat=deprecation")
+  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)
+    TestUtils.createTopic(zkClient, topic1, 1, 1, servers)
+
+    val topic2 = "invalid-alter-configs-topic-2"
+    val topicResource2 = new ConfigResource(ConfigResource.Type.TOPIC, topic2)
+    TestUtils.createTopic(zkClient, topic2, 1, 1, servers)
+
+    val topicConfigEntries1 = Seq(
+      new ConfigEntry(LogConfig.MinCleanableDirtyRatioProp, "1.1"), // this value is invalid as it's above 1.0
+      new ConfigEntry(LogConfig.CompressionTypeProp, "lz4")
+    ).asJava
+
+    var topicConfigEntries2 = Seq(new ConfigEntry(LogConfig.CompressionTypeProp, "snappy")).asJava
+
+    val brokerResource = new ConfigResource(ConfigResource.Type.BROKER, servers.head.config.brokerId.toString)
+    val brokerConfigEntries = Seq(new ConfigEntry(KafkaConfig.ZkConnectProp, "localhost:2181")).asJava
+
+    // Alter configs: first and third are invalid, second is valid
+    var alterResult = client.alterConfigs(Map(
+      topicResource1 -> new Config(topicConfigEntries1),
+      topicResource2 -> new Config(topicConfigEntries2),
+      brokerResource -> new Config(brokerConfigEntries)
+    ).asJava)
+
+    assertEquals(Set(topicResource1, topicResource2, brokerResource).asJava, alterResult.values.keySet)
+    assertTrue(intercept[ExecutionException](alterResult.values.get(topicResource1).get).getCause.isInstanceOf[InvalidRequestException])
+    alterResult.values.get(topicResource2).get
+    assertTrue(intercept[ExecutionException](alterResult.values.get(brokerResource).get).getCause.isInstanceOf[InvalidRequestException])
+
+    // Verify that first and third resources were not updated and second was updated
+    var describeResult = client.describeConfigs(Seq(topicResource1, topicResource2, brokerResource).asJava)
+    var configs = describeResult.all.get
+    assertEquals(3, configs.size)
+
+    assertEquals(Defaults.LogCleanerMinCleanRatio.toString,
+      configs.get(topicResource1).get(LogConfig.MinCleanableDirtyRatioProp).value)
+    assertEquals(Defaults.CompressionType.toString,
+      configs.get(topicResource1).get(LogConfig.CompressionTypeProp).value)
+
+    assertEquals("snappy", configs.get(topicResource2).get(LogConfig.CompressionTypeProp).value)
+
+    assertEquals(Defaults.CompressionType.toString, configs.get(brokerResource).get(KafkaConfig.CompressionTypeProp).value)
+
+    // Alter configs with validateOnly = true: first and third are invalid, second is valid
+    topicConfigEntries2 = Seq(new ConfigEntry(LogConfig.CompressionTypeProp, "gzip")).asJava
+
+    alterResult = client.alterConfigs(Map(
+      topicResource1 -> new Config(topicConfigEntries1),
+      topicResource2 -> new Config(topicConfigEntries2),
+      brokerResource -> new Config(brokerConfigEntries)
+    ).asJava, new AlterConfigsOptions().validateOnly(true))
+
+    assertEquals(Set(topicResource1, topicResource2, brokerResource).asJava, alterResult.values.keySet)
+    assertTrue(intercept[ExecutionException](alterResult.values.get(topicResource1).get).getCause.isInstanceOf[InvalidRequestException])
+    alterResult.values.get(topicResource2).get
+    assertTrue(intercept[ExecutionException](alterResult.values.get(brokerResource).get).getCause.isInstanceOf[InvalidRequestException])
+
+    // Verify that no resources are updated since validate_only = true
+    describeResult = client.describeConfigs(Seq(topicResource1, topicResource2, brokerResource).asJava)
+    configs = describeResult.all.get
+    assertEquals(3, configs.size)
+
+    assertEquals(Defaults.LogCleanerMinCleanRatio.toString,
+      configs.get(topicResource1).get(LogConfig.MinCleanableDirtyRatioProp).value)
+    assertEquals(Defaults.CompressionType,
+      configs.get(topicResource1).get(LogConfig.CompressionTypeProp).value)
+
+    assertEquals("snappy", configs.get(topicResource2).get(LogConfig.CompressionTypeProp).value)
+
+    assertEquals(Defaults.CompressionType, configs.get(brokerResource).get(KafkaConfig.CompressionTypeProp).value)
+  }
+
+}
diff --git a/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala b/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala
index aad9b6ad24a24..7625d042709bc 100644
--- a/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala
+++ b/core/src/test/scala/integration/kafka/api/PlaintextConsumerTest.scala
@@ -12,17 +12,18 @@
   */
 package kafka.api
 
+import java.time.Duration
 import java.util
+import java.util.Arrays.asList
 import java.util.regex.Pattern
-import java.util.{Collections, Locale, Properties}
+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}
 import org.apache.kafka.common.{MetricName, TopicPartition}
-import org.apache.kafka.common.errors.InvalidTopicException
+import org.apache.kafka.common.errors.{InvalidGroupIdException, InvalidTopicException}
 import org.apache.kafka.common.header.Headers
 import org.apache.kafka.common.record.{CompressionType, TimestampType}
 import org.apache.kafka.common.serialization._
@@ -30,30 +31,35 @@ 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.jdk.CollectionConverters._
 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)
-    
+
     record.headers().add("headerKey", "headerValue".getBytes)
-    
-    this.producers.head.send(record)
-    
-    assertEquals(0, this.consumers.head.assignment.size)
-    this.consumers.head.assign(List(tp).asJava)
-    assertEquals(1, this.consumers.head.assignment.size)
 
-    this.consumers.head.seek(tp, 0)
-    val records = consumeRecords(consumer = this.consumers.head, numRecords = numRecords)
+    val producer = createProducer()
+    producer.send(record)
+
+    val consumer = createConsumer()
+    assertEquals(0, consumer.assignment.size)
+    consumer.assign(List(tp).asJava)
+    assertEquals(1, consumer.assignment.size)
+
+    consumer.seek(tp, 0)
+    val records = consumeRecords(consumer = consumer, numRecords = numRecords)
 
     assertEquals(numRecords, records.size)
 
@@ -63,164 +69,181 @@ class PlaintextConsumerTest extends BaseConsumerTest {
       assertEquals("headerValue", if (header == null) null else new String(header.value()))
     }
   }
-  
-  @Test
-  def testHeadersExtendedSerializerDeserializer() {
-    val numRecords = 1
-    val record = new ProducerRecord(tp.topic, tp.partition, null, "key".getBytes, "value".getBytes)
 
-    val extendedSerializer = new ExtendedSerializer[Array[Byte]] {
-      
-      var serializer = new ByteArraySerializer()
-      
-      override def serialize(topic: String, headers: Headers, data: Array[Byte]): Array[Byte] = {
-        headers.add("content-type", "application/octet-stream".getBytes)
-        serializer.serialize(topic, data)
-      }
-
-      override def configure(configs: util.Map[String, _], isKey: Boolean): Unit = serializer.configure(configs, isKey)
-      
-      override def close(): Unit = serializer.close()
+  trait SerializerImpl extends Serializer[Array[Byte]]{
+    var serializer = new ByteArraySerializer()
 
-      override def serialize(topic: String, data: Array[Byte]): Array[Byte] = {
-        fail("method should not be invoked")
-        null
-      }
+    override def serialize(topic: String, headers: Headers, data: Array[Byte]): Array[Byte] = {
+      headers.add("content-type", "application/octet-stream".getBytes)
+      serializer.serialize(topic, data)
     }
 
+    override def configure(configs: util.Map[String, _], isKey: Boolean): Unit = serializer.configure(configs, isKey)
 
-    val extendedDeserializer = new ExtendedDeserializer[Array[Byte]] {
-      
-      var deserializer = new ByteArrayDeserializer()
-      
-      override def deserialize(topic: String, headers: Headers, data: Array[Byte]): Array[Byte] = {
-        val header = headers.lastHeader("content-type")
-        assertEquals("application/octet-stream", if (header == null) null else new String(header.value()))
-        deserializer.deserialize(topic, data)
-      }
+    override def close(): Unit = serializer.close()
 
-      override def configure(configs: util.Map[String, _], isKey: Boolean): Unit = deserializer.configure(configs, isKey)
+    override def serialize(topic: String, data: Array[Byte]): Array[Byte] = {
+      fail("method should not be invoked")
+      null
+    }
+  }
 
+  trait DeserializerImpl extends Deserializer[Array[Byte]]{
+    var deserializer = new ByteArrayDeserializer()
 
-      override def close(): Unit = deserializer.close()
+    override def deserialize(topic: String, headers: Headers, data: Array[Byte]): Array[Byte] = {
+      val header = headers.lastHeader("content-type")
+      assertEquals("application/octet-stream", if (header == null) null else new String(header.value()))
+      deserializer.deserialize(topic, data)
+    }
 
-      override def deserialize(topic: String, data: Array[Byte]): Array[Byte] = {
-        fail("method should not be invoked")
-        null
-      }
+    override def configure(configs: util.Map[String, _], isKey: Boolean): Unit = deserializer.configure(configs, isKey)
+
+    override def close(): Unit = deserializer.close()
 
+    override def deserialize(topic: String, data: Array[Byte]): Array[Byte] = {
+      fail("method should not be invoked")
+      null
     }
-    
-    val producer0 = new KafkaProducer(this.producerConfig, new ByteArraySerializer(), extendedSerializer)
-    producers += producer0
-    producer0.send(record)
+  }
 
-    val consumer0 = new KafkaConsumer(this.consumerConfig, new ByteArrayDeserializer(), extendedDeserializer)
-    consumers += consumer0
+  private def testHeadersSerializeDeserialize(serializer: Serializer[Array[Byte]], deserializer: Deserializer[Array[Byte]]): Unit = {
+    val numRecords = 1
+    val record = new ProducerRecord(tp.topic, tp.partition, null, "key".getBytes, "value".getBytes)
+
+    val producer = createProducer(
+      keySerializer = new ByteArraySerializer,
+      valueSerializer = serializer)
+    producer.send(record)
 
-    assertEquals(0, consumer0.assignment.size)
-    consumer0.assign(List(tp).asJava)
-    assertEquals(1, consumer0.assignment.size)
+    val consumer = createConsumer(
+      keyDeserializer = new ByteArrayDeserializer,
+      valueDeserializer = deserializer)
+    assertEquals(0, consumer.assignment.size)
+    consumer.assign(List(tp).asJava)
+    assertEquals(1, consumer.assignment.size)
 
-    consumer0.seek(tp, 0)
-    val records = consumeRecords(consumer = consumer0, numRecords = numRecords)
+    consumer.seek(tp, 0)
+    val records = consumeRecords(consumer = consumer, numRecords = numRecords)
 
     assertEquals(numRecords, records.size)
   }
-  
+
+  @deprecated("poll(Duration) is the replacement", since = "2.0")
   @Test
-  def testMaxPollRecords() {
+  def testDeprecatedPollBlocksForAssignment(): Unit = {
+    val consumer = createConsumer()
+    consumer.subscribe(Set(topic).asJava)
+    consumer.poll(0)
+    assertEquals(Set(tp, tp2), consumer.assignment().asScala)
+  }
+
+  @deprecated("Serializer now includes a default method that provides the headers", since = "2.1")
+  @Test
+  def testHeadersExtendedSerializerDeserializer(): Unit = {
+    val extendedSerializer = new ExtendedSerializer[Array[Byte]] with SerializerImpl
+
+    val extendedDeserializer = new ExtendedDeserializer[Array[Byte]] with DeserializerImpl
+
+    testHeadersSerializeDeserialize(extendedSerializer, extendedDeserializer)
+  }
+
+  @Test
+  def testHeadersSerializerDeserializer(): Unit = {
+    val extendedSerializer = new Serializer[Array[Byte]] with SerializerImpl
+
+    val extendedDeserializer = new Deserializer[Array[Byte]] with DeserializerImpl
+
+    testHeadersSerializeDeserialize(extendedSerializer, extendedDeserializer)
+  }
+
+  @Test
+  def testMaxPollRecords(): Unit = {
     val maxPollRecords = 2
     val numRecords = 10000
 
-    sendRecords(numRecords)
+    val producer = createProducer()
+    sendRecords(producer, numRecords, tp)
 
     this.consumerConfig.setProperty(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, maxPollRecords.toString)
-    val consumer0 = new KafkaConsumer(this.consumerConfig, new ByteArrayDeserializer(), new ByteArrayDeserializer())
-    consumers += consumer0
-
-    consumer0.assign(List(tp).asJava)
-
-    consumeAndVerifyRecords(consumer0, numRecords = numRecords, startingOffset = 0,
-      maxPollRecords = maxPollRecords)
+    val consumer = createConsumer()
+    consumer.assign(List(tp).asJava)
+    consumeAndVerifyRecords(consumer, numRecords = numRecords, startingOffset = 0, maxPollRecords = maxPollRecords)
   }
 
   @Test
-  def testMaxPollIntervalMs() {
-    this.consumerConfig.setProperty(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG, 3000.toString)
+  def testMaxPollIntervalMs(): Unit = {
+    this.consumerConfig.setProperty(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG, 1000.toString)
     this.consumerConfig.setProperty(ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG, 500.toString)
     this.consumerConfig.setProperty(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG, 2000.toString)
 
-    val consumer0 = new KafkaConsumer(this.consumerConfig, new ByteArrayDeserializer(), new ByteArrayDeserializer())
-    consumers += consumer0
+    val consumer = createConsumer()
 
     val listener = new TestConsumerReassignmentListener()
-    consumer0.subscribe(List(topic).asJava, listener)
+    consumer.subscribe(List(topic).asJava, listener)
 
-    // poll once to get the initial assignment
-    consumer0.poll(0)
+    // rebalance to get the initial assignment
+    awaitRebalance(consumer, listener)
     assertEquals(1, listener.callsToAssigned)
-    assertEquals(1, listener.callsToRevoked)
+    assertEquals(0, listener.callsToRevoked)
 
-    Thread.sleep(3500)
+    // after we extend longer than max.poll a rebalance should be triggered
+    // NOTE we need to have a relatively much larger value than max.poll to let heartbeat expired for sure
+    Thread.sleep(3000)
 
-    // we should fall out of the group and need to rebalance
-    consumer0.poll(0)
+    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)
     this.consumerConfig.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, false.toString)
 
-    val consumer0 = new KafkaConsumer(this.consumerConfig, new ByteArrayDeserializer(), new ByteArrayDeserializer())
-    consumers += consumer0
-
+    val consumer = createConsumer()
     var commitCompleted = false
     var committedPosition: Long = -1
 
     val listener = new TestConsumerReassignmentListener {
+      override def onPartitionsLost(partitions: util.Collection[TopicPartition]): Unit = {}
       override def onPartitionsRevoked(partitions: util.Collection[TopicPartition]): Unit = {
-        if (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
           Utils.sleep(1500)
-          committedPosition = consumer0.position(tp)
-          consumer0.commitSync(Map(tp -> new OffsetAndMetadata(committedPosition)).asJava)
+          committedPosition = consumer.position(tp)
+          consumer.commitSync(Map(tp -> new OffsetAndMetadata(committedPosition)).asJava)
           commitCompleted = true
         }
         super.onPartitionsRevoked(partitions)
       }
     }
 
-    consumer0.subscribe(List(topic).asJava, listener)
+    consumer.subscribe(List(topic).asJava, listener)
 
-    // poll once to join the group and get the initial assignment
-    consumer0.poll(0)
+    // rebalance to get the initial assignment
+    awaitRebalance(consumer, listener)
 
     // force a rebalance to trigger an invocation of the revocation callback while in the group
-    consumer0.subscribe(List("otherTopic").asJava, listener)
-    consumer0.poll(0)
+    consumer.subscribe(List("otherTopic").asJava, listener)
+    awaitRebalance(consumer, listener)
 
     assertEquals(0, committedPosition)
     assertTrue(commitCompleted)
   }
 
   @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)
     this.consumerConfig.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, false.toString)
 
-    val consumer0 = new KafkaConsumer(this.consumerConfig, new ByteArrayDeserializer(), new ByteArrayDeserializer())
-    consumers += consumer0
-
+    val consumer = createConsumer()
     val listener = new TestConsumerReassignmentListener {
       override def onPartitionsAssigned(partitions: util.Collection[TopicPartition]): Unit = {
         // sleep longer than the session timeout, we should still be in the group after invocation
@@ -228,86 +251,83 @@ class PlaintextConsumerTest extends BaseConsumerTest {
         super.onPartitionsAssigned(partitions)
       }
     }
-    consumer0.subscribe(List(topic).asJava, listener)
-
-    // poll once to join the group and get the initial assignment
-    consumer0.poll(0)
+    consumer.subscribe(List(topic).asJava, listener)
 
-    // we should still be in the group after this invocation
-    consumer0.poll(0)
+    // rebalance to get the initial assignment
+    awaitRebalance(consumer, listener)
 
-    assertEquals(1, listener.callsToAssigned)
-    assertEquals(1, listener.callsToRevoked)
+    // We should still be in the group after this invocation
+    ensureNoRebalance(consumer, listener)
   }
 
   @Test
-  def testAutoCommitOnClose() {
+  def testAutoCommitOnClose(): Unit = {
     this.consumerConfig.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true")
-    val consumer0 = new KafkaConsumer(this.consumerConfig, new ByteArrayDeserializer(), new ByteArrayDeserializer())
+    val consumer = createConsumer()
 
     val numRecords = 10000
-    sendRecords(numRecords)
+    val producer = createProducer()
+    sendRecords(producer, numRecords, tp)
 
-    consumer0.subscribe(List(topic).asJava)
-
-    val assignment = Set(tp, tp2)
-    TestUtils.waitUntilTrue(() => {
-      consumer0.poll(50)
-      consumer0.assignment() == assignment.asJava
-    }, s"Expected partitions ${assignment.asJava} but actually got ${consumer0.assignment()}")
+    consumer.subscribe(List(topic).asJava)
+    awaitAssignment(consumer, Set(tp, tp2))
 
     // should auto-commit seeked positions before closing
-    consumer0.seek(tp, 300)
-    consumer0.seek(tp2, 500)
-    consumer0.close()
+    consumer.seek(tp, 300)
+    consumer.seek(tp2, 500)
+    consumer.close()
 
     // now we should see the committed positions from another consumer
-    assertEquals(300, this.consumers.head.committed(tp).offset)
-    assertEquals(500, this.consumers.head.committed(tp2).offset)
+    val anotherConsumer = createConsumer()
+    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 consumer0 = new KafkaConsumer(this.consumerConfig, new ByteArrayDeserializer(), new ByteArrayDeserializer())
+    val consumer = createConsumer()
 
     val numRecords = 10000
-    sendRecords(numRecords)
-
-    consumer0.subscribe(List(topic).asJava)
+    val producer = createProducer()
+    sendRecords(producer, numRecords, tp)
 
-    val assignment = Set(tp, tp2)
-    TestUtils.waitUntilTrue(() => {
-      consumer0.poll(50)
-      consumer0.assignment() == assignment.asJava
-    }, s"Expected partitions ${assignment.asJava} but actually got ${consumer0.assignment()}")
+    consumer.subscribe(List(topic).asJava)
+    awaitAssignment(consumer, Set(tp, tp2))
 
     // should auto-commit seeked positions before closing
-    consumer0.seek(tp, 300)
-    consumer0.seek(tp2, 500)
+    consumer.seek(tp, 300)
+    consumer.seek(tp2, 500)
 
     // wakeup the consumer before closing to simulate trying to break a poll
     // loop from another thread
-    consumer0.wakeup()
-    consumer0.close()
+    consumer.wakeup()
+    consumer.close()
 
     // now we should see the committed positions from another consumer
-    assertEquals(300, this.consumers.head.committed(tp).offset)
-    assertEquals(500, this.consumers.head.committed(tp2).offset)
+    val anotherConsumer = createConsumer()
+    assertEquals(300, anotherConsumer.committed(Set(tp).asJava).get(tp).offset)
+    assertEquals(500, anotherConsumer.committed(Set(tp2).asJava).get(tp2).offset)
   }
 
   @Test
-  def testAutoOffsetReset() {
-    sendRecords(1)
-    this.consumers.head.assign(List(tp).asJava)
-    consumeAndVerifyRecords(consumer = this.consumers.head, numRecords = 1, startingOffset = 0)
+  def testAutoOffsetReset(): Unit = {
+    val producer = createProducer()
+    sendRecords(producer, numRecords = 1, tp)
+
+    val consumer = createConsumer()
+    consumer.assign(List(tp).asJava)
+    consumeAndVerifyRecords(consumer = consumer, numRecords = 1, startingOffset = 0)
   }
 
   @Test
-  def testGroupConsumption() {
-    sendRecords(10)
-    this.consumers.head.subscribe(List(topic).asJava)
-    consumeAndVerifyRecords(consumer = this.consumers.head, numRecords = 1, startingOffset = 0)
+  def testGroupConsumption(): Unit = {
+    val producer = createProducer()
+    sendRecords(producer, numRecords = 10, tp)
+
+    val consumer = createConsumer()
+    consumer.subscribe(List(topic).asJava)
+    consumeAndVerifyRecords(consumer = consumer, numRecords = 1, startingOffset = 0)
   }
 
   /**
@@ -320,59 +340,51 @@ class PlaintextConsumerTest extends BaseConsumerTest {
    * of that topic are assigned to it.
    */
   @Test
-  def testPatternSubscription() {
+  def testPatternSubscription(): Unit = {
     val numRecords = 10000
-    sendRecords(numRecords)
+    val producer = createProducer()
+    sendRecords(producer, numRecords, tp)
 
     val topic1 = "tblablac" // matches subscribed pattern
-    TestUtils.createTopic(this.zkUtils, topic1, 2, serverCount, this.servers)
-    sendRecords(1000, new TopicPartition(topic1, 0))
-    sendRecords(1000, new TopicPartition(topic1, 1))
+    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
-    TestUtils.createTopic(this.zkUtils, topic2, 2, serverCount, this.servers)
-    sendRecords(1000, new TopicPartition(topic2, 0))
-    sendRecords(1000, new TopicPartition(topic2, 1))
+    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
-    TestUtils.createTopic(this.zkUtils, topic3, 2, serverCount, this.servers)
-    sendRecords(1000, new TopicPartition(topic3, 0))
-    sendRecords(1000, new TopicPartition(topic3, 1))
+    createTopic(topic3, 2, brokerCount)
+    sendRecords(producer, numRecords = 1000, new TopicPartition(topic3, 0))
+    sendRecords(producer, numRecords = 1000, new TopicPartition(topic3, 1))
 
-    assertEquals(0, this.consumers.head.assignment().size)
+    val consumer = createConsumer()
+    assertEquals(0, consumer.assignment().size)
 
     val pattern = Pattern.compile("t.*c")
-    this.consumers.head.subscribe(pattern, new TestConsumerReassignmentListener)
-    this.consumers.head.poll(50)
+    consumer.subscribe(pattern, new TestConsumerReassignmentListener)
 
-    var subscriptions = Set(
+    var assignment = Set(
       new TopicPartition(topic, 0),
       new TopicPartition(topic, 1),
       new TopicPartition(topic1, 0),
       new TopicPartition(topic1, 1))
-
-    TestUtils.waitUntilTrue(() => {
-      this.consumers.head.poll(50)
-      this.consumers.head.assignment() == subscriptions.asJava
-    }, s"Expected partitions ${subscriptions.asJava} but actually got ${this.consumers.head.assignment()}")
+    awaitAssignment(consumer, assignment)
 
     val topic4 = "tsomec" // matches subscribed pattern
-    TestUtils.createTopic(this.zkUtils, topic4, 2, serverCount, this.servers)
-    sendRecords(1000, new TopicPartition(topic4, 0))
-    sendRecords(1000, new TopicPartition(topic4, 1))
+    createTopic(topic4, 2, brokerCount)
+    sendRecords(producer, numRecords = 1000, new TopicPartition(topic4, 0))
+    sendRecords(producer, numRecords = 1000, new TopicPartition(topic4, 1))
 
-    subscriptions ++= Set(
+    assignment ++= Set(
       new TopicPartition(topic4, 0),
       new TopicPartition(topic4, 1))
+    awaitAssignment(consumer, assignment)
 
-
-    TestUtils.waitUntilTrue(() => {
-      this.consumers.head.poll(50)
-      this.consumers.head.assignment() == subscriptions.asJava
-    }, s"Expected partitions ${subscriptions.asJava} but actually got ${this.consumers.head.assignment()}")
-
-    this.consumers.head.unsubscribe()
-    assertEquals(0, this.consumers.head.assignment().size)
+    consumer.unsubscribe()
+    assertEquals(0, consumer.assignment().size)
   }
 
   /**
@@ -385,58 +397,46 @@ 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 consumer0 = new KafkaConsumer(this.consumerConfig, new ByteArrayDeserializer(), new ByteArrayDeserializer())
-    consumers += consumer0
+    val consumer = createConsumer()
 
     val numRecords = 10000
-    sendRecords(numRecords)
+    val producer = createProducer()
+    sendRecords(producer, numRecords = numRecords, tp)
 
     // the first topic ('topic')  matches first subscription pattern only
 
     val fooTopic = "foo" // matches both subscription patterns
-    TestUtils.createTopic(this.zkUtils, fooTopic, 1, serverCount, this.servers)
-    sendRecords(1000, new TopicPartition(fooTopic, 0))
+    createTopic(fooTopic, 1, brokerCount)
+    sendRecords(producer, numRecords = 1000, new TopicPartition(fooTopic, 0))
 
-    assertEquals(0, consumer0.assignment().size)
+    assertEquals(0, consumer.assignment().size)
 
     val pattern1 = Pattern.compile(".*o.*") // only 'topic' and 'foo' match this
-    consumer0.subscribe(pattern1, new TestConsumerReassignmentListener)
-    consumer0.poll(50)
+    consumer.subscribe(pattern1, new TestConsumerReassignmentListener)
 
-    var subscriptions = Set(
+    var assignment = Set(
       new TopicPartition(topic, 0),
       new TopicPartition(topic, 1),
       new TopicPartition(fooTopic, 0))
-
-    TestUtils.waitUntilTrue(() => {
-      consumer0.poll(50)
-      consumer0.assignment() == subscriptions.asJava
-    }, s"Expected partitions ${subscriptions.asJava} but actually got ${consumer0.assignment()}")
+    awaitAssignment(consumer, assignment)
 
     val barTopic = "bar" // matches the next subscription pattern
-    TestUtils.createTopic(this.zkUtils, barTopic, 1, serverCount, this.servers)
-    sendRecords(1000, new TopicPartition(barTopic, 0))
+    createTopic(barTopic, 1, brokerCount)
+    sendRecords(producer, numRecords = 1000, new TopicPartition(barTopic, 0))
 
     val pattern2 = Pattern.compile("...") // only 'foo' and 'bar' match this
-    consumer0.subscribe(pattern2, new TestConsumerReassignmentListener)
-    consumer0.poll(50)
-
-    subscriptions --= Set(
+    consumer.subscribe(pattern2, new TestConsumerReassignmentListener)
+    assignment --= Set(
       new TopicPartition(topic, 0),
       new TopicPartition(topic, 1))
-
-    subscriptions ++= Set(
+    assignment ++= Set(
       new TopicPartition(barTopic, 0))
+    awaitAssignment(consumer, assignment)
 
-    TestUtils.waitUntilTrue(() => {
-      consumer0.poll(50)
-      consumer0.assignment() == subscriptions.asJava
-    }, s"Expected partitions ${subscriptions.asJava} but actually got ${consumer0.assignment()}")
-
-    consumer0.unsubscribe()
-    assertEquals(0, consumer0.assignment().size)
+    consumer.unsubscribe()
+    assertEquals(0, consumer.assignment().size)
   }
 
   /**
@@ -448,146 +448,139 @@ class PlaintextConsumerTest extends BaseConsumerTest {
    * assignments are cleared right away.
    */
   @Test
-  def testPatternUnsubscription() {
+  def testPatternUnsubscription(): Unit = {
     val numRecords = 10000
-    sendRecords(numRecords)
+    val producer = createProducer()
+    sendRecords(producer, numRecords, tp)
 
     val topic1 = "tblablac" // matches the subscription pattern
-    TestUtils.createTopic(this.zkUtils, topic1, 2, serverCount, this.servers)
-    sendRecords(1000, new TopicPartition(topic1, 0))
-    sendRecords(1000, new TopicPartition(topic1, 1))
+    createTopic(topic1, 2, brokerCount)
+    sendRecords(producer, numRecords = 1000, new TopicPartition(topic1, 0))
+    sendRecords(producer, numRecords = 1000, new TopicPartition(topic1, 1))
 
-    assertEquals(0, this.consumers.head.assignment().size)
-
-    this.consumers.head.subscribe(Pattern.compile("t.*c"), new TestConsumerReassignmentListener)
-    this.consumers.head.poll(50)
+    val consumer = createConsumer()
+    assertEquals(0, consumer.assignment().size)
 
-    val subscriptions = Set(
+    consumer.subscribe(Pattern.compile("t.*c"), new TestConsumerReassignmentListener)
+    val assignment = Set(
       new TopicPartition(topic, 0),
       new TopicPartition(topic, 1),
       new TopicPartition(topic1, 0),
       new TopicPartition(topic1, 1))
+    awaitAssignment(consumer, assignment)
 
-    TestUtils.waitUntilTrue(() => {
-      this.consumers.head.poll(50)
-      this.consumers.head.assignment() == subscriptions.asJava
-    }, s"Expected partitions ${subscriptions.asJava} but actually got ${this.consumers.head.assignment()}")
-
-    this.consumers.head.unsubscribe()
-    assertEquals(0, this.consumers.head.assignment().size)
+    consumer.unsubscribe()
+    assertEquals(0, consumer.assignment().size)
   }
 
   @Test
-  def testCommitMetadata() {
-    this.consumers.head.assign(List(tp).asJava)
+  def testCommitMetadata(): Unit = {
+    val consumer = createConsumer()
+    consumer.assign(List(tp).asJava)
 
     // sync commit
-    val syncMetadata = new OffsetAndMetadata(5, "foo")
-    this.consumers.head.commitSync(Map((tp, syncMetadata)).asJava)
-    assertEquals(syncMetadata, this.consumers.head.committed(tp))
+    val syncMetadata = new OffsetAndMetadata(5, Optional.of(15), "foo")
+    consumer.commitSync(Map((tp, syncMetadata)).asJava)
+    assertEquals(syncMetadata, consumer.committed(Set(tp).asJava).get(tp))
 
     // async commit
     val asyncMetadata = new OffsetAndMetadata(10, "bar")
-    val callback = new CountConsumerCommitCallback
-    this.consumers.head.commitAsync(Map((tp, asyncMetadata)).asJava, callback)
-    awaitCommitCallback(this.consumers.head, callback)
-    assertEquals(asyncMetadata, this.consumers.head.committed(tp))
+    sendAndAwaitAsyncCommit(consumer, Some(Map(tp -> asyncMetadata)))
+    assertEquals(asyncMetadata, consumer.committed(Set(tp).asJava).get(tp))
 
     // handle null metadata
     val nullMetadata = new OffsetAndMetadata(5, null)
-    this.consumers.head.commitSync(Map((tp, nullMetadata)).asJava)
-    assertEquals(nullMetadata, this.consumers.head.committed(tp))
+    consumer.commitSync(Map(tp -> nullMetadata).asJava)
+    assertEquals(nullMetadata, consumer.committed(Set(tp).asJava).get(tp))
   }
 
   @Test
-  def testAsyncCommit() {
-    val consumer = this.consumers.head
+  def testAsyncCommit(): Unit = {
+    val consumer = createConsumer()
     consumer.assign(List(tp).asJava)
-    consumer.poll(0)
 
     val callback = new CountConsumerCommitCallback
     val count = 5
+
     for (i <- 1 to count)
       consumer.commitAsync(Map(tp -> new OffsetAndMetadata(i)).asJava, callback)
 
-    awaitCommitCallback(consumer, callback, count=count)
-    assertEquals(new OffsetAndMetadata(count), consumer.committed(tp))
+    TestUtils.pollUntilTrue(consumer, () => callback.successCount >= count || callback.lastError.isDefined,
+      "Failed to observe commit callback before timeout", waitTimeMs = 10000)
+
+    assertEquals(None, callback.lastError)
+    assertEquals(count, callback.successCount)
+    assertEquals(new OffsetAndMetadata(count), consumer.committed(Set(tp).asJava).get(tp))
   }
 
   @Test
-  def testExpandingTopicSubscriptions() {
+  def testExpandingTopicSubscriptions(): Unit = {
     val otherTopic = "other"
-    val subscriptions = Set(new TopicPartition(topic, 0), new TopicPartition(topic, 1))
-    val expandedSubscriptions = subscriptions ++ Set(new TopicPartition(otherTopic, 0), new TopicPartition(otherTopic, 1))
-    this.consumers.head.subscribe(List(topic).asJava)
-    TestUtils.waitUntilTrue(() => {
-      this.consumers.head.poll(50)
-      this.consumers.head.assignment == subscriptions.asJava
-    }, s"Expected partitions ${subscriptions.asJava} but actually got ${this.consumers.head.assignment}")
-
-    TestUtils.createTopic(this.zkUtils, otherTopic, 2, serverCount, this.servers)
-    this.consumers.head.subscribe(List(topic, otherTopic).asJava)
-    TestUtils.waitUntilTrue(() => {
-      this.consumers.head.poll(50)
-      this.consumers.head.assignment == expandedSubscriptions.asJava
-    }, s"Expected partitions ${expandedSubscriptions.asJava} but actually got ${this.consumers.head.assignment}")
+    val initialAssignment = Set(new TopicPartition(topic, 0), new TopicPartition(topic, 1))
+    val consumer = createConsumer()
+    consumer.subscribe(List(topic).asJava)
+    awaitAssignment(consumer, initialAssignment)
+
+    createTopic(otherTopic, 2, brokerCount)
+    val expandedAssignment = initialAssignment ++ Set(new TopicPartition(otherTopic, 0), new TopicPartition(otherTopic, 1))
+    consumer.subscribe(List(topic, otherTopic).asJava)
+    awaitAssignment(consumer, expandedAssignment)
   }
 
   @Test
-  def testShrinkingTopicSubscriptions() {
+  def testShrinkingTopicSubscriptions(): Unit = {
     val otherTopic = "other"
-    TestUtils.createTopic(this.zkUtils, otherTopic, 2, serverCount, this.servers)
-    val subscriptions = Set(new TopicPartition(topic, 0), new TopicPartition(topic, 1), new TopicPartition(otherTopic, 0), new TopicPartition(otherTopic, 1))
-    val shrunkenSubscriptions = Set(new TopicPartition(topic, 0), new TopicPartition(topic, 1))
-    this.consumers.head.subscribe(List(topic, otherTopic).asJava)
-    TestUtils.waitUntilTrue(() => {
-      this.consumers.head.poll(50)
-      this.consumers.head.assignment == subscriptions.asJava
-    }, s"Expected partitions ${subscriptions.asJava} but actually got ${this.consumers.head.assignment}")
-
-    this.consumers.head.subscribe(List(topic).asJava)
-    TestUtils.waitUntilTrue(() => {
-      this.consumers.head.poll(50)
-      this.consumers.head.assignment == shrunkenSubscriptions.asJava
-    }, s"Expected partitions ${shrunkenSubscriptions.asJava} but actually got ${this.consumers.head.assignment}")
+    createTopic(otherTopic, 2, brokerCount)
+    val initialAssignment = Set(new TopicPartition(topic, 0), new TopicPartition(topic, 1), new TopicPartition(otherTopic, 0), new TopicPartition(otherTopic, 1))
+    val consumer = createConsumer()
+    consumer.subscribe(List(topic, otherTopic).asJava)
+    awaitAssignment(consumer, initialAssignment)
+
+    val shrunkenAssignment = Set(new TopicPartition(topic, 0), new TopicPartition(topic, 1))
+    consumer.subscribe(List(topic).asJava)
+    awaitAssignment(consumer, shrunkenAssignment)
   }
 
   @Test
-  def testPartitionsFor() {
+  def testPartitionsFor(): Unit = {
     val numParts = 2
-    TestUtils.createTopic(this.zkUtils, "part-test", numParts, 1, this.servers)
-    val parts = this.consumers.head.partitionsFor("part-test")
+    createTopic("part-test", numParts, 1)
+    val consumer = createConsumer()
+    val parts = consumer.partitionsFor("part-test")
     assertNotNull(parts)
     assertEquals(2, parts.size)
   }
 
   @Test
-  def testPartitionsForAutoCreate() {
-    val partitions = this.consumers.head.partitionsFor("non-exist-topic")
+  def testPartitionsForAutoCreate(): Unit = {
+    val consumer = createConsumer()
+    val partitions = consumer.partitionsFor("non-exist-topic")
     assertFalse(partitions.isEmpty)
   }
 
   @Test(expected = classOf[InvalidTopicException])
-  def testPartitionsForInvalidTopic() {
-    this.consumers.head.partitionsFor(";3# ads,{234")
+  def testPartitionsForInvalidTopic(): Unit = {
+    val consumer = createConsumer()
+    consumer.partitionsFor(";3# ads,{234")
   }
 
   @Test
-  def testSeek() {
-    val consumer = this.consumers.head
+  def testSeek(): Unit = {
+    val consumer = createConsumer()
     val totalRecords = 50L
     val mid = totalRecords / 2
 
     // Test seek non-compressed message
-    sendRecords(totalRecords.toInt, tp)
+    val producer = createProducer()
+    sendRecords(producer, totalRecords.toInt, tp)
     consumer.assign(List(tp).asJava)
 
     consumer.seekToEnd(List(tp).asJava)
     assertEquals(totalRecords, consumer.position(tp))
-    assertFalse(consumer.poll(totalRecords).iterator().hasNext)
+    assertTrue(consumer.poll(Duration.ofMillis(50)).isEmpty)
 
     consumer.seekToBeginning(List(tp).asJava)
-    assertEquals(0, consumer.position(tp), 0)
+    assertEquals(0L, consumer.position(tp))
     consumeAndVerifyRecords(consumer, numRecords = 1, startingOffset = 0)
 
     consumer.seek(tp, mid)
@@ -602,10 +595,10 @@ class PlaintextConsumerTest extends BaseConsumerTest {
 
     consumer.seekToEnd(List(tp2).asJava)
     assertEquals(totalRecords, consumer.position(tp2))
-    assertFalse(consumer.poll(totalRecords).iterator().hasNext)
+    assertTrue(consumer.poll(Duration.ofMillis(50)).isEmpty)
 
     consumer.seekToBeginning(List(tp2).asJava)
-    assertEquals(0, consumer.position(tp2), 0)
+    assertEquals(0L, consumer.position(tp2))
     consumeAndVerifyRecords(consumer, numRecords = 1, startingOffset = 0, tp = tp2)
 
     consumer.seek(tp2, mid)
@@ -614,12 +607,11 @@ 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, Long.MaxValue.toString)
-    val producer = TestUtils.createNewProducer(brokerList, securityProtocol = securityProtocol, trustStoreFile = trustStoreFile,
-        saslProperties = clientSaslProperties, retries = 0, lingerMs = Long.MaxValue, props = Some(producerProps))
+    producerProps.setProperty(ProducerConfig.LINGER_MS_CONFIG, Int.MaxValue.toString)
+    val producer = createProducer(configOverrides = producerProps)
     (0 until numRecords).foreach { i =>
       producer.send(new ProducerRecord(tp.topic, tp.partition, i.toLong, s"key $i".getBytes, s"value $i".getBytes))
     }
@@ -627,68 +619,76 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   @Test
-  def testPositionAndCommit() {
-    sendRecords(5)
+  def testPositionAndCommit(): Unit = {
+    val producer = createProducer()
+    sendRecords(producer, numRecords = 5, tp)
 
-    assertNull(this.consumers.head.committed(new TopicPartition(topic, 15)))
+    val topicPartition = new TopicPartition(topic, 15)
+    val consumer = createConsumer()
+    assertNull(consumer.committed(Set(topicPartition).asJava).get(topicPartition))
 
     // position() on a partition that we aren't subscribed to throws an exception
-    intercept[IllegalArgumentException] {
-      this.consumers.head.position(new TopicPartition(topic, 15))
+    intercept[IllegalStateException] {
+      consumer.position(topicPartition)
     }
 
-    this.consumers.head.assign(List(tp).asJava)
+    consumer.assign(List(tp).asJava)
 
-    assertEquals("position() on a partition that we are subscribed to should reset the offset", 0L, this.consumers.head.position(tp))
-    this.consumers.head.commitSync()
-    assertEquals(0L, this.consumers.head.committed(tp).offset)
+    assertEquals("position() on a partition that we are subscribed to should reset the offset", 0L, consumer.position(tp))
+    consumer.commitSync()
+    assertEquals(0L, consumer.committed(Set(tp).asJava).get(tp).offset)
 
-    consumeAndVerifyRecords(consumer = this.consumers.head, numRecords = 5, startingOffset = 0)
-    assertEquals("After consuming 5 records, position should be 5", 5L, this.consumers.head.position(tp))
-    this.consumers.head.commitSync()
-    assertEquals("Committed offset should be returned", 5L, this.consumers.head.committed(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(Set(tp).asJava).get(tp).offset)
 
-    sendRecords(1)
+    sendRecords(producer, numRecords = 1, tp)
 
     // another consumer in the same group should get the same position
-    this.consumers(1).assign(List(tp).asJava)
-    consumeAndVerifyRecords(consumer = this.consumers(1), numRecords = 1, startingOffset = 5)
+    val otherConsumer = createConsumer()
+    otherConsumer.assign(List(tp).asJava)
+    consumeAndVerifyRecords(consumer = otherConsumer, numRecords = 1, startingOffset = 5)
   }
 
   @Test
-  def testPartitionPauseAndResume() {
+  def testPartitionPauseAndResume(): Unit = {
     val partitions = List(tp).asJava
-    sendRecords(5)
-    this.consumers.head.assign(partitions)
-    consumeAndVerifyRecords(consumer = this.consumers.head, numRecords = 5, startingOffset = 0)
-    this.consumers.head.pause(partitions)
-    sendRecords(5)
-    assertTrue(this.consumers.head.poll(0).isEmpty)
-    this.consumers.head.resume(partitions)
-    consumeAndVerifyRecords(consumer = this.consumers.head, numRecords = 5, startingOffset = 5)
+    val producer = createProducer()
+    sendRecords(producer, numRecords = 5, tp)
+
+    val consumer = createConsumer()
+    consumer.assign(partitions)
+    consumeAndVerifyRecords(consumer = consumer, numRecords = 5, startingOffset = 0)
+    consumer.pause(partitions)
+    sendRecords(producer, numRecords = 5, tp)
+    assertTrue(consumer.poll(Duration.ofMillis(100)).isEmpty)
+    consumer.resume(partitions)
+    consumeAndVerifyRecords(consumer = consumer, numRecords = 5, startingOffset = 5)
   }
 
   @Test
-  def testFetchInvalidOffset() {
+  def testFetchInvalidOffset(): Unit = {
     this.consumerConfig.setProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "none")
-    val consumer0 = new KafkaConsumer(this.consumerConfig, new ByteArrayDeserializer(), new ByteArrayDeserializer())
-    consumers += consumer0
+    val consumer = createConsumer()
 
     // produce one record
     val totalRecords = 2
-    sendRecords(totalRecords, tp)
-    consumer0.assign(List(tp).asJava)
+    val producer = createProducer()
+    sendRecords(producer, totalRecords, tp)
+    consumer.assign(List(tp).asJava)
 
-    // poll should fail because there is no offset reset strategy set
+    // poll should fail because there is no offset reset strategy set.
+    // we fail only when resetting positions after coordinator is known, so using a long timeout.
     intercept[NoOffsetForPartitionException] {
-      consumer0.poll(50)
+      consumer.poll(Duration.ofMillis(15000))
     }
 
     // seek to out of range position
     val outOfRangePos = totalRecords + 1
-    consumer0.seek(tp, outOfRangePos)
+    consumer.seek(tp, outOfRangePos)
     val e = intercept[OffsetOutOfRangeException] {
-      consumer0.poll(20000)
+      consumer.poll(Duration.ofMillis(20000))
     }
     val outOfRangePartitions = e.offsetOutOfRangePartitions()
     assertNotNull(outOfRangePartitions)
@@ -697,24 +697,24 @@ 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)
   }
 
   private def checkLargeRecord(producerRecordSize: Int): Unit = {
-    val consumer0 = new KafkaConsumer(this.consumerConfig, new ByteArrayDeserializer(), new ByteArrayDeserializer())
-    consumers += consumer0
+    val consumer = createConsumer()
 
     // produce a record that is larger than the configured fetch size
     val record = new ProducerRecord(tp.topic(), tp.partition(), "key".getBytes,
       new Array[Byte](producerRecordSize))
-    this.producers.head.send(record)
+    val producer = createProducer()
+    producer.send(record)
 
     // consuming a record that is too large should succeed since KIP-74
-    consumer0.assign(List(tp).asJava)
-    val records = consumer0.poll(20000)
+    consumer.assign(List(tp).asJava)
+    val records = consumer.poll(Duration.ofMillis(20000))
     assertEquals(1, records.count)
     val consumerRecord = records.iterator().next()
     assertEquals(0L, consumerRecord.offset)
@@ -733,20 +733,20 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   private def checkFetchHonoursSizeIfLargeRecordNotFirst(largeProducerRecordSize: Int): Unit = {
-    val consumer0 = new KafkaConsumer(this.consumerConfig, new ByteArrayDeserializer(), new ByteArrayDeserializer())
-    consumers += consumer0
+    val consumer = createConsumer()
 
     val smallRecord = new ProducerRecord(tp.topic(), tp.partition(), "small".getBytes,
       "value".getBytes)
     val largeRecord = new ProducerRecord(tp.topic(), tp.partition(), "large".getBytes,
       new Array[Byte](largeProducerRecordSize))
 
-    this.producers.head.send(smallRecord).get
-    this.producers.head.send(largeRecord).get
+    val producer = createProducer()
+    producer.send(smallRecord).get
+    producer.send(largeRecord).get
 
     // we should only get the small record in the first `poll`
-    consumer0.assign(List(tp).asJava)
-    val records = consumer0.poll(20000)
+    consumer.assign(List(tp).asJava)
+    val records = consumer.poll(Duration.ofMillis(20000))
     assertEquals(1, records.count)
     val consumerRecord = records.iterator().next()
     assertEquals(0L, consumerRecord.offset)
@@ -765,7 +765,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)
@@ -779,8 +779,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     // this behaves a little different than when remaining limit bytes is 0 and it's important to test it
     this.consumerConfig.setProperty(ConsumerConfig.FETCH_MAX_BYTES_CONFIG, "500")
     this.consumerConfig.setProperty(ConsumerConfig.MAX_PARTITION_FETCH_BYTES_CONFIG, "100")
-    val consumer0 = new KafkaConsumer(this.consumerConfig, new ByteArrayDeserializer(), new ByteArrayDeserializer())
-    consumers += consumer0
+    val consumer = createConsumer()
 
     val topic1 = "topic1"
     val topic2 = "topic2"
@@ -788,24 +787,22 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     val partitionCount = 30
     val topics = Seq(topic1, topic2, topic3)
     topics.foreach { topicName =>
-      TestUtils.createTopic(zkUtils, topicName, partitionCount, serverCount, servers)
+      createTopic(topicName, partitionCount, brokerCount)
     }
 
     val partitions = topics.flatMap { topic =>
       (0 until partitionCount).map(new TopicPartition(topic, _))
     }
 
-    assertEquals(0, consumer0.assignment().size)
+    assertEquals(0, consumer.assignment().size)
 
-    consumer0.subscribe(List(topic1, topic2, topic3).asJava)
+    consumer.subscribe(List(topic1, topic2, topic3).asJava)
 
-    TestUtils.waitUntilTrue(() => {
-      consumer0.poll(50)
-      consumer0.assignment() == partitions.toSet.asJava
-    }, s"Expected partitions ${partitions.asJava} but actually got ${consumer0.assignment}")
+    awaitAssignment(consumer, partitions.toSet)
 
-    val producerRecords = partitions.flatMap(sendRecords(partitionCount, _))
-    val consumerRecords = consumeRecords(consumer0, producerRecords.size)
+    val producer = createProducer()
+    val producerRecords = partitions.flatMap(sendRecords(producer, numRecords = 15, _))
+    val consumerRecords = consumeRecords(consumer, producerRecords.size)
 
     val expected = producerRecords.map { record =>
       (record.topic, record.partition, new String(record.key), new String(record.value), record.timestamp)
@@ -819,75 +816,67 @@ 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)
-    val consumer0 = new KafkaConsumer(this.consumerConfig, new ByteArrayDeserializer(), new ByteArrayDeserializer())
-    consumers += consumer0
+    val consumer = createConsumer()
 
     // create two new topics, each having 2 partitions
     val topic1 = "topic1"
     val topic2 = "topic2"
-    val expectedAssignment = createTopicAndSendRecords(topic1, 2, 100) ++ createTopicAndSendRecords(topic2, 2, 100)
+    val producer = createProducer()
+    val expectedAssignment = createTopicAndSendRecords(producer, topic1, 2, 100) ++
+      createTopicAndSendRecords(producer, topic2, 2, 100)
 
-    assertEquals(0, consumer0.assignment().size)
+    assertEquals(0, consumer.assignment().size)
 
     // subscribe to two topics
-    consumer0.subscribe(List(topic1, topic2).asJava)
-    TestUtils.waitUntilTrue(() => {
-      consumer0.poll(50)
-      consumer0.assignment() == expectedAssignment.asJava
-    }, s"Expected partitions ${expectedAssignment.asJava} but actually got ${consumer0.assignment()}")
+    consumer.subscribe(List(topic1, topic2).asJava)
+    awaitAssignment(consumer, expectedAssignment)
 
     // add one more topic with 2 partitions
     val topic3 = "topic3"
-    createTopicAndSendRecords(topic3, 2, 100)
+    createTopicAndSendRecords(producer, topic3, 2, 100)
 
     val newExpectedAssignment = expectedAssignment ++ Set(new TopicPartition(topic3, 0), new TopicPartition(topic3, 1))
-    consumer0.subscribe(List(topic1, topic2, topic3).asJava)
-    TestUtils.waitUntilTrue(() => {
-      consumer0.poll(50)
-      consumer0.assignment() == newExpectedAssignment.asJava
-    }, s"Expected partitions ${newExpectedAssignment.asJava} but actually got ${consumer0.assignment()}")
+    consumer.subscribe(List(topic1, topic2, topic3).asJava)
+    awaitAssignment(consumer, newExpectedAssignment)
 
     // remove the topic we just added
-    consumer0.subscribe(List(topic1, topic2).asJava)
-    TestUtils.waitUntilTrue(() => {
-      consumer0.poll(50)
-      consumer0.assignment() == expectedAssignment.asJava
-    }, s"Expected partitions ${expectedAssignment.asJava} but actually got ${consumer0.assignment()}")
+    consumer.subscribe(List(topic1, topic2).asJava)
+    awaitAssignment(consumer, expectedAssignment)
 
-    consumer0.unsubscribe()
-    assertEquals(0, consumer0.assignment().size)
+    consumer.unsubscribe()
+    assertEquals(0, consumer.assignment().size)
   }
 
   @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)
 
     // create two new topics, total number of partitions must be greater than number of consumers
     val topic1 = "topic1"
     val topic2 = "topic2"
-    val subscriptions = createTopicAndSendRecords(topic1, 5, 100) ++ createTopicAndSendRecords(topic2, 8, 100)
+    val producer = createProducer()
+    val subscriptions = createTopicAndSendRecords(producer, topic1, 5, 100) ++
+      createTopicAndSendRecords(producer, topic2, 8, 100)
 
     // create a group of consumers, subscribe the consumers to all the topics and start polling
     // for the topic partition assignment
-    val (_, consumerPollers) = createConsumerGroupAndWaitForAssignment(10, List(topic1, topic2), subscriptions)
+    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, consumers, consumerPollers, List(topic1, topic2), subscriptions)
+      addConsumersToGroupAndWaitForGroupAssignment(1, consumerGroup, consumerPollers,
+        List(topic1, topic2), subscriptions, "roundrobin-group")
     } finally {
       consumerPollers.foreach(_.shutdown())
     }
   }
 
-  def reverse(m: Map[Long, Set[TopicPartition]]) =
-    m.values.toSet.flatten.map(v => (v, m.keys.filter(m(_).contains(v)).head)).toMap
-
   /**
    * This test runs the following scenario to verify sticky assignor behavior.
    * Topics: single-topic, with random number of partitions, where #par is 10, 20, 30, 40, 50, 60, 70, 80, 90, or 100
@@ -900,24 +889,28 @@ class PlaintextConsumerTest extends BaseConsumerTest {
    *    will move to consumer #10, leading to a total of (#par mod 9) partition movement
    */
   @Test
-  def testMultiConsumerStickyAssignment() {
-    this.consumers.clear()
+  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
+
     this.consumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "sticky-group")
     this.consumerConfig.setProperty(ConsumerConfig.PARTITION_ASSIGNMENT_STRATEGY_CONFIG, classOf[StickyAssignor].getName)
 
     // create one new topic
     val topic = "single-topic"
     val rand = 1 + scala.util.Random.nextInt(10)
-    val partitions = createTopicAndSendRecords(topic, rand * 10, 100)
+    val producer = createProducer()
+    val partitions = createTopicAndSendRecords(producer, topic, rand * 10, 100)
 
     // create a group of consumers, subscribe the consumers to the single topic and start polling
     // for the topic partition assignment
-    val (_, consumerPollers) = createConsumerGroupAndWaitForAssignment(9, List(topic), partitions)
-    validateGroupAssignment(consumerPollers, partitions, s"Did not get valid initial assignment for partitions ${partitions.asJava}")
+    val (consumerGroup, consumerPollers) = createConsumerGroupAndWaitForAssignment(9, List(topic), partitions)
+    validateGroupAssignment(consumerPollers, partitions)
     val prePartition2PollerId = reverse(consumerPollers.map(poller => (poller.getId, poller.consumerAssignment())).toMap)
 
     // add one more consumer and validate re-assignment
-    addConsumersToGroupAndWaitForGroupAssignment(1, consumers, 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)
@@ -942,25 +935,30 @@ 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
-    sendRecords(100, tp)
-    sendRecords(100, tp2)
+    val producer = createProducer()
+    sendRecords(producer, numRecords = 100, tp)
+    sendRecords(producer, numRecords = 100, tp2)
     val topic1 = "topic1"
-    val subscriptions = Set(tp, tp2) ++ createTopicAndSendRecords(topic1, 5, 100)
+    val subscriptions = Set(tp, tp2) ++ createTopicAndSendRecords(producer, topic1, 5, 100)
 
     // subscribe all consumers to all topics and validate the assignment
-    val consumerPollers = subscribeConsumers(consumers, List(topic, topic1))
 
+    val consumersInGroup = Buffer[KafkaConsumer[Array[Byte], Array[Byte]]]()
+    consumersInGroup += createConsumer()
+    consumersInGroup += createConsumer()
+
+    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, consumers, consumerPollers, List(topic, topic1), subscriptions)
+      addConsumersToGroupAndWaitForGroupAssignment(2, consumersInGroup, consumerPollers, List(topic, topic1), subscriptions)
 
       // add one more topic and validate partition re-assignment
       val topic2 = "topic2"
-      val expandedSubscriptions = subscriptions ++ createTopicAndSendRecords(topic2, 3, 100)
+      val expandedSubscriptions = subscriptions ++ createTopicAndSendRecords(producer, topic2, 3, 100)
       changeConsumerGroupSubscriptionAndValidateAssignment(consumerPollers, List(topic, topic1, topic2), expandedSubscriptions)
 
       // remove the topic we just added and validate re-assignment
@@ -982,17 +980,18 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   @Test
-  def testInterceptors() {
+  def testInterceptors(): Unit = {
     val appendStr = "mock"
     MockConsumerInterceptor.resetCounters()
     MockProducerInterceptor.resetCounters()
 
     // create producer with interceptor
     val producerProps = new Properties()
-    producerProps.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList)
-    producerProps.put(ProducerConfig.INTERCEPTOR_CLASSES_CONFIG, "org.apache.kafka.test.MockProducerInterceptor")
+    producerProps.put(ProducerConfig.INTERCEPTOR_CLASSES_CONFIG, classOf[MockProducerInterceptor].getName)
     producerProps.put("mock.interceptor.append", appendStr)
-    val testProducer = new KafkaProducer(producerProps, new StringSerializer, new StringSerializer)
+    val testProducer = createProducer(keySerializer = new StringSerializer,
+      valueSerializer = new StringSerializer,
+      configOverrides = producerProps)
 
     // produce records
     val numRecords = 10
@@ -1013,7 +1012,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
 
     // create consumer with interceptor
     this.consumerConfig.setProperty(ConsumerConfig.INTERCEPTOR_CLASSES_CONFIG, "org.apache.kafka.test.MockConsumerInterceptor")
-    val testConsumer = new KafkaConsumer(this.consumerConfig, new StringDeserializer, new StringDeserializer)
+    val testConsumer = createConsumer(keyDeserializer = new StringDeserializer, valueDeserializer = new StringDeserializer)
     testConsumer.assign(List(tp).asJava)
     testConsumer.seek(tp, 0)
 
@@ -1028,14 +1027,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
-    val commitCallback = new CountConsumerCommitCallback()
-    testConsumer.commitAsync(Map[TopicPartition, OffsetAndMetadata]((tp, new OffsetAndMetadata(5L))).asJava, commitCallback)
-    awaitCommitCallback(testConsumer, commitCallback)
-    assertEquals(5, testConsumer.committed(tp).offset)
+    sendAndAwaitAsyncCommit(testConsumer, Some(Map(tp -> new OffsetAndMetadata(5L))))
+    assertEquals(5, testConsumer.committed(Set(tp).asJava).get(tp).offset)
     assertEquals(commitCountBefore + 2, MockConsumerInterceptor.ON_COMMIT_COUNT.intValue)
 
     testConsumer.close()
@@ -1047,13 +1044,13 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   @Test
-  def testAutoCommitIntercept() {
+  def testAutoCommitIntercept(): Unit = {
     val topic2 = "topic2"
-    TestUtils.createTopic(this.zkUtils, topic2, 2, serverCount, this.servers)
+    createTopic(topic2, 2, brokerCount)
 
     // produce records
     val numRecords = 100
-    val testProducer = new KafkaProducer[String, String](this.producerConfig, new StringSerializer, new StringSerializer)
+    val testProducer = createProducer(keySerializer = new StringSerializer, valueSerializer = new StringSerializer)
     (0 until numRecords).map { i =>
       testProducer.send(new ProducerRecord(tp.topic(), tp.partition(), s"key $i", s"value $i"))
     }.foreach(_.get)
@@ -1061,7 +1058,7 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     // create consumer with interceptor
     this.consumerConfig.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true")
     this.consumerConfig.setProperty(ConsumerConfig.INTERCEPTOR_CLASSES_CONFIG, "org.apache.kafka.test.MockConsumerInterceptor")
-    val testConsumer = new KafkaConsumer[String, String](this.consumerConfig, new StringDeserializer(), new StringDeserializer())
+    val testConsumer = createConsumer(keyDeserializer = new StringDeserializer, valueDeserializer = new StringDeserializer)
     val rebalanceListener = new ConsumerRebalanceListener {
       override def onPartitionsAssigned(partitions: util.Collection[TopicPartition]) = {
         // keep partitions paused in this test so that we can verify the commits based on specific seeks
@@ -1082,8 +1079,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
@@ -1097,23 +1094,21 @@ 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()
     producerProps.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList)
     producerProps.put(ProducerConfig.INTERCEPTOR_CLASSES_CONFIG, "org.apache.kafka.test.MockProducerInterceptor")
     producerProps.put("mock.interceptor.append", appendStr)
-    val testProducer = new KafkaProducer(producerProps, new ByteArraySerializer(), new ByteArraySerializer())
-    producers += testProducer
+    val testProducer = createProducer()
 
     // producing records should succeed
     testProducer.send(new ProducerRecord(tp.topic(), tp.partition(), s"key".getBytes, s"value will not be modified".getBytes))
 
     // create consumer with interceptor that has different key and value types from the consumer
     this.consumerConfig.setProperty(ConsumerConfig.INTERCEPTOR_CLASSES_CONFIG, "org.apache.kafka.test.MockConsumerInterceptor")
-    val testConsumer = new KafkaConsumer[Array[Byte], Array[Byte]](this.consumerConfig, new ByteArrayDeserializer(), new ByteArrayDeserializer())
-    consumers += testConsumer
+    val testConsumer = createConsumer()
 
     testConsumer.assign(List(tp).asJava)
     testConsumer.seek(tp, 0)
@@ -1125,57 +1120,63 @@ class PlaintextConsumerTest extends BaseConsumerTest {
   }
 
   @Test
-  def testConsumeMessagesWithCreateTime() {
+  def testConsumeMessagesWithCreateTime(): Unit = {
     val numRecords = 50
     // Test non-compressed messages
-    sendRecords(numRecords, tp)
-    this.consumers.head.assign(List(tp).asJava)
-    consumeAndVerifyRecords(consumer = this.consumers.head, numRecords = numRecords, startingOffset = 0, startingKeyAndValueIndex = 0,
+    val producer = createProducer()
+    sendRecords(producer, numRecords, tp)
+    val consumer = createConsumer()
+    consumer.assign(List(tp).asJava)
+    consumeAndVerifyRecords(consumer = consumer, numRecords = numRecords, startingOffset = 0, startingKeyAndValueIndex = 0,
       startingTimestamp = 0)
 
     // Test compressed messages
     sendCompressedMessages(numRecords, tp2)
-    this.consumers.head.assign(List(tp2).asJava)
-    consumeAndVerifyRecords(consumer = this.consumers.head, numRecords = numRecords, tp = tp2, startingOffset = 0, startingKeyAndValueIndex = 0,
+    consumer.assign(List(tp2).asJava)
+    consumeAndVerifyRecords(consumer = consumer, numRecords = numRecords, tp = tp2, startingOffset = 0, startingKeyAndValueIndex = 0,
       startingTimestamp = 0)
   }
 
   @Test
-  def testConsumeMessagesWithLogAppendTime() {
+  def testConsumeMessagesWithLogAppendTime(): Unit = {
     val topicName = "testConsumeMessagesWithLogAppendTime"
     val topicProps = new Properties()
     topicProps.setProperty(LogConfig.MessageTimestampTypeProp, "LogAppendTime")
-    TestUtils.createTopic(zkUtils, topicName, 2, 2, servers, topicProps)
+    createTopic(topicName, 2, 2, topicProps)
 
     val startTime = System.currentTimeMillis()
     val numRecords = 50
 
     // Test non-compressed messages
     val tp1 = new TopicPartition(topicName, 0)
-    sendRecords(numRecords, tp1)
-    this.consumers.head.assign(List(tp1).asJava)
-    consumeAndVerifyRecords(consumer = this.consumers.head, numRecords = numRecords, tp = tp1, startingOffset = 0, startingKeyAndValueIndex = 0,
+    val producer = createProducer()
+    sendRecords(producer, numRecords, tp1)
+
+    val consumer = createConsumer()
+    consumer.assign(List(tp1).asJava)
+    consumeAndVerifyRecords(consumer = consumer, numRecords = numRecords, tp = tp1, startingOffset = 0, startingKeyAndValueIndex = 0,
       startingTimestamp = startTime, timestampType = TimestampType.LOG_APPEND_TIME)
 
     // Test compressed messages
     val tp2 = new TopicPartition(topicName, 1)
     sendCompressedMessages(numRecords, tp2)
-    this.consumers.head.assign(List(tp2).asJava)
-    consumeAndVerifyRecords(consumer = this.consumers.head, numRecords = numRecords, tp = tp2, startingOffset = 0, startingKeyAndValueIndex = 0,
+    consumer.assign(List(tp2).asJava)
+    consumeAndVerifyRecords(consumer = consumer, numRecords = numRecords, tp = tp2, startingOffset = 0, startingKeyAndValueIndex = 0,
       startingTimestamp = startTime, timestampType = TimestampType.LOG_APPEND_TIME)
   }
 
   @Test
-  def testListTopics() {
+  def testListTopics(): Unit = {
     val numParts = 2
     val topic1 = "part-test-topic-1"
     val topic2 = "part-test-topic-2"
     val topic3 = "part-test-topic-3"
-    TestUtils.createTopic(this.zkUtils, topic1, numParts, 1, this.servers)
-    TestUtils.createTopic(this.zkUtils, topic2, numParts, 1, this.servers)
-    TestUtils.createTopic(this.zkUtils, topic3, numParts, 1, this.servers)
+    createTopic(topic1, numParts, 1)
+    createTopic(topic2, numParts, 1)
+    createTopic(topic3, numParts, 1)
 
-    val topics = this.consumers.head.listTopics()
+    val consumer = createConsumer()
+    val topics = consumer.listTopics()
     assertNotNull(topics)
     assertEquals(5, topics.size())
     assertEquals(5, topics.keySet().size())
@@ -1185,31 +1186,32 @@ 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"
     val topic3 = "part-test-topic-3"
     val props = new Properties()
     props.setProperty(LogConfig.MessageFormatVersionProp, "0.9.0")
-    TestUtils.createTopic(this.zkUtils, topic1, numParts, 1, this.servers)
+    createTopic(topic1, numParts, 1)
     // Topic2 is in old message format.
-    TestUtils.createTopic(this.zkUtils, topic2, numParts, 1, this.servers, props)
-    TestUtils.createTopic(this.zkUtils, topic3, numParts, 1, this.servers)
+    createTopic(topic2, numParts, 1, props)
+    createTopic(topic3, numParts, 1)
 
-    val consumer = this.consumers.head
+    val consumer = createConsumer()
 
     // Test negative target time
     intercept[IllegalArgumentException](
       consumer.offsetsForTimes(Collections.singletonMap(new TopicPartition(topic1, 0), -1)))
 
+    val producer = createProducer()
     val timestampsToSearch = new util.HashMap[TopicPartition, java.lang.Long]()
     var i = 0
     for (topic <- List(topic1, topic2, topic3)) {
       for (part <- 0 until numParts) {
         val tp = new TopicPartition(topic, part)
         // In sendRecords(), each message will have key, value and timestamp equal to the sequence number.
-        sendRecords(100, tp)
+        sendRecords(producer, numRecords = 100, tp)
         timestampsToSearch.put(tp, i * 20)
         i += 1
       }
@@ -1222,34 +1224,46 @@ class PlaintextConsumerTest extends BaseConsumerTest {
     //  topic3Partition0 -> 80,
     //  topic3Partition1 -> 100)
     val timestampOffsets = consumer.offsetsForTimes(timestampsToSearch)
-    assertEquals(0, timestampOffsets.get(new TopicPartition(topic1, 0)).offset())
-    assertEquals(0, timestampOffsets.get(new TopicPartition(topic1, 0)).timestamp())
-    assertEquals(20, timestampOffsets.get(new TopicPartition(topic1, 1)).offset())
-    assertEquals(20, timestampOffsets.get(new TopicPartition(topic1, 1)).timestamp())
+
+    val timestampTopic1P0 = timestampOffsets.get(new TopicPartition(topic1, 0))
+    assertEquals(0, timestampTopic1P0.offset)
+    assertEquals(0, timestampTopic1P0.timestamp)
+    assertEquals(Optional.of(0), timestampTopic1P0.leaderEpoch)
+
+    val timestampTopic1P1 = timestampOffsets.get(new TopicPartition(topic1, 1))
+    assertEquals(20, timestampTopic1P1.offset)
+    assertEquals(20, timestampTopic1P1.timestamp)
+    assertEquals(Optional.of(0), timestampTopic1P1.leaderEpoch)
+
     assertEquals("null should be returned when message format is 0.9.0",
       null, timestampOffsets.get(new TopicPartition(topic2, 0)))
     assertEquals("null should be returned when message format is 0.9.0",
       null, timestampOffsets.get(new TopicPartition(topic2, 1)))
-    assertEquals(80, timestampOffsets.get(new TopicPartition(topic3, 0)).offset())
-    assertEquals(80, timestampOffsets.get(new TopicPartition(topic3, 0)).timestamp())
+
+    val timestampTopic3P0 = timestampOffsets.get(new TopicPartition(topic3, 0))
+    assertEquals(80, timestampTopic3P0.offset)
+    assertEquals(80, timestampTopic3P0.timestamp)
+    assertEquals(Optional.of(0), timestampTopic3P0.leaderEpoch)
+
     assertEquals(null, timestampOffsets.get(new TopicPartition(topic3, 1)))
   }
 
   @Test
-  def testEarliestOrLatestOffsets() {
+  def testEarliestOrLatestOffsets(): Unit = {
     val topic0 = "topicWithNewMessageFormat"
     val topic1 = "topicWithOldMessageFormat"
-    createTopicAndSendRecords(topicName = topic0, numPartitions = 2, recordsPerPartition = 100)
+    val producer = createProducer()
+    createTopicAndSendRecords(producer, topicName = topic0, numPartitions = 2, recordsPerPartition = 100)
     val props = new Properties()
     props.setProperty(LogConfig.MessageFormatVersionProp, "0.9.0")
-    TestUtils.createTopic(this.zkUtils, topic1, numPartitions = 1, replicationFactor = 1, this.servers, props)
-    sendRecords(100, new TopicPartition(topic1, 0))
+    createTopic(topic1, numPartitions = 1, replicationFactor = 1, props)
+    sendRecords(producer, numRecords = 100, new TopicPartition(topic1, 0))
 
     val t0p0 = new TopicPartition(topic0, 0)
     val t0p1 = new TopicPartition(topic0, 1)
     val t1p0 = new TopicPartition(topic1, 0)
     val partitions = Set(t0p0, t0p1, t1p0).asJava
-    val consumer = this.consumers.head
+    val consumer = createConsumer()
 
     val earliests = consumer.beginningOffsets(partitions)
     assertEquals(0L, earliests.get(t0p0))
@@ -1263,233 +1277,332 @@ 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 consumer0 = new KafkaConsumer(this.consumerConfig, new ByteArrayDeserializer(), new ByteArrayDeserializer())
-    consumers += consumer0
+    val consumer = createConsumer()
 
     val listener = new TestConsumerReassignmentListener()
-    consumer0.subscribe(List(topic).asJava, listener)
+    consumer.subscribe(List(topic).asJava, listener)
 
     // the initial subscription should cause a callback execution
-    while (listener.callsToAssigned == 0)
-      consumer0.poll(50)
+    awaitRebalance(consumer, listener)
 
-    consumer0.subscribe(List[String]().asJava)
-    assertEquals(0, consumer0.assignment.size())
+    consumer.subscribe(List[String]().asJava)
+    assertEquals(0, consumer.assignment.size())
   }
 
   @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 consumer0 = new KafkaConsumer(this.consumerConfig, new ByteArrayDeserializer(), new ByteArrayDeserializer())
-    consumers += consumer0
+    val consumer = createConsumer()
 
-    sendRecords(5)
-    consumer0.subscribe(List(topic).asJava)
-    consumeAndVerifyRecords(consumer = consumer0, numRecords = 5, startingOffset = 0)
-    consumer0.pause(List(tp).asJava)
+    val producer = createProducer()
+    sendRecords(producer, numRecords = 5, tp)
+    consumer.subscribe(List(topic).asJava)
+    consumeAndVerifyRecords(consumer = consumer, numRecords = 5, startingOffset = 0)
+    consumer.pause(List(tp).asJava)
 
     // subscribe to a new topic to trigger a rebalance
-    consumer0.subscribe(List("topic2").asJava)
+    consumer.subscribe(List("topic2").asJava)
 
     // after rebalance, our position should be reset and our pause state lost,
     // so we should be able to consume from the beginning
-    consumeAndVerifyRecords(consumer = consumer0, numRecords = 0, startingOffset = 5)
+    consumeAndVerifyRecords(consumer = consumer, numRecords = 0, startingOffset = 5)
   }
 
   @Test
-  def testCommitSpecifiedOffsets() {
-    sendRecords(5, tp)
-    sendRecords(7, tp2)
+  def testCommitSpecifiedOffsets(): Unit = {
+    val producer = createProducer()
+    sendRecords(producer, numRecords = 5, tp)
+    sendRecords(producer, numRecords = 7, tp2)
 
-    this.consumers.head.assign(List(tp, tp2).asJava)
+    val consumer = createConsumer()
+    consumer.assign(List(tp, tp2).asJava)
 
-    // Need to poll to join the group
-    this.consumers.head.poll(50)
-    val pos1 = this.consumers.head.position(tp)
-    val pos2 = this.consumers.head.position(tp2)
-    this.consumers.head.commitSync(Map[TopicPartition, OffsetAndMetadata]((tp, new OffsetAndMetadata(3L))).asJava)
-    assertEquals(3, this.consumers.head.committed(tp).offset)
-    assertNull(this.consumers.head.committed(tp2))
+    val pos1 = consumer.position(tp)
+    val pos2 = consumer.position(tp2)
+    consumer.commitSync(Map[TopicPartition, OffsetAndMetadata]((tp, new OffsetAndMetadata(3L))).asJava)
+    assertEquals(3, consumer.committed(Set(tp).asJava).get(tp).offset)
+    assertNull(consumer.committed(Set(tp2).asJava).get(tp2))
 
     // Positions should not change
-    assertEquals(pos1, this.consumers.head.position(tp))
-    assertEquals(pos2, this.consumers.head.position(tp2))
-    this.consumers.head.commitSync(Map[TopicPartition, OffsetAndMetadata]((tp2, new OffsetAndMetadata(5L))).asJava)
-    assertEquals(3, this.consumers.head.committed(tp).offset)
-    assertEquals(5, this.consumers.head.committed(tp2).offset)
+    assertEquals(pos1, consumer.position(tp))
+    assertEquals(pos2, consumer.position(tp2))
+    consumer.commitSync(Map[TopicPartition, OffsetAndMetadata]((tp2, new OffsetAndMetadata(5L))).asJava)
+    assertEquals(3, consumer.committed(Set(tp).asJava).get(tp).offset)
+    assertEquals(5, consumer.committed(Set(tp2).asJava).get(tp2).offset)
 
     // Using async should pick up the committed changes after commit completes
-    val commitCallback = new CountConsumerCommitCallback()
-    this.consumers.head.commitAsync(Map[TopicPartition, OffsetAndMetadata]((tp2, new OffsetAndMetadata(7L))).asJava, commitCallback)
-    awaitCommitCallback(this.consumers.head, commitCallback)
-    assertEquals(7, this.consumers.head.committed(tp2).offset)
+    sendAndAwaitAsyncCommit(consumer, Some(Map(tp2 -> new OffsetAndMetadata(7L))))
+    assertEquals(7, consumer.committed(Set(tp2).asJava).get(tp2).offset)
   }
 
   @Test
-  def testAutoCommitOnRebalance() {
+  def testAutoCommitOnRebalance(): Unit = {
     val topic2 = "topic2"
-    TestUtils.createTopic(this.zkUtils, topic2, 2, serverCount, this.servers)
+    createTopic(topic2, 2, brokerCount)
 
     this.consumerConfig.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true")
-    val consumer0 = new KafkaConsumer(this.consumerConfig, new ByteArrayDeserializer(), new ByteArrayDeserializer())
-    consumers += consumer0
+    val consumer = createConsumer()
 
     val numRecords = 10000
-    sendRecords(numRecords)
+    val producer = createProducer()
+    sendRecords(producer, numRecords, tp)
 
     val rebalanceListener = new ConsumerRebalanceListener {
       override def onPartitionsAssigned(partitions: util.Collection[TopicPartition]) = {
         // keep partitions paused in this test so that we can verify the commits based on specific seeks
-        consumer0.pause(partitions)
+        consumer.pause(partitions)
       }
 
       override def onPartitionsRevoked(partitions: util.Collection[TopicPartition]) = {}
     }
 
-    consumer0.subscribe(List(topic).asJava, rebalanceListener)
+    consumer.subscribe(List(topic).asJava, rebalanceListener)
 
-    val assignment = Set(tp, tp2)
-    TestUtils.waitUntilTrue(() => {
-      consumer0.poll(50)
-      consumer0.assignment() == assignment.asJava
-    }, s"Expected partitions ${assignment.asJava} but actually got ${consumer0.assignment()}")
+    awaitAssignment(consumer, Set(tp, tp2))
 
-    consumer0.seek(tp, 300)
-    consumer0.seek(tp2, 500)
+    consumer.seek(tp, 300)
+    consumer.seek(tp2, 500)
 
     // change subscription to trigger rebalance
-    consumer0.subscribe(List(topic, topic2).asJava, rebalanceListener)
+    consumer.subscribe(List(topic, topic2).asJava, rebalanceListener)
 
     val newAssignment = Set(tp, tp2, new TopicPartition(topic2, 0), new TopicPartition(topic2, 1))
-    TestUtils.waitUntilTrue(() => {
-      consumer0.poll(50)
-      consumer0.assignment() == newAssignment.asJava
-    }, s"Expected partitions ${newAssignment.asJava} but actually got ${consumer0.assignment()}")
+    awaitAssignment(consumer, newAssignment)
 
     // after rebalancing, we should have reset to the committed positions
-    assertEquals(300, consumer0.committed(tp).offset)
-    assertEquals(500, consumer0.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(): Unit = {
+    val numMessages = 1000
+    val topic2 = "topic2"
+    createTopic(topic2, 2, brokerCount)
+    // send some messages.
+    val producer = createProducer()
+    sendRecords(producer, numMessages, tp)
+    // Test subscribe
+    // Create a consumer and consumer some messages.
+    consumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "testPerPartitionLeadMetricsCleanUpWithSubscribe")
+    consumerConfig.setProperty(ConsumerConfig.CLIENT_ID_CONFIG, "testPerPartitionLeadMetricsCleanUpWithSubscribe")
+    val consumer = createConsumer()
+    val listener = new TestConsumerReassignmentListener
+    consumer.subscribe(List(topic, topic2).asJava, listener)
+    val records = awaitNonEmptyRecords(consumer, tp)
+    assertEquals("should be assigned once", 1, listener.callsToAssigned)
+    // Verify the metric exist.
+    val tags1 = new util.HashMap[String, String]()
+    tags1.put("client-id", "testPerPartitionLeadMetricsCleanUpWithSubscribe")
+    tags1.put("topic", tp.topic())
+    tags1.put("partition", String.valueOf(tp.partition()))
+
+    val tags2 = new util.HashMap[String, String]()
+    tags2.put("client-id", "testPerPartitionLeadMetricsCleanUpWithSubscribe")
+    tags2.put("topic", tp2.topic())
+    tags2.put("partition", String.valueOf(tp2.partition()))
+    val fetchLead0 = consumer.metrics.get(new MetricName("records-lead", "consumer-fetch-manager-metrics", "", tags1))
+    assertNotNull(fetchLead0)
+    assertEquals(s"The lead should be ${records.count}", records.count.toDouble, fetchLead0.metricValue())
+
+    // Remove topic from subscription
+    consumer.subscribe(List(topic2).asJava, listener)
+    awaitRebalance(consumer, listener)
+    // Verify the metric has gone
+    assertNull(consumer.metrics.get(new MetricName("records-lead", "consumer-fetch-manager-metrics", "", tags1)))
+    assertNull(consumer.metrics.get(new MetricName("records-lead", "consumer-fetch-manager-metrics", "", tags2)))
+  }
 
   @Test
-  def testPerPartitionLagMetricsCleanUpWithSubscribe() {
+  def testPerPartitionLagMetricsCleanUpWithSubscribe(): Unit = {
     val numMessages = 1000
     val topic2 = "topic2"
-    TestUtils.createTopic(this.zkUtils, topic2, 2, serverCount, this.servers)
+    createTopic(topic2, 2, brokerCount)
     // send some messages.
-    sendRecords(numMessages, tp)
+    val producer = createProducer()
+    sendRecords(producer, numMessages, tp)
     // Test subscribe
     // Create a consumer and consumer some messages.
     consumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "testPerPartitionLagMetricsCleanUpWithSubscribe")
     consumerConfig.setProperty(ConsumerConfig.CLIENT_ID_CONFIG, "testPerPartitionLagMetricsCleanUpWithSubscribe")
-    val consumer = new KafkaConsumer(this.consumerConfig, new ByteArrayDeserializer(), new ByteArrayDeserializer())
-    try {
-      val listener0 = new TestConsumerReassignmentListener
-      consumer.subscribe(List(topic, topic2).asJava, listener0)
-      var records: ConsumerRecords[Array[Byte], Array[Byte]] = ConsumerRecords.empty()
-      TestUtils.waitUntilTrue(() => {
-        records = consumer.poll(100)
-        !records.records(tp).isEmpty
-      }, "Consumer did not consume any message before timeout.")
-      assertEquals("should be assigned once", 1, listener0.callsToAssigned)
-      // Verify the metric exist.
-      val tags = Collections.singletonMap("client-id", "testPerPartitionLagMetricsCleanUpWithSubscribe")
-      val fetchLag0 = consumer.metrics.get(new MetricName(tp + ".records-lag", "consumer-fetch-manager-metrics", "", tags))
-      assertNotNull(fetchLag0)
-      val expectedLag = numMessages - records.count
-      assertEquals(s"The lag should be $expectedLag", expectedLag, fetchLag0.value, epsilon)
-
-      // Remove topic from subscription
-      consumer.subscribe(List(topic2).asJava, listener0)
-      TestUtils.waitUntilTrue(() => {
-        consumer.poll(100)
-        listener0.callsToAssigned >= 2
-      }, "Expected rebalance did not occur.")
-      // Verify the metric has gone
-      assertNull(consumer.metrics.get(new MetricName(tp + ".records-lag", "consumer-fetch-manager-metrics", "", tags)))
-      assertNull(consumer.metrics.get(new MetricName(tp2 + ".records-lag", "consumer-fetch-manager-metrics", "", tags)))
-    } finally {
-      consumer.close()
-    }
+    val consumer = createConsumer()
+    val listener = new TestConsumerReassignmentListener
+    consumer.subscribe(List(topic, topic2).asJava, listener)
+    val records = awaitNonEmptyRecords(consumer, tp)
+    assertEquals("should be assigned once", 1, listener.callsToAssigned)
+    // Verify the metric exist.
+    val tags1 = new util.HashMap[String, String]()
+    tags1.put("client-id", "testPerPartitionLagMetricsCleanUpWithSubscribe")
+    tags1.put("topic", tp.topic())
+    tags1.put("partition", String.valueOf(tp.partition()))
+
+    val tags2 = new util.HashMap[String, String]()
+    tags2.put("client-id", "testPerPartitionLagMetricsCleanUpWithSubscribe")
+    tags2.put("topic", tp2.topic())
+    tags2.put("partition", String.valueOf(tp2.partition()))
+    val fetchLag0 = consumer.metrics.get(new MetricName("records-lag", "consumer-fetch-manager-metrics", "", tags1))
+    assertNotNull(fetchLag0)
+    val expectedLag = numMessages - records.count
+    assertEquals(s"The lag should be $expectedLag", expectedLag, fetchLag0.metricValue.asInstanceOf[Double], epsilon)
+
+    // Remove topic from subscription
+    consumer.subscribe(List(topic2).asJava, listener)
+    awaitRebalance(consumer, listener)
+    // Verify the metric has gone
+    assertNull(consumer.metrics.get(new MetricName("records-lag", "consumer-fetch-manager-metrics", "", tags1)))
+    assertNull(consumer.metrics.get(new MetricName("records-lag", "consumer-fetch-manager-metrics", "", tags2)))
   }
 
   @Test
-  def testPerPartitionLagMetricsCleanUpWithAssign() {
+  def testPerPartitionLeadMetricsCleanUpWithAssign(): Unit = {
     val numMessages = 1000
     // Test assign
     // send some messages.
-    sendRecords(numMessages, tp)
-    sendRecords(numMessages, tp2)
+    val producer = createProducer()
+    sendRecords(producer, numMessages, tp)
+    sendRecords(producer, numMessages, tp2)
+
+    consumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "testPerPartitionLeadMetricsCleanUpWithAssign")
+    consumerConfig.setProperty(ConsumerConfig.CLIENT_ID_CONFIG, "testPerPartitionLeadMetricsCleanUpWithAssign")
+    val consumer = createConsumer()
+    consumer.assign(List(tp).asJava)
+    val records = awaitNonEmptyRecords(consumer, tp)
+    // Verify the metric exist.
+    val tags = new util.HashMap[String, String]()
+    tags.put("client-id", "testPerPartitionLeadMetricsCleanUpWithAssign")
+    tags.put("topic", tp.topic())
+    tags.put("partition", String.valueOf(tp.partition()))
+    val fetchLead = consumer.metrics.get(new MetricName("records-lead", "consumer-fetch-manager-metrics", "", tags))
+    assertNotNull(fetchLead)
+
+    assertTrue(s"The lead should be ${records.count}", records.count == fetchLead.metricValue())
+
+    consumer.assign(List(tp2).asJava)
+    awaitNonEmptyRecords(consumer ,tp2)
+    assertNull(consumer.metrics.get(new MetricName("records-lead", "consumer-fetch-manager-metrics", "", tags)))
+  }
+
+  @Test
+  def testPerPartitionLagMetricsCleanUpWithAssign(): Unit = {
+    val numMessages = 1000
+    // Test assign
+    // send some messages.
+    val producer = createProducer()
+    sendRecords(producer, numMessages, tp)
+    sendRecords(producer, numMessages, tp2)
+
     consumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "testPerPartitionLagMetricsCleanUpWithAssign")
     consumerConfig.setProperty(ConsumerConfig.CLIENT_ID_CONFIG, "testPerPartitionLagMetricsCleanUpWithAssign")
-    val consumer = new KafkaConsumer(this.consumerConfig, new ByteArrayDeserializer(), new ByteArrayDeserializer())
-    try {
-      consumer.assign(List(tp).asJava)
-      var records: ConsumerRecords[Array[Byte], Array[Byte]] = ConsumerRecords.empty()
-      TestUtils.waitUntilTrue(() => {
-        records = consumer.poll(100)
-        !records.records(tp).isEmpty
-      }, "Consumer did not consume any message before timeout.")
-      // Verify the metric exist.
-      val tags = Collections.singletonMap("client-id", "testPerPartitionLagMetricsCleanUpWithAssign")
-      val fetchLag = consumer.metrics.get(new MetricName(tp + ".records-lag", "consumer-fetch-manager-metrics", "", tags))
-      assertNotNull(fetchLag)
-      val expectedLag = numMessages - records.count
-      assertEquals(s"The lag should be $expectedLag", expectedLag, fetchLag.value, epsilon)
-
-      consumer.assign(List(tp2).asJava)
-      TestUtils.waitUntilTrue(() => !consumer.poll(100).isEmpty, "Consumer did not consume any message before timeout.")
-      assertNull(consumer.metrics.get(new MetricName(tp + ".records-lag", "consumer-fetch-manager-metrics", "", tags)))
-    } finally {
-      consumer.close()
-    }
+    val consumer = createConsumer()
+    consumer.assign(List(tp).asJava)
+    val records = awaitNonEmptyRecords(consumer, tp)
+    // Verify the metric exist.
+    val tags = new util.HashMap[String, String]()
+    tags.put("client-id", "testPerPartitionLagMetricsCleanUpWithAssign")
+    tags.put("topic", tp.topic())
+    tags.put("partition", String.valueOf(tp.partition()))
+    val fetchLag = consumer.metrics.get(new MetricName("records-lag", "consumer-fetch-manager-metrics", "", tags))
+    assertNotNull(fetchLag)
+
+    val expectedLag = numMessages - records.count
+    assertEquals(s"The lag should be $expectedLag", expectedLag, fetchLag.metricValue.asInstanceOf[Double], epsilon)
+
+    consumer.assign(List(tp2).asJava)
+    awaitNonEmptyRecords(consumer, tp2)
+    assertNull(consumer.metrics.get(new MetricName(tp.toString + ".records-lag", "consumer-fetch-manager-metrics", "", tags)))
+    assertNull(consumer.metrics.get(new MetricName("records-lag", "consumer-fetch-manager-metrics", "", tags)))
   }
 
   @Test
-  def testPerPartitionLagWithMaxPollRecords() {
+  def testPerPartitionLagMetricsWhenReadCommitted(): Unit = {
+    val numMessages = 1000
+    // send some messages.
+    val producer = createProducer()
+    sendRecords(producer, numMessages, tp)
+    sendRecords(producer, numMessages, tp2)
+
+    consumerConfig.setProperty(ConsumerConfig.ISOLATION_LEVEL_CONFIG, "read_committed")
+    consumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "testPerPartitionLagMetricsCleanUpWithAssign")
+    consumerConfig.setProperty(ConsumerConfig.CLIENT_ID_CONFIG, "testPerPartitionLagMetricsCleanUpWithAssign")
+    val consumer = createConsumer()
+    consumer.assign(List(tp).asJava)
+    awaitNonEmptyRecords(consumer, tp)
+    // Verify the metric exist.
+    val tags = new util.HashMap[String, String]()
+    tags.put("client-id", "testPerPartitionLagMetricsCleanUpWithAssign")
+    tags.put("topic", tp.topic())
+    tags.put("partition", String.valueOf(tp.partition()))
+    val fetchLag = consumer.metrics.get(new MetricName("records-lag", "consumer-fetch-manager-metrics", "", tags))
+    assertNotNull(fetchLag)
+  }
+
+  @Test
+  def testPerPartitionLeadWithMaxPollRecords(): Unit = {
+    val numMessages = 1000
+    val maxPollRecords = 10
+    val producer = createProducer()
+    sendRecords(producer, numMessages, tp)
+
+    consumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "testPerPartitionLeadWithMaxPollRecords")
+    consumerConfig.setProperty(ConsumerConfig.CLIENT_ID_CONFIG, "testPerPartitionLeadWithMaxPollRecords")
+    consumerConfig.setProperty(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, maxPollRecords.toString)
+    val consumer = createConsumer()
+    consumer.assign(List(tp).asJava)
+    awaitNonEmptyRecords(consumer, tp)
+
+    val tags = new util.HashMap[String, String]()
+    tags.put("client-id", "testPerPartitionLeadWithMaxPollRecords")
+    tags.put("topic", tp.topic())
+    tags.put("partition", String.valueOf(tp.partition()))
+    val lead = consumer.metrics.get(new MetricName("records-lead", "consumer-fetch-manager-metrics", "", tags))
+    assertTrue(s"The lead should be $maxPollRecords", lead.metricValue() == maxPollRecords)
+  }
+
+  @Test
+  def testPerPartitionLagWithMaxPollRecords(): Unit = {
     val numMessages = 1000
     val maxPollRecords = 10
-    sendRecords(numMessages, tp)
+    val producer = createProducer()
+    sendRecords(producer, numMessages, tp)
+
     consumerConfig.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "testPerPartitionLagWithMaxPollRecords")
     consumerConfig.setProperty(ConsumerConfig.CLIENT_ID_CONFIG, "testPerPartitionLagWithMaxPollRecords")
     consumerConfig.setProperty(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, maxPollRecords.toString)
-    val consumer = new KafkaConsumer(this.consumerConfig, new ByteArrayDeserializer(), new ByteArrayDeserializer())
+    val consumer = createConsumer()
     consumer.assign(List(tp).asJava)
-    try {
-      var records: ConsumerRecords[Array[Byte], Array[Byte]] = ConsumerRecords.empty()
-      TestUtils.waitUntilTrue(() => {
-        records = consumer.poll(100)
-        !records.isEmpty
-      }, "Consumer did not consume any message before timeout.")
-      val tags = Collections.singletonMap("client-id", "testPerPartitionLagWithMaxPollRecords")
-      val lag = consumer.metrics.get(new MetricName(tp + ".records-lag", "consumer-fetch-manager-metrics", "", tags))
-      assertEquals(s"The lag should be ${numMessages - records.count}", numMessages - records.count, lag.value, epsilon)
-    } finally {
-      consumer.close()
-    }
+    val records = awaitNonEmptyRecords(consumer, tp)
+
+    val tags = new util.HashMap[String, String]()
+    tags.put("client-id", "testPerPartitionLagWithMaxPollRecords")
+    tags.put("topic", tp.topic())
+    tags.put("partition", String.valueOf(tp.partition()))
+    val lag = consumer.metrics.get(new MetricName("records-lag", "consumer-fetch-manager-metrics", "", tags))
+
+    assertEquals(s"The lag should be ${numMessages - records.count}", numMessages - records.count, lag.metricValue.asInstanceOf[Double], epsilon)
   }
 
   @Test
-  def testQuotaMetricsNotCreatedIfNoQuotasConfigured() {
+  def testQuotaMetricsNotCreatedIfNoQuotasConfigured(): Unit = {
     val numRecords = 1000
-    sendRecords(numRecords)
+    val producer = createProducer()
+    sendRecords(producer, numRecords, tp)
 
-    this.consumers.head.assign(List(tp).asJava)
-    this.consumers.head.seek(tp, 0)
-    consumeAndVerifyRecords(consumer = this.consumers.head, numRecords = numRecords, startingOffset = 0)
+    val consumer = createConsumer()
+    consumer.assign(List(tp).asJava)
+    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,
                                   "",
                                   "user", "",
                                   "client-id", clientId)
-        assertNull("Metric should not hanve been created " + metricName, broker.metrics.metric(metricName))
+        assertNull("Metric should not have been created " + metricName, broker.metrics.metric(metricName))
     }
     servers.foreach(assertNoMetric(_, "byte-rate", QuotaType.Produce, producerClientId))
     servers.foreach(assertNoMetric(_, "throttle-time", QuotaType.Produce, producerClientId))
@@ -1501,175 +1614,109 @@ 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))
+        assertNull("Metric should not have been created " + metricName, broker.metrics.metric(metricName))
     }
-    servers.foreach(assertNoExemptRequestMetric(_))
+    servers.foreach(assertNoExemptRequestMetric)
   }
 
   def runMultiConsumerSessionTimeoutTest(closeConsumer: Boolean): Unit = {
     // use consumers defined in this class plus one additional consumer
     // Use topic defined in this class + one additional topic
-    sendRecords(100, tp)
-    sendRecords(100, tp2)
+    val producer = createProducer()
+    sendRecords(producer, numRecords = 100, tp)
+    sendRecords(producer, numRecords = 100, tp2)
     val topic1 = "topic1"
-    val subscriptions = Set(tp, tp2) ++ createTopicAndSendRecords(topic1, 6, 100)
+    val subscriptions = Set(tp, tp2) ++ createTopicAndSendRecords(producer, topic1, 6, 100)
 
     // first subscribe consumers that are defined in this class
     val consumerPollers = Buffer[ConsumerAssignmentPoller]()
-    for (consumer <- consumers)
-      consumerPollers += subscribeConsumerAndStartPolling(consumer, List(topic, topic1))
+    consumerPollers += subscribeConsumerAndStartPolling(createConsumer(), List(topic, topic1))
+    consumerPollers += subscribeConsumerAndStartPolling(createConsumer(), List(topic, topic1))
 
     // create one more consumer and add it to the group; we will timeout this consumer
-    val timeoutConsumer = new KafkaConsumer[Array[Byte], Array[Byte]](this.consumerConfig)
-    // Close the consumer on test teardown, unless this test will manually
-    if(!closeConsumer)
-      consumers += timeoutConsumer
+    val timeoutConsumer = createConsumer()
     val timeoutPoller = subscribeConsumerAndStartPolling(timeoutConsumer, List(topic, topic1))
     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
    */
-  def createTopicAndSendRecords(topicName: String, numPartitions: Int, recordsPerPartition: Int): Set[TopicPartition] = {
-    TestUtils.createTopic(this.zkUtils, topicName, numPartitions, serverCount, this.servers)
+  def createTopicAndSendRecords(producer: KafkaProducer[Array[Byte], Array[Byte]],
+                                topicName: String,
+                                numPartitions: Int,
+                                recordsPerPartition: Int): Set[TopicPartition] = {
+    createTopic(topicName, numPartitions, brokerCount)
     var parts = Set[TopicPartition]()
     for (partition <- 0 until numPartitions) {
       val tp = new TopicPartition(topicName, partition)
-      sendRecords(recordsPerPartition, tp)
+      sendRecords(producer, recordsPerPartition, tp)
       parts = parts + tp
     }
     parts
   }
 
   /**
-   * 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]) = {
     assertTrue(consumerCount <= subscriptions.size)
     val consumerGroup = Buffer[KafkaConsumer[Array[Byte], Array[Byte]]]()
     for (_ <- 0 until consumerCount)
-      consumerGroup += new KafkaConsumer[Array[Byte], Array[Byte]](this.consumerConfig)
-    consumers ++= consumerGroup
+      consumerGroup += createConsumer()
 
     // create consumer pollers, wait for assignment and validate it
     val consumerPollers = subscribeConsumers(consumerGroup, topicsToSubscribe)
-
     (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 newConsumer = new KafkaConsumer[Array[Byte], Array[Byte]](this.consumerConfig)
-      consumerGroup += newConsumer
-      consumerPollers += subscribeConsumerAndStartPolling(newConsumer, 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 = {
@@ -1679,22 +1726,159 @@ 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],
                                                             topicsToSubscribe: List[String],
-                                                            subscriptions: Set[TopicPartition],
+                                                            expectedAssignment: Set[TopicPartition],
                                                             rebalanceListener: ConsumerRebalanceListener): Unit = {
     consumer.subscribe(topicsToSubscribe.asJava, rebalanceListener)
+    awaitAssignment(consumer, expectedAssignment)
+  }
+
+  private def awaitNonEmptyRecords[K, V](consumer: Consumer[K, V], partition: TopicPartition): ConsumerRecords[K, V] = {
+    TestUtils.pollRecordsUntilTrue(consumer, (polledRecords: ConsumerRecords[K, V]) => {
+      if (polledRecords.records(partition).asScala.nonEmpty)
+        return polledRecords
+      false
+    }, s"Consumer did not consume any messages for partition $partition before timeout.")
+    throw new IllegalStateException("Should have timed out before reaching here")
+  }
+
+  private def awaitAssignment(consumer: Consumer[_, _], expectedAssignment: Set[TopicPartition]): Unit = {
+    TestUtils.pollUntilTrue(consumer, () => consumer.assignment() == expectedAssignment.asJava,
+      s"Timed out while awaiting expected assignment $expectedAssignment. " +
+        s"The current assignment is ${consumer.assignment()}")
+  }
+
+  @Test
+  def testConsumingWithNullGroupId(): Unit = {
+    val topic = "test_topic"
+    val partition = 0;
+    val tp = new TopicPartition(topic, partition)
+    createTopic(topic, 1, 1)
+
     TestUtils.waitUntilTrue(() => {
-      consumer.poll(50)
-      consumer.assignment() == subscriptions.asJava
-    }, s"Expected partitions ${subscriptions.asJava} but actually got ${consumer.assignment()}")
+      this.zkClient.topicExists(topic)
+    }, "Failed to create topic")
+
+    val producer = createProducer()
+    producer.send(new ProducerRecord(topic, partition, "k1".getBytes, "v1".getBytes)).get()
+    producer.send(new ProducerRecord(topic, partition, "k2".getBytes, "v2".getBytes)).get()
+    producer.send(new ProducerRecord(topic, partition, "k3".getBytes, "v3".getBytes)).get()
+    producer.close()
+
+    // consumer 1 uses the default group id and consumes from earliest offset
+    val consumer1Config = new Properties(consumerConfig)
+    consumer1Config.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest")
+    consumer1Config.put(ConsumerConfig.CLIENT_ID_CONFIG, "consumer1")
+    val consumer1 = createConsumer(
+      configOverrides = consumer1Config,
+      configsToRemove = List(ConsumerConfig.GROUP_ID_CONFIG))
+
+    // consumer 2 uses the default group id and consumes from latest offset
+    val consumer2Config = new Properties(consumerConfig)
+    consumer2Config.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "latest")
+    consumer2Config.put(ConsumerConfig.CLIENT_ID_CONFIG, "consumer2")
+    val consumer2 = createConsumer(
+      configOverrides = consumer2Config,
+      configsToRemove = List(ConsumerConfig.GROUP_ID_CONFIG))
+
+    // consumer 3 uses the default group id and starts from an explicit offset
+    val consumer3Config = new Properties(consumerConfig)
+    consumer3Config.put(ConsumerConfig.CLIENT_ID_CONFIG, "consumer3")
+    val consumer3 = createConsumer(
+      configOverrides = consumer3Config,
+      configsToRemove = List(ConsumerConfig.GROUP_ID_CONFIG))
+
+    consumer1.assign(asList(tp))
+    consumer2.assign(asList(tp))
+    consumer3.assign(asList(tp))
+    consumer3.seek(tp, 1)
+
+    val numRecords1 = consumer1.poll(Duration.ofMillis(5000)).count()
+
+    try {
+      consumer1.commitSync()
+      fail("Expected offset commit to fail due to null group id")
+    } catch {
+      case e: InvalidGroupIdException => // OK
+    }
+
+    try {
+      consumer2.committed(Set(tp).asJava)
+      fail("Expected committed offset fetch to fail due to null group id")
+    } catch {
+      case e: InvalidGroupIdException => // OK
+    }
+
+    val numRecords2 = consumer2.poll(Duration.ofMillis(5000)).count()
+    val numRecords3 = consumer3.poll(Duration.ofMillis(5000)).count()
+
+    consumer1.unsubscribe()
+    consumer2.unsubscribe()
+    consumer3.unsubscribe()
+
+    consumer1.close()
+    consumer2.close()
+    consumer3.close()
+
+    assertEquals("Expected consumer1 to consume from earliest offset", 3, numRecords1)
+    assertEquals("Expected consumer2 to consume from latest offset", 0, numRecords2)
+    assertEquals("Expected consumer3 to consume from offset 1", 2, numRecords3)
   }
 
+  @Test
+  def testConsumingWithEmptyGroupId(): Unit = {
+    val topic = "test_topic"
+    val partition = 0;
+    val tp = new TopicPartition(topic, partition)
+    createTopic(topic, 1, 1)
+
+    TestUtils.waitUntilTrue(() => {
+      this.zkClient.topicExists(topic)
+    }, "Failed to create topic")
+
+    val producer = createProducer()
+    producer.send(new ProducerRecord(topic, partition, "k1".getBytes, "v1".getBytes)).get()
+    producer.send(new ProducerRecord(topic, partition, "k2".getBytes, "v2".getBytes)).get()
+    producer.close()
+
+    // consumer 1 uses the empty group id
+    val consumer1Config = new Properties(consumerConfig)
+    consumer1Config.put(ConsumerConfig.GROUP_ID_CONFIG, "")
+    consumer1Config.put(ConsumerConfig.CLIENT_ID_CONFIG, "consumer1")
+    consumer1Config.put(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, "1")
+    val consumer1 = createConsumer(configOverrides = consumer1Config)
+
+    // consumer 2 uses the empty group id and consumes from latest offset if there is no committed offset
+    val consumer2Config = new Properties(consumerConfig)
+    consumer2Config.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "latest")
+    consumer2Config.put(ConsumerConfig.GROUP_ID_CONFIG, "")
+    consumer2Config.put(ConsumerConfig.CLIENT_ID_CONFIG, "consumer2")
+    consumer2Config.put(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, "1")
+    val consumer2 = createConsumer(configOverrides = consumer2Config)
+
+    consumer1.assign(asList(tp))
+    consumer2.assign(asList(tp))
+
+    val records1 = consumer1.poll(Duration.ofMillis(5000))
+    consumer1.commitSync()
+
+    val records2 = consumer2.poll(Duration.ofMillis(5000))
+    consumer2.commitSync()
+
+    consumer1.close()
+    consumer2.close()
+
+    assertTrue("Expected consumer1 to consume one message from offset 0",
+      records1.count() == 1 && records1.records(tp).asScala.head.offset == 0)
+    assertTrue("Expected consumer2 to consume one message from offset 1, which is the committed offset of consumer1",
+      records2.count() == 1 && records2.records(tp).asScala.head.offset == 1)
+  }
 }
diff --git a/core/src/test/scala/integration/kafka/api/PlaintextEndToEndAuthorizationTest.scala b/core/src/test/scala/integration/kafka/api/PlaintextEndToEndAuthorizationTest.scala
index 627934035c230..fdff774d73ad9 100644
--- a/core/src/test/scala/integration/kafka/api/PlaintextEndToEndAuthorizationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/PlaintextEndToEndAuthorizationTest.scala
@@ -19,13 +19,21 @@ package kafka.api
 import org.apache.kafka.common.config.internals.BrokerSecurityConfigs
 import org.apache.kafka.common.network.ListenerName
 import org.apache.kafka.common.security.auth._
-import org.junit.Before
+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
 object PlaintextEndToEndAuthorizationTest {
+  @volatile
+  private var clientListenerName = None: Option[String]
+  @volatile
+  private var serverListenerName = None: Option[String]
   class TestClientPrincipalBuilder extends KafkaPrincipalBuilder {
     override def build(context: AuthenticationContext): KafkaPrincipal = {
+      clientListenerName = Some(context.listenerName)
       context match {
         case ctx: PlaintextAuthenticationContext if ctx.clientAddress != null =>
           new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "client")
@@ -37,6 +45,7 @@ object PlaintextEndToEndAuthorizationTest {
 
   class TestServerPrincipalBuilder extends KafkaPrincipalBuilder {
     override def build(context: AuthenticationContext): KafkaPrincipal = {
+      serverListenerName = Some(context.listenerName)
       context match {
         case ctx: PlaintextAuthenticationContext =>
           new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "server")
@@ -58,13 +67,23 @@ class PlaintextEndToEndAuthorizationTest extends EndToEndAuthorizationTest {
     classOf[TestClientPrincipalBuilder].getName)
   this.serverConfig.setProperty("listener.name.server." + BrokerSecurityConfigs.PRINCIPAL_BUILDER_CLASS_CONFIG,
     classOf[TestServerPrincipalBuilder].getName)
-  override val clientPrincipal = "client"
-  override val kafkaPrincipal = "server"
+  override val clientPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "client")
+  override val kafkaPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "server")
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     startSasl(jaasSections(List.empty, None, ZkSasl))
     super.setUp()
   }
 
+  @Test
+  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))
+
+    assertEquals(Some("CLIENT"), PlaintextEndToEndAuthorizationTest.clientListenerName)
+    assertEquals(Some("SERVER"), PlaintextEndToEndAuthorizationTest.serverListenerName)
+  }
+
 }
diff --git a/core/src/test/scala/integration/kafka/api/PlaintextProducerSendTest.scala b/core/src/test/scala/integration/kafka/api/PlaintextProducerSendTest.scala
index 10063a99801d3..3b91c2c13fa49 100644
--- a/core/src/test/scala/integration/kafka/api/PlaintextProducerSendTest.scala
+++ b/core/src/test/scala/integration/kafka/api/PlaintextProducerSendTest.scala
@@ -18,48 +18,54 @@
 package kafka.api
 
 import java.util.Properties
-import java.util.concurrent.ExecutionException
+import java.util.concurrent.{ExecutionException, Future, TimeUnit}
 
 import kafka.log.LogConfig
+import kafka.server.Defaults
 import kafka.utils.TestUtils
-import org.apache.kafka.clients.producer.{KafkaProducer, ProducerConfig, ProducerRecord}
-import org.apache.kafka.common.errors.{InvalidTimestampException, SerializationException}
-import org.apache.kafka.common.record.TimestampType
+import org.apache.kafka.clients.producer.{BufferExhaustedException, KafkaProducer, ProducerConfig, ProducerRecord, RecordMetadata}
+import org.apache.kafka.common.errors.{InvalidTimestampException, RecordTooLargeException, SerializationException, TimeoutException}
+import org.apache.kafka.common.record.{DefaultRecord, DefaultRecordBatch, Records, TimestampType}
+import org.apache.kafka.common.serialization.ByteArraySerializer
 import org.junit.Assert._
 import org.junit.Test
+import org.scalatest.Assertions.intercept
+
 
 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")
     val producer = registerProducer(new KafkaProducer(producerProps))
-    val record = new ProducerRecord[Array[Byte], Array[Byte]](topic, new Integer(0), "key".getBytes, "value".getBytes)
+    val record = new ProducerRecord[Array[Byte], Array[Byte]](topic, 0, "key".getBytes, "value".getBytes)
     producer.send(record)
   }
 
   @Test
-  def testBatchSizeZero() {
-    val producerProps = new Properties()
-    producerProps.setProperty(ProducerConfig.BATCH_SIZE_CONFIG, "0")
-    val producer = createProducer(brokerList = brokerList, lingerMs = Long.MaxValue, props = Some(producerProps))
+  def testBatchSizeZero(): Unit = {
+    val producer = createProducer(brokerList = brokerList,
+      lingerMs = Int.MaxValue,
+      deliveryTimeoutMs = Int.MaxValue,
+      batchSize = 0)
     sendAndVerify(producer)
   }
 
   @Test
-  def testSendCompressedMessageWithLogAppendTime() {
-    val producerProps = new Properties()
-    producerProps.setProperty(ProducerConfig.COMPRESSION_TYPE_CONFIG, "gzip")
-    val producer = createProducer(brokerList = brokerList, lingerMs = Long.MaxValue, props = Some(producerProps))
+  def testSendCompressedMessageWithLogAppendTime(): Unit = {
+    val producer = createProducer(brokerList = brokerList,
+      compressionType = "gzip",
+      lingerMs = Int.MaxValue,
+      deliveryTimeoutMs = Int.MaxValue)
     sendAndVerifyTimestamp(producer, TimestampType.LOG_APPEND_TIME)
   }
 
   @Test
-  def testSendNonCompressedMessageWithLogAppendTime() {
-    val producer = createProducer(brokerList = brokerList, lingerMs = Long.MaxValue)
+  def testSendNonCompressedMessageWithLogAppendTime(): Unit = {
+    val producer = createProducer(brokerList = brokerList, lingerMs = Int.MaxValue, deliveryTimeoutMs = Int.MaxValue)
     sendAndVerifyTimestamp(producer, TimestampType.LOG_APPEND_TIME)
   }
 
@@ -69,16 +75,15 @@ class PlaintextProducerSendTest extends BaseProducerSendTest {
    * The topic should be created upon sending the first message
    */
   @Test
-  def testAutoCreateTopic() {
-    val producer = createProducer(brokerList, retries = 5)
-
+  def testAutoCreateTopic(): Unit = {
+    val producer = createProducer(brokerList)
     try {
       // Send a message to auto-create the topic
       val record = new ProducerRecord(topic, null, "key".getBytes, "value".getBytes)
       assertEquals("Should have offset 0", 0L, producer.send(record).get.offset)
 
       // double check that the topic is created with leader elected
-      TestUtils.waitUntilLeaderIsElectedOrChanged(zkUtils, topic, 0)
+      TestUtils.waitUntilLeaderIsElectedOrChanged(zkClient, topic, 0)
 
     } finally {
       producer.close()
@@ -86,33 +91,111 @@ class PlaintextProducerSendTest extends BaseProducerSendTest {
   }
 
   @Test
-  def testSendWithInvalidCreateTime() {
+  def testSendWithInvalidCreateTime(): Unit = {
     val topicProps = new Properties()
     topicProps.setProperty(LogConfig.MessageTimestampDifferenceMaxMsProp, "1000")
-    TestUtils.createTopic(zkUtils, topic, 1, 2, servers, topicProps)
+    createTopic(topic, 1, 2, topicProps)
 
     val producer = createProducer(brokerList = brokerList)
     try {
-      producer.send(new ProducerRecord(topic, 0, System.currentTimeMillis() - 1001, "key".getBytes, "value".getBytes)).get()
-      fail("Should throw CorruptedRecordException")
-    } catch {
-      case e: ExecutionException => assertTrue(e.getCause.isInstanceOf[InvalidTimestampException])
+      val e = intercept[ExecutionException] {
+        producer.send(new ProducerRecord(topic, 0, System.currentTimeMillis() - 1001, "key".getBytes, "value".getBytes)).get()
+      }.getCause
+      assertTrue(e.isInstanceOf[InvalidTimestampException])
+      assertEquals("One or more records have been rejected due to invalid timestamp", e.getMessage)
     } finally {
       producer.close()
     }
 
     // Test compressed messages.
-    val producerProps = new Properties()
-    producerProps.setProperty(ProducerConfig.COMPRESSION_TYPE_CONFIG, "gzip")
-    val compressedProducer = createProducer(brokerList = brokerList, props = Some(producerProps))
+    val compressedProducer = createProducer(brokerList = brokerList, compressionType = "gzip")
     try {
-      compressedProducer.send(new ProducerRecord(topic, 0, System.currentTimeMillis() - 1001, "key".getBytes, "value".getBytes)).get()
-      fail("Should throw CorruptedRecordException")
-    } catch {
-      case e: ExecutionException => assertTrue(e.getCause.isInstanceOf[InvalidTimestampException])
+      val e = intercept[ExecutionException] {
+        compressedProducer.send(new ProducerRecord(topic, 0, System.currentTimeMillis() - 1001, "key".getBytes, "value".getBytes)).get()
+      }.getCause
+      assertTrue(e.isInstanceOf[InvalidTimestampException])
+      assertEquals("One or more records have been rejected due to invalid timestamp", e.getMessage)
     } finally {
       compressedProducer.close()
     }
   }
 
+  // Test that producer with max.block.ms=0 can be used to send in non-blocking mode
+  // where requests are failed immediately without blocking if metadata is not available
+  // or buffer is full.
+  @Test
+  def testNonBlockingProducer(): Unit = {
+
+    def send(producer: KafkaProducer[Array[Byte],Array[Byte]]): Future[RecordMetadata] = {
+      producer.send(new ProducerRecord(topic, 0, "key".getBytes, new Array[Byte](1000)))
+    }
+
+    def sendUntilQueued(producer: KafkaProducer[Array[Byte],Array[Byte]]): Future[RecordMetadata] = {
+      val (future, _) = TestUtils.computeUntilTrue(send(producer))(future => {
+        if (future.isDone) {
+          try {
+            future.get
+            true  // Send was queued and completed successfully
+          } catch {
+            case _: ExecutionException => false
+          }
+        } else
+          true    // Send future not yet complete, so it has been queued to be sent
+      })
+      future
+    }
+
+    def verifySendSuccess(future: Future[RecordMetadata]): Unit = {
+      val recordMetadata = future.get(30, TimeUnit.SECONDS)
+      assertEquals(topic, recordMetadata.topic)
+      assertEquals(0, recordMetadata.partition)
+      assertTrue(s"Invalid offset $recordMetadata", recordMetadata.offset >= 0)
+    }
+
+    def verifyMetadataNotAvailable(future: Future[RecordMetadata]): Unit = {
+      assertTrue(future.isDone)  // verify future was completed immediately
+      assertEquals(classOf[TimeoutException], intercept[ExecutionException](future.get).getCause.getClass)
+    }
+
+    def verifyBufferExhausted(future: Future[RecordMetadata]): Unit = {
+      assertTrue(future.isDone)  // verify future was completed immediately
+      assertEquals(classOf[BufferExhaustedException], intercept[ExecutionException](future.get).getCause.getClass)
+    }
+
+    // Topic metadata not available, send should fail without blocking
+    val producer = createProducer(brokerList = brokerList, maxBlockMs = 0)
+    verifyMetadataNotAvailable(send(producer))
+
+    // Test that send starts succeeding once metadata is available
+    val future = sendUntilQueued(producer)
+    verifySendSuccess(future)
+
+    // Verify that send fails immediately without blocking when there is no space left in the buffer
+    val producer2 = createProducer(brokerList = brokerList, maxBlockMs = 0,
+                                   lingerMs = 15000, batchSize = 1100, bufferSize = 1500)
+    val future2 = sendUntilQueued(producer2) // wait until metadata is available and one record is queued
+    verifyBufferExhausted(send(producer2))       // should fail send since buffer is full
+    verifySendSuccess(future2)               // previous batch should be completed and sent now
+  }
+
+  @Test
+  def testSendRecordBatchWithMaxRequestSizeAndHigher(): Unit = {
+    val producerProps = new Properties()
+    producerProps.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList)
+    val producer = registerProducer(new KafkaProducer(producerProps, new ByteArraySerializer, new ByteArraySerializer))
+
+    val keyLengthSize = 1
+    val headerLengthSize = 1
+    val valueLengthSize = 3
+    val overhead = Records.LOG_OVERHEAD + DefaultRecordBatch.RECORD_BATCH_OVERHEAD + DefaultRecord.MAX_RECORD_OVERHEAD +
+      keyLengthSize + headerLengthSize + valueLengthSize
+    val valueSize = Defaults.MessageMaxBytes - overhead
+
+    val record0 = new ProducerRecord(topic, new Array[Byte](0), new Array[Byte](valueSize))
+    assertEquals(record0.value.length, producer.send(record0).get.serializedValueSize)
+
+    val record1 = new ProducerRecord(topic, new Array[Byte](0), new Array[Byte](valueSize + 1))
+    assertEquals(classOf[RecordTooLargeException], intercept[ExecutionException](producer.send(record1).get).getCause.getClass)
+  }
+
 }
diff --git a/core/src/test/scala/integration/kafka/api/ProducerBounceTest.scala b/core/src/test/scala/integration/kafka/api/ProducerBounceTest.scala
deleted file mode 100644
index ab13b0a557eef..0000000000000
--- a/core/src/test/scala/integration/kafka/api/ProducerBounceTest.scala
+++ /dev/null
@@ -1,170 +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.Properties
-import java.util.concurrent.Future
-
-import kafka.consumer.SimpleConsumer
-import kafka.integration.KafkaServerTestHarness
-import kafka.server.KafkaConfig
-import kafka.utils.{ShutdownableThread, TestUtils}
-import kafka.utils.Implicits._
-import org.apache.kafka.clients.producer._
-import org.apache.kafka.clients.producer.internals.ErrorLoggingCallback
-import org.junit.Assert._
-import org.junit.{Ignore, Test}
-
-import scala.collection.mutable.ArrayBuffer
-
-class ProducerBounceTest extends KafkaServerTestHarness {
-  private val producerBufferSize =  65536
-  private val serverMessageMaxBytes =  producerBufferSize/2
-
-  val numServers = 4
-
-  val overridingProps = new Properties()
-  overridingProps.put(KafkaConfig.AutoCreateTopicsEnableProp, false.toString)
-  overridingProps.put(KafkaConfig.MessageMaxBytesProp, serverMessageMaxBytes.toString)
-  // Set a smaller value for the number of partitions for the offset commit topic (__consumer_offset topic)
-  // so that the creation of that topic/partition(s) and subsequent leader assignment doesn't take relatively long
-  overridingProps.put(KafkaConfig.OffsetsTopicPartitionsProp, 1.toString)
-  overridingProps.put(KafkaConfig.ControlledShutdownEnableProp, true.toString)
-  overridingProps.put(KafkaConfig.UncleanLeaderElectionEnableProp, false.toString)
-  overridingProps.put(KafkaConfig.AutoLeaderRebalanceEnableProp, false.toString)
-  // This is the one of the few tests we currently allow to preallocate ports, despite the fact that this can result in transient
-  // failures due to ports getting reused. We can't use random ports because of bad behavior that can result from bouncing
-  // brokers too quickly when they get new, random ports. If we're not careful, the client can end up in a situation
-  // where metadata is not refreshed quickly enough, and by the time it's actually trying to, all the servers have
-  // been bounced and have new addresses. None of the bootstrap nodes or current metadata can get them connected to a
-  // running server.
-  //
-  // Since such quick rotation of servers is incredibly unrealistic, we allow this one test to preallocate ports, leaving
-  // a small risk of hitting errors due to port conflicts. Hopefully this is infrequent enough to not cause problems.
-  override def generateConfigs = {
-    FixedPortTestUtils.createBrokerConfigs(numServers, zkConnect,enableControlledShutdown = true)
-      .map(KafkaConfig.fromProps(_, overridingProps))
-  }
-
-  private val topic1 = "topic-1"
-
-  /**
-   * With replication, producer should able to find new leader after it detects broker failure
-   */
-  @Ignore // To be re-enabled once we can make it less flaky (KAFKA-2837)
-  @Test
-  def testBrokerFailure() {
-    val numPartitions = 3
-    val topicConfig = new Properties()
-    topicConfig.put(KafkaConfig.MinInSyncReplicasProp, 2.toString)
-    TestUtils.createTopic(zkUtils, topic1, numPartitions, numServers, servers, topicConfig)
-
-    val scheduler = new ProducerScheduler()
-    scheduler.start
-
-    // rolling bounce brokers
-
-    for (_ <- 0 until numServers) {
-      for (server <- servers) {
-        info("Shutting down server : %s".format(server.config.brokerId))
-        server.shutdown()
-        server.awaitShutdown()
-        info("Server %s shut down. Starting it up again.".format(server.config.brokerId))
-        server.startup()
-        info("Restarted server: %s".format(server.config.brokerId))
-      }
-
-      // Make sure the producer do not see any exception in returned metadata due to broker failures
-      assertFalse(scheduler.failed)
-
-      // Make sure the leader still exists after bouncing brokers
-      (0 until numPartitions).foreach(partition => TestUtils.waitUntilLeaderIsElectedOrChanged(zkUtils, topic1, partition))
-    }
-
-    scheduler.shutdown
-
-    // Make sure the producer do not see any exception
-    // when draining the left messages on shutdown
-    assertFalse(scheduler.failed)
-
-    // double check that the leader info has been propagated after consecutive bounces
-    val newLeaders = (0 until numPartitions).map(i => TestUtils.waitUntilMetadataIsPropagated(servers, topic1, i))
-    val fetchResponses = newLeaders.zipWithIndex.map { case (leader, partition) =>
-      // Consumers must be instantiated after all the restarts since they use random ports each time they start up
-      val consumer = new SimpleConsumer("localhost", boundPort(servers(leader)), 30000, 1024 * 1024, "")
-      val response = consumer.fetch(new FetchRequestBuilder().addFetch(topic1, partition, 0, Int.MaxValue).build()).messageSet(topic1, partition)
-      consumer.close
-      response
-    }
-    val messages = fetchResponses.flatMap(r => r.iterator.toList.map(_.message))
-    val uniqueMessages = messages.toSet
-    val uniqueMessageSize = uniqueMessages.size
-    info(s"number of unique messages sent: ${uniqueMessageSize}")
-    assertEquals(s"Found ${messages.size - uniqueMessageSize} duplicate messages.", uniqueMessageSize, messages.size)
-    assertEquals("Should have fetched " + scheduler.sent + " unique messages", scheduler.sent, messages.size)
-  }
-
-  private class ProducerScheduler extends ShutdownableThread("daemon-producer", false) {
-    val numRecords = 1000
-    var sent = 0
-    var failed = false
-
-    val producerConfig = new Properties()
-    producerConfig.put(ProducerConfig.ENABLE_IDEMPOTENCE_CONFIG, "true")
-    producerConfig.put(ProducerConfig.MAX_IN_FLIGHT_REQUESTS_PER_CONNECTION, "5")
-    val producerConfigWithCompression = new Properties()
-    producerConfigWithCompression ++= producerConfig
-    producerConfigWithCompression.put(ProducerConfig.COMPRESSION_TYPE_CONFIG, "lz4")
-    val producers = List(
-      TestUtils.createNewProducer(brokerList, bufferSize = producerBufferSize / 4, retries = 10, props = Some(producerConfig)),
-      TestUtils.createNewProducer(brokerList, bufferSize = producerBufferSize / 2, retries = 10, lingerMs = 5000, props = Some(producerConfig)),
-      TestUtils.createNewProducer(brokerList, bufferSize = producerBufferSize, retries = 10, lingerMs = 10000, props = Some(producerConfigWithCompression))
-    )
-
-    override def doWork(): Unit = {
-      info("Starting to send messages..")
-      var producerId = 0
-      val responses = new ArrayBuffer[IndexedSeq[Future[RecordMetadata]]]()
-      for (producer <- producers) {
-        val response =
-          for (i <- sent+1 to sent+numRecords)
-            yield producer.send(new ProducerRecord[Array[Byte],Array[Byte]](topic1, null, null, ((producerId + 1) * i).toString.getBytes),
-              new ErrorLoggingCallback(topic1, null, null, true))
-        responses.append(response)
-        producerId += 1
-      }
-
-      try {
-        for (response <- responses) {
-          val futures = response.toList
-          futures.map(_.get)
-          sent += numRecords
-        }
-        info(s"Sent $sent records")
-      } catch {
-        case e : Exception =>
-          error(s"Got exception ${e.getMessage}")
-          e.printStackTrace()
-          failed = true
-      }
-    }
-
-    override def shutdown(){
-      super.shutdown()
-      for (producer <- producers) {
-        producer.close()
-      }
-    }
-  }
-}
diff --git a/core/src/test/scala/integration/kafka/api/ProducerCompressionTest.scala b/core/src/test/scala/integration/kafka/api/ProducerCompressionTest.scala
index 8cbdd93da1a7a..a8788a447249d 100755
--- a/core/src/test/scala/integration/kafka/api/ProducerCompressionTest.scala
+++ b/core/src/test/scala/integration/kafka/api/ProducerCompressionTest.scala
@@ -19,7 +19,7 @@ package kafka.api.test
 
 import java.util.{Collection, Collections, Properties}
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 import org.junit.runners.Parameterized
 import org.junit.runner.RunWith
 import org.junit.runners.Parameterized.Parameters
@@ -30,7 +30,6 @@ import kafka.server.{KafkaConfig, KafkaServer}
 import kafka.zk.ZooKeeperTestHarness
 import kafka.utils.TestUtils
 import org.apache.kafka.common.TopicPartition
-import org.apache.kafka.common.security.auth.SecurityProtocol
 import org.apache.kafka.common.serialization.ByteArraySerializer
 
 @RunWith(value = classOf[Parameterized])
@@ -43,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()
   }
@@ -61,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))
@@ -70,11 +69,11 @@ class ProducerCompressionTest(compression: String) extends ZooKeeperTestHarness
     producerProps.put(ProducerConfig.BATCH_SIZE_CONFIG, "66000")
     producerProps.put(ProducerConfig.LINGER_MS_CONFIG, "200")
     val producer = new KafkaProducer(producerProps, new ByteArraySerializer, new ByteArraySerializer)
-    val consumer = TestUtils.createNewConsumer(bootstrapServers, securityProtocol = SecurityProtocol.PLAINTEXT)
+    val consumer = TestUtils.createConsumer(bootstrapServers)
 
     try {
       // create topic
-      TestUtils.createTopic(zkUtils, topic, 1, 1, List(server))
+      TestUtils.createTopic(zkClient, topic, 1, 1, List(server))
       val partition = 0
 
       // prepare the messages
@@ -114,7 +113,8 @@ object ProducerCompressionTest {
       Array("none"),
       Array("gzip"),
       Array("snappy"),
-      Array("lz4")
+      Array("lz4"),
+      Array("zstd")
     ).asJava
   }
 }
diff --git a/core/src/test/scala/integration/kafka/api/ProducerFailureHandlingTest.scala b/core/src/test/scala/integration/kafka/api/ProducerFailureHandlingTest.scala
index e1b1c9d6d146e..fb4d839c0351f 100644
--- a/core/src/test/scala/integration/kafka/api/ProducerFailureHandlingTest.scala
+++ b/core/src/test/scala/integration/kafka/api/ProducerFailureHandlingTest.scala
@@ -17,7 +17,7 @@
 
 package kafka.api
 
-import java.util.concurrent.{ExecutionException, TimeoutException}
+import java.util.concurrent.ExecutionException
 import java.util.Properties
 
 import kafka.integration.KafkaServerTestHarness
@@ -25,12 +25,12 @@ import kafka.log.LogConfig
 import kafka.server.KafkaConfig
 import kafka.utils.TestUtils
 import org.apache.kafka.clients.producer._
-import org.apache.kafka.common.KafkaException
 import org.apache.kafka.common.errors._
 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
@@ -61,19 +61,19 @@ class ProducerFailureHandlingTest extends KafkaServerTestHarness {
   private val topic2 = "topic-2"
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
 
-    producer1 = TestUtils.createNewProducer(brokerList, acks = 0, requestTimeoutMs = 30000L, maxBlockMs = 10000L,
+    producer1 = TestUtils.createProducer(brokerList, acks = 0, retries = 0, requestTimeoutMs = 30000, maxBlockMs = 10000L,
       bufferSize = producerBufferSize)
-    producer2 = TestUtils.createNewProducer(brokerList, acks = 1, requestTimeoutMs = 30000L, maxBlockMs = 10000L,
+    producer2 = TestUtils.createProducer(brokerList, acks = 1, retries = 0, requestTimeoutMs = 30000, maxBlockMs = 10000L,
       bufferSize = producerBufferSize)
-    producer3 = TestUtils.createNewProducer(brokerList, acks = -1, requestTimeoutMs = 30000L, maxBlockMs = 10000L,
+    producer3 = TestUtils.createProducer(brokerList, acks = -1, retries = 0, requestTimeoutMs = 30000, maxBlockMs = 10000L,
       bufferSize = producerBufferSize)
   }
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     if (producer1 != null) producer1.close()
     if (producer2 != null) producer2.close()
     if (producer3 != null) producer3.close()
@@ -86,9 +86,9 @@ 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
-    TestUtils.createTopic(zkUtils, topic1, 1, numServers, servers)
+    createTopic(topic1, replicationFactor = numServers)
 
     // send a too-large record
     val record = new ProducerRecord(topic1, null, "key".getBytes, new Array[Byte](serverMessageMaxBytes + 1))
@@ -103,9 +103,9 @@ class ProducerFailureHandlingTest extends KafkaServerTestHarness {
    * With ack == 1 the future metadata will throw ExecutionException caused by RecordTooLargeException
    */
   @Test
-  def testTooLargeRecordWithAckOne() {
+  def testTooLargeRecordWithAckOne(): Unit = {
     // create topic
-    TestUtils.createTopic(zkUtils, topic1, 1, numServers, servers)
+    createTopic(topic1, replicationFactor = numServers)
 
     // send a too-large record
     val record = new ProducerRecord(topic1, null, "key".getBytes, new Array[Byte](serverMessageMaxBytes + 1))
@@ -114,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)
@@ -122,7 +122,7 @@ class ProducerFailureHandlingTest extends KafkaServerTestHarness {
 
     // create topic
     val topic10 = "topic10"
-    TestUtils.createTopic(zkUtils, topic10, servers.size, numServers, servers, topicConfig)
+    createTopic(topic10, numPartitions = servers.size, replicationFactor = numServers, topicConfig)
 
     // send a record that is too large for replication, but within the broker max message limit
     val value = new Array[Byte](maxMessageSize - DefaultRecordBatch.RECORD_BATCH_OVERHEAD - DefaultRecord.MAX_RECORD_OVERHEAD)
@@ -134,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)
   }
 
@@ -148,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] {
@@ -161,18 +161,18 @@ class ProducerFailureHandlingTest extends KafkaServerTestHarness {
    *
    * TODO: other exceptions that can be thrown in ExecutionException:
    *    UnknownTopicOrPartitionException
-   *    NotLeaderForPartitionException
+   *    NotLeaderOrFollowerException
    *    LeaderNotAvailableException
    *    CorruptRecordException
    *    TimeoutException
    */
   @Test
-  def testWrongBrokerList() {
+  def testWrongBrokerList(): Unit = {
     // create topic
-    TestUtils.createTopic(zkUtils, topic1, 1, numServers, servers)
+    createTopic(topic1, replicationFactor = numServers)
 
     // producer with incorrect broker list
-    producer4 = TestUtils.createNewProducer("localhost:8686,localhost:4242", acks = 1, maxBlockMs = 10000L, bufferSize = producerBufferSize)
+    producer4 = TestUtils.createProducer("localhost:8686,localhost:4242", acks = 1, maxBlockMs = 10000L, bufferSize = producerBufferSize)
 
     // send a record with incorrect broker list
     val record = new ProducerRecord(topic1, null, "key".getBytes, "value".getBytes)
@@ -182,18 +182,21 @@ class ProducerFailureHandlingTest extends KafkaServerTestHarness {
   }
 
   /**
-    * Send with invalid partition id should throw KafkaException when partition is higher than the upper bound of
-    * partitions.
+    * Send with invalid partition id should return ExecutionException caused by TimeoutException
+    * when partition is higher than the upper bound of partitions.
     */
   @Test
-  def testInvalidPartition() {
+  def testInvalidPartition(): Unit = {
     // create topic with a single partition
-    TestUtils.createTopic(zkUtils, topic1, 1, numServers, servers)
+    createTopic(topic1, numPartitions = 1, replicationFactor = numServers)
 
     // create a record with incorrect partition id (higher than the number of partitions), send should fail
     val higherRecord = new ProducerRecord(topic1, 1, "key".getBytes, "value".getBytes)
-    intercept[KafkaException] {
-      producer1.send(higherRecord)
+    intercept[ExecutionException] {
+      producer1.send(higherRecord).get
+    }.getCause match {
+      case _: TimeoutException => // this is ok
+      case ex => throw new Exception("Sending to a partition not present in the metadata should result in a TimeoutException", ex)
     }
   }
 
@@ -201,11 +204,11 @@ class ProducerFailureHandlingTest extends KafkaServerTestHarness {
    * The send call after producer closed should throw IllegalStateException
    */
   @Test
-  def testSendAfterClosed() {
+  def testSendAfterClosed(): Unit = {
     // create topic
-    TestUtils.createTopic(zkUtils, topic1, 1, numServers, servers)
+    createTopic(topic1, replicationFactor = numServers)
 
-    val record = new ProducerRecord[Array[Byte],Array[Byte]](topic1, null, "key".getBytes, "value".getBytes)
+    val record = new ProducerRecord[Array[Byte], Array[Byte]](topic1, null, "key".getBytes, "value".getBytes)
 
     // first send a message to make sure the metadata is refreshed
     producer1.send(record).get
@@ -227,8 +230,8 @@ class ProducerFailureHandlingTest extends KafkaServerTestHarness {
   }
 
   @Test
-  def testCannotSendToInternalTopic() {
-    TestUtils.createOffsetsTopic(zkUtils, servers)
+  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
     }
@@ -236,12 +239,12 @@ 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)
 
-    TestUtils.createTopic(zkUtils, topicName, 1, numServers, servers, topicProps)
+    createTopic(topicName, replicationFactor = numServers, topicConfig = topicProps)
 
     val record = new ProducerRecord(topicName, null, "key".getBytes, "value".getBytes)
     try {
@@ -256,12 +259,12 @@ class ProducerFailureHandlingTest extends KafkaServerTestHarness {
   }
 
   @Test
-  def testNotEnoughReplicasAfterBrokerShutdown() {
+  def testNotEnoughReplicasAfterBrokerShutdown(): Unit = {
     val topicName = "minisrtest2"
     val topicProps = new Properties()
     topicProps.put("min.insync.replicas", numServers.toString)
 
-    TestUtils.createTopic(zkUtils, topicName, 1, numServers, servers,topicProps)
+    createTopic(topicName, replicationFactor = numServers, topicConfig = topicProps)
 
     val record = new ProducerRecord(topicName, null, "key".getBytes, "value".getBytes)
     // this should work with all brokers up and running
diff --git a/core/src/test/scala/integration/kafka/api/RackAwareAutoTopicCreationTest.scala b/core/src/test/scala/integration/kafka/api/RackAwareAutoTopicCreationTest.scala
index 0ae9d1798c8a2..4ba3fd5eac945 100644
--- a/core/src/test/scala/integration/kafka/api/RackAwareAutoTopicCreationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/RackAwareAutoTopicCreationTest.scala
@@ -43,16 +43,16 @@ class RackAwareAutoTopicCreationTest extends KafkaServerTestHarness with RackAwa
   private val topic = "topic"
 
   @Test
-  def testAutoCreateTopic() {
-    val producer = TestUtils.createNewProducer(brokerList, retries = 5)
+  def testAutoCreateTopic(): Unit = {
+    val producer = TestUtils.createProducer(brokerList)
     try {
       // Send a message to auto-create the topic
       val record = new ProducerRecord(topic, null, "key".getBytes, "value".getBytes)
       assertEquals("Should have offset 0", 0L, producer.send(record).get.offset)
 
       // double check that the topic is created with leader elected
-      TestUtils.waitUntilLeaderIsElectedOrChanged(zkUtils, topic, 0)
-      val assignment = zkUtils.getReplicaAssignmentForTopics(Seq(topic)).map { case (topicPartition, replicas) =>
+      TestUtils.waitUntilLeaderIsElectedOrChanged(zkClient, topic, 0)
+      val assignment = zkClient.getReplicaAssignmentForTopics(Set(topic)).map { case (topicPartition, replicas) =>
         topicPartition.partition -> replicas
       }
       val brokerMetadatas = adminZkClient.getBrokerMetadatas(RackAwareMode.Enforced)
diff --git a/core/src/test/scala/integration/kafka/api/SaslClientsWithInvalidCredentialsTest.scala b/core/src/test/scala/integration/kafka/api/SaslClientsWithInvalidCredentialsTest.scala
index b309b809e9760..87940e0c72621 100644
--- a/core/src/test/scala/integration/kafka/api/SaslClientsWithInvalidCredentialsTest.scala
+++ b/core/src/test/scala/integration/kafka/api/SaslClientsWithInvalidCredentialsTest.scala
@@ -12,22 +12,23 @@
   */
 package kafka.api
 
-import java.io.FileOutputStream
+import java.nio.file.Files
+import java.time.Duration
 import java.util.Collections
 import java.util.concurrent.{ExecutionException, TimeUnit}
 
-import scala.collection.JavaConverters._
-import org.apache.kafka.clients.admin.{AdminClient, AdminClientConfig}
+import scala.jdk.CollectionConverters._
+import org.apache.kafka.clients.admin.{Admin, AdminClientConfig}
 import org.apache.kafka.clients.consumer.{ConsumerConfig, KafkaConsumer}
 import org.apache.kafka.clients.producer.{KafkaProducer, ProducerConfig, ProducerRecord}
 import org.apache.kafka.common.{KafkaException, TopicPartition}
 import org.apache.kafka.common.errors.SaslAuthenticationException
-import org.apache.kafka.common.serialization.ByteArrayDeserializer
 import org.junit.{After, Before, Test}
 import org.junit.Assert._
-import kafka.admin.ConsumerGroupCommand.{ConsumerGroupCommandOptions, KafkaConsumerGroupService}
+import kafka.admin.ConsumerGroupCommand.{ConsumerGroupCommandOptions, ConsumerGroupService}
 import kafka.server.KafkaConfig
-import kafka.utils.{JaasTestUtils, TestUtils, ZkUtils}
+import kafka.utils.{JaasTestUtils, TestUtils}
+import kafka.zk.ConfigEntityChangeNotificationZNode
 import org.apache.kafka.common.security.auth.SecurityProtocol
 
 class SaslClientsWithInvalidCredentialsTest extends IntegrationTestHarness with SaslSetup {
@@ -38,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")
@@ -49,19 +50,24 @@ class SaslClientsWithInvalidCredentialsTest extends IntegrationTestHarness with
   val numPartitions = 1
   val tp = new TopicPartition(topic, 0)
 
-  override def configureSecurityBeforeServersStart() {
+  override def configureSecurityBeforeServersStart(): Unit = {
     super.configureSecurityBeforeServersStart()
-    zkUtils.makeSurePersistentPathExists(ZkUtils.ConfigChangesPath)
+    zkClient.makeSurePersistentPathExists(ConfigEntityChangeNotificationZNode.path)
     // Create broker credentials before starting brokers
     createScramCredentials(zkConnect, JaasTestUtils.KafkaScramAdmin, JaasTestUtils.KafkaScramAdminPassword)
   }
 
+  override def createPrivilegedAdminClient() = {
+    createAdminClient(brokerList, securityProtocol, trustStoreFile, clientSaslProperties,
+      kafkaClientSaslMechanism, JaasTestUtils.KafkaScramAdmin, JaasTestUtils.KafkaScramAdminPassword)
+  }
+
   @Before
   override def setUp(): Unit = {
     startSasl(jaasSections(kafkaServerSaslMechanisms, Some(kafkaClientSaslMechanism), Both,
       JaasTestUtils.KafkaServerContextName))
     super.setUp()
-    TestUtils.createTopic(this.zkUtils, topic, numPartitions, serverCount, this.servers)
+    createTopic(topic, numPartitions, brokerCount)
   }
 
   @After
@@ -71,16 +77,17 @@ class SaslClientsWithInvalidCredentialsTest extends IntegrationTestHarness with
   }
 
   @Test
-  def testProducerWithAuthenticationFailure() {
-    verifyAuthenticationException(sendOneRecord(10000))
-    verifyAuthenticationException(producers.head.partitionsFor(topic))
+  def testProducerWithAuthenticationFailure(): Unit = {
+    val producer = createProducer()
+    verifyAuthenticationException(sendOneRecord(producer, maxWaitMs = 10000))
+    verifyAuthenticationException(producer.partitionsFor(topic))
 
     createClientCredential()
-    verifyWithRetry(sendOneRecord())
+    verifyWithRetry(sendOneRecord(producer))
   }
 
   @Test
-  def testTransactionalProducerWithAuthenticationFailure() {
+  def testTransactionalProducerWithAuthenticationFailure(): Unit = {
     val txProducer = createTransactionalProducer()
     verifyAuthenticationException(txProducer.initTransactions())
 
@@ -94,50 +101,49 @@ class SaslClientsWithInvalidCredentialsTest extends IntegrationTestHarness with
   }
 
   @Test
-  def testConsumerWithAuthenticationFailure() {
-    val consumer = this.consumers.head
+  def testConsumerWithAuthenticationFailure(): Unit = {
+    val consumer = createConsumer()
     consumer.subscribe(List(topic).asJava)
     verifyConsumerWithAuthenticationFailure(consumer)
   }
 
   @Test
-  def testManualAssignmentConsumerWithAuthenticationFailure() {
-    val consumer = this.consumers.head
+  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 = new KafkaConsumer(this.consumerConfig, new ByteArrayDeserializer(), new ByteArrayDeserializer())
-    consumers += consumer
+    val consumer = createConsumer()
     consumer.assign(List(tp).asJava)
     consumer.seek(tp, 0)
-
     verifyConsumerWithAuthenticationFailure(consumer)
   }
 
-  private def verifyConsumerWithAuthenticationFailure(consumer: KafkaConsumer[Array[Byte], Array[Byte]]) {
-    verifyAuthenticationException(consumer.poll(10000))
+  private def verifyConsumerWithAuthenticationFailure(consumer: KafkaConsumer[Array[Byte], Array[Byte]]): Unit = {
+    verifyAuthenticationException(consumer.poll(Duration.ofMillis(1000)))
     verifyAuthenticationException(consumer.partitionsFor(topic))
 
     createClientCredential()
-    verifyWithRetry(sendOneRecord())
-    verifyWithRetry(assertEquals(1, consumer.poll(1000).count))
+    val producer = createProducer()
+    verifyWithRetry(sendOneRecord(producer))
+    verifyWithRetry(assertEquals(1, consumer.poll(Duration.ofMillis(1000)).count))
   }
 
   @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)
+    val adminClient = Admin.create(props)
 
     def describeTopic(): Unit = {
       try {
         val response = adminClient.describeTopics(Collections.singleton(topic)).all.get
         assertEquals(1, response.size)
-        response.asScala.foreach { case (topic, description) =>
+        response.forEach { (topic, description) =>
           assertEquals(numPartitions, description.partitions.size)
         }
       } catch {
@@ -151,40 +157,60 @@ class SaslClientsWithInvalidCredentialsTest extends IntegrationTestHarness with
       createClientCredential()
       verifyWithRetry(describeTopic())
     } finally {
-      adminClient.close
+      adminClient.close()
     }
   }
 
   @Test
-  def testConsumerGroupServiceWithAuthenticationFailure() {
+  def testConsumerGroupServiceWithAuthenticationFailure(): Unit = {
+    val consumerGroupService: ConsumerGroupService = prepareConsumerGroupService
+
+    val consumer = createConsumer()
+    try {
+      consumer.subscribe(List(topic).asJava)
+
+      verifyAuthenticationException(consumerGroupService.listGroups())
+    } finally consumerGroupService.close()
+  }
+
+  @Test
+  def testConsumerGroupServiceWithAuthenticationSuccess(): Unit = {
+    createClientCredential()
+    val consumerGroupService: ConsumerGroupService = prepareConsumerGroupService
+
+    val consumer = createConsumer()
+    try {
+      consumer.subscribe(List(topic).asJava)
+
+      verifyWithRetry(consumer.poll(Duration.ofMillis(1000)))
+      assertEquals(1, consumerGroupService.listConsumerGroups().size)
+    }
+    finally consumerGroupService.close()
+  }
+
+  private def prepareConsumerGroupService = {
     val propsFile = TestUtils.tempFile()
-    val propsStream = new FileOutputStream(propsFile)
-    propsStream.write("security.protocol=SASL_PLAINTEXT\n".getBytes())
-    propsStream.write(s"sasl.mechanism=$kafkaClientSaslMechanism".getBytes())
-    propsStream.close()
+    val propsStream = Files.newOutputStream(propsFile.toPath)
+    try {
+      propsStream.write("security.protocol=SASL_PLAINTEXT\n".getBytes())
+      propsStream.write(s"sasl.mechanism=$kafkaClientSaslMechanism".getBytes())
+    }
+    finally propsStream.close()
 
     val cgcArgs = Array("--bootstrap-server", brokerList,
                         "--describe",
                         "--group", "test.group",
                         "--command-config", propsFile.getAbsolutePath)
     val opts = new ConsumerGroupCommandOptions(cgcArgs)
-    val consumerGroupService = new KafkaConsumerGroupService(opts)
-
-    val consumer = consumers.head
-    consumer.subscribe(List(topic).asJava)
-
-    verifyAuthenticationException(consumerGroupService.listGroups)
-    createClientCredential()
-    verifyWithRetry(consumer.poll(1000))
-    assertEquals(1, consumerGroupService.listGroups.size)
+    val consumerGroupService = new ConsumerGroupService(opts)
+    consumerGroupService
   }
 
   private def createClientCredential(): Unit = {
-    createScramCredentials(zkConnect, JaasTestUtils.KafkaScramUser2, JaasTestUtils.KafkaScramPassword2)
+    createScramCredentialsViaPrivilegedAdminClient(JaasTestUtils.KafkaScramUser2, JaasTestUtils.KafkaScramPassword2)
   }
 
-  private def sendOneRecord(maxWaitMs: Long = 15000): Unit = {
-    val producer = this.producers.head
+  private def sendOneRecord(producer: KafkaProducer[Array[Byte], Array[Byte]], maxWaitMs: Long = 15000): Unit = {
     val record = new ProducerRecord(tp.topic(), tp.partition(), 0L, "key".getBytes, "value".getBytes)
     val future = producer.send(record)
     producer.flush()
@@ -202,11 +228,10 @@ class SaslClientsWithInvalidCredentialsTest extends IntegrationTestHarness with
       action
       fail("Expected an authentication exception")
     } catch {
-      case e: SaslAuthenticationException =>
+      case e : Exception =>
         // expected exception
         val elapsedMs = System.currentTimeMillis - startMs
         assertTrue(s"Poll took too long, elapsed=$elapsedMs", elapsedMs <= 5000)
-        assertTrue(s"Exception message not useful: $e", e.getMessage.contains("invalid credentials"))
     }
   }
 
@@ -226,13 +251,6 @@ class SaslClientsWithInvalidCredentialsTest extends IntegrationTestHarness with
   private def createTransactionalProducer(): KafkaProducer[Array[Byte], Array[Byte]] = {
     producerConfig.setProperty(ProducerConfig.TRANSACTIONAL_ID_CONFIG, "txclient-1")
     producerConfig.put(ProducerConfig.ENABLE_IDEMPOTENCE_CONFIG, "true")
-    val txProducer = TestUtils.createNewProducer(brokerList,
-                                  securityProtocol = this.securityProtocol,
-                                  saslProperties = this.clientSaslProperties,
-                                  retries = 1000,
-                                  acks = -1,
-                                  props = Some(producerConfig))
-    producers += txProducer
-    txProducer
+    createProducer()
   }
 }
diff --git a/core/src/test/scala/integration/kafka/api/SaslEndToEndAuthorizationTest.scala b/core/src/test/scala/integration/kafka/api/SaslEndToEndAuthorizationTest.scala
index fb6bee857d7df..66a1a15284415 100644
--- a/core/src/test/scala/integration/kafka/api/SaslEndToEndAuthorizationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/SaslEndToEndAuthorizationTest.scala
@@ -16,17 +16,15 @@
   */
 package kafka.api
 
-import java.util.Properties
-
-import kafka.utils.TestUtils
-import kafka.utils.Implicits._
 import org.apache.kafka.common.config.SaslConfigs
 import org.apache.kafka.common.security.auth.SecurityProtocol
-import org.apache.kafka.common.errors.TopicAuthorizationException
+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._
+import scala.jdk.CollectionConverters._
 
 abstract class SaslEndToEndAuthorizationTest extends EndToEndAuthorizationTest {
   override protected def securityProtocol = SecurityProtocol.SASL_SSL
@@ -37,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
@@ -45,7 +43,8 @@ abstract class SaslEndToEndAuthorizationTest extends EndToEndAuthorizationTest {
     val clientLoginContext = jaasClientLoginModule(kafkaClientSaslMechanism)
     producerConfig.put(SaslConfigs.SASL_JAAS_CONFIG, clientLoginContext)
     consumerConfig.put(SaslConfigs.SASL_JAAS_CONFIG, clientLoginContext)
-    super.setUp
+    adminClientConfig.put(SaslConfigs.SASL_JAAS_CONFIG, clientLoginContext)
+    super.setUp()
   }
 
   /**
@@ -55,21 +54,14 @@ abstract class SaslEndToEndAuthorizationTest extends EndToEndAuthorizationTest {
     */
   @Test(timeout = 15000)
   def testTwoConsumersWithDifferentSaslCredentials(): Unit = {
-    setAclsAndProduce()
-    val consumer1 = consumers.head
+    setAclsAndProduce(tp)
+    val consumer1 = createConsumer()
 
-    val consumer2Config = new Properties
-    consumer2Config ++= consumerConfig
     // consumer2 retrieves its credentials from the static JAAS configuration, so we test also this path
-    consumer2Config.remove(SaslConfigs.SASL_JAAS_CONFIG)
-
-    val consumer2 = TestUtils.createNewConsumer(brokerList,
-                                                securityProtocol = securityProtocol,
-                                                trustStoreFile = trustStoreFile,
-                                                saslProperties = clientSaslProperties,
-                                                props = Some(consumer2Config))
-    consumers += consumer2
+    consumerConfig.remove(SaslConfigs.SASL_JAAS_CONFIG)
+    consumerConfig.remove(SaslConfigs.SASL_CLIENT_CALLBACK_HANDLER_CLASS)
 
+    val consumer2 = createConsumer()
     consumer1.assign(List(tp).asJava)
     consumer2.assign(List(tp).asJava)
 
@@ -77,9 +69,13 @@ abstract class SaslEndToEndAuthorizationTest extends EndToEndAuthorizationTest {
 
     try {
       consumeRecords(consumer2)
-      fail("Expected exception as consumer2 has no access to topic")
+      fail("Expected exception as consumer2 has no access to topic or group")
     } catch {
-      case _: TopicAuthorizationException => //expected
+      // Either exception is possible depending on the order that the first Metadata
+      // and FindCoordinator requests are received
+      case e: TopicAuthorizationException => assertTrue(e.unauthorizedTopics.contains(topic))
+      case e: GroupAuthorizationException => assertEquals(group, e.groupId)
     }
+    confirmReauthenticationMetrics()
   }
 }
diff --git a/core/src/test/scala/integration/kafka/api/SaslGssapiSslEndToEndAuthorizationTest.scala b/core/src/test/scala/integration/kafka/api/SaslGssapiSslEndToEndAuthorizationTest.scala
index a9b4a60ac7268..55c28e4d023a9 100644
--- a/core/src/test/scala/integration/kafka/api/SaslGssapiSslEndToEndAuthorizationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/SaslGssapiSslEndToEndAuthorizationTest.scala
@@ -16,18 +16,25 @@
   */
 package kafka.api
 
+import kafka.security.auth.SimpleAclAuthorizer
 import kafka.server.KafkaConfig
 import kafka.utils.JaasTestUtils
+import org.apache.kafka.common.security.auth.KafkaPrincipal
 
 import scala.collection.immutable.List
 
+// Note: this test currently uses the deprecated SimpleAclAuthorizer to ensure we have test coverage
+// It must be replaced with the new AclAuthorizer when SimpleAclAuthorizer is removed
 class SaslGssapiSslEndToEndAuthorizationTest extends SaslEndToEndAuthorizationTest {
-  override val clientPrincipal = JaasTestUtils.KafkaClientPrincipalUnqualifiedName
-  override val kafkaPrincipal = JaasTestUtils.KafkaServerPrincipalUnqualifiedName
-  
+  override val clientPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE,
+    JaasTestUtils.KafkaClientPrincipalUnqualifiedName)
+  override val kafkaPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE,
+    JaasTestUtils.KafkaServerPrincipalUnqualifiedName)
+
   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
   // authentication or SASL authentication with SSL as the transport layer (but not both).
diff --git a/core/src/test/scala/integration/kafka/api/SaslMultiMechanismConsumerTest.scala b/core/src/test/scala/integration/kafka/api/SaslMultiMechanismConsumerTest.scala
index 6ae99df7e31ef..3f30f131452b0 100644
--- a/core/src/test/scala/integration/kafka/api/SaslMultiMechanismConsumerTest.scala
+++ b/core/src/test/scala/integration/kafka/api/SaslMultiMechanismConsumerTest.scala
@@ -16,10 +16,10 @@ import java.io.File
 
 import kafka.server.KafkaConfig
 import org.junit.{After, Before, Test}
-import kafka.utils.{JaasTestUtils, TestUtils}
+import kafka.utils.JaasTestUtils
 import org.apache.kafka.common.security.auth.SecurityProtocol
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 
 class SaslMultiMechanismConsumerTest extends BaseConsumerTest with SaslSetup {
   private val kafkaClientSaslMechanism = "PLAIN"
@@ -44,22 +44,13 @@ class SaslMultiMechanismConsumerTest extends BaseConsumerTest with SaslSetup {
   }
 
   @Test
-  def testMultipleBrokerMechanisms() {
-
-    val plainSaslProducer = producers.head
-    val plainSaslConsumer = consumers.head
+  def testMultipleBrokerMechanisms(): Unit = {
+    val plainSaslProducer = createProducer()
+    val plainSaslConsumer = createConsumer()
 
     val gssapiSaslProperties = kafkaClientSaslProperties("GSSAPI", dynamicJaasConfig = true)
-    val gssapiSaslProducer = TestUtils.createNewProducer(brokerList,
-                                                         securityProtocol = this.securityProtocol,
-                                                         trustStoreFile = this.trustStoreFile,
-                                                         saslProperties = Some(gssapiSaslProperties))
-    producers += gssapiSaslProducer
-    val gssapiSaslConsumer = TestUtils.createNewConsumer(brokerList,
-                                                         securityProtocol = this.securityProtocol,
-                                                         trustStoreFile = this.trustStoreFile,
-                                                         saslProperties = Some(gssapiSaslProperties))
-    consumers += gssapiSaslConsumer
+    val gssapiSaslProducer = createProducer(configOverrides = gssapiSaslProperties)
+    val gssapiSaslConsumer = createConsumer(configOverrides = gssapiSaslProperties)
     val numRecords = 1000
     var startingOffset = 0
 
@@ -68,9 +59,7 @@ class SaslMultiMechanismConsumerTest extends BaseConsumerTest with SaslSetup {
     plainSaslConsumer.assign(List(tp).asJava)
     plainSaslConsumer.seek(tp, 0)
     consumeAndVerifyRecords(consumer = plainSaslConsumer, numRecords = numRecords, startingOffset = startingOffset)
-    val plainCommitCallback = new CountConsumerCommitCallback()
-    plainSaslConsumer.commitAsync(plainCommitCallback)
-    awaitCommitCallback(plainSaslConsumer, plainCommitCallback)
+    sendAndAwaitAsyncCommit(plainSaslConsumer)
     startingOffset += numRecords
 
     // Test SASL/GSSAPI producer and consumer
@@ -78,9 +67,7 @@ class SaslMultiMechanismConsumerTest extends BaseConsumerTest with SaslSetup {
     gssapiSaslConsumer.assign(List(tp).asJava)
     gssapiSaslConsumer.seek(tp, startingOffset)
     consumeAndVerifyRecords(consumer = gssapiSaslConsumer, numRecords = numRecords, startingOffset = startingOffset)
-    val gssapiCommitCallback = new CountConsumerCommitCallback()
-    gssapiSaslConsumer.commitAsync(gssapiCommitCallback)
-    awaitCommitCallback(gssapiSaslConsumer, gssapiCommitCallback)
+    sendAndAwaitAsyncCommit(gssapiSaslConsumer)
     startingOffset += numRecords
 
     // Test SASL/PLAIN producer and SASL/GSSAPI consumer
@@ -95,6 +82,6 @@ class SaslMultiMechanismConsumerTest extends BaseConsumerTest with SaslSetup {
     plainSaslConsumer.assign(List(tp).asJava)
     plainSaslConsumer.seek(tp, startingOffset)
     consumeAndVerifyRecords(consumer = plainSaslConsumer, numRecords = numRecords, startingOffset = startingOffset)
-
   }
+
 }
diff --git a/core/src/test/scala/integration/kafka/api/SaslOAuthBearerSslEndToEndAuthorizationTest.scala b/core/src/test/scala/integration/kafka/api/SaslOAuthBearerSslEndToEndAuthorizationTest.scala
new file mode 100644
index 0000000000000..4baa61d4c915d
--- /dev/null
+++ b/core/src/test/scala/integration/kafka/api/SaslOAuthBearerSslEndToEndAuthorizationTest.scala
@@ -0,0 +1,27 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package kafka.api
+
+import kafka.utils.JaasTestUtils
+import org.apache.kafka.common.security.auth.KafkaPrincipal
+
+class SaslOAuthBearerSslEndToEndAuthorizationTest extends SaslEndToEndAuthorizationTest {
+  override protected def kafkaClientSaslMechanism = "OAUTHBEARER"
+  override protected def kafkaServerSaslMechanisms = List(kafkaClientSaslMechanism)
+  override val clientPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, JaasTestUtils.KafkaOAuthBearerUser)
+  override val kafkaPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, JaasTestUtils.KafkaOAuthBearerAdmin)
+}
diff --git a/core/src/test/scala/integration/kafka/api/SaslPlainPlaintextConsumerTest.scala b/core/src/test/scala/integration/kafka/api/SaslPlainPlaintextConsumerTest.scala
index 99ddcc32493ee..088f0ac0103f0 100644
--- a/core/src/test/scala/integration/kafka/api/SaslPlainPlaintextConsumerTest.scala
+++ b/core/src/test/scala/integration/kafka/api/SaslPlainPlaintextConsumerTest.scala
@@ -28,6 +28,8 @@ class SaslPlainPlaintextConsumerTest extends BaseConsumerTest with SaslSetup {
   private val kafkaServerJaasEntryName =
     s"${listenerName.value.toLowerCase(Locale.ROOT)}.${JaasTestUtils.KafkaServerContextName}"
   this.serverConfig.setProperty(KafkaConfig.ZkEnableSecureAclsProp, "false")
+  // disable secure acls of zkClient in ZooKeeperTestHarness
+  override protected def zkAclsEnabled = Some(false)
   override protected def securityProtocol = SecurityProtocol.SASL_PLAINTEXT
   override protected lazy val trustStoreFile = Some(File.createTempFile("truststore", ".jks"))
   override protected val serverSaslProperties = Some(kafkaServerSaslProperties(kafkaServerSaslMechanisms, kafkaClientSaslMechanism))
@@ -35,7 +37,7 @@ class SaslPlainPlaintextConsumerTest extends BaseConsumerTest with SaslSetup {
 
   @Before
   override def setUp(): Unit = {
-    startSasl(jaasSections(kafkaServerSaslMechanisms, Some(kafkaClientSaslMechanism), KafkaSasl, kafkaServerJaasEntryName))
+    startSasl(jaasSections(kafkaServerSaslMechanisms, Some(kafkaClientSaslMechanism), Both, kafkaServerJaasEntryName))
     super.setUp()
   }
 
@@ -46,11 +48,11 @@ class SaslPlainPlaintextConsumerTest extends BaseConsumerTest with SaslSetup {
   }
 
   /**
-   * 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() {
-    TestUtils.verifyUnsecureZkAcls(zkUtils)
+  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 c28de3e9fe6ad..356de957e2236 100644
--- a/core/src/test/scala/integration/kafka/api/SaslPlainSslEndToEndAuthorizationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/SaslPlainSslEndToEndAuthorizationTest.scala
@@ -16,46 +16,124 @@
   */
 package kafka.api
 
-import kafka.utils.{JaasTestUtils, TestUtils}
+import java.security.AccessController
+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._
+import org.apache.kafka.common.config.SaslConfigs
 import org.apache.kafka.common.config.internals.BrokerSecurityConfigs
-import org.apache.kafka.common.security.auth.{AuthenticationContext, KafkaPrincipal, KafkaPrincipalBuilder, SaslAuthenticationContext}
+import org.apache.kafka.common.network.ListenerName
+import org.apache.kafka.common.security.auth._
+import org.apache.kafka.common.security.plain.PlainAuthenticateCallback
 import org.junit.Test
 
 object SaslPlainSslEndToEndAuthorizationTest {
+
   class TestPrincipalBuilder extends KafkaPrincipalBuilder {
 
     override def build(context: AuthenticationContext): KafkaPrincipal = {
-      context match {
-        case ctx: SaslAuthenticationContext =>
-          ctx.server.getAuthorizationID match {
-            case JaasTestUtils.KafkaPlainAdmin =>
-              new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "admin")
-            case JaasTestUtils.KafkaPlainUser =>
-              new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "user")
-            case _ =>
-              KafkaPrincipal.ANONYMOUS
-          }
+      context.asInstanceOf[SaslAuthenticationContext].server.getAuthorizationID match {
+        case KafkaPlainAdmin =>
+          new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "admin")
+        case KafkaPlainUser =>
+          new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "user")
+        case _ =>
+          KafkaPrincipal.ANONYMOUS
+      }
+    }
+  }
+
+  object Credentials {
+    val allUsers = Map(KafkaPlainUser -> "user1-password",
+      KafkaPlainUser2 -> KafkaPlainPassword2,
+      KafkaPlainAdmin -> "broker-password")
+  }
+
+  class TestServerCallbackHandler extends AuthenticateCallbackHandler {
+    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])
+          username = callback.asInstanceOf[NameCallback].getDefaultName
+        else if (callback.isInstanceOf[PlainAuthenticateCallback]) {
+          val plainCallback = callback.asInstanceOf[PlainAuthenticateCallback]
+          plainCallback.authenticated(Credentials.allUsers(username) == new String(plainCallback.password))
+        } else
+          throw new UnsupportedCallbackException(callback)
       }
     }
+    def close(): Unit = {}
+  }
+
+  class TestClientCallbackHandler extends AuthenticateCallbackHandler {
+    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) {
+        if (callback.isInstanceOf[NameCallback])
+          callback.asInstanceOf[NameCallback].setName(username)
+        else if (callback.isInstanceOf[PasswordCallback]) {
+          if (username == KafkaPlainUser || username == KafkaPlainAdmin)
+            callback.asInstanceOf[PasswordCallback].setPassword(Credentials.allUsers(username).toCharArray)
+        } else
+          throw new UnsupportedCallbackException(callback)
+      }
+    }
+    def close(): Unit = {}
   }
 }
 
+
+// This test uses SASL callback handler overrides for server connections of Kafka broker
+// and client connections of Kafka producers and consumers. Client connections from Kafka brokers
+// used for inter-broker communication also use custom callback handlers. The second client used in
+// the multi-user test SaslEndToEndAuthorizationTest#testTwoConsumersWithDifferentSaslCredentials uses
+// static JAAS configuration with default callback handlers to test those code paths as well.
 class SaslPlainSslEndToEndAuthorizationTest extends SaslEndToEndAuthorizationTest {
-  import SaslPlainSslEndToEndAuthorizationTest.TestPrincipalBuilder
+  import SaslPlainSslEndToEndAuthorizationTest._
 
   this.serverConfig.setProperty(BrokerSecurityConfigs.PRINCIPAL_BUILDER_CLASS_CONFIG, classOf[TestPrincipalBuilder].getName)
+  this.serverConfig.put(KafkaConfig.SaslClientCallbackHandlerClassProp, classOf[TestClientCallbackHandler].getName)
+  val mechanismPrefix = ListenerName.forSecurityProtocol(SecurityProtocol.SASL_SSL).saslMechanismConfigPrefix("PLAIN")
+  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")
-  override val clientPrincipal = "user"
-  override val kafkaPrincipal = "admin"
+
+  override val clientPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "user")
+  override val kafkaPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "admin")
+
+  override def jaasSections(kafkaServerSaslMechanisms: Seq[String],
+                            kafkaClientSaslMechanism: Option[String],
+                            mode: SaslSetupMode,
+                            kafkaServerEntryName: String): Seq[JaasSection] = {
+    val brokerLogin = PlainLoginModule(KafkaPlainAdmin, "") // Password provided by callback handler
+    val clientLogin = PlainLoginModule(KafkaPlainUser2, KafkaPlainPassword2)
+    Seq(JaasSection(kafkaServerEntryName, Seq(brokerLogin)),
+      JaasSection(KafkaClientContextName, Seq(clientLogin))) ++ zkSections
+  }
 
   /**
    * Checks that secure paths created by broker and acl paths created by AclCommand
    * have expected ACLs.
    */
   @Test
-  def testAcls() {
-    TestUtils.verifySecureZkAcls(zkUtils, 1)
+  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 2f50706517669..eec16239e147d 100644
--- a/core/src/test/scala/integration/kafka/api/SaslScramSslEndToEndAuthorizationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/SaslScramSslEndToEndAuthorizationTest.scala
@@ -16,32 +16,46 @@
   */
 package kafka.api
 
-import org.apache.kafka.common.security.scram.ScramMechanism
+import java.util.Properties
+
 import kafka.utils.JaasTestUtils
-import kafka.utils.ZkUtils
-import scala.collection.JavaConverters._
+import kafka.zk.ConfigEntityChangeNotificationZNode
+import org.apache.kafka.common.security.auth.KafkaPrincipal
+import org.apache.kafka.common.security.scram.internals.ScramMechanism
+import org.apache.kafka.test.TestSslUtils
+
+import scala.jdk.CollectionConverters._
 import org.junit.Before
 
 class SaslScramSslEndToEndAuthorizationTest extends SaslEndToEndAuthorizationTest {
   override protected def kafkaClientSaslMechanism = "SCRAM-SHA-256"
   override protected def kafkaServerSaslMechanisms = ScramMechanism.mechanismNames.asScala.toList
-  override val clientPrincipal = JaasTestUtils.KafkaScramUser
-  override val kafkaPrincipal = JaasTestUtils.KafkaScramAdmin
-  private val clientPassword = JaasTestUtils.KafkaScramPassword
+  override val clientPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, JaasTestUtils.KafkaScramUser)
+  override val kafkaPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, JaasTestUtils.KafkaScramAdmin)
   private val kafkaPassword = JaasTestUtils.KafkaScramAdminPassword
 
-  override def configureSecurityBeforeServersStart() {
+  override def configureSecurityBeforeServersStart(): Unit = {
     super.configureSecurityBeforeServersStart()
-    zkUtils.makeSurePersistentPathExists(ZkUtils.ConfigChangesPath)
+    zkClient.makeSurePersistentPathExists(ConfigEntityChangeNotificationZNode.path)
     // Create broker credentials before starting brokers
-    createScramCredentials(zkConnect, kafkaPrincipal, kafkaPassword)
+    createScramCredentials(zkConnect, kafkaPrincipal.getName, kafkaPassword)
+    TestSslUtils.convertToPemWithoutFiles(producerConfig)
+    TestSslUtils.convertToPemWithoutFiles(consumerConfig)
+    TestSslUtils.convertToPemWithoutFiles(adminClientConfig)
   }
 
+  override def configureListeners(props: collection.Seq[Properties]): Unit = {
+    props.foreach(TestSslUtils.convertToPemWithoutFiles)
+    super.configureListeners(props)
+  }
+
+  override def createPrivilegedAdminClient() = createScramAdminClient(kafkaClientSaslMechanism, kafkaPrincipal.getName, kafkaPassword)
+
   @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, clientPrincipal, clientPassword)
-    createScramCredentials(zkConnect, JaasTestUtils.KafkaScramUser2, JaasTestUtils.KafkaScramPassword2)
+    createScramCredentialsViaPrivilegedAdminClient(JaasTestUtils.KafkaScramUser, JaasTestUtils.KafkaScramPassword)
+    createScramCredentialsViaPrivilegedAdminClient(JaasTestUtils.KafkaScramUser2, JaasTestUtils.KafkaScramPassword2)
   }
 }
diff --git a/core/src/test/scala/integration/kafka/api/SaslSetup.scala b/core/src/test/scala/integration/kafka/api/SaslSetup.scala
index 51deb859fbceb..2bd55b855a929 100644
--- a/core/src/test/scala/integration/kafka/api/SaslSetup.scala
+++ b/core/src/test/scala/integration/kafka/api/SaslSetup.scala
@@ -18,19 +18,28 @@
 package kafka.api
 
 import java.io.File
+import java.util
 import java.util.Properties
+
 import javax.security.auth.login.Configuration
 
-import kafka.admin.ConfigCommand
+import scala.collection.Seq
 import kafka.security.minikdc.MiniKdc
-import kafka.server.KafkaConfig
+import kafka.server.{ConfigType, KafkaConfig}
 import kafka.utils.JaasTestUtils.{JaasSection, Krb5LoginModule, ZkDigestModule}
 import kafka.utils.{JaasTestUtils, TestUtils}
+import kafka.zk.{AdminZkClient, KafkaZkClient}
+import org.apache.kafka.clients.admin.{Admin, AdminClientConfig, ScramCredentialInfo, UserScramCredentialAlteration, UserScramCredentialUpsertion, ScramMechanism => PublicScramMechanism}
 import org.apache.kafka.common.config.SaslConfigs
 import org.apache.kafka.common.config.internals.BrokerSecurityConfigs
 import org.apache.kafka.common.security.JaasUtils
+import org.apache.kafka.common.security.auth.SecurityProtocol
 import org.apache.kafka.common.security.authenticator.LoginManager
-import org.apache.kafka.common.security.scram.ScramMechanism
+import org.apache.kafka.common.security.scram.internals.{ScramCredentialUtils, ScramFormatter, ScramMechanism}
+import org.apache.kafka.common.utils.Time
+import org.apache.zookeeper.client.ZKClientConfig
+
+import scala.jdk.CollectionConverters._
 
 /*
  * Implements an enumeration for the modes enabled here:
@@ -51,7 +60,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 {
@@ -59,12 +68,7 @@ trait SaslSetup {
       case _ => false
     })
     if (hasKerberos) {
-      val (serverKeytabFile, clientKeytabFile) = maybeCreateEmptyKeytabFiles()
-      kdc = new MiniKdc(kdcConf, workDir)
-      kdc.start()
-      kdc.createPrincipal(serverKeytabFile, JaasTestUtils.KafkaServerPrincipalUnqualifiedName + "/localhost")
-      kdc.createPrincipal(clientKeytabFile,
-        JaasTestUtils.KafkaClientPrincipalUnqualifiedName, JaasTestUtils.KafkaClientPrincipalUnqualifiedName2)
+      initializeKerberos()
     }
     writeJaasConfigurationToFile(jaasSections)
     val hasZk = jaasSections.exists(_.modules.exists {
@@ -75,6 +79,15 @@ trait SaslSetup {
       System.setProperty("zookeeper.authProvider.1", "org.apache.zookeeper.server.auth.SASLAuthenticationProvider")
   }
 
+  protected def initializeKerberos(): Unit = {
+    val (serverKeytabFile, clientKeytabFile) = maybeCreateEmptyKeytabFiles()
+    kdc = new MiniKdc(kdcConf, workDir)
+    kdc.start()
+    kdc.createPrincipal(serverKeytabFile, JaasTestUtils.KafkaServerPrincipalUnqualifiedName + "/localhost")
+    kdc.createPrincipal(clientKeytabFile,
+      JaasTestUtils.KafkaClientPrincipalUnqualifiedName, JaasTestUtils.KafkaClientPrincipalUnqualifiedName2)
+  }
+
   /** Return a tuple with the path to the server keytab file and client keytab file */
   protected def maybeCreateEmptyKeytabFiles(): (File, File) = {
     if (serverKeytabFile.isEmpty)
@@ -89,7 +102,7 @@ trait SaslSetup {
                              mode: SaslSetupMode = Both,
                              kafkaServerEntryName: String = JaasTestUtils.KafkaServerContextName): Seq[JaasSection] = {
     val hasKerberos = mode != ZkSasl &&
-      (kafkaServerSaslMechanisms.contains("GSSAPI") || kafkaClientSaslMechanism.exists(_ == "GSSAPI"))
+      (kafkaServerSaslMechanisms.contains("GSSAPI") || kafkaClientSaslMechanism.contains("GSSAPI"))
     if (hasKerberos)
       maybeCreateEmptyKeytabFiles()
     mode match {
@@ -102,14 +115,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
@@ -134,16 +147,69 @@ trait SaslSetup {
     props
   }
 
-  def jaasClientLoginModule(clientSaslMechanism: String): String =
-    JaasTestUtils.clientLoginModule(clientSaslMechanism, clientKeytabFile)
+  def jaasClientLoginModule(clientSaslMechanism: String, serviceName: Option[String] = None): String = {
+    if (serviceName.isDefined)
+      JaasTestUtils.clientLoginModule(clientSaslMechanism, clientKeytabFile, serviceName.get)
+    else
+      JaasTestUtils.clientLoginModule(clientSaslMechanism, clientKeytabFile)
+  }
+
+  def jaasScramClientLoginModule(clientSaslScramMechanism: String, scramUser: String, scramPassword: String): String = {
+    JaasTestUtils.scramClientLoginModule(clientSaslScramMechanism, scramUser, scramPassword)
+  }
+
+  def createPrivilegedAdminClient(): Admin = {
+    // create an admin client instance that is authorized to create credentials
+    throw new UnsupportedOperationException("Must implement this if a test needs to use it")
+  }
+
+  def createAdminClient(brokerList: String, securityProtocol: SecurityProtocol, trustStoreFile: Option[File],
+                        clientSaslProperties: Option[Properties], scramMechanism: String, user: String, password: String) : Admin = {
+    val config = new util.HashMap[String, Object]
+    config.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList)
+    val securityProps: util.Map[Object, Object] =
+      TestUtils.adminClientSecurityConfigs(securityProtocol, trustStoreFile, clientSaslProperties)
+    securityProps.forEach { (key, value) => config.put(key.asInstanceOf[String], value) }
+    config.put(SaslConfigs.SASL_JAAS_CONFIG, jaasScramClientLoginModule(scramMechanism, user, password))
+    Admin.create(config)
+  }
+
+  def createScramCredentialsViaPrivilegedAdminClient(userName: String, password: String): Unit = {
+    val privilegedAdminClient = createPrivilegedAdminClient() // must explicitly implement this method
+    try {
+      // create the SCRAM credential for the given user
+      createScramCredentials(privilegedAdminClient, userName, password)
+    } finally {
+      privilegedAdminClient.close()
+    }
+  }
+
+    def createScramCredentials(adminClient: Admin, userName: String, password: String): Unit = {
+    val results = adminClient.alterUserScramCredentials(PublicScramMechanism.values().filter(_ != PublicScramMechanism.UNKNOWN).map(mechanism =>
+      new UserScramCredentialUpsertion(userName, new ScramCredentialInfo(mechanism, 4096), password)
+        .asInstanceOf[UserScramCredentialAlteration]).toList.asJava)
+    results.all.get
+  }
 
   def createScramCredentials(zkConnect: String, userName: String, password: String): Unit = {
-    val credentials = ScramMechanism.values.map(m => s"${m.mechanismName}=[iterations=4096,password=$password]")
-    val args = Array("--zookeeper", zkConnect,
-      "--alter", "--add-config", credentials.mkString(","),
-      "--entity-type", "users",
-      "--entity-name", userName)
-    ConfigCommand.main(args)
+    val zkClientConfig = new ZKClientConfig()
+    val zkClient = KafkaZkClient(
+      zkConnect, JaasUtils.isZkSaslEnabled || KafkaConfig.zkTlsClientAuthEnabled(zkClientConfig), 30000, 30000,
+      Int.MaxValue, Time.SYSTEM, zkClientConfig = Some(zkClientConfig))
+    val adminZkClient = new AdminZkClient(zkClient)
+
+    val entityType = ConfigType.User
+    val entityName = userName
+    val configs = adminZkClient.fetchEntityConfig(entityType, entityName)
+
+    ScramMechanism.values().foreach(mechanism => {
+      val credential = new ScramFormatter(mechanism).generateCredential(password, 4096)
+      val credentialString = ScramCredentialUtils.credentialToString(credential)
+      configs.setProperty(mechanism.mechanismName, credentialString)
+    })
+
+    adminZkClient.changeConfigs(entityType, entityName, configs)
+    zkClient.close()
   }
 
 }
diff --git a/core/src/test/scala/integration/kafka/api/SaslSslAdminClientIntegrationTest.scala b/core/src/test/scala/integration/kafka/api/SaslSslAdminClientIntegrationTest.scala
deleted file mode 100644
index 099af5247d366..0000000000000
--- a/core/src/test/scala/integration/kafka/api/SaslSslAdminClientIntegrationTest.scala
+++ /dev/null
@@ -1,279 +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.io.File
-
-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.utils.{CoreUtils, JaasTestUtils, TestUtils}
-import org.apache.kafka.clients.admin.{AdminClient, CreateAclsOptions, DeleteAclsOptions}
-import org.apache.kafka.common.acl.{AccessControlEntry, AccessControlEntryFilter, AclBinding, AclBindingFilter, AclOperation, AclPermissionType}
-import org.apache.kafka.common.errors.{ClusterAuthorizationException, InvalidRequestException}
-import org.apache.kafka.common.resource.{Resource, ResourceFilter, ResourceType}
-import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
-import org.junit.Assert.assertEquals
-import org.junit.{After, Assert, Before, Test}
-
-import scala.collection.JavaConverters._
-import scala.util.{Failure, Success, Try}
-
-class SaslSslAdminClientIntegrationTest extends AdminClientIntegrationTest with SaslSetup {
-  this.serverConfig.setProperty(KafkaConfig.ZkEnableSecureAclsProp, "true")
-  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() {
-    val authorizer = CoreUtils.createObject[Authorizer](classOf[SimpleAclAuthorizer].getName)
-    try {
-      authorizer.configure(this.configs.head.originals())
-      authorizer.addAcls(Set(new AuthAcl(AuthAcl.WildCardPrincipal, Allow,
-                             AuthAcl.WildCardHost, All)), new AuthResource(Topic, "*"))
-      authorizer.addAcls(Set(new AuthAcl(AuthAcl.WildCardPrincipal, Allow,
-                             AuthAcl.WildCardHost, All)), new AuthResource(Group, "*"))
-
-      authorizer.addAcls(Set(clusterAcl(Allow, Create),
-                             clusterAcl(Allow, Delete),
-                             clusterAcl(Allow, ClusterAction),
-                             clusterAcl(Allow, AlterConfigs),
-                             clusterAcl(Allow, Alter)),
-                             AuthResource.ClusterResource)
-    } finally {
-      authorizer.close()
-    }
-  }
-
-  @Before
-  override def setUp(): Unit = {
-    startSasl(jaasSections(Seq("GSSAPI"), Some("GSSAPI"), Both, JaasTestUtils.KafkaServerContextName))
-    super.setUp()
-  }
-
-  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 = {
-    val acls = Set(clusterAcl(permissionType, operation))
-    val authorizer = servers.head.apis.authorizer.get
-    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 = {
-    val acls = Set(clusterAcl(permissionType, operation))
-    val authorizer = servers.head.apis.authorizer.get
-    val prevAcls = authorizer.getAcls(AuthResource.ClusterResource)
-    Assert.assertTrue(authorizer.removeAcls(acls, AuthResource.ClusterResource))
-    TestUtils.waitAndVerifyAcls(prevAcls -- acls, authorizer, AuthResource.ClusterResource)
-  }
-
-  @After
-  override def tearDown(): Unit = {
-    super.tearDown()
-    closeSasl()
-  }
-
-  val acl2 = new AclBinding(new Resource(ResourceType.TOPIC, "mytopic2"),
-    new AccessControlEntry("User:ANONYMOUS", "*", AclOperation.WRITE, AclPermissionType.ALLOW))
-  val acl3 = new AclBinding(new Resource(ResourceType.TOPIC, "mytopic3"),
-    new AccessControlEntry("User:ANONYMOUS", "*", AclOperation.READ, AclPermissionType.ALLOW))
-  val fooAcl = new AclBinding(new Resource(ResourceType.TOPIC, "foobar"),
-    new AccessControlEntry("User:ANONYMOUS", "*", AclOperation.READ, AclPermissionType.ALLOW))
-  val transactionalIdAcl = new AclBinding(new Resource(ResourceType.TRANSACTIONAL_ID, "transactional_id"),
-    new AccessControlEntry("User:ANONYMOUS", "*", AclOperation.WRITE, AclPermissionType.ALLOW))
-  val groupAcl = new AclBinding(new Resource(ResourceType.GROUP, "*"),
-    new AccessControlEntry("User:*", "*", AclOperation.ALL, AclPermissionType.ALLOW))
-
-  @Test
-  override def testAclOperations(): Unit = {
-    client = AdminClient.create(createConfig())
-    assertEquals(7, client.describeAcls(AclBindingFilter.ANY).values.get().size)
-    val results = client.createAcls(List(acl2, acl3).asJava)
-    assertEquals(Set(acl2, acl3), results.values.keySet().asScala)
-    results.values.values().asScala.foreach(value => value.get)
-    val aclUnknown = new AclBinding(new Resource(ResourceType.TOPIC, "mytopic3"),
-      new AccessControlEntry("User:ANONYMOUS", "*", AclOperation.UNKNOWN, AclPermissionType.ALLOW))
-    val results2 = client.createAcls(List(aclUnknown).asJava)
-    assertEquals(Set(aclUnknown), results2.values.keySet().asScala)
-    assertFutureExceptionTypeEquals(results2.all, classOf[InvalidRequestException])
-    val results3 = client.deleteAcls(List(ACL1.toFilter, acl2.toFilter, acl3.toFilter).asJava).values
-    assertEquals(Set(ACL1.toFilter, acl2.toFilter, acl3.toFilter), results3.keySet.asScala)
-    assertEquals(0, results3.get(ACL1.toFilter).get.values.size())
-    assertEquals(Set(acl2), results3.get(acl2.toFilter).get.values.asScala.map(_.binding).toSet)
-    assertEquals(Set(acl3), results3.get(acl3.toFilter).get.values.asScala.map(_.binding).toSet)
-  }
-
-  def waitForDescribeAcls(client: AdminClient, filter: AclBindingFilter, acls: Set[AclBinding]): Unit = {
-    TestUtils.waitUntilTrue(() => {
-      val results = client.describeAcls(filter).values.get()
-      acls == results.asScala.toSet
-    }, s"timed out waiting for ACLs $acls")
-  }
-
-  @Test
-  def testAclOperations2(): Unit = {
-    client = AdminClient.create(createConfig())
-    val results = client.createAcls(List(acl2, acl2, transactionalIdAcl).asJava)
-    assertEquals(Set(acl2, acl2, transactionalIdAcl), results.values.keySet.asScala)
-    results.all.get()
-    waitForDescribeAcls(client, acl2.toFilter, Set(acl2))
-    waitForDescribeAcls(client, transactionalIdAcl.toFilter, Set(transactionalIdAcl))
-
-    val filterA = new AclBindingFilter(new ResourceFilter(ResourceType.GROUP, null), AccessControlEntryFilter.ANY)
-    val filterB = new AclBindingFilter(new ResourceFilter(ResourceType.TOPIC, "mytopic2"), AccessControlEntryFilter.ANY)
-    val filterC = new AclBindingFilter(new ResourceFilter(ResourceType.TRANSACTIONAL_ID, null), AccessControlEntryFilter.ANY)
-
-    waitForDescribeAcls(client, filterA, Set(groupAcl))
-    waitForDescribeAcls(client, filterC, Set(transactionalIdAcl))
-
-    val results2 = client.deleteAcls(List(filterA, filterB, filterC).asJava, new DeleteAclsOptions())
-    assertEquals(Set(filterA, filterB, filterC), results2.values.keySet.asScala)
-    assertEquals(Set(groupAcl), results2.values.get(filterA).get.values.asScala.map(_.binding).toSet)
-    assertEquals(Set(transactionalIdAcl), results2.values.get(filterC).get.values.asScala.map(_.binding).toSet)
-    assertEquals(Set(acl2), results2.values.get(filterB).get.values.asScala.map(_.binding).toSet)
-
-    waitForDescribeAcls(client, filterB, Set())
-    waitForDescribeAcls(client, filterC, Set())
-  }
-
-  @Test
-  def testAttemptToCreateInvalidAcls(): Unit = {
-    client = AdminClient.create(createConfig())
-    val clusterAcl = new AclBinding(new Resource(ResourceType.CLUSTER, "foobar"),
-      new AccessControlEntry("User:ANONYMOUS", "*", AclOperation.READ, AclPermissionType.ALLOW))
-    val emptyResourceNameAcl = new AclBinding(new Resource(ResourceType.TOPIC, ""),
-      new AccessControlEntry("User:ANONYMOUS", "*", AclOperation.READ, AclPermissionType.ALLOW))
-    val results = client.createAcls(List(clusterAcl, emptyResourceNameAcl).asJava, new CreateAclsOptions())
-    assertEquals(Set(clusterAcl, emptyResourceNameAcl), results.values.keySet().asScala)
-    assertFutureExceptionTypeEquals(results.values.get(clusterAcl), classOf[InvalidRequestException])
-    assertFutureExceptionTypeEquals(results.values.get(emptyResourceNameAcl), classOf[InvalidRequestException])
-  }
-
-  private def verifyCauseIsClusterAuth(e: Throwable): Unit = {
-    if (!e.getCause.isInstanceOf[ClusterAuthorizationException]) {
-      throw e.getCause
-    }
-  }
-
-  private def testAclCreateGetDelete(expectAuth: Boolean): Unit = {
-    TestUtils.waitUntilTrue(() => {
-      val result = client.createAcls(List(fooAcl, transactionalIdAcl).asJava, new CreateAclsOptions)
-      if (expectAuth) {
-        Try(result.all.get) match {
-          case Failure(e) =>
-            verifyCauseIsClusterAuth(e)
-            false
-          case Success(_) => true
-        }
-      } else {
-        Try(result.all.get) match {
-          case Failure(e) =>
-            verifyCauseIsClusterAuth(e)
-            true
-          case Success(_) => false
-        }
-      }
-    }, "timed out waiting for createAcls to " + (if (expectAuth) "succeed" else "fail"))
-    if (expectAuth) {
-      waitForDescribeAcls(client, fooAcl.toFilter, Set(fooAcl))
-      waitForDescribeAcls(client, transactionalIdAcl.toFilter, Set(transactionalIdAcl))
-    }
-    TestUtils.waitUntilTrue(() => {
-      val result = client.deleteAcls(List(fooAcl.toFilter, transactionalIdAcl.toFilter).asJava, new DeleteAclsOptions)
-      if (expectAuth) {
-        Try(result.all.get) match {
-          case Failure(e) =>
-            verifyCauseIsClusterAuth(e)
-            false
-          case Success(_) => true
-        }
-      } else {
-        Try(result.all.get) match {
-          case Failure(e) =>
-            verifyCauseIsClusterAuth(e)
-            true
-          case Success(_) =>
-            assertEquals(Set(fooAcl, transactionalIdAcl), result.values.keySet)
-            assertEquals(Set(fooAcl), result.values.get(fooAcl.toFilter).get.values.asScala.map(_.binding).toSet)
-            assertEquals(Set(transactionalIdAcl),
-              result.values.get(transactionalIdAcl.toFilter).get.values.asScala.map(_.binding).toSet)
-            true
-        }
-      }
-    }, "timed out waiting for deleteAcls to " + (if (expectAuth) "succeed" else "fail"))
-    if (expectAuth) {
-      waitForDescribeAcls(client, fooAcl.toFilter, Set.empty)
-      waitForDescribeAcls(client, transactionalIdAcl.toFilter, Set.empty)
-    }
-  }
-
-  private def testAclGet(expectAuth: Boolean): Unit = {
-    TestUtils.waitUntilTrue(() => {
-      val userAcl = new AclBinding(new Resource(ResourceType.TOPIC, "*"),
-        new AccessControlEntry("User:*", "*", AclOperation.ALL, AclPermissionType.ALLOW))
-      val results = client.describeAcls(userAcl.toFilter)
-      if (expectAuth) {
-        Try(results.values.get) match {
-          case Failure(e) =>
-            verifyCauseIsClusterAuth(e)
-            false
-          case Success(acls) => Set(userAcl).equals(acls.asScala.toSet)
-        }
-      } else {
-        Try(results.values.get) match {
-          case Failure(e) =>
-            verifyCauseIsClusterAuth(e)
-            true
-          case Success(_) => false
-        }
-      }
-    }, "timed out waiting for describeAcls to " + (if (expectAuth) "succeed" else "fail"))
-  }
-
-  @Test
-  def testAclAuthorizationDenied(): Unit = {
-    client = AdminClient.create(createConfig())
-
-    // Test that we cannot create or delete ACLs when Alter is denied.
-    addClusterAcl(Deny, Alter)
-    testAclGet(expectAuth = true)
-    testAclCreateGetDelete(expectAuth = false)
-
-    // Test that we cannot do anything with ACLs when Describe and Alter are denied.
-    addClusterAcl(Deny, Describe)
-    testAclGet(expectAuth = false)
-    testAclCreateGetDelete(expectAuth = false)
-
-    // Test that we can create, delete, and get ACLs with the default ACLs.
-    removeClusterAcl(Deny, Describe)
-    removeClusterAcl(Deny, Alter)
-    testAclGet(expectAuth = true)
-    testAclCreateGetDelete(expectAuth = true)
-
-    // Test that we can't do anything with ACLs without the Allow Alter ACL in place.
-    removeClusterAcl(Allow, Alter)
-    removeClusterAcl(Allow, Delete)
-    testAclGet(expectAuth = false)
-    testAclCreateGetDelete(expectAuth = false)
-
-    // Test that we can describe, but not alter ACLs, with only the Allow Describe ACL in place.
-    addClusterAcl(Allow, Describe)
-    testAclGet(expectAuth = true)
-    testAclCreateGetDelete(expectAuth = false)
-  }
-}
diff --git a/core/src/test/scala/integration/kafka/api/SaslSslAdminIntegrationTest.scala b/core/src/test/scala/integration/kafka/api/SaslSslAdminIntegrationTest.scala
new file mode 100644
index 0000000000000..c4e291f3ab92b
--- /dev/null
+++ b/core/src/test/scala/integration/kafka/api/SaslSslAdminIntegrationTest.scala
@@ -0,0 +1,533 @@
+/**
+  * Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE
+  * file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file
+  * to You under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the
+  * License. You may obtain a copy of the License at
+  *
+  * http://www.apache.org/licenses/LICENSE-2.0
+  *
+  * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+  * an "AS IS" BASIS, WITHOUT WARRANTIES 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 kafka.log.LogConfig
+import kafka.server.{Defaults, KafkaConfig}
+import kafka.utils.{CoreUtils, JaasTestUtils, TestUtils}
+import kafka.utils.TestUtils._
+import org.apache.kafka.clients.admin._
+import org.apache.kafka.common.acl._
+import org.apache.kafka.common.acl.AclOperation.{ALTER, ALTER_CONFIGS, CLUSTER_ACTION, CREATE, DELETE, DESCRIBE}
+import org.apache.kafka.common.acl.AclPermissionType.{ALLOW, DENY}
+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, assertTrue}
+import org.junit.{After, Assert, Before, Test}
+
+import scala.annotation.nowarn
+import scala.jdk.CollectionConverters._
+import scala.collection.Seq
+import scala.compat.java8.OptionConverters._
+import scala.concurrent.ExecutionException
+import scala.util.{Failure, Success, Try}
+
+abstract class AuthorizationAdmin {
+  def authorizerClassName: String
+  def initializeAcls(): Unit
+  def addClusterAcl(permissionType: AclPermissionType, operation: AclOperation): Unit
+  def removeClusterAcl(permissionType: AclPermissionType, operation: AclOperation): Unit
+}
+
+// Note: this test currently uses the deprecated SimpleAclAuthorizer to ensure we have test coverage
+// It must be replaced with the new AclAuthorizer when SimpleAclAuthorizer is removed
+class SaslSslAdminIntegrationTest extends BaseAdminIntegrationTest with SaslSetup {
+  @nowarn("cat=deprecation")
+  val authorizationAdmin: AuthorizationAdmin = new LegacyAuthorizationAdmin
+  this.serverConfig.setProperty(KafkaConfig.ZkEnableSecureAclsProp, "true")
+
+  override protected def securityProtocol = SecurityProtocol.SASL_SSL
+  override protected lazy val trustStoreFile = Some(File.createTempFile("truststore", ".jks"))
+
+  override def generateConfigs: Seq[KafkaConfig] = {
+    this.serverConfig.setProperty(KafkaConfig.AuthorizerClassNameProp, authorizationAdmin.authorizerClassName)
+    super.generateConfigs
+  }
+
+  override def configureSecurityBeforeServersStart(): Unit = {
+    authorizationAdmin.initializeAcls()
+  }
+
+  @Before
+  override def setUp(): Unit = {
+    setUpSasl()
+    super.setUp()
+  }
+
+  def setUpSasl(): Unit = {
+    startSasl(jaasSections(Seq("GSSAPI"), Some("GSSAPI"), Both, JaasTestUtils.KafkaServerContextName))
+  }
+
+  @After
+  override def tearDown(): Unit = {
+    super.tearDown()
+    closeSasl()
+  }
+
+  val anyAcl = new AclBinding(new ResourcePattern(ResourceType.TOPIC, "*", PatternType.LITERAL),
+    new AccessControlEntry("User:*", "*", AclOperation.ALL, AclPermissionType.ALLOW))
+  val acl2 = new AclBinding(new ResourcePattern(ResourceType.TOPIC, "mytopic2", PatternType.LITERAL),
+    new AccessControlEntry("User:ANONYMOUS", "*", AclOperation.WRITE, AclPermissionType.ALLOW))
+  val acl3 = new AclBinding(new ResourcePattern(ResourceType.TOPIC, "mytopic3", PatternType.LITERAL),
+    new AccessControlEntry("User:ANONYMOUS", "*", AclOperation.READ, AclPermissionType.ALLOW))
+  val fooAcl = new AclBinding(new ResourcePattern(ResourceType.TOPIC, "foobar", PatternType.LITERAL),
+    new AccessControlEntry("User:ANONYMOUS", "*", AclOperation.READ, AclPermissionType.ALLOW))
+  val prefixAcl = new AclBinding(new ResourcePattern(ResourceType.TOPIC, "mytopic", PatternType.PREFIXED),
+    new AccessControlEntry("User:ANONYMOUS", "*", AclOperation.READ, AclPermissionType.ALLOW))
+  val transactionalIdAcl = new AclBinding(new ResourcePattern(ResourceType.TRANSACTIONAL_ID, "transactional_id", PatternType.LITERAL),
+    new AccessControlEntry("User:ANONYMOUS", "*", AclOperation.WRITE, AclPermissionType.ALLOW))
+  val groupAcl = new AclBinding(new ResourcePattern(ResourceType.GROUP, "*", PatternType.LITERAL),
+    new AccessControlEntry("User:*", "*", AclOperation.ALL, AclPermissionType.ALLOW))
+
+  @Test
+  def testAclOperations(): Unit = {
+    client = Admin.create(createConfig)
+    val acl = new AclBinding(new ResourcePattern(ResourceType.TOPIC, "mytopic3", PatternType.LITERAL),
+      new AccessControlEntry("User:ANONYMOUS", "*", AclOperation.DESCRIBE, AclPermissionType.ALLOW))
+    assertEquals(7, getAcls(AclBindingFilter.ANY).size)
+    val results = client.createAcls(List(acl2, acl3).asJava)
+    assertEquals(Set(acl2, acl3), results.values.keySet().asScala)
+    results.values.values.forEach(value => value.get)
+    val aclUnknown = new AclBinding(new ResourcePattern(ResourceType.TOPIC, "mytopic3", PatternType.LITERAL),
+      new AccessControlEntry("User:ANONYMOUS", "*", AclOperation.UNKNOWN, AclPermissionType.ALLOW))
+    val results2 = client.createAcls(List(aclUnknown).asJava)
+    assertEquals(Set(aclUnknown), results2.values.keySet().asScala)
+    assertFutureExceptionTypeEquals(results2.all, classOf[InvalidRequestException])
+    val results3 = client.deleteAcls(List(acl.toFilter, acl2.toFilter, acl3.toFilter).asJava).values
+    assertEquals(Set(acl.toFilter, acl2.toFilter, acl3.toFilter), results3.keySet.asScala)
+    assertEquals(0, results3.get(acl.toFilter).get.values.size())
+    assertEquals(Set(acl2), results3.get(acl2.toFilter).get.values.asScala.map(_.binding).toSet)
+    assertEquals(Set(acl3), results3.get(acl3.toFilter).get.values.asScala.map(_.binding).toSet)
+  }
+
+  @Test
+  def testAclOperations2(): Unit = {
+    client = Admin.create(createConfig)
+    val results = client.createAcls(List(acl2, acl2, transactionalIdAcl).asJava)
+    assertEquals(Set(acl2, acl2, transactionalIdAcl), results.values.keySet.asScala)
+    results.all.get()
+    waitForDescribeAcls(client, acl2.toFilter, Set(acl2))
+    waitForDescribeAcls(client, transactionalIdAcl.toFilter, Set(transactionalIdAcl))
+
+    val filterA = new AclBindingFilter(new ResourcePatternFilter(ResourceType.GROUP, null, PatternType.LITERAL), AccessControlEntryFilter.ANY)
+    val filterB = new AclBindingFilter(new ResourcePatternFilter(ResourceType.TOPIC, "mytopic2", PatternType.LITERAL), AccessControlEntryFilter.ANY)
+    val filterC = new AclBindingFilter(new ResourcePatternFilter(ResourceType.TRANSACTIONAL_ID, null, PatternType.LITERAL), AccessControlEntryFilter.ANY)
+
+    waitForDescribeAcls(client, filterA, Set(groupAcl))
+    waitForDescribeAcls(client, filterC, Set(transactionalIdAcl))
+
+    val results2 = client.deleteAcls(List(filterA, filterB, filterC).asJava, new DeleteAclsOptions())
+    assertEquals(Set(filterA, filterB, filterC), results2.values.keySet.asScala)
+    assertEquals(Set(groupAcl), results2.values.get(filterA).get.values.asScala.map(_.binding).toSet)
+    assertEquals(Set(transactionalIdAcl), results2.values.get(filterC).get.values.asScala.map(_.binding).toSet)
+    assertEquals(Set(acl2), results2.values.get(filterB).get.values.asScala.map(_.binding).toSet)
+
+    waitForDescribeAcls(client, filterB, Set())
+    waitForDescribeAcls(client, filterC, Set())
+  }
+
+  @Test
+  def testAclDescribe(): Unit = {
+    client = Admin.create(createConfig)
+    ensureAcls(Set(anyAcl, acl2, fooAcl, prefixAcl))
+
+    val allTopicAcls = new AclBindingFilter(new ResourcePatternFilter(ResourceType.TOPIC, null, PatternType.ANY), AccessControlEntryFilter.ANY)
+    val allLiteralTopicAcls = new AclBindingFilter(new ResourcePatternFilter(ResourceType.TOPIC, null, PatternType.LITERAL), AccessControlEntryFilter.ANY)
+    val allPrefixedTopicAcls = new AclBindingFilter(new ResourcePatternFilter(ResourceType.TOPIC, null, PatternType.PREFIXED), AccessControlEntryFilter.ANY)
+    val literalMyTopic2Acls = new AclBindingFilter(new ResourcePatternFilter(ResourceType.TOPIC, "mytopic2", PatternType.LITERAL), AccessControlEntryFilter.ANY)
+    val prefixedMyTopicAcls = new AclBindingFilter(new ResourcePatternFilter(ResourceType.TOPIC, "mytopic", PatternType.PREFIXED), AccessControlEntryFilter.ANY)
+    val allMyTopic2Acls = new AclBindingFilter(new ResourcePatternFilter(ResourceType.TOPIC, "mytopic2", PatternType.MATCH), AccessControlEntryFilter.ANY)
+    val allFooTopicAcls = new AclBindingFilter(new ResourcePatternFilter(ResourceType.TOPIC, "foobar", PatternType.MATCH), AccessControlEntryFilter.ANY)
+
+    assertEquals(Set(anyAcl), getAcls(anyAcl.toFilter))
+    assertEquals(Set(prefixAcl), getAcls(prefixAcl.toFilter))
+    assertEquals(Set(acl2), getAcls(acl2.toFilter))
+    assertEquals(Set(fooAcl), getAcls(fooAcl.toFilter))
+
+    assertEquals(Set(acl2), getAcls(literalMyTopic2Acls))
+    assertEquals(Set(prefixAcl), getAcls(prefixedMyTopicAcls))
+    assertEquals(Set(anyAcl, acl2, fooAcl), getAcls(allLiteralTopicAcls))
+    assertEquals(Set(prefixAcl), getAcls(allPrefixedTopicAcls))
+    assertEquals(Set(anyAcl, acl2, prefixAcl), getAcls(allMyTopic2Acls))
+    assertEquals(Set(anyAcl, fooAcl), getAcls(allFooTopicAcls))
+    assertEquals(Set(anyAcl, acl2, fooAcl, prefixAcl), getAcls(allTopicAcls))
+  }
+
+  @Test
+  def testAclDelete(): Unit = {
+    client = Admin.create(createConfig)
+    ensureAcls(Set(anyAcl, acl2, fooAcl, prefixAcl))
+
+    val allTopicAcls = new AclBindingFilter(new ResourcePatternFilter(ResourceType.TOPIC, null, PatternType.MATCH), AccessControlEntryFilter.ANY)
+    val allLiteralTopicAcls = new AclBindingFilter(new ResourcePatternFilter(ResourceType.TOPIC, null, PatternType.LITERAL), AccessControlEntryFilter.ANY)
+    val allPrefixedTopicAcls = new AclBindingFilter(new ResourcePatternFilter(ResourceType.TOPIC, null, PatternType.PREFIXED), AccessControlEntryFilter.ANY)
+
+    // Delete only ACLs on literal 'mytopic2' topic
+    var deleted = client.deleteAcls(List(acl2.toFilter).asJava).all().get().asScala.toSet
+    assertEquals(Set(acl2), deleted)
+    assertEquals(Set(anyAcl, fooAcl, prefixAcl), getAcls(allTopicAcls))
+
+    ensureAcls(deleted)
+
+    // Delete only ACLs on literal '*' topic
+    deleted = client.deleteAcls(List(anyAcl.toFilter).asJava).all().get().asScala.toSet
+    assertEquals(Set(anyAcl), deleted)
+    assertEquals(Set(acl2, fooAcl, prefixAcl), getAcls(allTopicAcls))
+
+    ensureAcls(deleted)
+
+    // Delete only ACLs on specific prefixed 'mytopic' topics:
+    deleted = client.deleteAcls(List(prefixAcl.toFilter).asJava).all().get().asScala.toSet
+    assertEquals(Set(prefixAcl), deleted)
+    assertEquals(Set(anyAcl, acl2, fooAcl), getAcls(allTopicAcls))
+
+    ensureAcls(deleted)
+
+    // Delete all literal ACLs:
+    deleted = client.deleteAcls(List(allLiteralTopicAcls).asJava).all().get().asScala.toSet
+    assertEquals(Set(anyAcl, acl2, fooAcl), deleted)
+    assertEquals(Set(prefixAcl), getAcls(allTopicAcls))
+
+    ensureAcls(deleted)
+
+    // Delete all prefixed ACLs:
+    deleted = client.deleteAcls(List(allPrefixedTopicAcls).asJava).all().get().asScala.toSet
+    assertEquals(Set(prefixAcl), deleted)
+    assertEquals(Set(anyAcl, acl2, fooAcl), getAcls(allTopicAcls))
+
+    ensureAcls(deleted)
+
+    // Delete all topic ACLs:
+    deleted = client.deleteAcls(List(allTopicAcls).asJava).all().get().asScala.toSet
+    assertEquals(Set(), getAcls(allTopicAcls))
+  }
+
+  //noinspection ScalaDeprecation - test explicitly covers clients using legacy / deprecated constructors
+  @Test
+  def testLegacyAclOpsNeverAffectOrReturnPrefixed(): Unit = {
+    client = Admin.create(createConfig)
+    ensureAcls(Set(anyAcl, acl2, fooAcl, prefixAcl))  // <-- prefixed exists, but should never be returned.
+
+    val allTopicAcls = new AclBindingFilter(new ResourcePatternFilter(ResourceType.TOPIC, null, PatternType.MATCH), AccessControlEntryFilter.ANY)
+    val legacyAllTopicAcls = new AclBindingFilter(new ResourcePatternFilter(ResourceType.TOPIC, null, PatternType.LITERAL), AccessControlEntryFilter.ANY)
+    val legacyMyTopic2Acls = new AclBindingFilter(new ResourcePatternFilter(ResourceType.TOPIC, "mytopic2", PatternType.LITERAL), AccessControlEntryFilter.ANY)
+    val legacyAnyTopicAcls = new AclBindingFilter(new ResourcePatternFilter(ResourceType.TOPIC, "*", PatternType.LITERAL), AccessControlEntryFilter.ANY)
+    val legacyFooTopicAcls = new AclBindingFilter(new ResourcePatternFilter(ResourceType.TOPIC, "foobar", PatternType.LITERAL), AccessControlEntryFilter.ANY)
+
+    assertEquals(Set(anyAcl, acl2, fooAcl), getAcls(legacyAllTopicAcls))
+    assertEquals(Set(acl2), getAcls(legacyMyTopic2Acls))
+    assertEquals(Set(anyAcl), getAcls(legacyAnyTopicAcls))
+    assertEquals(Set(fooAcl), getAcls(legacyFooTopicAcls))
+
+    // Delete only (legacy) ACLs on 'mytopic2' topic
+    var deleted = client.deleteAcls(List(legacyMyTopic2Acls).asJava).all().get().asScala.toSet
+    assertEquals(Set(acl2), deleted)
+    assertEquals(Set(anyAcl, fooAcl, prefixAcl), getAcls(allTopicAcls))
+
+    ensureAcls(deleted)
+
+    // Delete only (legacy) ACLs on '*' topic
+    deleted = client.deleteAcls(List(legacyAnyTopicAcls).asJava).all().get().asScala.toSet
+    assertEquals(Set(anyAcl), deleted)
+    assertEquals(Set(acl2, fooAcl, prefixAcl), getAcls(allTopicAcls))
+
+    ensureAcls(deleted)
+
+    // Delete all (legacy) topic ACLs:
+    deleted = client.deleteAcls(List(legacyAllTopicAcls).asJava).all().get().asScala.toSet
+    assertEquals(Set(anyAcl, acl2, fooAcl), deleted)
+    assertEquals(Set(), getAcls(legacyAllTopicAcls))
+    assertEquals(Set(prefixAcl), getAcls(allTopicAcls))
+  }
+
+  @Test
+  def testAttemptToCreateInvalidAcls(): Unit = {
+    client = Admin.create(createConfig)
+    val clusterAcl = new AclBinding(new ResourcePattern(ResourceType.CLUSTER, "foobar", PatternType.LITERAL),
+      new AccessControlEntry("User:ANONYMOUS", "*", AclOperation.READ, AclPermissionType.ALLOW))
+    val emptyResourceNameAcl = new AclBinding(new ResourcePattern(ResourceType.TOPIC, "", PatternType.LITERAL),
+      new AccessControlEntry("User:ANONYMOUS", "*", AclOperation.READ, AclPermissionType.ALLOW))
+    val results = client.createAcls(List(clusterAcl, emptyResourceNameAcl).asJava, new CreateAclsOptions())
+    assertEquals(Set(clusterAcl, emptyResourceNameAcl), results.values.keySet().asScala)
+    assertFutureExceptionTypeEquals(results.values.get(clusterAcl), classOf[InvalidRequestException])
+    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
+    }
+  }
+
+  private def testAclCreateGetDelete(expectAuth: Boolean): Unit = {
+    TestUtils.waitUntilTrue(() => {
+      val result = client.createAcls(List(fooAcl, transactionalIdAcl).asJava, new CreateAclsOptions)
+      if (expectAuth) {
+        Try(result.all.get) match {
+          case Failure(e) =>
+            verifyCauseIsClusterAuth(e)
+            false
+          case Success(_) => true
+        }
+      } else {
+        Try(result.all.get) match {
+          case Failure(e) =>
+            verifyCauseIsClusterAuth(e)
+            true
+          case Success(_) => false
+        }
+      }
+    }, "timed out waiting for createAcls to " + (if (expectAuth) "succeed" else "fail"))
+    if (expectAuth) {
+      waitForDescribeAcls(client, fooAcl.toFilter, Set(fooAcl))
+      waitForDescribeAcls(client, transactionalIdAcl.toFilter, Set(transactionalIdAcl))
+    }
+    TestUtils.waitUntilTrue(() => {
+      val result = client.deleteAcls(List(fooAcl.toFilter, transactionalIdAcl.toFilter).asJava, new DeleteAclsOptions)
+      if (expectAuth) {
+        Try(result.all.get) match {
+          case Failure(e) =>
+            verifyCauseIsClusterAuth(e)
+            false
+          case Success(_) => true
+        }
+      } else {
+        Try(result.all.get) match {
+          case Failure(e) =>
+            verifyCauseIsClusterAuth(e)
+            true
+          case Success(_) =>
+            assertEquals(Set(fooAcl, transactionalIdAcl), result.values.keySet)
+            assertEquals(Set(fooAcl), result.values.get(fooAcl.toFilter).get.values.asScala.map(_.binding).toSet)
+            assertEquals(Set(transactionalIdAcl),
+              result.values.get(transactionalIdAcl.toFilter).get.values.asScala.map(_.binding).toSet)
+            true
+        }
+      }
+    }, "timed out waiting for deleteAcls to " + (if (expectAuth) "succeed" else "fail"))
+    if (expectAuth) {
+      waitForDescribeAcls(client, fooAcl.toFilter, Set.empty)
+      waitForDescribeAcls(client, transactionalIdAcl.toFilter, Set.empty)
+    }
+  }
+
+  private def testAclGet(expectAuth: Boolean): Unit = {
+    TestUtils.waitUntilTrue(() => {
+      val userAcl = new AclBinding(new ResourcePattern(ResourceType.TOPIC, "*", PatternType.LITERAL),
+        new AccessControlEntry("User:*", "*", AclOperation.ALL, AclPermissionType.ALLOW))
+      val results = client.describeAcls(userAcl.toFilter)
+      if (expectAuth) {
+        Try(results.values.get) match {
+          case Failure(e) =>
+            verifyCauseIsClusterAuth(e)
+            false
+          case Success(acls) => Set(userAcl).equals(acls.asScala.toSet)
+        }
+      } else {
+        Try(results.values.get) match {
+          case Failure(e) =>
+            verifyCauseIsClusterAuth(e)
+            true
+          case Success(_) => false
+        }
+      }
+    }, "timed out waiting for describeAcls to " + (if (expectAuth) "succeed" else "fail"))
+  }
+
+  @Test
+  def testAclAuthorizationDenied(): Unit = {
+    client = Admin.create(createConfig)
+
+    // Test that we cannot create or delete ACLs when ALTER is denied.
+    authorizationAdmin.addClusterAcl(DENY, ALTER)
+    testAclGet(expectAuth = true)
+    testAclCreateGetDelete(expectAuth = false)
+
+    // Test that we cannot do anything with ACLs when DESCRIBE and ALTER are denied.
+    authorizationAdmin.addClusterAcl(DENY, DESCRIBE)
+    testAclGet(expectAuth = false)
+    testAclCreateGetDelete(expectAuth = false)
+
+    // Test that we can create, delete, and get ACLs with the default ACLs.
+    authorizationAdmin.removeClusterAcl(DENY, DESCRIBE)
+    authorizationAdmin.removeClusterAcl(DENY, ALTER)
+    testAclGet(expectAuth = true)
+    testAclCreateGetDelete(expectAuth = true)
+
+    // Test that we can't do anything with ACLs without the ALLOW ALTER ACL in place.
+    authorizationAdmin.removeClusterAcl(ALLOW, ALTER)
+    authorizationAdmin.removeClusterAcl(ALLOW, DELETE)
+    testAclGet(expectAuth = false)
+    testAclCreateGetDelete(expectAuth = false)
+
+    // Test that we can describe, but not alter ACLs, with only the ALLOW DESCRIBE ACL in place.
+    authorizationAdmin.addClusterAcl(ALLOW, DESCRIBE)
+    testAclGet(expectAuth = true)
+    testAclCreateGetDelete(expectAuth = false)
+  }
+
+  @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 = Admin.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()
+      acls == lastResults.asScala.toSet
+    }, s"timed out waiting for ACLs $acls.\nActual $lastResults")
+  }
+
+  private def ensureAcls(bindings: Set[AclBinding]): Unit = {
+    client.createAcls(bindings.asJava).all().get()
+
+    bindings.foreach(binding => waitForDescribeAcls(client, binding.toFilter, Set(binding)))
+  }
+
+  private def getAcls(allTopicAcls: AclBindingFilter) = {
+    client.describeAcls(allTopicAcls).values.get().asScala.toSet
+  }
+
+  @deprecated("Use kafka.security.authorizer.AclAuthorizer", "Since 2.5")
+  class LegacyAuthorizationAdmin extends AuthorizationAdmin {
+    import kafka.security.auth._
+    import kafka.security.authorizer.AuthorizerWrapper
+
+    override def authorizerClassName: String = classOf[SimpleAclAuthorizer].getName
+
+    override def initializeAcls(): Unit = {
+      val authorizer = CoreUtils.createObject[Authorizer](classOf[SimpleAclAuthorizer].getName)
+      try {
+        authorizer.configure(configs.head.originals())
+        authorizer.addAcls(Set(new Acl(Acl.WildCardPrincipal, Allow,
+          Acl.WildCardHost, All)), new Resource(Topic, "*", PatternType.LITERAL))
+        authorizer.addAcls(Set(new Acl(Acl.WildCardPrincipal, Allow,
+          Acl.WildCardHost, All)), new Resource(Group, "*", PatternType.LITERAL))
+
+        authorizer.addAcls(Set(clusterAcl(ALLOW, CREATE),
+          clusterAcl(ALLOW, DELETE),
+          clusterAcl(ALLOW, CLUSTER_ACTION),
+          clusterAcl(ALLOW, ALTER_CONFIGS),
+          clusterAcl(ALLOW, ALTER)),
+          Resource.ClusterResource)
+      } finally {
+        authorizer.close()
+      }
+    }
+
+    override def addClusterAcl(permissionType: AclPermissionType, operation: AclOperation): Unit = {
+      val acls = Set(clusterAcl(permissionType, operation))
+      val authorizer = simpleAclAuthorizer
+      val prevAcls = authorizer.getAcls(Resource.ClusterResource)
+      authorizer.addAcls(acls, Resource.ClusterResource)
+      TestUtils.waitAndVerifyAcls(prevAcls ++ acls, authorizer, Resource.ClusterResource)
+    }
+
+    override def removeClusterAcl(permissionType: AclPermissionType, operation: AclOperation): Unit = {
+      val acls = Set(clusterAcl(permissionType, operation))
+      val authorizer = simpleAclAuthorizer
+      val prevAcls = authorizer.getAcls(Resource.ClusterResource)
+      Assert.assertTrue(authorizer.removeAcls(acls, Resource.ClusterResource))
+      TestUtils.waitAndVerifyAcls(prevAcls -- acls, authorizer, Resource.ClusterResource)
+    }
+
+
+    private def clusterAcl(permissionType: AclPermissionType, operation: AclOperation): Acl = {
+      new Acl(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "*"), PermissionType.fromJava(permissionType),
+        Acl.WildCardHost, Operation.fromJava(operation))
+    }
+
+    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/SslAdminIntegrationTest.scala b/core/src/test/scala/integration/kafka/api/SslAdminIntegrationTest.scala
new file mode 100644
index 0000000000000..4ab20fd411158
--- /dev/null
+++ b/core/src/test/scala/integration/kafka/api/SslAdminIntegrationTest.scala
@@ -0,0 +1,319 @@
+/**
+  * Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE
+  * file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file
+  * to You under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the
+  * License. You may obtain a copy of the License at
+  *
+  * http://www.apache.org/licenses/LICENSE-2.0
+  *
+  * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+  * an "AS IS" BASIS, WITHOUT WARRANTIES 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 com.yammer.metrics.core.Gauge
+import kafka.metrics.KafkaYammerMetrics
+import kafka.security.authorizer.AclAuthorizer
+import kafka.security.authorizer.AclEntry.{WildcardHost, WildcardPrincipalString}
+import kafka.server.KafkaConfig
+import kafka.utils.{CoreUtils, TestUtils}
+import org.apache.kafka.clients.admin.{Admin, 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.{PatternType, Resource, ResourcePattern, ResourceType}
+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.jdk.CollectionConverters._
+import scala.collection.mutable
+
+object SslAdminIntegrationTest {
+  @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 { (result, exception) =>
+                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 SslAdminIntegrationTest extends SaslSslAdminIntegrationTest {
+  override val authorizationAdmin = new AclAuthorizationAdmin
+  val clusterResourcePattern = new ResourcePattern(ResourceType.CLUSTER, Resource.CLUSTER_NAME, PatternType.LITERAL)
+
+  this.serverConfig.setProperty(KafkaConfig.ZkEnableSecureAclsProp, "true")
+
+  override protected def securityProtocol = SecurityProtocol.SSL
+  override protected lazy val trustStoreFile = Some(File.createTempFile("truststore", ".jks"))
+  private val adminClients = mutable.Buffer.empty[Admin]
+
+  override def setUpSasl(): Unit = {
+    SslAdminIntegrationTest.semaphore = None
+    SslAdminIntegrationTest.executor = None
+    SslAdminIntegrationTest.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 = SslAdminIntegrationTest.semaphore
+    SslAdminIntegrationTest.semaphore = None
+    semaphore.foreach(s => s.release(s.getQueueLength))
+
+    adminClients.foreach(_.close())
+    super.tearDown()
+  }
+
+  @Test
+  def testAclUpdatesUsingSynchronousAuthorizer(): Unit = {
+    verifyAclUpdates()
+  }
+
+  @Test
+  def testAclUpdatesUsingAsynchronousAuthorizer(): Unit = {
+    SslAdminIntegrationTest.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)
+    SslAdminIntegrationTest.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 = {
+    SslAdminIntegrationTest.executor = Some(Executors.newSingleThreadExecutor)
+    val testSemaphore = new Semaphore(0)
+    SslAdminIntegrationTest.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 = {
+    val acl = new AclBinding(new ResourcePattern(ResourceType.TOPIC, "mytopic3", PatternType.LITERAL),
+      new AccessControlEntry("User:ANONYMOUS", "*", AclOperation.DESCRIBE, AclPermissionType.ALLOW))
+
+    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)
+    SslAdminIntegrationTest.semaphore = Some(testSemaphore)
+
+    client = Admin.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.forEach(_.get)
+    validateRequestContext(SslAdminIntegrationTest.lastUpdateRequestContext.get, ApiKeys.CREATE_ACLS)
+
+    testSemaphore.acquire()
+    val results2 = client.deleteAcls(List(acl.toFilter, acl2.toFilter, acl3.toFilter).asJava).values
+    assertEquals(Set(acl.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.forEach(_.get)
+    assertEquals(0, results2.get(acl.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(SslAdminIntegrationTest.lastUpdateRequestContext.get, ApiKeys.DELETE_ACLS)
+  }
+
+  private def createAdminClient: Admin = {
+    val config = createConfig
+    config.put(AdminClientConfig.DEFAULT_API_TIMEOUT_MS_CONFIG, "40000")
+    val client = Admin.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 = KafkaYammerMetrics.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
+  }
+
+  class AclAuthorizationAdmin extends AuthorizationAdmin {
+
+    override def authorizerClassName: String = classOf[SslAdminIntegrationTest.TestableAclAuthorizer].getName
+
+    override def initializeAcls(): Unit = {
+      val authorizer = CoreUtils.createObject[Authorizer](classOf[AclAuthorizer].getName)
+      try {
+        authorizer.configure(configs.head.originals())
+        val ace = new AccessControlEntry(WildcardPrincipalString, 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 addClusterAcl(permissionType: AclPermissionType, operation: AclOperation): Unit = {
+      val ace = clusterAcl(permissionType, operation)
+      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: AclPermissionType, operation: AclOperation): Unit = {
+      val ace = clusterAcl(permissionType, operation)
+      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)
+    }
+  }
+}
diff --git a/core/src/test/scala/integration/kafka/api/SslEndToEndAuthorizationTest.scala b/core/src/test/scala/integration/kafka/api/SslEndToEndAuthorizationTest.scala
index b304f29497187..b705050c1b111 100644
--- a/core/src/test/scala/integration/kafka/api/SslEndToEndAuthorizationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/SslEndToEndAuthorizationTest.scala
@@ -17,41 +17,67 @@
 
 package kafka.api
 
+import java.util.Properties
+
+import kafka.utils.TestUtils
+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.security.auth._
+import org.apache.kafka.common.utils.Java
 import org.junit.Before
 
 object SslEndToEndAuthorizationTest {
   class TestPrincipalBuilder extends KafkaPrincipalBuilder {
-    private val Pattern = "O=A (.*?),CN=localhost".r
+    private val Pattern = "O=A (.*?),CN=(.*?)".r
 
+    // Use full DN as client principal to test special characters in principal
+    // Use field from DN as server principal to test custom PrincipalBuilder
     override def build(context: AuthenticationContext): KafkaPrincipal = {
-      context match {
-        case ctx: SslAuthenticationContext =>
-          ctx.session.getPeerPrincipal.getName match {
-            case Pattern(name) =>
-              new KafkaPrincipal(KafkaPrincipal.USER_TYPE, name)
-            case _ =>
-              KafkaPrincipal.ANONYMOUS
-          }
+      val peerPrincipal = context.asInstanceOf[SslAuthenticationContext].session.getPeerPrincipal.getName
+      peerPrincipal match {
+        case Pattern(name, _) =>
+          val principal = if (name == "server") name else peerPrincipal
+          new KafkaPrincipal(KafkaPrincipal.USER_TYPE, principal)
+        case _ =>
+          KafkaPrincipal.ANONYMOUS
       }
     }
   }
 }
 
 class SslEndToEndAuthorizationTest extends EndToEndAuthorizationTest {
+
   import kafka.api.SslEndToEndAuthorizationTest.TestPrincipalBuilder
 
   override protected def securityProtocol = SecurityProtocol.SSL
+  // Since there are other E2E tests that enable SSL, running this test with TLSv1.3 if supported
+  private  val tlsProtocol = if (Java.IS_JAVA11_COMPATIBLE) "TLSv1.3" else "TLSv1.2"
+
   this.serverConfig.setProperty(BrokerSecurityConfigs.SSL_CLIENT_AUTH_CONFIG, "required")
   this.serverConfig.setProperty(BrokerSecurityConfigs.PRINCIPAL_BUILDER_CLASS_CONFIG, classOf[TestPrincipalBuilder].getName)
-  override val clientPrincipal = "client"
-  override val kafkaPrincipal = "server"
-
+  this.serverConfig.setProperty(SslConfigs.SSL_PROTOCOL_CONFIG, tlsProtocol)
+  this.serverConfig.setProperty(SslConfigs.SSL_ENABLED_PROTOCOLS_CONFIG, tlsProtocol)
+  // Escaped characters in DN attribute values: from http://www.ietf.org/rfc/rfc2253.txt
+  // - a space or "#" character occurring at the beginning of the string
+  // - a space character occurring at the end of the string
+  // - one of the characters ",", "+", """, "\", "<", ">" or ";"
+  //
+  // Leading and trailing spaces in Kafka principal dont work with ACLs, but we can workaround by using
+  // a PrincipalBuilder that removes/replaces them.
+  private val clientCn = """\#A client with special chars in CN : (\, \+ \" \\ \< \> \; ')"""
+  override val clientPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, s"O=A client,CN=$clientCn")
+  override val kafkaPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "server")
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     startSasl(jaasSections(List.empty, None, ZkSasl))
     super.setUp()
   }
 
+  override def clientSecurityProps(certAlias: String): Properties = {
+    val props = TestUtils.securityConfigs(Mode.CLIENT, securityProtocol, trustStoreFile,
+      certAlias, clientCn, clientSaslProperties, tlsProtocol)
+    props.remove(SslConfigs.SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG)
+    props
+  }
 }
diff --git a/core/src/test/scala/integration/kafka/api/TransactionsBounceTest.scala b/core/src/test/scala/integration/kafka/api/TransactionsBounceTest.scala
index 1712ef872847a..6203fa2f51eb8 100644
--- a/core/src/test/scala/integration/kafka/api/TransactionsBounceTest.scala
+++ b/core/src/test/scala/integration/kafka/api/TransactionsBounceTest.scala
@@ -19,26 +19,22 @@ package kafka.api
 
 import java.util.Properties
 
-import kafka.integration.KafkaServerTestHarness
 import kafka.server.KafkaConfig
 import kafka.utils.{ShutdownableThread, TestUtils}
-import org.apache.kafka.clients.consumer.ConsumerConfig
+import org.apache.kafka.clients.consumer.{ConsumerConfig, KafkaConsumer}
+import org.apache.kafka.clients.producer.{KafkaProducer, ProducerConfig}
 import org.apache.kafka.clients.producer.internals.ErrorLoggingCallback
 import org.apache.kafka.common.TopicPartition
-import org.apache.kafka.common.security.auth.SecurityProtocol
 import org.junit.Assert._
 import org.junit.Test
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 import scala.collection.mutable
 
-
-class TransactionsBounceTest extends KafkaServerTestHarness {
+class TransactionsBounceTest extends IntegrationTestHarness {
   private val producerBufferSize =  65536
   private val serverMessageMaxBytes =  producerBufferSize/2
   private val numPartitions = 3
-
-  val numServers = 4
   private val outputTopic = "output-topic"
   private val inputTopic = "input-topic"
 
@@ -58,7 +54,6 @@ class TransactionsBounceTest extends KafkaServerTestHarness {
   overridingProps.put(KafkaConfig.GroupMinSessionTimeoutMsProp, "10") // set small enough session timeout
   overridingProps.put(KafkaConfig.GroupInitialRebalanceDelayMsProp, "0")
 
-
   // This is the one of the few tests we currently allow to preallocate ports, despite the fact that this can result in transient
   // failures due to ports getting reused. We can't use random ports because of bad behavior that can result from bouncing
   // brokers too quickly when they get new, random ports. If we're not careful, the client can end up in a situation
@@ -69,12 +64,26 @@ class TransactionsBounceTest extends KafkaServerTestHarness {
   // Since such quick rotation of servers is incredibly unrealistic, we allow this one test to preallocate ports, leaving
   // a small risk of hitting errors due to port conflicts. Hopefully this is infrequent enough to not cause problems.
   override def generateConfigs = {
-    FixedPortTestUtils.createBrokerConfigs(numServers, zkConnect,enableControlledShutdown = true)
+    FixedPortTestUtils.createBrokerConfigs(brokerCount, zkConnect, enableControlledShutdown = true)
       .map(KafkaConfig.fromProps(_, overridingProps))
   }
 
+  override protected def brokerCount: Int = 4
+
+  @Test
+  def testWithGroupId(): Unit = {
+    testBrokerFailure((producer, groupId, consumer) =>
+      producer.sendOffsetsToTransaction(TestUtils.consumerPositions(consumer).asJava, groupId))
+  }
+
   @Test
-  def testBrokerFailure() {
+  def testWithGroupMetadata(): Unit = {
+    testBrokerFailure((producer, _, consumer) =>
+      producer.sendOffsetsToTransaction(TestUtils.consumerPositions(consumer).asJava, consumer.groupMetadata()))
+  }
+
+  private def testBrokerFailure(commit: (KafkaProducer[Array[Byte], Array[Byte]],
+    String, KafkaConsumer[Array[Byte], Array[Byte]]) => Unit): 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"
@@ -82,17 +91,18 @@ class TransactionsBounceTest extends KafkaServerTestHarness {
     createTopics()
 
     TestUtils.seedTopicWithNumberedRecords(inputTopic, numInputRecords, servers)
-    val consumer = createConsumerAndSubscribeToTopics(consumerGroup, List(inputTopic))
-    val producer = TestUtils.createTransactionalProducer("test-txn", servers, 512)
+    val consumer = createConsumerAndSubscribe(consumerGroup, List(inputTopic))
+    val producer = createTransactionalProducer("test-txn")
 
     producer.initTransactions()
 
     val scheduler = new BounceScheduler
     scheduler.start()
 
-    var numMessagesProcessed = 0
-    var iteration = 0
     try {
+      var numMessagesProcessed = 0
+      var iteration = 0
+
       while (numMessagesProcessed < numInputRecords) {
         val toRead = Math.min(200, numInputRecords - numMessagesProcessed)
         trace(s"$iteration: About to read $toRead messages, processed $numMessagesProcessed so far..")
@@ -102,11 +112,10 @@ class TransactionsBounceTest extends KafkaServerTestHarness {
         producer.beginTransaction()
         val shouldAbort = iteration % 3 == 0
         records.foreach { record =>
-          producer.send(TestUtils.producerRecordWithExpectedTransactionStatus(outputTopic, record.key, record.value,
-            !shouldAbort), new ErrorLoggingCallback(outputTopic, record.key, record.value, true))
+          producer.send(TestUtils.producerRecordWithExpectedTransactionStatus(outputTopic, null, record.key, record.value, !shouldAbort), new ErrorLoggingCallback(outputTopic, record.key, record.value, true))
         }
         trace(s"Sent ${records.size} messages. Committing offsets.")
-        producer.sendOffsetsToTransaction(TestUtils.consumerPositions(consumer).asJava, consumerGroup)
+        commit(producer, consumerGroup, consumer)
 
         if (shouldAbort) {
           trace(s"Committed offsets. Aborting transaction of ${records.size} messages.")
@@ -120,13 +129,10 @@ class TransactionsBounceTest extends KafkaServerTestHarness {
         iteration += 1
       }
     } finally {
-      producer.close()
-      consumer.close()
+      scheduler.shutdown()
     }
 
-    scheduler.shutdown()
-
-    val verifyingConsumer = createConsumerAndSubscribeToTopics("randomGroup", List(outputTopic), readCommitted = true)
+    val verifyingConsumer = createConsumerAndSubscribe("randomGroup", List(outputTopic), readCommitted = true)
     val recordsByPartition = new mutable.HashMap[TopicPartition, mutable.ListBuffer[Int]]()
     TestUtils.pollUntilAtLeastNumRecords(verifyingConsumer, numInputRecords).foreach { record =>
       val value = TestUtils.assertCommittedAndGetValue(record).toInt
@@ -136,7 +142,7 @@ class TransactionsBounceTest extends KafkaServerTestHarness {
     }
 
     val outputRecords = new mutable.ListBuffer[Int]()
-    recordsByPartition.values.foreach { case (partitionValues) =>
+    recordsByPartition.values.foreach { partitionValues =>
       assertEquals("Out of order messages detected", partitionValues, partitionValues.sorted)
       outputRecords.appendAll(partitionValues)
     }
@@ -146,22 +152,26 @@ class TransactionsBounceTest extends KafkaServerTestHarness {
 
     val expectedValues = (0 until numInputRecords).toSet
     assertEquals(s"Missing messages: ${expectedValues -- recordSet}", expectedValues, recordSet)
-
-    verifyingConsumer.close()
   }
 
-  private def createConsumerAndSubscribeToTopics(groupId: String, topics: List[String], readCommitted: Boolean = false) = {
+  private def createTransactionalProducer(transactionalId: String) = {
     val props = new Properties()
-    if (readCommitted)
-      props.put(ConsumerConfig.ISOLATION_LEVEL_CONFIG, "read_committed")
-    props.put(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, "2000")
-    props.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false")
-    props.put(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG, "10000")
-    props.put(ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG, "3000")
-    props.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest")
-
-    val consumer = TestUtils.createNewConsumer(TestUtils.getBrokerListStrFromServers(servers), groupId = groupId,
-      securityProtocol = SecurityProtocol.PLAINTEXT, props = Some(props))
+    props.put(ProducerConfig.ACKS_CONFIG, "all")
+    props.put(ProducerConfig.BATCH_SIZE_CONFIG, "512")
+    props.put(ProducerConfig.TRANSACTIONAL_ID_CONFIG, transactionalId)
+    props.put(ProducerConfig.ENABLE_IDEMPOTENCE_CONFIG, "true")
+    createProducer(configOverrides = props)
+  }
+
+  private def createConsumerAndSubscribe(groupId: String,
+                                         topics: List[String],
+                                         readCommitted: Boolean = false) = {
+    val consumerProps = new Properties
+    consumerProps.put(ConsumerConfig.GROUP_ID_CONFIG, groupId)
+    consumerProps.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false")
+    consumerProps.put(ConsumerConfig.ISOLATION_LEVEL_CONFIG,
+      if (readCommitted) "read_committed" else "read_uncommitted")
+    val consumer = createConsumer(configOverrides = consumerProps)
     consumer.subscribe(topics.asJava)
     consumer
   }
@@ -169,8 +179,8 @@ class TransactionsBounceTest extends KafkaServerTestHarness {
   private def createTopics() =  {
     val topicConfig = new Properties()
     topicConfig.put(KafkaConfig.MinInSyncReplicasProp, 2.toString)
-    TestUtils.createTopic(zkUtils, inputTopic, numPartitions, 3, servers, topicConfig)
-    TestUtils.createTopic(zkUtils, outputTopic, numPartitions, 3, servers, topicConfig)
+    createTopic(inputTopic, numPartitions, 3, topicConfig)
+    createTopic(outputTopic, numPartitions, 3, topicConfig)
   }
 
   private class BounceScheduler extends ShutdownableThread("daemon-broker-bouncer", false) {
@@ -186,11 +196,12 @@ class TransactionsBounceTest extends KafkaServerTestHarness {
         Thread.sleep(500)
       }
 
-      (0 until numPartitions).foreach(partition => TestUtils.waitUntilLeaderIsElectedOrChanged(zkUtils, outputTopic, partition))
+      (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/TransactionsExpirationTest.scala b/core/src/test/scala/integration/kafka/api/TransactionsExpirationTest.scala
new file mode 100644
index 0000000000000..9251750081b50
--- /dev/null
+++ b/core/src/test/scala/integration/kafka/api/TransactionsExpirationTest.scala
@@ -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 kafka.api
+
+import java.util.Properties
+
+import kafka.integration.KafkaServerTestHarness
+import kafka.server.KafkaConfig
+import kafka.utils.TestUtils
+import kafka.utils.TestUtils.consumeRecords
+import org.apache.kafka.clients.consumer.KafkaConsumer
+import org.apache.kafka.clients.producer.KafkaProducer
+import org.apache.kafka.common.errors.InvalidPidMappingException
+import org.junit.{After, Before, Test}
+
+import scala.jdk.CollectionConverters._
+import scala.collection.Seq
+
+// Test class that uses a very small transaction timeout to trigger InvalidPidMapping errors
+class TransactionsExpirationTest extends KafkaServerTestHarness {
+  val topic1 = "topic1"
+  val topic2 = "topic2"
+  val numPartitions = 4
+  val replicationFactor = 3
+
+  var producer: KafkaProducer[Array[Byte], Array[Byte]] = _
+  var consumer: KafkaConsumer[Array[Byte], Array[Byte]] = _
+
+  override def generateConfigs: Seq[KafkaConfig] = {
+    TestUtils.createBrokerConfigs(3, zkConnect).map(KafkaConfig.fromProps(_, serverProps()))
+  }
+
+  @Before
+  override def setUp(): Unit = {
+    super.setUp()
+
+    producer = TestUtils.createTransactionalProducer("transactionalProducer", servers)
+    consumer = TestUtils.createConsumer(TestUtils.getBrokerListStrFromServers(servers),
+      enableAutoCommit = false,
+      readCommitted = true)
+
+    TestUtils.createTopic(zkClient, topic1, numPartitions, 3, servers, new Properties())
+    TestUtils.createTopic(zkClient, topic2, numPartitions, 3, servers, new Properties())
+  }
+
+  @After
+  override def tearDown(): Unit = {
+    producer.close()
+    consumer.close()
+
+    super.tearDown()
+  }
+
+  @Test
+  def testBumpTransactionalEpochAfterInvalidProducerIdMapping(): Unit = {
+    producer.initTransactions()
+
+    // Start and then abort a transaction to allow the transactional ID to expire
+    producer.beginTransaction()
+    producer.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic1, 0, "2", "2", willBeCommitted = false))
+    producer.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic2, 0, "4", "4", willBeCommitted = false))
+    producer.abortTransaction()
+
+    // Wait for the transactional ID to expire
+    Thread.sleep(3000)
+
+    // Start a new transaction and attempt to send, which will trigger an AddPartitionsToTxnRequest, which will fail due to the expired producer ID
+    producer.beginTransaction()
+    val failedFuture = producer.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic1, 3, "1", "1", willBeCommitted = false))
+    Thread.sleep(500)
+
+    org.apache.kafka.test.TestUtils.assertFutureThrows(failedFuture, classOf[InvalidPidMappingException])
+    producer.abortTransaction()
+
+    producer.beginTransaction()
+    producer.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic2, null, "2", "2", willBeCommitted = true))
+    producer.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic1, 2, "4", "4", willBeCommitted = true))
+    producer.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic2, null, "1", "1", willBeCommitted = true))
+    producer.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic1, 3, "3", "3", willBeCommitted = true))
+    producer.commitTransaction()
+
+    consumer.subscribe(List(topic1, topic2).asJava)
+
+    val records = consumeRecords(consumer, 4)
+    records.foreach { record =>
+      TestUtils.assertCommittedAndGetValue(record)
+    }
+  }
+  private def serverProps() = {
+    val serverProps = new Properties()
+    serverProps.put(KafkaConfig.AutoCreateTopicsEnableProp, false.toString)
+    // Set a smaller value for the number of partitions for the __consumer_offsets topic
+    // so that the creation of that topic/partition(s) and subsequent leader assignment doesn't take relatively long
+    serverProps.put(KafkaConfig.OffsetsTopicPartitionsProp, 1.toString)
+    serverProps.put(KafkaConfig.TransactionsTopicPartitionsProp, 3.toString)
+    serverProps.put(KafkaConfig.TransactionsTopicReplicationFactorProp, 2.toString)
+    serverProps.put(KafkaConfig.TransactionsTopicMinISRProp, 2.toString)
+    serverProps.put(KafkaConfig.ControlledShutdownEnableProp, true.toString)
+    serverProps.put(KafkaConfig.UncleanLeaderElectionEnableProp, false.toString)
+    serverProps.put(KafkaConfig.AutoLeaderRebalanceEnableProp, false.toString)
+    serverProps.put(KafkaConfig.GroupInitialRebalanceDelayMsProp, "0")
+    serverProps.put(KafkaConfig.TransactionsAbortTimedOutTransactionCleanupIntervalMsProp, "200")
+    serverProps.put(KafkaConfig.TransactionalIdExpirationMsProp, "2000")
+    serverProps.put(KafkaConfig.TransactionsRemoveExpiredTransactionalIdCleanupIntervalMsProp, "500")
+    serverProps
+  }
+}
diff --git a/core/src/test/scala/integration/kafka/api/TransactionsTest.scala b/core/src/test/scala/integration/kafka/api/TransactionsTest.scala
index bb6b52024cbbf..073cbcf15f8d9 100644
--- a/core/src/test/scala/integration/kafka/api/TransactionsTest.scala
+++ b/core/src/test/scala/integration/kafka/api/TransactionsTest.scala
@@ -18,8 +18,10 @@
 package kafka.api
 
 import java.lang.{Long => JLong}
-import java.util.Properties
-import java.util.concurrent.{ExecutionException, TimeUnit}
+import java.nio.charset.StandardCharsets
+import java.time.Duration
+import java.util.concurrent.TimeUnit
+import java.util.{Optional, Properties}
 
 import kafka.integration.KafkaServerTestHarness
 import kafka.server.KafkaConfig
@@ -27,13 +29,14 @@ import kafka.utils.TestUtils
 import kafka.utils.TestUtils.consumeRecords
 import org.apache.kafka.clients.consumer.{ConsumerConfig, KafkaConsumer, OffsetAndMetadata}
 import org.apache.kafka.clients.producer.{KafkaProducer, ProducerRecord}
-import org.apache.kafka.common.TopicPartition
-import org.apache.kafka.common.errors.ProducerFencedException
-import org.apache.kafka.common.security.auth.SecurityProtocol
-import org.junit.{After, Before, Test}
+import org.apache.kafka.common.errors.{InvalidProducerEpochException, ProducerFencedException, TimeoutException}
+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.jdk.CollectionConverters._
+import scala.collection.Seq
 import scala.collection.mutable.Buffer
 import scala.concurrent.ExecutionException
 
@@ -45,6 +48,7 @@ class TransactionsTest extends KafkaServerTestHarness {
 
   val topic1 = "topic1"
   val topic2 = "topic2"
+  val numPartitions = 4
 
   val transactionalProducers = Buffer[KafkaProducer[Array[Byte], Array[Byte]]]()
   val transactionalConsumers = Buffer[KafkaConsumer[Array[Byte], Array[Byte]]]()
@@ -57,11 +61,10 @@ class TransactionsTest extends KafkaServerTestHarness {
   @Before
   override def setUp(): Unit = {
     super.setUp()
-    val numPartitions = 4
     val topicConfig = new Properties()
     topicConfig.put(KafkaConfig.MinInSyncReplicasProp, 2.toString)
-    TestUtils.createTopic(zkUtils, topic1, numPartitions, numServers, servers, topicConfig)
-    TestUtils.createTopic(zkUtils, topic2, numPartitions, numServers, servers, topicConfig)
+    createTopic(topic1, numPartitions, numServers, topicConfig)
+    createTopic(topic2, numPartitions, numServers, topicConfig)
 
     for (_ <- 0 until transactionalProducerCount)
       createTransactionalProducer("transactional-producer")
@@ -88,14 +91,14 @@ class TransactionsTest extends KafkaServerTestHarness {
     producer.initTransactions()
 
     producer.beginTransaction()
-    producer.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic2, "2", "2", willBeCommitted = false))
-    producer.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic1, "4", "4", willBeCommitted = false))
+    producer.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic2, null, "2", "2", willBeCommitted = false))
+    producer.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic1, null, "4", "4", willBeCommitted = false))
     producer.flush()
     producer.abortTransaction()
 
     producer.beginTransaction()
-    producer.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic1, "1", "1", willBeCommitted = true))
-    producer.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic2, "3", "3", willBeCommitted = true))
+    producer.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic1, null, "1", "1", willBeCommitted = true))
+    producer.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic2, null, "3", "3", willBeCommitted = true))
     producer.commitTransaction()
 
     consumer.subscribe(List(topic1, topic2).asJava)
@@ -167,7 +170,7 @@ class TransactionsTest extends KafkaServerTestHarness {
     // even if we seek to the end, we should not be able to see the undecided data
     assertEquals(2, readCommittedConsumer.assignment.size)
     readCommittedConsumer.seekToEnd(readCommittedConsumer.assignment)
-    readCommittedConsumer.assignment.asScala.foreach { tp =>
+    readCommittedConsumer.assignment.forEach { tp =>
       assertEquals(1L, readCommittedConsumer.position(tp))
     }
 
@@ -210,7 +213,7 @@ class TransactionsTest extends KafkaServerTestHarness {
     val readCommittedConsumer = createReadCommittedConsumer(props = consumerProps)
 
     readCommittedConsumer.assign(Set(new TopicPartition(topic1, 0)).asJava)
-    val records = consumeRecords(readCommittedConsumer, numMessages = 2)
+    val records = consumeRecords(readCommittedConsumer, numRecords = 2)
     assertEquals(2, records.size)
 
     val first = records.head
@@ -225,7 +228,20 @@ class TransactionsTest extends KafkaServerTestHarness {
   }
 
   @Test
-  def testSendOffsets() = {
+  def testSendOffsetsWithGroupId() = {
+    sendOffset((producer, groupId, consumer) =>
+      producer.sendOffsetsToTransaction(TestUtils.consumerPositions(consumer).asJava, groupId))
+  }
+
+  @Test
+  def testSendOffsetsWithGroupMetadata() = {
+    sendOffset((producer, _, consumer) =>
+      producer.sendOffsetsToTransaction(TestUtils.consumerPositions(consumer).asJava, consumer.groupMetadata()))
+  }
+
+  private def sendOffset(commit: (KafkaProducer[Array[Byte], Array[Byte]],
+    String, KafkaConsumer[Array[Byte], Array[Byte]]) => Unit) = {
+
     // The basic plan for the test is as follows:
     //  1. Seed topic1 with 1000 unique, numbered, messages.
     //  2. Run a consume/process/produce loop to transactionally copy messages from topic1 to topic2 and commit
@@ -240,7 +256,7 @@ class TransactionsTest extends KafkaServerTestHarness {
 
     TestUtils.seedTopicWithNumberedRecords(topic1, numSeedMessages, servers)
 
-    val producer = transactionalProducers(0)
+    val producer = transactionalProducers.head
 
     val consumer = createReadCommittedConsumer(consumerGroupId, maxPollRecords = numSeedMessages / 4)
     consumer.subscribe(List(topic1).asJava)
@@ -256,23 +272,23 @@ class TransactionsTest extends KafkaServerTestHarness {
         shouldCommit = !shouldCommit
 
         records.foreach { record =>
-          val key = new String(record.key(), "UTF-8")
-          val value = new String(record.value(), "UTF-8")
-          producer.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic2, key, value, willBeCommitted = shouldCommit))
+          val key = new String(record.key(), StandardCharsets.UTF_8)
+          val value = new String(record.value(), StandardCharsets.UTF_8)
+          producer.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic2, null, key, value, willBeCommitted = shouldCommit))
         }
 
-        producer.sendOffsetsToTransaction(TestUtils.consumerPositions(consumer).asJava, consumerGroupId)
+        commit(producer, consumerGroupId, consumer)
         if (shouldCommit) {
           producer.commitTransaction()
           recordsProcessed += records.size
-          debug(s"committed transaction.. Last committed record: ${new String(records.last.value(), "UTF-8")}. Num " +
+          debug(s"committed transaction.. Last committed record: ${new String(records.last.value(), StandardCharsets.UTF_8)}. Num " +
             s"records written to $topic2: $recordsProcessed")
         } else {
           producer.abortTransaction()
-          debug(s"aborted transaction Last committed record: ${new String(records.last.value(), "UTF-8")}. Num " +
+          debug(s"aborted transaction Last committed record: ${new String(records.last.value(), StandardCharsets.UTF_8)}. Num " +
             s"records written to $topic2: $recordsProcessed")
           TestUtils.resetToCommittedPositions(consumer)
-       }
+        }
       }
     } finally {
       consumer.close()
@@ -301,13 +317,13 @@ class TransactionsTest extends KafkaServerTestHarness {
     producer1.initTransactions()
 
     producer1.beginTransaction()
-    producer1.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic1, "1", "1", willBeCommitted = false))
-    producer1.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic2, "3", "3", willBeCommitted = false))
+    producer1.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic1, null, "1", "1", willBeCommitted = false))
+    producer1.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic2, null, "3", "3", willBeCommitted = false))
 
     producer2.initTransactions()  // ok, will abort the open transaction.
     producer2.beginTransaction()
-    producer2.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic1, "2", "4", willBeCommitted = true))
-    producer2.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic2, "2", "4", willBeCommitted = true))
+    producer2.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic1, null, "2", "4", willBeCommitted = true))
+    producer2.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic2, null, "2", "4", willBeCommitted = true))
 
     try {
       producer1.commitTransaction()
@@ -338,13 +354,13 @@ class TransactionsTest extends KafkaServerTestHarness {
     producer1.initTransactions()
 
     producer1.beginTransaction()
-    producer1.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic1, "1", "1", willBeCommitted = false))
-    producer1.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic2, "3", "3", willBeCommitted = false))
+    producer1.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic1, null, "1", "1", willBeCommitted = false))
+    producer1.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic2, null, "3", "3", willBeCommitted = false))
 
     producer2.initTransactions()  // ok, will abort the open transaction.
     producer2.beginTransaction()
-    producer2.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic1, "2", "4", willBeCommitted = true))
-    producer2.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic2, "2", "4", willBeCommitted = true))
+    producer2.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic1, null, "2", "4", willBeCommitted = true))
+    producer2.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic2, null, "2", "4", willBeCommitted = true))
 
     try {
       producer1.sendOffsetsToTransaction(Map(new TopicPartition("foobartopic", 0) -> new OffsetAndMetadata(110L)).asJava,
@@ -366,7 +382,74 @@ class TransactionsTest extends KafkaServerTestHarness {
   }
 
   @Test
-  def testFencingOnSend() {
+  def testOffsetMetadataInSendOffsetsToTransaction() = {
+    val tp = new TopicPartition(topic1, 0)
+    val groupId = "group"
+
+    val producer = transactionalProducers.head
+    val consumer = createReadCommittedConsumer(groupId)
+
+    consumer.subscribe(List(topic1).asJava)
+
+    producer.initTransactions()
+
+    producer.beginTransaction()
+    val offsetAndMetadata = new OffsetAndMetadata(110L, Optional.of(15), "some metadata")
+    producer.sendOffsetsToTransaction(Map(tp -> offsetAndMetadata).asJava, groupId)
+    producer.commitTransaction()  // ok
+
+    // The call to commit the transaction may return before all markers are visible, so we initialize a second
+    // producer to ensure the transaction completes and the committed offsets are visible.
+    val producer2 = transactionalProducers(1)
+    producer2.initTransactions()
+
+    TestUtils.waitUntilTrue(() => offsetAndMetadata.equals(consumer.committed(Set(tp).asJava).get(tp)), "cannot read committed offset")
+  }
+
+  @Test(expected = classOf[TimeoutException])
+  def testInitTransactionsTimeout(): Unit = {
+    testTimeout(false, producer => producer.initTransactions())
+  }
+
+  @Test(expected = classOf[TimeoutException])
+  def testSendOffsetsToTransactionTimeout(): Unit = {
+    testTimeout(true, producer => producer.sendOffsetsToTransaction(
+      Map(new TopicPartition(topic1, 0) -> new OffsetAndMetadata(0)).asJava, "test-group"))
+  }
+
+  @Test(expected = classOf[TimeoutException])
+  def testCommitTransactionTimeout(): Unit = {
+    testTimeout(true, producer => producer.commitTransaction())
+  }
+
+  @Test(expected = classOf[TimeoutException])
+  def testAbortTransactionTimeout(): Unit = {
+    testTimeout(true, producer => producer.abortTransaction())
+  }
+
+  def testTimeout(needInitAndSendMsg: Boolean,
+                  timeoutProcess: KafkaProducer[Array[Byte], Array[Byte]] => Unit): Unit = {
+    val producer = createTransactionalProducer("transactionProducer", maxBlockMs =  1000)
+
+    if (needInitAndSendMsg) {
+      producer.initTransactions()
+      producer.beginTransaction()
+      producer.send(new ProducerRecord[Array[Byte], Array[Byte]](topic1, "foo".getBytes, "bar".getBytes))
+    }
+
+    for  (i <- servers.indices)
+      killBroker(i)
+
+    try {
+      timeoutProcess(producer)
+      fail("Should raise a TimeoutException")
+    } finally {
+      producer.close(Duration.ZERO)
+    }
+  }
+
+  @Test
+  def testFencingOnSend(): Unit = {
     val producer1 = transactionalProducers(0)
     val producer2 = transactionalProducers(1)
     val consumer = transactionalConsumers(0)
@@ -376,16 +459,16 @@ class TransactionsTest extends KafkaServerTestHarness {
     producer1.initTransactions()
 
     producer1.beginTransaction()
-    producer1.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic1, "1", "1", willBeCommitted = false))
-    producer1.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic2, "3", "3", willBeCommitted = false))
+    producer1.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic1, null, "1", "1", willBeCommitted = false))
+    producer1.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic2, null, "3", "3", willBeCommitted = false))
 
     producer2.initTransactions()  // ok, will abort the open transaction.
     producer2.beginTransaction()
-    producer2.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic1, "2", "4", willBeCommitted = true)).get()
-    producer2.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic2, "2", "4", willBeCommitted = true)).get()
+    producer2.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic1, null, "2", "4", willBeCommitted = true)).get()
+    producer2.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic2, null, "2", "4", willBeCommitted = true)).get()
 
     try {
-      val result =  producer1.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic1, "1", "5", willBeCommitted = false))
+      val result = producer1.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic1, null, "1", "5", willBeCommitted = false))
       val recordMetadata = result.get()
       error(s"Missed a producer fenced exception when writing to ${recordMetadata.topic}-${recordMetadata.partition}. Grab the logs!!")
       servers.foreach { server =>
@@ -396,7 +479,7 @@ class TransactionsTest extends KafkaServerTestHarness {
       case _: ProducerFencedException =>
         producer1.close()
       case e: ExecutionException =>
-        assertTrue(e.getCause.isInstanceOf[ProducerFencedException])
+        assertTrue(e.getCause.isInstanceOf[InvalidProducerEpochException])
       case e: Exception =>
         fail("Got an unexpected exception from a fenced producer.", e)
     }
@@ -419,23 +502,23 @@ class TransactionsTest extends KafkaServerTestHarness {
 
     producer1.initTransactions()
     producer1.beginTransaction()
-    producer1.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic1, "1", "1", willBeCommitted = false))
-    producer1.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic2, "3", "3", willBeCommitted = false))
+    producer1.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic1, null, "1", "1", willBeCommitted = false))
+    producer1.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic2, null, "3", "3", willBeCommitted = false))
     producer1.abortTransaction()
 
     producer2.initTransactions()  // ok, will abort the open transaction.
     producer2.beginTransaction()
-    producer2.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic1, "2", "4", willBeCommitted = true))
+    producer2.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic1, null, "2", "4", willBeCommitted = true))
       .get(20, TimeUnit.SECONDS)
-    producer2.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic2, "2", "4", willBeCommitted = true))
+    producer2.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic2, null, "2", "4", willBeCommitted = true))
       .get(20, TimeUnit.SECONDS)
 
     try {
       producer1.beginTransaction()
-      val result =  producer1.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic1, "1", "5", willBeCommitted = false))
+      val result =  producer1.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic1, null, "1", "5", willBeCommitted = false))
       val recordMetadata = result.get()
       error(s"Missed a producer fenced exception when writing to ${recordMetadata.topic}-${recordMetadata.partition}. Grab the logs!!")
-      servers.foreach { case (server) =>
+      servers.foreach { server =>
         error(s"log dirs: ${server.logManager.liveLogDirs.map(_.getAbsolutePath).head}")
       }
       fail("Should not be able to send messages from a fenced producer.")
@@ -463,7 +546,7 @@ class TransactionsTest extends KafkaServerTestHarness {
     producer.beginTransaction()
 
     // The first message and hence the first AddPartitions request should be successfully sent.
-    val firstMessageResult = producer.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic1, "1", "1", willBeCommitted = false)).get()
+    val firstMessageResult = producer.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic1, null, "1", "1", willBeCommitted = false)).get()
     assertTrue(firstMessageResult.hasOffset)
 
     // Wait for the expiration cycle to kick in.
@@ -471,7 +554,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, null, "2", "2", willBeCommitted = false)).get()
       fail("should have raised a ProducerFencedException since the transaction has expired")
     } catch {
       case _: ProducerFencedException =>
@@ -480,9 +563,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))
 
@@ -503,8 +590,8 @@ class TransactionsTest extends KafkaServerTestHarness {
     val topicConfig = new Properties()
     topicConfig.put(KafkaConfig.MinInSyncReplicasProp, 2.toString)
 
-    TestUtils.createTopic(zkUtils, topicWith10Partitions, 10, numServers, servers, topicConfig)
-    TestUtils.createTopic(zkUtils, topicWith10PartitionsAndOneReplica, 10, 1, servers, new Properties())
+    createTopic(topicWith10Partitions, 10, numServers, topicConfig)
+    createTopic(topicWith10PartitionsAndOneReplica, 10, 1, new Properties())
 
     firstProducer.initTransactions()
 
@@ -532,10 +619,136 @@ class TransactionsTest extends KafkaServerTestHarness {
     }
   }
 
+  @Test(expected = classOf[KafkaException])
+  def testConsecutivelyRunInitTransactions(): Unit = {
+    val producer = createTransactionalProducer(transactionalId = "normalProducer")
+
+    producer.initTransactions()
+    producer.initTransactions()
+    fail("Should have raised a KafkaException")
+  }
+
+  @Test
+  def testBumpTransactionalEpoch(): Unit = {
+    val producer = createTransactionalProducer("transactionalProducer",
+      deliveryTimeoutMs = 5000, requestTimeoutMs = 5000)
+    val consumer = transactionalConsumers.head
+    try {
+      // Create a topic with RF=1 so that a single broker failure will render it unavailable
+      val testTopic = "test-topic"
+      createTopic(testTopic, numPartitions, 1, new Properties)
+      val partitionLeader = TestUtils.waitUntilLeaderIsKnown(servers, new TopicPartition(testTopic, 0))
+
+      producer.initTransactions()
+
+      producer.beginTransaction()
+      producer.send(TestUtils.producerRecordWithExpectedTransactionStatus(testTopic, 0, "4", "4", willBeCommitted = true))
+      producer.commitTransaction()
+
+      var producerStateEntry =
+        servers(partitionLeader).logManager.getLog(new TopicPartition(testTopic, 0)).get.producerStateManager.activeProducers.head._2
+      val producerId = producerStateEntry.producerId
+      val initialProducerEpoch = producerStateEntry.producerEpoch
+
+      producer.beginTransaction()
+      producer.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic1, null, "2", "2", willBeCommitted = false))
+
+      killBroker(partitionLeader) // kill the partition leader to prevent the batch from being submitted
+      val failedFuture = producer.send(TestUtils.producerRecordWithExpectedTransactionStatus(testTopic, 0, "3", "3", willBeCommitted = false))
+      Thread.sleep(6000) // Wait for the record to time out
+      restartDeadBrokers()
+
+      org.apache.kafka.test.TestUtils.assertFutureThrows(failedFuture, classOf[TimeoutException])
+      producer.abortTransaction()
+
+      producer.beginTransaction()
+      producer.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic2, null, "2", "2", willBeCommitted = true))
+      producer.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic1, null, "4", "4", willBeCommitted = true))
+      producer.send(TestUtils.producerRecordWithExpectedTransactionStatus(testTopic, 0, "1", "1", willBeCommitted = true))
+      producer.send(TestUtils.producerRecordWithExpectedTransactionStatus(testTopic, 0, "3", "3", willBeCommitted = true))
+      producer.commitTransaction()
+
+      consumer.subscribe(List(topic1, topic2, testTopic).asJava)
+
+      val records = consumeRecords(consumer, 5)
+      records.foreach { record =>
+        TestUtils.assertCommittedAndGetValue(record)
+      }
+
+      // Producers can safely abort and continue after the last record of a transaction timing out, so it's possible to
+      // get here without having bumped the epoch. If bumping the epoch is possible, the producer will attempt to, so
+      // check there that the epoch has actually increased
+      producerStateEntry =
+        servers(partitionLeader).logManager.getLog(new TopicPartition(testTopic, 0)).get.producerStateManager.activeProducers(producerId)
+      assertTrue(producerStateEntry.producerEpoch > initialProducerEpoch)
+    } finally {
+      producer.close(Duration.ZERO)
+    }
+  }
+
+  @Test
+  def testFailureToFenceEpoch(): Unit = {
+    val producer1 = transactionalProducers.head
+    val producer2 = createTransactionalProducer("transactional-producer", maxBlockMs = 1000)
+
+    producer1.initTransactions()
+
+    producer1.beginTransaction()
+    producer1.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic1, 0, "4", "4", willBeCommitted = true))
+    producer1.commitTransaction()
+
+    val partitionLeader = TestUtils.waitUntilLeaderIsKnown(servers, new TopicPartition(topic1, 0))
+    var producerStateEntry =
+      servers(partitionLeader).logManager.getLog(new TopicPartition(topic1, 0)).get.producerStateManager.activeProducers.head._2
+    val producerId = producerStateEntry.producerId
+    val initialProducerEpoch = producerStateEntry.producerEpoch
+
+    // Kill two brokers to bring the transaction log under min-ISR
+    killBroker(0)
+    killBroker(1)
+
+    try {
+      producer2.initTransactions()
+    } catch {
+      case _: TimeoutException =>
+        // good!
+      case e: Exception =>
+        fail("Got an unexpected exception from initTransactions", e)
+    } finally {
+      producer2.close()
+    }
+
+    restartDeadBrokers()
+
+    // Because the epoch was bumped in memory, attempting to begin a transaction with producer 1 should fail
+    try {
+      producer1.beginTransaction()
+    } catch {
+      case _: ProducerFencedException =>
+        // good!
+      case e: Exception =>
+        fail("Got an unexpected exception from commitTransaction", e)
+    } finally {
+      producer1.close()
+    }
+
+    val producer3 = createTransactionalProducer("transactional-producer", maxBlockMs = 5000)
+    producer3.initTransactions()
+
+    producer3.beginTransaction()
+    producer3.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic1, 0, "4", "4", willBeCommitted = true))
+    producer3.commitTransaction()
+
+    // Check that the epoch only increased by 1
+    producerStateEntry =
+      servers(partitionLeader).logManager.getLog(new TopicPartition(topic1, 0)).get.producerStateManager.activeProducers(producerId)
+    assertEquals((initialProducerEpoch + 1).toShort, producerStateEntry.producerEpoch)
+  }
+
   private def sendTransactionalMessagesWithValueRange(producer: KafkaProducer[Array[Byte], Array[Byte]], topic: String,
                                                       start: Int, end: Int, willBeCommitted: Boolean): Unit = {
     for (i <- start until end) {
-      producer.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic, i.toString, i.toString, willBeCommitted))
+      producer.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic, null, value = i.toString, willBeCommitted = willBeCommitted, key = i.toString))
     }
     producer.flush()
   }
@@ -557,32 +770,37 @@ class TransactionsTest extends KafkaServerTestHarness {
     serverProps
   }
 
-  private def createReadCommittedConsumer(group: String = "group", maxPollRecords: Int = 500,
+  private def createReadCommittedConsumer(group: String = "group",
+                                          maxPollRecords: Int = 500,
                                           props: Properties = new Properties) = {
-    props.put(ConsumerConfig.ISOLATION_LEVEL_CONFIG, "read_committed")
-    props.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false")
-    props.put(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, maxPollRecords.toString)
-    val consumer = TestUtils.createNewConsumer(TestUtils.getBrokerListStrFromServers(servers),
-      groupId = group, securityProtocol = SecurityProtocol.PLAINTEXT, props = Some(props))
+    val consumer = TestUtils.createConsumer(TestUtils.getBrokerListStrFromServers(servers),
+      groupId = group,
+      enableAutoCommit = false,
+      readCommitted = true,
+      maxPollRecords = maxPollRecords)
     transactionalConsumers += consumer
     consumer
   }
 
   private def createReadUncommittedConsumer(group: String) = {
-    val props = new Properties()
-    props.put(ConsumerConfig.ISOLATION_LEVEL_CONFIG, "read_uncommitted")
-    props.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false")
-    val consumer = TestUtils.createNewConsumer(TestUtils.getBrokerListStrFromServers(servers),
-      groupId = group, securityProtocol = SecurityProtocol.PLAINTEXT, props = Some(props))
+    val consumer = TestUtils.createConsumer(TestUtils.getBrokerListStrFromServers(servers),
+      groupId = group,
+      enableAutoCommit = false)
     nonTransactionalConsumers += consumer
     consumer
   }
 
-  private def createTransactionalProducer(transactionalId: String, transactionTimeoutMs: Long = 60000): KafkaProducer[Array[Byte], Array[Byte]] = {
+  private def createTransactionalProducer(transactionalId: String,
+                                          transactionTimeoutMs: Long = 60000,
+                                          maxBlockMs: Long = 60000,
+                                          deliveryTimeoutMs: Int = 120000,
+                                          requestTimeoutMs: Int = 30000): KafkaProducer[Array[Byte], Array[Byte]] = {
     val producer = TestUtils.createTransactionalProducer(transactionalId, servers,
-      transactionTimeoutMs = transactionTimeoutMs)
+      transactionTimeoutMs = transactionTimeoutMs,
+      maxBlockMs = maxBlockMs,
+      deliveryTimeoutMs = deliveryTimeoutMs,
+      requestTimeoutMs = requestTimeoutMs)
     transactionalProducers += producer
     producer
   }
-
 }
diff --git a/core/src/test/scala/integration/kafka/api/TransactionsWithMaxInFlightOneTest.scala b/core/src/test/scala/integration/kafka/api/TransactionsWithMaxInFlightOneTest.scala
new file mode 100644
index 0000000000000..4d00a2cf20433
--- /dev/null
+++ b/core/src/test/scala/integration/kafka/api/TransactionsWithMaxInFlightOneTest.scala
@@ -0,0 +1,131 @@
+/**
+  * Licensed to the Apache Software Foundation (ASF) under one or more
+  * contributor license agreements.  See the NOTICE file distributed with
+  * this work for additional information regarding copyright ownership.
+  * The ASF licenses this file to You under the Apache License, Version 2.0
+  * (the "License"); you may not use this file except in compliance with
+  * the License.  You may obtain a copy of the License at
+  *
+  *    http://www.apache.org/licenses/LICENSE-2.0
+  *
+  * Unless required by applicable law or agreed to in writing, software
+  * distributed under the License is distributed on an "AS IS" BASIS,
+  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  * See the License for the specific language governing permissions and
+  * limitations under the License.
+  */
+
+package integration.kafka.api
+
+import java.util.Properties
+
+import kafka.integration.KafkaServerTestHarness
+import kafka.server.KafkaConfig
+import kafka.utils.TestUtils
+import kafka.utils.TestUtils.consumeRecords
+import org.apache.kafka.clients.consumer.KafkaConsumer
+import org.apache.kafka.clients.producer.KafkaProducer
+import org.junit.{After, Before, Test}
+import org.junit.Assert.assertEquals
+
+import scala.collection.Seq
+import scala.collection.mutable.Buffer
+import scala.jdk.CollectionConverters._
+
+/**
+ * This is used to test transactions with one broker and `max.in.flight.requests.per.connection=1`.
+ * A single broker is used to verify edge cases where different requests are queued on the same connection.
+ */
+class TransactionsWithMaxInFlightOneTest extends KafkaServerTestHarness {
+  val numServers = 1
+
+  val topic1 = "topic1"
+  val topic2 = "topic2"
+  val numPartitions = 4
+
+  val transactionalProducers = Buffer[KafkaProducer[Array[Byte], Array[Byte]]]()
+  val transactionalConsumers = Buffer[KafkaConsumer[Array[Byte], Array[Byte]]]()
+
+  override def generateConfigs: Seq[KafkaConfig] = {
+    TestUtils.createBrokerConfigs(numServers, zkConnect).map(KafkaConfig.fromProps(_, serverProps()))
+  }
+
+  @Before
+  override def setUp(): Unit = {
+    super.setUp()
+    val topicConfig = new Properties()
+    topicConfig.put(KafkaConfig.MinInSyncReplicasProp, 1.toString)
+    createTopic(topic1, numPartitions, numServers, topicConfig)
+    createTopic(topic2, numPartitions, numServers, topicConfig)
+
+    createTransactionalProducer("transactional-producer")
+    createReadCommittedConsumer("transactional-group")
+  }
+
+  @After
+  override def tearDown(): Unit = {
+    transactionalProducers.foreach(_.close())
+    transactionalConsumers.foreach(_.close())
+    super.tearDown()
+  }
+
+  @Test
+  def testTransactionalProducerSingleBrokerMaxInFlightOne() = {
+    // We want to test with one broker to verify multiple requests queued on a connection
+    assertEquals(1, servers.size)
+
+    val producer = transactionalProducers.head
+    val consumer = transactionalConsumers.head
+
+    producer.initTransactions()
+
+    producer.beginTransaction()
+    producer.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic2, null, "2", "2", willBeCommitted = false))
+    producer.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic1, null, "4", "4", willBeCommitted = false))
+    producer.flush()
+    producer.abortTransaction()
+
+    producer.beginTransaction()
+    producer.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic1, null, "1", "1", willBeCommitted = true))
+    producer.send(TestUtils.producerRecordWithExpectedTransactionStatus(topic2, null, "3", "3", willBeCommitted = true))
+    producer.commitTransaction()
+
+    consumer.subscribe(List(topic1, topic2).asJava)
+
+    val records = consumeRecords(consumer, 2)
+    records.foreach { record =>
+      TestUtils.assertCommittedAndGetValue(record)
+    }
+  }
+
+  private def serverProps() = {
+    val serverProps = new Properties()
+    serverProps.put(KafkaConfig.AutoCreateTopicsEnableProp, false.toString)
+    serverProps.put(KafkaConfig.OffsetsTopicPartitionsProp, 1.toString)
+    serverProps.put(KafkaConfig.OffsetsTopicReplicationFactorProp, 1.toString)
+    serverProps.put(KafkaConfig.TransactionsTopicPartitionsProp, 1.toString)
+    serverProps.put(KafkaConfig.TransactionsTopicReplicationFactorProp, 1.toString)
+    serverProps.put(KafkaConfig.TransactionsTopicMinISRProp, 1.toString)
+    serverProps.put(KafkaConfig.ControlledShutdownEnableProp, true.toString)
+    serverProps.put(KafkaConfig.UncleanLeaderElectionEnableProp, false.toString)
+    serverProps.put(KafkaConfig.AutoLeaderRebalanceEnableProp, false.toString)
+    serverProps.put(KafkaConfig.GroupInitialRebalanceDelayMsProp, "0")
+    serverProps.put(KafkaConfig.TransactionsAbortTimedOutTransactionCleanupIntervalMsProp, "200")
+    serverProps
+  }
+
+  private def createReadCommittedConsumer(group: String) = {
+    val consumer = TestUtils.createConsumer(TestUtils.getBrokerListStrFromServers(servers),
+      groupId = group,
+      enableAutoCommit = false,
+      readCommitted = true)
+    transactionalConsumers += consumer
+    consumer
+  }
+
+  private def createTransactionalProducer(transactionalId: String): KafkaProducer[Array[Byte], Array[Byte]] = {
+    val producer = TestUtils.createTransactionalProducer(transactionalId, servers, maxInFlight = 1)
+    transactionalProducers += producer
+    producer
+  }
+}
diff --git a/core/src/test/scala/integration/kafka/api/UserClientIdQuotaTest.scala b/core/src/test/scala/integration/kafka/api/UserClientIdQuotaTest.scala
index 453ac910269b1..99fa16c48c934 100644
--- a/core/src/test/scala/integration/kafka/api/UserClientIdQuotaTest.scala
+++ b/core/src/test/scala/integration/kafka/api/UserClientIdQuotaTest.scala
@@ -15,10 +15,9 @@
 package kafka.api
 
 import java.io.File
-import java.util.Properties
 
 import kafka.server._
-import org.apache.kafka.common.security.auth.SecurityProtocol
+import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
 import org.apache.kafka.common.utils.Sanitizer
 import org.junit.Before
 
@@ -27,42 +26,61 @@ class UserClientIdQuotaTest extends BaseQuotaTest {
   override protected def securityProtocol = SecurityProtocol.SSL
   override protected lazy val trustStoreFile = Some(File.createTempFile("truststore", ".jks"))
 
-  override val userPrincipal = "O=A client,CN=localhost"
   override def producerClientId = "QuotasTestProducer-!@#$%^&*()"
   override def consumerClientId = "QuotasTestConsumer-!@#$%^&*()"
-  override def producerQuotaId = QuotaId(Some(Sanitizer.sanitize(userPrincipal)), Some(producerClientId), Some(Sanitizer.sanitize(producerClientId)))
-  override def consumerQuotaId = QuotaId(Some(Sanitizer.sanitize(userPrincipal)), Some(consumerClientId), Some(Sanitizer.sanitize(consumerClientId)))
 
   @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)
     super.setUp()
-    val defaultProps = quotaProperties(defaultProducerQuota, defaultConsumerQuota, defaultRequestQuota)
-    adminZkClient.changeUserOrUserClientIdConfig(ConfigEntityName.Default + "/clients/" + ConfigEntityName.Default, defaultProps)
-    waitForQuotaUpdate(defaultProducerQuota, defaultConsumerQuota, defaultRequestQuota)
+    quotaTestClients.alterClientQuotas(
+      quotaTestClients.clientQuotaAlteration(
+        quotaTestClients.clientQuotaEntity(Some(QuotaTestClients.DefaultEntity), Some(QuotaTestClients.DefaultEntity)),
+        Some(defaultProducerQuota), Some(defaultConsumerQuota), Some(defaultRequestQuota)
+      )
+    )
+    quotaTestClients.waitForQuotaUpdate(defaultProducerQuota, defaultConsumerQuota, defaultRequestQuota)
   }
 
-  override def overrideQuotas(producerQuota: Long, consumerQuota: Long, requestQuota: Double) {
-    val producerProps = new Properties()
-    producerProps.setProperty(DynamicConfig.Client.ProducerByteRateOverrideProp, producerQuota.toString)
-    producerProps.setProperty(DynamicConfig.Client.RequestPercentageOverrideProp, requestQuota.toString)
-    updateQuotaOverride(userPrincipal, producerClientId, producerProps)
+  override def createQuotaTestClients(topic: String, leaderNode: KafkaServer): QuotaTestClients = {
+    val producer = createProducer()
+    val consumer = createConsumer()
+    val adminClient = createAdminClient()
 
-    val consumerProps = new Properties()
-    consumerProps.setProperty(DynamicConfig.Client.ConsumerByteRateOverrideProp, consumerQuota.toString)
-    consumerProps.setProperty(DynamicConfig.Client.RequestPercentageOverrideProp, requestQuota.toString)
-    updateQuotaOverride(userPrincipal, consumerClientId, consumerProps)
-  }
+    new QuotaTestClients(topic, leaderNode, producerClientId, consumerClientId, producer, consumer, adminClient) {
+      override def userPrincipal: KafkaPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "O=A client,CN=localhost")
 
-  override def removeQuotaOverrides() {
-    val emptyProps = new Properties
-    adminZkClient.changeUserOrUserClientIdConfig(Sanitizer.sanitize(userPrincipal) + "/clients/" + Sanitizer.sanitize(producerClientId), emptyProps)
-    adminZkClient.changeUserOrUserClientIdConfig(Sanitizer.sanitize(userPrincipal) + "/clients/" + Sanitizer.sanitize(consumerClientId), emptyProps)
-  }
+      override def quotaMetricTags(clientId: String): Map[String, String] = {
+        Map("user" -> Sanitizer.sanitize(userPrincipal.getName), "client-id" -> clientId)
+      }
+
+      override def overrideQuotas(producerQuota: Long, consumerQuota: Long, requestQuota: Double): Unit = {
+        alterClientQuotas(
+          clientQuotaAlteration(
+            clientQuotaEntity(Some(userPrincipal.getName), Some(producerClientId)),
+            Some(producerQuota), None, Some(requestQuota)
+          ),
+          clientQuotaAlteration(
+            clientQuotaEntity(Some(userPrincipal.getName), Some(consumerClientId)),
+            None, Some(consumerQuota), Some(requestQuota)
+          )
+        )
+      }
 
-  private def updateQuotaOverride(userPrincipal: String, clientId: String, properties: Properties) {
-    adminZkClient.changeUserOrUserClientIdConfig(Sanitizer.sanitize(userPrincipal) + "/clients/" + Sanitizer.sanitize(clientId), properties)
+      override def removeQuotaOverrides(): Unit = {
+        alterClientQuotas(
+          clientQuotaAlteration(
+            clientQuotaEntity(Some(userPrincipal.getName), Some(producerClientId)),
+            None, None, None
+          ),
+          clientQuotaAlteration(
+            clientQuotaEntity(Some(userPrincipal.getName), Some(consumerClientId)),
+            None, None, None
+          )
+        )
+      }
+    }
   }
 }
diff --git a/core/src/test/scala/integration/kafka/api/UserQuotaTest.scala b/core/src/test/scala/integration/kafka/api/UserQuotaTest.scala
index 91a92faf68f5c..47493914dfb75 100644
--- a/core/src/test/scala/integration/kafka/api/UserQuotaTest.scala
+++ b/core/src/test/scala/integration/kafka/api/UserQuotaTest.scala
@@ -15,12 +15,10 @@
 package kafka.api
 
 import java.io.File
-import java.util.Properties
 
-import kafka.server.{ConfigEntityName, KafkaConfig, QuotaId}
+import kafka.server.{KafkaConfig, KafkaServer}
 import kafka.utils.JaasTestUtils
-import org.apache.kafka.common.security.auth.SecurityProtocol
-import org.apache.kafka.common.utils.Sanitizer
+import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
 import org.junit.{After, Before}
 
 class UserQuotaTest extends BaseQuotaTest with SaslSetup {
@@ -32,20 +30,19 @@ class UserQuotaTest extends BaseQuotaTest with SaslSetup {
   override protected val serverSaslProperties = Some(kafkaServerSaslProperties(kafkaServerSaslMechanisms, kafkaClientSaslMechanism))
   override protected val clientSaslProperties = Some(kafkaClientSaslProperties(kafkaClientSaslMechanism))
 
-  override val userPrincipal = JaasTestUtils.KafkaClientPrincipalUnqualifiedName2
-  override val producerQuotaId = QuotaId(Some(userPrincipal), None, None)
-  override val consumerQuotaId = QuotaId(Some(userPrincipal), None, None)
-
-
   @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)
     super.setUp()
-    val defaultProps = quotaProperties(defaultProducerQuota, defaultConsumerQuota, defaultRequestQuota)
-    adminZkClient.changeUserOrUserClientIdConfig(ConfigEntityName.Default, defaultProps)
-    waitForQuotaUpdate(defaultProducerQuota, defaultConsumerQuota, defaultRequestQuota)
+    quotaTestClients.alterClientQuotas(
+      quotaTestClients.clientQuotaAlteration(
+        quotaTestClients.clientQuotaEntity(Some(QuotaTestClients.DefaultEntity), None),
+        Some(defaultProducerQuota), Some(defaultConsumerQuota), Some(defaultRequestQuota)
+      )
+    )
+    quotaTestClients.waitForQuotaUpdate(defaultProducerQuota, defaultConsumerQuota, defaultRequestQuota)
   }
 
   @After
@@ -54,18 +51,35 @@ class UserQuotaTest extends BaseQuotaTest with SaslSetup {
     closeSasl()
   }
 
-  override def overrideQuotas(producerQuota: Long, consumerQuota: Long, requestQuota: Double) {
-    val props = quotaProperties(producerQuota, consumerQuota, requestQuota)
-    updateQuotaOverride(props)
-  }
+  override def createQuotaTestClients(topic: String, leaderNode: KafkaServer): QuotaTestClients = {
+    val producer = createProducer()
+    val consumer = createConsumer()
+    val adminClient = createAdminClient()
 
-  override def removeQuotaOverrides() {
-    val emptyProps = new Properties
-    updateQuotaOverride(emptyProps)
-    updateQuotaOverride(emptyProps)
-  }
+    new QuotaTestClients(topic, leaderNode, producerClientId, consumerClientId, producer, consumer, adminClient) {
+      override val userPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, JaasTestUtils.KafkaClientPrincipalUnqualifiedName2)
+
+      override def quotaMetricTags(clientId: String): Map[String, String] = {
+        Map("user" -> userPrincipal.getName, "client-id" -> "")
+      }
+
+      override def overrideQuotas(producerQuota: Long, consumerQuota: Long, requestQuota: Double): Unit = {
+        alterClientQuotas(
+          clientQuotaAlteration(
+            clientQuotaEntity(Some(userPrincipal.getName), None),
+            Some(producerQuota), Some(consumerQuota), Some(requestQuota)
+          )
+        )
+      }
 
-  private def updateQuotaOverride(properties: Properties) {
-    adminZkClient.changeUserOrUserClientIdConfig(Sanitizer.sanitize(userPrincipal), properties)
+      override def removeQuotaOverrides(): Unit = {
+        alterClientQuotas(
+          clientQuotaAlteration(
+            clientQuotaEntity(Some(userPrincipal.getName), None),
+            None, None, None
+          )
+        )
+      }
+    }
   }
 }
diff --git a/core/src/test/scala/integration/kafka/network/DynamicConnectionQuotaTest.scala b/core/src/test/scala/integration/kafka/network/DynamicConnectionQuotaTest.scala
new file mode 100644
index 0000000000000..ac8f49ee9b7cf
--- /dev/null
+++ b/core/src/test/scala/integration/kafka/network/DynamicConnectionQuotaTest.scala
@@ -0,0 +1,416 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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.network
+
+import java.io.IOException
+import java.net.{InetAddress, Socket}
+import java.util.concurrent._
+import java.util.{Collections, Properties}
+
+import kafka.server.{BaseRequestTest, DynamicConfig, KafkaConfig}
+import kafka.utils.TestUtils
+import org.apache.kafka.clients.admin.{Admin, AdminClientConfig}
+import org.apache.kafka.common.message.ProduceRequestData
+import org.apache.kafka.common.network.ListenerName
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.quota.ClientQuotaEntity
+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.apache.kafka.common.{KafkaException, requests}
+import org.junit.Assert._
+import org.junit.{After, Before, Test}
+import org.scalatest.Assertions.intercept
+
+import scala.annotation.nowarn
+import scala.jdk.CollectionConverters._
+
+class DynamicConnectionQuotaTest extends BaseRequestTest {
+
+  override def brokerCount = 1
+
+  val topic = "test"
+  val listener = ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT)
+  val localAddress = InetAddress.getByName("127.0.0.1")
+  val unknownHost = "255.255.0.1"
+  val plaintextListenerDefaultQuota = 30
+  var executor: ExecutorService = _
+
+  override def brokerPropertyOverrides(properties: Properties): Unit = {
+    properties.put(KafkaConfig.NumQuotaSamplesProp, "2".toString)
+    properties.put("listener.name.plaintext.max.connection.creation.rate", plaintextListenerDefaultQuota.toString)
+  }
+
+  @Before
+  override def setUp(): Unit = {
+    super.setUp()
+    TestUtils.createTopic(zkClient, topic, brokerCount, brokerCount, servers)
+  }
+
+  @After
+  override def tearDown(): Unit = {
+    try {
+      if (executor != null) {
+        executor.shutdownNow()
+        assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS))
+      }
+    } finally {
+      super.tearDown()
+    }
+  }
+
+  @Test
+  def testDynamicConnectionQuota(): Unit = {
+    val maxConnectionsPerIP = 5
+
+    def connectAndVerify(): Unit = {
+      val socket = connect()
+      try {
+        sendAndReceive[ProduceResponse](produceRequest, socket)
+      } finally {
+        socket.close()
+      }
+    }
+
+    val props = new Properties
+    props.put(KafkaConfig.MaxConnectionsPerIpProp, maxConnectionsPerIP.toString)
+    reconfigureServers(props, perBrokerConfig = false, (KafkaConfig.MaxConnectionsPerIpProp, maxConnectionsPerIP.toString))
+
+    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"))
+
+    verifyMaxConnections(maxConnectionsPerIPOverride, connectAndVerify)
+  }
+
+  @Test
+  def testDynamicListenerConnectionQuota(): Unit = {
+    val initialConnectionCount = connectionCount
+
+    def connectAndVerify(): Unit = {
+      val socket = connect("PLAINTEXT")
+      socket.setSoTimeout(1000)
+      try {
+        sendAndReceive[ProduceResponse](produceRequest, socket)
+      } finally {
+        socket.close()
+      }
+    }
+
+    // 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((() => createAndVerifyConnection()): Runnable)
+    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[ProduceResponse](produceRequest, socket)
+    })
+    plaintextConns.foreach(_.close())
+    internalConns.foreach(_.close())
+    TestUtils.waitUntilTrue(() => initialConnectionCount == connectionCount, "Connections not closed")
+  }
+
+  @Test
+  def testDynamicListenerConnectionCreationRateQuota(): Unit = {
+    // Create another listener. PLAINTEXT is an inter-broker listener
+    // keep default limits
+    val newListenerNames = Seq("PLAINTEXT", "EXTERNAL")
+    val newListeners = "PLAINTEXT://localhost:0,EXTERNAL://localhost:0"
+    val props = new Properties
+    props.put(KafkaConfig.ListenersProp, newListeners)
+    props.put(KafkaConfig.ListenerSecurityProtocolMapProp, "PLAINTEXT:PLAINTEXT,EXTERNAL:PLAINTEXT")
+    reconfigureServers(props, perBrokerConfig = true, (KafkaConfig.ListenersProp, newListeners))
+    waitForListener("EXTERNAL")
+
+    // The expected connection count after each test run
+    val initialConnectionCount = connectionCount
+
+    // new broker-wide connection rate limit
+    val connRateLimit = 9
+
+    // before setting connection rate to 10, verify we can do at least double that by default (no limit)
+    verifyConnectionRate(2 * connRateLimit, plaintextListenerDefaultQuota, "PLAINTEXT", ignoreIOExceptions = false)
+    waitForConnectionCount(initialConnectionCount)
+
+    // Reduce total broker connection rate limit to 9 at the cluster level and verify the limit is enforced
+    props.clear()  // so that we do not pass security protocol map which cannot be set at the cluster level
+    props.put(KafkaConfig.MaxConnectionCreationRateProp, connRateLimit.toString)
+    reconfigureServers(props, perBrokerConfig = false, (KafkaConfig.MaxConnectionCreationRateProp, connRateLimit.toString))
+    // verify EXTERNAL listener is capped by broker-wide quota (PLAINTEXT is not capped by broker-wide limit, since it
+    // has limited quota set and is a protected listener)
+    verifyConnectionRate(8, connRateLimit, "EXTERNAL", ignoreIOExceptions = false)
+    waitForConnectionCount(initialConnectionCount)
+
+    // Set 4 conn/sec rate limit for each listener and verify it gets enforced
+    val listenerConnRateLimit = 4
+    val plaintextListenerProp = s"${listener.configPrefix}${KafkaConfig.MaxConnectionCreationRateProp}"
+    props.put(s"listener.name.external.${KafkaConfig.MaxConnectionCreationRateProp}", listenerConnRateLimit.toString)
+    props.put(plaintextListenerProp, listenerConnRateLimit.toString)
+    reconfigureServers(props, perBrokerConfig = true, (plaintextListenerProp, listenerConnRateLimit.toString))
+
+    executor = Executors.newFixedThreadPool(newListenerNames.size)
+    val futures = newListenerNames.map { listener =>
+      executor.submit((() => verifyConnectionRate(3, listenerConnRateLimit, listener, ignoreIOExceptions = false)): Runnable)
+    }
+    futures.foreach(_.get(40, TimeUnit.SECONDS))
+    waitForConnectionCount(initialConnectionCount)
+
+    // increase connection rate limit on PLAINTEXT (inter-broker) listener to 12 and verify that it will be able to
+    // achieve this rate even though total connection rate may exceed broker-wide rate limit, while EXTERNAL listener
+    // should not exceed its listener limit
+    val newPlaintextRateLimit = 12
+    props.put(plaintextListenerProp, newPlaintextRateLimit.toString)
+    reconfigureServers(props, perBrokerConfig = true, (plaintextListenerProp, newPlaintextRateLimit.toString))
+
+    val plaintextFuture = executor.submit((() =>
+      verifyConnectionRate(10, newPlaintextRateLimit, "PLAINTEXT", ignoreIOExceptions = false)): Runnable)
+    val externalFuture = executor.submit((() =>
+      verifyConnectionRate(3, listenerConnRateLimit, "EXTERNAL", ignoreIOExceptions = false)): Runnable)
+
+    plaintextFuture.get(40, TimeUnit.SECONDS)
+    externalFuture.get(40, TimeUnit.SECONDS)
+    waitForConnectionCount(initialConnectionCount)
+  }
+
+  @Test
+  def testDynamicIpConnectionRateQuota(): Unit = {
+    val connRateLimit = 10
+    val initialConnectionCount = connectionCount
+    // before setting connection rate to 10, verify we can do at least double that by default (no limit)
+    verifyConnectionRate(2 * connRateLimit, plaintextListenerDefaultQuota, "PLAINTEXT", ignoreIOExceptions = false)
+    waitForConnectionCount(initialConnectionCount)
+    // set default IP connection rate quota, verify that we don't exceed the limit
+    updateIpConnectionRate(None, connRateLimit)
+    verifyConnectionRate(8, connRateLimit, "PLAINTEXT", ignoreIOExceptions = true)
+    waitForConnectionCount(initialConnectionCount)
+    // set a higher IP connection rate quota override, verify that the higher limit is now enforced
+    val newRateLimit = 18
+    updateIpConnectionRate(Some(localAddress.getHostAddress), newRateLimit)
+    verifyConnectionRate(14, newRateLimit, "PLAINTEXT", ignoreIOExceptions = true)
+    waitForConnectionCount(initialConnectionCount)
+  }
+
+  private def reconfigureServers(newProps: Properties, perBrokerConfig: Boolean, aPropToVerify: (String, String)): Unit = {
+    val initialConnectionCount = connectionCount
+    val adminClient = createAdminClient()
+    TestUtils.incrementalAlterConfigs(servers, adminClient, newProps, perBrokerConfig).all.get()
+    waitForConfigOnServer(aPropToVerify._1, aPropToVerify._2)
+    adminClient.close()
+    TestUtils.waitUntilTrue(() => initialConnectionCount == connectionCount,
+      s"Admin client connection not closed (initial = $initialConnectionCount, current = $connectionCount)")
+  }
+
+  private def updateIpConnectionRate(ip: Option[String], updatedRate: Int): Unit = {
+    val initialConnectionCount = connectionCount
+    val adminClient = createAdminClient()
+    try {
+      val entity = new ClientQuotaEntity(Map(ClientQuotaEntity.IP -> ip.orNull).asJava)
+      val request = Map(entity -> Map(DynamicConfig.Ip.IpConnectionRateOverrideProp -> Some(updatedRate.toDouble)))
+      TestUtils.alterClientQuotas(adminClient, request).all.get()
+      // use a random throwaway address if ip isn't specified to get the default value
+      TestUtils.waitUntilTrue(() => servers.head.socketServer.connectionQuotas.
+        connectionRateForIp(InetAddress.getByName(ip.getOrElse(unknownHost))) == updatedRate,
+        s"Timed out waiting for connection rate update to propagate"
+      )
+    } finally {
+      adminClient.close()
+    }
+    TestUtils.waitUntilTrue(() => initialConnectionCount == connectionCount,
+      s"Admin client connection not closed (initial = $initialConnectionCount, current = $connectionCount)")
+  }
+
+  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)
+    config.put(AdminClientConfig.METADATA_MAX_AGE_CONFIG, "10")
+    val adminClient = Admin.create(config)
+    adminClient
+  }
+
+  private def waitForConfigOnServer(propName: String, propValue: String, maxWaitMs: Long = 10000): Unit = {
+    TestUtils.retry(maxWaitMs) {
+      assertEquals(propValue, servers.head.config.originals.get(propName))
+    }
+  }
+
+  private def produceRequest: ProduceRequest =
+    requests.ProduceRequest.forCurrentMagic(new ProduceRequestData()
+      .setTopicData(new ProduceRequestData.TopicProduceDataCollection(
+        Collections.singletonList(new ProduceRequestData.TopicProduceData()
+          .setName(topic)
+          .setPartitionData(Collections.singletonList(new ProduceRequestData.PartitionProduceData()
+            .setIndex(0)
+            .setRecords(MemoryRecords.withRecords(CompressionType.NONE,
+              new SimpleRecord(System.currentTimeMillis(), "key".getBytes, "value".getBytes))))))
+        .iterator))
+      .setAcks((-1).toShort)
+      .setTimeoutMs(3000)
+      .setTransactionalId(null))
+      .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()
+    }
+  }
+
+  @nowarn("cat=deprecation")
+  private def verifyConnection(socket: Socket): Unit = {
+    val produceResponse = sendAndReceive[ProduceResponse](produceRequest, socket)
+    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")
+  }
+
+  private def connectAndVerify(listener: String, ignoreIOExceptions: Boolean): Unit = {
+    val socket = connect(listener)
+    try {
+      sendAndReceive[ProduceResponse](produceRequest, socket)
+    } catch {
+      // IP rate throttling can lead to disconnected sockets on client's end
+      case e: IOException => if (!ignoreIOExceptions) throw e
+    } finally {
+      socket.close()
+    }
+  }
+
+  private def waitForConnectionCount(expectedConnectionCount: Int): Unit = {
+    TestUtils.waitUntilTrue(() => expectedConnectionCount == connectionCount,
+      s"Connections not closed (expected = $expectedConnectionCount current = $connectionCount)")
+  }
+
+  /**
+   * this method simulates a workload that creates connection, sends produce request, closes connection,
+   * and verifies that rate does not exceed the given maximum limit `maxConnectionRate`
+   *
+   * Since producing a request and closing a connection also takes time, this method does not verify that the lower bound
+   * of actual rate is close to `maxConnectionRate`. Instead, use `minConnectionRate` parameter to verify that the rate
+   * is at least certain value. Note that throttling is tested and verified more accurately in ConnectionQuotasTest
+   */
+  private def verifyConnectionRate(minConnectionRate: Int, maxConnectionRate: Int, listener: String, ignoreIOExceptions: Boolean): Unit = {
+    // duration such that the maximum rate should be at most 20% higher than the rate limit. Since all connections
+    // can fall in the beginning of quota window, it is OK to create extra 2 seconds (window size) worth of connections
+    val runTimeMs = TimeUnit.SECONDS.toMillis(13)
+    val startTimeMs = System.currentTimeMillis
+    val endTimeMs = startTimeMs + runTimeMs
+
+    var connCount = 0
+    while (System.currentTimeMillis < endTimeMs) {
+      connectAndVerify(listener, ignoreIOExceptions)
+      connCount += 1
+    }
+    val elapsedMs = System.currentTimeMillis - startTimeMs
+    val actualRate = (connCount.toDouble / elapsedMs) * 1000
+    val rateCap = if (maxConnectionRate < Int.MaxValue) 1.2 * maxConnectionRate.toDouble else Int.MaxValue.toDouble
+    assertTrue(s"Listener $listener connection rate $actualRate must be below $rateCap", actualRate <= rateCap)
+    assertTrue(s"Listener $listener connection rate $actualRate must be above $minConnectionRate", actualRate >= minConnectionRate)
+  }
+}
diff --git a/core/src/test/scala/integration/kafka/server/DelayedFetchTest.scala b/core/src/test/scala/integration/kafka/server/DelayedFetchTest.scala
new file mode 100644
index 0000000000000..efbaef83ce11f
--- /dev/null
+++ b/core/src/test/scala/integration/kafka/server/DelayedFetchTest.scala
@@ -0,0 +1,216 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package kafka.server
+
+import java.util.Optional
+
+import scala.collection.Seq
+import kafka.cluster.Partition
+import kafka.log.LogOffsetSnapshot
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.errors.{FencedLeaderEpochException, NotLeaderOrFollowerException}
+import org.apache.kafka.common.message.OffsetForLeaderEpochResponseData.EpochEndOffset
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.record.MemoryRecords
+import org.apache.kafka.common.requests.FetchRequest
+import org.easymock.{EasyMock, EasyMockSupport}
+import org.junit.Test
+import org.junit.Assert._
+
+class DelayedFetchTest extends EasyMockSupport {
+  private val maxBytes = 1024
+  private val replicaManager: ReplicaManager = mock(classOf[ReplicaManager])
+  private val replicaQuota: ReplicaQuota = mock(classOf[ReplicaQuota])
+
+  @Test
+  def testFetchWithFencedEpoch(): 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))
+        .andReturn(partition)
+    EasyMock.expect(partition.fetchOffsetSnapshot(
+        currentLeaderEpoch,
+        fetchOnlyFromLeader = true))
+        .andThrow(new FencedLeaderEpochException("Requested epoch has been fenced"))
+    EasyMock.expect(replicaManager.isAddingReplica(EasyMock.anyObject(), EasyMock.anyInt())).andReturn(false)
+
+    expectReadFromReplica(replicaId, topicPartition, fetchStatus.fetchInfo, Errors.FENCED_LEADER_EPOCH)
+
+    replayAll()
+
+    assertTrue(delayedFetch.tryComplete())
+    assertTrue(delayedFetch.isCompleted)
+    assertTrue(fetchResultOpt.isDefined)
+
+    val fetchResult = fetchResultOpt.get
+    assertEquals(Errors.FENCED_LEADER_EPOCH, fetchResult.error)
+  }
+
+  @Test
+  def testNotLeaderOrFollower(): 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)
+
+    EasyMock.expect(replicaManager.getPartitionOrException(topicPartition))
+      .andThrow(new NotLeaderOrFollowerException(s"Replica for $topicPartition not available"))
+    expectReadFromReplica(replicaId, topicPartition, fetchStatus.fetchInfo, Errors.NOT_LEADER_OR_FOLLOWER)
+    EasyMock.expect(replicaManager.isAddingReplica(EasyMock.anyObject(), EasyMock.anyInt())).andReturn(false)
+
+    replayAll()
+
+    assertTrue(delayedFetch.tryComplete())
+    assertTrue(delayedFetch.isCompleted)
+    assertTrue(fetchResultOpt.isDefined)
+  }
+
+  @Test
+  def testDivergingEpoch(): Unit = {
+    val topicPartition = new TopicPartition("topic", 0)
+    val fetchOffset = 500L
+    val logStartOffset = 0L
+    val currentLeaderEpoch = Optional.of[Integer](10)
+    val lastFetchedEpoch = Optional.of[Integer](9)
+    val replicaId = 1
+
+    val fetchStatus = FetchPartitionStatus(
+      startOffsetMetadata = LogOffsetMetadata(fetchOffset),
+      fetchInfo = new FetchRequest.PartitionData(fetchOffset, logStartOffset, maxBytes, currentLeaderEpoch, lastFetchedEpoch))
+    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)).andReturn(partition)
+    val endOffsetMetadata = LogOffsetMetadata(messageOffset = 500L, segmentBaseOffset = 0L, relativePositionInSegment = 500)
+    EasyMock.expect(partition.fetchOffsetSnapshot(
+      currentLeaderEpoch,
+      fetchOnlyFromLeader = true))
+      .andReturn(LogOffsetSnapshot(0L, endOffsetMetadata, endOffsetMetadata, endOffsetMetadata))
+    EasyMock.expect(partition.lastOffsetForLeaderEpoch(currentLeaderEpoch, lastFetchedEpoch.get, fetchOnlyFromLeader = false))
+      .andReturn(new EpochEndOffset()
+        .setPartition(topicPartition.partition)
+        .setErrorCode(Errors.NONE.code)
+        .setLeaderEpoch(lastFetchedEpoch.get)
+        .setEndOffset(fetchOffset - 1))
+    EasyMock.expect(replicaManager.isAddingReplica(EasyMock.anyObject(), EasyMock.anyInt())).andReturn(false)
+    expectReadFromReplica(replicaId, topicPartition, fetchStatus.fetchInfo, Errors.NONE)
+    replayAll()
+
+    assertTrue(delayedFetch.tryComplete())
+    assertTrue(delayedFetch.isCompleted)
+    assertTrue(fetchResultOpt.isDefined)
+  }
+
+  private def buildFetchMetadata(replicaId: Int,
+                                 topicPartition: TopicPartition,
+                                 fetchStatus: FetchPartitionStatus): FetchMetadata = {
+    FetchMetadata(fetchMinBytes = 1,
+      fetchMaxBytes = maxBytes,
+      hardMaxBytesLimit = false,
+      fetchOnlyLeader = true,
+      fetchIsolation = FetchLogEnd,
+      isFromFollower = true,
+      replicaId = replicaId,
+      fetchPartitionStatus = Seq((topicPartition, fetchStatus)))
+  }
+
+  private def expectReadFromReplica(replicaId: Int,
+                                    topicPartition: TopicPartition,
+                                    fetchPartitionData: FetchRequest.PartitionData,
+                                    error: Errors): Unit = {
+    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, buildReadResult(error))))
+  }
+
+  private def buildReadResult(error: Errors): LogReadResult = {
+    LogReadResult(
+      exception = if (error != Errors.NONE) Some(error.exception) else None,
+      info = FetchDataInfo(LogOffsetMetadata.UnknownOffsetMetadata, MemoryRecords.EMPTY),
+      divergingEpoch = None,
+      highWatermark = -1L,
+      leaderLogStartOffset = -1L,
+      leaderLogEndOffset = -1L,
+      followerLogStartOffset = -1L,
+      fetchTimeMs = -1L,
+      lastStableOffset = None)
+  }
+
+}
diff --git a/core/src/test/scala/integration/kafka/server/DynamicBrokerReconfigurationTest.scala b/core/src/test/scala/integration/kafka/server/DynamicBrokerReconfigurationTest.scala
new file mode 100644
index 0000000000000..5c38fe27d3984
--- /dev/null
+++ b/core/src/test/scala/integration/kafka/server/DynamicBrokerReconfigurationTest.scala
@@ -0,0 +1,1887 @@
+/**
+  * Licensed to the Apache Software Foundation (ASF) under one
+  * or more contributor license agreements.  See the NOTICE file
+  * distributed with this work for additional information
+  * regarding copyright ownership.  The ASF licenses this file
+  * to you under the Apache License, Version 2.0 (the
+  * "License"); you may not use this file except in compliance
+  * with the License.  You may obtain a copy of the License at
+  *
+  * http://www.apache.org/licenses/LICENSE-2.0
+  *
+  * Unless required by applicable law or agreed to in writing, software
+  * distributed under the License is distributed on an "AS IS" BASIS,
+  * WITHOUT WARRANTIES 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.{Closeable, File, FileWriter, IOException, Reader, StringReader}
+import java.nio.file.{Files, Paths, StandardCopyOption}
+import java.lang.management.ManagementFactory
+import java.security.KeyStore
+import java.time.Duration
+import java.util
+import java.util.{Collections, Properties}
+import java.util.concurrent._
+
+import javax.management.ObjectName
+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.metrics.KafkaYammerMetrics
+import kafka.network.{Processor, RequestChannel}
+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}
+import org.apache.kafka.clients.producer.{KafkaProducer, ProducerConfig, ProducerRecord}
+import org.apache.kafka.common.{ClusterResource, ClusterResourceListener, Reconfigurable, TopicPartition, TopicPartitionInfo}
+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}
+import org.apache.kafka.common.network.{ListenerName, Mode}
+import org.apache.kafka.common.network.CertStores.{KEYSTORE_PROPS, TRUSTSTORE_PROPS}
+import org.apache.kafka.common.record.TimestampType
+import org.apache.kafka.common.security.auth.SecurityProtocol
+import org.apache.kafka.common.security.scram.ScramCredential
+import org.apache.kafka.common.serialization.{StringDeserializer, StringSerializer}
+import org.apache.kafka.test.{TestSslUtils, TestUtils => JTestUtils}
+import org.junit.Assert._
+import org.junit.{After, Before, Ignore, Test}
+import org.scalatest.Assertions.intercept
+
+import scala.annotation.nowarn
+import scala.collection._
+import scala.collection.mutable.ArrayBuffer
+import scala.jdk.CollectionConverters._
+import scala.collection.Seq
+
+object DynamicBrokerReconfigurationTest {
+  val SecureInternal = "INTERNAL"
+  val SecureExternal = "EXTERNAL"
+}
+
+class DynamicBrokerReconfigurationTest extends ZooKeeperTestHarness with SaslSetup {
+
+  import DynamicBrokerReconfigurationTest._
+
+  private val servers = new ArrayBuffer[KafkaServer]
+  private val numServers = 3
+  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[Admin]()
+  private val clientThreads = new ArrayBuffer[ShutdownableThread]()
+  private val executors = new ArrayBuffer[ExecutorService]
+  private val topic = "testtopic"
+
+  private val kafkaClientSaslMechanism = "PLAIN"
+  private val kafkaServerSaslMechanisms = List("PLAIN")
+
+  private val trustStoreFile1 = File.createTempFile("truststore", ".jks")
+  private val trustStoreFile2 = File.createTempFile("truststore", ".jks")
+  private val sslProperties1 = TestUtils.sslConfigs(Mode.SERVER, clientCert = false, Some(trustStoreFile1), "kafka")
+  private val sslProperties2 = TestUtils.sslConfigs(Mode.SERVER, clientCert = false, Some(trustStoreFile2), "kafka")
+  private val invalidSslProperties = invalidSslConfigs
+
+  def addExtraProps(props: Properties): Unit = {
+  }
+
+  @Before
+  override def setUp(): Unit = {
+    startSasl(jaasSections(kafkaServerSaslMechanisms, Some(kafkaClientSaslMechanism)))
+    super.setUp()
+
+    clearLeftOverProcessorMetrics() // clear metrics left over from other tests so that new ones can be tested
+
+    (0 until numServers).foreach { brokerId =>
+
+      val props = TestUtils.createBrokerConfig(brokerId, zkConnect)
+      props ++= securityProps(sslProperties1, TRUSTSTORE_PROPS)
+      // Ensure that we can support multiple listeners per security protocol and multiple security protocols
+      props.put(KafkaConfig.ListenersProp, s"$SecureInternal://localhost:0, $SecureExternal://localhost:0")
+      props.put(KafkaConfig.ListenerSecurityProtocolMapProp, s"$SecureInternal:SSL, $SecureExternal:SASL_SSL")
+      props.put(KafkaConfig.InterBrokerListenerNameProp, SecureInternal)
+      props.put(KafkaConfig.SslClientAuthProp, "requested")
+      props.put(KafkaConfig.SaslMechanismInterBrokerProtocolProp, "PLAIN")
+      props.put(KafkaConfig.ZkEnableSecureAclsProp, "true")
+      props.put(KafkaConfig.SaslEnabledMechanismsProp, kafkaServerSaslMechanisms.mkString(","))
+      props.put(KafkaConfig.LogSegmentBytesProp, "2000") // low value to test log rolling on config update
+      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)
+      addExtraProps(props)
+
+      props ++= sslProperties1
+      props ++= securityProps(sslProperties1, KEYSTORE_PROPS, listenerPrefix(SecureInternal))
+
+      // Set invalid top-level properties to ensure that listener config is used
+      // Don't set any dynamic configs here since they get overridden in tests
+      props ++= invalidSslProperties
+      props ++= securityProps(invalidSslProperties, KEYSTORE_PROPS, "")
+      props ++= securityProps(sslProperties1, KEYSTORE_PROPS, listenerPrefix(SecureExternal))
+
+      val kafkaConfig = KafkaConfig.fromProps(props)
+      configureDynamicKeystoreInZooKeeper(kafkaConfig, sslProperties1)
+
+      servers += TestUtils.createServer(kafkaConfig)
+    }
+
+    TestUtils.createTopic(zkClient, topic, numPartitions, replicationFactor = numServers, servers)
+    TestUtils.createTopic(zkClient, Topic.GROUP_METADATA_TOPIC_NAME, servers.head.config.offsetsTopicPartitions,
+      replicationFactor = numServers, servers, servers.head.groupCoordinator.offsetsTopicConfigs)
+
+    createAdminClient(SecurityProtocol.SSL, SecureInternal)
+
+    TestMetricsReporter.testReporters.clear()
+  }
+
+  @After
+  override def tearDown(): Unit = {
+    clientThreads.foreach(_.interrupt())
+    clientThreads.foreach(_.initiateShutdown())
+    clientThreads.foreach(_.join(5 * 1000))
+    executors.foreach(_.shutdownNow())
+    producers.foreach(_.close(Duration.ZERO))
+    consumers.foreach(_.close(Duration.ofMillis(0)))
+    adminClients.foreach(_.close())
+    TestUtils.shutdownServers(servers)
+    super.tearDown()
+    closeSasl()
+  }
+
+  @Test
+  def testConfigDescribeUsingAdminClient(): 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)
+      } else {
+        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,
+                      expectedPrefix: String, expectedSource: ConfigSource, expectedProps: Properties): Unit = {
+      if (isSensitive)
+        assertNull(s"Sensitive value returned for $configName", synonym.value)
+      else
+        assertEquals(expectedProps.getProperty(configName), synonym.value)
+      assertTrue(s"Expected listener config, got $synonym", synonym.name.startsWith(expectedPrefix))
+      assertEquals(expectedSource, synonym.source)
+    }
+
+    def verifySynonyms(configName: String, synonyms: util.List[ConfigSynonym], isSensitive: Boolean,
+                       prefix: String, defaultValue: Option[String]): Unit = {
+
+      val overrideCount = if (prefix.isEmpty) 0 else 2
+      assertEquals(s"Wrong synonyms for $configName: $synonyms", 1 + overrideCount + defaultValue.size, synonyms.size)
+      if (overrideCount > 0) {
+        val listenerPrefix = "listener.name.external.ssl."
+        verifySynonym(configName, synonyms.get(0), isSensitive, listenerPrefix, ConfigSource.DYNAMIC_BROKER_CONFIG, sslProperties1)
+        verifySynonym(configName, synonyms.get(1), isSensitive, listenerPrefix, ConfigSource.STATIC_BROKER_CONFIG, sslProperties1)
+      }
+      verifySynonym(configName, synonyms.get(overrideCount), isSensitive, "ssl.", ConfigSource.STATIC_BROKER_CONFIG, invalidSslProperties)
+      defaultValue.foreach { value =>
+        val defaultProps = new Properties
+        defaultProps.setProperty(configName, value)
+        verifySynonym(configName, synonyms.get(overrideCount + 1), isSensitive, "ssl.", ConfigSource.DEFAULT_CONFIG, defaultProps)
+      }
+    }
+
+    def verifySslConfig(prefix: String, expectedProps: Properties, configDesc: Config): Unit = {
+      // Validate file-based SSL keystore configs
+      val keyStoreProps = new util.HashSet[String](KEYSTORE_PROPS)
+      keyStoreProps.remove(SSL_KEYSTORE_KEY_CONFIG)
+      keyStoreProps.remove(SSL_KEYSTORE_CERTIFICATE_CHAIN_CONFIG)
+      keyStoreProps.forEach { configName =>
+        val desc = configEntry(configDesc, s"$prefix$configName")
+        val isSensitive = configName.contains("password")
+        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)
+      }
+    }
+
+    val adminClient = adminClients.head
+    alterSslKeystoreUsingConfigCommand(sslProperties1, SecureExternal)
+
+    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 = f"$${file:polling.interval:interval}"
+    val PollingIntervalUpdateVal = f"$${file:polling.interval:updinterval}"
+    val SslTruststoreTypeVal = f"$${file:ssl.truststore.type:storetype}"
+    val SslKeystorePasswordVal = f"$${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
+  def testKeyStoreAlter(): Unit = {
+    val topic2 = "testtopic2"
+    TestUtils.createTopic(zkClient, topic2, numPartitions, replicationFactor = numServers, servers)
+
+    // Start a producer and consumer that work with the current broker keystore.
+    // This should continue working while changes are made
+    val (producerThread, consumerThread) = startProduceConsume(retries = 0)
+    TestUtils.waitUntilTrue(() => consumerThread.received >= 10, "Messages not received")
+
+    // Producer with new truststore should fail to connect before keystore update
+    val producer1 = ProducerBuilder().trustStoreProps(sslProperties2).maxRetries(0).build()
+    verifyAuthenticationFailure(producer1)
+
+    // Update broker keystore for external listener
+    alterSslKeystoreUsingConfigCommand(sslProperties2, SecureExternal)
+
+    // New producer with old truststore should fail to connect
+    val producer2 = ProducerBuilder().trustStoreProps(sslProperties1).maxRetries(0).build()
+    verifyAuthenticationFailure(producer2)
+
+    // Produce/consume should work with new truststore with new producer/consumer
+    val producer = ProducerBuilder().trustStoreProps(sslProperties2).maxRetries(0).build()
+    val consumer = ConsumerBuilder("group1").trustStoreProps(sslProperties2).topic(topic2).build()
+    verifyProduceConsume(producer, consumer, 10, topic2)
+
+    // Broker keystore update for internal listener with incompatible keystore should fail without update
+    val adminClient = adminClients.head
+    alterSslKeystore(adminClient, sslProperties2, SecureInternal, expectFailure = true)
+    verifyProduceConsume(producer, consumer, 10, topic2)
+
+    // Broker keystore update for internal listener with compatible keystore should succeed
+    val sslPropertiesCopy = sslProperties1.clone().asInstanceOf[Properties]
+    val oldFile = new File(sslProperties1.getProperty(SSL_KEYSTORE_LOCATION_CONFIG))
+    val newFile = File.createTempFile("keystore", ".jks")
+    Files.copy(oldFile.toPath, newFile.toPath, StandardCopyOption.REPLACE_EXISTING)
+    sslPropertiesCopy.setProperty(SSL_KEYSTORE_LOCATION_CONFIG, newFile.getPath)
+    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)
+  }
+
+  @Test
+  def testTrustStoreAlter(): Unit = {
+    val producerBuilder = ProducerBuilder().listenerName(SecureInternal).securityProtocol(SecurityProtocol.SSL)
+
+    // Producer with new keystore should fail to connect before truststore update
+    verifyAuthenticationFailure(producerBuilder.keyStoreProps(sslProperties2).build())
+
+    // Update broker truststore for SSL listener with both certificates
+    val combinedStoreProps = mergeTrustStores(sslProperties1, sslProperties2)
+    val prefix = listenerPrefix(SecureInternal)
+    val existingDynamicProps = new Properties
+    servers.head.config.dynamicConfig.currentDynamicBrokerConfigs.foreach { case (k, v) =>
+      existingDynamicProps.put(k, v)
+    }
+    val newProps = new Properties
+    newProps ++= existingDynamicProps
+    newProps ++= securityProps(combinedStoreProps, TRUSTSTORE_PROPS, prefix)
+    reconfigureServers(newProps, perBrokerConfig = true,
+      (s"$prefix$SSL_TRUSTSTORE_LOCATION_CONFIG", combinedStoreProps.getProperty(SSL_TRUSTSTORE_LOCATION_CONFIG)))
+
+    def verifySslProduceConsume(keyStoreProps: Properties, group: String): Unit = {
+      val producer = producerBuilder.keyStoreProps(keyStoreProps).build()
+      val consumer = ConsumerBuilder(group)
+        .listenerName(SecureInternal)
+        .securityProtocol(SecurityProtocol.SSL)
+        .keyStoreProps(keyStoreProps)
+        .autoOffsetReset("latest")
+        .build()
+      verifyProduceConsume(producer, consumer, 10, topic)
+    }
+
+    // Produce/consume should work with old as well as new client keystore
+    verifySslProduceConsume(sslProperties1, "alter-truststore-1")
+    verifySslProduceConsume(sslProperties2, "alter-truststore-2")
+
+    // Revert to old truststore with only one certificate and update. Clients should connect only with old keystore.
+    val oldTruststoreProps = new Properties
+    oldTruststoreProps ++= existingDynamicProps
+    oldTruststoreProps ++= securityProps(sslProperties1, TRUSTSTORE_PROPS, prefix)
+    reconfigureServers(oldTruststoreProps, perBrokerConfig = true,
+      (s"$prefix$SSL_TRUSTSTORE_LOCATION_CONFIG", sslProperties1.getProperty(SSL_TRUSTSTORE_LOCATION_CONFIG)))
+    verifyAuthenticationFailure(producerBuilder.keyStoreProps(sslProperties2).build())
+    verifySslProduceConsume(sslProperties1, "alter-truststore-3")
+
+    // Update same truststore file to contain both certificates without changing any configs.
+    // Clients should connect successfully with either keystore after admin client AlterConfigsRequest completes.
+    Files.copy(Paths.get(combinedStoreProps.getProperty(SSL_TRUSTSTORE_LOCATION_CONFIG)),
+      Paths.get(sslProperties1.getProperty(SSL_TRUSTSTORE_LOCATION_CONFIG)),
+      StandardCopyOption.REPLACE_EXISTING)
+    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
+  def testLogCleanerConfig(): Unit = {
+    val (producerThread, consumerThread) = startProduceConsume(retries = 0)
+
+    verifyThreads("kafka-log-cleaner-thread-", countPerBroker = 1)
+
+    val props = new Properties
+    props.put(KafkaConfig.LogCleanerThreadsProp, "2")
+    props.put(KafkaConfig.LogCleanerDedupeBufferSizeProp, "20000000")
+    props.put(KafkaConfig.LogCleanerDedupeBufferLoadFactorProp, "0.8")
+    props.put(KafkaConfig.LogCleanerIoBufferSizeProp, "300000")
+    props.put(KafkaConfig.MessageMaxBytesProp, "40000")
+    props.put(KafkaConfig.LogCleanerIoMaxBytesPerSecondProp, "50000000")
+    props.put(KafkaConfig.LogCleanerBackoffMsProp, "6000")
+    reconfigureServers(props, perBrokerConfig = false, (KafkaConfig.LogCleanerThreadsProp, "2"))
+
+    // Verify cleaner config was updated. Wait for one of the configs to be updated and verify
+    // that all other others were updated at the same time since they are reconfigured together
+    val newCleanerConfig = servers.head.logManager.cleaner.currentConfig
+    TestUtils.waitUntilTrue(() => newCleanerConfig.numThreads == 2, "Log cleaner not reconfigured")
+    assertEquals(20000000, newCleanerConfig.dedupeBufferSize)
+    assertEquals(0.8, newCleanerConfig.dedupeBufferLoadFactor, 0.001)
+    assertEquals(300000, newCleanerConfig.ioBufferSize)
+    assertEquals(40000, newCleanerConfig.maxMessageSize)
+    assertEquals(50000000, newCleanerConfig.maxIoBytesPerSecond, 50000000)
+    assertEquals(6000, newCleanerConfig.backOffMs)
+
+    // Verify thread count
+    verifyThreads("kafka-log-cleaner-thread-", countPerBroker = 2)
+
+    // Stop a couple of threads and verify they are recreated if any config is updated
+    def cleanerThreads = Thread.getAllStackTraces.keySet.asScala.filter(_.getName.startsWith("kafka-log-cleaner-thread-"))
+    cleanerThreads.take(2).foreach(_.interrupt())
+    TestUtils.waitUntilTrue(() => cleanerThreads.size == (2 * numServers) - 2, "Threads did not exit")
+    props.put(KafkaConfig.LogCleanerBackoffMsProp, "8000")
+    reconfigureServers(props, perBrokerConfig = false, (KafkaConfig.LogCleanerBackoffMsProp, "8000"))
+    verifyThreads("kafka-log-cleaner-thread-", countPerBroker = 2)
+
+    // Verify that produce/consume worked throughout this test without any retries in producer
+    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.forEach { (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)
+
+    val props = new Properties
+    props.put(KafkaConfig.LogSegmentBytesProp, "4000")
+    props.put(KafkaConfig.LogRollTimeMillisProp, TimeUnit.HOURS.toMillis(2).toString)
+    props.put(KafkaConfig.LogRollTimeJitterMillisProp, TimeUnit.HOURS.toMillis(1).toString)
+    props.put(KafkaConfig.LogIndexSizeMaxBytesProp, "100000")
+    props.put(KafkaConfig.LogFlushIntervalMessagesProp, "1000")
+    props.put(KafkaConfig.LogFlushIntervalMsProp, "60000")
+    props.put(KafkaConfig.LogRetentionBytesProp, "10000000")
+    props.put(KafkaConfig.LogRetentionTimeMillisProp, TimeUnit.DAYS.toMillis(1).toString)
+    props.put(KafkaConfig.MessageMaxBytesProp, "100000")
+    props.put(KafkaConfig.LogIndexIntervalBytesProp, "10000")
+    props.put(KafkaConfig.LogCleanerDeleteRetentionMsProp, TimeUnit.DAYS.toMillis(1).toString)
+    props.put(KafkaConfig.LogCleanerMinCompactionLagMsProp, "60000")
+    props.put(KafkaConfig.LogDeleteDelayMsProp, "60000")
+    props.put(KafkaConfig.LogCleanerMinCleanRatioProp, "0.3")
+    props.put(KafkaConfig.LogCleanupPolicyProp, "delete")
+    props.put(KafkaConfig.UncleanLeaderElectionEnableProp, "false")
+    props.put(KafkaConfig.MinInSyncReplicasProp, "2")
+    props.put(KafkaConfig.CompressionTypeProp, "gzip")
+    props.put(KafkaConfig.LogPreAllocateProp, true.toString)
+    props.put(KafkaConfig.LogMessageTimestampTypeProp, TimestampType.LOG_APPEND_TIME.toString)
+    props.put(KafkaConfig.LogMessageTimestampDifferenceMaxMsProp, "1000")
+    props.put(KafkaConfig.LogMessageDownConversionEnableProp, "false")
+    reconfigureServers(props, perBrokerConfig = false, (KafkaConfig.LogSegmentBytesProp, "4000"))
+
+    // Verify that all broker defaults have been updated
+    servers.foreach { server =>
+      props.forEach { (k, v) =>
+        assertEquals(s"Not reconfigured $k", server.config.originals.get(k).toString, v)
+      }
+    }
+
+    // Verify that configs of existing logs have been updated
+    val newLogConfig = LogConfig(KafkaServer.copyKafkaConfigToLog(servers.head.config))
+    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) =>
+      val logConfigName = DynamicLogConfig.KafkaConfigToLogConfigName(k)
+      val expectedValue = if (k == KafkaConfig.LogCleanupPolicyProp) s"[$v]" else v
+      assertEquals(s"Not reconfigured $logConfigName for existing log", expectedValue,
+        log.config.originals.get(logConfigName).toString)
+    }
+    consumerThread.waitForMatchingRecords(record => record.timestampType == TimestampType.LOG_APPEND_TIME)
+
+    // Verify that the new config is actually used for new segments of existing logs
+    TestUtils.waitUntilTrue(() => log.logSegments.exists(_.size > 3000), "Log segment size increase not applied")
+
+    // Verify that overridden topic configs are not updated when broker default is updated
+    val log2 = servers.head.logManager.getLog(new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, 0))
+      .getOrElse(throw new IllegalStateException("Log not found"))
+    assertFalse("Overridden clean up policy should not be updated", log2.config.delete)
+    assertEquals(ProducerCompressionCodec.name, log2.config.compressionType)
+
+    // Verify that we can alter subset of log configs
+    props.clear()
+    props.put(KafkaConfig.LogMessageTimestampTypeProp, TimestampType.CREATE_TIME.toString)
+    props.put(KafkaConfig.LogMessageTimestampDifferenceMaxMsProp, "1000")
+    reconfigureServers(props, perBrokerConfig = false, (KafkaConfig.LogMessageTimestampTypeProp, TimestampType.CREATE_TIME.toString))
+    consumerThread.waitForMatchingRecords(record => record.timestampType == TimestampType.CREATE_TIME)
+    // Verify that invalid configs are not applied
+    val invalidProps = Map(
+      KafkaConfig.LogMessageTimestampDifferenceMaxMsProp -> "abc", // Invalid type
+      KafkaConfig.LogMessageTimestampTypeProp -> "invalid", // Invalid value
+      KafkaConfig.LogRollTimeMillisProp -> "0" // Fails KafkaConfig validation
+    )
+    invalidProps.foreach { case (k, v) =>
+      val newProps = new Properties
+      newProps ++= props
+      props.put(k, v)
+      reconfigureServers(props, perBrokerConfig = false, (k, props.getProperty(k)), expectFailure = true)
+    }
+
+    // Verify that even though broker defaults can be defined at default cluster level for consistent
+    // configuration across brokers, they can also be defined at per-broker level for testing
+    props.clear()
+    props.put(KafkaConfig.LogIndexSizeMaxBytesProp, "500000")
+    props.put(KafkaConfig.LogRetentionTimeMillisProp, TimeUnit.DAYS.toMillis(2).toString)
+    alterConfigsOnServer(servers.head, props)
+    assertEquals(500000, servers.head.config.values.get(KafkaConfig.LogIndexSizeMaxBytesProp))
+    assertEquals(TimeUnit.DAYS.toMillis(2), servers.head.config.values.get(KafkaConfig.LogRetentionTimeMillisProp))
+    servers.tail.foreach { server =>
+      assertEquals(Defaults.LogIndexSizeMaxBytes, server.config.values.get(KafkaConfig.LogIndexSizeMaxBytesProp))
+      assertEquals(1680000000L, server.config.values.get(KafkaConfig.LogRetentionTimeMillisProp))
+    }
+
+    // Verify that produce/consume worked throughout this test without any retries in producer
+    stopAndVerifyProduceConsume(producerThread, consumerThread)
+
+    // Verify that configuration at both per-broker level and default cluster level could be deleted and
+    // the default value should be restored
+    props.clear()
+    props.put(KafkaConfig.LogRetentionTimeMillisProp, "")
+    props.put(KafkaConfig.LogIndexSizeMaxBytesProp, "")
+    TestUtils.incrementalAlterConfigs(servers.take(1), adminClients.head, props, perBrokerConfig = true, opType = OpType.DELETE).all.get
+    TestUtils.incrementalAlterConfigs(servers, adminClients.head, props, perBrokerConfig = false, opType = OpType.DELETE).all.get
+    servers.foreach { server =>
+      waitForConfigOnServer(server, KafkaConfig.LogRetentionTimeMillisProp, 1680000000.toString)
+    }
+    servers.foreach { server =>
+      val log = server.logManager.getLog(new TopicPartition(topic, 0)).getOrElse(throw new IllegalStateException("Log not found"))
+      // Verify default values for these two configurations are restored on all brokers
+      TestUtils.waitUntilTrue(() => log.config.maxIndexSize == Defaults.LogIndexSizeMaxBytes && log.config.retentionMs == 1680000000L,
+        "Existing topic config using defaults not updated")
+    }
+  }
+
+  @Test
+  def testUncleanLeaderElectionEnable(): Unit = {
+    val controller = servers.find(_.config.brokerId == TestUtils.waitUntilControllerElected(zkClient)).get
+    val controllerId = controller.config.brokerId
+
+    // Create a topic with two replicas on brokers other than the controller
+    val topic = "testtopic2"
+    val assignment = Map(0 -> Seq((controllerId + 1) % servers.size, (controllerId + 2) % servers.size))
+    TestUtils.createTopic(zkClient, topic, assignment, servers)
+
+    val producer = ProducerBuilder().acks(1).build()
+    val consumer = ConsumerBuilder("unclean-leader-test").enableAutoCommit(false).topic(topic).build()
+    verifyProduceConsume(producer, consumer, numRecords = 10, topic)
+    consumer.commitSync()
+
+    def partitionInfo: TopicPartitionInfo =
+      adminClients.head.describeTopics(Collections.singleton(topic)).values.get(topic).get().partitions().get(0)
+
+    val partitionInfo0 = partitionInfo
+    assertEquals(partitionInfo0.replicas.get(0), partitionInfo0.leader)
+    val leaderBroker = servers.find(_.config.brokerId == partitionInfo0.replicas.get(0).id).get
+    val followerBroker = servers.find(_.config.brokerId == partitionInfo0.replicas.get(1).id).get
+
+    // Stop follower
+    followerBroker.shutdown()
+    followerBroker.awaitShutdown()
+
+    // Produce and consume some messages when the only follower is down, this should succeed since MinIsr is 1
+    verifyProduceConsume(producer, consumer, numRecords = 10, topic)
+    consumer.commitSync()
+
+    // Shutdown leader and startup follower
+    leaderBroker.shutdown()
+    leaderBroker.awaitShutdown()
+    followerBroker.startup()
+
+    // Verify that new leader is not elected with unclean leader disabled since there are no ISRs
+    TestUtils.waitUntilTrue(() => partitionInfo.leader == null, "Unclean leader elected")
+
+    // Enable unclean leader election
+    val newProps = new Properties
+    newProps.put(KafkaConfig.UncleanLeaderElectionEnableProp, "true")
+    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
+    val (newLeader, elected) = TestUtils.computeUntilTrue(partitionInfo.leader)(leader => leader != null)
+    assertTrue("Unclean leader not elected", elected)
+    assertEquals(followerBroker.config.brokerId, newLeader.id)
+
+    // New leader doesn't have the last 10 records committed on the old leader that have already been consumed.
+    // With unclean leader election enabled, we should be able to produce to the new leader. The first 10 records
+    // produced will not be consumed since they have offsets less than the consumer's committed offset.
+    // Next 10 records produced should be consumed.
+    (1 to 10).map(i => new ProducerRecord(topic, s"key$i", s"value$i"))
+      .map(producer.send)
+      .map(_.get(10, TimeUnit.SECONDS))
+    verifyProduceConsume(producer, consumer, numRecords = 10, topic)
+    consumer.commitSync()
+  }
+
+  @Test
+  def testThreadPoolResize(): Unit = {
+    val requestHandlerPrefix = "data-plane-kafka-request-handler-"
+    val networkThreadPrefix = "data-plane-kafka-network-thread-"
+    val fetcherThreadPrefix = "ReplicaFetcherThread-"
+    // Executor threads and recovery threads are not verified since threads may not be running
+    // For others, thread count should be configuredCount * threadMultiplier * numBrokers
+    val threadMultiplier = Map(
+      requestHandlerPrefix -> 1,
+      networkThreadPrefix -> 2, // 2 endpoints
+      fetcherThreadPrefix -> (servers.size - 1)
+    )
+
+    // Tolerate threads left over from previous tests
+    def leftOverThreadCount(prefix: String, perBrokerCount: Int): Int = {
+      val count = matchingThreads(prefix).size - perBrokerCount * servers.size * threadMultiplier(prefix)
+      if (count > 0) count else 0
+    }
+
+    val leftOverThreads = Map(
+      requestHandlerPrefix -> leftOverThreadCount(requestHandlerPrefix, servers.head.config.numIoThreads),
+      networkThreadPrefix -> leftOverThreadCount(networkThreadPrefix, servers.head.config.numNetworkThreads),
+      fetcherThreadPrefix -> leftOverThreadCount(fetcherThreadPrefix, servers.head.config.numReplicaFetchers)
+    )
+
+    def maybeVerifyThreadPoolSize(propName: String, size: Int, threadPrefix: String): Unit = {
+      val ignoreCount = leftOverThreads.getOrElse(threadPrefix, 0)
+      val expectedCountPerBroker = threadMultiplier.getOrElse(threadPrefix, 0) * size
+      if (expectedCountPerBroker > 0)
+        verifyThreads(threadPrefix, expectedCountPerBroker, ignoreCount)
+    }
+
+    def reducePoolSize(propName: String, currentSize: => Int, threadPrefix: String): Int = {
+      val newSize = if (currentSize / 2 == 0) 1 else currentSize / 2
+      resizeThreadPool(propName, newSize, threadPrefix)
+      newSize
+    }
+
+    def increasePoolSize(propName: String, currentSize: => Int, threadPrefix: String): Int = {
+      val newSize = if (currentSize == 1) currentSize * 2 else currentSize * 2 - 1
+      resizeThreadPool(propName, newSize, threadPrefix)
+      newSize
+    }
+
+    def resizeThreadPool(propName: String, newSize: Int, threadPrefix: String): Unit = {
+      val props = new Properties
+      props.put(propName, newSize.toString)
+      reconfigureServers(props, perBrokerConfig = false, (propName, newSize.toString))
+      maybeVerifyThreadPoolSize(propName, newSize, threadPrefix)
+    }
+
+    def verifyThreadPoolResize(propName: String, currentSize: => Int, threadPrefix: String, mayReceiveDuplicates: Boolean): Unit = {
+      maybeVerifyThreadPoolSize(propName, currentSize, threadPrefix)
+      val numRetries = if (mayReceiveDuplicates) 100 else 0
+      val (producerThread, consumerThread) = startProduceConsume(retries = numRetries)
+      var threadPoolSize = currentSize
+      (1 to 2).foreach { _ =>
+        threadPoolSize = reducePoolSize(propName, threadPoolSize, threadPrefix)
+        Thread.sleep(100)
+        threadPoolSize = increasePoolSize(propName, threadPoolSize, threadPrefix)
+        Thread.sleep(100)
+      }
+      stopAndVerifyProduceConsume(producerThread, consumerThread, mayReceiveDuplicates)
+      // Verify that all threads are alive
+      maybeVerifyThreadPoolSize(propName, threadPoolSize, threadPrefix)
+    }
+
+    val config = servers.head.config
+    verifyThreadPoolResize(KafkaConfig.NumIoThreadsProp, config.numIoThreads,
+      requestHandlerPrefix, mayReceiveDuplicates = false)
+    verifyThreadPoolResize(KafkaConfig.NumReplicaFetchersProp, config.numReplicaFetchers,
+      fetcherThreadPrefix, mayReceiveDuplicates = false)
+    verifyThreadPoolResize(KafkaConfig.BackgroundThreadsProp, config.backgroundThreads,
+      "kafka-scheduler-", mayReceiveDuplicates = false)
+    verifyThreadPoolResize(KafkaConfig.NumRecoveryThreadsPerDataDirProp, config.numRecoveryThreadsPerDataDir,
+      "", mayReceiveDuplicates = false)
+    verifyThreadPoolResize(KafkaConfig.NumNetworkThreadsProp, config.numNetworkThreads,
+      networkThreadPrefix, mayReceiveDuplicates = true)
+    verifyThreads("data-plane-kafka-socket-acceptor-", config.listeners.size)
+
+    verifyProcessorMetrics()
+    verifyMarkPartitionsForTruncation()
+  }
+
+  private def isProcessorMetric(metricName: MetricName): Boolean = {
+    val mbeanName = metricName.getMBeanName
+    mbeanName.contains(s"${Processor.NetworkProcessorMetricTag}=") || mbeanName.contains(s"${RequestChannel.ProcessorMetricTag}=")
+  }
+
+  private def clearLeftOverProcessorMetrics(): Unit = {
+    val metricsFromOldTests = KafkaYammerMetrics.defaultRegistry.allMetrics.keySet.asScala.filter(isProcessorMetric)
+    metricsFromOldTests.foreach(KafkaYammerMetrics.defaultRegistry.removeMetric)
+  }
+
+  // Verify that metrics from processors that were removed have been deleted.
+  // Since processor ids are not reused, it is sufficient to check metrics count
+  // based on the current number of processors
+  private def verifyProcessorMetrics(): Unit = {
+    val numProcessors = servers.head.config.numNetworkThreads * 2 // 2 listeners
+
+    val kafkaMetrics = servers.head.metrics.metrics().keySet.asScala
+      .filter(_.tags.containsKey(Processor.NetworkProcessorMetricTag))
+      .groupBy(_.tags.get(Processor.NetworkProcessorMetricTag))
+    assertEquals(numProcessors, kafkaMetrics.size)
+
+    KafkaYammerMetrics.defaultRegistry.allMetrics.keySet.asScala
+      .filter(isProcessorMetric)
+      .groupBy(_.getName)
+      .foreach { case (name, set) => assertEquals(s"Metrics not deleted $name", numProcessors, set.size) }
+  }
+
+  // Verify that replicaFetcherManager.markPartitionsForTruncation uses the current fetcher thread size
+  // to obtain partition assignment
+  private def verifyMarkPartitionsForTruncation(): Unit = {
+    val leaderId = 0
+    val partitions = (0 until numPartitions).map(i => new TopicPartition(topic, i)).filter { tp =>
+      zkClient.getLeaderForPartition(tp).contains(leaderId)
+    }
+    assertTrue(s"Partitions not found with leader $leaderId", partitions.nonEmpty)
+    partitions.foreach { tp =>
+      (1 to 2).foreach { i =>
+        val replicaFetcherManager = servers(i).replicaManager.replicaFetcherManager
+        val truncationOffset = tp.partition
+        replicaFetcherManager.markPartitionsForTruncation(leaderId, tp, truncationOffset)
+        val fetcherThreads = replicaFetcherManager.fetcherThreadMap.filter(_._2.fetchState(tp).isDefined)
+        assertEquals(1, fetcherThreads.size)
+        assertEquals(replicaFetcherManager.getFetcherId(tp), fetcherThreads.head._1.fetcherId)
+        val thread = fetcherThreads.head._2
+        assertEquals(Some(truncationOffset), thread.fetchState(tp).map(_.fetchOffset))
+        assertEquals(Some(Truncating), thread.fetchState(tp).map(_.state))
+      }
+    }
+  }
+
+  @Test
+  def testMetricsReporterUpdate(): Unit = {
+    // Add a new metrics reporter
+    val newProps = new Properties
+    newProps.put(TestMetricsReporter.PollingIntervalProp, "100")
+    configureMetricsReporters(Seq(classOf[TestMetricsReporter]), newProps)
+
+    val reporters = TestMetricsReporter.waitForReporters(servers.size)
+    reporters.foreach { reporter =>
+      reporter.verifyState(reconfigureCount = 0, deleteCount = 0, pollingInterval = 100)
+      assertFalse("No metrics found", reporter.kafkaMetrics.isEmpty)
+      reporter.verifyMetricValue("request-total", "socket-server-metrics")
+    }
+    assertEquals(servers.map(_.config.brokerId).toSet, TestMetricsReporter.configuredBrokers.toSet)
+
+    val clientId = "test-client-1"
+    val (producerThread, consumerThread) = startProduceConsume(retries = 0, clientId)
+    TestUtils.waitUntilTrue(() => consumerThread.received >= 5, "Messages not sent")
+
+    // Verify that JMX reporter is still active (test a metric registered after the dynamic reporter update)
+    val mbeanServer = ManagementFactory.getPlatformMBeanServer
+    val byteRate = mbeanServer.getAttribute(new ObjectName(s"kafka.server:type=Produce,client-id=$clientId"), "byte-rate")
+    assertTrue("JMX attribute not updated", byteRate.asInstanceOf[Double] > 0)
+
+    // Property not related to the metrics reporter config should not reconfigure reporter
+    newProps.setProperty("some.prop", "some.value")
+    reconfigureServers(newProps, perBrokerConfig = false, (TestMetricsReporter.PollingIntervalProp, "100"))
+    reporters.foreach(_.verifyState(reconfigureCount = 0, deleteCount = 0, pollingInterval = 100))
+
+    // Update of custom config of metrics reporter should reconfigure reporter
+    newProps.put(TestMetricsReporter.PollingIntervalProp, "1000")
+    reconfigureServers(newProps, perBrokerConfig = false, (TestMetricsReporter.PollingIntervalProp, "1000"))
+    reporters.foreach(_.verifyState(reconfigureCount = 1, deleteCount = 0, pollingInterval = 1000))
+
+    // Verify removal of metrics reporter
+    configureMetricsReporters(Seq.empty[Class[_]], newProps)
+    reporters.foreach(_.verifyState(reconfigureCount = 1, deleteCount = 1, pollingInterval = 1000))
+    TestMetricsReporter.testReporters.clear()
+
+    // Verify recreation of metrics reporter
+    newProps.put(TestMetricsReporter.PollingIntervalProp, "2000")
+    configureMetricsReporters(Seq(classOf[TestMetricsReporter]), newProps)
+    val newReporters = TestMetricsReporter.waitForReporters(servers.size)
+    newReporters.foreach(_.verifyState(reconfigureCount = 0, deleteCount = 0, pollingInterval = 2000))
+
+    // Verify that validation failure of metrics reporter fails reconfiguration and leaves config unchanged
+    newProps.put(KafkaConfig.MetricReporterClassesProp, "unknownMetricsReporter")
+    reconfigureServers(newProps, perBrokerConfig = false, (TestMetricsReporter.PollingIntervalProp, "2000"), expectFailure = true)
+    servers.foreach { server =>
+      assertEquals(classOf[TestMetricsReporter].getName, server.config.originals.get(KafkaConfig.MetricReporterClassesProp))
+    }
+    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 = 0, deleteCount = 0, pollingInterval = 2000))
+
+    // Delete reporters
+    configureMetricsReporters(Seq.empty[Class[_]], newProps)
+    TestMetricsReporter.testReporters.clear()
+
+    // Verify that even though metrics reporters can be defined at default cluster level for consistent
+    // configuration across brokers, they can also be defined at per-broker level for testing
+    newProps.put(KafkaConfig.MetricReporterClassesProp, classOf[TestMetricsReporter].getName)
+    newProps.put(TestMetricsReporter.PollingIntervalProp, "4000")
+    alterConfigsOnServer(servers.head, newProps)
+    TestUtils.waitUntilTrue(() => !TestMetricsReporter.testReporters.isEmpty, "Metrics reporter not created")
+    val perBrokerReporter = TestMetricsReporter.waitForReporters(1).head
+    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
+    stopAndVerifyProduceConsume(producerThread, consumerThread)
+  }
+
+  @Test
+  def testAdvertisedListenerUpdate(): Unit = {
+    val adminClient = adminClients.head
+    val externalAdminClient = createAdminClient(SecurityProtocol.SASL_SSL, SecureExternal)
+
+    // Ensure connections are made to brokers before external listener is made inaccessible
+    describeConfig(externalAdminClient)
+
+    // Update broker external listener to use invalid listener address
+    // any address other than localhost is sufficient to fail (either connection or host name verification failure)
+    val invalidHost = "192.168.0.1"
+    alterAdvertisedListener(adminClient, externalAdminClient, "localhost", invalidHost)
+
+    def validateEndpointsInZooKeeper(server: KafkaServer, endpointMatcher: String => Boolean): Unit = {
+      val brokerInfo = zkClient.getBroker(server.config.brokerId)
+      assertTrue("Broker not registered", brokerInfo.nonEmpty)
+      val endpoints = brokerInfo.get.endPoints.toString
+      assertTrue(s"Endpoint update not saved $endpoints", endpointMatcher(endpoints))
+    }
+
+    // Verify that endpoints have been updated in ZK for all brokers
+    servers.foreach(validateEndpointsInZooKeeper(_, endpoints => endpoints.contains(invalidHost)))
+
+    // Trigger session expiry and ensure that controller registers new advertised listener after expiry
+    val controllerEpoch = zkClient.getControllerEpoch
+    val controllerServer = servers(zkClient.getControllerId.getOrElse(throw new IllegalStateException("No controller")))
+    val controllerZkClient = controllerServer.zkClient
+    val sessionExpiringClient = createZooKeeperClientToTriggerSessionExpiry(controllerZkClient.currentZooKeeper)
+    sessionExpiringClient.close()
+    TestUtils.waitUntilTrue(() => zkClient.getControllerEpoch != controllerEpoch,
+      "Controller not re-elected after ZK session expiry")
+    TestUtils.retry(10000)(validateEndpointsInZooKeeper(controllerServer, endpoints => endpoints.contains(invalidHost)))
+
+    // Verify that producer connections fail since advertised listener is invalid
+    val bootstrap = TestUtils.bootstrapServers(servers, new ListenerName(SecureExternal))
+      .replaceAll(invalidHost, "localhost") // allow bootstrap connection to succeed
+    val producer1 = ProducerBuilder()
+      .trustStoreProps(sslProperties1)
+      .maxRetries(0)
+      .requestTimeoutMs(1000)
+      .deliveryTimeoutMs(1000)
+      .bootstrapServers(bootstrap)
+      .build()
+
+    assertTrue(intercept[ExecutionException] {
+      val future = producer1.send(new ProducerRecord(topic, "key", "value"))
+      future.get(2, TimeUnit.SECONDS)
+    }.getCause.isInstanceOf[org.apache.kafka.common.errors.TimeoutException])
+
+    alterAdvertisedListener(adminClient, externalAdminClient, invalidHost, "localhost")
+    servers.foreach(validateEndpointsInZooKeeper(_, endpoints => !endpoints.contains(invalidHost)))
+
+    // Verify that produce/consume work now
+    val topic2 = "testtopic2"
+    TestUtils.createTopic(zkClient, topic2, numPartitions, replicationFactor = numServers, servers)
+    val producer = ProducerBuilder().trustStoreProps(sslProperties1).maxRetries(0).build()
+    val consumer = ConsumerBuilder("group2").trustStoreProps(sslProperties1).topic(topic2).build()
+    verifyProduceConsume(producer, consumer, 10, topic2)
+
+    // Verify updating inter-broker listener
+    val props = new Properties
+    props.put(KafkaConfig.InterBrokerListenerNameProp, SecureExternal)
+    try {
+      reconfigureServers(props, perBrokerConfig = true, (KafkaConfig.InterBrokerListenerNameProp, SecureExternal))
+      fail("Inter-broker listener cannot be dynamically updated")
+    } catch {
+      case e: ExecutionException =>
+        assertTrue(s"Unexpected exception ${e.getCause}", e.getCause.isInstanceOf[InvalidRequestException])
+        servers.foreach(server => assertEquals(SecureInternal, server.config.interBrokerListenerName.value))
+    }
+  }
+
+  @Test
+  @Ignore // Re-enable once we make it less flaky (KAFKA-6824)
+  def testAddRemoveSslListener(): Unit = {
+    verifyAddListener("SSL", SecurityProtocol.SSL, Seq.empty)
+
+    // Restart servers and check secret rotation
+    servers.foreach(_.shutdown())
+    servers.foreach(_.awaitShutdown())
+    adminClients.foreach(_.close())
+    adminClients.clear()
+
+    // All passwords are currently encoded with password.encoder.secret. Encode with password.encoder.old.secret
+    // and update ZK. When each server is started, it should decode using password.encoder.old.secret and update
+    // ZK with newly encoded values using password.encoder.secret.
+    servers.foreach { server =>
+      val props = adminZkClient.fetchEntityConfig(ConfigType.Broker, server.config.brokerId.toString)
+      val propsEncodedWithOldSecret = props.clone().asInstanceOf[Properties]
+      val config = server.config
+      val oldSecret = "old-dynamic-config-secret"
+      config.dynamicConfig.staticBrokerConfigs.put(KafkaConfig.PasswordEncoderOldSecretProp, oldSecret)
+      val passwordConfigs = props.asScala.filter { case (k, _) => DynamicBrokerConfig.isPasswordConfig(k) }
+      assertTrue("Password configs not found", passwordConfigs.nonEmpty)
+      val passwordDecoder = createPasswordEncoder(config, config.passwordEncoderSecret)
+      val passwordEncoder = createPasswordEncoder(config, Some(new Password(oldSecret)))
+      passwordConfigs.foreach { case (name, value) =>
+        val decoded = passwordDecoder.decode(value).value
+        propsEncodedWithOldSecret.put(name, passwordEncoder.encode(new Password(decoded)))
+      }
+      val brokerId = server.config.brokerId
+      adminZkClient.changeBrokerConfig(Seq(brokerId), propsEncodedWithOldSecret)
+      val updatedProps = adminZkClient.fetchEntityConfig(ConfigType.Broker, brokerId.toString)
+      passwordConfigs.foreach { case (name, value) => assertNotEquals(props.get(value), updatedProps.get(name)) }
+
+      server.startup()
+      TestUtils.retry(10000) {
+        val newProps = adminZkClient.fetchEntityConfig(ConfigType.Broker, brokerId.toString)
+        passwordConfigs.foreach { case (name, value) =>
+          assertEquals(passwordDecoder.decode(value), passwordDecoder.decode(newProps.getProperty(name))) }
+      }
+    }
+
+    verifyListener(SecurityProtocol.SSL, None, "add-ssl-listener-group2")
+    createAdminClient(SecurityProtocol.SSL, SecureInternal)
+    verifyRemoveListener("SSL", SecurityProtocol.SSL, Seq.empty)
+  }
+
+  @Test
+  def testAddRemoveSaslListeners(): Unit = {
+    createScramCredentials(adminClients.head, JaasTestUtils.KafkaScramUser, JaasTestUtils.KafkaScramPassword)
+    createScramCredentials(adminClients.head, JaasTestUtils.KafkaScramAdmin, JaasTestUtils.KafkaScramAdminPassword)
+    initializeKerberos()
+    // make sure each server's credential cache has all the created credentials
+    // (check after initializing Kerberos to minimize delays)
+    List(JaasTestUtils.KafkaScramUser, JaasTestUtils.KafkaScramAdmin).foreach { scramUser =>
+      servers.foreach { server =>
+        ScramMechanism.values().filter(_ != ScramMechanism.UNKNOWN).foreach(mechanism =>
+          TestUtils.waitUntilTrue(() => server.credentialProvider.credentialCache.cache(
+            mechanism.mechanismName(), classOf[ScramCredential]).get(scramUser) != null,
+            s"$mechanism credentials not created for $scramUser"))
+      }}
+
+    //verifyAddListener("SASL_SSL", SecurityProtocol.SASL_SSL, Seq("SCRAM-SHA-512", "SCRAM-SHA-256", "PLAIN"))
+    verifyAddListener("SASL_PLAINTEXT", SecurityProtocol.SASL_PLAINTEXT, Seq("GSSAPI"))
+    //verifyRemoveListener("SASL_SSL", SecurityProtocol.SASL_SSL, Seq("SCRAM-SHA-512", "SCRAM-SHA-256", "PLAIN"))
+    verifyRemoveListener("SASL_PLAINTEXT", SecurityProtocol.SASL_PLAINTEXT, Seq("GSSAPI"))
+
+    // Verify that a listener added to a subset of servers doesn't cause any issues
+    // when metadata is processed by the client.
+    addListener(servers.tail, "SCRAM_LISTENER", SecurityProtocol.SASL_PLAINTEXT, Seq("SCRAM-SHA-256"))
+    val bootstrap = TestUtils.bootstrapServers(servers.tail, new ListenerName("SCRAM_LISTENER"))
+    val producer = ProducerBuilder().bootstrapServers(bootstrap)
+      .securityProtocol(SecurityProtocol.SASL_PLAINTEXT)
+      .saslMechanism("SCRAM-SHA-256")
+      .maxRetries(1000)
+      .build()
+    val partitions = producer.partitionsFor(topic).asScala
+    assertEquals(0, partitions.count(p => p.leader != null && p.leader.id == servers.head.config.brokerId))
+    assertTrue("Did not find partitions with no leader", partitions.exists(_.leader == null))
+  }
+
+  private def addListener(servers: Seq[KafkaServer], listenerName: String, securityProtocol: SecurityProtocol,
+                          saslMechanisms: Seq[String]): Unit = {
+    val config = servers.head.config
+    val existingListenerCount = config.listeners.size
+    val listeners = config.listeners
+      .map(e => s"${e.listenerName.value}://${e.host}:${e.port}")
+      .mkString(",") + s",$listenerName://localhost:0"
+    val listenerMap = config.listenerSecurityProtocolMap
+      .map { case (name, protocol) => s"${name.value}:${protocol.name}" }
+      .mkString(",") + s",$listenerName:${securityProtocol.name}"
+
+    val props = fetchBrokerConfigsFromZooKeeper(servers.head)
+    props.put(KafkaConfig.ListenersProp, listeners)
+    props.put(KafkaConfig.ListenerSecurityProtocolMapProp, listenerMap)
+    securityProtocol match {
+      case SecurityProtocol.SSL =>
+        addListenerPropsSsl(listenerName, props)
+      case SecurityProtocol.SASL_PLAINTEXT =>
+        addListenerPropsSasl(listenerName, saslMechanisms, props)
+      case SecurityProtocol.SASL_SSL =>
+        addListenerPropsSasl(listenerName, saslMechanisms, props)
+        addListenerPropsSsl(listenerName, props)
+      case SecurityProtocol.PLAINTEXT => // no additional props
+    }
+
+    // Add a config to verify that configs whose types are not known are not returned by describeConfigs()
+    val unknownConfig = "some.config"
+    props.put(unknownConfig, "some.config.value")
+
+    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")
+    TestUtils.waitUntilTrue(() => servers.forall(server => {
+      try {
+        server.socketServer.boundPort(new ListenerName(listenerName)) > 0
+      } catch {
+        case _: Exception => false
+      }
+    }), "Listener not created")
+
+    val brokerConfigs = describeConfig(adminClients.head, servers).entries.asScala
+    props.asScala.foreach { case (name, value) =>
+      val entry = brokerConfigs.find(_.name == name).getOrElse(throw new IllegalArgumentException(s"Config not found $name"))
+      if (DynamicBrokerConfig.isPasswordConfig(name) || name == unknownConfig)
+        assertNull(s"Password or unknown config returned $entry", entry.value)
+      else
+        assertEquals(value, entry.value)
+    }
+  }
+
+  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")
+      }
+    else
+      verifyListener(securityProtocol, None, s"add-listener-group-$securityProtocol")
+  }
+
+  private def verifyRemoveListener(listenerName: String, securityProtocol: SecurityProtocol,
+                                   saslMechanisms: Seq[String]): Unit = {
+    val saslMechanism = if (saslMechanisms.isEmpty) "" else saslMechanisms.head
+    val producer1 = ProducerBuilder().listenerName(listenerName)
+      .securityProtocol(securityProtocol)
+      .saslMechanism(saslMechanism)
+      .maxRetries(1000)
+      .build()
+    val consumer1 = ConsumerBuilder(s"remove-listener-group-$securityProtocol")
+      .listenerName(listenerName)
+      .securityProtocol(securityProtocol)
+      .saslMechanism(saslMechanism)
+      .autoOffsetReset("latest")
+      .build()
+    verifyProduceConsume(producer1, consumer1, numRecords = 10, topic)
+
+    val config = servers.head.config
+    val existingListenerCount = config.listeners.size
+    val listeners = config.listeners
+      .filter(e => e.listenerName.value != securityProtocol.name)
+      .map(e => s"${e.listenerName.value}://${e.host}:${e.port}")
+      .mkString(",")
+    val listenerMap = config.listenerSecurityProtocolMap
+      .filter { case (listenerName, _) => listenerName.value != securityProtocol.name }
+      .map { case (listenerName, protocol) => s"${listenerName.value}:${protocol.name}" }
+      .mkString(",")
+
+    val props = fetchBrokerConfigsFromZooKeeper(servers.head)
+    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.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)
+    val consumerFuture = verifyConnectionFailure(consumer1)
+
+    // Test that other listeners still work
+    val topic2 = "testtopic2"
+    TestUtils.createTopic(zkClient, topic2, numPartitions, replicationFactor = numServers, servers)
+    val producer2 = ProducerBuilder().trustStoreProps(sslProperties1).maxRetries(0).build()
+    val consumer2 = ConsumerBuilder(s"remove-listener-group2-$securityProtocol")
+      .trustStoreProps(sslProperties1)
+      .topic(topic2)
+      .autoOffsetReset("latest")
+      .build()
+    verifyProduceConsume(producer2, consumer2, numRecords = 10, topic2)
+
+    // Verify that producer/consumer using old listener don't work
+    verifyTimeout(producerFuture)
+    verifyTimeout(consumerFuture)
+  }
+
+  private def verifyListener(securityProtocol: SecurityProtocol, saslMechanism: Option[String], groupId: String): Unit = {
+    val mechanism = saslMechanism.getOrElse("")
+    val retries = 1000 // since it may take time for metadata to be updated on all brokers
+    val producer = ProducerBuilder().listenerName(securityProtocol.name)
+      .securityProtocol(securityProtocol)
+      .saslMechanism(mechanism)
+      .maxRetries(retries)
+      .build()
+    val consumer = ConsumerBuilder(s"add-listener-group-$securityProtocol-$mechanism")
+      .listenerName(securityProtocol.name)
+      .securityProtocol(securityProtocol)
+      .saslMechanism(mechanism)
+      .autoOffsetReset("latest")
+      .build()
+    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)
+  }
+
+  private def awaitInitialPositions(consumer: KafkaConsumer[_, _]): Unit = {
+    TestUtils.pollUntilTrue(consumer, () => !consumer.assignment.isEmpty, "Timed out while waiting for assignment")
+    consumer.assignment.forEach(consumer.position(_))
+  }
+
+  private def clientProps(securityProtocol: SecurityProtocol, saslMechanism: Option[String] = None): Properties = {
+    val props = new Properties
+    props.put(CommonClientConfigs.SECURITY_PROTOCOL_CONFIG, securityProtocol.name)
+    props.put(SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG, "HTTPS")
+    if (securityProtocol == SecurityProtocol.SASL_PLAINTEXT || securityProtocol == SecurityProtocol.SASL_SSL)
+      props ++= kafkaClientSaslProperties(saslMechanism.getOrElse(kafkaClientSaslMechanism), dynamicJaasConfig = true)
+    props ++= sslProperties1
+    securityProps(props, props.keySet)
+  }
+
+  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)
+    config.put(AdminClientConfig.METADATA_MAX_AGE_CONFIG, "10")
+    val adminClient = Admin.create(config)
+    adminClients += adminClient
+    adminClient
+  }
+
+  private def verifyProduceConsume(producer: KafkaProducer[String, String],
+                                   consumer: KafkaConsumer[String, String],
+                                   numRecords: Int,
+                                   topic: String): Unit = {
+    val producerRecords = (1 to numRecords).map(i => new ProducerRecord(topic, s"key$i", s"value$i"))
+    producerRecords.map(producer.send).map(_.get(10, TimeUnit.SECONDS))
+    TestUtils.pollUntilAtLeastNumRecords(consumer, numRecords)
+  }
+
+  private def verifyAuthenticationFailure(producer: KafkaProducer[_, _]): Unit = {
+    try {
+      producer.partitionsFor(topic)
+      fail("Producer connection did not fail with invalid keystore")
+    } catch {
+      case _:AuthenticationException => // expected exception
+    }
+  }
+
+  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)
+    }
+    val describeOptions = new DescribeConfigsOptions().includeSynonyms(true)
+    val describeResult = adminClient.describeConfigs(configResources.asJava, describeOptions).all.get
+    assertEquals(servers.size, describeResult.values.size)
+    val configDescription = describeResult.values.iterator.next
+    assertFalse("Configs are empty", configDescription.entries.isEmpty)
+    configDescription
+  }
+
+  private def securityProps(srcProps: Properties, propNames: util.Set[_], listenerPrefix: String = ""): Properties = {
+    val resultProps = new Properties
+    propNames.asScala.filter(srcProps.containsKey).foreach { propName =>
+      resultProps.setProperty(s"$listenerPrefix$propName", configValueAsString(srcProps.get(propName)))
+    }
+    resultProps
+  }
+
+  // Creates a new truststore with certificates from the provided stores and returns the properties of the new store
+  private def mergeTrustStores(trustStore1Props: Properties, trustStore2Props: Properties): Properties = {
+
+    def load(props: Properties): KeyStore = {
+      val ks = KeyStore.getInstance("JKS")
+      val password = props.get(SSL_TRUSTSTORE_PASSWORD_CONFIG).asInstanceOf[Password].value
+      val in = Files.newInputStream(Paths.get(props.getProperty(SSL_TRUSTSTORE_LOCATION_CONFIG)))
+      try {
+        ks.load(in, password.toCharArray)
+        ks
+      } finally {
+        in.close()
+      }
+    }
+    val cert1 = load(trustStore1Props).getCertificate("kafka")
+    val cert2 = load(trustStore2Props).getCertificate("kafka")
+    val certs = Map("kafka1" -> cert1, "kafka2" -> cert2)
+
+    val combinedStorePath = File.createTempFile("truststore", ".jks").getAbsolutePath
+    val password = trustStore1Props.get(SSL_TRUSTSTORE_PASSWORD_CONFIG).asInstanceOf[Password]
+    TestSslUtils.createTrustStore(combinedStorePath, password, certs.asJava)
+    val newStoreProps = new Properties
+    newStoreProps.put(SSL_TRUSTSTORE_LOCATION_CONFIG, combinedStorePath)
+    newStoreProps.put(SSL_TRUSTSTORE_PASSWORD_CONFIG, password)
+    newStoreProps.put(SSL_TRUSTSTORE_TYPE_CONFIG, "JKS")
+    newStoreProps
+  }
+
+  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,
+      (s"$configPrefix$SSL_KEYSTORE_LOCATION_CONFIG", props.getProperty(SSL_KEYSTORE_LOCATION_CONFIG)), expectFailure)
+  }
+
+  private def alterSslKeystoreUsingConfigCommand(props: Properties, listener: String): Unit = {
+    val configPrefix = listenerPrefix(listener)
+    val newProps = securityProps(props, KEYSTORE_PROPS, configPrefix)
+    alterConfigsUsingConfigCommand(newProps)
+    waitForConfig(s"$configPrefix$SSL_KEYSTORE_LOCATION_CONFIG", props.getProperty(SSL_KEYSTORE_LOCATION_CONFIG))
+  }
+
+  private def serverEndpoints(adminClient: Admin): String = {
+    val nodes = adminClient.describeCluster().nodes().get
+    nodes.asScala.map { node =>
+      s"${node.host}:${node.port}"
+    }.mkString(",")
+  }
+
+  @nowarn("cat=deprecation")
+  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 =>
+        if (e.listenerName.value == SecureExternal)
+          s"${e.listenerName.value}://$newHost:${server.boundPort(e.listenerName)}"
+        else
+          s"${e.listenerName.value}://${e.host}:${server.boundPort(e.listenerName)}"
+      }.mkString(",")
+      val configEntry = new ConfigEntry(KafkaConfig.AdvertisedListenersProp, newListeners)
+      (resource, new Config(Collections.singleton(configEntry)))
+    }.toMap.asJava
+    adminClient.alterConfigs(configs).all.get
+    servers.foreach { server =>
+      TestUtils.retry(10000) {
+        val externalListener = server.config.advertisedListeners.find(_.listenerName.value == SecureExternal)
+          .getOrElse(throw new IllegalStateException("External listener not found"))
+        assertTrue("Config not updated", externalListener.host == newHost)
+      }
+    }
+    val (endpoints, altered) = TestUtils.computeUntilTrue(serverEndpoints(externalAdminClient)) { endpoints =>
+      !endpoints.contains(oldHost)
+    }
+    assertTrue(s"Advertised listener update not propagated by controller: $endpoints", altered)
+  }
+
+  @nowarn("cat=deprecation")
+  private def alterConfigsOnServer(server: KafkaServer, props: Properties): Unit = {
+    val configEntries = props.asScala.map { case (k, v) => new ConfigEntry(k, v) }.toList.asJava
+    val newConfig = new Config(configEntries)
+    val configs = Map(new ConfigResource(ConfigResource.Type.BROKER, server.config.brokerId.toString) -> newConfig).asJava
+    adminClients.head.alterConfigs(configs).all.get
+    props.asScala.foreach { case (k, v) => waitForConfigOnServer(server, k, v) }
+  }
+
+  @nowarn("cat=deprecation")
+  private 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)
+    val configs = if (perBrokerConfig) {
+      servers.map { server =>
+        val resource = new ConfigResource(ConfigResource.Type.BROKER, server.config.brokerId.toString)
+        (resource, newConfig)
+      }.toMap.asJava
+    } else {
+      Map(new ConfigResource(ConfigResource.Type.BROKER, "") -> newConfig).asJava
+    }
+    adminClient.alterConfigs(configs)
+  }
+
+  private def reconfigureServers(newProps: Properties, perBrokerConfig: Boolean, aPropToVerify: (String, String), expectFailure: Boolean = false): Unit = {
+    val alterResult = alterConfigs(servers, adminClients.head, newProps, perBrokerConfig)
+    if (expectFailure) {
+      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
+        Seq(new ConfigResource(ConfigResource.Type.BROKER, ""))
+      brokerResources.foreach { brokerResource =>
+        val exception = intercept[ExecutionException](alterResult.values.get(brokerResource).get)
+        assertTrue(exception.getCause.isInstanceOf[InvalidRequestException])
+      }
+      servers.foreach { server =>
+        assertEquals(oldProps, server.config.values.asScala.filter { case (k, _) => newProps.containsKey(k) })
+      }
+    } else {
+      alterResult.all.get
+      waitForConfig(aPropToVerify._1, aPropToVerify._2)
+    }
+  }
+
+  private def configEntry(configDesc: Config, configName: String): ConfigEntry = {
+    configDesc.entries.asScala.find(cfg => cfg.name == configName)
+      .getOrElse(throw new IllegalStateException(s"Config not found $configName"))
+  }
+
+  private def listenerPrefix(name: String): String = new ListenerName(name).configPrefix
+
+  private def configureDynamicKeystoreInZooKeeper(kafkaConfig: KafkaConfig, sslProperties: Properties): Unit = {
+    val externalListenerPrefix = listenerPrefix(SecureExternal)
+    val sslStoreProps = new Properties
+    sslStoreProps ++= securityProps(sslProperties, KEYSTORE_PROPS, externalListenerPrefix)
+    sslStoreProps.put(KafkaConfig.PasswordEncoderSecretProp, kafkaConfig.passwordEncoderSecret.map(_.value).orNull)
+    zkClient.makeSurePersistentPathExists(ConfigEntityChangeNotificationZNode.path)
+
+    val entityType = ConfigType.Broker
+    val entityName = kafkaConfig.brokerId.toString
+
+    val passwordConfigs = sslStoreProps.asScala.keySet.filter(DynamicBrokerConfig.isPasswordConfig)
+    val passwordEncoder = createPasswordEncoder(kafkaConfig, kafkaConfig.passwordEncoderSecret)
+
+    if (passwordConfigs.nonEmpty) {
+      passwordConfigs.foreach { configName =>
+        val encodedValue = passwordEncoder.encode(new Password(sslStoreProps.getProperty(configName)))
+        sslStoreProps.setProperty(configName, encodedValue)
+      }
+    }
+    sslStoreProps.remove(KafkaConfig.PasswordEncoderSecretProp)
+    adminZkClient.changeConfigs(entityType, entityName, sslStoreProps)
+
+    val brokerProps = adminZkClient.fetchEntityConfig("brokers", kafkaConfig.brokerId.toString)
+    assertEquals(4, brokerProps.size)
+    assertEquals(sslProperties.get(SSL_KEYSTORE_TYPE_CONFIG),
+      brokerProps.getProperty(s"$externalListenerPrefix$SSL_KEYSTORE_TYPE_CONFIG"))
+    assertEquals(sslProperties.get(SSL_KEYSTORE_LOCATION_CONFIG),
+      brokerProps.getProperty(s"$externalListenerPrefix$SSL_KEYSTORE_LOCATION_CONFIG"))
+    assertEquals(sslProperties.get(SSL_KEYSTORE_PASSWORD_CONFIG),
+      passwordEncoder.decode(brokerProps.getProperty(s"$externalListenerPrefix$SSL_KEYSTORE_PASSWORD_CONFIG")))
+    assertEquals(sslProperties.get(SSL_KEY_PASSWORD_CONFIG),
+      passwordEncoder.decode(brokerProps.getProperty(s"$externalListenerPrefix$SSL_KEY_PASSWORD_CONFIG")))
+  }
+
+  private def createPasswordEncoder(config: KafkaConfig, secret: Option[Password]): PasswordEncoder = {
+    val encoderSecret = secret.getOrElse(throw new IllegalStateException("Password encoder secret not configured"))
+    new PasswordEncoder(encoderSecret,
+      config.passwordEncoderKeyFactoryAlgorithm,
+      config.passwordEncoderCipherAlgorithm,
+      config.passwordEncoderKeyLength,
+      config.passwordEncoderIterations)
+  }
+
+  private def waitForConfig(propName: String, propValue: String, maxWaitMs: Long = 10000): Unit = {
+    servers.foreach { server => waitForConfigOnServer(server, propName, propValue, maxWaitMs) }
+  }
+
+  private def waitForConfigOnServer(server: KafkaServer, propName: String, propValue: String, maxWaitMs: Long = 10000): Unit = {
+    TestUtils.retry(maxWaitMs) {
+      assertEquals(propValue, server.config.originals.get(propName))
+    }
+  }
+
+  private def configureMetricsReporters(reporters: Seq[Class[_]], props: Properties,
+                                        perBrokerConfig: Boolean = false): Unit = {
+    val reporterStr = reporters.map(_.getName).mkString(",")
+    props.put(KafkaConfig.MetricReporterClassesProp, reporterStr)
+    reconfigureServers(props, perBrokerConfig, (KafkaConfig.MetricReporterClassesProp, reporterStr))
+  }
+
+  private def invalidSslConfigs: Properties = {
+    val props = new Properties
+    props.put(SSL_KEYSTORE_LOCATION_CONFIG, "invalid/file/path")
+    props.put(SSL_KEYSTORE_PASSWORD_CONFIG, new Password("invalid"))
+    props.put(SSL_KEY_PASSWORD_CONFIG, new Password("invalid"))
+    props.put(SSL_KEYSTORE_TYPE_CONFIG, "PKCS12")
+    props
+  }
+
+  private def currentThreads: List[String] = {
+    Thread.getAllStackTraces.keySet.asScala.toList.map(_.getName)
+  }
+
+  private def matchingThreads(threadPrefix: String): List[String] = {
+    currentThreads.filter(_.startsWith(threadPrefix))
+  }
+
+  private def verifyThreads(threadPrefix: String, countPerBroker: Int, leftOverThreads: Int = 0): Unit = {
+    val expectedCount = countPerBroker * servers.size
+    val (threads, resized) = TestUtils.computeUntilTrue(matchingThreads(threadPrefix)) { matching =>
+      matching.size >= expectedCount &&  matching.size <= expectedCount + leftOverThreads
+    }
+    assertTrue(s"Invalid threads: expected $expectedCount, got ${threads.size}: $threads", resized)
+  }
+
+  private def startProduceConsume(retries: Int, producerClientId: String = "test-producer"): (ProducerThread, ConsumerThread) = {
+    val producerThread = new ProducerThread(producerClientId, retries)
+    clientThreads += producerThread
+    val consumerThread = new ConsumerThread(producerThread)
+    clientThreads += consumerThread
+    consumerThread.start()
+    producerThread.start()
+    TestUtils.waitUntilTrue(() => producerThread.sent >= 10, "Messages not sent")
+    (producerThread, consumerThread)
+  }
+
+  private def stopAndVerifyProduceConsume(producerThread: ProducerThread, consumerThread: ConsumerThread,
+                                          mayReceiveDuplicates: Boolean = false): Unit = {
+    TestUtils.waitUntilTrue(() => producerThread.sent >= 10, "Messages not sent")
+    producerThread.shutdown()
+    consumerThread.initiateShutdown()
+    consumerThread.awaitShutdown()
+    assertEquals(producerThread.lastSent, consumerThread.lastReceived)
+    assertEquals(0, consumerThread.missingRecords.size)
+    if (!mayReceiveDuplicates)
+      assertFalse("Duplicates not expected", consumerThread.duplicates)
+    assertFalse("Some messages received out of order", consumerThread.outOfOrder)
+  }
+
+  private def verifyConnectionFailure(producer: KafkaProducer[String, String]): Future[_] = {
+    val executor = Executors.newSingleThreadExecutor
+    executors += executor
+    val future = executor.submit(new Runnable() {
+      def run(): Unit = {
+        producer.send(new ProducerRecord(topic, "key", "value")).get
+      }
+    })
+    verifyTimeout(future)
+    future
+  }
+
+  private def verifyConnectionFailure(consumer: KafkaConsumer[String, String]): Future[_] = {
+    val executor = Executors.newSingleThreadExecutor
+    executors += executor
+    val future = executor.submit(new Runnable() {
+      def run(): Unit = {
+        consumer.commitSync()
+      }
+    })
+    verifyTimeout(future)
+    future
+  }
+
+  private def verifyTimeout(future: Future[_]): Unit = {
+    try {
+      future.get(100, TimeUnit.MILLISECONDS)
+      fail("Operation should not have completed")
+    } catch {
+      case _: TimeoutException => // expected exception
+    }
+  }
+
+  private def configValueAsString(value: Any): String = {
+    value match {
+      case password: Password => password.value
+      case list: util.List[_] => list.asScala.map(_.toString).mkString(",")
+      case _ => value.toString
+    }
+  }
+
+  private def addListenerPropsSsl(listenerName: String, props: Properties): Unit = {
+    props ++= securityProps(sslProperties1, KEYSTORE_PROPS, listenerPrefix(listenerName))
+    props ++= securityProps(sslProperties1, TRUSTSTORE_PROPS, listenerPrefix(listenerName))
+  }
+
+  private def addListenerPropsSasl(listener: String, mechanisms: Seq[String], props: Properties): Unit = {
+    val listenerName = new ListenerName(listener)
+    val prefix = listenerName.configPrefix
+    props.put(prefix + KafkaConfig.SaslEnabledMechanismsProp, mechanisms.mkString(","))
+    props.put(prefix + KafkaConfig.SaslKerberosServiceNameProp, "kafka")
+    mechanisms.foreach { mechanism =>
+      val jaasSection = jaasSections(Seq(mechanism), None, KafkaSasl, "").head
+      val jaasConfig = jaasSection.modules.head.toString
+      props.put(listenerName.saslMechanismConfigPrefix(mechanism) + KafkaConfig.SaslJaasConfigProp, jaasConfig)
+    }
+  }
+
+  private def alterConfigsUsingConfigCommand(props: Properties): Unit = {
+    val propsFile = TestUtils.tempFile()
+    val propsWriter = new FileWriter(propsFile)
+    try {
+      clientProps(SecurityProtocol.SSL).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
+    protected var _securityProtocol = SecurityProtocol.SASL_SSL
+    protected var _saslMechanism = kafkaClientSaslMechanism
+    protected var _clientId = "test-client"
+    protected val _propsOverride: Properties = new Properties
+
+    def bootstrapServers(bootstrap: String): this.type = { _bootstrapServers = Some(bootstrap); this }
+    def listenerName(listener: String): this.type = { _listenerName = listener; this }
+    def securityProtocol(protocol: SecurityProtocol): this.type = { _securityProtocol = protocol; this }
+    def saslMechanism(mechanism: String): this.type = { _saslMechanism = mechanism; this }
+    def clientId(id: String): this.type = { _clientId = id; this }
+    def keyStoreProps(props: Properties): this.type = { _propsOverride ++= securityProps(props, KEYSTORE_PROPS); this }
+    def trustStoreProps(props: Properties): this.type = { _propsOverride ++= securityProps(props, TRUSTSTORE_PROPS); this }
+
+    def bootstrapServers: String =
+      _bootstrapServers.getOrElse(TestUtils.bootstrapServers(servers, new ListenerName(_listenerName)))
+
+    def propsOverride: Properties = {
+      val props = clientProps(_securityProtocol, Some(_saslMechanism))
+      props.put(CommonClientConfigs.CLIENT_ID_CONFIG, _clientId)
+      props ++= _propsOverride
+      props
+    }
+
+    def build(): T
+  }
+
+  private case class ProducerBuilder() extends ClientBuilder[KafkaProducer[String, String]] {
+    private var _retries = Int.MaxValue
+    private var _acks = -1
+    private var _requestTimeoutMs = 30000
+    private var _deliveryTimeoutMs = 30000
+
+    def maxRetries(retries: Int): ProducerBuilder = { _retries = retries; this }
+    def acks(acks: Int): ProducerBuilder = { _acks = acks; this }
+    def requestTimeoutMs(timeoutMs: Int): ProducerBuilder = { _requestTimeoutMs = timeoutMs; this }
+    def deliveryTimeoutMs(timeoutMs: Int): ProducerBuilder = { _deliveryTimeoutMs= timeoutMs; this }
+
+    override def build(): KafkaProducer[String, String] = {
+      val producerProps = propsOverride
+      producerProps.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers)
+      producerProps.put(ProducerConfig.ACKS_CONFIG, _acks.toString)
+      producerProps.put(ProducerConfig.RETRIES_CONFIG, _retries.toString)
+      producerProps.put(ProducerConfig.DELIVERY_TIMEOUT_MS_CONFIG, _deliveryTimeoutMs.toString)
+      producerProps.put(ProducerConfig.REQUEST_TIMEOUT_MS_CONFIG, _requestTimeoutMs.toString)
+
+      val producer = new KafkaProducer[String, String](producerProps, new StringSerializer, new StringSerializer)
+      producers += producer
+      producer
+    }
+  }
+
+  private case class ConsumerBuilder(group: String) extends ClientBuilder[KafkaConsumer[String, String]] {
+    private var _autoOffsetReset = "earliest"
+    private var _enableAutoCommit = false
+    private var _topic = DynamicBrokerReconfigurationTest.this.topic
+
+    def autoOffsetReset(reset: String): ConsumerBuilder = { _autoOffsetReset = reset; this }
+    def enableAutoCommit(enable: Boolean): ConsumerBuilder = { _enableAutoCommit = enable; this }
+    def topic(topic: String): ConsumerBuilder = { _topic = topic; this }
+
+    override def build(): KafkaConsumer[String, String] = {
+      val consumerProps = propsOverride
+      consumerProps.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers)
+      consumerProps.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, _autoOffsetReset)
+      consumerProps.put(ConsumerConfig.GROUP_ID_CONFIG, group)
+      consumerProps.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, _enableAutoCommit.toString)
+
+      val consumer = new KafkaConsumer[String, String](consumerProps, new StringDeserializer, new StringDeserializer)
+      consumers += consumer
+
+      consumer.subscribe(Collections.singleton(_topic))
+      if (_autoOffsetReset == "latest")
+        awaitInitialPositions(consumer)
+      consumer
+    }
+  }
+
+  private class ProducerThread(clientId: String, retries: Int)
+    extends ShutdownableThread(clientId, isInterruptible = false) {
+
+    private val producer = ProducerBuilder().maxRetries(retries).clientId(clientId).build()
+    val lastSent = new ConcurrentHashMap[Int, Int]()
+    @volatile var sent = 0
+    override def doWork(): Unit = {
+      try {
+        while (isRunning) {
+          val key = sent.toString
+          val partition = sent % numPartitions
+          val record = new ProducerRecord(topic, partition, key, s"value$sent")
+          producer.send(record).get(10, TimeUnit.SECONDS)
+          lastSent.put(partition, sent)
+          sent += 1
+        }
+      } finally {
+        producer.close()
+      }
+    }
+  }
+
+  private class ConsumerThread(producerThread: ProducerThread) extends ShutdownableThread("test-consumer", isInterruptible = false) {
+    private val consumer = ConsumerBuilder("group1").enableAutoCommit(true).build()
+    val lastReceived = new ConcurrentHashMap[Int, Int]()
+    val missingRecords = new ConcurrentLinkedQueue[Int]()
+    @volatile var outOfOrder = false
+    @volatile var duplicates = false
+    @volatile var lastBatch: ConsumerRecords[String, String] = _
+    @volatile private var endTimeMs = Long.MaxValue
+    @volatile var received = 0
+    override def doWork(): Unit = {
+      try {
+        while (isRunning || (lastReceived != producerThread.lastSent && System.currentTimeMillis < endTimeMs)) {
+          val records = consumer.poll(Duration.ofMillis(50L))
+          received += records.count
+          if (!records.isEmpty) {
+            lastBatch = records
+            records.partitions.forEach { tp =>
+              val partition = tp.partition
+              records.records(tp).asScala.map(_.key.toInt).foreach { key =>
+                val prevKey = lastReceived.asScala.getOrElse(partition, partition - numPartitions)
+                val expectedKey = prevKey + numPartitions
+                if (key < prevKey)
+                  outOfOrder = true
+                else if (key == prevKey)
+                  duplicates = true
+                else {
+                  for (i <- expectedKey until key by numPartitions)
+                    missingRecords.add(i)
+                }
+                lastReceived.put(partition, key)
+                missingRecords.remove(key)
+              }
+            }
+          }
+        }
+      } finally {
+        consumer.close()
+      }
+    }
+
+    override def initiateShutdown(): Boolean = {
+      endTimeMs = System.currentTimeMillis + 10 * 1000
+      super.initiateShutdown()
+    }
+
+    def waitForMatchingRecords(predicate: ConsumerRecord[String, String] => Boolean): Unit = {
+      TestUtils.waitUntilTrue(() => {
+        val records = lastBatch
+        if (records == null || records.isEmpty)
+          false
+        else
+          records.asScala.toList.exists(predicate)
+      }, "Received records did not match")
+    }
+  }
+}
+
+object TestMetricsReporter {
+  val PollingIntervalProp = "polling.interval"
+  val testReporters = new ConcurrentLinkedQueue[TestMetricsReporter]()
+  val configuredBrokers = mutable.Set[Int]()
+
+  def waitForReporters(count: Int): List[TestMetricsReporter] = {
+    TestUtils.waitUntilTrue(() => testReporters.size == count, msg = "Metrics reporters not created")
+
+    val reporters = testReporters.asScala.toList
+    TestUtils.waitUntilTrue(() => reporters.forall(_.configureCount == 1), msg = "Metrics reporters not configured")
+    reporters
+  }
+}
+
+class TestMetricsReporter extends MetricsReporter with Reconfigurable with Closeable with ClusterResourceListener {
+  import TestMetricsReporter._
+  val kafkaMetrics = ArrayBuffer[KafkaMetric]()
+  @volatile var initializeCount = 0
+  @volatile var configureCount = 0
+  @volatile var reconfigureCount = 0
+  @volatile var closeCount = 0
+  @volatile var clusterUpdateCount = 0
+  @volatile var pollingInterval: Int = -1
+  testReporters.add(this)
+
+  override def init(metrics: util.List[KafkaMetric]): Unit = {
+    kafkaMetrics ++= metrics.asScala
+    initializeCount += 1
+  }
+
+  override def configure(configs: util.Map[String, _]): Unit = {
+    configuredBrokers += configs.get(KafkaConfig.BrokerIdProp).toString.toInt
+    configureCount += 1
+    pollingInterval = configs.get(PollingIntervalProp).toString.toInt
+  }
+
+  override def metricChange(metric: KafkaMetric): Unit = {
+  }
+
+  override def metricRemoval(metric: KafkaMetric): Unit = {
+    kafkaMetrics -= metric
+  }
+
+  override def onUpdate(clusterResource: ClusterResource): Unit = {
+    assertNotNull("Cluster id not set", clusterResource.clusterId)
+    clusterUpdateCount += 1
+  }
+
+  override def reconfigurableConfigs(): util.Set[String] = {
+    Set(PollingIntervalProp).asJava
+  }
+
+  override def validateReconfiguration(configs: util.Map[String, _]): Unit = {
+    val pollingInterval = configs.get(PollingIntervalProp).toString.toInt
+    if (pollingInterval <= 0)
+      throw new ConfigException(s"Invalid polling interval $pollingInterval")
+  }
+
+  override def reconfigure(configs: util.Map[String, _]): Unit = {
+    reconfigureCount += 1
+    pollingInterval = configs.get(PollingIntervalProp).toString.toInt
+  }
+
+  override def close(): Unit = {
+    closeCount += 1
+  }
+
+  def verifyState(reconfigureCount: Int, deleteCount: Int, pollingInterval: Int): Unit = {
+    assertEquals(1, initializeCount)
+    assertEquals(1, configureCount)
+    assertEquals(reconfigureCount, this.reconfigureCount)
+    assertEquals(deleteCount, closeCount)
+    assertEquals(1, clusterUpdateCount)
+    assertEquals(pollingInterval, this.pollingInterval)
+  }
+
+  def verifyMetricValue(name: String, group: String): Unit = {
+    val matchingMetrics = kafkaMetrics.filter(metric => metric.metricName.name == name && metric.metricName.group == group)
+    assertTrue("Metric not found", matchingMetrics.nonEmpty)
+    val total = matchingMetrics.foldLeft(0.0)((total, metric) => total + metric.metricValue.asInstanceOf[Double])
+    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
new file mode 100644
index 0000000000000..fa21a947d7b5f
--- /dev/null
+++ b/core/src/test/scala/integration/kafka/server/GssapiAuthenticationTest.scala
@@ -0,0 +1,276 @@
+/**
+  * Licensed to the Apache Software Foundation (ASF) under one
+  * or more contributor license agreements.  See the NOTICE file
+  * distributed with this work for additional information
+  * regarding copyright ownership.  The ASF licenses this file
+  * to you under the Apache License, Version 2.0 (the
+  * "License"); you may not use this file except in compliance
+  * with the License.  You may obtain a copy of the License at
+  *
+  * http://www.apache.org/licenses/LICENSE-2.0
+  *
+  * Unless required by applicable law or agreed to in writing, software
+  * distributed under the License is distributed on an "AS IS" BASIS,
+  * WITHOUT WARRANTIES 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.net.InetSocketAddress
+import java.time.Duration
+import java.util.{Collections, Properties}
+import java.util.concurrent.{CountDownLatch, Executors, TimeUnit}
+
+import kafka.api.{Both, IntegrationTestHarness, SaslSetup}
+import kafka.utils.TestUtils
+import org.apache.kafka.clients.CommonClientConfigs
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.config.SaslConfigs
+import org.apache.kafka.common.errors.SaslAuthenticationException
+import org.apache.kafka.common.network._
+import org.apache.kafka.common.security.{JaasContext, TestSecurityConfig}
+import org.apache.kafka.common.security.auth.{Login, SecurityProtocol}
+import org.apache.kafka.common.security.kerberos.KerberosLogin
+import org.apache.kafka.common.utils.{LogContext, MockTime}
+import org.junit.Assert._
+import org.junit.{After, Before, Test}
+
+import scala.jdk.CollectionConverters._
+
+class GssapiAuthenticationTest extends IntegrationTestHarness with SaslSetup {
+  override val brokerCount = 1
+  override protected def securityProtocol = SecurityProtocol.SASL_PLAINTEXT
+
+  private val kafkaClientSaslMechanism = "GSSAPI"
+  private val kafkaServerSaslMechanisms = List("GSSAPI")
+
+  private val numThreads = 10
+  private val executor = Executors.newFixedThreadPool(numThreads)
+  private val clientConfig: Properties = new Properties
+  private var serverAddr: InetSocketAddress = _
+  private val time = new MockTime(10)
+  val topic = "topic"
+  val part = 0
+  val tp = new TopicPartition(topic, part)
+  private val failedAuthenticationDelayMs = 2000
+
+  @Before
+  override def setUp(): Unit = {
+    TestableKerberosLogin.reset()
+    startSasl(jaasSections(kafkaServerSaslMechanisms, Option(kafkaClientSaslMechanism), Both))
+    serverConfig.put(KafkaConfig.SslClientAuthProp, "required")
+    serverConfig.put(KafkaConfig.FailedAuthenticationDelayMsProp, failedAuthenticationDelayMs.toString)
+    super.setUp()
+    serverAddr = new InetSocketAddress("localhost",
+      servers.head.boundPort(ListenerName.forSecurityProtocol(SecurityProtocol.SASL_PLAINTEXT)))
+
+    clientConfig.put(CommonClientConfigs.SECURITY_PROTOCOL_CONFIG, SecurityProtocol.SASL_PLAINTEXT.name)
+    clientConfig.put(SaslConfigs.SASL_MECHANISM, kafkaClientSaslMechanism)
+    clientConfig.put(SaslConfigs.SASL_JAAS_CONFIG, jaasClientLoginModule(kafkaClientSaslMechanism))
+    clientConfig.put(CommonClientConfigs.CONNECTIONS_MAX_IDLE_MS_CONFIG, "5000")
+
+    // create the test topic with all the brokers as replicas
+    createTopic(topic, 2, brokerCount)
+  }
+
+  @After
+  override def tearDown(): Unit = {
+    executor.shutdownNow()
+    super.tearDown()
+    closeSasl()
+    TestableKerberosLogin.reset()
+  }
+
+  /**
+   * Tests that Kerberos replay error `Request is a replay (34)` is not handled as an authentication exception
+   * since replay detection used to detect DoS attacks may occasionally reject valid concurrent requests.
+   */
+  @Test
+  def testRequestIsAReplay(): Unit = {
+    val successfulAuthsPerThread = 10
+    val futures = (0 until numThreads).map(_ => executor.submit(new Runnable {
+      override def run(): Unit = verifyRetriableFailuresDuringAuthentication(successfulAuthsPerThread)
+    }))
+    futures.foreach(_.get(60, TimeUnit.SECONDS))
+    assertEquals(0, TestUtils.totalMetricValue(servers.head, "failed-authentication-total"))
+    val successfulAuths = TestUtils.totalMetricValue(servers.head, "successful-authentication-total")
+    assertTrue("Too few authentications: " + successfulAuths, successfulAuths > successfulAuthsPerThread * numThreads)
+  }
+
+  /**
+   * Verifies that there are no authentication failures during Kerberos re-login. If authentication
+   * is performed when credentials are unavailable between logout and login, we handle it as a
+   * transient error and not an authentication failure so that clients may retry.
+   */
+  @Test
+  def testReLogin(): Unit = {
+    val selector = createSelectorWithRelogin()
+    try {
+      val login = TestableKerberosLogin.instance
+      assertNotNull(login)
+      executor.submit(() => login.reLogin(), 0)
+
+      val node1 = "1"
+      selector.connect(node1, serverAddr, 1024, 1024)
+      login.logoutResumeLatch.countDown()
+      login.logoutCompleteLatch.await(15, TimeUnit.SECONDS)
+      assertFalse("Authenticated during re-login", pollUntilReadyOrDisconnected(selector, node1))
+
+      login.reLoginResumeLatch.countDown()
+      login.reLoginCompleteLatch.await(15, TimeUnit.SECONDS)
+      val node2 = "2"
+      selector.connect(node2, serverAddr, 1024, 1024)
+      assertTrue("Authenticated failed after re-login", pollUntilReadyOrDisconnected(selector, node2))
+    } finally {
+      selector.close()
+    }
+  }
+
+  /**
+   * Tests that Kerberos error `Server not found in Kerberos database (7)` is handled
+   * as a fatal authentication failure.
+   */
+  @Test
+  def testServerNotFoundInKerberosDatabase(): Unit = {
+    val jaasConfig = clientConfig.getProperty(SaslConfigs.SASL_JAAS_CONFIG)
+    val invalidServiceConfig = jaasConfig.replace("serviceName=\"kafka\"", "serviceName=\"invalid-service\"")
+    clientConfig.put(SaslConfigs.SASL_JAAS_CONFIG, invalidServiceConfig)
+    clientConfig.put(SaslConfigs.SASL_KERBEROS_SERVICE_NAME, "invalid-service")
+    verifyNonRetriableAuthenticationFailure()
+  }
+
+  /**
+   * Test that when client fails to verify authenticity of the server, the resulting failed authentication exception
+   * is thrown immediately, and is not affected by connection.failed.authentication.delay.ms.
+   */
+  @Test
+  def testServerAuthenticationFailure(): Unit = {
+    // Setup client with a non-existent service principal, so that server authentication fails on the client
+    val clientLoginContext = jaasClientLoginModule(kafkaClientSaslMechanism, Some("another-kafka-service"))
+    val configOverrides = new Properties()
+    configOverrides.setProperty(SaslConfigs.SASL_JAAS_CONFIG, clientLoginContext)
+    val consumer = createConsumer(configOverrides = configOverrides)
+    consumer.assign(List(tp).asJava)
+
+    val startMs = System.currentTimeMillis()
+    try {
+      consumer.poll(Duration.ofMillis(50))
+      fail()
+    } catch {
+      case _: SaslAuthenticationException =>
+    }
+    val endMs = System.currentTimeMillis()
+    require(endMs - startMs < failedAuthenticationDelayMs, "Failed authentication must not be delayed on the client")
+    consumer.close()
+  }
+
+  /**
+   * Verifies that any exceptions during authentication with the current `clientConfig` are
+   * notified with disconnect state `AUTHENTICATE` (and not `AUTHENTICATION_FAILED`). This
+   * is to ensure that NetworkClient doesn't handle this as a fatal authentication failure,
+   * but as a transient I/O exception. So Producer/Consumer/AdminClient will retry
+   * any operation based on their configuration until timeout and will not propagate
+   * the exception to the application.
+   */
+  private def verifyRetriableFailuresDuringAuthentication(numSuccessfulAuths: Int): Unit = {
+    val selector = createSelector()
+    try {
+      var actualSuccessfulAuths = 0
+      while (actualSuccessfulAuths < numSuccessfulAuths) {
+        val nodeId = actualSuccessfulAuths.toString
+        selector.connect(nodeId, serverAddr, 1024, 1024)
+        val isReady = pollUntilReadyOrDisconnected(selector, nodeId)
+        if (isReady)
+          actualSuccessfulAuths += 1
+        selector.close(nodeId)
+      }
+    } finally {
+      selector.close()
+    }
+  }
+
+  private def pollUntilReadyOrDisconnected(selector: Selector, nodeId: String): Boolean = {
+    TestUtils.waitUntilTrue(() => {
+      selector.poll(100)
+      val disconnectState = selector.disconnected().get(nodeId)
+      // Verify that disconnect state is not AUTHENTICATION_FAILED
+      if (disconnectState != null) {
+        assertEquals(s"Authentication failed with exception ${disconnectState.exception()}",
+          ChannelState.State.AUTHENTICATE, disconnectState.state())
+      }
+      selector.isChannelReady(nodeId) || disconnectState != null
+    }, "Client not ready or disconnected within timeout")
+    val isReady = selector.isChannelReady(nodeId)
+    selector.close(nodeId)
+    isReady
+  }
+
+  /**
+   * Verifies that authentication with the current `clientConfig` results in disconnection and that
+   * the disconnection is notified with disconnect state `AUTHENTICATION_FAILED`. This is to ensure
+   * that NetworkClient handles this as a fatal authentication failure that is propagated to
+   * applications by Producer/Consumer/AdminClient without retrying and waiting for timeout.
+   */
+  private def verifyNonRetriableAuthenticationFailure(): Unit = {
+    val selector = createSelector()
+    val nodeId = "1"
+    selector.connect(nodeId, serverAddr, 1024, 1024)
+    TestUtils.waitUntilTrue(() => {
+      selector.poll(100)
+      val disconnectState = selector.disconnected().get(nodeId)
+      if (disconnectState != null)
+        assertEquals(ChannelState.State.AUTHENTICATION_FAILED, disconnectState.state())
+      disconnectState != null
+    }, "Client not disconnected within timeout")
+  }
+
+  private def createSelector(): Selector = {
+    val channelBuilder = ChannelBuilders.clientChannelBuilder(securityProtocol,
+      JaasContext.Type.CLIENT, new TestSecurityConfig(clientConfig), null, kafkaClientSaslMechanism,
+      time, true, new LogContext())
+    NetworkTestUtils.createSelector(channelBuilder, time)
+  }
+
+  private def createSelectorWithRelogin(): Selector = {
+    clientConfig.setProperty(SaslConfigs.SASL_KERBEROS_MIN_TIME_BEFORE_RELOGIN, "0")
+    val config = new TestSecurityConfig(clientConfig)
+    val jaasContexts = Collections.singletonMap("GSSAPI", JaasContext.loadClientContext(config.values()))
+    val channelBuilder = new SaslChannelBuilder(Mode.CLIENT, jaasContexts, securityProtocol,
+      null, false, kafkaClientSaslMechanism, true, null, null, time, new LogContext()) {
+      override protected def defaultLoginClass(): Class[_ <: Login] = classOf[TestableKerberosLogin]
+    }
+    channelBuilder.configure(config.values())
+    NetworkTestUtils.createSelector(channelBuilder, time)
+  }
+}
+
+object TestableKerberosLogin {
+  @volatile var instance: TestableKerberosLogin = _
+  def reset(): Unit = {
+    instance = null
+  }
+}
+
+class TestableKerberosLogin extends KerberosLogin {
+  val logoutResumeLatch = new CountDownLatch(1)
+  val logoutCompleteLatch = new CountDownLatch(1)
+  val reLoginResumeLatch = new CountDownLatch(1)
+  val reLoginCompleteLatch = new CountDownLatch(1)
+
+  assertNull(TestableKerberosLogin.instance)
+  TestableKerberosLogin.instance = this
+
+  override def reLogin(): Unit = {
+    super.reLogin()
+    reLoginCompleteLatch.countDown()
+  }
+
+  override protected def logout(): Unit = {
+    logoutResumeLatch.await(15, TimeUnit.SECONDS)
+    super.logout()
+    logoutCompleteLatch.countDown()
+    reLoginResumeLatch.await(15, TimeUnit.SECONDS)
+  }
+}
diff --git a/core/src/test/scala/integration/kafka/server/MultipleListenersWithAdditionalJaasContextTest.scala b/core/src/test/scala/integration/kafka/server/MultipleListenersWithAdditionalJaasContextTest.scala
index 666ba4b92a63d..39aa3c3639ca8 100644
--- a/core/src/test/scala/integration/kafka/server/MultipleListenersWithAdditionalJaasContextTest.scala
+++ b/core/src/test/scala/integration/kafka/server/MultipleListenersWithAdditionalJaasContextTest.scala
@@ -19,27 +19,27 @@ package kafka.server
 
 import java.util.Properties
 
+import scala.collection.Seq
+
 import kafka.utils.JaasTestUtils
 import kafka.utils.JaasTestUtils.JaasSection
-import org.apache.kafka.common.network.ListenerName
 
 class MultipleListenersWithAdditionalJaasContextTest extends MultipleListenersWithSameSecurityProtocolBaseTest {
 
   import MultipleListenersWithSameSecurityProtocolBaseTest._
 
-  override def saslProperties(listenerName: ListenerName): Properties = {
-    listenerName.value match {
-      case SecureInternal => kafkaClientSaslProperties(Plain, dynamicJaasConfig = true)
-      case SecureExternal => kafkaClientSaslProperties(GssApi, dynamicJaasConfig = true)
-      case _ => throw new IllegalArgumentException(s"Unexpected listener name $listenerName")
-    }
+  override def staticJaasSections: Seq[JaasSection] = {
+    val (serverKeytabFile, _) = maybeCreateEmptyKeytabFiles()
+    JaasTestUtils.zkSections :+
+      JaasTestUtils.kafkaServerSection("secure_external.KafkaServer", kafkaServerSaslMechanisms(SecureExternal), Some(serverKeytabFile))
   }
 
-  override def jaasSections: Seq[JaasSection] = {
-    val (serverKeytabFile, _) = maybeCreateEmptyKeytabFiles()
-    JaasTestUtils.zkSections ++ Seq(
-      JaasTestUtils.kafkaServerSection("secure_external.KafkaServer", Seq(GssApi), Some(serverKeytabFile)),
-      JaasTestUtils.kafkaServerSection("secure_internal.KafkaServer", Seq(Plain), None)
-    )
+  override protected def dynamicJaasSections: Properties = {
+    val props = new Properties
+    kafkaServerSaslMechanisms(SecureInternal).foreach { mechanism =>
+      addDynamicJaasSection(props, SecureInternal, mechanism,
+        JaasTestUtils.kafkaServerSection("secure_internal.KafkaServer", Seq(mechanism), None))
+    }
+    props
   }
 }
diff --git a/core/src/test/scala/integration/kafka/server/MultipleListenersWithDefaultJaasContextTest.scala b/core/src/test/scala/integration/kafka/server/MultipleListenersWithDefaultJaasContextTest.scala
index f3e1b8b9fcb72..23746c495579b 100644
--- a/core/src/test/scala/integration/kafka/server/MultipleListenersWithDefaultJaasContextTest.scala
+++ b/core/src/test/scala/integration/kafka/server/MultipleListenersWithDefaultJaasContextTest.scala
@@ -19,19 +19,16 @@ package kafka.server
 
 import java.util.Properties
 
+import scala.collection.Seq
+
 import kafka.api.Both
 import kafka.utils.JaasTestUtils.JaasSection
-import org.apache.kafka.common.network.ListenerName
-
 
 class MultipleListenersWithDefaultJaasContextTest extends MultipleListenersWithSameSecurityProtocolBaseTest {
 
-  import MultipleListenersWithSameSecurityProtocolBaseTest._
-
-  override def saslProperties(listenerName: ListenerName): Properties =
-    kafkaClientSaslProperties(Plain, dynamicJaasConfig = true)
+  override def staticJaasSections: Seq[JaasSection] =
+    jaasSections(kafkaServerSaslMechanisms.values.flatMap(identity).toSeq, Some(kafkaClientSaslMechanism), Both)
 
-  override def jaasSections: Seq[JaasSection] =
-    jaasSections(kafkaServerSaslMechanisms, Some(kafkaClientSaslMechanism), Both)
+  override protected def dynamicJaasSections: Properties = new Properties
 
 }
diff --git a/core/src/test/scala/integration/kafka/server/MultipleListenersWithSameSecurityProtocolBaseTest.scala b/core/src/test/scala/integration/kafka/server/MultipleListenersWithSameSecurityProtocolBaseTest.scala
index 40ec29343ec1d..6e811886ac40a 100644
--- a/core/src/test/scala/integration/kafka/server/MultipleListenersWithSameSecurityProtocolBaseTest.scala
+++ b/core/src/test/scala/integration/kafka/server/MultipleListenersWithSameSecurityProtocolBaseTest.scala
@@ -19,16 +19,16 @@
 package kafka.server
 
 import java.io.File
-import java.util.{Collections, Properties}
+import java.util.{Collections, Objects, Properties}
 import java.util.concurrent.TimeUnit
 
 import kafka.api.SaslSetup
 import kafka.coordinator.group.OffsetConfig
 import kafka.utils.JaasTestUtils.JaasSection
-import kafka.utils.TestUtils
+import kafka.utils.{JaasTestUtils, TestUtils}
 import kafka.utils.Implicits._
 import kafka.zk.ZooKeeperTestHarness
-import org.apache.kafka.clients.consumer.{ConsumerRecord, KafkaConsumer}
+import org.apache.kafka.clients.consumer.KafkaConsumer
 import org.apache.kafka.clients.producer.{KafkaProducer, ProducerRecord}
 import org.apache.kafka.common.config.SslConfigs
 import org.apache.kafka.common.internals.Topic
@@ -38,7 +38,7 @@ import org.junit.{After, Before, Test}
 
 import scala.collection.mutable
 import scala.collection.mutable.ArrayBuffer
-import scala.collection.JavaConverters._
+import scala.collection.Seq
 
 object MultipleListenersWithSameSecurityProtocolBaseTest {
   val SecureInternal = "SECURE_INTERNAL"
@@ -55,18 +55,20 @@ abstract class MultipleListenersWithSameSecurityProtocolBaseTest extends ZooKeep
 
   private val trustStoreFile = File.createTempFile("truststore", ".jks")
   private val servers = new ArrayBuffer[KafkaServer]
-  private val producers = mutable.Map[ListenerName, KafkaProducer[Array[Byte], Array[Byte]]]()
-  private val consumers = mutable.Map[ListenerName, KafkaConsumer[Array[Byte], Array[Byte]]]()
+  private val producers = mutable.Map[ClientMetadata, KafkaProducer[Array[Byte], Array[Byte]]]()
+  private val consumers = mutable.Map[ClientMetadata, KafkaConsumer[Array[Byte], Array[Byte]]]()
 
   protected val kafkaClientSaslMechanism = Plain
-  protected val kafkaServerSaslMechanisms = List(GssApi, Plain)
+  protected val kafkaServerSaslMechanisms = Map(
+    SecureExternal -> Seq("SCRAM-SHA-256", GssApi),
+    SecureInternal -> Seq(Plain, "SCRAM-SHA-512"))
 
-  protected def saslProperties(listenerName: ListenerName): Properties
-  protected def jaasSections: Seq[JaasSection]
+  protected def staticJaasSections: Seq[JaasSection]
+  protected def dynamicJaasSections: Properties
 
   @Before
   override def setUp(): Unit = {
-    startSasl(jaasSections)
+    startSasl(staticJaasSections)
     super.setUp()
     // 2 brokers so that we can test that the data propagates correctly via UpdateMetadadaRequest
     val numServers = 2
@@ -82,8 +84,12 @@ abstract class MultipleListenersWithSameSecurityProtocolBaseTest extends ZooKeep
       props.put(KafkaConfig.InterBrokerListenerNameProp, Internal)
       props.put(KafkaConfig.ZkEnableSecureAclsProp, "true")
       props.put(KafkaConfig.SaslMechanismInterBrokerProtocolProp, kafkaClientSaslMechanism)
-      props.put(KafkaConfig.SaslEnabledMechanismsProp, kafkaServerSaslMechanisms.mkString(","))
+      props.put(s"${new ListenerName(SecureInternal).configPrefix}${KafkaConfig.SaslEnabledMechanismsProp}",
+        kafkaServerSaslMechanisms(SecureInternal).mkString(","))
+      props.put(s"${new ListenerName(SecureExternal).configPrefix}${KafkaConfig.SaslEnabledMechanismsProp}",
+        kafkaServerSaslMechanisms(SecureExternal).mkString(","))
       props.put(KafkaConfig.SaslKerberosServiceNameProp, "kafka")
+      props ++= dynamicJaasSections
 
       props ++= TestUtils.sslConfigs(Mode.SERVER, false, Some(trustStoreFile), s"server$brokerId")
 
@@ -104,34 +110,45 @@ abstract class MultipleListenersWithSameSecurityProtocolBaseTest extends ZooKeep
         Internal, config.interBrokerListenerName.value)
     }
 
-    TestUtils.createTopic(zkUtils, Topic.GROUP_METADATA_TOPIC_NAME, OffsetConfig.DefaultOffsetsTopicNumPartitions,
+    TestUtils.createTopic(zkClient, Topic.GROUP_METADATA_TOPIC_NAME, OffsetConfig.DefaultOffsetsTopicNumPartitions,
       replicationFactor = 2, servers, servers.head.groupCoordinator.offsetsTopicConfigs)
 
+    createScramCredentials(zkConnect, JaasTestUtils.KafkaScramUser, JaasTestUtils.KafkaScramPassword)
+
     servers.head.config.listeners.foreach { endPoint =>
       val listenerName = endPoint.listenerName
 
-      TestUtils.createTopic(zkUtils, listenerName.value, 2, 2, servers)
-
       val trustStoreFile =
         if (TestUtils.usesSslTransportLayer(endPoint.securityProtocol)) Some(this.trustStoreFile)
         else None
 
-      val saslProps =
-        if (TestUtils.usesSaslAuthentication(endPoint.securityProtocol)) Some(saslProperties(listenerName))
-        else None
-
       val bootstrapServers = TestUtils.bootstrapServers(servers, listenerName)
 
-      producers(listenerName) = TestUtils.createNewProducer(bootstrapServers, acks = -1,
-        securityProtocol = endPoint.securityProtocol, trustStoreFile = trustStoreFile, saslProperties = saslProps)
+      def addProducerConsumer(listenerName: ListenerName, mechanism: String, saslProps: Option[Properties]): Unit = {
+
+        val topic = s"${listenerName.value}${producers.size}"
+        TestUtils.createTopic(zkClient, topic, 2, 2, servers)
+        val clientMetadata = ClientMetadata(listenerName, mechanism, topic)
+
+        producers(clientMetadata) = TestUtils.createProducer(bootstrapServers, acks = -1,
+          securityProtocol = endPoint.securityProtocol, trustStoreFile = trustStoreFile, saslProperties = saslProps)
+
+        consumers(clientMetadata) = TestUtils.createConsumer(bootstrapServers, groupId = clientMetadata.toString,
+          securityProtocol = endPoint.securityProtocol, trustStoreFile = trustStoreFile, saslProperties = saslProps)
+      }
 
-      consumers(listenerName) = TestUtils.createNewConsumer(bootstrapServers, groupId = listenerName.value,
-        securityProtocol = endPoint.securityProtocol, trustStoreFile = trustStoreFile, saslProperties = saslProps)
+      if (TestUtils.usesSaslAuthentication(endPoint.securityProtocol)) {
+        kafkaServerSaslMechanisms(endPoint.listenerName.value).foreach { mechanism =>
+          addProducerConsumer(listenerName, mechanism, Some(kafkaClientSaslProperties(mechanism, dynamicJaasConfig = true)))
+        }
+      } else {
+        addProducerConsumer(listenerName, "", saslProps = None)
+      }
     }
   }
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     producers.values.foreach(_.close())
     consumers.values.foreach(_.close())
     TestUtils.shutdownServers(servers)
@@ -145,18 +162,30 @@ abstract class MultipleListenersWithSameSecurityProtocolBaseTest extends ZooKeep
     */
   @Test
   def testProduceConsume(): Unit = {
-    producers.foreach { case (listenerName, producer) =>
-      val producerRecords = (1 to 10).map(i => new ProducerRecord(listenerName.value, s"key$i".getBytes,
+    producers.foreach { case (clientMetadata, producer) =>
+      val producerRecords = (1 to 10).map(i => new ProducerRecord(clientMetadata.topic, s"key$i".getBytes,
         s"value$i".getBytes))
       producerRecords.map(producer.send).map(_.get(10, TimeUnit.SECONDS))
 
-      val consumer = consumers(listenerName)
-      consumer.subscribe(Collections.singleton(listenerName.value))
-      val records = new ArrayBuffer[ConsumerRecord[Array[Byte], Array[Byte]]]
-      TestUtils.waitUntilTrue(() => {
-        records ++= consumer.poll(50).asScala
-        records.size == producerRecords.size
-      }, s"Consumed ${records.size} records until timeout instead of the expected ${producerRecords.size} records")
+      val consumer = consumers(clientMetadata)
+      consumer.subscribe(Collections.singleton(clientMetadata.topic))
+      TestUtils.consumeRecords(consumer, producerRecords.size)
+    }
+  }
+
+  protected def addDynamicJaasSection(props: Properties, listener: String, mechanism: String, jaasSection: JaasSection): Unit = {
+    val listenerName = new ListenerName(listener)
+    val prefix = listenerName.saslMechanismConfigPrefix(mechanism)
+    val jaasConfig = jaasSection.modules.head.toString
+    props.put(s"${prefix}${KafkaConfig.SaslJaasConfigProp}", jaasConfig)
+  }
+
+  case class ClientMetadata(val listenerName: ListenerName, val saslMechanism: String, topic: String) {
+    override def hashCode: Int = Objects.hash(listenerName, saslMechanism)
+    override def equals(obj: Any): Boolean = obj match {
+      case other: ClientMetadata => listenerName == other.listenerName && saslMechanism == other.saslMechanism && topic == other.topic
+      case _ => false
     }
+    override def toString: String = s"${listenerName.value}:$saslMechanism:$topic"
   }
 }
diff --git a/core/src/test/scala/integration/kafka/server/ReplicaFetcherThreadFatalErrorTest.scala b/core/src/test/scala/integration/kafka/server/ReplicaFetcherThreadFatalErrorTest.scala
deleted file mode 100644
index 46badf0b625a8..0000000000000
--- a/core/src/test/scala/integration/kafka/server/ReplicaFetcherThreadFatalErrorTest.scala
+++ /dev/null
@@ -1,145 +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.util.concurrent.atomic.AtomicBoolean
-
-import kafka.cluster.BrokerEndPoint
-import kafka.server.ReplicaFetcherThread.{FetchRequest, PartitionData}
-import kafka.utils.{Exit, TestUtils, ZkUtils}
-import kafka.utils.TestUtils.createBrokerConfigs
-import kafka.zk.ZooKeeperTestHarness
-import org.apache.kafka.common.TopicPartition
-import org.apache.kafka.common.internals.FatalExitError
-import org.apache.kafka.common.metrics.Metrics
-import org.apache.kafka.common.protocol.Errors
-import org.apache.kafka.common.requests.FetchResponse
-import org.apache.kafka.common.utils.Time
-import org.junit.{After, Test}
-
-import scala.collection.Map
-import scala.collection.JavaConverters._
-import scala.concurrent.Future
-
-class ReplicaFetcherThreadFatalErrorTest extends ZooKeeperTestHarness {
-
-  private var brokers: Seq[KafkaServer] = null
-  @volatile private var shutdownCompleted = false
-
-  @After
-  override def tearDown() {
-    Exit.resetExitProcedure()
-    TestUtils.shutdownServers(brokers)
-    super.tearDown()
-  }
-
-  /**
-    * Verifies that a follower shuts down if the offset for an `added partition` is out of range and if a fatal
-    * exception is thrown from `handleOffsetOutOfRange`. It's a bit tricky to ensure that there are no deadlocks
-    * when the shutdown hook is invoked and hence this test.
-    */
-  @Test
-  def testFatalErrorInAddPartitions(): Unit = {
-
-    // Unlike `TestUtils.createTopic`, this doesn't wait for metadata propagation as the broker shuts down before
-    // the metadata is propagated.
-    def createTopic(zkUtils: ZkUtils, topic: String): Unit = {
-      adminZkClient.createTopic(topic, partitions = 1, replicationFactor = 2)
-      TestUtils.waitUntilLeaderIsElectedOrChanged(zkUtils, topic, 0)
-    }
-
-    val props = createBrokerConfigs(2, zkConnect)
-    brokers = props.map(KafkaConfig.fromProps).map(config => createServer(config, { params =>
-      import params._
-      new ReplicaFetcherThread(threadName, fetcherId, sourceBroker, config, replicaManager, metrics, time, quotaManager) {
-        override def handleOffsetOutOfRange(topicPartition: TopicPartition): Long = throw new FatalExitError
-        override def addPartitions(partitionAndOffsets: Map[TopicPartition, Long]): Unit =
-          super.addPartitions(partitionAndOffsets.mapValues(_ => -1))
-      }
-    }))
-    createTopic(zkUtils, "topic")
-    TestUtils.waitUntilTrue(() => shutdownCompleted, "Shutdown of follower did not complete")
-  }
-
-  /**
-    * Verifies that a follower shuts down if the offset of a partition in the fetch response is out of range and if a
-    * fatal exception is thrown from `handleOffsetOutOfRange`. It's a bit tricky to ensure that there are no deadlocks
-    * when the shutdown hook is invoked and hence this test.
-    */
-  @Test
-  def testFatalErrorInProcessFetchRequest(): Unit = {
-    val props = createBrokerConfigs(2, zkConnect)
-    brokers = props.map(KafkaConfig.fromProps).map(config => createServer(config, { params =>
-      import params._
-      new ReplicaFetcherThread(threadName, fetcherId, sourceBroker, config, replicaManager, metrics, time, quotaManager) {
-        override def handleOffsetOutOfRange(topicPartition: TopicPartition): Long = throw new FatalExitError
-        override protected def fetch(fetchRequest: FetchRequest): Seq[(TopicPartition, PartitionData)] = {
-          fetchRequest.underlying.fetchData.asScala.keys.toSeq.map { tp =>
-            (tp, new PartitionData(new FetchResponse.PartitionData(Errors.OFFSET_OUT_OF_RANGE,
-              FetchResponse.INVALID_HIGHWATERMARK, FetchResponse.INVALID_LAST_STABLE_OFFSET, FetchResponse.INVALID_LOG_START_OFFSET, null, null)))
-          }
-        }
-      }
-    }))
-    TestUtils.createTopic(zkUtils, "topic", numPartitions = 1, replicationFactor = 2, servers = brokers)
-    TestUtils.waitUntilTrue(() => shutdownCompleted, "Shutdown of follower did not complete")
-  }
-
-  private case class FetcherThreadParams(threadName: String, fetcherId: Int, sourceBroker: BrokerEndPoint,
-                                         replicaManager: ReplicaManager, metrics: Metrics, time: Time,
-                                         quotaManager: ReplicationQuotaManager)
-
-  private def createServer(config: KafkaConfig, fetcherThread: FetcherThreadParams => ReplicaFetcherThread): KafkaServer = {
-    val time = Time.SYSTEM
-    val server = new KafkaServer(config, time) {
-
-      override def createReplicaManager(isShuttingDown: AtomicBoolean): ReplicaManager = {
-        new ReplicaManager(config, metrics, time, zkClient, kafkaScheduler, logManager, isShuttingDown,
-          quotaManagers, new BrokerTopicStats, metadataCache, logDirFailureChannel) {
-
-          override protected def createReplicaFetcherManager(metrics: Metrics, time: Time, threadNamePrefix: Option[String],
-                                                             quotaManager: ReplicationQuotaManager) =
-            new ReplicaFetcherManager(config, this, metrics, time, threadNamePrefix, quotaManager) {
-              override def createFetcherThread(fetcherId: Int, sourceBroker: BrokerEndPoint): AbstractFetcherThread = {
-                val prefix = threadNamePrefix.map(tp => s"$tp:").getOrElse("")
-                val threadName = s"${prefix}ReplicaFetcherThread-$fetcherId-${sourceBroker.id}"
-                fetcherThread(FetcherThreadParams(threadName, fetcherId, sourceBroker, replicaManager, metrics,
-                  time, quotaManager))
-              }
-            }
-        }
-      }
-
-    }
-
-    Exit.setExitProcedure { (_, _) =>
-      import scala.concurrent.ExecutionContext.Implicits._
-      // Run in a separate thread like shutdown hooks
-      Future {
-        server.shutdown()
-        shutdownCompleted = true
-      }
-      // Sleep until interrupted to emulate the fact that `System.exit()` never returns
-      Thread.sleep(Long.MaxValue)
-      throw new AssertionError
-    }
-    server.startup()
-    server
-  }
-
-}
diff --git a/core/src/test/scala/integration/kafka/server/ScramServerStartupTest.scala b/core/src/test/scala/integration/kafka/server/ScramServerStartupTest.scala
new file mode 100644
index 0000000000000..15bbd54d0cfaa
--- /dev/null
+++ b/core/src/test/scala/integration/kafka/server/ScramServerStartupTest.scala
@@ -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 kafka.server
+
+import java.util.Collections
+
+import kafka.api.{IntegrationTestHarness, KafkaSasl, SaslSetup}
+import kafka.utils._
+import kafka.zk.ConfigEntityChangeNotificationZNode
+import org.apache.kafka.common.security.auth.SecurityProtocol
+import org.junit.Assert._
+import org.junit.Test
+
+import scala.jdk.CollectionConverters._
+
+/**
+ * Tests that there are no failed authentications during broker startup. This is to verify
+ * that SCRAM credentials are loaded by brokers before client connections can be made.
+ * For simplicity of testing, this test verifies authentications of controller connections.
+ */
+class ScramServerStartupTest extends IntegrationTestHarness with SaslSetup {
+
+  override val brokerCount = 1
+
+  private val kafkaClientSaslMechanism = "SCRAM-SHA-256"
+  private val kafkaServerSaslMechanisms = Collections.singletonList("SCRAM-SHA-256").asScala
+
+  override protected def securityProtocol = SecurityProtocol.SASL_PLAINTEXT
+
+  override protected val serverSaslProperties = Some(kafkaServerSaslProperties(kafkaServerSaslMechanisms, kafkaClientSaslMechanism))
+  override protected val clientSaslProperties = Some(kafkaClientSaslProperties(kafkaClientSaslMechanism))
+
+  override def configureSecurityBeforeServersStart(): Unit = {
+    super.configureSecurityBeforeServersStart()
+    zkClient.makeSurePersistentPathExists(ConfigEntityChangeNotificationZNode.path)
+    // Create credentials before starting brokers
+    createScramCredentials(zkConnect, JaasTestUtils.KafkaScramAdmin, JaasTestUtils.KafkaScramAdminPassword)
+
+    startSasl(jaasSections(kafkaServerSaslMechanisms, Option(kafkaClientSaslMechanism), KafkaSasl))
+  }
+
+  @Test
+  def testAuthentications(): Unit = {
+    val successfulAuths = TestUtils.totalMetricValue(servers.head, "successful-authentication-total")
+    assertTrue("No successful authentications", successfulAuths > 0)
+    val failedAuths = TestUtils.totalMetricValue(servers.head, "failed-authentication-total")
+    assertEquals(0, failedAuths)
+  }
+}
diff --git a/core/src/test/scala/integration/kafka/tools/MirrorMakerIntegrationTest.scala b/core/src/test/scala/integration/kafka/tools/MirrorMakerIntegrationTest.scala
index 1f9851dadef56..2e32f0a799534 100644
--- a/core/src/test/scala/integration/kafka/tools/MirrorMakerIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/tools/MirrorMakerIntegrationTest.scala
@@ -17,22 +17,75 @@
 package kafka.tools
 
 import java.util.Properties
+import java.util.concurrent.atomic.AtomicBoolean
 
-import kafka.consumer.ConsumerTimeoutException
+import scala.collection.Seq
 import kafka.integration.KafkaServerTestHarness
 import kafka.server.KafkaConfig
-import kafka.tools.MirrorMaker.{MirrorMakerNewConsumer, MirrorMakerProducer}
+import kafka.tools.MirrorMaker.{ConsumerWrapper, MirrorMakerProducer, NoRecordsException}
 import kafka.utils.TestUtils
 import org.apache.kafka.clients.consumer.{ConsumerConfig, KafkaConsumer}
 import org.apache.kafka.clients.producer.{ProducerConfig, ProducerRecord}
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.errors.TimeoutException
 import org.apache.kafka.common.serialization.{ByteArrayDeserializer, ByteArraySerializer}
+import org.apache.kafka.common.utils.Exit
+import org.junit.After
 import org.junit.Test
+import org.junit.Assert._
+import org.junit.Before
 
 class MirrorMakerIntegrationTest extends KafkaServerTestHarness {
 
   override def generateConfigs: Seq[KafkaConfig] =
     TestUtils.createBrokerConfigs(1, zkConnect).map(KafkaConfig.fromProps(_, new Properties()))
 
+  val exited = new AtomicBoolean(false)
+
+  @Before
+  override def setUp(): Unit = {
+    Exit.setExitProcedure((_, _) => exited.set(true))
+    super.setUp()
+  }
+
+  @After
+  override def tearDown(): Unit = {
+    super.tearDown()
+    try {
+      assertFalse(exited.get())
+    } finally {
+      Exit.resetExitProcedure()
+    }
+  }
+
+  @Test(expected = classOf[TimeoutException])
+  def testCommitOffsetsThrowTimeoutException(): Unit = {
+    val consumerProps = new Properties
+    consumerProps.put(ConsumerConfig.GROUP_ID_CONFIG, "test-group")
+    consumerProps.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest")
+    consumerProps.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList)
+    consumerProps.put(ConsumerConfig.DEFAULT_API_TIMEOUT_MS_CONFIG, "1")
+    val consumer = new KafkaConsumer(consumerProps, new ByteArrayDeserializer, new ByteArrayDeserializer)
+    val mirrorMakerConsumer = new ConsumerWrapper(consumer, None, whitelistOpt = Some("any"))
+    mirrorMakerConsumer.offsets.put(new TopicPartition("test", 0), 0L)
+    mirrorMakerConsumer.commit()
+  }
+
+  @Test
+  def testCommitOffsetsRemoveNonExistentTopics(): Unit = {
+    val consumerProps = new Properties
+    consumerProps.put(ConsumerConfig.GROUP_ID_CONFIG, "test-group")
+    consumerProps.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest")
+    consumerProps.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList)
+    consumerProps.put(ConsumerConfig.DEFAULT_API_TIMEOUT_MS_CONFIG, "2000")
+    val consumer = new KafkaConsumer(consumerProps, new ByteArrayDeserializer, new ByteArrayDeserializer)
+    val mirrorMakerConsumer = new ConsumerWrapper(consumer, None, whitelistOpt = Some("any"))
+    mirrorMakerConsumer.offsets.put(new TopicPartition("nonexistent-topic1", 0), 0L)
+    mirrorMakerConsumer.offsets.put(new TopicPartition("nonexistent-topic2", 0), 0L)
+    MirrorMaker.commitOffsets(mirrorMakerConsumer)
+    assertTrue("Offsets for non-existent topics should be removed", mirrorMakerConsumer.offsets.isEmpty)
+  }
+
   @Test
   def testCommaSeparatedRegex(): Unit = {
     val topic = "new-topic"
@@ -54,7 +107,7 @@ class MirrorMakerIntegrationTest extends KafkaServerTestHarness {
     consumerProps.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList)
     val consumer = new KafkaConsumer(consumerProps, new ByteArrayDeserializer, new ByteArrayDeserializer)
 
-    val mirrorMakerConsumer = new MirrorMakerNewConsumer(consumer, None, whitelistOpt = Some("another_topic,new.*,foo"))
+    val mirrorMakerConsumer = new ConsumerWrapper(consumer, None, whitelistOpt = Some("another_topic,new.*,foo"))
     mirrorMakerConsumer.init()
     try {
       TestUtils.waitUntilTrue(() => {
@@ -62,8 +115,8 @@ class MirrorMakerIntegrationTest extends KafkaServerTestHarness {
           val data = mirrorMakerConsumer.receive()
           data.topic == topic && new String(data.value) == msg
         } catch {
-          // this exception is thrown if no record is returned within a short timeout, so safe to ignore
-          case _: ConsumerTimeoutException => false
+          // these exceptions are thrown if no records are returned within the timeout, so safe to ignore
+          case _: NoRecordsException => false
         }
       }, "MirrorMaker consumer should read the expected message from the expected topic within the timeout")
     } finally consumer.close()
diff --git a/core/src/test/scala/kafka/common/InterBrokerSendThreadTest.scala b/core/src/test/scala/kafka/common/InterBrokerSendThreadTest.scala
index c6ebdd17c3612..f5110bf0b03e0 100644
--- a/core/src/test/scala/kafka/common/InterBrokerSendThreadTest.scala
+++ b/core/src/test/scala/kafka/common/InterBrokerSendThreadTest.scala
@@ -16,37 +16,42 @@
  */
 package kafka.common
 
-import org.junit.{Assert, Test}
+import java.util
+
 import kafka.utils.MockTime
 import org.apache.kafka.clients.{ClientRequest, ClientResponse, NetworkClient, RequestCompletionHandler}
 import org.apache.kafka.common.Node
+import org.apache.kafka.common.errors.AuthenticationException
 import org.apache.kafka.common.protocol.ApiKeys
 import org.apache.kafka.common.requests.AbstractRequest
-import org.apache.kafka.common.utils.Utils
 import org.easymock.EasyMock
+import org.junit.{Assert, Test}
 
 import scala.collection.mutable
 
 class InterBrokerSendThreadTest {
   private val time = new MockTime()
-  private val networkClient = EasyMock.createMock(classOf[NetworkClient])
+  private val networkClient: NetworkClient = EasyMock.createMock(classOf[NetworkClient])
   private val completionHandler = new StubCompletionHandler
+  private val requestTimeoutMs = 1000
 
   @Test
   def shouldNotSendAnythingWhenNoRequests(): Unit = {
     val sendThread = new InterBrokerSendThread("name", networkClient, time) {
+      override val requestTimeoutMs: Int = InterBrokerSendThreadTest.this.requestTimeoutMs
       override def generateRequests() = mutable.Iterable.empty
     }
 
     // poll is always called but there should be no further invocations on NetworkClient
     EasyMock.expect(networkClient.poll(EasyMock.anyLong(), EasyMock.anyLong()))
-    .andReturn(Utils.mkList())
+      .andReturn(new util.ArrayList())
 
     EasyMock.replay(networkClient)
 
     sendThread.doWork()
 
     EasyMock.verify(networkClient)
+    Assert.assertFalse(completionHandler.executedWithDisconnectedResponse)
   }
 
   @Test
@@ -55,49 +60,55 @@ class InterBrokerSendThreadTest {
     val node = new Node(1, "", 8080)
     val handler = RequestAndCompletionHandler(node, request, completionHandler)
     val sendThread = new InterBrokerSendThread("name", networkClient, time) {
+      override val requestTimeoutMs: Int = InterBrokerSendThreadTest.this.requestTimeoutMs
       override def generateRequests() = List[RequestAndCompletionHandler](handler)
     }
 
-    val clientRequest = new ClientRequest("dest", request, 0, "1", 0, true, handler.handler)
+    val clientRequest = new ClientRequest("dest", request, 0, "1", 0, true, requestTimeoutMs, handler.handler)
 
-    EasyMock.expect(networkClient.newClientRequest(EasyMock.eq("1"),
+    EasyMock.expect(networkClient.newClientRequest(
+      EasyMock.eq("1"),
       EasyMock.same(handler.request),
       EasyMock.anyLong(),
       EasyMock.eq(true),
+      EasyMock.eq(requestTimeoutMs),
       EasyMock.same(handler.handler)))
-    .andReturn(clientRequest)
+      .andReturn(clientRequest)
 
     EasyMock.expect(networkClient.ready(node, time.milliseconds()))
-    .andReturn(true)
+      .andReturn(true)
 
     EasyMock.expect(networkClient.send(clientRequest, time.milliseconds()))
 
     EasyMock.expect(networkClient.poll(EasyMock.anyLong(), EasyMock.anyLong()))
-      .andReturn(Utils.mkList())
+      .andReturn(new util.ArrayList())
 
     EasyMock.replay(networkClient)
 
     sendThread.doWork()
 
     EasyMock.verify(networkClient)
+    Assert.assertFalse(completionHandler.executedWithDisconnectedResponse)
   }
 
-
   @Test
   def shouldCallCompletionHandlerWithDisconnectedResponseWhenNodeNotReady(): Unit = {
     val request = new StubRequestBuilder
     val node = new Node(1, "", 8080)
     val requestAndCompletionHandler = RequestAndCompletionHandler(node, request, completionHandler)
     val sendThread = new InterBrokerSendThread("name", networkClient, time) {
+      override val requestTimeoutMs: Int = InterBrokerSendThreadTest.this.requestTimeoutMs
       override def generateRequests() = List[RequestAndCompletionHandler](requestAndCompletionHandler)
     }
 
-    val clientRequest = new ClientRequest("dest", request, 0, "1", 0, true, requestAndCompletionHandler.handler)
+    val clientRequest = new ClientRequest("dest", request, 0, "1", 0, true, requestTimeoutMs, requestAndCompletionHandler.handler)
 
-    EasyMock.expect(networkClient.newClientRequest(EasyMock.eq("1"),
+    EasyMock.expect(networkClient.newClientRequest(
+      EasyMock.eq("1"),
       EasyMock.same(requestAndCompletionHandler.request),
       EasyMock.anyLong(),
       EasyMock.eq(true),
+      EasyMock.eq(requestTimeoutMs),
       EasyMock.same(requestAndCompletionHandler.handler)))
       .andReturn(clientRequest)
 
@@ -105,27 +116,86 @@ class InterBrokerSendThreadTest {
       .andReturn(false)
 
     EasyMock.expect(networkClient.connectionDelay(EasyMock.anyObject(), EasyMock.anyLong()))
-    .andReturn(0)
+      .andReturn(0)
 
     EasyMock.expect(networkClient.poll(EasyMock.anyLong(), EasyMock.anyLong()))
-      .andReturn(Utils.mkList())
+      .andReturn(new util.ArrayList())
+
+    EasyMock.expect(networkClient.connectionFailed(node))
+      .andReturn(true)
+
+    EasyMock.expect(networkClient.authenticationException(node))
+      .andReturn(new AuthenticationException(""))
 
     EasyMock.replay(networkClient)
 
     sendThread.doWork()
 
     EasyMock.verify(networkClient)
-    Assert.assertTrue(completionHandler.response.wasDisconnected())
+    Assert.assertTrue(completionHandler.executedWithDisconnectedResponse)
   }
 
+  @Test
+  def testFailingExpiredRequests(): Unit = {
+    val request = new StubRequestBuilder()
+    val node = new Node(1, "", 8080)
+    val handler = RequestAndCompletionHandler(node, request, completionHandler)
+    val sendThread = new InterBrokerSendThread("name", networkClient, time) {
+      override val requestTimeoutMs: Int = InterBrokerSendThreadTest.this.requestTimeoutMs
+      override def generateRequests() = List[RequestAndCompletionHandler](handler)
+    }
+
+    val clientRequest = new ClientRequest("dest",
+      request,
+      0,
+      "1",
+      time.milliseconds(),
+      true,
+      requestTimeoutMs,
+      handler.handler)
+    time.sleep(1500)
+
+    EasyMock.expect(networkClient.newClientRequest(
+      EasyMock.eq("1"),
+      EasyMock.same(handler.request),
+      EasyMock.eq(time.milliseconds()),
+      EasyMock.eq(true),
+      EasyMock.eq(requestTimeoutMs),
+      EasyMock.same(handler.handler)))
+      .andReturn(clientRequest)
+
+    // make the node unready so the request is not cleared
+    EasyMock.expect(networkClient.ready(node, time.milliseconds()))
+      .andReturn(false)
+
+    EasyMock.expect(networkClient.connectionDelay(EasyMock.anyObject(), EasyMock.anyLong()))
+      .andReturn(0)
+
+    EasyMock.expect(networkClient.poll(EasyMock.anyLong(), EasyMock.anyLong()))
+      .andReturn(new util.ArrayList())
+
+    // rule out disconnects so the request stays for the expiry check
+    EasyMock.expect(networkClient.connectionFailed(node))
+      .andReturn(false)
+
+    EasyMock.replay(networkClient)
+
+    sendThread.doWork()
+
+    EasyMock.verify(networkClient)
+    Assert.assertFalse(sendThread.hasUnsentRequests)
+    Assert.assertTrue(completionHandler.executedWithDisconnectedResponse)
+  }
 
   private class StubRequestBuilder extends AbstractRequest.Builder(ApiKeys.END_TXN) {
     override def build(version: Short): Nothing = ???
   }
 
   private class StubCompletionHandler extends RequestCompletionHandler {
+    var executedWithDisconnectedResponse = false
     var response: ClientResponse = _
     override def onComplete(response: ClientResponse): Unit = {
+      this.executedWithDisconnectedResponse = response.wasDisconnected()
       this.response = response
     }
   }
diff --git a/core/src/test/scala/kafka/metrics/LinuxIoMetricsCollectorTest.scala b/core/src/test/scala/kafka/metrics/LinuxIoMetricsCollectorTest.scala
new file mode 100644
index 0000000000000..3cf0284df97de
--- /dev/null
+++ b/core/src/test/scala/kafka/metrics/LinuxIoMetricsCollectorTest.scala
@@ -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 kafka.metrics
+
+import java.nio.charset.StandardCharsets
+import java.nio.file.Files
+
+import kafka.utils.{Logging, MockTime}
+import org.apache.kafka.test.TestUtils
+import org.junit.Assert.{assertEquals, assertFalse, assertTrue}
+import org.junit.{Rule, Test}
+import org.junit.rules.Timeout
+
+class LinuxIoMetricsCollectorTest extends Logging {
+  @Rule
+  def globalTimeout = Timeout.millis(120000)
+
+  class TestDirectory() {
+    val baseDir = TestUtils.tempDirectory()
+    val selfDir = Files.createDirectories(baseDir.toPath.resolve("self"))
+
+    def writeProcFile(readBytes: Long, writeBytes: Long) = {
+      val bld = new StringBuilder()
+      bld.append("rchar: 0%n".format())
+      bld.append("wchar: 0%n".format())
+      bld.append("syschr: 0%n".format())
+      bld.append("syscw: 0%n".format())
+      bld.append("read_bytes: %d%n".format(readBytes))
+      bld.append("write_bytes: %d%n".format(writeBytes))
+      bld.append("cancelled_write_bytes: 0%n".format())
+      Files.write(selfDir.resolve("io"), bld.toString().getBytes(StandardCharsets.UTF_8))
+    }
+  }
+
+  @Test
+  def testReadProcFile(): Unit = {
+    val testDirectory = new TestDirectory()
+    val time = new MockTime(100, 1000)
+    testDirectory.writeProcFile(123L, 456L)
+    val collector = new LinuxIoMetricsCollector(testDirectory.baseDir.getAbsolutePath,
+      time, logger.underlying)
+
+    // Test that we can read the values we wrote.
+    assertTrue(collector.usable())
+    assertEquals(123L, collector.readBytes())
+    assertEquals(456L, collector.writeBytes())
+    testDirectory.writeProcFile(124L, 457L)
+
+    // The previous values should still be cached.
+    assertEquals(123L, collector.readBytes())
+    assertEquals(456L, collector.writeBytes())
+
+    // Update the time, and the values should be re-read.
+    time.sleep(1)
+    assertEquals(124L, collector.readBytes())
+    assertEquals(457L, collector.writeBytes())
+  }
+
+  @Test
+  def testUnableToReadNonexistentProcFile(): Unit = {
+    val testDirectory = new TestDirectory()
+    val time = new MockTime(100, 1000)
+    val collector = new LinuxIoMetricsCollector(testDirectory.baseDir.getAbsolutePath,
+      time, logger.underlying)
+
+    // Test that we can't read the file, since it hasn't been written.
+    assertFalse(collector.usable())
+  }
+}
diff --git a/core/src/test/scala/kafka/security/auth/ResourceTest.scala b/core/src/test/scala/kafka/security/auth/ResourceTest.scala
new file mode 100644
index 0000000000000..deed61b9361f8
--- /dev/null
+++ b/core/src/test/scala/kafka/security/auth/ResourceTest.scala
@@ -0,0 +1,68 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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.auth
+
+import kafka.common.KafkaException
+import org.apache.kafka.common.resource.PatternType.{LITERAL, PREFIXED}
+import org.junit.Test
+import org.junit.Assert._
+
+@deprecated("Use org.apache.kafka.common.resource.ResourcePattern", "Since 2.5")
+class ResourceTest {
+  @Test(expected = classOf[KafkaException])
+  def shouldThrowOnTwoPartStringWithUnknownResourceType(): Unit = {
+    Resource.fromString("Unknown:fred")
+  }
+
+  @Test(expected = classOf[KafkaException])
+  def shouldThrowOnBadResourceTypeSeparator(): Unit = {
+    Resource.fromString("Topic-fred")
+  }
+
+  @Test
+  def shouldParseOldTwoPartString(): Unit = {
+    assertEquals(Resource(Group, "fred", LITERAL), Resource.fromString("Group:fred"))
+    assertEquals(Resource(Topic, "t", LITERAL), Resource.fromString("Topic:t"))
+  }
+
+  @Test
+  def shouldParseOldTwoPartWithEmbeddedSeparators(): Unit = {
+    assertEquals(Resource(Group, ":This:is:a:weird:group:name:", LITERAL), Resource.fromString("Group::This:is:a:weird:group:name:"))
+  }
+
+  @Test
+  def shouldParseThreePartString(): Unit = {
+    assertEquals(Resource(Group, "fred", PREFIXED), Resource.fromString("Group:PREFIXED:fred"))
+    assertEquals(Resource(Topic, "t", LITERAL), Resource.fromString("Topic:LITERAL:t"))
+  }
+
+  @Test
+  def shouldParseThreePartWithEmbeddedSeparators(): Unit = {
+    assertEquals(Resource(Group, ":This:is:a:weird:group:name:", PREFIXED), Resource.fromString("Group:PREFIXED::This:is:a:weird:group:name:"))
+    assertEquals(Resource(Group, ":This:is:a:weird:group:name:", LITERAL), Resource.fromString("Group:LITERAL::This:is:a:weird:group:name:"))
+  }
+
+  @Test
+  def shouldRoundTripViaString(): Unit = {
+    val expected = Resource(Group, "fred", PREFIXED)
+
+    val actual = Resource.fromString(expected.toString)
+
+    assertEquals(expected, actual)
+  }
+}
\ No newline at end of file
diff --git a/core/src/test/scala/kafka/security/minikdc/MiniKdc.scala b/core/src/test/scala/kafka/security/minikdc/MiniKdc.scala
index 9b3058122aaa2..ea8815ddcccaf 100644
--- a/core/src/test/scala/kafka/security/minikdc/MiniKdc.scala
+++ b/core/src/test/scala/kafka/security/minikdc/MiniKdc.scala
@@ -27,7 +27,7 @@ import java.util.{Locale, Properties, UUID}
 
 import kafka.utils.{CoreUtils, Exit, Logging}
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 import org.apache.commons.lang.text.StrSubstitutor
 import org.apache.directory.api.ldap.model.entry.{DefaultEntry, Entry}
 import org.apache.directory.api.ldap.model.ldif.LdifReader
@@ -94,7 +94,7 @@ class MiniKdc(config: Properties, workDir: File) extends Logging {
 
   info("Configuration:")
   info("---------------------------------------------------------------")
-  config.asScala.foreach { case (key, value) =>
+  config.forEach { (key, value) =>
     info(s"\t$key: $value")
   }
   info("---------------------------------------------------------------")
@@ -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)
@@ -359,7 +359,7 @@ object MiniKdc {
           throw new RuntimeException(s"Specified configuration does not exist: ${configFile.getAbsolutePath}")
 
         val userConfig = Utils.loadProps(configFile.getAbsolutePath)
-        userConfig.asScala.foreach { case (key, value) =>
+        userConfig.forEach { (key, value) =>
           config.put(key, value)
         }
         val keytabFile = new File(keytabPath).getAbsoluteFile
@@ -370,7 +370,7 @@ object MiniKdc {
     }
   }
 
-  private def start(workDir: File, config: Properties, keytabFile: File, principals: Seq[String]) {
+  private[minikdc] def start(workDir: File, config: Properties, keytabFile: File, principals: Seq[String]): MiniKdc = {
     val miniKdc = new MiniKdc(config, workDir)
     miniKdc.start()
     miniKdc.createPrincipal(keytabFile, principals: _*)
@@ -390,9 +390,8 @@ object MiniKdc {
       |
     """.stripMargin
     println(infoMessage)
-    Runtime.getRuntime.addShutdownHook(CoreUtils.newThread("minikdc-shutdown-hook", daemon = false) {
-      miniKdc.stop()
-    })
+    Exit.addShutdownHook("minikdc-shutdown-hook", miniKdc.stop())
+    miniKdc
   }
 
   val OrgName = "org.name"
diff --git a/core/src/test/scala/kafka/security/minikdc/MiniKdcTest.scala b/core/src/test/scala/kafka/security/minikdc/MiniKdcTest.scala
new file mode 100644
index 0000000000000..bfd973ddfbc84
--- /dev/null
+++ b/core/src/test/scala/kafka/security/minikdc/MiniKdcTest.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.security.minikdc
+
+import java.util.Properties
+
+import kafka.utils.TestUtils
+import org.junit.Test
+import org.junit.Assert._
+
+class MiniKdcTest {
+  @Test
+  def shouldNotStopImmediatelyWhenStarted(): Unit = {
+    val config = new Properties()
+    config.setProperty("kdc.bind.address", "0.0.0.0")
+    config.setProperty("transport", "TCP");
+    config.setProperty("max.ticket.lifetime", "86400000")
+    config.setProperty("org.name", "Example")
+    config.setProperty("kdc.port", "0")
+    config.setProperty("org.domain", "COM")
+    config.setProperty("max.renewable.lifetime", "604800000")
+    config.setProperty("instance", "DefaultKrbServer")
+    val minikdc = MiniKdc.start(TestUtils.tempDir(), config, TestUtils.tempFile(), List("foo"))
+    val running = System.getProperty(MiniKdc.JavaSecurityKrb5Conf) != null
+    try {
+      assertTrue("MiniKdc stopped immediately; it should not have", running)
+    } finally {
+      if (running) minikdc.stop()
+    }
+  }
+}
\ No newline at end of file
diff --git a/core/src/test/scala/kafka/server/BrokerMetadataCheckpointTest.scala b/core/src/test/scala/kafka/server/BrokerMetadataCheckpointTest.scala
new file mode 100644
index 0000000000000..480551a572900
--- /dev/null
+++ b/core/src/test/scala/kafka/server/BrokerMetadataCheckpointTest.scala
@@ -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 kafka.server
+
+import java.io.File
+
+import org.junit.Assert.assertEquals
+import org.junit.Test
+
+class BrokerMetadataCheckpointTest {
+  @Test
+  def testReadWithNonExistentFile(): Unit = {
+    assertEquals(None, new BrokerMetadataCheckpoint(new File("path/that/does/not/exist")).read())
+  }
+}
diff --git a/core/src/test/scala/kafka/server/BrokerToControllerRequestThreadTest.scala b/core/src/test/scala/kafka/server/BrokerToControllerRequestThreadTest.scala
new file mode 100644
index 0000000000000..b68e48aa0a4a6
--- /dev/null
+++ b/core/src/test/scala/kafka/server/BrokerToControllerRequestThreadTest.scala
@@ -0,0 +1,257 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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.{CountDownLatch, LinkedBlockingDeque, TimeUnit}
+import java.util.Collections
+import kafka.cluster.{Broker, EndPoint}
+import kafka.utils.TestUtils
+import org.apache.kafka.clients.{ClientResponse, ManualMetadataUpdater, Metadata, MockClient}
+import org.apache.kafka.common.feature.Features
+import org.apache.kafka.common.feature.Features.emptySupportedFeatures
+import org.apache.kafka.common.utils.{MockTime, SystemTime}
+import org.apache.kafka.common.message.MetadataRequestData
+import org.apache.kafka.common.network.ListenerName
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.requests.{AbstractRequest, MetadataRequest, MetadataResponse, RequestTestUtils}
+import org.apache.kafka.common.security.auth.SecurityProtocol
+import org.junit.Assert.{assertEquals, assertFalse, assertTrue}
+import org.junit.Test
+import org.mockito.Mockito._
+
+class BrokerToControllerRequestThreadTest {
+
+  @Test
+  def testRequestsSent(): Unit = {
+    // just a simple test that tests whether the request from 1 -> 2 is sent and the response callback is called
+    val time = new SystemTime
+    val config = new KafkaConfig(TestUtils.createBrokerConfig(1, "localhost:2181"))
+    val controllerId = 2
+
+    val metadata = mock(classOf[Metadata])
+    val mockClient = new MockClient(time, metadata)
+
+    val requestQueue = new LinkedBlockingDeque[BrokerToControllerQueueItem]()
+    val metadataCache = mock(classOf[MetadataCache])
+    val listenerName = ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT)
+    val activeController = new Broker(controllerId,
+      Seq(new EndPoint("host", 1234, listenerName, SecurityProtocol.PLAINTEXT)), None, emptySupportedFeatures)
+
+    when(metadataCache.getControllerId).thenReturn(Some(controllerId))
+    when(metadataCache.getAliveBrokers).thenReturn(Seq(activeController))
+    when(metadataCache.getAliveBroker(controllerId)).thenReturn(Some(activeController))
+
+    val expectedResponse = RequestTestUtils.metadataUpdateWith(2, Collections.singletonMap("a", new Integer(2)))
+    val testRequestThread = new BrokerToControllerRequestThread(mockClient, new ManualMetadataUpdater(), requestQueue, metadataCache,
+      config, listenerName, time, "")
+    mockClient.prepareResponse(expectedResponse)
+
+    val responseLatch = new CountDownLatch(1)
+    val queueItem = BrokerToControllerQueueItem(
+      new MetadataRequest.Builder(new MetadataRequestData()),
+      new TestRequestCompletionHandler(expectedResponse, responseLatch),
+      Long.MaxValue)
+    requestQueue.put(queueItem)
+    // initialize to the controller
+    testRequestThread.doWork()
+    // send and process the request
+    testRequestThread.doWork()
+
+    assertTrue(responseLatch.await(10, TimeUnit.SECONDS))
+  }
+
+  @Test
+  def testControllerChanged(): Unit = {
+    // in this test the current broker is 1, and the controller changes from 2 -> 3 then back: 3 -> 2
+    val time = new SystemTime
+    val config = new KafkaConfig(TestUtils.createBrokerConfig(1, "localhost:2181"))
+    val oldControllerId = 1
+    val newControllerId = 2
+
+    val metadata = mock(classOf[Metadata])
+    val mockClient = new MockClient(time, metadata)
+
+    val requestQueue = new LinkedBlockingDeque[BrokerToControllerQueueItem]()
+    val metadataCache = mock(classOf[MetadataCache])
+    val listenerName = ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT)
+    val oldController = new Broker(oldControllerId,
+      Seq(new EndPoint("host1", 1234, listenerName, SecurityProtocol.PLAINTEXT)), None, Features.emptySupportedFeatures)
+    val oldControllerNode = oldController.node(listenerName)
+    val newController = new Broker(newControllerId,
+      Seq(new EndPoint("host2", 1234, listenerName, SecurityProtocol.PLAINTEXT)), None, Features.emptySupportedFeatures)
+
+    when(metadataCache.getControllerId).thenReturn(Some(oldControllerId), Some(newControllerId))
+    when(metadataCache.getAliveBroker(oldControllerId)).thenReturn(Some(oldController))
+    when(metadataCache.getAliveBroker(newControllerId)).thenReturn(Some(newController))
+    when(metadataCache.getAliveBrokers).thenReturn(Seq(oldController, newController))
+
+    val expectedResponse = RequestTestUtils.metadataUpdateWith(3, Collections.singletonMap("a", new Integer(2)))
+    val testRequestThread = new BrokerToControllerRequestThread(mockClient, new ManualMetadataUpdater(),
+      requestQueue, metadataCache, config, listenerName, time, "")
+
+    val responseLatch = new CountDownLatch(1)
+
+    val queueItem = BrokerToControllerQueueItem(
+      new MetadataRequest.Builder(new MetadataRequestData()),
+      new TestRequestCompletionHandler(expectedResponse, responseLatch),
+      Long.MaxValue)
+    requestQueue.put(queueItem)
+    mockClient.prepareResponse(expectedResponse)
+    // initialize the thread with oldController
+    testRequestThread.doWork()
+    // assert queue correctness
+    assertFalse(requestQueue.isEmpty)
+    assertEquals(1, requestQueue.size())
+    assertEquals(queueItem, requestQueue.peek())
+    // disconnect the node
+    mockClient.setUnreachable(oldControllerNode, time.milliseconds() + 5000)
+    // verify that the client closed the connection to the faulty controller
+    testRequestThread.doWork()
+    assertFalse(requestQueue.isEmpty)
+    assertEquals(1, requestQueue.size())
+    // should connect to the new controller
+    testRequestThread.doWork()
+    // should send the request and process the response
+    testRequestThread.doWork()
+
+    assertTrue(responseLatch.await(10, TimeUnit.SECONDS))
+  }
+
+  @Test
+  def testNotController(): Unit = {
+    val time = new SystemTime
+    val config = new KafkaConfig(TestUtils.createBrokerConfig(1, "localhost:2181"))
+    val oldControllerId = 1
+    val newControllerId = 2
+
+    val metadata = mock(classOf[Metadata])
+    val mockClient = new MockClient(time, metadata)
+
+    val requestQueue = new LinkedBlockingDeque[BrokerToControllerQueueItem]()
+    val metadataCache = mock(classOf[MetadataCache])
+    val listenerName = ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT)
+    val oldController = new Broker(oldControllerId,
+      Seq(new EndPoint("host1", 1234, listenerName, SecurityProtocol.PLAINTEXT)), None, Features.emptySupportedFeatures)
+    val newController = new Broker(2,
+      Seq(new EndPoint("host2", 1234, listenerName, SecurityProtocol.PLAINTEXT)), None, Features.emptySupportedFeatures)
+
+    when(metadataCache.getControllerId).thenReturn(Some(oldControllerId), Some(newControllerId))
+    when(metadataCache.getAliveBrokers).thenReturn(Seq(oldController, newController))
+    when(metadataCache.getAliveBroker(oldControllerId)).thenReturn(Some(oldController))
+    when(metadataCache.getAliveBroker(newControllerId)).thenReturn(Some(newController))
+
+    val responseWithNotControllerError = RequestTestUtils.metadataUpdateWith("cluster1", 2,
+      Collections.singletonMap("a", Errors.NOT_CONTROLLER),
+      Collections.singletonMap("a", new Integer(2)))
+    val expectedResponse = RequestTestUtils.metadataUpdateWith(3, Collections.singletonMap("a", new Integer(2)))
+    val testRequestThread = new BrokerToControllerRequestThread(mockClient, new ManualMetadataUpdater(), requestQueue, metadataCache,
+      config, listenerName, time, "")
+
+    val responseLatch = new CountDownLatch(1)
+    val queueItem = BrokerToControllerQueueItem(
+      new MetadataRequest.Builder(new MetadataRequestData()
+        .setAllowAutoTopicCreation(true)),
+      new TestRequestCompletionHandler(expectedResponse, responseLatch),
+      Long.MaxValue)
+    requestQueue.put(queueItem)
+    // initialize to the controller
+    testRequestThread.doWork()
+    // send and process the request
+    mockClient.prepareResponse((body: AbstractRequest) => {
+      body.isInstanceOf[MetadataRequest] &&
+      body.asInstanceOf[MetadataRequest].allowAutoTopicCreation()
+    }, responseWithNotControllerError)
+    testRequestThread.doWork()
+    // reinitialize the controller to a different node
+    testRequestThread.doWork()
+    // process the request again
+    mockClient.prepareResponse(expectedResponse)
+    testRequestThread.doWork()
+
+    assertTrue(responseLatch.await(10, TimeUnit.SECONDS))
+  }
+
+  @Test
+  def testRequestTimeout(): Unit = {
+    val time = new MockTime()
+    val config = new KafkaConfig(TestUtils.createBrokerConfig(1, "localhost:2181"))
+    val controllerId = 1
+
+    val metadata = mock(classOf[Metadata])
+    val mockClient = new MockClient(time, metadata)
+
+    val requestQueue = new LinkedBlockingDeque[BrokerToControllerQueueItem]()
+    val metadataCache = mock(classOf[MetadataCache])
+    val listenerName = ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT)
+    val controller = new Broker(controllerId,
+      Seq(new EndPoint("host1", 1234, listenerName, SecurityProtocol.PLAINTEXT)), None, Features.emptySupportedFeatures)
+
+    when(metadataCache.getControllerId).thenReturn(Some(controllerId))
+    when(metadataCache.getAliveBrokers).thenReturn(Seq(controller))
+    when(metadataCache.getAliveBroker(controllerId)).thenReturn(Some(controller))
+
+    val responseWithNotControllerError = RequestTestUtils.metadataUpdateWith("cluster1", 2,
+      Collections.singletonMap("a", Errors.NOT_CONTROLLER),
+      Collections.singletonMap("a", new Integer(2)))
+    val testRequestThread = new BrokerToControllerRequestThread(mockClient, new ManualMetadataUpdater(), requestQueue, metadataCache,
+      config, listenerName, time, "")
+
+    val responseLatch = new CountDownLatch(1)
+    val requestTimeout = config.requestTimeoutMs.longValue()
+    val queueItem = BrokerToControllerQueueItem(
+      new MetadataRequest.Builder(new MetadataRequestData()
+        .setAllowAutoTopicCreation(true)), new ControllerRequestCompletionHandler {
+        override def onComplete(response: ClientResponse): Unit = {}
+
+        override def onTimeout(): Unit = {
+          responseLatch.countDown()
+        }
+      }, requestTimeout + time.milliseconds())
+    requestQueue.put(queueItem)
+
+    // initialize to the controller
+    testRequestThread.doWork()
+    // send and process the request
+    mockClient.prepareResponse((body: AbstractRequest) => {
+      // Advance time to timeout the response
+      time.sleep(requestTimeout + 1)
+
+      body.isInstanceOf[MetadataRequest] &&
+        body.asInstanceOf[MetadataRequest].allowAutoTopicCreation()
+    }, responseWithNotControllerError)
+
+    testRequestThread.doWork()
+
+    // The queued item should be timed out, instead of
+    // re-enqueue by NOT_CONTROLLER error.
+    assertEquals(0, requestQueue.size())
+
+    assertTrue(responseLatch.await(10, TimeUnit.SECONDS))
+  }
+
+  class TestRequestCompletionHandler(expectedResponse: MetadataResponse,
+                                     responseLatch: CountDownLatch) extends ControllerRequestCompletionHandler {
+    override def onComplete(response: ClientResponse): Unit = {
+      assertEquals(expectedResponse, response.responseBody())
+      responseLatch.countDown()
+    }
+
+    override def onTimeout(): Unit = {
+    }
+  }
+}
diff --git a/core/src/test/scala/kafka/tools/CustomDeserializerTest.scala b/core/src/test/scala/kafka/tools/CustomDeserializerTest.scala
new file mode 100644
index 0000000000000..b7041720c1625
--- /dev/null
+++ b/core/src/test/scala/kafka/tools/CustomDeserializerTest.scala
@@ -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 kafka.tools
+
+import java.io.PrintStream
+
+import kafka.utils.TestUtils
+import org.apache.kafka.clients.consumer.ConsumerRecord
+import org.apache.kafka.common.header.Headers
+import org.apache.kafka.common.serialization.Deserializer
+import org.hamcrest.CoreMatchers
+import org.hamcrest.MatcherAssert._
+import org.junit.Test
+import org.mockito.Mockito._
+
+class CustomDeserializer extends Deserializer[String] {
+
+  override def deserialize(topic: String, data: Array[Byte]): String = {
+    assertThat("topic must not be null", topic, CoreMatchers.notNullValue)
+    new String(data)
+  }
+
+  override def deserialize(topic: String, headers: Headers, data: Array[Byte]): String = {
+    println("WITH HEADERS")
+    new String(data)
+  }
+}
+
+class CustomDeserializerTest {
+
+  @Test
+  def checkFormatterCallDeserializerWithHeaders(): Unit = {
+    val formatter = new DefaultMessageFormatter()
+    formatter.valueDeserializer = Some(new CustomDeserializer)
+    val output = TestUtils.grabConsoleOutput(formatter.writeTo(
+      new ConsumerRecord("topic_test", 1, 1L, "key".getBytes, "value".getBytes), mock(classOf[PrintStream])))
+    assertThat("DefaultMessageFormatter should call `deserialize` method with headers.", output, CoreMatchers.containsString("WITH HEADERS"))
+    formatter.close()
+  }
+
+  @Test
+  def checkDeserializerTopicIsNotNull(): Unit = {
+    val formatter = new DefaultMessageFormatter()
+    formatter.keyDeserializer = Some(new CustomDeserializer)
+
+    formatter.writeTo(new ConsumerRecord("topic_test", 1, 1L, "key".getBytes, "value".getBytes),
+      mock(classOf[PrintStream]))
+
+    formatter.close()
+  }
+}
diff --git a/core/src/test/scala/kafka/tools/DefaultMessageFormatterTest.scala b/core/src/test/scala/kafka/tools/DefaultMessageFormatterTest.scala
new file mode 100644
index 0000000000000..27442856515ab
--- /dev/null
+++ b/core/src/test/scala/kafka/tools/DefaultMessageFormatterTest.scala
@@ -0,0 +1,237 @@
+/**
+  * Licensed to the Apache Software Foundation (ASF) under one or more
+  * contributor license agreements.  See the NOTICE file distributed with
+  * this work for additional information regarding copyright ownership.
+  * The ASF licenses this file to You under the Apache License, Version 2.0
+  * (the "License"); you may not use this file except in compliance with
+  * the License.  You may obtain a copy of the License at
+  *
+  * http://www.apache.org/licenses/LICENSE-2.0
+  *
+  * Unless required by applicable law or agreed to in writing, software
+  * distributed under the License is distributed on an "AS IS" BASIS,
+  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  * See the License for the specific language governing permissions and
+  * limitations under the License.
+  */
+
+package unit.kafka.tools
+
+import java.io.{ByteArrayOutputStream, Closeable, PrintStream}
+import java.nio.charset.StandardCharsets
+import java.util
+
+import kafka.tools.DefaultMessageFormatter
+import org.apache.kafka.clients.consumer.ConsumerRecord
+import org.apache.kafka.common.header.Header
+import org.apache.kafka.common.header.internals.{RecordHeader, RecordHeaders}
+import org.apache.kafka.common.record.TimestampType
+import org.apache.kafka.common.serialization.Deserializer
+import org.junit.Assert._
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.Parameterized
+import org.junit.runners.Parameterized.Parameters
+
+import scala.jdk.CollectionConverters._
+
+@RunWith(value = classOf[Parameterized])
+class DefaultMessageFormatterTest(name: String, record: ConsumerRecord[Array[Byte], Array[Byte]], properties: Map[String, String], expected: String) {
+  import DefaultMessageFormatterTest._
+
+  @Test
+  def testWriteRecord()= {
+    withResource(new ByteArrayOutputStream()) { baos =>
+      withResource(new PrintStream(baos)) { ps =>
+        val formatter = buildFormatter(properties)
+        formatter.writeTo(record, ps)
+        val actual = new String(baos.toByteArray(), StandardCharsets.UTF_8)
+        assertEquals(expected, actual)
+
+      }
+    }
+  }
+}
+
+object DefaultMessageFormatterTest {
+  @Parameters(name = "Test {index} - {0}")
+  def parameters: java.util.Collection[Array[Object]] = {
+    Seq(
+      Array(
+        "print nothing",
+        consumerRecord(),
+        Map("print.value" -> "false"),
+        ""),
+      Array(
+        "print key",
+        consumerRecord(),
+        Map("print.key" -> "true",
+          "print.value" -> "false"),
+        "someKey\n"),
+      Array(
+        "print value",
+        consumerRecord(),
+        Map(),
+        "someValue\n"),
+      Array(
+        "print empty timestamp",
+        consumerRecord(timestampType = TimestampType.NO_TIMESTAMP_TYPE),
+        Map("print.timestamp" -> "true",
+            "print.value" -> "false"),
+        "NO_TIMESTAMP\n"),
+      Array(
+        "print log append time timestamp",
+        consumerRecord(timestampType = TimestampType.LOG_APPEND_TIME),
+        Map("print.timestamp" -> "true",
+            "print.value" -> "false"),
+        "LogAppendTime:1234\n"),
+      Array(
+        "print create time timestamp",
+        consumerRecord(timestampType = TimestampType.CREATE_TIME),
+        Map("print.timestamp" -> "true",
+            "print.value" -> "false"),
+        "CreateTime:1234\n"),
+      Array(
+        "print partition",
+        consumerRecord(),
+        Map("print.partition" -> "true",
+            "print.value" -> "false"),
+        "Partition:9\n"),
+      Array(
+        "print offset",
+        consumerRecord(),
+        Map("print.offset" -> "true",
+            "print.value" -> "false"),
+        "Offset:9876\n"),
+      Array(
+        "print headers",
+        consumerRecord(),
+        Map("print.headers" -> "true",
+            "print.value" -> "false"),
+        "h1:v1,h2:v2\n"),
+      Array(
+        "print empty headers",
+        consumerRecord(headers = Nil),
+        Map("print.headers" -> "true",
+            "print.value" -> "false"),
+        "NO_HEADERS\n"),
+      Array(
+        "print all possible fields with default delimiters",
+        consumerRecord(),
+        Map("print.key" -> "true",
+            "print.timestamp" -> "true",
+            "print.partition" -> "true",
+            "print.offset" -> "true",
+            "print.headers" -> "true",
+            "print.value" -> "true"),
+        "CreateTime:1234\tPartition:9\tOffset:9876\th1:v1,h2:v2\tsomeKey\tsomeValue\n"),
+      Array(
+        "print all possible fields with custom delimiters",
+        consumerRecord(),
+        Map("key.separator" -> "|",
+            "line.separator" -> "^",
+            "headers.separator" -> "#",
+            "print.key" -> "true",
+            "print.timestamp" -> "true",
+            "print.partition" -> "true",
+            "print.offset" -> "true",
+            "print.headers" -> "true",
+            "print.value" -> "true"),
+        "CreateTime:1234|Partition:9|Offset:9876|h1:v1#h2:v2|someKey|someValue^"),
+      Array(
+        "print key with custom deserializer",
+        consumerRecord(),
+        Map("print.key" -> "true",
+            "print.headers" -> "true",
+            "print.value" -> "true",
+            "key.deserializer" -> "unit.kafka.tools.UpperCaseDeserializer"),
+        "h1:v1,h2:v2\tSOMEKEY\tsomeValue\n"),
+      Array(
+        "print value with custom deserializer",
+        consumerRecord(),
+        Map("print.key" -> "true",
+            "print.headers" -> "true",
+            "print.value" -> "true",
+            "value.deserializer" -> "unit.kafka.tools.UpperCaseDeserializer"),
+        "h1:v1,h2:v2\tsomeKey\tSOMEVALUE\n"),
+      Array(
+        "print headers with custom deserializer",
+        consumerRecord(),
+        Map("print.key" -> "true",
+            "print.headers" -> "true",
+            "print.value" -> "true",
+            "headers.deserializer" -> "unit.kafka.tools.UpperCaseDeserializer"),
+        "h1:V1,h2:V2\tsomeKey\tsomeValue\n"),
+      Array(
+        "print key and value",
+        consumerRecord(),
+        Map("print.key" -> "true",
+            "print.value" -> "true"),
+        "someKey\tsomeValue\n"),
+      Array(
+        "print fields in the beginning, middle and the end",
+        consumerRecord(),
+        Map("print.key" -> "true",
+            "print.value" -> "true",
+            "print.partition" -> "true"),
+        "Partition:9\tsomeKey\tsomeValue\n"),
+      Array(
+        "null value without custom null literal",
+        consumerRecord(value = null),
+        Map("print.key" -> "true"),
+        "someKey\tnull\n"),
+      Array(
+        "null value with custom null literal",
+        consumerRecord(value = null),
+        Map("print.key" -> "true",
+            "null.literal" -> "NULL"),
+        "someKey\tNULL\n"),
+    ).asJava
+  }
+
+  private def buildFormatter(propsToSet: Map[String, String]): DefaultMessageFormatter = {
+    val formatter = new DefaultMessageFormatter()
+    formatter.configure(propsToSet.asJava)
+    formatter
+  }
+
+
+  private def header(key: String, value: String) = {
+    new RecordHeader(key, value.getBytes(StandardCharsets.UTF_8))
+  }
+
+  private def consumerRecord(key: String = "someKey",
+                             value: String = "someValue",
+                             headers: Iterable[Header] = Seq(header("h1", "v1"), header("h2", "v2")),
+                             partition: Int = 9,
+                             offset: Long = 9876,
+                             timestamp: Long = 1234,
+                             timestampType: TimestampType = TimestampType.CREATE_TIME) = {
+    new ConsumerRecord[Array[Byte], Array[Byte]](
+      "someTopic",
+      partition,
+      offset,
+      timestamp,
+      timestampType,
+      0L,
+      0,
+      0,
+      if (key == null) null else key.getBytes(StandardCharsets.UTF_8),
+      if (value == null) null else value.getBytes(StandardCharsets.UTF_8),
+      new RecordHeaders(headers.asJava))
+  }
+
+  private def withResource[Resource <: Closeable, Result](resource: Resource)(handler: Resource => Result): Result = {
+    try {
+      handler(resource)
+    } finally {
+      resource.close()
+    }
+  }
+}
+
+class UpperCaseDeserializer extends Deserializer[String] {
+  override def configure(configs: util.Map[String, _], isKey: Boolean): Unit = {}
+  override def deserialize(topic: String, data: Array[Byte]): String = new String(data, StandardCharsets.UTF_8).toUpperCase
+  override def close(): Unit = {}
+}
diff --git a/core/src/test/scala/kafka/tools/LogCompactionTester.scala b/core/src/test/scala/kafka/tools/LogCompactionTester.scala
new file mode 100755
index 0000000000000..da8a3c03a176f
--- /dev/null
+++ b/core/src/test/scala/kafka/tools/LogCompactionTester.scala
@@ -0,0 +1,348 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package kafka.tools
+
+import java.io._
+import java.nio.ByteBuffer
+import java.nio.charset.StandardCharsets.UTF_8
+import java.nio.file.{Files, Path}
+import java.time.Duration
+import java.util.{Properties, Random}
+
+import joptsimple.OptionParser
+import kafka.utils._
+import org.apache.kafka.clients.admin.{Admin, NewTopic}
+import org.apache.kafka.clients.CommonClientConfigs
+import org.apache.kafka.clients.consumer.{ConsumerConfig, KafkaConsumer}
+import org.apache.kafka.clients.producer.{KafkaProducer, ProducerConfig, ProducerRecord}
+import org.apache.kafka.common.config.TopicConfig
+import org.apache.kafka.common.serialization.{ByteArraySerializer, StringDeserializer}
+import org.apache.kafka.common.utils.{AbstractIterator, Utils}
+
+import scala.jdk.CollectionConverters._
+
+/**
+ * This is a torture test that runs against an existing broker
+ *
+ * Here is how it works:
+ *
+ * It produces a series of specially formatted messages to one or more partitions. Each message it produces
+ * it logs out to a text file. The messages have a limited set of keys, so there is duplication in the key space.
+ *
+ * The broker will clean its log as the test runs.
+ *
+ * When the specified number of messages have been produced we create a consumer and consume all the messages in the topic
+ * and write that out to another text file.
+ *
+ * Using a stable unix sort we sort both the producer log of what was sent and the consumer log of what was retrieved by the message key.
+ * Then we compare the final message in both logs for each key. If this final message is not the same for all keys we
+ * print an error and exit with exit code 1, otherwise we print the size reduction and exit with exit code 0.
+ */
+object LogCompactionTester {
+
+  //maximum line size while reading produced/consumed record text file
+  private val ReadAheadLimit = 4906
+
+  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
+      .describedAs("count")
+      .ofType(classOf[java.lang.Long])
+      .defaultsTo(Long.MaxValue)
+    val messageCompressionOpt = parser.accepts("compression-type", "message compression type")
+      .withOptionalArg
+      .describedAs("compressionType")
+      .ofType(classOf[java.lang.String])
+      .defaultsTo("none")
+    val numDupsOpt = parser.accepts("duplicates", "The number of duplicates for each key.")
+      .withRequiredArg
+      .describedAs("count")
+      .ofType(classOf[java.lang.Integer])
+      .defaultsTo(5)
+    val brokerOpt = parser.accepts("bootstrap-server", "The server(s) to connect to.")
+      .withRequiredArg
+      .describedAs("url")
+      .ofType(classOf[String])
+    val topicsOpt = parser.accepts("topics", "The number of topics to test.")
+      .withRequiredArg
+      .describedAs("count")
+      .ofType(classOf[java.lang.Integer])
+      .defaultsTo(1)
+    val percentDeletesOpt = parser.accepts("percent-deletes", "The percentage of updates that are deletes.")
+      .withRequiredArg
+      .describedAs("percent")
+      .ofType(classOf[java.lang.Integer])
+      .defaultsTo(0)
+    val sleepSecsOpt = parser.accepts("sleep", "Time in milliseconds to sleep between production and consumption.")
+      .withRequiredArg
+      .describedAs("ms")
+      .ofType(classOf[java.lang.Integer])
+      .defaultsTo(0)
+
+    val options = parser.parse(args: _*)
+
+    if (args.length == 0)
+      CommandLineUtils.printUsageAndDie(parser, "A tool to test log compaction. Valid options are: ")
+
+    CommandLineUtils.checkRequiredArgs(parser, options, brokerOpt, numMessagesOpt)
+
+    // parse options
+    val messages = options.valueOf(numMessagesOpt).longValue
+    val compressionType = options.valueOf(messageCompressionOpt)
+    val percentDeletes = options.valueOf(percentDeletesOpt).intValue
+    val dups = options.valueOf(numDupsOpt).intValue
+    val brokerUrl = options.valueOf(brokerOpt)
+    val topicCount = options.valueOf(topicsOpt).intValue
+    val sleepSecs = options.valueOf(sleepSecsOpt).intValue
+
+    val testId = new Random().nextLong
+    val topics = (0 until topicCount).map("log-cleaner-test-" + testId + "-" + _).toArray
+    createTopics(brokerUrl, topics.toSeq)
+
+    println(s"Producing $messages messages..to topics ${topics.mkString(",")}")
+    val producedDataFilePath = produceMessages(brokerUrl, topics, messages, compressionType, dups, percentDeletes)
+    println(s"Sleeping for $sleepSecs seconds...")
+    Thread.sleep(sleepSecs * 1000)
+    println("Consuming messages...")
+    val consumedDataFilePath = consumeMessages(brokerUrl, topics)
+
+    val producedLines = lineCount(producedDataFilePath)
+    val consumedLines = lineCount(consumedDataFilePath)
+    val reduction = 100 * (1.0 - consumedLines.toDouble / producedLines.toDouble)
+    println(f"$producedLines%d rows of data produced, $consumedLines%d rows of data consumed ($reduction%.1f%% reduction).")
+
+    println("De-duplicating and validating output files...")
+    validateOutput(producedDataFilePath.toFile, consumedDataFilePath.toFile)
+    Utils.delete(producedDataFilePath.toFile)
+    Utils.delete(consumedDataFilePath.toFile)
+    //if you change this line, we need to update test_log_compaction_tool.py system test
+    println("Data verification is completed")
+  }
+
+  def createTopics(brokerUrl: String, topics: Seq[String]): Unit = {
+    val adminConfig = new Properties
+    adminConfig.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, brokerUrl)
+    val adminClient = Admin.create(adminConfig)
+
+    try {
+      val topicConfigs = Map(TopicConfig.CLEANUP_POLICY_CONFIG -> TopicConfig.CLEANUP_POLICY_COMPACT)
+      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()
+      TestUtils.waitUntilTrue(() => {
+        val allTopics = adminClient.listTopics.names.get.asScala.toSeq
+        pendingTopics = topics.filter(topicName => !allTopics.contains(topicName))
+        pendingTopics.isEmpty
+      }, s"timed out waiting for topics : $pendingTopics")
+
+    } finally adminClient.close()
+  }
+
+  def lineCount(filPath: Path): Int = Files.readAllLines(filPath).size
+
+  def validateOutput(producedDataFile: File, consumedDataFile: File): Unit = {
+    val producedReader = externalSort(producedDataFile)
+    val consumedReader = externalSort(consumedDataFile)
+    val produced = valuesIterator(producedReader)
+    val consumed = valuesIterator(consumedReader)
+
+    val producedDedupedFile = new File(producedDataFile.getAbsolutePath + ".deduped")
+    val producedDeduped : BufferedWriter = Files.newBufferedWriter(producedDedupedFile.toPath, UTF_8)
+
+    val consumedDedupedFile = new File(consumedDataFile.getAbsolutePath + ".deduped")
+    val consumedDeduped : BufferedWriter = Files.newBufferedWriter(consumedDedupedFile.toPath, UTF_8)
+    var total = 0
+    var mismatched = 0
+    while (produced.hasNext && consumed.hasNext) {
+      val p = produced.next()
+      producedDeduped.write(p.toString)
+      producedDeduped.newLine()
+      val c = consumed.next()
+      consumedDeduped.write(c.toString)
+      consumedDeduped.newLine()
+      if (p != c)
+        mismatched += 1
+      total += 1
+    }
+    producedDeduped.close()
+    consumedDeduped.close()
+    println(s"Validated $total values, $mismatched mismatches.")
+    require(!produced.hasNext, "Additional values produced not found in consumer log.")
+    require(!consumed.hasNext, "Additional values consumed not found in producer log.")
+    require(mismatched == 0, "Non-zero number of row mismatches.")
+    // if all the checks worked out we can delete the deduped files
+    Utils.delete(producedDedupedFile)
+    Utils.delete(consumedDedupedFile)
+  }
+
+  def require(requirement: Boolean, message: => Any): Unit = {
+    if (!requirement) {
+      System.err.println(s"Data validation failed : $message")
+      Exit.exit(1)
+    }
+  }
+
+  def valuesIterator(reader: BufferedReader): Iterator[TestRecord] = {
+    new AbstractIterator[TestRecord] {
+      def makeNext(): TestRecord = {
+        var next = readNext(reader)
+        while (next != null && next.delete)
+          next = readNext(reader)
+        if (next == null)
+          allDone()
+        else
+          next
+      }
+    }.asScala
+  }
+
+  def readNext(reader: BufferedReader): TestRecord = {
+    var line = reader.readLine()
+    if (line == null)
+      return null
+    var curr = TestRecord.parse(line)
+    while (true) {
+      line = peekLine(reader)
+      if (line == null)
+        return curr
+      val next = TestRecord.parse(line)
+      if (next == null || next.topicAndKey != curr.topicAndKey)
+        return curr
+      curr = next
+      reader.readLine()
+    }
+    null
+  }
+
+  def peekLine(reader: BufferedReader) = {
+    reader.mark(ReadAheadLimit)
+    val line = reader.readLine
+    reader.reset()
+    line
+  }
+
+  def externalSort(file: File): BufferedReader = {
+    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(): Unit = {
+        val exitCode = process.waitFor()
+        if (exitCode != 0) {
+          System.err.println("Process exited abnormally.")
+          while (process.getErrorStream.available > 0) {
+            System.err.write(process.getErrorStream().read())
+          }
+        }
+      }
+    }.start()
+    new BufferedReader(new InputStreamReader(process.getInputStream(), UTF_8), 10 * 1024 * 1024)
+  }
+
+  def produceMessages(brokerUrl: String,
+                      topics: Array[String],
+                      messages: Long,
+                      compressionType: String,
+                      dups: Int,
+                      percentDeletes: Int): Path = {
+    val producerProps = new Properties
+    producerProps.setProperty(ProducerConfig.MAX_BLOCK_MS_CONFIG, Long.MaxValue.toString)
+    producerProps.setProperty(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, brokerUrl)
+    producerProps.setProperty(ProducerConfig.COMPRESSION_TYPE_CONFIG, compressionType)
+    val producer = new KafkaProducer(producerProps, new ByteArraySerializer, new ByteArraySerializer)
+    try {
+      val rand = new Random(1)
+      val keyCount = (messages / dups).toInt
+      val producedFilePath = Files.createTempFile("kafka-log-cleaner-produced-", ".txt")
+      println(s"Logging produce requests to $producedFilePath")
+      val producedWriter: BufferedWriter = Files.newBufferedWriter(producedFilePath, UTF_8)
+      for (i <- 0L until (messages * topics.length)) {
+        val topic = topics((i % topics.length).toInt)
+        val key = rand.nextInt(keyCount)
+        val delete = (i % 100) < percentDeletes
+        val msg =
+          if (delete)
+            new ProducerRecord[Array[Byte], Array[Byte]](topic, key.toString.getBytes(UTF_8), null)
+          else
+            new ProducerRecord(topic, key.toString.getBytes(UTF_8), i.toString.getBytes(UTF_8))
+        producer.send(msg)
+        producedWriter.write(TestRecord(topic, key, i, delete).toString)
+        producedWriter.newLine()
+      }
+      producedWriter.close()
+      producedFilePath
+    } finally {
+      producer.close()
+    }
+  }
+
+  def createConsumer(brokerUrl: String): KafkaConsumer[String, String] = {
+    val consumerProps = new Properties
+    consumerProps.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "log-cleaner-test-" + new Random().nextInt(Int.MaxValue))
+    consumerProps.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, brokerUrl)
+    consumerProps.setProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest")
+    new KafkaConsumer(consumerProps, new StringDeserializer, new StringDeserializer)
+  }
+
+  def consumeMessages(brokerUrl: String, topics: Array[String]): Path = {
+    val consumer = createConsumer(brokerUrl)
+    consumer.subscribe(topics.toSeq.asJava)
+    val consumedFilePath = Files.createTempFile("kafka-log-cleaner-consumed-", ".txt")
+    println(s"Logging consumed messages to $consumedFilePath")
+    val consumedWriter: BufferedWriter = Files.newBufferedWriter(consumedFilePath, UTF_8)
+
+    try {
+      var done = false
+      while (!done) {
+        val consumerRecords = consumer.poll(Duration.ofSeconds(20))
+        if (!consumerRecords.isEmpty) {
+          for (record <- consumerRecords.asScala) {
+            val delete = record.value == null
+            val value = if (delete) -1L else record.value.toLong
+            consumedWriter.write(TestRecord(record.topic, record.key.toInt, value, delete).toString)
+            consumedWriter.newLine
+          }
+        } else {
+          done = true
+        }
+      }
+      consumedFilePath
+    } finally {
+      consumedWriter.close()
+      consumer.close()
+    }
+  }
+
+  def readString(buffer: ByteBuffer): String = {
+    Utils.utf8(buffer)
+  }
+
+}
+
+case class TestRecord(topic: String, key: Int, value: Long, delete: Boolean) {
+  override def toString = topic + "\t" + key + "\t" + value + "\t" + (if (delete) "d" else "u")
+  def topicAndKey = topic + key
+}
+
+object TestRecord {
+  def parse(line: String): TestRecord = {
+    val components = line.split("\t")
+    new TestRecord(components(0), components(1).toInt, components(2).toLong, components(3) == "d")
+  }
+}
diff --git a/core/src/test/scala/kafka/tools/ReplicaVerificationToolTest.scala b/core/src/test/scala/kafka/tools/ReplicaVerificationToolTest.scala
index 211413a341fe4..25c15cf3da4c1 100644
--- a/core/src/test/scala/kafka/tools/ReplicaVerificationToolTest.scala
+++ b/core/src/test/scala/kafka/tools/ReplicaVerificationToolTest.scala
@@ -1,27 +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.
-  */
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
 
 package kafka.tools
 
-import kafka.api.FetchResponsePartitionData
-import kafka.common.TopicAndPartition
-import kafka.message.ByteBufferMessageSet
+import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.protocol.Errors
-import org.apache.kafka.common.record.{CompressionType, SimpleRecord, MemoryRecords}
+import org.apache.kafka.common.record.{CompressionType, MemoryRecords, SimpleRecord}
+import org.apache.kafka.common.requests.FetchResponse
 import org.junit.Test
 import org.junit.Assert.assertTrue
 
@@ -32,12 +31,12 @@ class ReplicaVerificationToolTest {
     val sb = new StringBuilder
 
     val expectedReplicasPerTopicAndPartition = Map(
-      TopicAndPartition("a", 0) -> 3,
-      TopicAndPartition("a", 1) -> 3,
-      TopicAndPartition("b", 0) -> 2
+      new TopicPartition("a", 0) -> 3,
+      new TopicPartition("a", 1) -> 3,
+      new TopicPartition("b", 0) -> 2
     )
 
-    val replicaBuffer = new ReplicaBuffer(expectedReplicasPerTopicAndPartition, Map.empty, 2, Map.empty, 0, 0)
+    val replicaBuffer = new ReplicaBuffer(expectedReplicasPerTopicAndPartition, Map.empty, 2, 0)
     expectedReplicasPerTopicAndPartition.foreach { case (tp, numReplicas) =>
       (0 until numReplicas).foreach { replicaId =>
         val records = (0 to 5).map { index =>
@@ -45,8 +44,9 @@ class ReplicaVerificationToolTest {
         }
         val initialOffset = 4
         val memoryRecords = MemoryRecords.withRecords(initialOffset, CompressionType.NONE, records: _*)
-        replicaBuffer.addFetchedData(tp, replicaId, new FetchResponsePartitionData(Errors.NONE, hw = 20,
-          new ByteBufferMessageSet(memoryRecords.buffer)))
+        val partitionData = new FetchResponse.PartitionData(Errors.NONE, 20, 20, 0L, null, memoryRecords)
+
+        replicaBuffer.addFetchedData(tp, replicaId, partitionData)
       }
     }
 
@@ -55,7 +55,7 @@ class ReplicaVerificationToolTest {
 
     // If you change this assertion, you should verify that the replica_verification_test.py system test still passes
     assertTrue(s"Max lag information should be in output: `$output`",
-      output.endsWith(": max lag is 10 for partition a-0 at offset 10 among 3 partitions"))
+      output.endsWith(": max lag is 10 for partition a-1 at offset 10 among 3 partitions"))
   }
 
 }
diff --git a/core/src/test/scala/kafka/tools/TestLogCleaning.scala b/core/src/test/scala/kafka/tools/TestLogCleaning.scala
deleted file mode 100755
index 4ad6629eae1d8..0000000000000
--- a/core/src/test/scala/kafka/tools/TestLogCleaning.scala
+++ /dev/null
@@ -1,323 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package kafka.tools
-
-import joptsimple.OptionParser
-import java.util.Properties
-import java.util.Random
-import java.io._
-
-import kafka.consumer._
-import kafka.serializer._
-import kafka.utils._
-import kafka.log.Log
-import org.apache.kafka.clients.producer.{KafkaProducer, ProducerConfig, ProducerRecord}
-import org.apache.kafka.common.record.FileRecords
-
-import scala.collection.JavaConverters._
-
-/**
- * This is a torture test that runs against an existing broker. Here is how it works:
- *
- * It produces a series of specially formatted messages to one or more partitions. Each message it produces
- * it logs out to a text file. The messages have a limited set of keys, so there is duplication in the key space.
- *
- * The broker will clean its log as the test runs.
- *
- * When the specified number of messages have been produced we create a consumer and consume all the messages in the topic
- * and write that out to another text file.
- *
- * Using a stable unix sort we sort both the producer log of what was sent and the consumer log of what was retrieved by the message key.
- * Then we compare the final message in both logs for each key. If this final message is not the same for all keys we
- * print an error and exit with exit code 1, otherwise we print the size reduction and exit with exit code 0.
- */
-object TestLogCleaning {
-
-  def main(args: Array[String]) {
-    val parser = new OptionParser(false)
-    val numMessagesOpt = parser.accepts("messages", "The number of messages to send or consume.")
-                               .withRequiredArg
-                               .describedAs("count")
-                               .ofType(classOf[java.lang.Long])
-                               .defaultsTo(Long.MaxValue)
-    val messageCompressionOpt = parser.accepts("compression-type", "message compression type")
-                               .withOptionalArg()
-                               .describedAs("compressionType")
-                               .ofType(classOf[java.lang.String])
-                               .defaultsTo("none")
-    val numDupsOpt = parser.accepts("duplicates", "The number of duplicates for each key.")
-                           .withRequiredArg
-                           .describedAs("count")
-                           .ofType(classOf[java.lang.Integer])
-                           .defaultsTo(5)
-    val brokerOpt = parser.accepts("broker", "Url to connect to.")
-                          .withRequiredArg
-                          .describedAs("url")
-                          .ofType(classOf[String])
-    val topicsOpt = parser.accepts("topics", "The number of topics to test.")
-                          .withRequiredArg
-                          .describedAs("count")
-                          .ofType(classOf[java.lang.Integer])
-                          .defaultsTo(1)
-    val percentDeletesOpt = parser.accepts("percent-deletes", "The percentage of updates that are deletes.")
-    		                      .withRequiredArg
-    		                      .describedAs("percent")
-    		                      .ofType(classOf[java.lang.Integer])
-    		                      .defaultsTo(0)
-    val zkConnectOpt = parser.accepts("zk", "Zk url.")
-                             .withRequiredArg
-                             .describedAs("url")
-                             .ofType(classOf[String])
-    val sleepSecsOpt = parser.accepts("sleep", "Time to sleep between production and consumption.")
-                             .withRequiredArg
-                             .describedAs("ms")
-                             .ofType(classOf[java.lang.Integer])
-                             .defaultsTo(0)
-    val dumpOpt = parser.accepts("dump", "Dump the message contents of a topic partition that contains test data from this test to standard out.")
-                        .withRequiredArg
-                        .describedAs("directory")
-                        .ofType(classOf[String])
-
-    val options = parser.parse(args:_*)
-
-    if(args.length == 0)
-      CommandLineUtils.printUsageAndDie(parser, "An integration test for log cleaning.")
-
-    if(options.has(dumpOpt)) {
-      dumpLog(new File(options.valueOf(dumpOpt)))
-      Exit.exit(0)
-    }
-
-    CommandLineUtils.checkRequiredArgs(parser, options, brokerOpt, zkConnectOpt, numMessagesOpt)
-
-    // parse options
-    val messages = options.valueOf(numMessagesOpt).longValue
-    val compressionType = options.valueOf(messageCompressionOpt)
-    val percentDeletes = options.valueOf(percentDeletesOpt).intValue
-    val dups = options.valueOf(numDupsOpt).intValue
-    val brokerUrl = options.valueOf(brokerOpt)
-    val topicCount = options.valueOf(topicsOpt).intValue
-    val zkUrl = options.valueOf(zkConnectOpt)
-    val sleepSecs = options.valueOf(sleepSecsOpt).intValue
-
-    val testId = new Random().nextInt(Int.MaxValue)
-    val topics = (0 until topicCount).map("log-cleaner-test-" + testId + "-" + _).toArray
-
-    println("Producing %d messages...".format(messages))
-    val producedDataFile = produceMessages(brokerUrl, topics, messages, compressionType, dups, percentDeletes)
-    println("Sleeping for %d seconds...".format(sleepSecs))
-    Thread.sleep(sleepSecs * 1000)
-    println("Consuming messages...")
-    val consumedDataFile = consumeMessages(zkUrl, topics)
-
-    val producedLines = lineCount(producedDataFile)
-    val consumedLines = lineCount(consumedDataFile)
-    val reduction = 1.0 - consumedLines.toDouble/producedLines.toDouble
-    println("%d rows of data produced, %d rows of data consumed (%.1f%% reduction).".format(producedLines, consumedLines, 100 * reduction))
-
-    println("De-duplicating and validating output files...")
-    validateOutput(producedDataFile, consumedDataFile)
-    producedDataFile.delete()
-    consumedDataFile.delete()
-  }
-
-  def dumpLog(dir: File) {
-    require(dir.exists, "Non-existent directory: " + dir.getAbsolutePath)
-    for (file <- dir.list.sorted; if file.endsWith(Log.LogFileSuffix)) {
-      val fileRecords = FileRecords.open(new File(dir, file))
-      for (entry <- fileRecords.records.asScala) {
-        val key = TestUtils.readString(entry.key)
-        val content =
-          if (!entry.hasValue)
-            null
-          else
-            TestUtils.readString(entry.value)
-        println("offset = %s, key = %s, content = %s".format(entry.offset, key, content))
-      }
-    }
-  }
-
-  def lineCount(file: File): Int = io.Source.fromFile(file).getLines.size
-
-  def validateOutput(producedDataFile: File, consumedDataFile: File) {
-    val producedReader = externalSort(producedDataFile)
-    val consumedReader = externalSort(consumedDataFile)
-    val produced = valuesIterator(producedReader)
-    val consumed = valuesIterator(consumedReader)
-    val producedDedupedFile = new File(producedDataFile.getAbsolutePath + ".deduped")
-    val producedDeduped = new BufferedWriter(new FileWriter(producedDedupedFile), 1024*1024)
-    val consumedDedupedFile = new File(consumedDataFile.getAbsolutePath + ".deduped")
-    val consumedDeduped = new BufferedWriter(new FileWriter(consumedDedupedFile), 1024*1024)
-    var total = 0
-    var mismatched = 0
-    while(produced.hasNext && consumed.hasNext) {
-      val p = produced.next()
-      producedDeduped.write(p.toString)
-      producedDeduped.newLine()
-      val c = consumed.next()
-      consumedDeduped.write(c.toString)
-      consumedDeduped.newLine()
-      if(p != c)
-        mismatched += 1
-      total += 1
-    }
-    producedDeduped.close()
-    consumedDeduped.close()
-    println("Validated " + total + " values, " + mismatched + " mismatches.")
-    require(!produced.hasNext, "Additional values produced not found in consumer log.")
-    require(!consumed.hasNext, "Additional values consumed not found in producer log.")
-    require(mismatched == 0, "Non-zero number of row mismatches.")
-    // if all the checks worked out we can delete the deduped files
-    producedDedupedFile.delete()
-    consumedDedupedFile.delete()
-  }
-
-  def valuesIterator(reader: BufferedReader) = {
-    new IteratorTemplate[TestRecord] {
-      def makeNext(): TestRecord = {
-        var next = readNext(reader)
-        while(next != null && next.delete)
-          next = readNext(reader)
-        if(next == null)
-          allDone()
-        else
-          next
-      }
-    }
-  }
-
-  def readNext(reader: BufferedReader): TestRecord = {
-    var line = reader.readLine()
-    if(line == null)
-      return null
-    var curr = new TestRecord(line)
-    while(true) {
-      line = peekLine(reader)
-      if(line == null)
-        return curr
-      val next = new TestRecord(line)
-      if(next == null || next.topicAndKey != curr.topicAndKey)
-        return curr
-      curr = next
-      reader.readLine()
-    }
-    null
-  }
-
-  def peekLine(reader: BufferedReader) = {
-    reader.mark(4096)
-    val line = reader.readLine
-    reader.reset()
-    line
-  }
-
-  def externalSort(file: File): BufferedReader = {
-    val builder = new ProcessBuilder("sort", "--key=1,2", "--stable", "--buffer-size=20%", "--temporary-directory=" + System.getProperty("java.io.tmpdir"), file.getAbsolutePath)
-    val process = builder.start()
-    new Thread() {
-      override def run() {
-        val exitCode = process.waitFor()
-        if(exitCode != 0) {
-          System.err.println("Process exited abnormally.")
-          while(process.getErrorStream.available > 0) {
-            System.err.write(process.getErrorStream().read())
-          }
-        }
-      }
-    }.start()
-    new BufferedReader(new InputStreamReader(process.getInputStream()), 10*1024*1024)
-  }
-
-  def produceMessages(brokerUrl: String,
-                      topics: Array[String],
-                      messages: Long,
-                      compressionType: String,
-                      dups: Int,
-                      percentDeletes: Int): File = {
-    val producerProps = new Properties
-    producerProps.setProperty(ProducerConfig.MAX_BLOCK_MS_CONFIG, Long.MaxValue.toString)
-    producerProps.setProperty(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, brokerUrl)
-    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")
-    producerProps.setProperty(ProducerConfig.COMPRESSION_TYPE_CONFIG, compressionType)
-    val producer = new KafkaProducer[Array[Byte],Array[Byte]](producerProps)
-    val rand = new Random(1)
-    val keyCount = (messages / dups).toInt
-    val producedFile = File.createTempFile("kafka-log-cleaner-produced-", ".txt")
-    println("Logging produce requests to " + producedFile.getAbsolutePath)
-    val producedWriter = new BufferedWriter(new FileWriter(producedFile), 1024*1024)
-    for(i <- 0L until (messages * topics.length)) {
-      val topic = topics((i % topics.length).toInt)
-      val key = rand.nextInt(keyCount)
-      val delete = i % 100 < percentDeletes
-      val msg =
-        if(delete)
-          new ProducerRecord[Array[Byte],Array[Byte]](topic, key.toString.getBytes(), null)
-        else
-          new ProducerRecord[Array[Byte],Array[Byte]](topic, key.toString.getBytes(), i.toString.getBytes())
-      producer.send(msg)
-      producedWriter.write(TestRecord(topic, key, i, delete).toString)
-      producedWriter.newLine()
-    }
-    producedWriter.close()
-    producer.close()
-    producedFile
-  }
-
-  @deprecated("This method has been deprecated and will be removed in a future release.", "0.11.0.0")
-  def makeConsumer(zkUrl: String, topics: Array[String]): ZookeeperConsumerConnector = {
-    val consumerProps = new Properties
-    consumerProps.setProperty("group.id", "log-cleaner-test-" + new Random().nextInt(Int.MaxValue))
-    consumerProps.setProperty("zookeeper.connect", zkUrl)
-    consumerProps.setProperty("consumer.timeout.ms", (20*1000).toString)
-    consumerProps.setProperty("auto.offset.reset", "smallest")
-    new ZookeeperConsumerConnector(new ConsumerConfig(consumerProps))
-  }
-
-  def consumeMessages(zkUrl: String, topics: Array[String]): File = {
-    val connector = makeConsumer(zkUrl, topics)
-    val streams = connector.createMessageStreams(topics.map(topic => (topic, 1)).toMap, new StringDecoder, new StringDecoder)
-    val consumedFile = File.createTempFile("kafka-log-cleaner-consumed-", ".txt")
-    println("Logging consumed messages to " + consumedFile.getAbsolutePath)
-    val consumedWriter = new BufferedWriter(new FileWriter(consumedFile))
-    for(topic <- topics) {
-      val stream = streams(topic).head
-      try {
-        for(item <- stream) {
-          val delete = item.message == null
-          val value = if(delete) -1L else item.message.toLong
-          consumedWriter.write(TestRecord(topic, item.key.toInt, value, delete).toString)
-          consumedWriter.newLine()
-        }
-      } catch {
-        case _: ConsumerTimeoutException =>
-      }
-    }
-    consumedWriter.close()
-    connector.shutdown()
-    consumedFile
-  }
-
-}
-
-case class TestRecord(topic: String, key: Int, value: Long, delete: Boolean) {
-  def this(pieces: Array[String]) = this(pieces(0), pieces(1).toInt, pieces(2).toLong, pieces(3) == "d")
-  def this(line: String) = this(line.split("\t"))
-  override def toString = topic + "\t" +  key + "\t" + value + "\t" + (if(delete) "d" else "u")
-  def topicAndKey = topic + key
-}
diff --git a/core/src/test/scala/kafka/utils/ExitTest.scala b/core/src/test/scala/kafka/utils/ExitTest.scala
new file mode 100644
index 0000000000000..b95a721cab275
--- /dev/null
+++ b/core/src/test/scala/kafka/utils/ExitTest.scala
@@ -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 kafka.utils
+
+import java.io.IOException
+
+import org.junit.Assert.assertEquals
+import org.junit.Test
+import org.scalatest.Assertions.intercept
+
+class ExitTest {
+  @Test
+  def shouldHaltImmediately(): Unit = {
+    val array:Array[Any] = Array("a", "b")
+    def haltProcedure(exitStatus: Int, message: Option[String]) : Nothing = {
+      array(0) = exitStatus
+      array(1) = message
+      throw new IOException()
+    }
+    Exit.setHaltProcedure(haltProcedure)
+    val statusCode = 0
+    val message = Some("message")
+    try {
+      intercept[IOException] {Exit.halt(statusCode)}
+      assertEquals(statusCode, array(0))
+      assertEquals(None, array(1))
+
+      intercept[IOException] {Exit.halt(statusCode, message)}
+      assertEquals(statusCode, array(0))
+      assertEquals(message, array(1))
+    } finally {
+      Exit.resetHaltProcedure()
+    }
+  }
+
+  @Test
+  def shouldExitImmediately(): Unit = {
+    val array:Array[Any] = Array("a", "b")
+    def exitProcedure(exitStatus: Int, message: Option[String]) : Nothing = {
+      array(0) = exitStatus
+      array(1) = message
+      throw new IOException()
+    }
+    Exit.setExitProcedure(exitProcedure)
+    val statusCode = 0
+    val message = Some("message")
+    try {
+      intercept[IOException] {Exit.exit(statusCode)}
+      assertEquals(statusCode, array(0))
+      assertEquals(None, array(1))
+
+      intercept[IOException] {Exit.exit(statusCode, message)}
+      assertEquals(statusCode, array(0))
+      assertEquals(message, array(1))
+    } finally {
+      Exit.resetExitProcedure()
+    }
+  }
+
+  @Test
+  def shouldAddShutdownHookImmediately(): Unit = {
+    val name = "name"
+    val array:Array[Any] = Array("", 0)
+    // immediately invoke the shutdown hook to mutate the data when a hook is added
+    def shutdownHookAdder(name: String, shutdownHook: => Unit) : Unit = {
+      // mutate the first element
+      array(0) = array(0).toString + name
+      // invoke the shutdown hook (see below, it mutates the second element)
+      shutdownHook
+    }
+    Exit.setShutdownHookAdder(shutdownHookAdder)
+    def sideEffect(): Unit = {
+      // mutate the second element
+      array(1) = array(1).asInstanceOf[Int] + 1
+    }
+    try {
+      Exit.addShutdownHook(name, sideEffect()) // by-name parameter, only invoked due to above shutdownHookAdder
+      assertEquals(1, array(1))
+      assertEquals(name * array(1).asInstanceOf[Int], array(0).toString)
+      Exit.addShutdownHook(name, array(1) = array(1).asInstanceOf[Int] + 1) // by-name parameter, only invoked due to above shutdownHookAdder
+      assertEquals(2, array(1))
+      assertEquals(name * array(1).asInstanceOf[Int], array(0).toString)
+    } finally {
+      Exit.resetShutdownHookAdder()
+    }
+  }
+
+  @Test
+  def shouldNotInvokeShutdownHookImmediately(): Unit = {
+    val name = "name"
+    val array:Array[String] = Array(name)
+
+    def sideEffect(): Unit = {
+      // mutate the first element
+      array(0) = array(0) + name
+    }
+    Exit.addShutdownHook(name, sideEffect()) // by-name parameter, not invoked
+    // make sure the first element wasn't mutated
+    assertEquals(name, array(0))
+    Exit.addShutdownHook(name, sideEffect()) // by-name parameter, not invoked
+    // again make sure the first element wasn't mutated
+    assertEquals(name, array(0))
+    Exit.addShutdownHook(name, array(0) = array(0) + name) // by-name parameter, not invoked
+    // again make sure the first element wasn't mutated
+    assertEquals(name, array(0))
+  }
+}
diff --git a/core/src/test/scala/kafka/utils/LoggingTest.scala b/core/src/test/scala/kafka/utils/LoggingTest.scala
index c0600f8d03dc7..e339b15593a89 100644
--- a/core/src/test/scala/kafka/utils/LoggingTest.scala
+++ b/core/src/test/scala/kafka/utils/LoggingTest.scala
@@ -18,14 +18,23 @@
 package kafka.utils
 
 import java.lang.management.ManagementFactory
-import javax.management.ObjectName
 
+import javax.management.ObjectName
 import org.junit.Test
 import org.junit.Assert.{assertEquals, assertTrue}
+import org.slf4j.LoggerFactory
 
 
 class LoggingTest extends Logging {
 
+  @Test
+  def testTypeOfGetLoggers(): Unit = {
+    val log4jController = new Log4jController
+    // the return object of getLoggers must be a collection instance from java standard library.
+    // That enables mbean client to deserialize it without extra libraries.
+    assertEquals(classOf[java.util.ArrayList[String]], log4jController.getLoggers.getClass)
+  }
+
   @Test
   def testLog4jControllerIsRegistered(): Unit = {
     val mbs = ManagementFactory.getPlatformMBeanServer()
@@ -34,4 +43,46 @@ class LoggingTest extends Logging {
     val instance = mbs.getObjectInstance(log4jControllerName)
     assertEquals("kafka.utils.Log4jController", instance.getClassName)
   }
+
+  @Test
+  def testLogNameOverride(): Unit = {
+    class TestLogging(overriddenLogName: String) extends Logging {
+      // Expose logger
+      def log = logger
+      override def loggerName = overriddenLogName
+    }
+    val overriddenLogName = "OverriddenLogName"
+    val logging = new TestLogging(overriddenLogName)
+
+    assertEquals(overriddenLogName, logging.log.underlying.getName)
+  }
+
+  @Test
+  def testLogName(): Unit = {
+    class TestLogging extends Logging {
+      // Expose logger
+      def log = logger
+    }
+    val logging = new TestLogging
+
+    assertEquals(logging.getClass.getName, logging.log.underlying.getName)
+  }
+
+  @Test
+  def testLoggerLevelIsResolved(): Unit = {
+    val controller = new Log4jController()
+    val previousLevel = controller.getLogLevel("kafka")
+    try {
+      controller.setLogLevel("kafka", "TRACE")
+      // Do some logging so that the Logger is created within the hierarchy
+      // (until loggers are used only loggers in the config file exist)
+      LoggerFactory.getLogger("kafka.utils.Log4jControllerTest").trace("test")
+      assertEquals("TRACE", controller.getLogLevel("kafka"))
+      assertEquals("TRACE", controller.getLogLevel("kafka.utils.Log4jControllerTest"))
+      assertTrue(controller.getLoggers.contains("kafka=TRACE"))
+      assertTrue(controller.getLoggers.contains("kafka.utils.Log4jControllerTest=TRACE"))
+    } finally {
+      controller.setLogLevel("kafka", previousLevel)
+    }
+  }
 }
diff --git a/core/src/test/scala/kafka/utils/ToolsUtilsTest.scala b/core/src/test/scala/kafka/utils/ToolsUtilsTest.scala
new file mode 100644
index 0000000000000..edc114e51bfed
--- /dev/null
+++ b/core/src/test/scala/kafka/utils/ToolsUtilsTest.scala
@@ -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 kafka.utils
+
+import java.io.ByteArrayOutputStream
+import java.util.Collections
+
+import org.apache.kafka.common.MetricName
+import org.apache.kafka.common.metrics.Metrics
+import org.apache.kafka.common.metrics.internals.IntGaugeSuite
+import org.junit.Assert.assertTrue
+import org.junit.Test
+import org.slf4j.LoggerFactory
+
+import scala.jdk.CollectionConverters._
+
+class ToolsUtilsTest {
+  private val log = LoggerFactory.getLogger(classOf[ToolsUtilsTest])
+
+  @Test def testIntegerMetric(): Unit = {
+    val outContent = new ByteArrayOutputStream()
+    val metrics = new Metrics
+    val suite = new IntGaugeSuite[String](log, "example", metrics, (k: String) => new MetricName(k + "-bar", "test", "A test metric", Collections.singletonMap("key", "value")), 1)
+    suite.increment("foo")
+    Console.withOut(outContent) {
+      ToolsUtils.printMetrics(metrics.metrics.asScala)
+      assertTrue(outContent.toString.split("\n").exists(line => line.trim.matches("^test:foo-bar:\\{key=value\\}     : 1$")))
+    }
+  }
+
+}
diff --git a/core/src/test/scala/kafka/zk/ExtendedAclStoreTest.scala b/core/src/test/scala/kafka/zk/ExtendedAclStoreTest.scala
new file mode 100644
index 0000000000000..14e6a9d2957a9
--- /dev/null
+++ b/core/src/test/scala/kafka/zk/ExtendedAclStoreTest.scala
@@ -0,0 +1,68 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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.zk
+
+import org.apache.kafka.common.resource.PatternType.{LITERAL, PREFIXED}
+import org.apache.kafka.common.resource.ResourcePattern
+import org.apache.kafka.common.resource.ResourceType.TOPIC
+import org.junit.Assert.assertEquals
+import org.junit.Test
+
+class ExtendedAclStoreTest {
+  private val literalResource = new ResourcePattern(TOPIC, "some-topic", LITERAL)
+  private val prefixedResource = new ResourcePattern(TOPIC, "some-topic", PREFIXED)
+  private val store = new ExtendedAclStore(PREFIXED)
+
+  @Test
+  def shouldHaveCorrectPaths(): Unit = {
+    assertEquals("/kafka-acl-extended/prefixed", store.aclPath)
+    assertEquals("/kafka-acl-extended/prefixed/Topic", store.path(TOPIC))
+    assertEquals("/kafka-acl-extended-changes", store.changeStore.aclChangePath)
+  }
+
+  @Test
+  def shouldHaveCorrectPatternType(): Unit = {
+    assertEquals(PREFIXED, store.patternType)
+  }
+
+  @Test(expected = classOf[IllegalArgumentException])
+  def shouldThrowIfConstructedWithLiteral(): Unit = {
+    new ExtendedAclStore(LITERAL)
+  }
+
+  @Test(expected = classOf[IllegalArgumentException])
+  def shouldThrowFromEncodeOnLiteral(): Unit = {
+    store.changeStore.createChangeNode(literalResource)
+  }
+
+  @Test
+  def shouldWriteChangesToTheWritePath(): Unit = {
+    val changeNode = store.changeStore.createChangeNode(prefixedResource)
+
+    assertEquals("/kafka-acl-extended-changes/acl_changes_", changeNode.path)
+  }
+
+  @Test
+  def shouldRoundTripChangeNode(): Unit = {
+    val changeNode = store.changeStore.createChangeNode(prefixedResource)
+
+    val actual = store.changeStore.decode(changeNode.bytes)
+
+    assertEquals(prefixedResource, actual)
+  }
+}
\ No newline at end of file
diff --git a/core/src/test/scala/kafka/zk/FeatureZNodeTest.scala b/core/src/test/scala/kafka/zk/FeatureZNodeTest.scala
new file mode 100644
index 0000000000000..c6157d63943e2
--- /dev/null
+++ b/core/src/test/scala/kafka/zk/FeatureZNodeTest.scala
@@ -0,0 +1,123 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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.zk
+
+import java.nio.charset.StandardCharsets
+
+import org.apache.kafka.common.feature.{Features, FinalizedVersionRange}
+import org.apache.kafka.common.feature.Features._
+import org.junit.Assert.{assertEquals, assertThrows}
+import org.junit.Test
+
+import scala.jdk.CollectionConverters._
+
+class FeatureZNodeTest {
+
+  @Test
+  def testEncodeDecode(): Unit = {
+    val featureZNode = FeatureZNode(
+      FeatureZNodeStatus.Enabled,
+      Features.finalizedFeatures(
+        Map[String, FinalizedVersionRange](
+          "feature1" -> new FinalizedVersionRange(1, 2),
+          "feature2" -> new FinalizedVersionRange(2, 4)).asJava))
+    val decoded = FeatureZNode.decode(FeatureZNode.encode(featureZNode))
+    assertEquals(featureZNode, decoded)
+  }
+
+  @Test
+  def testDecodeSuccess(): Unit = {
+    val featureZNodeStrTemplate = """{
+      "version":1,
+      "status":1,
+      "features":%s
+    }"""
+
+    val validFeatures = """{"feature1": {"min_version_level": 1, "max_version_level": 2}, "feature2": {"min_version_level": 2, "max_version_level": 4}}"""
+    val node1 = FeatureZNode.decode(featureZNodeStrTemplate.format(validFeatures).getBytes(StandardCharsets.UTF_8))
+    assertEquals(FeatureZNodeStatus.Enabled, node1.status)
+    assertEquals(
+      Features.finalizedFeatures(
+        Map[String, FinalizedVersionRange](
+          "feature1" -> new FinalizedVersionRange(1, 2),
+          "feature2" -> new FinalizedVersionRange(2, 4)).asJava), node1.features)
+
+    val emptyFeatures = "{}"
+    val node2 = FeatureZNode.decode(featureZNodeStrTemplate.format(emptyFeatures).getBytes(StandardCharsets.UTF_8))
+    assertEquals(FeatureZNodeStatus.Enabled, node2.status)
+    assertEquals(emptyFinalizedFeatures, node2.features)
+  }
+
+  @Test
+  def testDecodeFailOnInvalidVersionAndStatus(): Unit = {
+    val featureZNodeStrTemplate =
+      """{
+      "version":%d,
+      "status":%d,
+      "features":{"feature1": {"min_version_level": 1, "max_version_level": 2}, "feature2": {"min_version_level": 2, "max_version_level": 4}}
+    }"""
+    assertThrows(
+      classOf[IllegalArgumentException],
+      () => FeatureZNode.decode(
+        featureZNodeStrTemplate.format(FeatureZNode.V1 - 1, 1).getBytes(StandardCharsets.UTF_8)))
+    val invalidStatus = FeatureZNodeStatus.Enabled.id + 1
+    assertThrows(
+      classOf[IllegalArgumentException],
+      () => FeatureZNode.decode(
+        featureZNodeStrTemplate.format(FeatureZNode.CurrentVersion, invalidStatus).getBytes(StandardCharsets.UTF_8)))
+  }
+
+  @Test
+  def testDecodeFailOnInvalidFeatures(): Unit = {
+    val featureZNodeStrTemplate =
+      """{
+      "version":1,
+      "status":1%s
+    }"""
+
+    val missingFeatures = ""
+    assertThrows(
+      classOf[IllegalArgumentException],
+      () => FeatureZNode.decode(
+        featureZNodeStrTemplate.format(missingFeatures).getBytes(StandardCharsets.UTF_8)))
+
+    val malformedFeatures = ""","features":{"feature1": {"min_version_level": 1, "max_version_level": 2}, "partial"}"""
+    assertThrows(
+      classOf[IllegalArgumentException],
+      () => FeatureZNode.decode(
+        featureZNodeStrTemplate.format(malformedFeatures).getBytes(StandardCharsets.UTF_8)))
+
+    val invalidFeaturesMinVersionLevel = ""","features":{"feature1": {"min_version_level": 0, "max_version_level": 2}}"""
+    assertThrows(
+      classOf[IllegalArgumentException],
+      () => FeatureZNode.decode(
+        featureZNodeStrTemplate.format(invalidFeaturesMinVersionLevel).getBytes(StandardCharsets.UTF_8)))
+
+    val invalidFeaturesMaxVersionLevel = ""","features":{"feature1": {"min_version_level": 2, "max_version_level": 1}}"""
+    assertThrows(
+      classOf[IllegalArgumentException],
+      () => FeatureZNode.decode(
+        featureZNodeStrTemplate.format(invalidFeaturesMaxVersionLevel).getBytes(StandardCharsets.UTF_8)))
+
+    val invalidFeaturesMissingMinVersionLevel = ""","features":{"feature1": {"max_version_level": 1}}"""
+    assertThrows(
+      classOf[IllegalArgumentException],
+      () => FeatureZNode.decode(
+        featureZNodeStrTemplate.format(invalidFeaturesMissingMinVersionLevel).getBytes(StandardCharsets.UTF_8)))
+  }
+}
diff --git a/core/src/test/scala/kafka/zk/LiteralAclStoreTest.scala b/core/src/test/scala/kafka/zk/LiteralAclStoreTest.scala
new file mode 100644
index 0000000000000..2f07cc64c5d46
--- /dev/null
+++ b/core/src/test/scala/kafka/zk/LiteralAclStoreTest.scala
@@ -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 kafka.zk
+
+import java.nio.charset.StandardCharsets.UTF_8
+
+import kafka.security.authorizer.AclEntry
+import org.apache.kafka.common.resource.PatternType.{LITERAL, PREFIXED}
+import org.apache.kafka.common.resource.ResourcePattern
+import org.apache.kafka.common.resource.ResourceType.{GROUP, TOPIC}
+import org.junit.Assert.assertEquals
+import org.junit.Test
+
+class LiteralAclStoreTest {
+  private val literalResource = new ResourcePattern(TOPIC, "some-topic", LITERAL)
+  private val prefixedResource = new ResourcePattern(TOPIC, "some-topic", PREFIXED)
+  private val store = LiteralAclStore
+
+  @Test
+  def shouldHaveCorrectPaths(): Unit = {
+    assertEquals("/kafka-acl", store.aclPath)
+    assertEquals("/kafka-acl/Topic", store.path(TOPIC))
+    assertEquals("/kafka-acl-changes", store.changeStore.aclChangePath)
+  }
+
+  @Test
+  def shouldHaveCorrectPatternType(): Unit = {
+    assertEquals(LITERAL, store.patternType)
+  }
+
+  @Test(expected = classOf[IllegalArgumentException])
+  def shouldThrowFromEncodeOnNoneLiteral(): Unit = {
+    store.changeStore.createChangeNode(prefixedResource)
+  }
+
+  @Test
+  def shouldWriteChangesToTheWritePath(): Unit = {
+    val changeNode = store.changeStore.createChangeNode(literalResource)
+
+    assertEquals("/kafka-acl-changes/acl_changes_", changeNode.path)
+  }
+
+  @Test
+  def shouldRoundTripChangeNode(): Unit = {
+    val changeNode = store.changeStore.createChangeNode(literalResource)
+
+    val actual = store.changeStore.decode(changeNode.bytes)
+
+    assertEquals(literalResource, actual)
+  }
+
+  @Test
+  def shouldDecodeResourceUsingTwoPartLogic(): Unit = {
+    val resource = new ResourcePattern(GROUP, "PREFIXED:this, including the PREFIXED part, is a valid two part group name", LITERAL)
+    val encoded = (resource.resourceType.toString + AclEntry.ResourceSeparator + resource.name).getBytes(UTF_8)
+
+    val actual = store.changeStore.decode(encoded)
+
+    assertEquals(resource, actual)
+  }
+}
diff --git a/core/src/test/scala/other/kafka/ReplicationQuotasTestRig.scala b/core/src/test/scala/other/kafka/ReplicationQuotasTestRig.scala
index 389cb8f0ac396..3f72afd240489 100644
--- a/core/src/test/scala/other/kafka/ReplicationQuotasTestRig.scala
+++ b/core/src/test/scala/other/kafka/ReplicationQuotasTestRig.scala
@@ -15,26 +15,29 @@
   * limitations under the License.
   */
 
-package other.kafka
+package kafka
 
-import java.io.{File, FileOutputStream, PrintWriter}
-import javax.imageio.ImageIO
+import java.io.{File, PrintWriter}
+import java.nio.charset.StandardCharsets
+import java.nio.file.{Files, StandardOpenOption}
 
+import javax.imageio.ImageIO
 import kafka.admin.ReassignPartitionsCommand
-import kafka.admin.ReassignPartitionsCommand.Throttle
-import kafka.common.TopicAndPartition
-import org.apache.kafka.common.TopicPartition
 import kafka.server.{KafkaConfig, KafkaServer, QuotaType}
 import kafka.utils.TestUtils._
-import kafka.utils.ZkUtils._
-import kafka.utils.{Exit, Logging, TestUtils, ZkUtils}
-import kafka.zk.ZooKeeperTestHarness
+import kafka.utils.{Exit, Logging, TestUtils}
+import kafka.zk.{ReassignPartitionsZNode, ZooKeeperTestHarness}
+import org.apache.kafka.clients.admin.{Admin, AdminClientConfig}
 import org.apache.kafka.clients.producer.ProducerRecord
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.network.ListenerName
+import org.apache.kafka.common.security.auth.SecurityProtocol
+import org.apache.kafka.common.utils.Utils
 import org.jfree.chart.plot.PlotOrientation
 import org.jfree.chart.{ChartFactory, ChartFrame, JFreeChart}
 import org.jfree.data.xy.{XYSeries, XYSeriesCollection}
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 import scala.collection.{Map, Seq, mutable}
 
 /**
@@ -74,10 +77,10 @@ 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
+      experiment.setUp()
       experiment.run(config, journal, displayChartsOnScreen)
       journal.footer()
     }
@@ -85,7 +88,7 @@ object ReplicationQuotasTestRig {
       case e: Exception => e.printStackTrace()
     }
     finally {
-      experiment.tearDown
+      experiment.tearDown()
     }
   }
 
@@ -100,23 +103,32 @@ object ReplicationQuotasTestRig {
     var servers: Seq[KafkaServer] = null
     val leaderRates = mutable.Map[Int, Array[Double]]()
     val followerRates = mutable.Map[Int, Array[Double]]()
+    var adminClient: Admin = null
 
-    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)))
+
+      TestUtils.waitUntilBrokerMetadataIsPropagated(servers)
+      val brokerList = TestUtils.bootstrapServers(servers,
+        ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT))
+      adminClient = Admin.create(Map[String, Object](
+        AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG -> brokerList
+      ).asJava)
     }
 
-    override def tearDown() {
+    override def tearDown(): Unit = {
+      Utils.closeQuietly(adminClient, "adminClient")
       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
-      val shift = Math.round(config.brokers / 2)
+      val shift = Math.round(config.brokers / 2f)
 
       def nextReplicaRoundRobin(): Int = {
         count = count + 1
@@ -125,21 +137,25 @@ object ReplicationQuotasTestRig {
       val replicas = (0 to config.partitions).map(partition => partition -> Seq(nextReplicaRoundRobin())).toMap
 
       startBrokers(brokers)
-      createTopic(zkUtils, topicName, replicas, servers)
+      createTopic(zkClient, topicName, replicas, servers)
 
       println("Writing Data")
-      val producer = TestUtils.createNewProducer(TestUtils.getBrokerListStrFromServers(servers), retries = 5, acks = 0)
+      val producer = TestUtils.createProducer(TestUtils.getBrokerListStrFromServers(servers), acks = 0)
       (0 until config.msgsPerPartition).foreach { x =>
         (0 until config.partitions).foreach { partition =>
           producer.send(new ProducerRecord(topicName, partition, null, new Array[Byte](config.msgSize)))
         }
       }
 
-      println("Starting Reassignment")
-      val newAssignment = ReassignPartitionsCommand.generateAssignment(zkUtils, brokers, json(topicName), true)._1
+      println("Generating Reassignment")
+      val (newAssignment, _) = ReassignPartitionsCommand.generateAssignment(adminClient,
+        json(topicName), brokers.mkString(","), true)
 
+      println("Starting Reassignment")
       val start = System.currentTimeMillis()
-      ReassignPartitionsCommand.executeAssignment(zkUtils, None, ZkUtils.formatAsReassignmentJson(newAssignment), Throttle(config.throttle))
+      ReassignPartitionsCommand.executeAssignment(adminClient, false,
+        new String(ReassignPartitionsZNode.encode(newAssignment), StandardCharsets.UTF_8),
+        config.throttle)
 
       //Await completion
       waitForReassignmentToComplete()
@@ -167,14 +183,14 @@ object ReplicationQuotasTestRig {
     }
   }
 
-    def logOutput(config: ExperimentDef, replicas: Map[Int, Seq[Int]], newAssignment: Map[TopicAndPartition, Seq[Int]]): Unit = {
-      val actual = zkUtils.getPartitionAssignmentForTopics(Seq(topicName))(topicName)
+    def logOutput(config: ExperimentDef, replicas: Map[Int, Seq[Int]], newAssignment: Map[TopicPartition, Seq[Int]]): Unit = {
+      val actual = zkClient.getPartitionAssignmentForTopics(Set(topicName))(topicName)
 
       //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.map { case (k, v) => k -> v.replicas })
       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.map { case (k, v) => k -> v.replicas })
 
       //Test Stats
       println(s"numBrokers: ${config.brokers}")
@@ -186,11 +202,11 @@ object ReplicationQuotasTestRig {
       println(s"Worst case duration is ${config.targetBytesPerBrokerMB * 1000 * 1000/ config.throttle}")
     }
 
-    def waitForReassignmentToComplete() {
+    def waitForReassignmentToComplete(): Unit = {
       waitUntilTrue(() => {
         printRateMetrics()
-        !zkUtils.pathExists(ReassignPartitionsPath)
-      }, s"Znode ${ZkUtils.ReassignPartitionsPath} wasn't deleted", 60 * 60 * 1000, pause = 1000L)
+        adminClient.listPartitionReassignments().reassignments().get().isEmpty
+      }, s"Partition reassignments didn't complete.", 60 * 60 * 1000, pause = 1000L)
     }
 
     def renderChart(data: mutable.Map[Int, Array[Double]], name: String, journal: Journal, displayChartsOnScreen: Boolean): Unit = {
@@ -247,7 +263,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)
@@ -266,7 +282,7 @@ object ReplicationQuotasTestRig {
     private def measuredRate(broker: KafkaServer, repType: QuotaType): Double = {
       val metricName = broker.metrics.metricName("byte-rate", repType.toString)
       if (broker.metrics.metrics.asScala.contains(metricName))
-        broker.metrics.metrics.asScala(metricName).value
+        broker.metrics.metrics.asScala(metricName).metricValue.asInstanceOf[Double]
       else -1
     }
 
@@ -286,7 +302,7 @@ object ReplicationQuotasTestRig {
       val message = s"\n\n
                      ${config.name}
                      " +
         s"
                      - BrokerCount: ${config.brokers}" +
         s"
                      - PartitionCount: ${config.partitions}" +
-        f"
                      - Throttle: ${config.throttle}%,.0f MB/s" +
+        f"
                      - Throttle: ${config.throttle.toDouble}%,.0f MB/s" +
         f"
                      - MsgCount: ${config.msgsPerPartition}%,.0f " +
         f"
                      - MsgSize: ${config.msgSize}%,.0f" +
         s"
                      - TargetBytesPerBrokerMB: ${config.targetBytesPerBrokerMB}
                      "
@@ -312,7 +328,7 @@ object ReplicationQuotasTestRig {
     }
 
     def append(message: String): Unit = {
-      val stream = new FileOutputStream(log, true)
+      val stream = Files.newOutputStream(log.toPath, StandardOpenOption.CREATE, StandardOpenOption.APPEND)
       new PrintWriter(stream) {
         append(message)
         close
diff --git a/core/src/test/scala/other/kafka/StressTestLog.scala b/core/src/test/scala/other/kafka/StressTestLog.scala
index 1710da7b9c3bf..b411519fd4513 100755
--- a/core/src/test/scala/other/kafka/StressTestLog.scala
+++ b/core/src/test/scala/other/kafka/StressTestLog.scala
@@ -21,11 +21,10 @@ 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
-import org.apache.kafka.common.requests.IsolationLevel
 import org.apache.kafka.common.utils.Utils
 
 /**
@@ -35,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()
@@ -58,57 +57,87 @@ object StressTestLog {
     val reader = new ReaderThread(log)
     reader.start()
 
-    Runtime.getRuntime().addShutdownHook(new Thread() {
-      override def run() = {
+    Exit.addShutdownHook("stress-test-shutdown-hook", {
         running.set(false)
         writer.join()
         reader.join()
         Utils.delete(dir)
-      }
     })
 
     while(running.get) {
-      println("Reader offset = %d, writer offset = %d".format(reader.offset, writer.offset))
       Thread.sleep(1000)
+      println("Reader offset = %d, writer offset = %d".format(reader.currentOffset, writer.currentOffset))
+      writer.checkProgress()
+      reader.checkProgress()
     }
   }
 
   abstract class WorkerThread extends Thread {
-    override def run() {
+    val threadInfo = "Thread: " + Thread.currentThread.getName + " Class: " + getClass.getName
+
+    override def run(): Unit = {
       try {
         while(running.get)
           work()
       } catch {
-        case e: Exception =>
+        case e: Exception => {
           e.printStackTrace()
-          running.set(false)
+        }
+      } finally {
+        running.set(false)
+      }
+    }
+
+    def work(): Unit
+    def isMakingProgress(): Boolean
+  }
+
+  trait LogProgress {
+    @volatile var currentOffset = 0
+    private var lastOffsetCheckpointed = currentOffset
+    private var lastProgressCheckTime = System.currentTimeMillis
+
+    def isMakingProgress(): Boolean = {
+      if (currentOffset > lastOffsetCheckpointed) {
+        lastOffsetCheckpointed = currentOffset
+        return true
+      }
+
+      false
+    }
+
+    def checkProgress(): Unit = {
+      // Check if we are making progress every 500ms
+      val curTime = System.currentTimeMillis
+      if ((curTime - lastProgressCheckTime) > 500) {
+        require(isMakingProgress(), "Thread not making progress")
+        lastProgressCheckTime = curTime
       }
-      println(getClass.getName + " exiting...")
     }
-    def work()
   }
 
-  class WriterThread(val log: Log) extends WorkerThread {
-    @volatile var offset = 0
-    override def work() {
-      val logAppendInfo = log.appendAsFollower(TestUtils.singletonRecords(offset.toString.getBytes))
-      require(logAppendInfo.firstOffset == offset && logAppendInfo.lastOffset == offset)
-      offset += 1
-      if(offset % 1000 == 0)
-        Thread.sleep(500)
+  class WriterThread(val log: Log) extends WorkerThread with LogProgress {
+    override def work(): Unit = {
+      val logAppendInfo = log.appendAsLeader(TestUtils.singletonRecords(currentOffset.toString.getBytes), 0)
+      require(logAppendInfo.firstOffset.forall(_ == currentOffset) && logAppendInfo.lastOffset == currentOffset)
+      currentOffset += 1
+      if (currentOffset % 1000 == 0)
+        Thread.sleep(50)
     }
   }
 
-  class ReaderThread(val log: Log) extends WorkerThread {
-    @volatile var offset = 0
-    override def work() {
+  class ReaderThread(val log: Log) extends WorkerThread with LogProgress {
+    override def work(): Unit = {
       try {
-        log.read(offset, 1024, Some(offset+1), isolationLevel = IsolationLevel.READ_UNCOMMITTED).records match {
+        log.read(currentOffset,
+          maxLength = 1,
+          isolation = FetchLogEnd,
+          minOneMessage = true).records match {
           case read: FileRecords if read.sizeInBytes > 0 => {
             val first = read.batches.iterator.next()
-            require(first.lastOffset == offset, "We should either read nothing or the message we asked for.")
+            require(first.lastOffset == currentOffset, "We should either read nothing or the message we asked for.")
             require(first.sizeInBytes == read.sizeInBytes, "Expected %d but got %d.".format(first.sizeInBytes, read.sizeInBytes))
-            offset += 1
+            currentOffset += 1
           }
           case _ =>
         }
diff --git a/core/src/test/scala/other/kafka/TestCrcPerformance.scala b/core/src/test/scala/other/kafka/TestCrcPerformance.scala
deleted file mode 100755
index daeecbd303044..0000000000000
--- a/core/src/test/scala/other/kafka/TestCrcPerformance.scala
+++ /dev/null
@@ -1,44 +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.log
-
-import java.util.Random
-import kafka.message._
-import org.apache.kafka.common.utils.Utils
-
-object TestCrcPerformance {
-
-  def main(args: Array[String]): Unit = {
-    if(args.length < 2)
-      Utils.croak("USAGE: java " + getClass().getName() + " num_messages message_size")
-    val numMessages = args(0).toInt
-    val messageSize = args(1).toInt
-
-    val content = new Array[Byte](messageSize)
-    new Random(1).nextBytes(content)
-
-    // create message test
-    val start = System.nanoTime
-    for (_ <- 0 until numMessages)
-      new Message(content)
-
-    val elapsed = System.nanoTime - start
-    println("%d messages created in %.2f seconds + (%.2f ns per message).".format(numMessages, elapsed / (1000.0*1000.0*1000.0),
-      elapsed / numMessages.toDouble))
-
-  }
-}
diff --git a/core/src/test/scala/other/kafka/TestKafkaAppender.scala b/core/src/test/scala/other/kafka/TestKafkaAppender.scala
deleted file mode 100644
index ecdcac056fe15..0000000000000
--- a/core/src/test/scala/other/kafka/TestKafkaAppender.scala
+++ /dev/null
@@ -1,50 +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
-
-import org.apache.log4j.PropertyConfigurator
-import kafka.utils.{Exit, Logging}
-import serializer.Encoder
-
-object TestKafkaAppender extends Logging {
-  
-  def main(args:Array[String]) {
-    
-    if(args.length < 1) {
-      println("USAGE: " + TestKafkaAppender.getClass.getName + " log4j_config")
-      Exit.exit(1)
-    }
-
-    try {
-      PropertyConfigurator.configure(args(0))
-    } catch {
-      case e: Exception =>
-        System.err.println("KafkaAppender could not be initialized ! Exiting..")
-        e.printStackTrace()
-        Exit.exit(1)
-    }
-
-    for (_ <- 1 to 10)
-      info("test")    
-  }
-}
-
-class AppenderStringSerializer(encoding: String = "UTF-8") extends Encoder[AnyRef] {
-  def toBytes(event: AnyRef): Array[Byte] = event.toString.getBytes(encoding)
-}
-
diff --git a/core/src/test/scala/other/kafka/TestLinearWriteSpeed.scala b/core/src/test/scala/other/kafka/TestLinearWriteSpeed.scala
index 16325ee4ec53f..24a49c32d50cf 100755
--- a/core/src/test/scala/other/kafka/TestLinearWriteSpeed.scala
+++ b/core/src/test/scala/other/kafka/TestLinearWriteSpeed.scala
@@ -20,6 +20,7 @@ package kafka
 import java.io._
 import java.nio._
 import java.nio.channels._
+import java.nio.file.StandardOpenOption
 import java.util.{Properties, Random}
 
 import joptsimple._
@@ -102,7 +103,7 @@ object TestLinearWriteSpeed {
     val compressionCodec = CompressionCodec.getCompressionCodec(options.valueOf(compressionCodecOpt))
     val rand = new Random
     rand.nextBytes(buffer.array)
-    val numMessages = bufferSize / (messageSize + MessageSet.LogOverhead)
+    val numMessages = bufferSize / (messageSize + Records.LOG_OVERHEAD)
     val createTime = System.currentTimeMillis
     val messageSet = {
       val compressionType = CompressionType.forId(compressionCodec.codec)
@@ -167,13 +168,13 @@ object TestLinearWriteSpeed {
       }
     }
     val elapsedSecs = (System.nanoTime - beginTest) / (1000.0*1000.0*1000.0)
-    println(bytesToWrite / (1024.0 * 1024.0 * elapsedSecs) + " MB per sec")
+    println((bytesToWrite / (1024.0 * 1024.0 * elapsedSecs)).toString + " MB per sec")
     scheduler.shutdown()
   }
 
   trait Writable {
     def write(): Int
-    def close()
+    def close(): Unit
   }
 
   class MmapWritable(val file: File, size: Long, val content: ByteBuffer) extends Writable {
@@ -186,22 +187,24 @@ object TestLinearWriteSpeed {
       content.rewind()
       content.limit()
     }
-    def close() {
+    def close(): Unit = {
       raf.close()
+      Utils.delete(file)
     }
   }
 
   class ChannelWritable(val file: File, val content: ByteBuffer) extends Writable {
     file.deleteOnExit()
-    val raf = new RandomAccessFile(file, "rw")
-    val channel = raf.getChannel
+    val channel = FileChannel.open(file.toPath, StandardOpenOption.CREATE, StandardOpenOption.READ,
+      StandardOpenOption.WRITE)
     def write(): Int = {
       channel.write(content)
       content.rewind()
       content.limit()
     }
-    def close() {
-      raf.close()
+    def close(): Unit = {
+      channel.close()
+      Utils.delete(file)
     }
   }
 
@@ -213,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/other/kafka/TestOffsetManager.scala b/core/src/test/scala/other/kafka/TestOffsetManager.scala
deleted file mode 100644
index c8f9397353e4d..0000000000000
--- a/core/src/test/scala/other/kafka/TestOffsetManager.scala
+++ /dev/null
@@ -1,309 +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 other.kafka
-
-import kafka.api._
-import kafka.utils.{Exit, ShutdownableThread, ZkUtils}
-import org.apache.kafka.common.protocol.Errors
-
-import scala.collection._
-import kafka.client.ClientUtils
-import joptsimple.OptionParser
-import kafka.common.{OffsetAndMetadata, TopicAndPartition}
-import kafka.network.BlockingChannel
-
-import scala.util.Random
-import java.io.IOException
-
-import kafka.metrics.{KafkaMetricsGroup, KafkaTimer}
-import java.util.concurrent.TimeUnit
-import java.util.concurrent.atomic.AtomicInteger
-import java.nio.channels.ClosedByInterruptException
-
-
-object TestOffsetManager {
-
-  val random = new Random
-  val SocketTimeoutMs = 10000
-
-  class StatsThread(reportingIntervalMs: Long, commitThreads: Seq[CommitThread], fetchThread: FetchThread)
-        extends ShutdownableThread("stats-thread") {
-
-    def printStats() {
-      println("--------------------------------------------------------------------------------")
-      println("Aggregate stats for commits:")
-      println("Error count: %d; Max:%f; Min: %f; Mean: %f; Commit count: %d".format(
-        commitThreads.map(_.numErrors.get).sum,
-        commitThreads.map(_.timer.max()).max,
-        commitThreads.map(_.timer.min()).min,
-        commitThreads.map(_.timer.mean()).sum / commitThreads.size,
-        commitThreads.map(_.numCommits.get).sum))
-      println("--------------------------------------------------------------------------------")
-      commitThreads.foreach(t => println(t.stats))
-      println(fetchThread.stats)
-    }
-
-    override def doWork() {
-      printStats()
-      Thread.sleep(reportingIntervalMs)
-    }
-
-  }
-
-  class CommitThread(id: Int, partitionCount: Int, commitIntervalMs: Long, zkUtils: ZkUtils)
-        extends ShutdownableThread("commit-thread")
-        with KafkaMetricsGroup {
-
-    private val groupId = "group-" + id
-    private val metadata = "Metadata from commit thread " + id
-    private var offsetsChannel = ClientUtils.channelToOffsetManager(groupId, zkUtils, SocketTimeoutMs)
-    private var offset = 0L
-    val numErrors = new AtomicInteger(0)
-    val numCommits = new AtomicInteger(0)
-    val timer = newTimer("commit-thread", TimeUnit.MILLISECONDS, TimeUnit.SECONDS)
-    private val commitTimer = new KafkaTimer(timer)
-    val shutdownLock = new Object
-
-    private def ensureConnected() {
-      if (!offsetsChannel.isConnected)
-        offsetsChannel = ClientUtils.channelToOffsetManager(groupId, zkUtils, SocketTimeoutMs)
-    }
-
-    override def doWork() {
-      val commitRequest = OffsetCommitRequest(groupId, immutable.Map((1 to partitionCount).map(TopicAndPartition("topic-" + id, _) -> OffsetAndMetadata(offset, metadata)):_*))
-      try {
-        ensureConnected()
-        offsetsChannel.send(commitRequest)
-        numCommits.getAndIncrement
-        commitTimer.time {
-          val response = OffsetCommitResponse.readFrom(offsetsChannel.receive().payload())
-          if (response.commitStatus.exists(_._2 != Errors.NONE)) numErrors.getAndIncrement
-        }
-        offset += 1
-      }
-      catch {
-        case _: ClosedByInterruptException =>
-          offsetsChannel.disconnect()
-        case e2: IOException =>
-          println("Commit thread %d: Error while committing offsets to %s:%d for group %s due to %s.".format(id, offsetsChannel.host, offsetsChannel.port, groupId, e2))
-          offsetsChannel.disconnect()
-      }
-      finally {
-        Thread.sleep(commitIntervalMs)
-      }
-    }
-
-    override def shutdown() {
-      super.shutdown()
-      awaitShutdown()
-      offsetsChannel.disconnect()
-      println("Commit thread %d ended. Last committed offset: %d.".format(id, offset))
-    }
-
-    def stats = {
-      "Commit thread %d :: Error count: %d; Max:%f; Min: %f; Mean: %f; Commit count: %d"
-      .format(id, numErrors.get(), timer.max(), timer.min(), timer.mean(), numCommits.get())
-    }
-  }
-
-  class FetchThread(numGroups: Int, fetchIntervalMs: Long, zkUtils: ZkUtils)
-        extends ShutdownableThread("fetch-thread")
-        with KafkaMetricsGroup {
-
-    private val timer = newTimer("fetch-thread", TimeUnit.MILLISECONDS, TimeUnit.SECONDS)
-    private val fetchTimer = new KafkaTimer(timer)
-
-    private val channels = mutable.Map[Int, BlockingChannel]()
-    private var metadataChannel = ClientUtils.channelToAnyBroker(zkUtils, SocketTimeoutMs)
-
-    private val numErrors = new AtomicInteger(0)
-
-    override def doWork() {
-      val id = random.nextInt().abs % numGroups
-      val group = "group-" + id
-      try {
-        metadataChannel.send(GroupCoordinatorRequest(group))
-        val coordinatorId = GroupCoordinatorResponse.readFrom(metadataChannel.receive().payload()).coordinatorOpt.map(_.id).getOrElse(-1)
-
-        val channel = if (channels.contains(coordinatorId))
-          channels(coordinatorId)
-        else {
-          val newChannel = ClientUtils.channelToOffsetManager(group, zkUtils, SocketTimeoutMs)
-          channels.put(coordinatorId, newChannel)
-          newChannel
-        }
-
-        try {
-          // send the offset fetch request
-          val fetchRequest = OffsetFetchRequest(group, Seq(TopicAndPartition("topic-"+id, 1)))
-          channel.send(fetchRequest)
-
-          fetchTimer.time {
-            val response = OffsetFetchResponse.readFrom(channel.receive().payload())
-            if (response.requestInfo.exists(_._2.error != Errors.NONE)) {
-              numErrors.getAndIncrement
-            }
-          }
-        }
-        catch {
-          case _: ClosedByInterruptException =>
-            channel.disconnect()
-            channels.remove(coordinatorId)
-          case e2: IOException =>
-            println("Error while fetching offset from %s:%d due to %s.".format(channel.host, channel.port, e2))
-            channel.disconnect()
-            channels.remove(coordinatorId)
-        }
-      }
-      catch {
-        case _: IOException =>
-          println("Error while querying %s:%d - shutting down query channel.".format(metadataChannel.host, metadataChannel.port))
-          metadataChannel.disconnect()
-          println("Creating new query channel.")
-          metadataChannel = ClientUtils.channelToAnyBroker(zkUtils, SocketTimeoutMs)
-      }
-      finally {
-        Thread.sleep(fetchIntervalMs)
-      }
-
-    }
-
-    override def shutdown() {
-      super.shutdown()
-      awaitShutdown()
-      channels.foreach(_._2.disconnect())
-      metadataChannel.disconnect()
-    }
-
-    def stats = {
-      "Fetch thread :: Error count: %d; Max:%f; Min: %f; Mean: %f; Fetch count: %d"
-      .format(numErrors.get(), timer.max(), timer.min(), timer.mean(), timer.count())
-    }
-  }
-
-  def main(args: Array[String]) {
-    val parser = new OptionParser(false)
-    val zookeeperOpt = parser.accepts("zookeeper", "The ZooKeeper connection URL.")
-      .withRequiredArg
-      .describedAs("ZooKeeper URL")
-      .ofType(classOf[java.lang.String])
-      .defaultsTo("localhost:2181")
-
-    val commitIntervalOpt = parser.accepts("commit-interval-ms", "Offset commit interval.")
-      .withRequiredArg
-      .describedAs("interval")
-      .ofType(classOf[java.lang.Integer])
-      .defaultsTo(100)
-
-    val fetchIntervalOpt = parser.accepts("fetch-interval-ms", "Offset fetch interval.")
-      .withRequiredArg
-      .describedAs("interval")
-      .ofType(classOf[java.lang.Integer])
-      .defaultsTo(1000)
-
-    val numPartitionsOpt = parser.accepts("partition-count", "Number of partitions per commit.")
-      .withRequiredArg
-      .describedAs("interval")
-      .ofType(classOf[java.lang.Integer])
-      .defaultsTo(1)
-
-    val numThreadsOpt = parser.accepts("thread-count", "Number of commit threads.")
-      .withRequiredArg
-      .describedAs("threads")
-      .ofType(classOf[java.lang.Integer])
-      .defaultsTo(1)
-
-    val reportingIntervalOpt = parser.accepts("reporting-interval-ms", "Interval at which stats are reported.")
-      .withRequiredArg
-      .describedAs("interval (ms)")
-      .ofType(classOf[java.lang.Integer])
-      .defaultsTo(3000)
-
-    val helpOpt = parser.accepts("help", "Print this message.")
-
-    val options = parser.parse(args : _*)
-
-    if (options.has(helpOpt)) {
-      parser.printHelpOn(System.out)
-      Exit.exit(0)
-    }
-
-    val commitIntervalMs = options.valueOf(commitIntervalOpt).intValue()
-    val fetchIntervalMs = options.valueOf(fetchIntervalOpt).intValue()
-    val threadCount = options.valueOf(numThreadsOpt).intValue()
-    val partitionCount = options.valueOf(numPartitionsOpt).intValue()
-    val zookeeper = options.valueOf(zookeeperOpt)
-    val reportingIntervalMs = options.valueOf(reportingIntervalOpt).intValue()
-    println("Commit thread count: %d; Partition count: %d, Commit interval: %d ms; Fetch interval: %d ms; Reporting interval: %d ms"
-            .format(threadCount, partitionCount, commitIntervalMs, fetchIntervalMs, reportingIntervalMs))
-
-    var zkUtils: ZkUtils = null
-    var commitThreads: Seq[CommitThread] = Seq()
-    var fetchThread: FetchThread = null
-    var statsThread: StatsThread = null
-    try {
-      zkUtils = ZkUtils(zookeeper, 6000, 2000, false)
-      commitThreads = (0 until threadCount).map { threadId =>
-        new CommitThread(threadId, partitionCount, commitIntervalMs, zkUtils)
-      }
-
-      fetchThread = new FetchThread(threadCount, fetchIntervalMs, zkUtils)
-      statsThread = new StatsThread(reportingIntervalMs, commitThreads, fetchThread)
-
-      Runtime.getRuntime.addShutdownHook(new Thread() {
-        override def run() {
-          cleanShutdown()
-          statsThread.printStats()
-        }
-      })
-
-      commitThreads.foreach(_.start())
-
-      fetchThread.start()
-
-      statsThread.start()
-
-      commitThreads.foreach(_.join())
-      fetchThread.join()
-      statsThread.join()
-    }
-    catch {
-      case e: Throwable =>
-        println("Error: ", e)
-    }
-    finally {
-      cleanShutdown()
-    }
-
-    def cleanShutdown() {
-      commitThreads.foreach(_.shutdown())
-      commitThreads.foreach(_.join())
-      if (fetchThread != null) {
-        fetchThread.shutdown()
-        fetchThread.join()
-      }
-      if (statsThread != null) {
-        statsThread.shutdown()
-        statsThread.join()
-      }
-      zkUtils.close()
-    }
-
-  }
-}
-
diff --git a/core/src/test/scala/other/kafka/TestPurgatoryPerformance.scala b/core/src/test/scala/other/kafka/TestPurgatoryPerformance.scala
index 4b540a4f1a829..507c3ff0973ef 100644
--- a/core/src/test/scala/other/kafka/TestPurgatoryPerformance.scala
+++ b/core/src/test/scala/other/kafka/TestPurgatoryPerformance.scala
@@ -28,7 +28,7 @@ import kafka.utils._
 import org.apache.kafka.common.utils.Time
 
 import scala.math._
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 
 /**
  * This is a benchmark test of the purgatory.
diff --git a/core/src/test/scala/other/kafka/TestTruncate.scala b/core/src/test/scala/other/kafka/TestTruncate.scala
index e90d57ded6e0a..ba9c799cb7bed 100644
--- a/core/src/test/scala/other/kafka/TestTruncate.scala
+++ b/core/src/test/scala/other/kafka/TestTruncate.scala
@@ -19,6 +19,8 @@ package kafka
 
 import java.io._
 import java.nio._
+import java.nio.channels.FileChannel
+import java.nio.file.StandardOpenOption
 
 /* This code tests the correct function of java's FileChannel.truncate--some platforms don't work. */
 object TestTruncate {
@@ -26,7 +28,7 @@ object TestTruncate {
   def main(args: Array[String]): Unit = {
     val name = File.createTempFile("kafka", ".test")
     name.deleteOnExit()
-    val file = new RandomAccessFile(name, "rw").getChannel()
+    val file = FileChannel.open(name.toPath, StandardOpenOption.READ, StandardOpenOption.WRITE)
     val buffer = ByteBuffer.allocate(12)
     buffer.putInt(4).putInt(4).putInt(4)
     buffer.rewind()
diff --git a/core/src/test/scala/unit/kafka/KafkaConfigTest.scala b/core/src/test/scala/unit/kafka/KafkaConfigTest.scala
index c934c4a142809..e0d058115e774 100644
--- a/core/src/test/scala/unit/kafka/KafkaConfigTest.scala
+++ b/core/src/test/scala/unit/kafka/KafkaConfigTest.scala
@@ -16,7 +16,8 @@
  */
 package kafka
 
-import java.io.{File, FileOutputStream}
+import java.io.File
+import java.nio.file.Files
 import java.util
 
 import kafka.server.KafkaConfig
@@ -25,6 +26,9 @@ import org.apache.kafka.common.config.types.Password
 import org.apache.kafka.common.internals.FatalExitError
 import org.junit.{After, Before, Test}
 import org.junit.Assert._
+import org.apache.kafka.common.config.internals.BrokerSecurityConfigs
+
+import scala.jdk.CollectionConverters._
 
 class KafkaTest {
 
@@ -57,12 +61,6 @@ class KafkaTest {
     assertEquals(util.Arrays.asList("compact","delete"), config4.logCleanupPolicy)
   }
 
-  @Test(expected = classOf[FatalExitError])
-  def testGetKafkaConfigFromArgsWrongSetValue(): Unit = {
-    val propertiesFile = prepareDefaultConfig()
-    KafkaConfig.fromProps(Kafka.getPropsFromArgs(Array(propertiesFile, "--override", "a=b=c")))
-  }
-
   @Test(expected = classOf[FatalExitError])
   def testGetKafkaConfigFromArgsNonArgsAtTheEnd(): Unit = {
     val propertiesFile = prepareDefaultConfig()
@@ -86,14 +84,210 @@ class KafkaTest {
     val propertiesFile = prepareDefaultConfig()
     val config = KafkaConfig.fromProps(Kafka.getPropsFromArgs(Array(propertiesFile, "--override", "ssl.keystore.password=keystore_password",
                                                                                     "--override", "ssl.key.password=key_password",
-                                                                                    "--override", "ssl.truststore.password=truststore_password")))
-    assertEquals(Password.HIDDEN, config.sslKeyPassword.toString)
-    assertEquals(Password.HIDDEN, config.sslKeystorePassword.toString)
-    assertEquals(Password.HIDDEN, config.sslTruststorePassword.toString)
+                                                                                    "--override", "ssl.truststore.password=truststore_password",
+                                                                                    "--override", "ssl.keystore.certificate.chain=certificate_chain",
+                                                                                    "--override", "ssl.keystore.key=private_key",
+                                                                                    "--override", "ssl.truststore.certificates=truststore_certificates")))
+    assertEquals(Password.HIDDEN, config.getPassword(KafkaConfig.SslKeyPasswordProp).toString)
+    assertEquals(Password.HIDDEN, config.getPassword(KafkaConfig.SslKeystorePasswordProp).toString)
+    assertEquals(Password.HIDDEN, config.getPassword(KafkaConfig.SslTruststorePasswordProp).toString)
+    assertEquals(Password.HIDDEN, config.getPassword(KafkaConfig.SslKeystoreKeyProp).toString)
+    assertEquals(Password.HIDDEN, config.getPassword(KafkaConfig.SslKeystoreCertificateChainProp).toString)
+    assertEquals(Password.HIDDEN, config.getPassword(KafkaConfig.SslTruststoreCertificatesProp).toString)
+
+    assertEquals("key_password", config.getPassword(KafkaConfig.SslKeyPasswordProp).value)
+    assertEquals("keystore_password", config.getPassword(KafkaConfig.SslKeystorePasswordProp).value)
+    assertEquals("truststore_password", config.getPassword(KafkaConfig.SslTruststorePasswordProp).value)
+    assertEquals("private_key", config.getPassword(KafkaConfig.SslKeystoreKeyProp).value)
+    assertEquals("certificate_chain", config.getPassword(KafkaConfig.SslKeystoreCertificateChainProp).value)
+    assertEquals("truststore_certificates", config.getPassword(KafkaConfig.SslTruststoreCertificatesProp).value)
+  }
+
+  @Test
+  def testKafkaSslPasswordsWithSymbols(): Unit = {
+    val password = "=!#-+!?*/\"\'^%$=\\.,@:;="
+    val propertiesFile = prepareDefaultConfig()
+    val config = KafkaConfig.fromProps(Kafka.getPropsFromArgs(Array(propertiesFile,
+      "--override", "ssl.keystore.password=" + password,
+      "--override", "ssl.key.password=" + password,
+      "--override", "ssl.truststore.password=" + password)))
+    assertEquals(Password.HIDDEN, config.getPassword(KafkaConfig.SslKeyPasswordProp).toString)
+    assertEquals(Password.HIDDEN, config.getPassword(KafkaConfig.SslKeystorePasswordProp).toString)
+    assertEquals(Password.HIDDEN, config.getPassword(KafkaConfig.SslTruststorePasswordProp).toString)
+
+    assertEquals(password, config.getPassword(KafkaConfig.SslKeystorePasswordProp).value)
+    assertEquals(password, config.getPassword(KafkaConfig.SslKeyPasswordProp).value)
+    assertEquals(password, config.getPassword(KafkaConfig.SslTruststorePasswordProp).value)
+  }
+
+  private val booleanPropValueToSet = true
+  private val stringPropValueToSet = "foo"
+  private val passwordPropValueToSet = "ThePa$$word!"
+  private val listPropValueToSet = List("A", "B")
+
+  @Test
+  def testZkSslClientEnable(): Unit = {
+    testZkConfig(KafkaConfig.ZkSslClientEnableProp, "zookeeper.ssl.client.enable",
+      "zookeeper.client.secure", booleanPropValueToSet, config => Some(config.zkSslClientEnable), booleanPropValueToSet, Some(false))
+  }
+
+  @Test
+  def testZkSslKeyStoreLocation(): Unit = {
+    testZkConfig(KafkaConfig.ZkSslKeyStoreLocationProp, "zookeeper.ssl.keystore.location",
+      "zookeeper.ssl.keyStore.location", stringPropValueToSet, config => config.zkSslKeyStoreLocation, stringPropValueToSet)
+  }
+
+  @Test
+  def testZkSslTrustStoreLocation(): Unit = {
+    testZkConfig(KafkaConfig.ZkSslTrustStoreLocationProp, "zookeeper.ssl.truststore.location",
+      "zookeeper.ssl.trustStore.location", stringPropValueToSet, config => config.zkSslTrustStoreLocation, stringPropValueToSet)
+  }
 
-    assertEquals("key_password", config.sslKeyPassword.value)
-    assertEquals("keystore_password", config.sslKeystorePassword.value)
-    assertEquals("truststore_password", config.sslTruststorePassword.value)
+  @Test
+  def testZookeeperKeyStorePassword(): Unit = {
+    testZkConfig(KafkaConfig.ZkSslKeyStorePasswordProp, "zookeeper.ssl.keystore.password",
+      "zookeeper.ssl.keyStore.password", passwordPropValueToSet, config => config.zkSslKeyStorePassword, new Password(passwordPropValueToSet))
+  }
+
+  @Test
+  def testZookeeperTrustStorePassword(): Unit = {
+    testZkConfig(KafkaConfig.ZkSslTrustStorePasswordProp, "zookeeper.ssl.truststore.password",
+      "zookeeper.ssl.trustStore.password", passwordPropValueToSet, config => config.zkSslTrustStorePassword, new Password(passwordPropValueToSet))
+  }
+
+  @Test
+  def testZkSslKeyStoreType(): Unit = {
+    testZkConfig(KafkaConfig.ZkSslKeyStoreTypeProp, "zookeeper.ssl.keystore.type",
+      "zookeeper.ssl.keyStore.type", stringPropValueToSet, config => config.zkSslKeyStoreType, stringPropValueToSet)
+  }
+
+  @Test
+  def testZkSslTrustStoreType(): Unit = {
+    testZkConfig(KafkaConfig.ZkSslTrustStoreTypeProp, "zookeeper.ssl.truststore.type",
+      "zookeeper.ssl.trustStore.type", stringPropValueToSet, config => config.zkSslTrustStoreType, stringPropValueToSet)
+  }
+
+  @Test
+  def testZkSslProtocol(): Unit = {
+    testZkConfig(KafkaConfig.ZkSslProtocolProp, "zookeeper.ssl.protocol",
+      "zookeeper.ssl.protocol", stringPropValueToSet, config => Some(config.ZkSslProtocol), stringPropValueToSet, Some("TLSv1.2"))
+  }
+
+  @Test
+  def testZkSslEnabledProtocols(): Unit = {
+    testZkConfig(KafkaConfig.ZkSslEnabledProtocolsProp, "zookeeper.ssl.enabled.protocols",
+      "zookeeper.ssl.enabledProtocols", listPropValueToSet.mkString(","), config => config.ZkSslEnabledProtocols, listPropValueToSet.asJava)
+  }
+
+  @Test
+  def testZkSslCipherSuites(): Unit = {
+    testZkConfig(KafkaConfig.ZkSslCipherSuitesProp, "zookeeper.ssl.cipher.suites",
+      "zookeeper.ssl.ciphersuites", listPropValueToSet.mkString(","), config => config.ZkSslCipherSuites, listPropValueToSet.asJava)
+  }
+
+  @Test
+  def testZkSslEndpointIdentificationAlgorithm(): Unit = {
+    // this property is different than the others
+    // because the system property values and the Kafka property values don't match
+    val kafkaPropName = KafkaConfig.ZkSslEndpointIdentificationAlgorithmProp
+    assertEquals("zookeeper.ssl.endpoint.identification.algorithm", kafkaPropName)
+    val sysProp = "zookeeper.ssl.hostnameVerification"
+    val expectedDefaultValue = "HTTPS"
+    val propertiesFile = prepareDefaultConfig()
+    // first make sure there is the correct default value
+    val emptyConfig = KafkaConfig.fromProps(Kafka.getPropsFromArgs(Array(propertiesFile)))
+    assertNull(emptyConfig.originals.get(kafkaPropName)) // doesn't appear in the originals
+    assertEquals(expectedDefaultValue, emptyConfig.values.get(kafkaPropName)) // but default value appears in the values
+    assertEquals(expectedDefaultValue, emptyConfig.ZkSslEndpointIdentificationAlgorithm) // and has the correct default value
+    // next set system property alone
+    Map("true" -> "HTTPS", "false" -> "").foreach { case (sysPropValue, expected) => {
+      try {
+        System.setProperty(sysProp, sysPropValue)
+        val config = KafkaConfig.fromProps(Kafka.getPropsFromArgs(Array(propertiesFile)))
+        assertNull(config.originals.get(kafkaPropName)) // doesn't appear in the originals
+        assertEquals(expectedDefaultValue, config.values.get(kafkaPropName)) // default value appears in the values
+        assertEquals(expected, config.ZkSslEndpointIdentificationAlgorithm) // system property impacts the ultimate value of the property
+      } finally {
+        System.clearProperty(sysProp)
+      }
+    }}
+    // finally set Kafka config alone
+    List("https", "").foreach(expected => {
+      val config = KafkaConfig.fromProps(Kafka.getPropsFromArgs(Array(propertiesFile, "--override", s"$kafkaPropName=${expected}")))
+      assertEquals(expected, config.originals.get(kafkaPropName)) // appears in the originals
+      assertEquals(expected, config.values.get(kafkaPropName)) // appears in the values
+      assertEquals(expected, config.ZkSslEndpointIdentificationAlgorithm) // is the ultimate value
+    })
+  }
+
+  @Test
+  def testZkSslCrlEnable(): Unit = {
+    testZkConfig(KafkaConfig.ZkSslCrlEnableProp, "zookeeper.ssl.crl.enable",
+      "zookeeper.ssl.crl", booleanPropValueToSet, config => Some(config.ZkSslCrlEnable), booleanPropValueToSet, Some(false))
+  }
+
+  @Test
+  def testZkSslOcspEnable(): Unit = {
+    testZkConfig(KafkaConfig.ZkSslOcspEnableProp, "zookeeper.ssl.ocsp.enable",
+      "zookeeper.ssl.ocsp", booleanPropValueToSet, config => Some(config.ZkSslOcspEnable), booleanPropValueToSet, Some(false))
+  }
+
+  @Test
+  def testConnectionsMaxReauthMsDefault(): Unit = {
+    val propertiesFile = prepareDefaultConfig()
+    val config = KafkaConfig.fromProps(Kafka.getPropsFromArgs(Array(propertiesFile)))
+    assertEquals(0L, config.valuesWithPrefixOverride("sasl_ssl.oauthbearer.").get(BrokerSecurityConfigs.CONNECTIONS_MAX_REAUTH_MS).asInstanceOf[Long])
+  }
+
+  @Test
+  def testConnectionsMaxReauthMsExplicit(): Unit = {
+    val propertiesFile = prepareDefaultConfig()
+    val expected = 3600000
+    val config = KafkaConfig.fromProps(Kafka.getPropsFromArgs(Array(propertiesFile, "--override", s"sasl_ssl.oauthbearer.connections.max.reauth.ms=${expected}")))
+    assertEquals(expected, config.valuesWithPrefixOverride("sasl_ssl.oauthbearer.").get(BrokerSecurityConfigs.CONNECTIONS_MAX_REAUTH_MS).asInstanceOf[Long])
+  }
+
+  private def testZkConfig[T, U](kafkaPropName: String,
+                                 expectedKafkaPropName: String,
+                                 sysPropName: String,
+                                 propValueToSet: T,
+                                 getPropValueFrom: (KafkaConfig) => Option[T],
+                                 expectedPropertyValue: U,
+                                 expectedDefaultValue: Option[T] = None): Unit = {
+    assertEquals(expectedKafkaPropName, kafkaPropName)
+    val propertiesFile = prepareDefaultConfig()
+    // first make sure there is the correct default value (if any)
+    val emptyConfig = KafkaConfig.fromProps(Kafka.getPropsFromArgs(Array(propertiesFile)))
+    assertNull(emptyConfig.originals.get(kafkaPropName)) // doesn't appear in the originals
+    if (expectedDefaultValue.isDefined) {
+      // confirm default value behavior
+      assertEquals(expectedDefaultValue.get, emptyConfig.values.get(kafkaPropName)) // default value appears in the values
+      assertEquals(expectedDefaultValue.get, getPropValueFrom(emptyConfig).get) // default value appears in the property
+    } else {
+      // confirm no default value behavior
+      assertNull(emptyConfig.values.get(kafkaPropName)) // doesn't appear in the values
+      assertEquals(None, getPropValueFrom(emptyConfig)) // has no default value
+    }
+    // next set system property alone
+    try {
+      System.setProperty(sysPropName, s"$propValueToSet")
+      // need to create a new Kafka config for the system property to be recognized
+      val config = KafkaConfig.fromProps(Kafka.getPropsFromArgs(Array(propertiesFile)))
+      assertNull(config.originals.get(kafkaPropName)) // doesn't appear in the originals
+      // confirm default value (if any) overridden by system property
+      if (expectedDefaultValue.isDefined)
+        assertEquals(expectedDefaultValue.get, config.values.get(kafkaPropName)) // default value (different from system property) appears in the values
+      else
+        assertNull(config.values.get(kafkaPropName)) // doesn't appear in the values
+      // confirm system property appears in the property
+      assertEquals(Some(expectedPropertyValue), getPropValueFrom(config))
+    } finally {
+      System.clearProperty(sysPropName)
+    }
+    // finally set Kafka config alone
+    val config = KafkaConfig.fromProps(Kafka.getPropsFromArgs(Array(propertiesFile, "--override", s"$kafkaPropName=${propValueToSet}")))
+    assertEquals(expectedPropertyValue, config.values.get(kafkaPropName)) // appears in the values
+    assertEquals(Some(expectedPropertyValue), getPropValueFrom(config)) // appears in the property
   }
 
   def prepareDefaultConfig(): String = {
@@ -104,14 +298,13 @@ class KafkaTest {
     val file = File.createTempFile("kafkatest", ".properties")
     file.deleteOnExit()
 
-    val writer = new FileOutputStream(file)
-    lines.foreach { l =>
-      writer.write(l.getBytes)
-      writer.write("\n".getBytes)
-    }
-
-    writer.close
-
-    file.getAbsolutePath
+    val writer = Files.newOutputStream(file.toPath)
+    try {
+      lines.foreach { l =>
+        writer.write(l.getBytes)
+        writer.write("\n".getBytes)
+      }
+      file.getAbsolutePath
+    } finally writer.close()
   }
 }
diff --git a/core/src/test/scala/unit/kafka/admin/AclCommandTest.scala b/core/src/test/scala/unit/kafka/admin/AclCommandTest.scala
index ecf3427ff4d51..fdb5d1635a263 100644
--- a/core/src/test/scala/unit/kafka/admin/AclCommandTest.scala
+++ b/core/src/test/scala/unit/kafka/admin/AclCommandTest.scala
@@ -16,138 +16,317 @@
  */
 package kafka.admin
 
+import java.io.{File, PrintWriter}
 import java.util.Properties
 
+import javax.management.InstanceAlreadyExistsException
 import kafka.admin.AclCommand.AclCommandOptions
-import kafka.security.auth._
-import kafka.server.KafkaConfig
-import kafka.utils.{Logging, TestUtils}
+import kafka.security.authorizer.{AclAuthorizer, AclEntry}
+import kafka.server.{KafkaConfig, KafkaServer}
+import kafka.utils.{Exit, LogCaptureAppender, Logging, TestUtils}
 import kafka.zk.ZooKeeperTestHarness
-import org.apache.kafka.common.security.auth.KafkaPrincipal
-import org.junit.Test
+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.{AppInfoParser, SecurityUtils}
+import org.apache.kafka.server.authorizer.Authorizer
+import org.apache.log4j.Level
+import org.junit.Assert.assertFalse
+import org.junit.{After, Assert, Before, Test}
+import org.scalatest.Assertions.intercept
 
 class AclCommandTest extends ZooKeeperTestHarness with Logging {
 
-  private val Users = Set(KafkaPrincipal.fromString("User:CN=writeuser,OU=Unknown,O=Unknown,L=Unknown,ST=Unknown,C=Unknown"), KafkaPrincipal.fromString("User:test2"))
+  var servers: Seq[KafkaServer] = Seq()
+
+  private val principal: KafkaPrincipal = SecurityUtils.parseKafkaPrincipal("User:test2")
+  private val Users = Set(SecurityUtils.parseKafkaPrincipal("User:CN=writeuser,OU=Unknown,O=Unknown,L=Unknown,ST=Unknown,C=Unknown"),
+    principal, SecurityUtils.parseKafkaPrincipal("""User:CN=\#User with special chars in CN : (\, \+ \" \\ \< \> \; ')"""))
   private val Hosts = Set("host1", "host2")
   private val AllowHostCommand = Array("--allow-host", "host1", "--allow-host", "host2")
   private val DenyHostCommand = Array("--deny-host", "host1", "--deny-host", "host2")
 
-  private val TopicResources = Set(new Resource(Topic, "test-1"), new Resource(Topic, "test-2"))
-  private val GroupResources = Set(new Resource(Group, "testGroup-1"), new Resource(Group, "testGroup-2"))
-  private val TransactionalIdResources = Set(new Resource(TransactionalId, "t0"), new Resource(TransactionalId, "t1"))
+  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")
+    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, Describe, Delete, DescribeConfigs, AlterConfigs),
-      Array("--operation", "Read" , "--operation", "Write", "--operation", "Describe", "--operation", "Delete",
-        "--operation", "DescribeConfigs", "--operation", "AlterConfigs")),
-    Set(Resource.ClusterResource) -> (Set(Create, ClusterAction, DescribeConfigs, AlterConfigs, IdempotentWrite),
+  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(ClusterResource) -> (Set(CREATE, CLUSTER_ACTION, DESCRIBE_CONFIGS, ALTER_CONFIGS, IDEMPOTENT_WRITE, ALTER, DESCRIBE),
       Array("--operation", "Create", "--operation", "ClusterAction", "--operation", "DescribeConfigs",
-        "--operation", "AlterConfigs", "--operation", "IdempotentWrite")),
-    GroupResources -> (Set(Read, Describe), Array("--operation", "Read", "--operation", "Describe")),
-    TransactionalIdResources -> (Set(Describe, Write), Array("--operation", "Describe", "--operation", "Write"))
+        "--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"))
   )
 
-  private def ProducerResourceToAcls(enableIdempotence: Boolean = false) = Map[Set[Resource], Set[Acl]](
-    TopicResources -> AclCommand.getAcls(Users, Allow, Set(Write, Describe), Hosts),
-    TransactionalIdResources -> AclCommand.getAcls(Users, Allow, Set(Write, Describe), Hosts),
-    Set(Resource.ClusterResource) -> AclCommand.getAcls(Users, Allow, Set(Some(Create),
-      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 = _
+  private var zkArgs: Array[String] = _
+  private var adminArgs: Array[String] = _
+
+  @Before
+  override def setUp(): Unit = {
+    super.setUp()
+
+    brokerProps = TestUtils.createBrokerConfig(0, zkConnect)
+    brokerProps.put(KafkaConfig.AuthorizerClassNameProp, classOf[AclAuthorizer].getName)
+    brokerProps.put(AclAuthorizer.SuperUsersProp, "User:ANONYMOUS")
+
+    zkArgs = Array("--authorizer-properties", "zookeeper.connect=" + zkConnect)
+  }
+
+  @After
+  override def tearDown(): Unit = {
+    TestUtils.shutdownServers(servers)
+    super.tearDown()
+  }
+
   @Test
-  def testAclCli() {
-    val brokerProps = TestUtils.createBrokerConfig(0, zkConnect)
-    brokerProps.put(KafkaConfig.AuthorizerClassNameProp, "kafka.security.auth.SimpleAclAuthorizer")
-    val args = Array("--authorizer-properties", "zookeeper.connect=" + zkConnect)
+  def testAclCliWithAuthorizer(): Unit = {
+    testAclCli(zkArgs)
+  }
+
+  @Test
+  def testAclCliWithAdminAPI(): Unit = {
+    createServer()
+    testAclCli(adminArgs)
+  }
+
+  private def createServer(commandConfig: Option[File] = None): Unit = {
+    servers = Seq(TestUtils.createServer(KafkaConfig.fromProps(brokerProps)))
+    val listenerName = ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT)
+
+    var adminArgs = Array("--bootstrap-server", TestUtils.bootstrapServers(servers, listenerName))
+    if (commandConfig.isDefined) {
+      adminArgs ++= Array("--command-config", commandConfig.get.getAbsolutePath)
+    }
+    this.adminArgs = adminArgs
+  }
 
+  private def callMain(args: Array[String]): (String, String) = {
+    TestUtils.grabConsoleOutputAndError(AclCommand.main(args))
+  }
+
+  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(args ++ cmd ++ resourceCmd ++ operationToCmd._2 :+ "--add")
-          for (resource <- resources) {
-            withAuthorizer(brokerProps) { authorizer =>
-              TestUtils.waitAndVerifyAcls(acls, authorizer, resource)
-            }
+        val (addOut, addErr) = callMain(cmdArgs ++ cmd ++ resourceCmd ++ operationToCmd._2 :+ "--add")
+        assertOutputContains("Adding ACLs", resources, resourceCmd, addOut)
+        assertOutputContains("Current ACLs", resources, resourceCmd, addOut)
+        Assert.assertEquals("", addErr)
+
+        for (resource <- resources) {
+          withAuthorizer() { authorizer =>
+            TestUtils.waitAndVerifyAcls(acls, authorizer, resource)
           }
+        }
 
-          testRemove(resources, resourceCmd, args, brokerProps)
+        val (listOut, listErr) = callMain(cmdArgs :+ "--list")
+        assertOutputContains("Current ACLs", resources, resourceCmd, listOut)
+        Assert.assertEquals("", listErr)
+
+        testRemove(cmdArgs, resources, resourceCmd)
       }
     }
   }
 
+  private def assertOutputContains(prefix: String, resources: Set[ResourcePattern], resourceCmd: Array[String], output: String): Unit = {
+    resources.foreach { resource =>
+      val resourceType = resource.resourceType.toString
+      (if (resource == ClusterResource) Array("kafka-cluster") else resourceCmd.filter(!_.startsWith("--"))).foreach { name =>
+        val expected = s"$prefix for resource `ResourcePattern(resourceType=$resourceType, name=$name, patternType=LITERAL)`:"
+        Assert.assertTrue(s"Substring ${expected} not in output:\n$output",
+          output.contains(expected))
+      }
+    }
+  }
+
+  @Test
+  def testProducerConsumerCliWithAuthorizer(): Unit = {
+    testProducerConsumerCli(zkArgs)
+  }
+
+  @Test
+  def testProducerConsumerCliWithAdminAPI(): Unit = {
+    createServer()
+    testProducerConsumerCli(adminArgs)
+  }
+
   @Test
-  def testProducerConsumerCli() {
-    val brokerProps = TestUtils.createBrokerConfig(0, zkConnect)
-    brokerProps.put(KafkaConfig.AuthorizerClassNameProp, "kafka.security.auth.SimpleAclAuthorizer")
-    val args = Array("--authorizer-properties", "zookeeper.connect=" + zkConnect)
+  def testAclCliWithClientId(): Unit = {
+    val adminClientConfig = TestUtils.tempFile()
+    val pw = new PrintWriter(adminClientConfig)
+    pw.println("client.id=my-client")
+    pw.close()
+
+    createServer(Some(adminClientConfig))
+
+    val appender = LogCaptureAppender.createAndRegister()
+    val previousLevel = LogCaptureAppender.setClassLoggerLevel(classOf[AppInfoParser], Level.WARN)
+    try {
+        testAclCli(adminArgs)
+    } finally {
+      LogCaptureAppender.setClassLoggerLevel(classOf[AppInfoParser], previousLevel)
+      LogCaptureAppender.unregister(appender)
+    }
+    val warning = appender.getMessages.find(e => e.getLevel == Level.WARN &&
+      e.getThrowableInformation != null &&
+      e.getThrowableInformation.getThrowable.getClass.getName == classOf[InstanceAlreadyExistsException].getName)
+    assertFalse("There should be no warnings about multiple registration of mbeans", warning.isDefined)
+
+  }
 
+  private def testProducerConsumerCli(cmdArgs: Array[String]): Unit = {
     for ((cmd, resourcesToAcls) <- CmdToResourcesToAcl) {
       val resourceCommand: Array[String] = resourcesToAcls.keys.map(ResourceToCommand).foldLeft(Array[String]())(_ ++ _)
-      AclCommand.main(args ++ getCmd(Allow) ++ resourceCommand ++ cmd :+ "--add")
+      callMain(cmdArgs ++ getCmd(ALLOW) ++ resourceCommand ++ cmd :+ "--add")
       for ((resources, acls) <- resourcesToAcls) {
         for (resource <- resources) {
-          withAuthorizer(brokerProps) { authorizer =>
+          withAuthorizer() { authorizer =>
             TestUtils.waitAndVerifyAcls(acls, authorizer, resource)
           }
         }
       }
-      testRemove(resourcesToAcls.keys.flatten.toSet, resourceCommand ++ cmd, args, brokerProps)
+      testRemove(cmdArgs, resourcesToAcls.keys.flatten.toSet, resourceCommand ++ cmd)
+    }
+  }
+
+  @Test
+  def testAclsOnPrefixedResourcesWithAuthorizer(): Unit = {
+    testAclsOnPrefixedResources(zkArgs)
+  }
+
+  @Test
+  def testAclsOnPrefixedResourcesWithAdminAPI(): Unit = {
+    createServer()
+    testAclsOnPrefixedResources(adminArgs)
+  }
+
+  private def testAclsOnPrefixedResources(cmdArgs: Array[String]): Unit = {
+    val cmd = Array("--allow-principal", principal.toString, "--producer", "--topic", "Test-", "--resource-pattern-type", "Prefixed")
+
+    callMain(cmdArgs ++ cmd :+ "--add")
+
+    withAuthorizer() { authorizer =>
+      val writeAcl = new AccessControlEntry(principal.toString, AclEntry.WildcardHost, WRITE, ALLOW)
+      val describeAcl = new AccessControlEntry(principal.toString, AclEntry.WildcardHost, DESCRIBE, ALLOW)
+      val createAcl = new AccessControlEntry(principal.toString, AclEntry.WildcardHost, CREATE, ALLOW)
+      TestUtils.waitAndVerifyAcls(Set(writeAcl, describeAcl, createAcl), authorizer,
+        new ResourcePattern(TOPIC, "Test-", PREFIXED))
+    }
+
+    callMain(cmdArgs ++ cmd :+ "--remove" :+ "--force")
+
+    withAuthorizer() { authorizer =>
+      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)
-    AclCommand.withAuthorizer(new AclCommandOptions(args))(null)
+    val aclCommandService = new AclCommand.AuthorizerService(classOf[AclAuthorizer].getName,
+      new AclCommandOptions(args))
+    aclCommandService.listAcls()
+  }
+
+  @Test
+  def testPatternTypes(): Unit = {
+    Exit.setExitProcedure { (status, _) =>
+      if (status == 1)
+        throw new RuntimeException("Exiting command")
+      else
+        throw new AssertionError(s"Unexpected exit with status $status")
+    }
+    def verifyPatternType(cmd: Array[String], isValid: Boolean): Unit = {
+      if (isValid)
+        callMain(cmd)
+      else
+        intercept[RuntimeException](callMain(cmd))
+    }
+    try {
+      PatternType.values.foreach { patternType =>
+        val addCmd = zkArgs ++ Array("--allow-principal", principal.toString, "--producer", "--topic", "Test",
+          "--add", "--resource-pattern-type", patternType.toString)
+        verifyPatternType(addCmd, isValid = patternType.isSpecific)
+        val listCmd = zkArgs ++ Array("--topic", "Test", "--list", "--resource-pattern-type", patternType.toString)
+        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(resources: Set[Resource], resourceCmd: Array[String], args: Array[String], brokerProps: Properties) {
+  private def testRemove(cmdArgs: Array[String], resources: Set[ResourcePattern], resourceCmd: Array[String]): Unit = {
+    val (out, err) = callMain(cmdArgs ++ resourceCmd :+ "--remove" :+ "--force")
+    Assert.assertEquals("", out)
+    Assert.assertEquals("", err)
     for (resource <- resources) {
-      AclCommand.main(args ++ resourceCmd :+ "--remove" :+ "--force")
-      withAuthorizer(brokerProps) { authorizer =>
-        TestUtils.waitAndVerifyAcls(Set.empty[Acl], authorizer, resource)
+      withAuthorizer() { authorizer =>
+        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))
   }
 
-  def withAuthorizer(props: Properties)(f: Authorizer => Unit) {
-    val kafkaConfig = KafkaConfig.fromProps(props, doLog = false)
-    val authZ = new SimpleAclAuthorizer
+  private def withAuthorizer()(f: Authorizer => Unit): Unit = {
+    val kafkaConfig = KafkaConfig.fromProps(brokerProps, doLog = false)
+    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 d79ba321d940c..b5a497759a522 100755
--- a/core/src/test/scala/unit/kafka/admin/AddPartitionsTest.scala
+++ b/core/src/test/scala/unit/kafka/admin/AddPartitionsTest.scala
@@ -17,57 +17,46 @@
 
 package kafka.admin
 
-import kafka.api.TopicMetadata
-import org.junit.Assert._
-import kafka.zk.ZooKeeperTestHarness
-import kafka.utils.TestUtils._
+import java.util.Optional
+
+import kafka.controller.ReplicaAssignment
+import kafka.server.BaseRequestTest
 import kafka.utils.TestUtils
-import kafka.cluster.Broker
-import kafka.client.ClientUtils
-import kafka.server.{KafkaConfig, KafkaServer}
+import kafka.utils.TestUtils._
+import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.errors.InvalidReplicaAssignmentException
-import org.apache.kafka.common.network.ListenerName
-import org.apache.kafka.common.security.auth.SecurityProtocol
-import org.junit.{After, Before, Test}
+import org.apache.kafka.common.requests.MetadataResponse.TopicMetadata
+import org.apache.kafka.common.requests.{MetadataRequest, MetadataResponse}
+import org.junit.Assert._
+import org.junit.{Before, Test}
 
-class AddPartitionsTest extends ZooKeeperTestHarness {
-  var configs: Seq[KafkaConfig] = null
-  var servers: Seq[KafkaServer] = Seq.empty[KafkaServer]
-  var brokers: Seq[Broker] = Seq.empty[Broker]
+import scala.jdk.CollectionConverters._
+
+class AddPartitionsTest extends BaseRequestTest {
+
+  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()
 
-    configs = (0 until 4).map(i => KafkaConfig.fromProps(TestUtils.createBrokerConfig(i, zkConnect, enableControlledShutdown = false)))
-    // start all the servers
-    servers = configs.map(c => TestUtils.createServer(c))
-    brokers = servers.map(s => TestUtils.createBroker(s.config.brokerId, s.config.hostName, TestUtils.boundPort(s)))
-
-    // create topics first
-    createTopic(zkUtils, topic1, partitionReplicaAssignment = topic1Assignment, servers = servers)
-    createTopic(zkUtils, topic2, partitionReplicaAssignment = topic2Assignment, servers = servers)
-    createTopic(zkUtils, topic3, partitionReplicaAssignment = topic3Assignment, servers = servers)
-    createTopic(zkUtils, topic4, partitionReplicaAssignment = topic4Assignment, servers = servers)
-  }
-
-  @After
-  override def tearDown() {
-    TestUtils.shutdownServers(servers)
-    super.tearDown()
+    createTopic(topic1, partitionReplicaAssignment = topic1Assignment.map { case (k, v) => k -> v.replicas })
+    createTopic(topic2, partitionReplicaAssignment = topic2Assignment.map { case (k, v) => k -> v.replicas })
+    createTopic(topic3, partitionReplicaAssignment = topic3Assignment.map { case (k, v) => k -> v.replicas })
+    createTopic(topic4, partitionReplicaAssignment = topic4Assignment.map { case (k, v) => k -> v.replicas })
   }
 
   @Test
@@ -97,27 +86,31 @@ class AddPartitionsTest extends ZooKeeperTestHarness {
   def testIncrementPartitions(): Unit = {
     adminZkClient.addPartitions(topic1, topic1Assignment, adminZkClient.getBrokerMetadatas(), 3)
     // wait until leader is elected
-    val leader1 = waitUntilLeaderIsElectedOrChanged(zkUtils, topic1, 1)
-    val leader2 = waitUntilLeaderIsElectedOrChanged(zkUtils, topic1, 2)
-    val leader1FromZk = zkUtils.getLeaderForPartition(topic1, 1).get
-    val leader2FromZk = zkUtils.getLeaderForPartition(topic1, 2).get
+    val leader1 = waitUntilLeaderIsElectedOrChanged(zkClient, topic1, 1)
+    val leader2 = waitUntilLeaderIsElectedOrChanged(zkClient, topic1, 2)
+    val leader1FromZk = zkClient.getLeaderForPartition(new TopicPartition(topic1, 1)).get
+    val leader2FromZk = zkClient.getLeaderForPartition(new TopicPartition(topic1, 2)).get
     assertEquals(leader1, leader1FromZk)
     assertEquals(leader2, leader2FromZk)
 
     // read metadata from a broker and verify the new topic partitions exist
     TestUtils.waitUntilMetadataIsPropagated(servers, topic1, 1)
     TestUtils.waitUntilMetadataIsPropagated(servers, topic1, 2)
-    val listenerName = ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT)
-    val metadata = ClientUtils.fetchTopicMetadata(Set(topic1), brokers.map(_.getBrokerEndPoint(listenerName)),
-      "AddPartitionsTest-testIncrementPartitions", 2000, 0).topicsMetadata
-    val metaDataForTopic1 = metadata.filter(p => p.topic.equals(topic1))
-    val partitionDataForTopic1 = metaDataForTopic1.head.partitionsMetadata.sortBy(_.partitionId)
-    assertEquals(partitionDataForTopic1.size, 3)
-    assertEquals(partitionDataForTopic1(1).partitionId, 1)
-    assertEquals(partitionDataForTopic1(2).partitionId, 2)
-    val replicas = partitionDataForTopic1(1).replicas
-    assertEquals(replicas.size, 2)
-    assert(replicas.contains(partitionDataForTopic1(1).leader.get))
+    val response = connectAndReceive[MetadataResponse](
+      new MetadataRequest.Builder(Seq(topic1).asJava, false).build)
+    assertEquals(1, response.topicMetadata.size)
+    val partitions = response.topicMetadata.asScala.head.partitionMetadata.asScala.sortBy(_.partition)
+    assertEquals(partitions.size, 3)
+    assertEquals(1, partitions(1).partition)
+    assertEquals(2, partitions(2).partition)
+
+    for (partition <- partitions) {
+      val replicas = partition.replicaIds
+      assertEquals(2, replicas.size)
+      assertTrue(partition.leaderId.isPresent)
+      val leaderId = partition.leaderId.get
+      assertTrue(replicas.contains(leaderId))
+    }
   }
 
   @Test
@@ -126,28 +119,28 @@ class AddPartitionsTest extends ZooKeeperTestHarness {
     adminZkClient.addPartitions(topic2, topic2Assignment, adminZkClient.getBrokerMetadatas(), 3,
       Some(Map(0 -> Seq(1, 2), 1 -> Seq(0, 1), 2 -> Seq(2, 3))))
     // wait until leader is elected
-    val leader1 = waitUntilLeaderIsElectedOrChanged(zkUtils, topic2, 1)
-    val leader2 = waitUntilLeaderIsElectedOrChanged(zkUtils, topic2, 2)
-    val leader1FromZk = zkUtils.getLeaderForPartition(topic2, 1).get
-    val leader2FromZk = zkUtils.getLeaderForPartition(topic2, 2).get
+    val leader1 = waitUntilLeaderIsElectedOrChanged(zkClient, topic2, 1)
+    val leader2 = waitUntilLeaderIsElectedOrChanged(zkClient, topic2, 2)
+    val leader1FromZk = zkClient.getLeaderForPartition(new TopicPartition(topic2, 1)).get
+    val leader2FromZk = zkClient.getLeaderForPartition(new TopicPartition(topic2, 2)).get
     assertEquals(leader1, leader1FromZk)
     assertEquals(leader2, leader2FromZk)
 
     // read metadata from a broker and verify the new topic partitions exist
     TestUtils.waitUntilMetadataIsPropagated(servers, topic2, 1)
     TestUtils.waitUntilMetadataIsPropagated(servers, topic2, 2)
-    val metadata = ClientUtils.fetchTopicMetadata(Set(topic2),
-      brokers.map(_.getBrokerEndPoint(ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT))),
-      "AddPartitionsTest-testManualAssignmentOfReplicas", 2000, 0).topicsMetadata
-    val metaDataForTopic2 = metadata.filter(_.topic == topic2)
-    val partitionDataForTopic2 = metaDataForTopic2.head.partitionsMetadata.sortBy(_.partitionId)
-    assertEquals(3, partitionDataForTopic2.size)
-    assertEquals(1, partitionDataForTopic2(1).partitionId)
-    assertEquals(2, partitionDataForTopic2(2).partitionId)
-    val replicas = partitionDataForTopic2(1).replicas
+    val response = connectAndReceive[MetadataResponse](
+      new MetadataRequest.Builder(Seq(topic2).asJava, false).build)
+    assertEquals(1, response.topicMetadata.size)
+    val topicMetadata = response.topicMetadata.asScala.head
+    val partitionMetadata = topicMetadata.partitionMetadata.asScala.sortBy(_.partition)
+    assertEquals(3, topicMetadata.partitionMetadata.size)
+    assertEquals(0, partitionMetadata(0).partition)
+    assertEquals(1, partitionMetadata(1).partition)
+    assertEquals(2, partitionMetadata(2).partition)
+    val replicas = partitionMetadata(1).replicaIds
     assertEquals(2, replicas.size)
-    assertTrue(replicas.head.id == 0 || replicas.head.id == 1)
-    assertTrue(replicas(1).id == 0 || replicas(1).id == 1)
+    assertEquals(Set(0, 1), replicas.asScala.toSet)
   }
 
   @Test
@@ -162,19 +155,17 @@ class AddPartitionsTest extends ZooKeeperTestHarness {
     TestUtils.waitUntilMetadataIsPropagated(servers, topic3, 5)
     TestUtils.waitUntilMetadataIsPropagated(servers, topic3, 6)
 
-    val metadata = ClientUtils.fetchTopicMetadata(Set(topic3),
-      brokers.map(_.getBrokerEndPoint(ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT))),
-      "AddPartitionsTest-testReplicaPlacementAllServers", 2000, 0).topicsMetadata
-
-    val metaDataForTopic3 = metadata.find(p => p.topic == topic3).get
-
-    validateLeaderAndReplicas(metaDataForTopic3, 0, 2, Set(2, 3, 0, 1))
-    validateLeaderAndReplicas(metaDataForTopic3, 1, 3, Set(3, 2, 0, 1))
-    validateLeaderAndReplicas(metaDataForTopic3, 2, 0, Set(0, 3, 1, 2))
-    validateLeaderAndReplicas(metaDataForTopic3, 3, 1, Set(1, 0, 2, 3))
-    validateLeaderAndReplicas(metaDataForTopic3, 4, 2, Set(2, 3, 0, 1))
-    validateLeaderAndReplicas(metaDataForTopic3, 5, 3, Set(3, 0, 1, 2))
-    validateLeaderAndReplicas(metaDataForTopic3, 6, 0, Set(0, 1, 2, 3))
+    val response = connectAndReceive[MetadataResponse](
+      new MetadataRequest.Builder(Seq(topic3).asJava, false).build)
+    assertEquals(1, response.topicMetadata.size)
+    val topicMetadata = response.topicMetadata.asScala.head
+    validateLeaderAndReplicas(topicMetadata, 0, 2, Set(2, 3, 0, 1))
+    validateLeaderAndReplicas(topicMetadata, 1, 3, Set(3, 2, 0, 1))
+    validateLeaderAndReplicas(topicMetadata, 2, 0, Set(0, 3, 1, 2))
+    validateLeaderAndReplicas(topicMetadata, 3, 1, Set(1, 0, 2, 3))
+    validateLeaderAndReplicas(topicMetadata, 4, 2, Set(2, 3, 0, 1))
+    validateLeaderAndReplicas(topicMetadata, 5, 3, Set(3, 0, 1, 2))
+    validateLeaderAndReplicas(topicMetadata, 6, 0, Set(0, 1, 2, 3))
   }
 
   @Test
@@ -185,25 +176,23 @@ class AddPartitionsTest extends ZooKeeperTestHarness {
     TestUtils.waitUntilMetadataIsPropagated(servers, topic2, 1)
     TestUtils.waitUntilMetadataIsPropagated(servers, topic2, 2)
 
-    val metadata = ClientUtils.fetchTopicMetadata(Set(topic2),
-      brokers.map(_.getBrokerEndPoint(ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT))),
-      "AddPartitionsTest-testReplicaPlacementPartialServers", 2000, 0).topicsMetadata
-
-    val metaDataForTopic2 = metadata.find(p => p.topic == topic2).get
-
-    validateLeaderAndReplicas(metaDataForTopic2, 0, 1, Set(1, 2))
-    validateLeaderAndReplicas(metaDataForTopic2, 1, 2, Set(0, 2))
-    validateLeaderAndReplicas(metaDataForTopic2, 2, 3, Set(1, 3))
+    val response = connectAndReceive[MetadataResponse](
+      new MetadataRequest.Builder(Seq(topic2).asJava, false).build)
+    assertEquals(1, response.topicMetadata.size)
+    val topicMetadata = response.topicMetadata.asScala.head
+    validateLeaderAndReplicas(topicMetadata, 0, 1, Set(1, 2))
+    validateLeaderAndReplicas(topicMetadata, 1, 2, Set(0, 2))
+    validateLeaderAndReplicas(topicMetadata, 2, 3, Set(1, 3))
   }
 
-  def validateLeaderAndReplicas(metadata: TopicMetadata, partitionId: Int, expectedLeaderId: Int, expectedReplicas: Set[Int]) = {
-    val partitionOpt = metadata.partitionsMetadata.find(_.partitionId == partitionId)
+  def validateLeaderAndReplicas(metadata: TopicMetadata, partitionId: Int, expectedLeaderId: Int,
+                                expectedReplicas: Set[Int]): Unit = {
+    val partitionOpt = metadata.partitionMetadata.asScala.find(_.partition == partitionId)
     assertTrue(s"Partition $partitionId should exist", partitionOpt.isDefined)
     val partition = partitionOpt.get
 
-    assertTrue("Partition leader should exist", partition.leader.isDefined)
-    assertEquals("Partition leader id should match", expectedLeaderId, partition.leader.get.id)
-
-    assertEquals("Replica set should match", expectedReplicas, partition.replicas.map(_.id).toSet)
+    assertEquals("Partition leader id should match", Optional.of(expectedLeaderId), partition.leaderId)
+    assertEquals("Replica set should match", expectedReplicas, partition.replicaIds.asScala.toSet)
   }
+
 }
diff --git a/core/src/test/scala/unit/kafka/admin/AdminRackAwareTest.scala b/core/src/test/scala/unit/kafka/admin/AdminRackAwareTest.scala
index ef85c6dc8ee02..3f2c06d830e07 100644
--- a/core/src/test/scala/unit/kafka/admin/AdminRackAwareTest.scala
+++ b/core/src/test/scala/unit/kafka/admin/AdminRackAwareTest.scala
@@ -17,19 +17,21 @@
 package kafka.admin
 
 import kafka.utils.Logging
+import org.apache.kafka.common.errors.InvalidReplicationFactorException
 import org.junit.Assert._
 import org.junit.Test
+import org.scalatest.Assertions._
 
 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)
-    val anotherList = AdminUtils.getRackAlternatedBrokerList(rackMap - 5)
+    val anotherList = AdminUtils.getRackAlternatedBrokerList(rackMap.toMap - 5)
     assertEquals(List(0, 3, 1, 4, 2), anotherList)
     val assignment = AdminUtils.assignReplicasToBrokers(toBrokerMetadata(rackMap), 7, 3, 0, 0)
     val expected = Map(0 -> List(0, 3, 1),
@@ -43,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
@@ -54,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
@@ -65,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
@@ -76,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
@@ -87,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
@@ -98,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
@@ -109,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
@@ -120,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
@@ -129,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",
@@ -141,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)
@@ -180,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
@@ -192,4 +194,35 @@ class AdminRackAwareTest extends RackAwareTest with Logging {
     checkReplicaDistribution(assignment, rackInfo, 5, 6, 4,
       verifyRackAware = false, verifyLeaderDistribution = false, verifyReplicasDistribution = false)
   }
+
+  @Test
+  def testReplicaAssignment(): Unit = {
+    val brokerMetadatas = (0 to 4).map(new BrokerMetadata(_, None))
+
+    // test 0 replication factor
+    intercept[InvalidReplicationFactorException] {
+      AdminUtils.assignReplicasToBrokers(brokerMetadatas, 10, 0)
+    }
+
+    // test wrong replication factor
+    intercept[InvalidReplicationFactorException] {
+      AdminUtils.assignReplicasToBrokers(brokerMetadatas, 10, 6)
+    }
+
+    // correct assignment
+    val expectedAssignment = 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))
+
+    val actualAssignment = AdminUtils.assignReplicasToBrokers(brokerMetadatas, 10, 3, 0)
+    assertEquals(expectedAssignment, actualAssignment)
+  }
 }
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 0c9bd6e4c6dd8..0000000000000
--- a/core/src/test/scala/unit/kafka/admin/AdminTest.scala
+++ /dev/null
@@ -1,579 +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 kafka.server.DynamicConfig.Broker._
-import kafka.server.KafkaConfig._
-import org.apache.kafka.common.errors.{InvalidReplicaAssignmentException, InvalidReplicationFactorException, InvalidTopicException, TopicExistsException}
-import org.apache.kafka.common.metrics.Quota
-import org.easymock.EasyMock
-import org.junit.Assert._
-import org.junit.{After, Test}
-import java.util.Properties
-
-import kafka.utils._
-import kafka.log._
-import kafka.zk.ZooKeeperTestHarness
-import kafka.utils.{Logging, TestUtils, ZkUtils}
-import kafka.common.TopicAndPartition
-import kafka.server.{ConfigType, KafkaConfig, KafkaServer}
-import java.io.File
-import java.util
-import java.util.concurrent.LinkedBlockingQueue
-
-import kafka.utils.TestUtils._
-import kafka.admin.AdminUtils._
-
-import scala.collection.{Map, Set, immutable}
-import kafka.utils.CoreUtils._
-import org.apache.kafka.common.TopicPartition
-
-import scala.collection.JavaConverters._
-import scala.util.Try
-
-class AdminTest extends ZooKeeperTestHarness with Logging with RackAwareTest {
-
-  var servers: Seq[KafkaServer] = Seq()
-
-  @After
-  override def tearDown() {
-    TestUtils.shutdownServers(servers)
-    super.tearDown()
-  }
-
-  @Test
-  def testReplicaAssignment() {
-    val brokerMetadatas = (0 to 4).map(new BrokerMetadata(_, None))
-
-    // test 0 replication factor
-    intercept[InvalidReplicationFactorException] {
-      AdminUtils.assignReplicasToBrokers(brokerMetadatas, 10, 0)
-    }
-
-    // test wrong replication factor
-    intercept[InvalidReplicationFactorException] {
-      AdminUtils.assignReplicasToBrokers(brokerMetadatas, 10, 6)
-    }
-
-    // correct assignment
-    val expectedAssignment = 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))
-
-    val actualAssignment = AdminUtils.assignReplicasToBrokers(brokerMetadatas, 10, 3, 0)
-    assertEquals(expectedAssignment, actualAssignment)
-  }
-
-  @Test
-  def testManualReplicaAssignment() {
-    val brokers = List(0, 1, 2, 3, 4)
-    TestUtils.createBrokersInZk(zkUtils, 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(zkUtils, List(0, 1, 2, 3, 4))
-    // create the topic
-    AdminUtils.createOrUpdateTopicPartitionAssignmentPathInZK(zkUtils, topic, expectedReplicaAssignment)
-    // create leaders for all partitions
-    TestUtils.makeLeaderForPartition(zkUtils, 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(zkUtils, 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 = 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)
-    }
-  }
-
-  private def getBrokersWithPartitionDir(servers: Iterable[KafkaServer], topic: String, partitionId: Int): Set[Int] = {
-    servers.filter(server => new File(server.config.logDirs.head, topic + "-" + partitionId).exists)
-           .map(_.config.brokerId)
-           .toSet
-  }
-
-  @Test
-  def testPartitionReassignmentWithLeaderInNewReplicas() {
-    val expectedReplicaAssignment = Map(0  -> List(0, 1, 2))
-    val topic = "test"
-    // create brokers
-    servers = TestUtils.createBrokerConfigs(4, zkConnect, false).map(b => TestUtils.createServer(KafkaConfig.fromProps(b)))
-    // create the topic
-    AdminUtils.createOrUpdateTopicPartitionAssignmentPathInZK(zkUtils, topic, expectedReplicaAssignment)
-    // reassign partition 0
-    val newReplicas = Seq(0, 2, 3)
-    val partitionToBeReassigned = 0
-    val topicAndPartition = TopicAndPartition(topic, partitionToBeReassigned)
-    val reassignPartitionsCommand = new ReassignPartitionsCommand(zkUtils, None, Map(topicAndPartition -> newReplicas))
-    assertTrue("Partition reassignment attempt failed for [test, 0]", reassignPartitionsCommand.reassignPartitions())
-    // wait until reassignment is completed
-    TestUtils.waitUntilTrue(() => {
-        val partitionsBeingReassigned = zkUtils.getPartitionsBeingReassigned().mapValues(_.newReplicas)
-        ReassignPartitionsCommand.checkIfPartitionReassignmentSucceeded(zkUtils, topicAndPartition,
-        Map(topicAndPartition -> newReplicas), partitionsBeingReassigned) == ReassignmentCompleted
-      },
-      "Partition reassignment should complete")
-    val assignedReplicas = zkUtils.getReplicasForPartition(topic, partitionToBeReassigned)
-    // in sync replicas should not have any replica that is not in the new assigned replicas
-    checkForPhantomInSyncReplicas(zkUtils, topic, partitionToBeReassigned, assignedReplicas)
-    assertEquals("Partition should have been reassigned to 0, 2, 3", newReplicas, assignedReplicas)
-    ensureNoUnderReplicatedPartitions(zkUtils, topic, partitionToBeReassigned, assignedReplicas, servers)
-    TestUtils.waitUntilTrue(() => getBrokersWithPartitionDir(servers, topic, 0) == newReplicas.toSet,
-                            "New replicas should exist on brokers")
-  }
-
-  @Test
-  def testPartitionReassignmentWithLeaderNotInNewReplicas() {
-    val expectedReplicaAssignment = Map(0  -> List(0, 1, 2))
-    val topic = "test"
-    // create brokers
-    servers = TestUtils.createBrokerConfigs(4, zkConnect, false).map(b => TestUtils.createServer(KafkaConfig.fromProps(b)))
-    // create the topic
-    AdminUtils.createOrUpdateTopicPartitionAssignmentPathInZK(zkUtils, topic, expectedReplicaAssignment)
-    // reassign partition 0
-    val newReplicas = Seq(1, 2, 3)
-    val partitionToBeReassigned = 0
-    val topicAndPartition = TopicAndPartition(topic, partitionToBeReassigned)
-    val reassignPartitionsCommand = new ReassignPartitionsCommand(zkUtils, None, Map(topicAndPartition -> newReplicas))
-    assertTrue("Partition reassignment failed for test, 0", reassignPartitionsCommand.reassignPartitions())
-    // wait until reassignment is completed
-    TestUtils.waitUntilTrue(() => {
-        val partitionsBeingReassigned = zkUtils.getPartitionsBeingReassigned().mapValues(_.newReplicas)
-        ReassignPartitionsCommand.checkIfPartitionReassignmentSucceeded(zkUtils, topicAndPartition,
-          Map(topicAndPartition -> newReplicas), partitionsBeingReassigned) == ReassignmentCompleted
-      },
-      "Partition reassignment should complete")
-    val assignedReplicas = zkUtils.getReplicasForPartition(topic, partitionToBeReassigned)
-    assertEquals("Partition should have been reassigned to 0, 2, 3", newReplicas, assignedReplicas)
-    checkForPhantomInSyncReplicas(zkUtils, topic, partitionToBeReassigned, assignedReplicas)
-    ensureNoUnderReplicatedPartitions(zkUtils, topic, partitionToBeReassigned, assignedReplicas, servers)
-    TestUtils.waitUntilTrue(() => getBrokersWithPartitionDir(servers, topic, 0) == newReplicas.toSet,
-                            "New replicas should exist on brokers")
-  }
-
-  @Test
-  def testPartitionReassignmentNonOverlappingReplicas() {
-    val expectedReplicaAssignment = Map(0  -> List(0, 1))
-    val topic = "test"
-    // create brokers
-    servers = TestUtils.createBrokerConfigs(4, zkConnect, false).map(b => TestUtils.createServer(KafkaConfig.fromProps(b)))
-    // create the topic
-    AdminUtils.createOrUpdateTopicPartitionAssignmentPathInZK(zkUtils, topic, expectedReplicaAssignment)
-    // reassign partition 0
-    val newReplicas = Seq(2, 3)
-    val partitionToBeReassigned = 0
-    val topicAndPartition = TopicAndPartition(topic, partitionToBeReassigned)
-    val reassignPartitionsCommand = new ReassignPartitionsCommand(zkUtils, None, Map(topicAndPartition -> newReplicas))
-    assertTrue("Partition reassignment failed for test, 0", reassignPartitionsCommand.reassignPartitions())
-    // wait until reassignment is completed
-    TestUtils.waitUntilTrue(() => {
-        val partitionsBeingReassigned = zkUtils.getPartitionsBeingReassigned().mapValues(_.newReplicas)
-        ReassignPartitionsCommand.checkIfPartitionReassignmentSucceeded(zkUtils, topicAndPartition,
-          Map(topicAndPartition -> newReplicas), partitionsBeingReassigned) == ReassignmentCompleted
-      },
-      "Partition reassignment should complete")
-    val assignedReplicas = zkUtils.getReplicasForPartition(topic, partitionToBeReassigned)
-    assertEquals("Partition should have been reassigned to 2, 3", newReplicas, assignedReplicas)
-    checkForPhantomInSyncReplicas(zkUtils, topic, partitionToBeReassigned, assignedReplicas)
-    ensureNoUnderReplicatedPartitions(zkUtils, topic, partitionToBeReassigned, assignedReplicas, servers)
-    TestUtils.waitUntilTrue(() => getBrokersWithPartitionDir(servers, topic, 0) == newReplicas.toSet,
-                            "New replicas should exist on brokers")
-  }
-
-  @Test
-  def testReassigningNonExistingPartition() {
-    val topic = "test"
-    // create brokers
-    servers = TestUtils.createBrokerConfigs(4, zkConnect, false).map(b => TestUtils.createServer(KafkaConfig.fromProps(b)))
-    // reassign partition 0
-    val newReplicas = Seq(2, 3)
-    val partitionToBeReassigned = 0
-    val topicAndPartition = TopicAndPartition(topic, partitionToBeReassigned)
-    val reassignPartitionsCommand = new ReassignPartitionsCommand(zkUtils, None, Map(topicAndPartition -> newReplicas))
-    assertFalse("Partition reassignment failed for test, 0", reassignPartitionsCommand.reassignPartitions())
-    val reassignedPartitions = zkUtils.getPartitionsBeingReassigned()
-    assertFalse("Partition should not be reassigned", reassignedPartitions.contains(topicAndPartition))
-  }
-
-  @Test
-  def testResumePartitionReassignmentThatWasCompleted() {
-    val expectedReplicaAssignment = Map(0  -> List(0, 1))
-    val topic = "test"
-    // create the topic
-    AdminUtils.createOrUpdateTopicPartitionAssignmentPathInZK(zkUtils, topic, expectedReplicaAssignment)
-    // put the partition in the reassigned path as well
-    // reassign partition 0
-    val newReplicas = Seq(0, 1)
-    val partitionToBeReassigned = 0
-    val topicAndPartition = TopicAndPartition(topic, partitionToBeReassigned)
-    val reassignPartitionsCommand = new ReassignPartitionsCommand(zkUtils, None, Map(topicAndPartition -> newReplicas))
-    reassignPartitionsCommand.reassignPartitions()
-    // create brokers
-    servers = TestUtils.createBrokerConfigs(2, zkConnect, false).map(b => TestUtils.createServer(KafkaConfig.fromProps(b)))
-
-    // wait until reassignment completes
-    TestUtils.waitUntilTrue(() => !checkIfReassignPartitionPathExists(zkUtils),
-                            "Partition reassignment should complete")
-    val assignedReplicas = zkUtils.getReplicasForPartition(topic, partitionToBeReassigned)
-    assertEquals("Partition should have been reassigned to 0, 1", newReplicas, assignedReplicas)
-    checkForPhantomInSyncReplicas(zkUtils, topic, partitionToBeReassigned, assignedReplicas)
-    // ensure that there are no under replicated partitions
-    ensureNoUnderReplicatedPartitions(zkUtils, topic, partitionToBeReassigned, assignedReplicas, servers)
-    TestUtils.waitUntilTrue(() => getBrokersWithPartitionDir(servers, topic, 0) == newReplicas.toSet,
-                            "New replicas should exist on brokers")
-  }
-
-  @Test
-  def testPreferredReplicaJsonData() {
-    // write preferred replica json data to zk path
-    val partitionsForPreferredReplicaElection = Set(TopicAndPartition("test", 1), TopicAndPartition("test2", 1))
-    PreferredReplicaLeaderElectionCommand.writePreferredReplicaElectionData(zkUtils, partitionsForPreferredReplicaElection)
-    // try to read it back and compare with what was written
-    val preferredReplicaElectionZkData = zkUtils.readData(ZkUtils.PreferredReplicaLeaderElectionPath)._1
-    val partitionsUndergoingPreferredReplicaElection =
-      PreferredReplicaLeaderElectionCommand.parsePreferredReplicaElectionData(preferredReplicaElectionZkData)
-    assertEquals("Preferred replica election ser-de failed", partitionsForPreferredReplicaElection,
-      partitionsUndergoingPreferredReplicaElection)
-  }
-
-  @Test
-  def testBasicPreferredReplicaElection() {
-    val expectedReplicaAssignment = Map(1  -> List(0, 1, 2))
-    val topic = "test"
-    val partition = 1
-    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
-    AdminUtils.createOrUpdateTopicPartitionAssignmentPathInZK(zkUtils, topic, expectedReplicaAssignment)
-    servers = serverConfigs.reverseMap(s => TestUtils.createServer(s))
-    // broker 2 should be the leader since it was started first
-    val currentLeader = TestUtils.waitUntilLeaderIsElectedOrChanged(zkUtils, topic, partition, oldLeaderOpt = None)
-    // trigger preferred replica election
-    val preferredReplicaElection = new PreferredReplicaLeaderElectionCommand(zkUtils, Set(TopicAndPartition(topic, partition)))
-    preferredReplicaElection.moveLeaderToPreferredReplica()
-    val newLeader = TestUtils.waitUntilLeaderIsElectedOrChanged(zkUtils, topic, partition, oldLeaderOpt = Some(currentLeader))
-    assertEquals("Preferred replica election failed", preferredReplica, newLeader)
-  }
-
-  @Test
-  def testControlledShutdown() {
-    val expectedReplicaAssignment = Map(1  -> List(0, 1, 2))
-    val topic = "test"
-    val partition = 1
-    // create brokers
-    val serverConfigs = TestUtils.createBrokerConfigs(3, zkConnect, false).map(KafkaConfig.fromProps)
-    servers = serverConfigs.reverseMap(s => TestUtils.createServer(s))
-    // create the topic
-    TestUtils.createTopic(zkUtils, topic, partitionReplicaAssignment = expectedReplicaAssignment, servers = servers)
-
-    val controllerId = zkUtils.getController()
-    val controller = servers.find(p => p.config.brokerId == controllerId).get.kafkaController
-    val resultQueue = new LinkedBlockingQueue[Try[Set[TopicPartition]]]()
-    val controlledShutdownCallback = (controlledShutdownResult: Try[Set[TopicPartition]]) => resultQueue.put(controlledShutdownResult)
-    controller.controlledShutdown(2, controlledShutdownCallback)
-    var partitionsRemaining = resultQueue.take().get
-    var activeServers = servers.filter(s => s.config.brokerId != 2)
-    // wait for the update metadata request to trickle to the brokers
-    TestUtils.waitUntilTrue(() =>
-      activeServers.forall(_.apis.metadataCache.getPartitionInfo(topic,partition).get.basePartitionState.isr.size != 3),
-      "Topic test not created after timeout")
-    assertEquals(0, partitionsRemaining.size)
-    var partitionStateInfo = activeServers.head.apis.metadataCache.getPartitionInfo(topic,partition).get
-    var leaderAfterShutdown = partitionStateInfo.basePartitionState.leader
-    assertEquals(0, leaderAfterShutdown)
-    assertEquals(2, partitionStateInfo.basePartitionState.isr.size)
-    assertEquals(List(0,1), partitionStateInfo.basePartitionState.isr.asScala)
-
-    controller.controlledShutdown(1, controlledShutdownCallback)
-    partitionsRemaining = resultQueue.take().get
-    assertEquals(0, partitionsRemaining.size)
-    activeServers = servers.filter(s => s.config.brokerId == 0)
-    partitionStateInfo = activeServers.head.apis.metadataCache.getPartitionInfo(topic,partition).get
-    leaderAfterShutdown = partitionStateInfo.basePartitionState.leader
-    assertEquals(0, leaderAfterShutdown)
-
-    assertTrue(servers.forall(_.apis.metadataCache.getPartitionInfo(topic,partition).get.basePartitionState.leader == 0))
-    controller.controlledShutdown(0, controlledShutdownCallback)
-    partitionsRemaining = resultQueue.take().get
-    assertEquals(1, partitionsRemaining.size)
-    // leader doesn't change since all the replicas are shut down
-    assertTrue(servers.forall(_.apis.metadataCache.getPartitionInfo(topic,partition).get.basePartitionState.leader == 0))
-  }
-
-  /**
-   * This test creates a topic with a few config overrides and checks that the configs are applied to the new topic
-   * then changes the config and checks that the new values take effect.
-   */
-  @Test
-  def testTopicConfigChange() {
-    val partitions = 3
-    val topic = "my-topic"
-    val server = TestUtils.createServer(KafkaConfig.fromProps(TestUtils.createBrokerConfig(0, zkConnect)))
-    servers = Seq(server)
-
-    def makeConfig(messageSize: Int, retentionMs: Long, throttledLeaders: String, throttledFollowers: String) = {
-      val props = new Properties()
-      props.setProperty(LogConfig.MaxMessageBytesProp, messageSize.toString)
-      props.setProperty(LogConfig.RetentionMsProp, retentionMs.toString)
-      props.setProperty(LogConfig.LeaderReplicationThrottledReplicasProp, throttledLeaders)
-      props.setProperty(LogConfig.FollowerReplicationThrottledReplicasProp, throttledFollowers)
-      props
-    }
-
-    def checkConfig(messageSize: Int, retentionMs: Long, throttledLeaders: String, throttledFollowers: String, quotaManagerIsThrottled: Boolean) {
-      def checkList(actual: util.List[String], expected: String): Unit = {
-        assertNotNull(actual)
-        if (expected == "")
-          assertTrue(actual.isEmpty)
-        else
-          assertEquals(expected.split(",").toSeq, actual.asScala)
-      }
-      TestUtils.retry(10000) {
-        for (part <- 0 until partitions) {
-          val tp = new TopicPartition(topic, part)
-          val log = server.logManager.getLog(tp)
-          assertTrue(log.isDefined)
-          assertEquals(retentionMs, log.get.config.retentionMs)
-          assertEquals(messageSize, log.get.config.maxMessageSize)
-          checkList(log.get.config.LeaderReplicationThrottledReplicas, throttledLeaders)
-          checkList(log.get.config.FollowerReplicationThrottledReplicas, throttledFollowers)
-          assertEquals(quotaManagerIsThrottled, server.quotaManagers.leader.isThrottled(tp))
-        }
-      }
-    }
-
-    // create a topic with a few config overrides and check that they are applied
-    val maxMessageSize = 1024
-    val retentionMs = 1000 * 1000
-    AdminUtils.createTopic(server.zkUtils, topic, partitions, 1, makeConfig(maxMessageSize, retentionMs, "0:0,1:0,2:0", "0:1,1:1,2:1"))
-
-    //Standard topic configs will be propagated at topic creation time, but the quota manager will not have been updated.
-    checkConfig(maxMessageSize, retentionMs, "0:0,1:0,2:0", "0:1,1:1,2:1", false)
-
-    //Update dynamically and all properties should be applied
-    AdminUtils.changeTopicConfig(server.zkUtils, topic, makeConfig(maxMessageSize, retentionMs, "0:0,1:0,2:0", "0:1,1:1,2:1"))
-
-    checkConfig(maxMessageSize, retentionMs, "0:0,1:0,2:0", "0:1,1:1,2:1", true)
-
-    // now double the config values for the topic and check that it is applied
-    val newConfig = makeConfig(2 * maxMessageSize, 2 * retentionMs, "*", "*")
-    AdminUtils.changeTopicConfig(server.zkUtils, topic, makeConfig(2 * maxMessageSize, 2 * retentionMs, "*", "*"))
-    checkConfig(2 * maxMessageSize, 2 * retentionMs, "*", "*", quotaManagerIsThrottled = true)
-
-    // Verify that the same config can be read from ZK
-    val configInZk = AdminUtils.fetchEntityConfig(server.zkUtils, ConfigType.Topic, topic)
-    assertEquals(newConfig, configInZk)
-
-    //Now delete the config
-    AdminUtils.changeTopicConfig(server.zkUtils, topic, new Properties)
-    checkConfig(Defaults.MaxMessageSize, Defaults.RetentionMs, "", "", quotaManagerIsThrottled = false)
-
-    //Add config back
-    AdminUtils.changeTopicConfig(server.zkUtils, topic, makeConfig(maxMessageSize, retentionMs, "0:0,1:0,2:0", "0:1,1:1,2:1"))
-    checkConfig(maxMessageSize, retentionMs, "0:0,1:0,2:0", "0:1,1:1,2:1", quotaManagerIsThrottled = true)
-
-    //Now ensure updating to "" removes the throttled replica list also
-    AdminUtils.changeTopicConfig(server.zkUtils, topic, propsWith((LogConfig.FollowerReplicationThrottledReplicasProp, ""), (LogConfig.LeaderReplicationThrottledReplicasProp, "")))
-    checkConfig(Defaults.MaxMessageSize, Defaults.RetentionMs, "", "",  quotaManagerIsThrottled = false)
-  }
-
-  @Test
-  def shouldPropagateDynamicBrokerConfigs() {
-    val brokerIds = Seq(0, 1, 2)
-    servers = createBrokerConfigs(3, zkConnect).map(fromProps).map(createServer(_))
-
-    def checkConfig(limit: Long) {
-      retry(10000) {
-        for (server <- servers) {
-          assertEquals("Leader Quota Manager was not updated", limit, server.quotaManagers.leader.upperBound)
-          assertEquals("Follower Quota Manager was not updated", limit, server.quotaManagers.follower.upperBound)
-        }
-      }
-    }
-
-    val limit: Long = 1000000
-
-    // Set the limit & check it is applied to the log
-    changeBrokerConfig(zkUtils, brokerIds, propsWith(
-      (LeaderReplicationThrottledRateProp, limit.toString),
-      (FollowerReplicationThrottledRateProp, limit.toString)))
-    checkConfig(limit)
-
-    // Now double the config values for the topic and check that it is applied
-    val newLimit = 2 * limit
-    changeBrokerConfig(zkUtils, brokerIds,  propsWith(
-      (LeaderReplicationThrottledRateProp, newLimit.toString),
-      (FollowerReplicationThrottledRateProp, newLimit.toString)))
-    checkConfig(newLimit)
-
-    // Verify that the same config can be read from ZK
-    for (brokerId <- brokerIds) {
-      val configInZk = AdminUtils.fetchEntityConfig(servers(brokerId).zkUtils, ConfigType.Broker, brokerId.toString)
-      assertEquals(newLimit, configInZk.getProperty(LeaderReplicationThrottledRateProp).toInt)
-      assertEquals(newLimit, configInZk.getProperty(FollowerReplicationThrottledRateProp).toInt)
-    }
-
-    //Now delete the config
-    changeBrokerConfig(servers(0).zkUtils, brokerIds, new Properties)
-    checkConfig(DefaultReplicationThrottledRate)
-  }
-
-  /**
-   * 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)
-    zkUtils.updatePersistentPath(ZkUtils.getEntityConfigPath(ConfigType.Client, clientId), Json.encode(map))
-
-    val configInZk: Map[String, Properties] = AdminUtils.fetchAllEntityConfigs(zkUtils, ConfigType.Client)
-    assertEquals("Must have 1 overriden 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.apis.quotas.produce.quota("ANONYMOUS", clientId))
-    assertEquals(new Quota(2000, true), server.apis.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, zkUtils)
-
-    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 acac907417c00..1eb18c9815c7f 100644
--- a/core/src/test/scala/unit/kafka/admin/ConfigCommandTest.scala
+++ b/core/src/test/scala/unit/kafka/admin/ConfigCommandTest.scala
@@ -16,71 +16,210 @@
  */
 package kafka.admin
 
+import java.util
 import java.util.Properties
 
 import kafka.admin.ConfigCommand.ConfigCommandOptions
-import kafka.common.InvalidConfigException
-import kafka.server.ConfigEntityName
-import kafka.utils.Logging
-import kafka.zk.{AdminZkClient, KafkaZkClient, ZooKeeperTestHarness}
-import org.apache.kafka.common.security.scram.ScramCredentialUtils
+import kafka.api.ApiVersion
+import kafka.cluster.{Broker, EndPoint}
+import kafka.server.{ConfigEntityName, ConfigType, KafkaConfig}
+import kafka.utils.{Exit, Logging}
+import kafka.zk.{AdminZkClient, BrokerInfo, KafkaZkClient, ZooKeeperTestHarness}
+import org.apache.kafka.clients.admin._
+import org.apache.kafka.common.Node
+import org.apache.kafka.common.config.{ConfigException, ConfigResource}
+import org.apache.kafka.common.errors.InvalidConfigurationException
+import org.apache.kafka.common.internals.KafkaFutureImpl
+import org.apache.kafka.common.network.ListenerName
+import org.apache.kafka.common.quota.{ClientQuotaAlteration, ClientQuotaEntity, ClientQuotaFilter, ClientQuotaFilterComponent}
+import org.apache.kafka.common.security.auth.SecurityProtocol
+import org.apache.kafka.common.security.scram.internals.ScramCredentialUtils
 import org.apache.kafka.common.utils.Sanitizer
+import org.apache.kafka.test.TestUtils
 import org.easymock.EasyMock
 import org.junit.Assert._
 import org.junit.Test
+import org.scalatest.Assertions.intercept
 
-import scala.collection.mutable
-import scala.collection.JavaConverters._
+import scala.collection.{Seq, mutable}
+import scala.jdk.CollectionConverters._
 
 class ConfigCommandTest extends ZooKeeperTestHarness with Logging {
+
+  @Test
+  def shouldExitWithNonZeroStatusOnArgError(): Unit = {
+    assertNonZeroStatusExit(Array("--blah"))
+  }
+
+  @Test
+  def shouldExitWithNonZeroStatusOnZkCommandError(): Unit = {
+    assertNonZeroStatusExit(Array(
+      "--zookeeper", zkConnect,
+      "--entity-name", "1",
+      "--entity-type", "brokers",
+      "--alter",
+      "--add-config", "security.inter.broker.protocol=PLAINTEXT"))
+  }
+
+  @Test
+  def shouldExitWithNonZeroStatusOnBrokerCommandError(): Unit = {
+    assertNonZeroStatusExit(Array(
+      "--bootstrap-server", "invalid host",
+      "--entity-type", "brokers",
+      "--entity-name", "1",
+      "--describe"))
+  }
+
+  private def assertNonZeroStatusExit(args: Array[String]): Unit = {
+    var exitStatus: Option[Int] = None
+    Exit.setExitProcedure { (status, _) =>
+      exitStatus = Some(status)
+      throw new RuntimeException
+    }
+
+    try {
+      ConfigCommand.main(args)
+    } catch {
+      case e: RuntimeException =>
+    } finally {
+      Exit.resetExitProcedure()
+    }
+
+    assertEquals(Some(1), exitStatus)
+  }
+
+  @Test
+  def shouldParseArgumentsForClientsEntityTypeUsingZookeeper(): Unit = {
+    testArgumentParse("clients", zkConfig = true)
+  }
+
   @Test
-  def shouldParseArgumentsForClientsEntityType() {
-    testArgumentParse("clients")
+  def shouldParseArgumentsForClientsEntityType(): Unit = {
+    testArgumentParse("clients", zkConfig = false)
   }
 
   @Test
-  def shouldParseArgumentsForTopicsEntityType() {
-    testArgumentParse("topics")
+  def shouldParseArgumentsForUsersEntityTypeUsingZookeeper(): Unit = {
+    testArgumentParse("users", zkConfig = true)
   }
 
   @Test
-  def shouldParseArgumentsForBrokersEntityType() {
-    testArgumentParse("brokers")
+  def shouldParseArgumentsForUsersEntityType(): Unit = {
+    testArgumentParse("users", zkConfig = false)
   }
 
-  def testArgumentParse(entityType: String) = {
+  @Test
+  def shouldParseArgumentsForTopicsEntityTypeUsingZookeeper(): Unit = {
+    testArgumentParse("topics", zkConfig = true)
+  }
+
+  @Test
+  def shouldParseArgumentsForTopicsEntityType(): Unit = {
+    testArgumentParse("topics", zkConfig = false)
+  }
+
+  @Test
+  def shouldParseArgumentsForBrokersEntityTypeUsingZookeeper(): Unit = {
+    testArgumentParse("brokers", zkConfig = true)
+  }
+
+  @Test
+  def shouldParseArgumentsForBrokersEntityType(): Unit = {
+    testArgumentParse("brokers", zkConfig = false)
+  }
+
+  @Test
+  def shouldParseArgumentsForBrokerLoggersEntityType(): Unit = {
+    testArgumentParse("broker-loggers", zkConfig = false)
+  }
+
+  @Test
+  def shouldParseArgumentsForIpEntityType(): Unit = {
+    testArgumentParse("ips", zkConfig = false)
+  }
+
+  @Test
+  def shouldParseArgumentsForIpEntityTypeUsingZookeeper(): Unit = {
+    testArgumentParse("ips", zkConfig = true)
+  }
+
+  def testArgumentParse(entityType: String, zkConfig: Boolean): Unit = {
+    val shortFlag: String = s"--${entityType.dropRight(1)}"
+
+    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()
 
+    createOpts = new ConfigCommandOptions(Array(connectOpts._1, connectOpts._2,
+      shortFlag, "1",
+      "--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()
+
+    createOpts = new ConfigCommandOptions(Array(connectOpts._1, connectOpts._2,
+      "--entity-name", "1",
       "--entity-type", entityType,
       "--alter",
+      "--add-config-file", "/tmp/new.properties"))
+    createOpts.checkArgs()
+
+    createOpts = new ConfigCommandOptions(Array(connectOpts._1, connectOpts._2,
+      shortFlag, "1",
+      "--alter",
       "--add-config", "a=b,c=d"))
     createOpts.checkArgs()
 
+    createOpts = new ConfigCommandOptions(Array(connectOpts._1, connectOpts._2,
+      shortFlag, "1",
+      "--alter",
+      "--add-config-file", "/tmp/new.properties"))
+    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()
 
+    createOpts = new ConfigCommandOptions(Array(connectOpts._1, connectOpts._2,
+      shortFlag, "1",
+      "--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",
       "--delete-config", "a"))
     createOpts.checkArgs()
 
+    createOpts = new ConfigCommandOptions(Array(connectOpts._1, connectOpts._2,
+      shortFlag, "1",
+      "--alter",
+      "--add-config", "a=b,c=d",
+      "--delete-config", "a"))
+    createOpts.checkArgs()
+
     val addedProps = ConfigCommand.parseConfigsToBeAdded(createOpts)
     assertEquals(2, addedProps.size())
     assertEquals("b", addedProps.getProperty("a"))
@@ -89,24 +228,207 @@ class ConfigCommandTest extends ZooKeeperTestHarness with Logging {
     val deletedProps = ConfigCommand.parseConfigsToBeDeleted(createOpts)
     assertEquals(1, deletedProps.size)
     assertEquals("a", deletedProps.head)
+
+    createOpts = new ConfigCommandOptions(Array(connectOpts._1, connectOpts._2,
+      "--entity-name", "1",
+      "--entity-type", entityType,
+      "--alter",
+      "--add-config", "a=b,c=,d=e,f="))
+    createOpts.checkArgs()
+
+    createOpts = new ConfigCommandOptions(Array(connectOpts._1, connectOpts._2,
+      shortFlag, "1",
+      "--alter",
+      "--add-config", "a=b,c=,d=e,f="))
+    createOpts.checkArgs()
+
+    val addedProps2 = ConfigCommand.parseConfigsToBeAdded(createOpts)
+    assertEquals(4, addedProps2.size())
+    assertEquals("b", addedProps2.getProperty("a"))
+    assertEquals("e", addedProps2.getProperty("d"))
+    assertTrue(addedProps2.getProperty("c").isEmpty)
+    assertTrue(addedProps2.getProperty("f").isEmpty)
   }
 
   @Test(expected = classOf[IllegalArgumentException])
-  def shouldFailIfUnrecognisedEntityType(): Unit = {
+  def shouldFailIfAddAndAddFile(): Unit = {
+    // Should not parse correctly
+    val createOpts = new ConfigCommandOptions(Array("--bootstrap-server", "localhost:9092",
+      "--entity-name", "1",
+      "--entity-type", "brokers",
+      "--alter",
+      "--add-config", "a=b,c=d",
+      "--add-config-file", "/tmp/new.properties"
+    ))
+    createOpts.checkArgs()
+  }
+
+  @Test
+  def testParseConfigsToBeAddedForAddConfigFile(): Unit = {
+    val fileContents =
+      """a=b
+        |c = d
+        |json = {"key": "val"}
+        |nested = [[1, 2], [3, 4]]
+        |""".stripMargin
+
+    val file = TestUtils.tempFile(fileContents)
+
+    val addConfigFileArgs = Array("--add-config-file", file.getPath)
+
+    val createOpts = new ConfigCommandOptions(Array("--bootstrap-server", "localhost:9092",
+      "--entity-name", "1",
+      "--entity-type", "brokers",
+      "--alter")
+      ++ addConfigFileArgs)
+    createOpts.checkArgs()
+
+    val addedProps = ConfigCommand.parseConfigsToBeAdded(createOpts)
+    assertEquals(4, addedProps.size())
+    assertEquals("b", addedProps.getProperty("a"))
+    assertEquals("d", addedProps.getProperty("c"))
+    assertEquals("{\"key\": \"val\"}", addedProps.getProperty("json"))
+    assertEquals("[[1, 2], [3, 4]]", addedProps.getProperty("nested"))
+  }
+
+  def doTestOptionEntityTypeNames(zkConfig: Boolean): Unit = {
+    val connectOpts = if (zkConfig)
+      ("--zookeeper", zkConnect)
+    else
+      ("--bootstrap-server", "localhost:9092")
+
+    def testExpectedEntityTypeNames(expectedTypes: List[String], expectedNames: List[String], args: String*): Unit = {
+      val createOpts = new ConfigCommandOptions(Array(connectOpts._1, connectOpts._2, "--describe") ++ args)
+      createOpts.checkArgs()
+      assertEquals(createOpts.entityTypes, expectedTypes)
+      assertEquals(createOpts.entityNames, expectedNames)
+    }
+
+    testExpectedEntityTypeNames(List(ConfigType.Topic), List("A"), "--entity-type", "topics", "--entity-name", "A")
+    testExpectedEntityTypeNames(List(ConfigType.Broker), List("0"), "--entity-name", "0", "--entity-type", "brokers")
+    testExpectedEntityTypeNames(List(ConfigType.Ip), List("1.2.3.4"), "--entity-name", "1.2.3.4", "--entity-type", "ips")
+    testExpectedEntityTypeNames(List(ConfigType.User, ConfigType.Client), List("A", ""),
+      "--entity-type", "users", "--entity-type", "clients", "--entity-name", "A", "--entity-default")
+    testExpectedEntityTypeNames(List(ConfigType.User, ConfigType.Client), List("", "B"),
+      "--entity-default", "--entity-name", "B", "--entity-type", "users", "--entity-type", "clients")
+
+    testExpectedEntityTypeNames(List(ConfigType.Topic), List("A"), "--topic", "A")
+    testExpectedEntityTypeNames(List(ConfigType.Broker), List("0"), "--broker", "0")
+    testExpectedEntityTypeNames(List(ConfigType.Ip), List("1.2.3.4"), "--ip", "1.2.3.4")
+    testExpectedEntityTypeNames(List(ConfigType.Client, ConfigType.User), List("B", "A"), "--client", "B", "--user", "A")
+    testExpectedEntityTypeNames(List(ConfigType.Client, ConfigType.User), List("B", ""), "--client", "B", "--user-defaults")
+    testExpectedEntityTypeNames(List(ConfigType.Client, ConfigType.User), List("A"),
+      "--entity-type", "clients", "--entity-type", "users", "--entity-name", "A")
+
+    testExpectedEntityTypeNames(List(ConfigType.Topic), List.empty, "--entity-type", "topics")
+    testExpectedEntityTypeNames(List(ConfigType.User), List.empty, "--entity-type", "users")
+    testExpectedEntityTypeNames(List(ConfigType.Broker), List.empty, "--entity-type", "brokers")
+    testExpectedEntityTypeNames(List(ConfigType.Ip), List.empty, "--entity-type", "ips")
+  }
+
+  @Test
+  def testOptionEntityTypeNamesUsingZookeeper(): Unit = {
+    doTestOptionEntityTypeNames(zkConfig = true)
+  }
+
+  @Test
+  def testOptionEntityTypeNames(): Unit = {
+    doTestOptionEntityTypeNames(zkConfig = false)
+  }
+
+  @Test(expected = classOf[IllegalArgumentException])
+  def shouldFailIfUnrecognisedEntityTypeUsingZookeeper(): Unit = {
     val createOpts = new ConfigCommandOptions(Array("--zookeeper", zkConnect,
       "--entity-name", "client", "--entity-type", "not-recognised", "--alter", "--add-config", "a=b,c=d"))
-    ConfigCommand.alterConfig(null, createOpts, new DummyAdminZkClient(zkClient))
+    ConfigCommand.alterConfigWithZk(null, createOpts, new DummyAdminZkClient(zkClient))
+  }
+
+  @Test(expected = classOf[IllegalArgumentException])
+  def shouldFailIfUnrecognisedEntityType(): Unit = {
+    val createOpts = new ConfigCommandOptions(Array("--bootstrap-server", "localhost:9092",
+      "--entity-name", "client", "--entity-type", "not-recognised", "--alter", "--add-config", "a=b,c=d"))
+    ConfigCommand.alterConfig(new DummyAdminClient(new Node(1, "localhost", 9092)), createOpts)
+  }
+
+  @Test(expected = classOf[IllegalArgumentException])
+  def shouldFailIfBrokerEntityTypeIsNotAnIntegerUsingZookeeper(): Unit = {
+    val createOpts = new ConfigCommandOptions(Array("--zookeeper", zkConnect,
+      "--entity-name", "A", "--entity-type", "brokers", "--alter", "--add-config", "a=b,c=d"))
+    ConfigCommand.alterConfigWithZk(null, createOpts, new DummyAdminZkClient(zkClient))
+  }
+
+  @Test(expected = classOf[IllegalArgumentException])
+  def shouldFailIfBrokerEntityTypeIsNotAnInteger(): Unit = {
+    val createOpts = new ConfigCommandOptions(Array("--bootstrap-server", "localhost:9092",
+      "--entity-name", "A", "--entity-type", "brokers", "--alter", "--add-config", "a=b,c=d"))
+    ConfigCommand.alterConfig(new DummyAdminClient(new Node(1, "localhost", 9092)), createOpts)
+  }
+
+  @Test(expected = classOf[IllegalArgumentException])
+  def shouldFailIfShortBrokerEntityTypeIsNotAnIntegerUsingZookeeper(): Unit = {
+    val createOpts = new ConfigCommandOptions(Array("--zookeeper", zkConnect,
+      "--broker", "A", "--alter", "--add-config", "a=b,c=d"))
+    ConfigCommand.alterConfigWithZk(null, createOpts, new DummyAdminZkClient(zkClient))
+  }
+
+  @Test(expected = classOf[IllegalArgumentException])
+  def shouldFailIfShortBrokerEntityTypeIsNotAnInteger(): Unit = {
+    val createOpts = new ConfigCommandOptions(Array("--bootstrap-server", "localhost:9092",
+      "--broker", "A", "--alter", "--add-config", "a=b,c=d"))
+    ConfigCommand.alterConfig(new DummyAdminClient(new Node(1, "localhost", 9092)), createOpts)
+  }
+
+  @Test(expected = classOf[IllegalArgumentException])
+  def shouldFailIfMixedEntityTypeFlagsUsingZookeeper(): Unit = {
+    val createOpts = new ConfigCommandOptions(Array("--zookeeper", zkConnect,
+      "--entity-name", "A", "--entity-type", "users", "--client", "B", "--describe"))
+    createOpts.checkArgs()
+  }
+
+  @Test(expected = classOf[IllegalArgumentException])
+  def shouldFailIfMixedEntityTypeFlags(): Unit = {
+    val createOpts = new ConfigCommandOptions(Array("--bootstrap-server", "localhost:9092",
+      "--entity-name", "A", "--entity-type", "users", "--client", "B", "--describe"))
+    createOpts.checkArgs()
+  }
+
+  @Test(expected = classOf[IllegalArgumentException])
+  def shouldFailIfInvalidHost(): Unit = {
+    val createOpts = new ConfigCommandOptions(Array("--bootstrap-server", "localhost:9092",
+      "--entity-name", "A,B", "--entity-type", "ips", "--describe"))
+    createOpts.checkArgs()
+  }
+
+  @Test(expected = classOf[IllegalArgumentException])
+  def shouldFailIfInvalidHostUsingZookeeper(): Unit = {
+    val createOpts = new ConfigCommandOptions(Array("--zookeeper", zkConnect,
+      "--entity-name", "A,B", "--entity-type", "ips", "--describe"))
+    createOpts.checkArgs()
+  }
+
+  @Test(expected = classOf[IllegalArgumentException])
+  def shouldFailIfUnresolvableHost(): Unit = {
+    val createOpts = new ConfigCommandOptions(Array("--bootstrap-server", "localhost:9092",
+      "--entity-name", "admin", "--entity-type", "ips", "--describe"))
+    createOpts.checkArgs()
+  }
+
+  @Test(expected = classOf[IllegalArgumentException])
+  def shouldFailIfUnresolvableHostUsingZookeeper(): Unit = {
+    val createOpts = new ConfigCommandOptions(Array("--zookeeper", zkConnect,
+      "--entity-name", "admin", "--entity-type", "ips", "--describe"))
+    createOpts.checkArgs()
   }
 
   @Test
-  def shouldAddClientConfig(): Unit = {
+  def shouldAddClientConfigUsingZookeeper(): Unit = {
     val createOpts = new ConfigCommandOptions(Array("--zookeeper", zkConnect,
       "--entity-name", "my-client-id",
       "--entity-type", "clients",
       "--alter",
       "--add-config", "a=b,c=d"))
 
-    case class TestAdminZkClient(val zkClient: KafkaZkClient) extends AdminZkClient(zkClient) {
+    class TestAdminZkClient(zkClient: KafkaZkClient) extends AdminZkClient(zkClient) {
       override def changeClientIdConfig(clientId: String, configChange: Properties): Unit = {
         assertEquals("my-client-id", clientId)
         assertEquals("b", configChange.get("a"))
@@ -114,18 +436,299 @@ class ConfigCommandTest extends ZooKeeperTestHarness with Logging {
       }
     }
 
-    ConfigCommand.alterConfig(null, createOpts, new TestAdminZkClient(zkClient))
+    ConfigCommand.alterConfigWithZk(null, createOpts, new TestAdminZkClient(zkClient))
   }
 
   @Test
-  def shouldAddTopicConfig(): Unit = {
+  def shouldAddIpConfigsUsingZookeeper(): Unit = {
+    val createOpts = new ConfigCommandOptions(Array("--zookeeper", zkConnect,
+      "--entity-name", "1.2.3.4",
+      "--entity-type", "ips",
+      "--alter",
+      "--add-config", "a=b,c=d"))
+
+    class TestAdminZkClient(zkClient: KafkaZkClient) extends AdminZkClient(zkClient) {
+      override def changeIpConfig(ip: String, configChange: Properties): Unit = {
+        assertEquals("1.2.3.4", ip)
+        assertEquals("b", configChange.get("a"))
+        assertEquals("d", configChange.get("c"))
+      }
+    }
+
+    ConfigCommand.alterConfigWithZk(null, createOpts, new TestAdminZkClient(zkClient))
+  }
+
+  private def toValues(entityName: Option[String], entityType: String): (Array[String], Map[String, String]) = {
+    val command = entityType match {
+      case ClientQuotaEntity.USER => "users"
+      case ClientQuotaEntity.CLIENT_ID => "clients"
+      case ClientQuotaEntity.IP => "ips"
+    }
+    entityName match {
+      case Some(null) =>
+        (Array("--entity-type", command, "--entity-default"), Map(entityType -> null))
+      case Some(name) =>
+        (Array("--entity-type", command, "--entity-name", name), Map(entityType -> name))
+      case None => (Array.empty, Map.empty)
+    }
+  }
+
+  private def verifyAlterCommandFails(expectedErrorMessage: String, alterOpts: Seq[String]): Unit = {
+    val mockAdminClient: Admin = EasyMock.createStrictMock(classOf[Admin])
+    val opts = new ConfigCommandOptions(Array("--bootstrap-server", "localhost:9092",
+      "--alter") ++ alterOpts)
+    val e = intercept[IllegalArgumentException] {
+      ConfigCommand.alterConfig(mockAdminClient, opts)
+    }
+    assertTrue(s"Unexpected exception: $e", e.getMessage.contains(expectedErrorMessage))
+  }
+
+  @Test
+  def shouldNotAlterNonQuotaIpConfigsUsingBootstrapServer(): Unit = {
+    // when using --bootstrap-server, it should be illegal to alter anything that is not a connection quota
+    // for ip entities
+    val ipEntityOpts = List("--entity-type", "ips", "--entity-name", "127.0.0.1")
+    val invalidProp = "some_config"
+    verifyAlterCommandFails(invalidProp, ipEntityOpts ++ List("--add-config", "connection_creation_rate=10000,some_config=10"))
+    verifyAlterCommandFails(invalidProp, ipEntityOpts ++ List("--add-config", "some_config=10"))
+    verifyAlterCommandFails(invalidProp, ipEntityOpts ++ List("--delete-config", "connection_creation_rate=10000,some_config=10"))
+    verifyAlterCommandFails(invalidProp, ipEntityOpts ++ List("--delete-config", "some_config=10"))
+  }
+
+  private def verifyDescribeQuotas(describeArgs: List[String], expectedFilter: ClientQuotaFilter): Unit = {
+    val describeOpts = new ConfigCommandOptions(Array("--bootstrap-server", "localhost:9092",
+      "--describe") ++ describeArgs)
+    val describeFuture = new KafkaFutureImpl[util.Map[ClientQuotaEntity, util.Map[String, java.lang.Double]]]
+    describeFuture.complete(Map.empty[ClientQuotaEntity, util.Map[String, java.lang.Double]].asJava)
+    val describeResult: DescribeClientQuotasResult = EasyMock.createNiceMock(classOf[DescribeClientQuotasResult])
+    EasyMock.expect(describeResult.entities()).andReturn(describeFuture)
+
+    var describedConfigs = false
+    val node = new Node(1, "localhost", 9092)
+    val mockAdminClient = new MockAdminClient(util.Collections.singletonList(node), node) {
+      override def describeClientQuotas(filter: ClientQuotaFilter, options: DescribeClientQuotasOptions): DescribeClientQuotasResult = {
+        assertTrue(filter.strict)
+        assertEquals(expectedFilter.components().asScala.toSet, filter.components.asScala.toSet)
+        describedConfigs = true
+        describeResult
+      }
+    }
+    EasyMock.replay(describeResult)
+    ConfigCommand.describeConfig(mockAdminClient, describeOpts)
+    assertTrue(describedConfigs)
+  }
+
+  @Test
+  def testDescribeIpConfigs(): Unit = {
+    val entityType = ClientQuotaEntity.IP
+    val knownHost = "1.2.3.4"
+    val defaultIpFilter = ClientQuotaFilter.containsOnly(List(ClientQuotaFilterComponent.ofDefaultEntity(entityType)).asJava)
+    val singleIpFilter = ClientQuotaFilter.containsOnly(List(ClientQuotaFilterComponent.ofEntity(entityType, knownHost)).asJava)
+    val allIpsFilter = ClientQuotaFilter.containsOnly(List(ClientQuotaFilterComponent.ofEntityType(entityType)).asJava)
+    verifyDescribeQuotas(List("--entity-default", "--entity-type", "ips"), defaultIpFilter)
+    verifyDescribeQuotas(List("--ip-defaults"), defaultIpFilter)
+    verifyDescribeQuotas(List("--entity-type", "ips", "--entity-name", knownHost), singleIpFilter)
+    verifyDescribeQuotas(List("--ip", knownHost), singleIpFilter)
+    verifyDescribeQuotas(List("--entity-type", "ips"), allIpsFilter)
+  }
+
+  def verifyAlterQuotas(alterOpts: Seq[String], expectedAlterEntity: ClientQuotaEntity,
+                        expectedProps: Map[String, java.lang.Double], expectedAlterOps: Set[ClientQuotaAlteration.Op]): Unit = {
+    val createOpts = new ConfigCommandOptions(Array("--bootstrap-server", "localhost:9092",
+      "--alter") ++ alterOpts)
+
+    var describedConfigs = false
+    val describeFuture = new KafkaFutureImpl[util.Map[ClientQuotaEntity, util.Map[String, java.lang.Double]]]
+    describeFuture.complete(Map(expectedAlterEntity -> expectedProps.asJava).asJava)
+    val describeResult: DescribeClientQuotasResult = EasyMock.createNiceMock(classOf[DescribeClientQuotasResult])
+    EasyMock.expect(describeResult.entities()).andReturn(describeFuture)
+
+    val expectedFilterComponents = expectedAlterEntity.entries.asScala.map { case (entityType, entityName) =>
+      if (entityName == null)
+        ClientQuotaFilterComponent.ofDefaultEntity(entityType)
+      else
+        ClientQuotaFilterComponent.ofEntity(entityType, entityName)
+    }.toSet
+
+    var alteredConfigs = false
+    val alterFuture = new KafkaFutureImpl[Void]
+    alterFuture.complete(null)
+    val alterResult: AlterClientQuotasResult = EasyMock.createNiceMock(classOf[AlterClientQuotasResult])
+    EasyMock.expect(alterResult.all()).andReturn(alterFuture)
+
+    val node = new Node(1, "localhost", 9092)
+    val mockAdminClient = new MockAdminClient(util.Collections.singletonList(node), node) {
+      override def describeClientQuotas(filter: ClientQuotaFilter, options: DescribeClientQuotasOptions): DescribeClientQuotasResult = {
+        assertTrue(filter.strict)
+        assertEquals(expectedFilterComponents, filter.components().asScala.toSet)
+        describedConfigs = true
+        describeResult
+      }
+
+      override def alterClientQuotas(entries: util.Collection[ClientQuotaAlteration], options: AlterClientQuotasOptions): AlterClientQuotasResult = {
+        assertFalse(options.validateOnly)
+        assertEquals(1, entries.size)
+        val alteration = entries.asScala.head
+        assertEquals(expectedAlterEntity, alteration.entity)
+        val ops = alteration.ops.asScala
+        assertEquals(expectedAlterOps, ops.toSet)
+        alteredConfigs = true
+        alterResult
+      }
+    }
+    EasyMock.replay(alterResult, describeResult)
+    ConfigCommand.alterConfig(mockAdminClient, createOpts)
+    assertTrue(describedConfigs)
+    assertTrue(alteredConfigs)
+  }
+
+  @Test
+  def testAlterIpConfig(): Unit = {
+    val (singleIpArgs, singleIpEntry) = toValues(Some("1.2.3.4"), ClientQuotaEntity.IP)
+    val singleIpEntity = new ClientQuotaEntity(singleIpEntry.asJava)
+    val (defaultIpArgs, defaultIpEntry) = toValues(Some(null), ClientQuotaEntity.IP)
+    val defaultIpEntity = new ClientQuotaEntity(defaultIpEntry.asJava)
+
+    val deleteArgs = List("--delete-config", "connection_creation_rate")
+    val deleteAlterationOps = Set(new ClientQuotaAlteration.Op("connection_creation_rate", null))
+    val propsToDelete = Map("connection_creation_rate" -> Double.box(50.0))
+
+    val addArgs = List("--add-config", "connection_creation_rate=100")
+    val addAlterationOps = Set(new ClientQuotaAlteration.Op("connection_creation_rate", 100.0))
+
+    verifyAlterQuotas(singleIpArgs ++ deleteArgs, singleIpEntity, propsToDelete, deleteAlterationOps)
+    verifyAlterQuotas(singleIpArgs ++ addArgs, singleIpEntity, Map.empty, addAlterationOps)
+    verifyAlterQuotas(defaultIpArgs ++ deleteArgs, defaultIpEntity, propsToDelete, deleteAlterationOps)
+    verifyAlterQuotas(defaultIpArgs ++ addArgs, defaultIpEntity, Map.empty, addAlterationOps)
+  }
+
+  @Test
+  def shouldAddClientConfig(): Unit = {
+    val alterArgs = List("--add-config", "consumer_byte_rate=20000,producer_byte_rate=10000",
+      "--delete-config", "request_percentage")
+    val propsToDelete = Map("request_percentage" -> Double.box(50.0))
+
+    val alterationOps = Set(
+      new ClientQuotaAlteration.Op("consumer_byte_rate", Double.box(20000)),
+      new ClientQuotaAlteration.Op("producer_byte_rate", Double.box(10000)),
+      new ClientQuotaAlteration.Op("request_percentage", null)
+    )
+
+    def verifyAlterUserClientQuotas(userOpt: Option[String], clientOpt: Option[String]): Unit = {
+      val (userArgs, userEntry) = toValues(userOpt, ClientQuotaEntity.USER)
+      val (clientArgs, clientEntry) = toValues(clientOpt, ClientQuotaEntity.CLIENT_ID)
+
+      val commandArgs = alterArgs ++ userArgs ++ clientArgs
+      val clientQuotaEntity = new ClientQuotaEntity((userEntry ++ clientEntry).asJava)
+      verifyAlterQuotas(commandArgs, clientQuotaEntity, propsToDelete, alterationOps)
+    }
+    verifyAlterUserClientQuotas(Some("test-user-1"), Some("test-client-1"))
+    verifyAlterUserClientQuotas(Some("test-user-2"), Some(null))
+    verifyAlterUserClientQuotas(Some("test-user-3"), None)
+    verifyAlterUserClientQuotas(Some(null), Some("test-client-2"))
+    verifyAlterUserClientQuotas(Some(null), Some(null))
+    verifyAlterUserClientQuotas(Some(null), None)
+    verifyAlterUserClientQuotas(None, Some("test-client-3"))
+    verifyAlterUserClientQuotas(None, Some(null))
+  }
+
+  private val userEntityOpts = List("--entity-type", "users", "--entity-name", "admin")
+  private val clientEntityOpts = List("--entity-type", "clients", "--entity-name", "admin")
+  private val addScramOpts = List("--add-config", "SCRAM-SHA-256=[iterations=8192,password=foo-secret]")
+  private val deleteScramOpts = List("--delete-config", "SCRAM-SHA-256")
+
+  @Test
+  def shouldNotAlterNonQuotaNonScramUserOrClientConfigUsingBootstrapServer(): Unit = {
+    // when using --bootstrap-server, it should be illegal to alter anything that is not a quota and not a SCRAM credential
+    // for both user and client entities
+    val invalidProp = "some_config"
+    verifyAlterCommandFails(invalidProp, userEntityOpts ++
+      List("-add-config", "consumer_byte_rate=20000,producer_byte_rate=10000,some_config=10"))
+    verifyAlterCommandFails(invalidProp, userEntityOpts ++
+      List("--add-config", "consumer_byte_rate=20000,producer_byte_rate=10000,some_config=10"))
+    verifyAlterCommandFails(invalidProp, clientEntityOpts ++ List("--add-config", "some_config=10"))
+    verifyAlterCommandFails(invalidProp, userEntityOpts ++ List("--delete-config", "consumer_byte_rate,some_config"))
+    verifyAlterCommandFails(invalidProp, userEntityOpts ++ List("--delete-config", "SCRAM-SHA-256,some_config"))
+    verifyAlterCommandFails(invalidProp, clientEntityOpts ++ List("--delete-config", "some_config"))
+  }
+
+  @Test
+  def shouldNotAlterScramClientConfigUsingBootstrapServer(): Unit = {
+    // when using --bootstrap-server, it should be illegal to alter SCRAM credentials for client entities
+    verifyAlterCommandFails("SCRAM-SHA-256", clientEntityOpts ++ addScramOpts)
+    verifyAlterCommandFails("SCRAM-SHA-256", clientEntityOpts ++ deleteScramOpts)
+  }
+
+  @Test
+  def shouldNotCreateUserScramCredentialConfigWithUnderMinimumIterationsUsingBootstrapServer(): Unit = {
+    // when using --bootstrap-server, it should be illegal to create a SCRAM credential for a user
+    // with an iterations value less than the minimum
+    verifyAlterCommandFails("SCRAM-SHA-256", userEntityOpts ++ List("--add-config", "SCRAM-SHA-256=[iterations=100,password=foo-secret]"))
+  }
+
+  @Test
+  def shouldNotAlterUserScramCredentialAndClientQuotaConfigsSimultaneouslyUsingBootstrapServer(): Unit = {
+    // when using --bootstrap-server, it should be illegal to alter both SCRAM credentials and quotas for user entities
+    val expectedErrorMessage = "SCRAM-SHA-256"
+    val secondUserEntityOpts = List("--entity-type", "users", "--entity-name", "admin1")
+    val addQuotaOpts = List("--add-config", "consumer_byte_rate=20000")
+    val deleteQuotaOpts = List("--delete-config", "consumer_byte_rate")
+
+    verifyAlterCommandFails(expectedErrorMessage, userEntityOpts ++ addScramOpts ++ userEntityOpts ++ deleteQuotaOpts)
+    verifyAlterCommandFails(expectedErrorMessage, userEntityOpts ++ addScramOpts ++ secondUserEntityOpts ++ deleteQuotaOpts)
+    verifyAlterCommandFails(expectedErrorMessage, userEntityOpts ++ deleteScramOpts ++ userEntityOpts ++ addQuotaOpts)
+    verifyAlterCommandFails(expectedErrorMessage, userEntityOpts ++ deleteScramOpts ++ secondUserEntityOpts ++ addQuotaOpts)
+
+    // change order of quota/SCRAM commands, verify alter still fails
+    verifyAlterCommandFails(expectedErrorMessage, userEntityOpts ++ deleteQuotaOpts ++ userEntityOpts ++ addScramOpts)
+    verifyAlterCommandFails(expectedErrorMessage, secondUserEntityOpts ++ deleteQuotaOpts ++ userEntityOpts ++ addScramOpts)
+    verifyAlterCommandFails(expectedErrorMessage, userEntityOpts ++ addQuotaOpts ++ userEntityOpts ++ deleteScramOpts)
+    verifyAlterCommandFails(expectedErrorMessage, secondUserEntityOpts ++ addQuotaOpts ++ userEntityOpts ++ deleteScramOpts)
+  }
+
+  @Test
+  def shouldNotDescribeUserScramCredentialsWithEntityDefaultUsingBootstrapServer(): Unit = {
+    def verifyUserScramCredentialsNotDescribed(requestOpts: List[String]): Unit = {
+      // User SCRAM credentials should not be described when specifying
+      // --describe --entity-type users --entity-default (or --user-defaults) with --bootstrap-server
+      val describeFuture = new KafkaFutureImpl[util.Map[ClientQuotaEntity, util.Map[String, java.lang.Double]]]
+      describeFuture.complete(Map((new ClientQuotaEntity(Map("" -> "").asJava) -> Map(("request_percentage" -> Double.box(50.0))).asJava)).asJava)
+      val describeClientQuotasResult: DescribeClientQuotasResult = EasyMock.createNiceMock(classOf[DescribeClientQuotasResult])
+      EasyMock.expect(describeClientQuotasResult.entities()).andReturn(describeFuture)
+      EasyMock.replay(describeClientQuotasResult)
+      val node = new Node(1, "localhost", 9092)
+      val mockAdminClient = new MockAdminClient(util.Collections.singletonList(node), node) {
+        override def describeClientQuotas(filter: ClientQuotaFilter, options: DescribeClientQuotasOptions):  DescribeClientQuotasResult = {
+          describeClientQuotasResult
+        }
+        override def describeUserScramCredentials(users: util.List[String], options: DescribeUserScramCredentialsOptions): DescribeUserScramCredentialsResult = {
+          throw new IllegalStateException("Incorrectly described SCRAM credentials when specifying --entity-default with --bootstrap-server")
+        }
+      }
+      val opts = new ConfigCommandOptions(Array("--bootstrap-server", "localhost:9092", "--describe") ++ requestOpts)
+      ConfigCommand.describeConfig(mockAdminClient, opts) // fails if describeUserScramCredentials() is invoked
+    }
+
+    val expectedMsg = "The use of --entity-default or --user-defaults is not allowed with User SCRAM Credentials using --bootstrap-server."
+    val defaultUserOpt = List("--user-defaults")
+    val verboseDefaultUserOpts = List("--entity-type", "users", "--entity-default")
+    verifyAlterCommandFails(expectedMsg, verboseDefaultUserOpts ++ addScramOpts)
+    verifyAlterCommandFails(expectedMsg, verboseDefaultUserOpts ++ deleteScramOpts)
+    verifyUserScramCredentialsNotDescribed(verboseDefaultUserOpts)
+    verifyAlterCommandFails(expectedMsg, defaultUserOpt ++ addScramOpts)
+    verifyAlterCommandFails(expectedMsg, defaultUserOpt ++ deleteScramOpts)
+    verifyUserScramCredentialsNotDescribed(defaultUserOpt)
+  }
+
+  @Test
+  def shouldAddTopicConfigUsingZookeeper(): Unit = {
     val createOpts = new ConfigCommandOptions(Array("--zookeeper", zkConnect,
       "--entity-name", "my-topic",
       "--entity-type", "topics",
       "--alter",
       "--add-config", "a=b,c=d"))
 
-    case class TestAdminZkClient(val zkClient: KafkaZkClient) extends AdminZkClient(zkClient) {
+    class TestAdminZkClient(zkClient: KafkaZkClient) extends AdminZkClient(zkClient) {
       override def changeTopicConfig(topic: String, configChange: Properties): Unit = {
         assertEquals("my-topic", topic)
         assertEquals("b", configChange.get("a"))
@@ -133,88 +736,592 @@ class ConfigCommandTest extends ZooKeeperTestHarness with Logging {
       }
     }
 
-    ConfigCommand.alterConfig(null, createOpts, new TestAdminZkClient(zkClient))
+    ConfigCommand.alterConfigWithZk(null, createOpts, new TestAdminZkClient(zkClient))
   }
 
   @Test
-  def shouldAddBrokerConfig(): Unit = {
-    val createOpts = new ConfigCommandOptions(Array("--zookeeper", zkConnect,
+  def shouldAlterTopicConfig(): Unit = {
+    doShouldAlterTopicConfig(false)
+  }
+
+  @Test
+  def shouldAlterTopicConfigFile(): Unit = {
+    doShouldAlterTopicConfig(true)
+  }
+
+  def doShouldAlterTopicConfig(file: Boolean): Unit = {
+    var filePath = ""
+    val addedConfigs = Seq("delete.retention.ms=1000000", "min.insync.replicas=2")
+    if (file) {
+      val file = TestUtils.tempFile(addedConfigs.mkString("\n"))
+      filePath = file.getPath
+    }
+
+    val resourceName = "my-topic"
+    val alterOpts = new ConfigCommandOptions(Array("--bootstrap-server", "localhost:9092",
+      "--entity-name", resourceName,
+      "--entity-type", "topics",
+      "--alter",
+      if (file) "--add-config-file" else "--add-config",
+      if (file) filePath else addedConfigs.mkString(","),
+      "--delete-config", "unclean.leader.election.enable"))
+    var alteredConfigs = false
+
+    def newConfigEntry(name: String, value: String): ConfigEntry =
+      ConfigTest.newConfigEntry(name, value, ConfigEntry.ConfigSource.DYNAMIC_TOPIC_CONFIG, false, false, List.empty[ConfigEntry.ConfigSynonym].asJava)
+
+    val resource = new ConfigResource(ConfigResource.Type.TOPIC, resourceName)
+    val configEntries = List(newConfigEntry("min.insync.replicas", "1"), newConfigEntry("unclean.leader.election.enable", "1")).asJava
+    val future = new KafkaFutureImpl[util.Map[ConfigResource, Config]]
+    future.complete(util.Collections.singletonMap(resource, new Config(configEntries)))
+    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 node = new Node(1, "localhost", 9092)
+    val mockAdminClient = new MockAdminClient(util.Collections.singletonList(node), node) {
+      override def describeConfigs(resources: util.Collection[ConfigResource], options: DescribeConfigsOptions): DescribeConfigsResult = {
+        assertFalse("Config synonyms requested unnecessarily", options.includeSynonyms())
+        assertEquals(1, resources.size)
+        val resource = resources.iterator.next
+        assertEquals(resource.`type`, ConfigResource.Type.TOPIC)
+        assertEquals(resource.name, resourceName)
+        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.TOPIC, resource.`type`)
+        assertEquals(3, alterConfigOps.size)
+
+        val expectedConfigOps = Set(
+          new AlterConfigOp(newConfigEntry("delete.retention.ms", "1000000"), AlterConfigOp.OpType.SET),
+          new AlterConfigOp(newConfigEntry("min.insync.replicas", "2"), AlterConfigOp.OpType.SET),
+          new AlterConfigOp(newConfigEntry("unclean.leader.election.enable", ""), AlterConfigOp.OpType.DELETE)
+        )
+        assertEquals(expectedConfigOps, alterConfigOps.asScala.toSet)
+        alteredConfigs = true
+        alterResult
+      }
+    }
+    EasyMock.replay(alterResult, describeResult)
+    ConfigCommand.alterConfig(mockAdminClient, alterOpts)
+    assertTrue(alteredConfigs)
+    EasyMock.reset(alterResult, describeResult)
+  }
+
+  @Test
+  def shouldDescribeConfigSynonyms(): Unit = {
+    val resourceName = "my-topic"
+    val describeOpts = new ConfigCommandOptions(Array("--bootstrap-server", "localhost:9092",
+      "--entity-name", resourceName,
+      "--entity-type", "topics",
+      "--describe",
+      "--all"))
+
+    val resource = new ConfigResource(ConfigResource.Type.TOPIC, resourceName)
+    val future = new KafkaFutureImpl[util.Map[ConfigResource, Config]]
+    future.complete(util.Collections.singletonMap(resource, new Config(util.Collections.emptyList[ConfigEntry])))
+    val describeResult: DescribeConfigsResult = EasyMock.createNiceMock(classOf[DescribeConfigsResult])
+    EasyMock.expect(describeResult.all()).andReturn(future).once()
+
+    val node = new Node(1, "localhost", 9092)
+    val mockAdminClient = new MockAdminClient(util.Collections.singletonList(node), node) {
+      override def describeConfigs(resources: util.Collection[ConfigResource], options: DescribeConfigsOptions): DescribeConfigsResult = {
+        assertTrue("Synonyms not requested", options.includeSynonyms())
+        assertEquals(Set(resource), resources.asScala.toSet)
+        describeResult
+      }
+    }
+    EasyMock.replay(describeResult)
+    ConfigCommand.describeConfig(mockAdminClient, describeOpts)
+    EasyMock.reset(describeResult)
+  }
+
+  @Test
+  def shouldAddBrokerQuotaConfig(): Unit = {
+    val alterOpts = new ConfigCommandOptions(Array("--zookeeper", zkConnect,
       "--entity-name", "1",
       "--entity-type", "brokers",
       "--alter",
-      "--add-config", "a=b,c=d"))
+      "--add-config", "leader.replication.throttled.rate=10,follower.replication.throttled.rate=20"))
 
-    case class TestAdminZkClient(val zkClient: KafkaZkClient) extends AdminZkClient(zkClient) {
+    class TestAdminZkClient(zkClient: KafkaZkClient) extends AdminZkClient(zkClient) {
       override def changeBrokerConfig(brokerIds: Seq[Int], configChange: Properties): Unit = {
         assertEquals(Seq(1), brokerIds)
-        assertEquals("b", configChange.get("a"))
-        assertEquals("d", configChange.get("c"))
+        assertEquals("10", configChange.get("leader.replication.throttled.rate"))
+        assertEquals("20", configChange.get("follower.replication.throttled.rate"))
+      }
+    }
+
+    ConfigCommand.alterConfigWithZk(null, alterOpts, new TestAdminZkClient(zkClient))
+  }
+
+  @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", zkConnect,
+      "--entity-type", ConfigType.Broker,
+      "--describe"
+    )
+
+    new ConfigCommandOptions(optsList.toArray).checkArgs()
+  }
+
+  @Test
+  def testNoSpecifiedEntityOptionWithDescribeBrokersInBootstrapServerIsAllowed(): Unit = {
+    val optsList = List("--bootstrap-server", "localhost:9092",
+      "--entity-type", ConfigType.Broker,
+      "--describe"
+    )
+
+    new ConfigCommandOptions(optsList.toArray).checkArgs()
+  }
+
+  @Test
+  def testDescribeAllBrokerConfig(): Unit = {
+    val optsList = List("--bootstrap-server", "localhost:9092",
+      "--entity-type", ConfigType.Broker,
+      "--entity-name", "1",
+      "--describe",
+      "--all")
+
+    new ConfigCommandOptions(optsList.toArray).checkArgs()
+  }
+
+  @Test
+  def testDescribeAllTopicConfig(): Unit = {
+    val optsList = List("--bootstrap-server", "localhost:9092",
+      "--entity-type", ConfigType.Topic,
+      "--entity-name", "foo",
+      "--describe",
+      "--all")
+
+    new ConfigCommandOptions(optsList.toArray).checkArgs()
+  }
+
+  @Test(expected = classOf[IllegalArgumentException])
+  def testDescribeAllBrokerConfigBootstrapServerRequired(): Unit = {
+    val optsList = List("--zookeeper", zkConnect,
+      "--entity-type", ConfigType.Broker,
+      "--entity-name", "1",
+      "--describe",
+      "--all")
+
+    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)
+    verifyAlterBrokerConfig(node, "", List("--entity-default"))
+  }
+
+  @Test
+  def shouldAddBrokerDynamicConfig(): Unit = {
+    val node = new Node(1, "localhost", 9092)
+    verifyAlterBrokerConfig(node, "1", List("--entity-name", "1"))
+  }
+
+  def verifyAlterBrokerConfig(node: Node, resourceName: String, resourceOpts: List[String]): Unit = {
+    val optsList = List("--bootstrap-server", "localhost:9092",
+      "--entity-type", "brokers",
+      "--alter",
+      "--add-config", "message.max.bytes=10") ++ resourceOpts
+    val alterOpts = new ConfigCommandOptions(optsList.toArray)
+    val brokerConfigs = mutable.Map[String, String]("num.io.threads" -> "5")
+
+    val resource = new ConfigResource(ConfigResource.Type.BROKER, resourceName)
+    val configEntries = util.Collections.singletonList(new ConfigEntry("num.io.threads", "5"))
+    val future = new KafkaFutureImpl[util.Map[ConfigResource, Config]]
+    future.complete(util.Collections.singletonMap(resource, new Config(configEntries)))
+    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], options: DescribeConfigsOptions): DescribeConfigsResult = {
+        assertFalse("Config synonyms requested unnecessarily", options.includeSynonyms())
+        assertEquals(1, resources.size)
+        val resource = resources.iterator.next
+        assertEquals(ConfigResource.Type.BROKER, resource.`type`)
+        assertEquals(resourceName, resource.name)
+        describeResult
+      }
+
+      override def alterConfigs(configs: util.Map[ConfigResource, Config], options: AlterConfigsOptions): AlterConfigsResult = {
+        assertEquals(1, configs.size)
+        val entry = configs.entrySet.iterator.next
+        val resource = entry.getKey
+        val config = entry.getValue
+        assertEquals(ConfigResource.Type.BROKER, resource.`type`)
+        config.entries.forEach { e => brokerConfigs.put(e.name, e.value) }
+        alterResult
+      }
+    }
+    EasyMock.replay(alterResult, describeResult)
+    ConfigCommand.alterConfig(mockAdminClient, alterOpts)
+    assertEquals(Map("message.max.bytes" -> "10", "num.io.threads" -> "5"), brokerConfigs.toMap)
+    EasyMock.reset(alterResult, describeResult)
+  }
+
+  @Test
+  def shouldDescribeConfigBrokerWithoutEntityName(): Unit = {
+    val describeOpts = new ConfigCommandOptions(Array("--bootstrap-server", "localhost:9092",
+      "--entity-type", "brokers",
+      "--describe"))
+
+    val BrokerDefaultEntityName = ""
+    val resourceCustom = new ConfigResource(ConfigResource.Type.BROKER, "1")
+    val resourceDefault = new ConfigResource(ConfigResource.Type.BROKER, BrokerDefaultEntityName)
+    val future = new KafkaFutureImpl[util.Map[ConfigResource, Config]]
+    val emptyConfig = new Config(util.Collections.emptyList[ConfigEntry])
+    val resultMap = Map(resourceCustom -> emptyConfig, resourceDefault -> emptyConfig).asJava
+    future.complete(resultMap)
+    val describeResult: DescribeConfigsResult = EasyMock.createNiceMock(classOf[DescribeConfigsResult])
+    // make sure it will be called 2 times: (1) for broker "1" (2) for default broker ""
+    EasyMock.expect(describeResult.all()).andReturn(future).times(2)
+
+    val node = new Node(1, "localhost", 9092)
+    val mockAdminClient = new MockAdminClient(util.Collections.singletonList(node), node) {
+      override def describeConfigs(resources: util.Collection[ConfigResource], options: DescribeConfigsOptions): DescribeConfigsResult = {
+        assertTrue("Synonyms not requested", options.includeSynonyms())
+        val resource = resources.iterator.next
+        assertEquals(ConfigResource.Type.BROKER, resource.`type`)
+        assertTrue(resourceCustom.name == resource.name || resourceDefault.name == resource.name)
+        assertEquals(1, resources.size)
+        describeResult
       }
     }
+    EasyMock.replay(describeResult)
+    ConfigCommand.describeConfig(mockAdminClient, describeOpts)
+    EasyMock.verify(describeResult)
+    EasyMock.reset(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], options: DescribeConfigsOptions): DescribeConfigsResult = {
+        assertEquals(1, resources.size)
+        val resource = resources.iterator.next
+        assertEquals(ConfigResource.Type.BROKER_LOGGER, resource.`type`)
+        assertEquals(resourceName, resource.name)
+        describeResult
+      }
 
-    ConfigCommand.alterConfig(null, createOpts, new TestAdminZkClient(zkClient))
+      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.alterConfig(mockAdminClient, alterOpts)
+    assertTrue(alteredConfigs)
+    EasyMock.reset(alterResult, describeResult)
   }
 
   @Test
-  def shouldSupportCommaSeparatedValues(): Unit = {
+  def shouldSupportCommaSeparatedValuesUsingZookeeper(): Unit = {
     val createOpts = new ConfigCommandOptions(Array("--zookeeper", zkConnect,
       "--entity-name", "my-topic",
       "--entity-type", "topics",
       "--alter",
       "--add-config", "a=b,c=[d,e ,f],g=[h,i]"))
 
-    case class TestAdminZkClient(val zkClient: KafkaZkClient) extends AdminZkClient(zkClient) {
-      override def changeBrokerConfig(brokerIds: Seq[Int], configChange: Properties): Unit = {
-        assertEquals(Seq(1), brokerIds)
+    class TestAdminZkClient(zkClient: KafkaZkClient) extends AdminZkClient(zkClient) {
+      override def changeTopicConfig(topic: String, configChange: Properties): Unit = {
+        assertEquals("my-topic", topic)
         assertEquals("b", configChange.get("a"))
         assertEquals("d,e ,f", configChange.get("c"))
         assertEquals("h,i", configChange.get("g"))
       }
-
-      override def changeTopicConfig(topic: String, configs: Properties): Unit = {}
     }
 
-    ConfigCommand.alterConfig(null, createOpts, new TestAdminZkClient(zkClient))
+    ConfigCommand.alterConfigWithZk(null, createOpts, new TestAdminZkClient(zkClient))
   }
 
   @Test (expected = classOf[IllegalArgumentException])
-  def shouldNotUpdateBrokerConfigIfMalformedEntityName(): Unit = {
+  def shouldNotUpdateBrokerConfigIfMalformedEntityNameUsingZookeeper(): Unit = {
     val createOpts = new ConfigCommandOptions(Array("--zookeeper", zkConnect,
       "--entity-name", "1,2,3", //Don't support multiple brokers currently
       "--entity-type", "brokers",
       "--alter",
-      "--add-config", "a=b"))
-    ConfigCommand.alterConfig(null, createOpts, new DummyAdminZkClient(zkClient))
+      "--add-config", "leader.replication.throttled.rate=10"))
+    ConfigCommand.alterConfigWithZk(null, createOpts, new DummyAdminZkClient(zkClient))
+  }
+
+  @Test (expected = classOf[IllegalArgumentException])
+  def shouldNotUpdateBrokerConfigIfMalformedEntityName(): Unit = {
+    val createOpts = new ConfigCommandOptions(Array("--bootstrap-server", "localhost:9092",
+      "--entity-name", "1,2,3", //Don't support multiple brokers currently
+      "--entity-type", "brokers",
+      "--alter",
+      "--add-config", "leader.replication.throttled.rate=10"))
+    ConfigCommand.alterConfig(new DummyAdminClient(new Node(1, "localhost", 9092)), createOpts)
+  }
+
+  @Test
+  def testDynamicBrokerConfigUpdateUsingZooKeeper(): Unit = {
+    val brokerId = "1"
+    val adminZkClient = new AdminZkClient(zkClient)
+    val alterOpts = Array("--zookeeper", zkConnect, "--entity-type", "brokers", "--alter")
+
+    def entityOpt(brokerId: Option[String]): Array[String] = {
+      brokerId.map(id => Array("--entity-name", id)).getOrElse(Array("--entity-default"))
+    }
+
+    def alterConfigWithZk(configs: Map[String, String], brokerId: Option[String],
+                          encoderConfigs: Map[String, String] = Map.empty): Unit = {
+      val configStr = (configs ++ encoderConfigs).map { case (k, v) => s"$k=$v" }.mkString(",")
+      val addOpts = new ConfigCommandOptions(alterOpts ++ entityOpt(brokerId) ++ Array("--add-config", configStr))
+      ConfigCommand.alterConfigWithZk(zkClient, addOpts, adminZkClient)
+    }
+
+    def verifyConfig(configs: Map[String, String], brokerId: Option[String]): Unit = {
+      val entityConfigs = zkClient.getEntityConfigs("brokers", brokerId.getOrElse(ConfigEntityName.Default))
+      assertEquals(configs, entityConfigs.asScala)
+    }
+
+    def alterAndVerifyConfig(configs: Map[String, String], brokerId: Option[String]): Unit = {
+      alterConfigWithZk(configs, brokerId)
+      verifyConfig(configs, brokerId)
+    }
+
+    def deleteAndVerifyConfig(configNames: Set[String], brokerId: Option[String]): Unit = {
+      val deleteOpts = new ConfigCommandOptions(alterOpts ++ entityOpt(brokerId) ++
+        Array("--delete-config", configNames.mkString(",")))
+      ConfigCommand.alterConfigWithZk(zkClient, deleteOpts, adminZkClient)
+      verifyConfig(Map.empty, brokerId)
+    }
+
+    // Add config
+    alterAndVerifyConfig(Map("message.max.size" -> "110000"), Some(brokerId))
+    alterAndVerifyConfig(Map("message.max.size" -> "120000"), None)
+
+    // Change config
+    alterAndVerifyConfig(Map("message.max.size" -> "130000"), Some(brokerId))
+    alterAndVerifyConfig(Map("message.max.size" -> "140000"), None)
+
+    // Delete config
+    deleteAndVerifyConfig(Set("message.max.size"), Some(brokerId))
+    deleteAndVerifyConfig(Set("message.max.size"), None)
+
+    // Listener configs: should work only with listener name
+    alterAndVerifyConfig(Map("listener.name.external.ssl.keystore.location" -> "/tmp/test.jks"), Some(brokerId))
+    intercept[ConfigException](alterConfigWithZk(Map("ssl.keystore.location" -> "/tmp/test.jks"), Some(brokerId)))
+
+    // Per-broker config configured at default cluster-level should fail
+    intercept[ConfigException](alterConfigWithZk(Map("listener.name.external.ssl.keystore.location" -> "/tmp/test.jks"), None))
+    deleteAndVerifyConfig(Set("listener.name.external.ssl.keystore.location"), Some(brokerId))
+
+    // Password config update without encoder secret should fail
+    intercept[IllegalArgumentException](alterConfigWithZk(Map("listener.name.external.ssl.keystore.password" -> "secret"), Some(brokerId)))
+
+    // Password config update with encoder secret should succeed and encoded password must be stored in ZK
+    val configs = Map("listener.name.external.ssl.keystore.password" -> "secret", "log.cleaner.threads" -> "2")
+    val encoderConfigs = Map(KafkaConfig.PasswordEncoderSecretProp -> "encoder-secret")
+    alterConfigWithZk(configs, Some(brokerId), encoderConfigs)
+    val brokerConfigs = zkClient.getEntityConfigs("brokers", brokerId)
+    assertFalse("Encoder secret stored in ZooKeeper", brokerConfigs.contains(KafkaConfig.PasswordEncoderSecretProp))
+    assertEquals("2", brokerConfigs.getProperty("log.cleaner.threads")) // not encoded
+    val encodedPassword = brokerConfigs.getProperty("listener.name.external.ssl.keystore.password")
+    val passwordEncoder = ConfigCommand.createPasswordEncoder(encoderConfigs)
+    assertEquals("secret", passwordEncoder.decode(encodedPassword).value)
+    assertEquals(configs.size, brokerConfigs.size)
+
+    // Password config update with overrides for encoder parameters
+    val configs2 = Map("listener.name.internal.ssl.keystore.password" -> "secret2")
+    val encoderConfigs2 = Map(KafkaConfig.PasswordEncoderSecretProp -> "encoder-secret",
+      KafkaConfig.PasswordEncoderCipherAlgorithmProp -> "DES/CBC/PKCS5Padding",
+      KafkaConfig.PasswordEncoderIterationsProp -> "1024",
+      KafkaConfig.PasswordEncoderKeyFactoryAlgorithmProp -> "PBKDF2WithHmacSHA1",
+      KafkaConfig.PasswordEncoderKeyLengthProp -> "64")
+    alterConfigWithZk(configs2, Some(brokerId), encoderConfigs2)
+    val brokerConfigs2 = zkClient.getEntityConfigs("brokers", brokerId)
+    val encodedPassword2 = brokerConfigs2.getProperty("listener.name.internal.ssl.keystore.password")
+    assertEquals("secret2", ConfigCommand.createPasswordEncoder(encoderConfigs).decode(encodedPassword2).value)
+    assertEquals("secret2", ConfigCommand.createPasswordEncoder(encoderConfigs2).decode(encodedPassword2).value)
+
+
+    // Password config update at default cluster-level should fail
+    intercept[ConfigException](alterConfigWithZk(configs, None, encoderConfigs))
+
+    // Dynamic config updates using ZK should fail if broker is running.
+    registerBrokerInZk(brokerId.toInt)
+    intercept[IllegalArgumentException](alterConfigWithZk(Map("message.max.size" -> "210000"), Some(brokerId)))
+    intercept[IllegalArgumentException](alterConfigWithZk(Map("message.max.size" -> "220000"), None))
+
+    // Dynamic config updates using ZK should for a different broker that is not running should succeed
+    alterAndVerifyConfig(Map("message.max.size" -> "230000"), Some("2"))
+  }
+
+  @Test (expected = classOf[IllegalArgumentException])
+  def shouldNotUpdateBrokerConfigIfMalformedConfigUsingZookeeper(): Unit = {
+    val createOpts = new ConfigCommandOptions(Array("--zookeeper", zkConnect,
+      "--entity-name", "1",
+      "--entity-type", "brokers",
+      "--alter",
+      "--add-config", "a=="))
+    ConfigCommand.alterConfigWithZk(null, createOpts, new DummyAdminZkClient(zkClient))
   }
 
   @Test (expected = classOf[IllegalArgumentException])
   def shouldNotUpdateBrokerConfigIfMalformedConfig(): Unit = {
+    val createOpts = new ConfigCommandOptions(Array("--bootstrap-server", "localhost:9092",
+      "--entity-name", "1",
+      "--entity-type", "brokers",
+      "--alter",
+      "--add-config", "a=="))
+    ConfigCommand.alterConfig(new DummyAdminClient(new Node(1, "localhost", 9092)), createOpts)
+  }
+
+  @Test (expected = classOf[IllegalArgumentException])
+  def shouldNotUpdateBrokerConfigIfMalformedBracketConfigUsingZookeeper(): Unit = {
     val createOpts = new ConfigCommandOptions(Array("--zookeeper", zkConnect,
       "--entity-name", "1",
       "--entity-type", "brokers",
       "--alter",
-      "--add-config", "a="))
-    ConfigCommand.alterConfig(null, createOpts, new DummyAdminZkClient(zkClient))
+      "--add-config", "a=[b,c,d=e"))
+    ConfigCommand.alterConfigWithZk(null, createOpts, new DummyAdminZkClient(zkClient))
   }
 
   @Test (expected = classOf[IllegalArgumentException])
   def shouldNotUpdateBrokerConfigIfMalformedBracketConfig(): Unit = {
-    val createOpts = new ConfigCommandOptions(Array("--zookeeper", zkConnect,
+    val createOpts = new ConfigCommandOptions(Array("--bootstrap-server", "localhost:9092",
       "--entity-name", "1",
       "--entity-type", "brokers",
       "--alter",
       "--add-config", "a=[b,c,d=e"))
-    ConfigCommand.alterConfig(null, createOpts, new DummyAdminZkClient(zkClient))
+    ConfigCommand.alterConfig(new DummyAdminClient(new Node(1, "localhost", 9092)), createOpts)
   }
 
-  @Test (expected = classOf[InvalidConfigException])
-  def shouldNotUpdateBrokerConfigIfNonExistingConfigIsDeleted(): Unit = {
+  @Test (expected = classOf[InvalidConfigurationException])
+  def shouldNotUpdateConfigIfNonExistingConfigIsDeletedUsingZookeper(): Unit = {
     val createOpts = new ConfigCommandOptions(Array("--zookeeper", zkConnect,
       "--entity-name", "my-topic",
       "--entity-type", "topics",
       "--alter",
       "--delete-config", "missing_config1, missing_config2"))
-    ConfigCommand.alterConfig(null, createOpts, new DummyAdminZkClient(zkClient))
+    ConfigCommand.alterConfigWithZk(null, createOpts, new DummyAdminZkClient(zkClient))
+  }
+
+  @Test (expected = classOf[InvalidConfigurationException])
+  def shouldNotUpdateConfigIfNonExistingConfigIsDeleted(): Unit = {
+    val resourceName = "my-topic"
+    val createOpts = new ConfigCommandOptions(Array("--bootstrap-server", "localhost:9092",
+      "--entity-name", resourceName,
+      "--entity-type", "topics",
+      "--alter",
+      "--delete-config", "missing_config1, missing_config2"))
+
+    val resource = new ConfigResource(ConfigResource.Type.TOPIC, resourceName)
+    val configEntries = List.empty[ConfigEntry].asJava
+    val future = new KafkaFutureImpl[util.Map[ConfigResource, Config]]
+    future.complete(util.Collections.singletonMap(resource, new Config(configEntries)))
+    val describeResult: DescribeConfigsResult = EasyMock.createNiceMock(classOf[DescribeConfigsResult])
+    EasyMock.expect(describeResult.all()).andReturn(future).once()
+
+    val node = new Node(1, "localhost", 9092)
+    val mockAdminClient = new MockAdminClient(util.Collections.singletonList(node), node) {
+      override def describeConfigs(resources: util.Collection[ConfigResource], options: DescribeConfigsOptions): DescribeConfigsResult = {
+        assertEquals(1, resources.size)
+        val resource = resources.iterator.next
+        assertEquals(resource.`type`, ConfigResource.Type.TOPIC)
+        assertEquals(resource.name, resourceName)
+        describeResult
+      }
+    }
+
+    EasyMock.replay(describeResult)
+    ConfigCommand.alterConfig(mockAdminClient, createOpts)
+    EasyMock.reset(describeResult)
   }
 
   @Test
@@ -225,7 +1332,7 @@ class ConfigCommandTest extends ZooKeeperTestHarness with Logging {
       "--alter",
       "--delete-config", "a,c"))
 
-    case class TestAdminZkClient(val zkClient: KafkaZkClient) extends AdminZkClient(zkClient) {
+    class TestAdminZkClient(zkClient: KafkaZkClient) extends AdminZkClient(zkClient) {
       override def fetchEntityConfig(entityType: String, entityName: String): Properties = {
         val properties: Properties = new Properties
         properties.put("a", "b")
@@ -240,7 +1347,7 @@ class ConfigCommandTest extends ZooKeeperTestHarness with Logging {
       }
     }
 
-    ConfigCommand.alterConfig(null, createOpts, new TestAdminZkClient(zkClient))
+    ConfigCommand.alterConfigWithZk(null, createOpts, new TestAdminZkClient(zkClient))
   }
 
   @Test
@@ -260,7 +1367,7 @@ class ConfigCommandTest extends ZooKeeperTestHarness with Logging {
         "--delete-config", mechanism))
 
     val credentials = mutable.Map[String, Properties]()
-    case class CredentialChange(val user: String, val mechanisms: Set[String], val iterations: Int) extends AdminZkClient(zkClient) {
+    case class CredentialChange(user: String, mechanisms: Set[String], iterations: Int) extends AdminZkClient(zkClient) {
       override def fetchEntityConfig(entityType: String, entityName: String): Properties = {
         credentials.getOrElse(entityName, new Properties())
       }
@@ -278,22 +1385,24 @@ class ConfigCommandTest extends ZooKeeperTestHarness with Logging {
       }
     }
     val optsA = createOpts("userA", "SCRAM-SHA-256=[iterations=8192,password=abc, def]")
-    ConfigCommand.alterConfig(null, optsA, CredentialChange("userA", Set("SCRAM-SHA-256"), 8192))
+    ConfigCommand.alterConfigWithZk(null, optsA, CredentialChange("userA", Set("SCRAM-SHA-256"), 8192))
     val optsB = createOpts("userB", "SCRAM-SHA-256=[iterations=4096,password=abc, def],SCRAM-SHA-512=[password=1234=abc]")
-    ConfigCommand.alterConfig(null, optsB, CredentialChange("userB", Set("SCRAM-SHA-256", "SCRAM-SHA-512"), 4096))
+    ConfigCommand.alterConfigWithZk(null, optsB, CredentialChange("userB", Set("SCRAM-SHA-256", "SCRAM-SHA-512"), 4096))
 
     val del256 = deleteOpts("userB", "SCRAM-SHA-256")
-    ConfigCommand.alterConfig(null, del256, CredentialChange("userB", Set("SCRAM-SHA-512"), 4096))
+    ConfigCommand.alterConfigWithZk(null, del256, CredentialChange("userB", Set("SCRAM-SHA-512"), 4096))
     val del512 = deleteOpts("userB", "SCRAM-SHA-512")
-    ConfigCommand.alterConfig(null, del512, CredentialChange("userB", Set(), 4096))
+    ConfigCommand.alterConfigWithZk(null, del512, CredentialChange("userB", Set(), 4096))
   }
 
-  @Test
-  def testQuotaConfigEntity() {
+  def doTestQuotaConfigEntity(zkConfig: Boolean): Unit = {
+    val connectOpts = if (zkConfig)
+      ("--zookeeper", zkConnect)
+    else
+      ("--bootstrap-server", "localhost:9092")
 
     def createOpts(entityType: String, entityName: Option[String], otherArgs: Array[String]) : ConfigCommandOptions = {
-      val optArray = Array("--zookeeper", zkConnect,
-                           "--entity-type", entityType)
+      val optArray = Array(connectOpts._1, connectOpts._2, "--entity-type", entityType)
       val nameArray = entityName match {
         case Some(name) => Array("--entity-name", name)
         case None => Array[String]()
@@ -301,7 +1410,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)
@@ -309,7 +1418,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()
@@ -361,9 +1470,23 @@ class ConfigCommandTest extends ZooKeeperTestHarness with Logging {
   }
 
   @Test
-  def testUserClientQuotaOpts() {
-    def checkEntity(expectedEntityType: String, expectedEntityName: String, args: String*) {
-      val opts = new ConfigCommandOptions(Array("--zookeeper", zkConnect) ++ args)
+  def testQuotaConfigEntityUsingZookeeper(): Unit = {
+    doTestQuotaConfigEntity(zkConfig = true)
+  }
+
+  @Test
+  def testQuotaConfigEntity(): Unit = {
+    doTestQuotaConfigEntity(zkConfig = false)
+  }
+
+  def doTestUserClientQuotaOpts(zkConfig: Boolean): Unit = {
+    val connectOpts = if (zkConfig)
+      ("--zookeeper", zkConnect)
+    else
+      ("--bootstrap-server", "localhost:9092")
+
+    def checkEntity(expectedEntityType: String, expectedEntityName: String, args: String*): Unit = {
+      val opts = new ConfigCommandOptions(Array(connectOpts._1, connectOpts._2) ++ args)
       opts.checkArgs()
       val entity = ConfigCommand.parseEntity(opts)
       assertEquals(expectedEntityType, entity.root.entityType)
@@ -378,7 +1501,6 @@ class ConfigCommandTest extends ZooKeeperTestHarness with Logging {
         "--entity-type", "clients", "--entity-name", "",
         "--alter", "--add-config", "a=b,c=d")
 
-
     checkEntity("users", Sanitizer.sanitize("CN=user1") + "/clients/client1",
         "--entity-type", "users", "--entity-name", "CN=user1", "--entity-type", "clients", "--entity-name", "client1",
         "--alter", "--add-config", "a=b,c=d")
@@ -403,10 +1525,20 @@ class ConfigCommandTest extends ZooKeeperTestHarness with Logging {
   }
 
   @Test
-  def testQuotaDescribeEntities() {
-    val zkClient = EasyMock.createNiceMock(classOf[KafkaZkClient])
+  def testUserClientQuotaOptsUsingZookeeper(): Unit = {
+    doTestUserClientQuotaOpts(zkConfig = true)
+  }
+
+  @Test
+  def testUserClientQuotaOpts(): Unit = {
+    doTestUserClientQuotaOpts(zkConfig = false)
+  }
+
+  @Test
+  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)
@@ -464,7 +1596,15 @@ class ConfigCommandTest extends ZooKeeperTestHarness with Logging {
         Seq("/clients/client-3", sanitizedPrincipal + "/clients/client-2"))
   }
 
-  case class DummyAdminZkClient(zkClient: KafkaZkClient) extends AdminZkClient(zkClient) {
+  private def registerBrokerInZk(id: Int): Unit = {
+    zkClient.createTopLevelPaths()
+    val securityProtocol = SecurityProtocol.PLAINTEXT
+    val endpoint = new EndPoint("localhost", 9092, ListenerName.forSecurityProtocol(securityProtocol), securityProtocol)
+    val brokerInfo = BrokerInfo(Broker(id, Seq(endpoint), rack = None), ApiVersion.latestVersion, jmxPort = 9192)
+    zkClient.registerBroker(brokerInfo)
+  }
+
+  class DummyAdminZkClient(zkClient: KafkaZkClient) extends AdminZkClient(zkClient) {
     override def changeBrokerConfig(brokerIds: Seq[Int], configs: Properties): Unit = {}
     override def fetchEntityConfig(entityType: String, entityName: String): Properties = {new Properties}
     override def changeClientIdConfig(clientId: String, configs: Properties): Unit = {}
@@ -472,4 +1612,17 @@ class ConfigCommandTest extends ZooKeeperTestHarness with Logging {
     override def changeTopicConfig(topic: String, configs: Properties): Unit = {}
   }
 
+  class DummyAdminClient(node: Node) extends MockAdminClient(util.Collections.singletonList(node), node) {
+    override def describeConfigs(resources: util.Collection[ConfigResource], options: DescribeConfigsOptions): DescribeConfigsResult =
+      EasyMock.createNiceMock(classOf[DescribeConfigsResult])
+    override def incrementalAlterConfigs(configs: util.Map[ConfigResource, util.Collection[AlterConfigOp]],
+      options: AlterConfigsOptions): AlterConfigsResult = EasyMock.createNiceMock(classOf[AlterConfigsResult])
+    override def alterConfigs(configs: util.Map[ConfigResource, Config], options: AlterConfigsOptions): AlterConfigsResult =
+      EasyMock.createNiceMock(classOf[AlterConfigsResult])
+    override def describeClientQuotas(filter: ClientQuotaFilter, options: DescribeClientQuotasOptions): DescribeClientQuotasResult =
+      EasyMock.createNiceMock(classOf[DescribeClientQuotasResult])
+    override def alterClientQuotas(entries: util.Collection[ClientQuotaAlteration],
+      options: AlterClientQuotasOptions): AlterClientQuotasResult =
+      EasyMock.createNiceMock(classOf[AlterClientQuotasResult])
+  }
 }
diff --git a/core/src/test/scala/unit/kafka/admin/ConsumerGroupCommandTest.scala b/core/src/test/scala/unit/kafka/admin/ConsumerGroupCommandTest.scala
new file mode 100644
index 0000000000000..3a6405234eaab
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/admin/ConsumerGroupCommandTest.scala
@@ -0,0 +1,211 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package kafka.admin
+
+import java.time.Duration
+import java.util.concurrent.{ExecutorService, Executors, TimeUnit}
+import java.util.{Collections, Properties}
+
+import kafka.admin.ConsumerGroupCommand.{ConsumerGroupCommandOptions, ConsumerGroupService}
+import kafka.integration.KafkaServerTestHarness
+import kafka.server.KafkaConfig
+import kafka.utils.TestUtils
+import org.apache.kafka.clients.admin.AdminClientConfig
+import org.apache.kafka.clients.consumer.{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.jdk.CollectionConverters._
+import scala.collection.mutable.ArrayBuffer
+
+class ConsumerGroupCommandTest extends KafkaServerTestHarness {
+  import ConsumerGroupCommandTest._
+
+  val topic = "foo"
+  val group = "test.group"
+
+  private var consumerGroupService: List[ConsumerGroupService] = List()
+  private var consumerGroupExecutors: List[AbstractConsumerGroupExecutor] = List()
+
+  // configure the servers and clients
+  override def generateConfigs = {
+    TestUtils.createBrokerConfigs(1, zkConnect, enableControlledShutdown = false).map { props =>
+      KafkaConfig.fromProps(props)
+    }
+  }
+
+  @Before
+  override def setUp(): Unit = {
+    super.setUp()
+    createTopic(topic, 1, 1)
+  }
+
+  @After
+  override def tearDown(): Unit = {
+    consumerGroupService.foreach(_.close())
+    consumerGroupExecutors.foreach(_.shutdown())
+    super.tearDown()
+  }
+
+  def committedOffsets(topic: String = topic, group: String = group): collection.Map[TopicPartition, Long] = {
+    val consumer = createNoAutoCommitConsumer(group)
+    try {
+      val partitions: Set[TopicPartition] = consumer.partitionsFor(topic)
+        .asScala.toSet.map {partitionInfo : PartitionInfo => new TopicPartition(partitionInfo.topic, partitionInfo.partition)}
+      consumer.committed(partitions.asJava).asScala.filter(_._2 != null).map { case (k, v) => k -> v.offset }
+    } finally {
+      consumer.close()
+    }
+  }
+
+  def createNoAutoCommitConsumer(group: String): 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, Map(AdminClientConfig.RETRIES_CONFIG -> Int.MaxValue.toString))
+    consumerGroupService = service :: consumerGroupService
+    service
+  }
+
+  def addConsumerGroupExecutor(numConsumers: Int,
+                               topic: String = topic,
+                               group: String = group,
+                               strategy: String = classOf[RangeAssignor].getName,
+                               customPropsOpt: Option[Properties] = None,
+                               syncCommit: Boolean = false): ConsumerGroupExecutor = {
+    val executor = new ConsumerGroupExecutor(brokerList, numConsumers, group, topic, strategy, customPropsOpt, syncCommit)
+    addExecutor(executor)
+    executor
+  }
+
+  def addSimpleGroupExecutor(partitions: Iterable[TopicPartition] = Seq(new TopicPartition(topic, 0)),
+                             group: String = group): SimpleConsumerGroupExecutor = {
+    val executor = new SimpleConsumerGroupExecutor(brokerList, group, partitions)
+    addExecutor(executor)
+    executor
+  }
+
+  private def addExecutor(executor: AbstractConsumerGroupExecutor): AbstractConsumerGroupExecutor = {
+    consumerGroupExecutors = executor :: consumerGroupExecutors
+    executor
+  }
+
+}
+
+object ConsumerGroupCommandTest {
+
+  abstract class AbstractConsumerRunnable(broker: String, groupId: String, customPropsOpt: Option[Properties] = None,
+                                          syncCommit: Boolean = false) extends Runnable {
+    val props = new Properties
+    configure(props)
+    customPropsOpt.foreach(props.asScala ++= _.asScala)
+    val consumer = new KafkaConsumer(props)
+
+    def configure(props: Properties): Unit = {
+      props.put("bootstrap.servers", broker)
+      props.put("group.id", groupId)
+      props.put("key.deserializer", classOf[StringDeserializer].getName)
+      props.put("value.deserializer", classOf[StringDeserializer].getName)
+    }
+
+    def subscribe(): Unit
+
+    def run(): Unit = {
+      try {
+        subscribe()
+        while (true) {
+          consumer.poll(Duration.ofMillis(Long.MaxValue))
+          if (syncCommit)
+            consumer.commitSync()
+        }
+      } catch {
+        case _: WakeupException => // OK
+      } finally {
+        consumer.close()
+      }
+    }
+
+    def shutdown(): Unit = {
+      consumer.wakeup()
+    }
+  }
+
+  class ConsumerRunnable(broker: String, groupId: String, 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)
+      props.put("partition.assignment.strategy", strategy)
+    }
+
+    override def subscribe(): Unit = {
+      consumer.subscribe(Collections.singleton(topic))
+    }
+  }
+
+  class SimpleConsumerRunnable(broker: String, groupId: String, partitions: Iterable[TopicPartition])
+    extends AbstractConsumerRunnable(broker, groupId) {
+
+    override def subscribe(): Unit = {
+      consumer.assign(partitions.toList.asJava)
+    }
+  }
+
+  class AbstractConsumerGroupExecutor(numThreads: Int) {
+    private val executor: ExecutorService = Executors.newFixedThreadPool(numThreads)
+    private val consumers = new ArrayBuffer[AbstractConsumerRunnable]()
+
+    def submit(consumerThread: AbstractConsumerRunnable): Unit = {
+      consumers += consumerThread
+      executor.submit(consumerThread)
+    }
+
+    def shutdown(): Unit = {
+      consumers.foreach(_.shutdown())
+      executor.shutdown()
+      executor.awaitTermination(5000, TimeUnit.MILLISECONDS)
+    }
+  }
+
+  class ConsumerGroupExecutor(broker: String, numConsumers: Int, groupId: String, topic: String, strategy: String,
+                              customPropsOpt: Option[Properties] = None, syncCommit: Boolean = false)
+    extends AbstractConsumerGroupExecutor(numConsumers) {
+
+    for (_ <- 1 to numConsumers) {
+      submit(new ConsumerRunnable(broker, groupId, topic, strategy, customPropsOpt, syncCommit))
+    }
+
+  }
+
+  class SimpleConsumerGroupExecutor(broker: String, groupId: String, partitions: Iterable[TopicPartition])
+    extends AbstractConsumerGroupExecutor(1) {
+
+    submit(new SimpleConsumerRunnable(broker, groupId, partitions))
+  }
+
+}
+
diff --git a/core/src/test/scala/unit/kafka/admin/DelegationTokenCommandTest.scala b/core/src/test/scala/unit/kafka/admin/DelegationTokenCommandTest.scala
new file mode 100644
index 0000000000000..df68ca6a52c9f
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/admin/DelegationTokenCommandTest.scala
@@ -0,0 +1,147 @@
+/**
+  * Licensed to the Apache Software Foundation (ASF) under one or more
+  * contributor license agreements.  See the NOTICE file distributed with
+  * this work for additional information regarding copyright ownership.
+  * The ASF licenses this file to You under the Apache License, Version 2.0
+  * (the "License"); you may not use this file except in compliance with
+  * the License.  You may obtain a copy of the License at
+  *
+  * http://www.apache.org/licenses/LICENSE-2.0
+  *
+  * Unless required by applicable law or agreed to in writing, software
+  * distributed under the License is distributed on an "AS IS" BASIS,
+  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  * See the License for the specific language governing permissions and
+  * limitations under the License.
+  */
+package kafka.admin
+
+import java.util
+
+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.{Admin, 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.mutable.ListBuffer
+import scala.concurrent.ExecutionException
+
+class DelegationTokenCommandTest extends BaseRequestTest with SaslSetup {
+  override protected def securityProtocol = SecurityProtocol.SASL_PLAINTEXT
+  private val kafkaClientSaslMechanism = "PLAIN"
+  private val kafkaServerSaslMechanisms = List("PLAIN")
+  protected override val serverSaslProperties = Some(kafkaServerSaslProperties(kafkaServerSaslMechanisms, kafkaClientSaslMechanism))
+  protected override val clientSaslProperties = Some(kafkaClientSaslProperties(kafkaClientSaslMechanism))
+  var adminClient: Admin = null
+
+  override def brokerCount = 1
+
+  @Before
+  override def setUp(): Unit = {
+    startSasl(jaasSections(kafkaServerSaslMechanisms, Some(kafkaClientSaslMechanism), KafkaSasl, JaasTestUtils.KafkaServerContextName))
+    super.setUp()
+  }
+
+  override def generateConfigs = {
+    val props = TestUtils.createBrokerConfigs(brokerCount, zkConnect,
+      enableControlledShutdown = false,
+      interBrokerSecurityProtocol = Some(securityProtocol),
+      trustStoreFile = trustStoreFile, saslProperties = serverSaslProperties, enableToken = true)
+    props.foreach(brokerPropertyOverrides)
+    props.map(KafkaConfig.fromProps)
+  }
+
+  private def createAdminConfig: util.Map[String, Object] = {
+    val config = new util.HashMap[String, Object]
+    config.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList)
+    val securityProps: util.Map[Object, Object] =
+      TestUtils.adminClientSecurityConfigs(securityProtocol, trustStoreFile, clientSaslProperties)
+    securityProps.forEach { (key, value) => config.put(key.asInstanceOf[String], value) }
+    config
+  }
+
+  @Test
+  def testDelegationTokenRequests(): Unit = {
+    adminClient = Admin.create(createAdminConfig)
+    val renewer1 = "User:renewer1"
+    val renewer2 = "User:renewer2"
+
+    // create token1 with renewer1
+    val tokenCreated = DelegationTokenCommand.createToken(adminClient, getCreateOpts(List(renewer1)))
+
+    var tokens = DelegationTokenCommand.describeToken(adminClient, getDescribeOpts(List()))
+    assertTrue(tokens.size == 1)
+    val token1 = tokens.head
+    assertEquals(token1, tokenCreated)
+
+    // create token2 with renewer2
+    val token2 = DelegationTokenCommand.createToken(adminClient, getCreateOpts(List(renewer2)))
+
+    tokens = DelegationTokenCommand.describeToken(adminClient, getDescribeOpts(List()))
+    assertTrue(tokens.size == 2)
+    assertEquals(Set(token1, token2), tokens.toSet)
+
+    //get tokens for renewer2
+    tokens = DelegationTokenCommand.describeToken(adminClient, getDescribeOpts(List(renewer2)))
+    assertTrue(tokens.size == 1)
+    assertEquals(Set(token2), tokens.toSet)
+
+    //test renewing tokens
+    val expiryTimestamp = DelegationTokenCommand.renewToken(adminClient, getRenewOpts(token1.hmacAsBase64String()))
+    val renewedToken = DelegationTokenCommand.describeToken(adminClient, getDescribeOpts(List(renewer1))).head
+    assertEquals(expiryTimestamp, renewedToken.tokenInfo().expiryTimestamp())
+
+    //test expire tokens
+    DelegationTokenCommand.expireToken(adminClient, getExpireOpts(token1.hmacAsBase64String()))
+    DelegationTokenCommand.expireToken(adminClient, getExpireOpts(token2.hmacAsBase64String()))
+
+    tokens = DelegationTokenCommand.describeToken(adminClient, getDescribeOpts(List()))
+    assertTrue(tokens.size == 0)
+
+    //create token with invalid renewer principal type
+    intercept[ExecutionException](DelegationTokenCommand.createToken(adminClient, getCreateOpts(List("Group:Renewer3"))))
+
+    // try describing tokens for unknown owner
+    assertTrue(DelegationTokenCommand.describeToken(adminClient, getDescribeOpts(List("User:Unknown"))).isEmpty)
+  }
+
+  private def getCreateOpts(renewers: List[String]): DelegationTokenCommandOptions = {
+    val opts = ListBuffer("--bootstrap-server", brokerList, "--max-life-time-period", "-1",
+      "--command-config", "testfile", "--create")
+    renewers.foreach(renewer => opts ++= ListBuffer("--renewer-principal", renewer))
+    new DelegationTokenCommandOptions(opts.toArray)
+  }
+
+  private def getDescribeOpts(owners: List[String]): DelegationTokenCommandOptions = {
+    val opts = ListBuffer("--bootstrap-server", brokerList, "--command-config", "testfile", "--describe")
+    owners.foreach(owner => opts ++= ListBuffer("--owner-principal", owner))
+    new DelegationTokenCommandOptions(opts.toArray)
+  }
+
+  private def getRenewOpts(hmac: String): DelegationTokenCommandOptions = {
+    val opts = Array("--bootstrap-server", brokerList, "--command-config", "testfile", "--renew",
+      "--renew-time-period", "-1",
+      "--hmac", hmac)
+    new DelegationTokenCommandOptions(opts)
+  }
+
+  private def getExpireOpts(hmac: String): DelegationTokenCommandOptions = {
+    val opts = Array("--bootstrap-server", brokerList, "--command-config", "testfile", "--expire",
+      "--expiry-time-period", "-1",
+      "--hmac", hmac)
+    new DelegationTokenCommandOptions(opts)
+  }
+
+  @After
+  override def tearDown(): Unit = {
+    if (adminClient != null)
+      adminClient.close()
+    super.tearDown()
+    closeSasl()
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/admin/DeleteConsumerGroupTest.scala b/core/src/test/scala/unit/kafka/admin/DeleteConsumerGroupTest.scala
deleted file mode 100644
index ba11471e488e0..0000000000000
--- a/core/src/test/scala/unit/kafka/admin/DeleteConsumerGroupTest.scala
+++ /dev/null
@@ -1,220 +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.nio.charset.StandardCharsets
-
-import kafka.utils._
-import kafka.server.KafkaConfig
-import org.junit.Test
-import kafka.consumer._
-import org.apache.kafka.clients.producer.{KafkaProducer, ProducerRecord}
-import kafka.integration.KafkaServerTestHarness
-
-
-@deprecated("This test has been deprecated and will be removed in a future release.", "0.11.0.0")
-class DeleteConsumerGroupTest extends KafkaServerTestHarness {
-  def generateConfigs = TestUtils.createBrokerConfigs(3, zkConnect, false, true).map(KafkaConfig.fromProps)
-
-  @Test
-  def testGroupWideDeleteInZK() {
-    val topic = "test"
-    val groupToDelete = "groupToDelete"
-    val otherGroup = "otherGroup"
-
-    TestUtils.createTopic(zkUtils, topic, 1, 3, servers)
-    fillInConsumerGroupInfo(topic, groupToDelete, "consumer", 0, 10, false)
-    fillInConsumerGroupInfo(topic, otherGroup, "consumer", 0, 10, false)
-
-    AdminUtils.deleteConsumerGroupInZK(zkUtils, groupToDelete)
-
-    TestUtils.waitUntilTrue(() => !groupDirExists(new ZKGroupDirs(groupToDelete)),
-      "DeleteConsumerGroupInZK should delete the provided consumer group's directory")
-    TestUtils.waitUntilTrue(() => groupDirExists(new ZKGroupDirs(otherGroup)),
-      "DeleteConsumerGroupInZK should not delete unrelated consumer group directories")
-  }
-
-  @Test
-  def testGroupWideDeleteInZKDoesNothingForActiveConsumerGroup() {
-    val topic = "test"
-    val groupToDelete = "groupToDelete"
-    val otherGroup = "otherGroup"
-
-    TestUtils.createTopic(zkUtils, topic, 1, 3, servers)
-    fillInConsumerGroupInfo(topic, groupToDelete, "consumer", 0, 10, true)
-    fillInConsumerGroupInfo(topic, otherGroup, "consumer", 0, 10, false)
-
-    AdminUtils.deleteConsumerGroupInZK(zkUtils, groupToDelete)
-
-    TestUtils.waitUntilTrue(() => groupDirExists(new ZKGroupDirs(groupToDelete)),
-      "DeleteConsumerGroupInZK should not delete the provided consumer group's directory if the consumer group is still active")
-    TestUtils.waitUntilTrue(() => groupDirExists(new ZKGroupDirs(otherGroup)),
-      "DeleteConsumerGroupInZK should not delete unrelated consumer group directories")
-  }
-
-  @Test
-  def testGroupTopicWideDeleteInZKForGroupConsumingOneTopic() {
-    val topic = "test"
-    val groupToDelete = "groupToDelete"
-    val otherGroup = "otherGroup"
-    TestUtils.createTopic(zkUtils, topic, 1, 3, servers)
-    fillInConsumerGroupInfo(topic, groupToDelete, "consumer", 0, 10, false)
-    fillInConsumerGroupInfo(topic, otherGroup, "consumer", 0, 10, false)
-
-    AdminUtils.deleteConsumerGroupInfoForTopicInZK(zkUtils, groupToDelete, topic)
-
-    TestUtils.waitUntilTrue(() => !groupDirExists(new ZKGroupDirs(groupToDelete)),
-      "DeleteConsumerGroupInfoForTopicInZK should delete the provided consumer group's directory if it just consumes from one topic")
-    TestUtils.waitUntilTrue(() => groupTopicOffsetAndOwnerDirsExist(new ZKGroupTopicDirs(otherGroup, topic)),
-      "DeleteConsumerGroupInfoForTopicInZK should not delete unrelated consumer group owner and offset directories")
-  }
-
-  @Test
-  def testGroupTopicWideDeleteInZKForGroupConsumingMultipleTopics() {
-    val topicToDelete = "topicToDelete"
-    val otherTopic = "otherTopic"
-    val groupToDelete = "groupToDelete"
-    val otherGroup = "otherGroup"
-    TestUtils.createTopic(zkUtils, topicToDelete, 1, 3, servers)
-    TestUtils.createTopic(zkUtils, otherTopic, 1, 3, servers)
-
-    fillInConsumerGroupInfo(topicToDelete, groupToDelete, "consumer", 0, 10, false)
-    fillInConsumerGroupInfo(otherTopic, groupToDelete, "consumer", 0, 10, false)
-    fillInConsumerGroupInfo(topicToDelete, otherGroup, "consumer", 0, 10, false)
-
-    AdminUtils.deleteConsumerGroupInfoForTopicInZK(zkUtils, groupToDelete, topicToDelete)
-
-    TestUtils.waitUntilTrue(() => !groupTopicOffsetAndOwnerDirsExist(new ZKGroupTopicDirs(groupToDelete, topicToDelete)),
-      "DeleteConsumerGroupInfoForTopicInZK should delete the provided consumer group's owner and offset directories for the given topic")
-    TestUtils.waitUntilTrue(() => groupTopicOffsetAndOwnerDirsExist(new ZKGroupTopicDirs(groupToDelete, otherTopic)),
-      "DeleteConsumerGroupInfoForTopicInZK should not delete the provided consumer group's owner and offset directories for unrelated topics")
-    TestUtils.waitUntilTrue(() => groupTopicOffsetAndOwnerDirsExist(new ZKGroupTopicDirs(otherGroup, topicToDelete)),
-      "DeleteConsumerGroupInfoForTopicInZK should not delete unrelated consumer group owner and offset directories")
-  }
-
-  @Test
-  def testGroupTopicWideDeleteInZKDoesNothingForActiveGroupConsumingMultipleTopics() {
-    val topicToDelete = "topicToDelete"
-    val otherTopic = "otherTopic"
-    val group = "group"
-    TestUtils.createTopic(zkUtils, topicToDelete, 1, 3, servers)
-    TestUtils.createTopic(zkUtils, otherTopic, 1, 3, servers)
-
-    fillInConsumerGroupInfo(topicToDelete, group, "consumer", 0, 10, true)
-    fillInConsumerGroupInfo(otherTopic, group, "consumer", 0, 10, true)
-
-    AdminUtils.deleteConsumerGroupInfoForTopicInZK(zkUtils, group, topicToDelete)
-
-    TestUtils.waitUntilTrue(() => groupTopicOffsetAndOwnerDirsExist(new ZKGroupTopicDirs(group, topicToDelete)),
-      "DeleteConsumerGroupInfoForTopicInZK should not delete the provided consumer group's owner and offset directories for the given topic if the consumer group is still active")
-    TestUtils.waitUntilTrue(() => groupTopicOffsetAndOwnerDirsExist(new ZKGroupTopicDirs(group, otherTopic)),
-      "DeleteConsumerGroupInfoForTopicInZK should not delete the provided consumer group's owner and offset directories for unrelated topics")
-  }
-
-  @Test
-  def testTopicWideDeleteInZK() {
-    val topicToDelete = "topicToDelete"
-    val otherTopic = "otherTopic"
-    val groups = Seq("group1", "group2")
-
-    TestUtils.createTopic(zkUtils, topicToDelete, 1, 3, servers)
-    TestUtils.createTopic(zkUtils, otherTopic, 1, 3, servers)
-    val groupTopicDirsForTopicToDelete = groups.map(group => new ZKGroupTopicDirs(group, topicToDelete))
-    val groupTopicDirsForOtherTopic = groups.map(group => new ZKGroupTopicDirs(group, otherTopic))
-    groupTopicDirsForTopicToDelete.foreach(dir => fillInConsumerGroupInfo(topicToDelete, dir.group, "consumer", 0, 10, false))
-    groupTopicDirsForOtherTopic.foreach(dir => fillInConsumerGroupInfo(otherTopic, dir.group, "consumer", 0, 10, false))
-
-    AdminUtils.deleteAllConsumerGroupInfoForTopicInZK(zkUtils, topicToDelete)
-
-    TestUtils.waitUntilTrue(() => !groupTopicDirsForTopicToDelete.exists(groupTopicOffsetAndOwnerDirsExist),
-      "Consumer group info on deleted topic should be deleted by DeleteAllConsumerGroupInfoForTopicInZK")
-    TestUtils.waitUntilTrue(() => groupTopicDirsForOtherTopic.forall(groupTopicOffsetAndOwnerDirsExist),
-      "Consumer group info on unrelated topics should not be deleted by DeleteAllConsumerGroupInfoForTopicInZK")
-  }
-
-  @Test
-  def testConsumptionOnRecreatedTopicAfterTopicWideDeleteInZK() {
-    val topic = "topic"
-    val group = "group"
-
-    TestUtils.createTopic(zkUtils, topic, 1, 3, servers)
-    val dir = new ZKGroupTopicDirs(group, topic)
-    fillInConsumerGroupInfo(topic, dir.group, "consumer", 0, 10, false)
-
-    AdminUtils.deleteTopic(zkUtils, topic)
-    TestUtils.verifyTopicDeletion(zkUtils, topic, 1, servers)
-    AdminUtils.deleteAllConsumerGroupInfoForTopicInZK(zkUtils, topic)
-
-    TestUtils.waitUntilTrue(() => !groupDirExists(dir),
-      "Consumer group info on related topics should be deleted by DeleteAllConsumerGroupInfoForTopicInZK")
-    //produce events
-    val producer = TestUtils.createNewProducer(brokerList)
-    try {
-      produceEvents(producer, topic, List.fill(10)("test"))
-    } finally {
-      producer.close()
-    }
-
-    //consume events
-    val consumerProps = TestUtils.createConsumerProperties(zkConnect, group, "consumer")
-    consumerProps.put("auto.commit.enable", "false")
-    consumerProps.put("auto.offset.reset", "smallest")
-    consumerProps.put("consumer.timeout.ms", "2000")
-    consumerProps.put("fetch.wait.max.ms", "0")
-    val consumerConfig = new ConsumerConfig(consumerProps)
-    val consumerConnector: ConsumerConnector = Consumer.create(consumerConfig)
-    try {
-      val messageStream = consumerConnector.createMessageStreams(Map(topic -> 1))(topic).head
-      consumeEvents(messageStream, 5)
-      consumerConnector.commitOffsets(false)
-    } finally {
-      consumerConnector.shutdown()
-    }
-
-    TestUtils.waitUntilTrue(() => groupTopicOffsetAndOwnerDirsExist(dir),
-      "Consumer group info should exist after consuming from a recreated topic")
-  }
-
-  private def fillInConsumerGroupInfo(topic: String, group: String, consumerId: String, partition: Int, offset: Int, registerConsumer: Boolean) {
-    val consumerProps = TestUtils.createConsumerProperties(zkConnect, group, consumerId)
-    val consumerConfig = new ConsumerConfig(consumerProps)
-    val dir = new ZKGroupTopicDirs(group, topic)
-    TestUtils.updateConsumerOffset(consumerConfig, dir.consumerOffsetDir + "/" + partition, offset)
-    zkUtils.createEphemeralPathExpectConflict(zkUtils.getConsumerPartitionOwnerPath(group, topic, partition), "")
-    zkUtils.makeSurePersistentPathExists(dir.consumerRegistryDir)
-    if (registerConsumer) {
-      zkUtils.createEphemeralPathExpectConflict(dir.consumerRegistryDir + "/" + consumerId, "")
-    }
-  }
-
-  private def groupDirExists(dir: ZKGroupDirs) = {
-    zkUtils.pathExists(dir.consumerGroupDir)
-  }
-
-  private def groupTopicOffsetAndOwnerDirsExist(dir: ZKGroupTopicDirs) = {
-    zkUtils.pathExists(dir.consumerOffsetDir) && zkUtils.pathExists(dir.consumerOwnerDir)
-  }
-
-  private def produceEvents(producer: KafkaProducer[Array[Byte], Array[Byte]], topic: String, messages: List[String]) {
-    messages.foreach(message => producer.send(new ProducerRecord(topic, message.getBytes(StandardCharsets.UTF_8))))
-  }
-
-  private def consumeEvents(messageStream: KafkaStream[Array[Byte], Array[Byte]], n: Int) {
-    val iter = messageStream.iterator
-    (0 until n).foreach(_ => iter.next)
-  }
-}
diff --git a/core/src/test/scala/unit/kafka/admin/DeleteConsumerGroupsTest.scala b/core/src/test/scala/unit/kafka/admin/DeleteConsumerGroupsTest.scala
new file mode 100644
index 0000000000000..a3fa8f511e137
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/admin/DeleteConsumerGroupsTest.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.admin
+
+import joptsimple.OptionException
+import kafka.utils.TestUtils
+import org.apache.kafka.common.errors.{GroupIdNotFoundException, GroupNotEmptyException}
+import org.apache.kafka.common.protocol.Errors
+import org.junit.Assert._
+import org.junit.Test
+
+class DeleteConsumerGroupsTest extends ConsumerGroupCommandTest {
+
+  @Test(expected = classOf[OptionException])
+  def testDeleteWithTopicOption(): Unit = {
+    TestUtils.createOffsetsTopic(zkClient, servers)
+    val cgcArgs = Array("--bootstrap-server", brokerList, "--delete", "--group", group, "--topic")
+    getConsumerGroupService(cgcArgs)
+  }
+
+  @Test
+  def testDeleteCmdNonExistingGroup(): Unit = {
+    TestUtils.createOffsetsTopic(zkClient, servers)
+    val missingGroup = "missing.group"
+
+    val cgcArgs = Array("--bootstrap-server", brokerList, "--delete", "--group", missingGroup)
+    val service = getConsumerGroupService(cgcArgs)
+
+    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))
+  }
+
+  @Test
+  def testDeleteNonExistingGroup(): Unit = {
+    TestUtils.createOffsetsTopic(zkClient, servers)
+    val missingGroup = "missing.group"
+
+    // note the group to be deleted is a different (non-existing) group
+    val cgcArgs = Array("--bootstrap-server", brokerList, "--delete", "--group", missingGroup)
+    val service = getConsumerGroupService(cgcArgs)
+
+    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(missingGroup).getCause
+        .isInstanceOf[GroupIdNotFoundException])
+  }
+
+  @Test
+  def testDeleteCmdNonEmptyGroup(): Unit = {
+    TestUtils.createOffsetsTopic(zkClient, servers)
+
+    // run one consumer in the group
+    addConsumerGroupExecutor(numConsumers = 1)
+    val cgcArgs = Array("--bootstrap-server", brokerList, "--delete", "--group", group)
+    val service = getConsumerGroupService(cgcArgs)
+
+    TestUtils.waitUntilTrue(() => {
+      service.collectGroupMembers(group, false)._2.get.size == 1
+    }, "The group did not initialize as expected.")
+
+    val output = TestUtils.grabConsoleOutput(service.deleteGroups())
+    assertTrue(s"The expected error (${Errors.NON_EMPTY_GROUP}) was not detected while deleting consumer group. Output was: (${output})",
+      output.contains(s"Group '$group' could not be deleted due to:") && output.contains(Errors.NON_EMPTY_GROUP.message))
+  }
+
+  @Test
+  def testDeleteNonEmptyGroup(): Unit = {
+    TestUtils.createOffsetsTopic(zkClient, servers)
+
+    // run one consumer in the group
+    addConsumerGroupExecutor(numConsumers = 1)
+    val cgcArgs = Array("--bootstrap-server", brokerList, "--delete", "--group", group)
+    val service = getConsumerGroupService(cgcArgs)
+
+    TestUtils.waitUntilTrue(() => {
+      service.collectGroupMembers(group, false)._2.get.size == 1
+    }, "The group did not initialize as expected.")
+
+    val result = service.deleteGroups()
+    assertNotNull(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(group).getCause.isInstanceOf[GroupNotEmptyException])
+  }
+
+  @Test
+  def testDeleteCmdEmptyGroup(): Unit = {
+    TestUtils.createOffsetsTopic(zkClient, servers)
+
+    // run one consumer in the group
+    val executor = addConsumerGroupExecutor(numConsumers = 1)
+    val cgcArgs = Array("--bootstrap-server", brokerList, "--delete", "--group", group)
+    val service = getConsumerGroupService(cgcArgs)
+
+    TestUtils.waitUntilTrue(() => {
+      service.listConsumerGroups().contains(group) && service.collectGroupState(group).state == "Stable"
+    }, "The group did not initialize as expected.")
+
+    executor.shutdown()
+
+    TestUtils.waitUntilTrue(() => {
+      service.collectGroupState(group).state == "Empty"
+    }, "The group did not become empty as expected.")
+
+    val output = TestUtils.grabConsoleOutput(service.deleteGroups())
+    assertTrue(s"The consumer group could not be deleted as expected",
+      output.contains(s"Deletion of requested consumer groups ('$group') was successful."))
+  }
+
+  @Test
+  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.listConsumerGroups().toSet == groups.keySet &&
+        groups.keySet.forall(groupId => service.collectGroupState(groupId).state == "Stable")
+    }, "The group did not initialize as expected.")
+
+    // Shutdown consumers to empty out groups
+    groups.values.foreach(executor => executor.shutdown())
+
+    TestUtils.waitUntilTrue(() => {
+      groups.keySet.forall(groupId => service.collectGroupState(groupId).state == "Empty")
+    }, "The group did not become empty as expected.")
+
+    val output = TestUtils.grabConsoleOutput(service.deleteGroups()).trim
+    val expectedGroupsForDeletion = groups.keySet
+    val deletedGroupsGrepped = output.substring(output.indexOf('(') + 1, output.indexOf(')')).split(',')
+      .map(_.replaceAll("'", "").trim).toSet
+
+    assertTrue(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
+    val executor = addConsumerGroupExecutor(numConsumers = 1)
+    val cgcArgs = Array("--bootstrap-server", brokerList, "--delete", "--group", group)
+    val service = getConsumerGroupService(cgcArgs)
+
+    TestUtils.waitUntilTrue(() => {
+      service.listConsumerGroups().contains(group) && service.collectGroupState(group).state == "Stable"
+    }, "The group did not initialize as expected.")
+
+    executor.shutdown()
+
+    TestUtils.waitUntilTrue(() => {
+      service.collectGroupState(group).state == "Empty"
+    }, "The group did not become empty as expected.")
+
+    val result = service.deleteGroups()
+    assertTrue(s"The consumer group could not be deleted as expected",
+      result.size == 1 && result.keySet.contains(group) && result(group) == null)
+  }
+
+  @Test
+  def testDeleteCmdWithMixOfSuccessAndError(): Unit = {
+    TestUtils.createOffsetsTopic(zkClient, servers)
+    val missingGroup = "missing.group"
+
+    // run one consumer in the group
+    val executor = addConsumerGroupExecutor(numConsumers = 1)
+    val cgcArgs = Array("--bootstrap-server", brokerList, "--delete", "--group", group)
+    val service = getConsumerGroupService(cgcArgs)
+
+    TestUtils.waitUntilTrue(() => {
+      service.listConsumerGroups().contains(group) && service.collectGroupState(group).state == "Stable"
+    }, "The group did not initialize as expected.")
+
+    executor.shutdown()
+
+    TestUtils.waitUntilTrue(() => {
+      service.collectGroupState(group).state == "Empty"
+    }, "The group did not become empty as expected.")
+
+    val service2 = getConsumerGroupService(cgcArgs ++ Array("--group", missingGroup))
+    val output = TestUtils.grabConsoleOutput(service2.deleteGroups())
+    assertTrue(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'"))
+  }
+
+  @Test
+  def testDeleteWithMixOfSuccessAndError(): Unit = {
+    TestUtils.createOffsetsTopic(zkClient, servers)
+    val missingGroup = "missing.group"
+
+    // run one consumer in the group
+    val executor = addConsumerGroupExecutor(numConsumers = 1)
+    val cgcArgs = Array("--bootstrap-server", brokerList, "--delete", "--group", group)
+    val service = getConsumerGroupService(cgcArgs)
+
+    TestUtils.waitUntilTrue(() => {
+      service.listConsumerGroups().contains(group) && service.collectGroupState(group).state == "Stable"
+    }, "The group did not initialize as expected.")
+
+    executor.shutdown()
+
+    TestUtils.waitUntilTrue(() => {
+      service.collectGroupState(group).state == "Empty"
+    }, "The group did not become empty as expected.")
+
+    val service2 = getConsumerGroupService(cgcArgs ++ Array("--group", missingGroup))
+    val result = service2.deleteGroups()
+    assertTrue(s"The consumer group deletion did not work as expected",
+      result.size == 2 &&
+        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(): 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 1922aaf0f9197..63597537bbd40 100644
--- a/core/src/test/scala/unit/kafka/admin/DeleteTopicTest.scala
+++ b/core/src/test/scala/unit/kafka/admin/DeleteTopicTest.scala
@@ -16,47 +16,56 @@
  */
 package kafka.admin
 
+import java.util
+import java.util.concurrent.ExecutionException
+import java.util.{Collections, Optional, Properties}
+
+import scala.collection.Seq
 import kafka.log.Log
-import kafka.zk.ZooKeeperTestHarness
+import kafka.zk.{TopicPartitionZNode, ZooKeeperTestHarness}
 import kafka.utils.TestUtils
-import kafka.utils.ZkUtils._
 import kafka.server.{KafkaConfig, KafkaServer}
 import org.junit.Assert._
 import org.junit.{After, Test}
-import java.util.Properties
-
-import kafka.common.{TopicAndPartition, TopicAlreadyMarkedForDeletionException}
+import kafka.common.TopicAlreadyMarkedForDeletionException
+import kafka.controller.{OfflineReplica, PartitionAndReplica, ReplicaAssignment, ReplicaDeletionSuccessful}
+import org.apache.kafka.clients.admin.{Admin, AdminClientConfig, NewPartitionReassignment, NewPartitions}
 import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.errors.UnknownTopicOrPartitionException
+import org.scalatest.Assertions.fail
+import scala.jdk.CollectionConverters._
 
 class DeleteTopicTest extends ZooKeeperTestHarness {
 
   var servers: Seq[KafkaServer] = Seq()
 
   val expectedReplicaAssignment = Map(0 -> List(0, 1, 2))
+  val expectedReplicaFullAssignment = expectedReplicaAssignment.map { case (k, v) =>
+    k -> ReplicaAssignment(v, List(), List())
+  }
 
   @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
     adminZkClient.deleteTopic(topic)
-    TestUtils.verifyTopicDeletion(zkUtils, topic, 1, servers)
+    TestUtils.verifyTopicDeletion(zkClient, topic, 1, servers)
   }
 
   @Test
-  def testResumeDeleteTopicWithRecoveredFollower() {
+  def testResumeDeleteTopicWithRecoveredFollower(): Unit = {
     val topicPartition = new TopicPartition("test", 0)
     val topic = topicPartition.topic
     servers = createTestTopicAndCluster(topic)
     // shut down one follower replica
-    val leaderIdOpt = zkUtils.getLeaderForPartition(topic, 0)
+    val leaderIdOpt = zkClient.getLeaderForPartition(new TopicPartition(topic, 0))
     assertTrue("Leader should exist for partition [test,0]", leaderIdOpt.isDefined)
     val follower = servers.filter(s => s.config.brokerId != leaderIdOpt.get).last
     follower.shutdown()
@@ -65,23 +74,23 @@ class DeleteTopicTest extends ZooKeeperTestHarness {
     // check if all replicas but the one that is shut down has deleted the log
     TestUtils.waitUntilTrue(() =>
       servers.filter(s => s.config.brokerId != follower.config.brokerId)
-        .forall(_.getLogManager().getLog(topicPartition).isEmpty), "Replicas 0,1 have not deleted log.")
+        .forall(_.getLogManager.getLog(topicPartition).isEmpty), "Replicas 0,1 have not deleted log.")
     // ensure topic deletion is halted
-    TestUtils.waitUntilTrue(() => zkUtils.isTopicMarkedForDeletion(topic),
-      "Admin path /admin/delete_topic/test path deleted even when a follower replica is down")
+    TestUtils.waitUntilTrue(() => zkClient.isTopicMarkedForDeletion(topic),
+      "Admin path /admin/delete_topics/test path deleted even when a follower replica is down")
     // restart follower replica
     follower.startup()
-    TestUtils.verifyTopicDeletion(zkUtils, topic, 1, servers)
+    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)
-    val controllerId = zkUtils.getController()
+    val controllerId = zkClient.getControllerId.getOrElse(fail("Controller doesn't exist"))
     val controller = servers.filter(s => s.config.brokerId == controllerId).head
-    val leaderIdOpt = zkUtils.getLeaderForPartition(topic, 0)
+    val leaderIdOpt = zkClient.getLeaderForPartition(new TopicPartition(topic, 0))
     val follower = servers.filter(s => s.config.brokerId != leaderIdOpt.get && s.config.brokerId != controllerId).last
     follower.shutdown()
 
@@ -91,18 +100,17 @@ class DeleteTopicTest extends ZooKeeperTestHarness {
     controller.shutdown()
 
     // ensure topic deletion is halted
-    TestUtils.waitUntilTrue(() => zkUtils.isTopicMarkedForDeletion(topic),
-      "Admin path /admin/delete_topic/test path deleted even when a replica is down")
+    TestUtils.waitUntilTrue(() => zkClient.isTopicMarkedForDeletion(topic),
+      "Admin path /admin/delete_topics/test path deleted even when a replica is down")
 
     controller.startup()
     follower.startup()
 
-    TestUtils.verifyTopicDeletion(zkUtils, topic, 1, servers)
+    TestUtils.verifyTopicDeletion(zkClient, topic, 1, servers)
   }
 
   @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)
@@ -112,45 +120,146 @@ class DeleteTopicTest extends ZooKeeperTestHarness {
     this.servers = allServers
     val servers = allServers.filter(s => expectedReplicaAssignment(0).contains(s.config.brokerId))
     // create the topic
-    adminZkClient.createOrUpdateTopicPartitionAssignmentPathInZK(topic, expectedReplicaAssignment)
+    TestUtils.createTopic(zkClient, topic, expectedReplicaAssignment, servers)
     // wait until replica log is created on every broker
-    TestUtils.waitUntilTrue(() => servers.forall(_.getLogManager().getLog(topicPartition).isDefined),
+    TestUtils.waitUntilTrue(() => servers.forall(_.getLogManager.getLog(topicPartition).isDefined),
       "Replicas for topic test not created.")
-    val leaderIdOpt = zkUtils.getLeaderForPartition(topic, 0)
+    val leaderIdOpt = zkClient.getLeaderForPartition(new TopicPartition(topic, 0))
     assertTrue("Leader should exist for partition [test,0]", leaderIdOpt.isDefined)
     val follower = servers.filter(s => s.config.brokerId != leaderIdOpt.get).last
     follower.shutdown()
     // start topic deletion
     adminZkClient.deleteTopic(topic)
-    // start partition reassignment at the same time right after delete topic. In this case, reassignment will fail since
-    // the topic is being deleted
-    // reassign partition 0
-    val oldAssignedReplicas = zkUtils.getReplicasForPartition(topic, 0)
-    val newReplicas = Seq(1, 2, 3)
-    val reassignPartitionsCommand = new ReassignPartitionsCommand(zkUtils, None,
-      Map(new TopicAndPartition(topicPartition) -> newReplicas))
-    assertTrue("Partition reassignment should fail for [test,0]", reassignPartitionsCommand.reassignPartitions())
-    // wait until reassignment is completed
+    // verify that a partition from the topic cannot be reassigned
+    val props = new Properties()
+    props.setProperty(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, TestUtils.getBrokerListStrFromServers(servers))
+    val adminClient = Admin.create(props)
+    try {
+      waitUntilTopicGone(adminClient, "test")
+      verifyReassignmentFailsForMissing(adminClient, new TopicPartition(topic, 0),
+        new NewPartitionReassignment(util.Arrays.asList(1, 2, 3)))
+    } finally {
+      adminClient.close()
+    }
+    follower.startup()
+    TestUtils.verifyTopicDeletion(zkClient, topic, 1, servers)
+  }
+
+  private def waitUntilTopicGone(adminClient: Admin, topicName: String): Unit = {
     TestUtils.waitUntilTrue(() => {
-      val partitionsBeingReassigned = zkUtils.getPartitionsBeingReassigned().mapValues(_.newReplicas)
-      ReassignPartitionsCommand.checkIfPartitionReassignmentSucceeded(zkUtils, new TopicAndPartition(topicPartition),
-        Map(new TopicAndPartition(topicPartition) -> newReplicas), partitionsBeingReassigned) == ReassignmentFailed
-    }, "Partition reassignment shouldn't complete.")
-    val controllerId = zkUtils.getController()
-    val controller = servers.filter(s => s.config.brokerId == controllerId).head
-    assertFalse("Partition reassignment should fail",
-      controller.kafkaController.controllerContext.partitionsBeingReassigned.contains(topicPartition))
-    val assignedReplicas = zkUtils.getReplicasForPartition(topic, 0)
-    assertEquals("Partition should not be reassigned to 0, 1, 2", oldAssignedReplicas, assignedReplicas)
+      try {
+        adminClient.describeTopics(util.Collections.singletonList(topicName)).all().get()
+        false
+      } catch {
+        case e: ExecutionException =>
+          classOf[UnknownTopicOrPartitionException].equals(e.getCause.getClass)
+      }
+    }, s"Topic ${topicName} should be deleted.")
+  }
+
+  private def verifyReassignmentFailsForMissing(adminClient: Admin,
+                                                partition: TopicPartition,
+                                                reassignment: NewPartitionReassignment): Unit = {
+    try {
+      adminClient.alterPartitionReassignments(Collections.singletonMap(partition,
+        Optional.of(new NewPartitionReassignment(util.Arrays.asList(1, 2, 3))))).all().get()
+      fail("expected partition reassignment to fail for [test,0]")
+    } catch {
+      case e: ExecutionException =>
+        assertEquals(classOf[UnknownTopicOrPartitionException], e.getCause.getClass)
+    }
+  }
+
+  private def getController() : (KafkaServer, Int) = {
+    val controllerId = zkClient.getControllerId.getOrElse(fail("Controller doesn't exist"))
+    val controller = servers.find(s => s.config.brokerId == controllerId).get
+    (controller, controllerId)
+  }
+
+  private def ensureControllerExists() = {
+    TestUtils.waitUntilTrue(() => {
+      try {
+        getController()
+        true
+      } catch {
+        case _: Throwable  => false
+      }
+    }, "Controller should eventually exist")
+  }
+
+  private def getAllReplicasFromAssignment(topic : String, assignment : Map[Int, Seq[Int]]) : Set[PartitionAndReplica] = {
+    assignment.flatMap { case (partition, replicas) =>
+      replicas.map {r => new PartitionAndReplica(new TopicPartition(topic, partition), r)}
+    }.toSet
+  }
+
+  @Test
+  def testIncreasePartitionCountDuringDeleteTopic(): Unit = {
+    val topic = "test"
+    val topicPartition = new TopicPartition(topic, 0)
+    val brokerConfigs = TestUtils.createBrokerConfigs(4, zkConnect, false)
+    brokerConfigs.foreach(p => p.setProperty("delete.topic.enable", "true"))
+    // create brokers
+    val allServers = brokerConfigs.map(b => TestUtils.createServer(KafkaConfig.fromProps(b)))
+    this.servers = allServers
+    val servers = allServers.filter(s => expectedReplicaAssignment(0).contains(s.config.brokerId))
+    // create the topic
+    TestUtils.createTopic(zkClient, topic, expectedReplicaAssignment, servers)
+    // wait until replica log is created on every broker
+    TestUtils.waitUntilTrue(() => servers.forall(_.getLogManager.getLog(topicPartition).isDefined),
+      "Replicas for topic test not created.")
+    // shutdown a broker to make sure the following topic deletion will be suspended
+    val leaderIdOpt = zkClient.getLeaderForPartition(topicPartition)
+    assertTrue("Leader should exist for partition [test,0]", leaderIdOpt.isDefined)
+    val follower = servers.filter(s => s.config.brokerId != leaderIdOpt.get).last
+    follower.shutdown()
+    // start topic deletion
+    adminZkClient.deleteTopic(topic)
+
+    // make sure deletion of all of the topic's replicas have been tried
+    ensureControllerExists()
+    val (controller, controllerId) = getController()
+    val allReplicasForTopic = getAllReplicasFromAssignment(topic, expectedReplicaAssignment)
+    TestUtils.waitUntilTrue(() => {
+      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")
+
+    // increase the partition count for topic
+    val props = new Properties()
+    props.setProperty(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, TestUtils.getBrokerListStrFromServers(servers))
+    val adminClient = Admin.create(props)
+    try {
+      adminClient.createPartitions(Map(topic -> NewPartitions.increaseTo(2)).asJava).all().get()
+    } catch {
+      case _: ExecutionException =>
+    }
+    // trigger a controller switch now
+    val previousControllerId = controllerId
+
+    controller.shutdown()
+
+    ensureControllerExists()
+    // wait until a new controller to show up
+    TestUtils.waitUntilTrue(() => {
+      val (newController, newControllerId) = getController()
+      newControllerId != previousControllerId
+    }, "The new controller should not have the failed controller id")
+
+    // bring back the failed brokers
     follower.startup()
-    TestUtils.verifyTopicDeletion(zkUtils, topic, 1, servers)
+    controller.startup()
+    TestUtils.verifyTopicDeletion(zkClient, topic, 2, servers)
+    adminClient.close()
   }
 
+
   @Test
-  def testDeleteTopicDuringAddPartition() {
+  def testDeleteTopicDuringAddPartition(): Unit = {
     val topic = "test"
     servers = createTestTopicAndCluster(topic)
-    val leaderIdOpt = zkUtils.getLeaderForPartition(topic, 0)
+    val leaderIdOpt = zkClient.getLeaderForPartition(new TopicPartition(topic, 0))
     assertTrue("Leader should exist for partition [test,0]", leaderIdOpt.isDefined)
     val follower = servers.filter(_.config.brokerId != leaderIdOpt.get).last
     val newPartition = new TopicPartition(topic, 1)
@@ -158,24 +267,25 @@ class DeleteTopicTest extends ZooKeeperTestHarness {
     val brokers = adminZkClient.getBrokerMetadatas()
     follower.shutdown()
     // wait until the broker has been removed from ZK to reduce non-determinism
-    TestUtils.waitUntilTrue(() => zkUtils.getBrokerInfo(follower.config.brokerId).isEmpty,
+    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)
     follower.startup()
     // test if topic deletion is resumed
-    TestUtils.verifyTopicDeletion(zkUtils, topic, 1, servers)
+    TestUtils.verifyTopicDeletion(zkClient, topic, 1, servers)
     // verify that new partition doesn't exist on any broker either
     TestUtils.waitUntilTrue(() =>
-      servers.forall(_.getLogManager().getLog(newPartition).isEmpty),
+      servers.forall(_.getLogManager.getLog(newPartition).isEmpty),
       "Replica logs not for new partition [test,1] not deleted after delete topic is complete.")
   }
 
   @Test
-  def testAddPartitionDuringDeleteTopic() {
+  def testAddPartitionDuringDeleteTopic(): Unit = {
+    zkClient.createTopLevelPaths()
     val topic = "test"
     servers = createTestTopicAndCluster(topic)
     val brokers = adminZkClient.getBrokerMetadatas()
@@ -183,34 +293,32 @@ 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(zkUtils, topic, 1, servers)
+    TestUtils.verifyTopicDeletion(zkClient, topic, 1, servers)
     // verify that new partition doesn't exist on any broker either
     assertTrue("Replica logs not deleted after delete topic is complete",
-      servers.forall(_.getLogManager().getLog(newPartition).isEmpty))
+      servers.forall(_.getLogManager.getLog(newPartition).isEmpty))
   }
 
   @Test
-  def testRecreateTopicAfterDeletion() {
+  def testRecreateTopicAfterDeletion(): Unit = {
     val expectedReplicaAssignment = Map(0 -> List(0, 1, 2))
     val topic = "test"
     val topicPartition = new TopicPartition(topic, 0)
     servers = createTestTopicAndCluster(topic)
     // start topic deletion
     adminZkClient.deleteTopic(topic)
-    TestUtils.verifyTopicDeletion(zkUtils, topic, 1, servers)
+    TestUtils.verifyTopicDeletion(zkClient, topic, 1, servers)
     // re-create topic on same replicas
-    adminZkClient.createOrUpdateTopicPartitionAssignmentPathInZK(topic, expectedReplicaAssignment)
-    // wait until leader is elected
-    TestUtils.waitUntilLeaderIsElectedOrChanged(zkUtils, topic, 0, 1000)
+    TestUtils.createTopic(zkClient, topic, expectedReplicaAssignment, servers)
     // check if all replica logs are created
-    TestUtils.waitUntilTrue(() => servers.forall(_.getLogManager().getLog(topicPartition).isDefined),
+    TestUtils.waitUntilTrue(() => servers.forall(_.getLogManager.getLog(topicPartition).isDefined),
       "Replicas for topic test not created.")
   }
 
   @Test
-  def testDeleteNonExistingTopic() {
+  def testDeleteNonExistingTopic(): Unit = {
     val topicPartition = new TopicPartition("test", 0)
     val topic = topicPartition.topic
     servers = createTestTopicAndCluster(topic)
@@ -222,18 +330,18 @@ class DeleteTopicTest extends ZooKeeperTestHarness {
       case _: UnknownTopicOrPartitionException => // expected exception
     }
     // verify delete topic path for test2 is removed from ZooKeeper
-    TestUtils.verifyTopicDeletion(zkUtils, "test2", 1, servers)
+    TestUtils.verifyTopicDeletion(zkClient, "test2", 1, servers)
     // verify that topic test is untouched
-    TestUtils.waitUntilTrue(() => servers.forall(_.getLogManager().getLog(topicPartition).isDefined),
+    TestUtils.waitUntilTrue(() => servers.forall(_.getLogManager.getLog(topicPartition).isDefined),
       "Replicas for topic test not created")
     // test the topic path exists
-    assertTrue("Topic test mistakenly deleted", zkUtils.pathExists(getTopicPath(topic)))
+    assertTrue("Topic test mistakenly deleted", zkClient.topicExists(topic))
     // topic test should have a leader
-    TestUtils.waitUntilLeaderIsElectedOrChanged(zkUtils, topic, 0, 1000)
+    TestUtils.waitUntilLeaderIsElectedOrChanged(zkClient, topic, 0, 1000)
   }
 
   @Test
-  def testDeleteTopicWithCleaner() {
+  def testDeleteTopicWithCleaner(): Unit = {
     val topicName = "test"
     val topicPartition = new TopicPartition(topicName, 0)
     val topic = topicPartition.topic
@@ -245,7 +353,7 @@ class DeleteTopicTest extends ZooKeeperTestHarness {
     brokerConfigs.head.setProperty("log.segment.bytes","100")
     brokerConfigs.head.setProperty("log.cleaner.dedupe.buffer.size","1048577")
 
-    servers = createTestTopicAndCluster(topic,brokerConfigs)
+    servers = createTestTopicAndCluster(topic, brokerConfigs, expectedReplicaAssignment)
 
     // for simplicity, we are validating cleaner offsets on a single broker
     val server = servers.head
@@ -259,11 +367,11 @@ class DeleteTopicTest extends ZooKeeperTestHarness {
 
     // delete topic
     adminZkClient.deleteTopic("test")
-    TestUtils.verifyTopicDeletion(zkUtils, "test", 1, servers)
+    TestUtils.verifyTopicDeletion(zkClient, "test", 1, servers)
   }
 
   @Test
-  def testDeleteTopicAlreadyMarkedAsDeleted() {
+  def testDeleteTopicAlreadyMarkedAsDeleted(): Unit = {
     val topicPartition = new TopicPartition("test", 0)
     val topic = topicPartition.topic
     servers = createTestTopicAndCluster(topic)
@@ -279,23 +387,23 @@ class DeleteTopicTest extends ZooKeeperTestHarness {
       case _: TopicAlreadyMarkedForDeletionException => // expected exception
     }
 
-    TestUtils.verifyTopicDeletion(zkUtils, topic, 1, servers)
+    TestUtils.verifyTopicDeletion(zkClient, topic, 1, servers)
   }
 
-  private def createTestTopicAndCluster(topic: String, deleteTopicEnabled: Boolean = true): Seq[KafkaServer] = {
+  private def createTestTopicAndCluster(topic: String, deleteTopicEnabled: Boolean = true, replicaAssignment: Map[Int, List[Int]] = expectedReplicaAssignment): Seq[KafkaServer] = {
     val brokerConfigs = TestUtils.createBrokerConfigs(3, zkConnect, enableControlledShutdown = false)
     brokerConfigs.foreach(_.setProperty("delete.topic.enable", deleteTopicEnabled.toString))
-    createTestTopicAndCluster(topic, brokerConfigs)
+    createTestTopicAndCluster(topic, brokerConfigs, replicaAssignment)
   }
 
-  private def createTestTopicAndCluster(topic: String, brokerConfigs: Seq[Properties]): Seq[KafkaServer] = {
+  private def createTestTopicAndCluster(topic: String, brokerConfigs: Seq[Properties], replicaAssignment: Map[Int, List[Int]]): Seq[KafkaServer] = {
     val topicPartition = new TopicPartition(topic, 0)
     // create brokers
     val servers = brokerConfigs.map(b => TestUtils.createServer(KafkaConfig.fromProps(b)))
     // create the topic
-    adminZkClient.createOrUpdateTopicPartitionAssignmentPathInZK(topic, expectedReplicaAssignment)
+    TestUtils.createTopic(zkClient, topic, expectedReplicaAssignment, servers)
     // wait until replica log is created on every broker
-    TestUtils.waitUntilTrue(() => servers.forall(_.getLogManager().getLog(topicPartition).isDefined),
+    TestUtils.waitUntilTrue(() => servers.forall(_.getLogManager.getLog(topicPartition).isDefined),
       "Replicas for topic test not created")
     servers
   }
@@ -311,20 +419,51 @@ 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(() => !zkUtils.isTopicMarkedForDeletion(topic),
-      "Admin path /admin/delete_topic/%s path not deleted even if deleteTopic is disabled".format(topic))
+    TestUtils.waitUntilTrue(() => !zkClient.isTopicMarkedForDeletion(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))
+    assertTrue(servers.forall(_.getLogManager.getLog(topicPartition).isDefined))
     // test the topic path exists
-    assertTrue("Topic path disappeared", zkUtils.pathExists(getTopicPath(topic)))
+    assertTrue("Topic path disappeared", zkClient.topicExists(topic))
     // topic test should have a leader
-    val leaderIdOpt = zkUtils.getLeaderForPartition(topic, 0)
+    val leaderIdOpt = zkClient.getLeaderForPartition(new TopicPartition(topic, 0))
     assertTrue("Leader should exist for topic test", leaderIdOpt.isDefined)
   }
+
+  @Test
+  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.
+      */
+
+    val replicaAssignment = Map(0 -> List(0, 1, 2), 1 -> List(0, 1, 2))
+    val topic = "test"
+    servers = createTestTopicAndCluster(topic, true, replicaAssignment)
+
+    /**
+      * shutdown all brokers in order to create a partially deleted topic on ZK
+      */
+    servers.foreach(_.shutdown())
+
+    /**
+      * delete the partition znode at /brokers/topics/test/partition/0
+      * to simulate the case that a previous controller crashed right after deleting the partition znode
+      */
+    zkClient.deleteRecursive(TopicPartitionZNode.path(new TopicPartition(topic, 0)))
+    adminZkClient.deleteTopic(topic)
+
+    /**
+      * start up all brokers and verify that topic deletion eventually finishes.
+      */
+    servers.foreach(_.startup())
+    TestUtils.waitUntilTrue(() => servers.exists(_.kafkaController.isActive), "No controller is elected")
+    TestUtils.verifyTopicDeletion(zkClient, topic, 2, servers)
+  }
 }
diff --git a/core/src/test/scala/unit/kafka/admin/DescribeConsumerGroupTest.scala b/core/src/test/scala/unit/kafka/admin/DescribeConsumerGroupTest.scala
index e3673725bf7ca..2c28ccf8733ad 100644
--- a/core/src/test/scala/unit/kafka/admin/DescribeConsumerGroupTest.scala
+++ b/core/src/test/scala/unit/kafka/admin/DescribeConsumerGroupTest.scala
@@ -16,194 +16,348 @@
  */
 package kafka.admin
 
-import java.util.concurrent.ExecutorService
-import java.util.concurrent.Executors
-import java.util.concurrent.TimeUnit
-import java.util.Collections
 import java.util.Properties
 
-import org.junit.Assert._
-import org.junit.{After, Before, Test}
-import kafka.admin.ConsumerGroupCommand.{ConsumerGroupCommandOptions, ConsumerGroupService, KafkaConsumerGroupService, ZkConsumerGroupService}
-import kafka.consumer.OldConsumer
-import kafka.consumer.Whitelist
-import kafka.integration.KafkaServerTestHarness
-import kafka.server.KafkaConfig
-import kafka.utils.TestUtils
-import org.apache.kafka.clients.consumer.KafkaConsumer
+import kafka.utils.{Exit, 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.WakeupException
-import org.apache.kafka.common.serialization.StringDeserializer
+import org.junit.Assert._
+import org.junit.Test
 
-import scala.collection.mutable.ArrayBuffer
+import scala.concurrent.ExecutionException
+import scala.util.Random
 
-class DescribeConsumerGroupTest extends KafkaServerTestHarness {
-  private val topic = "foo"
-  private val group = "test.group"
+class DescribeConsumerGroupTest extends ConsumerGroupCommandTest {
 
-  @deprecated("This field will be removed in a future release", "0.11.0.0")
-  private val oldConsumers = new ArrayBuffer[OldConsumer]
-  private var consumerGroupService: ConsumerGroupService = _
-  private var consumerGroupExecutor: ConsumerGroupExecutor = _
+  private val describeTypeOffsets = Array(Array(""), Array("--offsets"))
+  private val describeTypeMembers = Array(Array("--members"), Array("--members", "--verbose"))
+  private val describeTypeState = Array(Array("--state"))
+  private val describeTypes = describeTypeOffsets ++ describeTypeMembers ++ describeTypeState
 
-  // configure the servers and clients
-  override def generateConfigs = {
-    TestUtils.createBrokerConfigs(1, zkConnect, enableControlledShutdown = false).map { props =>
-      KafkaConfig.fromProps(props)
+  @Test
+  def testDescribeNonExistingGroup(): Unit = {
+    TestUtils.createOffsetsTopic(zkClient, servers)
+    val missingGroup = "missing.group"
+
+    for (describeType <- describeTypes) {
+      // note the group to be queried is a different (non-existing) group
+      val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--group", missingGroup) ++ describeType
+      val service = getConsumerGroupService(cgcArgs)
+
+      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."))
     }
   }
 
-  @Before
-  override def setUp() {
-    super.setUp()
-    adminZkClient.createTopic(topic, 1, 1)
+  @Test
+  def testDescribeWithMultipleSubActions(): Unit = {
+    var exitStatus: Option[Int] = None
+    var exitMessage: Option[String] = None
+    Exit.setExitProcedure { (status, err) =>
+      exitStatus = Some(status)
+      exitMessage = err
+      throw new RuntimeException
+    }
+    val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--group", group, "--members", "--state")
+    try {
+      ConsumerGroupCommand.main(cgcArgs)
+    } catch {
+      case e: RuntimeException => //expected
+    } finally {
+      Exit.resetExitProcedure()
+    }
+    assertEquals(Some(1), exitStatus)
+    assertTrue(exitMessage.get.contains("Option [describe] takes at most one of these options"))
   }
 
-  @After
-  override def tearDown(): Unit = {
-    if (consumerGroupService != null)
-      consumerGroupService.close()
-    if (consumerGroupExecutor != null)
-      consumerGroupExecutor.shutdown()
-    oldConsumers.foreach(_.stop())
-    super.tearDown()
+  @Test
+  def testDescribeWithStateValue(): Unit = {
+    var exitStatus: Option[Int] = None
+    var exitMessage: Option[String] = None
+    Exit.setExitProcedure { (status, err) =>
+      exitStatus = Some(status)
+      exitMessage = err
+      throw new RuntimeException
+    }
+    val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--all-groups", "--state", "Stable")
+    try {
+      ConsumerGroupCommand.main(cgcArgs)
+    } catch {
+      case e: RuntimeException => //expected
+    } finally {
+      Exit.resetExitProcedure()
+    }
+    assertEquals(Some(1), exitStatus)
+    assertTrue(exitMessage.get.contains("Option [describe] does not take a value for [state]"))
   }
 
   @Test
-  @deprecated("This test has been deprecated and will be removed in a future release.", "0.11.0.0")
-  def testDescribeNonExistingGroup() {
-    TestUtils.createOffsetsTopic(zkUtils, servers)
-    createOldConsumer()
-    val opts = new ConsumerGroupCommandOptions(Array("--zookeeper", zkConnect, "--describe", "--group", "missing.group"))
-    consumerGroupService = new ZkConsumerGroupService(opts)
-    TestUtils.waitUntilTrue(() => consumerGroupService.describeGroup()._2.isEmpty, "Expected no rows in describe group results.")
+  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", group)
+    val service = getConsumerGroupService(cgcArgs)
+
+    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
-  @deprecated("This test has been deprecated and will be removed in a future release.", "0.11.0.0")
-  def testDescribeExistingGroup() {
-    TestUtils.createOffsetsTopic(zkUtils, servers)
-    createOldConsumer()
-    val opts = new ConsumerGroupCommandOptions(Array("--zookeeper", zkConnect, "--describe", "--group", group))
-    consumerGroupService = new ZkConsumerGroupService(opts)
-    TestUtils.waitUntilTrue(() => {
-      val (_, assignments) = consumerGroupService.describeGroup()
-      assignments.isDefined &&
-      assignments.get.count(_.group == group) == 1 &&
-      assignments.get.filter(_.group == group).head.consumerId.exists(_.trim != ConsumerGroupCommand.MISSING_COLUMN_VALUE)
-    }, "Expected rows and a consumer id column in describe group results.")
+  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", group)
+    val service = getConsumerGroupService(cgcArgs)
+
+    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(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
-  @deprecated("This test has been deprecated and will be removed in a future release.", "0.11.0.0")
-  def testDescribeExistingGroupWithNoMembers() {
-    TestUtils.createOffsetsTopic(zkUtils, servers)
-    createOldConsumer()
-    val opts = new ConsumerGroupCommandOptions(Array("--zookeeper", zkConnect, "--describe", "--group", group))
-    consumerGroupService = new ZkConsumerGroupService(opts)
+  def testDescribeStateOfNonExistingGroup(): Unit = {
+    val group = "missing.group"
+    TestUtils.createOffsetsTopic(zkClient, servers)
 
-    TestUtils.waitUntilTrue(() => {
-      val (_, assignments) = consumerGroupService.describeGroup()
-      assignments.isDefined &&
-      assignments.get.count(_.group == group) == 1 &&
-      assignments.get.filter(_.group == group).head.consumerId.exists(_.trim != ConsumerGroupCommand.MISSING_COLUMN_VALUE)
-    }, "Expected rows and a consumer id column in describe group results.")
-    oldConsumers.head.stop()
+    // 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", group)
+    val service = getConsumerGroupService(cgcArgs)
 
-    TestUtils.waitUntilTrue(() => {
-      val (_, assignments) = consumerGroupService.describeGroup()
-      assignments.isDefined &&
-      assignments.get.count(_.group == group) == 1 &&
-      assignments.get.filter(_.group == group).head.consumerId.exists(_.trim == ConsumerGroupCommand.MISSING_COLUMN_VALUE) // the member should be gone
-    }, "Expected no active member in describe group results.")
+    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)
+    )
   }
 
   @Test
-  @deprecated("This test has been deprecated and will be removed in a future release.", "0.11.0.0")
-  def testDescribeConsumersWithNoAssignedPartitions() {
-    TestUtils.createOffsetsTopic(zkUtils, servers)
-    createOldConsumer()
-    createOldConsumer()
-    val opts = new ConsumerGroupCommandOptions(Array("--zookeeper", zkConnect, "--describe", "--group", group))
-    consumerGroupService = new ZkConsumerGroupService(opts)
-    TestUtils.waitUntilTrue(() => {
-      val (_, assignments) = consumerGroupService.describeGroup()
-      assignments.isDefined &&
-      assignments.get.count(_.group == group) == 2 &&
-      assignments.get.count { x => x.group == group && x.partition.isDefined } == 1 &&
-      assignments.get.count { x => x.group == group && x.partition.isEmpty } == 1
-    }, "Expected rows for consumers with no assigned partitions in describe group results.")
+  def testDescribeExistingGroup(): Unit = {
+    TestUtils.createOffsetsTopic(zkClient, servers)
+
+    for (describeType <- describeTypes) {
+      val group = this.group + describeType.mkString("")
+      // run one consumer in the group consuming from a single-partition topic
+      addConsumerGroupExecutor(numConsumers = 1, group = group)
+      val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--group", group) ++ describeType
+      val service = getConsumerGroupService(cgcArgs)
+
+      TestUtils.waitUntilTrue(() => {
+        val (output, error) = TestUtils.grabConsoleOutputAndError(service.describeGroups())
+        output.trim.split("\n").length == 2 && error.isEmpty
+      }, s"Expected a data row and no error in describe results with describe type ${describeType.mkString(" ")}.")
+    }
   }
 
   @Test
-  def testDescribeNonExistingGroupWithNewConsumer() {
-    TestUtils.createOffsetsTopic(zkUtils, servers)
-    // run one consumer in the group consuming from a single-partition topic
-    consumerGroupExecutor = new ConsumerGroupExecutor(brokerList, 1, group, topic)
+  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(" ")}.")
+    }
+  }
 
-    // note the group to be queried is a different (non-existing) group
-    val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--group", "missing.group")
-    val opts = new ConsumerGroupCommandOptions(cgcArgs)
-    consumerGroupService = new KafkaConsumerGroupService(opts)
+  @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 (state, assignments) = consumerGroupService.describeGroup()
-    assertTrue("Expected the state to be 'Dead' with no members in the group.", state == Some("Dead") && assignments == Some(List()))
+    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 testDescribeExistingGroupWithNewConsumer() {
-    TestUtils.createOffsetsTopic(zkUtils, servers)
+  def testDescribeOffsetsOfExistingGroup(): Unit = {
+    TestUtils.createOffsetsTopic(zkClient, servers)
+
     // run one consumer in the group consuming from a single-partition topic
-    consumerGroupExecutor = new ConsumerGroupExecutor(brokerList, 1, group, topic)
+    addConsumerGroupExecutor(numConsumers = 1)
 
     val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--group", group)
-    val opts = new ConsumerGroupCommandOptions(cgcArgs)
-    consumerGroupService = new KafkaConsumerGroupService(opts)
+    val service = getConsumerGroupService(cgcArgs)
 
     TestUtils.waitUntilTrue(() => {
-        val (state, assignments) = consumerGroupService.describeGroup()
-        state == Some("Stable") &&
+      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)
-    }, "Expected a 'Stable' group status, rows and valid values for consumer id / client id / host columns in describe group results.")
+    }, 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(): Unit = {
+    TestUtils.createOffsetsTopic(zkClient, servers)
+
+    // run one consumer in the group consuming from a single-partition topic
+    addConsumerGroupExecutor(numConsumers = 1)
+    val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--group", group)
+    val service = getConsumerGroupService(cgcArgs)
+
+    TestUtils.waitUntilTrue(() => {
+      val (state, assignments) = service.collectGroupMembers(group, false)
+      state.contains("Stable") &&
+        (assignments match {
+          case Some(memberAssignments) =>
+            memberAssignments.count(_.group == group) == 1 &&
+              memberAssignments.filter(_.group == group).head.consumerId != ConsumerGroupCommand.MISSING_COLUMN_VALUE &&
+              memberAssignments.filter(_.group == group).head.clientId != ConsumerGroupCommand.MISSING_COLUMN_VALUE &&
+              memberAssignments.filter(_.group == group).head.host != ConsumerGroupCommand.MISSING_COLUMN_VALUE
+          case None =>
+            false
+        })
+    }, s"Expected a 'Stable' group status, rows and valid member information for group $group.")
+
+    val (_, assignments) = service.collectGroupMembers(group, true)
+    assignments match {
+      case None =>
+        fail(s"Expected partition assignments for members of group $group")
+      case Some(memberAssignments) =>
+        assertTrue(s"Expected a topic partition assigned to the single group member for group $group",
+          memberAssignments.size == 1 &&
+          memberAssignments.head.assignment.size == 1)
+    }
   }
 
   @Test
-  def testDescribeExistingGroupWithNoMembersWithNewConsumer() {
-    TestUtils.createOffsetsTopic(zkUtils, servers)
+  def testDescribeStateOfExistingGroup(): Unit = {
+    TestUtils.createOffsetsTopic(zkClient, servers)
+
     // run one consumer in the group consuming from a single-partition topic
-    consumerGroupExecutor = new ConsumerGroupExecutor(brokerList, 1, group, topic)
+    addConsumerGroupExecutor(numConsumers = 1)
+    val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--group", group)
+    val service = getConsumerGroupService(cgcArgs)
+
+    TestUtils.waitUntilTrue(() => {
+      val state = service.collectGroupState(group)
+      state.state == "Stable" &&
+        state.numMembers == 1 &&
+        state.assignmentStrategy == "range" &&
+        state.coordinator != null &&
+        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.")
+  }
+
+  @Test
+  def testDescribeStateOfExistingGroupWithRoundRobinAssignor(): Unit = {
+    TestUtils.createOffsetsTopic(zkClient, servers)
 
+    // run one consumer in the group consuming from a single-partition topic
+    addConsumerGroupExecutor(numConsumers = 1, strategy = classOf[RoundRobinAssignor].getName)
     val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--group", group)
-    val opts = new ConsumerGroupCommandOptions(cgcArgs)
-    consumerGroupService = new KafkaConsumerGroupService(opts)
+    val service = getConsumerGroupService(cgcArgs)
 
     TestUtils.waitUntilTrue(() => {
-      val (state, _) = consumerGroupService.describeGroup()
-      state == Some("Stable")
-    }, "Expected the group to initially become stable.")
+      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.")
+  }
+
+  @Test
+  def testDescribeExistingGroupWithNoMembers(): Unit = {
+    TestUtils.createOffsetsTopic(zkClient, servers)
+
+    for (describeType <- describeTypes) {
+      val group = this.group + describeType.mkString("")
+      // run one consumer in the group consuming from a single-partition topic
+      val executor = addConsumerGroupExecutor(numConsumers = 1, group = group)
+      val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--group", group) ++ describeType
+      val service = getConsumerGroupService(cgcArgs)
+
+      TestUtils.waitUntilTrue(() => {
+        val (output, error) = TestUtils.grabConsoleOutputAndError(service.describeGroups())
+        output.trim.split("\n").length == 2 && error.isEmpty
+      }, s"Expected describe group results with one data row for describe type '${describeType.mkString(" ")}'")
+
+      // stop the consumer so the group has no active member anymore
+      executor.shutdown()
+      TestUtils.waitUntilTrue(() => {
+        TestUtils.grabConsoleError(service.describeGroups()).contains(s"Consumer group '$group' has no active members.")
+      }, s"Expected no active member in describe group results with describe type ${describeType.mkString(" ")}")
+    }
+  }
+
+  @Test
+  def testDescribeOffsetsOfExistingGroupWithNoMembers(): Unit = {
+    TestUtils.createOffsetsTopic(zkClient, servers)
+
+    // run one consumer in the group consuming from a single-partition topic
+    val executor = addConsumerGroupExecutor(numConsumers = 1)
+
+    val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--group", group)
+    val service = getConsumerGroupService(cgcArgs)
 
-    // Group assignments in describeGroup rely on finding committed consumer offsets.
-    // Wait for an offset commit before shutting down the group executor.
     TestUtils.waitUntilTrue(() => {
-      val (_, assignments) = consumerGroupService.describeGroup()
-      assignments.exists(_.exists(_.group == group))
-    }, "Expected to find group in assignments after initial offset commit")
+      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.")
 
     // stop the consumer so the group has no active member anymore
-    consumerGroupExecutor.shutdown()
-
-    val (result, succeeded) = TestUtils.computeUntilTrue(consumerGroupService.describeGroup()) { case (state, assignments) =>
-      val testGroupAssignments = assignments.toSeq.flatMap(_.filter(_.group == group))
-      def assignment = testGroupAssignments.head
-      state == Some("Empty") &&
-        testGroupAssignments.size == 1 &&
-        assignment.consumerId.exists(_.trim == ConsumerGroupCommand.MISSING_COLUMN_VALUE) && // the member should be gone
-        assignment.clientId.exists(_.trim == ConsumerGroupCommand.MISSING_COLUMN_VALUE) &&
-        assignment.host.exists(_.trim == ConsumerGroupCommand.MISSING_COLUMN_VALUE)
+    executor.shutdown()
+
+    val (result, succeeded) = TestUtils.computeUntilTrue(service.collectGroupOffsets(group)) {
+      case (state, assignments) =>
+        val testGroupAssignments = assignments.toSeq.flatMap(_.filter(_.group == group))
+        def assignment = testGroupAssignments.head
+        state.contains("Empty") &&
+          testGroupAssignments.size == 1 &&
+          assignment.consumerId.exists(_.trim == ConsumerGroupCommand.MISSING_COLUMN_VALUE) && // the member should be gone
+          assignment.clientId.exists(_.trim == ConsumerGroupCommand.MISSING_COLUMN_VALUE) &&
+          assignment.host.exists(_.trim == ConsumerGroupCommand.MISSING_COLUMN_VALUE)
     }
     val (state, assignments) = result
     assertTrue(s"Expected no active member in describe group results, state: $state, assignments: $assignments",
@@ -211,122 +365,357 @@ class DescribeConsumerGroupTest extends KafkaServerTestHarness {
   }
 
   @Test
-  def testDescribeConsumersWithNoAssignedPartitionsWithNewConsumer() {
-    TestUtils.createOffsetsTopic(zkUtils, servers)
+  def testDescribeMembersOfExistingGroupWithNoMembers(): Unit = {
+    TestUtils.createOffsetsTopic(zkClient, servers)
+
+    // run one consumer in the group consuming from a single-partition topic
+    val executor = addConsumerGroupExecutor(numConsumers = 1)
+
+    val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--group", group)
+    val service = getConsumerGroupService(cgcArgs)
+
+    TestUtils.waitUntilTrue(() => {
+      val (state, assignments) = service.collectGroupMembers(group, false)
+      state.contains("Stable") && assignments.exists(_.exists(_.group == group))
+    }, "Expected the group to initially become stable, and to find group in assignments after initial offset commit.")
+
+    // stop the consumer so the group has no active member anymore
+    executor.shutdown()
+
+    TestUtils.waitUntilTrue(() => {
+      val (state, assignments) = service.collectGroupMembers(group, false)
+      state.contains("Empty") && assignments.isDefined && assignments.get.isEmpty
+    }, s"Expected no member in describe group members results for group '$group'")
+  }
+
+  @Test
+  def testDescribeStateOfExistingGroupWithNoMembers(): Unit = {
+    TestUtils.createOffsetsTopic(zkClient, servers)
+
+    // run one consumer in the group consuming from a single-partition topic
+    val executor = addConsumerGroupExecutor(numConsumers = 1)
+
+    val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--group", group)
+    val service = getConsumerGroupService(cgcArgs)
+
+    TestUtils.waitUntilTrue(() => {
+      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.")
+
+    // stop the consumer so the group has no active member anymore
+    executor.shutdown()
+
+    TestUtils.waitUntilTrue(() => {
+      val state = service.collectGroupState(group)
+      state.state == "Empty" && state.numMembers == 0 && state.assignmentStrategy == ""
+    }, s"Expected the group '$group' to become empty after the only member leaving.")
+  }
+
+  @Test
+  def testDescribeWithConsumersWithoutAssignedPartitions(): Unit = {
+    TestUtils.createOffsetsTopic(zkClient, servers)
+
+    for (describeType <- describeTypes) {
+      val group = this.group + describeType.mkString("")
+      // run two consumers in the group consuming from a single-partition topic
+      addConsumerGroupExecutor(numConsumers = 2, group = group)
+      val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--group", group) ++ describeType
+      val service = getConsumerGroupService(cgcArgs)
+
+      TestUtils.waitUntilTrue(() => {
+        val (output, error) = TestUtils.grabConsoleOutputAndError(service.describeGroups())
+        val expectedNumRows = if (describeTypeMembers.contains(describeType)) 3 else 2
+        error.isEmpty && output.trim.split("\n").size == expectedNumRows
+      }, s"Expected a single data row in describe group result with describe type '${describeType.mkString(" ")}'")
+    }
+  }
+
+  @Test
+  def testDescribeOffsetsWithConsumersWithoutAssignedPartitions(): Unit = {
+    TestUtils.createOffsetsTopic(zkClient, servers)
+
     // run two consumers in the group consuming from a single-partition topic
-    consumerGroupExecutor = new ConsumerGroupExecutor(brokerList, 2, group, topic)
+    addConsumerGroupExecutor(numConsumers = 2)
 
     val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--group", group)
-    val opts = new ConsumerGroupCommandOptions(cgcArgs)
-    consumerGroupService = new KafkaConsumerGroupService(opts)
+    val service = getConsumerGroupService(cgcArgs)
 
     TestUtils.waitUntilTrue(() => {
-      val (state, assignments) = consumerGroupService.describeGroup()
-      state == Some("Stable") &&
+      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")
+  }
+
+  @Test
+  def testDescribeMembersWithConsumersWithoutAssignedPartitions(): Unit = {
+    TestUtils.createOffsetsTopic(zkClient, servers)
+
+    // run two consumers in the group consuming from a single-partition topic
+    addConsumerGroupExecutor(numConsumers = 2)
+
+    val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--group", group)
+    val service = getConsumerGroupService(cgcArgs)
+
+    TestUtils.waitUntilTrue(() => {
+      val (state, assignments) = service.collectGroupMembers(group, false)
+      state.contains("Stable") &&
         assignments.isDefined &&
         assignments.get.count(_.group == group) == 2 &&
-        assignments.get.count { x => x.group == group && x.partition.isDefined } == 1 &&
-        assignments.get.count { x => x.group == group && x.partition.isEmpty } == 1
+        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")
+
+    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(): Unit = {
+    TestUtils.createOffsetsTopic(zkClient, servers)
+
+    // run two consumers in the group consuming from a single-partition topic
+    addConsumerGroupExecutor(numConsumers = 2)
+
+    val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--group", group)
+    val service = getConsumerGroupService(cgcArgs)
+
+    TestUtils.waitUntilTrue(() => {
+      val state = service.collectGroupState(group)
+      state.state == "Stable" && state.numMembers == 2
+    }, "Expected two consumers in describe group results")
+  }
+
+  @Test
+  def testDescribeWithMultiPartitionTopicAndMultipleConsumers(): Unit = {
+    TestUtils.createOffsetsTopic(zkClient, servers)
+    val topic2 = "foo2"
+    createTopic(topic2, 2, 1)
+
+    for (describeType <- describeTypes) {
+      val group = this.group + describeType.mkString("")
+      // run two consumers in the group consuming from a two-partition topic
+      addConsumerGroupExecutor(2, topic2, group = group)
+      val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--group", group) ++ describeType
+      val service = getConsumerGroupService(cgcArgs)
+
+      TestUtils.waitUntilTrue(() => {
+        val (output, error) = TestUtils.grabConsoleOutputAndError(service.describeGroups())
+        val expectedNumRows = if (describeTypeState.contains(describeType)) 2 else 3
+        error.isEmpty && output.trim.split("\n").size == expectedNumRows
+      }, s"Expected a single data row in describe group result with describe type '${describeType.mkString(" ")}'")
+    }
   }
 
   @Test
-  def testDescribeWithMultiPartitionTopicAndMultipleConsumersWithNewConsumer() {
-    TestUtils.createOffsetsTopic(zkUtils, servers)
+  def testDescribeOffsetsWithMultiPartitionTopicAndMultipleConsumers(): Unit = {
+    TestUtils.createOffsetsTopic(zkClient, servers)
     val topic2 = "foo2"
-    adminZkClient.createTopic(topic2, 2, 1)
+    createTopic(topic2, 2, 1)
 
     // run two consumers in the group consuming from a two-partition topic
-    consumerGroupExecutor = new ConsumerGroupExecutor(brokerList, 2, group, topic2)
+    addConsumerGroupExecutor(numConsumers = 2, topic2)
 
     val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--group", group)
-    val opts = new ConsumerGroupCommandOptions(cgcArgs)
-    consumerGroupService = new KafkaConsumerGroupService(opts)
+    val service = getConsumerGroupService(cgcArgs)
 
     TestUtils.waitUntilTrue(() => {
-      val (state, assignments) = consumerGroupService.describeGroup()
-      state == Some("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
+      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.")
   }
 
   @Test
-  def testDescribeGroupWithNewConsumerWithShortInitializationTimeout() {
-    // Let creation of the offsets topic happen during group initialisation to ensure that initialization doesn't
+  def testDescribeMembersWithMultiPartitionTopicAndMultipleConsumers(): Unit = {
+    TestUtils.createOffsetsTopic(zkClient, servers)
+    val topic2 = "foo2"
+    createTopic(topic2, 2, 1)
+
+    // run two consumers in the group consuming from a two-partition topic
+    addConsumerGroupExecutor(numConsumers = 2, topic2)
+
+    val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--group", group)
+    val service = getConsumerGroupService(cgcArgs)
+
+    TestUtils.waitUntilTrue(() => {
+      val (state, assignments) = service.collectGroupMembers(group, false)
+      state.contains("Stable") &&
+        assignments.isDefined &&
+        assignments.get.count(_.group == group) == 2 &&
+        assignments.get.count{ x => x.group == group && x.numPartitions == 1 } == 2 &&
+        assignments.get.count{ x => x.group == group && x.numPartitions == 0 } == 0
+    }, "Expected two rows (one row per consumer) in describe group members results.")
+
+    val (state, assignments) = service.collectGroupMembers(group, true)
+    assertTrue("Expected additional columns in verbose version of describe members",
+        state.contains("Stable") && assignments.get.count(_.assignment.isEmpty) == 0)
+  }
+
+  @Test
+  def testDescribeStateWithMultiPartitionTopicAndMultipleConsumers(): Unit = {
+    TestUtils.createOffsetsTopic(zkClient, servers)
+    val topic2 = "foo2"
+    createTopic(topic2, 2, 1)
+
+    // run two consumers in the group consuming from a two-partition topic
+    addConsumerGroupExecutor(numConsumers = 2, topic2)
+
+    val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--group", group)
+    val service = getConsumerGroupService(cgcArgs)
+
+    TestUtils.waitUntilTrue(() => {
+      val state = service.collectGroupState(group)
+      state.state == "Stable" && state.group == group && state.numMembers == 2
+    }, "Expected a stable group with two members in describe group state result.")
+  }
+
+  @Test
+  def testDescribeSimpleConsumerGroup(): Unit = {
+    // Ensure that the offsets of consumers which don't use group management are still displayed
+
+    TestUtils.createOffsetsTopic(zkClient, servers)
+    val topic2 = "foo2"
+    createTopic(topic2, 2, 1)
+    addSimpleGroupExecutor(Seq(new TopicPartition(topic2, 0), new TopicPartition(topic2, 1)))
+
+    val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--group", group)
+    val service = getConsumerGroupService(cgcArgs)
+
+    TestUtils.waitUntilTrue(() => {
+      val (state, assignments) = service.collectGroupOffsets(group)
+      state.contains("Empty") && assignments.isDefined && assignments.get.count(_.group == group) == 2
+    }, "Expected a stable group with two members in describe group state result.")
+  }
+
+  @Test
+  def testDescribeGroupWithShortInitializationTimeout(): Unit = {
+    // Let creation of the offsets topic happen during group initialization to ensure that initialization doesn't
     // complete before the timeout expires
 
+    val describeType = describeTypes(Random.nextInt(describeTypes.length))
+    val group = this.group + describeType.mkString("")
     // run one consumer in the group consuming from a single-partition topic
-    consumerGroupExecutor = new ConsumerGroupExecutor(brokerList, 1, group, topic)
+    addConsumerGroupExecutor(numConsumers = 1)
+    // set the group initialization timeout too low for the group to stabilize
+    val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--timeout", "1", "--group", group) ++ describeType
+    val service = getConsumerGroupService(cgcArgs)
+
+    try {
+      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 => assertEquals(classOf[TimeoutException], e.getCause.getClass)
+    }
+  }
+
+  @Test
+  def testDescribeGroupOffsetsWithShortInitializationTimeout(): Unit = {
+    // Let creation of the offsets topic happen during group initialization to ensure that initialization doesn't
+    // complete before the timeout expires
+
+    // run one consumer in the group consuming from a single-partition topic
+    addConsumerGroupExecutor(numConsumers = 1)
 
     // set the group initialization timeout too low for the group to stabilize
-    val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--group", "group", "--timeout", "1")
-    val opts = new ConsumerGroupCommandOptions(cgcArgs)
-    consumerGroupService = new KafkaConsumerGroupService(opts)
+    val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--group", group, "--timeout", "1")
+    val service = getConsumerGroupService(cgcArgs)
 
     try {
-      consumerGroupService.describeGroup()
+      service.collectGroupOffsets(group)
       fail("The consumer group command should fail due to low initialization timeout")
     } catch {
-      case _: TimeoutException => // OK
+      case e: ExecutionException => assertEquals(classOf[TimeoutException], e.getCause.getClass)
     }
   }
 
-  @deprecated("This test has been deprecated and will be removed in a future release.", "0.11.1.0")
-  private def createOldConsumer(): Unit = {
-    val consumerProps = new Properties
-    consumerProps.setProperty("group.id", group)
-    consumerProps.setProperty("zookeeper.connect", zkConnect)
-    oldConsumers += new OldConsumer(Whitelist(topic), consumerProps)
-  }
-}
+  @Test
+  def testDescribeGroupMembersWithShortInitializationTimeout(): Unit = {
+    // Let creation of the offsets topic happen during group initialization to ensure that initialization doesn't
+    // complete before the timeout expires
 
+    // run one consumer in the group consuming from a single-partition topic
+    addConsumerGroupExecutor(numConsumers = 1)
 
-class ConsumerThread(broker: String, id: Int, groupId: String, topic: String) extends Runnable {
-  val props = new Properties
-  props.put("bootstrap.servers", broker)
-  props.put("group.id", groupId)
-  props.put("key.deserializer", classOf[StringDeserializer].getName)
-  props.put("value.deserializer", classOf[StringDeserializer].getName)
-  val consumer = new KafkaConsumer(props)
+    // set the group initialization timeout too low for the group to stabilize
+    val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--group", group, "--timeout", "1")
+    val service = getConsumerGroupService(cgcArgs)
 
-  def run() {
     try {
-      consumer.subscribe(Collections.singleton(topic))
-      while (true)
-        consumer.poll(Long.MaxValue)
+      service.collectGroupMembers(group, false)
+      fail("The consumer group command should fail due to low initialization timeout")
     } catch {
-      case _: WakeupException => // OK
-    } finally {
-      consumer.close()
+      case e: ExecutionException => assertEquals(classOf[TimeoutException], e.getCause.getClass)
+        try {
+          service.collectGroupMembers(group, true)
+          fail("The consumer group command should fail due to low initialization timeout (verbose)")
+        } catch {
+          case e: ExecutionException => assertEquals(classOf[TimeoutException], e.getCause.getClass)
+        }
     }
   }
 
-  def shutdown() {
-    consumer.wakeup()
-  }
-}
+  @Test
+  def testDescribeGroupStateWithShortInitializationTimeout(): Unit = {
+    // Let creation of the offsets topic happen during group initialization to ensure that initialization doesn't
+    // complete before the timeout expires
 
+    // run one consumer in the group consuming from a single-partition topic
+    addConsumerGroupExecutor(numConsumers = 1)
 
-class ConsumerGroupExecutor(broker: String, numConsumers: Int, groupId: String, topic: String) {
-  val executor: ExecutorService = Executors.newFixedThreadPool(numConsumers)
-  private val consumers = new ArrayBuffer[ConsumerThread]()
-  for (i <- 1 to numConsumers) {
-    val consumer = new ConsumerThread(broker, i, groupId, topic)
-    consumers += consumer
-    executor.submit(consumer)
-  }
+    // set the group initialization timeout too low for the group to stabilize
+    val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--group", group, "--timeout", "1")
+    val service = getConsumerGroupService(cgcArgs)
 
-  def shutdown() {
-    consumers.foreach(_.shutdown)
-    executor.shutdown()
     try {
-      executor.awaitTermination(5000, TimeUnit.MILLISECONDS)
+      service.collectGroupState(group)
+      fail("The consumer group command should fail due to low initialization timeout")
     } catch {
-      case e: InterruptedException =>
-        e.printStackTrace()
+      case e: ExecutionException => assertEquals(classOf[TimeoutException], e.getCause.getClass)
     }
   }
+
+  @Test(expected = classOf[joptsimple.OptionException])
+  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(): Unit = {
+    TestUtils.createOffsetsTopic(zkClient, servers)
+
+    val customProps = new Properties
+    // create a consumer group that never commits offsets
+    customProps.setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false")
+    // run one consumer in the group consuming from a single-partition topic
+    addConsumerGroupExecutor(numConsumers = 1, customPropsOpt = Some(customProps))
+
+    val cgcArgs = Array("--bootstrap-server", brokerList, "--describe", "--group", group)
+    val service = getConsumerGroupService(cgcArgs)
+
+    TestUtils.waitUntilTrue(() => {
+      val (state, assignments) = service.collectGroupOffsets(group)
+      state.contains("Stable") &&
+        assignments.isDefined &&
+        assignments.get.count(_.group == group) == 1 &&
+        assignments.get.filter(_.group == group).head.consumerId.exists(_.trim != ConsumerGroupCommand.MISSING_COLUMN_VALUE) &&
+        assignments.get.filter(_.group == group).head.clientId.exists(_.trim != ConsumerGroupCommand.MISSING_COLUMN_VALUE) &&
+        assignments.get.filter(_.group == group).head.host.exists(_.trim != ConsumerGroupCommand.MISSING_COLUMN_VALUE)
+    }, s"Expected a 'Stable' group status, rows and valid values for consumer id / client id / host columns in describe results for non-offset-committing group $group.")
+  }
+
 }
+
diff --git a/core/src/test/scala/unit/kafka/admin/FeatureCommandTest.scala b/core/src/test/scala/unit/kafka/admin/FeatureCommandTest.scala
new file mode 100644
index 0000000000000..f548af6c2a86d
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/admin/FeatureCommandTest.scala
@@ -0,0 +1,244 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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 kafka.api.KAFKA_2_7_IV0
+import kafka.server.{BaseRequestTest, KafkaConfig, KafkaServer}
+import kafka.utils.TestUtils
+import kafka.utils.TestUtils.waitUntilTrue
+import org.apache.kafka.common.feature.{Features, SupportedVersionRange}
+import org.apache.kafka.common.utils.Utils
+
+import java.util.Properties
+
+import org.junit.Assert.{assertEquals, assertTrue}
+import org.junit.Test
+import org.scalatest.Assertions.intercept
+
+class FeatureCommandTest extends BaseRequestTest {
+  override def brokerCount: Int = 3
+
+  override def brokerPropertyOverrides(props: Properties): Unit = {
+    props.put(KafkaConfig.InterBrokerProtocolVersionProp, KAFKA_2_7_IV0.toString)
+  }
+
+  private val defaultSupportedFeatures: Features[SupportedVersionRange] =
+    Features.supportedFeatures(Utils.mkMap(Utils.mkEntry("feature_1", new SupportedVersionRange(1, 3)),
+                                           Utils.mkEntry("feature_2", new SupportedVersionRange(1, 5))))
+
+  private def updateSupportedFeatures(features: Features[SupportedVersionRange],
+                                      targetServers: Set[KafkaServer]): Unit = {
+    targetServers.foreach(s => {
+      s.brokerFeatures.setSupportedFeatures(features)
+      s.zkClient.updateBrokerInfo(s.createBrokerInfo)
+    })
+
+    // Wait until updates to all BrokerZNode supported features propagate to the controller.
+    val brokerIds = targetServers.map(s => s.config.brokerId)
+    waitUntilTrue(
+      () => servers.exists(s => {
+        if (s.kafkaController.isActive) {
+          s.kafkaController.controllerContext.liveOrShuttingDownBrokers
+            .filter(b => brokerIds.contains(b.id))
+            .forall(b => {
+              b.features.equals(features)
+            })
+        } else {
+          false
+        }
+      }),
+      "Controller did not get broker updates")
+  }
+
+  private def updateSupportedFeaturesInAllBrokers(features: Features[SupportedVersionRange]): Unit = {
+    updateSupportedFeatures(features, Set[KafkaServer]() ++ servers)
+  }
+
+  /**
+   * Tests if the FeatureApis#describeFeatures API works as expected when describing features before and
+   * after upgrading features.
+   */
+  @Test
+  def testDescribeFeaturesSuccess(): Unit = {
+    updateSupportedFeaturesInAllBrokers(defaultSupportedFeatures)
+    val featureApis = new FeatureApis(new FeatureCommandOptions(Array("--bootstrap-server", brokerList, "--describe")))
+    featureApis.setSupportedFeatures(defaultSupportedFeatures)
+    try {
+      val initialDescribeOutput = TestUtils.grabConsoleOutput(featureApis.describeFeatures())
+      val expectedInitialDescribeOutput =
+        "Feature: feature_1\tSupportedMinVersion: 1\tSupportedMaxVersion: 3\tFinalizedMinVersionLevel: -\tFinalizedMaxVersionLevel: -\tEpoch: 0\n" +
+        "Feature: feature_2\tSupportedMinVersion: 1\tSupportedMaxVersion: 5\tFinalizedMinVersionLevel: -\tFinalizedMaxVersionLevel: -\tEpoch: 0\n"
+      assertEquals(expectedInitialDescribeOutput, initialDescribeOutput)
+      featureApis.upgradeAllFeatures()
+      val finalDescribeOutput = TestUtils.grabConsoleOutput(featureApis.describeFeatures())
+      val expectedFinalDescribeOutput =
+        "Feature: feature_1\tSupportedMinVersion: 1\tSupportedMaxVersion: 3\tFinalizedMinVersionLevel: 1\tFinalizedMaxVersionLevel: 3\tEpoch: 1\n" +
+        "Feature: feature_2\tSupportedMinVersion: 1\tSupportedMaxVersion: 5\tFinalizedMinVersionLevel: 1\tFinalizedMaxVersionLevel: 5\tEpoch: 1\n"
+      assertEquals(expectedFinalDescribeOutput, finalDescribeOutput)
+    } finally {
+      featureApis.close()
+    }
+  }
+
+  /**
+   * Tests if the FeatureApis#upgradeAllFeatures API works as expected during a success case.
+   */
+  @Test
+  def testUpgradeAllFeaturesSuccess(): Unit = {
+    val upgradeOpts = new FeatureCommandOptions(Array("--bootstrap-server", brokerList, "--upgrade-all"))
+    val featureApis = new FeatureApis(upgradeOpts)
+    try {
+      // Step (1):
+      // - Update the supported features across all brokers.
+      // - Upgrade non-existing feature_1 to maxVersionLevel: 2.
+      // - Verify results.
+      val initialSupportedFeatures = Features.supportedFeatures(Utils.mkMap(Utils.mkEntry("feature_1", new SupportedVersionRange(1, 2))))
+      updateSupportedFeaturesInAllBrokers(initialSupportedFeatures)
+      featureApis.setSupportedFeatures(initialSupportedFeatures)
+      var output = TestUtils.grabConsoleOutput(featureApis.upgradeAllFeatures())
+      var expected =
+        "      [Add]\tFeature: feature_1\tExistingFinalizedMaxVersion: -\tNewFinalizedMaxVersion: 2\tResult: OK\n"
+      assertEquals(expected, output)
+
+      // Step (2):
+      // - Update the supported features across all brokers.
+      // - Upgrade existing feature_1 to maxVersionLevel: 3.
+      // - Upgrade non-existing feature_2 to maxVersionLevel: 5.
+      // - Verify results.
+      updateSupportedFeaturesInAllBrokers(defaultSupportedFeatures)
+      featureApis.setSupportedFeatures(defaultSupportedFeatures)
+      output = TestUtils.grabConsoleOutput(featureApis.upgradeAllFeatures())
+      expected =
+        "  [Upgrade]\tFeature: feature_1\tExistingFinalizedMaxVersion: 2\tNewFinalizedMaxVersion: 3\tResult: OK\n" +
+        "      [Add]\tFeature: feature_2\tExistingFinalizedMaxVersion: -\tNewFinalizedMaxVersion: 5\tResult: OK\n"
+      assertEquals(expected, output)
+
+      // Step (3):
+      // - Perform an upgrade of all features again.
+      // - Since supported features have not changed, expect that the above action does not yield
+      //   any results.
+      output = TestUtils.grabConsoleOutput(featureApis.upgradeAllFeatures())
+      assertTrue(output.isEmpty)
+      featureApis.setOptions(upgradeOpts)
+      output = TestUtils.grabConsoleOutput(featureApis.upgradeAllFeatures())
+      assertTrue(output.isEmpty)
+    } finally {
+      featureApis.close()
+    }
+  }
+
+  /**
+   * Tests if the FeatureApis#downgradeAllFeatures API works as expected during a success case.
+   */
+  @Test
+  def testDowngradeFeaturesSuccess(): Unit = {
+    val downgradeOpts = new FeatureCommandOptions(Array("--bootstrap-server", brokerList, "--downgrade-all"))
+    val upgradeOpts = new FeatureCommandOptions(Array("--bootstrap-server", brokerList, "--upgrade-all"))
+    val featureApis = new FeatureApis(upgradeOpts)
+    try {
+      // Step (1):
+      // - Update the supported features across all brokers.
+      // - Upgrade non-existing feature_1 to maxVersionLevel: 3.
+      // - Upgrade non-existing feature_2 to maxVersionLevel: 5.
+      updateSupportedFeaturesInAllBrokers(defaultSupportedFeatures)
+      featureApis.setSupportedFeatures(defaultSupportedFeatures)
+      featureApis.upgradeAllFeatures()
+
+      // Step (2):
+      // - Downgrade existing feature_1 to maxVersionLevel: 2.
+      // - Delete feature_2 since it is no longer supported by the FeatureApis object.
+      // - Verify results.
+      val downgradedFeatures = Features.supportedFeatures(Utils.mkMap(Utils.mkEntry("feature_1", new SupportedVersionRange(1, 2))))
+      featureApis.setSupportedFeatures(downgradedFeatures)
+      featureApis.setOptions(downgradeOpts)
+      var output = TestUtils.grabConsoleOutput(featureApis.downgradeAllFeatures())
+      var expected =
+        "[Downgrade]\tFeature: feature_1\tExistingFinalizedMaxVersion: 3\tNewFinalizedMaxVersion: 2\tResult: OK\n" +
+        "   [Delete]\tFeature: feature_2\tExistingFinalizedMaxVersion: 5\tNewFinalizedMaxVersion: -\tResult: OK\n"
+      assertEquals(expected, output)
+
+      // Step (3):
+      // - Perform a downgrade of all features again.
+      // - Since supported features have not changed, expect that the above action does not yield
+      //   any results.
+      updateSupportedFeaturesInAllBrokers(downgradedFeatures)
+      output = TestUtils.grabConsoleOutput(featureApis.downgradeAllFeatures())
+      assertTrue(output.isEmpty)
+
+      // Step (4):
+      // - Delete feature_1 since it is no longer supported by the FeatureApis object.
+      // - Verify results.
+      featureApis.setSupportedFeatures(Features.emptySupportedFeatures())
+      output = TestUtils.grabConsoleOutput(featureApis.downgradeAllFeatures())
+      expected =
+        "   [Delete]\tFeature: feature_1\tExistingFinalizedMaxVersion: 2\tNewFinalizedMaxVersion: -\tResult: OK\n"
+      assertEquals(expected, output)
+    } finally {
+      featureApis.close()
+    }
+  }
+
+  /**
+   * Tests if the FeatureApis#upgradeAllFeatures API works as expected during a partial failure case.
+   */
+  @Test
+  def testUpgradeFeaturesFailure(): Unit = {
+    val upgradeOpts = new FeatureCommandOptions(Array("--bootstrap-server", brokerList, "--upgrade-all"))
+    val featureApis = new FeatureApis(upgradeOpts)
+    try {
+      // Step (1): Update the supported features across all brokers.
+      updateSupportedFeaturesInAllBrokers(defaultSupportedFeatures)
+
+      // Step (2):
+      // - Intentionally setup the FeatureApis object such that it contains incompatible target
+      //   features (viz. feature_2 and feature_3).
+      // - Upgrade non-existing feature_1 to maxVersionLevel: 4. Expect the operation to fail with
+      //   an incompatibility failure.
+      // - Upgrade non-existing feature_2 to maxVersionLevel: 5. Expect the operation to succeed.
+      // - Upgrade non-existing feature_3 to maxVersionLevel: 3. Expect the operation to fail
+      //   since the feature is not supported.
+      val targetFeaturesWithIncompatibilities =
+        Features.supportedFeatures(
+          Utils.mkMap(Utils.mkEntry("feature_1", new SupportedVersionRange(1, 4)),
+                      Utils.mkEntry("feature_2", new SupportedVersionRange(1, 5)),
+                      Utils.mkEntry("feature_3", new SupportedVersionRange(1, 3))))
+      featureApis.setSupportedFeatures(targetFeaturesWithIncompatibilities)
+      val output = TestUtils.grabConsoleOutput({
+        val exception = intercept[UpdateFeaturesException] {
+          featureApis.upgradeAllFeatures()
+        }
+        assertEquals("2 feature updates failed!", exception.getMessage)
+      })
+      val expected =
+        "      [Add]\tFeature: feature_1\tExistingFinalizedMaxVersion: -" +
+        "\tNewFinalizedMaxVersion: 4\tResult: FAILED due to" +
+        " org.apache.kafka.common.errors.InvalidRequestException: Could not apply finalized" +
+        " feature update because brokers were found to have incompatible versions for the" +
+        " feature.\n" +
+        "      [Add]\tFeature: feature_2\tExistingFinalizedMaxVersion: -" +
+        "\tNewFinalizedMaxVersion: 5\tResult: OK\n" +
+        "      [Add]\tFeature: feature_3\tExistingFinalizedMaxVersion: -" +
+        "\tNewFinalizedMaxVersion: 3\tResult: FAILED due to" +
+        " org.apache.kafka.common.errors.InvalidRequestException: Could not apply finalized" +
+        " feature update because the provided feature is not supported.\n"
+      assertEquals(expected, output)
+    } finally {
+      featureApis.close()
+    }
+  }
+}
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..3cfba5edfebc8
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/admin/LeaderElectionCommandTest.scala
@@ -0,0 +1,331 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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.{Admin, 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.jdk.CollectionConverters._
+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(Admin.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(Admin.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(Admin.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(Admin.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.DEFAULT_API_TIMEOUT_MS_CONFIG -> "20000",
+      AdminClientConfig.REQUEST_TIMEOUT_MS_CONFIG -> "10000"
+    )
+  }
+
+  def bootstrapServers(servers: Seq[KafkaServer]): String = {
+    TestUtils.bootstrapServers(servers, new ListenerName("PLAINTEXT"))
+  }
+
+  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 4b2289890f5b2..3997c9ef53709 100644
--- a/core/src/test/scala/unit/kafka/admin/ListConsumerGroupTest.scala
+++ b/core/src/test/scala/unit/kafka/admin/ListConsumerGroupTest.scala
@@ -16,76 +16,121 @@
  */
 package kafka.admin
 
-import java.util.Properties
-
-import org.easymock.EasyMock
-import org.junit.Before
+import joptsimple.OptionException
+import org.junit.Assert._
 import org.junit.Test
-
-import kafka.admin.ConsumerGroupCommand.ConsumerGroupCommandOptions
-import kafka.admin.ConsumerGroupCommand.ZkConsumerGroupService
-import kafka.consumer.OldConsumer
-import kafka.consumer.Whitelist
-import kafka.integration.KafkaServerTestHarness
-import kafka.server.KafkaConfig
 import kafka.utils.TestUtils
+import org.apache.kafka.common.ConsumerGroupState
+import org.apache.kafka.clients.admin.ConsumerGroupListing
+import java.util.Optional
+import org.scalatest.Assertions.assertThrows
 
-class ListConsumerGroupTest extends KafkaServerTestHarness {
+class ListConsumerGroupTest extends ConsumerGroupCommandTest {
 
-  val overridingProps = new Properties()
-  val topic = "foo"
-  val topicFilter = Whitelist(topic)
-  val group = "test.group"
-  val props = new Properties
+  @Test
+  def testListConsumerGroups(): Unit = {
+    val simpleGroup = "simple-group"
+    addSimpleGroupExecutor(group = simpleGroup)
+    addConsumerGroupExecutor(numConsumers = 1)
 
-  // configure the servers and clients
-  override def generateConfigs =
-    TestUtils.createBrokerConfigs(1, zkConnect, enableControlledShutdown = false).map(KafkaConfig.fromProps(_, overridingProps))
+    val cgcArgs = Array("--bootstrap-server", brokerList, "--list")
+    val service = getConsumerGroupService(cgcArgs)
 
-  @Before
-  override def setUp() {
-    super.setUp()
+    val expectedGroups = Set(group, simpleGroup)
+    var foundGroups = Set.empty[String]
+    TestUtils.waitUntilTrue(() => {
+      foundGroups = service.listConsumerGroups().toSet
+      expectedGroups == foundGroups
+    }, s"Expected --list to show groups $expectedGroups, but found $foundGroups.")
+  }
 
-    adminZkClient.createTopic(topic, 1, 1)
-    props.setProperty("group.id", group)
-    props.setProperty("zookeeper.connect", zkConnect)
+  @Test(expected = classOf[OptionException])
+  def testListWithUnrecognizedNewConsumerOption(): Unit = {
+    val cgcArgs = Array("--new-consumer", "--bootstrap-server", brokerList, "--list")
+    getConsumerGroupService(cgcArgs)
   }
 
   @Test
-  def testListGroupWithNoExistingGroup() {
-    val opts = new ConsumerGroupCommandOptions(Array("--zookeeper", zkConnect))
-    val consumerGroupCommand = new ZkConsumerGroupService(opts)
-    try {
-      assert(consumerGroupCommand.listGroups().isEmpty)
-    } finally {
-      consumerGroupCommand.close()
+  def testListConsumerGroupsWithStates(): Unit = {
+    val simpleGroup = "simple-group"
+    addSimpleGroupExecutor(group = simpleGroup)
+    addConsumerGroupExecutor(numConsumers = 1)
+
+    val cgcArgs = Array("--bootstrap-server", brokerList, "--list", "--state")
+    val service = getConsumerGroupService(cgcArgs)
+
+    val expectedListing = Set(
+      new ConsumerGroupListing(simpleGroup, true, Optional.of(ConsumerGroupState.EMPTY)),
+      new ConsumerGroupListing(group, false, Optional.of(ConsumerGroupState.STABLE)))
+
+    var foundListing = Set.empty[ConsumerGroupListing]
+    TestUtils.waitUntilTrue(() => {
+      foundListing = service.listConsumerGroupsWithState(ConsumerGroupState.values.toSet).toSet
+      expectedListing == foundListing
+    }, s"Expected to show groups $expectedListing, but found $foundListing")
+
+    val expectedListingStable = Set(
+      new ConsumerGroupListing(group, false, Optional.of(ConsumerGroupState.STABLE)))
+
+    foundListing = Set.empty[ConsumerGroupListing]
+    TestUtils.waitUntilTrue(() => {
+      foundListing = service.listConsumerGroupsWithState(Set(ConsumerGroupState.STABLE)).toSet
+      expectedListingStable == foundListing
+    }, s"Expected to show groups $expectedListingStable, but found $foundListing")
+  }
+
+  @Test
+  def testConsumerGroupStatesFromString(): Unit = {
+    var result = ConsumerGroupCommand.consumerGroupStatesFromString("Stable")
+    assertEquals(Set(ConsumerGroupState.STABLE), result)
+
+    result = ConsumerGroupCommand.consumerGroupStatesFromString("Stable, PreparingRebalance")
+    assertEquals(Set(ConsumerGroupState.STABLE, ConsumerGroupState.PREPARING_REBALANCE), result)
+
+    result = ConsumerGroupCommand.consumerGroupStatesFromString("Dead,CompletingRebalance,")
+    assertEquals(Set(ConsumerGroupState.DEAD, ConsumerGroupState.COMPLETING_REBALANCE), result)
+
+    assertThrows[IllegalArgumentException] {
+      ConsumerGroupCommand.consumerGroupStatesFromString("bad, wrong")
+    }
+
+    assertThrows[IllegalArgumentException] {
+      ConsumerGroupCommand.consumerGroupStatesFromString("stable")
+    }
+
+    assertThrows[IllegalArgumentException] {
+      ConsumerGroupCommand.consumerGroupStatesFromString("  bad, Stable")
+    }
+
+    assertThrows[IllegalArgumentException] {
+      ConsumerGroupCommand.consumerGroupStatesFromString("   ,   ,")
     }
   }
 
   @Test
-  def testListGroupWithSomeGroups() {
-    // mocks
-    val consumer1Mock = EasyMock.createMockBuilder(classOf[OldConsumer]).withConstructor(topicFilter, props).createMock()
-    props.setProperty("group.id", "some.other.group")
-    val consumer2Mock = EasyMock.createMockBuilder(classOf[OldConsumer]).withConstructor(topicFilter, props).createMock()
+  def testListGroupCommand(): Unit = {
+    val simpleGroup = "simple-group"
+    addSimpleGroupExecutor(group = simpleGroup)
+    addConsumerGroupExecutor(numConsumers = 1)
+    var out = ""
 
-    // stubs
-    val opts = new ConsumerGroupCommandOptions(Array("--zookeeper", zkConnect))
-    val consumerGroupCommand = new ZkConsumerGroupService(opts)
+    var cgcArgs = Array("--bootstrap-server", brokerList, "--list")
+    TestUtils.waitUntilTrue(() => {
+      out = TestUtils.grabConsoleOutput(ConsumerGroupCommand.main(cgcArgs))
+      !out.contains("STATE") && out.contains(simpleGroup) && out.contains(group)
+    }, s"Expected to find $simpleGroup, $group and no header, but found $out")
 
-    // simulation
-    EasyMock.replay(consumer1Mock)
-    EasyMock.replay(consumer2Mock)
+    cgcArgs = Array("--bootstrap-server", brokerList, "--list", "--state")
+    TestUtils.waitUntilTrue(() => {
+      out = TestUtils.grabConsoleOutput(ConsumerGroupCommand.main(cgcArgs))
+      out.contains("STATE") && out.contains(simpleGroup) && out.contains(group)
+    }, s"Expected to find $simpleGroup, $group and the header, but found $out")
 
-    // action/test
+    cgcArgs = Array("--bootstrap-server", brokerList, "--list", "--state", "Stable")
     TestUtils.waitUntilTrue(() => {
-        val groups = consumerGroupCommand.listGroups()
-        groups.size == 2 && groups.contains(group)
-      }, "Expected a different list group results.")
-
-    // cleanup
-    consumerGroupCommand.close()
-    consumer1Mock.stop()
-    consumer2Mock.stop()
+      out = TestUtils.grabConsoleOutput(ConsumerGroupCommand.main(cgcArgs))
+      out.contains("STATE") && out.contains(group) && out.contains("Stable")
+    }, s"Expected to find $group in state Stable and the header, but found $out")
   }
+
 }
diff --git a/core/src/test/scala/unit/kafka/admin/PreferredReplicaLeaderElectionCommandTest.scala b/core/src/test/scala/unit/kafka/admin/PreferredReplicaLeaderElectionCommandTest.scala
new file mode 100644
index 0000000000000..1b715eb5ebd85
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/admin/PreferredReplicaLeaderElectionCommandTest.scala
@@ -0,0 +1,389 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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, Paths}
+import java.util
+import java.util.Properties
+
+import scala.collection.Seq
+import kafka.common.AdminCommandFailedException
+import kafka.security.authorizer.AclAuthorizer
+import kafka.server.{KafkaConfig, KafkaServer}
+import kafka.utils.{Logging, TestUtils}
+import kafka.zk.ZooKeeperTestHarness
+import org.apache.kafka.common.{TopicPartition, Uuid}
+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.jdk.CollectionConverters._
+
+class PreferredReplicaLeaderElectionCommandTest extends ZooKeeperTestHarness with Logging {
+  var servers: Seq[KafkaServer] = Seq()
+
+  @After
+  override def tearDown(): Unit = {
+    TestUtils.shutdownServers(servers)
+    super.tearDown()
+  }
+
+  private def createTestTopicAndCluster(topicPartition: Map[TopicPartition, List[Int]],
+                                        authorizer: Option[String] = None): Unit = {
+    val brokerConfigs = TestUtils.createBrokerConfigs(3, zkConnect, false)
+    brokerConfigs.foreach(p => p.setProperty("auto.leader.rebalance.enable", "false"))
+    authorizer match {
+      case Some(className) =>
+        brokerConfigs.foreach(p => p.setProperty("authorizer.class.name", className))
+      case None =>
+    }
+    createTestTopicAndCluster(topicPartition, brokerConfigs)
+  }
+
+  private def createTestTopicAndCluster(partitionsAndAssignments: Map[TopicPartition, List[Int]],
+                                        brokerConfigs: Seq[Properties]): Unit = {
+    // create brokers
+    servers = brokerConfigs.map(b => TestUtils.createServer(KafkaConfig.fromProps(b)))
+    // create the topic
+    partitionsAndAssignments.foreach { case (tp, assignment) =>
+      zkClient.createTopicAssignment(tp.topic, Uuid.randomUuid(),
+      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"
+    )
+  }
+
+  /** Bounce the given targetServer and wait for all servers to get metadata for the given partition */
+  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.isr.contains(targetServer)
+      }
+    },
+      s"Replicas for partition $partition not created")
+  }
+
+  private def getController() = {
+    servers.find(p => p.kafkaController.isActive)
+  }
+
+  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 = {
+    val port = servers(broker).socketServer.boundPort(ListenerName.normalised("PLAINTEXT"))
+    debug("Server bound to port "+port)
+    s"localhost:$port"
+  }
+
+  val testPartition = new TopicPartition("test", 0)
+  val testPartitionAssignment = List(1, 2, 0)
+  val testPartitionPreferredLeader = testPartitionAssignment.head
+  val testPartitionAndAssignment = Map(testPartition -> testPartitionAssignment)
+
+  /** Test the case multiple values are given for --bootstrap-broker */
+  @Test
+  def testMultipleBrokersGiven(): Unit = {
+    createTestTopicAndCluster(testPartitionAndAssignment)
+    bounceServer(testPartitionPreferredLeader, testPartition)
+    // Check the leader for the partition is not the preferred one
+    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, awaitLeader(testPartition))
+  }
+
+  /** Test the case when an invalid broker is given for --bootstrap-broker */
+  @Test
+  def testInvalidBrokerGiven(): Unit = {
+    try {
+      PreferredReplicaLeaderElectionCommand.run(Array(
+        "--bootstrap-server", "example.com:1234"),
+        timeout = 1000)
+      fail()
+    } catch {
+      case e: AdminCommandFailedException =>
+        assertTrue(e.getCause.isInstanceOf[TimeoutException])
+    }
+  }
+
+  /** Test the case where no partitions are given (=> elect all partitions) */
+  @Test
+  def testNoPartitionsGiven(): Unit = {
+    createTestTopicAndCluster(testPartitionAndAssignment)
+    bounceServer(testPartitionPreferredLeader, testPartition)
+    // Check the leader for the partition is not the preferred one
+    assertNotEquals(testPartitionPreferredLeader, awaitLeader(testPartition))
+    PreferredReplicaLeaderElectionCommand.run(Array(
+      "--bootstrap-server", bootstrapServer()))
+    // Check the leader for the partition IS the preferred one
+    assertEquals(testPartitionPreferredLeader, awaitLeader(testPartition))
+  }
+
+  private def toJsonFile(partitions: Set[TopicPartition]): File = {
+    val jsonFile = File.createTempFile("preferredreplicaelection", ".js")
+    jsonFile.deleteOnExit()
+    val jsonString = TestUtils.stringifyTopicPartitions(partitions)
+    debug("Using json: "+jsonString)
+    Files.write(Paths.get(jsonFile.getAbsolutePath), jsonString.getBytes(StandardCharsets.UTF_8))
+    jsonFile
+  }
+
+  /** Test the case where a list of partitions is given */
+  @Test
+  def testSingletonPartitionGiven(): Unit = {
+    createTestTopicAndCluster(testPartitionAndAssignment)
+    bounceServer(testPartitionPreferredLeader, testPartition)
+    // Check the leader for the partition is not the preferred one
+    assertNotEquals(testPartitionPreferredLeader, awaitLeader(testPartition))
+    val jsonFile = toJsonFile(testPartitionAndAssignment.keySet)
+    try {
+      PreferredReplicaLeaderElectionCommand.run(Array(
+        "--bootstrap-server", bootstrapServer(),
+        "--path-to-json-file", jsonFile.getAbsolutePath))
+    } finally {
+      jsonFile.delete()
+    }
+    // Check the leader for the partition IS the preferred one
+    assertEquals(testPartitionPreferredLeader, awaitLeader(testPartition))
+  }
+
+  /** Test the case where a topic does not exist */
+  @Test
+  def testTopicDoesNotExist(): Unit = {
+    val nonExistentPartition = new TopicPartition("does.not.exist", 0)
+    val nonExistentPartitionAssignment = List(1, 2, 0)
+    val nonExistentPartitionAndAssignment = Map(nonExistentPartition -> nonExistentPartitionAssignment)
+
+    createTestTopicAndCluster(testPartitionAndAssignment)
+    val jsonFile = toJsonFile(nonExistentPartitionAndAssignment.keySet)
+    try {
+      PreferredReplicaLeaderElectionCommand.run(Array(
+        "--bootstrap-server", bootstrapServer(),
+        "--path-to-json-file", jsonFile.getAbsolutePath))
+    } catch {
+      case e: AdminCommandFailedException =>
+        val suppressed = e.getSuppressed()(0)
+        assertTrue(suppressed.isInstanceOf[UnknownTopicOrPartitionException])
+      case e: Throwable =>
+        e.printStackTrace()
+        throw e
+    } finally {
+      jsonFile.delete()
+    }
+  }
+
+  /** Test the case where several partitions are given */
+  @Test
+  def testMultiplePartitionsSameAssignment(): Unit = {
+    val testPartitionA = new TopicPartition("testA", 0)
+    val testPartitionB = new TopicPartition("testB", 0)
+    val testPartitionAssignment = List(1, 2, 0)
+    val testPartitionPreferredLeader = testPartitionAssignment.head
+    val testPartitionAndAssignment = Map(testPartitionA -> testPartitionAssignment, testPartitionB -> testPartitionAssignment)
+
+    createTestTopicAndCluster(testPartitionAndAssignment)
+    bounceServer(testPartitionPreferredLeader, testPartitionA)
+    // Check the leader for the partition is not the preferred one
+    assertNotEquals(testPartitionPreferredLeader, awaitLeader(testPartitionA))
+    assertNotEquals(testPartitionPreferredLeader, awaitLeader(testPartitionB))
+    val jsonFile = toJsonFile(testPartitionAndAssignment.keySet)
+    try {
+      PreferredReplicaLeaderElectionCommand.run(Array(
+        "--bootstrap-server", bootstrapServer(),
+        "--path-to-json-file", jsonFile.getAbsolutePath))
+    } finally {
+      jsonFile.delete()
+    }
+    // Check the leader for the partition IS the preferred one
+    assertEquals(testPartitionPreferredLeader, awaitLeader(testPartitionA))
+    assertEquals(testPartitionPreferredLeader, awaitLeader(testPartitionB))
+  }
+
+  /** What happens when the preferred replica is already the leader? */
+  @Test
+  def testNoopElection(): Unit = {
+    createTestTopicAndCluster(testPartitionAndAssignment)
+    // 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
+      PreferredReplicaLeaderElectionCommand.run(Array(
+        "--bootstrap-server", bootstrapServer(),
+        "--path-to-json-file", jsonFile.getAbsolutePath))
+      // Check the leader for the partition still is the preferred one
+      assertEquals(testPartitionPreferredLeader, awaitLeader(testPartition))
+    } finally {
+      jsonFile.delete()
+    }
+  }
+
+  /** What happens if the preferred replica is offline? */
+  @Test
+  def testWithOfflinePreferredReplica(): Unit = {
+    createTestTopicAndCluster(testPartitionAndAssignment)
+    bounceServer(testPartitionPreferredLeader, testPartition)
+    // Check the leader for the partition is not the preferred one
+    val leader = awaitLeader(testPartition)
+    assertNotEquals(testPartitionPreferredLeader, leader)
+    // Now kill the preferred one
+    servers(testPartitionPreferredLeader).shutdown()
+    // Now try to elect the preferred one
+    val jsonFile = toJsonFile(testPartitionAndAssignment.keySet)
+    try {
+      PreferredReplicaLeaderElectionCommand.run(Array(
+        "--bootstrap-server", bootstrapServer(),
+        "--path-to-json-file", jsonFile.getAbsolutePath))
+      fail();
+    } catch {
+      case e: AdminCommandFailedException =>
+        assertEquals("1 preferred replica(s) could not be elected", e.getMessage)
+        val suppressed = e.getSuppressed()(0)
+        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, awaitLeader(testPartition))
+    } finally {
+      jsonFile.delete()
+    }
+  }
+
+  /** What happens if the controller gets killed just before an election? */
+  @Test
+  def testTimeout(): Unit = {
+    createTestTopicAndCluster(testPartitionAndAssignment)
+    bounceServer(testPartitionPreferredLeader, testPartition)
+    // Check the leader for the partition is not the preferred one
+    val leader = awaitLeader(testPartition)
+    assertNotEquals(testPartitionPreferredLeader, leader)
+    // Now kill the controller just before we trigger the election
+    val controller = getController().get.config.brokerId
+    servers(controller).shutdown()
+    val jsonFile = toJsonFile(testPartitionAndAssignment.keySet)
+    try {
+      PreferredReplicaLeaderElectionCommand.run(Array(
+        "--bootstrap-server", bootstrapServer(controller),
+        "--path-to-json-file", jsonFile.getAbsolutePath),
+        timeout = 2000)
+      fail();
+    } catch {
+      case e: AdminCommandFailedException =>
+        assertEquals("Timeout waiting for election results", e.getMessage)
+        // Check we still have the same leader
+        assertEquals(leader, awaitLeader(testPartition))
+    } finally {
+      jsonFile.delete()
+    }
+  }
+
+  /** Test the case where client is not authorized */
+  @Test
+  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 = awaitLeader(testPartition)
+    assertNotEquals(testPartitionPreferredLeader, leader)
+    // Check the leader for the partition is not the preferred one
+    assertNotEquals(testPartitionPreferredLeader, awaitLeader(testPartition))
+    val jsonFile = toJsonFile(testPartitionAndAssignment.keySet)
+    try {
+      PreferredReplicaLeaderElectionCommand.run(Array(
+        "--bootstrap-server", bootstrapServer(),
+        "--path-to-json-file", jsonFile.getAbsolutePath))
+      fail()
+    } catch {
+      case e: AdminCommandFailedException =>
+        assertEquals("Not authorized to perform leader election", e.getMessage)
+        assertTrue(e.getCause.isInstanceOf[ClusterAuthorizationException])
+        // Check we still have the same leader
+        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.reverse.map(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 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
deleted file mode 100644
index 11f75e22c9321..0000000000000
--- a/core/src/test/scala/unit/kafka/admin/ReassignPartitionsClusterTest.scala
+++ /dev/null
@@ -1,577 +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.Collections
-import java.util.Properties
-
-import kafka.admin.ReassignPartitionsCommand._
-import kafka.common.{AdminCommandFailedException, TopicAndPartition}
-import kafka.server.{KafkaConfig, KafkaServer}
-import kafka.utils.TestUtils._
-import kafka.utils.ZkUtils._
-import kafka.utils.{Logging, TestUtils, ZkUtils}
-import kafka.zk.ZooKeeperTestHarness
-import org.junit.Assert.{assertEquals, 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.common.TopicPartitionReplica
-
-import scala.collection.JavaConverters._
-import scala.collection.Map
-import scala.collection.Seq
-import scala.util.Random
-
-import java.io.File
-
-class ReassignPartitionsClusterTest extends ZooKeeperTestHarness with Logging {
-  val partitionId = 0
-  var servers: Seq[KafkaServer] = null
-  val topicName = "my-topic"
-  val delayMs = 1000
-  var adminClient: JAdminClient = null
-
-  def zkUpdateDelay(): Unit = Thread.sleep(delayMs)
-
-  @Before
-  override def setUp() {
-    super.setUp()
-  }
-
-  def startBrokers(brokerIds: Seq[Int]) {
-    servers = brokerIds.map(i => createBrokerConfig(i, zkConnect, logDirCount = 3))
-      .map(c => createServer(KafkaConfig.fromProps(c)))
-  }
-
-  def createAdminClient(servers: Seq[KafkaServer]): JAdminClient = {
-    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 = {
-    val server = servers.find(_.config.brokerId == brokerId).get
-    val logDirs = server.config.logDirs
-    new File(logDirs(Random.nextInt(logDirs.size))).getAbsolutePath
-  }
-
-  @After
-  override def tearDown() {
-    if (adminClient != null) {
-      adminClient.close()
-      adminClient = null
-    }
-    TestUtils.shutdownServers(servers)
-    super.tearDown()
-  }
-
-  @Test
-  def shouldMoveSinglePartition(): Unit = {
-    //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(zkUtils, topicName, Map(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"]}]}"""
-    ReassignPartitionsCommand.executeAssignment(zkUtils, Some(adminClient), topicJson, NoThrottle)
-    waitForReassignmentToComplete()
-
-    //Then the replica should be on 101
-    assertEquals(Seq(101), zkUtils.getPartitionAssignmentForTopics(Seq(topicName)).get(topicName).get(partition))
-    // 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() {
-    // Given a single replica on server 100
-    startBrokers(Seq(100, 101))
-    adminClient = createAdminClient(servers)
-    val expectedLogDir = getRandomLogDirAssignment(100)
-    createTopic(zkUtils, topicName, Map(0 -> Seq(100)), servers = servers)
-
-    // 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"]}]}"""
-    ReassignPartitionsCommand.executeAssignment(zkUtils, Some(adminClient), topicJson, NoThrottle)
-    val replica = new TopicPartitionReplica(topicName, 0, 100)
-    TestUtils.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() {
-    val brokers = Array(100, 101, 102)
-    startBrokers(brokers)
-    adminClient = createAdminClient(servers)
-    createTopic(zkUtils, topicName, Map(
-      0 -> Seq(100, 101),
-      1 -> Seq(100, 101),
-      2 -> Seq(100, 101)
-    ), servers = servers)
-
-    //When rebalancing
-    val newAssignment = generateAssignment(zkUtils, brokers, json(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 (replica, brokerIds) => brokerIds.contains(102)}.get._1.partition
-    val replica1 = new TopicPartitionReplica(topicName, partition1, 102)
-    val expectedLogDir1 = getRandomLogDirAssignment(102)
-    // Find a partition in the new assignment on broker 100 and a random log directory on broker 100,
-    // which currently has partition for this topic
-    val partition2 = newAssignment.find { case (replica, brokerIds) => brokerIds.contains(100)}.get._1.partition
-    val replica2 = new TopicPartitionReplica(topicName, partition2, 100)
-    val expectedLogDir2 = getRandomLogDirAssignment(100)
-    // Generate a replica assignment to reassign replicas on broker 100 and 102 respectively to a random log directory on the same broker.
-    // Before this reassignment, the replica already exists on broker 100 but does not exist on broker 102
-    val newReplicaAssignment = Map(replica1 -> expectedLogDir1, replica2 -> expectedLogDir2)
-    ReassignPartitionsCommand.executeAssignment(zkUtils, Some(adminClient),
-      ReassignPartitionsCommand.formatAsReassignmentJson(newAssignment, newReplicaAssignment), NoThrottle)
-    waitForReassignmentToComplete()
-
-    // Then the replicas should span all three brokers
-    val actual = zkUtils.getPartitionAssignmentForTopics(Seq(topicName))(topicName)
-    assertEquals(Seq(100, 101, 102), actual.values.flatten.toSeq.distinct.sorted)
-    // 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
-    }, "Partition should have been moved to the expected log directory on broker 102", 1000)
-    waitUntilTrue(() => {
-      expectedLogDir2 == adminClient.describeReplicaLogDirs(Collections.singleton(replica2)).all().get.get(replica2).getCurrentReplicaLogDir
-    }, "Partition should have been moved to the expected log directory on broker 100", 1000)
-  }
-
-  @Test
-  def shouldShrinkCluster() {
-    //Given partitions on 3 of 3 brokers
-    val brokers = Array(100, 101, 102)
-    startBrokers(brokers)
-    createTopic(zkUtils, topicName, Map(
-      0 -> Seq(100, 101),
-      1 -> Seq(101, 102),
-      2 -> Seq(102, 100)
-    ), servers = servers)
-
-    //When rebalancing
-    val newAssignment = generateAssignment(zkUtils, Array(100, 101), json(topicName), true)._1
-    ReassignPartitionsCommand.executeAssignment(zkUtils, None,
-      ReassignPartitionsCommand.formatAsReassignmentJson(newAssignment, Map.empty), NoThrottle)
-    waitForReassignmentToComplete()
-
-    //Then replicas should only span the first two brokers
-    val actual = zkUtils.getPartitionAssignmentForTopics(Seq(topicName))(topicName)
-    assertEquals(Seq(100, 101), actual.values.flatten.toSeq.distinct.sorted)
-  }
-
-  @Test
-  def shouldMoveSubsetOfPartitions() {
-    //Given partitions on 3 of 3 brokers
-    val brokers = Array(100, 101, 102)
-    startBrokers(brokers)
-    adminClient = createAdminClient(servers)
-    createTopic(zkUtils, "topic1", Map(
-      0 -> Seq(100, 101),
-      1 -> Seq(101, 102),
-      2 -> Seq(102, 100)
-    ), servers = servers)
-    createTopic(zkUtils, "topic2", Map(
-      0 -> Seq(100, 101),
-      1 -> Seq(101, 102),
-      2 -> Seq(102, 100)
-    ), servers = servers)
-
-    val proposed: Map[TopicAndPartition, Seq[Int]] = Map(
-      TopicAndPartition("topic1", 0) -> Seq(100, 102),
-      TopicAndPartition("topic1", 2) -> Seq(100, 102),
-      TopicAndPartition("topic2", 1) -> Seq(101, 100),
-      TopicAndPartition("topic2", 2) -> Seq(100, 102)
-    )
-
-    val replica1 = new TopicPartitionReplica("topic1", 0, 102)
-    val replica2 = new TopicPartitionReplica("topic2", 1, 100)
-    val proposedReplicaAssignment: Map[TopicPartitionReplica, String] = Map(
-      replica1 -> getRandomLogDirAssignment(102),
-      replica2 -> getRandomLogDirAssignment(100)
-    )
-
-    //When rebalancing
-    ReassignPartitionsCommand.executeAssignment(zkUtils, Some(adminClient),
-      ReassignPartitionsCommand.formatAsReassignmentJson(proposed, proposedReplicaAssignment), NoThrottle)
-    waitForReassignmentToComplete()
-
-    //Then the proposed changes should have been made
-    val actual = zkUtils.getPartitionAssignmentForTopics(Seq("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
-
-    // The replicas should be in the expected log directories
-    val replicaDirs = adminClient.describeReplicaLogDirs(List(replica1, replica2).asJavaCollection).all().get()
-    assertEquals(proposedReplicaAssignment(replica1), replicaDirs.get(replica1).getCurrentReplicaLogDir)
-    assertEquals(proposedReplicaAssignment(replica2), replicaDirs.get(replica2).getCurrentReplicaLogDir)
-  }
-
-  @Test
-  def shouldExecuteThrottledReassignment() {
-
-    //Given partitions on 3 of 3 brokers
-    val brokers = Array(100, 101, 102)
-    startBrokers(brokers)
-    createTopic(zkUtils, topicName, Map(
-      0 -> Seq(100, 101)
-    ), servers = servers)
-
-    //Given throttle set so replication will take a certain number of secs
-    val initialThrottle = Throttle(10 * 1000 * 1000, -1, () => zkUpdateDelay)
-    val expectedDurationSecs = 5
-    val numMessages: Int = 500
-    val msgSize: Int = 100 * 1000
-    produceMessages(servers, topicName, numMessages, acks = 0, msgSize)
-    assertEquals(expectedDurationSecs, numMessages * msgSize / initialThrottle.interBrokerLimit)
-
-    //Start rebalance which will move replica on 100 -> replica on 102
-    val newAssignment = generateAssignment(zkUtils, Array(101, 102), json(topicName), true)._1
-
-    val start = System.currentTimeMillis()
-    ReassignPartitionsCommand.executeAssignment(zkUtils, None,
-      ReassignPartitionsCommand.formatAsReassignmentJson(newAssignment, Map.empty), initialThrottle)
-
-    //Check throttle config. Should be throttling replica 0 on 100 and 102 only.
-    checkThrottleConfigAddedToZK(initialThrottle.interBrokerLimit, servers, topicName, "0:100,0:101", "0:102")
-
-    //Await completion
-    waitForReassignmentToComplete()
-    val took = System.currentTimeMillis() - start - delayMs
-
-    //Check move occurred
-    val actual = zkUtils.getPartitionAssignmentForTopics(Seq(topicName))(topicName)
-    assertEquals(Seq(101, 102), actual.values.flatten.toSeq.distinct.sorted)
-
-    //Then command should have taken longer than the throttle rate
-    assertTrue(s"Expected replication to be > ${expectedDurationSecs * 0.9 * 1000} but was $took",
-      took > expectedDurationSecs * 0.9 * 1000)
-    assertTrue(s"Expected replication to be < ${expectedDurationSecs * 2 * 1000} but was $took",
-      took < expectedDurationSecs * 2 * 1000)
-  }
-
-
-  @Test
-  def shouldOnlyThrottleMovingReplicas() {
-    //Given 6 brokers, two topics
-    val brokers = Array(100, 101, 102, 103, 104, 105)
-    startBrokers(brokers)
-    createTopic(zkUtils, "topic1", Map(
-      0 -> Seq(100, 101),
-      1 -> Seq(100, 101),
-      2 -> Seq(103, 104) //will leave in place
-    ), servers = servers)
-
-    createTopic(zkUtils, "topic2", Map(
-      0 -> Seq(104, 105),
-      1 -> Seq(104, 105),
-      2 -> Seq(103, 104)//will leave in place
-    ), servers = servers)
-
-    //Given throttle set so replication will take a while
-    val throttle: Long = 1000 * 1000
-    produceMessages(servers, "topic1", 100, acks = 0, 100 * 1000)
-    produceMessages(servers, "topic2", 100, acks = 0, 100 * 1000)
-
-    //Start rebalance
-    val newAssignment = Map(
-      TopicAndPartition("topic1", 0) -> Seq(100, 102),//moved 101=>102
-      TopicAndPartition("topic1", 1) -> Seq(100, 102),//moved 101=>102
-      TopicAndPartition("topic2", 0) -> Seq(103, 105),//moved 104=>103
-      TopicAndPartition("topic2", 1) -> Seq(103, 105),//moved 104=>103
-      TopicAndPartition("topic1", 2) -> Seq(103, 104), //didn't move
-      TopicAndPartition("topic2", 2) -> Seq(103, 104)  //didn't move
-    )
-    ReassignPartitionsCommand.executeAssignment(zkUtils, None,
-      ReassignPartitionsCommand.formatAsReassignmentJson(newAssignment, Map.empty), Throttle(throttle))
-
-    //Check throttle config. Should be throttling specific replicas for each topic.
-    checkThrottleConfigAddedToZK(throttle, servers, "topic1",
-      "1:100,1:101,0:100,0:101", //All replicas for moving partitions should be leader-throttled
-      "1:102,0:102" //Move destinations should be follower throttled.
-    )
-    checkThrottleConfigAddedToZK(throttle, servers, "topic2",
-      "1:104,1:105,0:104,0:105", //All replicas for moving partitions should be leader-throttled
-      "1:103,0:103" //Move destinations should be follower throttled.
-    )
-  }
-
-  @Test
-  def shouldChangeThrottleOnRerunAndRemoveOnVerify() {
-    //Given partitions on 3 of 3 brokers
-    val brokers = Array(100, 101, 102)
-    startBrokers(brokers)
-    createTopic(zkUtils, topicName, Map(
-      0 -> Seq(100, 101)
-    ), servers = servers)
-
-    //Given throttle set so replication will take at least 20 sec (we won't wait this long)
-    val initialThrottle: Long = 1000 * 1000
-    produceMessages(servers, topicName, numMessages = 200, acks = 0, valueBytes = 100 * 1000)
-
-    //Start rebalance
-    val newAssignment = generateAssignment(zkUtils, Array(101, 102), json(topicName), true)._1
-
-    ReassignPartitionsCommand.executeAssignment(zkUtils, None,
-      ReassignPartitionsCommand.formatAsReassignmentJson(newAssignment, Map.empty), Throttle(initialThrottle))
-
-    //Check throttle config
-    checkThrottleConfigAddedToZK(initialThrottle, servers, topicName, "0:100,0:101", "0:102")
-
-    //Ensure that running Verify, whilst the command is executing, should have no effect
-    verifyAssignment(zkUtils, None, ReassignPartitionsCommand.formatAsReassignmentJson(newAssignment, Map.empty))
-
-    //Check throttle config again
-    checkThrottleConfigAddedToZK(initialThrottle, servers, topicName, "0:100,0:101", "0:102")
-
-    //Now re-run the same assignment with a larger throttle, which should only act to increase the throttle and make progress
-    val newThrottle = initialThrottle * 1000
-
-    ReassignPartitionsCommand.executeAssignment(zkUtils, None,
-      ReassignPartitionsCommand.formatAsReassignmentJson(newAssignment, Map.empty), Throttle(newThrottle))
-
-    //Check throttle was changed
-    checkThrottleConfigAddedToZK(newThrottle, servers, topicName, "0:100,0:101", "0:102")
-
-    //Await completion
-    waitForReassignmentToComplete()
-
-    //Verify should remove the throttle
-    verifyAssignment(zkUtils, None, ReassignPartitionsCommand.formatAsReassignmentJson(newAssignment, Map.empty))
-
-    //Check removed
-    checkThrottleConfigRemovedFromZK(topicName, servers)
-
-    //Check move occurred
-    val actual = zkUtils.getPartitionAssignmentForTopics(Seq(topicName))(topicName)
-    assertEquals(Seq(101, 102), actual.values.flatten.toSeq.distinct.sorted)
-  }
-
-  @Test(expected = classOf[AdminCommandFailedException])
-  def shouldFailIfProposedDoesNotMatchExisting() {
-    //Given a single replica on server 100
-    startBrokers(Seq(100, 101))
-    createTopic(zkUtils, topicName, Map(0 -> 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]}]}"""
-    ReassignPartitionsCommand.executeAssignment(zkUtils, None, topicJson, NoThrottle)
-  }
-
-  @Test(expected = classOf[AdminCommandFailedException])
-  def shouldFailIfProposedHasEmptyReplicaList() {
-    //Given a single replica on server 100
-    startBrokers(Seq(100, 101))
-    createTopic(zkUtils, topicName, Map(0 -> 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":[]}]}"""
-    ReassignPartitionsCommand.executeAssignment(zkUtils, None, topicJson, NoThrottle)
-  }
-
-  @Test(expected = classOf[AdminCommandFailedException])
-  def shouldFailIfProposedHasInvalidBrokerID() {
-    //Given a single replica on server 100
-    startBrokers(Seq(100, 101))
-    createTopic(zkUtils, topicName, Map(0 -> 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]}]}"""
-    ReassignPartitionsCommand.executeAssignment(zkUtils, None, topicJson, NoThrottle)
-  }
-
-  @Test(expected = classOf[AdminCommandFailedException])
-  def shouldFailIfProposedHasInvalidLogDir() {
-    // Given a single replica on server 100
-    startBrokers(Seq(100, 101))
-    adminClient = createAdminClient(servers)
-    createTopic(zkUtils, topicName, Map(0 -> 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"]}]}"""
-    ReassignPartitionsCommand.executeAssignment(zkUtils, Some(adminClient), topicJson, NoThrottle)
-  }
-
-  @Test(expected = classOf[AdminCommandFailedException])
-  def shouldFailIfProposedHasInconsistentReplicasAndLogDirs() {
-    // Given a single replica on server 100
-    startBrokers(Seq(100, 101))
-    adminClient = createAdminClient(servers)
-    val logDir = getRandomLogDirAssignment(100)
-    createTopic(zkUtils, topicName, Map(0 -> 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"]}]}"""
-    ReassignPartitionsCommand.executeAssignment(zkUtils, Some(adminClient), topicJson, NoThrottle)
-  }
-
-  @Test
-  def shouldPerformThrottledReassignmentOverVariousTopics() {
-    val throttle = Throttle(1000L)
-
-    //Given four brokers
-    servers = TestUtils.createBrokerConfigs(4, zkConnect, false).map(conf => TestUtils.createServer(KafkaConfig.fromProps(conf)))
-
-    //With up several small topics
-    createTopic(zkUtils, "orders", Map(0 -> List(0, 1, 2), 1 -> List(0, 1, 2)), servers)
-    createTopic(zkUtils, "payments", Map(0 -> List(0, 1), 1 -> List(0, 1)), servers)
-    createTopic(zkUtils, "deliveries", Map(0 -> List(0)), servers)
-    createTopic(zkUtils, "customers", Map(0 -> List(0), 1 -> List(1), 2 -> List(2), 3 -> List(3)), servers)
-
-    //Define a move for some of them
-    val move = Map(
-      TopicAndPartition("orders", 0) -> Seq(0, 2, 3),//moves
-      TopicAndPartition("orders", 1) -> Seq(0, 1, 2),//stays
-      TopicAndPartition("payments", 1) -> Seq(1, 2), //only define one partition as moving
-      TopicAndPartition("deliveries", 0) -> Seq(1, 2) //increase replication factor
-    )
-
-    //When we run a throttled reassignment
-    new ReassignPartitionsCommand(zkUtils, None, move).reassignPartitions(throttle)
-
-    waitForReassignmentToComplete()
-
-    //Check moved replicas did move
-    assertEquals(Seq(0, 2, 3), zkUtils.getReplicasForPartition("orders", 0))
-    assertEquals(Seq(0, 1, 2), zkUtils.getReplicasForPartition("orders", 1))
-    assertEquals(Seq(1, 2), zkUtils.getReplicasForPartition("payments", 1))
-    assertEquals(Seq(1, 2), zkUtils.getReplicasForPartition("deliveries", 0))
-
-    //Check untouched replicas are still there
-    assertEquals(Seq(0, 1), zkUtils.getReplicasForPartition("payments", 0))
-    assertEquals(Seq(0), zkUtils.getReplicasForPartition("customers", 0))
-    assertEquals(Seq(1), zkUtils.getReplicasForPartition("customers", 1))
-    assertEquals(Seq(2), zkUtils.getReplicasForPartition("customers", 2))
-    assertEquals(Seq(3), zkUtils.getReplicasForPartition("customers", 3))
-  }
-
-  /**
-   * Verifies that the Controller sets a watcher for the reassignment znode after reassignment completion.
-   * This includes the case where the znode is set immediately after it's deleted (i.e. before the watch is set).
-   * This case relies on the scheduling of the operations, so it won't necessarily fail every time, but it fails
-   * often enough to detect a regression.
-   */
-  @Test
-  def shouldPerformMultipleReassignmentOperationsOverVariousTopics() {
-    servers = TestUtils.createBrokerConfigs(4, zkConnect, false).map(conf => TestUtils.createServer(KafkaConfig.fromProps(conf)))
-
-    createTopic(zkUtils, "orders", Map(0 -> List(0, 1, 2), 1 -> List(0, 1, 2)), servers)
-    createTopic(zkUtils, "payments", Map(0 -> List(0, 1), 1 -> List(0, 1)), servers)
-    createTopic(zkUtils, "deliveries", Map(0 -> List(0)), servers)
-    createTopic(zkUtils, "customers", Map(0 -> List(0), 1 -> List(1), 2 -> List(2), 3 -> List(3)), servers)
-
-    val firstMove = Map(
-      TopicAndPartition("orders", 0) -> Seq(0, 2, 3), //moves
-      TopicAndPartition("orders", 1) -> Seq(0, 1, 2), //stays
-      TopicAndPartition("payments", 1) -> Seq(1, 2), //only define one partition as moving
-      TopicAndPartition("deliveries", 0) -> Seq(1, 2) //increase replication factor
-    )
-
-    new ReassignPartitionsCommand(zkUtils, None, firstMove).reassignPartitions()
-    waitForReassignmentToComplete()
-
-    // Check moved replicas did move
-    assertEquals(Seq(0, 2, 3), zkUtils.getReplicasForPartition("orders", 0))
-    assertEquals(Seq(0, 1, 2), zkUtils.getReplicasForPartition("orders", 1))
-    assertEquals(Seq(1, 2), zkUtils.getReplicasForPartition("payments", 1))
-    assertEquals(Seq(1, 2), zkUtils.getReplicasForPartition("deliveries", 0))
-
-    // Check untouched replicas are still there
-    assertEquals(Seq(0, 1), zkUtils.getReplicasForPartition("payments", 0))
-    assertEquals(Seq(0), zkUtils.getReplicasForPartition("customers", 0))
-    assertEquals(Seq(1), zkUtils.getReplicasForPartition("customers", 1))
-    assertEquals(Seq(2), zkUtils.getReplicasForPartition("customers", 2))
-    assertEquals(Seq(3), zkUtils.getReplicasForPartition("customers", 3))
-
-    // Define a move for some of them
-    val secondMove = Map(
-      TopicAndPartition("orders", 0) -> Seq(0, 2, 3), // stays
-      TopicAndPartition("orders", 1) -> Seq(3, 1, 2), // moves
-      TopicAndPartition("payments", 1) -> Seq(2, 1), // changed preferred leader
-      TopicAndPartition("deliveries", 0) -> Seq(1, 2, 3) //increase replication factor
-    )
-
-    new ReassignPartitionsCommand(zkUtils, None, secondMove).reassignPartitions()
-    waitForReassignmentToComplete()
-
-    // Check moved replicas did move
-    assertEquals(Seq(0, 2, 3), zkUtils.getReplicasForPartition("orders", 0))
-    assertEquals(Seq(3, 1, 2), zkUtils.getReplicasForPartition("orders", 1))
-    assertEquals(Seq(2, 1), zkUtils.getReplicasForPartition("payments", 1))
-    assertEquals(Seq(1, 2, 3), zkUtils.getReplicasForPartition("deliveries", 0))
-
-    //Check untouched replicas are still there
-    assertEquals(Seq(0, 1), zkUtils.getReplicasForPartition("payments", 0))
-    assertEquals(Seq(0), zkUtils.getReplicasForPartition("customers", 0))
-    assertEquals(Seq(1), zkUtils.getReplicasForPartition("customers", 1))
-    assertEquals(Seq(2), zkUtils.getReplicasForPartition("customers", 2))
-    assertEquals(Seq(3), zkUtils.getReplicasForPartition("customers", 3))
-
-    // We set the znode and then continuously attempt to set it again to exercise the case where the znode is set
-    // immediately after deletion (i.e. before we set the watcher again)
-
-    val thirdMove = Map(TopicAndPartition("orders", 0) -> Seq(1, 2, 3))
-
-    new ReassignPartitionsCommand(zkUtils, None, thirdMove).reassignPartitions()
-
-    val fourthMove = Map(TopicAndPartition("payments", 1) -> Seq(2, 3))
-
-    // Continuously attempt to set the reassignment znode with `fourthMove` until it succeeds. It will only succeed
-    // after `thirdMove` completes.
-    Iterator.continually {
-      try new ReassignPartitionsCommand(zkUtils, None, fourthMove).reassignPartitions()
-      catch {
-        case _: AdminCommandFailedException => false
-      }
-    }.exists(identity)
-
-    waitForReassignmentToComplete()
-
-    // Check moved replicas for thirdMove and fourthMove
-    assertEquals(Seq(1, 2, 3), zkUtils.getReplicasForPartition("orders", 0))
-    assertEquals(Seq(2, 3), zkUtils.getReplicasForPartition("payments", 1))
-
-    //Check untouched replicas are still there
-    assertEquals(Seq(3, 1, 2), zkUtils.getReplicasForPartition("orders", 1))
-    assertEquals(Seq(1, 2, 3), zkUtils.getReplicasForPartition("deliveries", 0))
-    assertEquals(Seq(0, 1), zkUtils.getReplicasForPartition("payments", 0))
-    assertEquals(Seq(0), zkUtils.getReplicasForPartition("customers", 0))
-    assertEquals(Seq(1), zkUtils.getReplicasForPartition("customers", 1))
-    assertEquals(Seq(2), zkUtils.getReplicasForPartition("customers", 2))
-    assertEquals(Seq(3), zkUtils.getReplicasForPartition("customers", 3))
-  }
-
-  def waitForReassignmentToComplete() {
-    waitUntilTrue(() => !zkUtils.pathExists(ReassignPartitionsPath), s"Znode ${ZkUtils.ReassignPartitionsPath} wasn't deleted")
-  }
-
-  def json(topic: String*): String = {
-    val topicStr = topic.map { t => "{\"topic\": \"" + t + "\"}" }.mkString(",")
-    s"""{"topics": [$topicStr],"version":1}"""
-  }
-}
diff --git a/core/src/test/scala/unit/kafka/admin/ReassignPartitionsCommandArgsTest.scala b/core/src/test/scala/unit/kafka/admin/ReassignPartitionsCommandArgsTest.scala
index a85acf918622d..adeec20bb0022 100644
--- a/core/src/test/scala/unit/kafka/admin/ReassignPartitionsCommandArgsTest.scala
+++ b/core/src/test/scala/unit/kafka/admin/ReassignPartitionsCommandArgsTest.scala
@@ -18,29 +18,28 @@ package kafka.admin
 
 import kafka.utils.Exit
 import org.junit.Assert._
-import org.junit.{After, Before, Test}
-import org.scalatest.junit.JUnitSuite
+import org.junit.rules.Timeout
+import org.junit.{After, Before, Rule, Test}
 
-class ReassignPartitionsCommandArgsTest extends JUnitSuite {
+class ReassignPartitionsCommandArgsTest {
+  @Rule
+  def globalTimeout: Timeout = Timeout.millis(60000)
 
   @Before
-  def setUp() {
+  def setUp(): Unit = {
     Exit.setExitProcedure((_, message) => throw new IllegalArgumentException(message.orNull))
   }
 
   @After
-  def tearDown() {
+  def tearDown(): Unit = {
     Exit.resetExitProcedure()
   }
 
-  /**
-    * HAPPY PATH
-    */
-
+  ///// Test valid argument parsing
   @Test
   def shouldCorrectlyParseValidMinimumGenerateOptions(): Unit = {
     val args = Array(
-      "--zookeeper", "localhost:1234",
+      "--bootstrap-server", "localhost:1234",
       "--generate",
       "--broker-list", "101,102",
       "--topics-to-move-json-file", "myfile.json")
@@ -49,6 +48,15 @@ class ReassignPartitionsCommandArgsTest extends JUnitSuite {
 
   @Test
   def shouldCorrectlyParseValidMinimumExecuteOptions(): Unit = {
+    val args = Array(
+      "--bootstrap-server", "localhost:1234",
+      "--execute",
+      "--reassignment-json-file", "myfile.json")
+    ReassignPartitionsCommand.validateAndParseArgs(args)
+  }
+
+  @Test
+  def shouldCorrectlyParseValidMinimumLegacyExecuteOptions(): Unit = {
     val args = Array(
       "--zookeeper", "localhost:1234",
       "--execute",
@@ -58,6 +66,15 @@ class ReassignPartitionsCommandArgsTest extends JUnitSuite {
 
   @Test
   def shouldCorrectlyParseValidMinimumVerifyOptions(): Unit = {
+    val args = Array(
+      "--bootstrap-server", "localhost:1234",
+      "--verify",
+      "--reassignment-json-file", "myfile.json")
+    ReassignPartitionsCommand.validateAndParseArgs(args)
+  }
+
+  @Test
+  def shouldCorrectlyParseValidMinimumLegacyVerifyOptions(): Unit = {
     val args = Array(
       "--zookeeper", "localhost:1234",
       "--verify",
@@ -68,7 +85,7 @@ class ReassignPartitionsCommandArgsTest extends JUnitSuite {
   @Test
   def shouldAllowThrottleOptionOnExecute(): Unit = {
     val args = Array(
-      "--zookeeper", "localhost:1234",
+      "--bootstrap-server", "localhost:1234",
       "--execute",
       "--throttle", "100",
       "--reassignment-json-file", "myfile.json")
@@ -78,7 +95,7 @@ class ReassignPartitionsCommandArgsTest extends JUnitSuite {
   @Test
   def shouldUseDefaultsIfEnabled(): Unit = {
     val args = Array(
-      "--zookeeper", "localhost:1234",
+      "--bootstrap-server", "localhost:1234",
       "--execute",
       "--reassignment-json-file", "myfile.json")
     val opts = ReassignPartitionsCommand.validateAndParseArgs(args)
@@ -86,147 +103,183 @@ class ReassignPartitionsCommandArgsTest extends JUnitSuite {
     assertEquals(-1L, opts.options.valueOf(opts.interBrokerThrottleOpt))
   }
 
-  /**
-    * NO ARGS
-    */
+  @Test
+  def testList(): Unit = {
+    val args = Array(
+      "--list",
+      "--bootstrap-server", "localhost:1234")
+    ReassignPartitionsCommand.validateAndParseArgs(args)
+  }
 
+  @Test
+  def testCancelWithPreserveThrottlesOption(): Unit = {
+    val args = Array(
+      "--cancel",
+      "--bootstrap-server", "localhost:1234",
+      "--reassignment-json-file", "myfile.json",
+      "--preserve-throttles")
+    ReassignPartitionsCommand.validateAndParseArgs(args)
+  }
+
+  ///// Test handling missing or invalid actions
   @Test
   def shouldFailIfNoArgs(): Unit = {
     val args: Array[String]= Array()
-    shouldFailWith("This command moves topic partitions between replicas.", args)
+    shouldFailWith(ReassignPartitionsCommand.helpText, args)
   }
 
   @Test
   def shouldFailIfBlankArg(): Unit = {
     val args = Array(" ")
-    shouldFailWith("Command must include exactly one action: --generate, --execute or --verify", args)
+    shouldFailWith("Command must include exactly one action", args)
   }
 
-  /**
-    * UNHAPPY PATH: EXECUTE ACTION
-    */
+  @Test
+  def shouldFailIfMultipleActions(): Unit = {
+    val args = Array(
+      "--bootstrap-server", "localhost:1234",
+      "--execute",
+      "--verify",
+      "--reassignment-json-file", "myfile.json"
+    )
+    shouldFailWith("Command must include exactly one action", args)
+  }
 
+  ///// Test --execute
   @Test
   def shouldNotAllowExecuteWithTopicsOption(): Unit = {
     val args = Array(
-      "--zookeeper", "localhost:1234",
+      "--bootstrap-server", "localhost:1234",
       "--execute",
       "--reassignment-json-file", "myfile.json",
       "--topics-to-move-json-file", "myfile.json")
-    shouldFailWith("Option \"[execute]\" can't be used with option\"[topics-to-move-json-file]\"", args)
+    shouldFailWith("Option \"[topics-to-move-json-file]\" can't be used with action \"[execute]\"", args)
   }
 
   @Test
-  def shouldNotAllowExecuteWithBrokers(): Unit = {
+  def shouldNotAllowExecuteWithBrokerList(): Unit = {
     val args = Array(
-      "--zookeeper", "localhost:1234",
+      "--bootstrap-server", "localhost:1234",
       "--execute",
       "--reassignment-json-file", "myfile.json",
       "--broker-list", "101,102"
     )
-    shouldFailWith("Option \"[execute]\" can't be used with option\"[broker-list]\"", args)
+    shouldFailWith("Option \"[broker-list]\" can't be used with action \"[execute]\"", args)
   }
 
   @Test
   def shouldNotAllowExecuteWithoutReassignmentOption(): Unit = {
     val args = Array(
-      "--zookeeper", "localhost:1234",
+      "--bootstrap-server", "localhost:1234",
       "--execute")
-    shouldFailWith("If --execute option is used, command must include --reassignment-json-file that was output during the --generate option", args)
+    shouldFailWith("Missing required argument \"[reassignment-json-file]\"", args)
   }
 
-  /**
-    * UNHAPPY PATH: GENERATE ACTION
-    */
+  @Test
+  def testMissingBootstrapServerArgumentForExecute(): Unit = {
+    val args = Array(
+      "--execute")
+    shouldFailWith("Please specify --bootstrap-server", args)
+  }
 
+  ///// Test --generate
   @Test
   def shouldNotAllowGenerateWithoutBrokersAndTopicsOptions(): Unit = {
     val args = Array(
-      "--zookeeper", "localhost:1234",
+      "--bootstrap-server", "localhost:1234",
       "--generate")
-    shouldFailWith("If --generate option is used, command must include both --topics-to-move-json-file and --broker-list options", args)
+    shouldFailWith("Missing required argument \"[topics-to-move-json-file]\"", args)
   }
 
   @Test
   def shouldNotAllowGenerateWithoutBrokersOption(): Unit = {
     val args = Array(
-      "--zookeeper", "localhost:1234",
+      "--bootstrap-server", "localhost:1234",
       "--topics-to-move-json-file", "myfile.json",
       "--generate")
-    shouldFailWith("If --generate option is used, command must include both --topics-to-move-json-file and --broker-list options", args)
+    shouldFailWith("Missing required argument \"[broker-list]\"", args)
   }
 
   @Test
   def shouldNotAllowGenerateWithoutTopicsOption(): Unit = {
     val args = Array(
-      "--zookeeper", "localhost:1234",
+      "--bootstrap-server", "localhost:1234",
       "--broker-list", "101,102",
       "--generate")
-    shouldFailWith("If --generate option is used, command must include both --topics-to-move-json-file and --broker-list options", args)
+    shouldFailWith("Missing required argument \"[topics-to-move-json-file]\"", args)
   }
 
   @Test
   def shouldNotAllowGenerateWithThrottleOption(): Unit = {
     val args = Array(
-      "--zookeeper", "localhost:1234",
+      "--bootstrap-server", "localhost:1234",
       "--generate",
       "--broker-list", "101,102",
       "--throttle", "100",
       "--topics-to-move-json-file", "myfile.json")
-    shouldFailWith("Option \"[generate]\" can't be used with option\"[throttle]\"", args)
+    shouldFailWith("Option \"[throttle]\" can't be used with action \"[generate]\"", args)
   }
 
   @Test
   def shouldNotAllowGenerateWithReassignmentOption(): Unit = {
     val args = Array(
-      "--zookeeper", "localhost:1234",
+      "--bootstrap-server", "localhost:1234",
       "--generate",
       "--broker-list", "101,102",
       "--topics-to-move-json-file", "myfile.json",
       "--reassignment-json-file", "myfile.json")
-    shouldFailWith("Option \"[generate]\" can't be used with option\"[reassignment-json-file]\"", args)
+    shouldFailWith("Option \"[reassignment-json-file]\" can't be used with action \"[generate]\"", args)
   }
 
-  /**
-    * UNHAPPY PATH: VERIFY ACTION
-    */
+  @Test
+  def testInvalidCommandConfigArgumentForLegacyGenerate(): Unit = {
+    val args = Array(
+      "--zookeeper", "localhost:1234",
+      "--generate",
+      "--broker-list", "101,102",
+      "--topics-to-move-json-file", "myfile.json",
+      "--command-config", "/tmp/command-config.properties"
+    )
+    shouldFailWith("You must specify --bootstrap-server when using \"[command-config]\"", args)
+  }
 
+  ///// Test --verify
   @Test
   def shouldNotAllowVerifyWithoutReassignmentOption(): Unit = {
     val args = Array(
-      "--zookeeper", "localhost:1234",
+      "--bootstrap-server", "localhost:1234",
       "--verify")
-    shouldFailWith("If --verify option is used, command must include --reassignment-json-file that was used during the --execute option", args)
+    shouldFailWith("Missing required argument \"[reassignment-json-file]\"", args)
   }
 
   @Test
   def shouldNotAllowBrokersListWithVerifyOption(): Unit = {
     val args = Array(
-      "--zookeeper", "localhost:1234",
+      "--bootstrap-server", "localhost:1234",
       "--verify",
       "--broker-list", "100,101",
       "--reassignment-json-file", "myfile.json")
-    shouldFailWith("Option \"[verify]\" can't be used with option\"[broker-list]\"", args)
+    shouldFailWith("Option \"[broker-list]\" can't be used with action \"[verify]\"", args)
   }
 
   @Test
   def shouldNotAllowThrottleWithVerifyOption(): Unit = {
     val args = Array(
-      "--zookeeper", "localhost:1234",
+      "--bootstrap-server", "localhost:1234",
       "--verify",
       "--throttle", "100",
       "--reassignment-json-file", "myfile.json")
-    shouldFailWith("Option \"[verify]\" can't be used with option\"[throttle]\"", args)
+    shouldFailWith("Option \"[throttle]\" can't be used with action \"[verify]\"", args)
   }
 
   @Test
   def shouldNotAllowTopicsOptionWithVerify(): Unit = {
     val args = Array(
-      "--zookeeper", "localhost:1234",
+      "--bootstrap-server", "localhost:1234",
       "--verify",
       "--reassignment-json-file", "myfile.json",
       "--topics-to-move-json-file", "myfile.json")
-    shouldFailWith("Option \"[verify]\" can't be used with option\"[topics-to-move-json-file]\"", args)
+    shouldFailWith("Option \"[topics-to-move-json-file]\" can't be used with action \"[verify]\"", args)
   }
 
   def shouldFailWith(msg: String, args: Array[String]): Unit = {
@@ -237,4 +290,30 @@ class ReassignPartitionsCommandArgsTest extends JUnitSuite {
       case e: Exception => assertTrue(s"Expected exception with message:\n[$msg]\nbut was\n[${e.getMessage}]", e.getMessage.startsWith(msg))
     }
   }
+
+  ///// Test --cancel
+  @Test
+  def shouldNotAllowCancelWithoutBootstrapServerOption(): Unit = {
+    val args = Array(
+      "--cancel")
+    shouldFailWith("Please specify --bootstrap-server", args)
+  }
+
+  @Test
+  def shouldNotAllowCancelWithoutReassignmentJsonFile(): Unit = {
+    val args = Array(
+      "--cancel",
+      "--bootstrap-server", "localhost:1234",
+      "--preserve-throttles")
+    shouldFailWith("Missing required argument \"[reassignment-json-file]\"", args)
+  }
+
+  ///// Test --list
+  @Test
+  def shouldNotAllowZooKeeperWithListOption(): Unit = {
+    val args = Array(
+      "--list",
+      "--zookeeper", "localhost:1234")
+    shouldFailWith("Option \"[zookeeper]\" can't be used with action \"[list]\"", args)
+  }
 }
diff --git a/core/src/test/scala/unit/kafka/admin/ReassignPartitionsCommandTest.scala b/core/src/test/scala/unit/kafka/admin/ReassignPartitionsCommandTest.scala
deleted file mode 100644
index 6c5a2ad45c08c..0000000000000
--- a/core/src/test/scala/unit/kafka/admin/ReassignPartitionsCommandTest.scala
+++ /dev/null
@@ -1,415 +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.admin.ReassignPartitionsCommand.Throttle
-import kafka.common.TopicAndPartition
-import kafka.log.LogConfig
-import kafka.log.LogConfig._
-import kafka.server.{ConfigType, DynamicConfig}
-import kafka.utils.CoreUtils._
-import kafka.utils.TestUtils._
-import kafka.utils.{CoreUtils, Logging, TestUtils, ZkUtils}
-import org.easymock.EasyMock._
-import org.easymock.{Capture, CaptureType, EasyMock}
-import org.junit.{Before, Test}
-import org.junit.Assert.{assertEquals, assertNull, fail}
-
-import scala.collection.{Seq, mutable}
-import scala.collection.JavaConversions._
-
-class ReassignPartitionsCommandTest extends Logging {
-  var calls = 0
-
-  @Test
-  def shouldFindMovingReplicas() {
-    val control = TopicAndPartition("topic1", 1) -> Seq(100, 102)
-    val assigner = new ReassignPartitionsCommand(null, null, null, null)
-
-    //Given partition 0 moves from broker 100 -> 102. Partition 1 does not move.
-    val existing = Map(TopicAndPartition("topic1", 0) -> Seq(100, 101), control)
-    val proposed = Map(TopicAndPartition("topic1", 0) -> Seq(101, 102), control)
-
-
-    val mock = new TestAdminUtils {
-      override def changeTopicConfig(zkUtils: ZkUtils, topic: String, configChange: Properties): Unit = {
-        assertEquals("0:102", configChange.get(FollowerReplicationThrottledReplicasProp)) //Should only be follower-throttle the moving replica
-        assertEquals("0:100,0:101", configChange.get(LeaderReplicationThrottledReplicasProp)) //Should leader-throttle all existing (pre move) replicas
-        calls += 1
-      }
-    }
-
-    assigner.assignThrottledReplicas(existing, proposed, mock)
-    assertEquals(1, calls)
-  }
-
-  @Test
-  def shouldFindMovingReplicasWhenProposedIsSubsetOfExisting() {
-    val assigner = new ReassignPartitionsCommand(null, null, null, null)
-
-    //Given we have more existing partitions than we are proposing
-    val existingSuperset = Map(
-      TopicAndPartition("topic1", 0) -> Seq(100, 101),
-      TopicAndPartition("topic1", 1) -> Seq(100, 102),
-      TopicAndPartition("topic1", 2) -> Seq(100, 101),
-      TopicAndPartition("topic2", 0) -> Seq(100, 101, 102),
-      TopicAndPartition("topic3", 0) -> Seq(100, 101, 102)
-    )
-    val proposedSubset = Map(
-      TopicAndPartition("topic1", 0) -> Seq(101, 102),
-      TopicAndPartition("topic1", 1) -> Seq(102),
-      TopicAndPartition("topic1", 2) -> Seq(100, 101, 102)
-    )
-
-    val mock = new TestAdminUtils {
-      override def changeTopicConfig(zkUtils: ZkUtils, topic: String, configChange: Properties): Unit = {
-        assertEquals("0:102,2:102", configChange.get(FollowerReplicationThrottledReplicasProp))
-        assertEquals("0:100,0:101,2:100,2:101", configChange.get(LeaderReplicationThrottledReplicasProp))
-        assertEquals("topic1", topic)
-        calls += 1
-      }
-    }
-
-    //Then replicas should assign correctly (based on the proposed map)
-    assigner.assignThrottledReplicas(existingSuperset, proposedSubset, mock)
-    assertEquals(1, calls)
-  }
-
-  @Test
-  def shouldFindMovingReplicasMultiplePartitions() {
-    val control = TopicAndPartition("topic1", 2) -> Seq(100, 102)
-    val assigner = new ReassignPartitionsCommand(null, null, null, null)
-
-    //Given partitions 0 & 1 moves from broker 100 -> 102. Partition 2 does not move.
-    val existing = Map(TopicAndPartition("topic1", 0) -> Seq(100, 101), TopicAndPartition("topic1", 1) -> Seq(100, 101), control)
-    val proposed = Map(TopicAndPartition("topic1", 0) -> Seq(101, 102), TopicAndPartition("topic1", 1) -> Seq(101, 102), control)
-
-    // Then
-    val mock = new TestAdminUtils {
-      override def changeTopicConfig(zkUtils: ZkUtils, topic: String, configChange: Properties): Unit = {
-        assertEquals("0:102,1:102", configChange.get(FollowerReplicationThrottledReplicasProp)) //Should only be follower-throttle the moving replica
-        assertEquals("0:100,0:101,1:100,1:101", configChange.get(LeaderReplicationThrottledReplicasProp)) //Should leader-throttle all existing (pre move) replicas
-        calls += 1
-      }
-    }
-
-    //When
-    assigner.assignThrottledReplicas(existing, proposed, mock)
-    assertEquals(1, calls)
-  }
-
-  @Test
-  def shouldFindMovingReplicasMultipleTopics() {
-    val control = TopicAndPartition("topic1", 1) -> Seq(100, 102)
-    val assigner = new ReassignPartitionsCommand(null, null, null, null)
-
-    //Given topics 1 -> move from broker 100 -> 102, topics 2 -> move from broker 101 -> 100
-    val existing = Map(TopicAndPartition("topic1", 0) -> Seq(100, 101), TopicAndPartition("topic2", 0) -> Seq(101, 102), control)
-    val proposed = Map(TopicAndPartition("topic1", 0) -> Seq(101, 102), TopicAndPartition("topic2", 0) -> Seq(100, 102), control)
-
-    //Then
-    val mock = new TestAdminUtils {
-      override def changeTopicConfig(zkUtils: ZkUtils, topic: String, configChange: Properties): Unit = {
-        topic match {
-          case "topic1" =>
-            assertEquals("0:100,0:101", configChange.get(LeaderReplicationThrottledReplicasProp))
-            assertEquals("0:102", configChange.get(FollowerReplicationThrottledReplicasProp))
-          case "topic2" =>
-            assertEquals("0:101,0:102", configChange.get(LeaderReplicationThrottledReplicasProp))
-            assertEquals("0:100", configChange.get(FollowerReplicationThrottledReplicasProp))
-          case _ => fail(s"Unexpected topic $topic")
-        }
-        calls += 1
-      }
-    }
-
-    //When
-    assigner.assignThrottledReplicas(existing, proposed, mock)
-    assertEquals(2, calls)
-  }
-
-  @Test
-  def shouldFindMovingReplicasMultipleTopicsAndPartitions() {
-    val assigner = new ReassignPartitionsCommand(null, null, null, null)
-
-    //Given
-    val existing = Map(
-      TopicAndPartition("topic1", 0) -> Seq(100, 101),
-      TopicAndPartition("topic1", 1) -> Seq(100, 101),
-      TopicAndPartition("topic2", 0) -> Seq(101, 102),
-      TopicAndPartition("topic2", 1) -> Seq(101, 102)
-    )
-    val proposed = Map(
-      TopicAndPartition("topic1", 0) -> Seq(101, 102), //moves to 102
-      TopicAndPartition("topic1", 1) -> Seq(101, 102), //moves to 102
-      TopicAndPartition("topic2", 0) -> Seq(100, 102), //moves to 100
-      TopicAndPartition("topic2", 1) -> Seq(101, 100)  //moves to 100
-    )
-
-    //Then
-    val mock = new TestAdminUtils {
-      override def changeTopicConfig(zkUtils: ZkUtils, topic: String, configChange: Properties): Unit = {
-        topic match {
-          case "topic1" =>
-            assertEquals("0:102,1:102", configChange.get(FollowerReplicationThrottledReplicasProp))
-            assertEquals("0:100,0:101,1:100,1:101", configChange.get(LeaderReplicationThrottledReplicasProp))
-          case "topic2" =>
-            assertEquals("0:100,1:100", configChange.get(FollowerReplicationThrottledReplicasProp))
-            assertEquals("0:101,0:102,1:101,1:102", configChange.get(LeaderReplicationThrottledReplicasProp))
-          case _ => fail(s"Unexpected topic $topic")
-        }
-        calls += 1
-      }
-    }
-
-    //When
-    assigner.assignThrottledReplicas(existing, proposed, mock)
-    assertEquals(2, calls)
-  }
-
-  @Test
-  def shouldFindTwoMovingReplicasInSamePartition() {
-    val control = TopicAndPartition("topic1", 1) -> Seq(100, 102)
-    val assigner = new ReassignPartitionsCommand(null, null, null, null)
-
-    //Given partition 0 has 2 moves from broker 102 -> 104 & 103 -> 105
-    val existing = Map(TopicAndPartition("topic1", 0) -> Seq(100, 101, 102, 103), control)
-    val proposed = Map(TopicAndPartition("topic1", 0) -> Seq(100, 101, 104, 105), control)
-
-    // Then
-    val mock = new TestAdminUtils {
-      override def changeTopicConfig(zkUtils: ZkUtils, topic: String, configChange: Properties) = {
-        assertEquals("0:104,0:105", configChange.get(FollowerReplicationThrottledReplicasProp)) //Should only be follower-throttle the moving replicas
-        assertEquals("0:100,0:101,0:102,0:103", configChange.get(LeaderReplicationThrottledReplicasProp)) //Should leader-throttle all existing (pre move) replicas
-        calls += 1
-      }
-    }
-
-    //When
-    assigner.assignThrottledReplicas(existing, proposed, mock)
-    assertEquals(1, calls)
-  }
-
-  @Test
-  def shouldNotOverwriteEntityConfigsWhenUpdatingThrottledReplicas(): Unit = {
-    val control = TopicAndPartition("topic1", 1) -> Seq(100, 102)
-    val assigner = new ReassignPartitionsCommand(null, null, null, null)
-    val existing = Map(TopicAndPartition("topic1", 0) -> Seq(100, 101), control)
-    val proposed = Map(TopicAndPartition("topic1", 0) -> Seq(101, 102), control)
-
-    //Given partition there are existing properties
-    val existingProperties = propsWith("some-key", "some-value")
-
-    //Then the dummy property should still be there
-    val mock = new TestAdminUtils {
-      override def changeTopicConfig(zkUtils: ZkUtils, topic: String, configChange: Properties): Unit = {
-        assertEquals("some-value", configChange.getProperty("some-key"))
-        calls += 1
-      }
-
-      override def fetchEntityConfig(zkUtils: ZkUtils, entityType: String, entityName: String): Properties = {
-        existingProperties
-      }
-    }
-
-    //When
-    assigner.assignThrottledReplicas(existing, proposed, mock)
-    assertEquals(1, calls)
-  }
-
-  @Test
-  def shouldSetQuotaLimit(): Unit = {
-    //Given
-    val existing = mutable.Map(TopicAndPartition("topic1", 0) -> Seq(100, 101))
-    val proposed = mutable.Map(TopicAndPartition("topic1", 0) -> Seq(101, 102))
-
-    //Setup
-    val zk = stubZK(existing)
-    val admin = createMock(classOf[AdminUtilities])
-    val propsCapture: Capture[Properties] = newCapture(CaptureType.ALL)
-    val assigner = new ReassignPartitionsCommand(zk, None, proposed, Map.empty, admin)
-    expect(admin.fetchEntityConfig(is(zk), anyString(), anyString())).andStubReturn(new Properties)
-    expect(admin.changeBrokerConfig(is(zk), anyObject().asInstanceOf[List[Int]], capture(propsCapture))).anyTimes()
-    replay(admin)
-
-    //When
-    assigner.maybeLimit(Throttle(1000))
-
-    //Then
-    for (actual <- propsCapture.getValues) {
-      assertEquals("1000", actual.getProperty(DynamicConfig.Broker.LeaderReplicationThrottledRateProp))
-      assertEquals("1000", actual.getProperty(DynamicConfig.Broker.FollowerReplicationThrottledRateProp))
-    }
-    assertEquals(3, propsCapture.getValues.size) //3 brokers
-  }
-
-  @Test
-  def shouldUpdateQuotaLimit(): Unit = {
-    //Given
-    val existing = mutable.Map(TopicAndPartition("topic1", 0) -> Seq(100, 101))
-    val proposed = mutable.Map(TopicAndPartition("topic1", 0) -> Seq(101, 102))
-
-    //Setup
-    val zk = stubZK(existing)
-    val admin = createMock(classOf[AdminUtilities])
-    val propsCapture: Capture[Properties] = newCapture(CaptureType.ALL)
-    val assigner = new ReassignPartitionsCommand(zk, None, proposed, Map.empty, admin)
-    expect(admin.changeBrokerConfig(is(zk), anyObject().asInstanceOf[List[Int]], capture(propsCapture))).anyTimes()
-
-    //Expect the existing broker config to be changed from 10/100 to 1000
-    val existingConfigs = CoreUtils.propsWith(
-      (DynamicConfig.Broker.FollowerReplicationThrottledRateProp, "10"),
-      (DynamicConfig.Broker.LeaderReplicationThrottledRateProp, "100")
-    )
-    expect(admin.fetchEntityConfig(is(zk), is(ConfigType.Broker), is("100"))).andReturn(copyOf(existingConfigs))
-    expect(admin.fetchEntityConfig(is(zk), is(ConfigType.Broker), is("101"))).andReturn(copyOf(existingConfigs))
-    expect(admin.fetchEntityConfig(is(zk), is(ConfigType.Broker), is("102"))).andReturn(copyOf(existingConfigs))
-    replay(admin)
-
-    //When
-    assigner.maybeLimit(Throttle(1000))
-
-    //Then
-    for (actual <- propsCapture.getValues) {
-      assertEquals("1000", actual.getProperty(DynamicConfig.Broker.LeaderReplicationThrottledRateProp))
-      assertEquals("1000", actual.getProperty(DynamicConfig.Broker.FollowerReplicationThrottledRateProp))
-    }
-    assertEquals(3, propsCapture.getValues.size) //three brokers
-  }
-
-  @Test
-  def shouldNotOverwriteExistingPropertiesWhenLimitIsAdded(): Unit = {
-    //Given
-    val existing = mutable.Map(TopicAndPartition("topic1", 0) -> Seq(100, 101))
-    val proposed = mutable.Map(TopicAndPartition("topic1", 0) -> Seq(101, 102))
-
-    //Setup
-    val zk = stubZK(existing)
-    val admin = createMock(classOf[AdminUtilities])
-    val propsCapture: Capture[Properties] = newCapture(CaptureType.ALL)
-    val assigner = new ReassignPartitionsCommand(zk, None, proposed, Map.empty, admin)
-    expect(admin.changeBrokerConfig(is(zk), anyObject().asInstanceOf[List[Int]], capture(propsCapture))).anyTimes()
-
-    //Given there is some existing config
-    expect(admin.fetchEntityConfig(is(zk), is(ConfigType.Broker), anyString())).andReturn(
-      propsWith("useful.key", "useful.value")).atLeastOnce()
-
-    replay(admin)
-
-    //When
-    assigner.maybeLimit(Throttle(1000))
-
-    //Then other property remains
-    for (actual <- propsCapture.getValues) {
-      assertEquals("useful.value", actual.getProperty("useful.key"))
-      assertEquals("1000", actual.getProperty(DynamicConfig.Broker.LeaderReplicationThrottledRateProp))
-      assertEquals("1000", actual.getProperty(DynamicConfig.Broker.FollowerReplicationThrottledRateProp))
-    }
-    assertEquals(3, propsCapture.getValues.size) //3 brokers
-  }
-
-  @Test
-  def shouldRemoveThrottleLimitFromAllBrokers(): Unit = {
-    //Given 3 brokers, but with assignment only covering 2 of them
-    val brokers = Seq(100, 101, 102)
-    val status = mutable.Map(TopicAndPartition("topic1", 0) -> ReassignmentCompleted)
-    val existingBrokerConfigs = propsWith(
-      (DynamicConfig.Broker.FollowerReplicationThrottledRateProp, "10"),
-      (DynamicConfig.Broker.LeaderReplicationThrottledRateProp, "100"),
-      ("useful.key", "value")
-    )
-
-    //Setup
-    val zk = stubZK(brokers = brokers)
-    val admin = createMock(classOf[AdminUtilities])
-    val propsCapture: Capture[Properties] = newCapture(CaptureType.ALL)
-    expect(admin.fetchEntityConfig(is(zk), is(ConfigType.Topic), anyString())).andStubReturn(new Properties)
-    expect(admin.changeBrokerConfig(is(zk), anyObject().asInstanceOf[Seq[Int]], capture(propsCapture))).anyTimes()
-    //Stub each invocation as EasyMock caches the return value which can be mutated
-    expect(admin.fetchEntityConfig(is(zk), is(ConfigType.Broker), is("100"))).andReturn(copyOf(existingBrokerConfigs))
-    expect(admin.fetchEntityConfig(is(zk), is(ConfigType.Broker), is("101"))).andReturn(copyOf(existingBrokerConfigs))
-    expect(admin.fetchEntityConfig(is(zk), is(ConfigType.Broker), is("102"))).andReturn(copyOf(existingBrokerConfigs))
-    replay(admin)
-
-    //When
-    ReassignPartitionsCommand.removeThrottle(zk, status, Map.empty, admin)
-
-    //Then props should have gone (dummy remains)
-    for (capture <- propsCapture.getValues) {
-      assertEquals("value", capture.get("useful.key"))
-      assertNull(capture.get(DynamicConfig.Broker.FollowerReplicationThrottledRateProp))
-      assertNull(capture.get(DynamicConfig.Broker.LeaderReplicationThrottledRateProp))
-    }
-    assertEquals(3, propsCapture.getValues.size) //3 brokers
-  }
-
-  @Test
-  def shouldRemoveThrottleReplicaListBasedOnProposedAssignment(): Unit = {
-    //Given two topics with existing config
-    val status = mutable.Map(TopicAndPartition("topic1", 0) -> ReassignmentCompleted,
-                             TopicAndPartition("topic2", 0) -> ReassignmentCompleted)
-    val existingConfigs = CoreUtils.propsWith(
-      (LogConfig.LeaderReplicationThrottledReplicasProp, "1:100:2:100"),
-      (LogConfig.FollowerReplicationThrottledReplicasProp, "1:101,2:101"),
-      ("useful.key", "value")
-    )
-
-    //Setup
-    val zk = stubZK(brokers = Seq(100, 101))
-    val admin = createMock(classOf[AdminUtilities])
-    val propsCapture: Capture[Properties] = newCapture(CaptureType.ALL)
-    expect(admin.fetchEntityConfig(is(zk), is(ConfigType.Broker), anyString())).andStubReturn(new Properties)
-    expect(admin.fetchEntityConfig(is(zk), is(ConfigType.Topic), is("topic1"))).andStubReturn(copyOf(existingConfigs))
-    expect(admin.fetchEntityConfig(is(zk), is(ConfigType.Topic), is("topic2"))).andStubReturn(copyOf(existingConfigs))
-
-    //Should change both topics
-    expect(admin.changeTopicConfig(is(zk), is("topic1"), capture(propsCapture)))
-    expect(admin.changeTopicConfig(is(zk), is("topic2"), capture(propsCapture)))
-
-    replay(admin)
-
-    //When
-    ReassignPartitionsCommand.removeThrottle(zk, status, Map.empty, admin)
-
-    //Then props should have gone (dummy remains)
-    for (actual <- propsCapture.getValues) {
-      assertEquals("value", actual.getProperty("useful.key"))
-      assertNull(actual.getProperty(LogConfig.LeaderReplicationThrottledReplicasProp))
-      assertNull(actual.getProperty(LogConfig.FollowerReplicationThrottledReplicasProp))
-    }
-    assertEquals(2, propsCapture.getValues.size) //2 topics
-  }
-
-  //Override eq as is for brevity
-  def is[T](v: T): T = EasyMock.eq(v)
-
-  @Before
-  def setup(): Unit = {
-    calls = 0
-  }
-
-  def stubZK(existingAssignment: mutable.Map[TopicAndPartition, Seq[Int]] = mutable.Map[TopicAndPartition, Seq[Int]](),
-             brokers: Seq[Int] = Seq[Int]()): ZkUtils = {
-    val zk = createMock(classOf[ZkUtils])
-    expect(zk.getReplicaAssignmentForTopics(anyObject().asInstanceOf[Seq[String]])).andStubReturn(existingAssignment)
-    expect(zk.getAllBrokersInCluster()).andStubReturn(brokers.map(TestUtils.createBroker(_, "", 1)))
-    replay(zk)
-    zk
-  }
-}
diff --git a/core/src/test/scala/unit/kafka/admin/ReassignPartitionsUnitTest.scala b/core/src/test/scala/unit/kafka/admin/ReassignPartitionsUnitTest.scala
new file mode 100644
index 0000000000000..1429a41ad2f40
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/admin/ReassignPartitionsUnitTest.scala
@@ -0,0 +1,662 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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.concurrent.ExecutionException
+import java.util.{Arrays, Collections}
+
+import kafka.admin.ReassignPartitionsCommand._
+import kafka.common.AdminCommandFailedException
+import kafka.utils.Exit
+import org.apache.kafka.clients.admin.{Config, MockAdminClient, PartitionReassignment}
+import org.apache.kafka.common.config.ConfigResource
+import org.apache.kafka.common.errors.{InvalidReplicationFactorException, UnknownTopicOrPartitionException}
+import org.apache.kafka.common.{Node, TopicPartition, TopicPartitionInfo, TopicPartitionReplica}
+import org.junit.Assert.{assertEquals, assertFalse, assertThrows, assertTrue}
+import org.junit.function.ThrowingRunnable
+import org.junit.rules.Timeout
+import org.junit.{After, Assert, Before, Rule, Test}
+
+import scala.collection.mutable
+import scala.jdk.CollectionConverters._
+
+class ReassignPartitionsUnitTest {
+  @Rule
+  def globalTimeout: Timeout = Timeout.millis(60000)
+
+  @Before
+  def setUp(): Unit = {
+    Exit.setExitProcedure((_, message) => throw new IllegalArgumentException(message.orNull))
+  }
+
+  @After
+  def tearDown(): Unit = {
+    Exit.resetExitProcedure()
+  }
+
+  @Test
+  def testCompareTopicPartitions(): Unit = {
+    assertTrue(compareTopicPartitions(new TopicPartition("abc", 0),
+      new TopicPartition("abc", 1)))
+    assertFalse(compareTopicPartitions(new TopicPartition("def", 0),
+      new TopicPartition("abc", 1)))
+  }
+
+  @Test
+  def testCompareTopicPartitionReplicas(): Unit = {
+    assertTrue(compareTopicPartitionReplicas(new TopicPartitionReplica("def", 0, 0),
+      new TopicPartitionReplica("abc", 0, 1)))
+    assertFalse(compareTopicPartitionReplicas(new TopicPartitionReplica("def", 0, 0),
+      new TopicPartitionReplica("cde", 0, 0)))
+  }
+
+  @Test
+  def testPartitionReassignStatesToString(): Unit = {
+    assertEquals(Seq(
+      "Status of partition reassignment:",
+      "Reassignment of partition bar-0 is still in progress.",
+      "Reassignment of partition foo-0 is complete.",
+      "Reassignment of partition foo-1 is still in progress.").
+        mkString(System.lineSeparator()),
+      partitionReassignmentStatesToString(Map(
+        new TopicPartition("foo", 0) ->
+          new PartitionReassignmentState(Seq(1, 2, 3), Seq(1, 2, 3), true),
+        new TopicPartition("foo", 1) ->
+          new PartitionReassignmentState(Seq(1, 2, 3), Seq(1, 2, 4), false),
+        new TopicPartition("bar", 0) ->
+          new PartitionReassignmentState(Seq(1, 2, 3), Seq(1, 2, 4), false),
+      )))
+  }
+
+  private def addTopics(adminClient: MockAdminClient): Unit = {
+    val b = adminClient.brokers()
+    adminClient.addTopic(false, "foo", Arrays.asList(
+      new TopicPartitionInfo(0, b.get(0),
+        Arrays.asList(b.get(0), b.get(1), b.get(2)),
+        Arrays.asList(b.get(0), b.get(1))),
+      new TopicPartitionInfo(1, b.get(1),
+        Arrays.asList(b.get(1), b.get(2), b.get(3)),
+        Arrays.asList(b.get(1), b.get(2), b.get(3)))
+    ), Collections.emptyMap())
+    adminClient.addTopic(false, "bar", Arrays.asList(
+      new TopicPartitionInfo(0, b.get(2),
+        Arrays.asList(b.get(2), b.get(3), b.get(0)),
+        Arrays.asList(b.get(2), b.get(3), b.get(0)))
+    ), Collections.emptyMap())
+  }
+
+  @Test
+  def testFindPartitionReassignmentStates(): Unit = {
+    val adminClient = new MockAdminClient.Builder().numBrokers(4).build()
+    try {
+      addTopics(adminClient)
+      // Create a reassignment and test findPartitionReassignmentStates.
+      val reassignmentResult: Map[TopicPartition, Class[_ <: Throwable]] = alterPartitionReassignments(adminClient, Map(
+        new TopicPartition("foo", 0) -> Seq(0,1,3),
+        new TopicPartition("quux", 0) -> Seq(1,2,3))).map { case (k, v) => k -> v.getClass }.toMap
+      assertEquals(Map(new TopicPartition("quux", 0) -> classOf[UnknownTopicOrPartitionException]),
+        reassignmentResult)
+      assertEquals((Map(
+          new TopicPartition("foo", 0) -> PartitionReassignmentState(Seq(0,1,2), Seq(0,1,3), false),
+          new TopicPartition("foo", 1) -> PartitionReassignmentState(Seq(1,2,3), Seq(1,2,3), true)
+        ), true),
+        findPartitionReassignmentStates(adminClient, Seq(
+          (new TopicPartition("foo", 0), Seq(0,1,3)),
+          (new TopicPartition("foo", 1), Seq(1,2,3))
+        )))
+      // Cancel the reassignment and test findPartitionReassignmentStates again.
+      val cancelResult: Map[TopicPartition, Class[_ <: Throwable]] = cancelPartitionReassignments(adminClient,
+        Set(new TopicPartition("foo", 0), new TopicPartition("quux", 2))).map { case (k, v) =>
+          k -> v.getClass
+        }.toMap
+      assertEquals(Map(new TopicPartition("quux", 2) -> classOf[UnknownTopicOrPartitionException]),
+        cancelResult)
+      assertEquals((Map(
+          new TopicPartition("foo", 0) -> PartitionReassignmentState(Seq(0,1,2), Seq(0,1,3), true),
+          new TopicPartition("foo", 1) -> PartitionReassignmentState(Seq(1,2,3), Seq(1,2,3), true)
+        ), false),
+          findPartitionReassignmentStates(adminClient, Seq(
+            (new TopicPartition("foo", 0), Seq(0,1,3)),
+            (new TopicPartition("foo", 1), Seq(1,2,3))
+          )))
+    } finally {
+      adminClient.close()
+    }
+  }
+
+  @Test
+  def testFindLogDirMoveStates(): Unit = {
+    val adminClient = new MockAdminClient.Builder().
+      numBrokers(4).
+      brokerLogDirs(Arrays.asList(
+          Arrays.asList("/tmp/kafka-logs0", "/tmp/kafka-logs1"),
+          Arrays.asList("/tmp/kafka-logs0", "/tmp/kafka-logs1"),
+          Arrays.asList("/tmp/kafka-logs0", "/tmp/kafka-logs1"),
+          Arrays.asList("/tmp/kafka-logs0", null))).
+      build();
+    try {
+      addTopics(adminClient)
+      val b = adminClient.brokers()
+      adminClient.addTopic(false, "quux", Arrays.asList(
+        new TopicPartitionInfo(0, b.get(2),
+            Arrays.asList(b.get(1), b.get(2), b.get(3)),
+            Arrays.asList(b.get(1), b.get(2), b.get(3)))),
+        Collections.emptyMap())
+      adminClient.alterReplicaLogDirs(Map(
+        new TopicPartitionReplica("foo", 0, 0) -> "/tmp/kafka-logs1",
+        new TopicPartitionReplica("quux", 0, 0) -> "/tmp/kafka-logs1"
+      ).asJava).all().get()
+      assertEquals(Map(
+        new TopicPartitionReplica("bar", 0, 0) -> new CompletedMoveState("/tmp/kafka-logs0"),
+        new TopicPartitionReplica("foo", 0, 0) -> new ActiveMoveState("/tmp/kafka-logs0",
+            "/tmp/kafka-logs1", "/tmp/kafka-logs1"),
+        new TopicPartitionReplica("foo", 1, 0) -> new CancelledMoveState("/tmp/kafka-logs0",
+          "/tmp/kafka-logs1"),
+        new TopicPartitionReplica("quux", 1, 0) -> new MissingLogDirMoveState("/tmp/kafka-logs1"),
+        new TopicPartitionReplica("quuz", 0, 0) -> new MissingReplicaMoveState("/tmp/kafka-logs0")
+      ), findLogDirMoveStates(adminClient, Map(
+        new TopicPartitionReplica("bar", 0, 0) -> "/tmp/kafka-logs0",
+        new TopicPartitionReplica("foo", 0, 0) -> "/tmp/kafka-logs1",
+        new TopicPartitionReplica("foo", 1, 0) -> "/tmp/kafka-logs1",
+        new TopicPartitionReplica("quux", 1, 0) -> "/tmp/kafka-logs1",
+        new TopicPartitionReplica("quuz", 0, 0) -> "/tmp/kafka-logs0"
+      )))
+    } finally {
+      adminClient.close()
+    }
+  }
+
+  @Test
+  def testReplicaMoveStatesToString(): Unit = {
+    assertEquals(Seq(
+      "Reassignment of replica bar-0-0 completed successfully.",
+      "Reassignment of replica foo-0-0 is still in progress.",
+      "Partition foo-1 on broker 0 is not being moved from log dir /tmp/kafka-logs0 to /tmp/kafka-logs1.",
+      "Partition quux-0 cannot be found in any live log directory on broker 0.",
+      "Partition quux-1 on broker 1 is being moved to log dir /tmp/kafka-logs2 instead of /tmp/kafka-logs1.",
+      "Partition quux-2 is not found in any live log dir on broker 1. " +
+          "There is likely an offline log directory on the broker.").mkString(System.lineSeparator()),
+        replicaMoveStatesToString(Map(
+          new TopicPartitionReplica("bar", 0, 0) -> new CompletedMoveState("/tmp/kafka-logs0"),
+          new TopicPartitionReplica("foo", 0, 0) -> new ActiveMoveState("/tmp/kafka-logs0",
+            "/tmp/kafka-logs1", "/tmp/kafka-logs1"),
+          new TopicPartitionReplica("foo", 1, 0) -> new CancelledMoveState("/tmp/kafka-logs0",
+            "/tmp/kafka-logs1"),
+          new TopicPartitionReplica("quux", 0, 0) -> new MissingReplicaMoveState("/tmp/kafka-logs1"),
+          new TopicPartitionReplica("quux", 1, 1) -> new ActiveMoveState("/tmp/kafka-logs0",
+            "/tmp/kafka-logs1", "/tmp/kafka-logs2"),
+          new TopicPartitionReplica("quux", 2, 1) -> new MissingLogDirMoveState("/tmp/kafka-logs1")
+        )))
+  }
+
+  @Test
+  def testGetReplicaAssignments(): Unit = {
+    val adminClient = new MockAdminClient.Builder().numBrokers(4).build()
+    try {
+      addTopics(adminClient)
+      assertEquals(Map(
+          new TopicPartition("foo", 0) -> Seq(0, 1, 2),
+          new TopicPartition("foo", 1) -> Seq(1, 2, 3),
+        ),
+        getReplicaAssignmentForTopics(adminClient, Seq("foo")))
+      assertEquals(Map(
+          new TopicPartition("foo", 0) -> Seq(0, 1, 2),
+          new TopicPartition("bar", 0) -> Seq(2, 3, 0),
+        ),
+        getReplicaAssignmentForPartitions(adminClient, Set(
+          new TopicPartition("foo", 0), new TopicPartition("bar", 0))))
+    } finally {
+      adminClient.close()
+    }
+  }
+
+  @Test
+  def testGetBrokerRackInformation(): Unit = {
+    val adminClient = new MockAdminClient.Builder().
+      brokers(Arrays.asList(new Node(0, "locahost", 9092, "rack0"),
+        new Node(1, "locahost", 9093, "rack1"),
+        new Node(2, "locahost", 9094, null))).
+      build()
+    try {
+      assertEquals(Seq(
+        new BrokerMetadata(0, Some("rack0")),
+        new BrokerMetadata(1, Some("rack1"))
+      ), getBrokerMetadata(adminClient, Seq(0, 1), true))
+      assertEquals(Seq(
+        new BrokerMetadata(0, None),
+        new BrokerMetadata(1, None)
+      ), getBrokerMetadata(adminClient, Seq(0, 1), false))
+      assertStartsWith("Not all brokers have rack information",
+        assertThrows(classOf[AdminOperationException], new ThrowingRunnable {
+          override def run(): Unit = getBrokerMetadata(adminClient, Seq(1, 2), true)
+        }).getMessage)
+      assertEquals(Seq(
+        new BrokerMetadata(1, None),
+        new BrokerMetadata(2, None)
+      ), getBrokerMetadata(adminClient, Seq(1, 2), false))
+    } finally {
+      adminClient.close()
+    }
+  }
+
+  @Test
+  def testParseGenerateAssignmentArgs(): Unit = {
+    assertStartsWith("Broker list contains duplicate entries",
+      assertThrows("Expected to detect duplicate broker list entries",
+        classOf[AdminCommandFailedException], new ThrowingRunnable {
+          override def run():Unit = parseGenerateAssignmentArgs(
+            """{"topics": [{"topic": "foo"}], "version":1}""", "1,1,2")
+        }).getMessage)
+    assertStartsWith("Broker list contains duplicate entries",
+      assertThrows("Expected to detect duplicate broker list entries",
+        classOf[AdminCommandFailedException], new ThrowingRunnable {
+          override def run():Unit = parseGenerateAssignmentArgs(
+            """{"topics": [{"topic": "foo"}], "version":1}""", "5,2,3,4,5")
+        }).getMessage)
+    assertEquals((Seq(5,2,3,4),Seq("foo")),
+      parseGenerateAssignmentArgs("""{"topics": [{"topic": "foo"}], "version":1}""",
+        "5,2,3,4"))
+    assertStartsWith("List of topics to reassign contains duplicate entries",
+      assertThrows("Expected to detect duplicate topic entries",
+        classOf[AdminCommandFailedException], new ThrowingRunnable {
+          override def run():Unit = parseGenerateAssignmentArgs(
+            """{"topics": [{"topic": "foo"},{"topic": "foo"}], "version":1}""", "5,2,3,4")
+        }).getMessage)
+    assertEquals((Seq(5,3,4),Seq("foo","bar")),
+      parseGenerateAssignmentArgs(
+        """{"topics": [{"topic": "foo"},{"topic": "bar"}], "version":1}""",
+        "5,3,4"))
+  }
+
+  @Test
+  def testGenerateAssignmentFailsWithoutEnoughReplicas(): Unit = {
+    val adminClient = new MockAdminClient.Builder().numBrokers(4).build()
+    try {
+      addTopics(adminClient)
+      assertStartsWith("Replication factor: 3 larger than available brokers: 2",
+        assertThrows("Expected generateAssignment to fail",
+          classOf[InvalidReplicationFactorException], new ThrowingRunnable {
+            override def run():Unit = {
+              generateAssignment(adminClient,
+                """{"topics":[{"topic":"foo"},{"topic":"bar"}]}""", "0,1", false)
+            }
+          }).getMessage)
+    } finally {
+      adminClient.close()
+    }
+  }
+
+  @Test
+  def testGenerateAssignmentWithInconsistentRacks(): Unit = {
+    val adminClient = new MockAdminClient.Builder().
+      brokers(Arrays.asList(
+        new Node(0, "locahost", 9092, "rack0"),
+        new Node(1, "locahost", 9093, "rack0"),
+        new Node(2, "locahost", 9094, null),
+        new Node(3, "locahost", 9095, "rack1"),
+        new Node(4, "locahost", 9096, "rack1"),
+        new Node(5, "locahost", 9097, "rack2"))).
+      build()
+    try {
+      addTopics(adminClient)
+      assertStartsWith("Not all brokers have rack information.",
+        assertThrows("Expected generateAssignment to fail",
+          classOf[AdminOperationException], new ThrowingRunnable {
+            override def run():Unit = {
+              generateAssignment(adminClient,
+                """{"topics":[{"topic":"foo"}]}""", "0,1,2,3", true)
+            }
+          }).getMessage)
+      // It should succeed when --disable-rack-aware is used.
+      val (_, current) = generateAssignment(adminClient,
+        """{"topics":[{"topic":"foo"}]}""", "0,1,2,3", false)
+      assertEquals(Map(
+        new TopicPartition("foo", 0) -> Seq(0, 1, 2),
+        new TopicPartition("foo", 1) -> Seq(1, 2, 3),
+      ), current)
+    } finally {
+      adminClient.close()
+    }
+  }
+
+  @Test
+  def testGenerateAssignmentWithFewerBrokers(): Unit = {
+    val adminClient = new MockAdminClient.Builder().numBrokers(4).build()
+    try {
+      addTopics(adminClient)
+      val goalBrokers = Set(0,1,3)
+      val (proposed, current) = generateAssignment(adminClient,
+        """{"topics":[{"topic":"foo"},{"topic":"bar"}]}""",
+        goalBrokers.mkString(","), false)
+      assertEquals(Map(
+        new TopicPartition("foo", 0) -> Seq(0, 1, 2),
+        new TopicPartition("foo", 1) -> Seq(1, 2, 3),
+        new TopicPartition("bar", 0) -> Seq(2, 3, 0)
+      ), current)
+
+      // The proposed assignment should only span the provided brokers
+      proposed.values.foreach {
+        case replicas => {
+          if (!replicas.forall(goalBrokers.contains(_))) {
+            Assert.fail(s"Proposed assignment ${proposed} puts replicas on brokers " +
+              s"other than ${goalBrokers}")
+          }
+        }
+      }
+    } finally {
+      adminClient.close()
+    }
+  }
+
+  @Test
+  def testCurrentPartitionReplicaAssignmentToString(): Unit = {
+    assertEquals(Seq(
+        """Current partition replica assignment""",
+        """""",
+        """{"version":1,"partitions":""" +
+          """[{"topic":"bar","partition":0,"replicas":[7,8],"log_dirs":["any","any"]},""" +
+          """{"topic":"foo","partition":1,"replicas":[4,5,6],"log_dirs":["any","any","any"]}]""" +
+        """}""",
+        """""",
+        """Save this to use as the --reassignment-json-file option during rollback"""
+      ).mkString(System.lineSeparator()),
+      currentPartitionReplicaAssignmentToString(Map(
+          new TopicPartition("foo", 1) -> Seq(1,2,3),
+          new TopicPartition("bar", 0) -> Seq(7,8,9)
+        ),
+        Map(
+          new TopicPartition("foo", 0) -> Seq(1,2,3),
+          new TopicPartition("foo", 1) -> Seq(4,5,6),
+          new TopicPartition("bar", 0) -> Seq(7,8),
+          new TopicPartition("baz", 0) -> Seq(10,11,12)
+        ),
+      ))
+  }
+
+  @Test
+  def testMoveMap(): Unit = {
+    val moveMap = calculateProposedMoveMap(Map(
+      new TopicPartition("foo", 0) -> new PartitionReassignment(
+        Arrays.asList(1,2,3),Arrays.asList(4),Arrays.asList(3)),
+      new TopicPartition("foo", 1) -> new PartitionReassignment(
+        Arrays.asList(4,5,6),Arrays.asList(7, 8),Arrays.asList(4, 5))
+    ), Map(
+      new TopicPartition("foo", 0) -> Seq(1,2,5),
+      new TopicPartition("bar", 0) -> Seq(1,2,3)
+    ), Map(
+      new TopicPartition("foo", 0) -> Seq(1,2,3),
+      new TopicPartition("foo", 1) -> Seq(4,5,6),
+      new TopicPartition("bar", 0) -> Seq(2,3,4),
+      new TopicPartition("baz", 0) -> Seq(1,2,3)
+    ))
+    assertEquals(
+      mutable.Map("foo" -> mutable.Map(
+          0 -> new PartitionMove(mutable.Set(1,2,3), mutable.Set(5)),
+          1 -> new PartitionMove(mutable.Set(4,5,6), mutable.Set(7, 8))
+        ),
+        "bar" -> mutable.Map(
+          0 -> new PartitionMove(mutable.Set(2,3,4), mutable.Set(1)),
+        )
+      ), moveMap)
+    assertEquals(Map(
+        "foo" -> "0:1,0:2,0:3,1:4,1:5,1:6",
+        "bar" -> "0:2,0:3,0:4"
+      ), calculateLeaderThrottles(moveMap))
+    assertEquals(Map(
+        "foo" -> "0:5,1:7,1:8",
+        "bar" -> "0:1"
+      ), calculateFollowerThrottles(moveMap))
+    assertEquals(Set(1,2,3,4,5,6,7,8), calculateReassigningBrokers(moveMap))
+    assertEquals(Set(0,2), calculateMovingBrokers(
+      Set(new TopicPartitionReplica("quux", 0, 0),
+          new TopicPartitionReplica("quux", 1, 2))))
+  }
+
+  @Test
+  def testParseExecuteAssignmentArgs(): Unit = {
+    assertStartsWith("Partition reassignment list cannot be empty",
+      assertThrows("Expected to detect empty partition reassignment list",
+        classOf[AdminCommandFailedException], new ThrowingRunnable {
+          override def run():Unit =
+            parseExecuteAssignmentArgs("""{"version":1,"partitions":[]}""")
+        }).getMessage)
+    assertStartsWith("Partition reassignment contains duplicate topic partitions",
+      assertThrows("Expected to detect a partition list with duplicate entries",
+        classOf[AdminCommandFailedException], new ThrowingRunnable {
+          override def run():Unit =
+            parseExecuteAssignmentArgs(
+              """{"version":1,"partitions":""" +
+                """[{"topic":"foo","partition":0,"replicas":[0,1],"log_dirs":["any","any"]},""" +
+                """{"topic":"foo","partition":0,"replicas":[2,3,4],"log_dirs":["any","any","any"]}""" +
+                """]}""")
+        }).getMessage)
+    assertStartsWith("Partition reassignment contains duplicate topic partitions",
+      assertThrows("Expected to detect a partition replica list with duplicate entries",
+        classOf[AdminCommandFailedException], new ThrowingRunnable {
+          override def run():Unit =
+            parseExecuteAssignmentArgs(
+              """{"version":1,"partitions":""" +
+                """[{"topic":"foo","partition":0,"replicas":[0,1],"log_dirs":["/abc","/def"]},""" +
+                """{"topic":"foo","partition":0,"replicas":[2,3],"log_dirs":["/abc","/def"]}""" +
+                """]}""")
+        }).getMessage)
+    assertStartsWith("Partition replica lists may not contain duplicate entries",
+      assertThrows("Expected to detect a partition replica list with duplicate entries",
+        classOf[AdminCommandFailedException], new ThrowingRunnable {
+          override def run():Unit =
+            parseExecuteAssignmentArgs(
+              """{"version":1,"partitions":""" +
+                """[{"topic":"foo","partition":0,"replicas":[0,0],"log_dirs":["/abc","/def"]},""" +
+                """{"topic":"foo","partition":1,"replicas":[2,3],"log_dirs":["/abc","/def"]}""" +
+                """]}""")
+        }).getMessage)
+    assertEquals((Map(
+        new TopicPartition("foo", 0) -> Seq(1, 2, 3),
+        new TopicPartition("foo", 1) -> Seq(3, 4, 5),
+      ), Map(
+      )),
+      parseExecuteAssignmentArgs(
+        """{"version":1,"partitions":""" +
+          """[{"topic":"foo","partition":0,"replicas":[1,2,3],"log_dirs":["any","any","any"]},""" +
+          """{"topic":"foo","partition":1,"replicas":[3,4,5],"log_dirs":["any","any","any"]}""" +
+          """]}"""))
+    assertEquals((Map(
+      new TopicPartition("foo", 0) -> Seq(1, 2, 3),
+    ), Map(
+      new TopicPartitionReplica("foo", 0, 1) -> "/tmp/a",
+      new TopicPartitionReplica("foo", 0, 2) -> "/tmp/b",
+      new TopicPartitionReplica("foo", 0, 3) -> "/tmp/c"
+    )),
+      parseExecuteAssignmentArgs(
+        """{"version":1,"partitions":""" +
+          """[{"topic":"foo","partition":0,"replicas":[1,2,3],"log_dirs":["/tmp/a","/tmp/b","/tmp/c"]}""" +
+          """]}"""))
+  }
+
+  @Test
+  def testExecuteWithInvalidPartitionsFails(): Unit = {
+    val adminClient = new MockAdminClient.Builder().numBrokers(5).build()
+    try {
+      addTopics(adminClient)
+      assertStartsWith("Topic quux not found",
+        assertThrows("Expected reassignment with non-existent topic to fail",
+          classOf[ExecutionException], new ThrowingRunnable {
+            override def run():Unit =
+              executeAssignment(adminClient, false,
+                """{"version":1,"partitions":""" +
+                """[{"topic":"foo","partition":0,"replicas":[0,1],"log_dirs":["any","any"]},""" +
+                """{"topic":"quux","partition":0,"replicas":[2,3,4],"log_dirs":["any","any","any"]}""" +
+                """]}""")
+          }).getCause.getMessage)
+    } finally {
+      adminClient.close()
+    }
+  }
+
+  @Test
+  def testExecuteWithInvalidBrokerIdFails(): Unit = {
+    val adminClient = new MockAdminClient.Builder().numBrokers(4).build()
+    try {
+      addTopics(adminClient)
+      assertStartsWith("Unknown broker id 4",
+        assertThrows("Expected reassignment with non-existent broker id to fail",
+          classOf[AdminCommandFailedException], new ThrowingRunnable {
+            override def run():Unit =
+              executeAssignment(adminClient, false,
+                """{"version":1,"partitions":""" +
+                  """[{"topic":"foo","partition":0,"replicas":[0,1],"log_dirs":["any","any"]},""" +
+                  """{"topic":"foo","partition":1,"replicas":[2,3,4],"log_dirs":["any","any","any"]}""" +
+                  """]}""")
+          }).getMessage)
+    } finally {
+      adminClient.close()
+    }
+  }
+
+  @Test
+  def testModifyBrokerInterBrokerThrottle(): Unit = {
+    val adminClient = new MockAdminClient.Builder().numBrokers(4).build()
+    try {
+      modifyInterBrokerThrottle(adminClient, Set(0, 1, 2), 1000)
+      modifyInterBrokerThrottle(adminClient, Set(0, 3), 100)
+      val brokers = Seq(0, 1, 2, 3).map(
+        id => new ConfigResource(ConfigResource.Type.BROKER, id.toString))
+      val results = adminClient.describeConfigs(brokers.asJava).all().get()
+      verifyBrokerThrottleResults(results.get(brokers(0)), 100, -1)
+      verifyBrokerThrottleResults(results.get(brokers(1)), 1000, -1)
+      verifyBrokerThrottleResults(results.get(brokers(2)), 1000, -1)
+      verifyBrokerThrottleResults(results.get(brokers(3)), 100, -1)
+    } finally {
+      adminClient.close()
+    }
+  }
+
+  @Test
+  def testModifyLogDirThrottle(): Unit = {
+    val adminClient = new MockAdminClient.Builder().numBrokers(4).build()
+    try {
+      modifyLogDirThrottle(adminClient, Set(0, 1, 2), 2000)
+      modifyLogDirThrottle(adminClient, Set(0, 3), -1)
+      val brokers = Seq(0, 1, 2, 3).map(
+        id => new ConfigResource(ConfigResource.Type.BROKER, id.toString))
+      val results = adminClient.describeConfigs(brokers.asJava).all().get()
+      verifyBrokerThrottleResults(results.get(brokers(0)), -1, 2000)
+      verifyBrokerThrottleResults(results.get(brokers(1)), -1, 2000)
+      verifyBrokerThrottleResults(results.get(brokers(2)), -1, 2000)
+      verifyBrokerThrottleResults(results.get(brokers(3)), -1, -1)
+    } finally {
+      adminClient.close()
+    }
+  }
+
+  @Test
+  def testCurReassignmentsToString(): Unit = {
+    val adminClient = new MockAdminClient.Builder().numBrokers(4).build()
+    try {
+      addTopics(adminClient)
+      assertEquals("No partition reassignments found.", curReassignmentsToString(adminClient))
+      val reassignmentResult: Map[TopicPartition, Class[_ <: Throwable]] = alterPartitionReassignments(adminClient,
+        Map(
+          new TopicPartition("foo", 1) -> Seq(4,5,3),
+          new TopicPartition("foo", 0) -> Seq(0,1,4,2),
+          new TopicPartition("bar", 0) -> Seq(2,3)
+        )
+      ).map { case (k, v) => k -> v.getClass }.toMap
+      assertEquals(Map(), reassignmentResult)
+      assertEquals(Seq("Current partition reassignments:",
+                       "bar-0: replicas: 2,3,0. removing: 0.",
+                       "foo-0: replicas: 0,1,2. adding: 4.",
+                       "foo-1: replicas: 1,2,3. adding: 4,5. removing: 1,2.").mkString(System.lineSeparator()),
+                  curReassignmentsToString(adminClient))
+    } finally {
+      adminClient.close()
+    }
+  }
+
+  private def verifyBrokerThrottleResults(config: Config,
+                                          expectedInterBrokerThrottle: Long,
+                                          expectedReplicaAlterLogDirsThrottle: Long): Unit = {
+    val configs = new mutable.HashMap[String, String]
+    config.entries.forEach(entry => configs.put(entry.name, entry.value))
+    if (expectedInterBrokerThrottle >= 0) {
+      assertEquals(expectedInterBrokerThrottle.toString,
+        configs.getOrElse(brokerLevelLeaderThrottle, ""))
+      assertEquals(expectedInterBrokerThrottle.toString,
+        configs.getOrElse(brokerLevelFollowerThrottle, ""))
+    }
+    if (expectedReplicaAlterLogDirsThrottle >= 0) {
+      assertEquals(expectedReplicaAlterLogDirsThrottle.toString,
+        configs.getOrElse(brokerLevelLogDirThrottle, ""))
+    }
+  }
+
+  @Test
+  def testModifyTopicThrottles(): Unit = {
+    val adminClient = new MockAdminClient.Builder().numBrokers(4).build()
+    try {
+      addTopics(adminClient)
+      modifyTopicThrottles(adminClient,
+        Map("foo" -> "leaderFoo", "bar" -> "leaderBar"),
+        Map("bar" -> "followerBar"))
+      val topics = Seq("bar", "foo").map(
+        id => new ConfigResource(ConfigResource.Type.TOPIC, id.toString))
+      val results = adminClient.describeConfigs(topics.asJava).all().get()
+      verifyTopicThrottleResults(results.get(topics(0)), "leaderBar", "followerBar")
+      verifyTopicThrottleResults(results.get(topics(1)), "leaderFoo", "")
+    } finally {
+      adminClient.close()
+    }
+  }
+
+  private def verifyTopicThrottleResults(config: Config,
+                                         expectedLeaderThrottle: String,
+                                         expectedFollowerThrottle: String): Unit = {
+    val configs = new mutable.HashMap[String, String]
+    config.entries.forEach(entry => configs.put(entry.name, entry.value))
+    assertEquals(expectedLeaderThrottle.toString,
+      configs.getOrElse(topicLevelLeaderThrottle, ""))
+    assertEquals(expectedFollowerThrottle.toString,
+      configs.getOrElse(topicLevelFollowerThrottle, ""))
+  }
+
+  @Test
+  def testAlterReplicaLogDirs(): Unit = {
+    val adminClient = new MockAdminClient.Builder().
+      numBrokers(4).
+      brokerLogDirs(Collections.nCopies(4,
+        Arrays.asList("/tmp/kafka-logs0", "/tmp/kafka-logs1"))).
+      build()
+    try {
+      addTopics(adminClient)
+      assertEquals(Set(
+          new TopicPartitionReplica("foo", 0, 0)
+        ),
+        alterReplicaLogDirs(adminClient, Map(
+          new TopicPartitionReplica("foo", 0, 0) -> "/tmp/kafka-logs1",
+          new TopicPartitionReplica("quux", 1, 0) -> "/tmp/kafka-logs1"
+        )))
+    } finally {
+      adminClient.close()
+    }
+  }
+
+  def assertStartsWith(prefix: String, str: String): Unit = {
+    assertTrue("Expected the string to start with %s, but it was %s".format(prefix, str),
+      str.startsWith(prefix))
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/admin/ReplicationQuotaUtils.scala b/core/src/test/scala/unit/kafka/admin/ReplicationQuotaUtils.scala
index bdabd679e25f7..4ac64fe1a58a8 100644
--- a/core/src/test/scala/unit/kafka/admin/ReplicationQuotaUtils.scala
+++ b/core/src/test/scala/unit/kafka/admin/ReplicationQuotaUtils.scala
@@ -13,41 +13,42 @@
 package kafka.admin
 
 import kafka.log.LogConfig
-import kafka.server.{DynamicConfig, ConfigType, KafkaServer}
+import kafka.server.{ConfigType, DynamicConfig, KafkaServer}
 import kafka.utils.TestUtils
+import kafka.zk.AdminZkClient
 
 import scala.collection.Seq
 
 object ReplicationQuotaUtils {
 
-  def checkThrottleConfigRemovedFromZK(topic: String, servers: Seq[KafkaServer]): Unit = {
+  def checkThrottleConfigRemovedFromZK(adminZkClient: AdminZkClient, topic: String, servers: Seq[KafkaServer]): Unit = {
     TestUtils.waitUntilTrue(() => {
       val hasRateProp = servers.forall { server =>
-        val brokerConfig = AdminUtils.fetchEntityConfig(server.zkUtils, ConfigType.Broker, server.config.brokerId.toString)
+        val brokerConfig = adminZkClient.fetchEntityConfig(ConfigType.Broker, server.config.brokerId.toString)
         brokerConfig.contains(DynamicConfig.Broker.LeaderReplicationThrottledRateProp) ||
           brokerConfig.contains(DynamicConfig.Broker.FollowerReplicationThrottledRateProp)
       }
-      val topicConfig = AdminUtils.fetchEntityConfig(servers(0).zkUtils, ConfigType.Topic, topic)
+      val topicConfig = adminZkClient.fetchEntityConfig(ConfigType.Topic, topic)
       val hasReplicasProp = topicConfig.contains(LogConfig.LeaderReplicationThrottledReplicasProp) ||
         topicConfig.contains(LogConfig.FollowerReplicationThrottledReplicasProp)
       !hasRateProp && !hasReplicasProp
     }, "Throttle limit/replicas was not unset")
   }
 
-  def checkThrottleConfigAddedToZK(expectedThrottleRate: Long, servers: Seq[KafkaServer], topic: String, throttledLeaders: String, throttledFollowers: String): Unit = {
+  def checkThrottleConfigAddedToZK(adminZkClient: AdminZkClient, expectedThrottleRate: Long, servers: Seq[KafkaServer], topic: String, throttledLeaders: Set[String], throttledFollowers: Set[String]): Unit = {
     TestUtils.waitUntilTrue(() => {
       //Check for limit in ZK
       val brokerConfigAvailable = servers.forall { server =>
-        val configInZk = AdminUtils.fetchEntityConfig(server.zkUtils, ConfigType.Broker, server.config.brokerId.toString)
+        val configInZk = adminZkClient.fetchEntityConfig(ConfigType.Broker, server.config.brokerId.toString)
         val zkLeaderRate = configInZk.getProperty(DynamicConfig.Broker.LeaderReplicationThrottledRateProp)
         val zkFollowerRate = configInZk.getProperty(DynamicConfig.Broker.FollowerReplicationThrottledRateProp)
         zkLeaderRate != null && expectedThrottleRate == zkLeaderRate.toLong &&
           zkFollowerRate != null && expectedThrottleRate == zkFollowerRate.toLong
       }
       //Check replicas assigned
-      val topicConfig = AdminUtils.fetchEntityConfig(servers(0).zkUtils, ConfigType.Topic, topic)
-      val leader = topicConfig.getProperty(LogConfig.LeaderReplicationThrottledReplicasProp)
-      val follower = topicConfig.getProperty(LogConfig.FollowerReplicationThrottledReplicasProp)
+      val topicConfig = adminZkClient.fetchEntityConfig(ConfigType.Topic, topic)
+      val leader = topicConfig.getProperty(LogConfig.LeaderReplicationThrottledReplicasProp).split(",").toSet
+      val follower = topicConfig.getProperty(LogConfig.FollowerReplicationThrottledReplicasProp).split(",").toSet
       val topicConfigAvailable = leader == throttledLeaders && follower == throttledFollowers
       brokerConfigAvailable && topicConfigAvailable
     }, "throttle limit/replicas was not set")
diff --git a/core/src/test/scala/unit/kafka/admin/ResetConsumerGroupOffsetTest.scala b/core/src/test/scala/unit/kafka/admin/ResetConsumerGroupOffsetTest.scala
index 82274871af1e1..e6dd9407ef0af 100644
--- a/core/src/test/scala/unit/kafka/admin/ResetConsumerGroupOffsetTest.scala
+++ b/core/src/test/scala/unit/kafka/admin/ResetConsumerGroupOffsetTest.scala
@@ -13,16 +13,23 @@
 package kafka.admin
 
 import java.io.{BufferedWriter, File, FileWriter}
-import java.text.{ParseException, SimpleDateFormat}
+import java.text.{SimpleDateFormat}
 import java.util.{Calendar, Date, Properties}
 
-import kafka.admin.ConsumerGroupCommand.{ConsumerGroupCommandOptions, KafkaConsumerGroupService}
-import kafka.integration.KafkaServerTestHarness
+import joptsimple.OptionException
+import kafka.admin.ConsumerGroupCommand.ConsumerGroupService
 import kafka.server.KafkaConfig
 import kafka.utils.TestUtils
-import org.junit.{After, Before, Test}
+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.jdk.CollectionConverters._
+import scala.collection.Seq
+
 
-import scala.collection.mutable.ArrayBuffer
 
 /**
   * Test cases by:
@@ -34,569 +41,466 @@ import scala.collection.mutable.ArrayBuffer
   * - scope=topics+partitions, scenario=to-earliest
   * - export/import
   */
-class ResetConsumerGroupOffsetTest extends KafkaServerTestHarness {
+class ResetConsumerGroupOffsetTest extends ConsumerGroupCommandTest {
 
   val overridingProps = new Properties()
   val topic1 = "foo1"
   val topic2 = "foo2"
-  val group = "test.group"
-  val props = new Properties
-  val consumerGroupServices = new ArrayBuffer[KafkaConsumerGroupService]
-  val executors = new ArrayBuffer[ConsumerGroupExecutor]
-
-  /**
-    * Implementations must override this method to return a set of KafkaConfigs. This method will be invoked for every
-    * test and should not reuse previous configurations unless they select their ports randomly when servers are started.
-    */
-  override def generateConfigs: Seq[KafkaConfig] = TestUtils.createBrokerConfigs(1, zkConnect, enableControlledShutdown = false).map(KafkaConfig.fromProps(_, overridingProps))
 
-  @Before
-  override def setUp() {
-    super.setUp()
+  override def generateConfigs: Seq[KafkaConfig] = {
+    TestUtils.createBrokerConfigs(1, zkConnect, enableControlledShutdown = false)
+      .map(KafkaConfig.fromProps(_, overridingProps))
+  }
 
-    props.setProperty("group.id", group)
+  private def basicArgs: Array[String] = {
+    Array("--reset-offsets",
+      "--bootstrap-server", brokerList,
+      "--timeout", test.TestUtils.DEFAULT_MAX_WAIT_MS.toString)
   }
 
-  @After
-  override def tearDown() {
-    try {
-      executors.foreach(_.shutdown())
-      consumerGroupServices.foreach(_.close())
-    } finally {
-      super.tearDown()
-    }
+  private def buildArgsForGroups(groups: Seq[String], args: String*): Array[String] = {
+    val groupArgs = groups.flatMap(group => Seq("--group", group)).toArray
+    basicArgs ++ groupArgs ++ args
   }
 
-  @Test
-  def testResetOffsetsNotExistingGroup() {
-    createConsumerGroupExecutor(brokerList, 1, group, topic1)
+  private def buildArgsForGroup(group: String, args: String*): Array[String] = {
+    buildArgsForGroups(Seq(group), args: _*)
+  }
 
-    val cgcArgs = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", "missing.group", "--all-topics", "--to-current")
-    val opts = new ConsumerGroupCommandOptions(cgcArgs)
-    val consumerGroupCommand = createConsumerGroupService(opts)
+  private def buildArgsForAllGroups(args: String*): Array[String] = {
+    basicArgs ++ Array("--all-groups") ++ args
+  }
 
+  @Test
+  def testResetOffsetsNotExistingGroup(): Unit = {
+    val group = "missing.group"
+    val args = buildArgsForGroup(group, "--all-topics", "--to-current", "--execute")
+    val consumerGroupCommand = getConsumerGroupService(args)
+    // Make sure we got a coordinator
     TestUtils.waitUntilTrue(() => {
-      val assignmentsToReset = consumerGroupCommand.resetOffsets()
-      assignmentsToReset == Map.empty
-    }, "Expected to have an empty assignations map.")
-
+      consumerGroupCommand.collectGroupState(group).coordinator.host() == "localhost"
+    }, "Can't find a coordinator")
+    val resetOffsets = consumerGroupCommand.resetOffsets()(group)
+    assertEquals(Map.empty, resetOffsets)
+    assertEquals(resetOffsets, committedOffsets(group = group))
   }
 
   @Test
-  def testResetOffsetsNewConsumerExistingTopic(): Unit = {
-    val cgcArgs = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", "new.group", "--topic", topic1, "--to-offset", "50", "--execute")
-    val opts = new ConsumerGroupCommandOptions(cgcArgs)
-    val consumerGroupCommand = new KafkaConsumerGroupService(opts)
-
-    adminZkClient.createTopic(topic1, 1, 1)
-
-    TestUtils.produceMessages(servers, topic1, 100, acks = 1, 100 * 1000)
+  def testResetOffsetsExistingTopic(): Unit = {
+    val group = "new.group"
+    val args = buildArgsForGroup(group, "--topic", topic, "--to-offset", "50")
+    produceMessages(topic, 100)
+    resetAndAssertOffsets(args, expectedOffset = 50, dryRun = true)
+    resetAndAssertOffsets(args ++ Array("--dry-run"), expectedOffset = 50, dryRun = true)
+    resetAndAssertOffsets(args ++ Array("--execute"), expectedOffset = 50)
+  }
 
-    TestUtils.waitUntilTrue(() => {
-      val assignmentsToReset = consumerGroupCommand.resetOffsets()
-      assignmentsToReset.exists({ assignment => assignment._2.offset() == 50 })
-    }, "Expected the consumer group to reset to offset 1 (specific offset).")
+  @Test
+  def testResetOffsetsExistingTopicSelectedGroups(): Unit = {
+    produceMessages(topic, 100)
+    val groups =
+      for (id <- 1 to 3) yield {
+        val group = this.group + id
+        val executor = addConsumerGroupExecutor(numConsumers = 1, topic = topic, group = group)
+        awaitConsumerProgress(count = 100L, group = group)
+        executor.shutdown()
+        group
+      }
+    val args = buildArgsForGroups(groups,"--topic", topic, "--to-offset", "50")
+    resetAndAssertOffsets(args, expectedOffset = 50, dryRun = true)
+    resetAndAssertOffsets(args ++ Array("--dry-run"), expectedOffset = 50, dryRun = true)
+    resetAndAssertOffsets(args ++ Array("--execute"), expectedOffset = 50)
+  }
 
-    printConsumerGroup("new.group")
-    adminZkClient.deleteTopic(topic1)
-    consumerGroupCommand.close()
+  @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 testResetOffsetsToLocalDateTime() {
-    adminZkClient.createTopic(topic1, 1, 1)
+  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)
 
-    TestUtils.produceMessages(servers, topic1, 100, acks = 1, 100 * 1000)
-
-    val cgcArgs = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--all-topics")
-    val opts = new ConsumerGroupCommandOptions(cgcArgs)
-    val consumerGroupCommand = createConsumerGroupService(opts)
-
-    val executor = createConsumerGroupExecutor(brokerList, 1, group, topic1)
-
-    TestUtils.waitUntilTrue(() => {
-      val (_, assignmentsOption) = consumerGroupCommand.describeGroup()
-      assignmentsOption match {
-        case Some(assignments) =>
-          val sumOffset = assignments.filter(_.topic.exists(_ == topic1))
-            .filter(_.offset.isDefined)
-            .map(assignment => assignment.offset.get)
-            .foldLeft(0.toLong)(_ + _)
-          sumOffset == 100
-        case _ => false
-      }
-    }, "Expected that consumer group has consumed all messages from topic/partition.")
+    produceMessages(topic, 100)
 
+    val executor = addConsumerGroupExecutor(numConsumers = 1, topic)
+    awaitConsumerProgress(count = 100L)
     executor.shutdown()
 
-    val cgcArgs1 = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--all-topics", "--to-datetime", format.format(calendar.getTime), "--execute")
-    val opts1 = new ConsumerGroupCommandOptions(cgcArgs1)
-    val consumerGroupCommand1 = createConsumerGroupService(opts1)
-
-    TestUtils.waitUntilTrue(() => {
-      val assignmentsToReset = consumerGroupCommand1.resetOffsets()
-      assignmentsToReset.exists { assignment => assignment._2.offset() == 0 }
-    }, "Expected the consumer group to reset to when offset was 50.")
-
-    printConsumerGroup()
-
-    adminZkClient.deleteTopic(topic1)
+    val args = buildArgsForGroup(group, "--all-topics", "--to-datetime", format.format(calendar.getTime), "--execute")
+    resetAndAssertOffsets(args, expectedOffset = 0)
   }
 
   @Test
-  def testResetOffsetsToZonedDateTime() {
-    adminZkClient.createTopic(topic1, 1, 1)
-    TestUtils.produceMessages(servers, topic1, 50, acks = 1, 100 * 1000)
-
+  def testResetOffsetsToZonedDateTime(): Unit = {
     val format = new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss.SSSXXX")
-    val checkpoint = new Date()
 
-    TestUtils.produceMessages(servers, topic1, 50, acks = 1, 100 * 1000)
-
-    val cgcArgs = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--all-topics")
-    val opts = new ConsumerGroupCommandOptions(cgcArgs)
-    val consumerGroupCommand = createConsumerGroupService(opts)
-
-    val executor = createConsumerGroupExecutor(brokerList, 1, group, topic1)
-
-    TestUtils.waitUntilTrue(() => {
-      val (_, assignmentsOption) = consumerGroupCommand.describeGroup()
-      assignmentsOption match {
-        case Some(assignments) =>
-          val sumOffset = (assignments.filter(_.topic.exists(_ == topic1))
-            .filter(_.offset.isDefined)
-            .map(assignment => assignment.offset.get) foldLeft 0.toLong)(_ + _)
-          sumOffset == 100
-        case _ => false
-      }
-    }, "Expected that consumer group has consumed all messages from topic/partition.")
+    produceMessages(topic, 50)
+    val checkpoint = new Date()
+    produceMessages(topic, 50)
 
+    val executor = addConsumerGroupExecutor(numConsumers = 1, topic)
+    awaitConsumerProgress(count = 100L)
     executor.shutdown()
 
-    val cgcArgs1 = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--all-topics", "--to-datetime", format.format(checkpoint), "--execute")
-    val opts1 = new ConsumerGroupCommandOptions(cgcArgs1)
-    val consumerGroupCommand1 = createConsumerGroupService(opts1)
-
-    TestUtils.waitUntilTrue(() => {
-      val assignmentsToReset = consumerGroupCommand1.resetOffsets()
-      assignmentsToReset.exists { assignment => assignment._2.offset() == 50 }
-    }, "Expected the consumer group to reset to when offset was 50.")
-
-    printConsumerGroup()
-
-    adminZkClient.deleteTopic(topic1)
+    val args = buildArgsForGroup(group, "--all-topics", "--to-datetime", format.format(checkpoint), "--execute")
+    resetAndAssertOffsets(args, expectedOffset = 50)
   }
 
   @Test
-  def testDateTimeFormats() {
-    //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"))
-    invokeGetDateTimeMethod(new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss.SSSX"))
-    invokeGetDateTimeMethod(new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss.SSSXX"))
-    invokeGetDateTimeMethod(new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss.SSSXXX"))
-
-    //check some invalid formats
-    try {
-      invokeGetDateTimeMethod(new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss"))
-      fail("Call to getDateTime should fail")
-    } catch {
-      case _: ParseException =>
-    }
-
-    try {
-      invokeGetDateTimeMethod(new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss.X"))
-      fail("Call to getDateTime should fail")
-    } catch {
-      case _: ParseException =>
-    }
+  def testResetOffsetsByDuration(): Unit = {
+    val args = buildArgsForGroup(group, "--all-topics", "--by-duration", "PT1M", "--execute")
+    produceConsumeAndShutdown(topic, group, totalMessages = 100)
+    resetAndAssertOffsets(args, expectedOffset = 0)
   }
 
-  private def invokeGetDateTimeMethod(format: SimpleDateFormat) {
-    val checkpoint = new Date()
-    val cgcArgs = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--all-topics", "--to-datetime", format.format(checkpoint), "--execute")
-    val opts  = new ConsumerGroupCommandOptions(cgcArgs)
-    val consumerGroupCommand = createConsumerGroupService(opts)
-    consumerGroupCommand.getDateTime
+  @Test
+  def testResetOffsetsByDurationToEarliest(): Unit = {
+    val args = buildArgsForGroup(group, "--all-topics", "--by-duration", "PT0.1S", "--execute")
+    produceConsumeAndShutdown(topic, group, totalMessages = 100)
+    resetAndAssertOffsets(args, expectedOffset = 100)
   }
 
   @Test
-  def testResetOffsetsByDuration() {
-    val cgcArgs = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--all-topics", "--by-duration", "PT1M", "--execute")
-    val opts = new ConsumerGroupCommandOptions(cgcArgs)
-    val consumerGroupCommand = createConsumerGroupService(opts)
-
-    adminZkClient.createTopic(topic1, 1, 1)
-
-    produceConsumeAndShutdown(consumerGroupCommand, 1, topic1, 100)
-
-    TestUtils.waitUntilTrue(() => {
-        val assignmentsToReset = consumerGroupCommand.resetOffsets()
-        assignmentsToReset.exists { assignment => assignment._2.offset() == 0 }
-    }, "Expected the consumer group to reset to offset 0 (earliest by duration).")
-
-    printConsumerGroup()
-
-    adminZkClient.deleteTopic(topic1)
+  def testResetOffsetsToEarliest(): Unit = {
+    val args = buildArgsForGroup(group, "--all-topics", "--to-earliest", "--execute")
+    produceConsumeAndShutdown(topic, group, totalMessages = 100)
+    resetAndAssertOffsets(args, expectedOffset = 0)
   }
 
   @Test
-  def testResetOffsetsByDurationToEarliest() {
-    val cgcArgs = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--all-topics", "--by-duration", "PT0.1S", "--execute")
-    val opts = new ConsumerGroupCommandOptions(cgcArgs)
-    val consumerGroupCommand = createConsumerGroupService(opts)
-
-    adminZkClient.createTopic(topic1, 1, 1)
-
-    produceConsumeAndShutdown(consumerGroupCommand, 1, topic1, 100)
-
-    TestUtils.waitUntilTrue(() => {
-      val assignmentsToReset = consumerGroupCommand.resetOffsets()
-      assignmentsToReset.exists { assignment => assignment._2.offset() == 100 }
-    }, "Expected the consumer group to reset to offset 100 (latest by duration).")
-
-    printConsumerGroup()
-    adminZkClient.deleteTopic(topic1)
+  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 testResetOffsetsToEarliest() {
-    val cgcArgs = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--all-topics", "--to-earliest", "--execute")
-    val opts = new ConsumerGroupCommandOptions(cgcArgs)
-    val consumerGroupCommand = createConsumerGroupService(opts)
-
-    adminZkClient.createTopic(topic1, 1, 1)
-
-    produceConsumeAndShutdown(consumerGroupCommand, 1, topic1, 100)
-
-    TestUtils.waitUntilTrue(() => {
-      val assignmentsToReset = consumerGroupCommand.resetOffsets()
-      assignmentsToReset.exists { assignment => assignment._2.offset() == 0 }
-    }, "Expected the consumer group to reset to offset 0 (earliest).")
-
-    printConsumerGroup()
-    adminZkClient.deleteTopic(topic1)
+  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 testResetOffsetsToLatest() {
-    val cgcArgs = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--all-topics", "--to-latest", "--execute")
-    val opts = new ConsumerGroupCommandOptions(cgcArgs)
-    val consumerGroupCommand = createConsumerGroupService(opts)
-
-    adminZkClient.createTopic(topic1, 1, 1)
-
-    produceConsumeAndShutdown(consumerGroupCommand, 1, topic1, 100)
-
-    TestUtils.produceMessages(servers, topic1, 100, acks = 1, 100 * 1000)
-
-
-    TestUtils.waitUntilTrue(() => {
-      val assignmentsToReset = consumerGroupCommand.resetOffsets()
-      assignmentsToReset.exists({ assignment => assignment._2.offset() == 200 })
-    }, "Expected the consumer group to reset to offset 200 (latest).")
-
-    printConsumerGroup()
-    adminZkClient.deleteTopic(topic1)
+  def testResetOffsetsToSpecificOffset(): Unit = {
+    val args = buildArgsForGroup(group, "--all-topics", "--to-offset", "1", "--execute")
+    produceConsumeAndShutdown(topic, group, totalMessages = 100)
+    resetAndAssertOffsets(args, expectedOffset = 1)
   }
 
   @Test
-  def testResetOffsetsToCurrentOffset() {
-    val cgcArgs = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--all-topics", "--to-current", "--execute")
-    val opts = new ConsumerGroupCommandOptions(cgcArgs)
-    val consumerGroupCommand = createConsumerGroupService(opts)
-
-    adminZkClient.createTopic(topic1, 1, 1)
-
-    produceConsumeAndShutdown(consumerGroupCommand, 1, topic1, 100)
-
-    TestUtils.produceMessages(servers, topic1, 100, acks = 1, 100 * 1000)
-
-    TestUtils.waitUntilTrue(() => {
-      val assignmentsToReset = consumerGroupCommand.resetOffsets()
-      assignmentsToReset.exists({ assignment => assignment._2.offset() == 100 })
-    }, "Expected the consumer group to reset to offset 100 (current).")
-
-    printConsumerGroup()
-    adminZkClient.deleteTopic(topic1)
+  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)
   }
 
-  private def produceConsumeAndShutdown(consumerGroupCommand: KafkaConsumerGroupService, numConsumers: Int, topic: String, totalMessages: Int) {
-    TestUtils.produceMessages(servers, topic, totalMessages, acks = 1, 100 * 1000)
-    val executor = createConsumerGroupExecutor(brokerList, numConsumers, group, topic)
-
-    TestUtils.waitUntilTrue(() => {
-      val (_, assignmentsOption) = consumerGroupCommand.describeGroup()
-      assignmentsOption match {
-        case Some(assignments) =>
-          val sumOffset = assignments.filter(_.topic.exists(_ == topic))
-            .filter(_.offset.isDefined)
-            .map(assignment => assignment.offset.get)
-            .foldLeft(0.toLong)(_ + _)
-          sumOffset == totalMessages
-        case _ => false
-      }
-    }, "Expected the consumer group to consume all messages from topic.")
-
-    executor.shutdown()
+  @Test
+  def testResetOffsetsShiftMinus(): Unit = {
+    val args = buildArgsForGroup(group, "--all-topics", "--shift-by", "-50", "--execute")
+    produceConsumeAndShutdown(topic, group, totalMessages = 100)
+    produceMessages(topic, 100)
+    resetAndAssertOffsets(args, expectedOffset = 50)
   }
 
   @Test
-  def testResetOffsetsToSpecificOffset() {
-    val cgcArgs = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--all-topics", "--to-offset", "1", "--execute")
-    val opts = new ConsumerGroupCommandOptions(cgcArgs)
-    val consumerGroupCommand = createConsumerGroupService(opts)
-
-    adminZkClient.createTopic(topic1, 1, 1)
-
-    produceConsumeAndShutdown(consumerGroupCommand, 1, topic1, 100)
-
-
-    TestUtils.waitUntilTrue(() => {
-      val assignmentsToReset = consumerGroupCommand.resetOffsets()
-      assignmentsToReset.exists({ assignment => assignment._2.offset() == 1 })
-    }, "Expected the consumer group to reset to offset 1 (specific offset).")
-
-    printConsumerGroup()
-    adminZkClient.deleteTopic(topic1)
+  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 testResetOffsetsShiftPlus() {
-    val cgcArgs = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--all-topics", "--shift-by", "50", "--execute")
-    val opts = new ConsumerGroupCommandOptions(cgcArgs)
-    val consumerGroupCommand = createConsumerGroupService(opts)
-
-    adminZkClient.createTopic(topic1, 1, 1)
-
-    produceConsumeAndShutdown(consumerGroupCommand, 1, topic1, 100)
-
-    TestUtils.produceMessages(servers, topic1, 100, acks = 1, 100 * 1000)
-
-    TestUtils.waitUntilTrue(() => {
-      val assignmentsToReset = consumerGroupCommand.resetOffsets()
-      assignmentsToReset.exists({ assignment => assignment._2.offset() == 150 })
-    }, "Expected the consumer group to reset to offset 150 (current + 50).")
-
-    printConsumerGroup()
-    adminZkClient.deleteTopic(topic1)
+  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 testResetOffsetsShiftMinus() {
-    val cgcArgs = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--all-topics", "--shift-by", "-50", "--execute")
-    val opts = new ConsumerGroupCommandOptions(cgcArgs)
-    val consumerGroupCommand = createConsumerGroupService(opts)
-
-    adminZkClient.createTopic(topic1, 1, 1)
-
-    produceConsumeAndShutdown(consumerGroupCommand, 1, topic1, 100)
-
-    TestUtils.produceMessages(servers, topic1, 100, acks = 1, 100 * 1000)
-
-
-    TestUtils.waitUntilTrue(() => {
-      val assignmentsToReset = consumerGroupCommand.resetOffsets()
-      assignmentsToReset.exists({ assignment => assignment._2.offset() == 50 })
-    }, "Expected the consumer group to reset to offset 50 (current - 50).")
-
-    printConsumerGroup()
-    adminZkClient.deleteTopic(topic1)
+  def testResetOffsetsToEarliestOnOneTopic(): Unit = {
+    val args = buildArgsForGroup(group, "--topic", topic, "--to-earliest", "--execute")
+    produceConsumeAndShutdown(topic, group, totalMessages = 100)
+    resetAndAssertOffsets(args, expectedOffset = 0)
   }
 
   @Test
-  def testResetOffsetsShiftByLowerThanEarliest() {
-    val cgcArgs = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--all-topics", "--shift-by", "-150", "--execute")
-    val opts = new ConsumerGroupCommandOptions(cgcArgs)
-    val consumerGroupCommand = createConsumerGroupService(opts)
+  def testResetOffsetsToEarliestOnOneTopicAndPartition(): Unit = {
+    val topic = "bar"
+    createTopic(topic, 2, 1)
 
-    adminZkClient.createTopic(topic1, 1, 1)
+    val args = buildArgsForGroup(group, "--topic", s"$topic:1", "--to-earliest", "--execute")
+    val consumerGroupCommand = getConsumerGroupService(args)
 
-    produceConsumeAndShutdown(consumerGroupCommand, 1, topic1, 100)
+    produceConsumeAndShutdown(topic, group, totalMessages = 100, numConsumers = 2)
+    val priorCommittedOffsets = committedOffsets(topic = topic)
 
-    TestUtils.produceMessages(servers, topic1, 100, acks = 1, 100 * 1000)
+    val tp0 = new TopicPartition(topic, 0)
+    val tp1 = new TopicPartition(topic, 1)
+    val expectedOffsets = Map(tp0 -> priorCommittedOffsets(tp0), tp1 -> 0L)
+    resetAndAssertOffsetsCommitted(consumerGroupCommand, expectedOffsets, topic)
 
-    TestUtils.waitUntilTrue(() => {
-      val assignmentsToReset = consumerGroupCommand.resetOffsets()
-      assignmentsToReset.exists({ assignment => assignment._2.offset() == 0 })
-    }, "Expected the consumer group to reset to offset 0 (earliest by shift).")
-
-    printConsumerGroup()
-    adminZkClient.deleteTopic(topic1)
+    adminZkClient.deleteTopic(topic)
   }
 
   @Test
-  def testResetOffsetsShiftByHigherThanLatest() {
-    val cgcArgs = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--all-topics", "--shift-by", "150", "--execute")
-    val opts = new ConsumerGroupCommandOptions(cgcArgs)
-    val consumerGroupCommand = createConsumerGroupService(opts)
+  def testResetOffsetsToEarliestOnTopics(): Unit = {
+    val topic1 = "topic1"
+    val topic2 = "topic2"
+    createTopic(topic1, 1, 1)
+    createTopic(topic2, 1, 1)
 
-    adminZkClient.createTopic(topic1, 1, 1)
+    val args = buildArgsForGroup(group, "--topic", topic1, "--topic", topic2, "--to-earliest", "--execute")
+    val consumerGroupCommand = getConsumerGroupService(args)
 
-    produceConsumeAndShutdown(consumerGroupCommand, 1, topic1, 100)
+    produceConsumeAndShutdown(topic1, group, 100, 1)
+    produceConsumeAndShutdown(topic2, group, 100, 1)
 
-    TestUtils.produceMessages(servers, topic1, 100, acks = 1, 100 * 1000)
+    val tp1 = new TopicPartition(topic1, 0)
+    val tp2 = new TopicPartition(topic2, 0)
 
-    TestUtils.waitUntilTrue(() => {
-      val assignmentsToReset = consumerGroupCommand.resetOffsets()
-      assignmentsToReset.exists({ assignment => assignment._2.offset() == 200 })
-    }, "Expected the consumer group to reset to offset 200 (latest by shift).")
+    val allResetOffsets = resetOffsets(consumerGroupCommand)(group).map { case (k, v) => k -> v.offset }
+    assertEquals(Map(tp1 -> 0L, tp2 -> 0L), allResetOffsets)
+    assertEquals(Map(tp1 -> 0L), committedOffsets(topic1))
+    assertEquals(Map(tp2 -> 0L), committedOffsets(topic2))
 
-    printConsumerGroup()
     adminZkClient.deleteTopic(topic1)
+    adminZkClient.deleteTopic(topic2)
   }
 
   @Test
-  def testResetOffsetsToEarliestOnOneTopic() {
-    val cgcArgs = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--topic", topic1, "--to-earliest", "--execute")
-    val opts = new ConsumerGroupCommandOptions(cgcArgs)
-    val consumerGroupCommand = createConsumerGroupService(opts)
+  def testResetOffsetsToEarliestOnTopicsAndPartitions(): Unit = {
+    val topic1 = "topic1"
+    val topic2 = "topic2"
 
-    adminZkClient.createTopic(topic1, 1, 1)
+    createTopic(topic1, 2, 1)
+    createTopic(topic2, 2, 1)
 
-    produceConsumeAndShutdown(consumerGroupCommand, 1, topic1, 100)
+    val args = buildArgsForGroup(group, "--topic", s"$topic1:1", "--topic", s"$topic2:1", "--to-earliest", "--execute")
+    val consumerGroupCommand = getConsumerGroupService(args)
 
-    TestUtils.waitUntilTrue(() => {
-      val assignmentsToReset = consumerGroupCommand.resetOffsets()
-      assignmentsToReset.exists { assignment => assignment._2.offset() == 0 }
-    }, "Expected the consumer group to reset to offset 0 (earliest).")
+    produceConsumeAndShutdown(topic1, group, 100, 2)
+    produceConsumeAndShutdown(topic2, group, 100, 2)
 
-    printConsumerGroup()
-    adminZkClient.deleteTopic(topic1)
-  }
+    val priorCommittedOffsets1 = committedOffsets(topic1)
+    val priorCommittedOffsets2 = committedOffsets(topic2)
 
-  @Test
-  def testResetOffsetsToEarliestOnOneTopicAndPartition() {
-    val cgcArgs = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--topic", String.format("%s:1", topic1), "--to-earliest", "--execute")
-    val opts = new ConsumerGroupCommandOptions(cgcArgs)
-    val consumerGroupCommand = createConsumerGroupService(opts)
-
-    adminZkClient.createTopic(topic1, 2, 1)
-
-    produceConsumeAndShutdown(consumerGroupCommand, 2, topic1, 100)
+    val tp1 = new TopicPartition(topic1, 1)
+    val tp2 = new TopicPartition(topic2, 1)
+    val allResetOffsets = resetOffsets(consumerGroupCommand)(group).map { case (k, v) => k -> v.offset }
+    assertEquals(Map(tp1 -> 0, tp2 -> 0), allResetOffsets)
 
-    TestUtils.waitUntilTrue(() => {
-      val assignmentsToReset = consumerGroupCommand.resetOffsets()
-      assignmentsToReset.exists { assignment => assignment._2.offset() == 0 && assignment._1.partition() == 1 }
-    }, "Expected the consumer group to reset to offset 0 (earliest) in partition 1.")
+    assertEquals(priorCommittedOffsets1.toMap + (tp1 -> 0L), committedOffsets(topic1))
+    assertEquals(priorCommittedOffsets2.toMap + (tp2 -> 0L), committedOffsets(topic2))
 
-    printConsumerGroup()
     adminZkClient.deleteTopic(topic1)
+    adminZkClient.deleteTopic(topic2)
   }
 
   @Test
-  def testResetOffsetsToEarliestOnTopics() {
-    val cgcArgs = Array("--bootstrap-server", brokerList, "--reset-offsets",
-      "--group", group,
-      "--topic", topic1,
-      "--topic", topic2,
-      "--to-earliest", "--execute")
-    val opts = new ConsumerGroupCommandOptions(cgcArgs)
-    val consumerGroupCommand = createConsumerGroupService(opts)
-
-    adminZkClient.createTopic(topic1, 1, 1)
-    adminZkClient.createTopic(topic2, 1, 1)
+  // 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)
 
-    produceConsumeAndShutdown(consumerGroupCommand, 1, topic1, 100)
-    produceConsumeAndShutdown(consumerGroupCommand, 1, topic2, 100)
+    val cgcArgs = buildArgsForGroup(group, "--all-topics", "--to-offset", "2", "--export")
+    val consumerGroupCommand = getConsumerGroupService(cgcArgs)
 
-    TestUtils.waitUntilTrue(() => {
-      val assignmentsToReset = consumerGroupCommand.resetOffsets()
-      assignmentsToReset.exists { assignment => assignment._2.offset() == 0 && assignment._1.topic() == topic1 } &&
-        assignmentsToReset.exists { assignment => assignment._2.offset() == 0 && assignment._1.topic() == topic2 }
-    }, "Expected the consumer group to reset to offset 0 (earliest).")
-
-    printConsumerGroup()
-    adminZkClient.deleteTopic(topic1)
-    adminZkClient.deleteTopic(topic2)
-  }
+    produceConsumeAndShutdown(topic = topic, group = group, totalMessages = 100, numConsumers = 2)
 
-  @Test
-  def testResetOffsetsToEarliestOnTopicsAndPartitions() {
-    val cgcArgs = Array("--bootstrap-server", brokerList, "--reset-offsets",
-      "--group", group,
-      "--topic", String.format("%s:1", topic1),
-      "--topic", String.format("%s:1", topic2),
-      "--to-earliest", "--execute")
-    val opts = new ConsumerGroupCommandOptions(cgcArgs)
-    val consumerGroupCommand = createConsumerGroupService(opts)
-
-    adminZkClient.createTopic(topic1, 2, 1)
-    adminZkClient.createTopic(topic2, 2, 1)
+    val file = File.createTempFile("reset", ".csv")
+    file.deleteOnExit()
 
-    produceConsumeAndShutdown(consumerGroupCommand, 2, topic1, 100)
-    produceConsumeAndShutdown(consumerGroupCommand, 2, topic2, 100)
+    val exportedOffsets = consumerGroupCommand.resetOffsets()
+    val bw = new BufferedWriter(new FileWriter(file))
+    bw.write(consumerGroupCommand.exportOffsetsToCsv(exportedOffsets))
+    bw.close()
+    assertEquals(Map(tp0 -> 2L, tp1 -> 2L), exportedOffsets(group).map { case (k, v) => k -> v.offset })
 
-    TestUtils.waitUntilTrue(() => {
-      val assignmentsToReset = consumerGroupCommand.resetOffsets()
-      assignmentsToReset.exists { assignment => assignment._2.offset() == 0 && assignment._1.partition() == 1 && assignment._1.topic() == topic1 }
-      assignmentsToReset.exists { assignment => assignment._2.offset() == 0 && assignment._1.partition() == 1 && assignment._1.topic() == topic2 }
-    }, "Expected the consumer group to reset to offset 0 (earliest) in partition 1.")
+    val cgcArgsExec = buildArgsForGroup(group, "--all-topics", "--from-file", file.getCanonicalPath, "--dry-run")
+    val consumerGroupCommandExec = getConsumerGroupService(cgcArgsExec)
+    val importedOffsets = consumerGroupCommandExec.resetOffsets()
+    assertEquals(Map(tp0 -> 2L, tp1 -> 2L), importedOffsets(group).map { case (k, v) => k -> v.offset })
 
-    printConsumerGroup()
-    adminZkClient.deleteTopic(topic1)
-    adminZkClient.deleteTopic(topic2)
+    adminZkClient.deleteTopic(topic)
   }
 
   @Test
-  def testResetOffsetsExportImportPlan() {
-    val cgcArgs = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--all-topics", "--to-offset","2", "--export")
-    val opts = new ConsumerGroupCommandOptions(cgcArgs)
-    val consumerGroupCommand = createConsumerGroupService(opts)
-
-    adminZkClient.createTopic(topic1, 2, 1)
-
-    produceConsumeAndShutdown(consumerGroupCommand, 2, topic1, 100)
+  // This one deals with universal CSV export/import file format "group,topic,partition,offset",
+  // supporting multiple --group args or --all-groups arg
+  def testResetOffsetsExportImportPlan(): Unit = {
+    val group1 = group + "1"
+    val group2 = group + "2"
+    val topic1 = "bar1"
+    val topic2 = "bar2"
+    val t1p0 = new TopicPartition(topic1, 0)
+    val t1p1 = new TopicPartition(topic1, 1)
+    val t2p0 = new TopicPartition(topic2, 0)
+    val t2p1 = new TopicPartition(topic2, 1)
+    createTopic(topic1, 2, 1)
+    createTopic(topic2, 2, 1)
+
+    val cgcArgs = buildArgsForGroups(Seq(group1, group2), "--all-topics", "--to-offset", "2", "--export")
+    val consumerGroupCommand = getConsumerGroupService(cgcArgs)
+
+    produceConsumeAndShutdown(topic = topic1, group = group1, totalMessages = 100)
+    produceConsumeAndShutdown(topic = topic2, group = group2, totalMessages = 100)
+
+    awaitConsumerGroupInactive(consumerGroupCommand, group1)
+    awaitConsumerGroupInactive(consumerGroupCommand, group2)
 
     val file = File.createTempFile("reset", ".csv")
+    file.deleteOnExit()
 
-    TestUtils.waitUntilTrue(() => {
-      val assignmentsToReset = consumerGroupCommand.resetOffsets()
-      val bw = new BufferedWriter(new FileWriter(file))
-      bw.write(consumerGroupCommand.exportOffsetsToReset(assignmentsToReset))
-      bw.close()
-      assignmentsToReset.exists { assignment => assignment._2.offset() == 2 } && file.exists()
-    }, "Expected the consume all messages and save reset offsets plan to file")
+    val exportedOffsets = consumerGroupCommand.resetOffsets()
+    val bw = new BufferedWriter(new FileWriter(file))
+    bw.write(consumerGroupCommand.exportOffsetsToCsv(exportedOffsets))
+    bw.close()
+    assertEquals(Map(t1p0 -> 2L, t1p1 -> 2L), exportedOffsets(group1).map { case (k, v) => k -> v.offset })
+    assertEquals(Map(t2p0 -> 2L, t2p1 -> 2L), exportedOffsets(group2).map { case (k, v) => k -> v.offset })
+
+    // Multiple --group's offset import
+    val cgcArgsExec = buildArgsForGroups(Seq(group1, group2), "--all-topics", "--from-file", file.getCanonicalPath, "--dry-run")
+    val consumerGroupCommandExec = getConsumerGroupService(cgcArgsExec)
+    val importedOffsets = consumerGroupCommandExec.resetOffsets()
+    assertEquals(Map(t1p0 -> 2L, t1p1 -> 2L), importedOffsets(group1).map { case (k, v) => k -> v.offset })
+    assertEquals(Map(t2p0 -> 2L, t2p1 -> 2L), importedOffsets(group2).map { case (k, v) => k -> v.offset })
+
+    // Single --group offset import using "group,topic,partition,offset" csv format
+    val cgcArgsExec2 = buildArgsForGroup(group1, "--all-topics", "--from-file", file.getCanonicalPath, "--dry-run")
+    val consumerGroupCommandExec2 = getConsumerGroupService(cgcArgsExec2)
+    val importedOffsets2 = consumerGroupCommandExec2.resetOffsets()
+    assertEquals(Map(t1p0 -> 2L, t1p1 -> 2L), importedOffsets2(group1).map { case (k, v) => k -> v.offset })
+
+    adminZkClient.deleteTopic(topic)
+  }
 
+  @Test(expected = classOf[OptionException])
+  def testResetWithUnrecognizedNewConsumerOption(): Unit = {
+    val cgcArgs = Array("--new-consumer", "--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--all-topics",
+      "--to-offset", "2", "--export")
+    getConsumerGroupService(cgcArgs)
+  }
 
-    val cgcArgsExec = Array("--bootstrap-server", brokerList, "--reset-offsets", "--group", group, "--all-topics", "--from-file", file.getCanonicalPath)
-    val optsExec = new ConsumerGroupCommandOptions(cgcArgsExec)
-    val consumerGroupCommandExec = createConsumerGroupService(optsExec)
+  private def produceMessages(topic: String, numMessages: Int): Unit = {
+    val records = (0 until numMessages).map(_ => new ProducerRecord[Array[Byte], Array[Byte]](topic,
+      new Array[Byte](100 * 1000)))
+    TestUtils.produceMessages(servers, records, acks = 1)
+  }
 
-    TestUtils.waitUntilTrue(() => {
-        val assignmentsToReset = consumerGroupCommandExec.resetOffsets()
-        assignmentsToReset.exists { assignment => assignment._2.offset() == 2 }
-    }, "Expected the consumer group to reset to offset 2 according to the plan in the file.")
+  private def produceConsumeAndShutdown(topic: String, group: String, totalMessages: Int, numConsumers: Int = 1): Unit = {
+    produceMessages(topic, totalMessages)
+    val executor = addConsumerGroupExecutor(numConsumers = numConsumers, topic = topic, group = group)
+    awaitConsumerProgress(topic, group, totalMessages)
+    executor.shutdown()
+  }
 
-    file.deleteOnExit()
+  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()
+    }
 
-    printConsumerGroup()
-    adminZkClient.deleteTopic(topic1)
   }
 
-  private def printConsumerGroup() {
-    val cgcArgs = Array("--bootstrap-server", brokerList, "--group", group, "--describe")
-    ConsumerGroupCommand.main(cgcArgs)
+  private def awaitConsumerGroupInactive(consumerGroupService: ConsumerGroupService, group: String): Unit = {
+    TestUtils.waitUntilTrue(() => {
+      val state = consumerGroupService.collectGroupState(group).state
+      state == "Empty" || state == "Dead"
+    }, s"Expected that consumer group is inactive. Actual state: ${consumerGroupService.collectGroupState(group).state}")
   }
 
-  private def printConsumerGroup(group: String) {
-    val cgcArgs = Array("--bootstrap-server", brokerList, "--group", group, "--describe")
-    ConsumerGroupCommand.main(cgcArgs)
+  private def resetAndAssertOffsets(args: Array[String],
+                                    expectedOffset: Long,
+                                    dryRun: Boolean = false,
+                                    topics: Seq[String] = Seq(topic)): Unit = {
+    val consumerGroupCommand = getConsumerGroupService(args)
+    val expectedOffsets = topics.map(topic => topic -> Map(new TopicPartition(topic, 0) -> expectedOffset)).toMap
+    val resetOffsetsResultByGroup = resetOffsets(consumerGroupCommand)
+
+    try {
+      for {
+        topic <- topics
+        (group, partitionInfo) <- resetOffsetsResultByGroup
+      } {
+        val priorOffsets = committedOffsets(topic = topic, group = group)
+        assertEquals(expectedOffsets(topic),
+          partitionInfo.filter(partitionInfo => partitionInfo._1.topic() == topic).map { case (k, v) => k -> v.offset })
+        assertEquals(if (dryRun) priorOffsets else expectedOffsets(topic), committedOffsets(topic = topic, group = group))
+      }
+    } finally {
+      consumerGroupCommand.close()
+    }
   }
 
-  private def createConsumerGroupExecutor(brokerList: String, numConsumers: Int, groupId: String, topic: String): ConsumerGroupExecutor = {
-    val executor = new ConsumerGroupExecutor(brokerList, numConsumers, groupId, topic)
-    executors += executor
-    executor
+  private def resetAndAssertOffsetsCommitted(consumerGroupService: ConsumerGroupService,
+                                             expectedOffsets: Map[TopicPartition, Long],
+                                             topic: String): Unit = {
+    val allResetOffsets = resetOffsets(consumerGroupService)
+    for {
+      (group, offsetsInfo) <- allResetOffsets
+      (tp, offsetMetadata) <- offsetsInfo
+    } {
+      assertEquals(offsetMetadata.offset(), expectedOffsets(tp))
+      assertEquals(expectedOffsets, committedOffsets(topic, group))
+    }
   }
 
-  private def createConsumerGroupService(opts: ConsumerGroupCommandOptions): KafkaConsumerGroupService = {
-    val service = new KafkaConsumerGroupService(opts)
-    consumerGroupServices += service
-    service
+  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 55fe587f74db0..0000000000000
--- a/core/src/test/scala/unit/kafka/admin/TestAdminUtils.scala
+++ /dev/null
@@ -1,28 +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
-
-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 180f257204cf8..1a117222a7a56 100644
--- a/core/src/test/scala/unit/kafka/admin/TopicCommandTest.scala
+++ b/core/src/test/scala/unit/kafka/admin/TopicCommandTest.scala
@@ -16,213 +16,123 @@
  */
 package kafka.admin
 
+import kafka.admin.TopicCommand.{PartitionDescription, TopicCommandOptions}
+import kafka.common.AdminCommandFailedException
+import kafka.utils.Exit
+import org.apache.kafka.clients.admin.PartitionReassignment
+import org.apache.kafka.common.Node
+import org.apache.kafka.common.TopicPartitionInfo
 import org.junit.Assert._
 import org.junit.Test
-import kafka.utils.Logging
-import kafka.utils.TestUtils
-import kafka.zk.ZooKeeperTestHarness
-import kafka.server.ConfigType
-import kafka.admin.TopicCommand.TopicCommandOptions
-import kafka.utils.ZkUtils.ConfigChangesPath
-import kafka.utils.ZkUtils.getDeleteTopicPath
-import org.apache.kafka.common.errors.TopicExistsException
-import org.apache.kafka.common.internals.Topic
+import org.scalatest.Assertions.intercept
 
-class TopicCommandTest extends ZooKeeperTestHarness with Logging with RackAwareTest {
+import scala.jdk.CollectionConverters._
 
-  @Test
-  def testConfigPreservationAcrossPartitionAlteration() {
-    val topic = "test"
-    val numPartitionsOriginal = 1
-    val cleanupKey = "cleanup.policy"
-    val cleanupVal = "compact"
-    // create brokers
-    val brokers = List(0, 1, 2)
-    TestUtils.createBrokersInZk(zkUtils, brokers)
-    // create the topic
-    val createOpts = new TopicCommandOptions(Array("--partitions", numPartitionsOriginal.toString,
-      "--replication-factor", "1",
-      "--config", cleanupKey + "=" + cleanupVal,
-      "--topic", topic))
-    TopicCommand.createTopic(zkClient, createOpts)
-    val props = adminZkClient.fetchEntityConfig(ConfigType.Topic, topic)
-    assertTrue("Properties after creation don't contain " + cleanupKey, props.containsKey(cleanupKey))
-    assertTrue("Properties after creation have incorrect value", props.getProperty(cleanupKey).equals(cleanupVal))
-
-    // pre-create the topic config changes path to avoid a NoNodeException
-    zkUtils.createPersistentPath(ConfigChangesPath)
+class TopicCommandTest {
 
-    // modify the topic to add new partitions
-    val numPartitionsModified = 3
-    val alterOpts = new TopicCommandOptions(Array("--partitions", numPartitionsModified.toString, "--topic", topic))
-    TopicCommand.alterTopic(zkClient, alterOpts)
-    val newProps = adminZkClient.fetchEntityConfig(ConfigType.Topic, topic)
-    assertTrue("Updated properties do not contain " + cleanupKey, newProps.containsKey(cleanupKey))
-    assertTrue("Updated properties have incorrect value", newProps.getProperty(cleanupKey).equals(cleanupVal))
-  }
+  private[this] val brokerList = "localhost:9092"
+  private[this] val topicName = "topicName"
 
   @Test
-  def testTopicDeletion() {
-    val normalTopic = "test"
-
-    val numPartitionsOriginal = 1
-
-    // create brokers
-    val brokers = List(0, 1, 2)
-    TestUtils.createBrokersInZk(zkUtils, brokers)
-
-    // create the NormalTopic
-    val createOpts = new TopicCommandOptions(Array("--partitions", numPartitionsOriginal.toString,
-      "--replication-factor", "1",
-      "--topic", normalTopic))
-    TopicCommand.createTopic(zkClient, createOpts)
-
-    // delete the NormalTopic
-    val deleteOpts = new TopicCommandOptions(Array("--topic", normalTopic))
-    val deletePath = getDeleteTopicPath(normalTopic)
-    assertFalse("Delete path for topic shouldn't exist before deletion.", zkUtils.pathExists(deletePath))
-    TopicCommand.deleteTopic(zkClient, deleteOpts)
-    assertTrue("Delete path for topic should exist after deletion.", zkUtils.pathExists(deletePath))
-
-    // create the offset topic
-    val createOffsetTopicOpts = new TopicCommandOptions(Array("--partitions", numPartitionsOriginal.toString,
-      "--replication-factor", "1",
-      "--topic", Topic.GROUP_METADATA_TOPIC_NAME))
-    TopicCommand.createTopic(zkClient, createOffsetTopicOpts)
-
-    // try to delete the Topic.GROUP_METADATA_TOPIC_NAME and make sure it doesn't
-    val deleteOffsetTopicOpts = new TopicCommandOptions(Array("--topic", Topic.GROUP_METADATA_TOPIC_NAME))
-    val deleteOffsetTopicPath = getDeleteTopicPath(Topic.GROUP_METADATA_TOPIC_NAME)
-    assertFalse("Delete path for topic shouldn't exist before deletion.", zkUtils.pathExists(deleteOffsetTopicPath))
-    intercept[AdminOperationException] {
-      TopicCommand.deleteTopic(zkClient, deleteOffsetTopicOpts)
+  def testIsNotUnderReplicatedWhenAdding(): Unit = {
+    val replicaIds = List(1, 2)
+    val replicas = replicaIds.map { id =>
+      new Node(id, "localhost", 9090 + id)
     }
-    assertFalse("Delete path for topic shouldn't exist after deletion.", zkUtils.pathExists(deleteOffsetTopicPath))
+
+    val partitionDescription = PartitionDescription(
+      "test-topic",
+      new TopicPartitionInfo(
+        0,
+        new Node(1, "localhost", 9091),
+        replicas.asJava,
+        List(new Node(1, "localhost", 9091)).asJava
+      ),
+      None,
+      markedForDeletion = false,
+      Some(
+        new PartitionReassignment(
+          replicaIds.map(id => id: java.lang.Integer).asJava,
+          List(2: java.lang.Integer).asJava,
+          List.empty.asJava
+        )
+      )
+    )
+
+    assertFalse(partitionDescription.isUnderReplicated)
   }
 
   @Test
-  def testDeleteIfExists() {
-    // create brokers
-    val brokers = List(0, 1, 2)
-    TestUtils.createBrokersInZk(zkUtils, brokers)
-
-    // delete a topic that does not exist without --if-exists
-    val deleteOpts = new TopicCommandOptions(Array("--topic", "test"))
-    intercept[IllegalArgumentException] {
-      TopicCommand.deleteTopic(zkClient, deleteOpts)
-    }
-
-    // delete a topic that does not exist with --if-exists
-    val deleteExistsOpts = new TopicCommandOptions(Array("--topic", "test", "--if-exists"))
-    TopicCommand.deleteTopic(zkClient, deleteExistsOpts)
+  def testAlterWithUnspecifiedPartitionCount(): Unit = {
+    assertCheckArgsExitCode(1, new TopicCommandOptions(
+      Array("--bootstrap-server", brokerList ,"--alter", "--topic", topicName)))
   }
 
   @Test
-  def testAlterIfExists() {
-    // create brokers
-    val brokers = List(0, 1, 2)
-    TestUtils.createBrokersInZk(zkUtils, brokers)
-
-    // alter a topic that does not exist without --if-exists
-    val alterOpts = new TopicCommandOptions(Array("--topic", "test", "--partitions", "1"))
-    intercept[IllegalArgumentException] {
-      TopicCommand.alterTopic(zkClient, alterOpts)
-    }
-
-    // alter a topic that does not exist with --if-exists
-    val alterExistsOpts = new TopicCommandOptions(Array("--topic", "test", "--partitions", "1", "--if-exists"))
-    TopicCommand.alterTopic(zkClient, alterExistsOpts)
+  def testConfigOptWithBootstrapServers(): Unit = {
+    assertCheckArgsExitCode(1,
+      new TopicCommandOptions(Array("--bootstrap-server", brokerList ,"--alter", "--topic", topicName, "--partitions", "3", "--config", "cleanup.policy=compact")))
+    assertCheckArgsExitCode(1,
+      new TopicCommandOptions(Array("--bootstrap-server", brokerList ,"--alter", "--topic", topicName, "--partitions", "3", "--delete-config", "cleanup.policy")))
+    val opts =
+      new TopicCommandOptions(Array("--bootstrap-server", brokerList ,"--create", "--topic", topicName, "--partitions", "3", "--replication-factor", "3", "--config", "cleanup.policy=compact"))
+    opts.checkArgs()
+    assertTrue(opts.hasCreateOption)
+    assertEquals(brokerList, opts.bootstrapServer.get)
+    assertEquals("cleanup.policy=compact", opts.topicConfig.get.get(0))
   }
 
   @Test
-  def testCreateIfNotExists() {
-    // create brokers
-    val brokers = List(0, 1, 2)
-    TestUtils.createBrokersInZk(zkUtils, brokers)
-
-    val topic = "test"
-    val numPartitions = 1
-
-    // create the topic
-    val createOpts = new TopicCommandOptions(
-      Array("--partitions", numPartitions.toString, "--replication-factor", "1", "--topic", topic))
-    TopicCommand.createTopic(zkClient, createOpts)
-
-    // try to re-create the topic without --if-not-exists
-    intercept[TopicExistsException] {
-      TopicCommand.createTopic(zkClient, createOpts)
-    }
-
-    // try to re-create the topic with --if-not-exists
-    val createNotExistsOpts = new TopicCommandOptions(
-      Array("--partitions", numPartitions.toString, "--replication-factor", "1", "--topic", topic, "--if-not-exists"))
-    TopicCommand.createTopic(zkClient, createNotExistsOpts)
+  def testCreateWithAssignmentAndPartitionCount(): Unit = {
+    assertCheckArgsExitCode(1,
+      new TopicCommandOptions(
+        Array("--bootstrap-server", brokerList,
+          "--create",
+          "--replica-assignment", "3:0,5:1",
+          "--partitions", "2",
+          "--topic", "testTopic")))
   }
 
   @Test
-  def testCreateAlterTopicWithRackAware() {
-    val rackInfo = Map(0 -> "rack1", 1 -> "rack2", 2 -> "rack2", 3 -> "rack1", 4 -> "rack3", 5 -> "rack3")
-    TestUtils.createBrokersInZk(toBrokerMetadata(rackInfo), zkUtils)
-
-    val numPartitions = 18
-    val replicationFactor = 3
-    val createOpts = new TopicCommandOptions(Array(
-      "--partitions", numPartitions.toString,
-      "--replication-factor", replicationFactor.toString,
-      "--topic", "foo"))
-    TopicCommand.createTopic(zkClient, createOpts)
-
-    var assignment = zkUtils.getReplicaAssignmentForTopics(Seq("foo")).map { case (tp, replicas) =>
-      tp.partition -> replicas
-    }
-    checkReplicaDistribution(assignment, rackInfo, rackInfo.size, numPartitions, replicationFactor)
+  def testCreateWithAssignmentAndReplicationFactor(): Unit = {
+    assertCheckArgsExitCode(1,
+      new TopicCommandOptions(
+        Array("--bootstrap-server", brokerList,
+          "--create",
+          "--replica-assignment", "3:0,5:1",
+          "--replication-factor", "2",
+          "--topic", "testTopic")))
+  }
 
-    val alteredNumPartitions = 36
-    // verify that adding partitions will also be rack aware
-    val alterOpts = new TopicCommandOptions(Array(
-      "--partitions", alteredNumPartitions.toString,
-      "--topic", "foo"))
-    TopicCommand.alterTopic(zkClient, alterOpts)
-    assignment = zkUtils.getReplicaAssignmentForTopics(Seq("foo")).map { case (tp, replicas) =>
-      tp.partition -> replicas
+  @Test
+  def testParseAssignmentDuplicateEntries(): Unit = {
+    intercept[AdminCommandFailedException] {
+      TopicCommand.parseReplicaAssignment("5:5")
     }
-    checkReplicaDistribution(assignment, rackInfo, rackInfo.size, alteredNumPartitions, replicationFactor)
   }
 
   @Test
-  def testDescribeAndListTopicsMarkedForDeletion() {
-    val brokers = List(0)
-    val topic = "testtopic"
-    val markedForDeletionDescribe = "MarkedForDeletion"
-    val markedForDeletionList = "marked for deletion"
-    TestUtils.createBrokersInZk(zkUtils, brokers)
-
-    val createOpts = new TopicCommandOptions(Array("--partitions", "1", "--replication-factor", "1", "--topic", topic))
-    TopicCommand.createTopic(zkClient, createOpts)
-
-    // delete the broker first, so when we attempt to delete the topic it gets into "marked for deletion"
-    TestUtils.deleteBrokersInZk(zkUtils, brokers)
-    TopicCommand.deleteTopic(zkClient, new TopicCommandOptions(Array("--topic", topic)))
-
-    // Test describe topics
-    def describeTopicsWithConfig() {
-      TopicCommand.describeTopic(zkClient, new TopicCommandOptions(Array("--describe")))
+  def testParseAssignmentPartitionsOfDifferentSize(): Unit = {
+    intercept[AdminOperationException] {
+      TopicCommand.parseReplicaAssignment("5:4:3,2:1")
     }
-    val outputWithConfig = TestUtils.grabConsoleOutput(describeTopicsWithConfig)
-    assertTrue(outputWithConfig.contains(topic) && outputWithConfig.contains(markedForDeletionDescribe))
+  }
 
-    def describeTopicsNoConfig() {
-      TopicCommand.describeTopic(zkClient, new TopicCommandOptions(Array("--describe", "--unavailable-partitions")))
-    }
-    val outputNoConfig = TestUtils.grabConsoleOutput(describeTopicsNoConfig)
-    assertTrue(outputNoConfig.contains(topic) && outputNoConfig.contains(markedForDeletionDescribe))
+  @Test
+  def testParseAssignment(): Unit = {
+    val actualAssignment = TopicCommand.parseReplicaAssignment("5:4,3:2,1:0")
+    val expectedAssignment = Map(0 -> List(5, 4), 1 -> List(3, 2), 2 -> List(1, 0))
+    assertEquals(expectedAssignment, actualAssignment)
+  }
 
-    // Test list topics
-    def listTopics() {
-      TopicCommand.listTopics(zkClient, new TopicCommandOptions(Array("--list")))
+  private[this] def assertCheckArgsExitCode(expected: Int, options: TopicCommandOptions): Unit = {
+    Exit.setExitProcedure {
+      (exitCode: Int, _: Option[String]) =>
+        assertEquals(expected, exitCode)
+        throw new RuntimeException
     }
-    val output = TestUtils.grabConsoleOutput(listTopics)
-    assertTrue(output.contains(topic) && output.contains(markedForDeletionList))
+    try intercept[RuntimeException] {
+      options.checkArgs()
+    } finally Exit.resetExitProcedure()
   }
-
 }
diff --git a/core/src/test/scala/unit/kafka/admin/TopicCommandWithAdminClientTest.scala b/core/src/test/scala/unit/kafka/admin/TopicCommandWithAdminClientTest.scala
new file mode 100644
index 0000000000000..2b19f47ca2911
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/admin/TopicCommandWithAdminClientTest.scala
@@ -0,0 +1,835 @@
+/**
+  * Licensed to the Apache Software Foundation (ASF) under one or more
+  * contributor license agreements.  See the NOTICE file distributed with
+  * this work for additional information regarding copyright ownership.
+  * The ASF licenses this file to You under the Apache License, Version 2.0
+  * (the "License"); you may not use this file except in compliance with
+  * the License.  You may obtain a copy of the License at
+  *
+  *    http://www.apache.org/licenses/LICENSE-2.0
+  *
+  * Unless required by applicable law or agreed to in writing, software
+  * distributed under the License is distributed on an "AS IS" BASIS,
+  * WITHOUT WARRANTIES 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.{Collections, Optional, Properties}
+
+import kafka.admin.TopicCommand.{AdminClientTopicService, TopicCommandOptions}
+import kafka.integration.KafkaServerTestHarness
+import kafka.server.{ConfigType, KafkaConfig}
+import kafka.utils.{Logging, TestUtils}
+import kafka.zk.{ConfigEntityChangeNotificationZNode, DeleteTopicsTopicZNode}
+import org.apache.kafka.clients.CommonClientConfigs
+import org.apache.kafka.clients.admin._
+import org.apache.kafka.common.Node
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.TopicPartitionInfo
+import org.apache.kafka.common.config.{ConfigException, ConfigResource, TopicConfig}
+import org.apache.kafka.common.errors.ClusterAuthorizationException
+import org.apache.kafka.common.errors.ThrottlingQuotaExceededException
+import org.apache.kafka.common.errors.TopicExistsException
+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.assertThrows
+import org.junit.Assert.{assertEquals, assertFalse, assertTrue}
+import org.junit.rules.TestName
+import org.junit.{After, Before, Rule, Test}
+import org.mockito.ArgumentMatcher
+import org.mockito.ArgumentMatchers.{eq => eqThat, _}
+import org.mockito.Mockito._
+import org.scalatest.Assertions.{fail, intercept}
+
+import scala.jdk.CollectionConverters._
+import scala.collection.Seq
+import scala.concurrent.ExecutionException
+import scala.util.Random
+
+class TopicCommandWithAdminClientTest extends KafkaServerTestHarness with Logging with RackAwareTest {
+
+  /**
+    * Implementations must override this method to return a set of KafkaConfigs. This method will be invoked for every
+    * test and should not reuse previous configurations unless they select their ports randomly when servers are started.
+    *
+    * Note the replica fetch max bytes is set to `1` in order to throttle the rate of replication for test
+    * `testDescribeUnderReplicatedPartitionsWhenReassignmentIsInProgress`.
+    */
+  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"),
+    numPartitions = numPartitions,
+    defaultReplicationFactor = defaultReplicationFactor,
+  ).map { props =>
+    props.put(KafkaConfig.ReplicaFetchMaxBytesProp, "1")
+    KafkaConfig.fromProps(props)
+  }
+
+  private val numPartitions = 1
+  private val defaultReplicationFactor = 1.toShort
+
+  private var topicService: AdminClientTopicService = _
+  private var adminClient: Admin = _
+  private var testTopicName: String = _
+
+  private val _testName = new TestName
+  @Rule def testName: TestName = _testName
+
+  private[this] def createAndWaitTopic(opts: TopicCommandOptions): Unit = {
+    topicService.createTopic(opts)
+    waitForTopicCreated(opts.topic.get)
+  }
+
+  private[this] def waitForTopicCreated(topicName: String, timeout: Int = 10000): Unit = {
+    TestUtils.waitUntilMetadataIsPropagated(servers, topicName, partition = 0, timeout)
+  }
+
+  @Before
+  def setup(): Unit = {
+    // create adminClient
+    val props = new Properties()
+    props.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, brokerList)
+    adminClient = Admin.create(props)
+    topicService = AdminClientTopicService(adminClient)
+    testTopicName = s"${testName.getMethodName}-${Random.alphanumeric.take(10).mkString}"
+  }
+
+  @After
+  def close(): Unit = {
+    // adminClient is closed by topicService
+    if (topicService != null)
+      topicService.close()
+  }
+
+  @Test
+  def testCreate(): Unit = {
+    createAndWaitTopic(new TopicCommandOptions(
+      Array("--partitions", "2", "--replication-factor", "1", "--topic", testTopicName)))
+
+    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)
+    createAndWaitTopic(new TopicCommandOptions(
+      Array("--partitions", "2", "--replication-factor", "2", "--topic", testTopicName, "--config", "delete.retention.ms=1000")))
+
+    val configs = adminClient
+      .describeConfigs(Collections.singleton(configResource))
+      .all().get().get(configResource)
+    assertEquals(1000, Integer.valueOf(configs.get("delete.retention.ms").value()))
+  }
+
+  @Test
+  def testCreateWhenAlreadyExists(): Unit = {
+    val numPartitions = 1
+
+    // create the topic
+    val createOpts = new TopicCommandOptions(
+      Array("--partitions", numPartitions.toString, "--replication-factor", "1", "--topic", testTopicName))
+    createAndWaitTopic(createOpts)
+
+    // try to re-create the topic
+    intercept[TopicExistsException] {
+      topicService.createTopic(createOpts)
+    }
+  }
+
+  @Test
+  def testCreateWhenAlreadyExistsWithIfNotExists(): Unit = {
+    val createOpts = new TopicCommandOptions(Array("--topic", testTopicName, "--if-not-exists"))
+    createAndWaitTopic(createOpts)
+    topicService.createTopic(createOpts)
+  }
+
+  @Test
+  def testCreateWithReplicaAssignment(): Unit = {
+    // create the topic
+    val createOpts = new TopicCommandOptions(
+      Array("--replica-assignment", "5:4,3:2,1:0", "--topic", testTopicName))
+    createAndWaitTopic(createOpts)
+
+    val partitions = adminClient
+      .describeTopics(Collections.singletonList(testTopicName))
+      .all()
+      .get()
+      .get(testTopicName)
+      .partitions()
+    assertEquals(3, partitions.size())
+    assertEquals(List(5, 4), partitions.get(0).replicas().asScala.map(_.id()))
+    assertEquals(List(3, 2), partitions.get(1).replicas().asScala.map(_.id()))
+    assertEquals(List(1, 0), partitions.get(2).replicas().asScala.map(_.id()))
+  }
+
+  @Test
+  def testCreateWithInvalidReplicationFactor(): Unit = {
+    intercept[IllegalArgumentException] {
+      topicService.createTopic(new TopicCommandOptions(
+        Array("--partitions", "2", "--replication-factor", (Short.MaxValue+1).toString, "--topic", testTopicName)))
+    }
+  }
+
+  @Test
+  def testCreateWithNegativeReplicationFactor(): Unit = {
+    intercept[IllegalArgumentException] {
+      topicService.createTopic(new TopicCommandOptions(
+        Array("--partitions", "2", "--replication-factor", "-1", "--topic", testTopicName)))
+    }
+  }
+
+  @Test
+  def testCreateWithNegativePartitionCount(): Unit = {
+    intercept[IllegalArgumentException] {
+      topicService.createTopic(new TopicCommandOptions(
+        Array("--partitions", "-1", "--replication-factor", "1", "--topic", testTopicName)))
+    }
+  }
+
+  @Test
+  def testInvalidTopicLevelConfig(): Unit = {
+    val createOpts = new TopicCommandOptions(
+      Array("--partitions", "1", "--replication-factor", "1", "--topic", testTopicName,
+        "--config", "message.timestamp.type=boom"))
+    intercept[ConfigException] {
+      topicService.createTopic(createOpts)
+    }
+  }
+
+  @Test
+  def testListTopics(): Unit = {
+    createAndWaitTopic(new TopicCommandOptions(
+      Array("--partitions", "1", "--replication-factor", "1", "--topic", testTopicName)))
+
+    val output = TestUtils.grabConsoleOutput(
+      topicService.listTopics(new TopicCommandOptions(Array())))
+
+    assertTrue(output.contains(testTopicName))
+  }
+
+  @Test
+  def testListTopicsWithIncludeList(): Unit = {
+    val topic1 = "kafka.testTopic1"
+    val topic2 = "kafka.testTopic2"
+    val topic3 = "oooof.testTopic1"
+    adminClient.createTopics(
+      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)
+    waitForTopicCreated(topic3)
+
+    val output = TestUtils.grabConsoleOutput(
+      topicService.listTopics(new TopicCommandOptions(Array("--topic", "kafka.*"))))
+
+    assertTrue(output.contains(topic1))
+    assertTrue(output.contains(topic2))
+    assertFalse(output.contains(topic3))
+  }
+
+  @Test
+  def testListTopicsWithExcludeInternal(): Unit = {
+    val topic1 = "kafka.testTopic1"
+    adminClient.createTopics(
+      List(new NewTopic(topic1, 2, 2.toShort),
+        new NewTopic(Topic.GROUP_METADATA_TOPIC_NAME, 2, 2.toShort)).asJavaCollection)
+      .all().get()
+    waitForTopicCreated(topic1)
+
+    val output = TestUtils.grabConsoleOutput(
+      topicService.listTopics(new TopicCommandOptions(Array("--exclude-internal"))))
+
+    assertTrue(output.contains(topic1))
+    assertFalse(output.contains(Topic.GROUP_METADATA_TOPIC_NAME))
+  }
+
+  @Test
+  def testAlterPartitionCount(): Unit = {
+    adminClient.createTopics(
+      List(new NewTopic(testTopicName, 2, 2.toShort)).asJavaCollection).all().get()
+    waitForTopicCreated(testTopicName)
+
+    topicService.alterTopic(new TopicCommandOptions(
+      Array("--topic", testTopicName, "--partitions", "3")))
+
+    val topicDescription = adminClient.describeTopics(Collections.singletonList(testTopicName)).values().get(testTopicName).get()
+    assertTrue(topicDescription.partitions().size() == 3)
+  }
+
+  @Test
+  def testAlterAssignment(): Unit = {
+    adminClient.createTopics(
+      Collections.singletonList(new NewTopic(testTopicName, 2, 2.toShort))).all().get()
+    waitForTopicCreated(testTopicName)
+
+    topicService.alterTopic(new TopicCommandOptions(
+      Array("--topic", testTopicName, "--replica-assignment", "5:3,3:1,4:2", "--partitions", "3")))
+
+    val topicDescription = adminClient.describeTopics(Collections.singletonList(testTopicName)).values().get(testTopicName).get()
+    assertTrue(topicDescription.partitions().size() == 3)
+    assertEquals(List(4,2), topicDescription.partitions().get(2).replicas().asScala.map(_.id()))
+  }
+
+  @Test
+  def testAlterAssignmentWithMoreAssignmentThanPartitions(): Unit = {
+    adminClient.createTopics(
+      List(new NewTopic(testTopicName, 2, 2.toShort)).asJavaCollection).all().get()
+    waitForTopicCreated(testTopicName)
+
+    intercept[ExecutionException] {
+      topicService.alterTopic(new TopicCommandOptions(
+        Array("--topic", testTopicName, "--replica-assignment", "5:3,3:1,4:2,3:2", "--partitions", "3")))
+    }
+  }
+
+  @Test
+  def testAlterAssignmentWithMorePartitionsThanAssignment(): Unit = {
+    adminClient.createTopics(
+      List(new NewTopic(testTopicName, 2, 2.toShort)).asJavaCollection).all().get()
+    waitForTopicCreated(testTopicName)
+
+    intercept[ExecutionException] {
+      topicService.alterTopic(new TopicCommandOptions(
+        Array("--topic", testTopicName, "--replica-assignment", "5:3,3:1,4:2", "--partitions", "6")))
+    }
+  }
+
+  @Test
+  def testAlterWithInvalidPartitionCount(): Unit = {
+    createAndWaitTopic(new TopicCommandOptions(
+      Array("--partitions", "1", "--replication-factor", "1", "--topic", testTopicName)))
+
+    intercept[ExecutionException] {
+      topicService.alterTopic(new TopicCommandOptions(
+        Array("--partitions", "-1", "--topic", testTopicName)))
+    }
+  }
+
+  @Test
+  def testAlterWhenTopicDoesntExist(): Unit = {
+    // alter a topic that does not exist without --if-exists
+    val alterOpts = new TopicCommandOptions(Array("--topic", testTopicName, "--partitions", "1"))
+    val topicService = AdminClientTopicService(adminClient)
+    intercept[IllegalArgumentException] {
+      topicService.alterTopic(alterOpts)
+    }
+  }
+
+  @Test
+  def testAlterWhenTopicDoesntExistWithIfExists(): Unit = {
+    topicService.alterTopic(new TopicCommandOptions(
+      Array("--topic", testTopicName, "--partitions", "1", "--if-exists")))
+  }
+
+  @Test
+  def testCreateAlterTopicWithRackAware(): Unit = {
+    val rackInfo = Map(0 -> "rack1", 1 -> "rack2", 2 -> "rack2", 3 -> "rack1", 4 -> "rack3", 5 -> "rack3")
+
+    val numPartitions = 18
+    val replicationFactor = 3
+    val createOpts = new TopicCommandOptions(Array(
+      "--partitions", numPartitions.toString,
+      "--replication-factor", replicationFactor.toString,
+      "--topic", testTopicName))
+    createAndWaitTopic(createOpts)
+
+    var assignment = zkClient.getReplicaAssignmentForTopics(Set(testTopicName)).map { case (tp, replicas) =>
+      tp.partition -> replicas
+    }
+    checkReplicaDistribution(assignment, rackInfo, rackInfo.size, numPartitions, replicationFactor)
+
+    val alteredNumPartitions = 36
+    // verify that adding partitions will also be rack aware
+    val alterOpts = new TopicCommandOptions(Array(
+      "--partitions", alteredNumPartitions.toString,
+      "--topic", testTopicName))
+    topicService.alterTopic(alterOpts)
+    assignment = zkClient.getReplicaAssignmentForTopics(Set(testTopicName)).map { case (tp, replicas) =>
+      tp.partition -> replicas
+    }
+    checkReplicaDistribution(assignment, rackInfo, rackInfo.size, alteredNumPartitions, replicationFactor)
+  }
+
+  @Test
+  def testConfigPreservationAcrossPartitionAlteration(): Unit = {
+    val numPartitionsOriginal = 1
+    val cleanupKey = "cleanup.policy"
+    val cleanupVal = "compact"
+
+    // create the topic
+    val createOpts = new TopicCommandOptions(Array("--partitions", numPartitionsOriginal.toString,
+      "--replication-factor", "1",
+      "--config", cleanupKey + "=" + cleanupVal,
+      "--topic", testTopicName))
+    createAndWaitTopic(createOpts)
+    val props = adminZkClient.fetchEntityConfig(ConfigType.Topic, testTopicName)
+    assertTrue("Properties after creation don't contain " + cleanupKey, props.containsKey(cleanupKey))
+    assertTrue("Properties after creation have incorrect value", props.getProperty(cleanupKey).equals(cleanupVal))
+
+    // pre-create the topic config changes path to avoid a NoNodeException
+    zkClient.makeSurePersistentPathExists(ConfigEntityChangeNotificationZNode.path)
+
+    // modify the topic to add new partitions
+    val numPartitionsModified = 3
+    val alterOpts = new TopicCommandOptions(Array("--partitions", numPartitionsModified.toString, "--topic", testTopicName))
+    topicService.alterTopic(alterOpts)
+    val newProps = adminZkClient.fetchEntityConfig(ConfigType.Topic, testTopicName)
+    assertTrue("Updated properties do not contain " + cleanupKey, newProps.containsKey(cleanupKey))
+    assertTrue("Updated properties have incorrect value", newProps.getProperty(cleanupKey).equals(cleanupVal))
+  }
+
+  @Test
+  def testTopicDeletion(): Unit = {
+    // create the NormalTopic
+    val createOpts = new TopicCommandOptions(Array("--partitions", "1",
+      "--replication-factor", "1",
+      "--topic", testTopicName))
+    createAndWaitTopic(createOpts)
+
+    // delete the NormalTopic
+    val deleteOpts = new TopicCommandOptions(Array("--topic", testTopicName))
+
+    val deletePath = DeleteTopicsTopicZNode.path(testTopicName)
+    assertFalse("Delete path for topic shouldn't exist before deletion.", zkClient.pathExists(deletePath))
+    topicService.deleteTopic(deleteOpts)
+    TestUtils.verifyTopicDeletion(zkClient, testTopicName, 1, servers)
+  }
+
+  @Test
+  def testDeleteInternalTopic(): Unit = {
+    // create the offset topic
+    val createOffsetTopicOpts = new TopicCommandOptions(Array("--partitions", "1",
+      "--replication-factor", "1",
+      "--topic", Topic.GROUP_METADATA_TOPIC_NAME))
+    createAndWaitTopic(createOffsetTopicOpts)
+
+    // Try to delete the Topic.GROUP_METADATA_TOPIC_NAME which is allowed by default.
+    // This is a difference between the new and the old command as the old one didn't allow internal topic deletion.
+    // If deleting internal topics is not desired, ACLS should be used to control it.
+    val deleteOffsetTopicOpts = new TopicCommandOptions(Array("--topic", Topic.GROUP_METADATA_TOPIC_NAME))
+    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)
+    TestUtils.verifyTopicDeletion(zkClient, Topic.GROUP_METADATA_TOPIC_NAME, 1, servers)
+  }
+
+  @Test
+  def testDeleteWhenTopicDoesntExist(): Unit = {
+    // delete a topic that does not exist
+    val deleteOpts = new TopicCommandOptions(Array("--topic", testTopicName))
+    intercept[IllegalArgumentException] {
+      topicService.deleteTopic(deleteOpts)
+    }
+  }
+
+  @Test
+  def testDeleteWhenTopicDoesntExistWithIfExists(): Unit = {
+    topicService.deleteTopic(new TopicCommandOptions(Array("--topic", testTopicName, "--if-exists")))
+  }
+
+  @Test
+  def testDescribe(): Unit = {
+    adminClient.createTopics(
+      Collections.singletonList(new NewTopic(testTopicName, 2, 2.toShort))).all().get()
+    waitForTopicCreated(testTopicName)
+
+    val output = TestUtils.grabConsoleOutput(
+      topicService.describeTopic(new TopicCommandOptions(Array("--topic", testTopicName))))
+    val rows = output.split("\n")
+    assertEquals(3, rows.size)
+    rows(0).startsWith(s"Topic:$testTopicName\tPartitionCount:2")
+  }
+
+  @Test
+  def testDescribeWhenTopicDoesntExist(): Unit = {
+    intercept[IllegalArgumentException] {
+      topicService.describeTopic(new TopicCommandOptions(Array("--topic", testTopicName)))
+    }
+  }
+
+  @Test
+  def testDescribeWhenTopicDoesntExistWithIfExists(): Unit = {
+    topicService.describeTopic(new TopicCommandOptions(Array("--topic", testTopicName, "--if-exists")))
+  }
+
+  @Test
+  def testDescribeUnavailablePartitions(): Unit = {
+    adminClient.createTopics(
+      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(_.name.equals(testTopicName)).get.partitions.asScala.find(_.partitionIndex == partitionOnBroker0)
+              testPartitionMetadata match {
+                case None => fail(s"Partition metadata is not found in metadata cache")
+                case Some(metadata) => result && metadata.errorCode == Errors.LEADER_NOT_AVAILABLE.code
+              }
+            }
+          }
+      }, 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"))))
+      val rows = output.split("\n")
+      assertTrue(rows(0).startsWith(s"\tTopic: $testTopicName"))
+      assertTrue(rows(0).contains("Leader: none\tReplicas: 0\tIsr:"))
+    } finally {
+      restartDeadBrokers()
+    }
+  }
+
+  @Test
+  def testDescribeUnderReplicatedPartitions(): Unit = {
+    adminClient.createTopics(
+      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")
+      assertTrue(rows(0).startsWith(s"\tTopic: $testTopicName"))
+    } finally {
+      restartDeadBrokers()
+    }
+  }
+
+  @Test
+  def testDescribeUnderMinIsrPartitions(): Unit = {
+    val configMap = new java.util.HashMap[String, String]()
+    configMap.put(TopicConfig.MIN_IN_SYNC_REPLICAS_CONFIG, "6")
+
+    adminClient.createTopics(
+      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")
+      assertTrue(rows(0).startsWith(s"\tTopic: $testTopicName"))
+    } finally {
+      restartDeadBrokers()
+    }
+  }
+
+  @Test
+  def testDescribeUnderReplicatedPartitionsWhenReassignmentIsInProgress(): Unit = {
+    val configMap = new java.util.HashMap[String, String]()
+    val replicationFactor: Short = 1
+    val partitions = 1
+    val tp = new TopicPartition(testTopicName, 0)
+
+    adminClient.createTopics(
+      Collections.singletonList(new NewTopic(testTopicName, partitions, replicationFactor).configs(configMap))).all().get()
+    waitForTopicCreated(testTopicName)
+
+    // Produce multiple batches.
+    TestUtils.generateAndProduceMessages(servers, testTopicName, numMessages = 10, acks = -1)
+    TestUtils.generateAndProduceMessages(servers, testTopicName, numMessages = 10, acks = -1)
+
+    // Enable throttling. Note the broker config sets the replica max fetch bytes to `1` upon to minimize replication
+    // throughput so the reassignment doesn't complete quickly.
+    val brokerIds = servers.map(_.config.brokerId)
+    TestUtils.setReplicationThrottleForPartitions(adminClient, brokerIds, Set(tp), throttleBytes = 1)
+
+    val testTopicDesc = adminClient.describeTopics(Collections.singleton(testTopicName)).all().get().get(testTopicName)
+    val firstPartition = testTopicDesc.partitions().asScala.head
+
+    val replicasOfFirstPartition = firstPartition.replicas().asScala.map(_.id())
+    val targetReplica = brokerIds.diff(replicasOfFirstPartition).head
+
+    adminClient.alterPartitionReassignments(Collections.singletonMap(tp,
+      Optional.of(new NewPartitionReassignment(Collections.singletonList(targetReplica))))).all().get()
+
+    // let's wait until the LAIR is propagated
+    TestUtils.waitUntilTrue(() => {
+      val reassignments = adminClient.listPartitionReassignments(Collections.singleton(tp)).reassignments().get()
+      !reassignments.get(tp).addingReplicas().isEmpty
+    }, "Reassignment didn't add the second node")
+
+    // describe the topic and test if it's under-replicated
+    val simpleDescribeOutput = TestUtils.grabConsoleOutput(
+      topicService.describeTopic(new TopicCommandOptions(Array("--topic", testTopicName))))
+    val simpleDescribeOutputRows = simpleDescribeOutput.split("\n")
+    assertTrue(simpleDescribeOutputRows(0).startsWith(s"Topic: $testTopicName"))
+    assertEquals(2, simpleDescribeOutputRows.size)
+
+    val underReplicatedOutput = TestUtils.grabConsoleOutput(
+      topicService.describeTopic(new TopicCommandOptions(Array("--under-replicated-partitions"))))
+    assertEquals(s"--under-replicated-partitions shouldn't return anything: '$underReplicatedOutput'", "", underReplicatedOutput)
+
+    // Verify reassignment is still ongoing.
+    val reassignments = adminClient.listPartitionReassignments(Collections.singleton(tp)).reassignments.get().get(tp)
+    assertFalse(Option(reassignments).forall(_.addingReplicas.isEmpty))
+
+    TestUtils.removeReplicationThrottleForPartitions(adminClient, brokerIds, Set(tp))
+    TestUtils.waitForAllReassignmentsToComplete(adminClient)
+  }
+
+  @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
+    *   (2) topic with under-replicated partition (but not under min ISR count)
+    *   (3) topic with offline partition
+    *   (4) topic with fully replicated partition
+    *
+    * Output should only display the (1) topic with partition under min ISR count and (3) topic with offline partition
+    */
+  @Test
+  def testDescribeUnderMinIsrPartitionsMixed(): Unit = {
+    val underMinIsrTopic = "under-min-isr-topic"
+    val notUnderMinIsrTopic = "not-under-min-isr-topic"
+    val offlineTopic = "offline-topic"
+    val fullyReplicatedTopic = "fully-replicated-topic"
+
+    val configMap = new java.util.HashMap[String, String]()
+    configMap.put(TopicConfig.MIN_IN_SYNC_REPLICAS_CONFIG, "6")
+
+    adminClient.createTopics(
+      java.util.Arrays.asList(
+        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()
+
+    waitForTopicCreated(underMinIsrTopic)
+    waitForTopicCreated(notUnderMinIsrTopic)
+    waitForTopicCreated(offlineTopic)
+    waitForTopicCreated(fullyReplicatedTopic)
+
+    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")
+      assertTrue(rows(0).startsWith(s"\tTopic: $underMinIsrTopic"))
+      assertTrue(rows(1).startsWith(s"\tTopic: $offlineTopic"))
+      assertEquals(2, rows.length)
+    } finally {
+      restartDeadBrokers()
+    }
+  }
+
+  @Test
+  def testDescribeReportOverriddenConfigs(): Unit = {
+    val config = "file.delete.delay.ms=1000"
+    createAndWaitTopic(new TopicCommandOptions(
+      Array("--partitions", "2", "--replication-factor", "2", "--topic", testTopicName, "--config", config)))
+    val output = TestUtils.grabConsoleOutput(
+      topicService.describeTopic(new TopicCommandOptions(Array())))
+    assertTrue(s"Describe output should have contained $config", output.contains(config))
+  }
+
+  @Test
+  def testDescribeAndListTopicsWithoutInternalTopics(): Unit = {
+    createAndWaitTopic(
+      new TopicCommandOptions(Array("--partitions", "1", "--replication-factor", "1", "--topic", testTopicName)))
+    // create a internal topic
+    createAndWaitTopic(
+      new TopicCommandOptions(Array("--partitions", "1", "--replication-factor", "1", "--topic", Topic.GROUP_METADATA_TOPIC_NAME)))
+
+    // test describe
+    var output = TestUtils.grabConsoleOutput(topicService.describeTopic(new TopicCommandOptions(Array("--describe", "--exclude-internal"))))
+    assertTrue(s"Output should have contained $testTopicName", output.contains(testTopicName))
+    assertFalse(output.contains(Topic.GROUP_METADATA_TOPIC_NAME))
+
+    // test list
+    output = TestUtils.grabConsoleOutput(topicService.listTopics(new TopicCommandOptions(Array("--list", "--exclude-internal"))))
+    assertTrue(output.contains(testTopicName))
+    assertFalse(output.contains(Topic.GROUP_METADATA_TOPIC_NAME))
+  }
+
+  @Test
+  def testDescribeDoesNotFailWhenListingReassignmentIsUnauthorized(): Unit = {
+    adminClient = spy(adminClient)
+    topicService = AdminClientTopicService(adminClient)
+
+    val result = AdminClientTestUtils.listPartitionReassignmentsResult(
+      new ClusterAuthorizationException("Unauthorized"))
+
+    // Passing `null` here to help the compiler disambiguate the `doReturn` methods,
+    // compilation for scala 2.12 fails otherwise.
+    doReturn(result, null).when(adminClient).listPartitionReassignments(
+      Set(new TopicPartition(testTopicName, 0)).asJava
+    )
+
+    adminClient.createTopics(
+      Collections.singletonList(new NewTopic(testTopicName, 1, 1.toShort))
+    ).all().get()
+    waitForTopicCreated(testTopicName)
+
+    val output = TestUtils.grabConsoleOutput(
+      topicService.describeTopic(new TopicCommandOptions(Array("--topic", testTopicName))))
+    val rows = output.split("\n")
+    assertEquals(2, rows.size)
+    rows(0).startsWith(s"Topic:$testTopicName\tPartitionCount:1")
+  }
+
+  @Test
+  def testCreateTopicDoesNotRetryThrottlingQuotaExceededException(): Unit = {
+    val adminClient = mock(classOf[Admin])
+    val topicService = AdminClientTopicService(adminClient)
+
+    val result = AdminClientTestUtils.createTopicsResult(testTopicName, Errors.THROTTLING_QUOTA_EXCEEDED.exception())
+    when(adminClient.createTopics(any(), any())).thenReturn(result)
+
+    assertThrows(classOf[ThrottlingQuotaExceededException],
+      () => topicService.createTopic(new TopicCommandOptions(Array("--topic", testTopicName))))
+
+    val expectedNewTopic = new NewTopic(testTopicName, Optional.empty[Integer](), Optional.empty[java.lang.Short]())
+      .configs(Map.empty[String, String].asJava)
+
+    verify(adminClient, times(1)).createTopics(
+      eqThat(Set(expectedNewTopic).asJava),
+      argThat((_.shouldRetryOnQuotaViolation() == false): ArgumentMatcher[CreateTopicsOptions])
+    )
+  }
+
+  @Test
+  def testDeleteTopicDoesNotRetryThrottlingQuotaExceededException(): Unit = {
+    val adminClient = mock(classOf[Admin])
+    val topicService = AdminClientTopicService(adminClient)
+
+    val listResult = AdminClientTestUtils.listTopicsResult(testTopicName)
+    when(adminClient.listTopics(any())).thenReturn(listResult)
+
+    val result = AdminClientTestUtils.deleteTopicsResult(testTopicName, Errors.THROTTLING_QUOTA_EXCEEDED.exception())
+    when(adminClient.deleteTopics(any(), any())).thenReturn(result)
+
+    val exception = assertThrows(classOf[ExecutionException],
+      () => topicService.deleteTopic(new TopicCommandOptions(Array("--topic", testTopicName))))
+    assertTrue(exception.getCause.isInstanceOf[ThrottlingQuotaExceededException])
+
+    verify(adminClient, times(1)).deleteTopics(
+      eqThat(Seq(testTopicName).asJavaCollection),
+      argThat((_.shouldRetryOnQuotaViolation() == false): ArgumentMatcher[DeleteTopicsOptions])
+    )
+  }
+
+  @Test
+  def testCreatePartitionsDoesNotRetryThrottlingQuotaExceededException(): Unit = {
+    val adminClient = mock(classOf[Admin])
+    val topicService = AdminClientTopicService(adminClient)
+
+    val listResult = AdminClientTestUtils.listTopicsResult(testTopicName)
+    when(adminClient.listTopics(any())).thenReturn(listResult)
+
+    val topicPartitionInfo = new TopicPartitionInfo(0, new Node(0, "", 0),
+      Collections.emptyList(), Collections.emptyList())
+    val describeResult = AdminClientTestUtils.describeTopicsResult(testTopicName, new TopicDescription(
+      testTopicName, false, Collections.singletonList(topicPartitionInfo)))
+    when(adminClient.describeTopics(any())).thenReturn(describeResult)
+
+    val result = AdminClientTestUtils.createPartitionsResult(testTopicName, Errors.THROTTLING_QUOTA_EXCEEDED.exception())
+    when(adminClient.createPartitions(any(), any())).thenReturn(result)
+
+    val exception = assertThrows(classOf[ExecutionException],
+      () => topicService.alterTopic(new TopicCommandOptions(Array("--topic", testTopicName, "--partitions", "3"))))
+    assertTrue(exception.getCause.isInstanceOf[ThrottlingQuotaExceededException])
+
+    verify(adminClient, times(1)).createPartitions(
+      argThat((_.get(testTopicName).totalCount() == 3): ArgumentMatcher[java.util.Map[String, NewPartitions]]),
+      argThat((_.shouldRetryOnQuotaViolation() == false): ArgumentMatcher[CreatePartitionsOptions])
+    )
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/admin/TopicCommandWithZKClientTest.scala b/core/src/test/scala/unit/kafka/admin/TopicCommandWithZKClientTest.scala
new file mode 100644
index 0000000000000..c43240f0797ea
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/admin/TopicCommandWithZKClientTest.scala
@@ -0,0 +1,622 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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 kafka.admin.TopicCommand.{TopicCommandOptions, ZookeeperTopicService}
+import kafka.server.ConfigType
+import kafka.utils.{Exit, Logging, TestUtils}
+import kafka.zk.{ConfigEntityChangeNotificationZNode, DeleteTopicsTopicZNode, ZooKeeperTestHarness}
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.config.{ConfigException, ConfigResource}
+import org.apache.kafka.common.errors.{InvalidPartitionsException, InvalidReplicationFactorException, TopicExistsException}
+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
+
+class TopicCommandWithZKClientTest extends ZooKeeperTestHarness with Logging with RackAwareTest {
+
+  private var topicService: ZookeeperTopicService = _
+  private var testTopicName: String = _
+
+  private val _testName = new TestName
+  @Rule def testName = _testName
+
+  @Before
+  def setup(): Unit = {
+    topicService = ZookeeperTopicService(zkClient)
+    testTopicName = s"${testName.getMethodName}-${Random.alphanumeric.take(10).mkString}"
+  }
+
+  @After
+  def teardown(): Unit = {
+    if (topicService != null)
+      topicService.close()
+  }
+
+  @Test
+  def testCreate(): Unit = {
+    val brokers = List(0, 1, 2)
+    TestUtils.createBrokersInZk(zkClient, brokers)
+
+    topicService.createTopic(new TopicCommandOptions(
+      Array("--partitions", "2", "--replication-factor", "1", "--topic", testTopicName)))
+
+    assertTrue(zkClient.getAllTopicsInCluster().contains(testTopicName))
+  }
+
+  @Test
+  def testCreateWithConfigs(): Unit = {
+    val brokers = List(0, 1, 2)
+    TestUtils.createBrokersInZk(zkClient, brokers)
+
+    val configResource = new ConfigResource(ConfigResource.Type.TOPIC, testTopicName)
+    topicService.createTopic(new TopicCommandOptions(
+      Array("--partitions", "2", "--replication-factor", "2", "--topic", configResource.name(), "--config", "delete.retention.ms=1000")))
+
+    val configs = zkClient.getEntityConfigs(ConfigType.Topic, testTopicName)
+    assertEquals(1000, Integer.valueOf(configs.getProperty("delete.retention.ms")))
+  }
+
+  @Test
+  def testCreateIfNotExists(): Unit = {
+    // create brokers
+    val brokers = List(0, 1, 2)
+    TestUtils.createBrokersInZk(zkClient, brokers)
+
+    val numPartitions = 1
+
+    // create the topic
+    val createOpts = new TopicCommandOptions(
+      Array("--partitions", numPartitions.toString, "--replication-factor", "1", "--topic", testTopicName))
+    topicService.createTopic(createOpts)
+
+    // try to re-create the topic without --if-not-exists
+    intercept[TopicExistsException] {
+      topicService.createTopic(createOpts)
+    }
+
+    // try to re-create the topic with --if-not-exists
+    val createNotExistsOpts = new TopicCommandOptions(
+      Array("--partitions", numPartitions.toString, "--replication-factor", "1", "--topic", testTopicName, "--if-not-exists"))
+    topicService.createTopic(createNotExistsOpts)
+  }
+
+  @Test
+  def testCreateWithReplicaAssignment(): Unit = {
+    // create the topic
+    val createOpts = new TopicCommandOptions(
+      Array("--replica-assignment", "5:4,3:2,1:0", "--topic", testTopicName))
+    topicService.createTopic(createOpts)
+
+    val replicas0 = zkClient.getReplicasForPartition(new TopicPartition(testTopicName, 0))
+    assertEquals(List(5, 4), replicas0)
+    val replicas1 = zkClient.getReplicasForPartition(new TopicPartition(testTopicName, 1))
+    assertEquals(List(3, 2), replicas1)
+    val replicas2 = zkClient.getReplicasForPartition(new TopicPartition(testTopicName, 2))
+    assertEquals(List(1, 0), replicas2)
+  }
+
+  @Test
+  def testCreateWithInvalidReplicationFactor(): Unit = {
+    val brokers = List(0)
+    TestUtils.createBrokersInZk(zkClient, brokers)
+
+    intercept[InvalidReplicationFactorException] {
+      topicService.createTopic(new TopicCommandOptions(
+        Array("--partitions", "2", "--replication-factor", (Short.MaxValue+1).toString, "--topic", testTopicName)))
+    }
+  }
+
+  @Test
+  def testCreateWithNegativeReplicationFactor(): Unit = {
+    val brokers = List(0)
+    TestUtils.createBrokersInZk(zkClient, brokers)
+
+    intercept[InvalidReplicationFactorException] {
+      topicService.createTopic(new TopicCommandOptions(
+        Array("--partitions", "2", "--replication-factor", "-1", "--topic", testTopicName)))
+    }
+  }
+
+  @Test
+  def testCreateWithNegativePartitionCount(): Unit = {
+    val brokers = List(0)
+    TestUtils.createBrokersInZk(zkClient, brokers)
+
+    intercept[InvalidPartitionsException] {
+      topicService.createTopic(new TopicCommandOptions(
+        Array("--partitions", "-1", "--replication-factor", "1", "--topic", testTopicName)))
+    }
+  }
+
+  @Test
+  def testInvalidTopicLevelConfig(): Unit = {
+    val brokers = List(0)
+    TestUtils.createBrokersInZk(zkClient, brokers)
+
+    val createOpts = new TopicCommandOptions(
+      Array("--partitions", "1", "--replication-factor", "1", "--topic", testTopicName,
+        "--config", "message.timestamp.type=boom"))
+    intercept[ConfigException] {
+      topicService.createTopic(createOpts)
+    }
+
+    // try to create the topic with another invalid config
+    val createOpts2 = new TopicCommandOptions(
+      Array("--partitions", "1", "--replication-factor", "1", "--topic", testTopicName,
+        "--config", "message.format.version=boom"))
+    intercept[ConfigException] {
+      topicService.createTopic(createOpts2)
+    }
+  }
+
+  @Test
+  def testListTopics(): Unit = {
+    val brokers = List(0)
+    TestUtils.createBrokersInZk(zkClient, brokers)
+
+    topicService.createTopic(new TopicCommandOptions(
+      Array("--partitions", "1", "--replication-factor", "1", "--topic", testTopicName)))
+
+    val output = TestUtils.grabConsoleOutput(
+      topicService.listTopics(new TopicCommandOptions(Array())))
+
+    assertTrue(output.contains(testTopicName))
+  }
+
+  @Test
+  def testListTopicsWithIncludeList(): Unit = {
+    val brokers = List(0, 1, 2)
+    TestUtils.createBrokersInZk(zkClient, brokers)
+
+    val topic1 = "kafka.testTopic1"
+    val topic2 = "kafka.testTopic2"
+    val topic3 = "oooof.testTopic1"
+    adminZkClient.createTopic(topic1, 2, 2)
+    adminZkClient.createTopic(topic2, 2, 2)
+    adminZkClient.createTopic(topic3, 2, 2)
+
+    val output = TestUtils.grabConsoleOutput(
+      topicService.listTopics(new TopicCommandOptions(Array("--topic", "kafka.*"))))
+
+    assertTrue(output.contains(topic1))
+    assertTrue(output.contains(topic2))
+    assertFalse(output.contains(topic3))
+  }
+
+  @Test
+  def testListTopicsWithExcludeInternal(): Unit = {
+    val brokers = List(0, 1, 2)
+    TestUtils.createBrokersInZk(zkClient, brokers)
+
+    adminZkClient.createTopic(testTopicName, 2, 2)
+    adminZkClient.createTopic(Topic.GROUP_METADATA_TOPIC_NAME, 2, 2)
+
+    val output = TestUtils.grabConsoleOutput(
+      topicService.listTopics(new TopicCommandOptions(Array("--exclude-internal"))))
+
+    assertTrue(output.contains(testTopicName))
+    assertFalse(output.contains(Topic.GROUP_METADATA_TOPIC_NAME))
+  }
+
+  @Test
+  def testAlterPartitionCount(): Unit = {
+    val brokers = List(0, 1, 2)
+    TestUtils.createBrokersInZk(zkClient, brokers)
+
+    adminZkClient.createTopic(testTopicName, 2, 2)
+
+    topicService.alterTopic(new TopicCommandOptions(
+      Array("--topic", testTopicName, "--partitions", "3")))
+
+    assertEquals(3, zkClient.getPartitionsForTopics(Set(testTopicName))(testTopicName).size)
+  }
+
+  @Test
+  def testAlterAssignment(): Unit = {
+    val brokers = List(0, 1, 2, 3, 4, 5)
+    TestUtils.createBrokersInZk(zkClient, brokers)
+
+    adminZkClient.createTopic(testTopicName, 2, 2)
+
+    topicService.alterTopic(new TopicCommandOptions(
+      Array("--topic", testTopicName, "--replica-assignment", "5:3,3:1,4:2", "--partitions", "3")))
+
+    val replicas0 = zkClient.getReplicasForPartition(new TopicPartition(testTopicName, 2))
+    assertEquals(3, zkClient.getPartitionsForTopics(Set(testTopicName))(testTopicName).size)
+    assertEquals(List(4,2), replicas0)
+  }
+
+  @Test
+  def testAlterWithInvalidPartitionCount(): Unit = {
+    val brokers = List(0, 1, 2)
+    TestUtils.createBrokersInZk(zkClient, brokers)
+
+    topicService.createTopic(new TopicCommandOptions(
+      Array("--partitions", "1", "--replication-factor", "1", "--topic", testTopicName)))
+
+    intercept[InvalidPartitionsException] {
+      topicService.alterTopic(new TopicCommandOptions(
+        Array("--partitions", "-1", "--topic", testTopicName)))
+    }
+  }
+
+  @Test
+  def testAlterIfExists(): Unit = {
+    // create brokers
+    val brokers = List(0, 1, 2)
+    TestUtils.createBrokersInZk(zkClient, brokers)
+
+    // alter a topic that does not exist without --if-exists
+    val alterOpts = new TopicCommandOptions(Array("--topic", testTopicName, "--partitions", "1"))
+    intercept[IllegalArgumentException] {
+      topicService.alterTopic(alterOpts)
+    }
+
+    // alter a topic that does not exist with --if-exists
+    val alterExistsOpts = new TopicCommandOptions(Array("--topic", testTopicName, "--partitions", "1", "--if-exists"))
+    topicService.alterTopic(alterExistsOpts)
+  }
+
+  @Test
+  def testAlterConfigs(): Unit = {
+    // create brokers
+    val brokers = List(0, 1, 2)
+    TestUtils.createBrokersInZk(zkClient, brokers)
+
+    topicService.createTopic(new TopicCommandOptions(
+      Array("--partitions", "1", "--replication-factor", "1", "--topic", testTopicName)))
+
+    topicService.alterTopic(new TopicCommandOptions(
+      Array("--topic", testTopicName, "--config", "cleanup.policy=compact")))
+
+    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"))
+
+    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"))
+  }
+
+  @Test
+  def testConfigPreservationAcrossPartitionAlteration(): Unit = {
+    val numPartitionsOriginal = 1
+    val cleanupKey = "cleanup.policy"
+    val cleanupVal = "compact"
+    // create brokers
+    val brokers = List(0, 1, 2)
+    TestUtils.createBrokersInZk(zkClient, brokers)
+    // create the topic
+    val createOpts = new TopicCommandOptions(Array("--partitions", numPartitionsOriginal.toString,
+      "--replication-factor", "1",
+      "--config", cleanupKey + "=" + cleanupVal,
+      "--topic", testTopicName))
+    topicService.createTopic(createOpts)
+    val props = adminZkClient.fetchEntityConfig(ConfigType.Topic, testTopicName)
+    assertTrue("Properties after creation don't contain " + cleanupKey, props.containsKey(cleanupKey))
+    assertTrue("Properties after creation have incorrect value", props.getProperty(cleanupKey).equals(cleanupVal))
+
+    // pre-create the topic config changes path to avoid a NoNodeException
+    zkClient.makeSurePersistentPathExists(ConfigEntityChangeNotificationZNode.path)
+
+    // modify the topic to add new partitions
+    val numPartitionsModified = 3
+    val alterOpts = new TopicCommandOptions(Array("--partitions", numPartitionsModified.toString, "--topic", testTopicName))
+    topicService.alterTopic(alterOpts)
+    val newProps = adminZkClient.fetchEntityConfig(ConfigType.Topic, testTopicName)
+    assertTrue("Updated properties do not contain " + cleanupKey, newProps.containsKey(cleanupKey))
+    assertTrue("Updated properties have incorrect value", newProps.getProperty(cleanupKey).equals(cleanupVal))
+  }
+
+  @Test
+  def testTopicDeletion(): Unit = {
+
+    val numPartitionsOriginal = 1
+
+    // create brokers
+    val brokers = List(0, 1, 2)
+    TestUtils.createBrokersInZk(zkClient, brokers)
+
+    // create the NormalTopic
+    val createOpts = new TopicCommandOptions(Array("--partitions", numPartitionsOriginal.toString,
+      "--replication-factor", "1",
+      "--topic", testTopicName))
+    topicService.createTopic(createOpts)
+
+    // delete the NormalTopic
+    val deleteOpts = new TopicCommandOptions(Array("--topic", testTopicName))
+    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))
+
+    // create the offset topic
+    val createOffsetTopicOpts = new TopicCommandOptions(Array("--partitions", numPartitionsOriginal.toString,
+      "--replication-factor", "1",
+      "--topic", Topic.GROUP_METADATA_TOPIC_NAME))
+    topicService.createTopic(createOffsetTopicOpts)
+
+    // try to delete the Topic.GROUP_METADATA_TOPIC_NAME and make sure it doesn't
+    val deleteOffsetTopicOpts = new TopicCommandOptions(Array("--topic", Topic.GROUP_METADATA_TOPIC_NAME))
+    val deleteOffsetTopicPath = DeleteTopicsTopicZNode.path(Topic.GROUP_METADATA_TOPIC_NAME)
+    assertFalse("Delete path for topic shouldn't exist before deletion.", zkClient.pathExists(deleteOffsetTopicPath))
+    intercept[AdminOperationException] {
+      topicService.deleteTopic(deleteOffsetTopicOpts)
+    }
+    assertFalse("Delete path for topic shouldn't exist after deletion.", zkClient.pathExists(deleteOffsetTopicPath))
+  }
+
+  @Test
+  def testDeleteIfExists(): Unit = {
+    // create brokers
+    val brokers = List(0, 1, 2)
+    TestUtils.createBrokersInZk(zkClient, brokers)
+
+    // delete a topic that does not exist without --if-exists
+    val deleteOpts = new TopicCommandOptions(Array("--topic", testTopicName))
+    intercept[IllegalArgumentException] {
+      topicService.deleteTopic(deleteOpts)
+    }
+
+    // delete a topic that does not exist with --if-exists
+    val deleteExistsOpts = new TopicCommandOptions(Array("--topic", testTopicName, "--if-exists"))
+    topicService.deleteTopic(deleteExistsOpts)
+  }
+
+  @Test
+  def testDeleteInternalTopic(): Unit = {
+    val brokers = List(0, 1, 2)
+    TestUtils.createBrokersInZk(zkClient, brokers)
+
+    // create the offset topic
+    val createOffsetTopicOpts = new TopicCommandOptions(Array("--partitions", "1",
+      "--replication-factor", "1",
+      "--topic", Topic.GROUP_METADATA_TOPIC_NAME))
+    topicService.createTopic(createOffsetTopicOpts)
+
+    val deleteOffsetTopicOpts = new TopicCommandOptions(Array("--topic", Topic.GROUP_METADATA_TOPIC_NAME))
+    val deleteOffsetTopicPath = DeleteTopicsTopicZNode.path(Topic.GROUP_METADATA_TOPIC_NAME)
+    assertFalse("Delete path for topic shouldn't exist before deletion.", zkClient.pathExists(deleteOffsetTopicPath))
+    intercept[AdminOperationException] {
+      topicService.deleteTopic(deleteOffsetTopicOpts)
+    }
+  }
+
+  @Test
+  def testDescribeIfTopicNotExists(): Unit = {
+    // create brokers
+    val brokers = List(0, 1, 2)
+    TestUtils.createBrokersInZk(zkClient, brokers)
+
+    // describe topic that does not exist
+    val describeOpts = new TopicCommandOptions(Array("--topic", testTopicName))
+    intercept[IllegalArgumentException] {
+      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
+    topicService.describeTopic(describeOptsWithExists)
+  }
+
+  @Test
+  def testCreateAlterTopicWithRackAware(): Unit = {
+    val rackInfo = Map(0 -> "rack1", 1 -> "rack2", 2 -> "rack2", 3 -> "rack1", 4 -> "rack3", 5 -> "rack3")
+    TestUtils.createBrokersInZk(toBrokerMetadata(rackInfo), zkClient)
+
+    val numPartitions = 18
+    val replicationFactor = 3
+    val createOpts = new TopicCommandOptions(Array(
+      "--partitions", numPartitions.toString,
+      "--replication-factor", replicationFactor.toString,
+      "--topic", testTopicName))
+    topicService.createTopic(createOpts)
+
+    var assignment = zkClient.getReplicaAssignmentForTopics(Set(testTopicName)).map { case (tp, replicas) =>
+      tp.partition -> replicas
+    }
+    checkReplicaDistribution(assignment, rackInfo, rackInfo.size, numPartitions, replicationFactor)
+
+    val alteredNumPartitions = 36
+    // verify that adding partitions will also be rack aware
+    val alterOpts = new TopicCommandOptions(Array(
+      "--partitions", alteredNumPartitions.toString,
+      "--topic", testTopicName))
+    topicService.alterTopic(alterOpts)
+    assignment = zkClient.getReplicaAssignmentForTopics(Set(testTopicName)).map { case (tp, replicas) =>
+      tp.partition -> replicas
+    }
+    checkReplicaDistribution(assignment, rackInfo, rackInfo.size, alteredNumPartitions, replicationFactor)
+  }
+
+  @Test
+  def testDescribe(): Unit = {
+    val brokers = List(0, 1, 2)
+    TestUtils.createBrokersInZk(zkClient, brokers)
+
+    adminZkClient.createTopic(testTopicName, 2, 2)
+    val output = TestUtils.grabConsoleOutput(
+      topicService.describeTopic(new TopicCommandOptions(Array("--topic", testTopicName))))
+    val rows = output.split("\n")
+    assertEquals(3, rows.size)
+    rows(0).startsWith("Topic:testTopic\tPartitionCount:2")
+  }
+
+  @Test
+  def testDescribeReportOverriddenConfigs(): Unit = {
+    val brokers = List(0, 1, 2)
+    TestUtils.createBrokersInZk(zkClient, brokers)
+
+    val config = "file.delete.delay.ms=1000"
+    val configResource = new ConfigResource(ConfigResource.Type.TOPIC, testTopicName)
+    topicService.createTopic(new TopicCommandOptions(
+      Array("--partitions", "2", "--replication-factor", "2", "--topic", configResource.name(), "--config", config)))
+    val output = TestUtils.grabConsoleOutput(
+      topicService.describeTopic(new TopicCommandOptions(Array())))
+    assertTrue(output.contains(config))
+  }
+
+  @Test
+  def testDescribeAndListTopicsMarkedForDeletion(): Unit = {
+    val brokers = List(0)
+    val markedForDeletionDescribe = "MarkedForDeletion"
+    val markedForDeletionList = "marked for deletion"
+    TestUtils.createBrokersInZk(zkClient, brokers)
+
+    val createOpts = new TopicCommandOptions(Array("--partitions", "1", "--replication-factor", "1", "--topic", testTopicName))
+    topicService.createTopic(createOpts)
+
+    // delete the broker first, so when we attempt to delete the topic it gets into "marked for deletion"
+    TestUtils.deleteBrokersInZk(zkClient, brokers)
+    topicService.deleteTopic(new TopicCommandOptions(Array("--topic", testTopicName)))
+
+    // Test describe topics
+    def describeTopicsWithConfig(): Unit = {
+      topicService.describeTopic(new TopicCommandOptions(Array("--describe")))
+    }
+    val outputWithConfig = TestUtils.grabConsoleOutput(describeTopicsWithConfig())
+    assertTrue(outputWithConfig.contains(testTopicName) && outputWithConfig.contains(markedForDeletionDescribe))
+
+    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(): Unit = {
+      topicService.listTopics(new TopicCommandOptions(Array("--list")))
+    }
+    val output = TestUtils.grabConsoleOutput(listTopics())
+    assertTrue(output.contains(testTopicName) && output.contains(markedForDeletionList))
+  }
+
+  @Test
+  def testDescribeAndListTopicsWithoutInternalTopics(): Unit = {
+    val brokers = List(0)
+    TestUtils.createBrokersInZk(zkClient, brokers)
+
+    topicService.createTopic(
+      new TopicCommandOptions(Array("--partitions", "1", "--replication-factor", "1", "--topic", testTopicName)))
+    // create a internal topic
+    topicService.createTopic(
+      new TopicCommandOptions(Array("--partitions", "1", "--replication-factor", "1", "--topic", Topic.GROUP_METADATA_TOPIC_NAME)))
+
+    // test describe
+    var output = TestUtils.grabConsoleOutput(topicService.describeTopic(
+      new TopicCommandOptions(Array("--describe", "--exclude-internal"))))
+    assertTrue(output.contains(testTopicName))
+    assertFalse(output.contains(Topic.GROUP_METADATA_TOPIC_NAME))
+
+    // test list
+    output = TestUtils.grabConsoleOutput(topicService.listTopics(
+      new TopicCommandOptions(Array("--list", "--exclude-internal"))))
+    assertTrue(output.contains(testTopicName))
+    assertFalse(output.contains(Topic.GROUP_METADATA_TOPIC_NAME))
+  }
+
+  @Test
+  def testTopicOperationsWithRegexSymbolInTopicName(): Unit = {
+    val topic1 = "test.topic"
+    val topic2 = "test-topic"
+    val escapedTopic = "\"test\\.topic\""
+    val unescapedTopic = "test.topic"
+    val numPartitionsOriginal = 1
+
+    // create brokers
+    val brokers = List(0, 1, 2)
+    TestUtils.createBrokersInZk(zkClient, brokers)
+
+    // create the topics
+    val createOpts = new TopicCommandOptions(Array("--partitions", numPartitionsOriginal.toString,
+      "--replication-factor", "1", "--topic", topic1))
+    topicService.createTopic(createOpts)
+    val createOpts2 = new TopicCommandOptions(Array("--partitions", numPartitionsOriginal.toString,
+      "--replication-factor", "1", "--topic", topic2))
+    topicService.createTopic(createOpts2)
+
+    val escapedCommandOpts = new TopicCommandOptions(Array("--topic", escapedTopic))
+    val unescapedCommandOpts = new TopicCommandOptions(Array("--topic", unescapedTopic))
+
+    // topic actions with escaped regex do not affect 'test-topic'
+    // topic actions with unescaped topic affect 'test-topic'
+
+    assertFalse(TestUtils.grabConsoleOutput(topicService.describeTopic(escapedCommandOpts)).contains(topic2))
+    assertTrue(TestUtils.grabConsoleOutput(topicService.describeTopic(unescapedCommandOpts)).contains(topic2))
+
+    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)
+  }
+
+  @Test
+  def testCreateWithUnspecifiedReplicationFactorAndPartitionsWithZkClient(): Unit = {
+    assertExitCode(1, () =>
+      new TopicCommandOptions(Array("--create", "--zookeeper", "zk", "--topic", testTopicName)).checkArgs()
+    )
+  }
+
+  def assertExitCode(expected: Int, method: () => Unit): Unit = {
+    def mockExitProcedure(exitCode: Int, exitMessage: Option[String]): Nothing = {
+      assertEquals(expected, exitCode)
+      throw new RuntimeException
+    }
+    Exit.setExitProcedure(mockExitProcedure)
+    try {
+      intercept[RuntimeException] {
+        method()
+      }
+    } finally {
+      Exit.resetExitProcedure()
+    }
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/admin/UserScramCredentialsCommandTest.scala b/core/src/test/scala/unit/kafka/admin/UserScramCredentialsCommandTest.scala
new file mode 100644
index 0000000000000..5abbc0a823cc6
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/admin/UserScramCredentialsCommandTest.scala
@@ -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 kafka.admin
+
+import java.io.{ByteArrayOutputStream, PrintStream}
+import java.nio.charset.StandardCharsets
+
+import kafka.server.BaseRequestTest
+import kafka.utils.Exit
+import org.junit.Assert._
+import org.junit.Test
+
+class UserScramCredentialsCommandTest extends BaseRequestTest {
+  override def brokerCount = 1
+  var exitStatus: Option[Int] = None
+  var exitMessage: Option[String] = None
+
+  case class ConfigCommandResult(stdout: String, exitStatus: Option[Int] = None)
+
+  private def runConfigCommandViaBroker(args: Array[String]) : ConfigCommandResult = {
+    val byteArrayOutputStream = new ByteArrayOutputStream()
+    val utf8 = StandardCharsets.UTF_8.name
+    val printStream = new PrintStream(byteArrayOutputStream, true, utf8)
+    var exitStatus: Option[Int] = None
+    Exit.setExitProcedure { (status, _) =>
+      exitStatus = Some(status)
+      throw new RuntimeException
+    }
+    val commandArgs = Array("--bootstrap-server", brokerList) ++ args
+    try {
+      Console.withOut(printStream) {
+        ConfigCommand.main(commandArgs)
+      }
+      ConfigCommandResult(byteArrayOutputStream.toString(utf8))
+    } catch {
+      case e: Exception => {
+        debug(s"Exception running ConfigCommand ${commandArgs.mkString(" ")}", e)
+        ConfigCommandResult("", exitStatus)
+      }
+    } finally {
+      printStream.close
+      Exit.resetExitProcedure()
+    }
+  }
+
+  @Test
+  def testUserScramCredentialsRequests(): Unit = {
+    val user1 = "user1"
+    // create and describe a credential
+    var result = runConfigCommandViaBroker(Array("--user", user1, "--alter", "--add-config", "SCRAM-SHA-256=[iterations=4096,password=foo-secret]"))
+    val alterConfigsUser1Out = s"Completed updating config for user $user1.\n"
+    assertEquals(alterConfigsUser1Out, result.stdout)
+    result = runConfigCommandViaBroker(Array("--user", user1, "--describe"))
+    val scramCredentialConfigsUser1Out = s"SCRAM credential configs for user-principal '$user1' are SCRAM-SHA-256=iterations=4096\n"
+    assertEquals(scramCredentialConfigsUser1Out, result.stdout)
+    // create a user quota and describe the user again
+    result = runConfigCommandViaBroker(Array("--user", user1, "--alter", "--add-config", "consumer_byte_rate=20000"))
+    assertEquals(alterConfigsUser1Out, result.stdout)
+    result = runConfigCommandViaBroker(Array("--user", user1, "--describe"))
+    val quotaConfigsUser1Out = s"Quota configs for user-principal '$user1' are consumer_byte_rate=20000.0\n"
+    assertEquals(s"$quotaConfigsUser1Out$scramCredentialConfigsUser1Out", result.stdout)
+
+    // now do the same thing for user2
+    val user2 = "user2"
+    // create and describe a credential
+    result = runConfigCommandViaBroker(Array("--user", user2, "--alter", "--add-config", "SCRAM-SHA-256=[iterations=4096,password=foo-secret]"))
+    val alterConfigsUser2Out = s"Completed updating config for user $user2.\n"
+    assertEquals(alterConfigsUser2Out, result.stdout)
+    result = runConfigCommandViaBroker(Array("--user", user2, "--describe"))
+    val scramCredentialConfigsUser2Out = s"SCRAM credential configs for user-principal '$user2' are SCRAM-SHA-256=iterations=4096\n"
+    assertEquals(scramCredentialConfigsUser2Out, result.stdout)
+    // create a user quota and describe the user again
+    result = runConfigCommandViaBroker(Array("--user", user2, "--alter", "--add-config", "consumer_byte_rate=20000"))
+    assertEquals(alterConfigsUser2Out, result.stdout)
+    result = runConfigCommandViaBroker(Array("--user", user2, "--describe"))
+    val quotaConfigsUser2Out = s"Quota configs for user-principal '$user2' are consumer_byte_rate=20000.0\n"
+    assertEquals(s"$quotaConfigsUser2Out$scramCredentialConfigsUser2Out", result.stdout)
+
+    // describe both
+    result = runConfigCommandViaBroker(Array("--entity-type", "users", "--describe"))
+    // we don't know the order that quota or scram users come out, so we have 2 possibilities for each, 4 total
+    val quotaPossibilityAOut = s"$quotaConfigsUser1Out$quotaConfigsUser2Out"
+    val quotaPossibilityBOut = s"$quotaConfigsUser2Out$quotaConfigsUser1Out"
+    val scramPossibilityAOut = s"$scramCredentialConfigsUser1Out$scramCredentialConfigsUser2Out"
+    val scramPossibilityBOut = s"$scramCredentialConfigsUser2Out$scramCredentialConfigsUser1Out"
+    assertTrue(result.stdout.equals(s"$quotaPossibilityAOut$scramPossibilityAOut")
+      || result.stdout.equals(s"$quotaPossibilityAOut$scramPossibilityBOut")
+      || result.stdout.equals(s"$quotaPossibilityBOut$scramPossibilityAOut")
+      || result.stdout.equals(s"$quotaPossibilityBOut$scramPossibilityBOut"))
+
+    // now delete configs, in opposite order, for user1 and user2, and describe
+    result = runConfigCommandViaBroker(Array("--user", user1, "--alter", "--delete-config", "consumer_byte_rate"))
+    assertEquals(alterConfigsUser1Out, result.stdout)
+    result = runConfigCommandViaBroker(Array("--user", user2, "--alter", "--delete-config", "SCRAM-SHA-256"))
+    assertEquals(alterConfigsUser2Out, result.stdout)
+    result = runConfigCommandViaBroker(Array("--entity-type", "users", "--describe"))
+    assertEquals(s"$quotaConfigsUser2Out$scramCredentialConfigsUser1Out", result.stdout)
+
+    // now delete the rest of the configs, for user1 and user2, and describe
+    result = runConfigCommandViaBroker(Array("--user", user1, "--alter", "--delete-config", "SCRAM-SHA-256"))
+    assertEquals(alterConfigsUser1Out, result.stdout)
+    result = runConfigCommandViaBroker(Array("--user", user2, "--alter", "--delete-config", "consumer_byte_rate"))
+    assertEquals(alterConfigsUser2Out, result.stdout)
+    result = runConfigCommandViaBroker(Array("--entity-type", "users", "--describe"))
+    assertEquals("", result.stdout)
+  }
+
+  @Test
+  def testAlterWithEmptyPassword(): Unit = {
+    val user1 = "user1"
+    val result = runConfigCommandViaBroker(Array("--user", user1, "--alter", "--add-config", "SCRAM-SHA-256=[iterations=4096,password=]"))
+    assertTrue("Expected System.exit() to be called with an empty password", result.exitStatus.isDefined)
+    assertEquals("Expected empty password to cause failure with exit status=1", 1, result.exitStatus.get)
+  }
+
+  @Test
+  def testDescribeUnknownUser(): Unit = {
+    val unknownUser = "unknownUser"
+    val result = runConfigCommandViaBroker(Array("--user", unknownUser, "--describe"))
+    assertTrue("Expected System.exit() to not be called with an unknown user", result.exitStatus.isEmpty)
+    assertEquals("", result.stdout)
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/api/ApiUtilsTest.scala b/core/src/test/scala/unit/kafka/api/ApiUtilsTest.scala
index b71b00bc05a5c..08994c87ad01a 100644
--- a/core/src/test/scala/unit/kafka/api/ApiUtilsTest.scala
+++ b/core/src/test/scala/unit/kafka/api/ApiUtilsTest.scala
@@ -18,21 +18,22 @@
 package kafka.api
 
 import org.junit._
-import org.scalatest.junit.JUnitSuite
 import org.junit.Assert._
+
 import scala.util.Random
 import java.nio.ByteBuffer
-import kafka.common.KafkaException
+
 import kafka.utils.TestUtils
+import org.apache.kafka.common.KafkaException
 
 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 
@@ -45,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 6fc69747f5d55..8a5d4ee856e9d 100644
--- a/core/src/test/scala/unit/kafka/api/ApiVersionTest.scala
+++ b/core/src/test/scala/unit/kafka/api/ApiVersionTest.scala
@@ -17,8 +17,17 @@
 
 package kafka.api
 
-import org.junit.Test
+import java.util
+
+import org.apache.kafka.common.feature.{Features, FinalizedVersionRange, SupportedVersionRange}
+import org.apache.kafka.common.protocol.ApiKeys
+import org.apache.kafka.common.record.{RecordBatch, RecordVersion}
+import org.apache.kafka.common.requests.{AbstractResponse, ApiVersionsResponse}
+import org.apache.kafka.common.utils.Utils
 import org.junit.Assert._
+import org.junit.Test
+
+import scala.jdk.CollectionConverters._
 
 class ApiVersionTest {
 
@@ -72,6 +81,184 @@ class ApiVersionTest {
     assertEquals(KAFKA_1_0_IV0, ApiVersion("1.0.0"))
     assertEquals(KAFKA_1_0_IV0, ApiVersion("1.0.0-IV0"))
     assertEquals(KAFKA_1_0_IV0, ApiVersion("1.0.1"))
+
+    assertEquals(KAFKA_1_1_IV0, ApiVersion("1.1-IV0"))
+
+    assertEquals(KAFKA_2_0_IV1, ApiVersion("2.0"))
+    assertEquals(KAFKA_2_0_IV0, ApiVersion("2.0-IV0"))
+    assertEquals(KAFKA_2_0_IV1, ApiVersion("2.0-IV1"))
+
+    assertEquals(KAFKA_2_1_IV2, ApiVersion("2.1"))
+    assertEquals(KAFKA_2_1_IV0, ApiVersion("2.1-IV0"))
+    assertEquals(KAFKA_2_1_IV1, ApiVersion("2.1-IV1"))
+    assertEquals(KAFKA_2_1_IV2, ApiVersion("2.1-IV2"))
+
+    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_IV1, ApiVersion("2.3"))
+    assertEquals(KAFKA_2_3_IV0, ApiVersion("2.3-IV0"))
+    assertEquals(KAFKA_2_3_IV1, ApiVersion("2.3-IV1"))
+
+    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"))
+
+    assertEquals(KAFKA_2_5_IV0, ApiVersion("2.5"))
+    assertEquals(KAFKA_2_5_IV0, ApiVersion("2.5-IV0"))
+
+    assertEquals(KAFKA_2_6_IV0, ApiVersion("2.6"))
+    assertEquals(KAFKA_2_6_IV0, ApiVersion("2.6-IV0"))
+
+    assertEquals(KAFKA_2_7_IV2, ApiVersion("2.7"))
+    assertEquals(KAFKA_2_7_IV0, ApiVersion("2.7-IV0"))
+    assertEquals(KAFKA_2_7_IV1, ApiVersion("2.7-IV1"))
+    assertEquals(KAFKA_2_7_IV2, ApiVersion("2.7-IV2"))
+
+    assertEquals(KAFKA_2_8_IV0, ApiVersion("2.8"))
+    assertEquals(KAFKA_2_8_IV0, ApiVersion("2.8-IV0"))
+  }
+
+  @Test
+  def testApiVersionUniqueIds(): Unit = {
+    val allIds: Seq[Int] = ApiVersion.allVersions.map(apiVersion => {
+      apiVersion.id
+    })
+
+    val uniqueIds: Set[Int] = allIds.toSet
+
+    assertEquals(allIds.size, uniqueIds.size)
+  }
+
+  @Test
+  def testMinSupportedVersionFor(): Unit = {
+    assertEquals(KAFKA_0_8_0, ApiVersion.minSupportedFor(RecordVersion.V0))
+    assertEquals(KAFKA_0_10_0_IV0, ApiVersion.minSupportedFor(RecordVersion.V1))
+    assertEquals(KAFKA_0_11_0_IV0, ApiVersion.minSupportedFor(RecordVersion.V2))
+
+    // Ensure that all record versions have a defined min version so that we remember to update the method
+    for (recordVersion <- RecordVersion.values)
+      assertNotNull(ApiVersion.minSupportedFor(recordVersion))
   }
 
+  @Test
+  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)
+    assertEquals("2.5", KAFKA_2_5_IV0.shortVersion)
+    assertEquals("2.6", KAFKA_2_6_IV0.shortVersion)
+    assertEquals("2.7", KAFKA_2_7_IV2.shortVersion)
+    assertEquals("2.8", KAFKA_2_8_IV0.shortVersion)
+  }
+
+  @Test
+  def testApiVersionValidator(): Unit = {
+    val str = ApiVersionValidator.toString
+    val apiVersions = str.slice(1, str.length).split(",")
+    assertEquals(ApiVersion.allVersions.size, apiVersions.length)
+  }
+
+  @Test
+  def shouldCreateApiResponseOnlyWithKeysSupportedByMagicValue(): Unit = {
+    val response = ApiVersion.apiVersionsResponse(
+      10,
+      RecordBatch.MAGIC_VALUE_V1,
+      Features.emptySupportedFeatures
+    )
+    verifyApiKeysForMagic(response, RecordBatch.MAGIC_VALUE_V1)
+    assertEquals(10, response.throttleTimeMs)
+    assertTrue(response.data.supportedFeatures.isEmpty)
+    assertTrue(response.data.finalizedFeatures.isEmpty)
+    assertEquals(ApiVersionsResponse.UNKNOWN_FINALIZED_FEATURES_EPOCH, response.data.finalizedFeaturesEpoch)
+  }
+
+  @Test
+  def shouldReturnFeatureKeysWhenMagicIsCurrentValueAndThrottleMsIsDefaultThrottle(): Unit = {
+    val response = ApiVersion.apiVersionsResponse(
+      10,
+      RecordBatch.MAGIC_VALUE_V1,
+      Features.supportedFeatures(
+        Utils.mkMap(Utils.mkEntry("feature", new SupportedVersionRange(1.toShort, 4.toShort)))),
+      Features.finalizedFeatures(
+        Utils.mkMap(Utils.mkEntry("feature", new FinalizedVersionRange(2.toShort, 3.toShort)))),
+      10
+    )
+
+    verifyApiKeysForMagic(response, RecordBatch.MAGIC_VALUE_V1)
+    assertEquals(10, response.throttleTimeMs)
+    assertEquals(1, response.data.supportedFeatures.size)
+    val sKey = response.data.supportedFeatures.find("feature")
+    assertNotNull(sKey)
+    assertEquals(1, sKey.minVersion)
+    assertEquals(4, sKey.maxVersion)
+    assertEquals(1, response.data.finalizedFeatures.size)
+    val fKey = response.data.finalizedFeatures.find("feature")
+    assertNotNull(fKey)
+    assertEquals(2, fKey.minVersionLevel)
+    assertEquals(3, fKey.maxVersionLevel)
+    assertEquals(10, response.data.finalizedFeaturesEpoch)
+  }
+
+  private def verifyApiKeysForMagic(response: ApiVersionsResponse, maxMagic: Byte): Unit = {
+    for (version <- response.data.apiKeys.asScala) {
+      assertTrue(ApiKeys.forId(version.apiKey).minRequiredInterBrokerMagic <= maxMagic)
+    }
+  }
+
+  @Test
+  def shouldReturnAllKeysWhenMagicIsCurrentValueAndThrottleMsIsDefaultThrottle(): Unit = {
+    val response = ApiVersion.apiVersionsResponse(
+      AbstractResponse.DEFAULT_THROTTLE_TIME,
+      RecordBatch.CURRENT_MAGIC_VALUE,
+      Features.emptySupportedFeatures
+    )
+    assertEquals(new util.HashSet[ApiKeys](ApiKeys.enabledApis), apiKeysInResponse(response))
+    assertEquals(AbstractResponse.DEFAULT_THROTTLE_TIME, response.throttleTimeMs)
+    assertTrue(response.data.supportedFeatures.isEmpty)
+    assertTrue(response.data.finalizedFeatures.isEmpty)
+    assertEquals(ApiVersionsResponse.UNKNOWN_FINALIZED_FEATURES_EPOCH, response.data.finalizedFeaturesEpoch)
+  }
+
+  @Test
+  def testMetadataQuorumApisAreDisabled(): Unit = {
+    val response = ApiVersion.apiVersionsResponse(
+      AbstractResponse.DEFAULT_THROTTLE_TIME,
+      RecordBatch.CURRENT_MAGIC_VALUE,
+      Features.emptySupportedFeatures
+    )
+
+    // Ensure that APIs needed for the internal metadata quorum (KIP-500)
+    // are not exposed through ApiVersions until we are ready for them
+    val exposedApis = apiKeysInResponse(response)
+    assertFalse(exposedApis.contains(ApiKeys.ENVELOPE))
+    assertFalse(exposedApis.contains(ApiKeys.VOTE))
+    assertFalse(exposedApis.contains(ApiKeys.BEGIN_QUORUM_EPOCH))
+    assertFalse(exposedApis.contains(ApiKeys.END_QUORUM_EPOCH))
+    assertFalse(exposedApis.contains(ApiKeys.DESCRIBE_QUORUM))
+  }
+
+  private def apiKeysInResponse(apiVersions: ApiVersionsResponse) = {
+    val apiKeys = new util.HashSet[ApiKeys]
+    for (version <- apiVersions.data.apiKeys.asScala) {
+      apiKeys.add(ApiKeys.forId(version.apiKey))
+    }
+    apiKeys
+  }
 }
diff --git a/core/src/test/scala/unit/kafka/api/FetchRequestTest.scala b/core/src/test/scala/unit/kafka/api/FetchRequestTest.scala
deleted file mode 100644
index c2bdf49517f7f..0000000000000
--- a/core/src/test/scala/unit/kafka/api/FetchRequestTest.scala
+++ /dev/null
@@ -1,63 +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 kafka.common.TopicAndPartition
-import org.junit.Assert.{assertEquals, assertNotEquals}
-import org.junit.Test
-
-class FetchRequestTest {
-
-  @Test
-  def testShuffle() {
-    val seq = (0 to 100).map { i =>
-      val topic = s"topic${i % 10}"
-      (TopicAndPartition(topic, i / 10), PartitionFetchInfo(i, 50))
-    }
-    val shuffled = FetchRequest.shuffle(seq)
-    assertEquals(seq.size, shuffled.size)
-    assertNotEquals(seq, shuffled)
-
-    seq.foreach { case (tp1, fetchInfo1) =>
-      shuffled.foreach { case (tp2, fetchInfo2) =>
-        if (tp1 == tp2)
-          assertEquals(fetchInfo1, fetchInfo2)
-      }
-    }
-
-    val topics = seq.map { case (TopicAndPartition(t, _), _) => t }.distinct
-    topics.foreach { topic =>
-      val startIndex = shuffled.indexWhere { case (tp, _) => tp.topic == topic }
-      val endIndex = shuffled.lastIndexWhere { case (tp, _) => tp.topic == topic }
-      // all partitions for a given topic should appear in sequence
-      assertEquals(Set(topic), shuffled.slice(startIndex, endIndex + 1).map { case (tp, _) => tp.topic }.toSet)
-    }
-
-    val shuffled2 = FetchRequest.shuffle(seq)
-    assertNotEquals(shuffled, shuffled2)
-    assertNotEquals(seq, shuffled2)
-  }
-
-  @Test
-  def testShuffleWithSingleTopic() {
-    val seq = (0 to 50).map(i => (TopicAndPartition("topic", i), PartitionFetchInfo(i, 70)))
-    val shuffled = FetchRequest.shuffle(seq)
-    assertEquals(seq.size, shuffled.size)
-    assertNotEquals(seq, shuffled)
-  }
-
-}
diff --git a/core/src/test/scala/unit/kafka/api/RequestResponseSerializationTest.scala b/core/src/test/scala/unit/kafka/api/RequestResponseSerializationTest.scala
deleted file mode 100644
index 026786e978686..0000000000000
--- a/core/src/test/scala/unit/kafka/api/RequestResponseSerializationTest.scala
+++ /dev/null
@@ -1,239 +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 kafka.common.{OffsetAndMetadata, OffsetMetadataAndError}
-import kafka.common._
-import kafka.message.{ByteBufferMessageSet, Message}
-import kafka.common.TopicAndPartition
-import kafka.utils.TestUtils
-import TestUtils.createBroker
-import java.nio.ByteBuffer
-
-import org.apache.kafka.common.network.ListenerName
-import org.apache.kafka.common.protocol.Errors
-import org.apache.kafka.common.security.auth.SecurityProtocol
-import org.apache.kafka.common.utils.Time
-import org.junit._
-import org.scalatest.junit.JUnitSuite
-import org.junit.Assert._
-
-
-object SerializationTestUtils {
-  private val topic1 = "test1"
-  private val topic2 = "test2"
-  private val partitionDataFetchResponse0 = new FetchResponsePartitionData(messages = new ByteBufferMessageSet(new Message("first message".getBytes)))
-  private val partitionDataFetchResponse1 = new FetchResponsePartitionData(messages = new ByteBufferMessageSet(new Message("second message".getBytes)))
-  private val partitionDataFetchResponse2 = new FetchResponsePartitionData(messages = new ByteBufferMessageSet(new Message("third message".getBytes)))
-  private val partitionDataFetchResponse3 = new FetchResponsePartitionData(messages = new ByteBufferMessageSet(new Message("fourth message".getBytes)))
-  private val partitionDataFetchResponseMap = Map((0, partitionDataFetchResponse0), (1, partitionDataFetchResponse1), (2, partitionDataFetchResponse2), (3, partitionDataFetchResponse3))
-
-  private val topicDataFetchResponse = {
-    val groupedData = Array(topic1, topic2).flatMap(topic =>
-      partitionDataFetchResponseMap.map(partitionAndData =>
-        (TopicAndPartition(topic, partitionAndData._1), partitionAndData._2)))
-    collection.immutable.Map(groupedData:_*)
-  }
-
-  private val partitionDataMessage0 = new ByteBufferMessageSet(new Message("first message".getBytes))
-  private val partitionDataMessage1 = new ByteBufferMessageSet(new Message("second message".getBytes))
-  private val partitionDataMessage2 = new ByteBufferMessageSet(new Message("third message".getBytes))
-  private val partitionDataMessage3 = new ByteBufferMessageSet(new Message("fourth message".getBytes))
-  private val partitionDataProducerRequestArray = Array(partitionDataMessage0, partitionDataMessage1, partitionDataMessage2, partitionDataMessage3)
-
-  val topicDataProducerRequest = {
-    val groupedData = Array(topic1, topic2).flatMap(topic =>
-      partitionDataProducerRequestArray.zipWithIndex.map
-      {
-        case(partitionDataMessage, partition) =>
-          (TopicAndPartition(topic, partition), partitionDataMessage)
-      })
-    collection.mutable.Map(groupedData:_*)
-  }
-
-  private val requestInfos = collection.immutable.Map(
-    TopicAndPartition(topic1, 0) -> PartitionFetchInfo(1000, 100),
-    TopicAndPartition(topic1, 1) -> PartitionFetchInfo(2000, 100),
-    TopicAndPartition(topic1, 2) -> PartitionFetchInfo(3000, 100),
-    TopicAndPartition(topic1, 3) -> PartitionFetchInfo(4000, 100),
-    TopicAndPartition(topic2, 0) -> PartitionFetchInfo(1000, 100),
-    TopicAndPartition(topic2, 1) -> PartitionFetchInfo(2000, 100),
-    TopicAndPartition(topic2, 2) -> PartitionFetchInfo(3000, 100),
-    TopicAndPartition(topic2, 3) -> PartitionFetchInfo(4000, 100)
-  )
-
-  private val brokers = List(createBroker(0, "localhost", 1011), createBroker(0, "localhost", 1012),
-    createBroker(0, "localhost", 1013))
-
-  def createTestProducerRequest: ProducerRequest = {
-    new ProducerRequest(1, "client 1", 0, 1000, topicDataProducerRequest)
-  }
-
-  def createTestProducerResponse: ProducerResponse =
-    ProducerResponse(1, Map(
-      TopicAndPartition(topic1, 0) -> ProducerResponseStatus(Errors.forCode(0.toShort), 10001),
-      TopicAndPartition(topic2, 0) -> ProducerResponseStatus(Errors.forCode(0.toShort), 20001)
-    ), ProducerRequest.CurrentVersion, 100)
-
-  def createTestFetchRequest: FetchRequest = new FetchRequest(requestInfo = requestInfos.toVector)
-
-  def createTestFetchResponse: FetchResponse = FetchResponse(1, topicDataFetchResponse.toVector)
-
-  def createTestOffsetRequest = new OffsetRequest(
-      collection.immutable.Map(TopicAndPartition(topic1, 1) -> PartitionOffsetRequestInfo(1000, 200)),
-      replicaId = 0
-  )
-
-  def createTestOffsetResponse: OffsetResponse = {
-    new OffsetResponse(0, collection.immutable.Map(
-      TopicAndPartition(topic1, 1) -> PartitionOffsetsResponse(Errors.NONE, Seq(1000l, 2000l, 3000l, 4000l)))
-    )
-  }
-
-  def createTestOffsetCommitRequestV2: OffsetCommitRequest = {
-    new OffsetCommitRequest(
-      groupId = "group 1",
-      retentionMs = Time.SYSTEM.milliseconds,
-      requestInfo=collection.immutable.Map(
-      TopicAndPartition(topic1, 0) -> OffsetAndMetadata(42L, "some metadata"),
-      TopicAndPartition(topic1, 1) -> OffsetAndMetadata(100L, OffsetMetadata.NoMetadata)
-    ))
-  }
-
-  def createTestOffsetCommitRequestV1: OffsetCommitRequest = {
-    new OffsetCommitRequest(
-      versionId = 1,
-      groupId = "group 1",
-      requestInfo = collection.immutable.Map(
-      TopicAndPartition(topic1, 0) -> OffsetAndMetadata(42L, "some metadata", Time.SYSTEM.milliseconds),
-      TopicAndPartition(topic1, 1) -> OffsetAndMetadata(100L, OffsetMetadata.NoMetadata, Time.SYSTEM.milliseconds)
-    ))
-  }
-
-  def createTestOffsetCommitRequestV0: OffsetCommitRequest = {
-    new OffsetCommitRequest(
-      versionId = 0,
-      groupId = "group 1",
-      requestInfo = collection.immutable.Map(
-        TopicAndPartition(topic1, 0) -> OffsetAndMetadata(42L, "some metadata"),
-        TopicAndPartition(topic1, 1) -> OffsetAndMetadata(100L, OffsetMetadata.NoMetadata)
-      ))
-  }
-
-  def createTestOffsetCommitResponse: OffsetCommitResponse = {
-    new OffsetCommitResponse(collection.immutable.Map(TopicAndPartition(topic1, 0) -> Errors.NONE,
-                                 TopicAndPartition(topic1, 1) -> Errors.NONE))
-  }
-
-  def createTestOffsetFetchRequest: OffsetFetchRequest = {
-    new OffsetFetchRequest("group 1", Seq(
-      TopicAndPartition(topic1, 0),
-      TopicAndPartition(topic1, 1)
-    ))
-  }
-
-  def createTestOffsetFetchResponse: OffsetFetchResponse = {
-    new OffsetFetchResponse(collection.immutable.Map(
-      TopicAndPartition(topic1, 0) -> OffsetMetadataAndError(42L, "some metadata", Errors.NONE),
-      TopicAndPartition(topic1, 1) -> OffsetMetadataAndError(100L, OffsetMetadata.NoMetadata, Errors.UNKNOWN_TOPIC_OR_PARTITION)
-    ), error = Errors.NONE)
-  }
-
-  def createConsumerMetadataRequest: GroupCoordinatorRequest = GroupCoordinatorRequest("group 1", clientId = "client 1")
-
-  def createConsumerMetadataResponse: GroupCoordinatorResponse = {
-    GroupCoordinatorResponse(Some(
-      brokers.head.getBrokerEndPoint(ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT))), Errors.NONE, 0)
-  }
-}
-
-class RequestResponseSerializationTest extends JUnitSuite {
-  private val producerRequest = SerializationTestUtils.createTestProducerRequest
-  private val producerResponse = SerializationTestUtils.createTestProducerResponse
-  private val fetchRequest = SerializationTestUtils.createTestFetchRequest
-  private val offsetRequest = SerializationTestUtils.createTestOffsetRequest
-  private val offsetResponse = SerializationTestUtils.createTestOffsetResponse
-  private val offsetCommitRequestV0 = SerializationTestUtils.createTestOffsetCommitRequestV0
-  private val offsetCommitRequestV1 = SerializationTestUtils.createTestOffsetCommitRequestV1
-  private val offsetCommitRequestV2 = SerializationTestUtils.createTestOffsetCommitRequestV2
-  private val offsetCommitResponse = SerializationTestUtils.createTestOffsetCommitResponse
-  private val offsetFetchRequest = SerializationTestUtils.createTestOffsetFetchRequest
-  private val offsetFetchResponse = SerializationTestUtils.createTestOffsetFetchResponse
-  private val consumerMetadataRequest = SerializationTestUtils.createConsumerMetadataRequest
-  private val consumerMetadataResponse = SerializationTestUtils.createConsumerMetadataResponse
-  private val consumerMetadataResponseNoCoordinator = GroupCoordinatorResponse(None, Errors.COORDINATOR_NOT_AVAILABLE, 0)
-
-  @Test
-  def testSerializationAndDeserialization() {
-
-    val requestsAndResponses =
-      collection.immutable.Seq(producerRequest, producerResponse,
-                               fetchRequest, offsetRequest, offsetResponse,
-                               offsetCommitRequestV0, offsetCommitRequestV1, offsetCommitRequestV2, offsetCommitResponse,
-                               offsetFetchRequest, offsetFetchResponse,
-                               consumerMetadataRequest, consumerMetadataResponse,
-                               consumerMetadataResponseNoCoordinator)
-
-    requestsAndResponses.foreach { original =>
-      val buffer = ByteBuffer.allocate(original.sizeInBytes)
-      original.writeTo(buffer)
-      buffer.rewind()
-      val deserializer = original.getClass.getDeclaredMethod("readFrom", classOf[ByteBuffer])
-      val deserialized = deserializer.invoke(null, buffer)
-      assertFalse("All serialized bytes in " + original.getClass.getSimpleName + " should have been consumed",
-                  buffer.hasRemaining)
-      assertEquals("The original and deserialized for " + original.getClass.getSimpleName + " should be the same.", original, deserialized)
-    }
-  }
-
-  @Test
-  def testProduceResponseVersion() {
-    val oldClientResponse = ProducerResponse(1, Map(
-      TopicAndPartition("t1", 0) -> ProducerResponseStatus(Errors.NONE, 10001),
-      TopicAndPartition("t2", 0) -> ProducerResponseStatus(Errors.NONE, 20001)
-    ))
-
-    val newClientResponse = ProducerResponse(1, Map(
-      TopicAndPartition("t1", 0) -> ProducerResponseStatus(Errors.NONE, 10001),
-      TopicAndPartition("t2", 0) -> ProducerResponseStatus(Errors.NONE, 20001)
-    ), 1, 100)
-
-    // new response should have 4 bytes more than the old response since delayTime is an INT32
-    assertEquals(oldClientResponse.sizeInBytes + 4, newClientResponse.sizeInBytes)
-
-    val buffer = ByteBuffer.allocate(newClientResponse.sizeInBytes)
-    newClientResponse.writeTo(buffer)
-    buffer.rewind()
-    assertEquals(ProducerResponse.readFrom(buffer).throttleTime, 100)
-  }
-
-  @Test
-  def testFetchResponseVersion() {
-    val oldClientResponse = FetchResponse(1, Map(
-      TopicAndPartition("t1", 0) -> new FetchResponsePartitionData(messages = new ByteBufferMessageSet(new Message("first message".getBytes)))
-    ).toVector, 0)
-
-    val newClientResponse = FetchResponse(1, Map(
-      TopicAndPartition("t1", 0) -> new FetchResponsePartitionData(messages = new ByteBufferMessageSet(new Message("first message".getBytes)))
-    ).toVector, 1, 100)
-
-    // new response should have 4 bytes more than the old response since delayTime is an INT32
-    assertEquals(oldClientResponse.sizeInBytes + 4, newClientResponse.sizeInBytes)
-  }
-
-}
diff --git a/core/src/test/scala/unit/kafka/cluster/AbstractPartitionTest.scala b/core/src/test/scala/unit/kafka/cluster/AbstractPartitionTest.scala
new file mode 100644
index 0000000000000..603598e117a14
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/cluster/AbstractPartitionTest.scala
@@ -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 kafka.cluster
+
+import java.io.File
+import java.util.Properties
+
+import kafka.api.ApiVersion
+import kafka.log.{CleanerConfig, LogConfig, LogManager}
+import kafka.server.{Defaults, MetadataCache}
+import kafka.server.checkpoints.OffsetCheckpoints
+import kafka.utils.TestUtils.{MockAlterIsrManager, MockIsrChangeListener}
+import kafka.utils.{MockTime, TestUtils}
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.utils.Utils
+import org.junit.{After, Before}
+import org.mockito.ArgumentMatchers
+import org.mockito.Mockito.{mock, when}
+
+class AbstractPartitionTest {
+
+  val brokerId = 101
+  val topicPartition = new TopicPartition("test-topic", 0)
+  val time = new MockTime()
+  var tmpDir: File = _
+  var logDir1: File = _
+  var logDir2: File = _
+  var logManager: LogManager = _
+  var alterIsrManager: MockAlterIsrManager = _
+  var isrChangeListener: MockIsrChangeListener = _
+  var logConfig: LogConfig = _
+  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)
+
+    tmpDir = TestUtils.tempDir()
+    logDir1 = TestUtils.randomPartitionLogDir(tmpDir)
+    logDir2 = TestUtils.randomPartitionLogDir(tmpDir)
+    logManager = TestUtils.createLogManager(
+      logDirs = Seq(logDir1, logDir2), defaultConfig = logConfig, CleanerConfig(enableCleaner = false), time)
+    logManager.startup()
+
+    alterIsrManager = TestUtils.createAlterIsrManager()
+    isrChangeListener = TestUtils.createIsrChangeListener()
+    partition = new Partition(topicPartition,
+      replicaLagTimeMaxMs = Defaults.ReplicaLagTimeMaxMs,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      localBrokerId = brokerId,
+      time,
+      stateStore,
+      isrChangeListener,
+      delayedOperations,
+      metadataCache,
+      logManager,
+      alterIsrManager)
+
+    when(stateStore.fetchTopicConfig()).thenReturn(createLogProperties(Map.empty))
+    when(offsetCheckpoints.fetch(ArgumentMatchers.anyString, ArgumentMatchers.eq(topicPartition)))
+      .thenReturn(None)
+  }
+
+  def createLogProperties(overrides: Map[String, String]): Properties = {
+    val logProps = new Properties()
+    logProps.put(LogConfig.SegmentBytesProp, 512: java.lang.Integer)
+    logProps.put(LogConfig.SegmentIndexBytesProp, 1000: java.lang.Integer)
+    logProps.put(LogConfig.RetentionMsProp, 999: java.lang.Integer)
+    overrides.foreach { case (k, v) => logProps.put(k, v) }
+    logProps
+  }
+
+  @After
+  def tearDown(): Unit = {
+    if (tmpDir.exists()) {
+      logManager.shutdown()
+      Utils.delete(tmpDir)
+      TestUtils.clearYammerMetrics()
+    }
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/cluster/AssignmentStateTest.scala b/core/src/test/scala/unit/kafka/cluster/AssignmentStateTest.scala
new file mode 100644
index 0000000000000..83c0589096f22
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/cluster/AssignmentStateTest.scala
@@ -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 kafka.cluster
+
+import org.apache.kafka.common.message.LeaderAndIsrRequestData.LeaderAndIsrPartitionState
+import org.junit.Assert.{assertEquals, assertFalse, assertTrue}
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.Parameterized
+import org.junit.runners.Parameterized.Parameters
+
+import scala.jdk.CollectionConverters._
+
+object AssignmentStateTest extends AbstractPartitionTest {
+
+  @Parameters
+  def data: Array[Array[Any]] = Seq[Array[Any]](
+    Array(
+      List[Integer](brokerId, brokerId + 1, brokerId + 2),
+      List[Integer](brokerId, brokerId + 1, brokerId + 2),
+      List.empty[Integer], List.empty[Integer], Seq.empty[Int], false),
+    Array(
+      List[Integer](brokerId, brokerId + 1),
+      List[Integer](brokerId, brokerId + 1, brokerId + 2),
+      List.empty[Integer], List.empty[Integer], Seq.empty[Int], true),
+    Array(
+      List[Integer](brokerId, brokerId + 1, brokerId + 2),
+      List[Integer](brokerId, brokerId + 1, brokerId + 2),
+      List[Integer](brokerId + 3, brokerId + 4),
+      List[Integer](brokerId + 1),
+      Seq(brokerId, brokerId + 1, brokerId + 2), false),
+    Array(
+      List[Integer](brokerId, brokerId + 1, brokerId + 2),
+      List[Integer](brokerId, brokerId + 1, brokerId + 2),
+      List[Integer](brokerId + 3, brokerId + 4),
+      List.empty[Integer],
+      Seq(brokerId, brokerId + 1, brokerId + 2), false),
+    Array(
+      List[Integer](brokerId, brokerId + 1, brokerId + 2),
+      List[Integer](brokerId, brokerId + 1, brokerId + 2),
+      List.empty[Integer],
+      List[Integer](brokerId + 1),
+      Seq(brokerId, brokerId + 1, brokerId + 2), false),
+    Array(
+      List[Integer](brokerId + 1, brokerId + 2),
+      List[Integer](brokerId + 1, brokerId + 2),
+      List[Integer](brokerId),
+      List.empty[Integer],
+      Seq(brokerId + 1, brokerId + 2), false),
+    Array(
+      List[Integer](brokerId + 2, brokerId + 3, brokerId + 4),
+      List[Integer](brokerId, brokerId + 1, brokerId + 2),
+      List[Integer](brokerId + 3, brokerId + 4, brokerId + 5),
+      List.empty[Integer],
+      Seq(brokerId, brokerId + 1, brokerId + 2), false),
+    Array(
+      List[Integer](brokerId + 2, brokerId + 3, brokerId + 4),
+      List[Integer](brokerId, brokerId + 1, brokerId + 2),
+      List[Integer](brokerId + 3, brokerId + 4, brokerId + 5),
+      List.empty[Integer],
+      Seq(brokerId, brokerId + 1, brokerId + 2), false),
+    Array(
+      List[Integer](brokerId + 2, brokerId + 3),
+      List[Integer](brokerId, brokerId + 1, brokerId + 2),
+      List[Integer](brokerId + 3, brokerId + 4, brokerId + 5),
+      List.empty[Integer],
+      Seq(brokerId, brokerId + 1, brokerId + 2), true)
+  ).toArray
+}
+
+@RunWith(classOf[Parameterized])
+class AssignmentStateTest(isr: List[Integer], replicas: List[Integer],
+                          adding: List[Integer], removing: List[Integer],
+                          original: Seq[Int], isUnderReplicated: Boolean) extends AbstractPartitionTest {
+
+  @Test
+  def testPartitionAssignmentStatus(): Unit = {
+    val controllerEpoch = 3
+
+    val leaderState = new LeaderAndIsrPartitionState()
+      .setControllerEpoch(controllerEpoch)
+      .setLeader(brokerId)
+      .setLeaderEpoch(6)
+      .setIsr(isr.asJava)
+      .setZkVersion(1)
+      .setReplicas(replicas.asJava)
+      .setIsNew(false)
+    if (adding.nonEmpty)
+      leaderState.setAddingReplicas(adding.asJava)
+    if (removing.nonEmpty)
+      leaderState.setRemovingReplicas(removing.asJava)
+
+    val isReassigning = adding.nonEmpty || removing.nonEmpty
+
+    // set the original replicas as the URP calculation will need them
+    if (original.nonEmpty)
+      partition.assignmentState = SimpleAssignmentState(original)
+    // do the test
+    partition.makeLeader(leaderState, offsetCheckpoints)
+    assertEquals(isReassigning, partition.isReassigning)
+    if (adding.nonEmpty)
+      adding.foreach(r => assertTrue(partition.isAddingReplica(r)))
+    if (adding.contains(brokerId))
+      assertTrue(partition.isAddingLocalReplica)
+    else
+      assertFalse(partition.isAddingLocalReplica)
+
+    assertEquals(isUnderReplicated, partition.isUnderReplicated)
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/cluster/BrokerEndPointTest.scala b/core/src/test/scala/unit/kafka/cluster/BrokerEndPointTest.scala
index a563c03ce59ec..3708f73aedb64 100644
--- a/core/src/test/scala/unit/kafka/cluster/BrokerEndPointTest.scala
+++ b/core/src/test/scala/unit/kafka/cluster/BrokerEndPointTest.scala
@@ -17,12 +17,19 @@
 
 package kafka.cluster
 
+import java.nio.charset.StandardCharsets
+
 import kafka.utils.TestUtils
+import kafka.zk.BrokerIdZNode
+import org.apache.kafka.common.feature.{Features, SupportedVersionRange}
+import org.apache.kafka.common.feature.Features._
 import org.apache.kafka.common.network.ListenerName
 import org.apache.kafka.common.security.auth.SecurityProtocol
 import org.junit.Assert.{assertEquals, assertNotEquals, assertNull}
 import org.junit.Test
 
+import scala.jdk.CollectionConverters._
+
 class BrokerEndPointTest {
 
   @Test
@@ -44,7 +51,7 @@ class BrokerEndPointTest {
 
   @Test
   def testFromJsonFutureVersion(): Unit = {
-    // `createBroker` should support future compatible versions, we use a hypothetical future version here
+    // Future compatible versions should be supported, we use a hypothetical future version here
     val brokerInfoStr = """{
       "foo":"bar",
       "version":100,
@@ -54,9 +61,9 @@ class BrokerEndPointTest {
       "timestamp":"1416974968782",
       "endpoints":["SSL://localhost:9093"]
     }"""
-    val broker = Broker.createBroker(1, brokerInfoStr)
+    val broker = parseBrokerJson(1, brokerInfoStr)
     assertEquals(1, broker.id)
-    val brokerEndPoint = broker.getBrokerEndPoint(ListenerName.forSecurityProtocol(SecurityProtocol.SSL))
+    val brokerEndPoint = broker.brokerEndPoint(ListenerName.forSecurityProtocol(SecurityProtocol.SSL))
     assertEquals("localhost", brokerEndPoint.host)
     assertEquals(9093, brokerEndPoint.port)
   }
@@ -71,9 +78,9 @@ class BrokerEndPointTest {
       "timestamp":"1416974968782",
       "endpoints":["PLAINTEXT://localhost:9092"]
     }"""
-    val broker = Broker.createBroker(1, brokerInfoStr)
+    val broker = parseBrokerJson(1, brokerInfoStr)
     assertEquals(1, broker.id)
-    val brokerEndPoint = broker.getBrokerEndPoint(ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT))
+    val brokerEndPoint = broker.brokerEndPoint(ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT))
     assertEquals("localhost", brokerEndPoint.host)
     assertEquals(9092, brokerEndPoint.port)
   }
@@ -81,9 +88,9 @@ class BrokerEndPointTest {
   @Test
   def testFromJsonV1(): Unit = {
     val brokerInfoStr = """{"jmx_port":-1,"timestamp":"1420485325400","host":"172.16.8.243","version":1,"port":9091}"""
-    val broker = Broker.createBroker(1, brokerInfoStr)
+    val broker = parseBrokerJson(1, brokerInfoStr)
     assertEquals(1, broker.id)
-    val brokerEndPoint = broker.getBrokerEndPoint(ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT))
+    val brokerEndPoint = broker.brokerEndPoint(ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT))
     assertEquals("172.16.8.243", brokerEndPoint.host)
     assertEquals(9091, brokerEndPoint.port)
   }
@@ -99,9 +106,9 @@ class BrokerEndPointTest {
       "endpoints":["PLAINTEXT://host1:9092", "SSL://host1:9093"],
       "rack":"dc1"
     }"""
-    val broker = Broker.createBroker(1, json)
+    val broker = parseBrokerJson(1, json)
     assertEquals(1, broker.id)
-    val brokerEndPoint = broker.getBrokerEndPoint(ListenerName.forSecurityProtocol(SecurityProtocol.SSL))
+    val brokerEndPoint = broker.brokerEndPoint(ListenerName.forSecurityProtocol(SecurityProtocol.SSL))
     assertEquals("host1", brokerEndPoint.host)
     assertEquals(9093, brokerEndPoint.port)
     assertEquals(Some("dc1"), broker.rack)
@@ -119,9 +126,9 @@ class BrokerEndPointTest {
       "listener_security_protocol_map":{"CLIENT":"SSL", "REPLICATION":"PLAINTEXT"},
       "rack":null
     }"""
-    val broker = Broker.createBroker(1, json)
+    val broker = parseBrokerJson(1, json)
     assertEquals(1, broker.id)
-    val brokerEndPoint = broker.getBrokerEndPoint(new ListenerName("CLIENT"))
+    val brokerEndPoint = broker.brokerEndPoint(new ListenerName("CLIENT"))
     assertEquals("host1", brokerEndPoint.host)
     assertEquals(9092, brokerEndPoint.port)
     assertEquals(None, broker.rack)
@@ -138,14 +145,61 @@ class BrokerEndPointTest {
       "endpoints":["CLIENT://host1:9092", "REPLICATION://host1:9093"],
       "listener_security_protocol_map":{"CLIENT":"SSL", "REPLICATION":"PLAINTEXT"}
     }"""
-    val broker = Broker.createBroker(1, json)
+    val broker = parseBrokerJson(1, json)
     assertEquals(1, broker.id)
-    val brokerEndPoint = broker.getBrokerEndPoint(new ListenerName("CLIENT"))
+    val brokerEndPoint = broker.brokerEndPoint(new ListenerName("CLIENT"))
     assertEquals("host1", brokerEndPoint.host)
     assertEquals(9092, brokerEndPoint.port)
     assertEquals(None, broker.rack)
   }
 
+  @Test
+  def testFromJsonV4WithNoFeatures(): Unit = {
+    val json = """{
+      "version":4,
+      "host":"localhost",
+      "port":9092,
+      "jmx_port":9999,
+      "timestamp":"2233345666",
+      "endpoints":["CLIENT://host1:9092", "REPLICATION://host1:9093"],
+      "listener_security_protocol_map":{"CLIENT":"SSL", "REPLICATION":"PLAINTEXT"},
+      "rack":"dc1"
+    }"""
+    val broker = parseBrokerJson(1, json)
+    assertEquals(1, broker.id)
+    val brokerEndPoint = broker.brokerEndPoint(new ListenerName("CLIENT"))
+    assertEquals("host1", brokerEndPoint.host)
+    assertEquals(9092, brokerEndPoint.port)
+    assertEquals(Some("dc1"), broker.rack)
+    assertEquals(emptySupportedFeatures, broker.features)
+  }
+
+  @Test
+  def testFromJsonV5(): Unit = {
+    val json = """{
+      "version":5,
+      "host":"localhost",
+      "port":9092,
+      "jmx_port":9999,
+      "timestamp":"2233345666",
+      "endpoints":["CLIENT://host1:9092", "REPLICATION://host1:9093"],
+      "listener_security_protocol_map":{"CLIENT":"SSL", "REPLICATION":"PLAINTEXT"},
+      "rack":"dc1",
+      "features": {"feature1": {"min_version": 1, "max_version": 2}, "feature2": {"min_version": 2, "max_version": 4}}
+    }"""
+    val broker = parseBrokerJson(1, json)
+    assertEquals(1, broker.id)
+    val brokerEndPoint = broker.brokerEndPoint(new ListenerName("CLIENT"))
+    assertEquals("host1", brokerEndPoint.host)
+    assertEquals(9092, brokerEndPoint.port)
+    assertEquals(Some("dc1"), broker.rack)
+    assertEquals(Features.supportedFeatures(
+      Map[String, SupportedVersionRange](
+        "feature1" -> new SupportedVersionRange(1, 2),
+        "feature2" -> new SupportedVersionRange(2, 4)).asJava),
+      broker.features)
+  }
+
   @Test
   def testBrokerEndpointFromUri(): Unit = {
     var connectionString = "localhost:9092"
@@ -212,4 +266,7 @@ class BrokerEndPointTest {
     assertEquals(9092, endpoint.port)
     assertEquals("PLAINTEXT://MyHostname:9092", endpoint.connectionString)
   }
+
+  private def parseBrokerJson(id: Int, jsonString: String): Broker =
+    BrokerIdZNode.decode(id, jsonString.getBytes(StandardCharsets.UTF_8)).broker
 }
diff --git a/core/src/test/scala/unit/kafka/cluster/PartitionLockTest.scala b/core/src/test/scala/unit/kafka/cluster/PartitionLockTest.scala
new file mode 100644
index 0000000000000..e5dbec7745d9f
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/cluster/PartitionLockTest.scala
@@ -0,0 +1,368 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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.cluster
+
+import java.util.Properties
+import java.util.concurrent.atomic.AtomicBoolean
+import java.util.concurrent._
+
+import kafka.api.{ApiVersion, LeaderAndIsr}
+import kafka.log._
+import kafka.server._
+import kafka.server.checkpoints.OffsetCheckpoints
+import kafka.utils._
+import org.apache.kafka.common.message.LeaderAndIsrRequestData.LeaderAndIsrPartitionState
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.record.{MemoryRecords, SimpleRecord}
+import org.apache.kafka.common.utils.Utils
+import org.junit.Assert.{assertEquals, assertFalse, assertTrue}
+import org.junit.{After, Before, Test}
+import org.mockito.ArgumentMatchers
+import org.mockito.Mockito.{mock, when}
+
+import scala.jdk.CollectionConverters._
+import scala.concurrent.duration._
+
+/**
+ * Verifies that slow appends to log don't block request threads processing replica fetch requests.
+ *
+ * Test simulates:
+ *   1) Produce request handling by performing append to log as leader.
+ *   2) Replica fetch request handling by processing update of ISR and HW based on log read result.
+ *   3) Some tests also simulate a scheduler thread that checks or updates ISRs when replica falls out of ISR
+ */
+class PartitionLockTest extends Logging {
+
+  val numReplicaFetchers = 2
+  val numProducers = 3
+  val numRecordsPerProducer = 5
+
+  val mockTime = new MockTime() // Used for check to shrink ISR
+  val tmpDir = TestUtils.tempDir()
+  val logDir = TestUtils.randomPartitionLogDir(tmpDir)
+  val executorService = Executors.newFixedThreadPool(numReplicaFetchers + numProducers + 1)
+  val appendSemaphore = new Semaphore(0)
+  val shrinkIsrSemaphore = new Semaphore(0)
+  val followerQueues = (0 until numReplicaFetchers).map(_ => new ArrayBlockingQueue[MemoryRecords](2))
+
+  var logManager: LogManager = _
+  var partition: Partition = _
+
+  @Before
+  def setUp(): Unit = {
+    val logConfig = new LogConfig(new Properties)
+    logManager = TestUtils.createLogManager(Seq(logDir), logConfig, CleanerConfig(enableCleaner = false), mockTime)
+    partition = setupPartitionWithMocks(logManager, logConfig)
+  }
+
+  @After
+  def tearDown(): Unit = {
+    executorService.shutdownNow()
+    logManager.liveLogDirs.foreach(Utils.delete)
+    Utils.delete(tmpDir)
+  }
+
+  /**
+   * Verifies that delays in appending to leader while processing produce requests has no impact on timing
+   * of update of log read result when processing replica fetch request if no ISR update is required.
+   */
+  @Test
+  def testNoLockContentionWithoutIsrUpdate(): Unit = {
+    concurrentProduceFetchWithReadLockOnly()
+  }
+
+  /**
+   * Verifies that delays in appending to leader while processing produce requests has no impact on timing
+   * of update of log read result when processing replica fetch request even if a scheduler thread is checking
+   * for ISR shrink conditions if no ISR update is required.
+   */
+  @Test
+  def testAppendReplicaFetchWithSchedulerCheckForShrinkIsr(): Unit = {
+    val active = new AtomicBoolean(true)
+
+    val future = scheduleShrinkIsr(active, mockTimeSleepMs = 0)
+    concurrentProduceFetchWithReadLockOnly()
+    active.set(false)
+    future.get(15, TimeUnit.SECONDS)
+  }
+
+  /**
+   * Verifies concurrent produce and replica fetch log read result update with ISR updates. This
+   * can result in delays in processing produce and replica fetch requets since write lock is obtained,
+   * but it should complete without any failures.
+   */
+  @Test
+  def testAppendReplicaFetchWithUpdateIsr(): Unit = {
+    val active = new AtomicBoolean(true)
+
+    val future = scheduleShrinkIsr(active, mockTimeSleepMs = 10000)
+    TestUtils.waitUntilTrue(() => shrinkIsrSemaphore.hasQueuedThreads, "shrinkIsr not invoked")
+    concurrentProduceFetchWithWriteLock()
+    active.set(false)
+    future.get(15, TimeUnit.SECONDS)
+  }
+
+  /**
+   * Concurrently calling updateAssignmentAndIsr should always ensure that non-lock access
+   * to the inner remoteReplicaMap (accessed by getReplica) cannot see an intermediate state
+   * where replicas present both in the old and new assignment are missing
+   */
+  @Test
+  def testGetReplicaWithUpdateAssignmentAndIsr(): Unit = {
+    val active = new AtomicBoolean(true)
+    val replicaToCheck = 3
+    val firstReplicaSet = Seq[Integer](3, 4, 5).asJava
+    val secondReplicaSet = Seq[Integer](1, 2, 3).asJava
+    def partitionState(replicas: java.util.List[Integer]) = new LeaderAndIsrPartitionState()
+      .setControllerEpoch(1)
+      .setLeader(replicas.get(0))
+      .setLeaderEpoch(1)
+      .setIsr(replicas)
+      .setZkVersion(1)
+      .setReplicas(replicas)
+      .setIsNew(true)
+    val offsetCheckpoints: OffsetCheckpoints = mock(classOf[OffsetCheckpoints])
+    // Update replica set synchronously first to avoid race conditions
+    partition.makeLeader(partitionState(secondReplicaSet), offsetCheckpoints)
+    assertTrue(s"Expected replica $replicaToCheck to be defined", partition.getReplica(replicaToCheck).isDefined)
+
+    val future = executorService.submit((() => {
+      var i = 0
+      // Flip assignment between two replica sets
+      while (active.get) {
+        val replicas = if (i % 2 == 0) {
+          firstReplicaSet
+        } else {
+          secondReplicaSet
+        }
+
+        partition.makeLeader(partitionState(replicas), offsetCheckpoints)
+
+        i += 1
+        Thread.sleep(1) // just to avoid tight loop
+      }
+    }): Runnable)
+
+    val deadline = 1.seconds.fromNow
+    while (deadline.hasTimeLeft()) {
+      assertTrue(s"Expected replica $replicaToCheck to be defined", partition.getReplica(replicaToCheck).isDefined)
+    }
+    active.set(false)
+    future.get(5, TimeUnit.SECONDS)
+    assertTrue(s"Expected replica $replicaToCheck to be defined", partition.getReplica(replicaToCheck).isDefined)
+  }
+
+  /**
+   * Perform concurrent appends and replica fetch requests that don't require write lock to
+   * update follower state. Release sufficient append permits to complete all except one append.
+   * Verify that follower state updates complete even though an append holding read lock is in progress.
+   * Then release the permit for the final append and verify that all appends and follower updates complete.
+   */
+  private def concurrentProduceFetchWithReadLockOnly(): Unit = {
+    val appendFutures = scheduleAppends()
+    val stateUpdateFutures = scheduleUpdateFollowers(numProducers * numRecordsPerProducer - 1)
+
+    appendSemaphore.release(numProducers * numRecordsPerProducer - 1)
+    stateUpdateFutures.foreach(_.get(15, TimeUnit.SECONDS))
+
+    appendSemaphore.release(1)
+    scheduleUpdateFollowers(1).foreach(_.get(15, TimeUnit.SECONDS)) // just to make sure follower state update still works
+    appendFutures.foreach(_.get(15, TimeUnit.SECONDS))
+  }
+
+  /**
+   * Perform concurrent appends and replica fetch requests that may require write lock to update
+   * follower state. Threads waiting for write lock to update follower state while append thread is
+   * holding read lock will prevent other threads acquiring the read or write lock. So release sufficient
+   * permits for all appends to complete before verifying state updates.
+   */
+  private def concurrentProduceFetchWithWriteLock(): Unit = {
+
+    val appendFutures = scheduleAppends()
+    val stateUpdateFutures = scheduleUpdateFollowers(numProducers * numRecordsPerProducer)
+
+    assertFalse(stateUpdateFutures.exists(_.isDone))
+    appendSemaphore.release(numProducers * numRecordsPerProducer)
+    assertFalse(appendFutures.exists(_.isDone))
+
+    shrinkIsrSemaphore.release()
+    stateUpdateFutures.foreach(_.get(15, TimeUnit.SECONDS))
+    appendFutures.foreach(_.get(15, TimeUnit.SECONDS))
+  }
+
+  private def scheduleAppends(): Seq[Future[_]] = {
+    (0 until numProducers).map { _ =>
+      executorService.submit((() => {
+        try {
+          append(partition, numRecordsPerProducer, followerQueues)
+        } catch {
+          case e: Throwable =>
+            error("Exception during append", e)
+            throw e
+        }
+      }): Runnable)
+    }
+  }
+
+  private def scheduleUpdateFollowers(numRecords: Int): Seq[Future[_]] = {
+    (1 to numReplicaFetchers).map { index =>
+      executorService.submit((() => {
+        try {
+          updateFollowerFetchState(partition, index, numRecords, followerQueues(index - 1))
+        } catch {
+          case e: Throwable =>
+            error("Exception during updateFollowerFetchState", e)
+            throw e
+        }
+      }): Runnable)
+    }
+  }
+
+  private def scheduleShrinkIsr(activeFlag: AtomicBoolean, mockTimeSleepMs: Long): Future[_] = {
+    executorService.submit((() => {
+      while (activeFlag.get) {
+        if (mockTimeSleepMs > 0)
+          mockTime.sleep(mockTimeSleepMs)
+        partition.maybeShrinkIsr()
+        Thread.sleep(1) // just to avoid tight loop
+      }
+    }): Runnable)
+  }
+
+  private def setupPartitionWithMocks(logManager: LogManager, logConfig: LogConfig): Partition = {
+    val leaderEpoch = 1
+    val brokerId = 0
+    val topicPartition = new TopicPartition("test-topic", 0)
+    val stateStore: PartitionStateStore = mock(classOf[PartitionStateStore])
+    val isrChangeListener: IsrChangeListener = mock(classOf[IsrChangeListener])
+    val delayedOperations: DelayedOperations = mock(classOf[DelayedOperations])
+    val metadataCache: MetadataCache = mock(classOf[MetadataCache])
+    val offsetCheckpoints: OffsetCheckpoints = mock(classOf[OffsetCheckpoints])
+    val alterIsrManager: AlterIsrManager = mock(classOf[AlterIsrManager])
+
+    logManager.startup()
+    val partition = new Partition(topicPartition,
+      replicaLagTimeMaxMs = kafka.server.Defaults.ReplicaLagTimeMaxMs,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      localBrokerId = brokerId,
+      mockTime,
+      stateStore,
+      isrChangeListener,
+      delayedOperations,
+      metadataCache,
+      logManager,
+      alterIsrManager) {
+
+      override def shrinkIsr(newIsr: Set[Int]): Unit = {
+        shrinkIsrSemaphore.acquire()
+        try {
+          super.shrinkIsr(newIsr)
+        } finally {
+          shrinkIsrSemaphore.release()
+        }
+      }
+
+      override def createLog(isNew: Boolean, isFutureReplica: Boolean, offsetCheckpoints: OffsetCheckpoints): Log = {
+        val log = super.createLog(isNew, isFutureReplica, offsetCheckpoints)
+        new SlowLog(log, mockTime, appendSemaphore)
+      }
+    }
+    when(stateStore.fetchTopicConfig()).thenReturn(createLogProperties(Map.empty))
+    when(offsetCheckpoints.fetch(ArgumentMatchers.anyString, ArgumentMatchers.eq(topicPartition)))
+      .thenReturn(None)
+    when(stateStore.shrinkIsr(ArgumentMatchers.anyInt, ArgumentMatchers.any[LeaderAndIsr]))
+      .thenReturn(Some(2))
+    when(stateStore.expandIsr(ArgumentMatchers.anyInt, ArgumentMatchers.any[LeaderAndIsr]))
+      .thenReturn(Some(2))
+    when(alterIsrManager.enqueue(ArgumentMatchers.any[AlterIsrItem]))
+      .thenReturn(true)
+
+    partition.createLogIfNotExists(isNew = false, isFutureReplica = false, offsetCheckpoints)
+
+    val controllerEpoch = 0
+    val replicas = (0 to numReplicaFetchers).map(i => Integer.valueOf(brokerId + i)).toList.asJava
+    val isr = replicas
+
+    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))
+
+    partition
+  }
+
+  private def createLogProperties(overrides: Map[String, String]): Properties = {
+    val logProps = new Properties()
+    logProps.put(LogConfig.SegmentBytesProp, 512: java.lang.Integer)
+    logProps.put(LogConfig.SegmentIndexBytesProp, 1000: java.lang.Integer)
+    logProps.put(LogConfig.RetentionMsProp, 999: java.lang.Integer)
+    overrides.foreach { case (k, v) => logProps.put(k, v) }
+    logProps
+  }
+
+  private def append(partition: Partition, numRecords: Int, followerQueues: Seq[ArrayBlockingQueue[MemoryRecords]]): Unit = {
+    (0 until numRecords).foreach { _ =>
+      val batch = TestUtils.records(records = List(new SimpleRecord("k1".getBytes, "v1".getBytes),
+        new SimpleRecord("k2".getBytes, "v2".getBytes)))
+      partition.appendRecordsToLeader(batch, origin = AppendOrigin.Client, requiredAcks = 0)
+      followerQueues.foreach(_.put(batch))
+    }
+  }
+
+  private def updateFollowerFetchState(partition: Partition, followerId: Int, numRecords: Int, followerQueue: ArrayBlockingQueue[MemoryRecords]): Unit = {
+    (1 to numRecords).foreach { i =>
+      val batch = followerQueue.poll(15, TimeUnit.SECONDS)
+      if (batch == null)
+        throw new RuntimeException(s"Timed out waiting for next batch $i")
+      val batches = batch.batches.iterator.asScala.toList
+      assertEquals(1, batches.size)
+      val recordBatch = batches.head
+      partition.updateFollowerFetchState(
+        followerId,
+        followerFetchOffsetMetadata = LogOffsetMetadata(recordBatch.lastOffset + 1),
+        followerStartOffset = 0L,
+        followerFetchTimeMs = mockTime.milliseconds(),
+        leaderEndOffset = partition.localLogOrException.logEndOffset)
+    }
+  }
+
+  private class SlowLog(log: Log, mockTime: MockTime, appendSemaphore: Semaphore) extends Log(
+    log.dir,
+    log.config,
+    log.logStartOffset,
+    log.recoveryPoint,
+    mockTime.scheduler,
+    new BrokerTopicStats,
+    log.time,
+    log.maxProducerIdExpirationMs,
+    log.producerIdExpirationCheckIntervalMs,
+    log.topicPartition,
+    log.producerStateManager,
+    new LogDirFailureChannel(1)) {
+
+    override def appendAsLeader(records: MemoryRecords, leaderEpoch: Int, origin: AppendOrigin, interBrokerProtocolVersion: ApiVersion): LogAppendInfo = {
+      val appendInfo = super.appendAsLeader(records, leaderEpoch, origin, interBrokerProtocolVersion)
+      appendSemaphore.acquire()
+      appendInfo
+    }
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/cluster/PartitionTest.scala b/core/src/test/scala/unit/kafka/cluster/PartitionTest.scala
new file mode 100644
index 0000000000000..93aed43892b94
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/cluster/PartitionTest.scala
@@ -0,0 +1,1820 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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.cluster
+
+import java.nio.ByteBuffer
+import java.util.Optional
+import java.util.concurrent.{CountDownLatch, Semaphore}
+
+import com.yammer.metrics.core.Metric
+import kafka.api.{ApiVersion, LeaderAndIsr}
+import kafka.common.UnexpectedAppendOffsetException
+import kafka.log.{Defaults => _, _}
+import kafka.metrics.KafkaYammerMetrics
+import kafka.server._
+import kafka.server.checkpoints.OffsetCheckpoints
+import kafka.utils._
+import org.apache.kafka.common.errors.{ApiException, NotLeaderOrFollowerException, OffsetNotAvailableException, OffsetOutOfRangeException}
+import org.apache.kafka.common.message.FetchResponseData
+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.record._
+import org.apache.kafka.common.requests.OffsetsForLeaderEpochResponse.{UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET}
+import org.apache.kafka.common.requests.ListOffsetRequest
+import org.apache.kafka.common.utils.SystemTime
+import org.apache.kafka.common.{IsolationLevel, TopicPartition}
+import org.junit.Assert._
+import org.junit.Test
+import org.mockito.ArgumentMatchers
+import org.mockito.Mockito._
+import org.mockito.invocation.InvocationOnMock
+import org.scalatest.Assertions.assertThrows
+
+import scala.jdk.CollectionConverters._
+
+class PartitionTest extends AbstractPartitionTest {
+
+  @Test
+  def testLastFetchedOffsetValidation(): Unit = {
+    val log = logManager.getOrCreateLog(topicPartition, () => logConfig)
+    def append(leaderEpoch: Int, count: Int): Unit = {
+      val recordArray = (1 to count).map { i =>
+        new SimpleRecord(s"$i".getBytes)
+      }
+      val records = MemoryRecords.withRecords(0L, CompressionType.NONE, leaderEpoch,
+        recordArray: _*)
+      log.appendAsLeader(records, leaderEpoch = leaderEpoch)
+    }
+
+    append(leaderEpoch = 0, count = 2) // 0
+    append(leaderEpoch = 3, count = 3) // 2
+    append(leaderEpoch = 3, count = 3) // 5
+    append(leaderEpoch = 4, count = 5) // 8
+    append(leaderEpoch = 7, count = 1) // 13
+    append(leaderEpoch = 9, count = 3) // 14
+    assertEquals(17L, log.logEndOffset)
+
+    val leaderEpoch = 10
+    val partition = setupPartitionWithMocks(leaderEpoch = leaderEpoch, isLeader = true, log = log)
+
+    def epochEndOffset(epoch: Int, endOffset: Long): FetchResponseData.EpochEndOffset = {
+      new FetchResponseData.EpochEndOffset()
+        .setEpoch(epoch)
+        .setEndOffset(endOffset)
+    }
+
+    def read(lastFetchedEpoch: Int, fetchOffset: Long): LogReadInfo = {
+      partition.readRecords(
+        Optional.of(lastFetchedEpoch),
+        fetchOffset,
+        currentLeaderEpoch = Optional.of(leaderEpoch),
+        maxBytes = Int.MaxValue,
+        fetchIsolation = FetchLogEnd,
+        fetchOnlyFromLeader = true,
+        minOneMessage = true
+      )
+    }
+
+    def assertDivergence(
+      divergingEpoch: FetchResponseData.EpochEndOffset,
+      readInfo: LogReadInfo
+    ): Unit = {
+      assertEquals(Some(divergingEpoch), readInfo.divergingEpoch)
+      assertEquals(0, readInfo.fetchedData.records.sizeInBytes)
+    }
+
+    def assertNoDivergence(readInfo: LogReadInfo): Unit = {
+      assertEquals(None, readInfo.divergingEpoch)
+    }
+
+    assertDivergence(epochEndOffset(epoch = 0, endOffset = 2), read(lastFetchedEpoch = 2, fetchOffset = 5))
+    assertDivergence(epochEndOffset(epoch = 0, endOffset= 2), read(lastFetchedEpoch = 0, fetchOffset = 4))
+    assertDivergence(epochEndOffset(epoch = 4, endOffset = 13), read(lastFetchedEpoch = 6, fetchOffset = 6))
+    assertDivergence(epochEndOffset(epoch = 4, endOffset = 13), read(lastFetchedEpoch = 5, fetchOffset = 9))
+    assertDivergence(epochEndOffset(epoch = 10, endOffset = 17), read(lastFetchedEpoch = 10, fetchOffset = 18))
+    assertNoDivergence(read(lastFetchedEpoch = 0, fetchOffset = 2))
+    assertNoDivergence(read(lastFetchedEpoch = 7, fetchOffset = 14))
+    assertNoDivergence(read(lastFetchedEpoch = 9, fetchOffset = 17))
+    assertNoDivergence(read(lastFetchedEpoch = 10, fetchOffset = 17))
+
+    // Reads from epochs larger than we know about should cause an out of range error
+    assertThrows[OffsetOutOfRangeException] {
+      read(lastFetchedEpoch = 11, fetchOffset = 5)
+    }
+
+    // Move log start offset to the middle of epoch 3
+    log.updateHighWatermark(log.logEndOffset)
+    log.maybeIncrementLogStartOffset(newLogStartOffset = 5L, ClientRecordDeletion)
+
+    assertDivergence(epochEndOffset(epoch = 2, endOffset = 5), read(lastFetchedEpoch = 2, fetchOffset = 8))
+    assertNoDivergence(read(lastFetchedEpoch = 0, fetchOffset = 5))
+    assertNoDivergence(read(lastFetchedEpoch = 3, fetchOffset = 5))
+
+    assertThrows[OffsetOutOfRangeException] {
+      read(lastFetchedEpoch = 0, fetchOffset = 0)
+    }
+  }
+
+  @Test
+  def testMakeLeaderUpdatesEpochCache(): Unit = {
+    val leaderEpoch = 8
+
+    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)
+    ), leaderEpoch = 0)
+    log.appendAsLeader(MemoryRecords.withRecords(0L, CompressionType.NONE, 5,
+      new SimpleRecord("k3".getBytes, "v3".getBytes),
+      new SimpleRecord("k4".getBytes, "v4".getBytes)
+    ), leaderEpoch = 5)
+    assertEquals(4, log.logEndOffset)
+
+    val partition = setupPartitionWithMocks(leaderEpoch = leaderEpoch, isLeader = true, log = log)
+    assertEquals(Some(4), partition.leaderLogIfLocal.map(_.logEndOffset))
+
+    val epochEndOffset = partition.lastOffsetForLeaderEpoch(currentLeaderEpoch = Optional.of[Integer](leaderEpoch),
+      leaderEpoch = leaderEpoch, fetchOnlyFromLeader = true)
+    assertEquals(4, epochEndOffset.endOffset)
+    assertEquals(leaderEpoch, epochEndOffset.leaderEpoch)
+  }
+
+  @Test
+  def testMakeLeaderDoesNotUpdateEpochCacheForOldFormats(): Unit = {
+    val leaderEpoch = 8
+
+    val logConfig = LogConfig(createLogProperties(Map(
+      LogConfig.MessageFormatVersionProp -> kafka.api.KAFKA_0_10_2_IV0.shortVersion)))
+    val log = logManager.getOrCreateLog(topicPartition, () => logConfig)
+    log.appendAsLeader(TestUtils.records(List(
+      new SimpleRecord("k1".getBytes, "v1".getBytes),
+      new SimpleRecord("k2".getBytes, "v2".getBytes)),
+      magicValue = RecordVersion.V1.value
+    ), leaderEpoch = 0)
+    log.appendAsLeader(TestUtils.records(List(
+      new SimpleRecord("k3".getBytes, "v3".getBytes),
+      new SimpleRecord("k4".getBytes, "v4".getBytes)),
+      magicValue = RecordVersion.V1.value
+    ), leaderEpoch = 5)
+    assertEquals(4, log.logEndOffset)
+
+    val partition = setupPartitionWithMocks(leaderEpoch = leaderEpoch, isLeader = true, log = log)
+    assertEquals(Some(4), partition.leaderLogIfLocal.map(_.logEndOffset))
+    assertEquals(None, log.latestEpoch)
+
+    val epochEndOffset = partition.lastOffsetForLeaderEpoch(currentLeaderEpoch = Optional.of(leaderEpoch),
+      leaderEpoch = leaderEpoch, fetchOnlyFromLeader = true)
+    assertEquals(UNDEFINED_EPOCH_OFFSET, epochEndOffset.endOffset)
+    assertEquals(UNDEFINED_EPOCH, epochEndOffset.leaderEpoch)
+  }
+
+  @Test
+  // Verify that partition.removeFutureLocalReplica() and partition.maybeReplaceCurrentWithFutureReplica() can run concurrently
+  def testMaybeReplaceCurrentWithFutureReplica(): Unit = {
+    val latch = new CountDownLatch(1)
+
+    logManager.maybeUpdatePreferredLogDir(topicPartition, logDir1.getAbsolutePath)
+    partition.createLogIfNotExists(isNew = true, isFutureReplica = false, offsetCheckpoints)
+    logManager.maybeUpdatePreferredLogDir(topicPartition, logDir2.getAbsolutePath)
+    partition.maybeCreateFutureReplica(logDir2.getAbsolutePath, offsetCheckpoints)
+
+    val thread1 = new Thread {
+      override def run(): Unit = {
+        latch.await()
+        partition.removeFutureLocalReplica()
+      }
+    }
+
+    val thread2 = new Thread {
+      override def run(): Unit = {
+        latch.await()
+        partition.maybeReplaceCurrentWithFutureReplica()
+      }
+    }
+
+    thread1.start()
+    thread2.start()
+
+    latch.countDown()
+    thread1.join()
+    thread2.join()
+    assertEquals(None, partition.futureLog)
+  }
+
+  // Verify that partition.makeFollower() and partition.appendRecordsToFollowerOrFutureReplica() can run concurrently
+  @Test
+  def testMakeFollowerWithWithFollowerAppendRecords(): Unit = {
+    val appendSemaphore = new Semaphore(0)
+    val mockTime = new MockTime()
+
+    partition = new Partition(
+      topicPartition,
+      replicaLagTimeMaxMs = Defaults.ReplicaLagTimeMaxMs,
+      interBrokerProtocolVersion = ApiVersion.latestVersion,
+      localBrokerId = brokerId,
+      time,
+      stateStore,
+      isrChangeListener,
+      delayedOperations,
+      metadataCache,
+      logManager,
+      alterIsrManager) {
+
+      override def createLog(isNew: Boolean, isFutureReplica: Boolean, offsetCheckpoints: OffsetCheckpoints): Log = {
+        val log = super.createLog(isNew, isFutureReplica, offsetCheckpoints)
+        new SlowLog(log, mockTime, appendSemaphore)
+      }
+    }
+
+    partition.createLogIfNotExists(isNew = true, isFutureReplica = false, offsetCheckpoints)
+
+    val appendThread = new Thread {
+      override def run(): Unit = {
+        val records = createRecords(List(new SimpleRecord("k1".getBytes, "v1".getBytes),
+          new SimpleRecord("k2".getBytes, "v2".getBytes)),
+          baseOffset = 0)
+        partition.appendRecordsToFollowerOrFutureReplica(records, isFuture = false)
+      }
+    }
+    appendThread.start()
+    TestUtils.waitUntilTrue(() => appendSemaphore.hasQueuedThreads, "follower log append is not called.")
+
+    val partitionState = new LeaderAndIsrPartitionState()
+      .setControllerEpoch(0)
+      .setLeader(2)
+      .setLeaderEpoch(1)
+      .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))
+
+    appendSemaphore.release()
+    appendThread.join()
+
+    assertEquals(2L, partition.localLogOrException.logEndOffset)
+    assertEquals(2L, partition.leaderReplicaIdOpt.get)
+  }
+
+  @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 = {
+    logManager.maybeUpdatePreferredLogDir(topicPartition, logDir1.getAbsolutePath)
+    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),
+      new SimpleRecord("k2".getBytes, "v4".getBytes),
+      new SimpleRecord("k2".getBytes, "v5".getBytes),
+      new SimpleRecord("k2".getBytes, "v6".getBytes)
+    ), leaderEpoch = 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
+
+    val buffer = ByteBuffer.allocate(1024)
+    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))
+
+    val futureLog = partition.futureLocalLogOrException
+    futureLog.appendAsFollower(builder.build())
+
+    assertTrue(partition.maybeReplaceCurrentWithFutureReplica())
+  }
+
+  @Test
+  def testFetchOffsetSnapshotEpochValidationForLeader(): Unit = {
+    val leaderEpoch = 5
+    val partition = setupPartitionWithMocks(leaderEpoch, isLeader = true)
+
+    def assertSnapshotError(expectedError: Errors, currentLeaderEpoch: Optional[Integer]): Unit = {
+      try {
+        partition.fetchOffsetSnapshot(currentLeaderEpoch, fetchOnlyFromLeader = true)
+        assertEquals(Errors.NONE, expectedError)
+      } catch {
+        case error: ApiException => assertEquals(expectedError, Errors.forException(error))
+      }
+    }
+
+    assertSnapshotError(Errors.FENCED_LEADER_EPOCH, Optional.of(leaderEpoch - 1))
+    assertSnapshotError(Errors.UNKNOWN_LEADER_EPOCH, Optional.of(leaderEpoch + 1))
+    assertSnapshotError(Errors.NONE, Optional.of(leaderEpoch))
+    assertSnapshotError(Errors.NONE, Optional.empty())
+  }
+
+  @Test
+  def testFetchOffsetSnapshotEpochValidationForFollower(): Unit = {
+    val leaderEpoch = 5
+    val partition = setupPartitionWithMocks(leaderEpoch, isLeader = false)
+
+    def assertSnapshotError(expectedError: Errors,
+                            currentLeaderEpoch: Optional[Integer],
+                            fetchOnlyLeader: Boolean): Unit = {
+      try {
+        partition.fetchOffsetSnapshot(currentLeaderEpoch, fetchOnlyFromLeader = fetchOnlyLeader)
+        assertEquals(Errors.NONE, expectedError)
+      } catch {
+        case error: ApiException => assertEquals(expectedError, Errors.forException(error))
+      }
+    }
+
+    assertSnapshotError(Errors.NONE, Optional.of(leaderEpoch), fetchOnlyLeader = false)
+    assertSnapshotError(Errors.NONE, Optional.empty(), fetchOnlyLeader = false)
+    assertSnapshotError(Errors.FENCED_LEADER_EPOCH, Optional.of(leaderEpoch - 1), fetchOnlyLeader = false)
+    assertSnapshotError(Errors.UNKNOWN_LEADER_EPOCH, Optional.of(leaderEpoch + 1), fetchOnlyLeader = false)
+
+    assertSnapshotError(Errors.NOT_LEADER_OR_FOLLOWER, Optional.of(leaderEpoch), fetchOnlyLeader = true)
+    assertSnapshotError(Errors.NOT_LEADER_OR_FOLLOWER, Optional.empty(), fetchOnlyLeader = true)
+    assertSnapshotError(Errors.FENCED_LEADER_EPOCH, Optional.of(leaderEpoch - 1), fetchOnlyLeader = true)
+    assertSnapshotError(Errors.UNKNOWN_LEADER_EPOCH, Optional.of(leaderEpoch + 1), fetchOnlyLeader = true)
+  }
+
+  @Test
+  def testOffsetForLeaderEpochValidationForLeader(): Unit = {
+    val leaderEpoch = 5
+    val partition = setupPartitionWithMocks(leaderEpoch, isLeader = true)
+
+    def assertLastOffsetForLeaderError(error: Errors, currentLeaderEpochOpt: Optional[Integer]): Unit = {
+      val endOffset = partition.lastOffsetForLeaderEpoch(currentLeaderEpochOpt, 0,
+        fetchOnlyFromLeader = true)
+      assertEquals(error.code, endOffset.errorCode)
+    }
+
+    assertLastOffsetForLeaderError(Errors.NONE, Optional.empty())
+    assertLastOffsetForLeaderError(Errors.NONE, Optional.of(leaderEpoch))
+    assertLastOffsetForLeaderError(Errors.FENCED_LEADER_EPOCH, Optional.of(leaderEpoch - 1))
+    assertLastOffsetForLeaderError(Errors.UNKNOWN_LEADER_EPOCH, Optional.of(leaderEpoch + 1))
+  }
+
+  @Test
+  def testOffsetForLeaderEpochValidationForFollower(): Unit = {
+    val leaderEpoch = 5
+    val partition = setupPartitionWithMocks(leaderEpoch, isLeader = false)
+
+    def assertLastOffsetForLeaderError(error: Errors,
+                                       currentLeaderEpochOpt: Optional[Integer],
+                                       fetchOnlyLeader: Boolean): Unit = {
+      val endOffset = partition.lastOffsetForLeaderEpoch(currentLeaderEpochOpt, 0,
+        fetchOnlyFromLeader = fetchOnlyLeader)
+      assertEquals(error.code, endOffset.errorCode)
+    }
+
+    assertLastOffsetForLeaderError(Errors.NONE, Optional.empty(), fetchOnlyLeader = false)
+    assertLastOffsetForLeaderError(Errors.NONE, Optional.of(leaderEpoch), fetchOnlyLeader = false)
+    assertLastOffsetForLeaderError(Errors.FENCED_LEADER_EPOCH, Optional.of(leaderEpoch - 1), fetchOnlyLeader = false)
+    assertLastOffsetForLeaderError(Errors.UNKNOWN_LEADER_EPOCH, Optional.of(leaderEpoch + 1), fetchOnlyLeader = false)
+
+    assertLastOffsetForLeaderError(Errors.NOT_LEADER_OR_FOLLOWER, Optional.empty(), fetchOnlyLeader = true)
+    assertLastOffsetForLeaderError(Errors.NOT_LEADER_OR_FOLLOWER, Optional.of(leaderEpoch), fetchOnlyLeader = true)
+    assertLastOffsetForLeaderError(Errors.FENCED_LEADER_EPOCH, Optional.of(leaderEpoch - 1), fetchOnlyLeader = true)
+    assertLastOffsetForLeaderError(Errors.UNKNOWN_LEADER_EPOCH, Optional.of(leaderEpoch + 1), fetchOnlyLeader = true)
+  }
+
+  @Test
+  def testReadRecordEpochValidationForLeader(): Unit = {
+    val leaderEpoch = 5
+    val partition = setupPartitionWithMocks(leaderEpoch, isLeader = true)
+
+    def assertReadRecordsError(error: Errors,
+                               currentLeaderEpochOpt: Optional[Integer]): Unit = {
+      try {
+        partition.readRecords(
+          lastFetchedEpoch = Optional.empty(),
+          fetchOffset = 0L,
+          currentLeaderEpoch = currentLeaderEpochOpt,
+          maxBytes = 1024,
+          fetchIsolation = FetchLogEnd,
+          fetchOnlyFromLeader = true,
+          minOneMessage = false)
+        if (error != Errors.NONE)
+          fail(s"Expected readRecords to fail with error $error")
+      } catch {
+        case e: Exception =>
+          assertEquals(error, Errors.forException(e))
+      }
+    }
+
+    assertReadRecordsError(Errors.NONE, Optional.empty())
+    assertReadRecordsError(Errors.NONE, Optional.of(leaderEpoch))
+    assertReadRecordsError(Errors.FENCED_LEADER_EPOCH, Optional.of(leaderEpoch - 1))
+    assertReadRecordsError(Errors.UNKNOWN_LEADER_EPOCH, Optional.of(leaderEpoch + 1))
+  }
+
+  @Test
+  def testReadRecordEpochValidationForFollower(): Unit = {
+    val leaderEpoch = 5
+    val partition = setupPartitionWithMocks(leaderEpoch, isLeader = false)
+
+    def assertReadRecordsError(error: Errors,
+                               currentLeaderEpochOpt: Optional[Integer],
+                               fetchOnlyLeader: Boolean): Unit = {
+      try {
+        partition.readRecords(
+          lastFetchedEpoch = Optional.empty(),
+          fetchOffset = 0L,
+          currentLeaderEpoch = currentLeaderEpochOpt,
+          maxBytes = 1024,
+          fetchIsolation = FetchLogEnd,
+          fetchOnlyFromLeader = fetchOnlyLeader,
+          minOneMessage = false)
+        if (error != Errors.NONE)
+          fail(s"Expected readRecords to fail with error $error")
+      } catch {
+        case e: Exception =>
+          assertEquals(error, Errors.forException(e))
+      }
+    }
+
+    assertReadRecordsError(Errors.NONE, Optional.empty(), fetchOnlyLeader = false)
+    assertReadRecordsError(Errors.NONE, Optional.of(leaderEpoch), fetchOnlyLeader = false)
+    assertReadRecordsError(Errors.FENCED_LEADER_EPOCH, Optional.of(leaderEpoch - 1), fetchOnlyLeader = false)
+    assertReadRecordsError(Errors.UNKNOWN_LEADER_EPOCH, Optional.of(leaderEpoch + 1), fetchOnlyLeader = false)
+
+    assertReadRecordsError(Errors.NOT_LEADER_OR_FOLLOWER, Optional.empty(), fetchOnlyLeader = true)
+    assertReadRecordsError(Errors.NOT_LEADER_OR_FOLLOWER, Optional.of(leaderEpoch), fetchOnlyLeader = true)
+    assertReadRecordsError(Errors.FENCED_LEADER_EPOCH, Optional.of(leaderEpoch - 1), fetchOnlyLeader = true)
+    assertReadRecordsError(Errors.UNKNOWN_LEADER_EPOCH, Optional.of(leaderEpoch + 1), fetchOnlyLeader = true)
+  }
+
+  @Test
+  def testFetchOffsetForTimestampEpochValidationForLeader(): Unit = {
+    val leaderEpoch = 5
+    val partition = setupPartitionWithMocks(leaderEpoch, isLeader = true)
+
+    def assertFetchOffsetError(error: Errors,
+                               currentLeaderEpochOpt: Optional[Integer]): Unit = {
+      try {
+        partition.fetchOffsetForTimestamp(0L,
+          isolationLevel = None,
+          currentLeaderEpoch = currentLeaderEpochOpt,
+          fetchOnlyFromLeader = true)
+        if (error != Errors.NONE)
+          fail(s"Expected readRecords to fail with error $error")
+      } catch {
+        case e: Exception =>
+          assertEquals(error, Errors.forException(e))
+      }
+    }
+
+    assertFetchOffsetError(Errors.NONE, Optional.empty())
+    assertFetchOffsetError(Errors.NONE, Optional.of(leaderEpoch))
+    assertFetchOffsetError(Errors.FENCED_LEADER_EPOCH, Optional.of(leaderEpoch - 1))
+    assertFetchOffsetError(Errors.UNKNOWN_LEADER_EPOCH, Optional.of(leaderEpoch + 1))
+  }
+
+  @Test
+  def testFetchOffsetForTimestampEpochValidationForFollower(): Unit = {
+    val leaderEpoch = 5
+    val partition = setupPartitionWithMocks(leaderEpoch, isLeader = false)
+
+    def assertFetchOffsetError(error: Errors,
+                               currentLeaderEpochOpt: Optional[Integer],
+                               fetchOnlyLeader: Boolean): Unit = {
+      try {
+        partition.fetchOffsetForTimestamp(0L,
+          isolationLevel = None,
+          currentLeaderEpoch = currentLeaderEpochOpt,
+          fetchOnlyFromLeader = fetchOnlyLeader)
+        if (error != Errors.NONE)
+          fail(s"Expected readRecords to fail with error $error")
+      } catch {
+        case e: Exception =>
+          assertEquals(error, Errors.forException(e))
+      }
+    }
+
+    assertFetchOffsetError(Errors.NONE, Optional.empty(), fetchOnlyLeader = false)
+    assertFetchOffsetError(Errors.NONE, Optional.of(leaderEpoch), fetchOnlyLeader = false)
+    assertFetchOffsetError(Errors.FENCED_LEADER_EPOCH, Optional.of(leaderEpoch - 1), fetchOnlyLeader = false)
+    assertFetchOffsetError(Errors.UNKNOWN_LEADER_EPOCH, Optional.of(leaderEpoch + 1), fetchOnlyLeader = false)
+
+    assertFetchOffsetError(Errors.NOT_LEADER_OR_FOLLOWER, Optional.empty(), fetchOnlyLeader = true)
+    assertFetchOffsetError(Errors.NOT_LEADER_OR_FOLLOWER, Optional.of(leaderEpoch), fetchOnlyLeader = true)
+    assertFetchOffsetError(Errors.FENCED_LEADER_EPOCH, Optional.of(leaderEpoch - 1), fetchOnlyLeader = true)
+    assertFetchOffsetError(Errors.UNKNOWN_LEADER_EPOCH, Optional.of(leaderEpoch + 1), fetchOnlyLeader = true)
+  }
+
+  @Test
+  def testFetchLatestOffsetIncludesLeaderEpoch(): Unit = {
+    val leaderEpoch = 5
+    val partition = setupPartitionWithMocks(leaderEpoch, isLeader = true)
+
+    val timestampAndOffsetOpt = partition.fetchOffsetForTimestamp(ListOffsetRequest.LATEST_TIMESTAMP,
+      isolationLevel = None,
+      currentLeaderEpoch = Optional.empty(),
+      fetchOnlyFromLeader = true)
+
+    assertTrue(timestampAndOffsetOpt.isDefined)
+
+    val timestampAndOffset = timestampAndOffsetOpt.get
+    assertEquals(leaderEpoch, timestampAndOffset.leaderEpoch.get)
+  }
+
+  /**
+    * This test checks that after a new leader election, we don't answer any ListOffsetsRequest until
+    * the HW of the new leader has caught up to its startLogOffset for this epoch. From a client
+    * perspective this helps guarantee monotonic offsets
+    *
+    * @see KIP-207
+    */
+  @Test
+  def testMonotonicOffsetsAfterLeaderChange(): Unit = {
+    val controllerEpoch = 3
+    val leader = brokerId
+    val follower1 = brokerId + 1
+    val follower2 = brokerId + 2
+    val replicas = List(leader, follower1, follower2)
+    val isr = List[Integer](leader, follower2).asJava
+    val leaderEpoch = 8
+    val batch1 = TestUtils.records(records = List(
+      new SimpleRecord(10, "k1".getBytes, "v1".getBytes),
+      new SimpleRecord(11,"k2".getBytes, "v2".getBytes)))
+    val batch2 = TestUtils.records(records = List(new SimpleRecord("k3".getBytes, "v1".getBytes),
+      new SimpleRecord(20,"k4".getBytes, "v2".getBytes),
+      new SimpleRecord(21,"k5".getBytes, "v3".getBytes)))
+
+    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(leaderState, offsetCheckpoints))
+    assertEquals("Current leader epoch", leaderEpoch, partition.getLeaderEpoch)
+    assertEquals("ISR", Set[Integer](leader, follower2), partition.isrState.isr)
+
+    // after makeLeader(() call, partition should know about all the replicas
+    // append records with initial leader epoch
+    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 updateFollowerFetchState(followerId: Int, fetchOffsetMetadata: LogOffsetMetadata): Unit = {
+      partition.updateFollowerFetchState(
+        followerId,
+        followerFetchOffsetMetadata = fetchOffsetMetadata,
+        followerStartOffset = 0L,
+        followerFetchTimeMs = time.milliseconds(),
+        leaderEndOffset = partition.localLogOrException.logEndOffset)
+    }
+
+    def fetchOffsetsForTimestamp(timestamp: Long, isolation: Option[IsolationLevel]): Either[ApiException, Option[TimestampAndOffset]] = {
+      try {
+        Right(partition.fetchOffsetForTimestamp(
+          timestamp = timestamp,
+          isolationLevel = isolation,
+          currentLeaderEpoch = Optional.of(partition.getLeaderEpoch),
+          fetchOnlyFromLeader = true
+        ))
+      } catch {
+        case e: ApiException => Left(e)
+      }
+    }
+
+    when(stateStore.expandIsr(controllerEpoch, new LeaderAndIsr(leader, leaderEpoch,
+      List(leader, follower2, follower1), 1)))
+      .thenReturn(Some(2))
+
+    updateFollowerFetchState(follower1, LogOffsetMetadata(0))
+    updateFollowerFetchState(follower1, LogOffsetMetadata(2))
+
+    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.localLogOrException.highWatermark)
+
+    // Get the LEO
+    fetchOffsetsForTimestamp(ListOffsetRequest.LATEST_TIMESTAMP, None) match {
+      case Right(Some(offsetAndTimestamp)) => assertEquals(5, offsetAndTimestamp.offset)
+      case Right(None) => fail("Should have seen some offsets")
+      case Left(e) => fail("Should not have seen an error")
+    }
+
+    // Get the HW
+    fetchOffsetsForTimestamp(ListOffsetRequest.LATEST_TIMESTAMP, Some(IsolationLevel.READ_UNCOMMITTED)) match {
+      case Right(Some(offsetAndTimestamp)) => assertEquals(2, offsetAndTimestamp.offset)
+      case Right(None) => fail("Should have seen some offsets")
+      case Left(e) => fail("Should not have seen an error")
+    }
+
+    // Get a offset beyond the HW by timestamp, get a None
+    assertEquals(Right(None), fetchOffsetsForTimestamp(30, Some(IsolationLevel.READ_UNCOMMITTED)))
+
+    // Make into a follower
+    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
+    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 {
+      case Right(Some(offsetAndTimestamp)) => fail("Should have failed with OffsetNotAvailable")
+      case Right(None) => fail("Should have seen an error")
+      case Left(e: OffsetNotAvailableException) => // ok
+      case Left(e: ApiException) => fail(s"Expected OffsetNotAvailableException, got $e")
+    }
+
+    // If request is not from a client, we skip the check
+    fetchOffsetsForTimestamp(ListOffsetRequest.LATEST_TIMESTAMP, None) match {
+      case Right(Some(offsetAndTimestamp)) => assertEquals(5, offsetAndTimestamp.offset)
+      case Right(None) => fail("Should have seen some offsets")
+      case Left(e: ApiException) => fail(s"Got ApiException $e")
+    }
+
+    // If we request the earliest timestamp, we skip the check
+    fetchOffsetsForTimestamp(ListOffsetRequest.EARLIEST_TIMESTAMP, Some(IsolationLevel.READ_UNCOMMITTED)) match {
+      case Right(Some(offsetAndTimestamp)) => assertEquals(0, offsetAndTimestamp.offset)
+      case Right(None) => fail("Should have seen some offsets")
+      case Left(e: ApiException) => fail(s"Got ApiException $e")
+    }
+
+    // If we request an offset by timestamp earlier than the HW, we are ok
+    fetchOffsetsForTimestamp(11, Some(IsolationLevel.READ_UNCOMMITTED)) match {
+      case Right(Some(offsetAndTimestamp)) =>
+        assertEquals(1, offsetAndTimestamp.offset)
+        assertEquals(11, offsetAndTimestamp.timestamp)
+      case Right(None) => fail("Should have seen some offsets")
+      case Left(e: ApiException) => fail(s"Got ApiException $e")
+    }
+
+    // Request an offset by timestamp beyond the HW, get an error now since we're in a bad state
+    fetchOffsetsForTimestamp(100, Some(IsolationLevel.READ_UNCOMMITTED)) match {
+      case Right(Some(offsetAndTimestamp)) => fail("Should have failed")
+      case Right(None) => fail("Should have failed")
+      case Left(e: OffsetNotAvailableException) => // ok
+      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)
+    updateFollowerFetchState(follower1, LogOffsetMetadata(5))
+    updateFollowerFetchState(follower2, LogOffsetMetadata(5))
+
+    // Error goes away
+    fetchOffsetsForTimestamp(ListOffsetRequest.LATEST_TIMESTAMP, Some(IsolationLevel.READ_UNCOMMITTED)) match {
+      case Right(Some(offsetAndTimestamp)) => assertEquals(5, offsetAndTimestamp.offset)
+      case Right(None) => fail("Should have seen some offsets")
+      case Left(e: ApiException) => fail(s"Got ApiException $e")
+    }
+
+    // Now we see None instead of an error for out of range timestamp
+    assertEquals(Right(None), fetchOffsetsForTimestamp(100, Some(IsolationLevel.READ_UNCOMMITTED)))
+  }
+
+  private def setupPartitionWithMocks(leaderEpoch: Int,
+                                      isLeader: Boolean,
+                                      log: Log = logManager.getOrCreateLog(topicPartition, () => logConfig)): Partition = {
+    partition.createLogIfNotExists(isNew = false, isFutureReplica = false, offsetCheckpoints)
+
+    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(new LeaderAndIsrPartitionState()
+          .setControllerEpoch(controllerEpoch)
+          .setLeader(brokerId)
+          .setLeaderEpoch(leaderEpoch)
+          .setIsr(isr)
+          .setZkVersion(1)
+          .setReplicas(replicas)
+          .setIsNew(true), offsetCheckpoints))
+      assertEquals(leaderEpoch, partition.getLeaderEpoch)
+    } else {
+      assertTrue("Expected become follower transition to succeed",
+        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.leaderLogIfLocal)
+    }
+
+    partition
+  }
+
+  @Test
+  def testAppendRecordsAsFollowerBelowLogStartOffset(): Unit = {
+    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, log.logEndOffset)
+    assertEquals(s"Log start offset after truncate fully and start at $initialLogStartOffset:",
+                 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, 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
+    val newLogStartOffset = 4L
+    val records = createRecords(List(new SimpleRecord("k1".getBytes, "v1".getBytes),
+                                     new SimpleRecord("k2".getBytes, "v2".getBytes),
+                                     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, 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, 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] {
+      val records2 = createRecords(List(new SimpleRecord("k1".getBytes, "v1".getBytes),
+                                        new SimpleRecord("k2".getBytes, "v2".getBytes)),
+                                   baseOffset = 3L)
+      partition.appendRecordsToFollowerOrFutureReplica(records2, isFuture = false)
+    }
+    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, log.logEndOffset)
+    assertEquals(s"Log start offset not expected to change:", newLogStartOffset, log.logStartOffset)
+  }
+
+  @Test
+  def testListOffsetIsolationLevels(): Unit = {
+    val controllerEpoch = 0
+    val leaderEpoch = 5
+    val replicas = List[Integer](brokerId, brokerId + 1).asJava
+    val isr = replicas
+
+    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(leaderEpoch, partition.getLeaderEpoch)
+
+    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, origin = AppendOrigin.Client, requiredAcks = 0)
+
+    def fetchLatestOffset(isolationLevel: Option[IsolationLevel]): TimestampAndOffset = {
+      val res = partition.fetchOffsetForTimestamp(ListOffsetRequest.LATEST_TIMESTAMP,
+        isolationLevel = isolationLevel,
+        currentLeaderEpoch = Optional.empty(),
+        fetchOnlyFromLeader = true)
+      assertTrue(res.isDefined)
+      res.get
+    }
+
+    def fetchEarliestOffset(isolationLevel: Option[IsolationLevel]): TimestampAndOffset = {
+      val res = partition.fetchOffsetForTimestamp(ListOffsetRequest.EARLIEST_TIMESTAMP,
+        isolationLevel = isolationLevel,
+        currentLeaderEpoch = Optional.empty(),
+        fetchOnlyFromLeader = true)
+      assertTrue(res.isDefined)
+      res.get
+    }
+
+    assertEquals(3L, fetchLatestOffset(isolationLevel = None).offset)
+    assertEquals(0L, fetchLatestOffset(isolationLevel = Some(IsolationLevel.READ_UNCOMMITTED)).offset)
+    assertEquals(0L, fetchLatestOffset(isolationLevel = Some(IsolationLevel.READ_COMMITTED)).offset)
+
+    partition.log.get.updateHighWatermark(1L)
+
+    assertEquals(3L, fetchLatestOffset(isolationLevel = None).offset)
+    assertEquals(1L, fetchLatestOffset(isolationLevel = Some(IsolationLevel.READ_UNCOMMITTED)).offset)
+    assertEquals(0L, fetchLatestOffset(isolationLevel = Some(IsolationLevel.READ_COMMITTED)).offset)
+
+    assertEquals(0L, fetchEarliestOffset(isolationLevel = None).offset)
+    assertEquals(0L, fetchEarliestOffset(isolationLevel = Some(IsolationLevel.READ_UNCOMMITTED)).offset)
+    assertEquals(0L, fetchEarliestOffset(isolationLevel = Some(IsolationLevel.READ_COMMITTED)).offset)
+  }
+
+  @Test
+  def testGetReplica(): Unit = {
+    assertEquals(None, partition.log)
+    assertThrows[NotLeaderOrFollowerException] {
+      partition.localLogOrException
+    }
+  }
+
+  @Test
+  def testAppendRecordsToFollowerWithNoReplicaThrowsException(): Unit = {
+    assertThrows[NotLeaderOrFollowerException] {
+      partition.appendRecordsToFollowerOrFutureReplica(
+           createRecords(List(new SimpleRecord("k1".getBytes, "v1".getBytes)), baseOffset = 0L), isFuture = false)
+    }
+  }
+
+  @Test
+  def testMakeFollowerWithNoLeaderIdChange(): Unit = {
+    // Start off as follower
+    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
+    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
+    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
+  def testFollowerDoesNotJoinISRUntilCaughtUpToOffsetWithinCurrentLeaderEpoch(): 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, follower2).asJava
+    val leaderEpoch = 8
+    val batch1 = TestUtils.records(records = List(new SimpleRecord("k1".getBytes, "v1".getBytes),
+                                                  new SimpleRecord("k2".getBytes, "v2".getBytes)))
+    val batch2 = TestUtils.records(records = List(new SimpleRecord("k3".getBytes, "v1".getBytes),
+                                                  new SimpleRecord("k4".getBytes, "v2".getBytes),
+                                                  new SimpleRecord("k5".getBytes, "v3".getBytes)))
+    val batch3 = TestUtils.records(records = List(new SimpleRecord("k6".getBytes, "v1".getBytes),
+                                                  new SimpleRecord("k7".getBytes, "v2".getBytes)))
+
+    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(leaderState, offsetCheckpoints))
+    assertEquals("Current leader epoch", leaderEpoch, partition.getLeaderEpoch)
+    assertEquals("ISR", Set[Integer](leader, follower2), partition.isrState.isr)
+
+    // after makeLeader(() call, partition should know about all the replicas
+    // append records with initial leader epoch
+    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 updateFollowerFetchState(followerId: Int, fetchOffsetMetadata: LogOffsetMetadata): Unit = {
+      partition.updateFollowerFetchState(
+        followerId,
+        followerFetchOffsetMetadata = fetchOffsetMetadata,
+        followerStartOffset = 0L,
+        followerFetchTimeMs = time.milliseconds(),
+        leaderEndOffset = partition.localLogOrException.logEndOffset)
+    }
+
+    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)
+    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(newLeaderState, offsetCheckpoints))
+    val currentLeaderEpochStartOffset = partition.localLogOrException.logEndOffset
+
+    // append records with the latest leader epoch
+    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
+    updateFollowerFetchState(follower1, LogOffsetMetadata(0))
+    updateFollowerFetchState(follower1, LogOffsetMetadata(lastOffsetOfFirstBatch))
+    assertEquals("ISR", Set[Integer](leader, follower2), partition.isrState.isr)
+
+    // fetch from the follower not in ISR from start offset of the current leader epoch should
+    // add this follower to ISR
+    updateFollowerFetchState(follower1, LogOffsetMetadata(currentLeaderEpochStartOffset))
+
+    // Expansion does not affect the ISR
+    assertEquals("ISR", Set[Integer](leader, follower2), partition.isrState.isr)
+    assertEquals("ISR", Set[Integer](leader, follower1, follower2), partition.isrState.maximalIsr)
+    assertEquals("AlterIsr", alterIsrManager.isrUpdates.dequeue().leaderAndIsr.isr.toSet,
+      Set(leader, follower1, follower2))
+  }
+
+  def createRecords(records: Iterable[SimpleRecord], baseOffset: Long, partitionLeaderEpoch: Int = 0): MemoryRecords = {
+    val buf = ByteBuffer.allocate(DefaultRecordBatch.sizeInBytes(records.asJava))
+    val builder = MemoryRecords.builder(
+      buf, RecordBatch.CURRENT_MAGIC_VALUE, CompressionType.NONE, TimestampType.LOG_APPEND_TIME,
+      baseOffset, time.milliseconds, partitionLeaderEpoch)
+    records.foreach(builder.append)
+    builder.build()
+  }
+
+  def createTransactionalRecords(records: Iterable[SimpleRecord],
+                                 baseOffset: Long): MemoryRecords = {
+    val producerId = 1L
+    val producerEpoch = 0.toShort
+    val baseSequence = 0
+    val isTransactional = true
+    val buf = ByteBuffer.allocate(DefaultRecordBatch.sizeInBytes(records.asJava))
+    val builder = MemoryRecords.builder(buf, CompressionType.NONE, baseOffset, producerId,
+      producerEpoch, baseSequence, isTransactional)
+    records.foreach(builder.append)
+    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
+
+    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)
+
+    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)
+
+    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
+
+    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.isrState.isr)
+
+    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)
+
+    assertEquals(Set(brokerId), partition.isrState.isr)
+    assertEquals(3L, remoteReplica.logEndOffset)
+    assertEquals(0L, remoteReplica.logStartOffset)
+
+    partition.updateFollowerFetchState(remoteBrokerId,
+      followerFetchOffsetMetadata = LogOffsetMetadata(10),
+      followerStartOffset = 0L,
+      followerFetchTimeMs = time.milliseconds(),
+      leaderEndOffset = 6L)
+
+    assertEquals(alterIsrManager.isrUpdates.size, 1)
+    val isrItem = alterIsrManager.isrUpdates.dequeue()
+    assertEquals(isrItem.leaderAndIsr.isr, List(brokerId, remoteBrokerId))
+    assertEquals(Set(brokerId), partition.isrState.isr)
+    assertEquals(Set(brokerId, remoteBrokerId), partition.isrState.maximalIsr)
+    assertEquals(10L, remoteReplica.logEndOffset)
+    assertEquals(0L, remoteReplica.logStartOffset)
+
+    // Complete the ISR expansion
+    isrItem.callback.apply(Right(new LeaderAndIsr(brokerId, leaderEpoch, List(brokerId, remoteBrokerId), 2)))
+    assertEquals(Set(brokerId, remoteBrokerId), partition.isrState.isr)
+
+    assertEquals(isrChangeListener.expands.get, 1)
+    assertEquals(isrChangeListener.shrinks.get, 0)
+    assertEquals(isrChangeListener.failures.get, 0)
+  }
+
+  @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
+
+    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.isrState.isr)
+
+    val remoteReplica = partition.getReplica(remoteBrokerId).get
+    assertEquals(LogOffsetMetadata.UnknownOffsetMetadata.messageOffset, remoteReplica.logEndOffset)
+    assertEquals(Log.UnknownOffset, remoteReplica.logStartOffset)
+
+    partition.updateFollowerFetchState(remoteBrokerId,
+      followerFetchOffsetMetadata = LogOffsetMetadata(10),
+      followerStartOffset = 0L,
+      followerFetchTimeMs = time.milliseconds(),
+      leaderEndOffset = 10L)
+
+    // Follower state is updated, but the ISR has not expanded
+    assertEquals(Set(brokerId), partition.inSyncReplicaIds)
+    assertEquals(Set(brokerId, remoteBrokerId), partition.isrState.maximalIsr)
+    assertEquals(alterIsrManager.isrUpdates.size, 1)
+    assertEquals(10L, remoteReplica.logEndOffset)
+    assertEquals(0L, remoteReplica.logStartOffset)
+
+    // Simulate failure callback
+    val alterIsrItem = alterIsrManager.isrUpdates.dequeue()
+    alterIsrItem.callback.apply(Left(Errors.INVALID_UPDATE_VERSION))
+
+    // Still no ISR change
+    assertEquals(Set(brokerId), partition.inSyncReplicaIds)
+    assertEquals(Set(brokerId, remoteBrokerId), partition.isrState.maximalIsr)
+    assertEquals(alterIsrManager.isrUpdates.size, 0)
+
+    assertEquals(isrChangeListener.expands.get, 0)
+    assertEquals(isrChangeListener.shrinks.get, 0)
+    assertEquals(isrChangeListener.failures.get, 1)
+  }
+
+  @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
+
+    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.isrState.isr)
+    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()
+    assertEquals(Set(brokerId, remoteBrokerId), partition.isrState.isr)
+
+    // If enough time passes without a fetch update, the ISR should shrink
+    time.sleep(partition.replicaLagTimeMaxMs + 1)
+
+    // Shrink the ISR
+    partition.maybeShrinkIsr()
+    assertEquals(alterIsrManager.isrUpdates.size, 1)
+    assertEquals(alterIsrManager.isrUpdates.dequeue().leaderAndIsr.isr, List(brokerId))
+    assertEquals(Set(brokerId, remoteBrokerId), partition.isrState.isr)
+    assertEquals(Set(brokerId, remoteBrokerId), partition.isrState.maximalIsr)
+    assertEquals(0L, 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
+
+    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.isrState.isr)
+    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)
+    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)
+    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()
+    assertEquals(Set(brokerId, remoteBrokerId), partition.isrState.isr)
+    assertEquals(alterIsrManager.isrUpdates.size, 0)
+  }
+
+  @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
+
+    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.isrState.isr)
+    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)
+    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(30001)
+
+    // The ISR should not be shrunk because the follower is caught up to the leader's log end
+    partition.maybeShrinkIsr()
+    assertEquals(Set(brokerId, remoteBrokerId), partition.isrState.isr)
+    assertEquals(alterIsrManager.isrUpdates.size, 0)
+  }
+
+  @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
+
+    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(30001)
+
+    // Enqueue and AlterIsr that will fail
+    partition.maybeShrinkIsr()
+    assertEquals(Set(brokerId, remoteBrokerId), partition.inSyncReplicaIds)
+    assertEquals(alterIsrManager.isrUpdates.size, 1)
+    assertEquals(0L, partition.localLogOrException.highWatermark)
+
+    // Simulate failure callback
+    val alterIsrItem = alterIsrManager.isrUpdates.dequeue()
+    alterIsrItem.callback.apply(Left(Errors.INVALID_UPDATE_VERSION))
+
+    // Ensure ISR hasn't changed
+    assertEquals(partition.isrState.getClass, classOf[PendingShrinkIsr])
+    assertEquals(Set(brokerId, remoteBrokerId), partition.inSyncReplicaIds)
+    assertEquals(alterIsrManager.isrUpdates.size, 0)
+    assertEquals(0L, partition.localLogOrException.highWatermark)
+  }
+
+  @Test
+  def testAlterIsrUnknownTopic(): Unit = {
+    handleAlterIsrFailure(Errors.UNKNOWN_TOPIC_OR_PARTITION,
+      (brokerId: Int, remoteBrokerId: Int, partition: Partition) => {
+        assertEquals(partition.isrState.isr, Set(brokerId))
+        assertEquals(partition.isrState.maximalIsr, Set(brokerId, remoteBrokerId))
+        assertEquals(alterIsrManager.isrUpdates.size, 0)
+      })
+  }
+
+  @Test
+  def testAlterIsrInvalidVersion(): Unit = {
+    handleAlterIsrFailure(Errors.INVALID_UPDATE_VERSION,
+      (brokerId: Int, remoteBrokerId: Int, partition: Partition) => {
+        assertEquals(partition.isrState.isr, Set(brokerId))
+        assertEquals(partition.isrState.maximalIsr, Set(brokerId, remoteBrokerId))
+        assertEquals(alterIsrManager.isrUpdates.size, 0)
+      })
+  }
+
+  @Test
+  def testAlterIsrUnexpectedError(): Unit = {
+    handleAlterIsrFailure(Errors.UNKNOWN_SERVER_ERROR,
+      (brokerId: Int, remoteBrokerId: Int, partition: Partition) => {
+        // We retry these
+        assertEquals(partition.isrState.isr, Set(brokerId))
+        assertEquals(partition.isrState.maximalIsr, Set(brokerId, remoteBrokerId))
+        assertEquals(alterIsrManager.isrUpdates.size, 1)
+      })
+  }
+
+  def handleAlterIsrFailure(error: Errors, callback: (Int, Int, Partition) => Unit): 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
+
+    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.isrState.isr)
+
+    val remoteReplica = partition.getReplica(remoteBrokerId).get
+    assertEquals(LogOffsetMetadata.UnknownOffsetMetadata.messageOffset, remoteReplica.logEndOffset)
+    assertEquals(Log.UnknownOffset, remoteReplica.logStartOffset)
+
+    // This will attempt to expand the ISR
+    partition.updateFollowerFetchState(remoteBrokerId,
+      followerFetchOffsetMetadata = LogOffsetMetadata(10),
+      followerStartOffset = 0L,
+      followerFetchTimeMs = time.milliseconds(),
+      leaderEndOffset = 10L)
+
+    // Follower state is updated, but the ISR has not expanded
+    assertEquals(Set(brokerId), partition.inSyncReplicaIds)
+    assertEquals(Set(brokerId, remoteBrokerId), partition.isrState.maximalIsr)
+    assertEquals(alterIsrManager.isrUpdates.size, 1)
+    assertEquals(10L, remoteReplica.logEndOffset)
+    assertEquals(0L, remoteReplica.logStartOffset)
+
+    // Failure
+    alterIsrManager.isrUpdates.dequeue().callback(Left(error))
+    callback(brokerId, remoteBrokerId, partition)
+  }
+
+  @Test
+  def testSingleInFlightAlterIsr(): Unit = {
+    val log = logManager.getOrCreateLog(topicPartition, () => logConfig)
+    seedLogData(log, numRecords = 10, leaderEpoch = 4)
+
+    val controllerEpoch = 0
+    val leaderEpoch = 5
+    val follower1 = brokerId + 1
+    val follower2 = brokerId + 2
+    val follower3 = brokerId + 3
+    val replicas = List[Integer](brokerId, follower1, follower2, follower3).asJava
+    val isr = List[Integer](brokerId, follower1, follower2).asJava
+
+    doNothing().when(delayedOperations).checkAndCompleteAll()
+
+    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, follower1, follower2), partition.isrState.isr)
+    assertEquals(0L, partition.localLogOrException.highWatermark)
+
+    // Expand ISR
+    partition.expandIsr(follower3)
+    assertEquals(Set(brokerId, follower1, follower2), partition.isrState.isr)
+    assertEquals(Set(brokerId, follower1, follower2, follower3), partition.isrState.maximalIsr)
+
+    // One AlterIsr request in-flight
+    assertEquals(alterIsrManager.isrUpdates.size, 1)
+
+    // Try to modify ISR again, should do nothing
+    partition.shrinkIsr(Set(follower3))
+    assertEquals(alterIsrManager.isrUpdates.size, 1)
+  }
+
+  @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] = {
+      KafkaYammerMetrics.defaultRegistry().allMetrics().asScala.filter { case (metricName, _) =>
+        metricName.getName == metric && metricName.getType == "Partition"
+      }.headOption.map(_._2)
+    }
+
+    assertTrue(metricsToCheck.forall(getMetric(_).isDefined))
+
+    Partition.removeMetrics(topicPartition)
+
+    assertEquals(Set(), KafkaYammerMetrics.defaultRegistry().allMetrics().asScala.keySet.filter(_.getType == "Partition"))
+  }
+
+  @Test
+  def testUnderReplicatedPartitionsCorrectSemantics(): Unit = {
+    val controllerEpoch = 3
+    val replicas = List[Integer](brokerId, brokerId + 1, brokerId + 2).asJava
+    val isr = List[Integer](brokerId, brokerId + 1).asJava
+
+    var leaderState = new LeaderAndIsrPartitionState()
+      .setControllerEpoch(controllerEpoch)
+      .setLeader(brokerId)
+      .setLeaderEpoch(6)
+      .setIsr(isr)
+      .setZkVersion(1)
+      .setReplicas(replicas)
+      .setIsNew(false)
+    partition.makeLeader(leaderState, offsetCheckpoints)
+    assertTrue(partition.isUnderReplicated)
+
+    leaderState = leaderState.setIsr(replicas)
+    partition.makeLeader(leaderState, offsetCheckpoints)
+    assertFalse(partition.isUnderReplicated)
+  }
+
+  @Test
+  def testUpdateAssignmentAndIsr(): Unit = {
+    val topicPartition = new TopicPartition("test", 1)
+    val partition = new Partition(
+      topicPartition, 1000, ApiVersion.latestVersion, 0,
+      new SystemTime(), mock(classOf[PartitionStateStore]), mock(classOf[IsrChangeListener]), mock(classOf[DelayedOperations]),
+      mock(classOf[MetadataCache]), mock(classOf[LogManager]), mock(classOf[AlterIsrManager]))
+
+    val replicas = Seq(0, 1, 2, 3)
+    val isr = Set(0, 1, 2, 3)
+    val adding = Seq(4, 5)
+    val removing = Seq(1, 2)
+
+    // Test with ongoing reassignment
+    partition.updateAssignmentAndIsr(replicas, isr, adding, removing)
+
+    assertTrue("The assignmentState is not OngoingReassignmentState",
+      partition.assignmentState.isInstanceOf[OngoingReassignmentState])
+    assertEquals(replicas, partition.assignmentState.replicas)
+    assertEquals(isr, partition.isrState.isr)
+    assertEquals(adding, partition.assignmentState.asInstanceOf[OngoingReassignmentState].addingReplicas)
+    assertEquals(removing, partition.assignmentState.asInstanceOf[OngoingReassignmentState].removingReplicas)
+    assertEquals(Seq(1, 2, 3), partition.remoteReplicas.map(_.brokerId))
+
+    // Test with simple assignment
+    val replicas2 = Seq(0, 3, 4, 5)
+    val isr2 = Set(0, 3, 4, 5)
+    partition.updateAssignmentAndIsr(replicas2, isr2, Seq.empty, Seq.empty)
+
+    assertTrue("The assignmentState is not SimpleAssignmentState",
+      partition.assignmentState.isInstanceOf[SimpleAssignmentState])
+    assertEquals(replicas2, partition.assignmentState.replicas)
+    assertEquals(isr2, partition.isrState.isr)
+    assertEquals(Seq(3, 4, 5), partition.remoteReplicas.map(_.brokerId))
+  }
+
+  /**
+   * 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,
+      isrChangeListener,
+      delayedOperations,
+      metadataCache,
+      spyLogManager,
+      alterIsrManager)
+
+    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((invocation: InvocationOnMock) => {
+      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,
+      isrChangeListener,
+      delayedOperations,
+      metadataCache,
+      spyLogManager,
+      alterIsrManager)
+
+    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((invocation: InvocationOnMock) => {
+      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,
+      isrChangeListener,
+      delayedOperations,
+      metadataCache,
+      spyLogManager,
+      alterIsrManager)
+
+    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)
+    }
+  }
+
+  private class SlowLog(log: Log, mockTime: MockTime, appendSemaphore: Semaphore) extends Log(
+    log.dir,
+    log.config,
+    log.logStartOffset,
+    log.recoveryPoint,
+    mockTime.scheduler,
+    new BrokerTopicStats,
+    log.time,
+    log.maxProducerIdExpirationMs,
+    log.producerIdExpirationCheckIntervalMs,
+    log.topicPartition,
+    log.producerStateManager,
+    new LogDirFailureChannel(1)) {
+
+    override def appendAsFollower(records: MemoryRecords): LogAppendInfo = {
+      appendSemaphore.acquire()
+      val appendInfo = super.appendAsFollower(records)
+      appendInfo
+    }
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/cluster/ReplicaTest.scala b/core/src/test/scala/unit/kafka/cluster/ReplicaTest.scala
index b3d44682a7542..6241c0cd807f0 100644
--- a/core/src/test/scala/unit/kafka/cluster/ReplicaTest.scala
+++ b/core/src/test/scala/unit/kafka/cluster/ReplicaTest.scala
@@ -18,14 +18,13 @@ package kafka.cluster
 
 import java.util.Properties
 
-import kafka.log.{Log, LogConfig, LogManager}
-import kafka.server.{BrokerTopicStats, LogDirFailureChannel, LogOffsetMetadata}
+import kafka.log.{ClientRecordDeletion, Log, LogConfig, LogManager}
+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}
 import org.junit.Assert._
+import org.junit.{After, Before, Test}
 
 class ReplicaTest {
 
@@ -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.messageOffset)
+    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, ClientRecordDeletion)
   }
-
 }
diff --git a/core/src/test/scala/unit/kafka/common/ConfigTest.scala b/core/src/test/scala/unit/kafka/common/ConfigTest.scala
deleted file mode 100644
index 2d20b1e9c6486..0000000000000
--- a/core/src/test/scala/unit/kafka/common/ConfigTest.scala
+++ /dev/null
@@ -1,89 +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.common
-
-import org.junit.Assert._
-import collection.mutable.ArrayBuffer
-import org.junit.Test
-import kafka.producer.ProducerConfig
-import kafka.consumer.ConsumerConfig
-
-class ConfigTest {
-
-  @Test
-  @deprecated("This test is deprecated and it will be removed in a future release.", "0.10.0.0")
-  def testInvalidClientIds() {
-    val invalidClientIds = new ArrayBuffer[String]()
-    val badChars = Array('/', '\\', ',', '\u0000', ':', "\"", '\'', ';', '*', '?', ' ', '\t', '\r', '\n', '=')
-    for (weirdChar <- badChars) {
-      invalidClientIds += "Is" + weirdChar + "illegal"
-    }
-
-    for (i <- 0 until invalidClientIds.size) {
-      try {
-        ProducerConfig.validateClientId(invalidClientIds(i))
-        fail("Should throw InvalidClientIdException.")
-      }
-      catch {
-        case _: InvalidConfigException => // This is good
-      }
-    }
-
-    val validClientIds = new ArrayBuffer[String]()
-    validClientIds += ("valid", "CLIENT", "iDs", "ar6", "VaL1d", "_0-9_.", "")
-    for (i <- 0 until validClientIds.size) {
-      try {
-        ProducerConfig.validateClientId(validClientIds(i))
-      }
-      catch {
-        case _: Exception => fail("Should not throw exception.")
-      }
-    }
-  }
-
-  @Test
-  def testInvalidGroupIds() {
-    val invalidGroupIds = new ArrayBuffer[String]()
-    val badChars = Array('/', '\\', ',', '\u0000', ':', "\"", '\'', ';', '*', '?', ' ', '\t', '\r', '\n', '=')
-    for (weirdChar <- badChars) {
-      invalidGroupIds += "Is" + weirdChar + "illegal"
-    }
-
-    for (i <- 0 until invalidGroupIds.size) {
-      try {
-        ConsumerConfig.validateGroupId(invalidGroupIds(i))
-        fail("Should throw InvalidGroupIdException.")
-      }
-      catch {
-        case _: InvalidConfigException => // This is good
-      }
-    }
-
-    val validGroupIds = new ArrayBuffer[String]()
-    validGroupIds += ("valid", "GROUP", "iDs", "ar6", "VaL1d", "_0-9_.", "")
-    for (i <- 0 until validGroupIds.size) {
-      try {
-        ConsumerConfig.validateGroupId(validGroupIds(i))
-      }
-      catch {
-        case _: Exception => fail("Should not throw exception.")
-      }
-    }
-  }
-}
-
diff --git a/core/src/test/scala/unit/kafka/common/ZkNodeChangeNotificationListenerTest.scala b/core/src/test/scala/unit/kafka/common/ZkNodeChangeNotificationListenerTest.scala
index 99550d5a60821..8d15f0b89b7a5 100644
--- a/core/src/test/scala/unit/kafka/common/ZkNodeChangeNotificationListenerTest.scala
+++ b/core/src/test/scala/unit/kafka/common/ZkNodeChangeNotificationListenerTest.scala
@@ -17,50 +17,95 @@
 package kafka.common
 
 import kafka.utils.TestUtils
-import kafka.zk.{AclChangeNotificationSequenceZNode, AclChangeNotificationZNode, ZooKeeperTestHarness}
-import org.junit.Test
+import kafka.zk.{LiteralAclChangeStore, LiteralAclStore, ZkAclChangeStore, ZooKeeperTestHarness}
+import org.apache.kafka.common.resource.PatternType.LITERAL
+import org.apache.kafka.common.resource.ResourcePattern
+import org.apache.kafka.common.resource.ResourceType.GROUP
+import org.junit.{After, Before, Test}
+
+import scala.collection.mutable.ArrayBuffer
+import scala.collection.Seq
 
 class ZkNodeChangeNotificationListenerTest extends ZooKeeperTestHarness {
 
-  @Test
-  def testProcessNotification() {
-    @volatile var notification: String = null
-    @volatile var invocationCount = 0
-    val notificationHandler = new NotificationHandler {
-      override def processNotification(notificationMessage: String): Unit = {
-        notification = notificationMessage
-        invocationCount += 1
-      }
-    }
+  private val changeExpirationMs = 1000
+  private var notificationListener: ZkNodeChangeNotificationListener = _
+  private var notificationHandler: TestNotificationHandler = _
 
+  @Before
+  override def setUp(): Unit = {
+    super.setUp()
     zkClient.createAclPaths()
-    val notificationMessage1 = "message1"
-    val notificationMessage2 = "message2"
-    val changeExpirationMs = 1000
+    notificationHandler = new TestNotificationHandler()
+  }
 
-    val notificationListener = new ZkNodeChangeNotificationListener(zkClient,  AclChangeNotificationZNode.path,
-      AclChangeNotificationSequenceZNode.SequenceNumberPrefix, notificationHandler, changeExpirationMs)
+  @After
+  override def tearDown(): Unit = {
+    if (notificationListener != null) {
+      notificationListener.close()
+    }
+    super.tearDown()
+  }
+
+  @Test
+  def testProcessNotification(): Unit = {
+    val notificationMessage1 = new ResourcePattern(GROUP, "messageA", LITERAL)
+    val notificationMessage2 = new ResourcePattern(GROUP, "messageB", LITERAL)
+
+    notificationListener = new ZkNodeChangeNotificationListener(zkClient, LiteralAclChangeStore.aclChangePath,
+      ZkAclChangeStore.SequenceNumberPrefix, notificationHandler, changeExpirationMs)
     notificationListener.init()
 
     zkClient.createAclChangeNotification(notificationMessage1)
-    TestUtils.waitUntilTrue(() => invocationCount == 1 && notification == notificationMessage1,
+    TestUtils.waitUntilTrue(() => notificationHandler.received().size == 1 && notificationHandler.received().last == notificationMessage1,
       "Failed to send/process notification message in the timeout period.")
 
     /*
      * 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.
      */
 
     zkClient.createAclChangeNotification(notificationMessage2)
-    TestUtils.waitUntilTrue(() => invocationCount == 2 && notification == notificationMessage2,
+    TestUtils.waitUntilTrue(() => notificationHandler.received().size == 2 && notificationHandler.received().last == notificationMessage2,
       "Failed to send/process notification message in the timeout period.")
 
-    (3 to 10).foreach(i => zkClient.createAclChangeNotification("message" + i))
+    (3 to 10).foreach(i => zkClient.createAclChangeNotification(new ResourcePattern(GROUP, "message" + i, LITERAL)))
+
+    TestUtils.waitUntilTrue(() => notificationHandler.received().size == 10,
+      s"Expected 10 invocations of processNotifications, but there were ${notificationHandler.received()}")
+  }
+
+  @Test
+  def testSwallowsProcessorException(): Unit = {
+    notificationHandler.setThrowSize(2)
+    notificationListener = new ZkNodeChangeNotificationListener(zkClient, LiteralAclChangeStore.aclChangePath,
+      ZkAclChangeStore.SequenceNumberPrefix, notificationHandler, changeExpirationMs)
+    notificationListener.init()
+
+    zkClient.createAclChangeNotification(new ResourcePattern(GROUP, "messageA", LITERAL))
+    zkClient.createAclChangeNotification(new ResourcePattern(GROUP, "messageB", LITERAL))
+    zkClient.createAclChangeNotification(new ResourcePattern(GROUP, "messageC", LITERAL))
+
+    TestUtils.waitUntilTrue(() => notificationHandler.received().size == 3,
+      s"Expected 2 invocations of processNotifications, but there were ${notificationHandler.received()}")
+  }
+
+  private class TestNotificationHandler extends NotificationHandler {
+    private val messages = ArrayBuffer.empty[ResourcePattern]
+    @volatile private var throwSize = Option.empty[Int]
+
+    override def processNotification(notificationMessage: Array[Byte]): Unit = {
+      messages += LiteralAclStore.changeStore.decode(notificationMessage)
+
+      if (throwSize.contains(messages.size))
+        throw new RuntimeException("Oh no, my processing failed!")
+    }
+
+    def received(): Seq[ResourcePattern] = messages
 
-    TestUtils.waitUntilTrue(() => invocationCount == 10 ,
-      s"Expected 10 invocations of processNotifications, but there were $invocationCount")
+    def setThrowSize(index: Int): Unit = throwSize = Option(index)
   }
 }
diff --git a/core/src/test/scala/unit/kafka/consumer/ConsumerIteratorTest.scala b/core/src/test/scala/unit/kafka/consumer/ConsumerIteratorTest.scala
deleted file mode 100755
index 5571e03fa56f3..0000000000000
--- a/core/src/test/scala/unit/kafka/consumer/ConsumerIteratorTest.scala
+++ /dev/null
@@ -1,122 +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.consumer
-
-import java.util.concurrent._
-import java.util.concurrent.atomic._
-
-import kafka.common.LongRef
-
-import scala.collection._
-import org.junit.Assert._
-import kafka.message._
-import kafka.server._
-import kafka.utils.TestUtils._
-import kafka.utils._
-import org.junit.{Before, Test}
-import kafka.serializer._
-import kafka.integration.KafkaServerTestHarness
-
-@deprecated("This test has been deprecated and will be removed in a future release.", "0.11.0.0")
-class ConsumerIteratorTest extends KafkaServerTestHarness {
-
-  val numNodes = 1
-
-  def generateConfigs = TestUtils.createBrokerConfigs(numNodes, zkConnect).map(KafkaConfig.fromProps)
-
-  val messages = new mutable.HashMap[Int, Seq[Message]]
-  val topic = "topic"
-  val group = "group1"
-  val consumer0 = "consumer0"
-  val consumedOffset = 5
-  val queue = new LinkedBlockingQueue[FetchedDataChunk]
-  var topicInfos: Seq[PartitionTopicInfo] = null
-
-  def consumerConfig = new ConsumerConfig(TestUtils.createConsumerProperties(zkConnect, group, consumer0))
-
-  @Before
-  override def setUp() {
-    super.setUp()
-    topicInfos = configs.map(_ => new PartitionTopicInfo(topic,
-      0,
-      queue,
-      new AtomicLong(consumedOffset),
-      new AtomicLong(0),
-      new AtomicInteger(0),
-      ""))
-    createTopic(zkUtils, topic, partitionReplicaAssignment = Map(0 -> Seq(configs.head.brokerId)), servers = servers)
-  }
-
-  @Test
-  def testConsumerIteratorDeduplicationDeepIterator() {
-    val messageStrings = (0 until 10).map(_.toString).toList
-    val messages = messageStrings.map(s => new Message(s.getBytes))
-    val messageSet = new ByteBufferMessageSet(DefaultCompressionCodec, new LongRef(0), messages:_*)
-
-    topicInfos.head.enqueue(messageSet)
-    assertEquals(1, queue.size)
-    queue.put(ZookeeperConsumerConnector.shutdownCommand)
-
-    val iter = new ConsumerIterator[String, String](queue,
-                                                    consumerConfig.consumerTimeoutMs,
-                                                    new StringDecoder(),
-                                                    new StringDecoder(),
-                                                    clientId = "")
-    val receivedMessages = (0 until 5).map(_ => iter.next.message)
-
-    assertFalse(iter.hasNext)
-    assertEquals(0, queue.size) // Shutdown command has been consumed.
-    assertEquals(5, receivedMessages.size)
-    val unconsumed = messageSet.filter(_.offset >= consumedOffset).map(m => TestUtils.readString(m.message.payload))
-    assertEquals(unconsumed, receivedMessages)
-  }
-
-  @Test
-  def testConsumerIteratorDecodingFailure() {
-    val messageStrings = (0 until 10).map(_.toString).toList
-    val messages = messageStrings.map(s => new Message(s.getBytes))
-    val messageSet = new ByteBufferMessageSet(NoCompressionCodec, new LongRef(0), messages:_*)
-
-    topicInfos.head.enqueue(messageSet)
-    assertEquals(1, queue.size)
-
-    val iter = new ConsumerIterator[String, String](queue,
-      ConsumerConfig.ConsumerTimeoutMs,
-      new FailDecoder(),
-      new FailDecoder(),
-      clientId = "")
-
-    (0 until 5).foreach { i =>
-      assertTrue(iter.hasNext)
-      val message = iter.next
-      assertEquals(message.offset, i + consumedOffset)
-
-      try message.message // should fail
-      catch {
-        case _: UnsupportedOperationException => // this is ok
-      }
-    }
-  }
-
-  class FailDecoder(props: VerifiableProperties = null) extends Decoder[String] {
-    def fromBytes(bytes: Array[Byte]): String = {
-      throw new UnsupportedOperationException("This decoder does not work at all..")
-    }
-  }
-}
diff --git a/core/src/test/scala/unit/kafka/consumer/PartitionAssignorTest.scala b/core/src/test/scala/unit/kafka/consumer/PartitionAssignorTest.scala
deleted file mode 100644
index 3012112b58dd6..0000000000000
--- a/core/src/test/scala/unit/kafka/consumer/PartitionAssignorTest.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.consumer
-
-import org.easymock.EasyMock
-import org.I0Itec.zkclient.ZkClient
-import org.apache.zookeeper.data.Stat
-import kafka.utils.{TestUtils, Logging, ZkUtils, Json}
-import org.junit.Assert._
-import kafka.common.TopicAndPartition
-import kafka.consumer.PartitionAssignorTest.StaticSubscriptionInfo
-import kafka.consumer.PartitionAssignorTest.Scenario
-import kafka.consumer.PartitionAssignorTest.WildcardSubscriptionInfo
-import org.junit.Test
-
-@deprecated("This test has been deprecated and will be removed in a future release.", "0.11.0.0")
-class PartitionAssignorTest extends Logging {
-
-  @Test
-  def testRoundRobinPartitionAssignor() {
-    val assignor = new RoundRobinAssignor
-
-    /** various scenarios with only wildcard consumers */
-    (1 to PartitionAssignorTest.TestCaseCount).foreach { _ =>
-      val consumerCount = 1.max(TestUtils.random.nextInt(PartitionAssignorTest.MaxConsumerCount + 1))
-      val topicCount = PartitionAssignorTest.MinTopicCount.max(TestUtils.random.nextInt(PartitionAssignorTest.MaxTopicCount + 1))
-
-      val topicPartitionCounts = Map((1 to topicCount).map(topic => {
-        ("topic-" + topic, PartitionAssignorTest.MinPartitionCount.max(TestUtils.random.nextInt(PartitionAssignorTest.MaxPartitionCount)))
-      }):_*)
-
-      val subscriptions = Map((1 to consumerCount).map { consumer =>
-        val streamCount = 1.max(TestUtils.random.nextInt(PartitionAssignorTest.MaxStreamCount + 1))
-        ("g1c" + consumer, WildcardSubscriptionInfo(streamCount, ".*", isWhitelist = true))
-      }:_*)
-      val scenario = Scenario("g1", topicPartitionCounts, subscriptions)
-      val zkUtils = PartitionAssignorTest.setupZkClientMock(scenario)
-      EasyMock.replay(zkUtils.zkClient)
-      PartitionAssignorTest.assignAndVerify(scenario, assignor, zkUtils, verifyAssignmentIsUniform = true)
-    }
-  }
-
-  @Test
-  def testRoundRobinPartitionAssignorStaticSubscriptions() {
-    val assignor = new RoundRobinAssignor
-
-    /** test static subscription scenarios */
-    (1 to PartitionAssignorTest.TestCaseCount).foreach (testCase => {
-      val consumerCount = 1.max(TestUtils.random.nextInt(PartitionAssignorTest.MaxConsumerCount + 1))
-      val topicCount = PartitionAssignorTest.MinTopicCount.max(TestUtils.random.nextInt(PartitionAssignorTest.MaxTopicCount + 1))
-
-      val topicPartitionCounts = Map((1 to topicCount).map(topic => {
-        ("topic-" + topic, PartitionAssignorTest.MinPartitionCount.max(TestUtils.random.nextInt(PartitionAssignorTest.MaxPartitionCount)))
-      }).toSeq:_*)
-
-      val subscriptions = Map((1 to consumerCount).map(consumer => {
-        val streamCounts = Map((1 to topicCount).map(topic => {
-            ("topic-" + topic, 1)
-          }).toSeq:_*)
-        ("g1c" + consumer, StaticSubscriptionInfo(streamCounts))
-      }).toSeq:_*)
-      val scenario = Scenario("g1", topicPartitionCounts, subscriptions)
-      val zkUtils = PartitionAssignorTest.setupZkClientMock(scenario)
-      EasyMock.replay(zkUtils.zkClient)
-      PartitionAssignorTest.assignAndVerify(scenario, assignor, zkUtils, verifyAssignmentIsUniform = true)
-    })
-  }
-
-  @Test
-  def testRoundRobinPartitionAssignorUnbalancedStaticSubscriptions() {
-    val assignor = new RoundRobinAssignor
-    val minConsumerCount = 5
-
-    /** test unbalanced static subscription scenarios */
-    (1 to PartitionAssignorTest.TestCaseCount).foreach (testCase => {
-      val consumerCount = minConsumerCount.max(TestUtils.random.nextInt(PartitionAssignorTest.MaxConsumerCount + 1))
-      val topicCount = 10
-
-      val topicPartitionCounts = Map((1 to topicCount).map(topic => {
-        ("topic-" + topic, 10)
-      }).toSeq:_*)
-
-      val subscriptions = Map((1 to consumerCount).map(consumer => {
-        // Exclude some topics from some consumers
-        val topicRange = (1 to topicCount - consumer % minConsumerCount)
-        val streamCounts = Map(topicRange.map(topic => {
-            ("topic-" + topic, 3)
-          }).toSeq:_*)
-        ("g1c" + consumer, StaticSubscriptionInfo(streamCounts))
-      }).toSeq:_*)
-      val scenario = Scenario("g1", topicPartitionCounts, subscriptions)
-      val zkUtils = PartitionAssignorTest.setupZkClientMock(scenario)
-      EasyMock.replay(zkUtils.zkClient)
-      PartitionAssignorTest.assignAndVerify(scenario, assignor, zkUtils)
-    })
-  }
-
-  @Test
-  def testRangePartitionAssignor() {
-    val assignor = new RangeAssignor
-    (1 to PartitionAssignorTest.TestCaseCount).foreach { _ =>
-      val consumerCount = 1.max(TestUtils.random.nextInt(PartitionAssignorTest.MaxConsumerCount + 1))
-      val topicCount = PartitionAssignorTest.MinTopicCount.max(TestUtils.random.nextInt(PartitionAssignorTest.MaxTopicCount + 1))
-
-      val topicPartitionCounts = Map((1 to topicCount).map(topic => {
-        ("topic-" + topic, PartitionAssignorTest.MinPartitionCount.max(TestUtils.random.nextInt(PartitionAssignorTest.MaxPartitionCount)))
-      }):_*)
-
-      val subscriptions = Map((1 to consumerCount).map { consumer =>
-        val streamCounts = Map((1 to topicCount).map(topic => {
-            val streamCount = 1.max(TestUtils.random.nextInt(PartitionAssignorTest.MaxStreamCount + 1))
-            ("topic-" + topic, streamCount)
-          }):_*)
-        ("g1c" + consumer, StaticSubscriptionInfo(streamCounts))
-      }:_*)
-      val scenario = Scenario("g1", topicPartitionCounts, subscriptions)
-      val zkUtils = PartitionAssignorTest.setupZkClientMock(scenario)
-      EasyMock.replay(zkUtils.zkClient)
-
-      PartitionAssignorTest.assignAndVerify(scenario, assignor, zkUtils)
-    }
-  }
-}
-
-private object PartitionAssignorTest extends Logging {
-
-  private val TestCaseCount = 3
-  private val MaxConsumerCount = 10
-  private val MaxStreamCount = 8
-  private val MaxTopicCount = 100
-  private val MinTopicCount = 0
-  private val MaxPartitionCount = 120
-  private val MinPartitionCount = 8
-
-  private trait SubscriptionInfo {
-    def registrationString: String
-  }
-
-  private case class StaticSubscriptionInfo(streamCounts: Map[String, Int]) extends SubscriptionInfo {
-    def registrationString =
-      Json.encode(Map("version" -> 1,
-                      "subscription" -> streamCounts,
-                      "pattern" -> "static",
-                      "timestamp" -> 1234.toString))
-
-    override def toString = {
-      "Stream counts: " + streamCounts
-    }
-  }
-
-  private case class WildcardSubscriptionInfo(streamCount: Int, regex: String, isWhitelist: Boolean)
-          extends SubscriptionInfo {
-    def registrationString =
-      Json.encode(Map("version" -> 1,
-                      "subscription" -> Map(regex -> streamCount),
-                      "pattern" -> (if (isWhitelist) "white_list" else "black_list")))
-
-    override def toString = {
-      "\"%s\":%d (%s)".format(regex, streamCount, if (isWhitelist) "whitelist" else "blacklist")
-    }
-  }
-
-  private case class Scenario(group: String,
-                              topicPartitionCounts: Map[String, Int],
-                              /* consumerId -> SubscriptionInfo */
-                              subscriptions: Map[String, SubscriptionInfo]) {
-    override def toString = {
-      "\n" +
-      "Group                  : %s\n".format(group) +
-      "Topic partition counts : %s\n".format(topicPartitionCounts) +
-      "Consumer assignment : %s\n".format(subscriptions)
-    }
-  }
-
-  private def setupZkClientMock(scenario: Scenario) = {
-    val consumers = java.util.Arrays.asList(scenario.subscriptions.keys.toSeq:_*)
-
-    val zkClient = EasyMock.createStrictMock(classOf[ZkClient])
-    val zkUtils = ZkUtils(zkClient, false)
-    EasyMock.checkOrder(zkClient, false)
-
-    EasyMock.expect(zkClient.getChildren("/consumers/%s/ids".format(scenario.group))).andReturn(consumers)
-    EasyMock.expectLastCall().anyTimes()
-
-    scenario.subscriptions.foreach { case(consumerId, subscriptionInfo) =>
-      EasyMock.expect(zkClient.readData("/consumers/%s/ids/%s".format(scenario.group, consumerId), new Stat()))
-              .andReturn(subscriptionInfo.registrationString)
-      EasyMock.expectLastCall().anyTimes()
-    }
-
-    scenario.topicPartitionCounts.foreach { case(topic, partitionCount) =>
-      val replicaAssignment = Map((0 until partitionCount).map(partition => (partition.toString, Seq(0))):_*)
-      EasyMock.expect(zkClient.readData("/brokers/topics/%s".format(topic), new Stat()))
-              .andReturn(zkUtils.replicaAssignmentZkData(replicaAssignment))
-      EasyMock.expectLastCall().anyTimes()
-    }
-
-    EasyMock.expect(zkUtils.zkClient.getChildren("/brokers/topics")).andReturn(
-      java.util.Arrays.asList(scenario.topicPartitionCounts.keys.toSeq:_*))
-    EasyMock.expectLastCall().anyTimes()
-
-    zkUtils
-  }
-
-  private def assignAndVerify(scenario: Scenario, assignor: PartitionAssignor, zkUtils: ZkUtils,
-                              verifyAssignmentIsUniform: Boolean = false) {
-    val assignments = scenario.subscriptions.map { case (consumer, _)  =>
-      val ctx = new AssignmentContext("g1", consumer, excludeInternalTopics = true, zkUtils)
-      assignor.assign(ctx).get(consumer)
-    }
-
-    // check for uniqueness (i.e., any partition should be assigned to exactly one consumer stream)
-    val globalAssignment = collection.mutable.Map[TopicAndPartition, ConsumerThreadId]()
-    assignments.foreach(assignment => {
-      assignment.foreach { case(topicPartition, owner) =>
-        val previousOwnerOpt = globalAssignment.put(topicPartition, owner)
-        assertTrue("Scenario %s: %s is assigned to two owners.".format(scenario, topicPartition), previousOwnerOpt.isEmpty)
-      }
-    })
-
-    // check for coverage (i.e., all given partitions are owned)
-    val assignedPartitions = globalAssignment.keySet
-    val givenPartitions = scenario.topicPartitionCounts.flatMap{ case (topic, partitionCount) =>
-      (0 until partitionCount).map(partition => TopicAndPartition(topic, partition))
-    }.toSet
-    assertTrue("Scenario %s: the list of given partitions and assigned partitions are different.".format(scenario),
-      givenPartitions == assignedPartitions)
-
-    // check for uniform assignment
-    if (verifyAssignmentIsUniform) {
-      val partitionCountForStream = partitionCountPerStream(globalAssignment)
-      if (partitionCountForStream.nonEmpty) {
-        val maxCount = partitionCountForStream.valuesIterator.max
-        val minCount = partitionCountForStream.valuesIterator.min
-        assertTrue("Scenario %s: assignment is not uniform (partition counts per stream are in the range [%d, %d])"
-          .format(scenario, minCount, maxCount), (maxCount - minCount) <= 1)
-      }
-    }
-  }
-
-  /** For each consumer stream, count the number of partitions that it owns. */
-  private def partitionCountPerStream(assignment: collection.Map[TopicAndPartition, ConsumerThreadId]) = {
-    val ownedCounts = collection.mutable.Map[ConsumerThreadId, Int]()
-    assignment.foreach { case (_, owner) =>
-      val updatedCount = ownedCounts.getOrElse(owner, 0) + 1
-      ownedCounts.put(owner, updatedCount)
-    }
-    ownedCounts
-  }
-}
-
diff --git a/core/src/test/scala/unit/kafka/consumer/TopicFilterTest.scala b/core/src/test/scala/unit/kafka/consumer/TopicFilterTest.scala
deleted file mode 100644
index ca8a0d627179d..0000000000000
--- a/core/src/test/scala/unit/kafka/consumer/TopicFilterTest.scala
+++ /dev/null
@@ -1,85 +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.consumer
-
-
-import org.apache.kafka.common.internals.Topic
-import org.junit.Assert._
-import org.scalatest.junit.JUnitSuite
-import org.junit.Test
-
-@deprecated("This test has been deprecated and will be removed in a future release.", "0.11.0.0")
-class TopicFilterTest extends JUnitSuite {
-
-  @Test
-  def testWhitelists() {
-
-    val topicFilter1 = Whitelist("white1,white2")
-    assertTrue(topicFilter1.isTopicAllowed("white2", excludeInternalTopics = true))
-    assertTrue(topicFilter1.isTopicAllowed("white2", excludeInternalTopics = false))
-    assertFalse(topicFilter1.isTopicAllowed("black1", excludeInternalTopics = true))
-    assertFalse(topicFilter1.isTopicAllowed("black1", excludeInternalTopics = false))
-
-    val topicFilter2 = Whitelist(".+")
-    assertTrue(topicFilter2.isTopicAllowed("alltopics", excludeInternalTopics = true))
-    assertFalse(topicFilter2.isTopicAllowed(Topic.GROUP_METADATA_TOPIC_NAME, excludeInternalTopics = true))
-    assertTrue(topicFilter2.isTopicAllowed(Topic.GROUP_METADATA_TOPIC_NAME, excludeInternalTopics = false))
-
-    val topicFilter3 = Whitelist("white_listed-topic.+")
-    assertTrue(topicFilter3.isTopicAllowed("white_listed-topic1", excludeInternalTopics = true))
-    assertFalse(topicFilter3.isTopicAllowed("black1", excludeInternalTopics = true))
-
-    val topicFilter4 = Whitelist("test-(?!bad\\b)[\\w]+")
-    assertTrue(topicFilter4.isTopicAllowed("test-good", excludeInternalTopics = true))
-    assertFalse(topicFilter4.isTopicAllowed("test-bad", excludeInternalTopics = true))
-  }
-
-  @Test
-  def testBlacklists() {
-    val topicFilter1 = Blacklist("black1")
-    assertTrue(topicFilter1.isTopicAllowed("white2", excludeInternalTopics = true))
-    assertTrue(topicFilter1.isTopicAllowed("white2", excludeInternalTopics = false))
-    assertFalse(topicFilter1.isTopicAllowed("black1", excludeInternalTopics = true))
-    assertFalse(topicFilter1.isTopicAllowed("black1", excludeInternalTopics = false))
-
-    assertFalse(topicFilter1.isTopicAllowed(Topic.GROUP_METADATA_TOPIC_NAME, excludeInternalTopics = true))
-    assertTrue(topicFilter1.isTopicAllowed(Topic.GROUP_METADATA_TOPIC_NAME, excludeInternalTopics = false))
-  }
-
-  @Test
-  def testWildcardTopicCountGetTopicCountMapEscapeJson() {
-    def getTopicCountMapKey(regex: String): String = {
-      val topicCount = new WildcardTopicCount(null, "consumerId", new Whitelist(regex), 1, true)
-      topicCount.getTopicCountMap.head._1
-    }
-    //lets make sure that the JSON strings are escaping as we expect
-    //if they are not then when they get saved to ZooKeeper and read back out they will be broken on parse
-    assertEquals("-\\\"-", getTopicCountMapKey("-\"-"))
-    assertEquals("-\\\\-", getTopicCountMapKey("-\\-"))
-    assertEquals("-\\/-", getTopicCountMapKey("-/-"))
-    assertEquals("-\\\\b-", getTopicCountMapKey("-\\b-"))
-    assertEquals("-\\\\f-", getTopicCountMapKey("-\\f-"))
-    assertEquals("-\\\\n-", getTopicCountMapKey("-\\n-"))
-    assertEquals("-\\\\r-", getTopicCountMapKey("-\\r-"))
-    assertEquals("-\\\\t-", getTopicCountMapKey("-\\t-"))
-    assertEquals("-\\\\u0000-", getTopicCountMapKey("-\\u0000-"))
-    assertEquals("-\\\\u001f-", getTopicCountMapKey("-\\u001f-"))
-    assertEquals("-\\\\u007f-", getTopicCountMapKey("-\\u007f-"))
-    assertEquals("-\\\\u009f-", getTopicCountMapKey("-\\u009f-"))
-  }
-}
diff --git a/core/src/test/scala/unit/kafka/consumer/ZookeeperConsumerConnectorTest.scala b/core/src/test/scala/unit/kafka/consumer/ZookeeperConsumerConnectorTest.scala
deleted file mode 100644
index c36400b27a5e0..0000000000000
--- a/core/src/test/scala/unit/kafka/consumer/ZookeeperConsumerConnectorTest.scala
+++ /dev/null
@@ -1,442 +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.consumer
-
-import java.util.Properties
-
-import org.junit.Assert._
-import kafka.common.MessageStreamsExistException
-import kafka.integration.KafkaServerTestHarness
-import kafka.javaapi.consumer.ConsumerRebalanceListener
-import kafka.message._
-import kafka.serializer._
-import kafka.server._
-import kafka.utils.TestUtils._
-import kafka.utils._
-import org.apache.log4j.{Level, Logger}
-import org.junit.{Test, After, Before}
-
-import scala.collection._
-
-@deprecated("This test has been deprecated and it will be removed in a future release", "0.10.0.0")
-class ZookeeperConsumerConnectorTest extends KafkaServerTestHarness with Logging {
-
-  val RebalanceBackoffMs = 5000
-  var dirs : ZKGroupTopicDirs = null
-  val numNodes = 2
-  val numParts = 2
-  val topic = "topic1"
-  val overridingProps = new Properties()
-  overridingProps.put(KafkaConfig.NumPartitionsProp, numParts.toString)
-
-  override def generateConfigs =
-    TestUtils.createBrokerConfigs(numNodes, zkConnect).map(KafkaConfig.fromProps(_, overridingProps))
-
-  val group = "group1"
-  val consumer0 = "consumer0"
-  val consumer1 = "consumer1"
-  val consumer2 = "consumer2"
-  val consumer3 = "consumer3"
-  val nMessages = 2
-
-  @Before
-  override def setUp() {
-    super.setUp()
-    dirs = new ZKGroupTopicDirs(group, topic)
-  }
-
-  @After
-  override def tearDown() {
-    super.tearDown()
-  }
-
-  @Test
-  def testBasic() {
-    val requestHandlerLogger = Logger.getLogger(classOf[KafkaRequestHandler])
-    requestHandlerLogger.setLevel(Level.FATAL)
-
-    // test consumer timeout logic
-    val consumerConfig0 = new ConsumerConfig(
-      TestUtils.createConsumerProperties(zkConnect, group, consumer0)) {
-      override val consumerTimeoutMs = 200
-    }
-    val zkConsumerConnector0 = new ZookeeperConsumerConnector(consumerConfig0, true)
-    val topicMessageStreams0 = zkConsumerConnector0.createMessageStreams(Map(topic -> 1), new StringDecoder(), new StringDecoder())
-
-    // no messages to consume, we should hit timeout;
-    // also the iterator should support re-entrant, so loop it twice
-    for (_ <- 0 until 2) {
-      try {
-        getMessages(topicMessageStreams0, nMessages * 2)
-        fail("should get an exception")
-      } catch {
-        case _: ConsumerTimeoutException => // this is ok
-      }
-    }
-
-    zkConsumerConnector0.shutdown
-
-    // send some messages to each broker
-    val sentMessages1 = sendMessages(servers, topic, nMessages, 0) ++
-      sendMessages(servers, topic, nMessages, 1)
-
-    // wait to make sure the topic and partition have a leader for the successful case
-    waitUntilLeaderIsElectedOrChanged(zkUtils, topic, 0)
-    waitUntilLeaderIsElectedOrChanged(zkUtils, topic, 1)
-
-    TestUtils.waitUntilMetadataIsPropagated(servers, topic, 0)
-    TestUtils.waitUntilMetadataIsPropagated(servers, topic, 1)
-
-    // create a consumer
-    val consumerConfig1 = new ConsumerConfig(TestUtils.createConsumerProperties(zkConnect, group, consumer1))
-    val zkConsumerConnector1 = new ZookeeperConsumerConnector(consumerConfig1, true)
-    val topicMessageStreams1 = zkConsumerConnector1.createMessageStreams(Map(topic -> 1), new StringDecoder(), new StringDecoder())
-
-    val receivedMessages1 = getMessages(topicMessageStreams1, nMessages * 2)
-    assertEquals(sentMessages1.sorted, receivedMessages1.sorted)
-
-    // also check partition ownership
-    val actual_1 = getZKChildrenValues(dirs.consumerOwnerDir)
-    val expected_1 = List( ("0", "group1_consumer1-0"),
-                           ("1", "group1_consumer1-0"))
-    assertEquals(expected_1, actual_1)
-
-    // commit consumed offsets
-    zkConsumerConnector1.commitOffsets(true)
-
-    // create a consumer
-    val consumerConfig2 = new ConsumerConfig(TestUtils.createConsumerProperties(zkConnect, group, consumer2)) {
-      override val rebalanceBackoffMs = RebalanceBackoffMs
-    }
-    val zkConsumerConnector2 = new ZookeeperConsumerConnector(consumerConfig2, true)
-    val topicMessageStreams2 = zkConsumerConnector2.createMessageStreams(Map(topic -> 1), new StringDecoder(), new StringDecoder())
-    // send some messages to each broker
-    val sentMessages2 = sendMessages(servers, topic, nMessages, 0) ++
-                         sendMessages(servers, topic, nMessages, 1)
-
-    waitUntilLeaderIsElectedOrChanged(zkUtils, topic, 0)
-    waitUntilLeaderIsElectedOrChanged(zkUtils, topic, 1)
-
-    val receivedMessages2 = getMessages(topicMessageStreams1, nMessages) ++ getMessages(topicMessageStreams2, nMessages)
-    assertEquals(sentMessages2.sorted, receivedMessages2.sorted)
-
-    // also check partition ownership
-    val actual_2 = getZKChildrenValues(dirs.consumerOwnerDir)
-    val expected_2 = List( ("0", "group1_consumer1-0"),
-                           ("1", "group1_consumer2-0"))
-    assertEquals(expected_2, actual_2)
-
-    // create a consumer with empty map
-    val consumerConfig3 = new ConsumerConfig(
-      TestUtils.createConsumerProperties(zkConnect, group, consumer3))
-    val zkConsumerConnector3 = new ZookeeperConsumerConnector(consumerConfig3, true)
-    zkConsumerConnector3.createMessageStreams(new mutable.HashMap[String, Int]())
-    // send some messages to each broker
-    val sentMessages3 = sendMessages(servers, topic, nMessages, 0) ++
-                        sendMessages(servers, topic, nMessages, 1)
-
-    waitUntilLeaderIsElectedOrChanged(zkUtils, topic, 0)
-    waitUntilLeaderIsElectedOrChanged(zkUtils, topic, 1)
-
-    val receivedMessages3 = getMessages(topicMessageStreams1, nMessages) ++ getMessages(topicMessageStreams2, nMessages)
-    assertEquals(sentMessages3.sorted, receivedMessages3.sorted)
-
-    // also check partition ownership
-    val actual_3 = getZKChildrenValues(dirs.consumerOwnerDir)
-    assertEquals(expected_2, actual_3)
-
-    // call createMesssageStreams twice should throw MessageStreamsExistException
-    try {
-      zkConsumerConnector3.createMessageStreams(new mutable.HashMap[String, Int]())
-      fail("Should fail with MessageStreamsExistException")
-    } catch {
-      case _: MessageStreamsExistException => // expected
-    }
-
-    zkConsumerConnector1.shutdown
-    zkConsumerConnector2.shutdown
-    zkConsumerConnector3.shutdown
-    info("all consumer connectors stopped")
-    requestHandlerLogger.setLevel(Level.ERROR)
-  }
-
-  @Test
-  def testCompression() {
-    val requestHandlerLogger = Logger.getLogger(classOf[kafka.server.KafkaRequestHandler])
-    requestHandlerLogger.setLevel(Level.FATAL)
-
-    // send some messages to each broker
-    val sentMessages1 = sendMessages(servers, topic, nMessages, 0, GZIPCompressionCodec) ++
-                        sendMessages(servers, topic, nMessages, 1, GZIPCompressionCodec)
-
-    waitUntilLeaderIsElectedOrChanged(zkUtils, topic, 0)
-    waitUntilLeaderIsElectedOrChanged(zkUtils, topic, 1)
-
-    TestUtils.waitUntilMetadataIsPropagated(servers, topic, 0)
-    TestUtils.waitUntilMetadataIsPropagated(servers, topic, 1)
-
-    // create a consumer
-    val consumerConfig1 = new ConsumerConfig(
-      TestUtils.createConsumerProperties(zkConnect, group, consumer1))
-    val zkConsumerConnector1 = new ZookeeperConsumerConnector(consumerConfig1, true)
-    val topicMessageStreams1 = zkConsumerConnector1.createMessageStreams(Map(topic -> 1), new StringDecoder(), new StringDecoder())
-    val receivedMessages1 = getMessages(topicMessageStreams1, nMessages * 2)
-    assertEquals(sentMessages1.sorted, receivedMessages1.sorted)
-
-    // also check partition ownership
-    val actual_1 = getZKChildrenValues(dirs.consumerOwnerDir)
-    val expected_1 = List( ("0", "group1_consumer1-0"),
-                           ("1", "group1_consumer1-0"))
-    assertEquals(expected_1, actual_1)
-
-    // commit consumed offsets
-    zkConsumerConnector1.commitOffsets(true)
-
-    // create a consumer
-    val consumerConfig2 = new ConsumerConfig(TestUtils.createConsumerProperties(zkConnect, group, consumer2)) {
-      override val rebalanceBackoffMs = RebalanceBackoffMs
-    }
-    val zkConsumerConnector2 = new ZookeeperConsumerConnector(consumerConfig2, true)
-    val topicMessageStreams2 = zkConsumerConnector2.createMessageStreams(Map(topic -> 1), new StringDecoder(), new StringDecoder())
-    // send some messages to each broker
-    val sentMessages2 = sendMessages(servers, topic, nMessages, 0, GZIPCompressionCodec) ++
-                        sendMessages(servers, topic, nMessages, 1, GZIPCompressionCodec)
-
-    waitUntilLeaderIsElectedOrChanged(zkUtils, topic, 0)
-    waitUntilLeaderIsElectedOrChanged(zkUtils, topic, 1)
-
-    val receivedMessages2 = getMessages(topicMessageStreams1, nMessages) ++ getMessages(topicMessageStreams2, nMessages)
-    assertEquals(sentMessages2.sorted, receivedMessages2.sorted)
-
-    // also check partition ownership
-    val actual_2 = getZKChildrenValues(dirs.consumerOwnerDir)
-    val expected_2 = List( ("0", "group1_consumer1-0"),
-                           ("1", "group1_consumer2-0"))
-    assertEquals(expected_2, actual_2)
-
-    // create a consumer with empty map
-    val consumerConfig3 = new ConsumerConfig(
-      TestUtils.createConsumerProperties(zkConnect, group, consumer3))
-    val zkConsumerConnector3 = new ZookeeperConsumerConnector(consumerConfig3, true)
-    zkConsumerConnector3.createMessageStreams(new mutable.HashMap[String, Int](), new StringDecoder(), new StringDecoder())
-    // send some messages to each broker
-    val sentMessages3 = sendMessages(servers, topic, nMessages, 0, GZIPCompressionCodec) ++
-                        sendMessages(servers, topic, nMessages, 1, GZIPCompressionCodec)
-
-    waitUntilLeaderIsElectedOrChanged(zkUtils, topic, 0)
-    waitUntilLeaderIsElectedOrChanged(zkUtils, topic, 1)
-
-    val receivedMessages3 = getMessages(topicMessageStreams1, nMessages) ++ getMessages(topicMessageStreams2, nMessages)
-    assertEquals(sentMessages3.sorted, receivedMessages3.sorted)
-
-    // also check partition ownership
-    val actual_3 = getZKChildrenValues(dirs.consumerOwnerDir)
-    assertEquals(expected_2, actual_3)
-
-    zkConsumerConnector1.shutdown
-    zkConsumerConnector2.shutdown
-    zkConsumerConnector3.shutdown
-    info("all consumer connectors stopped")
-    requestHandlerLogger.setLevel(Level.ERROR)
-  }
-
-  @Test
-  def testCompressionSetConsumption() {
-    // send some messages to each broker
-    val sentMessages = sendMessages(servers, topic, 200, 0, DefaultCompressionCodec) ++
-                       sendMessages(servers, topic, 200, 1, DefaultCompressionCodec)
-
-    TestUtils.waitUntilMetadataIsPropagated(servers, topic, 0)
-    TestUtils.waitUntilMetadataIsPropagated(servers, topic, 1)
-
-    val consumerConfig1 = new ConsumerConfig(TestUtils.createConsumerProperties(zkConnect, group, consumer0))
-    val zkConsumerConnector1 = new ZookeeperConsumerConnector(consumerConfig1, true)
-    val topicMessageStreams1 = zkConsumerConnector1.createMessageStreams(Map(topic -> 1), new StringDecoder(), new StringDecoder())
-    val receivedMessages = getMessages(topicMessageStreams1, 400)
-    assertEquals(sentMessages.sorted, receivedMessages.sorted)
-
-    // also check partition ownership
-    val actual_2 = getZKChildrenValues(dirs.consumerOwnerDir)
-    val expected_2 = List( ("0", "group1_consumer0-0"),
-                           ("1", "group1_consumer0-0"))
-    assertEquals(expected_2, actual_2)
-
-    zkConsumerConnector1.shutdown
-  }
-
-  @Test
-  def testConsumerDecoder() {
-    val requestHandlerLogger = Logger.getLogger(classOf[kafka.server.KafkaRequestHandler])
-    requestHandlerLogger.setLevel(Level.FATAL)
-
-    // send some messages to each broker
-    val sentMessages = sendMessages(servers, topic, nMessages, 0, NoCompressionCodec) ++
-                       sendMessages(servers, topic, nMessages, 1, NoCompressionCodec)
-
-    TestUtils.waitUntilMetadataIsPropagated(servers, topic, 0)
-    TestUtils.waitUntilMetadataIsPropagated(servers, topic, 1)
-
-    val consumerConfig = new ConsumerConfig(TestUtils.createConsumerProperties(zkConnect, group, consumer1))
-
-    waitUntilLeaderIsElectedOrChanged(zkUtils, topic, 0)
-    waitUntilLeaderIsElectedOrChanged(zkUtils, topic, 1)
-
-    val zkConsumerConnector =
-      new ZookeeperConsumerConnector(consumerConfig, true)
-    val topicMessageStreams =
-      zkConsumerConnector.createMessageStreams(Map(topic -> 1), new StringDecoder(), new StringDecoder())
-
-    var receivedMessages: List[String] = Nil
-    for (messageStreams <- topicMessageStreams.values) {
-      for (messageStream <- messageStreams) {
-        val iterator = messageStream.iterator
-        for (_ <- 0 until nMessages * 2) {
-          assertTrue(iterator.hasNext())
-          val message = iterator.next().message
-          receivedMessages ::= message
-          debug("received message: " + message)
-        }
-      }
-    }
-    assertEquals(sentMessages.sorted, receivedMessages.sorted)
-
-    zkConsumerConnector.shutdown()
-    requestHandlerLogger.setLevel(Level.ERROR)
-  }
-
-  @Test
-  def testLeaderSelectionForPartition() {
-    val zkUtils = ZkUtils(zkConnect, 6000, 30000, false)
-
-    // create topic topic1 with 1 partition on broker 0
-    createTopic(zkUtils, topic, numPartitions = 1, replicationFactor = 1, servers = servers)
-
-    // send some messages to each broker
-    val sentMessages1 = sendMessages(servers, topic, nMessages)
-
-    // create a consumer
-    val consumerConfig1 = new ConsumerConfig(TestUtils.createConsumerProperties(zkConnect, group, consumer1))
-    val zkConsumerConnector1 = new ZookeeperConsumerConnector(consumerConfig1, true)
-    val topicMessageStreams1 = zkConsumerConnector1.createMessageStreams(Map(topic -> 1), new StringDecoder(), new StringDecoder())
-    val topicRegistry = zkConsumerConnector1.getTopicRegistry
-    assertEquals(1, topicRegistry.map(r => r._1).size)
-    assertEquals(topic, topicRegistry.map(r => r._1).head)
-    val topicsAndPartitionsInRegistry = topicRegistry.map(r => (r._1, r._2.map(p => p._2)))
-    val brokerPartition = topicsAndPartitionsInRegistry.head._2.head
-    assertEquals(0, brokerPartition.partitionId)
-
-    // also check partition ownership
-    val actual_1 = getZKChildrenValues(dirs.consumerOwnerDir)
-    val expected_1 = List( ("0", "group1_consumer1-0"))
-    assertEquals(expected_1, actual_1)
-
-    val receivedMessages1 = getMessages(topicMessageStreams1, nMessages)
-    assertEquals(sentMessages1, receivedMessages1)
-    zkConsumerConnector1.shutdown()
-    zkUtils.close()
-  }
-
-  @Test
-  def testConsumerRebalanceListener() {
-    // Send messages to create topic
-    sendMessages(servers, topic, nMessages, 0)
-    sendMessages(servers, topic, nMessages, 1)
-
-    val consumerConfig1 = new ConsumerConfig(TestUtils.createConsumerProperties(zkConnect, group, consumer1))
-    val zkConsumerConnector1 = new ZookeeperConsumerConnector(consumerConfig1, true)
-    // Register consumer rebalance listener
-    val rebalanceListener1 = new TestConsumerRebalanceListener()
-    zkConsumerConnector1.setConsumerRebalanceListener(rebalanceListener1)
-    val topicMessageStreams1 = zkConsumerConnector1.createMessageStreams(Map(topic -> 1), new StringDecoder(), new StringDecoder())
-
-    // Check if rebalance listener is fired
-    assertEquals(true, rebalanceListener1.beforeReleasingPartitionsCalled)
-    assertEquals(true, rebalanceListener1.beforeStartingFetchersCalled)
-    assertEquals(null, rebalanceListener1.partitionOwnership.get(topic))
-    // Check if partition assignment in rebalance listener is correct
-    assertEquals("group1_consumer1", rebalanceListener1.globalPartitionOwnership.get(topic).get(0).consumer)
-    assertEquals("group1_consumer1", rebalanceListener1.globalPartitionOwnership.get(topic).get(1).consumer)
-    assertEquals(0, rebalanceListener1.globalPartitionOwnership.get(topic).get(0).threadId)
-    assertEquals(0, rebalanceListener1.globalPartitionOwnership.get(topic).get(1).threadId)
-    assertEquals("group1_consumer1", rebalanceListener1.consumerId)
-    // reset the flag
-    rebalanceListener1.beforeReleasingPartitionsCalled = false
-    rebalanceListener1.beforeStartingFetchersCalled = false
-
-    val actual_1 = getZKChildrenValues(dirs.consumerOwnerDir)
-    val expected_1 = List(("0", "group1_consumer1-0"),
-                          ("1", "group1_consumer1-0"))
-    assertEquals(expected_1, actual_1)
-
-    val consumerConfig2 = new ConsumerConfig(TestUtils.createConsumerProperties(zkConnect, group, consumer2))
-    val zkConsumerConnector2 = new ZookeeperConsumerConnector(consumerConfig2, true)
-    // Register consumer rebalance listener
-    val rebalanceListener2 = new TestConsumerRebalanceListener()
-    zkConsumerConnector2.setConsumerRebalanceListener(rebalanceListener2)
-    zkConsumerConnector2.createMessageStreams(Map(topic -> 1), new StringDecoder(), new StringDecoder())
-
-    // Consume messages from consumer 1 to make sure it has finished rebalance
-    getMessages(topicMessageStreams1, nMessages)
-
-    val actual_2 = getZKChildrenValues(dirs.consumerOwnerDir)
-    val expected_2 = List(("0", "group1_consumer1-0"),
-                          ("1", "group1_consumer2-0"))
-    assertEquals(expected_2, actual_2)
-
-    // Check if rebalance listener is fired
-    assertEquals(true, rebalanceListener1.beforeReleasingPartitionsCalled)
-    assertEquals(true, rebalanceListener1.beforeStartingFetchersCalled)
-    assertEquals(Set[Int](0, 1), rebalanceListener1.partitionOwnership.get(topic))
-    // Check if global partition ownership in rebalance listener is correct
-    assertEquals("group1_consumer1", rebalanceListener1.globalPartitionOwnership.get(topic).get(0).consumer)
-    assertEquals("group1_consumer2", rebalanceListener1.globalPartitionOwnership.get(topic).get(1).consumer)
-    assertEquals(0, rebalanceListener1.globalPartitionOwnership.get(topic).get(0).threadId)
-    assertEquals(0, rebalanceListener1.globalPartitionOwnership.get(topic).get(1).threadId)
-    assertEquals("group1_consumer1", rebalanceListener1.consumerId)
-    assertEquals("group1_consumer2", rebalanceListener2.consumerId)
-    assertEquals(rebalanceListener1.globalPartitionOwnership, rebalanceListener2.globalPartitionOwnership)
-    zkConsumerConnector1.shutdown()
-    zkConsumerConnector2.shutdown()
-  }
-
-  def getZKChildrenValues(path : String) : Seq[Tuple2[String,String]] = {
-    val children = zkUtils.getChildren(path).sorted
-    children.map(partition =>
-      (partition, zkUtils.zkClient.readData(path + "/" + partition).asInstanceOf[String]))
-  }
-
-  private class TestConsumerRebalanceListener extends ConsumerRebalanceListener {
-    var beforeReleasingPartitionsCalled: Boolean = false
-    var beforeStartingFetchersCalled: Boolean = false
-    var consumerId: String = ""
-    var partitionOwnership: java.util.Map[String, java.util.Set[java.lang.Integer]] = null
-    var globalPartitionOwnership: java.util.Map[String, java.util.Map[java.lang.Integer, ConsumerThreadId]] = null
-
-    override def beforeReleasingPartitions(partitionOwnership: java.util.Map[String, java.util.Set[java.lang.Integer]]) {
-      beforeReleasingPartitionsCalled = true
-      this.partitionOwnership = partitionOwnership
-    }
-
-    override def beforeStartingFetchers(consumerId: String, globalPartitionOwnership: java.util.Map[String, java.util.Map[java.lang.Integer, ConsumerThreadId]]) {
-      beforeStartingFetchersCalled = true
-      this.consumerId = consumerId
-      this.globalPartitionOwnership = globalPartitionOwnership
-    }
-  }
-
-}
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..58c7203794196
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/controller/ControllerChannelManagerTest.scala
@@ -0,0 +1,923 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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_4_IV0, KAFKA_2_4_IV1, KAFKA_2_6_IV0, 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, UpdateMetadataResponseData}
+import org.apache.kafka.common.message.LeaderAndIsrResponseData.LeaderAndIsrPartitionError
+import org.apache.kafka.common.message.StopReplicaRequestData.StopReplicaPartitionState
+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, UpdateMetadataResponse}
+import org.apache.kafka.common.security.auth.SecurityProtocol
+import org.junit.Assert._
+import org.junit.Test
+import org.scalatest.Assertions
+
+import scala.jdk.CollectionConverters._
+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.putPartitionLeadershipInfo(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(leaderAndIsrResponse, brokerId) = batch.sentEvents.head
+    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.putPartitionLeadershipInfo(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.putPartitionLeadershipInfo(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) 4
+        else if (apiVersion >= KAFKA_2_4_IV0) 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.putPartitionLeadershipInfo(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.putPartitionLeadershipInfo(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)
+
+    applyUpdateMetadataResponseCallbacks(Errors.STALE_BROKER_EPOCH, batch.sentRequests(2).toList)
+    assertEquals(1, batch.sentEvents.size)
+
+    val UpdateMetadataResponseReceived(updateMetadataResponse, brokerId) = batch.sentEvents.head
+    assertEquals(2, brokerId)
+    assertEquals(Errors.STALE_BROKER_EPOCH, updateMetadataResponse.error)
+  }
+
+  @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.putPartitionLeadershipInfo(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) 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 = Map(
+      new TopicPartition("foo", 0) -> LeaderAndDelete(1, false),
+      new TopicPartition("foo", 1) -> LeaderAndDelete(2, false),
+      new TopicPartition("bar", 1) -> LeaderAndDelete(3, false)
+    )
+
+    batch.newBatch()
+    partitions.foreach { case (partition, LeaderAndDelete(leaderAndIsr, deletePartition)) =>
+      context.putPartitionLeadershipInfo(partition, LeaderIsrAndControllerEpoch(leaderAndIsr, controllerEpoch))
+      batch.addStopReplicaRequestForBrokers(Seq(2), partition, deletePartition)
+    }
+    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
+    assertEquals(partitionStates(partitions), stopReplicaRequest.partitionStates().asScala)
+
+    applyStopReplicaResponseCallbacks(Errors.NONE, batch.sentRequests(2).toList)
+    assertEquals(0, batch.sentEvents.size)
+  }
+
+  @Test
+  def testStopReplicaRequestWithAlreadyDefinedDeletedPartition(): 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))
+    context.putPartitionLeadershipInfo(partition, LeaderIsrAndControllerEpoch(leaderAndIsr, controllerEpoch))
+
+    batch.newBatch()
+    batch.addStopReplicaRequestForBrokers(Seq(2), partition, deletePartition = true)
+    batch.addStopReplicaRequestForBrokers(Seq(2), partition, deletePartition = false)
+    batch.sendRequestsToBrokers(controllerEpoch)
+
+    val sentStopReplicaRequests = batch.collectStopReplicaRequestsFor(2)
+    assertEquals(1, sentStopReplicaRequests.size)
+
+    val stopReplicaRequest = sentStopReplicaRequests.head
+    assertEquals(partitionStates(Map(partition -> LeaderAndDelete(leaderAndIsr, true))),
+      stopReplicaRequest.partitionStates().asScala)
+  }
+
+  @Test
+  def testStopReplicaRequestsWhileTopicQueuedForDeletion(): Unit = {
+    for (apiVersion <- ApiVersion.allVersions) {
+      testStopReplicaRequestsWhileTopicQueuedForDeletion(apiVersion)
+    }
+  }
+
+  private def testStopReplicaRequestsWhileTopicQueuedForDeletion(interBrokerProtocolVersion: ApiVersion): Unit = {
+    val context = initContext(Seq(1, 2, 3), Set("foo", "bar"), 2, 3)
+    val config = createConfig(interBrokerProtocolVersion)
+    val batch = new MockControllerBrokerRequestBatch(context, config)
+
+    val partitions = Map(
+      new TopicPartition("foo", 0) -> LeaderAndDelete(1, true),
+      new TopicPartition("foo", 1) -> LeaderAndDelete(2, true),
+      new TopicPartition("bar", 1) -> LeaderAndDelete(3, true)
+    )
+
+    // Topic deletion is queued, but has not begun
+    context.queueTopicDeletion(Set("foo"))
+
+    batch.newBatch()
+    partitions.foreach { case (partition, LeaderAndDelete(leaderAndIsr, deletePartition)) =>
+      context.putPartitionLeadershipInfo(partition, LeaderIsrAndControllerEpoch(leaderAndIsr, controllerEpoch))
+      batch.addStopReplicaRequestForBrokers(Seq(2), partition, deletePartition)
+    }
+    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
+    assertEquals(partitionStates(partitions, context.topicsQueuedForDeletion, stopReplicaRequest.version),
+      stopReplicaRequest.partitionStates().asScala)
+
+    // 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 = {
+    for (apiVersion <- ApiVersion.allVersions) {
+      testStopReplicaRequestsWhileTopicDeletionStarted(apiVersion)
+    }
+  }
+
+  private def testStopReplicaRequestsWhileTopicDeletionStarted(interBrokerProtocolVersion: ApiVersion): Unit = {
+    val context = initContext(Seq(1, 2, 3), Set("foo", "bar"), 2, 3)
+    val config = createConfig(interBrokerProtocolVersion)
+    val batch = new MockControllerBrokerRequestBatch(context, config)
+
+    val partitions = Map(
+      new TopicPartition("foo", 0) -> LeaderAndDelete(1, true),
+      new TopicPartition("foo", 1) -> LeaderAndDelete(2, true),
+      new TopicPartition("bar", 1) -> LeaderAndDelete(3, true)
+    )
+
+    context.queueTopicDeletion(Set("foo"))
+    context.beginTopicDeletion(Set("foo"))
+
+    batch.newBatch()
+    partitions.foreach { case (partition, LeaderAndDelete(leaderAndIsr, deletePartition)) =>
+      context.putPartitionLeadershipInfo(partition, LeaderIsrAndControllerEpoch(leaderAndIsr, controllerEpoch))
+      batch.addStopReplicaRequestForBrokers(Seq(2), partition, deletePartition)
+    }
+    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
+    assertEquals(partitionStates(partitions, context.topicsQueuedForDeletion, stopReplicaRequest.version),
+      stopReplicaRequest.partitionStates().asScala)
+
+    // 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.keys.filter(_.topic == "foo"), includedPartitions)
+  }
+
+  @Test
+  def testStopReplicaRequestWithoutDeletePartitionWhileTopicDeletionStarted(): Unit = {
+    for (apiVersion <- ApiVersion.allVersions) {
+      testStopReplicaRequestWithoutDeletePartitionWhileTopicDeletionStarted(apiVersion)
+    }
+  }
+
+  private def testStopReplicaRequestWithoutDeletePartitionWhileTopicDeletionStarted(interBrokerProtocolVersion: ApiVersion): Unit = {
+    val context = initContext(Seq(1, 2, 3), Set("foo", "bar"), 2, 3)
+    val config = createConfig(interBrokerProtocolVersion)
+    val batch = new MockControllerBrokerRequestBatch(context, config)
+
+    val partitions = Map(
+      new TopicPartition("foo", 0) -> LeaderAndDelete(1, false),
+      new TopicPartition("foo", 1) -> LeaderAndDelete(2, false),
+      new TopicPartition("bar", 1) -> LeaderAndDelete(3, false)
+    )
+
+    context.queueTopicDeletion(Set("foo"))
+    context.beginTopicDeletion(Set("foo"))
+
+    batch.newBatch()
+    partitions.foreach { case (partition, LeaderAndDelete(leaderAndIsr, deletePartition)) =>
+      context.putPartitionLeadershipInfo(partition, LeaderIsrAndControllerEpoch(leaderAndIsr, controllerEpoch))
+      batch.addStopReplicaRequestForBrokers(Seq(2), partition, deletePartition)
+    }
+    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
+    assertEquals(partitionStates(partitions, context.topicsQueuedForDeletion, stopReplicaRequest.version),
+      stopReplicaRequest.partitionStates().asScala)
+
+    // No events should be fired
+    applyStopReplicaResponseCallbacks(Errors.NONE, batch.sentRequests(2).toList)
+    assertEquals(0, batch.sentEvents.size)
+  }
+
+  @Test
+  def testMixedDeleteAndNotDeleteStopReplicaRequests(): Unit = {
+    testMixedDeleteAndNotDeleteStopReplicaRequests(ApiVersion.latestVersion,
+      ApiKeys.STOP_REPLICA.latestVersion)
+
+    for (apiVersion <- ApiVersion.allVersions) {
+      if (apiVersion < KAFKA_2_2_IV0)
+        testMixedDeleteAndNotDeleteStopReplicaRequests(apiVersion, 0.toShort)
+      else if (apiVersion < KAFKA_2_4_IV1)
+        testMixedDeleteAndNotDeleteStopReplicaRequests(apiVersion, 1.toShort)
+      else if (apiVersion < KAFKA_2_6_IV0)
+        testMixedDeleteAndNotDeleteStopReplicaRequests(apiVersion, 2.toShort)
+      else
+        testMixedDeleteAndNotDeleteStopReplicaRequests(apiVersion, 3.toShort)
+    }
+  }
+
+  private def testMixedDeleteAndNotDeleteStopReplicaRequests(interBrokerProtocolVersion: ApiVersion,
+                                                             expectedStopReplicaRequestVersion: 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 deletePartitions = Map(
+      new TopicPartition("foo", 0) -> LeaderAndDelete(1, true),
+      new TopicPartition("foo", 1) -> LeaderAndDelete(2, true)
+    )
+
+    val nonDeletePartitions = Map(
+      new TopicPartition("bar", 0) -> LeaderAndDelete(1, false),
+      new TopicPartition("bar", 1) -> LeaderAndDelete(2, false)
+    )
+
+    batch.newBatch()
+    deletePartitions.foreach { case (partition, LeaderAndDelete(leaderAndIsr, deletePartition)) =>
+      context.putPartitionLeadershipInfo(partition, LeaderIsrAndControllerEpoch(leaderAndIsr, controllerEpoch))
+      batch.addStopReplicaRequestForBrokers(Seq(2), partition, deletePartition)
+    }
+    nonDeletePartitions.foreach { case (partition, LeaderAndDelete(leaderAndIsr, deletePartition)) =>
+      context.putPartitionLeadershipInfo(partition, LeaderIsrAndControllerEpoch(leaderAndIsr, controllerEpoch))
+      batch.addStopReplicaRequestForBrokers(Seq(2), partition, deletePartition)
+    }
+    batch.sendRequestsToBrokers(controllerEpoch)
+
+    assertEquals(0, batch.sentEvents.size)
+    assertEquals(1, batch.sentRequests.size)
+    assertTrue(batch.sentRequests.contains(2))
+
+    // Since KAFKA_2_6_IV0, only one StopReplicaRequest is sent out
+    if (interBrokerProtocolVersion >= KAFKA_2_6_IV0) {
+      val sentRequests = batch.sentRequests(2)
+      assertEquals(1, sentRequests.size)
+
+      val sentStopReplicaRequests = batch.collectStopReplicaRequestsFor(2)
+      assertEquals(1, sentStopReplicaRequests.size)
+
+      val stopReplicaRequest = sentStopReplicaRequests.head
+      assertEquals(partitionStates(deletePartitions ++ nonDeletePartitions, version = stopReplicaRequest.version),
+        stopReplicaRequest.partitionStates().asScala)
+    } else {
+      val sentRequests = batch.sentRequests(2)
+      assertEquals(2, sentRequests.size)
+
+      val sentStopReplicaRequests = batch.collectStopReplicaRequestsFor(2)
+      assertEquals(2, sentStopReplicaRequests.size)
+
+      // StopReplicaRequest (deletePartitions = true) is sent first
+      val stopReplicaRequestWithDelete = sentStopReplicaRequests(0)
+      assertEquals(partitionStates(deletePartitions, version = stopReplicaRequestWithDelete.version),
+        stopReplicaRequestWithDelete.partitionStates().asScala)
+      val stopReplicaRequestWithoutDelete = sentStopReplicaRequests(1)
+      assertEquals(partitionStates(nonDeletePartitions, version = stopReplicaRequestWithoutDelete.version),
+        stopReplicaRequestWithoutDelete.partitionStates().asScala)
+    }
+  }
+
+  @Test
+  def testStopReplicaGroupsByBroker(): 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) -> LeaderAndDelete(1, false),
+      new TopicPartition("foo", 1) -> LeaderAndDelete(2, false),
+      new TopicPartition("bar", 1) -> LeaderAndDelete(3, false)
+    )
+
+    batch.newBatch()
+    partitions.foreach { case (partition, LeaderAndDelete(leaderAndIsr, deletePartition)) =>
+      context.putPartitionLeadershipInfo(partition, LeaderIsrAndControllerEpoch(leaderAndIsr, controllerEpoch))
+      batch.addStopReplicaRequestForBrokers(Seq(2, 3), partition, deletePartition)
+    }
+    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
+      assertEquals(partitionStates(partitions), stopReplicaRequest.partitionStates().asScala)
+
+      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 = Map(
+      new TopicPartition("foo", 0) -> LeaderAndDelete(1, false),
+      new TopicPartition("foo", 1) -> LeaderAndDelete(2, false),
+      new TopicPartition("bar", 1) -> LeaderAndDelete(3, false)
+    )
+
+    batch.newBatch()
+    partitions.foreach { case (partition, LeaderAndDelete(leaderAndIsr, deletePartition)) =>
+      context.putPartitionLeadershipInfo(partition, LeaderIsrAndControllerEpoch(leaderAndIsr, controllerEpoch))
+      batch.addStopReplicaRequestForBrokers(Seq(2, 3), partition, deletePartition)
+    }
+    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
+    assertEquals(partitionStates(partitions), stopReplicaRequest.partitionStates().asScala)
+  }
+
+  @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_4_IV1)
+        testStopReplicaFollowsInterBrokerProtocolVersion(apiVersion, 1.toShort)
+      else if (apiVersion < KAFKA_2_6_IV0)
+        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)
+    val leaderAndIsr = LeaderAndIsr(1, List(1, 2))
+
+    context.putPartitionLeadershipInfo(partition, LeaderIsrAndControllerEpoch(leaderAndIsr, controllerEpoch))
+
+    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 case class LeaderAndDelete(leaderAndIsr: LeaderAndIsr,
+                                     deletePartition: Boolean)
+
+  private object LeaderAndDelete {
+    def apply(leader: Int, deletePartition: Boolean): LeaderAndDelete =
+      new LeaderAndDelete(LeaderAndIsr(leader, List()), deletePartition)
+  }
+
+  private def partitionStates(partitions: Map[TopicPartition, LeaderAndDelete],
+                              topicsQueuedForDeletion: collection.Set[String] = Set.empty[String],
+                              version: Short = ApiKeys.STOP_REPLICA.latestVersion): Map[TopicPartition, StopReplicaPartitionState] = {
+    partitions.map { case (topicPartition, LeaderAndDelete(leaderAndIsr, deletePartition)) =>
+      topicPartition -> {
+        val partitionState = new StopReplicaPartitionState()
+          .setPartitionIndex(topicPartition.partition)
+          .setDeletePartition(deletePartition)
+
+        if (version >= 3) {
+          partitionState.setLeaderEpoch(if (topicsQueuedForDeletion.contains(topicPartition.topic))
+            LeaderAndIsr.EpochDuringDelete
+          else
+            leaderAndIsr.leaderEpoch)
+        }
+
+        partitionState
+      }
+    }
+  }
+
+  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.topicStates.asScala.flatMap { topic =>
+            topic.partitionStates.asScala.map { partition =>
+              new StopReplicaPartitionError()
+                .setTopicName(topic.topicName)
+                .setPartitionIndex(partition.partitionIndex)
+                .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 applyUpdateMetadataResponseCallbacks(error: Errors, sentRequests: List[SentRequest]): Unit = {
+    sentRequests.filter(_.request.apiKey == ApiKeys.UPDATE_METADATA).filter(_.responseCallback != null).foreach { sentRequest =>
+      val response = new UpdateMetadataResponse(new UpdateMetadataResponseData().setErrorCode(error.code))
+      sentRequest.responseCallback(response)
+    }
+  }
+
+  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.setLiveBrokers(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.updatePartitionFullReplicaAssignment(partition, ReplicaAssignment(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..66290d18e2773
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/controller/ControllerContextTest.scala
@@ -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 unit.kafka.controller
+
+import kafka.api.LeaderAndIsr
+import kafka.cluster.{Broker, EndPoint}
+import kafka.controller.LeaderIsrAndControllerEpoch
+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.setLiveBrokers(brokerEpochs)
+
+    // Simple round-robin replica assignment
+    var leaderIndex = 0
+    Seq(tp1, tp2, tp3).foreach { partition =>
+      val replicas = brokers.indices.map { i =>
+        brokers((i + leaderIndex) % brokers.size)
+      }
+      context.updatePartitionFullReplicaAssignment(partition, ReplicaAssignment(replicas))
+      leaderIndex += 1
+    }
+  }
+
+  @Test
+  def testUpdatePartitionFullReplicaAssignmentUpdatesReplicaAssignment(): Unit = {
+    val initialReplicas = Seq(4)
+    context.updatePartitionFullReplicaAssignment(tp1, ReplicaAssignment(initialReplicas))
+    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(Seq(1, 2, 3, 4)).reassignTo(Seq(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(Seq(1, 2, 3)).reassignTo(Seq(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(Seq(1, 2, 3)).reassignTo(Seq(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 testReassignToIdempotence(): Unit = {
+    val assignment1 = ReplicaAssignment(Seq(1, 2, 3))
+    assertEquals(assignment1, assignment1.reassignTo(assignment1.targetReplicas))
+
+    val assignment2 = ReplicaAssignment(Seq(4, 5, 6, 1, 2, 3),
+      addingReplicas = Seq(4, 5, 6), removingReplicas = Seq(1, 2, 3))
+    assertEquals(assignment2, assignment2.reassignTo(assignment2.targetReplicas))
+
+    val assignment3 = ReplicaAssignment(Seq(4, 2, 3, 1),
+      addingReplicas = Seq(4), removingReplicas = Seq(1))
+    assertEquals(assignment3, assignment3.reassignTo(assignment3.targetReplicas))
+  }
+
+  @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)))
+  }
+
+  @Test
+  def testPreferredReplicaImbalanceMetric(): Unit = {
+    context.updatePartitionFullReplicaAssignment(tp1, ReplicaAssignment(Seq(1, 2, 3)))
+    context.updatePartitionFullReplicaAssignment(tp2, ReplicaAssignment(Seq(1, 2, 3)))
+    context.updatePartitionFullReplicaAssignment(tp3, ReplicaAssignment(Seq(1, 2, 3)))
+    assertEquals(0, context.preferredReplicaImbalanceCount)
+
+    context.putPartitionLeadershipInfo(tp1, LeaderIsrAndControllerEpoch(LeaderAndIsr(1, List(1, 2, 3)), 0))
+    assertEquals(0, context.preferredReplicaImbalanceCount)
+
+    context.putPartitionLeadershipInfo(tp2, LeaderIsrAndControllerEpoch(LeaderAndIsr(2, List(2, 3, 1)), 0))
+    assertEquals(1, context.preferredReplicaImbalanceCount)
+
+    context.putPartitionLeadershipInfo(tp3, LeaderIsrAndControllerEpoch(LeaderAndIsr(3, List(3, 1, 2)), 0))
+    assertEquals(2, context.preferredReplicaImbalanceCount)
+
+    context.updatePartitionFullReplicaAssignment(tp1, ReplicaAssignment(Seq(2, 3, 1)))
+    context.updatePartitionFullReplicaAssignment(tp2, ReplicaAssignment(Seq(2, 3, 1)))
+    assertEquals(2, context.preferredReplicaImbalanceCount)
+
+    context.queueTopicDeletion(Set(tp3.topic))
+    assertEquals(1, context.preferredReplicaImbalanceCount)
+
+    context.putPartitionLeadershipInfo(tp3, LeaderIsrAndControllerEpoch(LeaderAndIsr(1, List(3, 1, 2)), 0))
+    assertEquals(1, context.preferredReplicaImbalanceCount)
+
+    context.removeTopic(tp1.topic)
+    context.removeTopic(tp2.topic)
+    context.removeTopic(tp3.topic)
+    assertEquals(0, context.preferredReplicaImbalanceCount)
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/controller/ControllerEventManagerTest.scala b/core/src/test/scala/unit/kafka/controller/ControllerEventManagerTest.scala
index ec9343e042560..7e4c23a5698c7 100644
--- a/core/src/test/scala/unit/kafka/controller/ControllerEventManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/controller/ControllerEventManagerTest.scala
@@ -20,13 +20,19 @@ package kafka.controller
 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.controller
+import kafka.metrics.KafkaYammerMetrics
 import kafka.utils.TestUtils
+import org.apache.kafka.common.message.UpdateMetadataResponseData
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.requests.UpdateMetadataResponse
+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._
+import scala.jdk.CollectionConverters._
+import scala.collection.mutable
 
 class ControllerEventManagerTest {
 
@@ -38,37 +44,168 @@ 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] = {
+      KafkaYammerMetrics.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 testEventWithoutRateMetrics(): Unit = {
+    val time = new MockTime()
+    val controllerStats = new ControllerStats
+    val processedEvents = mutable.Set.empty[ControllerEvent]
+
+    val eventProcessor = new ControllerEventProcessor {
+      override def process(event: ControllerEvent): Unit = { processedEvents += event }
+      override def preempt(event: ControllerEvent): Unit = {}
+    }
+
+    controllerEventManager = new ControllerEventManager(0, eventProcessor,
+      time, controllerStats.rateAndTimeMetrics)
+    controllerEventManager.start()
+
+    val updateMetadataResponse = new UpdateMetadataResponse(
+      new UpdateMetadataResponseData().setErrorCode(Errors.NONE.code)
+    )
+    val updateMetadataResponseEvent = controller.UpdateMetadataResponseReceived(updateMetadataResponse, brokerId = 1)
+    controllerEventManager.put(updateMetadataResponseEvent)
+    TestUtils.waitUntilTrue(() => processedEvents.size == 1,
+      "Failed to process expected event before timing out")
+    assertEquals(updateMetadataResponseEvent, processedEvents.head)
+  }
+
+  @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(KafkaYammerMetrics.defaultRegistry.allMetrics.asScala.filter { case (k, _) =>
+      k.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 = KafkaYammerMetrics.defaultRegistry.allMetrics.asScala.filter { case (k, _) =>
+      k.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 testEventQueueTimeResetOnTimeout(): Unit = {
+    val metricName = "kafka.controller:type=ControllerEventManager,name=EventQueueTimeMs"
+    val controllerStats = new ControllerStats
+    val time = new MockTime()
+    val processedEvents = new AtomicInteger()
+
+    val eventProcessor = new ControllerEventProcessor {
+      override def process(event: ControllerEvent): Unit = {
+        processedEvents.incrementAndGet()
+      }
+      override def preempt(event: ControllerEvent): Unit = {}
+    }
+
+    controllerEventManager = new ControllerEventManager(0, eventProcessor,
+      time, controllerStats.rateAndTimeMetrics, 1)
+    controllerEventManager.start()
+
+    controllerEventManager.put(TopicChange)
+    controllerEventManager.put(TopicChange)
+
+    TestUtils.waitUntilTrue(() => processedEvents.get() == 2,
+      "Timed out waiting for processing of all events")
+
+    val queueTimeHistogram = KafkaYammerMetrics.defaultRegistry.allMetrics.asScala.filter { case (k, _) =>
+      k.getMBeanName == metricName
+    }.values.headOption.getOrElse(fail(s"Unable to find metric $metricName")).asInstanceOf[Histogram]
+
+    TestUtils.waitUntilTrue(() => queueTimeHistogram.count == 0,
+      "Timed out on resetting queueTimeHistogram")
+    assertEquals(0, queueTimeHistogram.min, 0.1)
+    assertEquals(0, queueTimeHistogram.max, 0.1)
+  }
+
   @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(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,
@@ -79,8 +216,9 @@ class ControllerEventManagerTest {
   }
 
   private def timer(metricName: String): Timer = {
-    Metrics.defaultRegistry.allMetrics.asScala.filterKeys(_.getMBeanName == metricName).values.headOption
-      .getOrElse(fail(s"Unable to find metric $metricName")).asInstanceOf[Timer]
+    KafkaYammerMetrics.defaultRegistry.allMetrics.asScala.filter { case (k, _) =>
+      k.getMBeanName == metricName
+    }.values.headOption.getOrElse(fail(s"Unable to find metric $metricName")).asInstanceOf[Timer]
   }
 
 }
diff --git a/core/src/test/scala/unit/kafka/controller/ControllerFailoverTest.scala b/core/src/test/scala/unit/kafka/controller/ControllerFailoverTest.scala
index 32e23cc672c9a..f5af7c657e90d 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,34 +55,38 @@ 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")
     }
     val initialEpoch = initialController.epoch
     // Create topic with one partition
-    adminZkClient.createTopic(topic, 1, 1)
+    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()
-    })
+
+      override def preempt(): Unit = {}
+    }
     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 d917fe1fcbf9a..b186934f39053 100644
--- a/core/src/test/scala/unit/kafka/controller/ControllerIntegrationTest.scala
+++ b/core/src/test/scala/unit/kafka/controller/ControllerIntegrationTest.scala
@@ -17,97 +17,246 @@
 
 package kafka.controller
 
-import com.yammer.metrics.Metrics
+import java.util.Properties
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue}
+
 import com.yammer.metrics.core.Timer
-import kafka.api.LeaderAndIsr
-import kafka.common.TopicAndPartition
+import kafka.api.{ApiVersion, KAFKA_2_6_IV0, KAFKA_2_7_IV0, LeaderAndIsr}
+import kafka.metrics.KafkaYammerMetrics
 import kafka.server.{KafkaConfig, KafkaServer}
-import kafka.utils.{TestUtils, ZkUtils}
-import kafka.zk.ZooKeeperTestHarness
+import kafka.utils.{LogCaptureAppender, TestUtils}
+import kafka.zk.{FeatureZNodeStatus, _}
+import org.apache.kafka.common.errors.{ControllerMovedException, StaleBrokerEpochException}
+import org.apache.kafka.common.feature.Features
+import org.apache.kafka.common.metrics.KafkaMetric
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.{ElectionType, TopicPartition, Uuid}
+import org.apache.log4j.Level
+import org.junit.Assert.{assertEquals, assertNotEquals, assertTrue}
 import org.junit.{After, Before, Test}
-import org.junit.Assert.assertTrue
+import org.mockito.Mockito.{doAnswer, spy, verify}
+import org.mockito.invocation.InvocationOnMock
+import org.scalatest.Assertions.fail
 
-import scala.collection.JavaConverters._
+import scala.collection.{Map, Seq, mutable}
+import scala.jdk.CollectionConverters._
+import scala.util.{Failure, Success, Try}
 
 class ControllerIntegrationTest extends ZooKeeperTestHarness {
   var servers = Seq.empty[KafkaServer]
+  val firstControllerEpoch = KafkaController.InitialControllerEpoch + 1
+  val firstControllerEpochZkVersion = KafkaController.InitialControllerEpochZkVersion + 1
 
   @Before
-  override def setUp() {
-    super.setUp
+  override def setUp(): Unit = {
+    super.setUp()
     servers = Seq.empty[KafkaServer]
   }
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     TestUtils.shutdownServers(servers)
-    super.tearDown
+    super.tearDown()
   }
 
   @Test
   def testEmptyCluster(): Unit = {
     servers = makeServers(1)
-    TestUtils.waitUntilTrue(() => zkUtils.pathExists(ZkUtils.ControllerPath), "failed to elect a controller")
-    waitUntilControllerEpoch(KafkaController.InitialControllerEpoch, "broker failed to set controller epoch")
+    TestUtils.waitUntilTrue(() => zkClient.getControllerId.isDefined, "failed to elect a controller")
+    waitUntilControllerEpoch(firstControllerEpoch, "broker failed to set controller epoch")
   }
 
   @Test
   def testControllerEpochPersistsWhenAllBrokersDown(): Unit = {
     servers = makeServers(1)
-    TestUtils.waitUntilTrue(() => zkUtils.pathExists(ZkUtils.ControllerPath), "failed to elect a controller")
-    waitUntilControllerEpoch(KafkaController.InitialControllerEpoch, "broker failed to set controller epoch")
+    TestUtils.waitUntilTrue(() => zkClient.getControllerId.isDefined, "failed to elect a controller")
+    waitUntilControllerEpoch(firstControllerEpoch, "broker failed to set controller epoch")
     servers.head.shutdown()
     servers.head.awaitShutdown()
-    TestUtils.waitUntilTrue(() => !zkUtils.pathExists(ZkUtils.ControllerPath), "failed to kill controller")
-    waitUntilControllerEpoch(KafkaController.InitialControllerEpoch, "controller epoch was not persisted after broker failure")
+    TestUtils.waitUntilTrue(() => !zkClient.getControllerId.isDefined, "failed to kill controller")
+    waitUntilControllerEpoch(firstControllerEpoch, "controller epoch was not persisted after broker failure")
   }
 
   @Test
   def testControllerMoveIncrementsControllerEpoch(): Unit = {
     servers = makeServers(1)
-    TestUtils.waitUntilTrue(() => zkUtils.pathExists(ZkUtils.ControllerPath), "failed to elect a controller")
-    waitUntilControllerEpoch(KafkaController.InitialControllerEpoch, "broker failed to set controller epoch")
+    TestUtils.waitUntilTrue(() => zkClient.getControllerId.isDefined, "failed to elect a controller")
+    waitUntilControllerEpoch(firstControllerEpoch, "broker failed to set controller epoch")
     servers.head.shutdown()
     servers.head.awaitShutdown()
     servers.head.startup()
-    TestUtils.waitUntilTrue(() => zkUtils.pathExists(ZkUtils.ControllerPath), "failed to elect a controller")
-    waitUntilControllerEpoch(KafkaController.InitialControllerEpoch + 1, "controller epoch was not incremented after controller move")
+    TestUtils.waitUntilTrue(() => zkClient.getControllerId.isDefined, "failed to elect a controller")
+    waitUntilControllerEpoch(firstControllerEpoch + 1, "controller epoch was not incremented after controller move")
+  }
+
+  @Test
+  def testMetadataPropagationOnControlPlane(): Unit = {
+    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]()
+    val dataPlaneMetricMap = mutable.Map[String, KafkaMetric]()
+    servers.head.metrics.metrics.values.forEach { kafkaMetric =>
+      if (kafkaMetric.metricName.tags.values.contains("CONTROLLER")) {
+        controlPlaneMetricMap.put(kafkaMetric.metricName().name(), kafkaMetric)
+      }
+      if (kafkaMetric.metricName.tags.values.contains("PLAINTEXT")) {
+        dataPlaneMetricMap.put(kafkaMetric.metricName.name, kafkaMetric)
+      }
+    }
+    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
+  // UpdateMetadataRequest to all brokers in the cluster
+  @Test
+  def testMetadataPropagationOnBrokerChange(): Unit = {
+    servers = makeServers(3)
+    TestUtils.waitUntilBrokerMetadataIsPropagated(servers)
+    val controllerId = TestUtils.waitUntilControllerElected(zkClient)
+    // Need to make sure the broker we shutdown and startup are not the controller. Otherwise we will send out
+    // full UpdateMetadataReuqest to all brokers during controller failover.
+    val testBroker = servers.filter(e => e.config.brokerId != controllerId).head
+    val remainingBrokers = servers.filter(_.config.brokerId != testBroker.config.brokerId)
+    val topic = "topic1"
+    // Make sure shutdown the test broker will not require any leadership change to test avoid sending out full
+    // UpdateMetadataRequest on broker failure
+    val assignment = Map(
+      0 -> Seq(remainingBrokers(0).config.brokerId, testBroker.config.brokerId),
+      1 -> remainingBrokers.map(_.config.brokerId))
+
+    // Create topic
+    TestUtils.createTopic(zkClient, topic, assignment, servers)
+
+    // Shutdown the broker
+    testBroker.shutdown()
+    testBroker.awaitShutdown()
+    TestUtils.waitUntilBrokerMetadataIsPropagated(remainingBrokers)
+    remainingBrokers.foreach { server =>
+      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.replicas)
+      assertEquals(Seq(remainingBrokers.head.config.brokerId).asJava, offlineReplicaPartitionInfo.isr)
+      val onlinePartitionInfo = server.metadataCache.getPartitionInfo(topic, 1).get
+      assertEquals(assignment(1).asJava, onlinePartitionInfo.replicas)
+      assertTrue(onlinePartitionInfo.offlineReplicas.isEmpty)
+    }
+
+    // Startup the broker
+    testBroker.startup()
+    TestUtils.waitUntilTrue( () => {
+      !servers.exists { server =>
+        assignment.exists { case (partitionId, replicas) =>
+          val partitionInfoOpt = server.metadataCache.getPartitionInfo(topic, partitionId)
+          if (partitionInfoOpt.isDefined) {
+            val partitionInfo = partitionInfoOpt.get
+            !partitionInfo.offlineReplicas.isEmpty || !partitionInfo.replicas.asScala.equals(replicas)
+          } else {
+            true
+          }
+        }
+      }
+    }, "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)
-    val tp = TopicAndPartition("t", 0)
+    val tp = new TopicPartition("t", 0)
     val assignment = Map(tp.partition -> Seq(0))
-    TestUtils.createTopic(zkUtils, tp.topic, partitionReplicaAssignment = assignment, servers = servers)
-    waitForPartitionState(tp, KafkaController.InitialControllerEpoch, 0, LeaderAndIsr.initialLeaderEpoch,
+    TestUtils.createTopic(zkClient, tp.topic, partitionReplicaAssignment = assignment, servers = servers)
+    waitForPartitionState(tp, firstControllerEpoch, 0, LeaderAndIsr.initialLeaderEpoch,
       "failed to get expected partition state upon topic creation")
   }
 
   @Test
   def testTopicCreationWithOfflineReplica(): Unit = {
     servers = makeServers(2)
-    val controllerId = TestUtils.waitUntilControllerElected(zkUtils)
+    val controllerId = TestUtils.waitUntilControllerElected(zkClient)
     val otherBrokerId = servers.map(_.config.brokerId).filter(_ != controllerId).head
     servers(otherBrokerId).shutdown()
     servers(otherBrokerId).awaitShutdown()
-    val tp = TopicAndPartition("t", 0)
+    val tp = new TopicPartition("t", 0)
     val assignment = Map(tp.partition -> Seq(otherBrokerId, controllerId))
-    TestUtils.createTopic(zkUtils, tp.topic, partitionReplicaAssignment = assignment, servers = servers.take(1))
-    waitForPartitionState(tp, KafkaController.InitialControllerEpoch, controllerId, LeaderAndIsr.initialLeaderEpoch,
+    TestUtils.createTopic(zkClient, tp.topic, partitionReplicaAssignment = assignment, servers = servers.take(1))
+    waitForPartitionState(tp, firstControllerEpoch, controllerId, LeaderAndIsr.initialLeaderEpoch,
       "failed to get expected partition state upon topic creation")
   }
 
   @Test
   def testTopicPartitionExpansion(): Unit = {
     servers = makeServers(1)
-    val tp0 = TopicAndPartition("t", 0)
-    val tp1 = TopicAndPartition("t", 1)
+    val tp0 = new TopicPartition("t", 0)
+    val tp1 = new TopicPartition("t", 1)
     val assignment = Map(tp0.partition -> Seq(0))
-    val expandedAssignment = Map(tp0.partition -> Seq(0), tp1.partition -> Seq(0))
-    TestUtils.createTopic(zkUtils, tp0.topic, partitionReplicaAssignment = assignment, servers = servers)
-    zkUtils.updatePersistentPath(ZkUtils.getTopicPath(tp0.topic), zkUtils.replicaAssignmentZkData(expandedAssignment.map(kv => kv._1.toString -> kv._2)))
-    waitForPartitionState(tp1, KafkaController.InitialControllerEpoch, 0, LeaderAndIsr.initialLeaderEpoch,
+    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, Uuid.randomUuid(), expandedAssignment, firstControllerEpochZkVersion)
+    waitForPartitionState(tp1, firstControllerEpoch, 0, LeaderAndIsr.initialLeaderEpoch,
       "failed to get expected partition state upon topic partition expansion")
     TestUtils.waitUntilMetadataIsPropagated(servers, tp1.topic, tp1.partition)
   }
@@ -115,17 +264,19 @@ class ControllerIntegrationTest extends ZooKeeperTestHarness {
   @Test
   def testTopicPartitionExpansionWithOfflineReplica(): Unit = {
     servers = makeServers(2)
-    val controllerId = TestUtils.waitUntilControllerElected(zkUtils)
+    val controllerId = TestUtils.waitUntilControllerElected(zkClient)
     val otherBrokerId = servers.map(_.config.brokerId).filter(_ != controllerId).head
-    val tp0 = TopicAndPartition("t", 0)
-    val tp1 = TopicAndPartition("t", 1)
+    val tp0 = new TopicPartition("t", 0)
+    val tp1 = new TopicPartition("t", 1)
     val assignment = Map(tp0.partition -> Seq(otherBrokerId, controllerId))
-    val expandedAssignment = Map(tp0.partition -> Seq(otherBrokerId, controllerId), tp1.partition -> Seq(otherBrokerId, controllerId))
-    TestUtils.createTopic(zkUtils, tp0.topic, partitionReplicaAssignment = assignment, servers = servers)
+    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()
-    zkUtils.updatePersistentPath(ZkUtils.getTopicPath(tp0.topic), zkUtils.replicaAssignmentZkData(expandedAssignment.map(kv => kv._1.toString -> kv._2)))
-    waitForPartitionState(tp1, KafkaController.InitialControllerEpoch, controllerId, LeaderAndIsr.initialLeaderEpoch,
+    zkClient.setTopicAssignment(tp0.topic, Uuid.randomUuid(), expandedAssignment, firstControllerEpochZkVersion)
+    waitForPartitionState(tp1, firstControllerEpoch, controllerId, LeaderAndIsr.initialLeaderEpoch,
       "failed to get expected partition state upon topic partition expansion")
     TestUtils.waitUntilMetadataIsPropagated(Seq(servers(controllerId)), tp1.topic, tp1.partition)
   }
@@ -133,22 +284,60 @@ class ControllerIntegrationTest extends ZooKeeperTestHarness {
   @Test
   def testPartitionReassignment(): Unit = {
     servers = makeServers(2)
-    val controllerId = TestUtils.waitUntilControllerElected(zkUtils)
+    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 = TopicAndPartition("t", 0)
+    val tp = new TopicPartition("t", 0)
     val assignment = Map(tp.partition -> Seq(controllerId))
-    val reassignment = Map(tp -> Seq(otherBrokerId))
-    TestUtils.createTopic(zkUtils, tp.topic, partitionReplicaAssignment = assignment, servers = servers)
-    zkUtils.createPersistentPath(ZkUtils.ReassignPartitionsPath, ZkUtils.formatAsReassignmentJson(reassignment))
-    waitForPartitionState(tp, KafkaController.InitialControllerEpoch, otherBrokerId, LeaderAndIsr.initialLeaderEpoch + 3,
+    val reassignment = Map(tp -> ReplicaAssignment(Seq(otherBrokerId), List(), List()))
+    TestUtils.createTopic(zkClient, tp.topic, partitionReplicaAssignment = assignment, servers = servers)
+    zkClient.createPartitionReassignment(reassignment.map { case (k, v) => k -> v.replicas })
+    waitForPartitionState(tp, firstControllerEpoch, otherBrokerId, LeaderAndIsr.initialLeaderEpoch + 3,
       "failed to get expected partition state after partition reassignment")
-    TestUtils.waitUntilTrue(() => zkUtils.getReplicaAssignmentForTopics(Seq(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")
+
+    val updatedTimerCount = timer(metricName).count
+    assertTrue(s"Timer count $updatedTimerCount should be greater than $timerCount", updatedTimerCount > timerCount)
+  }
+
+  @Test
+  def testPartitionReassignmentToBrokerWithOfflineLogDir(): Unit = {
+    servers = makeServers(2, logDirCount = 2)
+    val controllerId = TestUtils.waitUntilControllerElected(zkClient)
+
+    val metricName = s"kafka.controller:type=ControllerStats,name=${ControllerState.AlterPartitionReassignment.rateAndTimeMetricName.get}"
+    val timerCount = timer(metricName).count
+
+    val otherBroker = servers.filter(_.config.brokerId != controllerId).head
+    val otherBrokerId = otherBroker.config.brokerId
+
+    // To have an offline log dir, we need a topicPartition assigned to it
+    val topicPartitionToPutOffline = new TopicPartition("filler", 0)
+    TestUtils.createTopic(
+      zkClient,
+      topicPartitionToPutOffline.topic,
+      partitionReplicaAssignment = Map(topicPartitionToPutOffline.partition -> Seq(otherBrokerId)),
+      servers = servers
+    )
+
+    TestUtils.causeLogDirFailure(TestUtils.Checkpoint, otherBroker, topicPartitionToPutOffline)
+
+    val tp = new TopicPartition("t", 0)
+    val assignment = Map(tp.partition -> Seq(controllerId))
+    val reassignment = Map(tp -> ReplicaAssignment(Seq(otherBrokerId), List(), List()))
+    TestUtils.createTopic(zkClient, tp.topic, partitionReplicaAssignment = assignment, servers = servers)
+    zkClient.createPartitionReassignment(reassignment.map { case (k, v) => k -> v.replicas })
+    waitForPartitionState(tp, firstControllerEpoch, otherBrokerId, LeaderAndIsr.initialLeaderEpoch + 3,
+      "with an offline log directory on the target broker, the partition reassignment stalls")
+    TestUtils.waitUntilTrue(() =>  zkClient.getFullReplicaAssignmentForTopics(Set(tp.topic)) == reassignment,
       "failed to get updated partition assignment on topic znode after partition reassignment")
-    TestUtils.waitUntilTrue(() => !zkUtils.pathExists(ZkUtils.ReassignPartitionsPath),
+    TestUtils.waitUntilTrue(() => !zkClient.reassignPartitionsInProgress,
       "failed to remove reassign partitions path after completion")
 
     val updatedTimerCount = timer(metricName).count
@@ -158,63 +347,64 @@ class ControllerIntegrationTest extends ZooKeeperTestHarness {
   @Test
   def testPartitionReassignmentWithOfflineReplicaHaltingProgress(): Unit = {
     servers = makeServers(2)
-    val controllerId = TestUtils.waitUntilControllerElected(zkUtils)
+    val controllerId = TestUtils.waitUntilControllerElected(zkClient)
     val otherBrokerId = servers.map(_.config.brokerId).filter(_ != controllerId).head
-    val tp = TopicAndPartition("t", 0)
+    val tp = new TopicPartition("t", 0)
     val assignment = Map(tp.partition -> Seq(controllerId))
     val reassignment = Map(tp -> Seq(otherBrokerId))
-    TestUtils.createTopic(zkUtils, tp.topic, partitionReplicaAssignment = assignment, servers = servers)
+    TestUtils.createTopic(zkClient, tp.topic, partitionReplicaAssignment = assignment, servers = servers)
     servers(otherBrokerId).shutdown()
     servers(otherBrokerId).awaitShutdown()
-    zkUtils.createPersistentPath(ZkUtils.ReassignPartitionsPath, ZkUtils.formatAsReassignmentJson(reassignment))
-    waitForPartitionState(tp, KafkaController.InitialControllerEpoch, controllerId, LeaderAndIsr.initialLeaderEpoch + 1,
+    val controller = getController()
+    zkClient.setOrCreatePartitionReassignment(reassignment, controller.kafkaController.controllerContext.epochZkVersion)
+    waitForPartitionState(tp, firstControllerEpoch, controllerId, LeaderAndIsr.initialLeaderEpoch + 1,
       "failed to get expected partition state during partition reassignment with offline replica")
-    TestUtils.waitUntilTrue(() => zkUtils.pathExists(ZkUtils.ReassignPartitionsPath),
+    TestUtils.waitUntilTrue(() => zkClient.reassignPartitionsInProgress,
       "partition reassignment path should remain while reassignment in progress")
   }
 
   @Test
   def testPartitionReassignmentResumesAfterReplicaComesOnline(): Unit = {
     servers = makeServers(2)
-    val controllerId = TestUtils.waitUntilControllerElected(zkUtils)
+    val controllerId = TestUtils.waitUntilControllerElected(zkClient)
     val otherBrokerId = servers.map(_.config.brokerId).filter(_ != controllerId).head
-    val tp = TopicAndPartition("t", 0)
+    val tp = new TopicPartition("t", 0)
     val assignment = Map(tp.partition -> Seq(controllerId))
-    val reassignment = Map(tp -> Seq(otherBrokerId))
-    TestUtils.createTopic(zkUtils, tp.topic, partitionReplicaAssignment = assignment, servers = servers)
+    val reassignment = Map(tp -> ReplicaAssignment(Seq(otherBrokerId), List(), List()))
+    TestUtils.createTopic(zkClient, tp.topic, partitionReplicaAssignment = assignment, servers = servers)
     servers(otherBrokerId).shutdown()
     servers(otherBrokerId).awaitShutdown()
-    zkUtils.createPersistentPath(ZkUtils.ReassignPartitionsPath, ZkUtils.formatAsReassignmentJson(reassignment))
-    waitForPartitionState(tp, KafkaController.InitialControllerEpoch, controllerId, LeaderAndIsr.initialLeaderEpoch + 1,
+    zkClient.createPartitionReassignment(reassignment.map { case (k, v) => k -> v.replicas })
+    waitForPartitionState(tp, firstControllerEpoch, controllerId, LeaderAndIsr.initialLeaderEpoch + 1,
       "failed to get expected partition state during partition reassignment with offline replica")
     servers(otherBrokerId).startup()
-    waitForPartitionState(tp, KafkaController.InitialControllerEpoch, otherBrokerId, LeaderAndIsr.initialLeaderEpoch + 4,
+    waitForPartitionState(tp, firstControllerEpoch, otherBrokerId, LeaderAndIsr.initialLeaderEpoch + 4,
       "failed to get expected partition state after partition reassignment")
-    TestUtils.waitUntilTrue(() => zkUtils.getReplicaAssignmentForTopics(Seq(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(() => !zkUtils.pathExists(ZkUtils.ReassignPartitionsPath),
+    TestUtils.waitUntilTrue(() => !zkClient.reassignPartitionsInProgress,
       "failed to remove reassign partitions path after completion")
   }
 
   @Test
   def testPreferredReplicaLeaderElection(): Unit = {
     servers = makeServers(2)
-    val controllerId = TestUtils.waitUntilControllerElected(zkUtils)
+    val controllerId = TestUtils.waitUntilControllerElected(zkClient)
     val otherBroker = servers.find(_.config.brokerId != controllerId).get
-    val tp = TopicAndPartition("t", 0)
+    val tp = new TopicPartition("t", 0)
     val assignment = Map(tp.partition -> Seq(otherBroker.config.brokerId, controllerId))
-    TestUtils.createTopic(zkUtils, tp.topic, partitionReplicaAssignment = assignment, servers = servers)
+    TestUtils.createTopic(zkClient, tp.topic, partitionReplicaAssignment = assignment, servers = servers)
     preferredReplicaLeaderElection(controllerId, otherBroker, tp, assignment(tp.partition).toSet, LeaderAndIsr.initialLeaderEpoch)
   }
 
   @Test
   def testBackToBackPreferredReplicaLeaderElections(): Unit = {
     servers = makeServers(2)
-    val controllerId = TestUtils.waitUntilControllerElected(zkUtils)
+    val controllerId = TestUtils.waitUntilControllerElected(zkClient)
     val otherBroker = servers.find(_.config.brokerId != controllerId).get
-    val tp = TopicAndPartition("t", 0)
+    val tp = new TopicPartition("t", 0)
     val assignment = Map(tp.partition -> Seq(otherBroker.config.brokerId, controllerId))
-    TestUtils.createTopic(zkUtils, tp.topic, partitionReplicaAssignment = assignment, servers = servers)
+    TestUtils.createTopic(zkClient, tp.topic, partitionReplicaAssignment = assignment, servers = servers)
     preferredReplicaLeaderElection(controllerId, otherBroker, tp, assignment(tp.partition).toSet, LeaderAndIsr.initialLeaderEpoch)
     preferredReplicaLeaderElection(controllerId, otherBroker, tp, assignment(tp.partition).toSet, LeaderAndIsr.initialLeaderEpoch + 2)
   }
@@ -222,53 +412,53 @@ class ControllerIntegrationTest extends ZooKeeperTestHarness {
   @Test
   def testPreferredReplicaLeaderElectionWithOfflinePreferredReplica(): Unit = {
     servers = makeServers(2)
-    val controllerId = TestUtils.waitUntilControllerElected(zkUtils)
+    val controllerId = TestUtils.waitUntilControllerElected(zkClient)
     val otherBrokerId = servers.map(_.config.brokerId).filter(_ != controllerId).head
-    val tp = TopicAndPartition("t", 0)
+    val tp = new TopicPartition("t", 0)
     val assignment = Map(tp.partition -> Seq(otherBrokerId, controllerId))
-    TestUtils.createTopic(zkUtils, tp.topic, partitionReplicaAssignment = assignment, servers = servers)
+    TestUtils.createTopic(zkClient, tp.topic, partitionReplicaAssignment = assignment, servers = servers)
     servers(otherBrokerId).shutdown()
     servers(otherBrokerId).awaitShutdown()
-    zkUtils.createPersistentPath(ZkUtils.PreferredReplicaLeaderElectionPath, ZkUtils.preferredReplicaLeaderElectionZkData(Set(tp)))
-    TestUtils.waitUntilTrue(() => !zkUtils.pathExists(ZkUtils.PreferredReplicaLeaderElectionPath),
+    zkClient.createPreferredReplicaElection(Set(tp))
+    TestUtils.waitUntilTrue(() => !zkClient.pathExists(PreferredReplicaElectionZNode.path),
       "failed to remove preferred replica leader election path after giving up")
-    waitForPartitionState(tp, KafkaController.InitialControllerEpoch, controllerId, LeaderAndIsr.initialLeaderEpoch + 1,
+    waitForPartitionState(tp, firstControllerEpoch, controllerId, LeaderAndIsr.initialLeaderEpoch + 1,
       "failed to get expected partition state upon broker shutdown")
   }
 
   @Test
   def testAutoPreferredReplicaLeaderElection(): Unit = {
     servers = makeServers(2, autoLeaderRebalanceEnable = true)
-    val controllerId = TestUtils.waitUntilControllerElected(zkUtils)
+    val controllerId = TestUtils.waitUntilControllerElected(zkClient)
     val otherBrokerId = servers.map(_.config.brokerId).filter(_ != controllerId).head
-    val tp = TopicAndPartition("t", 0)
+    val tp = new TopicPartition("t", 0)
     val assignment = Map(tp.partition -> Seq(1, 0))
-    TestUtils.createTopic(zkUtils, tp.topic, partitionReplicaAssignment = assignment, servers = servers)
+    TestUtils.createTopic(zkClient, tp.topic, partitionReplicaAssignment = assignment, servers = servers)
     servers(otherBrokerId).shutdown()
     servers(otherBrokerId).awaitShutdown()
-    waitForPartitionState(tp, KafkaController.InitialControllerEpoch, controllerId, LeaderAndIsr.initialLeaderEpoch + 1,
+    waitForPartitionState(tp, firstControllerEpoch, controllerId, LeaderAndIsr.initialLeaderEpoch + 1,
       "failed to get expected partition state upon broker shutdown")
     servers(otherBrokerId).startup()
-    waitForPartitionState(tp, KafkaController.InitialControllerEpoch, otherBrokerId, LeaderAndIsr.initialLeaderEpoch + 2,
+    waitForPartitionState(tp, firstControllerEpoch, otherBrokerId, LeaderAndIsr.initialLeaderEpoch + 2,
       "failed to get expected partition state upon broker startup")
   }
 
   @Test
   def testLeaderAndIsrWhenEntireIsrOfflineAndUncleanLeaderElectionDisabled(): Unit = {
     servers = makeServers(2)
-    val controllerId = TestUtils.waitUntilControllerElected(zkUtils)
+    val controllerId = TestUtils.waitUntilControllerElected(zkClient)
     val otherBrokerId = servers.map(_.config.brokerId).filter(_ != controllerId).head
-    val tp = TopicAndPartition("t", 0)
+    val tp = new TopicPartition("t", 0)
     val assignment = Map(tp.partition -> Seq(otherBrokerId))
-    TestUtils.createTopic(zkUtils, tp.topic, partitionReplicaAssignment = assignment, servers = servers)
-    waitForPartitionState(tp, KafkaController.InitialControllerEpoch, otherBrokerId, LeaderAndIsr.initialLeaderEpoch,
+    TestUtils.createTopic(zkClient, tp.topic, partitionReplicaAssignment = assignment, servers = servers)
+    waitForPartitionState(tp, firstControllerEpoch, otherBrokerId, LeaderAndIsr.initialLeaderEpoch,
       "failed to get expected partition state upon topic creation")
     servers(otherBrokerId).shutdown()
     servers(otherBrokerId).awaitShutdown()
     TestUtils.waitUntilTrue(() => {
-      val leaderIsrAndControllerEpochMap = zkUtils.getPartitionLeaderAndIsrForTopics(Set(tp))
+      val leaderIsrAndControllerEpochMap = zkClient.getTopicPartitionStates(Seq(tp))
       leaderIsrAndControllerEpochMap.contains(tp) &&
-        isExpectedPartitionState(leaderIsrAndControllerEpochMap(tp), KafkaController.InitialControllerEpoch, LeaderAndIsr.NoLeader, LeaderAndIsr.initialLeaderEpoch + 1) &&
+        isExpectedPartitionState(leaderIsrAndControllerEpochMap(tp), firstControllerEpoch, LeaderAndIsr.NoLeader, LeaderAndIsr.initialLeaderEpoch + 1) &&
         leaderIsrAndControllerEpochMap(tp).leaderAndIsr.isr == List(otherBrokerId)
     }, "failed to get expected partition state after entire isr went offline")
   }
@@ -276,49 +466,558 @@ class ControllerIntegrationTest extends ZooKeeperTestHarness {
   @Test
   def testLeaderAndIsrWhenEntireIsrOfflineAndUncleanLeaderElectionEnabled(): Unit = {
     servers = makeServers(2, uncleanLeaderElectionEnable = true)
-    val controllerId = TestUtils.waitUntilControllerElected(zkUtils)
+    val controllerId = TestUtils.waitUntilControllerElected(zkClient)
     val otherBrokerId = servers.map(_.config.brokerId).filter(_ != controllerId).head
-    val tp = TopicAndPartition("t", 0)
+    val tp = new TopicPartition("t", 0)
     val assignment = Map(tp.partition -> Seq(otherBrokerId))
-    TestUtils.createTopic(zkUtils, tp.topic, partitionReplicaAssignment = assignment, servers = servers)
-    waitForPartitionState(tp, KafkaController.InitialControllerEpoch, otherBrokerId, LeaderAndIsr.initialLeaderEpoch,
+    TestUtils.createTopic(zkClient, tp.topic, partitionReplicaAssignment = assignment, servers = servers)
+    waitForPartitionState(tp, firstControllerEpoch, otherBrokerId, LeaderAndIsr.initialLeaderEpoch,
       "failed to get expected partition state upon topic creation")
-    servers(1).shutdown()
-    servers(1).awaitShutdown()
+    servers(otherBrokerId).shutdown()
+    servers(otherBrokerId).awaitShutdown()
     TestUtils.waitUntilTrue(() => {
-      val leaderIsrAndControllerEpochMap = zkUtils.getPartitionLeaderAndIsrForTopics(Set(tp))
+      val leaderIsrAndControllerEpochMap = zkClient.getTopicPartitionStates(Seq(tp))
       leaderIsrAndControllerEpochMap.contains(tp) &&
-        isExpectedPartitionState(leaderIsrAndControllerEpochMap(tp), KafkaController.InitialControllerEpoch, LeaderAndIsr.NoLeader, LeaderAndIsr.initialLeaderEpoch + 1) &&
+        isExpectedPartitionState(leaderIsrAndControllerEpochMap(tp), firstControllerEpoch, LeaderAndIsr.NoLeader, LeaderAndIsr.initialLeaderEpoch + 1) &&
         leaderIsrAndControllerEpochMap(tp).leaderAndIsr.isr == List(otherBrokerId)
     }, "failed to get expected partition state after entire isr went offline")
   }
 
-  private def preferredReplicaLeaderElection(controllerId: Int, otherBroker: KafkaServer, tp: TopicAndPartition,
+  @Test
+  def testControlledShutdown(): Unit = {
+    val expectedReplicaAssignment = Map(0  -> List(0, 1, 2))
+    val topic = "test"
+    val partition = 0
+    // create brokers
+    val serverConfigs = TestUtils.createBrokerConfigs(3, zkConnect, false).map(KafkaConfig.fromProps)
+    servers = serverConfigs.reverse.map(s => TestUtils.createServer(s))
+    // create the topic
+    TestUtils.createTopic(zkClient, topic, partitionReplicaAssignment = expectedReplicaAssignment, servers = servers)
+
+    val controllerId = zkClient.getControllerId.get
+    val controller = servers.find(p => p.config.brokerId == controllerId).get.kafkaController
+    val resultQueue = new LinkedBlockingQueue[Try[collection.Set[TopicPartition]]]()
+    val controlledShutdownCallback = (controlledShutdownResult: Try[collection.Set[TopicPartition]]) => resultQueue.put(controlledShutdownResult)
+    controller.controlledShutdown(2, servers.find(_.config.brokerId == 2).get.kafkaController.brokerEpoch, controlledShutdownCallback)
+    var partitionsRemaining = resultQueue.take().get
+    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.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.leader
+    assertEquals(0, leaderAfterShutdown)
+    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() 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.leader
+    assertEquals(0, leaderAfterShutdown)
+
+    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.leader == 0))
+  }
+
+  @Test
+  def testControllerRejectControlledShutdownRequestWithStaleBrokerEpoch(): Unit = {
+    // create brokers
+    val serverConfigs = TestUtils.createBrokerConfigs(2, zkConnect, false).map(KafkaConfig.fromProps)
+    servers = serverConfigs.reverse.map(s => TestUtils.createServer(s))
+
+    val controller = getController().kafkaController
+    val otherBroker = servers.find(e => e.config.brokerId != controller.config.brokerId).get
+    @volatile var staleBrokerEpochDetected = false
+    controller.controlledShutdown(otherBroker.config.brokerId, otherBroker.kafkaController.brokerEpoch - 1, {
+      case scala.util.Failure(exception) if exception.isInstanceOf[StaleBrokerEpochException] => staleBrokerEpochDetected = true
+      case _ =>
+    })
+
+    TestUtils.waitUntilTrue(() => staleBrokerEpochDetected, "Fail to detect stale broker epoch")
+  }
+
+  @Test
+  def testControllerMoveOnTopicCreation(): Unit = {
+    servers = makeServers(1)
+    TestUtils.waitUntilControllerElected(zkClient)
+    val tp = new TopicPartition("t", 0)
+    val assignment = Map(tp.partition -> Seq(0))
+
+    testControllerMove(() => {
+      val adminZkClient = new AdminZkClient(zkClient)
+      adminZkClient.createTopicWithAssignment(tp.topic, config = new Properties(), assignment)
+    })
+  }
+
+  @Test
+  def testControllerMoveOnTopicDeletion(): Unit = {
+    servers = makeServers(1)
+    TestUtils.waitUntilControllerElected(zkClient)
+    val tp = new TopicPartition("t", 0)
+    val assignment = Map(tp.partition -> Seq(0))
+    TestUtils.createTopic(zkClient, tp.topic(), assignment, servers)
+
+    testControllerMove(() => {
+      val adminZkClient = new AdminZkClient(zkClient)
+      adminZkClient.deleteTopic(tp.topic())
+    })
+  }
+
+  @Test
+  def testControllerMoveOnPreferredReplicaElection(): Unit = {
+    servers = makeServers(1)
+    val tp = new TopicPartition("t", 0)
+    val assignment = Map(tp.partition -> Seq(0))
+    TestUtils.createTopic(zkClient, tp.topic(), assignment, servers)
+
+    testControllerMove(() => zkClient.createPreferredReplicaElection(Set(tp)))
+  }
+
+  @Test
+  def testControllerMoveOnPartitionReassignment(): Unit = {
+    servers = makeServers(1)
+    TestUtils.waitUntilControllerElected(zkClient)
+    val tp = new TopicPartition("t", 0)
+    val assignment = Map(tp.partition -> Seq(0))
+    TestUtils.createTopic(zkClient, tp.topic(), assignment, servers)
+
+    val reassignment = Map(tp -> Seq(0))
+    testControllerMove(() => zkClient.createPartitionReassignment(reassignment))
+  }
+
+  @Test
+  def testControllerFeatureZNodeSetupWhenFeatureVersioningIsEnabledWithNonExistingFeatureZNode(): Unit = {
+    testControllerFeatureZNodeSetup(Option.empty, KAFKA_2_7_IV0)
+  }
+
+  @Test
+  def testControllerFeatureZNodeSetupWhenFeatureVersioningIsEnabledWithDisabledExistingFeatureZNode(): Unit = {
+    testControllerFeatureZNodeSetup(Some(new FeatureZNode(FeatureZNodeStatus.Disabled, Features.emptyFinalizedFeatures())), KAFKA_2_7_IV0)
+  }
+
+  @Test
+  def testControllerFeatureZNodeSetupWhenFeatureVersioningIsEnabledWithEnabledExistingFeatureZNode(): Unit = {
+    testControllerFeatureZNodeSetup(Some(new FeatureZNode(FeatureZNodeStatus.Enabled, Features.emptyFinalizedFeatures())), KAFKA_2_7_IV0)
+  }
+
+  @Test
+  def testControllerFeatureZNodeSetupWhenFeatureVersioningIsDisabledWithNonExistingFeatureZNode(): Unit = {
+    testControllerFeatureZNodeSetup(Option.empty, KAFKA_2_6_IV0)
+  }
+
+  @Test
+  def testControllerFeatureZNodeSetupWhenFeatureVersioningIsDisabledWithDisabledExistingFeatureZNode(): Unit = {
+    testControllerFeatureZNodeSetup(Some(new FeatureZNode(FeatureZNodeStatus.Disabled, Features.emptyFinalizedFeatures())), KAFKA_2_6_IV0)
+  }
+
+  @Test
+  def testControllerFeatureZNodeSetupWhenFeatureVersioningIsDisabledWithEnabledExistingFeatureZNode(): Unit = {
+    testControllerFeatureZNodeSetup(Some(new FeatureZNode(FeatureZNodeStatus.Enabled, Features.emptyFinalizedFeatures())), KAFKA_2_6_IV0)
+  }
+
+  @Test
+  def testControllerDetectsBouncedBrokers(): Unit = {
+    servers = makeServers(2, enableControlledShutdown = false)
+    val controller = getController().kafkaController
+    val otherBroker = servers.find(e => e.config.brokerId != controller.config.brokerId).get
+
+    // Create a topic
+    val tp = new TopicPartition("t", 0)
+    val assignment = Map(tp.partition -> Seq(0, 1))
+
+    TestUtils.createTopic(zkClient, tp.topic, partitionReplicaAssignment = assignment, servers = servers)
+    waitForPartitionState(tp, firstControllerEpoch, 0, LeaderAndIsr.initialLeaderEpoch,
+      "failed to get expected partition state upon topic creation")
+
+    // Wait until the event thread is idle
+    TestUtils.waitUntilTrue(() => {
+      controller.eventManager.state == ControllerState.Idle
+    }, "Controller event thread is still busy")
+
+    val latch = new CountDownLatch(1)
+
+    // Let the controller event thread await on a latch until broker bounce finishes.
+    // This is used to simulate fast broker bounce
+
+    controller.eventManager.put(new MockEvent(ControllerState.TopicChange) {
+      override def process(): Unit = latch.await()
+      override def preempt(): Unit = {}
+    })
+
+    otherBroker.shutdown()
+    otherBroker.startup()
+
+    assertEquals(0, otherBroker.replicaManager.partitionCount.value())
+
+    // Release the latch so that controller can process broker change event
+    latch.countDown()
+    TestUtils.waitUntilTrue(() => {
+      otherBroker.replicaManager.partitionCount.value() == 1 &&
+      otherBroker.replicaManager.metadataCache.getAllTopics().size == 1 &&
+      otherBroker.replicaManager.metadataCache.getAliveBrokers.size == 2
+    }, "Broker fail to initialize after restart")
+  }
+
+  @Test
+  def testPreemptionOnControllerShutdown(): Unit = {
+    servers = makeServers(1, enableControlledShutdown = false)
+    val controller = getController().kafkaController
+    var count = 2
+    val latch = new CountDownLatch(1)
+    val spyThread = spy(controller.eventManager.thread)
+    controller.eventManager.thread = spyThread
+    val processedEvent = new MockEvent(ControllerState.TopicChange) {
+      override def process(): Unit = latch.await()
+      override def preempt(): Unit = {}
+    }
+    val preemptedEvent = new MockEvent(ControllerState.TopicChange) {
+      override def process(): Unit = {}
+      override def preempt(): Unit = count -= 1
+    }
+
+    controller.eventManager.put(processedEvent)
+    controller.eventManager.put(preemptedEvent)
+    controller.eventManager.put(preemptedEvent)
+
+    doAnswer((_: InvocationOnMock) => {
+      latch.countDown()
+    }).doCallRealMethod().when(spyThread).awaitShutdown()
+    controller.shutdown()
+    TestUtils.waitUntilTrue(() => {
+      count == 0
+    }, "preemption was not fully completed before shutdown")
+
+    verify(spyThread).awaitShutdown()
+  }
+
+  @Test
+  def testPreemptionWithCallbacks(): Unit = {
+    servers = makeServers(1, enableControlledShutdown = false)
+    val controller = getController().kafkaController
+    val latch = new CountDownLatch(1)
+    val spyThread = spy(controller.eventManager.thread)
+    controller.eventManager.thread = spyThread
+    val processedEvent = new MockEvent(ControllerState.TopicChange) {
+      override def process(): Unit = latch.await()
+      override def preempt(): Unit = {}
+    }
+    val tp0 = new TopicPartition("t", 0)
+    val tp1 = new TopicPartition("t", 1)
+    val partitions = Set(tp0, tp1)
+    val event1 = ReplicaLeaderElection(Some(partitions), ElectionType.PREFERRED, ZkTriggered, partitionsMap => {
+      for (partition <- partitionsMap) {
+        partition._2 match {
+          case Left(e) => assertEquals(Errors.NOT_CONTROLLER, e.error())
+          case Right(_) => fail("replica leader election should error")
+        }
+      }
+    })
+    val event2 = ControlledShutdown(0, 0, {
+      case Success(_) => fail("controlled shutdown should error")
+      case Failure(e) =>
+        assertEquals(classOf[ControllerMovedException], e.getClass)
+    })
+    val event3  = ApiPartitionReassignment(Map(tp0 -> None, tp1 -> None), {
+      case Left(_) => fail("api partition reassignment should error")
+      case Right(e) => assertEquals(Errors.NOT_CONTROLLER, e.error())
+    })
+    val event4 = ListPartitionReassignments(Some(partitions), {
+      case Left(_) => fail("api partition reassignment should error")
+      case Right(e) => assertEquals(Errors.NOT_CONTROLLER, e.error())
+    })
+
+    controller.eventManager.put(processedEvent)
+    controller.eventManager.put(event1)
+    controller.eventManager.put(event2)
+    controller.eventManager.put(event3)
+    controller.eventManager.put(event4)
+
+    doAnswer((_: InvocationOnMock) => {
+      latch.countDown()
+    }).doCallRealMethod().when(spyThread).awaitShutdown()
+    controller.shutdown()
+  }
+
+  private def testControllerFeatureZNodeSetup(initialZNode: Option[FeatureZNode],
+                                              interBrokerProtocolVersion: ApiVersion): Unit = {
+    val versionBeforeOpt = initialZNode match {
+      case Some(node) =>
+        zkClient.createFeatureZNode(node)
+        Some(zkClient.getDataAndVersion(FeatureZNode.path)._2)
+      case None =>
+        Option.empty
+    }
+    servers = makeServers(1, interBrokerProtocolVersion = Some(interBrokerProtocolVersion))
+    TestUtils.waitUntilControllerElected(zkClient)
+    // Below we wait on a dummy event to finish processing in the controller event thread.
+    // We schedule this dummy event only after the controller is elected, which is a sign that the
+    // controller has already started processing the Startup event. Waiting on the dummy event is
+    // used to make sure that the event thread has completed processing Startup event, that triggers
+    // the setup of FeatureZNode.
+    val controller = getController().kafkaController
+    val latch = new CountDownLatch(1)
+    controller.eventManager.put(new MockEvent(ControllerState.TopicChange) {
+      override def process(): Unit = {
+        latch.countDown()
+      }
+      override def preempt(): Unit = {}
+    })
+    latch.await()
+
+    val (mayBeFeatureZNodeBytes, versionAfter) = zkClient.getDataAndVersion(FeatureZNode.path)
+    val newZNode = FeatureZNode.decode(mayBeFeatureZNodeBytes.get)
+    if (interBrokerProtocolVersion >= KAFKA_2_7_IV0) {
+      val emptyZNode = new FeatureZNode(FeatureZNodeStatus.Enabled, Features.emptyFinalizedFeatures)
+      initialZNode match {
+        case Some(node) => {
+          node.status match {
+            case FeatureZNodeStatus.Enabled =>
+              assertEquals(versionBeforeOpt.get, versionAfter)
+              assertEquals(node, newZNode)
+            case FeatureZNodeStatus.Disabled =>
+              assertEquals(versionBeforeOpt.get + 1, versionAfter)
+              assertEquals(emptyZNode, newZNode)
+          }
+        }
+        case None =>
+          assertEquals(0, versionAfter)
+          assertEquals(new FeatureZNode(FeatureZNodeStatus.Enabled, Features.emptyFinalizedFeatures), newZNode)
+      }
+    } else {
+      val emptyZNode = new FeatureZNode(FeatureZNodeStatus.Disabled, Features.emptyFinalizedFeatures)
+      initialZNode match {
+        case Some(node) => {
+          node.status match {
+            case FeatureZNodeStatus.Enabled =>
+              assertEquals(versionBeforeOpt.get + 1, versionAfter)
+              assertEquals(emptyZNode, newZNode)
+            case FeatureZNodeStatus.Disabled =>
+              assertEquals(versionBeforeOpt.get, versionAfter)
+              assertEquals(emptyZNode, newZNode)
+          }
+        }
+        case None =>
+          assertEquals(0, versionAfter)
+          assertEquals(new FeatureZNode(FeatureZNodeStatus.Disabled, Features.emptyFinalizedFeatures), newZNode)
+      }
+    }
+  }
+
+  @Test
+  def testIdempotentAlterIsr(): Unit = {
+    servers = makeServers(2)
+    val controllerId = TestUtils.waitUntilControllerElected(zkClient)
+    val otherBroker = servers.find(_.config.brokerId != controllerId).get
+    val tp = new TopicPartition("t", 0)
+    val assignment = Map(tp.partition -> Seq(otherBroker.config.brokerId, controllerId))
+    TestUtils.createTopic(zkClient, tp.topic, partitionReplicaAssignment = assignment, servers = servers)
+
+    val latch = new CountDownLatch(1)
+    val controller = getController().kafkaController
+
+    val leaderIsrAndControllerEpochMap = zkClient.getTopicPartitionStates(Seq(tp))
+    val newLeaderAndIsr = leaderIsrAndControllerEpochMap(tp).leaderAndIsr
+
+    val callback = (result: Either[Map[TopicPartition, Either[Errors, LeaderAndIsr]], Errors]) => {
+      result match {
+        case Left(partitionResults: Map[TopicPartition, Either[Errors, LeaderAndIsr]]) =>
+          partitionResults.get(tp) match {
+            case Some(Left(error: Errors)) => fail(s"Should not have seen error for $tp")
+            case Some(Right(leaderAndIsr: LeaderAndIsr)) => assertEquals("ISR should remain unchanged", leaderAndIsr, newLeaderAndIsr)
+            case None => fail(s"Should have seen $tp in result")
+          }
+        case Right(_: Errors) => fail("Should not have had top-level error here")
+      }
+      latch.countDown()
+    }
+
+    val brokerEpoch = controller.controllerContext.liveBrokerIdAndEpochs.get(otherBroker.config.brokerId).get
+    // When re-sending the current ISR, we should not get and error or any ISR changes
+    controller.eventManager.put(AlterIsrReceived(otherBroker.config.brokerId, brokerEpoch, Map(tp -> newLeaderAndIsr), callback))
+    latch.await()
+  }
+
+  @Test
+  def testTopicIdsAreAdded(): Unit = {
+    servers = makeServers(1)
+    TestUtils.waitUntilControllerElected(zkClient)
+    val controller = getController().kafkaController
+    val tp1 = new TopicPartition("t1", 0)
+    val assignment1 = Map(tp1.partition -> Seq(0))
+
+    // Before adding the topic, an attempt to get the ID should result in None.
+    assertTrue(controller.controllerContext.topicIds.get("t1") == None)
+
+    TestUtils.createTopic(zkClient, tp1.topic(), assignment1, servers)
+
+    // Test that the first topic has its ID added correctly
+    waitForPartitionState(tp1, firstControllerEpoch, 0, LeaderAndIsr.initialLeaderEpoch,
+      "failed to get expected partition state upon topic creation")
+    assertNotEquals(controller.controllerContext.topicIds.get("t1"), None)
+    val topicId1 = controller.controllerContext.topicIds("t1")
+    assertEquals("t1", controller.controllerContext.topicNames(topicId1))
+
+    val tp2 = new TopicPartition("t2", 0)
+    val assignment2 = Map(tp2.partition -> Seq(0))
+    TestUtils.createTopic(zkClient, tp2.topic(), assignment2, servers)
+
+    // Test that the second topic has its ID added correctly
+    waitForPartitionState(tp2, firstControllerEpoch, 0, LeaderAndIsr.initialLeaderEpoch,
+      "failed to get expected partition state upon topic creation")
+    assertNotEquals(controller.controllerContext.topicIds.get("t2"), None)
+    val topicId2 = controller.controllerContext.topicIds("t2")
+    assertEquals("t2", controller.controllerContext.topicNames(topicId2))
+
+    // The first topic ID has not changed
+    assertEquals(topicId1, controller.controllerContext.topicIds.get("t1").get)
+    assertNotEquals(topicId1, topicId2)
+  }
+
+
+  @Test
+  def testTopicIdMigrationAndHandling(): Unit = {
+    val tp = new TopicPartition("t", 0)
+    val assignment = Map(tp.partition -> ReplicaAssignment(Seq(0), List(), List()))
+    val adminZkClient = new AdminZkClient(zkClient)
+
+    servers = makeServers(1)
+    adminZkClient.createTopic(tp.topic, 1, 1)
+    waitForPartitionState(tp, firstControllerEpoch, 0, LeaderAndIsr.initialLeaderEpoch,
+      "failed to get expected partition state upon topic creation")
+    val topicIdAfterCreate = zkClient.getTopicIdsForTopics(Set(tp.topic())).get(tp.topic())
+    assertTrue(topicIdAfterCreate.isDefined)
+    assertEquals("correct topic ID cannot be found in the controller context",
+      topicIdAfterCreate, servers.head.kafkaController.controllerContext.topicIds.get(tp.topic))
+
+    adminZkClient.addPartitions(tp.topic, assignment, adminZkClient.getBrokerMetadatas(), 2)
+    val topicIdAfterAddition = zkClient.getTopicIdsForTopics(Set(tp.topic())).get(tp.topic())
+    assertEquals(topicIdAfterCreate, topicIdAfterAddition)
+    assertEquals("topic ID changed after partition additions",
+      topicIdAfterCreate, servers.head.kafkaController.controllerContext.topicIds.get(tp.topic))
+
+    adminZkClient.deleteTopic(tp.topic)
+    TestUtils.waitUntilTrue(() => servers.head.kafkaController.controllerContext.topicIds.get(tp.topic).isEmpty,
+      "topic ID for topic should have been removed from controller context after deletion")
+  }
+
+  @Test
+  def testTopicIdPersistsThroughControllerReelection(): Unit = {
+    servers = makeServers(2)
+    val controllerId = TestUtils.waitUntilControllerElected(zkClient)
+    val controller = getController().kafkaController
+    val tp = new TopicPartition("t", 0)
+    val assignment = Map(tp.partition -> Seq(controllerId))
+    TestUtils.createTopic(zkClient, tp.topic, partitionReplicaAssignment = assignment, servers = servers)
+    waitForPartitionState(tp, firstControllerEpoch, controllerId, LeaderAndIsr.initialLeaderEpoch,
+      "failed to get expected partition state upon topic creation")
+    val topicId = controller.controllerContext.topicIds.get("t").get
+
+    servers(controllerId).shutdown()
+    servers(controllerId).awaitShutdown()
+    TestUtils.waitUntilTrue(() => zkClient.getControllerId.isDefined, "failed to elect a controller")
+    val controller2 = getController().kafkaController
+    assertEquals(topicId, controller2.controllerContext.topicIds.get("t").get)
+  }
+
+  @Test
+  def testTopicIdPersistsThroughControllerRestart(): Unit = {
+    servers = makeServers(1)
+    val controllerId = TestUtils.waitUntilControllerElected(zkClient)
+    val controller = getController().kafkaController
+    val tp = new TopicPartition("t", 0)
+    val assignment = Map(tp.partition -> Seq(controllerId))
+    TestUtils.createTopic(zkClient, tp.topic, partitionReplicaAssignment = assignment, servers = servers)
+    waitForPartitionState(tp, firstControllerEpoch, controllerId, LeaderAndIsr.initialLeaderEpoch,
+      "failed to get expected partition state upon topic creation")
+    val topicId = controller.controllerContext.topicIds.get("t").get
+
+    servers(controllerId).shutdown()
+    servers(controllerId).awaitShutdown()
+    servers(controllerId).startup()
+    TestUtils.waitUntilTrue(() => zkClient.getControllerId.isDefined, "failed to elect a controller")
+    val controller2 = getController().kafkaController
+    assertEquals(topicId, controller2.controllerContext.topicIds.get("t").get)
+  }
+
+  private def testControllerMove(fun: () => Unit): Unit = {
+    val controller = getController().kafkaController
+    val appender = LogCaptureAppender.createAndRegister()
+    val previousLevel = LogCaptureAppender.setClassLoggerLevel(controller.getClass, Level.INFO)
+
+    try {
+      TestUtils.waitUntilTrue(() => {
+        controller.eventManager.state == ControllerState.Idle
+      }, "Controller event thread is still busy")
+
+      val latch = new CountDownLatch(1)
+
+      // 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(new MockEvent(ControllerState.TopicChange) {
+        override def process(): Unit = latch.await()
+        override def preempt(): Unit = {}
+      })
+
+      // Execute pre-defined logic. This can be topic creation/deletion, preferred leader election, etc.
+      fun()
+
+      // Delete the controller path, re-create /controller znode to emulate controller movement
+      zkClient.deleteController(controller.controllerContext.epochZkVersion)
+      zkClient.registerControllerAndIncrementControllerEpoch(servers.size)
+
+      // Resume the controller event thread. At this point, the controller should see mismatch controller epoch zkVersion and resign
+      latch.countDown()
+      TestUtils.waitUntilTrue(() => !controller.isActive, "Controller fails to resign")
+
+      // 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(classOf[ControllerMovedException].getName))
+      assertTrue(event.isDefined)
+
+    } finally {
+      LogCaptureAppender.unregister(appender)
+      LogCaptureAppender.setClassLoggerLevel(controller.eventManager.thread.getClass, previousLevel)
+    }
+  }
+
+  private def preferredReplicaLeaderElection(controllerId: Int, otherBroker: KafkaServer, tp: TopicPartition,
                                              replicas: Set[Int], leaderEpoch: Int): Unit = {
     otherBroker.shutdown()
     otherBroker.awaitShutdown()
-    waitForPartitionState(tp, KafkaController.InitialControllerEpoch, controllerId, leaderEpoch + 1,
+    waitForPartitionState(tp, firstControllerEpoch, controllerId, leaderEpoch + 1,
       "failed to get expected partition state upon broker shutdown")
     otherBroker.startup()
-    TestUtils.waitUntilTrue(() => zkUtils.getInSyncReplicasForPartition(tp.topic, tp.partition).toSet == replicas, "restarted broker failed to join in-sync replicas")
-    zkUtils.createPersistentPath(ZkUtils.PreferredReplicaLeaderElectionPath, ZkUtils.preferredReplicaLeaderElectionZkData(Set(tp)))
-    TestUtils.waitUntilTrue(() => !zkUtils.pathExists(ZkUtils.PreferredReplicaLeaderElectionPath),
+    TestUtils.waitUntilTrue(() => zkClient.getInSyncReplicasForPartition(new TopicPartition(tp.topic, tp.partition)).get.toSet == replicas, "restarted broker failed to join in-sync replicas")
+    zkClient.createPreferredReplicaElection(Set(tp))
+    TestUtils.waitUntilTrue(() => !zkClient.pathExists(PreferredReplicaElectionZNode.path),
       "failed to remove preferred replica leader election path after completion")
-    waitForPartitionState(tp, KafkaController.InitialControllerEpoch, otherBroker.config.brokerId, leaderEpoch + 2,
+    waitForPartitionState(tp, firstControllerEpoch, otherBroker.config.brokerId, leaderEpoch + 2,
       "failed to get expected partition state upon broker startup")
   }
 
   private def waitUntilControllerEpoch(epoch: Int, message: String): Unit = {
-    TestUtils.waitUntilTrue(() => zkUtils.readDataMaybeNull(ZkUtils.ControllerEpochPath)._1.map(_.toInt) == Some(epoch), message)
+    TestUtils.waitUntilTrue(() => zkClient.getControllerEpoch.map(_._1).contains(epoch) , message)
   }
 
-  private def waitForPartitionState(tp: TopicAndPartition,
+  private def waitForPartitionState(tp: TopicPartition,
                                     controllerEpoch: Int,
                                     leader: Int,
                                     leaderEpoch: Int,
                                     message: String): Unit = {
     TestUtils.waitUntilTrue(() => {
-      val leaderIsrAndControllerEpochMap = zkUtils.getPartitionLeaderAndIsrForTopics(Set(tp))
+      val leaderIsrAndControllerEpochMap = zkClient.getTopicPartitionStates(Seq(tp))
       leaderIsrAndControllerEpochMap.contains(tp) &&
         isExpectedPartitionState(leaderIsrAndControllerEpochMap(tp), controllerEpoch, leader, leaderEpoch)
     }, message)
@@ -332,19 +1031,37 @@ class ControllerIntegrationTest extends ZooKeeperTestHarness {
       leaderIsrAndControllerEpoch.leaderAndIsr.leader == leader &&
       leaderIsrAndControllerEpoch.leaderAndIsr.leaderEpoch == leaderEpoch
 
-  private def makeServers(numConfigs: Int, autoLeaderRebalanceEnable: Boolean = false, uncleanLeaderElectionEnable: Boolean = false) = {
-    val configs = TestUtils.createBrokerConfigs(numConfigs, zkConnect)
+  private def makeServers(numConfigs: Int,
+                          autoLeaderRebalanceEnable: Boolean = false,
+                          uncleanLeaderElectionEnable: Boolean = false,
+                          enableControlledShutdown: Boolean = true,
+                          listeners : Option[String] = None,
+                          listenerSecurityProtocolMap : Option[String] = None,
+                          controlPlaneListenerName : Option[String] = None,
+                          interBrokerProtocolVersion: Option[ApiVersion] = None,
+                          logDirCount: Int = 1) = {
+    val configs = TestUtils.createBrokerConfigs(numConfigs, zkConnect, enableControlledShutdown = enableControlledShutdown, logDirCount = logDirCount)
     configs.foreach { config =>
       config.setProperty(KafkaConfig.AutoLeaderRebalanceEnableProp, autoLeaderRebalanceEnable.toString)
       config.setProperty(KafkaConfig.UncleanLeaderElectionEnableProp, uncleanLeaderElectionEnable.toString)
       config.setProperty(KafkaConfig.LeaderImbalanceCheckIntervalSecondsProp, "1")
+      listeners.foreach(listener => config.setProperty(KafkaConfig.ListenersProp, listener))
+      listenerSecurityProtocolMap.foreach(listenerMap => config.setProperty(KafkaConfig.ListenerSecurityProtocolMapProp, listenerMap))
+      controlPlaneListenerName.foreach(controlPlaneListener => config.setProperty(KafkaConfig.ControlPlaneListenerNameProp, controlPlaneListener))
+      interBrokerProtocolVersion.foreach(ibp => config.setProperty(KafkaConfig.InterBrokerProtocolVersionProp, ibp.toString))
     }
     configs.map(config => TestUtils.createServer(KafkaConfig.fromProps(config)))
   }
 
   private def timer(metricName: String): Timer = {
-    Metrics.defaultRegistry.allMetrics.asScala.filterKeys(_.getMBeanName == metricName).values.headOption
-      .getOrElse(fail(s"Unable to find metric $metricName")).asInstanceOf[Timer]
+    KafkaYammerMetrics.defaultRegistry.allMetrics.asScala.filter { case (k, _) =>
+      k.getMBeanName == metricName
+    }.values.headOption.getOrElse(fail(s"Unable to find metric $metricName")).asInstanceOf[Timer]
+  }
+
+  private def getController(): KafkaServer = {
+    val controllerId = TestUtils.waitUntilControllerElected(zkClient)
+    servers.filter(s => s.config.brokerId == controllerId).head
   }
 
 }
diff --git a/core/src/test/scala/unit/kafka/controller/ControllerTestUtils.scala b/core/src/test/scala/unit/kafka/controller/ControllerTestUtils.scala
deleted file mode 100644
index 407297a871314..0000000000000
--- a/core/src/test/scala/unit/kafka/controller/ControllerTestUtils.scala
+++ /dev/null
@@ -1,35 +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.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 = EasyMock.createMock(classOf[ControllerEvent])
-    EasyMock.expect(mockEvent.state).andReturn(controllerState)
-    EasyMock.expect(mockEvent.process()).andAnswer(new IAnswer[Unit]() {
-      def answer(): Unit = {
-        process()
-      }
-    })
-    EasyMock.replay(mockEvent)
-    mockEvent
-  }
-}
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..b9a4d04198da0
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/controller/MockPartitionStateMachine.scala
@@ -0,0 +1,131 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package 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) {
+
+  var stateChangesByTargetState = mutable.Map.empty[PartitionState, Int].withDefaultValue(0)
+
+  def stateChangesCalls(targetState: PartitionState): Int = {
+    stateChangesByTargetState(targetState)
+  }
+
+  def clear(): Unit = {
+    stateChangesByTargetState.clear()
+  }
+
+  override def handleStateChanges(
+    partitions: Seq[TopicPartition],
+    targetState: PartitionState,
+    leaderElectionStrategy: Option[PartitionLeaderElectionStrategy]
+  ): Map[TopicPartition, Either[Throwable, LeaderAndIsr]] = {
+    stateChangesByTargetState(targetState) = stateChangesByTargetState(targetState) + 1
+
+    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).get
+      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.putPartitionLeadershipInfo(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..32bfc50d6bb17
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/controller/MockReplicaStateMachine.scala
@@ -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 kafka.controller
+
+import scala.collection.Seq
+import scala.collection.mutable
+
+class MockReplicaStateMachine(controllerContext: ControllerContext) extends ReplicaStateMachine(controllerContext) {
+  val stateChangesByTargetState = mutable.Map.empty[ReplicaState, Int].withDefaultValue(0)
+
+  def stateChangesCalls(targetState: ReplicaState): Int = {
+    stateChangesByTargetState(targetState)
+  }
+
+  def clear(): Unit = {
+    stateChangesByTargetState.clear()
+  }
+
+  override def handleStateChanges(replicas: Seq[PartitionAndReplica], targetState: ReplicaState): Unit = {
+    stateChangesByTargetState(targetState) = stateChangesByTargetState(targetState) + 1
+
+    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 f149fc93a4991..3fd419257c3bb 100644
--- a/core/src/test/scala/unit/kafka/controller/PartitionLeaderElectionAlgorithmsTest.scala
+++ b/core/src/test/scala/unit/kafka/controller/PartitionLeaderElectionAlgorithmsTest.scala
@@ -17,10 +17,16 @@
 package kafka.controller
 
 import org.junit.Assert._
-import org.junit.Test
-import org.scalatest.junit.JUnitSuite
+import org.junit.{Before, Test}
 
-class PartitionLeaderElectionAlgorithmsTest  extends JUnitSuite {
+class PartitionLeaderElectionAlgorithmsTest {
+  private var controllerContext: ControllerContext = null
+
+  @Before
+  def setUp(): Unit = {
+    controllerContext = new ControllerContext
+    controllerContext.stats.removeMetric("UncleanLeaderElectionsPerSec")
+  }
 
   @Test
   def testOfflinePartitionLeaderElection(): Unit = {
@@ -30,7 +36,8 @@ class PartitionLeaderElectionAlgorithmsTest  extends JUnitSuite {
     val leaderOpt = PartitionLeaderElectionAlgorithms.offlinePartitionLeaderElection(assignment,
       isr,
       liveReplicas,
-      uncleanLeaderElectionEnabled = false)
+      uncleanLeaderElectionEnabled = false,
+      controllerContext)
     assertEquals(Option(4), leaderOpt)
   }
 
@@ -42,9 +49,12 @@ class PartitionLeaderElectionAlgorithmsTest  extends JUnitSuite {
     val leaderOpt = PartitionLeaderElectionAlgorithms.offlinePartitionLeaderElection(assignment,
       isr,
       liveReplicas,
-      uncleanLeaderElectionEnabled = false)
+      uncleanLeaderElectionEnabled = false,
+      controllerContext)
     assertEquals(None, leaderOpt)
+    assertEquals(0, controllerContext.stats.uncleanLeaderElectionRate.count())
   }
+
   @Test
   def testOfflinePartitionLeaderElectionLastIsrOfflineUncleanLeaderElectionEnabled(): Unit = {
     val assignment = Seq(2, 4)
@@ -53,8 +63,10 @@ class PartitionLeaderElectionAlgorithmsTest  extends JUnitSuite {
     val leaderOpt = PartitionLeaderElectionAlgorithms.offlinePartitionLeaderElection(assignment,
       isr,
       liveReplicas,
-      uncleanLeaderElectionEnabled = true)
+      uncleanLeaderElectionEnabled = true,
+      controllerContext)
     assertEquals(Option(4), leaderOpt)
+    assertEquals(1, controllerContext.stats.uncleanLeaderElectionRate.count())
   }
 
   @Test
@@ -62,10 +74,9 @@ class PartitionLeaderElectionAlgorithmsTest  extends JUnitSuite {
     val reassignment = Seq(2, 4)
     val isr = Seq(2, 4)
     val liveReplicas = Set(4)
-    val leaderOpt = PartitionLeaderElectionAlgorithms.offlinePartitionLeaderElection(reassignment,
+    val leaderOpt = PartitionLeaderElectionAlgorithms.reassignPartitionLeaderElection(reassignment,
       isr,
-      liveReplicas,
-      uncleanLeaderElectionEnabled = false)
+      liveReplicas)
     assertEquals(Option(4), leaderOpt)
   }
 
diff --git a/core/src/test/scala/unit/kafka/controller/PartitionStateMachineTest.scala b/core/src/test/scala/unit/kafka/controller/PartitionStateMachineTest.scala
index 2e56c33ecacfb..894d56e8dd7f9 100644
--- a/core/src/test/scala/unit/kafka/controller/PartitionStateMachineTest.scala
+++ b/core/src/test/scala/unit/kafka/controller/PartitionStateMachineTest.scala
@@ -20,25 +20,21 @@ 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.zookeeper.{CreateResponse, GetDataResponse, ZooKeeperClientException}
+import kafka.zk.{KafkaZkClient, TopicPartitionStateZNode}
+import kafka.zookeeper._
 import org.apache.kafka.common.TopicPartition
 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
+import org.mockito.Mockito
 
-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,10 +49,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, mockTopicDeletionManager,
-      mockZkClient, partitionState, mockControllerBrokerRequestBatch)
+    partitionStateMachine = new ZkPartitionStateMachine(config, new StateChangeLogger(brokerId, true, None), controllerContext,
+      mockZkClient, mockControllerBrokerRequestBatch)
+  }
+
+  private def partitionState(partition: TopicPartition): PartitionState = {
+    controllerContext.partitionState(partition)
   }
 
   @Test
@@ -67,7 +65,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))
   }
 
@@ -79,89 +81,101 @@ class PartitionStateMachineTest extends JUnitSuite {
 
   @Test
   def testNewPartitionToOnlinePartitionTransition(): Unit = {
-    controllerContext.liveBrokers = Set(TestUtils.createBroker(brokerId, "host", 0))
-    controllerContext.partitionReplicaAssignment.put(partition, Seq(brokerId))
-    partitionState.put(partition, NewPartition)
+    controllerContext.setLiveBrokers(Map(TestUtils.createBrokerAndEpoch(brokerId, "host", 0)))
+    controllerContext.updatePartitionFullReplicaAssignment(partition, ReplicaAssignment(Seq(brokerId)))
+    controllerContext.putPartitionState(partition, NewPartition)
     val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(LeaderAndIsr(brokerId, List(brokerId)), controllerEpoch)
     EasyMock.expect(mockControllerBrokerRequestBatch.newBatch())
-    EasyMock.expect(mockZkClient.createTopicPartitionStatesRaw(Map(partition -> leaderIsrAndControllerEpoch)))
-      .andReturn(Seq(CreateResponse(Code.OK, null, Some(partition), null)))
+    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 = {
-    controllerContext.liveBrokers = Set(TestUtils.createBroker(brokerId, "host", 0))
-    controllerContext.partitionReplicaAssignment.put(partition, Seq(brokerId))
-    partitionState.put(partition, NewPartition)
+  def testNewPartitionToOnlinePartitionTransitionZooKeeperClientExceptionFromCreateStates(): Unit = {
+    controllerContext.setLiveBrokers(Map(TestUtils.createBrokerAndEpoch(brokerId, "host", 0)))
+    controllerContext.updatePartitionFullReplicaAssignment(partition, ReplicaAssignment(Seq(brokerId)))
+    controllerContext.putPartitionState(partition, NewPartition)
     val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(LeaderAndIsr(brokerId, List(brokerId)), controllerEpoch)
     EasyMock.expect(mockControllerBrokerRequestBatch.newBatch())
-    EasyMock.expect(mockZkClient.createTopicPartitionStatesRaw(Map(partition -> leaderIsrAndControllerEpoch)))
+    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))
   }
 
   @Test
   def testNewPartitionToOnlinePartitionTransitionErrorCodeFromCreateStates(): Unit = {
-    controllerContext.liveBrokers = Set(TestUtils.createBroker(brokerId, "host", 0))
-    controllerContext.partitionReplicaAssignment.put(partition, Seq(brokerId))
-    partitionState.put(partition, NewPartition)
+    controllerContext.setLiveBrokers(Map(TestUtils.createBrokerAndEpoch(brokerId, "host", 0)))
+    controllerContext.updatePartitionFullReplicaAssignment(partition, ReplicaAssignment(Seq(brokerId)))
+    controllerContext.putPartitionState(partition, NewPartition)
     val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(LeaderAndIsr(brokerId, List(brokerId)), controllerEpoch)
     EasyMock.expect(mockControllerBrokerRequestBatch.newBatch())
-    EasyMock.expect(mockZkClient.createTopicPartitionStatesRaw(Map(partition -> leaderIsrAndControllerEpoch)))
-      .andReturn(Seq(CreateResponse(Code.NODEEXISTS, null, Some(partition), null)))
+    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))
   }
 
   @Test
   def testOnlinePartitionToOnlineTransition(): Unit = {
-    controllerContext.liveBrokers = Set(TestUtils.createBroker(brokerId, "host", 0))
-    controllerContext.partitionReplicaAssignment.put(partition, Seq(brokerId))
-    partitionState.put(partition, OnlinePartition)
+    controllerContext.setLiveBrokers(Map(TestUtils.createBrokerAndEpoch(brokerId, "host", 0)))
+    controllerContext.updatePartitionFullReplicaAssignment(partition, ReplicaAssignment(Seq(brokerId)))
+    controllerContext.putPartitionState(partition, OnlinePartition)
     val leaderAndIsr = LeaderAndIsr(brokerId, List(brokerId))
     val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(leaderAndIsr, controllerEpoch)
-    controllerContext.partitionLeadershipInfo.put(partition, leaderIsrAndControllerEpoch)
+    controllerContext.putPartitionLeadershipInfo(partition, leaderIsrAndControllerEpoch)
 
     val stat = new Stat(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0)
     EasyMock.expect(mockControllerBrokerRequestBatch.newBatch())
     EasyMock.expect(mockZkClient.getTopicPartitionStatesRaw(partitions))
       .andReturn(Seq(GetDataResponse(Code.OK, null, Some(partition),
-        TopicPartitionStateZNode.encode(leaderIsrAndControllerEpoch), stat)))
+        TopicPartitionStateZNode.encode(leaderIsrAndControllerEpoch), stat, ResponseMetadata(0, 0))))
 
     val leaderAndIsrAfterElection = leaderAndIsr.newLeader(brokerId)
     val updatedLeaderAndIsr = leaderAndIsrAfterElection.withZkVersion(2)
-    EasyMock.expect(mockZkClient.updateLeaderAndIsr(Map(partition -> leaderAndIsrAfterElection), controllerEpoch))
-      .andReturn(UpdateLeaderAndIsrResult(Map(partition -> updatedLeaderAndIsr), Seq.empty, Map.empty))
+    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), Seq(brokerId), isNew = false))
+      partition, LeaderIsrAndControllerEpoch(updatedLeaderAndIsr, controllerEpoch), replicaAssignment(Seq(brokerId)), isNew = false))
     EasyMock.expect(mockControllerBrokerRequestBatch.sendRequestsToBrokers(controllerEpoch))
     EasyMock.replay(mockZkClient, mockControllerBrokerRequestBatch)
 
@@ -173,26 +187,33 @@ class PartitionStateMachineTest extends JUnitSuite {
   @Test
   def testOnlinePartitionToOnlineTransitionForControlledShutdown(): Unit = {
     val otherBrokerId = brokerId + 1
-    controllerContext.liveBrokers = Set(TestUtils.createBroker(brokerId, "host", 0), TestUtils.createBroker(otherBrokerId, "host", 0))
+    controllerContext.setLiveBrokers(Map(
+      TestUtils.createBrokerAndEpoch(brokerId, "host", 0),
+      TestUtils.createBrokerAndEpoch(otherBrokerId, "host", 0)))
     controllerContext.shuttingDownBrokerIds.add(brokerId)
-    controllerContext.partitionReplicaAssignment.put(partition, Seq(brokerId, otherBrokerId))
-    partitionState.put(partition, OnlinePartition)
+    controllerContext.updatePartitionFullReplicaAssignment(
+      partition,
+      ReplicaAssignment(Seq(brokerId, otherBrokerId))
+    )
+    controllerContext.putPartitionState(partition, OnlinePartition)
     val leaderAndIsr = LeaderAndIsr(brokerId, List(brokerId, otherBrokerId))
     val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(leaderAndIsr, controllerEpoch)
-    controllerContext.partitionLeadershipInfo.put(partition, leaderIsrAndControllerEpoch)
+    controllerContext.putPartitionLeadershipInfo(partition, leaderIsrAndControllerEpoch)
 
     val stat = new Stat(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0)
     EasyMock.expect(mockControllerBrokerRequestBatch.newBatch())
     EasyMock.expect(mockZkClient.getTopicPartitionStatesRaw(partitions))
       .andReturn(Seq(GetDataResponse(Code.OK, null, Some(partition),
-        TopicPartitionStateZNode.encode(leaderIsrAndControllerEpoch), stat)))
+        TopicPartitionStateZNode.encode(leaderIsrAndControllerEpoch), stat, ResponseMetadata(0, 0))))
 
     val leaderAndIsrAfterElection = leaderAndIsr.newLeaderAndIsr(otherBrokerId, List(otherBrokerId))
     val updatedLeaderAndIsr = leaderAndIsrAfterElection.withZkVersion(2)
-    EasyMock.expect(mockZkClient.updateLeaderAndIsr(Map(partition -> leaderAndIsrAfterElection), controllerEpoch))
-      .andReturn(UpdateLeaderAndIsrResult(Map(partition -> updatedLeaderAndIsr), Seq.empty, Map.empty))
-    EasyMock.expect(mockControllerBrokerRequestBatch.addLeaderAndIsrRequestForBrokers(Seq(otherBrokerId),
-      partition, LeaderIsrAndControllerEpoch(updatedLeaderAndIsr, controllerEpoch), Seq(brokerId, otherBrokerId),
+    EasyMock.expect(mockZkClient.updateLeaderAndIsr(Map(partition -> leaderAndIsrAfterElection), controllerEpoch, controllerContext.epochZkVersion))
+      .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), replicaAssignment(Seq(brokerId, otherBrokerId)),
       isNew = false))
     EasyMock.expect(mockControllerBrokerRequestBatch.sendRequestsToBrokers(controllerEpoch))
     EasyMock.replay(mockZkClient, mockControllerBrokerRequestBatch)
@@ -204,64 +225,138 @@ 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))
   }
 
   @Test
   def testOfflinePartitionToOnlinePartitionTransition(): Unit = {
-    controllerContext.liveBrokers = Set(TestUtils.createBroker(brokerId, "host", 0))
-    controllerContext.partitionReplicaAssignment.put(partition, Seq(brokerId))
-    partitionState.put(partition, OfflinePartition)
+    controllerContext.setLiveBrokers(Map(TestUtils.createBrokerAndEpoch(brokerId, "host", 0)))
+    controllerContext.updatePartitionFullReplicaAssignment(partition, ReplicaAssignment(Seq(brokerId)))
+    controllerContext.putPartitionState(partition, OfflinePartition)
     val leaderAndIsr = LeaderAndIsr(LeaderAndIsr.NoLeader, List(brokerId))
     val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(leaderAndIsr, controllerEpoch)
-    controllerContext.partitionLeadershipInfo.put(partition, leaderIsrAndControllerEpoch)
+    controllerContext.putPartitionLeadershipInfo(partition, leaderIsrAndControllerEpoch)
 
     val stat = new Stat(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0)
     EasyMock.expect(mockControllerBrokerRequestBatch.newBatch())
     EasyMock.expect(mockZkClient.getTopicPartitionStatesRaw(partitions))
       .andReturn(Seq(GetDataResponse(Code.OK, null, Some(partition),
-        TopicPartitionStateZNode.encode(leaderIsrAndControllerEpoch), stat)))
+        TopicPartitionStateZNode.encode(leaderIsrAndControllerEpoch), stat, ResponseMetadata(0, 0))))
 
-    EasyMock.expect(mockZkClient.getLogConfigs(Seq.empty, config.originals()))
+    EasyMock.expect(mockZkClient.getLogConfigs(Set.empty, config.originals()))
       .andReturn((Map(partition.topic -> LogConfig()), Map.empty))
     val leaderAndIsrAfterElection = leaderAndIsr.newLeader(brokerId)
     val updatedLeaderAndIsr = leaderAndIsrAfterElection.withZkVersion(2)
-    EasyMock.expect(mockZkClient.updateLeaderAndIsr(Map(partition -> leaderAndIsrAfterElection), controllerEpoch))
-      .andReturn(UpdateLeaderAndIsrResult(Map(partition -> updatedLeaderAndIsr), Seq.empty, Map.empty))
+    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), 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))
+    partitionStateMachine.handleStateChanges(
+      partitions,
+      OnlinePartition,
+      Option(OfflinePartitionLeaderElectionStrategy(false))
+    )
     EasyMock.verify(mockZkClient, mockControllerBrokerRequestBatch)
     assertEquals(OnlinePartition, partitionState(partition))
   }
 
   @Test
-  def testOfflinePartitionToOnlinePartitionTransitionZkUtilsExceptionFromStateLookup(): Unit = {
-    controllerContext.liveBrokers = Set(TestUtils.createBroker(brokerId, "host", 0))
-    controllerContext.partitionReplicaAssignment.put(partition, Seq(brokerId))
-    partitionState.put(partition, OfflinePartition)
+  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.setLiveBrokers(Map(TestUtils.createBrokerAndEpoch(brokerId, "host", 0)))
+    controllerContext.updatePartitionFullReplicaAssignment(
+      partition,
+      ReplicaAssignment(Seq(leaderBrokerId, brokerId))
+    )
+    controllerContext.putPartitionState(partition, OfflinePartition)
+
+    val leaderAndIsr = LeaderAndIsr(leaderBrokerId, List(leaderBrokerId))
+    val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(leaderAndIsr, controllerEpoch)
+    controllerContext.putPartitionLeadershipInfo(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.setLiveBrokers(Map(TestUtils.createBrokerAndEpoch(brokerId, "host", 0)))
+    controllerContext.updatePartitionFullReplicaAssignment(partition, ReplicaAssignment(Seq(brokerId)))
+    controllerContext.putPartitionState(partition, OfflinePartition)
     val leaderAndIsr = LeaderAndIsr(LeaderAndIsr.NoLeader, List(brokerId))
     val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(leaderAndIsr, controllerEpoch)
-    controllerContext.partitionLeadershipInfo.put(partition, leaderIsrAndControllerEpoch)
+    controllerContext.putPartitionLeadershipInfo(partition, leaderIsrAndControllerEpoch)
 
     EasyMock.expect(mockControllerBrokerRequestBatch.newBatch())
     EasyMock.expect(mockZkClient.getTopicPartitionStatesRaw(partitions))
@@ -270,46 +365,164 @@ 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 testOfflinePartitionToOnlinePartitionTransitionErrorCodeFromStateLookup(): Unit = {
-    controllerContext.liveBrokers = Set(TestUtils.createBroker(brokerId, "host", 0))
-    controllerContext.partitionReplicaAssignment.put(partition, Seq(brokerId))
-    partitionState.put(partition, OfflinePartition)
+    controllerContext.setLiveBrokers(Map(TestUtils.createBrokerAndEpoch(brokerId, "host", 0)))
+    controllerContext.updatePartitionFullReplicaAssignment(partition, ReplicaAssignment(Seq(brokerId)))
+    controllerContext.putPartitionState(partition, OfflinePartition)
     val leaderAndIsr = LeaderAndIsr(LeaderAndIsr.NoLeader, List(brokerId))
     val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(leaderAndIsr, controllerEpoch)
-    controllerContext.partitionLeadershipInfo.put(partition, leaderIsrAndControllerEpoch)
+    controllerContext.putPartitionLeadershipInfo(partition, leaderIsrAndControllerEpoch)
 
     val stat = new Stat(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0)
     EasyMock.expect(mockControllerBrokerRequestBatch.newBatch())
     EasyMock.expect(mockZkClient.getTopicPartitionStatesRaw(partitions))
       .andReturn(Seq(GetDataResponse(Code.NONODE, null, Some(partition),
-        TopicPartitionStateZNode.encode(leaderIsrAndControllerEpoch), stat)))
+        TopicPartitionStateZNode.encode(leaderIsrAndControllerEpoch), stat, 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(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))
   }
 
+  private def prepareMockToElectLeaderForPartitions(partitions: Seq[TopicPartition]): Unit = {
+    val leaderAndIsr = LeaderAndIsr(brokerId, List(brokerId))
+    def prepareMockToGetTopicPartitionsStatesRaw(): Unit = {
+      val stat = new Stat(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0)
+      val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(leaderAndIsr, controllerEpoch)
+      val getDataResponses = partitions.map {p => GetDataResponse(Code.OK, null, Some(p),
+        TopicPartitionStateZNode.encode(leaderIsrAndControllerEpoch), stat, ResponseMetadata(0, 0))}
+      EasyMock.expect(mockZkClient.getTopicPartitionStatesRaw(partitions))
+        .andReturn(getDataResponses)
+    }
+    prepareMockToGetTopicPartitionsStatesRaw()
+
+    def prepareMockToGetLogConfigs(): Unit = {
+      EasyMock.expect(mockZkClient.getLogConfigs(Set.empty, config.originals()))
+        .andReturn(Map.empty, Map.empty)
+    }
+    prepareMockToGetLogConfigs()
+
+    def prepareMockToUpdateLeaderAndIsr(): Unit = {
+      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.map { case (k, v) => k -> Right(v) }, Seq.empty))
+    }
+    prepareMockToUpdateLeaderAndIsr()
+  }
+
+  /**
+    * This method tests changing partitions' state to OfflinePartition increments the offlinePartitionCount,
+    * and changing their state back to OnlinePartition decrements the offlinePartitionCount
+    */
+  @Test
+  def testUpdatingOfflinePartitionsCount(): Unit = {
+    controllerContext.setLiveBrokers(Map(TestUtils.createBrokerAndEpoch(brokerId, "host", 0)))
+
+    val partitionIds = Seq(0, 1, 2, 3)
+    val topic = "test"
+    val partitions = partitionIds.map(new TopicPartition(topic, _))
+
+    partitions.foreach { partition =>
+      controllerContext.updatePartitionFullReplicaAssignment(partition, ReplicaAssignment(Seq(brokerId)))
+    }
+
+    prepareMockToElectLeaderForPartitions(partitions)
+    EasyMock.replay(mockZkClient)
+
+    partitionStateMachine.handleStateChanges(partitions, NewPartition)
+    partitionStateMachine.handleStateChanges(partitions, OfflinePartition)
+    assertEquals(s"There should be ${partitions.size} offline partition(s)", partitions.size, controllerContext.offlinePartitionCount)
+
+    partitionStateMachine.handleStateChanges(partitions, OnlinePartition, Some(OfflinePartitionLeaderElectionStrategy(false)))
+    assertEquals(s"There should be no offline partition(s)", 0, controllerContext.offlinePartitionCount)
+  }
+
+  /**
+    * This method tests if a topic is being deleted, then changing partitions' state to OfflinePartition makes no change
+    * to the offlinePartitionCount
+    */
+  @Test
+  def testNoOfflinePartitionsChangeForTopicsBeingDeleted() = {
+    val partitionIds = Seq(0, 1, 2, 3)
+    val topic = "test"
+    val partitions = partitionIds.map(new TopicPartition(topic, _))
+
+    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, controllerContext.offlinePartitionCount)
+  }
+
+  /**
+    * This method tests if some partitions are already in OfflinePartition state,
+    * then deleting their topic will decrement the offlinePartitionCount.
+    * For example, if partitions test-0, test-1, test-2, test-3 are in OfflinePartition state,
+    * and the offlinePartitionCount is 4, trying to delete the topic "test" means these
+    * partitions no longer qualify as offline-partitions, and the offlinePartitionCount
+    * should be decremented to 0.
+    */
+  @Test
+  def testUpdatingOfflinePartitionsCountDuringTopicDeletion() = {
+    val partitionIds = Seq(0, 1, 2, 3)
+    val topic = "test"
+    val partitions = partitionIds.map(new TopicPartition("test", _))
+    partitions.foreach { partition =>
+      controllerContext.updatePartitionFullReplicaAssignment(partition, ReplicaAssignment(Seq(brokerId)))
+    }
+
+    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)
+    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, 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 5d24d79ce166c..cbf6df3da86b6 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
+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.setLiveBrokers(liveBrokerEpochs)
+    controllerContext.updatePartitionFullReplicaAssignment(partition, ReplicaAssignment(Seq(brokerId)))
+    assertEquals(None, controllerContext.replicaStates.get(replica))
+    replicaStateMachine.startup()
+    assertEquals(OnlineReplica, replicaState(replica))
+  }
+
+  @Test
+  def testStartupOfflinePartition(): Unit = {
+    controllerContext.updatePartitionFullReplicaAssignment(partition, ReplicaAssignment(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.setLiveBrokers(liveBrokerEpochs)
+    controllerContext.updatePartitionFullReplicaAssignment(partition, ReplicaAssignment(Seq(shutdownBrokerId)))
+    assertEquals(None, controllerContext.replicaStates.get(offlineReplica))
+    replicaStateMachine.startup()
+    assertEquals(OfflineReplica, replicaState(offlineReplica))
   }
 
   @Test
@@ -103,23 +137,28 @@ class ReplicaStateMachineTest extends JUnitSuite {
 
   @Test
   def testNewReplicaToOnlineReplicaTransition(): Unit = {
-    replicaState.put(replica, NewReplica)
-    controllerContext.partitionReplicaAssignment.put(partition, Seq(brokerId))
+    controllerContext.putReplicaState(replica, NewReplica)
+    controllerContext.updatePartitionFullReplicaAssignment(partition, ReplicaAssignment(Seq(brokerId)))
     replicaStateMachine.handleStateChanges(replicas, OnlineReplica)
     assertEquals(OnlineReplica, replicaState(replica))
   }
 
   @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.setLiveBrokers(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)
-    assertEquals(NewReplica, replicaState(replica))
+    assertEquals(OfflineReplica, replicaState(replica))
   }
 
   @Test
@@ -149,13 +188,13 @@ class ReplicaStateMachineTest extends JUnitSuite {
 
   @Test
   def testOnlineReplicaToOnlineReplicaTransition(): Unit = {
-    replicaState.put(replica, OnlineReplica)
-    controllerContext.partitionReplicaAssignment.put(partition, Seq(brokerId))
+    controllerContext.putReplicaState(replica, OnlineReplica)
+    controllerContext.updatePartitionFullReplicaAssignment(partition, ReplicaAssignment(Seq(brokerId)))
     val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(LeaderAndIsr(brokerId, List(brokerId)), controllerEpoch)
-    controllerContext.partitionLeadershipInfo.put(partition, leaderIsrAndControllerEpoch)
+    controllerContext.putPartitionLeadershipInfo(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,33 +206,31 @@ class ReplicaStateMachineTest extends JUnitSuite {
   def testOnlineReplicaToOfflineReplicaTransition(): Unit = {
     val otherBrokerId = brokerId + 1
     val replicaIds = List(brokerId, otherBrokerId)
-    replicaState.put(replica, OnlineReplica)
-    controllerContext.partitionReplicaAssignment.put(partition, replicaIds)
+    controllerContext.putReplicaState(replica, OnlineReplica)
+    controllerContext.updatePartitionFullReplicaAssignment(partition, ReplicaAssignment(replicaIds))
     val leaderAndIsr = LeaderAndIsr(brokerId, replicaIds)
     val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(leaderAndIsr, controllerEpoch)
-    controllerContext.partitionLeadershipInfo.put(partition, leaderIsrAndControllerEpoch)
+    controllerContext.putPartitionLeadershipInfo(partition, leaderIsrAndControllerEpoch)
 
     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)
     EasyMock.expect(mockZkClient.getTopicPartitionStatesRaw(partitions)).andReturn(
       Seq(GetDataResponse(Code.OK, null, Some(partition),
-        TopicPartitionStateZNode.encode(leaderIsrAndControllerEpoch), stat)))
-    EasyMock.expect(mockZkClient.updateLeaderAndIsr(Map(partition -> adjustedLeaderAndIsr), controllerEpoch))
-      .andReturn(UpdateLeaderAndIsrResult(Map(partition -> updatedLeaderAndIsr), Seq.empty, Map.empty))
-    EasyMock.expect(mockTopicDeletionManager.isPartitionToBeDeleted(partition)).andReturn(false)
+        TopicPartitionStateZNode.encode(leaderIsrAndControllerEpoch), stat, ResponseMetadata(0, 0))))
+    EasyMock.expect(mockZkClient.updateLeaderAndIsr(Map(partition -> adjustedLeaderAndIsr), controllerEpoch, controllerContext.epochZkVersion))
+      .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)
-    assertEquals(updatedLeaderIsrAndControllerEpoch, controllerContext.partitionLeadershipInfo(partition))
+    EasyMock.verify(mockZkClient, mockControllerBrokerRequestBatch)
+    assertEquals(updatedLeaderIsrAndControllerEpoch, controllerContext.partitionLeadershipInfo(partition).get)
     assertEquals(OfflineReplica, replicaState(replica))
   }
 
@@ -224,13 +261,13 @@ class ReplicaStateMachineTest extends JUnitSuite {
 
   @Test
   def testOfflineReplicaToOnlineReplicaTransition(): Unit = {
-    replicaState.put(replica, OfflineReplica)
-    controllerContext.partitionReplicaAssignment.put(partition, Seq(brokerId))
+    controllerContext.putReplicaState(replica, OfflineReplica)
+    controllerContext.updatePartitionFullReplicaAssignment(partition, ReplicaAssignment(Seq(brokerId)))
     val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(LeaderAndIsr(brokerId, List(brokerId)), controllerEpoch)
-    controllerContext.partitionLeadershipInfo.put(partition, leaderIsrAndControllerEpoch)
+    controllerContext.putPartitionLeadershipInfo(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.partitionReplicaAssignment.put(partition, Seq(brokerId))
+    controllerContext.putReplicaState(replica, ReplicaDeletionSuccessful)
+    controllerContext.updatePartitionFullReplicaAssignment(partition, ReplicaAssignment(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.partitionReplicaAssignment.put(partition, Seq(brokerId))
+    controllerContext.putReplicaState(replica, ReplicaDeletionIneligible)
+    controllerContext.updatePartitionFullReplicaAssignment(partition, ReplicaAssignment(Seq(brokerId)))
     val leaderIsrAndControllerEpoch = LeaderIsrAndControllerEpoch(LeaderAndIsr(brokerId, List(brokerId)), controllerEpoch)
-    controllerContext.partitionLeadershipInfo.put(partition, leaderIsrAndControllerEpoch)
+    controllerContext.putPartitionLeadershipInfo(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..d9f52e100b817
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/controller/TopicDeletionManagerTest.scala
@@ -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 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
+    val barPartitions = controllerContext.partitionsForTopic("bar")
+    val barReplicas = controllerContext.replicasForPartition(barPartitions).toSet
+
+    // Clean up state changes before starting the deletion
+    replicaStateMachine.clear()
+    partitionStateMachine.clear()
+
+    // Queue the topic for deletion
+    deletionManager.enqueueTopicsForDeletion(Set("foo", "bar"))
+
+    assertEquals(fooPartitions, controllerContext.partitionsInState("foo", NonExistentPartition))
+    assertEquals(fooReplicas, controllerContext.replicasInState("foo", ReplicaDeletionStarted))
+    assertEquals(barPartitions, controllerContext.partitionsInState("bar", NonExistentPartition))
+    assertEquals(barReplicas, controllerContext.replicasInState("bar", ReplicaDeletionStarted))
+    verify(deletionClient).sendMetadataUpdate(fooPartitions ++ barPartitions)
+    assertEquals(Set("foo", "bar"), controllerContext.topicsToBeDeleted)
+    assertEquals(Set("foo", "bar"), controllerContext.topicsWithDeletionStarted)
+    assertEquals(Set(), controllerContext.topicsIneligibleForDeletion)
+
+    // Complete the deletion
+    deletionManager.completeReplicaDeletion(fooReplicas ++ barReplicas)
+
+    assertEquals(Set.empty, controllerContext.partitionsForTopic("foo"))
+    assertEquals(Set.empty[PartitionAndReplica], controllerContext.replicaStates.keySet.filter(_.topic == "foo"))
+    assertEquals(Set.empty, controllerContext.partitionsForTopic("bar"))
+    assertEquals(Set.empty[PartitionAndReplica], controllerContext.replicaStates.keySet.filter(_.topic == "bar"))
+    assertEquals(Set(), controllerContext.topicsToBeDeleted)
+    assertEquals(Set(), controllerContext.topicsWithDeletionStarted)
+    assertEquals(Set(), controllerContext.topicsIneligibleForDeletion)
+
+    assertEquals(1, partitionStateMachine.stateChangesCalls(OfflinePartition))
+    assertEquals(1, partitionStateMachine.stateChangesCalls(NonExistentPartition))
+
+    assertEquals(1, replicaStateMachine.stateChangesCalls(ReplicaDeletionIneligible))
+    assertEquals(1, replicaStateMachine.stateChangesCalls(OfflineReplica))
+    assertEquals(1, replicaStateMachine.stateChangesCalls(ReplicaDeletionStarted))
+    assertEquals(1, replicaStateMachine.stateChangesCalls(ReplicaDeletionSuccessful))
+  }
+
+  @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.addLiveBrokers(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.addLiveBrokers(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.setLiveBrokers(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.updatePartitionFullReplicaAssignment(partition, ReplicaAssignment(replicas))
+      leaderIndex += 1
+    }
+    context
+  }
+
+}
diff --git a/core/src/test/scala/unit/kafka/coordinator/AbstractCoordinatorConcurrencyTest.scala b/core/src/test/scala/unit/kafka/coordinator/AbstractCoordinatorConcurrencyTest.scala
new file mode 100644
index 0000000000000..42ce8302bc984
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/coordinator/AbstractCoordinatorConcurrencyTest.scala
@@ -0,0 +1,224 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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.coordinator
+
+import java.util.concurrent.{ConcurrentHashMap, Executors}
+import java.util.{Collections, Random}
+import java.util.concurrent.atomic.AtomicInteger
+import java.util.concurrent.locks.Lock
+
+import kafka.coordinator.AbstractCoordinatorConcurrencyTest._
+import kafka.log.{AppendOrigin, Log}
+import kafka.server._
+import kafka.utils._
+import kafka.utils.timer.MockTimer
+import kafka.zk.KafkaZkClient
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.record.{MemoryRecords, RecordBatch, RecordConversionStats}
+import org.apache.kafka.common.requests.ProduceResponse.PartitionResponse
+import org.easymock.EasyMock
+import org.junit.{After, Before}
+
+import scala.collection._
+import scala.jdk.CollectionConverters._
+
+abstract class AbstractCoordinatorConcurrencyTest[M <: CoordinatorMember] {
+
+  val nThreads = 5
+
+  val time = new MockTime
+  val timer = new MockTimer
+  val executor = Executors.newFixedThreadPool(nThreads)
+  val scheduler = new MockScheduler(time)
+  var replicaManager: TestReplicaManager = _
+  var zkClient: KafkaZkClient = _
+  val serverProps = TestUtils.createBrokerConfig(nodeId = 0, zkConnect = "")
+  val random = new Random
+
+  @Before
+  def setUp(): Unit = {
+
+    replicaManager = EasyMock.partialMockBuilder(classOf[TestReplicaManager]).createMock()
+    replicaManager.createDelayedProducePurgatory(timer)
+
+    zkClient = EasyMock.createNiceMock(classOf[KafkaZkClient])
+  }
+
+  @After
+  def tearDown(): Unit = {
+    EasyMock.reset(replicaManager)
+    if (executor != null)
+      executor.shutdownNow()
+  }
+
+  /**
+    * Verify that concurrent operations run in the normal sequence produce the expected results.
+    */
+  def verifyConcurrentOperations(createMembers: String => Set[M], operations: Seq[Operation]): Unit = {
+    OrderedOperationSequence(createMembers("verifyConcurrentOperations"), operations).run()
+  }
+
+  /**
+    * 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]): Unit = {
+    EasyMock.reset(replicaManager)
+    for (i <- 0 to 10) {
+      // Run some random operations
+      RandomOperationSequence(createMembers(s"random$i"), operations).run()
+
+      // Check that proper sequences still work correctly
+      OrderedOperationSequence(createMembers(s"ordered$i"), operations).run()
+    }
+  }
+
+  def verifyConcurrentActions(actions: Set[Action]): Unit = {
+    val futures = actions.map(executor.submit)
+    futures.map(_.get)
+    enableCompletion()
+    actions.foreach(_.await())
+  }
+
+  def enableCompletion(): Unit = {
+    replicaManager.tryCompleteActions()
+    scheduler.tick()
+  }
+
+  abstract class OperationSequence(members: Set[M], operations: Seq[Operation]) {
+    def actionSequence: Seq[Set[Action]]
+    def run(): Unit = {
+      actionSequence.foreach(verifyConcurrentActions)
+    }
+  }
+
+  case class OrderedOperationSequence(members: Set[M], operations: Seq[Operation])
+    extends OperationSequence(members, operations) {
+    override def actionSequence: Seq[Set[Action]] = {
+      operations.map { op =>
+        members.map(op.actionWithVerify)
+      }
+    }
+  }
+
+  case class RandomOperationSequence(members: Set[M], operations: Seq[Operation])
+    extends OperationSequence(members, operations) {
+    val opCount = operations.length
+    def actionSequence: Seq[Set[Action]] = {
+      (0 to opCount).map { _ =>
+        members.map { member =>
+          val op = operations(random.nextInt(opCount))
+          op.actionNoVerify(member) // Don't wait or verify since these operations may block
+        }
+      }
+    }
+  }
+
+  abstract class Operation {
+    def run(member: M): Unit
+    def awaitAndVerify(member: M): Unit
+    def actionWithVerify(member: M): Action = {
+      new Action() {
+        def run(): Unit = Operation.this.run(member)
+        def await(): Unit = awaitAndVerify(member)
+      }
+    }
+    def actionNoVerify(member: M): Action = {
+      new Action() {
+        def run(): Unit = Operation.this.run(member)
+        def await(): Unit = timer.advanceClock(100) // Don't wait since operation may block
+      }
+    }
+  }
+}
+
+object AbstractCoordinatorConcurrencyTest {
+
+  trait Action extends Runnable {
+    def await(): Unit
+  }
+
+  trait CoordinatorMember {
+  }
+
+  class TestReplicaManager extends ReplicaManager(
+    null, null, null, null, null, null, null, null, null, null, null, null, null, null, null, None, null) {
+
+    var producePurgatory: DelayedOperationPurgatory[DelayedProduce] = _
+    var watchKeys: mutable.Set[TopicPartitionOperationKey] = _
+    def createDelayedProducePurgatory(timer: MockTimer): Unit = {
+      producePurgatory = new DelayedOperationPurgatory[DelayedProduce]("Produce", timer, 1, reaperEnabled = false)
+      watchKeys = Collections.newSetFromMap(new ConcurrentHashMap[TopicPartitionOperationKey, java.lang.Boolean]()).asScala
+    }
+
+    override def tryCompleteActions(): Unit = watchKeys.map(producePurgatory.checkAndComplete)
+
+    override def appendRecords(timeout: Long,
+                               requiredAcks: Short,
+                               internalTopicsAllowed: Boolean,
+                               origin: AppendOrigin,
+                               entriesPerPartition: Map[TopicPartition, MemoryRecords],
+                               responseCallback: Map[TopicPartition, PartitionResponse] => Unit,
+                               delayedProduceLock: Option[Lock] = None,
+                               processingStatsCallback: Map[TopicPartition, RecordConversionStats] => Unit = _ => ()): Unit = {
+
+      if (entriesPerPartition.isEmpty)
+        return
+      val produceMetadata = ProduceMetadata(1, entriesPerPartition.map {
+        case (tp, _) =>
+          (tp, ProducePartitionStatus(0L, new PartitionResponse(Errors.NONE, 0L, RecordBatch.NO_TIMESTAMP, 0L)))
+      })
+      val delayedProduce = new DelayedProduce(5, produceMetadata, this, responseCallback, delayedProduceLock) {
+        // Complete produce requests after a few attempts to trigger delayed produce from different threads
+        val completeAttempts = new AtomicInteger
+        override def tryComplete(): Boolean = {
+          if (completeAttempts.incrementAndGet() >= 3)
+            forceComplete()
+          else
+            false
+        }
+        override def onComplete(): Unit = {
+          responseCallback(entriesPerPartition.map {
+            case (tp, _) =>
+              (tp, new PartitionResponse(Errors.NONE, 0L, RecordBatch.NO_TIMESTAMP, 0L))
+          })
+        }
+      }
+      val producerRequestKeys = entriesPerPartition.keys.map(TopicPartitionOperationKey(_)).toSeq
+      watchKeys ++= producerRequestKeys
+      producePurgatory.tryCompleteElseWatch(delayedProduce, producerRequestKeys)
+    }
+    override def getMagic(topicPartition: TopicPartition): Option[Byte] = {
+      Some(RecordBatch.MAGIC_VALUE_V2)
+    }
+    @volatile var logs: mutable.Map[TopicPartition, (Log, Long)] = _
+    def getOrCreateLogs(): mutable.Map[TopicPartition, (Log, Long)] = {
+      if (logs == null)
+        logs = mutable.Map[TopicPartition, (Log, Long)]()
+      logs
+    }
+    def updateLog(topicPartition: TopicPartition, log: Log, endOffset: Long): Unit = {
+      getOrCreateLogs().put(topicPartition, (log, endOffset))
+    }
+    override def getLog(topicPartition: TopicPartition): Option[Log] =
+      getOrCreateLogs().get(topicPartition).map(l => l._1)
+    override def getLogEndOffset(topicPartition: TopicPartition): Option[Long] =
+      getOrCreateLogs().get(topicPartition).map(l => l._2)
+  }
+}
diff --git a/core/src/test/scala/unit/kafka/coordinator/group/GroupCoordinatorConcurrencyTest.scala b/core/src/test/scala/unit/kafka/coordinator/group/GroupCoordinatorConcurrencyTest.scala
new file mode 100644
index 0000000000000..1f54bd51f6ab4
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/coordinator/group/GroupCoordinatorConcurrencyTest.scala
@@ -0,0 +1,406 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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.coordinator.group
+
+import java.util.Properties
+import java.util.concurrent.locks.{Lock, ReentrantLock}
+import java.util.concurrent.{ConcurrentHashMap, TimeUnit}
+
+import kafka.common.OffsetAndMetadata
+import kafka.coordinator.AbstractCoordinatorConcurrencyTest
+import kafka.coordinator.AbstractCoordinatorConcurrencyTest._
+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, OffsetFetchResponse}
+import org.apache.kafka.common.utils.Time
+import org.easymock.EasyMock
+import org.junit.Assert._
+import org.junit.{After, Before, Test}
+
+import scala.collection._
+import scala.concurrent.duration.Duration
+import scala.concurrent.{Await, Future, Promise, TimeoutException}
+
+class GroupCoordinatorConcurrencyTest extends AbstractCoordinatorConcurrencyTest[GroupMember] {
+
+  private val protocolType = "consumer"
+  private val protocolName = "range"
+  private val metadata = Array[Byte]()
+  private val protocols = List((protocolName, metadata))
+
+  private val nGroups = nThreads * 10
+  private val nMembersPerGroup = nThreads * 5
+  private val numPartitions = 2
+
+  private val allOperations = Seq(
+      new JoinGroupOperation,
+      new SyncGroupOperation,
+      new OffsetFetchOperation,
+      new CommitOffsetsOperation,
+      new HeartbeatOperation,
+      new LeaveGroupOperation
+  )
+
+  var heartbeatPurgatory: DelayedOperationPurgatory[DelayedHeartbeat] = _
+  var joinPurgatory: DelayedOperationPurgatory[DelayedJoin] = _
+  var groupCoordinator: GroupCoordinator = _
+
+  @Before
+  override def setUp(): Unit = {
+    super.setUp()
+
+    EasyMock.expect(zkClient.getTopicPartitionCount(Topic.GROUP_METADATA_TOPIC_NAME))
+      .andReturn(Some(numPartitions))
+      .anyTimes()
+    EasyMock.replay(zkClient)
+
+    serverProps.setProperty(KafkaConfig.GroupMinSessionTimeoutMsProp, ConsumerMinSessionTimeout.toString)
+    serverProps.setProperty(KafkaConfig.GroupMaxSessionTimeoutMsProp, ConsumerMaxSessionTimeout.toString)
+    serverProps.setProperty(KafkaConfig.GroupInitialRebalanceDelayMsProp, GroupInitialRebalanceDelay.toString)
+
+    val config = KafkaConfig.fromProps(serverProps)
+
+    heartbeatPurgatory = new DelayedOperationPurgatory[DelayedHeartbeat]("Heartbeat", timer, config.brokerId, reaperEnabled = false)
+    joinPurgatory = new DelayedOperationPurgatory[DelayedJoin]("Rebalance", timer, config.brokerId, reaperEnabled = false)
+
+    groupCoordinator = GroupCoordinator(config, zkClient, replicaManager, heartbeatPurgatory, joinPurgatory, timer.time, new Metrics())
+    groupCoordinator.startup(false)
+  }
+
+  @After
+  override def tearDown(): Unit = {
+    try {
+      if (groupCoordinator != null)
+        groupCoordinator.shutdown()
+    } finally {
+      super.tearDown()
+    }
+  }
+
+  def createGroupMembers(groupPrefix: String): Set[GroupMember] = {
+    (0 until nGroups).flatMap { i =>
+      new Group(s"$groupPrefix$i", nMembersPerGroup, groupCoordinator, replicaManager).members
+    }.toSet
+  }
+
+  @Test
+  def testConcurrentGoodPathSequence(): Unit = {
+    verifyConcurrentOperations(createGroupMembers, allOperations)
+  }
+
+  @Test
+  def testConcurrentTxnGoodPathSequence(): Unit = {
+    verifyConcurrentOperations(createGroupMembers, Seq(
+      new JoinGroupOperation,
+      new SyncGroupOperation,
+      new OffsetFetchOperation,
+      new CommitTxnOffsetsOperation,
+      new CompleteTxnOperation,
+      new HeartbeatOperation,
+      new LeaveGroupOperation
+    ))
+  }
+
+  @Test
+  def testConcurrentRandomSequence(): Unit = {
+    /**
+     * handleTxnCommitOffsets does not complete delayed requests now so it causes error if handleTxnCompletion is executed
+     * before completing delayed request. In random mode, we use this global lock to prevent such an error.
+     */
+    val lock = new ReentrantLock()
+    verifyConcurrentRandomSequences(createGroupMembers, Seq(
+      new JoinGroupOperation,
+      new SyncGroupOperation,
+      new OffsetFetchOperation,
+      new CommitTxnOffsetsOperation(lock = Some(lock)),
+      new CompleteTxnOperation(lock = Some(lock)),
+      new HeartbeatOperation,
+      new LeaveGroupOperation
+    ))
+  }
+
+  @Test
+  def testConcurrentJoinGroupEnforceGroupMaxSize(): Unit = {
+    val groupMaxSize = 1
+    val newProperties = new Properties
+    newProperties.put(KafkaConfig.GroupMaxSizeProp, groupMaxSize.toString)
+    val config = KafkaConfig.fromProps(serverProps, newProperties)
+
+    if (groupCoordinator != null)
+      groupCoordinator.shutdown()
+    groupCoordinator = GroupCoordinator(config, zkClient, replicaManager, heartbeatPurgatory,
+      joinPurgatory, timer.time, new Metrics())
+    groupCoordinator.startup(false)
+
+    val members = new Group(s"group", nMembersPerGroup, groupCoordinator, replicaManager)
+      .members
+    val joinOp = new JoinGroupOperation()
+
+    verifyConcurrentActions(members.toSet.map(joinOp.actionNoVerify))
+
+    val errors = members.map { member =>
+      val joinGroupResult = joinOp.await(member, DefaultRebalanceTimeout)
+      joinGroupResult.error
+    }
+
+    assertEquals(groupMaxSize, errors.count(_ == Errors.NONE))
+    assertEquals(members.size-groupMaxSize, errors.count(_ == Errors.GROUP_MAX_SIZE_REACHED))
+  }
+
+  abstract class GroupOperation[R, C] extends Operation {
+    val responseFutures = new ConcurrentHashMap[GroupMember, Future[R]]()
+
+    def setUpCallback(member: GroupMember): C = {
+      val responsePromise = Promise[R]()
+      val responseFuture = responsePromise.future
+      responseFutures.put(member, responseFuture)
+      responseCallback(responsePromise)
+    }
+    def responseCallback(responsePromise: Promise[R]): C
+
+    override def run(member: GroupMember): Unit = {
+      val responseCallback = setUpCallback(member)
+      runWithCallback(member, responseCallback)
+    }
+
+    def runWithCallback(member: GroupMember, responseCallback: C): Unit
+
+    def await(member: GroupMember, timeoutMs: Long): R = {
+      var retries = (timeoutMs + 10) / 10
+      val responseFuture = responseFutures.get(member)
+      while (retries > 0) {
+        timer.advanceClock(10)
+        try {
+          return Await.result(responseFuture, Duration(10, TimeUnit.MILLISECONDS))
+        } catch {
+          case _: TimeoutException =>
+        }
+        retries -= 1
+      }
+      throw new TimeoutException(s"Operation did not complete within $timeoutMs millis")
+    }
+  }
+
+  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, None, requireKnownMemberId = false, "clientId", "clientHost",
+       DefaultRebalanceTimeout, DefaultSessionTimeout,
+       protocolType, protocols, responseCallback)
+      replicaManager.tryCompleteActions()
+    }
+    override def awaitAndVerify(member: GroupMember): Unit = {
+       val joinGroupResult = await(member, DefaultRebalanceTimeout)
+       assertEquals(Errors.NONE, joinGroupResult.error)
+       member.memberId = joinGroupResult.memberId
+       member.generationId = joinGroupResult.generationId
+    }
+  }
+
+  class SyncGroupOperation extends GroupOperation[SyncGroupCallbackParams, SyncGroupCallback] {
+    override def responseCallback(responsePromise: Promise[SyncGroupCallbackParams]): SyncGroupCallback = {
+      val callback: SyncGroupCallback = syncGroupResult =>
+        responsePromise.success(syncGroupResult.error, syncGroupResult.memberAssignment)
+      callback
+    }
+    override def runWithCallback(member: GroupMember, responseCallback: SyncGroupCallback): Unit = {
+      if (member.leader) {
+        groupCoordinator.handleSyncGroup(member.groupId, member.generationId, member.memberId,
+          Some(protocolType), Some(protocolName), member.groupInstanceId, member.group.assignment, responseCallback)
+      } else {
+        groupCoordinator.handleSyncGroup(member.groupId, member.generationId, member.memberId,
+          Some(protocolType), Some(protocolName), member.groupInstanceId, Map.empty[String, Array[Byte]], responseCallback)
+      }
+      replicaManager.tryCompleteActions()
+    }
+    override def awaitAndVerify(member: GroupMember): Unit = {
+      val result = await(member, DefaultSessionTimeout)
+      assertEquals(Errors.NONE, result._1)
+      assertNotNull(result._2)
+      assertEquals(0, result._2.length)
+    }
+  }
+
+  class HeartbeatOperation extends GroupOperation[HeartbeatCallbackParams, HeartbeatCallback] {
+    override def responseCallback(responsePromise: Promise[HeartbeatCallbackParams]): HeartbeatCallback = {
+      val callback: HeartbeatCallback = error => responsePromise.success(error)
+      callback
+    }
+    override def runWithCallback(member: GroupMember, responseCallback: HeartbeatCallback): Unit = {
+      groupCoordinator.handleHeartbeat(member.groupId, member.memberId,
+        member.groupInstanceId, member.generationId, responseCallback)
+      replicaManager.tryCompleteActions()
+    }
+    override def awaitAndVerify(member: GroupMember): Unit = {
+       val error = await(member, DefaultSessionTimeout)
+       assertEquals(Errors.NONE, error)
+    }
+  }
+
+  class OffsetFetchOperation extends GroupOperation[OffsetFetchCallbackParams, OffsetFetchCallback] {
+    override def responseCallback(responsePromise: Promise[OffsetFetchCallbackParams]): OffsetFetchCallback = {
+      val callback: OffsetFetchCallback = (error, offsets) => responsePromise.success(error, offsets)
+      callback
+    }
+    override def runWithCallback(member: GroupMember, responseCallback: OffsetFetchCallback): Unit = {
+      val (error, partitionData) = groupCoordinator.handleFetchOffsets(member.groupId, requireStable = true, None)
+      replicaManager.tryCompleteActions()
+      responseCallback(error, partitionData)
+    }
+    override def awaitAndVerify(member: GroupMember): Unit = {
+      val result = await(member, 500)
+      assertEquals(Errors.NONE, result._1)
+      assertEquals(Map.empty, result._2)
+    }
+  }
+
+  class CommitOffsetsOperation extends GroupOperation[CommitOffsetCallbackParams, CommitOffsetCallback] {
+    override def responseCallback(responsePromise: Promise[CommitOffsetCallbackParams]): CommitOffsetCallback = {
+      val callback: CommitOffsetCallback = offsets => responsePromise.success(offsets)
+      callback
+    }
+    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.groupInstanceId, member.generationId, offsets, responseCallback)
+      replicaManager.tryCompleteActions()
+    }
+    override def awaitAndVerify(member: GroupMember): Unit = {
+       val offsets = await(member, 500)
+       offsets.foreach { case (_, error) => assertEquals(Errors.NONE, error) }
+    }
+  }
+
+  class CommitTxnOffsetsOperation(lock: Option[Lock] = None) extends CommitOffsetsOperation {
+    override def runWithCallback(member: GroupMember, responseCallback: CommitOffsetCallback): Unit = {
+      val tp = new TopicPartition("topic", 0)
+      val offsets = immutable.Map(tp -> OffsetAndMetadata(1, "", Time.SYSTEM.milliseconds()))
+      val producerId = 1000L
+      val producerEpoch : Short = 2
+      // When transaction offsets are appended to the log, transactions may be scheduled for
+      // completion. Since group metadata locks are acquired for transaction completion, include
+      // this in the callback to test that there are no deadlocks.
+      def callbackWithTxnCompletion(errors: Map[TopicPartition, Errors]): Unit = {
+        val offsetsPartitions = (0 to numPartitions).map(new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, _))
+        groupCoordinator.groupManager.scheduleHandleTxnCompletion(producerId,
+          offsetsPartitions.map(_.partition).toSet, isCommit = random.nextBoolean)
+        responseCallback(errors)
+      }
+      lock.foreach(_.lock())
+      try {
+        groupCoordinator.handleTxnCommitOffsets(member.group.groupId, producerId, producerEpoch,
+          JoinGroupRequest.UNKNOWN_MEMBER_ID, Option.empty, JoinGroupRequest.UNKNOWN_GENERATION_ID,
+          offsets, callbackWithTxnCompletion)
+        replicaManager.tryCompleteActions()
+      } finally lock.foreach(_.unlock())
+    }
+  }
+
+  class CompleteTxnOperation(lock: Option[Lock] = None) extends GroupOperation[CompleteTxnCallbackParams, CompleteTxnCallback] {
+    override def responseCallback(responsePromise: Promise[CompleteTxnCallbackParams]): CompleteTxnCallback = {
+      val callback: CompleteTxnCallback = error => responsePromise.success(error)
+      callback
+    }
+    override def runWithCallback(member: GroupMember, responseCallback: CompleteTxnCallback): Unit = {
+      val producerId = 1000L
+      val offsetsPartitions = (0 to numPartitions).map(new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, _))
+      lock.foreach(_.lock())
+      try {
+        groupCoordinator.groupManager.handleTxnCompletion(producerId,
+          offsetsPartitions.map(_.partition).toSet, isCommit = random.nextBoolean)
+        responseCallback(Errors.NONE)
+      } finally lock.foreach(_.unlock())
+
+    }
+    override def awaitAndVerify(member: GroupMember): Unit = {
+      val error = await(member, 500)
+      assertEquals(Errors.NONE, error)
+    }
+  }
+
+  class LeaveGroupOperation extends GroupOperation[LeaveGroupCallbackParams, LeaveGroupCallback] {
+    override def responseCallback(responsePromise: Promise[LeaveGroupCallbackParams]): LeaveGroupCallback = {
+      val callback: LeaveGroupCallback = result => responsePromise.success(result)
+      callback
+    }
+    override def runWithCallback(member: GroupMember, responseCallback: LeaveGroupCallback): Unit = {
+      val memberIdentity = new MemberIdentity()
+          .setMemberId(member.memberId)
+      groupCoordinator.handleLeaveGroup(member.group.groupId, List(memberIdentity), responseCallback)
+    }
+    override def awaitAndVerify(member: GroupMember): Unit = {
+      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 = (Errors, Array[Byte])
+  type SyncGroupCallback = SyncGroupResult => Unit
+  type HeartbeatCallbackParams = Errors
+  type HeartbeatCallback = Errors => Unit
+  type OffsetFetchCallbackParams = (Errors, Map[TopicPartition, OffsetFetchResponse.PartitionData])
+  type OffsetFetchCallback = (Errors, Map[TopicPartition, OffsetFetchResponse.PartitionData]) => Unit
+  type CommitOffsetCallbackParams = Map[TopicPartition, Errors]
+  type CommitOffsetCallback = Map[TopicPartition, Errors] => Unit
+  type LeaveGroupCallbackParams = LeaveGroupResult
+  type LeaveGroupCallback = LeaveGroupResult => Unit
+  type CompleteTxnCallbackParams = Errors
+  type CompleteTxnCallback = Errors => Unit
+
+  private val ConsumerMinSessionTimeout = 10
+  private val ConsumerMaxSessionTimeout = 120 * 1000
+  private val DefaultRebalanceTimeout = 60 * 1000
+  private val DefaultSessionTimeout = 60 * 1000
+  private val GroupInitialRebalanceDelay = 50
+
+  class Group(val groupId: String, nMembers: Int,
+      groupCoordinator: GroupCoordinator, replicaManager: TestReplicaManager) {
+    val groupPartitionId = groupCoordinator.partitionFor(groupId)
+    groupCoordinator.groupManager.addPartitionOwnership(groupPartitionId)
+    val members = (0 until nMembers).map { i =>
+      new GroupMember(this, groupPartitionId, i == 0)
+    }
+    def assignment: Map[String, Array[Byte]] = members.map { m => (m.memberId, Array[Byte]()) }.toMap
+  }
+
+  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 1c770a416c316..d9331f6ae687e 100644
--- a/core/src/test/scala/unit/kafka/coordinator/group/GroupCoordinatorTest.scala
+++ b/core/src/test/scala/unit/kafka/coordinator/group/GroupCoordinatorTest.scala
@@ -17,8 +17,10 @@
 
 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
@@ -30,31 +32,40 @@ import org.easymock.{Capture, EasyMock, IAnswer}
 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._
+import scala.jdk.CollectionConverters._
+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 ConsumerMinSessionTimeout = 10
-  val ConsumerMaxSessionTimeout = 1000
+  val GroupMinSessionTimeout = 10
+  val GroupMaxSessionTimeout = 10 * 60 * 1000
+  val GroupMaxSize = 4
   val DefaultRebalanceTimeout = 500
   val DefaultSessionTimeout = 500
   val GroupInitialRebalanceDelay = 50
@@ -66,19 +77,27 @@ class GroupCoordinatorTest extends JUnitSuite {
 
   private val groupId = "groupId"
   private val protocolType = "consumer"
+  private val protocolName = "range"
   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 protocols = List((protocolName, metadata))
+  private val protocolSuperset = List((protocolName, metadata), ("roundrobin", metadata))
+  private val requireStable = true
   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, ConsumerMinSessionTimeout.toString)
-    props.setProperty(KafkaConfig.GroupMaxSessionTimeoutMsProp, ConsumerMaxSessionTimeout.toString)
+    props.setProperty(KafkaConfig.GroupMinSessionTimeoutMsProp, GroupMinSessionTimeout.toString)
+    props.setProperty(KafkaConfig.GroupMaxSessionTimeoutMsProp, GroupMaxSessionTimeout.toString)
+    props.setProperty(KafkaConfig.GroupMaxSizeProp, GroupMaxSize.toString)
     props.setProperty(KafkaConfig.GroupInitialRebalanceDelayMsProp, GroupInitialRebalanceDelay.toString)
     // make two partitions of the group topic to make sure some partitions are not owned by the coordinator
     val ret = mutable.Map[String, Map[Int, Seq[Int]]]()
@@ -98,8 +117,8 @@ 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.startup(false)
+    groupCoordinator = GroupCoordinator(config, zkClient, replicaManager, heartbeatPurgatory, joinPurgatory, timer.time, new Metrics())
+    groupCoordinator.startup(enableMetadataExpiration = false)
 
     // add the partition into the owned partition list
     groupPartitionId = groupCoordinator.partitionFor(groupId)
@@ -107,14 +126,75 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   @After
-  def tearDown() {
+  def tearDown(): Unit = {
     EasyMock.reset(replicaManager)
     if (groupCoordinator != null)
       groupCoordinator.shutdown()
   }
 
   @Test
-  def testOffsetsRetentionMsIntegerOverflow() {
+  def testRequestHandlingWhileLoadingInProgress(): Unit = {
+    val otherGroupPartitionId = groupCoordinator.groupManager.partitionFor(otherGroupId)
+    assertTrue(otherGroupPartitionId != groupPartitionId)
+
+    groupCoordinator.groupManager.addLoadingPartition(otherGroupPartitionId)
+    assertTrue(groupCoordinator.groupManager.isGroupLoading(otherGroupId))
+
+    // Dynamic Member JoinGroup
+    var joinGroupResponse: Option[JoinGroupResult] = None
+    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, Some("consumer"), Some("range"), 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, 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, None, 1, error => { heartbeatError = Some(error) })
+    assertEquals(Some(Errors.NONE), heartbeatError)
+
+    // DescribeGroups
+    val (describeGroupError, _) = groupCoordinator.handleDescribeGroup(otherGroupId)
+    assertEquals(Errors.COORDINATOR_LOAD_IN_PROGRESS, describeGroupError)
+
+    // ListGroups
+    val (listGroupsError, _) = groupCoordinator.handleListGroups(Set())
+    assertEquals(Errors.COORDINATOR_LOAD_IN_PROGRESS, listGroupsError)
+
+    // DeleteGroups
+    val deleteGroupsErrors = groupCoordinator.handleDeleteGroups(Set(otherGroupId))
+    assertEquals(Some(Errors.COORDINATOR_LOAD_IN_PROGRESS), deleteGroupsErrors.get(otherGroupId))
+
+    // Check that non-loading groups are still accessible
+    assertEquals(Errors.NONE, groupCoordinator.handleDescribeGroup(groupId)._1)
+
+    // After loading, we should be able to access the group
+    val otherGroupMetadataTopicPartition = new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, otherGroupPartitionId)
+    EasyMock.reset(replicaManager)
+    EasyMock.expect(replicaManager.getLog(otherGroupMetadataTopicPartition)).andReturn(None)
+    EasyMock.replay(replicaManager)
+    groupCoordinator.groupManager.loadGroupsAndOffsets(otherGroupMetadataTopicPartition, group => {}, 0L)
+    assertEquals(Errors.NONE, groupCoordinator.handleDescribeGroup(otherGroupId)._1)
+  }
+
+  @Test
+  def testOffsetsRetentionMsIntegerOverflow(): Unit = {
     val props = TestUtils.createBrokerConfig(nodeId = 0, zkConnect = "")
     props.setProperty(KafkaConfig.OffsetsRetentionMinutesProp, Integer.MAX_VALUE.toString)
     val config = KafkaConfig.fromProps(props)
@@ -123,1125 +203,3506 @@ 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 testJoinGroupSessionTimeoutTooSmall() {
-    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
+  def testJoinGroupShouldReceiveErrorIfGroupOverMaxSize(): Unit = {
+    val futures = ArrayBuffer[Future[JoinGroupResult]]()
+    val rebalanceTimeout = GroupInitialRebalanceDelay * 2
 
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols, sessionTimeout = ConsumerMinSessionTimeout - 1)
-    val joinGroupError = joinGroupResult.error
-    assertEquals(Errors.INVALID_SESSION_TIMEOUT, joinGroupError)
+    for (i <- 1.to(GroupMaxSize)) {
+      futures += sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols, rebalanceTimeout = rebalanceTimeout)
+      if (i != 1)
+        timer.advanceClock(GroupInitialRebalanceDelay)
+      EasyMock.reset(replicaManager)
+    }
+    // advance clock beyond rebalanceTimeout
+    timer.advanceClock(GroupInitialRebalanceDelay + 1)
+    for (future <- futures) {
+      assertEquals(Errors.NONE, await(future, 1).error)
+    }
+
+    // Should receive an error since the group is full
+    val errorFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols, rebalanceTimeout = rebalanceTimeout)
+    assertEquals(Errors.GROUP_MAX_SIZE_REACHED, await(errorFuture, 1).error)
   }
 
   @Test
-  def testJoinGroupSessionTimeoutTooLarge() {
-    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
+  def testDynamicMembersJoinGroupWithMaxSizeAndRequiredKnownMember(): Unit = {
+    val requiredKnownMemberId = true
+    val nbMembers = GroupMaxSize + 1
+
+    // First JoinRequests
+    var futures = 1.to(nbMembers).map { _ =>
+      EasyMock.reset(replicaManager)
+      sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols,
+        None, DefaultSessionTimeout, DefaultRebalanceTimeout, requiredKnownMemberId)
+    }
 
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols, sessionTimeout = ConsumerMaxSessionTimeout + 1)
-    val joinGroupError = joinGroupResult.error
-    assertEquals(Errors.INVALID_SESSION_TIMEOUT, joinGroupError)
-  }
+    // Get back the assigned member ids
+    val memberIds = futures.map(await(_, 1).memberId)
 
-  @Test
-  def testJoinGroupUnknownConsumerNewGroup() {
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
-    val joinGroupError = joinGroupResult.error
-    assertEquals(Errors.UNKNOWN_MEMBER_ID, joinGroupError)
-  }
+    // Second JoinRequests
+    futures = memberIds.map { memberId =>
+      EasyMock.reset(replicaManager)
+      sendJoinGroup(groupId, memberId, protocolType, protocols,
+        None, DefaultSessionTimeout, DefaultRebalanceTimeout, requiredKnownMemberId)
+    }
 
-  @Test
-  def testInvalidGroupId() {
-    val groupId = ""
-    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
+    // advance clock by GroupInitialRebalanceDelay to complete first InitialDelayedJoin
+    timer.advanceClock(GroupInitialRebalanceDelay + 1)
+    // advance clock by GroupInitialRebalanceDelay to complete second InitialDelayedJoin
+    timer.advanceClock(GroupInitialRebalanceDelay + 1)
 
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
-    assertEquals(Errors.INVALID_GROUP_ID, joinGroupResult.error)
-  }
+    // Awaiting results
+    val errors = futures.map(await(_, DefaultRebalanceTimeout + 1).error)
 
-  @Test
-  def testValidJoinGroup() {
-    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
+    assertEquals(GroupMaxSize, errors.count(_ == Errors.NONE))
+    assertEquals(nbMembers-GroupMaxSize, errors.count(_ == Errors.GROUP_MAX_SIZE_REACHED))
 
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
-    val joinGroupError = joinGroupResult.error
-    assertEquals(Errors.NONE, joinGroupError)
+    // Members which were accepted can rejoin, others are rejected, while
+    // completing rebalance
+    futures = memberIds.map { memberId =>
+      EasyMock.reset(replicaManager)
+      sendJoinGroup(groupId, memberId, protocolType, protocols,
+        None, DefaultSessionTimeout, DefaultRebalanceTimeout, requiredKnownMemberId)
+    }
+
+    // Awaiting results
+    val rejoinErrors = futures.map(await(_, 1).error)
+
+    assertEquals(errors, rejoinErrors)
   }
 
   @Test
-  def testJoinGroupInconsistentProtocolType() {
-    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
-    val otherMemberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
+  def testDynamicMembersJoinGroupWithMaxSize(): Unit = {
+    val requiredKnownMemberId = false
+    val nbMembers = GroupMaxSize + 1
+
+    // JoinRequests
+    var futures = 1.to(nbMembers).map { _ =>
+      EasyMock.reset(replicaManager)
+      sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols,
+        None, DefaultSessionTimeout, DefaultRebalanceTimeout, requiredKnownMemberId)
+    }
 
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
-    assertEquals(Errors.NONE, joinGroupResult.error)
+    // advance clock by GroupInitialRebalanceDelay to complete first InitialDelayedJoin
+    timer.advanceClock(GroupInitialRebalanceDelay + 1)
+    // advance clock by GroupInitialRebalanceDelay to complete second InitialDelayedJoin
+    timer.advanceClock(GroupInitialRebalanceDelay + 1)
 
-    EasyMock.reset(replicaManager)
-    val otherJoinGroupResult = await(sendJoinGroup(groupId, otherMemberId, "connect", protocols), 1)
-    assertEquals(Errors.INCONSISTENT_GROUP_PROTOCOL, otherJoinGroupResult.error)
-  }
+    // Awaiting results
+    val joinGroupResults = futures.map(await(_, DefaultRebalanceTimeout + 1))
+    val errors = joinGroupResults.map(_.error)
 
-  @Test
-  def testJoinGroupWithEmptyProtocolType() {
-    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
+    assertEquals(GroupMaxSize, errors.count(_ == Errors.NONE))
+    assertEquals(nbMembers-GroupMaxSize, errors.count(_ == Errors.GROUP_MAX_SIZE_REACHED))
 
-    val joinGroupResult = joinGroup(groupId, memberId, "", protocols)
-    assertEquals(Errors.INCONSISTENT_GROUP_PROTOCOL, joinGroupResult.error)
-  }
+    // Members which were accepted can rejoin, others are rejected, while
+    // completing rebalance
+    val memberIds = joinGroupResults.map(_.memberId)
+    futures = memberIds.map { memberId =>
+      EasyMock.reset(replicaManager)
+      sendJoinGroup(groupId, memberId, protocolType, protocols,
+        None, DefaultSessionTimeout, DefaultRebalanceTimeout, requiredKnownMemberId)
+    }
 
-  @Test
-  def testJoinGroupWithEmptyGroupProtocol() {
-    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
+    // Awaiting results
+    val rejoinErrors = futures.map(await(_, 1).error)
 
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, List())
-    assertEquals(Errors.INCONSISTENT_GROUP_PROTOCOL, joinGroupResult.error)
+    assertEquals(errors, rejoinErrors)
   }
 
   @Test
-  def testJoinGroupInconsistentGroupProtocol() {
-    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
+  def testStaticMembersJoinGroupWithMaxSize(): Unit = {
+    val nbMembers = GroupMaxSize + 1
+    val instanceIds = 1.to(nbMembers).map(i => Some(s"instance-id-$i"))
+
+    // JoinRequests
+    var futures = instanceIds.map { instanceId =>
+      EasyMock.reset(replicaManager)
+      sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols,
+        instanceId, DefaultSessionTimeout, DefaultRebalanceTimeout)
+    }
 
-    val otherMemberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
+    // advance clock by GroupInitialRebalanceDelay to complete first InitialDelayedJoin
+    timer.advanceClock(GroupInitialRebalanceDelay + 1)
+    // advance clock by GroupInitialRebalanceDelay to complete second InitialDelayedJoin
+    timer.advanceClock(GroupInitialRebalanceDelay + 1)
 
-    val joinGroupFuture = sendJoinGroup(groupId, memberId, protocolType, List(("range", metadata)))
+    // Awaiting results
+    val joinGroupResults = futures.map(await(_, DefaultRebalanceTimeout + 1))
+    val errors = joinGroupResults.map(_.error)
 
-    EasyMock.reset(replicaManager)
-    val otherJoinGroupResult = joinGroup(groupId, otherMemberId, protocolType, List(("roundrobin", metadata)))
+    assertEquals(GroupMaxSize, errors.count(_ == Errors.NONE))
+    assertEquals(nbMembers-GroupMaxSize, errors.count(_ == Errors.GROUP_MAX_SIZE_REACHED))
 
-    val joinGroupResult = await(joinGroupFuture, 1)
-    assertEquals(Errors.NONE, joinGroupResult.error)
-    assertEquals(Errors.INCONSISTENT_GROUP_PROTOCOL, otherJoinGroupResult.error)
+    // Members which were accepted can rejoin, others are rejected, while
+    // completing rebalance
+    val memberIds = joinGroupResults.map(_.memberId)
+    futures = instanceIds.zip(memberIds).map { case (instanceId, memberId) =>
+      EasyMock.reset(replicaManager)
+      sendJoinGroup(groupId, memberId, protocolType, protocols,
+        instanceId, DefaultSessionTimeout, DefaultRebalanceTimeout)
+    }
+
+    // Awaiting results
+    val rejoinErrors = futures.map(await(_, 1).error)
+
+    assertEquals(errors, rejoinErrors)
   }
 
   @Test
-  def testJoinGroupUnknownConsumerExistingGroup() {
-    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
-    val otherMemberId = "memberId"
+  def testDynamicMembersCanReJoinGroupWithMaxSizeWhileRebalancing(): Unit = {
+    val requiredKnownMemberId = true
+    val nbMembers = GroupMaxSize + 1
+
+    // First JoinRequests
+    var futures = 1.to(nbMembers).map { _ =>
+      EasyMock.reset(replicaManager)
+      sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols,
+        None, DefaultSessionTimeout, DefaultRebalanceTimeout, requiredKnownMemberId)
+    }
 
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
-    assertEquals(Errors.NONE, joinGroupResult.error)
+    // Get back the assigned member ids
+    val memberIds = futures.map(await(_, 1).memberId)
 
-    EasyMock.reset(replicaManager)
-    val otherJoinGroupResult = await(sendJoinGroup(groupId, otherMemberId, protocolType, protocols), 1)
-    assertEquals(Errors.UNKNOWN_MEMBER_ID, otherJoinGroupResult.error)
-  }
+    // Second JoinRequests
+    memberIds.map { memberId =>
+      EasyMock.reset(replicaManager)
+      sendJoinGroup(groupId, memberId, protocolType, protocols,
+        None, DefaultSessionTimeout, DefaultRebalanceTimeout, requiredKnownMemberId)
+    }
 
-  @Test
-  def testHeartbeatWrongCoordinator() {
+    // Members can rejoin while rebalancing
+    futures = memberIds.map { memberId =>
+      EasyMock.reset(replicaManager)
+      sendJoinGroup(groupId, memberId, protocolType, protocols,
+        None, DefaultSessionTimeout, DefaultRebalanceTimeout, requiredKnownMemberId)
+    }
 
-    val heartbeatResult = heartbeat(otherGroupId, memberId, -1)
-    assertEquals(Errors.NOT_COORDINATOR, heartbeatResult)
-  }
+    // advance clock by GroupInitialRebalanceDelay to complete first InitialDelayedJoin
+    timer.advanceClock(GroupInitialRebalanceDelay + 1)
+    // advance clock by GroupInitialRebalanceDelay to complete second InitialDelayedJoin
+    timer.advanceClock(GroupInitialRebalanceDelay + 1)
 
-  @Test
-  def testHeartbeatUnknownGroup() {
+    // Awaiting results
+    val errors = futures.map(await(_, DefaultRebalanceTimeout + 1).error)
 
-    val heartbeatResult = heartbeat(groupId, memberId, -1)
-    assertEquals(Errors.UNKNOWN_MEMBER_ID, heartbeatResult)
+    assertEquals(GroupMaxSize, errors.count(_ == Errors.NONE))
+    assertEquals(nbMembers-GroupMaxSize, errors.count(_ == Errors.GROUP_MAX_SIZE_REACHED))
   }
 
   @Test
-  def testHeartbeatUnknownConsumerExistingGroup() {
-    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
-    val otherMemberId = "memberId"
+  def testLastJoiningMembersAreKickedOutWhenReJoiningGroupWithMaxSize(): Unit = {
+    val nbMembers = GroupMaxSize + 2
+    val group = new GroupMetadata(groupId, Stable, new MockTime())
+    val memberIds = 1.to(nbMembers).map(_ => group.generateMemberId(ClientId, None))
+
+    memberIds.foreach { memberId =>
+      group.add(new MemberMetadata(memberId, groupId, None, ClientId, ClientHost,
+        DefaultRebalanceTimeout, GroupMaxSessionTimeout, protocolType, protocols))
+    }
+    groupCoordinator.groupManager.addGroup(group)
 
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
-    val assignedMemberId = joinGroupResult.memberId
-    val joinGroupError = joinGroupResult.error
-    assertEquals(Errors.NONE, joinGroupError)
+    groupCoordinator.prepareRebalance(group, "")
 
-    EasyMock.reset(replicaManager)
-    val syncGroupResult = syncGroupLeader(groupId, joinGroupResult.generationId, assignedMemberId, Map(assignedMemberId -> Array[Byte]()))
-    val syncGroupError = syncGroupResult._2
-    assertEquals(Errors.NONE, syncGroupError)
+    val futures = memberIds.map { memberId =>
+      EasyMock.reset(replicaManager)
+      sendJoinGroup(groupId, memberId, protocolType, protocols,
+        None, GroupMaxSessionTimeout, DefaultRebalanceTimeout)
+    }
 
-    EasyMock.reset(replicaManager)
-    val heartbeatResult = heartbeat(groupId, otherMemberId, 1)
-    assertEquals(Errors.UNKNOWN_MEMBER_ID, heartbeatResult)
+    // advance clock by GroupInitialRebalanceDelay to complete first InitialDelayedJoin
+    timer.advanceClock(DefaultRebalanceTimeout + 1)
+
+    // Awaiting results
+    val errors = futures.map(await(_, DefaultRebalanceTimeout + 1).error)
+
+    assertEquals(Set(Errors.NONE), errors.take(GroupMaxSize).toSet)
+    assertEquals(Set(Errors.GROUP_MAX_SIZE_REACHED), errors.drop(GroupMaxSize).toSet)
+
+    memberIds.drop(GroupMaxSize).foreach { memberId =>
+      assertFalse(group.has(memberId))
+    }
   }
 
   @Test
-  def testHeartbeatRebalanceInProgress() {
+  def testJoinGroupSessionTimeoutTooSmall(): Unit = {
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
 
-    val joinGroupResult = joinGroup(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)
-    assertEquals(Errors.REBALANCE_IN_PROGRESS, heartbeatResult)
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols, sessionTimeout = GroupMinSessionTimeout - 1)
+    assertEquals(Errors.INVALID_SESSION_TIMEOUT, joinGroupResult.error)
   }
 
   @Test
-  def testHeartbeatIllegalGeneration() {
+  def testJoinGroupSessionTimeoutTooLarge(): Unit = {
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
 
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
-    val assignedMemberId = joinGroupResult.memberId
-    val joinGroupError = joinGroupResult.error
-    assertEquals(Errors.NONE, joinGroupError)
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols, sessionTimeout = GroupMaxSessionTimeout + 1)
+    assertEquals(Errors.INVALID_SESSION_TIMEOUT, joinGroupResult.error)
+  }
 
-    EasyMock.reset(replicaManager)
-    val syncGroupResult = syncGroupLeader(groupId, joinGroupResult.generationId, assignedMemberId, Map(assignedMemberId -> Array[Byte]()))
-    val syncGroupError = syncGroupResult._2
-    assertEquals(Errors.NONE, syncGroupError)
+  @Test
+  def testJoinGroupUnknownConsumerNewGroup(): Unit = {
+    var joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
+    assertEquals(Errors.UNKNOWN_MEMBER_ID, joinGroupResult.error)
 
     EasyMock.reset(replicaManager)
-    val heartbeatResult = heartbeat(groupId, assignedMemberId, 2)
-    assertEquals(Errors.ILLEGAL_GENERATION, heartbeatResult)
+    joinGroupResult = staticJoinGroup(groupId, memberId, groupInstanceId, protocolType, protocols)
+    assertEquals(Errors.UNKNOWN_MEMBER_ID, joinGroupResult.error)
   }
 
   @Test
-  def testValidHeartbeat() {
+  def testInvalidGroupId(): Unit = {
+    val groupId = ""
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
 
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
-    val assignedConsumerId = joinGroupResult.memberId
-    val generationId = joinGroupResult.generationId
-    val joinGroupError = joinGroupResult.error
-    assertEquals(Errors.NONE, joinGroupError)
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
+    assertEquals(Errors.INVALID_GROUP_ID, joinGroupResult.error)
+  }
 
-    EasyMock.reset(replicaManager)
-    val syncGroupResult = syncGroupLeader(groupId, generationId, assignedConsumerId, Map(assignedConsumerId -> Array[Byte]()))
-    val syncGroupError = syncGroupResult._2
-    assertEquals(Errors.NONE, syncGroupError)
+  @Test
+  def testValidJoinGroup(): Unit = {
+    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
 
-    EasyMock.reset(replicaManager)
-    val heartbeatResult = heartbeat(groupId, assignedConsumerId, 1)
-    assertEquals(Errors.NONE, heartbeatResult)
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
+    assertEquals(Errors.NONE, joinGroupResult.error)
   }
 
   @Test
-  def testSessionTimeout() {
+  def testJoinGroupInconsistentProtocolType(): Unit = {
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
+    val otherMemberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
 
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
-    val assignedConsumerId = joinGroupResult.memberId
-    val generationId = joinGroupResult.generationId
-    val joinGroupError = joinGroupResult.error
-    assertEquals(Errors.NONE, joinGroupError)
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
+    assertEquals(Errors.NONE, joinGroupResult.error)
 
     EasyMock.reset(replicaManager)
-    val (_, syncGroupError) = syncGroupLeader(groupId, generationId, assignedConsumerId, Map(assignedConsumerId -> Array[Byte]()))
-    assertEquals(Errors.NONE, syncGroupError)
+    val otherJoinGroupResult = await(sendJoinGroup(groupId, otherMemberId, "connect", protocols), 1)
+    assertEquals(Errors.INCONSISTENT_GROUP_PROTOCOL, otherJoinGroupResult.error)
+  }
 
-    EasyMock.reset(replicaManager)
-    EasyMock.expect(replicaManager.getPartition(new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, groupPartitionId))).andReturn(None)
-    EasyMock.expect(replicaManager.getMagic(EasyMock.anyObject())).andReturn(Some(RecordBatch.MAGIC_VALUE_V1)).anyTimes()
-    EasyMock.replay(replicaManager)
+  @Test
+  def testJoinGroupWithEmptyProtocolType(): Unit = {
+    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
 
-    timer.advanceClock(DefaultSessionTimeout + 100)
+    var joinGroupResult = dynamicJoinGroup(groupId, memberId, "", protocols)
+    assertEquals(Errors.INCONSISTENT_GROUP_PROTOCOL, joinGroupResult.error)
 
     EasyMock.reset(replicaManager)
-    val heartbeatResult = heartbeat(groupId, assignedConsumerId, 1)
-    assertEquals(Errors.UNKNOWN_MEMBER_ID, heartbeatResult)
+    joinGroupResult = staticJoinGroup(groupId, memberId, groupInstanceId, "", protocols)
+    assertEquals(Errors.INCONSISTENT_GROUP_PROTOCOL, joinGroupResult.error)
   }
 
   @Test
-  def testHeartbeatMaintainsSession() {
+  def testJoinGroupWithEmptyGroupProtocol(): Unit = {
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
-    val sessionTimeout = 1000
 
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols,
-      rebalanceTimeout = sessionTimeout, sessionTimeout = sessionTimeout)
-    val assignedConsumerId = joinGroupResult.memberId
-    val generationId = joinGroupResult.generationId
-    val joinGroupError = joinGroupResult.error
-    assertEquals(Errors.NONE, joinGroupError)
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, List())
+    assertEquals(Errors.INCONSISTENT_GROUP_PROTOCOL, joinGroupResult.error)
+  }
 
-    EasyMock.reset(replicaManager)
-    val (_, syncGroupError) = syncGroupLeader(groupId, generationId, assignedConsumerId, Map(assignedConsumerId -> Array[Byte]()))
-    assertEquals(Errors.NONE, syncGroupError)
+  @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(sessionTimeout / 2)
+    timer.advanceClock(GroupInitialRebalanceDelay + 1)
 
-    EasyMock.reset(replicaManager)
-    var heartbeatResult = heartbeat(groupId, assignedConsumerId, 1)
-    assertEquals(Errors.NONE, heartbeatResult)
+    val joinResult = Await.result(responseFuture, Duration(DefaultRebalanceTimeout + 100, TimeUnit.MILLISECONDS))
+    val group = groupCoordinator.groupManager.getGroup(groupId).get
+    val memberId = joinResult.memberId
 
-    timer.advanceClock(sessionTimeout / 2 + 100)
+    assertEquals(Errors.NONE, joinResult.error)
+    assertEquals(0, group.allMemberMetadata.count(_.isNew))
 
     EasyMock.reset(replicaManager)
-    heartbeatResult = heartbeat(groupId, assignedConsumerId, 1)
-    assertEquals(Errors.NONE, heartbeatResult)
+    val syncGroupResult = syncGroupLeader(groupId, joinResult.generationId, memberId, Map(memberId -> Array[Byte]()))
+    assertEquals(Errors.NONE, syncGroupResult.error)
+    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 testCommitMaintainsSession() {
-    val sessionTimeout = 1000
-    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
-    val tp = new TopicPartition("topic", 0)
-    val offset = OffsetAndMetadata(0)
+  def testNewMemberJoinExpiration(): Unit = {
+    // This tests new member expiration during a protracted rebalance. We first create a
+    // group with one member which uses a large value for session timeout and rebalance timeout.
+    // We then join with one new member and let the rebalance hang while we await the first member.
+    // The new member join timeout expires and its JoinGroup request is failed.
 
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols,
-      rebalanceTimeout = sessionTimeout, sessionTimeout = sessionTimeout)
-    val assignedConsumerId = joinGroupResult.memberId
-    val generationId = joinGroupResult.generationId
-    val joinGroupError = joinGroupResult.error
-    assertEquals(Errors.NONE, joinGroupError)
+    val sessionTimeout = GroupCoordinator.NewMemberJoinTimeoutMs + 5000
+    val rebalanceTimeout = GroupCoordinator.NewMemberJoinTimeoutMs * 2
 
-    EasyMock.reset(replicaManager)
-    val (_, syncGroupError) = syncGroupLeader(groupId, generationId, assignedConsumerId, Map(assignedConsumerId -> Array[Byte]()))
-    assertEquals(Errors.NONE, syncGroupError)
+    val firstJoinResult = dynamicJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols,
+      sessionTimeout, rebalanceTimeout)
+    val firstMemberId = firstJoinResult.memberId
+    assertEquals(firstMemberId, firstJoinResult.leaderId)
+    assertEquals(Errors.NONE, firstJoinResult.error)
 
-    timer.advanceClock(sessionTimeout / 2)
+    val groupOpt = groupCoordinator.groupManager.getGroup(groupId)
+    assertTrue(groupOpt.isDefined)
+    val group = groupOpt.get
+    assertEquals(0, group.allMemberMetadata.count(_.isNew))
 
     EasyMock.reset(replicaManager)
-    val commitOffsetResult = commitOffsets(groupId, assignedConsumerId, generationId, immutable.Map(tp -> offset))
-    assertEquals(Errors.NONE, commitOffsetResult(tp))
+    val responseFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols, None, sessionTimeout, rebalanceTimeout)
+    assertFalse(responseFuture.isCompleted)
 
-    timer.advanceClock(sessionTimeout / 2 + 100)
+    assertEquals(2, group.allMembers.size)
+    assertEquals(1, group.allMemberMetadata.count(_.isNew))
 
-    EasyMock.reset(replicaManager)
-    val heartbeatResult = heartbeat(groupId, assignedConsumerId, 1)
-    assertEquals(Errors.NONE, heartbeatResult)
+    val newMember = group.allMemberMetadata.find(_.isNew).get
+    assertNotEquals(firstMemberId, newMember.memberId)
+
+    timer.advanceClock(GroupCoordinator.NewMemberJoinTimeoutMs + 1)
+    assertTrue(responseFuture.isCompleted)
+
+    val response = Await.result(responseFuture, Duration(0, TimeUnit.MILLISECONDS))
+    assertEquals(Errors.UNKNOWN_MEMBER_ID, response.error)
+    assertEquals(1, group.allMembers.size)
+    assertEquals(0, group.allMemberMetadata.count(_.isNew))
+    assertEquals(firstMemberId, group.allMembers.head)
   }
 
   @Test
-  def testSessionTimeoutDuringRebalance() {
-    // create a group with a single member
-    val firstJoinResult = joinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols,
-      rebalanceTimeout = 2000, sessionTimeout = 1000)
+  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)
+    val firstSyncResult = syncGroupLeader(groupId, firstGenerationId, firstMemberId,
+      Map(firstMemberId -> Array[Byte]()))
+    assertEquals(Errors.NONE, firstSyncResult.error)
 
-    // now have a new member join to trigger a rebalance
     EasyMock.reset(replicaManager)
     val otherJoinFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
 
-    timer.advanceClock(500)
-
-    EasyMock.reset(replicaManager)
-    var heartbeatResult = heartbeat(groupId, firstMemberId, firstGenerationId)
-    assertEquals(Errors.REBALANCE_IN_PROGRESS, heartbeatResult)
-
-    // letting the session expire should make the member fall out of the group
-    timer.advanceClock(1100)
-
     EasyMock.reset(replicaManager)
-    heartbeatResult = heartbeat(groupId, firstMemberId, firstGenerationId)
-    assertEquals(Errors.UNKNOWN_MEMBER_ID, heartbeatResult)
+    val joinFuture = sendJoinGroup(groupId, firstMemberId, protocolType, protocols,
+      requireKnownMemberId = false)
 
-    // and the rebalance should complete with only the new member
+    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 testRebalanceCompletesBeforeMemberJoins() {
-    // create a group with a single member
-    val firstJoinResult = joinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols,
-      rebalanceTimeout = 1200, sessionTimeout = 1000)
+  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)
+    val firstSyncResult = syncGroupLeader(groupId, firstGenerationId, firstMemberId,
+      Map(firstMemberId -> Array[Byte]()))
+    assertEquals(Errors.NONE, firstSyncResult.error)
 
-    // now have a new member join to trigger a rebalance
     EasyMock.reset(replicaManager)
     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)
+    val joinFuture = sendJoinGroup(groupId, firstMemberId, protocolType, protocols,
+      requireKnownMemberId = false)
 
-    // now timeout the rebalance, which should kick the unjoined member out of the group
-    // and let the rebalance finish with only the new member
-    timer.advanceClock(500)
+    val joinResult = await(joinFuture, DefaultSessionTimeout+100)
     val otherJoinResult = await(otherJoinFuture, DefaultSessionTimeout+100)
+    assertEquals(Errors.NONE, joinResult.error)
     assertEquals(Errors.NONE, otherJoinResult.error)
-  }
-
-  @Test
-  def testSyncGroupEmptyAssignment() {
-    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
-
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
-    val assignedConsumerId = joinGroupResult.memberId
-    val generationId = joinGroupResult.generationId
-    val joinGroupError = joinGroupResult.error
-    assertEquals(Errors.NONE, joinGroupError)
+    val secondGenerationId = joinResult.generationId
+    val secondMemberId = otherJoinResult.memberId
 
     EasyMock.reset(replicaManager)
-    val syncGroupResult = syncGroupLeader(groupId, generationId, assignedConsumerId, Map())
-    val syncGroupError = syncGroupResult._2
-    assertEquals(Errors.NONE, syncGroupError)
-    assertTrue(syncGroupResult._1.isEmpty)
+    sendSyncGroupFollower(groupId, secondGenerationId, secondMemberId)
 
     EasyMock.reset(replicaManager)
-    val heartbeatResult = heartbeat(groupId, assignedConsumerId, 1)
-    assertEquals(Errors.NONE, heartbeatResult)
+    val syncGroupResult = syncGroupLeader(groupId, secondGenerationId, firstMemberId,
+      Map(firstMemberId -> Array.emptyByteArray, secondMemberId -> Array.emptyByteArray))
+    assertEquals(Errors.NONE, syncGroupResult.error)
+
+    verifySessionExpiration(groupId)
   }
 
-  @Test
-  def testSyncGroupNotCoordinator() {
-    val generation = 1
+  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)
 
-    val syncGroupResult = syncGroupFollower(otherGroupId, generation, memberId)
-    assertEquals(Errors.NOT_COORDINATOR, syncGroupResult._2)
+    timer.advanceClock(DefaultSessionTimeout + 1)
+
+    val groupMetadata = group(groupId)
+    assertEquals(Empty, groupMetadata.currentState)
+    assertTrue(groupMetadata.allMembers.isEmpty)
   }
 
   @Test
-  def testSyncGroupFromUnknownGroup() {
-    val generation = 1
+  def testJoinGroupInconsistentGroupProtocol(): Unit = {
+    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
+    val otherMemberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
 
-    val syncGroupResult = syncGroupFollower(groupId, generation, memberId)
-    assertEquals(Errors.UNKNOWN_MEMBER_ID, syncGroupResult._2)
+    val joinGroupFuture = sendJoinGroup(groupId, memberId, protocolType, List(("range", metadata)))
+
+    EasyMock.reset(replicaManager)
+    val otherJoinGroupResult = dynamicJoinGroup(groupId, otherMemberId, protocolType, List(("roundrobin", metadata)))
+    timer.advanceClock(GroupInitialRebalanceDelay + 1)
+
+    val joinGroupResult = await(joinGroupFuture, 1)
+    assertEquals(Errors.NONE, joinGroupResult.error)
+    assertEquals(Errors.INCONSISTENT_GROUP_PROTOCOL, otherJoinGroupResult.error)
   }
 
   @Test
-  def testSyncGroupFromUnknownMember() {
+  def testJoinGroupUnknownConsumerExistingGroup(): Unit = {
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
+    val otherMemberId = "memberId"
 
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
-    val assignedConsumerId = joinGroupResult.memberId
-    val generationId = joinGroupResult.generationId
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
     assertEquals(Errors.NONE, joinGroupResult.error)
 
     EasyMock.reset(replicaManager)
-    val syncGroupResult = syncGroupLeader(groupId, generationId, assignedConsumerId, Map(assignedConsumerId -> Array[Byte]()))
-    val syncGroupError = syncGroupResult._2
-    assertEquals(Errors.NONE, syncGroupError)
+    val otherJoinGroupResult = await(sendJoinGroup(groupId, otherMemberId, protocolType, protocols), 1)
+    assertEquals(Errors.UNKNOWN_MEMBER_ID, otherJoinGroupResult.error)
+  }
+
+  @Test
+  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 syncGroupResult = syncGroupFollower(deadGroupId, 1, memberId)
+    assertEquals(Errors.COORDINATOR_NOT_AVAILABLE, syncGroupResult.error)
+  }
+
+  @Test
+  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)
+    val memberId = joinGroupResult.memberId
 
+    // Sending an inconsistent protocol shall be refused
     EasyMock.reset(replicaManager)
-    val unknownMemberId = "blah"
-    val unknownMemberSyncResult = syncGroupFollower(groupId, generationId, unknownMemberId)
-    assertEquals(Errors.UNKNOWN_MEMBER_ID, unknownMemberSyncResult._2)
+    responseFuture = sendJoinGroup(groupId, memberId, protocolType, List(), requireKnownMemberId = true)
+    joinGroupResult = Await.result(responseFuture, Duration(DefaultRebalanceTimeout + 1, TimeUnit.MILLISECONDS))
+    assertEquals(Errors.INCONSISTENT_GROUP_PROTOCOL, joinGroupResult.error)
+
+    // Sending consistent protocol shall be accepted
+    EasyMock.reset(replicaManager)
+    responseFuture = sendJoinGroup(groupId, memberId, protocolType, protocols, requireKnownMemberId = true)
+    timer.advanceClock(GroupInitialRebalanceDelay + 1)
+    joinGroupResult = Await.result(responseFuture, Duration(DefaultRebalanceTimeout + 1, TimeUnit.MILLISECONDS))
+    assertEquals(Errors.NONE, joinGroupResult.error)
   }
 
   @Test
-  def testSyncGroupFromIllegalGeneration() {
-    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
+  def staticMemberJoinAsFirstMember(): Unit = {
+    val joinGroupResult = staticJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, groupInstanceId, protocolType, protocols)
+    assertEquals(Errors.NONE, joinGroupResult.error)
+  }
 
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
-    val assignedConsumerId = joinGroupResult.memberId
-    val generationId = joinGroupResult.generationId
+  @Test
+  def staticMemberReJoinWithExplicitUnknownMemberId(): Unit = {
+    var joinGroupResult = staticJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, groupInstanceId, protocolType, protocols)
     assertEquals(Errors.NONE, joinGroupResult.error)
 
     EasyMock.reset(replicaManager)
-    // send the sync group with an invalid generation
-    val syncGroupResult = syncGroupLeader(groupId, generationId+1, assignedConsumerId, Map(assignedConsumerId -> Array[Byte]()))
-    assertEquals(Errors.ILLEGAL_GENERATION, syncGroupResult._2)
+    val unknownMemberId = "unknown_member"
+    joinGroupResult = staticJoinGroup(groupId, unknownMemberId, groupInstanceId, protocolType, protocols)
+    assertEquals(Errors.FENCED_INSTANCE_ID, joinGroupResult.error)
   }
 
   @Test
-  def testJoinGroupFromUnchangedFollowerDoesNotRebalance() {
-    // 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
+  def staticMemberFenceDuplicateRejoinedFollower(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
 
-    val firstJoinResult = joinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
-    val firstMemberId = firstJoinResult.memberId
-    val firstGenerationId = firstJoinResult.generationId
-    assertEquals(firstMemberId, firstJoinResult.leaderId)
-    assertEquals(Errors.NONE, firstJoinResult.error)
+    // 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)
-    val firstSyncResult = syncGroupLeader(groupId, firstGenerationId, firstMemberId, Map(firstMemberId -> Array[Byte]()))
-    assertEquals(Errors.NONE, firstSyncResult._2)
+    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)
-    val otherJoinFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
+    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)
 
-    EasyMock.reset(replicaManager)
-    val joinFuture = sendJoinGroup(groupId, firstMemberId, protocolType, protocols)
+    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,
+      None)
+    verifyDelayedTaskNotCompleted(duplicateFollowerJoinFuture)
+  }
 
-    val joinResult = await(joinFuture, DefaultSessionTimeout+100)
-    val otherJoinResult = await(otherJoinFuture, DefaultSessionTimeout+100)
-    assertEquals(Errors.NONE, joinResult.error)
-    assertEquals(Errors.NONE, otherJoinResult.error)
-    assertTrue(joinResult.generationId == otherJoinResult.generationId)
+  @Test
+  def staticMemberFenceDuplicateSyncingFollowerAfterMemberIdChanged(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
 
-    assertEquals(firstMemberId, joinResult.leaderId)
-    assertEquals(firstMemberId, otherJoinResult.leaderId)
+    // Known leader rejoins will trigger rebalance.
+    val leaderJoinGroupFuture =
+      sendJoinGroup(groupId, rebalanceResult.leaderId, protocolType, protocols, groupInstanceId = leaderInstanceId)
+    timer.advanceClock(1)
+    assertTrue(getGroup(groupId).is(PreparingRebalance))
 
-    val nextGenerationId = joinResult.generationId
+    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)
 
-    // this shouldn't cause a rebalance since protocol information hasn't changed
+    timer.advanceClock(1)
+    val leaderJoinGroupResult = Await.result(leaderJoinGroupFuture, Duration(1, TimeUnit.MILLISECONDS))
+    checkJoinGroupResult(leaderJoinGroupResult,
+      Errors.NONE,
+      rebalanceResult.generation + 1,
+      Set(leaderInstanceId, followerInstanceId),
+      groupId,
+      CompletingRebalance,
+      Some(protocolType))
+    assertEquals(rebalanceResult.leaderId, leaderJoinGroupResult.memberId)
+    assertEquals(rebalanceResult.leaderId, leaderJoinGroupResult.leaderId)
+
+    // Old follower 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,
+      Some(protocolType),
+      expectedLeaderId = leaderJoinGroupResult.memberId)
+    assertEquals(rebalanceResult.followerId, oldFollowerJoinGroupResult.memberId)
+    assertEquals(rebalanceResult.leaderId, oldFollowerJoinGroupResult.leaderId)
+    assertTrue(getGroup(groupId).is(CompletingRebalance))
+
+    // Duplicate follower joins group with unknown member id will trigger member.id replacement,
+    // and will also trigger a rebalance under CompletingRebalance state; the old follower sync callback
+    // will return fenced exception while broker replaces the member identity with the duplicate follower joins.
     EasyMock.reset(replicaManager)
-    val followerJoinResult = await(sendJoinGroup(groupId, otherJoinResult.memberId, protocolType, protocols), 1)
+    val oldFollowerSyncGroupFuture = sendSyncGroupFollower(groupId, oldFollowerJoinGroupResult.generationId,
+      oldFollowerJoinGroupResult.memberId, Some(protocolType), Some(protocolName), followerInstanceId)
 
-    assertEquals(Errors.NONE, followerJoinResult.error)
-    assertEquals(nextGenerationId, followerJoinResult.generationId)
+    EasyMock.reset(replicaManager)
+    val duplicateFollowerJoinFuture =
+      sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols, groupInstanceId = followerInstanceId)
+    timer.advanceClock(1)
+
+    val oldFollowerSyncGroupResult = Await.result(oldFollowerSyncGroupFuture, Duration(1, TimeUnit.MILLISECONDS))
+    assertEquals(Errors.FENCED_INSTANCE_ID, oldFollowerSyncGroupResult.error)
+    assertTrue(getGroup(groupId).is(PreparingRebalance))
+
+    timer.advanceClock(GroupInitialRebalanceDelay + 1)
+    timer.advanceClock(DefaultRebalanceTimeout + 1)
+
+    val duplicateFollowerJoinGroupResult = Await.result(duplicateFollowerJoinFuture, Duration(1, TimeUnit.MILLISECONDS))
+    checkJoinGroupResult(duplicateFollowerJoinGroupResult,
+      Errors.NONE,
+      rebalanceResult.generation + 2,
+      Set(followerInstanceId),   // this follower will become the new leader, and hence it would have the member list
+      groupId,
+      CompletingRebalance,
+      Some(protocolType),
+      expectedLeaderId = duplicateFollowerJoinGroupResult.memberId)
+    assertTrue(getGroup(groupId).is(CompletingRebalance))
   }
 
   @Test
-  def testJoinGroupFromUnchangedLeaderShouldRebalance() {
-    val firstJoinResult = joinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
-    val firstMemberId = firstJoinResult.memberId
-    val firstGenerationId = firstJoinResult.generationId
-    assertEquals(firstMemberId, firstJoinResult.leaderId)
-    assertEquals(Errors.NONE, firstJoinResult.error)
+  def staticMemberFenceDuplicateRejoiningFollowerAfterMemberIdChanged(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
+
+    // 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)
-    val firstSyncResult = syncGroupLeader(groupId, firstGenerationId, firstMemberId, Map(firstMemberId -> Array[Byte]()))
-    assertEquals(Errors.NONE, firstSyncResult._2)
+    // Duplicate follower joins group will trigger member.id replacement.
+    val duplicateFollowerJoinGroupFuture =
+      sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols, groupInstanceId = followerInstanceId)
 
-    // join groups from the leader should force the group to rebalance, which allows the
-    // leader to push new assignments when local metadata changes
+    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,
+      Some(protocolType))
+
+    val duplicateFollowerJoinGroupResult = Await.result(duplicateFollowerJoinGroupFuture, Duration(1, TimeUnit.MILLISECONDS))
+    checkJoinGroupResult(duplicateFollowerJoinGroupResult,
+      Errors.NONE,
+      rebalanceResult.generation + 1,
+      Set.empty,
+      groupId,
+      CompletingRebalance,
+      Some(protocolType))
+    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,
+      None)
+  }
+
+  @Test
+  def staticMemberRejoinWithKnownMemberId(): Unit = {
+    var joinGroupResult = staticJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, groupInstanceId, protocolType, protocols)
+    assertEquals(Errors.NONE, joinGroupResult.error)
 
     EasyMock.reset(replicaManager)
-    val secondJoinResult = await(sendJoinGroup(groupId, firstMemberId, protocolType, protocols), 1)
+    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))
 
-    assertEquals(Errors.NONE, secondJoinResult.error)
-    assertNotEquals(firstGenerationId, secondJoinResult.generationId)
+    EasyMock.reset(replicaManager)
+    val syncGroupFuture = sendSyncGroupLeader(groupId, joinGroupResult.generationId, assignedMemberId, Some(protocolType), Some(protocolName), groupInstanceId, Map(assignedMemberId -> Array[Byte]()))
+    timer.advanceClock(1)
+    val syncGroupResult = Await.result(syncGroupFuture, Duration(1, TimeUnit.MILLISECONDS))
+    assertEquals(Errors.NONE, syncGroupResult.error)
+    assertTrue(getGroup(groupId).is(Stable))
   }
 
   @Test
-  def testLeaderFailureInSyncGroup() {
-    // 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
+  def staticMemberRejoinWithLeaderIdAndUnknownMemberId(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
 
-    val firstJoinResult = joinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
-    val firstMemberId = firstJoinResult.memberId
-    val firstGenerationId = firstJoinResult.generationId
-    assertEquals(firstMemberId, firstJoinResult.leaderId)
-    assertEquals(Errors.NONE, firstJoinResult.error)
+    // A static leader rejoin with unknown id will not trigger rebalance, and no assignment will be returned.
+    val joinGroupResult = staticJoinGroupWithPersistence(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,
+      Some(protocolType),
+      rebalanceResult.leaderId)
 
     EasyMock.reset(replicaManager)
-    val firstSyncResult = syncGroupLeader(groupId, firstGenerationId, firstMemberId, Map(firstMemberId -> Array[Byte]()))
-    assertEquals(Errors.NONE, firstSyncResult._2)
+    val oldLeaderJoinGroupResult = staticJoinGroup(groupId, rebalanceResult.leaderId, leaderInstanceId, protocolType, protocolSuperset, clockAdvance = 1)
+    assertEquals(Errors.FENCED_INSTANCE_ID, oldLeaderJoinGroupResult.error)
 
     EasyMock.reset(replicaManager)
-    val otherJoinFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
+    // Old leader will get fenced.
+    val oldLeaderSyncGroupResult = syncGroupLeader(groupId, rebalanceResult.generation, rebalanceResult.leaderId, Map.empty, None, None, leaderInstanceId)
+    assertEquals(Errors.FENCED_INSTANCE_ID, oldLeaderSyncGroupResult.error)
 
+    // Calling sync on old leader.id will fail because that leader.id is no longer valid and replaced.
     EasyMock.reset(replicaManager)
-    val joinFuture = sendJoinGroup(groupId, firstMemberId, protocolType, protocols)
+    val newLeaderSyncGroupResult = syncGroupLeader(groupId, rebalanceResult.generation, joinGroupResult.leaderId, Map.empty)
+    assertEquals(Errors.UNKNOWN_MEMBER_ID, newLeaderSyncGroupResult.error)
+  }
 
-    val joinResult = await(joinFuture, DefaultSessionTimeout+100)
-    val otherJoinResult = await(otherJoinFuture, DefaultSessionTimeout+100)
-    assertEquals(Errors.NONE, joinResult.error)
-    assertEquals(Errors.NONE, otherJoinResult.error)
-    assertTrue(joinResult.generationId == otherJoinResult.generationId)
+  @Test
+  def staticMemberRejoinWithLeaderIdAndKnownMemberId(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId,
+      sessionTimeout = DefaultRebalanceTimeout / 2)
+
+    // A static leader with known id rejoin will trigger rebalance.
+    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,
+      Some(protocolType),
+      rebalanceResult.leaderId,
+      rebalanceResult.leaderId)
+  }
 
-    assertEquals(firstMemberId, joinResult.leaderId)
-    assertEquals(firstMemberId, otherJoinResult.leaderId)
+  @Test
+  def staticMemberRejoinWithLeaderIdAndUnexpectedDeadGroup(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
 
-    val nextGenerationId = joinResult.generationId
+    getGroup(groupId).transitionTo(Dead)
 
-    // with no leader SyncGroup, the follower's request should failure with an error indicating
-    // that it should rejoin
-    EasyMock.reset(replicaManager)
-    val followerSyncFuture = sendSyncGroupFollower(groupId, nextGenerationId, otherJoinResult.memberId)
+    val joinGroupResult = staticJoinGroup(groupId, rebalanceResult.leaderId, leaderInstanceId, protocolType, protocols, clockAdvance = 1)
+    assertEquals(Errors.COORDINATOR_NOT_AVAILABLE, joinGroupResult.error)
+  }
 
-    timer.advanceClock(DefaultSessionTimeout + 100)
+  @Test
+  def staticMemberRejoinWithLeaderIdAndUnexpectedEmptyGroup(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
 
-    val followerSyncResult = await(followerSyncFuture, DefaultSessionTimeout+100)
-    assertEquals(Errors.REBALANCE_IN_PROGRESS, followerSyncResult._2)
+    getGroup(groupId).transitionTo(PreparingRebalance)
+    getGroup(groupId).transitionTo(Empty)
+
+    val joinGroupResult = staticJoinGroup(groupId, rebalanceResult.leaderId, leaderInstanceId, protocolType, protocols, clockAdvance = 1)
+    assertEquals(Errors.UNKNOWN_MEMBER_ID, joinGroupResult.error)
   }
 
   @Test
-  def testSyncGroupFollowerAfterLeader() {
-    // 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
+  def staticMemberRejoinWithFollowerIdAndChangeOfProtocol(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId, sessionTimeout = DefaultSessionTimeout * 2)
+
+    // A static follower rejoin with changed protocol will trigger rebalance.
+    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,
+      Some(protocolType),
+      rebalanceResult.followerId,
+      rebalanceResult.followerId)
+  }
 
-    val firstJoinResult = joinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
-    val firstMemberId = firstJoinResult.memberId
-    val firstGenerationId = firstJoinResult.generationId
-    assertEquals(firstMemberId, firstJoinResult.leaderId)
-    assertEquals(Errors.NONE, firstJoinResult.error)
+  @Test
+  def staticMemberRejoinWithUnknownMemberIdAndChangeOfProtocolWithSelectedProtocolChanged(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
+
+    // A static follower rejoin with protocol changed and also cause updated group's selectedProtocol changed
+    // should trigger rebalance.
+    val selectedProtocols = getGroup(groupId).selectProtocol
+    val newProtocols = List(("roundrobin", metadata))
+    assert(!newProtocols.map(_._1).contains(selectedProtocols))
+    // Old leader hasn't joined in the meantime, triggering a re-election.
+    val joinGroupResult = staticJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, followerInstanceId, protocolType, newProtocols, clockAdvance = DefaultSessionTimeout + 1)
+
+    checkJoinGroupResult(joinGroupResult,
+      Errors.NONE,
+      rebalanceResult.generation + 1,
+      Set(leaderInstanceId, followerInstanceId),
+      groupId,
+      CompletingRebalance,
+      Some(protocolType))
+
+    assertTrue(getGroup(groupId).isLeader(joinGroupResult.memberId))
+    assertNotEquals(rebalanceResult.followerId, joinGroupResult.memberId)
+    assertEquals(joinGroupResult.protocolName, Some("roundrobin"))
+  }
 
-    EasyMock.reset(replicaManager)
-    val firstSyncResult = syncGroupLeader(groupId, firstGenerationId, firstMemberId, Map(firstMemberId -> Array[Byte]()))
-    assertEquals(Errors.NONE, firstSyncResult._2)
+  @Test
+  def staticMemberRejoinWithUnknownMemberIdAndChangeOfProtocolWhileSelectProtocolUnchangedPersistenceFailure(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
+
+    val selectedProtocol = getGroup(groupId).selectProtocol
+    val newProtocols = List((selectedProtocol, metadata))
+    // Timeout old leader in the meantime.
+    val joinGroupResult = staticJoinGroupWithPersistence(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, followerInstanceId, protocolType, newProtocols, clockAdvance = 1, appendRecordError = Errors.MESSAGE_TOO_LARGE)
+
+    checkJoinGroupResult(joinGroupResult,
+      Errors.UNKNOWN_SERVER_ERROR,
+      rebalanceResult.generation,
+      Set.empty,
+      groupId,
+      Stable,
+      Some(protocolType))
 
     EasyMock.reset(replicaManager)
-    val otherJoinFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
+    // Join with old member id will not fail because the member id is not updated because of persistence failure
+    assertNotEquals(rebalanceResult.followerId, joinGroupResult.memberId)
+    val oldFollowerJoinGroupResult = staticJoinGroup(groupId, rebalanceResult.followerId, followerInstanceId, protocolType, newProtocols, clockAdvance = 1)
+    assertEquals(Errors.NONE, oldFollowerJoinGroupResult.error)
 
     EasyMock.reset(replicaManager)
-    val joinFuture = sendJoinGroup(groupId, firstMemberId, protocolType, protocols)
-
-    val joinResult = await(joinFuture, DefaultSessionTimeout+100)
-    val otherJoinResult = await(otherJoinFuture, DefaultSessionTimeout+100)
-    assertEquals(Errors.NONE, joinResult.error)
-    assertEquals(Errors.NONE, otherJoinResult.error)
-    assertTrue(joinResult.generationId == otherJoinResult.generationId)
+    // Sync with old member id will also not fail because the member id is not updated because of persistence failure
+    val syncGroupWithOldMemberIdResult = syncGroupFollower(groupId, rebalanceResult.generation, rebalanceResult.followerId, None, None, followerInstanceId)
+    assertEquals(Errors.NONE, syncGroupWithOldMemberIdResult.error)
+  }
 
-    assertEquals(firstMemberId, joinResult.leaderId)
-    assertEquals(firstMemberId, otherJoinResult.leaderId)
+  @Test
+  def staticMemberRejoinWithUnknownMemberIdAndChangeOfProtocolWhileSelectProtocolUnchanged(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
+
+    // A static follower rejoin with protocol changing to leader protocol subset won't trigger rebalance if updated
+    // group's selectProtocol remain unchanged.
+    val selectedProtocol = getGroup(groupId).selectProtocol
+    val newProtocols = List((selectedProtocol, metadata))
+    // Timeout old leader in the meantime.
+    val joinGroupResult = staticJoinGroupWithPersistence(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, followerInstanceId, protocolType, newProtocols, clockAdvance = 1)
+
+    checkJoinGroupResult(joinGroupResult,
+      Errors.NONE,
+      rebalanceResult.generation,
+      Set.empty,
+      groupId,
+      Stable,
+      Some(protocolType))
 
-    val nextGenerationId = joinResult.generationId
-    val leaderId = firstMemberId
-    val leaderAssignment = Array[Byte](0)
-    val followerId = otherJoinResult.memberId
-    val followerAssignment = Array[Byte](1)
+    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)
-    val leaderSyncResult = syncGroupLeader(groupId, nextGenerationId, leaderId,
-      Map(leaderId -> leaderAssignment, followerId -> followerAssignment))
-    assertEquals(Errors.NONE, leaderSyncResult._2)
-    assertEquals(leaderAssignment, leaderSyncResult._1)
+    // Sync with old member id will fail because the member id is updated
+    val syncGroupWithOldMemberIdResult = syncGroupFollower(groupId, rebalanceResult.generation, rebalanceResult.followerId, None, None, followerInstanceId)
+    assertEquals(Errors.FENCED_INSTANCE_ID, syncGroupWithOldMemberIdResult.error)
 
     EasyMock.reset(replicaManager)
-    val followerSyncResult = syncGroupFollower(groupId, nextGenerationId, otherJoinResult.memberId)
-    assertEquals(Errors.NONE, followerSyncResult._2)
-    assertEquals(followerAssignment, followerSyncResult._1)
+    val syncGroupWithNewMemberIdResult = syncGroupFollower(groupId, rebalanceResult.generation, joinGroupResult.memberId, None, None, followerInstanceId)
+    assertEquals(Errors.NONE, syncGroupWithNewMemberIdResult.error)
+    assertEquals(rebalanceResult.followerAssignment, syncGroupWithNewMemberIdResult.memberAssignment)
   }
 
   @Test
-  def testSyncGroupLeaderAfterFollower() {
-    // 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 firstMemberId = joinGroupResult.memberId
-    val firstGenerationId = joinGroupResult.generationId
-    assertEquals(firstMemberId, joinGroupResult.leaderId)
-    assertEquals(Errors.NONE, joinGroupResult.error)
+  def staticMemberRejoinWithKnownLeaderIdToTriggerRebalanceAndFollowerWithChangeofProtocol(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
 
+    // A static leader rejoin with known member id will trigger rebalance.
+    val leaderRejoinGroupFuture = sendJoinGroup(groupId, rebalanceResult.leaderId, protocolType, protocolSuperset, leaderInstanceId)
+    // Rebalance complete immediately after follower rejoin.
     EasyMock.reset(replicaManager)
-    val syncGroupResult = syncGroupLeader(groupId, firstGenerationId, firstMemberId, Map(firstMemberId -> Array[Byte]()))
-    val syncGroupError = syncGroupResult._2
-    assertEquals(Errors.NONE, syncGroupError)
+    val followerRejoinWithFuture = sendJoinGroup(groupId, rebalanceResult.followerId, protocolType, protocolSuperset, followerInstanceId)
 
-    EasyMock.reset(replicaManager)
-    val otherJoinFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
+    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,
+      Some(protocolType),
+      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,
+      Some(protocolType),
+      rebalanceResult.leaderId,
+      rebalanceResult.followerId)
 
     EasyMock.reset(replicaManager)
-    val joinFuture = sendJoinGroup(groupId, firstMemberId, protocolType, protocols)
+    // 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,
+      Some(protocolType),
+      rebalanceResult.followerId,
+      rebalanceResult.followerId)
+  }
 
-    val joinResult = await(joinFuture, DefaultSessionTimeout+100)
-    val otherJoinResult = await(otherJoinFuture, DefaultSessionTimeout+100)
-    assertEquals(Errors.NONE, joinResult.error)
-    assertEquals(Errors.NONE, otherJoinResult.error)
-    assertTrue(joinResult.generationId == otherJoinResult.generationId)
+  @Test
+  def staticMemberRejoinAsFollowerWithUnknownMemberId(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
 
-    val nextGenerationId = joinResult.generationId
-    val leaderId = joinResult.leaderId
-    val leaderAssignment = Array[Byte](0)
-    val followerId = otherJoinResult.memberId
-    val followerAssignment = Array[Byte](1)
+    // A static follower rejoin with no protocol change will not trigger rebalance.
+    val joinGroupResult = staticJoinGroupWithPersistence(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, followerInstanceId, protocolType, protocolSuperset, clockAdvance = 1)
 
-    assertEquals(firstMemberId, joinResult.leaderId)
-    assertEquals(firstMemberId, otherJoinResult.leaderId)
+    // 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,
+      Some(protocolType))
 
-    EasyMock.reset(replicaManager)
-    val followerSyncFuture = sendSyncGroupFollower(groupId, nextGenerationId, followerId)
+    assertNotEquals(rebalanceResult.followerId, joinGroupResult.memberId)
 
     EasyMock.reset(replicaManager)
-    val leaderSyncResult = syncGroupLeader(groupId, nextGenerationId, leaderId,
-      Map(leaderId -> leaderAssignment, followerId -> followerAssignment))
-    assertEquals(Errors.NONE, leaderSyncResult._2)
-    assertEquals(leaderAssignment, leaderSyncResult._1)
-
-    val followerSyncResult = await(followerSyncFuture, DefaultSessionTimeout+100)
-    assertEquals(Errors.NONE, followerSyncResult._2)
-    assertEquals(followerAssignment, followerSyncResult._1)
+    val syncGroupResult = syncGroupFollower(groupId, rebalanceResult.generation, joinGroupResult.memberId)
+    assertEquals(Errors.NONE, syncGroupResult.error)
+    assertEquals(rebalanceResult.followerAssignment, syncGroupResult.memberAssignment)
   }
 
   @Test
-  def testCommitOffsetFromUnknownGroup() {
-    val generationId = 1
-    val tp = new TopicPartition("topic", 0)
-    val offset = OffsetAndMetadata(0)
-
-    val commitOffsetResult = commitOffsets(groupId, memberId, generationId, immutable.Map(tp -> offset))
-    assertEquals(Errors.ILLEGAL_GENERATION, commitOffsetResult(tp))
+  def staticMemberRejoinAsFollowerWithKnownMemberIdAndNoProtocolChange(): Unit = {
+    val rebalanceResult  = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
+
+    // A static follower rejoin with no protocol change will not trigger rebalance.
+    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,
+      Some(protocolType),
+      rebalanceResult.leaderId,
+      rebalanceResult.followerId)
   }
 
   @Test
-  def testCommitOffsetWithDefaultGeneration() {
-    val tp = new TopicPartition("topic", 0)
-    val offset = OffsetAndMetadata(0)
+  def staticMemberRejoinAsFollowerWithMismatchedMemberId(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
 
-    val commitOffsetResult = commitOffsets(groupId, OffsetCommitRequest.DEFAULT_MEMBER_ID,
-      OffsetCommitRequest.DEFAULT_GENERATION_ID, immutable.Map(tp -> offset))
-    assertEquals(Errors.NONE, commitOffsetResult(tp))
+    val joinGroupResult = staticJoinGroup(groupId, rebalanceResult.followerId, leaderInstanceId, protocolType, protocolSuperset, clockAdvance = 1)
+    assertEquals(Errors.FENCED_INSTANCE_ID, joinGroupResult.error)
   }
 
   @Test
-  def testFetchOffsets() {
-    val tp = new TopicPartition("topic", 0)
-    val offset = OffsetAndMetadata(0)
-
-    val commitOffsetResult = commitOffsets(groupId, OffsetCommitRequest.DEFAULT_MEMBER_ID,
-      OffsetCommitRequest.DEFAULT_GENERATION_ID, immutable.Map(tp -> offset))
-    assertEquals(Errors.NONE, commitOffsetResult(tp))
+  def staticMemberRejoinAsLeaderWithMismatchedMemberId(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
 
-    val (error, partitionData) = groupCoordinator.handleFetchOffsets(groupId, Some(Seq(tp)))
-    assertEquals(Errors.NONE, error)
-    assertEquals(Some(0), partitionData.get(tp).map(_.offset))
+    val joinGroupResult = staticJoinGroup(groupId, rebalanceResult.leaderId, followerInstanceId, protocolType, protocolSuperset, clockAdvance = 1)
+    assertEquals(Errors.FENCED_INSTANCE_ID, joinGroupResult.error)
   }
 
   @Test
-  def testBasicFetchTxnOffsets() {
-    val tp = new TopicPartition("topic", 0)
-    val offset = OffsetAndMetadata(0)
-    val producerId = 1000L
-    val producerEpoch : Short = 2
+  def staticMemberSyncAsLeaderWithInvalidMemberId(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
 
-    val commitOffsetResult = commitTransactionalOffsets(groupId, producerId, producerEpoch, immutable.Map(tp -> offset))
-    assertEquals(Errors.NONE, commitOffsetResult(tp))
+    val syncGroupResult = syncGroupLeader(groupId, rebalanceResult.generation, "invalid", Map.empty, None, None, leaderInstanceId)
+    assertEquals(Errors.FENCED_INSTANCE_ID, syncGroupResult.error)
+  }
 
-    val (error, partitionData) = groupCoordinator.handleFetchOffsets(groupId, Some(Seq(tp)))
+  @Test
+  def staticMemberHeartbeatLeaderWithInvalidMemberId(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
 
-    // Validate that the offset isn't materialjzed yet.
-    assertEquals(Errors.NONE, error)
-    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), partitionData.get(tp).map(_.offset))
+    val syncGroupResult = syncGroupLeader(groupId, rebalanceResult.generation, rebalanceResult.leaderId, Map.empty)
+    assertEquals(Errors.NONE, syncGroupResult.error)
 
-    val offsetsTopic = new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, groupPartitionId)
+    EasyMock.reset(replicaManager)
+    val validHeartbeatResult = heartbeat(groupId, rebalanceResult.leaderId, rebalanceResult.generation)
+    assertEquals(Errors.NONE, validHeartbeatResult)
 
-    // Send commit marker.
-    groupCoordinator.handleTxnCompletion(producerId, List(offsetsTopic), TransactionResult.COMMIT)
+    EasyMock.reset(replicaManager)
+    val invalidHeartbeatResult = heartbeat(groupId, invalidMemberId, rebalanceResult.generation, leaderInstanceId)
+    assertEquals(Errors.FENCED_INSTANCE_ID, invalidHeartbeatResult)
+  }
 
-    // Validate that committed offset is materialized.
-    val (secondReqError, secondReqPartitionData) = groupCoordinator.handleFetchOffsets(groupId, Some(Seq(tp)))
-    assertEquals(Errors.NONE, secondReqError)
-    assertEquals(Some(0), secondReqPartitionData.get(tp).map(_.offset))
+  @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 = staticJoinGroupWithPersistence(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 testFetchTxnOffsetsWithAbort() {
+  def testOffsetCommitDeadGroup(): Unit = {
+    val memberId = "memberId"
+
+    val deadGroupId = "deadGroupId"
     val tp = new TopicPartition("topic", 0)
-    val offset = OffsetAndMetadata(0)
-    val producerId = 1000L
-    val producerEpoch : Short = 2
+    val offset = offsetAndMetadata(0)
 
-    val commitOffsetResult = commitTransactionalOffsets(groupId, producerId, producerEpoch, immutable.Map(tp -> offset))
-    assertEquals(Errors.NONE, commitOffsetResult(tp))
+    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))
+  }
 
-    val (error, partitionData) = groupCoordinator.handleFetchOffsets(groupId, Some(Seq(tp)))
-    assertEquals(Errors.NONE, error)
-    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), partitionData.get(tp).map(_.offset))
+  @Test
+  def staticMemberCommitOffsetWithInvalidMemberId(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
 
-    val offsetsTopic = new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, groupPartitionId)
+    val syncGroupResult = syncGroupLeader(groupId, rebalanceResult.generation, rebalanceResult.leaderId, Map.empty)
+    assertEquals(Errors.NONE, syncGroupResult.error)
 
-    // Validate that the pending commit is discarded.
-    groupCoordinator.handleTxnCompletion(producerId, List(offsetsTopic), TransactionResult.ABORT)
+    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))
 
-    val (secondReqError, secondReqPartitionData) = groupCoordinator.handleFetchOffsets(groupId, Some(Seq(tp)))
-    assertEquals(Errors.NONE, secondReqError)
+    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)
+
+    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,
+      Some(protocolType)
+    )
+    assertEquals(Set(leaderRejoinGroupResult.memberId), getGroup(groupId).allMembers)
+    assertNotEquals(null, getGroup(groupId).leaderOrNull)
+    assertEquals(3, getGroup(groupId).generationId)
+  }
+
+  def groupStuckInRebalanceTimeoutDueToNonjoinedStaticMember(): Unit = {
+    val longSessionTimeout = DefaultSessionTimeout * 2
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId, sessionTimeout = longSessionTimeout)
+
+    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)
+
+    assertEquals(Set(rebalanceResult.leaderId, rebalanceResult.followerId,
+      dynamicJoinResult.memberId), getGroup(groupId).allMembers)
+    assertEquals(Set(leaderInstanceId.get, followerInstanceId.get),
+      getGroup(groupId).allStaticMembers)
+    assertEquals(Set(dynamicJoinResult.memberId), getGroup(groupId).allDynamicMembers)
+
+    // 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
+    assertTrue(getGroup(groupId).allDynamicMembers.isEmpty)
+    assertEquals(Set(rebalanceResult.leaderId), getGroup(groupId).allMembers)
+    assertTrue(getGroup(groupId).allDynamicMembers.isEmpty)
+    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)
+
+    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,
+      Some(protocolType),
+      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,
+      Some(protocolType),
+      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)
+
+    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,
+      Some(protocolType))
+
+    checkJoinGroupResult(newFollowerResult,
+      Errors.NONE,
+      initialRebalanceResult.generation + 1,
+      Set.empty,
+      groupId,
+      CompletingRebalance,
+      Some(protocolType),
+      expectedLeaderId = newLeaderResult.memberId)
+  }
+
+  @Test
+  def testJoinGroupProtocolTypeIsNotProvidedWhenAnErrorOccurs(): Unit = {
+    // JoinGroup(leader)
+    EasyMock.reset(replicaManager)
+    val leaderResponseFuture = sendJoinGroup(groupId, "fake-id", protocolType,
+      protocolSuperset, leaderInstanceId, DefaultSessionTimeout)
+
+    // The Protocol Type is None when there is an error
+    val leaderJoinGroupResult = await(leaderResponseFuture, 1)
+    assertEquals(Errors.UNKNOWN_MEMBER_ID, leaderJoinGroupResult.error)
+    assertEquals(None, leaderJoinGroupResult.protocolType)
+  }
+
+  @Test
+  def testJoinGroupReturnsTheProtocolType(): Unit = {
+    // JoinGroup(leader)
+    EasyMock.reset(replicaManager)
+    val leaderResponseFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType,
+      protocolSuperset, leaderInstanceId, DefaultSessionTimeout)
+
+    // JoinGroup(follower)
+    EasyMock.reset(replicaManager)
+    val followerResponseFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType,
+      protocolSuperset, followerInstanceId, DefaultSessionTimeout)
+
+    timer.advanceClock(GroupInitialRebalanceDelay + 1)
+    timer.advanceClock(DefaultRebalanceTimeout + 1)
+
+    // The Protocol Type is Defined when there is not error
+    val leaderJoinGroupResult = await(leaderResponseFuture, 1)
+    assertEquals(Errors.NONE, leaderJoinGroupResult.error)
+    assertEquals(protocolType, leaderJoinGroupResult.protocolType.orNull)
+
+    // The Protocol Type is Defined when there is not error
+    val followerJoinGroupResult = await(followerResponseFuture, 1)
+    assertEquals(Errors.NONE, followerJoinGroupResult.error)
+    assertEquals(protocolType, followerJoinGroupResult.protocolType.orNull)
+  }
+
+  @Test
+  def testSyncGroupReturnsAnErrorWhenProtocolTypeIsInconsistent(): Unit = {
+    testSyncGroupProtocolTypeAndNameWith(Some("whatever"), None, Errors.INCONSISTENT_GROUP_PROTOCOL,
+      None, None)
+  }
+
+  @Test
+  def testSyncGroupReturnsAnErrorWhenProtocolNameIsInconsistent(): Unit = {
+    testSyncGroupProtocolTypeAndNameWith(None, Some("whatever"), Errors.INCONSISTENT_GROUP_PROTOCOL,
+      None, None)
+  }
+
+  @Test
+  def testSyncGroupSucceedWhenProtocolTypeAndNameAreNotProvided(): Unit = {
+    testSyncGroupProtocolTypeAndNameWith(None, None, Errors.NONE,
+      Some(protocolType), Some(protocolName))
+  }
+
+  @Test
+  def testSyncGroupSucceedWhenProtocolTypeAndNameAreConsistent(): Unit = {
+    testSyncGroupProtocolTypeAndNameWith(Some(protocolType), Some(protocolName),
+      Errors.NONE, Some(protocolType), Some(protocolName))
+  }
+
+  private def testSyncGroupProtocolTypeAndNameWith(protocolType: Option[String],
+                                                   protocolName: Option[String],
+                                                   expectedError: Errors,
+                                                   expectedProtocolType: Option[String],
+                                                   expectedProtocolName: Option[String]): Unit = {
+    // JoinGroup(leader) with the Protocol Type of the group
+    EasyMock.reset(replicaManager)
+    val leaderResponseFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, this.protocolType,
+      protocolSuperset, leaderInstanceId, DefaultSessionTimeout)
+
+    // JoinGroup(follower) with the Protocol Type of the group
+    EasyMock.reset(replicaManager)
+    val followerResponseFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, this.protocolType,
+      protocolSuperset, followerInstanceId, DefaultSessionTimeout)
+
+    timer.advanceClock(GroupInitialRebalanceDelay + 1)
+    timer.advanceClock(DefaultRebalanceTimeout + 1)
+
+    val leaderJoinGroupResult = await(leaderResponseFuture, 1)
+    val leaderId = leaderJoinGroupResult.memberId
+    val generationId = leaderJoinGroupResult.generationId
+    val followerJoinGroupResult = await(followerResponseFuture, 1)
+    val followerId = followerJoinGroupResult.memberId
+
+    // SyncGroup with the provided Protocol Type and Name
+    EasyMock.reset(replicaManager)
+    val leaderSyncGroupResult = syncGroupLeader(groupId, generationId, leaderId,
+      Map(leaderId -> Array.empty), protocolType, protocolName)
+    assertEquals(expectedError, leaderSyncGroupResult.error)
+    assertEquals(expectedProtocolType, leaderSyncGroupResult.protocolType)
+    assertEquals(expectedProtocolName, leaderSyncGroupResult.protocolName)
+
+    // SyncGroup with the provided Protocol Type and Name
+    EasyMock.reset(replicaManager)
+    val followerSyncGroupResult = syncGroupFollower(groupId, generationId, followerId,
+      protocolType, protocolName)
+    assertEquals(expectedError, followerSyncGroupResult.error)
+    assertEquals(expectedProtocolType, followerSyncGroupResult.protocolType)
+    assertEquals(expectedProtocolName, followerSyncGroupResult.protocolName)
+  }
+
+  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.error)
+    assertTrue(getGroup(groupId).is(Stable))
+
+    EasyMock.reset(replicaManager)
+    val followerId = followerJoinGroupResult.memberId
+    val followerSyncGroupResult = syncGroupFollower(groupId, leaderJoinGroupResult.generationId, followerId)
+    assertEquals(Errors.NONE, followerSyncGroupResult.error)
+    assertTrue(getGroup(groupId).is(Stable))
+
+    EasyMock.reset(replicaManager)
+    new RebalanceResult(newGeneration,
+      leaderId,
+      leaderSyncGroupResult.memberAssignment,
+      followerId,
+      followerSyncGroupResult.memberAssignment)
+  }
+
+  private def checkJoinGroupResult(joinGroupResult: JoinGroupResult,
+                                   expectedError: Errors,
+                                   expectedGeneration: Int,
+                                   expectedGroupInstanceIds: Set[Option[String]],
+                                   groupId: String,
+                                   expectedGroupState: GroupState,
+                                   expectedProtocolType: Option[String],
+                                   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)
+    assertEquals(expectedProtocolType, joinGroupResult.protocolType)
+
+    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(): Unit = {
+    val heartbeatResult = heartbeat(groupId, memberId, -1)
+    assertEquals(Errors.UNKNOWN_MEMBER_ID, heartbeatResult)
+  }
+
+  @Test
+  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 = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
+    val assignedMemberId = joinGroupResult.memberId
+    val joinGroupError = joinGroupResult.error
+    assertEquals(Errors.NONE, joinGroupError)
+
+    EasyMock.reset(replicaManager)
+    val syncGroupResult = syncGroupLeader(groupId, joinGroupResult.generationId, assignedMemberId, Map(assignedMemberId -> Array[Byte]()))
+    assertEquals(Errors.NONE, syncGroupResult.error)
+
+    EasyMock.reset(replicaManager)
+    val heartbeatResult = heartbeat(groupId, otherMemberId, 1)
+    assertEquals(Errors.UNKNOWN_MEMBER_ID, heartbeatResult)
+  }
+
+  @Test
+  def testHeartbeatRebalanceInProgress(): Unit = {
+    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
+
+    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, 1)
+    assertEquals(Errors.NONE, heartbeatResult)
+  }
+
+  @Test
+  def testHeartbeatIllegalGeneration(): Unit = {
+    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
+
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
+    val assignedMemberId = joinGroupResult.memberId
+    val joinGroupError = joinGroupResult.error
+    assertEquals(Errors.NONE, joinGroupError)
+
+    EasyMock.reset(replicaManager)
+    val syncGroupResult = syncGroupLeader(groupId, joinGroupResult.generationId, assignedMemberId, Map(assignedMemberId -> Array[Byte]()))
+    assertEquals(Errors.NONE, syncGroupResult.error)
+
+    EasyMock.reset(replicaManager)
+    val heartbeatResult = heartbeat(groupId, assignedMemberId, 2)
+    assertEquals(Errors.ILLEGAL_GENERATION, heartbeatResult)
+  }
+
+  @Test
+  def testValidHeartbeat(): Unit = {
+    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
+
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
+    val assignedConsumerId = joinGroupResult.memberId
+    val generationId = joinGroupResult.generationId
+    val joinGroupError = joinGroupResult.error
+    assertEquals(Errors.NONE, joinGroupError)
+
+    EasyMock.reset(replicaManager)
+    val syncGroupResult = syncGroupLeader(groupId, generationId, assignedConsumerId, Map(assignedConsumerId -> Array[Byte]()))
+    assertEquals(Errors.NONE, syncGroupResult.error)
+
+    EasyMock.reset(replicaManager)
+    val heartbeatResult = heartbeat(groupId, assignedConsumerId, 1)
+    assertEquals(Errors.NONE, heartbeatResult)
+  }
+
+  @Test
+  def testSessionTimeout(): Unit = {
+    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
+
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
+    val assignedConsumerId = joinGroupResult.memberId
+    val generationId = joinGroupResult.generationId
+    val joinGroupError = joinGroupResult.error
+    assertEquals(Errors.NONE, joinGroupError)
+
+    EasyMock.reset(replicaManager)
+    val syncGroupResult = syncGroupLeader(groupId, generationId, assignedConsumerId, Map(assignedConsumerId -> Array[Byte]()))
+    assertEquals(Errors.NONE, syncGroupResult.error)
+
+    EasyMock.reset(replicaManager)
+    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)
+
+    timer.advanceClock(DefaultSessionTimeout + 100)
+
+    EasyMock.reset(replicaManager)
+    val heartbeatResult = heartbeat(groupId, assignedConsumerId, 1)
+    assertEquals(Errors.UNKNOWN_MEMBER_ID, heartbeatResult)
+  }
+
+  @Test
+  def testHeartbeatMaintainsSession(): Unit = {
+    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
+    val sessionTimeout = 1000
+
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols,
+      rebalanceTimeout = sessionTimeout, sessionTimeout = sessionTimeout)
+    val assignedConsumerId = joinGroupResult.memberId
+    val generationId = joinGroupResult.generationId
+    val joinGroupError = joinGroupResult.error
+    assertEquals(Errors.NONE, joinGroupError)
+
+    EasyMock.reset(replicaManager)
+    val syncGroupResult = syncGroupLeader(groupId, generationId, assignedConsumerId, Map(assignedConsumerId -> Array[Byte]()))
+    assertEquals(Errors.NONE, syncGroupResult.error)
+
+    timer.advanceClock(sessionTimeout / 2)
+
+    EasyMock.reset(replicaManager)
+    var heartbeatResult = heartbeat(groupId, assignedConsumerId, 1)
+    assertEquals(Errors.NONE, heartbeatResult)
+
+    timer.advanceClock(sessionTimeout / 2 + 100)
+
+    EasyMock.reset(replicaManager)
+    heartbeatResult = heartbeat(groupId, assignedConsumerId, 1)
+    assertEquals(Errors.NONE, heartbeatResult)
+  }
+
+  @Test
+  def testCommitMaintainsSession(): Unit = {
+    val sessionTimeout = 1000
+    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
+    val tp = new TopicPartition("topic", 0)
+    val offset = offsetAndMetadata(0)
+
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols,
+      rebalanceTimeout = sessionTimeout, sessionTimeout = sessionTimeout)
+    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]()))
+    assertEquals(Errors.NONE, syncGroupResult.error)
+
+    timer.advanceClock(sessionTimeout / 2)
+
+    EasyMock.reset(replicaManager)
+    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, assignedMemberId, 1)
+    assertEquals(Errors.NONE, heartbeatResult)
+  }
+
+  @Test
+  def testSessionTimeoutDuringRebalance(): Unit = {
+    // create a group with a single member
+    val firstJoinResult = dynamicJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols,
+      rebalanceTimeout = 2000, sessionTimeout = 1000)
+    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.error)
+
+    // now have a new member join to trigger a rebalance
+    EasyMock.reset(replicaManager)
+    val otherJoinFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
+
+    timer.advanceClock(500)
+
+    EasyMock.reset(replicaManager)
+    var heartbeatResult = heartbeat(groupId, firstMemberId, firstGenerationId)
+    assertEquals(Errors.REBALANCE_IN_PROGRESS, heartbeatResult)
+
+    // letting the session expire should make the member fall out of the group
+    timer.advanceClock(1100)
+
+    EasyMock.reset(replicaManager)
+    heartbeatResult = heartbeat(groupId, firstMemberId, firstGenerationId)
+    assertEquals(Errors.UNKNOWN_MEMBER_ID, heartbeatResult)
+
+    // and the rebalance should complete with only the new member
+    val otherJoinResult = await(otherJoinFuture, DefaultSessionTimeout+100)
+    assertEquals(Errors.NONE, otherJoinResult.error)
+  }
+
+  @Test
+  def testRebalanceCompletesBeforeMemberJoins(): Unit = {
+    // create a group with a single member
+    val firstJoinResult = staticJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, leaderInstanceId, protocolType, protocols,
+      rebalanceTimeout = 1200, sessionTimeout = 1000)
+    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.error)
+
+    // 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
+    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(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.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)
+    val otherRejoinGenerationId = otherReJoinResult.generationId
+    val reSyncResult = syncGroupLeader(groupId, otherRejoinGenerationId, otherMemberId, Map(otherMemberId -> Array[Byte]()))
+    assertEquals(Errors.NONE, reSyncResult.error)
+
+    // 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, otherRejoinGenerationId)
+      assertEquals(Errors.NONE, heartbeatResult)
+    }
+  }
+
+  @Test
+  def testSyncGroupEmptyAssignment(): Unit = {
+    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
+
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
+    val assignedConsumerId = joinGroupResult.memberId
+    val generationId = joinGroupResult.generationId
+    val joinGroupError = joinGroupResult.error
+    assertEquals(Errors.NONE, joinGroupError)
+
+    EasyMock.reset(replicaManager)
+    val syncGroupResult = syncGroupLeader(groupId, generationId, assignedConsumerId, Map())
+    assertEquals(Errors.NONE, syncGroupResult.error)
+    assertTrue(syncGroupResult.memberAssignment.isEmpty)
+
+    EasyMock.reset(replicaManager)
+    val heartbeatResult = heartbeat(groupId, assignedConsumerId, 1)
+    assertEquals(Errors.NONE, heartbeatResult)
+  }
+
+  @Test
+  def testSyncGroupNotCoordinator(): Unit = {
+    val generation = 1
+
+    val syncGroupResult = syncGroupFollower(otherGroupId, generation, memberId)
+    assertEquals(Errors.NOT_COORDINATOR, syncGroupResult.error)
+  }
+
+  @Test
+  def testSyncGroupFromUnknownGroup(): Unit = {
+    val syncGroupResult = syncGroupFollower(groupId, 1, memberId)
+    assertEquals(Errors.UNKNOWN_MEMBER_ID, syncGroupResult.error)
+  }
+
+  @Test
+  def testSyncGroupFromUnknownMember(): Unit = {
+    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
+
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
+    val assignedConsumerId = joinGroupResult.memberId
+    val generationId = joinGroupResult.generationId
+    assertEquals(Errors.NONE, joinGroupResult.error)
+
+    EasyMock.reset(replicaManager)
+    val syncGroupResult = syncGroupLeader(groupId, generationId, assignedConsumerId, Map(assignedConsumerId -> Array[Byte]()))
+    val syncGroupError = syncGroupResult.error
+    assertEquals(Errors.NONE, syncGroupError)
+
+    EasyMock.reset(replicaManager)
+    val unknownMemberId = "blah"
+    val unknownMemberSyncResult = syncGroupFollower(groupId, generationId, unknownMemberId)
+    assertEquals(Errors.UNKNOWN_MEMBER_ID, unknownMemberSyncResult.error)
+  }
+
+  @Test
+  def testSyncGroupFromIllegalGeneration(): Unit = {
+    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
+
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
+    val assignedConsumerId = joinGroupResult.memberId
+    val generationId = joinGroupResult.generationId
+    assertEquals(Errors.NONE, joinGroupResult.error)
+
+    EasyMock.reset(replicaManager)
+    // send the sync group with an invalid generation
+    val syncGroupResult = syncGroupLeader(groupId, generationId+1, assignedConsumerId, Map(assignedConsumerId -> Array[Byte]()))
+    assertEquals(Errors.ILLEGAL_GENERATION, syncGroupResult.error)
+  }
+
+  @Test
+  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 = 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.error)
+
+    EasyMock.reset(replicaManager)
+    val otherJoinFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
+
+    EasyMock.reset(replicaManager)
+    val joinFuture = sendJoinGroup(groupId, firstMemberId, protocolType, protocols)
+
+    val joinResult = await(joinFuture, DefaultSessionTimeout+100)
+    val otherJoinResult = await(otherJoinFuture, DefaultSessionTimeout+100)
+    assertEquals(Errors.NONE, joinResult.error)
+    assertEquals(Errors.NONE, otherJoinResult.error)
+    assertTrue(joinResult.generationId == otherJoinResult.generationId)
+
+    assertEquals(firstMemberId, joinResult.leaderId)
+    assertEquals(firstMemberId, otherJoinResult.leaderId)
+
+    val nextGenerationId = joinResult.generationId
+
+    // this shouldn't cause a rebalance since protocol information hasn't changed
+    EasyMock.reset(replicaManager)
+    val followerJoinResult = await(sendJoinGroup(groupId, otherJoinResult.memberId, protocolType, protocols), 1)
+
+    assertEquals(Errors.NONE, followerJoinResult.error)
+    assertEquals(nextGenerationId, followerJoinResult.generationId)
+  }
+
+  @Test
+  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)
+    assertEquals(Errors.NONE, firstJoinResult.error)
+
+    EasyMock.reset(replicaManager)
+    val firstSyncResult = syncGroupLeader(groupId, firstGenerationId, firstMemberId, Map(firstMemberId -> Array[Byte]()))
+    assertEquals(Errors.NONE, firstSyncResult.error)
+
+    // join groups from the leader should force the group to rebalance, which allows the
+    // leader to push new assignments when local metadata changes
+
+    EasyMock.reset(replicaManager)
+    val secondJoinResult = await(sendJoinGroup(groupId, firstMemberId, protocolType, protocols), 1)
+
+    assertEquals(Errors.NONE, secondJoinResult.error)
+    assertNotEquals(firstGenerationId, secondJoinResult.generationId)
+  }
+
+  /**
+    * Test if the following scenario completes a rebalance correctly: A new member starts a JoinGroup request with
+    * an UNKNOWN_MEMBER_ID, attempting to join a stable group. But never initiates the second JoinGroup request with
+    * the provided member ID and times out. The test checks if original member remains the sole member in this group,
+    * which should remain stable throughout this test.
+    */
+  @Test
+  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)
+    assertEquals(Errors.NONE, firstJoinResult.error)
+
+    // Starting sync group leader
+    EasyMock.reset(replicaManager)
+    val firstSyncResult = syncGroupLeader(groupId, firstGenerationId, firstMemberId, Map(firstMemberId -> Array[Byte]()))
+    assertEquals(Errors.NONE, firstSyncResult.error)
+    timer.advanceClock(100)
+    assertEquals(Set(firstMemberId), groupCoordinator.groupManager.getGroup(groupId).get.allMembers)
+    assertEquals(groupCoordinator.groupManager.getGroup(groupId).get.allMembers,
+      groupCoordinator.groupManager.getGroup(groupId).get.allDynamicMembers)
+    assertEquals(0, groupCoordinator.groupManager.getGroup(groupId).get.numPending)
+    val group = groupCoordinator.groupManager.getGroup(groupId).get
+
+    // ensure the group is stable before a new member initiates join request
+    assertEquals(Stable, group.currentState)
+
+    // new member initiates join group
+    EasyMock.reset(replicaManager)
+    val secondJoinResult = joinGroupPartial(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
+    assertEquals(Errors.MEMBER_ID_REQUIRED, secondJoinResult.error)
+    assertEquals(1, group.numPending)
+    assertEquals(Stable, group.currentState)
+
+    EasyMock.reset(replicaManager)
+    EasyMock.expect(replicaManager.getMagic(EasyMock.anyObject())).andReturn(Some(RecordBatch.MAGIC_VALUE_V1)).anyTimes()
+    EasyMock.replay(replicaManager)
+
+    // advance clock to timeout the pending member
+    assertEquals(Set(firstMemberId), group.allMembers)
+    assertEquals(1, group.numPending)
+    timer.advanceClock(300)
+
+    // original (firstMember) member sends heartbeats to prevent session timeouts.
+    EasyMock.reset(replicaManager)
+    val heartbeatResult = heartbeat(groupId, firstMemberId, 1)
+    assertEquals(Errors.NONE, heartbeatResult)
+
+    // timeout the pending member
+    timer.advanceClock(300)
+
+    // at this point the second member should have been removed from pending list (session timeout),
+    // and the group should be in Stable state with only the first member in it.
+    assertEquals(Set(firstMemberId), group.allMembers)
+    assertEquals(0, group.numPending)
+    assertEquals(Stable, group.currentState)
+    assertTrue(group.has(firstMemberId))
+  }
+
+  /**
+    * Create a group with two members in Stable state. Create a third pending member by completing it's first JoinGroup
+    * request without a member id.
+    */
+  private def setupGroupWithPendingMember(): JoinGroupResult = {
+    // add the first member
+    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
+    EasyMock.reset(replicaManager)
+    val firstSyncResult = syncGroupLeader(groupId, joinResult1.generationId, joinResult1.memberId, Map(joinResult1.memberId -> Array[Byte]()))
+    assertEquals(Errors.NONE, firstSyncResult.error)
+    assertGroupState(groupState = Stable)
+
+    // start the join for the second member
+    EasyMock.reset(replicaManager)
+    val secondJoinFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
+
+    // rejoin the first member back into the group
+    EasyMock.reset(replicaManager)
+    val firstJoinFuture = sendJoinGroup(groupId, joinResult1.memberId, protocolType, protocols)
+    val firstMemberJoinResult = await(firstJoinFuture, DefaultSessionTimeout+100)
+    val secondMemberJoinResult = await(secondJoinFuture, DefaultSessionTimeout+100)
+    assertGroupState(groupState = CompletingRebalance)
+
+    // stabilize the group
+    EasyMock.reset(replicaManager)
+    val secondSyncResult = syncGroupLeader(groupId, firstMemberJoinResult.generationId, joinResult1.memberId, Map(joinResult1.memberId -> Array[Byte]()))
+    assertEquals(Errors.NONE, secondSyncResult.error)
+    assertGroupState(groupState = Stable)
+
+    // re-join an existing member, to transition the group to PreparingRebalance state.
+    EasyMock.reset(replicaManager)
+    sendJoinGroup(groupId, firstMemberJoinResult.memberId, protocolType, protocols)
+    assertGroupState(groupState = PreparingRebalance)
+
+    // create a pending member in the group
+    EasyMock.reset(replicaManager)
+    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
+    EasyMock.reset(replicaManager)
+    sendJoinGroup(groupId, secondMemberJoinResult.memberId, protocolType, protocols)
+    assertGroupState(groupState = PreparingRebalance)
+    assertEquals(1, groupCoordinator.groupManager.getGroup(groupId).get.numPending)
+
+    pendingMember
+  }
+
+  /**
+    * Setup a group in with a pending member. The test checks if the a pending member joining completes the rebalancing
+    * operation
+    */
+  @Test
+  def testJoinGroupCompletionWhenPendingMemberJoins(): Unit = {
+    val pendingMember = setupGroupWithPendingMember()
+
+    // compete join group for the pending member
+    EasyMock.reset(replicaManager)
+    val pendingMemberJoinFuture = sendJoinGroup(groupId, pendingMember.memberId, protocolType, protocols)
+    await(pendingMemberJoinFuture, DefaultSessionTimeout+100)
+
+    assertGroupState(groupState = CompletingRebalance)
+    assertEquals(3, group().allMembers.size)
+    assertEquals(0, group().numPending)
+  }
+
+  /**
+    * Setup a group in with a pending member. The test checks if the timeout of the pending member will
+    * cause the group to return to a CompletingRebalance state.
+    */
+  @Test
+  def testJoinGroupCompletionWhenPendingMemberTimesOut(): Unit = {
+    setupGroupWithPendingMember()
+
+    // Advancing Clock by > 100 (session timeout for third and fourth member)
+    // and < 500 (for first and second members). This will force the coordinator to attempt join
+    // completion on heartbeat expiration (since we are in PendingRebalance stage).
+    EasyMock.reset(replicaManager)
+    EasyMock.expect(replicaManager.getMagic(EasyMock.anyObject())).andReturn(Some(RecordBatch.MAGIC_VALUE_V1)).anyTimes()
+    EasyMock.replay(replicaManager)
+    timer.advanceClock(120)
+
+    assertGroupState(groupState = CompletingRebalance)
+    assertEquals(2, group().allMembers.size)
+    assertEquals(0, group().numPending)
+  }
+
+  @Test
+  def testPendingMembersLeavesGroup(): Unit = {
+    val pending = setupGroupWithPendingMember()
+
+    EasyMock.reset(replicaManager)
+    val leaveGroupResults = singleLeaveGroup(groupId, pending.memberId)
+    verifyLeaveGroupResult(leaveGroupResults)
+
+    assertGroupState(groupState = CompletingRebalance)
+    assertEquals(2, group().allMembers.size)
+    assertEquals(2, group().allDynamicMembers.size)
+    assertEquals(0, group().numPending)
+  }
+
+  private def group(groupId: String = groupId) = {
+    groupCoordinator.groupManager.getGroup(groupId) match {
+      case Some(g) => g
+      case None => null
+    }
+  }
+
+  private def assertGroupState(groupId: String = groupId,
+                               groupState: GroupState): Unit = {
+    groupCoordinator.groupManager.getGroup(groupId) match {
+      case Some(group) => assertEquals(groupState, group.currentState)
+      case None => fail(s"Group $groupId not found in coordinator")
+    }
+  }
+
+  private def joinGroupPartial(groupId: String,
+                               memberId: String,
+                               protocolType: String,
+                               protocols: List[(String, Array[Byte])],
+                               sessionTimeout: Int = DefaultSessionTimeout,
+                               rebalanceTimeout: Int = DefaultRebalanceTimeout): JoinGroupResult = {
+    val requireKnownMemberId = true
+    val responseFuture = sendJoinGroup(groupId, memberId, protocolType, protocols, None, sessionTimeout, rebalanceTimeout, requireKnownMemberId)
+    Await.result(responseFuture, Duration(rebalanceTimeout + 100, TimeUnit.MILLISECONDS))
+  }
+
+  @Test
+  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 = 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.error)
+
+    EasyMock.reset(replicaManager)
+    val otherJoinFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
+
+    EasyMock.reset(replicaManager)
+    val joinFuture = sendJoinGroup(groupId, firstMemberId, protocolType, protocols)
+
+    val joinResult = await(joinFuture, DefaultSessionTimeout+100)
+    val otherJoinResult = await(otherJoinFuture, DefaultSessionTimeout+100)
+    assertEquals(Errors.NONE, joinResult.error)
+    assertEquals(Errors.NONE, otherJoinResult.error)
+    assertTrue(joinResult.generationId == otherJoinResult.generationId)
+
+    assertEquals(firstMemberId, joinResult.leaderId)
+    assertEquals(firstMemberId, otherJoinResult.leaderId)
+
+    val nextGenerationId = joinResult.generationId
+
+    // 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, None, None, None)
+
+    timer.advanceClock(DefaultSessionTimeout + 100)
+
+    val followerSyncResult = await(followerSyncFuture, DefaultSessionTimeout+100)
+    assertEquals(Errors.REBALANCE_IN_PROGRESS, followerSyncResult.error)
+  }
+
+  @Test
+  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 = 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.error)
+
+    EasyMock.reset(replicaManager)
+    val otherJoinFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
+
+    EasyMock.reset(replicaManager)
+    val joinFuture = sendJoinGroup(groupId, firstMemberId, protocolType, protocols)
+
+    val joinResult = await(joinFuture, DefaultSessionTimeout+100)
+    val otherJoinResult = await(otherJoinFuture, DefaultSessionTimeout+100)
+    assertEquals(Errors.NONE, joinResult.error)
+    assertEquals(Errors.NONE, otherJoinResult.error)
+    assertTrue(joinResult.generationId == otherJoinResult.generationId)
+
+    assertEquals(firstMemberId, joinResult.leaderId)
+    assertEquals(firstMemberId, otherJoinResult.leaderId)
+
+    val nextGenerationId = joinResult.generationId
+    val leaderId = firstMemberId
+    val leaderAssignment = Array[Byte](0)
+    val followerId = otherJoinResult.memberId
+    val followerAssignment = Array[Byte](1)
+
+    EasyMock.reset(replicaManager)
+    val leaderSyncResult = syncGroupLeader(groupId, nextGenerationId, leaderId,
+      Map(leaderId -> leaderAssignment, followerId -> followerAssignment))
+    assertEquals(Errors.NONE, leaderSyncResult.error)
+    assertEquals(leaderAssignment, leaderSyncResult.memberAssignment)
+
+    EasyMock.reset(replicaManager)
+    val followerSyncResult = syncGroupFollower(groupId, nextGenerationId, otherJoinResult.memberId)
+    assertEquals(Errors.NONE, followerSyncResult.error)
+    assertEquals(followerAssignment, followerSyncResult.memberAssignment)
+  }
+
+  @Test
+  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 = dynamicJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
+    val firstMemberId = joinGroupResult.memberId
+    val firstGenerationId = joinGroupResult.generationId
+    assertEquals(firstMemberId, joinGroupResult.leaderId)
+    assertEquals(Errors.NONE, joinGroupResult.error)
+
+    EasyMock.reset(replicaManager)
+    val syncGroupResult = syncGroupLeader(groupId, firstGenerationId, firstMemberId, Map(firstMemberId -> Array[Byte]()))
+    assertEquals(Errors.NONE, syncGroupResult.error)
+
+    EasyMock.reset(replicaManager)
+    val otherJoinFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
+
+    EasyMock.reset(replicaManager)
+    val joinFuture = sendJoinGroup(groupId, firstMemberId, protocolType, protocols)
+
+    val joinResult = await(joinFuture, DefaultSessionTimeout+100)
+    val otherJoinResult = await(otherJoinFuture, DefaultSessionTimeout+100)
+    assertEquals(Errors.NONE, joinResult.error)
+    assertEquals(Errors.NONE, otherJoinResult.error)
+    assertTrue(joinResult.generationId == otherJoinResult.generationId)
+
+    val nextGenerationId = joinResult.generationId
+    val leaderId = joinResult.leaderId
+    val leaderAssignment = Array[Byte](0)
+    val followerId = otherJoinResult.memberId
+    val followerAssignment = Array[Byte](1)
+
+    assertEquals(firstMemberId, joinResult.leaderId)
+    assertEquals(firstMemberId, otherJoinResult.leaderId)
+
+    EasyMock.reset(replicaManager)
+    val followerSyncFuture = sendSyncGroupFollower(groupId, nextGenerationId, followerId, None, None, None)
+
+    EasyMock.reset(replicaManager)
+    val leaderSyncResult = syncGroupLeader(groupId, nextGenerationId, leaderId,
+      Map(leaderId -> leaderAssignment, followerId -> followerAssignment))
+    assertEquals(Errors.NONE, leaderSyncResult.error)
+    assertEquals(leaderAssignment, leaderSyncResult.memberAssignment)
+
+    val followerSyncResult = await(followerSyncFuture, DefaultSessionTimeout+100)
+    assertEquals(Errors.NONE, followerSyncResult.error)
+    assertEquals(followerAssignment, followerSyncResult.memberAssignment)
+  }
+
+  @Test
+  def testCommitOffsetFromUnknownGroup(): Unit = {
+    val generationId = 1
+    val tp = new TopicPartition("topic", 0)
+    val offset = offsetAndMetadata(0)
+
+    val commitOffsetResult = commitOffsets(groupId, memberId, generationId, Map(tp -> offset))
+    assertEquals(Errors.ILLEGAL_GENERATION, commitOffsetResult(tp))
+  }
+
+  @Test
+  def testCommitOffsetWithDefaultGeneration(): Unit = {
+    val tp = new TopicPartition("topic", 0)
+    val offset = offsetAndMetadata(0)
+
+    val commitOffsetResult = commitOffsets(groupId, OffsetCommitRequest.DEFAULT_MEMBER_ID,
+      OffsetCommitRequest.DEFAULT_GENERATION_ID, Map(tp -> offset))
+    assertEquals(Errors.NONE, commitOffsetResult(tp))
+  }
+
+  @Test
+  def testCommitOffsetsAfterGroupIsEmpty(): Unit = {
+    // Tests the scenario where the reset offset tool modifies the offsets
+    // of a group after it becomes empty
+
+    // A group member joins
+    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
+    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 leaveGroupResults = singleLeaveGroup(groupId, assignedMemberId)
+    verifyLeaveGroupResult(leaveGroupResults)
+
+    // The simple offset commit should now fail
+    EasyMock.reset(replicaManager)
+    val tp = new TopicPartition("topic", 0)
+    val offset = offsetAndMetadata(0)
+    val commitOffsetResult = commitOffsets(groupId, OffsetCommitRequest.DEFAULT_MEMBER_ID,
+      OffsetCommitRequest.DEFAULT_GENERATION_ID, Map(tp -> offset))
+    assertEquals(Errors.NONE, commitOffsetResult(tp))
+
+    val (error, partitionData) = groupCoordinator.handleFetchOffsets(groupId, requireStable, Some(Seq(tp)))
+    assertEquals(Errors.NONE, error)
+    assertEquals(Some(0), partitionData.get(tp).map(_.offset))
+  }
+
+  @Test
+  def testFetchOffsets(): Unit = {
+    val tp = new TopicPartition("topic", 0)
+    val offset = 97L
+    val metadata = "some metadata"
+    val leaderEpoch = Optional.of[Integer](15)
+    val offsetAndMetadata = OffsetAndMetadata(offset, leaderEpoch, metadata, timer.time.milliseconds(), None)
+
+    val commitOffsetResult = commitOffsets(groupId, OffsetCommitRequest.DEFAULT_MEMBER_ID,
+      OffsetCommitRequest.DEFAULT_GENERATION_ID, Map(tp -> offsetAndMetadata))
+    assertEquals(Errors.NONE, commitOffsetResult(tp))
+
+    val (error, partitionData) = groupCoordinator.handleFetchOffsets(groupId, requireStable, Some(Seq(tp)))
+    assertEquals(Errors.NONE, error)
+
+    val maybePartitionData = partitionData.get(tp)
+    assertTrue(maybePartitionData.isDefined)
+    assertEquals(offset, maybePartitionData.get.offset)
+    assertEquals(metadata, maybePartitionData.get.metadata)
+    assertEquals(leaderEpoch, maybePartitionData.get.leaderEpoch)
+  }
+
+  @Test
+  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
+
+    val tp = new TopicPartition("topic", 0)
+    val offset = offsetAndMetadata(0)
+    val groupId = ""
+
+    val commitOffsetResult = commitOffsets(groupId, OffsetCommitRequest.DEFAULT_MEMBER_ID,
+      OffsetCommitRequest.DEFAULT_GENERATION_ID, Map(tp -> offset))
+    assertEquals(Errors.NONE, commitOffsetResult(tp))
+
+    val (fetchError, partitionData) = groupCoordinator.handleFetchOffsets(groupId, requireStable, Some(Seq(tp)))
+    assertEquals(Errors.NONE, fetchError)
+    assertEquals(Some(0), partitionData.get(tp).map(_.offset))
+
+    val (describeError, summary) = groupCoordinator.handleDescribeGroup(groupId)
+    assertEquals(Errors.NONE, describeError)
+    assertEquals(Empty.toString, summary.state)
+
+    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 deleteErrors = groupCoordinator.handleDeleteGroups(Set(groupId))
+    assertEquals(Errors.NONE, deleteErrors(groupId))
+
+    val (err, data) = groupCoordinator.handleFetchOffsets(groupId, requireStable, Some(Seq(tp)))
+    assertEquals(Errors.NONE, err)
+    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), data.get(tp).map(_.offset))
+  }
+
+  @Test
+  def testBasicFetchTxnOffsets(): Unit = {
+    val tp = new TopicPartition("topic", 0)
+    val offset = offsetAndMetadata(0)
+    val producerId = 1000L
+    val producerEpoch : Short = 2
+
+    val commitOffsetResult = commitTransactionalOffsets(groupId, producerId, producerEpoch, Map(tp -> offset))
+    assertEquals(Errors.NONE, commitOffsetResult(tp))
+
+    val (error, partitionData) = groupCoordinator.handleFetchOffsets(groupId, requireStable, Some(Seq(tp)))
+
+    // Validate that the offset isn't materialjzed yet.
+    assertEquals(Errors.NONE, error)
+    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), partitionData.get(tp).map(_.offset))
+
+    val offsetsTopic = new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, groupPartitionId)
+
+    // Send commit marker.
+    handleTxnCompletion(producerId, List(offsetsTopic), TransactionResult.COMMIT)
+
+    // Validate that committed offset is materialized.
+    val (secondReqError, secondReqPartitionData) = groupCoordinator.handleFetchOffsets(groupId, requireStable, Some(Seq(tp)))
+    assertEquals(Errors.NONE, secondReqError)
+    assertEquals(Some(0), secondReqPartitionData.get(tp).map(_.offset))
+  }
+
+  @Test
+  def testFetchTxnOffsetsWithAbort(): Unit = {
+    val tp = new TopicPartition("topic", 0)
+    val offset = offsetAndMetadata(0)
+    val producerId = 1000L
+    val producerEpoch : Short = 2
+
+    val commitOffsetResult = commitTransactionalOffsets(groupId, producerId, producerEpoch, Map(tp -> offset))
+    assertEquals(Errors.NONE, commitOffsetResult(tp))
+
+    val (error, partitionData) = groupCoordinator.handleFetchOffsets(groupId, requireStable, Some(Seq(tp)))
+    assertEquals(Errors.NONE, error)
+    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), partitionData.get(tp).map(_.offset))
+
+    val offsetsTopic = new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, groupPartitionId)
+
+    // Validate that the pending commit is discarded.
+    handleTxnCompletion(producerId, List(offsetsTopic), TransactionResult.ABORT)
+
+    val (secondReqError, secondReqPartitionData) = groupCoordinator.handleFetchOffsets(groupId, requireStable, Some(Seq(tp)))
+    assertEquals(Errors.NONE, secondReqError)
+    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), secondReqPartitionData.get(tp).map(_.offset))
+  }
+
+  @Test
+  def testFetchPendingTxnOffsetsWithAbort(): Unit = {
+    val tp = new TopicPartition("topic", 0)
+    val offset = offsetAndMetadata(0)
+    val producerId = 1000L
+    val producerEpoch : Short = 2
+
+    val commitOffsetResult = commitTransactionalOffsets(groupId, producerId, producerEpoch, Map(tp -> offset))
+    assertEquals(Errors.NONE, commitOffsetResult(tp))
+
+    val nonExistTp = new TopicPartition("non-exist-topic", 0)
+    val (error, partitionData) = groupCoordinator.handleFetchOffsets(groupId, requireStable, Some(Seq(tp, nonExistTp)))
+    assertEquals(Errors.NONE, error)
+    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), partitionData.get(tp).map(_.offset))
+    assertEquals(Some(Errors.UNSTABLE_OFFSET_COMMIT), partitionData.get(tp).map(_.error))
+    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), partitionData.get(nonExistTp).map(_.offset))
+    assertEquals(Some(Errors.NONE), partitionData.get(nonExistTp).map(_.error))
+
+    val offsetsTopic = new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, groupPartitionId)
+
+    // Validate that the pending commit is discarded.
+    handleTxnCompletion(producerId, List(offsetsTopic), TransactionResult.ABORT)
+
+    val (secondReqError, secondReqPartitionData) = groupCoordinator.handleFetchOffsets(groupId, requireStable, Some(Seq(tp)))
+    assertEquals(Errors.NONE, secondReqError)
+    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), secondReqPartitionData.get(tp).map(_.offset))
+    assertEquals(Some(Errors.NONE), secondReqPartitionData.get(tp).map(_.error))
+  }
+
+  @Test
+  def testFetchPendingTxnOffsetsWithCommit(): Unit = {
+    val tp = new TopicPartition("topic", 0)
+    val offset = offsetAndMetadata(25)
+    val producerId = 1000L
+    val producerEpoch : Short = 2
+
+    val commitOffsetResult = commitTransactionalOffsets(groupId, producerId, producerEpoch, Map(tp -> offset))
+    assertEquals(Errors.NONE, commitOffsetResult(tp))
+
+    val (error, partitionData) = groupCoordinator.handleFetchOffsets(groupId, requireStable, Some(Seq(tp)))
+    assertEquals(Errors.NONE, error)
+    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), partitionData.get(tp).map(_.offset))
+    assertEquals(Some(Errors.UNSTABLE_OFFSET_COMMIT), partitionData.get(tp).map(_.error))
+
+    val offsetsTopic = new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, groupPartitionId)
+
+    // Validate that the pending commit is committed
+    handleTxnCompletion(producerId, List(offsetsTopic), TransactionResult.COMMIT)
+
+    val (secondReqError, secondReqPartitionData) = groupCoordinator.handleFetchOffsets(groupId, requireStable, Some(Seq(tp)))
+    assertEquals(Errors.NONE, secondReqError)
+    assertEquals(Some(25), secondReqPartitionData.get(tp).map(_.offset))
+    assertEquals(Some(Errors.NONE), secondReqPartitionData.get(tp).map(_.error))
+  }
+
+  @Test
+  def testFetchTxnOffsetsIgnoreSpuriousCommit(): Unit = {
+    val tp = new TopicPartition("topic", 0)
+    val offset = offsetAndMetadata(0)
+    val producerId = 1000L
+    val producerEpoch : Short = 2
+
+    val commitOffsetResult = commitTransactionalOffsets(groupId, producerId, producerEpoch, Map(tp -> offset))
+    assertEquals(Errors.NONE, commitOffsetResult(tp))
+
+    val (error, partitionData) = groupCoordinator.handleFetchOffsets(groupId, requireStable, Some(Seq(tp)))
+    assertEquals(Errors.NONE, error)
+    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), partitionData.get(tp).map(_.offset))
+
+    val offsetsTopic = new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, groupPartitionId)
+    handleTxnCompletion(producerId, List(offsetsTopic), TransactionResult.ABORT)
+
+    val (secondReqError, secondReqPartitionData) = groupCoordinator.handleFetchOffsets(groupId, requireStable, Some(Seq(tp)))
+    assertEquals(Errors.NONE, secondReqError)
     assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), secondReqPartitionData.get(tp).map(_.offset))
+
+    // Ignore spurious commit.
+    handleTxnCompletion(producerId, List(offsetsTopic), TransactionResult.COMMIT)
+
+    val (thirdReqError, thirdReqPartitionData) = groupCoordinator.handleFetchOffsets(groupId, requireStable, Some(Seq(tp)))
+    assertEquals(Errors.NONE, thirdReqError)
+    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), thirdReqPartitionData.get(tp).map(_.offset))
+  }
+
+  @Test
+  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.
+
+    val partitions = List(new TopicPartition("topic1", 0), new TopicPartition("topic2", 0))
+    val offsets = List(offsetAndMetadata(10), offsetAndMetadata(15))
+    val producerId = 1000L
+    val producerEpoch: Short = 3
+
+    val groupIds = List(groupId, otherGroupId)
+    val offsetTopicPartitions = List(new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, groupCoordinator.partitionFor(groupId)),
+      new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, groupCoordinator.partitionFor(otherGroupId)))
+
+    groupCoordinator.groupManager.addPartitionOwnership(offsetTopicPartitions(1).partition)
+    val errors = mutable.ArrayBuffer[Errors]()
+    val partitionData = mutable.ArrayBuffer[scala.collection.Map[TopicPartition, OffsetFetchResponse.PartitionData]]()
+
+    val commitOffsetResults = mutable.ArrayBuffer[CommitOffsetCallbackParams]()
+
+    // Ensure that the two groups map to different partitions.
+    assertNotEquals(offsetTopicPartitions(0), offsetTopicPartitions(1))
+
+    commitOffsetResults.append(commitTransactionalOffsets(groupId, producerId, producerEpoch, Map(partitions(0) -> offsets(0))))
+    assertEquals(Errors.NONE, commitOffsetResults(0)(partitions(0)))
+    commitOffsetResults.append(commitTransactionalOffsets(otherGroupId, producerId, producerEpoch, Map(partitions(1) -> offsets(1))))
+    assertEquals(Errors.NONE, commitOffsetResults(1)(partitions(1)))
+
+    // We got a commit for only one __consumer_offsets partition. We should only materialize it's group offsets.
+    handleTxnCompletion(producerId, List(offsetTopicPartitions(0)), TransactionResult.COMMIT)
+    groupCoordinator.handleFetchOffsets(groupIds(0), requireStable, Some(partitions)) match {
+      case (error, partData) =>
+        errors.append(error)
+        partitionData.append(partData)
+      case _ =>
+    }
+
+    groupCoordinator.handleFetchOffsets(groupIds(1), requireStable, Some(partitions)) match {
+      case (error, partData) =>
+        errors.append(error)
+        partitionData.append(partData)
+      case _ =>
+    }
+
+    assertEquals(2, errors.size)
+    assertEquals(Errors.NONE, errors(0))
+    assertEquals(Errors.NONE, errors(1))
+
+    // Exactly one offset commit should have been materialized.
+    assertEquals(Some(offsets(0).offset), partitionData(0).get(partitions(0)).map(_.offset))
+    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), partitionData(0).get(partitions(1)).map(_.offset))
+    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), partitionData(1).get(partitions(0)).map(_.offset))
+    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), partitionData(1).get(partitions(1)).map(_.offset))
+
+    // Now we receive the other marker.
+    handleTxnCompletion(producerId, List(offsetTopicPartitions(1)), TransactionResult.COMMIT)
+    errors.clear()
+    partitionData.clear()
+    groupCoordinator.handleFetchOffsets(groupIds(0), requireStable, Some(partitions)) match {
+      case (error, partData) =>
+        errors.append(error)
+        partitionData.append(partData)
+      case _ =>
+    }
+
+     groupCoordinator.handleFetchOffsets(groupIds(1), requireStable, Some(partitions)) match {
+      case (error, partData) =>
+        errors.append(error)
+        partitionData.append(partData)
+      case _ =>
+    }
+    // Two offsets should have been materialized
+    assertEquals(Some(offsets(0).offset), partitionData(0).get(partitions(0)).map(_.offset))
+    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), partitionData(0).get(partitions(1)).map(_.offset))
+    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), partitionData(1).get(partitions(0)).map(_.offset))
+    assertEquals(Some(offsets(1).offset), partitionData(1).get(partitions(1)).map(_.offset))
+  }
+
+  @Test
+  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.
+
+    val partitions = List(new TopicPartition("topic1", 0), new TopicPartition("topic2", 0))
+    val offsets = List(offsetAndMetadata(10), offsetAndMetadata(15))
+    val producerIds = List(1000L, 1005L)
+    val producerEpochs: Seq[Short] = List(3, 4)
+
+    val offsetTopicPartition = new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, groupCoordinator.partitionFor(groupId))
+
+    val errors = mutable.ArrayBuffer[Errors]()
+    val partitionData = mutable.ArrayBuffer[scala.collection.Map[TopicPartition, OffsetFetchResponse.PartitionData]]()
+
+    val commitOffsetResults = mutable.ArrayBuffer[CommitOffsetCallbackParams]()
+
+    // producer0 commits the offsets for partition0
+    commitOffsetResults.append(commitTransactionalOffsets(groupId, producerIds(0), producerEpochs(0), Map(partitions(0) -> offsets(0))))
+    assertEquals(Errors.NONE, commitOffsetResults(0)(partitions(0)))
+
+    // producer1 commits the offsets for partition1
+    commitOffsetResults.append(commitTransactionalOffsets(groupId, producerIds(1), producerEpochs(1), Map(partitions(1) -> offsets(1))))
+    assertEquals(Errors.NONE, commitOffsetResults(1)(partitions(1)))
+
+    // producer0 commits its transaction.
+    handleTxnCompletion(producerIds(0), List(offsetTopicPartition), TransactionResult.COMMIT)
+    groupCoordinator.handleFetchOffsets(groupId, requireStable, Some(partitions)) match {
+      case (error, partData) =>
+        errors.append(error)
+        partitionData.append(partData)
+      case _ =>
+    }
+
+    assertEquals(Errors.NONE, errors(0))
+
+    // We should only see the offset commit for producer0
+    assertEquals(Some(offsets(0).offset), partitionData(0).get(partitions(0)).map(_.offset))
+    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), partitionData(0).get(partitions(1)).map(_.offset))
+
+    // producer1 now commits its transaction.
+    handleTxnCompletion(producerIds(1), List(offsetTopicPartition), TransactionResult.COMMIT)
+
+    groupCoordinator.handleFetchOffsets(groupId, requireStable, Some(partitions)) match {
+      case (error, partData) =>
+        errors.append(error)
+        partitionData.append(partData)
+      case _ =>
+    }
+
+    assertEquals(Errors.NONE, errors(1))
+
+    // We should now see the offset commits for both producers.
+    assertEquals(Some(offsets(0).offset), partitionData(1).get(partitions(0)).map(_.offset))
+    assertEquals(Some(offsets(1).offset), partitionData(1).get(partitions(1)).map(_.offset))
+  }
+
+  @Test
+  def testFetchOffsetForUnknownPartition(): Unit = {
+    val tp = new TopicPartition("topic", 0)
+    val (error, partitionData) = groupCoordinator.handleFetchOffsets(groupId, requireStable, Some(Seq(tp)))
+    assertEquals(Errors.NONE, error)
+    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), partitionData.get(tp).map(_.offset))
+  }
+
+  @Test
+  def testFetchOffsetNotCoordinatorForGroup(): Unit = {
+    val tp = new TopicPartition("topic", 0)
+    val (error, partitionData) = groupCoordinator.handleFetchOffsets(otherGroupId, requireStable, Some(Seq(tp)))
+    assertEquals(Errors.NOT_COORDINATOR, error)
+    assertTrue(partitionData.isEmpty)
+  }
+
+  @Test
+  def testFetchAllOffsets(): Unit = {
+    val tp1 = new TopicPartition("topic", 0)
+    val tp2 = new TopicPartition("topic", 1)
+    val tp3 = new TopicPartition("other-topic", 0)
+    val offset1 = offsetAndMetadata(15)
+    val offset2 = offsetAndMetadata(16)
+    val offset3 = offsetAndMetadata(17)
+
+    assertEquals((Errors.NONE, Map.empty), groupCoordinator.handleFetchOffsets(groupId, requireStable))
+
+    val commitOffsetResult = commitOffsets(groupId, OffsetCommitRequest.DEFAULT_MEMBER_ID,
+      OffsetCommitRequest.DEFAULT_GENERATION_ID, Map(tp1 -> offset1, tp2 -> offset2, tp3 -> offset3))
+    assertEquals(Errors.NONE, commitOffsetResult(tp1))
+    assertEquals(Errors.NONE, commitOffsetResult(tp2))
+    assertEquals(Errors.NONE, commitOffsetResult(tp3))
+
+    val (error, partitionData) = groupCoordinator.handleFetchOffsets(groupId, requireStable)
+    assertEquals(Errors.NONE, error)
+    assertEquals(3, partitionData.size)
+    assertTrue(partitionData.forall(_._2.error == Errors.NONE))
+    assertEquals(Some(offset1.offset), partitionData.get(tp1).map(_.offset))
+    assertEquals(Some(offset2.offset), partitionData.get(tp2).map(_.offset))
+    assertEquals(Some(offset3.offset), partitionData.get(tp3).map(_.offset))
+  }
+
+  @Test
+  def testCommitOffsetInCompletingRebalance(): 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, Map(tp -> offset))
+    assertEquals(Errors.REBALANCE_IN_PROGRESS, commitOffsetResult(tp))
+  }
+
+  @Test
+  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 testFetchTxnOffsetsIgnoreSpuriousCommit() {
+  def testTxnCommitOffsetWithFencedInstanceId(): Unit = {
     val tp = new TopicPartition("topic", 0)
-    val offset = OffsetAndMetadata(0)
+    val offset = offsetAndMetadata(0)
     val producerId = 1000L
     val producerEpoch : Short = 2
 
-    val commitOffsetResult = commitTransactionalOffsets(groupId, producerId, producerEpoch, immutable.Map(tp -> offset))
-    assertEquals(Errors.NONE, commitOffsetResult(tp))
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
 
-    val (error, partitionData) = groupCoordinator.handleFetchOffsets(groupId, Some(Seq(tp)))
-    assertEquals(Errors.NONE, error)
-    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), partitionData.get(tp).map(_.offset))
+    val leaderNoMemberIdCommitOffsetResult = commitTransactionalOffsets(groupId, producerId, producerEpoch,
+      Map(tp -> offset), memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID, groupInstanceId = leaderInstanceId)
+    assertEquals(Errors.FENCED_INSTANCE_ID, leaderNoMemberIdCommitOffsetResult (tp))
 
-    val offsetsTopic = new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, groupPartitionId)
-    groupCoordinator.handleTxnCompletion(producerId, List(offsetsTopic), TransactionResult.ABORT)
+    val leaderInvalidMemberIdCommitOffsetResult = commitTransactionalOffsets(groupId, producerId, producerEpoch,
+      Map(tp -> offset), memberId = "invalid-member", groupInstanceId = leaderInstanceId)
+    assertEquals(Errors.FENCED_INSTANCE_ID, leaderInvalidMemberIdCommitOffsetResult (tp))
 
-    val (secondReqError, secondReqPartitionData) = groupCoordinator.handleFetchOffsets(groupId, Some(Seq(tp)))
-    assertEquals(Errors.NONE, secondReqError)
-    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), secondReqPartitionData.get(tp).map(_.offset))
+    val leaderCommitOffsetResult = commitTransactionalOffsets(groupId, producerId, producerEpoch,
+      Map(tp -> offset), rebalanceResult.leaderId, leaderInstanceId)
+    assertEquals(Errors.NONE, leaderCommitOffsetResult (tp))
+  }
 
-    // Ignore spurious commit.
-    groupCoordinator.handleTxnCompletion(producerId, List(offsetsTopic), TransactionResult.COMMIT)
+  @Test
+  def testTxnCommitOffsetWithInvalidMemberId(): Unit = {
+    val tp = new TopicPartition("topic", 0)
+    val offset = offsetAndMetadata(0)
+    val producerId = 1000L
+    val producerEpoch : Short = 2
 
-    val (thirdReqError, thirdReqPartitionData) = groupCoordinator.handleFetchOffsets(groupId, Some(Seq(tp)))
-    assertEquals(Errors.NONE, secondReqError)
-    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), thirdReqPartitionData.get(tp).map(_.offset))
+    val joinGroupResult = dynamicJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
+    val joinGroupError = joinGroupResult.error
+    assertEquals(Errors.NONE, joinGroupError)
+
+    EasyMock.reset(replicaManager)
+    val invalidIdCommitOffsetResult = commitTransactionalOffsets(groupId, producerId, producerEpoch,
+      Map(tp -> offset), "invalid-member")
+    assertEquals(Errors.UNKNOWN_MEMBER_ID, invalidIdCommitOffsetResult (tp))
   }
 
   @Test
-  def testFetchTxnOffsetsOneProducerMultipleGroups() {
-    // 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.
+  def testTxnCommitOffsetWithKnownMemberId(): Unit = {
+    val tp = new TopicPartition("topic", 0)
+    val offset = offsetAndMetadata(0)
+    val producerId = 1000L
+    val producerEpoch : Short = 2
 
-    val partitions = List(new TopicPartition("topic1", 0), new TopicPartition("topic2", 0))
-    val offsets = List(OffsetAndMetadata(10), OffsetAndMetadata(15))
+    val joinGroupResult = dynamicJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
+    val joinGroupError = joinGroupResult.error
+    assertEquals(Errors.NONE, joinGroupError)
+
+    EasyMock.reset(replicaManager)
+    val assignedConsumerId = joinGroupResult.memberId
+    val leaderCommitOffsetResult = commitTransactionalOffsets(groupId, producerId, producerEpoch,
+      Map(tp -> offset), assignedConsumerId)
+    assertEquals(Errors.NONE, leaderCommitOffsetResult (tp))
+  }
+
+  @Test
+  def testTxnCommitOffsetWithIllegalGeneration(): Unit = {
+    val tp = new TopicPartition("topic", 0)
+    val offset = offsetAndMetadata(0)
     val producerId = 1000L
-    val producerEpoch: Short = 3
+    val producerEpoch : Short = 2
 
-    val groupIds = List(groupId, otherGroupId)
-    val offsetTopicPartitions = List(new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, groupCoordinator.partitionFor(groupId)),
-      new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, groupCoordinator.partitionFor(otherGroupId)))
+    val joinGroupResult = dynamicJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
+    val joinGroupError = joinGroupResult.error
+    assertEquals(Errors.NONE, joinGroupError)
 
-    groupCoordinator.groupManager.addPartitionOwnership(offsetTopicPartitions(1).partition)
-    val errors = mutable.ArrayBuffer[Errors]()
-    val partitionData = mutable.ArrayBuffer[Map[TopicPartition, OffsetFetchResponse.PartitionData]]()
+    EasyMock.reset(replicaManager)
 
-    val commitOffsetResults = mutable.ArrayBuffer[CommitOffsetCallbackParams]()
+    val assignedConsumerId = joinGroupResult.memberId
+    val initialGenerationId = joinGroupResult.generationId
+    val illegalGenerationCommitOffsetResult = commitTransactionalOffsets(groupId, producerId, producerEpoch,
+      Map(tp -> offset), memberId = assignedConsumerId, generationId = initialGenerationId + 5)
+    assertEquals(Errors.ILLEGAL_GENERATION, illegalGenerationCommitOffsetResult (tp))
+  }
 
-    // Ensure that the two groups map to different partitions.
-    assertNotEquals(offsetTopicPartitions(0), offsetTopicPartitions(1))
+  @Test
+  def testTxnCommitOffsetWithLegalGeneration(): Unit = {
+    val tp = new TopicPartition("topic", 0)
+    val offset = offsetAndMetadata(0)
+    val producerId = 1000L
+    val producerEpoch : Short = 2
 
-    commitOffsetResults.append(commitTransactionalOffsets(groupId, producerId, producerEpoch, immutable.Map(partitions(0) -> offsets(0))))
-    assertEquals(Errors.NONE, commitOffsetResults(0)(partitions(0)))
-    commitOffsetResults.append(commitTransactionalOffsets(otherGroupId, producerId, producerEpoch, immutable.Map(partitions(1) -> offsets(1))))
-    assertEquals(Errors.NONE, commitOffsetResults(1)(partitions(1)))
+    val joinGroupResult = dynamicJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
+    val joinGroupError = joinGroupResult.error
+    assertEquals(Errors.NONE, joinGroupError)
 
-    // We got a commit for only one __consumer_offsets partition. We should only materialize it's group offsets.
-    groupCoordinator.handleTxnCompletion(producerId, List(offsetTopicPartitions(0)), TransactionResult.COMMIT)
-    groupCoordinator.handleFetchOffsets(groupIds(0), Some(partitions)) match {
-      case (error, partData) =>
-        errors.append(error)
-        partitionData.append(partData)
-      case _ =>
-    }
+    EasyMock.reset(replicaManager)
 
-     groupCoordinator.handleFetchOffsets(groupIds(1), Some(partitions)) match {
-      case (error, partData) =>
-        errors.append(error)
-        partitionData.append(partData)
-      case _ =>
-    }
+    val assignedConsumerId = joinGroupResult.memberId
+    val initialGenerationId = joinGroupResult.generationId
+    val leaderCommitOffsetResult = commitTransactionalOffsets(groupId, producerId, producerEpoch,
+      Map(tp -> offset), memberId = assignedConsumerId, generationId = initialGenerationId)
+    assertEquals(Errors.NONE, leaderCommitOffsetResult (tp))
+  }
 
-    assertEquals(2, errors.size)
-    assertEquals(Errors.NONE, errors(0))
-    assertEquals(Errors.NONE, errors(1))
+  @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 = dynamicJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
+    val assignedConsumerId = joinGroupResult.memberId
+    val initialGenerationId = joinGroupResult.generationId
+    val joinGroupError = joinGroupResult.error
+    assertEquals(Errors.NONE, joinGroupError)
 
-    // Exactly one offset commit should have been materialized.
-    assertEquals(Some(offsets(0).offset), partitionData(0).get(partitions(0)).map(_.offset))
-    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), partitionData(0).get(partitions(1)).map(_.offset))
-    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), partitionData(1).get(partitions(0)).map(_.offset))
-    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), partitionData(1).get(partitions(1)).map(_.offset))
+    // Then join with a new consumer to trigger a rebalance
+    EasyMock.reset(replicaManager)
+    sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
 
-    // Now we receive the other marker.
-    groupCoordinator.handleTxnCompletion(producerId, List(offsetTopicPartitions(1)), TransactionResult.COMMIT)
-    errors.clear()
-    partitionData.clear()
-    groupCoordinator.handleFetchOffsets(groupIds(0), Some(partitions)) match {
-      case (error, partData) =>
-        errors.append(error)
-        partitionData.append(partData)
-      case _ =>
-    }
+    // We should be in the middle of a rebalance, so the heartbeat should return rebalance in progress
+    EasyMock.reset(replicaManager)
+    val heartbeatResult = heartbeat(groupId, assignedConsumerId, initialGenerationId)
+    assertEquals(Errors.REBALANCE_IN_PROGRESS, heartbeatResult)
+  }
 
-     groupCoordinator.handleFetchOffsets(groupIds(1), Some(partitions)) match {
-      case (error, partData) =>
-        errors.append(error)
-        partitionData.append(partData)
-      case _ =>
-    }
-    // Two offsets should have been materialized
-    assertEquals(Some(offsets(0).offset), partitionData(0).get(partitions(0)).map(_.offset))
-    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), partitionData(0).get(partitions(1)).map(_.offset))
-    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), partitionData(1).get(partitions(0)).map(_.offset))
-    assertEquals(Some(offsets(1).offset), partitionData(1).get(partitions(1)).map(_.offset))
+  @Test
+  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
+    assertEquals(1, initialGenerationId)
+    assertEquals(Errors.NONE, joinGroupError)
+
+    EasyMock.reset(replicaManager)
+    val syncGroupResult = syncGroupLeader(groupId, initialGenerationId, memberId, Map(memberId -> Array[Byte]()))
+    assertEquals(Errors.NONE, syncGroupResult.error)
+
+    EasyMock.reset(replicaManager)
+    val joinGroupFuture = sendJoinGroup(groupId, memberId, protocolType, protocols)
+    val otherJoinGroupResult = await(joinGroupFuture, 1)
+
+    val nextGenerationId = otherJoinGroupResult.generationId
+    val otherJoinGroupError = otherJoinGroupResult.error
+    assertEquals(2, nextGenerationId)
+    assertEquals(Errors.NONE, otherJoinGroupError)
   }
 
   @Test
-  def testFetchTxnOffsetsMultipleProducersOneGroup() {
-    // 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.
+  def testLeaveGroupWrongCoordinator(): Unit = {
+    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
 
-    val partitions = List(new TopicPartition("topic1", 0), new TopicPartition("topic2", 0))
-    val offsets = List(OffsetAndMetadata(10), OffsetAndMetadata(15))
-    val producerIds = List(1000L, 1005L)
-    val producerEpochs: Seq[Short] = List(3, 4)
+    val leaveGroupResults = singleLeaveGroup(otherGroupId, memberId)
+    verifyLeaveGroupResult(leaveGroupResults, Errors.NOT_COORDINATOR)
+  }
 
-    val offsetTopicPartition = new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, groupCoordinator.partitionFor(groupId))
+  @Test
+  def testLeaveGroupUnknownGroup(): Unit = {
+    val leaveGroupResults = singleLeaveGroup(groupId, memberId)
+    verifyLeaveGroupResult(leaveGroupResults, Errors.NONE, List(Errors.UNKNOWN_MEMBER_ID))
+  }
 
-    val errors = mutable.ArrayBuffer[Errors]()
-    val partitionData = mutable.ArrayBuffer[Map[TopicPartition, OffsetFetchResponse.PartitionData]]()
+  @Test
+  def testLeaveGroupUnknownConsumerExistingGroup(): Unit = {
+    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
+    val otherMemberId = "consumerId"
 
-    val commitOffsetResults = mutable.ArrayBuffer[CommitOffsetCallbackParams]()
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
+    val joinGroupError = joinGroupResult.error
+    assertEquals(Errors.NONE, joinGroupError)
 
-    // producer0 commits the offsets for partition0
-    commitOffsetResults.append(commitTransactionalOffsets(groupId, producerIds(0), producerEpochs(0), immutable.Map(partitions(0) -> offsets(0))))
-    assertEquals(Errors.NONE, commitOffsetResults(0)(partitions(0)))
+    EasyMock.reset(replicaManager)
+    val leaveGroupResults = singleLeaveGroup(groupId, otherMemberId)
+    verifyLeaveGroupResult(leaveGroupResults, Errors.NONE, List(Errors.UNKNOWN_MEMBER_ID))
+  }
 
-    // producer1 commits the offsets for partition1
-    commitOffsetResults.append(commitTransactionalOffsets(groupId, producerIds(1), producerEpochs(1), immutable.Map(partitions(1) -> offsets(1))))
-    assertEquals(Errors.NONE, commitOffsetResults(1)(partitions(1)))
+  @Test
+  def testSingleLeaveDeadGroup(): Unit = {
+    val deadGroupId = "deadGroupId"
 
-    // producer0 commits its transaction.
-    groupCoordinator.handleTxnCompletion(producerIds(0), List(offsetTopicPartition), TransactionResult.COMMIT)
-    groupCoordinator.handleFetchOffsets(groupId, Some(partitions)) match {
-      case (error, partData) =>
-        errors.append(error)
-        partitionData.append(partData)
-      case _ =>
-    }
+    groupCoordinator.groupManager.addGroup(new GroupMetadata(deadGroupId, Dead, new MockTime()))
+    val leaveGroupResults = singleLeaveGroup(deadGroupId, memberId)
+    verifyLeaveGroupResult(leaveGroupResults, Errors.COORDINATOR_NOT_AVAILABLE)
+  }
 
-    assertEquals(Errors.NONE, errors(0))
+  @Test
+  def testBatchLeaveDeadGroup(): Unit = {
+    val deadGroupId = "deadGroupId"
 
-    // We should only see the offset commit for producer0
-    assertEquals(Some(offsets(0).offset), partitionData(0).get(partitions(0)).map(_.offset))
-    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), partitionData(0).get(partitions(1)).map(_.offset))
+    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)
+  }
 
-    // producer1 now commits its transaction.
-    groupCoordinator.handleTxnCompletion(producerIds(1), List(offsetTopicPartition), TransactionResult.COMMIT)
+  @Test
+  def testValidLeaveGroup(): Unit = {
+    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
 
-    groupCoordinator.handleFetchOffsets(groupId, Some(partitions)) match {
-      case (error, partData) =>
-        errors.append(error)
-        partitionData.append(partData)
-      case _ =>
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
+    val assignedMemberId = joinGroupResult.memberId
+    val joinGroupError = joinGroupResult.error
+    assertEquals(Errors.NONE, joinGroupError)
+
+    EasyMock.reset(replicaManager)
+    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)
+
+    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)
+
+    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)
+
+    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)
+
+    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)))
     }
 
-    assertEquals(Errors.NONE, errors(1))
+    val message = expectedException.getMessage
+    assertTrue(message.contains(instanceId))
+    assertTrue(message.contains(pendingMemberId))
+  }
 
-    // We should now see the offset commits for both producers.
-    assertEquals(Some(offsets(0).offset), partitionData(1).get(partitions(0)).map(_.offset))
-    assertEquals(Some(offsets(1).offset), partitionData(1).get(partitions(1)).map(_.offset))
+  @Test
+  def testListGroupsIncludesStableGroups(): Unit = {
+    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
+    val assignedMemberId = joinGroupResult.memberId
+    val generationId = joinGroupResult.generationId
+    assertEquals(Errors.NONE, joinGroupResult.error)
+
+    EasyMock.reset(replicaManager)
+    val syncGroupResult = syncGroupLeader(groupId, generationId, assignedMemberId, Map(assignedMemberId -> Array[Byte]()))
+    assertEquals(Errors.NONE, syncGroupResult.error)
+
+    val (error, groups) = groupCoordinator.handleListGroups(Set())
+    assertEquals(Errors.NONE, error)
+    assertEquals(1, groups.size)
+    assertEquals(GroupOverview("groupId", "consumer", Stable.toString), groups.head)
+  }
+
+  @Test
+  def testListGroupsIncludesRebalancingGroups(): Unit = {
+    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
+    assertEquals(Errors.NONE, joinGroupResult.error)
+
+    val (error, groups) = groupCoordinator.handleListGroups(Set())
+    assertEquals(Errors.NONE, error)
+    assertEquals(1, groups.size)
+    assertEquals(GroupOverview("groupId", "consumer", CompletingRebalance.toString), groups.head)
   }
 
   @Test
-  def testFetchOffsetForUnknownPartition(): Unit = {
-    val tp = new TopicPartition("topic", 0)
-    val (error, partitionData) = groupCoordinator.handleFetchOffsets(groupId, Some(Seq(tp)))
+  def testListGroupsWithStates(): Unit = {
+    val allStates = Set(PreparingRebalance, CompletingRebalance, Stable, Dead, Empty).map(s => s.toString)
+    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
+
+    // Member joins the group
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
+    val assignedMemberId = joinGroupResult.memberId
+    val generationId = joinGroupResult.generationId
+    assertEquals(Errors.NONE, joinGroupResult.error)
+
+    // The group should be in CompletingRebalance
+    val (error, groups) = groupCoordinator.handleListGroups(Set(CompletingRebalance.toString))
     assertEquals(Errors.NONE, error)
-    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), partitionData.get(tp).map(_.offset))
+    assertEquals(1, groups.size)
+    val (error2, groups2) = groupCoordinator.handleListGroups(allStates.filterNot(s => s == CompletingRebalance.toString))
+    assertEquals(Errors.NONE, error2)
+    assertEquals(0, groups2.size)
+
+    // Member syncs
+    EasyMock.reset(replicaManager)
+    val syncGroupResult = syncGroupLeader(groupId, generationId, assignedMemberId, Map(assignedMemberId -> Array[Byte]()))
+    assertEquals(Errors.NONE, syncGroupResult.error)
+
+    // The group is now stable
+    val (error3, groups3) = groupCoordinator.handleListGroups(Set(Stable.toString))
+    assertEquals(Errors.NONE, error3)
+    assertEquals(1, groups3.size)
+    val (error4, groups4) = groupCoordinator.handleListGroups(allStates.filterNot(s => s == Stable.toString))
+    assertEquals(Errors.NONE, error4)
+    assertEquals(0, groups4.size)
+
+    // Member leaves
+    EasyMock.reset(replicaManager)
+    val leaveGroupResults = singleLeaveGroup(groupId, assignedMemberId)
+    verifyLeaveGroupResult(leaveGroupResults)
+
+    // The group is now empty
+    val (error5, groups5) = groupCoordinator.handleListGroups(Set(Empty.toString))
+    assertEquals(Errors.NONE, error5)
+    assertEquals(1, groups5.size)
+    val (error6, groups6) = groupCoordinator.handleListGroups(allStates.filterNot(s => s == Empty.toString))
+    assertEquals(Errors.NONE, error6)
+    assertEquals(0, groups6.size)
   }
 
   @Test
-  def testFetchOffsetNotCoordinatorForGroup(): Unit = {
-    val tp = new TopicPartition("topic", 0)
-    val (error, partitionData) = groupCoordinator.handleFetchOffsets(otherGroupId, Some(Seq(tp)))
+  def testDescribeGroupWrongCoordinator(): Unit = {
+    EasyMock.reset(replicaManager)
+    val (error, _) = groupCoordinator.handleDescribeGroup(otherGroupId)
     assertEquals(Errors.NOT_COORDINATOR, error)
-    assertTrue(partitionData.isEmpty)
   }
 
   @Test
-  def testFetchAllOffsets() {
-    val tp1 = new TopicPartition("topic", 0)
-    val tp2 = new TopicPartition("topic", 1)
-    val tp3 = new TopicPartition("other-topic", 0)
-    val offset1 = OffsetAndMetadata(15)
-    val offset2 = OffsetAndMetadata(16)
-    val offset3 = OffsetAndMetadata(17)
-
-    assertEquals((Errors.NONE, Map.empty), groupCoordinator.handleFetchOffsets(groupId))
-
-    val commitOffsetResult = commitOffsets(groupId, OffsetCommitRequest.DEFAULT_MEMBER_ID,
-      OffsetCommitRequest.DEFAULT_GENERATION_ID, immutable.Map(tp1 -> offset1, tp2 -> offset2, tp3 -> offset3))
-    assertEquals(Errors.NONE, commitOffsetResult(tp1))
-    assertEquals(Errors.NONE, commitOffsetResult(tp2))
-    assertEquals(Errors.NONE, commitOffsetResult(tp3))
-
-    val (error, partitionData) = groupCoordinator.handleFetchOffsets(groupId)
+  def testDescribeGroupInactiveGroup(): Unit = {
+    EasyMock.reset(replicaManager)
+    val (error, summary) = groupCoordinator.handleDescribeGroup(groupId)
     assertEquals(Errors.NONE, error)
-    assertEquals(3, partitionData.size)
-    assertTrue(partitionData.forall(_._2.error == Errors.NONE))
-    assertEquals(Some(offset1.offset), partitionData.get(tp1).map(_.offset))
-    assertEquals(Some(offset2.offset), partitionData.get(tp2).map(_.offset))
-    assertEquals(Some(offset3.offset), partitionData.get(tp3).map(_.offset))
+    assertEquals(GroupCoordinator.DeadGroup, summary)
   }
 
   @Test
-  def testCommitOffsetInCompletingRebalance() {
-    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
-    val tp = new TopicPartition("topic", 0)
-    val offset = OffsetAndMetadata(0)
-
-    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 commitOffsetResult = commitOffsets(groupId, assignedMemberId, generationId, immutable.Map(tp -> offset))
-    assertEquals(Errors.REBALANCE_IN_PROGRESS, commitOffsetResult(tp))
+    val syncGroupResult = syncGroupLeader(groupId, generationId, assignedMemberId, Map(assignedMemberId -> Array[Byte]()))
+    assertEquals(Errors.NONE, syncGroupResult.error)
+
+    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 testHeartbeatDuringRebalanceCausesRebalanceInProgress() {
-    // 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 assignedConsumerId = joinGroupResult.memberId
-    val initialGenerationId = joinGroupResult.generationId
+  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
     assertEquals(Errors.NONE, joinGroupError)
 
-    // Then join with a new consumer to trigger a rebalance
     EasyMock.reset(replicaManager)
-    sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
+    val syncGroupResult = syncGroupLeader(groupId, generationId, assignedMemberId, Map(assignedMemberId -> Array[Byte]()))
+    assertEquals(Errors.NONE, syncGroupResult.error)
 
-    // We should be in the middle of a rebalance, so the heartbeat should return rebalance in progress
     EasyMock.reset(replicaManager)
-    val heartbeatResult = heartbeat(groupId, assignedConsumerId, initialGenerationId)
-    assertEquals(Errors.REBALANCE_IN_PROGRESS, heartbeatResult)
+    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))
+    assertEquals(List(leaderInstanceId), summary.members.map(_.groupInstanceId))
   }
 
   @Test
-  def testGenerationIdIncrementsOnRebalance() {
-    val joinGroupResult = joinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
-    val initialGenerationId = joinGroupResult.generationId
+  def testDescribeGroupRebalancing(): Unit = {
+    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
     val joinGroupError = joinGroupResult.error
-    val memberId = joinGroupResult.memberId
-    assertEquals(1, initialGenerationId)
     assertEquals(Errors.NONE, joinGroupError)
 
     EasyMock.reset(replicaManager)
-    val syncGroupResult = syncGroupLeader(groupId, initialGenerationId, memberId, Map(memberId -> Array[Byte]()))
-    val syncGroupError = syncGroupResult._2
-    assertEquals(Errors.NONE, syncGroupError)
-
-    EasyMock.reset(replicaManager)
-    val joinGroupFuture = sendJoinGroup(groupId, memberId, protocolType, protocols)
-    val otherJoinGroupResult = await(joinGroupFuture, 1)
-
-    val nextGenerationId = otherJoinGroupResult.generationId
-    val otherJoinGroupError = otherJoinGroupResult.error
-    assertEquals(2, nextGenerationId)
-    assertEquals(Errors.NONE, otherJoinGroupError)
+    val (error, summary) = groupCoordinator.handleDescribeGroup(groupId)
+    assertEquals(Errors.NONE, error)
+    assertEquals(protocolType, summary.protocolType)
+    assertEquals(GroupCoordinator.NoProtocol, summary.protocol)
+    assertEquals(CompletingRebalance.toString, summary.state)
+    assertTrue(summary.members.map(_.memberId).contains(joinGroupResult.memberId))
+    assertTrue(summary.members.forall(_.metadata.isEmpty))
+    assertTrue(summary.members.forall(_.assignment.isEmpty))
   }
 
   @Test
-  def testLeaveGroupWrongCoordinator() {
+  def testDeleteNonEmptyGroup(): Unit = {
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
+    dynamicJoinGroup(groupId, memberId, protocolType, protocols)
 
-    val leaveGroupResult = leaveGroup(otherGroupId, memberId)
-    assertEquals(Errors.NOT_COORDINATOR, leaveGroupResult)
+    val result = groupCoordinator.handleDeleteGroups(Set(groupId))
+    assert(result.size == 1 && result.contains(groupId) && result.get(groupId).contains(Errors.NON_EMPTY_GROUP))
   }
 
   @Test
-  def testLeaveGroupUnknownGroup() {
+  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))
+  }
 
-    val leaveGroupResult = leaveGroup(groupId, memberId)
-    assertEquals(Errors.UNKNOWN_MEMBER_ID, leaveGroupResult)
+  @Test
+  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 testLeaveGroupUnknownConsumerExistingGroup() {
+  def testDeleteEmptyGroup(): Unit = {
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
-    val otherMemberId = "consumerId"
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
 
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
-    val joinGroupError = joinGroupResult.error
-    assertEquals(Errors.NONE, joinGroupError)
+    EasyMock.reset(replicaManager)
+    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)
-    val leaveGroupResult = leaveGroup(groupId, otherMemberId)
-    assertEquals(Errors.UNKNOWN_MEMBER_ID, leaveGroupResult)
+    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 result = groupCoordinator.handleDeleteGroups(Set(groupId))
+    assert(result.size == 1 && result.contains(groupId) && result.get(groupId).contains(Errors.NONE))
   }
 
   @Test
-  def testValidLeaveGroup() {
+  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)
 
     EasyMock.reset(replicaManager)
-    val leaveGroupResult = leaveGroup(groupId, assignedMemberId)
-    assertEquals(Errors.NONE, leaveGroupResult)
+    val syncGroupResult = syncGroupLeader(groupId, joinGroupResult.generationId, assignedMemberId, Map(assignedMemberId -> Array[Byte]()))
+    assertEquals(Errors.NONE, syncGroupResult.error)
+
+    EasyMock.reset(replicaManager)
+    val tp = new TopicPartition("topic", 0)
+    val offset = offsetAndMetadata(0)
+    val commitOffsetResult = commitOffsets(groupId, assignedMemberId, joinGroupResult.generationId, Map(tp -> offset))
+    assertEquals(Errors.NONE, commitOffsetResult(tp))
+
+    val describeGroupResult = groupCoordinator.handleDescribeGroup(groupId)
+    assertEquals(Stable.toString, describeGroupResult._2.state)
+    assertEquals(assignedMemberId, describeGroupResult._2.members.head.memberId)
+
+    EasyMock.reset(replicaManager)
+    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(HostedPartition.Online(partition))
+    EasyMock.expect(replicaManager.nonOfflinePartition(groupTopicPartition)).andStubReturn(Some(partition))
+    EasyMock.replay(replicaManager, partition)
+
+    val result = groupCoordinator.handleDeleteGroups(Set(groupId))
+    assert(result.size == 1 && result.contains(groupId) && result.get(groupId).contains(Errors.NONE))
+
+    assertEquals(Dead.toString, groupCoordinator.handleDescribeGroup(groupId)._2.state)
   }
 
   @Test
-  def testListGroupsIncludesStableGroups() {
-    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
-    val assignedMemberId = joinGroupResult.memberId
-    val generationId = joinGroupResult.generationId
-    assertEquals(Errors.NONE, joinGroupResult.error)
+  def testDeleteOffsetOfNonExistingGroup(): Unit = {
+    val tp = new TopicPartition("foo", 0)
+    val (groupError, topics) = groupCoordinator.handleDeleteOffsets(groupId, Seq(tp))
 
-    EasyMock.reset(replicaManager)
-    val syncGroupResult = syncGroupLeader(groupId, generationId, assignedMemberId, Map(assignedMemberId -> Array[Byte]()))
-    val syncGroupError = syncGroupResult._2
-    assertEquals(Errors.NONE, syncGroupError)
+    assertEquals(Errors.GROUP_ID_NOT_FOUND, groupError)
+    assertTrue(topics.isEmpty)
+  }
 
-    val (error, groups) = groupCoordinator.handleListGroups()
-    assertEquals(Errors.NONE, error)
-    assertEquals(1, groups.size)
-    assertEquals(GroupOverview("groupId", "consumer"), groups.head)
+  @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 testListGroupsIncludesRebalancingGroups() {
+  def testDeleteOffsetOfEmptyNonConsumerGroup(): Unit = {
+    // join the group
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
+
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, "My Protocol", protocols)
     assertEquals(Errors.NONE, joinGroupResult.error)
 
-    val (error, groups) = groupCoordinator.handleListGroups()
-    assertEquals(Errors.NONE, error)
-    assertEquals(1, groups.size)
-    assertEquals(GroupOverview("groupId", "consumer"), groups.head)
+    EasyMock.reset(replicaManager)
+    val syncGroupResult = syncGroupLeader(groupId, joinGroupResult.generationId, joinGroupResult.leaderId, Map.empty)
+    assertEquals(Errors.NONE, syncGroupResult.error)
+
+    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, requireStable, 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 testDescribeGroupWrongCoordinator() {
+  def testDeleteOffsetOfConsumerGroupWithUnparsableProtocol(): Unit = {
+    val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
+
     EasyMock.reset(replicaManager)
-    val (error, _) = groupCoordinator.handleDescribeGroup(otherGroupId)
-    assertEquals(Errors.NOT_COORDINATOR, error)
+    val syncGroupResult = syncGroupLeader(groupId, joinGroupResult.generationId, joinGroupResult.leaderId, Map.empty)
+    assertEquals(Errors.NONE, syncGroupResult.error)
+
+    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 testDescribeGroupInactiveGroup() {
-    EasyMock.reset(replicaManager)
-    val (error, summary) = groupCoordinator.handleDescribeGroup(groupId)
-    assertEquals(Errors.NONE, error)
-    assertEquals(GroupCoordinator.DeadGroup, summary)
+  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 testDescribeGroupStable() {
+  def testDeleteOffsetOfEmptyConsumerGroup(): Unit = {
+    // join the group
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
-    val assignedMemberId = joinGroupResult.memberId
-    val generationId = joinGroupResult.generationId
-    val joinGroupError = joinGroupResult.error
-    assertEquals(Errors.NONE, joinGroupError)
+    val joinGroupResult = dynamicJoinGroup(groupId, memberId, protocolType, protocols)
+    assertEquals(Errors.NONE, joinGroupResult.error)
 
     EasyMock.reset(replicaManager)
-    val syncGroupResult = syncGroupLeader(groupId, generationId, assignedMemberId, Map(assignedMemberId -> Array[Byte]()))
+    val syncGroupResult = syncGroupLeader(groupId, joinGroupResult.generationId, joinGroupResult.leaderId, Map.empty)
+    assertEquals(Errors.NONE, syncGroupResult.error)
 
-    val syncGroupError = syncGroupResult._2
-    assertEquals(Errors.NONE, syncGroupError)
+    val t1p0 = new TopicPartition("foo", 0)
+    val t2p0 = new TopicPartition("bar", 0)
+    val offset = offsetAndMetadata(37)
 
     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))
+    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, requireStable, 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 testDescribeGroupRebalancing() {
+  def testDeleteOffsetOfStableConsumerGroup(): Unit = {
+    // join the group
     val memberId = JoinGroupRequest.UNKNOWN_MEMBER_ID
-    val joinGroupResult = joinGroup(groupId, memberId, protocolType, protocols)
-    val joinGroupError = joinGroupResult.error
-    assertEquals(Errors.NONE, joinGroupError)
+    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 (error, summary) = groupCoordinator.handleDescribeGroup(groupId)
-    assertEquals(Errors.NONE, error)
-    assertEquals(protocolType, summary.protocolType)
-    assertEquals(GroupCoordinator.NoProtocol, summary.protocol)
-    assertEquals(CompletingRebalance.toString, summary.state)
-    assertTrue(summary.members.map(_.memberId).contains(joinGroupResult.memberId))
-    assertTrue(summary.members.forall(_.metadata.isEmpty))
-    assertTrue(summary.members.forall(_.assignment.isEmpty))
+    val syncGroupResult = syncGroupLeader(groupId, joinGroupResult.generationId, joinGroupResult.leaderId, Map.empty)
+    assertEquals(Errors.NONE, syncGroupResult.error)
+
+    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, requireStable, 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() {
+  def shouldDelayInitialRebalanceByGroupInitialRebalanceDelayOnEmptyGroup(): Unit = {
     val firstJoinFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols)
     timer.advanceClock(GroupInitialRebalanceDelay / 2)
     verifyDelayedTaskNotCompleted(firstJoinFuture)
@@ -1260,12 +3721,12 @@ 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)
+    val firstMemberJoinFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols, rebalanceTimeout = rebalanceTimeout)
     EasyMock.reset(replicaManager)
     timer.advanceClock(GroupInitialRebalanceDelay - 1)
-    val secondMemberJoinFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols, rebalanceTimeout)
+    val secondMemberJoinFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols, rebalanceTimeout = rebalanceTimeout)
     EasyMock.reset(replicaManager)
     timer.advanceClock(2)
 
@@ -1284,14 +3745,14 @@ 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)
+    val firstMemberJoinFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols, rebalanceTimeout = rebalanceTimeout)
     EasyMock.reset(replicaManager)
-    val secondMemberJoinFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols, rebalanceTimeout)
+    val secondMemberJoinFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols, rebalanceTimeout = rebalanceTimeout)
     timer.advanceClock(GroupInitialRebalanceDelay + 1)
     EasyMock.reset(replicaManager)
-    val thirdMemberJoinFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols, rebalanceTimeout)
+    val thirdMemberJoinFuture = sendJoinGroup(groupId, JoinGroupRequest.UNKNOWN_MEMBER_ID, protocolType, protocols, rebalanceTimeout = rebalanceTimeout)
     timer.advanceClock(GroupInitialRebalanceDelay)
     EasyMock.reset(replicaManager)
 
@@ -1310,63 +3771,136 @@ class GroupCoordinatorTest extends JUnitSuite {
     assertEquals(Errors.NONE, thirdResult.error)
   }
 
+  @Test
+  def testCompleteHeartbeatWithGroupDead(): Unit = {
+    val rebalanceResult = staticMembersJoinAndRebalance(leaderInstanceId, followerInstanceId)
+    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)
+    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 responsePromise = Promise[JoinGroupResult]()
     val responseFuture = responsePromise.future
-    val responseCallback: JoinGroupCallback = responsePromise.success(_)
+    val responseCallback: JoinGroupCallback = responsePromise.success
     (responseFuture, responseCallback)
   }
 
-  private def setupSyncGroupCallback: (Future[SyncGroupCallbackParams], SyncGroupCallback) = {
-    val responsePromise = Promise[SyncGroupCallbackParams]
+  private def setupSyncGroupCallback: (Future[SyncGroupResult], SyncGroupCallback) = {
+    val responsePromise = Promise[SyncGroupResult]()
     val responseFuture = responsePromise.future
-    val responseCallback: SyncGroupCallback = (assignment, error) =>
-      responsePromise.success((assignment, error))
+    val responseCallback: SyncGroupCallback = responsePromise.success
     (responseFuture, responseCallback)
   }
 
   private def setupHeartbeatCallback: (Future[HeartbeatCallbackParams], HeartbeatCallback) = {
-    val responsePromise = Promise[HeartbeatCallbackParams]
+    val responsePromise = Promise[HeartbeatCallbackParams]()
     val responseFuture = responsePromise.future
     val responseCallback: HeartbeatCallback = error => responsePromise.success(error)
     (responseFuture, responseCallback)
   }
 
   private def setupCommitOffsetsCallback: (Future[CommitOffsetCallbackParams], CommitOffsetCallback) = {
-    val responsePromise = Promise[CommitOffsetCallbackParams]
+    val responsePromise = Promise[CommitOffsetCallbackParams]()
     val responseFuture = responsePromise.future
     val responseCallback: CommitOffsetCallback = offsets => responsePromise.success(offsets)
     (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,
-                            sessionTimeout: Int = DefaultSessionTimeout): Future[JoinGroupResult] = {
+                            requireKnownMemberId: Boolean = false): Future[JoinGroupResult] = {
     val (responseFuture, responseCallback) = setupJoinGroupCallback
 
+    EasyMock.expect(replicaManager.getMagic(EasyMock.anyObject())).andReturn(Some(RecordBatch.MAGIC_VALUE_V1)).anyTimes()
     EasyMock.replay(replicaManager)
 
-    groupCoordinator.handleJoinGroup(groupId, memberId, "clientId", "clientHost", rebalanceTimeout, sessionTimeout,
-      protocolType, protocols, responseCallback)
+    groupCoordinator.handleJoinGroup(groupId, memberId, groupInstanceId,
+      requireKnownMemberId, "clientId", "clientHost", rebalanceTimeout, sessionTimeout, protocolType, protocols, responseCallback)
     responseFuture
   }
 
+  private def sendStaticJoinGroupWithPersistence(groupId: String,
+                                                 memberId: String,
+                                                 protocolType: String,
+                                                 protocols: List[(String, Array[Byte])],
+                                                 groupInstanceId: Option[String],
+                                                 sessionTimeout: Int,
+                                                 rebalanceTimeout: Int,
+                                                 appendRecordError: Errors,
+                                                 requireKnownMemberId: Boolean = false): Future[JoinGroupResult] = {
+    val (responseFuture, responseCallback) = setupJoinGroupCallback
+
+    val capturedArgument: Capture[scala.collection.Map[TopicPartition, PartitionResponse] => Unit] = EasyMock.newCapture()
+
+    EasyMock.expect(replicaManager.appendRecords(EasyMock.anyLong(),
+      EasyMock.anyShort(),
+      internalTopicsAllowed = EasyMock.eq(true),
+      origin = EasyMock.eq(AppendOrigin.Coordinator),
+      EasyMock.anyObject().asInstanceOf[Map[TopicPartition, MemoryRecords]],
+      EasyMock.capture(capturedArgument),
+      EasyMock.anyObject().asInstanceOf[Option[ReentrantLock]],
+      EasyMock.anyObject())).andAnswer(new IAnswer[Unit] {
+      override def answer = capturedArgument.getValue.apply(
+        Map(new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, groupPartitionId) ->
+          new PartitionResponse(appendRecordError, 0L, RecordBatch.NO_TIMESTAMP, 0L)
+        )
+      )})
+    EasyMock.expect(replicaManager.getMagic(EasyMock.anyObject())).andReturn(Some(RecordBatch.MAGIC_VALUE_V1)).anyTimes()
+    EasyMock.replay(replicaManager)
+
+    groupCoordinator.handleJoinGroup(groupId, memberId, groupInstanceId,
+      requireKnownMemberId, "clientId", "clientHost", rebalanceTimeout, sessionTimeout, protocolType, protocols, responseCallback)
+    responseFuture
+  }
 
   private def sendSyncGroupLeader(groupId: String,
                                   generation: Int,
                                   leaderId: String,
-                                  assignment: Map[String, Array[Byte]]): Future[SyncGroupCallbackParams] = {
+                                  protocolType: Option[String],
+                                  protocolName: Option[String],
+                                  groupInstanceId: Option[String],
+                                  assignment: Map[String, Array[Byte]]): Future[SyncGroupResult] = {
     val (responseFuture, responseCallback) = setupSyncGroupCallback
 
-    val capturedArgument: Capture[Map[TopicPartition, PartitionResponse] => Unit] = EasyMock.newCapture()
+    val capturedArgument: Capture[scala.collection.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]],
@@ -1379,39 +3913,90 @@ 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, protocolType, protocolName,
+      groupInstanceId, assignment, responseCallback)
     responseFuture
   }
 
   private def sendSyncGroupFollower(groupId: String,
                                     generation: Int,
-                                    memberId: String): Future[SyncGroupCallbackParams] = {
+                                    memberId: String,
+                                    prototolType: Option[String] = None,
+                                    prototolName: Option[String] = None,
+                                    groupInstanceId: Option[String] = None): Future[SyncGroupResult] = {
     val (responseFuture, responseCallback) = setupSyncGroupCallback
 
     EasyMock.replay(replicaManager)
 
-    groupCoordinator.handleSyncGroup(groupId, generation, memberId, Map.empty[String, Array[Byte]], responseCallback)
+    groupCoordinator.handleSyncGroup(groupId, generation, memberId,
+      prototolType, prototolName, 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 = {
-    val responseFuture = sendJoinGroup(groupId, memberId, protocolType, protocols, rebalanceTimeout, sessionTimeout)
+  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, 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) {
+      val joinGroupResult = Await.result(responseFuture, Duration(rebalanceTimeout + 100, TimeUnit.MILLISECONDS))
+      // If some other error is triggered, return the error immediately for caller to handle.
+      if (joinGroupResult.error != Errors.MEMBER_ID_REQUIRED) {
+        return joinGroupResult
+      }
+      EasyMock.reset(replicaManager)
+      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 staticJoinGroupWithPersistence(groupId: String,
+                                             memberId: String,
+                                             groupInstanceId: Option[String],
+                                             protocolType: String,
+                                             protocols: List[(String, Array[Byte])],
+                                             clockAdvance: Int,
+                                             sessionTimeout: Int = DefaultSessionTimeout,
+                                             rebalanceTimeout: Int = DefaultRebalanceTimeout,
+                                             appendRecordError: Errors = Errors.NONE): JoinGroupResult = {
+    val responseFuture = sendStaticJoinGroupWithPersistence(groupId, memberId, protocolType, protocols, groupInstanceId, sessionTimeout, rebalanceTimeout, appendRecordError)
+
+    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,
-                                sessionTimeout: Int = DefaultSessionTimeout): SyncGroupCallbackParams = {
-    val responseFuture = sendSyncGroupFollower(groupId, generationId, memberId)
+                                protocolType: Option[String] = None,
+                                protocolName: Option[String] = None,
+                                groupInstanceId: Option[String] = None,
+                                sessionTimeout: Int = DefaultSessionTimeout): SyncGroupResult = {
+    val responseFuture = sendSyncGroupFollower(groupId, generationId, memberId, protocolType,
+      protocolName, groupInstanceId)
     Await.result(responseFuture, Duration(sessionTimeout + 100, TimeUnit.MILLISECONDS))
   }
 
@@ -1419,19 +4004,24 @@ class GroupCoordinatorTest extends JUnitSuite {
                               generationId: Int,
                               memberId: String,
                               assignment: Map[String, Array[Byte]],
-                              sessionTimeout: Int = DefaultSessionTimeout): SyncGroupCallbackParams = {
-    val responseFuture = sendSyncGroupLeader(groupId, generationId, memberId, assignment)
+                              protocolType: Option[String] = None,
+                              protocolName: Option[String] = None,
+                              groupInstanceId: Option[String] = None,
+                              sessionTimeout: Int = DefaultSessionTimeout): SyncGroupResult = {
+    val responseFuture = sendSyncGroupLeader(groupId, generationId, memberId, protocolType,
+      protocolName, 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))
   }
 
@@ -1440,17 +4030,18 @@ class GroupCoordinatorTest extends JUnitSuite {
   }
 
   private def commitOffsets(groupId: String,
-                            consumerId: String,
+                            memberId: String,
                             generationId: Int,
-                            offsets: immutable.Map[TopicPartition, OffsetAndMetadata]): CommitOffsetCallbackParams = {
+                            offsets: Map[TopicPartition, OffsetAndMetadata],
+                            groupInstanceId: Option[String] = None): CommitOffsetCallbackParams = {
     val (responseFuture, responseCallback) = setupCommitOffsetsCallback
 
-    val capturedArgument: Capture[Map[TopicPartition, PartitionResponse] => Unit] = EasyMock.newCapture()
+    val capturedArgument: Capture[scala.collection.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]],
@@ -1465,22 +4056,25 @@ 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))
   }
 
   private def commitTransactionalOffsets(groupId: String,
                                          producerId: Long,
                                          producerEpoch: Short,
-                                         offsets: immutable.Map[TopicPartition, OffsetAndMetadata]): CommitOffsetCallbackParams = {
+                                         offsets: Map[TopicPartition, OffsetAndMetadata],
+                                         memberId: String = JoinGroupRequest.UNKNOWN_MEMBER_ID,
+                                         groupInstanceId: Option[String] = Option.empty,
+                                         generationId: Int = JoinGroupRequest.UNKNOWN_GENERATION_ID) = {
     val (responseFuture, responseCallback) = setupCommitOffsetsCallback
 
-    val capturedArgument: Capture[Map[TopicPartition, PartitionResponse] => Unit] = EasyMock.newCapture()
+    val capturedArgument: Capture[scala.collection.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]],
@@ -1494,21 +4088,58 @@ class GroupCoordinatorTest extends JUnitSuite {
     EasyMock.expect(replicaManager.getMagic(EasyMock.anyObject())).andReturn(Some(RecordBatch.MAGIC_VALUE_V2)).anyTimes()
     EasyMock.replay(replicaManager)
 
-    groupCoordinator.handleTxnCommitOffsets(groupId, producerId, producerEpoch, offsets, responseCallback)
+    groupCoordinator.handleTxnCommitOffsets(groupId, producerId, producerEpoch,
+      memberId, groupInstanceId, generationId, offsets, responseCallback)
     val result = Await.result(responseFuture, Duration(40, TimeUnit.MILLISECONDS))
     EasyMock.reset(replicaManager)
     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))
   }
 
+  def handleTxnCompletion(producerId: Long,
+                          offsetsPartitions: Iterable[TopicPartition],
+                          transactionResult: TransactionResult): Unit = {
+    val isCommit = transactionResult == TransactionResult.COMMIT
+    groupCoordinator.groupManager.handleTxnCompletion(producerId, offsetsPartitions.map(_.partition).toSet, isCommit)
+  }
+
+  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 abedcb879c875..89de104ca4140 100644
--- a/core/src/test/scala/unit/kafka/coordinator/group/GroupMetadataManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/coordinator/group/GroupMetadataManagerTest.scala
@@ -17,32 +17,39 @@
 
 package kafka.coordinator.group
 
-import kafka.api.ApiVersion
+import java.lang.management.ManagementFactory
+import java.nio.ByteBuffer
+import java.util.concurrent.locks.ReentrantLock
+import java.util.{Collections, Optional}
+
+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.metrics.KafkaYammerMetrics
+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.common.TopicPartition
+import org.apache.kafka.common.internals.Topic
+import org.apache.kafka.common.metrics.{JmxReporter, KafkaMetricsContext, Metrics => kMetrics}
 import org.apache.kafka.common.protocol.Errors
 import org.apache.kafka.common.record._
-import org.apache.kafka.common.requests.{IsolationLevel, OffsetFetchResponse}
+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, assertTrue}
+import org.junit.Assert.{assertEquals, assertFalse, assertNull, assertTrue, assertThrows}
 import org.junit.{Before, Test}
-import java.nio.ByteBuffer
-
-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.jdk.CollectionConverters._
 import scala.collection._
-import java.util.concurrent.locks.ReentrantLock
-
-import kafka.zk.KafkaZkClient
 
 class GroupMetadataManagerTest {
 
@@ -50,21 +57,22 @@ class GroupMetadataManagerTest {
   var replicaManager: ReplicaManager = null
   var groupMetadataManager: GroupMetadataManager = null
   var scheduler: KafkaScheduler = null
-  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"
   val rebalanceTimeout = 60000
   val sessionTimeout = 10000
+  val defaultRequireStable = false
 
-  @Before
-  def setUp() {
+  private val offsetConfig = {
     val config = KafkaConfig.fromProps(TestUtils.createBrokerConfig(nodeId = 0, zkConnect = ""))
-
-    val offsetConfig = OffsetConfig(maxMetadataSize = config.offsetMetadataMaxSize,
+    OffsetConfig(maxMetadataSize = config.offsetMetadataMaxSize,
       loadBufferSize = config.offsetsLoadBufferSize,
       offsetsRetentionMs = config.offsetsRetentionMinutes * 60 * 1000L,
       offsetsRetentionCheckIntervalMs = config.offsetsRetentionCheckIntervalMs,
@@ -74,20 +82,49 @@ class GroupMetadataManagerTest {
       offsetsTopicCompressionCodec = config.offsetsTopicCompressionCodec,
       offsetCommitTimeoutMs = config.offsetCommitTimeoutMs,
       offsetCommitRequiredAcks = config.offsetCommitRequiredAcks)
+  }
 
+  private def mockKafkaZkClient: KafkaZkClient = {
     // make two partitions of the group topic to make sure some partitions are not owned by the coordinator
-    zkClient = EasyMock.createNiceMock(classOf[KafkaZkClient])
+    val zkClient: KafkaZkClient = EasyMock.createNiceMock(classOf[KafkaZkClient])
     EasyMock.expect(zkClient.getTopicPartitionCount(Topic.GROUP_METADATA_TOPIC_NAME)).andReturn(Some(2))
     EasyMock.replay(zkClient)
+    zkClient
+  }
 
+  @Before
+  def setUp(): Unit = {
+    defaultOffsetRetentionMs = offsetConfig.offsetsRetentionMs
+    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,
+      mockKafkaZkClient, time, metrics)
     partition = EasyMock.niceMock(classOf[Partition])
   }
 
   @Test
-  def testLoadOffsetsWithoutGroup() {
+  def testLogInfoFromCleanupGroupMetadata(): Unit = {
+    var expiredOffsets: Int = 0
+    var infoCount = 0
+    val gmm = new GroupMetadataManager(0, ApiVersion.latestVersion, offsetConfig, replicaManager, mockKafkaZkClient, time, metrics) {
+      override def cleanupGroupMetadata(groups: Iterable[GroupMetadata],
+                                        selector: GroupMetadata => Map[TopicPartition, OffsetAndMetadata]): Int = expiredOffsets
+
+      override def info(msg: => String): Unit = infoCount += 1
+    }
+
+    // if there are no offsets to expire, we skip to log
+    gmm.cleanupGroupMetadata()
+    assertEquals(0, infoCount)
+    // if there are offsets to expire, we should log info
+    expiredOffsets = 100
+    gmm.cleanupGroupMetadata()
+    assertEquals(1, infoCount)
+  }
+
+  @Test
+  def testLoadOffsetsWithoutGroup(): Unit = {
     val groupMetadataTopicPartition = groupTopicPartition
     val startOffset = 15L
 
@@ -98,12 +135,12 @@ 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)
 
-    groupMetadataManager.loadGroupsAndOffsets(groupMetadataTopicPartition, _ => ())
+    groupMetadataManager.loadGroupsAndOffsets(groupMetadataTopicPartition, _ => (), 0L)
 
     val group = groupMetadataManager.getGroup(groupId).getOrElse(fail("Group was not loaded into the cache"))
     assertEquals(groupId, group.groupId)
@@ -115,7 +152,42 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testLoadTransactionalOffsetsWithoutGroup() {
+  def testLoadEmptyGroupWithOffsets(): Unit = {
+    val groupMetadataTopicPartition = groupTopicPartition
+    val generation = 15
+    val protocolType = "consumer"
+    val startOffset = 15L
+    val committedOffsets = Map(
+      new TopicPartition("foo", 0) -> 23L,
+      new TopicPartition("foo", 1) -> 455L,
+      new TopicPartition("bar", 0) -> 8992L
+    )
+
+    val offsetCommitRecords = createCommittedOffsetRecords(committedOffsets)
+    val groupMetadataRecord = buildEmptyGroupRecord(generation, protocolType)
+    val records = MemoryRecords.withRecords(startOffset, CompressionType.NONE,
+      (offsetCommitRecords ++ Seq(groupMetadataRecord)).toArray: _*)
+
+    expectGroupMetadataLoad(groupMetadataTopicPartition, startOffset, records)
+
+    EasyMock.replay(replicaManager)
+
+    groupMetadataManager.loadGroupsAndOffsets(groupMetadataTopicPartition, _ => (), 0L)
+
+    val group = groupMetadataManager.getGroup(groupId).getOrElse(fail("Group was not loaded into the cache"))
+    assertEquals(groupId, group.groupId)
+    assertEquals(Empty, group.currentState)
+    assertEquals(generation, group.generationId)
+    assertEquals(Some(protocolType), group.protocolType)
+    assertNull(group.leaderOrNull)
+    assertNull(group.protocolName.orNull)
+    committedOffsets.foreach { case (topicPartition, offset) =>
+      assertEquals(Some(offset), group.offset(topicPartition).map(_.offset))
+    }
+  }
+
+  @Test
+  def testLoadTransactionalOffsetsWithoutGroup(): Unit = {
     val groupMetadataTopicPartition = groupTopicPartition
     val producerId = 1000L
     val producerEpoch: Short = 2
@@ -137,7 +209,7 @@ class GroupMetadataManagerTest {
 
     EasyMock.replay(replicaManager)
 
-    groupMetadataManager.loadGroupsAndOffsets(groupMetadataTopicPartition, _ => ())
+    groupMetadataManager.loadGroupsAndOffsets(groupMetadataTopicPartition, _ => (), 0L)
 
     val group = groupMetadataManager.getGroup(groupId).getOrElse(fail("Group was not loaded into the cache"))
     assertEquals(groupId, group.groupId)
@@ -149,7 +221,7 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testDoNotLoadAbortedTransactionalOffsetCommits() {
+  def testDoNotLoadAbortedTransactionalOffsetCommits(): Unit = {
     val groupMetadataTopicPartition = groupTopicPartition
     val producerId = 1000L
     val producerEpoch: Short = 2
@@ -171,7 +243,7 @@ class GroupMetadataManagerTest {
 
     EasyMock.replay(replicaManager)
 
-    groupMetadataManager.loadGroupsAndOffsets(groupMetadataTopicPartition, _ => ())
+    groupMetadataManager.loadGroupsAndOffsets(groupMetadataTopicPartition, _ => (), 0L)
 
     // Since there are no committed offsets for the group, and there is no other group metadata, we don't expect the
     // group to be loaded.
@@ -184,10 +256,13 @@ class GroupMetadataManagerTest {
     val producerId = 1000L
     val producerEpoch: Short = 2
 
+    val foo0 = new TopicPartition("foo", 0)
+    val foo1 = new TopicPartition("foo", 1)
+    val bar0 = new TopicPartition("bar", 0)
     val pendingOffsets = Map(
-      new TopicPartition("foo", 0) -> 23L,
-      new TopicPartition("foo", 1) -> 455L,
-      new TopicPartition("bar", 0) -> 8992L
+      foo0 -> 23L,
+      foo1 -> 455L,
+      bar0 -> 8992L
     )
 
     val buffer = ByteBuffer.allocate(1024)
@@ -200,7 +275,7 @@ class GroupMetadataManagerTest {
 
     EasyMock.replay(replicaManager)
 
-    groupMetadataManager.loadGroupsAndOffsets(groupMetadataTopicPartition, _ => ())
+    groupMetadataManager.loadGroupsAndOffsets(groupMetadataTopicPartition, _ => (), 0L)
 
     // The group should be loaded with pending offsets.
     val group = groupMetadataManager.getGroup(groupId).getOrElse(fail("Group was not loaded into the cache"))
@@ -210,10 +285,13 @@ class GroupMetadataManagerTest {
     assertEquals(0, group.allOffsets.size)
     assertTrue(group.hasOffsets)
     assertTrue(group.hasPendingOffsetCommitsFromProducer(producerId))
+    assertTrue(group.hasPendingOffsetCommitsForTopicPartition(foo0))
+    assertTrue(group.hasPendingOffsetCommitsForTopicPartition(foo1))
+    assertTrue(group.hasPendingOffsetCommitsForTopicPartition(bar0))
   }
 
   @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
@@ -244,7 +322,7 @@ class GroupMetadataManagerTest {
 
     EasyMock.replay(replicaManager)
 
-    groupMetadataManager.loadGroupsAndOffsets(groupMetadataTopicPartition, _ => ())
+    groupMetadataManager.loadGroupsAndOffsets(groupMetadataTopicPartition, _ => (), 0L)
 
     val group = groupMetadataManager.getGroup(groupId).getOrElse(fail("Group was not loaded into the cache"))
     assertEquals(groupId, group.groupId)
@@ -259,7 +337,7 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testLoadWithCommittedAndAbortedAndPendingTransactionalOffsetCommits() {
+  def testLoadWithCommittedAndAbortedAndPendingTransactionalOffsetCommits(): Unit = {
     val groupMetadataTopicPartition = groupTopicPartition
     val producerId = 1000L
     val producerEpoch: Short = 2
@@ -270,14 +348,16 @@ class GroupMetadataManagerTest {
       new TopicPartition("bar", 0) -> 8992L
     )
 
+    val foo3 = new TopicPartition("foo", 3)
+
     val abortedOffsets = Map(
       new TopicPartition("foo", 2) -> 231L,
-      new TopicPartition("foo", 3) -> 4551L,
+      foo3 -> 4551L,
       new TopicPartition("bar", 1) -> 89921L
     )
 
     val pendingOffsets = Map(
-      new TopicPartition("foo", 3) -> 2312L,
+      foo3 -> 2312L,
       new TopicPartition("foo", 4) -> 45512L,
       new TopicPartition("bar", 2) -> 899212L
     )
@@ -297,7 +377,7 @@ class GroupMetadataManagerTest {
 
     EasyMock.replay(replicaManager)
 
-    groupMetadataManager.loadGroupsAndOffsets(groupMetadataTopicPartition, _ => ())
+    groupMetadataManager.loadGroupsAndOffsets(groupMetadataTopicPartition, _ => (), 0L)
 
     val group = groupMetadataManager.getGroup(groupId).getOrElse(fail("Group was not loaded into the cache"))
     assertEquals(groupId, group.groupId)
@@ -313,6 +393,7 @@ class GroupMetadataManagerTest {
 
     // We should have pending commits.
     assertTrue(group.hasPendingOffsetCommitsFromProducer(producerId))
+    assertTrue(group.hasPendingOffsetCommitsForTopicPartition(foo3))
 
     // The loaded pending commits should materialize after a commit marker comes in.
     groupMetadataManager.handleTxnCompletion(producerId, List(groupMetadataTopicPartition.partition).toSet, isCommit = true)
@@ -359,7 +440,7 @@ class GroupMetadataManagerTest {
 
     EasyMock.replay(replicaManager)
 
-    groupMetadataManager.loadGroupsAndOffsets(groupMetadataTopicPartition, _ => ())
+    groupMetadataManager.loadGroupsAndOffsets(groupMetadataTopicPartition, _ => (), 0L)
 
     val group = groupMetadataManager.getGroup(groupId).getOrElse(fail("Group was not loaded into the cache"))
     assertEquals(groupId, group.groupId)
@@ -405,7 +486,7 @@ class GroupMetadataManagerTest {
 
     EasyMock.replay(replicaManager)
 
-    groupMetadataManager.loadGroupsAndOffsets(groupMetadataTopicPartition, _ => ())
+    groupMetadataManager.loadGroupsAndOffsets(groupMetadataTopicPartition, _ => (), 0L)
 
     // The group should be loaded with pending offsets.
     val group = groupMetadataManager.getGroup(groupId).getOrElse(fail("Group was not loaded into the cache"))
@@ -447,7 +528,7 @@ class GroupMetadataManagerTest {
 
     EasyMock.replay(replicaManager)
 
-    groupMetadataManager.loadGroupsAndOffsets(groupMetadataTopicPartition, _ => ())
+    groupMetadataManager.loadGroupsAndOffsets(groupMetadataTopicPartition, _ => (), 0L)
 
     // The group should be loaded with pending offsets.
     val group = groupMetadataManager.getGroup(groupId).getOrElse(fail("Group was not loaded into the cache"))
@@ -462,6 +543,26 @@ class GroupMetadataManagerTest {
     }
   }
 
+  @Test
+  def testGroupNotExists(): Unit = {
+    // group is not owned
+    assertFalse(groupMetadataManager.groupNotExists(groupId))
+
+    groupMetadataManager.addPartitionOwnership(groupPartitionId)
+    // group is owned but does not exist yet
+    assertTrue(groupMetadataManager.groupNotExists(groupId))
+
+    val group = new GroupMetadata(groupId, Empty, time)
+    groupMetadataManager.addGroup(group)
+
+    // group is owned but not Dead
+    assertFalse(groupMetadataManager.groupNotExists(groupId))
+
+    group.transitionTo(Dead)
+    // group is owned and Dead
+    assertTrue(groupMetadataManager.groupNotExists(groupId))
+  }
+
   private def appendConsumerOffsetCommit(buffer: ByteBuffer, baseOffset: Long, offsets: Map[TopicPartition, Long]) = {
     val builder = MemoryRecords.builder(buffer, CompressionType.NONE, TimestampType.LOG_APPEND_TIME, baseOffset)
     val commitRecords = createCommittedOffsetRecords(offsets)
@@ -491,7 +592,7 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testLoadOffsetsWithTombstones() {
+  def testLoadOffsetsWithTombstones(): Unit = {
     val groupMetadataTopicPartition = groupTopicPartition
     val startOffset = 15L
 
@@ -505,13 +606,13 @@ 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)
 
     EasyMock.replay(replicaManager)
 
-    groupMetadataManager.loadGroupsAndOffsets(groupMetadataTopicPartition, _ => ())
+    groupMetadataManager.loadGroupsAndOffsets(groupMetadataTopicPartition, _ => (), 0L)
 
     val group = groupMetadataManager.getGroup(groupId).getOrElse(fail("Group was not loaded into the cache"))
     assertEquals(groupId, group.groupId)
@@ -526,8 +627,11 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testLoadOffsetsAndGroup() {
+  def testLoadOffsetsAndGroup(): Unit = {
     val groupMetadataTopicPartition = groupTopicPartition
+    val generation = 935
+    val protocolType = "consumer"
+    val protocol = "range"
     val startOffset = 15L
     val committedOffsets = Map(
       new TopicPartition("foo", 0) -> 23L,
@@ -537,47 +641,107 @@ class GroupMetadataManagerTest {
 
     val offsetCommitRecords = createCommittedOffsetRecords(committedOffsets)
     val memberId = "98098230493"
-    val groupMetadataRecord = buildStableGroupRecordWithMember(memberId)
+    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)
 
     EasyMock.replay(replicaManager)
 
-    groupMetadataManager.loadGroupsAndOffsets(groupMetadataTopicPartition, _ => ())
+    groupMetadataManager.loadGroupsAndOffsets(groupMetadataTopicPartition, _ => (), 0L)
 
     val group = groupMetadataManager.getGroup(groupId).getOrElse(fail("Group was not loaded into the cache"))
     assertEquals(groupId, group.groupId)
     assertEquals(Stable, group.currentState)
-    assertEquals(memberId, group.leaderId)
+    assertEquals(memberId, group.leaderOrNull)
+    assertEquals(generation, group.generationId)
+    assertEquals(Some(protocolType), group.protocolType)
+    assertEquals(protocol, group.protocolName.orNull)
     assertEquals(Set(memberId), group.allMembers)
     assertEquals(committedOffsets.size, group.allOffsets.size)
     committedOffsets.foreach { case (topicPartition, offset) =>
       assertEquals(Some(offset), group.offset(topicPartition).map(_.offset))
+      assertTrue(group.offset(topicPartition).map(_.expireTimestamp).contains(None))
     }
   }
 
   @Test
-  def testLoadGroupWithTombstone() {
+  def testLoadGroupWithTombstone(): Unit = {
     val groupMetadataTopicPartition = groupTopicPartition
     val startOffset = 15L
-
     val memberId = "98098230493"
-    val groupMetadataRecord = buildStableGroupRecordWithMember(memberId)
+    val groupMetadataRecord = buildStableGroupRecordWithMember(generation = 15,
+      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)
 
     EasyMock.replay(replicaManager)
 
-    groupMetadataManager.loadGroupsAndOffsets(groupMetadataTopicPartition, _ => ())
+    groupMetadataManager.loadGroupsAndOffsets(groupMetadataTopicPartition, _ => (), 0L)
 
     assertEquals(None, groupMetadataManager.getGroup(groupId))
   }
 
+  @Test
+  def testLoadGroupWithLargeGroupMetadataRecord(): Unit = {
+    val groupMetadataTopicPartition = groupTopicPartition
+    val startOffset = 15L
+    val committedOffsets = Map(
+      new TopicPartition("foo", 0) -> 23L,
+      new TopicPartition("foo", 1) -> 455L,
+      new TopicPartition("bar", 0) -> 8992L
+    )
+
+    // create a GroupMetadata record larger then offsets.load.buffer.size (here at least 16 bytes larger)
+    val assignmentSize = OffsetConfig.DefaultLoadBufferSize + 16
+    val memberId = "98098230493"
+
+    val offsetCommitRecords = createCommittedOffsetRecords(committedOffsets)
+    val groupMetadataRecord = buildStableGroupRecordWithMember(generation = 15,
+      protocolType = "consumer", protocol = "range", memberId, new Array[Byte](assignmentSize))
+    val records = MemoryRecords.withRecords(startOffset, CompressionType.NONE,
+      (offsetCommitRecords ++ Seq(groupMetadataRecord)).toArray: _*)
+
+    expectGroupMetadataLoad(groupMetadataTopicPartition, startOffset, records)
+
+    EasyMock.replay(replicaManager)
+
+    groupMetadataManager.loadGroupsAndOffsets(groupMetadataTopicPartition, _ => (), 0L)
+
+    val group = groupMetadataManager.getGroup(groupId).getOrElse(fail("Group was not loaded into the cache"))
+    committedOffsets.foreach { case (topicPartition, offset) =>
+      assertEquals(Some(offset), group.offset(topicPartition).map(_.offset))
+    }
+  }
+
+  @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, _ => (), 0L)
+
+    EasyMock.verify(logMock)
+    EasyMock.verify(replicaManager)
+
+    assertFalse(groupMetadataManager.isPartitionLoading(groupTopicPartition.partition()))
+  }
+
   @Test
   def testOffsetWriteAfterGroupRemoved(): Unit = {
     // this test case checks the following scenario:
@@ -585,6 +749,9 @@ class GroupMetadataManagerTest {
     // 2. a "simple" consumer (i.e. not a consumer group) then uses the same groupId to commit some offsets
 
     val groupMetadataTopicPartition = groupTopicPartition
+    val generation = 293
+    val protocolType = "consumer"
+    val protocol = "range"
     val startOffset = 15L
 
     val committedOffsets = Map(
@@ -594,18 +761,18 @@ class GroupMetadataManagerTest {
     )
     val offsetCommitRecords = createCommittedOffsetRecords(committedOffsets)
     val memberId = "98098230493"
-    val groupMetadataRecord = buildStableGroupRecordWithMember(memberId)
+    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)
 
     EasyMock.replay(replicaManager)
 
-    groupMetadataManager.loadGroupsAndOffsets(groupMetadataTopicPartition, _ => ())
+    groupMetadataManager.loadGroupsAndOffsets(groupMetadataTopicPartition, _ => (), 0L)
 
-    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)
@@ -616,38 +783,43 @@ class GroupMetadataManagerTest {
 
   @Test
   def testLoadGroupAndOffsetsFromDifferentSegments(): Unit = {
+    val generation = 293
+    val protocolType = "consumer"
+    val protocol = "range"
     val startOffset = 15L
     val tp0 = new TopicPartition("foo", 0)
     val tp1 = new TopicPartition("foo", 1)
     val tp2 = new TopicPartition("bar", 0)
     val tp3 = new TopicPartition("xxx", 0)
 
-    val logMock =  EasyMock.mock(classOf[Log])
+    val logMock: Log = EasyMock.mock(classOf[Log])
     EasyMock.expect(replicaManager.getLog(groupTopicPartition)).andStubReturn(Some(logMock))
 
     val segment1MemberId = "a"
     val segment1Offsets = Map(tp0 -> 23L, tp1 -> 455L, tp3 -> 42L)
     val segment1Records = MemoryRecords.withRecords(startOffset, CompressionType.NONE,
-      createCommittedOffsetRecords(segment1Offsets) ++ Seq(buildStableGroupRecordWithMember(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(segment2MemberId)): _*)
+      (createCommittedOffsetRecords(segment2Offsets) ++ Seq(buildStableGroupRecordWithMember(
+        generation, protocolType, protocol, segment2MemberId))).toArray: _*)
     val segment2End = expectGroupMetadataLoad(logMock, segment1End, segment2Records)
 
     EasyMock.expect(replicaManager.getLogEndOffset(groupTopicPartition)).andStubReturn(Some(segment2End))
 
     EasyMock.replay(logMock, replicaManager)
 
-    groupMetadataManager.loadGroupsAndOffsets(groupTopicPartition, _ => ())
+    groupMetadataManager.loadGroupsAndOffsets(groupTopicPartition, _ => (), 0L)
 
     val group = groupMetadataManager.getGroup(groupId).getOrElse(fail("Group was not loaded into the cache"))
     assertEquals(groupId, group.groupId)
     assertEquals(Stable, group.currentState)
 
-    assertEquals("segment2 group record member should be elected", segment2MemberId, group.leaderId)
+    assertEquals("segment2 group record member should be elected", segment2MemberId, group.leaderOrNull)
     assertEquals("segment2 group record member should be only member", Set(segment2MemberId), group.allMembers)
 
     // offsets of segment1 should be overridden by segment2 offsets of the same topic partitions
@@ -658,55 +830,250 @@ class GroupMetadataManagerTest {
     }
   }
 
-
   @Test
-  def testAddGroup() {
-    val group = new GroupMetadata("foo")
+  def testAddGroup(): Unit = {
+    val group = new GroupMetadata("foo", Empty, time)
     assertEquals(group, groupMetadataManager.addGroup(group))
-    assertEquals(group, groupMetadataManager.addGroup(new GroupMetadata("foo")))
+    assertEquals(group, groupMetadataManager.addGroup(new GroupMetadata("foo", Empty, time)))
+  }
+
+  @Test
+  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.protocolName.orNull)
+    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.protocolName.orNull)
+    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.protocolName.orNull)
+    assertEquals(None, group.getSubscribedTopics)
+    assertTrue(group.has(memberId))
+  }
+
+  @Test
+  def testShouldThrowExceptionForUnsupportedGroupMetadataVersion(): Unit = {
+    val generation = 1
+    val protocol = "range"
+    val memberId = "memberId"
+    val unsupportedVersion = Short.MinValue
+
+    // put the unsupported version as the version value
+    val groupMetadataRecordValue = buildStableGroupRecordWithMember(generation, protocolType, protocol, memberId)
+      .value().putShort(unsupportedVersion)
+    // reset the position to the starting position 0 so that it can read the data in correct order
+    groupMetadataRecordValue.position(0)
+
+    val e = assertThrows(classOf[IllegalStateException],
+      () => GroupMetadataManager.readGroupMessageValue(groupId, groupMetadataRecordValue, time))
+    assertEquals(s"Unknown group metadata message version: ${unsupportedVersion}", e.getMessage)
+  }
+
+  @Test
+  def testCurrentStateTimestampForAllGroupMetadataVersions(): Unit = {
+    val generation = 1
+    val protocol = "range"
+    val memberId = "memberId"
+
+    for (apiVersion <- ApiVersion.allVersions) {
+      val groupMetadataRecord = buildStableGroupRecordWithMember(generation, protocolType, protocol, memberId, apiVersion = apiVersion)
+
+      val deserializedGroupMetadata = GroupMetadataManager.readGroupMessageValue(groupId, groupMetadataRecord.value(), time)
+      // GROUP_METADATA_VALUE_SCHEMA_V2 or higher should correctly set the currentStateTimestamp
+      if (apiVersion >= KAFKA_2_1_IV0)
+        assertEquals(s"the apiVersion $apiVersion doesn't set the currentStateTimestamp correctly.",
+          Some(time.milliseconds()), deserializedGroupMetadata.currentStateTimestamp)
+      else
+        assertTrue(s"the apiVersion $apiVersion should not set the currentStateTimestamp.",
+          deserializedGroupMetadata.currentStateTimestamp.isEmpty)
+    }
+  }
+
+  @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.protocolName.orNull)
+      assertEquals(1, deserializedGroupMetadata.allMembers.size)
+      assertEquals(deserializedGroupMetadata.allMembers, deserializedGroupMetadata.allDynamicMembers)
+      assertTrue(deserializedGroupMetadata.allMembers.contains(memberId))
+      assertTrue(deserializedGroupMetadata.allStaticMembers.isEmpty)
+    }
   }
 
   @Test
-  def testStoreEmptyGroup() {
-    val group = new GroupMetadata(groupId)
+  def testStoreEmptyGroup(): Unit = {
+    val generation = 27
+    val protocolType = "consumer"
+
+    val group = GroupMetadata.loadGroup(groupId, Empty, generation, protocolType, null, null, None, Seq.empty, time)
     groupMetadataManager.addGroup(group)
 
-    expectAppendMessage(Errors.NONE)
+    val capturedRecords = expectAppendMessage(Errors.NONE)
+    EasyMock.replay(replicaManager)
+
+    var maybeError: Option[Errors] = None
+    def callback(error: Errors): Unit = {
+      maybeError = Some(error)
+    }
+
+    groupMetadataManager.storeGroup(group, Map.empty, callback)
+    assertEquals(Some(Errors.NONE), maybeError)
+    assertTrue(capturedRecords.hasCaptured)
+    val records = capturedRecords.getValue()(new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, groupPartitionId))
+        .records.asScala.toList
+    assertEquals(1, records.size)
+
+    val record = records.head
+    val groupMetadata = GroupMetadataManager.readGroupMessageValue(groupId, record.value, time)
+    assertTrue(groupMetadata.is(Empty))
+    assertEquals(generation, groupMetadata.generationId)
+    assertEquals(Some(protocolType), groupMetadata.protocolType)
+  }
+
+  @Test
+  def testStoreEmptySimpleGroup(): Unit = {
+    val group = new GroupMetadata(groupId, Empty, time)
+    groupMetadataManager.addGroup(group)
+
+    val capturedRecords = expectAppendMessage(Errors.NONE)
     EasyMock.replay(replicaManager)
 
     var maybeError: Option[Errors] = None
-    def callback(error: Errors) {
+    def callback(error: Errors): Unit = {
       maybeError = Some(error)
     }
 
     groupMetadataManager.storeGroup(group, Map.empty, callback)
     assertEquals(Some(Errors.NONE), maybeError)
+    assertTrue(capturedRecords.hasCaptured)
+
+    assertTrue(capturedRecords.hasCaptured)
+    val records = capturedRecords.getValue()(new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, groupPartitionId))
+      .records.asScala.toList
+    assertEquals(1, records.size)
+
+    val record = records.head
+    val groupMetadata = GroupMetadataManager.readGroupMessageValue(groupId, record.value, time)
+    assertTrue(groupMetadata.is(Empty))
+    assertEquals(0, groupMetadata.generationId)
+    assertEquals(None, groupMetadata.protocolType)
   }
 
   @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)
     assertStoreGroupErrorMapping(Errors.NOT_ENOUGH_REPLICAS_AFTER_APPEND, Errors.COORDINATOR_NOT_AVAILABLE)
-    assertStoreGroupErrorMapping(Errors.NOT_LEADER_FOR_PARTITION, Errors.NOT_COORDINATOR)
+    assertStoreGroupErrorMapping(Errors.NOT_LEADER_OR_FOLLOWER, Errors.NOT_COORDINATOR)
     assertStoreGroupErrorMapping(Errors.MESSAGE_TOO_LARGE, Errors.UNKNOWN_SERVER_ERROR)
     assertStoreGroupErrorMapping(Errors.RECORD_LIST_TOO_LARGE, Errors.UNKNOWN_SERVER_ERROR)
     assertStoreGroupErrorMapping(Errors.INVALID_FETCH_SIZE, Errors.UNKNOWN_SERVER_ERROR)
     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)
+    val group = new GroupMetadata(groupId, Empty, time)
     groupMetadataManager.addGroup(group)
 
     expectAppendMessage(appendError)
     EasyMock.replay(replicaManager)
 
     var maybeError: Option[Errors] = None
-    def callback(error: Errors) {
+    def callback(error: Errors): Unit = {
       maybeError = Some(error)
     }
 
@@ -717,18 +1084,17 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testStoreNonEmptyGroup() {
+  def testStoreNonEmptyGroup(): Unit = {
     val memberId = "memberId"
     val clientId = "clientId"
     val clientHost = "localhost"
 
-    val group = new GroupMetadata(groupId)
+    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]())))
-    member.awaitingJoinCallback = _ => ()
-    group.add(member)
+    group.add(member, _ => ())
     group.transitionTo(PreparingRebalance)
     group.initNextGeneration()
 
@@ -736,7 +1102,7 @@ class GroupMetadataManagerTest {
     EasyMock.replay(replicaManager)
 
     var maybeError: Option[Errors] = None
-    def callback(error: Errors) {
+    def callback(error: Errors): Unit = {
       maybeError = Some(error)
     }
 
@@ -747,25 +1113,24 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testStoreNonEmptyGroupWhenCoordinatorHasMoved() {
+  def testStoreNonEmptyGroupWhenCoordinatorHasMoved(): Unit = {
     EasyMock.expect(replicaManager.getMagic(EasyMock.anyObject())).andReturn(None)
     val memberId = "memberId"
     val clientId = "clientId"
     val clientHost = "localhost"
 
-    val group = new GroupMetadata(groupId)
+    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]())))
-    member.awaitingJoinCallback = _ => ()
-    group.add(member)
+    group.add(member, _ => ())
     group.transitionTo(PreparingRebalance)
     group.initNextGeneration()
 
     EasyMock.replay(replicaManager)
 
     var maybeError: Option[Errors] = None
-    def callback(error: Errors) {
+    def callback(error: Errors): Unit = {
       maybeError = Some(error)
     }
 
@@ -776,23 +1141,23 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testCommitOffset() {
+  def testCommitOffset(): Unit = {
     val memberId = ""
     val topicPartition = new TopicPartition("foo", 0)
     val offset = 37
 
     groupMetadataManager.addPartitionOwnership(groupPartitionId)
 
-    val group = new GroupMetadata(groupId)
+    val group = new GroupMetadata(groupId, Empty, time)
     groupMetadataManager.addGroup(group)
 
-    val offsets = immutable.Map(topicPartition -> OffsetAndMetadata(offset))
+    val offsets = immutable.Map(topicPartition -> OffsetAndMetadata(offset, "", time.milliseconds()))
 
     expectAppendMessage(Errors.NONE)
     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)
     }
 
@@ -804,7 +1169,7 @@ class GroupMetadataManagerTest {
     assertEquals(Some(Errors.NONE), maybeError)
     assertTrue(group.hasOffsets)
 
-    val cachedOffsets = groupMetadataManager.getOffsets(groupId, Some(Seq(topicPartition)))
+    val cachedOffsets = groupMetadataManager.getOffsets(groupId, defaultRequireStable, Some(Seq(topicPartition)))
     val maybePartitionResponse = cachedOffsets.get(topicPartition)
     assertFalse(maybePartitionResponse.isEmpty)
 
@@ -816,7 +1181,7 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testTransactionalCommitOffsetCommitted() {
+  def testTransactionalCommitOffsetCommitted(): Unit = {
     val memberId = ""
     val topicPartition = new TopicPartition("foo", 0)
     val offset = 37
@@ -825,16 +1190,17 @@ class GroupMetadataManagerTest {
 
     groupMetadataManager.addPartitionOwnership(groupPartitionId)
 
-    val group = new GroupMetadata(groupId)
+    val group = new GroupMetadata(groupId, Empty, time)
     groupMetadataManager.addGroup(group)
 
-    val offsets = immutable.Map(topicPartition -> OffsetAndMetadata(offset))
+    val offsetAndMetadata = OffsetAndMetadata(offset, "", time.milliseconds())
+    val offsets = immutable.Map(topicPartition -> offsetAndMetadata)
 
     val capturedResponseCallback = appendAndCaptureCallback()
     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)
     }
 
@@ -850,13 +1216,13 @@ class GroupMetadataManagerTest {
     group.completePendingTxnOffsetCommit(producerId, isCommit = true)
     assertTrue(group.hasOffsets)
     assertFalse(group.allOffsets.isEmpty)
-    assertEquals(Some(OffsetAndMetadata(offset)), group.offset(topicPartition))
+    assertEquals(Some(offsetAndMetadata), group.offset(topicPartition))
 
     EasyMock.verify(replicaManager)
   }
 
   @Test
-  def testTransactionalCommitOffsetAppendFailure() {
+  def testTransactionalCommitOffsetAppendFailure(): Unit = {
     val memberId = ""
     val topicPartition = new TopicPartition("foo", 0)
     val offset = 37
@@ -865,16 +1231,16 @@ class GroupMetadataManagerTest {
 
     groupMetadataManager.addPartitionOwnership(groupPartitionId)
 
-    val group = new GroupMetadata(groupId)
+    val group = new GroupMetadata(groupId, Empty, time)
     groupMetadataManager.addGroup(group)
 
-    val offsets = immutable.Map(topicPartition -> OffsetAndMetadata(offset))
+    val offsets = immutable.Map(topicPartition -> OffsetAndMetadata(offset, "", time.milliseconds()))
 
     val capturedResponseCallback = appendAndCaptureCallback()
     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)
     }
 
@@ -895,7 +1261,7 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testTransactionalCommitOffsetAborted() {
+  def testTransactionalCommitOffsetAborted(): Unit = {
     val memberId = ""
     val topicPartition = new TopicPartition("foo", 0)
     val offset = 37
@@ -904,16 +1270,16 @@ class GroupMetadataManagerTest {
 
     groupMetadataManager.addPartitionOwnership(groupPartitionId)
 
-    val group = new GroupMetadata(groupId)
+    val group = new GroupMetadata(groupId, Empty, time)
     groupMetadataManager.addGroup(group)
 
-    val offsets = immutable.Map(topicPartition -> OffsetAndMetadata(offset))
+    val offsets = immutable.Map(topicPartition -> OffsetAndMetadata(offset, "", time.milliseconds()))
 
     val capturedResponseCallback = appendAndCaptureCallback()
     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)
     }
 
@@ -934,7 +1300,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)
@@ -942,15 +1308,15 @@ class GroupMetadataManagerTest {
 
     groupMetadataManager.addPartitionOwnership(groupPartitionId)
 
-    val group = new GroupMetadata(groupId)
+    val group = new GroupMetadata(groupId, Empty, time)
     groupMetadataManager.addGroup(group)
 
-    val offsets = immutable.Map(topicPartition -> OffsetAndMetadata(offset))
+    val offsets = immutable.Map(topicPartition -> OffsetAndMetadata(offset, "", time.milliseconds()))
 
     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,11 +1330,11 @@ 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)
-    assertCommitOffsetErrorMapping(Errors.NOT_LEADER_FOR_PARTITION, Errors.NOT_COORDINATOR)
+    assertCommitOffsetErrorMapping(Errors.NOT_LEADER_OR_FOLLOWER, Errors.NOT_COORDINATOR)
     assertCommitOffsetErrorMapping(Errors.MESSAGE_TOO_LARGE, Errors.INVALID_COMMIT_OFFSET_SIZE)
     assertCommitOffsetErrorMapping(Errors.RECORD_LIST_TOO_LARGE, Errors.INVALID_COMMIT_OFFSET_SIZE)
     assertCommitOffsetErrorMapping(Errors.INVALID_FETCH_SIZE, Errors.INVALID_COMMIT_OFFSET_SIZE)
@@ -984,16 +1350,16 @@ class GroupMetadataManagerTest {
 
     groupMetadataManager.addPartitionOwnership(groupPartitionId)
 
-    val group = new GroupMetadata(groupId)
+    val group = new GroupMetadata(groupId, Empty, time)
     groupMetadataManager.addGroup(group)
 
-    val offsets = immutable.Map(topicPartition -> OffsetAndMetadata(offset))
+    val offsets = immutable.Map(topicPartition -> OffsetAndMetadata(offset, "", time.milliseconds()))
 
     val capturedResponseCallback = appendAndCaptureCallback()
     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)
     }
 
@@ -1007,14 +1373,14 @@ class GroupMetadataManagerTest {
     assertEquals(Some(expectedError), maybeError)
     assertFalse(group.hasOffsets)
 
-    val cachedOffsets = groupMetadataManager.getOffsets(groupId, Some(Seq(topicPartition)))
+    val cachedOffsets = groupMetadataManager.getOffsets(groupId, defaultRequireStable, Some(Seq(topicPartition)))
     assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), cachedOffsets.get(topicPartition).map(_.offset))
 
     EasyMock.verify(replicaManager)
   }
 
   @Test
-  def testExpireOffset() {
+  def testExpireOffset(): Unit = {
     val memberId = ""
     val topicPartition1 = new TopicPartition("foo", 0)
     val topicPartition2 = new TopicPartition("foo", 1)
@@ -1022,7 +1388,7 @@ class GroupMetadataManagerTest {
 
     groupMetadataManager.addPartitionOwnership(groupPartitionId)
 
-    val group = new GroupMetadata(groupId)
+    val group = new GroupMetadata(groupId, Empty, time)
     groupMetadataManager.addGroup(group)
 
     // expire the offset after 1 millisecond
@@ -1036,7 +1402,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)
     }
 
@@ -1051,7 +1417,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)
 
@@ -1061,7 +1427,7 @@ class GroupMetadataManagerTest {
     assertEquals(None, group.offset(topicPartition1))
     assertEquals(Some(offset), group.offset(topicPartition2).map(_.offset))
 
-    val cachedOffsets = groupMetadataManager.getOffsets(groupId, Some(Seq(topicPartition1, topicPartition2)))
+    val cachedOffsets = groupMetadataManager.getOffsets(groupId, defaultRequireStable, Some(Seq(topicPartition1, topicPartition2)))
     assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), cachedOffsets.get(topicPartition1).map(_.offset))
     assertEquals(Some(offset), cachedOffsets.get(topicPartition2).map(_.offset))
 
@@ -1069,13 +1435,13 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testGroupMetadataRemoval() {
+  def testGroupMetadataRemoval(): Unit = {
     val topicPartition1 = new TopicPartition("foo", 0)
     val topicPartition2 = new TopicPartition("foo", 1)
 
     groupMetadataManager.addPartitionOwnership(groupPartitionId)
 
-    val group = new GroupMetadata(groupId)
+    val group = new GroupMetadata(groupId, Empty, time)
     groupMetadataManager.addGroup(group)
     group.generationId = 5
 
@@ -1086,7 +1452,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)
 
@@ -1095,7 +1461,7 @@ class GroupMetadataManagerTest {
     assertTrue(recordsCapture.hasCaptured)
 
     val records = recordsCapture.getValue.records.asScala.toList
-    recordsCapture.getValue.batches.asScala.foreach { batch =>
+    recordsCapture.getValue.batches.forEach { batch =>
       assertEquals(RecordBatch.CURRENT_MAGIC_VALUE, batch.magic)
       assertEquals(TimestampType.CREATE_TIME, batch.timestampType)
     }
@@ -1111,19 +1477,19 @@ class GroupMetadataManagerTest {
 
     // the full group should be gone since all offsets were removed
     assertEquals(None, groupMetadataManager.getGroup(groupId))
-    val cachedOffsets = groupMetadataManager.getOffsets(groupId, Some(Seq(topicPartition1, topicPartition2)))
+    val cachedOffsets = groupMetadataManager.getOffsets(groupId, defaultRequireStable, Some(Seq(topicPartition1, topicPartition2)))
     assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), cachedOffsets.get(topicPartition1).map(_.offset))
     assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), cachedOffsets.get(topicPartition2).map(_.offset))
   }
 
   @Test
-  def testGroupMetadataRemovalWithLogAppendTime() {
+  def testGroupMetadataRemovalWithLogAppendTime(): Unit = {
     val topicPartition1 = new TopicPartition("foo", 0)
     val topicPartition2 = new TopicPartition("foo", 1)
 
     groupMetadataManager.addPartitionOwnership(groupPartitionId)
 
-    val group = new GroupMetadata(groupId)
+    val group = new GroupMetadata(groupId, Empty, time)
     groupMetadataManager.addGroup(group)
     group.generationId = 5
 
@@ -1134,7 +1500,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)
 
@@ -1143,7 +1509,7 @@ class GroupMetadataManagerTest {
     assertTrue(recordsCapture.hasCaptured)
 
     val records = recordsCapture.getValue.records.asScala.toList
-    recordsCapture.getValue.batches.asScala.foreach { batch =>
+    recordsCapture.getValue.batches.forEach { batch =>
       assertEquals(RecordBatch.CURRENT_MAGIC_VALUE, batch.magic)
       // Use CREATE_TIME, like the producer. The conversion to LOG_APPEND_TIME (if necessary) happens automatically.
       assertEquals(TimestampType.CREATE_TIME, batch.timestampType)
@@ -1160,13 +1526,13 @@ class GroupMetadataManagerTest {
 
     // the full group should be gone since all offsets were removed
     assertEquals(None, groupMetadataManager.getGroup(groupId))
-    val cachedOffsets = groupMetadataManager.getOffsets(groupId, Some(Seq(topicPartition1, topicPartition2)))
+    val cachedOffsets = groupMetadataManager.getOffsets(groupId, defaultRequireStable, Some(Seq(topicPartition1, topicPartition2)))
     assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), cachedOffsets.get(topicPartition1).map(_.offset))
     assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), cachedOffsets.get(topicPartition2).map(_.offset))
   }
 
   @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
 
@@ -1177,13 +1543,13 @@ class GroupMetadataManagerTest {
 
     groupMetadataManager.addPartitionOwnership(groupPartitionId)
 
-    val group = new GroupMetadata(groupId)
+    val group = new GroupMetadata(groupId, Empty, time)
     groupMetadataManager.addGroup(group)
 
     // expire the offset after 1 millisecond
     val startMs = time.milliseconds
     val offsets = immutable.Map(
-      topicPartition1 -> OffsetAndMetadata(offset, "", startMs, startMs + 1),
+      topicPartition1 -> OffsetAndMetadata(offset, Optional.empty(), "", startMs, Some(startMs + 1)),
       topicPartition2 -> OffsetAndMetadata(offset, "", startMs, startMs + 3))
 
     mockGetPartition()
@@ -1191,7 +1557,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)
     }
 
@@ -1209,7 +1575,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)
 
@@ -1230,7 +1596,7 @@ class GroupMetadataManagerTest {
 
     // the full group should be gone since all offsets were removed
     assertEquals(None, groupMetadataManager.getGroup(groupId))
-    val cachedOffsets = groupMetadataManager.getOffsets(groupId, Some(Seq(topicPartition1, topicPartition2)))
+    val cachedOffsets = groupMetadataManager.getOffsets(groupId, defaultRequireStable, Some(Seq(topicPartition1, topicPartition2)))
     assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), cachedOffsets.get(topicPartition1).map(_.offset))
     assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), cachedOffsets.get(topicPartition2).map(_.offset))
 
@@ -1238,38 +1604,45 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testExpireOffsetsWithActiveGroup() {
+  def testOffsetExpirationSemantics(): Unit = {
     val memberId = "memberId"
     val clientId = "clientId"
     val clientHost = "localhost"
-    val topicPartition1 = new TopicPartition("foo", 0)
-    val topicPartition2 = new TopicPartition("foo", 1)
+    val topic = "foo"
+    val topicPartition1 = new TopicPartition(topic, 0)
+    val topicPartition2 = new TopicPartition(topic, 1)
+    val topicPartition3 = new TopicPartition(topic, 2)
     val offset = 37
 
     groupMetadataManager.addPartitionOwnership(groupPartitionId)
 
-    val group = new GroupMetadata(groupId)
+    val group = new GroupMetadata(groupId, Empty, time)
     groupMetadataManager.addGroup(group)
 
-    val member = new MemberMetadata(memberId, groupId, clientId, clientHost, rebalanceTimeout, sessionTimeout,
-      protocolType, List(("protocol", Array[Byte]())))
-    member.awaitingJoinCallback = _ => ()
-    group.add(member)
+    val subscription = new Subscription(List(topic).asJava)
+    val member = new MemberMetadata(memberId, groupId, groupInstanceId, clientId, clientHost, rebalanceTimeout, sessionTimeout,
+      protocolType, List(("protocol", ConsumerProtocol.serializeSubscription(subscription).array())))
+    group.add(member, _ => ())
     group.transitionTo(PreparingRebalance)
     group.initNextGeneration()
 
-    // expire the offset after 1 millisecond
     val startMs = time.milliseconds
+    // old clients, expiry timestamp is explicitly set
+    val tp1OffsetAndMetadata = OffsetAndMetadata(offset, "", startMs, startMs + 1)
+    val tp2OffsetAndMetadata = OffsetAndMetadata(offset, "", startMs, startMs + 3)
+    // new clients, no per-partition expiry timestamp, offsets of group expire together
+    val tp3OffsetAndMetadata = OffsetAndMetadata(offset, "", startMs)
     val offsets = immutable.Map(
-      topicPartition1 -> OffsetAndMetadata(offset, "", startMs, startMs + 1),
-      topicPartition2 -> OffsetAndMetadata(offset, "", startMs, startMs + 3))
+      topicPartition1 -> tp1OffsetAndMetadata,
+      topicPartition2 -> tp2OffsetAndMetadata,
+      topicPartition3 -> tp3OffsetAndMetadata)
 
     mockGetPartition()
     expectAppendMessage(Errors.NONE)
     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)
     }
 
@@ -1279,28 +1652,613 @@ class GroupMetadataManagerTest {
     assertFalse(commitErrors.isEmpty)
     assertEquals(Some(Errors.NONE), commitErrors.get.get(topicPartition1))
 
-    // expire all of the offsets
-    time.sleep(4)
+    // do not expire any offset even though expiration timestamp is reached for one (due to group still being active)
+    time.sleep(2)
+
+    groupMetadataManager.cleanupGroupMetadata()
+
+    // group and offsets should still be there
+    assertEquals(Some(group), groupMetadataManager.getGroup(groupId))
+    assertEquals(Some(tp1OffsetAndMetadata), group.offset(topicPartition1))
+    assertEquals(Some(tp2OffsetAndMetadata), group.offset(topicPartition2))
+    assertEquals(Some(tp3OffsetAndMetadata), group.offset(topicPartition3))
+
+    var cachedOffsets = groupMetadataManager.getOffsets(groupId, defaultRequireStable, Some(Seq(topicPartition1, topicPartition2, topicPartition3)))
+    assertEquals(Some(offset), cachedOffsets.get(topicPartition1).map(_.offset))
+    assertEquals(Some(offset), cachedOffsets.get(topicPartition2).map(_.offset))
+    assertEquals(Some(offset), cachedOffsets.get(topicPartition3).map(_.offset))
+
+    EasyMock.verify(replicaManager)
+
+    group.transitionTo(PreparingRebalance)
+    group.transitionTo(Empty)
 
     // 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)
 
     groupMetadataManager.cleanupGroupMetadata()
 
-    // group should still be there, but the offsets should be gone
+    // group is empty now, only one offset should expire
+    assertEquals(Some(group), groupMetadataManager.getGroup(groupId))
+    assertEquals(None, group.offset(topicPartition1))
+    assertEquals(Some(tp2OffsetAndMetadata), group.offset(topicPartition2))
+    assertEquals(Some(tp3OffsetAndMetadata), group.offset(topicPartition3))
+
+    cachedOffsets = groupMetadataManager.getOffsets(groupId, defaultRequireStable, Some(Seq(topicPartition1, topicPartition2, topicPartition3)))
+    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), cachedOffsets.get(topicPartition1).map(_.offset))
+    assertEquals(Some(offset), cachedOffsets.get(topicPartition2).map(_.offset))
+    assertEquals(Some(offset), cachedOffsets.get(topicPartition3).map(_.offset))
+
+    EasyMock.verify(replicaManager)
+
+    time.sleep(2)
+
+    // expect the offset tombstone
+    EasyMock.reset(partition)
+    EasyMock.expect(partition.appendRecordsToLeader(EasyMock.anyObject(classOf[MemoryRecords]),
+      origin = EasyMock.eq(AppendOrigin.Coordinator), requiredAcks = EasyMock.anyInt()))
+      .andReturn(LogAppendInfo.UnknownLogAppendInfo)
+    EasyMock.replay(partition)
+
+    groupMetadataManager.cleanupGroupMetadata()
+
+    // one more offset should expire
     assertEquals(Some(group), groupMetadataManager.getGroup(groupId))
     assertEquals(None, group.offset(topicPartition1))
     assertEquals(None, group.offset(topicPartition2))
+    assertEquals(Some(tp3OffsetAndMetadata), group.offset(topicPartition3))
 
-    val cachedOffsets = groupMetadataManager.getOffsets(groupId, Some(Seq(topicPartition1, topicPartition2)))
+    cachedOffsets = groupMetadataManager.getOffsets(groupId, defaultRequireStable, Some(Seq(topicPartition1, topicPartition2, topicPartition3)))
     assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), cachedOffsets.get(topicPartition1).map(_.offset))
     assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), cachedOffsets.get(topicPartition2).map(_.offset))
+    assertEquals(Some(offset), cachedOffsets.get(topicPartition3).map(_.offset))
+
+    EasyMock.verify(replicaManager)
+
+    // advance time to just before the offset of last partition is to be expired, no offset should expire
+    time.sleep(group.currentStateTimestamp.get + defaultOffsetRetentionMs - time.milliseconds() - 1)
+
+    groupMetadataManager.cleanupGroupMetadata()
+
+    // one more offset should expire
+    assertEquals(Some(group), groupMetadataManager.getGroup(groupId))
+    assertEquals(None, group.offset(topicPartition1))
+    assertEquals(None, group.offset(topicPartition2))
+    assertEquals(Some(tp3OffsetAndMetadata), group.offset(topicPartition3))
+
+    cachedOffsets = groupMetadataManager.getOffsets(groupId, defaultRequireStable, Some(Seq(topicPartition1, topicPartition2, topicPartition3)))
+    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), cachedOffsets.get(topicPartition1).map(_.offset))
+    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), cachedOffsets.get(topicPartition2).map(_.offset))
+    assertEquals(Some(offset), cachedOffsets.get(topicPartition3).map(_.offset))
+
+    EasyMock.verify(replicaManager)
+
+    // advance time enough for that last offset to expire
+    time.sleep(2)
+
+    // expect the offset tombstone
+    EasyMock.reset(partition)
+    EasyMock.expect(partition.appendRecordsToLeader(EasyMock.anyObject(classOf[MemoryRecords]),
+      origin = EasyMock.eq(AppendOrigin.Coordinator), requiredAcks = EasyMock.anyInt()))
+      .andReturn(LogAppendInfo.UnknownLogAppendInfo)
+    EasyMock.replay(partition)
+
+    groupMetadataManager.cleanupGroupMetadata()
+
+    // group and all its offsets should be gone now
+    assertEquals(None, groupMetadataManager.getGroup(groupId))
+    assertEquals(None, group.offset(topicPartition1))
+    assertEquals(None, group.offset(topicPartition2))
+    assertEquals(None, group.offset(topicPartition3))
+
+    cachedOffsets = groupMetadataManager.getOffsets(groupId, defaultRequireStable, Some(Seq(topicPartition1, topicPartition2, topicPartition3)))
+    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), cachedOffsets.get(topicPartition1).map(_.offset))
+    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), cachedOffsets.get(topicPartition2).map(_.offset))
+    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), cachedOffsets.get(topicPartition3).map(_.offset))
+
+    EasyMock.verify(replicaManager)
+
+    assert(group.is(Dead))
+  }
+
+  @Test
+  def testOffsetExpirationOfSimpleConsumer(): Unit = {
+    val memberId = "memberId"
+    val topic = "foo"
+    val topicPartition1 = new TopicPartition(topic, 0)
+    val offset = 37
+
+    groupMetadataManager.addPartitionOwnership(groupPartitionId)
+
+    val group = new GroupMetadata(groupId, Empty, time)
+    groupMetadataManager.addGroup(group)
+
+    // expire the offset after 1 and 3 milliseconds (old clients) and after default retention (new clients)
+    val startMs = time.milliseconds
+    // old clients, expiry timestamp is explicitly set
+    val tp1OffsetAndMetadata = OffsetAndMetadata(offset, "", startMs)
+    // new clients, no per-partition expiry timestamp, offsets of group expire together
+    val offsets = immutable.Map(
+      topicPartition1 -> tp1OffsetAndMetadata)
+
+    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(topicPartition1))
+
+    // do not expire offsets while within retention period since commit timestamp
+    val expiryTimestamp = offsets.get(topicPartition1).get.commitTimestamp + defaultOffsetRetentionMs
+    time.sleep(expiryTimestamp - time.milliseconds() - 1)
+
+    groupMetadataManager.cleanupGroupMetadata()
+
+    // group and offsets should still be there
+    assertEquals(Some(group), groupMetadataManager.getGroup(groupId))
+    assertEquals(Some(tp1OffsetAndMetadata), group.offset(topicPartition1))
+
+    var cachedOffsets = groupMetadataManager.getOffsets(groupId, defaultRequireStable, Some(Seq(topicPartition1)))
+    assertEquals(Some(offset), cachedOffsets.get(topicPartition1).map(_.offset))
+
+    EasyMock.verify(replicaManager)
+
+    // advance time to enough for offsets to expire
+    time.sleep(2)
+
+    // expect the offset tombstone
+    EasyMock.reset(partition)
+    EasyMock.expect(partition.appendRecordsToLeader(EasyMock.anyObject(classOf[MemoryRecords]),
+      origin = EasyMock.eq(AppendOrigin.Coordinator), requiredAcks = EasyMock.anyInt()))
+      .andReturn(LogAppendInfo.UnknownLogAppendInfo)
+    EasyMock.replay(partition)
+
+    groupMetadataManager.cleanupGroupMetadata()
+
+    // group and all its offsets should be gone now
+    assertEquals(None, groupMetadataManager.getGroup(groupId))
+    assertEquals(None, group.offset(topicPartition1))
+
+    cachedOffsets = groupMetadataManager.getOffsets(groupId, defaultRequireStable, Some(Seq(topicPartition1)))
+    assertEquals(Some(OffsetFetchResponse.INVALID_OFFSET), cachedOffsets.get(topicPartition1).map(_.offset))
 
     EasyMock.verify(replicaManager)
+
+    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, defaultRequireStable, 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, defaultRequireStable, 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
+    val generation = 935
+    val protocolType = "consumer"
+    val protocol = "range"
+    val startOffset = 15L
+    val committedOffsets = Map(
+      new TopicPartition("foo", 0) -> 23L,
+      new TopicPartition("foo", 1) -> 455L,
+      new TopicPartition("bar", 0) -> 8992L
+    )
+
+    val apiVersion = KAFKA_1_1_IV0
+    val offsetCommitRecords = createCommittedOffsetRecords(committedOffsets, apiVersion = apiVersion, retentionTimeOpt = Some(100))
+    val memberId = "98098230493"
+    val groupMetadataRecord = buildStableGroupRecordWithMember(generation, protocolType, protocol, memberId, apiVersion = apiVersion)
+    val records = MemoryRecords.withRecords(startOffset, CompressionType.NONE,
+      (offsetCommitRecords ++ Seq(groupMetadataRecord)).toArray: _*)
+
+    expectGroupMetadataLoad(groupMetadataTopicPartition, startOffset, records)
+
+    EasyMock.replay(replicaManager)
+
+    groupMetadataManager.loadGroupsAndOffsets(groupMetadataTopicPartition, _ => (), 0L)
+
+    val group = groupMetadataManager.getGroup(groupId).getOrElse(fail("Group was not loaded into the cache"))
+    assertEquals(groupId, group.groupId)
+    assertEquals(Stable, group.currentState)
+    assertEquals(memberId, group.leaderOrNull)
+    assertEquals(generation, group.generationId)
+    assertEquals(Some(protocolType), group.protocolType)
+    assertEquals(protocol, group.protocolName.orNull)
+    assertEquals(Set(memberId), group.allMembers)
+    assertEquals(committedOffsets.size, group.allOffsets.size)
+    committedOffsets.foreach { case (topicPartition, offset) =>
+      assertEquals(Some(offset), group.offset(topicPartition).map(_.offset))
+      assertTrue(group.offset(topicPartition).map(_.expireTimestamp).get.nonEmpty)
+    }
+  }
+
+  @Test
+  def testLoadOffsetWithExplicitRetention() = {
+    val groupMetadataTopicPartition = groupTopicPartition
+    val generation = 935
+    val protocolType = "consumer"
+    val protocol = "range"
+    val startOffset = 15L
+    val committedOffsets = Map(
+      new TopicPartition("foo", 0) -> 23L,
+      new TopicPartition("foo", 1) -> 455L,
+      new TopicPartition("bar", 0) -> 8992L
+    )
+
+    val offsetCommitRecords = createCommittedOffsetRecords(committedOffsets, retentionTimeOpt = Some(100))
+    val memberId = "98098230493"
+    val groupMetadataRecord = buildStableGroupRecordWithMember(generation, protocolType, protocol, memberId)
+    val records = MemoryRecords.withRecords(startOffset, CompressionType.NONE,
+      (offsetCommitRecords ++ Seq(groupMetadataRecord)).toArray: _*)
+
+    expectGroupMetadataLoad(groupMetadataTopicPartition, startOffset, records)
+
+    EasyMock.replay(replicaManager)
+
+    groupMetadataManager.loadGroupsAndOffsets(groupMetadataTopicPartition, _ => (), 0L)
+
+    val group = groupMetadataManager.getGroup(groupId).getOrElse(fail("Group was not loaded into the cache"))
+    assertEquals(groupId, group.groupId)
+    assertEquals(Stable, group.currentState)
+    assertEquals(memberId, group.leaderOrNull)
+    assertEquals(generation, group.generationId)
+    assertEquals(Some(protocolType), group.protocolType)
+    assertEquals(protocol, group.protocolName.orNull)
+    assertEquals(Set(memberId), group.allMembers)
+    assertEquals(committedOffsets.size, group.allOffsets.size)
+    committedOffsets.foreach { case (topicPartition, offset) =>
+      assertEquals(Some(offset), group.offset(topicPartition).map(_.offset))
+      assertTrue(group.offset(topicPartition).map(_.expireTimestamp).get.nonEmpty)
+    }
+  }
+
+  @Test
+  def testSerdeOffsetCommitValue(): Unit = {
+    val offsetAndMetadata = OffsetAndMetadata(
+      offset = 537L,
+      leaderEpoch = Optional.of(15),
+      metadata = "metadata",
+      commitTimestamp = time.milliseconds(),
+      expireTimestamp = None)
+
+    def verifySerde(apiVersion: ApiVersion, expectedOffsetCommitValueVersion: Int): Unit = {
+      val bytes = GroupMetadataManager.offsetCommitValue(offsetAndMetadata, apiVersion)
+      val buffer = ByteBuffer.wrap(bytes)
+
+      assertEquals(expectedOffsetCommitValueVersion, buffer.getShort(0).toInt)
+
+      val deserializedOffsetAndMetadata = GroupMetadataManager.readOffsetMessageValue(buffer)
+      assertEquals(offsetAndMetadata.offset, deserializedOffsetAndMetadata.offset)
+      assertEquals(offsetAndMetadata.metadata, deserializedOffsetAndMetadata.metadata)
+      assertEquals(offsetAndMetadata.commitTimestamp, deserializedOffsetAndMetadata.commitTimestamp)
+
+      // Serialization drops the leader epoch silently if an older inter-broker protocol is in use
+      val expectedLeaderEpoch = if (expectedOffsetCommitValueVersion >= 3)
+        offsetAndMetadata.leaderEpoch
+      else
+        Optional.empty()
+
+      assertEquals(expectedLeaderEpoch, deserializedOffsetAndMetadata.leaderEpoch)
+    }
+
+    for (version <- ApiVersion.allVersions) {
+      val expectedSchemaVersion = version match {
+        case v if v < KAFKA_2_1_IV0 => 1
+        case v if v < KAFKA_2_1_IV1 => 2
+        case _ => 3
+      }
+      verifySerde(version, expectedSchemaVersion)
+    }
+  }
+
+  @Test
+  def testSerdeOffsetCommitValueWithExpireTimestamp(): Unit = {
+    // If expire timestamp is set, we should always use version 1 of the offset commit
+    // value schema since later versions do not support it
+
+    val offsetAndMetadata = OffsetAndMetadata(
+      offset = 537L,
+      leaderEpoch = Optional.empty(),
+      metadata = "metadata",
+      commitTimestamp = time.milliseconds(),
+      expireTimestamp = Some(time.milliseconds() + 1000))
+
+    def verifySerde(apiVersion: ApiVersion): Unit = {
+      val bytes = GroupMetadataManager.offsetCommitValue(offsetAndMetadata, apiVersion)
+      val buffer = ByteBuffer.wrap(bytes)
+      assertEquals(1, buffer.getShort(0).toInt)
+
+      val deserializedOffsetAndMetadata = GroupMetadataManager.readOffsetMessageValue(buffer)
+      assertEquals(offsetAndMetadata, deserializedOffsetAndMetadata)
+    }
+
+    for (version <- ApiVersion.allVersions)
+      verifySerde(version)
+  }
+
+  @Test
+  def testSerdeOffsetCommitValueWithNoneExpireTimestamp(): Unit = {
+    val offsetAndMetadata = OffsetAndMetadata(
+      offset = 537L,
+      leaderEpoch = Optional.empty(),
+      metadata = "metadata",
+      commitTimestamp = time.milliseconds(),
+      expireTimestamp = None)
+
+    def verifySerde(apiVersion: ApiVersion): Unit = {
+      val bytes = GroupMetadataManager.offsetCommitValue(offsetAndMetadata, apiVersion)
+      val buffer = ByteBuffer.wrap(bytes)
+      val version = buffer.getShort(0).toInt
+      if (apiVersion < KAFKA_2_1_IV0)
+        assertEquals(1, version)
+      else if (apiVersion < KAFKA_2_1_IV1)
+        assertEquals(2, version)
+      else
+        assertEquals(3, version)
+
+      val deserializedOffsetAndMetadata = GroupMetadataManager.readOffsetMessageValue(buffer)
+      assertEquals(offsetAndMetadata, deserializedOffsetAndMetadata)
+    }
+
+    for (version <- ApiVersion.allVersions)
+      verifySerde(version)
+  }
+
+  @Test
+  def testLoadOffsetsWithEmptyControlBatch(): Unit = {
+    val groupMetadataTopicPartition = groupTopicPartition
+    val startOffset = 15L
+    val generation = 15
+
+    val committedOffsets = Map(
+      new TopicPartition("foo", 0) -> 23L,
+      new TopicPartition("foo", 1) -> 455L,
+      new TopicPartition("bar", 0) -> 8992L
+    )
+
+    val offsetCommitRecords = createCommittedOffsetRecords(committedOffsets)
+    val groupMetadataRecord = buildEmptyGroupRecord(generation, protocolType)
+    val records = MemoryRecords.withRecords(startOffset, CompressionType.NONE,
+      (offsetCommitRecords ++ Seq(groupMetadataRecord)).toArray: _*)
+
+    // Prepend empty control batch to valid records
+    val mockBatch: MutableRecordBatch = EasyMock.createMock(classOf[MutableRecordBatch])
+    EasyMock.expect(mockBatch.iterator).andReturn(Collections.emptyIterator[Record])
+    EasyMock.expect(mockBatch.isControlBatch).andReturn(true)
+    EasyMock.expect(mockBatch.isTransactional).andReturn(true)
+    EasyMock.expect(mockBatch.nextOffset).andReturn(16L)
+    EasyMock.replay(mockBatch)
+    val mockRecords: MemoryRecords = EasyMock.createMock(classOf[MemoryRecords])
+    EasyMock.expect(mockRecords.batches).andReturn((Iterable[MutableRecordBatch](mockBatch) ++ records.batches.asScala).asJava).anyTimes()
+    EasyMock.expect(mockRecords.records).andReturn(records.records()).anyTimes()
+    EasyMock.expect(mockRecords.sizeInBytes()).andReturn(DefaultRecordBatch.RECORD_BATCH_OVERHEAD + records.sizeInBytes()).anyTimes()
+    EasyMock.replay(mockRecords)
+
+    val logMock: Log = EasyMock.mock(classOf[Log])
+    EasyMock.expect(logMock.logStartOffset).andReturn(startOffset).anyTimes()
+    EasyMock.expect(logMock.read(EasyMock.eq(startOffset),
+      maxLength = EasyMock.anyInt(),
+      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))
+    EasyMock.replay(logMock)
+    EasyMock.replay(replicaManager)
+
+    groupMetadataManager.loadGroupsAndOffsets(groupMetadataTopicPartition, _ => (), 0L)
+
+    // Empty control batch should not have caused the load to fail
+    val group = groupMetadataManager.getGroup(groupId).getOrElse(fail("Group was not loaded into the cache"))
+    assertEquals(groupId, group.groupId)
+    assertEquals(Empty, group.currentState)
+    assertEquals(generation, group.generationId)
+    assertEquals(Some(protocolType), group.protocolType)
+    assertNull(group.leaderOrNull)
+    assertNull(group.protocolName.orNull)
+    committedOffsets.foreach { case (topicPartition, offset) =>
+      assertEquals(Some(offset), group.offset(topicPartition).map(_.offset))
+    }
+  }
+
+  @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] = {
@@ -1308,7 +2266,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.anyObject().asInstanceOf[Map[TopicPartition, MemoryRecords]],
       EasyMock.capture(capturedArgument),
       EasyMock.anyObject().asInstanceOf[Option[ReentrantLock]],
@@ -1318,44 +2276,53 @@ class GroupMetadataManagerTest {
     capturedArgument
   }
 
-  private def expectAppendMessage(error: Errors) {
-    val capturedArgument: Capture[Map[TopicPartition, PartitionResponse] => Unit] = EasyMock.newCapture()
+  private def expectAppendMessage(error: Errors): Capture[Map[TopicPartition, MemoryRecords]] = {
+    val capturedCallback: Capture[Map[TopicPartition, PartitionResponse] => Unit] = EasyMock.newCapture()
+    val capturedRecords: Capture[Map[TopicPartition, MemoryRecords]] = EasyMock.newCapture()
     EasyMock.expect(replicaManager.appendRecords(EasyMock.anyLong(),
       EasyMock.anyShort(),
       internalTopicsAllowed = EasyMock.eq(true),
-      isFromClient = EasyMock.eq(false),
-      EasyMock.anyObject().asInstanceOf[Map[TopicPartition, MemoryRecords]],
-      EasyMock.capture(capturedArgument),
+      origin = EasyMock.eq(AppendOrigin.Coordinator),
+      EasyMock.capture(capturedRecords),
+      EasyMock.capture(capturedCallback),
       EasyMock.anyObject().asInstanceOf[Option[ReentrantLock]],
       EasyMock.anyObject())
     ).andAnswer(new IAnswer[Unit] {
-      override def answer = capturedArgument.getValue.apply(
+      override def answer = capturedCallback.getValue.apply(
         Map(groupTopicPartition ->
           new PartitionResponse(error, 0L, RecordBatch.NO_TIMESTAMP, 0L)
         )
       )})
     EasyMock.expect(replicaManager.getMagic(EasyMock.anyObject())).andStubReturn(Some(RecordBatch.CURRENT_MAGIC_VALUE))
+    capturedRecords
   }
 
-  private def buildStableGroupRecordWithMember(memberId: String): SimpleRecord = {
-    val group = new GroupMetadata(groupId)
-    group.transitionTo(PreparingRebalance)
-    val memberProtocols = List(("roundrobin", Array.emptyByteArray))
-    val member = new MemberMetadata(memberId, groupId, "clientId", "clientHost", 30000, 10000, "consumer", memberProtocols)
-    group.add(member)
-    member.awaitingJoinCallback = _ => {}
-    group.initNextGeneration()
-    group.transitionTo(Stable)
+  private def buildStableGroupRecordWithMember(generation: Int,
+                                               protocolType: String,
+                                               protocol: String,
+                                               memberId: String,
+                                               assignmentBytes: Array[Byte] = Array.emptyByteArray,
+                                               apiVersion: ApiVersion = ApiVersion.latestVersion): SimpleRecord = {
+    val memberProtocols = List((protocol, Array.emptyByteArray))
+    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 -> assignmentBytes), apiVersion)
+    new SimpleRecord(groupMetadataKey, groupMetadataValue)
+  }
 
+  private def buildEmptyGroupRecord(generation: Int, protocolType: String): SimpleRecord = {
+    val group = GroupMetadata.loadGroup(groupId, Empty, generation, protocolType, null, null, None, Seq.empty, time)
     val groupMetadataKey = GroupMetadataManager.groupMetadataKey(groupId)
-    val groupMetadataValue = GroupMetadataManager.groupMetadataValue(group, Map(memberId -> Array.empty[Byte]))
+    val groupMetadataValue = GroupMetadataManager.groupMetadataValue(group, Map.empty, ApiVersion.latestVersion)
     new SimpleRecord(groupMetadataKey, groupMetadataValue)
   }
 
   private def expectGroupMetadataLoad(groupMetadataTopicPartition: TopicPartition,
                                       startOffset: Long,
                                       records: MemoryRecords): Unit = {
-    val logMock =  EasyMock.mock(classOf[Log])
+    val logMock: Log =  EasyMock.mock(classOf[Log])
     EasyMock.expect(replicaManager.getLog(groupMetadataTopicPartition)).andStubReturn(Some(logMock))
     val endOffset = expectGroupMetadataLoad(logMock, startOffset, records)
     EasyMock.expect(replicaManager.getLogEndOffset(groupMetadataTopicPartition)).andStubReturn(Some(endOffset))
@@ -1371,36 +2338,58 @@ class GroupMetadataManagerTest {
                                       startOffset: Long,
                                       records: MemoryRecords): Long = {
     val endOffset = startOffset + records.records.asScala.size
-    val fileRecordsMock = EasyMock.mock(classOf[FileRecords])
+    val fileRecordsMock: FileRecords = EasyMock.mock(classOf[FileRecords])
 
     EasyMock.expect(logMock.logStartOffset).andStubReturn(startOffset)
-    EasyMock.expect(logMock.read(EasyMock.eq(startOffset), EasyMock.anyInt(), EasyMock.eq(None),
-      EasyMock.eq(true), EasyMock.eq(IsolationLevel.READ_UNCOMMITTED)))
+    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(fileRecordsMock.readInto(EasyMock.anyObject(classOf[ByteBuffer]), EasyMock.anyInt()))
-      .andReturn(records.buffer)
+
+    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 = {
+        val buffer = bufferCapture.getValue
+        buffer.put(records.buffer.duplicate)
+        buffer.flip()
+      }
+    })
+
     EasyMock.replay(fileRecordsMock)
 
     endOffset
   }
 
   private def createCommittedOffsetRecords(committedOffsets: Map[TopicPartition, Long],
-                                           groupId: String = groupId): Seq[SimpleRecord] = {
+                                           groupId: String = groupId,
+                                           apiVersion: ApiVersion = ApiVersion.latestVersion,
+                                           retentionTimeOpt: Option[Long] = None): Seq[SimpleRecord] = {
     committedOffsets.map { case (topicPartition, offset) =>
-      val offsetAndMetadata = OffsetAndMetadata(offset)
+      val commitTimestamp = time.milliseconds()
+      val offsetAndMetadata = retentionTimeOpt match {
+        case Some(retentionTimeMs) =>
+          val expirationTime = commitTimestamp + retentionTimeMs
+          OffsetAndMetadata(offset, "", commitTimestamp, expirationTime)
+        case None =>
+          OffsetAndMetadata(offset, "", commitTimestamp)
+      }
       val offsetCommitKey = GroupMetadataManager.offsetCommitKey(groupId, topicPartition)
-      val offsetCommitValue = GroupMetadataManager.offsetCommitValue(offsetAndMetadata)
+      val offsetCommitValue = GroupMetadataManager.offsetCommitValue(offsetAndMetadata, apiVersion)
       new SimpleRecord(offsetCommitKey, offsetCommitValue)
     }.toSeq
   }
 
   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))
   }
 
   private def getGauge(manager: GroupMetadataManager, name: String): Gauge[Int]  = {
-    Metrics.defaultRegistry().allMetrics().get(manager.metricName(name, Map.empty)).asInstanceOf[Gauge[Int]]
+    KafkaYammerMetrics.defaultRegistry().allMetrics().get(manager.metricName(name, Map.empty)).asInstanceOf[Gauge[Int]]
   }
 
   private def expectMetrics(manager: GroupMetadataManager,
@@ -1413,10 +2402,10 @@ class GroupMetadataManagerTest {
   }
 
   @Test
-  def testMetrics() {
+  def testMetrics(): Unit = {
     groupMetadataManager.cleanupGroupMetadata()
     expectMetrics(groupMetadataManager, 0, 0, 0)
-    val group = new GroupMetadata("foo2", Stable)
+    val group = new GroupMetadata("foo2", Stable, time)
     groupMetadataManager.addGroup(group)
     expectMetrics(groupMetadataManager, 1, 0, 0)
     group.transitionTo(PreparingRebalance)
@@ -1424,4 +2413,48 @@ 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
+    val metricsContext = new KafkaMetricsContext("kafka.server")
+    reporter.contextChange(metricsContext)
+    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)
+
+    // When passed a specific start offset, assert that the measured values are in excess of that.
+    val now = time.milliseconds()
+    val diff = 1000
+    groupMetadataManager.loadGroupsAndOffsets(groupMetadataTopicPartition, _ => (), now - diff)
+    assertTrue(partitionLoadTime("partition-load-time-max") >= diff)
+    assertTrue(partitionLoadTime("partition-load-time-avg") >= diff)
+  }
 }
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 ca62bf8f04a5e..569a9cdf19986 100644
--- a/core/src/test/scala/unit/kafka/coordinator/group/GroupMetadataTest.scala
+++ b/core/src/test/scala/unit/kafka/coordinator/group/GroupMetadataTest.scala
@@ -18,51 +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.jdk.CollectionConverters._
 
 /**
  * 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() {
-    group = new GroupMetadata("groupId")
+  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)
@@ -70,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)
@@ -90,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)
@@ -112,7 +120,7 @@ class GroupMetadataTest extends JUnitSuite {
   }
 
   @Test
-  def testAwaitingRebalanceToStableTransition() {
+  def testAwaitingRebalanceToStableTransition(): Unit = {
     group.transitionTo(PreparingRebalance)
     group.transitionTo(CompletingRebalance)
     group.transitionTo(Stable)
@@ -120,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)
@@ -138,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)
@@ -169,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)
@@ -207,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)
@@ -216,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)
@@ -237,65 +245,111 @@ class GroupMetadataTest extends JUnitSuite {
   }
 
   @Test
-  def testSupportsProtocols() {
+  def testSupportsProtocols(): Unit = {
     // by default, the group supports everything
-    assertTrue(group.supportsProtocols(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])))
+    assertTrue(group.supportsProtocols(protocolType, Set("roundrobin", "range")))
 
     group.add(member)
-    assertTrue(group.supportsProtocols(Set("roundrobin", "foo")))
-    assertTrue(group.supportsProtocols(Set("range", "foo")))
-    assertFalse(group.supportsProtocols(Set("foo", "bar")))
+    group.transitionTo(PreparingRebalance)
+    assertTrue(group.supportsProtocols(protocolType, Set("roundrobin", "foo")))
+    assertTrue(group.supportsProtocols(protocolType, Set("range", "foo")))
+    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)
 
-    assertTrue(group.supportsProtocols(Set("roundrobin", "foo")))
-    assertFalse(group.supportsProtocols(Set("range", "foo")))
+    assertTrue(group.supportsProtocols(protocolType, Set("roundrobin", "foo")))
+    assertFalse(group.supportsProtocols("invalid_type", Set("roundrobin", "foo")))
+    assertFalse(group.supportsProtocols(protocolType, Set("range", "foo")))
   }
 
   @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)
-    member.awaitingJoinCallback = _ => ()
     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, _ => ())
+
     assertEquals(0, group.generationId)
-    assertNull(group.protocol)
+    assertNull(group.protocolName.orNull)
 
     group.initNextGeneration()
 
     assertEquals(1, group.generationId)
-    assertEquals("roundrobin", group.protocol)
+    assertEquals("roundrobin", group.protocolName.orNull)
   }
 
   @Test
-  def testInitNextGenerationEmptyGroup() {
+  def testInitNextGenerationEmptyGroup(): Unit = {
     assertEquals(Empty, group.currentState)
     assertEquals(0, group.generationId)
-    assertNull(group.protocol)
+    assertNull(group.protocolName.orNull)
 
     group.transitionTo(PreparingRebalance)
     group.initNextGeneration()
 
     assertEquals(1, group.generationId)
-    assertNull(group.protocol)
+    assertNull(group.protocolName.orNull)
   }
 
   @Test
   def testOffsetCommit(): Unit = {
     val partition = new TopicPartition("foo", 0)
-    val offset = OffsetAndMetadata(37)
+    val offset = offsetAndMetadata(37)
     val commitRecordOffset = 3
 
     group.prepareOffsetCommit(Map(partition -> offset))
@@ -310,7 +364,7 @@ class GroupMetadataTest extends JUnitSuite {
   @Test
   def testOffsetCommitFailure(): Unit = {
     val partition = new TopicPartition("foo", 0)
-    val offset = OffsetAndMetadata(37)
+    val offset = offsetAndMetadata(37)
 
     group.prepareOffsetCommit(Map(partition -> offset))
     assertTrue(group.hasOffsets)
@@ -324,8 +378,8 @@ class GroupMetadataTest extends JUnitSuite {
   @Test
   def testOffsetCommitFailureWithAnotherPending(): Unit = {
     val partition = new TopicPartition("foo", 0)
-    val firstOffset = OffsetAndMetadata(37)
-    val secondOffset = OffsetAndMetadata(57)
+    val firstOffset = offsetAndMetadata(37)
+    val secondOffset = offsetAndMetadata(57)
 
     group.prepareOffsetCommit(Map(partition -> firstOffset))
     assertTrue(group.hasOffsets)
@@ -346,8 +400,8 @@ class GroupMetadataTest extends JUnitSuite {
   @Test
   def testOffsetCommitWithAnotherPending(): Unit = {
     val partition = new TopicPartition("foo", 0)
-    val firstOffset = OffsetAndMetadata(37)
-    val secondOffset = OffsetAndMetadata(57)
+    val firstOffset = offsetAndMetadata(37)
+    val secondOffset = offsetAndMetadata(57)
 
     group.prepareOffsetCommit(Map(partition -> firstOffset))
     assertTrue(group.hasOffsets)
@@ -369,8 +423,8 @@ class GroupMetadataTest extends JUnitSuite {
   def testConsumerBeatsTransactionalOffsetCommit(): Unit = {
     val partition = new TopicPartition("foo", 0)
     val producerId = 13232L
-    val txnOffsetCommit = OffsetAndMetadata(37)
-    val consumerOffsetCommit = OffsetAndMetadata(57)
+    val txnOffsetCommit = offsetAndMetadata(37)
+    val consumerOffsetCommit = offsetAndMetadata(57)
 
     group.prepareTxnOffsetCommit(producerId, Map(partition -> txnOffsetCommit))
     assertTrue(group.hasOffsets)
@@ -394,8 +448,8 @@ class GroupMetadataTest extends JUnitSuite {
   def testTransactionBeatsConsumerOffsetCommit(): Unit = {
     val partition = new TopicPartition("foo", 0)
     val producerId = 13232L
-    val txnOffsetCommit = OffsetAndMetadata(37)
-    val consumerOffsetCommit = OffsetAndMetadata(57)
+    val txnOffsetCommit = offsetAndMetadata(37)
+    val consumerOffsetCommit = offsetAndMetadata(57)
 
     group.prepareTxnOffsetCommit(producerId, Map(partition -> txnOffsetCommit))
     assertTrue(group.hasOffsets)
@@ -421,8 +475,8 @@ class GroupMetadataTest extends JUnitSuite {
   def testTransactionalCommitIsAbortedAndConsumerCommitWins(): Unit = {
     val partition = new TopicPartition("foo", 0)
     val producerId = 13232L
-    val txnOffsetCommit = OffsetAndMetadata(37)
-    val consumerOffsetCommit = OffsetAndMetadata(57)
+    val txnOffsetCommit = offsetAndMetadata(37)
+    val consumerOffsetCommit = offsetAndMetadata(57)
 
     group.prepareTxnOffsetCommit(producerId, Map(partition -> txnOffsetCommit))
     assertTrue(group.hasOffsets)
@@ -449,7 +503,7 @@ class GroupMetadataTest extends JUnitSuite {
   def testFailedTxnOffsetCommitLeavesNoPendingState(): Unit = {
     val partition = new TopicPartition("foo", 0)
     val producerId = 13232L
-    val txnOffsetCommit = OffsetAndMetadata(37)
+    val txnOffsetCommit = offsetAndMetadata(37)
 
     group.prepareTxnOffsetCommit(producerId, Map(partition -> txnOffsetCommit))
     assertTrue(group.hasPendingOffsetCommitsFromProducer(producerId))
@@ -465,7 +519,110 @@ 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, JoinGroupResult(member.memberId, Errors.NONE))
+    assertTrue(invoked)
+    assertFalse(member.isAwaitingJoin)
+  }
+
+  @Test
+  def testNotInvokeJoinCallback(): Unit = {
+    group.add(member)
+
+    assertFalse(member.isAwaitingJoin)
+    group.maybeInvokeJoinCallback(member, JoinGroupResult(member.memberId, Errors.NONE))
+    assertFalse(member.isAwaitingJoin)
+  }
+
+  @Test
+  def testInvokeSyncCallback(): Unit = {
+    group.add(member)
+    member.awaitingSyncCallback = _ => {}
+
+    val invoked = group.maybeInvokeSyncCallback(member, SyncGroupResult(Errors.NONE))
+    assertTrue(invoked)
+    assertFalse(member.isAwaitingSync)
+  }
+
+  @Test
+  def testNotInvokeSyncCallback(): Unit = {
+    group.add(member)
+
+    val invoked = group.maybeInvokeSyncCallback(member, SyncGroupResult(Errors.NONE))
+    assertFalse(invoked)
+    assertFalse(member.isAwaitingSync)
+  }
+
+  @Test
+  def testHasPendingNonTxnOffsets(): Unit = {
+    val partition = new TopicPartition("foo", 0)
+    val offset = offsetAndMetadata(37)
+
+    group.prepareOffsetCommit(Map(partition -> offset))
+    assertTrue(group.hasPendingOffsetCommitsForTopicPartition(partition))
+  }
+
+  @Test
+  def testHasPendingTxnOffsets(): Unit = {
+    val txnPartition = new TopicPartition("foo", 1)
+    val offset = offsetAndMetadata(37)
+    val producerId = 5
+
+    group.prepareTxnOffsetCommit(producerId, Map(txnPartition -> offset))
+    assertTrue(group.hasPendingOffsetCommitsForTopicPartition(txnPartition))
+
+    assertFalse(group.hasPendingOffsetCommitsForTopicPartition(new TopicPartition("non-exist", 0)))
+  }
+
+  private def assertState(group: GroupMetadata, targetState: GroupState): Unit = {
     val states: Set[GroupState] = Set(Stable, PreparingRebalance, CompletingRebalance, Dead)
     val otherStates = states - targetState
     otherStates.foreach { otherState =>
@@ -473,4 +630,9 @@ class GroupMetadataTest extends JUnitSuite {
     }
     assertTrue(group.is(targetState))
   }
+
+  private def offsetAndMetadata(offset: Long): OffsetAndMetadata = {
+    OffsetAndMetadata(offset, "", Time.SYSTEM.milliseconds())
+  }
+
 }
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 c5b42d4391aca..d780925f3347b 100644
--- a/core/src/test/scala/unit/kafka/coordinator/transaction/ProducerIdManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/coordinator/transaction/ProducerIdManagerTest.scala
@@ -16,15 +16,15 @@
  */
 package kafka.coordinator.transaction
 
-import kafka.common.KafkaException
 import kafka.zk.KafkaZkClient
-import org.easymock.{Capture, EasyMock, IAnswer}
+import org.apache.kafka.common.KafkaException
+import org.easymock.{Capture, EasyMock}
 import org.junit.{After, Test}
 import org.junit.Assert._
 
 class ProducerIdManagerTest {
 
-  private val zkClient = EasyMock.createNiceMock(classOf[KafkaZkClient])
+  private val zkClient: KafkaZkClient = EasyMock.createNiceMock(classOf[KafkaZkClient])
 
   @After
   def tearDown(): Unit = {
@@ -32,26 +32,24 @@ class ProducerIdManagerTest {
   }
 
   @Test
-  def testGetProducerId() {
+  def testGetProducerId(): Unit = {
     var zkVersion: Option[Int] = None
-    var data: String = null
-    EasyMock.expect(zkClient.getDataAndVersion(EasyMock.anyString)).andAnswer(new IAnswer[(Option[String], Int)] {
-      override def answer(): (Option[String], Int) = zkVersion.map(Some(data) -> _).getOrElse(None, 0)
-    }).anyTimes()
+    var data: Array[Byte] = null
+    EasyMock.expect(zkClient.getDataAndVersion(EasyMock.anyString)).andAnswer(() =>
+      zkVersion.map(Some(data) -> _).getOrElse(None, 0)).anyTimes()
 
     val capturedVersion: Capture[Int] = EasyMock.newCapture()
-    val capturedData: Capture[String] = EasyMock.newCapture()
+    val capturedData: Capture[Array[Byte]] = EasyMock.newCapture()
     EasyMock.expect(zkClient.conditionalUpdatePath(EasyMock.anyString(),
       EasyMock.capture(capturedData),
       EasyMock.capture(capturedVersion),
-      EasyMock.anyObject[Option[(KafkaZkClient, String, String) => (Boolean, Int)]])).andAnswer(new IAnswer[(Boolean, Int)] {
-        override def answer(): (Boolean, Int) = {
-          val newZkVersion = capturedVersion.getValue + 1
-          zkVersion = Some(newZkVersion)
-          data = capturedData.getValue
-          (true, newZkVersion)
-        }
-      }).anyTimes()
+      EasyMock.anyObject[Option[(KafkaZkClient, String, Array[Byte]) => (Boolean, Int)]])
+    ).andAnswer(() => {
+      val newZkVersion = capturedVersion.getValue + 1
+      zkVersion = Some(newZkVersion)
+      data = capturedData.getValue
+      (true, newZkVersion)
+    }).anyTimes()
 
     EasyMock.replay(zkClient)
 
@@ -75,13 +73,11 @@ class ProducerIdManagerTest {
   }
 
   @Test(expected = classOf[KafkaException])
-  def testExceedProducerIdLimit() {
-    EasyMock.expect(zkClient.getDataAndVersion(EasyMock.anyString)).andAnswer(new IAnswer[(Option[String], Int)] {
-      override def answer(): (Option[String], Int) = {
-        val json = ProducerIdManager.generateProducerIdBlockJson(
-          ProducerIdBlock(0, Long.MaxValue - ProducerIdManager.PidBlockSize, Long.MaxValue))
-        (Some(json), 0)
-      }
+  def testExceedProducerIdLimit(): Unit = {
+    EasyMock.expect(zkClient.getDataAndVersion(EasyMock.anyString)).andAnswer(() => {
+      val json = ProducerIdManager.generateProducerIdBlockJson(
+        ProducerIdBlock(0, Long.MaxValue - ProducerIdManager.PidBlockSize, Long.MaxValue))
+      (Some(json), 0)
     }).anyTimes()
     EasyMock.replay(zkClient)
     new ProducerIdManager(0, zkClient)
diff --git a/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionCoordinatorConcurrencyTest.scala b/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionCoordinatorConcurrencyTest.scala
new file mode 100644
index 0000000000000..3788cb1d65325
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionCoordinatorConcurrencyTest.scala
@@ -0,0 +1,623 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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.coordinator.transaction
+
+import java.nio.ByteBuffer
+import java.util.concurrent.atomic.AtomicBoolean
+
+import kafka.coordinator.AbstractCoordinatorConcurrencyTest
+import kafka.coordinator.AbstractCoordinatorConcurrencyTest._
+import kafka.coordinator.transaction.TransactionCoordinatorConcurrencyTest._
+import kafka.log.Log
+import kafka.server.{FetchDataInfo, FetchLogEnd, KafkaConfig, LogOffsetMetadata, MetadataCache}
+import kafka.utils.{Pool, TestUtils}
+import org.apache.kafka.clients.{ClientResponse, NetworkClient}
+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, RecordBatch, SimpleRecord}
+import org.apache.kafka.common.requests._
+import org.apache.kafka.common.utils.{LogContext, MockTime, ProducerIdAndEpoch}
+import org.apache.kafka.common.{Node, TopicPartition}
+import org.easymock.{EasyMock, IAnswer}
+import org.junit.Assert._
+import org.junit.{After, Before, Test}
+
+import scala.jdk.CollectionConverters._
+import scala.collection.{Map, mutable}
+
+class TransactionCoordinatorConcurrencyTest extends AbstractCoordinatorConcurrencyTest[Transaction] {
+  private val nTransactions = nThreads * 10
+  private val coordinatorEpoch = 10
+  private val numPartitions = nThreads * 5
+
+  private val txnConfig = TransactionConfig()
+  private var transactionCoordinator: TransactionCoordinator = _
+  private var txnStateManager: TransactionStateManager = _
+  private var txnMarkerChannelManager: TransactionMarkerChannelManager = _
+
+  private val allOperations = Seq(
+      new InitProducerIdOperation,
+      new AddPartitionsToTxnOperation(Set(new TopicPartition("topic", 0))),
+      new EndTxnOperation)
+
+  private val allTransactions = mutable.Set[Transaction]()
+  private val txnRecordsByPartition: Map[Int, mutable.ArrayBuffer[SimpleRecord]] =
+    (0 until numPartitions).map { i => (i, mutable.ArrayBuffer[SimpleRecord]()) }.toMap
+
+  val producerId = 11
+  private var bumpProducerId = false
+
+  @Before
+  override def setUp(): Unit = {
+    super.setUp()
+
+    EasyMock.expect(zkClient.getTopicPartitionCount(TRANSACTION_STATE_TOPIC_NAME))
+      .andReturn(Some(numPartitions))
+      .anyTimes()
+    EasyMock.replay(zkClient)
+
+    txnStateManager = new TransactionStateManager(0, zkClient, scheduler, replicaManager, txnConfig, time,
+      new Metrics())
+    for (i <- 0 until numPartitions)
+      txnStateManager.addLoadedTransactionsToCache(i, coordinatorEpoch, new Pool[String, TransactionMetadata]())
+
+    val pidManager: ProducerIdManager = EasyMock.createNiceMock(classOf[ProducerIdManager])
+    EasyMock.expect(pidManager.generateProducerId())
+      .andAnswer(() => if (bumpProducerId) producerId + 1 else producerId)
+      .anyTimes()
+    val brokerNode = new Node(0, "host", 10)
+    val metadataCache: MetadataCache = EasyMock.createNiceMock(classOf[MetadataCache])
+    EasyMock.expect(metadataCache.getPartitionLeaderEndpoint(
+      EasyMock.anyString(),
+      EasyMock.anyInt(),
+      EasyMock.anyObject())
+    ).andReturn(Some(brokerNode)).anyTimes()
+    val networkClient: NetworkClient = EasyMock.createNiceMock(classOf[NetworkClient])
+    txnMarkerChannelManager = new TransactionMarkerChannelManager(
+      KafkaConfig.fromProps(serverProps),
+      metadataCache,
+      networkClient,
+      txnStateManager,
+      time)
+
+    transactionCoordinator = new TransactionCoordinator(brokerId = 0,
+      txnConfig,
+      scheduler,
+      pidManager,
+      txnStateManager,
+      txnMarkerChannelManager,
+      time,
+      new LogContext)
+    EasyMock.replay(pidManager)
+    EasyMock.replay(metadataCache)
+    EasyMock.replay(networkClient)
+  }
+
+  @After
+  override def tearDown(): Unit = {
+    try {
+      EasyMock.reset(zkClient, replicaManager)
+      transactionCoordinator.shutdown()
+    } finally {
+      super.tearDown()
+    }
+  }
+
+  @Test
+  def testConcurrentGoodPathWithConcurrentPartitionLoading(): Unit = {
+    // This is a somewhat contrived test case which reproduces the bug in KAFKA-9777.
+    // When a new partition needs to be loaded, we acquire the write lock in order to
+    // add the partition to the set of loading partitions. We should still be able to
+    // make progress with transactions even while this is ongoing.
+
+    val keepRunning = new AtomicBoolean(true)
+    val t = new Thread() {
+      override def run(): Unit = {
+        while (keepRunning.get()) {
+          txnStateManager.addLoadingPartition(numPartitions + 1, coordinatorEpoch)
+        }
+      }
+    }
+    t.start()
+
+    verifyConcurrentOperations(createTransactions, allOperations)
+    keepRunning.set(false)
+    t.join()
+  }
+
+  @Test
+  def testConcurrentGoodPathSequence(): Unit = {
+    verifyConcurrentOperations(createTransactions, allOperations)
+  }
+
+  @Test
+  def testConcurrentRandomSequences(): Unit = {
+    verifyConcurrentRandomSequences(createTransactions, allOperations)
+  }
+
+  /**
+    * Concurrently load one set of transaction state topic partitions and unload another
+    * set of partitions. This tests partition leader changes of transaction state topic
+    * that are handled by different threads concurrently. Verifies that the metadata of
+    * unloaded partitions are removed from the transaction manager and that the transactions
+    * from the newly loaded partitions are loaded correctly.
+    */
+  @Test
+  def testConcurrentLoadUnloadPartitions(): Unit = {
+    val partitionsToLoad = (0 until numPartitions / 2).toSet
+    val partitionsToUnload = (numPartitions / 2 until numPartitions).toSet
+    verifyConcurrentActions(loadUnloadActions(partitionsToLoad, partitionsToUnload))
+  }
+
+  /**
+    * Concurrently load one set of transaction state topic partitions, unload a second set
+    * of partitions and expire transactions on a third set of partitions. This tests partition
+    * leader changes of transaction state topic that are handled by different threads concurrently
+    * while expiry is performed on another thread. Verifies the state of transactions on all the partitions.
+    */
+  @Test
+  def testConcurrentTransactionExpiration(): Unit = {
+    val partitionsToLoad = (0 until numPartitions / 3).toSet
+    val partitionsToUnload = (numPartitions / 3 until numPartitions * 2 / 3).toSet
+    val partitionsWithExpiringTxn = (numPartitions * 2 / 3 until numPartitions).toSet
+    val expiringTransactions = allTransactions.filter { txn =>
+      partitionsWithExpiringTxn.contains(txnStateManager.partitionFor(txn.transactionalId))
+    }.toSet
+    val expireAction = new ExpireTransactionsAction(expiringTransactions)
+    verifyConcurrentActions(loadUnloadActions(partitionsToLoad, partitionsToUnload) + expireAction)
+  }
+
+  @Test
+  def testConcurrentNewInitProducerIdRequests(): Unit = {
+    val transactions = (1 to 100).flatMap(i => createTransactions(s"testConcurrentInitProducerID$i-"))
+    bumpProducerId = true
+    transactions.foreach { txn =>
+      val txnMetadata = prepareExhaustedEpochTxnMetadata(txn)
+      txnStateManager.putTransactionStateIfNotExists(txnMetadata)
+
+      // Test simultaneous requests from an existing producer trying to bump the epoch and a new producer initializing
+      val newProducerOp1 = new InitProducerIdOperation()
+      val newProducerOp2 = new InitProducerIdOperation()
+      verifyConcurrentActions(Set(newProducerOp1, newProducerOp2).map(_.actionNoVerify(txn)))
+
+      // If only one request succeeds, assert that the epoch was successfully increased
+      // If both requests succeed, the new producer must have run after the existing one and should have the higher epoch
+      (newProducerOp1.result.get.error, newProducerOp2.result.get.error) match {
+        case (Errors.NONE, Errors.NONE) =>
+          assertNotEquals(newProducerOp1.result.get.producerEpoch, newProducerOp2.result.get.producerEpoch)
+          // assertEquals(0, newProducerOp1.result.get.producerEpoch)
+          // assertEquals(0, newProducerOp2.result.get.producerEpoch)
+        case (Errors.NONE, _) =>
+          assertEquals(0, newProducerOp1.result.get.producerEpoch)
+        case (_, Errors.NONE) =>
+          assertEquals(0, newProducerOp2.result.get.producerEpoch)
+        case (_, _) => fail("One of two InitProducerId requests should succeed")
+      }
+    }
+  }
+
+  @Test
+  def testConcurrentInitProducerIdRequestsOneNewOneContinuing(): Unit = {
+    val transactions = (1 to 10).flatMap(i => createTransactions(s"testConcurrentInitProducerID$i-"))
+    transactions.foreach { txn =>
+      val firstInitReq = new InitProducerIdOperation()
+      firstInitReq.run(txn)
+      firstInitReq.awaitAndVerify(txn)
+
+      // Test simultaneous requests from an existing producer trying to bump the epoch and a new producer initializing
+      val producerIdAndEpoch = new ProducerIdAndEpoch(firstInitReq.result.get.producerId, firstInitReq.result.get.producerEpoch)
+      val bumpEpochOp = new InitProducerIdOperation(Some(producerIdAndEpoch))
+      val newProducerOp = new InitProducerIdOperation()
+      verifyConcurrentActions(Set(bumpEpochOp, newProducerOp).map(_.actionNoVerify(txn)))
+
+      // If only one request succeeds, assert that the epoch was successfully increased
+      // If both requests succeed, the new producer must have run after the existing one and should have the higher epoch
+      (bumpEpochOp.result.get.error, newProducerOp.result.get.error) match {
+        case (Errors.NONE, Errors.NONE) =>
+          assertEquals(producerIdAndEpoch.epoch + 2, newProducerOp.result.get.producerEpoch)
+          assertEquals(producerIdAndEpoch.epoch + 1, bumpEpochOp.result.get.producerEpoch)
+        case (Errors.NONE, _) =>
+          assertEquals(producerIdAndEpoch.epoch + 1, bumpEpochOp.result.get.producerEpoch)
+        case (_, Errors.NONE) =>
+          assertEquals(producerIdAndEpoch.epoch + 1, newProducerOp.result.get.producerEpoch)
+        case (_, _) => fail("One of two InitProducerId requests should succeed")
+      }
+    }
+  }
+
+  @Test
+  def testConcurrentContinuingInitProducerIdRequests(): Unit = {
+    val transactions = (1 to 100).flatMap(i => createTransactions(s"testConcurrentInitProducerID$i-"))
+    transactions.foreach { txn =>
+      // Test simultaneous requests from an existing producers trying to re-initialize when no state is present
+      val producerIdAndEpoch = new ProducerIdAndEpoch(producerId, 10)
+      val bumpEpochOp1 = new InitProducerIdOperation(Some(producerIdAndEpoch))
+      val bumpEpochOp2 = new InitProducerIdOperation(Some(producerIdAndEpoch))
+      verifyConcurrentActions(Set(bumpEpochOp1, bumpEpochOp2).map(_.actionNoVerify(txn)))
+
+      // If only one request succeeds, assert that the epoch was successfully increased
+      // If both requests succeed, the new producer must have run after the existing one and should have the higher epoch
+      (bumpEpochOp1.result.get.error, bumpEpochOp2.result.get.error) match {
+        case (Errors.NONE, Errors.NONE) =>
+          fail("One of two InitProducerId requests should fail due to concurrent requests or non-matching epochs")
+        case (Errors.NONE, _) =>
+          assertEquals(0, bumpEpochOp1.result.get.producerEpoch)
+        case (_, Errors.NONE) =>
+          assertEquals(0, bumpEpochOp2.result.get.producerEpoch)
+        case (_, _) => fail("One of two InitProducerId requests should succeed")
+      }
+    }
+  }
+
+  @Test
+  def testConcurrentInitProducerIdRequestsWithRetry(): Unit = {
+    val transactions = (1 to 10).flatMap(i => createTransactions(s"testConcurrentInitProducerID$i-"))
+    transactions.foreach { txn =>
+      val firstInitReq = new InitProducerIdOperation()
+      firstInitReq.run(txn)
+      firstInitReq.awaitAndVerify(txn)
+
+      val initialProducerIdAndEpoch = new ProducerIdAndEpoch(firstInitReq.result.get.producerId, firstInitReq.result.get.producerEpoch)
+      val bumpEpochReq = new InitProducerIdOperation(Some(initialProducerIdAndEpoch))
+      bumpEpochReq.run(txn)
+      bumpEpochReq.awaitAndVerify(txn)
+
+      // Test simultaneous requests from an existing producer retrying the epoch bump and a new producer initializing
+      val bumpedProducerIdAndEpoch = new ProducerIdAndEpoch(bumpEpochReq.result.get.producerId, bumpEpochReq.result.get.producerEpoch)
+      val retryBumpEpochOp = new InitProducerIdOperation(Some(initialProducerIdAndEpoch))
+      val newProducerOp = new InitProducerIdOperation()
+      verifyConcurrentActions(Set(retryBumpEpochOp, newProducerOp).map(_.actionNoVerify(txn)))
+
+      // If both requests succeed, the new producer must have run after the existing one and should have the higher epoch
+      // If the retry succeeds and the new producer doesn't, assert that the already-bumped epoch was returned
+      // If the new producer succeeds and the retry doesn't, assert the epoch was bumped
+      (retryBumpEpochOp.result.get.error, newProducerOp.result.get.error) match {
+        case (Errors.NONE, Errors.NONE) =>
+          assertEquals(bumpedProducerIdAndEpoch.epoch + 1, newProducerOp.result.get.producerEpoch)
+          assertEquals(bumpedProducerIdAndEpoch.epoch, retryBumpEpochOp.result.get.producerEpoch)
+        case (Errors.NONE, _) =>
+          assertEquals(bumpedProducerIdAndEpoch.epoch, retryBumpEpochOp.result.get.producerEpoch)
+        case (_, Errors.NONE) =>
+          assertEquals(bumpedProducerIdAndEpoch.epoch + 1, newProducerOp.result.get.producerEpoch)
+        case (_, _) => fail("At least one InitProducerId request should succeed")
+      }
+    }
+  }
+
+  @Test
+  def testConcurrentInitProducerRequestsAtPidBoundary(): Unit = {
+    val transactions = (1 to 10).flatMap(i => createTransactions(s"testConcurrentInitProducerID$i-"))
+    bumpProducerId = true
+    transactions.foreach { txn =>
+      val txnMetadata = prepareExhaustedEpochTxnMetadata(txn)
+      txnStateManager.putTransactionStateIfNotExists(txnMetadata)
+
+      // Test simultaneous requests from an existing producer attempting to bump the epoch and a new producer initializing
+      val bumpEpochOp = new InitProducerIdOperation(Some(new ProducerIdAndEpoch(producerId, (Short.MaxValue - 1).toShort)))
+      val newProducerOp = new InitProducerIdOperation()
+      verifyConcurrentActions(Set(bumpEpochOp, newProducerOp).map(_.actionNoVerify(txn)))
+
+      // If the retry succeeds and the new producer doesn't, assert that the already-bumped epoch was returned
+      // If the new producer succeeds and the retry doesn't, assert the epoch was bumped
+      // If both requests succeed, the new producer must have run after the existing one and should have the higher epoch
+      (bumpEpochOp.result.get.error, newProducerOp.result.get.error) match {
+        case (Errors.NONE, Errors.NONE) =>
+          assertEquals(0, bumpEpochOp.result.get.producerEpoch)
+          assertEquals(producerId + 1, bumpEpochOp.result.get.producerId)
+
+          assertEquals(1, newProducerOp.result.get.producerEpoch)
+          assertEquals(producerId + 1, newProducerOp.result.get.producerId)
+        case (Errors.NONE, _) =>
+          assertEquals(0, bumpEpochOp.result.get.producerEpoch)
+          assertEquals(producerId + 1, bumpEpochOp.result.get.producerId)
+        case (_, Errors.NONE) =>
+          assertEquals(0, newProducerOp.result.get.producerEpoch)
+          assertEquals(producerId + 1, newProducerOp.result.get.producerId)
+        case (_, _) => fail("One of two InitProducerId requests should succeed")
+      }
+    }
+
+    bumpProducerId = false
+  }
+
+  @Test
+  def testConcurrentInitProducerRequestsWithRetryAtPidBoundary(): Unit = {
+    val transactions = (1 to 10).flatMap(i => createTransactions(s"testConcurrentInitProducerID$i-"))
+    bumpProducerId = true
+    transactions.foreach { txn =>
+      val txnMetadata = prepareExhaustedEpochTxnMetadata(txn)
+      txnStateManager.putTransactionStateIfNotExists(txnMetadata)
+
+      val bumpEpochReq = new InitProducerIdOperation(Some(new ProducerIdAndEpoch(producerId, (Short.MaxValue - 1).toShort)))
+      bumpEpochReq.run(txn)
+      bumpEpochReq.awaitAndVerify(txn)
+
+      // Test simultaneous requests from an existing producer attempting to bump the epoch and a new producer initializing
+      val retryBumpEpochOp = new InitProducerIdOperation(Some(new ProducerIdAndEpoch(producerId, (Short.MaxValue - 1).toShort)))
+      val newProducerOp = new InitProducerIdOperation()
+      verifyConcurrentActions(Set(retryBumpEpochOp, newProducerOp).map(_.actionNoVerify(txn)))
+
+      // If the retry succeeds and the new producer doesn't, assert that the already-bumped epoch was returned
+      // If the new producer succeeds and the retry doesn't, assert the epoch was bumped
+      // If both requests succeed, the new producer must have run after the existing one and should have the higher epoch
+      (retryBumpEpochOp.result.get.error, newProducerOp.result.get.error) match {
+        case (Errors.NONE, Errors.NONE) =>
+          assertEquals(0, retryBumpEpochOp.result.get.producerEpoch)
+          assertEquals(producerId + 1, retryBumpEpochOp.result.get.producerId)
+
+          assertEquals(1, newProducerOp.result.get.producerEpoch)
+          assertEquals(producerId + 1, newProducerOp.result.get.producerId)
+        case (Errors.NONE, _) =>
+          assertEquals(0, retryBumpEpochOp.result.get.producerEpoch)
+          assertEquals(producerId + 1, retryBumpEpochOp.result.get.producerId)
+        case (_, Errors.NONE) =>
+          assertEquals(1, newProducerOp.result.get.producerEpoch)
+          assertEquals(producerId + 1, newProducerOp.result.get.producerId)
+        case (_, _) => fail("One of two InitProducerId requests should succeed")
+      }
+    }
+
+    bumpProducerId = false
+  }
+
+  override def enableCompletion(): Unit = {
+    super.enableCompletion()
+
+    def createResponse(request: WriteTxnMarkersRequest): WriteTxnMarkersResponse  = {
+      val pidErrorMap = request.markers.asScala.map { marker =>
+        (marker.producerId.asInstanceOf[java.lang.Long], marker.partitions.asScala.map { tp => (tp, Errors.NONE) }.toMap.asJava)
+      }.toMap.asJava
+      new WriteTxnMarkersResponse(pidErrorMap)
+    }
+    synchronized {
+      txnMarkerChannelManager.generateRequests().foreach { requestAndHandler =>
+        val request = requestAndHandler.request.asInstanceOf[WriteTxnMarkersRequest.Builder].build()
+        val response = createResponse(request)
+        requestAndHandler.handler.onComplete(new ClientResponse(new RequestHeader(ApiKeys.PRODUCE, 0, "client", 1),
+          null, null, 0, 0, false, null, null, response))
+      }
+    }
+  }
+
+  /**
+    * Concurrently load `partitionsToLoad` and unload `partitionsToUnload`. Before the concurrent operations
+    * are run `partitionsToLoad` must be unloaded first since all partitions were loaded during setUp.
+    */
+  private def loadUnloadActions(partitionsToLoad: Set[Int], partitionsToUnload: Set[Int]): Set[Action] = {
+    val transactions = (1 to 10).flatMap(i => createTransactions(s"testConcurrentLoadUnloadPartitions$i-")).toSet
+    transactions.foreach(txn => prepareTransaction(txn))
+    val unload = partitionsToLoad.map(new UnloadTxnPartitionAction(_))
+    unload.foreach(_.run())
+    unload.foreach(_.await())
+    partitionsToLoad.map(new LoadTxnPartitionAction(_)) ++ partitionsToUnload.map(new UnloadTxnPartitionAction(_))
+  }
+
+  private def createTransactions(txnPrefix: String): Set[Transaction] = {
+    val transactions = (0 until nTransactions).map { i => new Transaction(s"$txnPrefix$i", i, time) }
+    allTransactions ++= transactions
+    transactions.toSet
+  }
+
+  private def verifyTransaction(txn: Transaction, expectedState: TransactionState): Unit = {
+    val (metadata, success) = TestUtils.computeUntilTrue({
+      enableCompletion()
+      transactionMetadata(txn)
+    })(metadata => metadata.nonEmpty && metadata.forall(m => m.state == expectedState && m.pendingState.isEmpty))
+    assertTrue(s"Invalid metadata state $metadata", success)
+  }
+
+  private def transactionMetadata(txn: Transaction): Option[TransactionMetadata] = {
+    txnStateManager.getTransactionState(txn.transactionalId) match {
+      case Left(error) =>
+        if (error == Errors.NOT_COORDINATOR)
+          None
+        else
+          throw new AssertionError(s"Unexpected transaction error $error for $txn")
+      case Right(Some(metadata)) =>
+        Some(metadata.transactionMetadata)
+      case Right(None) =>
+        None
+    }
+  }
+
+  private def prepareTransaction(txn: Transaction): Unit = {
+    val partitionId = txnStateManager.partitionFor(txn.transactionalId)
+    val txnRecords = txnRecordsByPartition(partitionId)
+    val initPidOp = new InitProducerIdOperation()
+    val addPartitionsOp = new AddPartitionsToTxnOperation(Set(new TopicPartition("topic", 0)))
+    initPidOp.run(txn)
+    initPidOp.awaitAndVerify(txn)
+    addPartitionsOp.run(txn)
+    addPartitionsOp.awaitAndVerify(txn)
+
+    val txnMetadata = transactionMetadata(txn).getOrElse(throw new IllegalStateException(s"Transaction not found $txn"))
+    txnRecords += new SimpleRecord(txn.txnMessageKeyBytes, TransactionLog.valueToBytes(txnMetadata.prepareNoTransit()))
+
+    txnMetadata.state = PrepareCommit
+    txnRecords += new SimpleRecord(txn.txnMessageKeyBytes, TransactionLog.valueToBytes(txnMetadata.prepareNoTransit()))
+
+    prepareTxnLog(partitionId)
+  }
+
+  private def prepareTxnLog(partitionId: Int): Unit = {
+
+    val logMock: Log =  EasyMock.mock(classOf[Log])
+    val fileRecordsMock: FileRecords = EasyMock.mock(classOf[FileRecords])
+
+    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).toArray: _*)
+    val endOffset = startOffset + records.records.asScala.size
+
+    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(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 = {
+        val buffer = bufferCapture.getValue
+        buffer.put(records.buffer.duplicate)
+        buffer.flip()
+      }
+    })
+
+    EasyMock.replay(logMock, fileRecordsMock)
+    synchronized {
+      replicaManager.updateLog(topicPartition, logMock, endOffset)
+    }
+  }
+
+  private def prepareExhaustedEpochTxnMetadata(txn: Transaction): TransactionMetadata = {
+    new TransactionMetadata(transactionalId = txn.transactionalId,
+      producerId = producerId,
+      lastProducerId = RecordBatch.NO_PRODUCER_ID,
+      producerEpoch = (Short.MaxValue - 1).toShort,
+      lastProducerEpoch = RecordBatch.NO_PRODUCER_EPOCH,
+      txnTimeoutMs = 60000,
+      state = Empty,
+      topicPartitions = collection.mutable.Set.empty[TopicPartition],
+      txnLastUpdateTimestamp = time.milliseconds())
+  }
+
+  abstract class TxnOperation[R] extends Operation {
+    @volatile var result: Option[R] = None
+    def resultCallback(r: R): Unit = this.result = Some(r)
+  }
+
+  class InitProducerIdOperation(val producerIdAndEpoch: Option[ProducerIdAndEpoch] = None) extends TxnOperation[InitProducerIdResult] {
+    override def run(txn: Transaction): Unit = {
+      transactionCoordinator.handleInitProducerId(txn.transactionalId, 60000, producerIdAndEpoch, resultCallback)
+      replicaManager.tryCompleteActions()
+    }
+    override def awaitAndVerify(txn: Transaction): Unit = {
+      val initPidResult = result.getOrElse(throw new IllegalStateException("InitProducerId has not completed"))
+      assertEquals(Errors.NONE, initPidResult.error)
+      verifyTransaction(txn, Empty)
+    }
+  }
+
+  class AddPartitionsToTxnOperation(partitions: Set[TopicPartition]) extends TxnOperation[Errors] {
+    override def run(txn: Transaction): Unit = {
+      transactionMetadata(txn).foreach { txnMetadata =>
+        transactionCoordinator.handleAddPartitionsToTransaction(txn.transactionalId,
+            txnMetadata.producerId,
+            txnMetadata.producerEpoch,
+            partitions,
+            resultCallback)
+        replicaManager.tryCompleteActions()
+      }
+    }
+    override def awaitAndVerify(txn: Transaction): Unit = {
+      val error = result.getOrElse(throw new IllegalStateException("AddPartitionsToTransaction has not completed"))
+      assertEquals(Errors.NONE, error)
+      verifyTransaction(txn, Ongoing)
+    }
+  }
+
+  class EndTxnOperation extends TxnOperation[Errors] {
+    override def run(txn: Transaction): Unit = {
+      transactionMetadata(txn).foreach { txnMetadata =>
+        transactionCoordinator.handleEndTransaction(txn.transactionalId,
+          txnMetadata.producerId,
+          txnMetadata.producerEpoch,
+          transactionResult(txn),
+          resultCallback)
+      }
+    }
+    override def awaitAndVerify(txn: Transaction): Unit = {
+      val error = result.getOrElse(throw new IllegalStateException("EndTransaction has not completed"))
+      if (!txn.ended) {
+        txn.ended = true
+        assertEquals(Errors.NONE, error)
+        val expectedState = if (transactionResult(txn) == TransactionResult.COMMIT) CompleteCommit else CompleteAbort
+        verifyTransaction(txn, expectedState)
+      } else
+        assertEquals(Errors.INVALID_TXN_STATE, error)
+    }
+    // Test both commit and abort. Transactional ids used in the test have the format 
+    // Use the last digit of the index to decide between commit and abort.
+    private def transactionResult(txn: Transaction): TransactionResult = {
+      val txnId = txn.transactionalId
+      val lastDigit = txnId(txnId.length - 1).toInt
+      if (lastDigit % 2 == 0) TransactionResult.COMMIT else TransactionResult.ABORT
+    }
+  }
+
+  class LoadTxnPartitionAction(txnTopicPartitionId: Int) extends Action {
+    override def run(): Unit = {
+      transactionCoordinator.onElection(txnTopicPartitionId, coordinatorEpoch)
+    }
+    override def await(): Unit = {
+      allTransactions.foreach { txn =>
+        if (txnStateManager.partitionFor(txn.transactionalId) == txnTopicPartitionId) {
+          verifyTransaction(txn, CompleteCommit)
+        }
+      }
+    }
+  }
+
+  class UnloadTxnPartitionAction(txnTopicPartitionId: Int) extends Action {
+    val txnRecords: mutable.ArrayBuffer[SimpleRecord] = mutable.ArrayBuffer[SimpleRecord]()
+    override def run(): Unit = {
+      transactionCoordinator.onResignation(txnTopicPartitionId, Some(coordinatorEpoch))
+    }
+    override def await(): Unit = {
+      allTransactions.foreach { txn =>
+        if (txnStateManager.partitionFor(txn.transactionalId) == txnTopicPartitionId)
+          assertTrue("Transaction metadata not removed", transactionMetadata(txn).isEmpty)
+      }
+    }
+  }
+
+  class ExpireTransactionsAction(transactions: Set[Transaction]) extends Action {
+    override def run(): Unit = {
+      transactions.foreach { txn =>
+        transactionMetadata(txn).foreach { txnMetadata =>
+          txnMetadata.txnLastUpdateTimestamp = time.milliseconds() - txnConfig.transactionalIdExpirationMs
+        }
+      }
+      txnStateManager.enableTransactionalIdExpiration()
+      replicaManager.tryCompleteActions()
+      time.sleep(txnConfig.removeExpiredTransactionalIdsIntervalMs + 1)
+    }
+
+    override def await(): Unit = {
+      val (_, success) = TestUtils.computeUntilTrue({
+        replicaManager.tryCompleteActions()
+        transactions.forall(txn => transactionMetadata(txn).isEmpty)
+      })(identity)
+      assertTrue("Transaction not expired", success)
+    }
+  }
+}
+
+object TransactionCoordinatorConcurrencyTest {
+
+  class Transaction(val transactionalId: String, producerId: Long, time: MockTime) extends CoordinatorMember {
+    val txnMessageKeyBytes: Array[Byte] = TransactionLog.keyToBytes(transactionalId)
+    @volatile var ended = false
+    override def toString: String = transactionalId
+  }
+}
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..64e18e4b7ebd9 100644
--- a/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionCoordinatorTest.scala
+++ b/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionCoordinatorTest.scala
@@ -19,9 +19,10 @@ package kafka.coordinator.transaction
 import kafka.utils.MockScheduler
 import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.protocol.Errors
+import org.apache.kafka.common.record.RecordBatch
 import org.apache.kafka.common.requests.TransactionResult
-import org.apache.kafka.common.utils.{LogContext, MockTime}
-import org.easymock.{Capture, EasyMock, IAnswer}
+import org.apache.kafka.common.utils.{LogContext, MockTime, ProducerIdAndEpoch}
+import org.easymock.{Capture, EasyMock}
 import org.junit.Assert._
 import org.junit.Test
 
@@ -37,18 +38,19 @@ class TransactionCoordinatorTest {
   val transactionMarkerChannelManager: TransactionMarkerChannelManager = EasyMock.createNiceMock(classOf[TransactionMarkerChannelManager])
   val capturedTxn: Capture[TransactionMetadata] = EasyMock.newCapture()
   val capturedErrorsCallback: Capture[Errors => Unit] = EasyMock.newCapture()
+  val capturedTxnTransitMetadata: Capture[TxnTransitMetadata] = EasyMock.newCapture()
   val brokerId = 0
   val coordinatorEpoch = 0
   private val transactionalId = "known"
   private val producerId = 10
-  private val producerEpoch:Short = 1
+  private val producerEpoch: Short = 1
   private val txnTimeoutMs = 1
 
   private val partitions = mutable.Set[TopicPartition](new TopicPartition("topic1", 0))
   private val scheduler = new MockScheduler(time)
 
   val coordinator = new TransactionCoordinator(brokerId,
-    new TransactionConfig(),
+    TransactionConfig(),
     scheduler,
     pidManager,
     transactionManager,
@@ -60,14 +62,10 @@ class TransactionCoordinatorTest {
   var error: Errors = Errors.NONE
 
   private def mockPidManager(): Unit = {
-    EasyMock.expect(pidManager.generateProducerId())
-      .andAnswer(new IAnswer[Long] {
-        override def answer(): Long = {
-          nextPid += 1
-          nextPid - 1
-        }
-      })
-      .anyTimes()
+    EasyMock.expect(pidManager.generateProducerId()).andAnswer(() => {
+      nextPid += 1
+      nextPid - 1
+    }).anyTimes()
   }
 
   private def initPidGenericMocks(transactionalId: String): Unit = {
@@ -82,9 +80,9 @@ class TransactionCoordinatorTest {
     mockPidManager()
     EasyMock.replay(pidManager)
 
-    coordinator.handleInitProducerId("", txnTimeoutMs, initProducerIdMockCallback)
+    coordinator.handleInitProducerId("", txnTimeoutMs, None, initProducerIdMockCallback)
     assertEquals(InitProducerIdResult(-1L, -1, Errors.INVALID_REQUEST), result)
-    coordinator.handleInitProducerId("", txnTimeoutMs, initProducerIdMockCallback)
+    coordinator.handleInitProducerId("", txnTimeoutMs, None, initProducerIdMockCallback)
     assertEquals(InitProducerIdResult(-1L, -1, Errors.INVALID_REQUEST), result)
   }
 
@@ -93,9 +91,9 @@ class TransactionCoordinatorTest {
     mockPidManager()
     EasyMock.replay(pidManager)
 
-    coordinator.handleInitProducerId(null, txnTimeoutMs, initProducerIdMockCallback)
+    coordinator.handleInitProducerId(null, txnTimeoutMs, None, initProducerIdMockCallback)
     assertEquals(InitProducerIdResult(0L, 0, Errors.NONE), result)
-    coordinator.handleInitProducerId(null, txnTimeoutMs, initProducerIdMockCallback)
+    coordinator.handleInitProducerId(null, txnTimeoutMs, None, initProducerIdMockCallback)
     assertEquals(InitProducerIdResult(1L, 0, Errors.NONE), result)
   }
 
@@ -107,30 +105,50 @@ class TransactionCoordinatorTest {
       .andReturn(Right(None))
       .once()
 
-    EasyMock.expect(transactionManager.putTransactionStateIfNotExists(EasyMock.eq(transactionalId), EasyMock.capture(capturedTxn)))
-      .andAnswer(new IAnswer[Either[Errors, CoordinatorEpochAndTxnMetadata]] {
-        override def answer(): Either[Errors, CoordinatorEpochAndTxnMetadata] = {
-          assertTrue(capturedTxn.hasCaptured)
-          Right(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, capturedTxn.getValue))
-        }
-      })
+    EasyMock.expect(transactionManager.putTransactionStateIfNotExists(EasyMock.capture(capturedTxn)))
+      .andAnswer(() => {
+        assertTrue(capturedTxn.hasCaptured)
+        Right(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, capturedTxn.getValue))
+      }).once()
+
+    EasyMock.expect(transactionManager.appendTransactionToLog(
+      EasyMock.eq(transactionalId),
+      EasyMock.eq(coordinatorEpoch),
+      EasyMock.anyObject().asInstanceOf[TxnTransitMetadata],
+      EasyMock.capture(capturedErrorsCallback),
+      EasyMock.anyObject())
+    ).andAnswer(() => capturedErrorsCallback.getValue.apply(Errors.NONE)).anyTimes()
+    EasyMock.replay(pidManager, transactionManager)
+
+    coordinator.handleInitProducerId(transactionalId, txnTimeoutMs, None, initProducerIdMockCallback)
+    assertEquals(InitProducerIdResult(nextPid - 1, 0, Errors.NONE), result)
+  }
+
+  @Test
+  def shouldGenerateNewProducerIdIfNoStateAndProducerIdAndEpochProvided(): Unit = {
+    initPidGenericMocks(transactionalId)
+
+    EasyMock.expect(transactionManager.getTransactionState(EasyMock.eq(transactionalId)))
+      .andReturn(Right(None))
       .once()
 
+    EasyMock.expect(transactionManager.putTransactionStateIfNotExists(EasyMock.capture(capturedTxn)))
+      .andAnswer(() => {
+        assertTrue(capturedTxn.hasCaptured)
+        Right(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, capturedTxn.getValue))
+      }).once()
+
     EasyMock.expect(transactionManager.appendTransactionToLog(
       EasyMock.eq(transactionalId),
       EasyMock.eq(coordinatorEpoch),
       EasyMock.anyObject().asInstanceOf[TxnTransitMetadata],
       EasyMock.capture(capturedErrorsCallback),
-      EasyMock.anyObject()))
-      .andAnswer(new IAnswer[Unit] {
-        override def answer(): Unit = {
-          capturedErrorsCallback.getValue.apply(Errors.NONE)
-        }
-      })
-      .anyTimes()
+      EasyMock.anyObject())
+    ).andAnswer(() => capturedErrorsCallback.getValue.apply(Errors.NONE)).anyTimes()
     EasyMock.replay(pidManager, transactionManager)
 
-    coordinator.handleInitProducerId(transactionalId, txnTimeoutMs, initProducerIdMockCallback)
+    coordinator.handleInitProducerId(transactionalId, txnTimeoutMs, Some(new ProducerIdAndEpoch(producerId, producerEpoch)),
+      initProducerIdMockCallback)
     assertEquals(InitProducerIdResult(nextPid - 1, 0, Errors.NONE), result)
   }
 
@@ -138,8 +156,8 @@ class TransactionCoordinatorTest {
   def shouldGenerateNewProducerIdIfEpochsExhausted(): Unit = {
     initPidGenericMocks(transactionalId)
 
-    val txnMetadata = new TransactionMetadata(transactionalId, producerId, (Short.MaxValue - 1).toShort,
-      txnTimeoutMs, Empty, mutable.Set.empty, time.milliseconds(), time.milliseconds())
+    val txnMetadata = new TransactionMetadata(transactionalId, producerId, producerId, (Short.MaxValue - 1).toShort,
+      (Short.MaxValue - 2).toShort, txnTimeoutMs, Empty, mutable.Set.empty, time.milliseconds(), time.milliseconds())
 
     EasyMock.expect(transactionManager.getTransactionState(EasyMock.eq(transactionalId)))
       .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, txnMetadata))))
@@ -150,15 +168,11 @@ class TransactionCoordinatorTest {
       EasyMock.anyObject().asInstanceOf[TxnTransitMetadata],
       EasyMock.capture(capturedErrorsCallback),
       EasyMock.anyObject()
-    )).andAnswer(new IAnswer[Unit] {
-      override def answer(): Unit = {
-        capturedErrorsCallback.getValue.apply(Errors.NONE)
-      }
-    })
+    )).andAnswer(() => capturedErrorsCallback.getValue.apply(Errors.NONE))
 
     EasyMock.replay(pidManager, transactionManager)
 
-    coordinator.handleInitProducerId(transactionalId, txnTimeoutMs, initProducerIdMockCallback)
+    coordinator.handleInitProducerId(transactionalId, txnTimeoutMs, None, initProducerIdMockCallback)
     assertNotEquals(producerId, result.producerId)
     assertEquals(0, result.producerEpoch)
     assertEquals(Errors.NONE, result.error)
@@ -173,7 +187,7 @@ class TransactionCoordinatorTest {
       .andReturn(Left(Errors.NOT_COORDINATOR))
     EasyMock.replay(transactionManager)
 
-    coordinator.handleInitProducerId(transactionalId, txnTimeoutMs, initProducerIdMockCallback)
+    coordinator.handleInitProducerId(transactionalId, txnTimeoutMs, None, initProducerIdMockCallback)
     assertEquals(InitProducerIdResult(-1, -1, Errors.NOT_COORDINATOR), result)
   }
 
@@ -186,7 +200,7 @@ class TransactionCoordinatorTest {
       .andReturn(Left(Errors.COORDINATOR_LOAD_IN_PROGRESS))
     EasyMock.replay(transactionManager)
 
-    coordinator.handleInitProducerId(transactionalId, txnTimeoutMs, initProducerIdMockCallback)
+    coordinator.handleInitProducerId(transactionalId, txnTimeoutMs, None, initProducerIdMockCallback)
     assertEquals(InitProducerIdResult(-1, -1, Errors.COORDINATOR_LOAD_IN_PROGRESS), result)
   }
 
@@ -245,7 +259,8 @@ class TransactionCoordinatorTest {
 
   def validateConcurrentTransactions(state: TransactionState): Unit = {
     EasyMock.expect(transactionManager.getTransactionState(EasyMock.eq(transactionalId)))
-      .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, new TransactionMetadata(transactionalId, 0, 0, 0, state, mutable.Set.empty, 0, 0)))))
+      .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch,
+        new TransactionMetadata(transactionalId, 0, 0, 0, RecordBatch.NO_PRODUCER_EPOCH, 0, state, mutable.Set.empty, 0, 0)))))
 
     EasyMock.replay(transactionManager)
 
@@ -254,14 +269,15 @@ class TransactionCoordinatorTest {
   }
 
   @Test
-  def shouldRespondWithInvalidTnxProduceEpochOnAddPartitionsWhenEpochsAreDifferent(): Unit = {
+  def shouldRespondWithProducerFencedOnAddPartitionsWhenEpochsAreDifferent(): Unit = {
     EasyMock.expect(transactionManager.getTransactionState(EasyMock.eq(transactionalId)))
-      .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, new TransactionMetadata(transactionalId, 0, 10, 0, PrepareCommit, mutable.Set.empty, 0, 0)))))
+      .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch,
+        new TransactionMetadata(transactionalId, 0, 0, 10, 9, 0, PrepareCommit, mutable.Set.empty, 0, 0)))))
 
     EasyMock.replay(transactionManager)
 
     coordinator.handleAddPartitionsToTransaction(transactionalId, 0L, 0, partitions, errorsCallback)
-    assertEquals(Errors.INVALID_PRODUCER_EPOCH, error)
+    assertEquals(Errors.PRODUCER_FENCED, error)
   }
 
   @Test
@@ -285,8 +301,8 @@ class TransactionCoordinatorTest {
   }
 
   def validateSuccessfulAddPartitions(previousState: TransactionState): Unit = {
-    val txnMetadata = new TransactionMetadata(transactionalId, producerId, producerEpoch, txnTimeoutMs, previousState,
-      mutable.Set.empty, time.milliseconds(), time.milliseconds())
+    val txnMetadata = new TransactionMetadata(transactionalId, producerId, producerId, producerEpoch, (producerEpoch - 1).toShort,
+      txnTimeoutMs, previousState, mutable.Set.empty, time.milliseconds(), time.milliseconds())
 
     EasyMock.expect(transactionManager.getTransactionState(EasyMock.eq(transactionalId)))
       .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, txnMetadata))))
@@ -309,7 +325,8 @@ class TransactionCoordinatorTest {
   @Test
   def shouldRespondWithErrorsNoneOnAddPartitionWhenNoErrorsAndPartitionsTheSame(): Unit = {
     EasyMock.expect(transactionManager.getTransactionState(EasyMock.eq(transactionalId)))
-      .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, new TransactionMetadata(transactionalId, 0, 0, 0, Empty, partitions, 0, 0)))))
+      .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch,
+        new TransactionMetadata(transactionalId, 0, 0, 0, RecordBatch.NO_PRODUCER_EPOCH, 0, Empty, partitions, 0, 0)))))
 
     EasyMock.replay(transactionManager)
 
@@ -333,7 +350,8 @@ class TransactionCoordinatorTest {
   @Test
   def shouldReplyWithInvalidPidMappingOnEndTxnWhenPidDosentMatchMapped(): Unit = {
     EasyMock.expect(transactionManager.getTransactionState(EasyMock.eq(transactionalId)))
-      .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, new TransactionMetadata(transactionalId, 10, 0, 0, Ongoing, collection.mutable.Set.empty[TopicPartition], 0, time.milliseconds())))))
+      .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch,
+        new TransactionMetadata(transactionalId, 10, 10, 0, RecordBatch.NO_PRODUCER_EPOCH, 0, Ongoing, collection.mutable.Set.empty[TopicPartition], 0, time.milliseconds())))))
     EasyMock.replay(transactionManager)
 
     coordinator.handleEndTransaction(transactionalId, 0, 0, TransactionResult.COMMIT, errorsCallback)
@@ -344,18 +362,20 @@ class TransactionCoordinatorTest {
   @Test
   def shouldReplyWithProducerFencedOnEndTxnWhenEpochIsNotSameAsTransaction(): Unit = {
     EasyMock.expect(transactionManager.getTransactionState(EasyMock.eq(transactionalId)))
-      .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, new TransactionMetadata(transactionalId, producerId, 1, 1, Ongoing, collection.mutable.Set.empty[TopicPartition], 0, time.milliseconds())))))
+      .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch,
+        new TransactionMetadata(transactionalId, producerId, producerId, producerEpoch, (producerEpoch - 1).toShort, 1, Ongoing, collection.mutable.Set.empty[TopicPartition], 0, time.milliseconds())))))
     EasyMock.replay(transactionManager)
 
     coordinator.handleEndTransaction(transactionalId, producerId, 0, TransactionResult.COMMIT, errorsCallback)
-    assertEquals(Errors.INVALID_PRODUCER_EPOCH, error)
+    assertEquals(Errors.PRODUCER_FENCED, error)
     EasyMock.verify(transactionManager)
   }
 
   @Test
   def shouldReturnOkOnEndTxnWhenStatusIsCompleteCommitAndResultIsCommit(): Unit ={
     EasyMock.expect(transactionManager.getTransactionState(EasyMock.eq(transactionalId)))
-      .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, new TransactionMetadata(transactionalId, producerId, 1, 1, CompleteCommit, collection.mutable.Set.empty[TopicPartition], 0, time.milliseconds())))))
+      .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch,
+        new TransactionMetadata(transactionalId, producerId, producerId, producerEpoch, (producerEpoch - 1).toShort, 1, CompleteCommit, collection.mutable.Set.empty[TopicPartition], 0, time.milliseconds())))))
     EasyMock.replay(transactionManager)
 
     coordinator.handleEndTransaction(transactionalId, producerId, 1, TransactionResult.COMMIT, errorsCallback)
@@ -365,7 +385,7 @@ class TransactionCoordinatorTest {
 
   @Test
   def shouldReturnOkOnEndTxnWhenStatusIsCompleteAbortAndResultIsAbort(): Unit ={
-    val txnMetadata = new TransactionMetadata(transactionalId, producerId, 1, 1, CompleteAbort, collection.mutable.Set.empty[TopicPartition], 0, time.milliseconds())
+    val txnMetadata = new TransactionMetadata(transactionalId, producerId, producerId, producerEpoch, (producerEpoch - 1).toShort, 1, CompleteAbort, collection.mutable.Set.empty[TopicPartition], 0, time.milliseconds())
     EasyMock.expect(transactionManager.getTransactionState(EasyMock.eq(transactionalId)))
       .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, txnMetadata))))
     EasyMock.replay(transactionManager)
@@ -377,7 +397,7 @@ class TransactionCoordinatorTest {
 
   @Test
   def shouldReturnInvalidTxnRequestOnEndTxnRequestWhenStatusIsCompleteAbortAndResultIsNotAbort(): Unit = {
-    val txnMetadata = new TransactionMetadata(transactionalId, producerId, 1, 1, CompleteAbort, collection.mutable.Set.empty[TopicPartition], 0, time.milliseconds())
+    val txnMetadata = new TransactionMetadata(transactionalId, producerId, producerId, producerEpoch, (producerEpoch - 1).toShort, 1, CompleteAbort, collection.mutable.Set.empty[TopicPartition], 0, time.milliseconds())
     EasyMock.expect(transactionManager.getTransactionState(EasyMock.eq(transactionalId)))
       .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, txnMetadata))))
     EasyMock.replay(transactionManager)
@@ -389,7 +409,7 @@ class TransactionCoordinatorTest {
 
   @Test
   def shouldReturnInvalidTxnRequestOnEndTxnRequestWhenStatusIsCompleteCommitAndResultIsNotCommit(): Unit = {
-    val txnMetadata = new TransactionMetadata(transactionalId, producerId, 1, 1, CompleteCommit, collection.mutable.Set.empty[TopicPartition], 0, time.milliseconds())
+    val txnMetadata = new TransactionMetadata(transactionalId, producerId, producerId, producerEpoch, (producerEpoch - 1).toShort,1, CompleteCommit, collection.mutable.Set.empty[TopicPartition], 0, time.milliseconds())
     EasyMock.expect(transactionManager.getTransactionState(EasyMock.eq(transactionalId)))
       .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, txnMetadata))))
     EasyMock.replay(transactionManager)
@@ -402,7 +422,7 @@ class TransactionCoordinatorTest {
   @Test
   def shouldReturnConcurrentTxnRequestOnEndTxnRequestWhenStatusIsPrepareCommit(): Unit = {
     EasyMock.expect(transactionManager.getTransactionState(EasyMock.eq(transactionalId)))
-      .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, new TransactionMetadata(transactionalId, producerId, 1, 1, PrepareCommit, collection.mutable.Set.empty[TopicPartition], 0, time.milliseconds())))))
+      .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, new TransactionMetadata(transactionalId, producerId, producerId, producerEpoch, (producerEpoch - 1).toShort, 1, PrepareCommit, collection.mutable.Set.empty[TopicPartition], 0, time.milliseconds())))))
     EasyMock.replay(transactionManager)
 
     coordinator.handleEndTransaction(transactionalId, producerId, 1, TransactionResult.COMMIT, errorsCallback)
@@ -413,7 +433,7 @@ class TransactionCoordinatorTest {
   @Test
   def shouldReturnInvalidTxnRequestOnEndTxnRequestWhenStatusIsPrepareAbort(): Unit = {
     EasyMock.expect(transactionManager.getTransactionState(EasyMock.eq(transactionalId)))
-      .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, new TransactionMetadata(transactionalId, producerId, 1, 1, PrepareAbort, collection.mutable.Set.empty[TopicPartition], 0, time.milliseconds())))))
+      .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, new TransactionMetadata(transactionalId, producerId, producerId, 1, RecordBatch.NO_PRODUCER_EPOCH, 1, PrepareAbort, collection.mutable.Set.empty[TopicPartition], 0, time.milliseconds())))))
     EasyMock.replay(transactionManager)
 
     coordinator.handleEndTransaction(transactionalId, producerId, 1, TransactionResult.COMMIT, errorsCallback)
@@ -479,6 +499,29 @@ class TransactionCoordinatorTest {
     assertEquals(Errors.COORDINATOR_LOAD_IN_PROGRESS, error)
   }
 
+  @Test
+  def shouldReturnInvalidEpochOnEndTxnWhenEpochIsLarger(): Unit = {
+    val serverProducerEpoch = 1.toShort
+    verifyEndTxnEpoch(serverProducerEpoch, (serverProducerEpoch + 1).toShort)
+  }
+
+  @Test
+  def shouldReturnInvalidEpochOnEndTxnWhenEpochIsSmaller(): Unit = {
+    val serverProducerEpoch = 1.toShort
+    verifyEndTxnEpoch(serverProducerEpoch, (serverProducerEpoch - 1).toShort)
+  }
+
+  private def verifyEndTxnEpoch(metadataEpoch: Short, requestEpoch: Short): Unit = {
+    EasyMock.expect(transactionManager.getTransactionState(EasyMock.eq(transactionalId)))
+      .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch,
+        new TransactionMetadata(transactionalId, producerId, producerId, metadataEpoch, 0, 1, CompleteCommit, collection.mutable.Set.empty[TopicPartition], 0, time.milliseconds())))))
+    EasyMock.replay(transactionManager)
+
+    coordinator.handleEndTransaction(transactionalId, producerId, requestEpoch, TransactionResult.COMMIT, errorsCallback)
+    assertEquals(Errors.PRODUCER_FENCED, error)
+    EasyMock.verify(transactionManager)
+  }
+
   @Test
   def shouldIncrementEpochAndUpdateMetadataOnHandleInitPidWhenExistingEmptyTransaction(): Unit = {
     validateIncrementEpochAndUpdateMetadata(Empty)
@@ -506,13 +549,13 @@ class TransactionCoordinatorTest {
 
   @Test
   def shouldAbortTransactionOnHandleInitPidWhenExistingTransactionInOngoingState(): Unit = {
-    val txnMetadata = new TransactionMetadata(transactionalId, producerId, producerEpoch, txnTimeoutMs, Ongoing,
-      partitions, time.milliseconds(), time.milliseconds())
+    val txnMetadata = new TransactionMetadata(transactionalId, producerId, producerId, producerEpoch,
+      (producerEpoch - 1).toShort, txnTimeoutMs, Ongoing, partitions, time.milliseconds(), time.milliseconds())
 
     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()
 
@@ -520,80 +563,381 @@ class TransactionCoordinatorTest {
       .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, txnMetadata))))
       .anyTimes()
 
-    val originalMetadata = new TransactionMetadata(transactionalId, producerId, (producerEpoch + 1).toShort,
-      txnTimeoutMs, Ongoing, partitions, time.milliseconds(), time.milliseconds())
+    val originalMetadata = new TransactionMetadata(transactionalId, producerId, producerId, (producerEpoch + 1).toShort,
+      RecordBatch.NO_PRODUCER_EPOCH, txnTimeoutMs, Ongoing, partitions, time.milliseconds(), time.milliseconds())
     EasyMock.expect(transactionManager.appendTransactionToLog(
       EasyMock.eq(transactionalId),
       EasyMock.eq(coordinatorEpoch),
       EasyMock.eq(originalMetadata.prepareAbortOrCommit(PrepareAbort, time.milliseconds())),
       EasyMock.capture(capturedErrorsCallback),
-      EasyMock.anyObject()))
-      .andAnswer(new IAnswer[Unit] {
-        override def answer(): Unit = {
-          capturedErrorsCallback.getValue.apply(Errors.NONE)
-        }
-      })
+      EasyMock.anyObject())
+    ).andAnswer(() => capturedErrorsCallback.getValue.apply(Errors.NONE))
 
     EasyMock.replay(transactionManager)
 
-    coordinator.handleInitProducerId(transactionalId, txnTimeoutMs, initProducerIdMockCallback)
+    coordinator.handleInitProducerId(transactionalId, txnTimeoutMs, None, initProducerIdMockCallback)
 
     assertEquals(InitProducerIdResult(-1, -1, Errors.CONCURRENT_TRANSACTIONS), result)
+
+    EasyMock.verify(transactionManager)
+  }
+
+  @Test
+  def shouldFailToAbortTransactionOnHandleInitPidWhenProducerEpochIsSmaller(): Unit = {
+    val txnMetadata = new TransactionMetadata(transactionalId, producerId, producerId, producerEpoch,
+      (producerEpoch - 1).toShort, txnTimeoutMs, Ongoing, partitions, time.milliseconds(), time.milliseconds())
+
+    EasyMock.expect(transactionManager.validateTransactionTimeoutMs(EasyMock.anyInt()))
+      .andReturn(true)
+
+    EasyMock.expect(transactionManager.putTransactionStateIfNotExists(EasyMock.anyObject[TransactionMetadata]()))
+      .andReturn(Right(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, txnMetadata)))
+      .anyTimes()
+
+    EasyMock.expect(transactionManager.getTransactionState(EasyMock.eq(transactionalId)))
+      .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, txnMetadata))))
+      .times(1)
+
+    val bumpedTxnMetadata = new TransactionMetadata(transactionalId, producerId, producerId, (producerEpoch + 2).toShort,
+      (producerEpoch - 1).toShort, txnTimeoutMs, Ongoing, partitions, time.milliseconds(), time.milliseconds())
+
+    EasyMock.expect(transactionManager.getTransactionState(EasyMock.eq(transactionalId)))
+      .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, bumpedTxnMetadata))))
+      .times(1)
+
+    EasyMock.replay(transactionManager)
+
+    coordinator.handleInitProducerId(transactionalId, txnTimeoutMs, None, initProducerIdMockCallback)
+
+    assertEquals(InitProducerIdResult(-1, -1, Errors.PRODUCER_FENCED), result)
+
+    EasyMock.verify(transactionManager)
+  }
+
+  @Test
+  def shouldNotRepeatedlyBumpEpochDueToInitPidDuringOngoingTxnIfAppendToLogFails(): Unit = {
+    val txnMetadata = new TransactionMetadata(transactionalId, producerId, producerId, producerEpoch,
+      RecordBatch.NO_PRODUCER_EPOCH, txnTimeoutMs, Ongoing, partitions, time.milliseconds(), time.milliseconds())
+
+    EasyMock.expect(transactionManager.validateTransactionTimeoutMs(EasyMock.anyInt()))
+      .andReturn(true)
+      .anyTimes()
+
+    EasyMock.expect(transactionManager.putTransactionStateIfNotExists(EasyMock.anyObject[TransactionMetadata]()))
+      .andReturn(Right(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, txnMetadata)))
+      .anyTimes()
+
+    EasyMock.expect(transactionManager.getTransactionState(EasyMock.eq(transactionalId)))
+      .andAnswer(() => Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, txnMetadata))))
+      .anyTimes()
+
+    val originalMetadata = new TransactionMetadata(transactionalId, producerId, producerId, (producerEpoch + 1).toShort,
+      RecordBatch.NO_PRODUCER_EPOCH, txnTimeoutMs, Ongoing, partitions, time.milliseconds(), time.milliseconds())
+    val txnTransitMetadata = originalMetadata.prepareAbortOrCommit(PrepareAbort, time.milliseconds())
+    EasyMock.expect(transactionManager.appendTransactionToLog(
+      EasyMock.eq(transactionalId),
+      EasyMock.eq(coordinatorEpoch),
+      EasyMock.eq(txnTransitMetadata),
+      EasyMock.capture(capturedErrorsCallback),
+      EasyMock.anyObject())
+    ).andAnswer(() => {
+      capturedErrorsCallback.getValue.apply(Errors.NOT_ENOUGH_REPLICAS)
+      txnMetadata.pendingState = None
+    }).times(2)
+
+    EasyMock.expect(transactionManager.appendTransactionToLog(
+      EasyMock.eq(transactionalId),
+      EasyMock.eq(coordinatorEpoch),
+      EasyMock.eq(txnTransitMetadata),
+      EasyMock.capture(capturedErrorsCallback),
+      EasyMock.anyObject())
+    ).andAnswer(() => {
+      capturedErrorsCallback.getValue.apply(Errors.NONE)
+
+      // For the successful call, execute the state transitions that would happen in appendTransactionToLog()
+      txnMetadata.completeTransitionTo(txnTransitMetadata)
+      txnMetadata.prepareComplete(time.milliseconds())
+    }).once()
+
+    EasyMock.replay(transactionManager)
+
+    // For the first two calls, verify that the epoch was only bumped once
+    coordinator.handleInitProducerId(transactionalId, txnTimeoutMs, None, initProducerIdMockCallback)
+    assertEquals(InitProducerIdResult(-1, -1, Errors.NOT_ENOUGH_REPLICAS), result)
+
+    assertEquals((producerEpoch + 1).toShort, txnMetadata.producerEpoch)
+    assertTrue(txnMetadata.hasFailedEpochFence)
+
+    coordinator.handleInitProducerId(transactionalId, txnTimeoutMs, None, initProducerIdMockCallback)
+    assertEquals(InitProducerIdResult(-1, -1, Errors.NOT_ENOUGH_REPLICAS), result)
+
+    assertEquals((producerEpoch + 1).toShort, txnMetadata.producerEpoch)
+    assertTrue(txnMetadata.hasFailedEpochFence)
+
+    // For the last, successful call, verify that the epoch was not bumped further
+    coordinator.handleInitProducerId(transactionalId, txnTimeoutMs, None, initProducerIdMockCallback)
+    assertEquals(InitProducerIdResult(-1, -1, Errors.CONCURRENT_TRANSACTIONS), result)
+
+    assertEquals((producerEpoch + 1).toShort, txnMetadata.producerEpoch)
+    assertFalse(txnMetadata.hasFailedEpochFence)
+
     EasyMock.verify(transactionManager)
   }
 
   @Test
   def shouldUseLastEpochToFenceWhenEpochsAreExhausted(): Unit = {
-    val txnMetadata = new TransactionMetadata(transactionalId, producerId, (Short.MaxValue - 1).toShort,
-      txnTimeoutMs, Ongoing, partitions, time.milliseconds(), time.milliseconds())
+    val txnMetadata = new TransactionMetadata(transactionalId, producerId, producerId, (Short.MaxValue - 1).toShort,
+      (Short.MaxValue - 2).toShort, txnTimeoutMs, Ongoing, partitions, time.milliseconds(), time.milliseconds())
     assertTrue(txnMetadata.isProducerEpochExhausted)
 
     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()
 
     EasyMock.expect(transactionManager.getTransactionState(EasyMock.eq(transactionalId)))
       .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, txnMetadata))))
-      .anyTimes()
+      .times(2)
+
+    val postFenceTxnMetadata = new TransactionMetadata(transactionalId, producerId, producerId, Short.MaxValue,
+      RecordBatch.NO_PRODUCER_EPOCH, txnTimeoutMs, PrepareAbort, partitions, time.milliseconds(), time.milliseconds())
+
+    EasyMock.expect(transactionManager.getTransactionState(EasyMock.eq(transactionalId)))
+      .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, postFenceTxnMetadata))))
+      .once()
 
     EasyMock.expect(transactionManager.appendTransactionToLog(
       EasyMock.eq(transactionalId),
       EasyMock.eq(coordinatorEpoch),
       EasyMock.eq(TxnTransitMetadata(
         producerId = producerId,
+        lastProducerId = producerId,
         producerEpoch = Short.MaxValue,
+        lastProducerEpoch = RecordBatch.NO_PRODUCER_EPOCH,
         txnTimeoutMs = txnTimeoutMs,
         txnState = PrepareAbort,
         topicPartitions = partitions.toSet,
         txnStartTimestamp = time.milliseconds(),
         txnLastUpdateTimestamp = time.milliseconds())),
       EasyMock.capture(capturedErrorsCallback),
-      EasyMock.anyObject()))
-      .andAnswer(new IAnswer[Unit] {
-        override def answer(): Unit = {
-          capturedErrorsCallback.getValue.apply(Errors.NONE)
-        }
-      })
+      EasyMock.anyObject())
+    ).andAnswer(() => capturedErrorsCallback.getValue.apply(Errors.NONE))
 
     EasyMock.replay(transactionManager)
 
-    coordinator.handleInitProducerId(transactionalId, txnTimeoutMs, initProducerIdMockCallback)
+    coordinator.handleInitProducerId(transactionalId, txnTimeoutMs, None, initProducerIdMockCallback)
     assertEquals(Short.MaxValue, txnMetadata.producerEpoch)
 
     assertEquals(InitProducerIdResult(-1, -1, Errors.CONCURRENT_TRANSACTIONS), result)
     EasyMock.verify(transactionManager)
   }
 
+  @Test
+  def testInitProducerIdWithNoLastProducerData(): Unit = {
+    // If the metadata doesn't include the previous producer data (for example, if it was written to the log by a broker
+    // on an old version), the retry case should fail
+    val txnMetadata = new TransactionMetadata(transactionalId, producerId, RecordBatch.NO_PRODUCER_ID, (producerEpoch + 1).toShort,
+      RecordBatch.NO_PRODUCER_EPOCH, txnTimeoutMs, Empty, partitions, time.milliseconds, time.milliseconds)
+
+    EasyMock.expect(transactionManager.validateTransactionTimeoutMs(EasyMock.anyInt()))
+      .andReturn(true).anyTimes()
+    EasyMock.expect(transactionManager.getTransactionState(EasyMock.eq(transactionalId)))
+      .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, txnMetadata)))).once
+    EasyMock.replay(transactionManager)
+
+    // Simulate producer trying to continue after new producer has already been initialized
+    coordinator.handleInitProducerId(transactionalId, txnTimeoutMs, Some(new ProducerIdAndEpoch(producerId, producerEpoch)),
+      initProducerIdMockCallback)
+    assertEquals(InitProducerIdResult(RecordBatch.NO_PRODUCER_ID, RecordBatch.NO_PRODUCER_EPOCH, Errors.PRODUCER_FENCED), result)
+  }
+
+  @Test
+  def testFenceProducerWhenMappingExistsWithDifferentProducerId(): Unit = {
+    // Existing transaction ID maps to new producer ID
+    val txnMetadata = new TransactionMetadata(transactionalId, producerId + 1, producerId, producerEpoch,
+      (producerEpoch - 1).toShort, txnTimeoutMs, Empty, partitions, time.milliseconds, time.milliseconds)
+
+    EasyMock.expect(transactionManager.validateTransactionTimeoutMs(EasyMock.anyInt()))
+      .andReturn(true).anyTimes()
+    EasyMock.expect(transactionManager.getTransactionState(EasyMock.eq(transactionalId)))
+      .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, txnMetadata)))).once
+    EasyMock.replay(transactionManager)
+
+    // Simulate producer trying to continue after new producer has already been initialized
+    coordinator.handleInitProducerId(transactionalId, txnTimeoutMs, Some(new ProducerIdAndEpoch(producerId, producerEpoch)),
+      initProducerIdMockCallback)
+    assertEquals(InitProducerIdResult(RecordBatch.NO_PRODUCER_ID, RecordBatch.NO_PRODUCER_EPOCH, Errors.PRODUCER_FENCED), result)
+  }
+
+  @Test
+  def testInitProducerIdWithCurrentEpochProvided(): Unit = {
+    mockPidManager()
+
+    val txnMetadata = new TransactionMetadata(transactionalId, producerId, producerId, 10,
+      9, txnTimeoutMs, Empty, partitions, time.milliseconds, time.milliseconds)
+
+    EasyMock.expect(transactionManager.validateTransactionTimeoutMs(EasyMock.anyInt()))
+      .andReturn(true).anyTimes()
+    EasyMock.expect(transactionManager.getTransactionState(EasyMock.eq(transactionalId)))
+      .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, txnMetadata)))).times(2)
+
+    EasyMock.expect(transactionManager.appendTransactionToLog(
+      EasyMock.eq(transactionalId),
+      EasyMock.eq(coordinatorEpoch),
+      EasyMock.anyObject().asInstanceOf[TxnTransitMetadata],
+      EasyMock.capture(capturedErrorsCallback),
+      EasyMock.anyObject())
+    ).andAnswer(() => {
+      capturedErrorsCallback.getValue.apply(Errors.NONE)
+      txnMetadata.pendingState = None
+    }).times(2)
+
+    EasyMock.replay(pidManager, transactionManager)
+
+    // Re-initialization should succeed and bump the producer epoch
+    coordinator.handleInitProducerId(transactionalId, txnTimeoutMs, Some(new ProducerIdAndEpoch(producerId, 10)),
+      initProducerIdMockCallback)
+    assertEquals(InitProducerIdResult(producerId, 11, Errors.NONE), result)
+
+    // Simulate producer retrying after successfully re-initializing but failing to receive the response
+    coordinator.handleInitProducerId(transactionalId, txnTimeoutMs, Some(new ProducerIdAndEpoch(producerId, 10)),
+      initProducerIdMockCallback)
+    assertEquals(InitProducerIdResult(producerId, 11, Errors.NONE), result)
+  }
+
+  @Test
+  def testInitProducerIdStaleCurrentEpochProvided(): Unit = {
+    mockPidManager()
+
+    val txnMetadata = new TransactionMetadata(transactionalId, producerId, producerId, 10,
+      9, txnTimeoutMs, Empty, partitions, time.milliseconds, time.milliseconds)
+
+    EasyMock.expect(transactionManager.validateTransactionTimeoutMs(EasyMock.anyInt()))
+      .andReturn(true).anyTimes()
+    EasyMock.expect(transactionManager.getTransactionState(EasyMock.eq(transactionalId)))
+      .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, txnMetadata)))).times(2)
+
+    val capturedTxnTransitMetadata = Capture.newInstance[TxnTransitMetadata]
+    EasyMock.expect(transactionManager.appendTransactionToLog(
+      EasyMock.eq(transactionalId),
+      EasyMock.eq(coordinatorEpoch),
+      EasyMock.capture(capturedTxnTransitMetadata),
+      EasyMock.capture(capturedErrorsCallback),
+      EasyMock.anyObject())
+    ).andAnswer(() => {
+      capturedErrorsCallback.getValue.apply(Errors.NONE)
+      txnMetadata.pendingState = None
+      txnMetadata.producerEpoch = capturedTxnTransitMetadata.getValue.producerEpoch
+      txnMetadata.lastProducerEpoch = capturedTxnTransitMetadata.getValue.lastProducerEpoch
+    }).times(2)
+
+    EasyMock.replay(pidManager, transactionManager)
+
+    // With producer epoch at 10, new producer calls InitProducerId and should get epoch 11
+    coordinator.handleInitProducerId(transactionalId, txnTimeoutMs, None, initProducerIdMockCallback)
+    assertEquals(InitProducerIdResult(producerId, 11, Errors.NONE), result)
+
+    // Simulate old producer trying to continue from epoch 10
+    coordinator.handleInitProducerId(transactionalId, txnTimeoutMs, Some(new ProducerIdAndEpoch(producerId, 10)),
+      initProducerIdMockCallback)
+    assertEquals(InitProducerIdResult(RecordBatch.NO_PRODUCER_ID, RecordBatch.NO_PRODUCER_EPOCH, Errors.PRODUCER_FENCED), result)
+  }
+
+  @Test
+  def testRetryInitProducerIdAfterProducerIdRotation(): Unit = {
+    // Existing transaction ID maps to new producer ID
+    val txnMetadata = new TransactionMetadata(transactionalId, producerId, producerId, (Short.MaxValue - 1).toShort,
+      (Short.MaxValue - 2).toShort, txnTimeoutMs, Empty, partitions, time.milliseconds, time.milliseconds)
+
+    EasyMock.expect(pidManager.generateProducerId())
+      .andReturn(producerId + 1)
+      .anyTimes()
+
+    EasyMock.expect(transactionManager.validateTransactionTimeoutMs(EasyMock.anyInt()))
+      .andReturn(true).anyTimes()
+    EasyMock.expect(transactionManager.getTransactionState(EasyMock.eq(transactionalId)))
+      .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, txnMetadata)))).times(2)
+
+    EasyMock.expect(transactionManager.appendTransactionToLog(
+      EasyMock.eq(transactionalId),
+      EasyMock.eq(coordinatorEpoch),
+      EasyMock.capture(capturedTxnTransitMetadata),
+      EasyMock.capture(capturedErrorsCallback),
+      EasyMock.anyObject())
+    ).andAnswer(() => {
+      capturedErrorsCallback.getValue.apply(Errors.NONE)
+      txnMetadata.pendingState = None
+      txnMetadata.producerId = capturedTxnTransitMetadata.getValue.producerId
+      txnMetadata.lastProducerId = capturedTxnTransitMetadata.getValue.lastProducerId
+      txnMetadata.producerEpoch = capturedTxnTransitMetadata.getValue.producerEpoch
+      txnMetadata.lastProducerEpoch = capturedTxnTransitMetadata.getValue.lastProducerEpoch
+    }).once
+
+    EasyMock.replay(pidManager, transactionManager)
+
+    // Bump epoch and cause producer ID to be rotated
+    coordinator.handleInitProducerId(transactionalId, txnTimeoutMs, Some(new ProducerIdAndEpoch(producerId,
+      (Short.MaxValue - 1).toShort)), initProducerIdMockCallback)
+    assertEquals(InitProducerIdResult(producerId + 1, 0, Errors.NONE), result)
+
+    // Simulate producer retrying old request after producer bump
+    coordinator.handleInitProducerId(transactionalId, txnTimeoutMs, Some(new ProducerIdAndEpoch(producerId,
+      (Short.MaxValue - 1).toShort)), initProducerIdMockCallback)
+    assertEquals(InitProducerIdResult(producerId + 1, 0, Errors.NONE), result)
+  }
+
+  @Test
+  def testInitProducerIdWithInvalidEpochAfterProducerIdRotation(): Unit = {
+    // Existing transaction ID maps to new producer ID
+    val txnMetadata = new TransactionMetadata(transactionalId, producerId, producerId, (Short.MaxValue - 1).toShort,
+      (Short.MaxValue - 2).toShort, txnTimeoutMs, Empty, partitions, time.milliseconds, time.milliseconds)
+
+    EasyMock.expect(pidManager.generateProducerId())
+      .andReturn(producerId + 1)
+      .anyTimes()
+
+    EasyMock.expect(transactionManager.validateTransactionTimeoutMs(EasyMock.anyInt()))
+      .andReturn(true).anyTimes()
+    EasyMock.expect(transactionManager.getTransactionState(EasyMock.eq(transactionalId)))
+      .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, txnMetadata)))).times(2)
+
+    EasyMock.expect(transactionManager.appendTransactionToLog(
+      EasyMock.eq(transactionalId),
+      EasyMock.eq(coordinatorEpoch),
+      EasyMock.capture(capturedTxnTransitMetadata),
+      EasyMock.capture(capturedErrorsCallback),
+      EasyMock.anyObject())
+    ).andAnswer(() => {
+      capturedErrorsCallback.getValue.apply(Errors.NONE)
+      txnMetadata.pendingState = None
+      txnMetadata.producerId = capturedTxnTransitMetadata.getValue.producerId
+      txnMetadata.lastProducerId = capturedTxnTransitMetadata.getValue.lastProducerId
+      txnMetadata.producerEpoch = capturedTxnTransitMetadata.getValue.producerEpoch
+      txnMetadata.lastProducerEpoch = capturedTxnTransitMetadata.getValue.lastProducerEpoch
+    }).once
+
+    EasyMock.replay(pidManager, transactionManager)
+
+    // Bump epoch and cause producer ID to be rotated
+    coordinator.handleInitProducerId(transactionalId, txnTimeoutMs, Some(new ProducerIdAndEpoch(producerId,
+      (Short.MaxValue - 1).toShort)), initProducerIdMockCallback)
+    assertEquals(InitProducerIdResult(producerId + 1, 0, Errors.NONE), result)
+
+    // Validate that producer with old producer ID and stale epoch is fenced
+    coordinator.handleInitProducerId(transactionalId, txnTimeoutMs, Some(new ProducerIdAndEpoch(producerId,
+      (Short.MaxValue - 2).toShort)), initProducerIdMockCallback)
+    assertEquals(InitProducerIdResult(RecordBatch.NO_PRODUCER_ID, RecordBatch.NO_PRODUCER_EPOCH, Errors.PRODUCER_FENCED), result)
+  }
+
   @Test
   def shouldRemoveTransactionsForPartitionOnEmigration(): Unit = {
     EasyMock.expect(transactionManager.removeTransactionsForTxnTopicPartition(0, coordinatorEpoch))
     EasyMock.expect(transactionMarkerChannelManager.removeMarkersForTxnTopicPartition(0))
     EasyMock.replay(transactionManager, transactionMarkerChannelManager)
 
-    coordinator.handleTxnEmigration(0, coordinatorEpoch)
+    coordinator.onResignation(0, Some(coordinatorEpoch))
 
     EasyMock.verify(transactionManager, transactionMarkerChannelManager)
   }
@@ -601,9 +945,8 @@ class TransactionCoordinatorTest {
   @Test
   def shouldAbortExpiredTransactionsInOngoingStateAndBumpEpoch(): Unit = {
     val now = time.milliseconds()
-    val txnMetadata = new TransactionMetadata(transactionalId, producerId, producerEpoch, txnTimeoutMs, Ongoing,
-      partitions, now, now)
-
+    val txnMetadata = new TransactionMetadata(transactionalId, producerId, producerId, producerEpoch,
+      RecordBatch.NO_PRODUCER_EPOCH, txnTimeoutMs, Ongoing, partitions, now, now)
 
     EasyMock.expect(transactionManager.timedOutTransactions())
       .andReturn(List(TransactionalIdAndProducerIdEpoch(transactionalId, producerId, producerEpoch)))
@@ -611,19 +954,15 @@ class TransactionCoordinatorTest {
       .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, txnMetadata))))
       .times(2)
 
-    val bumpedEpoch = (producerEpoch + 1).toShort
-    val expectedTransition = TxnTransitMetadata(producerId, bumpedEpoch, txnTimeoutMs, PrepareAbort,
-      partitions.toSet, now, now + TransactionStateManager.DefaultAbortTimedOutTransactionsIntervalMs)
+    val expectedTransition = TxnTransitMetadata(producerId, producerId, (producerEpoch + 1).toShort, RecordBatch.NO_PRODUCER_EPOCH,
+      txnTimeoutMs, PrepareAbort, partitions.toSet, now, now + TransactionStateManager.DefaultAbortTimedOutTransactionsIntervalMs)
 
     EasyMock.expect(transactionManager.appendTransactionToLog(EasyMock.eq(transactionalId),
       EasyMock.eq(coordinatorEpoch),
       EasyMock.eq(expectedTransition),
       EasyMock.capture(capturedErrorsCallback),
-      EasyMock.anyObject()))
-      .andAnswer(new IAnswer[Unit] {
-        override def answer(): Unit = {}
-      })
-    .once()
+      EasyMock.anyObject())
+    ).andAnswer(() => {}).once()
 
     EasyMock.replay(transactionManager, transactionMarkerChannelManager)
 
@@ -633,10 +972,36 @@ class TransactionCoordinatorTest {
     EasyMock.verify(transactionManager)
   }
 
+  @Test
+  def shouldNotAcceptSmallerEpochDuringTransactionExpiration(): Unit = {
+    val now = time.milliseconds()
+    val txnMetadata = new TransactionMetadata(transactionalId, producerId, producerId, producerEpoch,
+      RecordBatch.NO_PRODUCER_EPOCH, txnTimeoutMs, Ongoing, partitions, now, now)
+
+    EasyMock.expect(transactionManager.timedOutTransactions())
+      .andReturn(List(TransactionalIdAndProducerIdEpoch(transactionalId, producerId, producerEpoch)))
+    EasyMock.expect(transactionManager.getTransactionState(EasyMock.eq(transactionalId)))
+      .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, txnMetadata))))
+
+    val bumpedTxnMetadata = new TransactionMetadata(transactionalId, producerId, producerId, (producerEpoch + 2).toShort,
+      RecordBatch.NO_PRODUCER_EPOCH, txnTimeoutMs, Ongoing, partitions, now, now)
+    EasyMock.expect(transactionManager.getTransactionState(EasyMock.eq(transactionalId)))
+      .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, bumpedTxnMetadata))))
+
+    EasyMock.replay(transactionManager, transactionMarkerChannelManager)
+
+    def checkOnEndTransactionComplete(txnIdAndPidEpoch: TransactionalIdAndProducerIdEpoch)(error: Errors): Unit = {
+      assertEquals(Errors.PRODUCER_FENCED, error)
+    }
+    coordinator.abortTimedOutTransactions(checkOnEndTransactionComplete)
+
+    EasyMock.verify(transactionManager)
+  }
+
   @Test
   def shouldNotAbortExpiredTransactionsThatHaveAPendingStateTransition(): Unit = {
-    val metadata = new TransactionMetadata(transactionalId, producerId, producerEpoch, txnTimeoutMs, Ongoing,
-      partitions, time.milliseconds(), time.milliseconds())
+    val metadata = new TransactionMetadata(transactionalId, producerId, producerId, producerEpoch,
+      RecordBatch.NO_PRODUCER_EPOCH, txnTimeoutMs, Ongoing, partitions, time.milliseconds(), time.milliseconds())
     metadata.prepareAbortOrCommit(PrepareCommit, time.milliseconds())
 
     EasyMock.expect(transactionManager.timedOutTransactions())
@@ -652,22 +1017,83 @@ class TransactionCoordinatorTest {
     EasyMock.verify(transactionManager)
   }
 
-  private def validateRespondsWithConcurrentTransactionsOnInitPidWhenInPrepareState(state: TransactionState) = {
+  @Test
+  def shouldNotBumpEpochWhenAbortingExpiredTransactionIfAppendToLogFails(): Unit = {
+    val now = time.milliseconds()
+    val txnMetadata = new TransactionMetadata(transactionalId, producerId, producerId, producerEpoch,
+      RecordBatch.NO_PRODUCER_EPOCH, txnTimeoutMs, Ongoing, partitions, now, now)
+
+
+    EasyMock.expect(transactionManager.timedOutTransactions())
+      .andReturn(List(TransactionalIdAndProducerIdEpoch(transactionalId, producerId, producerEpoch)))
+    EasyMock.expect(transactionManager.getTransactionState(EasyMock.eq(transactionalId)))
+      .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, txnMetadata))))
+      .times(2)
+
+    val txnMetadataAfterAppendFailure = new TransactionMetadata(transactionalId, producerId, producerId, (producerEpoch + 1).toShort,
+      RecordBatch.NO_PRODUCER_EPOCH, txnTimeoutMs, Ongoing, partitions, now, now)
+    EasyMock.expect(transactionManager.getTransactionState(EasyMock.eq(transactionalId)))
+      .andAnswer(() => Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, txnMetadataAfterAppendFailure))))
+      .once
+
+    val bumpedEpoch = (producerEpoch + 1).toShort
+    val expectedTransition = TxnTransitMetadata(producerId, producerId, bumpedEpoch, RecordBatch.NO_PRODUCER_EPOCH, txnTimeoutMs,
+      PrepareAbort, partitions.toSet, now, now + TransactionStateManager.DefaultAbortTimedOutTransactionsIntervalMs)
+
+    EasyMock.expect(transactionManager.appendTransactionToLog(EasyMock.eq(transactionalId),
+      EasyMock.eq(coordinatorEpoch),
+      EasyMock.eq(expectedTransition),
+      EasyMock.capture(capturedErrorsCallback),
+      EasyMock.anyObject())
+    ).andAnswer(() => capturedErrorsCallback.getValue.apply(Errors.NOT_ENOUGH_REPLICAS)).once()
+
+    EasyMock.replay(transactionManager, transactionMarkerChannelManager)
+
+    coordinator.startup(false)
+    time.sleep(TransactionStateManager.DefaultAbortTimedOutTransactionsIntervalMs)
+    scheduler.tick()
+    EasyMock.verify(transactionManager)
+
+    assertEquals((producerEpoch + 1).toShort, txnMetadataAfterAppendFailure.producerEpoch)
+    assertTrue(txnMetadataAfterAppendFailure.hasFailedEpochFence)
+  }
+
+  @Test
+  def shouldNotBumpEpochWithPendingTransaction(): Unit = {
+    val txnMetadata = new TransactionMetadata(transactionalId, producerId, producerId, producerEpoch,
+      RecordBatch.NO_PRODUCER_EPOCH, txnTimeoutMs, Ongoing, partitions, time.milliseconds(), time.milliseconds())
+    txnMetadata.prepareAbortOrCommit(PrepareCommit, time.milliseconds())
+
     EasyMock.expect(transactionManager.validateTransactionTimeoutMs(EasyMock.anyInt()))
       .andReturn(true).anyTimes()
+    EasyMock.expect(transactionManager.getTransactionState(EasyMock.eq(transactionalId)))
+      .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, txnMetadata))))
+
+    EasyMock.replay(transactionManager)
+
+    coordinator.handleInitProducerId(transactionalId, txnTimeoutMs, Some(new ProducerIdAndEpoch(producerId, 10)),
+      initProducerIdMockCallback)
+    assertEquals(InitProducerIdResult(RecordBatch.NO_PRODUCER_ID, RecordBatch.NO_PRODUCER_EPOCH, Errors.CONCURRENT_TRANSACTIONS), result)
 
-    val metadata = new TransactionMetadata(transactionalId, 0, 0, 0, state, mutable.Set[TopicPartition](new TopicPartition("topic", 1)), 0, 0)
+    EasyMock.verify(transactionManager)
+  }
+
+  private def validateRespondsWithConcurrentTransactionsOnInitPidWhenInPrepareState(state: TransactionState): Unit = {
+    EasyMock.expect(transactionManager.validateTransactionTimeoutMs(EasyMock.anyInt()))
+      .andReturn(true).anyTimes()
+
+    val metadata = new TransactionMetadata(transactionalId, 0, 0, 0, RecordBatch.NO_PRODUCER_EPOCH, 0, state, mutable.Set[TopicPartition](new TopicPartition("topic", 1)), 0, 0)
     EasyMock.expect(transactionManager.getTransactionState(EasyMock.eq(transactionalId)))
       .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, metadata)))).anyTimes()
 
     EasyMock.replay(transactionManager)
 
-    coordinator.handleInitProducerId(transactionalId, 10, initProducerIdMockCallback)
+    coordinator.handleInitProducerId(transactionalId, 10, None, initProducerIdMockCallback)
 
     assertEquals(InitProducerIdResult(-1, -1, Errors.CONCURRENT_TRANSACTIONS), result)
   }
 
-  private def validateIncrementEpochAndUpdateMetadata(state: TransactionState) = {
+  private def validateIncrementEpochAndUpdateMetadata(state: TransactionState): Unit = {
     EasyMock.expect(pidManager.generateProducerId())
       .andReturn(producerId)
       .anyTimes()
@@ -675,7 +1101,7 @@ class TransactionCoordinatorTest {
     EasyMock.expect(transactionManager.validateTransactionTimeoutMs(EasyMock.anyInt()))
       .andReturn(true)
 
-    val metadata = new TransactionMetadata(transactionalId, producerId, producerEpoch, txnTimeoutMs, state, mutable.Set.empty[TopicPartition], time.milliseconds(), time.milliseconds())
+    val metadata = new TransactionMetadata(transactionalId, producerId, producerId, producerEpoch, RecordBatch.NO_PRODUCER_EPOCH, txnTimeoutMs, state, mutable.Set.empty[TopicPartition], time.milliseconds(), time.milliseconds())
     EasyMock.expect(transactionManager.getTransactionState(EasyMock.eq(transactionalId)))
       .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, metadata))))
 
@@ -686,17 +1112,15 @@ class TransactionCoordinatorTest {
       EasyMock.capture(capturedNewMetadata),
       EasyMock.capture(capturedErrorsCallback),
       EasyMock.anyObject()
-    )).andAnswer(new IAnswer[Unit] {
-      override def answer(): Unit = {
-        metadata.completeTransitionTo(capturedNewMetadata.getValue)
-        capturedErrorsCallback.getValue.apply(Errors.NONE)
-      }
+    )).andAnswer(() => {
+      metadata.completeTransitionTo(capturedNewMetadata.getValue)
+      capturedErrorsCallback.getValue.apply(Errors.NONE)
     })
 
     EasyMock.replay(pidManager, transactionManager)
 
     val newTxnTimeoutMs = 10
-    coordinator.handleInitProducerId(transactionalId, newTxnTimeoutMs, initProducerIdMockCallback)
+    coordinator.handleInitProducerId(transactionalId, newTxnTimeoutMs, None, initProducerIdMockCallback)
 
     assertEquals(InitProducerIdResult(producerId, (producerEpoch + 1).toShort, Errors.NONE), result)
     assertEquals(newTxnTimeoutMs, metadata.txnTimeoutMs)
@@ -707,11 +1131,11 @@ class TransactionCoordinatorTest {
 
   private def mockPrepare(transactionState: TransactionState, runCallback: Boolean = false): TransactionMetadata = {
     val now = time.milliseconds()
-    val originalMetadata = new TransactionMetadata(transactionalId, producerId, producerEpoch, txnTimeoutMs,
-      Ongoing, partitions, now, now)
+    val originalMetadata = new TransactionMetadata(transactionalId, producerId, producerId, producerEpoch, RecordBatch.NO_PRODUCER_EPOCH,
+      txnTimeoutMs, Ongoing, partitions, now, now)
 
-    val transition = TxnTransitMetadata(producerId, producerEpoch, txnTimeoutMs, transactionState,
-      partitions.toSet, now, now)
+    val transition = TxnTransitMetadata(producerId, producerId, producerEpoch, RecordBatch.NO_PRODUCER_EPOCH, txnTimeoutMs,
+      transactionState, partitions.toSet, now, now)
 
     EasyMock.expect(transactionManager.getTransactionState(EasyMock.eq(transactionalId)))
       .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, originalMetadata))))
@@ -721,16 +1145,14 @@ class TransactionCoordinatorTest {
       EasyMock.eq(coordinatorEpoch),
       EasyMock.eq(transition),
       EasyMock.capture(capturedErrorsCallback),
-      EasyMock.anyObject()))
-      .andAnswer(new IAnswer[Unit] {
-        override def answer(): Unit = {
-          if (runCallback)
-            capturedErrorsCallback.getValue.apply(Errors.NONE)
-        }
-      }).once()
+      EasyMock.anyObject())
+    ).andAnswer(() => {
+      if (runCallback)
+        capturedErrorsCallback.getValue.apply(Errors.NONE)
+    }).once()
 
-    new TransactionMetadata(transactionalId, producerId, producerEpoch, txnTimeoutMs, transactionState, partitions,
-      time.milliseconds(), time.milliseconds())
+    new TransactionMetadata(transactionalId, producerId, producerId, producerEpoch, RecordBatch.NO_PRODUCER_EPOCH,
+      txnTimeoutMs, transactionState, partitions, time.milliseconds(), time.milliseconds())
   }
 
   def initProducerIdMockCallback(ret: InitProducerIdResult): Unit = {
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..52b0c33046b8f 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._
+import scala.jdk.CollectionConverters._
 
-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/TransactionMarkerChannelManagerTest.scala b/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionMarkerChannelManagerTest.scala
index a039c53b3de40..441b4e07ee100 100644
--- a/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionMarkerChannelManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionMarkerChannelManagerTest.scala
@@ -16,29 +16,33 @@
  */
 package kafka.coordinator.transaction
 
-import java.util.concurrent.locks.ReentrantReadWriteLock
+import java.util
+import java.util.Arrays.asList
+import java.util.Collections
+import java.util.concurrent.{Callable, Executors, Future}
 
-import kafka.server.{DelayedOperationPurgatory, KafkaConfig, MetadataCache}
-import kafka.utils.timer.MockTimer
+import kafka.common.RequestAndCompletionHandler
+import kafka.metrics.KafkaYammerMetrics
+import kafka.server.{KafkaConfig, MetadataCache}
 import kafka.utils.TestUtils
 import org.apache.kafka.clients.{ClientResponse, NetworkClient}
+import org.apache.kafka.common.protocol.{ApiKeys, Errors}
+import org.apache.kafka.common.record.RecordBatch
 import org.apache.kafka.common.requests.{RequestHeader, TransactionResult, WriteTxnMarkersRequest, WriteTxnMarkersResponse}
-import org.apache.kafka.common.utils.{MockTime, Utils}
+import org.apache.kafka.common.utils.MockTime
 import org.apache.kafka.common.{Node, TopicPartition}
-import org.easymock.{Capture, EasyMock, IAnswer}
+import org.easymock.{Capture, EasyMock}
 import org.junit.Assert._
 import org.junit.Test
-import com.yammer.metrics.Metrics
-import kafka.common.RequestAndCompletionHandler
-import org.apache.kafka.common.protocol.{ApiKeys, Errors}
 
-import scala.collection.JavaConverters._
+import scala.jdk.CollectionConverters._
 import scala.collection.mutable
+import scala.util.Try
 
 class TransactionMarkerChannelManagerTest {
-  private val metadataCache = EasyMock.createNiceMock(classOf[MetadataCache])
-  private val networkClient = EasyMock.createNiceMock(classOf[NetworkClient])
-  private val txnStateManager = EasyMock.createNiceMock(classOf[TransactionStateManager])
+  private val metadataCache: MetadataCache = EasyMock.createNiceMock(classOf[MetadataCache])
+  private val networkClient: NetworkClient = EasyMock.createNiceMock(classOf[NetworkClient])
+  private val txnStateManager: TransactionStateManager = EasyMock.mock(classOf[TransactionStateManager])
 
   private val partition1 = new TopicPartition("topic1", 0)
   private val partition2 = new TopicPartition("topic1", 1)
@@ -50,21 +54,18 @@ class TransactionMarkerChannelManagerTest {
   private val producerId1 = 0.asInstanceOf[Long]
   private val producerId2 = 1.asInstanceOf[Long]
   private val producerEpoch = 0.asInstanceOf[Short]
+  private val lastProducerEpoch = RecordBatch.NO_PRODUCER_EPOCH
   private val txnTopicPartition1 = 0
   private val txnTopicPartition2 = 1
   private val coordinatorEpoch = 0
   private val txnTimeoutMs = 0
   private val txnResult = TransactionResult.COMMIT
-  private val txnMetadata1 = new TransactionMetadata(transactionalId1, producerId1, producerEpoch, txnTimeoutMs,
-    PrepareCommit, mutable.Set[TopicPartition](partition1, partition2), 0L, 0L)
-  private val txnMetadata2 = new TransactionMetadata(transactionalId2, producerId2, producerEpoch, txnTimeoutMs,
-    PrepareCommit, mutable.Set[TopicPartition](partition1), 0L, 0L)
+  private val txnMetadata1 = new TransactionMetadata(transactionalId1, producerId1, producerId1, producerEpoch, lastProducerEpoch,
+    txnTimeoutMs, PrepareCommit, mutable.Set[TopicPartition](partition1, partition2), 0L, 0L)
+  private val txnMetadata2 = new TransactionMetadata(transactionalId2, producerId2, producerId2, producerEpoch, lastProducerEpoch,
+    txnTimeoutMs, PrepareCommit, mutable.Set[TopicPartition](partition1), 0L, 0L)
 
   private val capturedErrorsCallback: Capture[Errors => Unit] = EasyMock.newCapture()
-
-  private val txnMarkerPurgatory = new DelayedOperationPurgatory[DelayedTxnMarker]("txn-purgatory-name",
-    new MockTimer,
-    reaperEnabled = false)
   private val time = new MockTime
 
   private val channelManager = new TransactionMarkerChannelManager(
@@ -72,7 +73,6 @@ class TransactionMarkerChannelManagerTest {
     metadataCache,
     networkClient,
     txnStateManager,
-    txnMarkerPurgatory,
     time)
 
   private def mockCache(): Unit = {
@@ -88,10 +88,70 @@ class TransactionMarkerChannelManagerTest {
     EasyMock.expect(txnStateManager.getTransactionState(EasyMock.eq(transactionalId2)))
       .andReturn(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, txnMetadata2))))
       .anyTimes()
-    val stateLock = new ReentrantReadWriteLock
-    EasyMock.expect(txnStateManager.stateReadLock)
-      .andReturn(stateLock.readLock)
-      .anyTimes()
+  }
+
+  @Test
+  def shouldOnlyWriteTxnCompletionOnce(): Unit = {
+    mockCache()
+
+    val expectedTransition = txnMetadata2.prepareComplete(time.milliseconds())
+
+    EasyMock.expect(metadataCache.getPartitionLeaderEndpoint(
+      EasyMock.eq(partition1.topic),
+      EasyMock.eq(partition1.partition),
+      EasyMock.anyObject())
+    ).andReturn(Some(broker1)).anyTimes()
+
+    EasyMock.expect(txnStateManager.appendTransactionToLog(
+      EasyMock.eq(transactionalId2),
+      EasyMock.eq(coordinatorEpoch),
+      EasyMock.eq(expectedTransition),
+      EasyMock.capture(capturedErrorsCallback),
+      EasyMock.anyObject()))
+      .andAnswer(() => {
+        txnMetadata2.completeTransitionTo(expectedTransition)
+        capturedErrorsCallback.getValue.apply(Errors.NONE)
+      }).once()
+
+    EasyMock.replay(txnStateManager, metadataCache)
+
+    var addMarkerFuture: Future[Try[Unit]] = null
+    val executor = Executors.newFixedThreadPool(1)
+    txnMetadata2.lock.lock()
+    try {
+      addMarkerFuture = executor.submit((() => {
+        Try(channelManager.addTxnMarkersToSend(coordinatorEpoch, txnResult,
+            txnMetadata2, expectedTransition))
+      }): Callable[Try[Unit]])
+
+      val header = new RequestHeader(ApiKeys.WRITE_TXN_MARKERS, 0, "client", 1)
+      val response = new WriteTxnMarkersResponse(
+        Collections.singletonMap(producerId2: java.lang.Long, Collections.singletonMap(partition1, Errors.NONE)))
+      val clientResponse = new ClientResponse(header, null, null,
+        time.milliseconds(), time.milliseconds(), false, null, null,
+        response)
+
+      TestUtils.waitUntilTrue(() => {
+        val requests = channelManager.drainQueuedTransactionMarkers()
+        if (requests.nonEmpty) {
+          assertEquals(1, requests.size)
+          val request = requests.head
+          request.handler.onComplete(clientResponse)
+          true
+        } else {
+          false
+        }
+      }, "Timed out waiting for expected WriteTxnMarkers request")
+    } finally {
+      txnMetadata2.lock.unlock()
+      executor.shutdown()
+    }
+
+    assertNotNull(addMarkerFuture)
+    assertTrue("Add marker task failed with exception " + addMarkerFuture.get().get,
+      addMarkerFuture.get().isSuccess)
+
+    EasyMock.verify(txnStateManager)
   }
 
   @Test
@@ -117,10 +177,10 @@ class TransactionMarkerChannelManagerTest {
 
     EasyMock.replay(metadataCache)
 
-    channelManager.addTxnMarkersToSend(transactionalId1, coordinatorEpoch, txnResult, txnMetadata1, txnMetadata1.prepareComplete(time.milliseconds()))
-    channelManager.addTxnMarkersToSend(transactionalId2, coordinatorEpoch, txnResult, txnMetadata2, txnMetadata2.prepareComplete(time.milliseconds()))
+    channelManager.addTxnMarkersToSend(coordinatorEpoch, txnResult, txnMetadata1, txnMetadata1.prepareComplete(time.milliseconds()))
+    channelManager.addTxnMarkersToSend(coordinatorEpoch, txnResult, txnMetadata2, txnMetadata2.prepareComplete(time.milliseconds()))
 
-    assertEquals(2, txnMarkerPurgatory.watched)
+    assertEquals(2, channelManager.numTxnsWithPendingMarkers)
     assertEquals(2, channelManager.queueForBroker(broker1.id).get.totalNumMarkers)
     assertEquals(1, channelManager.queueForBroker(broker1.id).get.totalNumMarkers(txnTopicPartition1))
     assertEquals(1, channelManager.queueForBroker(broker1.id).get.totalNumMarkers(txnTopicPartition2))
@@ -128,11 +188,11 @@ class TransactionMarkerChannelManagerTest {
     assertEquals(1, channelManager.queueForBroker(broker2.id).get.totalNumMarkers(txnTopicPartition1))
     assertEquals(0, channelManager.queueForBroker(broker2.id).get.totalNumMarkers(txnTopicPartition2))
 
-    val expectedBroker1Request = new WriteTxnMarkersRequest.Builder(
-      Utils.mkList(new WriteTxnMarkersRequest.TxnMarkerEntry(producerId1, producerEpoch, coordinatorEpoch, txnResult, Utils.mkList(partition1)),
-        new WriteTxnMarkersRequest.TxnMarkerEntry(producerId2, producerEpoch, coordinatorEpoch, txnResult, Utils.mkList(partition1)))).build()
-    val expectedBroker2Request = new WriteTxnMarkersRequest.Builder(
-      Utils.mkList(new WriteTxnMarkersRequest.TxnMarkerEntry(producerId1, producerEpoch, coordinatorEpoch, txnResult, Utils.mkList(partition2)))).build()
+    val expectedBroker1Request = new WriteTxnMarkersRequest.Builder(ApiKeys.WRITE_TXN_MARKERS.latestVersion(),
+      asList(new WriteTxnMarkersRequest.TxnMarkerEntry(producerId1, producerEpoch, coordinatorEpoch, txnResult, asList(partition1)),
+        new WriteTxnMarkersRequest.TxnMarkerEntry(producerId2, producerEpoch, coordinatorEpoch, txnResult, asList(partition1)))).build()
+    val expectedBroker2Request = new WriteTxnMarkersRequest.Builder(ApiKeys.WRITE_TXN_MARKERS.latestVersion(),
+      asList(new WriteTxnMarkersRequest.TxnMarkerEntry(producerId1, producerEpoch, coordinatorEpoch, txnResult, asList(partition2)))).build()
 
     val requests: Map[Node, WriteTxnMarkersRequest] = channelManager.generateRequests().map { handler =>
       (handler.destination, handler.request.asInstanceOf[WriteTxnMarkersRequest.Builder].build())
@@ -160,10 +220,9 @@ class TransactionMarkerChannelManagerTest {
 
     EasyMock.replay(metadataCache)
 
-    channelManager.addTxnMarkersToSend(transactionalId1, coordinatorEpoch, txnResult, txnMetadata1, txnMetadata1.prepareComplete(time.milliseconds()))
-    channelManager.addTxnMarkersToSend(transactionalId2, coordinatorEpoch, txnResult, txnMetadata2, txnMetadata2.prepareComplete(time.milliseconds()))
+    channelManager.addTxnMarkersToSend(coordinatorEpoch, txnResult, txnMetadata1, txnMetadata1.prepareComplete(time.milliseconds()))
 
-    assertEquals(1, txnMarkerPurgatory.watched)
+    assertEquals(1, channelManager.numTxnsWithPendingMarkers)
     assertEquals(1, channelManager.queueForBroker(broker2.id).get.totalNumMarkers)
     assertTrue(channelManager.queueForBroker(broker1.id).isEmpty)
     assertEquals(1, channelManager.queueForBroker(broker2.id).get.totalNumMarkers(txnTopicPartition1))
@@ -193,10 +252,10 @@ class TransactionMarkerChannelManagerTest {
 
     EasyMock.replay(metadataCache)
 
-    channelManager.addTxnMarkersToSend(transactionalId1, coordinatorEpoch, txnResult, txnMetadata1, txnMetadata1.prepareComplete(time.milliseconds()))
-    channelManager.addTxnMarkersToSend(transactionalId2, coordinatorEpoch, txnResult, txnMetadata2, txnMetadata2.prepareComplete(time.milliseconds()))
+    channelManager.addTxnMarkersToSend(coordinatorEpoch, txnResult, txnMetadata1, txnMetadata1.prepareComplete(time.milliseconds()))
+    channelManager.addTxnMarkersToSend(coordinatorEpoch, txnResult, txnMetadata2, txnMetadata2.prepareComplete(time.milliseconds()))
 
-    assertEquals(2, txnMarkerPurgatory.watched)
+    assertEquals(2, channelManager.numTxnsWithPendingMarkers)
     assertEquals(1, channelManager.queueForBroker(broker2.id).get.totalNumMarkers)
     assertTrue(channelManager.queueForBroker(broker1.id).isEmpty)
     assertEquals(1, channelManager.queueForBroker(broker2.id).get.totalNumMarkers(txnTopicPartition1))
@@ -205,11 +264,11 @@ class TransactionMarkerChannelManagerTest {
     assertEquals(1, channelManager.queueForUnknownBroker.totalNumMarkers(txnTopicPartition1))
     assertEquals(1, channelManager.queueForUnknownBroker.totalNumMarkers(txnTopicPartition2))
 
-    val expectedBroker1Request = new WriteTxnMarkersRequest.Builder(
-      Utils.mkList(new WriteTxnMarkersRequest.TxnMarkerEntry(producerId1, producerEpoch, coordinatorEpoch, txnResult, Utils.mkList(partition1)),
-        new WriteTxnMarkersRequest.TxnMarkerEntry(producerId2, producerEpoch, coordinatorEpoch, txnResult, Utils.mkList(partition1)))).build()
-    val expectedBroker2Request = new WriteTxnMarkersRequest.Builder(
-      Utils.mkList(new WriteTxnMarkersRequest.TxnMarkerEntry(producerId1, producerEpoch, coordinatorEpoch, txnResult, Utils.mkList(partition2)))).build()
+    val expectedBroker1Request = new WriteTxnMarkersRequest.Builder(ApiKeys.WRITE_TXN_MARKERS.latestVersion(),
+      asList(new WriteTxnMarkersRequest.TxnMarkerEntry(producerId1, producerEpoch, coordinatorEpoch, txnResult, asList(partition1)),
+        new WriteTxnMarkersRequest.TxnMarkerEntry(producerId2, producerEpoch, coordinatorEpoch, txnResult, asList(partition1)))).build()
+    val expectedBroker2Request = new WriteTxnMarkersRequest.Builder(ApiKeys.WRITE_TXN_MARKERS.latestVersion(),
+      asList(new WriteTxnMarkersRequest.TxnMarkerEntry(producerId1, producerEpoch, coordinatorEpoch, txnResult, asList(partition2)))).build()
 
     val firstDrainedRequests: Map[Node, WriteTxnMarkersRequest] = channelManager.generateRequests().map { handler =>
       (handler.destination, handler.request.asInstanceOf[WriteTxnMarkersRequest.Builder].build())
@@ -242,10 +301,10 @@ class TransactionMarkerChannelManagerTest {
 
     EasyMock.replay(metadataCache)
 
-    channelManager.addTxnMarkersToSend(transactionalId1, coordinatorEpoch, txnResult, txnMetadata1, txnMetadata1.prepareComplete(time.milliseconds()))
-    channelManager.addTxnMarkersToSend(transactionalId2, coordinatorEpoch, txnResult, txnMetadata2, txnMetadata2.prepareComplete(time.milliseconds()))
+    channelManager.addTxnMarkersToSend(coordinatorEpoch, txnResult, txnMetadata1, txnMetadata1.prepareComplete(time.milliseconds()))
+    channelManager.addTxnMarkersToSend(coordinatorEpoch, txnResult, txnMetadata2, txnMetadata2.prepareComplete(time.milliseconds()))
 
-    assertEquals(2, txnMarkerPurgatory.watched)
+    assertEquals(2, channelManager.numTxnsWithPendingMarkers)
     assertEquals(2, channelManager.queueForBroker(broker1.id).get.totalNumMarkers)
     assertEquals(1, channelManager.queueForBroker(broker1.id).get.totalNumMarkers(txnTopicPartition1))
     assertEquals(1, channelManager.queueForBroker(broker1.id).get.totalNumMarkers(txnTopicPartition2))
@@ -255,7 +314,7 @@ class TransactionMarkerChannelManagerTest {
 
     channelManager.removeMarkersForTxnTopicPartition(txnTopicPartition1)
 
-    assertEquals(1, txnMarkerPurgatory.watched)
+    assertEquals(1, channelManager.numTxnsWithPendingMarkers)
     assertEquals(1, channelManager.queueForBroker(broker1.id).get.totalNumMarkers)
     assertEquals(0, channelManager.queueForBroker(broker1.id).get.totalNumMarkers(txnTopicPartition1))
     assertEquals(1, channelManager.queueForBroker(broker1.id).get.totalNumMarkers(txnTopicPartition2))
@@ -287,27 +346,25 @@ class TransactionMarkerChannelManagerTest {
       EasyMock.eq(txnTransitionMetadata2),
       EasyMock.capture(capturedErrorsCallback),
       EasyMock.anyObject()))
-      .andAnswer(new IAnswer[Unit] {
-        override def answer(): Unit = {
-          txnMetadata2.completeTransitionTo(txnTransitionMetadata2)
-          capturedErrorsCallback.getValue.apply(Errors.NONE)
-        }
+      .andAnswer(() => {
+        txnMetadata2.completeTransitionTo(txnTransitionMetadata2)
+        capturedErrorsCallback.getValue.apply(Errors.NONE)
       }).once()
     EasyMock.replay(txnStateManager, metadataCache)
 
-    channelManager.addTxnMarkersToSend(transactionalId2, coordinatorEpoch, txnResult, txnMetadata2, txnTransitionMetadata2)
+    channelManager.addTxnMarkersToSend(coordinatorEpoch, txnResult, txnMetadata2, txnTransitionMetadata2)
 
     val requestAndHandlers: Iterable[RequestAndCompletionHandler] = channelManager.generateRequests()
 
     val response = new WriteTxnMarkersResponse(createPidErrorMap(Errors.NONE))
     for (requestAndHandler <- requestAndHandlers) {
-      requestAndHandler.handler.onComplete(new ClientResponse(new RequestHeader(ApiKeys.PRODUCE, 0, "client", 1),
-        null, null, 0, 0, false, null, response))
+      requestAndHandler.handler.onComplete(new ClientResponse(new RequestHeader(ApiKeys.WRITE_TXN_MARKERS, 0, "client", 1),
+        null, null, 0, 0, false, null, null, response))
     }
 
     EasyMock.verify(txnStateManager)
 
-    assertEquals(0, txnMarkerPurgatory.watched)
+    assertEquals(0, channelManager.numTxnsWithPendingMarkers)
     assertEquals(0, channelManager.queueForBroker(broker1.id).get.totalNumMarkers)
     assertEquals(None, txnMetadata2.pendingState)
     assertEquals(CompleteCommit, txnMetadata2.state)
@@ -336,27 +393,25 @@ class TransactionMarkerChannelManagerTest {
       EasyMock.eq(txnTransitionMetadata2),
       EasyMock.capture(capturedErrorsCallback),
       EasyMock.anyObject()))
-      .andAnswer(new IAnswer[Unit] {
-        override def answer(): Unit = {
-          txnMetadata2.pendingState = None
-          capturedErrorsCallback.getValue.apply(Errors.NOT_COORDINATOR)
-        }
+      .andAnswer(() => {
+        txnMetadata2.pendingState = None
+        capturedErrorsCallback.getValue.apply(Errors.NOT_COORDINATOR)
       }).once()
     EasyMock.replay(txnStateManager, metadataCache)
 
-    channelManager.addTxnMarkersToSend(transactionalId2, coordinatorEpoch, txnResult, txnMetadata2, txnTransitionMetadata2)
+    channelManager.addTxnMarkersToSend(coordinatorEpoch, txnResult, txnMetadata2, txnTransitionMetadata2)
 
     val requestAndHandlers: Iterable[RequestAndCompletionHandler] = channelManager.generateRequests()
 
     val response = new WriteTxnMarkersResponse(createPidErrorMap(Errors.NONE))
     for (requestAndHandler <- requestAndHandlers) {
-      requestAndHandler.handler.onComplete(new ClientResponse(new RequestHeader(ApiKeys.PRODUCE, 0, "client", 1),
-        null, null, 0, 0, false, null, response))
+      requestAndHandler.handler.onComplete(new ClientResponse(new RequestHeader(ApiKeys.WRITE_TXN_MARKERS, 0, "client", 1),
+        null, null, 0, 0, false, null, null, response))
     }
 
     EasyMock.verify(txnStateManager)
 
-    assertEquals(0, txnMarkerPurgatory.watched)
+    assertEquals(0, channelManager.numTxnsWithPendingMarkers)
     assertEquals(0, channelManager.queueForBroker(broker1.id).get.totalNumMarkers)
     assertEquals(None, txnMetadata2.pendingState)
     assertEquals(PrepareCommit, txnMetadata2.state)
@@ -385,28 +440,22 @@ class TransactionMarkerChannelManagerTest {
       EasyMock.eq(txnTransitionMetadata2),
       EasyMock.capture(capturedErrorsCallback),
       EasyMock.anyObject()))
-      .andAnswer(new IAnswer[Unit] {
-        override def answer(): Unit = {
-          capturedErrorsCallback.getValue.apply(Errors.COORDINATOR_NOT_AVAILABLE)
-        }
-      })
-      .andAnswer(new IAnswer[Unit] {
-      override def answer(): Unit = {
+      .andAnswer(() => capturedErrorsCallback.getValue.apply(Errors.COORDINATOR_NOT_AVAILABLE))
+      .andAnswer(() => {
         txnMetadata2.completeTransitionTo(txnTransitionMetadata2)
         capturedErrorsCallback.getValue.apply(Errors.NONE)
-      }
-    })
+      })
 
     EasyMock.replay(txnStateManager, metadataCache)
 
-    channelManager.addTxnMarkersToSend(transactionalId2, coordinatorEpoch, txnResult, txnMetadata2, txnTransitionMetadata2)
+    channelManager.addTxnMarkersToSend(coordinatorEpoch, txnResult, txnMetadata2, txnTransitionMetadata2)
 
     val requestAndHandlers: Iterable[RequestAndCompletionHandler] = channelManager.generateRequests()
 
     val response = new WriteTxnMarkersResponse(createPidErrorMap(Errors.NONE))
     for (requestAndHandler <- requestAndHandlers) {
-      requestAndHandler.handler.onComplete(new ClientResponse(new RequestHeader(ApiKeys.PRODUCE, 0, "client", 1),
-        null, null, 0, 0, false, null, response))
+      requestAndHandler.handler.onComplete(new ClientResponse(new RequestHeader(ApiKeys.WRITE_TXN_MARKERS, 0, "client", 1),
+        null, null, 0, 0, false, null, null, response))
     }
 
     // call this again so that append log will be retried
@@ -414,13 +463,13 @@ class TransactionMarkerChannelManagerTest {
 
     EasyMock.verify(txnStateManager)
 
-    assertEquals(0, txnMarkerPurgatory.watched)
+    assertEquals(0, channelManager.numTxnsWithPendingMarkers)
     assertEquals(0, channelManager.queueForBroker(broker1.id).get.totalNumMarkers)
     assertEquals(None, txnMetadata2.pendingState)
     assertEquals(CompleteCommit, txnMetadata2.state)
   }
 
-  private def createPidErrorMap(errors: Errors) = {
+  private def createPidErrorMap(errors: Errors): util.HashMap[java.lang.Long, util.Map[TopicPartition, Errors]] = {
     val pidMap = new java.util.HashMap[java.lang.Long, java.util.Map[TopicPartition, Errors]]()
     val errorsMap = new java.util.HashMap[TopicPartition, Errors]()
     errorsMap.put(partition1, errors)
@@ -430,13 +479,13 @@ class TransactionMarkerChannelManagerTest {
 
   @Test
   def shouldCreateMetricsOnStarting(): Unit = {
-    val metrics = Metrics.defaultRegistry.allMetrics.asScala
-
-    assertEquals(1, metrics
-      .filterKeys(_.getMBeanName == "kafka.coordinator.transaction:type=TransactionMarkerChannelManager,name=UnknownDestinationQueueSize")
-      .size)
-    assertEquals(1, metrics
-      .filterKeys(_.getMBeanName == "kafka.coordinator.transaction:type=TransactionMarkerChannelManager,name=LogAppendRetryQueueSize")
-      .size)
+    val metrics = KafkaYammerMetrics.defaultRegistry.allMetrics.asScala
+
+    assertEquals(1, metrics.count { case (k, _) =>
+      k.getMBeanName == "kafka.coordinator.transaction:type=TransactionMarkerChannelManager,name=UnknownDestinationQueueSize"
+    })
+    assertEquals(1, metrics.count { case (k, _) =>
+      k.getMBeanName == "kafka.coordinator.transaction:type=TransactionMarkerChannelManager,name=LogAppendRetryQueueSize"
+    })
   }
 }
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 41ec15969b170..07b2fa8d6cd12 100644
--- a/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionMarkerRequestCompletionHandlerTest.scala
+++ b/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionMarkerRequestCompletionHandlerTest.scala
@@ -17,13 +17,14 @@
 package kafka.coordinator.transaction
 
 import java.{lang, util}
+import java.util.Arrays.asList
 
 import org.apache.kafka.clients.ClientResponse
 import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.protocol.{ApiKeys, Errors}
+import org.apache.kafka.common.record.RecordBatch
 import org.apache.kafka.common.requests.{RequestHeader, TransactionResult, WriteTxnMarkersRequest, WriteTxnMarkersResponse}
-import org.apache.kafka.common.utils.Utils
-import org.easymock.{EasyMock, IAnswer}
+import org.easymock.EasyMock
 import org.junit.Assert._
 import org.junit.Test
 
@@ -36,20 +37,21 @@ class TransactionMarkerRequestCompletionHandlerTest {
   private val transactionalId = "txnId1"
   private val producerId = 0.asInstanceOf[Long]
   private val producerEpoch = 0.asInstanceOf[Short]
+  private val lastProducerEpoch = RecordBatch.NO_PRODUCER_EPOCH
   private val txnTimeoutMs = 0
   private val coordinatorEpoch = 0
   private val txnResult = TransactionResult.COMMIT
   private val topicPartition = new TopicPartition("topic1", 0)
-  private val txnIdAndMarkers =
-    Utils.mkList(
-      TxnIdAndMarkerEntry(transactionalId, new WriteTxnMarkersRequest.TxnMarkerEntry(producerId, producerEpoch, coordinatorEpoch, txnResult, Utils.mkList(topicPartition))))
+  private val txnIdAndMarkers = asList(
+      TxnIdAndMarkerEntry(transactionalId, new WriteTxnMarkersRequest.TxnMarkerEntry(producerId, producerEpoch, coordinatorEpoch, txnResult, asList(topicPartition))))
 
-  private val txnMetadata = new TransactionMetadata(transactionalId, producerId, producerEpoch, txnTimeoutMs,
-    PrepareCommit, mutable.Set[TopicPartition](topicPartition), 0L, 0L)
+  private val txnMetadata = new TransactionMetadata(transactionalId, producerId, producerId, producerEpoch, lastProducerEpoch,
+    txnTimeoutMs, PrepareCommit, mutable.Set[TopicPartition](topicPartition), 0L, 0L)
 
-  private val markerChannelManager = EasyMock.createNiceMock(classOf[TransactionMarkerChannelManager])
+  private val markerChannelManager: TransactionMarkerChannelManager =
+    EasyMock.createNiceMock(classOf[TransactionMarkerChannelManager])
 
-  private val txnStateManager = EasyMock.createNiceMock(classOf[TransactionStateManager])
+  private val txnStateManager: TransactionStateManager = EasyMock.createNiceMock(classOf[TransactionStateManager])
 
   private val handler = new TransactionMarkerRequestCompletionHandler(brokerId, txnStateManager, markerChannelManager, txnIdAndMarkers)
 
@@ -72,7 +74,7 @@ class TransactionMarkerRequestCompletionHandlerTest {
     EasyMock.replay(markerChannelManager)
 
     handler.onComplete(new ClientResponse(new RequestHeader(ApiKeys.PRODUCE, 0, "client", 1),
-      null, null, 0, 0, true, null, null))
+      null, null, 0, 0, true, null, null, null))
 
     EasyMock.verify(markerChannelManager)
   }
@@ -86,7 +88,7 @@ class TransactionMarkerRequestCompletionHandlerTest {
 
     try {
       handler.onComplete(new ClientResponse(new RequestHeader(ApiKeys.PRODUCE, 0, "client", 1),
-        null, null, 0, 0, false, null, response))
+        null, null, 0, 0, false, null, null, response))
       fail("should have thrown illegal argument exception")
     } catch {
       case _: IllegalStateException => // ok
@@ -175,8 +177,8 @@ class TransactionMarkerRequestCompletionHandlerTest {
   }
 
   @Test
-  def shouldRetryPartitionWhenNotLeaderForPartitionError(): Unit = {
-    verifyRetriesPartitionOnError(Errors.NOT_LEADER_FOR_PARTITION)
+  def shouldRetryPartitionWhenNotLeaderOrFollowerError(): Unit = {
+    verifyRetriesPartitionOnError(Errors.NOT_LEADER_OR_FOLLOWER)
   }
 
   @Test
@@ -189,6 +191,11 @@ class TransactionMarkerRequestCompletionHandlerTest {
     verifyRetriesPartitionOnError(Errors.NOT_ENOUGH_REPLICAS_AFTER_APPEND)
   }
 
+  @Test
+  def shouldRetryPartitionWhenKafkaStorageError(): Unit = {
+    verifyRetriesPartitionOnError(Errors.KAFKA_STORAGE_ERROR)
+  }
+
   @Test
   def shouldRemoveTopicPartitionFromWaitingSetOnUnsupportedForMessageFormat(): Unit = {
     mockCache()
@@ -204,7 +211,7 @@ class TransactionMarkerRequestCompletionHandlerTest {
 
     val response = new WriteTxnMarkersResponse(createProducerIdErrorMap(error))
     handler.onComplete(new ClientResponse(new RequestHeader(ApiKeys.PRODUCE, 0, "client", 1),
-      null, null, 0, 0, false, null, response))
+      null, null, 0, 0, false, null, null, response))
 
     assertEquals(txnMetadata.topicPartitions, mutable.Set[TopicPartition](topicPartition))
     EasyMock.verify(markerChannelManager)
@@ -216,7 +223,7 @@ class TransactionMarkerRequestCompletionHandlerTest {
     val response = new WriteTxnMarkersResponse(createProducerIdErrorMap(error))
     try {
       handler.onComplete(new ClientResponse(new RequestHeader(ApiKeys.PRODUCE, 0, "client", 1),
-        null, null, 0, 0, false, null, response))
+        null, null, 0, 0, false, null, null, response))
       fail("should have thrown illegal state exception")
     } catch {
       case _: IllegalStateException => // ok
@@ -226,18 +233,14 @@ class TransactionMarkerRequestCompletionHandlerTest {
   private def verifyCompleteDelayedOperationOnError(error: Errors): Unit = {
 
     var completed = false
-    EasyMock.expect(markerChannelManager.completeSendMarkersForTxnId(transactionalId))
-      .andAnswer(new IAnswer[Unit] {
-        override def answer(): Unit = {
-          completed = true
-        }
-      })
+    EasyMock.expect(markerChannelManager.maybeWriteTxnCompletion(transactionalId))
+      .andAnswer(() => completed = true)
       .once()
     EasyMock.replay(markerChannelManager)
 
     val response = new WriteTxnMarkersResponse(createProducerIdErrorMap(error))
     handler.onComplete(new ClientResponse(new RequestHeader(ApiKeys.PRODUCE, 0, "client", 1),
-      null, null, 0, 0, false, null, response))
+      null, null, 0, 0, false, null, null, response))
 
     assertTrue(txnMetadata.topicPartitions.isEmpty)
     assertTrue(completed)
@@ -247,17 +250,13 @@ class TransactionMarkerRequestCompletionHandlerTest {
 
     var removed = false
     EasyMock.expect(markerChannelManager.removeMarkersForTxnId(transactionalId))
-      .andAnswer(new IAnswer[Unit] {
-        override def answer(): Unit = {
-          removed = true
-        }
-      })
+      .andAnswer(() => removed = true)
       .once()
     EasyMock.replay(markerChannelManager)
 
     val response = new WriteTxnMarkersResponse(createProducerIdErrorMap(error))
     handler.onComplete(new ClientResponse(new RequestHeader(ApiKeys.PRODUCE, 0, "client", 1),
-      null, null, 0, 0, false, null, response))
+      null, null, 0, 0, false, null, null, response))
 
     assertTrue(removed)
   }
diff --git a/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionMetadataTest.scala b/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionMetadataTest.scala
index 3a4390c75c17e..ca2d1dce87e1e 100644
--- a/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionMetadataTest.scala
+++ b/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionMetadataTest.scala
@@ -17,87 +17,285 @@
 package kafka.coordinator.transaction
 
 import kafka.utils.MockTime
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.protocol.Errors
 import org.apache.kafka.common.record.RecordBatch
 import org.junit.Assert._
 import org.junit.Test
+import org.scalatest.Assertions
 
 import scala.collection.mutable
 
 class TransactionMetadataTest {
 
   val time = new MockTime()
+  val producerId = 23423L
+  val transactionalId = "txnlId"
 
   @Test
   def testInitializeEpoch(): Unit = {
-    val transactionalId = "txnlId"
-    val producerId = 23423L
     val producerEpoch = RecordBatch.NO_PRODUCER_EPOCH
 
     val txnMetadata = new TransactionMetadata(
       transactionalId = transactionalId,
       producerId = producerId,
+      lastProducerId = RecordBatch.NO_PRODUCER_ID,
       producerEpoch = producerEpoch,
+      lastProducerEpoch = RecordBatch.NO_PRODUCER_EPOCH,
       txnTimeoutMs = 30000,
       state = Empty,
       topicPartitions = mutable.Set.empty,
       txnLastUpdateTimestamp = time.milliseconds())
 
-    val transitMetadata = txnMetadata.prepareIncrementProducerEpoch(30000, time.milliseconds())
+    val transitMetadata = prepareSuccessfulIncrementProducerEpoch(txnMetadata, None)
     txnMetadata.completeTransitionTo(transitMetadata)
     assertEquals(producerId, txnMetadata.producerId)
     assertEquals(0, txnMetadata.producerEpoch)
+    assertEquals(RecordBatch.NO_PRODUCER_EPOCH, txnMetadata.lastProducerEpoch)
   }
 
   @Test
   def testNormalEpochBump(): Unit = {
-    val transactionalId = "txnlId"
-    val producerId = 23423L
     val producerEpoch = 735.toShort
 
     val txnMetadata = new TransactionMetadata(
       transactionalId = transactionalId,
       producerId = producerId,
+      lastProducerId = RecordBatch.NO_PRODUCER_ID,
       producerEpoch = producerEpoch,
+      lastProducerEpoch = RecordBatch.NO_PRODUCER_EPOCH,
       txnTimeoutMs = 30000,
       state = Empty,
       topicPartitions = mutable.Set.empty,
       txnLastUpdateTimestamp = time.milliseconds())
 
-    val transitMetadata = txnMetadata.prepareIncrementProducerEpoch(30000, time.milliseconds())
+    val transitMetadata = prepareSuccessfulIncrementProducerEpoch(txnMetadata, None)
     txnMetadata.completeTransitionTo(transitMetadata)
     assertEquals(producerId, txnMetadata.producerId)
     assertEquals(producerEpoch + 1, txnMetadata.producerEpoch)
+    assertEquals(RecordBatch.NO_PRODUCER_EPOCH, txnMetadata.lastProducerEpoch)
   }
 
   @Test(expected = classOf[IllegalStateException])
   def testBumpEpochNotAllowedIfEpochsExhausted(): Unit = {
-    val transactionalId = "txnlId"
-    val producerId = 23423L
     val producerEpoch = (Short.MaxValue - 1).toShort
 
     val txnMetadata = new TransactionMetadata(
       transactionalId = transactionalId,
       producerId = producerId,
+      lastProducerId = RecordBatch.NO_PRODUCER_ID,
       producerEpoch = producerEpoch,
+      lastProducerEpoch = RecordBatch.NO_PRODUCER_EPOCH,
       txnTimeoutMs = 30000,
       state = Empty,
       topicPartitions = mutable.Set.empty,
       txnLastUpdateTimestamp = time.milliseconds())
     assertTrue(txnMetadata.isProducerEpochExhausted)
 
-    txnMetadata.prepareIncrementProducerEpoch(30000, time.milliseconds())
+    txnMetadata.prepareIncrementProducerEpoch(30000, None, time.milliseconds())
+  }
+
+  @Test
+  def testTolerateUpdateTimeShiftDuringEpochBump(): Unit = {
+    val producerEpoch: Short = 1
+    val txnMetadata = new TransactionMetadata(
+      transactionalId = transactionalId,
+      producerId = producerId,
+      lastProducerId = RecordBatch.NO_PRODUCER_ID,
+      producerEpoch = producerEpoch,
+      lastProducerEpoch = RecordBatch.NO_PRODUCER_EPOCH,
+      txnTimeoutMs = 30000,
+      state = Empty,
+      topicPartitions = mutable.Set.empty,
+      txnStartTimestamp = 1L,
+      txnLastUpdateTimestamp = time.milliseconds())
+
+    // let new time be smaller
+    val transitMetadata = prepareSuccessfulIncrementProducerEpoch(txnMetadata, Option(producerEpoch),
+      Some(time.milliseconds() - 1))
+    txnMetadata.completeTransitionTo(transitMetadata)
+    assertEquals(producerId, txnMetadata.producerId)
+    assertEquals(producerEpoch + 1, txnMetadata.producerEpoch)
+    assertEquals(producerEpoch, txnMetadata.lastProducerEpoch)
+    assertEquals(1L, txnMetadata.txnStartTimestamp)
+    assertEquals(time.milliseconds() - 1, txnMetadata.txnLastUpdateTimestamp)
+  }
+
+  @Test
+  def testTolerateUpdateTimeResetDuringProducerIdRotation(): Unit = {
+    val producerEpoch: Short = 1
+    val txnMetadata = new TransactionMetadata(
+      transactionalId = transactionalId,
+      producerId = producerId,
+      lastProducerId = RecordBatch.NO_PRODUCER_ID,
+      producerEpoch = producerEpoch,
+      lastProducerEpoch = RecordBatch.NO_PRODUCER_EPOCH,
+      txnTimeoutMs = 30000,
+      state = Empty,
+      topicPartitions = mutable.Set.empty,
+      txnStartTimestamp = 1L,
+      txnLastUpdateTimestamp = time.milliseconds())
+
+    // let new time be smaller
+    val transitMetadata = txnMetadata.prepareProducerIdRotation(producerId + 1, 30000, time.milliseconds() - 1, recordLastEpoch = true)
+    txnMetadata.completeTransitionTo(transitMetadata)
+    assertEquals(producerId + 1, txnMetadata.producerId)
+    assertEquals(producerEpoch, txnMetadata.lastProducerEpoch)
+    assertEquals(0, txnMetadata.producerEpoch)
+    assertEquals(1L, txnMetadata.txnStartTimestamp)
+    assertEquals(time.milliseconds() - 1, txnMetadata.txnLastUpdateTimestamp)
+  }
+
+  @Test
+  def testTolerateTimeShiftDuringAddPartitions(): Unit = {
+    val producerEpoch: Short = 1
+    val txnMetadata = new TransactionMetadata(
+      transactionalId = transactionalId,
+      producerId = producerId,
+      lastProducerId = RecordBatch.NO_PRODUCER_ID,
+      producerEpoch = producerEpoch,
+      lastProducerEpoch = RecordBatch.NO_PRODUCER_EPOCH,
+      txnTimeoutMs = 30000,
+      state = Empty,
+      topicPartitions = mutable.Set.empty,
+      txnStartTimestamp = time.milliseconds(),
+      txnLastUpdateTimestamp = time.milliseconds())
+
+    // let new time be smaller; when transting from Empty the start time would be updated to the update-time
+    var transitMetadata = txnMetadata.prepareAddPartitions(Set[TopicPartition](new TopicPartition("topic1", 0)), time.milliseconds() - 1)
+    txnMetadata.completeTransitionTo(transitMetadata)
+    assertEquals(Set[TopicPartition](new TopicPartition("topic1", 0)), txnMetadata.topicPartitions)
+    assertEquals(producerId, txnMetadata.producerId)
+    assertEquals(RecordBatch.NO_PRODUCER_EPOCH, txnMetadata.lastProducerEpoch)
+    assertEquals(producerEpoch, txnMetadata.producerEpoch)
+    assertEquals(time.milliseconds() - 1, txnMetadata.txnStartTimestamp)
+    assertEquals(time.milliseconds() - 1, txnMetadata.txnLastUpdateTimestamp)
+
+    // add another partition, check that in Ongoing state the start timestamp would not change to update time
+    transitMetadata = txnMetadata.prepareAddPartitions(Set[TopicPartition](new TopicPartition("topic2", 0)), time.milliseconds() - 2)
+    txnMetadata.completeTransitionTo(transitMetadata)
+    assertEquals(Set[TopicPartition](new TopicPartition("topic1", 0), new TopicPartition("topic2", 0)), txnMetadata.topicPartitions)
+    assertEquals(producerId, txnMetadata.producerId)
+    assertEquals(RecordBatch.NO_PRODUCER_EPOCH, txnMetadata.lastProducerEpoch)
+    assertEquals(producerEpoch, txnMetadata.producerEpoch)
+    assertEquals(time.milliseconds() - 1, txnMetadata.txnStartTimestamp)
+    assertEquals(time.milliseconds() - 2, txnMetadata.txnLastUpdateTimestamp)
+  }
+
+  @Test
+  def testTolerateTimeShiftDuringPrepareCommit(): Unit = {
+    val producerEpoch: Short = 1
+    val txnMetadata = new TransactionMetadata(
+      transactionalId = transactionalId,
+      producerId = producerId,
+      lastProducerId = RecordBatch.NO_PRODUCER_ID,
+      producerEpoch = producerEpoch,
+      lastProducerEpoch = RecordBatch.NO_PRODUCER_EPOCH,
+      txnTimeoutMs = 30000,
+      state = Ongoing,
+      topicPartitions = mutable.Set.empty,
+      txnStartTimestamp = 1L,
+      txnLastUpdateTimestamp = time.milliseconds())
+
+    // let new time be smaller
+    val transitMetadata = txnMetadata.prepareAbortOrCommit(PrepareCommit, time.milliseconds() - 1)
+    txnMetadata.completeTransitionTo(transitMetadata)
+    assertEquals(PrepareCommit, txnMetadata.state)
+    assertEquals(producerId, txnMetadata.producerId)
+    assertEquals(RecordBatch.NO_PRODUCER_EPOCH, txnMetadata.lastProducerEpoch)
+    assertEquals(producerEpoch, txnMetadata.producerEpoch)
+    assertEquals(1L, txnMetadata.txnStartTimestamp)
+    assertEquals(time.milliseconds() - 1, txnMetadata.txnLastUpdateTimestamp)
+  }
+
+  @Test
+  def testTolerateTimeShiftDuringPrepareAbort(): Unit = {
+    val producerEpoch: Short = 1
+    val txnMetadata = new TransactionMetadata(
+      transactionalId = transactionalId,
+      producerId = producerId,
+      lastProducerId = RecordBatch.NO_PRODUCER_ID,
+      producerEpoch = producerEpoch,
+      lastProducerEpoch = RecordBatch.NO_PRODUCER_EPOCH,
+      txnTimeoutMs = 30000,
+      state = Ongoing,
+      topicPartitions = mutable.Set.empty,
+      txnStartTimestamp = 1L,
+      txnLastUpdateTimestamp = time.milliseconds())
+
+    // let new time be smaller
+    val transitMetadata = txnMetadata.prepareAbortOrCommit(PrepareAbort, time.milliseconds() - 1)
+    txnMetadata.completeTransitionTo(transitMetadata)
+    assertEquals(PrepareAbort, txnMetadata.state)
+    assertEquals(producerId, txnMetadata.producerId)
+    assertEquals(RecordBatch.NO_PRODUCER_EPOCH, txnMetadata.lastProducerEpoch)
+    assertEquals(producerEpoch, txnMetadata.producerEpoch)
+    assertEquals(1L, txnMetadata.txnStartTimestamp)
+    assertEquals(time.milliseconds() - 1, txnMetadata.txnLastUpdateTimestamp)
+  }
+
+  @Test
+  def testTolerateTimeShiftDuringCompleteCommit(): Unit = {
+    val producerEpoch: Short = 1
+    val txnMetadata = new TransactionMetadata(
+      transactionalId = transactionalId,
+      producerId = producerId,
+      lastProducerId = RecordBatch.NO_PRODUCER_ID,
+      producerEpoch = producerEpoch,
+      lastProducerEpoch = RecordBatch.NO_PRODUCER_EPOCH,
+      txnTimeoutMs = 30000,
+      state = PrepareCommit,
+      topicPartitions = mutable.Set.empty,
+      txnStartTimestamp = 1L,
+      txnLastUpdateTimestamp = time.milliseconds())
+
+    // let new time be smaller
+    val transitMetadata = txnMetadata.prepareComplete(time.milliseconds() - 1)
+    txnMetadata.completeTransitionTo(transitMetadata)
+    assertEquals(CompleteCommit, txnMetadata.state)
+    assertEquals(producerId, txnMetadata.producerId)
+    assertEquals(RecordBatch.NO_PRODUCER_EPOCH, txnMetadata.lastProducerEpoch)
+    assertEquals(producerEpoch, txnMetadata.producerEpoch)
+    assertEquals(1L, txnMetadata.txnStartTimestamp)
+    assertEquals(time.milliseconds() - 1, txnMetadata.txnLastUpdateTimestamp)
+  }
+
+  @Test
+  def testTolerateTimeShiftDuringCompleteAbort(): Unit = {
+    val producerEpoch: Short = 1
+    val txnMetadata = new TransactionMetadata(
+      transactionalId = transactionalId,
+      producerId = producerId,
+      lastProducerId = RecordBatch.NO_PRODUCER_ID,
+      producerEpoch = producerEpoch,
+      lastProducerEpoch = RecordBatch.NO_PRODUCER_EPOCH,
+      txnTimeoutMs = 30000,
+      state = PrepareAbort,
+      topicPartitions = mutable.Set.empty,
+      txnStartTimestamp = 1L,
+      txnLastUpdateTimestamp = time.milliseconds())
+
+    // let new time be smaller
+    val transitMetadata = txnMetadata.prepareComplete(time.milliseconds() - 1)
+    txnMetadata.completeTransitionTo(transitMetadata)
+    assertEquals(CompleteAbort, txnMetadata.state)
+    assertEquals(producerId, txnMetadata.producerId)
+    assertEquals(RecordBatch.NO_PRODUCER_EPOCH, txnMetadata.lastProducerEpoch)
+    assertEquals(producerEpoch, txnMetadata.producerEpoch)
+    assertEquals(1L, txnMetadata.txnStartTimestamp)
+    assertEquals(time.milliseconds() - 1, txnMetadata.txnLastUpdateTimestamp)
   }
 
   @Test
   def testFenceProducerAfterEpochsExhausted(): Unit = {
-    val transactionalId = "txnlId"
-    val producerId = 23423L
     val producerEpoch = (Short.MaxValue - 1).toShort
 
     val txnMetadata = new TransactionMetadata(
       transactionalId = transactionalId,
       producerId = producerId,
+      lastProducerId = RecordBatch.NO_PRODUCER_ID,
       producerEpoch = producerEpoch,
+      lastProducerEpoch = RecordBatch.NO_PRODUCER_EPOCH,
       txnTimeoutMs = 30000,
       state = Ongoing,
       topicPartitions = mutable.Set.empty,
@@ -106,6 +304,7 @@ class TransactionMetadataTest {
 
     val fencingTransitMetadata = txnMetadata.prepareFenceProducerEpoch()
     assertEquals(Short.MaxValue, fencingTransitMetadata.producerEpoch)
+    assertEquals(RecordBatch.NO_PRODUCER_EPOCH, fencingTransitMetadata.lastProducerEpoch)
     assertEquals(Some(PrepareEpochFence), txnMetadata.pendingState)
 
     // We should reset the pending state to make way for the abort transition.
@@ -118,14 +317,14 @@ class TransactionMetadataTest {
 
   @Test(expected = classOf[IllegalStateException])
   def testFenceProducerNotAllowedIfItWouldOverflow(): Unit = {
-    val transactionalId = "txnlId"
-    val producerId = 23423L
     val producerEpoch = Short.MaxValue
 
     val txnMetadata = new TransactionMetadata(
       transactionalId = transactionalId,
       producerId = producerId,
+      lastProducerId = RecordBatch.NO_PRODUCER_ID,
       producerEpoch = producerEpoch,
+      lastProducerEpoch = RecordBatch.NO_PRODUCER_EPOCH,
       txnTimeoutMs = 30000,
       state = Ongoing,
       topicPartitions = mutable.Set.empty,
@@ -136,24 +335,26 @@ class TransactionMetadataTest {
 
   @Test
   def testRotateProducerId(): Unit = {
-    val transactionalId = "txnlId"
-    val producerId = 23423L
     val producerEpoch = (Short.MaxValue - 1).toShort
 
     val txnMetadata = new TransactionMetadata(
       transactionalId = transactionalId,
       producerId = producerId,
+      lastProducerId = RecordBatch.NO_PRODUCER_ID,
       producerEpoch = producerEpoch,
+      lastProducerEpoch = RecordBatch.NO_PRODUCER_EPOCH,
       txnTimeoutMs = 30000,
       state = Empty,
       topicPartitions = mutable.Set.empty,
       txnLastUpdateTimestamp = time.milliseconds())
 
     val newProducerId = 9893L
-    val transitMetadata = txnMetadata.prepareProducerIdRotation(newProducerId, 30000, time.milliseconds())
+    val transitMetadata = txnMetadata.prepareProducerIdRotation(newProducerId, 30000, time.milliseconds(), recordLastEpoch = true)
     txnMetadata.completeTransitionTo(transitMetadata)
     assertEquals(newProducerId, txnMetadata.producerId)
+    assertEquals(producerId, txnMetadata.lastProducerId)
     assertEquals(0, txnMetadata.producerEpoch)
+    assertEquals(producerEpoch, txnMetadata.lastProducerEpoch)
   }
 
   @Test(expected = classOf[IllegalStateException])
@@ -171,22 +372,117 @@ class TransactionMetadataTest {
     testRotateProducerIdInOngoingState(PrepareCommit)
   }
 
+  @Test
+  def testAttemptedEpochBumpWithNewlyCreatedMetadata(): Unit = {
+    val producerEpoch = 735.toShort
+
+    val txnMetadata = new TransactionMetadata(
+      transactionalId = transactionalId,
+      producerId = producerId,
+      lastProducerId = RecordBatch.NO_PRODUCER_ID,
+      producerEpoch = RecordBatch.NO_PRODUCER_EPOCH,
+      lastProducerEpoch = RecordBatch.NO_PRODUCER_EPOCH,
+      txnTimeoutMs = 30000,
+      state = Empty,
+      topicPartitions = mutable.Set.empty,
+      txnLastUpdateTimestamp = time.milliseconds())
+
+    val transitMetadata = prepareSuccessfulIncrementProducerEpoch(txnMetadata, Some(producerEpoch))
+    txnMetadata.completeTransitionTo(transitMetadata)
+    assertEquals(producerId, txnMetadata.producerId)
+    assertEquals(0, txnMetadata.producerEpoch)
+    assertEquals(RecordBatch.NO_PRODUCER_EPOCH, txnMetadata.lastProducerEpoch)
+  }
+
+  @Test
+  def testEpochBumpWithCurrentEpochProvided(): Unit = {
+    val producerEpoch = 735.toShort
+
+    val txnMetadata = new TransactionMetadata(
+      transactionalId = transactionalId,
+      producerId = producerId,
+      lastProducerId = RecordBatch.NO_PRODUCER_ID,
+      producerEpoch = producerEpoch,
+      lastProducerEpoch = RecordBatch.NO_PRODUCER_EPOCH,
+      txnTimeoutMs = 30000,
+      state = Empty,
+      topicPartitions = mutable.Set.empty,
+      txnLastUpdateTimestamp = time.milliseconds())
+
+    val transitMetadata = prepareSuccessfulIncrementProducerEpoch(txnMetadata, Some(producerEpoch))
+    txnMetadata.completeTransitionTo(transitMetadata)
+    assertEquals(producerId, txnMetadata.producerId)
+    assertEquals(producerEpoch + 1, txnMetadata.producerEpoch)
+    assertEquals(producerEpoch, txnMetadata.lastProducerEpoch)
+  }
+
+  @Test
+  def testAttemptedEpochBumpWithLastEpoch(): Unit = {
+    val producerEpoch = 735.toShort
+    val lastProducerEpoch = (producerEpoch - 1).toShort
+
+    val txnMetadata = new TransactionMetadata(
+      transactionalId = transactionalId,
+      producerId = producerId,
+      lastProducerId = RecordBatch.NO_PRODUCER_ID,
+      producerEpoch = producerEpoch,
+      lastProducerEpoch = lastProducerEpoch,
+      txnTimeoutMs = 30000,
+      state = Empty,
+      topicPartitions = mutable.Set.empty,
+      txnLastUpdateTimestamp = time.milliseconds())
+
+    val transitMetadata = prepareSuccessfulIncrementProducerEpoch(txnMetadata, Some(lastProducerEpoch))
+    txnMetadata.completeTransitionTo(transitMetadata)
+    assertEquals(producerId, txnMetadata.producerId)
+    assertEquals(producerEpoch, txnMetadata.producerEpoch)
+    assertEquals(lastProducerEpoch, txnMetadata.lastProducerEpoch)
+  }
+
+  @Test
+  def testAttemptedEpochBumpWithFencedEpoch(): Unit = {
+    val producerEpoch = 735.toShort
+    val lastProducerEpoch = (producerEpoch - 1).toShort
+
+    val txnMetadata = new TransactionMetadata(
+      transactionalId = transactionalId,
+      producerId = producerId,
+      lastProducerId = producerId,
+      producerEpoch = producerEpoch,
+      lastProducerEpoch = lastProducerEpoch,
+      txnTimeoutMs = 30000,
+      state = Empty,
+      topicPartitions = mutable.Set.empty,
+      txnLastUpdateTimestamp = time.milliseconds())
+
+    val result = txnMetadata.prepareIncrementProducerEpoch(30000, Some((lastProducerEpoch - 1).toShort),
+      time.milliseconds())
+    assertEquals(Left(Errors.PRODUCER_FENCED), result)
+  }
+
   private def testRotateProducerIdInOngoingState(state: TransactionState): Unit = {
-    val transactionalId = "txnlId"
-    val producerId = 23423L
     val producerEpoch = (Short.MaxValue - 1).toShort
 
     val txnMetadata = new TransactionMetadata(
       transactionalId = transactionalId,
       producerId = producerId,
+      lastProducerId = producerId,
       producerEpoch = producerEpoch,
+      lastProducerEpoch = RecordBatch.NO_PRODUCER_EPOCH,
       txnTimeoutMs = 30000,
       state = state,
       topicPartitions = mutable.Set.empty,
       txnLastUpdateTimestamp = time.milliseconds())
     val newProducerId = 9893L
-    txnMetadata.prepareProducerIdRotation(newProducerId, 30000, time.milliseconds())
+    txnMetadata.prepareProducerIdRotation(newProducerId, 30000, time.milliseconds(), recordLastEpoch = false)
   }
 
+  private def prepareSuccessfulIncrementProducerEpoch(txnMetadata: TransactionMetadata,
+                                                      expectedProducerEpoch: Option[Short],
+                                                      now: Option[Long] = None): TxnTransitMetadata = {
+    val result = txnMetadata.prepareIncrementProducerEpoch(30000, expectedProducerEpoch,
+      now.getOrElse(time.milliseconds()))
+    result.getOrElse(Assertions.fail(s"prepareIncrementProducerEpoch failed with $result"))
+  }
 
 }
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 20dfaa6ec1e4a..2a8d046de7fe0 100644
--- a/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionStateManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/coordinator/transaction/TransactionStateManagerTest.scala
@@ -16,29 +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, KafkaMetricsContext, Metrics}
 import org.apache.kafka.common.protocol.Errors
 import org.apache.kafka.common.record._
-import org.apache.kafka.common.requests.IsolationLevel
 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.jdk.CollectionConverters._
+import scala.collection.{Map, mutable}
 
 class TransactionStateManagerTest {
 
@@ -60,9 +62,11 @@ 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"
@@ -75,20 +79,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))
@@ -97,20 +101,112 @@ 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).toOption.flatten
+        .map(_.transactionMetadata.producerEpoch)
+    }
+
+    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(() => {
+      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
@@ -157,7 +253,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)
 
@@ -171,7 +267,7 @@ class TransactionStateManagerTest {
       _ => fail(transactionalId2 + "'s transaction state is already in the cache")
     )
 
-    transactionManager.loadTransactionsForTxnTopicPartition(partitionId, 0, (_, _, _, _, _) => ())
+    transactionManager.loadTransactionsForTxnTopicPartition(partitionId, 0, (_, _, _, _) => ())
 
     // let the time advance to trigger the background thread loading
     scheduler.tick()
@@ -214,7 +310,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
@@ -233,7 +329,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())
@@ -265,12 +361,12 @@ 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())
 
-    prepareForTxnMessageAppend(Errors.NOT_LEADER_FOR_PARTITION)
+    prepareForTxnMessageAppend(Errors.NOT_LEADER_OR_FOLLOWER)
     transactionManager.appendTransactionToLog(transactionalId1, coordinatorEpoch = 10, failedMetadata, assertCallback)
     assertEquals(Right(Some(CoordinatorEpochAndTxnMetadata(coordinatorEpoch, txnMetadata1))), transactionManager.getTransactionState(transactionalId1))
     assertTrue(txnMetadata1.pendingState.isEmpty)
@@ -283,7 +379,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)
@@ -295,7 +391,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())
@@ -307,9 +403,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())
@@ -327,9 +423,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())
@@ -341,11 +437,11 @@ class TransactionStateManagerTest {
   }
 
   @Test
-  def testAppendTransactionToLogWhileProducerFenced() = {
+  def testAppendTransactionToLogWhileProducerFenced(): Unit = {
     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
@@ -361,10 +457,10 @@ class TransactionStateManagerTest {
   }
 
   @Test(expected = classOf[IllegalStateException])
-  def testAppendTransactionToLogWhilePendingStateChanged() = {
+  def testAppendTransactionToLogWhilePendingStateChanged(): Unit = {
     // 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
@@ -393,12 +489,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()
@@ -419,68 +515,124 @@ class TransactionStateManagerTest {
   def shouldRemoveCompleteCommmitExpiredTransactionalIds(): Unit = {
     setupAndRunTransactionalIdExpiration(Errors.NONE, CompleteCommit)
     verifyMetadataDoesntExist(transactionalId1)
-    verifyMetadataDoesExist(transactionalId2)
+    verifyMetadataDoesExistAndIsUsable(transactionalId2)
   }
 
   @Test
   def shouldRemoveCompleteAbortExpiredTransactionalIds(): Unit = {
     setupAndRunTransactionalIdExpiration(Errors.NONE, CompleteAbort)
     verifyMetadataDoesntExist(transactionalId1)
-    verifyMetadataDoesExist(transactionalId2)
+    verifyMetadataDoesExistAndIsUsable(transactionalId2)
   }
 
   @Test
   def shouldRemoveEmptyExpiredTransactionalIds(): Unit = {
     setupAndRunTransactionalIdExpiration(Errors.NONE, Empty)
     verifyMetadataDoesntExist(transactionalId1)
-    verifyMetadataDoesExist(transactionalId2)
+    verifyMetadataDoesExistAndIsUsable(transactionalId2)
   }
 
   @Test
   def shouldNotRemoveExpiredTransactionalIdsIfLogAppendFails(): Unit = {
     setupAndRunTransactionalIdExpiration(Errors.NOT_ENOUGH_REPLICAS, CompleteAbort)
-    verifyMetadataDoesExist(transactionalId1)
-    verifyMetadataDoesExist(transactionalId2)
+    verifyMetadataDoesExistAndIsUsable(transactionalId1)
+    verifyMetadataDoesExistAndIsUsable(transactionalId2)
   }
 
   @Test
   def shouldNotRemoveOngoingTransactionalIds(): Unit = {
     setupAndRunTransactionalIdExpiration(Errors.NONE, Ongoing)
-    verifyMetadataDoesExist(transactionalId1)
-    verifyMetadataDoesExist(transactionalId2)
+    verifyMetadataDoesExistAndIsUsable(transactionalId1)
+    verifyMetadataDoesExistAndIsUsable(transactionalId2)
   }
 
   @Test
   def shouldNotRemovePrepareAbortTransactionalIds(): Unit = {
     setupAndRunTransactionalIdExpiration(Errors.NONE, PrepareAbort)
-    verifyMetadataDoesExist(transactionalId1)
-    verifyMetadataDoesExist(transactionalId2)
+    verifyMetadataDoesExistAndIsUsable(transactionalId1)
+    verifyMetadataDoesExistAndIsUsable(transactionalId2)
   }
 
   @Test
   def shouldNotRemovePrepareCommitTransactionalIds(): Unit = {
     setupAndRunTransactionalIdExpiration(Errors.NONE, PrepareCommit)
-    verifyMetadataDoesExist(transactionalId1)
-    verifyMetadataDoesExist(transactionalId2)
+    verifyMetadataDoesExistAndIsUsable(transactionalId1)
+    verifyMetadataDoesExistAndIsUsable(transactionalId2)
+  }
+
+  @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 verifyMetadataDoesExist(transactionalId: String) = {
+  private def verifyMetadataDoesExistAndIsUsable(transactionalId: String): Unit = {
     transactionManager.getTransactionState(transactionalId) match {
-      case Left(errors) => fail("shouldn't have been any errors")
+      case Left(_) => fail("shouldn't have been any errors")
       case Right(None) => fail("metadata should have been removed")
-      case Right(Some(metadata)) => // ok
+      case Right(Some(metadata)) =>
+        assertTrue("metadata shouldn't be in a pending state", metadata.transactionMetadata.pendingState.isEmpty)
     }
   }
 
-  private def verifyMetadataDoesntExist(transactionalId: String) = {
+  private def verifyMetadataDoesntExist(transactionalId: String): Unit = {
     transactionManager.getTransactionState(transactionalId) match {
-      case Left(errors) => fail("shouldn't have been any errors")
-      case Right(Some(metdata)) => fail("metadata should have been removed")
+      case Left(_) => fail("shouldn't have been any errors")
+      case Right(Some(_)) => fail("metadata should have been removed")
       case Right(None) => // ok
     }
   }
 
-  private def setupAndRunTransactionalIdExpiration(error: Errors, txnState: TransactionState) = {
+  private def setupAndRunTransactionalIdExpiration(error: Errors, txnState: TransactionState): Unit = {
     for (partitionId <- 0 until numPartitions) {
       transactionManager.addLoadedTransactionsToCache(partitionId, 0, new Pool[String, TransactionMetadata]())
     }
@@ -497,20 +649,14 @@ 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]],
           EasyMock.anyObject()
-        )).andAnswer(new IAnswer[Unit] {
-          override def answer(): Unit = {
-            capturedArgument.getValue.apply(
-              Map(partition ->
-                new PartitionResponse(error, 0L, RecordBatch.NO_TIMESTAMP, 0L)
-              )
-            )
-          }
-        })
+        )).andAnswer(() => capturedArgument.getValue.apply(
+          Map(partition -> new PartitionResponse(error, 0L, RecordBatch.NO_TIMESTAMP, 0L)))
+        )
       case _ => // shouldn't append
     }
 
@@ -518,10 +664,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,17 +684,16 @@ 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)
 
     var txnId: String = null
-    def rememberTxnMarkers(transactionalId: String,
-                           coordinatorEpoch: Int,
+    def rememberTxnMarkers(coordinatorEpoch: Int,
                            command: TransactionResult,
                            metadata: TransactionMetadata,
                            newMetadata: TxnTransitMetadata): Unit = {
-      txnId = transactionalId
+      txnId = metadata.transactionalId
     }
 
     transactionManager.loadTransactionsForTxnTopicPartition(partitionId, 0, rememberTxnMarkers)
@@ -573,8 +718,8 @@ class TransactionStateManagerTest {
                             records: MemoryRecords): Unit = {
     EasyMock.reset(replicaManager)
 
-    val logMock =  EasyMock.mock(classOf[Log])
-    val fileRecordsMock = EasyMock.mock(classOf[FileRecords])
+    val logMock: Log = EasyMock.mock(classOf[Log])
+    val fileRecordsMock: FileRecords = EasyMock.mock(classOf[FileRecords])
 
     val endOffset = startOffset + records.records.asScala.size
 
@@ -582,12 +727,23 @@ class TransactionStateManagerTest {
     EasyMock.expect(replicaManager.getLogEndOffset(topicPartition)).andStubReturn(Some(endOffset))
 
     EasyMock.expect(logMock.logStartOffset).andStubReturn(startOffset)
-    EasyMock.expect(logMock.read(EasyMock.eq(startOffset), EasyMock.anyInt(), EasyMock.eq(None),
-      EasyMock.eq(true), EasyMock.eq(IsolationLevel.READ_UNCOMMITTED)))
+    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(fileRecordsMock.readInto(EasyMock.anyObject(classOf[ByteBuffer]), EasyMock.anyInt()))
-      .andReturn(records.buffer)
 
+    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 = {
+        val buffer = bufferCapture.getValue
+        buffer.put(records.buffer.duplicate)
+        buffer.flip()
+      }
+    })
     EasyMock.replay(logMock, fileRecordsMock, replicaManager)
   }
 
@@ -598,22 +754,53 @@ 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]],
       EasyMock.anyObject())
-    ).andAnswer(new IAnswer[Unit] {
-        override def answer(): Unit = capturedArgument.getValue.apply(
-          Map(new TopicPartition(TRANSACTION_STATE_TOPIC_NAME, partitionId) ->
-            new PartitionResponse(error, 0L, RecordBatch.NO_TIMESTAMP, 0L)
-          )
-        )
-      }
+    ).andAnswer(() => capturedArgument.getValue.apply(
+      Map(new TopicPartition(TRANSACTION_STATE_TOPIC_NAME, partitionId) ->
+        new PartitionResponse(error, 0L, RecordBatch.NO_TIMESTAMP, 0L)))
     )
     EasyMock.expect(replicaManager.getMagic(EasyMock.anyObject()))
       .andStubReturn(Some(RecordBatch.MAGIC_VALUE_V1))
 
     EasyMock.replay(replicaManager)
   }
+
+  @Test
+  def testPartitionLoadMetric(): Unit = {
+    val server = ManagementFactory.getPlatformMBeanServer
+    val mBeanName = "kafka.server:type=transaction-coordinator-metrics"
+    val reporter = new JmxReporter
+    val metricsContext = new KafkaMetricsContext("kafka.server")
+    reporter.contextChange(metricsContext)
+    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/AutoOffsetResetTest.scala b/core/src/test/scala/unit/kafka/integration/AutoOffsetResetTest.scala
deleted file mode 100644
index cfc325d361b7a..0000000000000
--- a/core/src/test/scala/unit/kafka/integration/AutoOffsetResetTest.scala
+++ /dev/null
@@ -1,120 +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.integration
-
-import kafka.utils.{ZKGroupTopicDirs, Logging}
-import kafka.consumer.{ConsumerTimeoutException, ConsumerConfig, ConsumerConnector, Consumer}
-import kafka.server._
-import kafka.utils.TestUtils
-import kafka.serializer._
-import kafka.producer.{Producer, KeyedMessage}
-
-import org.junit.{After, Before, Test}
-import org.apache.log4j.{Level, Logger}
-import org.junit.Assert._
-
-@deprecated("This test has been deprecated and it will be removed in a future release", "0.10.0.0")
-class AutoOffsetResetTest extends KafkaServerTestHarness with Logging {
-
-  def generateConfigs = List(KafkaConfig.fromProps(TestUtils.createBrokerConfig(0, zkConnect)))
-
-  val topic = "test_topic"
-  val group = "default_group"
-  val testConsumer = "consumer"
-  val NumMessages = 10
-  val LargeOffset = 10000
-  val SmallOffset = -1
-  
-  val requestHandlerLogger = Logger.getLogger(classOf[kafka.server.KafkaRequestHandler])
-
-  @Before
-  override def setUp() {
-    super.setUp()
-    // temporarily set request handler logger to a higher level
-    requestHandlerLogger.setLevel(Level.FATAL)
-  }
-
-  @After
-  override def tearDown() {
-    // restore set request handler logger to a higher level
-    requestHandlerLogger.setLevel(Level.ERROR)
-    super.tearDown
-  }
-
-  @Test
-  def testResetToEarliestWhenOffsetTooHigh() =
-    assertEquals(NumMessages, resetAndConsume(NumMessages, "smallest", LargeOffset))
-
-  @Test
-  def testResetToEarliestWhenOffsetTooLow() =
-    assertEquals(NumMessages, resetAndConsume(NumMessages, "smallest", SmallOffset))
-
-  @Test
-  def testResetToLatestWhenOffsetTooHigh() =
-    assertEquals(0, resetAndConsume(NumMessages, "largest", LargeOffset))
-
-  @Test
-  def testResetToLatestWhenOffsetTooLow() =
-    assertEquals(0, resetAndConsume(NumMessages, "largest", SmallOffset))
-
-  /* Produce the given number of messages, create a consumer with the given offset policy, 
-   * then reset the offset to the given value and consume until we get no new messages. 
-   * Returns the count of messages received.
-   */
-  def resetAndConsume(numMessages: Int, resetTo: String, offset: Long): Int = {
-    TestUtils.createTopic(zkUtils, topic, 1, 1, servers)
-
-    val producer: Producer[String, Array[Byte]] = TestUtils.createProducer(
-      TestUtils.getBrokerListStrFromServers(servers),
-      keyEncoder = classOf[StringEncoder].getName)
-
-    for(_ <- 0 until numMessages)
-      producer.send(new KeyedMessage[String, Array[Byte]](topic, topic, "test".getBytes))
-
-    // update offset in ZooKeeper for consumer to jump "forward" in time
-    val dirs = new ZKGroupTopicDirs(group, topic)
-    val consumerProps = TestUtils.createConsumerProperties(zkConnect, group, testConsumer)
-    consumerProps.put("auto.offset.reset", resetTo)
-    consumerProps.put("consumer.timeout.ms", "2000")
-    consumerProps.put("fetch.wait.max.ms", "0")
-    val consumerConfig = new ConsumerConfig(consumerProps)
-
-    TestUtils.updateConsumerOffset(consumerConfig, dirs.consumerOffsetDir + "/" + "0", offset)
-    info("Updated consumer offset to " + offset)
-    
-    val consumerConnector: ConsumerConnector = Consumer.create(consumerConfig)
-    val messageStream = consumerConnector.createMessageStreams(Map(topic -> 1))(topic).head
-
-    var received = 0
-    val iter = messageStream.iterator
-    try {
-      for (_ <- 0 until numMessages) {
-        iter.next // will throw a timeout exception if the message isn't there
-        received += 1
-      }
-    } catch {
-      case _: ConsumerTimeoutException =>
-        info("consumer timed out after receiving " + received + " messages.")
-    } finally {
-      producer.close()
-      consumerConnector.shutdown
-    }
-    received
-  }
-  
-}
diff --git a/core/src/test/scala/unit/kafka/integration/FetcherTest.scala b/core/src/test/scala/unit/kafka/integration/FetcherTest.scala
deleted file mode 100644
index f23225cb258cb..0000000000000
--- a/core/src/test/scala/unit/kafka/integration/FetcherTest.scala
+++ /dev/null
@@ -1,93 +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.integration
-
-import java.util.concurrent._
-import java.util.concurrent.atomic._
-import org.junit.{Test, After, Before}
-
-import scala.collection._
-import org.junit.Assert._
-
-import kafka.cluster._
-import kafka.server._
-import kafka.consumer._
-import kafka.utils.TestUtils
-
-@deprecated("This test has been deprecated and will be removed in a future release.", "0.11.0.0")
-class FetcherTest extends KafkaServerTestHarness {
-  val numNodes = 1
-  def generateConfigs = TestUtils.createBrokerConfigs(numNodes, zkConnect).map(KafkaConfig.fromProps)
-
-  val messages = new mutable.HashMap[Int, Seq[Array[Byte]]]
-  val topic = "topic"
-  val queue = new LinkedBlockingQueue[FetchedDataChunk]
-
-  var fetcher: ConsumerFetcherManager = null
-
-  @Before
-  override def setUp() {
-    super.setUp
-    TestUtils.createTopic(zkUtils, topic, partitionReplicaAssignment = Map(0 -> Seq(configs.head.brokerId)), servers = servers)
-
-    val cluster = new Cluster(servers.map { s =>
-      new Broker(s.config.brokerId, "localhost", boundPort(s), listenerName, securityProtocol)
-    })
-
-    fetcher = new ConsumerFetcherManager("consumer1", new ConsumerConfig(TestUtils.createConsumerProperties("", "", "")), zkUtils)
-    fetcher.stopConnections()
-    val topicInfos = configs.map(_ =>
-      new PartitionTopicInfo(topic,
-        0,
-        queue,
-        new AtomicLong(0),
-        new AtomicLong(0),
-        new AtomicInteger(0),
-        ""))
-    fetcher.startConnections(topicInfos, cluster)
-  }
-
-  @After
-  override def tearDown() {
-    fetcher.stopConnections()
-    super.tearDown
-  }
-
-  @Test
-  def testFetcher() {
-    val perNode = 2
-    var count = TestUtils.produceMessages(servers, topic, perNode).size
-
-    fetch(count)
-    assertQueueEmpty()
-    count = TestUtils.produceMessages(servers, topic, perNode).size
-    fetch(count)
-    assertQueueEmpty()
-  }
-
-  def assertQueueEmpty(): Unit = assertEquals(0, queue.size)
-
-  def fetch(expected: Int) {
-    var count = 0
-    while (count < expected) {
-      val chunk = queue.poll(2L, TimeUnit.SECONDS)
-      assertNotNull("Timed out waiting for data chunk " + (count + 1), chunk)
-      count += chunk.messages.size
-    }
-  }
-}
diff --git a/core/src/test/scala/unit/kafka/integration/KafkaServerTestHarness.scala b/core/src/test/scala/unit/kafka/integration/KafkaServerTestHarness.scala
index fed78a5dd58d6..076dd594c5ce0 100755
--- a/core/src/test/scala/unit/kafka/integration/KafkaServerTestHarness.scala
+++ b/core/src/test/scala/unit/kafka/integration/KafkaServerTestHarness.scala
@@ -20,17 +20,20 @@ package kafka.integration
 import java.io.File
 import java.util.Arrays
 
-import kafka.common.KafkaException
 import kafka.server._
 import kafka.utils.TestUtils
 import kafka.zk.ZooKeeperTestHarness
-import org.apache.kafka.common.security.auth.{KafkaPrincipal, SecurityProtocol}
+import org.apache.kafka.common.security.auth.SecurityProtocol
 import org.junit.{After, Before}
 
+import scala.collection.Seq
 import scala.collection.mutable.{ArrayBuffer, Buffer}
 import java.util.Properties
 
+import org.apache.kafka.common.KafkaException
 import org.apache.kafka.common.network.ListenerName
+import org.apache.kafka.common.security.scram.ScramCredential
+import org.apache.kafka.common.utils.Time
 
 /**
  * A test harness that brings up some number of broker nodes
@@ -40,7 +43,6 @@ abstract class KafkaServerTestHarness extends ZooKeeperTestHarness {
   var servers: Buffer[KafkaServer] = new ArrayBuffer
   var brokerList: String = null
   var alive: Array[Boolean] = null
-  val kafkaPrincipalType = KafkaPrincipal.USER_TYPE
 
   /**
    * Implementations must override this method to return a set of KafkaConfigs. This method will be invoked for every
@@ -59,7 +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(): Unit = {}
 
   def configs: Seq[KafkaConfig] = {
     if (instanceConfigs == null)
@@ -76,9 +84,10 @@ abstract class KafkaServerTestHarness extends ZooKeeperTestHarness {
   protected def trustStoreFile: Option[File] = None
   protected def serverSaslProperties: Option[Properties] = None
   protected def clientSaslProperties: Option[Properties] = None
+  protected def brokerTime(brokerId: Int): Time = Time.SYSTEM
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
 
     if (configs.isEmpty)
@@ -90,20 +99,40 @@ abstract class KafkaServerTestHarness extends ZooKeeperTestHarness {
     // Add each broker to `servers` buffer as soon as it is created to ensure that brokers
     // are shutdown cleanly in tearDown even if a subsequent broker fails to start
     for (config <- configs)
-      servers += TestUtils.createServer(config)
+      servers += TestUtils.createServer(config, time = brokerTime(config.brokerId))
     brokerList = TestUtils.bootstrapServers(servers, listenerName)
     alive = new Array[Boolean](servers.length)
     Arrays.fill(alive, true)
+
+    // default implementation is a no-op, it is overridden by subclasses if required
+    configureSecurityAfterServersStart()
   }
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     if (servers != null) {
       TestUtils.shutdownServers(servers)
     }
-    super.tearDown
+    super.tearDown()
   }
 
+  /**
+   * Create a topic in ZooKeeper.
+   * Wait until the leader is elected and the metadata is propagated to all brokers.
+   * Return the leader for each partition.
+   */
+  def createTopic(topic: String, numPartitions: Int = 1, replicationFactor: Int = 1,
+                  topicConfig: Properties = new Properties): scala.collection.immutable.Map[Int, Int] =
+    TestUtils.createTopic(zkClient, topic, numPartitions, replicationFactor, servers, topicConfig)
+
+  /**
+   * Create a topic in ZooKeeper using a customized replica assignment.
+   * Wait until the leader is elected and the metadata is propagated to all brokers.
+   * Return the leader for each partition.
+   */
+  def createTopic(topic: String, partitionReplicaAssignment: collection.Map[Int, Seq[Int]]): scala.collection.immutable.Map[Int, Int] =
+    TestUtils.createTopic(zkClient, topic, partitionReplicaAssignment, servers)
+
   /**
    * Pick a broker at random and kill it if it isn't already dead
    * Return the id of the broker killed
@@ -114,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()
@@ -125,10 +154,18 @@ 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
     }
   }
+
+  def waitForUserScramCredentialToAppearOnAllBrokers(clientPrincipal: String, mechanismName: String): Unit = {
+    servers.foreach { server =>
+      val cache = server.credentialProvider.credentialCache.cache(mechanismName, classOf[ScramCredential])
+      TestUtils.waitUntilTrue(() => cache.get(clientPrincipal) != null, s"SCRAM credentials not created for $clientPrincipal")
+    }
+  }
+
 }
diff --git a/core/src/test/scala/unit/kafka/integration/MetricsDuringTopicCreationDeletionTest.scala b/core/src/test/scala/unit/kafka/integration/MetricsDuringTopicCreationDeletionTest.scala
index daa427683dc10..003a2c4b48a1b 100644
--- a/core/src/test/scala/unit/kafka/integration/MetricsDuringTopicCreationDeletionTest.scala
+++ b/core/src/test/scala/unit/kafka/integration/MetricsDuringTopicCreationDeletionTest.scala
@@ -18,13 +18,15 @@
 package kafka.integration
 
 import java.util.Properties
+
 import kafka.server.KafkaConfig
 import kafka.utils.{Logging, TestUtils}
-import scala.collection.JavaConverters.mapAsScalaMapConverter
 
+import scala.jdk.CollectionConverters._
+import org.scalatest.Assertions.fail
 import org.junit.{Before, Test}
-import com.yammer.metrics.Metrics
 import com.yammer.metrics.core.Gauge
+import kafka.metrics.KafkaYammerMetrics
 
 class MetricsDuringTopicCreationDeletionTest extends KafkaServerTestHarness with Logging {
 
@@ -34,43 +36,43 @@ class MetricsDuringTopicCreationDeletionTest extends KafkaServerTestHarness with
   private val replicationFactor = 3
   private val partitionNum = 3
   private val createDeleteIterations = 3
-  
+
   private val overridingProps = new Properties
   overridingProps.put(KafkaConfig.DeleteTopicEnableProp, "true")
   overridingProps.put(KafkaConfig.AutoCreateTopicsEnableProp, "false")
-  // speed up the test for UnderReplicatedPartitions 
+  // speed up the test for UnderReplicatedPartitions
   // which relies on the ISR expiry thread to execute concurrently with topic creation
-  overridingProps.put(KafkaConfig.ReplicaLagTimeMaxMsProp, "2000") 
+  overridingProps.put(KafkaConfig.ReplicaLagTimeMaxMsProp, "2000")
 
   private val testedMetrics = List("OfflinePartitionsCount","PreferredReplicaImbalanceCount","UnderReplicatedPartitions")
   private val topics = List.tabulate(topicNum) (n => topicName + n)
 
   @volatile private var running = true
-  
+
   override def generateConfigs = TestUtils.createBrokerConfigs(nodesNum, zkConnect)
     .map(KafkaConfig.fromProps(_, overridingProps))
 
   @Before
-  override def setUp {
-    // Do some Metrics Registry cleanup by removing the metrics that this test checks. 
+  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
     for (m <- testedMetrics) {
-        val metricName = Metrics.defaultRegistry.allMetrics.asScala.keys.find(_.getName.endsWith(m))
-        metricName.foreach(Metrics.defaultRegistry.removeMetric)
+        val metricName = KafkaYammerMetrics.defaultRegistry.allMetrics.asScala.keys.find(_.getName.endsWith(m))
+        metricName.foreach(KafkaYammerMetrics.defaultRegistry.removeMetric)
     }
-    
-    super.setUp
+
+    super.setUp()
   }
 
   /*
    * 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 
+    // we can't access the metrics value of each server. So instead we directly invoke the method
     // replicaManager.underReplicatedPartitionCount() that defines the metrics value.
     @volatile var underReplicatedPartitionCount = 0
 
@@ -86,25 +88,23 @@ class MetricsDuringTopicCreationDeletionTest extends KafkaServerTestHarness with
 
     // Thread checking the metric continuously
     running = true
-    val thread = new Thread(new Runnable {
-      def run() {
-        while (running) {
-          for ( s <- servers if running) {
-            underReplicatedPartitionCount = s.replicaManager.underReplicatedPartitionCount
-            if (underReplicatedPartitionCount > 0) {
-              running = false
-            }
+    val thread = new Thread(() => {
+      while (running) {
+        for (s <- servers if running) {
+          underReplicatedPartitionCount = s.replicaManager.underReplicatedPartitionCount
+          if (underReplicatedPartitionCount > 0) {
+            running = false
           }
+        }
 
-          preferredReplicaImbalanceCount = preferredReplicaImbalanceCountGauge.value
-          if (preferredReplicaImbalanceCount > 0) {
-             running = false
-          }
+        preferredReplicaImbalanceCount = preferredReplicaImbalanceCountGauge.value
+        if (preferredReplicaImbalanceCount > 0) {
+          running = false
+        }
 
-          offlinePartitionsCount = offlinePartitionsCountGauge.value
-          if (offlinePartitionsCount > 0) {
-             running = false
-          }
+        offlinePartitionsCount = offlinePartitionsCountGauge.value
+        if (offlinePartitionsCount > 0) {
+          running = false
         }
       }
     })
@@ -116,41 +116,40 @@ class MetricsDuringTopicCreationDeletionTest extends KafkaServerTestHarness with
     // if the thread checking the gauge is still run, stop it
     running = false;
     thread.join
-    
+
     assert(offlinePartitionsCount==0, "OfflinePartitionCount not 0: "+ offlinePartitionsCount)
     assert(preferredReplicaImbalanceCount==0, "PreferredReplicaImbalanceCount not 0: " + preferredReplicaImbalanceCount)
     assert(underReplicatedPartitionCount==0, "UnderReplicatedPartitionCount not 0: " + underReplicatedPartitionCount)
   }
 
   private def getGauge(metricName: String) = {
-    Metrics.defaultRegistry.allMetrics.asScala
-           .filterKeys(k => k.getName.endsWith(metricName))
-           .headOption
-           .getOrElse { fail( "Unable to find metric " + metricName ) }
-           ._2.asInstanceOf[Gauge[Int]]
+    KafkaYammerMetrics.defaultRegistry.allMetrics.asScala
+                      .filter { case (k, _) => k.getName.endsWith(metricName) }
+                      .headOption
+                      .getOrElse { fail( "Unable to find metric " + metricName ) }
+                      ._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) {
         try {
-          adminZkClient.createTopic(t, partitionNum, replicationFactor)
+          createTopic(t, partitionNum, replicationFactor)
         } catch {
           case e: Exception => e.printStackTrace
         }
       }
-      Thread.sleep(500)
 
       // Delete topics
       for (t <- topics if running) {
           try {
-              adminZkClient.deleteTopic(t)
+            adminZkClient.deleteTopic(t)
+            TestUtils.verifyTopicDeletion(zkClient, t, partitionNum, servers)
           } catch {
           case e: Exception => e.printStackTrace
           }
       }
-      Thread.sleep(500)
     }
   }
 }
diff --git a/core/src/test/scala/unit/kafka/integration/MinIsrConfigTest.scala b/core/src/test/scala/unit/kafka/integration/MinIsrConfigTest.scala
index 455bbde58ad1a..a1b2386163af7 100644
--- a/core/src/test/scala/unit/kafka/integration/MinIsrConfigTest.scala
+++ b/core/src/test/scala/unit/kafka/integration/MinIsrConfigTest.scala
@@ -30,8 +30,8 @@ class MinIsrConfigTest extends KafkaServerTestHarness {
   def generateConfigs = TestUtils.createBrokerConfigs(1, zkConnect).map(KafkaConfig.fromProps(_, overridingProps))
 
   @Test
-  def testDefaultKafkaConfig() {
-    assert(servers.head.getLogManager().defaultConfig.minInSyncReplicas == 5)
+  def testDefaultKafkaConfig(): Unit = {
+    assert(servers.head.getLogManager.initialDefaultConfig.minInSyncReplicas == 5)
   }
 
 }
diff --git a/core/src/test/scala/unit/kafka/integration/PrimitiveApiTest.scala b/core/src/test/scala/unit/kafka/integration/PrimitiveApiTest.scala
deleted file mode 100755
index 9c4d08f7d8e54..0000000000000
--- a/core/src/test/scala/unit/kafka/integration/PrimitiveApiTest.scala
+++ /dev/null
@@ -1,276 +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.integration
-
-import java.nio.ByteBuffer
-
-import org.junit.Assert._
-import kafka.api.{FetchRequest, FetchRequestBuilder, PartitionFetchInfo}
-import kafka.server.{KafkaConfig, KafkaRequestHandler}
-import kafka.producer.{KeyedMessage, Producer}
-import org.apache.log4j.{Level, Logger}
-import org.junit.Test
-
-import scala.collection._
-import kafka.common.{ErrorMapping, OffsetOutOfRangeException, TopicAndPartition, UnknownTopicOrPartitionException}
-import kafka.utils.{StaticPartitioner, TestUtils}
-import kafka.serializer.StringEncoder
-import java.util.Properties
-
-import org.apache.kafka.common.TopicPartition
-
-/**
- * End to end tests of the primitive apis against a local server
- */
-@deprecated("This test has been deprecated and it will be removed in a future release", "0.10.0.0")
-class PrimitiveApiTest extends ProducerConsumerTestHarness {
-  val requestHandlerLogger = Logger.getLogger(classOf[KafkaRequestHandler])
-
-  def generateConfigs = List(KafkaConfig.fromProps(TestUtils.createBrokerConfig(0, zkConnect)))
-
-  @Test
-  def testFetchRequestCanProperlySerialize() {
-    val request = new FetchRequestBuilder()
-      .clientId("test-client")
-      .maxWait(10001)
-      .minBytes(4444)
-      .addFetch("topic1", 0, 0, 10000)
-      .addFetch("topic2", 1, 1024, 9999)
-      .addFetch("topic1", 1, 256, 444)
-      .build()
-    val serializedBuffer = ByteBuffer.allocate(request.sizeInBytes)
-    request.writeTo(serializedBuffer)
-    serializedBuffer.rewind()
-    val deserializedRequest = FetchRequest.readFrom(serializedBuffer)
-    assertEquals(request, deserializedRequest)
-  }
-
-  @Test
-  def testEmptyFetchRequest() {
-    val partitionRequests = immutable.Map[TopicAndPartition, PartitionFetchInfo]()
-    val request = new FetchRequest(requestInfo = partitionRequests.toVector)
-    val fetched = consumer.fetch(request)
-    assertTrue(!fetched.hasError && fetched.data.isEmpty)
-  }
-
-  @Test
-  def testDefaultEncoderProducerAndFetch() {
-    val topic = "test-topic"
-
-    producer.send(new KeyedMessage[String, String](topic, "test-message"))
-
-    val replica = servers.head.replicaManager.getReplica(new TopicPartition(topic, 0)).get
-    assertTrue("HighWatermark should equal logEndOffset with just 1 replica",
-               replica.logEndOffset.messageOffset > 0 && replica.logEndOffset.equals(replica.highWatermark))
-
-    val request = new FetchRequestBuilder()
-      .clientId("test-client")
-      .addFetch(topic, 0, 0, 10000)
-      .build()
-    val fetched = consumer.fetch(request)
-    assertEquals("Returned correlationId doesn't match that in request.", 0, fetched.correlationId)
-
-    val messageSet = fetched.messageSet(topic, 0)
-    assertTrue(messageSet.iterator.hasNext)
-
-    val fetchedMessageAndOffset = messageSet.head
-    assertEquals("test-message", TestUtils.readString(fetchedMessageAndOffset.message.payload, "UTF-8"))
-  }
-
-  @Test
-  def testDefaultEncoderProducerAndFetchWithCompression() {
-    val topic = "test-topic"
-    val props = new Properties()
-    props.put("compression.codec", "gzip")
-
-    val stringProducer1 = TestUtils.createProducer[String, String](
-      TestUtils.getBrokerListStrFromServers(servers),
-      encoder = classOf[StringEncoder].getName,
-      keyEncoder = classOf[StringEncoder].getName,
-      partitioner = classOf[StaticPartitioner].getName,
-      producerProps = props)
-
-    stringProducer1.send(new KeyedMessage[String, String](topic, "test-message"))
-
-    val fetched = consumer.fetch(new FetchRequestBuilder().addFetch(topic, 0, 0, 10000).build())
-    val messageSet = fetched.messageSet(topic, 0)
-    assertTrue(messageSet.iterator.hasNext)
-
-    val fetchedMessageAndOffset = messageSet.head
-    assertEquals("test-message", TestUtils.readString(fetchedMessageAndOffset.message.payload, "UTF-8"))
-  }
-
-  private def produceAndMultiFetch(producer: Producer[String, String]) {
-    for(topic <- List("test1", "test2", "test3", "test4"))
-      TestUtils.createTopic(zkUtils, topic, servers = servers)
-
-    // send some messages
-    val topics = List(("test4", 0), ("test1", 0), ("test2", 0), ("test3", 0));
-    {
-      val messages = new mutable.HashMap[String, Seq[String]]
-      val builder = new FetchRequestBuilder()
-      for( (topic, partition) <- topics) {
-        val messageList = List("a_" + topic, "b_" + topic)
-        val producerData = messageList.map(new KeyedMessage[String, String](topic, topic, _))
-        messages += topic -> messageList
-        producer.send(producerData:_*)
-        builder.addFetch(topic, partition, 0, 10000)
-      }
-
-      val request = builder.build()
-      val response = consumer.fetch(request)
-      for((topic, partition) <- topics) {
-        val fetched = response.messageSet(topic, partition)
-        assertEquals(messages(topic), fetched.map(messageAndOffset => TestUtils.readString(messageAndOffset.message.payload)))
-      }
-    }
-
-    // temporarily set request handler logger to a higher level
-    requestHandlerLogger.setLevel(Level.FATAL)
-
-    {
-      // send some invalid offsets
-      val builder = new FetchRequestBuilder()
-      for((topic, partition) <- topics)
-        builder.addFetch(topic, partition, -1, 10000)
-
-      try {
-        val request = builder.build()
-        val response = consumer.fetch(request)
-        response.data.foreach(pdata => ErrorMapping.maybeThrowException(pdata._2.error.code))
-        fail("Expected exception when fetching message with invalid offset")
-      } catch {
-        case _: OffsetOutOfRangeException => // This is good.
-      }
-    }
-
-    {
-      // send some invalid partitions
-      val builder = new FetchRequestBuilder()
-      for((topic, _) <- topics)
-        builder.addFetch(topic, -1, 0, 10000)
-
-      try {
-        val request = builder.build()
-        val response = consumer.fetch(request)
-        response.data.foreach(pdata => ErrorMapping.maybeThrowException(pdata._2.error.code))
-        fail("Expected exception when fetching message with invalid partition")
-      } catch {
-        case _: UnknownTopicOrPartitionException => // This is good.
-      }
-    }
-
-    // restore set request handler logger to a higher level
-    requestHandlerLogger.setLevel(Level.ERROR)
-  }
-
-  @Test
-  def testProduceAndMultiFetch() {
-    produceAndMultiFetch(producer)
-  }
-
-  private def multiProduce(producer: Producer[String, String]) {
-    val topics = Map("test4" -> 0, "test1" -> 0, "test2" -> 0, "test3" -> 0)
-    topics.keys.map(topic => TestUtils.createTopic(zkUtils, topic, servers = servers))
-
-    val messages = new mutable.HashMap[String, Seq[String]]
-    val builder = new FetchRequestBuilder()
-    for((topic, partition) <- topics) {
-      val messageList = List("a_" + topic, "b_" + topic)
-      val producerData = messageList.map(new KeyedMessage[String, String](topic, topic, _))
-      messages += topic -> messageList
-      producer.send(producerData:_*)
-      builder.addFetch(topic, partition, 0, 10000)
-    }
-
-    val request = builder.build()
-    val response = consumer.fetch(request)
-    for((topic, partition) <- topics) {
-      val fetched = response.messageSet(topic, partition)
-      assertEquals(messages(topic), fetched.map(messageAndOffset => TestUtils.readString(messageAndOffset.message.payload)))
-    }
-  }
-
-  @Test
-  def testMultiProduce() {
-    multiProduce(producer)
-  }
-
-  @Test
-  def testConsumerEmptyTopic() {
-    val newTopic = "new-topic"
-    TestUtils.createTopic(zkUtils, newTopic, numPartitions = 1, replicationFactor = 1, servers = servers)
-
-    val fetchResponse = consumer.fetch(new FetchRequestBuilder().addFetch(newTopic, 0, 0, 10000).build())
-    assertFalse(fetchResponse.messageSet(newTopic, 0).iterator.hasNext)
-  }
-
-  @Test
-  def testPipelinedProduceRequests() {
-    val topics = Map("test4" -> 0, "test1" -> 0, "test2" -> 0, "test3" -> 0)
-    topics.keys.map(topic => TestUtils.createTopic(zkUtils, topic, servers = servers))
-    val props = new Properties()
-    props.put("request.required.acks", "0")
-    val pipelinedProducer: Producer[String, String] =
-      TestUtils.createProducer[String, String](
-        TestUtils.getBrokerListStrFromServers(servers),
-        encoder = classOf[StringEncoder].getName,
-        keyEncoder = classOf[StringEncoder].getName,
-        partitioner = classOf[StaticPartitioner].getName,
-        producerProps = props)
-
-    // send some messages
-    val messages = new mutable.HashMap[String, Seq[String]]
-    val builder = new FetchRequestBuilder()
-    for( (topic, partition) <- topics) {
-      val messageList = List("a_" + topic, "b_" + topic)
-      val producerData = messageList.map(new KeyedMessage[String, String](topic, topic, _))
-      messages += topic -> messageList
-      pipelinedProducer.send(producerData:_*)
-      builder.addFetch(topic, partition, 0, 10000)
-    }
-
-    // wait until the messages are published
-    TestUtils.waitUntilTrue(() => { servers.head.logManager.getLog(new TopicPartition("test1", 0)).get.logEndOffset == 2 },
-                            "Published messages should be in the log")
-    TestUtils.waitUntilTrue(() => { servers.head.logManager.getLog(new TopicPartition("test2", 0)).get.logEndOffset == 2 },
-                            "Published messages should be in the log")
-    TestUtils.waitUntilTrue(() => { servers.head.logManager.getLog(new TopicPartition("test3", 0)).get.logEndOffset == 2 },
-                            "Published messages should be in the log")
-    TestUtils.waitUntilTrue(() => { servers.head.logManager.getLog(new TopicPartition("test4", 0)).get.logEndOffset == 2 },
-                            "Published messages should be in the log")
-
-    val replicaId = servers.head.config.brokerId
-    TestUtils.waitUntilTrue(() => { servers.head.replicaManager.getReplica(new TopicPartition("test1", 0), replicaId).get.highWatermark.messageOffset == 2 },
-                            "High watermark should equal to log end offset")
-    TestUtils.waitUntilTrue(() => { servers.head.replicaManager.getReplica(new TopicPartition("test2", 0), replicaId).get.highWatermark.messageOffset == 2 },
-                            "High watermark should equal to log end offset")
-    TestUtils.waitUntilTrue(() => { servers.head.replicaManager.getReplica(new TopicPartition("test3", 0), replicaId).get.highWatermark.messageOffset == 2 },
-                            "High watermark should equal to log end offset")
-    TestUtils.waitUntilTrue(() => { servers.head.replicaManager.getReplica(new TopicPartition("test4", 0), replicaId).get.highWatermark.messageOffset == 2 },
-                            "High watermark should equal to log end offset")
-
-    // test if the consumer received the messages in the correct order when producer has enabled request pipelining
-    val request = builder.build()
-    val response = consumer.fetch(request)
-    for( (topic, partition) <- topics) {
-      val fetched = response.messageSet(topic, partition)
-      assertEquals(messages(topic), fetched.map(messageAndOffset => TestUtils.readString(messageAndOffset.message.payload)))
-    }
-  }
-}
diff --git a/core/src/test/scala/unit/kafka/integration/ProducerConsumerTestHarness.scala b/core/src/test/scala/unit/kafka/integration/ProducerConsumerTestHarness.scala
deleted file mode 100644
index e3115e1bd2dcd..0000000000000
--- a/core/src/test/scala/unit/kafka/integration/ProducerConsumerTestHarness.scala
+++ /dev/null
@@ -1,48 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package kafka.integration
-
-import kafka.consumer.SimpleConsumer
-import org.junit.{After, Before}
-import kafka.producer.Producer
-import kafka.utils.{StaticPartitioner, TestUtils}
-import kafka.serializer.StringEncoder
-
-@deprecated("This test has been deprecated and it will be removed in a future release", "0.10.0.0")
-trait ProducerConsumerTestHarness extends KafkaServerTestHarness {
-  val host = "localhost"
-  var producer: Producer[String, String] = null
-  var consumer: SimpleConsumer = null
-
-  @Before
-  override def setUp() {
-    super.setUp
-    producer = TestUtils.createProducer[String, String](TestUtils.getBrokerListStrFromServers(servers),
-      encoder = classOf[StringEncoder].getName,
-      keyEncoder = classOf[StringEncoder].getName,
-      partitioner = classOf[StaticPartitioner].getName)
-    consumer = new SimpleConsumer(host, TestUtils.boundPort(servers.head), 1000000, 64 * 1024, "")
-  }
-
-  @After
-  override def tearDown() {
-    producer.close()
-    consumer.close()
-    super.tearDown
-  }
-}
diff --git a/core/src/test/scala/unit/kafka/integration/TopicMetadataTest.scala b/core/src/test/scala/unit/kafka/integration/TopicMetadataTest.scala
deleted file mode 100644
index e93dcf6ccaf55..0000000000000
--- a/core/src/test/scala/unit/kafka/integration/TopicMetadataTest.scala
+++ /dev/null
@@ -1,289 +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.integration
-
-import kafka.api.TopicMetadataResponse
-import kafka.client.ClientUtils
-import kafka.cluster.BrokerEndPoint
-import kafka.server.{KafkaConfig, KafkaServer, NotRunning}
-import kafka.utils.TestUtils
-import kafka.utils.TestUtils._
-import kafka.zk.ZooKeeperTestHarness
-import org.apache.kafka.common.protocol.Errors
-import org.junit.Assert._
-import org.junit.{Test, After, Before}
-
-class TopicMetadataTest extends ZooKeeperTestHarness {
-  private var server1: KafkaServer = null
-  private var adHocServers: Seq[KafkaServer] = Seq()
-  var brokerEndPoints: Seq[BrokerEndPoint] = null
-  var adHocConfigs: Seq[KafkaConfig] = null
-  val numConfigs: Int = 4
-
-  @Before
-  override def setUp() {
-    super.setUp()
-    val props = createBrokerConfigs(numConfigs, zkConnect)
-    val configs: Seq[KafkaConfig] = props.map(KafkaConfig.fromProps)
-    adHocConfigs = configs.takeRight(configs.size - 1) // Started and stopped by individual test cases
-    server1 = TestUtils.createServer(configs.head)
-    brokerEndPoints = Seq(
-      // We are using the Scala clients and they don't support SSL. Once we move to the Java ones, we should use
-      // `securityProtocol` instead of PLAINTEXT below
-      new BrokerEndPoint(server1.config.brokerId, server1.config.hostName, TestUtils.boundPort(server1))
-    )
-  }
-
-  @After
-  override def tearDown() {
-    TestUtils.shutdownServers(adHocServers :+ server1)
-    super.tearDown()
-  }
-
-  @Test
-  def testBasicTopicMetadata(): Unit = {
-    // create topic
-    val topic = "test"
-    createTopic(zkUtils, topic, numPartitions = 1, replicationFactor = 1, servers = Seq(server1))
-
-    val topicsMetadata = ClientUtils.fetchTopicMetadata(Set(topic), brokerEndPoints, "TopicMetadataTest-testBasicTopicMetadata",
-      2000, 0).topicsMetadata
-    assertEquals(Errors.NONE, topicsMetadata.head.error)
-    assertEquals(Errors.NONE, topicsMetadata.head.partitionsMetadata.head.error)
-    assertEquals("Expecting metadata only for 1 topic", 1, topicsMetadata.size)
-    assertEquals("Expecting metadata for the test topic", "test", topicsMetadata.head.topic)
-    val partitionMetadata = topicsMetadata.head.partitionsMetadata
-    assertEquals("Expecting metadata for 1 partition", 1, partitionMetadata.size)
-    assertEquals("Expecting partition id to be 0", 0, partitionMetadata.head.partitionId)
-    assertEquals(1, partitionMetadata.head.replicas.size)
-  }
-
-  @Test
-  def testGetAllTopicMetadata(): Unit = {
-    // create topic
-    val topic1 = "testGetAllTopicMetadata1"
-    val topic2 = "testGetAllTopicMetadata2"
-    createTopic(zkUtils, topic1, numPartitions = 1, replicationFactor = 1, servers = Seq(server1))
-    createTopic(zkUtils, topic2, numPartitions = 1, replicationFactor = 1, servers = Seq(server1))
-
-    // issue metadata request with empty list of topics
-    val topicsMetadata = ClientUtils.fetchTopicMetadata(Set.empty, brokerEndPoints, "TopicMetadataTest-testGetAllTopicMetadata",
-      2000, 0).topicsMetadata
-    assertEquals(Errors.NONE, topicsMetadata.head.error)
-    assertEquals(2, topicsMetadata.size)
-    assertEquals(Errors.NONE, topicsMetadata.head.partitionsMetadata.head.error)
-    assertEquals(Errors.NONE, topicsMetadata.last.partitionsMetadata.head.error)
-    val partitionMetadataTopic1 = topicsMetadata.head.partitionsMetadata
-    val partitionMetadataTopic2 = topicsMetadata.last.partitionsMetadata
-    assertEquals("Expecting metadata for 1 partition", 1, partitionMetadataTopic1.size)
-    assertEquals("Expecting partition id to be 0", 0, partitionMetadataTopic1.head.partitionId)
-    assertEquals(1, partitionMetadataTopic1.head.replicas.size)
-    assertEquals("Expecting metadata for 1 partition", 1, partitionMetadataTopic2.size)
-    assertEquals("Expecting partition id to be 0", 0, partitionMetadataTopic2.head.partitionId)
-    assertEquals(1, partitionMetadataTopic2.head.replicas.size)
-  }
-
-  @Test
-  def testAutoCreateTopic(): Unit = {
-    // auto create topic
-    val topic = "testAutoCreateTopic"
-    var topicsMetadata = ClientUtils.fetchTopicMetadata(Set(topic), brokerEndPoints, "TopicMetadataTest-testAutoCreateTopic",
-      2000,0).topicsMetadata
-    assertEquals(Errors.LEADER_NOT_AVAILABLE, topicsMetadata.head.error)
-    assertEquals("Expecting metadata only for 1 topic", 1, topicsMetadata.size)
-    assertEquals("Expecting metadata for the test topic", topic, topicsMetadata.head.topic)
-    assertEquals(0, topicsMetadata.head.partitionsMetadata.size)
-
-    // wait for leader to be elected
-    TestUtils.waitUntilLeaderIsElectedOrChanged(zkUtils, topic, 0)
-    TestUtils.waitUntilMetadataIsPropagated(Seq(server1), topic, 0)
-
-    // retry the metadata for the auto created topic
-    topicsMetadata = ClientUtils.fetchTopicMetadata(Set(topic), brokerEndPoints, "TopicMetadataTest-testBasicTopicMetadata",
-      2000,0).topicsMetadata
-    assertEquals(Errors.NONE, topicsMetadata.head.error)
-    assertEquals(Errors.NONE, topicsMetadata.head.partitionsMetadata.head.error)
-    val partitionMetadata = topicsMetadata.head.partitionsMetadata
-    assertEquals("Expecting metadata for 1 partition", 1, partitionMetadata.size)
-    assertEquals("Expecting partition id to be 0", 0, partitionMetadata.head.partitionId)
-    assertEquals(1, partitionMetadata.head.replicas.size)
-    assertTrue(partitionMetadata.head.leader.isDefined)
-  }
-
-  @Test
-  def testAutoCreateTopicWithInvalidReplication(): Unit = {
-    val adHocProps = createBrokerConfig(2, zkConnect)
-    // Set default replication higher than the number of live brokers
-    adHocProps.setProperty(KafkaConfig.DefaultReplicationFactorProp, "3")
-    // start adHoc brokers with replication factor too high
-    val adHocServer = createServer(new KafkaConfig(adHocProps))
-    adHocServers = Seq(adHocServer)
-    // We are using the Scala clients and they don't support SSL. Once we move to the Java ones, we should use
-    // `securityProtocol` instead of PLAINTEXT below
-    val adHocEndpoint = new BrokerEndPoint(adHocServer.config.brokerId, adHocServer.config.hostName,
-      TestUtils.boundPort(adHocServer))
-
-    // auto create topic on "bad" endpoint
-    val topic = "testAutoCreateTopic"
-    val topicsMetadata = ClientUtils.fetchTopicMetadata(Set(topic), Seq(adHocEndpoint), "TopicMetadataTest-testAutoCreateTopic",
-      2000, 0).topicsMetadata
-    assertEquals(Errors.INVALID_REPLICATION_FACTOR, topicsMetadata.head.error)
-    assertEquals("Expecting metadata only for 1 topic", 1, topicsMetadata.size)
-    assertEquals("Expecting metadata for the test topic", topic, topicsMetadata.head.topic)
-    assertEquals(0, topicsMetadata.head.partitionsMetadata.size)
-  }
-
-  @Test
-  def testAutoCreateTopicWithCollision(): Unit = {
-    // auto create topic
-    val topic1 = "testAutoCreate_Topic"
-    val topic2 = "testAutoCreate.Topic"
-    var topicsMetadata = ClientUtils.fetchTopicMetadata(Set(topic1, topic2), brokerEndPoints, "TopicMetadataTest-testAutoCreateTopic",
-      2000, 0).topicsMetadata
-    assertEquals("Expecting metadata for 2 topics", 2, topicsMetadata.size)
-    assertEquals("Expecting metadata for topic1", topic1, topicsMetadata.head.topic)
-    assertEquals(Errors.LEADER_NOT_AVAILABLE, topicsMetadata.head.error)
-    assertEquals("Expecting metadata for topic2", topic2, topicsMetadata(1).topic)
-    assertEquals("Expecting InvalidTopicCode for topic2 metadata", Errors.INVALID_TOPIC_EXCEPTION, topicsMetadata(1).error)
-
-    // wait for leader to be elected
-    TestUtils.waitUntilLeaderIsElectedOrChanged(zkUtils, topic1, 0)
-    TestUtils.waitUntilMetadataIsPropagated(Seq(server1), topic1, 0)
-
-    // retry the metadata for the first auto created topic
-    topicsMetadata = ClientUtils.fetchTopicMetadata(Set(topic1), brokerEndPoints, "TopicMetadataTest-testBasicTopicMetadata",
-      2000, 0).topicsMetadata
-    assertEquals(Errors.NONE, topicsMetadata.head.error)
-    assertEquals(Errors.NONE, topicsMetadata.head.partitionsMetadata.head.error)
-    val partitionMetadata = topicsMetadata.head.partitionsMetadata
-    assertEquals("Expecting metadata for 1 partition", 1, partitionMetadata.size)
-    assertEquals("Expecting partition id to be 0", 0, partitionMetadata.head.partitionId)
-    assertEquals(1, partitionMetadata.head.replicas.size)
-    assertTrue(partitionMetadata.head.leader.isDefined)
-  }
-
-  private def checkIsr(servers: Seq[KafkaServer]): Unit = {
-    val activeBrokers: Seq[KafkaServer] = servers.filter(x => x.brokerState.currentState != NotRunning.state)
-    val expectedIsr: Seq[BrokerEndPoint] = activeBrokers.map { x =>
-      new BrokerEndPoint(x.config.brokerId,
-        if (x.config.hostName.nonEmpty) x.config.hostName else "localhost",
-        TestUtils.boundPort(x))
-    }
-
-    // Assert that topic metadata at new brokers is updated correctly
-    activeBrokers.foreach(x => {
-      var metadata: TopicMetadataResponse = new TopicMetadataResponse(Seq(), Seq(), -1)
-      waitUntilTrue(() => {
-        metadata = ClientUtils.fetchTopicMetadata(Set.empty,
-                                Seq(new BrokerEndPoint(x.config.brokerId,
-                                                       if (x.config.hostName.nonEmpty) x.config.hostName else "localhost",
-                                                       TestUtils.boundPort(x))),
-                                "TopicMetadataTest-testBasicTopicMetadata", 2000, 0)
-        metadata.topicsMetadata.nonEmpty &&
-          metadata.topicsMetadata.head.partitionsMetadata.nonEmpty &&
-          expectedIsr.sortBy(_.id) == metadata.topicsMetadata.head.partitionsMetadata.head.isr.sortBy(_.id)
-      },
-        "Topic metadata is not correctly updated for broker " + x + ".\n" +
-        "Expected ISR: " + expectedIsr + "\n" +
-        "Actual ISR  : " + (if (metadata.topicsMetadata.nonEmpty &&
-                                metadata.topicsMetadata.head.partitionsMetadata.nonEmpty)
-                              metadata.topicsMetadata.head.partitionsMetadata.head.isr
-                            else
-                              ""), 8000L)
-    })
-  }
-
-  @Test
-  def testIsrAfterBrokerShutDownAndJoinsBack(): Unit = {
-    val numBrokers = 2 //just 2 brokers are enough for the test
-
-    // start adHoc brokers
-    adHocServers = adHocConfigs.take(numBrokers - 1).map(p => createServer(p))
-    val allServers: Seq[KafkaServer] = Seq(server1) ++ adHocServers
-
-    // create topic
-    val topic: String = "test"
-    adminZkClient.createTopic(topic, 1, numBrokers)
-
-    // shutdown a broker
-    adHocServers.last.shutdown()
-    adHocServers.last.awaitShutdown()
-
-    // startup a broker
-    adHocServers.last.startup()
-
-    // check metadata is still correct and updated at all brokers
-    checkIsr(allServers)
-  }
-
-  private def checkMetadata(servers: Seq[KafkaServer], expectedBrokersCount: Int): Unit = {
-    var topicMetadata: TopicMetadataResponse = new TopicMetadataResponse(Seq(), Seq(), -1)
-
-    // Get topic metadata from old broker
-    // Wait for metadata to get updated by checking metadata from a new broker
-    waitUntilTrue(() => {
-    topicMetadata = ClientUtils.fetchTopicMetadata(
-      Set.empty, brokerEndPoints, "TopicMetadataTest-testBasicTopicMetadata", 2000, 0)
-    topicMetadata.brokers.size == expectedBrokersCount},
-      "Alive brokers list is not correctly propagated by coordinator to brokers"
-    )
-
-    // Assert that topic metadata at new brokers is updated correctly
-    servers.filter(x => x.brokerState.currentState != NotRunning.state).foreach(x =>
-      waitUntilTrue(() => {
-          val foundMetadata = ClientUtils.fetchTopicMetadata(
-            Set.empty,
-            Seq(new BrokerEndPoint(x.config.brokerId, x.config.hostName, TestUtils.boundPort(x))),
-            "TopicMetadataTest-testBasicTopicMetadata", 2000, 0)
-          topicMetadata.brokers.sortBy(_.id) == foundMetadata.brokers.sortBy(_.id) &&
-            topicMetadata.topicsMetadata.sortBy(_.topic) == foundMetadata.topicsMetadata.sortBy(_.topic)
-        },
-        s"Topic metadata is not correctly updated"))
-  }
-
-  @Test
-  def testAliveBrokerListWithNoTopics(): Unit = {
-    checkMetadata(Seq(server1), 1)
-  }
-
-  @Test
-  def testAliveBrokersListWithNoTopicsAfterNewBrokerStartup(): Unit = {
-    adHocServers = adHocConfigs.takeRight(adHocConfigs.size - 1).map(p => createServer(p))
-
-    checkMetadata(adHocServers, numConfigs - 1)
-
-    // Add a broker
-    adHocServers = adHocServers ++ Seq(createServer(adHocConfigs.head))
-
-    checkMetadata(adHocServers, numConfigs)
-  }
-
-
-  @Test
-  def testAliveBrokersListWithNoTopicsAfterABrokerShutdown(): Unit = {
-    adHocServers = adHocConfigs.map(p => createServer(p))
-
-    checkMetadata(adHocServers, numConfigs)
-
-    // Shutdown a broker
-    adHocServers.last.shutdown()
-    adHocServers.last.awaitShutdown()
-
-    checkMetadata(adHocServers, numConfigs - 1)
-  }
-}
diff --git a/core/src/test/scala/unit/kafka/integration/UncleanLeaderElectionTest.scala b/core/src/test/scala/unit/kafka/integration/UncleanLeaderElectionTest.scala
index 7732e3888ef41..ca59380fa7e23 100755
--- a/core/src/test/scala/unit/kafka/integration/UncleanLeaderElectionTest.scala
+++ b/core/src/test/scala/unit/kafka/integration/UncleanLeaderElectionTest.scala
@@ -17,22 +17,30 @@
 
 package kafka.integration
 
-import org.apache.kafka.common.config.ConfigException
-import org.junit.{After, Before, Ignore, Test}
+import org.apache.kafka.common.config.{ConfigException, ConfigResource}
+import org.junit.{After, Before, Test}
 
 import scala.util.Random
+import scala.jdk.CollectionConverters._
+import scala.collection.{Map, Seq}
 import org.apache.log4j.{Level, Logger}
 import java.util.Properties
 import java.util.concurrent.ExecutionException
 
-import kafka.consumer.{Consumer, ConsumerConfig}
-import kafka.serializer.StringDecoder
 import kafka.server.{KafkaConfig, KafkaServer}
-import kafka.utils.CoreUtils
+import kafka.utils.{CoreUtils, TestUtils}
 import kafka.utils.TestUtils._
 import kafka.zk.ZooKeeperTestHarness
+import org.apache.kafka.common.TopicPartition
 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.{Admin, AdminClientConfig, AlterConfigsResult, Config, ConfigEntry}
 import org.junit.Assert._
+import org.scalatest.Assertions.intercept
+
+import scala.annotation.nowarn
 
 class UncleanLeaderElectionTest extends ZooKeeperTestHarness {
   val brokerId1 = 0
@@ -49,16 +57,14 @@ class UncleanLeaderElectionTest extends ZooKeeperTestHarness {
   var servers: Seq[KafkaServer] = Seq.empty[KafkaServer]
 
   val random = new Random()
-  val topic = "topic" + random.nextLong
+  val topic = "topic" + random.nextLong()
   val partitionId = 0
 
   val kafkaApisLogger = Logger.getLogger(classOf[kafka.server.KafkaApis])
   val networkProcessorLogger = Logger.getLogger(classOf[kafka.network.Processor])
-  val syncProducerLogger = Logger.getLogger(classOf[kafka.producer.SyncProducer])
-  val eventHandlerLogger = Logger.getLogger(classOf[kafka.producer.async.DefaultEventHandler[Object, Object]])
 
   @Before
-  override def setUp() {
+  override def setUp(): Unit = {
     super.setUp()
 
     configProps1 = createBrokerConfig(brokerId1, zkConnect)
@@ -73,25 +79,21 @@ class UncleanLeaderElectionTest extends ZooKeeperTestHarness {
     // temporarily set loggers to a higher level so that tests run quietly
     kafkaApisLogger.setLevel(Level.FATAL)
     networkProcessorLogger.setLevel(Level.FATAL)
-    syncProducerLogger.setLevel(Level.FATAL)
-    eventHandlerLogger.setLevel(Level.FATAL)
   }
 
   @After
-  override def tearDown() {
+  override def tearDown(): Unit = {
     servers.foreach(server => shutdownServer(server))
     servers.foreach(server => CoreUtils.delete(server.config.logDirs))
 
     // restore log levels
     kafkaApisLogger.setLevel(Level.ERROR)
     networkProcessorLogger.setLevel(Level.ERROR)
-    syncProducerLogger.setLevel(Level.ERROR)
-    eventHandlerLogger.setLevel(Level.ERROR)
 
     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)
@@ -108,21 +110,20 @@ class UncleanLeaderElectionTest extends ZooKeeperTestHarness {
     startBrokers(Seq(configProps1, configProps2))
 
     // create topic with 1 partition, 2 replicas, one on each broker
-    adminZkClient.createOrUpdateTopicPartitionAssignmentPathInZK(topic, Map(partitionId -> Seq(brokerId1, brokerId2)))
+    TestUtils.createTopic(zkClient, topic, Map(partitionId -> Seq(brokerId1, brokerId2)), servers)
 
-    verifyUncleanLeaderElectionEnabled
+    verifyUncleanLeaderElectionEnabled()
   }
 
   @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))
 
     // create topic with 1 partition, 2 replicas, one on each broker
-    adminZkClient.createOrUpdateTopicPartitionAssignmentPathInZK(topic, Map(partitionId -> Seq(brokerId1, brokerId2)))
+    TestUtils.createTopic(zkClient, topic, Map(partitionId -> Seq(brokerId1, brokerId2)), servers)
 
-    verifyUncleanLeaderElectionDisabled
+    verifyUncleanLeaderElectionDisabled()
   }
 
   @Test
@@ -135,15 +136,13 @@ class UncleanLeaderElectionTest extends ZooKeeperTestHarness {
     // create topic with 1 partition, 2 replicas, one on each broker, and unclean leader election enabled
     val topicProps = new Properties()
     topicProps.put("unclean.leader.election.enable", "true")
-    adminZkClient.createOrUpdateTopicPartitionAssignmentPathInZK(topic, Map(partitionId -> Seq(brokerId1, brokerId2)),
-      topicProps)
+    TestUtils.createTopic(zkClient, topic, Map(partitionId -> Seq(brokerId1, brokerId2)), servers, topicProps)
 
-    verifyUncleanLeaderElectionEnabled
+    verifyUncleanLeaderElectionEnabled()
   }
 
   @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")
@@ -152,10 +151,9 @@ class UncleanLeaderElectionTest extends ZooKeeperTestHarness {
     // create topic with 1 partition, 2 replicas, one on each broker, and unclean leader election disabled
     val topicProps = new Properties()
     topicProps.put("unclean.leader.election.enable", "false")
-    adminZkClient.createOrUpdateTopicPartitionAssignmentPathInZK(topic, Map(partitionId -> Seq(brokerId1, brokerId2)),
-      topicProps)
+    TestUtils.createTopic(zkClient, topic, Map(partitionId -> Seq(brokerId1, brokerId2)), servers, topicProps)
 
-    verifyUncleanLeaderElectionDisabled
+    verifyUncleanLeaderElectionDisabled()
   }
 
   @Test
@@ -167,13 +165,13 @@ class UncleanLeaderElectionTest extends ZooKeeperTestHarness {
     topicProps.put("unclean.leader.election.enable", "invalid")
 
     intercept[ConfigException] {
-      adminZkClient.createOrUpdateTopicPartitionAssignmentPathInZK(topic, Map(partitionId -> Seq(brokerId1)), topicProps)
+      TestUtils.createTopic(zkClient, topic, Map(partitionId -> Seq(brokerId1)), servers, topicProps)
     }
   }
 
   def verifyUncleanLeaderElectionEnabled(): Unit = {
     // wait until leader is elected
-    val leaderId = waitUntilLeaderIsElectedOrChanged(zkUtils, topic, partitionId)
+    val leaderId = waitUntilLeaderIsElectedOrChanged(zkClient, topic, partitionId)
     debug("Leader for " + topic  + " is elected to be: %s".format(leaderId))
     assertTrue("Leader id is set to expected value for topic: " + topic, leaderId == brokerId1 || leaderId == brokerId2)
 
@@ -183,30 +181,35 @@ class UncleanLeaderElectionTest extends ZooKeeperTestHarness {
 
     produceMessage(servers, topic, "first")
     waitUntilMetadataIsPropagated(servers, topic, partitionId)
-    assertEquals(List("first"), consumeAllMessages(topic))
+    assertEquals(List("first"), consumeAllMessages(topic, 1))
 
     // shutdown follower server
     servers.filter(server => server.config.brokerId == followerId).map(server => shutdownServer(server))
 
     produceMessage(servers, topic, "second")
-    assertEquals(List("first", "second"), consumeAllMessages(topic))
+    assertEquals(List("first", "second"), consumeAllMessages(topic, 2))
+
+    //remove any previous unclean election metric
+    servers.map(_.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 == followerId).map(server => server.startup())
+    servers.filter(_.config.brokerId == leaderId).map(shutdownServer)
+    val followerServer = servers.find(_.config.brokerId == followerId).get
+    followerServer.startup()
 
     // wait until new leader is (uncleanly) elected
-    waitUntilLeaderIsElectedOrChanged(zkUtils, topic, partitionId, newLeaderOpt = Some(followerId))
+    waitUntilLeaderIsElectedOrChanged(zkClient, topic, partitionId, newLeaderOpt = Some(followerId))
+    assertEquals(1, followerServer.kafkaController.controllerContext.stats.uncleanLeaderElectionRate.count())
 
     produceMessage(servers, topic, "third")
 
     // second message was lost due to unclean election
-    assertEquals(List("first", "third"), consumeAllMessages(topic))
+    assertEquals(List("first", "third"), consumeAllMessages(topic, 2))
   }
 
   def verifyUncleanLeaderElectionDisabled(): Unit = {
     // wait until leader is elected
-    val leaderId = waitUntilLeaderIsElectedOrChanged(zkUtils, topic, partitionId)
+    val leaderId = waitUntilLeaderIsElectedOrChanged(zkClient, topic, partitionId)
     debug("Leader for " + topic  + " is elected to be: %s".format(leaderId))
     assertTrue("Leader id is set to expected value for topic: " + topic, leaderId == brokerId1 || leaderId == brokerId2)
 
@@ -216,44 +219,53 @@ class UncleanLeaderElectionTest extends ZooKeeperTestHarness {
 
     produceMessage(servers, topic, "first")
     waitUntilMetadataIsPropagated(servers, topic, partitionId)
-    assertEquals(List("first"), consumeAllMessages(topic))
+    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))
+    assertEquals(List("first", "second"), consumeAllMessages(topic, 2))
+
+    //remove any previous unclean election metric
+    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 == followerId).map(server => server.startup())
+    servers.filter(server => server.config.brokerId == leaderId).foreach(server => shutdownServer(server))
+    val followerServer = servers.find(_.config.brokerId == followerId).get
+    followerServer.startup()
 
     // verify that unclean election to non-ISR follower does not occur
-    waitUntilLeaderIsElectedOrChanged(zkUtils, topic, partitionId, newLeaderOpt = Some(-1))
+    waitUntilLeaderIsElectedOrChanged(zkClient, topic, partitionId, newLeaderOpt = Some(-1))
+    assertEquals(0, followerServer.kafkaController.controllerContext.stats.uncleanLeaderElectionRate.count())
 
     // message production and consumption should both fail while leader is down
     try {
-      produceMessage(servers, topic, "third")
+      produceMessage(servers, topic, "third", deliveryTimeoutMs = 1000, requestTimeoutMs = 1000)
       fail("Message produced while leader is down should fail, but it succeeded")
     } catch {
       case e: ExecutionException if e.getCause.isInstanceOf[TimeoutException] => // expected
     }
 
-    assertEquals(List.empty[String], consumeAllMessages(topic))
+    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())
-    waitUntilLeaderIsElectedOrChanged(zkUtils, topic, partitionId, newLeaderOpt = Some(leaderId))
+    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(zkUtils, topic, partitionId, newLeaderOpt = Some(followerId))
+    waitUntilLeaderIsElectedOrChanged(zkClient, topic, partitionId, newLeaderOpt = Some(followerId))
 
     // verify messages can be consumed from ISR follower that was just promoted to leader
-    assertEquals(List("first", "second", "third"), consumeAllMessages(topic))
+    assertEquals(List("first", "second", "third"), consumeAllMessages(topic, 3))
   }
 
   private def shutdownServer(server: KafkaServer) = {
@@ -261,16 +273,95 @@ class UncleanLeaderElectionTest extends ZooKeeperTestHarness {
     server.awaitShutdown()
   }
 
-  private def consumeAllMessages(topic: String) : List[String] = {
-    // use a fresh consumer group every time so that we don't need to mess with disabling auto-commit or
-    // resetting the ZK offset
-    val consumerProps = createConsumerProperties(zkConnect, "group" + random.nextLong, "id", 1000)
-    val consumerConnector = Consumer.create(new ConsumerConfig(consumerProps))
-    val messageStream = consumerConnector.createMessageStreams(Map(topic -> 1), new StringDecoder(), new StringDecoder())
+  private def consumeAllMessages(topic: String, numMessages: Int): Seq[String] = {
+    val brokerList = TestUtils.bootstrapServers(servers, ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT))
+    // Don't rely on coordinator as it may be down when this method is called
+    val consumer = TestUtils.createConsumer(brokerList,
+      groupId = "group" + random.nextLong(),
+      enableAutoCommit = false,
+      valueDeserializer = new StringDeserializer)
+    try {
+      val tp = new TopicPartition(topic, partitionId)
+      consumer.assign(Seq(tp).asJava)
+      consumer.seek(tp, 0)
+      TestUtils.consumeRecords(consumer, numMessages).map(_.value)
+    } finally consumer.close()
+  }
+
+  @Test
+  def testTopicUncleanLeaderElectionEnable(): Unit = {
+    // unclean leader election is disabled by default
+    startBrokers(Seq(configProps1, configProps2))
+
+    // create topic with 1 partition, 2 replicas, one on each broker
+    adminZkClient.createTopicWithAssignment(topic, config = new Properties(), Map(partitionId -> Seq(brokerId1, brokerId2)))
 
-    val messages = getMessages(messageStream)
-    consumerConnector.shutdown
+    // wait until leader is elected
+    val leaderId = waitUntilLeaderIsElectedOrChanged(zkClient, topic, partitionId)
+
+    // the non-leader broker is the follower
+    val followerId = if (leaderId == brokerId1) brokerId2 else brokerId1
+
+    produceMessage(servers, topic, "first")
+    waitUntilMetadataIsPropagated(servers, topic, partitionId)
+    assertEquals(List("first"), consumeAllMessages(topic, 1))
+
+    // shutdown follower server
+    servers.filter(server => server.config.brokerId == followerId).map(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"))
+
+    // shutdown leader and then restart follower
+    servers.filter(server => server.config.brokerId == leaderId).map(server => shutdownServer(server))
+    val followerServer = servers.find(_.config.brokerId == followerId).get
+    followerServer.startup()
+
+    assertEquals(0, followerServer.kafkaController.controllerContext.stats.uncleanLeaderElectionRate.count())
+
+    // message production and consumption should both fail while leader is down
+    try {
+      produceMessage(servers, topic, "third", deliveryTimeoutMs = 1000, requestTimeoutMs = 1000)
+      fail("Message produced while leader is down should fail, but it succeeded")
+    } catch {
+      case e: ExecutionException if e.getCause.isInstanceOf[TimeoutException] => // expected
+    }
+
+    assertEquals(List.empty[String], consumeAllMessages(topic, 0))
+
+    // Enable unclean leader election for topic
+    val adminClient = createAdminClient()
+    val newProps = new Properties
+    newProps.put(KafkaConfig.UncleanLeaderElectionEnableProp, "true")
+    alterTopicConfigs(adminClient, topic, newProps).all.get
+    adminClient.close()
+
+    // wait until new leader is (uncleanly) elected
+    waitUntilLeaderIsElectedOrChanged(zkClient, topic, partitionId, newLeaderOpt = Some(followerId))
+    assertEquals(1, followerServer.kafkaController.controllerContext.stats.uncleanLeaderElectionRate.count())
+
+    produceMessage(servers, topic, "third")
+
+    // second message was lost due to unclean election
+    assertEquals(List("first", "third"), consumeAllMessages(topic, 2))
+  }
+
+  @nowarn("cat=deprecation")
+  private 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)
+  }
 
-    messages
+  private def createAdminClient(): Admin = {
+    val config = new Properties
+    val bootstrapServers = TestUtils.bootstrapServers(servers, new ListenerName("PLAINTEXT"))
+    config.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers)
+    config.put(AdminClientConfig.METADATA_MAX_AGE_CONFIG, "10")
+    Admin.create(config)
   }
 }
diff --git a/core/src/test/scala/unit/kafka/javaapi/consumer/ZookeeperConsumerConnectorTest.scala b/core/src/test/scala/unit/kafka/javaapi/consumer/ZookeeperConsumerConnectorTest.scala
deleted file mode 100644
index ede4ff7257696..0000000000000
--- a/core/src/test/scala/unit/kafka/javaapi/consumer/ZookeeperConsumerConnectorTest.scala
+++ /dev/null
@@ -1,120 +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.javaapi.consumer
-
-import java.util.Properties
-
-import kafka.server._
-import kafka.serializer._
-import kafka.integration.KafkaServerTestHarness
-import kafka.producer.KeyedMessage
-import kafka.javaapi.producer.Producer
-import kafka.utils.IntEncoder
-import kafka.utils.{Logging, TestUtils}
-import kafka.consumer.{KafkaStream, ConsumerConfig}
-import kafka.common.MessageStreamsExistException
-import org.junit.Test
-
-import scala.collection.JavaConversions
-
-import org.apache.log4j.{Level, Logger}
-import org.junit.Assert._
-
-@deprecated("This test has been deprecated and it will be removed in a future release", "0.10.0.0")
-class ZookeeperConsumerConnectorTest extends KafkaServerTestHarness with Logging {
-  val numNodes = 2
-  val numParts = 2
-  val topic = "topic1"
-
-  val overridingProps = new Properties()
-  overridingProps.put(KafkaConfig.NumPartitionsProp, numParts.toString)
-
-  def generateConfigs =
-    TestUtils.createBrokerConfigs(numNodes, zkConnect).map(KafkaConfig.fromProps(_, overridingProps))
-
-  val group = "group1"
-  val consumer1 = "consumer1"
-  val nMessages = 2
-
-  @Test
-  def testBasic() {
-    val requestHandlerLogger = Logger.getLogger(classOf[KafkaRequestHandler])
-    requestHandlerLogger.setLevel(Level.FATAL)
-
-    // create the topic
-    TestUtils.createTopic(zkUtils, topic, numParts, 1, servers)
-
-    // send some messages to each broker
-    val sentMessages1 = sendMessages(servers, nMessages, "batch1")
-
-    // create a consumer
-    val consumerConfig1 = new ConsumerConfig(TestUtils.createConsumerProperties(zkConnect, group, consumer1))
-    val zkConsumerConnector1 = new ZookeeperConsumerConnector(consumerConfig1, true)
-    val topicMessageStreams1 = zkConsumerConnector1.createMessageStreams(toJavaMap(Map(topic -> numNodes*numParts/2)), new StringDecoder(), new StringDecoder())
-
-    val receivedMessages1 = getMessages(nMessages*2, topicMessageStreams1)
-    assertEquals(sentMessages1.sorted, receivedMessages1.sorted)
-
-    // call createMesssageStreams twice should throw MessageStreamsExistException
-    try {
-      zkConsumerConnector1.createMessageStreams(toJavaMap(Map(topic -> numNodes*numParts/2)), new StringDecoder(), new StringDecoder())
-      fail("Should fail with MessageStreamsExistException")
-    } catch {
-      case _: MessageStreamsExistException => // expected
-    }
-    zkConsumerConnector1.shutdown
-    info("all consumer connectors stopped")
-    requestHandlerLogger.setLevel(Level.ERROR)
-  }
-
-  def sendMessages(servers: Seq[KafkaServer],
-                   messagesPerNode: Int,
-                   header: String): List[String] = {
-    var messages: List[String] = Nil
-    for(server <- servers) {
-      val producer: kafka.producer.Producer[Int, String] =
-        TestUtils.createProducer(TestUtils.getBrokerListStrFromServers(servers),
-          encoder = classOf[StringEncoder].getName,
-          keyEncoder = classOf[IntEncoder].getName)
-      val javaProducer: Producer[Int, String] = new kafka.javaapi.producer.Producer(producer)
-      for (partition <- 0 until numParts) {
-        val ms = 0.until(messagesPerNode).map(x => header + server.config.brokerId + "-" + partition + "-" + x)
-        messages ++= ms
-        import JavaConversions._
-        javaProducer.send(ms.map(new KeyedMessage[Int, String](topic, partition, _)): java.util.List[KeyedMessage[Int, String]])
-      }
-      javaProducer.close
-    }
-    messages
-  }
-
-  def getMessages(nMessagesPerThread: Int,
-                  jTopicMessageStreams: java.util.Map[String, java.util.List[KafkaStream[String, String]]]): List[String] = {
-    var messages: List[String] = Nil
-    import scala.collection.JavaConversions._
-    val topicMessageStreams = jTopicMessageStreams.mapValues(_.toList)
-    messages = TestUtils.getMessages(topicMessageStreams, nMessagesPerThread)
-    messages
-  }
-
-  private def toJavaMap(scalaMap: Map[String, Int]): java.util.Map[String, java.lang.Integer] = {
-    val javaMap = new java.util.HashMap[String, java.lang.Integer]()
-    scalaMap.foreach(m => javaMap.put(m._1, m._2.asInstanceOf[java.lang.Integer]))
-    javaMap
-  }
-}
diff --git a/core/src/test/scala/unit/kafka/javaapi/message/BaseMessageSetTestCases.scala b/core/src/test/scala/unit/kafka/javaapi/message/BaseMessageSetTestCases.scala
deleted file mode 100644
index 199bbbdf4ca69..0000000000000
--- a/core/src/test/scala/unit/kafka/javaapi/message/BaseMessageSetTestCases.scala
+++ /dev/null
@@ -1,69 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package kafka.javaapi.message
-
-import org.junit.Assert._
-import org.scalatest.junit.JUnitSuite
-import org.junit.Test
-import kafka.message.{CompressionCodec, DefaultCompressionCodec, Message, NoCompressionCodec}
-import org.apache.kafka.test.TestUtils
-
-import scala.collection.JavaConverters._
-
-trait BaseMessageSetTestCases extends JUnitSuite {
-  
-  val messages = Array(new Message("abcd".getBytes()), new Message("efgh".getBytes()))
-  def createMessageSet(messages: Seq[Message], compressed: CompressionCodec = NoCompressionCodec): MessageSet
-
-  @Test
-  def testWrittenEqualsRead(): Unit = {
-    val messageSet = createMessageSet(messages)
-    assertEquals(messages.toSeq, messageSet.asScala.map(m => m.message))
-  }
-
-  @Test
-  def testIteratorIsConsistent() {
-    val m = createMessageSet(messages)
-    // two iterators over the same set should give the same results
-    TestUtils.checkEquals(m, m)
-  }
-
-  @Test
-  def testIteratorIsConsistentWithCompression() {
-    val m = createMessageSet(messages, DefaultCompressionCodec)
-    // two iterators over the same set should give the same results
-    TestUtils.checkEquals(m, m)
-  }
-
-  @Test
-  def testSizeInBytes() {
-    assertEquals("Empty message set should have 0 bytes.",
-                 0,
-                 createMessageSet(Array[Message]()).sizeInBytes)
-    assertEquals("Predicted size should equal actual size.", 
-                 kafka.message.MessageSet.messageSetSize(messages),
-                 createMessageSet(messages).sizeInBytes)
-  }
-
-  @Test
-  def testSizeInBytesWithCompression () {
-    assertEquals("Empty message set should have 0 bytes.",
-                 0,           // overhead of the GZIP output stream
-                 createMessageSet(Array[Message](), DefaultCompressionCodec).sizeInBytes)
-  }
-}
diff --git a/core/src/test/scala/unit/kafka/javaapi/message/ByteBufferMessageSetTest.scala b/core/src/test/scala/unit/kafka/javaapi/message/ByteBufferMessageSetTest.scala
deleted file mode 100644
index fbdb000a1d8c4..0000000000000
--- a/core/src/test/scala/unit/kafka/javaapi/message/ByteBufferMessageSetTest.scala
+++ /dev/null
@@ -1,46 +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.javaapi.message
-
-import org.junit.Assert._
-import org.junit.Test
-import kafka.message.{DefaultCompressionCodec, CompressionCodec, NoCompressionCodec, Message}
-
-class ByteBufferMessageSetTest extends kafka.javaapi.message.BaseMessageSetTestCases {
-  
-  override def createMessageSet(messages: Seq[Message], compressed: CompressionCodec = NoCompressionCodec): ByteBufferMessageSet =
-    new ByteBufferMessageSet(new kafka.message.ByteBufferMessageSet(compressed, messages: _*).buffer)
-
-  val msgSeq: Seq[Message] = Seq(new Message("hello".getBytes()), new Message("there".getBytes()))
-
-  @Test
-  def testEquals() {
-    val messageList = createMessageSet(msgSeq, NoCompressionCodec)
-    val moreMessages = createMessageSet(msgSeq, NoCompressionCodec)
-    assertEquals(messageList, moreMessages)
-    assertTrue(messageList.equals(moreMessages))
-  }
-
-  @Test
-  def testEqualsWithCompression () {
-    val messageList = createMessageSet(msgSeq, DefaultCompressionCodec)
-    val moreMessages = createMessageSet(msgSeq, DefaultCompressionCodec)
-    assertEquals(messageList, moreMessages)
-    assertTrue(messageList.equals(moreMessages))
-  }
-}
diff --git a/core/src/test/scala/unit/kafka/log/AbstractLogCleanerIntegrationTest.scala b/core/src/test/scala/unit/kafka/log/AbstractLogCleanerIntegrationTest.scala
index bff27006e287a..fe98ebfe64c57 100644
--- a/core/src/test/scala/unit/kafka/log/AbstractLogCleanerIntegrationTest.scala
+++ b/core/src/test/scala/unit/kafka/log/AbstractLogCleanerIntegrationTest.scala
@@ -24,11 +24,17 @@ import kafka.server.{BrokerTopicStats, LogDirFailureChannel}
 import kafka.utils.{MockTime, Pool, TestUtils}
 import kafka.utils.Implicits._
 import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.record.{CompressionType, MemoryRecords, RecordBatch}
 import org.apache.kafka.common.utils.Utils
+import org.apache.kafka.test.IntegrationTest
 import org.junit.After
+import org.junit.experimental.categories.Category
 
+import scala.collection.Seq
 import scala.collection.mutable.ListBuffer
+import scala.util.Random
 
+@Category(Array(classOf[IntegrationTest]))
 abstract class AbstractLogCleanerIntegrationTest {
 
   var cleaner: LogCleaner = _
@@ -37,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 = 256
+  private val defaultSegmentSize = 2048
+  private val defaultMaxCompactionLagMs = Long.MaxValue
 
   def time: MockTime
 
@@ -55,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)
@@ -66,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
   }
@@ -76,9 +85,11 @@ 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 = {
 
     val logMap = new Pool[TopicPartition, Log]()
@@ -89,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,
@@ -108,7 +120,7 @@ abstract class AbstractLogCleanerIntegrationTest {
 
     val cleanerConfig = CleanerConfig(
       numThreads = numThreads,
-      ioBufferSize = maxMessageSize / 2,
+      ioBufferSize = cleanerIoBufferSize.getOrElse(maxMessageSize / 2),
       maxMessageSize = maxMessageSize,
       backOffMs = backOffMs)
     new LogCleaner(cleanerConfig,
@@ -117,4 +129,31 @@ abstract class AbstractLogCleanerIntegrationTest {
       logDirFailureChannel = new LogDirFailureChannel(1),
       time = time)
   }
+
+  def codec: CompressionType
+  private var ctr = 0
+  def counter: Int = ctr
+  def incCounter(): Unit = ctr += 1
+
+  def writeDups(numKeys: Int, numDups: Int, log: Log, codec: CompressionType,
+                        startKey: Int = 0, magicValue: Byte = RecordBatch.CURRENT_MAGIC_VALUE): Seq[(Int, String, Long)] = {
+    for(_ <- 0 until numDups; key <- startKey until (startKey + numKeys)) yield {
+      val value = counter.toString
+      val appendInfo = log.appendAsLeader(TestUtils.singletonRecords(value = value.toString.getBytes, codec = codec,
+        key = key.toString.getBytes, magicValue = magicValue), leaderEpoch = 0)
+      incCounter()
+      (key, value, appendInfo.firstOffset.get)
+    }
+  }
+
+  def createLargeSingleMessageSet(key: Int, messageFormatVersion: Byte): (String, MemoryRecords) = {
+    def messageValue(length: Int): String = {
+      val random = new Random(0)
+      new String(random.alphanumeric.take(length).toArray)
+    }
+    val value = messageValue(128)
+    val messageSet = TestUtils.singletonRecords(value = value.getBytes, codec = codec, key = key.toString.getBytes,
+      magicValue = messageFormatVersion)
+    (value, messageSet)
+  }
 }
diff --git a/core/src/test/scala/unit/kafka/log/BrokerCompressionTest.scala b/core/src/test/scala/unit/kafka/log/BrokerCompressionTest.scala
index 1cf393e82aa85..6e7ae74fca6ab 100755
--- a/core/src/test/scala/unit/kafka/log/BrokerCompressionTest.scala
+++ b/core/src/test/scala/unit/kafka/log/BrokerCompressionTest.scala
@@ -19,22 +19,21 @@ 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
 import org.junit.runners.Parameterized
 import org.junit.runners.Parameterized.Parameters
-import org.apache.kafka.common.record.{CompressionType, MemoryRecords, SimpleRecord}
+import org.apache.kafka.common.record.{CompressionType, MemoryRecords, RecordBatch, SimpleRecord}
 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._
+import scala.jdk.CollectionConverters._
 
 @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)
@@ -64,7 +63,13 @@ class BrokerCompressionTest(messageCompression: String, brokerCompression: Strin
     log.appendAsLeader(MemoryRecords.withRecords(CompressionType.forId(messageCompressionCode.codec), 0,
           new SimpleRecord("hello".getBytes), new SimpleRecord("there".getBytes)), leaderEpoch = 0)
 
-    def readBatch(offset: Int) = log.readUncommitted(offset, 4096).records.batches.iterator.next()
+    def readBatch(offset: Int): RecordBatch = {
+      val fetchInfo = log.read(offset,
+        maxLength = 4096,
+        isolation = FetchLogEnd,
+        minOneMessage = true)
+      fetchInfo.records.batches.iterator.next()
+    }
 
     if (!brokerCompression.equals("producer")) {
       val brokerCompressionCode = BrokerCompressionCodec.getCompressionCodec(brokerCompression)
diff --git a/core/src/test/scala/unit/kafka/log/LogCleanerIntegrationTest.scala b/core/src/test/scala/unit/kafka/log/LogCleanerIntegrationTest.scala
old mode 100755
new mode 100644
index b20622f99ab29..1ac9e62a7b1b0
--- a/core/src/test/scala/unit/kafka/log/LogCleanerIntegrationTest.scala
+++ b/core/src/test/scala/unit/kafka/log/LogCleanerIntegrationTest.scala
@@ -1,298 +1,216 @@
 /**
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
+  * Licensed to the Apache Software Foundation (ASF) under one or more
+  * contributor license agreements.  See the NOTICE file distributed with
+  * this work for additional information regarding copyright ownership.
+  * The ASF licenses this file to You under the Apache License, Version 2.0
+  * (the "License"); you may not use this file except in compliance with
+  * the License.  You may obtain a copy of the License at
+  *
+  *    http://www.apache.org/licenses/LICENSE-2.0
+  *
+  * Unless required by applicable law or agreed to in writing, software
+  * distributed under the License is distributed on an "AS IS" BASIS,
+  * WITHOUT WARRANTIES 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.util.Properties
+import java.io.PrintWriter
 
-import kafka.api.KAFKA_0_11_0_IV0
-import kafka.api.{KAFKA_0_10_0_IV1, KAFKA_0_9_0}
-import kafka.server.checkpoints.OffsetCheckpointFile
-import kafka.utils._
+import com.yammer.metrics.core.{Gauge, MetricName}
+import kafka.metrics.{KafkaMetricsGroup, KafkaYammerMetrics}
+import kafka.utils.{MockTime, TestUtils}
 import org.apache.kafka.common.TopicPartition
-import org.apache.kafka.common.record._
+import org.apache.kafka.common.record.{CompressionType, RecordBatch}
+import org.apache.kafka.test.TestUtils.DEFAULT_MAX_WAIT_MS
 import org.junit.Assert._
-import org.junit._
-import org.junit.runner.RunWith
-import org.junit.runners.Parameterized
-import org.junit.runners.Parameterized.Parameters
+import org.junit.{After, Test}
 
-import scala.Seq
-import scala.collection._
-import scala.util.Random
+import scala.collection.{Iterable, Seq}
+import scala.jdk.CollectionConverters._
 
 /**
- * This is an integration test that tests the fully integrated log cleaner
- */
-@RunWith(value = classOf[Parameterized])
-class LogCleanerIntegrationTest(compressionCodec: String) extends AbstractLogCleanerIntegrationTest {
+  * This is an integration test that tests the fully integrated log cleaner
+  */
+class LogCleanerIntegrationTest extends AbstractLogCleanerIntegrationTest with KafkaMetricsGroup {
+
+  val codec: CompressionType = CompressionType.LZ4
 
-  val codec = CompressionType.forName(compressionCodec)
   val time = new MockTime()
-  var counter = 0
   val topicPartitions = Array(new TopicPartition("log", 0), new TopicPartition("log", 1), new TopicPartition("log", 2))
 
-  @Test
-  def cleanerTest() {
+  @After
+  def cleanup(): Unit = {
+    TestUtils.clearYammerMetrics()
+  }
+
+  @Test(timeout = DEFAULT_MAX_WAIT_MS)
+  def testMarksPartitionsAsOfflineAndPopulatesUncleanableMetrics(): Unit = {
     val largeMessageKey = 20
-    val (largeMessageValue, largeMessageSet) = createLargeSingleMessageSet(largeMessageKey, RecordBatch.CURRENT_MAGIC_VALUE)
+    val (_, largeMessageSet) = createLargeSingleMessageSet(largeMessageKey, RecordBatch.CURRENT_MAGIC_VALUE)
     val maxMessageSize = largeMessageSet.sizeInBytes
+    cleaner = makeCleaner(partitions = topicPartitions, maxMessageSize = maxMessageSize, backOffMs = 100)
 
-    cleaner = makeCleaner(partitions = topicPartitions, maxMessageSize = maxMessageSize)
-    val log = cleaner.logs.get(topicPartitions(0))
-
-    val appends = writeDups(numKeys = 100, numDups = 3, log = log, codec = codec)
-    val startSize = log.size
-    cleaner.startup()
+    def breakPartitionLog(tp: TopicPartition): Unit = {
+      val log = cleaner.logs.get(tp)
+      writeDups(numKeys = 20, numDups = 3, log = log, codec = codec)
 
-    val firstDirty = log.activeSegment.baseOffset
-    checkLastCleaned("log", 0, firstDirty)
-    val compactedSize = log.logSegments.map(_.size).sum
-    assertTrue(s"log should have been compacted: startSize=$startSize compactedSize=$compactedSize", startSize > compactedSize)
+      val partitionFile = log.logSegments.last.log.file()
+      val writer = new PrintWriter(partitionFile)
+      writer.write("jogeajgoea")
+      writer.close()
 
-    checkLogAfterAppendingDups(log, startSize, appends)
-
-    val appendInfo = log.appendAsLeader(largeMessageSet, leaderEpoch = 0)
-    val largeMessageOffset = appendInfo.firstOffset
+      writeDups(numKeys = 20, numDups = 3, log = log, codec = codec)
+    }
 
-    val dups = writeDups(startKey = largeMessageKey + 1, numKeys = 100, numDups = 3, log = log, codec = codec)
-    val appends2 = appends ++ Seq((largeMessageKey, largeMessageValue, largeMessageOffset)) ++ dups
-    val firstDirty2 = log.activeSegment.baseOffset
-    checkLastCleaned("log", 0, firstDirty2)
+    breakPartitionLog(topicPartitions(0))
+    breakPartitionLog(topicPartitions(1))
 
-    checkLogAfterAppendingDups(log, startSize, appends2)
+    cleaner.startup()
 
-    // simulate deleting a partition, by removing it from logs
-    // force a checkpoint
-    // and make sure its gone from checkpoint file
-    cleaner.logs.remove(topicPartitions(0))
-    cleaner.updateCheckpoints(logDir)
-    val checkpoints = new OffsetCheckpointFile(new File(logDir, cleaner.cleanerManager.offsetCheckpointFile)).read()
-    // we expect partition 0 to be gone
-    assertFalse(checkpoints.contains(topicPartitions(0)))
+    val log = cleaner.logs.get(topicPartitions(0))
+    val log2 = cleaner.logs.get(topicPartitions(1))
+    val uncleanableDirectory = log.dir.getParent
+    val uncleanablePartitionsCountGauge = getGauge[Int]("uncleanable-partitions-count", uncleanableDirectory)
+    val uncleanableBytesGauge = getGauge[Long]("uncleanable-bytes", uncleanableDirectory)
+
+    TestUtils.waitUntilTrue(() => uncleanablePartitionsCountGauge.value() == 2, "There should be 2 uncleanable partitions", 2000L)
+    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)
+
+    val uncleanablePartitions = cleaner.cleanerManager.uncleanablePartitions(uncleanableDirectory)
+    assertTrue(uncleanablePartitions.contains(topicPartitions(0)))
+    assertTrue(uncleanablePartitions.contains(topicPartitions(1)))
+    assertFalse(uncleanablePartitions.contains(topicPartitions(2)))
   }
 
-  @Test
-  def testCleansCombinedCompactAndDeleteTopic(): Unit = {
-    val logProps  = new Properties()
-    val retentionMs: Integer = 100000
-    logProps.put(LogConfig.RetentionMsProp, retentionMs: Integer)
-    logProps.put(LogConfig.CleanupPolicyProp, "compact,delete")
-
-    def runCleanerAndCheckCompacted(numKeys: Int): (Log, Seq[(Int, String, Long)]) = {
-      cleaner = makeCleaner(partitions = topicPartitions.take(1), propertyOverrides = logProps, backOffMs = 100L)
-      val log = cleaner.logs.get(topicPartitions(0))
-
-      val messages = writeDups(numKeys = numKeys, numDups = 3, log = log, codec = codec)
-      val startSize = log.size
-
-      log.onHighWatermarkIncremented(log.logEndOffset)
-
-      val firstDirty = log.activeSegment.baseOffset
-      cleaner.startup()
-
-      // should compact the log
-      checkLastCleaned("log", 0, firstDirty)
-      val compactedSize = log.logSegments.map(_.size).sum
-      assertTrue(s"log should have been compacted: startSize=$startSize compactedSize=$compactedSize", startSize > compactedSize)
-      (log, messages)
-    }
-
-    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)
-    assertEquals(1, log.numberOfSegments)
-
-    cleaner.shutdown()
+  private def getGauge[T](filter: MetricName => Boolean): Gauge[T] = {
+    KafkaYammerMetrics.defaultRegistry.allMetrics.asScala
+      .filter { case (k, _) => filter(k) }
+      .headOption
+      .getOrElse { fail(s"Unable to find metric") }
+      .asInstanceOf[(Any, Gauge[Any])]
+      ._2
+      .asInstanceOf[Gauge[T]]
+  }
 
-    // run the cleaner again to make sure if there are no issues post deletion
-    val (log2, messages) = runCleanerAndCheckCompacted(20)
-    val read = readFromLog(log2)
-    assertEquals("Contents of the map shouldn't change", toMap(messages), toMap(read))
+  private def getGauge[T](metricName: String): Gauge[T] = {
+    getGauge(mName => mName.getName.endsWith(metricName) && mName.getScope == null)
   }
 
-  // returns (value, ByteBufferMessageSet)
-  private def createLargeSingleMessageSet(key: Int, messageFormatVersion: Byte): (String, MemoryRecords) = {
-    def messageValue(length: Int): String = {
-      val random = new Random(0)
-      new String(random.alphanumeric.take(length).toArray)
-    }
-    val value = messageValue(128)
-    val messageSet = TestUtils.singletonRecords(value = value.getBytes, codec = codec, key = key.toString.getBytes,
-      magicValue = messageFormatVersion)
-    (value, messageSet)
+  private def getGauge[T](metricName: String, metricScope: String): Gauge[T] = {
+    getGauge(k => k.getName.endsWith(metricName) && k.getScope.endsWith(metricScope))
   }
 
   @Test
-  def testCleanerWithMessageFormatV0(): Unit = {
-    val largeMessageKey = 20
-    val (largeMessageValue, largeMessageSet) = createLargeSingleMessageSet(largeMessageKey, RecordBatch.MAGIC_VALUE_V0)
-    val maxMessageSize = codec match {
-      case CompressionType.NONE => largeMessageSet.sizeInBytes
-      case _ =>
-        // the broker assigns absolute offsets for message format 0 which potentially causes the compressed size to
-        // increase because the broker offsets are larger than the ones assigned by the client
-        // adding `5` to the message set size is good enough for this test: it covers the increased message size while
-        // still being less than the overhead introduced by the conversion from message format version 0 to 1
-        largeMessageSet.sizeInBytes + 5
-    }
-
-    cleaner = makeCleaner(partitions = topicPartitions, maxMessageSize = maxMessageSize)
-
+  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 props = logConfigProperties(maxMessageSize = maxMessageSize)
-    props.put(LogConfig.MessageFormatVersionProp, KAFKA_0_9_0.version)
-    log.config = new LogConfig(props)
-
-    val appends = writeDups(numKeys = 100, numDups = 3, log = log, codec = codec, magicValue = RecordBatch.MAGIC_VALUE_V0)
-    val startSize = log.size
-    cleaner.startup()
-
-    val firstDirty = log.activeSegment.baseOffset
-    checkLastCleaned("log", 0, firstDirty)
-    val compactedSize = log.logSegments.map(_.size).sum
-    assertTrue(s"log should have been compacted: startSize=$startSize compactedSize=$compactedSize", startSize > compactedSize)
 
-    checkLogAfterAppendingDups(log, startSize, appends)
+    val T0 = time.milliseconds
+    writeKeyDups(numKeys = 100, numDups = 3, log, CompressionType.NONE, timestamp = T0, startValue = 0, step = 1)
 
-    val appends2: Seq[(Int, String, Long)] = {
-      val dupsV0 = writeDups(numKeys = 40, numDups = 3, log = log, codec = codec, magicValue = RecordBatch.MAGIC_VALUE_V0)
-      val appendInfo = log.appendAsLeader(largeMessageSet, leaderEpoch = 0)
-      val largeMessageOffset = appendInfo.firstOffset
+    val startSizeBlock0 = log.size
 
-      // also add some messages with version 1 and version 2 to check that we handle mixed format versions correctly
-      props.put(LogConfig.MessageFormatVersionProp, KAFKA_0_11_0_IV0.version)
-      log.config = new LogConfig(props)
-      val dupsV1 = writeDups(startKey = 30, numKeys = 40, numDups = 3, log = log, codec = codec, magicValue = RecordBatch.MAGIC_VALUE_V1)
-      val dupsV2 = writeDups(startKey = 15, numKeys = 5, numDups = 3, log = log, codec = codec, magicValue = RecordBatch.MAGIC_VALUE_V2)
-      appends ++ dupsV0 ++ Seq((largeMessageKey, largeMessageValue, largeMessageOffset)) ++ dupsV1 ++ dupsV2
-    }
-    val firstDirty2 = log.activeSegment.baseOffset
-    checkLastCleaned("log", 0, firstDirty2)
-
-    checkLogAfterAppendingDups(log, startSize, appends2)
-  }
+    val activeSegAtT0 = log.activeSegment
 
-  @Test
-  def testCleaningNestedMessagesWithMultipleVersions(): Unit = {
-    val maxMessageSize = 192
-    cleaner = makeCleaner(partitions = topicPartitions, maxMessageSize = maxMessageSize)
-
-    val log = cleaner.logs.get(topicPartitions(0))
-    val props = logConfigProperties(maxMessageSize = maxMessageSize)
-    props.put(LogConfig.MessageFormatVersionProp, KAFKA_0_9_0.version)
-    log.config = new LogConfig(props)
-
-    // with compression enabled, these messages will be written as a single message containing
-    // all of the individual messages
-    var appendsV0 = writeDupsSingleMessageSet(numKeys = 2, numDups = 3, log = log, codec = codec, magicValue = RecordBatch.MAGIC_VALUE_V0)
-    appendsV0 ++= writeDupsSingleMessageSet(numKeys = 2, startKey = 3, numDups = 2, log = log, codec = codec, magicValue = RecordBatch.MAGIC_VALUE_V0)
+    cleaner.startup()
 
-    props.put(LogConfig.MessageFormatVersionProp, KAFKA_0_10_0_IV1.version)
-    log.config = new LogConfig(props)
+    // 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)
 
-    var appendsV1 = writeDupsSingleMessageSet(startKey = 4, numKeys = 2, numDups = 2, log = log, codec = codec, magicValue = RecordBatch.MAGIC_VALUE_V1)
-    appendsV1 ++= writeDupsSingleMessageSet(startKey = 4, numKeys = 2, numDups = 2, log = log, codec = codec, magicValue = RecordBatch.MAGIC_VALUE_V1)
-    appendsV1 ++= writeDupsSingleMessageSet(startKey = 6, numKeys = 2, numDups = 2, log = log, codec = codec, magicValue = RecordBatch.MAGIC_VALUE_V1)
+    // advance to time a bit more than one maxCompactionLagMs from start
+    time.sleep(maxCompactionLagMs/2 + 1)
+    val T1 = time.milliseconds
 
-    val appends = appendsV0 ++ appendsV1
+    // 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)
 
-    val startSize = log.size
-    cleaner.startup()
+    // roll the active segment
+    log.roll()
+    val activeSegAtT1 = log.activeSegment
+    val firstBlockCleanableSegmentOffset = activeSegAtT0.baseOffset
 
-    val firstDirty = log.activeSegment.baseOffset
-    assertTrue(firstDirty > appendsV0.size) // ensure we clean data from V0 and V1
+    // the first block should get cleaned
+    cleaner.awaitCleaned(new TopicPartition("log", 0), firstBlockCleanableSegmentOffset)
 
-    checkLastCleaned("log", 0, firstDirty)
-    val compactedSize = log.logSegments.map(_.size).sum
-    assertTrue(s"log should have been compacted: startSize=$startSize compactedSize=$compactedSize", startSize > compactedSize)
+    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)
 
-    checkLogAfterAppendingDups(log, startSize, appends)
-  }
+    //minCleanableDirtyRatio  will prevent second block of data from compacting
+    assertNotEquals(s"log should still contain non-zero keys", appends1, read1)
 
-  private def checkLastCleaned(topic: String, partitionId: Int, firstDirty: Long) {
-    // 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)
-    cleaner.awaitCleaned(topicPartition, firstDirty)
-    val lastCleaned = cleaner.cleanerManager.allCleanerCheckpoints(topicPartition)
-    assertTrue(s"log cleaner should have processed up to offset $firstDirty, but lastCleaned=$lastCleaned",
-      lastCleaned >= firstDirty)
-  }
+    time.sleep(maxCompactionLagMs + 1)
+    // the second block should get cleaned. only zero keys left
+    cleaner.awaitCleaned(new TopicPartition("log", 0), activeSegAtT1.baseOffset)
 
-  private def checkLogAfterAppendingDups(log: Log, startSize: Long, appends: Seq[(Int, String, Long)]) {
-    val read = readFromLog(log)
-    assertEquals("Contents of the map shouldn't change", toMap(appends), toMap(read))
-    assertTrue(startSize > log.size)
-  }
+    val read2 = readFromLog(log)
 
-  private def toMap(messages: Iterable[(Int, String, Long)]): Map[Int, (String, Long)] = {
-    messages.map { case (key, value, offset) => key -> (value, offset) }.toMap
-  }
+    assertEquals(s"log should only contains zero keys now", appends1, read2)
 
-  private def readFromLog(log: Log): Iterable[(Int, String, Long)] = {
-    import JavaConverters._
-    for (segment <- log.logSegments; deepLogEntry <- segment.log.records.asScala) yield {
-      val key = TestUtils.readString(deepLogEntry.key).toInt
-      val value = TestUtils.readString(deepLogEntry.value)
-      (key, value, deepLogEntry.offset)
-    }
+    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 writeDups(numKeys: Int, numDups: Int, log: Log, codec: CompressionType,
-                        startKey: Int = 0, magicValue: Byte = RecordBatch.CURRENT_MAGIC_VALUE): Seq[(Int, String, Long)] = {
-    for(_ <- 0 until numDups; key <- startKey until (startKey + numKeys)) yield {
-      val value = counter.toString
-      val appendInfo = log.appendAsLeader(TestUtils.singletonRecords(value = value.toString.getBytes, codec = codec,
-              key = key.toString.getBytes, magicValue = magicValue), leaderEpoch = 0)
-      counter += 1
-      (key, value, appendInfo.firstOffset)
+  private def readFromLog(log: Log): Iterable[(Int, Int)] = {
+    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 writeDupsSingleMessageSet(numKeys: Int, numDups: Int, log: Log, codec: CompressionType,
-                                        startKey: Int = 0, magicValue: Byte): Seq[(Int, String, Long)] = {
-    val kvs = for (_ <- 0 until numDups; key <- startKey until (startKey + numKeys)) yield {
-      val payload = counter.toString
-      counter += 1
-      (key, payload)
-    }
-
-    val records = kvs.map { case (key, payload) =>
-      new SimpleRecord(key.toString.getBytes, payload.toString.getBytes)
+  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)
     }
-
-    val appendInfo = log.appendAsLeader(MemoryRecords.withRecords(magicValue, codec, records: _*), leaderEpoch = 0)
-    val offsets = appendInfo.firstOffset to appendInfo.lastOffset
-
-    kvs.zip(offsets).map { case (kv, offset) => (kv._1, kv._2, offset) }
   }
 
-}
-
-object LogCleanerIntegrationTest {
-  @Parameters
-  def parameters: java.util.Collection[Array[String]] = {
-    val list = new java.util.ArrayList[Array[String]]()
-    for (codec <- CompressionType.values)
-      list.add(Array(codec.name))
-    list
+  @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 bf634d731dd42..485111164e99e 100644
--- a/core/src/test/scala/unit/kafka/log/LogCleanerLagIntegrationTest.scala
+++ b/core/src/test/scala/unit/kafka/log/LogCleanerLagIntegrationTest.scala
@@ -27,6 +27,7 @@ import org.junit.runners.Parameterized
 import org.junit.runners.Parameterized.Parameters
 
 import scala.collection._
+import scala.jdk.CollectionConverters._
 
 /**
   * This is an integration test that tests the fully integrated log cleaner
@@ -35,27 +36,28 @@ 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
   val segmentSize = 512
-  var counter = 0
+
+  override def codec: CompressionType = CompressionType.forName(compressionCodecName)
+
   val topicPartitions = Array(new TopicPartition("log", 0), new TopicPartition("log", 1), new TopicPartition("log", 2))
-  val compressionCodec = CompressionType.forName(compressionCodecName)
 
   @Test
   def cleanerTest(): Unit = {
     cleaner = makeCleaner(partitions = topicPartitions,
       backOffMs = cleanerBackOffMs,
-      compactionLag = compactionLag,
+      minCompactionLagMs = minCompactionLag,
       segmentSize = segmentSize)
     val log = cleaner.logs.get(topicPartitions(0))
 
     // t = T0
     val T0 = time.milliseconds
-    val appends0 = writeDups(numKeys = 100, numDups = 3, log, compressionCodec, timestamp = T0)
+    val appends0 = writeDups(numKeys = 100, numDups = 3, log, codec, timestamp = T0)
     val startSizeBlock0 = log.size
     debug(s"total log size at T0: $startSizeBlock0")
 
@@ -68,17 +70,17 @@ 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
-    val appends1 = appends0 ++ writeDups(numKeys = 100, numDups = 3, log, compressionCodec, timestamp = T1)
+    val appends1 = appends0 ++ writeDups(numKeys = 100, numDups = 3, log, codec, timestamp = T1)
     val firstBlock1SegmentBaseOffset = activeSegAtT0.baseOffset
 
     // the first block should get cleaned
@@ -97,8 +99,6 @@ class LogCleanerLagIntegrationTest(compressionCodecName: String) extends Abstrac
   }
 
   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
@@ -111,11 +111,10 @@ class LogCleanerLagIntegrationTest(compressionCodecName: String) extends Abstrac
       val count = counter
       log.appendAsLeader(TestUtils.singletonRecords(value = counter.toString.getBytes, codec = codec,
               key = key.toString.getBytes, timestamp = timestamp), leaderEpoch = 0)
-      counter += 1
+      incCounter()
       (key, count)
     }
   }
-
 }
 
 object LogCleanerLagIntegrationTest {
diff --git a/core/src/test/scala/unit/kafka/log/LogCleanerManagerTest.scala b/core/src/test/scala/unit/kafka/log/LogCleanerManagerTest.scala
index 114602919ea38..6a57e25f48b02 100644
--- a/core/src/test/scala/unit/kafka/log/LogCleanerManagerTest.scala
+++ b/core/src/test/scala/unit/kafka/log/LogCleanerManagerTest.scala
@@ -27,27 +27,254 @@ 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 tmpDir2 = TestUtils.tempDir()
   val logDir = TestUtils.randomPartitionLogDir(tmpDir)
+  val logDir2 = TestUtils.randomPartitionLogDir(tmpDir)
+  val topicPartition = new TopicPartition("log", 0)
+  val topicPartition2 = new TopicPartition("log2", 0)
   val logProps = new Properties()
   logProps.put(LogConfig.SegmentBytesProp, 1024: java.lang.Integer)
   logProps.put(LogConfig.SegmentIndexBytesProp, 1024: java.lang.Integer)
   logProps.put(LogConfig.CleanupPolicyProp, LogConfig.Compact)
   val logConfig = LogConfig(logProps)
   val time = new MockTime(1400000000000L, 1000L)  // Tue May 13 16:53:20 UTC 2014 for `currentTimeMs`
+  val offset = 999
+
+  val cleanerCheckpoints: mutable.Map[TopicPartition, Long] = mutable.Map[TopicPartition, Long]()
+
+  class LogCleanerManagerMock(logDirs: Seq[File],
+                              logs: Pool[TopicPartition, Log],
+                              logDirFailureChannel: LogDirFailureChannel) extends LogCleanerManager(logDirs, logs, logDirFailureChannel) {
+    override def allCleanerCheckpoints: Map[TopicPartition, Long] = {
+      cleanerCheckpoints.toMap
+    }
+
+    override def updateCheckpoints(dataDir: File, partitionToUpdateOrAdd: Option[(TopicPartition,Long)] = None,
+                                   partitionToRemove: Option[TopicPartition] = None): Unit = {
+      assert(partitionToRemove.isEmpty, "partitionToRemove argument with value not yet handled")
+      val (tp, offset) = partitionToUpdateOrAdd.getOrElse(
+        throw new IllegalArgumentException("partitionToUpdateOrAdd==None argument not yet handled"))
+      cleanerCheckpoints.put(tp, offset)
+    }
+  }
 
   @After
   def tearDown(): Unit = {
     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 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 filthiestLog: LogToClean = cleanerManager.grabFilthiestCompactedLog(time).get
+    assertEquals(tp2, filthiestLog.topicPartition)
+    assertEquals(tp2, filthiestLog.log.topicPartition)
+  }
+
+  @Test
+  def testGrabFilthiestCompactedLogIgnoresUncleanablePartitions(): Unit = {
+    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))
+
+    cleanerManager.markPartitionUncleanable(logs.get(tp2).dir.getParent, tp2)
+
+    val filthiestLog: LogToClean = cleanerManager.grabFilthiestCompactedLog(time).get
+    assertEquals(tp1, filthiestLog.topicPartition)
+    assertEquals(tp1, filthiestLog.log.topicPartition)
+  }
+
+  @Test
+  def testGrabFilthiestCompactedLogIgnoresInProgressPartitions(): Unit = {
+    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))
+
+    cleanerManager.setCleaningState(tp2, LogCleaningInProgress)
+
+    val filthiestLog: LogToClean = cleanerManager.grabFilthiestCompactedLog(time).get
+
+    assertEquals(tp1, filthiestLog.topicPartition)
+    assertEquals(tp1, filthiestLog.log.topicPartition)
+  }
+
+  @Test
+  def testGrabFilthiestCompactedLogIgnoresBothInProgressPartitionsAndUncleanablePartitions(): Unit = {
+    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))
+
+    cleanerManager.setCleaningState(tp2, LogCleaningInProgress)
+    cleanerManager.markPartitionUncleanable(logs.get(tp1).dir.getParent, tp1)
+
+    val filthiestLog: Option[LogToClean] = cleanerManager.grabFilthiestCompactedLog(time)
+    assertEquals(None, filthiestLog)
+  }
+
+  @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, ClientRecordDeletion)
+
+    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, ClientRecordDeletion)
+
+    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)
+  }
+
   /**
     * When checking for logs with segments ready for deletion
     * we shouldn't consider logs where cleanup.policy=delete
@@ -76,18 +303,210 @@ class LogCleanerManagerTest extends JUnitSuite with Logging {
     assertEquals("should have 1 logs ready to be deleted", 1, readyToDelete)
   }
 
+  /**
+    * When looking for logs with segments ready to be deleted we should consider
+    * logs with cleanup.policy=compact because they may have segments from before the log start offset
+    */
+  @Test
+  def testLogsWithSegmentsToDeleteShouldConsiderCleanupPolicyCompactLogs(): Unit = {
+    val records = TestUtils.singletonRecords("test".getBytes, key="test".getBytes)
+    val log: Log = createLog(records.sizeInBytes * 5, LogConfig.Compact)
+    val cleanerManager: LogCleanerManager = createCleanerManager(log)
+
+    val readyToDelete = cleanerManager.deletableLogs().size
+    assertEquals("should have 1 logs ready to be deleted", 1, readyToDelete)
+  }
+
+  /**
+    * log under cleanup should be ineligible for compaction
+    */
+  @Test
+  def testLogsUnderCleanupIneligibleForCompaction(): Unit = {
+    val records = TestUtils.singletonRecords("test".getBytes, key="test".getBytes)
+    val log: Log = createLog(records.sizeInBytes * 5, LogConfig.Delete)
+    val cleanerManager: LogCleanerManager = createCleanerManager(log)
+
+    log.appendAsLeader(records, leaderEpoch = 0)
+    log.roll()
+    log.appendAsLeader(records, leaderEpoch = 0)
+    log.updateHighWatermark(2L)
+
+    // simulate cleanup thread working on the log partition
+    val deletableLog = cleanerManager.pauseCleaningForNonCompactedPartitions()
+    assertEquals("should have 1 logs ready to be deleted", 1, deletableLog.size)
+
+    // change cleanup policy from delete to compact
+    val logProps = new Properties()
+    logProps.put(LogConfig.SegmentBytesProp, log.config.segmentSize)
+    logProps.put(LogConfig.RetentionMsProp, log.config.retentionMs)
+    logProps.put(LogConfig.CleanupPolicyProp, LogConfig.Compact)
+    logProps.put(LogConfig.MinCleanableDirtyRatioProp, 0: Integer)
+    val config = LogConfig(logProps)
+    log.config = config
+
+    // log cleanup inprogress, the log is not available for compaction
+    val cleanable = cleanerManager.grabFilthiestCompactedLog(time)
+    assertEquals("should have 0 logs ready to be compacted", 0, cleanable.size)
+
+    // log cleanup finished, and log can be picked up for compaction
+    cleanerManager.resumeCleaning(deletableLog.map(_._1))
+    val cleanable2 = cleanerManager.grabFilthiestCompactedLog(time)
+    assertEquals("should have 1 logs ready to be compacted", 1, cleanable2.size)
+
+    // update cleanup policy to delete
+    logProps.put(LogConfig.CleanupPolicyProp, LogConfig.Delete)
+    val config2 = LogConfig(logProps)
+    log.config = config2
+
+    // compaction in progress, should have 0 log eligible for log cleanup
+    val deletableLog2 = cleanerManager.pauseCleaningForNonCompactedPartitions()
+    assertEquals("should have 0 logs ready to be deleted", 0, deletableLog2.size)
+
+    // compaction done, should have 1 log eligible for log cleanup
+    cleanerManager.doneDeleting(Seq(cleanable2.get.topicPartition))
+    val deletableLog3 = cleanerManager.pauseCleaningForNonCompactedPartitions()
+    assertEquals("should have 1 logs ready to be deleted", 1, deletableLog3.size)
+  }
+
+  @Test
+  def testUpdateCheckpointsShouldAddOffsetToPartition(): Unit = {
+    val records = TestUtils.singletonRecords("test".getBytes, key="test".getBytes)
+    val log: Log = createLog(records.sizeInBytes * 5, LogConfig.Compact)
+    val cleanerManager: LogCleanerManager = createCleanerManager(log)
+
+    // expect the checkpoint offset is not the expectedOffset before doing updateCheckpoints
+    assertNotEquals(offset, cleanerManager.allCleanerCheckpoints.get(topicPartition).getOrElse(0))
+
+    cleanerManager.updateCheckpoints(logDir, partitionToUpdateOrAdd = Option(topicPartition, offset))
+    // expect the checkpoint offset is now updated to the expected offset after doing updateCheckpoints
+    assertEquals(offset, cleanerManager.allCleanerCheckpoints(topicPartition))
+  }
+
+  @Test
+  def testUpdateCheckpointsShouldRemovePartitionData(): Unit = {
+    val records = TestUtils.singletonRecords("test".getBytes, key="test".getBytes)
+    val log: Log = createLog(records.sizeInBytes * 5, LogConfig.Compact)
+    val cleanerManager: LogCleanerManager = createCleanerManager(log)
+
+    // write some data into the cleaner-offset-checkpoint file
+    cleanerManager.updateCheckpoints(logDir, partitionToUpdateOrAdd = Option(topicPartition, offset))
+    assertEquals(offset, cleanerManager.allCleanerCheckpoints(topicPartition))
+
+    // updateCheckpoints should remove the topicPartition data in the logDir
+    cleanerManager.updateCheckpoints(logDir, partitionToRemove = Option(topicPartition))
+    assertTrue(cleanerManager.allCleanerCheckpoints.get(topicPartition).isEmpty)
+  }
+
+  @Test
+  def testHandleLogDirFailureShouldRemoveDirAndData(): Unit = {
+    val records = TestUtils.singletonRecords("test".getBytes, key="test".getBytes)
+    val log: Log = createLog(records.sizeInBytes * 5, LogConfig.Compact)
+    val cleanerManager: LogCleanerManager = createCleanerManager(log)
+
+    // write some data into the cleaner-offset-checkpoint file in logDir and logDir2
+    cleanerManager.updateCheckpoints(logDir, partitionToUpdateOrAdd = Option(topicPartition, offset))
+    cleanerManager.updateCheckpoints(logDir2, partitionToUpdateOrAdd = Option(topicPartition2, offset))
+    assertEquals(offset, cleanerManager.allCleanerCheckpoints(topicPartition))
+    assertEquals(offset, cleanerManager.allCleanerCheckpoints(topicPartition2))
+
+    cleanerManager.handleLogDirFailure(logDir.getAbsolutePath)
+    // verify the partition data in logDir is gone, and data in logDir2 is still there
+    assertEquals(offset, cleanerManager.allCleanerCheckpoints(topicPartition2))
+    assertTrue(cleanerManager.allCleanerCheckpoints.get(topicPartition).isEmpty)
+  }
+
+  @Test
+  def testMaybeTruncateCheckpointShouldTruncateData(): Unit = {
+    val records = TestUtils.singletonRecords("test".getBytes, key="test".getBytes)
+    val log: Log = createLog(records.sizeInBytes * 5, LogConfig.Compact)
+    val cleanerManager: LogCleanerManager = createCleanerManager(log)
+    val lowerOffset = 1L
+    val higherOffset = 1000L
+
+    // write some data into the cleaner-offset-checkpoint file in logDir
+    cleanerManager.updateCheckpoints(logDir, partitionToUpdateOrAdd = Option(topicPartition, offset))
+    assertEquals(offset, cleanerManager.allCleanerCheckpoints(topicPartition))
+
+    // we should not truncate the checkpoint data for checkpointed offset <= the given offset (higherOffset)
+    cleanerManager.maybeTruncateCheckpoint(logDir, topicPartition, higherOffset)
+    assertEquals(offset, cleanerManager.allCleanerCheckpoints(topicPartition))
+    // we should truncate the checkpoint data for checkpointed offset > the given offset (lowerOffset)
+    cleanerManager.maybeTruncateCheckpoint(logDir, topicPartition, lowerOffset)
+    assertEquals(lowerOffset, cleanerManager.allCleanerCheckpoints(topicPartition))
+  }
+
+  @Test
+  def testAlterCheckpointDirShouldRemoveDataInSrcDirAndAddInNewDir(): Unit = {
+    val records = TestUtils.singletonRecords("test".getBytes, key="test".getBytes)
+    val log: Log = createLog(records.sizeInBytes * 5, LogConfig.Compact)
+    val cleanerManager: LogCleanerManager = createCleanerManager(log)
+
+    // write some data into the cleaner-offset-checkpoint file in logDir
+    cleanerManager.updateCheckpoints(logDir, partitionToUpdateOrAdd = Option(topicPartition, offset))
+    assertEquals(offset, cleanerManager.allCleanerCheckpoints(topicPartition))
+
+    cleanerManager.alterCheckpointDir(topicPartition, logDir, logDir2)
+    // verify we still can get the partition offset after alterCheckpointDir
+    // This data should locate in logDir2, not logDir
+    assertEquals(offset, cleanerManager.allCleanerCheckpoints(topicPartition))
+
+    // force delete the logDir2 from checkpoints, so that the partition data should also be deleted
+    cleanerManager.handleLogDirFailure(logDir2.getAbsolutePath)
+    assertTrue(cleanerManager.allCleanerCheckpoints.get(topicPartition).isEmpty)
+  }
+
+  /**
+    * log under cleanup should still be eligible for log truncation
+    */
+  @Test
+  def testConcurrentLogCleanupAndLogTruncation(): Unit = {
+    val records = TestUtils.singletonRecords("test".getBytes, key="test".getBytes)
+    val log: Log = createLog(records.sizeInBytes * 5, LogConfig.Delete)
+    val cleanerManager: LogCleanerManager = createCleanerManager(log)
+
+    // log cleanup starts
+    val pausedPartitions = cleanerManager.pauseCleaningForNonCompactedPartitions()
+    // Log truncation happens due to unclean leader election
+    cleanerManager.abortAndPauseCleaning(log.topicPartition)
+    cleanerManager.resumeCleaning(Seq(log.topicPartition))
+    // log cleanup finishes and pausedPartitions are resumed
+    cleanerManager.resumeCleaning(pausedPartitions.map(_._1))
+
+    assertEquals(None, cleanerManager.cleaningState(log.topicPartition))
+  }
+
+  /**
+    * log under cleanup should still be eligible for topic deletion
+    */
+  @Test
+  def testConcurrentLogCleanupAndTopicDeletion(): Unit = {
+    val records = TestUtils.singletonRecords("test".getBytes, key = "test".getBytes)
+    val log: Log = createLog(records.sizeInBytes * 5, LogConfig.Delete)
+    val cleanerManager: LogCleanerManager = createCleanerManager(log)
+
+    // log cleanup starts
+    val pausedPartitions = cleanerManager.pauseCleaningForNonCompactedPartitions()
+    // Broker processes StopReplicaRequest with delete=true
+    cleanerManager.abortCleaning(log.topicPartition)
+    // log cleanup finishes and pausedPartitions are resumed
+    cleanerManager.resumeCleaning(pausedPartitions.map(_._1))
+
+    assertEquals(None, cleanerManager.cleaningState(log.topicPartition))
+  }
+
   /**
     * When looking for logs with segments ready to be deleted we shouldn't consider
-    * logs with cleanup.policy=compact as they shouldn't have segments truncated.
+    * logs that have had their partition marked as uncleanable.
     */
   @Test
-  def testLogsWithSegmentsToDeleteShouldNotConsiderCleanupPolicyCompactLogs(): Unit = {
+  def testLogsWithSegmentsToDeleteShouldNotConsiderUncleanablePartitions(): Unit = {
     val records = TestUtils.singletonRecords("test".getBytes, key="test".getBytes)
     val log: Log = createLog(records.sizeInBytes * 5, LogConfig.Compact)
     val cleanerManager: LogCleanerManager = createCleanerManager(log)
+    cleanerManager.markPartitionUncleanable(log.dir.getParent, topicPartition)
 
     val readyToDelete = cleanerManager.deletableLogs().size
-    assertEquals("should have 1 logs ready to be deleted", 0, readyToDelete)
+    assertEquals("should have 0 logs ready to be deleted", 0, readyToDelete)
   }
 
   /**
@@ -103,11 +522,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 topicPartition = new TopicPartition("log", 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)
   }
 
   /**
@@ -134,11 +552,10 @@ class LogCleanerManagerTest extends JUnitSuite with Logging {
     while (log.numberOfSegments < 8)
       log.appendAsLeader(records(log.logEndOffset.toInt, log.logEndOffset.toInt, t1), leaderEpoch = 0)
 
-    val topicPartition = new TopicPartition("log", 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)
   }
 
   /**
@@ -160,16 +577,33 @@ class LogCleanerManagerTest extends JUnitSuite with Logging {
 
     time.sleep(compactionLag + 1)
 
-    val topicPartition = new TopicPartition("log", 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 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, ClientRecordDeletion)
+
+    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, ClientRecordDeletion)
+    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
   def testUndecidedTransactionalDataNotCleanable(): Unit = {
-    val topicPartition = new TopicPartition("log", 0)
     val compactionLag = 60 * 60 * 1000
     val logProps = new Properties()
     logProps.put(LogConfig.SegmentBytesProp, 1024: java.lang.Integer)
@@ -186,51 +620,136 @@ 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
+  def testDoneCleaning(): Unit = {
+    val logProps = new Properties()
+    logProps.put(LogConfig.SegmentBytesProp, 1024: java.lang.Integer)
+    val log = makeLog(config = LogConfig.fromProps(logConfig.originals, logProps))
+    while(log.numberOfSegments < 8)
+      log.appendAsLeader(records(log.logEndOffset.toInt, log.logEndOffset.toInt, time.milliseconds()), leaderEpoch = 0)
+
+    val cleanerManager: LogCleanerManager = createCleanerManager(log)
+
+    intercept[IllegalStateException](cleanerManager.doneCleaning(topicPartition, log.dir, 1))
+
+    cleanerManager.setCleaningState(topicPartition, LogCleaningPaused(1))
+    intercept[IllegalStateException](cleanerManager.doneCleaning(topicPartition, log.dir, 1))
+
+    cleanerManager.setCleaningState(topicPartition, LogCleaningInProgress)
+    cleanerManager.doneCleaning(topicPartition, log.dir, 1)
+    assertTrue(cleanerManager.cleaningState(topicPartition).isEmpty)
+    assertTrue(cleanerManager.allCleanerCheckpoints.get(topicPartition).nonEmpty)
+
+    cleanerManager.setCleaningState(topicPartition, LogCleaningAborted)
+    cleanerManager.doneCleaning(topicPartition, log.dir, 1)
+    assertEquals(LogCleaningPaused(1), cleanerManager.cleaningState(topicPartition).get)
+    assertTrue(cleanerManager.allCleanerCheckpoints.get(topicPartition).nonEmpty)
+  }
+
+  @Test
+  def testDoneDeleting(): Unit = {
+    val records = TestUtils.singletonRecords("test".getBytes, key="test".getBytes)
+    val log: Log = createLog(records.sizeInBytes * 5, LogConfig.Compact + "," + LogConfig.Delete)
+    val cleanerManager: LogCleanerManager = createCleanerManager(log)
+    val tp = new TopicPartition("log", 0)
+
+    intercept[IllegalStateException](cleanerManager.doneDeleting(Seq(tp)))
+
+    cleanerManager.setCleaningState(tp, LogCleaningPaused(1))
+    intercept[IllegalStateException](cleanerManager.doneDeleting(Seq(tp)))
+
+    cleanerManager.setCleaningState(tp, LogCleaningInProgress)
+    cleanerManager.doneDeleting(Seq(tp))
+    assertTrue(cleanerManager.cleaningState(tp).isEmpty)
+
+    cleanerManager.setCleaningState(tp, LogCleaningAborted)
+    cleanerManager.doneDeleting(Seq(tp))
+    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, ClientRecordDeletion)
+    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, ClientRecordDeletion)
+    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(new TopicPartition("log", 0), log)
-    val cleanerManager = new LogCleanerManager(Array(logDir), logs, null)
-    cleanerManager
+    logs.put(topicPartition, log)
+    new LogCleanerManager(Seq(logDir, logDir2), logs, null)
   }
 
-  private def createLog(segmentSize: Int, cleanupPolicy: String): Log = {
-    val logProps = new Properties()
-    logProps.put(LogConfig.SegmentBytesProp, segmentSize: Integer)
-    logProps.put(LogConfig.RetentionMsProp, 1: Integer)
-    logProps.put(LogConfig.CleanupPolicyProp, cleanupPolicy)
+  private def createCleanerManagerMock(pool: Pool[TopicPartition, Log]): LogCleanerManagerMock = {
+    new LogCleanerManagerMock(Seq(logDir), pool, null)
+  }
 
-    val config = LogConfig(logProps)
-    val partitionDir = new File(logDir, "log-0")
-    val log = Log(partitionDir,
+  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))
+
+    Log(partitionDir,
       config,
       logStartOffset = 0L,
       recoveryPoint = 0L,
@@ -240,7 +759,43 @@ class LogCleanerManagerTest extends JUnitSuite with Logging {
       maxProducerIdExpirationMs = 60 * 60 * 1000,
       producerIdExpirationCheckIntervalMs = LogManager.ProducerIdExpirationCheckIntervalMs,
       logDirFailureChannel = new LogDirFailureChannel(10))
-    log
+  }
+
+  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.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
new file mode 100755
index 0000000000000..8adb4d1f680ed
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/log/LogCleanerParameterizedIntegrationTest.scala
@@ -0,0 +1,335 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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.util.Properties
+
+import kafka.api.KAFKA_0_11_0_IV0
+import kafka.api.{KAFKA_0_10_0_IV1, KAFKA_0_9_0}
+import kafka.server.KafkaConfig
+import kafka.server.checkpoints.OffsetCheckpointFile
+import kafka.utils._
+import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.record._
+import org.junit.Assert._
+import org.junit._
+import org.junit.runner.RunWith
+import org.junit.runners.Parameterized
+import org.junit.runners.Parameterized.Parameters
+
+import scala.collection._
+import scala.jdk.CollectionConverters._
+
+/**
+ * This is an integration test that tests the fully integrated log cleaner
+ */
+@RunWith(value = classOf[Parameterized])
+class LogCleanerParameterizedIntegrationTest(compressionCodec: String) extends AbstractLogCleanerIntegrationTest {
+
+  val codec: CompressionType = CompressionType.forName(compressionCodec)
+  val time = new MockTime()
+
+  val topicPartitions = Array(new TopicPartition("log", 0), new TopicPartition("log", 1), new TopicPartition("log", 2))
+
+
+  @Test
+  def cleanerTest(): Unit = {
+    val largeMessageKey = 20
+    val (largeMessageValue, largeMessageSet) = createLargeSingleMessageSet(largeMessageKey, RecordBatch.CURRENT_MAGIC_VALUE)
+    val maxMessageSize = largeMessageSet.sizeInBytes
+
+    cleaner = makeCleaner(partitions = topicPartitions, maxMessageSize = maxMessageSize)
+    val log = cleaner.logs.get(topicPartitions(0))
+
+    val appends = writeDups(numKeys = 100, numDups = 3, log = log, codec = codec)
+    val startSize = log.size
+    cleaner.startup()
+
+    val firstDirty = log.activeSegment.baseOffset
+    checkLastCleaned("log", 0, firstDirty)
+    val compactedSize = log.logSegments.map(_.size).sum
+    assertTrue(s"log should have been compacted: startSize=$startSize compactedSize=$compactedSize", startSize > compactedSize)
+
+    checkLogAfterAppendingDups(log, startSize, appends)
+
+    val appendInfo = log.appendAsLeader(largeMessageSet, leaderEpoch = 0)
+    val largeMessageOffset = appendInfo.firstOffset.get
+
+    val dups = writeDups(startKey = largeMessageKey + 1, numKeys = 100, numDups = 3, log = log, codec = codec)
+    val appends2 = appends ++ Seq((largeMessageKey, largeMessageValue, largeMessageOffset)) ++ dups
+    val firstDirty2 = log.activeSegment.baseOffset
+    checkLastCleaned("log", 0, firstDirty2)
+
+    checkLogAfterAppendingDups(log, startSize, appends2)
+
+    // simulate deleting a partition, by removing it from logs
+    // force a checkpoint
+    // and make sure its gone from checkpoint file
+    cleaner.logs.remove(topicPartitions(0))
+    cleaner.updateCheckpoints(logDir, partitionToRemove = Option(topicPartitions(0)))
+    val checkpoints = new OffsetCheckpointFile(new File(logDir, cleaner.cleanerManager.offsetCheckpointFile)).read()
+    // we expect partition 0 to be gone
+    assertFalse(checkpoints.contains(topicPartitions(0)))
+  }
+
+  @Test
+  def testCleansCombinedCompactAndDeleteTopic(): Unit = {
+    val logProps  = new Properties()
+    val retentionMs: Integer = 100000
+    logProps.put(LogConfig.RetentionMsProp, retentionMs: Integer)
+    logProps.put(LogConfig.CleanupPolicyProp, "compact,delete")
+
+    def runCleanerAndCheckCompacted(numKeys: Int): (Log, Seq[(Int, String, Long)]) = {
+      cleaner = makeCleaner(partitions = topicPartitions.take(1), propertyOverrides = logProps, backOffMs = 100L)
+      val log = cleaner.logs.get(topicPartitions(0))
+
+      val messages = writeDups(numKeys = numKeys, numDups = 3, log = log, codec = codec)
+      val startSize = log.size
+
+      log.updateHighWatermark(log.logEndOffset)
+
+      val firstDirty = log.activeSegment.baseOffset
+      cleaner.startup()
+
+      // should compact the log
+      checkLastCleaned("log", 0, firstDirty)
+      val compactedSize = log.logSegments.map(_.size).sum
+      assertTrue(s"log should have been compacted: startSize=$startSize compactedSize=$compactedSize", startSize > compactedSize)
+      (log, messages)
+    }
+
+    val (log, _) = runCleanerAndCheckCompacted(100)
+
+    // 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()
+
+    // run the cleaner again to make sure if there are no issues post deletion
+    val (log2, messages) = runCleanerAndCheckCompacted(20)
+    val read = readFromLog(log2)
+    assertEquals("Contents of the map shouldn't change", toMap(messages), toMap(read))
+  }
+
+  @Test
+  def testCleanerWithMessageFormatV0(): Unit = {
+    // zstd compression is not supported with older message formats
+    if (codec == CompressionType.ZSTD)
+      return
+
+    val largeMessageKey = 20
+    val (largeMessageValue, largeMessageSet) = createLargeSingleMessageSet(largeMessageKey, RecordBatch.MAGIC_VALUE_V0)
+    val maxMessageSize = codec match {
+      case CompressionType.NONE => largeMessageSet.sizeInBytes
+      case _ =>
+        // the broker assigns absolute offsets for message format 0 which potentially causes the compressed size to
+        // increase because the broker offsets are larger than the ones assigned by the client
+        // adding `5` to the message set size is good enough for this test: it covers the increased message size while
+        // still being less than the overhead introduced by the conversion from message format version 0 to 1
+        largeMessageSet.sizeInBytes + 5
+    }
+
+    cleaner = makeCleaner(partitions = topicPartitions, maxMessageSize = maxMessageSize)
+
+    val log = cleaner.logs.get(topicPartitions(0))
+    val props = logConfigProperties(maxMessageSize = maxMessageSize)
+    props.put(LogConfig.MessageFormatVersionProp, KAFKA_0_9_0.version)
+    log.config = new LogConfig(props)
+
+    val appends = writeDups(numKeys = 100, numDups = 3, log = log, codec = codec, magicValue = RecordBatch.MAGIC_VALUE_V0)
+    val startSize = log.size
+    cleaner.startup()
+
+    val firstDirty = log.activeSegment.baseOffset
+    checkLastCleaned("log", 0, firstDirty)
+    val compactedSize = log.logSegments.map(_.size).sum
+    assertTrue(s"log should have been compacted: startSize=$startSize compactedSize=$compactedSize", startSize > compactedSize)
+
+    checkLogAfterAppendingDups(log, startSize, appends)
+
+    val appends2: Seq[(Int, String, Long)] = {
+      val dupsV0 = writeDups(numKeys = 40, numDups = 3, log = log, codec = codec, magicValue = RecordBatch.MAGIC_VALUE_V0)
+      val appendInfo = log.appendAsLeader(largeMessageSet, leaderEpoch = 0)
+      val largeMessageOffset = appendInfo.firstOffset.get
+
+      // also add some messages with version 1 and version 2 to check that we handle mixed format versions correctly
+      props.put(LogConfig.MessageFormatVersionProp, KAFKA_0_11_0_IV0.version)
+      log.config = new LogConfig(props)
+      val dupsV1 = writeDups(startKey = 30, numKeys = 40, numDups = 3, log = log, codec = codec, magicValue = RecordBatch.MAGIC_VALUE_V1)
+      val dupsV2 = writeDups(startKey = 15, numKeys = 5, numDups = 3, log = log, codec = codec, magicValue = RecordBatch.MAGIC_VALUE_V2)
+      appends ++ dupsV0 ++ Seq((largeMessageKey, largeMessageValue, largeMessageOffset)) ++ dupsV1 ++ dupsV2
+    }
+    val firstDirty2 = log.activeSegment.baseOffset
+    checkLastCleaned("log", 0, firstDirty2)
+
+    checkLogAfterAppendingDups(log, startSize, appends2)
+  }
+
+  @Test
+  def testCleaningNestedMessagesWithMultipleVersions(): Unit = {
+    // zstd compression is not supported with older message formats
+    if (codec == CompressionType.ZSTD)
+      return
+
+    val maxMessageSize = 192
+    cleaner = makeCleaner(partitions = topicPartitions, maxMessageSize = maxMessageSize, segmentSize = 256)
+
+    val log = cleaner.logs.get(topicPartitions(0))
+    val props = logConfigProperties(maxMessageSize = maxMessageSize, segmentSize = 256)
+    props.put(LogConfig.MessageFormatVersionProp, KAFKA_0_9_0.version)
+    log.config = new LogConfig(props)
+
+    // with compression enabled, these messages will be written as a single message containing
+    // all of the individual messages
+    var appendsV0 = writeDupsSingleMessageSet(numKeys = 2, numDups = 3, log = log, codec = codec, magicValue = RecordBatch.MAGIC_VALUE_V0)
+    appendsV0 ++= writeDupsSingleMessageSet(numKeys = 2, startKey = 3, numDups = 2, log = log, codec = codec, magicValue = RecordBatch.MAGIC_VALUE_V0)
+
+    props.put(LogConfig.MessageFormatVersionProp, KAFKA_0_10_0_IV1.version)
+    log.config = new LogConfig(props)
+
+    var appendsV1 = writeDupsSingleMessageSet(startKey = 4, numKeys = 2, numDups = 2, log = log, codec = codec, magicValue = RecordBatch.MAGIC_VALUE_V1)
+    appendsV1 ++= writeDupsSingleMessageSet(startKey = 4, numKeys = 2, numDups = 2, log = log, codec = codec, magicValue = RecordBatch.MAGIC_VALUE_V1)
+    appendsV1 ++= writeDupsSingleMessageSet(startKey = 6, numKeys = 2, numDups = 2, log = log, codec = codec, magicValue = RecordBatch.MAGIC_VALUE_V1)
+
+    val appends = appendsV0 ++ appendsV1
+
+    val startSize = log.size
+    cleaner.startup()
+
+    val firstDirty = log.activeSegment.baseOffset
+    assertTrue(firstDirty > appendsV0.size) // ensure we clean data from V0 and V1
+
+    checkLastCleaned("log", 0, firstDirty)
+    val compactedSize = log.logSegments.map(_.size).sum
+    assertTrue(s"log should have been compacted: startSize=$startSize compactedSize=$compactedSize", startSize > compactedSize)
+
+    checkLogAfterAppendingDups(log, startSize, appends)
+  }
+
+  @Test
+  def cleanerConfigUpdateTest(): Unit = {
+    val largeMessageKey = 20
+    val (largeMessageValue, largeMessageSet) = createLargeSingleMessageSet(largeMessageKey, RecordBatch.CURRENT_MAGIC_VALUE)
+    val maxMessageSize = largeMessageSet.sizeInBytes
+
+    cleaner = makeCleaner(partitions = topicPartitions, backOffMs = 1, maxMessageSize = maxMessageSize,
+      cleanerIoBufferSize = Some(1))
+    val log = cleaner.logs.get(topicPartitions(0))
+
+    writeDups(numKeys = 100, numDups = 3, log = log, codec = codec)
+    val startSize = log.size
+    cleaner.startup()
+    assertEquals(1, cleaner.cleanerCount)
+
+    // Verify no cleaning with LogCleanerIoBufferSizeProp=1
+    val firstDirty = log.activeSegment.baseOffset
+    val topicPartition = new TopicPartition("log", 0)
+    cleaner.awaitCleaned(topicPartition, firstDirty, maxWaitMs = 10)
+    assertTrue("Should not have cleaned", cleaner.cleanerManager.allCleanerCheckpoints.isEmpty)
+
+    def kafkaConfigWithCleanerConfig(cleanerConfig: CleanerConfig): KafkaConfig = {
+      val props = TestUtils.createBrokerConfig(0, "localhost:2181")
+      props.put(KafkaConfig.LogCleanerThreadsProp, cleanerConfig.numThreads.toString)
+      props.put(KafkaConfig.LogCleanerDedupeBufferSizeProp, cleanerConfig.dedupeBufferSize.toString)
+      props.put(KafkaConfig.LogCleanerDedupeBufferLoadFactorProp, cleanerConfig.dedupeBufferLoadFactor.toString)
+      props.put(KafkaConfig.LogCleanerIoBufferSizeProp, cleanerConfig.ioBufferSize.toString)
+      props.put(KafkaConfig.MessageMaxBytesProp, cleanerConfig.maxMessageSize.toString)
+      props.put(KafkaConfig.LogCleanerBackoffMsProp, cleanerConfig.backOffMs.toString)
+      props.put(KafkaConfig.LogCleanerIoMaxBytesPerSecondProp, cleanerConfig.maxIoBytesPerSecond.toString)
+      KafkaConfig.fromProps(props)
+    }
+
+    // Verify cleaning done with larger LogCleanerIoBufferSizeProp
+    val oldConfig = kafkaConfigWithCleanerConfig(cleaner.currentConfig)
+    val newConfig = kafkaConfigWithCleanerConfig(CleanerConfig(numThreads = 2,
+      dedupeBufferSize = cleaner.currentConfig.dedupeBufferSize,
+      dedupeBufferLoadFactor = cleaner.currentConfig.dedupeBufferLoadFactor,
+      ioBufferSize = 100000,
+      maxMessageSize = cleaner.currentConfig.maxMessageSize,
+      maxIoBytesPerSecond = cleaner.currentConfig.maxIoBytesPerSecond,
+      backOffMs = cleaner.currentConfig.backOffMs))
+    cleaner.reconfigure(oldConfig, newConfig)
+
+    assertEquals(2, cleaner.cleanerCount)
+    checkLastCleaned("log", 0, firstDirty)
+    val compactedSize = log.logSegments.map(_.size).sum
+    assertTrue(s"log should have been compacted: startSize=$startSize compactedSize=$compactedSize", startSize > compactedSize)
+  }
+
+  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)
+    cleaner.awaitCleaned(topicPartition, firstDirty)
+    val lastCleaned = cleaner.cleanerManager.allCleanerCheckpoints(topicPartition)
+    assertTrue(s"log cleaner should have processed up to offset $firstDirty, but lastCleaned=$lastCleaned",
+      lastCleaned >= firstDirty)
+  }
+
+  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)
+  }
+
+  private def toMap(messages: Iterable[(Int, String, Long)]): Map[Int, (String, Long)] = {
+    messages.map { case (key, value, offset) => key -> (value, offset) }.toMap
+  }
+
+  private def readFromLog(log: Log): Iterable[(Int, String, Long)] = {
+    for (segment <- log.logSegments; deepLogEntry <- segment.log.records.asScala) yield {
+      val key = TestUtils.readString(deepLogEntry.key).toInt
+      val value = TestUtils.readString(deepLogEntry.value)
+      (key, value, deepLogEntry.offset)
+    }
+  }
+
+  private def writeDupsSingleMessageSet(numKeys: Int, numDups: Int, log: Log, codec: CompressionType,
+                                        startKey: Int = 0, magicValue: Byte): Seq[(Int, String, Long)] = {
+    val kvs = for (_ <- 0 until numDups; key <- startKey until (startKey + numKeys)) yield {
+      val payload = counter.toString
+      incCounter()
+      (key, payload)
+    }
+
+    val records = kvs.map { case (key, payload) =>
+      new SimpleRecord(key.toString.getBytes, payload.toString.getBytes)
+    }
+
+    val appendInfo = log.appendAsLeader(MemoryRecords.withRecords(magicValue, codec, records: _*), leaderEpoch = 0)
+    val offsets = appendInfo.firstOffset.get to appendInfo.lastOffset
+
+    kvs.zip(offsets).map { case (kv, offset) => (kv._1, kv._2, offset) }
+  }
+
+}
+
+object LogCleanerParameterizedIntegrationTest {
+  @Parameters
+  def parameters: java.util.Collection[Array[String]] = {
+    val list = new java.util.ArrayList[Array[String]]()
+    for (codec <- CompressionType.values)
+      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 517e8760a870f..3b1584d8b8037 100755
--- a/core/src/test/scala/unit/kafka/log/LogCleanerTest.scala
+++ b/core/src/test/scala/unit/kafka/log/LogCleanerTest.scala
@@ -17,28 +17,31 @@
 
 package kafka.log
 
-import java.io.File
+import java.io.{File, RandomAccessFile}
 import java.nio._
+import java.nio.charset.StandardCharsets
 import java.nio.file.Paths
 import java.util.Properties
+import java.util.concurrent.{CountDownLatch, TimeUnit}
 
 import kafka.common._
 import kafka.server.{BrokerTopicStats, LogDirFailureChannel}
 import kafka.utils._
 import org.apache.kafka.common.TopicPartition
+import org.apache.kafka.common.errors.CorruptRecordException
 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._
+import scala.jdk.CollectionConverters._
 
 /**
  * Unit tests for the log cleaning logic
  */
-class LogCleanerTest extends JUnitSuite {
+class LogCleanerTest {
 
   val tmpdir = TestUtils.tempDir()
   val dir = TestUtils.randomPartitionLogDir(tmpdir)
@@ -70,11 +73,11 @@ class LogCleanerTest extends JUnitSuite {
     // append messages to the log until we have four segments
     while(log.numberOfSegments < 4)
       log.appendAsLeader(record(log.logEndOffset.toInt, log.logEndOffset.toInt), leaderEpoch = 0)
-    val keysFound = keysInLog(log)
+    val keysFound = LogTest.keysInLog(log)
     assertEquals(0L until log.logEndOffset, keysFound)
 
     // pretend we have the following keys
-    val keys = immutable.ListSet(1, 3, 5, 7, 9)
+    val keys = immutable.ListSet(1L, 3L, 5L, 7L, 9L)
     val map = new FakeOffsetMap(Int.MaxValue)
     keys.foreach(k => map.put(key(k), Long.MaxValue))
 
@@ -82,12 +85,81 @@ 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)
-    val shouldRemain = keysInLog(log).filter(!keys.contains(_))
-    assertEquals(shouldRemain, keysInLog(log))
+    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)
   }
 
+  @Test
+  def testCleanSegmentsWithConcurrentSegmentDeletion(): Unit = {
+    val deleteStartLatch = new CountDownLatch(1)
+    val deleteCompleteLatch = new CountDownLatch(1)
+
+    // Construct a log instance. The replaceSegments() method of the log instance is overridden so that
+    // it waits for another thread to execute deleteOldSegments()
+    val logProps = new Properties()
+    logProps.put(LogConfig.SegmentBytesProp, 1024 : java.lang.Integer)
+    logProps.put(LogConfig.CleanupPolicyProp, LogConfig.Compact + "," + LogConfig.Delete)
+    val topicPartition = Log.parseTopicPartitionName(dir)
+    val producerStateManager = new ProducerStateManager(topicPartition, dir)
+    val log = new Log(dir,
+                      config = LogConfig.fromProps(logConfig.originals, logProps),
+                      logStartOffset = 0L,
+                      recoveryPoint = 0L,
+                      scheduler = time.scheduler,
+                      brokerTopicStats = new BrokerTopicStats, time,
+                      maxProducerIdExpirationMs = 60 * 60 * 1000,
+                      producerIdExpirationCheckIntervalMs = LogManager.ProducerIdExpirationCheckIntervalMs,
+                      topicPartition = topicPartition,
+                      producerStateManager = producerStateManager,
+                      logDirFailureChannel = new LogDirFailureChannel(10)) {
+      override def replaceSegments(newSegments: Seq[LogSegment], oldSegments: Seq[LogSegment], isRecoveredSwapFile: Boolean = false): Unit = {
+        deleteStartLatch.countDown()
+        if (!deleteCompleteLatch.await(5000, TimeUnit.MILLISECONDS)) {
+          throw new IllegalStateException("Log segment deletion timed out")
+        }
+        super.replaceSegments(newSegments, oldSegments, isRecoveredSwapFile)
+      }
+    }
+
+    // Start a thread which execute log.deleteOldSegments() right before replaceSegments() is executed
+    val t = new Thread() {
+      override def run(): Unit = {
+        deleteStartLatch.await(5000, TimeUnit.MILLISECONDS)
+        log.updateHighWatermark(log.activeSegment.baseOffset)
+        log.maybeIncrementLogStartOffset(log.activeSegment.baseOffset, LeaderOffsetIncremented)
+        log.updateHighWatermark(log.activeSegment.baseOffset)
+        log.deleteOldSegments()
+        deleteCompleteLatch.countDown()
+      }
+    }
+    t.start()
+
+    // Append records so that segment number increase to 3
+    while (log.numberOfSegments < 3) {
+      log.appendAsLeader(record(key = 0, log.logEndOffset.toInt), leaderEpoch = 0)
+      log.roll()
+    }
+    assertEquals(3, log.numberOfSegments)
+
+    // Remember reference to the first log and determine its file name expected for async deletion
+    val firstLogFile = log.logSegments.head.log
+    val expectedFileName = CoreUtils.replaceSuffix(firstLogFile.file.getPath, "", Log.DeletedFileSuffix)
+
+    // Clean the log. This should trigger replaceSegments() and deleteOldSegments();
+    val offsetMap = new FakeOffsetMap(Int.MaxValue)
+    val cleaner = makeCleaner(Int.MaxValue)
+    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, new CleanedTransactionMetadata)
+
+    // Validate based on the file name that log segment file is renamed exactly once for async deletion
+    assertEquals(expectedFileName, firstLogFile.file().getPath)
+    assertEquals(2, log.numberOfSegments)
+  }
+
   @Test
   def testSizeTrimmedForPreallocatedAndCompactedTopic(): Unit = {
     val originalMaxFileSize = 1024;
@@ -110,7 +182,7 @@ class LogCleanerTest extends JUnitSuite {
     cleaner.clean(LogToClean(new TopicPartition("test", 0), log, 2, log.activeSegment.baseOffset))
 
     assertTrue("Cleaned segment file should be trimmed to its real size.",
-      log.logSegments.iterator.next.log.channel().size() < originalMaxFileSize)
+      log.logSegments.iterator.next().log.channel.size < originalMaxFileSize)
   }
 
   @Test
@@ -134,7 +206,7 @@ class LogCleanerTest extends JUnitSuite {
     cleaner.clean(LogToClean(new TopicPartition("test", 0), log, 0L, log.activeSegment.baseOffset))
     assertEquals(List(2, 5, 7), lastOffsetsPerBatchInLog(log))
     assertEquals(Map(pid1 -> 2, pid2 -> 2, pid3 -> 1), lastSequencesInLog(log))
-    assertEquals(List(2, 3, 1, 4), keysInLog(log))
+    assertEquals(List(2, 3, 1, 4), LogTest.keysInLog(log))
     assertEquals(List(1, 3, 6, 7), offsetsInLog(log))
 
     // we have to reload the log to validate that the cleaner maintained sequence numbers correctly
@@ -147,17 +219,17 @@ class LogCleanerTest extends JUnitSuite {
 
     // check duplicate append from producer 1
     var logAppendInfo = appendIdempotentAsLeader(log, pid1, producerEpoch)(Seq(1, 2, 3))
-    assertEquals(0L, logAppendInfo.firstOffset)
+    assertEquals(0L, logAppendInfo.firstOffset.get)
     assertEquals(2L, logAppendInfo.lastOffset)
 
     // check duplicate append from producer 3
     logAppendInfo = appendIdempotentAsLeader(log, pid3, producerEpoch)(Seq(1, 4))
-    assertEquals(6L, logAppendInfo.firstOffset)
+    assertEquals(6L, logAppendInfo.firstOffset.get)
     assertEquals(7L, logAppendInfo.lastOffset)
 
     // check duplicate append from producer 2
     logAppendInfo = appendIdempotentAsLeader(log, pid2, producerEpoch)(Seq(3, 1, 4))
-    assertEquals(3L, logAppendInfo.firstOffset)
+    assertEquals(3L, logAppendInfo.firstOffset.get)
     assertEquals(5L, logAppendInfo.lastOffset)
 
     // do one more append and a round of cleaning to force another deletion from producer 1's batch
@@ -166,14 +238,14 @@ class LogCleanerTest extends JUnitSuite {
     cleaner.clean(LogToClean(new TopicPartition("test", 0), log, 0L, log.activeSegment.baseOffset))
     assertEquals(Map(pid1 -> 2, pid2 -> 2, pid3 -> 1, pid4 -> 0), lastSequencesInLog(log))
     assertEquals(List(2, 5, 7, 8), lastOffsetsPerBatchInLog(log))
-    assertEquals(List(3, 1, 4, 2), keysInLog(log))
+    assertEquals(List(3, 1, 4, 2), LogTest.keysInLog(log))
     assertEquals(List(3, 6, 7, 8), offsetsInLog(log))
 
     reloadLog()
 
     // duplicate append from producer1 should still be fine
     logAppendInfo = appendIdempotentAsLeader(log, pid1, producerEpoch)(Seq(1, 2, 3))
-    assertEquals(0L, logAppendInfo.firstOffset)
+    assertEquals(0L, logAppendInfo.firstOffset.get)
     assertEquals(2L, logAppendInfo.lastOffset)
   }
 
@@ -194,16 +266,16 @@ 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)
 
     log.roll()
     cleaner.clean(LogToClean(new TopicPartition("test", 0), log, 0L, log.activeSegment.baseOffset))
-    assertEquals(List(3, 2), keysInLog(log))
+    assertEquals(List(3, 2), LogTest.keysInLog(log))
     assertEquals(List(3, 6, 7, 8, 9), offsetsInLog(log))
 
     // ensure the transaction index is still correct
@@ -235,15 +307,15 @@ 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
-    assertEquals(List(2, 3, 5, 6, 6, 7, 7, 8, 8, 9, 9, 10), keysInLog(log))
+    assertEquals(List(2, 3, 5, 6, 6, 7, 7, 8, 8, 9, 9, 10), LogTest.keysInLog(log))
 
     log.roll()
 
@@ -253,7 +325,7 @@ class LogCleanerTest extends JUnitSuite {
 
     // finally only the keys from pid3 should remain
     cleaner.clean(LogToClean(new TopicPartition("test", 0), log, dirtyOffset, log.activeSegment.baseOffset))
-    assertEquals(List(2, 3, 6, 7, 8, 9, 11, 12), keysInLog(log))
+    assertEquals(List(2, 3, 6, 7, 8, 9, 11, 12), LogTest.keysInLog(log))
   }
 
   @Test
@@ -270,36 +342,69 @@ 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
-    assertEquals(List(1, 3, 2), keysInLog(log))
+    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
-    assertEquals(List(2, 1, 3), keysInLog(log))
+    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
-    assertEquals(List(2, 1, 3), keysInLog(log))
+    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
-    assertEquals(List(2, 1, 3), keysInLog(log))
+    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))
   }
 
+  /**
+   * Tests log cleaning with batches that are deleted where no additional messages
+   * are available to read in the buffer. Cleaning should continue from the next offset.
+   */
+  @Test
+  def testDeletedBatchesWithNoMessagesRead(): Unit = {
+    val tp = new TopicPartition("test", 0)
+    val cleaner = makeCleaner(capacity = Int.MaxValue, maxMessageSize = 100)
+    val logProps = new Properties()
+    logProps.put(LogConfig.MaxMessageBytesProp, 100: java.lang.Integer)
+    logProps.put(LogConfig.SegmentBytesProp, 1000: java.lang.Integer)
+    val log = makeLog(config = LogConfig.fromProps(logConfig.originals, logProps))
+
+    val producerEpoch = 0.toShort
+    val producerId = 1L
+    val appendProducer = appendTransactionalAsLeader(log, producerId, producerEpoch)
+
+    appendProducer(Seq(1))
+    log.appendAsLeader(abortMarker(producerId, producerEpoch), leaderEpoch = 0, origin = AppendOrigin.Coordinator)
+    appendProducer(Seq(2))
+    appendProducer(Seq(2))
+    log.appendAsLeader(commitMarker(producerId, producerEpoch), leaderEpoch = 0, origin = AppendOrigin.Coordinator)
+    log.roll()
+
+    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, log.activeSegment.baseOffset), deleteHorizonMs = Long.MaxValue)
+    assertEquals(List(2), LogTest.keysInLog(log))
+    assertEquals(List(3, 4), offsetsInLog(log))
+  }
+
   @Test
   def testCommitMarkerRetentionWithEmptyBatch(): Unit = {
     val tp = new TopicPartition("test", 0)
@@ -309,42 +414,145 @@ class LogCleanerTest extends JUnitSuite {
     val log = makeLog(config = LogConfig.fromProps(logConfig.originals, logProps))
 
     val producerEpoch = 0.toShort
-    val producerId = 1L
-    val appendProducer = appendTransactionalAsLeader(log, producerId, producerEpoch)
+    val producer1 = appendTransactionalAsLeader(log, 1L, producerEpoch)
+    val producer2 = appendTransactionalAsLeader(log, 2L, producerEpoch)
 
-    appendProducer(Seq(2, 3)) // batch last offset is 1
-    log.appendAsLeader(commitMarker(producerId, producerEpoch), leaderEpoch = 0, isFromClient = false)
+    // [{Producer1: 2, 3}]
+    producer1(Seq(2, 3)) // offsets 0, 1
     log.roll()
 
-    log.appendAsLeader(record(2, 2), leaderEpoch = 0)
-    log.appendAsLeader(record(3, 3), leaderEpoch = 0)
+    // [{Producer1: 2, 3}], [{Producer2: 2, 3}, {Producer2: Commit}]
+    producer2(Seq(2, 3)) // offsets 2, 3
+    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,
+      origin = AppendOrigin.Coordinator) // offset 7
     log.roll()
 
     // first time through the records are removed
-    var dirtyOffset = cleaner.doClean(LogToClean(tp, log, 0L, 100L), deleteHorizonMs = Long.MaxValue)._1
-    assertEquals(List(2, 3), keysInLog(log))
-    assertEquals(List(2, 3, 4), offsetsInLog(log)) // commit marker is retained
-    assertEquals(List(1, 2, 3, 4), lastOffsetsPerBatchInLog(log)) // empty batch is retained
+    // 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, 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
-    dirtyOffset = cleaner.doClean(LogToClean(tp, log, dirtyOffset, 100L), deleteHorizonMs = Long.MaxValue)._1
-    assertEquals(List(2, 3), keysInLog(log))
-    assertEquals(List(2, 3, 4), offsetsInLog(log)) // commit marker is still retained
-    assertEquals(List(1, 2, 3, 4), lastOffsetsPerBatchInLog(log)) // empty batch is retained
+    // 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, 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
-    appendProducer(Seq(1))
+    // [{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,
+      origin = AppendOrigin.Coordinator) // offset 9
     log.roll()
 
-    dirtyOffset = cleaner.doClean(LogToClean(tp, log, dirtyOffset, 100L), deleteHorizonMs = Long.MaxValue)._1
-    assertEquals(List(2, 3, 1), keysInLog(log))
-    assertEquals(List(2, 3, 4, 5), offsetsInLog(log)) // commit marker is still retained
-    assertEquals(List(2, 3, 4, 5), lastOffsetsPerBatchInLog(log)) // empty batch should be gone
+    // 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, 7, 8, 9), offsetsInLog(log))
+    assertEquals(List(1, 4, 5, 6, 7, 8, 9), lastOffsetsPerBatchInLog(log))
+
+    // 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, 7, 8, 9), offsetsInLog(log))
+    assertEquals(List(1, 5, 6, 7, 8, 9), lastOffsetsPerBatchInLog(log))
+  }
 
-    dirtyOffset = cleaner.doClean(LogToClean(tp, log, dirtyOffset, 100L), deleteHorizonMs = Long.MaxValue)._1
-    assertEquals(List(2, 3, 1), keysInLog(log))
-    assertEquals(List(3, 4, 5), offsetsInLog(log)) // commit marker is gone
-    assertEquals(List(3, 4, 5), lastOffsetsPerBatchInLog(log)) // empty batch is gone
+  @Test
+  def testCleanEmptyControlBatch(): Unit = {
+    val tp = new TopicPartition("test", 0)
+    val cleaner = makeCleaner(Int.MaxValue)
+    val logProps = new Properties()
+    logProps.put(LogConfig.SegmentBytesProp, 256: java.lang.Integer)
+    val log = makeLog(config = LogConfig.fromProps(logConfig.originals, logProps))
+
+    val producerEpoch = 0.toShort
+
+    // [{Producer1: Commit}, {2}, {3}]
+    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, 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, 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
@@ -361,22 +569,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
-    assertEquals(List(3), keysInLog(log))
+    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)
-    assertEquals(List(3), keysInLog(log))
+    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)
@@ -390,7 +639,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 = {
@@ -404,16 +653,16 @@ 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(), keysInLog(log))
+    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(), keysInLog(log))
+    assertEquals(List(), LogTest.keysInLog(log))
     assertEquals(List(2), offsetsInLog(log)) // abort marker is still retained
     assertEquals(List(1, 2), lastOffsetsPerBatchInLog(log)) // empty batch is retained
 
@@ -421,14 +670,14 @@ 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), keysInLog(log))
+    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
-    assertEquals(List(1), keysInLog(log))
+    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))
 
@@ -440,7 +689,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)
@@ -452,19 +701,91 @@ class LogCleanerTest extends JUnitSuite {
 
     while(log.numberOfSegments < 2)
       log.appendAsLeader(record(log.logEndOffset.toInt, Array.fill(largeMessageSize)(0: Byte)), leaderEpoch = 0)
-    val keysFound = keysInLog(log)
+    val keysFound = LogTest.keysInLog(log)
     assertEquals(0L until log.logEndOffset, keysFound)
 
     // pretend we have the following keys
-    val keys = immutable.ListSet(1, 3, 5, 7, 9)
+    val keys = immutable.ListSet(1L, 3L, 5L, 7L, 9L)
     val map = new FakeOffsetMap(Int.MaxValue)
     keys.foreach(k => map.put(key(k), Long.MaxValue))
 
     // clean the log
     val stats = new CleanerStats()
-    cleaner.cleanSegments(log, Seq(log.logSegments.head), map, 0L, stats)
-    val shouldRemain = keysInLog(log).filter(!keys.contains(_))
-    assertEquals(shouldRemain, keysInLog(log))
+    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))
+  }
+
+  /**
+   * Test log cleaning with logs containing messages larger than topic's max message size
+   */
+  @Test
+  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, new CleanedTransactionMetadata)
+    val shouldRemain = LogTest.keysInLog(log).filter(k => !offsetMap.map.containsKey(k.toString))
+    assertEquals(shouldRemain, LogTest.keysInLog(log))
+  }
+
+  /**
+   * Test log cleaning with logs containing messages larger than topic's max message size
+   * where header is corrupt
+   */
+  @Test
+  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)
+    file.write(0xff)
+    file.close()
+
+    val cleaner = makeCleaner(Int.MaxValue, maxMessageSize=1024)
+    intercept[CorruptRecordException] {
+      cleaner.cleanSegments(log, Seq(log.logSegments.head), offsetMap, 0L, new CleanerStats, new CleanedTransactionMetadata)
+    }
+  }
+
+  /**
+   * Test log cleaning with logs containing messages larger than topic's max message size
+   * where message size is corrupt and larger than bytes available in log segment.
+   */
+  @Test
+  def testCorruptMessageSizeLargerThanBytesAvailable(): Unit = {
+    val (log, offsetMap) = createLogWithMessagesLargerThanMaxSize(largeMessageSize = 1024 * 1024)
+    val file = new RandomAccessFile(log.logSegments.head.log.file, "rw")
+    file.setLength(1024)
+    file.close()
+
+    val cleaner = makeCleaner(Int.MaxValue, maxMessageSize=1024)
+    intercept[CorruptRecordException] {
+      cleaner.cleanSegments(log, Seq(log.logSegments.head), offsetMap, 0L, new CleanerStats, new CleanedTransactionMetadata)
+    }
+  }
+
+  def createLogWithMessagesLargerThanMaxSize(largeMessageSize: Int): (Log, FakeOffsetMap) = {
+    val logProps = new Properties()
+    logProps.put(LogConfig.SegmentBytesProp, largeMessageSize * 16: java.lang.Integer)
+    logProps.put(LogConfig.MaxMessageBytesProp, largeMessageSize * 2: java.lang.Integer)
+
+    val log = makeLog(config = LogConfig.fromProps(logConfig.originals, logProps))
+
+    while(log.numberOfSegments < 2)
+      log.appendAsLeader(record(log.logEndOffset.toInt, Array.fill(largeMessageSize)(0: Byte)), leaderEpoch = 0)
+    val keysFound = LogTest.keysInLog(log)
+    assertEquals(0L until log.logEndOffset, keysFound)
+
+    // Decrease the log's max message size
+    logProps.put(LogConfig.MaxMessageBytesProp, largeMessageSize / 2: java.lang.Integer)
+    log.config = LogConfig.fromProps(logConfig.originals, logProps)
+
+    // pretend we have the following keys
+    val keys = immutable.ListSet(1, 3, 5, 7, 9)
+    val map = new FakeOffsetMap(Int.MaxValue)
+    keys.foreach(k => map.put(key(k), Long.MaxValue))
+
+    (log, map)
   }
 
   @Test
@@ -489,14 +810,15 @@ class LogCleanerTest extends JUnitSuite {
       log.appendAsLeader(record(log.logEndOffset.toInt, log.logEndOffset.toInt), leaderEpoch = 0)
 
     cleaner.clean(LogToClean(new TopicPartition("test", 0), log, 0, log.activeSegment.baseOffset))
-    val keys = keysInLog(log).toSet
+    val keys = LogTest.keysInLog(log).toSet
     assertTrue("None of the keys we deleted should still exist.",
                (0 until leo.toInt by 2).forall(!keys.contains(_)))
   }
 
+  @Test
   def testLogCleanerStats(): Unit = {
-    // because loadFactor is 0.75, this means we can fit 2 messages in the map
-    val cleaner = makeCleaner(2)
+    // because loadFactor is 0.75, this means we can fit 3 messages in the map
+    val cleaner = makeCleaner(4)
     val logProps = new Properties()
     logProps.put(LogConfig.SegmentBytesProp, 1024: java.lang.Integer)
 
@@ -541,7 +863,7 @@ class LogCleanerTest extends JUnitSuite {
     cleaner.clean(LogToClean(new TopicPartition("test", 0), log, 0L, log.activeSegment.baseOffset))
     assertEquals(List(1, 3, 4), lastOffsetsPerBatchInLog(log))
     assertEquals(Map(1L -> 0, 2L -> 1, 3L -> 0), lastSequencesInLog(log))
-    assertEquals(List(0, 1), keysInLog(log))
+    assertEquals(List(0, 1), LogTest.keysInLog(log))
     assertEquals(List(3, 4), offsetsInLog(log))
   }
 
@@ -558,13 +880,13 @@ 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))
     assertEquals(List(2, 3), lastOffsetsPerBatchInLog(log))
     assertEquals(Map(producerId -> 2), lastSequencesInLog(log))
-    assertEquals(List(), keysInLog(log))
+    assertEquals(List(), LogTest.keysInLog(log))
     assertEquals(List(3), offsetsInLog(log))
 
     // Append a new entry from the producer and verify that the empty batch is cleaned up
@@ -574,13 +896,13 @@ class LogCleanerTest extends JUnitSuite {
 
     assertEquals(List(3, 5), lastOffsetsPerBatchInLog(log))
     assertEquals(Map(producerId -> 4), lastSequencesInLog(log))
-    assertEquals(List(1, 5), keysInLog(log))
+    assertEquals(List(1, 5), LogTest.keysInLog(log))
     assertEquals(List(3, 4, 5), offsetsInLog(log))
   }
 
   @Test
   def testPartialSegmentClean(): Unit = {
-    // because loadFactor is 0.75, this means we can fit 2 messages in the map
+    // because loadFactor is 0.75, this means we can fit 1 message in the map
     val cleaner = makeCleaner(2)
     val logProps = new Properties()
     logProps.put(LogConfig.SegmentBytesProp, 1024: java.lang.Integer)
@@ -597,16 +919,16 @@ class LogCleanerTest extends JUnitSuite {
 
     // clean the log with only one message removed
     cleaner.clean(LogToClean(new TopicPartition("test", 0), log, 2, log.activeSegment.baseOffset))
-    assertEquals(List(1,0,1,0), keysInLog(log))
+    assertEquals(List(1,0,1,0), LogTest.keysInLog(log))
     assertEquals(List(1,2,3,4), offsetsInLog(log))
 
     // continue to make progress, even though we can only clean one message at a time
     cleaner.clean(LogToClean(new TopicPartition("test", 0), log, 3, log.activeSegment.baseOffset))
-    assertEquals(List(0,1,0), keysInLog(log))
+    assertEquals(List(0,1,0), LogTest.keysInLog(log))
     assertEquals(List(2,3,4), offsetsInLog(log))
 
     cleaner.clean(LogToClean(new TopicPartition("test", 0), log, 4, log.activeSegment.baseOffset))
-    assertEquals(List(1,0), keysInLog(log))
+    assertEquals(List(1,0), LogTest.keysInLog(log))
     assertEquals(List(3,4), offsetsInLog(log))
   }
 
@@ -729,14 +1051,6 @@ class LogCleanerTest extends JUnitSuite {
     assertEquals("Cleaner should have seen %d invalid messages.", numInvalidMessages, stats.invalidMessagesRead)
   }
 
-  /* extract all the keys from a log */
-  def keysInLog(log: Log): Iterable[Int] = {
-    for (segment <- log.logSegments;
-         batch <- segment.log.batches.asScala if !batch.isControlBatch;
-         record <- batch.asScala if record.hasValue && record.hasKey)
-      yield TestUtils.readString(record.key).toInt
-  }
-
   def lastOffsetsPerBatchInLog(log: Log): Iterable[Long] = {
     for (segment <- log.logSegments; batch <- segment.log.batches.asScala)
       yield batch.lastOffset
@@ -774,11 +1088,12 @@ class LogCleanerTest extends JUnitSuite {
     while(log.numberOfSegments < 4)
       log.appendAsLeader(record(log.logEndOffset.toInt, log.logEndOffset.toInt), leaderEpoch = 0)
 
-    val keys = keysInLog(log)
+    val keys = LogTest.keysInLog(log)
     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)
     }
   }
 
@@ -826,7 +1141,7 @@ class LogCleanerTest extends JUnitSuite {
     assertTrue("All but the last group should be the target size.", groups.dropRight(1).forall(_.size == groupSize))
 
     // check grouping by index size
-    val indexSize = log.logSegments.take(groupSize).map(_.index.sizeInBytes).sum + 1
+    val indexSize = log.logSegments.take(groupSize).map(_.offsetIndex.sizeInBytes).sum + 1
     groups = cleaner.groupSegmentsBySize(log.logSegments, maxSize = Int.MaxValue, maxIndexSize = indexSize, log.logEndOffset)
     checkSegmentOrder(groups)
     assertTrue("All but the last group should be the target size.", groups.dropRight(1).forall(_.size == groupSize))
@@ -856,7 +1171,7 @@ class LogCleanerTest extends JUnitSuite {
     val records = messageWithOffset("hello".getBytes, "hello".getBytes, Int.MaxValue - 1)
     log.appendAsFollower(records)
     log.appendAsLeader(TestUtils.singletonRecords(value = "hello".getBytes, key = "hello".getBytes), leaderEpoch = 0)
-    assertEquals(Int.MaxValue, log.activeSegment.index.lastOffset)
+    assertEquals(Int.MaxValue, log.activeSegment.offsetIndex.lastOffset)
 
     // grouping should result in a single group with maximum relative offset of Int.MaxValue
     var groups = cleaner.groupSegmentsBySize(log.logSegments, maxSize = Int.MaxValue, maxIndexSize = Int.MaxValue, log.logEndOffset)
@@ -877,7 +1192,7 @@ class LogCleanerTest extends JUnitSuite {
     groups = cleaner.groupSegmentsBySize(log.logSegments, maxSize = Int.MaxValue, maxIndexSize = Int.MaxValue, log.logEndOffset)
     assertEquals(log.numberOfSegments - 1, groups.size)
     for (group <- groups)
-      assertTrue("Relative offset greater than Int.MaxValue", group.last.index.lastOffset - group.head.index.baseOffset <= Int.MaxValue)
+      assertTrue("Relative offset greater than Int.MaxValue", group.last.offsetIndex.lastOffset - group.head.offsetIndex.baseOffset <= Int.MaxValue)
     checkSegmentOrder(groups)
   }
 
@@ -905,21 +1220,21 @@ class LogCleanerTest extends JUnitSuite {
     log.appendAsFollower(record1)
     val record2 = messageWithOffset("hello".getBytes, "hello".getBytes, 1)
     log.appendAsFollower(record2)
-    log.roll(Int.MaxValue/2) // starting a new log segment at offset Int.MaxValue/2
+    log.roll(Some(Int.MaxValue/2)) // starting a new log segment at offset Int.MaxValue/2
     val record3 = messageWithOffset("hello".getBytes, "hello".getBytes, Int.MaxValue/2)
     log.appendAsFollower(record3)
     val record4 = messageWithOffset("hello".getBytes, "hello".getBytes, Int.MaxValue.toLong + 1)
     log.appendAsFollower(record4)
 
     assertTrue("Actual offset range should be > Int.MaxValue", log.logEndOffset - 1 - log.logStartOffset > Int.MaxValue)
-    assertTrue("index.lastOffset is reporting the wrong last offset", log.logSegments.last.index.lastOffset - log.logStartOffset <= Int.MaxValue)
+    assertTrue("index.lastOffset is reporting the wrong last offset", log.logSegments.last.offsetIndex.lastOffset - log.logStartOffset <= Int.MaxValue)
 
     // grouping should result in two groups because the second segment takes the offset range > MaxInt
     val groups = cleaner.groupSegmentsBySize(log.logSegments, maxSize = Int.MaxValue, maxIndexSize = Int.MaxValue, log.logEndOffset)
     assertEquals(2, groups.size)
 
     for (group <- groups)
-      assertTrue("Relative offset greater than Int.MaxValue", group.last.nextOffset() - 1 - group.head.baseOffset <= Int.MaxValue)
+      assertTrue("Relative offset greater than Int.MaxValue", group.last.readNextOffset - 1 - group.head.baseOffset <= Int.MaxValue)
     checkSegmentOrder(groups)
   }
 
@@ -940,7 +1255,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
@@ -959,6 +1274,50 @@ class LogCleanerTest extends JUnitSuite {
     checkRange(map, segments(3).baseOffset.toInt, log.logEndOffset.toInt)
   }
 
+  @Test
+  def testSegmentWithOffsetOverflow(): Unit = {
+    val cleaner = makeCleaner(Int.MaxValue)
+    val logProps = new Properties()
+    logProps.put(LogConfig.IndexIntervalBytesProp, 1: java.lang.Integer)
+    logProps.put(LogConfig.FileDeleteDelayMsProp, 1000: java.lang.Integer)
+    val config = LogConfig.fromProps(logConfig.originals, logProps)
+
+    LogTest.initializeLogDirWithOverflowedSegment(dir)
+
+    val log = makeLog(config = config, recoveryPoint = Long.MaxValue)
+    val segmentWithOverflow = LogTest.firstOverflowSegment(log).getOrElse {
+      fail("Failed to create log with a segment which has overflowed offsets")
+    }
+
+    val numSegmentsInitial = log.logSegments.size
+    val allKeys = LogTest.keysInLog(log).toList
+    val expectedKeysAfterCleaning = new mutable.ArrayBuffer[Long]()
+
+    // pretend we want to clean every alternate key
+    val offsetMap = new FakeOffsetMap(Int.MaxValue)
+    for (k <- 1 until allKeys.size by 2) {
+      expectedKeysAfterCleaning += allKeys(k - 1)
+      offsetMap.put(key(allKeys(k)), Long.MaxValue)
+    }
+
+    // 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(),
+        new CleanedTransactionMetadata)
+    }
+    assertEquals(numSegmentsInitial + 1, log.logSegments.size)
+    assertEquals(allKeys, LogTest.keysInLog(log))
+    assertFalse(LogTest.hasOffsetOverflow(log))
+
+    // Clean each segment now that split is complete.
+    for (segmentToClean <- log.logSegments)
+      cleaner.cleanSegments(log, List(segmentToClean), offsetMap, 0L, new CleanerStats(),
+        new CleanedTransactionMetadata)
+    assertEquals(expectedKeysAfterCleaning, LogTest.keysInLog(log))
+    assertFalse(LogTest.hasOffsetOverflow(log))
+    log.close()
+  }
+
   /**
    * Tests recovery if broker crashes at the following stages during the cleaning sequence
    * 
                        
@@ -978,28 +1337,14 @@ class LogCleanerTest extends JUnitSuite {
 
     val config = LogConfig.fromProps(logConfig.originals, logProps)
 
-    def recoverAndCheck(config: LogConfig, expectedKeys : Iterable[Int]) : Log = {
-      // Recover log file and check that after recovery, keys are as expected
-      // and all temporary files have been deleted
-      val recoveredLog = makeLog(config = config)
-      time.sleep(config.fileDeleteDelayMs + 1)
-      for (file <- dir.listFiles) {
-        assertFalse("Unexpected .deleted file after recovery", file.getName.endsWith(Log.DeletedFileSuffix))
-        assertFalse("Unexpected .cleaned file after recovery", file.getName.endsWith(Log.CleanedFileSuffix))
-        assertFalse("Unexpected .swap file after recovery", file.getName.endsWith(Log.SwapFileSuffix))
-      }
-      assertEquals(expectedKeys, keysInLog(recoveredLog))
-      recoveredLog
-    }
-
-    // create a log and append some messages
+   // create a log and append some messages
     var log = makeLog(config = config)
     var messageCount = 0
-    while(log.numberOfSegments < 10) {
+    while (log.numberOfSegments < 10) {
       log.appendAsLeader(record(log.logEndOffset.toInt, log.logEndOffset.toInt), leaderEpoch = 0)
       messageCount += 1
     }
-    val allKeys = keysInLog(log)
+    val allKeys = LogTest.keysInLog(log)
 
     // pretend we have odd-numbered keys
     val offsetMap = new FakeOffsetMap(Int.MaxValue)
@@ -1007,8 +1352,12 @@ 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())
-    var cleanedKeys = keysInLog(log)
+    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)
+    log.close()
 
     // 1) Simulate recovery just after .cleaned file is created, before rename to .swap
     //    On recovery, clean operation is aborted. All messages should be present in the log
@@ -1019,8 +1368,12 @@ class LogCleanerTest extends JUnitSuite {
     log = recoverAndCheck(config, allKeys)
 
     // clean again
-    cleaner.cleanSegments(log, log.logSegments.take(9).toSeq, offsetMap, 0L, new CleanerStats())
-    cleanedKeys = keysInLog(log)
+    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)
+    log.close()
 
     // 2) Simulate recovery just after swap file is created, before old segment files are
     //    renamed to .deleted. Clean operation is resumed during recovery.
@@ -1031,14 +1384,17 @@ class LogCleanerTest extends JUnitSuite {
     log = recoverAndCheck(config, cleanedKeys)
 
     // add some more messages and clean the log again
-    while(log.numberOfSegments < 10) {
+    while (log.numberOfSegments < 10) {
       log.appendAsLeader(record(log.logEndOffset.toInt, log.logEndOffset.toInt), leaderEpoch = 0)
       messageCount += 1
     }
     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())
-    cleanedKeys = keysInLog(log)
+    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)
 
     // 3) Simulate recovery after swap file is created and old segments files are renamed
     //    to .deleted. Clean operation is resumed during recovery.
@@ -1046,18 +1402,23 @@ class LogCleanerTest extends JUnitSuite {
     log = recoverAndCheck(config, cleanedKeys)
 
     // add some more messages and clean the log again
-    while(log.numberOfSegments < 10) {
+    while (log.numberOfSegments < 10) {
       log.appendAsLeader(record(log.logEndOffset.toInt, log.logEndOffset.toInt), leaderEpoch = 0)
       messageCount += 1
     }
     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())
-    cleanedKeys = keysInLog(log)
+    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)
+    log.close()
 
     // 4) Simulate recovery after swap is complete, but async deletion
     //    is not yet complete. Clean operation is resumed during recovery.
-    recoverAndCheck(config, cleanedKeys)
+    log = recoverAndCheck(config, cleanedKeys)
+    log.close()
   }
 
   @Test
@@ -1070,16 +1431,17 @@ class LogCleanerTest extends JUnitSuite {
     val logConfig = LogConfig(logProps)
     val log = makeLog(config = logConfig)
     val cleaner = makeCleaner(Int.MaxValue)
-    val start = 0
-    val end = 2
-    val offsetSeq = Seq(0L, 7206178L)
-    writeToLog(log, (start until end) zip (start until end), offsetSeq)
-    cleaner.buildOffsetMap(log, start, end, map, new CleanerStats())
-    val endOffset = map.latestOffset
-    assertEquals("Last offset should be the end offset.", 7206178L, endOffset)
-    assertEquals("Should have the expected number of messages in the map.", end - start, map.size)
+    val keyStart = 0
+    val keyEnd = 2
+    val offsetStart = 0L
+    val offsetEnd = 7206178L
+    val offsetSeq = Seq(offsetStart, offsetEnd)
+    writeToLog(log, (keyStart until keyEnd) zip (keyStart until keyEnd), offsetSeq)
+    cleaner.buildOffsetMap(log, keyStart, offsetEnd + 1L, map, new CleanerStats())
+    assertEquals("Last offset should be the end offset.", offsetEnd, map.latestOffset)
+    assertEquals("Should have the expected number of messages in the map.", keyEnd - keyStart, map.size)
     assertEquals("Map should contain first value", 0L, map.get(key(0)))
-    assertEquals("Map should contain second value", 7206178L, map.get(key(1)))
+    assertEquals("Map should contain second value", offsetEnd, map.get(key(1)))
   }
 
   /**
@@ -1088,9 +1450,9 @@ class LogCleanerTest extends JUnitSuite {
   @Test
   def testBuildPartialOffsetMap(): Unit = {
     // because loadFactor is 0.75, this means we can fit 2 messages in the map
-    val map = new FakeOffsetMap(3)
     val log = makeLog()
-    val cleaner = makeCleaner(2)
+    val cleaner = makeCleaner(3)
+    val map = cleaner.offsetMap
 
     log.appendAsLeader(record(0,0), leaderEpoch = 0)
     log.appendAsLeader(record(1,1), leaderEpoch = 0)
@@ -1113,7 +1475,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()
@@ -1156,8 +1518,6 @@ class LogCleanerTest extends JUnitSuite {
    */
   @Test
   def testClientHandlingOfCorruptMessageSet(): Unit = {
-    import JavaConverters._
-
     val keys = 1 until 10
     val offset = 50
     val set = keys zip (offset until offset + keys.size)
@@ -1201,6 +1561,71 @@ 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)
+    }
+  }
+
+  @Test
+  def testMaxCleanTimeSecs(): Unit = {
+    val logCleaner = new LogCleaner(new CleanerConfig,
+      logDirs = Array(TestUtils.tempDir()),
+      logs = new Pool[TopicPartition, Log](),
+      logDirFailureChannel = new LogDirFailureChannel(1),
+      time = time)
+
+    def checkGauge(name: String): Unit = {
+      val gauge = logCleaner.newGauge(name, () => 999)
+      // if there is no cleaners, 0 is default value
+      assertEquals(0, gauge.value())
+    }
+
+    try {
+      checkGauge("max-buffer-utilization-percent")
+      checkGauge("max-clean-time-secs")
+      checkGauge("max-compaction-delay-secs")
+    } finally logCleaner.shutdown()
+  }
+
+
   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
@@ -1242,9 +1667,7 @@ class LogCleanerTest extends JUnitSuite {
       producerIdExpirationCheckIntervalMs = LogManager.ProducerIdExpirationCheckIntervalMs,
       logDirFailureChannel = new LogDirFailureChannel(10))
 
-  private def noOpCheckDone(topicPartition: TopicPartition) { /* do nothing */  }
-
-  private def makeCleaner(capacity: Int, checkDone: TopicPartition => Unit = noOpCheckDone, maxMessageSize: Int = 64*1024) =
+  private def makeCleaner(capacity: Int, checkDone: TopicPartition => Unit = _ => (), maxMessageSize: Int = 64*1024) =
     new Cleaner(id = 0,
                 offsetMap = new FakeOffsetMap(capacity),
                 ioBufferSize = maxMessageSize,
@@ -1255,11 +1678,10 @@ class LogCleanerTest extends JUnitSuite {
                 checkDone = checkDone)
 
   private def writeToLog(log: Log, seq: Iterable[(Int, Int)]): Iterable[Long] = {
-    for((key, value) <- seq)
-      yield log.appendAsLeader(record(key, value), leaderEpoch = 0).firstOffset
+    for ((key, value) <- seq) yield log.appendAsLeader(record(key, value), leaderEpoch = 0).firstOffset.get
   }
 
-  private def key(id: Int) = ByteBuffer.wrap(id.toString.getBytes)
+  private def key(id: Long) = ByteBuffer.wrap(id.toString.getBytes)
 
   private def record(key: Int, value: Int,
              producerId: Long = RecordBatch.NO_PRODUCER_ID,
@@ -1270,13 +1692,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 =>
@@ -1284,11 +1713,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)
     }
   }
 
@@ -1313,6 +1742,9 @@ class LogCleanerTest extends JUnitSuite {
 
   private def tombstoneRecord(key: Int): MemoryRecords = record(key, null)
 
+  private def recoverAndCheck(config: LogConfig, expectedKeys: Iterable[Long]): Log = {
+    LogTest.recoverAndCheck(dir, config, expectedKeys, new BrokerTopicStats(), time, time.scheduler)
+  }
 }
 
 class FakeOffsetMap(val slots: Int) extends OffsetMap {
@@ -1320,7 +1752,7 @@ class FakeOffsetMap(val slots: Int) extends OffsetMap {
   var lastOffset = -1L
 
   private def keyFor(key: ByteBuffer) =
-    new String(Utils.readBytes(key.duplicate), "UTF-8")
+    new String(Utils.readBytes(key.duplicate), StandardCharsets.UTF_8)
 
   override def put(key: ByteBuffer, offset: Long): Unit = {
     lastOffset = offset
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..0a9e71f6febae
--- /dev/null
+++ b/core/src/test/scala/unit/kafka/log/LogConcurrencyTest.scala
@@ -0,0 +1,183 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES 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.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().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.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..c3d8c2db298ac 100644
--- a/core/src/test/scala/unit/kafka/log/LogConfigTest.scala
+++ b/core/src/test/scala/unit/kafka/log/LogConfigTest.scala
@@ -17,26 +17,28 @@
 
 package kafka.log
 
-import java.util.Properties
+import java.util.{Collections, 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(""))
@@ -101,6 +113,56 @@ class LogConfigTest {
     assertFalse(isValid("100 :0,10:   "))
     assertFalse(isValid("100: 0,10:   "))
     assertFalse(isValid("100:0,10 :   "))
+    assertFalse(isValid("*,100:10"))
+    assertFalse(isValid("* ,100: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